diff --git a/.agents/skills/model-microbenchmark/SKILL.md b/.agents/skills/model-microbenchmark/SKILL.md
index 6da792aba4de..b1bfabeb578f 100644
--- a/.agents/skills/model-microbenchmark/SKILL.md
+++ b/.agents/skills/model-microbenchmark/SKILL.md
@@ -45,7 +45,8 @@ Key flags (plus all `EngineArgs`):
 
 - `--kv-len-jitter 0.5` (default): prefixes drawn within ±50%, same on every
   DP rank, capped at `max_model_len`. Use `0` for exact lengths.
-- `--routing auto|real|balanced`: `auto` = balanced with dummy weights.
+- `--routing auto|real|balanced|zipf`: `auto` = Zipf with dummy weights;
+  `zipf` uses fixed alpha=1.25 skew (target only, no PCP or grouped top-k).
 - `--iters 20 --samples 5`, `--output-json out.json` (per-rank latency stats,
   `samples_us`, graph-vs-eager check).
 - Defaults: `--load-format dummy`, prefix caching off, no torch.compile.
@@ -95,7 +96,7 @@ that the JSON summary excludes.
 ## Caveats
 
 - **Dummy weights lie for data-dependent kernels.** Dummy gates route every
-  token to the same experts (hence balanced routing). Cross-check suspicious
+  token to the same experts (hence synthetic routing). Cross-check suspicious
   kernels with real weights.
 - **torch.compile models (e.g. Qwen3/3.5) look slow**: no compile here. Use
   them for correctness only; trust perf for models under `vllm/models/`.
diff --git a/.agents/skills/model-recipes/SKILL.md b/.agents/skills/model-recipes/SKILL.md
index f6ee22b1946f..5100b3938a00 100644
--- a/.agents/skills/model-recipes/SKILL.md
+++ b/.agents/skills/model-recipes/SKILL.md
@@ -7,9 +7,9 @@ description: Find vLLM deployment recipes and configs for LLMs and Hugging Face
 
 Model serving recipes are maintained in this repository's
 [recipes/models/ directory](../../../recipes/models). Omni (image, video,
-audio) recipes stay in vllm-project/recipes for now (moving to vllm-omni),
-and Ascend NPU content is inline in these recipes (moving to vllm-ascend
-overlays). The published catalog, rendered by
+audio) recipes live in vllm-project/vllm-omni, and Ascend NPU content lives in
+vllm-project/vllm-ascend as overlays merged into these recipes. The published
+catalog, rendered by
 [vllm-project/recipes](https://github.com/vllm-project/recipes), is at
 [recipes.vllm.ai](https://recipes.vllm.ai).
 
diff --git a/.buildkite/amd-disagg/cluster.sh b/.buildkite/amd-disagg/cluster.sh
index a8906e6e578b..d71c9eb410df 100644
--- a/.buildkite/amd-disagg/cluster.sh
+++ b/.buildkite/amd-disagg/cluster.sh
@@ -156,6 +156,21 @@ export MORI_GPU_ARCHS="gfx950"
 # path fails to come up.
 export MORI_ENABLE_HOST_PROXY="${MORI_ENABLE_HOST_PROXY:-1}"
 
+# NIXL UCX defaults: RC on ionic_0, GPU-direct RDMA.
+if [[ "${KV_CONNECTOR:-}" == Nixl* ]]; then
+    export UCX_TLS="${UCX_TLS:-rc_verbs,ud_verbs,tcp,self,rocm}"
+    export UCX_NET_DEVICES="${UCX_NET_DEVICES:-ionic_0:1,lo}"
+    export UCX_IB_GID_INDEX="${UCX_IB_GID_INDEX:-${_IB_GID_INDEX}}"
+    export UCX_UD_VERBS_ETH_DGID_CHECK="${UCX_UD_VERBS_ETH_DGID_CHECK:-n}"
+    export UCX_IB_GPU_DIRECT_RDMA="${UCX_IB_GPU_DIRECT_RDMA:-yes}"
+    export UCX_RMA_PPLN_ENABLE="${UCX_RMA_PPLN_ENABLE:-y}"
+    export UCX_PROTO_EMULATION_ENABLE="${UCX_PROTO_EMULATION_ENABLE:-n}"
+    export UCX_KEEPALIVE_INTERVAL="${UCX_KEEPALIVE_INTERVAL:-inf}"
+    export UCX_UD_VERBS_RX_QUEUE_LEN="${UCX_UD_VERBS_RX_QUEUE_LEN:-1024}"
+    export UCX_RCACHE_MAX_UNRELEASED="${UCX_RCACHE_MAX_UNRELEASED:-1024}"
+    export ACCURACY_NUM_CONCURRENT="${ACCURACY_NUM_CONCURRENT:-1}"
+fi
+
 # ----------------------------------------------------------------- benchmark
 export BENCHMARK_COMBINATIONS="${BENCHMARK_COMBINATIONS:-1024/128 2048/128}"
 export BENCHMARK_CON="${BENCHMARK_CON:-32 64}"
diff --git a/.buildkite/amd-disagg/pipeline-disagg.yaml b/.buildkite/amd-disagg/pipeline-disagg.yaml
index 03e6d0eeccc4..71c7b8387d86 100644
--- a/.buildkite/amd-disagg/pipeline-disagg.yaml
+++ b/.buildkite/amd-disagg/pipeline-disagg.yaml
@@ -414,3 +414,58 @@ steps:
   commands:
   - MODEL_NAME=GLM-5.2-FP8 xP=2 yD=2 NODES=4 GPUS_PER_NODE=8 WIDE_EP_MODE=1 MORIIO_READ_MODE=0 RUN_AFTER_HEALTH=accuracy ROUTER_TYPE=vllm-router WAIT=1 bash .buildkite/amd-disagg/run-slurm-disagg-test.sh
 
+#--------------------------------------  mi350 · disagg NIXL (UCX RDMA) · SLURM  --------------------------------------#
+
+- label: "DeepSeek-V3-PD-1P1D-TP8-NIXL-vllm-router"
+  agents:
+    queue: amd_mi350_ainic
+  timeout_in_minutes: 270
+  concurrency: 2
+  concurrency_group: "amd-disagg/mi350/2node"
+  commands:
+  - MODEL_NAME=DeepSeek-V3 NODES=2 GPUS_PER_NODE=8 WIDE_EP_MODE=0 KV_CONNECTOR=NixlConnector RUN_AFTER_HEALTH=accuracy ROUTER_TYPE=vllm-router WAIT=1 bash .buildkite/amd-disagg/run-slurm-disagg-test.sh
+
+- label: "DeepSeek-V3-PD-2P2D-EP16/DP16-WideEP-NIXL-vllm-router"
+  agents:
+    queue: amd_mi350_ainic
+  timeout_in_minutes: 270
+  concurrency: 1
+  concurrency_group: "amd-disagg/mi350/4node"
+  commands:
+  - MODEL_NAME=DeepSeek-V3 xP=2 yD=2 NODES=4 GPUS_PER_NODE=8 WIDE_EP_MODE=1 KV_CONNECTOR=NixlConnector RUN_AFTER_HEALTH=accuracy ROUTER_TYPE=vllm-router WAIT=1 bash .buildkite/amd-disagg/run-slurm-disagg-test.sh
+
+- label: "DeepSeek-V3-PD-1P1D-EP8/DP8-WideEP-NIXL-vllm-router"
+  agents:
+    queue: amd_mi350_ainic
+  timeout_in_minutes: 270
+  concurrency: 2
+  concurrency_group: "amd-disagg/mi350/2node"
+  commands:
+  - MODEL_NAME=DeepSeek-V3 NODES=2 GPUS_PER_NODE=8 WIDE_EP_MODE=1 KV_CONNECTOR=NixlConnector RUN_AFTER_HEALTH=accuracy ROUTER_TYPE=vllm-router WAIT=1 bash .buildkite/amd-disagg/run-slurm-disagg-test.sh
+
+- label: "DeepSeek-V3-PD-1P1D-TP8-NIXL-PUSH-vllm-router"
+  agents:
+    queue: amd_mi350_ainic
+  timeout_in_minutes: 270
+  concurrency: 2
+  concurrency_group: "amd-disagg/mi350/2node"
+  commands:
+  - MODEL_NAME=DeepSeek-V3 NODES=2 GPUS_PER_NODE=8 WIDE_EP_MODE=0 KV_CONNECTOR=NixlPushConnector RUN_AFTER_HEALTH=accuracy ROUTER_TYPE=vllm-router WAIT=1 bash .buildkite/amd-disagg/run-slurm-disagg-test.sh
+
+- label: "DeepSeek-V3-PD-1P1D-EP8/DP8-WideEP-NIXL-PUSH-vllm-router"
+  agents:
+    queue: amd_mi350_ainic
+  timeout_in_minutes: 270
+  concurrency: 2
+  concurrency_group: "amd-disagg/mi350/2node"
+  commands:
+  - MODEL_NAME=DeepSeek-V3 NODES=2 GPUS_PER_NODE=8 WIDE_EP_MODE=1 KV_CONNECTOR=NixlPushConnector RUN_AFTER_HEALTH=accuracy ROUTER_TYPE=vllm-router WAIT=1 bash .buildkite/amd-disagg/run-slurm-disagg-test.sh
+
+- label: "DeepSeek-V3-PD-2P2D-EP16/DP16-WideEP-NIXL-PUSH-vllm-router"
+  agents:
+    queue: amd_mi350_ainic
+  timeout_in_minutes: 270
+  concurrency: 1
+  concurrency_group: "amd-disagg/mi350/4node"
+  commands:
+  - MODEL_NAME=DeepSeek-V3 xP=2 yD=2 NODES=4 GPUS_PER_NODE=8 WIDE_EP_MODE=1 KV_CONNECTOR=NixlPushConnector RUN_AFTER_HEALTH=accuracy ROUTER_TYPE=vllm-router WAIT=1 bash .buildkite/amd-disagg/run-slurm-disagg-test.sh
diff --git a/.buildkite/amd-disagg/run-slurm-disagg-test.sh b/.buildkite/amd-disagg/run-slurm-disagg-test.sh
index 641133b9ff45..720b11573159 100644
--- a/.buildkite/amd-disagg/run-slurm-disagg-test.sh
+++ b/.buildkite/amd-disagg/run-slurm-disagg-test.sh
@@ -47,6 +47,7 @@ SHARED_MOUNT="${SHARED_MOUNT:-/data}"
 LOG_ROOT="${LOG_ROOT:-${SHARED_MOUNT}/${USER:-$(whoami)}/disagg_logs}"
 DRY_RUN="${DRY_RUN:-0}"
 MORIIO_READ_MODE="${MORIIO_READ_MODE:-0}"
+KV_CONNECTOR="${KV_CONNECTOR:-MoRIIOConnector}"
 
 WAIT="${WAIT:-0}"
 
@@ -69,7 +70,7 @@ export DISAGG_SCRIPTS_DIR="${STAGED_DIR}"
 echo "[slurm-submit] staged scripts for compute nodes: ${DISAGG_SCRIPTS_DIR}" >&2
 
 export IMAGE MODEL_NAME WIDE_EP_MODE xP yD GPUS_PER_NODE RUN_AFTER_HEALTH HEALTH_TIMEOUT_S
-export SHARED_MOUNT LOG_ROOT DRY_RUN MORIIO_READ_MODE
+export SHARED_MOUNT LOG_ROOT DRY_RUN MORIIO_READ_MODE KV_CONNECTOR
 export ROUTER_TYPE ROUTER_PORT VLLM_ROUTER_IMAGE
 # Leave the image's own ionic userspace in place instead of mounting the host's.
 export SKIP_IONIC_MOUNTS
diff --git a/.buildkite/amd-disagg/run_xPyD_disagg.slurm b/.buildkite/amd-disagg/run_xPyD_disagg.slurm
index e9bf8014784f..19fbb54856b0 100644
--- a/.buildkite/amd-disagg/run_xPyD_disagg.slurm
+++ b/.buildkite/amd-disagg/run_xPyD_disagg.slurm
@@ -23,6 +23,7 @@ LOG_ROOT="${LOG_ROOT:-/data/${USER:-$(id -un)}/disagg_logs}"
 RUN_AFTER_HEALTH="${RUN_AFTER_HEALTH:-accuracy}"   # bench | accuracy
 DRY_RUN="${DRY_RUN:-0}"
 MORIIO_READ_MODE="${MORIIO_READ_MODE:-0}"   # 0 -> prefill writes (default); 1 -> decode reads KV (read_mode)
+KV_CONNECTOR="${KV_CONNECTOR:-MoRIIOConnector}"   # MoRIIOConnector | NixlConnector | NixlPushConnector
 
 # ---- router / gateway (host-side; also forwarded into the container) --------
 # vllm-router -> external `vllm/vllm-router` container started on the rank-0 node (default)
@@ -232,7 +233,7 @@ echo "[disagg-pd] LOG_PATH=$LOG_PATH"
 #   rank xP .. end   -> decode
 export IPADDRS
 
-export WIDE_EP_MODE xP yD GPUS_PER_NODE RUN_AFTER_HEALTH HEALTH_TIMEOUT_S DRY_RUN MORIIO_READ_MODE
+export WIDE_EP_MODE xP yD GPUS_PER_NODE RUN_AFTER_HEALTH HEALTH_TIMEOUT_S DRY_RUN MORIIO_READ_MODE KV_CONNECTOR
 export IMAGE SHARED_MOUNT CONTAINER_SCRIPTS DISAGG_SCRIPTS_DIR
 export CLUSTER_ENV="${CLUSTER_ENV_IN}"
 export DOCKER_CONT_NAME="vllm_disagg_${MODEL_NAME:-model}_${SLURM_JOB_ID}"
@@ -247,6 +248,21 @@ export PREFILL_PORT DECODE_PORT PROXY_PORT ROUTER_METRICS_PORT
 export SERVE_PORT RPC_PORT KV_PORT LOCAL_PING_PORT
 export PREFLIGHT PREFLIGHT_MAX_VRAM_PCT PREFLIGHT_GPU_SETTLE_S
 
+# NIXL has no discovery service, so pass the worker URLs.
+ROUTER_KV_ARGS="--kv-connector moriio --vllm-discovery-address 0.0.0.0:${PROXY_PING_PORT}"
+if [[ "${KV_CONNECTOR}" == Nixl* ]]; then
+    nixl_startup_timeout=1800
+    [[ "${WIDE_EP_MODE}" == "1" && "${xP}" -gt 1 ]] && nixl_startup_timeout=7200
+    ROUTER_KV_ARGS="--kv-connector nixl --worker-startup-timeout-secs ${nixl_startup_timeout}"
+    if [[ "${WIDE_EP_MODE}" == "1" ]]; then
+        ROUTER_KV_ARGS+=" --prefill http://${IPS[0]}:${SERVE_PORT} --decode http://${IPS[$xP]}:${SERVE_PORT}"
+    else
+        for ((i = 0; i < xP; i++)); do ROUTER_KV_ARGS+=" --prefill http://${IPS[$i]}:${PREFILL_PORT}"; done
+        for ((i = xP; i < xP + yD; i++)); do ROUTER_KV_ARGS+=" --decode http://${IPS[$i]}:${DECODE_PORT}"; done
+    fi
+fi
+export ROUTER_KV_ARGS
+
 echo "[disagg-pd] IPADDRS=${IPADDRS}"
 echo "[disagg-pd] launching ${NUM_NODES} containers (1/node) -> vllm_disagg.sh node"
 
@@ -479,7 +495,7 @@ run_node_container() {
                 -e USER="$(id -un)" \
                 -e LOGNAME="$(id -un)" \
                 "$VLLM_ROUTER_IMAGE" \
-                bash -lc "exec vllm-router --vllm-pd-disaggregation --kv-connector moriio --vllm-discovery-address 0.0.0.0:$PROXY_PING_PORT --port $ROUTER_PORT --host 0.0.0.0 --intra-node-data-parallel-size $ROUTER_DP_LOCAL --policy $ROUTER_POLICY --prefill-policy $ROUTER_POLICY --decode-policy $ROUTER_POLICY --log-level info >> $LOG_PATH/vllm_router.log 2>&1" \
+                bash -lc "exec vllm-router --vllm-pd-disaggregation $ROUTER_KV_ARGS --port $ROUTER_PORT --host 0.0.0.0 --intra-node-data-parallel-size $ROUTER_DP_LOCAL --policy $ROUTER_POLICY --prefill-policy $ROUTER_POLICY --decode-policy $ROUTER_POLICY --log-level info >> $LOG_PATH/vllm_router.log 2>&1" \
               && echo "vllm-router container started ok" \
               || echo "WARN: failed to start vllm-router container on $(hostname)"
         } 2>&1 | tee -a "$LOG_PATH/vllm_router_start.log"
@@ -528,6 +544,7 @@ run_node_container() {
         -e HEALTH_TIMEOUT_S="$HEALTH_TIMEOUT_S" \
         -e DRY_RUN="$DRY_RUN" \
         -e MORIIO_READ_MODE="$MORIIO_READ_MODE" \
+        -e KV_CONNECTOR="$KV_CONNECTOR" \
         -e ROUTER_TYPE="$ROUTER_TYPE" \
         -e ROUTER_PORT="$ROUTER_PORT" \
         -e ROUTER_POLICY="$ROUTER_POLICY" \
diff --git a/.buildkite/amd-disagg/vllm_disagg.sh b/.buildkite/amd-disagg/vllm_disagg.sh
index 23320aaf6b1b..37f8915c421e 100644
--- a/.buildkite/amd-disagg/vllm_disagg.sh
+++ b/.buildkite/amd-disagg/vllm_disagg.sh
@@ -102,6 +102,18 @@ moriio_read_mode_kv() {
     esac
 }
 
+# NIXL KV config. UCX backend.
+nixl_kv_cfg() {
+    cat <<EOF
+{
+  "kv_connector": "${KV_CONNECTOR}",
+  "kv_role": "${KV_ROLE}",
+  "kv_load_failure_policy": "recompute",
+  "kv_connector_extra_config": {"backends": ["UCX"], "kv_lease_duration": 180}
+}
+EOF
+}
+
 # ----------------------------------------------------------------- topology
 # Resolves MODEL_PATH, MODELS_YAML, the IP array, and the master/proxy addresses.
 # IP_ARRAY[0..xP-1] = prefill nodes, [xP..xP+yD-1] = decode nodes.
@@ -309,6 +321,7 @@ build_tp_cmd() {
 }
 EOF
 )
+    if [[ "${KV_CONNECTOR:-}" == Nixl* ]]; then kv_cfg="$(nixl_kv_cfg)"; fi
     CMD+=(--kv-transfer-config "${kv_cfg}")
     log "role=${ROLE} mode=tp host=${HOST_IP:-0.0.0.0}:${port} kv_role=${KV_ROLE} tp=${TP_SIZE} read_mode=${MORIIO_READ_MODE}"
 }
@@ -337,6 +350,10 @@ build_ep_cmd() {
         --distributed-timeout-seconds "${DISTRIBUTED_TIMEOUT_SECONDS}"
     )
     local _mc; read -ra _mc <<< "${MODEL_CONFIG}"; CMD+=("${_mc[@]}")
+    # 2P2D EP16 only. At 0.85 the AITER MoE sort buffer runs out of memory.
+    if [[ "${KV_CONNECTOR:-}" == Nixl* && "${ROLE}" == "prefill" && "${xP}" -gt 1 ]]; then
+        CMD+=(--gpu-memory-utilization 0.75)
+    fi
 
     if [[ "${IS_MASTER}" == "1" ]]; then
         CMD+=(--api-server-count="${GPUS_PER_NODE}")
@@ -359,6 +376,7 @@ build_ep_cmd() {
 }
 EOF
 )
+        if [[ "${KV_CONNECTOR:-}" == Nixl* ]]; then kv_cfg="$(nixl_kv_cfg)"; fi
         CMD+=(--kv-transfer-config "${kv_cfg}")
         log "role=${ROLE} mode=ep MASTER node_rank=${NODE_RANK} dp_size=${DP_GROUP_SIZE} dp_addr=${DP_MASTER_ADDR} kv_role=${KV_ROLE} all2all=${backend} read_mode=${MORIIO_READ_MODE}"
     else
diff --git a/.buildkite/ci_config_rocm.yaml b/.buildkite/ci_config_rocm.yaml
index ea2cc30904ca..f969393b6018 100644
--- a/.buildkite/ci_config_rocm.yaml
+++ b/.buildkite/ci_config_rocm.yaml
@@ -5,18 +5,17 @@ run_all_patterns:
   - ".dockerignore"
   - "docker/Dockerfile.rocm"
   - "docker/Dockerfile.rocm_base"
-  - "docker/ci-rocm.hcl"
   - "docker/docker-bake-rocm.hcl"
   - ".buildkite/hardware_tests/amd.yaml"
   - ".buildkite/ci_config_rocm.yaml"
-  - ".buildkite/scripts/ci-bake-rocm.sh"
-  - ".buildkite/scripts/rocm/"
+  - ".buildkite/scripts/rocm/build.sh"
+  - ".buildkite/scripts/rocm/select_aiter_nightly_wheel.py"
   - ".buildkite/scripts/hardware_ci/run-amd-test.sh"
+  - ".buildkite/scripts/hardware_ci/amd_ci_fast_exit.py"
   - "CMakeLists.txt"
   - "requirements/common.txt"
   - "requirements/rocm.txt"
   - "requirements/build/rust.txt"
-  - "requirements/build/rocm.txt"
   - "requirements/test/rocm.txt"
   - "pyproject.toml"
   - "setup.py"
@@ -25,6 +24,8 @@ run_all_patterns:
   - "rust-toolchain.toml"
   - "tools/build_rust.py"
   - "tools/install_torchcodec_rocm.sh"
+  - "tools/vllm-rocm/therock_wheels.py"
+  - "tools/ci/rocm_kernel_symbol_map.py"
   - "tests/vllm_test_utils/"
   - "vllm/envs.py"
   - "vllm/__init__.py"
diff --git a/.buildkite/hardware_tests/amd.yaml b/.buildkite/hardware_tests/amd.yaml
index 9649b948c83d..a2e2414fd1c5 100644
--- a/.buildkite/hardware_tests/amd.yaml
+++ b/.buildkite/hardware_tests/amd.yaml
@@ -1,83 +1,38 @@
 group: Hardware - AMD Build
 
-# ROCm image flow:
-# 1. Select the content-matching ROCm base, building only on a cache miss.
-# 2. Build ci_base from either the stable base or the newly selected base.
-# 3. Build the build-scoped ROCm CI image and smoke-test it before GPU jobs run.
+# ROCm images: .buildkite/scripts/rocm/build.sh + docker/docker-bake-rocm.hcl. ci-infra's AMD CI pipeline
+# (test-template-amd.j2) renders these same steps and sets ROCM_CI_REPO (trusted rocm/vllm-ci, else rocm/vllm-dev).
+# The base step ensures the content-keyed stack (+ -wheels), runtime and test-deps images; the build
+# step adds the per-commit test image. Base groups are per stack/repo/PR/branch.
 steps:
-  - label: ":docker: Refresh ROCm base"
+  - label: ":docker: Ensure ROCm stack + runtime"
     key: refresh-rocm-base-amd
     depends_on: []
     device: amd_cpu
     no_plugin: true
     commands:
-      - bash .buildkite/scripts/rocm/refresh-base-image.sh
+      - bash .buildkite/scripts/rocm/build.sh rocm base
+    timeout_in_minutes: 540   # a key miss rebuilds the stack (hours); a hit takes seconds
     concurrency: 1
-    concurrency_group: "vllm/rocm/base-build/$BUILDKITE_COMMIT"
-    env:
-      DOCKER_BUILDKIT: "1"
-      BUILDKIT_PROGRESS: "tty"
-      TERM: "xterm-256color"
+    concurrency_group: "vllm/rocm/base/rocm/$ROCM_CI_REPO/$BUILDKITE_PULL_REQUEST/$BUILDKITE_BRANCH"
     retry:
       automatic:
-        - exit_status: 1  # Transient Docker/BuildKit failure
-          limit: 1
-        - exit_status: -1  # Agent was lost
-          limit: 1
-        - exit_status: -10  # Agent was lost
-          limit: 1
-
-  # Ensure ci_base is up-to-date before building the test image.
-  # Compares a content hash of ci_base-affecting files against the remote
-  # image label. If hashes match the build is skipped (< 30 s); if they
-  # differ ci_base is rebuilt and pushed automatically.
-  - label: ":docker: Ensure ROCm ci_base"
-    key: ensure-ci-base-amd
-    soft_fail: false
-    depends_on:
-      - refresh-rocm-base-amd
-    device: amd_cpu
-    no_plugin: true
-    commands:
-      - bash .buildkite/scripts/rocm/build-ci-base.sh
-    concurrency: 1
-    concurrency_group: "vllm/rocm/ci-base-build/$BUILDKITE_COMMIT"
-    env:
-      DOCKER_BUILDKIT: "1"
-      BUILDKIT_PROGRESS: "tty"
-      TERM: "xterm-256color"
-      VLLM_BAKE_FILE: "docker/docker-bake-rocm.hcl"
-      PYTORCH_ROCM_ARCH: "gfx90a;gfx942;gfx950"
-      REMOTE_VLLM: "0"
-    retry:
-      automatic:
-        - exit_status: 1  # Transient Docker/BuildKit failure
-          limit: 1
-        - exit_status: 2  # Exact image found but registry handoff failed
+        - exit_status: 75  # Registry error (build.sh)
           limit: 1
         - exit_status: -1  # Agent was lost
           limit: 1
         - exit_status: -10  # Agent was lost
           limit: 1
 
-  - label: ":docker: Build ROCm test image and artifacts"
+  - label: ":docker: Build ROCm test image"
     key: image-build-amd
-    soft_fail: false
     depends_on:
-      - ensure-ci-base-amd
+      - refresh-rocm-base-amd
     device: amd_cpu
     no_plugin: true
     commands:
-      - bash .buildkite/scripts/rocm/build-test-image.sh
-      - bash .buildkite/scripts/rocm/smoke-test-image.sh
-    env:
-      DOCKER_BUILDKIT: "1"
-      BUILDKIT_PROGRESS: "tty"
-      TERM: "xterm-256color"
-      VLLM_BAKE_FILE: "docker/docker-bake-rocm.hcl"
-      PYTORCH_ROCM_ARCH: "gfx90a;gfx942;gfx950"
-      IMAGE_TAG: "rocm/vllm-ci:build-$BUILDKITE_BUILD_ID"
-      REMOTE_VLLM: "0"
+      - bash .buildkite/scripts/rocm/build.sh rocm ci
+    timeout_in_minutes: 180
     retry:
       automatic:
         - exit_status: 1  # Transient Docker/BuildKit failure
diff --git a/.buildkite/hardware_tests/cpu.yaml b/.buildkite/hardware_tests/cpu.yaml
index c085a6adb706..d4d90eeaa6eb 100644
--- a/.buildkite/hardware_tests/cpu.yaml
+++ b/.buildkite/hardware_tests/cpu.yaml
@@ -79,6 +79,7 @@ steps:
   - vllm/v1/worker/gpu/spec_decode/
   - vllm/v1/attention/backends/cpu_attn.py
   - vllm/config/speculative.py
+  - vllm/models/qwen3_5/
   - vllm/model_executor/models/qwen3_5.py
   - vllm/model_executor/models/qwen3_5_mtp.py
   - vllm/model_executor/models/llama_eagle3.py
diff --git a/.buildkite/hardware_tests/intel_xpu_ci/test-intel.yaml b/.buildkite/hardware_tests/intel_xpu_ci/test-intel.yaml
index 96ac93049694..ed98d4408512 100644
--- a/.buildkite/hardware_tests/intel_xpu_ci/test-intel.yaml
+++ b/.buildkite/hardware_tests/intel_xpu_ci/test-intel.yaml
@@ -15,29 +15,28 @@ steps:
           limit: 2
         - exit_status: -10  # Agent was lost
           limit: 2
-  # Disabled until tblib transitive dependencies are locked.
-  # - label: ":intel: (BMG) Example Test"
-  #   key: xpu-example-test
-  #   depends_on:
-  #     - image-build-xpu
-  #   timeout_in_minutes: 50
-  #   optional: true
-  #   device: intel_gpu
-  #   agent_tags:
-  #     label: production
-  #     gpu: 2+
-  #     mem: 24+
-  #   no_plugin: true
-  #   env:
-  #     REGISTRY: "public.ecr.aws/q9t5s3a7"
-  #     REPO: "vllm-ci-test-repo"
-  #   source_file_dependencies:
-  #     - .buildkite/hardware_tests/intel_xpu_ci/test-intel.yaml
-  #     - .buildkite/scripts/hardware_ci/run-intel-ci-test.sh
-  #   commands:
-  #     - >-
-  #       bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-  #       'bash .buildkite/scripts/hardware_ci/run-intel-ci-test.sh example'
+  - label: ":intel: (BMG) Example Test"
+    key: xpu-example-test
+    depends_on:
+      - image-build-xpu
+    timeout_in_minutes: 50
+    optional: true
+    device: intel_gpu
+    agent_tags:
+      label: production
+      gpu: 2+
+      mem: 24+
+    no_plugin: true
+    env:
+      REGISTRY: "public.ecr.aws/q9t5s3a7"
+      REPO: "vllm-ci-test-repo"
+    source_file_dependencies:
+      - .buildkite/hardware_tests/intel_xpu_ci/test-intel.yaml
+      - .buildkite/scripts/hardware_ci/run-intel-ci-test.sh
+    commands:
+      - >-
+        bash .buildkite/scripts/hardware_ci/run-intel-test.sh
+        'bash .buildkite/scripts/hardware_ci/run-intel-ci-test.sh example'
   - label: ":intel: (BMG) V1 test"
     key: xpu-v1-test
     depends_on:
@@ -84,29 +83,28 @@ steps:
       - >-
         bash .buildkite/scripts/hardware_ci/run-intel-test.sh
         'bash .buildkite/scripts/hardware_ci/run-intel-ci-test.sh w8a8-fp8-linear'
-  # Disabled until the av installation is locked.
-  # - label: ":intel: (BMG) Server test"
-  #   key: xpu-server-test
-  #   depends_on:
-  #     - image-build-xpu
-  #   timeout_in_minutes: 45
-  #   optional: true
-  #   device: intel_gpu
-  #   agent_tags:
-  #     label: production
-  #     gpu: 1+
-  #     mem: 16+
-  #   no_plugin: true
-  #   env:
-  #     REGISTRY: "public.ecr.aws/q9t5s3a7"
-  #     REPO: "vllm-ci-test-repo"
-  #   source_file_dependencies:
-  #     - .buildkite/hardware_tests/intel_xpu_ci/test-intel.yaml
-  #     - .buildkite/scripts/hardware_ci/run-intel-ci-test.sh
-  #   commands:
-  #     - >-
-  #       bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-  #       'bash .buildkite/scripts/hardware_ci/run-intel-ci-test.sh server'
+  - label: ":intel: (BMG) Server test"
+    key: xpu-server-test
+    depends_on:
+      - image-build-xpu
+    timeout_in_minutes: 45
+    optional: true
+    device: intel_gpu
+    agent_tags:
+      label: production
+      gpu: 1+
+      mem: 16+
+    no_plugin: true
+    env:
+      REGISTRY: "public.ecr.aws/q9t5s3a7"
+      REPO: "vllm-ci-test-repo"
+    source_file_dependencies:
+      - .buildkite/hardware_tests/intel_xpu_ci/test-intel.yaml
+      - .buildkite/scripts/hardware_ci/run-intel-ci-test.sh
+    commands:
+      - >-
+        bash .buildkite/scripts/hardware_ci/run-intel-test.sh
+        'bash .buildkite/scripts/hardware_ci/run-intel-ci-test.sh server'
   - label: ":intel: (BMG) Quantization test"
     key: xpu-quantization-test
     depends_on:
diff --git a/.buildkite/image_build/image_build.sh b/.buildkite/image_build/image_build.sh
index 6d22acd5b16c..53cb3bc044b1 100755
--- a/.buildkite/image_build/image_build.sh
+++ b/.buildkite/image_build/image_build.sh
@@ -203,6 +203,35 @@ resolve_parent_commit() {
     fi
 }
 
+resolve_merge_base_cache_commit() {
+    # PR builds: the commit the PR branched from, whose build pushed
+    # vllm-ci-test-cache:<commit>. A cache source only (see ci.hcl), unlike
+    # the VLLM_MERGE_BASE_COMMIT build arg, which would change the cache key.
+    if [[ "${BUILDKITE_PULL_REQUEST:-false}" == "false" ]]; then
+        return 0
+    fi
+    if [[ -n "${VLLM_CACHE_MERGE_BASE_COMMIT:-}" ]]; then
+        echo "Using provided VLLM_CACHE_MERGE_BASE_COMMIT: ${VLLM_CACHE_MERGE_BASE_COMMIT}"
+        return 0
+    fi
+    local base_branch="${BUILDKITE_PULL_REQUEST_BASE_BRANCH:-main}"
+    local base_ref="refs/remotes/vllm-cache-base/${base_branch}"
+    git fetch --quiet --no-tags --depth=200 origin \
+        "+refs/heads/${base_branch}:${base_ref}" 2>/dev/null || true
+    VLLM_CACHE_MERGE_BASE_COMMIT=$(git merge-base HEAD "${base_ref}" 2>/dev/null || echo "")
+    if [[ -z "${VLLM_CACHE_MERGE_BASE_COMMIT}" ]]; then
+        git fetch --quiet --no-tags --deepen=1000 origin \
+            "+refs/heads/${base_branch}:${base_ref}" 2>/dev/null || true
+        VLLM_CACHE_MERGE_BASE_COMMIT=$(git merge-base HEAD "${base_ref}" 2>/dev/null || echo "")
+    fi
+    if [[ -n "${VLLM_CACHE_MERGE_BASE_COMMIT}" ]]; then
+        echo "Computed merge base with ${base_branch} for cache fallback: ${VLLM_CACHE_MERGE_BASE_COMMIT}"
+        export VLLM_CACHE_MERGE_BASE_COMMIT
+    else
+        echo "Could not determine merge base with ${base_branch}; skipping that cache source"
+    fi
+}
+
 print_bake_config() {
     echo "--- :page_facing_up: Resolved bake configuration"
     # Write to a temp directory to avoid polluting the repo root (which is the
@@ -353,6 +382,7 @@ setup_buildx_builder
 
 resolve_parent_commit
 export PARENT_COMMIT
+resolve_merge_base_cache_commit
 
 print_bake_config
 
diff --git a/.buildkite/intel_jobs/basic_correctness_intel.yaml b/.buildkite/intel_jobs/basic_correctness_intel.yaml
index 69059228179b..da6d2ec0f41b 100644
--- a/.buildkite/intel_jobs/basic_correctness_intel.yaml
+++ b/.buildkite/intel_jobs/basic_correctness_intel.yaml
@@ -20,6 +20,8 @@ steps:
   - vllm/
   - tests/basic_correctness/memory/
   - tests/basic_correctness/cpu_offload/
+  # MRV2 model state remains in the weights pool on XPU (vllm-project/vllm#63548).
+  # Re-enable Mamba deep sleep once XPU nested memory pools are supported.
   commands:
     - >-
       bash .buildkite/scripts/hardware_ci/run-intel-test.sh
@@ -28,4 +30,5 @@ steps:
       pytest -v -s basic_correctness/cpu_offload &&
       pytest -v -s
       basic_correctness/memory
-      --deselect="tests/basic_correctness/memory/cumem/test_cumem.py::test_basic_cumem"'
+      --deselect="tests/basic_correctness/memory/cumem/test_cumem.py::test_basic_cumem"
+      --deselect="tests/basic_correctness/memory/sleep_mode/test_sleep_mode.py::test_deep_sleep[yujiepan/mamba2-codestral-v0.1-tiny-random]"'
diff --git a/.buildkite/intel_jobs/engine_intel.yaml b/.buildkite/intel_jobs/engine_intel.yaml
index fce83b7dec82..7c7c0f02c7e6 100644
--- a/.buildkite/intel_jobs/engine_intel.yaml
+++ b/.buildkite/intel_jobs/engine_intel.yaml
@@ -96,12 +96,15 @@ steps:
     - vllm/model_executor/models/qwen3_next
     - vllm/config/cache.py
     - vllm/platforms/xpu.py
+    - vllm/v1/engine/
     - tests/v1/e2e/general/test_mamba_prefix_cache.py
+    - tests/v1/e2e/general/test_streaming_input.py
   commands:
     - >-
       bash .buildkite/scripts/hardware_ci/run-intel-test.sh
       'cd tests &&
-      pytest -v -s v1/e2e/general/test_mamba_prefix_cache.py'
+      pytest -v -s v1/e2e/general/test_mamba_prefix_cache.py &&
+      pytest -v -s v1/e2e/general/test_streaming_input.py'
 
 - label: ":intel: (BMG) V1 e2e (2 GPUs)"
   key: v1-e2e-2-gpus
diff --git a/.buildkite/intel_jobs/kernels_intel.yaml b/.buildkite/intel_jobs/kernels_intel.yaml
index f3a9d58f4357..fdb087971991 100644
--- a/.buildkite/intel_jobs/kernels_intel.yaml
+++ b/.buildkite/intel_jobs/kernels_intel.yaml
@@ -48,3 +48,29 @@ steps:
       bash .buildkite/scripts/hardware_ci/run-intel-test.sh
       'cd tests &&
       pytest -v -s kernels/core/test_uva.py'
+
+- label: ":intel: (BMG) XPU Sparse MLA Tests"
+  key: xpu-sparse-mla-tests
+  timeout_in_minutes: 20
+  device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    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/attention/backends/mla/xpu_mla_sparse.py
+    - tests/kernels/attention/test_xpu_mla_sparse.py
+    - tests/v1/attention/test_xpu_mla_sparse_backend.py
+  commands:
+    - >-
+      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
+      'cd tests &&
+      pytest -v -s
+      kernels/attention/test_xpu_mla_sparse.py
+      v1/attention/test_xpu_mla_sparse_backend.py'
diff --git a/.buildkite/intel_jobs/lora_intel.yaml b/.buildkite/intel_jobs/lora_intel.yaml
index 366aaf7109aa..3ff650c99752 100644
--- a/.buildkite/intel_jobs/lora_intel.yaml
+++ b/.buildkite/intel_jobs/lora_intel.yaml
@@ -57,7 +57,7 @@ steps:
       bash .buildkite/scripts/hardware_ci/run-intel-test.sh
       'cd tests &&
       export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
-      pytest -v -s lora/test_fused_moe_lora_kernel.py && 
+      pytest -v -s lora/test_fused_moe_lora_kernel.py --deselect="lora/test_fused_moe_lora_kernel.py::test_fused_moe_lora_kernel_large_rank_fallback[xpu:0-dtype0-1-True-256]" &&
       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]"'
diff --git a/.buildkite/lm-eval-harness/test_lm_eval_correctness.py b/.buildkite/lm-eval-harness/test_lm_eval_correctness.py
index 6b2e293ba981..a3c8b648d777 100644
--- a/.buildkite/lm-eval-harness/test_lm_eval_correctness.py
+++ b/.buildkite/lm-eval-harness/test_lm_eval_correctness.py
@@ -21,6 +21,7 @@
 from lm_eval.utils import simple_parse_args_string
 
 from vllm import LLM
+from vllm.config import CacheConfig
 from vllm.distributed import cleanup_dist_env_and_memory
 from vllm.logger import init_logger
 from vllm.model_executor.model_loader.weight_utils import drop_checkpoint_cache
@@ -89,6 +90,7 @@ def _shutdown_lm(lm) -> None:
         return
     gpu_memory_utilization = (
         llm.llm_engine.vllm_config.cache_config.gpu_memory_utilization
+        or CacheConfig.DEFAULT_GPU_MEMORY_UTILIZATION
     )
     try:
         llm.llm_engine.engine_core.shutdown(timeout=ROCM_ENGINE_SHUTDOWN_TIMEOUT_S)
diff --git a/.buildkite/release-pipeline.yaml b/.buildkite/release-pipeline.yaml
index 2de3a4bfa747..d966f3ed65d8 100644
--- a/.buildkite/release-pipeline.yaml
+++ b/.buildkite/release-pipeline.yaml
@@ -272,10 +272,6 @@ steps:
           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"
-          # TODO(arpera):
-          # Remove BUILD_NIXL=true once NIXL has a wheel built against torch 2.15
-          # For details: https://github.com/vllm-project/vllm/pull/60090
-          # Tracking issue report: https://github.com/ai-dynamo/nixl/issues/2297
           - |
             DOCKER_BUILDKIT=1 docker buildx build --load \
               $(bash .buildkite/scripts/docker-build-metadata-args.sh cu134) \
@@ -287,7 +283,6 @@ steps:
               --build-arg torch_cuda_arch_list="${CUDA_ARCH_AARCH64}" \
               --build-arg INSTALL_KV_CONNECTORS=true \
               --build-arg INSTALL_RUBIN_PRERELEASE=true \
-              --build-arg BUILD_NIXL=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 \
@@ -597,252 +592,96 @@ steps:
   # =============================================================================
   # ROCm Release Pipelines (x86_64 only)
   # =============================================================================
-  #
-  # One group per ROCm stack, identified by its base Dockerfile (ROCM_BASE_DOCKERFILE).
-  # The scripts derive the ROCm version, wheel variant, cache keys and image tags from
-  # it (.buildkite/scripts/rocm/stack.sh), so both groups share one set of step
-  # definitions, anchored on the ROCm 10.0 group.
-  #
-  # - ROCm 10.0 (TheRock, docker/Dockerfile.rocm{,_base}) is the default: it owns the
-  #   :nightly / :latest images and the root wheel index.
-  # - ROCm 7.2 (docker/Dockerfile.rocm_72{,_base}) is kept for a transition period as
-  #   rocm72 images and wheels.
-  #
-  # vLLM version is determined by the Buildkite checkout (like the CUDA pipeline).
+  # One group per ROCm stack, on this pipeline's own agents. A serialized step builds the stack and
+  # runtime from source into ECR (once per content key, with only its own ECR build cache; it never
+  # reads rocm/vllm-ci); the build step compiles csrc for all release arches and pushes
+  # <ecr>:<stack>-<commit>{,-wheels,-base} (.buildkite/scripts/rocm/build.sh release); the steps
+  # after it only copy those (.buildkite/scripts/rocm/publish.sh).
+  # - ROCm 10.0 (TheRock, stack `rocm`) is the default: it owns the :nightly / :latest images and
+  #   the root wheel index.
   # =============================================================================
 
-  - group: "Build ROCm 10.0 Wheel / Image"
-    key: "build-rocm100"
+  - group: ":rocm: ROCm 10.0"
+    key: build-rocm100
     depends_on: ~
     steps:
-      - label: ":rocm: Build ROCm 10.0 Base Image & Wheels"
-        id: build-rocm100-base
+      - label: ":rocm: Ensure ROCm 10.0 release stack + runtime"
+        key: build-rocm100-release-base
         depends_on: ~
-        agents: &rocm_build_agents
-          queue: cpu_queue_release
-        timeout_in_minutes: 360
-        env: &rocm100_env
-          ROCM_BASE_DOCKERFILE: "docker/Dockerfile.rocm_base"
-          ROCM_DEFAULT_ARG: "--default"
-          DOCKER_BUILDKIT: "1"
-          S3_BUCKET: "vllm-wheels"
-          DOCKERHUB_USERNAME: "vllmbot"
-        commands: &rocm_base_commands
-          - |
-            set -euo pipefail
-            source .buildkite/scripts/rocm/stack.sh "$${ROCM_BASE_DOCKERFILE}"
-            SCCACHE_ARGS=(--build-arg USE_SCCACHE=1 --build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache
-              --build-arg SCCACHE_REGION_NAME=us-west-2 --build-arg SCCACHE_S3_NO_CREDENTIALS=0)
-            aws ecr-public get-login-password --region us-east-1 | \
-              docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7
-
-            # Base image, cached in ECR by the base Dockerfile's hash
-            if docker manifest inspect "$${ROCM_STACK_ECR_BASE}" > /dev/null 2>&1; then
-              echo "ECR image cache HIT: $${ROCM_STACK_ECR_BASE}"
-            else
-              docker buildx build --file "$${ROCM_BASE_DOCKERFILE}" --tag "$${ROCM_STACK_ECR_BASE}" \
-                "$${SCCACHE_ARGS[@]}" --push .
-            fi
-
-            # Base wheels, cached in S3 (see cache-rocm-base-wheels.sh)
-            CACHE=.buildkite/scripts/cache-rocm-base-wheels.sh
-            if [ "$$($${CACHE} "$${ROCM_BASE_DOCKERFILE}" check)" = "hit" ]; then
-              $${CACHE} "$${ROCM_BASE_DOCKERFILE}" download
-            else
-              DEBS="rocm-base-debs:$${BUILDKITE_BUILD_NUMBER}-$${ROCM_STACK_VARIANT}"
-              docker buildx build --file "$${ROCM_BASE_DOCKERFILE}" --tag "$${DEBS}" \
-                --target debs_wheel_release "$${SCCACHE_ARGS[@]}" --load .
-              mkdir -p "$${ROCM_STACK_BASE_WHEELS_DIR}"
-              cid=$$(docker create "$${DEBS}")
-              docker cp "$${cid}":/app/debs/. "$${ROCM_STACK_BASE_WHEELS_DIR}/"
-              docker rm "$${cid}"
-              $${CACHE} "$${ROCM_BASE_DOCKERFILE}" upload
-            fi
-            buildkite-agent artifact upload "$${ROCM_STACK_BASE_WHEELS_DIR}/*.whl"
-
-      - label: ":python: Build vLLM ROCm 10.0 Wheel - x86_64"
-        id: build-rocm100-vllm-wheel
-        depends_on: build-rocm100-base
-        agents: *rocm_build_agents
-        timeout_in_minutes: 180
-        env: *rocm100_env
-        commands: &rocm_vllm_wheel_commands
-          - |
-            set -euo pipefail
-            source .buildkite/scripts/rocm/stack.sh "$${ROCM_BASE_DOCKERFILE}"
-            # Buildkite's default fetch doesn't update moved tags; version detection needs them
-            git fetch --tags --force origin
-            buildkite-agent artifact download "$${ROCM_STACK_BASE_WHEELS_DIR}/*.whl" .
-            aws ecr-public get-login-password --region us-east-1 | \
-              docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7
-            docker pull "$${ROCM_STACK_ECR_BASE}"
-            mkdir -p docker/context/base-wheels
-            touch docker/context/base-wheels/.keep
-            cp "$${ROCM_STACK_BASE_WHEELS_DIR}"/*.whl docker/context/base-wheels/
-            docker build --file "$${ROCM_STACK_DOCKERFILE}" --target export_vllm_wheel_release \
-              --output type=local,dest=rocm-dist \
-              --build-arg BASE_IMAGE="$${ROCM_STACK_ECR_BASE}" --build-arg REMOTE_VLLM=0 \
-              --build-arg GIT_REPO_CHECK=1 --build-arg USE_SCCACHE=1 \
-              --build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \
-              --build-arg SCCACHE_REGION_NAME=us-west-2 --build-arg SCCACHE_S3_NO_CREDENTIALS=0 .
-            mkdir -p "$${ROCM_STACK_VLLM_WHEEL_DIR}"
-            cp rocm-dist/*.whl "$${ROCM_STACK_VLLM_WHEEL_DIR}/"
-            buildkite-agent artifact upload "$${ROCM_STACK_VLLM_WHEEL_DIR}/*.whl"
-
-      - label: ":s3: Upload ROCm 10.0 Wheels to S3"
-        id: upload-rocm100-wheels
-        depends_on: build-rocm100-vllm-wheel
-        agents: *rocm_build_agents
-        timeout_in_minutes: 60
-        env: *rocm100_env
-        commands: &rocm_upload_commands
-          - |
-            set -euo pipefail
-            source .buildkite/scripts/rocm/stack.sh "$${ROCM_BASE_DOCKERFILE}"
-            buildkite-agent artifact download "$${ROCM_STACK_BASE_WHEELS_DIR}/*.whl" .
-            buildkite-agent artifact download "$${ROCM_STACK_VLLM_WHEEL_DIR}/*.whl" .
-            bash .buildkite/scripts/upload-rocm-wheels.sh "$${ROCM_BASE_DOCKERFILE}"
-
-      - label: ":memo: Annotate ROCm 10.0 wheel release"
-        id: annotate-rocm100-release
-        depends_on: upload-rocm100-wheels
-        agents: *rocm_build_agents
-        env: *rocm100_env
-        commands: &rocm_annotate_commands
-          - bash .buildkite/scripts/annotate-rocm-release.sh "$${ROCM_BASE_DOCKERFILE}" $${ROCM_DEFAULT_ARG}
-
-      # Generate https://wheels.vllm.ai/rocm/ (the default stack's variant), so users can
-      # `uv pip install vllm --extra-index-url https://wheels.vllm.ai/rocm/` (release only)
-      - block: "Generate Root Index for ROCm Wheels for Release"
-        key: block-generate-root-index-rocm-wheels
-        depends_on:
-          - upload-rocm100-wheels
-          - upload-rocm72-wheels
-
-      - label: ":package: Generate Root Index for ROCm Wheels for Release"
-        depends_on: block-generate-root-index-rocm-wheels
-        id: generate-root-index-rocm-wheels
-        agents: *rocm_build_agents
-        env: *rocm100_env
+        agents: {queue: cpu_queue_release}
+        timeout_in_minutes: 720   # a new stack key (release never uses the CI cache) takes hours
+        concurrency_group: vllm/rocm/release/stack-rocm   # one run builds a new key; the others then reuse it
+        concurrency: 1
+        env: &rocm_release_env {ROCM_CI_REPO: &rocm_ecr public.ecr.aws/q9t5s3a7/vllm-release-repo, SCCACHE_BUCKET_NAME: vllm-build-sccache}
         commands:
-          - "bash tools/vllm-rocm/generate-rocm-wheels-root-index.sh"
-
-      - label: ":docker: Build release image - x86_64 - ROCm 10.0"
-        id: build-rocm100-release-image
-        depends_on:
-          - step: block-build-release-images
-            allow_failure: true
-          - step: build-rocm100-base
-            allow_failure: false
-        agents: *rocm_build_agents
-        timeout_in_minutes: 180
-        env: *rocm100_env
-        commands: &rocm_release_image_commands
-          - |
-            set -euo pipefail
-            source .buildkite/scripts/rocm/stack.sh "$${ROCM_BASE_DOCKERFILE}"
-            aws ecr-public get-login-password --region us-east-1 | \
-              docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7
-            docker pull "$${ROCM_STACK_ECR_BASE}"
-            docker build --file "$${ROCM_STACK_DOCKERFILE}" --target vllm-openai \
-              --tag "$${ROCM_STACK_ECR_IMAGE}" --progress plain \
-              --build-arg max_jobs=16 --build-arg BASE_IMAGE="$${ROCM_STACK_ECR_BASE}" \
-              --build-arg INSTALL_LMCACHE=true --build-arg USE_SCCACHE=1 \
-              --build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \
-              --build-arg SCCACHE_REGION_NAME=us-west-2 --build-arg SCCACHE_S3_NO_CREDENTIALS=0 .
-            docker push "$${ROCM_STACK_ECR_IMAGE}"
-
-      - label: "Publish nightly ROCm 10.0 image to DockerHub"
-        id: publish-rocm100-nightly
-        depends_on: build-rocm100-release-image
-        if: build.env("NIGHTLY") == "1"
-        agents: &rocm_publish_agents
-          queue: small_cpu_queue_release
-        env: *rocm100_env
-        commands: &rocm_publish_nightly_commands
-          - bash .buildkite/scripts/push-nightly-builds-rocm.sh "$${ROCM_BASE_DOCKERFILE}" $${ROCM_DEFAULT_ARG}
-        plugins: &rocm_dockerhub_login
+          - aws ecr-public get-login-password --region us-east-1 | docker login -u AWS --password-stdin $${ROCM_CI_REPO%/*}
+          - bash .buildkite/scripts/rocm/build.sh rocm release-base
+        retry: &rocm_release_retry {automatic: [{exit_status: 75, limit: 1}]}   # registry error (build.sh)
+        plugins: &rocm_hub_login   # authenticated Docker Hub pulls and pushes
           - docker-login#v3.0.0:
               username: vllmbot
               password-env: DOCKERHUB_TOKEN
 
-  - group: "Build ROCm 7.2 Wheel / Image"
-    key: "build-rocm72"
-    depends_on: ~
-    steps:
-      - label: ":rocm: Build ROCm 7.2 Base Image & Wheels"
-        id: build-rocm72-base
-        depends_on: ~
-        agents: *rocm_build_agents
-        timeout_in_minutes: 360
-        env: &rocm72_env
-          ROCM_BASE_DOCKERFILE: "docker/Dockerfile.rocm_72_base"
-          ROCM_DEFAULT_ARG: ""
-          DOCKER_BUILDKIT: "1"
-          S3_BUCKET: "vllm-wheels"
-          DOCKERHUB_USERNAME: "vllmbot"
-        commands: *rocm_base_commands
-
-      - label: ":python: Build vLLM ROCm 7.2 Wheel - x86_64"
-        id: build-rocm72-vllm-wheel
-        depends_on: build-rocm72-base
-        agents: *rocm_build_agents
-        timeout_in_minutes: 180
-        env: *rocm72_env
-        commands: *rocm_vllm_wheel_commands
-
-      - label: ":s3: Upload ROCm 7.2 Wheels to S3"
-        id: upload-rocm72-wheels
-        depends_on: build-rocm72-vllm-wheel
-        agents: *rocm_build_agents
+      - label: ":rocm: Build ROCm 10.0 release image + wheels"
+        key: build-rocm100-release
+        depends_on: build-rocm100-release-base
+        agents: {queue: cpu_queue_release}
+        timeout_in_minutes: 720   # an all-arch csrc compile without sccache hits takes ~3h
+        env: *rocm_release_env
+        commands:
+          - aws ecr-public get-login-password --region us-east-1 | docker login -u AWS --password-stdin $${ROCM_CI_REPO%/*}
+          - bash .buildkite/scripts/rocm/build.sh rocm release
+        retry: *rocm_release_retry
+        plugins: *rocm_hub_login
+
+      - label: ":s3: Upload ROCm 10.0 wheels"
+        key: upload-rocm100-wheels
+        depends_on: build-rocm100-release
+        agents: {queue: cpu_queue_release}
         timeout_in_minutes: 60
-        env: *rocm72_env
-        commands: *rocm_upload_commands
-
-      - label: ":memo: Annotate ROCm 7.2 wheel release"
-        id: annotate-rocm72-release
-        depends_on: upload-rocm72-wheels
-        agents: *rocm_build_agents
-        env: *rocm72_env
-        commands: *rocm_annotate_commands
-
-      - label: ":docker: Build release image - x86_64 - ROCm 7.2"
-        id: build-rocm72-release-image
+        env: {S3_BUCKET: vllm-wheels, ECR: *rocm_ecr}
+        command: bash .buildkite/scripts/rocm/publish.sh rocm wheels
+        plugins: *rocm_hub_login
+
+      - label: ":docker: ROCm 10.0 image: ECR tags (+ nightly Docker Hub)"
+        key: build-rocm100-release-image
         depends_on:
           - step: block-build-release-images
             allow_failure: true
-          - step: build-rocm72-base
-            allow_failure: false
-        agents: *rocm_build_agents
-        timeout_in_minutes: 180
-        env: *rocm72_env
-        commands: *rocm_release_image_commands
-
-      - label: "Publish nightly ROCm 7.2 image to DockerHub"
-        id: publish-rocm72-nightly
-        depends_on: build-rocm72-release-image
-        if: build.env("NIGHTLY") == "1"
-        agents: *rocm_publish_agents
-        env: *rocm72_env
-        commands: *rocm_publish_nightly_commands
-        plugins: *rocm_dockerhub_login
+          - build-rocm100-release
+        agents: {queue: small_cpu_queue_release}
+        env: {ECR: *rocm_ecr}
+        command: bash .buildkite/scripts/rocm/publish.sh rocm images
+        plugins: *rocm_hub_login
 
-      - label: "Publish nightly XPU image to DockerHub"
-        depends_on:
-          - create-manifest-xpu
-        if: build.env("NIGHTLY") == "1"
-        agents:
-          queue: small_cpu_queue_release
-        commands:
-          - "bash .buildkite/scripts/xpu/push-nightly-builds-xpu.sh"
-          - "bash .buildkite/scripts/cleanup-nightly-builds.sh nightly- vllm/vllm-openai-xpu"
-        plugins:
-          - docker-login#v3.0.0:
-              username: vllmbot
-              password-env: DOCKERHUB_TOKEN
-        env:
-          DOCKER_BUILDKIT: "1"
-          DOCKERHUB_USERNAME: "vllmbot"
+      # Generate https://wheels.vllm.ai/rocm/ (the default stack's variant).
+      - block: "Generate Root Index for ROCm Wheels for Release"
+        key: block-generate-root-index-rocm-wheels
+        depends_on: upload-rocm100-wheels
+
+      - label: ":package: Generate Root Index for ROCm Wheels for Release"
+        key: generate-root-index-rocm-wheels
+        depends_on: block-generate-root-index-rocm-wheels
+        agents: {queue: cpu_queue_release}
+        env: {S3_BUCKET: vllm-wheels}
+        command: VARIANT=$$(buildkite-agent meta-data get rocm-rocm-variant) bash tools/vllm-rocm/generate-rocm-wheels-root-index.sh
+
+  - label: "Publish nightly XPU image to DockerHub"
+    depends_on:
+      - create-manifest-xpu
+    if: build.env("NIGHTLY") == "1"
+    agents:
+      queue: small_cpu_queue_release
+    commands:
+      - "bash .buildkite/scripts/xpu/push-nightly-builds-xpu.sh"
+      - "bash .buildkite/scripts/cleanup-nightly-builds.sh nightly- vllm/vllm-openai-xpu"
+    plugins:
+      - docker-login#v3.0.0:
+          username: vllmbot
+          password-env: DOCKERHUB_TOKEN
+    env:
+      DOCKER_BUILDKIT: "1"
+      DOCKERHUB_USERNAME: "vllmbot"
 
   # =============================================================================
   # Publish to DockerHub and PyPI (at the end so all builds complete first)
@@ -854,7 +693,6 @@ steps:
       - create-multi-arch-manifest
       - create-manifest-xpu
       - build-rocm100-release-image
-      - build-rocm72-release-image
       - input-release-version
       # Wait for CPU builds if their block steps were unblocked, so publish
       # doesn't race the in-progress CPU build. allow_failure lets publish
@@ -882,7 +720,6 @@ steps:
       # AWS EFA variant; skipped if no EFA image was built, fails on a partial build.
       - "cuda-13-0-efa"
       - "rocm"
-      - "rocm72"
       - "xpu"
       - "cpu"
     plugins:
diff --git a/.buildkite/scripts/annotate-rocm-release.sh b/.buildkite/scripts/annotate-rocm-release.sh
deleted file mode 100755
index 1889db43b5ad..000000000000
--- a/.buildkite/scripts/annotate-rocm-release.sh
+++ /dev/null
@@ -1,87 +0,0 @@
-#!/bin/bash
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-#
-# Generate Buildkite annotation for a ROCm stack's wheel release.
-#
-# Usage: annotate-rocm-release.sh <base-dockerfile> [--default]
-#   base-dockerfile  e.g. docker/Dockerfile.rocm_base or docker/Dockerfile.rocm_72_base
-#   --default        the stack also owns the plain :latest / :v<ver> Docker Hub tags
-set -ex
-
-# shellcheck source=.buildkite/scripts/rocm/stack.sh
-source "$(dirname "${BASH_SOURCE[0]}")/rocm/stack.sh" "${1:?Usage: $0 <base-dockerfile> [--default]}"
-TAG_VARIANTS=("-${ROCM_STACK_VARIANT}")
-if [[ "${2:-}" == "--default" ]]; then
-  TAG_VARIANTS=("" "-${ROCM_STACK_VARIANT}")
-fi
-ROCM_VERSION="$ROCM_STACK_VERSION"
-VARIANT="$ROCM_STACK_VARIANT"
-PYTHON_VERSION=$(sed -nE 's/^ARG PYTHON_VERSION="?([^" ]+)"?.*/\1/p' "$ROCM_STACK_BASE_DOCKERFILE")
-PYTORCH_ROCM_ARCH=$(sed -nE 's/^ARG PYTORCH_ROCM_ARCH="?([^"]+)"?.*/\1/p' "$ROCM_STACK_BASE_DOCKERFILE")
-
-# Get release version, default to 1.0.0.dev for nightly/per-commit builds
-RELEASE_VERSION=$(buildkite-agent meta-data get release-version 2>/dev/null || echo "")
-if [ -z "${RELEASE_VERSION}" ]; then
-  RELEASE_VERSION="1.0.0.dev"
-fi
-
-# S3 URLs
-S3_BUCKET="${S3_BUCKET:-vllm-wheels}"
-S3_REGION="${AWS_DEFAULT_REGION:-us-west-2}"
-S3_HOST="${S3_BUCKET}.s3-website-${S3_REGION}.amazonaws.com"
-S3_URL="http://${S3_HOST}"
-WHEEL_PATH="rocm/${BUILDKITE_COMMIT}/${VARIANT}-wheels/"
-
-DOCKER_CMDS="docker pull ${ROCM_STACK_ECR_BASE}
-docker pull ${ROCM_STACK_ECR_IMAGE}"
-for suffix in "${TAG_VARIANTS[@]}"; do
-  for tag in "latest${suffix}" "v${RELEASE_VERSION}${suffix}"; do
-    DOCKER_CMDS+="
-docker tag ${ROCM_STACK_ECR_BASE} vllm/vllm-openai-rocm:${tag}-base
-docker tag ${ROCM_STACK_ECR_IMAGE} vllm/vllm-openai-rocm:${tag}
-docker push vllm/vllm-openai-rocm:${tag}-base
-docker push vllm/vllm-openai-rocm:${tag}"
-  done
-done
-
-buildkite-agent annotate --style 'success' --context "rocm-release-workflow-${VARIANT}" << EOF
-## ROCm ${ROCM_VERSION} Wheel and Docker Image Releases
-### Build Configuration
-| Setting | Value |
-|---------|-------|
-| **ROCm Version** | ${ROCM_VERSION} |
-| **Index Variant** | ${VARIANT} |
-| **Python Version** | ${PYTHON_VERSION} |
-| **GPU Architectures** | ${PYTORCH_ROCM_ARCH} |
-| **Branch** | \`${BUILDKITE_BRANCH}\` |
-| **Commit** | \`${BUILDKITE_COMMIT}\` |
-
-### :package: Installation
-
-\`\`\`bash
-# This build (by commit)
-pip install vllm --extra-index-url ${S3_URL}/rocm/${BUILDKITE_COMMIT}/${VARIANT}/ --trusted-host ${S3_HOST}
-# Nightly (if published)
-pip install vllm --extra-index-url ${S3_URL}/rocm/nightly/${VARIANT}/ --trusted-host ${S3_HOST}
-\`\`\`
-
-### :floppy_disk: Download Wheels Directly
-
-\`\`\`bash
-aws s3 ls s3://${S3_BUCKET}/${WHEEL_PATH}
-aws s3 cp --recursive --exclude "*" --include "*.whl" s3://${S3_BUCKET}/${WHEEL_PATH} .
-\`\`\`
-
-### :warning: Notes
-- These wheels are built for **ROCm ${ROCM_VERSION}** and will NOT work with CUDA GPUs
-- Supported GPU architectures: ${PYTORCH_ROCM_ARCH}
-- Platform: Linux x86_64 only
-
-### :package: Docker Image Release
-
-\`\`\`
-${DOCKER_CMDS}
-\`\`\`
-
-EOF
diff --git a/.buildkite/scripts/cache-rocm-base-wheels.sh b/.buildkite/scripts/cache-rocm-base-wheels.sh
deleted file mode 100755
index 02fbbf52cf54..000000000000
--- a/.buildkite/scripts/cache-rocm-base-wheels.sh
+++ /dev/null
@@ -1,125 +0,0 @@
-#!/usr/bin/env bash
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-#
-# Cache helper for ROCm base wheels
-#
-# This script manages caching of pre-built ROCm base wheels (torch, triton, etc.)
-# to avoid rebuilding them when the stack's base Dockerfile hasn't changed.
-#
-# Usage: cache-rocm-base-wheels.sh <base-dockerfile> <command>
-#   base-dockerfile  e.g. docker/Dockerfile.rocm_base or docker/Dockerfile.rocm_72_base
-# Commands:
-#   check    - Check if cache exists, outputs "hit" or "miss"
-#   upload   - Upload wheels to cache
-#   download - Download wheels from cache
-#   key      - Output the cache key
-#   path     - Output the full S3 cache path
-#
-# Environment variables:
-#   S3_BUCKET          - S3 bucket name (default: vllm-wheels)
-#
-# The cache key hashes the base Dockerfile and the repo files it COPYs in, which
-# capture the ROCm version and every component pin (see rocm/stack.sh).
-
-set -euo pipefail
-
-BUCKET="${S3_BUCKET:-vllm-wheels}"
-CACHE_PREFIX="rocm/cache"
-
-# shellcheck source=.buildkite/scripts/rocm/stack.sh
-source "$(dirname "${BASH_SOURCE[0]}")/rocm/stack.sh" "${1:?Usage: $0 <base-dockerfile> <command>}"
-shift
-CACHE_KEY="$ROCM_STACK_WHEEL_KEY"
-WHEELS_DIR="$ROCM_STACK_BASE_WHEELS_DIR"
-CACHE_PATH="s3://${BUCKET}/${CACHE_PREFIX}/${CACHE_KEY}/"
-
-case "${1:-}" in
-    check)
-        echo "Checking cache for key: ${CACHE_KEY}" >&2
-        echo "Cache path: ${CACHE_PATH}" >&2
-
-        # Check if cache exists by listing objects
-        # We look for at least one .whl file
-        echo "Running: aws s3 ls ${CACHE_PATH}" >&2
-        S3_OUTPUT=$(aws s3 ls "${CACHE_PATH}" 2>&1) || true
-        echo "S3 ls output:" >&2
-        echo "$S3_OUTPUT" | head -5 >&2
-
-        if echo "$S3_OUTPUT" | grep -q "\.whl"; then
-            echo "hit"
-        else
-            echo "miss"
-        fi
-        ;;
-
-    upload)
-        echo "========================================"
-        echo "Uploading wheels to cache"
-        echo "========================================"
-        echo "Cache key: ${CACHE_KEY}"
-        echo "Cache path: ${CACHE_PATH}"
-        echo ""
-
-        if [[ ! -d "${WHEELS_DIR}" ]]; then
-            echo "ERROR: ${WHEELS_DIR} directory not found" >&2
-            exit 1
-        fi
-
-        WHEEL_COUNT=$(find "${WHEELS_DIR}" -maxdepth 1 -name '*.whl' 2>/dev/null | wc -l)
-        if [[ "$WHEEL_COUNT" -eq 0 ]]; then
-            echo "ERROR: No wheels found in ${WHEELS_DIR}/" >&2
-            exit 1
-        fi
-
-        echo "Uploading $WHEEL_COUNT wheels..."
-        aws s3 cp --recursive "${WHEELS_DIR}/" "${CACHE_PATH}"
-
-        echo ""
-        echo "Cache upload complete!"
-        echo "========================================"
-        ;;
-
-    download)
-        echo "========================================"
-        echo "Downloading wheels from cache"
-        echo "========================================"
-        echo "Cache key: ${CACHE_KEY}"
-        echo "Cache path: ${CACHE_PATH}"
-        echo ""
-        mkdir -p "${WHEELS_DIR}"
-        
-        # Use sync with include/exclude to only download .whl files
-        aws s3 sync "${CACHE_PATH}" "${WHEELS_DIR}/" \
-            --exclude "*" \
-            --include "*.whl"
-        
-        echo ""
-        echo "Downloaded wheels:"
-        find "${WHEELS_DIR}" -maxdepth 1 -name '*.whl' -exec ls -lh {} \;
-        WHEEL_COUNT=$(find "${WHEELS_DIR}" -maxdepth 1 -name '*.whl' 2>/dev/null | wc -l)
-        echo ""
-        echo "Total: $WHEEL_COUNT wheels"
-        echo "========================================"
-        ;;
-
-    key)
-        echo "${CACHE_KEY}"
-        ;;
-
-    path)
-        echo "${CACHE_PATH}"
-        ;;
-
-    *)
-        echo "Usage: $0 <base-dockerfile> {check|upload|download|key|path}" >&2
-        echo "" >&2
-        echo "Commands:" >&2
-        echo "  check    - Check if cache exists, outputs 'hit' or 'miss'" >&2
-        echo "  upload   - Upload wheels from ${WHEELS_DIR}/ to cache" >&2
-        echo "  download - Download wheels from cache to ${WHEELS_DIR}/" >&2
-        echo "  key      - Output the cache key" >&2
-        echo "  path     - Output the full S3 cache path" >&2
-        exit 1
-        ;;
-esac
diff --git a/.buildkite/scripts/check-ray-compatibility.sh b/.buildkite/scripts/check-ray-compatibility.sh
index ec42834802d6..8e1eef6948fe 100644
--- a/.buildkite/scripts/check-ray-compatibility.sh
+++ b/.buildkite/scripts/check-ray-compatibility.sh
@@ -27,6 +27,8 @@ if python3 -c "import torch; assert torch.version.hip" 2>/dev/null; then
         echo ">>> WARNING: ROCm ${ROCM_VER} wheel index not found at ${CANDIDATE_URL}"
         echo ">>>          Falling back to default PyPI (resolution may be incomplete)"
         TORCH_INDEX_URL=""
+        # The image's own stack builds (TheRock torch, triton, amdsmi) are on no public index.
+        LOCAL_BUILDS=1
     fi
 else
     if [ "${TORCH_NIGHTLY:-0}" = "1" ]; then
@@ -110,6 +112,14 @@ with open(out_path, "w") as f:
 print(f"Wrote {len(lines)} constraints to {out_path}")
 PYEOF
 
+if [[ -n "${LOCAL_BUILDS:-}" ]]; then
+    # Compare them at their public release. Dropped: amdsmi has none (TheRock ships it as
+    # sources) and triton is the stack's own, which its torch is repacked to require.
+    echo ">>> Pinning local builds to their public release:"
+    grep -E '==[^ ;]*\+|\[device-' "${WORK_DIR}/vllm-constraints.txt" || true
+    sed -i -E -e '/^(amdsmi|triton)([=<>!~ ;[]|$)|\[device-/d' -e 's/(==[^ ;+]+)\+[^ ;]*/\1/' "${WORK_DIR}/vllm-constraints.txt"
+fi
+
 echo ">>> Installed vllm deps (first 20 lines):"
 head -20 "${WORK_DIR}/vllm-constraints.txt"
 
diff --git a/.buildkite/scripts/ci-bake-rocm.sh b/.buildkite/scripts/ci-bake-rocm.sh
deleted file mode 100644
index db3d0a4acfe4..000000000000
--- a/.buildkite/scripts/ci-bake-rocm.sh
+++ /dev/null
@@ -1,2803 +0,0 @@
-#!/bin/bash
-# ci-bake-rocm.sh - Docker buildx bake wrapper for ROCm CI builds.
-#
-# The wrapper keeps three build concerns separate:
-#   * ci_base builds are content-addressed by vllm.ci_base.content_hash.
-#   * test image contents record the commit while runtime tags stay build-scoped.
-#   * ROCm install artifacts are uploaded only for targets that export wheels.
-#
-# Usage:
-#   ci-bake-rocm.sh [TARGET]
-#
-# Set BAKE_PRINT_ONLY=1 to stop after docker buildx bake --print.
-
-set -euo pipefail
-
-DEFAULT_REPO_SLUG="vllm-project/vllm"
-DEFAULT_CI_HCL_SOURCE="docker/ci-rocm.hcl"
-DEFAULT_CI_BASE_CONTENT_FILES=".dockerignore requirements/common.txt requirements/rocm.txt requirements/test/rocm.txt tools/install_torchcodec_rocm.sh rust-toolchain.toml tests/vllm_test_utils"
-DEFAULT_CI_BASE_DOCKERFILE="docker/Dockerfile.rocm"
-DEFAULT_CI_BASE_DOCKERFILE_STAGES="base rust_toolchain_input_0 rust-toolchain-input rust-toolchain build_nixl lmcache_source build_lmcache build_deepep mori_base ci_base"
-DEFAULT_CI_BASE_METADATA_VERSION="3"
-# ROCm CI forces REMOTE_VLLM=0, so content identity covers only the selected
-# local-source stages rather than unreachable remote-fetch alternatives.
-DEFAULT_ROCM_CSRC_CONTENT_FILES=".dockerignore requirements/common.txt requirements/rocm.txt pyproject.toml setup.py CMakeLists.txt cmake csrc vllm/envs.py vllm/__init__.py tools/build_rust.py"
-DEFAULT_ROCM_CSRC_DOCKERFILE_STAGES="base fetch_vllm_0 fetch_vllm build_vllm_dependencies rocm-triton-kernels csrc-build"
-DEFAULT_ROCM_RUST_CONTENT_FILES=".dockerignore .git_archival.txt pyproject.toml requirements/build/rust.txt rust/Cargo.lock rust/Cargo.toml rust/proto rust/src rust-toolchain.toml tools/build_rust.py tools/build_rust.sh"
-DEFAULT_ROCM_RUST_DOCKERFILE_STAGES="base fetch_vllm_0 fetch_vllm vllm-version rust_toolchain_input_0 rust-toolchain-input rust_input_0 rust-input rust-toolchain rust-build"
-# Docker's 128-character tag limit minus the longest cache prefix
-# ("csrc-rocm-branch-" and "rust-rocm-branch-", both 17 characters).
-ROCM_CACHE_BRANCH_TAG_MAX_LEN=111
-CI_BASE_WRITE_SCOPE=""
-
-TARGET=""
-CI_HCL_SOURCE="${CI_HCL_SOURCE:-}"
-CI_HCL_PATH=""
-CI_BASE_LABEL_OVERRIDE_PATH=""
-CSRC_CACHE_OVERRIDE_PATH=""
-ROCM_ARG_OVERRIDE_PATH=""
-BUILD_CONTEXT_OVERRIDE_PATH=""
-SCRIPT_TMP_DIR=""
-BAKE_CONFIG_FILE=""
-ROCM_BUILD_CONTEXT_ROOT=""
-ROCM_BUILD_CONTEXT_INDEX=""
-ROCM_BUILD_CONTEXT_COMMIT=""
-BAKE_FILES=()
-BAKE_ALLOW_ARGS=()
-BAKE_TARGETS=()
-DEPENDENCY_CACHE_TARGETS=()
-
-cleanup() {
-    if [[ -n "${SCRIPT_TMP_DIR}" && -d "${SCRIPT_TMP_DIR}" ]]; then
-        rm -rf "${SCRIPT_TMP_DIR}"
-    fi
-}
-trap cleanup EXIT
-
-clean_docker_tag() {
-    local input="$1"
-    echo "${input}" | sed 's/[^a-zA-Z0-9._-]/_/g' | cut -c1-128
-}
-
-is_url_like() {
-    local value="${1:-}"
-    [[ "${value}" =~ ^[a-zA-Z][a-zA-Z0-9+.-]*:// || "${value}" == git@*:* ]]
-}
-
-is_full_git_sha() {
-    local value="${1:-}"
-    [[ "${value}" =~ ^[0-9a-fA-F]{40}$ ]]
-}
-
-select_cache_branch_name() {
-    local candidate=""
-    local var=""
-
-    for var in \
-        ROCM_CACHE_BRANCH_NAME \
-        BUILDKITE_PULL_REQUEST_HEAD_BRANCH \
-        BUILDKITE_HEAD_BRANCH \
-        BUILDKITE_BRANCH \
-        VLLM_BRANCH; do
-        candidate="${!var:-}"
-        [[ -n "${candidate}" ]] || continue
-        is_url_like "${candidate}" && continue
-        is_full_git_sha "${candidate}" && continue
-        printf '%s\n' "${candidate}"
-        return 0
-    done
-}
-
-cache_scope_suffix() {
-    local arch_hash=""
-    arch_hash=$(printf '%s' "${PYTORCH_ROCM_ARCH:-default}" | sha256sum | cut -c1-12)
-    printf 'arch-%s\n' "${arch_hash}"
-}
-
-compose_cache_branch_tag() {
-    local repo_slug="$1"
-    local branch="$2"
-    local suffix=""
-    local prefix=""
-    local max_prefix_len=0
-    local max_tag_len="${ROCM_CACHE_BRANCH_TAG_MAX_LEN}"
-
-    suffix="$(cache_scope_suffix)"
-    prefix="$(clean_docker_tag "${repo_slug}")-$(clean_docker_tag "${branch}")"
-    max_prefix_len=$((max_tag_len - ${#suffix} - 1))
-    if (( max_prefix_len < 1 )); then
-        max_prefix_len=1
-    fi
-    printf '%s-%s\n' "${prefix:0:${max_prefix_len}}" "${suffix}"
-}
-
-parse_repo_slug() {
-    local repo_url="${1:-}"
-    local repo_slug=""
-
-    if [[ -z "${repo_url}" ]]; then
-        printf '%s\n' "${DEFAULT_REPO_SLUG}"
-        return 0
-    fi
-
-    repo_slug=$(echo "${repo_url}" | sed -E 's#(git@|https?://)([^/:]+)[:/]([^/]+/[^/.]+)(\.git)?$#\3#')
-    if [[ "${repo_slug}" != */* ]]; then
-        repo_slug="${DEFAULT_REPO_SLUG}"
-    fi
-    printf '%s\n' "${repo_slug}"
-}
-
-normalize_repo_slug() {
-    local repo_slug="${1:-}"
-
-    repo_slug="${repo_slug%/}"
-    repo_slug="${repo_slug%.git}"
-    repo_slug="${repo_slug#https://github.com/}"
-    repo_slug="${repo_slug#http://github.com/}"
-    repo_slug="${repo_slug#ssh://git@github.com/}"
-    repo_slug="${repo_slug#git@github.com:}"
-    repo_slug="${repo_slug#github.com/}"
-    printf '%s\n' "${repo_slug}"
-}
-
-is_trusted_ci_cache_writer() {
-    local actual_repo=""
-    local trusted_repo=""
-
-    [[ "${BUILDKITE:-false}" == "true" ]] || return 1
-    [[ "${BUILDKITE_PULL_REQUEST:-false}" == "false" ]] || return 1
-    [[ "${BUILDKITE_BRANCH:-}" == "${CI_BASE_STABLE_BRANCH:-main}" ]] || return 1
-    actual_repo=$(normalize_repo_slug "${BUILDKITE_REPO:-}")
-    trusted_repo=$(normalize_repo_slug \
-        "${CI_BASE_STABLE_REPO_SLUG:-${DEFAULT_REPO_SLUG}}")
-    [[ -n "${actual_repo}" && "${actual_repo}" == "${trusted_repo}" ]]
-}
-
-ci_base_write_scope() {
-    local identity=""
-    local source_repo="${BUILDKITE_PULL_REQUEST_REPO:-${BUILDKITE_REPO:-local}}"
-
-    if is_trusted_ci_cache_writer; then
-        return 0
-    fi
-    identity=$(printf '%s\n' "${source_repo}" | sha256sum | cut -c1-12)
-    printf 'preview-%s\n' "${identity}"
-}
-
-configure_ci_base_write_scope() {
-    local scope=""
-
-    scope=$(ci_base_write_scope)
-    if [[ -n "${scope}" ]]; then
-        CI_BASE_WRITE_SCOPE=$(clean_docker_tag "${scope}")
-        echo "Non-canonical cache writes use source scope: ${CI_BASE_WRITE_SCOPE}"
-    else
-        CI_BASE_WRITE_SCOPE=""
-        echo "Trusted main build: publishing canonical ci_base refs"
-    fi
-    export CI_BASE_WRITE_SCOPE
-}
-
-get_buildkite_repo_slug() {
-    parse_repo_slug "${BUILDKITE_PULL_REQUEST_REPO:-${BUILDKITE_REPO:-}}"
-}
-
-get_buildkite_target_repo_slug() {
-    parse_repo_slug "${BUILDKITE_REPO:-}"
-}
-
-get_buildkite_target_repo_url() {
-    local repo_url="${BUILDKITE_REPO:-}"
-
-    if [[ -n "${repo_url}" ]] && is_url_like "${repo_url}"; then
-        printf '%s\n' "${repo_url}"
-        return 0
-    fi
-
-    printf 'https://github.com/%s.git\n' "${DEFAULT_REPO_SLUG}"
-}
-
-git_fetch_with_timeout() {
-    local timeout_secs="${ROCM_CACHE_GIT_FETCH_TIMEOUT:-60}"
-    local -a fetch_command=(git fetch --no-auto-maintenance)
-
-    # Detached maintenance can race a later shallow fetch on .git/shallow.
-    if command -v timeout >/dev/null 2>&1; then
-        timeout "${timeout_secs}s" "${fetch_command[@]}" "$@"
-    else
-        "${fetch_command[@]}" "$@"
-    fi
-}
-
-git_fetch_for_cache() {
-    git_fetch_with_timeout "$@" 2>/dev/null
-}
-
-hash_string_short() {
-    printf '%s' "$1" | sha256sum | cut -c1-16
-}
-
-list_content_files() {
-    # Hash the checkout inputs Docker can intentionally consume, not ignored
-    # compiler/test debris left behind on a reused worker.
-    if [[ -n "${ROCM_BUILD_CONTEXT_ROOT:-}" ]]; then
-        GIT_INDEX_FILE="${ROCM_BUILD_CONTEXT_INDEX}" \
-            git ls-files -z --cached -- "$1" | LC_ALL=C sort -z
-    else
-        git ls-files -z --cached --others --exclude-standard -- "$1" \
-            | LC_ALL=C sort -z
-    fi
-}
-
-content_regular_file() {
-    local file="$1"
-    local physical_file=""
-
-    physical_file="${file}"
-    if [[ -n "${ROCM_BUILD_CONTEXT_ROOT:-}" && "${file}" != /* ]]; then
-        physical_file="${ROCM_BUILD_CONTEXT_ROOT}/${file}"
-    fi
-    [[ -f "${physical_file}" ]]
-}
-
-hash_content_file() {
-    local file="$1"
-    local physical_file=""
-    local checksum=""
-    local file_mode=""
-
-    physical_file="${file}"
-    if [[ -n "${ROCM_BUILD_CONTEXT_ROOT:-}" && "${file}" != /* ]]; then
-        physical_file="${ROCM_BUILD_CONTEXT_ROOT}/${file}"
-    fi
-    if [[ -L "${physical_file}" ]]; then
-        printf 'symlink:%s\ntarget:' "${file}"
-        readlink -n -- "${physical_file}" || return $?
-        printf '\n'
-        return
-    fi
-    if [[ ! -f "${physical_file}" ]]; then
-        printf 'missing:%s\n' "${file}"
-        return
-    fi
-    file_mode=$(stat -c '%a' "${physical_file}") || return $?
-    printf 'file:%s\nmode:%s\n' "${file}" "${file_mode}"
-    checksum=$(sha256sum < "${physical_file}") || return $?
-    checksum="${checksum%% *}"
-    printf '%s  %s\n' "${checksum}" "${file}"
-}
-
-hash_content_directory() {
-    if ! list_content_files "$1" | while IFS= read -r -d '' file; do
-        hash_content_file "${file}" || exit $?
-    done; then
-        echo "Failed to hash content under $1" >&2
-        return 1
-    fi
-}
-
-compute_content_hash() {
-    local path=""
-    local physical_path=""
-
-    for path in "$@"; do
-        physical_path="${path}"
-        if [[ -n "${ROCM_BUILD_CONTEXT_ROOT:-}" && "${path}" != /* ]]; then
-            physical_path="${ROCM_BUILD_CONTEXT_ROOT}/${path}"
-        fi
-        if [[ -L "${physical_path}" || -f "${physical_path}" ]]; then
-            hash_content_file "${path}" || return $?
-        elif [[ -d "${physical_path}" ]]; then
-            hash_content_directory "${path}" || return $?
-        else
-            printf 'missing:%s\n' "${path}"
-        fi
-    done | sha256sum | cut -d' ' -f1
-}
-
-validate_ci_build_context_source() {
-    local source_root="$1"
-    local context_commit=""
-    local source_head=""
-    local entry=""
-    local metadata=""
-    local mode=""
-    local stage=""
-    local path=""
-    local worktree_diff_status=0
-    local staged_diff_status=0
-    local index_modes_file="${SCRIPT_TMP_DIR}/git-index-modes"
-    local info_attributes=""
-
-    if ! context_commit=$(git -C "${source_root}" rev-parse \
-        --verify "${BUILDKITE_COMMIT:-HEAD}^{commit}"); then
-        echo "Failed to resolve the CI Docker context revision" >&2
-        return 1
-    fi
-    if ! source_head=$(git -C "${source_root}" rev-parse --verify HEAD); then
-        echo "Failed to resolve the checked-out CI revision" >&2
-        return 1
-    fi
-    if [[ "${context_commit}" != "${source_head}" ]]; then
-        echo "BUILDKITE_COMMIT does not match the checked-out CI revision" >&2
-        return 1
-    fi
-    ROCM_BUILD_CONTEXT_COMMIT="${context_commit}"
-    # Shared AMD workspaces may present tracked files with inflated executable
-    # bits. Context modes come from the pinned Git tree, so ignore only that
-    # filesystem drift while continuing to reject content changes.
-    git -C "${source_root}" -c core.fileMode=false diff \
-        --quiet --no-ext-diff --ignore-submodules=none -- \
-        || worktree_diff_status=$?
-    if (( worktree_diff_status != 0 )); then
-        if (( worktree_diff_status > 1 )); then
-            printf 'Failed to inspect tracked CI worktree changes (git diff exited %s)\n' \
-                "${worktree_diff_status}" >&2
-            return 1
-        fi
-        echo "Tracked worktree changes cannot be omitted from the CI Docker context" \
-            >&2
-        echo "Tracked worktree diff (first 50 entries):" >&2
-        git -C "${source_root}" -c core.fileMode=false diff \
-            --no-ext-diff --ignore-submodules=none --name-status -- \
-            | sed -n '1,50p' >&2 || true
-        git -C "${source_root}" -c core.fileMode=false diff \
-            --no-ext-diff --ignore-submodules=none --summary -- \
-            | sed -n '1,50p' >&2 || true
-        return 1
-    fi
-    git -C "${source_root}" diff --cached \
-        --quiet --no-ext-diff --ignore-submodules=none HEAD -- \
-        || staged_diff_status=$?
-    if (( staged_diff_status != 0 )); then
-        if (( staged_diff_status > 1 )); then
-            printf 'Failed to inspect staged CI changes (git diff exited %s)\n' \
-                "${staged_diff_status}" >&2
-            return 1
-        fi
-        echo "Staged changes cannot be omitted from the CI Docker context" >&2
-        echo "Staged diff (first 50 entries):" >&2
-        git -C "${source_root}" diff --cached \
-            --no-ext-diff --ignore-submodules=none --name-status HEAD -- \
-            | sed -n '1,50p' >&2 || true
-        git -C "${source_root}" diff --cached \
-            --no-ext-diff --ignore-submodules=none --summary HEAD -- \
-            | sed -n '1,50p' >&2 || true
-        return 1
-    fi
-    # Untracked and ignored worker outputs are intentionally absent: the
-    # pinned Git tree, rather than mutable checkout contents, is the contract.
-    if ! info_attributes=$(git -C "${source_root}" \
-        rev-parse --path-format=absolute --git-path info/attributes); then
-        echo "Failed to locate repository-local Git attributes" >&2
-        return 1
-    fi
-    if [[ -s "${info_attributes}" ]]; then
-        echo "Repository-local Git attributes cannot define the CI Docker context" \
-            >&2
-        return 1
-    fi
-
-    ROCM_BUILD_CONTEXT_INDEX="${SCRIPT_TMP_DIR}/docker-context.index"
-    if ! GIT_INDEX_FILE="${ROCM_BUILD_CONTEXT_INDEX}" \
-        git -C "${source_root}" -c core.splitIndex=false \
-            read-tree "${context_commit}^{tree}"; then
-        echo "Failed to create the CI Docker context index" >&2
-        return 1
-    fi
-    if ! GIT_INDEX_FILE="${ROCM_BUILD_CONTEXT_INDEX}" \
-        git -C "${source_root}" ls-files --stage -z \
-            > "${index_modes_file}"; then
-        echo "Failed to read the Git index for the CI Docker context" >&2
-        return 1
-    fi
-    while IFS= read -r -d '' entry; do
-        metadata="${entry%%$'\t'*}"
-        mode="${metadata%% *}"
-        stage="${metadata##* }"
-        path="${entry#*$'\t'}"
-        if [[ "${stage}" != "0" ]]; then
-            echo "Unmerged Git entry cannot be used as Docker context: ${path}" >&2
-            return 1
-        fi
-        case "${mode}" in
-            100644|100755|120000)
-                ;;
-            160000)
-                echo "Git submodule cannot be materialized in the CI Docker context: ${path}" \
-                    >&2
-                return 1
-                ;;
-            *)
-                echo "Unsupported Git mode ${mode} for ${path}" >&2
-                return 1
-                ;;
-        esac
-    done < "${index_modes_file}"
-}
-
-describe_ci_revision() {
-    git -C "$1" describe --tags --long --abbrev=10 \
-        --match 'v[0-9]*' "$2" 2>/dev/null
-}
-
-write_ci_git_archival_metadata() {
-    local source_root="$1"
-    local context_root="$2"
-    local commit="${ROCM_BUILD_CONTEXT_COMMIT}"
-    local commit_date=""
-    local describe=""
-    local is_shallow="false"
-
-    commit_date=$(git -C "${source_root}" show -s --format=%cI "${commit}") \
-        || return $?
-    is_shallow=$(git -C "${source_root}" rev-parse --is-shallow-repository) \
-        || return $?
-    if git -C "${source_root}" remote get-url origin >/dev/null 2>&1; then
-        # AMD agents use shallow, no-tag clones. Reach a release tag so one
-        # commit cannot acquire different versions from different depths.
-        # Versioning only needs tags and commits, not historical source trees.
-        echo "Synchronizing version tags for the canonical CI Docker context"
-        (cd "${source_root}" && git_fetch_with_timeout --quiet \
-            --filter=tree:0 --prune origin '+refs/tags/*:refs/tags/*') \
-            || return $?
-        if [[ "${is_shallow}" == "true" ]] \
-            && ! describe_ci_revision "${source_root}" "${commit}" >/dev/null; then
-            echo "Deepening history to reach a version tag"
-            (cd "${source_root}" && git_fetch_with_timeout --quiet \
-                --filter=tree:0 --no-tags --deepen=1000 origin "${commit}") \
-                || return $?
-            is_shallow=$(git -C "${source_root}" \
-                rev-parse --is-shallow-repository) || return $?
-            if [[ "${is_shallow}" == "true" ]] \
-                && ! describe_ci_revision "${source_root}" "${commit}" >/dev/null; then
-                echo "No version tag within 1,000 commits; fetching full history"
-                (cd "${source_root}" && git_fetch_with_timeout --quiet \
-                    --filter=tree:0 --no-tags --unshallow origin "${commit}") \
-                    || return $?
-            fi
-        fi
-    fi
-    if ! describe=$(describe_ci_revision "${source_root}" "${commit}"); then
-        echo "No numeric version tag is reachable from the CI revision" >&2
-        return 1
-    fi
-    if [[ -e "${context_root}/.git_archival.txt" \
-        || -L "${context_root}/.git_archival.txt" ]]; then
-        echo "The source tree already contains .git_archival.txt" >&2
-        return 1
-    fi
-    printf 'node: %s\nnode-date: %s\ndescribe-name: %s\n' \
-        "${commit}" "${commit_date}" "${describe}" \
-        > "${context_root}/.git_archival.txt" \
-        && chmod 0644 -- "${context_root}/.git_archival.txt"
-}
-
-write_build_context_override() {
-    local escaped_context=""
-
-    [[ -n "${ROCM_BUILD_CONTEXT_ROOT:-}" ]] || return 0
-    escaped_context=$(hcl_escape_string "${ROCM_BUILD_CONTEXT_ROOT}") || return $?
-    if ! {
-        printf 'target "_common-rocm" {\n'
-        printf '  context = "%s"\n' "${escaped_context}"
-        printf '}\n'
-    } > "${BUILD_CONTEXT_OVERRIDE_PATH}"; then
-        echo "Failed to write the CI Docker context override" >&2
-        return 1
-    fi
-    BAKE_FILES+=(-f "${BUILD_CONTEXT_OVERRIDE_PATH}")
-}
-
-prepare_ci_build_context() {
-    local source_root=""
-    local context_root=""
-
-    [[ "${BUILDKITE:-false}" == "true" ]] || return 0
-    [[ "${REMOTE_VLLM:-0}" == "0" ]] || return 0
-
-    # BuildKit includes file modes in cache keys. Export the pinned revision to
-    # an owned context instead of changing modes in the shared checkout.
-    if ! source_root=$(git rev-parse --show-toplevel); then
-        echo "Failed to locate the CI source checkout" >&2
-        return 1
-    fi
-    validate_ci_build_context_source "${source_root}" || return $?
-    context_root="${SCRIPT_TMP_DIR}/docker-context"
-    if ! mkdir -m 0700 -- "${context_root}"; then
-        echo "Failed to create the owned CI Docker context" >&2
-        return 1
-    fi
-    if ! (
-        umask 0022
-        unset GIT_LFS_SKIP_SMUDGE
-        export GIT_ATTR_NOSYSTEM=1
-        export GIT_INDEX_FILE="${ROCM_BUILD_CONTEXT_INDEX}"
-        git -C "${source_root}" \
-            -c core.attributesFile=/dev/null \
-            -c core.autocrlf=false \
-            -c core.eol=lf \
-            -c core.symlinks=true \
-            --work-tree="${context_root}" \
-            checkout-index --all
-    ); then
-        echo "Failed to materialize the CI Docker context from the Git index" >&2
-        return 1
-    fi
-
-    # setuptools-scm understands Git's stable archive format, so wheels and
-    # Rust artifacts retain their exact version without copying Git history.
-    if ! is_ci_base_target; then
-        write_ci_git_archival_metadata "${source_root}" "${context_root}" \
-            || return $?
-    fi
-    if [[ -e "${context_root}/.git" ]]; then
-        echo "Canonical CI Docker context unexpectedly contains .git" >&2
-        return 1
-    fi
-    ROCM_BUILD_CONTEXT_ROOT="${context_root}"
-    BAKE_ALLOW_ARGS+=(--allow "fs.read=${ROCM_BUILD_CONTEXT_ROOT}")
-    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}}"
-}
-
-compose_dependency_cache_key() {
-    local prefix="$1"
-    local material="$2"
-    local cleaned_prefix=""
-
-    cleaned_prefix=$(clean_docker_tag "${prefix}" | cut -c1-96)
-    printf '%s-%s\n' "${cleaned_prefix}" "$(hash_string_short "${material}")"
-}
-
-hash_dockerfile_stages() {
-    local dockerfile="$1"
-    local stages="$2"
-    local physical_dockerfile=""
-
-    physical_dockerfile="${dockerfile}"
-    if [[ -n "${ROCM_BUILD_CONTEXT_ROOT:-}" && "${dockerfile}" != /* ]]; then
-        physical_dockerfile="${ROCM_BUILD_CONTEXT_ROOT}/${dockerfile}"
-    fi
-
-    awk -v wanted_stages="${stages}" '
-        BEGIN {
-            split(wanted_stages, stage_list, /[[:space:]]+/)
-            for (idx in stage_list) {
-                if (stage_list[idx] != "") {
-                    wanted[stage_list[idx]] = 1
-                }
-            }
-            emit = 0
-        }
-        toupper($1) == "FROM" {
-            stage = ""
-            for (idx = 1; idx <= NF; idx++) {
-                if (tolower($idx) == "as" && idx < NF) {
-                    stage = tolower($(idx + 1))
-                }
-            }
-            emit = (stage in wanted)
-        }
-        emit {
-            print
-        }
-    ' "${physical_dockerfile}"
-}
-
-discover_dockerfile_stage_args() {
-    local dockerfile="$1"
-    local stages="$2"
-    local physical_dockerfile=""
-
-    physical_dockerfile="${dockerfile}"
-    if [[ -n "${ROCM_BUILD_CONTEXT_ROOT:-}" && "${dockerfile}" != /* ]]; then
-        physical_dockerfile="${ROCM_BUILD_CONTEXT_ROOT}/${dockerfile}"
-    fi
-    [[ -f "${physical_dockerfile}" ]] || return 0
-
-    awk -v wanted_stages="${stages}" '
-        function add_arg(name) {
-            if (name != "" && !(name in seen)) {
-                seen[name] = 1
-                args[++arg_count] = name
-            }
-        }
-        BEGIN {
-            split(wanted_stages, stage_list, /[[:space:]]+/)
-            for (idx in stage_list) {
-                if (stage_list[idx] != "") {
-                    wanted[stage_list[idx]] = 1
-                }
-            }
-            emit = 1
-        }
-        {
-            line = $0
-            if (toupper($1) == "FROM") {
-                stage = ""
-                for (idx = 1; idx <= NF; idx++) {
-                    if (tolower($idx) == "as" && idx < NF) {
-                        stage = tolower($(idx + 1))
-                    }
-                }
-                emit = (stage in wanted)
-            }
-            if (emit) {
-                lines[++line_count] = line
-            }
-        }
-        END {
-            for (idx = 1; idx <= line_count; idx++) {
-                line = lines[idx]
-                arg_name = line
-                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_]*$/) {
-                        add_arg(arg_name)
-                    }
-                }
-            }
-
-            for (idx = 1; idx <= line_count; idx++) {
-                line = lines[idx]
-                for (arg_idx = 1; arg_idx <= arg_count; arg_idx++) {
-                    name = args[arg_idx]
-                    if (line ~ "\\$\\{" name "([}:][^}]*)?\\}" \
-                        || line ~ "\\$" name "([^A-Za-z0-9_]|$)") {
-                        used[name] = 1
-                    }
-                }
-            }
-
-            for (arg_idx = 1; arg_idx <= arg_count; arg_idx++) {
-                name = args[arg_idx]
-                if (used[name]) {
-                    print name
-                }
-            }
-        }
-    ' "${physical_dockerfile}"
-}
-
-get_content_arg_names() {
-    local dockerfile="$1"
-    local stages="$2"
-    local explicit_args="${3:-}"
-
-    if [[ -n "${explicit_args}" ]]; then
-        tr ' ' '\n' <<< "${explicit_args}"
-    else
-        discover_dockerfile_stage_args "${dockerfile}" "${stages}"
-    fi | awk 'NF && !seen[$0]++'
-}
-
-compute_ci_base_content_hash() {
-    local -a content_paths=()
-    local -a content_args=()
-    local content_files_hash=""
-    local dockerfile="${CI_BASE_DOCKERFILE:-}"
-    local stages="${CI_BASE_DOCKERFILE_STAGES:-}"
-
-    read -r -a content_paths <<< "${CI_BASE_CONTENT_FILES}"
-    mapfile -t content_args < <(
-        get_content_arg_names "${dockerfile}" "${stages}" "${CI_BASE_CONTENT_ARGS:-}"
-    )
-    if ! content_files_hash=$(compute_content_hash "${content_paths[@]}"); then
-        echo "Failed to hash ci_base content files" >&2
-        return 1
-    fi
-
-    {
-        printf 'content-files-hash:%s\n' "${content_files_hash}"
-        if [[ -n "${dockerfile}" ]]; then
-            printf 'dockerfile:%s\n' "${dockerfile}"
-            printf 'resolved-build-args:\n'
-            hash_dockerfile_arg_values "${dockerfile}" "${content_args[@]}" \
-                || return 1
-            if [[ -n "${stages}" ]]; then
-                printf 'dockerfile-stages:%s\n' "${stages}"
-                if content_regular_file "${dockerfile}"; then
-                    hash_dockerfile_stages "${dockerfile}" "${stages}"
-                else
-                    printf 'missing:%s\n' "${dockerfile}"
-                fi
-            fi
-        fi
-    } | sha256sum | cut -d' ' -f1
-}
-
-extract_dockerfile_arg_default() {
-    local dockerfile="$1"
-    local arg_name="$2"
-    local physical_dockerfile=""
-
-    physical_dockerfile="${dockerfile}"
-    if [[ -n "${ROCM_BUILD_CONTEXT_ROOT:-}" && "${dockerfile}" != /* ]]; then
-        physical_dockerfile="${ROCM_BUILD_CONTEXT_ROOT}/${dockerfile}"
-    fi
-    sed -n -E "s/^[[:space:]]*[Aa][Rr][Gg][[:space:]]+${arg_name}=\"?([^\"[:space:]]+)\"?.*/\\1/p" \
-        "${physical_dockerfile}" | head -1
-}
-
-resolve_image_digest() {
-    local image_ref="$1"
-    local attempts="${ROCM_IMAGE_DIGEST_ATTEMPTS:-4}"
-    local delay_secs="${ROCM_IMAGE_DIGEST_RETRY_DELAY:-2}"
-    local attempt=0
-    local digest=""
-    local output=""
-    local status=0
-
-    if [[ "${image_ref}" =~ @(sha256:[0-9a-f]{64})$ ]]; then
-        printf '%s\n' "${BASH_REMATCH[1]}"
-        return
-    fi
-    if [[ ! "${attempts}" =~ ^[1-9][0-9]*$ \
-        || ! "${delay_secs}" =~ ^[0-9]+$ ]]; then
-        echo "Invalid image digest retry configuration" >&2
-        return 1
-    fi
-
-    for ((attempt = 1; attempt <= attempts; attempt++)); do
-        status=0
-        output=$(docker buildx imagetools inspect "${image_ref}" 2>&1) || status=$?
-        digest=$(awk '$1 == "Digest:" { print $2; exit }' <<< "${output}")
-        if ((status == 0)) && [[ "${digest}" =~ ^sha256:[0-9a-f]{64}$ ]]; then
-            printf '%s\n' "${digest}"
-            return
-        fi
-        if ((attempt < attempts)); then
-            printf \
-                'Image digest lookup failed for %s (%d/%d, status %d); retrying\n' \
-                "${image_ref}" "${attempt}" "${attempts}" "${status}" >&2
-            sleep "${delay_secs}"
-        fi
-    done
-
-    printf 'Failed to resolve digest for %s (status %d)\n%s\n' \
-        "${image_ref}" "${status}" "${output:-<no output>}" >&2
-    return 1
-}
-
-resolve_dockerfile_arg_value() {
-    local dockerfile="$1"
-    local arg_name="$2"
-    local env_name="${arg_name}"
-    local value=""
-
-    case "${arg_name}" in
-        ARG_PYTORCH_ROCM_ARCH)
-            env_name="PYTORCH_ROCM_ARCH"
-            ;;
-        max_jobs)
-            env_name="CI_MAX_JOBS"
-            ;;
-    esac
-
-    value="${!env_name:-}"
-    if [[ -z "${value}" && "${env_name}" != "${arg_name}" ]]; then
-        value="${!arg_name:-}"
-    fi
-    if [[ -z "${value}" ]] && content_regular_file "${dockerfile}"; then
-        value=$(extract_dockerfile_arg_default "${dockerfile}" "${arg_name}")
-    fi
-
-    printf '%s\n' "${value}"
-}
-
-hash_dockerfile_arg_values() {
-    local dockerfile="$1"
-    local arg_name=""
-    local arg_value=""
-    local digest=""
-    shift || true
-
-    for arg_name in "$@"; do
-        [[ -n "${arg_name}" ]] || continue
-        arg_value=$(resolve_dockerfile_arg_value "${dockerfile}" "${arg_name}")
-        if [[ "${arg_name}" == "BASE_IMAGE" && -n "${arg_value}" ]]; then
-            if ! digest=$(resolve_image_digest "${arg_value}"); then
-                echo "Failed to resolve digest for BASE_IMAGE=${arg_value}" >&2
-                return 1
-            fi
-            printf 'arg:%s.digest=%s\n' "${arg_name}" "${digest}"
-        else
-            printf 'arg:%s=%s\n' "${arg_name}" "${arg_value:-<empty>}"
-        fi
-    done
-}
-
-pin_base_image() {
-    local dockerfile="${CI_BASE_DOCKERFILE:-${DEFAULT_CI_BASE_DOCKERFILE}}"
-    local base_image=""
-    local digest=""
-
-    base_image=$(resolve_dockerfile_arg_value "${dockerfile}" "BASE_IMAGE")
-    [[ -n "${base_image}" ]] || return 0
-    if ! digest=$(resolve_image_digest "${base_image}"); then
-        echo "Error: could not resolve base image digest for ${base_image}" >&2
-        echo "Refusing to compute content-addressed cache keys from a mutable tag." >&2
-        return 1
-    fi
-    BASE_IMAGE="${base_image%@*}@${digest}"
-    export BASE_IMAGE
-    echo "Pinned base image for this build: ${BASE_IMAGE}"
-}
-
-is_ci_base_target() {
-    [[ "${TARGET}" == *"ci-base-rocm"* ]]
-}
-
-is_commit_image_target() {
-    [[ -n "${IMAGE_TAG:-}" && -n "${BUILDKITE_COMMIT:-}" ]] || return 1
-    is_ci_base_target && return 1
-    return 0
-}
-
-should_upload_wheel_artifacts() {
-    [[ "${UPLOAD_ROCM_WHEEL_ARTIFACTS:-0}" == "1" ]] && return 0
-    [[ "${TARGET}" == *"with-wheel"* \
-        || "${TARGET}" == *"export-wheel"* \
-        || "${TARGET}" == *"artifact"* ]]
-}
-
-should_export_rocm_smoke() {
-    [[ "${TARGET}" == "test-rocm-ci-with-wheel" \
-        || "${TARGET}" == "smoke-test-rocm-ci" ]]
-}
-
-should_export_kernel_symbol_map() {
-    [[ "${VLLM_KERNEL_SYMBOL_MAP:-0}" == "1" ]] || return 1
-    case "${TARGET}" in
-        test-rocm-ci|test-rocm-ci-with-wheel|test-rocm-ci-with-artifacts|\
-        export-wheel-rocm|csrc-rocm-ci|smoke-test-rocm-ci|kernel-symbol-map-rocm)
-            return 0 ;;
-        *) return 1 ;;
-    esac
-}
-
-export_kernel_symbol_map() {
-    should_export_kernel_symbol_map || return 0
-    echo "--- :world_map: Exporting ROCm kernel symbol map"
-    rm -rf ./kernel-symbol-map-rocm || return 1
-    docker buildx bake \
-        "${BAKE_ALLOW_ARGS[@]}" "${BAKE_FILES[@]}" \
-        --progress "${BUILDKIT_PROGRESS:-plain}" \
-        kernel-symbol-map-rocm || return 1
-    upload_kernel_symbol_map_if_present
-}
-
-export_kernel_symbol_map_from_cache() {
-    should_export_kernel_symbol_map || return 0
-    # A separate shell keeps setup failures fatal inside the optional export.
-    FORCE_BUILD=1 VLLM_KERNEL_SYMBOL_MAP=1 UPLOAD_ROCM_WHEEL_ARTIFACTS=0 \
-        CI_HCL_SOURCE="${CI_HCL_PATH}" VLLM_BAKE_FILE="${VLLM_BAKE_FILE}" \
-        BUILDER_NAME="${BUILDER_NAME}" \
-        bash "${BASH_SOURCE[0]}" kernel-symbol-map-rocm \
-        || echo "ROCm kernel symbol map export from cache failed; continuing without it" >&2
-}
-
-upload_kernel_symbol_map_if_present() {
-    should_export_kernel_symbol_map || return 0
-    local map_dir="./kernel-symbol-map-rocm"
-    local map_name="kernel_symbol_map.rocm.json.gz"
-    if [[ ! -s "${map_dir}/${map_name}" ]]; then
-        echo "ROCm kernel symbol map export is missing: ${map_dir}/${map_name}" >&2
-        return 1
-    fi
-    if [[ "${BUILDKITE:-false}" == "true" ]]; then
-        (cd "${map_dir}" && buildkite-agent artifact upload "${map_name}")
-    fi
-}
-
-verify_rocm_smoke_export() {
-    local marker="./build/rocm-smoke-export/vllm-smoke-ok"
-    local expected_smoke_id="${BUILDKITE_BUILD_ID:-local}"
-    local actual_smoke_id=""
-
-    should_export_rocm_smoke || return 0
-    if [[ ! -f "${marker}" ]]; then
-        echo "ROCm BuildKit smoke marker is missing: ${marker}" >&2
-        return 1
-    fi
-    actual_smoke_id="$(< "${marker}")"
-    if [[ "${actual_smoke_id}" != "${expected_smoke_id}" ]]; then
-        echo "ROCm BuildKit smoke marker belongs to ${actual_smoke_id}, not ${expected_smoke_id}" \
-            >&2
-        return 1
-    fi
-}
-
-get_remote_image_label() {
-    local image_ref="$1"
-    local label_key="$2"
-    local format="{{ index .Image.Config.Labels \"${label_key}\" }}"
-
-    docker buildx imagetools inspect "${image_ref}" \
-        --format "${format}" 2>/dev/null | awk 'NF { print; exit }' || true
-}
-
-remote_ci_base_identity_is_current_with_retry() {
-    local image_ref="$1"
-    local attempts="${CI_BASE_LABEL_ATTEMPTS:-4}"
-    local delay_secs="${CI_BASE_LABEL_RETRY_DELAY:-2}"
-    local identity=""
-    local remote_hash=""
-    local remote_version=""
-    local content_files_hash=""
-    local base_digest=""
-    local attempt=0
-    local expected_version="${CI_BASE_METADATA_VERSION:-${DEFAULT_CI_BASE_METADATA_VERSION}}"
-    local format='{{ index .Image.Config.Labels "vllm.ci_base.content_hash" }}|{{ index .Image.Config.Labels "vllm.ci_base.metadata_version" }}|{{ index .Image.Config.Labels "vllm.ci_base.content_files_hash" }}|{{ index .Image.Config.Labels "vllm.rocm.base_image_digest" }}'
-
-    if [[ ! "${attempts}" =~ ^[1-9][0-9]*$ \
-        || ! "${delay_secs}" =~ ^[0-9]+$ ]]; then
-        echo "Invalid ci_base label retry configuration" >&2
-        return 1
-    fi
-
-    for ((attempt = 1; attempt <= attempts; attempt++)); do
-        identity=$(docker buildx imagetools inspect "${image_ref}" \
-            --format "${format}" 2>/dev/null || true)
-        IFS='|' read -r \
-            remote_hash remote_version content_files_hash base_digest <<< "${identity}"
-        if [[ "${remote_hash}" == "${CI_BASE_CONTENT_HASH:-}" \
-            && "${remote_version}" == "${expected_version}" \
-            && "${content_files_hash}" =~ ^[0-9a-f]{64}$ \
-            && "${base_digest}" =~ ^sha256:[0-9a-f]{64}$ ]]; then
-            return 0
-        fi
-        ((attempt == attempts)) || sleep "${delay_secs}"
-    done
-
-    echo "ci_base identity did not match after ${attempts} attempts: ${image_ref}" >&2
-    echo "  expected hash/version: ${CI_BASE_CONTENT_HASH:-<missing>}/${expected_version}" >&2
-    echo "  observed hash/version: ${remote_hash:-<missing>}/${remote_version:-<missing>}" >&2
-    return 1
-}
-
-registry_ref_exists_with_retry() {
-    local inspect_kind="$1"
-    local image_ref="$2"
-    local attempts="${ROCM_REGISTRY_PROBE_ATTEMPTS:-2}"
-    local delay_secs="${ROCM_REGISTRY_PROBE_RETRY_DELAY:-1}"
-    local output=""
-    local status=0
-    local attempt=0
-
-    if [[ ! "${attempts}" =~ ^[1-9][0-9]*$ \
-        || ! "${delay_secs}" =~ ^[0-9]+$ ]]; then
-        echo "Invalid registry probe retry configuration" >&2
-        return 1
-    fi
-
-    for ((attempt = 1; attempt <= attempts; attempt++)); do
-        status=0
-        if [[ "${inspect_kind}" == "manifest" ]]; then
-            output=$(docker manifest inspect "${image_ref}" 2>&1) || status=$?
-        else
-            output=$(docker buildx imagetools inspect "${image_ref}" 2>&1) \
-                || status=$?
-        fi
-        if ((status == 0)); then
-            return 0
-        fi
-        # A definitive registry miss is normal for a new content key. Retry
-        # transport/rate-limit failures, but do not add latency to a real 404.
-        if grep -Eiq \
-            'manifest unknown|no such manifest|name unknown|not found|does not exist' \
-            <<< "${output}"; then
-            return 1
-        fi
-        if ((attempt < attempts)); then
-            echo "Registry probe failed for ${image_ref} (${attempt}/${attempts}); retrying" >&2
-            sleep "${delay_secs}"
-        fi
-    done
-
-    echo "Registry probe failed after ${attempts} attempts: ${image_ref}" >&2
-    return 1
-}
-
-remote_image_exists() {
-    registry_ref_exists_with_retry manifest "$1"
-}
-
-use_existing_builder() {
-    echo "Using existing builder: ${BUILDER_NAME}"
-    docker buildx use "${BUILDER_NAME}"
-    docker buildx inspect --bootstrap
-}
-
-buildx_driver() {
-    local builder="${1:-}"
-
-    if [[ -n "${builder}" ]]; then
-        docker buildx inspect "${builder}" 2>/dev/null
-    else
-        docker buildx inspect 2>/dev/null
-    fi | awk -F': *' '$1 == "Driver" { print $2; exit }'
-}
-
-builder_supports_registry_cache() {
-    local driver="$1"
-
-    [[ -n "${driver}" && "${driver}" != "docker" ]]
-}
-
-create_and_bootstrap_builder() {
-    local driver="$1"
-    local endpoint="${2:-}"
-
-    echo "Creating builder '${BUILDER_NAME}' with ${driver} driver"
-    if [[ -n "${endpoint}" ]]; then
-        docker buildx create \
-            --name "${BUILDER_NAME}" \
-            --driver "${driver}" \
-            --use \
-            "${endpoint}"
-    else
-        docker buildx create --name "${BUILDER_NAME}" --driver "${driver}" --use
-    fi
-    docker buildx inspect --bootstrap
-}
-
-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}}}"
-    VLLM_BAKE_FILE="${VLLM_BAKE_FILE:-docker/docker-bake-rocm.hcl}"
-    BUILDER_NAME="${BUILDER_NAME:-vllm-builder}"
-    BUILDKIT_SOCKET="${BUILDKIT_SOCKET:-/run/buildkit/buildkitd.sock}"
-    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}}"
-    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 CI_BASE_DOCKERFILE
-
-    SCRIPT_TMP_DIR=$(mktemp -d -t ci-bake-rocm.XXXXXX)
-    CI_HCL_PATH="${SCRIPT_TMP_DIR}/ci.hcl"
-    CI_BASE_LABEL_OVERRIDE_PATH="${SCRIPT_TMP_DIR}/ci-base-label-override.hcl"
-    CSRC_CACHE_OVERRIDE_PATH="${SCRIPT_TMP_DIR}/rocm-csrc-cache-override.hcl"
-    ROCM_ARG_OVERRIDE_PATH="${SCRIPT_TMP_DIR}/rocm-arg-override.hcl"
-    BUILD_CONTEXT_OVERRIDE_PATH="${SCRIPT_TMP_DIR}/build-context-override.hcl"
-    BAKE_CONFIG_FILE="bake-config-build-${BUILDKITE_BUILD_NUMBER:-local}.json"
-}
-
-print_header() {
-    echo "--- :docker: Setting up Docker buildx bake"
-    echo "Target: ${TARGET}"
-    echo "CI HCL source: ${CI_HCL_SOURCE}"
-    echo "vLLM bake file: ${VLLM_BAKE_FILE}"
-    if is_ci_base_target; then
-        echo "Build mode: ci_base"
-    elif is_commit_image_target; then
-        echo "Build mode: build-scoped commit image"
-    else
-        echo "Build mode: generic"
-    fi
-    if [[ "${USE_SCCACHE:-0}" == "1" ]]; then
-        echo "Compiler cache: sccache enabled"
-    fi
-}
-
-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"
-        exit 1
-    fi
-
-    if [[ -n "${CI_HCL_SOURCE:-}" ]] && is_url_like "${CI_HCL_SOURCE}"; then
-        echo "Error: remote CI HCL sources are not supported: ${CI_HCL_SOURCE}"
-        echo "Use the vLLM-owned docker/ci-rocm.hcl or set CI_HCL_SOURCE to a local file."
-        exit 1
-    fi
-
-    if [[ -n "${CI_HCL_SOURCE:-}" && ! -f "${CI_HCL_SOURCE}" ]]; then
-        echo "Error: CI HCL file not found at ${CI_HCL_SOURCE}"
-        echo "Set CI_HCL_SOURCE to a local file if you need an override."
-        exit 1
-    fi
-}
-
-load_ci_hcl() {
-    echo "--- :page_facing_up: Loading ci.hcl"
-    cp "${CI_HCL_SOURCE}" "${CI_HCL_PATH}"
-    echo "Copied ${CI_HCL_SOURCE} to ${CI_HCL_PATH}"
-}
-
-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() {
-    if [[ -z "${CI_BASE_CONTENT_FILES:-}" ]]; then
-        return 0
-    fi
-    pin_base_image
-    if ! is_ci_base_target; then
-        return 0
-    fi
-    if [[ "${REMOTE_VLLM:-0}" != "0" ]]; then
-        echo "Error: content-addressed ci_base builds require REMOTE_VLLM=0" >&2
-        return 1
-    fi
-
-    CI_BASE_CONTENT_HASH=$(compute_ci_base_content_hash)
-    export CI_BASE_CONTENT_HASH
-    echo "ci_base content hash: ${CI_BASE_CONTENT_HASH:0:16}..."
-}
-
-wants_stable_ci_base_tag() {
-    if [[ "${BUILDKITE_PULL_REQUEST:-false}" != "false" ]]; then
-        return 1
-    fi
-
-    if [[ "${CI_BASE_PUSH_STABLE_TAG:-}" == "1" ]]; then
-        return 0
-    fi
-    if [[ "${CI_BASE_PUSH_STABLE_TAG:-}" == "0" ]]; then
-        return 1
-    fi
-
-    [[ "${NIGHTLY:-0}" == "1" && "${BUILDKITE_BRANCH:-}" == "${CI_BASE_STABLE_BRANCH:-main}" ]]
-}
-
-trusted_ci_base_tip_matches_build() {
-    local branch="${CI_BASE_STABLE_BRANCH:-main}"
-    local build_commit="${BUILDKITE_COMMIT:-}"
-    local remote_tip=""
-
-    is_trusted_ci_cache_writer || return 1
-    if [[ ! "${build_commit}" =~ ^[0-9a-fA-F]{40}$ ]]; then
-        echo "Skipping ci_base stable tag: Buildkite commit is missing or invalid" >&2
-        return 1
-    fi
-    remote_tip=$(git ls-remote --exit-code "${BUILDKITE_REPO}" \
-        "refs/heads/${branch}" 2>/dev/null | awk 'NR == 1 { print $1 }')
-    if [[ ! "${remote_tip}" =~ ^[0-9a-fA-F]{40}$ ]]; then
-        echo "Skipping ci_base stable tag: could not resolve remote ${branch} tip" >&2
-        return 1
-    fi
-    if [[ "${remote_tip,,}" != "${build_commit,,}" ]]; then
-        echo "Skipping ci_base stable tag: ${branch} advanced from ${build_commit} to ${remote_tip}" >&2
-        return 1
-    fi
-}
-
-should_push_stable_ci_base_tag() {
-    wants_stable_ci_base_tag \
-        && is_trusted_ci_cache_writer \
-        && trusted_ci_base_tip_matches_build
-}
-
-ci_base_tag_with_suffix() {
-    local base_tag="$1"
-    local suffix="$2"
-
-    printf '%s-%s\n' "${base_tag}" "$(clean_docker_tag "${suffix}")"
-}
-
-configure_ci_base_image_refs() {
-    local stable_tag="${CI_BASE_IMAGE_TAG:-rocm/vllm-dev:ci_base}"
-    local metadata_version="${CI_BASE_METADATA_VERSION:-${DEFAULT_CI_BASE_METADATA_VERSION}}"
-    local scope="${CI_BASE_WRITE_SCOPE:-}"
-    local trusted_content_tag=""
-    local content_tag=""
-    local build_tag=""
-    local primary_tag=""
-
-    if [[ ! "${metadata_version}" =~ ^[A-Za-z0-9][A-Za-z0-9_.-]{0,15}$ ]]; then
-        echo "Invalid ci_base metadata version: ${metadata_version}" >&2
-        return 1
-    fi
-
-    if [[ "${BUILDKITE:-false}" == "true" ]] \
-        && ! is_full_git_sha "${BUILDKITE_COMMIT:-}"; then
-        echo "Invalid Buildkite commit for ci_base handoff: ${BUILDKITE_COMMIT:-<empty>}" >&2
-        return 1
-    fi
-    if [[ "${BUILDKITE:-false}" == "true" \
-        && -z "${BUILDKITE_BUILD_ID:-}" ]]; then
-        echo "Buildkite build ID is required for the ci_base runtime handoff" >&2
-        return 1
-    fi
-    if [[ -n "${BUILDKITE_BUILD_ID:-}" ]]; then
-        build_tag=$(ci_base_tag_with_suffix \
-            "${stable_tag}" "build-${BUILDKITE_BUILD_ID}")
-    fi
-    CI_BASE_IMAGE_TAG_BUILD_REF="${build_tag}"
-    export CI_BASE_IMAGE_TAG_BUILD_REF
-
-    if [[ -z "${CI_BASE_CONTENT_HASH:-}" ]]; then
-        if is_ci_base_target; then
-            echo "Error: ci_base builds require a content hash" >&2
-            return 1
-        fi
-        CI_BASE_IMAGE="${CI_BASE_IMAGE:-${stable_tag}}"
-        export CI_BASE_IMAGE
-        return 0
-    fi
-
-    trusted_content_tag=$(ci_base_tag_with_suffix \
-        "${stable_tag}" "v${metadata_version}-${CI_BASE_CONTENT_HASH}")
-    if [[ -n "${scope}" ]]; then
-        content_tag=$(ci_base_tag_with_suffix \
-            "${stable_tag}" "v${metadata_version}-${scope}-${CI_BASE_CONTENT_HASH}")
-    else
-        content_tag="${trusted_content_tag}"
-    fi
-    CI_BASE_IMAGE_TAG_CONTENT_REF="${content_tag}"
-    CI_BASE_TRUSTED_CONTENT_REF="${trusted_content_tag}"
-    primary_tag="${build_tag:-${content_tag}}"
-
-    # Main writes canonical content refs. Other builds import those refs
-    # read-only and publish into a source-scoped preview namespace.
-    # A fresh CI build first publishes its unique runtime tag. Content and
-    # stable aliases are created only after that image passes identity checks.
-    CI_BASE_STABLE_PROMOTION_REF="${stable_tag}"
-    CI_BASE_IMAGE_TAG="${primary_tag}"
-    export CI_BASE_IMAGE_TAG
-    export CI_BASE_IMAGE_TAG_CONTENT_REF
-    export CI_BASE_TRUSTED_CONTENT_REF
-    export CI_BASE_STABLE_PROMOTION_REF
-
-    if is_ci_base_target; then
-        IMAGE_TAG="${primary_tag}"
-        CI_BASE_IMAGE="${primary_tag}"
-        export CI_BASE_IMAGE
-        export IMAGE_TAG
-
-        echo "ci_base primary image tag: ${CI_BASE_IMAGE_TAG}"
-        if [[ -n "${build_tag}" ]]; then
-            echo "ci_base build image tag: ${build_tag}"
-        fi
-        echo "ci_base content image tag: ${content_tag}"
-        if wants_stable_ci_base_tag && is_trusted_ci_cache_writer; then
-            echo "ci_base stable alias is eligible for post-build promotion: ${stable_tag}"
-        else
-            echo "ci_base stable alias will not be pushed for this build"
-            echo "Set NIGHTLY=1 on ${CI_BASE_STABLE_BRANCH:-main} to refresh ${stable_tag}"
-        fi
-        return 0
-    fi
-
-    if [[ -z "${CI_BASE_IMAGE:-}" || "${CI_BASE_IMAGE}" == "${stable_tag}" ]]; then
-        CI_BASE_IMAGE="${content_tag}"
-        export CI_BASE_IMAGE
-        echo "Using ci_base image: ${CI_BASE_IMAGE}"
-    else
-        echo "Using provided CI_BASE_IMAGE override: ${CI_BASE_IMAGE}"
-    fi
-}
-
-publish_ci_base_handoff_ref() {
-    local source_ref="${1:-${CI_BASE_IMAGE_TAG_BUILD_REF:-${CI_BASE_IMAGE_TAG_CONTENT_REF:-}}}"
-    local content_ref="${CI_BASE_IMAGE_TAG_CONTENT_REF:-}"
-    local digest=""
-    local handoff_ref=""
-
-    is_ci_base_target || return 0
-    if [[ -z "${source_ref}" || -z "${content_ref}" ]]; then
-        echo "Cannot publish ci_base handoff without source and content refs" >&2
-        return 1
-    fi
-    if ! digest=$(resolve_image_digest "${source_ref}"); then
-        echo "Could not resolve immutable ci_base handoff: ${source_ref}" >&2
-        return 1
-    fi
-
-    handoff_ref="${content_ref%@*}@${digest}"
-    if ! confirm_remote_image_push "${handoff_ref}"; then
-        echo "Could not validate immutable ci_base handoff: ${handoff_ref}" >&2
-        return 1
-    fi
-    if command -v buildkite-agent >/dev/null 2>&1; then
-        if ! buildkite-agent meta-data set "rocm-ci-base-image" "${handoff_ref}"; then
-            echo "Could not publish required ci_base handoff metadata" >&2
-            return 1
-        fi
-    elif [[ "${BUILDKITE:-false}" == "true" ]]; then
-        echo "buildkite-agent not found; cannot publish ci_base handoff" >&2
-        return 1
-    fi
-    echo "Published immutable ci_base handoff: ${handoff_ref}"
-}
-
-ci_base_output_refs() {
-    printf '%s\n' \
-        "${CI_BASE_IMAGE_TAG_CONTENT_REF:-}" \
-        "${CI_BASE_IMAGE_TAG_BUILD_REF:-}" \
-        | awk 'NF && !seen[$0]++'
-}
-
-ci_base_candidate_refs() {
-    printf '%s\n' \
-        "${CI_BASE_TRUSTED_CONTENT_REF:-}" \
-        "${CI_BASE_IMAGE_TAG_CONTENT_REF:-}" \
-        "${CI_BASE_STABLE_PROMOTION_REF:-}" \
-        | awk 'NF && !seen[$0]++'
-}
-
-find_matching_ci_base_ref() {
-    local candidate=""
-    local candidate_digest=""
-    local immutable_candidate=""
-
-    while IFS= read -r candidate; do
-        [[ -n "${candidate}" ]] || continue
-        # Missing content refs are normal after an input change. Avoid paying
-        # the full digest retry budget for each definitely absent candidate.
-        remote_image_exists "${candidate}" || continue
-        if ! candidate_digest=$(resolve_image_digest "${candidate}"); then
-            echo "Could not pin ci_base candidate: ${candidate}" >&2
-            continue
-        fi
-        immutable_candidate="${candidate%@*}@${candidate_digest}"
-        if remote_ci_base_identity_is_current_with_retry "${immutable_candidate}"; then
-            printf '%s\n' "${immutable_candidate}"
-            return 0
-        fi
-    done < <(ci_base_candidate_refs)
-
-    return 1
-}
-
-refresh_ci_base_tags_from_ref() {
-    local source_ref="$1"
-    local source_digest=""
-    local immutable_source=""
-    local tag=""
-    local tag_digest=""
-
-    if ! source_digest=$(resolve_image_digest "${source_ref}"); then
-        echo "Could not resolve selected ci_base image: ${source_ref}" >&2
-        return 1
-    fi
-    immutable_source="${source_ref%@*}@${source_digest}"
-
-    while IFS= read -r tag; do
-        [[ -n "${tag}" ]] || continue
-        [[ "${tag}" != "${source_ref}" ]] || continue
-        tag_digest=""
-        if remote_image_exists "${tag}"; then
-            tag_digest=$(resolve_image_digest "${tag}" || true)
-        fi
-        if [[ "${tag_digest}" == "${source_digest}" ]]; then
-            echo "ci_base tag is already current: ${tag}"
-            continue
-        fi
-        echo "Updating ci_base tag ${tag} -> ${immutable_source}"
-        if ! docker buildx imagetools create --prefer-index=false \
-            -t "${tag}" "${immutable_source}"; then
-            echo "Failed to update ci_base tag ${tag} from ${immutable_source}" >&2
-            return 1
-        fi
-        if ! tag_digest=$(resolve_image_digest "${tag}") \
-            || [[ "${tag_digest}" != "${source_digest}" ]]; then
-            echo "Updated ci_base tag does not resolve to the selected digest: ${tag}" >&2
-            return 1
-        fi
-    done < <(ci_base_output_refs)
-}
-
-promote_stable_ci_base_tag() {
-    local source_ref="${1:-${CI_BASE_IMAGE_TAG_BUILD_REF:-${CI_BASE_IMAGE_TAG_CONTENT_REF:-}}}"
-    local stable_ref="${CI_BASE_STABLE_PROMOTION_REF:-}"
-    local digest=""
-    local immutable_source=""
-
-    is_ci_base_target || return 0
-    wants_stable_ci_base_tag || return 0
-    if ! should_push_stable_ci_base_tag; then
-        echo "Skipping ci_base stable promotion: build is not the trusted current main tip"
-        return 0
-    fi
-    if [[ -z "${source_ref}" || -z "${stable_ref}" ]]; then
-        echo "Cannot promote ci_base stable tag without source and destination refs" >&2
-        return 1
-    fi
-    if ! digest=$(resolve_image_digest "${source_ref}"); then
-        echo "Could not pin ci_base source before stable promotion: ${source_ref}" >&2
-        return 1
-    fi
-    immutable_source="${source_ref%@*}@${digest}"
-    echo "Promoting ci_base stable tag from ${immutable_source}"
-    if ! docker buildx imagetools create --prefer-index=false \
-        -t "${stable_ref}" "${immutable_source}"; then
-        echo "Failed to promote ci_base stable tag" >&2
-        return 1
-    fi
-    if ! confirm_remote_image_push "${stable_ref}"; then
-        echo "Promoted ci_base stable tag did not become visible: ${stable_ref}" >&2
-        return 1
-    fi
-}
-
-maybe_reuse_matching_ci_base_ref() {
-    local matching_ref=""
-
-    matching_ref=$(find_matching_ci_base_ref) || return 1
-
-    echo "Found existing ci_base image with matching content hash: ${matching_ref}"
-    if ! refresh_ci_base_tags_from_ref "${matching_ref}"; then
-        echo "ci_base tag refresh failed after finding an exact image; aborting" >&2
-        return 2
-    fi
-    if ! promote_stable_ci_base_tag "${matching_ref}"; then
-        echo "ci_base stable promotion failed after finding an exact image; aborting" >&2
-        return 2
-    fi
-    if ! publish_ci_base_handoff_ref "${matching_ref}"; then
-        echo "ci_base handoff failed after finding an exact image; aborting" >&2
-        return 2
-    fi
-    echo "Content hashes match -- ci_base is current"
-    return 0
-}
-
-maybe_skip_existing_image() {
-    local remote_revision=""
-    local reuse_status=0
-
-    if [[ -z "${IMAGE_TAG:-}" ]]; then
-        return 0
-    fi
-
-    if [[ "${FORCE_BUILD:-0}" == "1" ]]; then
-        echo "FORCE_BUILD=1 set; skipping existing-image check"
-        return 0
-    fi
-    if ! is_ci_base_target \
-        && { should_upload_wheel_artifacts || should_export_rocm_smoke \
-            || [[ "${TARGET}" == "kernel-symbol-map-rocm" ]]; }; then
-        echo "Local-output targets always run for the current build"
-        return 0
-    fi
-
-    echo "--- :mag: Checking image tag"
-    echo "Image tag: ${IMAGE_TAG}"
-
-    if is_ci_base_target && [[ -n "${CI_BASE_CONTENT_HASH:-}" ]]; then
-        maybe_reuse_matching_ci_base_ref || reuse_status=$?
-        case "${reuse_status}" in
-            0)
-                echo "Skipping build"
-                exit 0
-                ;;
-            1)
-                echo "No current ci_base image matched the expected content hash"
-                echo "Proceeding with build"
-                return 0
-                ;;
-            *)
-                return "${reuse_status}"
-                ;;
-        esac
-    fi
-
-    if ! remote_image_exists "${IMAGE_TAG}"; then
-        echo "Image not found, proceeding with build"
-        return 0
-    fi
-
-    if is_ci_base_target; then
-        echo "ci_base image already exists and no content hash was configured"
-        echo "Skipping build"
-        exit 0
-    fi
-
-    if is_commit_image_target; then
-        remote_revision=$(get_remote_image_label "${IMAGE_TAG}" "org.opencontainers.image.revision")
-        if [[ -n "${remote_revision}" && "${remote_revision}" != "${BUILDKITE_COMMIT}" ]]; then
-            echo "Existing image revision does not match ${BUILDKITE_COMMIT}"
-            echo "  found revision: ${remote_revision}"
-            echo "Rebuilding image"
-            return 0
-        fi
-
-        echo "Commit image already exists: ${IMAGE_TAG}"
-        echo "Skipping build"
-        export_kernel_symbol_map_from_cache
-        exit 0
-    fi
-
-    echo "Image already exists: ${IMAGE_TAG}"
-    echo "Skipping build"
-    export_kernel_symbol_map_from_cache
-    exit 0
-}
-
-setup_builder() {
-    echo "--- :buildkite: Setting up buildx builder"
-
-    local setup_mode="${ROCM_SETUP_BUILDX_BUILDER:-auto}"
-    local current_driver=""
-    local named_driver=""
-
-    if [[ "${setup_mode}" == "0" || "${setup_mode}" == "false" ]]; then
-        echo "Using current Docker buildx builder"
-        echo "ROCM_SETUP_BUILDX_BUILDER=${setup_mode}; cache exporters may fail if the driver is docker"
-        docker buildx inspect --bootstrap
-        echo "Active builder:"
-        docker buildx ls | grep -E '^\*|^NAME' || docker buildx ls
-        return 0
-    fi
-
-    current_driver=$(buildx_driver || true)
-    if [[ "${setup_mode}" != "1" ]] && builder_supports_registry_cache "${current_driver}"; then
-        echo "Using current Docker buildx builder with ${current_driver} driver"
-        docker buildx inspect --bootstrap
-        echo "Active builder:"
-        docker buildx ls | grep -E '^\*|^NAME' || docker buildx ls
-        return 0
-    fi
-
-    if [[ "${setup_mode}" != "1" ]]; then
-        echo "Current buildx driver '${current_driver:-unknown}' cannot export registry caches"
-        echo "Creating or using a cache-capable builder: ${BUILDER_NAME}"
-    fi
-
-    if docker buildx inspect "${BUILDER_NAME}" >/dev/null 2>&1; then
-        named_driver=$(buildx_driver "${BUILDER_NAME}" || true)
-        if ! builder_supports_registry_cache "${named_driver}"; then
-            echo "Builder '${BUILDER_NAME}' uses ${named_driver:-unknown} driver; using ${BUILDER_NAME}-cache instead"
-            BUILDER_NAME="${BUILDER_NAME}-cache"
-        fi
-    fi
-
-    if [[ -S "${BUILDKIT_SOCKET}" ]]; then
-        echo "Found local buildkitd socket at ${BUILDKIT_SOCKET}"
-        echo "Using remote driver to connect to buildkitd"
-
-        if docker buildx inspect "${BUILDER_NAME}" >/dev/null 2>&1; then
-            use_existing_builder
-        else
-            create_and_bootstrap_builder remote "unix://${BUILDKIT_SOCKET}"
-        fi
-    elif docker buildx inspect "${BUILDER_NAME}" >/dev/null 2>&1; then
-        use_existing_builder
-    else
-        echo "No local buildkitd found, using docker-container driver"
-        create_and_bootstrap_builder docker-container
-    fi
-
-    echo "Active builder:"
-    docker buildx ls | grep -E '^\*|^NAME' || docker buildx ls
-}
-
-validate_cache_branch_tag() {
-    local name="$1"
-    local value="$2"
-
-    if [[ -n "${value}" \
-        && (! "${value}" =~ ^[A-Za-z0-9_][A-Za-z0-9_.-]*$ \
-            || ${#value} -gt ${ROCM_CACHE_BRANCH_TAG_MAX_LEN}) ]]; then
-        echo "Invalid ${name}; expected a Docker tag component of at most ${ROCM_CACHE_BRANCH_TAG_MAX_LEN} characters" >&2
-        return 1
-    fi
-}
-
-prepare_git_cache_metadata() {
-    local cache_branch_name=""
-    local cache_base_branch="${BUILDKITE_PULL_REQUEST_BASE_BRANCH:-main}"
-    local target_repo_slug=""
-    local target_repo_url=""
-    local merge_base_ref=""
-
-    if is_ci_base_target; then
-        echo "Skipping commit-cache ancestry lookup for content-addressed ci_base"
-        return 0
-    fi
-
-    if [[ -z "${PARENT_COMMIT:-}" || -z "${VLLM_MERGE_BASE_COMMIT:-}" ]] \
-        && git rev-parse --is-shallow-repository 2>/dev/null | grep -q "true"; then
-        echo "Shallow clone detected - deepening for cache key computation"
-        git_fetch_for_cache --filter=tree:0 --no-tags --deepen=1 \
-            origin "$(git rev-parse HEAD)" || true
-    fi
-
-    if [[ -z "${PARENT_COMMIT:-}" ]]; then
-        PARENT_COMMIT=$(git rev-parse HEAD~1 2>/dev/null || echo "")
-        if [[ -n "${PARENT_COMMIT}" ]]; then
-            export PARENT_COMMIT
-            echo "Computed parent commit for cache fallback: ${PARENT_COMMIT}"
-        else
-            echo "Could not determine parent commit"
-        fi
-    else
-        echo "Using provided PARENT_COMMIT: ${PARENT_COMMIT}"
-    fi
-
-    if [[ -z "${ROCM_CACHE_BRANCH_TAG:-}" ]]; then
-        cache_branch_name=$(select_cache_branch_name)
-        if [[ -z "${cache_branch_name}" && "${BUILDKITE_PULL_REQUEST:-false}" != "false" ]]; then
-            cache_branch_name="pr-${BUILDKITE_PULL_REQUEST}"
-            echo "Using pull request number for ROCm branch cache tag: ${cache_branch_name}"
-        fi
-    fi
-
-    if [[ -z "${ROCM_CACHE_BRANCH_TAG:-}" && -n "${cache_branch_name}" ]]; then
-        ROCM_CACHE_BRANCH_TAG=$(
-            compose_cache_branch_tag "$(get_buildkite_repo_slug)" "${cache_branch_name}"
-        )
-        export ROCM_CACHE_BRANCH_TAG
-        echo "Computed ROCm branch cache tag: ${ROCM_CACHE_BRANCH_TAG} (from ${cache_branch_name})"
-    elif [[ -n "${ROCM_CACHE_BRANCH_TAG:-}" ]]; then
-        echo "Using provided ROCM_CACHE_BRANCH_TAG: ${ROCM_CACHE_BRANCH_TAG}"
-    elif [[ -n "${BUILDKITE_BRANCH:-}" ]]; then
-        echo "Skipping ROCm branch cache tag: no usable branch name found"
-        echo "  BUILDKITE_BRANCH=${BUILDKITE_BRANCH}"
-    fi
-
-    if [[ -z "${ROCM_CACHE_UPSTREAM_BRANCH_TAG:-}" \
-          && -n "${BUILDKITE_PULL_REQUEST_BASE_BRANCH:-}" \
-          && "${BUILDKITE_PULL_REQUEST:-false}" != "false" ]]; then
-        target_repo_slug=$(get_buildkite_target_repo_slug)
-        ROCM_CACHE_UPSTREAM_BRANCH_TAG=$(
-            compose_cache_branch_tag "${target_repo_slug}" "${BUILDKITE_PULL_REQUEST_BASE_BRANCH}"
-        )
-        export ROCM_CACHE_UPSTREAM_BRANCH_TAG
-        echo "Computed ROCm upstream branch cache tag: ${ROCM_CACHE_UPSTREAM_BRANCH_TAG}"
-    elif [[ -n "${ROCM_CACHE_UPSTREAM_BRANCH_TAG:-}" ]]; then
-        echo "Using provided ROCM_CACHE_UPSTREAM_BRANCH_TAG: ${ROCM_CACHE_UPSTREAM_BRANCH_TAG}"
-    fi
-
-    validate_cache_branch_tag \
-        ROCM_CACHE_BRANCH_TAG "${ROCM_CACHE_BRANCH_TAG:-}"
-    validate_cache_branch_tag \
-        ROCM_CACHE_UPSTREAM_BRANCH_TAG "${ROCM_CACHE_UPSTREAM_BRANCH_TAG:-}"
-
-    if [[ -z "${VLLM_MERGE_BASE_COMMIT:-}" ]]; then
-        target_repo_url=$(get_buildkite_target_repo_url)
-        merge_base_ref="refs/remotes/vllm-cache-upstream/${cache_base_branch}"
-        git_fetch_for_cache --no-tags --depth=200 "${target_repo_url}" \
-            "+refs/heads/${cache_base_branch}:${merge_base_ref}" 2>/dev/null || true
-        VLLM_MERGE_BASE_COMMIT=$(git merge-base HEAD "${merge_base_ref}" 2>/dev/null || echo "")
-        if [[ -z "${VLLM_MERGE_BASE_COMMIT}" ]]; then
-            git_fetch_for_cache --no-tags --deepen=1000 "${target_repo_url}" \
-                "+refs/heads/${cache_base_branch}:${merge_base_ref}" 2>/dev/null || true
-            VLLM_MERGE_BASE_COMMIT=$(git merge-base HEAD "${merge_base_ref}" 2>/dev/null || echo "")
-        fi
-        if [[ -n "${VLLM_MERGE_BASE_COMMIT}" ]]; then
-            export VLLM_MERGE_BASE_COMMIT
-            echo "Computed merge base commit for cache fallback: ${VLLM_MERGE_BASE_COMMIT}"
-        else
-            echo "Could not determine merge base with ${cache_base_branch}"
-        fi
-    else
-        echo "Using provided VLLM_MERGE_BASE_COMMIT: ${VLLM_MERGE_BASE_COMMIT}"
-    fi
-}
-
-ci_base_metadata_pairs() {
-    local dockerfile="${CI_BASE_DOCKERFILE:-${DEFAULT_CI_BASE_DOCKERFILE}}"
-    local stages="${CI_BASE_DOCKERFILE_STAGES:-${DEFAULT_CI_BASE_DOCKERFILE_STAGES}}"
-    local content_files="${CI_BASE_CONTENT_FILES:-${DEFAULT_CI_BASE_CONTENT_FILES}}"
-    local content_files_hash=""
-    local base_image=""
-    local base_image_digest=""
-    local -a content_paths=()
-    local -a content_args=()
-
-    read -r -a content_paths <<< "${content_files}"
-    if [[ ${#content_paths[@]} -gt 0 ]]; then
-        if ! content_files_hash=$(compute_content_hash "${content_paths[@]}"); then
-            echo "Failed to hash ci_base metadata content files" >&2
-            return 1
-        fi
-    fi
-    mapfile -t content_args < <(
-        get_content_arg_names "${dockerfile}" "${stages}" "${CI_BASE_CONTENT_ARGS:-}"
-    )
-
-    base_image=$(resolve_dockerfile_arg_value "${dockerfile}" "BASE_IMAGE")
-    if [[ -n "${base_image}" ]]; then
-        if ! base_image_digest=$(resolve_image_digest "${base_image}"); then
-            echo "Failed to resolve ci_base metadata digest for ${base_image}" >&2
-            return 1
-        fi
-    fi
-
-    metadata_pair "vllm.ci_base.metadata_version" "${CI_BASE_METADATA_VERSION:-${DEFAULT_CI_BASE_METADATA_VERSION}}"
-    metadata_pair "vllm.ci_base.content_hash" "${CI_BASE_CONTENT_HASH:-}"
-    metadata_pair "vllm.ci_base.content_files_hash" "${content_files_hash}"
-    metadata_pair "vllm.ci_base.content_files" "${content_files}"
-    metadata_pair "vllm.ci_base.content_args" "$(join_words "${content_args[@]}")"
-    metadata_pair "vllm.ci_base.dockerfile" "${dockerfile}"
-    metadata_pair "vllm.ci_base.dockerfile_stages" "${stages}"
-
-    # The parent identity is digest-only: mutable aliases that resolve to the
-    # same image must produce byte-identical canonical metadata.
-    metadata_pair "vllm.rocm.base_image" "${base_image_digest}"
-    metadata_pair "vllm.rocm.base_image_digest" "${base_image_digest}"
-    metadata_pair "vllm.rocm.pytorch_rocm_arch" "${PYTORCH_ROCM_ARCH:-}"
-    metadata_pair "vllm.rocm.nic_backend" "$(resolve_dockerfile_arg_value "${dockerfile}" "NIC_BACKEND")"
-    metadata_pair "vllm.rocm.ainic_version" "$(resolve_dockerfile_arg_value "${dockerfile}" "AINIC_VERSION")"
-    metadata_pair "vllm.rocm.ubuntu_codename" "$(resolve_dockerfile_arg_value "${dockerfile}" "UBUNTU_CODENAME")"
-    metadata_pair "vllm.rocm.max_jobs" "$(resolve_dockerfile_arg_value "${dockerfile}" "max_jobs")"
-    metadata_pair "vllm.rocm.nixl_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "NIXL_REPO")"
-    metadata_pair "vllm.rocm.nixl_commit" "${NIXL_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "NIXL_BRANCH")}"
-    metadata_pair "vllm.rocm.ucx_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "UCX_REPO")"
-    metadata_pair "vllm.rocm.ucx_commit" "${UCX_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "UCX_BRANCH")}"
-    metadata_pair "vllm.rocm.deepep_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_REPO")"
-    metadata_pair "vllm.rocm.deepep_commit" "${DEEPEP_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_BRANCH")}"
-    metadata_pair "vllm.rocm.deepep_nic" "$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_NIC")"
-    metadata_pair "vllm.rocm.deepep_rocm_arch" "$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_ROCM_ARCH")"
-    metadata_pair "vllm.rocm.nixl_cache_key" "${NIXL_CACHE_KEY:-}"
-    metadata_pair "vllm.rocm.deepep_cache_key" "${DEEPEP_CACHE_KEY:-}"
-}
-
-write_ci_base_metadata_annotations() {
-    local metadata="$1"
-    local key=""
-    local value=""
-
-    [[ -n "${metadata}" ]] || return 0
-    while IFS=$'\t' read -r key value; do
-        [[ -n "${key}" && -n "${value}" ]] || continue
-        printf '    "%s",\n' \
-            "$(hcl_escape_string "manifest:${key}=${value}")"
-    done <<< "${metadata}"
-}
-
-write_ci_base_metadata_labels() {
-    local metadata="$1"
-    local key=""
-    local value=""
-
-    [[ -n "${metadata}" ]] || return 0
-    while IFS=$'\t' read -r key value; do
-        [[ -n "${key}" && -n "${value}" ]] || continue
-        printf '    "%s" = "%s"\n' \
-            "$(hcl_escape_string "${key}")" \
-            "$(hcl_escape_string "${value}")"
-    done <<< "${metadata}"
-}
-
-write_ci_base_label_override() {
-    local target_name=""
-    local metadata=""
-    local -a ci_base_targets=()
-
-    if [[ -z "${CI_BASE_CONTENT_HASH:-}" ]]; then
-        return 0
-    fi
-
-    mapfile -t ci_base_targets < <(
-        {
-            printf '%s\n' "ci-base-rocm"
-            sed -n -E 's/^target "(ci-base-rocm[^"]+)".*/\1/p' "${CI_HCL_PATH}" 2>/dev/null || true
-        } | awk '!seen[$0]++'
-    )
-
-    if [[ ${#ci_base_targets[@]} -eq 0 ]]; then
-        return 0
-    fi
-
-    metadata=$(ci_base_metadata_pairs)
-
-    : > "${CI_BASE_LABEL_OVERRIDE_PATH}"
-    for target_name in "${ci_base_targets[@]}"; do
-        cat >> "${CI_BASE_LABEL_OVERRIDE_PATH}" <<EOF
-target "${target_name}" {
-  annotations = [
-    "manifest:org.opencontainers.image.revision=",
-EOF
-        write_ci_base_metadata_annotations "${metadata}" >> "${CI_BASE_LABEL_OVERRIDE_PATH}"
-        cat >> "${CI_BASE_LABEL_OVERRIDE_PATH}" <<EOF
-  ]
-  labels = {
-    "org.opencontainers.image.revision" = ""
-EOF
-        write_ci_base_metadata_labels "${metadata}" >> "${CI_BASE_LABEL_OVERRIDE_PATH}"
-        cat >> "${CI_BASE_LABEL_OVERRIDE_PATH}" <<EOF
-  }
-}
-
-EOF
-    done
-
-    BAKE_FILES+=(-f "${CI_BASE_LABEL_OVERRIDE_PATH}")
-    echo "Appended ci_base metadata label override for targets: ${ci_base_targets[*]}"
-}
-
-uses_rocm_csrc_cache() {
-    case "${TARGET}" in
-        csrc-rocm-ci \
-            | kernel-symbol-map-rocm \
-            | test-rocm-ci \
-            | test-rocm-ci-with-wheel \
-            | test-rocm-ci-with-artifacts \
-            | export-wheel-rocm \
-            | smoke-test-rocm-ci)
-            return 0
-            ;;
-        *)
-            return 1
-            ;;
-    esac
-}
-
-uses_rocm_rust_cache() {
-    case "${TARGET}" in
-        rust-rocm-ci \
-            | test-rocm-ci \
-            | test-rocm-ci-with-wheel \
-            | test-rocm-ci-with-artifacts \
-            | export-wheel-rocm \
-            | smoke-test-rocm-ci)
-            return 0
-            ;;
-        *)
-            return 1
-            ;;
-    esac
-}
-
-compute_rocm_csrc_content_hash() {
-    local bake_dir=""
-    local dockerfile_rocm=""
-    local content_files="${ROCM_CSRC_CONTENT_FILES:-${DEFAULT_ROCM_CSRC_CONTENT_FILES}}"
-    local stages="${ROCM_CSRC_DOCKERFILE_STAGES:-${DEFAULT_ROCM_CSRC_DOCKERFILE_STAGES}}"
-    local -a content_paths=()
-    local -a content_args=()
-
-    bake_dir=$(dirname "${VLLM_BAKE_FILE}")
-    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:-}"
-    )
-
-    {
-        printf 'csrc-input-files-hash:%s\n' "$(compute_content_hash "${content_paths[@]}")"
-        printf 'dockerfile:%s\n' "${dockerfile_rocm}"
-        printf 'resolved-build-args:\n'
-        hash_dockerfile_arg_values "${dockerfile_rocm}" "${content_args[@]}"
-        printf 'dockerfile-stages:%s\n' "${stages}"
-        if content_regular_file "${dockerfile_rocm}"; then
-            hash_dockerfile_stages "${dockerfile_rocm}" "${stages}"
-        else
-            printf 'missing:%s\n' "${dockerfile_rocm}"
-        fi
-    } | sha256sum | cut -d' ' -f1
-}
-
-compute_rocm_csrc_content_hash_if_needed() {
-    local cache_repo="${DOCKERHUB_CACHE_REPO:-rocm/vllm-ci-cache}"
-    local write_scope="${CI_BASE_WRITE_SCOPE:-}"
-
-    if [[ "${ROCM_CSRC_CONTENT_CACHE:-1}" == "0" ]] || ! uses_rocm_csrc_cache; then
-        return 0
-    fi
-
-    ROCM_CSRC_CONTENT_HASH=$(compute_rocm_csrc_content_hash)
-    ROCM_CSRC_TRUSTED_CONTENT_CACHE_REF="${cache_repo}:csrc-rocm-input-${ROCM_CSRC_CONTENT_HASH}"
-    ROCM_CSRC_CONTENT_CACHE_REF="${ROCM_CSRC_TRUSTED_CONTENT_CACHE_REF}"
-    if [[ -n "${write_scope}" ]]; then
-        ROCM_CSRC_CONTENT_CACHE_REF="${ROCM_CSRC_CONTENT_CACHE_REF}-${write_scope}"
-    fi
-    export ROCM_CSRC_CONTENT_HASH
-    export ROCM_CSRC_TRUSTED_CONTENT_CACHE_REF
-    export ROCM_CSRC_CONTENT_CACHE_REF
-    echo "ROCm csrc content cache ref: ${ROCM_CSRC_CONTENT_CACHE_REF}"
-}
-
-compute_rocm_rust_content_hash() {
-    local bake_dir=""
-    local dockerfile_rocm=""
-    local content_files="${ROCM_RUST_CONTENT_FILES:-${DEFAULT_ROCM_RUST_CONTENT_FILES}}"
-    local stages="${ROCM_RUST_DOCKERFILE_STAGES:-${DEFAULT_ROCM_RUST_DOCKERFILE_STAGES}}"
-    local -a content_paths=()
-    local -a content_args=()
-
-    bake_dir=$(dirname "${VLLM_BAKE_FILE}")
-    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_RUST_CONTENT_ARGS:-}"
-    )
-
-    {
-        printf 'rust-input-files-hash:%s\n' "$(compute_content_hash "${content_paths[@]}")"
-        printf 'dockerfile:%s\n' "${dockerfile_rocm}"
-        printf 'resolved-build-args:\n'
-        hash_dockerfile_arg_values "${dockerfile_rocm}" "${content_args[@]}"
-        printf 'dockerfile-stages:%s\n' "${stages}"
-        if content_regular_file "${dockerfile_rocm}"; then
-            hash_dockerfile_stages "${dockerfile_rocm}" "${stages}"
-        else
-            printf 'missing:%s\n' "${dockerfile_rocm}"
-        fi
-    } | sha256sum | cut -d' ' -f1
-}
-
-compute_rocm_rust_content_hash_if_needed() {
-    local cache_repo="${DOCKERHUB_CACHE_REPO:-rocm/vllm-ci-cache}"
-    local write_scope="${CI_BASE_WRITE_SCOPE:-}"
-
-    if [[ "${ROCM_RUST_CONTENT_CACHE:-1}" == "0" ]] || ! uses_rocm_rust_cache; then
-        return 0
-    fi
-
-    ROCM_RUST_CONTENT_HASH=$(compute_rocm_rust_content_hash)
-    ROCM_RUST_TRUSTED_CONTENT_CACHE_REF="${cache_repo}:rust-rocm-input-${ROCM_RUST_CONTENT_HASH}"
-    ROCM_RUST_CONTENT_CACHE_REF="${ROCM_RUST_TRUSTED_CONTENT_CACHE_REF}"
-    if [[ -n "${write_scope}" ]]; then
-        ROCM_RUST_CONTENT_CACHE_REF="${ROCM_RUST_CONTENT_CACHE_REF}-${write_scope}"
-    fi
-    export ROCM_RUST_CONTENT_HASH
-    export ROCM_RUST_TRUSTED_CONTENT_CACHE_REF
-    export ROCM_RUST_CONTENT_CACHE_REF
-    echo "ROCm Rust content cache ref: ${ROCM_RUST_CONTENT_CACHE_REF}"
-}
-
-write_hcl_string_list_entries() {
-    local indent="$1"
-    local value=""
-    shift
-
-    for value in "$@"; do
-        value="${value//\\/\\\\}"
-        value="${value//\"/\\\"}"
-        printf '%s"%s",\n' "${indent}" "${value}"
-    done
-}
-
-hcl_escape_string() {
-    local value="$1"
-
-    value="${value//\\/\\\\}"
-    value="${value//\"/\\\"}"
-    printf '%s' "${value}"
-}
-
-join_words() {
-    local IFS=" "
-    printf '%s' "$*"
-}
-
-metadata_pair() {
-    local key="$1"
-    local value="${2:-}"
-
-    printf '%s\t%s\n' "${key}" "${value}"
-}
-
-write_hcl_string_list() {
-    local indent="$1"
-    shift
-
-    printf '%s[\n' "${indent}"
-    write_hcl_string_list_entries "${indent}  " "$@"
-    printf '%s]\n' "${indent}"
-}
-
-write_rocm_build_arg_override() {
-    local bake_dir=""
-    local dockerfile_rocm=""
-    local -a arg_names=()
-    local arg_name=""
-    local arg_value=""
-
-    bake_dir=$(dirname "${VLLM_BAKE_FILE}")
-    dockerfile_rocm="${CI_BASE_DOCKERFILE:-${bake_dir}/Dockerfile.rocm}"
-    mapfile -t arg_names < <(
-        {
-            get_content_arg_names \
-                "${dockerfile_rocm}" \
-                "${CI_BASE_DOCKERFILE_STAGES:-${DEFAULT_CI_BASE_DOCKERFILE_STAGES}}" \
-                "${CI_BASE_CONTENT_ARGS:-}"
-            get_content_arg_names \
-                "${dockerfile_rocm}" \
-                "${ROCM_CSRC_DOCKERFILE_STAGES:-${DEFAULT_ROCM_CSRC_DOCKERFILE_STAGES}}" \
-                "${ROCM_CSRC_CONTENT_ARGS:-}"
-            get_content_arg_names \
-                "${dockerfile_rocm}" \
-                "${ROCM_RUST_DOCKERFILE_STAGES:-${DEFAULT_ROCM_RUST_DOCKERFILE_STAGES}}" \
-                "${ROCM_RUST_CONTENT_ARGS:-}"
-        } | awk 'NF && !seen[$0]++'
-    )
-
-    {
-        cat <<EOF
-target "_common-rocm" {
-  args = {
-EOF
-        for arg_name in "${arg_names[@]}"; do
-            [[ -n "${arg_name}" ]] || continue
-            arg_value=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "${arg_name}")
-            [[ -n "${arg_value}" ]] || continue
-            printf '    %s = "%s"\n' "${arg_name}" "$(hcl_escape_string "${arg_value}")"
-        done
-        cat <<EOF
-  }
-}
-EOF
-    } > "${ROCM_ARG_OVERRIDE_PATH}"
-
-    BAKE_FILES+=(-f "${ROCM_ARG_OVERRIDE_PATH}")
-    echo "Appended resolved ROCm Docker ARG override"
-}
-
-write_hcl_string_list_attr() {
-    local indent="$1"
-    local attr="$2"
-    shift 2
-
-    printf '%s%s = [\n' "${indent}" "${attr}"
-    write_hcl_string_list_entries "${indent}  " "$@"
-    printf '%s]\n' "${indent}"
-}
-
-validate_cache_export_mode() {
-    local mode="$1"
-    local env_name="$2"
-
-    case "${mode}" in
-        min|max)
-            ;;
-        *)
-            echo "Error: ${env_name} must be one of: min, max"
-            exit 1
-            ;;
-    esac
-}
-
-validate_content_cache_export_mode() {
-    local mode="$1"
-    local env_name="$2"
-
-    case "${mode}" in
-        missing|always|never)
-            ;;
-        *)
-            echo "Error: ${env_name} must be one of: missing, always, never"
-            exit 1
-            ;;
-    esac
-}
-
-should_export_content_cache_ref() {
-    local cache_ref="$1"
-    local cache_name="$2"
-    local trusted_ref="${3:-${cache_ref}}"
-    local mode="${ROCM_CONTENT_CACHE_EXPORT_MODE:-missing}"
-
-    case "${mode}" in
-        always)
-            echo "${cache_name} content cache export mode is always"
-            return 0
-            ;;
-        never)
-            echo "${cache_name} content cache export mode is never"
-            return 1
-            ;;
-        missing|"")
-            if registry_ref_exists_with_retry imagetools "${trusted_ref}"; then
-                echo "${cache_name} trusted content cache is visible: ${trusted_ref}"
-                return 1
-            fi
-            if [[ "${cache_ref}" != "${trusted_ref}" ]] \
-                && registry_ref_exists_with_retry imagetools "${cache_ref}"; then
-                echo "${cache_name} scoped content cache is visible: ${cache_ref}"
-                return 1
-            fi
-            echo "${cache_name} content cache is missing: ${cache_ref}"
-            return 0
-            ;;
-        *)
-            echo "Error: ROCM_CONTENT_CACHE_EXPORT_MODE must be one of: missing, always, never"
-            exit 1
-            ;;
-    esac
-}
-
-write_rocm_cache_override() {
-    local cache_repo="${DOCKERHUB_CACHE_REPO:-rocm/vllm-ci-cache}"
-    local content_cache_export_mode="${ROCM_CONTENT_CACHE_EXPORT_MODE:-missing}"
-    local csrc_cache_to_mode="${ROCM_CSRC_CACHE_TO_MODE:-max}"
-    local rust_cache_to_mode="${ROCM_RUST_CACHE_TO_MODE:-max}"
-    local rocm_cache_to_mode="${ROCM_FINAL_CACHE_TO_MODE:-min}"
-    local -a csrc_content_cache_from=()
-    local -a rust_content_cache_from=()
-    local -a combined_content_cache_from=()
-    local -a csrc_cache_to=()
-    local -a rust_cache_to=()
-    local -a rocm_cache_to=()
-    local -a export_wheel_cache_to=()
-    local export_csrc_cache=1
-    local export_rust_cache=1
-
-    if ! uses_rocm_csrc_cache && ! uses_rocm_rust_cache; then
-        return 0
-    fi
-
-    validate_content_cache_export_mode \
-        "${content_cache_export_mode}" \
-        "ROCM_CONTENT_CACHE_EXPORT_MODE"
-    validate_cache_export_mode "${csrc_cache_to_mode}" "ROCM_CSRC_CACHE_TO_MODE"
-    validate_cache_export_mode "${rust_cache_to_mode}" "ROCM_RUST_CACHE_TO_MODE"
-    validate_cache_export_mode "${rocm_cache_to_mode}" "ROCM_FINAL_CACHE_TO_MODE"
-    echo "ROCm content cache export mode: ${content_cache_export_mode}"
-    echo "ROCm csrc cache export mode: ${csrc_cache_to_mode}"
-    echo "ROCm Rust fallback cache export mode: ${rust_cache_to_mode}"
-    echo "ROCm final image cache export mode: ${rocm_cache_to_mode}"
-
-    if [[ -n "${ROCM_CSRC_CONTENT_CACHE_REF:-}" ]]; then
-        csrc_content_cache_from+=(
-            "type=registry,ref=${ROCM_CSRC_TRUSTED_CONTENT_CACHE_REF}"
-        )
-        if [[ "${ROCM_CSRC_CONTENT_CACHE_REF}" != \
-            "${ROCM_CSRC_TRUSTED_CONTENT_CACHE_REF}" ]]; then
-            csrc_content_cache_from+=(
-                "type=registry,ref=${ROCM_CSRC_CONTENT_CACHE_REF}"
-            )
-        fi
-        if should_export_content_cache_ref \
-            "${ROCM_CSRC_CONTENT_CACHE_REF}" "ROCm csrc" \
-            "${ROCM_CSRC_TRUSTED_CONTENT_CACHE_REF}"; then
-            csrc_cache_to+=(
-                "type=registry,ref=${ROCM_CSRC_CONTENT_CACHE_REF},mode=${csrc_cache_to_mode},ignore-error=true"
-            )
-        else
-            export_csrc_cache=0
-        fi
-    fi
-
-    if [[ -n "${ROCM_RUST_CONTENT_CACHE_REF:-}" ]]; then
-        rust_content_cache_from+=(
-            "type=registry,ref=${ROCM_RUST_TRUSTED_CONTENT_CACHE_REF}"
-        )
-        if [[ "${ROCM_RUST_CONTENT_CACHE_REF}" != \
-            "${ROCM_RUST_TRUSTED_CONTENT_CACHE_REF}" ]]; then
-            rust_content_cache_from+=(
-                "type=registry,ref=${ROCM_RUST_CONTENT_CACHE_REF}"
-            )
-        fi
-        # Legacy commit/branch exports are only needed while the exact-input
-        # content ref is absent. The exact ref itself is refreshed below.
-        if ! should_export_content_cache_ref \
-            "${ROCM_RUST_CONTENT_CACHE_REF}" "ROCm Rust" \
-            "${ROCM_RUST_TRUSTED_CONTENT_CACHE_REF}"; then
-            export_rust_cache=0
-        fi
-        if [[ "${content_cache_export_mode}" != "never" ]]; then
-            # Refresh the exact-input ref in the original solve regardless of
-            # which local or remote cache supplied the Rust result.
-            rust_cache_to+=(
-                "type=registry,ref=${ROCM_RUST_CONTENT_CACHE_REF},mode=min,ignore-error=true"
-            )
-            echo "ROCm Rust exact-input cache will be refreshed (mode=min): ${ROCM_RUST_CONTENT_CACHE_REF}"
-        fi
-    fi
-
-    combined_content_cache_from=("${csrc_content_cache_from[@]}" "${rust_content_cache_from[@]}")
-
-    # Docker Hub cache exports are best-effort. A cache-only target failure can
-    # otherwise cancel the sibling image target before its manifest is pushed.
-    if [[ -n "${BUILDKITE_COMMIT:-}" ]]; then
-        if [[ ${export_csrc_cache} -eq 1 ]]; then
-            csrc_cache_to+=(
-                "type=registry,ref=${cache_repo}:csrc-rocm-${BUILDKITE_COMMIT},mode=${csrc_cache_to_mode},ignore-error=true"
-            )
-        fi
-        if [[ ${export_rust_cache} -eq 1 ]]; then
-            rust_cache_to+=(
-                "type=registry,ref=${cache_repo}:rust-rocm-${BUILDKITE_COMMIT},mode=${rust_cache_to_mode},ignore-error=true"
-            )
-        fi
-        rocm_cache_to+=(
-            "type=registry,ref=${cache_repo}:rocm-${BUILDKITE_COMMIT},mode=${rocm_cache_to_mode},ignore-error=true"
-        )
-    fi
-
-    if [[ -n "${ROCM_CACHE_BRANCH_TAG:-}" ]]; then
-        if [[ ${export_csrc_cache} -eq 1 ]]; then
-            csrc_cache_to+=(
-                "type=registry,ref=${cache_repo}:csrc-rocm-branch-${ROCM_CACHE_BRANCH_TAG},mode=${csrc_cache_to_mode},ignore-error=true"
-            )
-        fi
-        if [[ ${export_rust_cache} -eq 1 ]]; then
-            rust_cache_to+=(
-                "type=registry,ref=${cache_repo}:rust-rocm-branch-${ROCM_CACHE_BRANCH_TAG},mode=${rust_cache_to_mode},ignore-error=true"
-            )
-        fi
-        rocm_cache_to+=(
-            "type=registry,ref=${cache_repo}:rocm-branch-${ROCM_CACHE_BRANCH_TAG},mode=${rocm_cache_to_mode},ignore-error=true"
-        )
-    fi
-
-    # Standalone image/wheel targets reach rust-build but not the cache-only
-    # exporter unless it is requested explicitly.
-    if ((${#rust_cache_to[@]} > 0)); then
-        case "${TARGET}" in
-            test-rocm-ci|export-wheel-rocm|smoke-test-rocm-ci)
-                BAKE_TARGETS=("rust-rocm-ci" "${BAKE_TARGETS[@]}")
-                ;;
-        esac
-    fi
-
-    if [[ "${TARGET}" == "test-rocm-ci-with-wheel" ]]; then
-        export_wheel_cache_to=()
-    else
-        export_wheel_cache_to=("${rocm_cache_to[@]}")
-    fi
-
-    {
-        cat <<EOF
-target "csrc-rocm-ci" {
-  cache-from = concat(
-    get_cache_from_rocm_csrc(),
-EOF
-        write_hcl_string_list "    " "${csrc_content_cache_from[@]}"
-        cat <<EOF
-  )
-EOF
-        write_hcl_string_list_attr "  " "cache-to" "${csrc_cache_to[@]}"
-        cat <<EOF
-}
-
-target "rust-rocm-ci" {
-  cache-from = concat(
-    get_cache_from_rocm_rust(),
-EOF
-        write_hcl_string_list "    " "${rust_content_cache_from[@]}"
-        cat <<EOF
-  )
-EOF
-        write_hcl_string_list_attr "  " "cache-to" "${rust_cache_to[@]}"
-        cat <<EOF
-}
-
-target "kernel-symbol-map-rocm" {
-  cache-from = concat(
-    get_cache_from_rocm_csrc(),
-EOF
-        write_hcl_string_list "    " "${csrc_content_cache_from[@]}"
-        cat <<EOF
-  )
-}
-
-target "test-rocm-ci" {
-  cache-from = concat(
-    get_cache_from_rocm(),
-EOF
-        write_hcl_string_list "    " "${combined_content_cache_from[@]}"
-        cat <<EOF
-  )
-EOF
-        write_hcl_string_list_attr "  " "cache-to" "${rocm_cache_to[@]}"
-        cat <<EOF
-}
-
-target "smoke-test-rocm-ci" {
-  cache-from = concat(
-    get_cache_from_rocm(),
-EOF
-        write_hcl_string_list "    " "${combined_content_cache_from[@]}"
-        cat <<EOF
-  )
-}
-
-target "export-wheel-rocm" {
-  cache-from = concat(
-    get_cache_from_rocm(),
-EOF
-        write_hcl_string_list "    " "${combined_content_cache_from[@]}"
-        cat <<EOF
-  )
-EOF
-        write_hcl_string_list_attr "  " "cache-to" "${export_wheel_cache_to[@]}"
-        cat <<EOF
-}
-EOF
-    } > "${CSRC_CACHE_OVERRIDE_PATH}"
-
-    BAKE_FILES+=(-f "${CSRC_CACHE_OVERRIDE_PATH}")
-    echo "Appended ROCm cache override with non-fatal registry exports"
-}
-
-extract_dependency_pins() {
-    local bake_dir=""
-    local dockerfile_rocm=""
-    local physical_dockerfile=""
-    local var=""
-    local val=""
-
-    bake_dir=$(dirname "${VLLM_BAKE_FILE}")
-    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}"
-    fi
-    if [[ ! -f "${physical_dockerfile}" ]]; then
-        return 0
-    fi
-
-    for var in NIXL_BRANCH UCX_BRANCH DEEPEP_BRANCH; do
-        if [[ -n "${!var:-}" ]]; then
-            echo "Using provided ${var}: ${!var}"
-            continue
-        fi
-
-        val=$(
-            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}"
-        fi
-    done
-}
-
-compute_dependency_cache_keys() {
-    local bake_dir=""
-    local dockerfile_rocm=""
-    local nixl_branch=""
-    local ucx_branch=""
-    local deepep_branch=""
-    local nixl_material=""
-    local deepep_material=""
-
-    bake_dir=$(dirname "${VLLM_BAKE_FILE}")
-    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")
-    deepep_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "DEEPEP_BRANCH")
-
-    if [[ -n "${nixl_branch}" && -n "${ucx_branch}" ]]; then
-        nixl_material=$(compose_stage_cache_material "${dockerfile_rocm}" "base build_nixl")
-        NIXL_CACHE_KEY=$(
-            compose_dependency_cache_key \
-                "${nixl_branch}-ucx-${ucx_branch}" \
-                "${nixl_material}"
-        )
-        export NIXL_CACHE_KEY
-        echo "NIXL dependency cache key: ${NIXL_CACHE_KEY}"
-    fi
-
-    if [[ -n "${deepep_branch}" ]]; then
-        deepep_material=$(compose_stage_cache_material "${dockerfile_rocm}" "base build_deepep")
-        DEEPEP_CACHE_KEY=$(
-            compose_dependency_cache_key \
-                "${deepep_branch}" \
-                "${deepep_material}"
-        )
-        export DEEPEP_CACHE_KEY
-        echo "DeepEP dependency cache key: ${DEEPEP_CACHE_KEY}"
-    fi
-}
-
-compose_stage_cache_material() {
-    local dockerfile="$1"
-    local stages="$2"
-    local -a content_args=()
-
-    mapfile -t content_args < <(get_content_arg_names "${dockerfile}" "${stages}" "")
-    {
-        printf 'dockerfile:%s\n' "${dockerfile}"
-        printf 'dockerfile-stages:%s\n' "${stages}"
-        hash_dockerfile_stages "${dockerfile}" "${stages}"
-        printf 'resolved-build-args:\n'
-        hash_dockerfile_arg_values "${dockerfile}" "${content_args[@]}"
-    }
-}
-
-dependency_cache_ref_exists() {
-    local cache_ref="$1"
-    registry_ref_exists_with_retry imagetools "${cache_ref}"
-}
-
-dependency_cache_ref_for_target() {
-    local target="$1"
-    local cache_repo="${DOCKERHUB_CACHE_REPO:-rocm/vllm-ci-cache}"
-
-    case "${target}" in
-        nixl-rocm-ci)
-            if [[ -n "${NIXL_CACHE_KEY:-}" ]]; then
-                printf '%s\n' "${cache_repo}:nixl-rocm-${NIXL_CACHE_KEY}"
-            elif [[ -n "${NIXL_BRANCH:-}" ]]; then
-                printf '%s\n' "${cache_repo}:nixl-rocm-${NIXL_BRANCH}-ucx-${UCX_BRANCH:-}"
-            fi
-            ;;
-        deepep-rocm-ci)
-            if [[ -n "${DEEPEP_CACHE_KEY:-}" ]]; then
-                printf '%s\n' "${cache_repo}:deepep-rocm-${DEEPEP_CACHE_KEY}"
-            elif [[ -n "${DEEPEP_BRANCH:-}" ]]; then
-                printf '%s\n' "${cache_repo}:deepep-rocm-${DEEPEP_BRANCH}"
-            fi
-            ;;
-    esac
-}
-
-add_dependency_cache_target() {
-    local target="$1"
-
-    if printf '%s\n' "${DEPENDENCY_CACHE_TARGETS[@]}" | grep -qx "${target}"; then
-        return 0
-    fi
-    DEPENDENCY_CACHE_TARGETS+=("${target}")
-}
-
-resolve_ci_base_dependency_targets() {
-    local mode="${ROCM_DEP_CACHE_EXPORT_MODE:-missing}"
-    local nixl_ref=""
-    local deepep_ref=""
-
-    [[ "${TARGET}" == "ci-base-rocm-ci-with-deps" ]] || return 0
-
-    case "${mode}" in
-        always)
-            echo "ROCM_DEP_CACHE_EXPORT_MODE=always; exporting all dependency caches serially"
-            for target in nixl-rocm-ci deepep-rocm-ci; do
-                if [[ -n "$(dependency_cache_ref_for_target "${target}")" ]]; then
-                    add_dependency_cache_target "${target}"
-                fi
-            done
-            ;;
-        never)
-            BAKE_TARGETS=("ci-base-rocm-ci")
-            DEPENDENCY_CACHE_TARGETS=()
-            echo "ROCM_DEP_CACHE_EXPORT_MODE=never; building ci_base without dependency cache exports"
-            return 0
-            ;;
-        missing|"")
-            ;;
-        *)
-            echo "Error: ROCM_DEP_CACHE_EXPORT_MODE must be one of: missing, always, never"
-            exit 1
-            ;;
-    esac
-
-    if [[ "${mode}" != "always" && -n "${NIXL_CACHE_KEY:-}" ]]; then
-        nixl_ref=$(dependency_cache_ref_for_target "nixl-rocm-ci")
-        if dependency_cache_ref_exists "${nixl_ref}"; then
-            echo "NIXL dependency cache exists: ${nixl_ref}"
-        else
-            echo "NIXL dependency cache missing; will seed: ${nixl_ref}"
-            add_dependency_cache_target "nixl-rocm-ci"
-        fi
-    fi
-
-    if [[ "${mode}" != "always" && -n "${DEEPEP_CACHE_KEY:-}" ]]; then
-        deepep_ref=$(dependency_cache_ref_for_target "deepep-rocm-ci")
-        if dependency_cache_ref_exists "${deepep_ref}"; then
-            echo "DeepEP dependency cache exists: ${deepep_ref}"
-        else
-            echo "DeepEP dependency cache missing; will seed: ${deepep_ref}"
-            add_dependency_cache_target "deepep-rocm-ci"
-        fi
-    fi
-
-    BAKE_TARGETS=("ci-base-rocm-ci")
-    if [[ ${#DEPENDENCY_CACHE_TARGETS[@]} -eq 0 ]]; then
-        echo "All dependency caches exist; building ci_base without dependency cache exports"
-    else
-        echo "Resolved dependency cache seed targets: ${DEPENDENCY_CACHE_TARGETS[*]}"
-        echo "Resolved ci_base bake targets: ${BAKE_TARGETS[*]}"
-    fi
-}
-
-bake_config_targets() {
-    printf '%s\n' "${DEPENDENCY_CACHE_TARGETS[@]}" "${BAKE_TARGETS[@]}" \
-        | awk 'NF && !seen[$0]++'
-}
-
-print_bake_config() {
-    local -a print_targets=()
-
-    echo "--- :page_facing_up: Resolved bake configuration"
-    mapfile -t print_targets < <(bake_config_targets)
-    docker buildx bake "${BAKE_ALLOW_ARGS[@]}" \
-        "${BAKE_FILES[@]}" --print "${print_targets[@]}" | tee "${BAKE_CONFIG_FILE}"
-
-    if command -v buildkite-agent >/dev/null 2>&1 && [[ -n "${BUILDKITE_BUILD_NUMBER:-}" ]]; then
-        buildkite-agent artifact upload "${BAKE_CONFIG_FILE}" || true
-        echo "Uploaded ${BAKE_CONFIG_FILE} as Buildkite artifact"
-    else
-        echo "Saved bake config to ${BAKE_CONFIG_FILE} (not in Buildkite, skipping upload)"
-    fi
-}
-
-confirm_remote_image_push() {
-    local image_ref="$1"
-
-    if [[ -z "${CI_BASE_CONTENT_HASH:-}" ]]; then
-        remote_image_exists "${image_ref}"
-        return
-    fi
-
-    if remote_ci_base_identity_is_current_with_retry "${image_ref}"; then
-        return 0
-    fi
-
-    echo "Remote image does not have the expected complete ci_base identity."
-    return 1
-}
-
-verify_dependency_cache_ref() {
-    local cache_ref="$1"
-    local attempts="${ROCM_DEP_CACHE_VERIFY_ATTEMPTS:-6}"
-    local delay_secs="${ROCM_DEP_CACHE_VERIFY_DELAY:-5}"
-    local attempt
-
-    for ((attempt = 1; attempt <= attempts; attempt++)); do
-        if dependency_cache_ref_exists "${cache_ref}"; then
-            echo "Dependency cache confirmed: ${cache_ref}"
-            return 0
-        fi
-        if [[ ${attempt} -lt ${attempts} ]]; then
-            echo "Dependency cache not visible yet (${attempt}/${attempts}): ${cache_ref}"
-            sleep "${delay_secs}"
-        fi
-    done
-
-    echo "ERROR: dependency cache was not confirmed after upload: ${cache_ref}"
-    return 1
-}
-
-seed_dependency_caches_if_needed() {
-    local target=""
-    local cache_ref=""
-
-    if [[ "${TARGET}" != "ci-base-rocm-ci-with-deps" ]]; then
-        return 0
-    fi
-    if [[ ${#DEPENDENCY_CACHE_TARGETS[@]} -eq 0 ]]; then
-        return 0
-    fi
-
-    echo "--- :docker: Seeding ROCm dependency caches"
-    echo "Dependency cache uploads are required for this build."
-    echo "Seeding serially to avoid concurrent Docker Hub cache exporters."
-
-    for target in "${DEPENDENCY_CACHE_TARGETS[@]}"; do
-        cache_ref=$(dependency_cache_ref_for_target "${target}")
-        if [[ -z "${cache_ref}" ]]; then
-            echo "ERROR: could not resolve dependency cache ref for ${target}"
-            return 1
-        fi
-
-        echo "--- :docker: Seeding ${target}"
-        echo "Expected cache ref: ${cache_ref}"
-        docker buildx bake \
-            "${BAKE_ALLOW_ARGS[@]}" \
-            "${BAKE_FILES[@]}" \
-            --progress "${BUILDKIT_PROGRESS:-plain}" \
-            "${target}"
-        verify_dependency_cache_ref "${cache_ref}"
-    done
-}
-
-run_bake() {
-    local confirmation_ref="${IMAGE_TAG:-}"
-
-    if is_ci_base_target && [[ -n "${CI_BASE_IMAGE_TAG_BUILD_REF:-}" ]]; then
-        confirmation_ref="${CI_BASE_IMAGE_TAG_BUILD_REF}"
-    fi
-
-    echo "--- :docker: Building ${TARGET}"
-    docker buildx bake \
-        "${BAKE_ALLOW_ARGS[@]}" \
-        "${BAKE_FILES[@]}" \
-        --progress "${BUILDKIT_PROGRESS:-plain}" \
-        "${BAKE_TARGETS[@]}"
-
-    if is_ci_base_target; then
-        if ! confirm_remote_image_push "${confirmation_ref}"; then
-            echo "Fresh ci_base image failed identity validation: ${confirmation_ref}" >&2
-            return 1
-        fi
-        refresh_ci_base_tags_from_ref "${confirmation_ref}"
-    fi
-    echo "--- :white_check_mark: Build complete"
-}
-
-upload_wheel_artifacts_if_present() {
-    local wheel_dir="./wheel-export"
-    local artifact_dir="artifacts/vllm-rocm-install"
-    local archive_name="vllm-rocm-install.tar.gz"
-    local metadata_dir="${wheel_dir}/.vllm-ci-artifact"
-    local build_base_digest=""
-    local expected_native_base_image=""
-    local native_base_image=""
-    local native_base_digest=""
-    local whl=""
-    local whl_name=""
-    local -a wheels=()
-
-    if ! should_upload_wheel_artifacts; then
-        return 0
-    fi
-
-    if [[ -d "${wheel_dir}" ]]; then
-        mapfile -t wheels < <(find "${wheel_dir}" -maxdepth 1 -type f -name '*.whl' -print)
-    fi
-    if [[ ${#wheels[@]} -ne 1 ]]; then
-        echo "Expected exactly one ROCm wheel in ${wheel_dir}; found ${#wheels[@]}" >&2
-        return 1
-    fi
-    whl="${wheels[0]}"
-    whl_name=$(basename "${whl}")
-    native_base_image="${CI_BASE_IMAGE_TAG_BUILD_REF:-${CI_BASE_IMAGE:-}}"
-    if [[ -z "${native_base_image}" ]]; then
-        echo "Native ROCm artifact requires a ci_base image reference" >&2
-        return 1
-    fi
-    if [[ "${BUILDKITE:-false}" == "true" ]]; then
-        expected_native_base_image=$(ci_base_tag_with_suffix \
-            "${CI_BASE_IMAGE_TAG:-rocm/vllm-dev:ci_base}" \
-            "build-${BUILDKITE_BUILD_ID:-}")
-        if [[ "${native_base_image}" != "${expected_native_base_image}" ]]; then
-            echo "Native ROCm artifact requires the exact ci_base build handoff: ${native_base_image}" >&2
-            return 1
-        fi
-        if [[ ! "${CI_BASE_IMAGE:-}" =~ @sha256:[0-9a-f]{64}$ ]]; then
-            echo "ROCm artifact build base must be digest-pinned: ${CI_BASE_IMAGE:-<empty>}" >&2
-            return 1
-        fi
-        build_base_digest="${CI_BASE_IMAGE##*@}"
-        if ! native_base_digest=$(resolve_image_digest "${native_base_image}"); then
-            echo "Could not resolve native ci_base handoff: ${native_base_image}" >&2
-            return 1
-        fi
-        if [[ "${native_base_digest}" != "${build_base_digest}" ]]; then
-            echo "Native ci_base handoff does not match the artifact build base" >&2
-            echo "  native: ${native_base_image}@${native_base_digest}" >&2
-            echo "  build:  ${CI_BASE_IMAGE}" >&2
-            return 1
-        fi
-    fi
-
-    echo "--- :package: Uploading ROCm vLLM install artifact"
-    rm -rf "${artifact_dir}" "${metadata_dir}"
-    mkdir -p "${artifact_dir}" "${metadata_dir}"
-
-    printf '%s\n' "${BUILDKITE_COMMIT:-local}" > "${metadata_dir}/commit.txt"
-    printf '%s\n' "${native_base_image}" > "${metadata_dir}/native-base-image.txt"
-    printf '%s\n' "${CI_BASE_IMAGE:-}" > "${metadata_dir}/ci-base-image.txt"
-    printf '%s\n' "${IMAGE_TAG:-}" > "${metadata_dir}/fallback-image.txt"
-    printf '%s\n' "${whl_name}" > "${metadata_dir}/wheel-filename.txt"
-
-    tar -C "${wheel_dir}" -czf "${artifact_dir}/${archive_name}" .
-    (
-        cd "${artifact_dir}"
-        sha256sum "${archive_name}" > "${archive_name}.sha256"
-    )
-    echo "Created ${archive_name}: $(du -sh "${artifact_dir}/${archive_name}" | cut -f1)"
-    cp "${metadata_dir}"/*.txt "${artifact_dir}/"
-    cp "${whl}" "${artifact_dir}/${whl_name}"
-    echo "Copied ${whl_name}: $(du -sh "${artifact_dir}/${whl_name}" | cut -f1)"
-
-    if command -v buildkite-agent >/dev/null 2>&1; then
-        buildkite-agent artifact upload "${artifact_dir}/*" || return 1
-        echo "ROCm vLLM install artifacts uploaded to ${artifact_dir}/"
-    elif [[ "${BUILDKITE:-false}" == "true" ]]; then
-        echo "buildkite-agent not found; cannot upload required ROCm artifacts" >&2
-        return 1
-    else
-        echo "Not in Buildkite, skipping artifact upload"
-    fi
-
-    rm -rf "${wheel_dir}"
-}
-
-main() {
-    init_config "$@"
-    configure_ci_base_write_scope
-    print_header
-    validate_inputs
-    load_ci_hcl
-    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
-    maybe_skip_existing_image
-    setup_builder
-    prepare_git_cache_metadata
-    # Non-ci_base builds may deepen a shallow checkout above. Derive archival
-    # 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
-    compute_dependency_cache_keys
-    write_ci_base_label_override
-    compute_rocm_csrc_content_hash_if_needed
-    compute_rocm_rust_content_hash_if_needed
-    write_rocm_cache_override
-    # Keep the context override last so every bake target uses the owned tree.
-    write_build_context_override
-    resolve_ci_base_dependency_targets
-    print_bake_config
-    if [[ "${BAKE_PRINT_ONLY:-0}" == "1" ]]; then
-        echo "BAKE_PRINT_ONLY=1 set; skipping build"
-        return 0
-    fi
-    if should_upload_wheel_artifacts; then
-        # wheel-export is an output directory, not a BuildKit cache. Starting
-        # clean prevents a failed/retried export from packaging a stale wheel.
-        rm -rf ./wheel-export
-    fi
-    if should_export_rocm_smoke; then
-        # The marker is a build output, not a cache. Never accept stale output
-        # from an earlier build or retry.
-        rm -rf ./build/rocm-smoke-export
-    fi
-    if [[ "${TARGET}" == "kernel-symbol-map-rocm" ]]; then
-        rm -rf ./kernel-symbol-map-rocm
-    fi
-    seed_dependency_caches_if_needed
-    run_bake
-    verify_rocm_smoke_export
-    promote_stable_ci_base_tag
-    publish_ci_base_handoff_ref
-    upload_wheel_artifacts_if_present
-    if [[ "${TARGET}" == "kernel-symbol-map-rocm" ]]; then
-        upload_kernel_symbol_map_if_present
-    else
-        export_kernel_symbol_map \
-            || echo "ROCm kernel symbol map export or upload failed; continuing without it" >&2
-    fi
-}
-
-if [[ "${BASH_SOURCE[0]}" == "$0" ]]; then
-    main "$@"
-fi
diff --git a/.buildkite/scripts/hardware_ci/amd_ci_fast_exit.py b/.buildkite/scripts/hardware_ci/amd_ci_fast_exit.py
new file mode 100644
index 000000000000..c77a19fbc966
--- /dev/null
+++ b/.buildkite/scripts/hardware_ci/amd_ci_fast_exit.py
@@ -0,0 +1,66 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""pytest plugin for AMD CI jobs: leave without native teardown once pytest is done.
+
+A pytest process can crash in native library destructors after its summary:
+vllm/ci #96786 printed "236 passed" and then died with SIGSEGV (exit 139) during
+interpreter shutdown, past the point where faulthandler reports. Once the outermost
+session has finished, this runs the Python exit handlers (engine shutdown, Ray,
+multiprocessing cleanup), flushes, and exits with pytest's own status, skipping
+module teardown and native destructors. pytest's result is never changed, and a
+command chain still stops on a real failure.
+
+Off when a kernel recorder is loaded (it writes its trace during native teardown)
+or with VLLM_CI_FAST_EXIT=0.
+"""
+
+import atexit
+import os
+import sys
+
+import pytest
+
+_outer = None
+_session = None
+
+
+def _recorder_loaded() -> bool:
+    return bool(
+        os.environ.get("ROCP_TOOL_LIBRARIES")
+        or os.environ.get("CUDA_INJECTION64_PATH")
+        or "rocprofiler" in os.environ.get("LD_PRELOAD", "")
+    )
+
+
+def pytest_configure(config):
+    global _outer
+    if _outer is None:  # pytester runs nested sessions in-process; act on the outer one
+        _outer = config
+
+
+def pytest_sessionfinish(session, exitstatus):
+    global _session
+    if session.config is _outer:
+        _session = session
+
+
+@pytest.hookimpl(trylast=True)
+def pytest_unconfigure(config):
+    if config is not _outer or _session is None:
+        return
+    if os.environ.get("VLLM_CI_FAST_EXIT") == "0" or _recorder_loaded():
+        return
+    status = int(_session.exitstatus)
+    sys.stdout.flush()
+    sys.stderr.flush()
+    try:
+        atexit._run_exitfuncs()
+    except BaseException as exc:  # an exit handler failing must not mask the result
+        print(f"[amd-ci] exit handler failed: {exc!r}", file=sys.stderr)
+    print(
+        f"[amd-ci] pytest finished with exit {status}; leaving without native teardown",
+        file=sys.stderr,
+        flush=True,
+    )
+    sys.stdout.flush()
+    os._exit(status)
diff --git a/.buildkite/scripts/hardware_ci/run-amd-test.sh b/.buildkite/scripts/hardware_ci/run-amd-test.sh
index 59240b3d0a36..b649cfd7c847 100755
--- a/.buildkite/scripts/hardware_ci/run-amd-test.sh
+++ b/.buildkite/scripts/hardware_ci/run-amd-test.sh
@@ -1,1813 +1,212 @@
 #!/bin/bash
-
-# shellcheck disable=SC2329  # Every function in this file is only ever reached
-# transitively through handle_amd_runner_exit, the EXIT trap handler - shellcheck's
-# unused-function check doesn't do full reachability analysis from an indirectly
-# invoked entry point, so it flags the whole diagnostic-collection call chain as
-# dead code even though it runs live on every nonzero exit.
-
-# This script runs ROCm tests either directly in a native CI pod or inside the
-# corresponding Docker container. Multi-node tests continue to use Docker.
-#
-# Multi-node detection: Instead of matching on fragile group names, we detect
-# multi-node jobs structurally by looking for the bracket command syntax
-# "[node0_cmds] && [node1_cmds]" or via the NUM_NODES environment variable.
-#
-###############################################################################
-# QUOTING / COMMAND PASSING
-#
-# Passing commands as positional arguments ($*) is fragile when the command
-# string itself contains double quotes, e.g.:
-#
-#   bash run-amd-test.sh "export FLAGS="value" && pytest -m "not slow""
-#
-# The outer shell resolves the nested quotes *before* this script runs, so
-# the script receives mangled input it cannot fully recover.
-#
-# Preferred: pass commands via the VLLM_TEST_COMMANDS environment variable:
+# Run one AMD CI job in its Kubernetes pod.
 #
-#   export VLLM_TEST_COMMANDS='export FLAGS="value" && pytest -m "not slow"'
-#   bash run-amd-test.sh
-#
-# Single-quoted assignment preserves all inner double quotes verbatim.
-# The $* path is kept for backward compatibility but callers should migrate.
-###############################################################################
+# The container image ($ROCM_CI_REPO:build-<build-id>) already contains vLLM,
+# all test dependencies, and /vllm-workspace. This script:
+#   1. Sets up per-job local caches and validates the HF mount.
+#   2. Fetches this build's wheel artifact for python_only_compile.sh.
+#   3. Acquires the DPX GPU partition lock when running on a DPX agent.
+#   4. Runs a preflight check (imports + GPU count).
+#   5. Executes $VLLM_TEST_COMMANDS under a wall-clock timeout.
+#   6. Writes a diagnostics bundle on failure.
 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}"
-: "${CLICOLOR_FORCE:=1}"
-: "${PY_COLORS:=1}"
-: "${ROCM_DOCKER_TTY:=1}"
-: "${PYTHONFAULTHANDLER:=1}"
-: "${PYTEST_TIMEOUT:=2400}"
-
-# Preserve runner-only values before clear_ci_orchestration_env removes them
-# from the environment inherited by the test process.
-amd_diagnostics_expected_gpu_count="${VLLM_CI_EXPECTED_GPU_COUNT:-1}"
-amd_diagnostics_dir="${VLLM_CI_DIAGNOSTICS_DIR:-artifacts/amd-gpu-diagnostics}"
-amd_diagnostics_checkout_root="${BUILDKITE_BUILD_CHECKOUT_PATH:-$(pwd -P)}"
-amd_diagnostics_execution_mode="${VLLM_CI_EXECUTION_MODE:-single-node}"
-amd_diagnostics_test_group="${VLLM_TEST_GROUP_NAME:-${BUILDKITE_STEP_KEY:-unknown}}"
-amd_diagnostics_workspace="${VLLM_CI_WORKSPACE:-/vllm-workspace}"
-amd_diagnostics_pod_name="${VLLM_CI_K8S_POD_NAME:-${POD_NAME:-${HOSTNAME:-}}}"
-amd_diagnostics_k8s_namespace="${VLLM_CI_K8S_NAMESPACE:-${POD_NAMESPACE:-unknown}}"
-amd_diagnostics_k8s_node_name="${VLLM_CI_K8S_NODE_NAME:-${NODE_NAME:-}}"
-amd_diagnostics_collected=0
-amd_diagnostics_memory_events_path=""
-amd_diagnostics_probe_budget_seconds=25
-amd_diagnostics_command_timeout_seconds=5
-amd_diagnostics_upload_timeout_seconds=20
-if [[ " ${PYTEST_ADDOPTS:-} " != *" --color"* ]]; then
-  PYTEST_ADDOPTS="${PYTEST_ADDOPTS:+${PYTEST_ADDOPTS} }--color=yes"
-fi
-if [[ " ${PYTEST_ADDOPTS:-} " != *" --durations="* ]]; then
-  PYTEST_ADDOPTS="${PYTEST_ADDOPTS:+${PYTEST_ADDOPTS} }--durations=25"
-fi
-if [[ " ${PYTEST_ADDOPTS:-} " != *" --durations-min="* ]]; then
-  PYTEST_ADDOPTS="${PYTEST_ADDOPTS:+${PYTEST_ADDOPTS} }--durations-min=1.0"
-fi
-# Dump stacks after 25 minutes, then stop an individual test after 40 minutes.
-if [[ " ${PYTEST_ADDOPTS:-} " != *" faulthandler_timeout="* ]]; then
-  PYTEST_ADDOPTS="${PYTEST_ADDOPTS:+${PYTEST_ADDOPTS} }-o faulthandler_timeout=1500"
-fi
-if [[ " ${PYTEST_ADDOPTS:-} " != *" --timeout-method="* &&
-  " ${PYTEST_ADDOPTS:-} " != *" --timeout-method "* ]]; then
-  PYTEST_ADDOPTS="${PYTEST_ADDOPTS:+${PYTEST_ADDOPTS} }--timeout-method=thread"
-fi
-export BUILDKIT_PROGRESS TERM FORCE_COLOR CLICOLOR_FORCE PY_COLORS PYTEST_ADDOPTS PYTEST_TIMEOUT ROCM_DOCKER_TTY
-export PYTHONFAULTHANDLER
-
-# Export Python path for commands that run directly on the host. Containerized
-# tests set this to /vllm-workspace below so spawned Python processes do not
-# depend on their current working directory.
-export PYTHONPATH="${PYTHONPATH:-..}"
-
-ci_started_at=$SECONDS
-
-###############################################################################
-# Helper Functions
-###############################################################################
-
-report_docker_usage() {
-  echo "--- Docker usage"
-  docker system df || true
-}
-
-clear_ci_orchestration_env() {
-  unset -v \
-    VLLM_TEST_GROUP_NAME \
-    VLLM_CI_REQUIRE_PERSISTENT_HF_CACHE \
-    VLLM_CI_ARTIFACT_STEP \
-    VLLM_TEST_CACHE \
-    VLLM_CI_EXECUTION_MODE \
-    VLLM_CI_WORKSPACE \
-    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 \
-    VLLM_CI_DIAGNOSTICS_DIR \
-    VLLM_CI_ARTIFACT_GLOB \
-    VLLM_CI_ARTIFACT_CHECKSUM_GLOB \
-    VLLM_CI_EXPECTED_GPU_COUNT \
-    VLLM_CI_K8S_POD_NAME \
-    VLLM_CI_K8S_NAMESPACE \
-    VLLM_CI_K8S_NODE_NAME \
-    VLLM_CI_USE_ARTIFACTS \
-    VLLM_CI_RESULTS_ROOT \
-    VLLM_ALLOW_DEPRECATED_BEAM_SEARCH
-}
-
-amd_ci_teardown_log() {
-  local event=$1
-  shift
-
-  printf '[amd-ci-teardown] event=%s' "${event}" >&2
-  if (($#)); then
-    printf ' %s' "$@" >&2
-  fi
-  printf '\n' >&2
-}
-
-run_docker_with_ci_timeout() {
-  local raw_timeout=${CONTAINER_TIMEOUT_S:-0}
-  local configured_timeout=0
-  local elapsed=0
-  local remaining_timeout=0
-  local status=0
-
-  if [[ ! "${raw_timeout}" =~ ^(0|[1-9][0-9]{0,5})$ ]]; then
-    amd_ci_teardown_log configuration_error \
-      "reason=invalid_timeout" "expected=integer_0_to_604800"
-    return 2
-  fi
-  configured_timeout=$((10#${raw_timeout}))
-  if ((configured_timeout > 604800)); then
-    amd_ci_teardown_log configuration_error \
-      "reason=invalid_timeout" "expected=integer_0_to_604800"
-    return 2
-  fi
-
-  remaining_timeout=${configured_timeout}
-  if ((configured_timeout > 0)); then
-    if ! command -v timeout >/dev/null 2>&1; then
-      amd_ci_teardown_log configuration_error \
-        "reason=timeout_command_not_found"
-      return 127
-    fi
-    elapsed=$((SECONDS - ci_started_at))
-    if ((elapsed >= configured_timeout)); then
-      amd_ci_teardown_log deadline_exhausted \
-        "mode=docker" "configured_s=${configured_timeout}" "elapsed_s=${elapsed}"
-      return 124
-    fi
-    remaining_timeout=$((configured_timeout - elapsed))
-  fi
-
-  amd_ci_teardown_log workload_started \
-    "mode=docker" "timeout_s=${remaining_timeout}"
-  if ((remaining_timeout == 0)); then
-    "$@" || status=$?
-  else
-    # Docker runs interactively, so keep it in the foreground process group.
-    # --verbose records both TERM and any KILL escalation caused by timeout.
-    timeout --verbose --foreground --signal=TERM --kill-after=10s \
-      "${remaining_timeout}s" "$@" || status=$?
-  fi
-  amd_ci_teardown_log workload_finished "mode=docker" "status=${status}"
-  return "${status}"
-}
-
-prepare_artifact_image() {
-  if [[ "${VLLM_CI_USE_ARTIFACTS:-0}" != "1" ]]; then
-    return 1
-  fi
-  if ! command -v buildkite-agent >/dev/null 2>&1; then
-    echo "buildkite-agent not found; cannot download ROCm wheel artifact"
-    return 1
-  fi
-
-  local artifact_glob="${VLLM_CI_ARTIFACT_GLOB:-artifacts/vllm-rocm-install/vllm-rocm-install.tar.gz}"
-  local archive=""
-  local metadata_file=""
-  local base_image="${VLLM_CI_BASE_IMAGE:-rocm/vllm-dev:ci_base}"
-  local artifact_image=""
-  local artifact_key=""
-  local base_digest=""
-  local wheel_dir=""
-  local context_dir=""
-  local workspace_dir=""
-
-  artifact_work_dir=$(mktemp -d -t vllm-rocm-artifact.XXXXXX)
-  wheel_dir="${artifact_work_dir}/wheels"
-  context_dir="${artifact_work_dir}/context"
-  workspace_dir="${context_dir}/workspace"
-  mkdir -p "${wheel_dir}" "${context_dir}/wheels" "${workspace_dir}"
-
-  echo "--- Downloading ROCm wheel artifact"
-  if ! buildkite-agent artifact download "${artifact_glob}" "${artifact_work_dir}"; then
-    echo "Failed to download ${artifact_glob}"
-    return 1
-  fi
-  buildkite-agent artifact download \
-    "artifacts/vllm-rocm-install/ci-base-image.txt" \
-    "${artifact_work_dir}" >/dev/null 2>&1 || true
-
-  archive=$(find "${artifact_work_dir}" -name "vllm-rocm-install.tar.gz" -type f | head -1)
-  if [[ -z "${archive}" || ! -f "${archive}" ]]; then
-    echo "ROCm wheel artifact archive was not found"
-    return 1
-  fi
-
-  metadata_file=$(find "${artifact_work_dir}" -name "ci-base-image.txt" -type f | head -1)
-  if [[ -n "${metadata_file}" && -s "${metadata_file}" ]]; then
-    base_image=$(tr -d '[:space:]' < "${metadata_file}")
-  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"
-  echo "Base image: ${base_image}"
-  docker pull "${base_image}" || return 1
-  base_digest=$(
-    docker image inspect \
-      --format='{{if .RepoDigests}}{{index .RepoDigests 0}}{{else}}{{.Id}}{{end}}' \
-      "${base_image}" 2>/dev/null || printf '%s' "${base_image}"
-  )
-
-  artifact_key=$(
-    {
-      printf 'base-image:%s\n' "${base_digest}"
-      sha256sum "${archive}"
-    } | sha256sum | cut -c1-24
-  )
-  artifact_image="rocm/vllm-ci-artifact:${artifact_key}"
-
-  if docker image inspect "${artifact_image}" >/dev/null 2>&1; then
-    echo "Using existing local ROCm artifact image: ${artifact_image}"
-    image_name="${artifact_image}"
-    return 0
-  fi
-
-  tar -xzf "${archive}" -C "${wheel_dir}" || return 1
-  if ! ls "${wheel_dir}"/*.whl >/dev/null 2>&1; then
-    echo "ROCm wheel artifact did not contain a wheel"
-    return 1
-  fi
-  if [[ ! -d "${wheel_dir}/tests" ]]; then
-    echo "ROCm wheel artifact did not contain the test workspace"
-    return 1
-  fi
-
-  cp "${wheel_dir}"/*.whl "${context_dir}/wheels/" || return 1
-  tar -C "${wheel_dir}" \
-    --exclude='*.whl' \
-    --exclude='.vllm-ci-artifact' \
-    --exclude='./.vllm-ci-artifact' \
-    -cf - . \
-    | tar -C "${workspace_dir}" -xf - || return 1
-  cat > "${context_dir}/Dockerfile" <<'EOF'
-ARG BASE_IMAGE
-FROM ${BASE_IMAGE}
-COPY wheels/ /tmp/vllm-wheels/
-COPY workspace/ /vllm-workspace/
-RUN python3 -m pip install --no-deps --force-reinstall /tmp/vllm-wheels/*.whl \
-    && rm -rf /tmp/vllm-wheels
-WORKDIR /vllm-workspace
-EOF
-
-  echo "--- Building local ROCm test image"
-  docker build \
-    --pull=false \
-    --progress "${BUILDKIT_PROGRESS}" \
-    --build-arg "BASE_IMAGE=${base_image}" \
-    -t "${artifact_image}" \
-    "${context_dir}" || return 1
-  image_name="${artifact_image}"
-  return 0
-}
-
-is_native_runtime() {
-  [[ "${AMD_CI_RUNTIME:-}" == "native" || "${NATIVE_CI:-}" == "true" ]]
-}
-
-validate_native_workspace() {
-  local workspace_dir="${VLLM_CI_WORKSPACE:-/vllm-workspace}"
-  local workspace_real=""
-  local checkout_real=""
-  local workspace_mount=""
-
-  mkdir -p "${workspace_dir}" || return 1
-  workspace_real=$(readlink -m "${workspace_dir}") || return 1
-  if [[ -n "${BUILDKITE_BUILD_CHECKOUT_PATH:-}" ]]; then
-    checkout_real=$(readlink -m "${BUILDKITE_BUILD_CHECKOUT_PATH}") || return 1
-    if [[ "${checkout_real}" == "${workspace_real}" \
-      || "${checkout_real}" == "${workspace_real}/"* \
-      || "${workspace_real}" == "${checkout_real}/"* ]]; then
-      echo "Refusing to replace ${workspace_real}; it overlaps the Buildkite checkout ${checkout_real}" >&2
-      return 1
-    fi
-  fi
-  if [[ "${VLLM_CI_REQUIRE_WORKSPACE_MOUNT:-1}" == "1" ]]; then
-    if ! command -v findmnt >/dev/null 2>&1; then
-      echo "findmnt is required to verify the native workspace mount" >&2
-      return 1
-    fi
-    workspace_mount=$(findmnt -n -T "${workspace_real}" -o TARGET 2>/dev/null || true)
-    if [[ "$(readlink -m "${workspace_mount:-/}")" != "${workspace_real}" ]]; then
-      echo "Native CI requires a dedicated volume mounted at ${workspace_real}" >&2
-      return 1
-    fi
-  fi
-}
-
-prepare_native_workspace() {
-  local test_commands="${1:-}"
-
-  if [[ "${VLLM_CI_USE_ARTIFACTS:-0}" != "1" ]]; then
-    echo "Native CI requires VLLM_CI_USE_ARTIFACTS=1"
-    return 1
-  fi
-  if ! command -v buildkite-agent >/dev/null 2>&1; then
-    echo "buildkite-agent not found; cannot download ROCm wheel artifact"
-    return 1
-  fi
-  validate_native_workspace || return 1
-
-  local artifact_glob="${VLLM_CI_ARTIFACT_GLOB:-artifacts/vllm-rocm-install/vllm-rocm-install.tar.gz}"
-  local artifact_checksum_glob="${VLLM_CI_ARTIFACT_CHECKSUM_GLOB:-${artifact_glob}.sha256}"
-  local artifact_step="${VLLM_CI_ARTIFACT_STEP:-image-build-amd}"
-  local archive=""
-  local checksum=""
-  local download_dir=""
-  local metadata_dir=""
-  local recorded_base=""
-  local recorded_commit=""
-  local recorded_wheel=""
-  local checkout=""
-  local checkout_commit=""
-  local workspace_dir="${VLLM_CI_WORKSPACE:-/vllm-workspace}"
-  local wheel_dir=""
-  local attempt=0
-  local attempt_dir=""
-  local -a archives=()
-  local -a checksums=()
-  local -a wheels=()
-
-  artifact_work_dir=$(mktemp -d -t vllm-rocm-artifact.XXXXXX) || return 1
-  wheel_dir="${artifact_work_dir}/wheels"
-  mkdir -p "${wheel_dir}" || return 1
-
-  echo "--- Downloading ROCm wheel artifact from ${artifact_step} (native in-pod)"
-  for attempt in 1 2 3; do
-    attempt_dir="${artifact_work_dir}/download-${attempt}"
-    rm -rf "${attempt_dir}" || return 1
-    mkdir -p "${attempt_dir}" || return 1
-    if buildkite-agent artifact download \
-      "${artifact_glob}" "${attempt_dir}" --step "${artifact_step}" \
-      && buildkite-agent artifact download \
-        "${artifact_checksum_glob}" "${attempt_dir}" --step "${artifact_step}"; then
-      download_dir="${attempt_dir}"
-      break
-    fi
-    echo "Artifact download attempt ${attempt}/3 failed"
-    if [[ "${attempt}" -lt 3 ]]; then
-      sleep $((attempt * 2))
-    fi
-  done
-  if [[ -z "${download_dir}" ]]; then
-    echo "Failed to download ${artifact_glob} and ${artifact_checksum_glob} from ${artifact_step}"
-    return 1
-  fi
-
-  mapfile -t archives < <(
-    find "${download_dir}" -name "vllm-rocm-install.tar.gz" -type f -print
-  )
-  mapfile -t checksums < <(
-    find "${download_dir}" -name "vllm-rocm-install.tar.gz.sha256" -type f -print
-  )
-  if [[ ${#archives[@]} -ne 1 || ${#checksums[@]} -ne 1 ]]; then
-    echo "Expected exactly one ROCm archive and checksum; found ${#archives[@]} archive(s) and ${#checksums[@]} checksum(s)" >&2
-    return 1
-  fi
-  archive="${archives[0]}"
-  checksum="${checksums[0]}"
-  if [[ "$(dirname "${archive}")" != "$(dirname "${checksum}")" ]]; then
-    echo "ROCm archive and checksum were downloaded to different directories" >&2
-    return 1
-  fi
-  (
-    cd "$(dirname "${archive}")"
-    sha256sum -c "$(basename "${checksum}")"
-  ) || return 1
-
-  tar --no-same-owner -xzf "${archive}" -C "${wheel_dir}" || return 1
-  mapfile -t wheels < <(
-    find "${wheel_dir}" -maxdepth 1 -type f -name '*.whl' -print
-  )
-  if [[ ${#wheels[@]} -ne 1 ]]; then
-    echo "ROCm artifact must contain exactly one top-level wheel; found ${#wheels[@]}" >&2
-    return 1
-  fi
-  metadata_dir="${wheel_dir}/.vllm-ci-artifact"
-  for metadata_file in commit.txt native-base-image.txt wheel-filename.txt; do
-    if [[ ! -s "${metadata_dir}/${metadata_file}" ]]; then
-      echo "ROCm artifact metadata is missing ${metadata_file}" >&2
-      return 1
-    fi
-  done
-  for metadata_file in ci-base-image.txt fallback-image.txt; do
-    if [[ ! -f "${metadata_dir}/${metadata_file}" ]]; then
-      echo "ROCm artifact metadata is missing ${metadata_file}" >&2
-      return 1
-    fi
-  done
-
-  recorded_commit=$(tr -d '\r\n' < "${metadata_dir}/commit.txt")
-  recorded_base=$(tr -d '\r\n' < "${metadata_dir}/native-base-image.txt")
-  recorded_wheel=$(tr -d '\r\n' < "${metadata_dir}/wheel-filename.txt")
-  if [[ -z "${BUILDKITE_COMMIT:-}" || "${recorded_commit}" != "${BUILDKITE_COMMIT}" ]]; then
-    echo "ROCm artifact commit ${recorded_commit} does not match ${BUILDKITE_COMMIT:-unset}" >&2
-    return 1
-  fi
-  if [[ -z "${BUILDKITE_BUILD_ID:-}" \
-    || "${recorded_base}" != *":ci_base-build-${BUILDKITE_BUILD_ID}" ]]; then
-    echo "ROCm artifact base is not scoped to Buildkite build ${BUILDKITE_BUILD_ID:-unset}: ${recorded_base}" >&2
-    return 1
-  fi
-  if [[ -z "${VLLM_CI_BASE_IMAGE:-}" || "${recorded_base}" != "${VLLM_CI_BASE_IMAGE}" ]]; then
-    echo "ROCm artifact base ${recorded_base} does not match ${VLLM_CI_BASE_IMAGE:-unset}" >&2
-    return 1
-  fi
-  if [[ "${recorded_wheel}" != "$(basename "${wheels[0]}")" ]]; then
-    echo "ROCm artifact wheel manifest ${recorded_wheel} does not match $(basename "${wheels[0]}")" >&2
-    return 1
-  fi
-  for required_dir in tests .buildkite requirements; do
-    if [[ ! -d "${wheel_dir}/${required_dir}" ]]; then
-      echo "ROCm wheel artifact did not contain ${required_dir}/" >&2
-      return 1
-    fi
+# ---------------------------------------------------------------------------
+# Inputs — set by the pipeline / ci-infra template
+# ---------------------------------------------------------------------------
+commands=${VLLM_TEST_COMMANDS:?VLLM_TEST_COMMANDS must be set by the pipeline}
+gpus=${VLLM_CI_EXPECTED_GPU_COUNT:-1}
+[[ $gpus =~ ^(0|[1-9][0-9]*)$ ]] || { echo "Invalid VLLM_CI_EXPECTED_GPU_COUNT=$gpus" >&2; exit 1; }
+job=${BUILDKITE_JOB_ID:-local}
+job=${job//[^A-Za-z0-9_.-]/_}
+# Short: Ray puts AF_UNIX sockets under $TMPDIR, whose paths may not exceed 107 bytes.
+root=/tmp/vllm-${job##*-}
+root=${root:0:22}
+
+# ---------------------------------------------------------------------------
+# Cache directories — all under $root so they are job-local and isolated.
+# HF hub is the one exception: it lives on the node's persistent NFS mount.
+# ---------------------------------------------------------------------------
+export HF_HOME=${HF_HOME:-/home/buildkite-agent/huggingface}
+export TIKTOKEN_RS_CACHE_DIR=$HF_HOME/tiktoken-rs-cache
+export PYTHONPATH=/vllm-workspace
+export PYTORCH_ROCM_ARCH=""                        # let vLLM detect arches at runtime
+export TMPDIR=$root/tmp
+export VLLM_RPC_BASE_PATH=/tmp
+export TORCHINDUCTOR_CACHE_DIR=$root/inductor
+export TRITON_CACHE_DIR=$root/triton
+export VLLM_CACHE_ROOT=$root/vllm
+export XDG_CACHE_HOME=$root/xdg
+export HF_DATASETS_CACHE=$root/datasets
+export HF_HUB_DOWNLOAD_TIMEOUT=${HF_HUB_DOWNLOAD_TIMEOUT:-300}
+export HF_HUB_ETAG_TIMEOUT=${HF_HUB_ETAG_TIMEOUT:-60}
+
+mkdir -p \
+  "$TMPDIR" "$TORCHINDUCTOR_CACHE_DIR" "$TRITON_CACHE_DIR" \
+  "$VLLM_CACHE_ROOT" "$XDG_CACHE_HOME" "$HF_DATASETS_CACHE" \
+  "$TIKTOKEN_RS_CACHE_DIR" || exit 1
+
+# HF_HOME must be a persistent mount, not rootfs (catches misconfigured nodes).
+[[ $(findmnt -n -T "$HF_HOME" -o TARGET) != / ]] \
+  || { echo "HF_HOME is on rootfs — it must be a persistent node mount" >&2; exit 1; }
+
+# hf-xet state corrupts on NFS; redirect it to a local tmpdir when needed.
+if [[ $(findmnt -n -T "$HF_HOME" -o FSTYPE) == nfs* ]]; then
+  export HF_XET_CACHE=$root/hf-xet
+  export HF_XET_HIGH_PERFORMANCE=0
+  export HF_XET_RECONSTRUCTION_USE_VECTORED_WRITE=0
+fi
+
+# ---------------------------------------------------------------------------
+# Terminal and pytest settings
+# ---------------------------------------------------------------------------
+# Keep CI log output colorized.
+: "${TERM:=xterm-256color}" "${FORCE_COLOR:=1}" "${CLICOLOR_FORCE:=1}" "${PY_COLORS:=1}"
+export TERM FORCE_COLOR CLICOLOR_FORCE PY_COLORS
+
+# Fault-handler dumps stacks after 25 min; pytest-timeout kills a single test
+# after 40 min (PYTEST_TIMEOUT=2400 s). Both are overridable by the pipeline.
+export PYTHONFAULTHANDLER=1
+export PYTEST_TIMEOUT=${PYTEST_TIMEOUT:-2400}
+export PYTEST_ADDOPTS="${PYTEST_ADDOPTS:-} \
+  --color=yes \
+  --durations=25 \
+  --durations-min=1.0 \
+  -o faulthandler_timeout=1500 \
+  --timeout-method=thread"
+
+# After its summary, a pytest process leaves without native library teardown, where it
+# has crashed with SIGSEGV after every test passed (amd_ci_fast_exit.py).
+# Installed into site-packages, not put on PYTHONPATH: commands that set PYTHONPATH themselves
+# still inherit PYTEST_ADDOPTS and must find the plugin.
+ci_dir=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
+cp "$ci_dir/amd_ci_fast_exit.py" \
+  "$(python3 -P -c 'import sysconfig; print(sysconfig.get_paths()["purelib"])')/" || exit 1
+export PYTEST_ADDOPTS="$PYTEST_ADDOPTS -p amd_ci_fast_exit"
+
+[[ $gpus == 0 ]] && export VLLM_TARGET_DEVICE=cpu
+
+# ---------------------------------------------------------------------------
+# Wheel reinstall — python_only_compile.sh path
+# ---------------------------------------------------------------------------
+# python_only_compile.sh reinstalls vLLM around this build's wheel, which build.sh uploads as an
+# artifact of the image build step (kept out of the image: only this test needs it).
+if [[ $commands == *python_only_compile.sh* ]]; then
+  wheel_dir=$(mktemp -d)
+  for try in 1 2 3; do
+    buildkite-agent artifact download ".rocm-wheel/vllm-*.whl" "$wheel_dir" --step image-build-amd \
+      && break
+    ((try == 3)) && exit 1
+    sleep $((try * 5))
   done
-
-  echo "--- Installing ROCm wheel into pod environment"
-  python3 -m pip install --no-deps --force-reinstall "${wheels[0]}" || return 1
-
-  echo "--- Preparing ${workspace_dir} from artifact"
-  find "${workspace_dir}" -mindepth 1 -maxdepth 1 -exec rm -rf -- {} + || return 1
-  tar -C "${wheel_dir}" \
-    --exclude='*.whl' \
-    --exclude='.vllm-ci-artifact' \
-    --exclude='./.vllm-ci-artifact' \
-    -cf - . | tar --no-same-owner -C "${workspace_dir}" -xf - || return 1
-  if [[ ! -d "${workspace_dir}/tests" ]]; then
-    echo "Failed to stage the native test workspace" >&2
-    return 1
-  fi
-
-  # The ROCm artifact intentionally contains only the installed wheel and the
-  # test workspace. The Python-only compilation job also needs setup.py and the
-  # vllm source tree, so overlay the matching Buildkite checkout for that job.
-  if [[ "${test_commands}" == *python_only_compile.sh* ]]; then
-    checkout="${BUILDKITE_BUILD_CHECKOUT_PATH:-}"
-    if [[ -z "${checkout}" || ! -d "${checkout}" ]]; then
-      echo "Python-only native CI requires BUILDKITE_BUILD_CHECKOUT_PATH" >&2
-      return 1
-    fi
-    if ! git -c "safe.directory=${checkout}" -C "${checkout}" \
-      rev-parse --is-inside-work-tree >/dev/null 2>&1; then
-      echo "Buildkite checkout is not a Git worktree: ${checkout}" >&2
-      return 1
-    fi
-    checkout_commit=$(
-      git -c "safe.directory=${checkout}" -C "${checkout}" rev-parse HEAD
-    ) || return 1
-    if [[ "${checkout_commit}" != "${recorded_commit}" ]]; then
-      echo "Buildkite checkout ${checkout_commit} does not match ROCm artifact ${recorded_commit}" >&2
-      return 1
-    fi
-
-    # setup.py normally derives this from .git via setuptools-scm. The native
-    # source overlay deliberately excludes Git metadata, so preserve the exact
-    # version from the already installed, artifact-matched wheel.
-    VLLM_VERSION_OVERRIDE=$(
-      python3 -c 'import importlib.metadata as m; print(m.version("vllm"))'
-    ) || return 1
-    export VLLM_VERSION_OVERRIDE
-    VLLM_PRECOMPILED_WHEEL_LOCATION="${wheels[0]}"
-    export VLLM_PRECOMPILED_WHEEL_LOCATION
-    echo "INFO: native Python-only wheel=${VLLM_PRECOMPILED_WHEEL_LOCATION}"
-
-    echo "--- Overlaying full source checkout for Python-only compilation"
-    # Archive the verified commit instead of copying the worktree so dirty or
-    # untracked agent files cannot contaminate the artifact-matched workspace.
-    git -c "safe.directory=${checkout}" -C "${checkout}" \
-      archive --format=tar "${recorded_commit}" \
-      | tar --no-same-owner -C "${workspace_dir}" -xf - || return 1
-    for required_source in setup.py pyproject.toml vllm; do
-      if [[ ! -e "${workspace_dir}/${required_source}" ]]; then
-        echo "Full source checkout is missing ${required_source}" >&2
-        return 1
-      fi
-    done
-  fi
-
-  return 0
-}
-
-initialize_native_environment() {
-  local job_id="${BUILDKITE_JOB_ID:-${BUILDKITE_PARALLEL_JOB:-local}}"
-  local job_id_suffix=""
-  local native_root=""
-  local hf_fstype=""
-  local hf_mount=""
-
-  if [[ "$(id -u)" -ne 0 ]]; then
-    echo "Native ROCm CI currently requires the ci_base container to run as root" >&2
-    return 1
-  fi
-
-  job_id="${job_id//[^A-Za-z0-9_.-]/_}"
-  job_id_suffix="${job_id##*-}"
-  job_id_suffix="${job_id_suffix:0:12}"
-  native_root="/tmp/vllm-native-${job_id}"
-  TMPDIR="/tmp/vllm-${job_id_suffix}/tmp"
-  VLLM_RPC_BASE_PATH="/tmp"
-  TORCHINDUCTOR_CACHE_DIR="${native_root}/cache/torchinductor"
-  TRITON_CACHE_DIR="${native_root}/cache/triton"
-  VLLM_CACHE_ROOT="${native_root}/cache/vllm"
-  XDG_CACHE_HOME="${native_root}/cache/xdg"
-  : "${HF_HOME:=/home/buildkite-agent/huggingface}"
-  # datasets uses POSIX locks that are unsupported by the shared HF NFS cache.
-  # Keep processed datasets job-local while retaining the persistent Hub cache.
-  HF_DATASETS_CACHE="${native_root}/cache/huggingface/datasets"
-  # openai-harmony downloads its tiktoken vocab on first use and caches it under
-  # $TMPDIR by default, which is job-local; keep it with the persistent Hub cache.
-  TIKTOKEN_RS_CACHE_DIR="${HF_HOME}/tiktoken-rs-cache"
-  : "${HF_HUB_DOWNLOAD_TIMEOUT:=300}"
-  : "${HF_HUB_ETAG_TIMEOUT:=60}"
-  if [[ "${VLLM_CI_EXPECTED_GPU_COUNT:-1}" == "0" ]]; then
-    # CPU-only native jobs intentionally reuse the ROCm wheel. Make that target
-    # explicit so platform selection does not depend on wheel metadata.
-    VLLM_TARGET_DEVICE=cpu
-    export VLLM_TARGET_DEVICE
-  fi
-  export TMPDIR VLLM_RPC_BASE_PATH
-  export TORCHINDUCTOR_CACHE_DIR TRITON_CACHE_DIR VLLM_CACHE_ROOT XDG_CACHE_HOME
-  export HF_HOME HF_DATASETS_CACHE HF_HUB_DOWNLOAD_TIMEOUT HF_HUB_ETAG_TIMEOUT
-  export TIKTOKEN_RS_CACHE_DIR
-  export PYTORCH_ROCM_ARCH=""
-
-  mkdir -p "${TMPDIR}" \
-    "${TORCHINDUCTOR_CACHE_DIR}" \
-    "${TRITON_CACHE_DIR}" \
-    "${VLLM_CACHE_ROOT}" \
-    "${XDG_CACHE_HOME}" \
-    "${HF_HOME}" \
-    "${TIKTOKEN_RS_CACHE_DIR}" \
-    "${HF_DATASETS_CACHE}" || return 1
-
-  echo "Native compile caches: VLLM_CACHE_ROOT=${VLLM_CACHE_ROOT} TORCHINDUCTOR_CACHE_DIR=${TORCHINDUCTOR_CACHE_DIR}"
-
-  if [[ "${VLLM_CI_REQUIRE_PERSISTENT_HF_CACHE:-0}" == "1" ]]; then
-    if ! command -v findmnt >/dev/null 2>&1; then
-      echo "findmnt is required to verify the native Hugging Face cache mount" >&2
-      return 1
-    fi
-    hf_mount=$(findmnt -n -T "${HF_HOME}" -o TARGET 2>/dev/null || true)
-    if [[ -z "${hf_mount}" || "${hf_mount}" == "/" ]]; then
-      echo "Native CI requires a persistent volume mounted at or above ${HF_HOME}" >&2
-      return 1
-    fi
-  fi
-
-  if command -v findmnt >/dev/null 2>&1; then
-    hf_fstype=$(findmnt -n -T "${HF_HOME}" -o FSTYPE 2>/dev/null || true)
-  fi
-  if [[ "${hf_fstype}" == nfs || "${hf_fstype}" == nfs4 ]]; then
-    # Keep hf-xet state local and avoid vectored writes on shared NFS.
-    export HF_XET_CACHE="${native_root}/cache/hf-xet"
-    export HF_XET_HIGH_PERFORMANCE=0
-    export HF_XET_RECONSTRUCTION_USE_VECTORED_WRITE=0
-    mkdir -p "${HF_XET_CACHE}" || return 1
-    echo "Configured hf-xet for shared ${hf_fstype} cache at ${HF_HOME}"
-  fi
-}
-
-check_dpx_gpu_exclusivity() {
-  local devices=(/dev/dri/renderD*)
-  local device_id lock_status
-
-  if [[ "${#devices[@]}" -ne 1 || ! -c "${devices[0]}" ]]; then
-    echo "DPX guard requires exactly one render device; refusing to start tests." >&2
-    return 1
-  fi
-  # This queue mounts the same host-local HF cache into every pod on the node.
-  if ! mountpoint -q "${HF_HOME}"; then
-    echo "DPX guard requires the shared HF cache mount; refusing to start tests." >&2
-    return 1
-  fi
-  device_id=$(stat -Lc '%t-%T' "${devices[0]}") || return 1
-  mkdir -p "${HF_HOME}/.vllm-dpx-locks" || return 1
-  # Hold the descriptor through workload and teardown; never delete the file.
-  exec {dpx_gpu_lock_fd}>>"${HF_HOME}/.vllm-dpx-locks/${device_id}.lock" || return 1
-  flock -n -E 75 "${dpx_gpu_lock_fd}" && return 0
-  lock_status=$?
-  if [[ "${lock_status}" -eq 75 ]]; then
-    echo "DPX GPU collision: ${devices[0]} (${device_id}) is already locked by another CI job; refusing to start tests." >&2
-  else
-    echo "DPX GPU lock failed (status ${lock_status}); refusing to start tests." >&2
-  fi
-  return 1
+  wheels=("$wheel_dir"/.rocm-wheel/vllm-*.whl)
+  [[ ${#wheels[@]} == 1 && -f ${wheels[0]} ]] \
+    || { echo "Expected exactly one wheel artifact, found: ${wheels[*]}" >&2; exit 1; }
+  VLLM_PRECOMPILED_WHEEL_LOCATION=${wheels[0]}
+  VLLM_VERSION_OVERRIDE=$(python3 -P -c \
+    'import importlib.metadata as m; print(m.version("vllm"))') || exit 1
+  export VLLM_TARGET_DEVICE=rocm VLLM_VERSION_OVERRIDE VLLM_PRECOMPILED_WHEEL_LOCATION
+fi
+
+# ---------------------------------------------------------------------------
+# DPX GPU partition lock
+# ---------------------------------------------------------------------------
+# DPX nodes expose exactly one GPU partition per pod. A host-wide lock file
+# prevents a second job from landing on the same device simultaneously.
+_acquire_dpx_lock() {
+  local dev=(/dev/dri/renderD*)
+  [[ ${#dev[@]} == 1 ]] \
+    || { echo "DPX: expected exactly one render device, found ${#dev[@]}" >&2; return 1; }
+  mkdir -p "$HF_HOME/.vllm-dpx-locks"
+  exec 9>>"$HF_HOME/.vllm-dpx-locks/$(stat -Lc %t-%T "${dev[0]}").lock"
+  flock -n 9 || { echo "DPX GPU collision on ${dev[0]}" >&2; return 1; }
+}
+
+if [[ ${BUILDKITE_AGENT_META_DATA_QUEUE:-} == *dpx* && $gpus != 0 ]]; then
+  _acquire_dpx_lock || exit 1
+fi
+
+# ---------------------------------------------------------------------------
+# Lifecycle logging
+# ---------------------------------------------------------------------------
+# Structured events for infra dashboards keyed on [amd-ci-teardown].
+# Usage: _teardown_log <event> [key=val ...]
+_teardown_log() {
+  local event=$1; shift
+  printf '[amd-ci-teardown] event=%s' "$event" >&2
+  (($#)) && printf ' %s' "$@" >&2
+  printf '\n' >&2
 }
 
-run_native_preflight() {
-  local expected_gpus="${VLLM_CI_EXPECTED_GPU_COUNT:-1}"
-
-  if [[ ! "${expected_gpus}" =~ ^[0-9]+$ ]]; then
-    echo "Invalid VLLM_CI_EXPECTED_GPU_COUNT=${expected_gpus}" >&2
-    return 1
-  fi
-
-  python3 -c "import encodings, importlib.metadata as im, importlib.util as iu; [im.version(d) for d in ('transformers', 'torch', 'ray', 'sympy', 'markupsafe', 'vllm')]; missing=[m for m in ('torch.utils.model_zoo', 'transformers.models.nomic_bert', 'ray.dag', 'sympy.physics', 'markupsafe._speedups') if iu.find_spec(m) is None]; assert not missing, missing" || return 1
-
-  if [[ "${expected_gpus}" == "0" ]]; then
-    echo "Native CPU-only AMD job: skipping ROCm device validation"
-    return 0
-  fi
-
+# ---------------------------------------------------------------------------
+# Diagnostics bundle (always collected on failure)
+# ---------------------------------------------------------------------------
+# Writes to artifacts/amd-gpu-diagnostics/<job>/diagnostics.log, which the
+# pipeline's artifact_paths uploads even when the step fails.
+diagnostics() {
+  local d=artifacts/amd-gpu-diagnostics/$job
+  mkdir -p "$d"
+  {
+    date
+    echo "node=${VLLM_CI_K8S_NODE_NAME:-?}  pod=${VLLM_CI_K8S_POD_NAME:-?}"
+    echo "--- amd-smi ---"
+    timeout 10 amd-smi list
+    echo "--- rocm-smi pids ---"
+    timeout 10 rocm-smi --showpids
+    echo "--- top processes by RSS ---"
+    ps -eo pid,etime,rss,args --sort=-rss | head -30
+    echo "--- disk / memory ---"
+    df -h /tmp /dev/shm "$HF_HOME"
+    free -g
+    echo "--- devices ---"
+    ls -l /dev/kfd /dev/dri
+  } >"$d/diagnostics.log" 2>&1
+}
+
+# ---------------------------------------------------------------------------
+# Preflight: verify imports and GPU visibility
+# ---------------------------------------------------------------------------
+# -P: the checkout (cwd) has a vllm/ source tree that would shadow the installed package.
+python3 -P -c "import vllm, torch, transformers, ray" || exit 1
+
+if [[ $gpus != 0 ]]; then
   echo "--- ROCm info"
-  rocminfo || return 1
-  VLLM_CI_EXPECTED_GPU_COUNT="${expected_gpus}" python3 - <<'PY'
-import os
-
+  rocminfo || { diagnostics; exit 1; }
+  python3 - <<EOF || { diagnostics; exit 1; }
 import torch
-
-expected = int(os.environ["VLLM_CI_EXPECTED_GPU_COUNT"])
-assert torch.version.hip, "PyTorch is not a ROCm build"
-assert torch.cuda.is_available(), "ROCm GPU is not available to PyTorch"
-actual = torch.cuda.device_count()
-assert actual == expected, f"Expected {expected} ROCm GPU(s), found {actual}"
-PY
-}
-
-is_multi_node() {
-  local cmds="$1"
-  # Primary signal: NUM_NODES environment variable set by the pipeline
-  if [[ "${NUM_NODES:-1}" -gt 1 ]]; then
-    return 0
-  fi
-  # Fallback: detect the bracket syntax structurally
-  # Pattern: [...] && [...] (per-node command arrays)
-  if [[ "$cmds" == *'] && ['* ]]; then
-    return 0
-  fi
-  return 1
-}
-
-append_failure_diagnostic_section() {
-  local log_file=$1
-  local title=$2
-
-  printf '\n===============================================================================\n' \
-    >> "${log_file}"
-  printf '%s\n' "${title}" >> "${log_file}"
-  printf '===============================================================================\n' \
-    >> "${log_file}"
-}
-
-append_failure_diagnostic_note() {
-  local log_file=$1
-  shift
-  printf '%s\n' "$@" >> "${log_file}"
-}
-
-run_failure_diagnostic() {
-  local log_file=$1
-  local probe_deadline=$2
-  shift 2
-  local command_status=0
-  local command_timeout_seconds=${amd_diagnostics_command_timeout_seconds}
-  local remaining_seconds=$((probe_deadline - SECONDS))
-
-  {
-    printf '\n$'
-    printf ' %q' "$@"
-    printf '\n'
-  } >> "${log_file}"
-
-  if ! command -v timeout >/dev/null 2>&1; then
-    echo "[timeout unavailable; bounded diagnostic command skipped]" \
-      >> "${log_file}"
-    return 0
-  fi
-  if [[ "${remaining_seconds}" -le 0 ]]; then
-    echo "[diagnostic probe budget exhausted; command skipped]" \
-      >> "${log_file}"
-    return 0
-  fi
-  if [[ "${remaining_seconds}" -lt "${command_timeout_seconds}" ]]; then
-    command_timeout_seconds=${remaining_seconds}
-  fi
-
-  timeout --kill-after=1s "${command_timeout_seconds}s" "$@" \
-    >> "${log_file}" 2>&1
-  command_status=$?
-  if [[ "${command_status}" -ne 0 ]]; then
-    echo "[diagnostic command exited with status ${command_status}]" \
-      >> "${log_file}"
-  fi
-  return 0
-}
-
-upload_failure_diagnostics_artifact() {
-  local upload_root=$1 upload_path=$2 agent_bin=""
-  local agent_upload_help=""
-  local -a upload_options=()
-
-  if [[ -n "${BUILDKITE_BIN_PATH:-}" \
-    && -x "${BUILDKITE_BIN_PATH}/buildkite-agent" ]]; then
-    agent_bin="${BUILDKITE_BIN_PATH}/buildkite-agent"
-  else
-    agent_bin=$(command -v buildkite-agent 2>/dev/null || true)
-  fi
-  if [[ -z "${agent_bin}" && -x /workspace/buildkite-agent ]]; then
-    agent_bin=/workspace/buildkite-agent
-  fi
-  [[ -n "${agent_bin}" && -n "${BUILDKITE_JOB_ID:-}" ]] || return 1
-  command -v timeout >/dev/null 2>&1 || return 1
-
-  # Older AMD runners predate the literal upload options. A restricted path
-  # has no glob or delimiter characters, so it is safe with either agent.
-  if [[ ! "${upload_path}" =~ ^[A-Za-z0-9_][A-Za-z0-9_.-]*(/[A-Za-z0-9_][A-Za-z0-9_.-]*)*$ ]]; then
-    agent_upload_help=$("${agent_bin}" artifact upload --help 2>&1 || true)
-    if [[ "${agent_upload_help}" != *"--literal"* \
-      || "${agent_upload_help}" != *"--delimiter"* ]]; then
-      return 1
-    fi
-    upload_options=(--literal --delimiter "")
-  fi
-
-  (
-    cd "${upload_root}" || exit 1
-    BUILDKITE_AGENT_DEBUG=false \
-      BUILDKITE_AGENT_DEBUG_HTTP=false \
-      BUILDKITE_AGENT_TRACE_HTTP=false \
-      BUILDKITE_AGENT_LOG_LEVEL=error \
-      timeout --kill-after=2s "${amd_diagnostics_upload_timeout_seconds}s" \
-      "${agent_bin}" artifact upload \
-        "${upload_options[@]}" "./${upload_path}"
-  ) >/dev/null 2>&1
-}
-
-append_failure_diagnostic_file() {
-  local log_file=$1
-  local label=$2
-  local path=$3
-  local command_status=0
-
-  printf '\n%s:\n' "${label}" >> "${log_file}"
-  printf '$ cat %q\n' "${path}" >> "${log_file}"
-  if [[ ! -r "${path}" ]]; then
-    echo "[file unavailable]" >> "${log_file}"
-    return 0
-  fi
-
-  cat "${path}" >> "${log_file}" 2>&1
-  command_status=$?
-  if [[ "${command_status}" -ne 0 ]]; then
-    echo "[diagnostic file read exited with status ${command_status}]" \
-      >> "${log_file}"
-  fi
-  return 0
-}
-
-decode_mountinfo_path() {
-  local value=$1
-
-  value=${value//\\040/ }
-  value=${value//\\011/$'\t'}
-  value=${value//\\012/$'\n'}
-  value=${value//\\134/\\}
-  printf '%s\n' "${value}"
-}
-
-resolve_current_cgroup_v2_dir() {
-  local cgroup_path=""
-  local mount_root=""
-  local mount_point=""
-  local relative_path=""
-  local candidate=""
-
-  cgroup_path=$(awk -F: '$2 == "" { print $3; exit }' \
-    /proc/self/cgroup 2>/dev/null)
-  [[ -n "${cgroup_path}" ]] || return 1
-
-  while IFS=$'\t' read -r mount_root mount_point; do
-    mount_root=$(decode_mountinfo_path "${mount_root}")
-    mount_point=$(decode_mountinfo_path "${mount_point}")
-    if [[ "${cgroup_path}" == "/" ]]; then
-      relative_path=""
-    elif [[ "${mount_root}" == "/" ]]; then
-      relative_path="${cgroup_path}"
-    elif [[ "${cgroup_path}" == "${mount_root}" ]]; then
-      relative_path=""
-    elif [[ "${cgroup_path}" == "${mount_root}/"* ]]; then
-      relative_path="/${cgroup_path#"${mount_root}/"}"
-    else
-      continue
-    fi
-
-    candidate="${mount_point%/}${relative_path}"
-    if [[ -d "${candidate}" ]]; then
-      printf '%s\n' "${candidate}"
-      return 0
-    fi
-  done < <(
-    awk '
-      {
-        for (i = 6; i <= NF; i++) {
-          if ($i == "-" && $(i + 1) == "cgroup2") {
-            print $4 "\t" $5
-            break
-          }
-        }
-      }
-    ' /proc/self/mountinfo 2>/dev/null
-  )
-  return 1
-}
-
-collect_cgroup_files() {
-  local log_file=$1
-  local title=$2
-  local cgroup_dir=$3
-  shift 3
-  local cgroup_file=""
-  local files_collected=0
-
-  append_failure_diagnostic_section "${log_file}" "${title}"
-  if [[ -z "${cgroup_dir}" || ! -d "${cgroup_dir}" ]]; then
-    append_failure_diagnostic_note "${log_file}" \
-      "resolved_cgroup_path=unavailable"
-    return 0
-  fi
-
-  append_failure_diagnostic_note "${log_file}" \
-    "resolved_cgroup_path=${cgroup_dir}"
-  for cgroup_file in "$@"; do
-    if [[ -r "${cgroup_dir}/${cgroup_file}" ]]; then
-      append_failure_diagnostic_file "${log_file}" "${cgroup_file}" \
-        "${cgroup_dir}/${cgroup_file}"
-      files_collected=$((files_collected + 1))
-    fi
-  done
-  if [[ "${files_collected}" -eq 0 ]]; then
-    append_failure_diagnostic_note "${log_file}" \
-      "[no requested cgroup files were readable]"
-  fi
-}
-
-collect_cgroup_diagnostics() {
-  local log_file=$1
-  local cgroup_dir=""
-
-  append_failure_diagnostic_section "${log_file}" "Cgroup membership"
-  append_failure_diagnostic_file "${log_file}" "current_process" \
-    /proc/self/cgroup
-  append_failure_diagnostic_file "${log_file}" "pid_namespace_init_process" \
-    /proc/1/cgroup
-
-  cgroup_dir=$(resolve_current_cgroup_v2_dir 2>/dev/null || true)
-  collect_cgroup_files "${log_file}" \
-    "Cgroup v2 current-process resource state" "${cgroup_dir}" \
-    cgroup.type cgroup.events \
-    memory.current memory.peak memory.max memory.high memory.low memory.min \
-    memory.oom.group memory.swap.current memory.swap.max \
-    memory.events memory.events.local memory.swap.events memory.stat \
-    memory.pressure \
-    cpu.max cpu.max.burst cpu.weight cpu.stat cpu.pressure \
-    cpuset.cpus cpuset.cpus.effective cpuset.mems cpuset.mems.effective \
-    pids.current pids.max pids.events pids.events.local \
-    io.stat io.max io.weight io.pressure
-  if [[ -n "${cgroup_dir}" && -r "${cgroup_dir}/memory.events" ]]; then
-    amd_diagnostics_memory_events_path="${cgroup_dir}/memory.events"
-  fi
-}
-
-collect_process_diagnostics() {
-  local log_file=$1
-  local probe_deadline=$2
-  local shell_proc_dir="/proc/$$"
-
-  append_failure_diagnostic_section "${log_file}" \
-    "Runner process constraints and PID-namespace-visible processes"
-  append_failure_diagnostic_file "${log_file}" "diagnostic_shell_limits" \
-    "${shell_proc_dir}/limits"
-  if command -v awk >/dev/null 2>&1; then
-    # shellcheck disable=SC2016  # The expression is evaluated by awk.
-    run_failure_diagnostic "${log_file}" "${probe_deadline}" awk '
-      $1 ~ /^(Pid:|PPid:|Threads:|VmPeak:|VmSize:|VmHWM:|VmRSS:|RssAnon:|RssFile:|RssShmem:|Cpus_allowed_list:|Mems_allowed_list:|voluntary_ctxt_switches:|nonvoluntary_ctxt_switches:)$/ {
-        print
-      }
-    ' "${shell_proc_dir}/status"
-  fi
-  if command -v nproc >/dev/null 2>&1; then
-    run_failure_diagnostic "${log_file}" "${probe_deadline}" nproc
-  fi
-  run_failure_diagnostic "${log_file}" "${probe_deadline}" \
-    /bin/bash -c 'ulimit -a'
-  if command -v ps >/dev/null 2>&1 && command -v head >/dev/null 2>&1; then
-    append_failure_diagnostic_note "${log_file}" \
-      "process_snapshot=top_100_by_rss"
-    run_failure_diagnostic "${log_file}" "${probe_deadline}" \
-      /bin/bash -c \
-      'ps -eo pid,ppid,stat,etimes,nlwp,pcpu,rss,vsz,comm --sort=-rss | head -n 101'
-  fi
-  return 0
-}
-
-collect_mount_diagnostic() {
-  local log_file=$1
-  local probe_deadline=$2
-  local label=$3
-  local path=$4
-
-  printf '\n%s (%s):\n' "${label}" "${path}" >> "${log_file}"
-  if [[ ! -e "${path}" ]]; then
-    echo "[path unavailable]" >> "${log_file}"
-    return 0
-  fi
-  if command -v findmnt >/dev/null 2>&1; then
-    run_failure_diagnostic "${log_file}" "${probe_deadline}" findmnt \
-      -n -T "${path}" -o TARGET,FSTYPE,VFS-OPTIONS
-  fi
-  if command -v df >/dev/null 2>&1; then
-    run_failure_diagnostic "${log_file}" "${probe_deadline}" df -h \
-      --output=fstype,size,used,avail,pcent,itotal,iused,iavail,ipcent,target \
-      -- "${path}"
-  fi
-  return 0
-}
-
-collect_execution_resource_diagnostics() {
-  local log_file=$1
-  local probe_deadline=$2
-
-  # free and /proc/meminfo generally describe the node from a pod. Cgroups are
-  # the live source for this container's resource limits, usage, and events.
-  collect_cgroup_diagnostics "${log_file}"
-  collect_process_diagnostics "${log_file}" "${probe_deadline}"
-
-  append_failure_diagnostic_section "${log_file}" \
-    "Mount-visible filesystem capacity"
-  append_failure_diagnostic_note "${log_file}" \
-    "These values are filesystem views, not Kubernetes ephemeral-storage quotas."
-  collect_mount_diagnostic "${log_file}" "${probe_deadline}" \
-    "checkout" "${amd_diagnostics_checkout_root}"
-  if is_native_runtime; then
-    collect_mount_diagnostic "${log_file}" "${probe_deadline}" \
-      "native workspace emptyDir" "${amd_diagnostics_workspace}"
-  else
-    collect_mount_diagnostic "${log_file}" "${probe_deadline}" \
-      "test workspace (if mounted)" "${amd_diagnostics_workspace}"
-  fi
-  collect_mount_diagnostic "${log_file}" "${probe_deadline}" \
-    "temporary directory" "${TMPDIR:-/tmp}"
-  if [[ "${TMPDIR:-/tmp}" != "/tmp" ]]; then
-    collect_mount_diagnostic "${log_file}" "${probe_deadline}" \
-      "system temporary directory" /tmp
-  fi
-  if is_native_runtime; then
-    collect_mount_diagnostic "${log_file}" "${probe_deadline}" \
-      "shared-memory emptyDir" /dev/shm
-  else
-    collect_mount_diagnostic "${log_file}" "${probe_deadline}" \
-      "shared-memory mount" /dev/shm
-  fi
-  if [[ -n "${HF_HOME:-}" ]]; then
-    collect_mount_diagnostic "${log_file}" "${probe_deadline}" \
-      "Hugging Face cache" "${HF_HOME}"
-  fi
-  return 0
-}
-
-collect_amd_device_nodes() {
-  local log_file=$1
-  local probe_deadline=$2
-  local device=""
-  local device_count=0
-
-  for device in /dev/kfd /dev/dri/renderD*; do
-    if [[ ! -e "${device}" ]]; then
-      continue
-    fi
-    device_count=$((device_count + 1))
-    if command -v stat >/dev/null 2>&1; then
-      run_failure_diagnostic "${log_file}" "${probe_deadline}" stat -Lc \
-        '%n type=%F mode=%a owner=%u:%g device=%t:%T' -- "${device}"
-    else
-      append_failure_diagnostic_note "${log_file}" "device=${device}"
-    fi
-  done
-  if [[ "${device_count}" -eq 0 ]]; then
-    append_failure_diagnostic_note "${log_file}" \
-      "No /dev/kfd or /dev/dri/renderD* device nodes are visible."
-  fi
-  return 0
-}
-
-collect_rocm_failure_diagnostics() {
-  local exit_code=$1
-  local job_id="${BUILDKITE_JOB_ID:-local}"
-  local retry_count="${BUILDKITE_RETRY_COUNT:-0}"
-  local parallel_job="${BUILDKITE_PARALLEL_JOB:-0}"
-  local diagnostics_relative_path=""
-  local diagnostics_path=""
-  local diagnostics_parent=""
-  local diagnostics_fallback_root=""
-  local diagnostics_storage="checkout"
-  local diagnostics_artifact_summary=""
-  local checkout_real=""
-  local diagnostics_parent_real=""
-  local exit_signal=""
-  local probe_deadline=0
-  local runtime="single-node-docker"
-  local diagnostics_scope="outer-runner-after-test-container-exit"
-  local identity_label="runner"
-  local identity_value="${HOSTNAME:-unknown}"
-  local k8s_pod="unknown"
-  local k8s_namespace="${amd_diagnostics_k8s_namespace}"
-  local oom_kill_count=""
-  local summary_identity_label="Runner"
-  local -a summary_rows=()
-
-  if [[ "${amd_diagnostics_collected}" == "1" ]]; then
-    return 0
-  fi
-  amd_diagnostics_collected=1
-
-  if is_native_runtime; then
-    runtime="native-kubernetes"
-    diagnostics_scope="current-container-cgroup-and-namespaces"
-    if [[ -n "${amd_diagnostics_pod_name}" ]]; then
-      identity_label="pod"
-      identity_value="${amd_diagnostics_pod_name}"
-      k8s_pod="${amd_diagnostics_pod_name}"
-    else
-      identity_label="container_hostname"
-    fi
-  elif [[ "${amd_diagnostics_execution_mode}" == "multi-node" ]]; then
-    runtime="multi-node-docker"
-  fi
-  if [[ "${k8s_namespace}" == "unknown" \
-    && -r /var/run/secrets/kubernetes.io/serviceaccount/namespace ]]; then
-    k8s_namespace=$(
-      tr -d '\r\n' \
-        < /var/run/secrets/kubernetes.io/serviceaccount/namespace 2>/dev/null
-    )
-  fi
-
-  job_id="${job_id//[^A-Za-z0-9_.-]/_}"
-  retry_count="${retry_count//[^A-Za-z0-9_.-]/_}"
-  parallel_job="${parallel_job//[^A-Za-z0-9_.-]/_}"
-
-  if [[ "${exit_code}" -gt 128 && "${exit_code}" -le 192 ]]; then
-    exit_signal=$(kill -l "$((exit_code - 128))" 2>/dev/null || true)
-  fi
-
-  # Buildkite keeps the supplied artifact path. Restrict the configurable
-  # directory to a checkout-relative path so the UI shows a clean artifact key.
-  if [[ -z "${amd_diagnostics_dir}" \
-    || "${amd_diagnostics_dir}" == /* \
-    || "${amd_diagnostics_dir}" == *".."* ]]; then
-    echo "WARNING: ignoring unsafe VLLM_CI_DIAGNOSTICS_DIR"
-    amd_diagnostics_dir="artifacts/amd-gpu-diagnostics"
-  fi
-  diagnostics_relative_path="${amd_diagnostics_dir}/${job_id}/diagnostics.log"
-  diagnostics_path="${amd_diagnostics_checkout_root}/${diagnostics_relative_path}"
-  diagnostics_parent=$(dirname "${diagnostics_path}")
-  checkout_real=$(readlink -m "${amd_diagnostics_checkout_root}" 2>/dev/null || true)
-  diagnostics_parent_real=$(readlink -m "${diagnostics_parent}" 2>/dev/null || true)
-  if [[ -z "${checkout_real}" || -z "${diagnostics_parent_real}" \
-    || ("${diagnostics_parent_real}" != "${checkout_real}" \
-      && "${diagnostics_parent_real}" != "${checkout_real}/"*) ]] \
-    || ! mkdir -p "${diagnostics_parent}" \
-    || [[ -L "${diagnostics_path}" ]] \
-    || ! (set -o noclobber; : > "${diagnostics_path}") 2>/dev/null; then
-    echo "WARNING: unable to create AMD CI diagnostics in the checkout; using temporary storage."
-    diagnostics_path=""
-    if diagnostics_fallback_root=$(mktemp -d -t \
-      vllm-amd-diagnostics.XXXXXX 2>/dev/null); then
-      diagnostics_path="${diagnostics_fallback_root}/${diagnostics_relative_path}"
-      diagnostics_parent=$(dirname "${diagnostics_path}")
-      if (umask 077; mkdir -p "${diagnostics_parent}" \
-        && (set -o noclobber; : > "${diagnostics_path}")) 2>/dev/null; then
-        diagnostics_storage="temporary-fallback"
-      else
-        rm -rf -- "${diagnostics_fallback_root}" || true
-        diagnostics_fallback_root=""
-        diagnostics_path=""
-      fi
-    fi
-  fi
-  if [[ -z "${diagnostics_path}" ]]; then
-    echo "WARNING: unable to create AMD CI diagnostics at ${diagnostics_relative_path}."
-    printf '\nAMD CI failure summary\n'
-    printf '%-22s | %s\n' \
-      "Field" "Value" \
-      "----------------------" "-----" \
-      "Exit code" "${exit_code}" \
-      "Signal" "${exit_signal:-none}" \
-      "Runtime" "${runtime}" \
-      "Pod/container" "${identity_value}" \
-      "Diagnostics artifact" "unavailable: ${diagnostics_relative_path}"
-    return 0
-  fi
-
-  {
-    echo "AMD CI failure diagnostics"
-    echo "timestamp_utc=$(date -u +%Y-%m-%dT%H:%M:%SZ)"
-    echo "exit_code=${exit_code}"
-    echo "exit_signal=${exit_signal:-none}"
-    echo "build_url=${BUILDKITE_BUILD_URL:-unknown}"
-    echo "job_id=${BUILDKITE_JOB_ID:-unknown}"
-    echo "step_key=${BUILDKITE_STEP_KEY:-unknown}"
-    echo "step_label=${BUILDKITE_LABEL:-unknown}"
-    echo "test_group=${amd_diagnostics_test_group}"
-    echo "retry_count=${retry_count}"
-    echo "parallel_job=${parallel_job}"
-    echo "execution_mode=${amd_diagnostics_execution_mode}"
-    echo "runtime=${runtime}"
-    echo "diagnostics_scope=${diagnostics_scope}"
-    echo "expected_gpu_count=${amd_diagnostics_expected_gpu_count}"
-    echo "probe_budget_seconds=${amd_diagnostics_probe_budget_seconds}"
-    echo "command_timeout_seconds=${amd_diagnostics_command_timeout_seconds}"
-    echo "diagnostics_storage=${diagnostics_storage}"
-    echo "agent_name=${BUILDKITE_AGENT_NAME:-unknown}"
-    echo "container_hostname=${HOSTNAME:-unknown}"
-    echo "k8s_pod=${k8s_pod}"
-    echo "k8s_namespace=${k8s_namespace:-unknown}"
-    echo "k8s_node=${amd_diagnostics_k8s_node_name:-unknown}"
-    echo "kernel=$(uname -srmo 2>/dev/null || echo unknown)"
-  } > "${diagnostics_path}"
-
-  probe_deadline=$((SECONDS + amd_diagnostics_probe_budget_seconds))
-
-  # Cgroup, process, and mount state is fast and most useful for explaining OOM,
-  # throttling, PID exhaustion, and emptyDir pressure. Capture it before external
-  # probes can consume the shared deadline.
-  collect_execution_resource_diagnostics \
-    "${diagnostics_path}" "${probe_deadline}"
-
-  append_failure_diagnostic_section "${diagnostics_path}" \
-    "AMD GPU diagnostics (tool-visible devices)"
-  if [[ "${amd_diagnostics_expected_gpu_count}" == "0" ]]; then
-    append_failure_diagnostic_note "${diagnostics_path}" \
-      "CPU-only job; AMD GPU probes skipped."
-  else
-    if [[ "${runtime}" == "native-kubernetes" ]]; then
-      append_failure_diagnostic_note "${diagnostics_path}" \
-        "SMI visibility may be broader than the Kubernetes device allocation."
-    else
-      append_failure_diagnostic_note "${diagnostics_path}" \
-        "Outer-runner SMI visibility may be broader than the test container."
-    fi
-    collect_amd_device_nodes "${diagnostics_path}" "${probe_deadline}"
-  fi
-  if [[ "${amd_diagnostics_expected_gpu_count}" != "0" ]] \
-    && command -v amd-smi >/dev/null 2>&1; then
-    run_failure_diagnostic "${diagnostics_path}" "${probe_deadline}" \
-      amd-smi version
-    # Bus data identifies a card within the public node without publishing its
-    # persistent UUID, serial number, or process list.
-    run_failure_diagnostic "${diagnostics_path}" "${probe_deadline}" \
-      amd-smi static --bus --gpu all
-    run_failure_diagnostic "${diagnostics_path}" "${probe_deadline}" \
-      amd-smi metric --ecc --ecc-blocks --pcie --gpu all
-    run_failure_diagnostic "${diagnostics_path}" "${probe_deadline}" \
-      amd-smi static --ras --gpu all
-    run_failure_diagnostic "${diagnostics_path}" "${probe_deadline}" \
-      amd-smi metric --power --temperature --usage --mem-usage --gpu all
-    run_failure_diagnostic "${diagnostics_path}" "${probe_deadline}" \
-      amd-smi xgmi --link-status --gpu all
-  elif [[ "${amd_diagnostics_expected_gpu_count}" != "0" ]] \
-    && command -v rocm-smi >/dev/null 2>&1; then
-    run_failure_diagnostic "${diagnostics_path}" "${probe_deadline}" rocm-smi \
-      --showbus --showreplaycount --showrasinfo --showpagesinfo
-  elif [[ "${amd_diagnostics_expected_gpu_count}" != "0" ]]; then
-    append_failure_diagnostic_note "${diagnostics_path}" \
-      "Neither amd-smi nor rocm-smi is available."
-  fi
-
-  if [[ "${runtime}" == "native-kubernetes" ]]; then
-    append_failure_diagnostic_section "${diagnostics_path}" \
-      "Unauthenticated pod network reachability"
-  else
-    append_failure_diagnostic_section "${diagnostics_path}" \
-      "Unauthenticated outer-runner network reachability"
-  fi
-  if command -v curl >/dev/null 2>&1; then
-    run_failure_diagnostic "${diagnostics_path}" "${probe_deadline}" curl -q \
-      --silent --location --proto '=https' --proto-redir '=https' \
-      --max-redirs 3 --output /dev/null --connect-timeout 2 --max-time 4 \
-      --write-out 'target=huggingface http_code=%{http_code} dns_done_s=%{time_namelookup} connect_done_s=%{time_connect} tls_done_s=%{time_appconnect} first_byte_s=%{time_starttransfer} total_s=%{time_total}\n' \
-      https://huggingface.co/api/models/gpt2
-    run_failure_diagnostic "${diagnostics_path}" "${probe_deadline}" curl -q \
-      --silent --location --proto '=https' --proto-redir '=https' \
-      --max-redirs 3 --output /dev/null --connect-timeout 2 --max-time 4 \
-      --write-out 'target=github_git_smart_http http_code=%{http_code} dns_done_s=%{time_namelookup} connect_done_s=%{time_connect} tls_done_s=%{time_appconnect} first_byte_s=%{time_starttransfer} total_s=%{time_total}\n' \
-      'https://github.com/vllm-project/ci-infra.git/info/refs?service=git-upload-pack'
-  else
-    append_failure_diagnostic_note "${diagnostics_path}" \
-      "curl unavailable; reachability probes skipped."
-  fi
-
-  if [[ "${runtime}" == "native-kubernetes" \
-    && -n "${amd_diagnostics_memory_events_path}" \
-    && -r "${amd_diagnostics_memory_events_path}" ]]; then
-    oom_kill_count=$(
-      awk '$1 == "oom_kill" { print $2; exit }' \
-        "${amd_diagnostics_memory_events_path}" 2>/dev/null
-    )
-    if [[ ! "${oom_kill_count}" =~ ^[0-9]+$ ]]; then
-      oom_kill_count=""
-    fi
-  fi
-
-  case "${identity_label}" in
-    pod)
-      summary_identity_label="Kubernetes pod"
-      ;;
-    container_hostname)
-      summary_identity_label="Container hostname"
-      ;;
-  esac
-
-  summary_rows=(
-    "Field" "Value"
-    "----------------------" "-----"
-    "Exit code" "${exit_code}"
-    "Signal" "${exit_signal:-none}"
-    "Runtime" "${runtime}"
-    "${summary_identity_label}" "${identity_value}"
-  )
-  if [[ -n "${amd_diagnostics_k8s_node_name}" ]]; then
-    summary_rows+=("Kubernetes node" "${amd_diagnostics_k8s_node_name}")
-  fi
-  if [[ -n "${oom_kill_count}" ]]; then
-    summary_rows+=("Cgroup OOM kills" "${oom_kill_count}")
-  fi
-  diagnostics_artifact_summary="${diagnostics_relative_path}"
-  if [[ -n "${diagnostics_fallback_root}" ]]; then
-    if upload_failure_diagnostics_artifact \
-      "${diagnostics_fallback_root}" "${diagnostics_relative_path}"; then
-      echo "Uploaded AMD CI diagnostics artifact from temporary storage: ${diagnostics_relative_path}"
-      diagnostics_artifact_summary="${diagnostics_relative_path} (direct upload)"
-    else
-      echo "WARNING: failed to upload temporary AMD CI diagnostics artifact: ${diagnostics_relative_path}"
-      echo "--- AMD CI diagnostics (artifact upload failed)"
-      cat "${diagnostics_path}" || true
-      diagnostics_artifact_summary="job log only: direct upload failed"
-    fi
-    rm -rf -- "${diagnostics_fallback_root}" || \
-      echo "WARNING: unable to remove temporary AMD CI diagnostics storage."
-  fi
-  summary_rows+=("Diagnostics artifact" "${diagnostics_artifact_summary}")
-
-  printf '\nAMD CI failure summary\n'
-  printf '%-22s | %s\n' "${summary_rows[@]}"
-
-  return 0
-}
-
-handle_pytest_exit() {
-  local exit_code=$1
-  if [ "$exit_code" -eq 5 ]; then
-    echo "Pytest exit code 5 (no tests collected) - treating as success."
-    exit 0
-  fi
-  exit "$exit_code"
-}
-
-###############################################################################
-# Pytest marker/keyword re-quoting
-#
-# When commands are passed through Buildkite -> shell -> $* -> bash -c,
-# quotes around multi-word pytest -m/-k expressions get stripped:
-#   pytest -v -s -m 'not cpu_test' v1/core
-# becomes:
-#   pytest -v -s -m not cpu_test v1/core
-#
-# pytest then interprets "cpu_test" as a file path, not part of the marker.
-#
-# This function detects unquoted expressions after -m/-k and re-quotes them
-# by collecting tokens until a recognizable boundary is reached:
-#   - test path (contains '/')
-#   - test file (ends with '.py')
-#   - another pytest flag (--xxx or -x single-char flags)
-#   - command separator (&& || ; |)
-#   - environment variable assignment (FOO=bar)
-#
-# Single-word markers (e.g. -m cpu_test, -m hybrid_model) pass through
-# unquoted since they have no spaces and work fine.
-#
-# Already-quoted expressions (containing literal single quotes) are passed
-# through untouched to avoid double-quoting well-formed shell fragments.
-#
-# NOTE: This ONLY fixes -m/-k flags. It cannot recover arbitrary inner
-# double-quotes stripped by the calling shell (see header comment).
-# Use VLLM_TEST_COMMANDS to avoid the problem entirely.
-###############################################################################
-re_quote_pytest_markers() {
-  local input="$1"
-  local output=""
-  local collecting=false
-  local marker_buf=""
-
-  # Strip backslash-newline continuations, then flatten remaining newlines
-  local flat="${input//$'\\\n'/ }"
-  flat="${flat//$'\n'/ }"
-
-  # Disable globbing to prevent *.py etc. from expanding during read -ra
-  local restore_glob
-  restore_glob="$(shopt -p -o noglob 2>/dev/null || true)"
-  set -o noglob
-  local -a words
-  read -ra words <<< "$flat"
-  eval "$restore_glob"
-
-  for word in "${words[@]}"; do
-    if $collecting; then
-      # If the token we're about to collect already contains a literal
-      # single quote, the expression was already quoted upstream.
-      # Flush and stop collecting.
-      if [[ "$word" == *"'"* ]]; then
-        if [[ -n "$marker_buf" ]]; then
-          # Should not normally happen (partial buf + quote), flush raw
-          output+="${marker_buf} "
-          marker_buf=""
-        fi
-        output+="${word} "
-        collecting=false
-        continue
-      fi
-
-      local is_boundary=false
-      case "$word" in
-        # Line-continuation artifact
-        "\\")
-          is_boundary=true ;;
-        # Command separators
-        "&&"|"||"|";"|"|")
-          is_boundary=true ;;
-        # Long flags (--ignore, --shard-id, etc.)
-        --*)
-          is_boundary=true ;;
-        # Short flags (-v, -s, -x, etc.) but NOT negative marker tokens
-        # like "not" which don't start with "-". Also skip -k/-m which
-        # would start a new marker (handled below).
-        -[a-zA-Z])
-          is_boundary=true ;;
-        # Test path (contains /)
-        */*)
-          is_boundary=true ;;
-        # Test file (ends with .py, possibly with ::method)
-        *.py|*.py::*)
-          is_boundary=true ;;
-        # Environment variable assignment preceding a command (FOO=bar)
-        *=*)
-          # Only treat as boundary if it looks like VAR=value, not
-          # pytest filter expressions like num_gpus=2 inside markers
-          if [[ "$word" =~ ^[A-Z_][A-Z0-9_]*= ]]; then
-            is_boundary=true
-          fi
-          ;;
-      esac
-
-      if $is_boundary; then
-        # Strip surrounding double quotes if present (from upstream
-        # single-to-double conversion); without this, wrapping below
-        # would produce '"expr"' with literal double-quote characters.
-        if [[ "$marker_buf" == '"'*'"' ]]; then
-          marker_buf="${marker_buf#\"}"
-          marker_buf="${marker_buf%\"}"
-        fi
-        # Flush the collected marker expression
-        if [[ "$marker_buf" == *" "* || "$marker_buf" == *"("* ]]; then
-          output+="'${marker_buf}' "
-        else
-          output+="${marker_buf} "
-        fi
-        collecting=false
-        marker_buf=""
-        # Check if this boundary word itself starts a new -m/-k
-        if [[ "$word" == "-m" || "$word" == "-k" ]]; then
-          output+="${word} "
-          collecting=true
-        # Drop stray backslash tokens silently
-        elif [[ "$word" == "\\" ]]; then
-          :
-        else
-          output+="${word} "
-        fi
-      else
-        # Accumulate into marker buffer
-        if [[ -n "$marker_buf" ]]; then
-          marker_buf+=" ${word}"
-        else
-          marker_buf="${word}"
-        fi
-      fi
-    elif [[ "$word" == "-m" || "$word" == "-k" ]]; then
-      output+="${word} "
-      collecting=true
-      marker_buf=""
-    else
-      output+="${word} "
-    fi
-  done
-
-  # Flush any trailing marker expression (marker at end of command)
-  if $collecting && [[ -n "$marker_buf" ]]; then
-    # Strip surrounding double quotes (see mid-stream flush comment)
-    if [[ "$marker_buf" == '"'*'"' ]]; then
-      marker_buf="${marker_buf#\"}"
-      marker_buf="${marker_buf%\"}"
-    fi
-    if [[ "$marker_buf" == *" "* || "$marker_buf" == *"("* ]]; then
-      output+="'${marker_buf}'"
-    else
-      output+="${marker_buf}"
-    fi
-  fi
-
-  echo "${output% }"
-}
-
-# shellcheck disable=SC2317  # Called indirectly by the EXIT trap.
-handle_amd_runner_exit() {
-  local exit_code=${1:-$?}
-  if [[ "${exit_code}" -ne 0 ]]; then
-    collect_rocm_failure_diagnostics "${exit_code}"
-  fi
-  return 0
-}
-
-# Catch both test failures and wrapper/setup failures. Runtime-specific cleanup
-# traps below replace this trap and call the same handler after cleanup.
-trap handle_amd_runner_exit EXIT
-
-###############################################################################
-# Main
-###############################################################################
-
-if is_native_runtime; then
-  echo "--- Native in-pod ROCm CI (AMD_CI_RUNTIME=${AMD_CI_RUNTIME:-unset}, NATIVE_CI=${NATIVE_CI:-unset})"
-  artifact_work_dir=""
-
-  # shellcheck disable=SC2317  # Called indirectly by the EXIT trap.
-  cleanup_native_workspace() {
-    local exit_code=$?
-    if [[ -n "${artifact_work_dir}" ]]; then
-      rm -rf "${artifact_work_dir}"
-    fi
-    handle_amd_runner_exit "${exit_code}"
-  }
-  trap cleanup_native_workspace EXIT
-
-  if [[ -n "${VLLM_TEST_COMMANDS:-}" ]]; then
-    commands="${VLLM_TEST_COMMANDS}"
-    commands_source="env"
-  else
-    commands="$*"
-    commands_source="argv"
-    if [[ -z "$commands" ]]; then
-      echo "Error: No test commands provided for native CI." >&2
-      exit 1
-    fi
-  fi
-
-  if [[ "$commands_source" == "argv" ]]; then
-    commands=$(re_quote_pytest_markers "$commands")
-  fi
-
-  if is_multi_node "$commands"; then
-    echo "Native CI does not support multi-node jobs yet."
-    exit 1
-  fi
-
-  if ! initialize_native_environment; then
-    echo "Failed to initialize the native test environment"
-    exit 1
-  fi
-  if [[ "${BUILDKITE_AGENT_META_DATA_QUEUE:-}" == *dpx* \
-    && "${VLLM_CI_EXPECTED_GPU_COUNT:-1}" != "0" ]]; then
-    check_dpx_gpu_exclusivity || exit 1
-  fi
-  if [[ "${commands}" == *python_only_compile.sh* ]]; then
-    # This no-GPU job validates the ROCm precompiled/editable install path,
-    # rather than CPU runtime platform selection.
-    VLLM_TARGET_DEVICE=rocm
-    export VLLM_TARGET_DEVICE
-  fi
-  if ! prepare_native_workspace "${commands}"; then
-    echo "Failed to prepare native test workspace"
-    exit 1
-  fi
-
-  export PYTHONPATH="${VLLM_CI_WORKSPACE:-/vllm-workspace}"
-
-  echo "Native test commands: $commands"
-  run_native_preflight || exit 1
-  echo "--- Test log"
-  # Keep AMD CI orchestration variables out of vLLM's runtime environment.
-  clear_ci_orchestration_env
-  /bin/bash -o pipefail -c "${commands}"
-  handle_pytest_exit "$?"
-fi
-
-# --- GPU initialization for legacy Docker execution ---
-echo "--- ROCm info"
-rocminfo
-
-# --- Docker status ---
-report_docker_usage
-
-# --- Pull test image ---
-echo "--- Pulling container"
-image_name="${VLLM_CI_FALLBACK_IMAGE:-rocm/vllm-ci:${BUILDKITE_COMMIT:-local}}"
-artifact_work_dir=""
-container_name="rocm_${BUILDKITE_COMMIT}_$(tr -dc A-Za-z0-9 < /dev/urandom | head -c 10; echo)"
-
-# shellcheck disable=SC2317  # Called indirectly by the EXIT trap.
-remove_docker_container() {
-  local exit_code=$?
-  if docker container inspect "${container_name}" >/dev/null 2>&1; then
-    docker rm -f "${container_name}" || true
-  fi
-  if [[ "${VLLM_CI_REMOVE_TEST_IMAGE:-0}" == "1" ]]; then
-    docker image rm -f "${image_name}" || true
-  else
-    # Keep images by default so later jobs on the same AMD node can reuse layers.
-    echo "Keeping ROCm test image locally: ${image_name}"
-  fi
-  if [[ -n "${artifact_work_dir}" ]]; then
-    rm -rf "${artifact_work_dir}"
-  fi
-  handle_amd_runner_exit "${exit_code}"
-}
-trap remove_docker_container EXIT
-
-# python_only_compile.sh runs `python setup.py develop` and needs the full repo tree
-# under /vllm-workspace (Dockerfile.rocm test stage: mkdir src && mv vllm).
-# The ROCm wheel artifact tarball only ships a thin tree (tests, etc.), so
-# artifact images cannot satisfy that test — use the full rocm/vllm-ci image.
-_cmd_probe="${VLLM_TEST_COMMANDS:-}"
-if [[ -z "${_cmd_probe}" ]]; then
-  _cmd_probe="$*"
-fi
-if [[ "${VLLM_CI_USE_ARTIFACTS:-0}" == "1" && "${_cmd_probe}" == *python_only_compile.sh* ]]; then
-  echo "INFO: disabling VLLM_CI_USE_ARTIFACTS for python_only_compile (requires full /vllm-workspace tree)"
-  export VLLM_CI_USE_ARTIFACTS=0
-fi
-unset -v _cmd_probe
-
-if ! prepare_artifact_image; then
-  echo "Using full ROCm CI image: ${image_name}"
-  docker pull "${image_name}" || exit 1
-fi
-
-# --- Prepare commands ---
-echo "--- Running container"
-
-HF_CACHE="$(realpath ~)/huggingface"
-mkdir -p "${HF_CACHE}"
-HF_MOUNT="/root/.cache/huggingface"
-
-# Hugging Face Hub defaults to 10s request/download timeouts, while the ROCm
-# CI image currently raises downloads to 60s. AMD model-test jobs routinely
-# start from a cold or partially-populated shared cache, and the 60s read cap
-# has still timed out before pytest reached the vLLM behavior under test.
-# Keep the CI default explicit and overridable from the Buildkite environment.
-: "${HF_HUB_DOWNLOAD_TIMEOUT:=300}"
-: "${HF_HUB_ETAG_TIMEOUT:=60}"
-
-# ---- Command source selection ----
-# Prefer VLLM_TEST_COMMANDS (preserves all inner quoting intact).
-# Fall back to $* for backward compatibility, but warn that inner
-# double-quotes will have been stripped by the calling shell.
-if [[ -n "${VLLM_TEST_COMMANDS:-}" ]]; then
-  commands="${VLLM_TEST_COMMANDS}"
-  commands_source="env"
-  echo "Commands sourced from VLLM_TEST_COMMANDS (quoting preserved)"
-else
-  commands="$*"
-  commands_source="argv"
-  if [[ -z "$commands" ]]; then
-    echo "Error: No test commands provided." >&2
-    echo "Usage:" >&2
-    echo "  Preferred:  VLLM_TEST_COMMANDS='...' bash $0" >&2
-    echo "  Legacy:     bash $0 \"commands here\"" >&2
-    exit 1
-  fi
-  echo "Commands sourced from positional args (legacy mode)"
-  echo "WARNING: Inner double-quotes in the command string may have been"
-  echo "  stripped by the calling shell. If you see syntax errors, switch to:"
-  echo "  export VLLM_TEST_COMMANDS='your commands here'"
-  echo "  bash $0"
-fi
-
-echo "Raw commands: $commands"
-
-# Only try to repair stripped pytest -m/-k quoting in legacy argv mode.
-# VLLM_TEST_COMMANDS preserves inner quoting already, and re-quoting that path
-# can corrupt embedded echo strings or otherwise well-formed shell fragments.
-if [[ "$commands_source" == "argv" ]]; then
-  commands=$(re_quote_pytest_markers "$commands")
-  echo "After re-quoting: $commands"
-else
-  echo "Skipping re-quoting for VLLM_TEST_COMMANDS input"
-fi
-
-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
-  # wheels.vllm.ai commit from the agent checkout for this test only.
-  vllm_standalone_merge_base=""
-  checkout="${BUILDKITE_BUILD_CHECKOUT_PATH:-}"
-  if [[ -z "${checkout}" || ! -d "${checkout}" ]]; then
-    checkout="."
-  fi
-  # Pass safe.directory per-command because Buildkite uses mixed user IDs.
-  if git -c "safe.directory=${checkout}" -C "${checkout}" rev-parse --is-inside-work-tree >/dev/null 2>&1; then
-    vllm_standalone_merge_base="$(
-      git -c "safe.directory=${checkout}" -C "${checkout}" merge-base HEAD origin/main 2>/dev/null || true
-    )"
-  fi
-  if [[ -z "${vllm_standalone_merge_base}" ]]; then
-    vllm_standalone_merge_base="${BUILDKITE_COMMIT:-}"
-  fi
-  echo "INFO: passing CI_STANDALONE_MERGE_BASE into container: ${vllm_standalone_merge_base}"
-  standalone_merge_base_env=(-e "CI_STANDALONE_MERGE_BASE=${vllm_standalone_merge_base}")
-fi
-
-MYPYTHONPATH="/vllm-workspace"
-
-container_job_id="${BUILDKITE_JOB_ID:-${BUILDKITE_PARALLEL_JOB:-0}}"
-container_job_id="${container_job_id//[^A-Za-z0-9_.-]/_}"
-container_job_id_short="${container_job_id:0:8}"
-CONTAINER_TMPDIR="/tmp/vllm-${container_job_id_short}"
-CONTAINER_CACHE_ROOT="/tmp/vllm-buildkite-${container_job_id}/cache"
-CONTAINER_PREFLIGHT="mkdir -p \"\$TMPDIR\" \"\$TIKTOKEN_RS_CACHE_DIR\" \"\$TORCHINDUCTOR_CACHE_DIR\" \"\$TRITON_CACHE_DIR\" \"\$VLLM_CACHE_ROOT\" \"\$XDG_CACHE_HOME\" && python -c \"import encodings, importlib.metadata as im, importlib.util as iu; [im.version(d) for d in ('transformers', 'torch', 'ray', 'sympy', 'markupsafe', 'vllm')]; missing=[m for m in ('torch.utils.model_zoo', 'transformers.models.nomic_bert', 'ray.dag', 'sympy.physics', 'markupsafe._speedups') if iu.find_spec(m) is None]; assert not missing, missing\""
-
-# Verify GPU access
-render_gid=$(getent group render | cut -d: -f3)
-if [[ -z "$render_gid" ]]; then
-  echo "Error: 'render' group not found. This is required for GPU access." >&2
-  exit 1
+n = torch.cuda.device_count()
+assert torch.version.hip and torch.cuda.is_available() and n == $gpus, \
+    f"preflight: expected $gpus ROCm GPU(s), found {n}"
+EOF
 fi
 
-# --- RDMA device passthrough (conditional) ---
-# If the host has RDMA devices, pass them through so tests like
-# test_moriio_connector can access ibverbs. On hosts without RDMA
-# hardware the tests will gracefully skip via _rdma_available().
-RDMA_FLAGS=""
-if [ -d /dev/infiniband ]; then
-  echo "RDMA devices detected on host, enabling passthrough"
-  RDMA_FLAGS="--device /dev/infiniband --cap-add=IPC_LOCK"
+# ---------------------------------------------------------------------------
+# Run
+# ---------------------------------------------------------------------------
+# Scrub orchestration vars before handing control to test commands.
+# (VLLM_CI_* that are test knobs are intentionally left — see tests/ci_envs.py)
+unset VLLM_TEST_COMMANDS VLLM_CI_EXPECTED_GPU_COUNT VLLM_CI_DIAGNOSTICS_DIR
+
+echo "Test commands: $commands"
+echo "--- Test log"
+_teardown_log workload_started "timeout_s=${CONTAINER_TIMEOUT_S:-0}"
+
+rc=0
+if [[ ${CONTAINER_TIMEOUT_S:-0} -gt 0 ]]; then
+  # timeout fires TERM, then KILL after 10 s — mirrors the old docker run path.
+  # Not --foreground: the whole process group (pytest, engine workers) must stop, not just bash.
+  # That group is in the background of the job's PTY, where reading the terminal stops it
+  # (SIGTTIN: apt, transformers' trust_remote_code prompt), so stdin is /dev/null.
+  timeout --verbose --signal=TERM --kill-after=10s \
+    "${CONTAINER_TIMEOUT_S}s" bash -o pipefail -c "$commands" </dev/null || rc=$?
 else
-  echo "No RDMA devices found on host, RDMA tests will be skipped"
+  bash -o pipefail -c "$commands" </dev/null || rc=$?
 fi
 
-# --- Route: multi-node vs single-node ---
-clear_ci_orchestration_env
-if is_multi_node "$commands"; then
-  echo "--- Multi-node job detected"
-  DCKR_VER=$(docker --version | sed 's/Docker version \(.*\), build .*/\1/')
-  export DCKR_VER
-
-  # Parse the bracket syntax:  prefix ; [node0_cmds] && [node1_cmds]
-  #   BASH_REMATCH[1] = prefix (everything before first bracket)
-  #   BASH_REMATCH[2] = comma-separated node0 commands
-  #   BASH_REMATCH[3] = comma-separated node1 commands
-  if [[ "$commands" =~ ^(.*)\[(.*)"] && ["(.*)\]$ ]]; then
-    prefix=${BASH_REMATCH[1]//;/}
-    echo "PREFIX: ${prefix}"
-
-    export composite_command="(command rocm-smi || true)"
-    saved_IFS=$IFS
-    IFS=','
-    read -ra node0 <<< "${BASH_REMATCH[2]}"
-    read -ra node1 <<< "${BASH_REMATCH[3]}"
-    IFS=$saved_IFS
+[[ $rc == 5   ]] && rc=0   # pytest exit 5 = no tests collected, not a failure
+[[ $rc == 124 ]] && _teardown_log deadline_exhausted \
+  "configured_s=${CONTAINER_TIMEOUT_S:-0}" "status=$rc"
+_teardown_log workload_finished "status=$rc"
 
-    if [[ ${#node0[@]} -ne ${#node1[@]} ]]; then
-      echo "Warning: node0 has ${#node0[@]} commands, node1 has ${#node1[@]}. They will be paired by index."
-    fi
-
-    for i in "${!node0[@]}"; do
-      command_node_0=${node0[i]//\"/}
-      command_node_1=${node1[i]//\"/}
-
-      step_cmd="./.buildkite/scripts/run-multi-node-test.sh /vllm-workspace/tests 2 2 ${image_name} '${command_node_0}' '${command_node_1}'"
-      echo "COMMANDS: ${step_cmd}"
-      composite_command="${composite_command} && ${step_cmd}"
-    done
-
-    /bin/bash -c "${composite_command}"
-    exit_code=$?
-    handle_pytest_exit "$exit_code"
-  else
-    echo "Multi-node job detected but failed to parse bracket command syntax."
-    echo "Expected format: prefix ; [node0_cmd1, node0_cmd2] && [node1_cmd1, node1_cmd2]"
-    echo "Got: $commands"
-    exit 111
-  fi
-else
-  echo "--- Single-node job"
-  echo "Render devices: $BUILDKITE_AGENT_META_DATA_RENDER_DEVICES"
-  docker_run_terminal_args=(-i)
-  if [[ "${ROCM_DOCKER_TTY}" == "1" ]]; then
-    docker_run_terminal_args+=(-t)
-    echo "Docker interactive stdin: enabled; TTY allocation: enabled"
-  else
-    echo "Docker interactive stdin: enabled; TTY allocation: disabled"
-  fi
-
-  ulimit_core_hard=$(ulimit -H -c)
-  if [[ "$ulimit_core_hard" == "unlimited" ]]; then
-    # docker run can't pass "unlimited" to --ulimit
-    ulimit_core_hard="-1"
-  fi
-   # Disable core dumps in the ROCm test container unless the ROCm debug agent is enabled
-  coredump_flags=(--ulimit "core=0:$ulimit_core_hard")
-  if [[ "$commands" == *"ROCm debug agent enabled"* ]]; then
-    # Works around https://github.com/rocm/rocm-systems/issues/6206
-    coredump_flags=(-e 'HSA_COREDUMP_PATTERN="/tmp/gpucore.%p"')
-  else
-    echo "ROCm debug agent not enabled, coredumps are disabled in the test container."
-  fi
-
-  exit_code=0
-  # shellcheck disable=SC2086  # word splitting is intentional: both hold multiple docker flags
-  run_docker_with_ci_timeout docker run \
-    "${docker_run_terminal_args[@]}" \
-    --device /dev/kfd $BUILDKITE_AGENT_META_DATA_RENDER_DEVICES \
-    $RDMA_FLAGS \
-    --network=host \
-    --shm-size=16gb \
-    --group-add "$render_gid" \
-    --rm \
-    "${coredump_flags[@]}" \
-    -e HF_TOKEN \
-    -e "HF_HUB_DOWNLOAD_TIMEOUT=${HF_HUB_DOWNLOAD_TIMEOUT}" \
-    -e "HF_HUB_ETAG_TIMEOUT=${HF_HUB_ETAG_TIMEOUT}" \
-    -e AWS_ACCESS_KEY_ID \
-    -e AWS_SECRET_ACCESS_KEY \
-    -e BUILDKITE_PARALLEL_JOB \
-    -e BUILDKITE_PARALLEL_JOB_COUNT \
-    -e TERM \
-    -e FORCE_COLOR \
-    -e CLICOLOR_FORCE \
-    -e PY_COLORS \
-    -e PYTHONFAULTHANDLER \
-    -e PYTEST_ADDOPTS \
-    -e PYTEST_TIMEOUT \
-    -v "${HF_CACHE}:${HF_MOUNT}" \
-    -e "HF_HOME=${HF_MOUNT}" \
-    -e "PYTHONPATH=${MYPYTHONPATH}" \
-    -e "TMPDIR=${CONTAINER_TMPDIR}/tmp" \
-    -e "TIKTOKEN_RS_CACHE_DIR=${HF_MOUNT}/tiktoken-rs-cache" \
-    -e "TORCHINDUCTOR_CACHE_DIR=${CONTAINER_CACHE_ROOT}/torchinductor" \
-    -e "TRITON_CACHE_DIR=${CONTAINER_CACHE_ROOT}/triton" \
-    -e "VLLM_CACHE_ROOT=${CONTAINER_CACHE_ROOT}/vllm" \
-    -e "XDG_CACHE_HOME=${CONTAINER_CACHE_ROOT}/xdg" \
-    -e "PYTORCH_ROCM_ARCH=" \
-    "${cpu_platform_env[@]}" \
-    "${standalone_merge_base_env[@]}" \
-    --name "${container_name}" \
-    "${image_name}" \
-    /bin/bash -c "${CONTAINER_PREFLIGHT} && ${commands}" || exit_code=$?
-  handle_pytest_exit "${exit_code}"
-fi
+[[ $rc != 0 ]] && diagnostics
+exit $rc
diff --git a/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh b/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh
index 1f598b277bb2..a5eadbcc2320 100755
--- a/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh
+++ b/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh
@@ -51,7 +51,8 @@ kernel_tests() {
     pytest -x -v -s tests/kernels/moe/test_cpu_int4_moe.py
     pytest -x -v -s tests/kernels/mamba/test_cpu_short_conv.py
     pytest -x -v -s tests/kernels/mamba/test_causal_conv1d.py
-    pytest -x -v -s tests/kernels/mamba/test_mamba_ssm.py"
+    pytest -x -v -s tests/kernels/mamba/test_mamba_ssm.py
+    pytest -x -v -s tests/v1/worker/cpu"
 }
 
 model_tests() {
@@ -68,6 +69,7 @@ model_tests() {
   docker exec "$CONTAINER_NAME" bash -c "
     set -e
     pytest -x -v -s tests/models/multimodal/generation/test_whisper.py -m cpu_model
+    pytest -x -v -s tests/v1/e2e/test_cpu_linear_attn_chunked_prefix.py
     pytest -x -v -s 'tests/models/language/pooling/test_embedding.py::test_models[sentence-transformers/all-MiniLM-L12-v2]'"
 
   docker exec "$CONTAINER_NAME" bash -c "
diff --git a/.buildkite/scripts/install-kv-offload.sh b/.buildkite/scripts/install-kv-offload.sh
index eb8bab9e393f..034d8e76b1c3 100644
--- a/.buildkite/scripts/install-kv-offload.sh
+++ b/.buildkite/scripts/install-kv-offload.sh
@@ -10,7 +10,7 @@ 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.1"
+    "kvcr @ git+https://github.com/ai-dynamo/kvcr.git@8c6fb79be4dfd9d7bb797c7dc4656790273bc786"
 
 # Keep only the NIXL wheel matching the CI image's CUDA runtime.
 NIXL_VERSION=$(uv pip show --system nixl | sed -n 's/^Version: //p')
diff --git a/.buildkite/scripts/publish-release-images.sh b/.buildkite/scripts/publish-release-images.sh
index bdd24be4ee45..061f73aa876e 100755
--- a/.buildkite/scripts/publish-release-images.sh
+++ b/.buildkite/scripts/publish-release-images.sh
@@ -10,9 +10,9 @@ set -euo pipefail
 
 TARGET="${1:-all}"
 case "${TARGET}" in
-  cuda-13-0 | cuda-13-0-efa | rocm | rocm72 | xpu | cpu | all) ;;
+  cuda-13-0 | cuda-13-0-efa | rocm | xpu | cpu | all) ;;
   *)
-    echo "Usage: $0 {cuda-13-0|cuda-13-0-efa|rocm|rocm72|xpu|cpu|all}"
+    echo "Usage: $0 {cuda-13-0|cuda-13-0-efa|rocm|xpu|cpu|all}"
     exit 2
     ;;
 esac
@@ -117,35 +117,37 @@ if target_enabled cuda-13-0-efa; then
 fi
 
 # ---- ROCm ----
-# Each ROCm stack is identified by its base Dockerfile (see rocm/stack.sh). The
-# default stack also owns the plain :latest / :v<ver> tags.
+# Copied from ECR (publish.sh). The default stack (rocm) also owns the plain :latest / :v<ver> tags.
+
+# pin_rocm <ref>  — <ref>@<digest>, so what was checked is what is published
+pin_rocm() {
+  local d
+  d=$(docker buildx imagetools inspect "$1" | awk '$1 == "Digest:" { print $2; exit }')
+  [[ $d =~ ^sha256:[0-9a-f]{64}$ ]] || { echo "cannot resolve the digest of $1" >&2; return 1; }
+  echo "$1@$d"
+}
 
 publish_rocm_stack() {
-  local base_dockerfile="$1" default="$2" suffix
-  (
-    # shellcheck source=.buildkite/scripts/rocm/stack.sh
-    source .buildkite/scripts/rocm/stack.sh "$base_dockerfile"
-    local suffixes=("-${ROCM_STACK_VARIANT}")
-    [[ "$default" == "1" ]] && suffixes=("" "-${ROCM_STACK_VARIANT}")
-    docker pull "$ROCM_STACK_ECR_IMAGE"
-    docker pull "$ROCM_STACK_ECR_BASE"
-    for suffix in "${suffixes[@]}"; do
-      for tag in "latest${suffix}" "v${RELEASE_VERSION}${suffix}"; do
-        docker tag "$ROCM_STACK_ECR_IMAGE" "vllm/vllm-openai-rocm:${tag}"
-        docker tag "$ROCM_STACK_ECR_BASE" "vllm/vllm-openai-rocm:${tag}-base"
-        docker push "vllm/vllm-openai-rocm:${tag}"
-        docker push "vllm/vllm-openai-rocm:${tag}-base"
-      done
-    done
-  )
+  local v s sfx img base ver ecr=public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT} hub=vllm/vllm-openai-rocm
+  v=$(buildkite-agent meta-data get "rocm-$1-variant")
+  img=$(pin_rocm "$ecr-$v")
+  base=$(pin_rocm "$ecr-$v-base")
+  # Another build of this commit (e.g. a branch push) may have retagged it with its own version.
+  ver=$(docker buildx imagetools inspect "$img" \
+    --format '{{index .Image.Config.Labels "org.opencontainers.image.version"}}')
+  if [[ ${ver%%+*} != "$RELEASE_VERSION" ]]; then
+    echo "ERROR: $img is version '$ver', not $RELEASE_VERSION" >&2
+    exit 1
+  fi
+  sfx=("-$v"); if [[ $1 == rocm ]]; then sfx+=(""); fi
+  for s in "${sfx[@]}"; do
+    docker buildx imagetools create -t "$hub:latest$s" -t "$hub:v${RELEASE_VERSION}$s" "$img"
+    docker buildx imagetools create -t "$hub:latest$s-base" -t "$hub:v${RELEASE_VERSION}$s-base" "$base"
+  done
 }
 
 if target_enabled rocm; then
-  publish_rocm_stack docker/Dockerfile.rocm_base 1
-fi
-
-if target_enabled rocm72; then
-  publish_rocm_stack docker/Dockerfile.rocm_72_base 0
+  publish_rocm_stack rocm
 fi
 
 # ---- XPU ----
diff --git a/.buildkite/scripts/push-nightly-builds-rocm.sh b/.buildkite/scripts/push-nightly-builds-rocm.sh
deleted file mode 100644
index 8b120883eff4..000000000000
--- a/.buildkite/scripts/push-nightly-builds-rocm.sh
+++ /dev/null
@@ -1,86 +0,0 @@
-#!/bin/bash
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-#
-# Push a ROCm stack's nightly base image and nightly image from ECR to Docker Hub
-# under vllm/vllm-openai-rocm, as base-nightly-<variant>, nightly-<variant> and their
-# -<commit> forms, then prune old nightly tags. With --default, also push the plain
-# base-nightly and nightly tags for the same images.
-# Run when NIGHTLY=1 after the stack's release image has been pushed to ECR.
-#
-# Usage: push-nightly-builds-rocm.sh <base-dockerfile> [--default]
-#   e.g. "docker/Dockerfile.rocm_base --default" (ROCm 10.0, also :nightly)
-#   and "docker/Dockerfile.rocm_72_base" (ROCm 7.2, :nightly-rocm72).
-#
-# Local testing (no push to Docker Hub):
-#   BUILDKITE_COMMIT=<commit-with-image-in-ecr> DRY_RUN=1 \
-#     bash .buildkite/scripts/push-nightly-builds-rocm.sh docker/Dockerfile.rocm_base --default
-# Requires: AWS CLI configured (for ECR public login), Docker. For full run: Docker Hub login.
-
-set -euxo pipefail
-
-usage() {
-  echo "Usage: $0 <base-dockerfile> [--default]" >&2
-}
-
-PUSH_DEFAULT=0
-case "$#" in
-  1) ;;
-  2)
-    if [[ "$2" != "--default" ]]; then
-      usage
-      exit 2
-    fi
-    PUSH_DEFAULT=1
-    ;;
-  *)
-    usage
-    exit 2
-    ;;
-esac
-
-# Use BUILDKITE_COMMIT from env (required; set to a commit that has the image in ECR for local test)
-BUILDKITE_COMMIT="${BUILDKITE_COMMIT:?Set BUILDKITE_COMMIT to the commit SHA that has the image in ECR (e.g. from a previous release pipeline run)}"
-DRY_RUN="${DRY_RUN:-0}"
-
-# shellcheck source=.buildkite/scripts/rocm/stack.sh
-source "$(dirname "${BASH_SOURCE[0]}")/rocm/stack.sh" "$1"
-DOCKERHUB_REPO="vllm/vllm-openai-rocm"
-VARIANT="$ROCM_STACK_VARIANT"
-
-BASE_TAGS=("base-nightly-${VARIANT}" "base-nightly-${VARIANT}-${BUILDKITE_COMMIT}")
-TAGS=("nightly-${VARIANT}" "nightly-${VARIANT}-${BUILDKITE_COMMIT}")
-if [[ "$PUSH_DEFAULT" == "1" ]]; then
-  BASE_TAGS+=("base-nightly" "base-nightly-${BUILDKITE_COMMIT}")
-  TAGS+=("nightly" "nightly-${BUILDKITE_COMMIT}")
-fi
-
-echo "Pushing base image $ROCM_STACK_ECR_BASE as: ${BASE_TAGS[*]}"
-echo "Pushing release image $ROCM_STACK_ECR_IMAGE as: ${TAGS[*]}"
-
-# Login to ECR and pull the stack's images
-aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7
-docker pull "$ROCM_STACK_ECR_BASE"
-docker pull "$ROCM_STACK_ECR_IMAGE"
-
-for tag in "${BASE_TAGS[@]}"; do docker tag "$ROCM_STACK_ECR_BASE" "$DOCKERHUB_REPO:$tag"; done
-for tag in "${TAGS[@]}"; do docker tag "$ROCM_STACK_ECR_IMAGE" "$DOCKERHUB_REPO:$tag"; done
-
-if [[ "$DRY_RUN" == "1" ]]; then
-  echo "[DRY_RUN] Local tags created. Exiting without push."
-  exit 0
-fi
-
-# Push to Docker Hub (docker-login plugin runs before this step in CI)
-for tag in "${BASE_TAGS[@]}" "${TAGS[@]}"; do docker push "$DOCKERHUB_REPO:$tag"; done
-echo "Pushed $DOCKERHUB_REPO: ${BASE_TAGS[*]} ${TAGS[*]}"
-
-# Keep only the last 14 builds of each tag family. The plain nightly- families
-# exclude the variant ones sharing their prefix; commit SHAs never start with "rocm".
-CLEANUP=.buildkite/scripts/cleanup-nightly-builds.sh
-bash "$CLEANUP" "nightly-${VARIANT}-" "$DOCKERHUB_REPO"
-bash "$CLEANUP" "base-nightly-${VARIANT}-" "$DOCKERHUB_REPO"
-if [[ "$PUSH_DEFAULT" == "1" ]]; then
-  bash "$CLEANUP" "nightly-" "$DOCKERHUB_REPO" "nightly-rocm"
-  bash "$CLEANUP" "base-nightly-" "$DOCKERHUB_REPO" "base-nightly-rocm"
-fi
diff --git a/.buildkite/scripts/rocm/aiter-nightly-overlay.sh b/.buildkite/scripts/rocm/aiter-nightly-overlay.sh
deleted file mode 100755
index bed4a3e09aec..000000000000
--- a/.buildkite/scripts/rocm/aiter-nightly-overlay.sh
+++ /dev/null
@@ -1,116 +0,0 @@
-#!/bin/bash
-# Install the AITER nightly wheel over this build's ci_base.
-#
-# Runs in CI when AITER_NIGHTLY=1, once the stock ci_base it builds on is
-# published. In order:
-#   1. Select the newest nightly wheel that fits the stock ci_base
-#      (select_aiter_nightly_wheel.py, beside this script, run inside that image)
-#   2. Build a new image: pip install that wheel over the stock ci_base
-#   3. Push the new image as this build's ci_base-build tag, replacing the stock
-#      one, so later steps and native jobs use it; record its digest
-#
-# Buildkite meta-data, a key-value store shared by every job in this build:
-#   reads   rocm-ci-base-image       the stock ci_base
-#   r/w     aiter-nightly-wheel-url  the wheel step 1 chose; a retry reads it
-#                                    back and installs that same wheel
-#   writes  aiter-nightly-ci-base    the new image as tag@digest; the vLLM
-#                                    image build uses it as its ci_base
-#
-# Exit status:
-#   10  no wheel fits this ci_base
-#   11  the wheel does not install or its modules do not load
-#   Anything else is unexpected and passed through as is.
-#   Buildkite retries this step once on exit status 1 or 2, which it treats as
-#   transient Docker/registry failures (amd_infra_retry in the template).
-#   10 and 11 sit outside that range: they are real AITER failures, and a retry
-#   would only fail the same way.
-
-set -euo pipefail
-
-SELECT_WHEEL="$(dirname "${BASH_SOURCE[0]}")/select_aiter_nightly_wheel.py"
-BUILD_TAG="rocm/vllm-dev:ci_base-build-${BUILDKITE_BUILD_ID:?BUILDKITE_BUILD_ID is required}"
-
-# Scratch space for the build context and captured errors. The trap deletes it
-# however the script exits: success, fail, or set -e.
-scratch="$(mktemp -d)"
-trap 'rm -rf "${scratch}"' EXIT
-errors="${scratch}/errors"
-
-fail() {
-    local status="$1" title="$2" detail="$3"
-    echo "+++ :x: AITER nightly: ${title}"
-    echo "${detail}" >&2
-    # shellcheck disable=SC2016 # the backticks are a literal markdown fence
-    printf ':x: **AITER nightly: %s**\n\n```\n%s\n```\n' "${title}" "${detail}" \
-        | buildkite-agent annotate --style error --context aiter-nightly
-    exit "${status}"
-}
-
-# Set by ensure-ci-base-amd; this step's depends_on guarantees it is there.
-ci_base="$(buildkite-agent meta-data get rocm-ci-base-image)"
-# Set below on the first attempt. Meta-data is per build, so a retry reuses
-# that wheel, not a newer nightly.
-pinned="$(buildkite-agent meta-data get aiter-nightly-wheel-url --default "")"
-
-echo "--- :docker: Pulling ${ci_base}"
-docker pull "${ci_base}"
-
-echo "--- :mag: Selecting the AITER nightly for ${ci_base}"
-set +e
-out="$(docker run --rm -i --entrypoint python3 "${ci_base}" \
-    -P - ${pinned:+--wheel-url "${pinned}"} \
-    < "${SELECT_WHEEL}" 2> "${errors}")"
-status=$?
-set -e
-cat "${errors}" >&2
-case "${status}" in
-    0) ;;
-    10) fail 10 "no wheel for this image" "${out}" ;;
-    *) fail "${status}" "could not select a wheel" "$(tail -n 20 "${errors}")" ;;
-esac
-# The first line is the wheel's URL; the rest, the annotation naming it.
-wheel_url="$(head -n 1 <<< "${out}")"
-note="$(tail -n +2 <<< "${out}")"
-buildkite-agent meta-data set aiter-nightly-wheel-url "${wheel_url}"
-
-cat > "${scratch}/Dockerfile" <<'EOF'
-ARG CI_BASE_IMAGE
-FROM ${CI_BASE_IMAGE}
-ARG AITER_WHEEL_URL
-RUN python3 -m pip install --no-cache-dir "$AITER_WHEEL_URL"
-# A wheel built against a newer glibc or libstdc++ than ci_base installs fine
-# but fails every test that loads it; fail here instead.
-RUN aiter="$(python3 -c 'import importlib.util as u; print(u.find_spec("aiter").origin)')" \
-    && ! find "${aiter%/*}" -name '*.so' -exec ldd {} + 2>&1 | grep "version .* not found"
-EOF
-
-echo "--- :docker: Installing ${wheel_url##*/}"
-if ! docker buildx build \
-    --builder default \
-    --provenance=false \
-    --progress plain \
-    --build-arg "CI_BASE_IMAGE=${ci_base}" \
-    --build-arg "AITER_WHEEL_URL=${wheel_url}" \
-    -t "${BUILD_TAG}" \
-    "${scratch}"; then
-    fail 11 "the wheel does not install" \
-        "Installing or loading ${wheel_url##*/} over ${ci_base} failed; see this step's log."
-fi
-
-docker push "${BUILD_TAG}"
-# Handed off as tag@digest: build-test-image.sh rejects a ci_base that is not
-# digest-pinned.
-digest="$(docker buildx imagetools inspect "${BUILD_TAG}" | awk '$1 == "Digest:" { print $2; exit }')"
-if [[ ! "${digest}" =~ ^sha256:[0-9a-f]{64}$ ]]; then
-    echo "Could not resolve the pushed digest of ${BUILD_TAG}" >&2
-    exit 1
-fi
-# The vLLM image build uses this as its ci_base.
-buildkite-agent meta-data set aiter-nightly-ci-base "${BUILD_TAG}@${digest}"
-# The annotation starts with :warning: when the wheel is stale.
-style=info
-if [[ "${note}" == :warning:* ]]; then
-    style=warning
-fi
-echo "${note}" | buildkite-agent annotate --style "${style}" --context aiter-nightly
-echo "--- :white_check_mark: ${BUILD_TAG}@${digest}"
diff --git a/.buildkite/scripts/rocm/build-ci-base.sh b/.buildkite/scripts/rocm/build-ci-base.sh
deleted file mode 100755
index 7d4eb381bc27..000000000000
--- a/.buildkite/scripts/rocm/build-ci-base.sh
+++ /dev/null
@@ -1,48 +0,0 @@
-#!/usr/bin/env bash
-# Build the ROCm ci_base image, optionally from a freshly rebuilt ROCm base.
-
-set -euo pipefail
-
-metadata_get() {
-    local key="$1"
-    if command -v buildkite-agent >/dev/null 2>&1; then
-        buildkite-agent meta-data get "${key}" 2>/dev/null || true
-    fi
-}
-
-main() {
-    local base_refreshed=""
-    local selected_base=""
-
-    # The content identity below describes this checkout. Do not let legacy
-    # external pipeline settings switch the build to a remote source tree.
-    export REMOTE_VLLM=0
-    unset VLLM_BRANCH
-
-    base_refreshed="$(metadata_get rocm-base-refresh)"
-    selected_base="$(metadata_get rocm-base-image)"
-    if [[ -n "${selected_base}" ]]; then
-        if [[ ! "${selected_base}" =~ @sha256:[0-9a-f]{64}$ ]]; then
-            echo "Selected ROCm base handoff is not digest-pinned: ${selected_base}" >&2
-            return 1
-        fi
-        export BASE_IMAGE="${selected_base}"
-        echo "Using selected ROCm base image for ci_base: ${BASE_IMAGE}"
-    elif [[ "${BUILDKITE:-false}" == "true" ]]; then
-        echo "Required ROCm base handoff metadata is missing" >&2
-        return 1
-    fi
-
-    if [[ "${base_refreshed}" == "1" ]]; then
-        export CI_BASE_PUSH_STABLE_TAG
-
-        CI_BASE_PUSH_STABLE_TAG="$(metadata_get rocm-base-push-stable-tag)"
-        CI_BASE_PUSH_STABLE_TAG="${CI_BASE_PUSH_STABLE_TAG:-0}"
-
-        echo "Push stable ci_base tag: ${CI_BASE_PUSH_STABLE_TAG}"
-    fi
-
-    bash .buildkite/scripts/ci-bake-rocm.sh ci-base-rocm-ci-with-deps
-}
-
-main "$@"
diff --git a/.buildkite/scripts/rocm/build-config.sh b/.buildkite/scripts/rocm/build-config.sh
deleted file mode 100644
index f6c6ff56ae38..000000000000
--- a/.buildkite/scripts/rocm/build-config.sh
+++ /dev/null
@@ -1,85 +0,0 @@
-#!/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 the legacy ROCm 7.2 stack, set CI_ROCM_DOCKERFILE_BASE=docker/Dockerfile.rocm_72_base
-# and CI_ROCM_DOCKERFILE=docker/Dockerfile.rocm_72.
-# 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/build-test-image.sh b/.buildkite/scripts/rocm/build-test-image.sh
deleted file mode 100755
index 69c4f77d20ca..000000000000
--- a/.buildkite/scripts/rocm/build-test-image.sh
+++ /dev/null
@@ -1,77 +0,0 @@
-#!/usr/bin/env bash
-# Build the ROCm CI test image or wheel artifact.
-#
-# When base selection changes, build the full image so downstream ROCm tests
-# validate the selected base -> ci_base -> CI image chain.
-
-set -euo pipefail
-
-metadata_get() {
-    local key="$1"
-    if command -v buildkite-agent >/dev/null 2>&1; then
-        buildkite-agent meta-data get "${key}" 2>/dev/null || true
-    fi
-}
-
-load_digest_handoff() {
-    local metadata_key="$1"
-    local env_name="$2"
-    local description="$3"
-    local image_ref=""
-
-    image_ref="$(metadata_get "${metadata_key}")"
-    if [[ -z "${image_ref}" ]]; then
-        return 1
-    fi
-    if [[ ! "${image_ref}" =~ @sha256:[0-9a-f]{64}$ ]]; then
-        echo "${description} is not digest-pinned: ${image_ref}" >&2
-        return 1
-    fi
-
-    printf -v "${env_name}" '%s' "${image_ref}"
-    export "${env_name?}"
-    echo "Using ${description}: ${image_ref}"
-}
-
-main() {
-    local base_refreshed=0
-
-    # This job always builds the checked-out commit. Some externally generated
-    # pipeline templates still inject remote-fetch settings; do not let those
-    # settings make the source identity commit-specific or bypass local edits.
-    export REMOTE_VLLM=0
-    unset VLLM_BRANCH
-
-    if ! load_digest_handoff \
-        rocm-ci-base-image CI_BASE_IMAGE "ROCm ci_base handoff"; then
-        if [[ "${BUILDKITE:-false}" == "true" ]]; then
-            echo "Required ROCm ci_base handoff metadata is missing or invalid" >&2
-            return 1
-        fi
-        echo "No ROCm ci_base handoff metadata found; using the local default"
-    fi
-
-    if ! load_digest_handoff \
-        rocm-base-image BASE_IMAGE "ROCm base handoff"; then
-        if [[ "${BUILDKITE:-false}" == "true" ]]; then
-            echo "Required ROCm base handoff metadata is missing or invalid" >&2
-            return 1
-        fi
-        echo "No ROCm base handoff metadata found; using the local default"
-    fi
-
-    if [[ "$(metadata_get rocm-base-refresh)" == "1" ]]; then
-        echo "The selected ROCm base differs from the current stable base"
-        base_refreshed=1
-    fi
-
-    if [[ "${ROCM_CI_ARTIFACT_ONLY:-0}" == "1" && "${base_refreshed}" != "1" ]]; then
-        echo "ROCM_CI_ARTIFACT_ONLY=1; building ROCm wheel artifact only"
-        IMAGE_TAG="" bash .buildkite/scripts/ci-bake-rocm.sh test-rocm-ci-with-artifacts
-        return
-    fi
-
-    bash .buildkite/scripts/ci-bake-rocm.sh test-rocm-ci-with-wheel
-}
-
-main "$@"
diff --git a/.buildkite/scripts/rocm/build.sh b/.buildkite/scripts/rocm/build.sh
new file mode 100755
index 000000000000..b61dcb913f87
--- /dev/null
+++ b/.buildkite/scripts/rocm/build.sh
@@ -0,0 +1,330 @@
+#!/usr/bin/env bash
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+#
+# Build and promote ROCm images via docker/docker-bake-rocm.hcl.
+# Pushes only to $ROCM_CI_REPO and the BuildKit cache that goes with it.
+#
+# Usage: build.sh <stack> <mode>
+#
+# Modes:
+#   keys     Print current Ks/Kr keys and whether their images exist (on-call debug).
+#   base     Ensure stack, runtime and test-deps images exist; write their refs to Buildkite metadata.
+#   ci       Build and push this build's per-commit test image.
+#   release  Release image + wheels from source into the release repo (release pipeline, AWS agents).
+#   release-base  Only the release stack and runtime (a serialized step, so one run builds a new key).
+#
+# Stack: rocm (TheRock ROCm 10). Force a stack rebuild: bump ARG REFRESH
+#        in docker/Dockerfile.rocm_base
+#
+# Trust: ROCM_CI_REPO is set by ci-infra (CI) or release-pipeline.yaml (release).
+#   Trusted main builds: rocm/vllm-ci
+#   PR / fork builds:    rocm/vllm-dev    (pr<N>- tags and caches; trusted builds never read it)
+#   Other untrusted:     rocm/vllm-dev    (br-<branch>-<hash>- tags and caches: nightlies, branch builds)
+#   Release pipeline:    its ECR repo     (never reads rocm/vllm-ci or its cache: it builds the
+#                                          stack and runtime from source, with only its own ECR cache;
+#                                          csrc objects come from the shared S3 sccache, as in CUDA)
+# ROCM_CI_TRUSTED_REPO / ROCM_CI_CACHE_REPO only point a local run at a sandbox registry.
+set -euo pipefail
+
+# ---------------------------------------------------------------------------
+# Arguments and trust
+# ---------------------------------------------------------------------------
+STACK=$1
+MODE=$2
+REPO=${ROCM_CI_REPO:?ROCM_CI_REPO must be set (by ci-infra or release-pipeline.yaml)}
+TRUSTED=${ROCM_CI_TRUSTED_REPO:-rocm/vllm-ci}
+
+export STACK PUSH=true
+export CACHE=${ROCM_CI_CACHE_REPO:-rocm/vllm-ci-cache}  # every CI build reads the trusted cache
+export CACHE_TO=$CACHE:                     # ref prefix trusted builds write (HCL appends <stack>-<name>)
+export BUILDKIT_PROGRESS=${BUILDKIT_PROGRESS:-plain}  # TTY progress redraws bloat CI logs
+
+# PR builds write to a per-PR cache and prefix their image tags. The release pipeline's agents
+# cannot write rocm/vllm-ci or its cache: release writes its images and its own cache to $REPO,
+# so the next release reuses this one's all-arch compile.
+PFX=""
+if [[ $MODE == release* ]]; then
+  CACHE="" CACHE_TO=$REPO:cache-   # and reads only that cache: CI agents also run PR code
+elif [[ $REPO != "$TRUSTED" ]]; then
+  if [[ ${BUILDKITE_PULL_REQUEST:-false} != false ]]; then
+    PFX="pr$BUILDKITE_PULL_REQUEST-"
+  else   # not a PR (AITER nightly, branch or manual builds): one namespace per branch
+    br=${BUILDKITE_BRANCH:-local}   # the hash keeps e.g. feat/x and feat-x apart
+    PFX="br-$(printf %s "$br" | tr -c 'A-Za-z0-9.-' - | cut -c1-40)-$(printf %s "$br" | sha256sum | cut -c1-8)-"
+  fi
+  CACHE_TO=$REPO:cache-$PFX
+  # Only the compile stages' cache: everything else is cheap or in the trusted cache.
+  export CACHE_CI=false
+fi
+
+# Stack Dockerfile — literal path used by ci-infra ci_selector (images.py).
+case $STACK in
+  rocm) SDF=docker/Dockerfile.rocm_base ;;
+  *) echo "Unknown stack: $STACK" >&2; exit 2 ;;
+esac
+
+# ---------------------------------------------------------------------------
+# BuildKit builder
+# ---------------------------------------------------------------------------
+# Prefer the host buildkitd (it has its own GC policy and cache). Fall back
+# to a docker-container builder if no socket is present. Release RUN steps must reach EC2 IMDS
+# for sccache's S3 credentials: their fallback builder uses the host network, under its own name.
+B=vllm-rocm net=()
+[[ $MODE == release* ]] && B=vllm-rocm-release net=(--driver-opt network=host)
+if ! docker buildx inspect "$B" >/dev/null 2>&1; then
+  if [[ -S /run/buildkit/buildkitd.sock ]]; then
+    drv=(--driver remote unix:///run/buildkit/buildkitd.sock)
+  else
+    drv=(--driver docker-container "${net[@]}" \
+      ${BUILDKITD_CONFIG:+--buildkitd-config "$BUILDKITD_CONFIG"})
+  fi
+  # || inspect: tolerate a race where another job created it between the check and create.
+  docker buildx create --name "$B" "${drv[@]}" >/dev/null \
+    || docker buildx inspect "$B" >/dev/null
+fi
+export BUILDX_BUILDER=$B
+docker buildx version
+
+# ---------------------------------------------------------------------------
+# Build context — an export of HEAD
+# ---------------------------------------------------------------------------
+# Bake builds .rocm-src, not the checkout: the content keys below hash HEAD, and BuildKit keys
+# COPY on file modes, which drift on shared agents. umask 022 gives every agent git's modes.
+[[ $(git rev-parse HEAD) == $(git rev-parse "${BUILDKITE_COMMIT:-HEAD}^{commit}") ]] \
+  || { echo "HEAD is not BUILDKITE_COMMIT ($BUILDKITE_COMMIT)" >&2; exit 1; }
+umask 022
+rm -rf .rocm-src && mkdir .rocm-src
+git archive HEAD | tar -x --no-same-permissions -C .rocm-src -f -
+export CONTEXT=.rocm-src
+
+# ---------------------------------------------------------------------------
+# Git tags — required by the version stage
+# ---------------------------------------------------------------------------
+# setuptools_scm needs a reachable v* tag. Agents clone without tags, partially
+# and with their objects in a host mirror (--reference) that the build context
+# cannot see; --force updates tags that moved (e.g. a re-tag after a release
+# candidate). The context has no .git: the version stage reads .git_archival.txt.
+git fetch -q --force --tags --filter=tree:0 origin
+git describe --tags --match 'v[0-9]*' >/dev/null 2>&1 \
+  || git fetch -q --unshallow --filter=tree:0 origin
+printf 'node: %s\nnode-date: %s\ndescribe-name: %s\n' "$(git rev-parse HEAD)" \
+  "$(git log -1 --format=%cI)" "$(git describe --tags --long --abbrev=40 --match 'v[0-9]*')" \
+  > .rocm-src/.git_archival.txt
+
+# ---------------------------------------------------------------------------
+# Helper functions
+# ---------------------------------------------------------------------------
+
+# digest <ref>
+#   Print the sha256 digest of <ref>. Returns 1 if the tag does not exist.
+#   The "Digest:" line of the plain output is parsed and validated (as the old CI
+#   scripts did): --format output and stray stderr are not safe to splice into a ref.
+#   Any other registry error, after retries, exits 75 (the steps retry it); in a $()
+#   subshell that becomes the substitution's status, which ensure() and set -e pass on.
+digest() {
+  local out rc d try
+  for try in 1 2 3; do
+    rc=0
+    out=$(docker buildx imagetools inspect "$1" 2>&1) || rc=$?
+    d=$(awk '$1 == "Digest:" { print $2; exit }' <<<"$out")
+    if ((rc == 0)) && [[ $d =~ ^sha256:[0-9a-f]{64}$ ]]; then echo "$d"; return 0; fi
+    if ((rc != 0)) && grep -qiE 'manifest unknown|no such manifest|name unknown|not found|does not exist' <<<"$out"; then
+      return 1
+    fi
+    ((try == 3)) || sleep 5
+  done
+  echo "registry error for $1 (exit $rc): $out" >&2
+  exit 75
+}
+
+# label <ref> <key>  — read one image label
+label() { docker buildx imagetools inspect "$1" \
+            --format "{{index .Image.Config.Labels \"$2\"}}"; }
+
+# bake <targets...>  — invoke docker buildx bake with the project HCL
+bake()   { docker buildx bake -f docker/docker-bake-rocm.hcl "$@"; }
+
+# meta <get|set> <key> [value]  — Buildkite build metadata
+meta()   { buildkite-agent meta-data "$@"; }
+
+# pushed <bake-metadata-file> <target>  — the digest that bake pushed for <target>
+pushed() {
+  local d
+  d=$(awk -F'"' -v t="$2" '/^  "/ { k = $2 } k == t && $2 == "containerimage.digest" { print $4 }' "$1")
+  [[ $d =~ ^sha256:[0-9a-f]{64}$ ]] && echo "$d"
+}
+
+# ---------------------------------------------------------------------------
+# Content keys
+# ---------------------------------------------------------------------------
+# Ks (stack key): hash of the stack Dockerfile + therock_wheels.py.
+#   A Ks cache miss triggers a full stack rebuild (hours).
+# Kr (runtime key): Ks + the portion of Dockerfile.rocm above the per-commit
+#   marker + all pip requirement files.
+#   A Kr cache miss triggers a runtime rebuild (minutes).
+
+_trw=$( [[ $STACK == rocm ]] && echo tools/vllm-rocm/therock_wheels.py || true )
+# The layer format (compression, and mtimes as built: see the bake file) is part of every image's
+# identity: an image left in an older format would hand later builds a second copy of the same layers.
+Ks=$( { echo "layers: zstd-3, mtimes kept"; git ls-tree HEAD -- "$SDF" ${_trw:+"$_trw"}; } | sha256sum | cut -c1-16)
+Kr=$(
+  {
+    echo "$Ks"
+    sed '/^# ==== per-commit stages/q' .rocm-src/docker/Dockerfile.rocm
+    git ls-tree HEAD -- \
+      requirements/common.txt \
+      requirements/rocm.txt \
+      requirements/test/rocm.txt \
+      rust-toolchain.toml \
+      tests/vllm_test_utils \
+      tools/install_torchcodec_rocm.sh
+  } | sha256sum | cut -c1-16
+)
+echo "build.sh $STACK $MODE: repo=$REPO trusted=$TRUSTED Ks=$Ks Kr=$Kr"
+
+# ---------------------------------------------------------------------------
+# ensure <bake-target> <tag-suffix> [<sibling-suffix>]
+# ---------------------------------------------------------------------------
+# Resolve a content-keyed image without building if it already exists.
+# Resolution order:
+#   1. Trusted registry (rocm/vllm-ci:<tag>) — safe for any CI build to reuse. Not for release:
+#      CI agents also run PR code, so release never ships an image it did not build itself.
+#   2. This repo's own tag ($REPO:$PFX<tag>) — PR builds and the release pipeline write here.
+#   3. Neither found: build via bake and push.
+# A <sibling-suffix> means the image only counts as present when the sibling
+# tag also exists (used by stack: must have both the image and -wheels); then
+# both pinned refs are printed, space-separated.
+ensure() {
+  local target=$1 tag=$2 sibling=${3:-} ref d ds refs rc mf
+  refs=("$REPO:$PFX$tag")
+  [[ $MODE == release* || $REPO == "$TRUSTED" ]] || refs=("$TRUSTED:$tag" "${refs[@]}")
+
+  for ref in "${refs[@]}"; do
+    if d=$(digest "$ref") && { [[ -z $sibling ]] || ds=$(digest "$ref$sibling"); }; then
+      echo "Reusing $ref@$d" >&2
+      echo "$ref@$d${sibling:+ $ref$sibling@$ds}"
+      return 0
+    else
+      rc=$?; ((rc == 1)) || return $rc   # a registry error is not a miss
+    fi
+  done
+
+  # Neither copy found — build and push.
+  ref=$REPO:$PFX$tag
+  echo "--- :docker: Building $target as $ref" >&2
+  # Explicit checks: $(ensure ...) runs in a subshell, where set -e does not apply.
+  mf=$(mktemp)
+  TAG=$ref bake --metadata-file "$mf" "$target" >&2 || { rm -f "$mf"; echo "building $ref failed" >&2; return 1; }
+  # The digests this bake pushed (base: stack-<stack>, stack-wheels), not the tags', which by now
+  # may hold a concurrent build's push of the same key.
+  if ! d=$(pushed "$mf" "${target/#base/stack-$STACK}") \
+      || { [[ -n $sibling ]] && ! ds=$(pushed "$mf" stack-wheels); }; then
+    rm -f "$mf"; echo "bake recorded no digest for $ref" >&2; return 1
+  fi
+  rm -f "$mf"
+  echo "$ref@$d${sibling:+ $ref$sibling@$ds}"
+}
+
+# load — read stack/runtime refs from Buildkite metadata (written by base mode).
+load() {
+  STACK_IMAGE=$(meta get "rocm-$STACK-stack")
+  STACK_WHEELS=$(meta get "rocm-$STACK-stack-wheels")
+  RUNTIME=$(meta get "rocm-$STACK-runtime")
+  TEST_DEPS=$(meta get "rocm-$STACK-test-deps")
+  export STACK_IMAGE STACK_WHEELS RUNTIME TEST_DEPS
+  for ref in "$STACK_IMAGE" "$STACK_WHEELS" "$RUNTIME" "$TEST_DEPS"; do
+    [[ $ref =~ @sha256:[0-9a-f]{64}$ ]] || { echo "not a pinned image ref in the metadata: '$ref'" >&2; exit 1; }
+  done
+}
+
+# ---------------------------------------------------------------------------
+# Mode dispatch
+# ---------------------------------------------------------------------------
+case $MODE in
+
+  keys)
+    # Print current content keys and registry status (useful for on-call triage).
+    for t in "stack-$STACK-$Ks" "runtime-$STACK-$Kr" "test-deps-$STACK-$Kr"; do
+      d=$(digest "$TRUSTED:$t") || { (($? == 1)) || exit 75; d=MISSING; }
+      echo "$t  $d"
+    done
+    ;;
+
+  base)
+    # Ensure the stack and runtime images exist; record their digests in metadata
+    # so the ci mode can pull them without re-resolving.
+    # The "base" bake group builds the stack and its -wheels sibling together.
+    STACK_IMAGE=$(ensure base "stack-$STACK-$Ks" -wheels)
+    STACK_WHEELS=${STACK_IMAGE#* } STACK_IMAGE=${STACK_IMAGE%% *}
+    export STACK_IMAGE
+    RUNTIME=$(ensure runtime "runtime-$STACK-$Kr")
+    export RUNTIME
+    TEST_DEPS=$(ensure test-deps "test-deps-$STACK-$Kr")
+
+    meta set "rocm-$STACK-stack"        "$STACK_IMAGE"
+    meta set "rocm-$STACK-runtime"      "$RUNTIME"
+    meta set "rocm-$STACK-test-deps"    "$TEST_DEPS"
+    meta set "rocm-$STACK-stack-wheels" "$STACK_WHEELS"
+    echo "stack=$STACK_IMAGE runtime=$RUNTIME test-deps=$TEST_DEPS stack-wheels=$STACK_WHEELS"
+    ;;
+
+  ci)
+    # Build and push the per-commit test image for GPU test steps.
+    load
+    export TAG=$REPO:build-$BUILDKITE_BUILD_ID
+    # AITER nightly (AITER_NIGHTLY=1 from the schedule): the test stage installs the newest
+    # AITER nightly that fits the image; an AITER_SPEC in the environment pins one instead.
+    rm -rf .rocm-wheel .rocm-kernel-map
+    bake ci
+    # For the python-only install test (run-amd-test.sh downloads it); not in the image.
+    buildkite-agent artifact upload ".rocm-wheel/vllm-*.whl"
+    # CI test selection (VLLM_KERNEL_SYMBOL_MAP=1, set by ci-infra): the kernel symbol map of this
+    # build's csrc, from the build cache. Optional: a missing map never fails the image build.
+    if [[ ${VLLM_KERNEL_SYMBOL_MAP:-0} == 1 ]]; then
+      if bake kernel-symbol-map && [[ -s .rocm-kernel-map/kernel_symbol_map.rocm.json.gz ]] \
+          && (cd .rocm-kernel-map && buildkite-agent artifact upload kernel_symbol_map.rocm.json.gz); then
+        :
+      else
+        echo "ROCm kernel symbol map export or upload failed; continuing without it" >&2
+      fi
+    fi
+    ;;
+
+  release|release-base)
+    # Build and push the release image + wheels into $REPO. The stack and runtime come from this
+    # repo only (built here once per key, from the same Dockerfiles and keys CI tests), never from
+    # rocm/vllm-ci or its cache.
+    STACK_IMAGE=$(ensure base "stack-$STACK-$Ks" -wheels)
+    STACK_WHEELS=${STACK_IMAGE#* } STACK_IMAGE=${STACK_IMAGE%% *}
+    export STACK_IMAGE   # before ensure runtime: a runtime built here must use this stack
+    RUNTIME=$(ensure runtime "runtime-$STACK-$Kr")
+    [[ $MODE == release ]] || exit 0
+    export STACK_WHEELS RUNTIME
+    export TAG=$REPO:$STACK-$BUILDKITE_COMMIT
+
+    # Resolve the vLLM version from the commit (setuptools_scm).
+    PUSH=false bake version
+    VLLM_VERSION=$(cat .rocm-version/vllm-version.txt)
+    export VLLM_VERSION
+
+    # Skip if this exact version was already released (handles rc → final re-tags).
+    if digest "$TAG" >/dev/null \
+        && [[ $(label "$TAG" org.opencontainers.image.version) == "$VLLM_VERSION" ]]; then
+      echo "Release image $TAG already exists with version $VLLM_VERSION — skipping." >&2
+      exit 0
+    fi
+
+    # Wheels first: their all-arch csrc compile is the long part, and vllm-openai reuses it.
+    bake wheels openai-smoke
+    # Alias the runtime as the -base tag (the base-nightly consumer reference).
+    docker buildx imagetools create -t "$TAG-base" "$RUNTIME"
+    # Push the vllm-openai tag last — its presence marks the set as complete.
+    bake vllm-openai
+    ;;
+
+  *)
+    echo "Unknown mode: $MODE" >&2
+    exit 2
+    ;;
+esac
diff --git a/.buildkite/scripts/rocm/publish.sh b/.buildkite/scripts/rocm/publish.sh
new file mode 100755
index 000000000000..3593d00ada39
--- /dev/null
+++ b/.buildkite/scripts/rocm/publish.sh
@@ -0,0 +1,208 @@
+#!/usr/bin/env bash
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+#
+# Promote ROCm release artifacts without rebuilding anything.
+# The release pipeline's build step (build.sh release) pushed the source images to
+# ECR; this script only copies digests and uploads wheels.
+#
+# Usage: publish.sh <stack> <mode>
+#
+# Modes:
+#   wheels   Extract the vLLM wheel from the -wheels image and upload to S3.
+#   images   Tag the release image and its runtime base for their variant in ECR.
+#            With NIGHTLY=1, also push the nightly tags to Docker Hub.
+#
+# Required environment variables:
+#   BUILDKITE_COMMIT   — the commit being released
+#   ECR            — ECR repository prefix (release-pipeline.yaml's ROCM_CI_REPO)
+#
+# Optional:
+#   HUB            — Docker Hub repo for nightly images (default: vllm/vllm-openai-rocm)
+#   NIGHTLY        — set to 1 to push nightly Hub tags (default: 0)
+#   DOCKERHUB_USERNAME — Hub username for cleanup script (default: vllmbot)
+set -euo pipefail
+
+# ---------------------------------------------------------------------------
+# Arguments
+# ---------------------------------------------------------------------------
+STACK=$1
+MODE=$2
+COMMIT=${BUILDKITE_COMMIT:?BUILDKITE_COMMIT must be set}
+
+ECR=${ECR:?ECR must be set (by release-pipeline.yaml)}
+SRC=$ECR:${STACK}-${COMMIT}   # build.sh release, in the same build
+HUB=${HUB:-vllm/vllm-openai-rocm}
+
+# ---------------------------------------------------------------------------
+# Helper functions
+# ---------------------------------------------------------------------------
+
+# retry <cmd...>  — registry lookups: these steps have no retry rules
+retry() { "$@" || { sleep 5; "$@"; } || { sleep 5; "$@"; }; }
+
+# pin <ref>  — resolve <ref> to a pinned ref@digest (resolves once, stable for the rest of the script).
+#   Parses the validated "Digest:" line, as build.sh does.
+pin() {
+  local d
+  d=$(retry docker buildx imagetools inspect "$1" | awk '$1 == "Digest:" { print $2; exit }')
+  [[ $d =~ ^sha256:[0-9a-f]{64}$ ]] || { echo "cannot resolve the digest of $1" >&2; return 1; }
+  echo "$1@$d"
+}
+
+# variant <ref>  — read the ai.vllm.rocm.variant image label
+variant() {
+  retry docker buildx imagetools inspect "$1" \
+    --format '{{index .Image.Config.Labels "ai.vllm.rocm.variant"}}'
+}
+
+# ---------------------------------------------------------------------------
+# Resolve source image refs (done once; images step uses them too)
+# ---------------------------------------------------------------------------
+# Logged in, the ECR lookups below are not subject to the anonymous rate limits.
+aws ecr-public get-login-password --region us-east-1 \
+  | docker login -u AWS --password-stdin "${ECR%%/*}"
+IMG=$(pin "$SRC")
+V=$(variant "$IMG")      # e.g. rocm100 — used in S3 paths, ECR tags and Hub tags
+[[ $V =~ ^rocm[0-9]+$ ]] || { echo "$IMG has no valid ai.vllm.rocm.variant label: '$V'" >&2; exit 1; }
+# For the steps after either mode: root-index generation and publish-release-images.sh.
+buildkite-agent meta-data set "rocm-$STACK-variant" "$V"
+
+# ---------------------------------------------------------------------------
+# Mode dispatch
+# ---------------------------------------------------------------------------
+case $MODE in
+
+  # -------------------------------------------------------------------------
+  wheels)
+  # -------------------------------------------------------------------------
+  # Extract the vLLM wheel from this commit's -wheels image and upload it
+  # to S3 so users can pip-install it directly.
+  # -------------------------------------------------------------------------
+
+    WHEELS_IMG=$(pin "$SRC-wheels")
+
+    # Export wheel files out of the scratch image into ./dist/
+    rm -rf dist
+    echo "FROM $WHEELS_IMG" | docker buildx build -q -o dist -
+
+    # Sanity: the wheel's variant suffix must match the stack label.
+    # (setup.py's rocm_version_to_variant produces e.g. rocm100 for 10.0.x)
+    actual_variant=$(basename dist/vllm-*.whl | grep -oE 'rocm[0-9]+' | tail -1)
+    [[ $actual_variant == "$V" ]] \
+      || { echo "Wheel variant $actual_variant != image variant $V" >&2; exit 1; }
+
+    bash .buildkite/scripts/upload-rocm-wheels.sh "$STACK" dist
+
+    # Buildkite annotation: pip install + direct S3 download commands.
+    WHL_URL="https://wheels.vllm.ai/rocm/$COMMIT/$V"
+    buildkite-agent annotate --style success --context "rocm-$STACK-wheels" << EOF
+## :package: ROCm \`$V\` wheels published
+
+| | |
+|---|---|
+| **Variant** | \`$V\` |
+| **Commit**  | \`$COMMIT\` |
+| **Branch**  | \`${BUILDKITE_BRANCH:-?}\` |
+
+### Install (this commit)
+\`\`\`bash
+pip install vllm --extra-index-url $WHL_URL
+\`\`\`
+
+### Install (latest nightly)
+\`\`\`bash
+pip install vllm --extra-index-url https://wheels.vllm.ai/rocm/nightly/$V
+\`\`\`
+
+### Download wheels directly
+\`\`\`bash
+aws s3 ls  s3://${S3_BUCKET:-vllm-wheels}/rocm/$COMMIT/$V-wheels/
+aws s3 cp  --recursive --exclude "*" --include "*.whl" s3://${S3_BUCKET:-vllm-wheels}/rocm/$COMMIT/$V-wheels/ .
+\`\`\`
+EOF
+    ;;
+
+  # -------------------------------------------------------------------------
+  images)
+  # -------------------------------------------------------------------------
+  # Tag the release image and its runtime base for their variant in ECR.
+  # With NIGHTLY=1, also tag them on Docker Hub and prune old nightly tags.
+  # -------------------------------------------------------------------------
+
+    RT=$(pin "$SRC-base")    # runtime base alias (published as base-nightly*)
+
+    # --- ECR -----------------------------------------------------------------
+    docker buildx imagetools create -t "$ECR:$COMMIT-$V"       "$IMG"
+    docker buildx imagetools create -t "$ECR:$COMMIT-$V-base"  "$RT"
+
+    ECR_IMG="$ECR:$COMMIT-$V"
+    ECR_BASE="$ECR:$COMMIT-$V-base"
+
+    # --- Non-nightly annotation ----------------------------------------------
+    if [[ ${NIGHTLY:-0} != 1 ]]; then
+      buildkite-agent annotate --style success --context "rocm-$STACK-images" << EOF
+## :docker: ROCm \`$V\` images published to ECR
+
+| | |
+|---|---|
+| **Variant** | \`$V\` |
+| **Commit**  | \`$COMMIT\` |
+| **ECR image** | \`$ECR_IMG\` |
+| **ECR base**  | \`$ECR_BASE\` |
+
+\`\`\`bash
+docker pull $ECR_IMG
+\`\`\`
+EOF
+      exit 0
+    fi
+
+    # --- Nightly: Docker Hub -------------------------------------------------
+    export DOCKERHUB_USERNAME=${DOCKERHUB_USERNAME:-vllmbot}
+
+    # The default stack (rocm) owns the unqualified :nightly / :base-nightly tags.
+    # All stacks own the variant-qualified :nightly-<variant> tags.
+    suffixes=("-$V")
+    [[ $STACK == rocm ]] && suffixes+=("")
+
+    hub_tags=""
+    for sfx in "${suffixes[@]}"; do
+      docker buildx imagetools create \
+        -t "$HUB:nightly$sfx" -t "$HUB:nightly$sfx-$COMMIT" "$IMG"
+      docker buildx imagetools create \
+        -t "$HUB:base-nightly$sfx" -t "$HUB:base-nightly$sfx-$COMMIT" "$RT"
+      hub_tags="$hub_tags \`$HUB:nightly$sfx\` \`$HUB:nightly$sfx-$COMMIT\`"
+    done
+
+    # Prune old nightly tags (keep last 14 builds per tag family).
+    clean=.buildkite/scripts/cleanup-nightly-builds.sh
+    bash "$clean" "nightly-$V-"      "$HUB"
+    bash "$clean" "base-nightly-$V-" "$HUB"
+    if [[ $STACK == rocm ]]; then
+      bash "$clean" "nightly-"      "$HUB" "nightly-rocm"
+      bash "$clean" "base-nightly-" "$HUB" "base-nightly-rocm"
+    fi
+
+    buildkite-agent annotate --style success --context "rocm-$STACK-images" << EOF
+## :docker: ROCm \`$V\` nightly images published
+
+| | |
+|---|---|
+| **Variant**   | \`$V\` |
+| **Commit**    | \`$COMMIT\` |
+| **ECR image** | \`$ECR_IMG\` |
+| **ECR base**  | \`$ECR_BASE\` |
+| **Docker Hub** |$hub_tags |
+
+\`\`\`bash
+docker pull $HUB:nightly-$V
+\`\`\`
+EOF
+    ;;
+
+  *)
+    echo "Unknown mode: $MODE" >&2
+    exit 2
+    ;;
+esac
diff --git a/.buildkite/scripts/rocm/refresh-base-image.sh b/.buildkite/scripts/rocm/refresh-base-image.sh
deleted file mode 100755
index 10e594577a87..000000000000
--- a/.buildkite/scripts/rocm/refresh-base-image.sh
+++ /dev/null
@@ -1,613 +0,0 @@
-#!/usr/bin/env bash
-# Select the content-addressed ROCm base for this checkout, building it only
-# when the exact content ref is absent.
-
-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}}"
-BUILDER_NAME="${ROCM_BASE_BUILDER_NAME:-vllm-rocm-base-builder}"
-DEFAULT_ROCM_BASE_METADATA_VERSION="2"
-DEFAULT_ROCM_BASE_CONTENT_FILES="${DOCKERFILE}"
-
-ROCM_BASE_LAYER_CACHE_REF=""
-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=()
-
-metadata_set() {
-    local key="$1"
-    local value="$2"
-
-    [[ -n "${value}" ]] || return 0
-    if command -v buildkite-agent >/dev/null 2>&1; then
-        buildkite-agent meta-data set "${key}" "${value}"
-    elif [[ "${BUILDKITE:-false}" == "true" ]]; then
-        echo "buildkite-agent is required to publish ROCm base metadata" >&2
-        return 1
-    fi
-}
-
-compute_content_hash() {
-    local path=""
-    local file=""
-
-    for path in "$@"; do
-        if [[ -d "${path}" ]]; then
-            while IFS= read -r -d '' file; do
-                printf 'file:%s\n' "${file}"
-                sha256sum "${file}"
-            done < <(find "${path}" -type f -print0 | sort -z)
-        elif [[ -f "${path}" ]]; then
-            printf 'file:%s\n' "${path}"
-            sha256sum "${path}"
-        else
-            printf 'missing:%s\n' "${path}"
-        fi
-    done | sha256sum | cut -d' ' -f1
-}
-
-clean_docker_tag() {
-    local input="$1"
-    echo "${input}" | sed 's/[^a-zA-Z0-9._-]/_/g' | cut -c1-128
-}
-
-tag_component() {
-    local input="$1"
-    local max_chars="${2:-24}"
-
-    clean_docker_tag "${input:-unknown}" | cut -c1-"${max_chars}"
-}
-
-normalize_repo_slug() {
-    local repo_slug="${1:-}"
-
-    repo_slug="${repo_slug%/}"
-    repo_slug="${repo_slug%.git}"
-    repo_slug="${repo_slug#https://github.com/}"
-    repo_slug="${repo_slug#http://github.com/}"
-    repo_slug="${repo_slug#ssh://git@github.com/}"
-    repo_slug="${repo_slug#git@github.com:}"
-    repo_slug="${repo_slug#github.com/}"
-    printf '%s\n' "${repo_slug}"
-}
-
-is_trusted_main_build() {
-    local actual_repo=""
-    local trusted_repo=""
-
-    [[ "${BUILDKITE:-false}" == "true" ]] || return 1
-    [[ "${BUILDKITE_PULL_REQUEST:-false}" == "false" ]] || return 1
-    [[ "${BUILDKITE_BRANCH:-}" == "${ROCM_BASE_STABLE_BRANCH:-main}" ]] || return 1
-    actual_repo=$(normalize_repo_slug "${BUILDKITE_REPO:-}")
-    trusted_repo=$(normalize_repo_slug \
-        "${ROCM_BASE_STABLE_REPO_SLUG:-vllm-project/vllm}")
-    [[ -n "${actual_repo}" && "${actual_repo}" == "${trusted_repo}" ]]
-}
-
-rocm_base_layer_cache_scope() {
-    local pull_request="${BUILDKITE_PULL_REQUEST:-false}"
-    local branch="${BUILDKITE_PULL_REQUEST_HEAD_BRANCH:-${BUILDKITE_BRANCH:-local}}"
-    local identity=""
-    local repo_slug=""
-
-    if is_trusted_main_build; then
-        printf 'main\n'
-        return 0
-    fi
-    if [[ "${pull_request}" != "false" && -n "${pull_request}" ]]; then
-        repo_slug=$(normalize_repo_slug "${BUILDKITE_REPO:-local}")
-        identity=$(printf '%s\n' "${repo_slug:-local}" | sha256sum | cut -c1-12)
-        printf 'pr-%s-%s\n' \
-            "$(tag_component "${pull_request}" 32)" "${identity}"
-        return 0
-    fi
-
-    identity=$(printf '%s\n%s\n' "${BUILDKITE_REPO:-local}" "${branch}" \
-        | sha256sum | cut -c1-12)
-    printf 'preview-%s-%s\n' "$(tag_component "${branch}" 24)" "${identity}"
-}
-
-configure_rocm_base_layer_cache() {
-    local scope=""
-
-    ROCM_BASE_CACHE_ARGS=()
-    if [[ "${ROCM_BASE_NO_CACHE:-0}" == "1" ]]; then
-        ROCM_BASE_CACHE_ARGS+=(--no-cache)
-        ROCM_BASE_LAYER_CACHE_REF="disabled"
-        return 0
-    fi
-
-    scope=$(rocm_base_layer_cache_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}"
-    )
-    if [[ "${ROCM_BASE_LAYER_CACHE_REF}" != \
-        "${ROCM_BASE_TRUSTED_LAYER_CACHE_REF}" ]]; then
-        ROCM_BASE_CACHE_ARGS+=(
-            --cache-from "type=registry,ref=${ROCM_BASE_TRUSTED_LAYER_CACHE_REF}"
-        )
-    fi
-    ROCM_BASE_CACHE_ARGS+=(
-        --cache-to \
-        "type=registry,ref=${ROCM_BASE_LAYER_CACHE_REF},mode=max,ignore-error=true"
-    )
-}
-
-extract_arg_default() {
-    local arg_name="$1"
-
-    sed -n -E "s/^[[:space:]]*[Aa][Rr][Gg][[:space:]]+${arg_name}=\"?([^\"[:space:]]+)\"?.*/\\1/p" \
-        "${DOCKERFILE}" | head -1
-}
-
-resolve_image_digest() {
-    local image_ref="$1"
-    local attempts="${ROCM_IMAGE_DIGEST_ATTEMPTS:-4}"
-    local delay_secs="${ROCM_IMAGE_DIGEST_RETRY_DELAY:-2}"
-    local attempt=0
-    local output=""
-    local digest=""
-    local status=0
-
-    if [[ "${image_ref}" =~ @(sha256:[0-9a-f]{64})$ ]]; then
-        printf '%s\n' "${BASH_REMATCH[1]}"
-        return 0
-    fi
-    if [[ ! "${attempts}" =~ ^[1-9][0-9]*$ \
-        || ! "${delay_secs}" =~ ^[0-9]+$ ]]; then
-        echo "Invalid image digest retry configuration" >&2
-        return 1
-    fi
-
-    for ((attempt = 1; attempt <= attempts; attempt++)); do
-        status=0
-        output=$(docker buildx imagetools inspect "${image_ref}" 2>&1) || status=$?
-        digest=$(awk '$1 == "Digest:" { print $2; exit }' <<< "${output}")
-        if ((status == 0)) && [[ "${digest}" =~ ^sha256:[0-9a-f]{64}$ ]]; then
-            printf '%s\n' "${digest}"
-            return 0
-        fi
-        if ((attempt < attempts)); then
-            printf \
-                'Image digest lookup failed for %s (%d/%d, status %d); retrying\n' \
-                "${image_ref}" "${attempt}" "${attempts}" "${status}" >&2
-            sleep "${delay_secs}"
-        fi
-    done
-
-    printf 'Failed to resolve digest for %s (status %d)\n%s\n' \
-        "${image_ref}" "${status}" "${output:-<no output>}" >&2
-    return 1
-}
-
-trusted_base_content_ref() {
-    local base_hash="$1"
-    local metadata_version="$2"
-
-    printf '%s:base-v%s-%s\n' "${BASE_REPO}" "${metadata_version}" "${base_hash}"
-}
-
-scoped_base_content_ref() {
-    local base_hash="$1"
-    local metadata_version="$2"
-
-    printf '%s:base-v%s-preview-%s\n' \
-        "${BASE_REPO}" "${metadata_version}" "${base_hash}"
-}
-
-find_matching_base_content_ref() {
-    local expected_hash="$1"
-    local expected_version="$2"
-    local image_ref=""
-    local digest=""
-    local immutable_ref=""
-    local label_values=""
-    local remote_hash=""
-    local remote_version=""
-    local attempts="${ROCM_IMAGE_DIGEST_ATTEMPTS:-4}"
-    local delay_secs="${ROCM_IMAGE_DIGEST_RETRY_DELAY:-2}"
-    local attempt=0
-    local inspect_status=0
-    shift 2
-
-    if [[ ! "${attempts}" =~ ^[1-9][0-9]*$ \
-        || ! "${delay_secs}" =~ ^[0-9]+$ ]]; then
-        echo "Invalid image identity retry configuration" >&2
-        return 1
-    fi
-
-    for image_ref in "$@"; do
-        if ! digest=$(resolve_image_digest "${image_ref}"); then
-            continue
-        fi
-        immutable_ref="${image_ref}@${digest}"
-        label_values=""
-        for ((attempt = 1; attempt <= attempts; attempt++)); do
-            inspect_status=0
-            label_values=$(docker buildx imagetools inspect "${immutable_ref}" \
-                --format '{{ index .Image.Config.Labels "vllm.rocm_base.content_hash" }} {{ index .Image.Config.Labels "vllm.rocm_base.metadata_version" }}' \
-                2>/dev/null) || inspect_status=$?
-            remote_hash=""
-            remote_version=""
-            if ((inspect_status == 0)) && [[ -n "${label_values}" ]]; then
-                read -r remote_hash remote_version _ <<< "${label_values}"
-                if [[ "${remote_hash}" == "${expected_hash}" \
-                    && "${remote_version}" == "${expected_version}" ]]; then
-                    printf '%s\n' "${immutable_ref}"
-                    return 0
-                fi
-            fi
-            if ((attempt < attempts)); then
-                printf \
-                    'ROCm base identity lookup incomplete or mismatched for %s (%d/%d); retrying\n' \
-                    "${immutable_ref}" "${attempt}" "${attempts}" >&2
-                sleep "${delay_secs}"
-            fi
-        done
-        if ((inspect_status == 0)) \
-            && [[ -n "${remote_hash}" && -n "${remote_version}" ]]; then
-            echo "ROCm base content ref has unexpected identity: ${immutable_ref}" >&2
-            echo "  expected hash/version: ${expected_hash} / ${expected_version}" >&2
-            echo "  found hash/version:    ${remote_hash} / ${remote_version}" >&2
-            # A content-addressed ref with the wrong identity is corrupt. The
-            # stable alias is allowed to describe older content and is only a
-            # migration/recovery candidate.
-            if [[ "${image_ref}" != "${BASE_REPO}:base" ]]; then
-                return 2
-            fi
-        fi
-        echo "Failed to read complete ROCm base identity labels: ${immutable_ref}" >&2
-    done
-    return 1
-}
-
-rocm_version_from_base_image() {
-    local base_image="$1"
-    local version=""
-
-    version="$(sed -n -E 's/.*:([0-9]+\.[0-9]+(\.[0-9]+)?)-.*/\1/p' <<<"${base_image}")"
-    tag_component "${version:-${base_image}}" 16
-}
-
-should_push_stable_tag() {
-    if [[ "${BUILDKITE_PULL_REQUEST:-false}" != "false" ]]; then
-        return 1
-    fi
-
-    if [[ "${ROCM_BASE_PUSH_STABLE_TAG:-}" == "0" ]]; then
-        return 1
-    fi
-
-    is_trusted_main_build
-}
-
-trusted_main_tip_matches_build() {
-    local branch="${ROCM_BASE_STABLE_BRANCH:-main}"
-    local build_commit="${BUILDKITE_COMMIT:-}"
-    local remote_tip=""
-
-    is_trusted_main_build || return 1
-    if [[ ! "${build_commit}" =~ ^[0-9a-fA-F]{40}$ ]]; then
-        echo "Skipping ROCm stable tag: Buildkite commit is missing or invalid" >&2
-        return 1
-    fi
-    remote_tip=$(git ls-remote --exit-code "${BUILDKITE_REPO}" \
-        "refs/heads/${branch}" 2>/dev/null | awk 'NR == 1 { print $1 }')
-    if [[ ! "${remote_tip}" =~ ^[0-9a-fA-F]{40}$ ]]; then
-        echo "Skipping ROCm stable tag: could not resolve remote ${branch} tip" >&2
-        return 1
-    fi
-    if [[ "${remote_tip,,}" != "${build_commit,,}" ]]; then
-        echo "Skipping ROCm stable tag: ${branch} changed after ${build_commit} (${remote_tip})" >&2
-        return 1
-    fi
-    return 0
-}
-
-promote_stable_base_tag() {
-    local source_ref="$1"
-    local stable_tag="$2"
-
-    ROCM_BASE_STABLE_TAG_UPDATED=0
-    should_push_stable_tag && trusted_main_tip_matches_build || return 0
-
-    docker buildx imagetools create --prefer-index=false \
-        -t "${stable_tag}" "${source_ref}"
-    ROCM_BASE_STABLE_TAG_UPDATED=1
-}
-
-setup_builder() {
-    echo "--- :buildkite: Setting up buildx builder for ROCm base"
-    if docker buildx inspect "${BUILDER_NAME}" >/dev/null 2>&1; then
-        docker buildx use "${BUILDER_NAME}"
-    else
-        docker buildx create --name "${BUILDER_NAME}" --driver docker-container --use
-    fi
-    docker buildx inspect --bootstrap
-}
-
-compute_base_content_hash() {
-    local use_sccache="$1"
-    local base_image_digest="$2"
-    local content_files="${ROCM_BASE_CONTENT_FILES:-${DEFAULT_ROCM_BASE_CONTENT_FILES}}"
-    local sccache_arg=""
-    local sccache_value=""
-    local -a content_paths=()
-
-    read -r -a content_paths <<< "${content_files}"
-
-    {
-        printf 'content-files-hash:%s\n' "$(compute_content_hash "${content_paths[@]}")"
-        printf 'dockerfile:%s\n' "${DOCKERFILE}"
-        printf 'base-image-digest:%s\n' "${base_image_digest}"
-        printf 'use-sccache:%s\n' "${use_sccache}"
-        if [[ "${use_sccache}" == "1" ]]; then
-            # These values can change the installed binary or final image
-            # configuration. SCCACHE_ENDPOINT is transport-only and therefore
-            # deliberately excluded from image identity.
-            for sccache_arg in \
-                SCCACHE_DOWNLOAD_URL \
-                SCCACHE_BUCKET_NAME \
-                SCCACHE_REGION_NAME \
-                SCCACHE_S3_NO_CREDENTIALS; do
-                sccache_value="${!sccache_arg:-}"
-                if [[ -z "${sccache_value}" ]]; then
-                    sccache_value=$(extract_arg_default "${sccache_arg}")
-                fi
-                printf 'arg:%s=%s\n' \
-                    "${sccache_arg}" "${sccache_value:-<empty>}"
-            done
-        fi
-    } | sha256sum | cut -d' ' -f1
-}
-
-build_base_image() {
-    local use_sccache="${ROCM_BASE_USE_SCCACHE:-${USE_SCCACHE:-0}}"
-    local base_hash=""
-    local base_image_arg=""
-    local base_image_digest=""
-    local pinned_base_image=""
-    local rocm_version=""
-    local triton_arg=""
-    local pytorch_arg=""
-    local pytorch_vision_arg=""
-    local pytorch_audio_arg=""
-    local fa_arg=""
-    local aiter_arg=""
-    local mori_arg=""
-    local python_version_arg=""
-    local pytorch_rocm_arch_arg=""
-    local dependency_summary=""
-    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=""
-    local selected_digest=""
-    local stable_digest=""
-    local reuse_status=0
-    local content_files="${ROCM_BASE_CONTENT_FILES:-${DEFAULT_ROCM_BASE_CONTENT_FILES}}"
-    local content_files_hash=""
-    local metadata_version="${ROCM_BASE_METADATA_VERSION:-${DEFAULT_ROCM_BASE_METADATA_VERSION}}"
-    local -a sccache_args=()
-    local -a content_paths=()
-
-    if [[ ! -f "${DOCKERFILE}" ]]; then
-        echo "Error: ROCm base Dockerfile not found: ${DOCKERFILE}" >&2
-        exit 1
-    fi
-    if [[ ! "${metadata_version}" =~ ^[A-Za-z0-9][A-Za-z0-9_.-]{0,15}$ ]]; then
-        echo "Invalid ROCm base metadata version: ${metadata_version}" >&2
-        return 1
-    fi
-    if [[ "${use_sccache}" != "0" && "${use_sccache}" != "1" ]]; then
-        echo "ROCm base USE_SCCACHE must be 0 or 1: ${use_sccache}" >&2
-        return 1
-    fi
-
-    base_image_arg="$(extract_arg_default BASE_IMAGE)"
-    if ! base_image_digest=$(resolve_image_digest "${base_image_arg}"); then
-        echo "Failed to resolve ROCm base input image: ${base_image_arg}" >&2
-        return 1
-    fi
-    pinned_base_image="${base_image_arg%@*}@${base_image_digest}"
-    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="$(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)"
-    python_version_arg="$(extract_arg_default PYTHON_VERSION)"
-    pytorch_rocm_arch_arg="$(extract_arg_default PYTORCH_ROCM_ARCH)"
-    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
-        scoped_content_tag="${trusted_content_tag}"
-        writable_content_tag="${trusted_content_tag}"
-    else
-        writable_content_tag="${scoped_content_tag}"
-    fi
-
-    configure_rocm_base_layer_cache
-
-    if [[ "${use_sccache}" == "1" ]]; then
-        for env_name in \
-            SCCACHE_DOWNLOAD_URL \
-            SCCACHE_ENDPOINT \
-            SCCACHE_BUCKET_NAME \
-            SCCACHE_REGION_NAME \
-            SCCACHE_S3_NO_CREDENTIALS; do
-            if [[ -n "${!env_name:-}" ]]; then
-                sccache_args+=(--build-arg "${env_name}=${!env_name}")
-            fi
-        done
-    fi
-
-    echo "--- :docker: Preparing ROCm base image"
-    echo "Dockerfile: ${DOCKERFILE}"
-    echo "Trusted content tag: ${trusted_content_tag}"
-    echo "Writable content tag: ${writable_content_tag}"
-    echo "Stable tag: ${stable_tag} ($(should_push_stable_tag && echo enabled || echo disabled))"
-    echo "Content hash: ${base_hash}"
-    echo "Dependency summary: ${dependency_summary}"
-    echo "USE_SCCACHE: ${use_sccache}"
-    echo "BuildKit layer cache: ${ROCM_BASE_LAYER_CACHE_REF}"
-
-    if [[ "${ROCM_BASE_REFRESH_FORCE:-0}" != "1" \
-        && "${ROCM_BASE_NO_CACHE:-0}" != "1" ]]; then
-        reuse_status=0
-        if [[ "${trusted_content_tag}" == "${scoped_content_tag}" ]]; then
-            immutable_ref=$(find_matching_base_content_ref \
-                "${base_hash}" "${metadata_version}" \
-                "${trusted_content_tag}" "${preview_content_tag}" \
-                "${stable_tag}") || reuse_status=$?
-        else
-            immutable_ref=$(find_matching_base_content_ref \
-                "${base_hash}" "${metadata_version}" \
-                "${trusted_content_tag}" "${scoped_content_tag}" \
-                "${stable_tag}") \
-                || reuse_status=$?
-        fi
-    else
-        reuse_status=1
-        echo "Forced/no-cache ROCm base build; bypassing content image reuse"
-    fi
-
-    if ((reuse_status > 1)); then
-        return "${reuse_status}"
-    fi
-
-    if [[ ${reuse_status} -eq 0 && -n "${immutable_ref}" ]]; then
-        echo "Reusing ROCm base image with matching content: ${immutable_ref}"
-        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}"
-            if ! immutable_ref=$(find_matching_base_content_ref \
-                "${base_hash}" "${metadata_version}" \
-                "${writable_content_tag}"); then
-                echo "Backfilled ROCm base content tag failed identity validation" >&2
-                return 1
-            fi
-        fi
-    else
-        build_ref="${writable_content_tag}"
-        echo "No reusable ROCm base content image found; building ${build_ref}"
-        setup_builder
-        docker buildx build \
-            "${ROCM_BASE_CACHE_ARGS[@]}" \
-            --pull \
-            --provenance=false \
-            --progress "${BUILDKIT_PROGRESS:-plain}" \
-            --file "${DOCKERFILE}" \
-            --build-arg "BASE_IMAGE=${pinned_base_image}" \
-            --build-arg "USE_SCCACHE=${use_sccache}" \
-            "${sccache_args[@]}" \
-            --label "org.opencontainers.image.source=https://github.com/vllm-project/vllm" \
-            --label "org.opencontainers.image.vendor=vLLM" \
-            --label "org.opencontainers.image.title=vLLM ROCm base" \
-            --label "vllm.rocm_base.metadata_version=${metadata_version}" \
-            --label "vllm.rocm_base.content_hash=${base_hash}" \
-            --label "vllm.rocm_base.content_files_hash=${content_files_hash}" \
-            --label "vllm.rocm_base.dockerfile=${DOCKERFILE}" \
-            --label "vllm.rocm_base.dependency_summary=${dependency_summary}" \
-            --label "vllm.rocm_base.base_image=${pinned_base_image}" \
-            --label "vllm.rocm_base.base_image_digest=${base_image_digest}" \
-            --label "vllm.rocm_base.dependency.rocm=${rocm_version}" \
-            --label "vllm.rocm_base.dependency.python=${python_version_arg}" \
-            --label "vllm.rocm_base.dependency.pytorch=${pytorch_arg}" \
-            --label "vllm.rocm_base.dependency.torchvision=${pytorch_vision_arg}" \
-            --label "vllm.rocm_base.dependency.torchaudio=${pytorch_audio_arg}" \
-            --label "vllm.rocm_base.dependency.triton=${triton_arg}" \
-            --label "vllm.rocm_base.dependency.flash_attention=${fa_arg}" \
-            --label "vllm.rocm_base.dependency.aiter=${aiter_arg}" \
-            --label "vllm.rocm_base.dependency.mori=${mori_arg}" \
-            --label "vllm.rocm_base.pytorch_rocm_arch=${pytorch_rocm_arch_arg}" \
-            -t "${build_ref}" \
-            --push \
-            .
-        if ! immutable_ref=$(find_matching_base_content_ref \
-            "${base_hash}" "${metadata_version}" "${build_ref}"); then
-            echo "Published ROCm base image failed identity validation: ${build_ref}" >&2
-            return 1
-        fi
-    fi
-
-    selected_digest="${immutable_ref##*@}"
-    metadata_set "rocm-base-image" "${immutable_ref}"
-    metadata_set "rocm-base-push-stable-tag" "0"
-    stable_digest=$(resolve_image_digest "${stable_tag}" 2>/dev/null || true)
-    if [[ -n "${stable_digest}" && "${stable_digest}" == "${selected_digest}" ]]; then
-        echo "--- :white_check_mark: Existing ROCm base is current"
-        echo "Stable ${stable_tag} already resolves to ${selected_digest}"
-        return 0
-    fi
-
-    promote_stable_base_tag "${immutable_ref}" "${stable_tag}"
-
-    metadata_set "rocm-base-push-stable-tag" "${ROCM_BASE_STABLE_TAG_UPDATED}"
-    metadata_set "rocm-base-refresh" "1"
-
-    echo "--- :white_check_mark: ROCm base image published"
-    echo "Use BASE_IMAGE=${immutable_ref} for downstream ROCm CI builds"
-}
-
-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
-        local stable_ref="${BASE_REPO}:base"
-        local stable_digest=""
-
-        if ! stable_digest=$(resolve_image_digest "${stable_ref}"); then
-            echo "Could not pin explicitly skipped ROCm base: ${stable_ref}" >&2
-            return 1
-        fi
-        metadata_set "rocm-base-image" "${stable_ref}@${stable_digest}"
-        metadata_set "rocm-base-push-stable-tag" "0"
-        echo "ROCM_BASE_REFRESH_SKIP=1; pinned current stable base at ${stable_digest}"
-        return 0
-    fi
-    build_base_image
-}
-
-if [[ "${BASH_SOURCE[0]}" == "$0" ]]; then
-    main "$@"
-fi
diff --git a/.buildkite/scripts/rocm/select_aiter_nightly_wheel.py b/.buildkite/scripts/rocm/select_aiter_nightly_wheel.py
index af3de54f6020..3559afe404d4 100755
--- a/.buildkite/scripts/rocm/select_aiter_nightly_wheel.py
+++ b/.buildkite/scripts/rocm/select_aiter_nightly_wheel.py
@@ -3,9 +3,9 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Select the AITER nightly wheel for a vLLM ROCm image.
 
-aiter-nightly-overlay.sh runs this inside the stock ci_base, so it reads that
-image's ROCm, torch and Python. It prints the wheel's URL, then the Buildkite
-annotation naming it.
+docker/Dockerfile.rocm runs this in the test image's AITER nightly step, so it
+reads that image's ROCm, torch and Python. It prints the wheel's URL, then the
+Buildkite annotation naming it.
 """
 
 import argparse
@@ -21,17 +21,21 @@
 
 import regex as re
 
-INDEX_URL = "https://rocm.frameworks-nightlies.amd.com/whl-multi-arch/amd-aiter/"
+INDEX_URL = (
+    "https://rocm.frameworks-nightlies.amd.com/whl-multi-arch-staging/amd-aiter/"
+)
 AITER_COMMIT_URL = "https://github.com/ROCm/aiter/commit/"
 STALE_AFTER = datetime.timedelta(days=1)
 
 # Example:
-#   amd_aiter-0.1.25+rocm7.2.3.torch2.12.ae984f5.d20261008-cp312-cp312-linux_x86_64.whl
-# is AITER main at commit ae984f5, built 2026-10-08, for ROCm 7.2.3, torch 2.12
-# and Python 3.12. 0.1.25 is a guess at the next release, not a real one.
+#   amd_aiter-0.1.28+rocm10.0.0.torch2.13.f684e21.d20261029
+#     -cp312-cp312-manylinux_2_28_x86_64.whl
+# is AITER main at commit f684e21, built 2026-10-29, for ROCm 10.0.0, torch 2.13
+# and Python 3.12, and runs on glibc 2.28 or newer. 0.1.28 is a guess at the
+# next release, not a real one.
 WHEEL_NAME = re.compile(
     r"^amd_aiter-(?P<base_version>[^+-]+)\+(?P<local>[^-]+)"
-    r"-(?P<python>cp\d+)-(?P=python)-linux_x86_64\.whl$"
+    r"-(?P<python>cp\d+)-(?P=python)-(?:manylinux_\d+_\d+|linux)_x86_64\.whl$"
 )
 ROCM_FIELD = re.compile(r"(?:^|\.)rocm(\d+(?:\.\d+)*)(?:a\d+)?(?=\.|$)")
 TORCH_FIELD = re.compile(r"(?:^|\.)torch(\d+\.\d+)(?=\.|$)")
@@ -235,7 +239,9 @@ def _today() -> datetime.date:
     return datetime.datetime.now(datetime.UTC).date()
 
 
-# Not 1 or 2, which Buildkite retries; see aiter-nightly-overlay.sh.
+# Its own status for "no wheel for this image". Dockerfile.rocm fails the image
+# build on any failure here; bake reports that as exit 1, which the image-build
+# step retries once.
 EXIT_NO_WHEEL = 10
 
 
diff --git a/.buildkite/scripts/rocm/smoke-test-image.sh b/.buildkite/scripts/rocm/smoke-test-image.sh
deleted file mode 100755
index b155a0897d32..000000000000
--- a/.buildkite/scripts/rocm/smoke-test-image.sh
+++ /dev/null
@@ -1,84 +0,0 @@
-#!/usr/bin/env bash
-# Fast structural smoke test for the full ROCm CI image.
-
-set -euo pipefail
-
-run_smoke_checks() {
-    local required_dir=""
-
-    for required_dir in \
-        /vllm-workspace \
-        /vllm-workspace/tests \
-        /vllm-workspace/src/vllm; do
-        if [[ ! -d "${required_dir}" ]]; then
-            echo "Missing directory: ${required_dir}" >&2
-            return 1
-        fi
-    done
-    if [[ ! -x /vllm-workspace/src/vllm/vllm-rs ]]; then
-        echo "Missing executable: /vllm-workspace/src/vllm/vllm-rs" >&2
-        return 1
-    fi
-
-    command -v python3
-    command -v uv
-    command -v pytest
-
-    if ! command -v amd-smi >/dev/null 2>&1 \
-        && ! command -v rocminfo >/dev/null 2>&1; then
-        echo "No ROCm CLI found in image" >&2
-        return 1
-    fi
-
-    PYTHONDONTWRITEBYTECODE=1 python3 - <<'PY'
-import torch
-import vllm
-
-print(torch.__version__)
-print(vllm.__version__)
-PY
-
-    echo "AMD image smoke OK"
-}
-
-if [[ "${1:-}" == "--inside" ]]; then
-    run_smoke_checks
-    exit
-fi
-if (($#)); then
-    echo "Usage: $0 [--inside]" >&2
-    exit 2
-fi
-
-if [[ "${ROCM_CI_ARTIFACT_ONLY:-0}" == "1" ]]; then
-    base_refreshed=""
-    if command -v buildkite-agent >/dev/null 2>&1; then
-        base_refreshed="$(buildkite-agent meta-data get rocm-base-refresh 2>/dev/null || true)"
-    fi
-    if [[ "${base_refreshed}" != "1" ]]; then
-        echo "ROCM_CI_ARTIFACT_ONLY=1; no full image was built, skipping smoke test"
-        exit 0
-    fi
-fi
-
-smoke_marker="./build/rocm-smoke-export/vllm-smoke-ok"
-expected_smoke_id="${BUILDKITE_BUILD_ID:-local}"
-if [[ -f "${smoke_marker}" \
-    && ( -z "${VLLM_CI_SMOKE_IMAGE:-}" \
-        || "${VLLM_CI_SMOKE_IMAGE}" == "${IMAGE_TAG:-}" ) ]]; then
-    actual_smoke_id="$(< "${smoke_marker}")"
-    if [[ "${actual_smoke_id}" != "${expected_smoke_id}" ]]; then
-        echo "ROCm smoke marker belongs to ${actual_smoke_id}, not ${expected_smoke_id}" \
-            >&2
-        exit 1
-    fi
-    echo "AMD image smoke OK (verified inside BuildKit)"
-    rm -f -- "${smoke_marker}"
-    rmdir -- "$(dirname "${smoke_marker}")" 2>/dev/null || true
-    exit
-fi
-
-image_ref="${VLLM_CI_SMOKE_IMAGE:-${IMAGE_TAG:-rocm/vllm-ci:${BUILDKITE_COMMIT:?BUILDKITE_COMMIT is required}}}"
-
-docker run --rm -i --network=none --entrypoint /bin/bash "${image_ref}" \
-    -s -- --inside < "${BASH_SOURCE[0]}"
diff --git a/.buildkite/scripts/rocm/stack.sh b/.buildkite/scripts/rocm/stack.sh
deleted file mode 100644
index 14bfeb48a017..000000000000
--- a/.buildkite/scripts/rocm/stack.sh
+++ /dev/null
@@ -1,78 +0,0 @@
-#!/usr/bin/env bash
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-#
-# Derive everything the release pipeline needs about a ROCm stack from its base
-# Dockerfile, so the stack's image is the single input to every ROCm script.
-#
-# Usage: source .buildkite/scripts/rocm/stack.sh <base-dockerfile>
-#   e.g. docker/Dockerfile.rocm_base (ROCm 10.0) or docker/Dockerfile.rocm_72_base
-#
-# Exports:
-#   ROCM_STACK_BASE_DOCKERFILE  the base Dockerfile
-#   ROCM_STACK_DOCKERFILE       the matching vLLM Dockerfile (base name minus _base)
-#   ROCM_STACK_THEROCK          1 if the ROCm SDK comes from TheRock wheels
-#   ROCM_STACK_VERSION          ROCm version, e.g. 10.0.0 or 7.2.3
-#   ROCM_STACK_VARIANT          wheel variant and Docker Hub tag flavor from major.minor,
-#                               same rule as setup.py (rocm100, rocm72)
-#   ROCM_STACK_IMAGE_KEY        hash of the base Dockerfile (base image cache key)
-#   ROCM_STACK_WHEEL_KEY        hash of the base Dockerfile plus files it COPYs from
-#                               the repo (base wheel cache key)
-#   ROCM_STACK_ECR_BASE         ECR tag of the cached base image
-#   ROCM_STACK_ECR_IMAGE        ECR tag of this commit's release image
-#   ROCM_STACK_BASE_WHEELS_DIR  local dir for base wheels (also the artifact path)
-#   ROCM_STACK_VLLM_WHEEL_DIR   local dir for the vLLM wheel (also the artifact path)
-
-rocm_stack_init() {
-    local base="$1"
-    if [[ ! -f "$base" ]]; then
-        echo "ROCm base Dockerfile not found: $base" >&2
-        return 2
-    fi
-    ROCM_STACK_BASE_DOCKERFILE="$base"
-    ROCM_STACK_DOCKERFILE="${base%_base}"
-
-    local sdk
-    sdk=$(sed -nE 's/^ARG ROCM_SDK_VERSION=([0-9.]+).*/\1/p' "$base" | head -1)
-    if [[ -n "$sdk" ]]; then
-        ROCM_STACK_THEROCK=1
-        ROCM_STACK_VERSION="$sdk"
-    else
-        ROCM_STACK_THEROCK=0
-        # BASE_IMAGE format: rocm/dev-ubuntu-22.04:7.2.3-complete -> 7.2.3
-        ROCM_STACK_VERSION=$(sed -nE 's/^ARG BASE_IMAGE=.*:([0-9]+(\.[0-9]+)+).*/\1/p' "$base" | head -1)
-    fi
-    if [[ -z "$ROCM_STACK_VERSION" ]]; then
-        echo "Could not determine the ROCm version from $base" >&2
-        return 2
-    fi
-
-    local major minor
-    IFS=. read -r major minor _ <<< "$ROCM_STACK_VERSION"
-    ROCM_STACK_VARIANT="rocm${major}${minor}"
-
-    ROCM_STACK_IMAGE_KEY=$(sha256sum "$base" | cut -c1-16)
-    local inputs
-    inputs=$(grep -E '^COPY ' "$base" | grep -v -- '--from' | awk '{print $2}' || true)
-    if [[ -z "$inputs" ]]; then
-        ROCM_STACK_WHEEL_KEY="$ROCM_STACK_IMAGE_KEY"
-    else
-        # shellcheck disable=SC2086
-        ROCM_STACK_WHEEL_KEY=$(sha256sum "$base" $inputs | sha256sum | cut -c1-16)
-    fi
-
-    local ecr="public.ecr.aws/q9t5s3a7/vllm-release-repo"
-    ROCM_STACK_ECR_BASE="${ecr}:${ROCM_STACK_IMAGE_KEY}-${ROCM_STACK_VARIANT}-base"
-    ROCM_STACK_ECR_IMAGE="${ecr}:${BUILDKITE_COMMIT:-unknown}-${ROCM_STACK_VARIANT}"
-    ROCM_STACK_BASE_WHEELS_DIR="artifacts/${ROCM_STACK_VARIANT}-base-wheels"
-    ROCM_STACK_VLLM_WHEEL_DIR="artifacts/${ROCM_STACK_VARIANT}-vllm-wheel"
-
-    export ROCM_STACK_BASE_DOCKERFILE ROCM_STACK_DOCKERFILE ROCM_STACK_THEROCK \
-        ROCM_STACK_VERSION ROCM_STACK_VARIANT ROCM_STACK_IMAGE_KEY \
-        ROCM_STACK_WHEEL_KEY ROCM_STACK_ECR_BASE ROCM_STACK_ECR_IMAGE \
-        ROCM_STACK_BASE_WHEELS_DIR ROCM_STACK_VLLM_WHEEL_DIR
-}
-
-if [[ $# -gt 0 ]]; then
-    rocm_stack_init "$1"
-fi
diff --git a/.buildkite/scripts/upload-rocm-wheels.sh b/.buildkite/scripts/upload-rocm-wheels.sh
index 5cf32bcfb358..d303f16eae0e 100755
--- a/.buildkite/scripts/upload-rocm-wheels.sh
+++ b/.buildkite/scripts/upload-rocm-wheels.sh
@@ -4,8 +4,8 @@
 #
 # Upload ROCm wheels to S3 with proper index generation
 #
-# Usage: upload-rocm-wheels.sh <base-dockerfile>
-#   base-dockerfile  the ROCm stack, e.g. docker/Dockerfile.rocm_base (see rocm/stack.sh)
+# Usage: upload-rocm-wheels.sh <stack> <dir>
+#   stack  rocm (TheRock);  dir  the stack's wheels plus the vLLM wheel
 #
 # Required environment variables:
 #   AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY (or IAM role)
@@ -23,11 +23,11 @@ set -ex
 
 # ======== Configuration ========
 BUCKET="${S3_BUCKET:-vllm-wheels}"
-# shellcheck source=.buildkite/scripts/rocm/stack.sh
-source "$(dirname "${BASH_SOURCE[0]}")/rocm/stack.sh" "${1:?Usage: $0 <base-dockerfile>}"
+STACK="${1:?Usage: $0 <stack> <dir>}"
+WHEEL_DIR="${2:?Usage: $0 <stack> <dir>}"
+VARIANT=$(basename "$WHEEL_DIR"/vllm-*.whl | grep -oE 'rocm[0-9]+' | tail -1)
 ROCM_SUBPATH="rocm/${BUILDKITE_COMMIT}"
-WHEEL_DIRS="$ROCM_STACK_BASE_WHEELS_DIR $ROCM_STACK_VLLM_WHEEL_DIR"
-WHEEL_SUBPATH="$ROCM_SUBPATH/${ROCM_STACK_VARIANT}-wheels"
+WHEEL_SUBPATH="$ROCM_SUBPATH/${VARIANT:?no vLLM wheel in $WHEEL_DIR}-wheels"
 S3_WHEEL_PREFIX="s3://$BUCKET/$WHEEL_SUBPATH/"
 INDICES_OUTPUT_DIR="rocm-indices"
 
@@ -61,9 +61,7 @@ source .buildkite/scripts/lib/manylinux.sh
 # Collect all wheels
 rm -rf all-rocm-wheels "$INDICES_OUTPUT_DIR"
 mkdir -p all-rocm-wheels
-for dir in $WHEEL_DIRS; do
-    cp "$dir"/*.whl all-rocm-wheels/ 2>/dev/null || true
-done
+cp "$WHEEL_DIR"/*.whl all-rocm-wheels/ 2>/dev/null || true
 
 WHEEL_COUNT=$(find all-rocm-wheels -maxdepth 1 -name '*.whl' 2>/dev/null | wc -l)
 echo "Total wheels to upload: $WHEEL_COUNT"
@@ -119,10 +117,10 @@ sed -i 's/import pybase64 as base64/import base64/g' tools/vllm-rocm/therock_whe
 
 INDEX_ARGS=(--version "$ROCM_SUBPATH" --wheel-dir "${WHEEL_SUBPATH#rocm/}")
 # TheRock's SDK, device kernels and torchvision/torchaudio are linked, not hosted
-if [[ "$ROCM_STACK_THEROCK" == "1" ]]; then
+if [[ "$STACK" == "rocm" ]]; then
     $PYTHON tools/vllm-rocm/therock_wheels.py external-links \
-        --dockerfile "$ROCM_STACK_BASE_DOCKERFILE" \
-        --torch-wheel "$(ls "$ROCM_STACK_BASE_WHEELS_DIR"/torch-*.whl)" > therock-external-links.txt
+        --dockerfile docker/Dockerfile.rocm_base \
+        --torch-wheel "$(ls "$WHEEL_DIR"/torch-*.whl)" > therock-external-links.txt
     cat therock-external-links.txt
     INDEX_ARGS+=(--external-links therock-external-links.txt)
 fi
@@ -191,11 +189,11 @@ echo "Wheels available at:"
 echo "  s3://$BUCKET/$WHEEL_SUBPATH/"
 echo ""
 echo "Install command (by commit):"
-echo "  pip install vllm --extra-index-url https://${BUCKET}.s3.amazonaws.com/$ROCM_SUBPATH/$ROCM_STACK_VARIANT/"
+echo "  pip install vllm --extra-index-url https://${BUCKET}.s3.amazonaws.com/$ROCM_SUBPATH/$VARIANT/"
 echo ""
 if [[ "${NIGHTLY:-0}" == "1" ]]; then
     echo "Install command (nightly):"
-    echo "  pip install vllm --extra-index-url https://${BUCKET}.s3.amazonaws.com/rocm/nightly/$ROCM_STACK_VARIANT/"
+    echo "  pip install vllm --extra-index-url https://${BUCKET}.s3.amazonaws.com/rocm/nightly/$VARIANT/"
 fi
 echo ""
 echo "Wheel count: $WHEEL_COUNT"
diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml
index cfa5a1575888..173d8643eedd 100644
--- a/.buildkite/test-amd.yaml
+++ b/.buildkite/test-amd.yaml
@@ -14,12 +14,9 @@
 # command(str): the single command to run for tests. incompatible with commands.
 # commands(list): the list of commands to run for the test. incompatible with command.
 # mirror_hardwares(list): selector tags for AMD pipelines that should include the test.
-# dind(bool): when false, run the job directly in the AMD Kubernetes pod.
-#     When true or omitted, use the legacy Docker-in-Docker path.
+# dind(bool): must be false: the job runs directly in the AMD Kubernetes pod (no Docker-in-Docker).
 # num_gpus(int): override the number of GPUs for the test. defaults to 1 GPU. currently supports 2,4.
-# num_nodes(int): whether to simulate multi-node setup by launching multiple containers on one host,
-#     in this case, commands must be specified. the first command runs on the first host, the second
-#     command runs on the second host.
+# num_nodes(int): not supported on AMD (num_nodes >= 2 fails).
 # timeout_in_minutes(int): sets a timeout for the step in minutes. if not specified, uses the default timeout.
 # parallelism(int): number of parallel jobs to run for this step. enables test sharding using $$BUILDKITE_PARALLEL_JOB
 #     and $$BUILDKITE_PARALLEL_JOB_COUNT environment variables.
@@ -121,18 +118,26 @@ steps:
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
   - .buildkite/scripts/docker-build-metadata-args.sh
-  - .buildkite/scripts/ci-bake-rocm.sh
   - .buildkite/release-pipeline.yaml
   - docker/Dockerfile
   - docker/Dockerfile.cpu
   - docker/Dockerfile.rocm
   - docker/Dockerfile.rocm_base
-  - docker/ci-rocm.hcl
   - docker/docker-bake.hcl
   - docker/docker-bake-rocm.hcl
+  - .buildkite/scripts/rocm/build.sh
+  - .buildkite/scripts/rocm/publish.sh
+  - tools/vllm-rocm/therock_wheels.py
   - tests/tools/test_docker_build_metadata_args.py
+  - tests/tools/test_rocm_build.py
+  - tests/tools/test_amd_ci_fast_exit.py
+  - .buildkite/scripts/hardware_ci/amd_ci_fast_exit.py
+  - tests/tools/test_kernel_symbol_map.py
+  - tools/ci/kernel_symbol_map.py
+  - tools/ci/rocm_kernel_symbol_map.py
   commands:
-  - pytest -v -s tools/test_docker_build_metadata_args.py
+  - pytest -v -s tools/test_docker_build_metadata_args.py tools/test_rocm_build.py tools/test_amd_ci_fast_exit.py
+  - pytest -v -s tools/test_kernel_symbol_map.py
 
 # Disabled until torch-abi-audit transitive dependencies are locked.
 # - label: ":amd: (MI300) Torch Stable ABI Audit" # TBD
@@ -274,29 +279,27 @@ steps:
 
 #------------------------------------------------------------  mi250 · v1  -------------------------------------------------------------#
 
-# Disabled until the pytest plugins installation is locked.
-# - label: ":amd: (MI250) Batch Invariance" # TBD
-#   timeout_in_minutes: 60
-#   mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
-#   dind: false
-#   agent_pool: mi250_1
-#   optional: true
-#   working_dir: "/vllm-workspace/tests"
-#   source_file_dependencies:
-#   - vllm/v1/attention
-#   - vllm/model_executor/layers
-#   - vllm/model_executor/kernels/linear
-#   - vllm/model_executor/determinism
-#   - tests/v1/determinism/
-#   - vllm/_aiter_ops.py
-#   - vllm/platforms/rocm.py
-#   commands:
-#   - export VLLM_GPU_MEMORY_UTILIZATION=0.8
-#   - export VLLM_TEST_CLEAN_GPU_MEMORY=1
-#   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-#   - pip install pytest-timeout pytest-forked
-#   - pytest -v -s v1/determinism/test_batch_invariance.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
+- label: ":amd: (MI250) Batch Invariance" # TBD
+  timeout_in_minutes: 60
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
+  dind: false
+  agent_pool: mi250_1
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/v1/attention
+  - vllm/model_executor/layers
+  - vllm/model_executor/kernels/linear
+  - vllm/model_executor/determinism
+  - tests/v1/determinism/
+  - vllm/_aiter_ops.py
+  - vllm/platforms/rocm.py
+  commands:
+  - export VLLM_GPU_MEMORY_UTILIZATION=0.8
+  - export VLLM_TEST_CLEAN_GPU_MEMORY=1
+  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  - pytest -v -s v1/determinism/test_batch_invariance.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
 
 - label: ":amd: (MI250) CUDAGraph" # TBD
   timeout_in_minutes: 55
@@ -473,7 +476,9 @@ steps:
     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)' ;;
+    VLLM_ROCM_USE_AITER=1 pytest -v -s models/test_hyv4_rocm.py -m 'distributed(num_gpus=2)' &&
+    pytest -v -s models/test_mimo_v2_attention_backend.py &&
+    pytest -v -s models/multimodal/test_mimo_v2_omni.py ;;
     2)
     pytest models/multimodal -v -s -m 'distributed(num_gpus=2)'
     --ignore models/multimodal/generation/test_whisper.py &&
@@ -887,25 +892,23 @@ steps:
   - pytest -v -s -x lora/test_mixtral.py
   - pytest -v -s models/test_vision.py -m 'distributed(num_gpus=4)'
 
-# Disabled until the KV connectors installation is locked.
-# - label: ":amd: (MI300) Elastic EP Scaling" # TBD
-#   timeout_in_minutes: 55
-#   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/
-#   - vllm/engine/
-#   - vllm/v1/executor/
-#   - vllm/compilation/
-#   - tests/distributed/
-#   - vllm/platforms/rocm.py
-#   commands:
-#   - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
-#   - pytest -v -s distributed/test_elastic_ep.py
+- label: ":amd: (MI300) Elastic EP Scaling" # TBD
+  timeout_in_minutes: 55
+  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/
+  - vllm/engine/
+  - vllm/v1/executor/
+  - vllm/compilation/
+  - tests/distributed/
+  - vllm/platforms/rocm.py
+  commands:
+  - pytest -v -s distributed/test_elastic_ep.py
 
 - label: ":amd: (MI300) RayExecutorV2" # TBD
   timeout_in_minutes: 75
@@ -1568,8 +1571,6 @@ steps:
   - 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 '(not core_model) and (not hybrid_model)'
 
 #----------------------------------------------------  mi300 · models / multimodal  ----------------------------------------------------#
@@ -2020,33 +2021,31 @@ steps:
 
 #------------------------------------------------------------  mi300 · v1  -------------------------------------------------------------#
 
-# Disabled until the pytest plugins installation is locked.
-# - label: ":amd: (MI300) Batch Invariance" # TBD
-#   timeout_in_minutes: 60
-#   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
-#   dind: false
-#   agent_pool: mi300_1
-#   optional: true
-#   working_dir: "/vllm-workspace/tests"
-#   source_file_dependencies:
-#   - vllm/v1/attention
-#   - vllm/model_executor/layers
-#   - vllm/lora/ops/triton_ops
-#   - vllm/model_executor/kernels/linear
-#   - vllm/model_executor/determinism
-#   - 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
-#   - pytest -v -s v1/determinism/test_lora_shrink_batch_invariant.py
-#   - pytest -v -s v1/determinism/test_cutlass_batch_invariance.py
-#   - pytest -v -s v1/determinism/test_matmul_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 TRITON_ATTN
+- label: ":amd: (MI300) Batch Invariance" # TBD
+  timeout_in_minutes: 60
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+  dind: false
+  agent_pool: mi300_1
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/v1/attention
+  - vllm/model_executor/layers
+  - vllm/lora/ops/triton_ops
+  - vllm/model_executor/kernels/linear
+  - vllm/model_executor/determinism
+  - tests/v1/determinism/
+  - vllm/_aiter_ops.py
+  - vllm/platforms/rocm.py
+  commands:
+  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  - pytest -v -s v1/determinism/test_batch_invariance.py
+  - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py
+  - pytest -v -s v1/determinism/test_lora_shrink_batch_invariant.py
+  - pytest -v -s v1/determinism/test_cutlass_batch_invariance.py
+  - pytest -v -s v1/determinism/test_matmul_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 TRITON_ATTN
 
 - label: ":amd: (MI300) Batch Invariance (Distributed)" # TBD
   timeout_in_minutes: 45
@@ -2070,6 +2069,7 @@ steps:
   - VLLM_TEST_TP_SIZE=4 VLLM_TEST_DISABLE_CUSTOM_AR=1 pytest -v -s v1/determinism/test_batch_invariance.py -k bs1_vs_bsN
   - VLLM_TP_SIZE=4 VLLM_NEEDLE_TRIALS=2 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_ATTN
   - VLLM_TEST_MODEL=allenai/OLMoE-1B-7B-0924 VLLM_TEST_SERVER_ARGS="--data-parallel-size 4 --enable-expert-parallel --max-model-len 4096" pytest -v -s v1/determinism/test_online_batch_invariance.py
+  - VLLM_TEST_MODEL=ibm-granite/granite-3.0-1b-a400m-base VLLM_TEST_SERVER_ARGS="--data-parallel-size 4 --enable-expert-parallel --max-model-len 4096" pytest -v -s v1/determinism/test_online_batch_invariance.py
 
 - label: ":amd: (MI300) V1 Spec Decode Shard %N" # TBD
   timeout_in_minutes: 50
@@ -2159,7 +2159,9 @@ steps:
   - vllm/model_executor/model_loader/
   - vllm/model_executor/models/extract_hidden_states.py
   - vllm/model_executor/models/llama.py
+  - vllm/models/qwen3_5/
   - vllm/model_executor/models/qwen3_5.py
+  - vllm/model_executor/models/qwen3_5_mtp.py
   - vllm/model_executor/models/qwen3_next.py
   - vllm/model_executor/models/registry.py
   - vllm/transformers_utils/configs/extract_hidden_states.py
@@ -2204,39 +2206,33 @@ steps:
 
   - VLLM_ROCM_USE_AITER=1 pytest -v -s v1/attention/test_sparse_mla_backends.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
 
-# Disabled until the tracing packages installation is locked.
-# - label: ":amd: (MI300) Metrics, Tracing" # TBD
-#   timeout_in_minutes: 30
-#   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
-#   dind: false
-#   agent_pool: mi300_2
-#   optional: true
-#   num_gpus: 2
-#   working_dir: "/vllm-workspace/tests"
-#   source_file_dependencies:
-#   - 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' \
-#       'opentelemetry-api>=1.26.0' \
-#       'opentelemetry-exporter-otlp>=1.26.0' \
-#       'opentelemetry-semantic-conventions-ai>=0.4.1'"
-#   - pytest -v -s v1/tracing
-#   - pytest -v -s tracing
+- label: ":amd: (MI300) Metrics, Tracing" # TBD
+  timeout_in_minutes: 30
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+  dind: false
+  agent_pool: mi300_2
+  optional: true
+  num_gpus: 2
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - 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:
+  - pytest -v -s v1/tracing
+  - pytest -v -s tracing
 
 - label: ":amd: (MI300) V1 E2E" # TBD
   timeout_in_minutes: 30
@@ -2325,40 +2321,36 @@ steps:
   - export BASELINE_FILE='/tmp/mooncake-ec-e2e/baseline.json'
   - bash tests/v1/ec_connector/integration/run_epd_mooncake_ec_full_pipeline.sh
 
-# Disabled until the KV connectors installation is locked.
-# - label: ":amd: (MI300) NixlConnector PD edge case" # TBD
-#   timeout_in_minutes: 65
-#   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/kv_transfer/kv_connector/v1/nixl/
-#   - vllm/v1/core/sched/
-#   - tests/v1/kv_connector/nixl_integration/
-#   - vllm/platforms/rocm.py
-#   commands:
-#   - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
-#   - PREFILL_GPU_ID=0 DECODE_GPU_ID=1 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_edge_case_test.sh
+- label: ":amd: (MI300) NixlConnector PD edge case" # TBD
+  timeout_in_minutes: 65
+  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/kv_transfer/kv_connector/v1/nixl/
+  - vllm/v1/core/sched/
+  - tests/v1/kv_connector/nixl_integration/
+  - vllm/platforms/rocm.py
+  commands:
+  - PREFILL_GPU_ID=0 DECODE_GPU_ID=1 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_edge_case_test.sh
 
-# Disabled until the KV connectors installation is locked.
-# - label: ":amd: (MI300) CrossLayer KV layout Distributed NixlConnector PD accuracy" # TBD
-#   timeout_in_minutes: 60
-#   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
-#   - CROSS_LAYERS_BLOCKS=True ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
+- label: ":amd: (MI300) CrossLayer KV layout Distributed NixlConnector PD accuracy" # TBD
+  timeout_in_minutes: 60
+  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:
+  - 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: 75
@@ -2384,133 +2376,119 @@ steps:
   - pytest -v -s v1/engine/test_engine_core_client.py::test_kv_cache_events_dp
   - pytest -v -s distributed/test_utils.py
 
-# Disabled until the KV connectors installation is locked.
-# - 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) 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:
+  - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
 
-# Disabled until the KV connectors installation is locked.
-# - label: ":amd: (MI300) Push NixlConnector PP prefill PD accuracy" # TBD
-#   timeout_in_minutes: 55
-#   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/
-#   - tests/v1/kv_connector/nixl_push_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_push_integration/config_sweep_accuracy_test.sh
+- label: ":amd: (MI300) Push NixlConnector PP prefill PD accuracy" # TBD
+  timeout_in_minutes: 55
+  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/
+  - tests/v1/kv_connector/nixl_push_integration/
+  - vllm/platforms/rocm.py
+  commands:
+  - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_push_integration/config_sweep_accuracy_test.sh
 
-# Disabled until the KV connectors installation is locked.
-# - 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) 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:
+  - DP_EP=1 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
 
-# Disabled until the KV connectors installation is locked.
-# - label: ":amd: (MI300) Hybrid SSM NixlConnector PD accuracy" # TBD
-#   timeout_in_minutes: 80
-#   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
-#   - HYBRID_SSM=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: 80
+  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:
+  - HYBRID_SSM=1 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
 
-# Disabled until the KV connectors installation is locked.
-# - label: ":amd: (MI300) Hybrid SSM NixlConnector PD prefix cache" # TBD
-#   timeout_in_minutes: 25
-#   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/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: ":amd: (MI300) Hybrid SSM NixlConnector PD prefix cache" # TBD
+  timeout_in_minutes: 25
+  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/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:
+  - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh
 
-# Disabled until the KV connectors installation is locked.
-# - label: ":amd: (MI300) MultiConnector (Nixl+Offloading) PD accuracy" # TBD
-#   timeout_in_minutes: 30
-#   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/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_accuracy_test.sh
+- label: ":amd: (MI300) MultiConnector (Nixl+Offloading) PD accuracy" # TBD
+  timeout_in_minutes: 30
+  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/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:
+  - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh
 
-# Disabled until the KV connectors installation is locked.
-# - label: ":amd: (MI300) MultiConnector (Nixl+Offloading) PD edge cases" # TBD
-#   timeout_in_minutes: 25
-#   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/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
+- label: ":amd: (MI300) MultiConnector (Nixl+Offloading) PD edge cases" # TBD
+  timeout_in_minutes: 25
+  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/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:
+  - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh
 
 #########################################################################################################################################
 #                                                                                                                                       #
@@ -2928,132 +2906,117 @@ steps:
   - 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]'
 
-# Disabled until the KV connectors installation is locked.
-# - label: ":amd: (MI355) NixlConnector PD + Spec Decode acceptance" # TBD
-#   timeout_in_minutes: 45
-#   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:
-#   - 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:
-#   - 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) NixlConnector PD + Spec Decode acceptance" # TBD
+  timeout_in_minutes: 45
+  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:
+  - 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:
+  - KV_CACHE_MEMORY_BYTES=8G ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_spec_decode_test.sh
 
-# Disabled until the KV connectors installation is locked.
-# - label: ":amd: (MI355) Distributed NixlConnector PD accuracy" # TBD
-#   timeout_in_minutes: 60
-#   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/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: (MI355) Distributed NixlConnector PD accuracy" # TBD
+  timeout_in_minutes: 60
+  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/distributed/kv_transfer/kv_connector/v1/nixl/
+  - tests/v1/kv_connector/nixl_integration/
+  - vllm/platforms/rocm.py
+  commands:
+  - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
 
-# Disabled until the KV connectors installation is locked.
-# - label: ":amd: (MI355) Distributed AITER NixlConnector PD accuracy"
-#   timeout_in_minutes: 60
-#   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/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
-#   - pytest -v -s v1/kv_connector/nixl_integration/test_accuracy_rocm.py
+- label: ":amd: (MI355) Distributed AITER NixlConnector PD accuracy"
+  timeout_in_minutes: 60
+  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/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
+  commands:
+  - pytest -v -s v1/kv_connector/nixl_integration/test_accuracy_rocm.py
 
-# Disabled until the KV connectors installation is locked.
-# - label: ":amd: (MI355) Distributed MooncakeConnector HIP PD accuracy"
-#   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/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
-#   - pytest -v -s v1/kv_connector/mooncake_integration/test_accuracy_rocm.py
+- label: ":amd: (MI355) Distributed MooncakeConnector HIP PD accuracy"
+  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/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
+  commands:
+  - pytest -v -s v1/kv_connector/mooncake_integration/test_accuracy_rocm.py
 
-# Disabled until the KV connectors installation is locked.
-# - 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_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) DP EP Distributed NixlConnector PD accuracy" # TBD
+  timeout_in_minutes: 40
+  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/distributed/kv_transfer/kv_connector/v1/nixl/
+  - tests/v1/kv_connector/nixl_integration/
+  - vllm/platforms/rocm.py
+  commands:
+  - DP_EP=1 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
 
-# Disabled until the KV connectors installation is locked.
-# - 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,--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) 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/
+  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,--max-model-len 32768'
+  - bash v1/kv_connector/nixl_integration/run_accuracy_test.sh
 
 - label: ":amd: (MI355) Distributed Features" # TBD
   timeout_in_minutes: 90
@@ -3167,45 +3130,42 @@ steps:
     not +quant_fp8 and not +rms_norm and DeepSeek-Coder and
     ROCM_AITER_MLA'
 
-# Disabled until the KV connectors installation is locked.
-# - label: ":amd: (MI355) DSv4-Flash Disaggregated DP EP" # TBD
-#   timeout_in_minutes: 85
-#   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
-#   dind: false
-#   agent_pool: mi355_8
-#   num_gpus: 8
-#   optional: true
-#   working_dir: "/vllm-workspace/tests"
-#   source_file_dependencies:
-#   - requirements/kv_connectors_rocm.txt
-#   - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
-#   - vllm/model_executor/layers/fused_moe/
-#   - vllm/model_executor/layers/quantization/
-#   - vllm/models/deepseek_v4/
-#   - vllm/v1/attention/
-#   - tests/v1/kv_connector/nixl_integration/
-#   - vllm/_aiter_ops.py
-#   - vllm/platforms/rocm.py
-#   commands:
-#   - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
-#   # A mirror does not inherit the parent env, so it is repeated here in
-#   # the parent's order. The layouts must stay empty: unset, the script
-#   # defaults to HND, which the V4 indexer rejects.
-#   - >-
-#     VLLM_ROCM_USE_AITER=1
-#     ENABLE_HMA_FLAG=1
-#     DP_EP=1
-#     GPU_MEMORY_UTILIZATION=0.85
-#     PREFILLER_TP_SIZE=4
-#     DECODER_TP_SIZE=4
-#     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
-#     PREFILLER_KV_LAYOUT=
-#     DECODER_KV_LAYOUT=
-#     bash v1/kv_connector/nixl_integration/run_accuracy_test.sh
+- label: ":amd: (MI355) DSv4-Flash Disaggregated DP EP" # TBD
+  timeout_in_minutes: 85
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
+  agent_pool: mi355_8
+  num_gpus: 8
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
+  - vllm/model_executor/layers/fused_moe/
+  - vllm/model_executor/layers/quantization/
+  - vllm/models/deepseek_v4/
+  - vllm/v1/attention/
+  - tests/v1/kv_connector/nixl_integration/
+  - vllm/_aiter_ops.py
+  - vllm/platforms/rocm.py
+  commands:
+  # A mirror does not inherit the parent env, so it is repeated here in
+  # the parent's order. The layouts must stay empty: unset, the script
+  # defaults to HND, which the V4 indexer rejects.
+  - >-
+    VLLM_ROCM_USE_AITER=1
+    ENABLE_HMA_FLAG=1
+    DP_EP=1
+    GPU_MEMORY_UTILIZATION=0.85
+    PREFILLER_TP_SIZE=4
+    DECODER_TP_SIZE=4
+    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
+    PREFILLER_KV_LAYOUT=
+    DECODER_KV_LAYOUT=
+    bash v1/kv_connector/nixl_integration/run_accuracy_test.sh
 
 - label: ":amd: (MI355) Distributed Torchrun + Examples" # TBD
   timeout_in_minutes: 55
@@ -3621,6 +3581,7 @@ steps:
   optional: true
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
+  - vllm/models/qwen3_5/
   - vllm/model_executor/models/qwen3_5.py
   - vllm/model_executor/models/qwen3_5_mtp.py
   - vllm/model_executor/models/qwen2.py
@@ -4198,6 +4159,8 @@ steps:
   - vllm/v1/kv_offload/
   - vllm/v1/simple_kv_offload/
   - tests/evals/gsm8k/test_gsm8k_offloading.py
+  - vllm/distributed/kv_transfer/kv_connector/v1/umbp/
+  - tests/v1/kv_connector/umbp_test_utils.py
   - vllm/platforms/rocm.py
   commands:
   - pytest -s -v evals/gsm8k/test_gsm8k_offloading.py -k "nemotron-h-8b or gemma-4-e4b-it"
@@ -4410,6 +4373,7 @@ steps:
   - vllm/platforms/rocm.py
   commands:
   - pytest -v -s
+      kernels/attention/test_minimax_m3_nvfp4_index.py
       kernels/test_apply_repetition_penalties.py
       kernels/test_awq_int4_to_int8.py
       kernels/test_bf16_skinny_gemm.py
@@ -4597,11 +4561,11 @@ steps:
     --ignore=kernels/attention/test_rocm_aiter_mla_op_registration.py
     --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
 
-- label: ":amd: (MI355 DPX) MoE Kernels Shard %N" # TBD
+- label: ":amd: (MI355) MoE Kernels Shard %N" # TBD
   timeout_in_minutes: 55
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
   dind: false
-  agent_pool: mi355_dpx
+  agent_pool: mi355_1
   parallelism: 5
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
@@ -4876,7 +4840,7 @@ steps:
   - tests/models/test_hyv4_rocm.py
   - tests/models/test_vision.py
   - tests/models/test_deepseek_v4_vl_rocm.py
-  - tests/models/test_deepseek_v4_vl_stream_weights.py
+  - tests/models/test_deepseek_v4_vl.py
   commands:
   - >-
     pytest -v -s
@@ -4884,7 +4848,7 @@ steps:
     models/transformers/test_backend.py
     models/test_registry.py
     models/test_deepseek_v4_vl_rocm.py
-    models/test_deepseek_v4_vl_stream_weights.py
+    models/test_deepseek_v4_vl.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
 
@@ -4908,6 +4872,34 @@ steps:
   # PLE and QSA import conflicting HC registrations in a shared process.
   - pytest -v -s models/qwen4_exp/test_ple.py
 
+- label: ":amd: (MI355 DPX) Step5"
+  timeout_in_minutes: 30
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
+  agent_pool: mi355_dpx
+  num_gpus: 1
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/models/step5/
+  - tests/models/step5/
+  - vllm/model_executor/models/step3p7.py
+  - vllm/model_executor/models/step3_vl.py
+  - vllm/model_executor/models/step_vl.py
+  - vllm/config/compilation.py
+  - vllm/model_executor/layers/indexer_topk.py
+  - vllm/transformers_utils/configs/step5.py
+  - vllm/model_executor/warmup/step5_dsa_warmup.py
+  - vllm/v1/attention/backends/rocm_aiter_unified_attn.py
+  - vllm/triton_utils/
+  - vllm/platforms/rocm.py
+  - tests/v1/attention/utils.py
+  - tests/models/test_initialization.py
+  - tests/models/registry.py
+  commands:
+  - pytest -v -s models/step5
+  - pytest -v -s models/test_initialization.py -k Step5ForConditionalGeneration
+
 - label: ":amd: (MI355 DPX) GLM5Next"
   timeout_in_minutes: 35
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
@@ -5008,11 +5000,6 @@ steps:
   - 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 DPX) Granite Language Model Compatibility"
@@ -5028,11 +5015,6 @@ steps:
   - "!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 DPX) Language Models (Extended Generation)" # TBD
@@ -5046,8 +5028,6 @@ steps:
   - 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 DPX) Language Models (Extended Pooling) Shard %N"
@@ -5072,6 +5052,7 @@ steps:
   optional: true
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
+  - vllm/models/qwen3_5/
   - vllm/model_executor/models/qwen3_5.py
   - vllm/model_executor/models/qwen3_5_mtp.py
   - vllm/model_executor/models/qwen2.py
@@ -5276,37 +5257,35 @@ steps:
 
 #------------------------------------------------------------  mi355 · v1  -------------------------------------------------------------#
 
-# Disabled until the pytest plugins installation is locked.
-# - label: ":amd: (MI355) Batch Invariance" # TBD
-#   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/v1/attention
-#   - vllm/model_executor/layers
-#   - vllm/lora/ops/triton_ops
-#   - vllm/model_executor/kernels/linear
-#   - vllm/model_executor/determinism
-#   - 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
-#   - pytest -v -s v1/determinism/test_lora_shrink_batch_invariant.py
-#   - pytest -v -s v1/determinism/test_batch_invariance_vlm.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 TRITON_ATTN
-#   - pytest -v -s v1/determinism/test_nvfp4_batch_invariant.py
-#   - pytest -v -s v1/determinism/test_matmul_batch_invariant.py
-#   - pytest -v -s v1/determinism/test_cutlass_batch_invariance.py
-#   - pytest -v -s v1/determinism/test_online_batch_invariance.py
-#   - pytest -v -s v1/determinism/test_mx_linear_batch_invariant.py
+- label: ":amd: (MI355) Batch Invariance" # TBD
+  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/v1/attention
+  - vllm/model_executor/layers
+  - vllm/lora/ops/triton_ops
+  - vllm/model_executor/kernels/linear
+  - vllm/model_executor/determinism
+  - tests/v1/determinism/
+  - vllm/_aiter_ops.py
+  - vllm/platforms/rocm.py
+  commands:
+  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  - pytest -v -s v1/determinism/test_batch_invariance.py
+  - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py
+  - pytest -v -s v1/determinism/test_lora_shrink_batch_invariant.py
+  - pytest -v -s v1/determinism/test_batch_invariance_vlm.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 TRITON_ATTN
+  - pytest -v -s v1/determinism/test_nvfp4_batch_invariant.py
+  - pytest -v -s v1/determinism/test_matmul_batch_invariant.py
+  - pytest -v -s v1/determinism/test_cutlass_batch_invariance.py
+  - pytest -v -s v1/determinism/test_online_batch_invariance.py
+  - pytest -v -s v1/determinism/test_mx_linear_batch_invariant.py
 
 - label: ":amd: (MI355) Spec Decode AL DFlash2 Nightly"
   timeout_in_minutes: 60
@@ -5328,7 +5307,9 @@ steps:
   - vllm/v1/attention/
   - vllm/model_executor/models/qwen3_dflash.py
   - vllm/model_executor/models/qwen3_dflash2.py
+  - vllm/models/qwen3_5/
   - vllm/model_executor/models/qwen3_5.py
+  - vllm/model_executor/models/qwen3_5_mtp.py
   - vllm/model_executor/model_loader/
   - vllm/model_executor/layers/quantization/
   - vllm/model_executor/kernels/linear/
@@ -5637,6 +5618,9 @@ steps:
   - vllm/v1/kv_cache_interface.py
   - tests/v1/kv_connector/unit/test_umbp_
   - tests/v1/kv_connector/umbp_test_utils.py
+  - tests/v1/kv_connector/unit/utils.py
+  - tests/v1/kv_connector/conftest.py
+  - docker/Dockerfile.rocm_base
   commands:
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
   # Fail loudly instead of skipping every test if the image lacks mori.
@@ -5837,6 +5821,26 @@ steps:
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
   - pytest -v -s -m 'not distributed' tests/v1/kv_connector/extract_hidden_states_integration
 
+- label: ":amd: (MI355 DPX) UMBP Connector Integration"
+  timeout_in_minutes: 30
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
+  agent_pool: mi355_dpx
+  num_gpus: 1
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/distributed/kv_transfer/kv_connector/v1/umbp/
+  - vllm/distributed/kv_transfer/kv_connector/v1/base.py
+  - tests/v1/kv_connector/umbp_integration/
+  - tests/v1/kv_connector/umbp_test_utils.py
+  - tests/v1/kv_connector/conftest.py
+  - tests/v1/kv_connector/rocm_pd_accuracy_utils.py
+  - docker/Dockerfile.rocm_base
+  commands:
+  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  - pytest -v -s -rs v1/kv_connector/umbp_integration
+
 - label: ":amd: (MI355 DPX) V1 Attention Shard %N" # TBD
   timeout_in_minutes: 125
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
@@ -6217,3 +6221,4 @@ steps:
   - vllm/envs.py
   commands:
   - pytest -v -s tests/jit_monitor/test_no_runtime_jit_rocm.py --timeout=300
+  - pytest -v -s tests/jit_monitor/test_flydsl_contract.py
diff --git a/.buildkite/test_areas/attention.yaml b/.buildkite/test_areas/attention.yaml
index e5cbb528df01..41866227f342 100644
--- a/.buildkite/test_areas/attention.yaml
+++ b/.buildkite/test_areas/attention.yaml
@@ -82,6 +82,7 @@ steps:
 
 - label: ":nvidia: (DGX Spark) [Nightly] FlashInfer NVFP4 XQA"
   key: flashinfer-nvfp4-xqa-dgx-spark
+  automatic_shard: true
   timeout_in_minutes: 30
   device: dgx-spark
   optional: true
diff --git a/.buildkite/test_areas/basic_correctness.yaml b/.buildkite/test_areas/basic_correctness.yaml
index f294301a7b1f..7cf48f6467fa 100644
--- a/.buildkite/test_areas/basic_correctness.yaml
+++ b/.buildkite/test_areas/basic_correctness.yaml
@@ -26,6 +26,7 @@ steps:
 
 - label: ":nvidia: (H200 MIG 18GB) Basic Correctness CuMem"
   key: basic-correctness-cumem
+  automatic_shard: true
   timeout_in_minutes: 68
   device: h200_18gb
   source_file_dependencies: *basic-correctness-deps
@@ -34,6 +35,7 @@ steps:
   - pytest -v -s basic_correctness/memory/cumem
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355 DPX) Basic Correctness CuMem"
       dind: false
       device: mi355_dpx
@@ -43,6 +45,7 @@ steps:
 
 - label: ":nvidia: (H200 MIG 18GB) Basic Correctness Sleep Mode"
   key: basic-correctness-sleep-mode
+  automatic_shard: true
   timeout_in_minutes: 68
   device: h200_18gb
   source_file_dependencies: *basic-correctness-deps
@@ -51,6 +54,7 @@ steps:
   - pytest -v -s basic_correctness/memory/sleep_mode
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI300) Basic Correctness Sleep Mode"
       dind: false
       device: mi300_1
@@ -60,6 +64,7 @@ steps:
 
 - label: ":nvidia: (H200 MIG 18GB) Basic Correctness CPU Offload"
   key: basic-correctness-cpu-offload
+  automatic_shard: true
   timeout_in_minutes: 68
   device: h200_18gb
   source_file_dependencies: *basic-correctness-deps
@@ -68,6 +73,7 @@ steps:
   - pytest -v -s basic_correctness/cpu_offload
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355 DPX) Basic Correctness CPU Offload"
       dind: false
       device: mi355_dpx
@@ -77,6 +83,7 @@ steps:
 
 - label: ":nvidia: (H200 MIG 18GB) Basic Correctness Prefetch Offload"
   key: basic-correctness-prefetch-offload
+  automatic_shard: true
   timeout_in_minutes: 68
   device: h200_18gb
   source_file_dependencies: *basic-correctness-deps
@@ -85,6 +92,7 @@ steps:
   - pytest -v -s basic_correctness/prefetch_offload
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355 DPX) Basic Correctness Prefetch Offload"
       dind: false
       device: mi355_dpx
diff --git a/.buildkite/test_areas/benchmarks.yaml b/.buildkite/test_areas/benchmarks.yaml
index bd381080094f..1eca056cbb9c 100644
--- a/.buildkite/test_areas/benchmarks.yaml
+++ b/.buildkite/test_areas/benchmarks.yaml
@@ -16,6 +16,7 @@ steps:
   - pytest -v -s benchmarks/
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355 DPX) Benchmarks CLI"
       dind: false
       device: mi355_dpx
diff --git a/.buildkite/test_areas/compile.yaml b/.buildkite/test_areas/compile.yaml
index 163904e6aac2..46d980df398b 100644
--- a/.buildkite/test_areas/compile.yaml
+++ b/.buildkite/test_areas/compile.yaml
@@ -4,6 +4,7 @@ depends_on:
 steps:
 - label: ":nvidia: (H100) Distributed Compile"
   key: distributed-compile-unit-tests-2xh100
+  automatic_shard: true
   timeout_in_minutes: 45
   working_dir: "/vllm-workspace/"
   device: h100
@@ -17,6 +18,7 @@ steps:
   - pytest -s -v tests/compile/passes/distributed
   mirror:
     amd:
+      automatic_shard: true
       label: ':amd: (MI300) Distributed Compile'
       dind: false
       device: mi300_2
@@ -36,6 +38,7 @@ steps:
 
 - label: ":nvidia: (B200) Fusion and Compile"
   key: fusion-and-compile-unit-tests-2xb200
+  automatic_shard: true
   timeout_in_minutes: 30
   working_dir: "/vllm-workspace/"
   device: b200-k8s
@@ -65,6 +68,7 @@ steps:
     - pytest -v -s tests/compile/passes/distributed/test_fusion_all_reduce.py
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355) Fusion and Compile FP8"
       dind: false
       device: mi355_2
@@ -95,6 +99,7 @@ steps:
 
 - label: ":nvidia: (H100) Fusion E2E Quick"
   key: fusion-e2e-quick-h100
+  automatic_shard: true
   timeout_in_minutes: 25
   working_dir: "/vllm-workspace/"
   device: h100
@@ -113,6 +118,7 @@ steps:
     - pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k "inductor_partition and not +rms_norm and +quant_fp8 and (qwen3 or deepseek)"
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355 DPX) Fusion E2E Quick"
       dind: false
       device: mi355_dpx
@@ -135,6 +141,7 @@ steps:
 
 - label: ":nvidia: (H100) Fusion E2E Config Sweep"
   key: fusion-e2e-config-sweep-h100
+  automatic_shard: true
   timeout_in_minutes: 25
   working_dir: "/vllm-workspace/"
   device: h100
@@ -154,6 +161,7 @@ steps:
     - pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k "llama-3"
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355 DPX) Fusion E2E Config Sweep"
       dind: false
       device: mi355_dpx
@@ -212,6 +220,7 @@ steps:
 
 - label: ":nvidia: (H100) Fusion E2E TP2 Quick"
   key: fusion-e2e-tp2-quick-h100
+  automatic_shard: true
   timeout_in_minutes: 35
   working_dir: "/vllm-workspace/"
   device: h100
@@ -228,6 +237,7 @@ steps:
     - 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))"
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355) Fusion E2E TP2 Quick"
       dind: false
       device: mi355_2
@@ -270,6 +280,7 @@ steps:
 
 - label: ":nvidia: (H100) Fusion E2E TP2 AR-RMS Config Sweep"
   key: fusion-e2e-tp2-ar-rms-config-sweep-h100
+  automatic_shard: true
   timeout_in_minutes: 30
   working_dir: "/vllm-workspace/"
   device: h100
@@ -289,6 +300,7 @@ steps:
     - pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py -k "llama-3"
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355) Fusion E2E TP2 AR-RMS FP8 Config Sweep"
       dind: false
       device: mi355_2
@@ -315,6 +327,7 @@ steps:
 
 - label: ":nvidia: (B200) Fusion E2E TP2"
   key: fusion-e2e-tp2-b200
+  automatic_shard: true
   timeout_in_minutes: 45
   working_dir: "/vllm-workspace/"
   device: b200-k8s
@@ -333,6 +346,7 @@ steps:
     - pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py -k "(FLASHINFER and inductor_partition and not +rms_norm and (not +quant_fp8 or +quant_fp8 and (qwen3 or deepseek))) or Llama-3.1-8B-Instruct-FP4"
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355) Fusion E2E TP2 FP8 BF16 rewrites"
       dind: false
       device: mi355_2
@@ -360,6 +374,7 @@ steps:
 
 - label: ':amd: (MI355) Fusion E2E TP2 AR-RMS (Inductor Partition)'
   key: fusion-e2e-tp2-ar-rms-inductor-partition-amd
+  automatic_shard: true
   dind: false
   device: mi355_2
   num_devices: 2
@@ -381,6 +396,7 @@ steps:
 
 - label: ':amd: (MI355) Fusion E2E TP2 AR-RMS (Dynamo Partition)'
   key: fusion-e2e-tp2-ar-rms-dynamo-partition-amd
+  automatic_shard: true
   dind: false
   device: mi355_2
   num_devices: 2
diff --git a/.buildkite/test_areas/cuda.yaml b/.buildkite/test_areas/cuda.yaml
index b3bfa9f09708..bc4afc0d9d36 100644
--- a/.buildkite/test_areas/cuda.yaml
+++ b/.buildkite/test_areas/cuda.yaml
@@ -4,6 +4,7 @@ depends_on:
 steps:
 - label: ":nvidia: (H200 MIG 18GB) CUDA Platform"
   key: platform-tests
+  automatic_shard: true
   timeout_in_minutes: 20
   device: h200_18gb
   source_file_dependencies:
@@ -19,6 +20,7 @@ steps:
     - pytest -v -s cuda/test_cuda_compatibility_path.py
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355 DPX) CUDA Platform"
       dind: false
       device: mi355_dpx
@@ -54,6 +56,7 @@ steps:
     - pytest -v -s v1/cudagraph/test_encoder_cudagraph.py
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI250) CUDAGraph"
       dind: false
       device: mi250_1
diff --git a/.buildkite/test_areas/distributed.yaml b/.buildkite/test_areas/distributed.yaml
index 883aa03aba05..ef3d5e339f09 100644
--- a/.buildkite/test_areas/distributed.yaml
+++ b/.buildkite/test_areas/distributed.yaml
@@ -4,6 +4,7 @@ depends_on:
 steps:
 - label: ":nvidia: (L4) Distributed Comm Ops"
   key: distributed-comm-ops
+  automatic_shard: true
   timeout_in_minutes: 25
   working_dir: "/vllm-workspace/tests"
   device: l4
@@ -21,6 +22,7 @@ steps:
   - pytest -v -s distributed/test_shm_storage.py
   mirror:
     amd:
+      automatic_shard: true
       label: ':amd: (MI300) Distributed Comm Ops'
       dind: false
       device: mi300_2
@@ -40,6 +42,7 @@ steps:
 
 - label: ":nvidia: (L4) Sharded RDT Weight Transfer"
   key: sharded-rdt-weight-transfer
+  automatic_shard: true
   device: l4
   timeout_in_minutes: 20
   working_dir: "/vllm-workspace/tests"
@@ -55,6 +58,7 @@ steps:
   - pytest -v -s distributed/test_sharded_rdt_trainer.py
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355 DPX) Sharded RDT Weight Transfer"
       dind: false
       device: mi355_dpx
@@ -130,6 +134,7 @@ steps:
 
 - label: ":nvidia: (L4) Distributed Compile + RPC"
   key: distributed-compile-rpc-tests-2-gpus
+  automatic_shard: true
   timeout_in_minutes: 65
   working_dir: "/vllm-workspace/tests"
   device: l4
@@ -144,14 +149,17 @@ steps:
   - vllm/v1/engine/
   - vllm/v1/worker/
   - tests/compile/test_wrapper.py
+  - tests/compile/test_aot_compile.py
   - tests/entrypoints/llm/test_collective_rpc.py
   commands:
   # https://github.com/NVIDIA/nccl/issues/1838
   - export NCCL_CUMEM_HOST_ENABLE=0
   - pytest -v -s entrypoints/llm/test_collective_rpc.py
   - pytest -v -s ./compile/test_wrapper.py
+  - pytest -v -s ./compile/test_aot_compile.py::test_load_on_another_device
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355) Distributed Compile + RPC"
       dind: false
       device: mi355_2
@@ -176,6 +184,7 @@ steps:
 
 - label: ":nvidia: (L4) Distributed Torchrun + Shutdown"
   key: distributed-torchrun-shutdown-tests-2-gpus
+  automatic_shard: true
   timeout_in_minutes: 30
   working_dir: "/vllm-workspace/tests"
   device: l4
@@ -200,6 +209,7 @@ steps:
   - pytest -v -s v1/worker/test_worker_memory_snapshot.py
   mirror:
     amd:
+      automatic_shard: true
       label: ':amd: (MI300) Distributed Torchrun + Shutdown'
       dind: false
       device: mi300_2
@@ -369,6 +379,7 @@ steps:
 
 - label: ":nvidia: (L4) [Nightly] Distributed Compile + Comm"
   key: distributed-compile-comm-4-gpus
+  automatic_shard: true
   # Nightly only: 4-GPU L4 pods wait behind smaller jobs on per-PR runs.
   optional: true
   timeout_in_minutes: 70
@@ -392,6 +403,7 @@ steps:
   - pytest -v -s distributed/test_multiproc_executor.py::test_multiproc_executor_multi_node
   mirror:
     amd:
+      automatic_shard: true
       # Run on PRs even when the L4 step is optional (nightly only).
       optional: false
       label: ':amd: (MI355) Distributed Compile + Comm'
@@ -416,6 +428,7 @@ steps:
 
 - label: ":nvidia: (H100) [Nightly] Distributed DP + EP"
   key: distributed-tests-8xh100
+  automatic_shard: true
   timeout_in_minutes: 20
   optional: true
   device: h100
@@ -430,14 +443,18 @@ steps:
   - vllm/v1/executor/uniproc_executor.py
   - vllm/v1/worker/gpu_worker.py
   - tests/distributed/test_mnnvl_alltoall.py
+  - tests/distributed/test_eplb_execute.py
 
   commands:
   # https://github.com/NVIDIA/nccl/issues/1838
   - 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
+  # The L4 EPLB Execution job has no NVLink, so it skips symm_mem.
+  - pytest -v -s distributed/test_eplb_execute.py -k symm_mem
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355) Distributed DP + EP"
       dind: false
       device: mi355_8
@@ -462,43 +479,34 @@ steps:
       - export TORCH_CPP_LOG_LEVEL=INFO
       - torchrun --nproc-per-node=8 ../examples/features/torchrun/torchrun_dp_example_offline.py --tp-size=2 --pp-size=1 --dp-size=4 --enable-ep
 
-- label: ":nvidia: (A100) [Nightly] Distributed"
-  key: distributed-tests-4xa100
-  device: a100
+# Key kept from the former A100 step's mirror for job history and retries.
+- label: ':amd: (MI355) Distributed'
+  key: amd-distributed-tests-4xa100
+  dind: false
+  device: mi355_4
   optional: true
   num_devices: 4
+  timeout_in_minutes: 40
+  working_dir: /vllm-workspace/tests
+  depends_on:
+  - image-build-amd
   source_file_dependencies:
   - vllm/
   - "!vllm/distributed/kv_transfer/"
+  - tests/models/test_vision.py
+  - vllm/platforms/rocm.py
+  - tests/distributed/
+  - tests/basic_correctness/
+  - tests/lora/test_mixtral.py
   commands:
   # NOTE: don't test llama model here, it seems hf implementation is buggy
   # see https://github.com/vllm-project/vllm/pull/5689 for details
   - pytest -v -s distributed/test_custom_all_reduce.py
   - torchrun --nproc_per_node=2 distributed/test_ca_buffer_sharing.py
-  - TARGET_TEST_SUITE=A100 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)'
+  - VLLM_ALLOW_INSECURE_SERIALIZATION=1 pytest -v -s distributed/test_weight_transfer.py -k multiple_sources
+  - TARGET_TEST_SUITE=MI355 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)'
   - pytest -v -s -x lora/test_mixtral.py
-  mirror:
-    amd:
-      label: ':amd: (MI355) Distributed'
-      dind: false
-      device: mi355_4
-      timeout_in_minutes: 40
-      working_dir: /vllm-workspace/tests
-      depends_on:
-      - image-build-amd
-      source_file_dependencies:
-      - vllm/
-      - tests/models/test_vision.py
-      - vllm/platforms/rocm.py
-      - tests/distributed/
-      - tests/basic_correctness/
-      - tests/lora/test_mixtral.py
-      commands:
-      - pytest -v -s distributed/test_custom_all_reduce.py
-      - torchrun --nproc_per_node=2 distributed/test_ca_buffer_sharing.py
-      - TARGET_TEST_SUITE=MI355 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)'
-      - pytest -v -s -x lora/test_mixtral.py
-      - pytest -v -s models/test_vision.py -m 'distributed(num_gpus=4)'
+  - pytest -v -s models/test_vision.py -m 'distributed(num_gpus=4)'
 
 - label: ":nvidia: (H100) [Nightly] Distributed Features"
   key: distributed-tests-2xh100-2xmi300
@@ -572,6 +580,7 @@ steps:
 
 - label: ":nvidia: (B200) [Nightly] Distributed"
   key: distributed-tests-2xb200
+  automatic_shard: true
   device: b200-k8s
   optional: true
   working_dir: "/vllm-workspace/"
@@ -584,6 +593,7 @@ steps:
     - pytest -v -s tests/distributed/test_mnnvl_alltoall.py
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355) Distributed AgRs All2All"
       dind: false
       device: mi355_2
@@ -606,6 +616,7 @@ steps:
 
 - label: ":nvidia: (B200) GEMM-RS/AR"
   key: gemm-rs-ar-2xb200
+  automatic_shard: true
   timeout_in_minutes: 20
   device: b200-k8s
   working_dir: "/vllm-workspace/"
@@ -618,6 +629,7 @@ steps:
     - pytest -v -s tests/kernels/test_gemm_rs_ar.py
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355) GEMM-RS-AR"
       dind: false
       device: mi355_4
@@ -774,11 +786,13 @@ steps:
 
 - label: ":nvidia: (L4) Distributed Preload"
   key: preload-2-gpus
+  automatic_shard: true
   timeout_in_minutes: 15
   working_dir: "/vllm-workspace/tests"
   device: l4
   num_devices: 2
   source_file_dependencies:
+  - vllm/entrypoints/cli/preload.py
   - vllm/model_executor/model_loader/weight_cache/
   - vllm/model_executor/layers/fused_moe/expert_map_manager.py
   - tests/model_executor/model_loader/test_weight_cache.py
@@ -786,6 +800,7 @@ steps:
   - pytest -v -s model_executor/model_loader/test_weight_cache.py -m 'distributed(num_gpus=2)'
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355) Distributed Preload"
       dind: false
       device: mi355_2
@@ -794,6 +809,7 @@ steps:
       depends_on:
       - image-build-amd
       source_file_dependencies:
+      - vllm/entrypoints/cli/preload.py
       - vllm/model_executor/model_loader/weight_cache/
       - vllm/model_executor/layers/fused_moe/expert_map_manager.py
       - tests/model_executor/model_loader/test_weight_cache.py
diff --git a/.buildkite/test_areas/docker.yaml b/.buildkite/test_areas/docker.yaml
index 632c3aa1322d..8e08f158c64b 100644
--- a/.buildkite/test_areas/docker.yaml
+++ b/.buildkite/test_areas/docker.yaml
@@ -9,16 +9,24 @@ steps:
   source_file_dependencies:
     - .buildkite/release-pipeline.yaml
     - .buildkite/scripts/docker-build-metadata-args.sh
-    - .buildkite/scripts/rocm/build-config.sh
+    - .buildkite/scripts/rocm/build.sh
+    - .buildkite/scripts/rocm/publish.sh
+    - tools/vllm-rocm/therock_wheels.py
     - docker/Dockerfile
     - docker/Dockerfile.cpu
+    - docker/Dockerfile.rocm
+    - docker/Dockerfile.rocm_base
     - docker/docker-bake.hcl
+    - docker/docker-bake-rocm.hcl
     - tests/tools/test_docker_build_metadata_args.py
+    - tests/tools/test_rocm_build.py
+    - tests/tools/test_amd_ci_fast_exit.py
+    - .buildkite/scripts/hardware_ci/amd_ci_fast_exit.py
     - tests/tools/test_kernel_symbol_map.py
     - tools/ci/kernel_symbol_map.py
     - tools/ci/rocm_kernel_symbol_map.py
   commands:
-    - pytest -v -s tools/test_docker_build_metadata_args.py
+    - pytest -v -s tools/test_docker_build_metadata_args.py tools/test_rocm_build.py tools/test_amd_ci_fast_exit.py
     - pytest -v -s tools/test_kernel_symbol_map.py
   mirror:
     amd:
@@ -31,12 +39,16 @@ steps:
       - image-build-amd
       source_file_dependencies:
       - .buildkite/scripts/docker-build-metadata-args.sh
-      - .buildkite/scripts/ci-bake-rocm.sh
       - docker/Dockerfile.rocm
       - docker/Dockerfile.rocm_base
-      - docker/ci-rocm.hcl
       - docker/docker-bake-rocm.hcl
+      - .buildkite/scripts/rocm/build.sh
+      - .buildkite/scripts/rocm/publish.sh
+      - tools/vllm-rocm/therock_wheels.py
       - tests/tools/test_docker_build_metadata_args.py
+      - tests/tools/test_rocm_build.py
+      - tests/tools/test_amd_ci_fast_exit.py
+      - .buildkite/scripts/hardware_ci/amd_ci_fast_exit.py
       - tests/tools/test_kernel_symbol_map.py
       - tools/ci/kernel_symbol_map.py
       - tools/ci/rocm_kernel_symbol_map.py
diff --git a/.buildkite/test_areas/engine.yaml b/.buildkite/test_areas/engine.yaml
index f6f91e26f86d..f2dc9e6041fc 100644
--- a/.buildkite/test_areas/engine.yaml
+++ b/.buildkite/test_areas/engine.yaml
@@ -4,6 +4,7 @@ depends_on:
 steps:
 - label: ":nvidia: (H200 MIG 18GB) Engine Core"
   key: engine
+  automatic_shard: true
   timeout_in_minutes: 30
   device: h200_18gb
   source_file_dependencies:
@@ -39,6 +40,7 @@ steps:
     amd:
       label: ":amd: (MI355 DPX) Engine Core"
       dind: false
+      automatic_shard: true
       device: mi355_dpx
       timeout_in_minutes: 40
       depends_on:
@@ -63,6 +65,7 @@ steps:
     - pytest -v -s v1/test_tensor_ipc_queue.py
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI250) V1 Engine"
       dind: false
       device: mi250_1
@@ -142,6 +145,7 @@ steps:
       -k "not test_mamba_prefix_cache_mrv2"
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI250) E2E Core"
       dind: false
       commands:
@@ -165,6 +169,7 @@ steps:
 - label: ":nvidia: (H200 MIG 35GB) E2E Core Large Memory"
   device: h200_35gb
   key: e2e-core-large-memory
+  automatic_shard: true
   timeout_in_minutes: 50
   source_file_dependencies:
     - vllm/v1/
@@ -178,6 +183,7 @@ steps:
       v1/e2e/general/test_mamba_prefix_cache.py::test_mamba_prefix_cache_mrv2_async_flashinfer
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355 DPX) E2E Core Large Memory"
       dind: false
       device: mi355_dpx
@@ -220,6 +226,7 @@ steps:
     - pytest -v -s v1/e2e/test_replayssm_decode.py
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355) ReplaySSM E2E"
       dind: false
       device: mi355_2
@@ -232,6 +239,7 @@ steps:
 
 - label: ":nvidia: (L4) [Nightly] V1 E2E"
   key: v1-e2e-2-gpus
+  automatic_shard: true
   timeout_in_minutes: 25 # TODO: Fix timeout after we have more confidence in the test stability
   optional: true
   device: l4
@@ -262,8 +270,10 @@ steps:
       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
+      v1/e2e/spec_decode/draft_model/test_structured_output_pp.py::test_specdec_pp_structured_output
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI300) V1 E2E"
       dind: false
       device: mi300_2
@@ -273,6 +283,7 @@ steps:
 
 - label: ":nvidia: (B200) [Nightly] V1 E2E Hybrid Chunked Prefill"
   key: v1-e2e-hybrid-chunked-prefill-b200
+  automatic_shard: true
   timeout_in_minutes: 35
   device: b200-k8s
   num_devices: 1
@@ -286,6 +297,7 @@ steps:
     - pytest -v -s v1/e2e/test_hybrid_chunked_prefill.py
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355 DPX) V1 E2E Hybrid Chunked Prefill"
       dind: false
       device: mi355_dpx
diff --git a/.buildkite/test_areas/entrypoints.yaml b/.buildkite/test_areas/entrypoints.yaml
index 8a568046eebc..72b993425e12 100644
--- a/.buildkite/test_areas/entrypoints.yaml
+++ b/.buildkite/test_areas/entrypoints.yaml
@@ -156,6 +156,7 @@ steps:
   - pytest -v -s entrypoints/openai/chat_completion
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355 DPX) Entrypoints Integration (OpenAI API chat_completion)"
       dind: false
       device: mi355_dpx
@@ -185,6 +186,7 @@ steps:
   - pytest -v -s entrypoints/cohere
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355 DPX) Entrypoints Integration (API Server Generate)"
       dind: false
       device: mi355_dpx
@@ -254,6 +256,7 @@ steps:
     amd:
       label: ":amd: (MI355 DPX) Entrypoints Integration (Multimodal)"
       dind: false
+      automatic_shard: true
       device: mi355_dpx
       timeout_in_minutes: 55
       depends_on:
@@ -285,6 +288,7 @@ steps:
 
 - label: ":nvidia: (H200 MIG 18GB) OpenAI API Correctness"
   key: openai-api-correctness
+  automatic_shard: true
   timeout_in_minutes: 20
   device: h200_18gb
   source_file_dependencies:
@@ -294,6 +298,7 @@ steps:
   - pytest -s entrypoints/openai/correctness/
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355 DPX) OpenAI API Correctness"
       dind: false
       device: mi355_dpx
diff --git a/.buildkite/test_areas/expert_parallelism.yaml b/.buildkite/test_areas/expert_parallelism.yaml
index 1ca3eacd48cb..e9c13dad6348 100644
--- a/.buildkite/test_areas/expert_parallelism.yaml
+++ b/.buildkite/test_areas/expert_parallelism.yaml
@@ -4,6 +4,7 @@ depends_on:
 steps:
 - label: ":nvidia: (H200 MIG 18GB) EPLB Algorithm"
   key: eplb-algorithm
+  automatic_shard: true
   timeout_in_minutes: 20
   device: h200_18gb
   working_dir: "/vllm-workspace/tests"
@@ -18,6 +19,7 @@ steps:
   - pytest -v -s distributed/test_eplb_utils.py
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355 DPX) EPLB Algorithm"
       dind: false
       device: mi355_dpx
@@ -32,6 +34,7 @@ steps:
 
 - label: ":nvidia: (L4) [Nightly] EPLB Execution"
   key: eplb-execution
+  automatic_shard: true
   # Nightly only: 4-GPU L4 pods wait behind smaller jobs on per-PR runs.
   optional: true
   timeout_in_minutes: 25
@@ -46,6 +49,7 @@ steps:
   - pytest -v -s distributed/test_eplb_spec_decode.py
   mirror:
     amd:
+      automatic_shard: true
       # Run on PRs even when the L4 step is optional (nightly only).
       optional: false
       label: ":amd: (MI355) EPLB Execution"
diff --git a/.buildkite/test_areas/fault_tolerance.yaml b/.buildkite/test_areas/fault_tolerance.yaml
index f9439880d454..972b405822cd 100644
--- a/.buildkite/test_areas/fault_tolerance.yaml
+++ b/.buildkite/test_areas/fault_tolerance.yaml
@@ -4,10 +4,10 @@ depends_on:
 steps:
 # Disabled until the KV connectors installation is locked.
 # - label: ":nvidia: (H100) Fault Tolerance E2E"
-#   key: fault-tolerance-e2e-2xh100
+#   key: fault-tolerance-e2e-4xh100
 #   timeout_in_minutes: 35
 #   device: h100
-#   num_devices: 2
+#   num_devices: 4
 #   working_dir: "/vllm-workspace/tests"
 #   source_file_dependencies:
 #   - vllm/v1/fault_tolerance/
@@ -62,6 +62,7 @@ steps:
 # Keep the AMD job: its commands do not run this installation.
 - label: ':amd: (MI355) MoRI EP Healthy Baseline'
   key: amd-fault-tolerance-e2e-2xh100
+  automatic_shard: true
   timeout_in_minutes: 10
   device: mi355_2
   num_devices: 2
@@ -93,6 +94,7 @@ steps:
   - tests/v1/fault_tolerance/
   - tests/v1/distributed/test_external_lb_dp.py
   commands:
+  - pytest -v -s v1/fault_tolerance/test_expert_reload.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]'
   dind: false
diff --git a/.buildkite/test_areas/jit_monitor.yaml b/.buildkite/test_areas/jit_monitor.yaml
index 9e40c9d1affc..2e606a988ca0 100644
--- a/.buildkite/test_areas/jit_monitor.yaml
+++ b/.buildkite/test_areas/jit_monitor.yaml
@@ -4,6 +4,7 @@ depends_on:
 steps:
 - label: ":nvidia: (H200 MIG 35GB) No Runtime JITs E2E"
   key: jit-monitor-no-runtime-jit
+  automatic_shard: true
   device: h200_35gb
   timeout_in_minutes: 45
   source_file_dependencies:
@@ -22,6 +23,7 @@ steps:
     - pytest -v -s jit_monitor/test_no_runtime_jit.py --timeout=900 --timeout-method=thread
   mirror:
     amd:
+      automatic_shard: true
       # Focused MRV2 regression; the broader model/backend matrix is pending.
       label: ':amd: (MI355 DPX) MRV2 Sampler JIT Warmup'
       dind: false
@@ -41,3 +43,4 @@ steps:
       - vllm/envs.py
       commands:
       - pytest -v -s tests/jit_monitor/test_no_runtime_jit_rocm.py --timeout=300
+      - pytest -v -s tests/jit_monitor/test_flydsl_contract.py
diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml
index 6784ee1b28bd..9ca0ba999e2e 100644
--- a/.buildkite/test_areas/kernels.yaml
+++ b/.buildkite/test_areas/kernels.yaml
@@ -4,6 +4,7 @@ depends_on:
 steps:
 - label: ":nvidia: (H200 MIG 18GB) vLLM IR"
   key: vllm-ir-tests
+  automatic_shard: true
   timeout_in_minutes: 35
   device: h200_18gb
   working_dir: "/vllm-workspace/"
@@ -15,6 +16,7 @@ steps:
     - pytest -v -s tests/kernels/ir
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355 DPX) vLLM IR"
       dind: false
       device: mi355_dpx
@@ -38,8 +40,10 @@ steps:
   source_file_dependencies:
   - csrc/
   - tests/kernels/core
+  - vllm/kernels/triton/
   - tests/kernels/test_concat_mla_q.py
   - tests/kernels/test_fused_qk_norm_rope_gate.py
+  - vllm/model_executor/layers/fused_qk_norm_rope.py
   commands:
     - >-
       pytest -v -s
@@ -60,6 +64,7 @@ steps:
       source_file_dependencies:
       - csrc/
       - tests/kernels/core
+      - vllm/kernels/triton/
       - tests/kernels/test_top_k_per_row.py
       - tests/kernels/test_concat_mla_q.py
       - tests/kernels/test_rocm_fp32_router_gemm.py
@@ -69,6 +74,7 @@ steps:
       - vllm/_aiter_ops.py
       - vllm/platforms/rocm.py
       - tests/kernels/test_fused_qk_norm_rope_gate.py
+      - vllm/model_executor/layers/fused_qk_norm_rope.py
       commands:
       - >-
         pytest -v -s
@@ -81,6 +87,7 @@ steps:
 
 - label: ":nvidia: (H100) MiniMax Reduce RMS Kernels"
   key: kernels-minimax-reduce-rms-test-2-gpus
+  automatic_shard: true
   timeout_in_minutes: 20
   num_devices: 2
   device: h100
@@ -96,6 +103,7 @@ steps:
     - pytest -v -s kernels/core/test_minimax_reduce_rms.py
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355) MiniMax Reduce RMS Kernels"
       dind: false
       device: mi355_2
@@ -114,6 +122,7 @@ steps:
 
 - label: ":nvidia: (H100) DeepSeek V4 Kernels"
   key: deepseek-v4-kernel-test-h100
+  automatic_shard: true
   timeout_in_minutes: 30
   device: h100
   source_file_dependencies:
@@ -122,12 +131,13 @@ steps:
   - tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py
   - tests/kernels/test_top_k_per_row.py # it runs on Blackwell too - some kernels have arch-specific optimizations
   commands:
-    - pytest -v -s kernels/test_fused_deepseek_v4_*.py
+    - pytest -v -s kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py
     - pytest -v -s kernels/test_top_k_per_row.py
   mirror:
     amd:
       label: ":amd: (MI355 DPX) DeepSeek V4 Fused Kernels"
       dind: false
+      automatic_shard: true
       device: mi355_dpx
       num_devices: 1
       timeout_in_minutes: 60
@@ -144,6 +154,7 @@ steps:
 
 - label: ":nvidia: (H100) BF16x3 Router GEMM Kernels"
   key: kernels-bf16x3-router-gemm-h100
+  automatic_shard: true
   timeout_in_minutes: 20
   device: h100
   source_file_dependencies:
@@ -155,6 +166,7 @@ steps:
 
 - label: ":nvidia: (B200) DeepSeek V4 Kernels"
   key: deepseek-v4-kernel-test-b200
+  automatic_shard: true
   timeout_in_minutes: 20
   device: b200-k8s
   source_file_dependencies:
@@ -171,13 +183,14 @@ steps:
   - vllm/model_executor/layers/sparse_mqa_indexer.py
   - tests/model_executor/layers/test_mla_short_prefill_indexer.py
   commands:
-    - pytest -v -s kernels/test_fused_deepseek_v4_*.py
+    - pytest -v -s kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py
     - pytest -v -s kernels/moe/test_deep_gemm_mega_moe.py
     - pytest -v -s models/test_deepseek_v4_mega_moe.py
     # DeepSeek V4.1 sparse MQA logits (MXFP4 and FP8 indexer cache); SM100 only.
     - pytest -v -s model_executor/layers/test_mla_short_prefill_indexer.py -k sparse_mqa_logits
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355) DeepSeek V4 MoE Integration"
       dind: false
       device: mi355_1
@@ -210,6 +223,7 @@ steps:
 # ptxas override.
 - label: ":nvidia: (B200) mHC Kernels"
   key: kernels-mhc-test-b200
+  automatic_shard: true
   timeout_in_minutes: 25
   device: b200-k8s
   source_file_dependencies:
@@ -221,6 +235,7 @@ steps:
     - pytest -v -s kernels/mhc
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355) mHC Kernels"
       dind: false
       device: mi355_1
@@ -236,6 +251,7 @@ steps:
 
 - label: ":nvidia: (B200) FLA Ops Kernels"
   key: kernels-fla-ops-test-b200
+  automatic_shard: true
   timeout_in_minutes: 25
   device: b200-k8s
   source_file_dependencies:
@@ -247,6 +263,7 @@ steps:
     - pytest -v -s kernels/fla
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355) FLA Ops Kernels"
       dind: false
       device: mi355_1
@@ -260,6 +277,32 @@ steps:
       - tests/kernels/
       - vllm/platforms/rocm.py
 
+- label: ":nvidia: (B200) MiniMax M3 MegaMoE"
+  key: minimax-m3-mega-moe-test-b200
+  automatic_shard: true
+  timeout_in_minutes: 20
+  device: b200-k8s
+  source_file_dependencies:
+  - cmake/external_projects/deepgemm.cmake
+  - vllm/models/minimax_m3/nvidia/
+  - vllm/utils/deep_gemm.py
+  - tests/models/test_minimax_m3_mega_moe.py
+  commands:
+    - pytest -v -s models/test_minimax_m3_mega_moe.py
+
+- label: ":nvidia: (B200) MiniMax-M3 Model Tests"
+  key: minimax-m3-model-tests-b200
+  automatic_shard: true
+  timeout_in_minutes: 20
+  device: b200-k8s
+  source_file_dependencies:
+  - vllm/models/minimax_m3/
+  - vllm/cute_utils/
+  - vllm/utils/flashinfer.py
+  - tests/models/minimax_m3/
+  commands:
+    - pytest -v -s models/minimax_m3 --ignore=models/minimax_m3/amd_specialized_decode
+
 # Catch-all for test files at the tests/kernels root. This job collects
 # the whole root so new files are wired by default.
 # Files with dedicated jobs elsewhere in this file are excluded via --ignore
@@ -345,8 +388,9 @@ steps:
           --deselect=tests/kernels/test_bf16_skinny_gemm.py::test_every_dsv3_routed_shape_is_instantiated
       - pytest -v -s kernels/test_fused_muse_spark_qknorm_rope_kv_insert.py -k rocm
 
-- label: ":nvidia: (L4) Attention Kernels Shard %N"
+- label: ":nvidia: (L4) Attention Kernels"
   key: kernels-attention-test
+  automatic_shard: true
   device: l4
   num_devices: 1
   timeout_in_minutes: 30
@@ -359,15 +403,14 @@ steps:
   - vllm/utils/flashinfer.py
   - tests/kernels/attention
   commands:
-    - pytest -v -s kernels/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
-  parallelism: 7
+    - pytest -v -s kernels/attention
   mirror:
     amd:
-      label: ":amd: (MI300) Attention Kernels Shard %N"
+      label: ":amd: (MI300) Attention Kernels"
+      automatic_shard: true
       dind: false
       device: mi300_1
       timeout_in_minutes: 90
-      parallelism: 6
       depends_on:
       - image-build-amd
       source_file_dependencies:
@@ -380,18 +423,11 @@ steps:
       - 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
+      - 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
 
 - label: ":nvidia: (H100) Attention DiffKV Kernels"
   key: kernels-attention-diffkv-test-h100
+  automatic_shard: true
   timeout_in_minutes: 20
   device: h100
   num_devices: 1
@@ -404,6 +440,7 @@ steps:
     - pytest -v -s kernels/attention/test_triton_unified_attention_diffkv.py
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355 DPX) Attention DiffKV Kernels"
       dind: false
       device: mi355_dpx
@@ -423,6 +460,7 @@ steps:
 
 - label: ":nvidia: (H100) MLA Kernel Test"
   key: kernels-flashmla-test-h100
+  automatic_shard: true
   timeout_in_minutes: 25
   device: h100
   num_devices: 1
@@ -444,6 +482,7 @@ steps:
     - pytest -v -s kernels/attention/test_mla_cross_layer_kernel_equivalence.py
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355 DPX) MLA Kernels"
       dind: false
       device: mi355_dpx
@@ -512,8 +551,9 @@ steps:
       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"
+- label: ":nvidia: (L4) MoE Kernels"
   key: kernels-moe-test
+  automatic_shard: true
   device: l4
   num_devices: 1
   timeout_in_minutes: 50
@@ -522,20 +562,18 @@ steps:
   - csrc/moe/
   - tests/kernels/moe
   - vllm/model_executor/layers/fused_moe/
+  - vllm/models/muse_spark/common/
+  - vllm/models/muse_spark/nvidia/
   - vllm/distributed/device_communicators/
   - vllm/envs.py
   - vllm/config
   commands:
-    - >-
-      pytest -v -s
-      kernels/moe
-      --ignore=kernels/moe/test_modular_oai_triton_moe.py
-      --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
-    - pytest -v -s kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
-  parallelism: 5
+    - pytest -v -s kernels/moe --ignore=kernels/moe/test_modular_oai_triton_moe.py
+    - pytest -v -s kernels/moe/test_modular_oai_triton_moe.py
   mirror:
     amd:
-      label: ":amd: (MI300) MoE Kernels Shard %N"
+      label: ":amd: (MI300) MoE Kernels"
+      automatic_shard: true
       dind: false
       device: mi300_1
       timeout_in_minutes: 55
@@ -544,6 +582,8 @@ steps:
       - csrc/moe/
       - tests/kernels/moe
       - vllm/model_executor/layers/fused_moe/
+      - vllm/models/muse_spark/common/
+      - vllm/models/muse_spark/nvidia/
       - vllm/distributed/device_communicators/
       - vllm/envs.py
       - vllm/config
@@ -566,6 +606,7 @@ steps:
   - tests/kernels/mamba
   - vllm/model_executor/layers/mamba/checkpoint.py
   - vllm/model_executor/layers/mamba/gdn
+  - vllm/models/qwen3_5/
   - vllm/model_executor/layers/mamba/ops
   - vllm/v1/attention/backends/gdn_attn.py
   commands:
@@ -585,6 +626,7 @@ steps:
 
 - label: ":nvidia: (H100) DeepGEMM Kernels"
   key: kernels-deepgemm-test-h100
+  automatic_shard: true
   timeout_in_minutes: 35
   device: h100
   num_devices: 1
@@ -611,6 +653,7 @@ steps:
 
 - label: ":nvidia: (B200) MiniMax-M3 Attention Kernels"
   key: kernels-minimax-m3-attention-b200
+  automatic_shard: true
   timeout_in_minutes: 30
   working_dir: "/vllm-workspace/"
   device: b200-k8s
@@ -625,6 +668,7 @@ steps:
     - pytest -v -s tests/kernels/attention/test_minimax_m3.py
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355 DPX) MiniMax-M3 Attention Kernels"
       dind: false
       device: mi355_dpx
@@ -635,9 +679,14 @@ steps:
       source_file_dependencies:
       - vllm/_aiter_ops.py
       - vllm/platforms/rocm.py
+      - tests/kernels/attention/test_minimax_m3_nvfp4_index.py
+      commands:
+      - pytest -v -s tests/kernels/attention/test_minimax_m3.py
+          tests/kernels/attention/test_minimax_m3_nvfp4_index.py
 
 - label: ":nvidia: (H100) MiniMax-M3 Sparse Prefill Kernels"
   key: kernels-minimax-m3-sparse-prefill-h100
+  automatic_shard: true
   timeout_in_minutes: 15
   working_dir: "/vllm-workspace/"
   device: h100
@@ -699,6 +748,8 @@ steps:
   - tests/kernels/mamba/test_gdn_prefill_cutedsl.py
   - tests/kernels/test_bf16x3_router_gemm_cutedsl.py
   - tests/kernels/test_top_k_per_row.py
+  - csrc/fused_mqa_topk/
+  - tests/kernels/attention/test_fused_mqa_topk.py
   - tests/kernels/quantization/test_silu_mul_nvfp4_quant.py
   - tests/kernels/quantization/test_block_fp8.py
   - tests/models/quantization/test_nvfp4.py
@@ -710,6 +761,7 @@ steps:
     - pytest -v -s tests/kernels/attention/test_flashinfer_trtllm_attention.py
     - pytest -v -s tests/kernels/attention/test_flashinfer_mla_decode.py
     - pytest -v -s tests/kernels/test_top_k_per_row.py
+    - pytest -v -s tests/kernels/attention/test_fused_mqa_topk.py
     - pytest -v -s tests/kernels/quantization/test_silu_mul_nvfp4_quant.py
     - pytest -v -s tests/kernels/quantization/test_flashinfer_scaled_mm.py
     - pytest -v -s tests/kernels/moe/test_nvfp4_moe.py
@@ -735,6 +787,7 @@ steps:
     - pytest -v -s tests/kernels/mamba/test_gdn_prefill_cutedsl.py
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355) NVFP4 Dense Emulation"
       dind: false
       device: mi355_1
@@ -762,6 +815,7 @@ steps:
 
 - label: ":nvidia: (B200) Gate Linear CuTeDSL Kernels"
   key: kernels-gate-linear-cutedsl-b200
+  automatic_shard: true
   timeout_in_minutes: 15
   working_dir: "/vllm-workspace/"
   device: b200-k8s
@@ -783,6 +837,7 @@ steps:
     - pytest -v -s tests/kernels/test_ll_fp32w_gemm.py
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355) Gate Linear Router Kernels"
       dind: false
       device: mi355_1
@@ -873,6 +928,7 @@ steps:
  
 - label: ":nvidia: (H100) [Nightly] FP8 MoE Kernels"
   key: kernels-fp8-moe-test-1xh100
+  automatic_shard: true
   timeout_in_minutes: 40
   device: h100
   num_devices: 1
@@ -890,6 +946,7 @@ steps:
 
 - label: ":nvidia: (H100) [Nightly] DeepEP FP8 MoE Kernels"
   key: kernels-fp8-moe-test-2xh100
+  automatic_shard: true
   timeout_in_minutes: 45
   device: h100
   num_devices: 2
@@ -900,6 +957,7 @@ steps:
   mirror:
     amd:
       label: ":amd: (MI355) DeepEP FP8 MoE Kernels"
+      automatic_shard: true
       dind: false
       device: mi355_2
       timeout_in_minutes: 55
@@ -919,6 +977,7 @@ steps:
 
 - label: ":nvidia: (B200) [Nightly] FP4 MoE Kernels"
   key: kernels-fp4-moe-test-b200
+  automatic_shard: true
   timeout_in_minutes: 25
   device: b200-k8s
   num_devices: 1
@@ -930,6 +989,7 @@ steps:
     - pytest -v -s kernels/moe/test_ocp_mx_moe.py
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355 DPX) FP4 and INT4 MoE Kernels"
       dind: false
       device: mi355_dpx
@@ -961,6 +1021,7 @@ steps:
 
 - label: ":nvidia: (H100) FusedMoE Layer Kernels"
   key: kernels-fusedmoe-layer-test-2-h100s
+  automatic_shard: true
   timeout_in_minutes: 30
   device: h100
   num_devices: 2
@@ -978,6 +1039,7 @@ steps:
     amd:
       label: ":amd: (MI355) FusedMoE Layer Kernels"
       dind: false
+      automatic_shard: true
       device: mi355_2
       timeout_in_minutes: 35
       working_dir: "/vllm-workspace/tests"
@@ -1004,6 +1066,7 @@ steps:
 
 - label: ":nvidia: (B200) FusedMoE Layer Kernels"
   key: kernels-fusedmoe-layer-test-2-b200s
+  automatic_shard: true
   timeout_in_minutes: 45
   device: b200-k8s
   num_devices: 2
@@ -1020,6 +1083,7 @@ steps:
     - pytest -v -s kernels/moe/test_deepep_v2_moe.py
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355) MoRI EP Numerics"
       dind: false
       device: mi355_2
@@ -1053,6 +1117,7 @@ steps:
 
 - label: ":amd: (MI355 DPX) Native Quantization Kernels"
   key: amd-native-quantization-kernels-mi355
+  automatic_shard: true
   dind: false
   device: mi355_dpx
   num_devices: 1
@@ -1081,6 +1146,7 @@ steps:
 
 - label: ":amd: (MI355 DPX) FP8 MoE Kernels"
   key: amd-fp8-moe-kernels-mi355
+  automatic_shard: true
   dind: false
   device: mi355_dpx
   num_devices: 1
@@ -1111,6 +1177,7 @@ steps:
 
 - label: ":amd: (MI355 DPX) Kernels"
   key: amd-kernels-mi355
+  automatic_shard: true
   dind: false
   device: mi355_dpx
   num_devices: 1
diff --git a/.buildkite/test_areas/lm_eval.yaml b/.buildkite/test_areas/lm_eval.yaml
index 527203d0a818..6ecfd895cc4a 100644
--- a/.buildkite/test_areas/lm_eval.yaml
+++ b/.buildkite/test_areas/lm_eval.yaml
@@ -15,6 +15,7 @@ steps:
   - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small.yaml
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI300) LM Eval Small Models"
       dind: false
       device: mi300_1
@@ -33,6 +34,7 @@ steps:
 
 - label: ":amd: (MI300) LM Eval Small Models Harness"
   key: amd-lm-eval-small-models-harness
+  automatic_shard: true
   dind: false
   device: mi300_1
   timeout_in_minutes: 35
@@ -55,6 +57,7 @@ steps:
 - label: ":nvidia: (H200 MIG 35GB) LM Eval Watermarking"
   device: h200_35gb
   key: lm-eval-watermarking
+  automatic_shard: true
   # Recorded runs took 80–140 seconds
   timeout_in_minutes: 10
   source_file_dependencies:
@@ -69,6 +72,7 @@ steps:
   - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-watermark.txt
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355 DPX) LM Eval Watermarking"
       dind: false
       device: mi355_dpx
@@ -90,19 +94,6 @@ steps:
       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
-#   optional: true
-#   num_devices: 4
-#   working_dir: "/vllm-workspace/.buildkite/lm-eval-harness"
-#   source_file_dependencies:
-#   - csrc/
-#   - vllm/model_executor/layers/quantization
-#   commands:
-#   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-#   - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large.txt --tp-size=4
-
 - label: ":nvidia: (H100) [Nightly] LM Eval Large Models"
   key: lm-eval-large-models-4xh100
   automatic_shard: true
@@ -178,6 +169,7 @@ steps:
 
 - label: ":nvidia: (B200) [Nightly] LM Eval Small Models Distributed"
   key: lm-eval-small-models-distributed-2xb200
+  automatic_shard: true
   timeout_in_minutes: 120
   device: b200-k8s
   num_devices: 2
@@ -191,6 +183,7 @@ steps:
   mirror:
     amd:
       label: ":amd: (MI355) LM Eval Small Models Distributed"
+      automatic_shard: true
       dind: false
       device: mi355_2
       timeout_in_minutes: 145
@@ -265,6 +258,7 @@ steps:
 
 - label: ":nvidia: (H100) LM Eval Watermark Feature Combination"
   key: lm-eval-dspark-watermark-2xh100
+  automatic_shard: true
   device: h100
   num_devices: 2
   # Local TP2 runs took 84–379 seconds, including startup.
@@ -294,6 +288,7 @@ steps:
     watermarking/test_parallel_compatibility.py
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355) LM Eval Watermark Feature Combination"
       dind: false
       device: mi355_2
@@ -454,6 +449,7 @@ steps:
   optional: true
   num_devices: 2
   source_file_dependencies:
+  - vllm/models/qwen3_5/
   - vllm/model_executor/models/qwen3_5.py
   - vllm/model_executor/models/qwen3_5_mtp.py
   - vllm/model_executor/models/qwen3_next.py
@@ -471,6 +467,7 @@ steps:
       depends_on:
       - image-build-amd
       source_file_dependencies:
+      - vllm/models/qwen3_5/
       - vllm/model_executor/models/qwen3_5.py
       - vllm/model_executor/models/qwen3_5_mtp.py
       - vllm/model_executor/models/qwen2.py
@@ -674,6 +671,7 @@ steps:
 
 - label: ":nvidia: (H200) [Nightly] LM Eval Nemotron 3 Super"
   key: lm-eval-nemotron-3-super-8xh200
+  automatic_shard: true
   timeout_in_minutes: 35
   device: h200
   optional: true
@@ -682,6 +680,7 @@ steps:
     - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-h200-nemotron-3-super.txt
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355) LM Eval Nemotron 3 Super"
       dind: false
       device: mi355_4
@@ -718,6 +717,7 @@ steps:
 
 - label: ":amd: (MI355) [Nightly] LM Eval Large Models"
   key: amd-lm-eval-large-models
+  automatic_shard: true
   dind: false
   device: mi355_8
   optional: true
@@ -741,6 +741,7 @@ steps:
     - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor/config-h100.txt
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355) MoE Refactor Integration"
       dind: false
       device: mi355_2
@@ -770,6 +771,7 @@ steps:
     - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor/config-b200.txt
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355) MoE Refactor BF16 Integration"
       dind: false
       device: mi355_2
@@ -804,6 +806,7 @@ steps:
     - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor-dp-ep/config-b200.txt
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355) MoE Refactor BF16 DP Integration"
       dind: false
       device: mi355_2
@@ -831,6 +834,7 @@ steps:
 
 - label: ":nvidia: (L4) [Nightly] LM Eval Humming TEMPORARY"
   key: lm-eval-humming-l4
+  automatic_shard: true
   optional: true
   timeout_in_minutes: 90
   device: l4
@@ -848,45 +852,13 @@ steps:
   commands:
   - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-l4.txt
 
-- label: ":nvidia: (A100) [Nightly] LM Eval Humming Default TEMPORARY Shard %N"
-  key: lm-eval-humming-a100-default
-  optional: true
-  timeout_in_minutes: 75
-  device: a100
-  num_devices: 1
-  parallelism: 3
-  source_file_dependencies:
-  - 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-a100-shard-$$BUILDKITE_PARALLEL_JOB.txt
-
-- label: ":nvidia: (A100) [Nightly] LM Eval Humming Activation INT8 TEMPORARY Shard %N"
-  key: lm-eval-humming-a100-act
-  optional: true
-  timeout_in_minutes: 45
-  device: a100
-  num_devices: 1
-  parallelism: 2
-  source_file_dependencies:
-  - 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-a100-act-int8-shard-$$BUILDKITE_PARALLEL_JOB.txt
-
-- label: ":nvidia: (H100) [Nightly] LM Eval Humming Default TEMPORARY Shard %N"
+- label: ":nvidia: (H100) [Nightly] LM Eval Humming Default TEMPORARY"
   key: lm-eval-humming-h100-default
+  automatic_shard: true
   optional: true
   timeout_in_minutes: 70
   device: h100
   num_devices: 1
-  parallelism: 3
   source_file_dependencies:
   - vllm/model_executor/layers/quantization/humming.py
   - vllm/model_executor/layers/quantization/utils/humming/
@@ -894,10 +866,11 @@ steps:
   - 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-h100-shard-$$BUILDKITE_PARALLEL_JOB.txt
+  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-h100.txt
 
 - label: ":nvidia: (H100) [Nightly] LM Eval Humming Activation FP8/INT8 TEMPORARY"
   key: lm-eval-humming-h100-act
+  automatic_shard: true
   optional: true
   timeout_in_minutes: 90
   device: h100
@@ -914,6 +887,7 @@ steps:
 
 - label: ":nvidia: (H100) [Nightly] LM Eval Humming DeepEP TEMPORARY"
   key: lm-eval-humming-h100-deepep
+  automatic_shard: true
   optional: true
   timeout_in_minutes: 60
   device: h100
@@ -934,6 +908,7 @@ steps:
 
 - label: ":nvidia: (B200) [Nightly] LM Eval Humming Default TEMPORARY"
   key: lm-eval-humming-b200-default
+  automatic_shard: true
   optional: true
   timeout_in_minutes: 50
   device: b200-k8s
@@ -965,6 +940,7 @@ steps:
 
 - label: ":nvidia: (H200 MIG 18GB) LM Eval TurboQuant k8v4"
   key: lm-eval-turboquant-k8v4
+  automatic_shard: true
   timeout_in_minutes: 30
   device: h200_18gb
   source_file_dependencies:
@@ -976,6 +952,7 @@ steps:
   - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/models-turboquant-k8v4.txt
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355 DPX) LM Eval TurboQuant k8v4"
       dind: false
       device: mi355_dpx
@@ -1005,6 +982,7 @@ steps:
 
 - label: ":nvidia: (H200 MIG 18GB) LM Eval TurboQuant t4nc"
   key: lm-eval-turboquant-t4nc
+  automatic_shard: true
   timeout_in_minutes: 30
   device: h200_18gb
   source_file_dependencies:
@@ -1018,6 +996,7 @@ steps:
     amd:
       label: ":amd: (MI355 DPX) LM Eval TurboQuant t4nc"
       dind: false
+      automatic_shard: true
       device: mi355_dpx
       num_devices: 1
       timeout_in_minutes: 30
@@ -1045,6 +1024,7 @@ steps:
 
 - label: ":nvidia: (H200 MIG 18GB) LM Eval TurboQuant k3v4nc"
   key: lm-eval-turboquant-k3v4nc
+  automatic_shard: true
   timeout_in_minutes: 30
   device: h200_18gb
   source_file_dependencies:
@@ -1058,6 +1038,7 @@ steps:
     amd:
       label: ":amd: (MI355 DPX) LM Eval TurboQuant k3v4nc"
       dind: false
+      automatic_shard: true
       device: mi355_dpx
       num_devices: 1
       timeout_in_minutes: 30
@@ -1085,6 +1066,7 @@ steps:
 
 - label: ":nvidia: (H200 MIG 18GB) LM Eval TurboQuant t3nc"
   key: lm-eval-turboquant-t3nc
+  automatic_shard: true
   timeout_in_minutes: 30
   device: h200_18gb
   source_file_dependencies:
@@ -1098,6 +1080,7 @@ steps:
     amd:
       label: ":amd: (MI355 DPX) LM Eval TurboQuant t3nc"
       dind: false
+      automatic_shard: true
       device: mi355_dpx
       num_devices: 1
       timeout_in_minutes: 30
@@ -1126,6 +1109,7 @@ steps:
 # Octave kernels are built only for gfx942 and gfx1100, so this runs on MI300.
 - label: ":amd: (MI300) LM Eval Octave KV Cache"
   key: amd-lm-eval-octave-mi300
+  automatic_shard: true
   dind: false
   device: mi300_1
   num_devices: 1
@@ -1182,6 +1166,7 @@ steps:
 
 - label: ":nvidia: (B200) [Nightly] GPQA Eval (GPT-OSS)"
   key: gpqa-eval-gpt-oss-2xb200
+  automatic_shard: true
   timeout_in_minutes: 30
   device: b200-k8s
   optional: true
@@ -1194,6 +1179,7 @@ steps:
     - pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-b200.txt
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355) GPQA Eval (GPT-OSS)"
       dind: false
       device: mi355_2
@@ -1272,6 +1258,7 @@ steps:
 
 - label: ":nvidia: (DGX Spark) [Nightly] GPQA Eval (GPT-OSS)"
   key: gpqa-eval-gpt-oss-spark
+  automatic_shard: true
   timeout_in_minutes: 35
   device: dgx-spark
   optional: true
@@ -1286,6 +1273,7 @@ steps:
     - pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-spark.txt
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355 DPX) GPQA Eval (GPT-OSS TP1)"
       # The Spark config is a backend-neutral GPT-OSS TP1 baseline.
       dind: false
@@ -1309,6 +1297,7 @@ steps:
 
 - label: ":nvidia: (H200 MIG 35GB) LM Eval KV-Offload"
   key: kv-offload-small
+  automatic_shard: true
   timeout_in_minutes: 30
   device: h200_35gb
   source_file_dependencies:
@@ -1321,6 +1310,7 @@ steps:
   - pytest -s -v evals/gsm8k/test_gsm8k_offloading.py -k "nemotron-h-8b or gemma-4-e4b-it"
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355 DPX) LM Eval KV-Offload"
       dind: false
       device: mi355_dpx
@@ -1334,10 +1324,13 @@ steps:
       - vllm/v1/kv_offload/
       - vllm/v1/simple_kv_offload/
       - tests/evals/gsm8k/test_gsm8k_offloading.py
+      - vllm/distributed/kv_transfer/kv_connector/v1/umbp/
+      - tests/v1/kv_connector/umbp_test_utils.py
       - vllm/platforms/rocm.py
 
 - label: ":nvidia: (H100) LM Eval KV-Offload Medium"
   key: kv-offload-medium
+  automatic_shard: true
   timeout_in_minutes: 45
   device: h100
   num_devices: 2
@@ -1351,6 +1344,7 @@ steps:
   - pytest -s -v evals/gsm8k/test_gsm8k_offloading.py -k "qwen3.5-35b or deepseek-v2-lite"
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355) LM Eval KV-Offload Medium"
       dind: false
       device: mi355_2
@@ -1368,6 +1362,7 @@ steps:
 
 - label: ":nvidia: (H100) [Nightly] LM Eval KV-Offload Large"
   key: kv-offload-large
+  automatic_shard: true
   # Nightly only: 4-GPU H100 pods wait behind smaller jobs on per-PR runs.
   optional: true
   timeout_in_minutes: 40
@@ -1383,6 +1378,7 @@ steps:
   - pytest -s -v evals/gsm8k/test_gsm8k_offloading.py -k "deepseek-v4-flash"
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355) LM Eval KV-Offload Large"
       dind: false
       device: mi355_4
@@ -1402,6 +1398,7 @@ steps:
 
 - label: ":nvidia: (H200 MIG 35GB) MRCR Eval Small Models"
   key: mrcr-eval-small-models
+  automatic_shard: true
   device: h200_35gb
   timeout_in_minutes: 25
   source_file_dependencies:
@@ -1410,6 +1407,7 @@ steps:
     - pytest -s -v evals/mrcr/test_mrcr_correctness.py --config-list-file=evals/mrcr/configs/models-small.txt
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355 DPX) MRCR Eval Small Models"
       dind: false
       device: mi355_dpx
@@ -1427,6 +1425,7 @@ steps:
 
 - label: ":amd: (MI355) [Nightly] LM Eval Large Models Harness"
   key: amd-lm-eval-large-models-harness
+  automatic_shard: true
   dind: false
   device: mi355_8
   optional: true
@@ -1452,6 +1451,7 @@ steps:
 
 - label: ":amd: (MI355) [Nightly] LM Eval Large Models ROCm Harness"
   key: amd-lm-eval-large-models-rocm-harness
+  automatic_shard: true
   dind: false
   device: mi355_8
   optional: true
diff --git a/.buildkite/test_areas/lora.yaml b/.buildkite/test_areas/lora.yaml
index afd7d07220e5..74770b4a4d89 100644
--- a/.buildkite/test_areas/lora.yaml
+++ b/.buildkite/test_areas/lora.yaml
@@ -52,6 +52,7 @@ steps:
   source_file_dependencies:
   - vllm/lora
   - vllm/model_executor/layers/fused_moe/
+  - vllm/model_executor/models/transformers/moe.py
   - tests/lora
   commands:
     # FIXIT: find out which code initialize cuda before running the test
@@ -84,6 +85,7 @@ steps:
       source_file_dependencies:
       - vllm/lora
       - vllm/model_executor/layers/fused_moe/
+      - vllm/model_executor/models/transformers/moe.py
       - vllm/platforms/rocm.py
       - tests/lora
       commands:
diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml
index 0cbfe1d0a883..3f162810a3d4 100644
--- a/.buildkite/test_areas/misc.yaml
+++ b/.buildkite/test_areas/misc.yaml
@@ -40,6 +40,7 @@ steps:
 
 - label: ":nvidia: (H200 MIG 18GB) V1 Logits + Oracle"
   key: v1-logits-oracle
+  automatic_shard: true
   timeout_in_minutes: 45
   device: h200_18gb
   source_file_dependencies: *v1-sample-logits-deps
@@ -55,6 +56,7 @@ steps:
     amd:
       label: ":amd: (MI355 DPX) V1 Logits + Oracle"
       dind: false
+      automatic_shard: true
       device: mi355_dpx
       timeout_in_minutes: 70
       depends_on:
@@ -63,6 +65,7 @@ steps:
 - label: ":nvidia: (H200 MIG 35GB) V1 Core"
   device: h200_35gb
   key: v1-core
+  automatic_shard: true
   timeout_in_minutes: 45
   source_file_dependencies: &v1-core-kv-metrics-deps
     - vllm/config/
@@ -98,6 +101,7 @@ steps:
     - pytest -v -s -m 'not cpu_test' v1/core
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355 DPX) V1 Core"
       dind: false
       device: mi355_dpx
@@ -121,6 +125,7 @@ steps:
     amd:
       label: ":amd: (MI355 DPX) V1 Executor + Worker"
       dind: false
+      automatic_shard: true
       device: mi355_dpx
       timeout_in_minutes: 60
       depends_on:
@@ -197,6 +202,7 @@ steps:
 # the connector or its tests change.
 - label: ":amd: (MI355 DPX) UMBP Connector Unit Tests"
   key: amd-umbp-connector-unit-mi355
+  automatic_shard: true
   dind: false
   device: mi355_dpx
   num_devices: 1
@@ -210,9 +216,33 @@ steps:
   - vllm/v1/kv_cache_interface.py
   - tests/v1/kv_connector/unit/test_umbp_
   - tests/v1/kv_connector/umbp_test_utils.py
+  - tests/v1/kv_connector/unit/utils.py
+  - tests/v1/kv_connector/conftest.py
+  - docker/Dockerfile.rocm_base
   commands:
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s -rs v1/kv_connector/unit/test_umbp_*.py
+  - pytest -v -s -rs v1/kv_connector/unit/test_umbp_async_lookup.py v1/kv_connector/unit/test_umbp_core.py v1/kv_connector/unit/test_umbp_embedded.py v1/kv_connector/unit/test_umbp_events.py v1/kv_connector/unit/test_umbp_hybrid.py
+
+- label: ":amd: (MI355 DPX) UMBP Connector Integration"
+  key: amd-umbp-connector-integration-mi355
+  dind: false
+  device: mi355_dpx
+  num_devices: 1
+  timeout_in_minutes: 30
+  working_dir: /vllm-workspace/tests
+  depends_on:
+  - image-build-amd
+  source_file_dependencies:
+  - vllm/distributed/kv_transfer/kv_connector/v1/umbp/
+  - vllm/distributed/kv_transfer/kv_connector/v1/base.py
+  - tests/v1/kv_connector/umbp_integration/
+  - tests/v1/kv_connector/umbp_test_utils.py
+  - tests/v1/kv_connector/conftest.py
+  - tests/v1/kv_connector/rocm_pd_accuracy_utils.py
+  - docker/Dockerfile.rocm_base
+  commands:
+  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  - pytest -v -s -rs v1/kv_connector/umbp_integration
 
 # Single-node MoRIIO P/D on dummy weights, WRITE and READ mode: prefill and
 # decode on separate GPUs over xGMI, compared exactly against a direct serve.
@@ -271,6 +301,7 @@ steps:
 - label: ":nvidia: (H200 MIG 35GB) V1 Metrics + LM Eval"
   device: h200_35gb
   key: v1-metrics-lmeval
+  automatic_shard: true
   timeout_in_minutes: 45
   env:
     # GitHub intermittently rejects HTTP/2 upload-pack requests from the H200
@@ -288,6 +319,7 @@ steps:
     - pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355 DPX) V1 Metrics + LM Eval"
       dind: false
       device: mi355_dpx
@@ -319,13 +351,13 @@ steps:
     - tests/v1
     - tests/watermarking
     - examples/disaggregated/mooncake_connector/mooncake_connector_proxy.py
+    - rust/src/engine-core-client/src/tests/prompt_block_hash_vectors.json
   device: cpu-small
   commands:
     # split the test to avoid interference
     - pytest -v -s -m 'cpu_test' v1/core
     - pytest -v -s v1/structured_output
     - pytest -v -s v1/test_serial_utils.py
-    - pytest -v -s v1/worker/test_gpu_worker.py::test_cpu_worker_without_triton
     - pytest -v -s v1/worker/test_cpu_worker_kv_budget.py
     - pytest -v -s v1/test_kv_cache_spec_registry.py
     - pytest -v -s v1/cudagraph/test_cudagraph_manager.py
@@ -336,6 +368,7 @@ steps:
 
 - label: ":nvidia: (H200 MIG 18GB) Extract Hidden States Integration"
   key: extract-hidden-states-integration
+  automatic_shard: true
   timeout_in_minutes: 20
   device: h200_18gb
   working_dir: "/vllm-workspace"
@@ -350,6 +383,7 @@ steps:
     - pytest -v -s tests/v1/kv_connector/extract_hidden_states_integration
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355 DPX) Extract Hidden States Integration Single GPU"
       dind: false
       device: mi355_dpx
@@ -372,6 +406,7 @@ steps:
 
 - label: ":nvidia: (L4) Extract Hidden States Integration"
   key: extract-hidden-states-integration-2-gpus
+  automatic_shard: true
   timeout_in_minutes: 20
   device: l4
   num_devices: 2
@@ -387,6 +422,7 @@ steps:
     - pytest -v -s -m 'distributed' tests/v1/kv_connector/extract_hidden_states_integration
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355) Extract Hidden States Integration"
       dind: false
       device: mi355_2
@@ -401,7 +437,9 @@ steps:
       - vllm/model_executor/model_loader/
       - vllm/model_executor/models/extract_hidden_states.py
       - vllm/model_executor/models/llama.py
+      - vllm/models/qwen3_5/
       - vllm/model_executor/models/qwen3_5.py
+      - vllm/model_executor/models/qwen3_5_mtp.py
       - vllm/model_executor/models/qwen3_next.py
       - vllm/model_executor/models/registry.py
       - vllm/transformers_utils/configs/extract_hidden_states.py
@@ -493,44 +531,40 @@ steps:
       depends_on:
       - image-build-amd
 
-# Disabled until the tracing packages installation is locked.
-# - label: ":nvidia: (L4) Metrics, Tracing"
-#   key: metrics-tracing-2-gpus
-#   timeout_in_minutes: 25
-#   device: l4
-#   num_devices: 2
-#   source_file_dependencies:
-#   - vllm/config/
-#   - vllm/distributed/
-#   - "!vllm/distributed/kv_transfer/"
-#   - vllm/engine/
-#   - vllm/inputs/
-#   - vllm/model_executor/
-#   - vllm/multimodal/
-#   - vllm/platforms/
-#   - vllm/sampling_params.py
-#   - vllm/tracing/
-#   - vllm/transformers_utils/
-#   - vllm/utils/
-#   - vllm/v1/
-#   - tests/v1/tracing
-#   - tests/tracing/
-#   commands:
-#   - "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'"
-#   - pytest -v -s v1/tracing
-#   - pytest -v -s tracing
-#   mirror:
-#     amd:
-#       label: ":amd: (MI300) Metrics, Tracing"
-#       dind: false
-#       device: mi300_2
-#       timeout_in_minutes: 30
-#       depends_on:
-#       - image-build-amd
+- label: ":nvidia: (L4) Metrics, Tracing"
+  key: metrics-tracing-2-gpus
+  automatic_shard: true
+  timeout_in_minutes: 25
+  device: l4
+  num_devices: 2
+  source_file_dependencies:
+  - vllm/config/
+  - vllm/distributed/
+  - "!vllm/distributed/kv_transfer/"
+  - vllm/engine/
+  - vllm/inputs/
+  - vllm/model_executor/
+  - vllm/multimodal/
+  - vllm/platforms/
+  - vllm/sampling_params.py
+  - vllm/tracing/
+  - vllm/transformers_utils/
+  - vllm/utils/
+  - vllm/v1/
+  - tests/v1/tracing
+  - tests/tracing/
+  commands:
+  - pytest -v -s v1/tracing
+  - pytest -v -s tracing
+  mirror:
+    amd:
+      automatic_shard: true
+      label: ":amd: (MI300) Metrics, Tracing"
+      dind: false
+      device: mi300_2
+      timeout_in_minutes: 30
+      depends_on:
+      - image-build-amd
 
 - label: ":nvidia: (H200 MIG 18GB) [Nightly] Python-only Installation"
   key: python-only-installation
@@ -570,6 +604,7 @@ steps:
 - label: ":nvidia: (H200 MIG 35GB) Async Engine, Inputs, Utils, Worker"
   device: h200_35gb
   key: async-engine-inputs-utils-worker
+  automatic_shard: true
   timeout_in_minutes: 25
   source_file_dependencies:
   - vllm/assets/
@@ -597,6 +632,7 @@ steps:
     amd:
       label: ":amd: (MI355 DPX) Async Engine, Inputs, Utils, Worker"
       dind: false
+      automatic_shard: true
       device: mi355_dpx
       timeout_in_minutes: 50
       depends_on:
@@ -845,149 +881,133 @@ steps:
   commands:
   - pytest -v -s tool_parsers
 
-# Disabled until the pytest plugins installation is locked.
-# - label: ":nvidia: (A100) Batch Invariance"
-#   key: batch-invariance-a100
-#   timeout_in_minutes: 60
-#   device: a100
-#   source_file_dependencies:
-#     - vllm/v1/attention
-#     - vllm/model_executor/layers
-#     - vllm/model_executor/kernels/linear
-#     - vllm/model_executor/determinism
-#     - tests/v1/determinism/
-#   commands:
-#     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-#     - pip install pytest-timeout pytest-forked
-#     - pytest -v -s v1/determinism/test_batch_invariance.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
-#   mirror:
-#     amd:
-#       label: ":amd: (MI250) Batch Invariance"
-#       dind: false
-#       device: mi250_1
-#       timeout_in_minutes: 60
-#       env:
-#         VLLM_GPU_MEMORY_UTILIZATION: "0.8"
-#         VLLM_TEST_CLEAN_GPU_MEMORY: "1"
-#       depends_on:
-#       - image-build-amd
-#       source_file_dependencies:
-#       - vllm/v1/attention
-#       - vllm/model_executor/layers
-#       - vllm/model_executor/kernels/linear
-#       - vllm/model_executor/determinism
-#       - tests/v1/determinism/
-#       - vllm/_aiter_ops.py
-#       - vllm/platforms/rocm.py
+# Key kept from the former A100 step's mirror for job history and retries.
+- label: ":amd: (MI250) Batch Invariance"
+  key: amd-batch-invariance-a100
+  dind: false
+  device: mi250_1
+  timeout_in_minutes: 60
+  env:
+    VLLM_GPU_MEMORY_UTILIZATION: "0.8"
+    VLLM_TEST_CLEAN_GPU_MEMORY: "1"
+  depends_on:
+  - image-build-amd
+  source_file_dependencies:
+  - vllm/v1/attention
+  - vllm/model_executor/layers
+  - vllm/model_executor/kernels/linear
+  - vllm/model_executor/determinism
+  - tests/v1/determinism/
+  - vllm/_aiter_ops.py
+  - vllm/platforms/rocm.py
+  - docker/Dockerfile.rocm_base
+  commands:
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - pytest -v -s v1/determinism/test_batch_invariance.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
 
-# Disabled until the pytest plugins installation is locked.
-# - label: ":nvidia: (H100) Batch Invariance"
-#   key: batch-invariance-h100
-#   automatic_shard: true
-#   timeout_in_minutes: 60
-#   device: h100
-#   source_file_dependencies:
-#     - vllm/v1/attention
-#     - vllm/model_executor/layers
-#     - vllm/lora/ops/triton_ops
-#     - vllm/model_executor/kernels/linear
-#     - vllm/model_executor/determinism
-#     - tests/v1/determinism/
-#   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
-#     - pytest -v -s v1/determinism/test_lora_shrink_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: (MI300) Batch Invariance"
-#       dind: false
-#       automatic_shard: true
-#       device: mi300_1
-#       timeout_in_minutes: 60
-#       depends_on:
-#       - image-build-amd
-#       source_file_dependencies:
-#       - vllm/v1/attention
-#       - vllm/model_executor/layers
-#       - vllm/lora/ops/triton_ops
-#       - vllm/model_executor/kernels/linear
-#       - vllm/model_executor/determinism
-#       - 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
-#       - pytest -v -s v1/determinism/test_lora_shrink_batch_invariant.py
-#       - pytest -v -s v1/determinism/test_cutlass_batch_invariance.py
-#       - pytest -v -s v1/determinism/test_matmul_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 TRITON_ATTN
+- label: ":nvidia: (H100) Batch Invariance"
+  key: batch-invariance-h100
+  automatic_shard: true
+  timeout_in_minutes: 60
+  device: h100
+  source_file_dependencies:
+    - vllm/v1/attention
+    - vllm/model_executor/layers
+    - vllm/lora/ops/triton_ops
+    - vllm/model_executor/kernels/linear
+    - vllm/model_executor/determinism
+    - tests/v1/determinism/
+  commands:
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - pytest -v -s v1/determinism/test_batch_invariance.py
+    - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py
+    - pytest -v -s v1/determinism/test_lora_shrink_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: (MI300) Batch Invariance"
+      dind: false
+      automatic_shard: true
+      device: mi300_1
+      timeout_in_minutes: 60
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - vllm/v1/attention
+      - vllm/model_executor/layers
+      - vllm/lora/ops/triton_ops
+      - vllm/model_executor/kernels/linear
+      - vllm/model_executor/determinism
+      - tests/v1/determinism/
+      - vllm/_aiter_ops.py
+      - vllm/platforms/rocm.py
+      commands:
+      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+      - pytest -v -s v1/determinism/test_batch_invariance.py
+      - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py
+      - pytest -v -s v1/determinism/test_lora_shrink_batch_invariant.py
+      - pytest -v -s v1/determinism/test_cutlass_batch_invariance.py
+      - pytest -v -s v1/determinism/test_matmul_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 TRITON_ATTN
 
-# Disabled until the pytest plugins installation is locked.
-# - label: ":nvidia: (B200) Batch Invariance"
-#   key: batch-invariance-b200
-#   automatic_shard: true
-#   timeout_in_minutes: 45
-#   device: b200-k8s
-#   source_file_dependencies:
-#     - vllm/v1/attention
-#     - vllm/model_executor/layers
-#     - vllm/lora/ops/triton_ops
-#     - vllm/model_executor/kernels/linear
-#     - vllm/model_executor/determinism
-#     - tests/v1/determinism/
-#   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
-#     - pytest -v -s v1/determinism/test_lora_shrink_batch_invariant.py
-#     - pytest -v -s v1/determinism/test_batch_invariance_vlm.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
-#     - pytest -v -s v1/determinism/test_nvfp4_batch_invariant.py
-#     - pytest -v -s v1/determinism/test_matmul_batch_invariant.py
-#     - pytest -v -s v1/determinism/test_cutlass_batch_invariance.py
-#     - pytest -v -s v1/determinism/test_online_batch_invariance.py
-#   mirror:
-#     amd:
-#       label: ":amd: (MI355) Batch Invariance"
-#       dind: false
-#       device: mi355_1
-#       timeout_in_minutes: 60
-#       depends_on:
-#       - image-build-amd
-#       source_file_dependencies:
-#       - vllm/v1/attention
-#       - vllm/model_executor/layers
-#       - vllm/lora/ops/triton_ops
-#       - vllm/model_executor/kernels/linear
-#       - vllm/model_executor/determinism
-#       - 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
-#       - pytest -v -s v1/determinism/test_lora_shrink_batch_invariant.py
-#       - pytest -v -s v1/determinism/test_batch_invariance_vlm.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 TRITON_ATTN
-#       - pytest -v -s v1/determinism/test_nvfp4_batch_invariant.py
-#       - pytest -v -s v1/determinism/test_matmul_batch_invariant.py
-#       - pytest -v -s v1/determinism/test_cutlass_batch_invariance.py
-#       - pytest -v -s v1/determinism/test_online_batch_invariance.py
-#       - pytest -v -s v1/determinism/test_mx_linear_batch_invariant.py
+- label: ":nvidia: (B200) Batch Invariance"
+  key: batch-invariance-b200
+  automatic_shard: true
+  timeout_in_minutes: 45
+  device: b200-k8s
+  source_file_dependencies:
+    - vllm/v1/attention
+    - vllm/model_executor/layers
+    - vllm/lora/ops/triton_ops
+    - vllm/model_executor/kernels/linear
+    - vllm/model_executor/determinism
+    - tests/v1/determinism/
+  commands:
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - pytest -v -s v1/determinism/test_batch_invariance.py
+    - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py
+    - pytest -v -s v1/determinism/test_lora_shrink_batch_invariant.py
+    - pytest -v -s v1/determinism/test_batch_invariance_vlm.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
+    - pytest -v -s v1/determinism/test_nvfp4_batch_invariant.py
+    - pytest -v -s v1/determinism/test_matmul_batch_invariant.py
+    - pytest -v -s v1/determinism/test_cutlass_batch_invariance.py
+    - pytest -v -s v1/determinism/test_online_batch_invariance.py
+  mirror:
+    amd:
+      label: ":amd: (MI355) Batch Invariance"
+      dind: false
+      automatic_shard: true
+      device: mi355_1
+      timeout_in_minutes: 60
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - vllm/v1/attention
+      - vllm/model_executor/layers
+      - vllm/lora/ops/triton_ops
+      - vllm/model_executor/kernels/linear
+      - vllm/model_executor/determinism
+      - tests/v1/determinism/
+      - vllm/_aiter_ops.py
+      - vllm/platforms/rocm.py
+      commands:
+      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+      - pytest -v -s v1/determinism/test_batch_invariance.py
+      - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py
+      - pytest -v -s v1/determinism/test_lora_shrink_batch_invariant.py
+      - pytest -v -s v1/determinism/test_batch_invariance_vlm.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 TRITON_ATTN
+      - pytest -v -s v1/determinism/test_nvfp4_batch_invariant.py
+      - pytest -v -s v1/determinism/test_matmul_batch_invariant.py
+      - pytest -v -s v1/determinism/test_cutlass_batch_invariance.py
+      - pytest -v -s v1/determinism/test_online_batch_invariance.py
+      - pytest -v -s v1/determinism/test_mx_linear_batch_invariant.py
 
 - label: ":nvidia: (H100) [Nightly] Batch Invariance (Distributed)"
   key: batch-invariance-distributed-4xh100
@@ -1009,6 +1029,7 @@ steps:
     - VLLM_TEST_TP_SIZE=4 VLLM_TEST_DISABLE_CUSTOM_AR=1 pytest -v -s v1/determinism/test_batch_invariance.py -k bs1_vs_bsN
     - VLLM_TP_SIZE=4 VLLM_NEEDLE_TRIALS=2 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k FLASH_ATTN
     - VLLM_TEST_MODEL=allenai/OLMoE-1B-7B-0924 VLLM_TEST_SERVER_ARGS="--data-parallel-size 4 --enable-expert-parallel --max-model-len 4096" pytest -v -s v1/determinism/test_online_batch_invariance.py
+    - VLLM_TEST_MODEL=ibm-granite/granite-3.0-1b-a400m-base VLLM_TEST_SERVER_ARGS="--data-parallel-size 4 --enable-expert-parallel --max-model-len 4096" pytest -v -s v1/determinism/test_online_batch_invariance.py
   mirror:
     amd:
       label: ":amd: (MI300) Batch Invariance (Distributed)"
@@ -1031,10 +1052,12 @@ steps:
       - VLLM_TEST_TP_SIZE=4 VLLM_TEST_DISABLE_CUSTOM_AR=1 pytest -v -s v1/determinism/test_batch_invariance.py -k bs1_vs_bsN
       - VLLM_TP_SIZE=4 VLLM_NEEDLE_TRIALS=2 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_ATTN
       - VLLM_TEST_MODEL=allenai/OLMoE-1B-7B-0924 VLLM_TEST_SERVER_ARGS="--data-parallel-size 4 --enable-expert-parallel --max-model-len 4096" pytest -v -s v1/determinism/test_online_batch_invariance.py
+      - VLLM_TEST_MODEL=ibm-granite/granite-3.0-1b-a400m-base VLLM_TEST_SERVER_ARGS="--data-parallel-size 4 --enable-expert-parallel --max-model-len 4096" pytest -v -s v1/determinism/test_online_batch_invariance.py
 
 - label: ":nvidia: (H200 MIG 35GB) [Nightly] Acceptance Length (Large Models)"
   device: h200_35gb
   key: acceptance-length-test-large-models
+  automatic_shard: true
   timeout_in_minutes: 20
   gpu: h100
   optional: true
@@ -1048,6 +1071,7 @@ steps:
     - pytest -v -s v1/spec_decode/test_acceptance_length.py -m slow_test
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355 DPX) Acceptance Length (Large Models)"
       dind: false
       device: mi355_dpx
diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml
index b5c64ec7b3f0..03fbd875f782 100644
--- a/.buildkite/test_areas/models_basic.yaml
+++ b/.buildkite/test_areas/models_basic.yaml
@@ -4,6 +4,7 @@ depends_on:
 steps:
 - label: ":nvidia: (H200 MIG 18GB) Basic Models (Initialization)"
   key: basic-models-tests-initialization
+  automatic_shard: true
   timeout_in_minutes: 25
   device: h200_18gb
   source_file_dependencies:
@@ -19,6 +20,7 @@ steps:
     - pytest -v -s models/test_beam_construction.py
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355 DPX) Basic Models (Initialization)"
       dind: false
       device: mi355_dpx
@@ -70,16 +72,17 @@ steps:
   - tests/models/transformers/test_backend.py
   - tests/models/test_registry.py
   - tests/models/test_deepseek_v4_vl_rocm.py
+  - tests/models/test_deepseek_v4_vl.py
   - tests/models/test_deepseek_v41_decoder_replay_layers.py
   - tests/models/test_deepseek_v41_replay_batch.py
   - tests/models/test_deepseek_v41_pcp.py
   - tests/models/test_deepseek_v41_pipeline.py
   - tests/models/transformers/test_layer_registry_e2e.py
   - tests/models/test_deepseek_v41_replay_start.py
-  - tests/models/test_deepseek_v4_vl_stream_weights.py
   - tests/models/test_muse_spark_rocm_loading.py
   - tests/models/test_muse_spark_logits_local.py
   - tests/models/test_muse_spark_rocm_shared_rope.py
+  - tests/models/test_muse_spark_rocm_mono.py
   commands:
     - >-
       pytest -v -s
@@ -87,16 +90,17 @@ steps:
       models/transformers/test_backend.py
       models/test_registry.py
       models/test_deepseek_v4_vl_rocm.py
+      models/test_deepseek_v4_vl.py
       models/transformers/test_layer_registry_e2e.py
       models/test_deepseek_v41_replay_start.py
       models/test_deepseek_v41_decoder_replay_layers.py
       models/test_deepseek_v41_replay_batch.py
       models/test_deepseek_v41_pcp.py
       models/test_deepseek_v41_pipeline.py
-      models/test_deepseek_v4_vl_stream_weights.py
       models/test_muse_spark_rocm_loading.py
       models/test_muse_spark_logits_local.py
       models/test_muse_spark_rocm_shared_rope.py
+      models/test_muse_spark_rocm_mono.py
   mirror:
     amd:
       label: ":amd: (MI355) Basic Models (Other)"
@@ -110,6 +114,7 @@ steps:
       - vllm/models/deepseek_v4/
       - vllm/models/deepseek_v41/
       - tests/models/test_deepseek_v4_vl_rocm.py
+      - tests/models/test_deepseek_v4_vl.py
       - tests/models/test_hyv4_rocm.py
       - tests/models/transformers/test_layer_registry_e2e.py
       - tests/models/test_deepseek_v41_replay_start.py
@@ -122,11 +127,11 @@ steps:
       - tests/models/test_dsv41_mono_decode.py
       - tests/models/test_muse_spark_rocm_loading.py
       - tests/models/test_dsv4_mono_decode.py
-      - tests/models/test_deepseek_v4_vl_stream_weights.py
       - tests/models/hy_v4/test_triton_ihc.py
       - vllm/models/muse_spark/
       - tests/models/test_muse_spark_logits_local.py
       - tests/models/test_muse_spark_rocm_shared_rope.py
+      - tests/models/test_muse_spark_rocm_mono.py
       commands:
       - >-
         pytest -v -s
@@ -134,6 +139,7 @@ steps:
         models/transformers/test_backend.py
         models/test_registry.py
         models/test_deepseek_v4_vl_rocm.py
+        models/test_deepseek_v4_vl.py
         models/transformers/test_layer_registry_e2e.py
         models/test_deepseek_v41_replay_start.py
         models/test_deepseek_v41_decoder_replay_layers.py
@@ -144,15 +150,16 @@ steps:
         models/test_deepseek_v4_rocm_compressor_gemm_fusion.py
         models/test_dsv41_mono_decode.py
         models/test_muse_spark_rocm_loading.py
-        models/test_deepseek_v4_vl_stream_weights.py
         models/test_muse_spark_logits_local.py
         models/test_muse_spark_rocm_shared_rope.py
+        models/test_muse_spark_rocm_mono.py
       - VLLM_ROCM_USE_AITER=1 pytest -v -s models/test_hyv4_rocm.py -m 'not distributed'
       - pytest -v -s models/hy_v4/test_triton_ihc.py
       - pytest -v -s models/test_dsv4_mono_decode.py
 
 - label: ":amd: (MI355) MiniMax M3 MonoKernel"
   key: minimax-m3-monokernel-mi355
+  automatic_shard: true
   dind: false
   timeout_in_minutes: 45
   device: mi355_4
@@ -171,6 +178,7 @@ steps:
 
 - label: ":nvidia: (B200) Inkling"
   key: inkling-unit-tests-b200
+  automatic_shard: true
   timeout_in_minutes: 40
   device: b200-k8s
   source_file_dependencies:
@@ -183,6 +191,7 @@ steps:
     - pytest -v -s models/inkling
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355) Inkling"
       dind: false
       device: mi355_1
@@ -213,6 +222,7 @@ steps:
 
 - label: ":nvidia: (B200) Kimi K3"
   key: kimi-k3-unit-tests-b200
+  automatic_shard: true
   timeout_in_minutes: 40
   device: b200-k8s
   source_file_dependencies:
@@ -229,6 +239,7 @@ steps:
       kernels/attention/test_kimi_k3_mla_fused_epilogue.py
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355) Kimi K3"
       dind: false
       device: mi355_1
@@ -287,16 +298,47 @@ steps:
 
 - label: ":nvidia: (H200) Step5"
   key: step5-unit-tests
+  automatic_shard: true
   timeout_in_minutes: 20
   device: h200_35gb
   source_file_dependencies:
   - vllm/models/step5/
   - tests/models/step5/
+  - vllm/model_executor/warmup/step5_dsa_warmup.py
   commands:
     - pytest -v -s models/step5
+  mirror:
+    amd:
+      label: ":amd: (MI355 DPX) Step5"
+      dind: false
+      device: mi355_dpx
+      timeout_in_minutes: 30
+      working_dir: "/vllm-workspace/tests"
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - vllm/models/step5/
+      - tests/models/step5/
+      - vllm/model_executor/models/step3p7.py
+      - vllm/model_executor/models/step3_vl.py
+      - vllm/model_executor/models/step_vl.py
+      - vllm/config/compilation.py
+      - vllm/model_executor/layers/indexer_topk.py
+      - vllm/transformers_utils/configs/step5.py
+      - vllm/model_executor/warmup/step5_dsa_warmup.py
+      - vllm/v1/attention/backends/rocm_aiter_unified_attn.py
+      - vllm/triton_utils/
+      - vllm/platforms/rocm.py
+      - tests/v1/attention/utils.py
+      - tests/models/test_initialization.py
+      - tests/models/registry.py
+      commands:
+        - pytest -v -s models/step5
+        - pytest -v -s models/test_initialization.py -k Step5ForConditionalGeneration
 
 - label: ":nvidia: (H200) GLM5Next"
   key: glm5next-unit-tests
+  automatic_shard: true
   timeout_in_minutes: 20
   device: h200_35gb
   source_file_dependencies:
@@ -307,6 +349,7 @@ steps:
     - pytest -v -s models/glm5next
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355 DPX) GLM5Next"
       dind: false
       device: mi355_dpx
@@ -319,6 +362,7 @@ steps:
 
 - label: ":nvidia: (H200) Qwen4 Exp"
   key: qwen4-exp-unit-tests
+  automatic_shard: true
   timeout_in_minutes: 40
   device: h200_35gb
   source_file_dependencies:
@@ -336,6 +380,7 @@ steps:
     - pytest -v -s models/qwen4_exp/test_ple.py
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355 DPX) Qwen4 Exp"
       dind: false
       device: mi355_dpx
@@ -349,6 +394,26 @@ steps:
       # PLE and QSA import conflicting HC registrations in a shared process.
       - pytest -v -s models/qwen4_exp/test_ple.py
 
+- label: ":amd: (MI300) Qwen4 Exp"
+  key: qwen4-exp-unit-tests-mi300
+  automatic_shard: true
+  timeout_in_minutes: 40
+  dind: false
+  device: mi300_1
+  working_dir: "/vllm-workspace/tests"
+  depends_on:
+  - image-build-amd
+  source_file_dependencies:
+  - vllm/models/qwen4_exp/
+  - vllm/config/compilation.py
+  - vllm/v1/attention/backends/short_conv_attn.py
+  - vllm/model_executor/layers/mamba/checkpoint.py
+  - tests/models/qwen4_exp/
+  commands:
+  - pytest -v -s models/qwen4_exp/test_qsa_amd.py models/qwen4_exp/test_qsa_prepare_amd.py
+  # PLE and QSA import conflicting HC registrations in a shared process.
+  - pytest -v -s models/qwen4_exp/test_ple.py
+
 - label: ":computer: (CPU) Qwen4 Exp"
   key: qwen4-exp-unit-tests-cpu
   depends_on:
@@ -366,6 +431,20 @@ steps:
       models/qwen4_exp/test_config.py
       models/qwen4_exp/test_ple.py
 
+- label: ":nvidia: (B200) Qwen3.5"
+  key: qwen3-5-unit-tests-b200
+  timeout_in_minutes: 20
+  device: b200-k8s
+  source_file_dependencies:
+  - vllm/model_executor/layers/linear.py
+  - vllm/model_executor/layers/mamba/gdn
+  - vllm/model_executor/layers/utils.py
+  - vllm/models/qwen3_5/
+  - vllm/model_executor/models/qwen3_next.py
+  - tests/models/qwen3_5/
+  commands:
+  - pytest -v -s models/qwen3_5/
+
 - label: ":computer: (CPU) Basic Models Other"
   key: basic-models-test-other-cpu
   depends_on:
@@ -378,6 +457,7 @@ steps:
   - tests/models/test_vision.py
   - tests/models/test_adapters.py
   - tests/models/test_qwen3_5_mtp_config.py
+  - tests/models/test_qwen3_5.py
   - tests/models/test_dory.py
   - tests/models/test_alice_ai.py
   - tests/models/test_alice_ai_config.py
@@ -388,6 +468,7 @@ steps:
   - tests/model_executor/layers/test_activation.py
   - tests/models/test_language_model_cache_is_weak.py
   - tests/models/deepseek_v32/test_sequence_parallel.py
+  - tests/models/deepseek_v32/test_pp_topk.py
   - tests/models/test_deepseek_v4_fi_moe_ep.py
   - tests/models/test_dspark_mla.py
   - tests/models/transformers/test_layer_registry.py
@@ -400,11 +481,13 @@ steps:
       models/test_vision.py
       models/test_adapters.py
       models/test_qwen3_5_mtp_config.py
+      models/test_qwen3_5.py
       models/transformers/fusers/
       models/transformers/test_dynamic_rope.py
       model_executor/layers/test_activation.py
       models/test_language_model_cache_is_weak.py
       models/deepseek_v32/test_sequence_parallel.py
+      models/deepseek_v32/test_pp_topk.py
       models/test_deepseek_v4_fi_moe_ep.py
       models/test_dspark_mla.py
       models/transformers/test_layer_registry.py
@@ -417,6 +500,7 @@ steps:
 
 - label: ':amd: (MI355) DeepSeek-V4.1 Mono Decode Numerics'
   key: dsv41-mono-decode-numerics-amd
+  automatic_shard: true
   dind: false
   device: mi355_4
   num_devices: 4
@@ -434,8 +518,33 @@ steps:
   - pytest -v -s tests/models/test_dsv41_mono_numerics.py
   - pytest -v -s tests/models/test_dsv41_mono_layers.py
 
+- label: ':amd: (MI355) DeepSeek-V4.1 Mono Decode Numerics TP8'
+  key: dsv41-mono-decode-numerics-tp8-amd
+  automatic_shard: true
+  dind: false
+  device: mi355_8
+  num_devices: 8
+  timeout_in_minutes: 60
+  working_dir: /vllm-workspace/
+  depends_on:
+  - image-build-amd
+  source_file_dependencies:
+  - docker/Dockerfile.rocm_base
+  - vllm/models/deepseek_v41/amd/
+  - vllm/models/deepseek_v41/nvidia/model_state.py
+  - tests/models/test_dsv41_mono_numerics.py
+  - tests/models/test_dsv41_mono_layers.py
+  commands:
+  # only the TP8 cases; the 4-GPU step above runs TP2/TP4 (TP8 skips there)
+  - pytest -v -s "tests/models/test_dsv41_mono_numerics.py::test_mono_decode_layer_matches_vllm[8]"
+  - >-
+    pytest -v -s
+    "tests/models/test_dsv41_mono_layers.py::test_mono_decode_matches_the_model[8-8]"
+    "tests/models/test_dsv41_mono_layers.py::test_mono_decode_matches_the_model[8-28]"
+
 - label: ':amd: (MI355) DeepSeek-V4-Pro Mono Decode Numerics'
   key: dsv4-mono-decode-numerics-amd
+  automatic_shard: true
   dind: false
   device: mi355_8
   num_devices: 8
@@ -454,6 +563,7 @@ steps:
 
 - label: ':amd: (MI300) DeepSeek-V4.1 gfx942 Mono Decode Layers'
   key: dsv41-mono-decode-layers-gfx942-amd
+  automatic_shard: true
   dind: false
   device: mi300_4
   num_devices: 4
@@ -498,6 +608,7 @@ steps:
 
 - label: ':amd: (MI355) Kimi-K3 Specialized Decode'
   key: kimi-k3-specialized-decode-amd
+  automatic_shard: true
   dind: false
   device: mi355_8
   num_devices: 8
diff --git a/.buildkite/test_areas/models_distributed.yaml b/.buildkite/test_areas/models_distributed.yaml
index bdaef9ec5a1b..22e38ab1aa24 100644
--- a/.buildkite/test_areas/models_distributed.yaml
+++ b/.buildkite/test_areas/models_distributed.yaml
@@ -4,6 +4,7 @@ depends_on:
 steps:
 - label: ":nvidia: (L4) [Nightly] Qwen4 Exp SP"
   key: qwen4-exp-sp-tests-4-gpus
+  automatic_shard: true
   # Nightly only: 4-GPU L4 pods wait behind smaller jobs on per-PR runs.
   optional: true
   timeout_in_minutes: 30
@@ -22,6 +23,7 @@ steps:
   - pytest -v -s models/qwen4_exp/test_tpsp.py
   mirror:
     amd:
+      automatic_shard: true
       # Run on PRs even when the L4 step is optional (nightly only).
       optional: false
       label: ":amd: (MI355) Qwen4 Exp SP"
diff --git a/.buildkite/test_areas/models_language.yaml b/.buildkite/test_areas/models_language.yaml
index 909bb30e5e0f..a157d0ae76f8 100644
--- a/.buildkite/test_areas/models_language.yaml
+++ b/.buildkite/test_areas/models_language.yaml
@@ -4,6 +4,7 @@ depends_on:
 steps:
 - label: ":nvidia: (H200 MIG 18GB) Language Models (Standard)"
   key: language-models-tests-standard
+  automatic_shard: true
   timeout_in_minutes: 30
   device: h200_18gb
   source_file_dependencies:
@@ -11,12 +12,14 @@ steps:
   - "!vllm/distributed/kv_transfer/"
   - tests/models/language
   commands:
-    # Test standard language models, excluding a subset of slow tests
+    # Test standard language models, excluding a subset of slow tests.
+    # The exclusion sits in -k so the shard planner reads -m as picking tests.
     - pip freeze | grep -E 'torch'
-    - pytest -v -s models/language -m 'core_model and (not slow_test)'
+    - pytest -v -s models/language -m core_model -k 'not slow_test'
   mirror:
     amd:
       label: ":amd: (MI355 DPX) Language Models (Standard)"
+      automatic_shard: true
       dind: false
       device: mi355_dpx
       timeout_in_minutes: 60
@@ -90,8 +93,6 @@ steps:
       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 -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
@@ -127,11 +128,6 @@ steps:
       - "!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) [Nightly] Language Models (Extended Generation)"
@@ -159,12 +155,11 @@ steps:
       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 '(not core_model) and (not hybrid_model)'
 
 - label: ":nvidia: (H200 MIG 18GB) [Nightly] Language Models (PPL)"
   key: language-models-test-ppl
+  automatic_shard: true
   timeout_in_minutes: 30
   device: h200_18gb
   optional: true
@@ -178,11 +173,13 @@ steps:
     amd:
       label: ":amd: (MI250) Language Models (PPL)"
       dind: false
+      automatic_shard: true
       device: mi250_1
       timeout_in_minutes: 65
       depends_on:
       - image-build-amd
       source_file_dependencies:
+      - vllm/models/qwen3_5/
       - vllm/model_executor/models/qwen3_5.py
       - vllm/model_executor/models/qwen3_5_mtp.py
       - vllm/model_executor/models/qwen2.py
@@ -217,10 +214,10 @@ steps:
     - pytest -v -s models/language/pooling -m 'not core_model' --timeout=900 --timeout-method=thread
   mirror:
     amd:
-      parallelism: 4
+      automatic_shard: true
       commands:
-        - "pytest -v -s models/language/pooling -m 'not core_model' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB"
-      label: ":amd: (MI355 DPX) Language Models (Extended Pooling) Shard %N"
+        - "pytest -v -s models/language/pooling -m 'not core_model'"
+      label: ":amd: (MI355 DPX) Language Models (Extended Pooling)"
       dind: false
       device: mi355_dpx
       timeout_in_minutes: 95
diff --git a/.buildkite/test_areas/models_multimodal.yaml b/.buildkite/test_areas/models_multimodal.yaml
index 51cc3ef602ed..c47651641055 100644
--- a/.buildkite/test_areas/models_multimodal.yaml
+++ b/.buildkite/test_areas/models_multimodal.yaml
@@ -4,6 +4,7 @@ depends_on:
 steps:
 - label: ":nvidia: (H200 MIG 18GB) Multimodal Models (Standard) 1: qwen2"
   key: multi-modal-models-standard-1-qwen2
+  automatic_shard: true
   timeout_in_minutes: 68
   device: h200_18gb
   source_file_dependencies:
@@ -17,6 +18,7 @@ steps:
     amd:
       label: ":amd: (MI355 DPX) Multimodal Models (Standard) 1: qwen2"
       dind: false
+      automatic_shard: true
       device: mi355_dpx
       timeout_in_minutes: 65
       depends_on:
@@ -38,6 +40,7 @@ steps:
   mirror:
     amd:
       label: ":amd: (MI355 DPX) Multimodal Models (Standard) 2: qwen3 + gemma"
+      automatic_shard: true
       dind: false
       device: mi355_dpx
       timeout_in_minutes: 55
@@ -47,6 +50,7 @@ steps:
 - label: ":nvidia: (H200 MIG 35GB) Multimodal Models (Standard) 3: llava + qwen2_vl"
   device: h200_35gb
   key: multi-modal-models-standard-3-llava-qwen2-vl
+  automatic_shard: true
   timeout_in_minutes: 40
   source_file_dependencies:
   - vllm/
@@ -57,6 +61,7 @@ steps:
     - pytest -v -s models/multimodal/generation/test_qwen2_vl.py -m core_model
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI250) Multimodal Models (Standard) 3: llava + qwen2_vl"
       dind: false
       device: mi250_1
@@ -144,6 +149,7 @@ steps:
 - label: ":nvidia: (H200 MIG 35GB) Multimodal Accuracy Eval (Small Models)"
   device: h200_35gb
   key: multi-modal-accuracy-eval-small-models
+  automatic_shard: true
   timeout_in_minutes: 30
   working_dir: "/vllm-workspace/.buildkite/lm-eval-harness"
   source_file_dependencies:
@@ -154,6 +160,7 @@ steps:
   - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-mm-small.txt --tp-size=1
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355 DPX) Multimodal Accuracy Eval (Small Models)"
       dind: false
       device: mi355_dpx
@@ -247,6 +254,7 @@ steps:
     amd:
       label: ":amd: (MI355 DPX) Multimodal Models (Extended Generation 3)"
       dind: false
+      automatic_shard: true
       device: mi355_dpx
       timeout_in_minutes: 90
       depends_on:
diff --git a/.buildkite/test_areas/pytorch.yaml b/.buildkite/test_areas/pytorch.yaml
index 2d16e162e457..33638dd04f27 100644
--- a/.buildkite/test_areas/pytorch.yaml
+++ b/.buildkite/test_areas/pytorch.yaml
@@ -57,12 +57,14 @@ steps:
 - label: ":nvidia: (H200 MIG 35GB) PyTorch Compilation / Dynamic Shapes"
   device: h200_35gb
   key: pytorch-compilation-dynamic-shapes
+  automatic_shard: true
   timeout_in_minutes: 30
   source_file_dependencies: *pytorch-compilation-deps
   commands:
   - pytest -s -v compile/dynamic_shapes/
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI300) PyTorch Compilation / Dynamic Shapes"
       dind: false
       device: mi300_1
@@ -75,6 +77,7 @@ steps:
 
 - label: ":nvidia: (H200 MIG 18GB) PyTorch Compilation H100 Cases"
   key: pytorch-compilation-unit-tests-h100
+  automatic_shard: true
   timeout_in_minutes: 30
   device: h200_18gb
   source_file_dependencies:
@@ -108,6 +111,7 @@ steps:
   - "find compile/h100/ -name 'test_*.py' -print0 | xargs -0 -n1 -I{} pytest -s -v '{}'"
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355 DPX) PyTorch Compilation H100 Cases"
       dind: false
       device: mi355_dpx
@@ -182,6 +186,7 @@ steps:
 - label: ":nvidia: (H200 MIG 35GB) PyTorch Fullgraph"
   device: h200_35gb
   key: pytorch-fullgraph-test
+  automatic_shard: true
   timeout_in_minutes: 90
   source_file_dependencies:
   - vllm/__init__.py
@@ -212,13 +217,11 @@ steps:
   - 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 '{}'"
+  # The heavy cudagraph test is covered in other steps.
+  - pytest -v -s compile/fullgraph --ignore=compile/fullgraph/test_full_cudagraph.py
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355 DPX) PyTorch Fullgraph"
       dind: false
       device: mi355_dpx
@@ -235,6 +238,7 @@ steps:
 # those SM90 cases skip while the architecture-compatible cases still run.
 - label: ":nvidia: (L4) PyTorch Fullgraph CUDAGraph Compatibility"
   key: pytorch-fullgraph-cudagraph-l4-compatibility
+  automatic_shard: true
   device: l4
   num_devices: 1
   timeout_in_minutes: 60
@@ -269,6 +273,7 @@ steps:
   - pytest -s -v compile/fullgraph/test_full_cudagraph.py
   mirror:
     amd:
+      automatic_shard: true
       label: ':amd: (MI250) PyTorch Fullgraph CUDAGraph Compatibility'
       dind: false
       device: mi250_1
@@ -289,6 +294,7 @@ steps:
 
 - label: ":nvidia: (H200 MIG 18GB) PyTorch Nightly Dependency Override Check"
   key: pytorch-nightly-dependency-override-check
+  automatic_shard: true
   # if this test fails, it means the nightly torch version is not compatible with some
   # of the dependencies. Please check the error message and add the package to whitelist
   # in /vllm/tools/pre_commit/generate_nightly_torch_test.py
@@ -300,6 +306,7 @@ steps:
   - bash standalone_tests/pytorch_nightly_dependency.sh
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI300) PyTorch Nightly Dependency Override Check"
       dind: false
       device: mi300_1
diff --git a/.buildkite/test_areas/quantization.yaml b/.buildkite/test_areas/quantization.yaml
index 1d066c21f0a2..9349c767be9f 100644
--- a/.buildkite/test_areas/quantization.yaml
+++ b/.buildkite/test_areas/quantization.yaml
@@ -56,6 +56,7 @@ steps:
 - label: ":nvidia: (H200 MIG 35GB) Quantized Fusions"
   device: h200_35gb
   key: quantized-fusions
+  automatic_shard: true
   timeout_in_minutes: 20
   source_file_dependencies:
   - tests/fusion
@@ -67,6 +68,7 @@ steps:
     - pytest -v -s fusion/
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI300) Quantized Fusions"
       dind: false
       device: mi300_1
@@ -90,6 +92,9 @@ steps:
   device: b200-k8s
   source_file_dependencies:
   - tests/quantization/test_blackwell_moe.py
+  - tests/benchmarks/test_forward_cli.py
+  - vllm/benchmarks/forward/
+  - vllm/models/muse_spark/
   - vllm/model_executor/models/deepseek_v2.py
   - vllm/model_executor/models/gpt_oss.py
   - vllm/model_executor/models/llama4.py
@@ -99,9 +104,13 @@ steps:
   - vllm/model_executor/layers/quantization/mxfp4.py
   - vllm/v1/attention/backends/flashinfer.py
   commands:
-    - pytest -s -v tests/quantization/test_blackwell_moe.py
+    - >-
+      pytest -v -s
+      tests/quantization/test_blackwell_moe.py
+      tests/benchmarks/test_forward_cli.py::test_bench_forward_muse_balanced_routing
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355) Quantized MoE"
       dind: false
       device: mi355_1
diff --git a/.buildkite/test_areas/rust_frontend.yaml b/.buildkite/test_areas/rust_frontend.yaml
index 4e978d63d0b1..043429033932 100644
--- a/.buildkite/test_areas/rust_frontend.yaml
+++ b/.buildkite/test_areas/rust_frontend.yaml
@@ -156,6 +156,7 @@ steps:
 
 - label: ":nvidia: (H200 MIG 18GB) Rust Frontend Core Correctness"
   key: rust-frontend-core-correctness
+  automatic_shard: true
   timeout_in_minutes: 20
   device: h200_18gb
   working_dir: "/vllm-workspace/tests"
@@ -170,6 +171,7 @@ steps:
   - pytest -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355 DPX) Rust Frontend Core Correctness"
       dind: false
       device: mi355_dpx
@@ -185,6 +187,7 @@ steps:
 
 - label: ":nvidia: (H200 MIG 35GB) Rust Frontend Tool Use"
   key: rust-frontend-tool-use
+  automatic_shard: true
   device: h200_35gb
   timeout_in_minutes: 35
   working_dir: "/vllm-workspace/tests"
@@ -213,6 +216,7 @@ steps:
   - pytest -v -s entrypoints/tool_parsers/test_hermes_tool_parser.py -k granite4
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355 DPX) Rust Frontend Tool Use"
       dind: false
       device: mi355_dpx
diff --git a/.buildkite/test_areas/samplers.yaml b/.buildkite/test_areas/samplers.yaml
index 9d66e785ea8a..68055d6a0fd7 100644
--- a/.buildkite/test_areas/samplers.yaml
+++ b/.buildkite/test_areas/samplers.yaml
@@ -5,6 +5,7 @@ steps:
 - label: ":nvidia: (H200 MIG 18GB) Samplers"
   device: h200_18gb
   key: samplers-test
+  automatic_shard: true
   timeout_in_minutes: 60
   source_file_dependencies:
   - vllm/model_executor/layers
@@ -19,6 +20,7 @@ steps:
     - VLLM_USE_FLASHINFER_SAMPLER=1 pytest -v -s samplers -k "not test_beam_search_passes_multimodal_data"
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI250) Samplers"
       dind: false
       device: mi250_1
@@ -41,6 +43,7 @@ steps:
 - label: ":nvidia: (H200 MIG 35GB) Samplers Multimodal Beam Search"
   device: h200_35gb
   key: samplers-multimodal-beam-search
+  automatic_shard: true
   timeout_in_minutes: 30
   source_file_dependencies:
   - vllm/model_executor/layers
@@ -52,6 +55,7 @@ steps:
     - VLLM_USE_FLASHINFER_SAMPLER=1 pytest -v -s samplers/test_beam_search.py::test_beam_search_passes_multimodal_data
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355 DPX) Samplers Multimodal Beam Search"
       dind: false
       device: mi355_dpx
diff --git a/.buildkite/test_areas/spec_decode.yaml b/.buildkite/test_areas/spec_decode.yaml
index 61b6677ae586..32d377ca5efa 100644
--- a/.buildkite/test_areas/spec_decode.yaml
+++ b/.buildkite/test_areas/spec_decode.yaml
@@ -36,6 +36,7 @@ steps:
 
 - label: ":nvidia: (H200 MIG 35GB) Spec Decode Eagle 1: DeepSeek + Qwen"
   key: spec-decode-eagle-1-deepseek-qwen
+  automatic_shard: true
   timeout_in_minutes: 25
   device: h200_35gb
   source_file_dependencies:
@@ -46,6 +47,7 @@ steps:
     - pytest -v -s v1/e2e/spec_decode/eagle/ -k "deepseek_eagle or qwen3_eagle3"
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355 DPX) Spec Decode Eagle 1: DeepSeek + Qwen"
       dind: false
       device: mi355_dpx
@@ -63,6 +65,7 @@ steps:
 
 - label: ":nvidia: (H200 MIG 35GB) Spec Decode Eagle 2: Llama 3 + Qwen VL + Other"
   key: spec-decode-eagle-2-llama3-qwen-vl-other
+  automatic_shard: true
   timeout_in_minutes: 25
   device: h200_35gb
   source_file_dependencies:
@@ -73,6 +76,7 @@ steps:
     - pytest -v -s v1/e2e/spec_decode/eagle/ -k "not deepseek_eagle and not qwen3_eagle3"
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355 DPX) Spec Decode Eagle 2: Llama 3 + Qwen VL + Other"
       dind: false
       device: mi355_dpx
@@ -127,6 +131,7 @@ steps:
 # each job below runs whole directories.
 - label: ":nvidia: (H200 MIG 35GB) Spec Decode Speculators"
   key: spec-decode-speculators
+  automatic_shard: true
   timeout_in_minutes: 30
   device: h200_35gb
   source_file_dependencies:
@@ -139,6 +144,7 @@ steps:
     - pytest -v -s v1/e2e/spec_decode/speculators/
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355 DPX) Spec Decode Speculators"
       dind: false
       device: mi355_dpx
@@ -157,6 +163,7 @@ steps:
 
 - label: ":nvidia: (H200 MIG 35GB) Spec Decode MTP Gemma4"
   key: spec-decode-mtp-gemma4
+  automatic_shard: true
   timeout_in_minutes: 30
   device: h200_35gb
   source_file_dependencies:
@@ -169,6 +176,7 @@ steps:
     - pytest -v -s v1/e2e/spec_decode/mtp/gemma4/
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355 DPX) Spec Decode MTP Gemma4"
       dind: false
       device: mi355_dpx
@@ -187,6 +195,7 @@ steps:
 
 - label: ":nvidia: (H200 MIG 35GB) Spec Decode MTP DeepSeek + MiMo"
   key: spec-decode-mtp-deepseek-mimo
+  automatic_shard: true
   timeout_in_minutes: 30
   device: h200_35gb
   source_file_dependencies:
@@ -200,6 +209,7 @@ steps:
     - pytest -v -s v1/e2e/spec_decode/mtp/mimo/
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355 DPX) Spec Decode MTP DeepSeek + MiMo"
       dind: false
       device: mi355_dpx
@@ -218,6 +228,7 @@ steps:
 
 - label: ":nvidia: (H200 MIG 35GB) Spec Decode MTP Qwen3.5"
   key: spec-decode-mtp-qwen3-5
+  automatic_shard: true
   timeout_in_minutes: 30
   device: h200_35gb
   source_file_dependencies:
@@ -225,11 +236,16 @@ steps:
     - vllm/v1/worker/gpu/spec_decode/
     - vllm/v1/attention/backends/
     - vllm/transformers_utils/configs/speculators/
+    - vllm/model_executor/models/qwen3_5_mtp.py
+    - vllm/models/qwen3_5/nvidia/mtp.py
+    - vllm/model_executor/models/qwen3_next_mtp.py
+    - vllm/config/speculative.py
     - tests/v1/e2e/spec_decode/
   commands:
     - pytest -v -s v1/e2e/spec_decode/mtp/qwen3_5/
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355 DPX) Spec Decode MTP Qwen3.5"
       dind: false
       device: mi355_dpx
@@ -243,6 +259,9 @@ steps:
       - vllm/v1/sample/
       - vllm/model_executor/layers/
       - vllm/transformers_utils/configs/speculators/
+      - vllm/model_executor/models/qwen3_5_mtp.py
+      - vllm/model_executor/models/qwen3_next_mtp.py
+      - vllm/config/speculative.py
       - tests/v1/e2e/spec_decode/
       - vllm/platforms/rocm.py
 
@@ -286,6 +305,7 @@ steps:
 
 - label: ":nvidia: (H200 MIG 35GB) Spec Decode N-Gram + Suffix"
   key: spec-decode-ngram-suffix
+  automatic_shard: true
   timeout_in_minutes: 20
   device: h200_35gb
   source_file_dependencies:
@@ -298,6 +318,7 @@ steps:
     - python3 spec_decode/test_custom_proposer.py
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355 DPX) Spec Decode N-Gram + Suffix"
       dind: false
       device: mi355_dpx
@@ -352,8 +373,11 @@ steps:
     - vllm/v1/spec_decode/
     - vllm/v1/worker/gpu/spec_decode/
     - tests/v1/e2e/spec_decode/
+    - vllm/model_executor/models/qwen3_dflash.py
+    - tests/models/test_qwen3_dflash_mxfp8_context_kv.py
   commands:
     - pytest -v -s v1/e2e/spec_decode/draft_model/
+    - pytest -v -s models/test_qwen3_dflash_mxfp8_context_kv.py
   mirror:
     amd:
       label: ":amd: (MI355) Spec Decode Draft Model Nightly"
@@ -378,6 +402,7 @@ steps:
 
 - label: ":nvidia: (H100) [Nightly] Speculators Correctness"
   key: speculators-correctness
+  automatic_shard: true
   timeout_in_minutes: 30
   device: h100
   optional: true
@@ -392,6 +417,7 @@ steps:
     - pytest -v -s v1/spec_decode/test_speculators_correctness.py -m slow_test
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355 DPX) Speculators Correctness Nightly"
       dind: false
       device: mi355_dpx
@@ -404,6 +430,7 @@ steps:
 
 - label: ":nvidia: (B200) [Nightly] Spec Decode DeepSeek MTP Parallel Load"
   key: spec-decode-deepseek-mtp-parallel-load-2xb200-2xmi300
+  automatic_shard: true
   timeout_in_minutes: 30
   device: b200-k8s
   optional: true
@@ -419,6 +446,7 @@ steps:
     - pytest -v -s v1/e2e/spec_decode/test_mtp_parallel_load.py
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355) Spec Decode DeepSeek MTP Parallel Load"
       dind: false
       device: mi355_2
@@ -450,12 +478,15 @@ steps:
     - vllm/v1/worker/gpu/spec_decode/
     - vllm/model_executor/models/qwen3_dflash.py
     - vllm/model_executor/models/qwen3_dflash2.py
+    - vllm/models/qwen3_5/
     - vllm/model_executor/models/qwen3_5.py
+    - vllm/model_executor/models/qwen3_5_mtp.py
     - vllm/model_executor/models/laguna_dflash.py
     - tests/v1/e2e/spec_decode/
   commands:
     # DFlash2 NVFP4 (27B) does not fit a 35GB MIG slice; runs on B200 below.
     - pytest -v -s v1/e2e/spec_decode/acceptance_rates/dflash/ -k "not dflash2"
+    - pytest -v -s v1/e2e/spec_decode/test_distribution_shift.py
   mirror:
     amd:
       label: ":amd: (MI355 DPX) Spec Decode AL DFlash Nightly"
@@ -465,6 +496,7 @@ steps:
       timeout_in_minutes: 120
       commands:
       - pytest -v -s v1/e2e/spec_decode/acceptance_rates/dflash/ -k "not dflash2"
+      - pytest -v -s v1/e2e/spec_decode/test_distribution_shift.py
       depends_on:
       - image-build-amd
       source_file_dependencies:
@@ -484,6 +516,7 @@ steps:
 
 - label: ":nvidia: (B200) [Nightly] Spec Decode AL DFlash2"
   key: spec-decode-dflash2-nightly-b200
+  automatic_shard: true
   timeout_in_minutes: 60
   device: b200-k8s
   optional: true
@@ -492,12 +525,15 @@ steps:
     - vllm/v1/worker/gpu/spec_decode/
     - vllm/model_executor/models/qwen3_dflash.py
     - vllm/model_executor/models/qwen3_dflash2.py
+    - vllm/models/qwen3_5/
     - vllm/model_executor/models/qwen3_5.py
+    - vllm/model_executor/models/qwen3_5_mtp.py
     - tests/v1/e2e/spec_decode/
   commands:
     - pytest -v -s v1/e2e/spec_decode/acceptance_rates/dflash/ -k "dflash2"
   mirror:
     amd:
+      automatic_shard: true
       label: ":amd: (MI355) Spec Decode AL DFlash2 Nightly"
       dind: false
       device: mi355_1
@@ -517,7 +553,9 @@ steps:
       - vllm/v1/attention/
       - vllm/model_executor/models/qwen3_dflash.py
       - vllm/model_executor/models/qwen3_dflash2.py
+      - vllm/models/qwen3_5/
       - vllm/model_executor/models/qwen3_5.py
+      - vllm/model_executor/models/qwen3_5_mtp.py
       - vllm/model_executor/model_loader/
       - vllm/model_executor/layers/quantization/
       - vllm/model_executor/kernels/linear/
@@ -525,6 +563,7 @@ steps:
 
 - label: ":nvidia: (H100) [Nightly] Spec Decode AL XPress"
   key: spec-decode-xpress-nightly
+  automatic_shard: true
   timeout_in_minutes: 30
   device: h100
   optional: true
diff --git a/.buildkite/test_areas/torch_abi.yaml b/.buildkite/test_areas/torch_abi.yaml
index fb2ff09cdbf4..4c38ba0ff2e3 100644
--- a/.buildkite/test_areas/torch_abi.yaml
+++ b/.buildkite/test_areas/torch_abi.yaml
@@ -4,6 +4,7 @@ depends_on:
 steps:
 - label: ":nvidia: (L4) Torch Stable ABI Audit"
   key: torch-stable-abi-audit
+  automatic_shard: true
   timeout_in_minutes: 10
   source_file_dependencies:
   - .buildkite/check-torch-abi.py
diff --git a/.dockerignore b/.dockerignore
index e0284129eefa..7ed1895898df 100644
--- a/.dockerignore
+++ b/.dockerignore
@@ -1,8 +1,9 @@
 /.venv
 /build
-/artifacts/vllm-rocm-install/
-/bake-config-build-*.json
-/wheel-export/
+/.rocm-version/
+/.rocm-wheel/
+/.rocm-kernel-map/
+/.rocm-src/
 dist
 vllm/*.so
 vllm/vllm-rs
diff --git a/.github/mergify.yml b/.github/mergify.yml
index 0f545a4a0784..7ece42cdecfe 100644
--- a/.github/mergify.yml
+++ b/.github/mergify.yml
@@ -752,6 +752,28 @@ pull_request_rules:
         - "stmatengss"
         - "LCAIZJ"
 
+- name: label-nixl
+  description: Automatically apply nixl label
+  conditions:
+    - label != stale
+    - or:
+      # Any NIXL-named file: KV connector and shared utils, NIXL EP
+      # dispatch, EC-transfer NIXL data, the offload p2p data layer,
+      # tests, docs, examples.
+      - files~=(?i)nixl
+      # KV offloading paths (multi-tiered KVCR tiering and simple), plus
+      # their tests and docs.
+      - files~=kv_offload
+      - title~=(?i)nixl
+  actions:
+    label:
+      add:
+        - nixl
+    assign:
+      users:
+        - "ovidiusm"
+        - "iyastreb"
+
 - name: label-scheduler
   description: Automatically apply scheduler label
   conditions:
diff --git a/.github/workflows/buf.yml b/.github/workflows/buf.yml
index dd695d28d529..15c7ce0e1179 100644
--- a/.github/workflows/buf.yml
+++ b/.github/workflows/buf.yml
@@ -17,8 +17,8 @@ jobs:
   proto:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
-      - uses: bufbuild/buf-action@8c6a16e16f12ba20b6470afa9c2ba9b5ba8c97c3 # v1.5.0
+      - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
+      - uses: bufbuild/buf-action@85aebf73123b5c15fd5528aaecbf9129cddf7fa7 # v1.6.0
         with:
           version: 1.72.0
           input: rust/proto
diff --git a/.github/workflows/issue_autolabel.yml b/.github/workflows/issue_autolabel.yml
index 43a458df7678..03eeb71a5bda 100644
--- a/.github/workflows/issue_autolabel.yml
+++ b/.github/workflows/issue_autolabel.yml
@@ -15,7 +15,7 @@ jobs:
     steps:
       - name: Label issues based on keywords
         id: label-step
-        uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3  # v9.0.0
+        uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0
         with:
           script: |
             // Configuration: Add new labels and keywords here
@@ -633,7 +633,7 @@ jobs:
 
       - name: CC users for labeled issues
         if: steps.label-step.outputs.labels_added != '[]'
-        uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3  # v9.0.0
+        uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0
         with:
           script: |
             // Configuration: Map labels to GitHub users to CC
@@ -714,7 +714,7 @@ jobs:
 
       - name: Request missing ROCm info from issue author
         if: contains(steps.label-step.outputs.labels_added, 'rocm') && contains(toJSON(github.event.issue.labels.*.name), 'bug')
-        uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3  # v9.0.0
+        uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0
         with:
           script: |
             const body = (context.payload.issue.body || '').toLowerCase();
diff --git a/.github/workflows/macos-smoke-test.yml b/.github/workflows/macos-smoke-test.yml
index 2e63c5676dd8..8de27d848e8b 100644
--- a/.github/workflows/macos-smoke-test.yml
+++ b/.github/workflows/macos-smoke-test.yml
@@ -28,9 +28,9 @@ jobs:
     timeout-minutes: 30
 
     steps:
-      - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+      - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
 
-      - uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0
+      - uses: astral-sh/setup-uv@68559e99cc807c1eb7e9238fd43a15bba460eec2 # v10.3.3
         with:
           enable-cache: true
           cache-dependency-glob: |
diff --git a/.github/workflows/notify-ci-authorized.yml b/.github/workflows/notify-ci-authorized.yml
index c260cc67dba0..1396032d51cb 100644
--- a/.github/workflows/notify-ci-authorized.yml
+++ b/.github/workflows/notify-ci-authorized.yml
@@ -34,11 +34,11 @@ jobs:
       github.event.workflow_run.conclusion == 'success'))
     runs-on: [self-hosted, linux, x64, vllm-runners]
     steps:
-      - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+      - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
         with:
           ref: ${{ github.event.repository.default_branch }}
           persist-credentials: false
-      - uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0
+      - uses: astral-sh/setup-uv@68559e99cc807c1eb7e9238fd43a15bba460eec2 # v10.3.3
         with:
           python-version: "3.12"
       - name: Notify PR author that CI is available
diff --git a/.github/workflows/pre-commit.yml b/.github/workflows/pre-commit.yml
index 861d4a5e498b..1410a8944828 100644
--- a/.github/workflows/pre-commit.yml
+++ b/.github/workflows/pre-commit.yml
@@ -49,14 +49,18 @@ jobs:
     if: github.repository_owner == 'vllm-project' && always() && (needs.pre-run-check.result == 'success' || needs.pre-run-check.result == 'skipped')
     runs-on: [self-hosted, linux, x64, vllm-runners]
     steps:
-    - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
-    - uses: actions/setup-python@83679a892e2d95755f2dac6acb0bfd1e9ac5d548 # v6.1.0
+    - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
+    - uses: astral-sh/setup-uv@68559e99cc807c1eb7e9238fd43a15bba460eec2 # v10.3.3
       with:
         python-version: "3.12"
     - run: echo "::add-matcher::.github/workflows/matchers/actionlint.json"
     - run: echo "::add-matcher::.github/workflows/matchers/markdownlint.json"
     - run: echo "::add-matcher::.github/workflows/matchers/mypy.json"
     - run: echo "::add-matcher::.github/workflows/matchers/shellcheck.json"
-    - uses: pre-commit/action@2c7b3805fd2a0fd8c1884dcaebf91fc102a13ecd # v3.0.1
+    # Replaces pre-commit/action, whose latest release (v3.0.1) uses the node20 actions/cache@v4
+    - uses: actions/cache@55cc8345863c7cc4c66a329aec7e433d2d1c52a9 # v6.1.0
       with:
-        extra_args: --all-files --hook-stage manual
+        path: ~/.cache/pre-commit
+        key: pre-commit-uv|${{ runner.os }}|${{ env.UV_PYTHON }}|${{ hashFiles('.pre-commit-config.yaml') }}
+    - run: >-
+        uvx pre-commit run --show-diff-on-failure --color=always --all-files --hook-stage manual
diff --git a/.github/workflows/proto-crate.yml b/.github/workflows/proto-crate.yml
index 43be5c8c9214..ad20db39fa66 100644
--- a/.github/workflows/proto-crate.yml
+++ b/.github/workflows/proto-crate.yml
@@ -24,8 +24,8 @@ jobs:
   semver:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
-      - uses: obi1kenobi/cargo-semver-checks-action@6b69fcf40e9b5fb17adeb57e4b6ecd020649a239 # v2
+      - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
+      - uses: obi1kenobi/cargo-semver-checks-action@6b69fcf40e9b5fb17adeb57e4b6ecd020649a239 # v2.9
         with:
           manifest-path: rust/proto/Cargo.toml
           package: vllm-proto
@@ -33,7 +33,7 @@ jobs:
   check:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+      - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
       - run: rustup show active-toolchain
       - name: Check release version
         if: github.event_name == 'push'
@@ -56,13 +56,13 @@ jobs:
       contents: read
       id-token: write
     steps:
-      - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+      - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
         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
+      - uses: rust-lang/crates-io-auth-action@3268eb413133e2960fa4ae0e963ca76380c54cd9 # v1.0.6
         id: auth
       - name: Publish
         env:
diff --git a/.github/workflows/recipes-preview.yml b/.github/workflows/recipes-preview.yml
index 42127e99b79e..efaca4681a67 100644
--- a/.github/workflows/recipes-preview.yml
+++ b/.github/workflows/recipes-preview.yml
@@ -46,7 +46,7 @@ jobs:
     steps:
       - name: Create a token for vllm-project/recipes
         id: token
-        uses: actions/create-github-app-token@bcd2ba49218906704ab6c1aa796996da409d3eb1 # v3.2.0
+        uses: actions/create-github-app-token@64ab729a6294000f5aa4a4bc0b73d1f32bc45510 # v3.3.0
         with:
           app-id: ${{ vars.RECIPES_APP_ID }}
           private-key: ${{ secrets.RECIPES_APP_PRIVATE_KEY }}
diff --git a/.github/workflows/run-ci-command.yml b/.github/workflows/run-ci-command.yml
index d3336772717f..877c2d91b94b 100644
--- a/.github/workflows/run-ci-command.yml
+++ b/.github/workflows/run-ci-command.yml
@@ -22,11 +22,11 @@ jobs:
       startsWith(github.event.comment.body, '/amd-ci'))
     runs-on: [self-hosted, linux, x64, vllm-runners]
     steps:
-      - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+      - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
         with:
           ref: ${{ github.event.repository.default_branch }}
           persist-credentials: false
-      - uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0
+      - uses: astral-sh/setup-uv@68559e99cc807c1eb7e9238fd43a15bba460eec2 # v10.3.3
         with:
           python-version: "3.12"
       - name: Authorize and run CI command
diff --git a/.github/workflows/scripts/run_ci_command.py b/.github/workflows/scripts/run_ci_command.py
index 64ec8cb7623b..f249c08ce5c2 100644
--- a/.github/workflows/scripts/run_ci_command.py
+++ b/.github/workflows/scripts/run_ci_command.py
@@ -76,11 +76,9 @@
 }
 RETRY_STATES = "failed,timed_out,expired"
 CANCELABLE_BUILD_STATES = ("scheduled", "running", "failing")
-SETUP_STEP_KEYS = {
-    "ensure-ci-base-amd",
-    "pre-commit",
-    "refresh-rocm-base-amd",
-}
+# ensure-ci-base-amd: the old ROCm build chain, which ci-infra still renders
+# for branches that predate .buildkite/scripts/rocm/build.sh.
+SETUP_STEP_KEYS = {"ensure-ci-base-amd", "refresh-rocm-base-amd", "pre-commit"}
 
 
 class ApiError(RuntimeError):
diff --git a/.github/workflows/stale.yml b/.github/workflows/stale.yml
index ba807fab7c3a..d2c5533d31f6 100644
--- a/.github/workflows/stale.yml
+++ b/.github/workflows/stale.yml
@@ -15,7 +15,7 @@ jobs:
       actions: write
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/stale@eb5cf3af3ac0a1aa4c9c45633dd1ae542a27a899 # v10.3.0
+      - uses: actions/stale@4391f3da665fdf50b6810c1a66712fb9ba21aa93 # v11.0.0
         with:
           # Increasing this value ensures that changes to this workflow
           # propagate to all issues and PRs in days rather than months
diff --git a/.gitignore b/.gitignore
index 57d0d54fc87b..533f9cfa4921 100644
--- a/.gitignore
+++ b/.gitignore
@@ -251,3 +251,9 @@ vllm/grpc/vllm_engine_pb2.pyi
 
 # Ignore generated cpu headers 
 csrc/cpu/cpu_attn_dispatch_generated.h
+
+# Written by .buildkite/scripts/rocm/build.sh
+/.rocm-src/
+/.rocm-version/
+/.rocm-wheel/
+/.rocm-kernel-map/
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index c940db4ae844..b42d38d93a31 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -5,42 +5,43 @@ default_stages:
   - pre-commit # Run locally
   - manual # Run in CI
 exclude: 'vllm/third_party/.*|vllm/models/kimi_k3/nvidia/ops/third_party/.*|vllm/models/kimi_k3/amd/ops/third_party/.*'
+# Revs are frozen to commit SHAs, bump them with `pre-commit autoupdate --freeze [--repo <url>]`
 repos:
 - repo: https://github.com/astral-sh/ruff-pre-commit
-  rev: v0.14.0
+  rev: f9351c924055bf6c7b4a4670237d3ce141e9f57c  # frozen: v0.14.0
   hooks:
   - id: ruff-check
     args: [--output-format, github, --fix]
   - id: ruff-format
 - repo: https://github.com/crate-ci/typos
-  rev: v1.43.5
+  rev: 9023c1f779400877cafdd98052113fc777e2cfdc  # frozen: v1.51.1
   hooks:
   - id: typos
     args: [--force-exclude]
 - repo: https://github.com/pre-commit/mirrors-clang-format
-  rev: v21.1.2
+  rev: 98ad4ddb9c057de2f5b9e93e44e6db959807d68c  # frozen: v23.1.3
   hooks:
   - id: clang-format
     exclude: 'csrc/moe/topk_softmax_kernels.cu|vllm/third_party/.*'
     types_or: [c++, cuda]
     args: [--style=file, --verbose]
 - repo: https://github.com/DavidAnson/markdownlint-cli2
-  rev: v0.21.0
+  rev: 341e2a331ba578029e66c2e01183e7a6b5ca9ab4  # frozen: v0.24.0
   hooks:
   - id: markdownlint-cli2
-    language_version: lts
+    language_version: 24.22.1 # Node 26 (current lts) needs libatomic, which the CI runners lack
     args: [--config, tools/pre_commit/.markdownlint.yaml, --fix]
 - repo: https://github.com/rhysd/actionlint
-  rev: v1.7.7
+  rev: 914e7df21a07ef503a81201c76d2b11c789d3fca  # frozen: v1.7.12
   hooks:
   - id: actionlint
 - repo: https://github.com/shellcheck-py/shellcheck-py
-  rev: 745eface02aef23e168a8afb6b5737818efbea95 # v0.11.0.1
+  rev: a6564f5d57444e123433a5c9af733e3975c77229  # frozen: v0.11.0.1-1
   hooks:
   - id: shellcheck
     args: ["-s", "bash", "--format=gcc", "--rcfile=tools/pre_commit/.shellcheckrc"]
 - repo: https://github.com/astral-sh/uv-pre-commit
-  rev: 0.11.1
+  rev: 46787482b9fb2003c56e0f7592e312ef381829d9  # frozen: 0.13.6
   hooks:
     - id: pip-compile
       args: [
@@ -74,6 +75,7 @@ repos:
         --no-emit-package, cupy-cuda12x,
         # nvidia packages (unsuffixed / unified naming)
         --no-emit-package, nvidia-cublas,
+        --no-emit-package, nvidia-cuda-cccl,
         --no-emit-package, nvidia-cuda-cupti,
         --no-emit-package, nvidia-cuda-nvrtc,
         --no-emit-package, nvidia-cuda-runtime,
@@ -158,7 +160,7 @@ repos:
       ]
       files: ^requirements/docs\.(in|txt)$
 - repo: https://github.com/pre-commit/pre-commit-hooks
-  rev: 3e8a8703264a2f4a69428a0aa4dcb512790b2c8c
+  rev: 3e8a8703264a2f4a69428a0aa4dcb512790b2c8c  # frozen: v6.0.0
   hooks:
   - id: check-json
     files: ^\.buildkite/.*\.json$
diff --git a/CMakeLists.txt b/CMakeLists.txt
index 7694e6f8bc46..32f95b697069 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -1232,6 +1232,25 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
     message(STATUS "Building hadacore")
   endif()
 
+  # Fused MQA logits + top-k prefill kernel (SM100 family). It compiles against
+  # DeepGEMM's headers, so DeepGEMM is fetched here.
+  include(cmake/external_projects/deepgemm.cmake)
+  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9)
+    cuda_archs_loose_intersection(FUSED_MQA_TOPK_ARCHS "10.0f" "${CUDA_ARCHS}")
+  endif()
+  if(FUSED_MQA_TOPK_ARCHS)
+    set(FUSED_MQA_TOPK_SRC "csrc/fused_mqa_topk/fused_mqa_topk.cu")
+    set_gencode_flags_for_srcs(
+      SRCS "${FUSED_MQA_TOPK_SRC}"
+      CUDA_ARCHS "${FUSED_MQA_TOPK_ARCHS}")
+    set_property(SOURCE ${FUSED_MQA_TOPK_SRC} APPEND PROPERTY INCLUDE_DIRECTORIES
+      "${deepgemm_SOURCE_DIR}/deep_gemm/include"
+      "${deepgemm_SOURCE_DIR}/third-party/cutlass/include")
+    list(APPEND VLLM_STABLE_EXT_SRC "${FUSED_MQA_TOPK_SRC}")
+    message(STATUS
+      "Building fused MQA top-k for archs: ${FUSED_MQA_TOPK_ARCHS}")
+  endif()
+
   # if CUDA endif
   endif()
 
@@ -1296,6 +1315,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
       target_compile_definitions(_C_stable_libtorch PRIVATE
         VLLM_ENABLE_FUSED_KDA_DECODE=1)
     endif()
+    if(FUSED_MQA_TOPK_ARCHS)
+      target_compile_definitions(_C_stable_libtorch PRIVATE
+        VLLM_ENABLE_FUSED_MQA_TOPK=1)
+    endif()
     if(FUSED_GDN_DECODE_ARCHS)
       target_compile_definitions(_C_stable_libtorch PRIVATE
         VLLM_ENABLE_FUSED_GDN_DECODE=1)
@@ -1372,6 +1395,19 @@ set_gencode_flags_for_srcs(
   SRCS "${VLLM_MOE_EXT_SRC}"
   CUDA_ARCHS "${CUDA_ARCHS}")
 
+if(VLLM_GPU_LANG STREQUAL "CUDA")
+  set(MUSE_SPARK_ROUTING_ARCHS "${CUDA_ARCHS}")
+  if(CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 13.0)
+    # Native FP32 warp-max reduction requires the Blackwell family target.
+    list(TRANSFORM MUSE_SPARK_ROUTING_ARCHS
+      REPLACE "^10[.]0([+]PTX)?$" "10.0f")
+  endif()
+  set_gencode_flags_for_srcs(
+    SRCS "csrc/moe/muse_spark_routing.cu"
+    CUDA_ARCHS "${MUSE_SPARK_ROUTING_ARCHS}")
+  list(APPEND VLLM_MOE_EXT_SRC "csrc/moe/muse_spark_routing.cu")
+endif()
+
 if(VLLM_GPU_LANG STREQUAL "CUDA")
   set(VLLM_MOE_WNA16_SRC
     "csrc/moe/moe_wna16.cu")
@@ -1550,61 +1586,101 @@ if(VLLM_GPU_LANG STREQUAL "HIP")
   #
   # _rocm_C extension
   #
-  # The paged attention *instantiations* are split over one translation unit
-  # per (dtype, kv-cache dtype, head size) combination: compiling them all in a
-  # single file serializes ~10 min of device code generation onto one core.
-  # attention_dispatch.cu only holds the runtime dispatch, attention_kernels.cuh
-  # the kernels and launchers. A file listed here is dead weight unless
-  # attention_dispatch.cu also dispatches its entry point (guarded below).
+  # The paged attention instantiations are split over one generated translation
+  # unit per (dtype, kv-cache dtype, head size, block size); compiling them in a
+  # single file serialized ~10 min of device code generation onto one core. The
+  # same script emits the entry point prototypes and paged_attention()'s dispatch
+  # arms, so a unit cannot be compiled without being dispatched.
+  # attention_dispatch.cu only holds argument validation and is_navi detection,
+  # attention_kernels.cuh the kernels and launchers.
+  #
+  set(ROCM_ATTN_GEN_SCRIPT
+    ${CMAKE_CURRENT_SOURCE_DIR}/csrc/rocm/generate_paged_attn.py)
+  # Editing the script has to re-run configure, otherwise the units on disk stay
+  # stale until something else triggers a reconfigure.
+  set_property(DIRECTORY APPEND PROPERTY CMAKE_CONFIGURE_DEPENDS
+    ${ROCM_ATTN_GEN_SCRIPT})
+  file(MD5 ${ROCM_ATTN_GEN_SCRIPT} ROCM_ATTN_GEN_SCRIPT_HASH)
+  # The script records this key plus the exact list of files it wrote, in a stamp
+  # next to the outputs; the list is what CMake uses as the source list, so a
+  # stray `attention_paged_attn_*.cu` or a partially deleted set cannot slip in.
+  # The stamp lives in the source tree (not in a per-build-directory cache)
+  # because the outputs do. Whether the Navi kernel set is instantiated is a
+  # compile definition, not part of the generated text, so the units are identical
+  # for every build directory and arch list.
+  set(ROCM_ATTN_GEN_STAMP
+    ${CMAKE_CURRENT_SOURCE_DIR}/csrc/rocm/.paged_attn_generated)
+
+  # Sets OUT_SRCS to the .cu files the stamp lists, or to nothing when the stamp
+  # is missing, written by another version of the script, or lists a missing file
+  # (all of which mean "regenerate").
+  function (vllm_rocm_attn_generated_srcs out_srcs)
+    set(${out_srcs} "" PARENT_SCOPE)
+    if (NOT EXISTS ${ROCM_ATTN_GEN_STAMP})
+      return()
+    endif()
+    file(STRINGS ${ROCM_ATTN_GEN_STAMP} _lines)
+    if (NOT _lines)
+      return()
+    endif()
+    list(GET _lines 0 _key)
+    list(REMOVE_AT _lines 0)
+    if (NOT _key STREQUAL "${ROCM_ATTN_GEN_SCRIPT_HASH}")
+      return()
+    endif()
+    set(_srcs)
+    foreach (_f ${_lines})
+      if (NOT EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/csrc/rocm/${_f})
+        return()
+      endif()
+      if (_f MATCHES "\\.cu$")
+        list(APPEND _srcs csrc/rocm/${_f})
+      endif()
+    endforeach()
+    set(${out_srcs} ${_srcs} PARENT_SCOPE)
+  endfunction()
+
+  vllm_rocm_attn_generated_srcs(VLLM_ATTENTION_PAGED_ATTN_SRC)
+  if (NOT VLLM_ATTENTION_PAGED_ATTN_SRC)
+    execute_process(
+      COMMAND ${Python_EXECUTABLE} ${ROCM_ATTN_GEN_SCRIPT}
+              ${ROCM_ATTN_GEN_SCRIPT_HASH}
+      RESULT_VARIABLE rocm_attn_generation_result
+      OUTPUT_FILE ${CMAKE_CURRENT_BINARY_DIR}/rocm_attn_generation.log
+      ERROR_FILE ${CMAKE_CURRENT_BINARY_DIR}/rocm_attn_generation.log
+    )
+
+    if (NOT rocm_attn_generation_result EQUAL 0)
+      message(FATAL_ERROR "Paged attention generation failed."
+                          " Result: \"${rocm_attn_generation_result}\""
+                          "\nCheck the log for details: "
+                          "${CMAKE_CURRENT_BINARY_DIR}/rocm_attn_generation.log")
+    endif()
+    vllm_rocm_attn_generated_srcs(VLLM_ATTENTION_PAGED_ATTN_SRC)
+    if (NOT VLLM_ATTENTION_PAGED_ATTN_SRC)
+      message(FATAL_ERROR "Paged attention generation produced no sources; see "
+                          "${CMAKE_CURRENT_BINARY_DIR}/rocm_attn_generation.log")
+    endif()
+    message(STATUS "Paged attention generation completed successfully.")
+  else()
+    message(STATUS "Paged attention generation is up to date.")
+  endif()
+
+  # Configure deps too, so a plain `ninja` on an out-of-tree build dir regenerates
+  # a missing unit (e.g. after `git clean -X`) instead of failing on a missing
+  # source; see docs/contributing/incremental_build.md. pip install reruns anyway.
+  list(TRANSFORM VLLM_ATTENTION_PAGED_ATTN_SRC
+    PREPEND ${CMAKE_CURRENT_SOURCE_DIR}/ OUTPUT_VARIABLE _attn_generated_abs)
+  set_property(DIRECTORY APPEND PROPERTY CMAKE_CONFIGURE_DEPENDS
+    ${ROCM_ATTN_GEN_STAMP} ${_attn_generated_abs})
+
   set(VLLM_ROCM_EXT_SRC
     "csrc/rocm/torch_bindings.cpp"
     "csrc/rocm/skinny_gemms.cu"
     "csrc/rocm/skinny_gemms_int4.cu"
     "csrc/rocm/attention_dispatch.cu"
-    "csrc/rocm/attention_paged_attn_h64_f16_auto.cu"
-    "csrc/rocm/attention_paged_attn_h128_f16_auto.cu"
-    "csrc/rocm/attention_paged_attn_h64_bf16_auto.cu"
-    "csrc/rocm/attention_paged_attn_h128_bf16_auto.cu"
-    "csrc/rocm/attention_paged_attn_h64_f16_fp8_f16.cu"
-    "csrc/rocm/attention_paged_attn_h128_f16_fp8_f16.cu"
-    "csrc/rocm/attention_paged_attn_h64_f16_fp8_fp8.cu"
-    "csrc/rocm/attention_paged_attn_h128_f16_fp8_fp8.cu"
-    "csrc/rocm/attention_paged_attn_h64_bf16_fp8_f16.cu"
-    "csrc/rocm/attention_paged_attn_h128_bf16_fp8_f16.cu"
-    "csrc/rocm/attention_paged_attn_h64_bf16_fp8_fp8.cu"
-    "csrc/rocm/attention_paged_attn_h128_bf16_fp8_fp8.cu")
-
-  # Each attention_paged_attn_*.cu file below instantiates one kernel set and
-  # exposes it as one entry point. That entry point is only reachable if
-  # paged_attention() in csrc/rocm/attention_dispatch.cu names it in one of its
-  # ATTN_HEAD_SWITCH() calls (the dispatch arms); otherwise the file is compiled
-  # and linked unnecessarily. Fail at configure time instead.
-  file(READ "${CMAKE_SOURCE_DIR}/csrc/rocm/attention_dispatch.cu" _attn_dispatch)
-  string(FIND "${_attn_dispatch}" "#define ATTN_HEAD_SWITCH" _attn_arms_pos)
-  if (_attn_arms_pos EQUAL -1)
-    message(FATAL_ERROR
-      "csrc/rocm/attention_dispatch.cu no longer defines ATTN_HEAD_SWITCH; "
-      "update the instantiation guard in CMakeLists.txt accordingly.")
-  endif()
-  string(SUBSTRING "${_attn_dispatch}" ${_attn_arms_pos} -1 _attn_arms)
-  foreach (_attn_src ${VLLM_ROCM_EXT_SRC})
-    if (_attn_src MATCHES "^csrc/rocm/attention_paged_attn_.*\\.cu$")
-      file(READ "${CMAKE_SOURCE_DIR}/${_attn_src}" _attn_tu)
-      string(REGEX MATCHALL "ATTN_ENTRY\\([A-Za-z0-9_]+\\)" _attn_entries
-        "${_attn_tu}")
-      foreach (_attn_entry ${_attn_entries})
-        string(REGEX REPLACE "ATTN_ENTRY\\((.*)\\)" "\\1" _attn_name
-          "${_attn_entry}")
-        if (NOT _attn_arms MATCHES "${_attn_name}[,)]")
-          message(FATAL_ERROR
-            "${_attn_src} instantiates ${_attn_name}, but "
-            "csrc/rocm/attention_dispatch.cu never dispatches it: the file "
-            "would be compiled and linked for nothing. Add it to one of the "
-            "ATTN_HEAD_SWITCH() calls in paged_attention().")
-        endif()
-      endforeach()
-    endif()
-  endforeach()
+    ${VLLM_ATTENTION_PAGED_ATTN_SRC})
+
   # Compress the offloaded device binaries. The 9-arch release build packs
   # ~2 GiB of .hip_fatbin into _rocm_C.abi3.so, which pushes the module past
   # the 2 GiB reach of the 32-bit PC-relative relocations in crtbegin's .text
@@ -1671,6 +1747,21 @@ if(VLLM_GPU_LANG STREQUAL "HIP")
     target_compile_definitions(_rocm_C PRIVATE VLLM_SKIP_SKINNY_GEMMS)
   endif()
 
+  # The paged attention units are identical whatever the arch list is; whether the
+  # Navi (gfx11/gfx12) kernel set is needed is decided here, as a compile
+  # definition, so no generated file has to be rewritten for a different build
+  # directory. Keep this predicate in sync with is_navi_gpu() in
+  # csrc/rocm/attention_dispatch.cu.
+  set(_ROCM_NAVI_ARCHS ${VLLM_GPU_ARCHES})
+  list(FILTER _ROCM_NAVI_ARCHS INCLUDE REGEX "gfx1[12]")
+  set(_ROCM_NO_NAVI_ARCHS ${VLLM_GPU_ARCHES})
+  list(FILTER _ROCM_NO_NAVI_ARCHS EXCLUDE REGEX "gfx1[12]")
+  if(_ROCM_NO_NAVI_ARCHS AND NOT _ROCM_NAVI_ARCHS)
+    message(STATUS "paged attention: CDNA only (${VLLM_GPU_ARCHES}), "
+                   "Navi kernels not instantiated")
+    target_compile_definitions(_rocm_C PRIVATE VLLM_ROCM_ATTN_NO_NAVI)
+  endif()
+
   if(VLLM_ROCM_HAS_GFX1100)
     target_compile_definitions(_rocm_C PRIVATE VLLM_ROCM_GFX1100)
   endif()
@@ -1701,7 +1792,6 @@ 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)
diff --git a/benchmarks/kernels/benchmark_dsv41_mega_attn.py b/benchmarks/kernels/benchmark_dsv41_mega_attn.py
index 1b3b5d9c3848..ede9021e0a36 100644
--- a/benchmarks/kernels/benchmark_dsv41_mega_attn.py
+++ b/benchmarks/kernels/benchmark_dsv41_mega_attn.py
@@ -118,7 +118,6 @@ def bench_decode(s_q, device, topk_extra, extra_bytes, local_heads):
     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(
@@ -177,7 +176,7 @@ def mega():
             ex_len,
             False,
             0.0,
-            pos32,
+            positions,
             False,
             64,
             cos_sin,
diff --git a/benchmarks/kernels/benchmark_gemm_rs_ar.py b/benchmarks/kernels/benchmark_gemm_rs_ar.py
index 3e47f61c74f0..9c06e3ba21f3 100644
--- a/benchmarks/kernels/benchmark_gemm_rs_ar.py
+++ b/benchmarks/kernels/benchmark_gemm_rs_ar.py
@@ -18,6 +18,11 @@
 
     torchrun --nproc-per-node=4 -- benchmarks/kernels/benchmark_gemm_rs_ar.py \
         --model deepseek_v41 --backend flashinfer_cutedsl
+
+or the MiniMax-M3 MXFP8 attention O-proj and dense down-proj at TP4 with:
+
+    torchrun --nproc-per-node=4 -- benchmarks/kernels/benchmark_gemm_rs_ar.py \
+        --model minimax_m3 --backend flashinfer_cutedsl
 """
 
 import argparse
@@ -45,10 +50,12 @@
 
 # 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).
+# ``wo_b`` output projection (o_groups * o_lora_rank). MiniMax-M3 lists the
+# attention O-proj and the dense-layer down-proj.
 _MODEL_PROJECTIONS = {
     "kimi_k3": (7168, (6144, 12288)),
     "deepseek_v41": (5120, (8192,)),
+    "minimax_m3": (6144, (8192, 12288)),
 }
 
 
diff --git a/benchmarks/kernels/benchmark_mxfp8_tp4.py b/benchmarks/kernels/benchmark_mxfp8_tp4.py
new file mode 100644
index 000000000000..e9ee03c8a431
--- /dev/null
+++ b/benchmarks/kernels/benchmark_mxfp8_tp4.py
@@ -0,0 +1,208 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Validate MXFP8 TP4 linears and time HIP graph replay over rotating weights.
+
+Example (record the baseline and candidate SHAs):
+    baseline=$(git rev-parse origin/main)
+    git rev-parse "$baseline" HEAD
+    git show "$baseline":vllm/model_executor/kernels/linear/mxfp8/rocm_native.py \
+        > /tmp/baseline_rocm_native.py
+    .venv/bin/python benchmarks/kernels/benchmark_mxfp8_tp4.py \
+        --baseline /tmp/baseline_rocm_native.py
+
+Reports cold (>400 MiB rotating weights) and warm timings separately. Each timing
+is the median of 30 device-event samples; Python wall time is never ranked.
+"""
+
+import argparse
+import importlib.util
+import json
+import math
+import statistics
+from functools import partial
+from pathlib import Path
+
+import torch
+
+from vllm.model_executor.kernels.linear.mxfp8 import rocm_native
+from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
+    dequant_mxfp8_to_bf16,
+    mxfp8_e4m3_quantize,
+)
+from vllm.triton_utils import triton
+
+SHAPES = [
+    (m, n, k)
+    for n, k in [(2560, 6144), (6144, 2048), (1536, 6144), (6144, 768)]
+    for m in [1, 8, 16]
+]
+
+
+def time_graph(calls):
+    for call in calls:
+        call()
+    torch.accelerator.synchronize()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        for _ in range(2):
+            for call in calls:
+                call()
+    for _ in range(5):
+        graph.replay()
+    events = [
+        (torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True))
+        for _ in range(30)
+    ]
+    for start, end in events:
+        start.record()
+        for _ in range(5):
+            graph.replay()
+        end.record()
+        # Bound outstanding HIP graph nodes on short-K cases with many slots.
+        end.synchronize()
+    torch.accelerator.synchronize()
+    return statistics.median(
+        start.elapsed_time(end) / (10 * len(calls)) for start, end in events
+    )
+
+
+def launch(module, xq, xs, w, ws, out, config):
+    bm, bn, bk, nw, ns = config
+    m, k = xq.shape
+    n = w.shape[0]
+    return module._mxfp8_linear_kernel[(triton.cdiv(m, bm), triton.cdiv(n, bn))](
+        xq,
+        xs,
+        w,
+        ws,
+        out,
+        m,
+        n,
+        k,
+        *xq.stride(),
+        *xs.stride(),
+        *w.stride(),
+        *ws.stride(),
+        *out.stride(),
+        BLOCK_M=bm,
+        BLOCK_N=bn,
+        BLOCK_K=bk,
+        num_warps=nw,
+        num_stages=ns,
+    )
+
+
+def emit(**data):
+    print(json.dumps(data), flush=True)
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--baseline", type=Path)
+    parser.add_argument("--device", type=int, default=0)
+    parser.add_argument("--shape", type=int, choices=range(len(SHAPES)))
+    parser.add_argument("--tune", action="store_true")
+    parser.add_argument("--profile", action="store_true")
+    parser.add_argument("--asm-dir", type=Path)
+    args = parser.parse_args()
+    torch.accelerator.set_device_index(args.device)
+    torch.manual_seed(42)
+    variants = {"candidate": rocm_native}
+    if args.baseline:
+        spec = importlib.util.spec_from_file_location(
+            "vllm.model_executor.kernels.linear.mxfp8.baseline_rocm_native",
+            args.baseline,
+        )
+        baseline = importlib.util.module_from_spec(spec)
+        spec.loader.exec_module(baseline)
+        variants = {"baseline": baseline, **variants}
+    emit(
+        torch=torch.__version__,
+        triton=triton.__version__,
+        device=str(torch.cuda.get_device_properties(args.device)),
+    )
+    for index, (m, n, k) in enumerate(SHAPES):
+        if args.shape is not None and index != args.shape:
+            continue
+        x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16) * 0.5
+        w, ws = mxfp8_e4m3_quantize(
+            torch.randn(n, k, device="cuda", dtype=torch.bfloat16) * 0.05
+        )
+        xq, xs = mxfp8_e4m3_quantize(x)
+        reference = (
+            dequant_mxfp8_to_bf16(xq, xs).float()
+            @ dequant_mxfp8_to_bf16(w, ws).float().T
+        )
+        out = torch.empty(m, n, device="cuda", dtype=x.dtype)
+        slots = math.ceil(400 * 1024**2 / (n * k))
+        weights = [(w.clone(), ws.clone()) for _ in range(slots)]
+        calls = {}
+        for name, module in variants.items():
+            configs = [module._select_cfg(m, n, k)]
+            if args.tune:
+                configs += [
+                    (16, 16, 1024, 2, 2),
+                    (16, 16, 1024, 2, 3),
+                    (16, 16, 1024, 2, 4),
+                    (16, 16, 512, 2, 2),
+                    (16, 32, 256, 4, 3),
+                    (16, 16, 256, 2, 3),
+                    (128, 128, 128, 8, 2),
+                ]
+            for config in dict.fromkeys(configs):
+                if k % config[2]:
+                    continue
+                fn = partial(launch, module, xq, xs, config=config, out=out)
+                obj = fn(w=w, ws=ws)
+                rel = ((out.float() - reference).norm() / reference.norm()).item()
+                snr = -20 * math.log10(rel) if rel else float("inf")
+                assert rel < 0.005 and snr >= 30, (name, config, rel, snr)
+                emit(
+                    shape=[m, n, k],
+                    variant=name,
+                    config=config,
+                    rel=rel,
+                    snr_db=snr,
+                    registers=obj.n_regs,
+                    spills=obj.n_spills,
+                    lds=obj.metadata.shared,
+                )
+                if args.asm_dir:
+                    args.asm_dir.mkdir(parents=True, exist_ok=True)
+                    (args.asm_dir / f"{name}-{m}-{n}-{k}-{config}.amdgcn").write_text(
+                        obj.asm["amdgcn"]
+                    )
+                if args.profile:
+                    for _ in range(10):
+                        fn(w=w, ws=ws)
+                    torch.accelerator.synchronize()
+                    continue
+                calls[f"{name}-kernel-{config}"] = [
+                    partial(fn, w=wi, ws=si) for wi, si in weights
+                ]
+            if not args.profile and not args.tune:
+                wrapper = partial(module._mxfp8_dot_scaled_linear, x)
+                actual = wrapper(w, ws)
+                assert ((actual.float() - reference).norm() / reference.norm()) < 0.005
+                calls[f"{name}-wrapper"] = [
+                    partial(wrapper, wi, si) for wi, si in weights
+                ]
+        for repeat in range(2):
+            order = list(calls.items())
+            if repeat:
+                order.reverse()
+            for name, functions in order:
+                cold = time_graph(functions)
+                warm = time_graph([functions[0]] * 32)
+                emit(
+                    shape=[m, n, k],
+                    variant=name,
+                    repeat=repeat,
+                    cold_ms=cold,
+                    warm_ms=warm,
+                    slots=slots,
+                )
+
+
+if __name__ == "__main__":
+    main()
diff --git a/benchmarks/kernels/benchmark_reshape_and_cache_flash.py b/benchmarks/kernels/benchmark_reshape_and_cache_flash.py
index 84d9cc707c40..229f5385a4c2 100644
--- a/benchmarks/kernels/benchmark_reshape_and_cache_flash.py
+++ b/benchmarks/kernels/benchmark_reshape_and_cache_flash.py
@@ -22,10 +22,45 @@
 
 
 @torch.inference_mode()
+def _make_packed_kv_caches(
+    num_blocks: int,
+    block_size: int,
+    num_heads: int,
+    head_size: int,
+    v_head_size: int,
+    kv_cache_dtype: str,
+    dtype: torch.dtype,
+    kv_cache_layout: str,
+    device: str,
+):
+    """K and V as views of one [.., head_size + v_head_size] buffer.
+
+    The asymmetric kernel needs both sides to share a head stride, which is what
+    a packed K|V buffer gives (and what backends such as ROCm AITER unified
+    attention produce).
+    """
+    if kv_cache_dtype not in ("auto", "bfloat16", "float16", "fp8"):
+        raise ValueError(f"asymmetric benchmark does not support {kv_cache_dtype}.")
+    content = head_size + v_head_size
+    if kv_cache_layout == "NHD":
+        physical = (num_blocks, block_size, num_heads, content)
+        permute = None
+    else:
+        physical = (num_blocks, num_heads, block_size, content)
+        permute = (0, 2, 1, 3)
+
+    cache_dtype = torch.uint8 if kv_cache_dtype == "fp8" else dtype
+    packed = torch.zeros(physical, dtype=cache_dtype, device=device)
+    view = packed if permute is None else packed.permute(permute)
+    key_cache, value_cache = view.split([head_size, v_head_size], dim=-1)
+    return key_cache, value_cache
+
+
 def run_benchmark(
     num_tokens: int,
     num_heads: int,
     head_size: int,
+    v_head_size: int | None,
     block_size: int,
     num_blocks: int,
     dtype: torch.dtype,
@@ -37,8 +72,14 @@ 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.")
+    if v_head_size is None:
+        v_head_size = head_size
+    asymmetric = v_head_size != head_size
+
+    if kv_cache_dtype == "fp8" and (head_size % 16 or v_head_size % 16):
+        raise ValueError(
+            "fp8 kv-cache requires head_size and v_head_size to be multiples of 16."
+        )
 
     if implementation not in ("cuda", "triton"):
         raise ValueError(
@@ -47,13 +88,23 @@ def run_benchmark(
         )
     if implementation == "triton" and kv_cache_layout == "HND":
         return float("nan")  # Triton does not support HND layout yet.
+    if asymmetric and implementation != "cuda":
+        raise ValueError(
+            "asymmetric head sizes (v_head_size != head_size) are only supported "
+            "by the 'cuda' implementation."
+        )
 
     set_random_seed(42)
     torch.set_default_device(device)
 
     # create random key / value tensors [T, H, D].
     key = torch.randn(num_tokens, num_heads, head_size, dtype=dtype, device=device)
-    value = torch.randn_like(key)
+    if asymmetric:
+        value = torch.randn(
+            num_tokens, num_heads, v_head_size, dtype=dtype, device=device
+        )
+    else:
+        value = torch.randn_like(key)
 
     # prepare the slot mapping.
     # each token is assigned a unique slot in the KV-cache.
@@ -63,20 +114,36 @@ def run_benchmark(
     slot_mapping_lst = random.sample(range(num_slots), num_tokens)
     slot_mapping = torch.tensor(slot_mapping_lst, dtype=torch.long, device=device)
 
-    key_caches, value_caches = create_kv_caches_with_random_flash(
-        num_blocks,
-        block_size,
-        1,  # num_layers
-        num_heads,
-        head_size,
-        kv_cache_dtype,
-        dtype,
-        device=device,
-        cache_layout=kv_cache_layout,
-    )
-    key_cache, value_cache = key_caches[0], value_caches[0]
-    # to free unused memory
-    del key_caches, value_caches
+    if asymmetric:
+        # K and V are views of one buffer, so both sides share the head stride;
+        # that is the only layout the asymmetric path accepts (see
+        # reshape_and_cache_flash).
+        key_cache, value_cache = _make_packed_kv_caches(
+            num_blocks,
+            block_size,
+            num_heads,
+            head_size,
+            v_head_size,
+            kv_cache_dtype,
+            dtype,
+            kv_cache_layout,
+            device,
+        )
+    else:
+        key_caches, value_caches = create_kv_caches_with_random_flash(
+            num_blocks,
+            block_size,
+            1,  # num_layers
+            num_heads,
+            head_size,
+            kv_cache_dtype,
+            dtype,
+            device=device,
+            cache_layout=kv_cache_layout,
+        )
+        key_cache, value_cache = key_caches[0], value_caches[0]
+        # to free unused memory
+        del key_caches, value_caches
 
     # compute per-kernel scaling factors for fp8 conversion (if used).
     k_scale = (key.amax() / 64.0).to(torch.float32)
@@ -142,6 +209,7 @@ def main(args):
                 num_tokens=n_tok,
                 num_heads=args.num_heads,
                 head_size=args.head_size,
+                v_head_size=args.v_head_size,
                 block_size=args.block_size,
                 num_blocks=args.num_blocks,
                 dtype=STR_DTYPE_TO_TORCH_DTYPE[args.dtype],
@@ -154,9 +222,14 @@ def main(args):
             )
             rows.append([n_tok, layout, f"{lat * 1e6:.3f}"])
 
+    v_note = (
+        ""
+        if args.v_head_size is None
+        else f", head_size={args.head_size} v_head_size={args.v_head_size}"
+    )
     print(
         f"Benchmark results for implementation {args.implementation}"
-        f" (measuring with {args.mode}):"
+        f" (measuring with {args.mode}{v_note}):"
     )
     print(tabulate(rows, headers=["num_tokens", "layout", "latency (µs)"]))
 
@@ -171,6 +244,13 @@ def main(args):
         choices=[64, 80, 96, 112, 120, 128, 192, 256],
         default=128,
     )
+    parser.add_argument(
+        "--v-head-size",
+        type=int,
+        default=None,
+        help="V head dim; defaults to --head-size. Setting it differently "
+        "benchmarks the asymmetric K/V path (cuda implementation only).",
+    )
     parser.add_argument("--block-size", type=int, choices=[16, 32], default=16)
     parser.add_argument("--num-blocks", type=int, default=128 * 512)
 
diff --git a/cmake/cpu_extension.cmake b/cmake/cpu_extension.cmake
index d92ad58a9cf3..39b327cb76b7 100644
--- a/cmake/cpu_extension.cmake
+++ b/cmake/cpu_extension.cmake
@@ -492,6 +492,8 @@ set(VLLM_EXT_SRC
     "csrc/cpu/activation.cpp"
     "csrc/cpu/sgl-kernels/fla.cpp"
     "csrc/cpu/utils.cpp"
+    "csrc/cpu/block_table.cpp"
+    "csrc/cpu/input_batch.cpp"
     "csrc/cpu/spec_decode_utils.cpp"
     "csrc/cpu/layernorm.cpp"
     "csrc/cpu/mla_decode.cpp"
@@ -693,6 +695,7 @@ else()
         USE_SABI 3
         WITH_SOABI
     )
+    target_compile_definitions(_C PRIVATE VLLM_STANDALONE_CPU_MRV2)
     if (VLLM_OPENBLAS_LIB)
         target_compile_definitions(_C PRIVATE VLLM_HAS_OPENBLAS)
     endif()
diff --git a/cmake/external_projects/deepgemm.cmake b/cmake/external_projects/deepgemm.cmake
index 2fbfa157de88..e9d0bbbd2803 100644
--- a/cmake/external_projects/deepgemm.cmake
+++ b/cmake/external_projects/deepgemm.cmake
@@ -37,7 +37,7 @@ else()
   # (vllm-project/DeepGEMM#39), and the MX MegaMoE alignment fix
   # (vllm-project/DeepGEMM#43).
   # Uses the PyTorch 2.11 stable ABI and CPython 3.10 limited API.
-  set(_DEEPGEMM_UPSTREAM_TAG "0820030756945b2daec208e8e8255268360f337d")
+  set(_DEEPGEMM_UPSTREAM_TAG "5cdb145704ec4672ecf4b61441599521ce635df6")
 
   set(_deepgemm_fc_root "${FETCHCONTENT_BASE_DIR}")
   if(NOT _deepgemm_fc_root)
diff --git a/cmake/external_projects/flashkda.cmake b/cmake/external_projects/flashkda.cmake
index fadee58efbbd..f34efd577cca 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 17a037d98da546deb4591e967cf961a43c034d8b
+    GIT_TAG 9d9e928edc42475a67a523dd37814540d105f51c
     GIT_PROGRESS TRUE
     GIT_SUBMODULES cutlass
   )
diff --git a/cmake/external_projects/flashmla.cmake b/cmake/external_projects/flashmla.cmake
index d60fd84e914b..6286e8e16303 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 e2afa5fd8c220fbd05fe03fc24c9d5e64c24641b
+        GIT_TAG 847ef308bb7a6828c8117301f5a5a84a095a90e8
         GIT_PROGRESS TRUE
         CONFIGURE_COMMAND ""
         BUILD_COMMAND ""
diff --git a/cmake/external_projects/fmha_sm100.cmake b/cmake/external_projects/fmha_sm100.cmake
index 6b84b35c9049..2336d66d23dc 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 968061ab16cb14991a5c7d4c0e8eabf531be0dcf
+    GIT_TAG 25c144700f727b638067e98066a60507081b63ad
     GIT_PROGRESS TRUE
     CONFIGURE_COMMAND ""
     BUILD_COMMAND ""
diff --git a/csrc/cache_kernels.cu b/csrc/cache_kernels.cu
index 9567d87f443a..bd4a080c25ff 100644
--- a/csrc/cache_kernels.cu
+++ b/csrc/cache_kernels.cu
@@ -420,6 +420,73 @@ __global__ void reshape_and_cache_flash_kernel(
   }
 }
 
+// Asymmetric K/V head widths. K and V are views of one packed buffer, so they
+// share a head stride and the head-contiguous path below cannot apply; this
+// kernel therefore walks heads explicitly and takes V's width as its one extra
+// argument. It is kept separate so the symmetric kernel's codegen is unchanged.
+template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt>
+__global__ void reshape_and_cache_flash_kernel_asym(
+    const scalar_t* __restrict__ key,    // [num_tokens, num_heads, head_size]
+    const scalar_t* __restrict__ value,  // [num_tokens, num_heads, v_head_size]
+    // NHD or HND: [num_blocks, block_size, num_heads, head_size], addressed
+    // through the strides passed by the launcher
+    cache_t* __restrict__ key_cache,
+    cache_t* __restrict__ value_cache,         // same, with v_head_size
+    const int64_t* __restrict__ slot_mapping,  // [num_tokens]
+    const int64_t block_stride, const int64_t page_stride,
+    const int64_t head_stride, const int64_t key_stride,
+    const int64_t value_stride, const int num_heads, const int head_size,
+    const int v_head_size, const int block_size, const float* k_scale,
+    const float* v_scale, const int kv_scale_stride) {
+  const int64_t token_idx = blockIdx.x;
+  const int64_t slot_idx = slot_mapping[token_idx];
+  // NOTE: slot_idx can be -1 if the token is padded
+  if (slot_idx < 0) {
+    return;
+  }
+  const int64_t block_idx = slot_idx / block_size;
+  const int64_t block_offset = slot_idx % block_size;
+
+  const scalar_t* __restrict__ key_src = key + token_idx * key_stride;
+  const scalar_t* __restrict__ value_src = value + token_idx * value_stride;
+
+  cache_t* __restrict__ key_dst =
+      key_cache + block_idx * block_stride + block_offset * page_stride;
+  cache_t* __restrict__ value_dst =
+      value_cache + block_idx * block_stride + block_offset * page_stride;
+
+  constexpr int VEC_SIZE = (sizeof(scalar_t) == 2) ? 8 : 4;
+
+  const int lane = threadIdx.x & 31;     // 0..31 within warp
+  const int warp_id = threadIdx.x >> 5;  // warp index within block
+  const int warps_per_block = blockDim.x >> 5;
+
+  for (int head = warp_id; head < num_heads; head += warps_per_block) {
+    const scalar_t* __restrict__ k_src_h = key_src + head * head_size;
+    const scalar_t* __restrict__ v_src_h = value_src + head * v_head_size;
+
+    cache_t* __restrict__ k_dst_h =
+        key_dst + static_cast<int64_t>(head) * head_stride;
+    cache_t* __restrict__ v_dst_h =
+        value_dst + static_cast<int64_t>(head) * head_stride;
+
+    float k_scale_val = (kv_dt == Fp8KVCacheDataType::kAuto)
+                            ? 0.f
+                            : k_scale[head * kv_scale_stride];
+    float v_scale_val = (kv_dt == Fp8KVCacheDataType::kAuto)
+                            ? 0.f
+                            : v_scale[head * kv_scale_stride];
+
+    CopyWithScaleOp<cache_t, scalar_t, kv_dt> k_op{k_scale_val};
+    CopyWithScaleOp<cache_t, scalar_t, kv_dt> v_op{v_scale_val};
+
+    vectorize_with_alignment<VEC_SIZE>(k_src_h, k_dst_h, head_size, lane, 32,
+                                       k_op);
+    vectorize_with_alignment<VEC_SIZE>(v_src_h, v_dst_h, v_head_size, lane, 32,
+                                       v_op);
+  }
+}
+
 template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt,
           bool skip_scale>
 __global__ void concat_and_cache_mla_kernel(
@@ -540,6 +607,7 @@ __global__ void concat_and_cache_ds_mla_kernel(
     const int kv_lora_rank,                    //
     const int pe_dim,                          //
     const int block_size,                      //
+    const bool has_rope_tail,                  //
     const float* scale                         //
 ) {
   const int64_t token_idx = blockIdx.x;
@@ -556,15 +624,15 @@ __global__ void concat_and_cache_ds_mla_kernel(
   // For the NoPE part, each tile of 128 elements is handled by half of one warp
   // (16 threads). There are 4 total tiles, so 2 warps (64 threads).
   // Lanes 0 and 16 of each warp write the scale values for that warp's tiles.
-  // The RoPE part (last 64 elements) is handled by another 1 warp (32 threads).
+  // The optional RoPE tail is handled by another 1 warp (32 threads).
   // So in total, we use 3 warps (96 threads) per block.
 
-  // Cast kv_cache to 16_bit for RoPE values
-  scalar_t* kv_cache_16bit =
-      reinterpret_cast<scalar_t*>(&kv_cache[dst_idx_start]);
-
-  // Zero the reserved RoPE tail for NoPE rows.
   if (threadIdx.x >= 64) {
+    if (!has_rope_tail) {
+      return;
+    }
+    scalar_t* kv_cache_16bit =
+        reinterpret_cast<scalar_t*>(&kv_cache[dst_idx_start]);
     // Each thread handles two elements of RoPE
     const int8_t pe_idx_start = (threadIdx.x - 64) * 2;
     // RoPE values start after the packed 8-bit NoPE values and the
@@ -826,6 +894,20 @@ void reshape_and_cache(
 // KV_T is the data type of key and value tensors.
 // CACHE_T is the stored data type of kv-cache.
 // KV_DTYPE is the real data type of kv-cache.
+#define CALL_RESHAPE_AND_CACHE_FLASH_ASYM(KV_T, CACHE_T, KV_DTYPE)           \
+  vllm::reshape_and_cache_flash_kernel_asym<KV_T, CACHE_T, KV_DTYPE>         \
+      <<<grid, block, 0, stream>>>(                                          \
+          reinterpret_cast<KV_T*>(key.data_ptr()),                           \
+          reinterpret_cast<KV_T*>(value.data_ptr()),                         \
+          reinterpret_cast<CACHE_T*>(key_cache.data_ptr()),                  \
+          reinterpret_cast<CACHE_T*>(value_cache.data_ptr()),                \
+          slot_mapping.const_data_ptr<int64_t>(), block_stride, page_stride, \
+          head_stride, key_stride, value_stride, num_heads, head_size,       \
+          v_head_size, block_size,                                           \
+          reinterpret_cast<const float*>(k_scale.data_ptr()),                \
+          reinterpret_cast<const float*>(v_scale.data_ptr()),                \
+          kv_scale_stride);
+
 #define CALL_RESHAPE_AND_CACHE_FLASH(KV_T, CACHE_T, KV_DTYPE)                \
   vllm::reshape_and_cache_flash_kernel<KV_T, CACHE_T, KV_DTYPE>              \
       <<<grid, block, 0, stream>>>(                                          \
@@ -863,6 +945,7 @@ void reshape_and_cache_flash(
   int num_tokens = slot_mapping.size(0);
   int num_heads = key.size(1);
   int head_size = key.size(2);
+  int v_head_size = value.size(2);
 
   const torch::stable::accelerator::DeviceGuard device_guard(
       key.get_device_index());
@@ -896,8 +979,8 @@ void reshape_and_cache_flash(
   int64_t block_stride = key_cache.stride(0);
   int64_t page_stride = key_cache.stride(1);
   int64_t head_stride = key_cache.stride(2);
-  STD_TORCH_CHECK(key_cache.stride(0) == value_cache.stride(0));
 
+  STD_TORCH_CHECK(key_cache.stride(0) == value_cache.stride(0));
   STD_TORCH_CHECK(k_scale.sizes().equals(v_scale.sizes()),
                   "k_scale and v_scale must have the same shape");
   STD_TORCH_CHECK(k_scale.numel() == 1 || k_scale.numel() == num_heads,
@@ -905,10 +988,31 @@ void reshape_and_cache_flash(
   int kv_scale_stride = (k_scale.numel() > 1) ? 1 : 0;
 
   dim3 grid(num_tokens);
-  dim3 block(std::min(num_heads * head_size, 512));
-
-  DISPATCH_BY_KV_CACHE_DTYPE(key.scalar_type(), kv_cache_dtype,
-                             CALL_RESHAPE_AND_CACHE_FLASH);
+  // Only models with v_head_size != head_size take the kernel below; everything
+  // else keeps using the kernel above, so its codegen is untouched.
+  if (v_head_size == head_size) {
+    dim3 block(std::min(num_heads * head_size, 512));
+    DISPATCH_BY_KV_CACHE_DTYPE(key.scalar_type(), kv_cache_dtype,
+                               CALL_RESHAPE_AND_CACHE_FLASH);
+  } else {
+    // V is a view of K's buffer, so the two must agree on geometry; these
+    // checks stay out of the symmetric path so its host-side work is unchanged.
+    STD_TORCH_CHECK(value_cache.stride(1) == page_stride,
+                    "key_cache and value_cache must share the page stride");
+    STD_TORCH_CHECK(key_cache.size(0) == value_cache.size(0) &&
+                        key_cache.size(1) == value_cache.size(1),
+                    "key_cache and value_cache must share block geometry");
+    STD_TORCH_CHECK(key.size(1) == value.size(1),
+                    "key and value must have the same number of heads");
+    STD_TORCH_CHECK(value_cache.stride(2) == head_stride,
+                    "key_cache and value_cache must share a head stride");
+    // Only the strided walk applies here, and V's width is the single extra
+    // argument. One warp per head, and at least one warp so the head loop
+    // always advances.
+    dim3 block(32 * std::min(num_heads, 16));
+    DISPATCH_BY_KV_CACHE_DTYPE(key.scalar_type(), kv_cache_dtype,
+                               CALL_RESHAPE_AND_CACHE_FLASH_ASYM);
+  }
 }
 
 // KV_T is the data type of key and value tensors.
@@ -941,7 +1045,7 @@ void reshape_and_cache_flash(
           reinterpret_cast<CACHE_T*>(kv_cache.data_ptr()),                    \
           slot_mapping.const_data_ptr<int64_t>(), block_stride, entry_stride, \
           kv_c_stride, k_pe_stride, kv_lora_rank, pe_dim, block_size,         \
-          reinterpret_cast<const float*>(scale.data_ptr()));
+          has_rope_tail, reinterpret_cast<const float*>(scale.data_ptr()));
 
 void concat_and_cache_mla(
     torch::stable::Tensor& kv_c,      // [num_tokens, kv_lora_rank]
@@ -965,6 +1069,8 @@ void concat_and_cache_mla(
   int kv_lora_rank = kv_c.size(1);
   int pe_dim = k_pe.size(1);
   int block_size = kv_cache.size(1);
+  const int64_t row_bytes = kv_cache.size(2) * kv_cache.element_size();
+  const bool has_rope_tail = row_bytes == 656;
 
   const bool is_nvfp4_ds_mla = kv_cache_dtype == "nvfp4_ds_mla";
   if (kv_cache_dtype == "fp8_ds_mla") {
@@ -972,8 +1078,9 @@ void concat_and_cache_mla(
                     "kv_lora_rank must be 512 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(has_rope_tail || (pe_dim == 0 && row_bytes == 528),
+                    "fp8_ds_mla requires 656-byte rows, or 528-byte rows with "
+                    "pe_dim == 0");
     STD_TORCH_CHECK(kv_c.element_size() == 2,
                     "kv_c.element_size() must be 2 for fp8_ds_mla");
     STD_TORCH_CHECK(k_pe.element_size() == 2,
@@ -1011,8 +1118,8 @@ void concat_and_cache_mla(
     // For the NoPE part, each tile of 128 elements is handled by half of one
     // warp (16 threads). There are 4 total tiles, so 2 warps (64 threads).
     // Lanes 0 and 16 of each warp write the scale values for that warp's tiles.
-    // The RoPE part (last 64 elements) is handled by another 1 warp (32
-    // threads). So in total, we use 3 warps (96 threads) per block.
+    // The optional RoPE tail is handled by another 1 warp (32 threads).
+    // So in total, we use 3 warps (96 threads) per block.
     dim3 block(96);
     DISPATCH_BY_KV_CACHE_DTYPE(kv_c.scalar_type(), kv_cache_dtype,
                                CALL_CONCAT_AND_CACHE_DS_MLA);
@@ -1837,6 +1944,8 @@ void cp_gather_and_upconvert_fp8_kv_cache(
   STD_TORCH_CHECK(dst.scalar_type() == torch::headeronly::ScalarType::BFloat16,
                   "dst must be bfloat16");
   STD_TORCH_CHECK(head_dim == 576, "head_dim must be 576 for MLA");
+  STD_TORCH_CHECK(src_cache.size(2) * src_cache.element_size() >= 656,
+                  "src_cache rows must contain at least 656 bytes for MLA");
 
   int64_t block_table_stride = block_table.stride(0);
   int64_t cache_block_stride = src_cache.stride(0);
diff --git a/csrc/cpu/block_table.cpp b/csrc/cpu/block_table.cpp
new file mode 100644
index 000000000000..ac8f21d44d6e
--- /dev/null
+++ b/csrc/cpu/block_table.cpp
@@ -0,0 +1,73 @@
+#include "cpu_types.hpp"
+
+#include <algorithm>
+#include <cstring>
+
+namespace cpu_utils {
+
+void gather_block_table_cpu(const torch::Tensor& block_table,
+                            torch::Tensor& output,
+                            const torch::Tensor& idx_mapping,
+                            const torch::Tensor& num_blocks,
+                            const int64_t num_reqs_padded) {
+  const int32_t* block_table_ptr = block_table.data_ptr<int32_t>();
+  int32_t* output_ptr = output.data_ptr<int32_t>();
+  const int64_t* idx_mapping_ptr = idx_mapping.data_ptr<int64_t>();
+  const int32_t* num_blocks_ptr = num_blocks.data_ptr<int32_t>();
+
+  const int64_t num_reqs = idx_mapping.size(0);
+  const int64_t block_table_stride = block_table.stride(0);
+  const int64_t output_stride = output.stride(0);
+  for (int64_t batch_idx = 0; batch_idx < num_reqs; ++batch_idx) {
+    const int64_t req_idx = idx_mapping_ptr[batch_idx];
+    const int64_t request_num_blocks = num_blocks_ptr[req_idx];
+    std::memcpy(output_ptr + batch_idx * output_stride,
+                block_table_ptr + req_idx * block_table_stride,
+                request_num_blocks * sizeof(int32_t));
+  }
+
+  for (int64_t batch_idx = num_reqs; batch_idx < num_reqs_padded; ++batch_idx) {
+    std::memset(output_ptr + batch_idx * output_stride, 0,
+                output.size(1) * sizeof(int32_t));
+  }
+}
+
+void compute_slot_mapping_cpu(
+    const torch::Tensor& block_table, torch::Tensor& slot_mapping,
+    const torch::Tensor& idx_mapping, const torch::Tensor& query_start_loc,
+    const torch::Tensor& positions, const int64_t block_size,
+    const bool mapping_enabled, const int64_t pad_slot_id) {
+  const int32_t* block_table_ptr = block_table.data_ptr<int32_t>();
+  int64_t* slot_mapping_ptr = slot_mapping.data_ptr<int64_t>();
+  const int64_t* idx_mapping_ptr = idx_mapping.data_ptr<int64_t>();
+  const int32_t* query_start_loc_ptr = query_start_loc.data_ptr<int32_t>();
+  const int64_t* positions_ptr = positions.data_ptr<int64_t>();
+
+  const int64_t num_reqs = idx_mapping.size(0);
+  const int64_t num_tokens = query_start_loc_ptr[num_reqs];
+  const int64_t block_table_stride = block_table.stride(0);
+  std::fill(slot_mapping_ptr + num_tokens,
+            slot_mapping_ptr + slot_mapping.numel(), pad_slot_id);
+
+  for (int64_t batch_idx = 0; batch_idx < num_reqs; ++batch_idx) {
+    const int64_t req_state_idx = idx_mapping_ptr[batch_idx];
+    const int64_t start = query_start_loc_ptr[batch_idx];
+    const int64_t end = query_start_loc_ptr[batch_idx + 1];
+    if (!mapping_enabled || req_state_idx < 0) {
+      std::fill(slot_mapping_ptr + start, slot_mapping_ptr + end, pad_slot_id);
+      continue;
+    }
+
+    for (int64_t token_idx = start; token_idx < end; ++token_idx) {
+      const int64_t position = positions_ptr[token_idx];
+      const int64_t block_idx = position / block_size;
+      const int64_t block_offset = position % block_size;
+      const int32_t block_number =
+          block_table_ptr[req_state_idx * block_table_stride + block_idx];
+      slot_mapping_ptr[token_idx] =
+          static_cast<int64_t>(block_number) * block_size + block_offset;
+    }
+  }
+}
+
+}  // namespace cpu_utils
diff --git a/csrc/cpu/cpu_types_scalar.hpp b/csrc/cpu/cpu_types_scalar.hpp
index d94cdbc076e8..76b3aec5d45e 100644
--- a/csrc/cpu/cpu_types_scalar.hpp
+++ b/csrc/cpu/cpu_types_scalar.hpp
@@ -502,7 +502,7 @@ inline FP16Vec16::FP16Vec16(const FP32Vec16& v) {
       [&v, this](int i) { reg.val[i] = float_to_fp16(v.reg.val[i]); });
 }
 
-inline FP16Vec8 ::FP16Vec8(const FP32Vec8& v) {
+inline FP16Vec8::FP16Vec8(const FP32Vec8& v) {
   unroll_loop<int, FP16Vec8::VEC_ELEM_NUM>(
       [&v, this](int i) { reg.val[i] = float_to_fp16(v.reg.val[i]); });
 }
diff --git a/csrc/cpu/input_batch.cpp b/csrc/cpu/input_batch.cpp
new file mode 100644
index 000000000000..8f04e05b4d2b
--- /dev/null
+++ b/csrc/cpu/input_batch.cpp
@@ -0,0 +1,295 @@
+#include "cpu_types.hpp"
+
+#include <algorithm>
+#include <cstring>
+#include <numeric>
+
+namespace cpu_utils {
+
+void prepare_prefill_inputs_cpu(torch::Tensor& input_ids,
+                                torch::Tensor& next_prefill_tokens,
+                                const torch::Tensor& idx_mapping,
+                                const torch::Tensor& query_start_loc,
+                                const torch::Tensor& all_token_ids,
+                                const torch::Tensor& prefill_len,
+                                const torch::Tensor& num_computed_tokens) {
+  int32_t* input_ids_ptr = input_ids.data_ptr<int32_t>();
+  int32_t* next_prefill_tokens_ptr = next_prefill_tokens.data_ptr<int32_t>();
+  const int64_t* idx_mapping_ptr = idx_mapping.data_ptr<int64_t>();
+  const int32_t* query_start_loc_ptr = query_start_loc.data_ptr<int32_t>();
+  const int32_t* all_token_ids_ptr = all_token_ids.data_ptr<int32_t>();
+  const int32_t* prefill_len_ptr = prefill_len.data_ptr<int32_t>();
+  const int32_t* num_computed_tokens_ptr =
+      num_computed_tokens.data_ptr<int32_t>();
+
+  const int64_t num_reqs = idx_mapping.size(0);
+  const int64_t num_lookahead = next_prefill_tokens.size(0);
+  const int64_t next_prefill_tokens_stride = next_prefill_tokens.stride(0);
+  const int64_t all_token_ids_stride = all_token_ids.stride(0);
+
+  for (int64_t batch_idx = 0; batch_idx < num_reqs; ++batch_idx) {
+    const int64_t req_state_idx = idx_mapping_ptr[batch_idx];
+    const int32_t request_prefill_len = prefill_len_ptr[req_state_idx];
+    const int32_t num_computed = num_computed_tokens_ptr[req_state_idx];
+    if (num_computed >= request_prefill_len) {
+      continue;
+    }
+
+    const int32_t query_start = query_start_loc_ptr[batch_idx];
+    const int32_t query_end = query_start_loc_ptr[batch_idx + 1];
+    const int32_t query_len = query_end - query_start;
+    const int32_t* request_tokens =
+        all_token_ids_ptr + req_state_idx * all_token_ids_stride;
+    std::memcpy(input_ids_ptr + query_start, request_tokens + num_computed,
+                query_len * sizeof(int32_t));
+
+    const int64_t lookahead_start =
+        static_cast<int64_t>(num_computed) + query_len;
+    for (int64_t lookahead = 0; lookahead < num_lookahead; ++lookahead) {
+      const int64_t pos = lookahead_start + lookahead;
+      next_prefill_tokens_ptr[lookahead * next_prefill_tokens_stride +
+                              req_state_idx] =
+          pos < request_prefill_len ? request_tokens[pos] : 0;
+    }
+  }
+}
+
+torch::Tensor get_prompt_logprobs_token_ids_cpu(
+    const int64_t num_tokens, const torch::Tensor& query_start_loc,
+    const torch::Tensor& idx_mapping, const torch::Tensor& num_computed_tokens,
+    const torch::Tensor& all_token_ids) {
+  torch::Tensor token_ids =
+      torch::empty({num_tokens}, idx_mapping.options().dtype(torch::kInt64));
+
+  int64_t* token_ids_ptr = token_ids.data_ptr<int64_t>();
+  const int32_t* query_start_loc_ptr = query_start_loc.data_ptr<int32_t>();
+  const int64_t* idx_mapping_ptr = idx_mapping.data_ptr<int64_t>();
+  const int32_t* num_computed_tokens_ptr =
+      num_computed_tokens.data_ptr<int32_t>();
+  const int32_t* all_token_ids_ptr = all_token_ids.data_ptr<int32_t>();
+  const int64_t all_token_ids_stride = all_token_ids.stride(0);
+
+  for (int64_t batch_idx = 0; batch_idx < idx_mapping.size(0); ++batch_idx) {
+    const int64_t req_state_idx = idx_mapping_ptr[batch_idx];
+    const int32_t query_start = query_start_loc_ptr[batch_idx];
+    const int32_t query_end = query_start_loc_ptr[batch_idx + 1];
+    const int32_t target_start = num_computed_tokens_ptr[req_state_idx] + 1;
+    const int32_t* request_tokens =
+        all_token_ids_ptr + req_state_idx * all_token_ids_stride;
+
+    std::transform(request_tokens + target_start,
+                   request_tokens + target_start + query_end - query_start,
+                   token_ids_ptr + query_start, [](int32_t token_id) {
+                     return static_cast<int64_t>(token_id);
+                   });
+  }
+  return token_ids;
+}
+
+void prepare_pos_seq_lens_cpu(const torch::Tensor& idx_mapping,
+                              const torch::Tensor& query_start_loc,
+                              const torch::Tensor& num_computed_tokens,
+                              torch::Tensor& pos, torch::Tensor& seq_lens) {
+  const int64_t* idx_mapping_ptr = idx_mapping.data_ptr<int64_t>();
+  const int32_t* query_start_loc_ptr = query_start_loc.data_ptr<int32_t>();
+  const int32_t* num_computed_tokens_ptr =
+      num_computed_tokens.data_ptr<int32_t>();
+  int64_t* pos_ptr = pos.data_ptr<int64_t>();
+  int32_t* seq_lens_ptr = seq_lens.data_ptr<int32_t>();
+
+  const int64_t num_reqs = idx_mapping.size(0);
+  std::fill(seq_lens_ptr + num_reqs, seq_lens_ptr + seq_lens.numel(), 0);
+
+  for (int64_t req_idx = 0; req_idx < num_reqs; ++req_idx) {
+    const int64_t req_state_idx = idx_mapping_ptr[req_idx];
+    const int32_t num_computed = num_computed_tokens_ptr[req_state_idx];
+    const int64_t start = query_start_loc_ptr[req_idx];
+    const int64_t end = query_start_loc_ptr[req_idx + 1];
+
+    seq_lens_ptr[req_idx] = num_computed + static_cast<int32_t>(end - start);
+    std::iota(pos_ptr + start, pos_ptr + end,
+              static_cast<int64_t>(num_computed));
+  }
+}
+
+void prepare_mrope_positions_cpu(torch::Tensor& positions,
+                                 const torch::Tensor& prefill_positions,
+                                 const torch::Tensor& prefill_delta,
+                                 const torch::Tensor& idx_mapping,
+                                 const torch::Tensor& query_start_loc,
+                                 const torch::Tensor& prefill_lens,
+                                 const torch::Tensor& num_computed_tokens) {
+  int64_t* positions_ptr = positions.data_ptr<int64_t>();
+  const int32_t* prefill_positions_ptr = prefill_positions.data_ptr<int32_t>();
+  const int32_t* prefill_delta_ptr = prefill_delta.data_ptr<int32_t>();
+  const int64_t* idx_mapping_ptr = idx_mapping.data_ptr<int64_t>();
+  const int32_t* query_start_loc_ptr = query_start_loc.data_ptr<int32_t>();
+  const int32_t* prefill_lens_ptr = prefill_lens.data_ptr<int32_t>();
+  const int32_t* num_computed_tokens_ptr =
+      num_computed_tokens.data_ptr<int32_t>();
+
+  const int64_t num_reqs = idx_mapping.size(0);
+  const int64_t num_dims = positions.size(0);
+  const int64_t positions_stride0 = positions.stride(0);
+  const int64_t positions_stride1 = positions.stride(1);
+  const int64_t prefill_positions_stride0 = prefill_positions.stride(0);
+  const int64_t prefill_positions_stride1 = prefill_positions.stride(1);
+
+  for (int64_t batch_idx = 0; batch_idx < num_reqs; ++batch_idx) {
+    const int64_t req_state_idx = idx_mapping_ptr[batch_idx];
+    const int32_t query_start = query_start_loc_ptr[batch_idx];
+    const int32_t query_end = query_start_loc_ptr[batch_idx + 1];
+    const int32_t query_len = query_end - query_start;
+    const int32_t num_computed = num_computed_tokens_ptr[req_state_idx];
+    const int32_t prompt_len = std::min(
+        query_len, std::max(prefill_lens_ptr[req_state_idx] - num_computed, 0));
+    const int32_t delta = prefill_delta_ptr[req_state_idx];
+
+    for (int64_t dim = 0; dim < num_dims; ++dim) {
+      int64_t* output = positions_ptr + dim * positions_stride0 +
+                        query_start * positions_stride1;
+      const int32_t* prompt =
+          prefill_positions_ptr +
+          (req_state_idx * num_dims + dim) * prefill_positions_stride0;
+
+      for (int32_t token_idx = 0; token_idx < prompt_len; ++token_idx) {
+        output[token_idx * positions_stride1] =
+            prompt[(num_computed + token_idx) * prefill_positions_stride1];
+      }
+      for (int32_t token_idx = prompt_len; token_idx < query_len; ++token_idx) {
+        output[token_idx * positions_stride1] =
+            num_computed + token_idx + delta;
+      }
+    }
+  }
+}
+
+torch::Tensor combine_sampled_and_draft_tokens_cpu(
+    torch::Tensor& input_ids, const torch::Tensor& idx_mapping,
+    const torch::Tensor& last_sampled_tokens,
+    const torch::Tensor& query_start_loc, const torch::Tensor& seq_lens,
+    const torch::Tensor& prefill_len, const torch::Tensor& draft_tokens,
+    const torch::Tensor& cu_num_logits, const int64_t num_logits,
+    const int64_t num_new_sampled_tokens) {
+  const int64_t* idx_mapping_ptr = idx_mapping.data_ptr<int64_t>();
+  const int64_t* last_sampled_tokens_ptr =
+      last_sampled_tokens.data_ptr<int64_t>();
+  const int32_t* query_start_loc_ptr = query_start_loc.data_ptr<int32_t>();
+  const int32_t* seq_lens_ptr = seq_lens.data_ptr<int32_t>();
+  const int32_t* prefill_len_ptr = prefill_len.data_ptr<int32_t>();
+  const int64_t* draft_tokens_ptr = draft_tokens.data_ptr<int64_t>();
+  const int32_t* cu_num_logits_ptr = cu_num_logits.data_ptr<int32_t>();
+  int32_t* input_ids_ptr = input_ids.data_ptr<int32_t>();
+
+  torch::Tensor logits_indices =
+      torch::empty({num_logits}, input_ids.options().dtype(torch::kInt64));
+  int64_t* logits_indices_ptr = logits_indices.data_ptr<int64_t>();
+
+  const int64_t num_reqs = idx_mapping.size(0);
+  const int64_t draft_tokens_stride = draft_tokens.stride(0);
+  for (int64_t req_idx = 0; req_idx < num_reqs; ++req_idx) {
+    const int64_t req_state_idx = idx_mapping_ptr[req_idx];
+    const int64_t logits_offset = cu_num_logits_ptr[req_idx];
+    const int64_t logits_offset_end = cu_num_logits_ptr[req_idx + 1];
+    const int64_t request_num_logits = logits_offset_end - logits_offset;
+    const int64_t num_draft_tokens =
+        request_num_logits - num_new_sampled_tokens;
+    const int64_t query_end = query_start_loc_ptr[req_idx + 1];
+    const int64_t logits_start = query_end - request_num_logits;
+
+    std::iota(logits_indices_ptr + logits_offset,
+              logits_indices_ptr + logits_offset_end, logits_start);
+
+    const int32_t seq_len = seq_lens_ptr[req_idx];
+    const int32_t request_prefill_len = prefill_len_ptr[req_state_idx];
+    if (seq_len <= request_prefill_len) {
+      continue;
+    }
+
+    const int64_t first_logit_seq_pos = seq_len - request_num_logits;
+    if (num_new_sampled_tokens > 0 &&
+        first_logit_seq_pos >= request_prefill_len) {
+      input_ids_ptr[logits_start] =
+          static_cast<int32_t>(last_sampled_tokens_ptr[req_state_idx]);
+    }
+
+    const int64_t draft_start = query_end - num_draft_tokens;
+    for (int64_t i = 0; i < num_draft_tokens; ++i) {
+      input_ids_ptr[draft_start + i] = static_cast<int32_t>(
+          draft_tokens_ptr[req_state_idx * draft_tokens_stride + i]);
+    }
+  }
+
+  return logits_indices;
+}
+
+void post_update_cpu(const torch::Tensor& idx_mapping,
+                     torch::Tensor& num_computed_tokens,
+                     torch::Tensor& last_sampled_tokens,
+                     const c10::optional<torch::Tensor>& output_bin_counts,
+                     const torch::Tensor& sampled_tokens,
+                     const torch::Tensor& num_sampled,
+                     const torch::Tensor& num_rejected,
+                     const c10::optional<torch::Tensor>& query_start_loc,
+                     torch::Tensor& all_token_ids, torch::Tensor& total_len) {
+  const int64_t* idx_mapping_ptr = idx_mapping.data_ptr<int64_t>();
+  int32_t* num_computed_tokens_ptr = num_computed_tokens.data_ptr<int32_t>();
+  int64_t* last_sampled_tokens_ptr = last_sampled_tokens.data_ptr<int64_t>();
+  int32_t* output_bin_counts_ptr =
+      output_bin_counts ? output_bin_counts->data_ptr<int32_t>() : nullptr;
+  const int64_t* sampled_tokens_ptr = sampled_tokens.data_ptr<int64_t>();
+  const int32_t* num_sampled_ptr = num_sampled.data_ptr<int32_t>();
+  const int32_t* num_rejected_ptr = num_rejected.data_ptr<int32_t>();
+  const int32_t* query_start_loc_ptr =
+      query_start_loc ? query_start_loc->data_ptr<int32_t>() : nullptr;
+  int32_t* all_token_ids_ptr = all_token_ids.data_ptr<int32_t>();
+  int32_t* total_len_ptr = total_len.data_ptr<int32_t>();
+
+  const int64_t num_reqs = idx_mapping.size(0);
+  const int64_t last_sampled_tokens_stride = last_sampled_tokens.stride(0);
+  const int64_t output_bin_counts_stride =
+      output_bin_counts ? output_bin_counts->stride(0) : 0;
+  const int64_t sampled_tokens_stride = sampled_tokens.stride(0);
+  const int64_t all_token_ids_stride = all_token_ids.stride(0);
+
+  for (int64_t req_idx = 0; req_idx < num_reqs; ++req_idx) {
+    const int64_t req_state_idx = idx_mapping_ptr[req_idx];
+    if (req_state_idx < 0) {
+      continue;
+    }
+
+    const int32_t request_total_len = total_len_ptr[req_state_idx];
+    const int32_t request_num_sampled = num_sampled_ptr[req_idx];
+    const int64_t* request_sampled_tokens =
+        sampled_tokens_ptr + req_idx * sampled_tokens_stride;
+    int32_t* request_all_token_ids =
+        all_token_ids_ptr + req_state_idx * all_token_ids_stride;
+
+    if (request_num_sampled > 0) {
+      last_sampled_tokens_ptr[req_state_idx * last_sampled_tokens_stride] =
+          request_sampled_tokens[request_num_sampled - 1];
+      total_len_ptr[req_state_idx] = request_total_len + request_num_sampled;
+    }
+
+    for (int32_t i = 0; i < request_num_sampled; ++i) {
+      const int64_t token_id = request_sampled_tokens[i];
+      request_all_token_ids[request_total_len + i] =
+          static_cast<int32_t>(token_id);
+      if (output_bin_counts_ptr != nullptr) {
+        ++output_bin_counts_ptr[req_state_idx * output_bin_counts_stride +
+                                token_id];
+      }
+    }
+
+    const int32_t query_len =
+        query_start_loc_ptr == nullptr
+            ? 0
+            : query_start_loc_ptr[req_idx + 1] - query_start_loc_ptr[req_idx];
+    const int32_t computed_delta = query_len - num_rejected_ptr[req_idx];
+    if (computed_delta != 0) {
+      num_computed_tokens_ptr[req_state_idx] += computed_delta;
+    }
+  }
+}
+
+}  // namespace cpu_utils
diff --git a/csrc/cpu/sampling/common.h b/csrc/cpu/sampling/common.h
new file mode 100644
index 000000000000..d18da0079c1e
--- /dev/null
+++ b/csrc/cpu/sampling/common.h
@@ -0,0 +1,279 @@
+#pragma once
+
+#include "../utils.hpp"
+
+#include <algorithm>
+#include <cmath>
+#include <cstdint>
+#include <limits>
+#include <type_traits>
+#include <vector>
+
+namespace cpu_sampling {
+
+// Draft proposals and target verification must draw from disjoint noise
+// streams even when they use the same request seed and token position.
+inline constexpr int64_t DRAFT_NOISE_SALT = 1 << 30;
+
+inline uint32_t murmur3_rotl32(uint32_t value, int shift) {
+  return (value << shift) | (value >> (32 - shift));
+}
+
+inline uint32_t murmur3_mix(uint32_t hash, uint32_t key) {
+  key *= 0xcc9e2d51;
+  key = murmur3_rotl32(key, 15);
+  key *= 0x1b873593;
+  hash ^= key;
+  hash = murmur3_rotl32(hash, 13);
+  return hash * 5 + 0xe6546b64;
+}
+
+inline uint32_t murmur3_fmix32(uint32_t hash) {
+  hash ^= hash >> 16;
+  hash *= 0x85ebca6b;
+  hash ^= hash >> 13;
+  hash *= 0xc2b2ae35;
+  return hash ^ (hash >> 16);
+}
+
+inline uint32_t murmur3_hash32_prefix(int64_t seed, int64_t position,
+                                      uint32_t domain) {
+  // Seed and position are shared by every vocabulary entry in a row. Mix this
+  // prefix once so the inner loop only needs to mix the token ID and finalize.
+  const uint64_t seed_bits = static_cast<uint64_t>(seed);
+  uint32_t hash = domain;
+  hash = murmur3_mix(hash, static_cast<uint32_t>(seed_bits));
+  hash = murmur3_mix(hash, static_cast<uint32_t>(seed_bits >> 32));
+  hash = murmur3_mix(hash, static_cast<uint32_t>(position));
+  return hash;
+}
+
+inline uint32_t murmur3_hash32(int64_t seed, int64_t position,
+                               uint32_t token_id, uint32_t domain = 0) {
+  uint32_t hash = murmur3_hash32_prefix(seed, position, domain);
+  hash = murmur3_mix(hash, token_id);
+  return murmur3_fmix32(hash ^ 16);
+}
+
+inline float murmur3_uniform32(int64_t seed, int64_t position,
+                               uint32_t token_id) {
+  const uint32_t random = murmur3_hash32(seed, position, token_id);
+  // Splitting the hash makes both integer-to-float conversions exact:
+  //   u = round_fp32((random + 0.5) / 2^32).
+  // Converting random directly to FP32 could discard up to eight low bits.
+  const int32_t high = static_cast<int32_t>(random >> 16);
+  const int32_t low = static_cast<int32_t>(random & 0xffff);
+  return static_cast<float>(high) * 1.52587890625e-05f +
+         (static_cast<float>(low) + 0.5f) * 2.3283064365386963e-10f;
+}
+
+inline double murmur3_uniform64(int64_t seed, int64_t position,
+                                uint32_t token_id) {
+  const uint64_t low = murmur3_hash32(seed, position, token_id);
+  const uint64_t high = murmur3_hash32(seed, position, token_id, 0x9e3779b9);
+  // Combine two hashes and retain the 53 bits representable by an FP64
+  // significand: u = (random53 + 0.5) / 2^53.
+  const uint64_t random53 = ((high << 32) | low) >> 11;
+  double uniform = (static_cast<double>(random53) + 0.5) * 0x1p-53;
+  if (uniform == 1.0) uniform -= 0x1p-53;
+  return uniform;
+}
+
+inline float log1p_neg_stable(float value) {
+  // This is inherited from the triton impl:
+  // For small x, evaluate the degree-8 expansion
+  //   log(1-x) = -x * (1 + x/2 + x^2/3 + ... + x^7/8)
+  // instead of forming 1-x, which loses the low bits of x.
+  float polynomial = 1.0f / 8.0f;
+  polynomial = 1.0f / 7.0f + value * polynomial;
+  polynomial = 1.0f / 6.0f + value * polynomial;
+  polynomial = 1.0f / 5.0f + value * polynomial;
+  polynomial = 1.0f / 4.0f + value * polynomial;
+  polynomial = 1.0f / 3.0f + value * polynomial;
+  polynomial = 1.0f / 2.0f + value * polynomial;
+  polynomial = 1.0f + value * polynomial;
+  const float series = -value * polynomial;
+  const float direct = std::log(std::max(1.0f - value, 5.960464477539063e-08f));
+  return value < 0.25f ? series : direct;
+}
+
+inline float gumbel_noise32(int64_t seed, int64_t position, uint32_t token_id) {
+  const float uniform = murmur3_uniform32(seed, position, token_id);
+  return -std::log(-log1p_neg_stable(uniform));
+}
+
+inline double gumbel_noise64(int64_t seed, int64_t position,
+                             uint32_t token_id) {
+  const double uniform = murmur3_uniform64(seed, position, token_id);
+  return -std::log(-std::log(uniform));
+}
+
+template <typename scalar_t>
+void scalar_greedy_argmax(const scalar_t* logits, int64_t token_begin,
+                          int64_t token_end, float& best_score,
+                          int64_t& best_idx) {
+  best_score = -std::numeric_limits<float>::infinity();
+  best_idx = token_begin;
+  for (int64_t token_id = token_begin; token_id < token_end; ++token_id) {
+    const float score = static_cast<float>(logits[token_id]);
+    if (score > best_score) {
+      best_score = score;
+      best_idx = token_id;
+    }
+  }
+}
+
+template <typename scalar_t, typename score_t, bool APPLY_TEMPERATURE>
+void scalar_gumbel_argmax(const scalar_t* logits, int64_t token_begin,
+                          int64_t token_end, int64_t seed, int64_t position,
+                          float temperature, score_t& best_score,
+                          int64_t& best_idx) {
+  best_score = -std::numeric_limits<score_t>::infinity();
+  best_idx = token_begin;
+
+  for (int64_t token_id = token_begin; token_id < token_end; ++token_id) {
+    float logit = static_cast<float>(logits[token_id]);
+    if constexpr (APPLY_TEMPERATURE) logit /= temperature;
+
+    score_t noise;
+    if constexpr (std::is_same_v<score_t, double>) {
+      noise = gumbel_noise64(seed, position, token_id);
+    } else {
+      noise = gumbel_noise32(seed, position, token_id);
+    }
+    const score_t score = static_cast<score_t>(logit) + noise;
+    if (score > best_score) {
+      best_score = score;
+      best_idx = token_id;
+    }
+  }
+}
+
+template <cpu_utils::ISA isa, typename scalar_t>
+class SamplingImpl;
+
+template <typename sampling_impl_t>
+class SamplingMainLoop {
+ public:
+  using scalar_t = typename sampling_impl_t::scalar_type;
+
+  void operator()(int64_t* output, const scalar_t* logits,
+                  int64_t logits_stride, const int64_t* expanded_idx_mapping,
+                  const float* temperature, const int64_t* seeds,
+                  const int64_t* positions, int64_t num_tokens,
+                  int64_t vocab_size, bool apply_temperature, bool is_drafting,
+                  bool use_fp64) const {
+    if (num_tokens == 0) return;
+
+    // Resolve runtime options before entering the vocabulary loop so its
+    // temperature and score-type branches are compile-time specializations.
+    if (apply_temperature) {
+      if (use_fp64) {
+        execute<double, true>(output, logits, logits_stride,
+                              expanded_idx_mapping, temperature, seeds,
+                              positions, num_tokens, vocab_size, is_drafting);
+      } else {
+        execute<float, true>(output, logits, logits_stride,
+                             expanded_idx_mapping, temperature, seeds,
+                             positions, num_tokens, vocab_size, is_drafting);
+      }
+    } else if (use_fp64) {
+      execute<double, false>(output, logits, logits_stride,
+                             expanded_idx_mapping, temperature, seeds,
+                             positions, num_tokens, vocab_size, is_drafting);
+    } else {
+      execute<float, false>(output, logits, logits_stride, expanded_idx_mapping,
+                            temperature, seeds, positions, num_tokens,
+                            vocab_size, is_drafting);
+    }
+  }
+
+ private:
+  template <typename score_t, bool APPLY_TEMPERATURE>
+  void execute(int64_t* output, const scalar_t* logits, int64_t logits_stride,
+               const int64_t* expanded_idx_mapping, const float* temperature,
+               const int64_t* seeds, const int64_t* positions,
+               int64_t num_tokens, int64_t vocab_size, bool is_drafting) const {
+    // Parallelize across logit rows first; num_tokens is not necessarily the
+    // request count. When there are fewer rows than workers, split each row to
+    // expose roughly one task per worker. Keep ranges large enough to amortize
+    // queue overhead and align them to an ISA-chosen block multiple (32 on
+    // NEON) so range boundaries do not create extra scalar tails.
+    constexpr int64_t MIN_RANGE_SIZE = 2048;
+    constexpr int64_t RANGE_ALIGNMENT = sampling_impl_t::RangeAlignment;
+    static_assert(RANGE_ALIGNMENT > 0);
+
+    const int max_workers = cpu_utils::get_max_threads();
+    const int64_t target_ranges_per_row =
+        std::max<int64_t>(1, (max_workers + num_tokens - 1) / num_tokens);
+    const int64_t unaligned_range_size = std::max<int64_t>(
+        MIN_RANGE_SIZE,
+        (vocab_size + target_ranges_per_row - 1) / target_ranges_per_row);
+    const int64_t range_size = (unaligned_range_size + RANGE_ALIGNMENT - 1) /
+                               RANGE_ALIGNMENT * RANGE_ALIGNMENT;
+    const int64_t num_ranges = (vocab_size + range_size - 1) / range_size;
+    const int64_t num_tasks = num_tokens * num_ranges;
+
+    std::vector<score_t> range_scores(num_tasks);
+    std::vector<int64_t> range_indices(num_tasks);
+    alignas(64) cpu_utils::Counter task_counter;
+    const sampling_impl_t sampling_impl;
+
+#pragma omp parallel
+    {
+      // Phase 1: find one winner per vocabulary range of each row.
+      for (;;) {
+        const int64_t task_id = task_counter.acquire_counter();
+        if (task_id >= num_tasks) break;
+
+        const int64_t token_idx = task_id / num_ranges;
+        const int64_t range_idx = task_id % num_ranges;
+        const int64_t token_begin = range_idx * range_size;
+        const int64_t token_end =
+            std::min(token_begin + range_size, vocab_size);
+        const int64_t req_state_idx = expanded_idx_mapping[token_idx];
+        const bool valid_request = req_state_idx >= 0;
+        const float temp = valid_request ? temperature[req_state_idx] : 0.0f;
+        const scalar_t* row = logits + token_idx * logits_stride;
+
+        if (temp == 0.0f) {
+          float best_score;
+          sampling_impl.greedy_argmax(row, token_begin, token_end, best_score,
+                                      range_indices[task_id]);
+          range_scores[task_id] = best_score;
+        } else {
+          int64_t position = positions[token_idx];
+          if (is_drafting) position += DRAFT_NOISE_SALT;
+          sampling_impl.template gumbel_argmax<score_t, APPLY_TEMPERATURE>(
+              row, token_begin, token_end, seeds[req_state_idx], position, temp,
+              range_scores[task_id], range_indices[task_id]);
+        }
+      }
+
+      // Every range winner must be visible before its row is reduced.
+#pragma omp barrier
+      // Phase 1 found the best token in each vocabulary range. Phase 2 assigns
+      // each row to one worker, which compares its range winners to find the
+      // best token for the complete vocabulary. Ranges are visited in order,
+      // and the strict comparison retains the first token when scores tie.
+#pragma omp for schedule(static)
+      for (int64_t token_idx = 0; token_idx < num_tokens; ++token_idx) {
+        score_t best_score = -std::numeric_limits<score_t>::infinity();
+        int64_t best_idx = 0;
+        const int64_t first_range = token_idx * num_ranges;
+        // Vectorizing this isn't worthwhile because each row has few ranges.
+        for (int64_t range_idx = 0; range_idx < num_ranges; ++range_idx) {
+          const int64_t task_id = first_range + range_idx;
+          const score_t score = range_scores[task_id];
+          if (score > best_score) {
+            best_score = score;
+            best_idx = range_indices[task_id];
+          }
+        }
+        output[token_idx] = best_idx;
+      }
+    }
+  }
+};
+
+}  // namespace cpu_sampling
diff --git a/csrc/cpu/sampling/neon.h b/csrc/cpu/sampling/neon.h
new file mode 100644
index 000000000000..b4808d5c3c09
--- /dev/null
+++ b/csrc/cpu/sampling/neon.h
@@ -0,0 +1,395 @@
+#pragma once
+
+#include "common.h"
+
+#include <arm_neon.h>
+#include <torch/types.h>
+
+#include <cstdint>
+#include <limits>
+#include <type_traits>
+
+namespace cpu_sampling::neon {
+
+FORCE_INLINE uint32x4_t murmur3_hash32x4(uint32x4_t hash,
+                                         uint32x4_t token_ids) {
+  // `hash` contains the seed/position prefix. Mix four token IDs and finalize
+  // four independent hashes, lane-for-lane matching murmur3_hash32().
+  token_ids = vmulq_n_u32(token_ids, 0xcc9e2d51);
+  token_ids = vorrq_u32(vshlq_n_u32(token_ids, 15), vshrq_n_u32(token_ids, 17));
+  token_ids = vmulq_n_u32(token_ids, 0x1b873593);
+  hash = veorq_u32(hash, token_ids);
+  hash = vorrq_u32(vshlq_n_u32(hash, 13), vshrq_n_u32(hash, 19));
+  hash = vaddq_u32(vmulq_n_u32(hash, 5), vdupq_n_u32(0xe6546b64));
+  hash = veorq_u32(hash, vdupq_n_u32(16));
+  hash = veorq_u32(hash, vshrq_n_u32(hash, 16));
+  hash = vmulq_n_u32(hash, 0x85ebca6b);
+  hash = veorq_u32(hash, vshrq_n_u32(hash, 13));
+  hash = vmulq_n_u32(hash, 0xc2b2ae35);
+  return veorq_u32(hash, vshrq_n_u32(hash, 16));
+}
+
+FORCE_INLINE float32x4_t murmur3_uniform32x4(uint32x4_t hash,
+                                             uint32x4_t token_ids) {
+  const uint32x4_t random = murmur3_hash32x4(hash, token_ids);
+  // Evaluate u = round_fp32((random + 0.5) / 2^32) without first rounding the
+  // entire 32-bit integer to FP32.
+  const uint32x4_t high = vshrq_n_u32(random, 16);
+  const uint32x4_t low = vandq_u32(random, vdupq_n_u32(0xffff));
+  return vmlaq_n_f32(vmulq_n_f32(vcvtq_f32_u32(high), 1.52587890625e-05f),
+                     vaddq_f32(vcvtq_f32_u32(low), vdupq_n_f32(0.5f)),
+                     2.3283064365386963e-10f);
+}
+
+// Adapted from data and inline_logf in Arm Optimized Routines:
+// math/aarch64/advsimd/logf.c at commit
+// 59afb6fed4275e2fc249df49fce43ee5e75af178.
+// https://github.com/ARM-software/optimized-routines
+struct PositiveNormalLogData {
+  float32x4_t c2;
+  float32x4_t c4;
+  float32x4_t c6;
+  float32x4_t ln2;
+  uint32x4_t offset;
+  uint32x4_t mantissa_mask;
+  float32x4_t odd_coefficients;
+};
+
+constexpr PositiveNormalLogData kPositiveNormalLogData = {
+    .c2 = {-0x1.4f9934p-3f, -0x1.4f9934p-3f, -0x1.4f9934p-3f, -0x1.4f9934p-3f},
+    .c4 = {-0x1.00187cp-2f, -0x1.00187cp-2f, -0x1.00187cp-2f, -0x1.00187cp-2f},
+    .c6 = {-0x1.ffffc8p-2f, -0x1.ffffc8p-2f, -0x1.ffffc8p-2f, -0x1.ffffc8p-2f},
+    .ln2 = {0x1.62e43p-1f, 0x1.62e43p-1f, 0x1.62e43p-1f, 0x1.62e43p-1f},
+    .offset = {0x3f2aaaab, 0x3f2aaaab, 0x3f2aaaab, 0x3f2aaaab},
+    .mantissa_mask = {0x007fffff, 0x007fffff, 0x007fffff, 0x007fffff},
+    .odd_coefficients = {0x1.5a9aa2p-3f, 0x1.961348p-3f, 0x1.555d7cp-2f,
+                         -0x1.3e737cp-3f},
+};
+
+// The midpoint conversion produces u = round_fp32((hash + 0.5) / 2^32), so
+// 2^-33 <= u <= 1. The direct inner call is clamped to [2^-24, 1], and the
+// value passed to the outer call is either the positive series result or
+// -log(max(1-u, 2^-24)), approximately in [2^-33, 24*ln(2)]. Both inputs are
+// therefore finite positive normal FP32 values (the smallest normal is
+// 2^-126), so no special-value path is needed.
+FORCE_INLINE float32x4_t positive_normal_log(float32x4_t value) {
+  const PositiveNormalLogData& data = kPositiveNormalLogData;
+  // Reduce x to x = 2^n * (1+r), evaluate the log1p polynomial in r, then
+  // reconstruct log(x) = n*log(2) + log1p(r).
+  const uint32x4_t offset_value =
+      vsubq_u32(vreinterpretq_u32_f32(value), data.offset);
+  const float32x4_t exponent =
+      vcvtq_f32_s32(vshrq_n_s32(vreinterpretq_s32_u32(offset_value), 23));
+  const uint32x4_t mantissa =
+      vaddq_u32(vandq_u32(offset_value, data.mantissa_mask), data.offset);
+  const float32x4_t reduced =
+      vsubq_f32(vreinterpretq_f32_u32(mantissa), vdupq_n_f32(1.0f));
+  const float32x4_t reduced_squared = vmulq_f32(reduced, reduced);
+
+  float32x4_t first =
+      vfmaq_laneq_f32(data.c2, reduced, data.odd_coefficients, 0);
+  float32x4_t second =
+      vfmaq_laneq_f32(data.c4, reduced, data.odd_coefficients, 1);
+  float32x4_t third =
+      vfmaq_laneq_f32(data.c6, reduced, data.odd_coefficients, 2);
+  first = vfmaq_laneq_f32(first, reduced_squared, data.odd_coefficients, 3);
+  second = vfmaq_f32(second, first, reduced_squared);
+  third = vfmaq_f32(third, second, reduced_squared);
+  const float32x4_t leading = vfmaq_f32(reduced, data.ln2, exponent);
+  return vfmaq_f32(leading, third, reduced_squared);
+}
+
+FORCE_INLINE float32x4_t log1p_neg_stable(float32x4_t value) {
+  // Each lane evaluates log(1-x). Use the degree-8 series near zero and the
+  // direct logarithm elsewhere; vbsl selects per lane without scalar branches.
+  float32x4_t polynomial = vdupq_n_f32(1.0f / 8.0f);
+  polynomial = vfmaq_f32(vdupq_n_f32(1.0f / 7.0f), value, polynomial);
+  polynomial = vfmaq_f32(vdupq_n_f32(1.0f / 6.0f), value, polynomial);
+  polynomial = vfmaq_f32(vdupq_n_f32(1.0f / 5.0f), value, polynomial);
+  polynomial = vfmaq_f32(vdupq_n_f32(1.0f / 4.0f), value, polynomial);
+  polynomial = vfmaq_f32(vdupq_n_f32(1.0f / 3.0f), value, polynomial);
+  polynomial = vfmaq_f32(vdupq_n_f32(1.0f / 2.0f), value, polynomial);
+  polynomial = vfmaq_f32(vdupq_n_f32(1.0f), value, polynomial);
+  const float32x4_t series = vnegq_f32(vmulq_f32(value, polynomial));
+  const float32x4_t direct =
+      positive_normal_log(vmaxq_f32(vsubq_f32(vdupq_n_f32(1.0f), value),
+                                    vdupq_n_f32(5.960464477539063e-08f)));
+  return vbslq_f32(vcltq_f32(value, vdupq_n_f32(0.25f)), series, direct);
+}
+
+FORCE_INLINE float32x4x2_t load_logits(const float* logits) {
+  return {vld1q_f32(logits), vld1q_f32(logits + 4)};
+}
+
+FORCE_INLINE float32x4x2_t load_logits(const c10::Half* logits) {
+  const auto values = vld1q_f16(reinterpret_cast<const float16_t*>(logits));
+  return {vcvt_f32_f16(vget_low_f16(values)),
+          vcvt_f32_f16(vget_high_f16(values))};
+}
+
+FORCE_INLINE float32x4x2_t load_logits(const c10::BFloat16* logits) {
+#ifdef ARM_BF16_SUPPORT
+  const auto values = vld1q_bf16(reinterpret_cast<const bfloat16_t*>(logits));
+  return {vcvt_f32_bf16(vget_low_bf16(values)),
+          vcvt_f32_bf16(vget_high_bf16(values))};
+#else
+  const auto values = vld1q_u16(reinterpret_cast<const uint16_t*>(logits));
+  return {
+      vreinterpretq_f32_u32(vshlq_n_u32(vmovl_u16(vget_low_u16(values)), 16)),
+      vreinterpretq_f32_u32(vshlq_n_u32(vmovl_u16(vget_high_u16(values)), 16)),
+  };
+#endif
+}
+
+template <typename scalar_t>
+FORCE_INLINE void greedy_argmax(const scalar_t* logits, int64_t token_begin,
+                                int64_t token_end, float& best_score,
+                                int64_t& best_idx) {
+  constexpr int64_t VEC_SIZE = 4;
+  constexpr int64_t UNROLLED_SIZE = VEC_SIZE * 4;
+  const uint32x4_t first_offsets = {0, 1, 2, 3};
+  const uint32x4_t second_offsets =
+      vaddq_u32(first_offsets, vdupq_n_u32(VEC_SIZE));
+  const uint32x4_t third_offsets =
+      vaddq_u32(second_offsets, vdupq_n_u32(VEC_SIZE));
+  const uint32x4_t fourth_offsets =
+      vaddq_u32(third_offsets, vdupq_n_u32(VEC_SIZE));
+
+  float32x4_t maximum = vdupq_n_f32(-std::numeric_limits<float>::infinity());
+  uint32x4_t maximum_indices = vdupq_n_u32(static_cast<uint32_t>(token_begin));
+
+  int64_t token_id = token_begin;
+  const int64_t vector_end =
+      token_end - (token_end - token_begin) % UNROLLED_SIZE;
+  for (; token_id < vector_end; token_id += UNROLLED_SIZE) {
+    // Reduce 16 logits into four lanes. Lane j carries the winner among token
+    // offsets j, j+4, j+8, and j+12, together with its absolute token ID.
+    const float32x4x2_t first = load_logits(logits + token_id);
+    const float32x4x2_t second = load_logits(logits + token_id + 2 * VEC_SIZE);
+
+    const float32x4_t first_half = vmaxnmq_f32(first.val[0], first.val[1]);
+    const uint32x4_t first_half_offsets = vbslq_u32(
+        vceqq_f32(first.val[0], first_half), first_offsets, second_offsets);
+    const float32x4_t second_half = vmaxnmq_f32(second.val[0], second.val[1]);
+    const uint32x4_t second_half_offsets = vbslq_u32(
+        vceqq_f32(second.val[0], second_half), third_offsets, fourth_offsets);
+    const float32x4_t block_maximum = vmaxnmq_f32(first_half, second_half);
+    const uint32x4_t block_offsets =
+        vbslq_u32(vceqq_f32(first_half, block_maximum), first_half_offsets,
+                  second_half_offsets);
+    const uint32x4_t block_indices =
+        vaddq_u32(vdupq_n_u32(static_cast<uint32_t>(token_id)), block_offsets);
+    const uint32x4_t update = vcgtq_f32(block_maximum, maximum);
+    maximum = vbslq_f32(update, block_maximum, maximum);
+    maximum_indices = vbslq_u32(update, block_indices, maximum_indices);
+  }
+
+  if (token_end - token_id >= 2 * VEC_SIZE) {
+    const float32x4x2_t values = load_logits(logits + token_id);
+    const float32x4_t block_maximum = vmaxnmq_f32(values.val[0], values.val[1]);
+    const uint32x4_t block_offsets = vbslq_u32(
+        vceqq_f32(values.val[0], block_maximum), first_offsets, second_offsets);
+    const uint32x4_t block_indices =
+        vaddq_u32(vdupq_n_u32(static_cast<uint32_t>(token_id)), block_offsets);
+    const uint32x4_t update = vcgtq_f32(block_maximum, maximum);
+    maximum = vbslq_f32(update, block_maximum, maximum);
+    maximum_indices = vbslq_u32(update, block_indices, maximum_indices);
+    token_id += 2 * VEC_SIZE;
+  }
+
+  best_score = vmaxvq_f32(maximum);
+  // More than one lane can contain the global maximum. Select the smallest
+  // associated token ID to preserve first-occurrence argmax semantics.
+  const uint32x4_t matching_indices =
+      vbslq_u32(vceqq_f32(maximum, vdupq_n_f32(best_score)), maximum_indices,
+                vdupq_n_u32(std::numeric_limits<uint32_t>::max()));
+  best_idx = vminvq_u32(matching_indices);
+
+  for (; token_id < token_end; ++token_id) {
+    const float score = static_cast<float>(logits[token_id]);
+    if (score > best_score) {
+      best_score = score;
+      best_idx = token_id;
+    }
+  }
+}
+
+template <typename scalar_t, bool APPLY_TEMPERATURE>
+FORCE_INLINE void score_gumbel_block(const scalar_t* logits, int64_t token_id,
+                                     uint32_t token_offset, uint32x4_t hash,
+                                     uint32x4_t offsets, float temperature,
+                                     float32x4_t& winning_scores,
+                                     uint32x4_t& winning_token_offsets) {
+  constexpr int64_t VEC_SIZE = 8;
+  // For each of eight tokens, compute
+  //
+  //   u = (Murmur3(seed, position, token) + 0.5) / 2^32
+  //   g = -log(-log(1-u))
+  //   scaled_logit = APPLY_TEMPERATURE ? logit / temperature : logit
+  //   score = scaled_logit + g.
+  //
+  // The two four-lane score vectors are folded lane-wise. Output lane j is the
+  // winner between token offsets j and j+4; token_offset makes that ID relative
+  // to the caller's outer block (0, 8, 16, or 24 in the 32-token loop).
+  const uint32x4_t first_tokens =
+      vaddq_u32(vdupq_n_u32(static_cast<uint32_t>(token_id)), offsets);
+  const uint32x4_t second_tokens =
+      vaddq_u32(first_tokens, vdupq_n_u32(VEC_SIZE / 2));
+  const float32x4x2_t uniforms = {murmur3_uniform32x4(hash, first_tokens),
+                                  murmur3_uniform32x4(hash, second_tokens)};
+  const float32x4x2_t inner_log = {
+      vnegq_f32(log1p_neg_stable(uniforms.val[0])),
+      vnegq_f32(log1p_neg_stable(uniforms.val[1]))};
+  const float32x4x2_t noise = {
+      vnegq_f32(positive_normal_log(inner_log.val[0])),
+      vnegq_f32(positive_normal_log(inner_log.val[1]))};
+
+  float32x4x2_t logit = load_logits(logits + token_id);
+  if constexpr (APPLY_TEMPERATURE) {
+    // Match Triton's temperature division. Its cached-logits path divides by
+    // the same temperature on reload and relies on reproducing this value
+    // bitwise.
+    const float32x4_t divisor = vdupq_n_f32(temperature);
+    logit.val[0] = vdivq_f32(logit.val[0], divisor);
+    logit.val[1] = vdivq_f32(logit.val[1], divisor);
+  }
+  const float32x4_t first_scores = vaddq_f32(logit.val[0], noise.val[0]);
+  const float32x4_t second_scores = vaddq_f32(logit.val[1], noise.val[1]);
+  winning_scores = vmaxnmq_f32(first_scores, second_scores);
+  const uint32x4_t first_token_offsets =
+      vaddq_u32(offsets, vdupq_n_u32(token_offset));
+  const uint32x4_t second_token_offsets =
+      vaddq_u32(first_token_offsets, vdupq_n_u32(VEC_SIZE / 2));
+  winning_token_offsets = vbslq_u32(vceqq_f32(first_scores, winning_scores),
+                                    first_token_offsets, second_token_offsets);
+}
+
+template <typename scalar_t, bool APPLY_TEMPERATURE>
+void gumbel_argmax(const scalar_t* logits, int64_t token_begin,
+                   int64_t token_end, int64_t seed, int64_t position,
+                   float temperature, float& best_score, int64_t& best_idx) {
+  constexpr int64_t VEC_SIZE = 8;
+  constexpr int64_t UNROLLED_SIZE = VEC_SIZE * 4;
+  const int64_t vector_end =
+      token_end - (token_end - token_begin) % UNROLLED_SIZE;
+  const uint32x4_t hash = vdupq_n_u32(murmur3_hash32_prefix(seed, position, 0));
+  const uint32x4_t offsets = {0, 1, 2, 3};
+
+  best_score = -std::numeric_limits<float>::infinity();
+  best_idx = token_begin;
+
+  int64_t token_id = token_begin;
+  for (; token_id < vector_end; token_id += UNROLLED_SIZE) {
+    // Keep four independent eight-token pipelines in flight to hide the log
+    // latency.
+    float32x4_t first_maximum;
+    uint32x4_t first_token_offsets;
+    score_gumbel_block<scalar_t, APPLY_TEMPERATURE>(
+        logits, token_id, 0, hash, offsets, temperature, first_maximum,
+        first_token_offsets);
+    float32x4_t second_maximum;
+    uint32x4_t second_token_offsets;
+    score_gumbel_block<scalar_t, APPLY_TEMPERATURE>(
+        logits, token_id + VEC_SIZE, VEC_SIZE, hash, offsets, temperature,
+        second_maximum, second_token_offsets);
+    float32x4_t third_maximum;
+    uint32x4_t third_token_offsets;
+    score_gumbel_block<scalar_t, APPLY_TEMPERATURE>(
+        logits, token_id + 2 * VEC_SIZE, 2 * VEC_SIZE, hash, offsets,
+        temperature, third_maximum, third_token_offsets);
+    float32x4_t fourth_maximum;
+    uint32x4_t fourth_token_offsets;
+    score_gumbel_block<scalar_t, APPLY_TEMPERATURE>(
+        logits, token_id + 3 * VEC_SIZE, 3 * VEC_SIZE, hash, offsets,
+        temperature, fourth_maximum, fourth_token_offsets);
+
+    // reduce the four 8 token pipelines to find the block's argmax
+    const float32x4_t first_half_maximum =
+        vmaxnmq_f32(first_maximum, second_maximum);
+    const uint32x4_t first_half_token_offsets =
+        vbslq_u32(vceqq_f32(first_maximum, first_half_maximum),
+                  first_token_offsets, second_token_offsets);
+    const float32x4_t second_half_maximum =
+        vmaxnmq_f32(third_maximum, fourth_maximum);
+    const uint32x4_t second_half_token_offsets =
+        vbslq_u32(vceqq_f32(third_maximum, second_half_maximum),
+                  third_token_offsets, fourth_token_offsets);
+    const float32x4_t block_maximum =
+        vmaxnmq_f32(first_half_maximum, second_half_maximum);
+    const uint32x4_t block_token_offsets =
+        vbslq_u32(vceqq_f32(first_half_maximum, block_maximum),
+                  first_half_token_offsets, second_half_token_offsets);
+
+    const float score = vmaxnmvq_f32(block_maximum);
+    if (score > best_score) {
+      const uint32x4_t matching_offsets = vbslq_u32(
+          vceqq_f32(block_maximum, vdupq_n_f32(score)), block_token_offsets,
+          vdupq_n_u32(std::numeric_limits<uint32_t>::max()));
+      best_score = score;
+      best_idx = token_id + vminvq_u32(matching_offsets);
+    }
+  }
+
+  const int64_t vector_tail_end = token_end - (token_end - token_id) % VEC_SIZE;
+  for (; token_id < vector_tail_end; token_id += VEC_SIZE) {
+    float32x4_t maximum;
+    uint32x4_t token_offsets;
+    score_gumbel_block<scalar_t, APPLY_TEMPERATURE>(logits, token_id, 0, hash,
+                                                    offsets, temperature,
+                                                    maximum, token_offsets);
+    const float score = vmaxnmvq_f32(maximum);
+    if (score > best_score) {
+      const uint32x4_t matching_offsets =
+          vbslq_u32(vceqq_f32(maximum, vdupq_n_f32(score)), token_offsets,
+                    vdupq_n_u32(std::numeric_limits<uint32_t>::max()));
+      best_score = score;
+      best_idx = token_id + vminvq_u32(matching_offsets);
+    }
+  }
+
+  for (; token_id < token_end; ++token_id) {
+    float logit = static_cast<float>(logits[token_id]);
+    if constexpr (APPLY_TEMPERATURE) logit /= temperature;
+    const float score = logit + gumbel_noise32(seed, position, token_id);
+    if (score > best_score) {
+      best_score = score;
+      best_idx = token_id;
+    }
+  }
+}
+
+}  // namespace cpu_sampling::neon
+
+namespace cpu_sampling {
+
+template <typename scalar_t>
+class SamplingImpl<cpu_utils::ISA::NEON, scalar_t> {
+ public:
+  using scalar_type = scalar_t;
+  static constexpr int64_t RangeAlignment = 32;
+
+  void greedy_argmax(const scalar_t* logits, int64_t token_begin,
+                     int64_t token_end, float& best_score,
+                     int64_t& best_idx) const {
+    neon::greedy_argmax(logits, token_begin, token_end, best_score, best_idx);
+  }
+
+  template <typename score_t, bool APPLY_TEMPERATURE>
+  void gumbel_argmax(const scalar_t* logits, int64_t token_begin,
+                     int64_t token_end, int64_t seed, int64_t position,
+                     float temperature, score_t& best_score,
+                     int64_t& best_idx) const {
+    if constexpr (std::is_same_v<score_t, double>) {
+      // TODO: do we need to vectorize this?
+      scalar_gumbel_argmax<scalar_t, double, APPLY_TEMPERATURE>(
+          logits, token_begin, token_end, seed, position, temperature,
+          best_score, best_idx);
+    } else {
+      neon::gumbel_argmax<scalar_t, APPLY_TEMPERATURE>(
+          logits, token_begin, token_end, seed, position, temperature,
+          best_score, best_idx);
+    }
+  }
+};
+
+}  // namespace cpu_sampling
diff --git a/csrc/cpu/sampling/vec.h b/csrc/cpu/sampling/vec.h
new file mode 100644
index 000000000000..bb9f28355d78
--- /dev/null
+++ b/csrc/cpu/sampling/vec.h
@@ -0,0 +1,30 @@
+#pragma once
+
+#include "common.h"
+
+namespace cpu_sampling {
+
+template <typename scalar_t>
+class SamplingImpl<cpu_utils::ISA::VEC, scalar_t> {
+ public:
+  using scalar_type = scalar_t;
+  static constexpr int64_t RangeAlignment = 1;
+
+  void greedy_argmax(const scalar_t* logits, int64_t token_begin,
+                     int64_t token_end, float& best_score,
+                     int64_t& best_idx) const {
+    scalar_greedy_argmax(logits, token_begin, token_end, best_score, best_idx);
+  }
+
+  template <typename score_t, bool APPLY_TEMPERATURE>
+  void gumbel_argmax(const scalar_t* logits, int64_t token_begin,
+                     int64_t token_end, int64_t seed, int64_t position,
+                     float temperature, score_t& best_score,
+                     int64_t& best_idx) const {
+    scalar_gumbel_argmax<scalar_t, score_t, APPLY_TEMPERATURE>(
+        logits, token_begin, token_end, seed, position, temperature, best_score,
+        best_idx);
+  }
+};
+
+}  // namespace cpu_sampling
diff --git a/csrc/cpu/sampling_kernels.cpp b/csrc/cpu/sampling_kernels.cpp
index 1d30b92a8d9a..ccca5b12ce9a 100644
--- a/csrc/cpu/sampling_kernels.cpp
+++ b/csrc/cpu/sampling_kernels.cpp
@@ -1,7 +1,23 @@
 #include "cpu_types.hpp"
 
+#if defined(__aarch64__)
+  #include "sampling/neon.h"
+
+template <typename scalar_t>
+using SelectedSamplingImpl =
+    cpu_sampling::SamplingImpl<cpu_utils::ISA::NEON, scalar_t>;
+#else
+  #include "sampling/vec.h"
+
+template <typename scalar_t>
+using SelectedSamplingImpl =
+    cpu_sampling::SamplingImpl<cpu_utils::ISA::VEC, scalar_t>;
+#endif
+
 #include <torch/all.h>
 
+#include <algorithm>
+#include <array>
 #include <cmath>
 #include <cstdint>
 #include <limits>
@@ -166,3 +182,133 @@ torch::Tensor greedy_argmax(const torch::Tensor& logits) {
                        logits_contig.size(1));
   return output;
 }
+
+torch::Tensor gumbel_sample_cpu(const torch::Tensor& logits,
+                                const torch::Tensor& expanded_idx_mapping,
+                                const torch::Tensor& temperature,
+                                const torch::Tensor& seeds,
+                                const torch::Tensor& positions,
+                                bool apply_temperature, bool is_drafting,
+                                bool use_fp64) {
+  TORCH_CHECK(logits.dim() == 2, "logits must be 2-D [num_tokens, vocab]");
+  TORCH_CHECK(logits.stride(1) == 1,
+              "the logits vocabulary dimension must be contiguous");
+
+  auto output =
+      torch::empty({logits.size(0)}, logits.options().dtype(torch::kInt64));
+  VLLM_DISPATCH_FLOATING_TYPES(logits.scalar_type(), "gumbel_sample_cpu", [&] {
+    using sampling_impl_t = SelectedSamplingImpl<scalar_t>;
+    cpu_sampling::SamplingMainLoop<sampling_impl_t> mainloop;
+    mainloop(output.data_ptr<int64_t>(), logits.data_ptr<scalar_t>(),
+             logits.stride(0), expanded_idx_mapping.data_ptr<int64_t>(),
+             temperature.data_ptr<float>(), seeds.data_ptr<int64_t>(),
+             positions.data_ptr<int64_t>(), logits.size(0), logits.size(1),
+             apply_temperature, is_drafting, use_fp64);
+  });
+  return output;
+}
+
+FORCE_INLINE void fill_logits(float* logits, int64_t size, float value) {
+  constexpr int64_t VEC_SIZE = vec_op::FP32Vec16::VEC_ELEM_NUM;
+  const vec_op::FP32Vec16 values(value);
+  const int64_t vectorized_size = size - size % VEC_SIZE;
+  int64_t offset = 0;
+  for (; offset < vectorized_size; offset += VEC_SIZE) {
+    values.save(logits + offset);
+  }
+  std::fill(logits + offset, logits + size, value);
+}
+
+void apply_logit_bias_cpu(
+    torch::Tensor& logits, const torch::Tensor& expanded_idx_mapping,
+    const torch::Tensor& positions, const torch::Tensor& num_allowed_token_ids,
+    const torch::Tensor& allowed_token_ids, const torch::Tensor& num_logit_bias,
+    const torch::Tensor& logit_bias_token_ids, const torch::Tensor& logit_bias,
+    const torch::Tensor& min_lens, const torch::Tensor& num_stop_token_ids,
+    const torch::Tensor& restore_when_all_masked,
+    const torch::Tensor& stop_token_ids, bool check_all_masked_rows) {
+  constexpr int64_t MAX_ALLOWED_TOKEN_IDS = 1024;
+  constexpr int64_t MAX_STOP_TOKEN_IDS = 128;
+  const float negative_infinity = -std::numeric_limits<float>::infinity();
+  const int64_t num_tokens = logits.size(0);
+  const int64_t vocab_size = logits.size(1);
+
+  float* logits_ptr = logits.data_ptr<float>();
+  const int64_t* idx_mapping_ptr = expanded_idx_mapping.data_ptr<int64_t>();
+  const int64_t* positions_ptr = positions.data_ptr<int64_t>();
+  const int32_t* num_allowed_ptr = num_allowed_token_ids.data_ptr<int32_t>();
+  const int32_t* allowed_ptr = allowed_token_ids.data_ptr<int32_t>();
+  const int32_t* num_bias_ptr = num_logit_bias.data_ptr<int32_t>();
+  const int32_t* bias_token_ids_ptr = logit_bias_token_ids.data_ptr<int32_t>();
+  const float* bias_ptr = logit_bias.data_ptr<float>();
+  const int32_t* min_lens_ptr = min_lens.data_ptr<int32_t>();
+  const int32_t* num_stop_ptr = num_stop_token_ids.data_ptr<int32_t>();
+  const int32_t* restore_ptr = restore_when_all_masked.data_ptr<int32_t>();
+  const int32_t* stop_token_ids_ptr = stop_token_ids.data_ptr<int32_t>();
+
+  const int64_t logits_stride = logits.stride(0);
+  const int64_t allowed_stride = allowed_token_ids.stride(0);
+  const int64_t bias_ids_stride = logit_bias_token_ids.stride(0);
+  const int64_t bias_stride = logit_bias.stride(0);
+  const int64_t stop_stride = stop_token_ids.stride(0);
+
+#pragma omp parallel for schedule(static) if (num_tokens > 1)
+  for (int64_t token_idx = 0; token_idx < num_tokens; ++token_idx) {
+    const int64_t req_idx = idx_mapping_ptr[token_idx];
+    float* row = logits_ptr + token_idx * logits_stride;
+
+    const int32_t num_allowed = num_allowed_ptr[req_idx];
+    if (num_allowed > 0) {
+      std::array<float, MAX_ALLOWED_TOKEN_IDS> saved_logits;
+      const int32_t* row_allowed = allowed_ptr + req_idx * allowed_stride;
+      for (int32_t i = 0; i < num_allowed; ++i) {
+        saved_logits[i] = row[row_allowed[i]];
+      }
+      fill_logits(row, vocab_size, negative_infinity);
+      for (int32_t i = 0; i < num_allowed; ++i) {
+        row[row_allowed[i]] = saved_logits[i];
+      }
+    }
+
+    const int32_t num_bias = num_bias_ptr[req_idx];
+    const int32_t* row_bias_ids =
+        bias_token_ids_ptr + req_idx * bias_ids_stride;
+    const float* row_bias = bias_ptr + req_idx * bias_stride;
+    for (int32_t i = 0; i < num_bias; ++i) {
+      row[row_bias_ids[i]] += row_bias[i];
+    }
+
+    const int32_t num_stop = num_stop_ptr[req_idx];
+    if (num_stop == 0 ||
+        positions_ptr[token_idx] + 1 >= min_lens_ptr[req_idx]) {
+      continue;
+    }
+
+    const int32_t* row_stop_ids = stop_token_ids_ptr + req_idx * stop_stride;
+    if (!check_all_masked_rows || restore_ptr[req_idx] == 0) {
+      for (int32_t i = 0; i < num_stop; ++i) {
+        row[row_stop_ids[i]] = negative_infinity;
+      }
+      continue;
+    }
+
+    std::array<float, MAX_STOP_TOKEN_IDS> saved_stop_logits;
+    for (int32_t i = 0; i < num_stop; ++i) {
+      saved_stop_logits[i] = row[row_stop_ids[i]];
+      row[row_stop_ids[i]] = negative_infinity;
+    }
+
+    const bool all_masked =
+        std::all_of(row, row + vocab_size, [negative_infinity](float value) {
+          return value == negative_infinity;
+        });
+    if (all_masked) {
+      for (int32_t i = 0; i < num_stop; ++i) {
+        const float value = saved_stop_logits[i];
+        if (std::isfinite(value)) {
+          row[row_stop_ids[i]] = value;
+        }
+      }
+    }
+  }
+}
diff --git a/csrc/cpu/torch_bindings.cpp b/csrc/cpu/torch_bindings.cpp
index c68d5c6317df..cffa908841e3 100644
--- a/csrc/cpu/torch_bindings.cpp
+++ b/csrc/cpu/torch_bindings.cpp
@@ -406,9 +406,77 @@ void dynamic_per_token_scaled_fp8_quant(
     torch::Tensor& out, const torch::Tensor& input, torch::Tensor& scale,
     const std::optional<torch::Tensor>& scale_ub);
 
+torch::Tensor gumbel_sample_cpu(const torch::Tensor& logits,
+                                const torch::Tensor& expanded_idx_mapping,
+                                const torch::Tensor& temperature,
+                                const torch::Tensor& seeds,
+                                const torch::Tensor& positions,
+                                bool apply_temperature, bool is_drafting,
+                                bool use_fp64);
+
+void apply_logit_bias_cpu(
+    torch::Tensor& logits, const torch::Tensor& expanded_idx_mapping,
+    const torch::Tensor& positions, const torch::Tensor& num_allowed_token_ids,
+    const torch::Tensor& allowed_token_ids, const torch::Tensor& num_logit_bias,
+    const torch::Tensor& logit_bias_token_ids, const torch::Tensor& logit_bias,
+    const torch::Tensor& min_lens, const torch::Tensor& num_stop_token_ids,
+    const torch::Tensor& restore_when_all_masked,
+    const torch::Tensor& stop_token_ids, bool check_all_masked_rows);
+
 void init_cpu_memory_env(std::vector<int64_t> node_ids);
 
 namespace cpu_utils {
+// Kernels used only by standalone CPU MRV2.
+#ifdef VLLM_STANDALONE_CPU_MRV2
+void gather_block_table_cpu(const torch::Tensor& block_table,
+                            torch::Tensor& output,
+                            const torch::Tensor& idx_mapping,
+                            const torch::Tensor& num_blocks,
+                            const int64_t num_reqs_padded);
+void compute_slot_mapping_cpu(
+    const torch::Tensor& block_table, torch::Tensor& slot_mapping,
+    const torch::Tensor& idx_mapping, const torch::Tensor& query_start_loc,
+    const torch::Tensor& positions, const int64_t block_size,
+    const bool mapping_enabled, const int64_t pad_slot_id);
+void prepare_prefill_inputs_cpu(torch::Tensor& input_ids,
+                                torch::Tensor& next_prefill_tokens,
+                                const torch::Tensor& idx_mapping,
+                                const torch::Tensor& query_start_loc,
+                                const torch::Tensor& all_token_ids,
+                                const torch::Tensor& prefill_len,
+                                const torch::Tensor& num_computed_tokens);
+torch::Tensor get_prompt_logprobs_token_ids_cpu(
+    int64_t num_tokens, const torch::Tensor& query_start_loc,
+    const torch::Tensor& idx_mapping, const torch::Tensor& num_computed_tokens,
+    const torch::Tensor& all_token_ids);
+void prepare_pos_seq_lens_cpu(const torch::Tensor& idx_mapping,
+                              const torch::Tensor& query_start_loc,
+                              const torch::Tensor& num_computed_tokens,
+                              torch::Tensor& pos, torch::Tensor& seq_lens);
+void prepare_mrope_positions_cpu(torch::Tensor& positions,
+                                 const torch::Tensor& prefill_positions,
+                                 const torch::Tensor& prefill_delta,
+                                 const torch::Tensor& idx_mapping,
+                                 const torch::Tensor& query_start_loc,
+                                 const torch::Tensor& prefill_lens,
+                                 const torch::Tensor& num_computed_tokens);
+torch::Tensor combine_sampled_and_draft_tokens_cpu(
+    torch::Tensor& input_ids, const torch::Tensor& idx_mapping,
+    const torch::Tensor& last_sampled_tokens,
+    const torch::Tensor& query_start_loc, const torch::Tensor& seq_lens,
+    const torch::Tensor& prefill_len, const torch::Tensor& draft_tokens,
+    const torch::Tensor& cu_num_logits, const int64_t num_logits,
+    const int64_t num_new_sampled_tokens);
+void post_update_cpu(const torch::Tensor& idx_mapping,
+                     torch::Tensor& num_computed_tokens,
+                     torch::Tensor& last_sampled_tokens,
+                     const c10::optional<torch::Tensor>& output_bin_counts,
+                     const torch::Tensor& sampled_tokens,
+                     const torch::Tensor& num_sampled,
+                     const torch::Tensor& num_rejected,
+                     const c10::optional<torch::Tensor>& query_start_loc,
+                     torch::Tensor& all_token_ids, torch::Tensor& total_len);
+#endif
 void eagle_prepare_inputs_padded_kernel_impl(
     const torch::Tensor& cu_num_draft_tokens,
     const torch::Tensor& valid_sampled_tokens_count,
@@ -999,6 +1067,62 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
 
   ops.def("init_cpu_memory_env(SymInt[] node_ids) -> ()", &init_cpu_memory_env);
 
+#ifdef VLLM_STANDALONE_CPU_MRV2
+  ops.def(
+      "gather_block_table_cpu(Tensor block_table, Tensor(a1!) output, "
+      "Tensor idx_mapping, Tensor num_blocks, SymInt num_reqs_padded) -> ()",
+      &cpu_utils::gather_block_table_cpu);
+
+  ops.def(
+      "compute_slot_mapping_cpu("
+      "Tensor block_table, Tensor(a1!) slot_mapping, Tensor idx_mapping, "
+      "Tensor query_start_loc, Tensor positions, SymInt block_size, "
+      "bool mapping_enabled, SymInt pad_slot_id) -> ()",
+      &cpu_utils::compute_slot_mapping_cpu);
+
+  ops.def(
+      "prepare_prefill_inputs_cpu(Tensor(a0!) input_ids, "
+      "Tensor(a1!) next_prefill_tokens, Tensor idx_mapping, "
+      "Tensor query_start_loc, Tensor all_token_ids, Tensor prefill_len, "
+      "Tensor num_computed_tokens) -> ()",
+      &cpu_utils::prepare_prefill_inputs_cpu);
+
+  ops.def(
+      "get_prompt_logprobs_token_ids_cpu("
+      "SymInt num_tokens, Tensor query_start_loc, Tensor idx_mapping, "
+      "Tensor num_computed_tokens, Tensor all_token_ids) -> Tensor",
+      &cpu_utils::get_prompt_logprobs_token_ids_cpu);
+
+  ops.def(
+      "prepare_pos_seq_lens_cpu(Tensor idx_mapping, Tensor query_start_loc, "
+      "Tensor num_computed_tokens, Tensor(a3!) pos, "
+      "Tensor(a4!) seq_lens) -> ()",
+      &cpu_utils::prepare_pos_seq_lens_cpu);
+
+  ops.def(
+      "prepare_mrope_positions_cpu("
+      "Tensor(a0!) positions, Tensor prefill_positions, Tensor prefill_delta, "
+      "Tensor idx_mapping, Tensor query_start_loc, Tensor prefill_lens, "
+      "Tensor num_computed_tokens) -> ()",
+      &cpu_utils::prepare_mrope_positions_cpu);
+
+  ops.def(
+      "combine_sampled_and_draft_tokens_cpu("
+      "Tensor(a0!) input_ids, Tensor idx_mapping, Tensor last_sampled_tokens, "
+      "Tensor query_start_loc, Tensor seq_lens, Tensor prefill_len, "
+      "Tensor draft_tokens, Tensor cu_num_logits, SymInt num_logits, "
+      "SymInt num_new_sampled_tokens) -> Tensor",
+      &cpu_utils::combine_sampled_and_draft_tokens_cpu);
+
+  ops.def(
+      "post_update_cpu(Tensor idx_mapping, "
+      "Tensor(a1!) num_computed_tokens, Tensor(a2!) last_sampled_tokens, "
+      "Tensor(a3!)? output_bin_counts, Tensor sampled_tokens, "
+      "Tensor num_sampled, Tensor num_rejected, Tensor? query_start_loc, "
+      "Tensor(a8!) all_token_ids, Tensor(a9!) total_len) -> ()",
+      &cpu_utils::post_update_cpu);
+#endif
+
   // Fused sampling kernels
   ops.def("fused_gumbel_argmax(Tensor logits, Tensor seeds) -> Tensor");
   ops.impl("fused_gumbel_argmax", torch::kCPU, &fused_gumbel_argmax);
@@ -1006,6 +1130,22 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
   ops.def("greedy_argmax(Tensor logits) -> Tensor");
   ops.impl("greedy_argmax", torch::kCPU, &greedy_argmax);
 
+  ops.def(
+      "gumbel_sample_cpu(Tensor logits, Tensor expanded_idx_mapping, "
+      "Tensor temperature, Tensor seeds, Tensor positions, "
+      "bool apply_temperature, bool is_drafting, bool use_fp64) -> Tensor");
+  ops.impl("gumbel_sample_cpu", torch::kCPU, &gumbel_sample_cpu);
+
+  ops.def(
+      "apply_logit_bias_cpu(Tensor(a0!) logits, "
+      "Tensor expanded_idx_mapping, Tensor positions, "
+      "Tensor num_allowed_token_ids, Tensor allowed_token_ids, "
+      "Tensor num_logit_bias, Tensor logit_bias_token_ids, "
+      "Tensor logit_bias, Tensor min_lens, Tensor num_stop_token_ids, "
+      "Tensor restore_when_all_masked, Tensor stop_token_ids, "
+      "bool check_all_masked_rows) -> ()",
+      &apply_logit_bias_cpu);
+
   // Speculative decoding kernels
   ops.def(
       "eagle_prepare_inputs_padded_kernel_impl(Tensor cu_num_draft_tokens, "
diff --git a/csrc/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu b/csrc/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu
index 09484daf3315..11e76e6d9873 100644
--- a/csrc/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu
+++ b/csrc/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu
@@ -48,9 +48,11 @@
  * ``qkv`` tensor.  Caches (bf16) are scatter-written by slot.
  */
 
+#include <algorithm>
 #include <cmath>
 #include <cuda_runtime.h>
 #include <type_traits>
+#include <limits>
 
 #include "torch_utils.h"
 
@@ -271,11 +273,43 @@ __device__ __forceinline__ void storeElemsFp8(
 
 // Match scaled_fp8_quant(q_out): materialize q in scalar_t before applying the
 // inverse dequantization scale and converting it to E4M3.
+#ifndef USE_ROCM
+// storeScaledQElemsFp8 with inv_scale = 1, from the packed model-dtype values:
+// E4M3 of a scalar_t equals E4M3 of its float, the max takes NaN and values
+// below -448 to -448 as storeElemsFp8's fmaxf does, and the cvt saturates the
+// top at 448.
+template <typename scalar_t>
+__device__ __forceinline__ void storeUnitScaleQElemsFp8(
+    uint8_t* __restrict__ dst, float const (&elems)[kElemsPerLane]) {
+  using Converter = vllm::_typeConvert<scalar_t>;
+  __nv_fp8x2_storage_t out2[kElemsPerLane / 2];
+  #pragma unroll
+  for (int i = 0; i < kElemsPerLane / 2; i++) {
+    auto const v =
+        Converter::convert(make_float2(elems[2 * i], elems[2 * i + 1]));
+    if constexpr (std::is_same_v<scalar_t, c10::BFloat16>) {
+      out2[i] = __nv_cvt_bfloat16raw2_to_fp8x2(
+          __hmax2(v, __float2bfloat162_rn(-448.0f)), __NV_SATFINITE, __NV_E4M3);
+    } else {
+      out2[i] = __nv_cvt_halfraw2_to_fp8x2(
+          __hmax2(v, __float2half2_rn(-448.0f)), __NV_SATFINITE, __NV_E4M3);
+    }
+  }
+  *reinterpret_cast<uint32_t*>(dst) = *reinterpret_cast<uint32_t const*>(out2);
+}
+#endif
+
 template <typename scalar_t>
 __device__ __forceinline__ void storeScaledQElemsFp8(
     uint8_t* __restrict__ dst, float const (&elems)[kElemsPerLane],
     float const inv_scale) {
   using Converter = vllm::_typeConvert<scalar_t>;
+#ifndef USE_ROCM
+  if (inv_scale == 1.0f) {
+    storeUnitScaleQElemsFp8<scalar_t>(dst, elems);
+    return;
+  }
+#endif
   float scaled[kElemsPerLane];
 #pragma unroll
   for (int i = 0; i < kElemsPerLane; i++) {
@@ -424,27 +458,14 @@ __device__ __forceinline__ uint16_t clearNegZeroNibbles(uint16_t packed) {
   return packed & (((packed & 0x7777u) + 0x7777u) | 0x7777u);
 }
 
-// Quantize one K or V head row of one token into its NVFP4 page, where
-// ``scale`` is the dequantization (global) scale. Every lane of the warp must
-// call this (the scale-group reductions shuffle).
+// The E2M1 codes (low nibble first) of four values x, already rounded to the
+// model dtype, of a 16-value group with absolute maximum group_max, and the
+// group's E4M3 scale in ``sf_byte``, where ``scale`` is the dequantization
+// (global) scale.
 template <typename scalar_t>
-__device__ __forceinline__ void storeNvfp4CacheElems(
-    uint8_t* __restrict__ page, int const kv, int const head, int const token,
-    int const block_size, int const laneId, float const (&elems)[kElemsPerLane],
-    float const scale) {
-  // Quantize the model-dtype value that the unfused path reads back from qkv.
-  using Converter = vllm::_typeConvert<scalar_t>;
-  float x[kElemsPerLane];
-  float group_max = 0.0f;
-  #pragma unroll
-  for (int i = 0; i < kElemsPerLane; i++) {
-    x[i] = static_cast<float>(Converter::convert(elems[i]));
-    group_max = fmaxf(group_max, fabsf(x[i]));
-  }
-  group_max = fmaxf(group_max, __shfl_xor_sync(FINAL_MASK, group_max, 1, 32));
-  group_max = fmaxf(group_max, __shfl_xor_sync(FINAL_MASK, group_max, 2, 32));
-
-  uint8_t sf_byte;
+__device__ __forceinline__ uint16_t
+nvfp4EncodeElems(float const (&x)[kElemsPerLane], float const group_max,
+                 float const scale, uint8_t& sf_byte) {
   float sf_rcp;
   float q[kElemsPerLane];  // (x / scale) / sf_value, E2M1-encoded below
   if (scale == 1.0f) {
@@ -516,8 +537,49 @@ __device__ __forceinline__ void storeNvfp4CacheElems(
   for (int i = 0; i < kElemsPerLane; i++) codes |= e2m1Code(q[i]) << (4 * i);
   packed = static_cast<uint16_t>(codes);
   #endif
-  packed = clearNegZeroNibbles(packed);
+  return clearNegZeroNibbles(packed);
+}
+
+// Offset, within a K (kv 0) or V (kv 1) head's scale block, of the E4M3 scale
+// of a token's 16-value group: linear for K, the SM100 trtllm-gen 4x4
+// token-quad swizzle for V.
+__device__ __forceinline__ int nvfp4ScaleOffset(int const kv, int const token,
+                                                int const group) {
+  if (kv == 0) return token * kNvfp4ScaleBytes + group;
+  constexpr int kGroupsPerQuad = kNvfp4ScaleBytes / 4;
+  int const swizzled_t = (token / 4) * 4 + group / kGroupsPerQuad;
+  int const swizzled_s = (group % kGroupsPerQuad) * 4 + token % 4;
+  return swizzled_t * kNvfp4ScaleBytes + swizzled_s;
+}
+
+// The E2M1 codes of this lane's four values of a K or V head row, rounded to
+// the model dtype, and in ``sf_byte`` its 16-value group's E4M3 scale, where
+// ``scale`` is the dequantization (global) scale. Every lane of the warp must
+// call this (the scale-group reductions shuffle).
+template <typename scalar_t>
+__device__ __forceinline__ uint16_t nvfp4QuantizeElems(
+    float const (&elems)[kElemsPerLane], float const scale, uint8_t& sf_byte) {
+  // Quantize the model-dtype value that the unfused path reads back from qkv.
+  using Converter = vllm::_typeConvert<scalar_t>;
+  float x[kElemsPerLane];
+  float group_max = 0.0f;
+  #pragma unroll
+  for (int i = 0; i < kElemsPerLane; i++) {
+    x[i] = static_cast<float>(Converter::convert(elems[i]));
+    group_max = fmaxf(group_max, fabsf(x[i]));
+  }
+  group_max = fmaxf(group_max, __shfl_xor_sync(FINAL_MASK, group_max, 1, 32));
+  group_max = fmaxf(group_max, __shfl_xor_sync(FINAL_MASK, group_max, 2, 32));
+
+  return nvfp4EncodeElems<scalar_t>(x, group_max, scale, sf_byte);
+}
 
+// Store the codes and scale from nvfp4QuantizeElems of lane laneId's four
+// values of a token's K (kv 0) or V (kv 1) head row.
+__device__ __forceinline__ void nvfp4StoreElems(
+    uint8_t* __restrict__ page, int const kv, int const head, int const token,
+    int const block_size, int const laneId, uint16_t const packed,
+    uint8_t const sf_byte) {
   uint8_t* const slot = page + static_cast<int64_t>(2 * head + kv) *
                                    block_size *
                                    (kNvfp4DataBytes + kNvfp4ScaleBytes);
@@ -526,19 +588,57 @@ __device__ __forceinline__ void storeNvfp4CacheElems(
   *reinterpret_cast<uint16_t*>(data) = packed;
 
   if (laneId % kNvfp4LanesPerGroup == 0) {
-    int const group = laneId / kNvfp4LanesPerGroup;
-    int offset;
-    if (kv == 0) {
-      offset = token * kNvfp4ScaleBytes + group;
-    } else {
-      constexpr int kGroupsPerQuad = kNvfp4ScaleBytes / 4;
-      int const swizzled_t = (token / 4) * 4 + group / kGroupsPerQuad;
-      int const swizzled_s = (group % kGroupsPerQuad) * 4 + token % 4;
-      offset = swizzled_t * kNvfp4ScaleBytes + swizzled_s;
-    }
-    slot[block_size * kNvfp4DataBytes + offset] = sf_byte;
+    slot[block_size * kNvfp4DataBytes +
+         nvfp4ScaleOffset(kv, token, laneId / kNvfp4LanesPerGroup)] = sf_byte;
   }
 }
+
+// nvfp4QuantizeElems for one lane holding a whole 16-value group x (already
+// rounded to the model dtype): the group max is order-independent, so this is
+// the same codes and scale without the shuffles. Low nibble first.
+template <typename scalar_t>
+__device__ __forceinline__ uint2 nvfp4QuantizeGroup(float const (&x)[16],
+                                                    float const scale,
+                                                    uint8_t& sf_byte) {
+  float group_max = 0.0f;
+  #pragma unroll
+  for (int i = 0; i < 16; i++) group_max = fmaxf(group_max, fabsf(x[i]));
+  uint32_t codes[4];
+  #pragma unroll
+  for (int k = 0; k < 4; k++) {
+    float const xk[kElemsPerLane] = {x[4 * k], x[4 * k + 1], x[4 * k + 2],
+                                     x[4 * k + 3]};
+    codes[k] = nvfp4EncodeElems<scalar_t>(xk, group_max, scale, sf_byte);
+  }
+  return make_uint2(codes[0] | codes[1] << 16, codes[2] | codes[3] << 16);
+}
+
+// Store group ``group`` (its 8 code bytes and its scale) of a token's K
+// (kv 0) or V (kv 1) head row into its NVFP4 page.
+__device__ __forceinline__ void nvfp4StoreGroup(
+    uint8_t* __restrict__ page, int const kv, int const head, int const token,
+    int const block_size, int const group, uint2 const codes,
+    uint8_t const sf_byte) {
+  uint8_t* const slot = page + static_cast<int64_t>(2 * head + kv) *
+                                   block_size *
+                                   (kNvfp4DataBytes + kNvfp4ScaleBytes);
+  *reinterpret_cast<uint2*>(slot + token * kNvfp4DataBytes + 8 * group) = codes;
+  slot[block_size * kNvfp4DataBytes + nvfp4ScaleOffset(kv, token, group)] =
+      sf_byte;
+}
+
+// Quantize one K or V head row of one token into its NVFP4 page, where
+// ``scale`` is the dequantization (global) scale. Every lane of the warp must
+// call this (the scale-group reductions shuffle).
+template <typename scalar_t>
+__device__ __forceinline__ void storeNvfp4CacheElems(
+    uint8_t* __restrict__ page, int const kv, int const head, int const token,
+    int const block_size, int const laneId, float const (&elems)[kElemsPerLane],
+    float const scale) {
+  uint8_t sf_byte;
+  uint16_t const packed = nvfp4QuantizeElems<scalar_t>(elems, scale, sf_byte);
+  nvfp4StoreElems(page, kv, head, token, block_size, laneId, packed, sf_byte);
+}
 #endif  // USE_ROCM
 
 // ────────────────────────────────────────────────────────────────────────────
@@ -553,18 +653,21 @@ __device__ __forceinline__ void storeNvfp4CacheElems(
 // Slots per token:
 //     Q : nq                            (always — norm+RoPE)
 //     K : nkv                           (always — norm+RoPE; +K-cache insert)
-//     V : nkv  only if kInsertKV        (V-cache insert; no warps in dense)
+//     V : nkv  only if kInsertKV        (V-cache insert)
 //     IQ: niq  only if kProcessIndex    (norm+RoPE)
 //     IK: 1    only if kProcessIndex    (norm+RoPE; +index-cache insert)
 // cache_t/kv_dt: main attention KV-cache dtype (auto/fp8). out_idx_t/kFp8Idx:
 // indexer index-K cache + index-Q output dtype (scalar_t or e4m3 byte).
+// kNvfp4Idx (implies kFp8Idx): the index-K cache is instead packed NVFP4 pages
+// [nb, bs, 72] (bs*64 B E2M1, then bs*8 B linear E4M3 scales, no global scale),
+// the layout of the MSA Q8KV4 indexer; index-Q stays e4m3.
 // kHasIndex means the qkv row is laid out as sparse [q|k|v|index_q|index_k].
 // kProcessIndex controls whether this launch actually norms/ropes the index
 // branch and writes index_q/index_k outputs. Skip-index-topk reuse layers keep
 // kHasIndex=true but set kProcessIndex=false.
 template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt,
           bool kNvfp4, typename out_idx_t, bool kHasIndex, bool kInsertKV,
-          bool kProcessIndex, bool kFp8Idx>
+          bool kProcessIndex, bool kFp8Idx, bool kNvfp4Idx>
 __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
     scalar_t* __restrict__ qkv,  // [N, qkv_row] in/out (packs index if sparse)
     scalar_t* __restrict__ q_out,         // [N, nq*128] contiguous, or nullptr
@@ -588,6 +691,8 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
     // NHD/HND layout choice is captured by the head/token strides.
     int64_t const kv_s_block, int64_t const kv_s_head, int64_t const kv_s_token,
     int64_t const kv_s_dim,
+    // NVFP4 index cache only: page stride in bytes.
+    int64_t const idx_s_block,
     // NVFP4 main cache only: per-layer K/V dequantization (global) scales.
     float const* __restrict__ kv_k_scale,
     float const* __restrict__ kv_v_scale) {
@@ -604,8 +709,11 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
 
     static_assert(!kProcessIndex || kHasIndex,
                   "index processing requires sparse row layout");
+    static_assert(!kNvfp4Idx || (kFp8Idx && kInsertKV),
+                  "NVFP4 index cache requires fp8 index-Q and cache insert");
 
-    // Slot layout (compile-time gated: dense has neither V nor index slots).
+    // Slot layout (compile-time gated: V slots only when inserting, index
+    // slots only on the sparse layer).
     int const v_slots = kInsertKV ? nkv : 0;
     int const idx_slots = kProcessIndex ? niq + 1 : 0;
     int const slots_per_token = nq + nkv + v_slots + idx_slots;
@@ -743,7 +851,7 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
       }
     }
 
-    // ── Cache inserts (sparse serving only). ───────────────────────────────
+    // ── Cache inserts. ─────────────────────────────────────────────────────
     if constexpr (kInsertKV) {
       // Guard (not early-return) so every thread reaches the PDL trigger below.
       int64_t sm = -1;
@@ -754,7 +862,16 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
       }
       if (sm >= 0) {  // skip padded / unscheduled tokens
         if (isIK) {
-          if constexpr (kFp8Idx) {
+          if constexpr (kNvfp4Idx) {
+#ifndef USE_ROCM
+            // A single-head K side: data [bs, 64] then linear scales [bs, 8].
+            storeNvfp4CacheElems<scalar_t>(
+                reinterpret_cast<uint8_t*>(index_cache) +
+                    (sm / block_size) * idx_s_block,
+                /*kv=*/0, /*head=*/0, static_cast<int>(sm % block_size),
+                block_size, laneId, elems, /*scale=*/1.0f);
+#endif
+          } else if constexpr (kFp8Idx) {
             storeElemsFp8(index_cache + sm * kHeadDim + dim_base, elems);
           } else {
             storeElems<scalar_t>(index_cache + sm * kHeadDim + dim_base, elems);
@@ -779,6 +896,15 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
           int64_t const off = b * kv_s_block + head * kv_s_head +
                               t * kv_s_token +
                               (kv * kHeadDim + dim_base) * kv_s_dim;
+          if constexpr (!kHasIndex) {
+            // Dense layers stand in for reshape_and_cache_flash, which
+            // quantizes the k already materialized in the model dtype.
+            using Converter = vllm::_typeConvert<scalar_t>;
+#pragma unroll
+            for (int i = 0; i < kElemsPerLane; i++) {
+              elems[i] = Converter::convert(Converter::convert(elems[i]));
+            }
+          }
           storeCacheElems<scalar_t, cache_t, kv_dt>(kv_cache + off, elems);
         }
       }
@@ -794,6 +920,754 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
 #endif
 }
 
+#ifndef USE_ROCM
+// ────────────────────────────────────────────────────────────────────────────
+// Persistent row-group kernel
+// ────────────────────────────────────────────────────────────────────────────
+// Writes the same bytes as the kernel above, for every instantiation of it,
+// when rotary_dim == 64. A token's work is a group of warps with fixed roles:
+// the q heads kQHeads (4 or 8) per warp, then K, V and index_q four heads per
+// warp, then index_k. A branch of fewer than four heads, and index_k, takes
+// one warp per head in the layout of the kernel above (lane l holds dims
+// 4l .. 4l + 3); the others run four-row steps, where lane (g, v) holds dims
+// 32c + 4v .. 32c + 4v + 3 (c = 0..3) of row g, the dims lane 8c + v holds
+// above, so the stores call the same helpers with that lane id. A CTA runs
+// some or all warps of the warp groups of tokens x, x + gridDim.x, ..., and
+// each warp copies the qkv rows and cos/sin row of its next tokens into a
+// kStages-deep shared-memory ring with cp.async.
+namespace rg {
+
+constexpr int kRows = 4;
+constexpr int kChunks = 4;
+constexpr int kRotaryDim = 64;
+constexpr int kStages = 3;  // ring depth: two tokens in flight per warp
+// A ring row: 256 B padded to 320 B, so that rows 1 apart start 16 banks apart
+// and a 4-row step's 8-byte loads are conflict-free.
+constexpr int kRowPitch = kHeadDim * 2 + 64;
+constexpr int kCosSinBytes = kRotaryDim * 2;
+enum Role : int { kQ = 0, kK, kV, kIQ, kIK };
+
+// SM100's packed FP32 instructions where they exist; the same RN results.
+__device__ __forceinline__ float2 mul2(float2 const a, float2 const b) {
+  #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000
+  return __fmul2_rn(a, b);
+  #else
+  return make_float2(__fmul_rn(a.x, b.x), __fmul_rn(a.y, b.y));
+  #endif
+}
+
+__device__ __forceinline__ float2 fma2(float2 const a, float2 const b,
+                                       float2 const c) {
+  #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000
+  return __ffma2_rn(a, b, c);
+  #else
+  return make_float2(__fmaf_rn(a.x, b.x, c.x), __fmaf_rn(a.y, b.y, c.y));
+  #endif
+}
+
+__device__ __forceinline__ float2 add2(float2 const a, float2 const b) {
+  #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000
+  return __fadd2_rn(a, b);
+  #else
+  return make_float2(__fadd_rn(a.x, b.x), __fadd_rn(a.y, b.y));
+  #endif
+}
+
+// A head row as one lane holds it, paired for packed FP32 math:
+// x[k][i] = element i of chunks 2k and 2k + 1.
+using Row = float2[2][4];
+
+template <typename scalar_t>
+__device__ __forceinline__ void loadRow(uint8_t const* row, int const v,
+                                        Row& x) {
+  float e[kChunks][kElemsPerLane];
+  #pragma unroll
+  for (int c = 0; c < kChunks; c++) {
+    loadElems<scalar_t>(reinterpret_cast<scalar_t const*>(row + 64 * c + 8 * v),
+                        e[c]);
+  }
+  #pragma unroll
+  for (int k = 0; k < 2; k++) {
+  #pragma unroll
+    for (int i = 0; i < kElemsPerLane; i++) {
+      x[k][i] = make_float2(e[2 * k][i], e[2 * k + 1][i]);
+    }
+  }
+}
+
+__device__ __forceinline__ void chunkElems(Row const& x, int const c,
+                                           float (&e)[kElemsPerLane]) {
+  #pragma unroll
+  for (int i = 0; i < kElemsPerLane; i++) {
+    e[i] = c % 2 ? x[c / 2][i].y : x[c / 2][i].x;
+  }
+}
+
+// normAndRope (Gemma RMSNorm, then NeoX RoPE on dims [0, 64)) with its
+// rounding as compiled: an FFMA sum of squares per lane of the one-warp-per-row
+// layout (one chunk here), the xor-16/8/4/2/1 tree (16 pairs chunks c and
+// c ^ 2, 8 pairs c and c ^ 1), FFMA(sum, 1/128, eps), rsqrtf, x * rms * w, and
+// x * cos -/+ RN(partner * sin) as one FFMA.
+__device__ __forceinline__ void normRope(Row& x, Row const& w1, float const eps,
+                                         float const (&cos4)[4],
+                                         float const (&sin4)[4]) {
+  float2 sq[2];
+  #pragma unroll
+  for (int k = 0; k < 2; k++) {
+    sq[k] = mul2(x[k][0], x[k][0]);  // = FFMA(x, x, 0)
+  #pragma unroll
+    for (int i = 1; i < 4; i++) sq[k] = fma2(x[k][i], x[k][i], sq[k]);
+  }
+  float2 const xor16 = add2(sq[0], sq[1]);
+  float sum = __fadd_rn(xor16.x, xor16.y);
+  sum = __fadd_rn(sum, __shfl_xor_sync(FINAL_MASK, sum, 4, 32));
+  sum = __fadd_rn(sum, __shfl_xor_sync(FINAL_MASK, sum, 2, 32));
+  sum = __fadd_rn(sum, __shfl_xor_sync(FINAL_MASK, sum, 1, 32));
+  float const rms_rcp =
+      rsqrtf(__fmaf_rn(sum, 1.0f / static_cast<float>(kHeadDim), eps));
+  float2 const r2 = make_float2(rms_rcp, rms_rcp);
+  #pragma unroll
+  for (int k = 0; k < 2; k++) {
+  #pragma unroll
+    for (int i = 0; i < 4; i++) x[k][i] = mul2(mul2(x[k][i], r2), w1[k][i]);
+  }
+  // x[0][i] = (x1, x2): dims 4v + i and 32 + 4v + i.
+  #pragma unroll
+  for (int i = 0; i < 4; i++) {
+    float2 const x_sin = mul2(x[0][i], make_float2(sin4[i], sin4[i]));
+    x[0][i] = make_float2(__fmaf_rn(x[0][i].x, cos4[i], -x_sin.y),
+                          __fmaf_rn(x[0][i].y, cos4[i], x_sin.x));
+  }
+}
+
+__device__ __forceinline__ void cpAsync16(uint32_t const dst, void const* src) {
+  asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" ::"r"(dst),
+               "l"(src)
+               : "memory");
+}
+
+__device__ __forceinline__ void cpAsyncCommit() {
+  asm volatile("cp.async.commit_group;\n" ::: "memory");
+}
+
+template <int kPending>
+__device__ __forceinline__ void cpAsyncWait() {
+  asm volatile("cp.async.wait_group %0;\n" ::"n"(kPending) : "memory");
+}
+
+// Warp roles of a token's warp group: the q heads kQHeads at a time, then K,
+// V and index_q four heads at a time, then index_k.
+// A branch of fewer than kRows heads (and index_k) takes one warp per head in
+// the layout of the kernel above instead of a part-empty four-row step.
+struct Roles {
+  int nq, nkv, niq;
+  bool kv_one, iq_one;  // one warp per head
+  int q_warps, kv_warps, v_warps, iq_warps, ik_warps;
+  int warps;  // one token's warp group
+
+  __host__ __device__ Roles(int const q_heads, int const nq_, int const nkv_,
+                            int const niq_, bool const insert_kv,
+                            bool const process_index)
+      : nq(nq_), nkv(nkv_), niq(niq_) {
+    kv_one = nkv < kRows;
+    iq_one = niq < kRows;
+    q_warps = (nq + q_heads - 1) / q_heads;
+    kv_warps = kv_one ? nkv : (nkv + kRows - 1) / kRows;
+    v_warps = insert_kv ? kv_warps : 0;
+    iq_warps = process_index ? (iq_one ? niq : (niq + kRows - 1) / kRows) : 0;
+    ik_warps = process_index ? 1 : 0;
+    warps = q_warps + kv_warps + v_warps + iq_warps + ik_warps;
+  }
+
+  // Shared-memory offset of warp w's ring: the q warps' rings, then the
+  // others'.
+  __host__ __device__ int ringOffset(int const w, int const q_ring,
+                                     int const side_ring) const {
+    return min(w, q_warps) * q_ring + max(w - q_warps, 0) * side_ring;
+  }
+};
+
+// CTA size limit, for 80 registers per thread.
+constexpr int kMaxWarps = 24;
+
+template <int kQHeads>
+constexpr int stageBytes() {
+  return kQHeads * kRowPitch + kCosSinBytes;
+}
+
+// What a warp writes besides the caches, fixed per launch so that the stores
+// of the other modes are not compiled into the loop: q (kQFp8Unit: only
+// q_fp8_out, unit scale; kQFp8: only q_fp8_out; kQAny: decided per store),
+// index_q (kIQOut: index_q_out; kIQInPlace: back into qkv; kNone), and the
+// in-place k / index_k rewrite (kInPlace or kNone).
+enum Out : int {
+  kNone = 0,
+  kInPlace,
+  kQFp8Unit,
+  kQFp8,
+  kQAny,
+  kIQOut,
+  kIQInPlace
+};
+
+// One warp's loop over its CTA's tokens: rows [head0, head0 + nrows) of one
+// branch, kSteps four-row steps per token.
+template <int kRole, int kSteps, int kOut, typename scalar_t, typename cache_t,
+          Fp8KVCacheDataType kv_dt, bool kNvfp4, typename out_idx_t,
+          bool kHasIndex, bool kInsertKV, bool kFp8Idx, bool kNvfp4Idx,
+          bool kOneRow = false>
+__device__ __forceinline__ void runWarp(
+    int const head0, int const nrows, int const qkv_head0, uint8_t* const ring,
+    int const stage_bytes, scalar_t* __restrict__ qkv,
+    scalar_t* __restrict__ q_out, uint8_t* __restrict__ q_fp8_out,
+    out_idx_t* __restrict__ index_q_out, scalar_t const* __restrict__ norm_w,
+    scalar_t const* __restrict__ cos_sin_cache,
+    int64_t const* __restrict__ positions,
+    int64_t const* __restrict__ slot_mapping, cache_t* __restrict__ kv_cache,
+    out_idx_t* __restrict__ index_cache, float const eps,
+    float const q_fp8_inv_scale, int const num_tokens, int const nq,
+    int const niq, int const qkv_row, int const block_size,
+    int64_t const kv_s_block, int64_t const kv_s_head, int64_t const kv_s_token,
+    int64_t const kv_s_dim, int64_t const idx_s_block, float const kv_scale,
+    bool const write_qkv) {
+  constexpr bool kNorm = kRole != kV;
+  constexpr bool kSlots =
+      kInsertKV && (kRole == kK || kRole == kV || kRole == kIK);
+  // kOneRow: a single row in the layout of the kernel above (lane l holds
+  // dims 4l .. 4l + 3), for branches of one head.
+  constexpr int kStepRows = kOneRow ? 1 : kRows;
+  constexpr int kLaneChunks = kOneRow ? 1 : kChunks;
+  int const lane = threadIdx.x % 32;
+  int const g = kOneRow ? 0 : lane / 8;
+  int const v = kOneRow ? lane : lane % 8;
+  constexpr int kCopies = (kSteps * kStepRows * 16 + 31) / 32;
+  uint32_t const ring_smem =
+      static_cast<uint32_t>(__cvta_generic_to_shared(ring));
+
+  // Not produced by the predecessor: loaded before the PDL wait, used after.
+  uint2 w_raw[kChunks];
+  if constexpr (kNorm && !kOneRow) {
+  #pragma unroll
+    for (int c = 0; c < kChunks; c++) {
+      w_raw[c] = __ldg(reinterpret_cast<uint2 const*>(norm_w + 32 * c + 4 * v));
+    }
+  }
+
+  // PDL: qkv comes from the predecessor (the qkv GEMM).
+  #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+  cudaGridDependencySynchronize();
+  #endif
+
+  Row w1;
+  if constexpr (kNorm && !kOneRow) {
+    float e[kChunks][kElemsPerLane];
+  #pragma unroll
+    for (int c = 0; c < kChunks; c++) {
+      loadElems<scalar_t>(reinterpret_cast<scalar_t const*>(&w_raw[c]), e[c]);
+    }
+  #pragma unroll
+    for (int k = 0; k < 2; k++) {
+  #pragma unroll
+      for (int i = 0; i < kElemsPerLane; i++) {
+        w1[k][i] = make_float2(1.0f + e[2 * k][i], 1.0f + e[2 * k + 1][i]);
+      }
+    }
+  }
+
+  int const stride = gridDim.x;
+  int const t0 = blockIdx.x;
+  int const n = (num_tokens - t0 + stride - 1) / stride;
+  // Lane j holds the position (slot) of the token of iteration base + j.
+  auto fetch = [&](int64_t const* src, int const base, int64_t const fill) {
+    int64_t const t = t0 + static_cast<int64_t>(base + lane) * stride;
+    return t < num_tokens ? static_cast<int64_t>(__ldg(
+                                reinterpret_cast<long long const*>(src) + t))
+                          : fill;
+  };
+  int pos_base = 0, slot_base = 0;
+  int64_t pos_l = 0, slot_l = -1;
+  if constexpr (kNorm) pos_l = fetch(positions, 0, 0);
+  if constexpr (kSlots) slot_l = fetch(slot_mapping, 0, -1);
+
+  scalar_t const* const src0 =
+      qkv + static_cast<int64_t>(t0) * qkv_row + qkv_head0 * kHeadDim;
+  int64_t const src_step = static_cast<int64_t>(stride) * qkv_row;
+  auto issue_rows = [&](int const it) {
+    if (it >= n) return;
+    scalar_t const* const src = src0 + it * src_step;
+    uint32_t const stage = ring_smem + (it % kStages) * stage_bytes;
+  #pragma unroll
+    for (int k = 0; k < kCopies; k++) {
+      int const r = lane / 16 + 2 * k;
+      if (r < nrows && (!kOneRow || lane < 16)) {
+        cpAsync16(stage + r * kRowPitch + lane % 16 * 16,
+                  src + r * kHeadDim + lane % 16 * 8);
+      }
+    }
+  };
+  // The cos/sin row of iteration it, then the commit that closes its group.
+  auto issue_cos_sin = [&](int const it) {
+    if constexpr (kNorm) {
+      if (it < n) {
+        if (it - pos_base == 32) {
+          pos_base = it;
+          pos_l = fetch(positions, pos_base, 0);
+        }
+        int64_t const pos = __shfl_sync(
+            FINAL_MASK, static_cast<long long>(pos_l), it - pos_base, 32);
+        if (lane < kCosSinBytes / 16) {
+          cpAsync16(ring_smem + (it % kStages) * stage_bytes +
+                        kSteps * kStepRows * kRowPitch + lane * 16,
+                    cos_sin_cache + pos * kRotaryDim + lane * 8);
+        }
+      }
+    }
+    cpAsyncCommit();
+  };
+
+  // The prologue's rows do not wait for the positions.
+  #pragma unroll
+  for (int it = 0; it < kStages - 1; it++) issue_rows(it);
+  #pragma unroll
+  for (int it = 0; it < kStages - 1; it++) issue_cos_sin(it);
+
+  // This lane's dims 4v .. 4v + 3 (+ 32c) of row g of the first step, at the
+  // first token; each iteration advances them by stride tokens. Dims
+  // 32c + 4v of a row are lane 8c + v's dims in the kernel above.
+  scalar_t* in_place = qkv + static_cast<int64_t>(t0) * qkv_row +
+                       (qkv_head0 + g) * kHeadDim + 4 * v;
+  int const head_rows = kRole == kQ ? nq : kRole == kIQ ? niq : 0;
+  int64_t const out_lane =
+      (static_cast<int64_t>(t0) * head_rows + head0 + g) * kHeadDim + 4 * v;
+  int64_t const out_step = static_cast<int64_t>(stride) * head_rows * kHeadDim;
+  int64_t out = out_lane;
+
+  for (int it = 0; it < n; it++) {
+    issue_rows(it + kStages - 1);
+    issue_cos_sin(it + kStages - 1);
+    cpAsyncWait<kStages - 1>();
+    __syncwarp();
+    uint8_t const* const stage = ring + (it % kStages) * stage_bytes;
+    float cos4[4], sin4[4];
+    if constexpr (kNorm) {
+      scalar_t const* const cs = reinterpret_cast<scalar_t const*>(
+          stage + kSteps * kStepRows * kRowPitch);
+      if constexpr (!kOneRow) {
+        loadElems<scalar_t>(cs + 4 * v, cos4);
+        loadElems<scalar_t>(cs + kRotaryDim / 2 + 4 * v, sin4);
+      }
+    }
+    // Slot of the token: block and offset with one 32-bit division (the host
+    // checks that every slot fits).
+    int sm = -1;
+    uint32_t blk = 0, tok = 0;
+    if constexpr (kSlots) {
+      if (it - slot_base == 32) {
+        slot_base = it;
+        slot_l = fetch(slot_mapping, slot_base, -1);
+      }
+      sm = static_cast<int>(__shfl_sync(
+          FINAL_MASK, static_cast<long long>(slot_l), it - slot_base, 32));
+      blk = static_cast<uint32_t>(sm) / static_cast<uint32_t>(block_size);
+      tok = static_cast<uint32_t>(sm) - blk * block_size;
+    }
+
+  #pragma unroll
+    for (int step = 0; step < kSteps; step++) {
+      int const r = step * kStepRows + g;  // row of the warp
+      int const head = head0 + r;
+      bool const row_ok = r < nrows;
+      int const so = step * kStepRows * kHeadDim;  // step offset, elements
+      float e[kChunks][kElemsPerLane];
+      if constexpr (kOneRow) {
+        loadElems<scalar_t>(
+            reinterpret_cast<scalar_t const*>(stage + r * kRowPitch) + 4 * v,
+            e[0]);
+        if constexpr (kNorm) {
+          normAndRope<scalar_t>(
+              e[0], lane, eps, norm_w, /*do_rope=*/true, kRotaryDim,
+              reinterpret_cast<scalar_t const*>(stage + kSteps * kRowPitch),
+              /*apply_norm=*/true);
+        }
+      } else {
+        Row x;
+        loadRow<scalar_t>(stage + r * kRowPitch, v, x);
+        if constexpr (kNorm) normRope(x, w1, eps, cos4, sin4);
+  #pragma unroll
+        for (int c = 0; c < kChunks; c++) chunkElems(x, c, e[c]);
+      }
+      if constexpr (kRole == kQ) {
+        if (row_ok) {
+  #pragma unroll
+          for (int c = 0; c < kLaneChunks; c++) {
+            if constexpr (kOut == kQFp8Unit) {
+              storeUnitScaleQElemsFp8<scalar_t>(q_fp8_out + out + so + 32 * c,
+                                                e[c]);
+            } else if constexpr (kOut == kQFp8) {
+              storeScaledQElemsFp8<scalar_t>(q_fp8_out + out + so + 32 * c,
+                                             e[c], q_fp8_inv_scale);
+            } else {
+              if (q_out != nullptr) {
+                storeElems<scalar_t>(q_out + out + so + 32 * c, e[c]);
+              } else if (q_fp8_out == nullptr && write_qkv) {
+                storeElems<scalar_t>(in_place + so + 32 * c, e[c]);
+              }
+              if (q_fp8_out != nullptr) {
+                storeScaledQElemsFp8<scalar_t>(q_fp8_out + out + so + 32 * c,
+                                               e[c], q_fp8_inv_scale);
+              }
+            }
+          }
+        }
+      } else if constexpr (kRole == kIQ) {
+        if (row_ok) {
+  #pragma unroll
+          for (int c = 0; c < kLaneChunks; c++) {
+            if constexpr (kOut == kIQOut) {
+              if constexpr (kFp8Idx) {
+                storeIndexQElemsFp8<scalar_t>(
+                    reinterpret_cast<uint8_t*>(index_q_out + out + so + 32 * c),
+                    e[c]);
+              } else {
+                storeElems<scalar_t>(index_q_out + out + so + 32 * c, e[c]);
+              }
+            } else if constexpr (kOut == kIQInPlace) {
+              storeElems<scalar_t>(in_place + so + 32 * c, e[c]);
+            }
+          }
+        }
+      } else if constexpr (kRole == kK || kRole == kIK) {
+        if constexpr (kOut == kInPlace) {
+          if (row_ok) {
+  #pragma unroll
+            for (int c = 0; c < kLaneChunks; c++) {
+              storeElems<scalar_t>(in_place + so + 32 * c, e[c]);
+            }
+          }
+        }
+      }
+
+      if constexpr (kInsertKV && kRole == kIK) {
+        if constexpr (kNvfp4Idx) {
+          uint8_t* const page =
+              reinterpret_cast<uint8_t*>(index_cache) + blk * idx_s_block;
+  #pragma unroll
+          for (int c = 0; c < kLaneChunks; c++) {
+            uint8_t sf;
+            uint16_t const packed =
+                nvfp4QuantizeElems<scalar_t>(e[c], 1.0f, sf);
+            if (sm >= 0 && row_ok) {
+              nvfp4StoreElems(page, /*kv=*/0, /*head=*/0, tok, block_size,
+                              kOneRow ? lane : 8 * c + v, packed, sf);
+            }
+          }
+        } else if (sm >= 0 && row_ok) {
+          int64_t const off = static_cast<int64_t>(sm) * kHeadDim + 4 * v;
+  #pragma unroll
+          for (int c = 0; c < kLaneChunks; c++) {
+            if constexpr (kFp8Idx) {
+              storeElemsFp8(
+                  reinterpret_cast<uint8_t*>(index_cache) + off + 32 * c, e[c]);
+            } else {
+              storeElems<scalar_t>(index_cache + off + 32 * c, e[c]);
+            }
+          }
+        }
+      } else if constexpr (kInsertKV && (kRole == kK || kRole == kV)) {
+        constexpr int kv = kRole == kK ? 0 : 1;
+        if constexpr (kNvfp4 && !kOneRow) {
+          // One 16-value scale group per lane, regrouped through the ring
+          // row: lane (g, v) takes dims 16v .. 16v + 15 of row g.
+          uint8_t* const row =
+              ring + (it % kStages) * stage_bytes + r * kRowPitch;
+          if constexpr (kRole == kK) {
+            // The k in the model dtype, as the helper would quantize it (V
+            // rows already hold their model-dtype values).
+  #pragma unroll
+            for (int c = 0; c < kChunks; c++) {
+              storeElems<scalar_t>(
+                  reinterpret_cast<scalar_t*>(row + 64 * c + 8 * v), e[c]);
+            }
+          }
+          __syncwarp();
+          float x[16];
+  #pragma unroll
+          for (int k = 0; k < 4; k++) {
+            float q[kElemsPerLane];
+            loadElems<scalar_t>(
+                reinterpret_cast<scalar_t const*>(row + 32 * v + 8 * k), q);
+  #pragma unroll
+            for (int i = 0; i < kElemsPerLane; i++) x[4 * k + i] = q[i];
+          }
+          uint8_t sf;
+          uint2 const codes = nvfp4QuantizeGroup<scalar_t>(x, kv_scale, sf);
+          if (sm >= 0 && row_ok) {
+            nvfp4StoreGroup(
+                reinterpret_cast<uint8_t*>(kv_cache) + blk * kv_s_block, kv,
+                head, tok, block_size, v, codes, sf);
+          }
+        } else if constexpr (kNvfp4) {
+          uint8_t* const page =
+              reinterpret_cast<uint8_t*>(kv_cache) + blk * kv_s_block;
+  #pragma unroll
+          for (int c = 0; c < kLaneChunks; c++) {
+            uint8_t sf;
+            uint16_t const packed =
+                nvfp4QuantizeElems<scalar_t>(e[c], kv_scale, sf);
+            if (sm >= 0 && row_ok) {
+              nvfp4StoreElems(page, kv, head, tok, block_size,
+                              kOneRow ? lane : 8 * c + v, packed, sf);
+            }
+          }
+        } else if (sm >= 0 && row_ok) {
+          cache_t* const dst = kv_cache + blk * kv_s_block + head * kv_s_head +
+                               tok * kv_s_token +
+                               (kv * kHeadDim + 4 * v) * kv_s_dim;
+  #pragma unroll
+          for (int c = 0; c < kLaneChunks; c++) {
+            if constexpr (!kHasIndex) {
+              using Converter = vllm::_typeConvert<scalar_t>;
+  #pragma unroll
+              for (int i = 0; i < kElemsPerLane; i++) {
+                e[c][i] = Converter::convert(Converter::convert(e[c][i]));
+              }
+            }
+            storeCacheElems<scalar_t, cache_t, kv_dt>(dst + 32 * c * kv_s_dim,
+                                                      e[c]);
+          }
+        }
+      }
+    }
+    in_place += src_step;
+    out += out_step;
+    __syncwarp();
+  }
+}
+
+template <int kQHeads, typename scalar_t, typename cache_t,
+          Fp8KVCacheDataType kv_dt, bool kNvfp4, typename out_idx_t,
+          bool kHasIndex, bool kInsertKV, bool kProcessIndex, bool kFp8Idx,
+          bool kNvfp4Idx>
+__global__ void __launch_bounds__(kMaxWarps * 32) fusedMiniMaxM3RowGroupKernel(
+    scalar_t* __restrict__ qkv, scalar_t* __restrict__ q_out,
+    uint8_t* __restrict__ q_fp8_out, out_idx_t* __restrict__ index_q_out,
+    scalar_t const* __restrict__ q_norm_w,
+    scalar_t const* __restrict__ k_norm_w,
+    scalar_t const* __restrict__ iq_norm_w,
+    scalar_t const* __restrict__ ik_norm_w,
+    scalar_t const* __restrict__ cos_sin_cache,
+    int64_t const* __restrict__ positions,
+    int64_t const* __restrict__ slot_mapping,
+    int64_t const* __restrict__ index_slot_mapping,
+    cache_t* __restrict__ kv_cache, out_idx_t* __restrict__ index_cache,
+    float const eps, float const q_fp8_inv_scale, int const rotary_dim,
+    int const num_tokens, int const nq, int const nkv, int const niq,
+    int const block_size, int64_t const kv_s_block, int64_t const kv_s_head,
+    int64_t const kv_s_token, int64_t const kv_s_dim, int64_t const idx_s_block,
+    float const* __restrict__ kv_k_scale, float const* __restrict__ kv_v_scale,
+    bool const write_qkv) {
+  #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
+  extern __shared__ __align__(16) uint8_t smem[];
+  Roles const roles(kQHeads, nq, nkv, niq, kInsertKV, kProcessIndex);
+  int const qkv_row = (nq + 2 * nkv + (kHasIndex ? niq + 1 : 0)) * kHeadDim;
+  // CTA (x, y) runs warps [y * cta_warps, (y + 1) * cta_warps) of the warp
+  // groups of tokens x, x + gridDim.x, ...
+  int const cta_warps = blockDim.x / 32;
+  int w = blockIdx.y * cta_warps + threadIdx.x / 32;
+  if (w >= roles.warps) return;
+  constexpr int kQStage = stageBytes<kQHeads>();
+  constexpr int kSideStage = stageBytes<kRows>();
+  int const q_ring = kStages * kQStage;
+  int const side_ring = kStages * kSideStage;
+  uint8_t* const ring_w =
+      smem + roles.ringOffset(w, q_ring, side_ring) -
+      roles.ringOffset(blockIdx.y * cta_warps, q_ring, side_ring);
+    #define RG_ARGS(NORM_W, SLOTS, CACHE_SCALE)                               \
+      qkv, q_out, q_fp8_out, index_q_out, NORM_W, cos_sin_cache, positions,   \
+          SLOTS, kv_cache, index_cache, eps, q_fp8_inv_scale, num_tokens, nq, \
+          niq, qkv_row, block_size, kv_s_block, kv_s_head, kv_s_token,        \
+          kv_s_dim, idx_s_block, CACHE_SCALE, write_qkv
+    #define RG_RUN(ROLE, STEPS, OUT)                                         \
+      runWarp<ROLE, STEPS, OUT, scalar_t, cache_t, kv_dt, kNvfp4, out_idx_t, \
+              kHasIndex, kInsertKV, kFp8Idx, kNvfp4Idx>
+    #define RG_RUN1(ROLE, OUT)                                           \
+      runWarp<ROLE, 1, OUT, scalar_t, cache_t, kv_dt, kNvfp4, out_idx_t, \
+              kHasIndex, kInsertKV, kFp8Idx, kNvfp4Idx, true>
+  if (w < roles.q_warps) {
+    int const h0 = w * kQHeads;
+    constexpr int kSteps = kQHeads / kRows;
+    uint8_t* const ring = ring_w;
+    int const rows = min(kQHeads, nq - h0);
+    if (q_out == nullptr && q_fp8_out != nullptr) {
+      if (q_fp8_inv_scale == 1.0f) {
+        RG_RUN(kQ, kSteps, kQFp8Unit)
+        (h0, rows, h0, ring, kQStage, RG_ARGS(q_norm_w, nullptr, 1.0f));
+      } else {
+        RG_RUN(kQ, kSteps, kQFp8)
+        (h0, rows, h0, ring, kQStage, RG_ARGS(q_norm_w, nullptr, 1.0f));
+      }
+    } else {
+      RG_RUN(kQ, kSteps, kQAny)
+      (h0, rows, h0, ring, kQStage, RG_ARGS(q_norm_w, nullptr, 1.0f));
+    }
+  } else {
+    uint8_t* const ring = ring_w;
+    w -= roles.q_warps;
+    if (w < roles.kv_warps) {
+      int const h0 = roles.kv_one ? w : w * kRows;
+      float const scale = kNvfp4 ? *kv_k_scale : 1.0f;
+      int const rows = roles.kv_one ? 1 : min(kRows, nkv - h0);
+      if (rows == 1) {
+        if (write_qkv) {
+          RG_RUN1(kK, kInPlace)
+          (h0, 1, nq + h0, ring, kSideStage,
+           RG_ARGS(k_norm_w, slot_mapping, scale));
+        } else {
+          RG_RUN1(kK, kNone)
+          (h0, 1, nq + h0, ring, kSideStage,
+           RG_ARGS(k_norm_w, slot_mapping, scale));
+        }
+      } else if (write_qkv) {
+        RG_RUN(kK, 1, kInPlace)
+        (h0, rows, nq + h0, ring, kSideStage,
+         RG_ARGS(k_norm_w, slot_mapping, scale));
+      } else {
+        RG_RUN(kK, 1, kNone)
+        (h0, rows, nq + h0, ring, kSideStage,
+         RG_ARGS(k_norm_w, slot_mapping, scale));
+      }
+    } else if ((w -= roles.kv_warps) < roles.v_warps) {
+      if constexpr (kInsertKV) {
+        int const h0 = roles.kv_one ? w : w * kRows;
+        int const rows = roles.kv_one ? 1 : min(kRows, nkv - h0);
+        float const scale = kNvfp4 ? *kv_v_scale : 1.0f;
+        if (rows == 1) {
+          RG_RUN1(kV, kNone)
+          (h0, 1, nq + nkv + h0, ring, kSideStage,
+           RG_ARGS(nullptr, slot_mapping, scale));
+        } else {
+          RG_RUN(kV, 1, kNone)
+          (h0, rows, nq + nkv + h0, ring, kSideStage,
+           RG_ARGS(nullptr, slot_mapping, scale));
+        }
+      }
+    } else if ((w -= roles.v_warps) < roles.iq_warps) {
+      if constexpr (kProcessIndex) {
+        int const h0 = roles.iq_one ? w : w * kRows;
+        int const rows = roles.iq_one ? 1 : min(kRows, niq - h0);
+        int const qh0 = nq + 2 * nkv + h0;
+        if (rows == 1) {
+          if (index_q_out != nullptr) {
+            RG_RUN1(kIQ, kIQOut)
+            (h0, 1, qh0, ring, kSideStage, RG_ARGS(iq_norm_w, nullptr, 1.0f));
+          } else if (write_qkv) {
+            RG_RUN1(kIQ, kIQInPlace)
+            (h0, 1, qh0, ring, kSideStage, RG_ARGS(iq_norm_w, nullptr, 1.0f));
+          } else {
+            RG_RUN1(kIQ, kNone)
+            (h0, 1, qh0, ring, kSideStage, RG_ARGS(iq_norm_w, nullptr, 1.0f));
+          }
+        } else if (index_q_out != nullptr) {
+          RG_RUN(kIQ, 1, kIQOut)
+          (h0, rows, qh0, ring, kSideStage, RG_ARGS(iq_norm_w, nullptr, 1.0f));
+        } else if (write_qkv) {
+          RG_RUN(kIQ, 1, kIQInPlace)
+          (h0, rows, qh0, ring, kSideStage, RG_ARGS(iq_norm_w, nullptr, 1.0f));
+        } else {
+          RG_RUN(kIQ, 1, kNone)
+          (h0, rows, qh0, ring, kSideStage, RG_ARGS(iq_norm_w, nullptr, 1.0f));
+        }
+      }
+    } else if constexpr (kProcessIndex) {
+      if (write_qkv) {
+        RG_RUN1(kIK, kInPlace)
+        (0, 1, nq + 2 * nkv + niq, ring, kSideStage,
+         RG_ARGS(ik_norm_w, index_slot_mapping, 1.0f));
+      } else {
+        RG_RUN1(kIK, kNone)
+        (0, 1, nq + 2 * nkv + niq, ring, kSideStage,
+         RG_ARGS(ik_norm_w, index_slot_mapping, 1.0f));
+      }
+    }
+  }
+    #undef RG_RUN
+    #undef RG_RUN1
+    #undef RG_ARGS
+    #if __CUDA_ARCH__ >= 900
+  cudaTriggerProgrammaticLaunchCompletion();
+    #endif
+  #endif
+}
+
+// Shared memory of the largest CTA of cta_warps warps.
+template <int kQHeads>
+__host__ inline int smemBytes(Roles const& r, int const cta_warps) {
+  int const q_ring = kStages * stageBytes<kQHeads>();
+  int const side_ring = kStages * stageBytes<kRows>();
+  int bytes = 0;
+  for (int w0 = 0; w0 < r.warps; w0 += cta_warps) {
+    int const w1 = std::min(w0 + cta_warps, r.warps);
+    bytes = std::max(bytes, r.ringOffset(w1, q_ring, side_ring) -
+                                r.ringOffset(w0, q_ring, side_ring));
+  }
+  return bytes;
+}
+
+// Launch with the grid this instantiation's occupancy allows: the CTAs walk
+// the tokens with stride gridDim.x.
+template <int kQHeads, typename scalar_t, typename cache_t,
+          Fp8KVCacheDataType kv_dt, bool kNvfp4, typename out_idx_t,
+          bool kHasIndex, bool kInsertKV, bool kProcessIndex, bool kFp8Idx,
+          bool kNvfp4Idx, typename... Args>
+void launch(cudaLaunchConfig_t config, int const num_tokens, Args... args) {
+  auto* const kernel =
+      fusedMiniMaxM3RowGroupKernel<kQHeads, scalar_t, cache_t, kv_dt, kNvfp4,
+                                   out_idx_t, kHasIndex, kInsertKV,
+                                   kProcessIndex, kFp8Idx, kNvfp4Idx>;
+  // Per CTA shape (warps, shared memory), computed once: the resident CTAs
+  // per SM.
+  static bool const opted_in =
+      cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize,
+                           get_device_prop()->sharedMemPerBlockOptin) ==
+      cudaSuccess;
+  static int smem_of[kMaxWarps + 1] = {};
+  static int resident_of[kMaxWarps + 1] = {};
+  int const warps = config.blockDim.x / 32;
+  int const smem = static_cast<int>(config.dynamicSmemBytes);
+  if (smem_of[warps] != smem) {
+    (void)opted_in;
+    int n = 0;
+    cudaOccupancyMaxActiveBlocksPerMultiprocessor(&n, kernel, warps * 32, smem);
+    resident_of[warps] = n > 0 ? n : 1;
+    smem_of[warps] = smem;
+  }
+  int const ctas_per_token = config.gridDim.y;
+  config.gridDim.x = std::max(
+      1, std::min(num_tokens, get_device_prop()->multiProcessorCount *
+                                  resident_of[warps] / ctas_per_token));
+  cudaLaunchKernelEx(&config, kernel, args...);
+}
+
+}  // namespace rg
+#endif  // USE_ROCM
+
+#ifndef USE_ROCM
+// q heads per q warp: four (one step, shorter warps) for a launch of at most
+// 128 q rows, or while a token has at most 32 q heads and there are few q rows
+// per SM; eight otherwise (fewer, longer warps keep more CTAs resident).
+inline int rgQHeads(int const num_tokens, int const nq) {
+  static int const num_sms = get_device_prop()->multiProcessorCount;
+  int64_t const q_rows = static_cast<int64_t>(num_tokens) * nq;
+  return q_rows <= 128 || (nq <= 32 && q_rows < 64 * num_sms) ? 4 : 8;
+}
+// Warps per CTA: a token's whole warp group once there are enough tokens to
+// give every SM CTAs; before that, eight, which spreads a token over SMs.
+inline int rgCtaWarps(int const num_tokens) {
+  static int const num_sms = get_device_prop()->multiProcessorCount;
+  return 2 * num_tokens >= 3 * num_sms ? rg::kMaxWarps : 8;
+}
+#endif
+
 // ────────────────────────────────────────────────────────────────────────────
 // Launch wrapper
 // ────────────────────────────────────────────────────────────────────────────
@@ -809,9 +1683,10 @@ void launchFusedMiniMaxM3(
     float const q_fp8_inv_scale, int const rotary_dim, int const num_tokens,
     int const nq, int const nkv, int const niq, int const block_size,
     int64_t const kv_s_block, int64_t const kv_s_head, int64_t const kv_s_token,
-    int64_t const kv_s_dim, bool const has_index, bool const insert_kv,
-    bool const process_index, bool const fp8_idx, float const* kv_k_scale,
-    float const* kv_v_scale, cudaStream_t stream) {
+    int64_t const kv_s_dim, int64_t const idx_s_block, bool const has_index,
+    bool const insert_kv, bool const process_index, bool const fp8_idx,
+    bool const nvfp4_idx, float const* kv_k_scale, float const* kv_v_scale,
+    bool const write_qkv, bool const slots_fit_int32, cudaStream_t stream) {
   // Index outputs are scalar_t (bf16) or e4m3 bytes (uint8_t); reinterpret the
   // void* pointers per instantiation in the LAUNCH macro.
   // Slot count must match the kernel's compile-time gating.
@@ -843,73 +1718,437 @@ void launchFusedMiniMaxM3(
   config.attrs = attrs;
   config.numAttrs = (sm_version >= 90) ? 1 : 0;
 
-  #define LAUNCH(HAS_INDEX, INSERT, PROCESS_INDEX, FP8, OUT_T)              \
-    cudaLaunchKernelEx(                                                     \
-        &config,                                                            \
-        fusedMiniMaxM3QNormRopeKVInsertKernel<scalar_t, cache_t, kv_dt,     \
-                                              kNvfp4, OUT_T, HAS_INDEX,     \
-                                              INSERT, PROCESS_INDEX, FP8>,  \
-        qkv, q_out, q_fp8_out, reinterpret_cast<OUT_T*>(index_q_out),       \
-        q_norm_w, k_norm_w, iq_norm_w, ik_norm_w, cos_sin_cache, positions, \
+  // The row-group kernel: same bytes, for rotary_dim == 64 on SM80+ (cp.async).
+  int const q_heads = rgQHeads(num_tokens, nq);
+  rg::Roles const roles(q_heads, nq, nkv, niq, insert_kv, process_index);
+  auto const rg_smem = [&](int const warps) {
+    return q_heads == 8 ? rg::smemBytes<8>(roles, warps)
+                        : rg::smemBytes<4>(roles, warps);
+  };
+  int cta_warps =
+      std::min({roles.warps, rgCtaWarps(num_tokens), rg::kMaxWarps});
+  // Eight warps take at most 63 KB; a whole group can exceed the 99 KB per
+  // block of SM 8.6 / 8.9 / 12.x.
+  static int const max_block_smem = get_device_prop()->sharedMemPerBlockOptin;
+  if (rg_smem(cta_warps) > max_block_smem) cta_warps = std::min(roles.warps, 8);
+  auto const aligned = [](void const* p, uintptr_t const n) {
+    return reinterpret_cast<uintptr_t>(p) % n == 0;
+  };
+  bool const row_groups =
+      sm_version >= 80 && rotary_dim == rg::kRotaryDim && slots_fit_int32 &&
+      aligned(qkv, 16) && aligned(cos_sin_cache, 16) && aligned(q_norm_w, 8) &&
+      aligned(k_norm_w, 8) && aligned(iq_norm_w, 8) && aligned(ik_norm_w, 8);
+  cudaLaunchConfig_t rg_config = config;
+  rg_config.blockDim = dim3(cta_warps * 32);
+  rg_config.gridDim = dim3(1, (roles.warps + cta_warps - 1) / cta_warps);
+  rg_config.dynamicSmemBytes = rg_smem(cta_warps);
+
+  #define LAUNCH_ARGS(OUT_T)                                                \
+    qkv, q_out, q_fp8_out, reinterpret_cast<OUT_T*>(index_q_out), q_norm_w, \
+        k_norm_w, iq_norm_w, ik_norm_w, cos_sin_cache, positions,           \
         slot_mapping, index_slot_mapping, kv_cache,                         \
         reinterpret_cast<OUT_T*>(index_cache), eps, q_fp8_inv_scale,        \
         rotary_dim, num_tokens, nq, nkv, niq, block_size, kv_s_block,       \
-        kv_s_head, kv_s_token, kv_s_dim, kv_k_scale, kv_v_scale)
+        kv_s_head, kv_s_token, kv_s_dim, idx_s_block, kv_k_scale, kv_v_scale
+  #define LAUNCH_RG(QH, HAS_INDEX, INSERT, PROCESS_INDEX, FP8, NVFP4_IDX,      \
+                    OUT_T)                                                     \
+    rg::launch<QH, scalar_t, cache_t, kv_dt, kNvfp4, OUT_T, HAS_INDEX, INSERT, \
+               PROCESS_INDEX, FP8, NVFP4_IDX>(rg_config, num_tokens,           \
+                                              LAUNCH_ARGS(OUT_T), write_qkv)
+  #define LAUNCH(HAS_INDEX, INSERT, PROCESS_INDEX, FP8, NVFP4_IDX, OUT_T)      \
+    do {                                                                       \
+      if (!row_groups) {                                                       \
+        cudaLaunchKernelEx(                                                    \
+            &config,                                                           \
+            fusedMiniMaxM3QNormRopeKVInsertKernel<                             \
+                scalar_t, cache_t, kv_dt, kNvfp4, OUT_T, HAS_INDEX, INSERT,    \
+                PROCESS_INDEX, FP8, NVFP4_IDX>,                                \
+            LAUNCH_ARGS(OUT_T));                                               \
+      } else if (q_heads == 8) {                                               \
+        LAUNCH_RG(8, HAS_INDEX, INSERT, PROCESS_INDEX, FP8, NVFP4_IDX, OUT_T); \
+      } else {                                                                 \
+        LAUNCH_RG(4, HAS_INDEX, INSERT, PROCESS_INDEX, FP8, NVFP4_IDX, OUT_T); \
+      }                                                                        \
+    } while (0)
 #else
   // ROCm: standard kernel launch syntax (no PDL/stream serialization).
   // clang-format off
-  #define LAUNCH(HAS_INDEX, INSERT, PROCESS_INDEX, FP8, OUT_T)              \
+  #define LAUNCH(HAS_INDEX, INSERT, PROCESS_INDEX, FP8, NVFP4_IDX, OUT_T)   \
     fusedMiniMaxM3QNormRopeKVInsertKernel<                                  \
         scalar_t, cache_t, kv_dt, kNvfp4, OUT_T, HAS_INDEX, INSERT,         \
-        PROCESS_INDEX, FP8><<<grid, kBlockSize, 0, stream>>>(               \
+        PROCESS_INDEX, FP8, NVFP4_IDX><<<grid, kBlockSize, 0, stream>>>(    \
         qkv, q_out, q_fp8_out, reinterpret_cast<OUT_T*>(index_q_out),       \
         q_norm_w, k_norm_w, iq_norm_w, ik_norm_w, cos_sin_cache, positions, \
         slot_mapping, index_slot_mapping, kv_cache,                         \
         reinterpret_cast<OUT_T*>(index_cache), eps, q_fp8_inv_scale,        \
         rotary_dim, num_tokens, nq, nkv, niq, block_size, kv_s_block,       \
-        kv_s_head, kv_s_token, kv_s_dim, kv_k_scale, kv_v_scale)
+        kv_s_head, kv_s_token, kv_s_dim, idx_s_block, kv_k_scale,           \
+        kv_v_scale)
   // clang-format on
 #endif
 
   if constexpr (kNvfp4) {
     // NVFP4 main cache: sparse serving only (the host checks insert_kv).
     if (!process_index) {
-      LAUNCH(true, true, false, false, scalar_t);
+      LAUNCH(true, true, false, false, false, scalar_t);
+    } else if (nvfp4_idx) {
+      LAUNCH(true, true, true, true, true, uint8_t);
     } else if (fp8_idx) {
-      LAUNCH(true, true, true, true, uint8_t);
+      LAUNCH(true, true, true, true, false, uint8_t);
     } else {
-      LAUNCH(true, true, true, false, scalar_t);
+      LAUNCH(true, true, true, false, false, scalar_t);
     }
   } else if (has_index) {
     if (!process_index) {
       if (insert_kv) {
-        LAUNCH(true, true, false, false, scalar_t);
+        LAUNCH(true, true, false, false, false, scalar_t);
       } else {
-        LAUNCH(true, false, false, false, scalar_t);
+        LAUNCH(true, false, false, false, false, scalar_t);
       }
     } else if (insert_kv) {
-      if (fp8_idx) {
-        LAUNCH(true, true, true, true,
+      if (nvfp4_idx) {
+        LAUNCH(true, true, true, true, true,
+               uint8_t);  // sparse serving, NVFP4 index-K, fp8 index_q
+      } else if (fp8_idx) {
+        LAUNCH(true, true, true, true, false,
                uint8_t);  // sparse serving, fp8 index outputs
       } else {
-        LAUNCH(true, true, true, false, scalar_t);  // sparse serving, bf16
+        LAUNCH(true, true, true, false, false,
+               scalar_t);  // sparse serving, bf16
       }
     } else {
       if (fp8_idx) {
-        LAUNCH(true, false, true, true,
+        LAUNCH(true, false, true, true, false,
                uint8_t);  // sparse profiling, fp8 index_q
       } else {
-        LAUNCH(true, false, true, false, scalar_t);  // sparse profiling, bf16
+        LAUNCH(true, false, true, false, false,
+               scalar_t);  // sparse profiling, bf16
       }
     }
+  } else if (insert_kv) {
+    LAUNCH(false, true, false, false, false, scalar_t);  // dense, KV insert
   } else {
-    // Dense layer: never has an index branch and never inserts here (the
-    // generic Attention layer owns the KV insert).
-    LAUNCH(false, false, false, false, scalar_t);
+    LAUNCH(false, false, false, false, false, scalar_t);  // dense
   }
 #undef LAUNCH
+#undef LAUNCH_RG
+#undef LAUNCH_ARGS
+}
+
+#ifndef USE_ROCM
+// ────────────────────────────────────────────────────────────────────────────
+// DSpark draft context K/V: every draft layer in one launch (bf16 in).
+// ────────────────────────────────────────────────────────────────────────────
+// The draft projects the target's context states to the K/V of all its layers
+// with one GEMM, ``all_kv [N, L, 2, nkv, 128]``. blockIdx.z picks the layer,
+// blockIdx.y the head, blockIdx.x a run of tokens. Eight lanes own one token;
+// lane l holds two 8-dim chunks of K and of V, A = dims [8l, 8l + 8) and
+// B = dims [64 + 8l, 64 + 8l + 8), so with rotary_dim == 64 every lane rotates
+// its A chunk (partner lane l ^ 4) and no lane idles through the RoPE. K gets
+// the layer's Gemma k-norm (rounded to bf16 like the unfused norm) and partial
+// NeoX RoPE (fp32 math rounded once, like rotary_embedding); K and V go into
+// the token's cache row like reshape_and_cache_flash (fp8: e4m3(bf16(x) *
+// (1 / scale)); a unit scale converts the bf16 pairs directly).
+// Layer caches are separate allocations of one layout, logical shape
+// [nb, nkv, bs, 2*128] with K at [..., :128] and V at [..., 128:]; they are
+// addressed as layer 0's cache plus a per-layer element offset (likewise the
+// fp8 k/v scales from layer 0's k scale), so the launch is CUDA-graph safe.
+constexpr int kCtxLanesPerRow = 8;
+constexpr int kCtxElems = 8;  // bf16 per 16-byte chunk
+constexpr int kCtxHalfVecs = kHeadDim / 2 / kCtxElems;  // 8: chunk B offset
+constexpr int kCtxBlockSize = 256;
+constexpr int kCtxRowsPerBlock = kCtxBlockSize / kCtxLanesPerRow;  // 32
+
+__device__ __forceinline__ float bf16lo(uint32_t v) {
+  return __uint_as_float(v << 16);
+}
+__device__ __forceinline__ float bf16hi(uint32_t v) {
+  return __uint_as_float(v & 0xffff0000u);
+}
+__device__ __forceinline__ uint32_t packBf16(float lo, float hi) {
+  __nv_bfloat162 const r = __floats2bfloat162_rn(lo, hi);
+  return *reinterpret_cast<uint32_t const*>(&r);
+}
+// Loads of what the PDL predecessor may still be writing: coherent (an .nc
+// load may be hoisted above griddepcontrol.wait).
+__device__ __forceinline__ uint4 ldgV4(uint4 const* p) {
+  uint4 r;
+  asm volatile("ld.global.v4.u32 {%0, %1, %2, %3}, [%4];"
+               : "=r"(r.x), "=r"(r.y), "=r"(r.z), "=r"(r.w)
+               : "l"(p));
+  return r;
+}
+// Low 32 bits of an int64 (positions and slots fit).
+__device__ __forceinline__ int ldgLo32(int64_t const* p) {
+  int r;
+  asm volatile("ld.global.b32 %0, [%1];" : "=r"(r) : "l"(p));
+  return r;
+}
+// Same approximation as flashinfer's gemma_rmsnorm.
+__device__ __forceinline__ float rsqrtApproxFtz(float x) {
+  float r;
+  asm("rsqrt.approx.ftz.f32 %0, %1;" : "=f"(r) : "f"(x));
+  return r;
+}
+// e4m3 of four floats, low byte first.
+__device__ __forceinline__ uint32_t packE4m3(float a, float b, float c,
+                                             float d) {
+  #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 890
+  uint32_t r;
+  asm("{\n"
+      "  .reg .b16 lo, hi;\n"
+      "  cvt.rn.satfinite.e4m3x2.f32 lo, %2, %1;\n"
+      "  cvt.rn.satfinite.e4m3x2.f32 hi, %4, %3;\n"
+      "  mov.b32 %0, {lo, hi};\n"
+      "}"
+      : "=r"(r)
+      : "f"(a), "f"(b), "f"(c), "f"(d));
+  return r;
+  #else
+  uint32_t const lo =
+      __nv_cvt_float2_to_fp8x2(make_float2(a, b), __NV_SATFINITE, __NV_E4M3);
+  uint32_t const hi =
+      __nv_cvt_float2_to_fp8x2(make_float2(c, d), __NV_SATFINITE, __NV_E4M3);
+  return lo | (hi << 16);
+  #endif
+}
+// Eight bf16 (4 packed pairs) * inv_scale -> e4m3, low byte first.
+__device__ __forceinline__ uint2 bf16x8ToE4m3(uint4 const p,
+                                              float const inv_scale) {
+  uint32_t const w[4] = {p.x, p.y, p.z, p.w};
+  uint32_t r[2];
+  #pragma unroll
+  for (int i = 0; i < 2; i++) {
+    r[i] = packE4m3(bf16lo(w[2 * i]) * inv_scale, bf16hi(w[2 * i]) * inv_scale,
+                    bf16lo(w[2 * i + 1]) * inv_scale,
+                    bf16hi(w[2 * i + 1]) * inv_scale);
+  }
+  return make_uint2(r[0], r[1]);
+}
+  #if ((defined(__CUDA_ARCH_FAMILY_SPECIFIC__) &&                              \
+        __CUDA_ARCH_FAMILY_SPECIFIC__ >= 1000) ||                              \
+       (defined(__CUDA_ARCH_SPECIFIC__) && __CUDA_ARCH_SPECIFIC__ >= 1000)) && \
+      defined(CUDART_VERSION) && CUDART_VERSION >= 13010
+    #define VLLM_CTX_KV_BF16X2_TO_E4M3X2
+  #endif
+// Eight bf16 -> e4m3 with a unit scale (exact); one convert per pair where
+// the arch and toolkit (PTX ISA 9.1+) convert bf16x2 -> e4m3x2.
+__device__ __forceinline__ uint2 bf16x8ToE4m3(uint4 const p) {
+  #ifdef VLLM_CTX_KV_BF16X2_TO_E4M3X2
+  uint32_t r[2];
+  uint32_t const w[4] = {p.x, p.y, p.z, p.w};
+    #pragma unroll
+  for (int i = 0; i < 2; i++) {
+    asm("{\n"
+        "  .reg .b16 lo, hi;\n"
+        "  cvt.rn.satfinite.e4m3x2.bf16x2 lo, %1;\n"
+        "  cvt.rn.satfinite.e4m3x2.bf16x2 hi, %2;\n"
+        "  mov.b32 %0, {lo, hi};\n"
+        "}"
+        : "=r"(r[i])
+        : "r"(w[2 * i]), "r"(w[2 * i + 1]));
+  }
+  return make_uint2(r[0], r[1]);
+  #else
+  return bf16x8ToE4m3(p, 1.0f);
+  #endif
 }
 
+template <bool kFp8>
+__global__ void __launch_bounds__(kCtxBlockSize)
+    fusedMiniMaxM3DSparkContextKVKernel(
+        uint4 const* __restrict__ all_kv,           // bf16 [N, L, 2, nkv, 128]
+        uint4 const* __restrict__ k_norm_w,         // bf16 [L, 128]
+        uint4 const* __restrict__ cos_sin_cache,    // bf16 [max_pos, rot]
+        int64_t const* __restrict__ positions,      // [N]
+        int64_t const* __restrict__ slot_mapping,   // [N], < 0: no write
+        uint8_t* __restrict__ kv_cache,             // layer 0's cache, bytes
+        int64_t const* __restrict__ cache_offsets,  // [L] elements from it
+        float const* __restrict__ scales,           // layer 0's k scale
+        int64_t const* __restrict__ scale_offsets,  // [2, L] k, v; floats
+        float const eps, int const rotary_dim, int const num_tokens,
+        int const num_layers, int const block_shift, int const s_token,
+        int const s_layer, int const s_side, int const s_head,
+        int const kv_s_block, int const kv_s_head, int const kv_s_token) {
+  // Strides are in 16-byte (8 x bf16) units for all_kv / cos_sin, elements
+  // for the cache.
+  constexpr int kElemBytes = kFp8 ? 1 : 2;
+  int const row0 = blockIdx.x * kCtxRowsPerBlock;
+  // Warps with no token exit; the others run their tail rows too (the
+  // shuffles need every lane) and skip the store.
+  if (row0 + threadIdx.x / 32 * (32 / kCtxLanesPerRow) >= num_tokens) return;
+  int const head = blockIdx.y;
+  int const layer = blockIdx.z;
+  int const lane = threadIdx.x % kCtxLanesPerRow;
+  int const dim0 = lane * kCtxElems;  // chunk A; chunk B is 64 dims later
+  int const t = row0 + threadIdx.x / kCtxLanesPerRow;
+
+  int const half = rotary_dim / 2;
+  bool const in_rope = dim0 < rotary_dim;
+  bool const first_half = dim0 < half;
+  int const cs_vec =
+      in_rope ? (first_half ? dim0 : dim0 - half) / kCtxElems : 0;
+  int const lane_xor = half / kCtxElems;
+  int const rot_vecs = rotary_dim / kCtxElems;
+  // x * cos - partner * sin in the first half, + in the second.
+  uint32_t const sin_sign = first_half ? 0x80008000u : 0u;
+
+  uint8_t* const cache =
+      kv_cache + (cache_offsets[layer] + head * kv_s_head + dim0) * kElemBytes;
+  float k_scale = 1.0f, v_scale = 1.0f;
+  if constexpr (kFp8) {
+    int64_t const k_off = scale_offsets[layer];
+    int64_t const v_off = scale_offsets[num_layers + layer];
+    k_scale = scales[k_off];
+    v_scale = scales[v_off];
+  }
+  float g[2][kCtxElems];  // 1 + w, chunks A and B
+  #pragma unroll
+  for (int c = 0; c < 2; c++) {
+    uint4 const w =
+        k_norm_w[layer * (kHeadDim / kCtxElems) + c * kCtxHalfVecs + lane];
+    uint32_t const wp[4] = {w.x, w.y, w.z, w.w};
+  #pragma unroll
+    for (int i = 0; i < 4; i++) {
+      g[c][2 * i] = 1.0f + bf16lo(wp[i]);
+      g[c][2 * i + 1] = 1.0f + bf16hi(wp[i]);
+    }
+  }
+
+  #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+  cudaGridDependencySynchronize();
+  #endif
+
+  bool const valid = t < num_tokens;
+  int const ts = valid ? t : 0;
+  int const pos = ldgLo32(positions + ts);
+  int const slot = valid ? ldgLo32(slot_mapping + ts) : -1;
+  uint4 const* const src =
+      all_kv + ts * s_token + layer * s_layer + head * s_head + lane;
+  uint4 kr[2], vr[2];
+  #pragma unroll
+  for (int c = 0; c < 2; c++) {
+    kr[c] = ldgV4(src + c * kCtxHalfVecs);
+    vr[c] = ldgV4(src + s_side + c * kCtxHalfVecs);
+  }
+  // Lanes past the rotary dims load (and later discard) a valid row too.
+  uint4 const* const cs = cos_sin_cache + pos * rot_vecs + cs_vec;
+  uint4 const cr = cs[0];
+  uint4 const sr = cs[lane_xor];
+
+  bool const store = slot >= 0;
+  // Power-of-two block size: shift and mask, not division.
+  uint8_t* const dst =
+      cache + (static_cast<int64_t>(slot >> block_shift) * kv_s_block +
+               (slot & ((1 << block_shift) - 1)) * kv_s_token) *
+                  kElemBytes;
+  constexpr int kChunkB = kHeadDim / 2 * kElemBytes;
+  constexpr int kV = kHeadDim * kElemBytes;
+  float const k_inv = __fdividef(1.0f, k_scale);
+  float const v_inv = __fdividef(1.0f, v_scale);
+  #ifdef VLLM_CTX_KV_BF16X2_TO_E4M3X2
+  bool const unit_scale = k_scale == 1.0f && v_scale == 1.0f;
+  #else
+  constexpr bool unit_scale = false;
+  #endif
+  // Eight bf16 -> the cache row at byte offset `off`.
+  auto put = [&](int const off, uint4 const x, float const inv_scale) {
+    if constexpr (kFp8) {
+      uint2 const y = unit_scale ? bf16x8ToE4m3(x) : bf16x8ToE4m3(x, inv_scale);
+      if (store) *reinterpret_cast<uint2*>(dst + off) = y;
+    } else {
+      if (store) *reinterpret_cast<uint4*>(dst + off) = x;
+    }
+  };
+
+  // V first, so its registers free up before the K math.
+  #pragma unroll
+  for (int c = 0; c < 2; c++) put(kV + c * kChunkB, vr[c], v_inv);
+
+  // K: Gemma RMSNorm over the 8 lanes of the row.
+  float x[2][kCtxElems];
+  float sumsq = 0.0f;
+  #pragma unroll
+  for (int c = 0; c < 2; c++) {
+    uint32_t const kp[4] = {kr[c].x, kr[c].y, kr[c].z, kr[c].w};
+  #pragma unroll
+    for (int j = 0; j < 4; j++) {
+      x[c][2 * j] = bf16lo(kp[j]);
+      x[c][2 * j + 1] = bf16hi(kp[j]);
+      sumsq = fmaf(x[c][2 * j], x[c][2 * j], sumsq);
+      sumsq = fmaf(x[c][2 * j + 1], x[c][2 * j + 1], sumsq);
+    }
+  }
+  #pragma unroll
+  for (int mask = kCtxLanesPerRow / 2; mask > 0; mask >>= 1) {
+    sumsq += __shfl_xor_sync(FINAL_MASK, sumsq, mask, 32);
+  }
+  float const rr = rsqrtApproxFtz(sumsq * (1.0f / kHeadDim) + eps);
+  // bf16-rounded norm output, packed.
+  uint32_t kn[2][4];
+  #pragma unroll
+  for (int c = 0; c < 2; c++) {
+  #pragma unroll
+    for (int j = 0; j < 4; j++) {
+      kn[c][j] = packBf16(x[c][2 * j] * rr * g[c][2 * j],
+                          x[c][2 * j + 1] * rr * g[c][2 * j + 1]);
+    }
+  }
+  // RoPE on chunk A; the partner chunk is half / 8 lanes away.
+  uint32_t const cp[4] = {cr.x, cr.y, cr.z, cr.w};
+  uint32_t const sp[4] = {sr.x ^ sin_sign, sr.y ^ sin_sign, sr.z ^ sin_sign,
+                          sr.w ^ sin_sign};
+  #pragma unroll
+  for (int j = 0; j < 4; j++) {
+    uint32_t const pn = __shfl_xor_sync(FINAL_MASK, kn[0][j], lane_xor, 32);
+    uint32_t const rot =
+        packBf16(__fmaf_rn(bf16lo(kn[0][j]), bf16lo(cp[j]),
+                           __fmul_rn(bf16lo(pn), bf16lo(sp[j]))),
+                 __fmaf_rn(bf16hi(kn[0][j]), bf16hi(cp[j]),
+                           __fmul_rn(bf16hi(pn), bf16hi(sp[j]))));
+    kn[0][j] = in_rope ? rot : kn[0][j];
+  }
+  #pragma unroll
+  for (int c = 0; c < 2; c++) {
+    put(c * kChunkB, make_uint4(kn[c][0], kn[c][1], kn[c][2], kn[c][3]), k_inv);
+  }
+
+  #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+  cudaTriggerProgrammaticLaunchCompletion();
+  #endif
+}
+
+template <bool kFp8>
+void launchFusedMiniMaxM3DSparkContextKV(
+    uint4 const* all_kv, uint4 const* k_norm_w, uint4 const* cos_sin_cache,
+    int64_t const* positions, int64_t const* slot_mapping, uint8_t* kv_cache,
+    int64_t const* cache_offsets, float const* scales,
+    int64_t const* scale_offsets, float const eps, int const rotary_dim,
+    int const num_tokens, int const num_layers, int const nkv,
+    int const block_shift, int const s_token, int const s_layer,
+    int const s_side, int const s_head, int const kv_s_block,
+    int const kv_s_head, int const kv_s_token, cudaStream_t stream) {
+  if (num_tokens == 0) return;
+  static int const sm_version = getSMVersion();
+  cudaLaunchConfig_t config;
+  config.gridDim = dim3((num_tokens + kCtxRowsPerBlock - 1) / kCtxRowsPerBlock,
+                        nkv, num_layers);
+  config.blockDim = dim3(kCtxBlockSize);
+  config.dynamicSmemBytes = 0;
+  config.stream = stream;
+  cudaLaunchAttribute attrs[1];
+  attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
+  attrs[0].val.programmaticStreamSerializationAllowed = 1;
+  config.attrs = attrs;
+  config.numAttrs = (sm_version >= 90) ? 1 : 0;
+  cudaLaunchKernelEx(&config, fusedMiniMaxM3DSparkContextKVKernel<kFp8>, all_kv,
+                     k_norm_w, cos_sin_cache, positions, slot_mapping, kv_cache,
+                     cache_offsets, scales, scale_offsets, eps, rotary_dim,
+                     num_tokens, num_layers, block_shift, s_token, s_layer,
+                     s_side, s_head, kv_s_block, kv_s_head, kv_s_token);
+}
+#endif  // !USE_ROCM
+
 }  // namespace minimax_m3_fused_ops
 }  // namespace vllm
 
@@ -948,8 +2187,9 @@ void launchFusedMiniMaxM3(
       static_cast<float>(eps), 1.0f / static_cast<float>(q_fp8_scale),          \
       static_cast<int>(rotary_dim), num_tokens, nq, nkv, niq,                  \
       static_cast<int>(block_size), kv_s_block, kv_s_head, kv_s_token,         \
-      kv_s_dim, has_index, insert_kv, process_index, fp8_idx, kv_k_scale_ptr,  \
-      kv_v_scale_ptr, stream)
+      kv_s_dim, idx_s_block, has_index, insert_kv, process_index, fp8_idx,     \
+      nvfp4_idx, kv_k_scale_ptr, kv_v_scale_ptr, write_qkv, slots_fit_int32,  \
+      stream)
 #define CALL_FUSED_MINIMAX_M3(_RAW_T, CACHE_T, KV_DTYPE)                       \
   CALL_FUSED_MINIMAX_M3_IMPL(CACHE_T, KV_DTYPE, false)
 // clang-format on
@@ -980,7 +2220,7 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
         q_fp8_out,  // [N, nq*128] contiguous E4M3
     double q_fp8_scale,
     std::optional<torch::stable::Tensor> kv_k_scale,  // [] f32, NVFP4 only
-    std::optional<torch::stable::Tensor> kv_v_scale) {
+    std::optional<torch::stable::Tensor> kv_v_scale, bool write_qkv) {
   STD_TORCH_CHECK(qkv.is_cuda() && qkv.is_contiguous(),
                   "qkv must be contiguous CUDA");
   STD_TORCH_CHECK(
@@ -1032,11 +2272,8 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
                   " + num_index_heads + 1) * 128 for sparse, "
                   "(num_heads + 2*num_kv_heads) * 128 for dense");
 
-  // Only the sparse layer inserts here (dense lets the generic Attention layer
-  // own the KV write); there is no dense+insert kernel instantiation.
-  STD_TORCH_CHECK(
-      !insert_kv || has_index,
-      "insert mode (kv_cache) requires the index branch (sparse layer)");
+  STD_TORCH_CHECK(!nvfp4_kv || has_index,
+                  "nvfp4 kv_cache requires the index branch (sparse layer)");
   STD_TORCH_CHECK(has_index || !skip_index_branch,
                   "skip_index_branch requires sparse qkv rows");
   if (process_index) {
@@ -1084,14 +2321,18 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
           kv_cache->scalar_type() == torch::headeronly::ScalarType::Byte,
           "fp8 kv_cache must use uint8 storage");
     }
-    // Indexer index-K cache: independent dtype -- qkv dtype or fp8 e4m3.
+    // Indexer index-K cache: independent dtype -- qkv dtype, fp8 e4m3, or
+    // uint8 packed NVFP4 pages.
     if (process_index) {
       STD_TORCH_CHECK(
           index_cache.has_value() &&
               (index_cache->scalar_type() == qkv.scalar_type() ||
                index_cache->scalar_type() ==
-                   torch::headeronly::ScalarType::Float8_e4m3fn),
-          "insert mode requires index_cache matching qkv dtype or fp8 e4m3");
+                   torch::headeronly::ScalarType::Float8_e4m3fn ||
+               index_cache->scalar_type() ==
+                   torch::headeronly::ScalarType::Byte),
+          "insert mode requires index_cache matching qkv dtype, fp8 e4m3, "
+          "or uint8 (nvfp4)");
     }
     STD_TORCH_CHECK(kv_cache->dim() == 4 && kv_cache->stride(3) == 1,
                     "kv_cache must be [nb,nkv,bs,2*head_dim] with contiguous "
@@ -1150,6 +2391,9 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
     STD_TORCH_CHECK(std::isfinite(q_fp8_scale) && q_fp8_scale > 0.0,
                     "q_fp8_scale must be finite and positive");
   }
+  STD_TORCH_CHECK(write_qkv || (!q_out.has_value() && q_fp8_out.has_value()),
+                  "write_qkv=False writes q only to q_fp8_out: pass q_fp8_out "
+                  "and no q_out");
   if (index_q_out.has_value()) {
     STD_TORCH_CHECK(process_index,
                     "index_q_out requires index branch processing");
@@ -1166,16 +2410,36 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
   }
 
   // fp8 index path: the index-K cache and index-Q outputs are e4m3 bytes while
-  // q/k/v + q_out stay qkv dtype. Both index outputs must agree.
+  // q/k/v + q_out stay qkv dtype. Both index outputs must agree. The NVFP4
+  // index path keeps e4m3 index-Q and stores index-K as packed NVFP4 pages.
   auto const kFp8 = torch::headeronly::ScalarType::Float8_e4m3fn;
+  bool const nvfp4_idx =
+      process_index && insert_kv &&
+      index_cache->scalar_type() == torch::headeronly::ScalarType::Byte;
   bool const fp8_idx =
       process_index &&
-      ((index_cache.has_value() && index_cache->scalar_type() == kFp8) ||
+      (nvfp4_idx ||
+       (index_cache.has_value() && index_cache->scalar_type() == kFp8) ||
        (index_q_out.has_value() && index_q_out->scalar_type() == kFp8));
-  if (fp8_idx) {
+  int64_t idx_s_block = 0;
+  if (nvfp4_idx) {
+#ifdef USE_ROCM
+    STD_TORCH_CHECK(false, "nvfp4 index_cache is not supported on ROCm");
+#endif
+    int64_t const page_width = kHeadDim / 2 + kHeadDim / 16;
     STD_TORCH_CHECK(
-        !index_cache.has_value() || index_cache->scalar_type() == kFp8,
-        "fp8 index path: index_cache must be fp8 e4m3");
+        index_cache->dim() == 3 && index_cache->size(1) == block_size &&
+            index_cache->size(2) == page_width && index_cache->stride(2) == 1 &&
+            index_cache->stride(1) == page_width &&
+            index_cache->stride(0) >= block_size * page_width,
+        "nvfp4 index_cache must be [nb, block_size, 72] uint8 with "
+        "contiguous pages");
+    idx_s_block = index_cache->stride(0);
+  }
+  if (fp8_idx) {
+    STD_TORCH_CHECK(!index_cache.has_value() || nvfp4_idx ||
+                        index_cache->scalar_type() == kFp8,
+                    "fp8 index path: index_cache must be fp8 e4m3");
     STD_TORCH_CHECK(
         !index_q_out.has_value() || index_q_out->scalar_type() == kFp8,
         "fp8 index path: index_q_out must be fp8 e4m3");
@@ -1188,6 +2452,12 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
   float const* kv_v_scale_ptr =
       nvfp4_kv ? kv_v_scale->const_data_ptr<float>() : nullptr;
 
+  // The row-group kernel indexes cache slots in 32 bits.
+  bool const slots_fit_int32 =
+      !insert_kv ||
+      (kv_cache->size(0) * block_size <= INT32_MAX &&
+       (!process_index || index_cache->size(0) * block_size <= INT32_MAX));
+
   const torch::stable::accelerator::DeviceGuard device_guard(
       qkv.get_device_index());
   auto stream = get_current_cuda_stream(qkv.get_device_index());
@@ -1216,3 +2486,122 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
 
 #undef CALL_FUSED_MINIMAX_M3
 #undef CALL_FUSED_MINIMAX_M3_IMPL
+
+// ────────────────────────────────────────────────────────────────────────────
+// DSpark context K/V torch op
+// ────────────────────────────────────────────────────────────────────────────
+void fused_minimax_m3_dspark_context_kv(
+    torch::stable::Tensor const& all_kv,         // bf16 [N, L, 2, nkv, 128]
+    torch::stable::Tensor const& k_norm_weight,  // bf16 [L, 128]
+    torch::stable::Tensor const& cos_sin_cache,  // bf16 [max_pos, rotary_dim]
+    torch::stable::Tensor const& positions,      // [N] i64
+    torch::stable::Tensor const& slot_mapping,   // [N] i64
+    torch::stable::Tensor& kv_cache,  // layer 0's [nb, nkv, bs, 2*128]
+    torch::stable::Tensor const& cache_offsets,  // [L] i64, elements
+    torch::stable::Tensor const& k_scale,        // layer 0's [] f32 k scale
+    torch::stable::Tensor const& scale_offsets,  // [2, L] i64, floats
+    int64_t rotary_dim, double eps, const std::string& kv_cache_dtype) {
+#ifdef USE_ROCM
+  STD_TORCH_CHECK(false, "fused_minimax_m3_dspark_context_kv is CUDA-only");
+#else
+  using torch::headeronly::ScalarType;
+  int const kHeadDim = vllm::minimax_m3_fused_ops::kHeadDim;
+  STD_TORCH_CHECK(all_kv.is_cuda() && all_kv.dim() == 5 &&
+                      all_kv.size(2) == 2 && all_kv.size(4) == kHeadDim &&
+                      all_kv.stride(4) == 1 &&
+                      all_kv.scalar_type() == ScalarType::BFloat16,
+                  "all_kv must be bf16 CUDA [N, L, 2, nkv, 128], last dim "
+                  "contiguous");
+  for (int d = 0; d < 4; d++) {
+    STD_TORCH_CHECK(all_kv.stride(d) % 8 == 0,
+                    "all_kv strides must be multiples of 8 elements");
+  }
+  STD_TORCH_CHECK(reinterpret_cast<uintptr_t>(all_kv.data_ptr()) % 16 == 0,
+                  "all_kv must be 16-byte aligned");
+  int const num_tokens = static_cast<int>(all_kv.size(0));
+  int const num_layers = static_cast<int>(all_kv.size(1));
+  int const nkv = static_cast<int>(all_kv.size(3));
+  STD_TORCH_CHECK(k_norm_weight.is_contiguous() &&
+                      k_norm_weight.scalar_type() == ScalarType::BFloat16 &&
+                      k_norm_weight.numel() == num_layers * kHeadDim,
+                  "k_norm_weight must be contiguous bf16 [L, 128]");
+  STD_TORCH_CHECK(cos_sin_cache.is_contiguous() &&
+                      cos_sin_cache.scalar_type() == ScalarType::BFloat16 &&
+                      cos_sin_cache.dim() == 2 &&
+                      cos_sin_cache.size(1) == rotary_dim,
+                  "cos_sin_cache must be contiguous bf16 [max_pos, "
+                  "rotary_dim]");
+  STD_TORCH_CHECK(rotary_dim == 16 || rotary_dim == 32 || rotary_dim == 64,
+                  "rotary_dim must be 16, 32 or 64");
+  for (auto const* t : {&positions, &slot_mapping}) {
+    STD_TORCH_CHECK(t->is_cuda() && t->is_contiguous() &&
+                        t->scalar_type() == ScalarType::Long &&
+                        t->numel() == num_tokens,
+                    "positions and slot_mapping must be contiguous int64 [N]");
+  }
+  STD_TORCH_CHECK(cache_offsets.is_cuda() && cache_offsets.is_contiguous() &&
+                      cache_offsets.scalar_type() == ScalarType::Long &&
+                      cache_offsets.numel() == num_layers,
+                  "cache_offsets must be int64 [L]");
+  STD_TORCH_CHECK(scale_offsets.is_cuda() && scale_offsets.is_contiguous() &&
+                      scale_offsets.scalar_type() == ScalarType::Long &&
+                      scale_offsets.numel() == 2 * num_layers,
+                  "scale_offsets must be int64 [2, L]");
+  STD_TORCH_CHECK(
+      k_scale.is_cuda() && k_scale.scalar_type() == ScalarType::Float,
+      "k_scale must be a float32 CUDA tensor");
+  bool const fp8 = kv_cache_dtype == "fp8" || kv_cache_dtype == "fp8_e4m3";
+  STD_TORCH_CHECK(fp8 || kv_cache_dtype == "auto",
+                  "kv_cache_dtype must be auto, fp8 or fp8_e4m3");
+  STD_TORCH_CHECK(fp8 ? kv_cache.scalar_type() == ScalarType::Byte
+                      : kv_cache.scalar_type() == ScalarType::BFloat16,
+                  "kv_cache must be uint8 (fp8) or bf16 (auto)");
+  STD_TORCH_CHECK(
+      kv_cache.dim() == 4 && kv_cache.size(1) == nkv &&
+          kv_cache.size(3) == 2 * kHeadDim && kv_cache.stride(3) == 1 &&
+          kv_cache.stride(0) % 16 == 0 && kv_cache.stride(1) % 16 == 0 &&
+          kv_cache.stride(2) % 16 == 0,
+      "kv_cache must be [nb, nkv, bs, 256] with a contiguous last dim");
+  int64_t const block_size = kv_cache.size(2);
+  STD_TORCH_CHECK(block_size > 0 && (block_size & (block_size - 1)) == 0,
+                  "kv_cache block size must be a power of two");
+  int block_shift = 0;
+  while ((int64_t{1} << block_shift) < block_size) block_shift++;
+  int64_t constexpr kIntMax = std::numeric_limits<int>::max();
+  STD_TORCH_CHECK(all_kv.numel() / 8 <= kIntMax &&
+                      kv_cache.stride(0) <= kIntMax &&
+                      kv_cache.size(0) * block_size <= kIntMax,
+                  "all_kv / kv_cache too large for 32-bit indexing");
+
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      all_kv.get_device_index());
+  auto stream = get_current_cuda_stream(all_kv.get_device_index());
+  // all_kv strides in 16-byte vectors.
+  auto launch = [&](auto fp8_tag) {
+    constexpr bool kFp8 = decltype(fp8_tag)::value;
+    vllm::minimax_m3_fused_ops::launchFusedMiniMaxM3DSparkContextKV<kFp8>(
+        reinterpret_cast<uint4 const*>(all_kv.data_ptr()),
+        reinterpret_cast<uint4 const*>(k_norm_weight.data_ptr()),
+        reinterpret_cast<uint4 const*>(cos_sin_cache.data_ptr()),
+        reinterpret_cast<int64_t const*>(positions.data_ptr()),
+        reinterpret_cast<int64_t const*>(slot_mapping.data_ptr()),
+        reinterpret_cast<uint8_t*>(kv_cache.data_ptr()),
+        reinterpret_cast<int64_t const*>(cache_offsets.data_ptr()),
+        reinterpret_cast<float const*>(k_scale.data_ptr()),
+        reinterpret_cast<int64_t const*>(scale_offsets.data_ptr()),
+        static_cast<float>(eps), static_cast<int>(rotary_dim), num_tokens,
+        num_layers, nkv, block_shift, static_cast<int>(all_kv.stride(0) / 8),
+        static_cast<int>(all_kv.stride(1) / 8),
+        static_cast<int>(all_kv.stride(2) / 8),
+        static_cast<int>(all_kv.stride(3) / 8),
+        static_cast<int>(kv_cache.stride(0)),
+        static_cast<int>(kv_cache.stride(1)),
+        static_cast<int>(kv_cache.stride(2)), stream);
+  };
+  if (fp8) {
+    launch(std::true_type{});
+  } else {
+    launch(std::false_type{});
+  }
+#endif
+}
diff --git a/csrc/fused_mqa_topk/candidates.cuh b/csrc/fused_mqa_topk/candidates.cuh
new file mode 100644
index 000000000000..30afdec28efd
--- /dev/null
+++ b/csrc/fused_mqa_topk/candidates.cuh
@@ -0,0 +1,81 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+//
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2025 DeepSeek
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+// SOFTWARE.
+//
+// Candidate records of the gated scan: a uint16 value (the low 16 bits of the
+// fp24 score code) and an int32 word (key index in bits 0-19, the code's high
+// 8 bits in bits 20-27).
+#pragma once
+#include <cstdint>
+#include <cutlass/cutlass.h>
+
+#include "common.cuh"
+#include "constants.h"
+
+namespace fused_mqa_topk {
+using CandidateValue = uint16_t;
+constexpr uint32_t kCandidateIndexBits = 20;
+constexpr uint32_t kCandidateIndexMask = (1u << kCandidateIndexBits) - 1u;
+static_assert(kMaxKeys <= (1u << kCandidateIndexBits),
+              "key indices must fit the candidate index field");
+
+CUTLASS_DEVICE uint32_t candidate_pack_index(const uint32_t payload,
+                                             const uint32_t kv_index) {
+  return (kv_index & kCandidateIndexMask) |
+         ((payload >> 16) << kCandidateIndexBits);
+}
+
+CUTLASS_DEVICE uint32_t candidate_fp24_code(const float value) {
+  return ordered_code(value) >> 8;
+}
+
+CUTLASS_DEVICE uint32_t candidate_load_score_code(const CandidateValue value,
+                                                  const int32_t packed_idx) {
+  return (static_cast<uint32_t>(packed_idx) >> kCandidateIndexBits) << 16 |
+         static_cast<uint32_t>(value);
+}
+
+CUTLASS_DEVICE int32_t candidate_decode_index(const int32_t packed_idx) {
+  return static_cast<int32_t>(static_cast<uint32_t>(packed_idx) &
+                              kCandidateIndexMask);
+}
+
+// Stores a candidate as its fp24 code (low 16 bits in value_dst, high 8 bits
+// in bits 20-27 of the index word) and key index.
+CUTLASS_DEVICE void store_candidate_record(CandidateValue* value_dst,
+                                           int32_t* index_dst,
+                                           const uint32_t payload,
+                                           const uint32_t kv_index) {
+  __stcg(value_dst, static_cast<CandidateValue>(payload));
+  __stcg(index_dst,
+         static_cast<int32_t>(candidate_pack_index(payload, kv_index)));
+}
+
+CUTLASS_DEVICE void store_candidate(CandidateValue* value_dst,
+                                    int32_t* index_dst, const float bq,
+                                    const uint32_t kv_index) {
+  store_candidate_record(value_dst, index_dst, candidate_fp24_code(bq),
+                         kv_index);
+}
+
+}  // namespace fused_mqa_topk
diff --git a/csrc/fused_mqa_topk/common.cuh b/csrc/fused_mqa_topk/common.cuh
new file mode 100644
index 000000000000..860bed7c3dcc
--- /dev/null
+++ b/csrc/fused_mqa_topk/common.cuh
@@ -0,0 +1,130 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+//
+// Small device helpers shared by the persistent kernel, its gate sample and
+// the exact rescan.
+#pragma once
+
+#include <cuda_fp16.h>
+#include <cutlass/arch/barrier.h>
+
+#include <cstdint>
+
+#include "constants.h"
+
+namespace fused_mqa_topk {
+
+// Warp roles of both pipelines (persistent kernel and dense rescan): warps 0-7
+// are the math warps (two warpgroups of 4 query rows each), warps 8-11 the
+// producer / UMMA / drain warps.
+constexpr uint32_t kMathThreads = 256, kMathWgThreads = 128,
+                   kMathWarps = kMathThreads / 32;
+// Query rows of a block each math warpgroup scores.
+constexpr uint32_t kWgRows = kBlockRows / 2;
+// Register split: ptxas gives each of the 384 threads 168 registers
+// (65536 / 384, rounded down to a multiple of 8); setmaxnreg moves them from
+// the producer / UMMA / drain warpgroup to the two math warpgroups.
+constexpr uint32_t kMathRegs = 224, kSpecRegs = 168 * 3 - 2 * kMathRegs;
+
+// Named barrier ids (cutlass::arch::NamedBarrier offsets them past the ids it
+// reserves). kBarMath / kBarSelect: the 256 math threads; kBarDrain: the math
+// threads and the two drain warps; kBarDenseMath: the dense rescan's math
+// threads.
+constexpr uint32_t kBarDenseMath = 0, kBarMath = 2, kBarDrain = 4,
+                   kBarSelect = 5;
+constexpr uint32_t kDrainSyncThreads = kMathThreads + 2 * 32;
+
+// Barrier of the 256 math threads (warps 0-7); the compiler barriers keep
+// shared-memory accesses on their side of it.
+__device__ __forceinline__ void math_wg_sync() {
+  asm volatile("" ::: "memory");
+  cutlass::arch::NamedBarrier(kMathThreads, kBarMath).sync();
+  asm volatile("" ::: "memory");
+}
+
+// Shared-memory histograms pad one word per 32 bins (bank conflicts):
+// bin b lives at padded_bin(b), a histogram of n bins takes padded_bins(n).
+__host__ __device__ constexpr uint32_t padded_bins(uint32_t n) {
+  return n + n / 32;
+}
+__device__ __forceinline__ uint32_t padded_bin(uint32_t b) {
+  return b + b / 32;
+}
+
+// Mask of all 32 lanes for the warp-synchronous intrinsics.
+constexpr uint32_t kFullMask = 0xffffffffu;
+
+// Inclusive prefix sum of v over the warp's lanes.
+__device__ __forceinline__ uint32_t warp_inclusive_scan(uint32_t v,
+                                                        uint32_t lane) {
+#pragma unroll
+  for (uint32_t d = 1; d < 32; d *= 2) {
+    const uint32_t u = __shfl_up_sync(kFullMask, v, d);
+    if (lane >= d) v += u;
+  }
+  return v;
+}
+
+__device__ __forceinline__ void mbar_inval(const void* b) {
+  asm volatile("mbarrier.inval.shared::cta.b64 [%0];" ::"r"(
+                   static_cast<uint32_t>(__cvta_generic_to_shared(b)))
+               : "memory");
+}
+
+// Reads `x` with a volatile load, so the compiler re-reads it from (shared)
+// memory instead of keeping it live in a register.
+template <class T>
+__device__ __forceinline__ T volatile_load(const T& x) {
+  return *reinterpret_cast<const volatile T*>(&x);
+}
+
+// Order-preserving uint32 code of a float (higher code = larger value) and
+// its inverse.
+__device__ __forceinline__ uint32_t ordered_code(float x) {
+  const uint32_t u = __float_as_uint(x);
+  return (u & 0x80000000u) ? ~u : (u ^ 0x80000000u);
+}
+__device__ __forceinline__ float ordered_float(uint32_t code) {
+  return __uint_as_float((code & 0x80000000u) ? (code ^ 0x80000000u) : ~code);
+}
+
+// Histogram bin of a score: its fp16 value (rounded down) as an
+// order-preserving 16-bit code, >> Shift. half_bin_floor is the smallest value
+// in bin `b`.
+template <uint32_t Shift>
+__device__ __forceinline__ uint32_t half_bin(float x) {
+  const uint32_t u = __half_as_ushort(__float2half_rd(x));
+  return ((u & 0x8000u) ? ((~u) & 65535u) : (u ^ 0x8000u)) >> Shift;
+}
+template <uint32_t Shift>
+__device__ __forceinline__ float half_bin_floor(uint32_t b) {
+  const uint32_t o = b << Shift,
+                 u = (o & 0x8000u) ? (o ^ 0x8000u) : ((~o) & 65535u);
+  return __half2float(__ushort_as_half(u));
+}
+
+// Weighted head reduction of one row's 32 head scores s_h in two FFMA2 chains
+// seeded with sum_0 / sum_1: sum_h w_h * |s_h| (kWithLinear = false: the
+// persistent kernel adds the linear part on the tensor core) or
+// sum_h w_h * (s_h + |s_h|) (kWithLinear = true); relu(s) = (s + |s|) / 2.
+// Adapted from DeepGEMM's sm100_fp8_mqa_logits.cuh at 88965b078186.
+template <bool kWithLinear>
+__device__ __forceinline__ float reduce_heads(const float* s, const float* w,
+                                              float2 sum_0 = {},
+                                              float2 sum_1 = {}) {
+  auto term = [&](uint32_t at, float2 sum) {
+    const auto abs_s = make_float2(fabsf(s[at]), fabsf(s[at + 1]));
+    const auto x =
+        kWithLinear ? __fadd2_rn(make_float2(s[at], s[at + 1]), abs_s) : abs_s;
+    return __ffma2_rn(x, make_float2(w[at], w[at + 1]), sum);
+  };
+#pragma unroll
+  for (uint32_t j = 0; j < kNumHeads; j += 4) {
+    sum_0 = term(j, sum_0);
+    sum_1 = term(j + 2, sum_1);
+  }
+  const auto sum = __fadd2_rn(sum_0, sum_1);
+  return sum.x + sum.y;
+}
+
+}  // namespace fused_mqa_topk
diff --git a/csrc/fused_mqa_topk/constants.h b/csrc/fused_mqa_topk/constants.h
new file mode 100644
index 000000000000..c831bbfcfa6e
--- /dev/null
+++ b/csrc/fused_mqa_topk/constants.h
@@ -0,0 +1,33 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+//
+// Shapes and workspace layout shared by the fused MQA-logits + top-k kernel
+// and its host wrapper. The Python side mirrors the workspace sizes in
+// vllm/model_executor/kernels/attention/dsa/fused_mqa_topk.py (the host wrapper
+// checks them).
+#pragma once
+
+#include <cstdint>
+
+namespace fused_mqa_topk {
+
+// Indexer shape: 32 heads of dimension 128.
+constexpr uint32_t kNumHeads = 32, kHeadDim = 128;
+// Query rows per block (the unit a CTA claims), threads per CTA, keys per
+// scan tile.
+constexpr uint32_t kBlockRows = 8, kNumThreads = 384, kTileKeys = 256;
+// Largest key range per call.
+constexpr uint32_t kMaxKeys = 1u << 20;
+// Per-row workspace: the gated candidate slab (a row with more candidates is
+// rescanned exactly) and the exact rescan's retained set.
+constexpr uint32_t kCandidatesPerRow = 16384, kRetainedPerRow = 4096;
+// Candidate / retained slots (of kBlockRows rows each) per persistent CTA: one
+// block in selection, one claimed ahead.
+constexpr uint32_t kSlotsPerCta = 2;
+
+// Top-k sizes the kernel is compiled for.
+constexpr bool is_supported_topk(int64_t k) {
+  return k == 512 || k == 1024 || k == 2048;
+}
+
+}  // namespace fused_mqa_topk
diff --git a/csrc/fused_mqa_topk/deep_gemm_compat.cuh b/csrc/fused_mqa_topk/deep_gemm_compat.cuh
new file mode 100644
index 000000000000..ae13a12c3d57
--- /dev/null
+++ b/csrc/fused_mqa_topk/deep_gemm_compat.cuh
@@ -0,0 +1,180 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+//
+// Adapted from DeepGEMM (https://github.com/deepseek-ai/DeepGEMM) at
+// 88965b078186ee7510ab9fc4f1d5ebc19adfa8d1, the last release before its SM100
+// MQA-logits kernels went MX-only. FP8 MQA-logits building blocks: the
+// shared-storage layout with fp32 per-key scales and fp32 weights, the
+// grid-stride contiguous-KV scheduler and the 2D/3D TMA load. Trimmed to what
+// the fused kernel instantiates (FP8 only, no metadata schedule, no multicast,
+// no paging).
+#pragma once
+
+#include <cute/arch/copy_sm90_tma.hpp>
+#include <cutlass/arch/barrier.h>
+
+#include <deep_gemm/common/math.cuh>
+#include <deep_gemm/common/packing.cuh>
+#include <deep_gemm/common/types.cuh>
+#include <deep_gemm/common/utils.cuh>
+
+namespace fused_mqa_topk::dg {
+namespace math = deep_gemm::math;
+
+// FP8 Q/K (one byte per element), fp32 per-key scales, fp32 weights: the
+// DeepGEMM layout without its MX / FP4 variants.
+template <uint32_t kNumHeads, uint32_t kHeadDim, uint32_t BLOCK_Q,
+          uint32_t SPLIT_KV, uint32_t kNumQStages, uint32_t kNumKVStages,
+          uint32_t kNumTmemStages, typename qk_dtype_t,
+          typename reduce_dtype_t = float>
+struct MQALogitsSharedStorage {
+  using Barrier = cutlass::arch::ClusterTransactionBarrier;
+  DG_STATIC_ASSERT(sizeof(qk_dtype_t) == 1, "FP8 Q/K only");
+
+  static constexpr uint32_t BLOCK_QH = BLOCK_Q * kNumHeads;
+  static constexpr uint32_t UMMA_N = math::constexpr_align(BLOCK_QH, 8u);
+  // Align to one 8-row Q/K swizzle tile (head_dim bytes per FP8 token).
+  static constexpr uint32_t kSwizzleAlignment = 8 * kHeadDim;
+  static constexpr uint32_t kNumQElementsPerStage = UMMA_N * kHeadDim;
+  static constexpr uint32_t kNumKVElementsPerStage = SPLIT_KV * kHeadDim;
+  // TMA destinations in shared memory must be 128-byte aligned.
+  static constexpr uint32_t kTmaAlignment = 128;
+  static constexpr uint32_t kNumWeightBytesPerRow = math::constexpr_align(
+      kNumHeads * static_cast<uint32_t>(sizeof(reduce_dtype_t)), 16u);
+  static constexpr uint32_t kNumWeightBytesPerStage =
+      math::constexpr_align(BLOCK_Q * kNumWeightBytesPerRow, kTmaAlignment);
+  static constexpr uint32_t kNumWeightElementsPerStage =
+      kNumWeightBytesPerStage / static_cast<uint32_t>(sizeof(reduce_dtype_t));
+
+  DG_STATIC_ASSERT(kNumQElementsPerStage % kSwizzleAlignment == 0,
+                   "Unaligned TMA swizzling");
+  DG_STATIC_ASSERT(kNumKVElementsPerStage % kSwizzleAlignment == 0,
+                   "Unaligned TMA swizzling");
+  DG_STATIC_ASSERT(kSwizzleAlignment % 128 == 0,
+                   "TMA destination must be 128-byte aligned");
+
+  alignas(kSwizzleAlignment)
+      qk_dtype_t smem_q[kNumQStages][kNumQElementsPerStage];
+  alignas(kSwizzleAlignment)
+      qk_dtype_t smem_kv[kNumKVStages][kNumKVElementsPerStage];
+  alignas(kTmaAlignment) float smem_sf_kv[kNumKVStages][SPLIT_KV];
+  alignas(kTmaAlignment)
+      reduce_dtype_t smem_weights[kNumQStages][kNumWeightElementsPerStage];
+  // Barriers require 8-byte alignment, already guaranteed by the preceding
+  // TMA-aligned arrays.
+  Barrier full_q_barriers[kNumQStages];
+  Barrier empty_q_barriers[kNumQStages];
+  Barrier full_kv_barriers[kNumKVStages];
+  Barrier empty_kv_barriers[kNumKVStages];
+  Barrier full_tmem_barriers[kNumTmemStages];
+  Barrier empty_tmem_barriers[kNumTmemStages];
+  uint32_t tmem_ptr_in_smem;
+};
+
+// Grid-stride contiguous-KV scheduler: SM `s` owns Q blocks `s`, `s + kNumSMs`,
+// ..., each visiting the KV range [min k_start, max k_end) of its rows.
+template <uint32_t BLOCK_Q, uint32_t SPLIT_KV, uint32_t kNumSMs>
+struct SM100MQALogitsScheduler {
+  uint32_t current_q_block_idx;
+  uint32_t num_q_blocks;
+  uint32_t num_q_tokens;
+  uint32_t num_kv_tokens;
+  const uint32_t* cu_seq_len_k_start;
+  const uint32_t* cu_seq_len_k_end;
+
+  CUTLASS_DEVICE SM100MQALogitsScheduler(const uint32_t& sm_idx,
+                                         const uint32_t& num_q_tokens,
+                                         const uint32_t& num_kv_tokens,
+                                         const uint32_t* cu_seq_len_k_start,
+                                         const uint32_t* cu_seq_len_k_end)
+      : current_q_block_idx(sm_idx),
+        num_q_blocks(math::ceil_div(num_q_tokens, BLOCK_Q)),
+        num_q_tokens(num_q_tokens),
+        num_kv_tokens(num_kv_tokens),
+        cu_seq_len_k_start(cu_seq_len_k_start),
+        cu_seq_len_k_end(cu_seq_len_k_end) {
+    DG_STATIC_ASSERT(kNumSMs > 0, "Invalid SM count");
+  }
+
+  // NOTE: Out-of-range rows are clamped to the last valid row
+  CUTLASS_DEVICE bool next_q_block(uint32_t& q_block_idx,
+                                   uint32_t& kv_token_base,
+                                   uint32_t& num_kv_splits) {
+    if (current_q_block_idx >= num_q_blocks) return false;
+    q_block_idx = current_q_block_idx;
+    current_q_block_idx += kNumSMs;
+    uint32_t start = cute::numeric_limits<uint32_t>::max();
+    uint32_t end = cute::numeric_limits<uint32_t>::min();
+#pragma unroll 8
+    for (uint32_t token_idx = 0; token_idx < BLOCK_Q; ++token_idx) {
+      const auto row_idx =
+          cute::min(q_block_idx * BLOCK_Q + token_idx, num_q_tokens - 1);
+      start = cute::min(start,
+                        cute::min(cu_seq_len_k_start[row_idx], num_kv_tokens));
+      end = cute::max(end, cute::min(cu_seq_len_k_end[row_idx], num_kv_tokens));
+    }
+    kv_token_base = start / 4 * 4;
+    num_kv_splits = math::ceil_div(end - kv_token_base, SPLIT_KV);
+    return true;
+  }
+
+  CUTLASS_DEVICE uint32_t
+  get_q_tma_token_base(const uint32_t& q_block_idx) const {
+    return q_block_idx * BLOCK_Q;
+  }
+
+  CUTLASS_DEVICE static uint32_t get_kv_tma_offset(
+      const uint32_t& kv_token_base, const uint32_t& kv_split_idx) {
+    return kv_token_base + kv_split_idx * SPLIT_KV;
+  }
+
+  CUTLASS_DEVICE static uint32_t get_logits_row(const uint32_t& q_block_idx,
+                                                const uint32_t& token_idx) {
+    return q_block_idx * BLOCK_Q + token_idx;
+  }
+
+  CUTLASS_DEVICE static uint32_t get_logits_col(
+      const uint32_t& kv_token_base, const uint32_t& kv_split_idx,
+      const uint32_t& math_thread_idx) {
+    return kv_token_base + kv_split_idx * SPLIT_KV + math_thread_idx;
+  }
+};
+
+// Single-CTA TMA load of a BLOCK_OUTER x BLOCK_INNER box (3D adds a batch
+// coordinate).
+template <uint32_t BLOCK_INNER, uint32_t BLOCK_OUTER, uint32_t kSwizzleMode,
+          typename dtype_t, bool kIs3DTMA = false>
+CUTLASS_DEVICE void tma_copy(
+    void const* desc_ptr, cutlass::arch::ClusterTransactionBarrier* barrier_ptr,
+    dtype_t* smem_ptr, const uint32_t& inner_idx, const uint32_t& outer_idx,
+    const uint32_t& batch_idx = 0) {
+  constexpr auto kCacheHint =
+      static_cast<uint64_t>(cute::TMA::CacheHintSm90::EVICT_NORMAL);
+  constexpr uint32_t kPackFactor = deep_gemm::get_smem_pack_factor<dtype_t>();
+  DG_STATIC_ASSERT(kPackFactor == 1 or sizeof(dtype_t) == 1,
+                   "Packing expects a 1-byte storage type");
+  constexpr uint32_t BLOCK_INNER_ATOM_STORAGE =
+      kSwizzleMode == 0 ? BLOCK_INNER / kPackFactor
+                        : kSwizzleMode / sizeof(dtype_t);
+  constexpr uint32_t BLOCK_INNER_ATOM_LOGICAL =
+      BLOCK_INNER_ATOM_STORAGE * kPackFactor;
+  DG_STATIC_ASSERT(BLOCK_INNER % BLOCK_INNER_ATOM_LOGICAL == 0,
+                   "TMA inner block must contain whole atoms");
+
+#pragma unroll
+  for (uint32_t i = 0; i < BLOCK_INNER / BLOCK_INNER_ATOM_LOGICAL; ++i) {
+    if constexpr (kIs3DTMA) {
+      cute::SM90_TMA_LOAD_3D::copy(
+          desc_ptr, reinterpret_cast<uint64_t*>(barrier_ptr), kCacheHint,
+          smem_ptr + i * BLOCK_OUTER * BLOCK_INNER_ATOM_STORAGE,
+          inner_idx + i * BLOCK_INNER_ATOM_LOGICAL, outer_idx, batch_idx);
+    } else {
+      cute::SM90_TMA_LOAD_2D::copy(
+          desc_ptr, reinterpret_cast<uint64_t*>(barrier_ptr), kCacheHint,
+          smem_ptr + i * BLOCK_OUTER * BLOCK_INNER_ATOM_STORAGE,
+          inner_idx + i * BLOCK_INNER_ATOM_LOGICAL, outer_idx);
+    }
+  }
+}
+
+}  // namespace fused_mqa_topk::dg
diff --git a/csrc/fused_mqa_topk/dense_scan.cuh b/csrc/fused_mqa_topk/dense_scan.cuh
new file mode 100644
index 000000000000..395cb84f183e
--- /dev/null
+++ b/csrc/fused_mqa_topk/dense_scan.cuh
@@ -0,0 +1,289 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+//
+// Adapted from DeepGEMM (https://github.com/deepseek-ai/DeepGEMM) at
+// 88965b078186ee7510ab9fc4f1d5ebc19adfa8d1:
+// deep_gemm/include/deep_gemm/impls/sm100_fp8_mqa_logits.cuh.
+//
+// Dense FP8 MQA scoring of one 8-row block over its whole KV range: DeepGEMM's
+// SM100 MQA logits pipeline (contiguous KV, 32 heads, D = 128) with the logits
+// store replaced by a per-score policy. Used by the persistent kernel's exact
+// recovery rescan (rescan.cuh).
+//
+// Roles (384 threads): warps 0-7 reduce the 32 heads (4 rows per warpgroup),
+// warp 8 loads Q and the weights, warp 9 the KV tiles and their scales, warp 10
+// issues the UMMAs, warp 11 is idle.
+//
+// Policy interface (all called by the math warps):
+//   kWeightScale: score = (sum_h kWeightScale * w_h * (s_h + |s_h|)) * scale_kv
+//   begin_block(sched, q_block, query_base): after the block's Q / weights
+//       landed
+//   begin_split(split): before the split's scores
+//   score(i, kv_offset, kv_lane, split, score): row query_base + i of the
+//       warpgroup
+//   end_split(split, num_splits, kv_base): after the split's scores
+//   end_block(sched, q_block): before the block's Q stage is released
+#pragma once
+
+#include <cutlass/arch/reg_reconfig.h>
+#include <cute/arch/cluster_sm90.hpp>
+#include <cute/arch/copy_sm90_desc.hpp>
+
+#include <deep_gemm/common/cute_tie.cuh>
+#include <deep_gemm/common/ring_pipeline.cuh>
+#include <deep_gemm/common/utils.cuh>
+#include <deep_gemm/mma/sm100.cuh>
+#include <deep_gemm/ptx/ld_st.cuh>
+#include <deep_gemm/ptx/tcgen05.cuh>
+#include <deep_gemm/ptx/utils.cuh>
+
+#include "common.cuh"
+#include "constants.h"
+#include "deep_gemm_compat.cuh"
+
+namespace fused_mqa_topk::dense {
+using namespace deep_gemm;
+
+constexpr uint32_t BLOCK_Q = kBlockRows, SPLIT_KV = kTileKeys;
+constexpr uint32_t kNumKVStages = 2, kNumTmemStages = 2;
+constexpr uint32_t UMMA_M = 128, UMMA_N = BLOCK_Q * kNumHeads, UMMA_K = 32;
+// The rescan scores exactly one 8-row block: the scheduler is built with that
+// block as its `sm_idx`, and its grid stride (kNumSMs) exceeds any block count
+// (rows < 2^26), so next_q_block yields that block only. A smaller stride would
+// also score block + stride, outside the rescan's row bounds and output rows.
+constexpr uint32_t kSingleBlockStride = 1u << 31;
+using Scheduler =
+    dg::SM100MQALogitsScheduler<BLOCK_Q, SPLIT_KV, kSingleBlockStride>;
+using SharedStorage =
+    dg::MQALogitsSharedStorage<kNumHeads, kHeadDim, BLOCK_Q, SPLIT_KV, 1,
+                               kNumKVStages, kNumTmemStages,
+                               cutlass::float_e4m3_t, float>;
+
+// The policy's state lives in dynamic shared memory right after the pipeline
+// storage. Barrier words are zeroed before init: the persistent kernel reuses
+// this memory between rescans.
+template <typename Policy>
+CUTLASS_DEVICE void score_block(const Scheduler& block, Policy& policy,
+                                const cute::TmaDescriptor& tensor_map_q,
+                                const cute::TmaDescriptor& tensor_map_kv,
+                                const cute::TmaDescriptor& tensor_map_sf_kv,
+                                const cute::TmaDescriptor& tensor_map_weights) {
+  const auto warp_idx = cutlass::canonical_warp_idx_sync();
+  const auto lane_idx = ptx::get_lane_idx();
+  constexpr uint32_t kSpecWarpStart = 8;
+  constexpr uint32_t SMEM_Q_SIZE_PER_STAGE = BLOCK_Q * kNumHeads * kHeadDim;
+  constexpr uint32_t SMEM_KV_SIZE_PER_STAGE = SPLIT_KV * kHeadDim;
+  constexpr uint32_t SMEM_SF_KV_SIZE_PER_STAGE = SPLIT_KV * sizeof(float);
+  constexpr uint32_t SMEM_WEIGHT_SIZE_PER_STAGE =
+      BLOCK_Q * kNumHeads * sizeof(float);
+  constexpr uint32_t kNumTmemCols =
+      utils::get_num_aligned_tmem_cols<UMMA_N * kNumTmemStages>();
+
+  if (warp_idx == kSpecWarpStart) {
+    cute::prefetch_tma_descriptor(&tensor_map_q);
+    cute::prefetch_tma_descriptor(&tensor_map_weights);
+    cute::prefetch_tma_descriptor(&tensor_map_kv);
+    cute::prefetch_tma_descriptor(&tensor_map_sf_kv);
+  }
+
+  extern __shared__ __align__(SharedStorage::kSwizzleAlignment)
+      uint8_t smem_buffer[];
+  auto& smem = *reinterpret_cast<SharedStorage*>(smem_buffer);
+  if (warp_idx == kSpecWarpStart + 1 and cute::elect_one_sync()) {
+    const auto init = [](auto& barrier, uint32_t count) {
+      *reinterpret_cast<volatile uint64_t*>(&barrier) = 0;
+      barrier.init(count);
+    };
+    init(smem.full_q_barriers[0], 1);
+    init(smem.empty_q_barriers[0], kMathThreads + 32);  // + the UMMA warp
+#pragma unroll
+    for (uint32_t i = 0; i < kNumKVStages; ++i) {
+      init(smem.full_kv_barriers[i], 1);
+      init(smem.empty_kv_barriers[i], kMathThreads);
+    }
+#pragma unroll
+    for (uint32_t i = 0; i < kNumTmemStages; ++i) {
+      init(smem.full_tmem_barriers[i], 1);
+      init(smem.empty_tmem_barriers[i], kMathThreads);
+    }
+    cutlass::arch::fence_barrier_init();
+  }
+  __syncwarp();
+  if (warp_idx == kSpecWarpStart + 2)
+    cute::TMEM::Allocator1Sm().allocate(kNumTmemCols, &smem.tmem_ptr_in_smem);
+  __syncthreads();
+
+  RingPipeline<1> q_pipeline;
+  RingPipeline<kNumKVStages> kv_pipeline;
+  RingPipeline<kNumTmemStages> tmem_pipeline;
+
+  uint32_t q_block_idx, kv_base, num_kv_splits;
+  if (warp_idx == kSpecWarpStart) {
+    cutlass::arch::warpgroup_reg_dealloc<kSpecRegs>();
+    if (cute::elect_one_sync()) {
+      Scheduler scheduler = block;
+      while (scheduler.next_q_block(q_block_idx, kv_base, num_kv_splits)) {
+        CUTE_TIE_DECL(q_pipeline.advance(), q_stage_idx, q_phase);
+        smem.empty_q_barriers[q_stage_idx].wait(q_phase ^ 1);
+        const uint32_t q_token_base =
+            scheduler.get_q_tma_token_base(q_block_idx);
+        dg::tma_copy<kHeadDim, BLOCK_Q * kNumHeads, 0>(
+            &tensor_map_q, &smem.full_q_barriers[q_stage_idx],
+            smem.smem_q[q_stage_idx], 0, q_token_base * kNumHeads);
+        dg::tma_copy<kNumHeads, BLOCK_Q, 0>(
+            &tensor_map_weights, &smem.full_q_barriers[q_stage_idx],
+            smem.smem_weights[q_stage_idx], 0, q_token_base);
+        smem.full_q_barriers[q_stage_idx].arrive_and_expect_tx(
+            SMEM_Q_SIZE_PER_STAGE + SMEM_WEIGHT_SIZE_PER_STAGE);
+      }
+    }
+    __syncwarp();
+  } else if (warp_idx == kSpecWarpStart + 1) {
+    cutlass::arch::warpgroup_reg_dealloc<kSpecRegs>();
+    Scheduler scheduler = block;
+    while (scheduler.next_q_block(q_block_idx, kv_base, num_kv_splits)) {
+#pragma unroll 1
+      for (uint32_t kv_split_idx = 0; kv_split_idx < num_kv_splits;
+           ++kv_split_idx) {
+        if (cute::elect_one_sync()) {
+          CUTE_TIE_DECL(kv_pipeline.advance(), kv_stage_idx, kv_phase);
+          smem.empty_kv_barriers[kv_stage_idx].wait(kv_phase ^ 1);
+          const uint32_t kv_tma_offset =
+              scheduler.get_kv_tma_offset(kv_base, kv_split_idx);
+          dg::tma_copy<kHeadDim, SPLIT_KV, 0>(
+              &tensor_map_kv, &smem.full_kv_barriers[kv_stage_idx],
+              smem.smem_kv[kv_stage_idx], 0, kv_tma_offset);
+          dg::tma_copy<SPLIT_KV, 1, 0>(
+              &tensor_map_sf_kv, &smem.full_kv_barriers[kv_stage_idx],
+              smem.smem_sf_kv[kv_stage_idx], kv_tma_offset, 0);
+          smem.full_kv_barriers[kv_stage_idx].arrive_and_expect_tx(
+              SMEM_KV_SIZE_PER_STAGE + SMEM_SF_KV_SIZE_PER_STAGE);
+        }
+        __syncwarp();
+      }
+    }
+  } else if (warp_idx == kSpecWarpStart + 2) {
+    cutlass::arch::warpgroup_reg_dealloc<kSpecRegs>();
+    DG_TRAP_ONLY_DEVICE_ASSERT(ptx::ld_shared(&smem.tmem_ptr_in_smem) == 0);
+    auto instr_desc = cute::UMMA::make_instr_desc<
+        cutlass::float_e4m3_t, cutlass::float_e4m3_t, float, UMMA_M, UMMA_N,
+        cute::UMMA::Major::K, cute::UMMA::Major::K>();
+    auto runtime_instr_desc = cute::UMMA::make_runtime_instr_desc(instr_desc);
+    Scheduler scheduler = block;
+    while (scheduler.next_q_block(q_block_idx, kv_base, num_kv_splits)) {
+      CUTE_TIE_DECL(q_pipeline.advance(), q_stage_idx, q_phase);
+      smem.full_q_barriers[q_stage_idx].wait(q_phase);
+      for (uint32_t kv_split_idx = 0; kv_split_idx < num_kv_splits;
+           ++kv_split_idx) {
+        CUTE_TIE_DECL(kv_pipeline.advance(), kv_stage_idx, kv_phase);
+        smem.full_kv_barriers[kv_stage_idx].wait(kv_phase);
+        if (cute::elect_one_sync()) {
+          // One UMMA per math warpgroup: KV rows [128 i, 128 i + 128) x all 256
+          // Q rows.
+#pragma unroll
+          for (uint32_t i = 0; i < 2; ++i) {
+            CUTE_TIE_DECL(tmem_pipeline.advance(), tmem_stage_idx, tmem_phase);
+            const uint32_t tmem_addr = tmem_stage_idx * UMMA_N;
+            smem.empty_tmem_barriers[tmem_stage_idx].wait(tmem_phase ^ 1);
+            ptx::tcgen05_after_thread_sync();
+#pragma unroll
+            for (uint32_t k = 0; k < kHeadDim / UMMA_K; ++k) {
+              auto a_desc = mma::sm100::make_umma_desc<cute::UMMA::Major::K, 0,
+                                                       kHeadDim, kHeadDim>(
+                  smem.smem_kv[kv_stage_idx], i * UMMA_M, k * UMMA_K);
+              auto b_desc = mma::sm100::make_umma_desc<cute::UMMA::Major::K, 0,
+                                                       kHeadDim, kHeadDim>(
+                  smem.smem_q[q_stage_idx], 0, k * UMMA_K);
+              ptx::SM100_MMA_F8F6F4_SS::fma(a_desc, b_desc, tmem_addr, k,
+                                            runtime_instr_desc);
+            }
+            ptx::umma_arrive_no_elect(smem.full_tmem_barriers[tmem_stage_idx]);
+          }
+        }
+        __syncwarp();
+      }
+      smem.empty_q_barriers[q_stage_idx].arrive();
+    }
+  } else if (warp_idx == kSpecWarpStart + 3) {
+    cutlass::arch::warpgroup_reg_dealloc<kSpecRegs>();
+  } else {
+    cutlass::arch::warpgroup_reg_alloc<kMathRegs>();
+    const uint32_t math_thread_idx = warp_idx * 32 + lane_idx;
+    const uint32_t query_base = (warp_idx / 4) * 4;
+    float weights[4][kNumHeads];
+    Scheduler scheduler = block;
+    while (scheduler.next_q_block(q_block_idx, kv_base, num_kv_splits)) {
+      CUTE_TIE_DECL(q_pipeline.advance(), q_stage_idx, q_phase);
+      smem.full_q_barriers[q_stage_idx].wait(q_phase);
+#pragma unroll
+      for (uint32_t i = 0; i < 4; ++i)
+#pragma unroll
+        for (uint32_t j = 0; j < kNumHeads; ++j)
+          weights[i][j] = Policy::kWeightScale *
+                          ptx::ld_shared(smem.smem_weights[q_stage_idx] +
+                                         (query_base + i) * kNumHeads + j);
+      policy.begin_block(scheduler, q_block_idx, query_base);
+      for (uint32_t kv_split_idx = 0; kv_split_idx < num_kv_splits;
+           ++kv_split_idx) {
+        policy.begin_split(kv_split_idx);
+        CUTE_TIE_DECL(kv_pipeline.advance(), kv_stage_idx, kv_phase);
+        smem.full_kv_barriers[kv_stage_idx].wait(kv_phase);
+#pragma unroll
+        for (uint32_t kv_half = 0; kv_half < 2; ++kv_half) {
+          const uint32_t kv_lane =
+              (math_thread_idx % kMathWgThreads) + kv_half * kMathWgThreads;
+          const uint32_t kv_offset =
+              scheduler.get_logits_col(kv_base, kv_split_idx, kv_lane);
+          const float scale_kv =
+              ptx::ld_shared(smem.smem_sf_kv[kv_stage_idx] + kv_lane);
+          CUTE_TIE_DECL(tmem_pipeline.advance(), tmem_stage_idx, tmem_phase);
+          smem.full_tmem_barriers[tmem_stage_idx].wait(tmem_phase);
+          ptx::tcgen05_after_thread_sync();
+          // Both KV half-MMAs have committed before the shared stage is freed.
+          if (kv_half == 1) {
+            cutlass::arch::fence_view_async_shared();
+            smem.empty_kv_barriers[kv_stage_idx].arrive();
+          }
+          float accum_pair[2][kNumHeads];
+#pragma unroll
+          for (uint32_t pair = 0; pair < 4; pair += 2) {
+            // Two query rows at a time keep the live accumulator registers low.
+            ptx::tmem_load_32dp32b<2 * kNumHeads>(
+                tmem_stage_idx * UMMA_N + (query_base + pair) * kNumHeads,
+                reinterpret_cast<uint32_t*>(accum_pair));
+            cutlass::arch::fence_view_async_tmem_load();
+            if (pair == 2) {
+              ptx::tcgen05_before_thread_sync();
+              smem.empty_tmem_barriers[tmem_stage_idx].arrive();
+            }
+#pragma unroll
+            for (uint32_t r = 0; r < 2; ++r)
+              policy.score(
+                  pair + r, kv_offset, kv_lane, kv_split_idx,
+                  reduce_heads<true>(accum_pair[r], weights[pair + r]) *
+                      scale_kv);
+          }
+        }
+        policy.end_split(kv_split_idx, num_kv_splits, kv_base);
+      }
+      policy.end_block(scheduler, q_block_idx);
+      cutlass::arch::fence_view_async_shared();
+      smem.empty_q_barriers[q_stage_idx].arrive();
+    }
+    cutlass::arch::NamedBarrier(kMathThreads, kBarDenseMath).sync();
+    if (warp_idx == 0) cute::TMEM::Allocator1Sm().free(0, kNumTmemCols);
+  }
+}
+
+// Invalidates the pipeline's barriers so the next score_block can re-initialise
+// the same shared memory (mbarrier.inval before reuse).
+CUTLASS_DEVICE void invalidate_barriers(SharedStorage& smem) {
+  for (auto& b : smem.full_q_barriers) mbar_inval(&b);
+  for (auto& b : smem.empty_q_barriers) mbar_inval(&b);
+  for (auto& b : smem.full_kv_barriers) mbar_inval(&b);
+  for (auto& b : smem.empty_kv_barriers) mbar_inval(&b);
+  for (auto& b : smem.full_tmem_barriers) mbar_inval(&b);
+  for (auto& b : smem.empty_tmem_barriers) mbar_inval(&b);
+}
+
+}  // namespace fused_mqa_topk::dense
diff --git a/csrc/fused_mqa_topk/fused_mqa_topk.cu b/csrc/fused_mqa_topk/fused_mqa_topk.cu
new file mode 100644
index 000000000000..652453494fbe
--- /dev/null
+++ b/csrc/fused_mqa_topk/fused_mqa_topk.cu
@@ -0,0 +1,394 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+//
+// Single-launch fused FP8 MQA-logits + top-k (k = 512, 1024, 2048) prefill
+// selection (SM100 family): one persistent kernel (persistent_kernel.cuh, one
+// CTA per SM; sampled gate, candidate emission and selection, exact rescan of
+// out-of-range blocks) scores every visible key of a row and selects its top-k
+// keys without materializing the [rows, keys] logits matrix.
+// This file is the libtorch-stable host wrapper. All scratch memory is passed
+// in by the caller (no allocation here), so the op is CUDA-graph friendly.
+
+#include "../torch_utils.h"
+
+#include <torch/headeronly/core/ScalarType.h>
+
+#include <cuda.h>
+#include <cudaTypedefs.h>
+#include <cuda_runtime.h>
+
+#include <algorithm>
+#include <cstdint>
+#include <map>
+#include <mutex>
+#include <utility>
+
+#include "constants.h"
+#include "persistent_kernel.cuh"
+
+namespace {
+
+using fused_mqa_topk::kBlockRows;
+using fused_mqa_topk::kCandidatesPerRow;
+using fused_mqa_topk::kHeadDim;
+using fused_mqa_topk::kMaxKeys;
+using fused_mqa_topk::kNumHeads;
+using fused_mqa_topk::kNumThreads;
+using fused_mqa_topk::kRetainedPerRow;
+using fused_mqa_topk::kSlotsPerCta;
+using fused_mqa_topk::kTileKeys;
+using torch::headeronly::ScalarType;
+namespace persist = fused_mqa_topk::persist;
+
+void check_cuda(cudaError_t err) {
+  STD_TORCH_CHECK(err == cudaSuccess, cudaGetErrorString(err));
+}
+
+// cuTensorMapEncodeTiled from the driver, resolved once through the runtime
+// (thread-safe static initialisation; no link-time libcuda dependency).
+PFN_cuTensorMapEncodeTiled_v12000 encode_tiled() {
+  static const PFN_cuTensorMapEncodeTiled_v12000 fn = [] {
+    void* ptr = nullptr;
+    cudaDriverEntryPointQueryResult status;
+    const cudaError_t err = cudaGetDriverEntryPointByVersion(
+        "cuTensorMapEncodeTiled", &ptr, 12000, cudaEnableDefault, &status);
+    STD_TORCH_CHECK(err == cudaSuccess && status == cudaDriverEntryPointSuccess,
+                    "failed to resolve cuTensorMapEncodeTiled");
+    return reinterpret_cast<PFN_cuTensorMapEncodeTiled_v12000>(ptr);
+  }();
+  return fn;
+}
+
+// SM count of `device`, queried once per device.
+int sm_count(int device) {
+  static std::mutex mutex;
+  static std::map<int, int> cache;
+  const std::lock_guard<std::mutex> lock(mutex);
+  auto [it, inserted] = cache.try_emplace(device, 0);
+  if (inserted) {
+    const cudaError_t err = cudaDeviceGetAttribute(
+        &it->second, cudaDevAttrMultiProcessorCount, device);
+    if (err != cudaSuccess) cache.erase(it);
+    check_cuda(err);
+  }
+  return it->second;
+}
+
+// Adapted from DeepGEMM's make_tma_2d_desc
+// (csrc/jit_kernels/impls/runtime_utils.hpp at 88965b078186)
+CUtensorMap make_2d(void* ptr, CUtensorMapDataType dt, int elem_size,
+                    int gmem_inner, int gmem_outer, int smem_inner,
+                    int smem_outer, long gmem_outer_stride, int swizzle_mode) {
+  if (swizzle_mode != 0) smem_inner = swizzle_mode / elem_size;
+  CUtensorMap tm;
+  const cuuint64_t gdims[2] = {(cuuint64_t)gmem_inner, (cuuint64_t)gmem_outer};
+  const cuuint32_t sdims[2] = {(cuuint32_t)smem_inner, (cuuint32_t)smem_outer};
+  const cuuint64_t gstrides[1] = {(cuuint64_t)(gmem_outer_stride * elem_size)};
+  const cuuint32_t estrides[2] = {1, 1};
+  CUtensorMapSwizzle swizzle = swizzle_mode == 128  ? CU_TENSOR_MAP_SWIZZLE_128B
+                               : swizzle_mode == 64 ? CU_TENSOR_MAP_SWIZZLE_64B
+                               : swizzle_mode == 32
+                                   ? CU_TENSOR_MAP_SWIZZLE_32B
+                                   : CU_TENSOR_MAP_SWIZZLE_NONE;
+  CUresult r = encode_tiled()(&tm, dt, 2, ptr, gdims, gstrides, sdims, estrides,
+                              CU_TENSOR_MAP_INTERLEAVE_NONE, swizzle,
+                              CU_TENSOR_MAP_L2_PROMOTION_L2_256B,
+                              CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE);
+  STD_TORCH_CHECK(r == CUDA_SUCCESS, "cuTensorMapEncodeTiled failed: ", (int)r);
+  return tm;
+}
+
+// KV [keys,128] bytes viewed as {128, groups, stride} with strides
+// {128*stride, 128}: box {128, 256, 1} at (0, 0, o) loads keys o, o+stride, ...
+// with 128B swizzle. The view spans groups * stride keys, which can exceed
+// `keys` by < stride / 2; that is safe only because the kernel loads offset
+// o = stride / 2, whose last key ((groups - 1) * stride + stride / 2) is
+// always < keys (launch_config sizes `groups` for that).
+CUtensorMap make_strided_kv(void* ptr, uint32_t groups, uint32_t stride) {
+  CUtensorMap tm;
+  constexpr uint64_t kRowBytes = kHeadDim;
+  const cuuint64_t gdims[3] = {kRowBytes, groups, stride};
+  const cuuint64_t gstrides[2] = {kRowBytes * stride, kRowBytes};
+  const cuuint32_t box[3] = {kRowBytes, kTileKeys, 1};
+  const cuuint32_t estrides[3] = {1, 1, 1};
+  CUresult r = encode_tiled()(
+      &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 3, ptr, gdims, gstrides, box,
+      estrides, CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B,
+      CU_TENSOR_MAP_L2_PROMOTION_L2_256B, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE);
+  STD_TORCH_CHECK(r == CUDA_SUCCESS,
+                  "strided KV cuTensorMapEncodeTiled failed: ", (int)r);
+  return tm;
+}
+
+// -------- launch configuration --------
+// Keys per row above which the long-context instantiations (Interval 64, 128)
+// and gate settings apply.
+constexpr uint32_t kLongContextKeys = 131072;
+
+using PersistentKernel =
+    decltype(&persist::fused_mqa_topk_fp8_prefill_kernel<16, 2048>);
+
+// vLLM dispatches chunks of >= 32K keys per row here (fused_mqa_topk_min_keys);
+// smaller key counts are supported but take the same <= 32K configuration.
+template <uint32_t TopK>
+PersistentKernel persistent_kernel(uint32_t keys) {
+  return keys <= 32768 ? persist::fused_mqa_topk_fp8_prefill_kernel<16, TopK>
+         : keys <= kLongContextKeys
+             ? persist::fused_mqa_topk_fp8_prefill_kernel<32, TopK>
+         : keys <= 524288
+             ? persist::fused_mqa_topk_fp8_prefill_kernel<64, TopK>
+             : persist::fused_mqa_topk_fp8_prefill_kernel<128, TopK>;
+}
+
+struct LaunchConfig {
+  PersistentKernel kernel;
+  int smem;              // dynamic shared memory bytes
+  uint32_t phase_bytes;  // selection scratch bytes past the Q / weights region
+  // Gate sample: sample_tiles tiles of 256 keys on a uniform grid (stride
+  // sample_stride).
+  uint32_t sample_target, sample_stride, sample_groups, sample_tiles;
+};
+
+// Dynamic shared memory of `kernel` on `device`: everything past its static
+// shared memory up to the device's opt-in limit (232448 bytes on SM100; the
+// staged selection uses it all), at least the exact rescan's layout. Sets
+// the opt-in attribute; both are computed once per device and kernel.
+cudaError_t kernel_smem(int device, PersistentKernel kernel, int& smem) {
+  // The exact rescan reuses the same dynamic shared memory with its own layout.
+  constexpr int kRescanSmem = int(sizeof(fused_mqa_topk::dense::SharedStorage) +
+                                  sizeof(persist::RescanShared));
+  constexpr int kMinSmem = int(
+      sizeof(persist::KernelSmem) +
+      std::max(sizeof(persist::HandoffRing), sizeof(persist::SampleShared)));
+  static std::mutex mutex;
+  static std::map<std::pair<int, PersistentKernel>, int> cache;
+  const std::lock_guard<std::mutex> lock(mutex);
+  if (const auto it = cache.find({device, kernel}); it != cache.end()) {
+    smem = it->second;
+    return cudaSuccess;
+  }
+  int max_smem = 0;
+  if (const cudaError_t e = cudaDeviceGetAttribute(
+          &max_smem, cudaDevAttrMaxSharedMemoryPerBlockOptin, device);
+      e != cudaSuccess)
+    return e;
+  cudaFuncAttributes attr;
+  if (const cudaError_t e =
+          cudaFuncGetAttributes(&attr, reinterpret_cast<const void*>(kernel));
+      e != cudaSuccess)
+    return e;
+  const int bytes = std::max(kRescanSmem, max_smem - int(attr.sharedSizeBytes));
+  if (bytes < kMinSmem) return cudaErrorInvalidConfiguration;
+  if (const cudaError_t e = cudaFuncSetAttribute(
+          reinterpret_cast<const void*>(kernel),
+          cudaFuncAttributeMaxDynamicSharedMemorySize, bytes);
+      e != cudaSuccess)
+    return e;
+  cache.emplace(std::make_pair(device, kernel), bytes);
+  smem = bytes;
+  return cudaSuccess;
+}
+
+// Picks the kernel instantiation, its shared memory and the gate sample for
+// `keys` keys and top-`topk` rows on `device` (the current device).
+cudaError_t launch_config(int device, uint32_t keys, uint32_t topk,
+                          LaunchConfig& c) {
+  // The gate aims for `target` candidates per row and sizes the sample so its
+  // quantile rank is `rank` (count error ~1/sqrt(rank)); a row whose count
+  // falls below the top-k is rescanned exactly. Up to 128K keys the candidate
+  // stream dominates, so smaller top-ks aim lower with a coarser sample. Above
+  // 128K the sample is clamped to 16-32 tiles (rank ~32 at 1M keys): a sample
+  // tile costs about two scan tiles, so a larger sample costs more than a
+  // smaller candidate set saves. Tuned on B300 real captures.
+  struct Gate {
+    uint32_t target, rank;
+  };
+  constexpr Gate kGates[3][2] = {
+      // <= 128K       > 128K
+      {{1024, 32}, {1536, 64}},  // top-512
+      {{1536, 48}, {2048, 48}},  // top-1024
+      {{3072, 64}, {4096, 64}},  // top-2048
+  };
+  const bool long_ctx = keys > kLongContextKeys;
+  const Gate gate =
+      kGates[topk == 512 ? 0 : (topk == 1024 ? 1 : 2)][long_ctx ? 1 : 0];
+  c.sample_target = gate.target;
+  const uint32_t rank = gate.rank;
+  const uint64_t tile = kTileKeys;
+  const uint32_t rank_tiles = std::max(
+      1u, uint32_t((uint64_t(rank) * keys + tile * c.sample_target - 1) /
+                   (tile * c.sample_target)));
+  c.sample_tiles =
+      long_ctx ? std::min(32u, std::max(16u, rank_tiles)) : rank_tiles;
+  c.sample_stride = std::max(1u, uint32_t((keys + tile * c.sample_tiles - 1) /
+                                          (tile * c.sample_tiles)));
+  c.sample_groups =
+      (keys - c.sample_stride / 2 + c.sample_stride - 1) / c.sample_stride;
+  c.kernel = topk == 512    ? persistent_kernel<512>(keys)
+             : topk == 1024 ? persistent_kernel<1024>(keys)
+                            : persistent_kernel<2048>(keys);
+  if (const cudaError_t e = kernel_smem(device, c.kernel, c.smem);
+      e != cudaSuccess)
+    return e;
+  c.phase_bytes = uint32_t(c.smem) - persist::kSelBase;
+  return cudaSuccess;
+}
+
+void check_tensor(const torch::stable::Tensor& t, int32_t device,
+                  ScalarType dtype, const char* name) {
+  STD_TORCH_CHECK(t.is_cuda() && t.get_device_index() == device, name,
+                  " must be on the same CUDA device as q");
+  STD_TORCH_CHECK(t.scalar_type() == dtype, name, " has the wrong dtype");
+  STD_TORCH_CHECK(t.is_contiguous(), name, " must be contiguous");
+}
+
+// TMA tensor maps need 16-byte aligned base addresses.
+void check_aligned(const torch::stable::Tensor& t, const char* name) {
+  STD_TORCH_CHECK(reinterpret_cast<uintptr_t>(t.data_ptr()) % 16 == 0, name,
+                  " must be 16-byte aligned");
+}
+
+}  // namespace
+
+// Inputs: q [rows,32,128] fp8e4m3, k [keys,128] fp8e4m3, scales [keys] fp32,
+// weights [rows,32] fp32, starts/ends [rows] int32 (the visible key range
+// [start,end) of each row), with 0 < rows < 2^26 and 0 < keys <= 1048576 (the
+// last 8-row block may be partial). starts/ends are expected in [0, keys] with
+// start <= end; they are read as uint32_t and clamped to keys, so a negative
+// value means `keys` and an end before its start means an empty row.
+//
+// Output: out [rows,K] int32 with K (the top-k) 512, 1024 or 2048: key indices
+// relative to the row's start (key - start, the top_k_per_row_prefill
+// convention), unordered; -1 padded when fewer than K keys are visible. Rows
+// that see at most K keys get them all, in order. Exact up to ties: the gated
+// path ranks fp24 score codes (the low 8 mantissa bits dropped), the exact
+// rescan ranks fp32 scores.
+//
+// All tensors are contiguous; q, k, scales, weights and out are 16-byte aligned
+// (TMA / vector stores).
+//
+// Workspace (caller-owned, reused across calls), with B = ceil(rows / 8) and
+// S >= 2 slots per CTA of the persistent grid (one CTA per SM):
+//   counter [2] int32 (zero between calls; the kernel leaves it zeroed, so a
+//     workspace must not be shared by launches that can run concurrently, e.g.
+//     on two streams),
+//   values [8S,16384] uint16 (low 16 bits of the fp24 candidate codes),
+//   indices [8S,16384] int32 (key index in bits 0-19, the code's high 8 bits
+//   in bits 20-27),
+//   retained [8S,4096] fp32, retained_indices [8S,4096] int32,
+//   counts [>=8B] int32, defer [>=B] int32.
+void fused_mqa_topk_prefill(
+    torch::stable::Tensor q, torch::stable::Tensor k,
+    torch::stable::Tensor scales, torch::stable::Tensor weights,
+    torch::stable::Tensor starts, torch::stable::Tensor ends,
+    torch::stable::Tensor out, torch::stable::Tensor counter,
+    torch::stable::Tensor values, torch::stable::Tensor indices,
+    torch::stable::Tensor retained, torch::stable::Tensor retained_indices,
+    torch::stable::Tensor counts, torch::stable::Tensor defer) {
+  STD_TORCH_CHECK(q.is_cuda(), "q must be a CUDA tensor");
+  const int32_t dev = q.get_device_index();
+  check_tensor(q, dev, ScalarType::Float8_e4m3fn, "q");
+  check_tensor(k, dev, ScalarType::Float8_e4m3fn, "k");
+  check_tensor(scales, dev, ScalarType::Float, "scales");
+  check_tensor(weights, dev, ScalarType::Float, "weights");
+  check_tensor(starts, dev, ScalarType::Int, "starts");
+  check_tensor(ends, dev, ScalarType::Int, "ends");
+  check_tensor(out, dev, ScalarType::Int, "out");
+  check_tensor(counter, dev, ScalarType::Int, "counter");
+  check_tensor(values, dev, ScalarType::UInt16, "values");
+  check_tensor(indices, dev, ScalarType::Int, "indices");
+  check_tensor(retained, dev, ScalarType::Float, "retained");
+  check_tensor(retained_indices, dev, ScalarType::Int, "retained_indices");
+  check_tensor(counts, dev, ScalarType::Int, "counts");
+  check_tensor(defer, dev, ScalarType::Int, "defer");
+
+  STD_TORCH_CHECK(
+      q.dim() == 3 && q.size(1) == kNumHeads && q.size(2) == kHeadDim,
+      "q must be [rows, 32, 128]");
+  STD_TORCH_CHECK(k.dim() == 2 && k.size(1) == kHeadDim,
+                  "k must be [keys, 128]");
+  const int64_t rows = q.size(0), keys = k.size(0);
+  // TMA row coordinates (rows * 32 Q rows) are int32.
+  STD_TORCH_CHECK(rows > 0 && rows < (int64_t(1) << 26),
+                  "rows must be in [1, 2^26)");
+  STD_TORCH_CHECK(keys > 0 && keys <= kMaxKeys, "keys must be in [1, 1048576]");
+  const int64_t blocks = (rows + kBlockRows - 1) / kBlockRows;
+  STD_TORCH_CHECK(scales.numel() == keys, "scales must be [keys]");
+  STD_TORCH_CHECK(weights.dim() == 2 && weights.size(0) == rows &&
+                      weights.size(1) == kNumHeads,
+                  "weights must be [rows, 32]");
+  STD_TORCH_CHECK(starts.numel() == rows && ends.numel() == rows,
+                  "starts/ends must be [rows]");
+  STD_TORCH_CHECK(out.dim() == 2 && out.size(0) == rows &&
+                      fused_mqa_topk::is_supported_topk(out.size(1)),
+                  "out must be [rows, K] with K in {512, 1024, 2048}");
+  check_aligned(q, "q");
+  check_aligned(k, "k");
+  check_aligned(scales, "scales");
+  check_aligned(weights, "weights");
+  check_aligned(out, "out");
+
+  torch::stable::accelerator::DeviceGuard guard(dev);
+  const cudaStream_t stream = get_current_cuda_stream(dev);
+  const uint32_t urows = static_cast<uint32_t>(rows),
+                 ukeys = static_cast<uint32_t>(keys),
+                 ublocks = static_cast<uint32_t>(blocks),
+                 topk = static_cast<uint32_t>(out.size(1));
+  const auto* starts_p = static_cast<const uint32_t*>(starts.data_ptr());
+  const auto* ends_p = static_cast<const uint32_t*>(ends.data_ptr());
+  auto* out_p = static_cast<int*>(out.data_ptr());
+  // Q: rows * 32 head rows of 128 bytes, a block's 256 per box; K: keys rows of
+  // 128 bytes, a tile's 256 per box (both 128B-swizzled); per-key scales: a
+  // tile's 256; weights: a block's 8 rows of 32.
+  const int head_rows = int(rows) * int(kNumHeads);
+  auto tq = make_2d(q.data_ptr(), CU_TENSOR_MAP_DATA_TYPE_UINT8, 1, kHeadDim,
+                    head_rows, kHeadDim, kBlockRows * kNumHeads, kHeadDim, 128);
+  auto tk = make_2d(k.data_ptr(), CU_TENSOR_MAP_DATA_TYPE_UINT8, 1, kHeadDim,
+                    int(keys), kHeadDim, kTileKeys, kHeadDim, 128);
+  auto ts = make_2d(scales.data_ptr(), CU_TENSOR_MAP_DATA_TYPE_FLOAT32, 4,
+                    int(keys), 1, kTileKeys, 1, 0, 0);
+  auto tw = make_2d(weights.data_ptr(), CU_TENSOR_MAP_DATA_TYPE_FLOAT32, 4,
+                    kNumHeads, int(rows), kNumHeads, kBlockRows, kNumHeads, 0);
+
+  STD_TORCH_CHECK(counter.numel() == 2, "counter must be [2]");
+  constexpr int64_t kCtaRows = kSlotsPerCta * kBlockRows;
+  STD_TORCH_CHECK(values.dim() == 2 && values.size(1) == kCandidatesPerRow &&
+                      values.size(0) % kCtaRows == 0 && indices.dim() == 2 &&
+                      indices.size(0) == values.size(0) &&
+                      indices.size(1) == kCandidatesPerRow,
+                  "values/indices must be [8 * slots, 16384], slots even");
+  STD_TORCH_CHECK(values.size(0) > 0, "the candidate workspace is empty");
+  STD_TORCH_CHECK(retained.dim() == 2 && retained.size(1) == kRetainedPerRow &&
+                      retained.size(0) == values.size(0) &&
+                      retained_indices.dim() == 2 &&
+                      retained_indices.size(0) == retained.size(0) &&
+                      retained_indices.size(1) == kRetainedPerRow,
+                  "retained/retained_indices must be [8 * slots, 4096]");
+  STD_TORCH_CHECK(
+      counts.numel() >= blocks * kBlockRows && defer.numel() >= blocks,
+      "per-row workspace too small");
+  // One persistent CTA per SM (capped by the blocks and the workspace slots);
+  // CTA c owns candidate slots 2c and 2c + 1.
+  const uint32_t grid = static_cast<uint32_t>(std::min<int64_t>(
+      {int64_t(sm_count(dev)), blocks, values.size(0) / kCtaRows}));
+  LaunchConfig c;
+  const cudaError_t config_err = launch_config(dev, ukeys, topk, c);
+  STD_TORCH_CHECK(config_err != cudaErrorInvalidConfiguration,
+                  "fused_mqa_topk_prefill needs more shared memory per block "
+                  "than this GPU allows");
+  check_cuda(config_err);
+  auto tks = make_strided_kv(k.data_ptr(), c.sample_groups, c.sample_stride);
+  const fused_mqa_topk::rescan::RescanBuffers rescan{
+      static_cast<float*>(retained.data_ptr()),
+      static_cast<int*>(retained_indices.data_ptr())};
+  c.kernel<<<grid, kNumThreads, c.smem, stream>>>(
+      urows, ukeys, starts_p, ends_p, out_p,
+      static_cast<uint16_t*>(values.data_ptr()),
+      static_cast<int*>(indices.data_ptr()),
+      static_cast<int*>(counts.data_ptr()), rescan, c.phase_bytes,
+      c.sample_target, c.sample_stride, c.sample_groups, c.sample_tiles,
+      static_cast<const float*>(scales.data_ptr()), tks, tq, ts, tk, tw,
+      ublocks, static_cast<uint32_t*>(counter.data_ptr()),
+      static_cast<int*>(defer.data_ptr()));
+  check_cuda(cudaGetLastError());
+}
diff --git a/csrc/fused_mqa_topk/persistent_epoch.inc b/csrc/fused_mqa_topk/persistent_epoch.inc
new file mode 100644
index 000000000000..1af6a1d2088b
--- /dev/null
+++ b/csrc/fused_mqa_topk/persistent_epoch.inc
@@ -0,0 +1,742 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+//
+// Body of one persistent epoch (included twice by persistent_kernel.cuh: the
+// first epoch is straight-line code so the hot loops are not nested in an outer
+// loop).
+if (warp_idx == 9 && cute::elect_one_sync()) {
+  if (epoch)
+    for (uint32_t i = 0; i < BarIdx::total; ++i) mbar_inval(bar_words + i);
+  for (uint32_t i = 0; i < BarIdx::total; ++i)
+    reinterpret_cast<volatile uint64_t*>(bar_words)[i] = 0;
+  full_q.init(1);
+  empty_q.init(kMathThreads + 32);  // math threads + the UMMA warp
+  for (uint32_t i = 0; i < kKVStages; ++i) {
+    full_kv[i].init(kFullKvArrivals);
+    empty_kv[i].init(kMathThreads);
+  }
+  for (uint32_t i = 0; i < BarIdx::kTmemSlots; ++i) {
+    full_tmem[i].init(1);
+    empty_tmem[i].init(kMathWgThreads);
+  }
+  full_lin->init(32);
+  // Every reader lane arrives on claim_empty.
+  for (uint32_t i = 0; i < 2; ++i) {
+    claim_full[i].init(32);
+    claim_empty[i].init(kClaimConsumers * 32);
+  }
+  sel_done.init(kMathWarps);  // one lane per math warp
+  for (uint32_t i = 0; i < 2 * kMathWarps; ++i) {
+    h_full[i].init(32);
+    h_empty[i].init(32);
+  }
+  cutlass::arch::fence_barrier_init();
+  unsafe_rows = 0;
+}
+__syncwarp();
+if (warp_idx == 10) cute::TMEM::Allocator1Sm().allocate(512, &tmem_ptr);
+__syncthreads();
+
+if (warp_idx == 8) {
+  // -------- claimer + Q/weights producer + drain of WG1 rows --------
+  cutlass::arch::warpgroup_reg_dealloc<kSpecRegs>();
+  // n_blk: blocks this CTA processed (all-short blocks excluded).
+  uint32_t n_blk = 0, drain_interval = 0;
+  for (uint32_t k = 0;; ++k) {
+    if (lane == 0) {
+      empty_q.wait((n_blk & 1) ^ 1);
+    }
+    __syncwarp();
+    // Every lane waits: lanes 0-7 write the entry.
+    claim_empty[k & 1].wait(((k >> 1) & 1) ^ 1);
+    uint32_t b = 0;
+    if (lane == 0)
+      b = atomicOr(&pause_req, 0u) ? kPause : atomicAdd(counter, 1u);
+    b = __shfl_sync(kFullMask, b, 0);
+    ClaimEntry& e = entries[k & 1];
+    if (b >= nblocks) {
+      if (lane == 0) {
+        const uint32_t kind = b == kPause ? kPause : kTerm;
+        e.vblock = kind;
+        epoch_kind = kind;
+      }
+      __syncwarp();
+      claim_full[k & 1].arrive();  // every writer lane arrives
+      break;
+    }
+    // Blocks are claimed heaviest (last rows, longest context) first.
+    const uint32_t vblock = nblocks - 1 - b, row_base = vblock * kBlockRows;
+    uint32_t st = 0, en = 0;
+    if (lane < kBlockRows) {
+      const uint32_t r = row_base + lane;
+      if (r < rows) {
+        st = min(starts_g[r], keys);
+        en = max(min(ends_g[r], keys), st);
+      }
+      e.starts[lane] = st;
+      e.ends[lane] = en;
+    }
+    const bool all_short =
+        __ballot_sync(kFullMask, lane < kBlockRows && en - st > topk) == 0;
+    const bool gated =
+        __ballot_sync(kFullMask,
+                      lane < kBlockRows && en - st > sample_target) != 0;
+    // KV range: rows past the end are clamped to the last valid row (DeepGEMM
+    // scheduler).
+    const uint32_t src = min(row_base + (lane & 7), rows - 1) - row_base;
+    uint32_t lo = __shfl_sync(kFullMask, st, src),
+             hi = __shfl_sync(kFullMask, en, src);
+#pragma unroll
+    for (int d = 4; d >= 1; d >>= 1) {
+      lo = min(lo, __shfl_xor_sync(kFullMask, lo, d));
+      hi = max(hi, __shfl_xor_sync(kFullMask, hi, d));
+    }
+    const uint32_t kv_base = lo / 4 * 4,
+                   splits = (hi - kv_base + kTileKeys - 1) / kTileKeys;
+    // CTA-private slot pair: block n of this CTA uses slot 2 * cta + (n & 1);
+    // at most two blocks (one in selection, one claimed ahead) are in flight,
+    // so no claim atomics.
+    const uint32_t slot = kSlotsPerCta * blockIdx.x + (n_blk & 1);
+    if (!all_short && lane < kBlockRows) row_counts[row_base + lane] = 0;
+    if (lane == 0) {
+      e.vblock = vblock;
+      e.all_short = all_short;
+      e.kv_base = kv_base;
+      e.splits = splits;
+      e.samples = gated ? sample_tiles : 0u;
+      e.slot = slot;
+    }
+    __syncwarp();
+    claim_full[k & 1].arrive();  // every writer lane arrives
+    if (all_short) continue;
+    if (lane == 0) {
+      dg::tma_copy<kHeadDim, kBlockRows * kNumHeads, 0>(&tq, &full_q, sm.q, 0,
+                                                        row_base * kNumHeads);
+      dg::tma_copy<kNumHeads, kBlockRows, 0>(&tw, &full_q, sm.weights, 0,
+                                             row_base);
+      full_q.arrive_and_expect_tx(kQBytes + kWBytes);
+    }
+    __syncwarp();
+    ++n_blk;
+    const int64_t row_delta = slot_row_delta(slot, row_base);
+    for (uint32_t base = 0; base < splits; base += Interval, ++drain_interval) {
+      for (uint32_t p = kMathWarps / 2; p < kMathWarps; ++p)
+        drain_handoff(hdata, h_full, h_empty, p, drain_interval, row_base + 4,
+                      row_counts, base + Interval >= splits,
+                      [&](uint32_t row, uint32_t slot2, uint32_t record,
+                          uint32_t batch_base, uint32_t source) {
+                        store_ring_record(vals, indices, row, row_delta, slot2,
+                                          record, kv_base, batch_base,
+                                          source % kMathWgThreads);
+                      });
+    }
+    cutlass::arch::NamedBarrier(kDrainSyncThreads, kBarDrain).arrive();
+  }
+} else if (warp_idx == 9) {
+  // -------- KV producer: sample tiles then scan tiles, one ring --------
+  cutlass::arch::warpgroup_reg_dealloc<kSpecRegs>();
+  RingPipeline<kKVStages> kv_pipe;
+  uint32_t n_blk = 0;
+  for (uint32_t k = 0;; ++k) {
+    ClaimHeader claim;
+    if (!consume_claim(k, claim_full, claim_empty, entries, claim)) break;
+    if (claim.all_short) continue;
+    const uint32_t kv_base = claim.kv_base, splits = claim.splits,
+                   samples = claim.samples;
+    {
+      // The whole warp stages the strided sample scales into sf_kv[stage]
+      // (prefetched one tile ahead in registers); lane 0 issues the KV TMA.
+      if (lane == 0 && n_blk > 0) sel_done.wait((n_blk - 1) & 1);
+      __syncwarp();
+      const uint32_t center = sample_stride / 2;
+      // Per-key scales of one sample tile, kScalesPerLane per lane.
+      constexpr uint32_t kScalesPerLane = kTileKeys / 32;
+      float pf[kScalesPerLane];
+#pragma unroll
+      for (uint32_t j = 0; j < kScalesPerLane; ++j) {
+        const uint32_t g = j * 32 + lane;
+        pf[j] = g < sample_groups ? __ldg(scales + g * sample_stride + center)
+                                  : 0.0f;
+      }
+#pragma unroll 1
+      for (uint32_t t = 0; t < samples; ++t) {
+        // Only lane 0's ring state is maintained across blocks (the scan loop
+        // runs on lane 0).
+        CUTE_TIE_DECL(kv_pipe.advance(), s0, ph);
+        const uint32_t s = __shfl_sync(kFullMask, s0, 0);
+        // Every writing lane acquires and releases the stage itself.
+        empty_kv[s].wait(__shfl_sync(kFullMask, ph, 0) ^ 1);
+#pragma unroll
+        for (uint32_t j = 0; j < kScalesPerLane; ++j)
+          sm.sf_kv[s][j * 32 + lane] = pf[j];
+        // Generic writes before later TMA writes of this stage.
+        cutlass::arch::fence_view_async_shared();
+        if (lane == 0) {
+          full_kv[s].expect_transaction(kKVBytes);
+          dg::tma_copy<kHeadDim, kTileKeys, 0, cutlass::float_e4m3_t, true>(
+              &tks, &full_kv[s], sm.kv[s], 0, t * kTileKeys, center);
+        }
+        full_kv[s].arrive();
+#pragma unroll
+        for (uint32_t j = 0; j < kScalesPerLane; ++j) {
+          const uint32_t g = (t + 1) * kTileKeys + j * 32 + lane;
+          pf[j] = g < sample_groups ? __ldg(scales + g * sample_stride + center)
+                                    : 0.0f;
+        }
+      }
+    }
+    if (lane == 0) {
+#pragma unroll 1
+      for (uint32_t t = 0; t < splits; ++t) {
+        CUTE_TIE_DECL(kv_pipe.advance(), s, ph);
+        empty_kv[s].wait(ph ^ 1);
+        const uint32_t off = kv_base + t * kTileKeys;
+        dg::tma_copy<kHeadDim, kTileKeys, 0>(&tk, &full_kv[s], sm.kv[s], 0,
+                                             off);
+        dg::tma_copy<kTileKeys, 1, 0>(&ts, &full_kv[s], sm.sf_kv[s], off, 0);
+        mbar_arrive_expect_tx(&full_kv[s], kKVBytes + kScaleBytes,
+                              kFullKvArrivals);
+      }
+    }
+    __syncwarp();
+    ++n_blk;
+  }
+} else if (warp_idx == 10) {
+  // -------- UMMA issuer --------
+  cutlass::arch::warpgroup_reg_dealloc<kSpecRegs>();
+  RingPipeline<1> q_pipe;
+  RingPipeline<kKVStages> kv_pipe;
+  RingPipeline<BarIdx::kTmemSlots> tmem_pipe;
+  auto instr_desc =
+      cute::UMMA::make_instr_desc<cutlass::float_e4m3_t, cutlass::float_e4m3_t,
+                                  float, 128, kUmmaN, cute::UMMA::Major::K,
+                                  cute::UMMA::Major::K>();
+  auto runtime_instr_desc = cute::UMMA::make_runtime_instr_desc(instr_desc);
+  for (uint32_t k = 0;; ++k) {
+    ClaimHeader claim;
+    if (!consume_claim(k, claim_full, claim_empty, entries, claim)) break;
+    if (claim.all_short) continue;
+    const uint32_t splits = claim.splits, samples = claim.samples;
+    CUTE_TIE_DECL(q_pipe.advance(), q_stage, q_phase);
+    (void)q_stage;
+    full_q.wait(q_phase);
+    full_lin->wait(q_phase);  // this block's linear-term rows (warp 11)
+    ptx::tcgen05_after_thread_sync();
+    const uint32_t total = samples + splits;
+#pragma unroll 1
+    for (uint32_t t = 0; t < total; ++t) {
+      CUTE_TIE_DECL(kv_pipe.advance(), s, ph);
+      full_kv[s].wait(ph);
+      if (cute::elect_one_sync()) {
+        // Per KV half, one N=144 MMA per math warpgroup into its own TMEM slot
+        // (ring of 3; WG0 takes even, WG1 odd positions of the issue sequence).
+#pragma unroll
+        for (uint32_t i = 0; i < 2; ++i) {
+#pragma unroll
+          for (uint32_t g = 0; g < 2; ++g) {
+            CUTE_TIE_DECL(tmem_pipe.advance(), ts_idx, ts_ph);
+            const uint32_t tmem_addr = ts_idx * kSlotCols;
+            empty_tmem[ts_idx].wait(ts_ph ^ 1);
+            ptx::tcgen05_after_thread_sync();
+            cutlass::float_e4m3_t* b_base =
+                g == 0 ? reinterpret_cast<cutlass::float_e4m3_t*>(sm.lin_lo)
+                       : sm.q + 4 * kNumHeads * kHeadDim;  // WG1's rows
+#pragma unroll
+            for (uint32_t kk = 0; kk < 4; ++kk) {
+              auto a_desc =
+                  mma::sm100::make_umma_desc<cute::UMMA::Major::K, 0, 128, 128>(
+                      sm.kv[s], i * 128, kk * 32);
+              auto b_desc =
+                  mma::sm100::make_umma_desc<cute::UMMA::Major::K, 0, 128, 128>(
+                      b_base, 0, kk * 32);
+              ptx::SM100_MMA_F8F6F4_SS::fma(a_desc, b_desc, tmem_addr, kk,
+                                            runtime_instr_desc);
+            }
+            ptx::umma_arrive_no_elect(full_tmem[ts_idx]);
+          }
+        }
+      }
+      __syncwarp();
+    }
+    empty_q.arrive();
+  }
+} else if (warp_idx == 11) {
+  // -------- linear-term rows + drain of WG0 rows --------
+  cutlass::arch::warpgroup_reg_dealloc<kSpecRegs>();
+  uint32_t drain_interval = 0, n_lin = 0;
+  for (uint32_t k = 0;; ++k) {
+    ClaimHeader claim;
+    if (!consume_claim(k, claim_full, claim_empty, entries, claim)) break;
+    if (claim.all_short) continue;
+    const uint32_t vblock = claim.vblock, kv_base = claim.kv_base,
+                   splits = claim.splits, slot = claim.slot;
+    // Q of block k lands while the math warps still select block k-1, so this
+    // is off the critical path except for a CTA's first block.
+    full_q.wait(n_lin & 1);
+    ++n_lin;
+    compute_lin_rows(sm, lane);
+    // Generic smem writes -> UMMA (async proxy).
+    cutlass::arch::fence_view_async_shared();
+    full_lin->arrive();  // every lane arrives (count 32)
+    const uint32_t row_base = vblock * kBlockRows;
+    const int64_t row_delta = slot_row_delta(slot, row_base);
+    for (uint32_t base = 0; base < splits; base += Interval, ++drain_interval) {
+      for (uint32_t p = 0; p < kMathWarps / 2; ++p)
+        drain_handoff(hdata, h_full, h_empty, p, drain_interval, row_base,
+                      row_counts, base + Interval >= splits,
+                      [&](uint32_t row, uint32_t slot2, uint32_t record,
+                          uint32_t batch_base, uint32_t source) {
+                        store_ring_record(vals, indices, row, row_delta, slot2,
+                                          record, kv_base, batch_base,
+                                          source % kMathWgThreads);
+                      });
+    }
+    cutlass::arch::NamedBarrier(kDrainSyncThreads, kBarDrain).arrive();
+  }
+} else {
+  // -------- math warps 0-7 --------
+  cutlass::arch::warpgroup_reg_alloc<kMathRegs>();
+  const uint32_t math_tid = threadIdx.x;
+  const uint32_t query_base = (warp_idx / 4) * kWgRows;  // 4 warps per WG
+  RingPipeline<1> q_pipe;
+  RingPipeline<kKVStages> kv_pipe;
+  RingPipeline<BarIdx::kTmemSlots> tmem_pipe;
+  // Each warpgroup consumes every second TMEM slot position (WG1 starts at 1).
+  constexpr uint32_t kTmemStep = 2;
+  if (warp_idx >= 4) tmem_pipe.advance(1);
+  const uint32_t heads_col = warp_idx < 4 ? 16u : 0u;
+  const uint32_t lin_col = warp_idx < 4 ? 0u : 128u;
+  uint32_t produce_interval = 0;
+  for (uint32_t k = 0;; ++k) {
+    claim_full[k & 1].wait((k >> 1) & 1);
+    const ClaimEntry& e = entries[k & 1];
+    const uint32_t vblock = e.vblock;
+    if (vblock >= kPause) break;
+    const uint32_t row_base = vblock * kBlockRows;
+    if (e.all_short) {
+      for (uint32_t r = 0; r < kBlockRows && row_base + r < rows; ++r)
+        write_all_keys(out + (uint64_t)(row_base + r) * topk, topk,
+                       e.ends[r] - e.starts[r], math_tid, kMathThreads);
+      __syncwarp();
+      claim_empty[k & 1].arrive();
+      continue;
+    }
+    // ---- block setup: Q, weights, sample histogram ----
+    {
+      CUTE_TIE_DECL(q_pipe.advance(), q_stage, q_phase);
+      (void)q_stage;
+      full_q.wait(q_phase);
+      full_lin->wait(q_phase);
+    }
+    static_assert(sizeof(SampleShared) % 16 == 0);
+    for (uint32_t j = math_tid; j < sizeof(SampleShared) / 16;
+         j += kMathThreads)
+      reinterpret_cast<uint4*>(sample_smem.hist)[j] = make_uint4(0, 0, 0, 0);
+    float weights[kWgRows][kNumHeads];
+    float row_s[kWgRows];  // pow2 row scale of the linear term
+    uint32_t seq_k_start[kWgRows], seq_k_end[kWgRows];
+#pragma unroll
+    for (uint32_t i = 0; i < kWgRows; ++i) {
+      seq_k_start[i] = e.starts[query_base + i];
+      seq_k_end[i] = e.ends[query_base + i];
+      // Weights in units of 1 / s_r (exact: s_r is a power of two).
+      row_s[i] = sm.lin_scale[query_base + i];
+      const float inv_s = 1.0f / row_s[i];
+#pragma unroll
+      for (uint32_t j = 0; j < kNumHeads; j += 4) {
+        const float4 w4 = *reinterpret_cast<const float4*>(
+            sm.weights + (query_base + i) * kNumHeads + j);
+        weights[i][j] = w4.x * 0.5f * inv_s;
+        weights[i][j + 1] = w4.y * 0.5f * inv_s;
+        weights[i][j + 2] = w4.z * 0.5f * inv_s;
+        weights[i][j + 3] = w4.w * 0.5f * inv_s;
+      }
+    }
+    {
+      // Weights that are subnormal, or become subnormal or infinite once
+      // scaled, take the exact recovery path.
+      const float w = sm.weights[warp_idx * kNumHeads + lane];  // lane = head
+      // 2^-125: the smallest magnitude whose half is still a normal float.
+      constexpr float kMinScalable = 0x1p-125f;
+      bool bad_w = w != 0.0f && fabsf(w) < kMinScalable;
+      const float wr = w * 0.5f * (1.0f / sm.lin_scale[warp_idx]);
+      bad_w |= (wr != 0.0f && fabsf(wr) < kMinScalable) || isinf(wr);
+      const uint32_t bad = __ballot_sync(kFullMask, bad_w);
+      if (lane == 0 && bad) atomicOr(&unsafe_rows, 1u << warp_idx);
+    }
+    math_wg_sync();
+    // ---- sample tiles ----
+    // The first tile's per-row max is the histogram reference (see
+    // SampleShared): tile 0 keeps its scores until the reference is known.
+    float ref[kWgRows] = {};
+#pragma unroll 1
+    for (uint32_t t = 0; t < e.samples; ++t) {
+      CUTE_TIE_DECL(kv_pipe.advance(), kv_stage, kv_phase);
+      full_kv[kv_stage].wait(kv_phase);
+      uint32_t first_code[kWgRows] = {};
+#pragma unroll
+      for (uint32_t kv_half = 0; kv_half < 2; ++kv_half) {
+        const uint32_t kv_lane =
+            (math_tid % kMathWgThreads) + kv_half * kMathWgThreads;
+        const uint32_t group = t * kTileKeys + kv_lane;
+        const uint32_t kv_offset = group * sample_stride + sample_stride / 2;
+        const bool in_view = group < sample_groups;
+        // 0 outside the view (staged by the KV producer).
+        const float scale_kv = ptx::ld_shared(sm.sf_kv[kv_stage] + kv_lane);
+        CUTE_TIE_DECL(tmem_pipe.advance(kTmemStep), tmem_stage, tmem_phase);
+        full_tmem[tmem_stage].wait(tmem_phase);
+        ptx::tcgen05_after_thread_sync();
+        if (kv_half == 1) {
+          cutlass::arch::fence_view_async_shared();
+          empty_kv[kv_stage].arrive();
+        }
+        // One query row (32 TMEM columns) at a time keeps the sample loop's
+        // live registers low enough to carry the scan weights through it.
+#pragma unroll
+        for (uint32_t i = 0; i < kWgRows; ++i) {
+          float acc[kNumHeads];
+          float lin[4];
+          ptx::tmem_load_32dp32b<32>(
+              tmem_stage * kSlotCols + heads_col + i * kNumHeads,
+              reinterpret_cast<uint32_t*>(acc));
+          ptx::tmem_load_32dp32b<4>(tmem_stage * kSlotCols + lin_col + i * 4,
+                                    reinterpret_cast<uint32_t*>(lin));
+          cutlass::arch::fence_view_async_tmem_load();
+          if (i == kWgRows - 1) {
+            ptx::tcgen05_before_thread_sync();
+            empty_tmem[tmem_stage].arrive();
+          }
+          const float bq = reduce_heads_with_lin(acc, weights[i], lin) *
+                           (row_s[i] * scale_kv);
+          const bool valid = in_view && kv_offset >= seq_k_start[i] &&
+                             kv_offset < seq_k_end[i] && isfinite(bq);
+          if (t == 0) {
+            sample_smem.first[kv_half][i][math_tid] =
+                valid ? bq : -CUDART_INF_F;
+            if (valid) first_code[i] = max(first_code[i], ordered_code(bq));
+          } else if (valid)
+            add_sample(sample_smem.hist[query_base + i], bq - ref[i]);
+        }
+      }
+      if (t == 0) {
+#pragma unroll
+        for (uint32_t i = 0; i < kWgRows; ++i) {
+          const uint32_t code = __reduce_max_sync(kFullMask, first_code[i]);
+          if (lane == 0) atomicMax(&sample_smem.ref_code[query_base + i], code);
+        }
+        math_wg_sync();
+#pragma unroll
+        for (uint32_t i = 0; i < kWgRows; ++i) {
+          ref[i] = sample_ref(sample_smem.ref_code[query_base + i]);
+#pragma unroll
+          for (uint32_t h = 0; h < 2; ++h) {
+            const float x = sample_smem.first[h][i][math_tid];
+            if (isfinite(x))
+              add_sample(sample_smem.hist[query_base + i], x - ref[i]);
+          }
+        }
+      }
+    }
+    math_wg_sync();
+    finish_sample_histogram(
+        sample_smem.hist[warp_idx], sample_ref(sample_smem.ref_code[warp_idx]),
+        gate_s + warp_idx, e.ends[warp_idx] - e.starts[warp_idx], sample_target,
+        lane);
+    math_wg_sync();
+    // The block index is re-read from the claim entry (volatile) instead of
+    // kept live across the sample loop.
+    const uint32_t scan_row_base = volatile_load(e.vblock) * kBlockRows;
+
+    // ---- scan tiles ----
+    const uint32_t kv_base = e.kv_base, splits = e.splits;
+    const int64_t cand_row_delta = slot_row_delta(e.slot, scan_row_base);
+    float gate_score[kWgRows];
+#pragma unroll
+    for (uint32_t i = 0; i < kWgRows; ++i) {
+      gate_score[i] = gate_s[query_base + i];
+      if ((volatile_load(unsafe_rows) >> (query_base + i)) & 1u) {
+        // Unsafe weights: no key passes (+inf gate), and a count above the
+        // candidate cap sends the block to the exact rescan.
+        gate_score[i] = CUDART_INF_F;
+        if (math_tid % kMathWgThreads == 0)
+          atomicAdd(row_counts + scan_row_base + query_base + i,
+                    static_cast<int>(kCandidatesPerRow) + 1);
+      }
+      // Compare S * scale_kv against gate / row_s (row_s is a power of two ->
+      // exact).
+      gate_score[i] *= 1.0f / row_s[i];
+    }
+    uint32_t lane_counts[kWgRows] = {};
+    uint32_t flush_base = 0;
+#pragma unroll 1
+    for (uint32_t kv_split_idx = 0; kv_split_idx < splits; ++kv_split_idx) {
+      if ((kv_split_idx & (Interval - 1)) == 0) {
+        h_empty[(produce_interval & 1) * kMathWarps + warp_idx].wait(
+            ((produce_interval >> 1) & 1) ^ 1);
+        __syncwarp();
+      }
+      CUTE_TIE_DECL(kv_pipe.advance(), kv_stage, kv_phase);
+      full_kv[kv_stage].wait(kv_phase);
+#pragma unroll
+      for (uint32_t kv_half = 0; kv_half < 2; ++kv_half) {
+        const uint32_t kv_lane =
+            (math_tid % kMathWgThreads) + kv_half * kMathWgThreads;
+        const uint32_t kv_offset = kv_base + kv_split_idx * kTileKeys + kv_lane;
+        const float scale_kv = ptx::ld_shared(sm.sf_kv[kv_stage] + kv_lane);
+        CUTE_TIE_DECL(tmem_pipe.advance(kTmemStep), tmem_stage, tmem_phase);
+        full_tmem[tmem_stage].wait(tmem_phase);
+        ptx::tcgen05_after_thread_sync();
+        if (kv_half == 1) {
+          cutlass::arch::fence_view_async_shared();
+          empty_kv[kv_stage].arrive();
+        }
+        float accum_pair[2][kNumHeads];
+        float lin_pair[2][4];
+        float
+            bp4[kWgRows];  // row scores of the slot, in units of the row scale
+#pragma unroll
+        for (uint32_t pair = 0; pair < kWgRows; pair += 2) {
+          const uint32_t slot_base = tmem_stage * kSlotCols;
+          ptx::tmem_load_32dp32b<2 * kNumHeads>(
+              slot_base + heads_col + pair * kNumHeads,
+              reinterpret_cast<uint32_t*>(accum_pair));
+          ptx::tmem_load_32dp32b<8>(slot_base + lin_col + pair * 4,
+                                    reinterpret_cast<uint32_t*>(lin_pair));
+          cutlass::arch::fence_view_async_tmem_load();
+          if (pair + 2 >= kWgRows) {
+            ptx::tcgen05_before_thread_sync();
+            empty_tmem[tmem_stage].arrive();
+          }
+          // The linear term (sum_t term_t . k, in units of the row scale; terms
+          // 2, 3 in units of s * 2^-8) seeds the sums.
+#pragma unroll
+          for (uint32_t r = 0; r < 2; ++r)
+            bp4[pair + r] = reduce_heads_with_lin(
+                                accum_pair[r], weights[pair + r], lin_pair[r]) *
+                            scale_kv;
+        }
+        // One warp vote per slot: most keys fail every gate, so the common path
+        // is 4 compares + 1 vote + 1 uniform branch.
+        const bool any_pass =
+            (bp4[0] > gate_score[0]) | (bp4[1] > gate_score[1]) |
+            (bp4[2] > gate_score[2]) | (bp4[3] > gate_score[3]);
+        if (__any_sync(kFullMask, any_pass)) [[unlikely]] {
+#pragma unroll
+          for (uint32_t i = 0; i < kWgRows; ++i) {
+            const float bp = bp4[i];
+            // Gate first; range and finiteness only for passing keys.
+            if (bp > gate_score[i]) {
+              const float bq = bp * row_s[i];
+              const uint32_t row = scan_row_base + query_base + i;
+              if (kv_offset >= seq_k_start[i] && kv_offset < seq_k_end[i] &&
+                  isfinite(bq)) {
+                const uint32_t count = lane_counts[i];
+                if (count < kRingSlots) {
+                  const uint32_t code = candidate_fp24_code(bq);
+                  hdata
+                      .records[produce_interval & 1][warp_idx][i][count][lane] =
+                      (code << 8) | ((kv_split_idx - flush_base) * 2 + kv_half);
+                  lane_counts[i] = count + 1;
+                } else {
+                  const int slot = atomicAdd(row_counts + row, 1);
+                  if (slot < static_cast<int>(kCandidatesPerRow)) {
+                    const uint64_t off =
+                        candidate_offset(row, cand_row_delta, slot);
+                    store_candidate(vals + off, indices + off, bq, kv_offset);
+                  }
+                }
+              }
+            }
+          }
+        }
+      }
+      if (((kv_split_idx + 1) & (Interval - 1)) == 0 &&
+          kv_split_idx + 1 != splits) {
+#pragma unroll
+        for (uint32_t i = 0; i < kWgRows; ++i) {
+          hdata.num_records[produce_interval & 1][warp_idx][i][lane] =
+              lane_counts[i];
+          lane_counts[i] = 0;
+        }
+        if (lane == 0) hdata.bases[produce_interval & 1][warp_idx] = flush_base;
+        __syncwarp();
+        h_full[(produce_interval & 1) * kMathWarps + warp_idx].arrive();
+        ++produce_interval;
+        flush_base = kv_split_idx + 1;
+      }
+    }
+    {
+      // Final interval: every math warp stores its own ring records (8 warps in
+      // parallel instead of 2 drain warps after the scan); the drain warps only
+      // cycle the barriers.
+      const uint32_t bank = produce_interval & 1,
+                     rbase = scan_row_base + query_base;
+      static_assert(sizeof(hdata.records[0][0]) ==
+                        kWgRows * kRingSlots * 32 * sizeof(uint32_t),
+                    "ring record layout");
+      self_drain_last(&hdata.records[bank][warp_idx][0][0][0], lane_counts[0],
+                      lane_counts[1], lane_counts[2], lane_counts[3],
+                      row_counts, rbase, kv_base + (math_tid % kMathWgThreads),
+                      flush_base, cand_row_delta, vals, indices);
+#pragma unroll
+      for (uint32_t row = 0; row < kWgRows; ++row) lane_counts[row] = 0;
+      __syncwarp();
+      h_full[bank * kMathWarps + warp_idx].arrive();
+      ++produce_interval;
+    }
+    cutlass::arch::fence_view_async_shared();
+    empty_q.arrive();
+    // Drains of this block done.
+    cutlass::arch::NamedBarrier(kDrainSyncThreads, kBarDrain).sync();
+    // ---- selection ----
+    // Re-read as well (not kept live across the scan).
+    const uint32_t sel_vblock = volatile_load(e.vblock),
+                   sel_row_base = sel_vblock * kBlockRows;
+    __threadfence_block();
+    cutlass::arch::NamedBarrier(kMathThreads, kBarSelect).sync();
+    const ClaimEntry& sel_entry = entries[k & 1];
+    if (math_tid == 0) {
+      uint32_t recover = 0, short_rows = 0;
+      row_offset[0] = 0;
+      for (uint32_t r = 0; r < kBlockRows; ++r) {
+        const bool is_short = sel_entry.ends[r] - sel_entry.starts[r] <= topk;
+        short_rows |= uint32_t(is_short) << r;
+        const int c = __ldcg(row_counts + sel_row_base + r);
+        recover |= !is_short && (c < int(topk) || c > int(kCandidatesPerRow));
+        // Rows start 16-byte aligned for select_staged_row's vector loads.
+        row_offset[r + 1] =
+            row_offset[r] + (is_short ? 0 : (uint32_t(max(c, 0)) + 3u) & ~3u);
+      }
+      // 32-bit byte arithmetic (sizeof is a 64-bit size_t).
+      constexpr uint32_t kWordBytes = sizeof(uint32_t);
+      // Stage all 8 rows' candidates at once (1), in two rounds of 4 (2), or,
+      // when they exceed the scratch, rescan exactly.
+      const uint32_t half = max(row_offset[kWgRows],
+                                row_offset[kBlockRows] - row_offset[kWgRows]);
+      const uint32_t staged =
+          recover ? 0u
+                  : ((kBlockRows * selection::kBins + row_offset[kBlockRows]) *
+                                 kWordBytes <=
+                             phase_bytes
+                         ? 1u
+                         : ((kWgRows * selection::kBins + half) * kWordBytes <=
+                                    phase_bytes
+                                ? 2u
+                                : 0u));
+      if (staged == 0) recover = 1;
+      if (recover) {
+        defer_next[sel_vblock] = int(defer_head);
+        defer_head = sel_vblock;
+        atomicExch(&pause_req, 1u);
+      }
+      sel_recover = recover;
+      sel_staged = staged;
+      sel_short = short_rows;
+      unsafe_rows = 0;
+    }
+    cutlass::arch::NamedBarrier(kMathThreads, kBarSelect).sync();
+    const uint32_t recover = sel_recover, short_rows = sel_short;
+    const auto write_short = [&](uint32_t r, uint32_t t, uint32_t stride) {
+      if (r >= rows) return;
+      write_all_keys(
+          out + (uint64_t)r * topk, topk,
+          sel_entry.ends[r - sel_row_base] - sel_entry.starts[r - sel_row_base],
+          t, stride);
+    };
+    if (!recover) {
+      const int64_t row_delta = slot_row_delta(sel_entry.slot, sel_row_base);
+      const uint32_t per_round = sel_staged == 1 ? kBlockRows : kWgRows;
+      for (uint32_t round = 0; round < kBlockRows / per_round; ++round) {
+        if (warp_idx < per_round) {
+          const uint32_t local = round * per_round + warp_idx,
+                         first = round * per_round;
+          uint32_t* hist = reinterpret_cast<uint32_t*>(smem_buffer + kSelBase) +
+                           warp_idx * selection::kBins;
+          uint32_t* codes =
+              reinterpret_cast<uint32_t*>(smem_buffer + kSelBase) +
+              per_round * selection::kBins +
+              (row_offset[local] - row_offset[first]);
+          const uint32_t r = sel_row_base + local;
+          const uint64_t base = candidate_offset(r, row_delta, 0);
+          if ((short_rows >> local) & 1u)
+            write_short(r, lane, 32);
+          else {
+            const uint32_t cnt = __ldcg(row_counts + r);
+            int* o = out + (uint64_t)r * topk;
+            if (!selection::select_staged_row(codes, cnt, vals + base,
+                                              indices + base, o, topk, hist))
+              selection::select_row(codes, cnt, vals + base, indices + base, o,
+                                    topk, hist);
+          }
+        }
+        cutlass::arch::NamedBarrier(kMathThreads, kBarSelect).sync();
+      }
+    }
+    // Every row is final here: the last selection round ends with a barrier,
+    // and a recovered block skipped the rounds right after the barrier above.
+    // Short rows were written relative already.
+    if (!recover)
+      make_relative(out, topk, rows, sel_row_base, starts_g, math_tid,
+                    kMathThreads, [&](uint32_t local) {
+                      return sel_entry.ends[local] - sel_entry.starts[local] <=
+                             topk;
+                    });
+    // The selection scratch aliases kv[] and sf_kv[]. Order its generic
+    // accesses before the next block's TMA writes, which have no proxy fence
+    // of their own when that block has no sample tiles.
+    cutlass::arch::fence_view_async_shared();
+    __syncwarp();
+    if (lane == 0) sel_done.arrive();
+    claim_empty[k & 1].arrive();
+  }
+}
+
+// -------- exact rescan of the blocks queued by the selection --------
+// The rescan's first TMA loads overwrite the epoch's generic shared memory.
+cutlass::arch::fence_view_async_shared();
+__syncthreads();
+if (warp_idx == 0) cute::TMEM::Allocator1Sm().free(0, 512);
+if (defer_head != kTerm) {
+  // The rescan expects the 224 / 56 register split of the math warps.
+  __syncthreads();
+  for (uint32_t cur = defer_head; cur != kTerm;) {
+    const uint32_t vblock = cur, row_base = vblock * kBlockRows;
+    if (threadIdx.x < kBlockRows) {
+      const uint32_t r = row_base + threadIdx.x;
+      row_bounds[0][threadIdx.x] = r < rows ? min(starts_g[r], keys) : 0u;
+      row_bounds[1][threadIdx.x] =
+          r < rows ? max(min(ends_g[r], keys), row_bounds[0][threadIdx.x]) : 0u;
+    }
+    __syncthreads();
+    const uint32_t* const starts = row_bounds[0] - row_base;
+    const uint32_t* const ends = row_bounds[1] - row_base;
+    // No block is in flight during the rescan: use the CTA's first slot.
+    const int64_t row_delta =
+        slot_row_delta(kSlotsPerCta * blockIdx.x, row_base);
+    const RescanBuffers block_buffers{
+        rescan.retained + row_delta * kRetainedPerRow,
+        rescan.retained_indices + row_delta * kRetainedPerRow};
+    rescan::rescan_block(dense::Scheduler(vblock, rows, keys, starts, ends),
+                         block_buffers, out, topk, tq, tk, ts, tw);
+    __syncthreads();
+    if (threadIdx.x == 0)
+      dense::invalidate_barriers(
+          *reinterpret_cast<dense::SharedStorage*>(smem_buffer));
+    __syncthreads();
+    // The rescan writes absolute indices for every row of the block.
+    make_relative(out, topk, rows, row_base, starts_g, threadIdx.x, kNumThreads,
+                  [](uint32_t) { return false; });
+    cur = uint32_t(__ldcg(defer_next + vblock));
+    // Order the generic mbarrier.inval above before the TMA writes of the
+    // next rescan block or the next epoch.
+    cutlass::arch::fence_view_async_shared();
+    __syncthreads();
+  }
+  // The loop ran at least once and ended with a barrier.
+  if (threadIdx.x == 0) {
+    defer_head = kTerm;
+    pause_req = 0;
+  }
+}  // deferred blocks
+more = epoch_kind != kTerm;
+__syncthreads();
diff --git a/csrc/fused_mqa_topk/persistent_kernel.cuh b/csrc/fused_mqa_topk/persistent_kernel.cuh
new file mode 100644
index 000000000000..86063b08fbac
--- /dev/null
+++ b/csrc/fused_mqa_topk/persistent_kernel.cuh
@@ -0,0 +1,592 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+//
+// Persistent fused FP8 MQA logits + per-row top-K (K = 512, 1024, 2048; SM100),
+// DeepGEMM/FA4-style: one CTA per SM, barriers and TMEM initialised once, every
+// warp role runs its own loop over blocks claimed (heaviest first) from a
+// global counter, and the sample and scan tiles of a block flow through ONE
+// continuous Q/KV/TMEM pipeline whose stage/phase state is carried across
+// blocks.
+//
+// Roles (384 threads):
+//   warp 8   claimer + Q/weights TMA + drain of math WG1 rows
+//   warp 9   KV TMA (sample tiles via the strided map, then the scan tiles)
+//   warp 10  UMMA issuer
+//   warp 11  linear-term rows + drain of math WG0 rows
+//   warps 0-7 math: sample histogram -> gate -> scan/emit -> selection
+// A block is rescanned exactly (rescan.cuh) when a row's candidate count
+// leaves [K, kCandidatesPerRow], a row's weights are unsafe to scale, or the
+// block's candidates exceed the selection scratch (~11K+ per row): the role
+// loops pause, the CTA rescans its queued blocks, and the loops resume with
+// freshly initialised barriers.
+#pragma once
+#include <math_constants.h>
+
+#include <cstddef>
+#include <cuda_fp8.h>
+#include <cutlass/arch/reg_reconfig.h>
+#include <cute/arch/cluster_sm90.hpp>
+#include <deep_gemm/common/cute_tie.cuh>
+#include <deep_gemm/common/ring_pipeline.cuh>
+#include <deep_gemm/common/utils.cuh>
+#include <deep_gemm/mma/sm100.cuh>
+#include <deep_gemm/ptx/ld_st.cuh>
+#include <deep_gemm/ptx/tcgen05.cuh>
+#include <deep_gemm/ptx/utils.cuh>
+
+#include "candidates.cuh"
+#include "common.cuh"
+#include "constants.h"
+#include "deep_gemm_compat.cuh"
+#include "rescan.cuh"
+#include "select.cuh"
+
+namespace fused_mqa_topk::persist {
+using namespace deep_gemm;
+using rescan::RescanBuffers;
+using rescan::RescanShared;
+using TBar = cutlass::arch::ClusterTransactionBarrier;
+
+constexpr uint32_t kTerm = 0xffffffffu, kPause = 0xfffffffeu;
+// Bytes of a block's Q, of one KV tile and its per-key scales, of a block's
+// weights.
+constexpr uint32_t kQBytes = kBlockRows * kNumHeads * kHeadDim,
+                   kKVBytes = kTileKeys * kHeadDim,
+                   kScaleBytes = kTileKeys * sizeof(float),
+                   kWBytes = kBlockRows * kNumHeads * sizeof(float);
+// KV (and per-key scale) pipeline stages.
+constexpr uint32_t kKVStages = 3;
+// Linear-term B rows per math warpgroup: 4 query rows x 4 FP8 terms.
+constexpr uint32_t kLinRows = 16;
+// Terms 2 and 3 of the linear term are stored in units of s * 2^-8.
+constexpr float kLinLoScale = 0x1p-8f;
+// TMEM: one UMMA per math warpgroup and KV half writes 4 rows x 32 heads plus
+// the kLinRows linear-term columns (N = 144). Slots are 160 columns apart (144
+// rounded up to a multiple of 32); the 3 slots take 480 of the 512 allocated
+// columns.
+constexpr uint32_t kUmmaN = 4 * kNumHeads + kLinRows, kSlotCols = 160;
+static_assert(kSlotCols >= kUmmaN && kSlotCols % 32 == 0);
+
+// Rebalanced scan: the linear-term rows of each math warpgroup bracket Q, so
+// WG0's N=144 UMMA reads B = [lin WG0 | Q rows 0..127] (D = [lin | heads]) and
+// WG1's reads B = [Q rows 128..255 | lin WG1] (D = [heads | lin]).
+struct KernelSmem {
+  alignas(1024) uint8_t lin_lo[kLinRows * kHeadDim];
+  cutlass::float_e4m3_t q[kQBytes];
+  uint8_t lin_hi[kLinRows * kHeadDim];
+  alignas(128) float weights[kBlockRows * kNumHeads];
+  float lin_scale[kBlockRows];
+  alignas(1024) cutlass::float_e4m3_t kv[kKVStages][kKVBytes];
+  alignas(128) float sf_kv[kKVStages][kTileKeys];
+};
+static_assert(offsetof(KernelSmem, q) == kLinRows * kHeadDim &&
+                  offsetof(KernelSmem, lin_hi) == kLinRows * kHeadDim + kQBytes,
+              "lin rows must bracket Q");
+// Selection scratch starts after lin + Q + lin + weights + row scales, so the
+// next block's Q and linear term land during this block's selection.
+constexpr uint32_t kSelBase = offsetof(KernelSmem, kv);
+
+// Byte offset of byte `col` of row `row` in a 128B-swizzled K-major tile of
+// 128-byte rows (the UMMA B layout of Q and the linear-term rows).
+__device__ __forceinline__ uint32_t swizzle128_offset(uint32_t row,
+                                                      uint32_t col) {
+  return row * 128u + (((col >> 4) ^ (row & 7u)) << 4) + (col & 15u);
+}
+
+// q'_r = sum_h (w_rh / 2) q_rh, split into 4 FP8 e4m3 terms: terms 0,1 scale
+// s_r, terms 2,3 s_r * 2^-8 (lin_scale[r]); term t of row r is B row (r % 4) *
+// 4 + t of lin_lo (rows 0-3) / lin_hi (4-7), 128B-swizzled K-major like Q. One
+// warp; lane l owns dims 4l..4l+3.
+__device__ __forceinline__ void compute_lin_rows(KernelSmem& sm,
+                                                 uint32_t lane) {
+  const uint8_t* qbytes = reinterpret_cast<const uint8_t*>(sm.q);
+#pragma unroll 1
+  for (uint32_t r = 0; r < kBlockRows; ++r) {
+    float acc[4] = {0.f, 0.f, 0.f, 0.f};
+#pragma unroll 8
+    for (uint32_t h = 0; h < kNumHeads; ++h) {
+      const uint32_t off = swizzle128_offset(r * kNumHeads + h, lane * 4u);
+      const uint32_t packed =
+          ptx::ld_shared(reinterpret_cast<const uint32_t*>(qbytes + off));
+      const float wv = ptx::ld_shared(sm.weights + r * kNumHeads + h) * 0.5f;
+      const __half2_raw lo = __nv_cvt_fp8x2_to_halfraw2(
+          static_cast<__nv_fp8x2_storage_t>(packed & 0xffffu), __NV_E4M3);
+      const __half2_raw hi = __nv_cvt_fp8x2_to_halfraw2(
+          static_cast<__nv_fp8x2_storage_t>(packed >> 16), __NV_E4M3);
+      const float2 f0 = __half22float2(*reinterpret_cast<const __half2*>(&lo));
+      const float2 f1 = __half22float2(*reinterpret_cast<const __half2*>(&hi));
+      acc[0] = fmaf(wv, f0.x, acc[0]);
+      acc[1] = fmaf(wv, f0.y, acc[1]);
+      acc[2] = fmaf(wv, f1.x, acc[2]);
+      acc[3] = fmaf(wv, f1.y, acc[3]);
+    }
+    float m = fmaxf(fmaxf(fabsf(acc[0]), fabsf(acc[1])),
+                    fmaxf(fabsf(acc[2]), fabsf(acc[3])));
+#pragma unroll
+    for (uint32_t o = 16; o > 0; o >>= 1)
+      m = fmaxf(m, __shfl_xor_sync(kFullMask, m, o));
+    int e = -120;  // pow2 scale s >= m / 448, clamped so 1/s stays finite
+    if (m > 0.f && isfinite(m)) {
+      frexpf(m * (1.0f / 448.0f), &e);
+      e = max(e, -120);
+    }
+    const float s = ldexpf(1.0f, e), inv_s = ldexpf(1.0f, -e);
+    float res[4];
+#pragma unroll
+    for (uint32_t dd = 0; dd < 4; ++dd) res[dd] = acc[dd] * inv_s;
+    uint8_t* lin_base = r < kWgRows ? sm.lin_lo : sm.lin_hi;
+#pragma unroll
+    for (uint32_t tt = 0; tt < 4; ++tt) {
+      const __nv_fp8x2_storage_t c01 = __nv_cvt_float2_to_fp8x2(
+          make_float2(res[0], res[1]), __NV_SATFINITE, __NV_E4M3);
+      const __nv_fp8x2_storage_t c23 = __nv_cvt_float2_to_fp8x2(
+          make_float2(res[2], res[3]), __NV_SATFINITE, __NV_E4M3);
+      const __half2_raw d01 = __nv_cvt_fp8x2_to_halfraw2(c01, __NV_E4M3);
+      const __half2_raw d23 = __nv_cvt_fp8x2_to_halfraw2(c23, __NV_E4M3);
+      const float2 g01 =
+          __half22float2(*reinterpret_cast<const __half2*>(&d01));
+      const float2 g23 =
+          __half22float2(*reinterpret_cast<const __half2*>(&d23));
+      res[0] -= g01.x;
+      res[1] -= g01.y;
+      res[2] -= g23.x;
+      res[3] -= g23.y;
+      // Terms 2, 3 in units of s * 2^-8: the residual after two terms is <=
+      // 0.5, so x256 <= 128 (no saturation), and the e4m3 subnormal floor drops
+      // from s * 2^-10 to s * 2^-18 per dim.
+      if (tt == 1) {
+        res[0] *= 1.0f / kLinLoScale;
+        res[1] *= 1.0f / kLinLoScale;
+        res[2] *= 1.0f / kLinLoScale;
+        res[3] *= 1.0f / kLinLoScale;
+      }
+      const uint32_t off =
+          swizzle128_offset((r % kWgRows) * 4u + tt, lane * 4u);
+      ptx::st_shared(
+          reinterpret_cast<uint32_t*>(lin_base + off),
+          static_cast<uint32_t>(c01) | (static_cast<uint32_t>(c23) << 16));
+    }
+    if (lane == 0) sm.lin_scale[r] = s;
+  }
+}
+
+// One row's weighted head reduction seeded with its linear-term columns
+// (terms 2, 3 in units of s * 2^-8).
+__device__ __forceinline__ float reduce_heads_with_lin(const float* s,
+                                                       const float* w,
+                                                       const float* lin) {
+  return reduce_heads<false>(s, w, make_float2(lin[0], lin[1]),
+                             __fmul2_rn(make_float2(lin[2], lin[3]),
+                                        make_float2(kLinLoScale, kLinLoScale)));
+}
+
+// Gate sample state: each math warp histograms its row's sampled scores
+// relative to the row's reference `ref` (its max over the first sample tile):
+// half_bin<5>(x - ref), kSampleBins bins, 32 per octave of the distance, so
+// the bins near the row's top stay fine whatever the scores' common offset
+// (absolute bins are ~2 wide at |x| ~ 100, too coarse for top tails that
+// narrow).
+constexpr uint32_t kSampleShift = 5, kSampleBins = 65536 >> kSampleShift;
+struct alignas(16) SampleShared {
+  uint32_t hist[kBlockRows][padded_bins(kSampleBins)];
+  // Ordered code of each row's max over the first sample tile (0: none).
+  uint32_t ref_code[kBlockRows];
+  // The first tile's scores, held until the reference is known (in registers
+  // they cost the hot loops spills).
+  float first[2][kWgRows][kMathThreads];
+};
+
+__device__ __forceinline__ void add_sample(uint32_t* hist, float distance) {
+  atomicAdd(hist + padded_bin(half_bin<kSampleShift>(distance)), 1u);
+}
+__device__ __forceinline__ uint32_t sample_count(const uint32_t* hist,
+                                                 uint32_t bin) {
+  return hist[padded_bin(bin)];
+}
+// The row reference from its ordered code; 0 when the row had no valid sample
+// in the first tile (then the bins are absolute).
+__device__ __forceinline__ float sample_ref(uint32_t code) {
+  const float ref = ordered_float(code);
+  return isfinite(ref) ? ref : 0.0f;
+}
+
+// The gate aims kGateSlack16 / 16 above the target count, so that few rows fall
+// below the top-k (and into the exact rescan) when the sampled estimate runs
+// low. On the B300 captures 17/16 cuts the rescanned blocks by ~85% at about
+// the candidate counts of the coarser absolute bins.
+constexpr uint32_t kGateSlack16 = 17;
+
+// Gate: the k-th largest sample (k = count * target / visible, count = the
+// histogram total) of one row from its histogram; one warp per row. -inf when
+// the row sees at most `target` keys or has no samples.
+__device__ __noinline__ inline void finish_sample_histogram(
+    const uint32_t* hist, float ref, float* gate, uint32_t visible,
+    uint32_t target, uint32_t lane) {
+  if (visible <= target) {
+    if (lane == 0) *gate = -CUDART_INF_F;
+    __syncwarp();
+    return;
+  }
+  // Lane l owns kChunksPerLane 32-bin chunks, descending with the lane: chunk
+  // kChunks - 1 - kChunksPerLane * l - j for j < kChunksPerLane.
+  constexpr uint32_t kChunks = kSampleBins / 32, kChunksPerLane = kChunks / 32;
+  uint32_t sum = 0;
+#pragma unroll
+  for (uint32_t j = 0; j < kChunksPerLane; ++j) {
+    const uint32_t c = kChunks - 1 - kChunksPerLane * lane - j;
+#pragma unroll 8
+    for (uint32_t b = 0; b < 32; ++b) sum += hist[padded_bin(c * 32 + b)];
+  }
+  const uint32_t incl = warp_inclusive_scan(sum, lane);
+  const uint32_t tot = __shfl_sync(kFullMask, incl, 31);
+  if (tot == 0) {
+    if (lane == 0) *gate = -CUDART_INF_F;
+    __syncwarp();
+    return;
+  }
+  const uint32_t rank = max(
+      1u, min(tot, (tot * target * kGateSlack16 / 16 + visible - 1) / visible));
+  const uint32_t w = __ffs(__ballot_sync(kFullMask, incl >= rank)) - 1;
+  uint32_t above = __shfl_sync(kFullMask, incl - sum, w), chosen = 0;
+#pragma unroll 1
+  for (uint32_t cc = 0; cc < kChunksPerLane; ++cc) {
+    const uint32_t c = kChunks - 1 - kChunksPerLane * w - cc;
+    // Lane 0 = highest bin of the chunk.
+    const uint32_t v =
+        warp_inclusive_scan(sample_count(hist, c * 32 + 31 - lane), lane);
+    const uint32_t mask = __ballot_sync(kFullMask, above + v >= rank);
+    if (mask) {
+      chosen = c * 32 + 31 - (__ffs(mask) - 1);
+      break;
+    }
+    above += __shfl_sync(kFullMask, v, 31);
+  }
+  if (lane == 0) *gate = ref + half_bin_floor<kSampleShift>(chosen);
+  __syncwarp();
+}
+
+// Candidate records each lane of a math warp buffers per row and handoff
+// interval before it falls back to direct global stores.
+constexpr uint32_t kRingSlots = 8;
+// The drain's prefix sum packs two 16-bit per-row counts into one word.
+static_assert(32 * kRingSlots < 65536,
+              "per-warp record count must fit 16 bits");
+
+// Handoff ring: [bank][math warp][row][slot][lane] records, per-lane record
+// counts and the interval's first scan tile.
+struct alignas(16) HandoffRing {
+  uint32_t records[2][kMathWarps][kWgRows][kRingSlots][32];
+  uint32_t num_records[2][kMathWarps][kWgRows][32];
+  uint32_t bases[2][kMathWarps];
+};
+
+// samples: gate sample tiles of the block (0 when no row sees more keys than
+// the gate target, so every gate is -inf).
+struct ClaimHeader {
+  uint32_t vblock, all_short, kv_base, splits, slot, samples;
+};
+struct ClaimEntry : ClaimHeader {
+  uint32_t starts[kBlockRows], ends[kBlockRows];
+};
+
+// Claim k in a producer role (warps 9-11): waits for the entry, copies its
+// header and releases the entry (every lane arrives). False once the claimer
+// ends the epoch (kTerm / kPause).
+__device__ __forceinline__ bool consume_claim(uint32_t k, TBar* claim_full,
+                                              TBar* claim_empty,
+                                              const ClaimEntry* entries,
+                                              ClaimHeader& claim) {
+  claim_full[k & 1].wait((k >> 1) & 1);
+  claim = entries[k & 1];
+  __syncwarp();
+  if (claim.vblock >= kPause) return false;
+  claim_empty[k & 1].arrive();
+  return true;
+}
+
+// Barrier words (static shared, never aliased with data).
+struct BarIdx {
+  // TMEM slots, shared by the two math warpgroups (WG0 even, WG1 odd
+  // positions).
+  static constexpr uint32_t kTmemSlots = 3;
+  static constexpr uint32_t full_q = 0, empty_q = 1, full_kv = 2,
+                            empty_kv = 2 + kKVStages,
+                            full_tmem = 2 + 2 * kKVStages,
+                            empty_tmem = full_tmem + kTmemSlots,
+                            claim_full = empty_tmem + kTmemSlots,
+                            claim_empty = claim_full + 2,
+                            sel_done = claim_empty + 2, h_full = sel_done + 1,
+                            h_empty = h_full + 2 * kMathWarps,
+                            full_lin = h_empty + 2 * kMathWarps,
+                            total = full_lin + 1;
+};
+
+// Expect `bytes` of TMA transactions and arrive `count` times on one barrier.
+__device__ __forceinline__ void mbar_arrive_expect_tx(const void* b,
+                                                      uint32_t bytes,
+                                                      uint32_t count) {
+  const uint32_t a = static_cast<uint32_t>(__cvta_generic_to_shared(b));
+  asm volatile(
+      "mbarrier.expect_tx.relaxed.cta.shared::cta.b64 [%0], %1;" ::"r"(a),
+      "r"(bytes)
+      : "memory");
+  asm volatile(
+      "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0], %1;" ::"r"(a),
+      "r"(count)
+      : "memory");
+}
+
+// Candidate slabs: a block's 8 rows live in its CTA slot's 8 slab rows, so row
+// r maps to slab row r + slot_row_delta(slot, row_base).
+__device__ __forceinline__ int64_t slot_row_delta(uint32_t slot,
+                                                  uint32_t row_base) {
+  return int64_t(slot) * kBlockRows - int64_t(row_base);
+}
+// Offset of entry `index` of row `row` in the candidate slabs.
+__device__ __forceinline__ uint64_t candidate_offset(uint32_t row,
+                                                     int64_t row_delta,
+                                                     uint32_t index) {
+  return static_cast<uint64_t>(int64_t(row) + row_delta) * kCandidatesPerRow +
+         index;
+}
+// Stores ring record `record` (fp24 code << 8 | tile offset * 2 + KV half,
+// tile offset relative to the interval's first tile `interval_base`) of KV
+// lane `kv_lane` as entry `index` of row `row`.
+__device__ __forceinline__ void store_ring_record(
+    uint16_t* vals, int32_t* indices, uint32_t row, int64_t row_delta,
+    uint32_t index, uint32_t record, uint32_t kv_base, uint32_t interval_base,
+    uint32_t kv_lane) {
+  const uint32_t key =
+      kv_base + (interval_base * 2 + (record & 255u)) * 128 + kv_lane;
+  const uint64_t off = candidate_offset(row, row_delta, index);
+  store_candidate_record(vals + off, indices + off, record >> 8, key);
+}
+
+// Output row `out` of a row that sees at most top-k keys: its `visible` keys
+// in order (0 .. visible - 1, relative to the row start), then -1 padding.
+// Threads t, t + stride, ... write the slots.
+__device__ __forceinline__ void write_all_keys(int* out, uint32_t topk,
+                                               uint32_t visible, uint32_t t,
+                                               uint32_t stride) {
+  for (uint32_t j = t; j < topk; j += stride)
+    out[j] = j < visible ? int(j) : -1;
+}
+
+// Rewrites the output rows of the block at row_base in place from absolute key
+// indices to indices relative to the row's start (the top_k_per_row_prefill
+// convention; -1 padding stays). Rows that start at 0 and rows for which
+// skip(local row) holds are left as they are. Threads t, t + stride, ... each
+// rewrite 4 slots.
+template <class Skip>
+__device__ __forceinline__ void make_relative(int* output, uint32_t topk,
+                                              uint32_t rows, uint32_t row_base,
+                                              const uint32_t* starts_g,
+                                              uint32_t t, uint32_t stride,
+                                              Skip skip) {
+  const uint32_t quads = topk / 4;
+  for (uint32_t i = t; i < kBlockRows * quads; i += stride) {
+    const uint32_t r = row_base + i / quads;
+    if (r >= rows || skip(i / quads)) continue;
+    const int s = int(starts_g[r]);
+    if (s == 0) continue;
+    int4* p = reinterpret_cast<int4*>(output + (uint64_t)r * topk) + i % quads;
+    int4 v = *p;
+    v.x = v.x < 0 ? -1 : v.x - s;
+    v.y = v.y < 0 ? -1 : v.y - s;
+    v.z = v.z < 0 ? -1 : v.z - s;
+    v.w = v.w < 0 ? -1 : v.w - s;
+    *p = v;
+  }
+}
+
+// Stores one warp's ring records of a handoff interval in its rows' candidate
+// slabs. recs = [4 rows][kRingSlots records][32 lanes]; lane l holds n[row]
+// records of each row. Each row's slots are reserved with one atomic on
+// row_counts[row_base + row]; records past kCandidatesPerRow are dropped (the
+// row's count then exceeds the cap, so its block is rescanned exactly).
+// store(row, slot, record) writes one record.
+// Written for a warpgroup's 4 rows (two 16-bit counts per word in the scan).
+static_assert(kWgRows == 4, "store_ring_records packs 4 row counts");
+template <class Store>
+__device__ __forceinline__ void store_ring_records(const uint32_t* recs,
+                                                   const uint32_t (&n)[4],
+                                                   int* row_counts,
+                                                   uint32_t row_base,
+                                                   Store store) {
+  const uint32_t lane = threadIdx.x & 31;
+  uint32_t my_total = 0;
+#pragma unroll
+  for (uint32_t row = 0; row < 4; ++row) {
+    const uint32_t total = __reduce_add_sync(kFullMask, n[row]);
+    if (lane == row) my_total = total;
+  }
+  uint32_t my_start = 0;
+  if (lane < 4 && my_total)
+    my_start = atomicAdd(row_counts + row_base + lane, my_total);
+  // Two 16-bit per-row counts per word: one scan covers two rows.
+  uint32_t p01 = n[0] | (n[1] << 16), p23 = n[2] | (n[3] << 16);
+#pragma unroll
+  for (int dd = 1; dd < 32; dd *= 2) {
+    const uint32_t a = __shfl_up_sync(kFullMask, p01, dd),
+                   b = __shfl_up_sync(kFullMask, p23, dd);
+    if (lane >= dd) {
+      p01 += a;
+      p23 += b;
+    }
+  }
+  const uint32_t prefix[4] = {p01 & 65535u, p01 >> 16, p23 & 65535u, p23 >> 16};
+#pragma unroll
+  for (uint32_t row = 0; row < 4; ++row) {
+    const uint32_t start = __shfl_sync(kFullMask, my_start, row);
+    for (uint32_t j = 0; j < n[row]; ++j) {
+      const uint32_t slot = start + prefix[row] - n[row] + j;
+      const uint32_t record = recs[(row * kRingSlots + j) * 32 + lane];
+      if (slot < kCandidatesPerRow) store(row_base + row, slot, record);
+    }
+  }
+}
+
+// Self-drain, out of line: stores this warp's ring records of the last
+// handoff interval. recs = records[bank][warp].
+__device__ __noinline__ inline void self_drain_last(
+    const uint32_t* __restrict__ recs, uint32_t n0, uint32_t n1, uint32_t n2,
+    uint32_t n3, int* __restrict__ row_counts, uint32_t rbase, uint32_t lane_kv,
+    uint32_t flush_base, int64_t row_delta, uint16_t* __restrict__ vals,
+    int32_t* __restrict__ indices) {
+  const uint32_t n[4] = {n0, n1, n2, n3};
+  store_ring_records(recs, n, row_counts, rbase,
+                     [&](uint32_t row, uint32_t slot, uint32_t record) {
+                       store_ring_record(vals, indices, row, row_delta, slot,
+                                         record, lane_kv, flush_base, 0);
+                     });
+  __syncwarp();
+}
+
+// Drain warp: waits for producer warp `producer`'s handoff interval, stores its
+// records (unless the producer self-drained it) and frees the ring bank.
+template <class Store>
+__device__ __forceinline__ void drain_handoff(
+    HandoffRing& d, TBar* hfull, TBar* hempty, int producer, uint32_t interval,
+    uint32_t row_base, int* row_counts, bool skip, Store store) {
+  const int lane = threadIdx.x & 31, bank = interval & 1;
+  hfull[bank * kMathWarps + producer].wait((interval >> 1) & 1);
+  __syncwarp();
+  if (skip) {  // self-drained by the producer
+    hempty[bank * kMathWarps + producer].arrive();
+    return;
+  }
+  const uint32_t base = d.bases[bank][producer];
+  uint32_t n[4];
+#pragma unroll
+  for (uint32_t row = 0; row < 4; ++row)
+    n[row] = d.num_records[bank][producer][row][lane];
+  store_ring_records(&d.records[bank][producer][0][0][0], n, row_counts,
+                     row_base,
+                     [&](uint32_t row, uint32_t slot, uint32_t record) {
+                       store(row, slot, record, base, producer * 32 + lane);
+                     });
+  __syncwarp();
+  hempty[bank * kMathWarps + producer].arrive();
+}
+
+// Interval: scan tiles per candidate handoff to the drain warps.
+template <uint32_t Interval, uint32_t TopK>
+__global__
+__launch_bounds__(kNumThreads, 1) void fused_mqa_topk_fp8_prefill_kernel(
+    uint32_t rows, uint32_t keys, const uint32_t* starts_g,
+    const uint32_t* ends_g, int* out, uint16_t* vals, int* indices,
+    int* row_counts, RescanBuffers rescan, uint32_t phase_bytes,
+    uint32_t sample_target, uint32_t sample_stride, uint32_t sample_groups,
+    uint32_t sample_tiles, const float* scales,
+    const __grid_constant__ cute::TmaDescriptor tks,
+    const __grid_constant__ cute::TmaDescriptor tq,
+    const __grid_constant__ cute::TmaDescriptor ts,
+    const __grid_constant__ cute::TmaDescriptor tk,
+    const __grid_constant__ cute::TmaDescriptor tw, uint32_t nblocks,
+    uint32_t* counter, int* defer_next) {
+  // Compile-time: a runtime top-k costs ~0.5% in the scan (codegen).
+  constexpr uint32_t topk = TopK;
+  static_assert(BarIdx::kTmemSlots * kSlotCols <= 512,
+                "TMEM slots exceed 512 columns");
+  // A ring record's low byte holds (tile - interval start) * 2 + half, and the
+  // handoff points use `& (Interval - 1)`.
+  static_assert((Interval & (Interval - 1)) == 0 && Interval <= 128,
+                "Interval must be a power of two <= 128");
+  extern __shared__ __align__(1024) uint8_t smem_buffer[];
+  auto& sm = *reinterpret_cast<KernelSmem*>(smem_buffer);
+  uint8_t* const region = smem_buffer + sizeof(KernelSmem);
+  auto& sample_smem = *reinterpret_cast<SampleShared*>(region);
+  auto& hdata = *reinterpret_cast<HandoffRing*>(region);
+
+  __shared__ alignas(8) uint64_t bar_words[BarIdx::total];
+  __shared__ uint32_t tmem_ptr;
+  __shared__ ClaimEntry entries[2];
+  __shared__ float gate_s[kBlockRows];
+  __shared__ uint32_t sel_recover, sel_staged, sel_short,
+      row_offset[kBlockRows + 1];
+  __shared__ uint32_t defer_head, unsafe_rows, pause_req, epoch_kind;
+
+  TBar* const bars = reinterpret_cast<TBar*>(bar_words);
+  TBar& full_q = bars[BarIdx::full_q];
+  TBar& empty_q = bars[BarIdx::empty_q];
+  TBar* const full_kv = bars + BarIdx::full_kv;
+  // The 32 producer lanes stage the sample scales, so each arrives on full_kv.
+  constexpr uint32_t kFullKvArrivals = 32;
+  TBar* const empty_kv = bars + BarIdx::empty_kv;
+  TBar* const full_tmem = bars + BarIdx::full_tmem;
+  TBar* const empty_tmem = bars + BarIdx::empty_tmem;
+  TBar* const claim_full = bars + BarIdx::claim_full;
+  TBar* const claim_empty = bars + BarIdx::claim_empty;
+  TBar& sel_done = bars[BarIdx::sel_done];
+  TBar* const h_full = bars + BarIdx::h_full;
+  TBar* const h_empty = bars + BarIdx::h_empty;
+  TBar* const full_lin = bars + BarIdx::full_lin;
+
+  const uint32_t warp_idx = cutlass::canonical_warp_idx_sync();
+  const uint32_t lane = ptx::get_lane_idx();
+  constexpr uint32_t kClaimConsumers = 11;  // warps 9, 10, 11, 0-7
+
+  if (warp_idx == 8 && lane == 0) {
+    cute::prefetch_tma_descriptor(&tq);
+    cute::prefetch_tma_descriptor(&tw);
+    cute::prefetch_tma_descriptor(&tk);
+    cute::prefetch_tma_descriptor(&ts);
+    cute::prefetch_tma_descriptor(&tks);
+  }
+  if (warp_idx == 9 && cute::elect_one_sync()) {
+    defer_head = kTerm;
+    pause_req = 0;
+  }
+  // Epochs: the persistent role loops run until the work runs out (kTerm) or a
+  // block needs the exact rescan (kPause); rare blocks are then rescanned
+  // exactly (rescan.cuh) and the loops resume with freshly initialised
+  // barriers. The first epoch is emitted outside any loop (nesting the role
+  // loops in an epoch loop costs ~1.5 percent in the scan: uniform address math
+  // is rematerialised every tile). Later epochs only run after a recovery
+  // pause.
+  __shared__ uint32_t row_bounds[2][kBlockRows];
+  bool more = false;
+  {
+    constexpr uint32_t epoch = 0;
+#include "persistent_epoch.inc"
+  }
+#pragma unroll 1
+  for (uint32_t epoch = 1; more; ++epoch) {
+#include "persistent_epoch.inc"
+  }
+  // counter = {next block, finished CTAs}: the last CTA to finish (every claim
+  // is done by then) zeroes both, so the caller's workspace is ready for the
+  // next launch / graph replay.
+  if (threadIdx.x == 0) {
+    __threadfence();
+    if (atomicAdd(counter + 1, 1u) == gridDim.x - 1) {
+      // Acquire: every other CTA's last claim on counter happens before the
+      // reset below.
+      __threadfence();
+      atomicExch(counter, 0u);
+      atomicExch(counter + 1, 0u);
+    }
+  }
+}
+
+}  // namespace fused_mqa_topk::persist
diff --git a/csrc/fused_mqa_topk/rescan.cuh b/csrc/fused_mqa_topk/rescan.cuh
new file mode 100644
index 000000000000..80d9dc57370f
--- /dev/null
+++ b/csrc/fused_mqa_topk/rescan.cuh
@@ -0,0 +1,350 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+//
+// Exact recovery for the blocks the persistent kernel cannot select from its
+// gated candidates (see persistent_kernel.cuh): the block is rescanned densely,
+// and every 1024-key segment of fp32 scores is merged into a running per-row
+// top set (kept <= kRetainedPerRow, compacted to K).
+#pragma once
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+#include <cutlass/arch/barrier.h>
+#include <math_constants.h>
+#include <cstdint>
+
+#include "common.cuh"
+#include "constants.h"
+#include "dense_scan.cuh"
+
+namespace fused_mqa_topk::rescan {
+using dense::Scheduler;
+
+// Keys whose scores are merged into the retained sets at a time.
+constexpr uint32_t kSegmentKeys = 1024;
+// Per-row retained top set of the exact rescan (scores and key indices).
+struct RescanBuffers {
+  float* retained;
+  int* retained_indices;
+};
+// Per-row segment histograms: half_bin<6> of a score (1024 bins); the radix
+// levels of compact use 256 bins.
+constexpr uint32_t kSegmentBins = 1024, kRadixBins = 256;
+struct alignas(16) RescanShared {
+  float ring[2][kBlockRows][kSegmentKeys];
+  uint32_t hist[kBlockRows][padded_bins(kSegmentBins)], gates[kBlockRows],
+      kept[kBlockRows];
+  uint32_t work[padded_bins(kRadixBins)], prefix, rank, emit_counts[kMathWarps];
+  __device__ void init(uint32_t tid) {
+    for (uint32_t j = tid; j < sizeof(hist) / sizeof(uint32_t);
+         j += kNumThreads)
+      reinterpret_cast<uint32_t*>(hist)[j] = 0;
+    if (tid < kBlockRows) {
+      gates[tid] = __float_as_uint(-CUDART_INF_F);
+      kept[tid] = 0;
+    }
+  }
+};
+// What compact_retained's emit pass keeps.
+enum class Emit { kAboveGate, kAboveCode, kTies };
+
+__device__ inline void compact_retained(RescanShared& shared,
+                                        RescanBuffers buffers, uint32_t topk,
+                                        uint32_t local, uint32_t row,
+                                        uint32_t tid,
+                                        bool compact_to_k = false) {
+  constexpr uint32_t kThreads = kMathThreads,
+                     kItems = kRetainedPerRow / kThreads;
+  uint32_t lane = tid & 31, count = shared.kept[local];
+  uint64_t base = (uint64_t)row * kRetainedPerRow;
+  float values[kItems];
+  int indices[kItems];
+#pragma unroll
+  for (uint32_t i = 0; i < kItems; ++i) {
+    uint32_t j = tid + i * kThreads;
+    values[i] = j < count ? buffers.retained[base + j] : -CUDART_INF_F;
+    indices[i] = j < count ? buffers.retained_indices[base + j] : -1;
+  }
+  // Emits the values >= gate, the codes above threshold, or the codes equal
+  // to it (ties, up to K) to the retained set from base_offset on.
+  auto emit = [&](float gate, uint32_t threshold, Emit mode,
+                  uint32_t base_offset) {
+    const auto passes = [&](float v) {
+      return isfinite(v) &&
+             (mode == Emit::kAboveGate
+                  ? v >= gate
+                  : (mode == Emit::kAboveCode ? ordered_code(v) > threshold
+                                              : ordered_code(v) == threshold));
+    };
+    uint32_t total = 0;
+#pragma unroll
+    for (uint32_t i = 0; i < kItems; ++i) total += passes(values[i]);
+    uint32_t warp_total = __reduce_add_sync(kFullMask, total);
+    if (lane == 0) shared.emit_counts[tid / 32] = warp_total;
+    math_wg_sync();
+    uint32_t start = base_offset;
+    for (uint32_t w = 0; w < tid / 32; ++w) start += shared.emit_counts[w];
+    uint32_t all = 0;
+    for (uint32_t w = 0; w < kMathWarps; ++w) all += shared.emit_counts[w];
+#pragma unroll
+    for (uint32_t i = 0; i < kItems; ++i) {
+      const float v = values[i];
+      const bool pass = passes(v);
+      uint32_t mask = __ballot_sync(kFullMask, pass),
+               pos = start + __popc(mask & ((1u << lane) - 1));
+      if (pass && (mode != Emit::kTies || pos < topk)) {
+        buffers.retained[base + pos] = v;
+        buffers.retained_indices[base + pos] = indices[i];
+      }
+      start += __popc(mask);
+    }
+    math_wg_sync();
+    return all;
+  };
+  if (!compact_to_k) {
+    if (tid == 0) {
+      shared.rank = 0;
+    }
+    math_wg_sync();
+    float gate = __uint_as_float(shared.gates[local]);
+    uint32_t kept = 0;
+#pragma unroll
+    for (uint32_t i = 0; i < kItems; ++i)
+      kept += isfinite(values[i]) && values[i] >= gate;
+    kept = __reduce_add_sync(kFullMask, kept);
+    if (lane == 0) atomicAdd(&shared.rank, kept);
+    math_wg_sync();
+    kept = shared.rank;
+    if (kept >= topk && kept <= kRetainedPerRow - kSegmentKeys) {
+      emit(gate, 0, Emit::kAboveGate, 0);
+      math_wg_sync();
+      if (tid == 0) shared.kept[local] = kept;
+      math_wg_sync();
+      return;
+    }
+  }
+  if (tid == 0) {
+    shared.prefix = 0;
+    shared.rank = topk;
+  }
+  math_wg_sync();
+  for (int shift = 24; shift >= 0; shift -= 8) {
+    for (uint32_t j = tid; j < padded_bins(kRadixBins); j += kThreads)
+      shared.work[j] = 0;
+    math_wg_sync();
+    uint32_t prefix = shared.prefix;
+#pragma unroll
+    for (uint32_t i = 0; i < kItems; ++i) {
+      uint32_t code = ordered_code(values[i]);
+      bool valid = isfinite(values[i]) &&
+                   (shift == 24 || (code >> (shift + 8)) == prefix);
+      uint32_t mask = __ballot_sync(kFullMask, valid);
+      if (valid) {
+        uint32_t bin = (code >> shift) & 255u,
+                 peers = __match_any_sync(mask, bin);
+        if (lane == uint32_t(__ffs(peers) - 1))
+          atomicAdd(shared.work + padded_bin(bin), __popc(peers));
+      }
+    }
+    math_wg_sync();
+    if (tid < 32) {
+      uint32_t sum = 0;
+      for (uint32_t j = 0; j < 8; ++j) {
+        uint32_t bin = 255 - lane * 8 - j;
+        sum += shared.work[padded_bin(bin)];
+      }
+      const uint32_t inclusive = warp_inclusive_scan(sum, lane);
+      uint32_t before = inclusive - sum, rank = shared.rank,
+               crossing =
+                   __ballot_sync(kFullMask, before < rank && inclusive >= rank),
+               chosen_lane = __ffs(crossing) - 1, chosen = 0, above = before;
+      if (lane == chosen_lane) {
+        for (uint32_t j = 0; j < 8; ++j) {
+          uint32_t bin = 255 - lane * 8 - j;
+          if (above + shared.work[padded_bin(bin)] >= rank) {
+            chosen = bin;
+            break;
+          }
+          above += shared.work[padded_bin(bin)];
+        }
+      }
+      chosen = __shfl_sync(kFullMask, chosen, chosen_lane);
+      above = __shfl_sync(kFullMask, above, chosen_lane);
+      __syncwarp();  // every lane read shared.rank above
+      if (lane == 0) {
+        shared.prefix = (prefix << 8) | chosen;
+        shared.rank = rank - above;
+      }
+    }
+    math_wg_sync();
+  }
+  uint32_t threshold = shared.prefix;
+  uint32_t greater = emit(0, threshold, Emit::kAboveCode, 0);
+  emit(0, threshold, Emit::kTies, greater);
+  if (tid == 0) {
+    shared.kept[local] = topk;
+    float gate =
+        fmaxf(__uint_as_float(shared.gates[local]), ordered_float(threshold));
+    atomicExch(shared.gates + local, __float_as_uint(gate));
+  }
+  math_wg_sync();
+}
+
+__device__ inline void merge_segment(RescanShared& shared,
+                                     RescanBuffers buffers, uint32_t topk,
+                                     uint32_t local, uint32_t row,
+                                     uint32_t segment, uint32_t first,
+                                     uint32_t n, uint32_t lane) {
+  constexpr uint32_t kItems = kSegmentKeys / 32;
+  const uint32_t bank = segment & 1;
+  float values[kItems], gate = __uint_as_float(shared.gates[local]);
+  uint32_t count = 0;
+#pragma unroll
+  for (uint32_t i = 0; i < kItems; ++i) {
+    uint32_t j = lane + i * 32;
+    float v = j < n ? shared.ring[bank][local][j] : -CUDART_INF_F;
+    if (v <= gate) v = -CUDART_INF_F;
+    values[i] = v;
+    bool valid = isfinite(v);
+    count += valid;
+    uint32_t mask = __ballot_sync(kFullMask, valid);
+    if (valid) {
+      uint32_t bin = half_bin<6>(v), peers = __match_any_sync(mask, bin);
+      if (lane == uint32_t(__ffs(peers) - 1))
+        atomicAdd(&shared.hist[local][padded_bin(bin)], __popc(peers));
+    }
+  }
+  count = __reduce_add_sync(kFullMask, count);
+  __syncwarp();
+  if (count) {
+    uint32_t sum = 0;
+#pragma unroll
+    for (uint32_t j = 0; j < 32; ++j)
+      sum += shared.hist[local][padded_bin((31 - lane) * 32 + j)];
+    uint32_t inclusive = warp_inclusive_scan(sum, lane);
+    uint32_t before = inclusive - sum,
+             mask =
+                 __ballot_sync(kFullMask, before < topk && inclusive >= topk);
+    if (mask) {
+      uint32_t chosen = __ffs(mask) - 1, group = 31 - chosen,
+               rank = topk - __shfl_sync(kFullMask, before, chosen),
+               bin = group * 32 + 31 - lane;
+      sum = shared.hist[local][padded_bin(bin)];
+      inclusive = warp_inclusive_scan(sum, lane);
+      before = inclusive - sum;
+      chosen =
+          __ffs(__ballot_sync(kFullMask, before < rank && inclusive >= rank)) -
+          1;
+      bin = __shfl_sync(kFullMask, bin, chosen);
+      gate = fmaxf(gate, half_bin_floor<6>(bin));
+      if (lane == 0) atomicExch(shared.gates + local, __float_as_uint(gate));
+    }
+    uint32_t written = shared.kept[local];
+    uint64_t base = (uint64_t)row * kRetainedPerRow;
+#pragma unroll
+    for (uint32_t i = 0; i < kItems; ++i) {
+      float v = values[i];
+      bool pass = isfinite(v) && v >= gate;
+      uint32_t mask = __ballot_sync(kFullMask, pass),
+               pos = written + __popc(mask & ((1u << lane) - 1));
+      if (pass) {
+        buffers.retained[base + pos] = v;
+        buffers.retained_indices[base + pos] = first + lane + i * 32;
+      }
+      written += __popc(mask);
+    }
+    __syncwarp();
+    if (lane == 0) shared.kept[local] = written;
+    __syncwarp();
+  }
+}
+
+// Streams each row's scores (above its running gate) into the segment ring and
+// merges every full segment into the row's retained top set; the final set is
+// written to the output.
+struct RescanPolicy {
+  static constexpr float kWeightScale = 0.5f;
+  static constexpr uint32_t kSplitsPerSegment = kSegmentKeys / kTileKeys;
+  RescanShared& shared;
+  RescanBuffers buffers;
+  int* out;  // [rows, topk]
+  uint32_t topk, query_base, row_base, segment = 0;
+  uint32_t seq_k_start[kWgRows], seq_k_end[kWgRows];
+  float gate[kWgRows];
+  CUTLASS_DEVICE void begin_block(const Scheduler& sched, uint32_t q_block,
+                                  uint32_t base) {
+    query_base = base;
+    row_base = sched.get_logits_row(q_block, 0);
+#pragma unroll
+    for (uint32_t i = 0; i < kWgRows; ++i) {
+      const uint32_t row = sched.get_logits_row(q_block, base + i);
+      gate[i] = -CUDART_INF_F;
+      seq_k_start[i] = min(sched.cu_seq_len_k_start[row], sched.num_kv_tokens);
+      seq_k_end[i] = min(sched.cu_seq_len_k_end[row], sched.num_kv_tokens);
+    }
+  }
+  CUTLASS_DEVICE void begin_split(uint32_t split) {
+    if (split % kSplitsPerSegment == 0) {
+      __syncwarp();
+      for (uint32_t i = 0; i < kWgRows; ++i)
+        gate[i] = __uint_as_float(atomicAdd(shared.gates + query_base + i, 0u));
+    }
+  }
+  CUTLASS_DEVICE void score(uint32_t i, uint32_t kv_offset, uint32_t kv_lane,
+                            uint32_t split, float bq) {
+    const bool valid = kv_offset >= seq_k_start[i] &&
+                       kv_offset < seq_k_end[i] && isfinite(bq) && bq > gate[i];
+    shared.ring[segment & 1][query_base + i]
+               [(split % kSplitsPerSegment) * kTileKeys + kv_lane] =
+        valid ? bq : -CUDART_INF_F;
+  }
+  CUTLASS_DEVICE void end_split(uint32_t split, uint32_t num_splits,
+                                uint32_t kv_base) {
+    if ((split + 1) % kSplitsPerSegment != 0 && split + 1 != num_splits) return;
+    __threadfence_block();
+    math_wg_sync();
+    const uint32_t tid = threadIdx.x;
+    for (uint32_t local = 0; local < kBlockRows; ++local)
+      if (shared.kept[local] > kRetainedPerRow - kSegmentKeys)
+        compact_retained(shared, buffers, topk, local, row_base + local, tid);
+    math_wg_sync();
+    const uint32_t local = tid / 32;
+    merge_segment(shared, buffers, topk, local, row_base + local, segment,
+                  kv_base + (split / kSplitsPerSegment) * kSegmentKeys,
+                  (split % kSplitsPerSegment + 1) * kTileKeys, tid % 32);
+    math_wg_sync();
+    if (split + 1 == num_splits)
+      for (uint32_t l = 0; l < kBlockRows; ++l)
+        if (shared.kept[l] > topk)
+          compact_retained(shared, buffers, topk, l, row_base + l, tid, true);
+    math_wg_sync();
+    ++segment;
+  }
+  CUTLASS_DEVICE void end_block(const Scheduler& sched, uint32_t q_block) {
+    const uint32_t local = threadIdx.x / 32, lane = threadIdx.x % 32;
+    const uint32_t row = sched.get_logits_row(q_block, local),
+                   count = shared.kept[local];
+    if (row < sched.num_q_tokens)
+      for (uint32_t j = lane; j < topk; j += 32)
+        out[(uint64_t)row * topk + j] =
+            j < count
+                ? buffers.retained_indices[(uint64_t)row * kRetainedPerRow + j]
+                : -1;
+  }
+};
+
+// Rescans one block exactly and writes its output rows of `out`; `buffers`
+// rows are addressed by absolute row (callers offset them). Out of line: the
+// rare rescan must not add register pressure to the caller's hot loops.
+__device__ __noinline__ inline void rescan_block(
+    const Scheduler& block, RescanBuffers buffers, int* out, uint32_t topk,
+    const cute::TmaDescriptor& tq, const cute::TmaDescriptor& tk,
+    const cute::TmaDescriptor& ts, const cute::TmaDescriptor& tw) {
+  extern __shared__ __align__(1024) uint8_t smem_buffer[];
+  auto& shared = *reinterpret_cast<RescanShared*>(smem_buffer +
+                                                  sizeof(dense::SharedStorage));
+  shared.init(threadIdx.x);
+  RescanPolicy policy{shared, buffers, out, topk};
+  dense::score_block(block, policy, tq, tk, ts, tw);
+}
+
+}  // namespace fused_mqa_topk::rescan
diff --git a/csrc/fused_mqa_topk/select.cuh b/csrc/fused_mqa_topk/select.cuh
new file mode 100644
index 000000000000..759ec7fce46a
--- /dev/null
+++ b/csrc/fused_mqa_topk/select.cuh
@@ -0,0 +1,452 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+//
+// Per-row top-K selection of the persistent kernel's candidates, one warp per
+// row (K = topk, the row's top-k: 512, 1024 or 2048, a runtime value).
+//
+// select_staged_row: bracket the K-th largest fp24 code with sampled codes,
+// place the codes above the bracket directly, and resolve the bracket with
+// histogram levels on register-resident codes.
+//
+// select_row: the fallback of select_staged_row when its sampled bracket
+// misses. Candidate codes (higher code = higher score, staged once in shared
+// memory) span a narrow range (gate .. row max), so each level histograms the
+// range-relative code into at most kBins bins: collisions are rare, plain
+// shared atomics suffice and each lane keeps kUnroll independent elements in
+// flight. Levels repeat inside the boundary bin until it is one code wide;
+// codes above it are emitted, then ties until K.
+#pragma once
+#include <cstdint>
+
+#include "candidates.cuh"
+#include "common.cuh"
+#include "constants.h"
+
+namespace fused_mqa_topk::selection {
+constexpr uint32_t kBinsLog2 = 10, kBins = 1u << kBinsLog2;
+// Elements each lane keeps in flight.
+constexpr uint32_t kUnroll = 8;
+
+// Bin (descending) where the cumulative count first reaches `rank`;
+// `above` = elements in higher bins. Lane l owns Bins / 32 bins, descending
+// from bin Bins - 1 - l * Bins / 32.
+template <uint32_t Bins>
+__device__ __forceinline__ uint32_t find_bin_n(const uint32_t* hist,
+                                               uint32_t rank, uint32_t& above) {
+  const uint32_t lane = threadIdx.x & 31;
+  constexpr uint32_t kPerLane = Bins / 32;
+  uint32_t sum = 0;
+#pragma unroll
+  for (uint32_t j = 0; j < kPerLane; ++j)
+    sum += hist[Bins - 1 - lane * kPerLane - j];
+  const uint32_t inclusive = warp_inclusive_scan(sum, lane);
+  const uint32_t before = inclusive - sum;
+  const uint32_t winner =
+      __ffs(__ballot_sync(kFullMask, before < rank && inclusive >= rank)) - 1;
+  uint32_t chosen = 0, acc = before;
+  if (lane == winner) {
+    for (uint32_t j = 0; j < kPerLane; ++j) {
+      const uint32_t bin = Bins - 1 - lane * kPerLane - j, h = hist[bin];
+      if (acc + h >= rank) {
+        chosen = bin;
+        break;
+      }
+      acc += h;
+    }
+  }
+  above = __shfl_sync(kFullMask, acc, winner);
+  return __shfl_sync(kFullMask, chosen, winner);
+}
+
+// Requires count >= topk (callers recover rows with fewer candidates).
+__device__ inline void select_row(uint32_t* codes, uint32_t count,
+                                  const uint16_t* values, const int32_t* packed,
+                                  int32_t* out, uint32_t topk, uint32_t* hist) {
+  const uint32_t lane = threadIdx.x & 31;
+  uint32_t lo = 0xffffffffu, hi = 0;
+  uint32_t j = lane;
+  for (; j + 32 * (kUnroll - 1) < count; j += 32 * kUnroll) {
+    uint16_t v[kUnroll];
+    int32_t p[kUnroll];
+#pragma unroll
+    for (uint32_t u = 0; u < kUnroll; ++u) {
+      v[u] = values[j + 32 * u];
+      p[u] = packed[j + 32 * u];
+    }
+#pragma unroll
+    for (uint32_t u = 0; u < kUnroll; ++u) {
+      const uint32_t c = candidate_load_score_code(v[u], p[u]);
+      codes[j + 32 * u] = c;
+      lo = min(lo, c);
+      hi = max(hi, c);
+    }
+  }
+  for (; j < count; j += 32) {
+    const uint32_t c = candidate_load_score_code(values[j], packed[j]);
+    codes[j] = c;
+    lo = min(lo, c);
+    hi = max(hi, c);
+  }
+  lo = __reduce_min_sync(kFullMask, lo);
+  hi = __reduce_max_sync(kFullMask, hi);
+  __syncwarp();
+
+  // Current bucket: codes in [lo, lo + width). rank = elements still needed
+  // from it.
+  uint32_t width = hi - lo + 1, rank = topk;
+#pragma unroll 1
+  while (width > 1) {
+    const uint32_t spread = (width - 1) >> kBinsLog2;
+    const uint32_t shift = spread ? 32 - __clz(spread) : 0;
+    for (uint32_t b = lane; b < kBins; b += 32) hist[b] = 0;
+    __syncwarp();
+    const uint32_t top = lo + width - 1;
+    uint32_t i = lane;
+    for (; i + 32 * (kUnroll - 1) < count; i += 32 * kUnroll) {
+      uint32_t c[kUnroll];
+#pragma unroll
+      for (uint32_t u = 0; u < kUnroll; ++u) c[u] = codes[i + 32 * u];
+#pragma unroll
+      for (uint32_t u = 0; u < kUnroll; ++u)
+        if (c[u] >= lo && c[u] <= top)
+          atomicAdd(hist + ((c[u] - lo) >> shift), 1u);
+    }
+    for (; i < count; i += 32) {
+      const uint32_t c = codes[i];
+      if (c >= lo && c <= top) atomicAdd(hist + ((c - lo) >> shift), 1u);
+    }
+    __syncwarp();
+    uint32_t above;
+    const uint32_t chosen = find_bin_n<kBins>(hist, rank, above);
+    rank -= above;
+    lo += chosen << shift;
+    width = min(1u << shift, top - lo + 1);
+    __syncwarp();
+  }
+
+  // lo is now the K-th largest code: emit codes above it, then ties until K.
+  const uint32_t threshold = lo, lt = (1u << lane) - 1u;
+  uint32_t pos = 0;
+#pragma unroll 1
+  for (uint32_t tie = 0; tie < 2; ++tie) {
+    for (uint32_t j0 = 0; j0 < count && pos < topk; j0 += 32 * kUnroll) {
+      bool pass[kUnroll];
+      int32_t p[kUnroll];
+#pragma unroll
+      for (uint32_t u = 0; u < kUnroll; ++u) {
+        const uint32_t e = j0 + 32 * u + lane;
+        const uint32_t c = e < count ? codes[e] : 0;
+        pass[u] = e < count && (tie ? c == threshold : c > threshold);
+        p[u] = pass[u] ? packed[e] : 0;
+      }
+#pragma unroll
+      for (uint32_t u = 0; u < kUnroll; ++u) {
+        const uint32_t mask = __ballot_sync(kFullMask, pass[u]),
+                       at = pos + __popc(mask & lt);
+        if (pass[u] && at < topk) out[at] = candidate_decode_index(p[u]);
+        pos += __popc(mask);
+      }
+    }
+  }
+}
+
+// Register samples per lane for the sampled bracket of select_staged_row.
+constexpr uint32_t kSamplesPerLane = 16;
+// Largest bracket (elements between the two sampled bounds) that
+// select_staged_row resolves in registers; wider brackets fall back to
+// select_row.
+constexpr uint32_t kMidCap = 1024;
+// Bins of the histogram levels that resolve the exact threshold inside a
+// bracket (select_staged_row).
+constexpr uint32_t kMidBins = 512;
+
+// Bracket element ids are stored as uint16 (midbuf).
+static_assert(kCandidatesPerRow <= 65536, "candidate ids must fit 16 bits");
+// Per-warp scratch of select_staged_row is kBins words: radix_bracket's 768,
+// or the midbuf (kMidCap uint16) plus the kMidBins-bin histogram.
+static_assert(768 <= kBins && kMidCap / 2 + kMidBins <= kBins,
+              "selection scratch exceeds kBins words");
+
+// r-th largest (1-based) of the kNumSamples * 32 register samples (kNumSamples
+// per lane) for two ranks at once
+// (r_hi <= r_lo), via a 256-bin hi8 level and two 256-bin lo8 levels. h8: 768
+// words of scratch.
+template <uint32_t kNumSamples = kSamplesPerLane>
+__device__ __forceinline__ void radix_bracket(
+    const uint32_t (&smp)[kNumSamples], uint32_t r_hi, uint32_t r_lo,
+    uint32_t* h8, uint32_t& h_hi, uint32_t& h_lo) {
+  const uint32_t lane = threadIdx.x & 31;
+#pragma unroll
+  for (uint32_t b = lane; b < 192; b += 32)
+    reinterpret_cast<uint4*>(h8)[b] = make_uint4(0, 0, 0, 0);
+  __syncwarp();
+#pragma unroll
+  for (uint32_t u = 0; u < kNumSamples; ++u) atomicAdd(h8 + (smp[u] >> 8), 1u);
+  __syncwarp();
+  uint32_t ab_hi, ab_lo;
+  const uint32_t b_hi = find_bin_n<256>(h8, r_hi, ab_hi);
+  const uint32_t b_lo = find_bin_n<256>(h8, r_lo, ab_lo);
+  uint32_t* const lh = h8 + 256;
+  uint32_t* const ll = h8 + 512;
+#pragma unroll
+  for (uint32_t u = 0; u < kNumSamples; ++u) {
+    const uint32_t b = smp[u] >> 8;
+    if (b == b_hi) atomicAdd(lh + (smp[u] & 255u), 1u);
+    if (b == b_lo) atomicAdd(ll + (smp[u] & 255u), 1u);
+  }
+  __syncwarp();
+  uint32_t d0, d1;
+  h_hi = b_hi << 8 | find_bin_n<256>(lh, r_hi - ab_hi, d0);
+  h_lo = b_lo << 8 | find_bin_n<256>(ll, r_lo - ab_lo, d1);
+  __syncwarp();  // h8 is reused as the output stage
+}
+
+// Brackets the topk-th largest of the `count` staged codes between the hi16
+// codes of sample ranks expect -+ margin, where the kPerLane * 32 samples sit
+// at evenly spaced quantiles and expect is the topk-th largest's expected rank
+// among them.
+template <uint32_t kPerLane>
+__device__ __forceinline__ void sampled_bracket(const uint32_t* codes,
+                                                uint32_t count, uint32_t topk,
+                                                uint32_t margin, uint32_t* hist,
+                                                uint32_t& h_hi,
+                                                uint32_t& h_lo) {
+  constexpr uint32_t kSamples = kPerLane * 32;
+  const uint32_t lane = threadIdx.x & 31;
+  uint32_t smp[kPerLane];
+#pragma unroll
+  for (uint32_t u = 0; u < kPerLane; ++u)
+    smp[u] = codes[((u * 32 + lane) * 2 + 1) * count / (2 * kSamples)] >> 8;
+  const uint32_t expect = (topk * kSamples + count / 2) / count;
+  const uint32_t r_hi = expect > margin ? expect - margin : 1,
+                 r_lo = min(expect + margin, kSamples);
+  radix_bracket<kPerLane>(smp, r_hi, r_lo, hist, h_hi, h_lo);
+}
+
+// find_bin_n with each lane's kPer contiguous bins loaded as 16-byte vectors
+// and searched in registers (needs a 16-byte aligned histogram and kPer % 4 ==
+// 0; find_bin_n reads the bins one at a time and has neither requirement).
+template <uint32_t Bins>
+__device__ __forceinline__ uint32_t find_bin_v(const uint32_t* hist,
+                                               uint32_t rank, uint32_t& above) {
+  const uint32_t lane = threadIdx.x & 31;
+  constexpr uint32_t kPer = Bins / 32, kVec = kPer / 4;
+  static_assert(kPer % 4 == 0);
+  // lane l owns bins [Bins - (l + 1) * kPer, Bins - l * kPer), highest first.
+  const uint4* hv =
+      reinterpret_cast<const uint4*>(hist + Bins - (lane + 1) * kPer);
+  uint32_t h[kPer];
+#pragma unroll
+  for (uint32_t v = 0; v < kVec; ++v) {
+    const uint4 x = hv[v];
+    h[4 * v] = x.x;
+    h[4 * v + 1] = x.y;
+    h[4 * v + 2] = x.z;
+    h[4 * v + 3] = x.w;
+  }
+  uint32_t sum = 0;
+#pragma unroll
+  for (uint32_t j = 0; j < kPer; ++j) sum += h[j];
+  const uint32_t inclusive = warp_inclusive_scan(sum, lane);
+  const uint32_t before = inclusive - sum;
+  const uint32_t winner =
+      __ffs(__ballot_sync(kFullMask, before < rank && inclusive >= rank)) - 1;
+  // Descending walk inside the lane: bin index Bins - lane * kPer - 1 - j  <->
+  // h[kPer - 1 - j].
+  uint32_t acc = before, chosen = 0, found = 0, acc_at = before;
+#pragma unroll
+  for (uint32_t j = 0; j < kPer; ++j) {
+    const uint32_t x = h[kPer - 1 - j];
+    const bool hit = !found && acc + x >= rank;
+    if (hit) {
+      chosen = Bins - lane * kPer - 1 - j;
+      acc_at = acc;
+      found = 1;
+    }
+    acc += x;
+  }
+  above = __shfl_sync(kFullMask, acc_at, winner);
+  return __shfl_sync(kFullMask, chosen, winner);
+}
+
+// One row's global-memory candidates, staged to shared memory as
+// fp24 codes (codes must be 16-byte aligned, i.e. row offsets rounded up to 4).
+// Register budget: ptxas compiles the whole kernel to 168 regs, so every pass
+// keeps <= ~32 regs of arrays live (more spills into the scan loop).
+// Passes: stage (16-byte global loads of values / packed indices, two chunks of
+// 8 in flight per lane, to fp24 codes); hi16 bracket (radix levels over 512
+// shared-memory samples); count above / inside the bracket (16-byte shared
+// loads); collect (codes above the bracket to out, bracket element ids to
+// midbuf); histogram levels on register codes; emit.
+// Returns false before writing out; the caller's select_row then restages and
+// runs the levels.
+__device__ __noinline__ inline bool select_staged_row(
+    uint32_t* __restrict__ codes, uint32_t count,
+    const uint16_t* __restrict__ values, const int32_t* __restrict__ packed,
+    int32_t* __restrict__ out, uint32_t topk, uint32_t* hist) {
+  const uint32_t lane = threadIdx.x & 31;
+  if (count <= topk) return false;
+  const uint32_t chunks = (count + 7) / 8;
+  constexpr uint32_t kChunksInFlight = 2;
+  const uint4* v4p = reinterpret_cast<const uint4*>(values);
+  const uint4* p4p = reinterpret_cast<const uint4*>(packed);
+#pragma unroll 1
+  for (uint32_t c0i = lane; c0i < chunks; c0i += 32 * kChunksInFlight) {
+    uint4 vv[kChunksInFlight], pp0[kChunksInFlight], pp1[kChunksInFlight];
+#pragma unroll
+    for (uint32_t u = 0; u < kChunksInFlight; ++u) {
+      const uint32_t c = min(c0i + 32 * u, chunks - 1);
+      vv[u] = __ldcg(v4p + c);
+      pp0[u] = __ldcg(p4p + 2 * c);
+      pp1[u] = __ldcg(p4p + 2 * c + 1);
+    }
+#pragma unroll
+    for (uint32_t u = 0; u < kChunksInFlight; ++u) {
+      const uint32_t c = c0i + 32 * u;
+      if (c < chunks) {
+        const uint32_t vw[4] = {vv[u].x, vv[u].y, vv[u].z, vv[u].w};
+        const uint32_t pw[8] = {pp0[u].x, pp0[u].y, pp0[u].z, pp0[u].w,
+                                pp1[u].x, pp1[u].y, pp1[u].z, pp1[u].w};
+        uint32_t cw[8];
+#pragma unroll
+        for (uint32_t k = 0; k < 8; ++k)
+          cw[k] = candidate_load_score_code(
+              uint16_t(vw[k / 2] >> (16 * (k & 1))), int32_t(pw[k]));
+        if (8 * c + 8 <= count) {
+          reinterpret_cast<uint4*>(codes)[2 * c] =
+              make_uint4(cw[0], cw[1], cw[2], cw[3]);
+          reinterpret_cast<uint4*>(codes)[2 * c + 1] =
+              make_uint4(cw[4], cw[5], cw[6], cw[7]);
+        } else {
+#pragma unroll
+          for (uint32_t k = 0; k < 8; ++k)
+            if (8 * c + k < count) codes[8 * c + k] = cw[k];
+        }
+      }
+    }
+  }
+  __syncwarp();
+  // The bracket is the code range between the sampled ranks expect -+ margin
+  // (expect = the K-th largest's expected rank among the samples). Rows up to
+  // kFewCandidates use 512 samples, +-40 ranks. Long-context rows (4-6K
+  // candidates) would get a bracket of ~16% of the row that way, over kMidCap
+  // for 10-20% of the rows at 512K-1M: 1024 samples, +-48 (~3 sd), give ~9%.
+  constexpr uint32_t kFewCandidates = 3584, kMargin = 40, kWideMargin = 48;
+  uint32_t h_hi, h_lo;
+  if (count <= kFewCandidates)
+    sampled_bracket<kSamplesPerLane>(codes, count, topk, kMargin, hist, h_hi,
+                                     h_lo);
+  else
+    sampled_bracket<2 * kSamplesPerLane>(codes, count, topk, kWideMargin, hist,
+                                         h_hi, h_lo);
+  const uint4* c4p = reinterpret_cast<const uint4*>(codes);
+  const uint32_t quads = (count + 3) / 4;
+  const uint32_t lt = (1u << lane) - 1u;
+  uint16_t* midbuf = reinterpret_cast<uint16_t*>(hist);
+  uint32_t* h512 = hist + kMidCap / 2;
+  // Collect with warp-ballot placement: passing lanes of one step write
+  // consecutive slots, so the global stores of the above elements are
+  // coalesced. Writes are bounded; the
+  // bracket check follows (out may be partially written on failure: the
+  // fallback select_row rewrites all K slots).
+  uint32_t above = 0, mid = 0;
+  constexpr uint32_t kQuadsInFlight = 4;
+#pragma unroll 1
+  for (uint32_t q0 = 0; q0 < quads; q0 += 32 * kQuadsInFlight) {
+    uint4 pp[kQuadsInFlight], cc[kQuadsInFlight];
+#pragma unroll
+    for (uint32_t u = 0; u < kQuadsInFlight; ++u) {
+      const uint32_t q = min(q0 + 32 * u + lane, quads - 1);
+      pp[u] = __ldcg(p4p + q);
+      cc[u] = c4p[q];
+    }
+#pragma unroll
+    for (uint32_t u = 0; u < kQuadsInFlight; ++u) {
+      const uint32_t q = q0 + 32 * u + lane;
+      const uint32_t w[4] = {cc[u].x, cc[u].y, cc[u].z, cc[u].w};
+      const uint32_t p[4] = {pp[u].x, pp[u].y, pp[u].z, pp[u].w};
+#pragma unroll
+      for (uint32_t k = 0; k < 4; ++k) {
+        const uint32_t e = 4 * q + k, h = w[k] >> 8;
+        const bool valid = e < count;
+        const bool a = valid && h > h_hi, m = valid && !a && h >= h_lo;
+        const uint32_t ma = __ballot_sync(kFullMask, a),
+                       mm = __ballot_sync(kFullMask, m);
+        const uint32_t at = above + __popc(ma & lt), mt = mid + __popc(mm & lt);
+        if (a && at < topk) out[at] = candidate_decode_index(int32_t(p[k]));
+        if (m && mt < kMidCap) midbuf[mt] = uint16_t(e);
+        above += __popc(ma);
+        mid += __popc(mm);
+      }
+    }
+  }
+  __syncwarp();
+  if (above >= topk || above + mid < topk || mid > kMidCap) return false;
+  // Bracket codes per lane (register-resident).
+  constexpr uint32_t kMidPerLane = kMidCap / 32;
+  uint32_t mc[kMidPerLane];
+#pragma unroll
+  for (uint32_t u = 0; u < kMidPerLane; ++u) {
+    const uint32_t i = lane + 32 * u;
+    mc[u] = 0;
+    if (i < mid) mc[u] = codes[midbuf[i]];
+  }
+  uint32_t rank = topk - above, base = h_lo << 8, top = h_hi << 8 | 0xffu;
+#pragma unroll 1
+  while (top > base) {
+    const uint32_t width = top - base + 1, spread = (width - 1) / kMidBins;
+    const uint32_t shift = spread ? 32 - __clz(spread) : 0;
+#pragma unroll
+    for (uint32_t b = lane; b < kMidBins / 4; b += 32)
+      reinterpret_cast<uint4*>(h512)[b] = make_uint4(0, 0, 0, 0);
+    __syncwarp();
+#pragma unroll
+    for (uint32_t u = 0; u < kMidPerLane; ++u) {
+      if (u * 32 >= mid) break;
+      if (lane + 32 * u < mid && mc[u] >= base && mc[u] <= top)
+        atomicAdd(h512 + ((mc[u] - base) >> shift), 1u);
+    }
+    __syncwarp();
+    uint32_t above_bin;
+    const uint32_t chosen = find_bin_v<kMidBins>(h512, rank, above_bin);
+    rank -= above_bin;
+    base += chosen << shift;
+    top = min(base + (1u << shift) - 1, top);
+    __syncwarp();
+  }
+  const uint32_t th = base;
+  {
+    // Ballot placement (coalesced stores): greater codes fill [above, K -
+    // rank), ties [K - rank, K).
+    uint32_t pg = above, pt = topk - rank;
+#pragma unroll
+    for (uint32_t h = 0; h < 2; ++h) {
+      if (h * 16 * 32 >= mid) break;
+      int32_t pk[16];
+#pragma unroll
+      for (uint32_t j = 0; j < 16; ++j) {
+        const uint32_t u = h * 16 + j;
+        pk[j] = (lane + 32 * u < mid && mc[u] >= th)
+                    ? __ldcg(packed + midbuf[lane + 32 * u])
+                    : 0;
+      }
+#pragma unroll
+      for (uint32_t j = 0; j < 16; ++j) {
+        const uint32_t u = h * 16 + j;
+        if (u * 32 >= mid) break;
+        const bool v = lane + 32 * u < mid;
+        const bool g = v && mc[u] > th, tq = v && mc[u] == th;
+        const uint32_t mg = __ballot_sync(kFullMask, g),
+                       mt = __ballot_sync(kFullMask, tq);
+        const uint32_t ag = pg + __popc(mg & lt), at = pt + __popc(mt & lt);
+        if (g) out[ag] = candidate_decode_index(pk[j]);
+        if (tq && at < topk) out[at] = candidate_decode_index(pk[j]);
+        pg += __popc(mg);
+        pt += __popc(mt);
+      }
+    }
+  }
+  return true;
+}
+
+}  // namespace fused_mqa_topk::selection
diff --git a/csrc/hisparse_kernels.cu b/csrc/hisparse_kernels.cu
index 74df6553d020..cae27cb44a48 100644
--- a/csrc/hisparse_kernels.cu
+++ b/csrc/hisparse_kernels.cu
@@ -156,13 +156,25 @@ __device__ __forceinline__ int warp_inclusive_scan(int32_t* s_data, int lane_id,
   return __shfl_sync(0xffffffff, val, 31);
 }
 
+// Non-negative x only. Block sizes are powers of two in practice, and integer
+// division would otherwise dominate the resolver's per-position instructions.
+__device__ __forceinline__ int32_t block_div(int32_t x, int32_t block_size) {
+  return (block_size & (block_size - 1)) == 0 ? x >> (__ffs(block_size) - 1)
+                                              : x / block_size;
+}
+
+__device__ __forceinline__ int32_t block_mod(int32_t x, int32_t block_size) {
+  return (block_size & (block_size - 1)) == 0 ? x & (block_size - 1)
+                                              : x % block_size;
+}
+
 __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;
+                      block_div(slot, block_size)];
+  return static_cast<int64_t>(block) * block_size + block_mod(slot, block_size);
 }
 
 __device__ __forceinline__ void store_hot_index(
@@ -174,8 +186,8 @@ __device__ __forceinline__ void store_hot_index(
     attention_indices[index] =
         physical_row < 0
             ? -1
-            : (physical_row / hot_block_size) * attention_block_stride +
-                  physical_row % hot_block_size;
+            : block_div(physical_row, hot_block_size) * attention_block_stride +
+                  block_mod(physical_row, hot_block_size);
   }
 }
 
@@ -189,6 +201,8 @@ constexpr int kResidencyCounters = 5;
 // Per-request table entries per referenced top-k position, within the
 // device's shared-memory budget.
 constexpr int kHashLoadInverse = 4;
+// Top-k positions per thread the resolver loads before its request lookups.
+constexpr int kPrefetchedTopK = 2;
 
 // Translate a request-relative top-k position to its resident GPU row (-1 when
 // the page is not resident) and returns its host row (-1 when not host-backed).
@@ -205,7 +219,7 @@ __device__ __forceinline__ int32_t translate_topk_entry(
   resident_row = -1;
   if (source_block_table != nullptr) {
     const int32_t source_block =
-        token_index >= 0 ? token_index / source_block_size : -1;
+        token_index >= 0 ? block_div(token_index, source_block_size) : -1;
     if (request_row >= 0 && request_row < source_num_reqs &&
         source_block >= 0 && source_block < source_num_blocks) {
       const int32_t physical_block =
@@ -213,14 +227,14 @@ __device__ __forceinline__ int32_t translate_topk_entry(
                                  source_bt_stride +
                              source_block];
       g = physical_block > 0 ? physical_block * source_block_size +
-                                   token_index % source_block_size
+                                   block_mod(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;
+          token_index >= 0 ? block_div(token_index, resident_block_size) : -1;
       if (state_row >= 0 && state_row < resident_num_reqs &&
           resident_block >= 0 && resident_block < resident_num_blocks) {
         const int32_t physical_block =
@@ -229,7 +243,7 @@ __device__ __forceinline__ int32_t translate_topk_entry(
                                  resident_block];
         if (physical_block != resident_null_block && physical_block >= 0) {
           resident_row = physical_block * resident_block_size +
-                         token_index % resident_block_size;
+                         block_mod(token_index, resident_block_size);
         }
       }
     }
@@ -298,7 +312,9 @@ __device__ __forceinline__ void write_entry_index(
 //   s_done[nbc]              hot slots holding a union row (phase 2), then
 //                            positions phase 3 resolved, one word per chunk
 //   s_lru_out[hot_size]      int16, compacted slots: [hits fwd | evict bwd]
-// Valid global ids must be unique within each row.
+// Valid global ids must be unique within each row. A nonzero kBlockSize fixes
+// every block size at compile time.
+template <int kBlockSize>
 __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
@@ -319,69 +335,93 @@ __global__ __launch_bounds__(1024) void hisparse_resolve_residency_kernel(
     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 int64_t hot_table_stride, const int32_t hot_block_size_arg,
     const int32_t top_k, const int32_t hot_size, const int32_t hash_size,
     const int32_t max_union_rows, 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 source_block_size_arg, 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 int32_t resident_block_size_arg, 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 int64_t valid_count_stride,
+    const int32_t* __restrict__ num_valid_rows) {  // [1] or nullptr
+  const int32_t hot_block_size =
+      kBlockSize > 0 ? kBlockSize : hot_block_size_arg;
+  const int32_t source_block_size =
+      kBlockSize > 0 ? kBlockSize : source_block_size_arg;
+  const int32_t resident_block_size =
+      kBlockSize > 0 ? kBlockSize : resident_block_size_arg;
   const int NUM_WARPS = blockDim.x / kWarpSize;
   const int num_buffer_chunks = (hot_size + kWarpSize - 1) / kWarpSize;
 
-  // The first row of each contiguous run of a request's rows leads it. FULL
-  // CUDA-graph padding rows map to request 0 after the real rows, forming a
-  // second run, but carry no valid top-k: that run returns before phase 2
-  // touches the request's shared state.
+  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 int first_row = blockIdx.x;
-  const int request_row =
-      request_ids != nullptr ? request_ids[first_row] : first_row;
-  if (request_ids != nullptr && first_row > 0 &&
-      request_ids[first_row - 1] == request_row) {
-    return;
-  }
-  int end_row = first_row + 1;
-  if (request_ids != nullptr) {
-    while (end_row < static_cast<int>(gridDim.x) &&
-           request_ids[end_row] == request_row) {
-      ++end_row;
-    }
+
+  // Issued before the request lookups resolve so that an all-resident step
+  // pays few dependent memory round trips.
+  int32_t first_topk[kPrefetchedTopK];
+#pragma unroll
+  for (int k = 0; k < kPrefetchedTopK; ++k) {
+    const int i = tid + k * blockDim.x;
+    first_topk[k] = i < top_k ? global_indices[static_cast<int64_t>(first_row) *
+                                                   input_row_stride +
+                                               i]
+                              : -1;
   }
+  // FULL CUDA-graph padding rows past num_valid_rows belong to no request.
+  const int valid_rows = num_valid_rows != nullptr
+                             ? min(*num_valid_rows, static_cast<int>(gridDim.x))
+                             : static_cast<int>(gridDim.x);
+  const auto request_of = [&](const int row) -> int32_t {
+    if (request_ids == nullptr) return row;
+    const int32_t request = request_ids[row];
+    return row < valid_rows ? request : -1;
+  };
+  const int request_row = request_of(first_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.
+  // V2 publishes -1 for CUDA-graph padding rows. They share no request state,
+  // so every block masks its own row.
   if (state_row < 0) {
-    for (int row = first_row; row < end_row; ++row) {
-      for (int i = threadIdx.x; i < top_k; i += blockDim.x) {
-        const int64_t index = static_cast<int64_t>(row) * top_k + i;
-        hot_indices[index] = -1;
-        if (attention_indices != nullptr) {
-          attention_indices[static_cast<int64_t>(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>(row) * valid_count_stride] = 0;
+    for (int i = tid; i < top_k; i += blockDim.x) {
+      const int64_t index = static_cast<int64_t>(first_row) * top_k + i;
+      hot_indices[index] = -1;
+      if (attention_indices != nullptr) {
+        attention_indices[static_cast<int64_t>(first_row) *
+                              attention_row_stride +
+                          i] = -1;
       }
-      if (swap_counts != nullptr && threadIdx.x == 0) {
-        swap_counts[row] = 0;
+      if (resolved_global_indices != nullptr) {
+        resolved_global_indices[index] = -1;
       }
     }
+    if (valid_counts != nullptr && tid == 0) {
+      valid_counts[static_cast<int64_t>(first_row) * valid_count_stride] = 0;
+    }
+    if (swap_counts != nullptr && tid == 0) {
+      swap_counts[first_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;
+  // The first row of each contiguous run of a request's rows leads it.
+  if (request_ids != nullptr && first_row > 0 &&
+      request_of(first_row - 1) == request_row) {
+    return;
+  }
+  int end_row = first_row + 1;
+  if (request_ids != nullptr) {
+    while (end_row < static_cast<int>(gridDim.x) &&
+           request_of(end_row) == request_row) {
+      ++end_row;
+    }
+  }
 
   int32_t* row_dgi =
       device_global_indices + static_cast<int64_t>(state_row) * region_stride;
@@ -411,7 +451,8 @@ __global__ __launch_bounds__(1024) void hisparse_resolve_residency_kernel(
   // s_chunk_off holds per-row valid counts until phase 2.
   int32_t* s_row_valid = s_chunk_off;
   const int num_positions = (end_row - first_row) * top_k;
-  for (int position = tid; position < num_positions; position += blockDim.x) {
+  const auto resolve_position = [&](const int position,
+                                    const int32_t token_index) {
     const int row = first_row + position / top_k;
     const int i = position % top_k;
     int32_t* row_out = hot_indices + static_cast<int64_t>(row) * top_k;
@@ -422,12 +463,11 @@ __global__ __launch_bounds__(1024) void hisparse_resolve_residency_kernel(
             : nullptr;
     int32_t resident_row;
     const int32_t g = translate_topk_entry(
-        global_indices[static_cast<int64_t>(row) * input_row_stride + i],
-        request_row, state_row, source_block_table, resident_block_table,
-        host_rows, source_bt_stride, source_num_reqs, source_num_blocks,
-        source_block_size, resident_bt_stride, resident_num_reqs,
-        resident_num_blocks, resident_block_size, resident_null_block,
-        resident_row);
+        token_index, request_row, state_row, source_block_table,
+        resident_block_table, host_rows, source_bt_stride, source_num_reqs,
+        source_num_blocks, source_block_size, resident_bt_stride,
+        resident_num_reqs, resident_num_blocks, resident_block_size,
+        resident_null_block, resident_row);
     if (resolved_global_indices != nullptr) {
       resolved_global_indices[static_cast<int64_t>(row) * top_k + i] = g;
     }
@@ -446,6 +486,27 @@ __global__ __launch_bounds__(1024) void hisparse_resolve_residency_kernel(
       row_out[i] = encode_host_row(g);
       atomicAdd(&s_counters[3], 1);
     }
+  };
+#pragma unroll
+  for (int k = 0; k < kPrefetchedTopK; ++k) {
+    const int position = tid + k * blockDim.x;
+    if (position < num_positions) {
+      resolve_position(position,
+                       position < top_k
+                           ? first_topk[k]
+                           : global_indices[static_cast<int64_t>(
+                                                first_row + position / top_k) *
+                                                input_row_stride +
+                                            position % top_k]);
+    }
+  }
+  for (int position = tid + kPrefetchedTopK * blockDim.x;
+       position < num_positions; position += blockDim.x) {
+    resolve_position(
+        position,
+        global_indices[static_cast<int64_t>(first_row + position / top_k) *
+                           input_row_stride +
+                       position % top_k]);
   }
   __syncthreads();
   for (int row = first_row + tid; row < end_row; row += blockDim.x) {
@@ -879,30 +940,42 @@ __global__ void hisparse_gather_compact_kernel(
   }
 }
 
+// One block per (written token, index group): drop the group's hot copy of
+// the token's rewritten host slot.
 __global__ void hisparse_invalidate_written_slots_kernel(
-    int32_t* __restrict__ device_global_indices,
+    int32_t* const* __restrict__ device_global_indices,  // [num_groups]
     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* __restrict__ written_slots, 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) {
+  if (state_idx < 0 || state_idx >= num_state_rows || written_slot < 0 ||
+      written_slot > INT32_MAX) {
+    return;
+  }
+  const int32_t slot = static_cast<int32_t>(written_slot);
+  int32_t* row = device_global_indices[blockIdx.y] + state_idx * region_stride;
+  if (region_stride % 4 == 0 && (reinterpret_cast<uintptr_t>(row) & 15) == 0) {
+    const int4* row4 = reinterpret_cast<int4*>(row);
+    for (int64_t offset = threadIdx.x; offset < region_stride / 4;
+         offset += blockDim.x) {
+      const int4 v = row4[offset];
+      if (v.x == slot) row[4 * offset] = -1;
+      if (v.y == slot) row[4 * offset + 1] = -1;
+      if (v.z == slot) row[4 * offset + 2] = -1;
+      if (v.w == slot) row[4 * offset + 3] = -1;
+    }
     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) {
+    if (row[offset] == slot) {
       row[offset] = -1;
     }
   }
@@ -920,20 +993,22 @@ int64_t check_2d_rows(const torch::stable::Tensor& t, const char* name,
 }  // namespace
 
 void hisparse_invalidate_written_slots(
-    torch::stable::Tensor& device_global_indices,
+    torch::stable::Tensor const& device_global_indices_ptrs,
     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() &&
+    torch::stable::Tensor const& written_slots, int64_t num_state_rows,
+    int64_t region_stride) {
+  STD_TORCH_CHECK(device_global_indices_ptrs.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");
+      device_global_indices_ptrs.scalar_type() ==
+              torch::headeronly::ScalarType::Long &&
+          device_global_indices_ptrs.dim() == 1 &&
+          device_global_indices_ptrs.is_contiguous(),
+      "device_global_indices_ptrs must be a contiguous 1D int64 CUDA tensor "
+      "of [max_num_reqs, region_stride] int32 table addresses");
   STD_TORCH_CHECK(
       request_state_indices.scalar_type() ==
               torch::headeronly::ScalarType::Int &&
@@ -950,7 +1025,7 @@ void hisparse_invalidate_written_slots(
       "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();
+  const int device_index = device_global_indices_ptrs.get_device_index();
   STD_TORCH_CHECK(
       request_state_indices.get_device_index() == device_index &&
           req_id_per_token.get_device_index() == device_index &&
@@ -958,24 +1033,25 @@ void hisparse_invalidate_written_slots(
       "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) {
+  const int64_t num_groups = device_global_indices_ptrs.numel();
+  if (num_tokens == 0 || num_groups == 0 || num_state_rows == 0 ||
+      region_stride == 0) {
     return;
   }
-  STD_TORCH_CHECK(num_tokens <= INT32_MAX,
+  STD_TORCH_CHECK(num_tokens <= INT32_MAX && num_groups <= 65535,
                   "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>(),
+  const dim3 grid(static_cast<unsigned int>(num_tokens),
+                  static_cast<unsigned int>(num_groups));
+  hisparse_invalidate_written_slots_kernel<<<grid, kBlockSize, 0, stream>>>(
+      static_cast<int32_t* const*>(device_global_indices_ptrs.const_data_ptr()),
       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);
+      written_slots.const_data_ptr<int64_t>(), 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: ",
@@ -1004,7 +1080,8 @@ void hisparse_resolve_residency(
     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) {
+    int64_t resident_block_size, int64_t resident_null_block,
+    std::optional<torch::stable::Tensor> const& num_valid_rows) {
   STD_TORCH_CHECK(
       host_cache.device().is_cpu() && is_pinned_cpu_tensor(host_cache),
       "host_cache must be pinned CPU memory");
@@ -1098,6 +1175,18 @@ void hisparse_resolve_residency(
     source_num_reqs = static_cast<int32_t>(table.size(0));
     source_num_blocks = static_cast<int32_t>(table.size(1));
   }
+  const int32_t* num_valid_rows_ptr = nullptr;
+  if (num_valid_rows.has_value()) {
+    auto const& valid_rows = num_valid_rows.value();
+    STD_TORCH_CHECK(request_ids.has_value(),
+                    "num_valid_rows requires request_ids");
+    STD_TORCH_CHECK(
+        valid_rows.is_cuda() &&
+            valid_rows.scalar_type() == torch::headeronly::ScalarType::Int &&
+            valid_rows.numel() == 1,
+        "num_valid_rows must be a one-element int32 CUDA tensor");
+    num_valid_rows_ptr = valid_rows.const_data_ptr<int32_t>();
+  }
   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,
@@ -1278,20 +1367,31 @@ void hisparse_resolve_residency(
       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;
+  // Block arithmetic dominates the all-resident path; specialize the
+  // ubiquitous 64-row block.
+  constexpr int kFixedBlockSize = 64;
+  const bool fixed_block_size = hot_block_size == kFixedBlockSize &&
+                                (!source_block_table.has_value() ||
+                                 source_block_size == kFixedBlockSize) &&
+                                (!resident_block_table.has_value() ||
+                                 resident_block_size == kFixedBlockSize);
+  auto kernel = fixed_block_size
+                    ? hisparse_resolve_residency_kernel<kFixedBlockSize>
+                    : hisparse_resolve_residency_kernel<0>;
   if (smem_bytes > 48 * 1024) {
-    static std::vector<size_t> configured_smem_bytes;
+    static std::vector<size_t> configured_smem_bytes[2];
+    std::vector<size_t>& configured = configured_smem_bytes[fixed_block_size];
     const size_t device_index = hot_cache.get_device_index();
-    if (configured_smem_bytes.size() <= device_index) {
-      configured_smem_bytes.resize(device_index + 1, 0);
+    if (configured.size() <= device_index) {
+      configured.resize(device_index + 1, 0);
     }
-    if (configured_smem_bytes[device_index] < smem_bytes) {
+    if (configured[device_index] < smem_bytes) {
       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));
-      configured_smem_bytes[device_index] = smem_bytes;
+      configured[device_index] = smem_bytes;
     }
   }
   kernel<<<launch_rows, kBlockSize, smem_bytes, stream>>>(
@@ -1311,7 +1411,7 @@ void hisparse_resolve_residency(
       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);
+      attention_row_stride, valid_count_stride, num_valid_rows_ptr);
   const cudaError_t launch_error = cudaGetLastError();
   STD_TORCH_CHECK(launch_error == cudaSuccess,
                   "HiSparse residency kernel launch failed: ",
diff --git a/csrc/moe/grouped_topk_kernels.cu b/csrc/moe/grouped_topk_kernels.cu
index ea89aa38a1f0..443cdd14dc69 100644
--- a/csrc/moe/grouped_topk_kernels.cu
+++ b/csrc/moe/grouped_topk_kernels.cu
@@ -479,9 +479,9 @@ __device__ void topk_with_k2(T* output, T const* input, BiasT const* bias,
 template <typename T, typename BiasT, typename IdxT, ScoringFunc SF>
 __global__ void grouped_topk_fused_kernel(
     T* scores, float* topk_values, IdxT* topk_indices, BiasT const* bias,
-    int64_t const num_tokens, int64_t const num_experts, int64_t const n_group,
-    int64_t const topk_group, int64_t const topk, bool renormalize,
-    double routed_scaling_factor) {
+    bool const* is_padding, int64_t const num_tokens, int64_t const num_experts,
+    int64_t const n_group, int64_t const topk_group, int64_t const topk,
+    bool renormalize, double routed_scaling_factor) {
   int32_t const token_id = static_cast<int32_t>(blockIdx.x);
   if (token_id >= num_tokens) {
     return;
@@ -564,8 +564,11 @@ __global__ void grouped_topk_fused_kernel(
 
   if (!proceed) {
     for (int i = lane_id; i < topk_i32; i += WARP_SIZE) {
-      topk_indices[i] = static_cast<IdxT>(i);
-      topk_values[i] = 1.0f / static_cast<float>(topk_i32);
+      topk_indices[i] =
+          is_padding && is_padding[token_id] ? -1 : static_cast<IdxT>(i);
+      topk_values[i] = is_padding && is_padding[token_id]
+                           ? 0.0f
+                           : 1.0f / static_cast<float>(topk_i32);
     }
 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
     cudaTriggerProgrammaticLaunchCompletion();
@@ -624,8 +627,9 @@ __global__ void grouped_topk_fused_kernel(
   }
 
   if (lane_id < topk_i32) {
-    topk_indices[lane_id] = lane_idx;
-    topk_values[lane_id] = lane_unbiased * scale;
+    topk_indices[lane_id] = is_padding && is_padding[token_id] ? -1 : lane_idx;
+    topk_values[lane_id] =
+        is_padding && is_padding[token_id] ? 0.0f : lane_unbiased * scale;
   }
 
 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
@@ -638,8 +642,8 @@ template <typename T, typename BiasT, typename IdxT, ScoringFunc SF,
           int MaxNumTopExperts = DefaultMaxNumTopExperts>
 __global__ void grouped_topk_fused_small_expert_count_kernel(
     T* scores, float* topkValues, IdxT* topkIndices, BiasT const* routingBias,
-    int64_t const numTokens, int64_t const numGroup, int64_t const topkGroup,
-    int64_t const topk, int64_t const numExperts,
+    bool const* is_padding, int64_t const numTokens, int64_t const numGroup,
+    int64_t const topkGroup, int64_t const topk, int64_t const numExperts,
     int64_t const numExpertsPerGroup, bool const renormalize,
     double const routedScalingFactor) {
 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
@@ -838,8 +842,10 @@ __global__ void grouped_topk_fused_small_expert_count_kernel(
     }
     // store the topk scores and experts to output
     if (laneIdx < topk) {
-      topkValues[laneIdx] = finalScore;
-      topkIndices[laneIdx] = expertIdx;
+      topkValues[laneIdx] =
+          is_padding && is_padding[blockIdx.x] ? 0.0f : finalScore;
+      topkIndices[laneIdx] =
+          is_padding && is_padding[blockIdx.x] ? -1 : expertIdx;
     }
   }
 
@@ -890,7 +896,8 @@ __device__ __forceinline__ void write_outputs(
     cg::thread_block_tile<WARP_SIZE> const& warp, float lane_selection_score,
     float lane_unbiased, int32_t lane_expert, int32_t lane, int32_t token,
     int32_t topk, float* topk_values, IdxT* topk_indices, bool renormalize,
-    float routed_scaling_factor) {
+    float routed_scaling_factor, bool const* is_padding) {
+  bool const padded = is_padding && is_padding[token];
   bool const finite_selection =
       lane < topk && lane_selection_score != InvalidScore;
   lane_unbiased = finite_selection ? lane_unbiased : 0.0F;
@@ -912,8 +919,8 @@ __device__ __forceinline__ void write_outputs(
     }
 
     int64_t const output_index = int64_t{token} * topk + lane;
-    topk_values[output_index] = output;
-    topk_indices[output_index] = static_cast<IdxT>(lane_expert);
+    topk_values[output_index] = padded ? 0.0F : output;
+    topk_indices[output_index] = static_cast<IdxT>(padded ? -1 : lane_expert);
   }
 }
 
@@ -922,8 +929,8 @@ template <typename T, typename BiasT, typename IdxT, ScoringFunc SF,
 __global__ void __launch_bounds__(BlockDim)
     single_group_topk_block_kernel(T const* scores, float* topk_values,
                                    IdxT* topk_indices, BiasT const* bias,
-                                   int64_t num_experts, int64_t topk,
-                                   bool renormalize,
+                                   bool const* is_padding, int64_t num_experts,
+                                   int64_t topk, bool renormalize,
                                    float routed_scaling_factor,
                                    bool enable_pdl) {
   static constexpr int NumChunks = (MaxNumExperts + WARP_SIZE - 1) / WARP_SIZE;
@@ -1032,7 +1039,7 @@ __global__ void __launch_bounds__(BlockDim)
             : 0.0F;
     write_outputs(warp, lane_score, lane_unbiased, lane_expert, lane, token,
                   topk_i32, topk_values, topk_indices, renormalize,
-                  routed_scaling_factor);
+                  routed_scaling_factor, is_padding);
   } else {
     if (warp_id != 0) {
       return;
@@ -1060,7 +1067,7 @@ __global__ void __launch_bounds__(BlockDim)
             : 0.0F;
     write_outputs(warp, lane_score, lane_unbiased, lane_expert, lane, token,
                   topk_i32, topk_values, topk_indices, renormalize,
-                  routed_scaling_factor);
+                  routed_scaling_factor, is_padding);
   }
 
 #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
@@ -1095,9 +1102,9 @@ template <typename T, typename BiasT, typename IdxT, ScoringFunc SF,
 __global__ void __launch_bounds__(WarpTopKLaunchConfig<MaxNumExperts>::BlockDim)
     single_group_topk_warp_kernel(T const* scores, float* topk_values,
                                   IdxT* topk_indices, BiasT const* bias,
-                                  int64_t num_tokens, int64_t num_experts,
-                                  int64_t topk, bool renormalize,
-                                  float routed_scaling_factor,
+                                  bool const* is_padding, int64_t num_tokens,
+                                  int64_t num_experts, int64_t topk,
+                                  bool renormalize, float routed_scaling_factor,
                                   bool enable_pdl) {
   static constexpr int NumChunks = (MaxNumExperts + WARP_SIZE - 1) / WARP_SIZE;
   static constexpr int WarpBlockDim =
@@ -1166,7 +1173,7 @@ __global__ void __launch_bounds__(WarpTopKLaunchConfig<MaxNumExperts>::BlockDim)
     }
     write_outputs(warp, lane_score, lane_unbiased, lane_expert, lane, token,
                   topk_i32, topk_values, topk_indices, renormalize,
-                  routed_scaling_factor);
+                  routed_scaling_factor, is_padding);
   }
 
 #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
@@ -1198,9 +1205,10 @@ using PrecomputedSoftmaxBiasTiers =
 template <typename T, typename BiasT, typename IdxT, ScoringFunc SF,
           int MaxNumExperts, int MaxNumTopExperts>
 void launch(T* scores, float* topk_values, IdxT* topk_indices,
-            BiasT const* bias, int64_t num_tokens, int64_t num_experts,
-            int64_t topk, bool renormalize, double routed_scaling_factor,
-            bool enable_pdl, cudaLaunchConfig_t& config) {
+            BiasT const* bias, bool const* is_padding, int64_t num_tokens,
+            int64_t num_experts, int64_t topk, bool renormalize,
+            double routed_scaling_factor, bool enable_pdl,
+            cudaLaunchConfig_t& config) {
   config.dynamicSmemBytes = 0;
   bool const use_block_kernel =
       UseTunedBlockPath<MaxNumExperts, MaxNumTopExperts> ||
@@ -1213,8 +1221,8 @@ void launch(T* scores, float* topk_values, IdxT* topk_indices,
         &config,
         &single_group_topk_block_kernel<T, BiasT, IdxT, SF, MaxNumExperts,
                                         MaxNumTopExperts>,
-        scores, topk_values, topk_indices, bias, num_experts, topk, renormalize,
-        static_cast<float>(routed_scaling_factor), enable_pdl);
+        scores, topk_values, topk_indices, bias, is_padding, num_experts, topk,
+        renormalize, static_cast<float>(routed_scaling_factor), enable_pdl);
   } else {
     using WarpConfig = WarpTopKLaunchConfig<MaxNumExperts>;
     config.gridDim = WarpConfig::grid_dim(num_tokens);
@@ -1223,54 +1231,58 @@ void launch(T* scores, float* topk_values, IdxT* topk_indices,
         &config,
         &single_group_topk_warp_kernel<T, BiasT, IdxT, SF, MaxNumExperts,
                                        MaxNumTopExperts>,
-        scores, topk_values, topk_indices, bias, num_tokens, num_experts, topk,
-        renormalize, static_cast<float>(routed_scaling_factor), enable_pdl);
+        scores, topk_values, topk_indices, bias, is_padding, num_tokens,
+        num_experts, topk, renormalize,
+        static_cast<float>(routed_scaling_factor), enable_pdl);
   }
 }
 
 template <typename T, typename BiasT, typename IdxT, ScoringFunc SF>
-bool dispatch(TierList<>*, T*, float*, IdxT*, BiasT const*, int64_t, int64_t,
-              int64_t, bool, double, bool, cudaLaunchConfig_t&) {
+bool dispatch(TierList<>*, T*, float*, IdxT*, BiasT const*, bool const*,
+              int64_t, int64_t, int64_t, bool, double, bool,
+              cudaLaunchConfig_t&) {
   return false;
 }
 
 template <typename T, typename BiasT, typename IdxT, ScoringFunc SF,
           typename First, typename... Rest>
 bool dispatch(TierList<First, Rest...>*, T* scores, float* topk_values,
-              IdxT* topk_indices, BiasT const* bias, int64_t num_tokens,
-              int64_t num_experts, int64_t topk, bool renormalize,
-              double routed_scaling_factor, bool enable_pdl,
+              IdxT* topk_indices, BiasT const* bias, bool const* is_padding,
+              int64_t num_tokens, int64_t num_experts, int64_t topk,
+              bool renormalize, double routed_scaling_factor, bool enable_pdl,
               cudaLaunchConfig_t& config) {
   if (num_experts <= First::kExperts && topk <= First::kTopK) {
     launch<T, BiasT, IdxT, SF, First::kExperts, First::kTopK>(
-        scores, topk_values, topk_indices, bias, num_tokens, num_experts, topk,
-        renormalize, routed_scaling_factor, enable_pdl, config);
+        scores, topk_values, topk_indices, bias, is_padding, num_tokens,
+        num_experts, topk, renormalize, routed_scaling_factor, enable_pdl,
+        config);
     return true;
   }
   return dispatch<T, BiasT, IdxT, SF>(
       static_cast<TierList<Rest...>*>(nullptr), scores, topk_values,
-      topk_indices, bias, num_tokens, num_experts, topk, renormalize,
-      routed_scaling_factor, enable_pdl, config);
+      topk_indices, bias, is_padding, num_tokens, num_experts, topk,
+      renormalize, routed_scaling_factor, enable_pdl, config);
 }
 
 }  // namespace detail
 
 template <typename T, typename BiasT, typename IdxT, ScoringFunc SF>
 bool invoke(T* scores, float* topk_values, IdxT* topk_indices,
-            BiasT const* bias, int64_t num_tokens, int64_t num_experts,
-            int64_t topk, bool renormalize, double routed_scaling_factor,
-            bool enable_pdl, cudaLaunchConfig_t& config) {
+            BiasT const* bias, bool const* is_padding, int64_t num_tokens,
+            int64_t num_experts, int64_t topk, bool renormalize,
+            double routed_scaling_factor, bool enable_pdl,
+            cudaLaunchConfig_t& config) {
   static_assert(SF == SCORING_NONE || SF == SCORING_SIGMOID);
   if constexpr (SF == SCORING_SIGMOID) {
     return detail::dispatch<T, BiasT, IdxT, SF>(
         static_cast<detail::SigmoidBiasTiers*>(nullptr), scores, topk_values,
-        topk_indices, bias, num_tokens, num_experts, topk, renormalize,
-        routed_scaling_factor, enable_pdl, config);
+        topk_indices, bias, is_padding, num_tokens, num_experts, topk,
+        renormalize, routed_scaling_factor, enable_pdl, config);
   } else {
     return detail::dispatch<T, BiasT, IdxT, SF>(
         static_cast<detail::PrecomputedSoftmaxBiasTiers*>(nullptr), scores,
-        topk_values, topk_indices, bias, num_tokens, num_experts, topk,
-        renormalize, routed_scaling_factor, enable_pdl, config);
+        topk_values, topk_indices, bias, is_padding, num_tokens, num_experts,
+        topk, renormalize, routed_scaling_factor, enable_pdl, config);
   }
 }
 
@@ -1278,10 +1290,11 @@ bool invoke(T* scores, float* topk_values, IdxT* topk_indices,
 
 template <typename T, typename BiasT, typename IdxT, ScoringFunc SF>
 void invokeNoAuxTc(T* scores, float* topk_values, IdxT* topk_indices,
-                   BiasT const* bias, int64_t const num_tokens,
-                   int64_t const num_experts, int64_t const n_group,
-                   int64_t const topk_group, int64_t const topk,
-                   bool const renormalize, double const routed_scaling_factor,
+                   BiasT const* bias, bool const* is_padding,
+                   int64_t const num_tokens, int64_t const num_experts,
+                   int64_t const n_group, int64_t const topk_group,
+                   int64_t const topk, bool const renormalize,
+                   double const routed_scaling_factor,
                    const bool enable_pdl = false,
                    cudaStream_t const stream = 0) {
   cudaLaunchConfig_t config;
@@ -1293,8 +1306,9 @@ void invokeNoAuxTc(T* scores, float* topk_values, IdxT* topk_indices,
   config.attrs = attrs;
   if (n_group == 1 && topk_group == 1 &&
       single_group_topk::invoke<T, BiasT, IdxT, SF>(
-          scores, topk_values, topk_indices, bias, num_tokens, num_experts,
-          topk, renormalize, routed_scaling_factor, enable_pdl, config)) {
+          scores, topk_values, topk_indices, bias, is_padding, num_tokens,
+          num_experts, topk, renormalize, routed_scaling_factor, enable_pdl,
+          config)) {
     return;
   }
 
@@ -1346,9 +1360,9 @@ void invokeNoAuxTc(T* scores, float* topk_values, IdxT* topk_indices,
     config.blockDim = num_threads;
     config.dynamicSmemBytes = 0;
     cudaLaunchKernelEx(&config, kernel_instance, scores, topk_values,
-                       topk_indices, bias, num_tokens, n_group, topk_group,
-                       topk, num_experts, num_experts / n_group, renormalize,
-                       routed_scaling_factor);
+                       topk_indices, bias, is_padding, num_tokens, n_group,
+                       topk_group, topk, num_experts, num_experts / n_group,
+                       renormalize, routed_scaling_factor);
   } else {
     auto* kernel_instance = &grouped_topk_fused_kernel<T, BiasT, IdxT, SF>;
     // One block per token; one warp per group.
@@ -1366,18 +1380,20 @@ void invokeNoAuxTc(T* scores, float* topk_values, IdxT* topk_indices,
     size_t const extra_bytes = 16 + static_cast<size_t>(n_group) * sizeof(T);
     config.dynamicSmemBytes = internal_bytes + extra_bytes;
     cudaLaunchKernelEx(&config, kernel_instance, scores, topk_values,
-                       topk_indices, bias, num_tokens, num_experts, n_group,
-                       topk_group, topk, renormalize, routed_scaling_factor);
+                       topk_indices, bias, is_padding, num_tokens, num_experts,
+                       n_group, topk_group, topk, renormalize,
+                       routed_scaling_factor);
   }
 }
 
 #define INSTANTIATE_NOAUX_TC(T, BiasT, IdxT, SF)                             \
   template void invokeNoAuxTc<T, BiasT, IdxT, SF>(                           \
       T * scores, float* topk_values, IdxT* topk_indices, BiasT const* bias, \
-      int64_t const num_tokens, int64_t const num_experts,                   \
-      int64_t const n_group, int64_t const topk_group, int64_t const topk,   \
-      bool const renormalize, double const routed_scaling_factor,            \
-      const bool enable_pdl, cudaStream_t const stream);
+      bool const* is_padding, int64_t const num_tokens,                      \
+      int64_t const num_experts, int64_t const n_group,                      \
+      int64_t const topk_group, int64_t const topk, bool const renormalize,  \
+      double const routed_scaling_factor, const bool enable_pdl,             \
+      cudaStream_t const stream);
 
 INSTANTIATE_NOAUX_TC(float, float, int32_t, SCORING_SIGMOID);
 INSTANTIATE_NOAUX_TC(float, half, int32_t, SCORING_SIGMOID);
@@ -1403,7 +1419,8 @@ INSTANTIATE_NOAUX_TC(__nv_bfloat16, __nv_bfloat16, int32_t, SCORING_NONE);
 std::tuple<torch::stable::Tensor, torch::stable::Tensor> grouped_topk(
     torch::stable::Tensor const& scores, int64_t n_group, int64_t topk_group,
     int64_t topk, bool renormalize, double routed_scaling_factor,
-    torch::stable::Tensor const& bias, int64_t scoring_func = 0) {
+    torch::stable::Tensor const& bias, int64_t scoring_func,
+    const std::optional<torch::stable::Tensor>& is_padding) {
   const auto data_type = scores.scalar_type();
   const auto bias_type = bias.scalar_type();
   STD_TORCH_CHECK(scores.dim() == 2, "scores must be a 2D Tensor");
@@ -1426,6 +1443,20 @@ std::tuple<torch::stable::Tensor, torch::stable::Tensor> grouped_topk(
           scoring_func == vllm::moe::SCORING_SIGMOID,
       "scoring_func must be SCORING_NONE (0) or SCORING_SIGMOID (1)");
 
+  const bool* is_padding_ptr = nullptr;
+  if (is_padding.has_value()) {
+    STD_TORCH_CHECK(
+        is_padding->scalar_type() == torch::headeronly::ScalarType::Bool,
+        "is_padding must be bool");
+    STD_TORCH_CHECK(is_padding->dim() == 1 && is_padding->size(0) == num_tokens,
+                    "is_padding must have one entry per token");
+    STD_TORCH_CHECK(is_padding->is_contiguous(),
+                    "is_padding must be contiguous");
+    STD_TORCH_CHECK(is_padding->device() == scores.device(),
+                    "is_padding must be on the scores device");
+    is_padding_ptr = is_padding->const_data_ptr<bool>();
+  }
+
   // Always output float32 for topk_values (eliminates Python-side conversion)
   auto topk_values = torch::stable::new_empty(
       scores, {num_tokens, topk}, torch::headeronly::ScalarType::Float);
@@ -1447,8 +1478,8 @@ std::tuple<torch::stable::Tensor, torch::stable::Tensor> grouped_topk(
             reinterpret_cast<T*>(scores.mutable_data_ptr()),                  \
             reinterpret_cast<float*>(topk_values.mutable_data_ptr()),         \
             reinterpret_cast<IdxT*>(topk_indices.mutable_data_ptr()),         \
-            reinterpret_cast<BiasT const*>(bias.data_ptr()), num_tokens,      \
-            num_experts, n_group, topk_group, topk, renormalize,              \
+            reinterpret_cast<BiasT const*>(bias.data_ptr()), is_padding_ptr,  \
+            num_tokens, num_experts, n_group, topk_group, topk, renormalize,  \
             routed_scaling_factor, pdl_flag, stream);                         \
         break;                                                                \
       case vllm::moe::SCORING_SIGMOID:                                        \
@@ -1456,8 +1487,8 @@ std::tuple<torch::stable::Tensor, torch::stable::Tensor> grouped_topk(
             reinterpret_cast<T*>(scores.mutable_data_ptr()),                  \
             reinterpret_cast<float*>(topk_values.mutable_data_ptr()),         \
             reinterpret_cast<IdxT*>(topk_indices.mutable_data_ptr()),         \
-            reinterpret_cast<BiasT const*>(bias.data_ptr()), num_tokens,      \
-            num_experts, n_group, topk_group, topk, renormalize,              \
+            reinterpret_cast<BiasT const*>(bias.data_ptr()), is_padding_ptr,  \
+            num_tokens, num_experts, n_group, topk_group, topk, renormalize,  \
             routed_scaling_factor, pdl_flag, stream);                         \
         break;                                                                \
       default:                                                                \
diff --git a/csrc/moe/moe_ops.h b/csrc/moe/moe_ops.h
index 8ab24fcb33b8..484ae2a86957 100644
--- a/csrc/moe/moe_ops.h
+++ b/csrc/moe/moe_ops.h
@@ -20,6 +20,15 @@ void topk_sigmoid(torch::stable::Tensor& topk_weights,
                   double routed_scaling_factor,
                   std::optional<torch::stable::Tensor> is_padding);
 
+#ifdef USE_CUDA
+void muse_spark_routing(torch::stable::Tensor& topk_weights,
+                        torch::stable::Tensor& topk_ids,
+                        const torch::stable::Tensor& logits,
+                        const torch::stable::Tensor& bias, double multiplier,
+                        double route_scale, double division_eps,
+                        const std::optional<torch::stable::Tensor>& is_padding);
+#endif
+
 void topk_softplus_sqrt(
     torch::stable::Tensor& topk_weights, torch::stable::Tensor& topk_indices,
     torch::stable::Tensor& token_expert_indices,
@@ -82,7 +91,8 @@ torch::stable::Tensor moe_wna16_gemm(
 std::tuple<torch::stable::Tensor, torch::stable::Tensor> grouped_topk(
     const torch::stable::Tensor& scores, int64_t n_group, int64_t topk_group,
     int64_t topk, bool renormalize, double routed_scaling_factor,
-    const torch::stable::Tensor& bias, int64_t scoring_func);
+    const torch::stable::Tensor& bias, int64_t scoring_func,
+    const std::optional<torch::stable::Tensor>& is_padding);
 #endif
 
 bool moe_permute_unpermute_supported();
diff --git a/csrc/moe/muse_spark_routing.cu b/csrc/moe/muse_spark_routing.cu
new file mode 100644
index 000000000000..0c77237be32f
--- /dev/null
+++ b/csrc/moe/muse_spark_routing.cu
@@ -0,0 +1,346 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+// Adapted from SGLang's Muse Spark CUDA routing implementation.
+
+#include <cuda_runtime.h>
+#include <torch/csrc/stable/accelerator.h>
+#include <torch/csrc/stable/tensor.h>
+#include <torch/headeronly/core/ScalarType.h>
+#include <torch/headeronly/util/Exception.h>
+#include <limits>
+#include <optional>
+
+#include "dispatch_utils.h"
+#include "torch_utils.h"
+
+namespace vllm {
+namespace moe {
+
+__device__ __forceinline__ float warp_max(float value) {
+#if defined(__CUDA_ARCH_FAMILY_SPECIFIC__) && __CUDA_ARCH__ >= 1000 && \
+    __CUDA_ARCH__ < 1100
+  float result;
+  asm("redux.sync.max.f32 %0, %1, %2;"
+      : "=f"(result)
+      : "f"(value), "r"(0xffffffffu));
+  return result;
+#else
+  #pragma unroll
+  for (int offset = 16; offset; offset /= 2)
+    value = fmaxf(value, __shfl_xor_sync(0xffffffffu, value, offset));
+  return value;
+#endif
+}
+
+inline constexpr float kNegInf = -__builtin_huge_valf();
+// fp32(log2(e)) == 0x3FB8AA3B, the constant Triton folds into tl.exp.
+inline constexpr float kLog2e = 1.4426950408889634f;
+
+__device__ __forceinline__ float ex2_approx(float x) {
+  float r;
+  asm("ex2.approx.f32 %0, %1;" : "=f"(r) : "f"(x));
+  return r;
+}
+
+// Triton lowers `a / b` on fp32 to div.full.f32, not the div.rn.f32 nvcc emits.
+__device__ __forceinline__ float div_full(float a, float b) {
+  float r;
+  asm("div.full.f32 %0, %1, %2;" : "=f"(r) : "f"(a), "f"(b));
+  return r;
+}
+
+// triton/language/standard.py:50 -- `1 / (1 + exp(-x))`, lowered as
+// sub.f32 / mul.f32 / ex2.approx.f32 / add.f32 / div.full.f32.
+__device__ __forceinline__ float triton_sigmoid(float x) {
+  return div_full(1.0f, ex2_approx((0.0f - x) * kLog2e) + 1.0f);
+}
+
+// kExpertsPerLane must be a multiple of 4 (128-bit row loads) and covers
+// kExpertsPerLane * 32 experts. Lane `l` holds expert v*128 + l*4 + j at
+// register slot v*4 + j, monotone in the slot -- both the stable sort and the
+// group-ballot tie-break rely on that to reproduce tl.argmax.
+template <int kWarpsPerBlock_>
+struct RoutingTrait {
+  static constexpr int kTopK = 8;
+  static constexpr int kExpertsPerLane = 8;
+  static constexpr int kVecsPerLane = 2;
+  static constexpr int kWarpsPerBlock = kWarpsPerBlock_;
+  static constexpr int kMaxExperts = 256;
+};
+
+template <typename Index>
+struct RoutingParams {
+  const float* __restrict__ gating_output;
+  const float* __restrict__ bias;
+  float* __restrict__ topk_weights;
+  Index* __restrict__ topk_ids;
+  const bool* __restrict__ is_padding;
+  int64_t padding_stride;
+  int num_tokens;
+  int num_experts;
+  int64_t stride_gm;
+  int64_t stride_wm;
+  int64_t stride_im;
+  float router_output_multiplier;
+  float route_scale;
+  float division_eps;
+};
+
+template <typename Trait, bool kUsePDL, typename Index>
+__global__ void routing_kernel(
+    const RoutingParams<Index> __grid_constant__ params) {
+  constexpr int kEPL = Trait::kExpertsPerLane;
+  constexpr int kTopK = Trait::kTopK;
+  constexpr int kGroupExperts = 4 * 32;
+
+  const int lane = threadIdx.x % 32;
+  const int token =
+      blockIdx.x * Trait::kWarpsPerBlock + static_cast<int>(threadIdx.x / 32);
+  if (token >= params.num_tokens) return;
+
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
+  if constexpr (kUsePDL) cudaGridDependencySynchronize();
+#endif
+  if (params.is_padding != nullptr &&
+      params.is_padding[token * params.padding_stride]) {
+    if (lane < kTopK) {
+      params.topk_ids[token * params.stride_im + lane] = -1;
+      params.topk_weights[token * params.stride_wm + lane] = 0.0f;
+    }
+    return;
+  }
+
+  const float* __restrict__ row =
+      params.gating_output + static_cast<int64_t>(token) * params.stride_gm;
+  float score[kEPL];
+  int expert[kEPL];
+
+#pragma unroll
+  for (int v = 0; v < Trait::kVecsPerLane; ++v) {
+    const int base = v * kGroupExperts + lane * 4;
+    float4 g = make_float4(kNegInf, kNegInf, kNegInf, kNegInf);
+    if (base < params.num_experts)
+      g = *reinterpret_cast<const float4*>(row + base);
+    score[v * 4 + 0] = g.x;
+    score[v * 4 + 1] = g.y;
+    score[v * 4 + 2] = g.z;
+    score[v * 4 + 3] = g.w;
+#pragma unroll
+    for (int j = 0; j < 4; ++j) expert[v * 4 + j] = base + j;
+  }
+
+#pragma unroll
+  for (int s = 0; s < kEPL; ++s) score[s] *= params.router_output_multiplier;
+
+#pragma unroll
+  for (int s = 0; s < kEPL; ++s) score[s] = triton_sigmoid(score[s]);
+
+  float biased[kEPL];
+#pragma unroll
+  for (int v = 0; v < Trait::kVecsPerLane; ++v) {
+    const int base = v * kGroupExperts + lane * 4;
+    float4 b = make_float4(0.0f, 0.0f, 0.0f, 0.0f);
+    if (base < params.num_experts)
+      b = *reinterpret_cast<const float4*>(params.bias + base);
+    biased[v * 4 + 0] = b.x;
+    biased[v * 4 + 1] = b.y;
+    biased[v * 4 + 2] = b.z;
+    biased[v * 4 + 3] = b.w;
+  }
+
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
+  if constexpr (kUsePDL) cudaTriggerProgrammaticLaunchCompletion();
+#endif
+
+  float key[kEPL];
+#pragma unroll
+  for (int s = 0; s < kEPL; ++s) {
+    if (expert[s] < params.num_experts) {
+      key[s] = biased[s] + score[s];
+    } else {
+      score[s] = 0.0f;
+      key[s] = kNegInf;
+    }
+  }
+
+  // Must be a stable sort: equal keys keep their slot order and slots are
+  // monotone in expert id, which is tl.argmax's lowest-index-wins tie-break.
+#pragma unroll
+  for (int r = 0; r < kEPL; ++r) {
+#pragma unroll
+    for (int i = r % 2; i + 1 < kEPL; i += 2) {
+      const bool swap = key[i + 1] > key[i];
+      const float k0 = swap ? key[i + 1] : key[i];
+      const float k1 = swap ? key[i] : key[i + 1];
+      const int e0 = swap ? expert[i + 1] : expert[i];
+      const int e1 = swap ? expert[i] : expert[i + 1];
+      const float s0 = swap ? score[i + 1] : score[i];
+      const float s1 = swap ? score[i] : score[i + 1];
+      key[i] = k0;
+      key[i + 1] = k1;
+      expert[i] = e0;
+      expert[i + 1] = e1;
+      score[i] = s0;
+      score[i + 1] = s1;
+    }
+  }
+
+  float score_sum = 0.0f;
+  int out_id = 0;
+  float out_score = 0.0f;
+
+#pragma unroll
+  for (int k = 0; k < kTopK; ++k) {
+    const float round_max = warp_max(key[0]);
+    const bool is_head = key[0] == round_max;
+
+    // Lowest expert id among the tied heads: within one kGroupExperts-wide
+    // group the id is monotone in the lane, so the lowest set bit of the lowest
+    // non-empty group's ballot owns the round.
+    unsigned owners = 0;
+#pragma unroll
+    for (int g = Trait::kVecsPerLane - 1; g >= 0; --g) {
+      const unsigned m =
+          __ballot_sync(0xffffffffu, is_head && expert[0] / kGroupExperts == g);
+      owners = (m != 0u) ? m : owners;
+    }
+    const int owner = __ffs(owners) - 1;
+    const int winner = __shfl_sync(0xffffffffu, expert[0], owner);
+    const float selected = __shfl_sync(0xffffffffu, score[0], owner);
+
+    // Tapered pop: after round k a lane can be popped at most kTopK-1-k more
+    // times, so slots past that can never surface and need not be maintained.
+    const bool pop = expert[0] == winner;
+#pragma unroll
+    for (int s = 0; s + 1 < kEPL; ++s)
+      if (s + k <= kTopK - 2) {
+        key[s] = pop ? key[s + 1] : key[s];
+        expert[s] = pop ? expert[s + 1] : expert[s];
+        score[s] = pop ? score[s + 1] : score[s];
+      }
+    // Only reachable when kTopK > kEPL: a lane popped kEPL times is exhausted
+    // and must stop winning rounds.
+    if (kEPL - 1 + k <= kTopK - 2)
+      key[kEPL - 1] = pop ? kNegInf : key[kEPL - 1];
+
+    score_sum += selected;
+    if (lane == k) {
+      out_id = winner;
+      out_score = selected;
+    }
+  }
+
+  score_sum += params.division_eps;
+  const float inv_sum_scaled = div_full(params.route_scale, score_sum);
+  if (lane < kTopK) {
+    const Index id = out_id;
+    const float weight = out_score * inv_sum_scaled;
+    params.topk_ids[token * params.stride_im + lane] = id;
+    params.topk_weights[token * params.stride_wm + lane] = weight;
+  }
+}
+
+}  // namespace moe
+}  // namespace vllm
+
+namespace {
+using Tensor = torch::stable::Tensor;
+using ScalarType = torch::headeronly::ScalarType;
+
+template <int Warps, bool PDL, typename Index>
+void launch_routing(const vllm::moe::RoutingParams<Index>& params,
+                    cudaStream_t stream) {
+  using Trait = vllm::moe::RoutingTrait<Warps>;
+  cudaLaunchConfig_t config{};
+  config.gridDim = dim3((params.num_tokens + Warps - 1) / Warps);
+  config.blockDim = dim3(32 * Warps);
+  config.stream = stream;
+  cudaLaunchAttribute attr{};
+  if constexpr (PDL) {
+    attr.id = cudaLaunchAttributeProgrammaticStreamSerialization;
+    attr.val.programmaticStreamSerializationAllowed = 1;
+    config.attrs = &attr;
+    config.numAttrs = 1;
+  }
+  auto err = cudaLaunchKernelEx(
+      &config, vllm::moe::routing_kernel<Trait, PDL, Index>, params);
+  STD_TORCH_CHECK(err == cudaSuccess, "Muse Spark routing launch failed: ",
+                  cudaGetErrorString(err));
+}
+}  // namespace
+
+void muse_spark_routing(Tensor& topk_weights, Tensor& topk_ids,
+                        const Tensor& logits, const Tensor& bias,
+                        double multiplier, double route_scale,
+                        double division_eps,
+                        const std::optional<Tensor>& is_padding) {
+  STD_TORCH_CHECK(logits.device().is_cuda() && logits.dim() == 2 &&
+                      logits.scalar_type() == ScalarType::Float &&
+                      logits.size(1) == 256 && logits.stride(1) == 1 &&
+                      logits.stride(0) >= 256 && logits.stride(0) % 4 == 0 &&
+                      reinterpret_cast<uintptr_t>(logits.data_ptr()) % 16 == 0,
+                  "logits must be aligned FP32 CUDA rows of 256 experts");
+  STD_TORCH_CHECK(logits.size(0) <= std::numeric_limits<int>::max(),
+                  "Too many routing tokens");
+  const int tokens = logits.size(0);
+  STD_TORCH_CHECK(bias.device() == logits.device() && bias.dim() == 1 &&
+                      bias.size(0) == 256 &&
+                      bias.scalar_type() == ScalarType::Float &&
+                      bias.stride(0) == 1 &&
+                      reinterpret_cast<uintptr_t>(bias.data_ptr()) % 16 == 0,
+                  "bias must be an aligned FP32 tensor of 256 experts on the "
+                  "logits device");
+  for (const auto* output : {&topk_weights, &topk_ids}) {
+    STD_TORCH_CHECK(
+        output->device() == logits.device() && output->dim() == 2 &&
+            output->size(0) == tokens && output->size(1) == 8 &&
+            output->stride(1) == 1 && output->stride(0) >= 8,
+        "routing outputs must have shape [tokens, 8] with contiguous rows");
+  }
+  STD_TORCH_CHECK(
+      topk_weights.scalar_type() == ScalarType::Float &&
+          (topk_ids.scalar_type() == ScalarType::Int ||
+           topk_ids.scalar_type() == ScalarType::Long),
+      "routing outputs must be FP32 weights and int32 or int64 IDs");
+  if (is_padding) {
+    STD_TORCH_CHECK(is_padding->device() == logits.device() &&
+                        is_padding->scalar_type() == ScalarType::Bool &&
+                        is_padding->dim() == 1 && is_padding->size(0) == tokens,
+                    "is_padding must be a bool tensor of shape [tokens] on the "
+                    "logits device");
+  }
+  if (tokens == 0) return;
+  torch::stable::accelerator::DeviceGuard guard(logits.get_device_index());
+  const cudaStream_t stream =
+      get_current_cuda_stream(logits.get_device_index());
+  VLLM_STABLE_DISPATCH_IDX_TYPES(
+      topk_ids.scalar_type(), "muse_spark_routing", [&] {
+        vllm::moe::RoutingParams<idx_t> params{
+            .gating_output = static_cast<const float*>(logits.data_ptr()),
+            .bias = static_cast<const float*>(bias.data_ptr()),
+            .topk_weights = static_cast<float*>(topk_weights.data_ptr()),
+            .topk_ids = static_cast<idx_t*>(topk_ids.data_ptr()),
+            .is_padding = is_padding
+                              ? static_cast<const bool*>(is_padding->data_ptr())
+                              : nullptr,
+            .padding_stride = is_padding ? is_padding->stride(0) : 0,
+            .num_tokens = tokens,
+            .num_experts = 256,
+            .stride_gm = logits.stride(0),
+            .stride_wm = topk_weights.stride(0),
+            .stride_im = topk_ids.stride(0),
+            .router_output_multiplier = static_cast<float>(multiplier),
+            .route_scale = static_cast<float>(route_scale),
+            .division_eps = static_cast<float>(division_eps),
+        };
+        if (tokens <= 2048) {
+          const bool pdl = get_device_prop()->major >= 9;
+          if (pdl)
+            launch_routing<1, true>(params, stream);
+          else
+            launch_routing<1, false>(params, stream);
+        } else {
+          launch_routing<4, false>(params, stream);
+        }
+      });
+}
diff --git a/csrc/moe/topk_softmax_kernels.cu b/csrc/moe/topk_softmax_kernels.cu
index abacf6502b1c..b687e95f75e8 100644
--- a/csrc/moe/topk_softmax_kernels.cu
+++ b/csrc/moe/topk_softmax_kernels.cu
@@ -627,6 +627,8 @@ void topkGatingLauncherHelper(const InputType* input, const bool* finished, floa
 
     dim3 block_dim(WARP_SIZE_PARAM, WARPS_PER_TB);
 #ifndef USE_ROCM
+    // Pre-SM90 cannot use PDL, but the attribute still adds a programmatic edge to captured CUDA graphs.
+    static const bool use_pdl = get_device_prop()->major >= 9;
     cudaLaunchConfig_t config{};
     config.gridDim = dim3(num_blocks);
     config.blockDim = block_dim;
@@ -635,7 +637,7 @@ void topkGatingLauncherHelper(const InputType* input, const bool* finished, floa
     attr.id = cudaLaunchAttributeProgrammaticStreamSerialization;
     attr.val.programmaticStreamSerializationAllowed = 1;
     config.attrs = &attr;
-    config.numAttrs = 1;
+    config.numAttrs = use_pdl ? 1 : 0;
     auto err = cudaLaunchKernelEx(
         &config,
         topkGating<VPT, EXPERTS, WARPS_PER_TB, BYTES_PER_LDG, WARP_SIZE_PARAM, IndType, InputType, SF>,
diff --git a/csrc/moe/torch_bindings.cpp b/csrc/moe/torch_bindings.cpp
index 9e8907afd714..3e2ee68a5eaa 100644
--- a/csrc/moe/torch_bindings.cpp
+++ b/csrc/moe/torch_bindings.cpp
@@ -23,6 +23,13 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_moe_C, m) {
       "bias, Tensor? input_ids, Tensor? tid2eid, Tensor? is_padding, "
       "Tensor? bias_vl=None, int image_sentinel_lo=0) -> ()");
 
+#ifdef USE_CUDA
+  m.def(
+      "muse_spark_routing(Tensor! topk_weights, Tensor! topk_ids, "
+      "Tensor logits, Tensor bias, float multiplier, float route_scale, "
+      "float division_eps, Tensor? is_padding) -> ()");
+#endif
+
   // Calculate the result of moe by summing up the partial results
   // from all selected experts. topk_ids/expert_map are optional and, when
   // both given, enable pad-aware reduce that skips (token, expert)
@@ -144,7 +151,8 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_moe_C, m) {
   m.def(
       "grouped_topk(Tensor scores, int n_group, int "
       "topk_group, int topk, bool renormalize, float "
-      "routed_scaling_factor, Tensor bias, int scoring_func) -> (Tensor, "
+      "routed_scaling_factor, Tensor bias, int scoring_func, Tensor? "
+      "is_padding=None) -> (Tensor, "
       "Tensor)");
 #endif
 }
@@ -152,6 +160,9 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_moe_C, m) {
 STABLE_TORCH_LIBRARY_IMPL(_moe_C, CUDA, m) {
   m.impl("topk_softmax", TORCH_BOX(&topk_softmax));
   m.impl("topk_sigmoid", TORCH_BOX(&topk_sigmoid));
+#ifdef USE_CUDA
+  m.impl("muse_spark_routing", TORCH_BOX(&muse_spark_routing));
+#endif
   m.impl("topk_softplus_sqrt", TORCH_BOX(&topk_softplus_sqrt));
   m.impl("moe_sum", TORCH_BOX(&moe_sum));
   m.impl("moe_align_block_size", TORCH_BOX(&moe_align_block_size));
diff --git a/csrc/ops.h b/csrc/ops.h
index 401f6b4996ab..802de7f550e7 100644
--- a/csrc/ops.h
+++ b/csrc/ops.h
@@ -413,7 +413,18 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
     const std::string& kv_cache_dtype, bool skip_index_branch,
     std::optional<torch::stable::Tensor> q_fp8_out, double q_fp8_scale,
     std::optional<torch::stable::Tensor> kv_k_scale,
-    std::optional<torch::stable::Tensor> kv_v_scale);
+    std::optional<torch::stable::Tensor> kv_v_scale, bool write_qkv);
+
+void fused_minimax_m3_dspark_context_kv(
+    torch::stable::Tensor const& all_kv,
+    torch::stable::Tensor const& k_norm_weight,
+    torch::stable::Tensor const& cos_sin_cache,
+    torch::stable::Tensor const& positions,
+    torch::stable::Tensor const& slot_mapping, torch::stable::Tensor& kv_cache,
+    torch::stable::Tensor const& cache_offsets,
+    torch::stable::Tensor const& k_scale,
+    torch::stable::Tensor const& scale_offsets, int64_t rotary_dim, double eps,
+    const std::string& kv_cache_dtype);
 
 #ifdef VLLM_ENABLE_FUSED_KDA_DECODE
 void fused_kda_decode(
@@ -746,13 +757,15 @@ void hisparse_resolve_residency(
     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);
+    int64_t resident_block_size, int64_t resident_null_block,
+    std::optional<torch::stable::Tensor> const& num_valid_rows);
 
 void hisparse_invalidate_written_slots(
-    torch::stable::Tensor& device_global_indices,
+    torch::stable::Tensor const& device_global_indices_ptrs,
     torch::stable::Tensor const& request_state_indices,
     torch::stable::Tensor const& req_id_per_token,
-    torch::stable::Tensor const& written_slots);
+    torch::stable::Tensor const& written_slots, int64_t num_state_rows,
+    int64_t region_stride);
 
 void hisparse_gather_plan(
     torch::stable::Tensor const& host_cache, torch::stable::Tensor& hot_cache,
@@ -871,3 +884,12 @@ void ngram_compute_n_gram_ids(
     torch::stable::Tensor& exclusive_req_len_sums,
     torch::stable::Tensor& ne_token_table, torch::stable::Tensor& row_indices,
     torch::stable::Tensor& column_starts, torch::stable::Tensor& n_gram_ids);
+
+void fused_mqa_topk_prefill(
+    torch::stable::Tensor q, torch::stable::Tensor k,
+    torch::stable::Tensor scales, torch::stable::Tensor weights,
+    torch::stable::Tensor starts, torch::stable::Tensor ends,
+    torch::stable::Tensor out, torch::stable::Tensor counter,
+    torch::stable::Tensor values, torch::stable::Tensor indices,
+    torch::stable::Tensor retained, torch::stable::Tensor retained_indices,
+    torch::stable::Tensor counts, torch::stable::Tensor defer);
diff --git a/csrc/persistent_topk.cuh b/csrc/persistent_topk.cuh
index bd18483b48d0..30ffed044561 100644
--- a/csrc/persistent_topk.cuh
+++ b/csrc/persistent_topk.cuh
@@ -65,8 +65,9 @@ struct RadixRowState {
 // Kernel parameters
 // ============================================================================
 
+template <typename InT>
 struct PersistentTopKParams {
-  const float* __restrict__ input;      // [num_rows, stride]
+  const InT* __restrict__ input;        // [num_rows, stride]
   int32_t* __restrict__ output;         // [num_rows, top_k]
   const int32_t* __restrict__ lengths;  // [num_rows]
   RadixRowState* row_states;            // large path: per-group state
@@ -150,9 +151,9 @@ __device__ __forceinline__ void select_coarse_bin(
 // kernel
 // Returns where exact recovery resumes when the threshold bin overflows the
 // candidate buffer; .needed is false otherwise.
-template <int TopK, uint32_t VecSize>
+template <int TopK, uint32_t VecSize, typename InT>
 __device__ __noinline__ overflow::MediumResume histogram_medium_topk(
-    const float* __restrict__ logits, int* __restrict__ output_indices,
+    const InT* __restrict__ logits, int* __restrict__ output_indices,
     int seq_len) {
   // All shared state lives in dynamic shared memory to avoid static
   extern __shared__ char medium_smem[];
@@ -291,7 +292,8 @@ __device__ __noinline__ overflow::MediumResume histogram_medium_topk(
     if (remaining_k == 0) {
       for (int i = thread_id; i < num_buffered; i += kThreadsPerBlock) {
         const int idx = buffered_indices[src_buffer][i];
-        const uint32_t fp32_bits = convert_to_uint32_v2(logits[idx]);
+        const uint32_t fp32_bits =
+            convert_to_uint32_v2(topk_histogram_4096::to_float(logits[idx]));
         const int bin = (fp32_bits >> bit_offset) & 0xFF;
         if (bin > threshold_bin) {
           const int output_pos = atomicAdd(&shared_output_count, 1);
@@ -310,7 +312,7 @@ __device__ __noinline__ overflow::MediumResume histogram_medium_topk(
 
     for (int i = thread_id; i < num_buffered; i += kThreadsPerBlock) {
       const int idx = buffered_indices[src_buffer][i];
-      const float logit_value = logits[idx];
+      const float logit_value = topk_histogram_4096::to_float(logits[idx]);
       const uint32_t fp32_bits = convert_to_uint32_v2(logit_value);
       const int bin = (fp32_bits >> bit_offset) & 0xFF;
       if (bin > threshold_bin) {
@@ -400,8 +402,8 @@ __device__ __forceinline__ void wait_ge(int* ptr, int target_val,
 // Adapted from https://github.com/flashinfer-ai/flashinfer/pull/2215
 // ============================================================================
 
-template <int TopK, uint32_t VEC_SIZE>
-__device__ void radix_topk(const float* __restrict__ row_input,
+template <int TopK, uint32_t VEC_SIZE, typename InT>
+__device__ void radix_topk(const InT* __restrict__ row_input,
                            int32_t* __restrict__ row_output, uint32_t seq_len,
                            uint32_t my_chunk_start, uint32_t chunk_size,
                            uint32_t* local_histogram, uint32_t* suffix_sum,
@@ -421,24 +423,26 @@ __device__ void radix_topk(const float* __restrict__ row_input,
 
     for (uint32_t i = tx * VEC_SIZE; i < aligned_size;
          i += kThreadsPerBlock * VEC_SIZE) {
-      const float* src = row_input + my_chunk_start + i;
+      const InT* src = row_input + my_chunk_start + i;
       if constexpr (VEC_SIZE == 4) {
-        float4 v = *reinterpret_cast<const float4*>(src);
+        float4 v = topk_histogram_4096::load_score4(src);
         shared_ordered[i] = convert_to_uint32_v2(v.x);
         shared_ordered[i + 1] = convert_to_uint32_v2(v.y);
         shared_ordered[i + 2] = convert_to_uint32_v2(v.z);
         shared_ordered[i + 3] = convert_to_uint32_v2(v.w);
       } else if constexpr (VEC_SIZE == 2) {
-        float2 v = *reinterpret_cast<const float2*>(src);
+        float2 v = topk_histogram_4096::load_score2(src);
         shared_ordered[i] = convert_to_uint32_v2(v.x);
         shared_ordered[i + 1] = convert_to_uint32_v2(v.y);
       } else {
-        shared_ordered[i] = convert_to_uint32_v2(*src);
+        shared_ordered[i] =
+            convert_to_uint32_v2(topk_histogram_4096::to_float(*src));
       }
     }
     for (uint32_t i = aligned_size + tx; i < actual_chunk_size;
          i += kThreadsPerBlock) {
-      shared_ordered[i] = convert_to_uint32_v2(row_input[my_chunk_start + i]);
+      shared_ordered[i] = convert_to_uint32_v2(
+          topk_histogram_4096::to_float(row_input[my_chunk_start + i]));
     }
   }
   __syncthreads();
@@ -609,9 +613,9 @@ __device__ void radix_topk(const float* __restrict__ row_input,
 // Persistent kernel — short/medium/large paths with RadixTopK.
 // ============================================================================
 
-template <int TopK = 2048, uint32_t VEC_SIZE = 1>
+template <int TopK = 2048, uint32_t VEC_SIZE = 1, typename InT = float>
 __global__ void __launch_bounds__(kThreadsPerBlock, 2)
-    persistent_topk_kernel(PersistentTopKParams params) {
+    persistent_topk_kernel(PersistentTopKParams<InT> params) {
   const uint32_t tx = threadIdx.x;
   extern __shared__ uint8_t smem_raw[];
 
@@ -675,7 +679,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 2)
     const uint32_t seq_len =
         non_negative_len < row_bound ? non_negative_len : row_bound;
     int32_t* row_output = params.output + row_idx * params.top_k;
-    const float* row_input = params.input + row_idx * params.stride;
+    const InT* row_input = params.input + row_idx * params.stride;
 
     if (seq_len <= RADIX_THRESHOLD) {
       if (cta_in_group == 0) {
@@ -736,10 +740,11 @@ struct vec_t {
   FLASHINFER_INLINE T& operator[](size_t i) { return data[i]; }
   FLASHINFER_INLINE const T& operator[](size_t i) const { return data[i]; }
 
-  FLASHINFER_INLINE void cast_load(const T* ptr) {
+  template <typename InT>
+  FLASHINFER_INLINE void cast_load(const InT* ptr) {
 #pragma unroll
     for (size_t i = 0; i < N; ++i) {
-      data[i] = ptr[i];
+      data[i] = topk_histogram_4096::to_float(ptr[i]);
     }
   }
 };
@@ -802,7 +807,7 @@ __device__ __forceinline__ bool filtered_topk_row(
   extern __shared__ __align__(16) uint8_t dynamic_smem[];
   constexpr uint32_t BLOCK_SIZE = FILTERED_TOPK_BLOCK_THREADS;
   constexpr int SMEM_INPUT_SIZE = FILTERED_TOPK_SMEM_INPUT_SIZE;
-  using Traits = FilteredTopKTraits<DType, UseWideCoarse>;
+  using Traits = FilteredTopKTraits<float, UseWideCoarse>;
   constexpr int RADIX = 1 << Traits::REFINE_BITS;
   static_assert(RADIX == FilteredTopKStorage<MAX_K, UseWideCoarse>::RADIX);
   constexpr int COARSE_RADIX = Traits::COARSE_RADIX;
@@ -853,7 +858,7 @@ __device__ __forceinline__ bool filtered_topk_row(
   }
   __syncthreads();
 
-  vec_t<DType, VEC_SIZE> score_vec;
+  vec_t<float, VEC_SIZE> score_vec;
 
   const int aligned_length = (length / VEC_SIZE) * VEC_SIZE;
 #pragma unroll 2
@@ -868,7 +873,7 @@ __device__ __forceinline__ bool filtered_topk_row(
   }
   // Handle tail
   for (int i = aligned_length + tx; i < length; i += BLOCK_SIZE) {
-    const auto bin = Traits::ToCoarseKey(score[i]);
+    const auto bin = Traits::ToCoarseKey(hist4096::to_float(score[i]));
     atomicAdd(&s_coarse_histogram[bin], 1);
   }
   __syncthreads();
@@ -920,7 +925,8 @@ __device__ __forceinline__ bool filtered_topk_row(
     }
     // Handle tail
     for (int i = aligned_length + tx; i < length; i += BLOCK_SIZE) {
-      const auto bin = static_cast<int>(Traits::ToCoarseKey(score[i]));
+      const auto bin =
+          static_cast<int>(Traits::ToCoarseKey(hist4096::to_float(score[i])));
       if (bin > threshold_bin) {
         const auto pos = atomicAdd(&s_counter, 1);
         s_indices[pos] = i;
@@ -972,7 +978,7 @@ __device__ __forceinline__ bool filtered_topk_row(
     }
     // Handle tail
     for (int i = aligned_length + tx; i < length; i += BLOCK_SIZE) {
-      filter_and_add_to_histogram(score[i], i);
+      filter_and_add_to_histogram(hist4096::to_float(score[i]), i);
     }
     __syncthreads();
 
@@ -1007,7 +1013,8 @@ __device__ __forceinline__ bool filtered_topk_row(
         for (int i = tx; i < num_input; i += BLOCK_SIZE) {
           const auto idx = s_input_idx[r_idx][i];
           const auto bin =
-              (Traits::ToOrdered(score[idx]) >> offset) & (RADIX - 1);
+              (Traits::ToOrdered(hist4096::to_float(score[idx])) >> offset) &
+              (RADIX - 1);
           if (static_cast<int>(bin) > threshold) {
             const auto pos = atomicAdd(&s_counter, 1);
             s_indices[pos] = idx;
@@ -1023,7 +1030,7 @@ __device__ __forceinline__ bool filtered_topk_row(
         __syncthreads();
         for (int i = tx; i < num_input; i += BLOCK_SIZE) {
           const auto idx = s_input_idx[r_idx][i];
-          const auto raw_input = score[idx];
+          const float raw_input = hist4096::to_float(score[idx]);
           const auto bin =
               (Traits::ToOrdered(raw_input) >> offset) & (RADIX - 1);
           if (static_cast<int>(bin) > threshold) {
@@ -1100,7 +1107,7 @@ constexpr uint32_t gcd(uint32_t a, uint32_t b) {
 // Returns 1, 2, 4, or 8
 template <typename DType>
 constexpr int ComputeFilteredTopKVecSize(uint32_t max_len) {
-  constexpr int MAX_VEC = 16 / sizeof(DType);  // 4 for float32, 8 for fp16/bf16
+  constexpr int MAX_VEC = 4;
   // Use GCD to find largest power-of-2 divisor
   const uint32_t g = gcd(max_len, static_cast<uint32_t>(MAX_VEC));
   return static_cast<int>(g);
@@ -1114,7 +1121,7 @@ cudaError_t FilteredTopKRaggedTransform(const DType* input,
                                         uint32_t max_len,
                                         cudaStream_t stream = 0) {
   constexpr size_t smem_size = FILTERED_TOPK_SMEM_DYNAMIC;
-  constexpr int MAX_VEC = 16 / sizeof(DType);
+  constexpr int MAX_VEC = 4;
 
   dim3 grid(num_rows);
   dim3 block(FILTERED_TOPK_BLOCK_THREADS);
diff --git a/csrc/persistent_topk_overflow.cuh b/csrc/persistent_topk_overflow.cuh
index 9ae116998e89..306dcfe263a5 100644
--- a/csrc/persistent_topk_overflow.cuh
+++ b/csrc/persistent_topk_overflow.cuh
@@ -19,10 +19,10 @@ struct MediumResume {
   uint32_t remaining = 0;
 };
 
-template <int TopK, int BlockSize, int VecSize>
-__device__ __noinline__ void recover_medium(const float* scores,
-                                            int32_t* output, uint32_t length,
-                                            void* smem, uint32_t prefix,
+template <int TopK, int BlockSize, int VecSize, typename InT>
+__device__ __noinline__ void recover_medium(const InT* scores, int32_t* output,
+                                            uint32_t length, void* smem,
+                                            uint32_t prefix,
                                             uint32_t remaining) {
   topk_histogram_4096::exact_topk_rescan<TopK, BlockSize, VecSize, 1, 4096>(
       scores, output, length, smem, prefix, remaining);
@@ -70,7 +70,9 @@ __device__ bool try_sampled_exact_topk(const DType* __restrict__ scores,
                                        IdType* __restrict__ output,
                                        uint32_t length, uint32_t top_k,
                                        void* dynamic_smem, int* scratch) {
-  using Traits = FilteredTopKTraits<DType, false>;
+  using Traits =
+      FilteredTopKTraits<decltype(topk_histogram_4096::to_float(DType{})),
+                         false>;
   using Radix = SampledRadixDigits;
   constexpr uint32_t kItemsPerThread = 2;
   constexpr uint32_t kSampleSize = BlockSize * kItemsPerThread;
@@ -92,7 +94,8 @@ __device__ bool try_sampled_exact_topk(const DType* __restrict__ scores,
     const uint32_t score_idx = static_cast<uint32_t>(
         (static_cast<uint64_t>(2 * sample_idx + 1) * length) /
         (2 * kSampleSize));
-    sample_keys[i] = Traits::ToOrdered(scores[score_idx]);
+    sample_keys[i] =
+        Traits::ToOrdered(topk_histogram_4096::to_float(scores[score_idx]));
   }
 
   auto* sort_smem =
@@ -127,14 +130,14 @@ __device__ bool try_sampled_exact_topk(const DType* __restrict__ scores,
       static_cast<uint64_t>(CandidateCapacity) * kSampleSize) {
     return false;
   }
-  const auto collect_candidate = [&](uint32_t idx, DType score) {
+  const auto collect_candidate = [&](uint32_t idx, float score) {
     if (Traits::ToOrdered(score) >= proposed_pivot) {
       const int pos = atomicAdd(candidate_count, 1);
       if (pos < static_cast<int>(CandidateCapacity)) candidates[pos] = idx;
     }
   };
 
-  vec_t<DType, VecSize> score_vec;
+  vec_t<float, VecSize> score_vec;
   const uint32_t aligned_length = length / VecSize * VecSize;
   for (uint32_t base = tx * VecSize; base < aligned_length;
        base += BlockSize * VecSize) {
@@ -145,7 +148,7 @@ __device__ bool try_sampled_exact_topk(const DType* __restrict__ scores,
     }
   }
   for (uint32_t idx = aligned_length + tx; idx < length; idx += BlockSize) {
-    collect_candidate(idx, scores[idx]);
+    collect_candidate(idx, topk_histogram_4096::to_float(scores[idx]));
   }
   __syncthreads();
 
@@ -157,8 +160,8 @@ __device__ bool try_sampled_exact_topk(const DType* __restrict__ scores,
 
   auto* candidate_keys = candidates + CandidateCapacity;
   for (int i = tx; i < num_candidates; i += BlockSize) {
-    candidate_keys[i] =
-        static_cast<int>(Traits::ToOrdered(scores[candidates[i]]));
+    candidate_keys[i] = static_cast<int>(Traits::ToOrdered(
+        topk_histogram_4096::to_float(scores[candidates[i]])));
   }
 
   if (tx == 0) {
diff --git a/csrc/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu b/csrc/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu
index 2064775111fb..269e50cfdf4a 100644
--- a/csrc/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu
+++ b/csrc/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu
@@ -38,7 +38,7 @@ namespace vllm {
 template <class Type, bool UE8M0_SF = false>
 __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
     silu_mul_cvt_fp16_to_fp4(int32_t numRows, int32_t numCols,
-                             int32_t num_packed_cols,
+                             int32_t num_packed_cols, int32_t num_padded_cols,
                              Type const* __restrict__ in,
                              float const* __restrict__ SFScale,
                              uint32_t* __restrict__ out,
@@ -49,9 +49,16 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
   static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD,
                 "Vec size is not matched.");
 
+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
+  cudaGridDependencySynchronize();
+  cudaTriggerProgrammaticLaunchCompletion();
+#endif
+
   // Precompute SF layout parameter (constant for entire kernel).
   int32_t const numKTiles = (numCols + 63) / 64;
 
+  int const sf_m = round_up<int>(numRows, 128);
+
   // Get the global scaling factor, which will be applied to the SF.
   // Note SFScale is the same as next GEMM's alpha, which is
   // (448.f / (Alpha_A / 6.f)).
@@ -60,9 +67,10 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
   int32_t const colIdx = blockDim.x * blockIdx.y + threadIdx.x;
   int elem_idx = colIdx * CVT_FP4_ELTS_PER_THREAD;
 
-  // Input tensor row/col loops.
-  for (int rowIdx = blockIdx.x; rowIdx < numRows; rowIdx += gridDim.x) {
-    if (colIdx < num_packed_cols) {
+  // Iterate over all rows and cols including padded ones -
+  //  ensures we visit every single scale factor address to initialize it.
+  for (int rowIdx = blockIdx.x; rowIdx < sf_m; rowIdx += gridDim.x) {
+    if (colIdx < num_padded_cols) {
       PackedVec in_vec;
       PackedVec in_vec2;
       int64_t inOffset =
@@ -70,7 +78,8 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
       int64_t inOffset2 = rowIdx * (numCols * 2 / CVT_FP4_ELTS_PER_THREAD) +
                           numCols / CVT_FP4_ELTS_PER_THREAD + colIdx;
 
-      bool valid = (rowIdx < numRows) && (elem_idx < numCols);
+      // If we are outside valid rows/columns, feed zeros
+      bool valid = (rowIdx < numRows) && (colIdx < num_packed_cols);
       if constexpr (CVT_FP4_PACK16) {
         ld256_cg_or_zero(reinterpret_cast<u32x8_t&>(in_vec),
                          &reinterpret_cast<const uint32_t*>(in)[inOffset * 8],
@@ -95,9 +104,17 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
                                              CVT_FP4_NUM_THREADS_PER_SF>(
               rowIdx, colIdx, numKTiles, SFout);
 
-      auto out_val =
-          cvt_warp_fp16_to_fp4<Type, CVT_FP4_NUM_THREADS_PER_SF, UE8M0_SF>(
-              out_silu_mul, SFScaleVal, sf_out);
+      fp4_packed_t out_val{};
+      if (valid) {
+        out_val =
+            cvt_warp_fp16_to_fp4<Type, CVT_FP4_NUM_THREADS_PER_SF, UE8M0_SF>(
+                out_silu_mul, SFScaleVal, sf_out);
+      } else if (sf_out != nullptr) {
+        // The swizzled scale tensor may be allocated with torch.empty.
+        // Explicitly initialize padded rows and columns in this kernel so
+        // downstream NVFP4 GEMMs never observe stale scale factors.
+        *sf_out = 0;
+      }
 
       if (valid) {
         if constexpr (CVT_FP4_PACK16) {
@@ -108,7 +125,7 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
         } else {
           int64_t outOffset =
               rowIdx * (numCols / CVT_FP4_ELTS_PER_THREAD) + colIdx;
-          out[outOffset] = out_val;
+          reinterpret_cast<fp4_packed_t*>(out)[outOffset] = out_val;
         }
       }
     }
@@ -146,18 +163,41 @@ void silu_and_mul_nvfp4_quant_sm1xxa(
 
   int num_packed_cols = int(n / CVT_FP4_ELTS_PER_THREAD);
 
-  int grid_y = vllm::div_round_up(num_packed_cols, static_cast<int>(block.x));
-  int grid_x = std::min(
-      int(m), std::max(1, (multiProcessorCount * numBlocksPerSM) / grid_y));
+  // The swizzled scale tensor is padded to 128 rows x 4 scale columns; the
+  // grid covers the padding so the kernel can initialize it.
+  int sf_n_int = int(vllm::round_up(n / CVT_FP4_SF_VEC_SIZE, 4) / 4);
+  int num_padded_cols =
+      sf_n_int * 4 * CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD;
+
+  int grid_y = vllm::div_round_up(num_padded_cols, static_cast<int>(block.x));
+  int grid_x =
+      std::min(vllm::computeEffectiveRows(m),
+               std::max(1, (multiProcessorCount * numBlocksPerSM) / grid_y));
   dim3 grid(grid_x, grid_y);
 
+  // PDL launch pays off at small M; at large M the early-launched blocks
+  // slow the chain down, so keep the serialized launch there.
+  bool const use_pdl =
+      int64_t(m) * num_packed_cols <= int64_t(multiProcessorCount) * 4096;
+
   VLLM_STABLE_DISPATCH_HALF_TYPES(
       input.scalar_type(), "silu_and_mul_nvfp4_quant_kernel", [&] {
         using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
         auto input_ptr = static_cast<cuda_type const*>(input.data_ptr());
-        vllm::silu_mul_cvt_fp16_to_fp4<cuda_type><<<grid, block, 0, stream>>>(
-            m, n, num_packed_cols, input_ptr, input_sf_ptr,
-            reinterpret_cast<uint32_t*>(output_ptr),
-            reinterpret_cast<uint32_t*>(sf_out));
+        cudaLaunchConfig_t config = {};
+        config.gridDim = grid;
+        config.blockDim = block;
+        config.dynamicSmemBytes = 0;
+        config.stream = stream;
+        cudaLaunchAttribute attrs[1];
+        attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
+        attrs[0].val.programmaticStreamSerializationAllowed = 1;
+        config.numAttrs = use_pdl ? 1 : 0;
+        config.attrs = attrs;
+        cudaLaunchKernelEx(&config, vllm::silu_mul_cvt_fp16_to_fp4<cuda_type>,
+                           m, n, num_packed_cols, num_padded_cols, input_ptr,
+                           input_sf_ptr,
+                           reinterpret_cast<uint32_t*>(output_ptr),
+                           reinterpret_cast<uint32_t*>(sf_out));
       });
 }
diff --git a/csrc/quantization/w8a8/fp8/common.cu b/csrc/quantization/w8a8/fp8/common.cu
index 471499786977..837948fdd2a7 100644
--- a/csrc/quantization/w8a8/fp8/common.cu
+++ b/csrc/quantization/w8a8/fp8/common.cu
@@ -95,7 +95,7 @@ __global__ void scaled_fp8_quant_kernel_strided_group_shape(
                                               int size, float inv_scale) {
     vectorize_with_alignment<VEC_SIZE>(
         in, out, size, tid, blockDim.x,
-        [=] __device__(fp8_type & dst, const scalar_t& src) {
+        [=] __device__(fp8_type& dst, const scalar_t& src) {
           dst = scaled_fp8_conversion<true, fp8_type>(static_cast<float>(src),
                                                       inv_scale);
         });
@@ -194,7 +194,7 @@ __global__ void scaled_fp8_quant_kernel_strided_dynamic(
   const float reciprocal_scale = 1.0f / (*scale);
   vectorize_with_alignment<16>(
       token_in, token_out, hidden_size, tid, blockDim.x,
-      [=] __device__(fp8_type & dst, const scalar_t& src) {
+      [=] __device__(fp8_type& dst, const scalar_t& src) {
         dst = scaled_fp8_conversion<true, fp8_type>(static_cast<float>(src),
                                                     reciprocal_scale);
       });
@@ -238,7 +238,7 @@ __global__ void dynamic_per_token_scaled_fp8_quant_kernel_strided(
   // 2) quantize
   vectorize_with_alignment<16>(
       token_in, token_out, hidden_size, tid, blockDim.x,
-      [=] __device__(fp8_type & dst, const scalar_t& src) {
+      [=] __device__(fp8_type& dst, const scalar_t& src) {
         dst = scaled_fp8_conversion<false, fp8_type>(static_cast<float>(src),
                                                      token_scale);
       });
diff --git a/csrc/quantization/w8a8/fp8/per_token_group_quant.cu b/csrc/quantization/w8a8/fp8/per_token_group_quant.cu
index 40d3f92efb39..162a2d879f67 100644
--- a/csrc/quantization/w8a8/fp8/per_token_group_quant.cu
+++ b/csrc/quantization/w8a8/fp8/per_token_group_quant.cu
@@ -49,7 +49,7 @@ __device__ __forceinline__ float ComputeGroupScale(
   constexpr int vec_size = 16 / sizeof(T);
 
   // copy global -> shared & compute absmax
-  auto scalar_op_cache = [&] __device__(T & dst, const T& src) {
+  auto scalar_op_cache = [&] __device__(T& dst, const T& src) {
     float abs_v = fabsf(static_cast<float>(src));
     local_absmax = fmaxf(local_absmax, abs_v);
     dst = src;
@@ -85,7 +85,7 @@ __device__ __forceinline__ void QuantizeGroup(
   const float inv_y = 1.0f / y_s;
 
   // quantize shared -> global 8-bit
-  auto scalar_op_quant = [&] __device__(DST_DTYPE & dst, const T& src) {
+  auto scalar_op_quant = [&] __device__(DST_DTYPE& dst, const T& src) {
     float q = fminf(fmaxf(static_cast<float>(src) * inv_y, min_8bit), max_8bit);
     dst = DST_DTYPE(q);
   };
diff --git a/csrc/rocm/.gitignore b/csrc/rocm/.gitignore
new file mode 100644
index 000000000000..afdfc475e0f1
--- /dev/null
+++ b/csrc/rocm/.gitignore
@@ -0,0 +1,7 @@
+# Generated at configure time by generate_paged_attn.py.
+attention_paged_attn_*.cu
+attention_paged_attn_entries.inc
+attention_paged_attn_dispatch.inc
+# Record of what is on disk (the script hash and the list of files written), so a
+# second build directory regenerates instead of silently reusing stale units.
+.paged_attn_generated
diff --git a/csrc/rocm/attention_dispatch.cu b/csrc/rocm/attention_dispatch.cu
index 53859c77e15e..3bb6a7ca6faf 100644
--- a/csrc/rocm/attention_dispatch.cu
+++ b/csrc/rocm/attention_dispatch.cu
@@ -17,31 +17,23 @@
 // Runtime dispatch for the ROCm paged attention kernels.
 //
 // The kernels themselves (and the expensive template instantiation they
-// trigger) live in attention_paged_attn_*.cu, one translation unit per
-// (dtype, kv-cache dtype, head size) combination, so ninja compiles them in
-// parallel instead of serializing the whole attention kernel set onto a single
-// core.
+// trigger) live in the generated attention_paged_attn_*.cu translation units,
+// one per (dtype, kv-cache dtype, head size, block size) combination, so ninja
+// compiles them in parallel instead of serializing the whole attention kernel
+// set onto a single core. Those translation units, their prototypes and the
+// dispatch below are all emitted from the variant list in
+// csrc/rocm/generate_paged_attn.py.
 #include <torch/all.h>
 
 #include "attention_kernels.cuh"
 #include "ops.h"
 
-// One prototype per entry point. Adding one here is useless unless
-// paged_attention() below also dispatches it (checked by CMakeLists.txt).
-ATTN_ENTRY(paged_attention_h64_f16_auto);
-ATTN_ENTRY(paged_attention_h128_f16_auto);
-ATTN_ENTRY(paged_attention_h64_bf16_auto);
-ATTN_ENTRY(paged_attention_h128_bf16_auto);
-ATTN_ENTRY(paged_attention_h64_f16_fp8_f16);
-ATTN_ENTRY(paged_attention_h128_f16_fp8_f16);
-ATTN_ENTRY(paged_attention_h64_f16_fp8_fp8);
-ATTN_ENTRY(paged_attention_h128_f16_fp8_fp8);
-ATTN_ENTRY(paged_attention_h64_bf16_fp8_f16);
-ATTN_ENTRY(paged_attention_h128_bf16_fp8_f16);
-ATTN_ENTRY(paged_attention_h64_bf16_fp8_fp8);
-ATTN_ENTRY(paged_attention_h128_bf16_fp8_fp8);
+// One prototype per generated entry point.
+#include "attention_paged_attn_entries.inc"
 
-// gfx11/gfx12 (RDNA) need a different kernel set than gfx9 (CDNA).
+// gfx11/gfx12 (RDNA) need a different kernel set than gfx9 (CDNA). Keep this
+// predicate in sync with the arch filter in CMakeLists.txt, which defines
+// VLLM_ROCM_ATTN_NO_NAVI when no Navi arch is targeted.
 bool is_navi_gpu() {
   static bool is_cached = false;
   static bool result;
@@ -60,22 +52,6 @@ bool is_navi_gpu() {
   return result;
 }
 
-// Each ATTN_HEAD_SWITCH() call below is a dispatch arm: it names the entry
-// point that paged_attention() calls, which is what makes the corresponding
-// attention_paged_attn_*.cu file reachable.
-#define ATTN_HEAD_SWITCH(FN64, FN128)                           \
-  switch (head_size) {                                          \
-    case 64:                                                    \
-      FN64(ATTN_FWD);                                           \
-      break;                                                    \
-    case 128:                                                   \
-      FN128(ATTN_FWD);                                          \
-      break;                                                    \
-    default:                                                    \
-      TORCH_CHECK(false, "Unsupported head size: ", head_size); \
-      break;                                                    \
-  }
-
 // clang-format off
 void paged_attention(
     torch::Tensor& out,         // [num_seqs, num_heads, head_size]
@@ -100,43 +76,11 @@ void paged_attention(
   // clang-format on
   bool is_navi = is_navi_gpu();
   const int head_size = query.size(2);
-  if (kv_cache_dtype == "auto") {
-    if (query.dtype() == at::ScalarType::Half) {
-      ATTN_HEAD_SWITCH(paged_attention_h64_f16_auto,
-                       paged_attention_h128_f16_auto);
-    } else if (query.dtype() == at::ScalarType::BFloat16) {
-      ATTN_HEAD_SWITCH(paged_attention_h64_bf16_auto,
-                       paged_attention_h128_bf16_auto);
-    } else {
-      TORCH_CHECK(false, "Unsupported data type: ", query.dtype());
-    }
-  } else if (kv_cache_dtype == "fp8" || kv_cache_dtype == "fp8_e4m3") {
-    if (query.dtype() == at::ScalarType::Half) {
-      if (mfma_type == "fp8") {
-        ATTN_HEAD_SWITCH(paged_attention_h64_f16_fp8_fp8,
-                         paged_attention_h128_f16_fp8_fp8);
-      } else {
-        ATTN_HEAD_SWITCH(paged_attention_h64_f16_fp8_f16,
-                         paged_attention_h128_f16_fp8_f16);
-      }
-    } else if (query.dtype() == at::ScalarType::BFloat16) {
-      if (mfma_type == "fp8") {
-        ATTN_HEAD_SWITCH(paged_attention_h64_bf16_fp8_fp8,
-                         paged_attention_h128_bf16_fp8_fp8);
-      } else {
-        ATTN_HEAD_SWITCH(paged_attention_h64_bf16_fp8_f16,
-                         paged_attention_h128_bf16_fp8_f16);
-      }
-    } else {
-      TORCH_CHECK(false, "Unsupported data type: ", query.dtype());
-    }
-  } else {
-    TORCH_CHECK(false, "Unsupported KV cache dtype: ", kv_cache_dtype);
-  }
+  // clang-format off
+  #include "attention_paged_attn_dispatch.inc"
+  // clang-format on
 }
 
-#undef ATTN_HEAD_SWITCH
 #undef ATTN_ENTRY
-#undef ATTN_LAUNCH
 #undef ATTN_FWD
 #undef ATTN_PARAMS
diff --git a/csrc/rocm/attention_kernels.cuh b/csrc/rocm/attention_kernels.cuh
index 42539faa0d73..d17acc74d808 100644
--- a/csrc/rocm/attention_kernels.cuh
+++ b/csrc/rocm/attention_kernels.cuh
@@ -16,9 +16,16 @@
 
 #pragma once
 
-#include <torch/all.h>
+// Only the ATen/c10 surface is needed here (at::Tensor, the CUDA stream/guard
+// helpers and TORCH_CHECK). <torch/all.h> also drags in autograd, nn and jit,
+// which is ~7 s of the ~9.5 s frontend of every attention translation unit.
+#include <ATen/DeviceGuard.h>
+#include <ATen/core/Tensor.h>
 #include <ATen/cuda/CUDAContext.h>
 #include <c10/cuda/CUDAGuard.h>
+#include <c10/util/Exception.h>
+
+#include <optional>
 #include <hip/hip_fp8.h>
 #include <hip/hip_bf16.h>
 #include "core/cuda_compat.h"
@@ -3433,13 +3440,13 @@ template <typename T, typename KVT, vllm::Fp8KVCacheDataType KV_DTYPE,
           int BLOCK_SIZE, int HEAD_SIZE, typename OUTT, int PARTITION_SIZE_OLD,
           bool ALIBI_ENABLED, MFMAType MFMA_TYPE, bool USE_INTERLEAVED_V_CACHE>
 void paged_attention_custom_launcher(
-    torch::Tensor& out, torch::Tensor& exp_sums, torch::Tensor& max_logits,
-    torch::Tensor& tmp_out, torch::Tensor& query, torch::Tensor& key_cache,
-    torch::Tensor& value_cache, const int num_kv_heads, float scale,
-    torch::Tensor& block_tables, torch::Tensor& seq_lens,
-    const std::optional<torch::Tensor>& query_start_loc, int max_seq_len,
-    const std::optional<torch::Tensor>& alibi_slopes, torch::Tensor& k_scale,
-    torch::Tensor& v_scale, const std::optional<torch::Tensor>& fp8_out_scale) {
+    at::Tensor& out, at::Tensor& exp_sums, at::Tensor& max_logits,
+    at::Tensor& tmp_out, at::Tensor& query, at::Tensor& key_cache,
+    at::Tensor& value_cache, const int num_kv_heads, float scale,
+    at::Tensor& block_tables, at::Tensor& seq_lens,
+    const std::optional<at::Tensor>& query_start_loc, int max_seq_len,
+    const std::optional<at::Tensor>& alibi_slopes, at::Tensor& k_scale,
+    at::Tensor& v_scale, const std::optional<at::Tensor>& fp8_out_scale) {
   int num_seqs = block_tables.size(0);
   int num_heads = query.size(1);
   int head_size = query.size(2);
@@ -3589,13 +3596,13 @@ template <typename T, typename KVT, vllm::Fp8KVCacheDataType KV_DTYPE,
           int BLOCK_SIZE, int HEAD_SIZE, typename OUTT, int PARTITION_SIZE_OLD,
           bool ALIBI_ENABLED, MFMAType MFMA_TYPE, bool USE_INTERLEAVED_V_CACHE>
 void paged_attention_custom_launcher_navi(
-    torch::Tensor& out, torch::Tensor& exp_sums, torch::Tensor& max_logits,
-    torch::Tensor& tmp_out, torch::Tensor& query, torch::Tensor& key_cache,
-    torch::Tensor& value_cache, const int num_kv_heads, float scale,
-    torch::Tensor& block_tables, torch::Tensor& seq_lens,
-    const std::optional<torch::Tensor>& query_start_loc, int max_seq_len,
-    const std::optional<torch::Tensor>& alibi_slopes, torch::Tensor& k_scale,
-    torch::Tensor& v_scale) {
+    at::Tensor& out, at::Tensor& exp_sums, at::Tensor& max_logits,
+    at::Tensor& tmp_out, at::Tensor& query, at::Tensor& key_cache,
+    at::Tensor& value_cache, const int num_kv_heads, float scale,
+    at::Tensor& block_tables, at::Tensor& seq_lens,
+    const std::optional<at::Tensor>& query_start_loc, int max_seq_len,
+    const std::optional<at::Tensor>& alibi_slopes, at::Tensor& k_scale,
+    at::Tensor& v_scale) {
   int num_seqs = block_tables.size(0);
   int num_heads = query.size(1);
   int head_size = query.size(2);
@@ -3766,19 +3773,46 @@ void paged_attention_custom_launcher_navi(
 // compatibility; a TORCH_CHECK prevents dispatch to NAVI with interleaved V.
 // TODO: implement interleaved V-cache support in the NAVI (GFX11/GFX12)
 // kernels if shuffle KV cache layout is needed on RDNA GPUs.
-#define CALL_CUSTOM_LAUNCHER(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, OUTT,    \
-                             PSIZE, ALIBI_ENABLED, MFMA_TYPE)                \
-  if (use_interleaved_v_cache) {                                             \
-    TORCH_CHECK(!is_navi,                                                    \
-                "Reading interleaved V-cache layout from paged decode "      \
-                "attention not supported on NAVI GPUs");                     \
-    paged_attention_custom_launcher<T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE,   \
-                                    OUTT, PSIZE, ALIBI_ENABLED, MFMA_TYPE,   \
-                                    true>(                                   \
-        out, exp_sums, max_logits, tmp_out, query, key_cache, value_cache,   \
-        num_kv_heads, scale, block_tables, seq_lens, query_start_loc,        \
-        max_seq_len, alibi_slopes, k_scale, v_scale, fp8_out_scale);         \
-  } else {                                                                   \
+//
+// BLK_SIZE/HEAD_SIZE are baked into the generated entry points (see
+// generate_paged_attn.py); ALIBI_ENABLED, the out dtype,
+// use_interleaved_v_cache and is_navi stay runtime branches, see the NOTE
+// below.
+#define CALL_CUSTOM_LAUNCHER(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, OUTT,  \
+                             PSIZE, ALIBI_ENABLED, MFMA_TYPE)              \
+  if (use_interleaved_v_cache) {                                           \
+    TORCH_CHECK(!is_navi,                                                  \
+                "Reading interleaved V-cache layout from paged decode "    \
+                "attention not supported on NAVI GPUs");                   \
+    paged_attention_custom_launcher<T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, \
+                                    OUTT, PSIZE, ALIBI_ENABLED, MFMA_TYPE, \
+                                    true>(                                 \
+        out, exp_sums, max_logits, tmp_out, query, key_cache, value_cache, \
+        num_kv_heads, scale, block_tables, seq_lens, query_start_loc,      \
+        max_seq_len, alibi_slopes, k_scale, v_scale, fp8_out_scale);       \
+  } else {                                                                 \
+    ATTN_CALL_NON_INTERLEAVED(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, OUTT, \
+                              PSIZE, ALIBI_ENABLED, MFMA_TYPE)             \
+  }
+
+// The non-interleaved path picks the CDNA or the Navi launcher at runtime. With
+// no gfx11/gfx12 arch CMake defines VLLM_ROCM_ATTN_NO_NAVI: is_navi_gpu() can
+// then never be true, so the Navi set is dead weight and is not instantiated.
+#ifdef VLLM_ROCM_ATTN_NO_NAVI
+  #define ATTN_CALL_NON_INTERLEAVED(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE,  \
+                                    OUTT, PSIZE, ALIBI_ENABLED, MFMA_TYPE)  \
+    TORCH_CHECK(!is_navi,                                                   \
+                "Navi (gfx11/gfx12) GPUs are not supported by this build: " \
+                "its arch list has no Navi code objects");                  \
+    paged_attention_custom_launcher<T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE,  \
+                                    OUTT, PSIZE, ALIBI_ENABLED, MFMA_TYPE,  \
+                                    false>(                                 \
+        out, exp_sums, max_logits, tmp_out, query, key_cache, value_cache,  \
+        num_kv_heads, scale, block_tables, seq_lens, query_start_loc,       \
+        max_seq_len, alibi_slopes, k_scale, v_scale, fp8_out_scale);
+#else
+  #define ATTN_CALL_NON_INTERLEAVED(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE,   \
+                                    OUTT, PSIZE, ALIBI_ENABLED, MFMA_TYPE)   \
     if (!is_navi) {                                                          \
       paged_attention_custom_launcher<T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, \
                                       OUTT, PSIZE, ALIBI_ENABLED, MFMA_TYPE, \
@@ -3793,8 +3827,8 @@ void paged_attention_custom_launcher_navi(
           out, exp_sums, max_logits, tmp_out, query, key_cache, value_cache, \
           num_kv_heads, scale, block_tables, seq_lens, query_start_loc,      \
           max_seq_len, alibi_slopes, k_scale, v_scale);                      \
-    }                                                                        \
-  }
+    }
+#endif
 
 #define CALL_CUSTOM_LAUNCHER_ALIBI(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE,    \
                                    OUTT, PSIZE, MFMA_TYPE)                   \
@@ -3827,18 +3861,11 @@ void paged_attention_custom_launcher_navi(
     }
 #endif
 
-#define CALL_CUSTOM_LAUNCHER_BLK(T, KVT, KV_DTYPE, HEAD_SIZE, MFMA_TYPE)    \
-  switch (block_size) {                                                     \
-    case 16:                                                                \
-      CALL_CUSTOM_LAUNCHER_OUT(T, KVT, KV_DTYPE, 16, HEAD_SIZE, MFMA_TYPE); \
-      break;                                                                \
-    case 32:                                                                \
-      CALL_CUSTOM_LAUNCHER_OUT(T, KVT, KV_DTYPE, 32, HEAD_SIZE, MFMA_TYPE); \
-      break;                                                                \
-    default:                                                                \
-      TORCH_CHECK(false, "Unsupported block size: ", block_size);           \
-      break;                                                                \
-  }
+// NOTE: ALIBI_ENABLED and the out dtype (fp8 vs query dtype) stay runtime
+// branches rather than getting their own units (4 per (variant, block size)
+// instead of 1). That halves the per-unit codegen, but every unit pays ~20 s to
+// parse this header, so _rocm_C's serial work roughly doubles and the -j16 wall
+// went from 142 s to 240 s.
 
 // ---------------------------------------------------------------------------
 // Fixed-head-size entry points.
@@ -3848,42 +3875,39 @@ void paged_attention_custom_launcher_navi(
 // size, block size, alibi, out dtype, v-cache layout) combination. To let ninja
 // compile them in parallel instead of serializing ~10 min onto a single core,
 // the instantiations are spread over one translation unit per (dtype, kv-cache
-// dtype, head size) combination, see attention_paged_attn_*.cu.
+// dtype, head size, block size) combination: the translation units, their
+// prototypes and the dispatch arms of paged_attention() are all emitted from
+// the variant list in csrc/rocm/generate_paged_attn.py.
 // ---------------------------------------------------------------------------
 
 // Parameter list shared by paged_attention() and every entry point defined in
-// attention_paged_attn_*.cu. The names must match what CALL_CUSTOM_LAUNCHER*
-// expect to find in scope.
-#define ATTN_PARAMS                                                            \
-  torch::Tensor &out, torch::Tensor &exp_sums, torch::Tensor &max_logits,      \
-      torch::Tensor &tmp_out, torch::Tensor &query, torch::Tensor &key_cache,  \
-      torch::Tensor &value_cache, int64_t num_kv_heads, double scale,          \
-      torch::Tensor &block_tables, torch::Tensor &seq_lens,                    \
-      const std::optional<torch::Tensor>&query_start_loc, int64_t max_seq_len, \
-      const std::optional<torch::Tensor>&alibi_slopes, torch::Tensor &k_scale, \
-      torch::Tensor &v_scale,                                                  \
-      const std::optional<torch::Tensor>&fp8_out_scale, int64_t block_size,    \
+// the generated attention_paged_attn_*.cu files. The names must match what
+// CALL_CUSTOM_LAUNCHER* expect to find in scope. block_size is not part of it:
+// it is part of the entry point name, so that each translation unit only
+// instantiates the kernels of one block size and paged_attention() switches on
+// it.
+#define ATTN_PARAMS                                                         \
+  at::Tensor &out, at::Tensor &exp_sums, at::Tensor &max_logits,            \
+      at::Tensor &tmp_out, at::Tensor &query, at::Tensor &key_cache,        \
+      at::Tensor &value_cache, int64_t num_kv_heads, double scale,          \
+      at::Tensor &block_tables, at::Tensor &seq_lens,                       \
+      const std::optional<at::Tensor>&query_start_loc, int64_t max_seq_len, \
+      const std::optional<at::Tensor>&alibi_slopes, at::Tensor &k_scale,    \
+      at::Tensor &v_scale, const std::optional<at::Tensor>&fp8_out_scale,   \
       bool is_navi, bool use_interleaved_v_cache
 
-#define ATTN_FWD                                                              \
-  out, exp_sums, max_logits, tmp_out, query, key_cache, value_cache,          \
-      num_kv_heads, scale, block_tables, seq_lens, query_start_loc,           \
-      max_seq_len, alibi_slopes, k_scale, v_scale, fp8_out_scale, block_size, \
-      is_navi, use_interleaved_v_cache
+#define ATTN_FWD                                                           \
+  out, exp_sums, max_logits, tmp_out, query, key_cache, value_cache,       \
+      num_kv_heads, scale, block_tables, seq_lens, query_start_loc,        \
+      max_seq_len, alibi_slopes, k_scale, v_scale, fp8_out_scale, is_navi, \
+      use_interleaved_v_cache
 
-// Declare/define one fixed-head-size entry point. A definition in one of the
-// attention_paged_attn_*.cu files is reachable only if attention_dispatch.cu
-// names it in one of its ATTN_HEAD_SWITCH() calls (CMakeLists.txt fails if one
-// is missing).
+// Declare/define one entry point. The generated attention_paged_attn_*.cu files
+// define one, and attention_paged_attn_entries.inc declares them; both are
+// emitted from the same variant list, so an entry point cannot be compiled
+// without being dispatched.
 #define ATTN_ENTRY(fn) void fn(ATTN_PARAMS)
 
-// Instantiate one entry point (one per attention_paged_attn_*.cu file), for
-// every supported block size. Kept as an alias of CALL_CUSTOM_LAUNCHER_BLK so
-// the block-size switch has a single definition; the head-size switch lives in
-// attention_dispatch.cu.
-#define ATTN_LAUNCH(T, KVT, KV_DTYPE, HEAD_SIZE, MFMA_TYPE) \
-  CALL_CUSTOM_LAUNCHER_BLK(T, KVT, KV_DTYPE, HEAD_SIZE, MFMA_TYPE)
-
 #undef WARP_SIZE
 #undef MAX
 #undef MIN
diff --git a/csrc/rocm/attention_paged_attn_h128_bf16_auto.cu b/csrc/rocm/attention_paged_attn_h128_bf16_auto.cu
deleted file mode 100644
index 536af4557185..000000000000
--- a/csrc/rocm/attention_paged_attn_h128_bf16_auto.cu
+++ /dev/null
@@ -1,9 +0,0 @@
-// SPDX-License-Identifier: Apache-2.0
-// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-#include "attention_kernels.cuh"
-
-ATTN_ENTRY(paged_attention_h128_bf16_auto) {
-  ATTN_LAUNCH(__hip_bfloat16, __hip_bfloat16, vllm::Fp8KVCacheDataType::kAuto,
-              128, MFMAType::F16);
-}
diff --git a/csrc/rocm/attention_paged_attn_h128_bf16_fp8_f16.cu b/csrc/rocm/attention_paged_attn_h128_bf16_fp8_f16.cu
deleted file mode 100644
index b183a378a861..000000000000
--- a/csrc/rocm/attention_paged_attn_h128_bf16_fp8_f16.cu
+++ /dev/null
@@ -1,9 +0,0 @@
-// SPDX-License-Identifier: Apache-2.0
-// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-#include "attention_kernels.cuh"
-
-ATTN_ENTRY(paged_attention_h128_bf16_fp8_f16) {
-  ATTN_LAUNCH(__hip_bfloat16, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3, 128,
-              MFMAType::F16);
-}
diff --git a/csrc/rocm/attention_paged_attn_h128_bf16_fp8_fp8.cu b/csrc/rocm/attention_paged_attn_h128_bf16_fp8_fp8.cu
deleted file mode 100644
index 913d9f018963..000000000000
--- a/csrc/rocm/attention_paged_attn_h128_bf16_fp8_fp8.cu
+++ /dev/null
@@ -1,9 +0,0 @@
-// SPDX-License-Identifier: Apache-2.0
-// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-#include "attention_kernels.cuh"
-
-ATTN_ENTRY(paged_attention_h128_bf16_fp8_fp8) {
-  ATTN_LAUNCH(__hip_bfloat16, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3, 128,
-              MFMAType::Fp8);
-}
diff --git a/csrc/rocm/attention_paged_attn_h128_f16_auto.cu b/csrc/rocm/attention_paged_attn_h128_f16_auto.cu
deleted file mode 100644
index 1a38303f4d7f..000000000000
--- a/csrc/rocm/attention_paged_attn_h128_f16_auto.cu
+++ /dev/null
@@ -1,9 +0,0 @@
-// SPDX-License-Identifier: Apache-2.0
-// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-#include "attention_kernels.cuh"
-
-ATTN_ENTRY(paged_attention_h128_f16_auto) {
-  ATTN_LAUNCH(_Float16, _Float16, vllm::Fp8KVCacheDataType::kAuto, 128,
-              MFMAType::F16);
-}
diff --git a/csrc/rocm/attention_paged_attn_h128_f16_fp8_f16.cu b/csrc/rocm/attention_paged_attn_h128_f16_fp8_f16.cu
deleted file mode 100644
index 62fb1928f2af..000000000000
--- a/csrc/rocm/attention_paged_attn_h128_f16_fp8_f16.cu
+++ /dev/null
@@ -1,9 +0,0 @@
-// SPDX-License-Identifier: Apache-2.0
-// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-#include "attention_kernels.cuh"
-
-ATTN_ENTRY(paged_attention_h128_f16_fp8_f16) {
-  ATTN_LAUNCH(_Float16, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3, 128,
-              MFMAType::F16);
-}
diff --git a/csrc/rocm/attention_paged_attn_h128_f16_fp8_fp8.cu b/csrc/rocm/attention_paged_attn_h128_f16_fp8_fp8.cu
deleted file mode 100644
index 21afad1cbebd..000000000000
--- a/csrc/rocm/attention_paged_attn_h128_f16_fp8_fp8.cu
+++ /dev/null
@@ -1,9 +0,0 @@
-// SPDX-License-Identifier: Apache-2.0
-// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-#include "attention_kernels.cuh"
-
-ATTN_ENTRY(paged_attention_h128_f16_fp8_fp8) {
-  ATTN_LAUNCH(_Float16, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3, 128,
-              MFMAType::Fp8);
-}
diff --git a/csrc/rocm/attention_paged_attn_h64_bf16_auto.cu b/csrc/rocm/attention_paged_attn_h64_bf16_auto.cu
deleted file mode 100644
index 9d33b250df8d..000000000000
--- a/csrc/rocm/attention_paged_attn_h64_bf16_auto.cu
+++ /dev/null
@@ -1,9 +0,0 @@
-// SPDX-License-Identifier: Apache-2.0
-// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-#include "attention_kernels.cuh"
-
-ATTN_ENTRY(paged_attention_h64_bf16_auto) {
-  ATTN_LAUNCH(__hip_bfloat16, __hip_bfloat16, vllm::Fp8KVCacheDataType::kAuto,
-              64, MFMAType::F16);
-}
diff --git a/csrc/rocm/attention_paged_attn_h64_bf16_fp8_f16.cu b/csrc/rocm/attention_paged_attn_h64_bf16_fp8_f16.cu
deleted file mode 100644
index 10c74d77b479..000000000000
--- a/csrc/rocm/attention_paged_attn_h64_bf16_fp8_f16.cu
+++ /dev/null
@@ -1,9 +0,0 @@
-// SPDX-License-Identifier: Apache-2.0
-// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-#include "attention_kernels.cuh"
-
-ATTN_ENTRY(paged_attention_h64_bf16_fp8_f16) {
-  ATTN_LAUNCH(__hip_bfloat16, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3, 64,
-              MFMAType::F16);
-}
diff --git a/csrc/rocm/attention_paged_attn_h64_bf16_fp8_fp8.cu b/csrc/rocm/attention_paged_attn_h64_bf16_fp8_fp8.cu
deleted file mode 100644
index 8992c7143232..000000000000
--- a/csrc/rocm/attention_paged_attn_h64_bf16_fp8_fp8.cu
+++ /dev/null
@@ -1,9 +0,0 @@
-// SPDX-License-Identifier: Apache-2.0
-// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-#include "attention_kernels.cuh"
-
-ATTN_ENTRY(paged_attention_h64_bf16_fp8_fp8) {
-  ATTN_LAUNCH(__hip_bfloat16, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3, 64,
-              MFMAType::Fp8);
-}
diff --git a/csrc/rocm/attention_paged_attn_h64_f16_auto.cu b/csrc/rocm/attention_paged_attn_h64_f16_auto.cu
deleted file mode 100644
index b65c197de732..000000000000
--- a/csrc/rocm/attention_paged_attn_h64_f16_auto.cu
+++ /dev/null
@@ -1,9 +0,0 @@
-// SPDX-License-Identifier: Apache-2.0
-// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-#include "attention_kernels.cuh"
-
-ATTN_ENTRY(paged_attention_h64_f16_auto) {
-  ATTN_LAUNCH(_Float16, _Float16, vllm::Fp8KVCacheDataType::kAuto, 64,
-              MFMAType::F16);
-}
diff --git a/csrc/rocm/attention_paged_attn_h64_f16_fp8_f16.cu b/csrc/rocm/attention_paged_attn_h64_f16_fp8_f16.cu
deleted file mode 100644
index 6c5897ece5e3..000000000000
--- a/csrc/rocm/attention_paged_attn_h64_f16_fp8_f16.cu
+++ /dev/null
@@ -1,9 +0,0 @@
-// SPDX-License-Identifier: Apache-2.0
-// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-#include "attention_kernels.cuh"
-
-ATTN_ENTRY(paged_attention_h64_f16_fp8_f16) {
-  ATTN_LAUNCH(_Float16, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3, 64,
-              MFMAType::F16);
-}
diff --git a/csrc/rocm/attention_paged_attn_h64_f16_fp8_fp8.cu b/csrc/rocm/attention_paged_attn_h64_f16_fp8_fp8.cu
deleted file mode 100644
index 573a58f19f47..000000000000
--- a/csrc/rocm/attention_paged_attn_h64_f16_fp8_fp8.cu
+++ /dev/null
@@ -1,9 +0,0 @@
-// SPDX-License-Identifier: Apache-2.0
-// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-#include "attention_kernels.cuh"
-
-ATTN_ENTRY(paged_attention_h64_f16_fp8_fp8) {
-  ATTN_LAUNCH(_Float16, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3, 64,
-              MFMAType::Fp8);
-}
diff --git a/csrc/rocm/generate_paged_attn.py b/csrc/rocm/generate_paged_attn.py
new file mode 100644
index 000000000000..802c48565b76
--- /dev/null
+++ b/csrc/rocm/generate_paged_attn.py
@@ -0,0 +1,435 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Generate the ROCm paged attention translation units and their dispatch.
+
+Instantiating these kernels dominates the _rocm_C build: one file holding every
+(dtype, kv-cache dtype, head size, block size, ...) combination serialized
+~10 min of device code generation onto one core.
+
+All outputs derive from VARIANTS:
+
+  * attention_paged_attn_<variant>_bs<block_size>.cu: one unit per (variant,
+    block size), each instantiating only its slice (~48 s per unit, vs ~95 s for
+    a whole set), so ninja runs them in parallel,
+  * attention_paged_attn_entries.inc: entry-point prototypes (file scope),
+  * attention_paged_attn_dispatch.inc: paged_attention()'s dispatch (function
+    body).
+
+So a unit cannot be compiled without being dispatched - what the configure-time
+guard in CMakeLists.txt used to regex for. Outputs land next to the kernels (like
+marlin's generate_kernels.py) and are gitignored.
+"""
+
+import glob
+import hashlib
+import os
+import sys
+
+# Block sizes the kernels are instantiated for (see CALL_CUSTOM_LAUNCHER in
+# attention_kernels.cuh). One translation unit is generated per (variant, block
+# size).
+BLOCK_SIZES = [16, 32]
+
+# ALIBI_ENABLED and the out dtype stay runtime branches (CALL_CUSTOM_LAUNCHER_OUT)
+# instead of getting units of their own: see the NOTE in attention_kernels.cuh.
+
+# One entry per supported (query dtype, kv-cache dtype, head size) kernel set:
+#   query_type/kv_type: C++ types of the query/out and of the key/value cache
+#                       (template parameters T and KVT), written with the HIP
+#                       spellings hipify produces; `__half` is a distinct type
+#                       from `_Float16` for the MFMA builtins.
+#   kv_cache_enum/head_size/mfma: Fp8KVCacheDataType, head size and MFMAType
+#   query_scalar_type:  at::ScalarType of the query this variant serves
+#   kv_cache_dtype/mfma_type: the runtime strings this variant serves; mfma_type
+#                       is None when the kv-cache dtype does not split on it.
+VARIANTS = [
+    {
+        "query_type": "_Float16",
+        "kv_type": "_Float16",
+        "kv_cache_enum": "vllm::Fp8KVCacheDataType::kAuto",
+        "head_size": 64,
+        "mfma": "MFMAType::F16",
+        "query_scalar_type": "Half",
+        "kv_cache_dtype": "auto",
+        "mfma_type": None,
+    },
+    {
+        "query_type": "_Float16",
+        "kv_type": "_Float16",
+        "kv_cache_enum": "vllm::Fp8KVCacheDataType::kAuto",
+        "head_size": 128,
+        "mfma": "MFMAType::F16",
+        "query_scalar_type": "Half",
+        "kv_cache_dtype": "auto",
+        "mfma_type": None,
+    },
+    {
+        "query_type": "__hip_bfloat16",
+        "kv_type": "__hip_bfloat16",
+        "kv_cache_enum": "vllm::Fp8KVCacheDataType::kAuto",
+        "head_size": 64,
+        "mfma": "MFMAType::F16",
+        "query_scalar_type": "BFloat16",
+        "kv_cache_dtype": "auto",
+        "mfma_type": None,
+    },
+    {
+        "query_type": "__hip_bfloat16",
+        "kv_type": "__hip_bfloat16",
+        "kv_cache_enum": "vllm::Fp8KVCacheDataType::kAuto",
+        "head_size": 128,
+        "mfma": "MFMAType::F16",
+        "query_scalar_type": "BFloat16",
+        "kv_cache_dtype": "auto",
+        "mfma_type": None,
+    },
+    {
+        "query_type": "_Float16",
+        "kv_type": "uint8_t",
+        "kv_cache_enum": "vllm::Fp8KVCacheDataType::kFp8E4M3",
+        "head_size": 64,
+        "mfma": "MFMAType::F16",
+        "query_scalar_type": "Half",
+        "kv_cache_dtype": "fp8",
+        "mfma_type": "f16",
+    },
+    {
+        "query_type": "_Float16",
+        "kv_type": "uint8_t",
+        "kv_cache_enum": "vllm::Fp8KVCacheDataType::kFp8E4M3",
+        "head_size": 128,
+        "mfma": "MFMAType::F16",
+        "query_scalar_type": "Half",
+        "kv_cache_dtype": "fp8",
+        "mfma_type": "f16",
+    },
+    {
+        "query_type": "_Float16",
+        "kv_type": "uint8_t",
+        "kv_cache_enum": "vllm::Fp8KVCacheDataType::kFp8E4M3",
+        "head_size": 64,
+        "mfma": "MFMAType::Fp8",
+        "query_scalar_type": "Half",
+        "kv_cache_dtype": "fp8",
+        "mfma_type": "fp8",
+    },
+    {
+        "query_type": "_Float16",
+        "kv_type": "uint8_t",
+        "kv_cache_enum": "vllm::Fp8KVCacheDataType::kFp8E4M3",
+        "head_size": 128,
+        "mfma": "MFMAType::Fp8",
+        "query_scalar_type": "Half",
+        "kv_cache_dtype": "fp8",
+        "mfma_type": "fp8",
+    },
+    {
+        "query_type": "__hip_bfloat16",
+        "kv_type": "uint8_t",
+        "kv_cache_enum": "vllm::Fp8KVCacheDataType::kFp8E4M3",
+        "head_size": 64,
+        "mfma": "MFMAType::F16",
+        "query_scalar_type": "BFloat16",
+        "kv_cache_dtype": "fp8",
+        "mfma_type": "f16",
+    },
+    {
+        "query_type": "__hip_bfloat16",
+        "kv_type": "uint8_t",
+        "kv_cache_enum": "vllm::Fp8KVCacheDataType::kFp8E4M3",
+        "head_size": 128,
+        "mfma": "MFMAType::F16",
+        "query_scalar_type": "BFloat16",
+        "kv_cache_dtype": "fp8",
+        "mfma_type": "f16",
+    },
+    {
+        "query_type": "__hip_bfloat16",
+        "kv_type": "uint8_t",
+        "kv_cache_enum": "vllm::Fp8KVCacheDataType::kFp8E4M3",
+        "head_size": 64,
+        "mfma": "MFMAType::Fp8",
+        "query_scalar_type": "BFloat16",
+        "kv_cache_dtype": "fp8",
+        "mfma_type": "fp8",
+    },
+    {
+        "query_type": "__hip_bfloat16",
+        "kv_type": "uint8_t",
+        "kv_cache_enum": "vllm::Fp8KVCacheDataType::kFp8E4M3",
+        "head_size": 128,
+        "mfma": "MFMAType::Fp8",
+        "query_scalar_type": "BFloat16",
+        "kv_cache_dtype": "fp8",
+        "mfma_type": "fp8",
+    },
+]
+
+
+def variant_name(variant):
+    """File and entry-point suffix, derived so it cannot disagree with the fields."""
+    parts = [
+        f"h{variant['head_size']}",
+        {"_Float16": "f16", "__hip_bfloat16": "bf16"}[variant["query_type"]],
+        variant["kv_cache_dtype"],
+    ]
+    if variant["mfma_type"] is not None:
+        parts.append(variant["mfma_type"])
+    return "_".join(parts)
+
+
+for variant in VARIANTS:
+    variant["name"] = variant_name(variant)
+
+# kv_cache_dtype values accepted per VARIANTS["kv_cache_dtype"], in dispatch
+# order (the first matching arm wins).
+KV_CACHE_DTYPE_PREDICATES = {
+    "auto": 'kv_cache_dtype == "auto"',
+    "fp8": 'kv_cache_dtype == "fp8" || kv_cache_dtype == "fp8_e4m3"',
+}
+
+OUTPUT_DIR = os.path.dirname(os.path.abspath(__file__))
+
+# MD5 of this script, recorded in the stamp so CMake can tell whether the files
+# on disk were produced by the current version (same MD5 as `file(MD5 ...)`).
+# CMake passes the key it compares against so the format is defined in one place;
+# a manual run falls back to the hash above.
+with open(__file__, "rb") as _self:
+    # A cache key only: a FIPS-enforcing host rejects md5 without this flag.
+    SCRIPT_HASH = hashlib.md5(_self.read(), usedforsecurity=False).hexdigest()
+KEY = sys.argv[1] if len(sys.argv) > 1 else SCRIPT_HASH
+STAMP_NAME = ".paged_attn_generated"
+
+# Whether the Navi (gfx11/gfx12) kernel set is instantiated is a compile
+# definition decided by CMake (`VLLM_ROCM_ATTN_NO_NAVI`, see CMakeLists.txt), so
+# the units written here are identical for every build directory and arch list.
+
+FILE_HEAD = """\
+// auto generated by generate_paged_attn.py
+// clang-format off
+"""
+
+TU_HEAD = (
+    FILE_HEAD
+    + """
+#include "attention_kernels.cuh"
+"""
+)
+
+TU_ENTRY = """
+ATTN_ENTRY(paged_attention_{name}_bs{block_size}) {{
+  CALL_CUSTOM_LAUNCHER_OUT({query_type}, {kv_type}, {kv_cache_enum},
+                           {block_size}, {head_size}, {mfma});
+}}
+"""
+
+ENTRIES_HEAD = (
+    FILE_HEAD
+    + """
+// Prototypes of the entry points defined in the generated
+// attention_paged_attn_*.cu translation units.
+"""
+)
+
+DISPATCH_HEAD = (
+    FILE_HEAD
+    + """
+  // Dispatch on kv_cache_dtype, query dtype, mfma_type, block size and head
+  // size, mirroring the template parameters of the generated entry points.
+"""
+)
+
+
+def entry_point(variant, block_size):
+    return f"paged_attention_{variant['name']}_bs{block_size}"
+
+
+def block_size_switch(variants, indent_level):
+    """Switch on block_size then head_size for one dispatch group.
+
+    `variants` is the list of variants sharing (kv_cache_dtype, query dtype,
+    mfma_type); they differ in head size only.
+    """
+    indent = " " * indent_level
+    inner = " " * (indent_level + 4)
+    lines = [f"{indent}switch (block_size) {{"]
+    for block_size in BLOCK_SIZES:
+        lines.append(f"{indent}  case {block_size}:")
+        lines.append(f"{inner}switch (head_size) {{")
+        for variant in variants:
+            lines.append(f"{inner}  case {variant['head_size']}:")
+            lines.append(f"{inner}    {entry_point(variant, block_size)}(ATTN_FWD);")
+            lines.append(f"{inner}    break;")
+        lines.append(f"{inner}  default:")
+        lines.append(
+            f'{inner}    TORCH_CHECK(false, "Unsupported head size: ", head_size);'
+        )
+        lines.append(f"{inner}    break;")
+        lines.append(f"{inner}}}")
+        lines.append(f"{indent}    break;")
+    lines.append(f"{indent}  default:")
+    lines.append(
+        f'{indent}    TORCH_CHECK(false, "Unsupported block size: ", block_size);'
+    )
+    lines.append(f"{indent}    break;")
+    lines.append(f"{indent}}}")
+    return lines
+
+
+def group_variants():
+    """Group VARIANTS by (kv_cache_dtype, query dtype, mfma_type), in order."""
+    groups = {}
+    for variant in VARIANTS:
+        key = (
+            variant["kv_cache_dtype"],
+            variant["query_scalar_type"],
+            variant["mfma_type"],
+        )
+        groups.setdefault(key, []).append(variant)
+    return groups
+
+
+def generate_dispatch():
+    groups = group_variants()
+    # kv_cache_dtype -> query dtype -> mfma_type -> variants
+    by_kv_cache_dtype = {}
+    for (kv_cache_dtype, query_scalar_type, mfma_type), variants in groups.items():
+        by_kv_cache_dtype.setdefault(kv_cache_dtype, {}).setdefault(
+            query_scalar_type, {}
+        )[mfma_type] = variants
+
+    lines = [DISPATCH_HEAD.rstrip("\n")]
+    kv_cache_dtypes = list(by_kv_cache_dtype)
+    for index, kv_cache_dtype in enumerate(kv_cache_dtypes):
+        keyword = "if" if index == 0 else "else if"
+        predicate = KV_CACHE_DTYPE_PREDICATES[kv_cache_dtype]
+        lines.append(f"  {keyword} ({predicate}) {{")
+
+        by_query_dtype = by_kv_cache_dtype[kv_cache_dtype]
+        for dtype_index, (query_scalar_type, by_mfma) in enumerate(
+            by_query_dtype.items()
+        ):
+            dtype_keyword = "if" if dtype_index == 0 else "else if"
+            lines.append(
+                f"    {dtype_keyword} (query.dtype() == "
+                f"at::ScalarType::{query_scalar_type}) {{"
+            )
+
+            mfma_types = list(by_mfma)
+            # A lone arm is the catch-all and takes the no-split path: None when
+            # mfma_type is not a dimension here, "f16" when it is the only arm left
+            # (a bare `else {` would otherwise open a block with no `if`).
+            if mfma_types in ([None], ["f16"]):
+                lines.extend(block_size_switch(by_mfma[mfma_types[0]], 6))
+            else:
+                # ``mfma_type`` is matched with a plain ``else`` on the last arm,
+                # as the hand-written dispatch used to: any value other than the
+                # recognised ones keeps selecting that arm instead of failing.
+                if "f16" not in mfma_types:
+                    raise SystemExit(
+                        f"the {mfma_types} group needs an f16 arm, which is "
+                        "emitted as the catch-all else"
+                    )
+                ordered = sorted(mfma_types, key=lambda m: m == "f16")  # f16 last
+                for mfma_index, mfma_type in enumerate(ordered):
+                    if mfma_index == len(ordered) - 1:
+                        lines.append("      else {")
+                    else:
+                        mfma_keyword = "if" if mfma_index == 0 else "else if"
+                        lines.append(
+                            f'      {mfma_keyword} (mfma_type == "{mfma_type}") {{'
+                        )
+                    lines.extend(block_size_switch(by_mfma[mfma_type], 8))
+                    lines.append("      }")
+            lines.append("    }")
+
+        lines.append("    else {")
+        lines.append(
+            '      TORCH_CHECK(false, "Unsupported data type: ", query.dtype());'
+        )
+        lines.append("    }")
+        lines.append("  }")
+
+    lines.append("  else {")
+    lines.append(
+        '    TORCH_CHECK(false, "Unsupported KV cache dtype: ", kv_cache_dtype);'
+    )
+    lines.append("  }")
+    return "\n".join(lines) + "\n"
+
+
+def generate_entries():
+    lines = [ENTRIES_HEAD.rstrip("\n")]
+    for variant in VARIANTS:
+        for block_size in BLOCK_SIZES:
+            lines.append(f"ATTN_ENTRY({entry_point(variant, block_size)});")
+    return "\n".join(lines) + "\n"
+
+
+def generate_tu(variant, block_size):
+    return TU_HEAD + TU_ENTRY.format(block_size=block_size, **variant)
+
+
+def remove_old_files():
+    # Only remove files this script produced: a name match alone would delete a
+    # developer's scratch copy in the same directory.
+    for path in glob.glob(os.path.join(OUTPUT_DIR, "attention_paged_attn_*.cu")):
+        with open(path) as fd:
+            if fd.readline() != FILE_HEAD.splitlines()[0] + "\n":
+                continue
+        os.remove(path)
+    for name in (
+        "attention_paged_attn_entries.inc",
+        "attention_paged_attn_dispatch.inc",
+    ):
+        path = os.path.join(OUTPUT_DIR, name)
+        if os.path.exists(path):
+            os.remove(path)
+
+
+def write(name, content):
+    with open(os.path.join(OUTPUT_DIR, name), "w") as fd:
+        fd.write(content)
+
+
+def check_dispatch_covers_entries(entries, dispatch):
+    # Cheap replacement for the configure-time guard that used to regex
+    # attention_dispatch.cu: every instantiated entry point must be named by the
+    # generated dispatch, otherwise its unit is compiled and linked for nothing.
+    for name in entries:
+        if f"{name}(" not in dispatch:
+            raise SystemExit(
+                f"{name} is instantiated but never dispatched: add it to "
+                "VARIANTS ordering/generate_dispatch(), or drop it."
+            )
+
+
+def main():
+    remove_old_files()
+    outputs = []
+    entries = []
+    for variant in VARIANTS:
+        for block_size in BLOCK_SIZES:
+            name = f"attention_paged_attn_{variant['name']}_bs{block_size}.cu"
+            write(name, generate_tu(variant, block_size))
+            outputs.append(name)
+            entries.append(entry_point(variant, block_size))
+    entries_inc = generate_entries()
+    dispatch_inc = generate_dispatch()
+    check_dispatch_covers_entries(entries, dispatch_inc)
+    for name, content in (
+        ("attention_paged_attn_entries.inc", entries_inc),
+        ("attention_paged_attn_dispatch.inc", dispatch_inc),
+    ):
+        write(name, content)
+        outputs.append(name)
+    # Record what is on disk next to the outputs: CMake takes this list as its
+    # source list and compares the key (this script's hash) so that units left by
+    # another version of the script are regenerated. The key is arch-independent.
+    write(STAMP_NAME, KEY + "\n" + "\n".join(outputs) + "\n")
+    print(f"Generated {len(outputs)} paged attention files.")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/csrc/sampled_topk.cuh b/csrc/sampled_topk.cuh
index c25f8ed8e731..9923dcbf417e 100644
--- a/csrc/sampled_topk.cuh
+++ b/csrc/sampled_topk.cuh
@@ -69,9 +69,8 @@ __device__ void select_bin(Storage& s) {
 
 // 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) {
+template <bool Buffered, typename InT>
+__device__ uint32_t threshold(Storage& s, const InT* row, int count, int rank) {
   const int tid = threadIdx.x;
   if (tid == 0) {
     s.remaining = rank;
@@ -81,7 +80,8 @@ __device__ uint32_t threshold(Storage& s, const float* row, int count,
   __syncthreads();
   for (int i = tid; i < count; i += kThreads) {
     const uint32_t key =
-        Buffered ? s.candidates[i].y : hist4096::convert_to_uint32_v2(row[i]);
+        Buffered ? s.candidates[i].y
+                 : hist4096::convert_to_uint32_v2(hist4096::to_float(row[i]));
     atomicAdd(&s.histogram[coarse(key)], 1);
   }
   __syncthreads();
@@ -94,7 +94,8 @@ __device__ uint32_t threshold(Storage& s, const float* row, int count,
     const uint32_t prefix = s.prefix;
     for (int i = tid; i < count; i += kThreads) {
       const uint32_t key =
-          Buffered ? s.candidates[i].y : hist4096::convert_to_uint32_v2(row[i]);
+          Buffered ? s.candidates[i].y
+                   : hist4096::convert_to_uint32_v2(hist4096::to_float(row[i]));
       if (coarse(key) == coarse_bin && (key & mask) == prefix) {
         atomicAdd(&s.histogram[(key >> shift) & 255], 1);
       }
@@ -118,8 +119,8 @@ __device__ __forceinline__ int reserve(bool hit, int* counter) {
   return base + __popc(mask & ((1u << lane) - 1));
 }
 
-template <bool Buffered>
-__device__ void emit(Storage& s, const float* row, int32_t* dst, int count,
+template <bool Buffered, typename InT>
+__device__ void emit(Storage& s, const InT* row, int32_t* dst, int count,
                      uint32_t cutoff) {
   const int tid = threadIdx.x;
   if (tid == 0) s.count = 0;
@@ -127,10 +128,10 @@ __device__ void emit(Storage& s, const float* row, int32_t* dst, int count,
   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, hist4096::convert_to_uint32_v2(row[i])))
-            : make_uint2(0, 0);
+        i < count ? (Buffered ? s.candidates[i]
+                              : make_uint2(i, hist4096::convert_to_uint32_v2(
+                                                  hist4096::to_float(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;
@@ -140,9 +141,9 @@ __device__ void emit(Storage& s, const float* row, int32_t* dst, int count,
   }
 }
 
-template <int K>
+template <int K, typename InT>
 __global__ void __launch_bounds__(kThreads)
-    sampled_topk_kernel(const float* __restrict__ input,
+    sampled_topk_kernel(const InT* __restrict__ input,
                         const int32_t* __restrict__ lengths,
                         int32_t* __restrict__ output, int64_t stride,
                         int max_length) {
@@ -150,7 +151,7 @@ __global__ void __launch_bounds__(kThreads)
   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;
+  const InT* 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>) {
@@ -158,11 +159,11 @@ __global__ void __launch_bounds__(kThreads)
     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>(
+          filtered_topk::filtered_topk_row<InT, 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>(
+          filtered_topk::filtered_topk_row<InT, int32_t, 1, K, true, true>(
               row, dst, length, K, fallback_storage);
     }
     if (!selected) {
@@ -177,7 +178,8 @@ __global__ void __launch_bounds__(kThreads)
   // 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(hist4096::convert_to_uint32_v2(row[index]))],
+    atomicAdd(&s.histogram[coarse(hist4096::convert_to_uint32_v2(
+                  hist4096::to_float(row[index])))],
               1);
   }
   __syncthreads();
@@ -200,7 +202,7 @@ __global__ void __launch_bounds__(kThreads)
 #pragma unroll
       for (int j = 0; j < kItems / 4; ++j) {
         const float4 values =
-            reinterpret_cast<const float4*>(row + base)[tid + j * kThreads];
+            hist4096::load_score4(row + base + (tid + j * kThreads) * 4);
         keys[4 * j] = hist4096::convert_to_uint32_v2(values.x);
         keys[4 * j + 1] = hist4096::convert_to_uint32_v2(values.y);
         keys[4 * j + 2] = hist4096::convert_to_uint32_v2(values.z);
@@ -210,7 +212,10 @@ __global__ void __launch_bounds__(kThreads)
 #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 ? hist4096::convert_to_uint32_v2(row[i]) : 0;
+        keys[j] =
+            i < length
+                ? hist4096::convert_to_uint32_v2(hist4096::to_float(row[i]))
+                : 0;
       }
     }
 #pragma unroll
diff --git a/csrc/topk.cu b/csrc/topk.cu
index d5ff2a8386eb..1322ca3b0dd4 100644
--- a/csrc/topk.cu
+++ b/csrc/topk.cu
@@ -15,7 +15,7 @@
 namespace {
 
 #ifndef USE_ROCM
-template <int TopK>
+template <int TopK, typename InT>
 void launch_persistent_topk(const torch::stable::Tensor& logits,
                             const torch::stable::Tensor& lengths,
                             torch::stable::Tensor& output,
@@ -28,6 +28,7 @@ void launch_persistent_topk(const torch::stable::Tensor& logits,
   const int64_t num_rows = logits.size(0);
   const int64_t stride = logits.stride(0);
   const cudaStream_t stream = get_current_cuda_stream();
+  const auto* input = reinterpret_cast<const InT*>(logits.const_data_ptr());
 
   static int num_sms = 0;
   static int max_smem_per_block = 0;
@@ -37,11 +38,15 @@ void launch_persistent_topk(const torch::stable::Tensor& logits,
     max_smem_per_block = device_prop->sharedMemPerBlockOptin;
   }
 
+  STD_TORCH_CHECK(
+      sizeof(InT) == sizeof(float) || max_smem_per_block >= 128 * 1024,
+      "bfloat16 persistent_topk requires 128 KiB shared memory");
+
   // 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>;
+    auto kernel = vllm::sampled_topk::sampled_topk_kernel<TopK, InT>;
     constexpr size_t smem_size =
         vllm::filtered_topk::FILTERED_TOPK_SMEM_DYNAMIC;
     cudaError_t status = cudaFuncSetAttribute(
@@ -49,15 +54,14 @@ void launch_persistent_topk(const torch::stable::Tensor& logits,
     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>(),
+        input, 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>(),
-            lengths.const_data_ptr<int32_t>(), static_cast<uint32_t>(num_rows),
-            static_cast<uint32_t>(TopK), static_cast<uint32_t>(stride), stream);
+    cudaError_t status = vllm::FilteredTopKRaggedTransform<InT, int32_t, TopK>(
+        input, output.mutable_data_ptr<int32_t>(),
+        lengths.const_data_ptr<int32_t>(), static_cast<uint32_t>(num_rows),
+        static_cast<uint32_t>(TopK), static_cast<uint32_t>(stride), stream);
     STD_TORCH_CHECK(status == cudaSuccess,
                     "FilteredTopK failed: ", cudaGetErrorString(status));
   } else {
@@ -83,8 +87,7 @@ void launch_persistent_topk(const torch::stable::Tensor& logits,
         static_cast<uint32_t>(available_for_ordered / sizeof(uint32_t));
 
     // Views like x[:, 1:] keep the stride but shift the base.
-    const auto base =
-        reinterpret_cast<uintptr_t>(logits.const_data_ptr<float>());
+    const auto base = reinterpret_cast<uintptr_t>(input);
     uint32_t vec_size = 1;
     if (stride % 4 == 0 && base % 16 == 0)
       vec_size = 4;
@@ -112,16 +115,16 @@ void launch_persistent_topk(const torch::stable::Tensor& logits,
     cudaError_t occ_err = cudaSuccess;
     if (vec_size == 4) {
       occ_err = cudaOccupancyMaxActiveBlocksPerMultiprocessor(
-          &occupancy, P::persistent_topk_kernel<TopK, 4>, P::kThreadsPerBlock,
-          smem_size);
+          &occupancy, P::persistent_topk_kernel<TopK, 4, InT>,
+          P::kThreadsPerBlock, smem_size);
     } else if (vec_size == 2) {
       occ_err = cudaOccupancyMaxActiveBlocksPerMultiprocessor(
-          &occupancy, P::persistent_topk_kernel<TopK, 2>, P::kThreadsPerBlock,
-          smem_size);
+          &occupancy, P::persistent_topk_kernel<TopK, 2, InT>,
+          P::kThreadsPerBlock, smem_size);
     } else {
       occ_err = cudaOccupancyMaxActiveBlocksPerMultiprocessor(
-          &occupancy, P::persistent_topk_kernel<TopK, 1>, P::kThreadsPerBlock,
-          smem_size);
+          &occupancy, P::persistent_topk_kernel<TopK, 1, InT>,
+          P::kThreadsPerBlock, smem_size);
     }
     STD_TORCH_CHECK(occ_err == cudaSuccess,
                     "persistent_topk occupancy query failed: ",
@@ -162,9 +165,8 @@ void launch_persistent_topk(const torch::stable::Tensor& logits,
         return;
       }
       cudaError_t status =
-          vllm::FilteredTopKRaggedTransform<float, int32_t, TopK>(
-              logits.const_data_ptr<float>(),
-              output.mutable_data_ptr<int32_t>(),
+          vllm::FilteredTopKRaggedTransform<InT, int32_t, TopK>(
+              input, output.mutable_data_ptr<int32_t>(),
               lengths.const_data_ptr<int32_t>(),
               static_cast<uint32_t>(num_rows), static_cast<uint32_t>(TopK),
               static_cast<uint32_t>(stride), stream);
@@ -209,8 +211,8 @@ void launch_persistent_topk(const torch::stable::Tensor& logits,
                       "row_states memset failed: ", cudaGetErrorString(mz_err));
     }
 
-    P::PersistentTopKParams params;
-    params.input = logits.const_data_ptr<float>();
+    P::PersistentTopKParams<InT> params;
+    params.input = input;
     params.output = output.mutable_data_ptr<int32_t>();
     params.lengths = lengths.const_data_ptr<int32_t>();
     params.num_rows = static_cast<uint32_t>(num_rows);
@@ -224,7 +226,7 @@ void launch_persistent_topk(const torch::stable::Tensor& logits,
 
   #define LAUNCH_PERSISTENT(TOPK_VAL, VS)                                     \
     do {                                                                      \
-      auto kernel = &P::persistent_topk_kernel<TOPK_VAL, VS>;                 \
+      auto kernel = &P::persistent_topk_kernel<TOPK_VAL, VS, InT>;            \
       cudaError_t err = cudaFuncSetAttribute(                                 \
           kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);    \
       STD_TORCH_CHECK(err == cudaSuccess,                                     \
@@ -246,6 +248,21 @@ void launch_persistent_topk(const torch::stable::Tensor& logits,
   STD_TORCH_CHECK(err == cudaSuccess,
                   "persistent_topk failed: ", cudaGetErrorString(err));
 }
+
+template <int TopK>
+void launch_persistent_topk_dtype(const torch::stable::Tensor& logits,
+                                  const torch::stable::Tensor& lengths,
+                                  torch::stable::Tensor& output,
+                                  torch::stable::Tensor& workspace,
+                                  int64_t max_seq_len) {
+  if (logits.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
+    launch_persistent_topk<TopK, __nv_bfloat16>(logits, lengths, output,
+                                                workspace, max_seq_len);
+  } else {
+    launch_persistent_topk<TopK, float>(logits, lengths, output, workspace,
+                                        max_seq_len);
+  }
+}
 #endif
 
 }  // anonymous namespace
@@ -259,8 +276,14 @@ void persistent_topk(const torch::stable::Tensor& logits,
   STD_TORCH_CHECK(logits.is_cuda(), "logits must be CUDA tensor");
   STD_TORCH_CHECK(lengths.is_cuda(), "lengths must be CUDA tensor");
   STD_TORCH_CHECK(output.is_cuda(), "output must be CUDA tensor");
-  STD_TORCH_CHECK(logits.scalar_type() == torch::headeronly::ScalarType::Float,
-                  "Only float32 supported");
+  STD_TORCH_CHECK(
+      logits.scalar_type() == torch::headeronly::ScalarType::Float ||
+          logits.scalar_type() == torch::headeronly::ScalarType::BFloat16,
+      "logits must be float32 or bfloat16");
+  STD_TORCH_CHECK(
+      logits.scalar_type() != torch::headeronly::ScalarType::BFloat16 ||
+          logits.stride(1) == 1,
+      "bfloat16 logits columns must be contiguous");
   STD_TORCH_CHECK(lengths.scalar_type() == torch::headeronly::ScalarType::Int,
                   "lengths must be int32");
   STD_TORCH_CHECK(output.scalar_type() == torch::headeronly::ScalarType::Int,
@@ -284,14 +307,14 @@ void persistent_topk(const torch::stable::Tensor& logits,
       logits.get_device_index());
 
   if (k == 512) {
-    launch_persistent_topk<512>(logits, lengths, output, workspace,
-                                max_seq_len);
+    launch_persistent_topk_dtype<512>(logits, lengths, output, workspace,
+                                      max_seq_len);
   } else if (k == 1024) {
-    launch_persistent_topk<1024>(logits, lengths, output, workspace,
-                                 max_seq_len);
+    launch_persistent_topk_dtype<1024>(logits, lengths, output, workspace,
+                                       max_seq_len);
   } else {
-    launch_persistent_topk<2048>(logits, lengths, output, workspace,
-                                 max_seq_len);
+    launch_persistent_topk_dtype<2048>(logits, lengths, output, workspace,
+                                       max_seq_len);
   }
 #else
   STD_TORCH_CHECK(false, "persistent_topk is not supported on ROCm");
diff --git a/csrc/topk_histogram_4096.cuh b/csrc/topk_histogram_4096.cuh
index b5aeb2d78cd6..86e7b680461e 100644
--- a/csrc/topk_histogram_4096.cuh
+++ b/csrc/topk_histogram_4096.cuh
@@ -5,6 +5,7 @@
 #ifndef TOPK_HISTOGRAM_4096_CUH_
 #define TOPK_HISTOGRAM_4096_CUH_
 
+#include <cuda_bf16.h>
 #include <cuda_fp16.h>
 #include <cuda_runtime.h>
 #include <cstdint>
@@ -66,6 +67,39 @@ __device__ __forceinline__ void load_float4_predicated(const float* ptr,
   v3 = __uint_as_float(r3);
 }
 
+__device__ __forceinline__ float to_float(float x) { return x; }
+__device__ __forceinline__ float to_float(__nv_bfloat16 x) {
+  return __bfloat162float(x);
+}
+
+__device__ __forceinline__ float4 load_score4(const float* ptr) {
+  return *reinterpret_cast<const float4*>(ptr);
+}
+__device__ __forceinline__ float4 load_score4(const __nv_bfloat16* ptr) {
+  const uint2 raw = *reinterpret_cast<const uint2*>(ptr);
+  return make_float4(
+      __uint_as_float(raw.x << 16), __uint_as_float(raw.x & 0xFFFF0000u),
+      __uint_as_float(raw.y << 16), __uint_as_float(raw.y & 0xFFFF0000u));
+}
+__device__ __forceinline__ float2 load_score2(const float* ptr) {
+  return *reinterpret_cast<const float2*>(ptr);
+}
+__device__ __forceinline__ float2 load_score2(const __nv_bfloat16* ptr) {
+  const uint32_t raw = *reinterpret_cast<const uint32_t*>(ptr);
+  return make_float2(__uint_as_float(raw << 16),
+                     __uint_as_float(raw & 0xFFFF0000u));
+}
+
+__device__ __forceinline__ void load_float4_predicated(const __nv_bfloat16* ptr,
+                                                       int base, int seq_len,
+                                                       float& v0, float& v1,
+                                                       float& v2, float& v3) {
+  v0 = base < seq_len ? to_float(ptr[0]) : __uint_as_float(0xFF800000u);
+  v1 = base + 1 < seq_len ? to_float(ptr[1]) : __uint_as_float(0xFF800000u);
+  v2 = base + 2 < seq_len ? to_float(ptr[2]) : __uint_as_float(0xFF800000u);
+  v3 = base + 3 < seq_len ? to_float(ptr[3]) : __uint_as_float(0xFF800000u);
+}
+
 // converts the float32 score to a 32-bit ordered unsigned integer — the full
 // precision key for radix sorting
 __device__ __forceinline__ auto convert_to_uint32_v2(float x) -> uint32_t {
@@ -333,9 +367,9 @@ struct Histogram4096Smem {
 
 // Defined in topk_histogram_4096_overflow.cuh. Only kFp32Rescan instantiations
 // call it, and their callers include that header.
-template <uint32_t TopK, uint32_t BlockSize, uint32_t VecSize>
+template <uint32_t TopK, uint32_t BlockSize, uint32_t VecSize, typename InT>
 __device__ __noinline__ void exact_topk_rescan_cold(
-    const float* __restrict__ scores, int32_t* __restrict__ output,
+    const InT* __restrict__ scores, int32_t* __restrict__ output,
     uint32_t length, void* smem);
 
 // kRegisterValues refines from the values already held in registers and
@@ -349,16 +383,17 @@ __device__ __forceinline__ bool exact_topk_refine_registers(
     uint32_t length, uint32_t coarse_bin, uint32_t num_above, void* _smem);
 
 template <uint32_t TopK, uint32_t BlockSize, uint32_t VecSize,
-          uint32_t StartRound, uint32_t FinalCandidateCapacity>
-__device__ void exact_topk_rescan(const float* __restrict__ scores,
+          uint32_t StartRound, uint32_t FinalCandidateCapacity, typename InT>
+__device__ void exact_topk_rescan(const InT* __restrict__ scores,
                                   int32_t* __restrict__ output, uint32_t length,
                                   void* _smem, uint32_t initial_prefix,
                                   uint32_t initial_remaining);
 
 template <uint32_t TopK, uint32_t HIST_BITS,
           uint32_t VECS_PER_THREAD = kHist4096VecsPerThread,
-          bool UsePredicatedLoads = false, OverflowRecovery Recovery>
-__device__ void histogram_4096_topk(const float* __restrict__ scores,
+          bool UsePredicatedLoads = false, OverflowRecovery Recovery,
+          typename InT>
+__device__ void histogram_4096_topk(const InT* __restrict__ scores,
                                     int32_t* __restrict__ output,
                                     uint32_t length, void* _smem) {
   constexpr uint32_t HIST_BINS = 1 << HIST_BITS;
@@ -388,14 +423,14 @@ __device__ void histogram_4096_topk(const float* __restrict__ scores,
     smem->counter_gt = 0;
     smem->counter_eq = 0;
   }
-  if constexpr (UsePredicatedLoads) {
+  if constexpr (UsePredicatedLoads || sizeof(InT) == 2) {
     const bool row_aligned = (reinterpret_cast<uintptr_t>(scores) & 0xFu) == 0;
 #pragma unroll
     for (uint32_t v = 0; v < VECS_PER_THREAD; v++) {
       const uint32_t base = (tx + v * kBlockSize) * 4;
       if (base < length) {
         if (row_aligned && base + 3 < length) {
-          vecs[v] = *reinterpret_cast<const float4*>(scores + base);
+          vecs[v] = load_score4(scores + base);
         } else {
           load_float4_predicated(scores + base, static_cast<int>(base),
                                  static_cast<int>(length), vecs[v].x, vecs[v].y,
@@ -408,7 +443,7 @@ __device__ void histogram_4096_topk(const float* __restrict__ scores,
     for (uint32_t v = 0; v < VECS_PER_THREAD; v++) {
       const uint32_t base = (tx + v * kBlockSize) * 4;
       if (base < length) {
-        vecs[v] = *reinterpret_cast<const float4*>(scores + base);
+        vecs[v] = load_score4(scores + base);
       }
     }
   }
@@ -596,9 +631,9 @@ __device__ void histogram_4096_topk(const float* __restrict__ scores,
 
 template <uint32_t TopK, uint32_t HIST_BITS,
           uint32_t VECS_PER_THREAD = kHist4096VecsPerThread,
-          OverflowRecovery Recovery>
+          OverflowRecovery Recovery, typename InT>
 __device__ __noinline__ void histogram_4096_topk_predicated(
-    const float* __restrict__ scores, int32_t* __restrict__ output,
+    const InT* __restrict__ scores, int32_t* __restrict__ output,
     uint32_t length, void* _smem) {
   histogram_4096_topk<TopK, HIST_BITS, VECS_PER_THREAD, true, Recovery>(
       scores, output, length, _smem);
diff --git a/csrc/topk_histogram_4096_overflow.cuh b/csrc/topk_histogram_4096_overflow.cuh
index 7433a0362acf..f788cadac2c9 100644
--- a/csrc/topk_histogram_4096_overflow.cuh
+++ b/csrc/topk_histogram_4096_overflow.cuh
@@ -9,15 +9,15 @@ namespace vllm {
 namespace topk_histogram_4096 {
 
 // Visit a row with vectorized loads and a scalar tail.
-template <uint32_t BlockSize, uint32_t VecSize, typename Visit>
-__device__ __forceinline__ void scan_scores(const float* __restrict__ scores,
+template <uint32_t BlockSize, uint32_t VecSize, typename InT, typename Visit>
+__device__ __forceinline__ void scan_scores(const InT* __restrict__ scores,
                                             uint32_t length, Visit visit) {
   static_assert(VecSize == 1 || VecSize == 2 || VecSize == 4);
   const uint32_t tx = threadIdx.x;
 
   if constexpr (VecSize == 1) {
     for (uint32_t idx = tx; idx < length; idx += BlockSize) {
-      visit(idx, scores[idx]);
+      visit(idx, to_float(scores[idx]));
     }
     return;
   }
@@ -27,20 +27,20 @@ __device__ __forceinline__ void scan_scores(const float* __restrict__ scores,
   for (uint32_t base = tx * VecSize; base < aligned_length;
        base += BlockSize * VecSize) {
     if constexpr (VecSize == 4) {
-      const float4 values = *reinterpret_cast<const float4*>(scores + base);
+      const float4 values = load_score4(scores + base);
       visit(base, values.x);
       visit(base + 1, values.y);
       visit(base + 2, values.z);
       visit(base + 3, values.w);
     } else {
-      const float2 values = *reinterpret_cast<const float2*>(scores + base);
+      const float2 values = load_score2(scores + base);
       visit(base, values.x);
       visit(base + 1, values.y);
     }
   }
 
   for (uint32_t idx = aligned_length + tx; idx < length; idx += BlockSize) {
-    visit(idx, scores[idx]);
+    visit(idx, to_float(scores[idx]));
   }
 }
 
@@ -70,8 +70,9 @@ constexpr uint32_t kExactCandidateOffset = 16 * 1024;
 // resume at StartRound from a known prefix, and with FinalCandidateCapacity
 // the last round stages its candidates so the output pass skips a row scan.
 template <uint32_t TopK, uint32_t BlockSize, uint32_t VecSize = 1,
-          uint32_t StartRound = 0, uint32_t FinalCandidateCapacity = 0>
-__device__ void exact_topk_rescan(const float* __restrict__ scores,
+          uint32_t StartRound = 0, uint32_t FinalCandidateCapacity = 0,
+          typename InT>
+__device__ void exact_topk_rescan(const InT* __restrict__ scores,
                                   int32_t* __restrict__ output, uint32_t length,
                                   void* _smem, uint32_t initial_prefix = 0,
                                   uint32_t initial_remaining = TopK) {
@@ -215,7 +216,7 @@ __device__ void exact_topk_rescan(const float* __restrict__ scores,
     if (candidate_count <= FinalCandidateCapacity) {
       for (uint32_t i = tx; i < candidate_count; i += BlockSize) {
         const uint32_t idx = final_candidates[i];
-        collect(idx, scores[idx]);
+        collect(idx, to_float(scores[idx]));
       }
     } else if (smem->pivot_count == candidate_count) {
       // Every key under the final prefix equals the pivot (heavy ties), so
@@ -240,9 +241,9 @@ __device__ void exact_topk_rescan(const float* __restrict__ scores,
   __syncthreads();
 }
 
-template <uint32_t TopK, uint32_t BlockSize, uint32_t VecSize>
+template <uint32_t TopK, uint32_t BlockSize, uint32_t VecSize, typename InT>
 __device__ __noinline__ void exact_topk_rescan_cold(
-    const float* __restrict__ scores, int32_t* __restrict__ output,
+    const InT* __restrict__ scores, int32_t* __restrict__ output,
     uint32_t length, void* smem) {
   exact_topk_rescan<TopK, BlockSize, VecSize>(scores, output, length, smem);
 }
diff --git a/csrc/torch_bindings.cpp b/csrc/torch_bindings.cpp
index fef3076056ee..9969f6309b8f 100644
--- a/csrc/torch_bindings.cpp
+++ b/csrc/torch_bindings.cpp
@@ -8,6 +8,17 @@
 // Note: We register under namespace "_C" so ops are accessible as
 // torch.ops._C.<op_name> for compatibility with existing code.
 STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
+#ifdef VLLM_ENABLE_FUSED_MQA_TOPK
+  // Single-launch fused FP8 MQA-logits + top-k (512/1024/2048) prefill
+  // selection (SM100 family). Workspace tensors are caller-owned; counter must
+  // be zero.
+  ops.def(
+      "fused_mqa_topk_prefill(Tensor q, Tensor k, Tensor scales, "
+      "Tensor weights, Tensor starts, Tensor ends, Tensor! out, "
+      "Tensor! counter, Tensor! values, Tensor! indices, Tensor! retained, "
+      "Tensor! retained_indices, Tensor! counts, Tensor! defer) -> ()");
+#endif
+
   // Compute per-token-group FP8 quantized tensor and scaling factor.
   // The dummy arguments are here so we can correctly fuse with RMSNorm.
   ops.def(
@@ -538,7 +549,17 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
       "int block_size, Tensor!? q_out, Tensor!? index_q_out, "
       "str kv_cache_dtype, bool skip_index_branch=False, "
       "Tensor!? q_fp8_out=None, float q_fp8_scale=1.0, "
-      "Tensor? kv_k_scale=None, Tensor? kv_v_scale=None) -> ()");
+      "Tensor? kv_k_scale=None, Tensor? kv_v_scale=None, "
+      "bool write_qkv=True) -> ()");
+
+  // MiniMax-M3 DSpark draft: Gemma k-norm + partial NeoX RoPE + K/V cache
+  // insert for every draft layer's context K/V in one launch.
+  ops.def(
+      "fused_minimax_m3_dspark_context_kv("
+      "Tensor all_kv, Tensor k_norm_weight, Tensor cos_sin_cache, "
+      "Tensor positions, Tensor slot_mapping, Tensor! kv_cache, "
+      "Tensor cache_offsets, Tensor k_scale, Tensor scale_offsets, "
+      "int rotary_dim, float eps, str kv_cache_dtype) -> ()");
 
 #ifdef VLLM_ENABLE_FUSED_KDA_DECODE
   ops.def(
@@ -782,6 +803,10 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
 }
 
 STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
+#ifdef VLLM_ENABLE_FUSED_MQA_TOPK
+  ops.impl("fused_mqa_topk_prefill", TORCH_BOX(&fused_mqa_topk_prefill));
+#endif
+
   // LongCat n-gram embedding index kernel.
   ops.impl("ngram_compute_n_gram_ids", TORCH_BOX(&ngram_compute_n_gram_ids));
   ops.impl("vocab_parallel_embedding", TORCH_BOX(&vocab_parallel_embedding));
@@ -879,6 +904,8 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
            TORCH_BOX(&fused_muse_spark_qknorm_rope_kv_insert));
   ops.impl("fused_minimax_m3_qknorm_rope_kv_insert",
            TORCH_BOX(&fused_minimax_m3_qknorm_rope_kv_insert));
+  ops.impl("fused_minimax_m3_dspark_context_kv",
+           TORCH_BOX(&fused_minimax_m3_dspark_context_kv));
 #ifdef VLLM_ENABLE_FUSED_KDA_DECODE
   ops.impl("fused_kda_decode", TORCH_BOX(&fused_kda_decode));
 #endif
@@ -1064,13 +1091,16 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C_cache_ops, ops) {
       "                 Tensor(h!)? swap_counts=None,"
       "                 Tensor? resident_block_table=None,"
       "                 int resident_block_size=0,"
-      "                 int resident_null_block=0) -> ()");
+      "                 int resident_null_block=0,"
+      "                 Tensor? num_valid_rows=None) -> ()");
 
   ops.def(
-      "hisparse_invalidate_written_slots(Tensor! device_global_indices,"
+      "hisparse_invalidate_written_slots(Tensor device_global_indices_ptrs,"
       "                                   Tensor request_state_indices,"
       "                                   Tensor req_id_per_token,"
-      "                                   Tensor written_slots) -> ()");
+      "                                   Tensor written_slots,"
+      "                                   int num_state_rows,"
+      "                                   int region_stride) -> ()");
 
   ops.def(
       "hisparse_gather_plan(Tensor host_cache,"
diff --git a/docker/Dockerfile b/docker/Dockerfile
index f3cac42652e9..ff5b273c56fd 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -93,11 +93,6 @@ ARG UV_KEYRING_PROVIDER=${PIP_KEYRING_PROVIDER}
 # Flag enables built-in KV-connector dependency libs into docker images
 ARG INSTALL_KV_CONNECTORS=false
 
-# Flag builds NIXL, including its NIXL EP extension for the installed PyTorch,
-# from source and installs it in place of the nixl packages from the KV-connector
-# requirements
-ARG BUILD_NIXL=false
-
 #################### BASE BUILD IMAGE ####################
 # prepare basic build environment
 FROM ${BUILD_BASE_IMAGE} AS base
@@ -253,8 +248,13 @@ RUN --mount=type=cache,target=/opt/uv/cache \
         && uv pip install --python /opt/venv/bin/python3 -r requirements/cuda.txt \
         --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/nightly/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
     else \
+        # Keep the CUDA-specific PyTorch index first. Using it as an extra
+        # index with unsafe-best-match can select a higher, incompatible
+        # torch build from the base image (for example cu130 in a cu129 image).
         uv pip install --python /opt/venv/bin/python3 -r requirements/cuda.txt \
-        --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
+        --index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \
+        --extra-index-url https://pypi.org/simple \
+        --index-strategy unsafe-first-match; \
     fi
 
 # Track PyTorch lib versions used during build and match in downstream instances.
@@ -372,7 +372,9 @@ RUN --mount=type=cache,target=/opt/uv/cache \
     else \
         echo "Installing build requirements..." \
         && uv pip install --python /opt/venv/bin/python3 -r requirements/build/cuda.txt \
-        --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
+        --index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \
+        --extra-index-url https://pypi.org/simple \
+        --index-strategy unsafe-first-match; \
     fi
 
 WORKDIR /workspace
@@ -561,24 +563,6 @@ RUN --mount=type=cache,target=/opt/uv/cache \
         (cd "${TRITON_WORKSPACE}/dist" \
             && sha256sum build-state > wheel.sha256); \
     fi
-
-# Build NIXL and its NIXL EP extension against the PyTorch installed here when
-# requested. The wheels are installed by the final image, which has no compiler.
-ARG BUILD_NIXL
-COPY tools/build_nixl_with_torch215_for_rubin.py /tmp/build_nixl_for_rubin.py
-RUN --mount=type=bind,source=requirements,target=/tmp/vllm-requirements,ro \
-    set -eux; \
-    NIXL_WORKSPACE=/tmp/nixl_workspace; \
-    mkdir -p "${NIXL_WORKSPACE}/dist"; \
-    if [ "${BUILD_NIXL}" = "true" ]; then \
-        dnf install -y --setopt=install_weak_deps=False autoconf automake libtool; \
-        dnf clean all; \
-        rm -rf /var/cache/dnf; \
-        NIXL_REQUIREMENTS_FILE=/tmp/vllm-requirements/kv_connectors.txt \
-        WHEELS_CACHE_HOME="${NIXL_WORKSPACE}/dist" \
-        MAX_JOBS="${max_jobs:-}" \
-            python3 /tmp/build_nixl_for_rubin.py --skip-system-deps --build-only; \
-    fi
 #################### EXTENSIONS BUILD IMAGE ####################
 
 #################### WHEEL BUILD IMAGE ####################
@@ -732,7 +716,9 @@ RUN --mount=type=cache,target=/opt/uv/cache \
     else \
         echo "Installing dev requirements..." \
         && uv pip install --python /opt/venv/bin/python3 -r requirements/dev.txt \
-        --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
+        --index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \
+        --extra-index-url https://pypi.org/simple \
+        --index-strategy unsafe-first-match; \
     fi
 
 #################### DEV IMAGE ####################
@@ -941,7 +927,9 @@ RUN --mount=type=cache,target=/opt/uv/cache \
     else \
         uv pip install --system \
             -r /tmp/requirements-cuda.txt \
-            --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
+            --index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \
+            --extra-index-url https://pypi.org/simple \
+            --index-strategy unsafe-first-match; \
     fi && \
     rm /tmp/requirements-cuda.txt /tmp/common.txt
 
@@ -976,7 +964,7 @@ RUN set -eux; \
 # From versions.json: .openai_server_extras.timm, .openai_server_extras.runai_model_streamer,
 # .openai_server_extras.gguf (NVFP4 GGUF checkpoints, e.g. RecursiveSI/flint-nvfp4-gguf)
 ARG TIMM_VERSION=">=1.0.17"
-ARG RUNAI_MODEL_STREAMER_VERSION=">=0.15.7"
+ARG RUNAI_MODEL_STREAMER_VERSION=">=0.17.1"
 ARG GGUF_VERSION=">=0.19.0"
 RUN --mount=type=cache,target=/opt/uv/cache \
     uv pip install --system accelerate 'modelscope<1.38' \
@@ -1017,7 +1005,9 @@ RUN --mount=type=bind,from=build,src=/workspace/dist,target=/vllm-workspace/dist
     else \
         echo "Installing vLLM..." \
         && uv pip install --system dist/*.whl --verbose \
-        --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
+        --index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \
+        --extra-index-url https://pypi.org/simple \
+        --index-strategy unsafe-first-match; \
     fi
 
 # Expose the Rust frontend bundled in the wheel on the container's PATH.
@@ -1140,7 +1130,9 @@ RUN --mount=type=bind,from=build,src=/workspace/dist,target=/vllm-workspace/dist
     else \
         echo "Installing vLLM..." \
         && uv pip install --system dist/*.whl --verbose \
-        --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
+        --index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \
+        --extra-index-url https://pypi.org/simple \
+        --index-strategy unsafe-first-match; \
     fi
 
 # Pytorch now installs NVSHMEM, setting LD_LIBRARY_PATH
@@ -1150,7 +1142,9 @@ COPY --from=build /tmp/ep_kernels_workspace/dist/wheels.sha256 /tmp/ep-kernels-w
 RUN --mount=type=bind,from=build,src=/tmp/ep_kernels_workspace/dist,target=/vllm-workspace/ep_kernels/dist \
     --mount=type=cache,target=/opt/uv/cache \
     uv pip install --system ep_kernels/dist/*.whl --verbose \
-        --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
+        --index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \
+        --extra-index-url https://pypi.org/simple \
+        --index-strategy unsafe-first-match
 
 # See the matching comment in the vllm-base stage.
 ENV LD_LIBRARY_PATH=/usr/local/nvidia/lib64:${LD_LIBRARY_PATH}
@@ -1293,25 +1287,6 @@ RUN --mount=type=cache,target=/opt/uv/cache \
         /tmp/install_triton_from_source.sh \
         /tmp/triton-source-wheel.sha256
 
-# Install source-built NIXL last, replacing the nixl packages installed above, so
-# the image has a NIXL EP build for its PyTorch version. There is no GPU driver
-# at build time, so check for the extension instead of importing it.
-ARG BUILD_NIXL
-RUN --mount=type=cache,target=/opt/uv/cache \
-    --mount=type=bind,from=extensions-build,source=/tmp/nixl_workspace/dist,target=/tmp/nixl-dist,ro \
-    set -eu; \
-    if [ "${BUILD_NIXL}" = "true" ]; then \
-        ls /tmp/nixl-dist/*.whl; \
-        uv pip uninstall --system nixl nixl-cu12 nixl-cu13; \
-        uv pip install --system --no-deps /tmp/nixl-dist/*.whl; \
-        python3 -c "import glob, importlib.util as u, os, torch; \
-            spec = u.find_spec('nixl_ep_cu' + os.environ['CUDA_VERSION'].split('.')[0]); \
-            mm = ''.join(torch.__version__.split('.')[:2]); \
-            ext = glob.glob(os.path.join(spec.submodule_search_locations[0], 'nixl_ep_cpp_torch' + mm + '.*')) if spec else []; \
-            assert ext, 'NIXL EP has no build for torch ' + torch.__version__; \
-            print('NIXL EP extension:', ext[0])"; \
-    fi
-
 ENV VLLM_USAGE_SOURCE production-docker-image
 ENV VLLM_BUILD_COMMIT=${VLLM_BUILD_COMMIT:-unknown} \
     VLLM_BUILD_PIPELINE=${VLLM_BUILD_PIPELINE:-local} \
diff --git a/docker/Dockerfile.cpu b/docker/Dockerfile.cpu
index db146475c830..b8fb7e07ccba 100644
--- a/docker/Dockerfile.cpu
+++ b/docker/Dockerfile.cpu
@@ -446,6 +446,10 @@ ADD ./tools/ ./tools/
 # ship either file, and this stage does not inherit the vllm-src tree.
 ADD ./setup.py ./setup.py
 ADD ./requirements/cpu.txt ./requirements/cpu.txt
+# tests/v1/core/test_prompt_block_hash_vectors.py checks the engine against the
+# prompt block hash vectors that the Rust frontend's unit tests also use.
+ADD ./rust/src/engine-core-client/src/tests/prompt_block_hash_vectors.json \
+    ./rust/src/engine-core-client/src/tests/prompt_block_hash_vectors.json
 
 # install development dependencies (for testing)
 RUN --mount=type=cache,target=/root/.cache/uv \
diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm
index d1a9c99803ac..30de51609dd0 100644
--- a/docker/Dockerfile.rocm
+++ b/docker/Dockerfile.rocm
@@ -1,409 +1,40 @@
-# default base image
-ARG REMOTE_VLLM="0"
-ARG VLLM_REPO="https://github.com/vllm-project/vllm.git"
-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=669b31acc1dd1b3fd93286afd5db67f65d9f7557
-ARG INSTALL_LMCACHE=false
-# NIC backend for MoRI RDMA support.
-# By default (all), drivers and userspace libraries for all supported NIC types
-# (ainic and bnxt) are installed; MoRI selects the appropriate one at runtime.
-# To install drivers for a single NIC type only, set NIC_BACKEND explicitly:
-#   --build-arg NIC_BACKEND=ainic   # AMD AINIC (Pensando) only
-#   --build-arg NIC_BACKEND=bnxt    # Broadcom Thor-2 only
-#   --build-arg NIC_BACKEND=none    # Install nothing.
-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 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
-# shared S3-compatible cache backend is available.
-ARG USE_SCCACHE
-ARG SCCACHE_DOWNLOAD_URL
-ARG SCCACHE_ENDPOINT
-ARG SCCACHE_BUCKET_NAME=vllm-build-sccache
-ARG SCCACHE_REGION_NAME=us-west-2
-ARG SCCACHE_S3_NO_CREDENTIALS=0
-
-FROM ${BASE_IMAGE} AS base
-
-ARG ARG_PYTORCH_ROCM_ARCH
-ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}}
-
-# Install build dependencies and utilities
-RUN apt-get update -q -y && apt-get install -q -y \
+# syntax=docker/dockerfile:1.27@sha256:4edf897a3ffa55b89f906fc8cc78afdb3f1834cc9c7083565e611a8a7d5fe99e
+# vLLM on ROCm. Build it with docker/docker-bake-rocm.hcl, which provides the
+# `stack` named context (the `final` stage of docker/Dockerfile.rocm_base)
+# and optionally prebuilt `runtime` and `test-deps` images.
+# Everything above the per-commit marker is hashed into the runtime key by .buildkite/scripts/rocm/build.sh.
+
+FROM stack AS dev
+# Build tools plus the runtime system packages (RDMA userspace, FFmpeg for torchcodec).
+# mold is installed but not the default linker: some packages JIT-compile at runtime
+# with flags mold does not support. Build stages opt in via LDFLAGS="-fuse-ld=mold".
+RUN apt-get update -q -y -o Acquire::Retries=3 && apt-get install -q -y -o Acquire::Retries=3 \
     sqlite3 libsqlite3-dev libfmt-dev libmsgpack-dev libsuitesparse-dev \
-    apt-transport-https ca-certificates wget curl \
-    libnuma-dev ccache mold
+    apt-transport-https ca-certificates wget curl libnuma-dev ccache mold \
+    librdmacm1 libibverbs1 ibverbs-providers ibverbs-utils util-linux unzip libsodium23 \
+    pkg-config ffmpeg libavcodec-dev libavformat-dev libavutil-dev \
+    libswscale-dev libavdevice-dev libavfilter-dev libswresample-dev \
+    && rm -rf /var/lib/apt/lists/*
 RUN --mount=type=cache,target=/root/.cache/pip \
     python3 -m pip install --upgrade pip
-# Note: mold is installed but not set as the system default linker because
-# some packages use JIT compilation at runtime with flags mold does not support.
-# Build stages opt in via LDFLAGS="-fuse-ld=mold".
-# Remove sccache only if not using sccache (it exists in base image from Dockerfile.rocm_base)
-ARG USE_SCCACHE
-RUN if [ "$USE_SCCACHE" != "1" ]; then \
-        apt-get purge -y sccache || true; \
-        python3 -m pip uninstall -y sccache || true; \
-        rm -f "$(which sccache)" || true; \
-    fi
-
-# Install UV — download first, then run, so a curl failure is not masked by the pipe
+# Download first, then run, so a curl failure is not masked by the pipe
 RUN curl -LsSf --retry 3 --retry-delay 5 https://astral.sh/uv/install.sh -o /tmp/uv-install.sh \
     && env UV_INSTALL_DIR="/usr/local/bin" sh /tmp/uv-install.sh \
     && rm -f /tmp/uv-install.sh \
     && uv --version
-
-# This timeout (in seconds) is necessary when installing some dependencies via uv since it's likely to time out
-# Reference: https://github.com/astral-sh/uv/pull/1694
+# Some dependencies time out with uv's default; see https://github.com/astral-sh/uv/pull/1694
 ENV UV_HTTP_TIMEOUT=500
 ENV UV_INDEX_STRATEGY="unsafe-best-match"
 # Use copy mode to avoid hardlink failures with Docker cache mounts
 ENV UV_LINK_MODE=copy
-# ccache directory - persisted across layer rebuilds via cache mounts.
 ENV CCACHE_DIR=/root/.cache/ccache
 ENV CCACHE_COMPILERCHECK=content
-# Empty by default so build steps fall back to $(nproc); CI can override.
-ARG max_jobs
-ENV MAX_JOBS=${max_jobs}
-
-# Install sccache if USE_SCCACHE is enabled (for release builds)
-ARG USE_SCCACHE
-ARG SCCACHE_DOWNLOAD_URL
-ARG SCCACHE_ENDPOINT
-ARG SCCACHE_BUCKET_NAME
-ARG SCCACHE_REGION_NAME
-ARG SCCACHE_S3_NO_CREDENTIALS
-RUN if [ "$USE_SCCACHE" = "1" ]; then \
-        if command -v sccache >/dev/null 2>&1; then \
-            echo "sccache already installed, skipping installation"; \
-            sccache --version; \
-        else \
-            echo "Installing sccache..." \
-            && SCCACHE_ARCH="x86_64" \
-            && SCCACHE_VERSION="v0.8.1" \
-            && SCCACHE_DL_URL="${SCCACHE_DOWNLOAD_URL:-https://github.com/mozilla/sccache/releases/download/${SCCACHE_VERSION}/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl.tar.gz}" \
-            && curl -L -o /tmp/sccache.tar.gz ${SCCACHE_DL_URL} \
-            && tar -xzf /tmp/sccache.tar.gz -C /tmp \
-            && mv /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl/sccache /usr/bin/sccache \
-            && chmod +x /usr/bin/sccache \
-            && rm -rf /tmp/sccache.tar.gz /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl \
-            && sccache --version; \
-        fi; \
-    fi
-
-# Set sccache environment variables only when USE_SCCACHE=1
-# This prevents S3 config from leaking into images when sccache is not used
-ARG USE_SCCACHE
-ENV SCCACHE_BUCKET=${USE_SCCACHE:+${SCCACHE_BUCKET_NAME}}
-ENV SCCACHE_REGION=${USE_SCCACHE:+${SCCACHE_REGION_NAME}}
-ENV SCCACHE_S3_NO_CREDENTIALS=${USE_SCCACHE:+${SCCACHE_S3_NO_CREDENTIALS}}
-ENV SCCACHE_IDLE_TIMEOUT=${USE_SCCACHE:+0}
-
-ARG COMMON_WORKDIR
-WORKDIR ${COMMON_WORKDIR}
-
-
-# -----------------------
-# vLLM fetch stages
-FROM base AS fetch_vllm_0
-ONBUILD COPY ./ vllm/
-FROM base AS fetch_vllm_1
-ARG VLLM_REPO
-ARG VLLM_BRANCH
-ONBUILD RUN git clone ${VLLM_REPO} \
-	    && cd vllm \
-	    && git fetch -v --prune -- origin ${VLLM_BRANCH} \
-	    && git checkout FETCH_HEAD \
-        && if [ ${VLLM_REPO} != "https://github.com/vllm-project/vllm.git" ] ; then \
-               git remote add upstream "https://github.com/vllm-project/vllm.git" \
-               && git fetch upstream ; fi
-FROM fetch_vllm_${REMOTE_VLLM} AS fetch_vllm
-
-# Resolve the device-independent source version once for every downstream
-# artifact. setuptools-scm reads Git metadata for local/remote sources and the
-# archive metadata generated for canonical CI contexts.
-FROM fetch_vllm AS vllm-version
-ARG COMMON_WORKDIR
 RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \
-    cd ${COMMON_WORKDIR}/vllm \
-    && uv pip install --system -r requirements/build/rust.txt \
-    && python3 -c \
-        "from setuptools_scm import get_version; print(get_version(git_describe_command=['git', 'describe', '--dirty', '--tags', '--long', '--abbrev=40', '--match', 'v[0-9]*']))" \
-        > /vllm-version.txt
-
-# -----------------------
-# Rust toolchain inputs
-#
-# Keep these separate from Rust source inputs so Rust source-only changes do not
-# invalidate the long-lived ci_base image that carries the toolchain.
-FROM scratch AS rust_toolchain_input_0
-COPY rust-toolchain.toml /rust-toolchain-input/vllm/rust-toolchain.toml
-
-FROM base AS rust_toolchain_input_1
-ARG VLLM_REPO
-ARG VLLM_BRANCH
-RUN git clone --no-checkout --filter=blob:none ${VLLM_REPO} /rust-toolchain-input/vllm \
-    && cd /rust-toolchain-input/vllm \
-    && git fetch -v --prune -- origin ${VLLM_BRANCH} \
-    && git sparse-checkout init --no-cone \
-    && git sparse-checkout set \
-        rust-toolchain.toml \
-    && git checkout FETCH_HEAD
-
-FROM rust_toolchain_input_${REMOTE_VLLM} AS rust-toolchain-input
-
-# -----------------------
-# Rust build inputs
-#
-# REMOTE_VLLM=0: use only the local files that can affect Rust artifacts.
-# REMOTE_VLLM=1: fetch only those same files from the requested remote ref.
-FROM scratch AS rust_input_0
-COPY requirements/build/rust.txt /rust-input/vllm/requirements/build/rust.txt
-COPY rust/Cargo.lock rust/Cargo.toml /rust-input/vllm/rust/
-COPY rust/proto /rust-input/vllm/rust/proto
-COPY rust/src /rust-input/vllm/rust/src
-COPY rust-toolchain.toml /rust-input/vllm/rust-toolchain.toml
-COPY tools/build_rust.py /rust-input/vllm/tools/build_rust.py
-COPY tools/build_rust.sh /rust-input/vllm/tools/build_rust.sh
-
-FROM base AS rust_input_1
-ARG VLLM_REPO
-ARG VLLM_BRANCH
-RUN git clone --no-checkout --filter=blob:none ${VLLM_REPO} /rust-input/vllm \
-    && cd /rust-input/vllm \
-    && git fetch -v --prune -- origin ${VLLM_BRANCH} \
-    && git sparse-checkout init --no-cone \
-    && git sparse-checkout set \
-        requirements/build/rust.txt \
-        rust/Cargo.lock \
-        rust/Cargo.toml \
-        rust/proto/** \
-        rust/src/** \
-        rust-toolchain.toml \
-        tools/build_rust.py \
-        tools/build_rust.sh \
-    && git checkout FETCH_HEAD
-
-FROM rust_input_${REMOTE_VLLM} AS rust-input
-
-# -----------------------
-# Rust toolchain
-FROM base AS rust-toolchain
-
-ENV CARGO_HOME=/root/.cargo
-ENV RUSTUP_HOME=/root/.rustup
-ENV PATH=${CARGO_HOME}/bin:${PATH}
-
-RUN apt-get update -q -y && apt-get install -q -y --no-install-recommends \
-        ca-certificates curl \
-    && rm -rf /var/lib/apt/lists/*
-
-COPY --from=rust-toolchain-input /rust-toolchain-input/vllm/rust-toolchain.toml /tmp/rust-toolchain.toml
-RUN TOOLCHAIN="$(grep '^channel' /tmp/rust-toolchain.toml | sed 's/.*= *"\(.*\)"/\1/')" \
-    && if ! command -v rustup >/dev/null 2>&1; then \
-        curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs \
-            | sh -s -- -y --default-toolchain none; \
-    fi \
-    && rustup toolchain install "${TOOLCHAIN}" \
-    && rustup default "${TOOLCHAIN}" \
-    && rustc --version \
-    && cargo --version
-
-# -----------------------
-# Rust build stage
-# Builds the `vllm-rs` frontend in a dedicated stage so the wheel build stages
-# don't need the rust toolchain.
-FROM rust-toolchain AS rust-build
-ARG COMMON_WORKDIR
-ARG USE_SCCACHE
-
-# Cap cargo parallelism to avoid exhausting the AMD CI host's open-file limit
-# (rustc spawns enough concurrent processes to hit RLIMIT_NOFILE otherwise).
-ENV CARGO_BUILD_JOBS=4
-ENV CARGO_NET_RETRY=10
-ENV RUSTUP_MAX_RETRIES=10
-
-WORKDIR ${COMMON_WORKDIR}/vllm
-
-COPY --from=rust-input /rust-input/vllm/requirements/build/rust.txt requirements/build/rust.txt
-
-# BuildKit can run this stage in parallel with ROCm native builds. Keep Rust on
-# a separate local sccache daemon while sharing the same remote cache backend.
-ENV SCCACHE_SERVER_PORT=4227
-
-RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \
-    uv pip install --system -r requirements/build/rust.txt
-
-COPY --from=rust-input /rust-input/vllm/rust rust
-COPY --from=rust-input /rust-input/vllm/rust-toolchain.toml rust-toolchain.toml
-COPY --from=rust-input /rust-input/vllm/tools/build_rust.py tools/build_rust.py
-COPY --from=rust-input /rust-input/vllm/tools/build_rust.sh tools/build_rust.sh
-COPY --from=vllm-version /vllm-version.txt /tmp/vllm-version.txt
-
-# Build the release binary. Cargo's registry/git caches can be written by
-# concurrent BuildKit jobs on shared workers, so lock those cache mounts while
-# keeping the cache benefit. Do not cache target/, because stale target metadata
-# can outlive source updates across BuildKit cache reuse.
-RUN --mount=type=cache,id=vllm-rocm-cargo-registry,target=/root/.cargo/registry,sharing=locked \
-    --mount=type=cache,id=vllm-rocm-cargo-git,target=/root/.cargo/git,sharing=locked \
-    if [ "$USE_SCCACHE" = "1" ]; then \
-        export RUSTC_WRAPPER=sccache \
-        && sccache --show-stats; \
-    fi \
-    && export SETUPTOOLS_SCM_PRETEND_VERSION="$(cat /tmp/vllm-version.txt)" \
-    && bash tools/build_rust.sh \
-    && test -x vllm/vllm-rs \
-    && if [ "$USE_SCCACHE" = "1" ]; then \
-        sccache --show-stats; \
-    fi
-
-# -----------------------
-# Shared vLLM build dependencies
-#
-# Install requirements once, before bringing source-specific inputs into either
-# consumer. Native compilation and wheel packaging then share this expensive
-# layer while retaining their independent source cache boundaries.
-FROM base AS build_vllm_dependencies
-ARG COMMON_WORKDIR
-ENV VLLM_TARGET_DEVICE=rocm
-WORKDIR ${COMMON_WORKDIR}/vllm
-
-COPY --from=fetch_vllm \
-    ${COMMON_WORKDIR}/vllm/requirements/common.txt \
-    ${COMMON_WORKDIR}/vllm/requirements/rocm.txt \
-    requirements/
-
-# Use constraints file to protect triton, pybind11, and torch dependencies.
-RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \
-    uv pip install --system -c /etc/rocm-constraints.txt -r requirements/rocm.txt
-
-# Fetch only the pinned ROCm Triton kernels package. Keeping this before the
-# vLLM source COPYs makes the download reusable when native sources change.
-FROM base AS rocm-triton-kernels
-ARG ROCM_TRITON_KERNELS_COMMIT
-ADD https://codeload.github.com/ROCm/triton/tar.gz/${ROCM_TRITON_KERNELS_COMMIT} \
-    /tmp/rocm-triton.tar.gz
-RUN mkdir -p /opt/rocm-triton-kernels \
-    && tar --extract --gzip --file /tmp/rocm-triton.tar.gz \
-        --directory /opt/rocm-triton-kernels --strip-components=4 \
-        --no-same-owner \
-        "triton-${ROCM_TRITON_KERNELS_COMMIT}/python/triton_kernels/triton_kernels" \
-    && test -f /opt/rocm-triton-kernels/__init__.py \
-    && rm /tmp/rocm-triton.tar.gz
-
-# -----------------------
-# vLLM native build stages
-#
-# csrc-build intentionally copies only files that affect ROCm native extension
-# compilation. That keeps unrelated CI/test/docs edits from invalidating the
-# expensive HIP/C++ build layer.
-FROM build_vllm_dependencies AS csrc-build
-ARG ROCM_TRITON_KERNELS_COMMIT
-COPY --from=rocm-triton-kernels \
-    /opt/rocm-triton-kernels /opt/rocm-triton-kernels
-ENV TRITON_KERNELS_SRC_DIR=/opt/rocm-triton-kernels
-
-# pyproject.toml is bind-mounted so it is not persisted in the csrc-build image,
-# but it still participates in the BuildKit cache key for this RUN step.
-COPY setup.py CMakeLists.txt ./
-COPY tools/build_rust.py tools/build_rust.py
-COPY cmake cmake/
-COPY csrc csrc/
-COPY vllm/envs.py vllm/envs.py
-COPY vllm/__init__.py vllm/__init__.py
-
-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}\")" \
-        cmake/external_projects/triton_kernels.cmake \
-        || echo "WARNING: Dockerfile.rocm 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) \
-    && (ccache --zero-stats || true) \
-    && EFFECTIVE_MAX_JOBS="${MAX_JOBS:-$(nproc)}" \
-    && echo "Building ROCm native extension wheel with MAX_JOBS=${EFFECTIVE_MAX_JOBS}" \
-    && PYTHONDONTWRITEBYTECODE=1 LDFLAGS="-fuse-ld=mold" MAX_JOBS="${EFFECTIVE_MAX_JOBS}" python3 setup.py bdist_wheel --dist-dir=dist \
-    && test -d dist \
-    && ls dist/*.whl >/dev/null \
-    && echo "=== ccache stats after ROCm native build ===" \
-    && (ccache --show-stats || true)
-
-FROM csrc-build AS kernel-symbol-map-build
-ARG VLLM_KERNEL_SYMBOL_MAP=0
-ARG buildkite_commit=""
-COPY tools/ci/kernel_symbol_map.py tools/ci/rocm_kernel_symbol_map.py tools/ci/
-RUN if [ "${VLLM_KERNEL_SYMBOL_MAP}" = "1" ]; then \
-        uv venv --system-site-packages /tmp/kernel-map/.venv \
-        && /tmp/kernel-map/.venv/bin/python tools/ci/kernel_symbol_map.py \
-            --backend rocm --build-root build --source-root . \
-            --commit "${buildkite_commit}" \
-            --out dist/kernel_symbol_map.rocm.json.gz \
-        && rm -rf /tmp/kernel-map; \
-    fi
-
-FROM scratch AS kernel-symbol-map-rocm
-ARG COMMON_WORKDIR
-COPY --from=kernel-symbol-map-build ${COMMON_WORKDIR}/vllm/dist/kernel_symbol_map.rocm.json.gz /
-
-# Build the full vLLM ROCm wheel by reusing the native extension wheel from
-# csrc-build. This stage still rebuilds for Python/package changes, but skips
-# the expensive dependency installation and HIP/C++ compilation when their
-# respective inputs are unchanged.
-FROM build_vllm_dependencies AS build_vllm
-COPY --from=fetch_vllm ${COMMON_WORKDIR}/vllm ${COMMON_WORKDIR}/vllm
-COPY --from=csrc-build ${COMMON_WORKDIR}/vllm/dist /precompiled-wheels
-COPY --from=vllm-version /vllm-version.txt /tmp/vllm-version.txt
-
-# Drop the pre-built Rust artifacts into the source tree. setup.py detects
-# them and ships them as-is, skipping the local Rust build.
-COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/vllm-rs ${COMMON_WORKDIR}/vllm/vllm/vllm-rs
-COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/_rust_*.so ${COMMON_WORKDIR}/vllm/vllm/
-
-RUN cd "${COMMON_WORKDIR}/vllm" \
-    && export SETUPTOOLS_SCM_PRETEND_VERSION="$(cat /tmp/vllm-version.txt)" \
-    && export VLLM_USE_PRECOMPILED=1 \
-    && export VLLM_PRECOMPILED_WHEEL_LOCATION="$(ls /precompiled-wheels/*.whl)" \
-    && export VLLM_DOCKER_BUILD_CONTEXT=1 \
-    && echo "Packaging vLLM ROCm wheel using precompiled extensions from ${VLLM_PRECOMPILED_WHEEL_LOCATION}" \
-    && PYTHONDONTWRITEBYTECODE=1 python3 setup.py bdist_wheel --dist-dir=dist \
-    && test -d dist \
-    && ls dist/*.whl >/dev/null
-FROM scratch AS export_vllm
-ARG COMMON_WORKDIR
-COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/dist/*.whl /
-COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/requirements /requirements
-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/tools/benchmark_helion_kernels.py /tools/benchmark_helion_kernels.py
-COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/tools/install_torchcodec_rocm.sh /tools/install_torchcodec_rocm.sh
-# tests/tools/test_kernel_symbol_map.py loads the symbol-map producers from here.
-COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/tools/ci /tools/ci
-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/
-COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm_base /docker/Dockerfile.rocm_base
-COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/ci-rocm.hcl /docker/ci-rocm.hcl
-COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/docker-bake.hcl /docker/docker-bake.hcl
-COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/docker-bake-rocm.hcl /docker/docker-bake-rocm.hcl
-COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/.buildkite /.buildkite
-COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/pyproject.toml /pyproject.toml
-COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/rust /rust
-COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/rust-toolchain.toml /rust-toolchain.toml
-COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1
+    uv pip install --system -c /etc/rocm-constraints.txt cmake ninja "setuptools<80" setuptools-scm \
+        setuptools-rust wheel jinja2 packaging regex
 
 # NIXL/UCX build stages
-FROM base AS build_nixl
+FROM dev AS build_nixl
 ARG NIXL_BRANCH="1683cf3b7f3d11674c03c5e861cea22339876c96"
 ARG NIXL_REPO="https://github.com/ai-dynamo/nixl.git"
 ARG UCX_BRANCH="2ff0b47ce7626584cb9b56c647d08f712c6c3edf"
@@ -428,11 +59,11 @@ RUN apt-get -y update && apt-get -y install autoconf libtool pkg-config \
     ibverbs-providers \
     && rm -rf /var/lib/apt/lists/*
 
-RUN --mount=type=cache,target=/root/.cache/uv \
+RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \
     uv pip install --system meson meson-python pyyaml types-PyYAML \
         auditwheel build patchelf pytest tomlkit "setuptools>=80.9.0"
 
-RUN --mount=type=cache,target=/root/.cache/ccache \
+RUN --mount=type=cache,id=vllm-rocm-ccache,target=/root/.cache/ccache \
     cd /usr/local/src && \
     git clone ${UCX_REPO} &&  \
     cd ucx  && \
@@ -457,7 +88,7 @@ RUN --mount=type=cache,target=/root/.cache/ccache \
 ENV PATH=/usr/local/ucx/bin:$PATH
 ENV LD_LIBRARY_PATH=${UCX_HOME}/lib:${LD_LIBRARY_PATH}
 
-RUN --mount=type=cache,target=/root/.cache/ccache \
+RUN --mount=type=cache,id=vllm-rocm-ccache,target=/root/.cache/ccache \
     git clone ${NIXL_REPO} /opt/nixl && \
     cd /opt/nixl && \
     git checkout ${NIXL_BRANCH} && \
@@ -506,12 +137,12 @@ RUN cd /opt/nixl && \
     cp /tmp/nixl_wheels/repaired/*.whl /app/install
 
 # Fetch and verify pinned LMCache source.
-FROM base AS lmcache_source
+FROM dev AS lmcache_source
 ARG LMCACHE_REPO="https://github.com/LMCache/LMCache.git"
 ARG LMCACHE_REF="140819c9d57a975dbc5678a6459a218e544cb58b"
 ARG LMCACHE_VERSION="0.5.3"
 
-# Remove NVIDIA-only dependencies and accept the release image's NumPy version.
+# Remove NVIDIA-only dependencies and accept the runtime's NumPy version.
 RUN printf '%s\n' "${LMCACHE_REF}" | grep -Eq '^[0-9a-f]{40}$' \
     && git clone --depth=1 --branch "v${LMCACHE_VERSION}" \
         "${LMCACHE_REPO}" /app/lmcache \
@@ -526,7 +157,7 @@ RUN printf '%s\n' "${LMCACHE_REF}" | grep -Eq '^[0-9a-f]{40}$' \
         /app/lmcache/requirements/common.txt \
     && rm -rf /app/lmcache/.git
 
-FROM base AS build_lmcache
+FROM dev AS build_lmcache
 COPY --from=lmcache_source \
     /app/lmcache/requirements/build.txt /tmp/lmcache-build-requirements.txt
 RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \
@@ -534,7 +165,6 @@ RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \
 
 ARG LMCACHE_VERSION="0.5.3"
 ARG LMCACHE_ROCM_ARCH="gfx942;gfx950"
-ARG USE_SCCACHE
 
 COPY --from=lmcache_source /app/lmcache /app/lmcache
 WORKDIR /app/lmcache
@@ -545,9 +175,6 @@ RUN --mount=type=cache,id=vllm-rocm-ccache,target=/root/.cache/ccache \
         | grep -Eq '^gfx[0-9a-z]+(;gfx[0-9a-z]+)*$'; \
     offload_flags="$(printf '%s' "${LMCACHE_ROCM_ARCH}" \
         | sed 's/;/ --offload-arch=/g; s/^/--offload-arch=/')"; \
-    if [ "${USE_SCCACHE:-0}" = "1" ]; then \
-        export HIP_CLANG_PATH=/opt/sccache-wrappers; \
-    fi; \
     mkdir -p /app/install; \
     SETUPTOOLS_SCM_PRETEND_VERSION="${LMCACHE_VERSION}" \
     PYTORCH_ROCM_ARCH="${LMCACHE_ROCM_ARCH}" \
@@ -570,10 +197,8 @@ RUN --mount=type=tmpfs,target=/tmp/lmcache-wheel \
         exit 1; \
     fi
 
-# DeepEP build stage. ROCShmem is provided by TheRock SDK
-# ROCSHMEM_DIR set to the SDK install path so DeepEP finds the headers, etc.
-FROM base AS build_deepep
-ENV ROCSHMEM_DIR=${ROCM_PATH}
+# DeepEP build stage. The stack sets ROCSHMEM_DIR via the TheRock SDK.
+FROM dev AS build_deepep
 ARG DEEPEP_BRANCH="0f63d3e9"
 ARG DEEPEP_REPO="https://github.com/ROCm/DeepEP.git"
 ARG DEEPEP_NIC="cx7"
@@ -583,20 +208,53 @@ ARG DEEPEP_NIC="cx7"
 # We pass --rocm-gfx942-fp8-fnuz-max 224 to align the DeepEP dispatch FP8
 # quantization bound with vLLM's own clamp
 # (vllm/model_executor/layers/quantization/utils/fp8_utils.py): 224 for E4M3FNUZ
-RUN --mount=type=cache,target=/root/.cache/ccache \
+RUN --mount=type=cache,id=vllm-rocm-ccache,target=/root/.cache/ccache \
     export PYTORCH_ROCM_ARCH="gfx942;gfx950" \
  && git clone ${DEEPEP_REPO} \
  && cd DeepEP \
  && git checkout ${DEEPEP_BRANCH} \
  && LDFLAGS="-fuse-ld=mold" MAX_JOBS="${MAX_JOBS:-$(nproc)}" python3 setup.py --variant rocm --rocm-explicit-ctx --nic ${DEEPEP_NIC} --rocm-gfx942-fp8-fnuz-max 224 bdist_wheel --dist-dir=/app/deep_install
 
-# MoRI runtime dependencies live in Dockerfile.rocm so NIC backend changes do
-# not force users to rebuild the long-lived Dockerfile.rocm_base image.
-FROM base AS mori_base
-ARG NIC_BACKEND
-ARG AINIC_VERSION
-ARG UBUNTU_CODENAME
-RUN /bin/bash -lc 'set -euo pipefail; \
+# mamba-ssm and causal-conv1d, the fast path transformers uses for the hybrid-model references in the
+# language model tests. CI only. Their setup.py targets HIP_ARCHITECTURES (default: the local GPU).
+FROM dev AS build_mamba
+ARG MAMBA_COMMIT="269cbe654d10883374736eb936e61d3ee5625695"
+ARG CAUSAL_CONV1D_COMMIT="da6dbaa9fd5a919967f14d3fd031da1288ad5025"
+ARG HIP_ARCHITECTURES="gfx942,gfx950"
+RUN MAMBA_FORCE_BUILD=TRUE CAUSAL_CONV1D_FORCE_BUILD=TRUE \
+    python3 -m pip wheel --no-build-isolation --no-deps -w /app/mamba_install \
+        "git+https://github.com/AndreasKaratzas/mamba@${MAMBA_COMMIT}" \
+        "git+https://github.com/Dao-AILab/causal-conv1d@${CAUSAL_CONV1D_COMMIT}" \
+    && for w in /app/mamba_install/*.whl; do python3 -m zipfile -e "$w" /tmp/whl; done \
+    && for so in /tmp/whl/*.so; do for a in $(echo "$HIP_ARCHITECTURES" | tr , ' '); do \
+         llvm-objdump --offloading "$so" | grep -qw "$a" || { echo "$so lacks $a" >&2; exit 1; }; done; done \
+    && rm -rf /tmp/whl
+
+# rdma-core recent enough for DeepEP's rocshmem. CI only: it replaces the distro libibverbs, which
+# the vendor NIC providers in the runtime (AINIC, Broadcom bnxt_re) are built against.
+FROM dev AS build_rdma_core
+RUN git clone --branch v62.0 --depth 1 https://github.com/linux-rdma/rdma-core.git /tmp/rdma-core \
+    && cd /tmp/rdma-core && mkdir build && cd build \
+    && cmake -GNinja -DCMAKE_INSTALL_PREFIX=/usr -DNO_MAN_PAGES=1 .. \
+    && ninja && DESTDIR=/out ninja install
+
+# The runtime: stack + every compiled component + all pip dependencies + all env. Published as base-nightly.
+FROM dev AS runtime
+# NIC backend for MoRI RDMA support.
+# By default (all), drivers and userspace libraries for all supported NIC types
+# (ainic and bnxt) are installed; MoRI selects the appropriate one at runtime.
+# To install drivers for a single NIC type only, set NIC_BACKEND explicitly:
+#   --build-arg NIC_BACKEND=ainic   # AMD AINIC (Pensando) only
+#   --build-arg NIC_BACKEND=bnxt    # Broadcom Thor-2 only
+#   --build-arg NIC_BACKEND=none    # Install nothing.
+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 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
+RUN --mount=type=cache,target=/root/.cache/pip /bin/bash -lc 'set -euo pipefail; \
  \
  install_ainic() { \
    apt-get update && apt-get install -y --no-install-recommends ca-certificates curl gnupg apt-transport-https; \
@@ -629,9 +287,7 @@ RUN /bin/bash -lc 'set -euo pipefail; \
    rm -rf /var/lib/apt/lists/*; \
  }; \
  \
- echo "[MORI] Install MoRI proxy deps"; \
- pip install --quiet --ignore-installed blinker && \
- pip install --quiet quart msgpack aiohttp pyzmq; \
+ pip install --quiet --ignore-installed blinker; \
  echo "[MORI] NIC_BACKEND=${NIC_BACKEND}"; \
  \
  # NIC backend deps — mori auto-detects NIC at runtime (MORI_DEVICE_NIC env var override). \
@@ -643,395 +299,297 @@ RUN /bin/bash -lc 'set -euo pipefail; \
    bnxt)  install_bnxt ;; \
    *)     echo "ERROR: unknown NIC_BACKEND=${NIC_BACKEND}. Use one of: none, ainic, bnxt, all"; exit 2 ;; \
  esac'
-
-# -----------------------
-# vLLM wheel release build stage (for building distributable wheels)
-# This stage pins dependencies to custom ROCm wheel versions and handles version detection
-FROM fetch_vllm AS build_vllm_wheel_release
-
-ARG COMMON_WORKDIR
-ARG ROCM_TRITON_KERNELS_COMMIT
-COPY --from=vllm-version /vllm-version.txt /tmp/vllm-version.txt
-COPY --from=rocm-triton-kernels \
-    /opt/rocm-triton-kernels /opt/rocm-triton-kernels
-ENV TRITON_KERNELS_SRC_DIR=/opt/rocm-triton-kernels
-
-# Drop the pre-built Rust artifacts into the source tree. setup.py detects
-# them and ships them as-is, skipping the local Rust build.
-COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/vllm-rs ${COMMON_WORKDIR}/vllm/vllm/vllm-rs
-COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/_rust_*.so ${COMMON_WORKDIR}/vllm/vllm/
-
-# Create /install directory for custom wheels
-RUN mkdir -p /install
-
-# Copy custom ROCm wheels from docker/context if they exist
-# COPY ensures Docker cache is invalidated when wheels change
-# .keep file ensures directory always exists for COPY to work
-COPY docker/context/base-wheels/ /tmp/base-wheels/
-# This is how we know if we are building for a wheel release or not.
-# If there are not wheels found there, we are not building for a wheel release. 
-# So we exit with an error. To skip this stage.
-RUN if [ -n "$(ls /tmp/base-wheels/*.whl 2>/dev/null)" ]; then \
-        echo "Found custom wheels - copying to /install"; \
-        cp /tmp/base-wheels/*.whl /install/ && \
-        echo "Copied custom wheels:"; \
-        ls -lh /install/; \
-    else \
-        echo "ERROR: No custom wheels found in docker/context/base-wheels/"; \
-        echo "Wheel releases require pre-built ROCm wheels."; \
-        exit 1; \
-    fi
-
-# GIT_REPO_CHECK: Verify repo is clean and tags are available (for release builds)
-# This matches CUDA's Dockerfile behavior for proper version detection via setuptools_scm
-ARG GIT_REPO_CHECK=0
-RUN if [ "$GIT_REPO_CHECK" != "0" ]; then \
-        echo "Running repository checks..."; \
-        cd vllm && bash tools/check_repo.sh; \
-    fi
-
-# 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}\")" \
-        vllm/cmake/external_projects/triton_kernels.cmake \
-        || echo "WARNING: Dockerfile.rocm 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
-
-# Fail if git-based package dependencies are found in requirements files
-# (uv doesn't handle git+ URLs well, and packages should be distributed on PyPI)
-# Extra notes: pip install is able to handle git+ URLs, but uv doesn't.
-RUN echo "Checking for git-based packages in requirements files..." \
-    && echo "Checking common.txt for git-based packages:" \
-    && if grep -q 'git+' ${COMMON_WORKDIR}/vllm/requirements/common.txt; then \
-         echo "ERROR: Git-based packages found in common.txt:"; \
-         grep 'git+' ${COMMON_WORKDIR}/vllm/requirements/common.txt; \
-         echo "Please publish these packages to PyPI instead of using git dependencies."; \
-         exit 1; \
-       else \
-         echo "  ✓ No git-based packages found in common.txt"; \
-       fi \
-    && echo "Checking rocm.txt for git-based packages:" \
-    && if grep -q 'git+' ${COMMON_WORKDIR}/vllm/requirements/rocm.txt; then \
-         echo "ERROR: Git-based packages found in rocm.txt:"; \
-         grep 'git+' ${COMMON_WORKDIR}/vllm/requirements/rocm.txt; \
-         echo "Please publish these packages to PyPI instead of using git dependencies."; \
-         exit 1; \
-       else \
-         echo "  ✓ No git-based packages found in rocm.txt"; \
-       fi \
-    && echo "All requirements files are clean - no git-based packages found"
-
-# Pin vLLM dependencies to exact versions of custom ROCm wheels
-# This ensures 'pip install vllm' automatically installs correct torch/triton/torchvision/amdsmi
-COPY tools/vllm-rocm/pin_rocm_dependencies.py /tmp/pin_rocm_dependencies.py
-RUN echo "Pinning vLLM dependencies to custom wheel versions..." \
-    && python3 /tmp/pin_rocm_dependencies.py /install ${COMMON_WORKDIR}/vllm/requirements/rocm.txt
-
-# Install dependencies using custom wheels from /install
-RUN --mount=type=cache,target=/root/.cache/uv \
-    cd vllm \
-    && echo "Building vLLM with custom wheels from /install" \
-    && uv pip install --system --find-links /install -r requirements/rocm.txt
-
-# Build wheel using the shared source version to avoid dirty state from modified requirements/rocm.txt
-# (setup.py auto-detects ccache/sccache in PATH)
-RUN --mount=type=cache,id=vllm-rocm-ccache,target=/root/.cache/ccache \
-    cd vllm \
-    && export CCACHE_BASEDIR="$PWD" \
-    && export SETUPTOOLS_SCM_PRETEND_VERSION="$(cat /tmp/vllm-version.txt)" \
-    && echo "Building wheel with version: ${SETUPTOOLS_SCM_PRETEND_VERSION}" \
-    && MAX_JOBS="${MAX_JOBS:-$(nproc)}" python3 setup.py bdist_wheel --dist-dir=dist
-
-FROM scratch AS export_vllm_wheel_release
-ARG COMMON_WORKDIR
-COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/dist/*.whl /
-COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/requirements /requirements
-COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/benchmarks /benchmarks
-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/tools/benchmark_helion_kernels.py /tools/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/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/
-COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm_base /docker/Dockerfile.rocm_base
-COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/ci-rocm.hcl /docker/ci-rocm.hcl
-COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/docker-bake.hcl /docker/docker-bake.hcl
-COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/docker-bake-rocm.hcl /docker/docker-bake-rocm.hcl
-COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/.buildkite /.buildkite
-COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/pyproject.toml /pyproject.toml
-COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/rust /rust
-COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/rust-toolchain.toml /rust-toolchain.toml
-COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1
-
-# -----------------------
-# CI base image (Tier 1) - stable, rarely changing CI dependencies.
-# Per-PR test builds pull this as CI_BASE_IMAGE so the test stage only layers
-# in the vLLM artifacts for the current commit.
-FROM mori_base AS ci_base
-ARG COMMON_WORKDIR
-
-ENV CARGO_HOME=/root/.cargo
-ENV RUSTUP_HOME=/root/.rustup
-ENV PATH=${CARGO_HOME}/bin:${PATH}
-
-# Bake the Rust toolchain into ci_base so Rust build/test steps do not
-# need to fetch it on each per-commit image build.
-COPY --from=rust-toolchain /root/.cargo /root/.cargo
-COPY --from=rust-toolchain /root/.rustup /root/.rustup
-
-# Update rdma-core to support latest rocshmem.
-ARG DEEPEP_NIC
-RUN if [ "${DEEPEP_NIC}" = "cx7" ] || [ "${DEEPEP_NIC}" = "io" ]; then \
-    git clone --branch v62.0 --depth 1 https://github.com/linux-rdma/rdma-core.git /tmp/rdma-core && \
-    cd /tmp/rdma-core && \
-    mkdir -p build && cd build && \
-    cmake -GNinja -DCMAKE_INSTALL_PREFIX=/usr -DNO_MAN_PAGES=1 .. && \
-    ninja && ninja install && ldconfig && rm -rf /tmp/rdma-core; \
-fi
-
-# Install NIXL + DeepEP wheels.
-RUN --mount=type=bind,from=build_nixl,src=/app/install,target=/nixl_install \
-    --mount=type=bind,from=build_deepep,src=/app/deep_install,target=/deep_install \
-    uv pip install --system /nixl_install/*.whl /deep_install/*.whl
-
-# RDMA userspace libraries plus FFmpeg dev libs needed by torchcodec.
-RUN timeout --kill-after=30s 5m apt-get update -q -y --error-on=any \
-    -o Acquire::Retries=3 \
-    -o Acquire::http::Timeout=30 \
-    -o Acquire::https::Timeout=30 \
-    && apt-get install -q -y --no-install-recommends \
-    -o Acquire::Retries=3 \
-    -o Acquire::http::Timeout=30 \
-    -o Acquire::https::Timeout=30 \
-    librdmacm1 \
-    libibverbs1 \
-    ibverbs-providers \
-    ibverbs-utils \
-    util-linux \
-    unzip \
-    libsodium23 \
-    pkg-config ffmpeg libavcodec-dev libavformat-dev libavutil-dev \
-    libswscale-dev libavdevice-dev libavfilter-dev libswresample-dev \
-    && rm -rf /var/lib/apt/lists/*
+RUN echo "MORI_NIC_BACKEND=${NIC_BACKEND}" >> /app/versions.txt \
+    && echo "AINIC_VERSION=${AINIC_VERSION}" >> /app/versions.txt
 
 # Install torchcodec from source for ROCm/torch ABI compatibility.
-COPY tools/install_torchcodec_rocm.sh /tmp/install_torchcodec.sh
-RUN --mount=type=cache,target=/root/.cache/uv \
+RUN --mount=type=bind,source=tools/install_torchcodec_rocm.sh,target=/tmp/install_torchcodec.sh \
+    --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \
     --mount=type=cache,target=/root/.cache/pip \
-    --mount=type=cache,target=/root/.cache/torchcodec-wheels \
-    bash /tmp/install_torchcodec.sh \
-    && rm /tmp/install_torchcodec.sh \
-    && apt-get clean && rm -rf /var/lib/apt/lists/*
-
-# Pre-install shared ROCm runtime dependencies.
-COPY requirements/common.txt requirements/rocm.txt /tmp/ci-base-requirements/
-RUN --mount=type=cache,target=/root/.cache/uv \
-    uv pip install --system -c /etc/rocm-constraints.txt \
-        -r /tmp/ci-base-requirements/rocm.txt \
-    && rm -rf /tmp/ci-base-requirements
-
-# Enable fast and less brittle model downloads in tests.
-ENV HF_XET_HIGH_PERFORMANCE=1
-ENV HF_HUB_DOWNLOAD_TIMEOUT=60
-
-# Keep torch.cuda.is_available() fork-safe (see vllm/env_override.py).
-ENV PYTORCH_NVML_BASED_CUDA_CHECK=1
-
-# Pre-install vLLM test dependencies.
-COPY requirements/test/rocm.txt /tmp/rocm-test-reqs.txt
-RUN --mount=type=cache,target=/root/.cache/uv \
-    uv pip install --system -r /tmp/rocm-test-reqs.txt
-
+    --mount=type=tmpfs,target=/root/.cache/torchcodec-wheels \
+    ldconfig && bash /tmp/install_torchcodec.sh
+# All Python dependencies, at the versions CI tests (the test lock is a constraint).
+# quart, msgpack, aiohttp, pyzmq: MoRI proxy dependencies (quart: not in the test lock).
+RUN --mount=type=bind,source=requirements/common.txt,target=/r/common.txt \
+    --mount=type=bind,source=requirements/rocm.txt,target=/r/rocm.txt \
+    --mount=type=bind,source=requirements/test/rocm.txt,target=/r/test.txt \
+    --mount=type=bind,from=build_nixl,src=/app/install,target=/nixl \
+    --mount=type=bind,from=build_lmcache,src=/app/install,target=/lmcache \
+    --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \
+    uv pip install --system -c /etc/rocm-constraints.txt -c /r/test.txt -r /r/rocm.txt \
+        /nixl/*.whl /lmcache/*.whl quart==0.22.0 msgpack aiohttp pyzmq \
+    && python3 -c "import lmcache, lmcache.c_ops; print('lmcache', lmcache.__version__)"
 # Rebuild fastsafetensors from source so its C++ extension is compiled with
 # USE_ROCM and can detect libamdhip64.so at runtime.
-RUN --mount=type=cache,target=/root/.cache/pip \
-    FASTSAFETENSORS_REQ="$(grep -E '^fastsafetensors(==| @ )' /tmp/rocm-test-reqs.txt | head -1)" \
+RUN --mount=type=bind,source=requirements/test/rocm.txt,target=/r/test.txt \
+    --mount=type=cache,target=/root/.cache/pip \
+    FASTSAFETENSORS_REQ="$(grep -E '^fastsafetensors(==| @ )' /r/test.txt | head -1)" \
     && test -n "${FASTSAFETENSORS_REQ}" \
     && python3 -m pip install --force-reinstall --no-deps \
-        --no-binary fastsafetensors "${FASTSAFETENSORS_REQ}" \
-    && rm /tmp/rocm-test-reqs.txt
+        --no-binary fastsafetensors "${FASTSAFETENSORS_REQ}"
 
+# ROCm profiler limits workaround.
+RUN echo "ROCTRACER_MAX_EVENTS=10000000" > /app/libkineto.conf
+ENV KINETO_CONFIG="/app/libkineto.conf"
 # Set MIOPEN ENVS to resolve performance regressions in MIOpen 3D convolution kernel.
 # See: https://github.com/pytorch/pytorch/issues/169857
 ENV MIOPEN_DEBUG_CONV_DIRECT=0
 ENV MIOPEN_DEBUG_CONV_GEMM=0
-
 # Use legacy IPC mode for HSA to avoid GPU memory pinning issues with UCX rocm_ipc.
 # See: https://github.com/ROCm/rocm-libraries/issues/6266
 ENV HSA_ENABLE_IPC_MODE_LEGACY=1
 ENV UCX_ROCM_COPY_SIGPOOL_MAX_ELEMS=inf
 ENV UCX_ROCM_IPC_SIGPOOL_MAX_ELEMS=inf
-
-# ROCm profiler limits workaround.
-RUN echo "ROCTRACER_MAX_EVENTS=10000000" > ${COMMON_WORKDIR}/libkineto.conf
-ENV KINETO_CONFIG="${COMMON_WORKDIR}/libkineto.conf"
-
-# Install vllm_test_utils in ci_base for ci_base + wheel parity.
-COPY tests/vllm_test_utils /tmp/vllm_test_utils
-RUN --mount=type=cache,target=/root/.cache/uv \
-    uv pip install --system /tmp/vllm_test_utils \
-    && rm -rf /tmp/vllm_test_utils
-
-RUN --mount=type=bind,from=build_lmcache,src=/app/install,target=/lmcache_install \
-    --mount=type=bind,source=requirements/test/rocm.txt,target=/tmp/lmcache-constraints.txt \
-    --mount=type=cache,target=/root/.cache/uv \
-    uv pip install --system --constraint /tmp/lmcache-constraints.txt \
-        /lmcache_install/*.whl \
-    && python3 -c "import lmcache, lmcache.c_ops; print('lmcache', lmcache.__version__)"
-
-# -----------------------
-# Test vLLM image (Tier 2) - vLLM-only layer on top of ci_base.
-FROM ${CI_BASE_IMAGE} AS test
-ARG COMMON_WORKDIR
-
-# Install the vLLM wheel (--no-deps: all deps already in ci_base).
-RUN --mount=type=bind,from=export_vllm,src=/,target=/install \
-    --mount=type=cache,target=/root/.cache/uv \
-    cd /install \
-    && uv pip install --system --no-deps *.whl
-
-# Store the vLLM wheel in the image for python-only install tests.
-COPY --from=export_vllm /*.whl /opt/vllm-wheels/
-
-WORKDIR /vllm-workspace
-COPY --from=build_vllm ${COMMON_WORKDIR}/vllm /vllm-workspace
-
-# Copy in the v1 package (for python-only install test group).
-COPY --from=export_vllm /vllm_v1 /usr/local/lib/python${PYTHON_VERSION}/dist-packages/vllm/v1
-
-# Hide source under src/ so it won't shadow the installed package in tests.
-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
-ENV UV_COMPILE_BYTECODE=1
-
-RUN python3 -m pip install --upgrade pip && rm -rf /var/lib/apt/lists/*
-
-# Clean up sccache from release image (not needed at runtime)
-# This removes the binary and wrappers that may have been installed during build
-RUN rm -f /usr/bin/sccache || true \
-    && rm -rf /opt/sccache-wrappers || true
-
-# Unset sccache environment variables for the release image
-# This prevents S3 bucket config from leaking into production images
-ENV SCCACHE_BUCKET=
-ENV SCCACHE_REGION=
-ENV SCCACHE_ENDPOINT=
-ENV SCCACHE_S3_NO_CREDENTIALS=
-ENV SCCACHE_IDLE_TIMEOUT=
-
-# Error related to odd state for numpy 1.20.3 where there is no METADATA etc, but an extra LICENSES_bundled.txt.
-# Manually remove it so that later steps of numpy upgrade can continue
-RUN case "$(which python3)" in \
-        *"/opt/conda/envs/py_3.9"*) \
-            rm -rf /opt/conda/envs/py_3.9/lib/python3.9/site-packages/numpy-1.20.3.dist-info/;; \
-        *) ;; esac
-
-RUN --mount=type=cache,target=/root/.cache/uv \
-    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
-# Keep OpenTelemetry within LMCache's cap.
-RUN --mount=type=bind,from=export_vllm,src=/,target=/install \
-    --mount=type=cache,target=/root/.cache/uv \
-    cd /install \
-    && echo "opentelemetry-api<=1.40.0" > /tmp/otel-constraints.txt \
-    && uv pip install --system -c /etc/rocm-constraints.txt -c /tmp/otel-constraints.txt -r requirements/rocm.txt \
-    && pip uninstall -y vllm \
-    && uv pip install --system -c /etc/rocm-constraints.txt -c /tmp/otel-constraints.txt *.whl \
-    && rm /tmp/otel-constraints.txt
-
-# Install NIXL ROCm wheel
-RUN --mount=type=bind,from=build_nixl,src=/app/install,target=/nixl_install \
-    uv pip install --system /nixl_install/*.whl
-
-# Recompute the base's triton fingerprint to detect resolver interference.
-RUN EXPECTED=$(grep '^TRITON_SHA256: ' /app/versions.txt | cut -d' ' -f2) \
-    && ACTUAL=$(find $(python3 -c 'import triton, os; print(os.path.dirname(triton.__file__))') \
-        -type f \( -name '*.py' -o -name '*.so' \) | sort | xargs sha256sum | sha256sum | cut -d' ' -f1) \
-    && echo "triton sha256 expected=${EXPECTED} actual=${ACTUAL}" \
-    && test -n "${EXPECTED}" \
-    && test "${EXPECTED}" = "${ACTUAL}" \
-    && python3 -c "import torch; assert torch.version.hip, f'non-ROCm torch: {torch.__version__}'; print(f'torch={torch.__version__}')"
-
-# Recompute the base's pybind11 fingerprint to detect resolver interference.
-RUN EXPECTED=$(grep '^PYBIND11_SHA256: ' /app/versions.txt | cut -d' ' -f2) \
-    && EXPECTED_ABI=$(grep '^PYBIND11_INTERNALS_VERSION: ' /app/versions.txt | cut -d' ' -f2) \
-    && INC=$(python3 -c 'import pybind11; print(pybind11.get_include())') \
-    && ACTUAL=$(find "${INC}" -type f -name '*.h' | sort | xargs sha256sum | sha256sum | cut -d' ' -f1) \
-    && ACTUAL_ABI=$(grep -rhE '^#[[:space:]]*define PYBIND11_INTERNALS_VERSION' "${INC}" | awk '{print $NF}') \
-    && echo "pybind11 $(pip show pybind11 | awk '/^Version:/ {print $2}')" \
-        "sha256 expected=${EXPECTED} actual=${ACTUAL}" \
-        "internals expected=${EXPECTED_ABI} actual=${ACTUAL_ABI}" \
-    && test -n "${EXPECTED}" \
-    && test "${EXPECTED}" = "${ACTUAL}" \
-    && test "${EXPECTED_ABI}" = "${ACTUAL_ABI}"
-
-ARG COMMON_WORKDIR
-ARG BASE_IMAGE
-ARG NIC_BACKEND
-ARG AINIC_VERSION
-
-# Copy over the benchmark scripts as well
-COPY --from=export_vllm /benchmarks ${COMMON_WORKDIR}/vllm/benchmarks
-COPY --from=export_vllm /examples ${COMMON_WORKDIR}/vllm/examples
-COPY --from=export_vllm /docker ${COMMON_WORKDIR}/vllm/docker
-
-# Use legacy IPC mode for HSA to avoid GPU memory pinning issues with UCX rocm_ipc
-# See: https://github.com/ROCm/rocm-libraries/issues/6266
-ENV HSA_ENABLE_IPC_MODE_LEGACY=1
-ENV UCX_ROCM_COPY_SIGPOOL_MAX_ELEMS=inf
-ENV UCX_ROCM_IPC_SIGPOOL_MAX_ELEMS=inf
-
 ENV TOKENIZERS_PARALLELISM=false
-
 # ENV that can improve safe tensor loading, and end-to-end time
 ENV SAFETENSORS_FAST_GPU=1
-
 # Performance environment variable.
 ENV HIP_FORCE_DEV_KERNARG=1
-
 # Keep torch.cuda.is_available() fork-safe (see vllm/env_override.py).
 ENV PYTORCH_NVML_BASED_CUDA_CHECK=1
+ARG STACK
+LABEL ai.vllm.rocm.stack=${STACK}
 
-# Workaround for ROCm profiler limits
-RUN echo "ROCTRACER_MAX_EVENTS=10000000" > ${COMMON_WORKDIR}/libkineto.conf
-ENV KINETO_CONFIG="${COMMON_WORKDIR}/libkineto.conf"
-RUN echo "VLLM_BASE_IMAGE=${BASE_IMAGE}" >> ${COMMON_WORKDIR}/versions.txt \
-    && echo "MORI_NIC_BACKEND=${NIC_BACKEND}" >> ${COMMON_WORKDIR}/versions.txt \
-    && echo "AINIC_VERSION=${AINIC_VERSION}" >> ${COMMON_WORKDIR}/versions.txt
+FROM dev AS rust-toolchain
+ENV CARGO_HOME=/root/.cargo
+ENV RUSTUP_HOME=/root/.rustup
+ENV PATH=${CARGO_HOME}/bin:${PATH}
+RUN --mount=type=bind,source=rust-toolchain.toml,target=/tmp/rust-toolchain.toml \
+    TOOLCHAIN="$(grep '^channel' /tmp/rust-toolchain.toml | sed 's/.*= *"\(.*\)"/\1/')" \
+    && curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --default-toolchain none \
+    && rustup toolchain install "${TOOLCHAIN}" \
+    && rustup default "${TOOLCHAIN}" \
+    && rustc --version \
+    && cargo --version
 
-CMD ["/bin/bash"]
+# Test-only dependencies, a content-keyed image like the runtime (build.sh: test-deps-<stack>-<Kr>),
+# so every build of a key reuses these layers by digest instead of rebuilding them from cache.
+FROM runtime AS test-deps
+COPY --from=rust-toolchain /root/.cargo /root/.cargo
+COPY --from=rust-toolchain /root/.rustup /root/.rustup
+ENV CARGO_HOME=/root/.cargo RUSTUP_HOME=/root/.rustup PATH=/root/.cargo/bin:$PATH
+# DeepEP and the rdma-core its rocshmem needs: CI only, as before (see build_rdma_core).
+COPY --from=build_rdma_core /out /
+RUN --mount=type=bind,from=build_deepep,src=/app/deep_install,target=/deepep \
+    --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \
+    ldconfig && uv pip install --system -c /etc/rocm-constraints.txt /deepep/*.whl
+RUN --mount=type=bind,source=requirements/test/rocm.txt,target=/t.txt --mount=type=bind,source=tests/vllm_test_utils,target=/tu,rw \
+    --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \
+    sp="$(python3 -c 'import sysconfig; print(sysconfig.get_paths()["purelib"])')" \
+    && pip freeze --path "$sp" | sort >/tmp/shipped && uv pip install --system -c /etc/rocm-constraints.txt -r /t.txt /tu \
+    && ! pip freeze --path "$sp" | sort | comm -23 /tmp/shipped - | grep . # a test dep changed a shipped package (apt's python3-* may be shadowed)
+RUN --mount=type=bind,from=build_mamba,src=/app/mamba_install,target=/mamba \
+    --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \
+    uv pip install --system --no-deps /mamba/*.whl
+ENV HF_XET_HIGH_PERFORMANCE=1
 
-FROM final_common AS final_lmcache_false
+# ==== per-commit stages: nothing below is part of the runtime key ====
 
-FROM final_common AS final_lmcache_true
-RUN --mount=type=bind,from=build_lmcache,src=/app/install,target=/lmcache_install \
-    --mount=type=cache,target=/root/.cache/uv \
-    uv pip install --system /lmcache_install/*.whl \
-    && test "$(python3 -c 'import numpy; print(numpy.__version__)')" = "2.3.5" \
-    && uv pip check --system \
-    && python3 -c "import lmcache, lmcache.c_ops; print('lmcache', lmcache.__version__)"
+# The pin lives below the marker: a bump rebuilds csrc, not the runtime.
+FROM dev AS triton-kernels
+ARG ROCM_TRITON_KERNELS_COMMIT=669b31acc1dd1b3fd93286afd5db67f65d9f7557
+ADD https://codeload.github.com/ROCm/triton/tar.gz/${ROCM_TRITON_KERNELS_COMMIT} \
+    /tmp/rocm-triton.tar.gz
+RUN mkdir -p /opt/rocm-triton-kernels \
+    && tar --extract --gzip --file /tmp/rocm-triton.tar.gz \
+        --directory /opt/rocm-triton-kernels --strip-components=4 \
+        --no-same-owner \
+        "triton-${ROCM_TRITON_KERNELS_COMMIT}/python/triton_kernels/triton_kernels" \
+    && test -f /opt/rocm-triton-kernels/__init__.py \
+    && rm /tmp/rocm-triton.tar.gz
+ENV TRITON_KERNELS_SRC_DIR=/opt/rocm-triton-kernels ROCM_TRITON_KERNELS_COMMIT=${ROCM_TRITON_KERNELS_COMMIT}
+
+# The version every per-commit artifact carries. build.sh builds from an export of HEAD (no .git)
+# and writes .git_archival.txt from the host's `git describe`; without it, e.g. locally, .git and a
+# v* tag are used.
+FROM dev AS version-calc
+RUN --mount=type=bind,target=/src,rw cd /src \
+    && if [ -f .git_archival.txt ]; then rm -rf .git; \
+       else git config --global --add safe.directory '*'; fi \
+    && python3 -c "from setuptools_scm import get_version; print(get_version(git_describe_command=['git', 'describe', '--dirty', '--tags', '--long', '--abbrev=40', '--match', 'v[0-9]*']))" \
+        > /vllm-version.txt \
+    && cat /vllm-version.txt \
+    && if grep -q '^0\.1\.dev' /vllm-version.txt; then echo "no v* tag reachable from HEAD: git fetch --tags" >&2; exit 1; fi
+FROM scratch AS version
+COPY --from=version-calc /vllm-version.txt /
+
+# Compiles the `vllm-rs` frontend once per change to the rust sources, so the wheel build does not
+# need the rust toolchain. The crates, their downloads and target/ stay in this stage's layers (no
+# cache mounts), so rust-build always continues from the target/ of these exact sources.
+FROM rust-toolchain AS rust-deps
+# Cap cargo parallelism to avoid exhausting the AMD CI host's open-file limit
+# (rustc spawns enough concurrent processes to hit RLIMIT_NOFILE otherwise).
+ENV CARGO_BUILD_JOBS=4
+ENV CARGO_NET_RETRY=10
+ENV RUSTUP_MAX_RETRIES=10
+WORKDIR /app/vllm
+COPY requirements/build/rust.txt requirements/build/rust.txt
+RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \
+    uv pip install --system -r requirements/build/rust.txt
+COPY rust/Cargo.lock rust/Cargo.toml rust/
+COPY rust/proto rust/proto
+COPY rust/src rust/src
+COPY rust-toolchain.toml rust-toolchain.toml
+COPY tools/build_rust.py tools/build_rust.sh tools/
+RUN SETUPTOOLS_SCM_PRETEND_VERSION=0.0.0 bash tools/build_rust.sh
+
+# Per commit: vllm-rs embeds the vLLM version, so cargo recompiles only the crates that read it and
+# relinks (~30 s). Those few rustc processes stay far from the open-file limit, and the job count
+# also caps the thin-LTO threads, so lift the cap.
+FROM rust-deps AS rust-build
+COPY --from=version /vllm-version.txt /tmp/vllm-version.txt
+RUN CARGO_BUILD_JOBS="$(nproc)" SETUPTOOLS_SCM_PRETEND_VERSION="$(cat /tmp/vllm-version.txt)" \
+    bash tools/build_rust.sh \
+    && test -x vllm/vllm-rs
+
+# Native extensions only: copies just the files that affect the HIP/C++ build (and the map tool).
+FROM triton-kernels AS csrc
+WORKDIR /app/vllm
+COPY setup.py CMakeLists.txt pyproject.toml ./
+COPY tools/build_rust.py tools/build_rust.py
+COPY cmake cmake/
+COPY csrc csrc/
+COPY vllm/envs.py vllm/__init__.py vllm/
+# Empty: the stack's full PYTORCH_ROCM_ARCH (release); CI passes the arches it can test.
+ARG ARCH
+# setup.py reads requirements/rocm.txt for install_requires. This wheel only carries the .so files
+# the wheel stage repackages with the real requirements, so empty files keep a requirements bump
+# from recompiling HIP.
+# CI builds (ARCH set) also write the kernel symbol map (CI test selection: which source file each
+# GPU kernel came from). It reads the build tree, so it runs here, before the tree is removed, and
+# an edit to its tool rebuilds this stage. It never fails the build.
+# Release (SCCACHE_BUCKET_NAME, us-west-2): its agents run one job each, so the ccache mount starts
+# empty; sccache shares per-file results in S3 through the agent's instance role. setup.py prefers
+# sccache when it is on PATH; without S3 access it is removed and ccache compiles as in CI.
+ARG SCCACHE_BUCKET_NAME
+RUN --mount=type=cache,id=vllm-rocm-ccache,target=/root/.cache/ccache \
+    --mount=type=bind,source=tools/ci/rocm_kernel_symbol_map.py,target=/tmp/rocm_kernel_symbol_map.py \
+    mkdir -p requirements && : > requirements/rocm.txt \
+    && { grep -Fq "set(TRITON_KERNELS_TAG \"${ROCM_TRITON_KERNELS_COMMIT}\")" \
+        cmake/external_projects/triton_kernels.cmake \
+        || echo "WARNING: Dockerfile.rocm 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"; } \
+    && { [ -z "${SCCACHE_BUCKET_NAME}" ] \
+         || { export SCCACHE_BUCKET="${SCCACHE_BUCKET_NAME}" SCCACHE_REGION=us-west-2 SCCACHE_IDLE_TIMEOUT=0 \
+              && curl -fsSL --retry 3 -o /tmp/sccache.tgz https://github.com/mozilla/sccache/releases/download/v0.8.1/sccache-v0.8.1-x86_64-unknown-linux-musl.tar.gz \
+              && echo "e0ee621fb16b6940666cd770b091c62cadafd3e062dd12e3a49d9caaff3b795f  /tmp/sccache.tgz" | sha256sum -c - \
+              && tar -xzf /tmp/sccache.tgz -C /usr/local/bin --strip-components=1 --wildcards '*/sccache' \
+              && rm /tmp/sccache.tgz && sccache --start-server; } \
+         || { echo "WARNING: no S3 sccache; compiling with ccache only"; rm -f /usr/local/bin/sccache /tmp/sccache.tgz; }; } \
+    && CCACHE_BASEDIR="$PWD" PYTORCH_ROCM_ARCH="${ARCH:-$PYTORCH_ROCM_ARCH}" \
+       SETUPTOOLS_SCM_PRETEND_VERSION=0.0.0+rocm.csrc.build PYTHONDONTWRITEBYTECODE=1 \
+       LDFLAGS="-fuse-ld=mold" MAX_JOBS="${MAX_JOBS:-$(nproc)}" python3 setup.py bdist_wheel --dist-dir=/dist \
+    && { sccache --stop-server 2>/dev/null || true; } \
+    && mkdir /kernel-map \
+    && { [ -z "${ARCH}" ] \
+         || timeout 30m python3 /tmp/rocm_kernel_symbol_map.py --build-root build --source-root . \
+              --out /kernel-map/kernel_symbol_map.rocm.json.gz \
+         || echo "WARNING: no kernel symbol map; CI test selection falls back to its static rule"; } \
+    && rm -rf build
+
+# The map, stamped with this build's commit (CI selection only uses a map of the commit it ran). A
+# csrc layer reused from another commit was built from the same sources, so its map holds here.
+FROM dev AS kernel-symbol-map-stamp
+ARG BUILDKITE_COMMIT
+RUN --mount=type=bind,from=csrc,src=/kernel-map,target=/in python3 - "${BUILDKITE_COMMIT}" <<'EOF'
+import gzip, json, os, sys
+os.makedirs("/out")
+if os.path.exists(src := "/in/kernel_symbol_map.rocm.json.gz"):
+    with gzip.open(src, "rt") as f:
+        m = json.load(f)
+    m["commit"] = sys.argv[1]
+    with gzip.open("/out/kernel_symbol_map.rocm.json.gz", "wt") as f:
+        json.dump(m, f, separators=(",", ":"))
+EOF
+FROM scratch AS kernel-symbol-map
+COPY --from=kernel-symbol-map-stamp /out/ /
+
+# The full vLLM wheel, repackaged around the csrc extensions and the rust artifacts.
+FROM dev AS wheel
+COPY --from=csrc /dist /precompiled
+COPY --from=rust-build /app/vllm/vllm/vllm-rs /app/vllm/vllm/_rust_*.so /rust/
+COPY --from=version /vllm-version.txt /
+# The published wheel pins the exact stack builds (torch, triton, ...) it was compiled against.
+RUN --mount=type=bind,target=/src,rw cd /src && cp /rust/* vllm/ \
+    && ! grep -n 'git+' requirements/common.txt requirements/rocm.txt \
+    && pip list --format=freeze \
+        | grep -iE '^(torch|triton|torchvision|torchaudio|amdsmi|flash[-_]attn|amd[-_]aiter)==' >> requirements/rocm.txt \
+    && SETUPTOOLS_SCM_PRETEND_VERSION="$(cat /vllm-version.txt)" VLLM_USE_PRECOMPILED=1 \
+       VLLM_PRECOMPILED_WHEEL_LOCATION="$(ls /precompiled/*.whl)" VLLM_DOCKER_BUILD_CONTEXT=1 \
+       PYTHONDONTWRITEBYTECODE=1 python3 setup.py bdist_wheel --dist-dir=/dist
+
+# The installed csrc extensions (~700 MB), as their own layer: it depends only on the csrc stage,
+# so a commit that leaves csrc alone reuses the layer and pulls only the Python one below.
+FROM dev AS native
+RUN --mount=type=bind,from=csrc,src=/dist,target=/csrc \
+    sp="$(python3 -c 'import sysconfig; print(sysconfig.get_paths()["purelib"])')" \
+    && mkdir -p "/out$sp" && unzip -q /csrc/*.whl 'vllm/*.so' -d "/out$sp"
+
+# Bytecode compiled as if installed at / (not /out), so tracebacks name the real paths.
+# The csrc extensions come from the native layer instead.
+FROM runtime AS install
+RUN --mount=type=bind,from=wheel,src=/dist,target=/dist \
+    --mount=type=bind,from=csrc,src=/dist,target=/csrc \
+    pip install --no-deps --no-compile --root /out /dist/*.whl \
+    && sp="$(python3 -c 'import sysconfig; print(sysconfig.get_paths()["purelib"])')" \
+    && unzip -Z1 /csrc/*.whl 'vllm/*.so' | sed "s|^|/out$sp/|" | xargs rm \
+    && python3 -m compileall -q -j 0 -s /out -p / /out
+
+FROM runtime AS vllm-openai
+COPY --link --from=native /out /
+COPY --link --from=install /out /
+COPY --link benchmarks /app/vllm/benchmarks
+COPY --link examples /app/vllm/examples
+ARG CI_ARCH
+LABEL ai.vllm.rocm.ci-arch=${CI_ARCH}
+ARG VLLM_VERSION
+LABEL org.opencontainers.image.version=${VLLM_VERSION}
+ENTRYPOINT ["vllm", "serve"]
 
-FROM final_lmcache_${INSTALL_LMCACHE} AS final
+# Release check (cacheonly): vLLM imports offline and _C carries every arch of PYTORCH_ROCM_ARCH.
+FROM vllm-openai AS openai-smoke
+RUN --network=none so="$(python3 -c 'import os, vllm, vllm._C; print(os.path.dirname(vllm.__file__))')/_C.abi3.so" \
+    && for a in $(echo "$PYTORCH_ROCM_ARCH" | tr ';' ' '); do \
+         llvm-objdump --offloading "$so" | grep -qw "$a" || { echo "_C lacks $a" >&2; exit 1; }; done
+
+FROM test-deps AS test-base
+# AITER nightly: replace the stack's aiter with the newest nightly wheel that fits this image's ROCm,
+# torch and Python (never cached, see the bake file). AITER_SPEC pins a pip spec instead.
+ARG AITER_NIGHTLY
+ARG AITER_SPEC
+RUN --mount=type=bind,source=.buildkite/scripts/rocm/select_aiter_nightly_wheel.py,target=/tmp/select.py \
+    set -e; spec="$AITER_SPEC"; \
+    if [ -z "$spec" ] && [ "$AITER_NIGHTLY" = 1 ]; then \
+        out="$(python3 -P /tmp/select.py)" || { echo "$out" >&2; exit 1; }; \
+        echo "$out"; spec="$(echo "$out" | head -n 1)"; \
+    fi; \
+    if [ -n "$spec" ]; then \
+        pip install --no-cache-dir --pre --no-deps --force-reinstall "$spec" || exit 1; \
+        # Locate aiter without importing it: import needs a GPU, which a build has not. A wheel built
+        # against a newer glibc or libstdc++ than this image installs fine but fails every test that
+        # loads its modules; fail here instead.
+        aiter="$(python3 -c 'import importlib.util as u; print(u.find_spec("aiter").origin)')"; \
+        ! find "${aiter%/*}" -name '*.so' -exec ldd {} + 2>&1 | grep "version .* not found"; \
+    fi
 
-#Set entrypoint for vllm-openai official images
-FROM final AS vllm-openai
-ENTRYPOINT ["vllm", "serve"]
+FROM test-base AS test
+COPY --link --from=native /out /
+COPY --link --from=install /out /
+COPY --link --exclude=vllm --exclude=.git . /vllm-workspace/
+# Source under src/ so it does not shadow the installed package in tests.
+COPY --link vllm /vllm-workspace/src/vllm
+WORKDIR /vllm-workspace
+
+FROM test AS test-smoke
+RUN --network=none python3 -c "import vllm, vllm._C, vllm._rocm_C; print(vllm.__version__)"
+
+FROM scratch AS wheels
+COPY --link --from=stack-wheels / /
+COPY --link --from=wheel /dist/ /
+
+# CI: the vLLM wheel alone (a build artifact for the python-only install test), and the expensive
+# compile stages, whose cache PR builds export.
+FROM scratch AS vllm-wheel
+COPY --link --from=wheel /dist/ /
+FROM scratch AS compile
+COPY --link --from=csrc /dist /csrc
+COPY --link --from=rust-deps /app/vllm/rust/Cargo.lock /rust/
diff --git a/docker/Dockerfile.rocm_72 b/docker/Dockerfile.rocm_72
deleted file mode 100644
index eeaa9d6c1b18..000000000000
--- a/docker/Dockerfile.rocm_72
+++ /dev/null
@@ -1,1018 +0,0 @@
-# default base image
-ARG REMOTE_VLLM="0"
-ARG VLLM_REPO="https://github.com/vllm-project/vllm.git"
-ARG VLLM_BRANCH="main"
-ARG COMMON_WORKDIR=/app
-ARG BASE_IMAGE=vllm/vllm-openai-rocm:base-nightly-rocm72
-ARG CI_BASE_IMAGE=rocm/vllm-dev:ci_base
-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
-# (ainic and bnxt) are installed; MoRI selects the appropriate one at runtime.
-# To install drivers for a single NIC type only, set NIC_BACKEND explicitly:
-#   --build-arg NIC_BACKEND=ainic   # AMD AINIC (Pensando) only
-#   --build-arg NIC_BACKEND=bnxt    # Broadcom Thor-2 only
-#   --build-arg NIC_BACKEND=none    # Install nothing.
-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 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
-# shared S3-compatible cache backend is available.
-ARG USE_SCCACHE
-ARG SCCACHE_DOWNLOAD_URL
-ARG SCCACHE_ENDPOINT
-ARG SCCACHE_BUCKET_NAME=vllm-build-sccache
-ARG SCCACHE_REGION_NAME=us-west-2
-ARG SCCACHE_S3_NO_CREDENTIALS=0
-
-FROM ${BASE_IMAGE} AS base
-
-ARG ARG_PYTORCH_ROCM_ARCH
-ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}}
-
-# Install build dependencies and utilities
-RUN apt-get update -q -y && apt-get install -q -y \
-    sqlite3 libsqlite3-dev libfmt-dev libmsgpack-dev libsuitesparse-dev \
-    apt-transport-https ca-certificates wget curl \
-    libnuma-dev ccache mold
-RUN --mount=type=cache,target=/root/.cache/pip \
-    python3 -m pip install --upgrade pip
-# Note: mold is installed but not set as the system default linker because
-# some packages use JIT compilation at runtime with flags mold does not support.
-# Build stages opt in via LDFLAGS="-fuse-ld=mold".
-# Remove sccache only if not using sccache (it exists in base image from Dockerfile.rocm_72_base)
-ARG USE_SCCACHE
-RUN if [ "$USE_SCCACHE" != "1" ]; then \
-        apt-get purge -y sccache || true; \
-        python3 -m pip uninstall -y sccache || true; \
-        rm -f "$(which sccache)" || true; \
-    fi
-
-# Install UV — download first, then run, so a curl failure is not masked by the pipe
-RUN curl -LsSf --retry 3 --retry-delay 5 https://astral.sh/uv/install.sh -o /tmp/uv-install.sh \
-    && env UV_INSTALL_DIR="/usr/local/bin" sh /tmp/uv-install.sh \
-    && rm -f /tmp/uv-install.sh \
-    && uv --version
-
-# This timeout (in seconds) is necessary when installing some dependencies via uv since it's likely to time out
-# Reference: https://github.com/astral-sh/uv/pull/1694
-ENV UV_HTTP_TIMEOUT=500
-ENV UV_INDEX_STRATEGY="unsafe-best-match"
-# Use copy mode to avoid hardlink failures with Docker cache mounts
-ENV UV_LINK_MODE=copy
-# ccache directory - persisted across layer rebuilds via cache mounts.
-ENV CCACHE_DIR=/root/.cache/ccache
-ENV CCACHE_COMPILERCHECK=content
-# Empty by default so build steps fall back to $(nproc); CI can override.
-ARG max_jobs
-ENV MAX_JOBS=${max_jobs}
-
-# Install sccache if USE_SCCACHE is enabled (for release builds)
-ARG USE_SCCACHE
-ARG SCCACHE_DOWNLOAD_URL
-ARG SCCACHE_ENDPOINT
-ARG SCCACHE_BUCKET_NAME
-ARG SCCACHE_REGION_NAME
-ARG SCCACHE_S3_NO_CREDENTIALS
-RUN if [ "$USE_SCCACHE" = "1" ]; then \
-        if command -v sccache >/dev/null 2>&1; then \
-            echo "sccache already installed, skipping installation"; \
-            sccache --version; \
-        else \
-            echo "Installing sccache..." \
-            && SCCACHE_ARCH="x86_64" \
-            && SCCACHE_VERSION="v0.8.1" \
-            && SCCACHE_DL_URL="${SCCACHE_DOWNLOAD_URL:-https://github.com/mozilla/sccache/releases/download/${SCCACHE_VERSION}/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl.tar.gz}" \
-            && curl -L -o /tmp/sccache.tar.gz ${SCCACHE_DL_URL} \
-            && tar -xzf /tmp/sccache.tar.gz -C /tmp \
-            && mv /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl/sccache /usr/bin/sccache \
-            && chmod +x /usr/bin/sccache \
-            && rm -rf /tmp/sccache.tar.gz /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl \
-            && sccache --version; \
-        fi; \
-    fi
-
-# Set sccache environment variables only when USE_SCCACHE=1
-# This prevents S3 config from leaking into images when sccache is not used
-ARG USE_SCCACHE
-ENV SCCACHE_BUCKET=${USE_SCCACHE:+${SCCACHE_BUCKET_NAME}}
-ENV SCCACHE_REGION=${USE_SCCACHE:+${SCCACHE_REGION_NAME}}
-ENV SCCACHE_S3_NO_CREDENTIALS=${USE_SCCACHE:+${SCCACHE_S3_NO_CREDENTIALS}}
-ENV SCCACHE_IDLE_TIMEOUT=${USE_SCCACHE:+0}
-
-ARG COMMON_WORKDIR
-WORKDIR ${COMMON_WORKDIR}
-
-
-# -----------------------
-# vLLM fetch stages
-FROM base AS fetch_vllm_0
-ONBUILD COPY ./ vllm/
-FROM base AS fetch_vllm_1
-ARG VLLM_REPO
-ARG VLLM_BRANCH
-ONBUILD RUN git clone ${VLLM_REPO} \
-	    && cd vllm \
-	    && git fetch -v --prune -- origin ${VLLM_BRANCH} \
-	    && git checkout FETCH_HEAD \
-        && if [ ${VLLM_REPO} != "https://github.com/vllm-project/vllm.git" ] ; then \
-               git remote add upstream "https://github.com/vllm-project/vllm.git" \
-               && git fetch upstream ; fi
-FROM fetch_vllm_${REMOTE_VLLM} AS fetch_vllm
-
-# Resolve the device-independent source version once for every downstream
-# artifact. setuptools-scm reads Git metadata for local/remote sources and the
-# archive metadata generated for canonical CI contexts.
-FROM fetch_vllm AS vllm-version
-ARG COMMON_WORKDIR
-RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \
-    cd ${COMMON_WORKDIR}/vllm \
-    && uv pip install --system -r requirements/build/rust.txt \
-    && python3 -c \
-        "from setuptools_scm import get_version; print(get_version(git_describe_command=['git', 'describe', '--dirty', '--tags', '--long', '--abbrev=40', '--match', 'v[0-9]*']))" \
-        > /vllm-version.txt
-
-# -----------------------
-# Rust toolchain inputs
-#
-# Keep these separate from Rust source inputs so Rust source-only changes do not
-# invalidate the long-lived ci_base image that carries the toolchain.
-FROM scratch AS rust_toolchain_input_0
-COPY rust-toolchain.toml /rust-toolchain-input/vllm/rust-toolchain.toml
-
-FROM base AS rust_toolchain_input_1
-ARG VLLM_REPO
-ARG VLLM_BRANCH
-RUN git clone --no-checkout --filter=blob:none ${VLLM_REPO} /rust-toolchain-input/vllm \
-    && cd /rust-toolchain-input/vllm \
-    && git fetch -v --prune -- origin ${VLLM_BRANCH} \
-    && git sparse-checkout init --no-cone \
-    && git sparse-checkout set \
-        rust-toolchain.toml \
-    && git checkout FETCH_HEAD
-
-FROM rust_toolchain_input_${REMOTE_VLLM} AS rust-toolchain-input
-
-# -----------------------
-# Rust build inputs
-#
-# REMOTE_VLLM=0: use only the local files that can affect Rust artifacts.
-# REMOTE_VLLM=1: fetch only those same files from the requested remote ref.
-FROM scratch AS rust_input_0
-COPY requirements/build/rust.txt /rust-input/vllm/requirements/build/rust.txt
-COPY rust/Cargo.lock rust/Cargo.toml /rust-input/vllm/rust/
-COPY rust/proto /rust-input/vllm/rust/proto
-COPY rust/src /rust-input/vllm/rust/src
-COPY rust-toolchain.toml /rust-input/vllm/rust-toolchain.toml
-COPY tools/build_rust.py /rust-input/vllm/tools/build_rust.py
-COPY tools/build_rust.sh /rust-input/vllm/tools/build_rust.sh
-
-FROM base AS rust_input_1
-ARG VLLM_REPO
-ARG VLLM_BRANCH
-RUN git clone --no-checkout --filter=blob:none ${VLLM_REPO} /rust-input/vllm \
-    && cd /rust-input/vllm \
-    && git fetch -v --prune -- origin ${VLLM_BRANCH} \
-    && git sparse-checkout init --no-cone \
-    && git sparse-checkout set \
-        requirements/build/rust.txt \
-        rust/Cargo.lock \
-        rust/Cargo.toml \
-        rust/proto/** \
-        rust/src/** \
-        rust-toolchain.toml \
-        tools/build_rust.py \
-        tools/build_rust.sh \
-    && git checkout FETCH_HEAD
-
-FROM rust_input_${REMOTE_VLLM} AS rust-input
-
-# -----------------------
-# Rust toolchain
-FROM base AS rust-toolchain
-
-ENV CARGO_HOME=/root/.cargo
-ENV RUSTUP_HOME=/root/.rustup
-ENV PATH=${CARGO_HOME}/bin:${PATH}
-
-RUN apt-get update -q -y && apt-get install -q -y --no-install-recommends \
-        ca-certificates curl \
-    && rm -rf /var/lib/apt/lists/*
-
-COPY --from=rust-toolchain-input /rust-toolchain-input/vllm/rust-toolchain.toml /tmp/rust-toolchain.toml
-RUN TOOLCHAIN="$(grep '^channel' /tmp/rust-toolchain.toml | sed 's/.*= *"\(.*\)"/\1/')" \
-    && if ! command -v rustup >/dev/null 2>&1; then \
-        curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs \
-            | sh -s -- -y --default-toolchain none; \
-    fi \
-    && rustup toolchain install "${TOOLCHAIN}" \
-    && rustup default "${TOOLCHAIN}" \
-    && rustc --version \
-    && cargo --version
-
-# -----------------------
-# Rust build stage
-# Builds the `vllm-rs` frontend in a dedicated stage so the wheel build stages
-# don't need the rust toolchain.
-FROM rust-toolchain AS rust-build
-ARG COMMON_WORKDIR
-ARG USE_SCCACHE
-
-# Cap cargo parallelism to avoid exhausting the AMD CI host's open-file limit
-# (rustc spawns enough concurrent processes to hit RLIMIT_NOFILE otherwise).
-ENV CARGO_BUILD_JOBS=4
-ENV CARGO_NET_RETRY=10
-ENV RUSTUP_MAX_RETRIES=10
-
-WORKDIR ${COMMON_WORKDIR}/vllm
-
-COPY --from=rust-input /rust-input/vllm/requirements/build/rust.txt requirements/build/rust.txt
-
-# BuildKit can run this stage in parallel with ROCm native builds. Keep Rust on
-# a separate local sccache daemon while sharing the same remote cache backend.
-ENV SCCACHE_SERVER_PORT=4227
-
-RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \
-    uv pip install --system -r requirements/build/rust.txt
-
-COPY --from=rust-input /rust-input/vllm/rust rust
-COPY --from=rust-input /rust-input/vllm/rust-toolchain.toml rust-toolchain.toml
-COPY --from=rust-input /rust-input/vllm/tools/build_rust.py tools/build_rust.py
-COPY --from=rust-input /rust-input/vllm/tools/build_rust.sh tools/build_rust.sh
-COPY --from=vllm-version /vllm-version.txt /tmp/vllm-version.txt
-
-# Build the release binary. Cargo's registry/git caches can be written by
-# concurrent BuildKit jobs on shared workers, so lock those cache mounts while
-# keeping the cache benefit. Do not cache target/, because stale target metadata
-# can outlive source updates across BuildKit cache reuse.
-RUN --mount=type=cache,id=vllm-rocm-cargo-registry,target=/root/.cargo/registry,sharing=locked \
-    --mount=type=cache,id=vllm-rocm-cargo-git,target=/root/.cargo/git,sharing=locked \
-    if [ "$USE_SCCACHE" = "1" ]; then \
-        export RUSTC_WRAPPER=sccache \
-        && sccache --show-stats; \
-    fi \
-    && export SETUPTOOLS_SCM_PRETEND_VERSION="$(cat /tmp/vllm-version.txt)" \
-    && bash tools/build_rust.sh \
-    && test -x vllm/vllm-rs \
-    && if [ "$USE_SCCACHE" = "1" ]; then \
-        sccache --show-stats; \
-    fi
-
-# -----------------------
-# Shared vLLM build dependencies
-#
-# Install requirements once, before bringing source-specific inputs into either
-# consumer. Native compilation and wheel packaging then share this expensive
-# layer while retaining their independent source cache boundaries.
-FROM base AS build_vllm_dependencies
-ARG COMMON_WORKDIR
-ENV VLLM_TARGET_DEVICE=rocm
-WORKDIR ${COMMON_WORKDIR}/vllm
-
-COPY --from=fetch_vllm \
-    ${COMMON_WORKDIR}/vllm/requirements/common.txt \
-    ${COMMON_WORKDIR}/vllm/requirements/rocm.txt \
-    requirements/
-
-RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \
-    uv pip install --system -r requirements/rocm.txt
-
-# Fetch only the pinned ROCm Triton kernels package. Keeping this before the
-# vLLM source COPYs makes the download reusable when native sources change.
-FROM base AS rocm-triton-kernels
-ARG ROCM_TRITON_KERNELS_COMMIT
-ADD https://codeload.github.com/ROCm/triton/tar.gz/${ROCM_TRITON_KERNELS_COMMIT} \
-    /tmp/rocm-triton.tar.gz
-RUN mkdir -p /opt/rocm-triton-kernels \
-    && tar --extract --gzip --file /tmp/rocm-triton.tar.gz \
-        --directory /opt/rocm-triton-kernels --strip-components=4 \
-        --no-same-owner \
-        "triton-${ROCM_TRITON_KERNELS_COMMIT}/python/triton_kernels/triton_kernels" \
-    && test -f /opt/rocm-triton-kernels/__init__.py \
-    && rm /tmp/rocm-triton.tar.gz
-
-# -----------------------
-# vLLM native build stages
-#
-# csrc-build intentionally copies only files that affect ROCm native extension
-# compilation. That keeps unrelated CI/test/docs edits from invalidating the
-# expensive HIP/C++ build layer.
-FROM build_vllm_dependencies AS csrc-build
-ARG ROCM_TRITON_KERNELS_COMMIT
-COPY --from=rocm-triton-kernels \
-    /opt/rocm-triton-kernels /opt/rocm-triton-kernels
-ENV TRITON_KERNELS_SRC_DIR=/opt/rocm-triton-kernels
-
-# pyproject.toml is bind-mounted so it is not persisted in the csrc-build image,
-# but it still participates in the BuildKit cache key for this RUN step.
-COPY setup.py CMakeLists.txt ./
-COPY tools/build_rust.py tools/build_rust.py
-COPY cmake cmake/
-COPY csrc csrc/
-COPY vllm/envs.py vllm/envs.py
-COPY vllm/__init__.py vllm/__init__.py
-
-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}\")" \
-        cmake/external_projects/triton_kernels.cmake \
-    && export CCACHE_BASEDIR="$PWD" \
-    && echo "=== ccache stats before ROCm native build ===" \
-    && (ccache --show-stats || true) \
-    && (ccache --zero-stats || true) \
-    && EFFECTIVE_MAX_JOBS="${MAX_JOBS:-$(nproc)}" \
-    && echo "Building ROCm native extension wheel with MAX_JOBS=${EFFECTIVE_MAX_JOBS}" \
-    && PYTHONDONTWRITEBYTECODE=1 LDFLAGS="-fuse-ld=mold" MAX_JOBS="${EFFECTIVE_MAX_JOBS}" python3 setup.py bdist_wheel --dist-dir=dist \
-    && test -d dist \
-    && ls dist/*.whl >/dev/null \
-    && echo "=== ccache stats after ROCm native build ===" \
-    && (ccache --show-stats || true)
-
-# Build the full vLLM ROCm wheel by reusing the native extension wheel from
-# csrc-build. This stage still rebuilds for Python/package changes, but skips
-# the expensive dependency installation and HIP/C++ compilation when their
-# respective inputs are unchanged.
-FROM build_vllm_dependencies AS build_vllm
-COPY --from=fetch_vllm ${COMMON_WORKDIR}/vllm ${COMMON_WORKDIR}/vllm
-COPY --from=csrc-build ${COMMON_WORKDIR}/vllm/dist /precompiled-wheels
-COPY --from=vllm-version /vllm-version.txt /tmp/vllm-version.txt
-
-# Drop the pre-built Rust artifacts into the source tree. setup.py detects
-# them and ships them as-is, skipping the local Rust build.
-COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/vllm-rs ${COMMON_WORKDIR}/vllm/vllm/vllm-rs
-COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/_rust_*.so ${COMMON_WORKDIR}/vllm/vllm/
-
-RUN cd "${COMMON_WORKDIR}/vllm" \
-    && export SETUPTOOLS_SCM_PRETEND_VERSION="$(cat /tmp/vllm-version.txt)" \
-    && export VLLM_USE_PRECOMPILED=1 \
-    && export VLLM_PRECOMPILED_WHEEL_LOCATION="$(ls /precompiled-wheels/*.whl)" \
-    && export VLLM_DOCKER_BUILD_CONTEXT=1 \
-    && echo "Packaging vLLM ROCm wheel using precompiled extensions from ${VLLM_PRECOMPILED_WHEEL_LOCATION}" \
-    && PYTHONDONTWRITEBYTECODE=1 python3 setup.py bdist_wheel --dist-dir=dist \
-    && test -d dist \
-    && ls dist/*.whl >/dev/null
-FROM scratch AS export_vllm
-ARG COMMON_WORKDIR
-COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/dist/*.whl /
-COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/requirements /requirements
-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/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_72 /docker/
-COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm_72_base /docker/
-COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/ci-rocm.hcl /docker/ci-rocm.hcl
-COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/docker-bake.hcl /docker/docker-bake.hcl
-COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/docker-bake-rocm.hcl /docker/docker-bake-rocm.hcl
-COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/.buildkite /.buildkite
-COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/pyproject.toml /pyproject.toml
-COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/rust /rust
-COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/rust-toolchain.toml /rust-toolchain.toml
-COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1
-
-# NIXL/UCX build stages
-FROM base AS build_nixl
-ARG NIXL_BRANCH="1683cf3b7f3d11674c03c5e861cea22339876c96"
-ARG NIXL_REPO="https://github.com/ai-dynamo/nixl.git"
-ARG UCX_BRANCH="2ff0b47ce7626584cb9b56c647d08f712c6c3edf"
-ARG UCX_REPO="https://github.com/openucx/ucx.git"
-ENV ROCM_PATH=/opt/rocm
-ENV UCX_HOME=/usr/local/ucx
-ENV NIXL_HOME=/usr/local/nixl
-
-# NIXL build system dependencies and RDMA support
-RUN apt-get -y update && apt-get -y install autoconf libtool pkg-config \
-    libgrpc-dev \
-    libgrpc++-dev \
-    libprotobuf-dev \
-    protobuf-compiler-grpc \
-    libcpprest-dev \
-    libaio-dev \
-    librdmacm1 \
-    librdmacm-dev \
-    libibverbs1 \
-    libibverbs-dev \
-    ibverbs-utils \
-    rdmacm-utils \
-    ibverbs-providers \
-    && rm -rf /var/lib/apt/lists/*
-
-RUN --mount=type=cache,target=/root/.cache/uv \
-    uv pip install --system meson meson-python pyyaml types-PyYAML \
-        auditwheel build patchelf pytest tomlkit "setuptools>=80.9.0"
-
-RUN --mount=type=cache,target=/root/.cache/ccache \
-    cd /usr/local/src && \
-    git clone ${UCX_REPO} &&  \
-    cd ucx  && \
-    git checkout ${UCX_BRANCH} && \
-    ./autogen.sh && \
-    mkdir build && cd build && \
-    CC="ccache gcc" CXX="ccache g++" \
-    ../configure \
-        --prefix=/usr/local/ucx \
-        --enable-shared \
-        --disable-static \
-        --disable-doxygen-doc \
-        --enable-optimizations \
-        --enable-devel-headers \
-        --with-rocm=${ROCM_PATH} \
-        --with-verbs \
-        --with-dm \
-        --enable-mt && \
-    make -j$(nproc) && \
-    make install
-
-ENV PATH=/usr/local/ucx/bin:$PATH
-ENV LD_LIBRARY_PATH=${UCX_HOME}/lib:${LD_LIBRARY_PATH}
-
-RUN --mount=type=cache,target=/root/.cache/ccache \
-    git clone ${NIXL_REPO} /opt/nixl && \
-    cd /opt/nixl && \
-    git checkout ${NIXL_BRANCH} && \
-    CC="ccache gcc" CXX="ccache g++" \
-    meson setup build --prefix=${NIXL_HOME} \
-                     -Ducx_path=${UCX_HOME} \
-                     -Dwheel_variant=rocm \
-                     -Dbuild_tests=false \
-                     -Dbuild_examples=false && \
-    cd build && \
-    ninja -j$(nproc) && \
-    ninja install && \
-    echo "${NIXL_HOME}/lib/$(uname -m)-linux-gnu" \
-        > /etc/ld.so.conf.d/nixl.conf && \
-    echo "${NIXL_HOME}/lib/$(uname -m)-linux-gnu/plugins" \
-        >> /etc/ld.so.conf.d/nixl.conf && \
-    ldconfig
-
-# Generate the ROCm NIXL wheel. Upstream's generic wheel helper detects CUDA,
-# so configure the ROCm wheel variant directly through Meson.
-# Exclude libcore and libpull from auditwheel: transitive dependencies
-# that are not shipped in the wheel and vary across base images.
-# Tag manylinux_2_35 to match the Ubuntu 22.04 base: NIXL >= 1.5 links
-# GLIBCXX_3.4.30 symbols, which manylinux_2_34 does not allow.
-RUN cd /opt/nixl && \
-    ./contrib/tomlutil.py --wheel-name nixl-rocm pyproject.toml && \
-    CC="ccache gcc" CXX="ccache g++" \
-    uv build --wheel --no-build-isolation --out-dir /tmp/nixl_wheels \
-        --python ${PYTHON_VERSION} \
-        -Csetup-args=-Ducx_path=${UCX_HOME} \
-        -Csetup-args=-Dwheel_variant=rocm \
-        -Csetup-args=-Dbuild_tests=false \
-        -Csetup-args=-Dbuild_examples=false && \
-    mkdir -p /tmp/nixl_wheels/repaired /app/install && \
-    auditwheel repair \
-        --exclude 'libamdhip64*' \
-        --exclude 'libcore*' \
-        --exclude 'libpull*' \
-        /tmp/nixl_wheels/nixl_rocm*.whl \
-        --plat manylinux_2_35_$(uname -m) \
-        --wheel-dir /tmp/nixl_wheels/repaired && \
-    ./contrib/wheel_add_ucx_plugins.py \
-        --ucx-plugins-dir ${UCX_HOME}/lib/ucx \
-        --nixl-plugins-dir ${NIXL_HOME}/lib/$(uname -m)-linux-gnu/plugins \
-        /tmp/nixl_wheels/repaired/*.whl && \
-    cp /tmp/nixl_wheels/repaired/*.whl /app/install
-
-# Fetch and verify pinned LMCache source.
-FROM base AS lmcache_source
-ARG LMCACHE_REPO="https://github.com/LMCache/LMCache.git"
-ARG LMCACHE_REF="140819c9d57a975dbc5678a6459a218e544cb58b"
-ARG LMCACHE_VERSION="0.5.3"
-
-# Remove NVIDIA-only dependencies and accept the release image's NumPy version.
-RUN printf '%s\n' "${LMCACHE_REF}" | grep -Eq '^[0-9a-f]{40}$' \
-    && git clone --depth=1 --branch "v${LMCACHE_VERSION}" \
-        "${LMCACHE_REPO}" /app/lmcache \
-    && test "$(git -C /app/lmcache rev-parse HEAD)" = "${LMCACHE_REF}" \
-    && test "$(grep -Ec '^(cufile-python|nvtx)([<=>[:space:]]|$)' \
-        /app/lmcache/requirements/common.txt)" -eq 2 \
-    && test "$(grep -Fxc 'numpy<=2.2.6' \
-        /app/lmcache/requirements/common.txt)" -eq 1 \
-    && sed -i -E \
-        -e '/^(cufile-python|nvtx)([<=>[:space:]]|$)/d' \
-        -e 's/^numpy<=2\.2\.6$/numpy<=2.3.5/' \
-        /app/lmcache/requirements/common.txt \
-    && rm -rf /app/lmcache/.git
-
-FROM base AS build_lmcache
-COPY --from=lmcache_source \
-    /app/lmcache/requirements/build.txt /tmp/lmcache-build-requirements.txt
-RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \
-    uv pip install --system -r /tmp/lmcache-build-requirements.txt
-
-ARG LMCACHE_VERSION="0.5.3"
-ARG LMCACHE_ROCM_ARCH="gfx942;gfx950"
-ARG USE_SCCACHE
-
-COPY --from=lmcache_source /app/lmcache /app/lmcache
-WORKDIR /app/lmcache
-
-RUN --mount=type=cache,id=vllm-rocm-ccache,target=/root/.cache/ccache \
-    set -eu; \
-    printf '%s\n' "${LMCACHE_ROCM_ARCH}" \
-        | grep -Eq '^gfx[0-9a-z]+(;gfx[0-9a-z]+)*$'; \
-    offload_flags="$(printf '%s' "${LMCACHE_ROCM_ARCH}" \
-        | sed 's/;/ --offload-arch=/g; s/^/--offload-arch=/')"; \
-    if [ "${USE_SCCACHE:-0}" = "1" ]; then \
-        export HIP_CLANG_PATH=/opt/sccache-wrappers; \
-    fi; \
-    mkdir -p /app/install; \
-    SETUPTOOLS_SCM_PRETEND_VERSION="${LMCACHE_VERSION}" \
-    PYTORCH_ROCM_ARCH="${LMCACHE_ROCM_ARCH}" \
-    HIPCC_COMPILE_FLAGS_APPEND="${offload_flags}" \
-    BUILD_WITH_HIP=1 CXX=hipcc \
-        python3 setup.py bdist_wheel --dist-dir=/app/install; \
-    test "$(find /app/install -maxdepth 1 -name '*.whl' | wc -l)" -eq 1
-
-RUN --mount=type=tmpfs,target=/tmp/lmcache-wheel \
-    set -eu; \
-    python3 -m zipfile -e /app/install/*.whl /tmp/lmcache-wheel; \
-    c_ops="$(find /tmp/lmcache-wheel -type f -name 'c_ops*.so')"; \
-    test "$(find /tmp/lmcache-wheel -type f -name 'c_ops*.so' | wc -l)" -eq 1; \
-    expected="$(printf '%s' "${LMCACHE_ROCM_ARCH}" | tr ';' '\n' | sort -u)"; \
-    actual="$(/opt/rocm/llvm/bin/llvm-objdump --offloading "${c_ops}" \
-        | grep -oE 'gfx[0-9a-z]+' | sort -u)"; \
-    if [ "${actual}" != "${expected}" ]; then \
-        printf 'LMCache code objects mismatch\nexpected:\n%s\nactual:\n%s\n' \
-            "${expected}" "${actual}" >&2; \
-        exit 1; \
-    fi
-
-# ROCShmem build stage - split from DeepEP so changing DEEPEP_BRANCH does not
-# invalidate the slow ROCShmem build.
-FROM base AS build_rocshmem
-ARG ROCSHMEM_BRANCH="f0acb0c6"
-ARG ROCSHMEM_REPO="https://github.com/ROCm/rocm-systems.git"
-# DeepEP only supports gfx942 and gfx950; build ROCShmem for the same set so
-# it can be linked against DeepEP without arch mismatches.
-ARG DEEPEP_ROCM_ARCH="gfx942;gfx950"
-ENV ROCM_PATH=/opt/rocm
-ENV ROCSHMEM_DIR=/opt/rocshmem
-
-RUN --mount=type=cache,target=/root/.cache/ccache \
-    git clone --no-checkout --filter=blob:none ${ROCSHMEM_REPO} \
- && cd rocm-systems \
- && git sparse-checkout set --cone projects/rocshmem \
- && git checkout ${ROCSHMEM_BRANCH} \
- && mkdir -p projects/rocshmem/build \
- && cd projects/rocshmem/build \
- && CC="ccache gcc" CXX="ccache g++" INSTALL_PREFIX=${ROCSHMEM_DIR} \
-    bash ../scripts/build_configs/all_backends \
-      -DROCM_PATH=${ROCM_PATH} \
-      -DGPU_TARGETS="${DEEPEP_ROCM_ARCH}" \
-      -DUSE_EXTERNAL_MPI=OFF
-
-# DeepEP build stage - depends on ROCShmem, builds the HIP kernel wheel.
-FROM build_rocshmem AS build_deepep
-ARG DEEPEP_BRANCH="0f63d3e9"
-ARG DEEPEP_REPO="https://github.com/ROCm/DeepEP.git"
-ARG DEEPEP_NIC="cx7"
-
-# Build DeepEP wheel. DeepEP looks for rocshmem at ROCSHMEM_DIR.
-# DeepEP only supports gfx942 and gfx950, so avoid gfx90a in the default list.
-# We pass --rocm-gfx942-fp8-fnuz-max 224 to align the DeepEP dispatch FP8
-# quantization bound with vLLM's own clamp
-# (vllm/model_executor/layers/quantization/utils/fp8_utils.py): 224 for E4M3FNUZ
-RUN --mount=type=cache,target=/root/.cache/ccache \
-    export PYTORCH_ROCM_ARCH="gfx942;gfx950" \
- && git clone ${DEEPEP_REPO} \
- && cd DeepEP \
- && git checkout ${DEEPEP_BRANCH} \
- && LDFLAGS="-fuse-ld=mold" MAX_JOBS="${MAX_JOBS:-$(nproc)}" python3 setup.py --variant rocm --rocm-explicit-ctx --nic ${DEEPEP_NIC} --rocm-gfx942-fp8-fnuz-max 224 bdist_wheel --dist-dir=/app/deep_install
-
-# MoRI runtime dependencies live in Dockerfile.rocm_72 so NIC backend changes do
-# not force users to rebuild the long-lived Dockerfile.rocm_72_base image.
-FROM base AS mori_base
-ARG NIC_BACKEND
-ARG AINIC_VERSION
-ARG UBUNTU_CODENAME
-RUN /bin/bash -lc 'set -euo pipefail; \
- \
- install_ainic() { \
-   apt-get update && apt-get install -y --no-install-recommends ca-certificates curl gnupg apt-transport-https; \
-   rm -rf /var/lib/apt/lists/*; \
-   mkdir -p /etc/apt/keyrings; \
-   curl -fsSL https://repo.radeon.com/rocm/rocm.gpg.key | gpg --dearmor > /etc/apt/keyrings/amdainic.gpg; \
-   echo "deb [arch=amd64 signed-by=/etc/apt/keyrings/amdainic.gpg] https://repo.radeon.com/amdainic/pensando/ubuntu/${AINIC_VERSION} ${UBUNTU_CODENAME} main" \
-     > /etc/apt/sources.list.d/amdainic.list; \
-   apt-get update && apt-get install -y --no-install-recommends \
-     libionic-dev \
-     ionic-common \
-   ; \
-   rm -rf /var/lib/apt/lists/*; \
- }; \
- \
- # NOTE: requires FW 235.2.86.0 and kernel drivers on the host: \
- #   bnxt-en-dkms=1.10.3.235.2.86.0 bnxt-re-dkms=235.2.86.0 (from packages.broadcom.com PPA) \
- install_bnxt() { \
-   install -m 0755 -d /etc/apt/keyrings; \
-   curl -fsSL https://packages.broadcom.com/artifactory/api/security/keypair/PackagesKey/public \
-     -o /etc/apt/keyrings/broadcom-nic.asc; \
-   chmod a+r /etc/apt/keyrings/broadcom-nic.asc; \
-   echo "deb [arch=amd64 signed-by=/etc/apt/keyrings/broadcom-nic.asc] https://packages.broadcom.com/artifactory/ethernet-nic-debian-public jammy main" \
-     > /etc/apt/sources.list.d/broadcom-nic.list; \
-   apt-get update && apt-get install -y --no-install-recommends \
-     bnxt-rocelib=235.2.86.0 \
-   ; \
-   cp -a /usr/local/lib/x86_64-linux-gnu/libbnxt_re* /usr/local/lib/; \
-   ldconfig; \
-   rm -rf /var/lib/apt/lists/*; \
- }; \
- \
- echo "[MORI] Install MoRI proxy deps"; \
- pip install --quiet --ignore-installed blinker && \
- pip install --quiet quart msgpack aiohttp pyzmq; \
- echo "[MORI] NIC_BACKEND=${NIC_BACKEND}"; \
- \
- # NIC backend deps — mori auto-detects NIC at runtime (MORI_DEVICE_NIC env var override). \
- # Only vendor packages are installed here for dlopen; no compile-time flags needed. \
- case "${NIC_BACKEND}" in \
-   none)  ;; \
-   all)   install_ainic; install_bnxt ;; \
-   ainic) install_ainic ;; \
-   bnxt)  install_bnxt ;; \
-   *)     echo "ERROR: unknown NIC_BACKEND=${NIC_BACKEND}. Use one of: none, ainic, bnxt, all"; exit 2 ;; \
- esac'
-
-# -----------------------
-# vLLM wheel release build stage (for building distributable wheels)
-# This stage pins dependencies to custom ROCm wheel versions and handles version detection
-FROM fetch_vllm AS build_vllm_wheel_release
-
-ARG COMMON_WORKDIR
-ARG ROCM_TRITON_KERNELS_COMMIT
-COPY --from=vllm-version /vllm-version.txt /tmp/vllm-version.txt
-COPY --from=rocm-triton-kernels \
-    /opt/rocm-triton-kernels /opt/rocm-triton-kernels
-ENV TRITON_KERNELS_SRC_DIR=/opt/rocm-triton-kernels
-
-# Drop the pre-built Rust artifacts into the source tree. setup.py detects
-# them and ships them as-is, skipping the local Rust build.
-COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/vllm-rs ${COMMON_WORKDIR}/vllm/vllm/vllm-rs
-COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/_rust_*.so ${COMMON_WORKDIR}/vllm/vllm/
-
-# Create /install directory for custom wheels
-RUN mkdir -p /install
-
-# Copy custom ROCm wheels from docker/context if they exist
-# COPY ensures Docker cache is invalidated when wheels change
-# .keep file ensures directory always exists for COPY to work
-COPY docker/context/base-wheels/ /tmp/base-wheels/
-# This is how we know if we are building for a wheel release or not.
-# If there are not wheels found there, we are not building for a wheel release. 
-# So we exit with an error. To skip this stage.
-RUN if [ -n "$(ls /tmp/base-wheels/*.whl 2>/dev/null)" ]; then \
-        echo "Found custom wheels - copying to /install"; \
-        cp /tmp/base-wheels/*.whl /install/ && \
-        echo "Copied custom wheels:"; \
-        ls -lh /install/; \
-    else \
-        echo "ERROR: No custom wheels found in docker/context/base-wheels/"; \
-        echo "Wheel releases require pre-built ROCm wheels."; \
-        exit 1; \
-    fi
-
-# GIT_REPO_CHECK: Verify repo is clean and tags are available (for release builds)
-# This matches CUDA's Dockerfile behavior for proper version detection via setuptools_scm
-ARG GIT_REPO_CHECK=0
-RUN if [ "$GIT_REPO_CHECK" != "0" ]; then \
-        echo "Running repository checks..."; \
-        cd vllm && bash tools/check_repo.sh; \
-    fi
-
-# 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}\")" \
-        vllm/cmake/external_projects/triton_kernels.cmake \
-    && cd vllm \
-    && uv pip install --system setuptools_scm regex
-
-# Fail if git-based package dependencies are found in requirements files
-# (uv doesn't handle git+ URLs well, and packages should be distributed on PyPI)
-# Extra notes: pip install is able to handle git+ URLs, but uv doesn't.
-RUN echo "Checking for git-based packages in requirements files..." \
-    && echo "Checking common.txt for git-based packages:" \
-    && if grep -q 'git+' ${COMMON_WORKDIR}/vllm/requirements/common.txt; then \
-         echo "ERROR: Git-based packages found in common.txt:"; \
-         grep 'git+' ${COMMON_WORKDIR}/vllm/requirements/common.txt; \
-         echo "Please publish these packages to PyPI instead of using git dependencies."; \
-         exit 1; \
-       else \
-         echo "  ✓ No git-based packages found in common.txt"; \
-       fi \
-    && echo "Checking rocm.txt for git-based packages:" \
-    && if grep -q 'git+' ${COMMON_WORKDIR}/vllm/requirements/rocm.txt; then \
-         echo "ERROR: Git-based packages found in rocm.txt:"; \
-         grep 'git+' ${COMMON_WORKDIR}/vllm/requirements/rocm.txt; \
-         echo "Please publish these packages to PyPI instead of using git dependencies."; \
-         exit 1; \
-       else \
-         echo "  ✓ No git-based packages found in rocm.txt"; \
-       fi \
-    && echo "All requirements files are clean - no git-based packages found"
-
-# Pin vLLM dependencies to exact versions of custom ROCm wheels
-# This ensures 'pip install vllm' automatically installs correct torch/triton/torchvision/amdsmi
-COPY tools/vllm-rocm/pin_rocm_dependencies.py /tmp/pin_rocm_dependencies.py
-RUN echo "Pinning vLLM dependencies to custom wheel versions..." \
-    && python3 /tmp/pin_rocm_dependencies.py /install ${COMMON_WORKDIR}/vllm/requirements/rocm.txt
-
-# Install dependencies using custom wheels from /install
-RUN --mount=type=cache,target=/root/.cache/uv \
-    cd vllm \
-    && echo "Building vLLM with custom wheels from /install" \
-    && uv pip install --system --find-links /install -r requirements/rocm.txt
-
-# Build wheel using the shared source version to avoid dirty state from modified requirements/rocm.txt
-# (setup.py auto-detects ccache/sccache in PATH)
-RUN --mount=type=cache,id=vllm-rocm-ccache,target=/root/.cache/ccache \
-    cd vllm \
-    && export CCACHE_BASEDIR="$PWD" \
-    && export SETUPTOOLS_SCM_PRETEND_VERSION="$(cat /tmp/vllm-version.txt)" \
-    && echo "Building wheel with version: ${SETUPTOOLS_SCM_PRETEND_VERSION}" \
-    && MAX_JOBS="${MAX_JOBS:-$(nproc)}" python3 setup.py bdist_wheel --dist-dir=dist
-
-FROM scratch AS export_vllm_wheel_release
-ARG COMMON_WORKDIR
-COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/dist/*.whl /
-COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/requirements /requirements
-COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/benchmarks /benchmarks
-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/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_72 /docker/
-COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm_72_base /docker/
-COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/ci-rocm.hcl /docker/ci-rocm.hcl
-COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/docker-bake.hcl /docker/docker-bake.hcl
-COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/docker-bake-rocm.hcl /docker/docker-bake-rocm.hcl
-COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/.buildkite /.buildkite
-COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/pyproject.toml /pyproject.toml
-COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/rust /rust
-COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/rust-toolchain.toml /rust-toolchain.toml
-COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1
-
-# -----------------------
-# CI base image (Tier 1) - stable, rarely changing CI dependencies.
-# Per-PR test builds pull this as CI_BASE_IMAGE so the test stage only layers
-# in the vLLM artifacts for the current commit.
-FROM mori_base AS ci_base
-ARG COMMON_WORKDIR
-
-ENV CARGO_HOME=/root/.cargo
-ENV RUSTUP_HOME=/root/.rustup
-ENV PATH=${CARGO_HOME}/bin:${PATH}
-
-# Bake the Rust toolchain into ci_base so Rust build/test steps do not
-# need to fetch it on each per-commit image build.
-COPY --from=rust-toolchain /root/.cargo /root/.cargo
-COPY --from=rust-toolchain /root/.rustup /root/.rustup
-
-# Update rdma-core to support latest rocshmem.
-ARG DEEPEP_NIC
-RUN if [ "${DEEPEP_NIC}" = "cx7" ] || [ "${DEEPEP_NIC}" = "io" ]; then \
-    git clone --branch v62.0 --depth 1 https://github.com/linux-rdma/rdma-core.git /tmp/rdma-core && \
-    cd /tmp/rdma-core && \
-    mkdir -p build && cd build && \
-    cmake -GNinja -DCMAKE_INSTALL_PREFIX=/usr -DNO_MAN_PAGES=1 .. && \
-    ninja && ninja install && ldconfig && rm -rf /tmp/rdma-core; \
-fi
-
-# Install NIXL + DeepEP wheels.
-RUN --mount=type=bind,from=build_nixl,src=/app/install,target=/nixl_install \
-    --mount=type=bind,from=build_deepep,src=/app/deep_install,target=/deep_install \
-    uv pip install --system /nixl_install/*.whl /deep_install/*.whl
-
-# Copy ROCShmem runtime libraries.
-COPY --from=build_rocshmem /opt/rocshmem /opt/rocshmem
-
-# RDMA userspace libraries plus FFmpeg dev libs needed by torchcodec.
-RUN timeout --kill-after=30s 5m apt-get update -q -y --error-on=any \
-    -o Acquire::Retries=3 \
-    -o Acquire::http::Timeout=30 \
-    -o Acquire::https::Timeout=30 \
-    && apt-get install -q -y --no-install-recommends \
-    -o Acquire::Retries=3 \
-    -o Acquire::http::Timeout=30 \
-    -o Acquire::https::Timeout=30 \
-    librdmacm1 \
-    libibverbs1 \
-    ibverbs-providers \
-    ibverbs-utils \
-    util-linux \
-    unzip \
-    libsodium23 \
-    pkg-config ffmpeg libavcodec-dev libavformat-dev libavutil-dev \
-    libswscale-dev libavdevice-dev libavfilter-dev libswresample-dev \
-    && rm -rf /var/lib/apt/lists/*
-
-# Install torchcodec from source for ROCm/torch ABI compatibility.
-COPY tools/install_torchcodec_rocm.sh /tmp/install_torchcodec.sh
-RUN --mount=type=cache,target=/root/.cache/uv \
-    --mount=type=cache,target=/root/.cache/pip \
-    --mount=type=cache,target=/root/.cache/torchcodec-wheels \
-    bash /tmp/install_torchcodec.sh \
-    && rm /tmp/install_torchcodec.sh \
-    && apt-get clean && rm -rf /var/lib/apt/lists/*
-
-# Pre-install shared ROCm runtime dependencies.
-COPY requirements/common.txt requirements/rocm.txt /tmp/ci-base-requirements/
-RUN --mount=type=cache,target=/root/.cache/uv \
-    uv pip install --system -r /tmp/ci-base-requirements/rocm.txt \
-    && rm -rf /tmp/ci-base-requirements
-
-# Enable fast and less brittle model downloads in tests.
-ENV HF_XET_HIGH_PERFORMANCE=1
-ENV HF_HUB_DOWNLOAD_TIMEOUT=60
-
-# Keep torch.cuda.is_available() fork-safe (see vllm/env_override.py).
-ENV PYTORCH_NVML_BASED_CUDA_CHECK=1
-
-# Pre-install vLLM test dependencies.
-COPY requirements/test/rocm.txt /tmp/rocm-test-reqs.txt
-RUN --mount=type=cache,target=/root/.cache/uv \
-    uv pip install --system -r /tmp/rocm-test-reqs.txt
-
-# Rebuild fastsafetensors from source so its C++ extension is compiled with
-# USE_ROCM and can detect libamdhip64.so at runtime.
-RUN --mount=type=cache,target=/root/.cache/pip \
-    FASTSAFETENSORS_REQ="$(grep -E '^fastsafetensors(==| @ )' /tmp/rocm-test-reqs.txt | head -1)" \
-    && test -n "${FASTSAFETENSORS_REQ}" \
-    && python3 -m pip install --force-reinstall --no-deps \
-        --no-binary fastsafetensors "${FASTSAFETENSORS_REQ}" \
-    && rm /tmp/rocm-test-reqs.txt
-
-# Set MIOPEN ENVS to resolve performance regressions in MIOpen 3D convolution kernel.
-# See: https://github.com/pytorch/pytorch/issues/169857
-ENV MIOPEN_DEBUG_CONV_DIRECT=0
-ENV MIOPEN_DEBUG_CONV_GEMM=0
-
-# Use legacy IPC mode for HSA to avoid GPU memory pinning issues with UCX rocm_ipc.
-# See: https://github.com/ROCm/rocm-libraries/issues/6266
-ENV HSA_ENABLE_IPC_MODE_LEGACY=1
-ENV UCX_ROCM_COPY_SIGPOOL_MAX_ELEMS=inf
-ENV UCX_ROCM_IPC_SIGPOOL_MAX_ELEMS=inf
-
-# ROCm profiler limits workaround.
-RUN echo "ROCTRACER_MAX_EVENTS=10000000" > ${COMMON_WORKDIR}/libkineto.conf
-ENV KINETO_CONFIG="${COMMON_WORKDIR}/libkineto.conf"
-
-# Install vllm_test_utils in ci_base for ci_base + wheel parity.
-COPY tests/vllm_test_utils /tmp/vllm_test_utils
-RUN --mount=type=cache,target=/root/.cache/uv \
-    uv pip install --system /tmp/vllm_test_utils \
-    && rm -rf /tmp/vllm_test_utils
-
-RUN --mount=type=bind,from=build_lmcache,src=/app/install,target=/lmcache_install \
-    --mount=type=bind,source=requirements/test/rocm.txt,target=/tmp/lmcache-constraints.txt \
-    --mount=type=cache,target=/root/.cache/uv \
-    uv pip install --system --constraint /tmp/lmcache-constraints.txt \
-        /lmcache_install/*.whl \
-    && python3 -c "import lmcache, lmcache.c_ops; print('lmcache', lmcache.__version__)"
-
-# -----------------------
-# Test vLLM image (Tier 2) - vLLM-only layer on top of ci_base.
-FROM ${CI_BASE_IMAGE} AS test
-ARG COMMON_WORKDIR
-
-# Install the vLLM wheel (--no-deps: all deps already in ci_base).
-RUN --mount=type=bind,from=export_vllm,src=/,target=/install \
-    --mount=type=cache,target=/root/.cache/uv \
-    cd /install \
-    && uv pip install --system --no-deps *.whl
-
-# Store the vLLM wheel in the image for python-only install tests.
-COPY --from=export_vllm /*.whl /opt/vllm-wheels/
-
-WORKDIR /vllm-workspace
-COPY --from=build_vllm ${COMMON_WORKDIR}/vllm /vllm-workspace
-
-# Copy in the v1 package (for python-only install test group).
-COPY --from=export_vllm /vllm_v1 /usr/local/lib/python${PYTHON_VERSION}/dist-packages/vllm/v1
-
-# Hide source under src/ so it won't shadow the installed package in tests.
-RUN mkdir src && mv vllm src/vllm
-
-# Catch GPU<->CPU syncs in execute_model/sample_tokens during tests.
-ENV VLLM_GPU_SYNC_CHECK=error
-
-# Run the structural smoke test inside BuildKit, where the test image layers
-# are already available. Export only a build-scoped success marker.
-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
-ENV UV_COMPILE_BYTECODE=1
-
-RUN python3 -m pip install --upgrade pip && rm -rf /var/lib/apt/lists/*
-
-# Clean up sccache from release image (not needed at runtime)
-# This removes the binary and wrappers that may have been installed during build
-RUN rm -f /usr/bin/sccache || true \
-    && rm -rf /opt/sccache-wrappers || true
-
-# Unset sccache environment variables for the release image
-# This prevents S3 bucket config from leaking into production images
-ENV SCCACHE_BUCKET=
-ENV SCCACHE_REGION=
-ENV SCCACHE_ENDPOINT=
-ENV SCCACHE_S3_NO_CREDENTIALS=
-ENV SCCACHE_IDLE_TIMEOUT=
-
-# Error related to odd state for numpy 1.20.3 where there is no METADATA etc, but an extra LICENSES_bundled.txt.
-# Manually remove it so that later steps of numpy upgrade can continue
-RUN case "$(which python3)" in \
-        *"/opt/conda/envs/py_3.9"*) \
-            rm -rf /opt/conda/envs/py_3.9/lib/python3.9/site-packages/numpy-1.20.3.dist-info/;; \
-        *) ;; esac
-
-RUN --mount=type=cache,target=/root/.cache/uv \
-    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
-# Keep OpenTelemetry within LMCache's cap.
-RUN --mount=type=bind,from=export_vllm,src=/,target=/install \
-    --mount=type=cache,target=/root/.cache/uv \
-    cd /install \
-    && echo "opentelemetry-api<=1.40.0" > /tmp/otel-constraints.txt \
-    && uv pip install --system -c /tmp/otel-constraints.txt -r requirements/rocm.txt \
-    && pip uninstall -y vllm \
-    && uv pip install --system -c /tmp/otel-constraints.txt *.whl \
-    && rm /tmp/otel-constraints.txt
-
-# Install NIXL ROCm wheel
-RUN --mount=type=bind,from=build_nixl,src=/app/install,target=/nixl_install \
-    uv pip install --system /nixl_install/*.whl
-
-ARG COMMON_WORKDIR
-ARG BASE_IMAGE
-ARG NIC_BACKEND
-ARG AINIC_VERSION
-
-# Copy over the benchmark scripts as well
-COPY --from=export_vllm /benchmarks ${COMMON_WORKDIR}/vllm/benchmarks
-COPY --from=export_vllm /examples ${COMMON_WORKDIR}/vllm/examples
-COPY --from=export_vllm /docker ${COMMON_WORKDIR}/vllm/docker
-
-# Use legacy IPC mode for HSA to avoid GPU memory pinning issues with UCX rocm_ipc
-# See: https://github.com/ROCm/rocm-libraries/issues/6266
-ENV HSA_ENABLE_IPC_MODE_LEGACY=1
-ENV UCX_ROCM_COPY_SIGPOOL_MAX_ELEMS=inf
-ENV UCX_ROCM_IPC_SIGPOOL_MAX_ELEMS=inf
-
-ENV TOKENIZERS_PARALLELISM=false
-
-# ENV that can improve safe tensor loading, and end-to-end time
-ENV SAFETENSORS_FAST_GPU=1
-
-# Performance environment variable.
-ENV HIP_FORCE_DEV_KERNARG=1
-
-# Keep torch.cuda.is_available() fork-safe (see vllm/env_override.py).
-ENV PYTORCH_NVML_BASED_CUDA_CHECK=1
-
-# Workaround for ROCm profiler limits
-RUN echo "ROCTRACER_MAX_EVENTS=10000000" > ${COMMON_WORKDIR}/libkineto.conf
-ENV KINETO_CONFIG="${COMMON_WORKDIR}/libkineto.conf"
-RUN echo "VLLM_BASE_IMAGE=${BASE_IMAGE}" >> ${COMMON_WORKDIR}/versions.txt \
-    && echo "MORI_NIC_BACKEND=${NIC_BACKEND}" >> ${COMMON_WORKDIR}/versions.txt \
-    && echo "AINIC_VERSION=${AINIC_VERSION}" >> ${COMMON_WORKDIR}/versions.txt
-
-CMD ["/bin/bash"]
-
-FROM final_common AS final_lmcache_false
-
-FROM final_common AS final_lmcache_true
-RUN --mount=type=bind,from=build_lmcache,src=/app/install,target=/lmcache_install \
-    --mount=type=cache,target=/root/.cache/uv \
-    uv pip install --system /lmcache_install/*.whl \
-    && test "$(python3 -c 'import numpy; print(numpy.__version__)')" = "2.3.5" \
-    && uv pip check --system \
-    && python3 -c "import lmcache, lmcache.c_ops; print('lmcache', lmcache.__version__)"
-
-FROM final_lmcache_${INSTALL_LMCACHE} AS final
-
-#Set entrypoint for vllm-openai official images
-FROM final AS vllm-openai
-ENTRYPOINT ["vllm", "serve"]
diff --git a/docker/Dockerfile.rocm_72_base b/docker/Dockerfile.rocm_72_base
deleted file mode 100644
index 22b75d2941ea..000000000000
--- a/docker/Dockerfile.rocm_72_base
+++ /dev/null
@@ -1,479 +0,0 @@
-ARG BASE_IMAGE=rocm/dev-ubuntu-22.04:7.2.3-complete
-ARG TRITON_BRANCH="669b31a" # release/internal/3.8.x as of 09/17
-ARG TRITON_REPO="https://github.com/ROCm/triton.git"
-ARG PYTORCH_BRANCH="733fca1" # release/2.13 as of 09/17
-ARG PYTORCH_REPO="https://github.com/ROCm/pytorch.git"
-ARG PYTORCH_VISION_BRANCH="v0.28.0"
-ARG PYTORCH_VISION_REPO="https://github.com/pytorch/vision.git"
-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.26"
-ARG AITER_REPO="https://github.com/ROCm/aiter.git"
-ARG MORI_BRANCH="v1.2.3.post1"
-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="af70edc91f1dc5d228fd4c076df99e65e49c4521"
-
-# Sccache configuration (only used in release pipeline)
-ARG USE_SCCACHE
-ARG SCCACHE_DOWNLOAD_URL
-ARG SCCACHE_ENDPOINT
-ARG SCCACHE_BUCKET_NAME=vllm-build-sccache
-ARG SCCACHE_REGION_NAME=us-west-2
-ARG SCCACHE_S3_NO_CREDENTIALS=0
-
-
-FROM ${BASE_IMAGE} AS base
-
-ENV PATH=/opt/rocm/llvm/bin:/opt/rocm/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin
-ENV ROCM_PATH=/opt/rocm
-ENV LD_LIBRARY_PATH=/opt/rocm/lib:/usr/local/lib:
-
-ARG PYTORCH_ROCM_ARCH=gfx90a;gfx942;gfx950;gfx1100;gfx1101;gfx1200;gfx1201;gfx1150;gfx1151
-ENV PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH}
-ENV AITER_ROCM_ARCH=gfx942;gfx950
-# Note: Do not set MORI_GPU_ARCHS here, it is automatically inferred at runtime
-
-# Required for RCCL in ROCm7.1
-ENV HSA_NO_SCRATCH_RECLAIM=1
-
-ARG PYTHON_VERSION=3.12
-ENV PYTHON_VERSION=${PYTHON_VERSION}
-
-RUN mkdir -p /app
-WORKDIR /app
-ENV DEBIAN_FRONTEND=noninteractive
-
-# Install Python and other dependencies (plus gcc-13, selected later in build_pytorch).
-RUN apt-get update -y \
-    && apt-get install -y software-properties-common git curl sudo vim less libgfortran5 libopenmpi-dev libpci-dev liblzma-dev pkg-config \
-    && for i in 1 2 3; do \
-        add-apt-repository -y ppa:deadsnakes/ppa && add-apt-repository -y ppa:ubuntu-toolchain-r/test && break || \
-        { echo "Attempt $i failed, retrying in 5s..."; sleep 5; }; \
-    done \
-    && apt-get update -y \
-    && apt-get install -y python${PYTHON_VERSION} python${PYTHON_VERSION}-dev python${PYTHON_VERSION}-venv \
-       python${PYTHON_VERSION}-lib2to3 python-is-python3 gcc-13 g++-13 \
-    && update-alternatives --install /usr/bin/python3 python3 /usr/bin/python${PYTHON_VERSION} 1 \
-    && update-alternatives --set python3 /usr/bin/python${PYTHON_VERSION} \
-    && ln -sf /usr/bin/python${PYTHON_VERSION}-config /usr/bin/python3-config \
-    && curl -sS https://bootstrap.pypa.io/get-pip.py | python${PYTHON_VERSION} \
-    && python3 --version && python3 -m pip --version
-
-RUN pip install -U packaging 'cmake<4' ninja wheel 'setuptools<80' 'pybind11<3.1.0' Cython
-RUN apt-get update && apt-get install -y libjpeg-dev libsox-dev libsox-fmt-all sox && rm -rf /var/lib/apt/lists/*
-
-# Install sccache if USE_SCCACHE is enabled (for release builds)
-ARG USE_SCCACHE
-ARG SCCACHE_DOWNLOAD_URL
-ARG SCCACHE_ENDPOINT
-ARG SCCACHE_BUCKET_NAME
-ARG SCCACHE_REGION_NAME
-ARG SCCACHE_S3_NO_CREDENTIALS
-RUN if [ "$USE_SCCACHE" = "1" ]; then \
-        echo "Installing sccache..." \
-        && SCCACHE_ARCH="x86_64" \
-        && SCCACHE_VERSION="v0.8.1" \
-        && SCCACHE_DL_URL="${SCCACHE_DOWNLOAD_URL:-https://github.com/mozilla/sccache/releases/download/${SCCACHE_VERSION}/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl.tar.gz}" \
-        && curl -L -o /tmp/sccache.tar.gz ${SCCACHE_DL_URL} \
-        && tar -xzf /tmp/sccache.tar.gz -C /tmp \
-        && mv /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl/sccache /usr/bin/sccache \
-        && chmod +x /usr/bin/sccache \
-        && rm -rf /tmp/sccache.tar.gz /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl \
-        && sccache --version; \
-    fi
-
-# Setup sccache for HIP compilation via HIP_CLANG_PATH
-# This creates wrapper scripts in a separate directory and points HIP to use them
-# This avoids modifying the original ROCm binaries which can break detection
-# NOTE: HIP_CLANG_PATH is NOT set as ENV to avoid affecting downstream images (Dockerfile.rocm_72)
-# Instead, each build stage should export HIP_CLANG_PATH=/opt/sccache-wrappers if USE_SCCACHE=1
-RUN if [ "$USE_SCCACHE" = "1" ]; then \
-        echo "Setting up sccache wrappers for HIP compilation..." \
-        && mkdir -p /opt/sccache-wrappers \
-        && printf '#!/bin/bash\nexec sccache /opt/rocm/lib/llvm/bin/clang++ "$@"\n' > /opt/sccache-wrappers/clang++ \
-        && chmod +x /opt/sccache-wrappers/clang++ \
-        && printf '#!/bin/bash\nexec sccache /opt/rocm/lib/llvm/bin/clang "$@"\n' > /opt/sccache-wrappers/clang \
-        && chmod +x /opt/sccache-wrappers/clang \
-        && echo "sccache wrappers created in /opt/sccache-wrappers"; \
-    fi
-
-# Set sccache environment variables only when USE_SCCACHE=1
-# This prevents S3 config from leaking into images when sccache is not used
-ARG USE_SCCACHE
-ENV SCCACHE_BUCKET=${USE_SCCACHE:+${SCCACHE_BUCKET_NAME}}
-ENV SCCACHE_REGION=${USE_SCCACHE:+${SCCACHE_REGION_NAME}}
-ENV SCCACHE_S3_NO_CREDENTIALS=${USE_SCCACHE:+${SCCACHE_S3_NO_CREDENTIALS}}
-ENV SCCACHE_IDLE_TIMEOUT=${USE_SCCACHE:+0}
-
-
-###
-### 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} \
-    && git config --global user.email "you@example.com" && git config --global user.name "Your Name" \
-    && if [ ! -f setup.py ]; then cd python; fi \
-    && MAX_JOBS=16 python3 setup.py bdist_wheel --dist-dir=dist \
-    && mkdir -p /app/install && cp dist/*.whl /app/install
-
-
-###
-### AMD SMI Build
-###
-FROM base AS build_amdsmi
-RUN cd /opt/rocm/share/amd_smi \
-    && pip wheel . --wheel-dir=dist
-RUN mkdir -p /app/install && cp /opt/rocm/share/amd_smi/dist/*.whl /app/install
-
-
-###
-### Pytorch build
-###
-FROM base AS build_pytorch
-ARG PYTORCH_BRANCH
-ARG PYTORCH_REPO
-ARG USE_SCCACHE
-
-RUN apt-get update && apt-get install -y pkg-config liblzma-dev
-# torch 2.12 emits float->__bf16 libcalls (__truncsfbf2) absent from gcc-11 libgcc; build with gcc-13.
-RUN update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-13 130 --slave /usr/bin/g++ g++ /usr/bin/g++-13
-RUN git clone ${PYTORCH_REPO} pytorch
-RUN cd pytorch && git checkout ${PYTORCH_BRANCH}
-RUN cd pytorch \
-    && pip install -r requirements.txt && git submodule update --init --recursive
-RUN cd pytorch && python3 tools/amd_build/build_amd.py \
-    && if [ "$USE_SCCACHE" = "1" ]; then \
-           export HIP_CLANG_PATH=/opt/sccache-wrappers \
-           && export CMAKE_C_COMPILER_LAUNCHER=sccache \
-           && export CMAKE_CXX_COMPILER_LAUNCHER=sccache \
-           && sccache --show-stats; \
-       fi \
-    && CMAKE_PREFIX_PATH=$(python3 -c 'import sys; print(sys.prefix)') python3 setup.py bdist_wheel --dist-dir=dist \
-    && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi
-RUN mkdir -p /app/install && cp /app/pytorch/dist/*.whl /app/install
-
-
-###
-### Shared PyTorch runtime for dependent builds
-###
-FROM base AS build_pytorch_runtime
-RUN update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-13 130 --slave /usr/bin/g++ g++ /usr/bin/g++-13
-RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install,ro \
-    pip install /install/*.whl
-
-
-###
-### TorchVision build
-###
-FROM build_pytorch_runtime AS build_torchvision
-ARG PYTORCH_VISION_BRANCH
-ARG PYTORCH_VISION_REPO
-ARG USE_SCCACHE
-
-RUN git clone ${PYTORCH_VISION_REPO} vision
-RUN cd vision && git checkout ${PYTORCH_VISION_BRANCH} \
-    && if [ "$USE_SCCACHE" = "1" ]; then \
-           export HIP_CLANG_PATH=/opt/sccache-wrappers \
-           && export CMAKE_C_COMPILER_LAUNCHER=sccache \
-           && export CMAKE_CXX_COMPILER_LAUNCHER=sccache; \
-       fi \
-    && python3 setup.py bdist_wheel --dist-dir=dist \
-    && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi
-RUN mkdir -p /app/install && cp /app/vision/dist/*.whl /app/install
-
-
-###
-### TorchAudio build
-###
-FROM build_pytorch_runtime AS build_torchaudio
-ARG PYTORCH_AUDIO_BRANCH
-ARG PYTORCH_AUDIO_REPO
-ARG USE_SCCACHE
-
-RUN git clone ${PYTORCH_AUDIO_REPO} audio
-RUN cd audio && git checkout ${PYTORCH_AUDIO_BRANCH} \
-    && git submodule update --init --recursive
-RUN cd audio && pip install -r requirements.txt
-RUN cd audio \
-    && if [ "$USE_SCCACHE" = "1" ]; then \
-           export HIP_CLANG_PATH=/opt/sccache-wrappers \
-           && export CMAKE_C_COMPILER_LAUNCHER=sccache \
-           && export CMAKE_CXX_COMPILER_LAUNCHER=sccache; \
-       fi \
-    && python3 setup.py bdist_wheel --dist-dir=dist \
-    && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi
-RUN mkdir -p /app/install && cp /app/audio/dist/*.whl /app/install
-
-
-###
-### MORI Build
-###
-FROM base AS build_mori
-ARG MORI_BRANCH
-ARG MORI_REPO
-RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \
-    pip install /install/*.whl
-RUN git clone ${MORI_REPO}
-RUN cd mori \
-    && git checkout ${MORI_BRANCH} \
-    && git submodule update --init --recursive \
-    && python3 setup.py bdist_wheel --dist-dir=dist && ls /app/mori/dist/*.whl
-RUN mkdir -p /app/install && cp /app/mori/dist/*.whl /app/install
-
-
-###
-### FlashAttention Build
-###
-FROM base AS build_fa
-ARG FA_BRANCH
-ARG FA_REPO
-ARG USE_SCCACHE
-RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \
-    pip install /install/*.whl
-RUN git clone ${FA_REPO}
-RUN cd flash-attention \
-    && git checkout ${FA_BRANCH} \
-    && git submodule update --init \
-    && if [ "$USE_SCCACHE" = "1" ]; then \
-           export HIP_CLANG_PATH=/opt/sccache-wrappers \
-           && sccache --show-stats; \
-       fi \
-    && GPU_ARCHS=$(echo ${PYTORCH_ROCM_ARCH} | sed -e 's/;gfx1[0-9]\{3\}//g') python3 setup.py bdist_wheel --dist-dir=dist \
-    && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi
-RUN mkdir -p /app/install && cp /app/flash-attention/dist/*.whl /app/install
-
-
-###
-### AITER Build
-###
-FROM base AS build_aiter
-ARG AITER_BRANCH
-ARG AITER_REPO
-ARG USE_SCCACHE
-RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \
-    pip install /install/*.whl
-RUN git clone --recursive --branch ${AITER_BRANCH} ${AITER_REPO}
-RUN cd aiter \
-    && git submodule update --init --recursive \
-    && pip install -r requirements.txt
-RUN pip install pyyaml && cd aiter \
-    && if [ "$USE_SCCACHE" = "1" ]; then \
-           export HIP_CLANG_PATH=/opt/sccache-wrappers \
-           && sccache --show-stats; \
-       fi \
-    && PREBUILD_KERNELS=1 AITER_USE_SYSTEM_TRITON=1 GPU_ARCHS=${AITER_ROCM_ARCH} python3 setup.py bdist_wheel --dist-dir=dist \
-    && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi \
-    && ls /app/aiter/dist/*.whl
-RUN mkdir -p /app/install && cp /app/aiter/dist/*.whl /app/install
-
-
-###
-### Final Build
-###
-
-# Wheel release stage - 
-# 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_amdsmi,src=/app/install/,target=/install \
-    cp /install/*.whl /app/debs
-RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \
-    cp /install/*.whl /app/debs
-RUN --mount=type=bind,from=build_torchvision,src=/app/install/,target=/install \
-    cp /install/*.whl /app/debs
-RUN --mount=type=bind,from=build_torchaudio,src=/app/install/,target=/install \
-    cp /install/*.whl /app/debs
-RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \
-    cp /install/*.whl /app/debs
-
-# Full debs stage - includes Mori (used by Docker releases)
-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_amdsmi,src=/app/install/,target=/install \
-    cp /install/*.whl /app/debs
-RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \
-    cp /install/*.whl /app/debs
-RUN --mount=type=bind,from=build_torchvision,src=/app/install/,target=/install \
-    cp /install/*.whl /app/debs
-RUN --mount=type=bind,from=build_torchaudio,src=/app/install/,target=/install \
-    cp /install/*.whl /app/debs
-RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \
-    cp /install/*.whl /app/debs
-RUN --mount=type=bind,from=build_mori,src=/app/install/,target=/install \
-    cp /install/*.whl /app/debs
-
-###
-### ROCr + CLR build
-###
-# Pull in updates to ROCr + clr from ROCm 7.2.4 and beyond.
-FROM base AS build_rocm_runtime
-ARG ROCM_SYSTEMS_REPO
-ARG ROCM_RUNTIME_COMMIT
-RUN apt-get update \
-    && apt-get install -y --no-install-recommends g++ \
-       libelf-dev libdrm-dev libnuma-dev libdw-dev xxd \
-    && rm -rf /var/lib/apt/lists/*
-RUN pip install --no-cache-dir CppHeaderParser
-RUN git init -q /src && cd /src \
-    && git remote add origin ${ROCM_SYSTEMS_REPO} \
-    && git config core.sparseCheckout true \
-    && git config remote.origin.promisor true \
-    && git config remote.origin.partialclonefilter blob:none \
-    && git sparse-checkout init --cone \
-    && git sparse-checkout set projects/rocr-runtime projects/clr projects/hip shared cmake \
-    && git fetch --filter=blob:none --depth 1 origin ${ROCM_RUNTIME_COMMIT} \
-    && git checkout -q FETCH_HEAD
-# ROCr's trap_handler and blit_shaders call find_package(Clang/LLVM REQUIRED) only to get
-# IMPORTED executables. The ROCm dev image ships the binaries but no CMake package config, so
-# supply one pointing at them.
-RUN mkdir -p /opt/rocm-cmake-shim \
-    && printf 'if(NOT TARGET clang)\n  add_executable(clang IMPORTED GLOBAL)\n  set_target_properties(clang PROPERTIES IMPORTED_LOCATION "/opt/rocm/llvm/bin/clang")\nendif()\nset(Clang_PACKAGE_VERSION "rocm-image-shim")\n' > /opt/rocm-cmake-shim/ClangConfig.cmake \
-    && printf 'if(NOT TARGET llvm-objcopy)\n  add_executable(llvm-objcopy IMPORTED GLOBAL)\n  set_target_properties(llvm-objcopy PROPERTIES IMPORTED_LOCATION "/opt/rocm/llvm/bin/llvm-objcopy")\nendif()\nset(LLVM_FOUND TRUE)\nset(LLVM_PACKAGE_VERSION "rocm-image-shim")\n' > /opt/rocm-cmake-shim/LLVMConfig.cmake
-# ROCr first, and installed into /opt/rocm, because the clr build below needs its headers
-RUN cd /src/projects/rocr-runtime \
-    && cmake -B build -G Ninja \
-        -DCMAKE_BUILD_TYPE=Release \
-        -DCMAKE_PREFIX_PATH=/opt/rocm \
-        -DCMAKE_INSTALL_PREFIX=/opt/rocm \
-        -DClang_DIR=/opt/rocm-cmake-shim \
-        -DLLVM_DIR=/opt/rocm-cmake-shim \
-    && cmake --build build --parallel 16 \
-    && cmake --install build \
-    && cmake --install build --prefix /rocr-install --strip
-RUN cd /src/projects/clr \
-    && cmake -B build -G Ninja \
-        -DCMAKE_BUILD_TYPE=Release \
-        -DCLR_BUILD_HIP=ON \
-        -DCLR_BUILD_OCL=OFF \
-        -DHIP_COMMON_DIR=/src/projects/hip \
-        -DROCM_PATH=/opt/rocm \
-        -DCMAKE_PREFIX_PATH=/opt/rocm \
-        -DCMAKE_INSTALL_PREFIX=/opt/rocm \
-        -DLLVM_DIR=/opt/rocm-cmake-shim \
-        -DHIP_LLVM_ROOT=/opt/rocm/llvm \
-    && cmake --build build --parallel 16 \
-    && cmake --install build --prefix /clr-install --strip
-RUN mkdir -p /staging/lib \
-    && cp -P /rocr-install/lib/libhsa-runtime64.so* /staging/lib/ \
-    && cp -P /clr-install/lib/libamdhip64.so* /staging/lib/ \
-    && ls -l /staging/lib
-
-# rocprofiler-sdk 1.3.2 fixes torch 2.12 profiler on ROCm 7.2.3
-FROM base AS build_rocprofiler_sdk
-ARG ROCM_SYSTEMS_REPO
-ARG ROCPROFILER_SDK_COMMIT
-ARG USE_SCCACHE
-RUN apt-get update \
-    && apt-get install -y --no-install-recommends libdw-dev libelf-dev libsqlite3-dev libzstd-dev \
-    && rm -rf /var/lib/apt/lists/*
-RUN git init -q /src && cd /src \
-    && git remote add origin ${ROCM_SYSTEMS_REPO} \
-    && git config core.sparseCheckout true \
-    && git config remote.origin.promisor true \
-    && git config remote.origin.partialclonefilter blob:none \
-    && git sparse-checkout init --cone \
-    && git sparse-checkout set projects/rocprofiler-sdk shared cmake \
-    && git fetch --filter=blob:none --depth 1 origin ${ROCPROFILER_SDK_COMMIT} \
-    && git checkout -q FETCH_HEAD
-# Patch the ROCm 7.2.3 build of rocprofiler-sdk to include queue-interposition fixes that are
-# absent from the pinned 7.14 build. This resolves hangs on torch.profiler on ROCm.
-# TODO: Drop these once ${ROCPROFILER_SDK_COMMIT} is bumped to a build that already 
-# contains them, such as ROCm 10.0. See https://github.com/vllm-project/vllm/issues/54087.
-ARG ROCPROFILER_SDK_PR_7796=fae7fd57ffb8a76e811d6410b8a0cee1e94d710a
-ARG ROCPROFILER_SDK_PR_7924=f8db22adbe17686a9b05ae03367dd1ec73011653
-RUN cd /src \
-    && HSA_PATHS='projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/hsa/*' \
-    && curl -fsSL -o /tmp/qi-7796.patch "${ROCM_SYSTEMS_REPO%.git}/commit/${ROCPROFILER_SDK_PR_7796}.patch" \
-    && curl -fsSL -o /tmp/qi-7924.patch "${ROCM_SYSTEMS_REPO%.git}/commit/${ROCPROFILER_SDK_PR_7924}.patch" \
-    && git apply -p1 --include="$HSA_PATHS" --verbose /tmp/qi-7796.patch \
-    && git apply -p1 --include="$HSA_PATHS" --verbose /tmp/qi-7924.patch
-RUN cd /src/projects/rocprofiler-sdk \
-    && if [ "$USE_SCCACHE" = "1" ]; then \
-           export CMAKE_C_COMPILER_LAUNCHER=sccache \
-           && export CMAKE_CXX_COMPILER_LAUNCHER=sccache; \
-       fi \
-    && cmake -B build -G Ninja \
-        -DCMAKE_BUILD_TYPE=Release \
-        -DCMAKE_PREFIX_PATH=/opt/rocm \
-        -DCMAKE_INSTALL_PREFIX=/opt/rocm \
-        -DROCPROFILER_BUILD_TESTS=OFF \
-        -DROCPROFILER_BUILD_SAMPLES=OFF \
-        -DROCPROFILER_BUILD_DOCS=OFF \
-    && cmake --build build --parallel "$(nproc)" \
-    && cmake --install build --prefix /staging --strip
-
-FROM base AS final
-RUN --mount=type=bind,from=debs,src=/app/debs,target=/install \
-    pip install /install/*.whl
-
-COPY --from=build_rocprofiler_sdk /staging/ /opt/rocm/
-ENV ROCPROFILER_QUEUE_INTERPOSITION=0
-
-COPY --from=build_rocm_runtime /staging/ /opt/rocm/
-# Hard link the patched library onto every stock file name so that both the
-# SONAME lookup and the hard coded paths resolve to it.
-RUN set -eux; \
-    cd /opt/rocm/lib; \
-    ours_hsa="$(readlink -f libhsa-runtime64.so.1)"; \
-    ours_hip="$(readlink -f libamdhip64.so.7)"; \
-    for f in libhsa-runtime64.so.1.*; do \
-        [ "$f" = "${ours_hsa##*/}" ] || ln -f "$ours_hsa" "$f"; \
-    done; \
-    for f in libamdhip64.so.7.*; do \
-        [ "$f" = "${ours_hip##*/}" ] || ln -f "$ours_hip" "$f"; \
-    done
-
-RUN ldconfig
-
-ARG BASE_IMAGE
-ARG TRITON_BRANCH
-ARG TRITON_REPO
-ARG PYTORCH_BRANCH
-ARG PYTORCH_VISION_BRANCH
-ARG PYTORCH_REPO
-ARG PYTORCH_VISION_REPO
-ARG PYTORCH_AUDIO_BRANCH
-ARG PYTORCH_AUDIO_REPO
-ARG FA_BRANCH
-ARG FA_REPO
-ARG AITER_BRANCH
-ARG AITER_REPO
-ARG MORI_BRANCH
-ARG MORI_REPO
-ARG ROCM_SYSTEMS_REPO
-ARG ROCPROFILER_SDK_COMMIT
-ARG ROCM_RUNTIME_COMMIT
-RUN echo "BASE_IMAGE: ${BASE_IMAGE}" > /app/versions.txt \
-    && echo "ROCM_SYSTEMS_REPO: ${ROCM_SYSTEMS_REPO}" >> /app/versions.txt \
-    && echo "ROCPROFILER_SDK_COMMIT: ${ROCPROFILER_SDK_COMMIT}" >> /app/versions.txt \
-    && echo "ROCM_RUNTIME_COMMIT: ${ROCM_RUNTIME_COMMIT}" >> /app/versions.txt \
-    && echo "TRITON_BRANCH: ${TRITON_BRANCH}" >> /app/versions.txt \
-    && echo "TRITON_REPO: ${TRITON_REPO}" >> /app/versions.txt \
-    && echo "PYTORCH_BRANCH: ${PYTORCH_BRANCH}" >> /app/versions.txt \
-    && echo "PYTORCH_VISION_BRANCH: ${PYTORCH_VISION_BRANCH}" >> /app/versions.txt \
-    && echo "PYTORCH_REPO: ${PYTORCH_REPO}" >> /app/versions.txt \
-    && echo "PYTORCH_VISION_REPO: ${PYTORCH_VISION_REPO}" >> /app/versions.txt \
-    && echo "PYTORCH_AUDIO_BRANCH: ${PYTORCH_AUDIO_BRANCH}" >> /app/versions.txt \
-    && echo "PYTORCH_AUDIO_REPO: ${PYTORCH_AUDIO_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 \
-    && echo "AITER_REPO: ${AITER_REPO}" >> /app/versions.txt \
-    && echo "MORI_BRANCH: ${MORI_BRANCH}" >> /app/versions.txt \
-    && echo "MORI_REPO: ${MORI_REPO}" >> /app/versions.txt
diff --git a/docker/Dockerfile.rocm_base b/docker/Dockerfile.rocm_base
index 997af6183805..9561b4a8e24c 100644
--- a/docker/Dockerfile.rocm_base
+++ b/docker/Dockerfile.rocm_base
@@ -1,4 +1,6 @@
-ARG BASE_IMAGE=ubuntu:22.04
+ARG BASE_IMAGE=ubuntu:22.04@sha256:5ec03bb3441e8b0bf3b4f9cd4629a1ae763010dc3035bb8da3ae6cf026486401
+# Bump to rebuild the stack (changes its key and busts every layer); upstream drift alone never does.
+ARG REFRESH=2026-10-12
 ARG PYTHON_VERSION=3.12
 ARG TORCH_VERSION=2.13.0+rocm10.0.0
 ARG TORCHVISION_VERSION=0.28.0+rocm10.0.0
@@ -8,7 +10,7 @@ ARG TRITON_VERSION=3.8.0+git669b31ac
 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.26"
+ARG AITER_BRANCH="v0.1.27"
 ARG AITER_REPO="https://github.com/ROCm/aiter.git"
 ARG MORI_BRANCH="v1.2.3.post1"
 ARG MORI_REPO="https://github.com/ROCm/mori.git"
@@ -26,14 +28,6 @@ ARG ROCM_SDK_VERSION=10.0.0
 ARG PYTORCH_ROCM_ARCH="gfx90a;gfx942;gfx950;gfx1100;gfx1101;gfx1200;gfx1201;gfx1150;gfx1151"
 ARG AITER_ROCM_ARCH="gfx942;gfx950"
 
-# Sccache configuration (only used in release pipeline)
-ARG USE_SCCACHE
-ARG SCCACHE_DOWNLOAD_URL
-ARG SCCACHE_ENDPOINT
-ARG SCCACHE_BUCKET_NAME=vllm-build-sccache
-ARG SCCACHE_REGION_NAME=us-west-2
-ARG SCCACHE_S3_NO_CREDENTIALS=0
-
 FROM ${BASE_IMAGE} AS base
 
 ARG PYTORCH_ROCM_ARCH
@@ -52,6 +46,7 @@ ENV ROCM_PATH="$SDK_DEV" \
     HIP_DEVICE_LIB_PATH="$SDK_CORE/lib/llvm/amdgcn/bitcode"
 ENV LD_LIBRARY_PATH="$SDK_DEV/lib:$SDK_DEV/lib/rocm_sysdeps/lib:$SDK_CORE/lib:$SDK_CORE/lib/rocm_sysdeps/lib"
 ENV PATH="$SDK_DEV/bin:$SDK_DEV/lib/llvm/bin:$SDK_CORE/bin:$PATH"
+ENV ROCSHMEM_DIR="$SDK_DEV"
 
 # Required for RCCL in ROCm7.1
 ENV HSA_NO_SCRATCH_RECLAIM=1
@@ -59,12 +54,16 @@ ENV HSA_NO_SCRATCH_RECLAIM=1
 ARG PYTHON_VERSION
 ENV PYTHON_VERSION=${PYTHON_VERSION}
 
-RUN mkdir -p /app
+ARG REFRESH
+RUN echo "$REFRESH" >/etc/vllm-rocm-refresh && mkdir -p /app
 WORKDIR /app
 ENV DEBIAN_FRONTEND=noninteractive
 
 # Install Python and other dependencies
-RUN apt-get update -y \
+# pip runs with a cache mount, so its downloads (several GB of SDK
+# and torch wheels) never land in an image layer.
+RUN --mount=type=cache,target=/root/.cache/pip \
+    apt-get update -y \
     && apt-get install -y software-properties-common build-essential git curl sudo vim less libgfortran5 libopenmpi-dev libpci-dev liblzma-dev pkg-config libdrm-dev libnuma-dev kmod \
     && for i in 1 2 3; do \
         add-apt-repository -y ppa:deadsnakes/ppa && break || \
@@ -79,14 +78,16 @@ RUN apt-get update -y \
     && curl -sS https://bootstrap.pypa.io/get-pip.py | python${PYTHON_VERSION} \
     && python3 --version && python3 -m pip --version
 
-RUN pip install -U packaging 'cmake<4' ninja wheel 'setuptools<80' pybind11==3.0.4 Cython
+RUN --mount=type=cache,target=/root/.cache/pip \
+    pip install -U packaging 'cmake<4' ninja wheel==0.48.0 'setuptools<80' pybind11==3.0.4 Cython
 RUN apt-get update && apt-get install -y libjpeg-dev libsox-dev libsox-fmt-all sox && rm -rf /var/lib/apt/lists/*
 
 # Install ROCm SDK via TheRock
 ARG ROCM_RELEASE_WHEELS_MULTIARCH_URL
 ARG ROCM_NIGHTLY_WHEELS_MULTIARCH_URL
 ARG ROCM_SDK_VERSION
-RUN pip install --index-url ${ROCM_RELEASE_WHEELS_MULTIARCH_URL} \
+RUN --mount=type=cache,target=/root/.cache/pip \
+    pip install --index-url ${ROCM_RELEASE_WHEELS_MULTIARCH_URL} \
     --extra-index-url ${ROCM_NIGHTLY_WHEELS_MULTIARCH_URL} \
     "rocm[libraries,devel]${ROCM_SDK_VERSION:+==${ROCM_SDK_VERSION}}" \
     && rocm-sdk init
@@ -95,7 +96,8 @@ RUN pip install --index-url ${ROCM_RELEASE_WHEELS_MULTIARCH_URL} \
 ARG TORCH_VERSION
 ARG TORCHAUDIO_VERSION
 ARG TORCHVISION_VERSION
-RUN pip install \
+RUN --mount=type=cache,target=/root/.cache/pip \
+    pip install \
     --index-url ${ROCM_RELEASE_WHEELS_MULTIARCH_URL} \
     --extra-index-url ${ROCM_NIGHTLY_WHEELS_MULTIARCH_URL} \
     "torch==${TORCH_VERSION}" \
@@ -112,7 +114,8 @@ RUN pip install \
 # For each arch in PYTORCH_ROCM_ARCH, install the precompiled GPU code that
 # PyTorch needs to run on that specific chip.
 SHELL ["/bin/bash", "-c"]
-RUN TORCH_VER=$(python3 -c "import torch; print(torch.__version__)") \
+RUN --mount=type=cache,target=/root/.cache/pip \
+    TORCH_VER=$(python3 -c "import torch; print(torch.__version__)") \
     && for arch in $(echo "${PYTORCH_ROCM_ARCH}" | tr ';' ' '); do \
         echo "Installing device packages for ${arch}" \
         && pip install --index-url ${ROCM_RELEASE_WHEELS_MULTIARCH_URL} \
@@ -126,7 +129,8 @@ RUN TORCH_VER=$(python3 -c "import torch; print(torch.__version__)") \
 RUN rm -f "${SITE_PACKAGES}/torch/.kpack/torch_gfx90a:xnack+.kpack" \
          "${SITE_PACKAGES}/torch/.kpack/torch_gfx90a:xnack-.kpack"
 
-RUN cd $SITE_PACKAGES/_rocm_sdk_core/share/amd_smi && pip install .
+RUN --mount=type=cache,target=/root/.cache/pip \
+    cd $SITE_PACKAGES/_rocm_sdk_core/share/amd_smi && pip install .
 ENV PYTHONPATH=$SITE_PACKAGES/_rocm_sdk_core/share/amd_smi
 
 # TheRock runtime fixes
@@ -156,61 +160,28 @@ RUN if [ -f "${SITE_PACKAGES}/rocm_sdk/__init__.py" ]; then \
     done \
     && ldconfig
 
-# Install sccache if USE_SCCACHE is enabled (for release builds)
-ARG USE_SCCACHE
-ARG SCCACHE_DOWNLOAD_URL
-ARG SCCACHE_ENDPOINT
-ARG SCCACHE_BUCKET_NAME
-ARG SCCACHE_REGION_NAME
-ARG SCCACHE_S3_NO_CREDENTIALS
-RUN if [ "$USE_SCCACHE" = "1" ]; then \
-        echo "Installing sccache..." \
-        && SCCACHE_ARCH="x86_64" \
-        && SCCACHE_VERSION="v0.8.1" \
-        && SCCACHE_DL_URL="${SCCACHE_DOWNLOAD_URL:-https://github.com/mozilla/sccache/releases/download/${SCCACHE_VERSION}/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl.tar.gz}" \
-        && curl -L -o /tmp/sccache.tar.gz ${SCCACHE_DL_URL} \
-        && tar -xzf /tmp/sccache.tar.gz -C /tmp \
-        && mv /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl/sccache /usr/bin/sccache \
-        && chmod +x /usr/bin/sccache \
-        && rm -rf /tmp/sccache.tar.gz /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl \
-        && sccache --version; \
-    fi
-
-# Setup sccache for HIP compilation via HIP_CLANG_PATH
-# This creates wrapper scripts in a separate directory and points HIP to use them
-# This avoids modifying the original ROCm binaries which can break detection
-# NOTE: HIP_CLANG_PATH is NOT set as ENV to avoid affecting downstream images (Dockerfile.rocm)
-# Instead, each build stage should export HIP_CLANG_PATH=/opt/sccache-wrappers if USE_SCCACHE=1
-RUN if [ "$USE_SCCACHE" = "1" ]; then \
-        echo "Setting up sccache wrappers for HIP compilation..." \
-        && mkdir -p /opt/sccache-wrappers \
-        && printf '#!/bin/bash\nexec sccache %s/lib/llvm/bin/clang++ "$@"\n' "${ROCM_PATH}" > /opt/sccache-wrappers/clang++ \
-        && chmod +x /opt/sccache-wrappers/clang++ \
-        && printf '#!/bin/bash\nexec sccache %s/lib/llvm/bin/clang "$@"\n' "${ROCM_PATH}" > /opt/sccache-wrappers/clang \
-        && chmod +x /opt/sccache-wrappers/clang \
-        && echo "sccache wrappers created in /opt/sccache-wrappers"; \
-    fi
-
-ENV SCCACHE_BUCKET=${USE_SCCACHE:+${SCCACHE_BUCKET_NAME}}
-ENV SCCACHE_REGION=${USE_SCCACHE:+${SCCACHE_REGION_NAME}}
-ENV SCCACHE_S3_NO_CREDENTIALS=${USE_SCCACHE:+${SCCACHE_S3_NO_CREDENTIALS}}
-ENV SCCACHE_IDLE_TIMEOUT=${USE_SCCACHE:+0}
-
 
 ###
 ### MORI Build
 ###
 FROM base AS build_mori
+# mori silently builds without UMBP when gRPC or Protobuf is missing.
+RUN apt-get update \
+    && apt-get install -y --no-install-recommends libgrpc++-dev libprotobuf-dev \
+       protobuf-compiler protobuf-compiler-grpc \
+    && rm -rf /var/lib/apt/lists/*
 ARG MORI_BRANCH
 ARG MORI_REPO
 RUN git clone ${MORI_REPO}
-RUN cd mori \
+RUN --mount=type=cache,target=/root/.cache/pip cd mori \
     && git checkout ${MORI_BRANCH} \
     && git submodule update --init --recursive \
     && ln -sfn ${ROCM_PATH} /opt/rocm \
     && pip install "setuptools_scm>=6.2" \
     && CMAKE_PREFIX_PATH="${ROCM_PATH}:${ROCM_PATH}/lib/rocm_sysdeps" \
-       python3 setup.py bdist_wheel --dist-dir=dist && ls /app/mori/dist/*.whl
+       python3 setup.py bdist_wheel --dist-dir=dist \
+    && grep -qx 'BUILD_UMBP:BOOL=ON' build/CMakeCache.txt \
+    && ls /app/mori/dist/*.whl
 RUN mkdir -p /app/install && cp /app/mori/dist/*.whl /app/install
 
 ###
@@ -219,17 +190,11 @@ RUN mkdir -p /app/install && cp /app/mori/dist/*.whl /app/install
 FROM base AS build_fa
 ARG FA_BRANCH
 ARG FA_REPO
-ARG USE_SCCACHE
 RUN git clone ${FA_REPO}
 RUN cd flash-attention \
     && git checkout ${FA_BRANCH} \
     && git submodule update --init \
-    && if [ "$USE_SCCACHE" = "1" ]; then \
-           export HIP_CLANG_PATH=/opt/sccache-wrappers \
-           && sccache --show-stats; \
-       fi \
-    && GPU_ARCHS=$(echo ${PYTORCH_ROCM_ARCH} | sed -e 's/;gfx1[0-9]\{3\}//g') python3 setup.py bdist_wheel --dist-dir=dist \
-    && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi
+    && GPU_ARCHS=$(echo ${PYTORCH_ROCM_ARCH} | sed -e 's/;gfx1[0-9]\{3\}//g') python3 setup.py bdist_wheel --dist-dir=dist
 RUN mkdir -p /app/install && cp /app/flash-attention/dist/*.whl /app/install
 
 
@@ -239,20 +204,14 @@ RUN mkdir -p /app/install && cp /app/flash-attention/dist/*.whl /app/install
 FROM base AS build_aiter
 ARG AITER_BRANCH
 ARG AITER_REPO
-ARG USE_SCCACHE
 RUN git clone --recursive --branch ${AITER_BRANCH} ${AITER_REPO}
-RUN cd aiter \
+RUN --mount=type=cache,target=/root/.cache/pip cd aiter \
     && git submodule update --init --recursive \
     && pip install -r requirements.txt
 # AITER_USE_SYSTEM_TRITON keeps setup.py from resolving its own triton, which would
 # displace the pinned one already installed in base.
-RUN pip install pyyaml && cd aiter \
-    && if [ "$USE_SCCACHE" = "1" ]; then \
-           export HIP_CLANG_PATH=/opt/sccache-wrappers \
-           && sccache --show-stats; \
-       fi \
+RUN --mount=type=cache,target=/root/.cache/pip pip install pyyaml && cd aiter \
     && PREBUILD_KERNELS=1 AITER_USE_SYSTEM_TRITON=1 GPU_ARCHS=${AITER_ROCM_ARCH} python3 setup.py bdist_wheel --dist-dir=dist \
-    && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi \
     && ls /app/aiter/dist/*.whl
 RUN mkdir -p /app/install && cp /app/aiter/dist/*.whl /app/install
 
@@ -268,7 +227,10 @@ 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
+    && python3 -m wheel unpack dist/*.whl -d /tmp/triton-wheel \
+    && find /tmp/triton-wheel -name '*.so' -exec strip --strip-debug {} + \
+    && mkdir -p /app/install \
+    && python3 -m wheel pack /tmp/triton-wheel/* -d /app/install
 
 
 ###
@@ -330,9 +292,8 @@ RUN mkdir -p /staging/lib \
 ### Final Build
 ###
 
-# Wheel release stage -
-# only includes dependencies used by wheel release pipeline
-FROM base AS debs_wheel_release
+# The stack's half of the release wheel index.
+FROM base AS wheels-build
 RUN mkdir /app/debs
 RUN --mount=type=bind,from=build_triton,src=/app/install/,target=/install \
     cp /install/*.whl /app/debs
@@ -348,7 +309,8 @@ ARG ROCM_RELEASE_WHEELS_MULTIARCH_URL
 ARG ROCM_NIGHTLY_WHEELS_MULTIARCH_URL
 COPY tools/vllm-rocm/therock_wheels.py /tmp/therock_wheels.py
 # regex and pybase64 aren't installed here; fall back to the stdlib modules.
-RUN sed -i -e 's/import regex as re/import re/' \
+RUN --mount=type=cache,target=/root/.cache/pip \
+    sed -i -e 's/import regex as re/import re/' \
         -e 's/import pybase64 as base64/import base64/' /tmp/therock_wheels.py \
     && pip download --no-deps --dest /tmp/torch-dl \
         --index-url ${ROCM_RELEASE_WHEELS_MULTIARCH_URL} \
@@ -363,19 +325,12 @@ RUN sed -i -e 's/import regex as re/import re/' \
     && rm -rf /tmp/torch-dl /tmp/amd_smi /tmp/therock_wheels.py \
     && ls -l /app/debs
 
-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 \
-    cp /install/*.whl /app/debs
-RUN --mount=type=bind,from=build_mori,src=/app/install/,target=/install \
-    cp /install/*.whl /app/debs
+FROM scratch AS wheels
+COPY --from=wheels-build /app/debs /
 
 FROM base AS final
 ARG TRITON_VERSION
+LABEL ai.vllm.rocm.variant=rocm100
 
 # 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.
@@ -409,8 +364,17 @@ RUN TORCH_DIST_INFO=$(ls -d ${SITE_PACKAGES}/torch-*.dist-info) \
     && 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 \
+# Runtime libraries of mori's UMBP extension.
+RUN apt-get update \
+    && apt-get install -y --no-install-recommends libgrpc++1 libprotobuf23 \
+    && rm -rf /var/lib/apt/lists/*
+
+RUN --mount=type=bind,from=build_triton,src=/app/install,target=/i/triton \
+    --mount=type=bind,from=build_fa,src=/app/install,target=/i/fa \
+    --mount=type=bind,from=build_aiter,src=/app/install,target=/i/aiter \
+    --mount=type=bind,from=build_mori,src=/app/install,target=/i/mori \
+    --mount=type=cache,target=/root/.cache/pip \
+    pip install /i/*/*.whl \
     && ! pip check 2>&1 | grep -iE '^(torch|triton) ' \
     && test "$(pip show triton | awk '/^Version:/ {print $2}')" = "${TRITON_VERSION}" \
     # Create a constraints file to protect the versions we need, especially
@@ -422,8 +386,16 @@ RUN --mount=type=bind,from=debs,src=/app/debs,target=/install \
     && grep -q '^pybind11==3\.0\.' /etc/rocm-constraints.txt \
     && python3 -c "import torch; assert torch.version.hip, f'non-ROCm torch: {torch.__version__}'; print(f'torch={torch.__version__}')" \
     && python3 -c "import triton; print(f'triton={triton.__version__}')" \
+    && python3 -c "from mori.cpp import UMBPClient" \
     && pip list 2>/dev/null | grep -iE '^(torch|triton) '
 
+# Fail if a pip call above ran without the cache mount and left its downloads in the image.
+RUN if [ -n "$(ls -A /root/.cache/pip 2>/dev/null)" ]; then \
+        du -sh /root/.cache/pip; \
+        echo "ERROR: /root/.cache/pip is not empty; add --mount=type=cache,target=/root/.cache/pip" >&2; \
+        exit 1; \
+    fi
+
 ARG BASE_IMAGE
 ARG PYTHON_VERSION
 ARG ROCM_SDK_VERSION
@@ -443,7 +415,6 @@ 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 \
     && echo "ROCM_SDK_VERSION: ${ROCM_SDK_VERSION}" >> /app/versions.txt \
@@ -462,15 +433,5 @@ RUN echo "BASE_IMAGE: ${BASE_IMAGE}" > /app/versions.txt \
     && echo "MORI_REPO: ${MORI_REPO}" >> /app/versions.txt \
     && echo "ROCM_SYSTEMS_REPO: ${ROCM_SYSTEMS_REPO}" >> /app/versions.txt \
     && echo "ROCM_RUNTIME_COMMIT: ${ROCM_RUNTIME_COMMIT}" >> /app/versions.txt \
-    # Fingerprint our triton so Dockerfile.rocm can prove it survived the build.
     && echo "TRITON_VERSION: $(pip show triton | awk '/^Version:/ {print $2}')" >> /app/versions.txt \
-    && echo "TRITON_SHA256: $(find $(python3 -c 'import triton, os; print(os.path.dirname(triton.__file__))') \
-        -type f \( -name '*.py' -o -name '*.so' \) | sort | xargs sha256sum | sha256sum | cut -d' ' -f1)" \
-        >> /app/versions.txt \
-    && echo "PYBIND11_VERSION: $(pip show pybind11 | awk '/^Version:/ {print $2}')" >> /app/versions.txt \
-    && echo "PYBIND11_INTERNALS_VERSION: $(grep -rhE '^#[[:space:]]*define PYBIND11_INTERNALS_VERSION' \
-        $(python3 -c 'import pybind11; print(pybind11.get_include())') | awk '{print $NF}')" \
-        >> /app/versions.txt \
-    && echo "PYBIND11_SHA256: $(find $(python3 -c 'import pybind11; print(pybind11.get_include())') \
-        -type f -name '*.h' | sort | xargs sha256sum | sha256sum | cut -d' ' -f1)" \
-        >> /app/versions.txt
+    && echo "PYBIND11_VERSION: $(pip show pybind11 | awk '/^Version:/ {print $2}')" >> /app/versions.txt
diff --git a/docker/Dockerfile.rocm_gfx1250 b/docker/Dockerfile.rocm_gfx1250
index 634e522e0d62..dd22d0c27685 100644
--- a/docker/Dockerfile.rocm_gfx1250
+++ b/docker/Dockerfile.rocm_gfx1250
@@ -2,7 +2,7 @@
 ARG REMOTE_VLLM="0"
 ARG COMMON_WORKDIR=/app
 ARG BASE_IMAGE=rocm/vllm-dev:base
-ARG CI_BASE_IMAGE=rocm/vllm-dev:ci_base
+ARG CI_BASE_IMAGE=ci_base
 # NIC backend for MoRI RDMA support.
 # By default (all), drivers and userspace libraries for all supported NIC types
 # (ainic and bnxt) are installed; MoRI selects the appropriate one at runtime.
diff --git a/docker/Dockerfile.xpu b/docker/Dockerfile.xpu
index 5e1e2984f4a4..fbde67b4e6d3 100644
--- a/docker/Dockerfile.xpu
+++ b/docker/Dockerfile.xpu
@@ -171,7 +171,7 @@ FROM vllm-base AS ucx-nixl-build
 # NIXL_VERSION (nixl contrib/Dockerfile.manylinux: ARG UCX_REF), so XPU runs the
 # same UCX commit as the released nixl-cu* wheels rather than an untested pairing.
 ARG UCX_VERSION=v1.23.x
-ARG NIXL_VERSION=v1.5.0
+ARG NIXL_VERSION=v1.5.2
 
 # Build-time only: compiler, autotools, and verbs dev headers
 RUN apt-get update -y && apt-get install -y --no-install-recommends \
diff --git a/docker/ci-rocm.hcl b/docker/ci-rocm.hcl
deleted file mode 100644
index 31c8ad6405c5..000000000000
--- a/docker/ci-rocm.hcl
+++ /dev/null
@@ -1,424 +0,0 @@
-# ci-rocm.hcl - CI-specific configuration for vLLM ROCm Docker builds
-#
-# This file lives in the vLLM repo at docker/ci-rocm.hcl so ROCm Docker
-# build mechanics can evolve with Dockerfile.rocm and docker-bake-rocm.hcl.
-# Used with: docker buildx bake -f docker/docker-bake-rocm.hcl -f docker/ci-rocm.hcl test-rocm-ci
-#
-# Registry cache: Docker Hub (rocm/vllm-ci-cache) is used exclusively.
-# AMD build agents already have Docker Hub credentials (they push the test
-# image to rocm/vllm-ci), so no additional credential setup is required.
-# ROCm CI uses Docker Hub for BuildKit layer cache by default. A separate
-# compiler cache can be enabled with USE_SCCACHE=1 when AMD provides a shared
-# S3-compatible cache endpoint.
-
-# CI metadata
-
-variable "BUILDKITE_COMMIT" {
-  default = ""
-}
-
-variable "BUILDKITE_BUILD_NUMBER" {
-  default = ""
-}
-
-variable "BUILDKITE_BUILD_ID" {
-  default = ""
-}
-
-variable "VLLM_KERNEL_SYMBOL_MAP" {
-  default = "0"
-}
-
-variable "PARENT_COMMIT" {
-  default = ""
-}
-
-# Merge-base of HEAD with main - provides a more stable cache fallback than
-# parent commit for long-lived PRs. Mirrors the VLLM_MERGE_BASE_COMMIT
-# pattern used in the shared ci.hcl file. Auto-computed by ci-bake-rocm.sh
-# when unset.
-variable "VLLM_MERGE_BASE_COMMIT" {
-  default = ""
-}
-
-# Bridge to vLLM's COMMIT variable for OCI labels
-variable "COMMIT" {
-  default = BUILDKITE_COMMIT
-}
-
-# Image tags (set by CI)
-
-variable "IMAGE_TAG" {
-  default = ""
-}
-
-variable "IMAGE_TAG_LATEST" {
-  default = ""
-}
-
-# ROCm-specific GPU architecture targets
-
-variable "PYTORCH_ROCM_ARCH" {
-  default = "gfx90a;gfx942;gfx950"
-}
-
-# Pre-built CI base image (Tier 1). Per-PR builds pull this instead of
-# rebuilding NIXL/DeepEP/torchcodec from scratch. The ci_base stage in
-# Dockerfile.rocm inherits from base, so CI_BASE_IMAGE only affects the test
-# stage and is irrelevant when building --target ci_base itself.
-variable "CI_BASE_IMAGE" {
-  default = "rocm/vllm-dev:ci_base"
-}
-
-# Leave CI_MAX_JOBS empty so the Dockerfile falls back to $(nproc) and uses
-# the full builder parallelism. Operators can still override this per build.
-variable "CI_MAX_JOBS" {
-  default = ""
-}
-
-# Upstream dependency commit pins -- extracted from Dockerfile.rocm by
-# ci-bake-rocm.sh at build time. Empty defaults are safe: the cache
-# functions produce no entries when the variable is empty.
-variable "NIXL_BRANCH" {
-  default = ""
-}
-
-variable "UCX_BRANCH" {
-  default = ""
-}
-
-variable "DEEPEP_BRANCH" {
-  default = ""
-}
-
-variable "NIXL_CACHE_KEY" {
-  default = ""
-}
-
-variable "DEEPEP_CACHE_KEY" {
-  default = ""
-}
-
-# Docker Hub registry cache for AMD builds.
-#
-# A separate repo (rocm/vllm-ci-cache) is used for BuildKit layer cache.
-# Final-image cache exports use mode=min to reduce the volume of data pushed.
-# Source-scoped csrc cache exports default to mode=max so fresh workers can
-# recover more of the native build graph when ROCm extension inputs change.
-# NOTE: mode=min still includes all layers referenced by the final image
-# manifest, including inherited base layers (~7.25GB ROCm runtime).
-# Docker Hub auto-creates the repo on first push.
-#
-# Final-image cache stays commit-scoped. Branch-to-branch reuse for the test
-# image comes from importing the parent and merge-base commit cache refs.
-#
-# The source-scoped native cache is exported both per-commit and per-branch so
-# ROCm extension rebuilds are shareable within the same commit reruns and across
-# consecutive commits on the same branch without depending on a single global
-# latest tag.
-
-variable "DOCKERHUB_CACHE_REPO" {
-  default = "rocm/vllm-ci-cache"
-}
-
-variable "DOCKERHUB_CACHE_TO" {
-  default = ""
-}
-
-variable "ROCM_CACHE_BRANCH_TAG" {
-  default = ""
-}
-
-variable "ROCM_CACHE_UPSTREAM_BRANCH_TAG" {
-  default = ""
-}
-
-variable "ROCM_CSRC_CACHE_TO_MODE" {
-  default = "max"
-}
-
-variable "ROCM_RUST_CACHE_TO_MODE" {
-  default = "max"
-}
-
-variable "ROCM_FINAL_CACHE_TO_MODE" {
-  default = "min"
-}
-
-# Functions
-
-function "get_cache_from_rocm" {
-  params = []
-  result = compact([
-    # Exact commit hit - fastest cache on re-runs of the same commit
-    BUILDKITE_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocm-${BUILDKITE_COMMIT}" : "",
-    # Parent commit - useful cache for incremental changes
-    PARENT_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocm-${PARENT_COMMIT}" : "",
-    # Merge-base with main - stable fallback for long-lived or rebased PRs;
-    # maps to a real main-branch commit whose cache layers are likely warm
-    VLLM_MERGE_BASE_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocm-${VLLM_MERGE_BASE_COMMIT}" : "",
-    # Import the source-scoped native build cache as well so builds whose
-    # Python/package layers changed can still reuse compiled ROCm objects.
-    BUILDKITE_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:csrc-rocm-${BUILDKITE_COMMIT}" : "",
-    PARENT_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:csrc-rocm-${PARENT_COMMIT}" : "",
-    VLLM_MERGE_BASE_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:csrc-rocm-${VLLM_MERGE_BASE_COMMIT}" : "",
-    ROCM_CACHE_BRANCH_TAG != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:csrc-rocm-branch-${ROCM_CACHE_BRANCH_TAG}" : "",
-    ROCM_CACHE_UPSTREAM_BRANCH_TAG != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:csrc-rocm-branch-${ROCM_CACHE_UPSTREAM_BRANCH_TAG}" : "",
-    # Import the source-scoped Rust frontend cache so non-Rust changes do not
-    # force a fresh cargo release build.
-    BUILDKITE_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rust-rocm-${BUILDKITE_COMMIT}" : "",
-    PARENT_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rust-rocm-${PARENT_COMMIT}" : "",
-    VLLM_MERGE_BASE_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rust-rocm-${VLLM_MERGE_BASE_COMMIT}" : "",
-    ROCM_CACHE_BRANCH_TAG != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rust-rocm-branch-${ROCM_CACHE_BRANCH_TAG}" : "",
-    ROCM_CACHE_UPSTREAM_BRANCH_TAG != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rust-rocm-branch-${ROCM_CACHE_UPSTREAM_BRANCH_TAG}" : "",
-    # Branch-scoped full image cache - fallback when parent-commit cache is evicted
-    ROCM_CACHE_BRANCH_TAG != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocm-branch-${ROCM_CACHE_BRANCH_TAG}" : "",
-    ROCM_CACHE_UPSTREAM_BRANCH_TAG != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocm-branch-${ROCM_CACHE_UPSTREAM_BRANCH_TAG}" : "",
-  ])
-}
-
-function "get_cache_to_rocm" {
-  params = []
-  result = compact([
-    # Commit-scoped cache for exact re-runs.
-    BUILDKITE_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocm-${BUILDKITE_COMMIT},mode=${ROCM_FINAL_CACHE_TO_MODE}" : "",
-    # Branch-scoped cache so later commits on the same branch can reuse the full
-    # image layers when the parent-commit cache is evicted. Unlike the old
-    # rocm-latest tag (which caused duplicate exporter 400s), this is per-branch.
-    ROCM_CACHE_BRANCH_TAG != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocm-branch-${ROCM_CACHE_BRANCH_TAG},mode=${ROCM_FINAL_CACHE_TO_MODE}" : "",
-  ])
-}
-
-function "get_cache_from_rocm_csrc" {
-  params = []
-  result = compact([
-    BUILDKITE_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:csrc-rocm-${BUILDKITE_COMMIT}" : "",
-    PARENT_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:csrc-rocm-${PARENT_COMMIT}" : "",
-    VLLM_MERGE_BASE_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:csrc-rocm-${VLLM_MERGE_BASE_COMMIT}" : "",
-    ROCM_CACHE_BRANCH_TAG != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:csrc-rocm-branch-${ROCM_CACHE_BRANCH_TAG}" : "",
-    ROCM_CACHE_UPSTREAM_BRANCH_TAG != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:csrc-rocm-branch-${ROCM_CACHE_UPSTREAM_BRANCH_TAG}" : "",
-  ])
-}
-
-function "get_cache_to_rocm_csrc" {
-  params = []
-  result = compact([
-    # Export the exact-commit native cache for same-commit reruns.
-    BUILDKITE_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:csrc-rocm-${BUILDKITE_COMMIT},mode=${ROCM_CSRC_CACHE_TO_MODE}" : "",
-    # Export the branch-scoped native cache so later commits on the same branch
-    # can reuse compiled ROCm objects even when the exact parent cache is absent.
-    ROCM_CACHE_BRANCH_TAG != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:csrc-rocm-branch-${ROCM_CACHE_BRANCH_TAG},mode=${ROCM_CSRC_CACHE_TO_MODE}" : "",
-  ])
-}
-
-function "get_cache_from_rocm_rust" {
-  params = []
-  result = compact([
-    BUILDKITE_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rust-rocm-${BUILDKITE_COMMIT}" : "",
-    PARENT_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rust-rocm-${PARENT_COMMIT}" : "",
-    VLLM_MERGE_BASE_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rust-rocm-${VLLM_MERGE_BASE_COMMIT}" : "",
-    ROCM_CACHE_BRANCH_TAG != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rust-rocm-branch-${ROCM_CACHE_BRANCH_TAG}" : "",
-    ROCM_CACHE_UPSTREAM_BRANCH_TAG != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rust-rocm-branch-${ROCM_CACHE_UPSTREAM_BRANCH_TAG}" : "",
-  ])
-}
-
-function "get_cache_to_rocm_rust" {
-  params = []
-  result = compact([
-    # Export exact-commit and branch-scoped Rust caches. A content-addressed
-    # cache ref is appended by ci-bake-rocm.sh when that wrapper is used.
-    BUILDKITE_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rust-rocm-${BUILDKITE_COMMIT},mode=${ROCM_RUST_CACHE_TO_MODE}" : "",
-    ROCM_CACHE_BRANCH_TAG != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rust-rocm-branch-${ROCM_CACHE_BRANCH_TAG},mode=${ROCM_RUST_CACHE_TO_MODE}" : "",
-  ])
-}
-
-# Cache functions for upstream dependency stages (NIXL/UCX, DeepEP).
-# These stages are pinned to specific upstream commit hashes, so cache keys use
-# those hashes rather than the Buildkite commit. This means the cache persists
-# across all vLLM commits as long as the upstream dependency pins don't change.
-
-function "get_cache_from_rocm_deps" {
-  params = []
-  result = compact([
-    NIXL_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:nixl-rocm-${NIXL_CACHE_KEY}" : (NIXL_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:nixl-rocm-${NIXL_BRANCH}-ucx-${UCX_BRANCH}" : ""),
-    DEEPEP_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:deepep-rocm-${DEEPEP_CACHE_KEY}" : (DEEPEP_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:deepep-rocm-${DEEPEP_BRANCH}" : ""),
-  ])
-}
-
-function "get_cache_to_rocm_nixl" {
-  params = []
-  result = compact([
-    NIXL_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:nixl-rocm-${NIXL_CACHE_KEY},mode=min" : (NIXL_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:nixl-rocm-${NIXL_BRANCH}-ucx-${UCX_BRANCH},mode=min" : ""),
-  ])
-}
-
-function "get_cache_to_rocm_deepep" {
-  params = []
-  result = compact([
-    DEEPEP_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:deepep-rocm-${DEEPEP_CACHE_KEY},mode=min" : (DEEPEP_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:deepep-rocm-${DEEPEP_BRANCH},mode=min" : ""),
-  ])
-}
-
-# CI targets
-
-target "_ci-rocm" {
-  annotations = [
-    "manifest:vllm.buildkite.build_number=${BUILDKITE_BUILD_NUMBER}",
-    "manifest:vllm.buildkite.build_id=${BUILDKITE_BUILD_ID}",
-  ]
-  args = {
-    ARG_PYTORCH_ROCM_ARCH = PYTORCH_ROCM_ARCH
-    CI_BASE_IMAGE         = CI_BASE_IMAGE
-    ROCM_SMOKE_ID         = BUILDKITE_BUILD_ID
-    max_jobs              = CI_MAX_JOBS
-    VLLM_KERNEL_SYMBOL_MAP = VLLM_KERNEL_SYMBOL_MAP
-    buildkite_commit      = BUILDKITE_COMMIT
-  }
-}
-
-target "test-rocm-ci" {
-  inherits   = ["_common-rocm", "_ci-rocm", "_labels"]
-  target     = "test"
-  cache-from = get_cache_from_rocm()
-  cache-to   = get_cache_to_rocm()
-  tags = compact([
-    IMAGE_TAG,
-    IMAGE_TAG_LATEST,
-  ])
-  output = ["type=registry"]
-}
-
-# Validate the test image in the shared BuildKit graph and export only the
-# success marker. This avoids pulling the multi-GB image into the host daemon.
-target "smoke-test-rocm-ci" {
-  inherits   = ["_common-rocm", "_ci-rocm"]
-  target     = "export_test_smoke"
-  cache-from = get_cache_from_rocm()
-  output     = ["type=local,dest=./build/rocm-smoke-export"]
-}
-
-# Cache-only target for the source-scoped ROCm native build stage.
-# This persists the csrc-build stage in the registry cache even though the
-# final test image only consumes it indirectly while packaging the wheel.
-target "csrc-rocm-ci" {
-  inherits   = ["_common-rocm", "_ci-rocm"]
-  target     = "csrc-build"
-  cache-from = get_cache_from_rocm_csrc()
-  cache-to   = get_cache_to_rocm_csrc()
-  output     = ["type=cacheonly"]
-}
-
-target "kernel-symbol-map-rocm" {
-  inherits   = ["_common-rocm", "_ci-rocm"]
-  target     = "kernel-symbol-map-rocm"
-  cache-from = get_cache_from_rocm_csrc()
-  output     = ["type=local,dest=./kernel-symbol-map-rocm"]
-}
-
-# Cache-only target for the Rust frontend build stage. Final-image cache
-# exports use mode=min and do not reliably persist intermediate cargo layers,
-# so Rust gets its own source-scoped cache target.
-target "rust-rocm-ci" {
-  inherits   = ["_common-rocm", "_ci-rocm"]
-  target     = "rust-build"
-  cache-from = get_cache_from_rocm_rust()
-  cache-to   = get_cache_to_rocm_rust()
-  output     = ["type=cacheonly"]
-}
-
-# Keep wheel export on the same CI graph as the test image build so the
-# shared build_vllm/export_vllm stages resolve identically within one bake
-# invocation. Without this, export-wheel-rocm uses the plain local target
-# args while test-rocm-ci uses CI-only args, which can lead to separate
-# cache lineages and inconsistent export_vllm results.
-target "export-wheel-rocm" {
-  inherits   = ["_common-rocm", "_ci-rocm"]
-  target     = "export_vllm"
-  cache-from = get_cache_from_rocm()
-  cache-to   = get_cache_to_rocm()
-  output     = ["type=local,dest=./wheel-export"]
-}
-
-# Artifact-only vLLM build. GPU test jobs consume this artifact on top of
-# ci_base, avoiding a per-commit multi-GB image push/pull.
-group "test-rocm-ci-with-artifacts" {
-  targets = ["rust-rocm-ci", "csrc-rocm-ci", "export-wheel-rocm"]
-}
-
-# Full test image + wheel export. Kept for fallback/debugging when a pushed,
-# build-scoped image is useful.
-group "test-rocm-ci-with-wheel" {
-  targets = [
-    "rust-rocm-ci",
-    "csrc-rocm-ci",
-    "test-rocm-ci",
-    "smoke-test-rocm-ci",
-    "export-wheel-rocm",
-  ]
-}
-
-# Primary output tag for the ci_base build. In Buildkite this is a unique,
-# build-scoped tag; ci-bake-rocm.sh validates it before creating content and
-# stable aliases.
-variable "CI_BASE_IMAGE_TAG" {
-  default = "rocm/vllm-dev:ci_base"
-}
-
-# Versioned, content-addressed trusted ref. Preview builds import it read-only;
-# the wrapper creates their source-scoped content alias after validation.
-variable "CI_BASE_TRUSTED_CONTENT_REF" {
-  default = ""
-}
-
-# Cache-only targets for upstream dependency stages. These persist each stage
-# in the registry cache keyed by its upstream commit hash. When ci_base rebuilds
-# (e.g., requirements change), these stages are cache hits if their upstream
-# pins haven't changed -- saving ~35min of compilation.
-target "nixl-rocm-ci" {
-  inherits   = ["_common-rocm", "_ci-rocm"]
-  target     = "build_nixl"
-  cache-from = get_cache_from_rocm_deps()
-  cache-to   = get_cache_to_rocm_nixl()
-  output     = ["type=cacheonly"]
-}
-
-target "deepep-rocm-ci" {
-  inherits   = ["_common-rocm", "_ci-rocm"]
-  target     = "build_deepep"
-  cache-from = get_cache_from_rocm_deps()
-  cache-to   = get_cache_to_rocm_deepep()
-  output     = ["type=cacheonly"]
-}
-
-# Builds only the ci_base stage (NIXL, DeepEP, torchcodec, etc.)
-# Invoked by the ensure-ci-base step when the content hash of ci_base-affecting
-# files drifts from the remote image label. Per-PR builds then pull the result
-# as CI_BASE_IMAGE instead of rebuilding those slow layers on every commit.
-# Uses inline cache metadata on the ci_base image itself instead of exporting a
-# separate registry cache artifact.
-target "ci-base-rocm-ci" {
-  inherits   = ["_common-rocm"]
-  target     = "ci_base"
-  args = {
-    max_jobs = CI_MAX_JOBS
-  }
-  cache-from = concat(
-    compact([
-      CI_BASE_IMAGE_TAG != "" ? "type=registry,ref=${CI_BASE_IMAGE_TAG}" : "",
-      CI_BASE_TRUSTED_CONTENT_REF != "" ? "type=registry,ref=${CI_BASE_TRUSTED_CONTENT_REF}" : "",
-    ]),
-    # Import upstream dependency caches so NIXL/DeepEP stages
-    # are cache hits even when ci_base itself needs rebuilding.
-    get_cache_from_rocm_deps(),
-  )
-  cache-to = ["type=inline"]
-  tags     = compact([CI_BASE_IMAGE_TAG])
-  attest   = ["type=provenance,disabled=true"]
-  output   = ["type=registry"]
-}
-
-# Group for ci_base builds -- exports dependency stage caches alongside the
-# ci_base image so future rebuilds can reuse them independently.
-group "ci-base-rocm-ci-with-deps" {
-  targets = ["nixl-rocm-ci", "deepep-rocm-ci", "ci-base-rocm-ci"]
-}
diff --git a/docker/context/base-wheels/.keep b/docker/context/base-wheels/.keep
new file mode 100644
index 000000000000..e69de29bb2d1
diff --git a/docker/docker-bake-rocm.hcl b/docker/docker-bake-rocm.hcl
index 2eb1d5a43076..75c0ebe0cb02 100644
--- a/docker/docker-bake-rocm.hcl
+++ b/docker/docker-bake-rocm.hcl
@@ -1,143 +1,162 @@
-# docker-bake-rocm.hcl - vLLM ROCm Docker build configuration
-#
-# This file lives in the vLLM repo at docker/docker-bake-rocm.hcl
-# Equivalent of docker-bake.hcl for ROCm builds.
-#
-# Usage:
-#   docker buildx bake -f docker/docker-bake-rocm.hcl              # Build test (default)
-#   docker buildx bake -f docker/docker-bake-rocm.hcl final-rocm   # Build final image
-#   docker buildx bake -f docker/docker-bake-rocm.hcl --print      # Show resolved config
-#
-# CI usage (with the vLLM-owned CI overlay):
-#   docker buildx bake -f docker/docker-bake-rocm.hcl -f docker/ci-rocm.hcl test-rocm-ci
-
-variable "MAX_JOBS" {
-  # Empty string lets the Dockerfile fall back to $(nproc) via
-  # MAX_JOBS="${MAX_JOBS:-$(nproc)}" in each RUN step, which uses all
-  # available cores on whatever machine the build runs on.
-  # Override with --set '*.args.max_jobs=8' for local builds on small machines.
-  default = ""
-}
-
-variable "PYTORCH_ROCM_ARCH" {
-  default = "gfx90a;gfx942;gfx950"
-}
-
-variable "COMMIT" {
-  default = ""
-}
-
-# Content hash of ci_base-affecting files. Computed by ci-bake-rocm.sh and
-# embedded as a label so future builds can compare without rebuilding.
-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" {
-  default = "0"
-}
-
-variable "VLLM_BRANCH" {
-  default = "main"
-}
-
-# CI_BASE_IMAGE: pre-built ci_base image for per-PR test builds.
-# Defaults to the local "ci_base" stage for standalone/local builds.
-# CI overrides this to "rocm/vllm-dev:ci_base" via environment variable.
-variable "CI_BASE_IMAGE" {
-  default = "rocm/vllm-dev:ci_base"
-}
-
-# Upstream dependency commit pins. Plain local bake builds use the Dockerfile
-# ARG defaults. ci-bake-rocm.sh resolves those defaults (plus any env
-# overrides) and writes a small HCL override before invoking CI targets.
-variable "NIXL_BRANCH" {
-  default = ""
-}
-
-variable "UCX_BRANCH" {
-  default = ""
-}
-
-variable "DEEPEP_BRANCH" {
-  default = ""
-}
-
-group "default" {
-  targets = ["test-rocm"]
-}
-
-target "_common-rocm" {
-  dockerfile = CI_BASE_DOCKERFILE
-  context    = "."
-  args = {
-    max_jobs                        = MAX_JOBS
-    ARG_PYTORCH_ROCM_ARCH           = PYTORCH_ROCM_ARCH
-    REMOTE_VLLM                     = REMOTE_VLLM
-    VLLM_BRANCH                     = VLLM_BRANCH
-    CI_BASE_IMAGE                   = CI_BASE_IMAGE
-  }
-}
-
-target "_labels" {
-  labels = {
-    "org.opencontainers.image.source"      = "https://github.com/vllm-project/vllm"
-    "org.opencontainers.image.vendor"      = "vLLM"
-    "org.opencontainers.image.title"       = "vLLM ROCm"
-    "org.opencontainers.image.description" = "vLLM: A high-throughput and memory-efficient inference and serving engine for LLMs (ROCm)"
-    "org.opencontainers.image.licenses"    = "Apache-2.0"
-    "org.opencontainers.image.revision"    = COMMIT
-  }
-  annotations = [
-    "manifest:org.opencontainers.image.revision=${COMMIT}",
-  ]
-}
-
-target "test-rocm" {
-  inherits = ["_common-rocm", "_labels"]
+# ROCm images (docker/Dockerfile.rocm). CI/release drive this via .buildkite/scripts/rocm/build.sh. Locally:
+#   docker buildx bake -f docker/docker-bake-rocm.hcl                    # test image, everything from source
+#   STACK_IMAGE=vllm/vllm-openai-rocm:base-nightly RUNTIME=vllm/vllm-openai-rocm:base-nightly \
+#     docker buildx bake -f docker/docker-bake-rocm.hcl vllm-openai        # compile only vLLM
+variable "STACK"        { default = "rocm" }   # rocm: TheRock ROCm 10
+variable "STACK_IMAGE"  { default = "" }       # prebuilt stack ref@digest; empty: build it in this graph
+variable "STACK_WHEELS" { default = "" }       # that stack's -wheels image; empty: build it in this graph
+variable "RUNTIME"      { default = "" }       # prebuilt runtime ref@digest; empty: build it in this graph
+variable "TEST_DEPS"    { default = "" }       # prebuilt test-deps ref@digest; empty: build it in this graph
+variable "CI_ARCH"      { default = "gfx90a;gfx942;gfx950" }  # MI250, MI300/325, MI355: what CI can test
+variable "TAG"          { default = "vllm-rocm:dev" }
+variable "PUSH"         { default = false }
+variable "CACHE"        { default = "" }       # trusted cache repo (rocm/vllm-ci-cache): every CI build reads it
+variable "CACHE_TO"     { default = "" }       # ref prefix this build writes ("" = none); build.sh sets it per repo
+variable "AITER_SPEC"   { default = "" }       # pins the AITER nightly wheel (pip spec)
+variable "AITER_NIGHTLY" { default = "" }      # "1": test the newest AITER nightly wheel
+variable "VLLM_VERSION" { default = "" }       # release: image version label (build.sh reuse check)
+variable "CACHE_CI"     { default = true }     # false (PR builds): skip the full ci cache; compile-cache still exports
+variable "BUILDKITE_COMMIT" { default = "" }   # the commit the kernel symbol map is stamped with
+variable "CONTEXT"      { default = "." }      # build.sh: .rocm-src, an export of HEAD
+variable "SCCACHE_BUCKET_NAME" { default = "" } # release: S3 sccache for the csrc compile (release-pipeline.yaml)
+
+function "cfrom" {   # trusted cache (CI only), plus this build's own cache when CACHE_TO points elsewhere
+                     # (without its "ci" when CACHE_CI is false: nothing writes that)
+  params = [names]
+  result = flatten([for p in distinct(compact([CACHE == "" ? "" : "${CACHE}:", CACHE_TO])) :
+                    [for n in names : "type=registry,ref=${p}${STACK}-${n}" if CACHE_CI || n != "ci" || p != CACHE_TO]])
+}
+function "cto" {
+  params = [name]
+  result = CACHE_TO == "" ? [] : ["type=registry,ref=${CACHE_TO}${STACK}-${name},mode=max,image-manifest=true,ignore-error=true,compression=zstd,compression-level=3,force-compression=true"]
+}
+# Images and cache layers are zstd: on the CI fleet, unpacking dominates pull time (vllm#60258,
+# vllm#28656). Pulling needs Docker >= 23 or containerd >= 1.5. force-compression keeps every layer
+# zstd: without it, a builder exports whichever copy of a layer it holds, e.g. a gzip one imported
+# from an older cache. File mtimes stay as built: clamping them would leave every timestamp-based
+# .pyc (pip's, and those written by imports during the build) stale in each new container.
+function "out" {
+  params = []
+  result = PUSH ? "type=registry,compression=zstd,compression-level=3,force-compression=true,oci-mediatypes=true" : "type=docker"
+}
+function "ctx" {   # named contexts: prebuilt refs when given, else in-graph targets
+  params = []
+  result = merge(
+    { stack = STACK_IMAGE == "" ? "target:stack-${STACK}" : "docker-image://${STACK_IMAGE}" },
+    RUNTIME == "" ? {} : { runtime = "docker-image://${RUNTIME}" },
+    TEST_DEPS == "" ? {} : { test-deps = "docker-image://${TEST_DEPS}" })
+}
+
+target "_rocm" {
+  context = CONTEXT
+  args    = { SOURCE_DATE_EPOCH = "0", STACK = STACK }
+  attest  = ["type=provenance,disabled=true"]   # provenance would change the pushed digest on every build
+  output  = [out()]
+}
+target "_stack" {
+  inherits   = ["_rocm"]
+  target     = "final"
+  tags       = [TAG]
+  cache-from = cfrom(["stack"])
+  cache-to   = cto("stack")
+}
+# literal dockerfile paths: ci-infra ci_selector (images.py) only follows `dockerfile = "..."` literals
+target "stack-rocm" {
+  inherits   = ["_stack"]
+  dockerfile = "docker/Dockerfile.rocm_base"
+}
+target "stack-wheels" {
+  inherits = ["stack-${STACK}"]
+  target   = "wheels"
+  tags     = ["${TAG}-wheels"]
+  cache-to = []
+}
+target "runtime" {
+  inherits   = ["_rocm"]
+  dockerfile = "docker/Dockerfile.rocm"
+  target     = "runtime"
+  contexts   = ctx()
+  tags       = [TAG]
+  cache-from = cfrom(["runtime"])
+  cache-to   = cto("runtime")
+}
+target "test-deps" {   # content-keyed like the runtime (build.sh base mode)
+  inherits   = ["_rocm"]
+  dockerfile = "docker/Dockerfile.rocm"
+  target     = "test-deps"
+  contexts   = ctx()
+  tags       = [TAG]
+  cache-from = cfrom(["test-deps", "runtime", "ci"])
+  cache-to   = cto("test-deps")
+}
+target "_app" {
+  inherits   = ["_rocm"]
+  dockerfile = "docker/Dockerfile.rocm"
+  contexts   = ctx()
+  args       = { ARCH = "", AITER_SPEC = AITER_SPEC, AITER_NIGHTLY = AITER_NIGHTLY, VLLM_VERSION = VLLM_VERSION,
+                 SCCACHE_BUCKET_NAME = SCCACHE_BUCKET_NAME, CI_ARCH = CI_ARCH }
+  # dev/csrc records are shared. Only release (CACHE = "") writes "release"; only CI "ci" and "compile".
+  cache-from = cfrom(CACHE == "" ? ["runtime", "release"] : ["runtime", "ci", "compile"])
+}
+target "test" {
+  inherits = ["_app"]
   target   = "test"
-  tags     = ["rocm/vllm:test"]
-  output   = ["type=docker"]
-}
-
-# CI base image target - builds only the ci_base stage (NIXL, DeepEP,
-# torchcodec, requirements, etc.). Used by the weekly scheduled build and
-# the auto-rebuild trigger when requirements change in a PR.
-target "ci-base-rocm" {
-  inherits = ["_common-rocm", "_labels"]
-  target   = "ci_base"
-  labels   = {
-    "vllm.ci_base.content_hash" = CI_BASE_CONTENT_HASH
-  }
-  tags     = ["rocm/vllm-dev:ci_base"]
-  output   = ["type=docker"]
-}
-
-# Wheel export target - extracts the built vLLM wheel + test workspace
-# to local disk. Used by CI to upload the wheel as a Buildkite artifact
-# so test jobs can assemble images locally from ci_base + wheel instead
-# of pulling the full large image from Docker Hub.
-#
-# Usage:
-#   docker buildx bake -f docker/docker-bake-rocm.hcl export-wheel-rocm
-#   # Creates ./wheel-export/*.whl, ./wheel-export/requirements/, etc.
-#
-# After a full bake build, BuildKit cache makes this nearly instant.
-target "export-wheel-rocm" {
-  inherits = ["_common-rocm"]
-  target   = "export_vllm"
-  output   = ["type=local,dest=./wheel-export"]
-}
-
-target "final-rocm" {
-  inherits = ["_common-rocm", "_labels"]
-  target   = "final"
-  tags     = ["rocm/vllm:latest"]
-  output   = ["type=docker"]
-}
+  args     = { ARCH = CI_ARCH }
+  tags     = [TAG]
+  cache-to = CACHE_CI ? cto("ci") : []
+  no-cache-filter = AITER_SPEC == "" && AITER_NIGHTLY != "1" ? [] : ["test-base"]   # AITER nightly: always fetch the newest wheel
+}
+target "test-smoke" {
+  inherits = ["test"]
+  target   = "test-smoke"
+  output   = ["type=cacheonly"]
+  cache-to = []
+}
+target "compile-cache" {   # csrc + rust-deps on their own: a fresh builder (or a new PR) reuses them
+  inherits = ["_app"]
+  target   = "compile"
+  args     = { ARCH = CI_ARCH }
+  output   = ["type=cacheonly"]
+  cache-to = cto("compile")
+}
+target "vllm-wheel" {      # the python-only install test's wheel, uploaded as a build artifact
+  inherits = ["test"]
+  target   = "vllm-wheel"
+  output   = ["type=local,dest=.rocm-wheel"]
+  cache-to = []
+}
+target "kernel-symbol-map" {   # CI test selection (build.sh ci, VLLM_KERNEL_SYMBOL_MAP=1): from test's csrc
+  inherits = ["test"]
+  target   = "kernel-symbol-map"
+  args     = { BUILDKITE_COMMIT = BUILDKITE_COMMIT }
+  tags     = []
+  output   = ["type=local,dest=.rocm-kernel-map"]
+  cache-to = []
+}
+target "vllm-openai" {
+  inherits = ["_app"]
+  target   = "vllm-openai"
+  tags     = [TAG]
+}
+# cacheonly checks
+target "openai-smoke" {
+  inherits = ["_app"]
+  target   = "openai-smoke"
+  output   = ["type=cacheonly"]
+}
+target "wheels" {
+  inherits = ["_app"]
+  target   = "wheels"
+  tags     = ["${TAG}-wheels"]
+  contexts = merge(ctx(), { stack-wheels = STACK_WHEELS == "" ? "target:stack-wheels" : "docker-image://${STACK_WHEELS}" })
+  cache-to = cto("release")   # the all-arch compile, for the next release
+}
+target "version" {
+  inherits = ["_app"]
+  target   = "version"
+  output   = ["type=local,dest=.rocm-version"]
+}
+
+group "default" { targets = ["test"] }
+group "base"    { targets = ["stack-${STACK}", "stack-wheels"] }   # build.sh pushes runtime separately (own key)
+group "ci"      { targets = ["test", "test-smoke", "compile-cache", "vllm-wheel"] }
diff --git a/docker/versions.json b/docker/versions.json
index 03a9babdc0d9..84e4ac65d915 100644
--- a/docker/versions.json
+++ b/docker/versions.json
@@ -37,9 +37,6 @@
     "INSTALL_KV_CONNECTORS": {
       "default": "false"
     },
-    "BUILD_NIXL": {
-      "default": "false"
-    },
     "SCCACHE_BUCKET_NAME": {
       "default": "vllm-build-sccache"
     },
@@ -83,7 +80,7 @@
       "default": ">=1.0.17"
     },
     "RUNAI_MODEL_STREAMER_VERSION": {
-      "default": ">=0.15.7"
+      "default": ">=0.17.1"
     },
     "GGUF_VERSION": {
       "default": ">=0.19.0"
diff --git a/docs/.nav.yml b/docs/.nav.yml
index 783c40bb5c6d..0b470f200019 100644
--- a/docs/.nav.yml
+++ b/docs/.nav.yml
@@ -30,7 +30,7 @@ nav:
       - models/extensions
       - Hardware Supported Models:
         - models/hardware_supported_models/*
-        - TPU: https://docs.vllm.ai/projects/tpu/en/latest/recommended_models_features/
+        - TPU: https://docs.vllm.ai/projects/tpu/en/latest/recommended_models/
     - Features: features
   - Developer Guide:
     - contributing/README.md
diff --git a/docs/benchmarking/cli.md b/docs/benchmarking/cli.md
index 1ea961fefa4f..533b89816971 100644
--- a/docs/benchmarking/cli.md
+++ b/docs/benchmarking/cli.md
@@ -14,6 +14,29 @@ It's a living document, updated as new features and datasets become available.
 For exact-batch model timing, including steady-state pipeline-parallel forwards,
 see [vllm bench forward](../cli/README.md#forward).
 
+For skewed synthetic MoE expert loads, use `--routing zipf`:
+
+```bash
+vllm bench forward --model Qwen/Qwen3-30B-A3B \
+  --routing zipf --case decode=6*1@8192 --case prefill=2048
+```
+
+Each token independently selects distinct experts without replacement, weighted
+by popularity rank to the power of `-1.25`. Expert popularity is randomly
+permuted per layer to avoid systematically assigning hot experts to the lowest
+EP ranks.
+The fixture is seeded by `--seed`, fixed across timing replays, and generated
+outside the timed forward. Sequence-parallel TP ranks use their corresponding
+token rows. Zipf allows EP load imbalance; `--routing balanced` gives equal
+load per EP rank. Neither synthetic mode supports grouped top-k routing;
+Zipf also does not support prefill context parallelism.
+
+The default `--routing auto` selects Zipf with dummy weights and real routing
+otherwise. When the model or parallel layout does not support the synthetic
+fixture, `auto` warns and falls back to real routing. Use `--routing balanced`
+to explicitly request equal EP load. Zipf is a synthetic workload, not a fit
+to a particular model's routing or speculative-token correlations.
+
 ## Attention Kernel Benchmark
 
 Run the attention kernel smoke benchmark from the repository root using the
diff --git a/docs/cli/README.md b/docs/cli/README.md
index a8880f493d3c..116cd12f9c58 100644
--- a/docs/cli/README.md
+++ b/docs/cli/README.md
@@ -183,8 +183,10 @@ Use `--exclude-drafter` to measure only the target forward. Verification assumes
 all draft tokens are accepted and uses a fixed bonus token; target sampling and
 rejection sampling are excluded. Both paths use CUDA graphs and restore the
 batch's state before each replay.
-Balanced MoE routing fixtures apply to the target model; the drafter uses its
-model routing, so use real weights when measuring a MoE drafter.
+Synthetic MoE routing fixtures (`--routing balanced` or `--routing zipf`) apply
+to the target model; the drafter uses its model routing, so use real weights
+when measuring a MoE drafter. For the Zipf distribution and its limitations,
+see [Model Forward Benchmark](../benchmarking/cli.md#model-forward-benchmark).
 
 ```bash
 vllm bench forward --model Qwen/Qwen3-1.7B \
diff --git a/docs/configuration/optimization.md b/docs/configuration/optimization.md
index 18122ea2c8bc..fea8977b2ede 100644
--- a/docs/configuration/optimization.md
+++ b/docs/configuration/optimization.md
@@ -156,10 +156,47 @@ GPU. vLLM can pin each worker with `numactl` before the Python subprocess starts
 so the interpreter, imports, and early allocator state are created with the
 desired NUMA policy from the beginning.
 
-Use `--numa-bind` to enable the feature. By default, vLLM auto-detects the
-GPU-to-NUMA mapping and uses `--cpunodebind=<node> --membind=<node>` for each
-worker. When you need a custom CPU policy, add `--numa-bind-cpus` and vLLM will
-switch to `--physcpubind=<cpu-list> --membind=<node>`.
+By default, vLLM binds automatically when the environment supports it:
+CUDA-alike platform, multiple NUMA nodes, `numactl` installed, permission to
+set NUMA memory policy, unrestricted CPU affinity, a default NUMA memory policy,
+the `spawn` multiprocessing start method (the CLI default), and a local `mp` or
+`uni` executor with multiprocessing DP, including supported prefill context
+parallelism (PCP) launches. Ray, external launchers and `--device-ids` remapping
+are excluded from automatic binding.
+Otherwise vLLM starts unbound. Use
+`--numa-bind` to make a missing `numactl` or undetectable GPU-to-NUMA topology a
+startup error, or `--no-numa-bind` to disable binding. Passing
+`--numa-bind-nodes` or `--numa-bind-cpus` explicitly enables binding even without
+`--numa-bind`: either list requests strict `--membind`, not automatic soft memory
+preference or CPU-only binding. Neither can be combined with `--no-numa-bind`.
+Automatic binding preserves an existing non-default memory policy, such as
+interleave, preferred-node, or bind. Explicit enablement may override it.
+
+Automatic binding uses `--cpunodebind=<node> --preferred=<node>` for each worker:
+memory allocations prefer the GPU-local node but can fall back to other nodes.
+EngineCore prefers its shard's nodes with `--preferred-many=<nodes>` when the
+shard spans multiple NUMA nodes, or `--preferred=<node>` for a single-node shard.
+Explicit `--numa-bind` or binding lists retain strict `--membind=<nodes>`.
+When you need a custom CPU policy, add `--numa-bind-cpus` and vLLM will switch
+to `--physcpubind=<cpu-list>`.
+
+Automatic eligibility and topology validation happen in each local launcher
+before starting its first bound subprocess, using freshly discovered topology.
+The mapping must cover all visible GPUs and local workers, and each mapped NUMA
+node must have CPUs. The resolved decision and mapping are passed to spawned
+children, so an already-bound EngineCore does not repeat detection under its
+restricted CPU affinity.
+PCP workers use the full TP × PCP × PP shard size for local DP GPU offsets;
+EngineCore's CPU affinity covers all workers in its shard.
+Multi-node multiprocessing launchers resolve their own host-local mapping;
+NUMA node IDs are never interpreted as cluster-wide IDs.
+
+Before each bound launch, vLLM probes the selected policy with
+`numactl <args> true`. If CPU-and-memory binding is rejected, it retries with
+CPU-only binding; if that is also rejected, the subprocess starts unbound.
+Logs report the selected policy or fallback. If `--preferred-many` is unsupported,
+automatic mode falls back to CPU-only binding, never strict memory binding.
+These probes test policy acceptance, not future allocation capacity.
 
 These `--numa-bind*` options only apply to GPU execution processes. They do not
 configure the CPU backend's separate thread-affinity controls. Automatic
@@ -174,7 +211,7 @@ order as the GPU indices. Each CPU list must use
 `numactl --physcpubind` syntax such as `0-3`, `0,2,4-7`, or `16-31,48-63`.
 
 ```bash
-# Auto-detect NUMA nodes for visible GPUs
+# Require binding, auto-detecting NUMA nodes for visible GPUs
 vllm serve meta-llama/Llama-3.1-8B-Instruct \
   --tensor-parallel-size 4 \
   --numa-bind
@@ -195,10 +232,12 @@ vllm serve meta-llama/Llama-3.1-8B-Instruct \
 
 Notes:
 
-- CLI usage forces multiprocessing to use the `spawn` method automatically. If you enable NUMA binding through the Python API, also set `VLLM_WORKER_MULTIPROC_METHOD=spawn`.
-- Automatic detection relies on NVML and NUMA support from the host. If it cannot determine the mapping reliably, pass `--numa-bind-nodes` explicitly.
+- The CLI defaults to `spawn` when `VLLM_WORKER_MULTIPROC_METHOD` is unset. Automatic NUMA binding never changes the start method: if it is `fork`, default binding stays off. Python API usage normally defaults to `fork`; set `VLLM_WORKER_MULTIPROC_METHOD=spawn` to allow binding and follow Python's `if __name__ == "__main__":` requirements.
+- Automatic memory preference permits allocations on other NUMA nodes when local memory is unavailable. Explicit `--numa-bind` still uses `--membind`, which restricts host allocations to the selected nodes; for host-memory-heavy setups (e.g. large CPU KV offload), avoid explicit strict binding if workers need memory from other nodes. The policy probes do not check allocation capacity or container memory limits.
+- Automatic detection relies on NVML (CUDA) or amdsmi (ROCm) and NUMA support from the host. If it cannot determine the mapping reliably, default mode starts unbound; pass `--numa-bind-nodes` to provide a mapping explicitly.
 - Explicit `--numa-bind-nodes` and `--numa-bind-cpus` values must be valid `numactl` inputs. vLLM does a small amount of validation, but the effective binding semantics are still determined by `numactl`.
 - The current implementation binds GPU execution processes such as `EngineCore` and multiprocessing workers. It does not apply NUMA binding to frontend API server processes or the DP coordinator.
+- Ray actors are not bound by the multiprocessing executable wrapper. Enabling automatic NUMA binding for Ray requires a separate actor-launch implementation.
 - In containerized environments, NUMA policy syscalls may require extra permissions, such as `--cap-add SYS_NICE` when running via `docker run`.
 
 ### CPU Backend Thread Affinity
@@ -380,6 +419,11 @@ A processed item larger than this budget is served uncached (with a warning)
 instead of failing engine startup; raise `mm_processor_cache_gb` if you want
 those items cached.
 
+With processor-only caching, a request can contain more processed multi-modal
+data than the cache budget. Cached tensors and prompt updates needed by the
+request remain referenced while its inputs are merged, even if they are evicted
+from the LRU. The budget limits cache storage, not transient per-request memory.
+
 If you do not benefit much from the cache, you can disable both IPC
 and processor caching completely via `mm_processor_cache_gb=0`.
 
diff --git a/docs/contributing/README.md b/docs/contributing/README.md
index d665116e47b5..c97b5bd6c17a 100644
--- a/docs/contributing/README.md
+++ b/docs/contributing/README.md
@@ -103,6 +103,13 @@ integration, use the configuration files in `tools/pre_commit/`:
     pre-commit run -a  # runs on all files (short for --all-files)
     ```
 
+    Hook revisions in `.pre-commit-config.yaml` are pinned to commit SHAs. To bump them, use:
+
+    ```bash
+    pre-commit autoupdate --freeze               # bumps all hooks
+    pre-commit autoupdate --freeze --repo <url>  # bumps one hook
+    ```
+
 ### Documentation
 
 MkDocs is a fast, simple and downright gorgeous static site generator that's geared towards building project documentation. Documentation source files are written in Markdown, and configured with a single YAML configuration file, [mkdocs.yaml](../../mkdocs.yaml).
diff --git a/docs/contributing/dockerfile/dockerfile.md b/docs/contributing/dockerfile/dockerfile.md
index b5c8b8001a2f..307d98b9a0eb 100644
--- a/docs/contributing/dockerfile/dockerfile.md
+++ b/docs/contributing/dockerfile/dockerfile.md
@@ -29,11 +29,6 @@ The `extensions-build` stage can also produce an optional source-built Triton
 wheel. `vllm-openai-base` installs that wheel after its other Python
 dependencies so dependency resolution cannot restore an older Triton version.
 
-With `BUILD_NIXL=true`, `extensions-build` also builds NIXL, including its NIXL
-EP extension for the installed PyTorch, with `tools/build_nixl_with_torch215_for_rubin.py`.
-`vllm-openai-base` then replaces the NIXL packages installed from the KV-connector
-requirements with these wheels.
-
 The AWS EFA images are built separately by `docker/Dockerfile.efa` on top of a
 `vllm-openai` image built with `INSTALL_KV_CONNECTORS=true`. It adds a minimal
 AWS EFA userspace install (libfabric, aws-ofi-nccl and hwloc), the EFA Mooncake
diff --git a/docs/contributing/model/transcription.md b/docs/contributing/model/transcription.md
index f82bafa44d52..4fdab121cd85 100644
--- a/docs/contributing/model/transcription.md
+++ b/docs/contributing/model/transcription.md
@@ -66,7 +66,11 @@ This is for controlling general behavior of the API when serving your model:
 
 See [Audio preprocessing and chunking](#audio-preprocessing-and-chunking) for what each field controls.
 
-Implement the prompt construction via [get_generation_prompt][vllm.model_executor.models.interfaces.SupportsTranscription.get_generation_prompt]. The server builds a [SpeechToTextParams][vllm.config.speech_to_text.SpeechToTextParams] object that bundles the resampled waveform, task parameters, and request-specific options. Your model receives this single object and returns a valid [PromptType][vllm.inputs.llm.PromptType]. There are two common patterns:
+Implement the prompt construction via [get_generation_prompt][vllm.model_executor.models.interfaces.SupportsTranscription.get_generation_prompt]. The server builds a [SpeechToTextParams][vllm.config.speech_to_text.SpeechToTextParams] object that bundles the resampled waveform, task parameters, and request-specific options. Your model receives this single object and returns a valid [PromptType][vllm.inputs.llm.PromptType].
+
+For transcription requests, `stt_params.extra_args` contains the optional `vllm_xargs` dictionary, which is also passed to `SamplingParams.extra_args`. Models can read request-specific frontend or prompt options here and should validate the options they consume. An omitted dictionary is `None`. Multipart clients can send `vllm_xargs` as a JSON string. These options apply to the request; do not use them to mutate the shared model configuration.
+
+There are two common patterns:
 
 #### Multimodal LLM with audio embeddings (e.g., Voxtral, Gemma3n)
 
@@ -194,7 +198,7 @@ Provide a fast duration→token estimate to improve streaming usage statistics:
 
 The API server takes care of basic audio I/O and optional chunking before building prompts:
 
-- Resampling: Input audio is resampled to `SpeechToTextConfig.sample_rate` using `AudioResampler`.
+- Resampling: Input audio is resampled to `SpeechToTextConfig.sample_rate`. Models can set specific `audio_resample_method` (e.g., `pyav`, `scipy`, `torchaudio`, etc) to decode at the native sample rate, then resample with `AudioResampler` before duration calculation and chunking. This can match a model's reference audio frontend without changing other models' defaults.
 - Chunking: If `SpeechToTextConfig.allow_audio_chunking` is True and the duration exceeds `max_audio_clip_s`, the server splits the audio into chunks and generates a prompt per chunk. There is no overlap between chunks, overlap_chunk_second controls the size of the search window used to find the split point.
 - Energy-aware splitting: When `min_energy_split_window_size` is set, the server finds low-energy regions to minimize cutting within words.
 
diff --git a/docs/deployment/integrations/dynamo.md b/docs/deployment/integrations/dynamo.md
index 8d0a0dcb0c84..f22e9f520cbc 100644
--- a/docs/deployment/integrations/dynamo.md
+++ b/docs/deployment/integrations/dynamo.md
@@ -2,6 +2,6 @@
 
 [NVIDIA Dynamo](https://github.com/ai-dynamo/dynamo) is an open-source framework for distributed LLM inference that can run vLLM on Kubernetes with flexible serving architectures (e.g. aggregated/disaggregated, optional router/planner).
 
-For Kubernetes deployment instructions and examples (including vLLM), see the [Deploying Dynamo on Kubernetes](https://github.com/ai-dynamo/dynamo/blob/main/docs/kubernetes/README.md) guide.
+For Kubernetes deployment instructions and examples (including vLLM), see the [Deploying Dynamo on Kubernetes](https://docs.nvidia.com/dynamo/kubernetes/getting-started/introduction) guide.
 
 Background reading: InfoQ news coverage — [NVIDIA Dynamo simplifies Kubernetes deployment for LLM inference](https://www.infoq.com/news/2025/12/nvidia-dynamo-kubernetes/).
diff --git a/docs/deployment/integrations/k8s-aibom.md b/docs/deployment/integrations/k8s-aibom.md
new file mode 100644
index 000000000000..8f1cf1ac25a9
--- /dev/null
+++ b/docs/deployment/integrations/k8s-aibom.md
@@ -0,0 +1,29 @@
+# k8s-aibom
+
+[k8s-aibom](https://github.com/GoogleCloudPlatform/k8s-aibom) is an open-source Kubernetes controller that keeps an inventory of the AI workloads actually serving in a cluster. For every vLLM deployment it produces a [CycloneDX ML-BOM](https://cyclonedx.org/capabilities/mlbom/): the vLLM image and resolved digest, the model the server was started with, and, when the model is signed, the verified signer identity. Each attribute carries the evidence it was read from and a confidence tier (declared, inferred, or unresolved).
+
+It is read-only toward your cluster: it watches workloads through the Kubernetes API, adds no sidecars or init containers, runs unprivileged, and never changes scheduling or serving behavior. vLLM is detected whether you run it as a plain Deployment or StatefulSet, as a LeaderWorkerSet or DisaggregatedSet, under [KubeAI](kubeai.md), [KServe](kserve.md), [KubeRay](kuberay.md), [llm-d](llm-d.md), or the [production stack](production-stack.md), which ships it as an optional chart dependency.
+
+## Install
+
+```bash
+helm install k8s-aibom oci://ghcr.io/googlecloudplatform/charts/k8s-aibom \
+  --version 1.6.0 \
+  --namespace k8s-aibom-system --create-namespace
+```
+
+Then opt in each namespace that serves models:
+
+```bash
+kubectl label namespace <your-serving-namespace> aibom.k8saibom.dev/enabled=true
+```
+
+Documents appear as `AIBOM` resources next to the workloads they describe. The `kubectl aibom` plugin (`kubectl krew install aibom`) summarizes and searches them:
+
+```bash
+kubectl aibom summary -A
+kubectl aibom find --runtime vllm -A
+kubectl aibom find --image vllm/vllm-openai -A
+```
+
+See the [k8s-aibom documentation](https://github.com/GoogleCloudPlatform/k8s-aibom#readme) for the document format, signature verification, external sinks, and the full list of detected runtimes.
diff --git a/docs/deployment/integrations/llm-d.md b/docs/deployment/integrations/llm-d.md
index dc335edb0399..d125a8dc034a 100644
--- a/docs/deployment/integrations/llm-d.md
+++ b/docs/deployment/integrations/llm-d.md
@@ -9,10 +9,10 @@ It is a [CNCF Sandbox project](https://www.cncf.io/blog/2026/03/24/welcome-llm-d
 A single vLLM server is fast, but at scale the picture changes: across many replicas, cache locality breaks under round-robin load balancing, long prompts inflate time-to-first-token, and accelerators sit underused. llm-d adds the cluster-level layer that vLLM does not aim to provide on its own:
 
 - **[Prefix-aware routing](https://llm-d.ai/docs/well-lit-paths/foundations/precise-prefix-cache-routing).** Instead of round-robin, llm-d reads vLLM's KV-cache events and routes each request to the replica that already holds its prefix, reusing cache instead of recomputing it.
-- **[Distributed KV-cache management](https://llm-d.ai/docs/architecture/advanced/kv-management).** A global index tracks which token blocks live on which replica, and [tiered offloading](https://llm-d.ai/docs/well-lit-paths/foundations/tiered-prefix-cache) spills cache to CPU memory or local SSD, extending the working set beyond accelerator HBM.
+- **[Distributed KV-cache management](https://llm-d.ai/docs/architecture/kv-management).** A global index tracks which token blocks live on which replica, and [tiered offloading](https://llm-d.ai/docs/well-lit-paths/foundations/tiered-prefix-cache) spills cache to CPU memory or local SSD, extending the working set beyond accelerator HBM.
 - **[Prefill/decode disaggregation](https://llm-d.ai/docs/well-lit-paths/foundations/pd-disaggregation).** Prompt processing and token generation run on separate vLLM workers, with KV-cache moved over the vLLM [NIXL connector](https://docs.vllm.ai/en/latest/features/nixl_connector_usage/), lowering TTFT and steadying per-token latency on long prompts.
 - **[Wide expert-parallelism](https://llm-d.ai/docs/well-lit-paths/foundations/wide-expert-parallelism).** Serve large Mixture-of-Experts models such as DeepSeek-R1 and GPT-OSS across nodes with combined data and expert parallelism, for more KV-cache capacity and throughput.
-- **SLO-aware [autoscaling](https://llm-d.ai/docs/well-lit-paths/foundations/workload-autoscaling) and [flow control](https://llm-d.ai/docs/well-lit-paths/foundations/flow-control).** Scale vLLM pools on real inference signals (queue depth, true demand) rather than raw GPU utilization, with multi-tenant fairness and priority dispatch.
+- **SLO-aware [autoscaling](https://llm-d.ai/docs/operations/autoscaling) and [flow control](https://llm-d.ai/docs/operations/traffic/flow-control).** Scale vLLM pools on real inference signals (queue depth, true demand) rather than raw GPU utilization, with multi-tenant fairness and priority dispatch.
 
 These are composable. Most teams start by adding prefix-aware routing over an existing vLLM pool, then layer in the rest as specific bottlenecks appear.
 
diff --git a/docs/deployment/k8s.md b/docs/deployment/k8s.md
index e7d0853e9f49..8b10baab23d4 100644
--- a/docs/deployment/k8s.md
+++ b/docs/deployment/k8s.md
@@ -431,6 +431,43 @@ You can also verify the health service manually with `grpcurl`:
 grpcurl -plaintext localhost:50051 grpc.health.v1.Health/Check
 ```
 
+### Rust frontend gRPC health checks
+
+The Rust frontend is a separate server from `--grpc` above. Start it with `VLLM_USE_RUST_FRONTEND=1 vllm serve <model> --grpc-port 50051` or `vllm-rs serve <model> --grpc-port 50051`. Its gRPC port serves `grpc.health.v1.Health` and reports engine liveness and inference readiness under separate service names, so a paused engine can leave the load balancer without being restarted:
+
+| Service | Reports `NOT_SERVING` when |
+| --- | --- |
+| `liveness` | The engine has failed. Pause, sleep, weight updates, and orderly shutdown do not affect it until the listener closes. |
+| `readiness` and `vllm.Inference` | The engine has failed, shutdown has started, the engine is paused (including for a process checkpoint), asleep, or partially awake, or a weight update has not finished successfully. |
+| `vllm.Control` and `""` | The engine has failed or shutdown has started. |
+
+Point the startup and liveness probes at `liveness` and the readiness probe at `readiness`. A liveness probe on `readiness` or `vllm.Inference` restarts an engine that stays paused longer than the probe's failure threshold. The listener stops accepting connections when shutdown starts, so probes see the endpoint as unavailable while `--shutdown-timeout` drains in-flight requests.
+
+```yaml
+startupProbe:
+  grpc:
+    port: 50051
+    service: liveness
+  periodSeconds: 10
+  failureThreshold: 120 # Allow for model loading.
+livenessProbe:
+  grpc:
+    port: 50051
+    service: liveness
+readinessProbe:
+  grpc:
+    port: 50051
+    service: readiness
+```
+
+A paused or sleeping pod leaves the endpoints of its Services, so send resume, wake-up, and weight-update calls to the pod directly, for example through a headless Service with `publishNotReadyAddresses: true`. After a process checkpoint fails, the engine stays paused and `liveness` stays `SERVING`, so whatever drives the checkpoint must restart the pod when `/checkpoint/status` reports the failure.
+
+This server does not enable gRPC reflection, so check a service manually with [`grpc_health_probe`](https://github.com/grpc-ecosystem/grpc-health-probe):
+
+```bash
+grpc_health_probe -addr=localhost:50051 -service=readiness
+```
+
 ## Troubleshooting
 
 ### Startup Probe or Readiness Probe Failure, container log contains "KeyboardInterrupt: terminated"
diff --git a/docs/design/arch_overview.md b/docs/design/arch_overview.md
index 711e9a2668eb..f7695d4a3d92 100644
--- a/docs/design/arch_overview.md
+++ b/docs/design/arch_overview.md
@@ -131,6 +131,58 @@ A data parallel deployment with 8 GPUs (`vllm serve -tp=2 -dp=4`) has:
 For CPU resource sizing recommendations, see
 [CPU Resources for GPU Deployments](../configuration/optimization.md#cpu-resources-for-gpu-deployments).
 
+## Scheduler
+
+The scheduler runs inside the **Engine Core Process** and decides which
+requests execute each step. The code can be found in
+[vllm/v1/core/sched/scheduler.py](../../vllm/v1/core/sched/scheduler.py).
+
+The scheduler is **token-budget-based**: each step, it fills a budget of
+`max_num_batched_tokens` across both running and waiting requests. There are no
+separate prefill and decode phases. Each request tracks how many of its tokens
+have been computed, and the scheduler assigns it tokens to close the gap. This
+single loop handles chunked prefill, prefix caching and speculative decoding.
+
+### Scheduling Policies
+
+Two policies are supported via `--scheduling-policy`:
+
+- **FCFS** (default): requests are scheduled first-come, first-served. When
+    the KV cache is exhausted, the last-admitted running request is preempted.
+- **Priority**: requests carry an integer `priority` field. When the KV cache
+    is exhausted, the running request with the lowest priority (highest
+    `priority` value, then latest arrival time) is preempted.
+
+### Preemption
+
+When a request cannot be allocated new KV cache blocks, the scheduler
+**preempts** running requests, chosen by the scheduling policy, until enough
+blocks are free:
+
+1. The preempted request's KV cache blocks are freed.
+2. Its progress is reset and it returns to the waiting queue. Under FCFS it
+    goes to the front of the queue; under Priority it is reinserted by
+    `(priority, arrival_time)`.
+3. When re-admitted, it is recomputed rather than swapped back from CPU
+    memory. Prefix caching or a KV connector can restore part of its progress
+    immediately.
+
+For how to reduce preemption, see
+[Preemption](../configuration/optimization.md#preemption).
+
+### Async Scheduling
+
+With async scheduling,
+[vllm/v1/core/sched/async_scheduler.py](../../vllm/v1/core/sched/async_scheduler.py)
+overlaps CPU scheduling with GPU execution: while the GPU executes step *N*,
+the CPU schedules step *N+1*, using placeholders for the tokens that step *N*
+will produce.
+
+If a request is preempted while it still has output in flight, that output is
+normally still delivered when it arrives. It is dropped instead when delivering
+it would be invalid, for example when `reset_prefix_cache` preempts and resumes
+requests within the same step.
+
 ## LLM Engine
 
 The `LLMEngine` and `AsyncLLMEngine` classes are central to the functioning of
diff --git a/docs/design/custom_op.md b/docs/design/custom_op.md
index f99c8043fa04..84faeaeff42e 100644
--- a/docs/design/custom_op.md
+++ b/docs/design/custom_op.md
@@ -264,7 +264,7 @@ Currently, thanks to [vLLM's hardware-plugin mechanism](./plugin_system.md), the
 
 - **Official device plugins:** [vllm-ascend](https://github.com/vllm-project/vllm-ascend) (for Huawei Ascend NPU), [vllm-spyre](https://github.com/vllm-project/vllm-spyre)
 (for Spyre), [vllm-gaudi](https://github.com/vllm-project/vllm-gaudi) (for Intel Gaudi), [vllm-neuron](https://github.com/vllm-project/vllm-neuron) (for AWS Neuron), [vllm-meta](https://github.com/vllm-project/vllm-metal) (for Apple Silicon), etc.
-- **Non-official device plugins:** [vllm-metax](https://github.com/MetaX-MACA/vLLM-metax) (for MetaX GPU), [vllm-kunlun](https://github.com/baidu/vLLM-Kunlun) (for Baidu Kunlun XPU), [vllm-musa](https://github.com/MooreThreads/vllm-musa) (for Moore Threads GPU), etc.
+- **Non-official device plugins:** [vllm-metax](https://github.com/MetaX-MACA/vLLM-metax) (for MetaX GPU), [vllm-kunlun](https://github.com/baidu/vLLM-Kunlun) (for Baidu Kunlun XPU), [vllm-musa](https://github.com/MooreThreads/vllm-musa) (for Moore Threads GPU), [vllm-iluvatar](https://github.com/Deep-Spark/vllm-iluvatar) (for Iluvatar CoreX GPU), etc.
 
 In this case, `CustomOp` can enable these hardware manufacturers to seamlessly replace vLLM's operations with their deep-optimized kernels for specific devices at runtime, by just registering an OOT `CustomOp` and implementing the `forward_oot()` method.
 
@@ -314,3 +314,46 @@ In addition, you can also register all your `CustomOp` at one place for better m
     for op_name, op_cls in REGISTERED_CUSTOM_OPS.items():
         CustomOp.register_oot(_decorated_op_cls=op_cls, name=op_name)
     ```
+
+#### Opt in to the hw-agnostic layers
+
+With `VLLM_USE_HW_AGNOSTIC=1`, the Transformers modeling backend builds layers from `vllm.model_executor.hw_agnostic.layers` where a hw-agnostic implementation exists, and from the in-tree `vllm.model_executor.layers` otherwise. These are separate classes with separate `op_registry_oot`s, so an override of the in-tree class does not reach the hw-agnostic one.
+
+To override the class the Transformers backend builds on either path, obtain it with `hw_agnostic.resolve` instead of importing it, and register through that class's `register_oot`:
+
+??? code
+
+    ```python
+    from vllm.model_executor import hw_agnostic
+
+    SiluAndMul = hw_agnostic.resolve("activation", "SiluAndMul")
+
+
+    @SiluAndMul.register_oot
+    class CustomSiluAndMul(SiluAndMul):
+
+        def forward_oot(...):
+            # Call optimized device-specific kernels.
+            ...
+    ```
+
+`resolve` returns the hw-agnostic class when `VLLM_USE_HW_AGNOSTIC` is set and one exists, and the in-tree class otherwise, so the plugin code is the same for both paths. Do not register through `vllm.model_executor.custom_op.CustomOp.register_oot`, which always writes to the in-tree registry. On the hw-agnostic path, vLLM warns about an override registered only for the in-tree class, and raises if an override does not derive from the class it replaces.
+
+`VLLM_USE_HW_AGNOSTIC` only affects the Transformers backend. Models with a native vLLM implementation build the in-tree class even when it is set, so a plugin that also serves those models registers its override for the in-tree class as well.
+
+The Transformers backend builds its norms as `TPAwareRMSNorm` and `TPAwareGemmaRMSNorm`, subclasses of `RMSNorm` and `GemmaRMSNorm`. Overrides are keyed on the exact class name, so an `RMSNorm` or `GemmaRMSNorm` override does not reach them yet.
+
+Type checkers do not accept a variable as a base class, so plugins that run mypy can import the in-tree class for type checking only:
+
+??? code
+
+    ```python
+    from typing import TYPE_CHECKING
+
+    from vllm.model_executor import hw_agnostic
+
+    if TYPE_CHECKING:
+        from vllm.model_executor.layers.activation import SiluAndMul
+    else:
+        SiluAndMul = hw_agnostic.resolve("activation", "SiluAndMul")
+    ```
diff --git a/docs/design/hisparse.md b/docs/design/hisparse.md
index f5fab3c0f5d2..64aab6f52715 100644
--- a/docs/design/hisparse.md
+++ b/docs/design/hisparse.md
@@ -221,9 +221,16 @@ hot row? ──────── yes ──► existing hot physical row + upda
 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.
+Device operations are selected by `Platform.get_hisparse_backend_cls()`. The base
+platform returns `None`; CUDA selects `CUDAHiSparseBackend` in `platforms/cuda.py`.
+Hardware plugins can provide a `HiSparseBackend` for host registration, cache
+operations, IPC events, capture queries, and transfers while reusing the common
+runtime and connector. HiSparse validates availability and cache formats before
+allocation. Streams, ordinary events, current-stream queries, and synchronization
+use `torch.Stream`, `torch.Event`, and `torch.accelerator` directly. The backend
+provides IPC event creation/import/export only for shared host pools.
+Without HiSparse, sparse attention keeps its normal stream and event path.
+ROCm does not currently provide a HiSparse backend.
 
 ## Main classes
 
@@ -260,6 +267,5 @@ the same command, output, and cache-resolution boundaries.
 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.
+kernel. Other accelerators can implement `HiSparseBackend` and return its class
+path from their platform override without changing the shared control flow.
diff --git a/docs/design/hybrid_kv_cache_manager.md b/docs/design/hybrid_kv_cache_manager.md
index 3984000c7d8f..50e87cbdd2da 100644
--- a/docs/design/hybrid_kv_cache_manager.md
+++ b/docs/design/hybrid_kv_cache_manager.md
@@ -131,13 +131,17 @@ Some architectures (e.g., Bamba, Jamba, Minimax) interleave standard attention l
 The current algorithm is:
 
 1. Increase the `block_size` of attention layers until
+
     $$
     \text{block_size} \times \text{kv_hidden_size}_{\text{att}} \ge \text{state_size}_{\text{mamba}}
     $$
+
 2. Pad the mamba state per layer to
+
     $$
     \text{block_size} \times \text{kv_hidden_size}_{\text{att}}
     $$
+
 3. Apply the grouping strategy in case 3.
 
 !!! note
diff --git a/docs/design/metrics.md b/docs/design/metrics.md
index 5d9fa42c11c3..b063939ed142 100644
--- a/docs/design/metrics.md
+++ b/docs/design/metrics.md
@@ -597,7 +597,7 @@ A common use case for our metrics is to support automated scaling of
 vLLM instances.
 
 For related discussion from the
-[Kubernetes Serving Working Group](https://github.com/kubernetes/community/tree/master/wg-serving),
+[Kubernetes Serving Working Group](https://github.com/kubernetes/community/tree/main/archive/wg-serving),
 see:
 
 - [Standardizing Large Model Server Metrics in Kubernetes](https://docs.google.com/document/d/1SpSp1E6moa4HSrJnS4x3NpLuj88sMXr2tbofKlzTZpk)
diff --git a/docs/design/mm_processing.md b/docs/design/mm_processing.md
index 20226171a912..530e1f903284 100644
--- a/docs/design/mm_processing.md
+++ b/docs/design/mm_processing.md
@@ -90,6 +90,9 @@ GPU-side fusion is controlled by `multimodal_config.mm_device_do_normalize`:
 | `llama-nemotron-vl-embed`  | `LlamaNemotronVLForEmbedding`              | `nvidia/llama-nemotron-embed-vl-1b-v2`                |
 | `llama-nemotron-vl-rerank` | `LlamaNemotronVLForSequenceClassification` | `nvidia/llama-nemotron-rerank-vl-1b-v2`               |
 | `mistral3`                 | `Mistral3ForConditionalGeneration`         | `mistralai/Mistral-Small-3.2-24B-Instruct-2506`, etc. |
+| `nano-nemotron-vl`         | `NemotronH_Nano_VL_V2`                     | `nvidia/NVIDIA-Nemotron-Nano-12B-v2-VL-BF16`          |
+| `nano-nemotron-omni`       | `NemotronH_Nano_Omni_Reasoning_V3`         | `nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-NVFP4` |
+| `nemotron-omni`            | `NemotronH_Omni_Reasoning_V3`              | `nvidia/NVIDIA-Nemotron-3.5-Super-VL-120B-A12B-NVFP4` |
 | `nemotron-parse`           | `NemotronParseForConditionalGeneration`    | `nvidia/NVIDIA-Nemotron-Parse-2.0`, etc.              |
 | `qwen2-vl`                 | `Qwen2VLForConditionalGeneration`          | `Qwen/Qwen2-VL-2B-Instruct`, etc.                     |
 | `qwen2.5-vl`               | `Qwen2_5_VLForConditionalGeneration`       | `Qwen/Qwen2.5-VL-3B-Instruct`, etc.                   |
diff --git a/docs/design/nixl_kv_push_connector.md b/docs/design/nixl_kv_push_connector.md
index 621b0f7c9dd7..afdf84a9bd02 100644
--- a/docs/design/nixl_kv_push_connector.md
+++ b/docs/design/nixl_kv_push_connector.md
@@ -6,6 +6,17 @@ prefill completes. `NixlPushConnector` adds a **push-based** alternative
 in which P writes the KV blocks directly into D's pre-allocated memory
 via `NIXL WRITE`.
 
+Enable push with `"kv_connector": "NixlPushConnector"`. Alternatively, use
+`"kv_connector": "NixlConnector"` and set
+`"kv_connector_extra_config": {"transfer_mode": "push"}`.
+
+For the built-in `NixlConnector`, `transfer_mode` defaults to `"pull"` and accepts
+only `"pull"` or `"push"`; other values raise `ValueError`.
+Explicit `NixlPullConnector` and `NixlPushConnector` names ignore this option.
+`kv_connector_module_path` takes precedence over built-in selection.
+Under `MultiConnector`, set `transfer_mode` in the NIXL child's
+`kv_connector_extra_config`.
+
 This document describes the threading, queues, and scheduling
 interactions specific to the push design. The pull-mode design is
 unchanged; the push connector reuses the same handshake, NIXL agent
@@ -183,31 +194,40 @@ 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 peer's KV regions do not line up with the local ones by index under
+pipeline parallelism (PP) or 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** 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 layer routing derives its transfer layer order 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
-layer order. Descriptor offsets use each layer's region capacity, so P and
-D need not allocate the same number of blocks. Physical allocations are still
-registered once, even when multiple layers share them.
+* with HMA, layers of several cache groups can share one region, and
+  which layers share a region can differ between P and D.
+
+Both connectors therefore pair regions **by layer**, for every peer. Each
+worker advertises which layers back each of its NIXL regions in the
+handshake metadata (`NixlAgentMetadata.region_members`). On the
+handshake, `_pair_regions` looks each layer up on both sides and records
+one pair (remote region, local region, cache group) per layer,
+deduplicated; a layer registered as two regions (K and V apart) pairs
+them in order. The remote descriptor list keeps only the remote regions
+that pair, and a transfer reads or writes each pair at its cache group's
+blocks. A region shared by several groups thus transfers once per group
+it holds a layer of, from one descriptor list.
+
+```text
+D (PP=1):  region 0 = [L0, L0.swa]   region 1 = [L1]   region 2 = [L2]
+P stage 1: region 0 = [L2]
+
+pairs (remote, local, group) on P: (2, 0, 0)  ->  D's region 2 only
+```
+
+Region capacities stay per region on both sides, so P and D need not
+allocate the same number of blocks. Physical allocations are still
+registered once, even when multiple layers share them. Mamba/SSM hybrids
+pair the same way: each SSM region's conv and SSM-state descriptors are
+built from the peer region its layers pair with, so P and D may group the
+same layers into regions differently. A layer in different cache groups on
+P and D is refused at the handshake for these models.
 
 Packed MLA caches (block-outermost layouts such as DeepSeek-V4's) store the
 pages of every layer for one block in a single row. A PP stage holds only its
@@ -219,24 +239,27 @@ P (stage 2), block b: | L3 | L4 | L5 |
 ```
 
 `PP=1` peers transfer whole rows. They also advertise where each layer's page
-sits in the row (`NixlAgentMetadata.packed_member_layouts`), so a PP producer
-writes each of its pages to the matching slot in D's row. Page sizes must
-match; offsets, row strides, and block counts may differ. Because PP stages can
-list the same attention backends in a different order, packed push compares the
-sorted set of backend names during the handshake. Pull mode is unchanged.
-
-Invariants enforced when the remote regions are aligned:
-
-* every locally owned layer must be advertised exactly once by the
-  remote; a missing layer fails the handshake rather than silently
-  leaving that layer's KV stale, and remote-only layers (owned by other
-  PP stages) are ignored;
-* a remote that omits layer metadata while the local layout requires
-  layer routing fails the handshake instead of falling back to
-  region-index routing;
-* layer 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.
+sits in the row (`NixlAgentMetadata.packed_member_layouts`). A PP producer
+registers one region per layer, so it views each of D's rows as one region per
+layer and pairs its pages with those. Page sizes must match; offsets, row
+strides, and block counts may differ. Because PP stages can list the same
+attention backends in a different order, packed push compares the sorted set of
+backend names during the handshake.
+
+Invariants enforced when pairing:
+
+* the remote must hold every local layer; a missing layer fails the
+  handshake rather than silently leaving that layer's KV stale. Layers
+  only the remote holds are skipped on a PP stage (other stages own
+  them) and fail the handshake otherwise;
+* a remote that does not advertise its layers fails the handshake
+  instead of falling back to pairing by index;
+* a layer must span the same number of regions on both sides;
+* every TP rank of one remote engine must pair the same way;
+* a layer in different cache groups on P and D is refused by push. Each
+  worker advertises every layer's cache group
+  (`NixlAgentMetadata.layer_group_ids`), so this is caught even when the
+  layer's region is shared by several groups.
 
 Decode-side PP is unsupported because completions are counted per
 consumer rank. Mamba/SSM hybrids are unsupported under PP. Packed PP push
diff --git a/docs/features/README.md b/docs/features/README.md
index 1df8960ee354..969d5f33a0d9 100644
--- a/docs/features/README.md
+++ b/docs/features/README.md
@@ -1,5 +1,9 @@
 # Features
 
+## KV Cache Offloading
+
+- [UMBPStoreConnector](umbp_store_connector_usage.md)
+
 ## Compatibility Matrix
 
 The tables below show mutually exclusive features and the support on some hardware.
@@ -78,4 +82,4 @@ th:not(:first-child) {
 | beam-search | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
 
 !!! note
-    For information on feature support on Google TPU, please refer to the [TPU-Inference Recommended Models and Features](https://docs.vllm.ai/projects/tpu/en/latest/recommended_models_features/) documentation.
+    For information on feature support on Google TPU, please refer to the [TPU-Inference Feature Matrix](https://docs.vllm.ai/projects/tpu/en/latest/features/) documentation.
diff --git a/docs/features/custom_logitsprocs.md b/docs/features/custom_logitsprocs.md
index 8c1c25939928..53bb258d6078 100644
--- a/docs/features/custom_logitsprocs.md
+++ b/docs/features/custom_logitsprocs.md
@@ -33,7 +33,7 @@ Custom logits processors must subclass `vllm.v1.worker.gpu.sample.logits_process
     * Return whether this processor modifies logits for the request. The sampler skips the logits-processing pipeline for batches in which no request needs it
 
 * `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
+    * Flush host-side writes staged by `add_request()` to the device; called before the forward pass of every step that admits new requests, so only state set in `add_request()` is flushed. Read data that changes every step from `ctx` in `apply()`
 
 * `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
@@ -70,7 +70,7 @@ The contrived example below implements a custom logits processor which masks out
 
         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().
+            # flushed to the device in apply_staged_writes().
             self.target_token = torch.full(
                 (req_states.max_num_reqs,), -1, dtype=torch.int64
             )
@@ -264,6 +264,6 @@ The examples below show how a user would pass a custom argument (`target_token`)
 
     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
 
-* 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
+* Stage host-side writes in `add_request()` and flush them in `apply_staged_writes()`, instead of writing device tensors directly; this is the discipline every built-in sampler state follows
 
 * 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/ec_cpu_connector.md b/docs/features/ec_cpu_connector.md
index 18274b6dd970..b895d8dcbb0c 100644
--- a/docs/features/ec_cpu_connector.md
+++ b/docs/features/ec_cpu_connector.md
@@ -71,7 +71,7 @@ vllm serve <model> --ec-transfer-config '{
 
 ### Protocol
 
-- **Control plane**: ZMQ. The producer binds a `ROUTER` socket on `VLLM_EC_SIDE_CHANNEL_HOST:VLLM_EC_SIDE_CHANNEL_PORT`; each consumer opens one `DEALER` connection per producer peer, with ZMQ heartbeating (2s interval, 4s timeout, 8s TTL) to detect a dead peer. `XferReq`/`XferAck` are `msgspec` msgpack structs, versioned by `EC_CONNECTOR_VERSION` (currently `1`) — a version mismatch is NACKed.
+- **Control plane**: ZMQ. By default, each producer binds a `ROUTER` socket on the local IP returned by vLLM's `get_ip()` with an OS-assigned port and advertises the bound port through `ec_transfer_params.peer_port`. Optional host and port overrides are listed below. Each consumer opens one `DEALER` connection per producer peer, with ZMQ heartbeating (2s interval, 4s timeout, 8s TTL) to detect a dead peer. `XferReq`/`XferAck` are `msgspec` msgpack structs, versioned by `EC_CONNECTOR_VERSION` (currently `1`) — a version mismatch is NACKed.
 - **Compatibility check**: every `XferReq` carries a SHA-256 hash over `(vllm_version, model, dtype, block_size_bytes)`; the producer NACKs (`NACK_INCOMPAT`) any peer whose hash differs.
 - **Ack statuses**: `OK`, `NACK_MISSING` (the producer no longer holds the encoding), `NACK_NOT_READY` (held, but its save has not landed), `NACK_INCOMPAT`, `NACK_VERSION`, `NACK_INTERNAL`. Only `NACK_NOT_READY` is retryable — `RETRYABLE_NACKS` in `cpu/protocol.py` is what both ends consult, so the classification lives with the wire vocabulary rather than at each call site.
 - **Data plane**: NIXL, `UCX` backend, consumer-initiated `READ` — the consumer pulls bytes directly out of the producer's registered mmap region; the producer never pushes.
@@ -98,14 +98,14 @@ EC transfer is configured via `--ec-transfer-config` (CLI) or the `ec_transfer_c
 | `consumer_ack_timeout_s` | `float` | No (default `2.0`) | How long a consumer waits for an `XferAck` before giving up on a read. The producer answers `XferReq`s from its own scheduler step, so its reply latency scales with the encoder's `--max-num-batched-tokens`: a loaded encoder whose steps run longer than this makes consumers abandon reads the producer is about to grant. Raise it for large encoder batches. |
 | `ec_cpu_bytes` | `int` | Yes | Total size, in bytes, of the shared CPU mmap region. `ECCPUConnector` raises `ValueError` if unset. Block count = `ec_cpu_bytes // block_size_bytes`, where each block holds one LLM input-embedding vector. Entries allocate enough blocks for their modality-specific output shape, including visual DeepStack features. |
 
-### Environment Variables
+The control address uses vLLM's standard IP discovery. Set `VLLM_HOST_IP` to select a specific network interface when automatic discovery is unsuitable.
 
-| Variable | Default | Description |
+| Environment variable | Default | Description |
 | --- | --- | --- |
-| `VLLM_EC_SIDE_CHANNEL_HOST` | `localhost` | Host the producer's ZMQ `ROUTER` socket binds to. Set to a routable address (e.g. the pod IP) for multi-instance/multi-node P2P — the default only works when producer and consumer share a host. |
-| `VLLM_EC_SIDE_CHANNEL_PORT` | `5601` | Port for the same ZMQ `ROUTER` socket. |
+| `VLLM_EC_SIDE_CHANNEL_HOST` | Unset | Override the EC control host instead of using `get_ip()`. Must be reachable by consumers; wildcard addresses are rejected. |
+| `VLLM_EC_SIDE_CHANNEL_PORT` | `0` | `0` selects an OS-assigned port per producer. A nonzero value binds `base_port + data_parallel_index`; an occupied port fails startup. |
 
-Both are read only when `ec_enable_nixl=True` on a producer (`ec_role="ec_producer"` or `"ec_both"`).
+For fixed ports, assign non-overlapping ranges to independent instances sharing a host. Consumers use the advertised address and do not need matching port configuration.
 
 ## Limitations
 
diff --git a/docs/features/engram.md b/docs/features/engram.md
index 023477b34559..e3bba1ca0806 100644
--- a/docs/features/engram.md
+++ b/docs/features/engram.md
@@ -96,6 +96,41 @@ vllm serve deepseek-ai/DeepSeek-V4.1-Flash \
 CPU offload requires a GPU with UVA support;
 vLLM fails fast when it is unavailable.
 
+## Disk offload (Qwen4Exp)
+
+`disk_offload` keeps Qwen4Exp PLE tables in the checkpoint files: no table
+is allocated in GPU or pinned host memory, and nothing is copied to an offload
+directory. Each step reads that step's rows from the safetensors files on the
+host. Use it where neither `cpu_offload` nor a resident table fits, for example
+on unified-memory GPUs such as DGX Spark, where pinned host memory comes out of
+the same pool as the weights and the KV cache:
+
+```bash
+vllm serve <qwen4exp-checkpoint> --engram-config '{"disk_offload": true}'
+```
+
+It applies only to Qwen4Exp on CUDA and takes precedence over `cpu_offload`.
+
+Rows are read with `preadv` and copied to fixed device buffers before the
+forward, so CUDA graphs replay unchanged. They stay in the checkpoint format
+(BF16, FP8 or NVFP4) and are decoded on the GPU, so an unquantized table must
+match `--dtype`. Rows that were read recently are served from the OS page
+cache, which is shared by co-located ranks, not counted in the worker's RSS,
+and reclaimed by the kernel under memory pressure.
+
+Costs and limits:
+
+- The read needs the step's token ids, so it waits for the previous step on
+  the GPU, and every TP rank reads all of its rows.
+- A row that is not in the page cache is read from disk. For a fast first
+  request, read the PLE shard files once after the server has started.
+- The checkpoint files must not change while the server runs. Truncation stops
+  the engine with an error naming the file; deleting or renaming is safe.
+- Weight reload and elastic EP are unsupported. Loading fails when tensors are
+  not file-backed safetensors views (`.bin`, runai, `fastsafetensors`,
+  `instanttensor`, or `--safetensors-load-strategy eager`);
+  `--load-format dummy` serves zeros.
+
 ## Data-parallel topologies (DeepSeek V4.1)
 
 By default (`embedding_across_dp: false`), each DP replica keeps its own
@@ -115,7 +150,8 @@ vllm serve deepseek-ai/DeepSeek-V4.1-Flash \
 ```
 
 `embedding_across_dp` is not supported with elastic expert parallelism yet.
-Qwen4Exp PLE tables are ETP-sharded instead and honor `cpu_offload` only.
+Qwen4Exp PLE tables are ETP-sharded instead and honor only `cpu_offload` and
+`disk_offload`.
 
 ## Sharing host tables across DP replicas
 
diff --git a/docs/features/kv_offloading_usage.md b/docs/features/kv_offloading_usage.md
index 7450e0dacc13..6a938dfc0cd8 100644
--- a/docs/features/kv_offloading_usage.md
+++ b/docs/features/kv_offloading_usage.md
@@ -95,6 +95,8 @@ vllm serve <model> \
   }'
 ```
 
+Secondary tiers are not supported yet on engines whose workers span several nodes; such an engine fails at startup.
+
 ## `kv_connector_extra_config` Reference
 
 | Key | Required | Default | Scope | Notes |
@@ -111,7 +113,7 @@ vllm serve <model> \
 | `tier_poll_interval_s` | no | `0.001` | multi-tier | Pause between tier-poll rounds for tiers that must be serviced between engine steps (currently `p2p`). See [Threading model](#threading-model). Set to `0` to disable the thread and service those tiers once per engine step instead. |
 | `offload_prompt_only` | no | `true` | both | If `true`, only prompt (prefill) blocks are offloaded; decode blocks are skipped. |
 | `sync_load_budget_bytes` | no | `0` | both | Per-step budget of KV bytes (per GPU) loaded within the same step, so those requests run immediately instead of waiting a step for an async load. Hits that don't fit the remaining budget, or would leave nothing to compute, stay async. `0` disables. |
-| `self_describing_kv_events` | no | `false` | both | Opt-in. When `true` *and* KV cache events are enabled (`--kv-events-config` with `enable_kv_cache_events`), the connector emits self-describing block-granular `BlockStored`/`BlockRemoved` payloads (constituent block hashes, whole-chunk `token_ids`, per-block `block_size`, parent hash, LoRA + group/cache-spec metadata) instead of the placeholder fallback, so external KV-event consumers can index offloaded blocks. Inert unless events are enabled. With `TieringOffloadingSpec`, a CPU promotion is self-describing when a local request observes its primary-tier `HIT` before event translation; otherwise its stored event may retain the placeholder, while a later `HIT` can backfill metadata for removal. Pending-removal/re-promotion races and externally initiated promotions may also produce placeholders, and consumers must ignore removals for unknown hashes. Partial recurrent tails emit the hash-aligned portion from the physical block start through the tail boundary. Other sliding-window/SSM chunks keep the placeholder fallback. In chunk mode (`block_size` > GPU block size, or `blocks_per_chunk` > 1), overlapping chunks re-announce shared per-block hashes, so consumers must reference-count (deduplicate) repeated store/remove announcements. |
+| `self_describing_kv_events` | no | `false` | both | Opt-in. When `true` *and* KV cache events are enabled (`--kv-events-config` with `enable_kv_cache_events`), the connector emits self-describing block-granular `BlockStored`/`BlockRemoved` payloads (constituent block hashes, whole-chunk `token_ids`, per-block `block_size`, parent hash, LoRA + group/cache-spec metadata) instead of the placeholder fallback, so external KV-event consumers can index offloaded blocks. Inert unless events are enabled. With `TieringOffloadingSpec`, a CPU promotion is self-describing when a local request observes its primary-tier `HIT` before event translation; otherwise its stored event may retain the placeholder, while a later `HIT` can backfill metadata for removal. Pending-removal/re-promotion races and externally initiated promotions may also produce placeholders, and consumers must ignore removals for unknown hashes. Recurrent (Mamba / linear-attention) state groups announce state checkpoints, with `kv_cache_spec_kind` `mamba`: an offloaded chunk or partial tail holds one recurrent state, valid at its last hash, so its `BlockStored` carries that single hash, its parent hash and an empty `token_ids`. Consumers align it by hash with the attention group's blocks; a hybrid prefix is reusable only up to a boundary where every group has a residency. Sliding-window and chunked-local attention chunks keep the placeholder fallback. In chunk mode (`block_size` > GPU block size, or `blocks_per_chunk` > 1), overlapping chunks re-announce shared per-block hashes, so consumers must reference-count (deduplicate) repeated store/remove announcements. |
 | `spec_module_path` | no | — | both | Python import path for a custom `OffloadingSpec` not in the built-in registry. Required only when `spec_name` is not built-in (advanced). |
 | `canonical_layout` | no | `false` | both | Store each offloaded block in a common canonical representation instead of each worker's local GPU page (enables sharing across TP sizes and single-copy storage of replicated KV). See [Canonical Layout](#canonical-layout). |
 
@@ -193,6 +195,11 @@ context provides each key's token position through `get_offload_key_position(key
 so a policy can establish prefix ordering during insertion. Policies that do not
 use request-specific insertion ordering can accept and ignore `req_context`.
 
+Custom policies must also implement `on_request_finished`, which receives the
+request's keys per KV cache group in head-to-tail order and applies the request's
+cache access once, when it finishes. `CachePolicy` no longer has a `touch` method;
+recency updates belong in this hook.
+
 `eviction_policy` is checked against the built-in registry first; if it isn't a registered name, vLLM imports `cache_policy_module_path` and looks up `eviction_policy` as a class name in that module — the same fallback `spec_module_path` provides for a custom `OffloadingSpec`. No import or registration call needs to run before the server starts.
 
 ### Registering a friendly short name (in-process only)
@@ -224,9 +231,9 @@ Most tiers only ever see the scheduler thread, which holds the lock for the leng
 Two consequences worth knowing:
 
 - Any tier can be reached from that thread indirectly, because serving a peer looks keys up through the tiering manager, which fans out to the other tiers. Tier state must therefore not assume a particular thread, only that the manager lock is held.
-- Work the thread deliberately leaves alone is batched per step: a promotion it initiates is submitted at the next `on_schedule_end()`, not immediately.
+- Work the thread deliberately leaves alone is batched per step: a promotion it initiates is submitted at the next `on_schedule_end()`, not immediately. Because only an engine step submits that promotion and collects its completion, a multi-tier setup that includes an opted-in tier (for example `p2p` alongside `fs`) keeps the engine stepping even when idle, so pause and graceful shutdown wait there too. A setup whose only secondary tier is the opted-in one still goes idle, since a peer's lookup cannot promote out of it. Work the thread can finish by itself, such as serving a peer's lookup and fetch, does not count as pending, so a producer that peers keep fetching from still goes idle. Pause does not stop the P2P server either: it stops scheduling and quiesces the device, and a paused engine keeps serving peers from its CPU tier.
 
-Setting `tier_poll_interval_s` to `0` disables the thread; opted-in tiers then get serviced once per engine step, as they were before it existed.
+Setting `tier_poll_interval_s` to `0` disables the thread; opted-in tiers then get serviced once per engine step, as they were before it existed. With the thread disabled the engine keeps stepping for such a tier's lifetime, so an idle engine does not report itself idle and pause and graceful shutdown do not complete on their own. A round that keeps failing does not stop the thread: after 30 s of consecutive failures it backs off to a slower interval — the larger of 1 s and `tier_poll_interval_s` — and keeps retrying, so peers are still served while the engine is idle. Failing rounds log at most one message every 30 s: a full traceback for the first failure of a run, a one-line summary for the rest.
 
 ### Filesystem (FS)
 
diff --git a/docs/features/mooncake_connector_usage.md b/docs/features/mooncake_connector_usage.md
index 42c1dcd676ce..e0df25436ca2 100644
--- a/docs/features/mooncake_connector_usage.md
+++ b/docs/features/mooncake_connector_usage.md
@@ -66,6 +66,7 @@ cache), see [ECMooncakeConnector Usage Guide](mooncake_ec_connector_usage.md).
 - `VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT`: Timeout (in seconds) for automatically releasing the prefiller’s KV cache for a particular request. (Optional)
     - Default: 480
     - If a request is aborted and the decoder has not yet notified the prefiller, the prefill instance will release its KV-cache blocks after this timeout to avoid holding them indefinitely.
+    - When `kv_lease_duration` is set, the lease decides when those blocks are released instead. This timeout still bounds how long a pull waits for a prefill that has not finished.
 
 ## KV Transfer Config
 
@@ -80,6 +81,10 @@ cache), see [ECMooncakeConnector Usage Guide](mooncake_ec_connector_usage.md).
 - **num_workers**: Size of thread pool for one prefiller worker to transfer KV caches by mooncake. (default 10)
 - **mooncake_protocol**: Mooncake connector protocol. (default "rdma")
 - **device_name**: Comma-separated whitelist of RDMA devices (e.g. `"mlx5_0,mlx5_1"`) to restrict topology discovery to. Empty discovers every device. Useful on hosts exposing a mix of InfiniBand and RoCE ports, where both peers must settle on the same link layer.
+- **kv_lease_duration**: Seconds the prefiller holds a finished request's KV blocks before the decoder must renew them. (default unset)
+    - Unset keeps the fixed `VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT`.
+    - When set, the decoder renews the lease every `kv_lease_duration / 6` seconds from the time the request arrives until its pull settles, and each renewal extends it to at least `kv_lease_duration * 2 / 3` seconds from then, as in [NixlConnector](../design/nixl_kv_cache_lease.md). Blocks of a decoder that crashed are freed within about `kv_lease_duration` seconds, while requests queued on a busy decoder keep their blocks.
+    - Set the same value on prefill and decode instances. A decoder without it does not renew, so its requests that wait longer than the lease lose their prefilled KV; a prefiller without it rejects renewals, and the decoder logs a warning. Prefill instances from releases without this option log an error for each renewal they receive, so upgrade them before enabling it on decode instances.
 - **bootstrap_server_address**: Optional `"host:port"` of an already running Mooncake bootstrap server. All prefiller ranks register with this endpoint instead of launching an embedded server. The caller owns the server lifetime and must configure the proxy to use the same endpoint. The example proxy assumes the bootstrap server is on the prefiller API host; a different host requires a proxy that supplies that address in `remote_bootstrap_addr`. Omit this option to use the embedded server and `VLLM_MOONCAKE_BOOTSTRAP_PORT`.
 
 ## Prometheus metrics
@@ -102,8 +107,9 @@ The following Prometheus metrics are exported when MooncakeConnector is active,
 
 !!! tip
     High `vllm:mooncake_num_kv_expired_reqs` indicates that the prefiller's
-    lease duration is too short for your network or workload. Increase it via
-    `VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT`.
+    lease duration is too short for your network or workload. Increase
+    `VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT`, or set `kv_lease_duration` on both
+    instances so that queued requests keep their blocks.
 
 ## Example Scripts/Code
 
diff --git a/docs/features/mooncake_store_connector_usage.md b/docs/features/mooncake_store_connector_usage.md
index ac5f432e1bc7..a46642f335cc 100644
--- a/docs/features/mooncake_store_connector_usage.md
+++ b/docs/features/mooncake_store_connector_usage.md
@@ -177,13 +177,56 @@ remain unchanged. A decoder configured with `"save_decode_cache": true` uses
 the same Store TP for decode KV from every prefiller.
 
 The list may contain positive integer TP sizes. Sharing requires a Store TP that
-is at least the local TP and divisible by it, an LBHNC or LBNHC local KV cache,
-and the existing topology and KV-head constraints. The Store namespace includes
-the attention backend's selected layout. Different layouts use separate Store
-entries. Malformed lists and unsupported endpoints use an isolated rank-local
-key layout. When
-`enable_store_tp_lcm` is absent or false, `prefill_tp_sizes` has no effect and
-the existing `store_tp_size` behavior is unchanged.
+is at least the local TP and divisible by it. LBHNC, LBNHC, BLHNC, BLNHC, LHBNC,
+and BHLNC local KV cache layouts are supported. The connector normalizes the
+four head-major layouts to one Store value order and the two token-major
+layouts to another. Physical layouts with the same value order can share
+entries.
+
+In hybrid models, local Attention block sizes can vary with TP size. Shared
+Store chunks have the same size on every endpoint; a local block can contain
+several chunks. Hybrid sharing requires the same explicit `--prefix-match-unit`
+on every endpoint, which also sets the default Attention Store chunk size.
+A larger `store_chunk_size` can reduce Store object and metadata overhead:
+
+```json
+{
+    "kv_connector_extra_config": {
+        "store_tp_size": 8,
+        "store_chunk_size": 512
+    }
+}
+```
+
+The configured size applies to Attention Store objects. It must be a multiple
+of `prefix_match_unit`, divide every local Attention page, and align with each
+layer's stored states. Sharing endpoints must use the same Store chunk size.
+
+Mamba and GDN entries contain complete state snapshots. Both `DS` and `SD`
+convolution state layouts are supported; sharing endpoints use the same layout.
+Incremental offloads save intervals aligned across all participating cache
+groups. Partial-tail checkpoints use the core's state-block handoff and matching
+Attention KV. Store jobs hold their source blocks until all workers finish.
+
+Lookup reports a reusable prefix only when every required cache group is
+present. For each Attention page, this includes every Store chunk covering its
+reusable tokens, across all Store shards.
+
+Attention groups are planned from their per-layer cache specs, including specs
+wrapped by `UniformTypeKVCacheSpecs`. The Store TP and global head-slot count
+must divide one another; MQA and replicated GQA use the smaller head-slot count
+as the physical Store shard count. `MambaSpec` state segments map in complete
+state units; a replicated local unit can serve multiple Store shards. All
+endpoints use the same PP size and Store TP; PCP, DCP, and cross-layer blocks
+remain rank-local. Each group namespace records its Store value format and
+normalized value schema. Incompatible configurations use an isolated rank-local
+namespace.
+
+Cache groups marked `prefix_cacheable=False` are excluded from shared-layout
+validation, Store transfers, and prefix lookup. Unsupported layouts in participating
+groups use the rank-local fallback.
+
+When `enable_store_tp_lcm` is absent or false, `prefill_tp_sizes` has no effect.
 
 **Proxy:**
 
@@ -285,7 +328,8 @@ Strict isolation requires a Mooncake master started with `--enable_multi_tenants
   incompatible with sleep mode and the CuMem allocator.
 - `cache_prefix` (str): Namespace prepended to every store key. Lets separate deployments share one Mooncake master without polluting each other — instances configured with different prefixes never see each other's cached blocks, even for identical prompts. All instances that should share a prefix cache must use the same value. Default: `""` (no prefix; keys are byte-identical to the unprefixed format).
 - `save_decode_cache` (bool): Enable offloading decode tokens' KV cache. A `kv_consumer` does not save during prefill; when decode starts, it fills any missing block-aligned prompt prefix before appending completed decode blocks. Default: `false`.
-- `store_tp_size` (int): Common Store TP for endpoints with different local TP sizes. It supports LBHNC and LBNHC local KV caches, with `store_tp_size >= local_tp_size` and `store_tp_size % local_tp_size == 0`. The current topology is one full-attention cache group, PCP/DCP disabled, and cross-layer blocks disabled. For GQA and MHA, the total KV-head count must be divisible by `store_tp_size`. Store shards contain fixed global KV-head ranges in the local layout. Shared endpoints use the same KV cache layout, pipeline-parallel size, and Store TP. The Store namespace includes the layout and PP size. Unsupported configurations use a topology-specific rank-local namespace.
+- `store_tp_size` (int): Common Store TP for divisible heterogeneous-TP
+  sharing. See [Sharing one Store across multiple Prefill TP sizes](#sharing-one-store-across-multiple-prefill-tp-sizes).
 
 LBHNC/HND is strongly recommended for TP-sharded Store when supported. LBNHC/NHD
 creates many transfer segments and may significantly reduce PUT/GET performance.
@@ -294,13 +338,10 @@ For example, with prefill TP 4, decode TP 2, and eight KV heads, set
 `store_tp_size` to 4 on both instances. Each decode rank reads and writes two
 of the four Store shards.
 
-MQA with one total KV head uses a replicated-head layout. For the supported
-prefill TP 4 to decode TP 2 case, every rank uses the
-same rank-0 key namespace. The four prefill replicas stripe block PUTs so each
-object is stored once, while both decode ranks GET every block into their local
-KV replica. `store_tp_size` does not appear in MQA keys, so identical MQA
-objects written at different store TP sizes share the same pool entry when PP
-sizes match.
+MQA with one total KV head uses one physical Store shard. For prefill TP 4 and
+decode TP 2, every rank uses the same shard key. Prefill replicas stripe block
+PUTs so each object is stored once, while both decode ranks GET every block into
+their local KV replica.
 
 Tensor-parallel collectives and low-precision arithmetic are not bitwise
 invariant across TP sizes, so heterogeneous-TP reuse does not guarantee the
diff --git a/docs/features/moriio_connector_usage.md b/docs/features/moriio_connector_usage.md
index 352022c7aa6b..d806e04cde8c 100644
--- a/docs/features/moriio_connector_usage.md
+++ b/docs/features/moriio_connector_usage.md
@@ -142,7 +142,7 @@ WRITE mode supports hybrid attention plus Mamba/KDA models, including DSpark. Us
 
 For DSpark, configure the same draft model on both instances through `--speculative-config`. Target and draft attention groups retain their own block IDs and token spans; only the running recurrent state is transferred, not speculative scratch slots. The producer transfers the state after the first N-1 prompt tokens, and the consumer recomputes the last prompt token. Convolution and SSM transfers must both finish before the layer is complete.
 
-Hybrid READ keeps its existing single-transferable-attention-group requirement. Plain-attention WRITE retains its existing notification format.
+Hybrid READ transfers every attention group, so drafters that share the target's attention groups (such as MTP) also work. Plain-attention WRITE retains its existing notification format.
 
 ### Kimi-K3 specialized decode
 
diff --git a/docs/features/nixl_connector_compatibility.md b/docs/features/nixl_connector_compatibility.md
index 05084e58966b..54462dd5b785 100644
--- a/docs/features/nixl_connector_compatibility.md
+++ b/docs/features/nixl_connector_compatibility.md
@@ -108,11 +108,10 @@ By default, a **compatibility hash** is checked during handshake. P and D instan
 
 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
+layout and raises at startup. Use the push connector
+(`NixlPushConnector`) for pipeline parallelism with hybrid KV caches.
+Both connectors pair KV regions with a peer's by the layers each region
+holds, 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:
@@ -120,7 +119,7 @@ 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.
-- Non-MLA attention-HMA layer routing requires decode TP to be no greater than prefill TP.
+- Non-MLA attention under HMA requires decode TP to be no greater than prefill TP.
 - Packed layouts under PP require MLA caches. Local and remote packed block strides may differ; layer pages must have equal sizes. MLA supports decode TP greater than prefill TP because its KV is replicated.
 
 ### Quantized KV cache
diff --git a/docs/features/nixl_connector_usage.md b/docs/features/nixl_connector_usage.md
index 4402095b2ec1..00f5860af4c2 100644
--- a/docs/features/nixl_connector_usage.md
+++ b/docs/features/nixl_connector_usage.md
@@ -63,6 +63,19 @@ vllm serve <MODEL> \
 !!! note
     Backend availability depends on how NIXL was built and what plugins are present in your environment. Refer to the [NIXL repository](https://github.com/ai-dynamo/nixl) for available backends and build instructions.
 
+#### UCX backend parameters
+
+`kv_connector_extra_config.ucx_backend_params` passes extra init parameters to NIXL's UCX backend. Values are strings. For example, NIXL creates UCX with `ucx_error_handling_mode=peer` by default; UCX then uses only transports that can detect peer failure, which can exclude intra-node GPU transports such as `rocm_ipc`. To trade that error handling for those transports:
+
+```bash
+vllm serve <MODEL> \
+  --kv-transfer-config '{
+    "kv_connector":"NixlConnector",
+    "kv_role":"kv_both",
+    "kv_connector_extra_config":{"ucx_backend_params":{"ucx_error_handling_mode":"none"}}
+  }'
+```
+
 ## Basic Usage (on the same host)
 
 ### Producer (Prefiller) Configuration
diff --git a/docs/features/preload.md b/docs/features/preload.md
index 5787fdb51c1c..92dab0418286 100644
--- a/docs/features/preload.md
+++ b/docs/features/preload.md
@@ -281,10 +281,18 @@ threat model.
 
 - Each GPU's socket path is guarded by an exclusive lock file, so a second
   daemon for the same GPU fails fast instead of clobbering the live socket.
-- `vllm preload` reports readiness once every rank is serving; if any rank
-  dies during startup, the remaining ranks are terminated and the command
-  exits with that rank's exit code.
+- `vllm preload` reports readiness once every rank is serving. If any rank
+  exits during startup or serving, the launcher terminates and reaps the
+  remaining ranks, stops its health endpoint, and exits with a nonzero status.
+  The failing rank's nonzero exit code is preserved. Signal failures use
+  `128 + signal number` (for example, 137 for `SIGKILL`); an unexpected clean
+  rank exit is reported as status 1. Restart the entire preload group before
+  attaching new engines.
 - `SIGINT`/`SIGTERM` to the `vllm preload` process terminates all daemon
-  ranks, which unlinks their sockets and releases the GPU memory.
+  ranks and exits with status 130/143.
+- In either shutdown path, ranks that do not stop within ten seconds are
+  killed and reaped, releasing their GPU memory. A forcibly
+  terminated rank may leave a stale socket; the next daemon replaces it
+  after acquiring the per-GPU lock.
 
 See [vllm preload](../cli/preload.md) for the full CLI reference.
diff --git a/docs/features/prompt_embeds.md b/docs/features/prompt_embeds.md
index 14b4e56afa5c..4721294a16c9 100644
--- a/docs/features/prompt_embeds.md
+++ b/docs/features/prompt_embeds.md
@@ -59,6 +59,8 @@ Prompt embeddings can be included as content parts in chat messages, interleaved
 
 Each `prompt_embeds` content part contains a `data` field with a base64-encoded `torch.Tensor` of shape `(num_tokens, hidden_size)`. Multiple `prompt_embeds` parts can appear in any message, in any position relative to text parts. The server expands each part into the correct number of placeholder tokens during chat template rendering, then splices the pre-computed embeddings into the model's input at the corresponding positions.
 
+For multimodal models, `prompt_embeds` parts can also be mixed with image content parts. Their placeholder token IDs are replaced with safe dummy IDs before model execution; the supplied embeddings replace the dummy embeddings at those positions.
+
 Unlike the Completions API, a `prompt_embeds` content part should encode **only** the content, not a templated conversation. The server wraps the chat template around the embedded content at request time, the same way it would for a plain text `content` string. Embedding a full templated conversation here would double-apply the template and produce incorrect inputs to the model.
 
 Assistant messages can also pass a single `prompt_embeds` part as `reasoning` instead of a string:
diff --git a/docs/features/quantization/README.md b/docs/features/quantization/README.md
index ba1a5b86fe45..87c7afe9cde6 100644
--- a/docs/features/quantization/README.md
+++ b/docs/features/quantization/README.md
@@ -83,7 +83,7 @@ th:not(:first-child) {
 - *Turing does not support Marlin MXFP4.
 
 !!! note
-    For information on quantization support on Google TPU, please refer to the [TPU-Inference Recommended Models and Features](https://docs.vllm.ai/projects/tpu/en/latest/recommended_models_features/) documentation.
+    For information on quantization support on Google TPU, please refer to the [TPU-Inference Feature Matrix](https://docs.vllm.ai/projects/tpu/en/latest/features/#quantization-support) documentation.
 
 !!! note
     This compatibility chart is subject to change as vLLM continues to evolve and expand its support for different hardware platforms and quantization methods.
diff --git a/docs/features/quantization/modelopt.md b/docs/features/quantization/modelopt.md
index 6d59a23565f0..36440e19c38c 100644
--- a/docs/features/quantization/modelopt.md
+++ b/docs/features/quantization/modelopt.md
@@ -49,12 +49,23 @@ module prefix and decoder layer indices before constructing quantized layers.
     use `--linear-backend flashinfer_trtllm` to select FlashInfer's TensorRT-LLM
     GEMM backend.
 
+For Muse Spark NVFP4 on NVIDIA, the prepared-input FlashInfer TRT-LLM path
+skips CUDA-graph padding rows when routing runs outside the expert kernel.
+The router writes expert IDs of `-1` and zero weights in its existing fused
+kernel. This requires a padding mask matching the routed token layout;
+other dispatch paths retain their backend-specific behavior.
+Set `VLLM_MOE_SKIP_PADDING=0` before starting the engine to disable this
+optimization for comparisons.
+
 ## Benchmarking with dummy weights
 
 For performance experiments with random weights, use `--load-format dummy`.
 ModelOpt NVFP4 MoE layers initialize positive quantization scales with matching
 gate/up global scales, including for FlashInfer MegaMoE
 (`--moe-backend flashinfer_moe_ep_cutedsl`).
+Muse Spark's NVIDIA `deep_gemm_mega_moe` backend also supports dummy loading:
+it initializes valid NVFP4 scales and finalizes its expert weights and residual
+constants before the first forward. Enable expert parallelism for this backend.
 Dummy weights are for benchmarking, not model quality evaluation. Normal
 checkpoint loading uses the checkpoint's scales.
 
diff --git a/docs/features/reasoning_outputs.md b/docs/features/reasoning_outputs.md
index 5e59644def9b..729325e29c36 100644
--- a/docs/features/reasoning_outputs.md
+++ b/docs/features/reasoning_outputs.md
@@ -308,6 +308,8 @@ curl http://localhost:8000/v1/chat/completions \
   }'
 ```
 
+The Responses API (`/v1/responses`) accepts the same `thinking_token_budget` field.
+
 ### Offline Inference
 
 ```python
diff --git a/docs/features/sleep_mode.md b/docs/features/sleep_mode.md
index bb6c6559228c..8c73dd16a624 100644
--- a/docs/features/sleep_mode.md
+++ b/docs/features/sleep_mode.md
@@ -242,3 +242,134 @@ curl -X POST 'http://localhost:8000/wake_up?tags=kv_cache'
 ## Limitation
 
 On ROCm, the virtual memory allocation on ROCm is done through chunked memory allocation. You can control the chunk size through `VLLM_ROCM_SLEEP_MEM_CHUNK_SIZE` (in MB). The default value is set at 256MB. The larger the chunk size the faster the performance. However, setting it too large will cause OOM. So if you encounter OOM when using sleep mode. Try reducing the chunk size. It is recommended to define the chunk size as a power of 2.
+
+## Experimental CUDA checkpoint executor
+
+This capability is experimental and disabled by default. Selecting the executor
+below checkpoints each CUDA worker after ordinary level-1 sleep,
+then restores all workers before waking allocator-managed memory. This preserves
+the remaining CUDA state, including captured graph executable objects, without
+recapturing graphs. Ordinary level-1 semantics still apply: KV cache contents are
+discarded. This is not a durable application checkpoint or model switching API.
+
+Set `VLLM_USE_V2_MODEL_RUNNER=1` before starting the engine:
+
+```python
+from vllm import LLM
+
+llm = LLM(
+    "Qwen/Qwen3-0.6B",
+    enable_sleep_mode=True,
+    distributed_executor_backend=(
+        "vllm.v1.executor.cuda_checkpoint_executor.CUDACheckpointExecutor"
+    ),
+)
+llm.sleep(level=1)
+llm.wake_up()
+```
+
+This experimental implementation requires Linux, NVIDIA driver checkpoint APIs,
+and dedicated worker processes. It uses an external helper to call
+`cuCheckpointProcessLock`, `cuCheckpointProcessCheckpoint`,
+`cuCheckpointProcessRestore`, and `cuCheckpointProcessUnlock`. CPU threads remain
+alive, but worker CUDA calls cannot execute while checkpointed. The helper uses
+the same process namespace and user as the executor. Driver permission and
+resource restrictions still apply; see the
+[CUDA checkpoint API documentation](https://docs.nvidia.com/cuda/archive/13.0.0/cuda-driver-api/group__CUDA__CHECKPOINT.html).
+
+The initial validation used NVIDIA H20 with driver 580.126.09, NCCL 2.29.7,
+and FlashInfer 0.6.18.post1. See [PR #58183](https://github.com/vllm-project/vllm/pull/58183)
+for versioned test results and communication performance comparisons. Support
+for other GPU, driver, and communication-library combinations is not established.
+
+Only single-node tensor parallelism is accepted. Tensor parallel sizes 1, 2,
+and 4 have been GPU-tested; other sizes remain unvalidated. Use
+`VLLM_USE_V2_MODEL_RUNNER=1`, the standard GPU worker and the `cumem` sleep
+backend. Pipeline, data, context and expert parallelism, external cache or
+weight transfer, custom worker classes, and sleep level 2 are rejected before
+worker startup (sleep level 2 is rejected when requested). Worker extensions
+that allocate additional CUDA IPC resources are outside the validated scope.
+Level 0 remains
+a scheduler pause and does not perform a CUDA checkpoint.
+
+For tensor parallel sizes greater than 1, use NCCL 2.29.7 or later with memory
+suspension support, set `enable_nccl_comm_suspend=True`, and pass
+`disable_custom_all_reduce=True`. The following explicit configuration retains
+NCCL GPU peer-to-peer or legacy shared-memory communication, selected by NCCL.
+`NCCL_P2P_DISABLE` and `NCCL_SHM_DISABLE` need not be set when communicator
+suspension is enabled. NCCL CUDA host-memory allocation remains disabled; it
+is not equivalent to legacy shared-memory transport. Before driver checkpoint,
+ordinary worker
+sleep suspends NCCL mappings; worker wake restores them at their reserved
+addresses. The executor checks each rank's actual suspension state.
+
+```bash
+export NCCL_IB_DISABLE=1
+export NCCL_CUMEM_ENABLE=1 NCCL_CUMEM_HOST_ENABLE=0 NCCL_NVLS_ENABLE=0
+export VLLM_ALLREDUCE_USE_SYMM_MEM=0 VLLM_USE_NCCL_SYMM_MEM=0
+export VLLM_ALLREDUCE_USE_FLASHINFER_PCIE_IPC=0
+export VLLM_ALLREDUCE_USE_FLASHINFER=1
+export VLLM_FLASHINFER_ALLREDUCE_BACKEND=trtllm
+```
+
+FlashInfer must provide stable-address workspace `checkpoint_prepare` and
+`checkpoint_restore` methods. The executor calls existing worker hooks around
+the driver checkpoint. Disabling standalone FlashInfer all-reduce with
+`VLLM_ALLREDUCE_USE_FLASHINFER=0` does not disable the compiler
+`fuse_allreduce_rms` pass, which can allocate FlashInfer workspaces independently.
+When that pass is enabled, the same backend restriction and workspace hooks
+apply. A configuration using only NCCL must also disable that compiler pass.
+The FlashInfer `mnnvl` backend is not accepted, and the NCCL suspend configuration
+requires NVLS to be disabled. Live multicast
+allocations can block the combined sleep/checkpoint flow even on newer
+user-mode drivers. Keep these restrictions when upgrading the driver. The
+helper must use the same user-mode driver as the workers; mixed versions are
+unsupported. Performance costs of the restrictions must be measured against
+the workload's default communication
+configuration; retaining GPU peer-to-peer communication alone does not establish
+unchanged inference performance.
+
+A more restrictive Socket-only configuration is also accepted without
+`enable_nccl_comm_suspend`: set `NCCL_P2P_DISABLE=1`, `NCCL_SHM_DISABLE=1`,
+`NCCL_CUMEM_ENABLE=0`, and
+`VLLM_ALLREDUCE_USE_FLASHINFER=0`, retaining the other communication restrictions
+above. This configuration is primarily useful for compatibility comparisons.
+Do not create additional CUDA IPC resources or fall back to an independently
+created PyTorch NCCL process group while using the checkpoint executor.
+
+Provide host memory for both the ordinary weight backup and the driver's
+remaining CUDA state backup. Restore needs sufficient free GPU memory; release
+any temporary GPU workload before waking. Checkpoint and restore add latency to
+ordinary sleep/wake. The helper itself can retain a small driver allocation, so
+worker offload does not imply a zero device-wide memory reading.
+
+All ranks are locked before any is checkpointed. All ranks are restored before
+any is unlocked. Operations within each stage run concurrently across workers.
+The whole stage has a 30-second response deadline (each lock also has a
+10-second driver timeout). Worker sleep/wake and FlashInfer workspace
+hooks have a 120-second response deadline. Failure terminates the workers rather than
+serving with a partially restored rank set; restart the engine after failure.
+Shutdown terminates the dedicated worker and helper processes without restoring
+their CUDA state or requiring free GPU memory for the model.
+
+Repeated level-1 sleep is idempotent. A first partial wake, such as
+`wake_up(tags=["weights"])`, restores the **whole CUDA process** before restoring
+the selected allocator tags. It cannot selectively restore graph or driver
+state. Ordinary worker RPCs are rejected while checkpointed. A checkpointed
+worker cannot execute another model until restored.
+
+Serialize sleep, wake, pause, and resume calls across all clients, including
+while level-1 sleep is draining requests. Wait for each call to complete before
+issuing the next. A partial wake does not automatically resume scheduling; wake
+all allocator tags before explicitly resuming or submitting new inference.
+Do not use an explicit resume call to bypass sleeping memory.
+
+If EngineCore exits during a management call, the client may remain waiting even
+after the workers terminate. Use an external timeout and restart the failed
+engine; a timeout does not establish that the operation completed. Do not retry
+control operations on an engine whose checkpoint outcome is unknown.
+
+As with the ordinary executor, `sleep()` is a no-op while weights remain asleep.
+For example, after `wake_up(tags=["kv_cache"])`, another `sleep()` does not
+checkpoint the process again. Wake the remaining weights before starting a new
+full sleep/checkpoint cycle.
diff --git a/docs/features/speculative_decoding/README.md b/docs/features/speculative_decoding/README.md
index 568308e4ef67..cf7d442b6f9a 100644
--- a/docs/features/speculative_decoding/README.md
+++ b/docs/features/speculative_decoding/README.md
@@ -16,7 +16,7 @@ vLLM supports a variety of methods of speculative decoding. Model-based methods
 - [N-Gram](n_gram.md)
 - [Suffix Decoding](suffix.md)
 - [Hidden State Extraction](extract_hidden_states.md)
-- [Custom Proposer Backend (Experimental)](#custom-proposer-backend-experimental)
+- [Custom Speculators (Experimental)](#custom-speculators-experimental)
 - [Dynamic Speculative Decoding](dynamic_speculative_decoding.md)
 - [Adaptive Verification](adaptive_verification.md)
 - [FLy Verification (Lossy)](fly.md)
@@ -43,10 +43,80 @@ For reproducible measurements in your environment, use
 [`examples/features/speculative_decoding/spec_decode_offline.py`](../../../examples/features/speculative_decoding/spec_decode_offline.py)
 or the [benchmark CLI guide](../../benchmarking/cli.md).
 
-## Custom Proposer Backend (Experimental)
+## Custom Speculators (Experimental)
+
+You can plug in your own speculator with the V2 model runner by registering it
+as a speculative method from a [vLLM plugin](../../design/plugin_system.md), and then using it
+with `"method": "<name>"`. Custom speculative methods require the V2 model
+runner.
+
+A custom speculator consists of:
+
+- The speculator itself: a subclass of
+  `vllm.v1.worker.gpu.spec_decode.speculator.BaseSpeculator`, typically of
+  `DraftModelSpeculator` (for a speculator that runs its own draft model, which
+  the model runner then loads and allocates draft KV cache for) or
+  `TargetDependentARSpeculator` (for EAGLE/MTP-style drafting from the target
+  model's hidden states). It's constructed with `(vllm_config, device)`.
+- Its config: a subclass of
+  `vllm.v1.worker.gpu.spec_decode.speculator_config.SpeculatorConfig`, which
+  names the speculator class and describes how it behaves (KV lookahead slots,
+  whether it needs the target's hidden states, which features it supports,
+  ...). It's used in every vLLM process, including those that never build the
+  speculator, so its module shouldn't import the speculator. A speculator that
+  behaves like a built-in one should subclass that speculator's config (e.g.
+  `Eagle3SpeculatorConfig`), which also determines how its draft model is
+  configured from `model`.
+
+```python
+# my_package/config.py
+from vllm.v1.worker.gpu.spec_decode.eagle.config import Eagle3SpeculatorConfig
+
+
+class MySpeculatorConfig(Eagle3SpeculatorConfig):
+    speculator_class = "my_package.speculator.MySpeculator"
+    # Accept options in `speculator_options` (validated in `apply_options`).
+    supported_options = frozenset({"speculator_options"})
+
+    @classmethod
+    def apply_options(cls, spec):
+        options = spec.speculator_options or {}
+        if unknown := options.keys() - {"my_option"}:
+            raise ValueError(f"Unknown speculator_options: {sorted(unknown)}")
+
+
+# my_package/__init__.py, registered as a `vllm.general_plugins` entry point
+def register():
+    from vllm.v1.worker.gpu.spec_decode.speculator_config import (
+        register_speculative_method,
+    )
+
+    register_speculative_method("my_spec", "my_package.config.MySpeculatorConfig")
+```
+
+```bash
+vllm serve <target-model> \
+  --speculative-config '{
+    "method": "my_spec",
+    "model": "<draft-model>",
+    "num_speculative_tokens": 3,
+    "speculator_options": {"my_option": 1}
+  }'
+```
+
+The speculator reads its options from
+`vllm_config.speculative_config.speculator_options`. Once set up,
+`speculative_config.custom_method` is the registered name (`"my_spec"`), and
+`speculative_config.method` is the built-in method its config is based on
+(`"eagle3"`), so vLLM treats it like that built-in speculator.
+
+### V1 Custom Proposer Backend
 
-You can plug in your own custom proposer class for speculative decoding by setting the method to `custom_class` and providing the full module path to your class.
-Your custom class must accept a `VllmConfig` upon instantiation and implement a `propose` method.
+With the V1 model runner, you can plug in your own custom proposer class by
+setting the method to `custom_class` and providing the full module path to your
+class. Your custom class must accept a `VllmConfig` upon instantiation and
+implement a `propose` method. This is not supported by the V2 model runner, and
+will be removed together with the V1 model runner.
 
 **Example configuration:**
 
diff --git a/docs/features/umbp_store_connector_usage.md b/docs/features/umbp_store_connector_usage.md
new file mode 100644
index 000000000000..fa6e1af69ba9
--- /dev/null
+++ b/docs/features/umbp_store_connector_usage.md
@@ -0,0 +1,75 @@
+# UMBPStoreConnector Usage Guide
+
+`UMBPStoreConnector` offloads KV cache blocks to host DRAM managed by UMBP,
+the unified memory buffer pool of [MORI](https://github.com/ROCm/mori). It
+restores them when a later request shares a prefix that has left GPU memory.
+
+## Requirements
+
+- MORI built with UMBP (`BUILD_UMBP=ON`), importable as `mori`.
+- Prefix caching enabled.
+
+## Embedded mode
+
+Embedded mode is the default. Each vLLM worker owns a private DRAM pool in its
+own process. The engine that stored a prefix can restore it; independent
+engines do not share these pools.
+
+```bash
+vllm serve <model> \
+  --enable-prefix-caching \
+  --kv-transfer-config '{
+    "kv_connector": "UMBPStoreConnector",
+    "kv_role": "kv_both",
+    "kv_load_failure_policy": "recompute",
+    "kv_connector_extra_config": {
+      "mode": "embedded",
+      "total_capacity_bytes": 68719476736
+    }
+  }'
+```
+
+### Embedded options
+
+Set these options in `kv_connector_extra_config`:
+
+| Option | Default | Meaning |
+| --- | --- | --- |
+| `capacity_bytes` | 64 GiB | Pool size per worker. |
+| `total_capacity_bytes` | | Pool size per engine, split evenly across its workers. Exclusive with `capacity_bytes`. |
+| `lookup_instance` | | Distinguishes independent engines serving the same model on one host. Data-parallel ranks are already distinguished. |
+| `lookup_dir` | `/tmp` | Directory of the per-worker lookup sockets. |
+| `dram_high_watermark`, `dram_low_watermark` | MORI defaults | Pool occupancy at which eviction starts and stops. |
+| `dram_use_hugepages`, `dram_hugepage_size`, `dram_numa_node`, `dram_prefault` | MORI defaults | Allocation of each worker's pool. |
+
+## Shared options
+
+Set these options in `kv_connector_extra_config`:
+
+| Option | Default | Meaning |
+| --- | --- | --- |
+| `load_async` | `true` | Load in the background while the request waits. Models with Mamba-style or sparse-attention indexer layers always load asynchronously, since those layers read restored state before a synchronous load could have landed. |
+| `lookup_async` | `false` | Look up the pool off the scheduler thread; the request waits until the lookup completes. |
+| `kv_event_max_objects` | `65536` | Maximum objects tracked for KV events. When full, new stored-object notifications are suppressed until tracked objects leave the pool. Storage and lookup continue normally. |
+| `enable_lookup` | `true` | Restore from the pool. With `false` the connector only stores. |
+| `save_decode_cache` | `false` | Also store blocks produced by decoding. |
+| `enable_partial_hash_hits` | `false` | Allow hits that end inside a block, at core's prefix-match unit (`cache_config.prefix_match_unit`). |
+| `num_workers`, `timeout_ms` | `4`, `30000` | Transfer threads per worker, and the transfer timeout. A synchronous wait that times out fails the step rather than reusing its blocks. |
+| `key_namespace` | derived | Overrides the namespace derived from the model, revision, KV layout, cache groups, parallel topology and draft model. |
+
+Use `kv_load_failure_policy: "recompute"` in `kv_transfer_config`, as in the
+example above. An object can be evicted between lookup and load; the default
+`fail` policy fails the request instead of recomputing it.
+
+With `kv_role: "kv_consumer"` the connector only restores and never stores.
+
+## Limitations
+
+- Pipeline parallelism is rejected.
+- Stores are submitted once per step after the forward pass.
+
+With KV events enabled, the scheduler verifies residency on every rank before
+publishing a stored-object event. It periodically rechecks tracked objects
+during engine steps to report removals. These checks do not delay the release
+of GPU source blocks after a completed transfer. With events disabled, no
+historical object identities are retained for event reporting.
diff --git a/docs/getting_started/installation/gpu.cuda.inc.md b/docs/getting_started/installation/gpu.cuda.inc.md
index 19a2432442fe..2166e1a856d2 100644
--- a/docs/getting_started/installation/gpu.cuda.inc.md
+++ b/docs/getting_started/installation/gpu.cuda.inc.md
@@ -458,11 +458,6 @@ BuildKit does not automatically invalidate cached layers when a mutable Git
 ref changes. Use `--no-cache-filter extensions-build` to refresh an empty,
 branch, or tag revision.
 
-Set `BUILD_NIXL=true` to build NIXL from source. NIXL's release wheels do not
-include the NIXL EP extension for the PyTorch nightly used by the Rubin build.
-The build uses the NIXL version pinned in `requirements/kv_connectors.txt` and
-replaces the NIXL packages installed from the KV-connector requirements.
-
 For `FINAL_BASE_IMAGE`, use the public, multi-arch
 `nvidia/cuda:13.4.1-base-ubuntu24.04` image. Set `NCCL_VERSION` to 2.32.3 or
 newer, the first NCCL release with Rubin (SM107) support.
diff --git a/docs/getting_started/installation/gpu.rocm.inc.md b/docs/getting_started/installation/gpu.rocm.inc.md
index d8debbab3ea3..69274c8ac350 100644
--- a/docs/getting_started/installation/gpu.rocm.inc.md
+++ b/docs/getting_started/installation/gpu.rocm.inc.md
@@ -1,7 +1,7 @@
 <!-- markdownlint-disable MD041 MD051 -->
 --8<-- [start:installation]
 
-vLLM supports AMD GPUs with ROCm 6.3 or above. Pre-built wheels are available for ROCm 10.0 (the default) and ROCm 7.2.
+vLLM supports AMD GPUs with ROCm 6.3 or above. Pre-built wheels are available for ROCm 10.0; ROCm 7.2 wheels remain available for earlier releases.
 
 #### Prebuilt Wheels
 
@@ -9,7 +9,7 @@ vLLM supports AMD GPUs with ROCm 6.3 or above. Pre-built wheels are available fo
 | ------------ | -------------- | ------------ | ----------------- | ------------------ |
 | `rocm700` | 3.12 | 7.0 | >= 2.35 | `0.14.0` to `0.18.0` |
 | `rocm721` | 3.12 | 7.2.1 | >= 2.35 | Nightly releases after commit `171775f306a333a9cf105bfd533bf3e113d401d9` |
-| `rocm72` | 3.12 | 7.2 | >= 2.39 | Nightly releases on the legacy ROCm 7.2 stack |
+| `rocm72` | 3.12 | 7.2 | >= 2.39 | Earlier nightly releases on the legacy ROCm 7.2 stack (no longer built) |
 | `rocm100` | 3.12 | 10.0 | >= 2.35 | Nightly releases (default) |
 
 --8<-- [end:installation]
@@ -51,7 +51,7 @@ uv pip install vllm --extra-index-url https://wheels.vllm.ai/rocm/ --upgrade
 !!! tip
     The ROCm 10.0 wheels install the ROCm SDK from pip ([TheRock](https://github.com/ROCm/TheRock)), including the GPU kernels for every supported architecture.
 
-    Wheels for the legacy ROCm 7.2 stack, which use a system ROCm installation, remain available at `https://wheels.vllm.ai/rocm/${VLLM_VERSION}/rocm72`.
+    Wheels for the legacy ROCm 7.2 stack, which use a system ROCm installation, remain available for earlier releases at `https://wheels.vllm.ai/rocm/${VLLM_VERSION}/rocm72`.
 
 !!! tip
     You can find out about which ROCm version the latest vLLM supports by checking the `vllm` package in index in extra-index-url <https://wheels.vllm.ai/rocm/> at [https://wheels.vllm.ai/rocm/vllm](https://wheels.vllm.ai/rocm/vllm) .
@@ -234,7 +234,6 @@ You can find more information about vLLM's wheels in
 
 !!! tip
     - The steps below follow [docker/Dockerfile.rocm_base](https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile.rocm_base), which pins the validated version of every component. If a step does not work for you, refer to it; a Dockerfile is a form of installation steps.
-    - The legacy ROCm 7.2 stack uses a system ROCm installation instead; its steps are in [docker/Dockerfile.rocm_72_base](https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile.rocm_72_base).
 
 0. Install the ROCm SDK and PyTorch (skip if you are already in an environment or docker image with them installed, e.g. `vllm/vllm-openai-rocm:base-nightly`).
 
@@ -371,7 +370,10 @@ The images can be used to run OpenAI compatible server and are available on Dock
 
 - `vllm/vllm-openai-rocm:latest` — stable release, built on ROCm 10.0 (also tagged `latest-rocm100`)
 - `vllm/vllm-openai-rocm:nightly` — preview build from the latest development branch on ROCm 10.0 (also tagged `nightly-rocm100`), use this if you want the latest features and fixes
-- `vllm/vllm-openai-rocm:latest-rocm72` / `nightly-rocm72` — the same builds on the legacy ROCm 7.2 stack, kept for a transition period
+- `vllm/vllm-openai-rocm:latest-rocm72` / `nightly-rocm72` — earlier builds on the legacy ROCm 7.2 stack, no longer updated
+- `vllm/vllm-openai-rocm:base-nightly` / `latest-base` (and their `-rocm100` forms) — the same image without vLLM itself: the ROCm stack plus all of vLLM's dependencies. Earlier builds of these tags held only the ROCm stack ([docker/Dockerfile.rocm_base](https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile.rocm_base))
+
+ROCm 10.0 images use zstd-compressed layers; pulling them needs Docker 23 or newer (or containerd 1.5 or newer).
 
 ```bash
 docker run --rm \
@@ -422,57 +424,19 @@ AMD also offered nightly prebuilt docker image from [Docker Hub](https://hub.doc
 --8<-- [end:pre-built-images]
 --8<-- [start:build-image-from-source]
 
-You can build and run vLLM from source via the provided [docker/Dockerfile.rocm](https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile.rocm).
-
-??? info "(Optional) Build an image with ROCm software stack"
-
-    Build a docker image from [docker/Dockerfile.rocm_base](https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile.rocm_base) which setup ROCm software stack needed by the vLLM.
-    **This step is optional as this rocm_base image is usually prebuilt and store at [Docker Hub](https://hub.docker.com/r/rocm/vllm-dev) under tag `rocm/vllm-dev:base` to speed up user experience.** For the legacy ROCm 7.2 stack, the base image is built from [docker/Dockerfile.rocm_72_base](https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile.rocm_72_base) and published as `vllm/vllm-openai-rocm:base-nightly-rocm72`.
-    If you choose to build this rocm_base image yourself, the steps are as follows.
+You can build and run vLLM from source via the provided [docker/Dockerfile.rocm](https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile.rocm), driven by [docker/docker-bake-rocm.hcl](https://github.com/vllm-project/vllm/blob/main/docker/docker-bake-rocm.hcl) (requires Docker Buildx).
+It uses ROCm 10.0 (installed from TheRock wheels, [docker/Dockerfile.rocm_base](https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile.rocm_base)). Older vLLM branches support ROCm 5.7 through 7.0.
 
-    It is important that the user kicks off the docker build using buildkit. Either the user put `DOCKER_BUILDKIT=1` as environment variable when calling docker build command, or the user needs to set up buildkit in the docker daemon configuration `/etc/docker/daemon.json` as follows and restart the daemon:
-
-    ```json
-    {
-        "features": {
-            "buildkit": true
-        }
-    }
-    ```
-
-    To build vllm on ROCm 10.0 for all supported GPUs, you can use the default:
-
-    ```bash
-    DOCKER_BUILDKIT=1 docker build \
-        -f docker/Dockerfile.rocm_base \
-        -t rocm/vllm-dev:base .
-    ```
-
-First, build a docker image from [docker/Dockerfile.rocm](https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile.rocm) and launch a docker container from the image.
-It is important that the user kicks off the docker build using buildkit. Either the user put `DOCKER_BUILDKIT=1` as environment variable when calling docker build command, or the user needs to set up buildkit in the docker daemon configuration /etc/docker/daemon.json as follows and restart the daemon:
-
-```json
-{
-    "features": {
-        "buildkit": true
-    }
-}
-```
-
-[docker/Dockerfile.rocm](https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile.rocm) uses ROCm 10.0 (installed from TheRock wheels) by default. The legacy ROCm 7.2 stack is available as [docker/Dockerfile.rocm_72](https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile.rocm_72) with [docker/Dockerfile.rocm_72_base](https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile.rocm_72_base); older vLLM branches support ROCm 5.7 through 7.0.
-It provides flexibility to customize the build of docker image using the following arguments:
-
-- `BASE_IMAGE`: specifies the base image used when running `docker build`. The default value `rocm/vllm-dev:base` is an image published and maintained by AMD. It is being built using [docker/Dockerfile.rocm_base](https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile.rocm_base). For [docker/Dockerfile.rocm_72](https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile.rocm_72), the default is `vllm/vllm-openai-rocm:base-nightly-rocm72`
-- `ARG_PYTORCH_ROCM_ARCH`: Allows to override the gfx architecture values from the base docker image
-
-Their values can be passed in when running `docker build` with `--build-arg` options.
-
-To build vllm on ROCm 10.0 for all supported GPUs, you can use the default (which build a docker image with `vllm serve` as entrypoint):
+To build vllm on ROCm 10.0 for all supported GPUs (an image with `vllm serve` as entrypoint), reusing the published ROCm stack and dependencies (`vllm/vllm-openai-rocm:base-nightly`) so that only vLLM itself is compiled:
 
 ```bash
-DOCKER_BUILDKIT=1 docker build -f docker/Dockerfile.rocm -t vllm/vllm-openai-rocm .
+STACK_IMAGE=vllm/vllm-openai-rocm:base-nightly RUNTIME=vllm/vllm-openai-rocm:base-nightly \
+    TAG=vllm/vllm-openai-rocm docker buildx bake -f docker/docker-bake-rocm.hcl vllm-openai
 ```
 
+Without `STACK_IMAGE` and `RUNTIME`, the ROCm stack and every dependency are built from source as well, which takes several hours.
+To compile vLLM only for your GPU, add for example `--set vllm-openai.args.ARCH=gfx942`.
+
 To run vLLM with the custom-built Docker image:
 
 ```bash
diff --git a/docs/models/extensions/tensorizer.md b/docs/models/extensions/tensorizer.md
index 1e0d5da81e6e..75639574ccfe 100644
--- a/docs/models/extensions/tensorizer.md
+++ b/docs/models/extensions/tensorizer.md
@@ -1,6 +1,6 @@
 # Loading models with CoreWeave's Tensorizer
 
-vLLM supports loading models with [CoreWeave's Tensorizer](https://docs.coreweave.com/coreweave-machine-learning-and-ai/inference/tensorizer).
+vLLM supports loading models with [CoreWeave's Tensorizer](https://github.com/coreweave/tensorizer).
 vLLM model tensors that have been serialized to disk, an HTTP/HTTPS endpoint, or S3 endpoint can be deserialized
 at runtime extremely quickly directly to the GPU, resulting in significantly
 shorter Pod startup times and CPU memory usage. Tensor encryption is also supported.
diff --git a/docs/models/pooling_models/token_embed.md b/docs/models/pooling_models/token_embed.md
index b0a09c0f887d..3823d5560527 100644
--- a/docs/models/pooling_models/token_embed.md
+++ b/docs/models/pooling_models/token_embed.md
@@ -82,6 +82,15 @@ If your model is not in the above list, we will try to automatically convert the
 
 jina-reranker-v3 is a listwise document reranker model with a novel `last but not late interaction` architecture. More information can be found at: [examples/pooling/token_embed/jina_reranker_v3_offline.py](../../../examples/pooling/token_embed/jina_reranker_v3_offline.py)
 
+For `jinaai/jina-reranker-v3`, `truncate_prompt_tokens` is unsupported. This applies to
+`LLM.encode(..., pooling_task="token_embed")`, `LLM.score`, `/pooling`, `/score`, and
+`/rerank`, even when the limit would leave the prompt unchanged. Requests that set
+it fail before inference; HTTP APIs return status 400 with
+`param="truncate_prompt_tokens"`. For scoring, use `max_tokens_per_query` and
+`max_tokens_per_doc` to truncate the query and documents separately. With `LLM.score`,
+pass these limits through `PoolingParams(extra_kwargs={...})`. For `/pooling` and
+`LLM.encode`, truncate each input text on the client before submitting the request.
+
 --8<-- [end:supported-token-embed-models]
 
 ## Offline Inference
diff --git a/docs/models/structured_decisions_models.md b/docs/models/structured_decisions_models.md
index 8d2d09f3b6d0..2ec032a3c3ba 100644
--- a/docs/models/structured_decisions_models.md
+++ b/docs/models/structured_decisions_models.md
@@ -200,6 +200,11 @@ The response adds a `joint` object, `null` when each question was read on its ow
 | `joint` | optional boolean, default `false`; see [Joint reads](#joint-reads). DiffusionGemma only, other models return 400 |
 | `images`, `audio`, `videos` | optional media URLs or data URLs, for models that take them ([Clef-Flash](#clef-flash), images/video; [Clef-Omni](#clef-omni); [Decision 3.0](#decision-30), images/video; DiffusionGemma, images only, with or without `joint`; [Liquid d1](#liquid-d1), images only; multimodal `NextTokenStrategy` models such as Qwen3.5, images only) |
 
+By default, `/v1/systemone` rejects `chat_template_kwargs.chat_template` with
+HTTP 400. Start the server with `--trust-request-chat-template` to allow
+request-provided chat templates. Other template kwargs, such as
+`enable_thinking`, do not require this flag.
+
 Unknown top-level fields, and question fields other than `type`,
 `instructions` and `criteria`, are rejected with a 400.
 
diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md
index 8520a3347fda..6719e661b8ff 100644
--- a/docs/models/supported_models.md
+++ b/docs/models/supported_models.md
@@ -200,6 +200,7 @@ Some model architectures are supported via vLLM plugins. These plugins extend vL
 | `BartForConditionalGeneration` | BART | [bart-plugin](https://github.com/vllm-project/bart-plugin) |
 | `Florence2ForConditionalGeneration` | Florence-2 | [bart-plugin](https://github.com/vllm-project/bart-plugin) |
 | `Nemotron3_5AsrForRNNT` | NVIDIA Nemotron 3.5 ASR | [vllm-nemotron-asr-plugin](https://github.com/Sohaib-Ahmed21/vllm-nemotron-asr-plugin) |
+| `ParakeetForTDT` | NVIDIA Parakeet TDT | [vllm-parakeet-plugin](https://github.com/Alex-Wengg/vllm-parakeet-plugin) |
 
 For other model architectures not natively supported, in particular for Encoder-Decoder models, we recommend following a similar pattern by implementing support through the plugin system.
 
@@ -451,7 +452,6 @@ th {
 | `NemotronForCausalLM` | Nemotron-3, Nemotron-4, Minitron | `nvidia/Minitron-8B-Base`, `mgoin/Nemotron-4-340B-Base-hf-FP8`, etc. | ✅︎ | ✅︎ |
 | `NemotronHForCausalLM` | Nemotron-H | `nvidia/Nemotron-H-8B-Base-8K`, `nvidia/Nemotron-H-47B-Base-8K`, `nvidia/Nemotron-H-56B-Base-8K`, etc. | ✅︎ | ✅︎ |
 | `OlmoHybridForCausalLM` | OLMo Hybrid | `allenai/Olmo-Hybrid-7B` | ✅︎ | ✅︎ |
-| `OlmoeForCausalLM` | OLMoE | `allenai/OLMoE-1B-7B-0924`, `allenai/OLMoE-1B-7B-0924-Instruct`, etc. | | ✅︎ |
 | `OPTForCausalLM` | OPT, OPT-IML | `facebook/opt-66b`, `facebook/opt-iml-max-30b`, etc. | ✅︎ | ✅︎ |
 | `OrionForCausalLM` | Orion | `OrionStarAI/Orion-14B-Base`, `OrionStarAI/Orion-14B-Chat`, etc. | | ✅︎ |
 | `PanguEmbeddedForCausalLM` | openPangu-Embedded-7B | `FreedomIntelligence/openPangu-Embedded-7B-V1.1` | ✅︎ | ✅︎ |
@@ -497,6 +497,7 @@ Some models are supported only via the [Transformers modeling backend](#transfor
 | `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. | ✅︎ | ✅︎ |
+| `OlmoeForCausalLM` | OLMoE | `allenai/OLMoE-1B-7B-0924`, `allenai/OLMoE-1B-7B-0924-Instruct`, etc. | ✅︎ | ✅︎ |
 | `PhiForCausalLM` | Phi | `microsoft/phi-1_5`, `microsoft/phi-2`, etc. | ✅︎ | ✅︎ |
 | `SeedOssForCausalLM` | SeedOss | `ByteDance-Seed/Seed-OSS-36B-Instruct`, etc. | ✅︎ | ✅︎ |
 | `SmolLM3ForCausalLM` | SmolLM3 | `HuggingFaceTB/SmolLM3-3B` | ✅︎ | ✅︎ |
@@ -581,8 +582,6 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
 | `InternVLChatModel` | InternVL 3.5, InternVL 3.0, InternVideo 2.5, InternVL 2.5, InternVL 2.0 | T + I<sup>E+</sup> + (V<sup>E+</sup>) | `OpenGVLab/InternVL3_5-14B`, `OpenGVLab/InternVL3-9B`, `OpenGVLab/InternVideo2_5_Chat_8B`, `OpenGVLab/InternVL2_5-4B`, `OpenGVLab/InternVL2-4B`, etc. | ✅︎ | ✅︎ |
 | `InternVLForConditionalGeneration` | InternVL 3.0 (HF format) | T + I<sup>E+</sup> + V<sup>E+</sup> | `OpenGVLab/InternVL3-1B-hf`, etc. | ✅︎ | ✅︎ |
 | `KananaVForConditionalGeneration` | Kanana-V | T + I<sup>+</sup> | `kakaocorp/kanana-1.5-v-3b-instruct`, etc. | | ✅︎ |
-| `KeyeForConditionalGeneration` | Keye-VL-8B-Preview | T + I<sup>E+</sup> + V<sup>E+</sup> | `Kwai-Keye/Keye-VL-8B-Preview` | ✅︎ | ✅︎ |
-| `KeyeVL1_5ForConditionalGeneration` | Keye-VL-1_5-8B | T + I<sup>E+</sup> + V<sup>E+</sup> | `Kwai-Keye/Keye-VL-1_5-8B` | ✅︎ | ✅︎ |
 | `KimiAudioForConditionalGeneration` | Kimi-Audio | T + A<sup>+</sup> | `moonshotai/Kimi-Audio-7B-Instruct` | | ✅︎ |
 | `KimiK25ForConditionalGeneration` | Kimi-K2.5 | T + I<sup>+</sup> | `moonshotai/Kimi-K2.5` | | ✅︎ |
 | `KimiK3ForConditionalGeneration` | Kimi-K3 | T + I<sup>+</sup> | `moonshotai/Kimi-K3` | | ✅︎ |
@@ -610,7 +609,9 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
 | `Moondream3ForCausalLM` | Moondream3 | T + I | `moondream/moondream3-preview` | | ✅︎ |
 | `MuseGlimmerForCausalLM`, `MuseGlimmerForConditionalGeneration` | Muse Glimmer | T + I<sup>+</sup> + V<sup>+</sup> | `meta-models/Muse-Glimmer-30B` | ✅︎ | ✅︎ |
 | `MuseSparkForConditionalGeneration` | Muse Spark Open | T + I<sup>+</sup> + V<sup>+</sup> + A<sup>+</sup> | `meta-models/Muse-Spark-Open-816B-A42B-NVFP4`, etc. | ✅︎ | ✅︎ |
+| `Nemotron3_5TranscribeForConditionalGeneration` | Nemotron-3.5-Transcribe | T + A | `nvidia/NVIDIA-Nemotron-3.5-Transcribe-30B-A4B`, `nvidia/NVIDIA-Nemotron-3.5-Transcribe-30B-A4B-FP8`, `nvidia/NVIDIA-Nemotron-3.5-Transcribe-30B-A4B-NVFP4` | | ✅︎ |
 | `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` | ✅︎ | |
+| `NemotronH_Omni_Reasoning_V3` | Nemotron 3.5 Super VL | T + I<sup>E+</sup> + V<sup>+</sup> | `nvidia/NVIDIA-Nemotron-3.5-Super-VL-120B-A12B-BF16`, `nvidia/NVIDIA-Nemotron-3.5-Super-VL-120B-A12B-NVFP4`, etc. | ✅︎ | |
 | `RadioMBartPrefixForCausalLM` | Nemotron Parse 3.5 | T + I | `nvidia/Nemotron-3.5-Parse` | | |
 | `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` | ✅︎ | ✅︎ |
@@ -723,10 +724,20 @@ Some models are supported only via the [Transformers modeling backend](#transfor
     its [DFlash](../features/speculative_decoding/README.md) path.
 
 !!! note
-    Both Nemotron Nano architecture names map to the same vLLM implementation.
+    Both Nemotron Nano architecture names and `NemotronH_Omni_Reasoning_V3` 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.
+    image and video only, as do the Nemotron 3.5 Super VL checkpoints.
+    `NemotronH_Omni_Reasoning_V3` checkpoints load without `--trust-remote-code`
+    on `transformers>=5.18.0`; older versions still require the flag.
+
+!!! note
+    On ROCm, `Step5ForConditionalGeneration` runs its MoE layers on Triton. With
+    `VLLM_ROCM_USE_AITER=1`, also pass `--moe-backend triton`: AITER's MoE does not
+    apply Step-5's SwiGLU clamp. An FP8 KV cache (`--kv-cache-dtype fp8`) needs
+    `--block-size 32`. The checkpoint ships three MTP layers, enabled with
+    `--speculative-config '{"method":"mtp","num_speculative_tokens":3}'`.
 
 !!! note
     The official `openbmb/MiniCPM-V-2` doesn't work yet, so we need to use a fork (`HwwwH/MiniCPM-V-2`) for now.
@@ -738,6 +749,7 @@ Speech2Text models trained specifically for Automatic Speech Recognition.
 
 | Architecture | Models | Example HF Models | [LoRA](../features/lora.md) | [PP](../serving/parallelism_scaling.md) |
 | ------------ | ------ | ----------------- | -------------------- | ------------------------- |
+| `CanaryForConditionalGeneration` | Canary | `nvidia/canary-1b-v2` | | |
 | `CohereAsrForConditionalGeneration` | Cohere-Transcribe | `CohereLabs/cohere-transcribe-03-2026` | | |
 | `FireRedASR2ForConditionalGeneration` | FireRedASR2 | `allendou/FireRedASR2-LLM-vllm`, etc. | | |
 | `FunASRForConditionalGeneration` | FunASR | `FunAudioLLM/Fun-ASR-Nano-2512-vllm`, etc. | | |
@@ -745,12 +757,33 @@ Speech2Text models trained specifically for Automatic Speech Recognition.
 | `GlmAsrForConditionalGeneration` | GLM-ASR | `zai-org/GLM-ASR-Nano-2512` | ✅︎ | ✅︎ |
 | `GraniteSpeechForConditionalGeneration` | Granite Speech | `ibm-granite/granite-4.0-1b-speech`, `ibm-granite/granite-speech-3.3-2b`, etc. | ✅︎ | ✅︎ |
 | `GraniteSpeechPlusForConditionalGeneration` | Granite Speech Plus | `ibm-granite/granite-speech-4.1-2b-plus` | ✅︎ | ✅︎ |
+| `MingV3ASRForConditionalGeneration` | Ming-3.0-ASR | `inclusionAI/Ming-3.0-asr-tiny` | | |
 | `MossTranscribeDiarizeForConditionalGeneration` | MOSS-Transcribe-Diarize | `OpenMOSS-Team/MOSS-Transcribe-Diarize` | | ✅︎ |
+| `Nemotron3_5TranscribeForConditionalGeneration` | Nemotron-3.5-Transcribe | `nvidia/NVIDIA-Nemotron-3.5-Transcribe-30B-A4B`, `nvidia/NVIDIA-Nemotron-3.5-Transcribe-30B-A4B-FP8`, `nvidia/NVIDIA-Nemotron-3.5-Transcribe-30B-A4B-NVFP4` | | ✅︎ |
 | `Qwen3ASRForConditionalGeneration` | Qwen3-ASR | `Qwen/Qwen3-ASR-1.7B`, etc. | ✅︎ | ✅︎ |
 | `Qwen3OmniMoeThinkerForConditionalGeneration` | Qwen3-Omni | `Qwen/Qwen3-Omni-30B-A3B-Instruct`, etc. | | ✅︎ |
 | `VoxtralForConditionalGeneration` | Voxtral (Mistral format) | `mistralai/Voxtral-Mini-3B-2507`, `mistralai/Voxtral-Small-24B-2507`, etc. | ✅︎ | ✅︎ |
 | `WhisperForConditionalGeneration` | Whisper | `openai/whisper-small`, `openai/whisper-large-v3-turbo`, etc. | | |
 
+!!! note
+    `Nemotron3_5TranscribeForConditionalGeneration` transcribes with speaker tags. Its chat template turns
+    thinking on by default, so transcription servers should set
+    `--default-chat-template-kwargs '{"enable_thinking": false}'`. It supports MTP and DFlash2 speculative
+    decoding.
+
+    Long audio: by default vLLM refuses audio longer than 600 s and transcription uploads larger than 25 MB
+    (`VLLM_MAX_AUDIO_DECODE_DURATION_S`, `VLLM_MAX_AUDIO_CLIP_FILESIZE_MB`). A chat request takes one clip of
+    up to one hour, the model's training limit (`--limit-mm-per-prompt '{"audio": {"length": <samples>}}'`
+    changes it), and `/v1/audio/transcriptions` splits longer audio into pieces of at most one hour, shorter
+    when `--max-model-len` is below 65,536.
+
+    Audio tower: the FP8 and NVFP4 checkpoints quantize the linear layers of its ASR encoder to FP8 through
+    their `quantization_config`. FP8 encoder attention (`--mm-encoder-attn-dtype fp8`) is not supported and
+    raises an error when the model is built. The tower attends through vLLM's multimodal encoder attention
+    backend (FlashAttention on CUDA GPUs that support it; `--mm-encoder-attn-backend` selects another). In
+    bfloat16, its output can depend on that backend and on the other clips in the batch: small rounding
+    differences can flip the thresholded speaker features on some frames.
+
 !!! note
     `VoxtralForConditionalGeneration` requires `mistral-common[audio]` to be installed.
 
diff --git a/docs/pre_run_check.sh b/docs/pre_run_check.sh
index 04f4a642ef07..9ccfc85b35e9 100644
--- a/docs/pre_run_check.sh
+++ b/docs/pre_run_check.sh
@@ -29,7 +29,7 @@ DOCS_PATHS=(
   requirements/docs.txt       # Affects build process
   requirements/docs.in        # Affects build process
 )
-if git diff --quiet origin/main -- "${DOCS_PATHS[@]}"; then
+if git diff --quiet origin/main...HEAD -- "${DOCS_PATHS[@]}"; then
   echo "No docs-affecting files changed vs origin/main; cancelling build."
   # See https://docs.readthedocs.com/platform/latest/guides/build/skip-build.html for info on exit code
   exit 183
diff --git a/docs/serving/expert_parallel_deployment.md b/docs/serving/expert_parallel_deployment.md
index f3b329610216..0c6d08560eb0 100644
--- a/docs/serving/expert_parallel_deployment.md
+++ b/docs/serving/expert_parallel_deployment.md
@@ -29,6 +29,7 @@ vLLM provides multiple communication backends for EP. Use `--all2all-backend` to
 | `flashinfer_nvlink_one_sided` | MNNVL systems | FlashInfer's one-sided A2A strategy for multi-node NVLink | High-throughput workloads |
 | `passthrough` | Any | Internal, not accepted on the CLI: bound automatically for MoE backends that dispatch and combine themselves, so no all-to-all runs in the framework | `--moe-backend flashinfer_moe_ep_*` (FlashInfer MoE-EP megakernels) |
 | `flashinfer_nvlink_two_sided` | MNNVL systems | FlashInfer's two-sided A2A strategy for multi-node NVLink | Systems with NVLink across nodes |
+| `mori` | ROCm | MoRI EP: EPv2 (dispatch/combine over a CCO symmetric-memory window) where supported, otherwise EPv1 | AMD Instinct; see [MoRI EP](#mori-ep-rocm) |
 
 ## Single Node Deployment
 
@@ -80,6 +81,48 @@ vllm serve deepseek-ai/DeepSeek-V3-0324 \
     --enable-expert-parallel         # Enable expert parallelism
 ```
 
+### MoRI EP (ROCm)
+
+`--all2all-backend mori` moves MoE tokens with MoRI. It runs MoRI EPv2 (`mori.ops.dispatch_combine_v2`, which reads and writes peer GPU memory directly through a MoRI CCO symmetric-memory window) where EPv2 is supported, and MoRI EPv1 otherwise.
+
+```bash
+VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_MOE=1 \
+VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=0 \
+vllm serve deepseek-ai/DeepSeek-V4-Pro \
+    --tensor-parallel-size 1 \
+    --data-parallel-size 8 \
+    --enable-expert-parallel \
+    --all2all-backend mori
+```
+
+`VLLM_MORI_EP_V2` selects the MoRI version:
+
+| `VLLM_MORI_EP_V2` | Behavior |
+| ----------------- | -------- |
+| unset | EPv2, falling back to EPv1 with a warning when MoRI is older than 1.2.4, the GPU is not gfx950 or gfx1250, EP spans hosts (except on gfx1250, whose UALink team can span hosts), `MORI_ENABLE_SDMA=1`, `--moe-backend aiter_mega_moe`, the model is not BF16, or MoE activations use per-tensor FP8 |
+| `1` | Always EPv2. The checks above only log a warning; startup fails if the EP ranks are not in one CCO-LSA team |
+| `0` | EPv1 |
+
+On EPv1, `mori` uses the IntraNode kernel within one node. Across nodes it uses InterNodeV1LL when the per-rank token budget (`--max-num-batched-tokens`) is at most `VLLM_MORI_DISPATCH_INTER_KERNEL_SWITCH_THRESHOLD` (default 256), and InterNodeV1 otherwise. `mori_high_throughput` and `mori_low_latency` always run EPv1 and pin the inter-node kernel to InterNodeV1 and InterNodeV1LL respectively.
+
+Choosing a MoRI backend: use EPv2 wherever it applies, and EPv1 for deployments EPv2 does not support. EPv2 cuts its receive buffers to the rows a step can actually receive and lets AITER fused MoE write its output straight into the combine buffer, so it does less work per step than EPv1.
+
+| Deployment | Backend |
+| ---------- | ------- |
+| EP within one node (or one gfx1250 UALink fabric), BF16 model, AITER fused MoE | `mori` (runs EPv2) |
+| EP within one node, with experts that take activations quantized before dispatch (e.g. FP8 or INT8 Triton experts) | EPv2 rejects these experts at startup. Enable AITER fused MoE (`VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_MOE=1`) so `mori` runs EPv2 with AITER experts; MoRI EPv1 (`VLLM_MORI_EP_V2=0`) also needs AITER fused MoE. To keep Triton experts without AITER, use a non-MoRI backend such as `allgather_reducescatter` |
+| EP across nodes | `mori` (falls back to EPv1; picks InterNodeV1LL when the per-rank token budget is at most `VLLM_MORI_DISPATCH_INTER_KERNEL_SWITCH_THRESHOLD`, else InterNodeV1) |
+
+Within one node, all EPv1 options run the same IntraNode kernel. To pin the inter-node kernel, set the threshold: `0` always picks InterNodeV1, and a value at least the per-rank token budget always picks InterNodeV1LL. `mori_high_throughput` and `mori_low_latency` do the same pinning but are deprecated and will be removed in a future release; use `mori` instead.
+
+EPv2 limitations:
+
+- All EP ranks must be in one MoRI CCO-LSA team, i.e. GPUs that can load and store each other's memory: one node, or one UALink fabric. Startup fails otherwise; there is no multi-node RDMA path yet.
+- Activations are dispatched in BF16, so the model dtype must be BF16 and the experts must take unquantized inputs (AITER fused MoE, which quantizes after the all-to-all, or BF16 experts). Experts that expect activations quantized before dispatch, such as FP8 or INT8 Triton experts, are rejected at startup. Block FP8 and MXFP4 weights work; per-tensor FP8 activations (`--quantization fp8` without blocks, ModelOpt FP8) do not, since AITER cannot quantize the received rows per tensor.
+- AITER shared-expert fusion is not supported: set `VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=0`.
+- Install a MoRI build that ships `mori.ops.dispatch_combine_v2`. On NICs that limit a registered memory region to 2 GiB, the window registration fails once the dispatch arena outgrows that limit (DeepSeek-V4-Pro at EP8 already does at the default 16k batched tokens).
+- Older MoRI builds leave a stale HIP error behind when the fabric-handle probe fails (e.g. on a single MI355X node), and the next PyTorch allocation raises `operation not supported`.
+
 ## Multi-Node Deployment
 
 For multi-node deployment, use the DeepEP communication kernel with one of two modes (see [Backend Selection Guide](#backend-selection-guide) above).
@@ -156,7 +199,7 @@ Configure EPLB with the `--eplb-config` argument, which accepts a JSON string. T
 | `num_redundant_experts` | Additional global experts per EP rank beyond equal distribution | `0` |
 | `use_async` | Use non-blocking EPLB for reduced latency overhead | `true` |
 | `policy` | The policy type for expert parallel load balancing | `"default"` |
-| `communicator` | Backend for expert weight transfers: `"torch_nccl"`, `"torch_gloo"`, `"pynccl"`, `"nixl"`,  or `null` (auto) | `null` |
+| `communicator` | Backend for expert weight transfers: `"torch_nccl"`, `"torch_gloo"`, `"pynccl"`, `"nixl"`, `"symm_mem"`, or `null` (auto) | `null` |
 
 For example:
 
@@ -189,6 +232,8 @@ EPLB uses redundant experts that need to fit in GPU memory. This means that EPLB
 This overhead equals `NUM_MOE_LAYERS * BYTES_PER_EXPERT * (NUM_TOTAL_EXPERTS + NUM_REDUNDANT_EXPERTS) ÷ NUM_EP_RANKS`.
 For DeepSeekV3, this is approximately `2.4 GB` for one redundant expert per EP rank.
 
+The `symm_mem` communicator requires all EP ranks in one NVLink domain (MNNVL across nodes). It allocates its one-layer receive buffer in symmetric memory instead of regular memory, and senders write into it directly, so it uses no more memory than the other backends.
+
 ### Example Command
 
 Single node deployment with EPLB enabled:
@@ -319,3 +364,44 @@ except Exception as e:
 - To simulate the decode deployment of disaggregated serving, pass `--kv-transfer-config '{"kv_connector":"DecodeBenchConnector","kv_role":"kv_both"}'` to the `vllm serve` invocation. The connector populates KV cache with random values so decode can be profiled in isolation. Add `"kv_connector_extra_config":{"startup_fill":true}` to fill the whole KV cache once at startup instead of in the decode steps that admit new requests (see the connector docstring for its limitations).
 
 - **CUDAGraph capture**: Use `--compilation_config '{"cudagraph_mode": "FULL_DECODE_ONLY"}'` to enable CUDA graph capture for decode only and save KV cache.
+
+## Fault tolerance and scale-down recovery
+
+Scale-down reloads every surviving physical replica using the recovered EPLB
+placement. The fused GPT-OSS MXFP4 checkpoint loader is not supported for this
+reload path; recovery rejects it rather than loading weights into the wrong slots.
+
+Fault tolerance lets vLLM recover from certain distributed engine failures without
+restarting the entire deployment. An external orchestrator detects failures and
+coordinates recovery across the engines.
+
+### Configuring recovery
+
+Both retry and scale-down require `--enable-fault-tolerance`,
+`--enable-expert-parallel`, an FT-capable `--all2all-backend` (`nixl_ep` or
+`deepep_low_latency`), and the v2 model runner (`VLLM_USE_V2_MODEL_RUNNER=1`).
+Retry can recover from a fault when all workers remain alive.
+
+For scale-down recovery, also enable `--enable-eplb` and configure enough
+redundant expert slots with `--eplb-config.num_redundant_experts` to fit all
+logical experts on the surviving ranks. The original checkpoint must remain
+accessible to survivors so they can reload expert weights.
+The orchestrator selects the recovery mode when it sends the recovery request.
+
+### Coordinating recovery
+
+Send recovery instructions to `POST /v1/fault_tolerance/apply`, then poll
+`GET /v1/fault_tolerance/status` on each participating engine. Use a new,
+nonempty `request_id` for each recovery round and send the same ID to every
+engine in that round. An HTTP 202 response means the instruction was accepted;
+wait until all participating engines report `healthy` before resuming traffic.
+
+Use `retry` when all workers are still alive. To remove a failed DP rank, stop
+routing traffic to it and send `scale_down` to every surviving engine once they
+report `unhealthy`. Include the original rank IDs in `removed_dp_ranks`. If the
+removed rank hosted the DP store, also provide a reachable surviving
+`dp_master_ip` and a free `dp_store_port`.
+
+Scale-down restores service with fewer ranks; it does not restart the removed
+engines or restore their capacity. The previous `/fault_tolerance/apply`
+endpoint is no longer registered.
diff --git a/docs/serving/integrations/claude_code.md b/docs/serving/integrations/claude_code.md
index 99a89a076769..6f8f8f387a49 100644
--- a/docs/serving/integrations/claude_code.md
+++ b/docs/serving/integrations/claude_code.md
@@ -77,4 +77,6 @@ If the model responds correctly, your setup is working. You can now use Claude C
 
 **Tool calls not working**: Verify that your model supports tool calling and that you've enabled it with the correct `--tool-call-parser` flag. See [Tool Calling](../../features/tool_calling.md).
 
+**WebSearch fails with `tool type 'web_search_...' is not supported`**: Claude Code's WebSearch tool relies on Anthropic's server-side web search tool, which runs on Anthropic's servers. vLLM does not run server tools and rejects requests that include them with a 400. Use an MCP server that provides web search instead.
+
 **Model not found**: Ensure the `--served-model-name` matches the model names in your environment variables. You cannot use model names with `/` in them, such as `openai/gpt-oss-120b` directly from Huggingface, so beware of that limitation with Claude Code.
diff --git a/docs/serving/online_serving/derenderer.md b/docs/serving/online_serving/derenderer.md
index 65b1d0daf8cf..09e96a4ba205 100644
--- a/docs/serving/online_serving/derenderer.md
+++ b/docs/serving/online_serving/derenderer.md
@@ -8,6 +8,8 @@ This closes the loop for a token-in / token-out engine in disaggregated serving:
 - **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
 - **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))
 
+For clients that already use `GenerateRequest`, `output_mode: "derender"` on `/inference/v1/generate` parses inline and supports `stream: true` with no per-chunk hop (see [Token In Token Out](token_in_token_out.md)).
+
 Both endpoints are hosted by the GPU less rendering server started with [`vllm launch render`](../../cli/launch/render.md), alongside the `/render`
 endpoints.
 
diff --git a/docs/serving/online_serving/openai_compatible_server.md b/docs/serving/online_serving/openai_compatible_server.md
index 7b4e84a17fee..b87903412d22 100644
--- a/docs/serving/online_serving/openai_compatible_server.md
+++ b/docs/serving/online_serving/openai_compatible_server.md
@@ -174,6 +174,8 @@ see our [Multimodal Inputs](../../features/multimodal_inputs.md) guide for more
 
 Code example: [examples/basic/online_serving/openai_chat_completion_client.py](../../../examples/basic/online_serving/openai_chat_completion_client.py)
 
+With `--enable-prompt-tokens-details`, `/v1/chat/completions/batch` sums `cached_tokens`, `created_cache_tokens` and per-modality `multimodal_tokens` across conversations in `usage.prompt_tokens_details`; an unreported cache count makes that total `null`, and disabling the flag leaves the details `null`.
+
 #### Extra parameters
 
 The following [sampling parameters](../../api/README.md#inference-parameters) are supported.
diff --git a/docs/serving/online_serving/token_in_token_out.md b/docs/serving/online_serving/token_in_token_out.md
index f71c5e78fb83..6e4d78bb4c69 100644
--- a/docs/serving/online_serving/token_in_token_out.md
+++ b/docs/serving/online_serving/token_in_token_out.md
@@ -93,9 +93,17 @@ print(response["choices"][0]["token_ids"])
 - `parse_context` is ignored at other levels, so a rendered request can go to a prefill leg unchanged.
 - `logprobs` follow coupled chat. `top_logprobs` is cut to `sampling_params.logprobs`, `logprob_token_ids` is honored. When a parser hides the reasoning with `include_reasoning: false`, `logprobs`, `token_ids` and `sampling_mask` are dropped so they don't expose it.
 - `finish_reason` follows coupled chat too. It's `tool_calls` when the message has tool calls except with a named `tool_choice` or a `length` cut.
-- `stream: true` isn't supported yet.
+- With `stream: true`, each choice carries a chat `delta` (`role`, `content`, `reasoning`, `tool_calls`) instead of `text`, and `finish_reason` follows coupled chat. `role` is set on the first chunk of each choice, `n > 1` is supported, and `null` fields are left out of `delta`. A chunk is sent whenever the engine output carries a parsed delta, visible token IDs, logprobs or a `finish_reason`. Parsing runs on the API server event loop, so `--api-server-count` scales it.
 - Cohere grounding citations aren't returned. The parser strips the citation spans from `content`, and the document map can't go over JSON in `parse_context`.
 
+A streamed tool call looks like this, with the usage chunk and `data: [DONE]` framed as at the other levels:
+
+```text
+data: {"request_id": "...", "output_mode": "derender", "choices": [{"index": 0, "token_ids": [151657], "delta": {"role": "assistant", "tool_calls": [{"index": 0, "id": "chatcmpl-tool-...", "type": "function", "function": {"name": "get_weather", "arguments": ""}}]}, "finish_reason": null, ...}]}
+
+data: {"request_id": "...", "output_mode": "derender", "choices": [{"index": 0, "token_ids": [], "delta": {}, "finish_reason": "tool_calls", ...}]}
+```
+
 ### Logprobs
 
 With `output_mode: "tokens"`, logprobs are `GenerateLogProbs` with an integer `token_id` and `rank` per entry and no `token` or `bytes` (see [Generate Output Logprobs](renderer.md#generate-output-logprobs)). With `output_mode: "text"`, they are `ChatCompletionLogProbs` carrying the decoded token strings and their UTF-8 `bytes` for the sampled token and every entry in `top_logprobs`; a server started with `--return-tokens-as-token-ids` returns `token_id:N` placeholders there instead, as `/v1/completions` does.
@@ -106,7 +114,6 @@ The server returns a 400 for:
 
 - `output_mode: "text"` or `"derender"` on a server without a tokenizer (`--tokens-only` or `--skip-tokenizer-init`).
 - `output_mode: "text"` or `"derender"` with `sampling_params.detokenize: false`.
-- `output_mode: "derender"` with `stream: true`.
 - `output_mode: "derender"` without `parse_context` on a server with a tool or reasoning parser or with a named `tool_choice` in `parse_context` that matches none of its `tools`. This is checked before generating.
 - An unsupported `output_mode` value.
 
diff --git a/docs/serving/online_serving/trace_replay.md b/docs/serving/online_serving/trace_replay.md
index 86ed13644c7a..9acd09b57ad6 100644
--- a/docs/serving/online_serving/trace_replay.md
+++ b/docs/serving/online_serving/trace_replay.md
@@ -44,6 +44,7 @@ When `trace_decode_token_ids` is set:
 - The trace is truncated if it does not fit within `max_model_len` given the prompt length.
 - All stop conditions (EOS, stop strings, stop token IDs) are disabled.
 - Generation produces exactly the trace tokens, then stops.
+- With speculative decoding, a draft token is accepted only if it equals the trace token at its position, so a step emits the matching drafts plus one trace token, up to the end of the trace. The returned logprobs are those of the trace tokens, but each verification step scores several positions in one forward pass, so they are bitwise equal to the logprobs without speculative decoding only with batch invariance (`VLLM_BATCH_INVARIANT=1`).
 
 ## Limitations
 
@@ -53,7 +54,6 @@ When `trace_decode_token_ids` is set:
 | --- | --- |
 | `n > 1` | Trace replay produces a single deterministic sequence |
 | `prompt_logprobs` | Expanded logit layout conflicts with trace kernel indexing |
-| Speculative decoding | Multi-token speculation conflicts with single-token trace stepping |
 | Structured outputs | Grammar constraints conflict with forced token injection |
 | `repetition_detection` | Would terminate the request on repeated patterns in the trace |
 | `thinking_token_budget` | Logit masking corrupts logprobs for trace tokens |
diff --git a/docs/training/weight_transfer/README.md b/docs/training/weight_transfer/README.md
index 9c946eba0f13..73b56c6cc2fc 100644
--- a/docs/training/weight_transfer/README.md
+++ b/docs/training/weight_transfer/README.md
@@ -131,10 +131,11 @@ parameter is usually itself a collective among its owners (an FSDP
 `full_tensor()` all-gather, a Megatron export) that would deadlock if some of
 them skipped it. Which ranks own which parameter is declared in
 [`ParamMeta.owners`](base.md#parammeta-ownership-and-layout).
-Only [sharded RDT](sharded_rdt.md) accepts ranks that own different parameters,
-such as pipeline stages. Pipeline-parallel support for NCCL and IPC is in
-progress; until then they, like NCCL M2N, need every parameter owned by rank 0
-or by every rank (see each backend's Limitations).
+Only [sharded RDT](sharded_rdt.md) and [NCCL](nccl.md#broadcasting-ranks)
+accept ranks that own different parameters, such as pipeline stages.
+Pipeline-parallel support for IPC is in progress; until then it, like NCCL M2N,
+needs every parameter owned by rank 0 or by every rank (see each backend's
+Limitations).
 
 Pass each process's own `rank` on the init info. It is explicit rather than read
 from a global process group, which is ambiguous once several groups (FSDP / TP /
diff --git a/docs/training/weight_transfer/nccl.md b/docs/training/weight_transfer/nccl.md
index 8ada8ed870f9..279ba3f21e93 100644
--- a/docs/training/weight_transfer/nccl.md
+++ b/docs/training/weight_transfer/nccl.md
@@ -10,7 +10,7 @@ The NCCL weight transfer engine uses [NCCL](https://developer.nvidia.com/nccl) b
 
 ## How It Works
 
-1. The trainer and all inference workers join a shared NCCL process group using `StatelessProcessGroup` (vLLM's torch.distributed-independent group abstraction). The trainer is rank 0; workers start at `rank_offset` (1).
+1. The trainer's broadcasting ranks and all inference workers join a shared NCCL process group using `StatelessProcessGroup` (vLLM's torch.distributed-independent group abstraction). Trainer rank 0 is group rank 0; with [several broadcasters](#broadcasting-ranks) they take `[0, rank_offset)` and workers follow them (`rank_offset` is 1 otherwise).
 2. The trainer broadcasts weights to all workers simultaneously. Each worker receives and loads the weights.
 3. Optionally, **packed tensor broadcasting** batches multiple small tensors into larger buffers with double/triple buffering and CUDA stream overlap for higher throughput. This implementation is based on [NeMo-RL's packed tensor](https://github.com/NVIDIA-NeMo/RL/blob/main/nemo_rl/utils/packed_tensor.py).
 
@@ -52,7 +52,7 @@ engine = WeightTransferTrainerFactory.trainer_init(
     init_info=NCCLTrainerInitInfo(
         master_address=master_address,
         master_port=master_port,
-        world_size=world_size,   # trainer + all inference workers
+        world_size=world_size,   # broadcasting trainer ranks + all inference workers
         rank=0,                  # this trainer rank; rank 0 is the sender
         packed=True,
     ),
@@ -66,7 +66,8 @@ engine.send_weights()   # once per sync
 `trainer_init` drives the whole handshake: it kicks off the inference side's
 `init_weight_transfer_engine` (through the client) on a side thread while opening
 the trainer's own rank-0 endpoint, since both ends must rendezvous together. It
-builds the worker's init info itself, with `rank_offset=1` and the same packed
+builds the worker's init info itself, with `rank_offset` set to the number of
+broadcasting ranks (1 unless the source is partitioned) and the same packed
 params, so the two sides cannot disagree.
 
 `send_weights()` then drives `start_weight_update`, `update_weights` — run
@@ -79,7 +80,7 @@ once every transfer has drained.
 | ----- | ------- | ----------- |
 | `master_address` | — | Rendezvous host |
 | `master_port` | — | Rendezvous port |
-| `world_size` | — | Full trainer + worker NCCL group size |
+| `world_size` | — | Transfer NCCL group size: the [broadcasting trainer ranks](#broadcasting-ranks) plus the inference workers (`1 + num_workers` for a single sender). Other trainer ranks are not counted |
 | `rank` | — | Keyword-only. This trainer process's rank; **0 is the sender** |
 | `packed` | `True` | Use packed broadcasting |
 | `packed_buffer_size_bytes` | 1 GiB | Packed buffer size |
@@ -124,13 +125,37 @@ rather than letting it reach the wire. `ModuleSource` satisfies this by
 construction. If you write a custom source — a Megatron export, an MoE re-fusing
 pass — this is the invariant to test first.
 
-### One sender
-
-Each parameter is sent by its first owner, and this backend has a single
-sender, trainer rank 0. Every parameter's `owners` must therefore start at rank
-0 or be unset; a source whose pipeline stages own different parameters is
-rejected at `trainer_init`. Other owners (FSDP peers of rank 0) only
-drive `iter_full()` to take part in the gathers.
+### Broadcasting ranks
+
+Each parameter is broadcast by its first owner (rank 0 when `owners` is unset).
+When that is rank 0 for every parameter, rank 0 is the only broadcaster. Other
+owners (FSDP peers of rank 0) only drive `iter_full()` to take part in the
+gathers.
+
+A partitioned source, such as pipeline stages declaring `owners=(stage
+ranks...)`, has several broadcasters. They join the transfer group at group
+ranks `[0, rank_offset)` in trainer-rank order, rank 0 always included. Every
+rank derives this plan from `metadata`, which is identical on every rank, so
+the broadcasters need no extra coordination. In `send_weights()`, rank 0 sends
+one `update_weights` per broadcaster, in order, each tagged with that
+broadcaster's group rank as `src`. Each broadcaster sends its own parameters
+during its turn and receives and discards the others' to stay in the
+collective. Initialize and call `send_weights()` on every trainer rank
+concurrently.
+
+Every parameter a rank owns must share its first owner, since each broadcaster
+gathers and sends in its own turn. So with several broadcasters every parameter
+must declare `owners`. Rank 0 is always in the transfer group and only gathers
+in its own turn, so if it owns a parameter it must be that parameter's first
+owner. `trainer_init` raises on every rank otherwise.
+
+A failure on any broadcaster during an update is fatal for the transfer. That
+rank raises immediately, but the other trainer ranks and the inference workers
+can stay blocked in that update's broadcasts, so tear down the trainer job and
+restart or re-initialize the inference engine before transferring again.
+
+[`rlhf_nccl_multi_source.py`](../../../examples/rl/rlhf_nccl_multi_source.py)
+shows two pipeline stages updating one tensor-parallel inference engine.
 
 ### Metadata is sent once, then reused
 
@@ -142,7 +167,7 @@ round. So the worker remembers the list:
 | `update_info` you send | What the worker loads |
 | ---------------------- | --------------------- |
 | `names`, `dtype_names`, `shapes` filled in | Exactly these parameters. The worker also saves this list. |
-| All three lists empty (or `{}`) | The list saved by the previous `update_weights` call on the same model. |
+| All three lists empty (or `{}`) | The list saved by the previous `update_weights` call on the same model from the same `src`. |
 | All three lists empty, nothing saved yet | Nothing. |
 
 Rules to remember:
@@ -195,6 +220,9 @@ while rank 0 broadcasts in `names` order, then `POST /finish_weight_update`.
 checkpoint shapes. From the second sync on, send `{"update_info": {}}` if the
 parameters did not change; see
 [Metadata is sent once, then reused](#metadata-is-sent-once-then-reused).
+An optional `src` (default `0`, below `rank_offset`) selects the trainer rank
+that broadcasts that update, so one start/finish lifecycle can carry updates
+from several trainer ranks.
 
 - Unpacked: one broadcast per tensor.
 - Packed: concatenate `t.contiguous().view(-1).view(torch.uint8)` in `names`
@@ -275,11 +303,9 @@ Sparse NCCL sends only `O(nnz)` indices and values. Checkpoint application still
 
 ## Limitations
 
-- **Pipeline-parallel trainers are not supported yet (work in progress).** Rank 0
-  is the only sender, so it must own every parameter (see
-  [One sender](#one-sender)). A source whose ranks own different parameters is
-  rejected at `trainer_init`.
-- **Full tensors on the wire.** Rank 0 materializes each parameter in full and
+- **A broadcaster failure is fatal.** See
+  [Broadcasting ranks](#broadcasting-ranks).
+- **Full tensors on the wire.** Each broadcaster materializes its parameters in full and
   every worker receives all of it, whatever its tensor-parallel shard. For
   large or expert-parallel models, see [sharded RDT](sharded_rdt.md) or
   [NCCL M2N](m2n.md).
@@ -288,6 +314,7 @@ Sparse NCCL sends only `O(nnz)` indices and values. Checkpoint application still
 ## Examples
 
 - [RLHF with NCCL weight syncing (`vllm serve`, HTTP)](../../../examples/rl/rlhf_http_nccl.py) - **Start here.** Trainer on one GPU, 2x tensor-parallel fp8 server on two others; HTTP control plane, NCCL data plane. Launches and tears down its own server
+- [Multi-source NCCL weight syncing](../../../examples/rl/rlhf_nccl_multi_source.py) - Two pipeline stages each broadcast their own parameters to one TP2 inference engine within a single lifecycle
 - [NCCL weight syncing from a trainer that does not import vLLM](../../../examples/rl/rlhf_http_nccl_no_vllm.py) - `torch` + `requests` only, `nccl_unique_id_b64` rendezvous; see [Trainer Without vLLM over HTTP](#trainer-without-vllm-over-http)
 - [RLHF with NCCL + FSDP2 and expert parallelism](../../../examples/rl/rlhf_nccl_fsdp_ep.py) - Multi-rank trainer: every FSDP rank builds an engine and joins the `full_tensor()` gather while only rank 0 touches the wire
 - [RLHF with sparse NCCL weight syncing (offline, Ray)](../../../examples/rl/rlhf_sparse_nccl.py) - Qwen3 MoE per-expert checkpoint updates with one trainer GPU and two TP2/EP2 inference GPUs
diff --git a/docs/usage/security.md b/docs/usage/security.md
index ed10de055890..88f248a38f1e 100644
--- a/docs/usage/security.md
+++ b/docs/usage/security.md
@@ -324,6 +324,7 @@ vLLM enforces several configurable limits to prevent this:
 | `VLLM_P2P_MAX_PEERS` | `128` | Maximum concurrent P2P peer sessions. New peers (both outbound and inbound) are rejected when this limit is reached. |
 | `VLLM_P2P_HANDSHAKE_TIMEOUT_S` | `30` | Seconds to wait for a newly connected peer to complete the P2P handshake. Sessions that remain connected but not ready after this deadline are torn down. |
 | `VLLM_P2P_IDLE_TIMEOUT_S` | `300` | Seconds a P2P session may remain idle (no pending work, no messages) before it is eligible for eviction. Idle eviction runs only when remaining peer slots are within 8 of `VLLM_P2P_MAX_PEERS`. Set to `0` to disable idle eviction. |
+| `VLLM_P2P_TRANSFER_STALL_TIMEOUT_S` | `20` | Seconds a P2P WRITE may stay in progress before its peer is treated as faulted: no new WRITEs to it, and the session is torn down. The stalled transfer's memory stays held until it stops. Set to `0` to disable. |
 
 Connections that fail (ZMQ socket errors, handshake failures) are handled gracefully — the request falls back to local prefill instead of crashing the engine.
 
@@ -494,7 +495,7 @@ driver-side credentials from worker-side processes within the same Ray
 cluster.
 
 This assumption is consistent with
-[Ray's own security model](https://docs.ray.io/en/latest/ray-core/security.html),
+[Ray's own security model](https://docs.ray.io/en/latest/ray-security/index.html),
 which states that any user who can connect to a Ray cluster can run arbitrary
 code on any node in that cluster. In other words, Ray cluster access already
 implies full code execution on worker nodes, so restricting environment
diff --git a/docs/usage/v1_guide.md b/docs/usage/v1_guide.md
index d16db468955f..bdf3b4b2a139 100644
--- a/docs/usage/v1_guide.md
+++ b/docs/usage/v1_guide.md
@@ -139,6 +139,7 @@ Whisper is supported natively. Other encoder-decoder models are supported via th
 - **BART**: `BartForConditionalGeneration` is supported via the official [bart-plugin](https://github.com/vllm-project/bart-plugin).
 - **Florence-2**: `Florence2ForConditionalGeneration` is supported via the official [bart-plugin](https://github.com/vllm-project/bart-plugin).
 - **Nemotron 3.5 ASR**: `Nemotron3_5AsrForRNNT` is supported via [vllm-nemotron-asr-plugin](https://github.com/Sohaib-Ahmed21/vllm-nemotron-asr-plugin).
+- **Parakeet TDT**: `ParakeetForTDT` is supported via [vllm-parakeet-plugin](https://github.com/Alex-Wengg/vllm-parakeet-plugin).
 
 For other encoder-decoder models (e.g., `MllamaForConditionalGeneration`), we recommend
 following a similar pattern by implementing support through the [plugin system](../design/plugin_system.md).
diff --git a/examples/features/openai_batch/README.md b/examples/features/openai_batch/README.md
index c1bb32937688..fff4acc2e601 100644
--- a/examples/features/openai_batch/README.md
+++ b/examples/features/openai_batch/README.md
@@ -13,9 +13,9 @@ The OpenAI batch file format consists of a series of json objects on new lines.
 Each line represents a separate request. See the [OpenAI package reference](https://platform.openai.com/docs/api-reference/batch/requestInput) for more details.
 
 ```{note}
-We currently support `/v1/chat/completions`, `/v1/embeddings`, `/score`, `/rerank`,
-`/v1/audio/transcriptions`, and `/v1/audio/translations` endpoints (`/v1/completions`
-coming soon). Score and rerank match any URL ending in `/score` or `/rerank`, such as
+We currently support `/v1/chat/completions`, `/v1/completions`, `/v1/embeddings`,
+`/score`, `/rerank`, `/v1/audio/transcriptions`, and `/v1/audio/translations`
+endpoints. Score and rerank match any URL ending in `/score` or `/rerank`, such as
 `/v1/score`.
 ```
 
diff --git a/examples/generate/multimodal/audio_language_offline.py b/examples/generate/multimodal/audio_language_offline.py
index 8b2215d8c1eb..5866101973d0 100644
--- a/examples/generate/multimodal/audio_language_offline.py
+++ b/examples/generate/multimodal/audio_language_offline.py
@@ -251,6 +251,26 @@ def run_midashenglm(question: str, audio_count: int):
     )
 
 
+# Ming3-ASR
+def run_ming_v3_asr(question: str, audio_count: int) -> ModelRequestData:
+    from vllm.model_executor.models.ming_v3_asr import ASR_INSTRUCTION
+
+    assert audio_count == 1, "Ming3-ASR only supports single audio input per prompt"
+    engine_args = EngineArgs(
+        model="inclusionAI/Ming-3.0-asr-tiny",
+        trust_remote_code=True,
+        dtype="bfloat16",
+        max_model_len=8192,
+        max_num_seqs=4,
+        limit_mm_per_prompt={"audio": audio_count},
+        enforce_eager=True,
+    )
+    return ModelRequestData(
+        engine_args=engine_args,
+        prompt=ASR_INSTRUCTION + "<audio><audioPatch></audio>",
+    )
+
+
 # MiniCPM-O
 def run_minicpmo(question: str, audio_count: int) -> ModelRequestData:
     model_name = "openbmb/MiniCPM-o-2_6"
@@ -503,6 +523,7 @@ def run_whisper(question: str, audio_count: int) -> ModelRequestData:
     "granite_speech": run_granite_speech,
     "kimi_audio": run_kimi_audio,
     "midashenglm": run_midashenglm,
+    "ming_v3_asr": run_ming_v3_asr,
     "minicpmo": run_minicpmo,
     "phi4_mm": run_phi4mm,
     "qwen2_audio": run_qwen2_audio,
diff --git a/examples/generate/multimodal/vision_language_multi_image_offline.py b/examples/generate/multimodal/vision_language_multi_image_offline.py
index 7565f1e42b51..d829503750dc 100644
--- a/examples/generate/multimodal/vision_language_multi_image_offline.py
+++ b/examples/generate/multimodal/vision_language_multi_image_offline.py
@@ -454,80 +454,6 @@ def load_internvl(question: str, image_urls: list[str]) -> ModelRequestData:
     )
 
 
-def load_keye_vl(question: str, image_urls: list[str]) -> ModelRequestData:
-    model_name = "Kwai-Keye/Keye-VL-8B-Preview"
-
-    engine_args = EngineArgs(
-        model=model_name,
-        trust_remote_code=True,
-        max_model_len=8192,
-        max_num_seqs=5,
-        limit_mm_per_prompt={"image": len(image_urls)},
-    )
-
-    placeholders = [{"type": "image", "image": url} for url in image_urls]
-    messages = [
-        {
-            "role": "user",
-            "content": [
-                *placeholders,
-                {"type": "text", "text": question},
-            ],
-        },
-    ]
-
-    processor = AutoProcessor.from_pretrained(model_name, trust_remote_code=True)
-
-    prompt = processor.apply_chat_template(
-        messages, tokenize=False, add_generation_prompt=True
-    )
-
-    image_data = [fetch_image(url) for url in image_urls]
-
-    return ModelRequestData(
-        engine_args=engine_args,
-        prompt=prompt,
-        image_data=image_data,
-    )
-
-
-def load_keye_vl1_5(question: str, image_urls: list[str]) -> ModelRequestData:
-    model_name = "Kwai-Keye/Keye-VL-1_5-8B"
-
-    engine_args = EngineArgs(
-        model=model_name,
-        trust_remote_code=True,
-        max_model_len=32768,
-        max_num_seqs=5,
-        limit_mm_per_prompt={"image": len(image_urls)},
-    )
-
-    placeholders = [{"type": "image", "image": url} for url in image_urls]
-    messages = [
-        {
-            "role": "user",
-            "content": [
-                *placeholders,
-                {"type": "text", "text": question},
-            ],
-        },
-    ]
-
-    processor = AutoProcessor.from_pretrained(model_name, trust_remote_code=True)
-
-    prompt = processor.apply_chat_template(
-        messages, tokenize=False, add_generation_prompt=True
-    )
-
-    image_data = [fetch_image(url) for url in image_urls]
-
-    return ModelRequestData(
-        engine_args=engine_args,
-        prompt=prompt,
-        image_data=image_data,
-    )
-
-
 def load_kimi_vl(question: str, image_urls: list[str]) -> ModelRequestData:
     model_name = "moonshotai/Kimi-VL-A3B-Instruct"
 
@@ -1349,8 +1275,6 @@ def load_molmo2(question: str, image_urls: list[str]) -> ModelRequestData:
     "idefics3": load_idefics3,
     "interns1": load_interns1,
     "internvl_chat": load_internvl,
-    "keye_vl": load_keye_vl,
-    "keye_vl1_5": load_keye_vl1_5,
     "kimi_vl": load_kimi_vl,
     "llama4": load_llama4,
     "llava": load_llava,
diff --git a/examples/generate/multimodal/vision_language_offline.py b/examples/generate/multimodal/vision_language_offline.py
index c1448f1cbe87..02602f3e4fda 100644
--- a/examples/generate/multimodal/vision_language_offline.py
+++ b/examples/generate/multimodal/vision_language_offline.py
@@ -21,6 +21,7 @@
 from vllm.lora.request import LoRARequest
 from vllm.multimodal.image import convert_image_mode
 from vllm.platforms import current_platform
+from vllm.transformers_utils.chat_templates import get_chat_template_fallback_path
 from vllm.utils.argparse_utils import FlexibleArgumentParser
 
 
@@ -935,80 +936,6 @@ def run_kanana_v(questions: list[str], modality: str) -> ModelRequestData:
     )
 
 
-# Keye-VL
-def run_keye_vl(questions: list[str], modality: str) -> ModelRequestData:
-    model_name = "Kwai-Keye/Keye-VL-8B-Preview"
-
-    mm_limit = {"image": 1, "video": 1} if modality == "image+video" else {modality: 1}
-    engine_args = EngineArgs(
-        model=model_name,
-        max_model_len=8192,
-        trust_remote_code=True,
-        limit_mm_per_prompt=mm_limit,
-    )
-
-    image_placeholder = "<|vision_start|><|image_pad|><|vision_end|>"
-    video_placeholder = "<|vision_start|><|video_pad|><|vision_end|>"
-
-    if modality == "image":
-        placeholder = image_placeholder
-    elif modality == "video":
-        placeholder = video_placeholder
-    elif modality == "image+video":
-        placeholder = image_placeholder + video_placeholder
-
-    prompts = [
-        (
-            f"<|im_start|>user\n{placeholder}"
-            f"{question}<|im_end|>\n"
-            "<|im_start|>assistant\n"
-        )
-        for question in questions
-    ]
-
-    return ModelRequestData(
-        engine_args=engine_args,
-        prompts=prompts,
-    )
-
-
-# Keye-VL-1.5
-def run_keye_vl1_5(questions: list[str], modality: str) -> ModelRequestData:
-    model_name = "Kwai-Keye/Keye-VL-1.5-8B"
-
-    mm_limit = {"image": 1, "video": 1} if modality == "image+video" else {modality: 1}
-    engine_args = EngineArgs(
-        model=model_name,
-        max_model_len=8192,
-        trust_remote_code=True,
-        limit_mm_per_prompt=mm_limit,
-    )
-
-    image_placeholder = "<|vision_start|><|image_pad|><|vision_end|>"
-    video_placeholder = "<|vision_start|><|video_pad|><|vision_end|>"
-
-    if modality == "image":
-        placeholder = image_placeholder
-    elif modality == "video":
-        placeholder = video_placeholder
-    elif modality == "image+video":
-        placeholder = image_placeholder + video_placeholder
-
-    prompts = [
-        (
-            f"<|im_start|>user\n{placeholder}"
-            f"{question}<|im_end|>\n"
-            "<|im_start|>assistant\n"
-        )
-        for question in questions
-    ]
-
-    return ModelRequestData(
-        engine_args=engine_args,
-        prompts=prompts,
-    )
-
-
 # Kimi-VL
 def run_kimi_vl(questions: list[str], modality: str) -> ModelRequestData:
     assert modality == "image"
@@ -1528,6 +1455,36 @@ def run_openpangu_vl(questions: list[str], modality: str) -> ModelRequestData:
     )
 
 
+# OpenVLA
+def run_openvla(questions: list[str], modality: str) -> ModelRequestData:
+    assert modality == "image"
+
+    model_name = "openvla/openvla-7b"
+    chat_template = get_chat_template_fallback_path("openvla", model_name)
+    assert chat_template is not None
+
+    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+    messages = [{"role": "user", "content": "pick up the red block"}]
+    prompt = tokenizer.apply_chat_template(
+        messages,
+        chat_template=chat_template.read_text(),
+        tokenize=False,
+    ).removeprefix(tokenizer.bos_token)
+    prompts = [prompt for _ in questions]
+    engine_args = EngineArgs(
+        model=model_name,
+        trust_remote_code=True,
+        max_model_len=512,
+        limit_mm_per_prompt={modality: 1},
+    )
+
+    return ModelRequestData(
+        engine_args=engine_args,
+        prompts=prompts,
+        sampling_params=SamplingParams(temperature=0.0, max_tokens=7),
+    )
+
+
 # Ovis
 def run_ovis(questions: list[str], modality: str) -> ModelRequestData:
     assert modality == "image"
@@ -2235,8 +2192,6 @@ def run_step_vl(questions: list[str], modality: str) -> ModelRequestData:
     "interns1_pro": run_interns1_pro,
     "internvl_chat": run_internvl,
     "kanana_v": run_kanana_v,
-    "keye_vl": run_keye_vl,
-    "keye_vl1_5": run_keye_vl1_5,
     "kimi_vl": run_kimi_vl,
     "kimi_k25": run_kimi_k25,
     "lightonocr": run_lightonocr,
@@ -2255,6 +2210,7 @@ def run_step_vl(questions: list[str], modality: str) -> ModelRequestData:
     "nemotron_vl": run_nemotron_vl,
     "NVLM_D": run_nvlm_d,
     "openpangu_vl": run_openpangu_vl,
+    "openvla": run_openvla,
     "ovis": run_ovis,
     "ovis2_5": run_ovis2_5,
     "paddleocr_vl": run_paddleocr_vl,
diff --git a/examples/pooling/embed/openai_embedding_long_text/README.md b/examples/pooling/embed/openai_embedding_long_text/README.md
index 2ed04f1d91f2..cafe4975a58b 100644
--- a/examples/pooling/embed/openai_embedding_long_text/README.md
+++ b/examples/pooling/embed/openai_embedding_long_text/README.md
@@ -46,7 +46,7 @@ The key parameters for chunked processing are in the `--pooler-config`:
 
 ```json
 {
-  "pooling_type": "auto",
+  "pooling_type": "MEAN",
   "use_activation": true,
   "enable_chunked_processing": true,
   "max_embed_len": 3072000
@@ -54,7 +54,7 @@ The key parameters for chunked processing are in the `--pooler-config`:
 ```
 
 !!! note
-    `pooling_type` sets the model's own pooling strategy for processing within each chunk. The cross-chunk aggregation automatically uses MEAN strategy when input exceeds the model's native maximum length.
+    `pooling_type` sets the model's own pooling strategy for processing within each chunk and must be one of `CLS`, `LAST` or `MEAN`. `service.sh` accepts `POOLING_TYPE=auto` and resolves it to one of these before building the config. The cross-chunk aggregation automatically uses MEAN strategy when input exceeds the model's native maximum length.
 
 #### Chunked Processing Behavior
 
diff --git a/examples/pooling/embed/openai_embedding_long_text/service.sh b/examples/pooling/embed/openai_embedding_long_text/service.sh
index 68950d04ee14..1c68b00f9177 100644
--- a/examples/pooling/embed/openai_embedding_long_text/service.sh
+++ b/examples/pooling/embed/openai_embedding_long_text/service.sh
@@ -20,8 +20,7 @@ API_KEY=${API_KEY:-"your-api-key"}
 
 # Enhanced pooling configuration with model-specific defaults
 POOLING_TYPE=${POOLING_TYPE:-"auto"}  # auto, MEAN, CLS, LAST
-export VLLM_ENABLE_CHUNKED_PROCESSING=true
-export CUDA_VISIBLE_DEVICES=2,3,4,5
+ENABLE_CHUNKED_PROCESSING=true
 
 echo "🚀 Starting vLLM Embedding Server with Enhanced Chunked Processing"
 echo "=================================================================="
@@ -68,7 +67,7 @@ echo "📋 Configuration:"
 echo "   - Model: $MODEL_NAME"
 echo "   - Port: $PORT"
 echo "   - GPU Count: $GPU_COUNT"
-echo "   - Enhanced Chunked Processing: ${VLLM_ENABLE_CHUNKED_PROCESSING}"
+echo "   - Enhanced Chunked Processing: ${ENABLE_CHUNKED_PROCESSING}"
 echo "   - Max Embed Length: ${MAX_EMBED_LEN} tokens"
 echo "   - Native Pooling Type: $POOLING_TYPE + Normalization"
 echo "   - Cross-chunk Aggregation: MEAN (automatic)"
@@ -96,7 +95,7 @@ echo ""
 echo "🔧 Starting server with enhanced chunked processing configuration..."
 
 # Build pooler config JSON
-POOLER_CONFIG="{\"pooling_type\": \"$POOLING_TYPE\", \"use_activation\": true, \"enable_chunked_processing\": ${VLLM_ENABLE_CHUNKED_PROCESSING}, \"max_embed_len\": ${MAX_EMBED_LEN}}"
+POOLER_CONFIG="{\"pooling_type\": \"$POOLING_TYPE\", \"use_activation\": true, \"enable_chunked_processing\": ${ENABLE_CHUNKED_PROCESSING}, \"max_embed_len\": ${MAX_EMBED_LEN}}"
 
 # Start vLLM server with enhanced chunked processing
 vllm serve "$MODEL_NAME" \
diff --git a/examples/rl/rlhf_nccl_multi_source.py b/examples/rl/rlhf_nccl_multi_source.py
new file mode 100644
index 000000000000..730d4443528c
--- /dev/null
+++ b/examples/rl/rlhf_nccl_multi_source.py
@@ -0,0 +1,207 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Dense NCCL weight transfer from two disjoint trainer sources.
+
+The example uses four GPUs:
+
+* GPUs 0-1 run two trainer actors. Each actor exposes the parameters owned by
+  one logical pipeline stage.
+* GPUs 2-3 run one tensor-parallel vLLM engine.
+
+Each stage declares itself the owner of its parameters in `ParamMeta.owners`,
+so both trainer ranks join the weight-transfer communicator and each
+broadcasts its own parameters. Rank 0 drives the control plane: one
+`update_weights` per stage, within one start/finish lifecycle.
+
+For simplicity, both trainer actors load the small model before filtering their
+parameter streams. A real pipeline-parallel trainer would already hold only its
+stage-local parameters; the transfer behavior is the same.
+
+This example was written for four GPUs with at least 8 GiB each.
+"""
+
+import os
+import time
+from functools import cached_property
+
+import ray
+import torch
+from huggingface_hub import snapshot_download
+from transformers import AutoModelForCausalLM
+
+from vllm import LLM, SamplingParams
+from vllm.config import WeightTransferConfig
+from vllm.distributed.weight_transfer import (
+    ParamMeta,
+    RayVLLMWeightSyncClient,
+    WeightSource,
+    WeightTransferTrainerFactory,
+)
+from vllm.distributed.weight_transfer.nccl_engine import NCCLTrainerInitInfo
+from vllm.utils.network_utils import get_ip, get_open_port
+
+MODEL_NAME = "Qwen/Qwen3-0.6B"
+NUM_SOURCES = 2
+INFERENCE_TP_SIZE = 2
+
+
+class StageSource(WeightSource):
+    """Expose the checkpoint parameters assigned to one logical stage."""
+
+    def __init__(
+        self,
+        model: torch.nn.Module,
+        stage_rank: int,
+        num_stages: int,
+        num_layers: int,
+    ) -> None:
+        self.model = model
+        self.stage_rank = stage_rank
+        self.num_stages = num_stages
+        self.num_layers = num_layers
+
+    def _owner(self, name: str) -> int:
+        parts = name.split(".")
+        for index, part in enumerate(parts[:-1]):
+            if part == "layers" and parts[index + 1].isdigit():
+                layer = int(parts[index + 1])
+                return min(
+                    layer * self.num_stages // self.num_layers, self.num_stages - 1
+                )
+        return 0 if "embed_tokens" in name else self.num_stages - 1
+
+    @cached_property
+    def metadata(self) -> list[ParamMeta]:
+        return [
+            ParamMeta(
+                name, param.dtype, tuple(param.shape), owners=(self._owner(name),)
+            )
+            for name, param in self.model.named_parameters()
+        ]
+
+    def iter_full(self):
+        for name, param in self.model.named_parameters():
+            if self._owner(name) == self.stage_rank:
+                yield name, param
+
+
+class MyLLM(LLM):
+    """Allow the vLLM Ray executor to allocate its own worker GPUs."""
+
+    def __init__(self, *args, **kwargs):
+        os.environ.pop("CUDA_VISIBLE_DEVICES", None)
+        super().__init__(*args, **kwargs)
+
+    def ready(self):
+        return True
+
+
+@ray.remote(num_cpus=1, num_gpus=1)
+class TrainerSource:
+    def __init__(self, model_path: str, source_rank: int):
+        self.source_rank = source_rank
+        torch.accelerator.set_device_index(0)
+        self.model = AutoModelForCausalLM.from_pretrained(
+            model_path, torch_dtype=torch.bfloat16
+        ).cuda()
+        self.engine = None
+
+    def ready(self):
+        return self.source_rank
+
+    def setup_engine(
+        self,
+        llm_handle,
+        master_address: str,
+        master_port: int,
+        world_size: int,
+    ):
+        source = StageSource(
+            self.model,
+            self.source_rank,
+            NUM_SOURCES,
+            self.model.config.num_hidden_layers,
+        )
+        self.engine = WeightTransferTrainerFactory.trainer_init(
+            init_info=NCCLTrainerInitInfo(
+                master_address=master_address,
+                master_port=master_port,
+                world_size=world_size,
+                rank=self.source_rank,
+                packed=True,
+            ),
+            client=RayVLLMWeightSyncClient(llm_handle),
+            source=source,
+        )
+
+    def send_weights(self):
+        self.engine.send_weights()
+
+
+def generate(llm, prompts: list[str]) -> list[list[int]]:
+    outputs = ray.get(
+        llm.generate.remote(
+            prompts,
+            SamplingParams(temperature=0, max_tokens=16),
+        )
+    )
+    return [output.outputs[0].token_ids for output in outputs]
+
+
+def main():
+    ray.init()
+    model_path = snapshot_download(MODEL_NAME)
+
+    sources = [
+        TrainerSource.remote(model_path, source_rank)
+        for source_rank in range(NUM_SOURCES)
+    ]
+    ray.get([source.ready.remote() for source in sources])
+
+    llm = ray.remote(num_cpus=0, num_gpus=0)(MyLLM).remote(
+        model=model_path,
+        enforce_eager=True,
+        tensor_parallel_size=INFERENCE_TP_SIZE,
+        distributed_executor_backend="ray",
+        weight_transfer_config=WeightTransferConfig(backend="nccl"),
+        load_format="dummy",
+    )
+    ray.get(llm.ready.remote())
+    num_workers = ray.get(llm.get_world_size.remote())
+
+    prompts = ["The capital of France is", "The future of AI is"]
+    before = generate(llm, prompts)
+
+    master_address = get_ip()
+    master_port = get_open_port()
+    world_size = NUM_SOURCES + num_workers
+    t0 = time.perf_counter()
+    ray.get(
+        [
+            source.setup_engine.remote(
+                llm,
+                master_address,
+                master_port,
+                world_size,
+            )
+            for source in sources
+        ]
+    )
+    print(f"[sync] engine setup took {time.perf_counter() - t0:.2f}s")
+
+    t0 = time.perf_counter()
+    ray.get([source.send_weights.remote() for source in sources])
+    print(f"[sync] iter 0 took {time.perf_counter() - t0:.3f}s")
+    after = generate(llm, prompts)
+    t0 = time.perf_counter()
+    ray.get([source.send_weights.remote() for source in sources])
+    print(f"[sync] iter 1 took {time.perf_counter() - t0:.3f}s")
+    repeated = generate(llm, prompts)
+
+    assert after != before, "generation did not change after weight transfer"
+    assert repeated == after, "repeating the same transfer changed generation"
+    print("Multi-source NCCL weight transfer completed successfully.")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/pyproject.toml b/pyproject.toml
index 46c836043ca2..386a029f1e74 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -226,8 +226,15 @@ FoPE = "FoPE"
 k_ot = "k_ot"
 view_seperator = "view_seperator"
 inverse_std_variences = "inverse_std_variences"
+# Muse Spark K3 monokernel: threads an expert
+TPE = "TPE"
+stride_bse = "stride_bse"
 
 [tool.typos.default.extend-words]
+# Muse Spark monokernel: shuffled-weight / scale tags, o_proj flag region
+shs = "shs"
+shw = "shw"
+opf = "opf"
 certifi = "certifi"
 MoNE = "MoNE"
 mone = "mone"
@@ -254,10 +261,15 @@ pard = "pard"
 AKS = "AKS"
 ba = "ba"
 fo = "fo"
+ot = "ot"
 nd = "nd"
 eles = "eles"
 datas = "datas"
 ser = "ser"
+# numpy flags.writeable
+writeable = "writeable"
+# HF CLIP weight name
+layrnorm = "layrnorm"
 ure = "ure"
 VALU = "VALU"
 # Walsh-Hadamard Transform
diff --git a/recipes/README.md b/recipes/README.md
index d96c9096ce96..5fd08935b4fc 100644
--- a/recipes/README.md
+++ b/recipes/README.md
@@ -9,9 +9,10 @@ The site, the command builder, and the JSON API are built by
 [vllm-project/recipes](https://github.com/vllm-project/recipes). That
 repository owns the schema and the shared vocabulary (hardware profiles,
 serving strategies, KV-offload deployments) and pins this directory to a
-commit. Omni (image, video, audio) recipes stay in vllm-project/recipes for
-now (moving to vllm-omni), and Ascend NPU content is inline in these recipes
-(moving to vllm-ascend overlays).
+commit. Omni (image, video, audio) recipes live in vllm-project/vllm-omni, and
+Ascend NPU content lives in vllm-project/vllm-ascend as overlays that merge into
+these recipes when the site is built. Don't add Ascend hardware here: open those
+changes against `recipes/overlays/` in vllm-ascend.
 
 ## Editing a recipe
 
diff --git a/recipes/models/MiniMaxAI/MiniMax-M3.1-Flash.yaml b/recipes/models/MiniMaxAI/MiniMax-M3.1-Flash.yaml
index 25ea0663eff1..14a88fdf6251 100644
--- a/recipes/models/MiniMaxAI/MiniMax-M3.1-Flash.yaml
+++ b/recipes/models/MiniMaxAI/MiniMax-M3.1-Flash.yaml
@@ -4,7 +4,7 @@ meta:
   provider: "MiniMax"
   description: "MiniMax M3.1-Flash vision-language MoE (~428B total / 23B active) with MSA sparse attention throughout and a built-in DSpark draft; quantization-aware trained, MXFP8 dense layers and NVFP4 experts. Served TP4 on NVIDIA Blackwell or AMD MI355X, prefill/decode disaggregated as TP4 + TP4 on MI355X or DEP2 + DEP4 on GB200/GB300, and TP8 on NVIDIA Hopper."
   date_added: 2026-10-23
-  date_updated: 2026-10-27
+  date_updated: 2026-10-30
   tasks:
     - text
     - multimodal
@@ -378,6 +378,625 @@ strategy_overrides:
             VLLM_KV_CACHE_LAYOUT: "BLHNC"
 
 presets:
+  # GB300 AgentX under Dynamo: aggregated TP4 workers, and DEP2 prefill
+  # workers feeding one DEP decode worker. Each preset restates its engine
+  # flags, so it does not depend on the pd_cluster GB300 block above.
+  gb300_agentx_agg:
+    label: "AgentX TP4 (Dynamo)"
+    hardware: gb300
+    strategy: single_node_tp
+    variants: [default]
+    workload: "Agentic (InferenceX AgentX)"
+    description: "Aggregated TP4 vLLM workers, one per 4-GPU GB300 tray, behind the Dynamo KV router with session affinity. Seven-token DSpark draft on FLASH_ATTN and prefill CUDA graphs up to 8192 tokens. KV stays in HBM, or spills to host DRAM through SimpleCPUOffloadConnector."
+    source: { label: "InferenceX #4631", href: "https://github.com/SemiAnalysisAI/InferenceX/pull/4631", key: "minimaxm31flash-fp4-gb300-dynamo-vllm-agentic-dspark-agg", revision: 5d3aa21b680eb8fb66a72386bd8e85e435e15a8c }
+    image:
+      vllm: "vllm/vllm-openai:minimax-m31-flash-arm64-cu134"
+      dynamo: "1.6.0.dev20261023"
+      pip: ["ai-dynamo-runtime==1.6.0.dev20261023", "ai-dynamo==1.6.0.dev20261023"]
+    tp: 4
+    select:
+      features: { text_only: true }
+    args:
+      - "--enable-prefix-caching"
+      - "--max-model-len"
+      - "1048576"
+      - "--max-num-batched-tokens"
+      - "32768"
+    env:
+      # Pins the prefix-block hashes, so they match across processes and restarts.
+      PYTHONHASHSEED: "0"
+      # A cold start (JIT builds, graph capture, autotune) takes more than 10 minutes.
+      VLLM_ENGINE_READY_TIMEOUT_S: "7200"
+      VLLM_HTTP_TIMEOUT_KEEP_ALIVE: "600"
+      VLLM_CHAT_TEMPLATE_RENDER_TIMEOUT: "600"
+    json_args:
+      "--compilation-config": { cudagraph_capture_sizes: [8, 16, 24, 32, 40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120, 128, 136, 144, 152, 160, 168, 176, 184, 192, 200, 208, 216, 224, 232, 240, 248, 256, 320, 384, 448, 512, 576, 640, 704, 768, 832, 896, 960, 1024, 1536, 2048, 3072, 4096, 6144, 8192] }
+      "--speculative-config": { attention_backend: "FLASH_ATTN" }
+    # --max-num-seqs is about twice the sessions a worker serves.
+    points:
+      gpu_kv:
+        label: "GPU-resident KV"
+        description: "KV cache in HBM only: --max-num-seqs 32 for about 16 sessions per worker, 64 for up to 32."
+        default: true
+        args: ["--max-num-seqs", "32", "--gpu-memory-utilization", "0.9"]
+        select: { kv_offload: "off" }
+      dram_kv:
+        label: "Host-DRAM KV offload"
+        description: "Adds a host-DRAM KV tier (860 GB per worker) for about 40-56 sessions per worker."
+        args: ["--max-num-seqs", "96", "--gpu-memory-utilization", "0.86"]
+        env: { VLLM_USE_SIMPLE_KV_OFFLOAD: "1" }
+        select: { kv_offload: simple }
+        json_args:
+          "--kv-transfer-config": { kv_connector_extra_config: { cpu_bytes_to_use_per_rank: null, cpu_bytes_to_use: 860000000000, lazy_offload: true } }
+    runbook: |
+      The GB300 aggregated AgentX lane in [SemiAnalysis InferenceX](https://inferencex.semianalysis.com/) (`minimaxm31flash-fp4-gb300-dynamo-vllm-agentic-dspark-agg`, [SemiAnalysisAI/InferenceX#4631](https://github.com/SemiAnalysisAI/InferenceX/pull/4631)) serves the default checkpoint text-only, with one TP4 `dynamo.vllm` worker per 4-GPU tray behind the Dynamo frontend. Add trays to scale out: the KV router spreads conversations across them. The lane measured 29.4k tok/s/GPU at 417 tok/s/user with 16 sessions per tray (GPU-resident KV, `--max-num-seqs 32`), and 96.6k tok/s/GPU at 185 tok/s/user with 56 sessions per tray (host-DRAM KV offload, `--max-num-seqs 96`).
+
+      **Image and versions.** The lane used an arm64 CUDA 13 build of the `vllm/vllm-openai:minimax-m31-flash` image (tag `minimax-m31-flash-arm64-cu134`); any build of that image family for GB300 will do. Dynamo `1.6.0.dev20261023` was installed on top: `pip install --extra-index-url https://pypi.nvidia.com ai-dynamo-runtime==1.6.0.dev20261023 ai-dynamo==1.6.0.dev20261023`.
+
+      **Tuning rules.**
+
+      - `--max-num-seqs` is about twice the sessions one worker serves: 32 up to 16 sessions, 64 at 32-48, 96 at 56.
+      - **GPU-resident KV** (`--gpu-memory-utilization 0.9`) is enough up to about 32 sessions per tray. Beyond that, multi-turn contexts outgrow HBM and evicted prefixes are recomputed.
+      - **Host-DRAM KV offload** adds `SimpleCPUOffloadConnector` with lazy offload (`VLLM_USE_SIMPLE_KV_OFFLOAD=1`). `cpu_bytes_to_use` is the budget for the whole worker, split across its four ranks: 860 GB pins 215 GB per rank. Size it from the host DRAM that is free after the engine processes are up. Lower `--gpu-memory-utilization` to 0.86 to leave room for the transfer buffers.
+      - The CUDA graph sizes go up to 8192 tokens, so prefill chunks of up to 8192 tokens also replay graphs. `--max-num-batched-tokens 32768` still allows larger mixed steps, which run piecewise.
+
+      **Running under Dynamo.** Start etcd and NATS (32 MB max payload) once:
+
+      ```bash
+      etcd --listen-client-urls http://0.0.0.0:2379 --advertise-client-urls http://<infra-host>:2379
+      printf 'max_payload: 33554432\n' > nats.conf && nats-server -c nats.conf
+      ```
+
+      On each tray, export `ETCD_ENDPOINTS=http://<infra-host>:2379` and `NATS_SERVER=nats://<infra-host>:4222` together with the generated environment and `DYN_TCP_CONNECT_TIMEOUT=120`, `DYN_BACKEND_ADMISSION_ENGINE_REQUEST_LIMIT=1024`, `DYN_BACKEND_ADMISSION_ENGINE_REQUEST_WAITING_LIMIT=1024` and `DYN_BACKEND_ADMISSION_REQUEST_QUEUE_LIMIT=2048`. Then run the generated flags with `python3 -m dynamo.vllm --model MiniMaxAI/MiniMax-M3.1-Flash --served-model-name MiniMaxAI/MiniMax-M3.1-Flash --trust-remote-code` in place of `vllm serve MiniMaxAI/MiniMax-M3.1-Flash`. `dynamo.vllm` rejects vLLM's API-server flags, so leave out `--tool-call-parser minimax_m3 --enable-auto-tool-choice` and add:
+
+      ```bash
+      --dyn-tool-call-parser minimax_m3 --dyn-reasoning-parser minimax_m3 \
+        --kv-events-config '{"publisher":"zmq","topic":"kv-events","enable_kv_cache_events":true,"endpoint":"tcp://*:5200"}'
+      ```
+
+      Start the frontend:
+
+      ```bash
+      export ETCD_ENDPOINTS=http://<infra-host>:2379 NATS_SERVER=nats://<infra-host>:4222
+      export DYN_ROUTER_ACTIVE_REQUEST_EXPIRY_SECS=3600 DYN_TCP_CONNECT_TIMEOUT=120 DYN_TOKENIZER=fastokens
+      python3 -m dynamo.frontend --http-port 8000 --router-mode kv --router-session-affinity-ttl-secs 14400
+      ```
+
+      `DYN_TOKENIZER=fastokens` switches the frontend to the fastokens tokenizer backend that ships with Dynamo. With 100K+ token agentic prompts, the default tokenizer spends about 150 ms per request on the frontend (up to 600 ms at p90), all of it added to TTFT; fastokens takes about 1 ms.
+
+      The KV router uses the workers' KV events to send each request to the worker that holds its prefix. Session affinity keeps a conversation on that worker for up to four hours; send a stable conversation ID in the `X-Dynamo-Session-ID` request header.
+
+      The InferenceX throughput points replaced block verification with synthetic rejection at the committed golden DSpark acceptance length (4.04), so every point is measured at the same acceptance rate. Production serving keeps the default rejection sampler (`rejection_sample_method: block`), as in the command above.
+
+  gb300_agentx_pd:
+    label: "AgentX PD (Dynamo)"
+    hardware: gb300
+    strategy: pd_cluster
+    variants: [default]
+    workload: "Agentic (InferenceX AgentX)"
+    description: "DEP2 prefill workers, two per GB300 tray, feed one DEP8 decode worker (throughput) or DEP16 (lower latency) over NIXL. Both roles share a Mooncake KV store in host DRAM. The Dynamo KV router routes with session affinity. Seven-token DSpark draft on FLASHINFER_TS, DeepGEMM MegaMoE on both roles."
+    source: { label: "InferenceX #4631", href: "https://github.com/SemiAnalysisAI/InferenceX/pull/4631", key: "minimaxm31flash-fp4-gb300-dynamo-vllm-agentic-dspark-disagg", revision: 5d3aa21b680eb8fb66a72386bd8e85e435e15a8c }
+    image:
+      vllm: "vllm/vllm-openai:minimax-m31-flash-arm64-cu134"
+      dynamo: "1.6.0.dev20261023"
+      pip: ["ai-dynamo-runtime==1.6.0.dev20261023", "ai-dynamo==1.6.0.dev20261023", "mooncake-transfer-engine-cuda13==0.3.13.post1"]
+    # Dynamo workers take engine flags only; the tool-call parser moves to
+    # --dyn-tool-call-parser (see the runbook).
+    select:
+      pd_router: dynamo
+      features: { text_only: true, tool_calling: false }
+    args:
+      - "--enable-expert-parallel"
+      - "--moe-backend"
+      - "deep_gemm_mega_moe"
+      - "--numa-bind"
+      - "--enable-prefix-caching"
+      - "--max-model-len"
+      - "1048576"
+    # The RDMA names are a GB300 tray's four NICs; use the host's own.
+    env:
+      PYTHONHASHSEED: "0"
+      VLLM_USE_V2_MODEL_RUNNER: "1"
+      VLLM_ENGINE_READY_TIMEOUT_S: "7200"
+      VLLM_HTTP_TIMEOUT_KEEP_ALIVE: "600"
+      NCCL_NET_GDR_C2C: "1"
+      NCCL_IB_HCA: "mlx5_0,mlx5_1,mlx5_2,mlx5_3"
+      NVSHMEM_REMOTE_TRANSPORT: "none"
+      NVSHMEM_ENABLE_NIC_PE_MAPPING: "1"
+      NVSHMEM_HCA_PE_MAPPING: "mlx5_0:1:1,mlx5_1:1:1,mlx5_2:1:1,mlx5_3:1:1"
+      UCX_TLS: "rc,cuda_copy"
+      UCX_NET_DEVICES: "mlx5_0:1,mlx5_1:1,mlx5_2:1,mlx5_3:1"
+      UCX_CUDA_IPC_ENABLE_MNNVL: "n"
+      UCX_MEMTYPE_CACHE: "n"
+      UCX_MEMTYPE_REG_WHOLE: "n"
+      UCX_RNDV_PIPELINE_ERROR_HANDLING: "y"
+      NIXL_UCX_SGL_ENABLE: "1"
+      WITH_NVIDIA_PEERMEM: "0"
+      MC_SLICE_SIZE: "1048576"
+      MC_MAX_MR_SIZE: "4294967296"
+      MC_ENABLE_DEST_DEVICE_AFFINITY: "1"
+      DYN_TCP_CONNECT_TIMEOUT: "120"
+      DYN_BACKEND_ADMISSION_ENGINE_REQUEST_LIMIT: "1024"
+      DYN_BACKEND_ADMISSION_ENGINE_REQUEST_WAITING_LIMIT: "1024"
+      DYN_BACKEND_ADMISSION_REQUEST_QUEUE_LIMIT: "2048"
+    roles:
+      prefill:
+        # One DEP2 worker on two GPUs; two per tray, as many trays as the load needs.
+        nodes: 1
+        parallelism: dep
+        args:
+          - "--data-parallel-size"
+          - "2"
+          - "--data-parallel-size-local"
+          - "2"
+          - "--max-num-seqs"
+          - "64"
+          - "--max-num-batched-tokens"
+          - "32768"
+          - "--long-prefill-token-threshold"
+          - "4096"
+          - "--gpu-memory-utilization"
+          - "0.85"
+          - "--compilation-config"
+          - '{"cudagraph_mode":"FULL_AND_PIECEWISE"}'
+          - "--max-cudagraph-capture-size"
+          - "2048"
+          - "--kv-events-config"
+          - '{"publisher":"zmq","topic":"kv-events","enable_kv_cache_events":true,"endpoint":"tcp://*:5200"}'
+          - "--kv-transfer-config"
+          - '{"kv_connector":"MultiConnector","kv_role":"kv_both","kv_load_failure_policy":"recompute","kv_connector_extra_config":{"connectors":[{"kv_connector":"NixlConnector","kv_role":"kv_both","kv_load_failure_policy":"fail","kv_buffer_device":"cuda","kv_connector_extra_config":{"backends":["UCX"],"enforce_handshake_compat":false,"kv_lease_duration":120}},{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_both","kv_load_failure_policy":"recompute","kv_connector_extra_config":{"load_async":true,"lookup_async":true,"enable_offload":false}}]}}'
+        env:
+          VLLM_MOONCAKE_LOAD_RECV_THREADS: "8"
+      decode:
+        nodes: 2
+        parallelism: dep
+        # 48 sequences per DP rank; decode graphs cover 48 x (1 + 7) tokens.
+        # Decode also writes the KV it generates into the Mooncake store.
+        remove_args: ["--no-enable-prefix-caching"]
+        args:
+          - "--max-num-seqs"
+          - "48"
+          - "--max-num-batched-tokens"
+          - "8192"
+          - "--watermark"
+          - "0.05"
+          - "--gpu-memory-utilization"
+          - "0.95"
+          - "--compilation-config"
+          - '{"cudagraph_mode":"FULL_DECODE_ONLY"}'
+          - "--max-cudagraph-capture-size"
+          - "384"
+          - "--numa-bind-nodes"
+          - "0"
+          - "0"
+          - "1"
+          - "1"
+          - "--kv-transfer-config"
+          - '{"kv_connector":"MultiConnector","kv_role":"kv_both","kv_load_failure_policy":"recompute","kv_connector_extra_config":{"connectors":[{"kv_connector":"NixlConnector","kv_role":"kv_both","kv_load_failure_policy":"fail","kv_buffer_device":"cuda","kv_connector_extra_config":{"backends":["UCX"],"enforce_handshake_compat":false,"kv_lease_duration":120}},{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_both","kv_load_failure_policy":"recompute","kv_connector_extra_config":{"load_async":true,"lookup_async":true,"enable_offload":false,"save_decode_cache":true}}]}}'
+        env:
+          VLLM_MOONCAKE_LOAD_RECV_THREADS: "4"
+          VLLM_MAX_ASYNC_KV_ADMITS_PER_STEP: "4"
+          VLLM_RANDOMIZE_DP_DUMMY_INPUTS: "1"
+          VLLM_ALLREDUCE_USE_SYMM_MEM: "0"
+          NCCL_CUMEM_ENABLE: "1"
+          NCCL_MNNVL_ENABLE: "1"
+          NCCL_NVLS_ENABLE: "1"
+    points:
+      dep8:
+        label: "DEP8 decode (throughput)"
+        description: "One DEP8 decode worker across two trays."
+        default: true
+      dep16:
+        label: "DEP16 decode (lower latency)"
+        description: "One DEP16 decode worker across four trays."
+        roles: { decode: { nodes: 4 } }
+    runbook: |
+      The GB300 disaggregated AgentX lane in [SemiAnalysis InferenceX](https://inferencex.semianalysis.com/) (`minimaxm31flash-fp4-gb300-dynamo-vllm-agentic-dspark-disagg`, [SemiAnalysisAI/InferenceX#4631](https://github.com/SemiAnalysisAI/InferenceX/pull/4631)) serves the default checkpoint text-only with Dynamo and vLLM. Six prefill workers, each DEP2 on two GPUs (two per 4-GPU tray, three trays), feed one decode worker that is DEP8 across two trays or DEP16 across four. KV moves from prefill to decode over NIXL (UCX over RDMA), and every rank contributes host DRAM to a shared Mooncake KV store, so a returning conversation's prefix is loaded from the store instead of recomputed. Both roles use expert parallelism with DeepGEMM MegaMoE, attention data parallelism, the V2 model runner and the seven-token DSpark draft on `FLASHINFER_TS`.
+
+      | Prefill workers (DEP2) | Decode | GPUs | Sessions | tok/s/GPU | tok/s/user | TTFT p90 | E2E latency p90 |
+      |---:|---|---:|---:|---:|---:|---:|---:|
+      | 6 | DEP8 | 20 | 280 | 93.1k | 240 | 2.7 s | 12.2 s |
+      | 6 | DEP8 | 20 | 320 | 100.5k | 229 | 3.3 s | 13.8 s |
+      | 6 | DEP8 | 20 | 360 | 107.9k | 224 | 4.2 s | 15.9 s |
+      | 6 | DEP8 | 20 | 400 | 114.5k | 219 | 5.9 s | 18.4 s |
+      | 6 | DEP16 | 28 | 300 | 75.0k | 296 | 1.8 s | 9.3 s |
+      | 6 | DEP16 | 28 | 360 | 85.8k | 285 | 2.4 s | 10.3 s |
+
+      The preset renders one prefill worker and the decode worker. Run more prefill workers with the same command. Two workers share a tray on GPUs 0-1 and 2-3 (`--device-ids 0,1` and `--device-ids 2,3` on `dynamo.vllm`). Give each its own `--data-parallel-rpc-port`, `VLLM_NIXL_SIDE_CHANNEL_PORT` and `--kv-events-config` endpoint port.
+
+      **Image and versions.** The lane used an arm64 CUDA 13 build of the `vllm/vllm-openai:minimax-m31-flash` image (tag `minimax-m31-flash-arm64-cu134`), with Dynamo `1.6.0.dev20261023` and Mooncake `mooncake-transfer-engine-cuda13==0.3.13.post1` installed on every worker and on the Mooncake master host: `pip install --extra-index-url https://pypi.nvidia.com ai-dynamo-runtime==1.6.0.dev20261023 ai-dynamo==1.6.0.dev20261023 mooncake-transfer-engine-cuda13==0.3.13.post1`.
+
+      ### Tuning rules
+
+      - **Prefill pool.** The prefill pool sets TTFT, so scale prefill workers with load. The lane's points kept TTFT p90 between 1.8 and 5.9 s with one DEP2 prefill worker per 47-67 concurrent agentic sessions, so 93-133 sessions per prefill tray (six workers at 280-400 sessions). Add prefill workers when TTFT p90 climbs while decode still has room. Two DEP2 workers per tray give the router twice as many prefill targets as one DEP4 worker.
+      - **DEP8 or DEP16 decode.** DEP8 gives the most tokens per GPU: up to 50 sessions per decode GPU (400 sessions) at 219 tok/s/user. DEP16 halves the batch per rank and gives about 23% more tok/s/user (296 at 300 sessions, against 240 at 280 on DEP8) at about 19% fewer tokens per GPU, and keeps E2E latency p90 near 10 s up to 360 sessions. `--max-num-seqs 48` per DP rank is a cap, not the operating point. The decode graph cap is `--max-num-seqs` x (1 + 7 draft tokens) = 384. `--watermark 0.05` keeps 5% of the decode KV cache free, so extra requests wait in the queue instead of preempting running ones. Keep decode `--max-num-batched-tokens` at 8192: a preempted request recomputes its context in chunks of that size.
+      - **Prefill batching.** `--long-prefill-token-threshold 4096` chunks long prompts into 4096-token pieces. Short prompts then share each step with them instead of waiting behind a 100K-token context, which lowers the TTFT tail. The recipe's 1P1D section above uses 8192. `--max-num-batched-tokens 32768` bounds a prefill step. Prefill captures CUDA graphs up to 2048 tokens.
+      - **Mooncake pool size.** `global_segment_size` is registered by every DP rank, prefill and decode alike, so a tray pins four times that value. The lane used 180 GB per rank (720 GB per tray), which leaves 80-100 GiB of a GB300 tray's host DRAM free with the engines up; 200 GB per rank ran trays out of host memory. A larger pool raises the hit rate at high session counts: going from 160 to 180 GB per rank, with the services moved off the engine trays, lowered TTFT p90 by 0.4-0.8 s at 280-360 sessions. Size it as (host DRAM free per tray after the engines are up - headroom for NIXL and the 4 GB local buffers) / ranks per tray. The decode ranks also write the KV they generate into the store (`save_decode_cache: true`), so the next turn's prefill can load the previous answer instead of recomputing it.
+      - **Routing.** The KV router (`--router-mode kv`) reads the prefill workers' KV events and sends each request to the worker that already holds its prefix. Session affinity (`--router-session-affinity-ttl-secs 14400`) keeps a conversation on that worker for up to four hours. Send a stable conversation ID in the `X-Dynamo-Session-ID` header. The Mooncake store catches the prefixes that miss after a reroute.
+      - **Transport.** Cross-tray NIXL needs `UCX_TLS=rc,cuda_copy` with `UCX_NET_DEVICES` set to the tray's RDMA ports. Without them, UCX can fall back to TCP and transfers stall. Point `UCX_MODULE_DIR` at the UCX bundled with the NIXL wheel (`nixl_cu13.libs/ucx` in the image's `dist-packages`). `kv_lease_duration: 120` holds a prefill's blocks for up to 120 s until decode reads them (the default is 30 s). `VLLM_MAX_ASYNC_KV_ADMITS_PER_STEP=4` caps how many requests with finished KV loads decode admits per step.
+      - **NUMA.** `--numa-bind` pins each rank to its GPU's NUMA node. Decode passes the tray's map, `--numa-bind-nodes 0 0 1 1`, explicitly.
+      - `--moe-backend deep_gemm_mega_moe` is the same as the lane's `--kernel-config '{"moe_backend":"deep_gemm_mega_moe"}'`. `--prefill-schedule-interval 1`, also set by the lane, is vLLM's default.
+
+      ### Services
+
+      Start etcd, NATS and the Mooncake master once, on a host that runs no engine workers, together with the frontend. The engine trays' host DRAM then goes to the Mooncake pool.
+
+      ```bash
+      etcd --listen-client-urls http://0.0.0.0:2379 --advertise-client-urls http://<infra-host>:2379
+      printf 'max_payload: 33554432\n' > nats.conf && nats-server -c nats.conf
+      mooncake_master --rpc_port=50051 \
+        --default_kv_lease_ttl=60000 --eviction_high_watermark_ratio=0.95 --eviction_ratio=0.10
+      ```
+
+      Every worker reads the store configuration from the JSON file that `MOONCAKE_CONFIG_PATH` points to (export it with the rest of the environment):
+
+      ```json
+      {"master_server_address": "<infra-host>:50051", "metadata_server": "P2PHANDSHAKE",
+       "global_segment_size": "180GB", "local_buffer_size": "4GB", "protocol": "rdma",
+       "device_name": "mlx5_0,mlx5_1,mlx5_2,mlx5_3", "enable_offload": false}
+      ```
+
+      ### Workers
+
+      Export `ETCD_ENDPOINTS=http://<infra-host>:2379`, `NATS_SERVER=nats://<infra-host>:4222` and the generated environment on every tray, then run the generated `python3 -m dynamo.vllm` commands. Add `--served-model-name MiniMaxAI/MiniMax-M3.1-Flash --trust-remote-code --dyn-tool-call-parser minimax_m3 --dyn-reasoning-parser minimax_m3` to both roles. `dynamo.vllm` rejects vLLM's API-server flags `--tool-call-parser` and `--enable-auto-tool-choice`. Multi-node DEP decode runs one process per tray with `--data-parallel-start-rank` set to 0, 4, 8 and so on, and a shared `--data-parallel-address` and `--data-parallel-rpc-port`. `--data-parallel-hybrid-lb` lets every tray's Dynamo runtime register its own DP ranks with the frontend.
+
+      ### Frontend
+
+      ```bash
+      export ETCD_ENDPOINTS=http://<infra-host>:2379 NATS_SERVER=nats://<infra-host>:4222
+      export DYN_ROUTER_ACTIVE_REQUEST_EXPIRY_SECS=3600 DYN_TCP_CONNECT_TIMEOUT=120 DYN_TOKENIZER=fastokens
+      python3 -m dynamo.frontend --http-port 8000 --router-mode kv --router-session-affinity-ttl-secs 14400
+      ```
+
+      `DYN_TOKENIZER=fastokens` switches the frontend to the fastokens tokenizer backend that ships with Dynamo. With 100K+ token agentic prompts, the default tokenizer spends about 150 ms per request on the frontend (up to 600 ms at p90), all of it added to TTFT; fastokens takes about 1 ms.
+
+      The InferenceX throughput points replaced block verification with synthetic rejection at the committed golden DSpark acceptance length (4.04), so every point is measured at the same acceptance rate. Production serving keeps the default rejection sampler (`rejection_sample_method: block`), as in the command above.
+
+  # GB300 AgentX on one 4-GPU tray without Dynamo: TP4 engines, a DEP4 engine
+  # behind vllm-router, and TP2 engines with host-DRAM KV offload.
+  gb300_agentx_tp4:
+    label: "AgentX TP4"
+    hardware: gb300
+    strategy: single_node_tp
+    variants: [default]
+    workload: "Agentic (InferenceX AgentX)"
+    description: "One TP4 engine on a 4-GPU GB300 tray with the seven-token DSpark draft on FLASHINFER_TS; GPU-resident KV, with per-concurrency scheduling and CUDA graph settings."
+    source: { label: "InferenceX #4662", href: "https://github.com/SemiAnalysisAI/InferenceX/pull/4662", key: "minimaxm31flash-fp4-gb300-vllm-agentic-dspark", revision: 85956a00e240c2b7428e1e9406f25c731ddb0470 }
+    image: "vllm/vllm-openai:minimax-m31-flash-arm64-cu134@sha256:5d76347628973fb80519a919a2f2740b6645b62cba3fe2bd87b25f111bd0e948"
+    tp: 4
+    select:
+      features: { text_only: true }
+    args:
+      - "--enable-prefix-caching"
+      - "--max-model-len"
+      - "1048576"
+      - "--max-num-batched-tokens"
+      - "8192"
+    env:
+      VLLM_ENGINE_READY_TIMEOUT_S: "7200"
+      VLLM_CHAT_TEMPLATE_RENDER_TIMEOUT: "600"
+    points:
+      c1: { label: "1", concurrency: 1, args: ["--max-num-seqs", "2"] }
+      c4: { label: "4", concurrency: 4, args: ["--max-num-seqs", "8"] }
+      c8:
+        label: "8"
+        concurrency: 8
+        args: ["--max-num-seqs", "16"]
+        json_args:
+          "--compilation-config": { cudagraph_capture_sizes: [1, 2, 4, 8, 16, 24, 32, 40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120, 128, 136, 144, 152, 160, 168, 176, 184, 192, 200, 208, 216, 224, 232, 240, 248, 256, 264, 272, 280, 288, 296, 304, 312, 320, 328, 336, 344, 352, 360, 368, 376, 384, 392, 400, 408, 416, 424, 432, 440, 448, 456, 464, 472, 480, 488, 496, 504, 512, 640, 768, 896, 1024, 1152, 1280, 1408, 1536, 1664, 1792, 1920, 2048, 2304, 2560, 2816, 3072, 3328, 3584, 3840, 4096, 4608, 5120, 5632, 6144, 6656, 7168, 7680, 8192] }
+      c16:
+        label: "16"
+        concurrency: 16
+        args: ["--max-num-seqs", "64"]
+        json_args:
+          "--compilation-config": { cudagraph_capture_sizes: [1, 2, 4, 8, 16, 24, 32, 40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120, 128, 136, 144, 152, 160, 168, 176, 184, 192, 200, 208, 216, 224, 232, 240, 248, 256, 264, 272, 280, 288, 296, 304, 312, 320, 328, 336, 344, 352, 360, 368, 376, 384, 392, 400, 408, 416, 424, 432, 440, 448, 456, 464, 472, 480, 488, 496, 504, 512, 640, 768, 896, 1024, 1152, 1280, 1408, 1536, 1664, 1792, 1920, 2048, 2304, 2560, 2816, 3072, 3328, 3584, 3840, 4096, 4608, 5120, 5632, 6144, 6656, 7168, 7680, 8192] }
+      c24:
+        label: "24"
+        concurrency: 24
+        args:
+          - "--max-num-seqs"
+          - "64"
+          - "--max-num-batched-tokens"
+          - "32768"
+          - "--long-prefill-token-threshold"
+          - "8192"
+          - "--prefill-schedule-interval"
+          - "16"
+          - "--cudagraph-capture-sizes"
+          - "1"
+          - "2"
+          - "4"
+          - "8"
+          - "16"
+          - "24"
+          - "32"
+          - "40"
+          - "48"
+          - "56"
+          - "64"
+          - "72"
+          - "80"
+          - "88"
+          - "96"
+          - "104"
+          - "112"
+          - "120"
+          - "128"
+          - "136"
+          - "144"
+          - "152"
+          - "160"
+          - "168"
+          - "176"
+          - "184"
+          - "192"
+          - "200"
+          - "208"
+          - "216"
+          - "224"
+          - "232"
+          - "240"
+          - "248"
+          - "256"
+          - "272"
+          - "288"
+          - "304"
+          - "320"
+          - "336"
+          - "352"
+          - "368"
+          - "384"
+          - "400"
+          - "416"
+          - "432"
+          - "448"
+          - "464"
+          - "480"
+          - "496"
+          - "512"
+          - "528"
+          - "544"
+          - "560"
+          - "576"
+          - "592"
+          - "608"
+          - "624"
+          - "640"
+          - "656"
+          - "672"
+          - "688"
+          - "704"
+          - "720"
+          - "736"
+          - "752"
+          - "768"
+          - "784"
+          - "800"
+          - "816"
+          - "832"
+          - "848"
+          - "864"
+          - "880"
+          - "896"
+          - "912"
+          - "928"
+          - "944"
+          - "960"
+          - "976"
+          - "992"
+          - "1008"
+          - "1024"
+          - "1280"
+          - "1536"
+          - "1792"
+          - "2048"
+          - "2304"
+          - "2560"
+          - "2816"
+          - "3072"
+          - "3328"
+          - "3584"
+          - "3840"
+          - "4096"
+          - "4608"
+          - "5120"
+          - "5632"
+          - "6144"
+          - "6656"
+          - "7168"
+          - "7680"
+          - "8192"
+        env: { VLLM_HTTP_TIMEOUT_KEEP_ALIVE: "600" }
+      c32: { label: "32", concurrency: 32, args: ["--max-num-seqs", "64"] }
+      c40:
+        label: "40"
+        concurrency: 40
+        args:
+          - "--max-num-seqs"
+          - "64"
+          - "--long-prefill-token-threshold"
+          - "8192"
+          - "--gpu-memory-utilization"
+          - "0.95"
+          - "--watermark"
+          - "0.05"
+          - "--cudagraph-capture-sizes"
+          - "1"
+          - "2"
+          - "4"
+          - "8"
+          - "16"
+          - "24"
+          - "32"
+          - "40"
+          - "48"
+          - "56"
+          - "64"
+          - "72"
+          - "80"
+          - "88"
+          - "96"
+          - "104"
+          - "112"
+          - "120"
+          - "128"
+          - "136"
+          - "144"
+          - "152"
+          - "160"
+          - "168"
+          - "176"
+          - "184"
+          - "192"
+          - "200"
+          - "208"
+          - "216"
+          - "224"
+          - "232"
+          - "240"
+          - "248"
+          - "256"
+          - "272"
+          - "288"
+          - "304"
+          - "320"
+          - "336"
+          - "352"
+          - "368"
+          - "384"
+          - "400"
+          - "416"
+          - "432"
+          - "448"
+          - "464"
+          - "480"
+          - "496"
+          - "512"
+          - "528"
+          - "544"
+          - "560"
+          - "576"
+          - "592"
+          - "608"
+          - "624"
+          - "640"
+          - "656"
+          - "672"
+          - "688"
+          - "704"
+          - "720"
+          - "736"
+          - "752"
+          - "768"
+          - "784"
+          - "800"
+          - "816"
+          - "832"
+          - "848"
+          - "864"
+          - "880"
+          - "896"
+          - "912"
+          - "928"
+          - "944"
+          - "960"
+          - "976"
+          - "992"
+          - "1008"
+          - "1024"
+          - "1280"
+          - "1536"
+          - "1792"
+          - "2048"
+          - "2304"
+          - "2560"
+          - "2816"
+          - "3072"
+          - "3328"
+          - "3584"
+          - "3840"
+          - "4096"
+          - "4608"
+          - "5120"
+          - "5632"
+          - "6144"
+          - "6656"
+          - "7168"
+          - "7680"
+          - "8192"
+        env: { VLLM_HTTP_TIMEOUT_KEEP_ALIVE: "600" }
+    runbook: |
+      The GB300 TP4 points of the AgentX lane in [SemiAnalysis InferenceX](https://inferencex.semianalysis.com/) (`minimaxm31flash-fp4-gb300-vllm-agentic-dspark`, [SemiAnalysisAI/InferenceX#4662](https://github.com/SemiAnalysisAI/InferenceX/pull/4662)) serve the default checkpoint text-only on one TP4 engine per 4-GPU tray, with `--max-num-seqs` set per concurrency.
+
+      - **Concurrency 1, 4 and 32** use the default CUDA graph capture sizes.
+      - **Concurrency 8 and 16** capture CUDA graphs for every size up to 8192 tokens.
+      - **Concurrency 24** raises `--max-num-batched-tokens` to 32768 and schedules prefill every 16 steps; **concurrency 40** keeps 8192 and raises `--gpu-memory-utilization` to 0.95. Both cap prefill chunks at 8192 tokens and capture graphs to 8192.
+
+  gb300_agentx_dep4:
+    label: "AgentX DEP4"
+    hardware: gb300
+    strategy: single_node_dep
+    variants: [default]
+    workload: "Agentic (InferenceX AgentX)"
+    description: "One data-parallel, expert-parallel engine across the four GPUs of a GB300 tray behind vllm-router with session affinity. Three-token DSpark draft on FLASHINFER_TS."
+    source: { label: "InferenceX #4662", href: "https://github.com/SemiAnalysisAI/InferenceX/pull/4662", key: "minimaxm31flash-fp4-gb300-vllm-agentic-dspark", revision: 85956a00e240c2b7428e1e9406f25c731ddb0470 }
+    image: "vllm/vllm-openai:minimax-m31-flash-arm64-cu134@sha256:5d76347628973fb80519a919a2f2740b6645b62cba3fe2bd87b25f111bd0e948"
+    select:
+      features: { text_only: true }
+    args:
+      - "--enable-prefix-caching"
+      - "--max-model-len"
+      - "1048576"
+      - "--max-num-batched-tokens"
+      - "8192"
+    env:
+      VLLM_ENGINE_READY_TIMEOUT_S: "7200"
+      VLLM_CHAT_TEMPLATE_RENDER_TIMEOUT: "600"
+    json_args:
+      "--speculative-config": { num_speculative_tokens: 3 }
+    points:
+      c64: { label: "64", concurrency: 64, args: ["--max-num-seqs", "32"] }
+      c96: { label: "96", concurrency: 96, args: ["--max-num-seqs", "48"] }
+      c128: { label: "128", concurrency: 128, args: ["--max-num-seqs", "64"] }
+    runbook: |
+      The DEP4 points of the GB300 AgentX lane in [SemiAnalysis InferenceX](https://inferencex.semianalysis.com/) (`minimaxm31flash-fp4-gb300-vllm-agentic-dspark`, [SemiAnalysisAI/InferenceX#4662](https://github.com/SemiAnalysisAI/InferenceX/pull/4662)) run one engine with `--data-parallel-size 4` and `--enable-expert-parallel`; `--max-num-seqs` is the per-rank limit.
+
+      The router uses the `consistent_hash` policy so a conversation stays on the rank that holds its prefix. The lane also runs it with `--request-timeout-secs 14400` and `--disable-retries`.
+
+  gb300_agentx_tp2:
+    label: "AgentX TP2 + KV offload"
+    hardware: gb300
+    strategy: single_node_tp
+    variants: [default]
+    workload: "Agentic (InferenceX AgentX)"
+    description: "One TP2 engine on two of the tray's four GB300 GPUs with the seven-token DSpark draft on FLASHINFER_TS and host-DRAM KV offload (482 GB per engine)."
+    source: { label: "InferenceX #4662", href: "https://github.com/SemiAnalysisAI/InferenceX/pull/4662", key: "minimaxm31flash-fp4-gb300-vllm-agentic-dspark", revision: 85956a00e240c2b7428e1e9406f25c731ddb0470 }
+    image: "vllm/vllm-openai:minimax-m31-flash-arm64-cu134@sha256:5d76347628973fb80519a919a2f2740b6645b62cba3fe2bd87b25f111bd0e948"
+    tp: 2
+    select:
+      features: { text_only: true }
+      kv_offload: simple
+    args:
+      - "--enable-prefix-caching"
+      - "--max-model-len"
+      - "1048576"
+      - "--max-num-batched-tokens"
+      - "32768"
+    env:
+      VLLM_ENGINE_READY_TIMEOUT_S: "7200"
+      VLLM_CHAT_TEMPLATE_RENDER_TIMEOUT: "600"
+      VLLM_HTTP_TIMEOUT_KEEP_ALIVE: "600"
+      VLLM_USE_SIMPLE_KV_OFFLOAD: "1"
+    json_args:
+      "--compilation-config": { cudagraph_capture_sizes: [1, 2, 4, 8, 16, 24, 32, 40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120, 128, 136, 144, 152, 160, 168, 176, 184, 192, 200, 208, 216, 224, 232, 240, 248, 256, 264, 272, 280, 288, 296, 304, 312, 320, 328, 336, 344, 352, 360, 368, 376, 384, 392, 400, 408, 416, 424, 432, 440, 448, 456, 464, 472, 480, 488, 496, 504, 512, 640, 768, 896, 1024, 1152, 1280, 1408, 1536, 1664, 1792, 1920, 2048, 2304, 2560, 2816, 3072, 3328, 3584, 3840, 4096, 4608, 5120, 5632, 6144, 6656, 7168, 7680, 8192] }
+      "--kv-transfer-config": { kv_connector_extra_config: { cpu_bytes_to_use_per_rank: null, cpu_bytes_to_use: 482000000000, lazy_offload: true } }
+    points:
+      c28: { label: "28", concurrency: 28, args: ["--max-num-seqs", "42"] }
+      c32: { label: "32", concurrency: 32, args: ["--max-num-seqs", "48"] }
+      c36: { label: "36", concurrency: 36, args: ["--max-num-seqs", "54"] }
+    runbook: |
+      The TP2 points of the GB300 AgentX lane in [SemiAnalysis InferenceX](https://inferencex.semianalysis.com/) (`minimaxm31flash-fp4-gb300-vllm-agentic-dspark`, [SemiAnalysisAI/InferenceX#4662](https://github.com/SemiAnalysisAI/InferenceX/pull/4662)) add a host-DRAM KV tier through `SimpleCPUOffloadConnector` with lazy offload. `cpu_bytes_to_use` (482 GB) is the budget for the whole engine, about half of the tray's host DRAM, so the engine needs that much free host RAM. `--max-num-seqs` is 1.5 times the concurrency.
+
   b300_agentx_tp4:
     label: "AgentX TP4"
     hardware: b300
@@ -488,6 +1107,386 @@ presets:
 
       The InferenceX throughput points replaced block verification with synthetic rejection at the committed golden DSpark acceptance length (4.04) so every point is measured at the same acceptance rate. The command above keeps block verification.
 
+  # GB200 AgentX under Dynamo: aggregated TP4 workers, and DEP4 prefill trays
+  # feeding one DEP decode worker. Each preset restates its engine flags, so it
+  # does not depend on the pd_cluster GB200 block above.
+  gb200_agentx_agg:
+    label: "AgentX TP4 (Dynamo)"
+    hardware: gb200
+    strategy: single_node_tp
+    variants: [default]
+    workload: "Agentic (InferenceX AgentX)"
+    description: "Aggregated TP4 vLLM workers, one per 4-GPU GB200 tray, behind the Dynamo KV router with session affinity. Seven-token DSpark draft on FLASH_ATTN and prefill CUDA graphs up to 8192 tokens. KV stays in HBM, or spills to host DRAM through SimpleCPUOffloadConnector."
+    source: { label: "InferenceX #4632", href: "https://github.com/SemiAnalysisAI/InferenceX/pull/4632", key: "minimaxm31flash-fp4-gb200-dynamo-vllm-agentic-dspark-agg", revision: 936213146b79b55e8a11ffbc1d6010145995c6a6 }
+    image:
+      vllm: "vllm/vllm-openai:minimax-m31-flash-arm64-cu134"
+      dynamo: "1.6.0.dev20261023"
+      pip: ["ai-dynamo-runtime==1.6.0.dev20261023", "ai-dynamo==1.6.0.dev20261023"]
+    tp: 4
+    select:
+      features: { text_only: true }
+    args:
+      - "--enable-prefix-caching"
+      - "--max-model-len"
+      - "1048576"
+      - "--max-num-batched-tokens"
+      - "32768"
+    env:
+      # Pins the prefix-block hashes, so they match across processes and restarts.
+      PYTHONHASHSEED: "0"
+      # A cold start (JIT builds, graph capture, autotune) takes more than 10 minutes.
+      VLLM_ENGINE_READY_TIMEOUT_S: "7200"
+    json_args:
+      "--compilation-config": { cudagraph_capture_sizes: [8, 16, 24, 32, 40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120, 128, 136, 144, 152, 160, 168, 176, 184, 192, 200, 208, 216, 224, 232, 240, 248, 256, 320, 384, 448, 512, 576, 640, 704, 768, 832, 896, 960, 1024, 1536, 2048, 3072, 4096, 6144, 8192] }
+      "--speculative-config": { attention_backend: "FLASH_ATTN" }
+    # --max-num-seqs is about twice the sessions a worker serves.
+    points:
+      gpu_kv:
+        label: "GPU-resident KV"
+        description: "KV cache in HBM only, for up to about 16 sessions per worker."
+        default: true
+        args: ["--max-num-seqs", "32", "--gpu-memory-utilization", "0.88"]
+        select: { kv_offload: "off" }
+      dram_kv:
+        label: "Host-DRAM KV offload"
+        description: "Adds a host-DRAM KV tier (360 GB per worker) for about 32-40 sessions per worker."
+        args: ["--max-num-seqs", "64", "--gpu-memory-utilization", "0.82", "--kv-cache-metrics"]
+        select: { kv_offload: simple }
+        json_args:
+          "--kv-transfer-config": { kv_connector_extra_config: { cpu_bytes_to_use_per_rank: null, cpu_bytes_to_use: 360000000000, lazy_offload: true } }
+    runbook: |
+      The GB200 aggregated AgentX lane in [SemiAnalysis InferenceX](https://inferencex.semianalysis.com/) (`minimaxm31flash-fp4-gb200-dynamo-vllm-agentic-dspark-agg`, [SemiAnalysisAI/InferenceX#4632](https://github.com/SemiAnalysisAI/InferenceX/pull/4632)) serves the default checkpoint text-only, with one TP4 `dynamo.vllm` worker per 4-GPU tray behind the Dynamo frontend. Add trays to scale out: the KV router spreads conversations across them. At 16 sessions per tray (GPU-resident KV) the lane measured 29.3k tok/s/GPU at 402 tok/s/user.
+
+      **Image and versions.** The lane used an arm64 CUDA 13 build of the `vllm/vllm-openai:minimax-m31-flash` image (tag `minimax-m31-flash-arm64-cu134`); any build of that image family for GB200 will do. Dynamo `1.6.0.dev20261023` was installed on top: `pip install --extra-index-url https://pypi.nvidia.com ai-dynamo-runtime==1.6.0.dev20261023 ai-dynamo==1.6.0.dev20261023`.
+
+      **Tuning rules.**
+
+      - `--max-num-seqs` is about twice the sessions one worker serves: 32 up to 16 sessions, 64 at 32-40.
+      - **GPU-resident KV** (`--gpu-memory-utilization 0.88`) is enough up to about 16 sessions per tray. Beyond that, multi-turn contexts outgrow HBM and evicted prefixes are recomputed.
+      - **Host-DRAM KV offload** adds `SimpleCPUOffloadConnector` with lazy offload. `cpu_bytes_to_use` is the budget for the whole worker, split across its four ranks: 360 GB pins 90 GB per rank. Size it from the host DRAM that is free after the engine processes are up. Lower `--gpu-memory-utilization` to 0.82 to leave room for the transfer buffers. `--kv-cache-metrics` exports KV residency metrics, which help with that sizing; it can be left out.
+      - The CUDA graph sizes go up to 8192 tokens, so prefill chunks of up to 8192 tokens also replay graphs. `--max-num-batched-tokens 32768` still allows larger mixed steps, which run piecewise.
+
+      **Running under Dynamo.** Start etcd and NATS (32 MB max payload) once:
+
+      ```bash
+      etcd --listen-client-urls http://0.0.0.0:2379 --advertise-client-urls http://<infra-host>:2379
+      printf 'max_payload: 33554432\n' > nats.conf && nats-server -c nats.conf
+      ```
+
+      On each tray, export `ETCD_ENDPOINTS=http://<infra-host>:2379` and `NATS_SERVER=nats://<infra-host>:4222` together with the generated environment and `DYN_TCP_CONNECT_TIMEOUT=120`, `DYN_BACKEND_ADMISSION_ENGINE_REQUEST_LIMIT=1024`, `DYN_BACKEND_ADMISSION_ENGINE_REQUEST_WAITING_LIMIT=1024` and `DYN_BACKEND_ADMISSION_REQUEST_QUEUE_LIMIT=2048`. Then run the generated flags with `python3 -m dynamo.vllm --model MiniMaxAI/MiniMax-M3.1-Flash --served-model-name MiniMaxAI/MiniMax-M3.1-Flash --trust-remote-code` in place of `vllm serve MiniMaxAI/MiniMax-M3.1-Flash`. `dynamo.vllm` rejects vLLM's API-server flags, so leave out `--tool-call-parser minimax_m3 --enable-auto-tool-choice` and add:
+
+      ```bash
+      --dyn-tool-call-parser minimax_m3 --dyn-reasoning-parser minimax_m3 \
+        --kv-events-config '{"publisher":"zmq","topic":"kv-events","enable_kv_cache_events":true,"endpoint":"tcp://*:5200"}'
+      ```
+
+      Start the frontend:
+
+      ```bash
+      export ETCD_ENDPOINTS=http://<infra-host>:2379 NATS_SERVER=nats://<infra-host>:4222
+      export DYN_ROUTER_ACTIVE_REQUEST_EXPIRY_SECS=3600 DYN_TCP_CONNECT_TIMEOUT=120 DYN_TOKENIZER=fastokens
+      python3 -m dynamo.frontend --http-port 8000 --router-mode kv --router-session-affinity-ttl-secs 14400
+      ```
+
+      `DYN_TOKENIZER=fastokens` switches the frontend to the fastokens tokenizer backend that ships with Dynamo. With 100K+ token agentic prompts, the default tokenizer spends about 150 ms per request on the frontend (up to 600 ms at p90), all of it added to TTFT; fastokens takes about 1 ms.
+
+      The KV router uses the workers' KV events to send each request to the worker that holds its prefix. Session affinity keeps a conversation on that worker for up to four hours; send a stable conversation ID in the `X-Dynamo-Session-ID` request header.
+
+      The InferenceX throughput points replaced block verification with synthetic rejection at the committed golden DSpark acceptance length (4.04), so every point is measured at the same acceptance rate. Production serving keeps the default rejection sampler (`rejection_sample_method: block`), as in the command above.
+
+  gb200_agentx_pd:
+    label: "AgentX PD (Dynamo)"
+    hardware: gb200
+    strategy: pd_cluster
+    variants: [default]
+    workload: "Agentic (InferenceX AgentX)"
+    description: "DEP4 prefill workers, one per GB200 tray, feed one DEP8 decode worker (throughput) or DEP16 (lower latency) over NIXL. Both roles share a Mooncake KV store in host DRAM. The Dynamo KV router routes with session affinity. Seven-token DSpark draft on FLASHINFER_TS, DeepGEMM MegaMoE on both roles."
+    source: { label: "InferenceX #4632", href: "https://github.com/SemiAnalysisAI/InferenceX/pull/4632", key: "minimaxm31flash-fp4-gb200-dynamo-vllm-agentic-dspark-disagg", revision: 936213146b79b55e8a11ffbc1d6010145995c6a6 }
+    image:
+      vllm: "vllm/vllm-openai:minimax-m31-flash-arm64-cu134"
+      dynamo: "1.6.0.dev20261023"
+      pip: ["ai-dynamo-runtime==1.6.0.dev20261023", "ai-dynamo==1.6.0.dev20261023", "mooncake-transfer-engine-cuda13==0.3.13.post1"]
+    # Dynamo workers take engine flags only; the tool-call parser moves to
+    # --dyn-tool-call-parser (see the runbook).
+    select:
+      pd_router: dynamo
+      features: { text_only: true, tool_calling: false }
+    args:
+      - "--enable-expert-parallel"
+      - "--moe-backend"
+      - "deep_gemm_mega_moe"
+      - "--numa-bind"
+      - "--enable-prefix-caching"
+      - "--max-model-len"
+      - "1048576"
+      - "--load-format"
+      - "fastsafetensors"
+    # The RDMA names are a GB200 tray's four NICs; use the host's own.
+    env:
+      PYTHONHASHSEED: "0"
+      VLLM_USE_V2_MODEL_RUNNER: "1"
+      VLLM_ENGINE_READY_TIMEOUT_S: "7200"
+      VLLM_HTTP_TIMEOUT_KEEP_ALIVE: "600"
+      VLLM_RANDOMIZE_DP_DUMMY_INPUTS: "1"
+      NCCL_NET_GDR_C2C: "1"
+      NCCL_IB_HCA: "mlx5_0,mlx5_1,mlx5_2,mlx5_3"
+      NVSHMEM_REMOTE_TRANSPORT: "none"
+      NVSHMEM_ENABLE_NIC_PE_MAPPING: "1"
+      NVSHMEM_HCA_PE_MAPPING: "mlx5_0:1:1,mlx5_1:1:1,mlx5_2:1:1,mlx5_3:1:1"
+      UCX_TLS: "rc,cuda_copy"
+      UCX_NET_DEVICES: "mlx5_0:1,mlx5_1:1,mlx5_2:1,mlx5_3:1"
+      UCX_CUDA_IPC_ENABLE_MNNVL: "n"
+      UCX_MEMTYPE_CACHE: "n"
+      UCX_MEMTYPE_REG_WHOLE: "n"
+      UCX_RNDV_PIPELINE_ERROR_HANDLING: "y"
+      NIXL_UCX_SGL_ENABLE: "1"
+      WITH_NVIDIA_PEERMEM: "0"
+      MC_SLICE_SIZE: "1048576"
+      MC_MAX_MR_SIZE: "4294967296"
+      MC_ENABLE_DEST_DEVICE_AFFINITY: "1"
+      DYN_TCP_CONNECT_TIMEOUT: "120"
+      DYN_BACKEND_ADMISSION_ENGINE_REQUEST_LIMIT: "1024"
+      DYN_BACKEND_ADMISSION_ENGINE_REQUEST_WAITING_LIMIT: "1024"
+      DYN_BACKEND_ADMISSION_REQUEST_QUEUE_LIMIT: "2048"
+    roles:
+      prefill:
+        # One DEP4 worker per tray; run as many trays as the load needs.
+        nodes: 1
+        parallelism: dep
+        args:
+          - "--data-parallel-size"
+          - "4"
+          - "--data-parallel-size-local"
+          - "4"
+          - "--max-num-seqs"
+          - "64"
+          - "--max-num-batched-tokens"
+          - "16384"
+          - "--long-prefill-token-threshold"
+          - "4096"
+          - "--gpu-memory-utilization"
+          - "0.92"
+          - "--compilation-config"
+          - '{"cudagraph_mode":"FULL_AND_PIECEWISE"}'
+          - "--max-cudagraph-capture-size"
+          - "2048"
+          - "--kv-events-config"
+          - '{"publisher":"zmq","topic":"kv-events","enable_kv_cache_events":true,"endpoint":"tcp://*:5200"}'
+          - "--kv-transfer-config"
+          - '{"kv_connector":"MultiConnector","kv_role":"kv_both","kv_load_failure_policy":"recompute","kv_connector_extra_config":{"connectors":[{"kv_connector":"NixlConnector","kv_role":"kv_both","kv_load_failure_policy":"fail","kv_buffer_device":"cuda","kv_connector_extra_config":{"backends":["UCX"],"enforce_handshake_compat":false,"kv_lease_duration":120}},{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_both","kv_load_failure_policy":"recompute","kv_connector_extra_config":{"load_async":true,"lookup_async":true,"enable_offload":false}}]}}'
+        env:
+          VLLM_MOONCAKE_LOAD_RECV_THREADS: "8"
+      decode:
+        nodes: 2
+        parallelism: dep
+        # 64 sequences per DP rank; decode graphs cover 64 x (1 + 7) tokens.
+        # The decode ranks only add capacity to the Mooncake store.
+        remove_args: ["--no-enable-prefix-caching"]
+        args:
+          - "--max-num-seqs"
+          - "64"
+          - "--max-num-batched-tokens"
+          - "8192"
+          - "--watermark"
+          - "0.05"
+          - "--gpu-memory-utilization"
+          - "0.88"
+          - "--compilation-config"
+          - '{"cudagraph_mode":"FULL_DECODE_ONLY"}'
+          - "--max-cudagraph-capture-size"
+          - "512"
+          - "--numa-bind-nodes"
+          - "0"
+          - "0"
+          - "1"
+          - "1"
+          - "--kv-transfer-config"
+          - '{"kv_connector":"MultiConnector","kv_role":"kv_both","kv_load_failure_policy":"recompute","kv_connector_extra_config":{"connectors":[{"kv_connector":"NixlConnector","kv_role":"kv_both","kv_load_failure_policy":"fail","kv_buffer_device":"cuda","kv_connector_extra_config":{"backends":["UCX"],"enforce_handshake_compat":false,"kv_lease_duration":120}},{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_consumer","kv_load_failure_policy":"recompute","kv_connector_extra_config":{"load_async":true,"lookup_async":true,"enable_offload":false,"save_decode_cache":false,"enable_lookup":false}}]}}'
+        env:
+          VLLM_MOONCAKE_LOAD_RECV_THREADS: "4"
+          VLLM_ALLREDUCE_USE_SYMM_MEM: "0"
+          NCCL_CUMEM_ENABLE: "1"
+          NCCL_MNNVL_ENABLE: "1"
+          NCCL_NVLS_ENABLE: "1"
+    points:
+      dep8:
+        label: "DEP8 decode (throughput)"
+        description: "One DEP8 decode worker across two trays."
+        default: true
+      dep16:
+        label: "DEP16 decode (lower latency)"
+        description: "One DEP16 decode worker across four trays."
+        roles: { decode: { nodes: 4 } }
+    runbook: |
+      The GB200 disaggregated AgentX lane in [SemiAnalysis InferenceX](https://inferencex.semianalysis.com/) (`minimaxm31flash-fp4-gb200-dynamo-vllm-agentic-dspark-disagg`, [SemiAnalysisAI/InferenceX#4632](https://github.com/SemiAnalysisAI/InferenceX/pull/4632)) serves the default checkpoint text-only with Dynamo and vLLM. Several prefill workers, each DEP4 on one 4-GPU tray, feed one decode worker that is DEP8 across two trays or DEP16 across four. KV moves from prefill to decode over NIXL (UCX over RDMA), and every rank contributes host DRAM to a shared Mooncake KV store, so a returning conversation's prefix is loaded from the store instead of recomputed. Both roles use expert parallelism with DeepGEMM MegaMoE, attention data parallelism, the V2 model runner and the seven-token DSpark draft on `FLASHINFER_TS`.
+
+      | Prefill workers (DEP4) | Decode | GPUs | Sessions | tok/s/GPU | tok/s/user | TTFT p90 | E2E latency p90 |
+      |---:|---|---:|---:|---:|---:|---:|---:|
+      | 2 | DEP8 | 16 | 224 | 94.2k | 250 | 3.0 s | 12.2 s |
+      | 2 | DEP8 | 16 | 256 | 98.8k | 244 | 5.3 s | 15.7 s |
+      | 3 | DEP8 | 20 | 256 | 86.4k | 240 | 2.4 s | 12.1 s |
+      | 3 | DEP8 | 20 | 288 | 93.2k | 230 | 3.5 s | 13.3 s |
+      | 3 | DEP16 | 28 | 280 | 70.4k | 302 | 1.9 s | 9.4 s |
+      | 3 | DEP16 | 28 | 320 | 77.3k | 291 | 2.2 s | 10.0 s |
+      | 3 | DEP16 | 28 | 360 | 85.4k | 284 | 2.5 s | 10.6 s |
+      | 4 | DEP16 | 32 | 352 | 74.5k | 284 | 1.8 s | 9.9 s |
+
+      The preset renders one prefill worker and the decode worker. Run more prefill workers with the same command, one per tray.
+
+      **Image and versions.** The lane used an arm64 CUDA 13 build of the `vllm/vllm-openai:minimax-m31-flash` image (tag `minimax-m31-flash-arm64-cu134`), with Dynamo `1.6.0.dev20261023` and Mooncake `mooncake-transfer-engine-cuda13==0.3.13.post1` installed on every worker and on the Mooncake master host: `pip install --extra-index-url https://pypi.nvidia.com ai-dynamo-runtime==1.6.0.dev20261023 ai-dynamo==1.6.0.dev20261023 mooncake-transfer-engine-cuda13==0.3.13.post1`.
+
+      ### Tuning rules
+
+      - **Prefill pool.** The prefill pool sets TTFT, so scale prefill workers with load. The lane's points kept TTFT p90 between 1.8 and 5.3 s with one DEP4 prefill tray per 85-128 concurrent agentic sessions (2 trays at 224-256 sessions, 3 at 256-360, 4 at 352). Two prefill trays per DEP8 decode give the most tokens per GPU; a third keeps TTFT p90 near 3.5 s up to 288 sessions. Add a prefill tray when TTFT p90 climbs while decode still has room.
+      - **DEP8 or DEP16 decode.** DEP8 gives the most tokens per GPU: 28-36 sessions per decode GPU at 230-250 tok/s/user. DEP16 halves the batch per rank and gives about 25% more tok/s/user (302 at 280 sessions) and holds E2E latency p90 near 10 s up to 360 sessions, at about 10-20% fewer tokens per GPU. `--max-num-seqs 64` per DP rank is a cap, not the operating point. The decode graph cap is `--max-num-seqs` x (1 + 7 draft tokens) = 512. `--watermark 0.05` keeps 5% of the decode KV cache free, so extra requests wait in the queue instead of preempting running ones. Keep decode `--max-num-batched-tokens` at 8192: a preempted request recomputes its context in chunks of that size.
+      - **Prefill batching.** `--long-prefill-token-threshold 4096` chunks long prompts into 4096-token pieces. Short prompts then share each step with them instead of waiting behind a 100K-token context, which lowers the TTFT tail. The recipe's 1P1D section above uses 8192. `--max-num-batched-tokens 16384` bounds a prefill step. Prefill captures CUDA graphs up to 2048 tokens. Prefill runs at `--gpu-memory-utilization 0.92`, decode at 0.88.
+      - **Mooncake pool size.** `global_segment_size` is registered by every DP rank, prefill and decode alike, so a tray pins four times that value. The lane used 190 GB per rank (760 GB per tray), which lets three DEP4 prefill trays serve up to 288 sessions with DEP8 and 360 with DEP16; 210 GB per rank ran the prefill trays out of host memory. Size it as (host DRAM free per tray after the engines are up - headroom for NIXL and the 4 GB local buffers) / ranks per tray. The decode ranks only add capacity to the pool (`kv_consumer` with `enable_lookup: false` and `save_decode_cache: false`): the prefill workers write and read every block.
+      - **Routing.** The KV router (`--router-mode kv`) reads the prefill workers' KV events and sends each request to the worker that already holds its prefix. Session affinity (`--router-session-affinity-ttl-secs 14400`) keeps a conversation on that worker for up to four hours. Send a stable conversation ID in the `X-Dynamo-Session-ID` header. The Mooncake store catches the prefixes that miss after a reroute.
+      - **Transport.** Cross-tray NIXL needs `UCX_TLS=rc,cuda_copy` with `UCX_NET_DEVICES` set to the tray's RDMA ports. Without them, UCX can fall back to TCP and transfers stall. Point `UCX_MODULE_DIR` at the UCX bundled with the NIXL wheel (`nixl_cu13.libs/ucx` in the image's `dist-packages`). `kv_lease_duration: 120` holds a prefill's blocks for up to 120 s until decode reads them (the default is 30 s).
+      - **NUMA.** `--numa-bind` pins each rank to its GPU's NUMA node. Decode passes the tray's map, `--numa-bind-nodes 0 0 1 1`, explicitly.
+      - `--moe-backend deep_gemm_mega_moe` is the same as the lane's `--kernel-config '{"moe_backend":"deep_gemm_mega_moe"}'`. `--prefill-schedule-interval 1`, also set by the lane, is vLLM's default.
+
+      ### Services
+
+      Start etcd, NATS and the Mooncake master once, on a host that runs no engine workers, together with the frontend. The engine trays' host DRAM then goes to the Mooncake pool.
+
+      ```bash
+      etcd --listen-client-urls http://0.0.0.0:2379 --advertise-client-urls http://<infra-host>:2379
+      printf 'max_payload: 33554432\n' > nats.conf && nats-server -c nats.conf
+      mooncake_master --rpc_port=50051 \
+        --default_kv_lease_ttl=60000 --eviction_high_watermark_ratio=0.95 --eviction_ratio=0.10
+      ```
+
+      Every worker reads the store configuration from the JSON file that `MOONCAKE_CONFIG_PATH` points to (export it with the rest of the environment):
+
+      ```json
+      {"master_server_address": "<infra-host>:50051", "metadata_server": "P2PHANDSHAKE",
+       "global_segment_size": "190GB", "local_buffer_size": "4GB", "protocol": "rdma",
+       "device_name": "mlx5_0,mlx5_1,mlx5_2,mlx5_3", "enable_offload": false}
+      ```
+
+      ### Workers
+
+      Export `ETCD_ENDPOINTS=http://<infra-host>:2379`, `NATS_SERVER=nats://<infra-host>:4222` and the generated environment on every tray, then run the generated `python3 -m dynamo.vllm` commands. Add `--served-model-name MiniMaxAI/MiniMax-M3.1-Flash --trust-remote-code --dyn-tool-call-parser minimax_m3 --dyn-reasoning-parser minimax_m3` to both roles. `dynamo.vllm` rejects vLLM's API-server flags `--tool-call-parser` and `--enable-auto-tool-choice`. Multi-node DEP decode runs one process per tray with `--data-parallel-start-rank` set to 0, 4, 8 and so on, and a shared `--data-parallel-address` and `--data-parallel-rpc-port`. `--data-parallel-hybrid-lb` lets every tray's Dynamo runtime register its own DP ranks with the frontend.
+
+      ### Frontend
+
+      ```bash
+      export ETCD_ENDPOINTS=http://<infra-host>:2379 NATS_SERVER=nats://<infra-host>:4222
+      export DYN_ROUTER_ACTIVE_REQUEST_EXPIRY_SECS=3600 DYN_TCP_CONNECT_TIMEOUT=120 DYN_TOKENIZER=fastokens
+      python3 -m dynamo.frontend --http-port 8000 --router-mode kv --router-session-affinity-ttl-secs 14400
+      ```
+
+      `DYN_TOKENIZER=fastokens` switches the frontend to the fastokens tokenizer backend that ships with Dynamo. With 100K+ token agentic prompts, the default tokenizer spends about 150 ms per request on the frontend (up to 600 ms at p90), all of it added to TTFT; fastokens takes about 1 ms.
+
+      The InferenceX throughput points replaced block verification with synthetic rejection at the committed golden DSpark acceptance length (4.04), so every point is measured at the same acceptance rate. Production serving keeps the default rejection sampler (`rejection_sample_method: block`), as in the command above.
+
+  mi355x_agentx:
+    label: "AgentX TP4"
+    hardware: mi355x
+    strategy: single_node_tp
+    variants: [mxfp4]
+    workload: "Agentic (InferenceX AgentX)"
+    description: "One TP4 engine on four of the node's eight MI355X GPUs, serving the Quark MXFP4 checkpoint: seven-token DSpark with mono decode and an FP8 KV cache up to concurrency 8, seven-token DSpark with DBO and an NVFP4 KV cache at 16 and 32, and three-token DSpark with DBO and a host-DRAM KV tier at 64."
+    source: { label: "InferenceX #4671", href: "https://github.com/SemiAnalysisAI/InferenceX/pull/4671", key: "minimaxm31flash-fp4-mi355x-vllm-agentic-dspark", revision: a328b443582c45e45eef7f0c4f133f3e173f0211 }
+    image: "vllm/vllm-openai-rocm:minimax-m31-flash@sha256:ae8681212aa5291d8c27f06e90016435e3fca144a68fd5c439231f6475f74569"
+    tp: 4
+    select:
+      features: { text_only: true, spec_decoding: true }
+    args:
+      - "--enable-prefix-caching"
+    env:
+      VLLM_HTTP_TIMEOUT_KEEP_ALIVE: "900"
+    json_args:
+      "--speculative-config": { num_speculative_tokens: 7, rejection_sample_method: block }
+    # The DBO points pass the DBO flags directly: the DBO / SP prefill feature's
+    # modes set their own thresholds (8192, or 16384 with SP on the default
+    # checkpoint only), and this lane runs the 16384-token threshold without SP.
+    points:
+      c1:
+        label: "1 (mono)"
+        concurrency: 1
+        select: { features: { mono_decode: mono } }
+        args: ["--max-num-seqs", "2", "--max-num-batched-tokens", "8192"]
+      c2:
+        label: "2 (mono)"
+        concurrency: 2
+        select: { features: { mono_decode: mono } }
+        args: ["--max-num-seqs", "4", "--max-num-batched-tokens", "8192"]
+      c4:
+        label: "4 (mono)"
+        concurrency: 4
+        select: { features: { mono_decode: mono } }
+        args: ["--max-num-seqs", "8", "--max-num-batched-tokens", "8192"]
+      c8:
+        label: "8 (mono)"
+        concurrency: 8
+        select: { features: { mono_decode: mono } }
+        args: ["--max-num-seqs", "16", "--max-num-batched-tokens", "8192"]
+      c16:
+        label: "16 (DBO)"
+        concurrency: 16
+        args: ["--max-num-seqs", "32", "--max-num-batched-tokens", "32768", "--gpu-memory-utilization", "0.85", "--enable-dbo", "--dbo-prefill-token-threshold", "16384", "--dbo-decode-token-threshold", "32768"]
+        env: { VLLM_USE_V2_MODEL_RUNNER: "1", DEBUG_HIP_DYNAMIC_QUEUES: "0" }
+      c32:
+        label: "32 (DBO)"
+        concurrency: 32
+        args: ["--max-num-seqs", "64", "--max-num-batched-tokens", "32768", "--gpu-memory-utilization", "0.85", "--enable-dbo", "--dbo-prefill-token-threshold", "16384", "--dbo-decode-token-threshold", "32768"]
+        env: { VLLM_USE_V2_MODEL_RUNNER: "1", DEBUG_HIP_DYNAMIC_QUEUES: "0" }
+      c64:
+        label: "64 (DBO, KV offload)"
+        concurrency: 64
+        select: { kv_offload: simple }
+        args: ["--max-num-seqs", "128", "--max-num-batched-tokens", "32768", "--gpu-memory-utilization", "0.9", "--enable-dbo", "--dbo-prefill-token-threshold", "16384", "--dbo-decode-token-threshold", "32768"]
+        env: { VLLM_USE_V2_MODEL_RUNNER: "1", DEBUG_HIP_DYNAMIC_QUEUES: "0", PYTHONHASHSEED: "42" }
+        json_args:
+          "--speculative-config": { num_speculative_tokens: 3 }
+          "--kv-transfer-config": { kv_connector_extra_config: { cpu_bytes_to_use_per_rank: null, cpu_bytes_to_use: 1349000000000, lazy_offload: false } }
+    runbook: |
+      The MI355X AgentX lane in [SemiAnalysis InferenceX](https://inferencex.semianalysis.com/)
+      (`minimaxm31flash-fp4-mi355x-vllm-agentic-dspark`,
+      [SemiAnalysisAI/InferenceX#4671](https://github.com/SemiAnalysisAI/InferenceX/pull/4671)),
+      on the Quark MXFP4 checkpoint with `--max-num-seqs` at twice the concurrency
+      (c64 admits 128). Throughput runs replaced block verification with synthetic
+      rejection at the committed golden DSpark acceptance length (4.04 for seven
+      tokens, 2.88 for three); the commands above keep block verification.
+
+      - **Concurrency 1-8** run mono decode with an FP8 KV cache and no DBO: a decode
+        step of up to 16 tokens runs each layer as one fused kernel, and prefill steps
+        stay at 8192 tokens.
+      - **Concurrency 16 and 32** run seven-token DSpark with DBO and an NVFP4 KV cache
+        at 0.85 memory utilization, leaving headroom for DBO prefill.
+      - **Concurrency 64** runs three-token DSpark with DBO at 0.90 and adds a host-DRAM
+        KV tier through `SimpleCPUOffloadConnector`, so cached prefixes are not lost
+        once the working set outgrows the GPU KV cache (14.6M tokens at 0.90).
+        `cpu_bytes_to_use` (1,349 GB) is the budget for the whole engine.
+
+      [InferenceX run 41003224750](https://github.com/SemiAnalysisAI/InferenceX/actions/runs/41003224750)
+      passed all seven performance points and GSM8K at concurrency 8, 32 and 64 at
+      `a328b4435` (single run):
+
+      | Concurrency | Interactivity p50 / p90 (tokens/s/user) | Total tokens/s per GPU |
+      | --- | ---: | ---: |
+      | 1 | 551.0 / 515.7 | 6,766 |
+      | 2 | 548.9 / 410.4 | 6,918 |
+      | 4 | 546.3 / 395.6 | 9,494 |
+      | 8 | 407.4 / 240.0 | 15,411 |
+      | 16 | 228.6 / 168.3 | 27,275 |
+      | 32 | 172.3 / 115.1 | 54,158 |
+      | 64 | 71.4 / 47.3 | 75,434 |
+
+      GSM8K exact match (strict) over 1319 questions was 0.964 at concurrency 8, 0.959 at
+      32 and 0.957 at 64.
+
+      - **Concurrency 64 needed a rerun.** The first attempt hung at its first load from
+        the host KV tier (23 fetches in progress, none completing; `sample_tokens` RPC
+        timeout) and was rerun once, which passed. Treat the offload point as less
+        proven than the others.
+      - **Sequence-parallel prefill does not apply to this checkpoint.** InferenceX also
+        set `VLLM_ROCM_MINIMAX_M3_SP_PREFILL=1` at concurrency 64, and the server logged
+        sequence-parallel prefill as off; it is omitted here. SP prefill is available on
+        the default checkpoint through the DBO / SP prefill feature.
+
 guide: |
   ## Overview
 
diff --git a/recipes/models/Qwen/Qwen3-30B-A3B.yaml b/recipes/models/Qwen/Qwen3-30B-A3B.yaml
index ae97eb6fa1f7..b07ad144c2ae 100644
--- a/recipes/models/Qwen/Qwen3-30B-A3B.yaml
+++ b/recipes/models/Qwen/Qwen3-30B-A3B.yaml
@@ -12,7 +12,6 @@ meta:
     - "Qwen/Qwen3-235B-A22B-Instruct-2507"
   hardware:
     xeon6: verified
-    ascend_910c: verified
     arc_pro_b70: verified
 
 model:
@@ -48,59 +47,12 @@ features:
 opt_in_features:
   - long_context
 
-hardware_opt_in_features:
-  ascend_910c:
-    - tool_calling
-    - reasoning
-
 variants:
   default:
     precision: bf16
     vram_minimum_gb: 74
     description: "Full precision BF16"
     supported_hardware: [xeon6, arc_pro_b70]
-  ascend_w8a8:
-    model_id: "Eco-Tech/Qwen3-30B-A3B-w8a8"
-    precision: int8
-    label: "W8A8 (Ascend A3)"
-    vram_minimum_gb: 37
-    description: "ModelSlim INT8 W8A8 checkpoint for Ascend 910C, verified with TP4 + expert parallelism."
-    supported_hardware: [ascend_910c]
-    tp: 4
-    min_vllm_version: "0.23.0"
-    extra_args:
-      - "--trust-remote-code"
-      - "--max-num-seqs"
-      - "100"
-      - "--max-model-len"
-      - "40960"
-      - "--max-num-batched-tokens"
-      - "16384"
-      - "--enable-expert-parallel"
-      - "--quantization"
-      - "ascend"
-      - "--distributed_executor_backend"
-      - "mp"
-      - "--no-enable-prefix-caching"
-      - "--async-scheduling"
-      - "--compilation-config"
-      - '{"cudagraph_mode": "FULL_DECODE_ONLY"}'
-      - "--additional-config"
-      - '{"enable_flashcomm1": true, "weight_nz_mode": 2}'
-      - "--gpu-memory-utilization"
-      - "0.95"
-    extra_env:
-      VLLM_USE_MODELSCOPE: "True"
-      HCCL_OP_EXPANSION_MODE: "AIV"
-      HCCL_BUFFSIZE: "1024"
-      OMP_PROC_BIND: "false"
-      OMP_NUM_THREADS: "1"
-      TASK_QUEUE_ENABLE: "1"
-      PYTORCH_NPU_ALLOC_CONF: "expandable_segments:True"
-      VLLM_USE_RUST_FRONTEND: "1"
-    hardware_overrides:
-      ascend_910c:
-        docker_image: "swr.cn-southwest-2.myhuaweicloud.com/base_image/ascend-ci/vllm-ascend:v0.23.0-a3"
 
 compatible_strategies:
   - single_node_tp
diff --git a/recipes/models/Qwen/Qwen3.5-397B-A17B.yaml b/recipes/models/Qwen/Qwen3.5-397B-A17B.yaml
index 061dfa53f14d..d46f798ea66d 100644
--- a/recipes/models/Qwen/Qwen3.5-397B-A17B.yaml
+++ b/recipes/models/Qwen/Qwen3.5-397B-A17B.yaml
@@ -4,7 +4,7 @@ meta:
   provider: "Qwen"
   description: "Multimodal MoE model with gated delta networks architecture, 397B total / 17B active parameters, up to 262K context"
   date_added: 2026-02-16
-  date_updated: 2026-10-07
+  date_updated: 2026-10-29
   tasks:
     - multimodal
     - text
@@ -18,7 +18,6 @@ meta:
     mi325x: verified
     mi350x: verified
     mi355x: verified
-    ascend_950dt: verified
 
 model:
   model_id: "Qwen/Qwen3.5-397B-A17B"
@@ -142,100 +141,6 @@ variants:
       - "ROCM_AITER_FA"
       - "--max-num-seqs"
       - "64"
-  ascend_w8a8_mtp:
-    model_id: "Eco-Tech/Qwen3.5-397B-A17B-w8a8-mtp"
-    precision: int8
-    label: "W8A8 (Ascend)"
-    vram_minimum_gb: 477
-    supported_hardware: [ascend_950dt]
-    tp: 8
-    min_vllm_version: "0.23.0"
-    default_modes:
-      spec_decoding: qwen3_5_mtp
-    extra_args:
-      - "--distributed-executor-backend"
-      - "mp"
-      - "--data-parallel-size"
-      - "1"
-      - "--quantization"
-      - "ascend"
-      - "--served-model-name"
-      - "qwen3.5"
-      - "--max-num-seqs"
-      - "128"
-      - "--max-model-len"
-      - "133000"
-      - "--max-num-batched-tokens"
-      - "8192"
-      - "--gpu-memory-utilization"
-      - "0.95"
-      - "--enable-prefix-caching"
-      - "--async-scheduling"
-      - "--compilation-config"
-      - '{"cudagraph_mode":"FULL_DECODE_ONLY"}'
-      - "--additional-config"
-      - '{"enable_cpu_binding":true}'
-    extra_env:
-      VLLM_USE_MODELSCOPE: "True"
-      PYTORCH_NPU_ALLOC_CONF: "expandable_segments:True"
-      HCCL_OP_EXPANSION_MODE: "AIV"
-      HCCL_BUFFSIZE: "400"
-      OMP_PROC_BIND: "false"
-      OMP_NUM_THREADS: "100"
-      TASK_QUEUE_ENABLE: "1"
-      VLLM_ASCEND_ENABLE_PREFETCH_MLP: "1"
-      HCCL_INTRA_PCIE_ENABLE: "1"
-      HCCL_INTRA_ROCE_ENABLE: "0"
-    hardware_overrides:
-      ascend_950dt:
-        docker_image: "quay.io/ascend/vllm-ascend:v0.23.0-950DT"
-  ascend_w4a4_mxfp4:
-    model_id: "Eco-Tech/Qwen3.5-397B-A17B-w4a4-mxfp4"
-    precision: mxfp4
-    label: "W4A4 (Ascend)"
-    vram_minimum_gb: 239
-    supported_hardware: [ascend_950dt]
-    tp: 8
-    min_vllm_version: "0.23.0"
-    default_modes:
-      spec_decoding: qwen3_5_mtp
-    extra_args:
-      - "--distributed-executor-backend"
-      - "mp"
-      - "--data-parallel-size"
-      - "1"
-      - "--quantization"
-      - "ascend"
-      - "--served-model-name"
-      - "qwen3.5"
-      - "--max-num-seqs"
-      - "128"
-      - "--max-model-len"
-      - "133000"
-      - "--max-num-batched-tokens"
-      - "8192"
-      - "--gpu-memory-utilization"
-      - "0.95"
-      - "--enable-prefix-caching"
-      - "--async-scheduling"
-      - "--compilation-config"
-      - '{"cudagraph_mode":"FULL_DECODE_ONLY"}'
-      - "--additional-config"
-      - '{"enable_cpu_binding":true}'
-    extra_env:
-      VLLM_USE_MODELSCOPE: "True"
-      PYTORCH_NPU_ALLOC_CONF: "expandable_segments:True"
-      HCCL_OP_EXPANSION_MODE: "AIV"
-      HCCL_BUFFSIZE: "400"
-      OMP_PROC_BIND: "false"
-      OMP_NUM_THREADS: "100"
-      TASK_QUEUE_ENABLE: "1"
-      VLLM_ASCEND_ENABLE_PREFETCH_MLP: "1"
-      HCCL_INTRA_PCIE_ENABLE: "1"
-      HCCL_INTRA_ROCE_ENABLE: "0"
-    hardware_overrides:
-      ascend_950dt:
-        docker_image: "quay.io/ascend/vllm-ascend:v0.23.0-950DT"
 
 compatible_strategies:
   - single_node_tp
@@ -265,6 +170,57 @@ strategy_overrides:
       env:
         VLLM_SSM_CONV_STATE_LAYOUT: "DS"
 
+presets:
+  b200_tp2_kv_offload:
+    label: "TP2 + KV offload (multi-turn agentic)"
+    hardware: b200
+    strategy: single_node_tp
+    variants: [nvfp4]
+    workload: "Multi-turn agentic coding, ~64k input, KV-reuse-bound (mean ITL under 20 ms)"
+    description: "One TP2 engine with expert parallelism on two of the node's eight B200 GPUs: MTP with three draft tokens, FP8 KV cache, float16 Mamba SSM cache, async scheduling, and a 160 GiB host KV tier so prefixes survive across turns; admission is capped at 48 sequences to hold mean ITL under 20 ms."
+    source: { label: "vLLM #64607", href: "https://github.com/vllm-project/vllm/pull/64607" }
+    image: "vllm/vllm-openai:v0.32.0"
+    tp: 2
+    select:
+      frontend: python
+      features: { spec_decoding: mtp, encoder_parallel: false }
+    args:
+      - "--enable-expert-parallel"
+      - "--async-scheduling"
+      - "--block-size"
+      - "64"
+      - "--mamba-ssm-cache-dtype"
+      - "float16"
+      - "--max-num-batched-tokens"
+      - "16384"
+      # Admission stays at 48 sequences so mean ITL holds under 20 ms up to 64 clients; the rest queue.
+      - "--max-num-seqs"
+      - "48"
+      # Qwen3.5 mixes Gated DeltaNet and attention layers; the hybrid manager keeps prefix caching across both.
+      - "--no-disable-hybrid-kv-cache-manager"
+      # vLLM sums --kv-offloading-size over the TP ranks: 80 GiB of host KV per rank.
+      - "--kv-offloading-backend"
+      - "native"
+      - "--kv-offloading-size"
+      - "160"
+    runbook: |
+      Measured by NVIDIA on vLLM v0.32.0 (`vllm/vllm-openai:v0.32.0`) on two B200 SXM 180 GB with a
+      942-request multi-turn agentic coding trace (about 64k mean input tokens, most prefix tokens
+      reused across turns, outputs capped at 12k tokens), replayed at fixed concurrency with real
+      MTP acceptance (mean acceptance length 3.2 to 3.3). Host: 512 GiB RAM and 200 GiB /dev/shm per
+      server for the 160 GiB offload tier.
+
+      | Concurrency | Output tok/s/GPU | Mean ITL | Mean TTFT |
+      |---|---:|---:|---:|
+      | 16 | 936 | 8.7 ms | 0.70 s |
+
+      Further concurrencies and the comparison against this recipe's baseline command on the same
+      image and trace are being collected and will follow in a separate update.
+
+      The 48-sequence cap trades time to first token for inter-token latency at 64 clients; raise
+      `--max-num-seqs` if TTFT matters more than ITL for your traffic. The offload tier only helps
+      workloads that revisit long prefixes; drop the `--kv-offloading-*` flags for single-turn traffic.
+
 guide: |
   ## Overview
 
diff --git a/recipes/models/Qwen/Qwen3.6-27B.yaml b/recipes/models/Qwen/Qwen3.6-27B.yaml
index 2129e22b2c63..63e4418af82c 100644
--- a/recipes/models/Qwen/Qwen3.6-27B.yaml
+++ b/recipes/models/Qwen/Qwen3.6-27B.yaml
@@ -16,8 +16,6 @@ meta:
     rtx_pro_6000: verified
     dgx_station_gb300: verified
     rtx_5090: verified
-    ascend_950dt: verified
-    ascend_910c: verified
     arc_pro_b60: verified
     arc_pro_b70: verified
     mi300x: verified
@@ -102,39 +100,11 @@ variants:
     precision: bf16
     vram_minimum_gb: 65
     description: "Full precision BF16 — fits on 1x H200, 2x H100, 4x Intel Arc Pro B60, 2x Intel Arc Pro B70, or an Ascend A3 node (TP2+DP8)"
-    tp:
-      ascend_910c: 2
     hardware_overrides:
       rtx_pro_6000:
         extra_args:
           - "--max-num-seqs"
           - "8"
-      ascend_910c:
-        docker_image: "swr.cn-southwest-2.myhuaweicloud.com/base_image/ascend-ci/vllm-ascend:v0.23.0-a3"
-        extra_args:
-          - "--data-parallel-size"
-          - "8"
-          - "--served-model-name"
-          - "qwen3.6"
-          - "--max-num-seqs"
-          - "32"
-          - "--max-model-len"
-          - "133000"
-          - "--max-num-batched-tokens"
-          - "16384"
-          - "--gpu-memory-utilization"
-          - "0.90"
-          - "--no-enable-prefix-caching"
-          - "--compilation-config"
-          - '{"cudagraph_mode":"FULL_DECODE_ONLY"}'
-          - "--async-scheduling"
-        extra_env:
-          VLLM_USE_MODELSCOPE: "True"
-          PYTORCH_NPU_ALLOC_CONF: "expandable_segments:True"
-          HCCL_BUFFSIZE: "512"
-          OMP_PROC_BIND: "false"
-          OMP_NUM_THREADS: "1"
-          TASK_QUEUE_ENABLE: "1"
   fp8:
     model_id: "Qwen/Qwen3.6-27B-FP8"
     precision: fp8
@@ -247,90 +217,6 @@ variants:
     label: "FP8 (RedHat)"
     vram_minimum_gb: 33
     description: "RedHatAI FP8 dynamic quantization"
-  ascend_w8a8_mxfp8:
-    model_id: "Eco-Tech/Qwen3.6-27B-w8a8-mxfp8"
-    precision: mxfp8
-    label: "W8A8 (Ascend)"
-    vram_minimum_gb: 34
-    supported_hardware: [ascend_950dt]
-    tp: 1
-    min_vllm_version: "0.23.0"
-    default_modes:
-      spec_decoding: qwen3_5_mtp
-    extra_args:
-      - "--distributed-executor-backend"
-      - "mp"
-      - "--data-parallel-size"
-      - "8"
-      - "--quantization"
-      - "ascend"
-      - "--served-model-name"
-      - "qwen3.6"
-      - "--max-num-seqs"
-      - "32"
-      - "--max-model-len"
-      - "262144"
-      - "--max-num-batched-tokens"
-      - "16384"
-      - "--gpu-memory-utilization"
-      - "0.90"
-      - "--no-enable-prefix-caching"
-      - "--compilation-config"
-      - '{"cudagraph_mode":"FULL_DECODE_ONLY"}'
-      - "--additional-config"
-      - '{"enable_cpu_binding":true}'
-    extra_env:
-      VLLM_USE_MODELSCOPE: "True"
-      PYTORCH_NPU_ALLOC_CONF: "expandable_segments:True"
-      HCCL_BUFFSIZE: "512"
-      OMP_PROC_BIND: "false"
-      OMP_NUM_THREADS: "1"
-      TASK_QUEUE_ENABLE: "1"
-    hardware_overrides:
-      ascend_950dt:
-        docker_image: "quay.io/ascend/vllm-ascend:v0.23.0-a5"
-  ascend_w8a8:
-    model_id: "Eco-Tech/Qwen3.6-27B-w8a8"
-    precision: int8
-    label: "W8A8 (Ascend A3)"
-    vram_minimum_gb: 33
-    supported_hardware: [ascend_910c]
-    tp: 2
-    min_vllm_version: "0.23.0"
-    default_modes:
-      spec_decoding: qwen3_5_mtp
-    extra_args:
-      - "--data-parallel-size"
-      - "8"
-      - "--quantization"
-      - "ascend"
-      - "--served-model-name"
-      - "qwen3.6"
-      - "--max-num-seqs"
-      - "32"
-      - "--max-model-len"
-      - "133000"
-      - "--max-num-batched-tokens"
-      - "16384"
-      - "--gpu-memory-utilization"
-      - "0.90"
-      - "--no-enable-prefix-caching"
-      - "--compilation-config"
-      - '{"cudagraph_mode":"FULL_DECODE_ONLY"}'
-      - "--additional-config"
-      - '{"enable_cpu_binding":true}'
-      - "--async-scheduling"
-    extra_env:
-      VLLM_USE_MODELSCOPE: "True"
-      PYTORCH_NPU_ALLOC_CONF: "expandable_segments:True"
-      HCCL_BUFFSIZE: "512"
-      OMP_PROC_BIND: "false"
-      OMP_NUM_THREADS: "1"
-      TASK_QUEUE_ENABLE: "1"
-    hardware_overrides:
-      ascend_910c:
-        docker_image: "swr.cn-southwest-2.myhuaweicloud.com/base_image/ascend-ci/vllm-ascend:v0.23.0-a3"
-    description: "ModelSlim INT8 W8A8 checkpoint for one Atlas 800 A3 node (16x Ascend 910C), using TP2+DP8. Requires --quantization ascend."
 
 compatible_strategies:
   - single_node_tp
diff --git a/recipes/models/Qwen/Qwen3.8-27B.yaml b/recipes/models/Qwen/Qwen3.8-27B.yaml
index 1f58d069472a..bed45a33b168 100644
--- a/recipes/models/Qwen/Qwen3.8-27B.yaml
+++ b/recipes/models/Qwen/Qwen3.8-27B.yaml
@@ -28,11 +28,6 @@ meta:
     rtx_5090_2x: verified
     rtx_pro_6000: verified
     dgx_spark_gb10: verified
-    # Huawei Ascend NPU, verified 2026-08-18 on vLLM Ascend 0.23.0 (image
-    # quay.io/ascend/vllm-ascend:qwen3.8-a5). The `ascend_w8a8` variant booted on a
-    # single 950PR at TP1 with the in-checkpoint MTP head and FULL_DECODE_ONLY
-    # ACLGraph and answered a long-form prompt correctly.
-    ascend_950pr: verified
     arc_pro_b60: verified
     arc_pro_b70: verified
 
@@ -189,14 +184,7 @@ variants:
     # (quantization_config.weight_block_size [128, 128]).
     vram_minimum_gb: 38
     description: "Official block-scaled FP8 checkpoint — 28.7 GiB: 1 GPU."
-    # ascend_950pr added 2026-08-26: vllm-project/vllm-ascend#14852 dispatches
-    # `quant_method: fp8` + `weight_block_size` in AscendFp8Config.get_quant_method and
-    # re-quantizes the block-scaled tiles to MXFP8 at load time on 950. No
-    # `--quantization ascend` — the scheme is picked up from the checkpoint config.
-    supported_hardware: [h100, h200, b200, b300, gb200, gb300, mi300x, mi325x, mi355x, rtx_5090, rtx_5090_2x, ascend_950pr]
-    hardware_overrides:
-      ascend_950pr:
-        docker_image: "quay.io/ascend/vllm-ascend:nightly-main-a5"
+    supported_hardware: [h100, h200, b200, b300, gb200, gb300, mi300x, mi325x, mi355x, rtx_5090, rtx_5090_2x]
   nvfp4:
     model_id: "Inferact/Qwen3.8-27B-NVFP4"
     precision: nvfp4
@@ -276,40 +264,6 @@ variants:
     supported_hardware: [h100, h200, b200, b300, gb200, gb300, rtx_5090, rtx_5090_2x]
     description: "Red Hat AI compressed-tensors MXFP4 build — 19.8 GB: 1 GPU. True W4A4 on Blackwell (SM100+); W4A16 via Marlin on Hopper and elsewhere."
 
-  ascend_w8a8:
-    model_id: "Eco-Tech/Qwen3.8-27B-w8a8"
-    precision: int8
-    label: "W8A8 (Ascend)"
-    # 32,128,768,176 B across 10 shards = 32.1 GB. ModelSlim keeps the vision tower in BF16
-    # and stores FP32 per-channel scales, so this is not a uniform 8-bit build.
-    vram_minimum_gb: 39
-    supported_hardware: [ascend_950pr]
-    default_modes:
-      spec_decoding: qwen3_5_mtp
-    # `--quantization ascend` is what selects the ModelSlim scheme; without it the loader
-    # treats the checkpoint as float. ACLGraph replays decode only. CPU binding cuts
-    # cross-core scheduling overhead on the Kunpeng host.
-    extra_args:
-      - "--quantization"
-      - "ascend"
-      - "--compilation-config"
-      - '{"cudagraph_mode":"FULL_DECODE_ONLY"}'
-      - "--additional-config"
-      - '{"enable_cpu_binding":true}'
-      - "--max-num-batched-tokens"
-      - "16384"
-    extra_env:
-      PYTORCH_NPU_ALLOC_CONF: "expandable_segments:True"
-      HCCL_BUFFSIZE: "512"
-      OMP_PROC_BIND: "false"
-      OMP_NUM_THREADS: "1"
-    # `computeDockerMeta` has no Huawei brand, so a bare string `docker_image` would only pin
-    # NVIDIA. The exact variant+hardware override lands regardless of brand.
-    hardware_overrides:
-      ascend_950pr:
-        docker_image: "quay.io/ascend/vllm-ascend:qwen3.8-a5"
-    description: "ModelSlim INT8 W8A8 for Ascend NPU — 32.1 GB: 1 NPU. Requires --quantization ascend."
-
 # Dense model: no routed experts, so none of the expert-parallel or PD strategies apply.
 #
 # `multi_node_tp` is deliberately ABSENT, unlike the 2.4T sibling. Multi-node TP sets
@@ -457,74 +411,7 @@ guide: |
   official vLLM XPU image `vllm/vllm-openai-xpu:latest`. `--enforce-eager` is required on this
   backend. Select Arc Pro B60/B70 in the hardware picker above for the exact command.
 
-  ### Huawei Ascend 950PR-w8a8 (vLLM Ascend)
-
-  Verified on a single 950PR, TP1, vLLM Ascend 0.23.0 from the
-  `quay.io/ascend/vllm-ascend:qwen3.8-a5` image.
-
-  The ModelSlim INT8 build is the checkpoint to use here. `--quantization ascend` is what
-  selects the ModelSlim scheme — without it the loader treats the weights as float. The MTP
-  head inside the checkpoint works and is spelled `qwen3_5_mtp` on this backend.
-
-  Both the weight and the image come from ModelScope, not HuggingFace: set
-  `VLLM_USE_MODELSCOPE=True` or pass a local path.
-
-  ```bash
-  export VLLM_USE_MODELSCOPE=True
-  export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True
-  export HCCL_BUFFSIZE=512
-  export OMP_PROC_BIND=false
-  export OMP_NUM_THREADS=1
-
-  vllm serve Eco-Tech/Qwen3.8-27B-w8a8 \
-    --tensor-parallel-size 1 \
-    --quantization ascend \
-    --max-num-batched-tokens 16384 \
-    --enable-prefix-caching \
-    --reasoning-parser qwen3 \
-    --tool-call-parser qwen3_coder \
-    --mm-encoder-tp-mode data \
-    --speculative-config '{"method":"qwen3_5_mtp","num_speculative_tokens":3,"enforce_eager":true}' \
-    --compilation-config '{"cudagraph_mode":"FULL_DECODE_ONLY"}' \
-    --additional-config '{"enable_cpu_binding":true}'
-  ```
-
-  `--reasoning-parser qwen3` is not optional in practice: the chat template opens every
-  assistant turn with `<think>`, so without it the entire reasoning block lands in
-  `message.content` and a 2048-token budget can be spent before the answer starts.
-  
-  Measured single-stream decode: 64 tok/s (2048-token completion, one request, MTP on).
-  Read MTP acceptance from `vllm:spec_decode_num_{accepted,draft}_tokens_total` before
-  tuning `num_speculative_tokens` — throughput alone cannot tell a working drafter from
-  one that loaded and was ignored.
-  
-  ### Huawei Ascend 950PR-Native block-scaled FP8 (vLLM Ascend)
-
-  Verified on a single 950PR, TP1, vLLM Ascend 0.23.0 from the
-  `quay.io/ascend/vllm-ascend:nightly-main-a5` image.
-
-  The official `Qwen/Qwen3.8-27B-FP8` checkpoint serves directly on a 950PR — no offline
-  re-quantization and **no `--quantization ascend`**. vLLM Ascend reads `quant_method: fp8`
-  and `weight_block_size` straight from the checkpoint, expands the `weight_scale_inv`
-  tiles at load time and re-quantizes to MXFP8 on 950.
-
-  Requires a vLLM Ascend build that includes
-  [vllm-ascend#14852](https://github.com/vllm-project/vllm-ascend/pull/14852) (merged
-  2026-08-26); no tagged release carries it yet. The W8A8 path above is unaffected.
-
-  ```bash
-  vllm serve Qwen/Qwen3.8-27B-FP8 \
-    --tensor-parallel-size 1 \
-    --max-model-len 128000 \
-    --max-num-seqs 16 \
-    --gpu-memory-utilization 0.85 \
-    --trust-remote-code \
-    --enable-prefix-caching \
-    --reasoning-parser qwen3 \
-    --tool-call-parser qwen3_coder \
-    --mm-encoder-tp-mode data \
-    --speculative-config '{"method": "qwen3_5_mtp", "num_speculative_tokens": 3, "enforce_eager": true}'
-  ```
+  <!-- overlay:vllm-ascend:ascend -->
   
   ### DFlash2 (opt-in)
   
diff --git a/recipes/models/autotrust/JEV-27B-VL.yaml b/recipes/models/autotrust/JEV-27B-VL.yaml
new file mode 100644
index 000000000000..61a0ae1327ac
--- /dev/null
+++ b/recipes/models/autotrust/JEV-27B-VL.yaml
@@ -0,0 +1,200 @@
+# yaml-language-server: $schema=https://recipes.vllm.ai/recipe.schema.json
+
+meta:
+  title: "JEV-27B-VL"
+  slug: "jev-27b-vl"
+  provider: "AutoTrust"
+  description: "Dense 27B vision-language decision model on a Qwen3.8-27B backbone, with a rank-32 LoRA (jev-decision) for calibrated choices over text and images."
+  date_added: 2026-10-28
+  date_updated: 2026-10-28
+  tasks:
+    - multimodal
+    - text
+  performance_headline: "Qwen3.8-27B VL backbone; System 1 decisions through the jev-decision LoRA"
+  related_recipes:
+    - "Qwen/Qwen3.8-27B"
+
+model:
+  model_id: "autotrust/JEV-27B-VL"
+  # Model author: this checkpoint only works on vLLM >= 0.32.0.
+  # Qwen3_5ForConditionalGeneration is registered in that release.
+  min_vllm_version: "0.32.0"
+  architecture: dense
+  # safetensors parameter total is 27,781,427,952 BF16 (vision tower included).
+  parameter_count: "27B"
+  active_parameters: "27B"
+  # text_config.max_position_embeddings. The validated serve line caps the
+  # window at 32768 via --max-model-len; 256K is the backbone native length.
+  context_length: 262144
+  base_args:
+    - "--served-model-name"
+    - "autotrust/JEV-27B-VL"
+    - "--enable-lora"
+    - "--max-lora-rank"
+    - "32"
+    - "--lora-modules"
+    - "jev-decision=JEV-27B-VL/adapter_vllm"
+    - "--logprobs-mode"
+    - "processed_logprobs"
+    - "--max-model-len"
+    - "32768"
+    - "--enable-prefix-caching"
+    - "--mamba-cache-mode"
+    - "align"
+    - "--limit-mm-per-prompt"
+    - '{"image": 8}'
+    - "--max-num-seqs"
+    - "8"
+    - "--trust-request-chat-template"
+  base_env: {}
+
+# Host checkout for the LoRA. vLLM serve still pulls weights by model_id;
+# --lora-modules needs adapter_vllm on disk. Same subdirectory on the
+# ModelScope checkout. install_modes includes docker because the image
+# cannot fetch this relative path for you.
+dependencies:
+  - note: "Local layout for LoRA module jev-decision. After this download the adapter is JEV-27B-VL/adapter_vllm; a ModelScope checkout uses the same adapter_vllm directory under whatever folder you chose."
+    command: "hf download autotrust/JEV-27B-VL --local-dir JEV-27B-VL"
+    install_modes: [pip, docker]
+
+features:
+  text_only:
+    description: "Skip the vision encoder for text-only workloads and free its memory for KV cache."
+    args:
+      - "--language-model-only"
+  encoder_parallel:
+    description: "Run the vision encoder data-parallel across the tensor-parallel ranks instead of sharding it."
+    args:
+      - "--mm-encoder-tp-mode"
+      - "data"
+
+opt_in_features:
+  - text_only
+  - encoder_parallel
+
+variants:
+  default:
+    precision: bf16
+    # 55,562,855,904 bytes on disk; ceil(55.563 GB × 1.2) = 67.
+    # Uniform BF16, so params × 2 × 1.2 lands on the same budget.
+    vram_minimum_gb: 67
+    description: "BF16 weights (about 52–56 GB). The model card's quick start targets one GPU with 80 GB or more; this recipe does not mark any GPU verified."
+
+compatible_strategies:
+  - single_node_tp
+  - multi_node_tp
+  - multi_node_tp_pp
+
+# serve.sh does not set tensor-parallel-size: one GPU. 67 GB fits a single
+# high-memory GPU, so the single-node default stays TP 1 instead of the
+# whole node.
+strategy_overrides:
+  single_node_tp:
+    tp: 1
+
+guide: |
+  ## Overview
+
+  [JEV-27B-VL](https://huggingface.co/autotrust/JEV-27B-VL) is AutoTrust's
+  vision-language decision model. The checkpoint is a dense BF16
+  `Qwen3_5ForConditionalGeneration` model (about 27.8B parameters, vision
+  tower included) continued from
+  [Qwen/Qwen3.8-27B](https://huggingface.co/Qwen/Qwen3.8-27B). System 2 is
+  the served checkpoint, named `autotrust/JEV-27B-VL`. System 1 is the
+  rank-32 LoRA `jev-decision` (`adapter_vllm/`, `r` 32, `lora_alpha` 64),
+  used for calibrated choice probabilities over text or images.
+
+  This recipe has no `verified` hardware entry. The model card's quick start
+  asks for one GPU with 80 GB or more, and the card's published latency
+  figures are from a B200. Those runs are the author's, not a badge on this
+  page.
+
+  ## Prerequisites
+
+  - **vLLM:** >= 0.32.0. Earlier releases are not supported for this
+    checkpoint.
+  - **Weights:** `hf download autotrust/JEV-27B-VL --local-dir JEV-27B-VL`
+    so the LoRA path below exists. The same `adapter_vllm` directory is in
+    the [ModelScope](https://modelscope.cn/models/autotrust/JEV-27B-VL)
+    checkout; if that folder is not named `JEV-27B-VL`, point the module at
+    `<checkout>/adapter_vllm` (the HF snapshot path ends in
+    `snapshots/<rev>/adapter_vllm`).
+  - **Docker:** the image only bind-mounts the Hugging Face cache. Mount the
+    checkout into the container and keep `jev-decision` pointed at that
+    `adapter_vllm` directory.
+  - **Context:** `text_config.max_position_embeddings` is 262144. The
+    validated command uses `--max-model-len 32768` (the `serve.sh` default).
+    The card says a full 256K prompt needs about 17 GB of KV cache on top of
+    the weights, and suggests `--max-model-len 131072` on an 80 GB GPU.
+
+  ## Decision LoRA and serve flags
+
+  These flags are the model card's validated `vllm serve` line. They are
+  base arguments, not optional toggles.
+
+  - **`jev-decision`** — `--enable-lora`, `--max-lora-rank 32`, and
+    `--lora-modules jev-decision=JEV-27B-VL/adapter_vllm`. The module name
+    stays `jev-decision`. The path is the local checkout
+    `JEV-27B-VL/adapter_vllm`, or the `adapter_vllm` directory inside an
+    HF or ModelScope checkout.
+  - **`--mamba-cache-mode align`** — the Qwen3.8 hybrid backbone mixes
+    full attention with Mamba-style layers; the card serves it in `align`
+    mode.
+  - **`--limit-mm-per-prompt '{"image": 8}'`** — at most 8 images per
+    request.
+  - **`--max-num-seqs 8`** — required. The card says that with more than 8
+    sequences in one batch, vLLM's LoRA path for this multimodal class
+    returns wrong System 1 probabilities. Extra requests queue.
+  - **`--logprobs-mode processed_logprobs`** — System 1 reads post-processor
+    logprobs. On `/v1/chat/completions` or `/v1/completions` for
+    `jev-decision`, the card also says to send `top_k: 0` and `top_p: 1.0`:
+    `generation_config.json` sets `top_k` 20 and `top_p` 0.95, and those
+    defaults would truncate the returned distribution.
+  - **`--trust-request-chat-template`** — image decisions send a raw
+    decision chat template on the request. System 2 keeps the checkpoint's
+    own template.
+  - **`--served-model-name autotrust/JEV-27B-VL`** — public name for
+    System 2. `--enable-prefix-caching` matches the card's serve line.
+
+  `POST /v1/decide` is not part of stock `vllm serve`. The checkpoint ships
+  `serve_decide.py`, which starts the OpenAI server with these same flags
+  and adds that route. Stock `vllm serve` still exposes System 2 as
+  `autotrust/JEV-27B-VL` and System 1 as the LoRA model `jev-decision`.
+
+  ## Client usage
+
+  System 2 (the base checkpoint), with an image:
+
+  ```python
+  client.chat.completions.create(
+      model="autotrust/JEV-27B-VL",
+      messages=[{
+          "role": "user",
+          "content": [
+              {"type": "image_url", "image_url": {"url": image_url}},
+              {"type": "text", "text": "What does this image show?"},
+          ],
+      }],
+  )
+  ```
+
+  System 1 through the checkpoint's `/v1/decide` route (`serve_decide.py`):
+
+  ```bash
+  curl localhost:8000/v1/decide -H 'Content-Type: application/json' -d '{
+    "kind": "choice",
+    "state": "Customer: my card was charged twice for one coffee.",
+    "question": "Which team should handle this?",
+    "options": ["billing", "shipping", "tech support"]}'
+  ```
+
+  Text-only System 1 calls the LoRA directly (`model` `jev-decision`). The
+  card says those text decisions match
+  [autotrust/JEV-27B](https://huggingface.co/autotrust/JEV-27B).
+
+  ## References
+
+  - [Hugging Face model card](https://huggingface.co/autotrust/JEV-27B-VL)
+  - [ModelScope (autotrust/JEV-27B-VL)](https://modelscope.cn/models/autotrust/JEV-27B-VL)
+  - [ModelScope mirror listed on the card (AutoTrust-Labs/JEV-27B-VL)](https://www.modelscope.ai/models/AutoTrust-Labs/JEV-27B-VL)
+  - [Base model: Qwen3.8-27B](https://huggingface.co/Qwen/Qwen3.8-27B)
diff --git a/recipes/models/deepseek-ai/DeepSeek-V4-Flash-Vision-Exp.yaml b/recipes/models/deepseek-ai/DeepSeek-V4-Flash-Vision-Exp.yaml
index 0337d41783ca..eb45f98f8129 100644
--- a/recipes/models/deepseek-ai/DeepSeek-V4-Flash-Vision-Exp.yaml
+++ b/recipes/models/deepseek-ai/DeepSeek-V4-Flash-Vision-Exp.yaml
@@ -17,10 +17,6 @@ meta:
   # NVIDIA hardware with a published end-to-end run on that branch: TP4 + EP
   # on one NVL4 tray, FP8 KV cache, block size 256, DSpark k=3, full
   # 1000-sample OCRBench at 83.5% with zero request errors.
-  # Atlas A5 8-card mixed serving of the native FP4+FP8 Vision-Exp checkpoint
-  # (DP4+EP, TP=1, no speculative decoding). 950PR verified 2026-09-10 on
-  # vLLM Ascend; 950DT uses the same command (4 of 8 NPUs). See the hardware
-  # note for the required hash-router bias fix.
   # AMD MI350X (gfx950): ROCm vision merged in vllm-project/vllm#55107 (2026-09-11)
   # on the #54566 base. Verified text+vision on stock ROCm nightly at single-node
   # TP2 and TP4, with DSpark k=3 (needs VLLM_USE_V2_MODEL_RUNNER=1). See
@@ -29,8 +25,6 @@ meta:
   hardware:
     gb200: verified
     mi350x: verified
-    ascend_950dt_8x: verified
-    ascend_950pr_8x: verified
 
 model:
   model_id: "deepseek-ai/DeepSeek-V4-Flash-Vision-Exp"
@@ -74,13 +68,6 @@ features:
       - "deepseek_v4"
       - "--reasoning-config"
       - '{"reasoning_parser":"deepseek_v4","reasoning_start_str":"","reasoning_end_str":""}'
-    # Empty start/end strings make the Ascend chat path emit content=null
-    # for a full completion. Keep the parser, drop the override.
-    hardware_overrides:
-      npu:
-        args:
-          - "--reasoning-parser"
-          - "deepseek_v4"
   # Single method, so the boolean form (not `modes`): the checkpoint's fused
   # DSpark module is the only drafter with a published run on this branch. Left
   # OUT of opt_in_features on purpose — the verified GB200 configuration had
@@ -92,9 +79,6 @@ features:
     args:
       - "--speculative-config"
       - '{"method":"dspark","model":"deepseek-ai/DeepSeek-V4-Flash-Vision-Exp","num_speculative_tokens":3,"draft_sample_method":"probabilistic","enable_adaptive_verification":true}'
-    # 950 has no working MTP or DSpark path yet. Empty args so a toggle-on
-    # still emits no --speculative-config. Default-off via
-    # hardware_opt_in_features below — GB200 keeps DSpark on.
     hardware_overrides:
       # ROCm: the DeepseekV4 indexer backend rejects adaptive verification, so it
       # is forced false (DSpark still drafts k=3). VLLM_USE_V2_MODEL_RUNNER=1
@@ -103,15 +87,6 @@ features:
         args:
           - "--speculative-config"
           - '{"method":"dspark","model":"deepseek-ai/DeepSeek-V4-Flash-Vision-Exp","num_speculative_tokens":3,"draft_sample_method":"probabilistic","enable_adaptive_verification":false}'
-      npu:
-        args: []
-
-# Spec decoding stays default-on everywhere except the A5 8-card pills.
-hardware_opt_in_features:
-  ascend_950pr_8x:
-    - spec_decoding
-  ascend_950dt_8x:
-    - spec_decoding
 
 variants:
   default:
@@ -120,18 +95,6 @@ variants:
     # underestimates: sized from the real checkpoint (167.8 GB x 1.2).
     vram_minimum_gb: 202
     description: "Native FP4+FP8 mixed checkpoint (MoE experts FP4, attention/norm/router FP8) with a BF16 vision tower, aligner, and the fused DSpark draft module."
-    hardware_overrides:
-      # Exact-GPU pin: computeDockerMeta has no Huawei brand, so the model-level
-      # docker_image string only affects NVIDIA. Hide pip — official wheels
-      # have no NPU kernels.
-      ascend_950pr_8x:
-        docker_image: "quay.io/ascend/vllm-ascend:v0.23.0-a5"
-        install:
-          pip: false
-      ascend_950dt_8x:
-        docker_image: "quay.io/ascend/vllm-ascend:v0.23.0-a5"
-        install:
-          pip: false
 
 compatible_strategies:
   - single_node_tp
@@ -143,23 +106,6 @@ compatible_strategies:
 # TP + EP is the verified layout (TP4 + EP on one GB200 NVL4 tray).
 default_strategy: single_node_tep
 
-# Atlas A5 mixed serving is DP+EP only. The native checkpoint has 64
-# attention heads; 950's FP8 DSA operator accepts local num_heads_q of
-# 64 or 128, so TP must stay 1. TEP/TP and intra-node PD are not offered.
-strategy_hardware:
-  single_node_tp:
-    npu: unsupported
-  single_node_tep:
-    npu: unsupported
-  multi_node_tep:
-    npu: unsupported
-  multi_node_dep:
-    npu: unsupported
-
-# Shared by 950DT 8x (8×96 GB) and 950PR 8x (8×128 GB). Mixed serving uses
-# 4 of the 8 NPUs (DP4+EP, TP=1). HCCL_BUFFSIZE=1024 is required for MC2
-# dispatch. last-wins shadows the DEP builder's --data-parallel-size 8
-# and the recipe-level --block-size 256 / --kv-cache-dtype fp8.
 hardware_overrides:
   amd:
     extra_args:
@@ -173,49 +119,6 @@ hardware_overrides:
       VLLM_ROCM_USE_AITER: "1"
       VLLM_USE_BREAKABLE_CUDAGRAPH: "1"
       VLLM_USE_V2_MODEL_RUNNER: "1"
-  npu:
-    extra_args:
-      - "--data-parallel-size"
-      - "4"
-      - "--block-size"
-      - "128"
-      - "--kv-cache-dtype"
-      - "auto"
-      - "--tokenizer-mode"
-      - "deepseek_v4"
-      - "--safetensors-load-strategy"
-      - "prefetch"
-      - "--async-scheduling"
-      - "--max-model-len"
-      - "8192"
-      - "--max-num-seqs"
-      - "16"
-      - "--max-num-batched-tokens"
-      - "2048"
-      - "--gpu-memory-utilization"
-      - "0.9"
-      - "--api-server-count"
-      - "1"
-      - "--compilation-config"
-      - '{"cudagraph_mode":"FULL_DECODE_ONLY"}'
-      - "--additional-config"
-      - '{"enable_cpu_binding": true, "multistream_overlap_shared_expert": true, "enable_shared_expert_dp": true}'
-    extra_env:
-      PYTORCH_NPU_ALLOC_CONF: "expandable_segments:True"
-      OMP_PROC_BIND: "false"
-      OMP_NUM_THREADS: "10"
-      HCCL_BUFFSIZE: "1024"
-      HCCL_OP_EXPANSION_MODE: "AIV"
-      TASK_QUEUE_ENABLE: "1"
-      ASCEND_RT_VISIBLE_DEVICES: "0,1,2,3"
-
-hardware_notes:
-  ascend_950dt_8x:
-    title: "vLLM Ascend A5 image required"
-    body: "Mixed serving on one 8-card 950DT node uses 4 of 8 NPUs (DP4+EP, TP=1). Same command as 950PR. Pin quay.io/ascend/vllm-ascend:v0.23.0-a5 — the NVIDIA pip wheel cannot run here. Set HCCL_BUFFSIZE=1024 before launch. Native Flash-Vision checkpoints need the hash-router bias skip in vllm-project/vllm-ascend#16259 — without it load_weights aborts on model.layers.0.mlp.gate.e_score_correction_bias."
-  ascend_950pr_8x:
-    title: "vLLM Ascend A5 image required"
-    body: "Mixed serving on one 8-card 950PR node uses 4 of 8 NPUs (DP4+EP, TP=1). Same command as 950DT. Pin quay.io/ascend/vllm-ascend:v0.23.0-a5 and set HCCL_BUFFSIZE=1024 before launch. Native Flash-Vision checkpoints need the hash-router bias skip in vllm-project/vllm-ascend#16259 — without it load_weights aborts on model.layers.0.mlp.gate.e_score_correction_bias."
 
 guide: |
   ## Overview
diff --git a/recipes/models/deepseek-ai/DeepSeek-V4-Flash.yaml b/recipes/models/deepseek-ai/DeepSeek-V4-Flash.yaml
index 2ccd6acacbaa..ccbbdf16c9ba 100644
--- a/recipes/models/deepseek-ai/DeepSeek-V4-Flash.yaml
+++ b/recipes/models/deepseek-ai/DeepSeek-V4-Flash.yaml
@@ -21,17 +21,6 @@ meta:
     mi300x: verified
     mi325x: verified
     mi355x: verified
-    # Atlas A5 8-card mixed serving of FP8 (Preview). 950PR verified 2026-08-29
-    # on vLLM Ascend 0.23.0 (DP4+EP, MTP k=1). 950DT matches the v0.23.0
-    # vLLM Ascend mixed recipe (same command, 4 of 8 NPUs).
-    # 950PR additionally verified 2026-09-09 on FP8 (0731) with DSpark k=5,
-    # same DP4+EP command; 950DT follows it. See the hardware note for the
-    # required fix.
-    ascend_950dt_8x: verified
-    ascend_950pr_8x: verified
-    # Atlas 800 A3 (16× Ascend 910C, 64 GB each). W8A8 quantized checkpoint,
-    # DP2+TP8 on all 16 NPUs. Verified 2026-09-23 on vLLM Ascend v0.23.0-a3.
-    ascend_910c: verified
 
 model:
   model_id: "deepseek-ai/DeepSeek-V4-Flash"
@@ -94,13 +83,6 @@ features:
             args:
               - "--speculative-config"
               - '{"method":"mtp","num_speculative_tokens":1}'
-          # Preview MTP head is next-1. vLLM Ascend 0.23.0 runs the draft eager
-          # while the target keeps FULL_DECODE_ONLY ACLGraph. DSpark is not in
-          # this release; do not raise k or switch method on the preview weights.
-          npu:
-            args:
-              - "--speculative-config"
-              - '{"method":"mtp","num_speculative_tokens":1,"enforce_eager":true}'
       dspark:
         label: "DSpark"
         description: "DSpark drafting, 7 draft tokens, probabilistic sampling. Requires the DSpark or 0731 checkpoint (Variant row)."
@@ -117,13 +99,6 @@ features:
             args:
               - "--speculative-config"
               - '{"method":"dspark","num_speculative_tokens":5,"draft_sample_method":"probabilistic","attention_backend":"B12X_MLA_SPARSE"}'
-          # 950PR was verified at k=5 with an eager draft while the target keeps
-          # its FULL_DECODE_ONLY ACLGraph. vLLM Ascend does not implement
-          # probabilistic draft sampling, so that key is left off.
-          npu:
-            args:
-              - "--speculative-config"
-              - '{"method":"dspark","num_speculative_tokens":5,"enforce_eager":true}'
 
       dflash:
         label: "DFlash"
@@ -145,12 +120,7 @@ variants:
     precision: fp8
     vram_minimum_gb: 200
     description: "Official DeepSeek-V4-Flash release (2026-07-31) superseding the preview — substantially stronger agentic performance, with the DSpark draft module attached."
-    # 950PR does load this checkpoint's DSpark draft module: the v0.23.0-a5 tag
-    # currently ships vLLM Ascend main (0.1.dev1+g6e6aad57e), not the 0.23.0
-    # release. 950DT is offered on the same grounds the FP8 (Preview) row
-    # already uses — one A5 profile, one command, 4 of 8 NPUs — and its 4x96 GB
-    # clears this variant's 200 GB minimum. 0731 itself was run on 950PR.
-    supported_hardware: [arch:sm90, arch:sm100, arch:sm103, arch:gfx942, mi355x, rtx_pro_6000_8x, dgx_spark_gb10, ascend_950pr_8x, ascend_950dt_8x]
+    supported_hardware: [arch:sm90, arch:sm100, arch:sm103, arch:gfx942, mi355x, rtx_pro_6000_8x, dgx_spark_gb10]
     min_vllm_version: "0.25.0"
     docker_image: "vllm/vllm-openai:v0.28.0"
     default_modes:
@@ -158,16 +128,6 @@ variants:
     hardware_overrides:
       rtx_pro_6000_8x:
         docker_image: "vllm/vllm-openai:nightly"
-      # Same pinning as the FP8 (Preview) variant below: the NVIDIA wheel
-      # carries no NPU kernels, so the pip tab is hidden here.
-      ascend_950pr_8x:
-        docker_image: "quay.io/ascend/vllm-ascend:v0.23.0-a5"
-        install:
-          pip: false
-      ascend_950dt_8x:
-        docker_image: "quay.io/ascend/vllm-ascend:v0.23.0-a5"
-        install:
-          pip: false
       dgx_spark_gb10:
         docker_image: "eugr/spark-vllm-b12x:latest"
         # No wheel path here: the B12X args and env below ship only in the
@@ -296,17 +256,6 @@ variants:
         extra_args:
           - "--moe-backend"
           - "deep_gemm_mega_moe"
-      # `computeDockerMeta` has no Huawei brand, so a bare string docker_image
-      # would only pin NVIDIA. Exact-GPU overrides land regardless of brand.
-      # The pip tab is hidden: the official vLLM wheel has no NPU kernels.
-      ascend_950dt_8x:
-        docker_image: "quay.io/ascend/vllm-ascend:v0.23.0-a5"
-        install:
-          pip: false
-      ascend_950pr_8x:
-        docker_image: "quay.io/ascend/vllm-ascend:v0.23.0-a5"
-        install:
-          pip: false
   nvfp4:
     model_id: "nvidia/DeepSeek-V4-Flash-NVFP4"
     precision: nvfp4
@@ -451,59 +400,6 @@ variants:
         extra_args:
           - "--moe-backend"
           - "deep_gemm_mega_moe"
-  ascend_w8a8:
-    # W8A8 quantized checkpoint for Atlas 800 A3 (16× Ascend 910C, 64 GB each).
-    # TP8 shards the 284B weights across 8 NPUs (35.5 GB each), fitting within
-    # the 64 GB per-NPU limit. DP2 uses all 16 NPUs (2 replicas × 8 NPUs each).
-    model_id: "Eco-Tech/DeepSeek-V4-Flash-w8a8-mtp"
-    label: "W8A8 (Ascend A3)"
-    precision: int8
-    vram_minimum_gb: 341
-    description: "ModelSlim INT8 W8A8 checkpoint for one Atlas 800 A3 node (16× Ascend 910C), using DP2+TP8. Requires --quantization ascend."
-    supported_hardware: [ascend_910c]
-    tp:
-      ascend_910c: 8
-    min_vllm_version: "0.23.0"
-    extra_args:
-      - "--data-parallel-size"
-      - "2"
-      - "--quantization"
-      - "ascend"
-      - "--max-num-seqs"
-      - "32"
-      - "--max-num-batched-tokens"
-      - "10240"
-      - "--gpu-memory-utilization"
-      - "0.90"
-      - "--block-size"
-      - "128"
-      - "--compilation-config"
-      - '{"cudagraph_mode":"FULL_DECODE_ONLY"}'
-      - "--async-scheduling"
-      - "--api-server-count"
-      - "1"
-      - "--safetensors-load-strategy"
-      - "prefetch"
-      - "--additional-config"
-      - '{"ascend_compilation_config":{"enable_npugraph_ex":true,"enable_static_kernel":false},"enable_cpu_binding":true,"multistream_overlap_shared_expert":true}'
-    extra_env:
-      HCCL_OP_EXPANSION_MODE: "AIV"
-      HCCL_BUFFSIZE: "2560"
-      PYTORCH_NPU_ALLOC_CONF: "expandable_segments:True"
-      OMP_PROC_BIND: "false"
-      OMP_NUM_THREADS: "10"
-      TASK_QUEUE_ENABLE: "1"
-      LD_PRELOAD: "/usr/lib/aarch64-linux-gnu/libjemalloc.so.2:$LD_PRELOAD"
-      # Override the npu generation-level "0,1,2,3" (for A5) — A3 uses all 16 NPUs.
-      ASCEND_RT_VISIBLE_DEVICES: "0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15"
-      # Pull W8A8 weights from ModelScope (Eco-Tech checkpoint).
-      VLLM_USE_MODELSCOPE: "True"
-    hardware_overrides:
-      ascend_910c:
-        docker_image: "quay.io/ascend/vllm-ascend:v0.23.0-a3"
-        # Ascend recommends Docker; no official pip wheel for NPU.
-        install:
-          pip: false
 
 compatible_strategies:
   - single_node_tp
@@ -534,21 +430,8 @@ strategy_hardware:
   # hardware; this makes the Strategy row agree. multi_node_dep needs no entry:
   # the Nodes row is pinned to 1 here, which already hides every multi_node
   # strategy.
-  # Atlas A5 mixed serving is DP+EP only (vLLM Ascend 0.23.0). TP-only does not
-  # partition the V4 MoE layers on 950; TEP and intra-node PD are documented
-  # separately in the vLLM Ascend 1P1D tutorial, not this builder.
-  single_node_tp:
-    npu: unsupported
-  single_node_tep:
-    npu: unsupported
-  multi_node_dep:
-    npu: unsupported
   pd_cluster:
     rtx_pro_6000_8x: unsupported
-    npu: unsupported
-    # Atlas 800 A3 PD separation (1P1D or multi-node) is documented in the
-    # vllm-ascend tutorial separately from the A5 single-node mixed path.
-    ascend_910c: supported
 
 hardware_overrides:
   blackwell:
@@ -571,16 +454,6 @@ hardware_overrides:
       - '{"mode": 3, "cudagraph_mode": "FULL_DECODE_ONLY"}'
     extra_env:
       VLLM_ROCM_USE_AITER: "1"
-  # Shared by 950DT 8x (8×96 GB) and 950PR 8x (8×128 GB). Mixed serving uses 4 of the
-  # 8 NPUs (DP4+EP). HCCL_BUFFSIZE=1024 is required for MC2 dispatch; 200 MB
-  # default fails profile_run with aclnnMoeDistributeDispatchV4 561002.
-  npu:
-    extra_env:
-      PYTORCH_NPU_ALLOC_CONF: "expandable_segments:True"
-      OMP_PROC_BIND: "false"
-      OMP_NUM_THREADS: "10"
-      HCCL_BUFFSIZE: "1024"
-      ASCEND_RT_VISIBLE_DEVICES: "0,1,2,3"
 
 strategy_overrides:
   single_node_tp:
@@ -645,38 +518,6 @@ strategy_overrides:
         extra_args:
           - "--gpu-memory-utilization"
           - "0.85"
-      npu:
-        # A per-strategy generation override replaces the recipe-level one, and
-        # DEP is the only strategy the A5 pills can select — so hardware_overrides.npu
-        # never reaches the command and its env has to be repeated here. Without
-        # HCCL_BUFFSIZE the MC2 dispatch fails profile_run with
-        # aclnnMoeDistributeDispatchV4 561002.
-        extra_env:
-          PYTORCH_NPU_ALLOC_CONF: "expandable_segments:True"
-          OMP_PROC_BIND: "false"
-          OMP_NUM_THREADS: "10"
-          HCCL_BUFFSIZE: "1024"
-          ASCEND_RT_VISIBLE_DEVICES: "0,1,2,3"
-        extra_args:
-          - "--block-size"
-          - "32"
-          - "--kv-cache-dtype"
-          - "auto"
-          - "--safetensors-load-strategy"
-          - "prefetch"
-          - "--async-scheduling"
-          - "--max-num-seqs"
-          - "64"
-          - "--max-num-batched-tokens"
-          - "4096"
-          - "--gpu-memory-utilization"
-          - "0.9"
-          - "--api-server-count"
-          - "1"
-          - "--compilation-config"
-          - '{"cudagraph_mode":"FULL_DECODE_ONLY"}'
-          - "--additional-config"
-          - '{"enable_cpu_binding": true, "multistream_overlap_shared_expert": true, "enable_shared_expert_dp": true}'
   multi_node_dep:
     extra_args:
       - "--compilation-config"
@@ -742,24 +583,6 @@ hardware_notes:
     link:
       label: "DGX Spark cluster guide"
       href: "#dgx-spark-gb10-cluster"
-  ascend_950dt_8x:
-    title: "vLLM Ascend A5 image required"
-    body: "Mixed serving on one 8-card 950DT node uses 4 of 8 NPUs (DP4+EP). FP8 (Preview) runs MTP k=1; FP8 (0731) runs DSpark k=5, verified on 950PR and offered here on the same A5 command. Pin quay.io/ascend/vllm-ascend:v0.23.0-a5 — the NVIDIA pip wheel cannot run here. Set HCCL_BUFFSIZE=1024 before launch. DeepSeek-V4 on 950 also needs the indexer quant_mode fix in vllm-project/vllm-ascend#16104 — without it the image the tag currently ships aborts during warmup with QuantLightningIndexerV2 EZ0021."
-    link:
-      label: "Ascend 950DT 8x mixed"
-      href: "#huawei-ascend-950dt-8x"
-  ascend_950pr_8x:
-    title: "vLLM Ascend A5 image required"
-    body: "Mixed serving on one 8-card 950PR node uses 4 of 8 NPUs (DP4+EP). FP8 (Preview) runs MTP k=1, same command as 950DT; FP8 (0731) runs DSpark k=5. Pin quay.io/ascend/vllm-ascend:v0.23.0-a5 and set HCCL_BUFFSIZE=1024 before launch. DeepSeek-V4 on 950 also needs the indexer quant_mode fix in vllm-project/vllm-ascend#16104 — without it the image the tag currently ships aborts during warmup with QuantLightningIndexerV2 EZ0021."
-    link:
-      label: "Ascend 950PR 8x mixed"
-      href: "#huawei-ascend-950pr-8x"
-  ascend_910c:
-    title: "vLLM Ascend A3 image required"
-    body: "Single-node DP2+TP8 on all 16 NPUs using the W8A8 quantized checkpoint from Eco-Tech. TP8 shards the 284B weights to 35.5 GB per NPU (within the 64 GB limit). Pin quay.io/ascend/vllm-ascend:v0.23.0-a3 — the NVIDIA pip wheel cannot run here. Set HCCL_BUFFSIZE=2560 before launch. The A3 path uses ModelScope to download the W8A8 weights (pip install modelscope). Multi-node PD separation is also available; see the vllm-ascend tutorial."
-    link:
-      label: "Ascend 910C A3 reference"
-      href: "https://docs.vllm.ai/projects/ascend/en/latest/tutorials/models/DeepSeek-V4-Flash.html"
 
 guide: |
   ## Overview
diff --git a/recipes/models/deepseek-ai/DeepSeek-V4-Pro.yaml b/recipes/models/deepseek-ai/DeepSeek-V4-Pro.yaml
index d57a3fc7c916..6b3f63e2fbc5 100644
--- a/recipes/models/deepseek-ai/DeepSeek-V4-Pro.yaml
+++ b/recipes/models/deepseek-ai/DeepSeek-V4-Pro.yaml
@@ -4,7 +4,7 @@ meta:
   provider: "DeepSeek"
   description: "DeepSeek V4 flagship MoE (1.6T total / 49B active) with hybrid CSA+HCA attention, manifold-constrained hyper-connections, Muon-trained on 32T+ tokens, and three-tier reasoning."
   date_added: 2026-04-24
-  date_updated: 2026-10-13
+  date_updated: 2026-10-28
   tasks:
     - text
   performance_headline: "Frontier 1.6T/49B reasoning MoE with native FP4+FP8 weights, MTP and DSpark speculative decoding, and 1M-token context"
@@ -171,6 +171,10 @@ compatible_strategies:
 # latency-oriented alternative.
 default_strategy: single_node_tep
 
+strategy_hardware:
+  single_node_dpa_tp:
+    mi355x: supported
+
 hardware_overrides:
   hopper:
     # H200 (141 GB × 8 = 1128 GB) leaves KV cache tight at 1M context with
@@ -275,6 +279,547 @@ strategy_overrides:
         - "--enable-sleep-mode"
       env: {}
 
+# Tuned presets (InferenceX AgentX lanes); see SCHEMA.md "presets".
+presets:
+  b200_agentx_tp:
+    label: "AgentX TP8"
+    hardware: b200
+    strategy: single_node_tp
+    tp: 8
+    variants: [default]
+    workload: "Agentic (InferenceX AgentX)"
+    description: "TP8 on one 8-GPU B200 node with DSpark (six draft tokens), FlashInfer sparse MLA and a fixed 50 GiB KV cache per GPU; no host KV offload."
+    source: { label: "InferenceX #4554", href: "https://github.com/SemiAnalysisAI/InferenceX/pull/4554", key: "dsv4-fp4-b200-vllm-agentic-mtp", revision: 2c044c3 }
+    image: "vllm/vllm-openai:nightly-8e2beaa3bff595577a069d49d96bde9074294ec2@sha256:3207cb99c0f59b92f7d47a74da8c6eb5467565d811a4f8d9df51df5188f59c70"
+    select:
+      features: { spec_decoding: dspark }
+    # The lane's FlashInfer sparse MLA attention config replaces the builder's
+    # standalone MXFP4 indexer flag; DSpark drafts six tokens, not seven; the pinned
+    # image predates --disable-flashinfer-autotune, so keep the --no-enable form.
+    remove_args: ["--attention_config.indexer_kv_dtype=mxfp4", "--disable-flashinfer-autotune"]
+    args:
+      - "--no-enable-flashinfer-autotune"
+      - "--max-model-len"
+      - "1048576"
+      - "--gpu-memory-utilization"
+      - "0.90"
+      - "--max-num-batched-tokens"
+      - "16384"
+      - "--numa-bind"
+      - "--enable-cumem-allocator"
+      - "--disable-uvicorn-access-log"
+      - "--no-disable-hybrid-kv-cache-manager"
+      - "--attention-config"
+      - '{"backend":"FLASHINFER_MLA_SPARSE_DSV4","use_prefill_query_quantization":true,"indexer_kv_dtype":"mxfp4"}'
+    json_args:
+      "--speculative-config": { num_speculative_tokens: 6 }
+    env:
+      VLLM_USE_V2_MODEL_RUNNER: "1"
+      VLLM_FLOAT32_MATMUL_PRECISION: "high"
+      TORCH_CUDA_ARCH_LIST: "10.0"
+      VLLM_ENGINE_READY_TIMEOUT_S: "7200"
+      VLLM_RPC_TIMEOUT: "600000"
+    # --max-num-seqs equals the concurrency. Graphs capture every decode batch of up to
+    # that many sequences (target plus six drafts each) and mixed batches of 100-500 tokens.
+    points:
+      c1:
+        label: "1"
+        concurrency: 1
+        args:
+          - "--max-num-seqs"
+          - "1"
+          - "--kv-cache-memory-bytes"
+          - "53687091200"
+          - "--compilation-config"
+          - '{"cudagraph_mode":"FULL_AND_PIECEWISE","cudagraph_capture_sizes":[7,100,200,300,400,500]}'
+        env:
+          VLLM_DSV4_MEGA_FP8_COMBINE: "1"
+      c4:
+        label: "4 (low latency)"
+        concurrency: 4
+        image: "vllm/vllm-openai@sha256:d428631bb917a33d22308265d1b8be91522f862bf9fd133bee62dde3313664ae"
+        args:
+          - "--max-num-seqs"
+          - "4"
+          - "--kv-cache-memory-bytes"
+          - "42949672960"
+          - "--compilation-config"
+          - '{"cudagraph_mode":"FULL_AND_PIECEWISE","cudagraph_capture_sizes":[7,14,21,28,100,200,300,400,500,640,800,1024,1280,1600,2048,2560,3200,4096,5120,6400,8192,10240,12800,16384],"max_cudagraph_capture_size":16384}'
+          - "--sparse-indexer-topk-backend"
+          - "deep_select"
+        env:
+          VLLM_DSV4_MEGA_FP8_COMBINE: "1"
+      c6:
+        label: "6"
+        concurrency: 6
+        args:
+          - "--max-num-seqs"
+          - "6"
+          - "--kv-cache-memory-bytes"
+          - "53687091200"
+          - "--compilation-config"
+          - '{"cudagraph_mode":"FULL_AND_PIECEWISE","cudagraph_capture_sizes":[7,14,21,28,35,42,100,200,300,400,500]}'
+        env:
+          VLLM_DSV4_MEGA_FP8_COMBINE: "1"
+      c10:
+        label: "10"
+        concurrency: 10
+        image: "vllm/vllm-openai@sha256:d428631bb917a33d22308265d1b8be91522f862bf9fd133bee62dde3313664ae"
+        args:
+          - "--max-num-seqs"
+          - "10"
+          - "--kv-cache-memory-bytes"
+          - "53687091200"
+          - "--compilation-config"
+          - '{"cudagraph_mode":"FULL_AND_PIECEWISE","cudagraph_capture_sizes":[7,14,21,28,35,42,49,56,63,70,100,200,300,400,500]}'
+          - "--prefix-cache-retention-interval"
+          - "32768"
+          - "--prefill-schedule-interval"
+          - "64"
+        json_args:
+          "--attention-config": { indexer_tp_shard_min_rows: 4096 }
+      c14:
+        label: "14"
+        concurrency: 14
+        args:
+          - "--max-num-seqs"
+          - "14"
+          - "--kv-cache-memory-bytes"
+          - "53687091200"
+          - "--compilation-config"
+          - '{"cudagraph_mode":"FULL_AND_PIECEWISE","cudagraph_capture_sizes":[7,14,21,28,35,42,49,56,63,70,77,84,91,98,100,200,300,400,500]}'
+        env:
+          VLLM_DSV4_MEGA_FP8_COMBINE: "1"
+    runbook: |
+      The TP8 points of the B200 AgentX lane (`dsv4-fp4-b200-vllm-agentic-mtp` and its low-concurrency
+      companion, [SemiAnalysisAI/InferenceX#4554](https://github.com/SemiAnalysisAI/InferenceX/pull/4554)) serve the
+      0813 checkpoint with DSpark (six draft tokens) as one aggregated vLLM worker on one 8-GPU B200 node: TP8, no
+      expert parallelism, no host KV offload. Each GPU holds a fixed 50 GiB KV cache and `--max-num-seqs` equals the
+      concurrency.
+
+      Concurrency 4 is the low-latency variant: a 40 GiB KV cache, the `deep_select` sparse indexer top-k backend and
+      graph capture up to 16384 tokens. Concurrency 10 sets `--prefix-cache-retention-interval 32768` and
+      `--prefill-schedule-interval 64`, shards the indexer at 4096 rows (`indexer_tp_shard_min_rows`) and leaves
+      `VLLM_DSV4_MEGA_FP8_COMBINE` unset. Concurrency 4 and 10 were measured on a second pinned image, shown on
+      those points.
+
+      Left out of the command: the lane also exported `VLLM_PREFIX_CACHE_RETENTION_INTERVAL`, which vLLM removed in
+      v0.30.0 ([vllm#55353](https://github.com/vllm-project/vllm/pull/55353)) in favour of the flag above, so it had
+      no effect on the image it ran on; and `--served-model-name`, which the builder replaces with the checkpoint ID.
+      The throughput points in InferenceX swap block verification for synthetic rejection at a committed
+      acceptance length so every point is measured at the same acceptance rate; the command does not set it.
+
+  b200_agentx_dep:
+    label: "AgentX DEP"
+    hardware: b200
+    strategy: single_node_dep
+    variants: [default]
+    workload: "Agentic (InferenceX AgentX)"
+    description: "Data-parallel ranks with expert parallelism and a Simple CPU KV tier behind a consistent-hash vLLM Router: DEP8 for concurrency 64 to 192, TP4 x DP2 at 32 and TP2 x DP4 at 48 and 64."
+    source: { label: "InferenceX #4554", href: "https://github.com/SemiAnalysisAI/InferenceX/pull/4554", key: "dsv4-fp4-b200-vllm-agentic-mtp", revision: 2c044c3 }
+    image: "vllm/vllm-openai:nightly-8e2beaa3bff595577a069d49d96bde9074294ec2@sha256:3207cb99c0f59b92f7d47a74da8c6eb5467565d811a4f8d9df51df5188f59c70"
+    select:
+      kv_offload: simple
+      features: { spec_decoding: dspark }
+    remove_args: ["--attention_config.indexer_kv_dtype=mxfp4"]
+    args:
+      - "--max-model-len"
+      - "1048576"
+      - "--gpu-memory-utilization"
+      - "0.90"
+      - "--numa-bind"
+      - "--enable-cumem-allocator"
+      - "--disable-uvicorn-access-log"
+      - "--enable-ep-weight-filter"
+      - "--no-disable-hybrid-kv-cache-manager"
+      - "--no-enable-flashinfer-autotune"
+      - "--prefill-schedule-interval"
+      - "16"
+      - "--long-prefill-token-threshold"
+      - "512"
+      - "--moe-backend"
+      - "deep_gemm_mega_moe"
+      - "--attention-config"
+      - '{"backend":"FLASHINFER_MLA_SPARSE_DSV4","use_prefill_query_quantization":true,"indexer_kv_dtype":"mxfp4"}'
+    json_args:
+      "--speculative-config": { num_speculative_tokens: 6 }
+      "--kv-transfer-config": { kv_connector_extra_config: { cpu_bytes_to_use_per_rank: 241500000000, enable_cross_layers_blocks: "true", lazy_offload: false } }
+    env:
+      VLLM_USE_V2_MODEL_RUNNER: "1"
+      VLLM_FLOAT32_MATMUL_PRECISION: "high"
+      TORCH_CUDA_ARCH_LIST: "10.0"
+      VLLM_ENGINE_READY_TIMEOUT_S: "7200"
+      VLLM_RPC_TIMEOUT: "600000"
+      PYTORCH_ALLOC_CONF: "expandable_segments:True"
+      PYTHONHASHSEED: "42"
+    # --max-num-seqs is per DP rank; --max-num-batched-tokens is an 8192-token prefill budget
+    # plus seven verification slots per sequence; graphs capture decode batches of 7 to 7 x seqs.
+    points:
+      dep8_c64:
+        label: "64 (DEP8)"
+        concurrency: 64
+        args:
+          - "--tensor-parallel-size"
+          - "1"
+          - "--data-parallel-size"
+          - "8"
+          - "--max-num-seqs"
+          - "16"
+          - "--max-num-batched-tokens"
+          - "8304"
+          - "--kv-cache-memory-bytes"
+          - "37580963840"
+          - "--compilation-config"
+          - '{"cudagraph_mode":"FULL_DECODE_ONLY","cudagraph_capture_sizes":[7,14,21,28,35,42,49,56,63,70,77,84,91,98,105,112],"mode":0}'
+        env:
+          VLLM_DSV4_MEGA_FP8_COMBINE: "1"
+      dep8_c96:
+        label: "96 (DEP8)"
+        concurrency: 96
+        args:
+          - "--tensor-parallel-size"
+          - "1"
+          - "--data-parallel-size"
+          - "8"
+          - "--max-num-seqs"
+          - "24"
+          - "--max-num-batched-tokens"
+          - "8360"
+          - "--kv-cache-memory-bytes"
+          - "33285996544"
+          - "--compilation-config"
+          - '{"cudagraph_mode":"FULL_DECODE_ONLY","cudagraph_capture_sizes":[7,14,21,28,35,42,49,56,63,70,77,84,91,98,105,112,119,126,133,140,147,154,161,168],"mode":0}'
+        env:
+          VLLM_DSV4_MEGA_FP8_COMBINE: "1"
+      dep8_c128:
+        label: "128 (DEP8)"
+        concurrency: 128
+        args:
+          - "--tensor-parallel-size"
+          - "1"
+          - "--data-parallel-size"
+          - "8"
+          - "--max-num-seqs"
+          - "32"
+          - "--max-num-batched-tokens"
+          - "8416"
+          - "--kv-cache-memory-bytes"
+          - "33285996544"
+          - "--compilation-config"
+          - '{"cudagraph_mode":"FULL_DECODE_ONLY","cudagraph_capture_sizes":[7,14,21,28,35,42,49,56,63,70,77,84,91,98,105,112,119,126,133,140,147,154,161,168,175,182,189,196,203,210,217,224],"mode":0}'
+        env:
+          VLLM_DSV4_MEGA_FP8_COMBINE: "1"
+      dep8_c160:
+        label: "160 (DEP8)"
+        concurrency: 160
+        args:
+          - "--tensor-parallel-size"
+          - "1"
+          - "--data-parallel-size"
+          - "8"
+          - "--max-num-seqs"
+          - "40"
+          - "--max-num-batched-tokens"
+          - "8472"
+          - "--kv-cache-memory-bytes"
+          - "32212254720"
+          - "--compilation-config"
+          - '{"cudagraph_mode":"FULL_DECODE_ONLY","cudagraph_capture_sizes":[7,14,21,28,35,42,49,56,63,70,77,84,91,98,105,112,119,126,133,140,147,154,161,168,175,182,189,196,203,210,217,224,231,238,245,252,259,266,273,280],"mode":0}'
+        env:
+          VLLM_DSV4_MEGA_FP8_COMBINE: "1"
+      dep8_c192:
+        label: "192 (DEP8)"
+        concurrency: 192
+        args:
+          - "--tensor-parallel-size"
+          - "1"
+          - "--data-parallel-size"
+          - "8"
+          - "--max-num-seqs"
+          - "48"
+          - "--max-num-batched-tokens"
+          - "8528"
+          - "--kv-cache-memory-bytes"
+          - "32212254720"
+          - "--compilation-config"
+          - '{"cudagraph_mode":"FULL_DECODE_ONLY","cudagraph_capture_sizes":[7,14,21,28,35,42,49,56,63,70,77,84,91,98,105,112,119,126,133,140,147,154,161,168,175,182,189,196,203,210,217,224,231,238,245,252,259,266,273,280,287,294,301,308,315,322,329,336],"mode":0}'
+        env:
+          VLLM_DSV4_MEGA_FP8_COMBINE: "1"
+      tp4dp2_c32:
+        label: "32 (TP4 x DP2)"
+        concurrency: 32
+        args:
+          - "--tensor-parallel-size"
+          - "4"
+          - "--data-parallel-size"
+          - "2"
+          - "--max-num-seqs"
+          - "32"
+          - "--max-num-batched-tokens"
+          - "8416"
+          - "--kv-cache-memory-bytes"
+          - "51539607552"
+          - "--compilation-config"
+          - '{"cudagraph_mode":"FULL_DECODE_ONLY","cudagraph_capture_sizes":[7,14,21,28,35,42,49,56,63,70,77,84,91,98,105,112,119,126,133,140,147,154,161,168,175,182,189,196,203,210,217,224],"mode":0}'
+      tp2dp4_c48:
+        label: "48 (TP2 x DP4)"
+        concurrency: 48
+        image: "vllm/vllm-openai@sha256:d428631bb917a33d22308265d1b8be91522f862bf9fd133bee62dde3313664ae"
+        args:
+          - "--tensor-parallel-size"
+          - "2"
+          - "--data-parallel-size"
+          - "4"
+          - "--max-num-seqs"
+          - "24"
+          - "--max-num-batched-tokens"
+          - "8360"
+          - "--kv-cache-memory-bytes"
+          - "37580963840"
+          - "--compilation-config"
+          - '{"cudagraph_mode":"FULL_DECODE_ONLY","cudagraph_capture_sizes":[7,14,21,28,35,42,49,56,63,70,77,84,91,98,105,112,119,126,133,140,147,154,161,168],"mode":0}'
+      tp2dp4_c64:
+        label: "64 (TP2 x DP4)"
+        concurrency: 64
+        image: "vllm/vllm-openai@sha256:d428631bb917a33d22308265d1b8be91522f862bf9fd133bee62dde3313664ae"
+        args:
+          - "--tensor-parallel-size"
+          - "2"
+          - "--data-parallel-size"
+          - "4"
+          - "--max-num-seqs"
+          - "32"
+          - "--max-num-batched-tokens"
+          - "8416"
+          - "--kv-cache-memory-bytes"
+          - "37580963840"
+          - "--compilation-config"
+          - '{"cudagraph_mode":"FULL_DECODE_ONLY","cudagraph_capture_sizes":[7,14,21,28,35,42,49,56,63,70,77,84,91,98,105,112,119,126,133,140,147,154,161,168,175,182,189,196,203,210,217,224],"mode":0}'
+    runbook: |
+      The data-parallel points of the B200 AgentX lane (`dsv4-fp4-b200-vllm-agentic-mtp` and its low-concurrency
+      companion, [SemiAnalysisAI/InferenceX#4554](https://github.com/SemiAnalysisAI/InferenceX/pull/4554)) serve the
+      0813 checkpoint with DSpark (six draft tokens) as one engine on one 8-GPU B200 node: DEP8 (TP1 x DP8) for
+      concurrency 64 to 192, TP4 x DP2 at 32 and TP2 x DP4 at 48 and 64, all with expert parallelism and the Simple
+      CPU KV tier. Each rank gets 241.5 GB of host memory (1932 GB across the node), so size it to the machine.
+      `--max-num-seqs` is per DP rank, and `--max-num-batched-tokens` is an 8192-token prefill budget plus seven
+      verification slots per sequence. The TP4 x DP2 and TP2 x DP4 points leave `VLLM_DSV4_MEGA_FP8_COMBINE` unset.
+      TP2 x DP4 was measured on a second pinned image, shown on those points.
+
+      The lane fronts the server with vLLM Router 0.1.14 (`--policy consistent_hash`, `--intra-node-data-parallel-size`
+      equal to the server's DP size, as on the router tab), plus `--request-timeout-secs 14400` so long agentic
+      requests outlive the default and `--disable-retries` so a failed request is not replayed. Clients send a stable
+      per-conversation `X-Session-ID` so the router keeps a conversation on the rank that holds its prefix.
+
+      Left out of the command: `VLLM_PREFIX_CACHE_RETENTION_INTERVAL` (removed in vLLM v0.30.0,
+      [vllm#55353](https://github.com/vllm-project/vllm/pull/55353), so it had no effect), `--served-model-name`
+      (the builder replaces it with the checkpoint ID) and the router's Prometheus port. The throughput points in
+      InferenceX swap block verification for synthetic rejection at a committed acceptance length so every point is
+      measured at the same acceptance rate; the command does not set it.
+
+  mi355x_agentx_tp:
+    label: "AgentX TP8 + monokernel"
+    hardware: mi355x
+    strategy: single_node_tp
+    variants: [default]
+    workload: "Agentic coding (InferenceX AgentX)"
+    description: "Source reproduction: DSpark K6 and monokernel on eight MI355X GPUs; draft communication precision requires resolution."
+    source:
+      label: "InferenceX #4526"
+      href: "https://github.com/SemiAnalysisAI/InferenceX/pull/4526"
+      key: dsv4-fp4-mi355x-vllm-agentic-mtp
+      revision: 5ec3c995988e5d3632c1ff53e8ae39da3e799892
+    image:
+      vllm: "vllm/vllm-openai-rocm:nightly-e4a5b5537344b103ed10de7764378ef09604ef3d"
+      model_revision: 72e1d3230f6c080a530b0a1d46f8eb4602340597
+      pip: [Pillow, fastapi, uvicorn]
+    tp: 8
+    select:
+      frontend: python
+      features: {tool_calling: true, reasoning: true, spec_decoding: dspark}
+      kv_offload: off
+    args:
+      - "--revision"
+      - "72e1d3230f6c080a530b0a1d46f8eb4602340597"
+      - "--tokenizer-revision"
+      - "72e1d3230f6c080a530b0a1d46f8eb4602340597"
+      - "--async-scheduling"
+      - "--data-parallel-size"
+      - "1"
+      - "--gpu-memory-utilization"
+      - "0.86"
+      - "--moe-backend"
+      - "aiter"
+      - "--monokernel"
+      - "--enable-prefix-caching"
+      - "--no-disable-hybrid-kv-cache-manager"
+    env:
+      VLLM_ENGINE_READY_TIMEOUT_S: "10800"
+      VLLM_USE_V2_MODEL_RUNNER: "1"
+      VLLM_ROCM_USE_AITER: "1"
+      VLLM_ROCM_QUICK_REDUCE_QUANTIZATION: "INT4"
+      VLLM_ROCM_USE_AITER_MOE: "1"
+      VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS: "0"
+      OMP_NUM_THREADS: "1"
+    json_args:
+      --compilation-config: {mode: 3, cudagraph_mode: FULL_AND_PIECEWISE}
+      --speculative-config: {num_speculative_tokens: 6}
+    points:
+      c1: {label: "1", concurrency: 1, args: ["--max-num-seqs", "2"]}
+      c4: {label: "4", concurrency: 4, args: ["--max-num-seqs", "8"]}
+      c8: {label: "8", concurrency: 8, args: ["--max-num-seqs", "16"]}
+      c16: {label: "16", concurrency: 16, args: ["--max-num-seqs", "32"]}
+      c32: {label: "32", concurrency: 32, args: ["--max-num-seqs", "64"]}
+      c40: {label: "40", concurrency: 40, args: ["--max-num-seqs", "80"]}
+      c48: {label: "48", concurrency: 48, args: ["--max-num-seqs", "96"]}
+      c56: {label: "56", concurrency: 56, args: ["--max-num-seqs", "112"]}
+      c64: {label: "64", concurrency: 64, args: ["--max-num-seqs", "128"]}
+    runbook: |
+      This opt-in profile serves [DeepSeek-V4-Pro-0813](https://huggingface.co/deepseek-ai/DeepSeek-V4-Pro-0813)
+      on one 8× MI355X node. The monokernel is enabled only for the TP8 arm;
+      [vLLM #61482](https://github.com/vllm-project/vllm/pull/61482) added this
+      DSpark-compatible path. These presets do not change the generic defaults.
+
+      **Qualification blocked:** the source enables nondefault INT4 QuickReduce
+      for draft collectives too. For example, 32 active K6 requests produce at
+      least 192 draft rows, making the BF16 attention-output reduction eligible
+      for INT4 communication in the pinned image. This conflicts with
+      [InferenceX's draft-precision rule](https://github.com/SemiAnalysisAI/InferenceX/blob/67de4b34ee1419b2a31c0e7200964a6e78015be4/CONTRIBUTING.md#draft-model-precision).
+      These are faithful source profiles, pending a source correction and
+      qualifying validation; they are not a production recommendation. Real
+      target verification and passing GSM8K do not resolve that precision issue.
+
+      InferenceX installs `Pillow fastapi uvicorn` before engine startup. When
+      using the pinned image interactively, install these prerequisites inside
+      the container, then run the generated Python command in that container:
+
+      ```bash
+      uv pip install Pillow fastapi uvicorn
+      ```
+
+      The generated Docker launch does not perform this setup step; provide the
+      dependencies in its container before using that launch form.
+
+      The source pins the target and tokenizer revisions, but supplies no separate
+      DSpark draft revision. DSpark loads the native draft from the same model ID
+      through the image's default mixed FP8/MXFP4 handling; no draft dtype or
+      quantization override is added. INT4 QuickReduce is shared by eligible
+      standard-path target and draft collectives, rather than being target-only.
+      The KV cache is FP8. The image is recorded as a tag, without a source-pinned
+      digest. Keep inherited environment settings consistent with the preset.
+
+      The generated serving commands use real target verification. For the
+      **InferenceX throughput sweep only**, merge the following keys into their
+      speculative config, retaining DSpark K6 and probabilistic drafting:
+
+      ```json
+      {"rejection_sample_method":"synthetic","synthetic_acceptance_length":3.77}
+      ```
+
+      `3.77` comes from the committed
+      [DeepSeek-V4-Pro thinking-on K6 golden curve](https://github.com/SemiAnalysisAI/InferenceX/blob/5ec3c995988e5d3632c1ff53e8ae39da3e799892/inferencex-e2e/infx/golden_al_distribution/dsv4-pro-0813-dspark.yaml).
+      Synthetic acceptance is a measurement control; do not use it for production
+      or accuracy evaluation. The generic Bench command is a smoke benchmark,
+      not an AgentX replay; reproduce AgentX with the linked InferenceX harness.
+
+      The TP8 scheduler caps are twice the listed client concurrency for AgentX
+      fan-out. No host-memory offload is enabled. The source's obsolete
+      `VLLM_PREFIX_CACHE_RETENTION_INTERVAL=32768` has no effect in this image;
+      the presets retain the effective CLI default of `0` (semantic checkpoints).
+
+      Evidence: [sweep 40260921924, attempt 3](https://github.com/SemiAnalysisAI/InferenceX/actions/runs/40260921924/attempts/3)
+      for source head `882db2060239f4d849f8b1c93f3b8b192ad3637e` covers all twelve TP8/DPA
+      throughput points and two real-verification GSM8K evals. The current-head
+      [run 40626241537](https://github.com/SemiAnalysisAI/InferenceX/actions/runs/40626241537)
+      only reuses that sweep; its performance/eval matrices are skipped.
+      Its actual CI merge `f301906060dcfd3b63ac6f77a735698f731f65b1` and the cited
+      head produce identical native recipe expansion for this lane.
+      This documentation change does not rerun GPU inference or establish private
+      agentic accuracy evaluation.
+  mi355x_agentx_dpa:
+    label: "AgentX DPA8 + TP experts"
+    hardware: mi355x
+    strategy: single_node_dpa_tp
+    variants: [default]
+    workload: "Agentic coding (InferenceX AgentX)"
+    description: "TP1 attention / DP8 with tensor-sharded experts, DSpark K6 and session-affine routing; monokernel disabled."
+    source:
+      label: "InferenceX #4526"
+      href: "https://github.com/SemiAnalysisAI/InferenceX/pull/4526"
+      key: dsv4-fp4-mi355x-vllm-agentic-mtp
+      revision: 5ec3c995988e5d3632c1ff53e8ae39da3e799892
+    image:
+      vllm: "vllm/vllm-openai-rocm:nightly-e4a5b5537344b103ed10de7764378ef09604ef3d"
+      model_revision: 72e1d3230f6c080a530b0a1d46f8eb4602340597
+      pip: [Pillow, fastapi, uvicorn, "vllm-router==0.1.14"]
+    select:
+      frontend: python
+      features: {tool_calling: true, reasoning: true, spec_decoding: dspark}
+      kv_offload: off
+    args:
+      - "--revision"
+      - "72e1d3230f6c080a530b0a1d46f8eb4602340597"
+      - "--tokenizer-revision"
+      - "72e1d3230f6c080a530b0a1d46f8eb4602340597"
+      - "--async-scheduling"
+      - "--gpu-memory-utilization"
+      - "0.86"
+      - "--moe-backend"
+      - "aiter"
+      - "--enable-prefix-caching"
+      - "--no-disable-hybrid-kv-cache-manager"
+      - "--prefill-schedule-interval"
+      - "8"
+      - "--long-prefill-token-threshold"
+      - "16384"
+    env:
+      VLLM_ENGINE_READY_TIMEOUT_S: "10800"
+      VLLM_USE_V2_MODEL_RUNNER: "1"
+      VLLM_ROCM_USE_AITER: "1"
+      VLLM_ROCM_QUICK_REDUCE_QUANTIZATION: "INT4"
+      VLLM_ROCM_USE_AITER_MOE: "1"
+      VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS: "0"
+      OMP_NUM_THREADS: "1"
+    json_args:
+      --compilation-config: {mode: 3, cudagraph_mode: FULL_AND_PIECEWISE}
+      --speculative-config: {num_speculative_tokens: 6}
+    points:
+      c64: {label: "64", concurrency: 64, args: ["--max-num-seqs", "64"]}
+      c128: {label: "128", concurrency: 128, args: ["--max-num-seqs", "128"]}
+      c192: {label: "192", concurrency: 192, args: ["--max-num-seqs", "192"]}
+    runbook: |
+      This is the DP-attention arm of the
+      [TP8 profile](#preset-mi355x_agentx_tp), using the same checkpoint, image,
+      runtime prerequisites and real-versus-synthetic verification distinction.
+      The source qualification blocker described there remains unresolved.
+      Runtime flags are TP1/DP8 with no expert-parallel flag: experts remain
+      tensor-sharded across the eight ranks. The source's matrix label TP8/EP1
+      describes the aggregate topology, not the attention TP argument.
+      Monokernel remains disabled (the image default); TP8-only monokernel
+      eligibility does not extend to this arm.
+
+      Each DP scheduler caps sequences at the listed client concurrency.
+      The measured prefill interval is `8`; the `16384` long-prefill threshold
+      exceeds the `8192` token budget and does not bind.
+
+      For exact source routing, replace the generic router companion with this
+      pinned setup in another shell on the engine's host/container network:
+
+      ```bash
+      uv pip install vllm-router==0.1.14
+      vllm-router \
+        --host 0.0.0.0 --port 8000 \
+        --worker-urls http://localhost:8100 \
+        --intra-node-data-parallel-size 8 \
+        --policy consistent_hash \
+        --request-timeout-secs 14400 \
+        --disable-retries
+      ```
+
+      Clients connect to the router on port `8000`, retaining an `X-Session-ID`
+      across turns of each conversation. InferenceX maps correlation IDs to
+      this header; `AIPERF_HTTP_X_SESSION_ID_FROM_CORRELATION_ID=1` is a benchmark
+      client setting, not an engine environment variable.
+
 guide: |
   ## Overview
 
@@ -726,81 +1271,25 @@ guide: |
     -d '{"model":"deepseek-ai/DeepSeek-V4-Pro-0813","messages":[{"role":"user","content":"Hello"}],"max_tokens":32}'
   ```
 
-  ## Agentic benchmark reproduction (InferenceX MI355X sweep)
-
-  The command above is the default serving configuration. The SemiAnalysis InferenceX
-  MI355X benchmark lane for this checkpoint is a separate high-concurrency benchmark
-  configuration — not the recipe default. To reproduce that sweep, start from the
-  default command above and apply these deltas:
-
-  **Source configuration:**
-
-  ```yaml
-  image: vllm/vllm-openai-rocm:nightly-eed1f3d0c6043bd494424a22443ee198dd56f657
-  ```
-
-  **Applied to both arms:**
+  ## MI355X AgentX profiles (InferenceX #4526)
 
-  - Enable `VLLM_ROCM_QUICK_REDUCE_QUANTIZATION=INT4` (previously commented out).
-  - Disable `VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=0` for this nightly.
-    DSv4-Pro is a mixed checkpoint (MXFP4 routed experts, FP8 shared expert), and
-    this image fails while loading incompatible scales and shapes when the fused
-    shared-expert path is enabled.
-  - Raise `--gpu-memory-utilization` from `0.8` to `0.86`.
-  - Pin `--max-num-batched-tokens 8192` instead of the nightly default of `16384`.
-    On the TP8 initialization check this raised GPU KV-cache capacity from `4,730,981`
-    to `8,524,228` tokens and reduced peak activation memory from `11.44 GiB` to
-    `8.9 GiB`.
+  The opt-in **AgentX TP8 + monokernel** and **AgentX DPA8 + TP experts** presets
+  cover the [InferenceX #4526](https://github.com/SemiAnalysisAI/InferenceX/pull/4526)
+  agentic-coding lane for
+  [DeepSeek-V4-Pro-0813](https://huggingface.co/deepseek-ai/DeepSeek-V4-Pro-0813)
+  on one eight-GPU MI355X node. Both use native DSpark K6, mixed FP8/MXFP4
+  weights and FP8 KV cache; only the TP8 arm enables the monokernel. They pin the
+  source nightly and target/tokenizer revisions, leaving the generic recipe and
+  the separate fixed-sequence lane unchanged.
 
-  **Applied to the DP-attention arm only:**
-
-  - Cap `--max-num-seqs` at `CONC` rather than `2*CONC`. The limit is per scheduler
-    and DP-attention runs one scheduler per rank. The pure TP8 arm keeps the existing
-    `2*CONC` headroom for AgentX subagent fan-out.
-  - Add `--prefill-schedule-interval 8` and `--long-prefill-token-threshold 16384`.
-
-  With `--max-num-batched-tokens 8192` now pinned for both arms,
-  `--long-prefill-token-threshold 16384` sits above the token budget and therefore
-  never binds. It is inert rather than harmful, and the DP-attention point was
-  measured with it present.
-
-  `--compilation-config '{"mode":3,"cudagraph_mode":"FULL_AND_PIECEWISE"}'` is
-  unchanged; it was already in the checked-in recipe.
-
-  Export the AITER and INT4 quantized all-reduce environment variables before launch:
-
-  ```bash
-  export VLLM_ROCM_USE_AITER=1
-  export VLLM_ROCM_QUICK_REDUCE_QUANTIZATION=INT4
-  export VLLM_ROCM_USE_AITER_MOE=1
-  export VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=0
-  export VLLM_ENGINE_READY_TIMEOUT_S=10800
-
-  vllm serve "$MODEL_PATH" --served-model-name "$MODEL" \
-    --host 0.0.0.0 \
-    --port "$VLLM_BACKEND_PORT" \
-    --trust-remote-code \
-    --async-scheduling \
-    --distributed-executor-backend mp \
-    --kv-cache-dtype fp8 \
-    --max-num-batched-tokens 8192 \
-    --gpu-memory-utilization 0.86 \
-    --moe-backend aiter \
-    --compilation-config '{"mode":3,"cudagraph_mode":"FULL_AND_PIECEWISE"}' \
-    --speculative-config '{"method": "mtp", "num_speculative_tokens": 3, "rejection_sample_method": "synthetic", "synthetic_acceptance_length": 2.49}' \
-    --tokenizer-mode deepseek_v4 \
-    --tool-call-parser deepseek_v4 \
-    --reasoning-parser deepseek_v4 \
-    --enable-auto-tool-choice \
-    --enable-prefix-caching \
-    --prefix-cache-retention-interval 32768 \
-    --no-disable-hybrid-kv-cache-manager \
-    --max-num-seqs "8"
-  ```
+  These source profiles await resolution of the shared INT4 draft-communication
+  precision issue described in their runbooks before production qualification.
 
-  `--prefix-cache-retention-interval` replaces the `VLLM_PREFIX_CACHE_RETENTION_INTERVAL`
-  env var, removed in vLLM 0.30.0
-  ([vllm-project/vllm#55353](https://github.com/vllm-project/vllm/pull/55353)).
+  Their generated commands retain real target verification for serving and eval.
+  The preset runbooks record runtime package prerequisites, the session-affine
+  DP router, all twelve concurrency points, and the separate throughput-only
+  synthetic acceptance overlay (`3.77`). The source sweep and real GSM8K evals
+  supply external validation; no GPU run was performed for this recipe update.
 
   ## Fixed-sequence benchmark reproduction (InferenceX MI355X 8K/1K)
 
@@ -947,8 +1436,8 @@ guide: |
   | Workload | ISL 8192 / OSL 1024, conc 4–512 | agentic-coding traces |
   | `--gpu-memory-utilization` | `0.8` | `0.86` |
   | Prefix caching | `--no-enable-prefix-caching` | `--enable-prefix-caching` |
-  | Speculation | none (STP) or real MTP, 2 tokens | synthetic MTP, 3 tokens |
-  | `--max-num-seqs` | not passed | `8` |
+  | Speculation | none (STP) or real MTP, 2 tokens | real DSpark K6 for serving/eval; synthetic AL 3.77 for throughput |
+  | `--max-num-seqs` | not passed | per point: twice concurrency for TP8, concurrency for DPA8 |
   | `--max-num-batched-tokens` | not passed | `8192` |
 
   The Agentic lane's synthetic acceptance length is a benchmark-comparability
diff --git a/recipes/models/deepseek-ai/DeepSeek-V4.1-Flash.yaml b/recipes/models/deepseek-ai/DeepSeek-V4.1-Flash.yaml
index ccf9b38f4f83..43d127bd6ad3 100644
--- a/recipes/models/deepseek-ai/DeepSeek-V4.1-Flash.yaml
+++ b/recipes/models/deepseek-ai/DeepSeek-V4.1-Flash.yaml
@@ -892,6 +892,213 @@ presets:
       plus `--request-timeout-secs 14400` so long agentic requests outlive the 1800 s default
       and `--disable-retries` so a failed request is not replayed.
 
+  mi325x_agentx:
+    label: "AgentX TP8 / TP4 / TP2"
+    hardware: mi325x
+    strategy: single_node_tp
+    variants: [fp8]
+    workload: "Agentic (InferenceX AgentX)"
+    description: "TP8 with resident Engram tables, TP4 with monokernel and host Engram, and TP2 with host Engram and smaller scheduler and graph budgets; six concurrency points per topology."
+    source: { label: "InferenceX #4617", href: "https://github.com/SemiAnalysisAI/InferenceX/pull/4617", key: "dsv41flash-fp4-mi325x-vllm-agentic-dspark", revision: 7c375027a3cd50f63153222fa6eb7069ee86d3b4 }
+    image: "vllm/vllm-openai-rocm:nightly-rocm100-edad8b03c5285a1c12e7e4f83c1af709b319a01b@sha256:009169c114af5b15e25eddc808afb3c1fc3f12c94e746472b16f39fcf39cd716"
+    select:
+      frontend: rust
+      features: { tool_calling: true, reasoning: true, spec_decoding: true, text_only: true, encoder_parallel: false }
+      dcp: 1
+      kv_offload: off
+    args:
+      - "--served-model-name"
+      - "deepseek-ai/DeepSeek-V4.1-Flash"
+      - "--max-model-len"
+      - "1048576"
+      - "--max-num-seqs"
+      - "128"
+      - "--max-num-batched-tokens"
+      - "16384"
+      - "--max-cudagraph-capture-size"
+      - "1024"
+      - "--disable-uvicorn-access-log"
+      - "--monokernel"
+    env:
+      OMP_NUM_THREADS: "1"
+      PYTHONUNBUFFERED: "1"
+    json_args:
+      "--speculative-config":
+        num_speculative_tokens: 5
+        draft_sample_method: probabilistic
+        rejection_sample_method: block
+        enable_adaptive_verification: false
+      "--compilation-config": { cudagraph_mode: FULL_DECODE_ONLY }
+    points:
+      tp8_c1:
+        label: "1 (TP8)"
+        concurrency: 1
+        tp: 8
+        remove_args: ["--monokernel"]
+        json_args:
+          "--engram-config": { cpu_offload: false }
+      tp8_c2:
+        label: "2 (TP8)"
+        concurrency: 2
+        tp: 8
+        remove_args: ["--monokernel"]
+        json_args:
+          "--engram-config": { cpu_offload: false }
+      tp8_c4:
+        label: "4 (TP8)"
+        concurrency: 4
+        tp: 8
+        remove_args: ["--monokernel"]
+        json_args:
+          "--engram-config": { cpu_offload: false }
+      tp8_c8:
+        label: "8 (TP8)"
+        concurrency: 8
+        tp: 8
+        remove_args: ["--monokernel"]
+        json_args:
+          "--engram-config": { cpu_offload: false }
+      tp8_c16:
+        label: "16 (TP8)"
+        concurrency: 16
+        tp: 8
+        remove_args: ["--monokernel"]
+        json_args:
+          "--engram-config": { cpu_offload: false }
+      tp8_c32:
+        label: "32 (TP8)"
+        concurrency: 32
+        tp: 8
+        remove_args: ["--monokernel"]
+        json_args:
+          "--engram-config": { cpu_offload: false }
+      tp4_c1:
+        label: "1 (TP4)"
+        concurrency: 1
+        tp: 4
+      tp4_c2:
+        label: "2 (TP4)"
+        concurrency: 2
+        tp: 4
+      tp4_c4:
+        label: "4 (TP4)"
+        concurrency: 4
+        default: true
+        tp: 4
+      tp4_c8:
+        label: "8 (TP4)"
+        concurrency: 8
+        tp: 4
+      tp4_c16:
+        label: "16 (TP4)"
+        concurrency: 16
+        tp: 4
+      tp4_c32:
+        label: "32 (TP4)"
+        concurrency: 32
+        tp: 4
+      tp2_c1:
+        label: "1 (TP2)"
+        concurrency: 1
+        tp: 2
+        remove_args: ["--monokernel"]
+        args:
+          - "--max-num-seqs"
+          - "16"
+          - "--max-num-batched-tokens"
+          - "8192"
+          - "--max-cudagraph-capture-size"
+          - "512"
+      tp2_c2:
+        label: "2 (TP2)"
+        concurrency: 2
+        tp: 2
+        remove_args: ["--monokernel"]
+        args:
+          - "--max-num-seqs"
+          - "16"
+          - "--max-num-batched-tokens"
+          - "8192"
+          - "--max-cudagraph-capture-size"
+          - "512"
+      tp2_c4:
+        label: "4 (TP2)"
+        concurrency: 4
+        tp: 2
+        remove_args: ["--monokernel"]
+        args:
+          - "--max-num-seqs"
+          - "16"
+          - "--max-num-batched-tokens"
+          - "8192"
+          - "--max-cudagraph-capture-size"
+          - "512"
+      tp2_c8:
+        label: "8 (TP2)"
+        concurrency: 8
+        tp: 2
+        remove_args: ["--monokernel"]
+        args:
+          - "--max-num-seqs"
+          - "16"
+          - "--max-num-batched-tokens"
+          - "8192"
+          - "--max-cudagraph-capture-size"
+          - "512"
+      tp2_c16:
+        label: "16 (TP2)"
+        concurrency: 16
+        tp: 2
+        remove_args: ["--monokernel"]
+        args:
+          - "--max-num-seqs"
+          - "32"
+          - "--max-num-batched-tokens"
+          - "8192"
+          - "--max-cudagraph-capture-size"
+          - "512"
+      tp2_c32:
+        label: "32 (TP2)"
+        concurrency: 32
+        tp: 2
+        remove_args: ["--monokernel"]
+        args:
+          - "--max-num-seqs"
+          - "64"
+          - "--max-num-batched-tokens"
+          - "8192"
+          - "--max-cudagraph-capture-size"
+          - "512"
+    runbook: |
+      These points serve the 0909 launch checkpoint with five-token DSpark and
+      real block verification. TP8 keeps Engram in HBM; TP4 and TP2 offload it
+      to host memory. TP2 uses smaller scheduler, prefill and graph budgets.
+      Decode uses full graphs, prefill runs eagerly, and decoder SWA bounded
+      replay keeps the pinned image's enabled default.
+
+      The TP4 monokernel runs eligible decode steps of at most 42 rows, then
+      falls back to the ordinary path for larger steps. On this image it also
+      quantizes DSpark draft FFN inputs from BF16 to MXFP8. The profile records
+      the measured configuration; its [InferenceX draft-precision qualification](https://github.com/SemiAnalysisAI/InferenceX/blob/main/CONTRIBUTING.md#draft-model-precision)
+      remains unresolved. TP8 and TP2 use the ordinary FFN path. Adaptive
+      verification remains disabled for the gfx942 startup issue tracked in
+      [vLLM #63816](https://github.com/vllm-project/vllm/pull/63816).
+
+      [Run 40744174186](https://github.com/SemiAnalysisAI/InferenceX/actions/runs/40744174186)
+      passed all 18 performance points and the TP8 c32 GSM8K evaluation at
+      `7c375027a3cd`. The separate [TP4 c4 evaluation](https://github.com/SemiAnalysisAI/InferenceX/actions/runs/40718348873)
+      at `3df335bfe629` confirmed active mono execution. Both GSM8K evaluations
+      scored 97.42% strict exact match over 1,319 samples. The effective serving
+      configuration is unchanged at reviewed source head `c8befe0ed1f0`.
+
+      InferenceX throughput measurement separately substitutes synthetic
+      acceptance length 3.51 from its committed DSpark thinking-on curve.
+      This control is excluded from the preset's serving and accuracy commands;
+      the generic Bench action is not an AgentX trace replay. Concurrency labels
+      identify the measured workload, while the generated sequence cap controls
+      admission and may queue additional sessions. No independent accuracy
+      evaluation was performed at every concurrency or on the 1017 checkpoint.
+
   mi355x_agentx:
     label: "AgentX TP2"
     hardware: mi355x
@@ -1067,6 +1274,55 @@ presets:
           - '{"cudagraph_mode":"FULL_DECODE_ONLY","cudagraph_capture_sizes":[6,12,18,24,30,36,42,48,54,60,66,72,78,84,90,96,102,108,114,120,126,132,138,144,150,156,162,168,174,180,186,192,198,204,210,216,222,228,234,240],"max_cudagraph_capture_size":240}'
           - "--no-enable-prefix-caching"
 
+  mi355x_agentx_tp8:
+    label: "AgentX TP8 conc 1"
+    hardware: mi355x
+    strategy: single_node_tp
+    # Same checkpoint note as mi355x_agentx: the lane served the 0909 launch
+    # checkpoint, and 1017 shares its format and memory layout.
+    variants: [fp8, default]
+    workload: "Agentic (InferenceX AgentX)"
+    description: "TP8 at concurrency 1 with mono decode and resident Engram tables: one sequence at a time and 32768-token prefill chunks, as the TP4 concurrency-1 point, on a nightly that carries mono decode at TP8."
+    source: { label: "InferenceX #4666", href: "https://github.com/SemiAnalysisAI/InferenceX/pull/4666", key: "dsv41flash-fp4-mi355x-vllm-agentic-dspark-tp8-c1", revision: fa9d48627d37b236ca6401505da2626bfe469cb2 }
+    image: "vllm/vllm-openai-rocm:nightly-335dd77b09161c5d63139b4e7ae51bb9928305b1@sha256:3c634cc2c51eda7eb840a236b1f09a5ecf6a4b0cb92c4b8486e2158770507291"
+    tp: 8
+    select:
+      frontend: rust
+      features: { spec_decoding: true, text_only: true, encoder_parallel: false }
+    args:
+      - "--max-num-seqs"
+      - "1"
+      - "--max-num-batched-tokens"
+      - "32768"
+      - "--max-cudagraph-capture-size"
+      - "8"
+    env:
+      OMP_NUM_THREADS: "1"
+    json_args:
+      "--engram-config": { cpu_offload: false }
+    points:
+      tp8_c1:
+        label: "1"
+        concurrency: 1
+    runbook: |
+      `--monokernel` runs at TP8 only on builds that carry
+      [vllm-project/vllm#64101](https://github.com/vllm-project/vllm/pull/64101)
+      (merged 2026-10-29); earlier builds refuse it outside TP2 and TP4. The pinned
+      nightly is 17 commits past that merge.
+
+      TP8 mono is a concurrency-1 configuration. #64101 measured one TP8 server at
+      concurrency 4 with 26% lower throughput than two TP4 replicas, so serve TP4 or
+      TP2 replicas, as the `mi355x_agentx` preset does, once more than one request is
+      in flight.
+
+      [InferenceX run 41098938196](https://github.com/SemiAnalysisAI/InferenceX/actions/runs/41098938196)
+      passed its concurrency-1 performance point and GSM8K evaluation at `fa9d48627`.
+      The performance point completed 362 of 373 requests with interactivity p50 696.9
+      and p90 685.5 tokens/s/user and TTFT p50 0.17 s (single run). The evaluation
+      scored 0.972 exact match (strict) over 1319 questions. Throughput runs use
+      DSpark's synthetic golden acceptance length (3.51); the preset keeps real block
+      rejection for ordinary serving and accuracy evaluation.
+
 hardware_notes:
   dgx_station_gb300:
     title: "Weights in Grace memory (UVA offload)"
diff --git a/recipes/models/inclusionAI/Ming-3.0-asr-tiny.yaml b/recipes/models/inclusionAI/Ming-3.0-asr-tiny.yaml
new file mode 100644
index 000000000000..0a43c7029aba
--- /dev/null
+++ b/recipes/models/inclusionAI/Ming-3.0-asr-tiny.yaml
@@ -0,0 +1,244 @@
+meta:
+  title: "Ming3-ASR-tiny"
+  slug: "ming-3-0-asr-tiny"
+  provider: "inclusionAI"
+  description: "Chinese, English, and dialect speech recognition with topic, hotword, and dialogue context."
+  date_added: 2026-10-26
+  date_updated: 2026-10-27
+  tasks:
+    - multimodal
+  performance_headline: "Context-aware Chinese, English, and dialect ASR; single-GPU serving"
+
+model:
+  model_id: "inclusionAI/Ming-3.0-asr-tiny"
+  min_vllm_version: "nightly"
+  nightly_required: true
+  install:
+    pip:
+      note: "Install from main for now."
+      command: |
+        git clone --branch main https://github.com/vllm-project/vllm.git
+        cd vllm
+        uv venv --python 3.12
+        source .venv/bin/activate
+        VLLM_USE_PRECOMPILED=1 uv pip install --editable ".[audio]" --torch-backend=auto
+        vllm download-kernels
+    docker: false
+  architecture: moe
+  parameter_count: "8.55B"
+  active_parameters: "1.3B"
+  context_length: 8192
+  base_args:
+    - "--trust-remote-code"
+  base_env: {}
+
+dependencies:
+  - note: "Ming uses torchaudio's sinc resampler; install a build compatible with PyTorch."
+    command: "uv pip install torchaudio --torch-backend=auto"
+
+features:
+  no_repeat_ngram:
+    description: "Load the Model Runner V2 plugin for request-level full-history n-gram blocking."
+    args:
+      - "--logits-processors"
+      - "vllm.model_executor.models.ming_v3_asr_logits_processor:NoRepeatNGramLogitsProcessor"
+
+opt_in_features: []
+
+variants:
+  default:
+    precision: bf16
+    vram_minimum_gb: 22
+    description: "BF16 checkpoint; the capacity estimate includes the audio tower and is not a measured minimum."
+
+compatible_strategies:
+  - single_node_tp
+
+hardware_overrides: {}
+
+strategy_overrides:
+  single_node_tp:
+    tp: 1
+
+guide: |
+  ## Overview
+
+  [Ming-3.0-asr-tiny](https://huggingface.co/inclusionAI/Ming-3.0-asr-tiny) supports Chinese, English, and dialect transcription with optional topic, hotword, and dialogue context. 
+  This guide covers single-GPU serving, transcription, streaming, Context ASR, and offline inference with vLLM.
+
+  ## Public example audio
+
+  Download the three clips from the official inference repository:
+
+  ```bash
+  MING_AUDIO_BASE=https://raw.githubusercontent.com/inclusionAI/Ming-3.0-asr/6141e8c552815db724d8df0f72c9d865c547a1ea/examples/cases/audio
+  curl -fL "$MING_AUDIO_BASE/english_common_voice.wav" -o english.wav
+  curl -fL "$MING_AUDIO_BASE/chinese_common_voice.wav" -o chinese.wav
+  curl -fL "$MING_AUDIO_BASE/context_same_sentence_pcm16.wav" -o context.wav
+  ```
+
+  | Clip | Observed transcription |
+  | --- | --- |
+  | English | `same way you did` |
+  | Mandarin | `漓江和湘江有秦代开凿的灵渠沟通` |
+  | Context clip with the hints below | Contains the film title `同囚` |
+
+  The Context clip illustrates correction of `铜球` to `同囚` with a relevant topic and hotwords. These are sample-specific observations, not a general Context ASR accuracy claim. Audio provenance and licenses are documented in the [official examples](https://github.com/inclusionAI/Ming-3.0-asr/blob/6141e8c552815db724d8df0f72c9d865c547a1ea/examples/cases/README.md).
+
+  For accuracy comparisons, keep batch settings and attention backends consistent; greedy transcription outputs can differ across these configurations.
+
+  ## Transcription API
+
+  Use the audio requests below to verify speech recognition. The catalog's generic text-only cURL and benchmark shortcuts do not exercise ASR.
+
+  ```bash
+  curl http://localhost:8000/v1/audio/transcriptions \
+    -F model=inclusionAI/Ming-3.0-asr-tiny \
+    -F file=@chinese.wav \
+    -F temperature=0 \
+    -F max_completion_tokens=256 \
+    -F 'vllm_xargs={"no_repeat_ngram_size":10}'
+  ```
+
+  Omit `language` for automatic recognition, or use `language=zh` / `language=en` to seed the known language. These are the supported language-hint codes; dialect transcription does not introduce additional hint codes. The response contains the transcript without the generated language label. `response_format=json` and `response_format=text` are supported.
+
+  ### OpenAI Python client
+
+  ```python
+  import json
+  from openai import OpenAI
+
+  client = OpenAI(base_url="http://localhost:8000/v1", api_key="EMPTY")
+  with open("english.wav", "rb") as audio:
+      result = client.audio.transcriptions.create(
+          model="inclusionAI/Ming-3.0-asr-tiny",
+          file=audio,
+          temperature=0,
+          extra_body={
+              "max_completion_tokens": 256,
+              "vllm_xargs": json.dumps({"no_repeat_ngram_size": 10}),
+          },
+      )
+  print(result.text)
+  ```
+
+  `vllm_xargs` is JSON-encoded for the multipart request, as in the cURL examples.
+
+  ### Context ASR
+
+  ```bash
+  curl http://localhost:8000/v1/audio/transcriptions \
+    -F model=inclusionAI/Ming-3.0-asr-tiny \
+    -F file=@context.wav \
+    -F language=zh \
+    -F temperature=0 \
+    -F max_completion_tokens=256 \
+    -F 'prompt=This is an audio about 同囚.' \
+    -F 'hotwords=同囚,黄国权,游学修,关楚耀,赵永洪,赵咏瑶,狱警,长镜头,吞玻璃戏,糖化玻璃,真实案例数据.' \
+    -F 'vllm_xargs={"no_repeat_ngram_size":10,"context_chunks":true}'
+  ```
+
+  `prompt` accepts topic or preceding dialogue appended to the official ASR instruction, or the full official instruction followed by context. Context and hotwords must be single-line text without audio placeholders. To reproduce the vendor's full prompt exactly, put its hotword text in `prompt` and omit the separate `hotwords` field.
+
+  ### Streaming
+
+  ```bash
+  curl -N http://localhost:8000/v1/audio/transcriptions \
+    -F model=inclusionAI/Ming-3.0-asr-tiny \
+    -F file=@english.wav \
+    -F language=en \
+    -F temperature=0 \
+    -F stream=true \
+    -F max_completion_tokens=256 \
+    -F 'vllm_xargs={"no_repeat_ngram_size":10}'
+  ```
+
+  Streaming sends transcript deltas after the complete audio upload. It is not live microphone input. The English clip produces `same way you did` without leaking the language prefix into the transcript.
+
+  ## Chunking and repetition control
+
+  Audio is downmixed and resampled to 16 kHz using torchaudio. The processor extracts the entire clip's log-mel features before splitting encoder chunks. Chunk settings are request-local and are included in multimodal cache keys.
+
+  | Option in `vllm_xargs` | Behavior |
+  | --- | --- |
+  | No chunk option | Use 3,000-mel-frame encoder chunks, approximately 30 seconds each. |
+  | `"context_chunks": true` | Use 3,000 frames up to 30 seconds and 6,000 frames for 30–60 seconds; reject longer clips. |
+  | `"audio_chunk_frames": 6000` | Explicit chunk size; any integer from 1 through 6,000 is accepted. |
+  | `"no_repeat_ngram_size": 10` | Block repeated 10-grams over prompt plus generated tokens. |
+
+  `context_chunks` and an explicit `audio_chunk_frames` are mutually exclusive. Ordinary chunking can handle clips beyond 60 seconds within the shared 8,192-token input/output budget and audio-upload size limit. Internal encoder chunking does not create separate transcription requests.
+
+  The no-repeat feature loads a plugin; each request must still set a nonzero `no_repeat_ngram_size` to enable blocking. Omitted values and `0` leave ordinary sampling in effect. Blocking uses the full prompt/output history, and requires Model Runner V2. Speculative decoding and MTP are unsupported with this plugin.
+
+  ## Offline inference
+
+  From the same vLLM source checkout, use the shared [audio language model example](https://github.com/vllm-project/vllm/blob/main/examples/generate/multimodal/audio_language_offline.py) with its built-in audio assets:
+
+  ```bash
+  python examples/generate/multimodal/audio_language_offline.py --model-type ming_v3_asr
+  ```
+
+  This uses the shared example's sampling defaults (`temperature=0.2`, `max_tokens=64`) and prints raw model output, including the language label. For the downloaded Context clip, use the native prompt builder with greedy decoding and 10-gram blocking:
+
+  ```python
+  from vllm import LLM, SamplingParams
+  from vllm.config.speech_to_text import SpeechToTextParams
+  from vllm.model_executor.models.ming_v3_asr import (
+      MingV3ASRForConditionalGeneration as MingASR,
+  )
+  from vllm.multimodal.audio import AudioResampler
+  from vllm.multimodal.media.audio import load_audio
+
+  def main():
+      llm = LLM(
+          model="inclusionAI/Ming-3.0-asr-tiny",
+          trust_remote_code=True,
+          dtype="bfloat16",
+          max_model_len=8192,
+          max_num_seqs=4,
+          enforce_eager=True,
+          logits_processors=[
+              "vllm.model_executor.models.ming_v3_asr_logits_processor:"
+              "NoRepeatNGramLogitsProcessor",
+          ],
+      )
+      audio, rate = load_audio("context.wav", sr=None, mono=True)
+      audio = AudioResampler(16000).resample(audio, orig_sr=rate)
+      params = SpeechToTextParams(
+          audio=audio,
+          stt_config=MingASR.get_speech_to_text_config(llm.model_config, "transcribe"),
+          model_config=llm.model_config,
+          language="zh",
+          request_prompt="This is an audio about 同囚.",
+          hotwords=(
+              "同囚,黄国权,游学修,关楚耀,赵永洪,赵咏瑶,狱警,长镜头,"
+              "吞玻璃戏,糖化玻璃,真实案例数据."
+          ),
+          extra_args={"context_chunks": True},
+      )
+      outputs = llm.generate(
+          MingASR.get_generation_prompt(params),
+          SamplingParams(
+              temperature=0,
+              max_tokens=256,
+              extra_args={"no_repeat_ngram_size": 10},
+          ),
+      )
+      print(MingASR.post_process_output(outputs[0].outputs[0].text))
+
+
+  if __name__ == "__main__":
+      main()
+  ```
+
+  ## Limitations
+
+  The [official model card](https://huggingface.co/inclusionAI/Ming-3.0-asr-tiny#usage-notes) describes transcription text output without word-level timestamps or speaker diarization.
+
+  The current vLLM integration does not support translation, segment timestamps, or live audio input. Translation and timestamp requests are rejected. The `context_chunks=true` policy accepts clips up to 60 seconds; ordinary chunking, including requests with context prompts, can accept longer clips subject to the model context and upload limits. Longer recordings can produce incomplete transcripts; accepting an upload does not guarantee complete transcription.
+
+  ## References
+
+  - [model card](https://huggingface.co/inclusionAI/Ming-3.0-asr-tiny)
+  - [Official inference repository](https://github.com/inclusionAI/Ming-3.0-asr)
+  - [vLLM speech-to-text API](https://docs.vllm.ai/en/latest/serving/online_serving/speech_to_text.html)
diff --git a/recipes/models/meta-models/Muse-Spark-Open-816B-A42B.yaml b/recipes/models/meta-models/Muse-Spark-Open-816B-A42B.yaml
index c69c65a57f59..dc6c11f7de50 100644
--- a/recipes/models/meta-models/Muse-Spark-Open-816B-A42B.yaml
+++ b/recipes/models/meta-models/Muse-Spark-Open-816B-A42B.yaml
@@ -4,7 +4,7 @@ meta:
   provider: "Muse (Meta)"
   description: "Sparse MoE reasoning model with 816B total / 42B active parameters and a 1M-token context, built for long-horizon agentic coding. Native image and video perception, channel-scoped reasoning and tool calls via the `muse` parsers, and an optional DSpark draft model for speculative decoding."
   date_added: 2026-10-15
-  date_updated: 2026-10-16
+  date_updated: 2026-10-23
   tasks:
     - text
     - multimodal
@@ -381,6 +381,19 @@ guide: |
   keep `--max-num-seqs 32`. For throughput beyond about 64 concurrent
   requests, use expert parallelism without DSpark.
 
+  ### Monokernel decode (MI355X, manual opt-in)
+
+  On MI355X, `--monokernel` runs decoder layers 1-60 of the DSpark verify
+  steps of 1 to 16 requests as persistent FlyDSL launches: one launch a layer
+  for one or two requests, three for 3-16 requests, with the verify attention
+  fused into the launches for up to four requests. Add it to the MXFP4
+  command with **Tensor + Expert Parallel** (TP8 + `--enable-expert-parallel`)
+  and the seven-token DSpark profile above (AITER FA, bf16 KV cache). vLLM
+  refuses other deployments at startup; prefill and larger steps keep the
+  regular path. The launches keep the expert activations in bf16 (MXFP4
+  weights x bf16 activations) instead of the checkpoint's W4A4, so outputs
+  are close to, but not bit-identical with, the regular path.
+
   ### Adaptive DSpark (manual opt-in, Triton attention)
 
   Adaptive verification captures variable-length decode CUDA graphs, which
diff --git a/recipes/models/mistralai/Mistral-Large-4-1T-A52B.yaml b/recipes/models/mistralai/Mistral-Large-4-1T-A52B.yaml
new file mode 100644
index 000000000000..c554a74c2792
--- /dev/null
+++ b/recipes/models/mistralai/Mistral-Large-4-1T-A52B.yaml
@@ -0,0 +1,956 @@
+# yaml-language-server: $schema=https://recipes.vllm.ai/recipe.schema.json
+#
+# Every flag below comes from a launcher that ran on GB300, B300, B200, H200 or
+# MI350X; untested options are left out.
+
+meta:
+  title: "Mistral-Large-4-1T-A52B"
+  slug: "mistral-large-4-1t-a52b"
+  provider: "Mistral AI"
+  description: "Mistral Large 4: a 1T-total / 49B-active multimodal MoE with sparse MLA attention, a built-in DSpark drafter, tool calling and reasoning, in FP8, NVFP4 and AMD MXFP4 checkpoints."
+  date_added: 2026-10-28
+  date_updated: 2026-10-30
+  tasks:
+    - multimodal
+  performance_headline: "8x GB300, NVFP4 + DSpark: TP8 c1 TPOT 3.1 ms; 1P1D (4+4) P DP4/EP → D TP4 reaches 241 output tok/s/GPU at 32k input"
+  related_recipes:
+    - mistralai/Mistral-Large-3-675B-Instruct-2512
+  default_hardware: gb300
+  # NVFP4 doesn't run on ROCm: AMD defaults to the FP8 checkpoint.
+  default_variant:
+    amd: fp8
+  hardware:
+    gb300: verified
+    b300: verified
+    b200: verified
+    h200: verified
+    mi350x: verified
+    mi355x: verified
+
+model:
+  model_id: "mistralai/Mistral-Large-4-1T-A52B"
+  min_vllm_version: "0.34.0"
+  docker_image:
+    nvidia: "vllm/vllm-openai:mistral-large-4"
+    amd: "vllm/vllm-openai-rocm:mistral-large-4"
+  install:
+    docker:
+      note: "Mistral Large 4 support ships in these images; no vLLM release or nightly wheel includes it yet. The ROCm image also carries the MXFP4 expert path and the AITER tuning tables."
+    pip: false
+  architecture: moe
+  parameter_count: "1T"
+  active_parameters: "49B"
+  context_length: 1048576
+  supports_dcp: true
+  base_args:
+    - "--tokenizer-mode"
+    - "mistral"
+    - "--config-format"
+    - "mistral"
+    - "--load-format"
+    - "mistral"
+    # Validated up to 256k; the checkpoint advertises 1,048,576 positions.
+    - "--max-model-len"
+    - "262144"
+    # Every validated run used these; current vLLM defaults on B200/GB300
+    # (16384 / 1024) are untested with this model.
+    - "--max-num-batched-tokens"
+    - "8192"
+    - "--max-num-seqs"
+    - "128"
+    - "--gpu-memory-utilization"
+    - "0.85"
+  base_env:
+    # Cold starts (weight load + DeepGEMM warmup + CUDA graph capture) took
+    # 1300-1800 s, past the 600 s default.
+    VLLM_ENGINE_READY_TIMEOUT_S: "6000"
+
+dependencies:
+  - note: "Mistral tokenizer and image processing for Mistral Large 4"
+    command: 'uv pip install -U "mistral_common[image]>=1.13.0"'
+  - note: "Only for the MegaMoE kernel (flashinfer_moe_ep_cutedsl) used by the P/D presets with expert parallelism"
+    command: 'uv pip install "nvshmem4py-cu13==0.4.1"'
+    optional: true
+    brand: NVIDIA
+
+features:
+  tool_calling:
+    description: "Mistral tool-call parser with automatic tool choice."
+    args:
+      - "--enable-auto-tool-choice"
+      - "--tool-call-parser"
+      - "mistral"
+  reasoning:
+    description: "Mistral reasoning parser: thinking goes to the `reasoning` field instead of `content`."
+    args:
+      - "--reasoning-parser"
+      - "mistral"
+  spec_decoding:
+    description: "DSpark speculative decoding with the drafter shipped in both checkpoints. Depth 3 is fixed by the checkpoint."
+    # Not validated with pipeline parallelism.
+    strategies:
+      - single_node_tp
+      - multi_node_tp
+      - multi_node_dep
+      - pd_cluster
+    default_mode: dspark
+    modes:
+      dspark:
+        label: "DSpark"
+        description: "Greedy drafting (default). Same acceptance as probabilistic at T ≤ 1 with top-p."
+        args:
+          - "--speculative-config"
+          - '{"method":"dspark","num_speculative_tokens":3}'
+      dspark_probabilistic:
+        label: "DSpark (probabilistic)"
+        description: "Probabilistic drafting for sampled serving at T ≥ 1 without top-p: AL 2.04 → 2.36–2.53 at T=1.5."
+        args:
+          - "--speculative-config"
+          - '{"method":"dspark","num_speculative_tokens":3,"draft_sample_method":"probabilistic"}'
+  text_only:
+    description: "Skip the Pixtral vision encoder for text-only serving. Required by prefill context parallelism (PCP)."
+    args:
+      - "--language-model-only"
+
+opt_in_features:
+  - text_only
+
+variants:
+  # NVFP4 is the default on Blackwell; H200 falls back to FP8.
+  default:
+    model_id: "mistralai/Mistral-Large-4-1T-A52B-NVFP4"
+    precision: nvfp4
+    # ModelOpt mixed precision: size from the 610.1 GB checkpoint, × 1.2
+    vram_minimum_gb: 733
+    supported_hardware: [gb300, b300, b200]
+    description: "ModelOpt NVFP4 experts with FP8 attention and shared experts (~76 GiB per GPU on 8 GPUs), with the DSpark drafter. ~1.5x the KV capacity of FP8. Validated on GB300, B300 and B200."
+  fp8:
+    precision: fp8
+    # ceil(1060.6 GB checkpoint × 1.2)
+    vram_minimum_gb: 1273
+    supported_hardware: [gb300, b300, h200, mi350x, mi355x]
+    description: "FP8 block-scaled weights with the DSpark drafter (~128 GiB of weights per GPU on 8 GPUs). On AMD the routed experts are re-quantized to MXFP4 at load, which also fits TP4. Validated on GB300, B300, H200 and MI350X."
+    hardware_overrides:
+      amd:
+        extra_env:
+          # FP8 → MXFP4 routed experts, layer by layer at load (peak memory:
+          # the MXFP4 model plus one layer of FP8 experts). a8w4 = FP8
+          # activations; a4w4 (MXFP4 activations) measured level with it at TP4.
+          VLLM_FP8_MOE_REQUANT: "mxfp4"
+          VLLM_FP8_MOE_REQUANT_ACT: "a8w4"
+      # The ROCm image carries ML4 kernels and AITER tables a wheel doesn't.
+      mi350x:
+        install:
+          pip: false
+        docker_serve_command: "vllm serve"
+      mi355x:
+        install:
+          pip: false
+        docker_serve_command: "vllm serve"
+  mxfp4:
+    model_id: "amd/Mistral-Large-4-1T-A52B-Quark-MXFP4"
+    precision: mxfp4
+    label: "MXFP4 (Quark)"
+    # ceil(578 GB checkpoint × 1.2)
+    vram_minimum_gb: 694
+    supported_hardware: [mi350x, mi355x]
+    hardware_overrides:
+      # The ROCm image carries ML4 kernels and AITER tables a wheel doesn't.
+      mi350x:
+        install:
+          pip: false
+        docker_serve_command: "vllm serve"
+      mi355x:
+        install:
+          pip: false
+        docker_serve_command: "vllm serve"
+    description: "AMD Quark OCP MXFP4 routed experts (static weight, dynamic activation scales); attention, shared experts, dense layers, router, head, vision encoder and the DSpark drafter keep the original precision. ~578 GB, fits four MI350 Series GPUs."
+
+# Only the validated layouts are offered:
+#   GB300: one 4-GPU tray (TP4, NVFP4), 2 trays (TP8, DP8/EP) and 1P1D with
+#          one tray per side
+#   B300:  one 8-GPU node (TP8, NVFP4 and FP8)
+#   B200:  one 8-GPU node (TP8) and 1P1D with one node per side
+#   H200:  2 x 8-GPU nodes (TP16, PP2 x TP8)
+compatible_strategies: []
+
+strategy_hardware:
+  single_node_tp:
+    gb300: supported
+    b300: supported
+    b200: supported
+    mi350x: supported
+    mi355x: supported
+  multi_node_tp:
+    gb300: supported
+    h200: supported
+  multi_node_dep:
+    gb300: supported
+  multi_node_tp_pp:
+    h200: supported
+  pd_cluster:
+    gb300: supported
+    b200: supported
+
+default_strategy: single_node_tp
+
+strategy_min_gpus:
+  h200: 16
+
+hardware_overrides:
+  amd:
+    extra_env:
+      VLLM_ROCM_USE_AITER: "1"
+      # AITER's Triton sparse MLA, and AITER custom all-reduce up to 128 MiB
+      # (large prefill steps): +2–5% output tok/s at TP4.
+      VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA: "1"
+      AITER_CUSTOM_AR_MAX_SIZE: "134217728"
+
+strategy_overrides:
+  multi_node_tp_pp:
+    env:
+      # Each stage must own layer 35 (the KV-sharing source) and the indexer
+      # layers its shared-index layers reuse. Valid for PP2 only.
+      VLLM_PP_LAYER_PARTITION: "35,37"
+  pd_cluster:
+    # Validated: P DP/EP on one node → D TP on one node
+    # (4+4 GB300, 8+8 B200).
+    prefill:
+      nodes: 1
+      parallelism: dep
+    decode:
+      nodes: 1
+      parallelism: tp
+
+presets:
+  gb300_dcp8:
+    label: "TP8 DCP8 (long context)"
+    hardware: gb300
+    strategy: multi_node_tp
+    workload: "Long context, KV-capacity-bound"
+    description: "Decode context parallelism shards the latent KV cache across the 8 ranks: ~8x KV capacity at a higher per-step latency. Works with DSpark."
+    select:
+      nodes: 2
+    args:
+      - "--decode-context-parallel-size"
+      - "8"
+    runbook: |
+      Measured on 2x GB300 (8 GPUs), NVFP4 revision 4359ec6, `--max-model-len 65536`.
+
+      | | TP8 | TP8 DCP8 |
+      |---|---:|---:|
+      | KV cache, NVFP4 (no spec) | 3.60M tokens | 28.76M tokens |
+      | KV cache, FP8 (no spec) | 2.38M tokens | 18.88M tokens |
+      | GSM8K-200 (NVFP4, no spec / DSpark) | 0.970 / 0.965 | 0.975 / 0.970 |
+      | c1 TPOT p50, 32k input, DSpark | 4.02 ms | 5.44 ms |
+
+      On NVFP4, 28.76M tokens is room for ~27 concurrent 1M-token requests
+      (KV only; serving at 1M context is not validated, the recipe caps
+      `--max-model-len` at 262144). With DSpark the draft's per-block
+      admission bound costs some of that capacity (~0.69x of plain DCP8 at
+      64k context; it doesn't bind at `--max-num-seqs 128` below ~99k
+      average context).
+
+      DCP8 only pays off when KV capacity binds: at `--max-num-seqs 128`
+      and ≤60k input its TPOT is worse than TP8 at every load we measured.
+      On the final stack, TP8 DCP8 FP8 at 60k input / 1k output measured
+      c1 TTFT 4.73 s, c1 TPOT 12.50 ms and 204 output tok/s at c64; with
+      DSpark c1 TPOT is 6.16 ms (AL 2.81).
+  gb300_pcp8_dcp8:
+    label: "TP1×PCP8×DCP8 + EP (lowest TTFT)"
+    hardware: gb300
+    strategy: multi_node_tp
+    workload: "Long prompts, low load"
+    description: "Prefill context parallelism splits each prompt across the 8 ranks: ~3x lower c1 TTFT than TP8 on long prompts, but the lowest throughput. Text-only for now."
+    select:
+      nodes: 2
+      features:
+        text_only: true
+    args:
+      - "--tensor-parallel-size"
+      - "1"
+      - "--prefill-context-parallel-size"
+      - "8"
+      - "--decode-context-parallel-size"
+      - "8"
+      - "--enable-expert-parallel"
+    runbook: |
+      Measured on 2x GB300 (8 GPUs), NVFP4 revision 4359ec6, DSpark, real-text
+      prompts, 1k output tokens, `--max-model-len 65536`, without
+      `--enable-expert-parallel`. The preset adds EP, as used by the PCP
+      prefill side of the P/D presets; it was not benchmarked in this layout.
+
+      | c1 TTFT p50 / TPOT p50 | 8k | 32k | 60k |
+      |---|---|---|---|
+      | TP1×PCP8×DCP8 + DSpark | 0.16 s / 8.75 ms | 0.54 s / 8.77 ms | 1.08 s / 9.35 ms |
+      | TP8 + DSpark | 0.45 s / 3.45 ms | 1.87 s / 4.02 ms | 3.65 s / 3.82 ms |
+
+      Under load (DSpark, p50 / p99 TTFT): 60k at ~0.1 req/s 3.30 / 5.91 s
+      (TP8: 3.70 / 7.28 s). Best throughput is 86.2 output tok/s/GPU at 8k
+      (c16, 47 tok/s/user), vs 178.3 for TP8 + DSpark. GSM8K-200 0.965–0.970
+      with DSpark, 0.980 without. KV capacity 24.48M tokens without spec.
+
+      PCP requires `--language-model-only` (text only) for now.
+  h200_dcp8:
+    label: "TP16 DCP8 (long context)"
+    hardware: h200
+    strategy: multi_node_tp
+    workload: "Long context, KV-capacity-bound"
+    description: "Decode context parallelism shards the latent KV cache across 8 ranks of the TP16 group. Validated with DSpark."
+    select:
+      nodes: 2
+    args:
+      - "--decode-context-parallel-size"
+      - "8"
+    runbook: |
+      Served correctly on 2x 8 H200 with the FP8 checkpoint, TP16 + DCP8 with
+      DSpark (3 tokens). TP16 without DCP scored 0.960 on GSM8K (500
+      questions, 5-shot, chat template, greedy). No performance numbers were
+      taken on H200.
+  gb300_pd_dep_tp:
+    label: "P DP4/EP → D TP4"
+    hardware: gb300
+    strategy: pd_cluster
+    variants: [default]
+    workload: "Throughput at 32k/60k input"
+    default: true
+    description: "Default 1P1D layout (4+4 GPUs): the best throughput at 32k and 60k input. MegaMoE kernel on the expert-parallel prefill side."
+    roles:
+      prefill:
+        nodes: 1
+        parallelism: dep
+        args:
+          - "--kernel-config"
+          - '{"moe_backend":"flashinfer_moe_ep_cutedsl"}'
+        json_args:
+          --kv-transfer-config: { NixlConnector: { kv_role: kv_both, kv_load_failure_policy: null } }
+      decode:
+        nodes: 1
+        parallelism: tp
+        remove_args: ["--compilation-config"]
+        json_args:
+          --kv-transfer-config: { NixlConnector: { kv_role: kv_both, kv_load_failure_policy: null } }
+    runbook: |
+      Measured 1P1D on 4+4 GB300 (one NVL72 rack), NVFP4 revision 4359ec6, DSpark
+      depth 3 on both sides, CUDA graphs, no prefix caching, `--max-model-len 65536`.
+      Workload: 320 unique real-text prompts per input length (8k / 32k / 60k
+      tokens), 1,024 output tokens, greedy, EOS ignored. Routed with the NIXL
+      toy proxy (`tests/v1/kv_connector/nixl_integration/toy_proxy_server.py`),
+      not vllm-router.
+
+      Best output tok/s per GPU (8 GPUs) at a minimum per-user speed:
+
+      | | 8k ≥40 | 8k ≥100 | 32k ≥40 | 32k ≥100 | 60k ≥40 / ≥100 |
+      |---|---:|---:|---:|---:|---:|
+      | P DP4/EP (MegaMoE) → D TP4 | 596 | 200 | 241 | 154 | 129 / 129 |
+      | P DP4/EP (default MoE) → D TP4 | 588 | 192 | 206 | 149 | not run |
+
+      c1 TTFT 0.43 / 1.78 / 3.48 s and c1 TPOT 3.9–4.2 ms at 8k / 32k / 60k.
+      Max 4.65 / 1.88 / 1.01 req/s. GSM8K-200 through the proxy ≥ 0.965.
+      DSpark acceptance ~2.98 / 2.85 / 2.79 at 8k / 32k / 60k.
+
+      MegaMoE (`flashinfer_moe_ep_cutedsl`) needs `nvshmem4py-cu13`; drop the
+      prefill `--kernel-config` to use the default MoE backend.
+  gb300_pd_pcp_tp:
+    label: "P PCP4×DCP4+EP → D TP4"
+    hardware: gb300
+    strategy: pd_cluster
+    variants: [default]
+    workload: "Lowest TTFT, ≥100 tok/s/user at 32k"
+    description: "Prefill with TP1×PCP4×DCP4 + EP and a 32k-token batch: 2–3.5x lower TTFT than P DP4/EP, within 7–9% of its throughput. Text only."
+    roles:
+      prefill:
+        nodes: 1
+        parallelism: tp
+        args:
+          - "--tensor-parallel-size"
+          - "1"
+          - "--prefill-context-parallel-size"
+          - "4"
+          - "--decode-context-parallel-size"
+          - "4"
+          - "--cp-kv-cache-interleave-size"
+          - "64"
+          - "--enable-expert-parallel"
+          - "--max-num-batched-tokens"
+          - "32768"
+          - "--kernel-config"
+          - '{"moe_backend":"flashinfer_moe_ep_cutedsl"}'
+          - "--language-model-only"
+        json_args:
+          --kv-transfer-config: { NixlConnector: { kv_role: kv_producer, kv_load_failure_policy: null } }
+      decode:
+        nodes: 1
+        parallelism: tp
+        remove_args: ["--compilation-config"]
+        json_args:
+          --kv-transfer-config: { NixlConnector: { kv_role: kv_both, kv_load_failure_policy: null } }
+    runbook: |
+      Same setup as the default P/D preset (4+4 GB300, NVFP4 4359ec6, DSpark on
+      both sides, toy proxy).
+
+      | | 8k ≥40 | 8k ≥100 | 32k ≥40 | 32k ≥100 | 60k ≥40 / ≥100 | c1 TTFT 8k / 32k / 60k |
+      |---|---:|---:|---:|---:|---:|---|
+      | P PCP4×DCP4+EP → D TP4 | 596 | 201 | 223 | 174 | 118 / 118 | 0.35 / 0.64 / 1.21 s |
+      | P DP4/EP → D TP4 (default) | 596 | 200 | 241 | 154 | 129 / 129 | 0.43 / 1.78 / 3.48 s |
+
+      Fixed-QPS TTFT p50 / p99: 0.76 / 1.73 s at 32k and 0.87 req/s (default
+      preset: 2.53 / 4.80 s); 1.35 / 4.47 s at 60k and 0.46 req/s (default:
+      4.77 / 8.33 s).
+
+      `--max-num-batched-tokens` was PCP's real limit: at 8192 the 32k / 60k
+      req/s drops to 0.74 / 0.40 and c1 TTFT rises to 1.43 / 2.62 s. 65536
+      adds nothing over 32768.
+  gb300_pd_dep_dep:
+    label: "P DP4/EP → D DP4/EP"
+    hardware: gb300
+    strategy: pd_cluster
+    variants: [default]
+    workload: "Short inputs, high concurrency"
+    description: "Expert-parallel decode: the highest throughput at 8k input and high concurrency, with higher per-user latency. MegaMoE on both sides; the decode side keeps a 256-token batch so the kernel uses its low-latency profile."
+    roles:
+      prefill:
+        nodes: 1
+        parallelism: dep
+        args:
+          - "--kernel-config"
+          - '{"moe_backend":"flashinfer_moe_ep_cutedsl"}'
+        json_args:
+          --kv-transfer-config: { NixlConnector: { kv_role: kv_both, kv_load_failure_policy: null } }
+      decode:
+        nodes: 1
+        parallelism: dep
+        args:
+          - "--max-num-batched-tokens"
+          - "256"
+          - "--kernel-config"
+          - '{"moe_backend":"flashinfer_moe_ep_cutedsl"}'
+        remove_args: ["--compilation-config"]
+        json_args:
+          --kv-transfer-config: { NixlConnector: { kv_role: kv_both, kv_load_failure_policy: null } }
+    runbook: |
+      Same setup as the default P/D preset (4+4 GB300, NVFP4 4359ec6, DSpark on
+      both sides, toy proxy).
+
+      | | 8k ≥40 | 8k ≥100 | 32k ≥40 | 32k ≥100 | 60k ≥40 / ≥100 | c1 TPOT |
+      |---|---:|---:|---:|---:|---:|---|
+      | P DP4/EP → D DP4/EP (D mnbt 256) | 766 | 98 | 238 | 83 | 128 / 68 | 6.0–6.5 ms |
+
+      Max 5.98 req/s at 8k. With the default MoE on both sides: 723 / 178
+      output tok/s/GPU at 8k (≥40 / ≥100 tok/s/user). MegaMoE on the decode
+      side picks its low-latency profile only when `--max-num-batched-tokens`
+      is below 512; at 8192 it lost 10% at 8k. Use session-sticky routing to
+      the DP ranks when prefix caching is on.
+  gb300_pd_agentic:
+    label: "Agentic: P PCP4×DCP4+EP → D TP4 DCP4"
+    hardware: gb300
+    strategy: pd_cluster
+    variants: [default]
+    workload: "Agentic, KV-capacity-bound (AgentX traces)"
+    description: "No spec decoding, prefix caching on: DCP4 on the decode node gives 4x its KV capacity, so long multi-turn sessions keep their cache. Text only."
+    select:
+      features:
+        spec_decoding: false
+    roles:
+      prefill:
+        nodes: 1
+        parallelism: tp
+        args:
+          - "--tensor-parallel-size"
+          - "1"
+          - "--prefill-context-parallel-size"
+          - "4"
+          - "--decode-context-parallel-size"
+          - "4"
+          - "--cp-kv-cache-interleave-size"
+          - "64"
+          - "--enable-expert-parallel"
+          - "--max-num-batched-tokens"
+          - "32768"
+          - "--kernel-config"
+          - '{"moe_backend":"flashinfer_moe_ep_cutedsl"}'
+          - "--language-model-only"
+        json_args:
+          --kv-transfer-config: { NixlConnector: { kv_role: kv_producer, kv_load_failure_policy: null } }
+      decode:
+        nodes: 1
+        parallelism: tp
+        args:
+          - "--decode-context-parallel-size"
+          - "4"
+          - "--cp-kv-cache-interleave-size"
+          - "64"
+        remove_args: ["--compilation-config"]
+        json_args:
+          --kv-transfer-config: { NixlConnector: { kv_role: kv_both, kv_load_failure_policy: null } }
+    runbook: |
+      Measured 1P1D on 4+4 GB300 with `agentx-harness==1.0.7` on
+      `semianalysisai/cc-traces-weka-062126-256k` (393 sessions, ISL mean
+      100.9k / max 255.8k, OSL p50 376), 900 s per point, NVFP4 4359ec6,
+      no spec decoding, prefix caching on, `--max-model-len 262144`.
+
+      | | KV tokens (D) | c16 out tok/s/GPU | c32 TTFT p50 | c32 out tok/s/GPU | c32 req/s |
+      |---|---:|---:|---:|---:|---:|
+      | P PCP4×DCP4 → D TP4 DCP4 | 7.81M | 57.6 | 0.72 s | 92.5 | 1.185 |
+      | P PCP4×DCP4 → D TP4 | 1.96M | 60.9 | 14.95 s | 78.3 | 0.939 |
+      | P DP4/EP → D TP4 | 1.96M | 61.0 | 12.66 s | 78.9 | 0.969 |
+      | P DP4/EP → D DP4/EP | 4 x 1.51M | 54.3 | 0.98 s | 89.7 | 1.155 |
+
+      At ~32 sessions D TP4 runs out of KV (local hit rate 17–33%) and TTFT
+      explodes; DCP4 on D fixes it. DCP4 decode costs 12–25% ITL, so at ≤16
+      sessions P PCP4×DCP4 → D TP4 is better (it is the "P PCP4×DCP4+EP → D
+      TP4" preset with spec decoding switched off). Everything is past the
+      knee at c64.
+
+      Route each session to the same DP rank on DP sides (our proxy set the
+      `X-data-parallel-rank` header): without it the hit rate drops 14–22
+      points and TTFT p99 is 2–3x worse. DSpark on a DCP decode node is not
+      supported yet.
+  b200_pd_dep_tp:
+    label: "P DP8/EP → D TP8"
+    hardware: b200
+    strategy: pd_cluster
+    variants: [default]
+    workload: "Throughput at 32k/60k input"
+    default: true
+    description: "1P1D on one 8-GPU B200 node per side: the best B200 layout we measured."
+    roles:
+      prefill:
+        nodes: 1
+        parallelism: dep
+        json_args:
+          --kv-transfer-config: { NixlConnector: { kv_role: kv_both, kv_load_failure_policy: null } }
+      decode:
+        nodes: 1
+        parallelism: tp
+        remove_args: ["--compilation-config"]
+        json_args:
+          --kv-transfer-config: { NixlConnector: { kv_role: kv_both, kv_load_failure_policy: null } }
+    runbook: |
+      Measured 1P1D on 8+8 B200 (one node per side, IB between them), NVFP4
+      revision 4359ec6, DSpark depth 3 on both sides, default MoE backend,
+      `--max-model-len 65536`, no prefix caching, NIXL toy proxy. Unique
+      real-text prompts of 32k / 60k tokens, 1,024 output tokens, greedy.
+
+      Best output tok/s per GPU over all 16 GPUs (concurrency, tok/s/user):
+
+      | P → D | 32k | 60k | c1 TPOT 32k / 60k |
+      |---|---|---|---|
+      | P DP8/EP → D TP8 | 164.4 (c256, 100) | 82.8 (c128, 150) | 3.95 / 4.11 ms |
+      | P DP8/EP → D DP8/EP | 162.9 (c256, 104) | 84.8 (c256, 112) | 6.77 / 5.87 ms |
+      | P TP8 → D TP8 | 31.7 (c32, 229) | 15.8 (c16, 243) | 4.19 / 3.75 ms |
+
+      Both DP8/EP prefill layouts are prefill-bound at 32k and 60k. GSM8K-200
+      (vLLM `tests/evals/gsm8k`, through the proxy) 0.855–0.900, within noise
+      of colocated TP8 + DSpark (0.890 / 0.885). B200 per-GPU throughput is
+      ~0.8x GB300 in the same layouts.
+
+  mi350x_tp4:
+    label: "TP4 (two replicas per node)"
+    hardware: mi350x
+    strategy: single_node_tp
+    variants: [fp8]
+    workload: "Interactive and mid concurrency"
+    select:
+      features:
+        spec_decoding: dspark
+    description: "MXFP4 experts fit four MI350X GPUs, so one 8-GPU node runs two TP4 replicas (run the second in its own container on GPUs 4-7, e.g. `docker run -e ROCR_VISIBLE_DEVICES=4,5,6,7`, on another port). Validated with DSpark and the FP8 KV cache."
+    tp: 4
+    runbook: |
+      Measured on MI350X, TP4, FP8 checkpoint with MXFP4 experts at load
+      (a4w4; the preset's default a8w4 measured level), FP8 KV cache, DSpark depth 3. 200 real prompts (80 MT-Bench
+      first turns + 120 GSM8K questions), greedy, ≤ 512 output tokens.
+
+      | DSpark vs no spec | c1 | c8 | c32 |
+      |---|---:|---:|---:|
+      | TPOT p50 speedup | 2.7x | 1.9x | 1.8x |
+      | Output tok/s | 2.1x | 1.6x | 1.4x |
+
+      AL 3.1 at every concurrency. a8w4 and a4w4 experts measured level
+      (within 3% tok/s, same AL). KV cache 1.79–1.90M tokens per replica
+      without spec, 1.02M with DSpark at `--max-model-len 32768`.
+
+      The prompt set is not the one behind the NVIDIA acceptance numbers, so
+      don't compare AL across hardware.
+  mi350x_tp4_specialized:
+    label: "TP4 specialized decode (lowest latency)"
+    hardware: mi350x
+    strategy: single_node_tp
+    variants: [fp8]
+    workload: "1–4 concurrent requests"
+    description: "Runs every MoE layer of a 1-, 2- or 4-request decode step as one persistent FlyDSL kernel (attention, indexer, router, MXFP4 experts and both all-reduces in one launch). Larger batches and prefill run the standard path. Needs TP4 and a4w4 experts; refuses to start otherwise."
+    tp: 4
+    args:
+      - "--additional-config"
+      - '{"decode_backend":"latency_specialized"}'
+      - "--max-num-seqs"
+      - "64"
+    env:
+      VLLM_FP8_MOE_REQUANT_ACT: "a4w4"
+    runbook: |
+      Measured on MI350X, TP4, a4w4 experts, FP8 KV cache, against the
+      standard decode on the same GPUs.
+
+      | TPOT speedup vs standard decode | c1 | c4 |
+      |---|---:|---:|
+      | Specialized decode (8k in / 1k out) | 1.5x | 1.1x |
+      | Specialized decode + DSpark, vs no spec (real prompts) | 2.9x | n/a |
+
+      With DSpark only single-request verify blocks run fused: at c1 it is
+      1.09x faster per token than DSpark on the standard decode, at c8 equal.
+  mi350x_mxfp4_tp4:
+    label: "TP4 (two replicas per node)"
+    hardware: mi350x
+    strategy: single_node_tp
+    variants: [mxfp4]
+    workload: "Interactive and mid concurrency"
+    select:
+      features:
+        spec_decoding: dspark
+    description: "MXFP4 experts fit four MI350X GPUs, so one 8-GPU node runs two TP4 replicas (run the second in its own container on GPUs 4-7, e.g. `docker run -e ROCR_VISIBLE_DEVICES=4,5,6,7`, on another port)."
+    tp: 4
+  mi350x_mxfp4_tp4_specialized:
+    label: "TP4 specialized decode (lowest latency)"
+    hardware: mi350x
+    strategy: single_node_tp
+    variants: [mxfp4]
+    workload: "1–4 concurrent requests"
+    description: "Runs every MoE layer of a 1-, 2- or 4-request decode step as one persistent FlyDSL kernel (attention, indexer, router, MXFP4 experts and both all-reduces in one launch). Larger batches and prefill run the standard path. Needs TP4 and a4w4 experts; refuses to start otherwise."
+    tp: 4
+    args:
+      - "--additional-config"
+      - '{"decode_backend":"latency_specialized"}'
+      - "--max-num-seqs"
+      - "64"
+  mi355x_tp4:
+    label: "TP4 (two replicas per node)"
+    hardware: mi355x
+    strategy: single_node_tp
+    variants: [fp8]
+    workload: "Interactive and mid concurrency"
+    select:
+      features:
+        spec_decoding: dspark
+    description: "MXFP4 experts fit four MI355X GPUs, so one 8-GPU node runs two TP4 replicas (run the second in its own container on GPUs 4-7, e.g. `docker run -e ROCR_VISIBLE_DEVICES=4,5,6,7`, on another port). Validated with DSpark and the FP8 KV cache."
+    tp: 4
+    runbook: |
+      Measured on MI355X, a8w4 experts, FP8 KV cache, DSpark: c1 / c4 TPOT p50
+      3.75 / 6.26 ms on real prompts.
+  mi355x_tp4_specialized:
+    label: "TP4 specialized decode (lowest latency)"
+    hardware: mi355x
+    strategy: single_node_tp
+    variants: [fp8]
+    workload: "1–4 concurrent requests"
+    description: "Runs every MoE layer of a 1-, 2- or 4-request decode step as one persistent FlyDSL kernel (attention, indexer, router, MXFP4 experts and both all-reduces in one launch). Larger batches and prefill run the standard path. Needs TP4 and a4w4 experts; refuses to start otherwise."
+    tp: 4
+    args:
+      - "--additional-config"
+      - '{"decode_backend":"latency_specialized"}'
+    env:
+      VLLM_FP8_MOE_REQUANT_ACT: "a4w4"
+    runbook: |
+      Measured on MI355X, a4w4 experts, FP8 KV cache, DSpark: c1 / c4 TPOT p50
+      3.56 / 6.14 ms on real prompts (1.05x / 1.02x vs the TP4 preset).
+  mi355x_mxfp4_tp4:
+    label: "TP4 (two replicas per node)"
+    hardware: mi355x
+    strategy: single_node_tp
+    variants: [mxfp4]
+    workload: "Interactive and mid concurrency"
+    select:
+      features:
+        spec_decoding: dspark
+    description: "MXFP4 experts fit four MI355X GPUs, so one 8-GPU node runs two TP4 replicas (run the second in its own container on GPUs 4-7, e.g. `docker run -e ROCR_VISIBLE_DEVICES=4,5,6,7`, on another port). Validated with DSpark and the FP8 KV cache."
+    tp: 4
+    runbook: |
+      Measured on MI355X, FP8 KV cache, DSpark: c1 / c4 TPOT p50 3.85 / 6.05 ms
+      on real prompts.
+  mi355x_mxfp4_tp4_specialized:
+    label: "TP4 specialized decode (lowest latency)"
+    hardware: mi355x
+    strategy: single_node_tp
+    variants: [mxfp4]
+    workload: "1–4 concurrent requests"
+    description: "Runs every MoE layer of a 1-, 2- or 4-request decode step as one persistent FlyDSL kernel (attention, indexer, router, MXFP4 experts and both all-reduces in one launch). Larger batches and prefill run the standard path. Needs TP4 and a4w4 experts; refuses to start otherwise."
+    tp: 4
+    args:
+      - "--additional-config"
+      - '{"decode_backend":"latency_specialized"}'
+    runbook: |
+      Measured on MI355X, FP8 KV cache, DSpark: c1 / c4 TPOT p50 3.60 / 6.03 ms
+      on real prompts (1.07x / 1.00x vs the TP4 preset).
+
+notes:
+  - title: "Two GB300 trays per 8-GPU instance"
+    when:
+      hardware: [gb300]
+      strategies: [multi_node_tp]
+    body: |
+      The colocated configurations were validated on 8 GPUs: two 4-GPU GB300
+      trays in one NVL72 rack. Node 1 runs `--headless`; send requests to node 0.
+  - title: "DP8/EP load balancing"
+    when:
+      hardware: [gb300]
+      strategies: [multi_node_dep]
+    body: |
+      DP8/EP over two GB300 trays was validated with vLLM's internal load
+      balancer: node 0 served the API for all 8 DP ranks, and node 1 ran with
+      `--headless --data-parallel-start-rank 4` instead of
+      `--data-parallel-hybrid-lb`. Measured with DSpark: 209.8 output tok/s/GPU
+      at 8k input (c32, 55 tok/s/user) and 50.6 at 32k (c8, 58 tok/s/user),
+      GSM8K-200 0.975, 1.97M KV tokens per DP rank.
+  - title: "Use NVFP4 for P/D on GB300"
+    when:
+      hardware: [gb300]
+      strategies: [pd_cluster]
+      variants: [fp8]
+    body: |
+      The P/D presets run one 4-GPU tray per side, which holds the NVFP4 weights
+      but not the FP8 ones (~128 GiB per GPU at 8 GPUs). Pick the NVFP4 variant.
+  - title: "Probabilistic drafting"
+    when:
+      features: [spec_decoding]
+    body: |
+      For sampled serving at T ≥ 1 without top-p, pick **DSpark (probabilistic)**:
+      at T=1.5 it raises AL from 2.04 to 2.36–2.53 (+12–22% throughput). At
+      T=1.0 with top-p 0.95 both modes give the same acceptance (AL 2.81 vs 2.80
+      at c8), and greedy is ~1% faster per step. Both are lossless.
+  - title: "HiSparse: set the TRTLLM all-reduce backend"
+    when:
+      hardware: [gb300, b300, b200]
+      strategies: [single_node_tp]
+    body: |
+      With HiSparse KV offload (`--kv-transfer-config '{"kv_connector":"HiSparseConnector",...}'`),
+      also set `VLLM_FLASHINFER_ALLREDUCE_BACKEND=trtllm`. The default single-node
+      FlashInfer all-reduce (MNNVL) crashes the engine with `CUDA error: unspecified
+      launch failure` once HiSparse spills KV to host
+      ([vllm#64327](https://github.com/vllm-project/vllm/issues/64327); fix in
+      [flashinfer#7782](https://github.com/flashinfer-ai/flashinfer/pull/7782)).
+      Not needed without HiSparse.
+  - title: "H200: two 8-GPU nodes"
+    when:
+      hardware: [h200]
+    body: |
+      FP8 only (NVFP4 needs Blackwell). The ~1.06 TB of FP8 weights need two
+      8x H200 nodes: TP16 is the primary layout, PP2 x TP8 the alternative.
+      FlashMLA sparse with the `fp8_ds_mla` KV cache runs on SM90 as is (DCP
+      top-k compaction is disabled automatically), so no extra flags are needed.
+      Node 1 runs `--headless`; send requests to node 0.
+
+      TP16 with DSpark (3 tokens), `fp8_ds_mla` and CUDA graphs scored 0.960 on
+      GSM8K (500 questions, 5-shot, chat template, greedy). No performance
+      numbers were taken on H200.
+  - title: "AMD: MXFP4 experts and AITER"
+    when:
+      hardware: [amd]
+    body: |
+      ROCm has no NVFP4 kernels. Serve the FP8 checkpoint, whose routed
+      experts are re-quantized to OCP MXFP4 at load
+      (`VLLM_FP8_MOE_REQUANT=mxfp4`; expert weights drop from ~125 to ~69 GiB
+      per GPU at TP8), or the AMD Quark MXFP4 checkpoint.
+      `VLLM_FP8_MOE_REQUANT_ACT` picks the activation format: `a8w4` (FP8, the
+      default) or `a4w4` (MXFP4, required by the specialized decode).
+      `VLLM_ROCM_USE_AITER=1` runs the MoE, GEMMs, MLA decode, RMSNorm and the
+      TP all-reduce on AITER; every AMD number here was measured with it on.
+      The latent KV cache is FP8 (E4M3) by default on ROCm, 1.86x the capacity
+      of bf16; `--kv-cache-dtype bfloat16` opts out.
+  - title: "AMD: TP8 with DSpark not yet validated"
+    when:
+      hardware: [amd]
+      features: [spec_decoding]
+    body: |
+      DSpark was measured at TP4 on MI350X (see the TP4 presets). TP8 serves
+      without spec decoding; TP8 with DSpark has not been run yet.
+  - title: "Pipeline parallelism needs a layer partition"
+    when:
+      strategies: [multi_node_tp_pp]
+    body: |
+      `VLLM_PP_LAYER_PARTITION=35,37` is required for PP2 so each stage owns
+      layer 35 (the source of the cross-layer KV sharing) and the indexer
+      layers its shared-index layers reuse. PP2 x TP8 on 2x 8 H200 scored
+      GSM8K 0.960, the same as TP16. DSpark is not offered with PP.
+
+troubleshooting:
+  - title: "NIXL_ERR_BACKEND / UCX \"Invalid value for SCOPY_\" at startup"
+    when:
+      strategies: [pd_cluster]
+    body: |
+      NCCL auto-loaded the HPC-X net plugin, which loads a second UCX into the
+      process next to NIXL's bundled UCX. Set `NCCL_NET_PLUGIN=none` on the
+      prefill and decode nodes (fine when each side is a single node), and keep
+      HPC-X out of `LD_LIBRARY_PATH`.
+  - title: "ModuleNotFoundError: nvshmem"
+    body: |
+      The MegaMoE backend (`--kernel-config '{"moe_backend":"flashinfer_moe_ep_cutedsl"}'`)
+      needs `nvshmem4py-cu13`. Install `nvshmem4py-cu13==0.4.1` (the optional
+      dependency above), or drop the `--kernel-config` to use the default MoE backend.
+  - title: "SIGBUS in FlashInfer JIT or autotune"
+    body: |
+      Seen on shared filesystems. Give each job a private
+      `FLASHINFER_WORKSPACE_BASE`, or add
+      `--kernel-config '{"enable_flashinfer_autotune": false}'`.
+  - title: "DSpark on a DCP decode node"
+    when:
+      strategies: [pd_cluster]
+      features: [spec_decoding]
+    body: |
+      Not supported yet: a P/D decode node with `--decode-context-parallel-size`
+      can't run DSpark. DCP prefill → TP decode with DSpark on both sides works,
+      and colocated DCP works with DSpark.
+  - title: "Mistral Large 4 specialized decode unsupported: …"
+    when:
+      hardware: [amd]
+    body: |
+      The specialized decode is default-deny: it starts only at TP4 on
+      gfx950, with MXFP4 a4w4 experts (`VLLM_FP8_MOE_REQUANT=mxfp4
+      VLLM_FP8_MOE_REQUANT_ACT=a4w4`), Model Runner V2, an FP8 (E4M3) or bf16
+      latent KV cache and no context, data, pipeline or expert parallelism.
+      The error lists every reason; fix them or drop `--additional-config`.
+
+guide: |
+  ## Overview
+
+  Mistral Large 4 is Mistral AI's frontier mixture-of-experts multimodal model: 1T
+  parameters in total, 49B active per token (52B including embeddings and output
+  layers). It has 72 layers, 128 routed experts (top-4) plus one shared expert, MLA
+  attention with a DSA sparse indexer (top-2048 keys), cross-layer KV sharing, a
+  Pixtral vision encoder and a built-in DSpark drafter. It supports tool calling and
+  reasoning (`reasoning_effort`).
+
+  | Variant | Checkpoint | Weights | Validated on |
+  |---|---|---|---|
+  | NVFP4 (default) | [`mistralai/Mistral-Large-4-1T-A52B-NVFP4`](https://huggingface.co/mistralai/Mistral-Large-4-1T-A52B-NVFP4) | ModelOpt NVFP4 experts, FP8 attention / shared experts, ~610 GB | GB300, B200 |
+  | FP8 (default on AMD) | [`mistralai/Mistral-Large-4-1T-A52B`](https://huggingface.co/mistralai/Mistral-Large-4-1T-A52B) | FP8 block-scaled, ~1.06 TB (on AMD: MXFP4 experts at load) | GB300, H200, MI350X |
+  | MXFP4 (Quark) | [`amd/Mistral-Large-4-1T-A52B-Quark-MXFP4`](https://huggingface.co/amd/Mistral-Large-4-1T-A52B-Quark-MXFP4) | AMD Quark OCP MXFP4 routed experts, rest unchanged, ~578 GB | MI350 Series |
+
+  Both checkpoints ship the DSpark drafter. The speculative depth is fixed at 3 by
+  the checkpoint.
+
+  vLLM serves the model on model runner V2 (the default) with breakable CUDA graphs,
+  FlashMLA sparse attention and an FP8 latent KV cache (`fp8_ds_mla`) by default, so
+  no attention or KV-cache flags are needed. The checkpoint advertises 1,048,576
+  positions; we validated up to 256k, which is the `--max-model-len` set here.
+
+  ## Prerequisites
+
+  - **vLLM**: the `vllm/vllm-openai:mistral-large-4` image. No vLLM release or
+    nightly wheel includes Mistral Large 4 support yet.
+  - **AMD**: the ROCm image `vllm/vllm-openai-rocm:mistral-large-4`.
+  - **mistral_common**: `mistral_common[image] >= 1.13.0` (included in both images).
+  - **MegaMoE** (P/D presets only): the image does not include `nvshmem4py-cu13`;
+    install it in the container, or drop the `--kernel-config` flag.
+  - **Hardware**: validated on NVIDIA GB300 (one 4-GPU tray with NVFP4, or 8 GPUs as
+    two trays with NVFP4 and FP8), 8x B300 (NVFP4 and FP8), 8x B200 (NVFP4) and
+    2x 8 H200 (FP8, TP16 or PP2 x TP8). FP8 needs ~128 GiB of weights per GPU on
+    8 GPUs, NVFP4 ~76 GiB. AMD Instinct MI350X (8 GPUs, TP8 or two TP4 replicas)
+    with the ROCm image.
+
+  ## Choosing a configuration
+
+  All numbers are NVFP4 on 8x GB300 with DSpark unless stated; the presets carry
+  the full tables. On H200 use TP16 (or TP16 DCP8 for long context), with PP2 x TP8
+  as the alternative; we took no performance numbers there.
+
+  | Goal | Configuration | Why |
+  |---|---|---|
+  | Default, best interactivity | TP8 + DSpark | c1 TPOT 3.13 ms (NVFP4) / 3.56 ms (FP8), AL ~2.9; 1.9–2.2x output throughput vs no spec |
+  | Throughput, short/medium inputs | DP8 + EP + DSpark | 209.8 output tok/s/GPU at 8k (c32, 55 tok/s/user), 50.6 at 32k |
+  | Long context, KV-bound | TP8 DCP8 (+ DSpark) | 28.76M KV tokens vs 3.60M on TP8 (~27 concurrent 1M-token requests by KV size) |
+  | Lowest TTFT on long prompts | TP1×PCP8×DCP8 + EP (+ DSpark) | c1 TTFT 0.16 / 0.54 / 1.08 s at 8k / 32k / 60k, ~3x faster than TP8; text only |
+  | Disaggregated throughput | P DP4/EP → D TP4 | 241 output tok/s/GPU at 32k (≥40 tok/s/user), 129 at 60k |
+  | Disaggregated, lowest TTFT | P PCP4×DCP4+EP → D TP4 | c1 TTFT 0.35 / 0.64 / 1.21 s at 8k / 32k / 60k |
+  | Disaggregated, short inputs | P DP4/EP → D DP4/EP | 766 output tok/s/GPU at 8k |
+  | Agentic, KV-bound, no spec | P PCP4×DCP4+EP → D TP4 DCP4 | at 32 sessions TTFT p50 0.72 s vs 14.95 s with D TP4 |
+
+  Colocated TP8 + DSpark at ≥100 tok/s/user: 110.2 output tok/s/GPU at 8k (c8),
+  30.5 at 32k (c2), 21.4 at 60k (c2). DCP8 doesn't pay off colocated until KV
+  capacity binds. Disaggregated 4+4 beats colocated 8-GPU serving by 1.8–2.2x at 8k
+  and 4–5x at 32k/60k, because colocated decode shares steps with prefill chunks.
+
+  Tuning knobs used by the presets:
+
+  - `--kernel-config '{"moe_backend":"flashinfer_moe_ep_cutedsl"}'` (MegaMoE) with
+    expert parallelism, +14–17% prefill req/s at 32k. Needs `nvshmem4py-cu13`. On EP
+    decode nodes only with `--max-num-batched-tokens` below 512 (256 here).
+  - `--cp-kv-cache-interleave-size 64` on the side that runs DCP in P/D.
+  - Prefix caching (on by default) for agentic workloads, with session-sticky routing
+    to the DP ranks of DP sides.
+
+  NVFP4 vs FP8 at TP8: 1.09–1.18x faster per token without spec and 1.14–1.22x with
+  DSpark, and ~1.5x the KV capacity.
+
+  ## AMD (ROCm)
+
+  Mistral Large 4 runs on AMD Instinct MI350 Series GPUs (gfx950) with the ROCm build
+  of vLLM (image `vllm/vllm-openai-rocm:mistral-large-4`). The routed experts run as
+  OCP MXFP4 on AITER's fused MoE, either from the
+  AMD Quark MXFP4 checkpoint or from the FP8 checkpoint re-quantized at load
+  (`VLLM_FP8_MOE_REQUANT=mxfp4`, layer by layer, so peak load memory stays low). With
+  MXFP4 experts the model fits four GPUs, so one 8-GPU node serves two TP4 replicas.
+  The sparse MLA runs on AITER's MLA decode kernels (`ROCM_AITER_MLA_SPARSE`), the
+  DSA indexer on AITER's gfx950 indexer kernels, and the latent KV cache is plain FP8
+  (E4M3) by default, 1.86x the capacity of bf16. The builder adds the AITER
+  environment on AMD.
+
+  | Goal | Configuration | Why |
+  |---|---|---|
+  | Default | TP8 | one replica per node; DSpark at TP8 not yet validated on ROCm |
+  | Interactive, more replicas | TP4 + DSpark (preset) | AL 3.1; TPOT 2.7x / 1.9x / 1.8x lower than no spec at c1 / c8 / c32 |
+  | Lowest latency, 1–4 users | TP4 specialized decode (preset) | one fused kernel per MoE layer: TPOT 1.5x lower at c1; 2.9x with DSpark |
+
+  MXFP4 vs FP8 experts at TP8: +22–71% output tok/s at 1–64 concurrent requests.
+
+  ## Accuracy
+
+  Validated against Mistral's reference implementation on 2x GB300 TP8: per-layer
+  parity and 256k long-context perplexity equal to the reference. GSM8K (1319
+  questions, 5-shot) scores 0.909 on the FP8 checkpoint and 0.918 with DSpark, which
+  sits at the base server's own noise floor. All P/D and colocated configurations
+  scored 0.965–0.980 on GSM8K-200 (zero-shot chat).
+
+  ## Client usage
+
+  ```bash
+  curl http://localhost:8000/v1/chat/completions -H "Content-Type: application/json" -d '{
+    "model": "mistralai/Mistral-Large-4-1T-A52B",
+    "messages": [{"role": "user", "content": "Explain sparse attention in two sentences."}],
+    "reasoning_effort": "high"
+  }'
+  ```
+
+  Image input (allow more images per prompt with `--limit-mm-per-prompt '{"image": N}'`):
+
+  ```python
+  from openai import OpenAI
+
+  client = OpenAI(api_key="EMPTY", base_url="http://localhost:8000/v1")
+  response = client.chat.completions.create(
+      model="mistralai/Mistral-Large-4-1T-A52B",
+      messages=[{
+          "role": "user",
+          "content": [
+              {"type": "text", "text": "What does this chart show?"},
+              {"type": "image_url", "image_url": {"url": "https://example.com/chart.png"}},
+          ],
+      }],
+  )
+  print(response.choices[0].message.reasoning, response.choices[0].message.content)
+  ```
+
+  ## Limitations
+
+  - Prefill context parallelism is text-only (`--language-model-only`).
+  - DSpark doesn't run on a DCP decode node in P/D yet.
+  - DSpark wasn't validated with pipeline parallelism.
+  - AMD: NVFP4 checkpoints, DCP / PCP, NIXL P/D and HiSparse aren't supported on
+    ROCm yet; TP8 with DSpark and MI300 Series (gfx942) are not validated.
+
+  ## References
+
+  - [Mistral-Large-4-1T-A52B (FP8)](https://huggingface.co/mistralai/Mistral-Large-4-1T-A52B)
+  - [Mistral-Large-4-1T-A52B-NVFP4](https://huggingface.co/mistralai/Mistral-Large-4-1T-A52B-NVFP4)
+  - [Mistral-Large-4-1T-A52B-Quark-MXFP4](https://huggingface.co/amd/Mistral-Large-4-1T-A52B-Quark-MXFP4)
+  - [Cross-layer MLA KV sharing (vllm#64095)](https://github.com/vllm-project/vllm/pull/64095)
+  - [DSpark sliding-window MLA drafts (vllm#63795)](https://github.com/vllm-project/vllm/pull/63795) and [DSpark under DCP (vllm#63796)](https://github.com/vllm-project/vllm/pull/63796)
+  - [NIXL connector](https://docs.vllm.ai/en/latest/features/nixl_connector_usage.html)
diff --git a/recipes/models/moonshotai/Kimi-K3.yaml b/recipes/models/moonshotai/Kimi-K3.yaml
index a4f61b63639f..fd8440e1cca3 100644
--- a/recipes/models/moonshotai/Kimi-K3.yaml
+++ b/recipes/models/moonshotai/Kimi-K3.yaml
@@ -4,7 +4,7 @@ meta:
   provider: "Moonshot AI"
   description: "Pre-release 2.8T-parameter native multimodal MoE with Kimi Delta Attention, Gated MLA, Attention Residuals, and a 1M-token context window"
   date_added: 2026-07-27
-  date_updated: 2026-10-26
+  date_updated: 2026-10-29
   tasks:
     - multimodal
     - text
@@ -20,7 +20,6 @@ meta:
     gb200: verified
     gb300: verified
     mi355x: verified
-    ascend_910c: verified
 
 model:
   model_id: "moonshotai/Kimi-K3"
@@ -94,10 +93,6 @@ features:
         args:
           - "--speculative-config"
           - '{"model":"RedHatAI/Kimi-K3-speculator.dspark", "num_speculative_tokens":8, "method": "dspark", "attention_backend":"TOKENSPEED_MLA", "draft_sample_method": "probabilistic", "rejection_sample_method": "block"}'
-      npu:
-        args:
-          - "--speculative-config"
-          - '{"method":"dspark","model":"RedHatAI/Kimi-K3-speculator.dspark","num_speculative_tokens":7,"enforce_eager":true}'
   text_only:
     description: "Skip the vision encoder for text-only workloads. Mutually exclusive with encoder_parallel."
     args:
@@ -127,49 +122,6 @@ variants:
     model_id: "RedHatAI/Kimi-K3-NVFP4"
     description: "RedHatAI compressed-tensors NVFP4 quantization; MoE layers quantized to NVFP4, attention and embeddings at full precision"
 
-  ascend_w4a8:
-    model_id: "Eco-Tech/Kimi-K3-w4a8"
-    precision: int4
-    label: "W4A8 (Ascend)"
-    vram_minimum_gb: 1711
-    supported_hardware: [ascend_910c]
-    extra_args:
-      - "--quantization"
-      - "ascend"
-      - "--enable-prefix-caching"
-      - "--enable-expert-parallel"
-      - "--max-num-seqs"
-      - "16"
-      - "--max-model-len"
-      - "131072"
-      - "--max-num-batched-tokens"
-      - "24576"
-      - "--gpu-memory-utilization"
-      - "0.9"
-      - "--compilation-config"
-      - '{"cudagraph_mode":"FULL_DECODE_ONLY"}'
-      - "--mm-processor-cache-gb"
-      - "0"
-      - "--additional-config"
-      - '{"enable_cpu_binding":true, "enable_flashcomm1":true}'
-      - "--mm-encoder-tp-mode"
-      - "data"
-      - "--limit-mm-per-prompt"
-      - '{"vision_chunk": 2}'
-      - "--allowed-local-media-path"
-      - "/"
-    extra_env:
-      PYTORCH_NPU_ALLOC_CONF: "expandable_segments:True"
-      OMP_PROC_BIND: "false"
-      OMP_NUM_THREADS: "1"
-      TASK_QUEUE_ENABLE: "1"
-      HCCL_BUFFSIZE: "800"
-      VLLM_ENGINE_READY_TIMEOUT_S: "7200"
-    hardware_overrides:
-      ascend_910c:
-        docker_image: "quay.io/ascend/vllm-ascend:kimi-k3-a3"
-    description: "ModelSlim INT4 W4A8 for Ascend NPU. Requires --quantization ascend. Validated on 4× 910C mixed P/D."
-
   fp8:
     model_id: "RedHatAI/Kimi-K3-FP8-BLOCK"
     precision: fp8
@@ -194,19 +146,9 @@ strategy_hardware:
     h200: supported
     gb200: supported
     gb300: supported
-    ascend_910c: supported
   multi_node_dep:
     h100: unsupported
     h200: unsupported
-    ascend_910c: unsupported
-  multi_node_tp:
-    ascend_910c: unsupported
-  multi_node_tep:
-    ascend_910c: unsupported
-  multi_node_tp_pp:
-    ascend_910c: unsupported
-  pd_cluster:
-    ascend_910c: unsupported
 
 default_strategy: single_node_tp
 
@@ -230,7 +172,6 @@ strategy_min_gpus:
   # https://vllm-project.github.io/2026/07/27/k3.html#performance-optimizations
   gb200:
     multi_node_tp: 16
-  ascend_910c: 64
 
 hardware_overrides:
   hopper:
@@ -705,10 +646,10 @@ presets:
     strategy: pd_cluster
     workload: "Agentic (InferenceX AgentX)"
     default: true
-    description: "TP8 + DCP8 prefill and decode workers across two GB300 trays each, served by Dynamo with the DSpark K4 draft."
-    source: { label: "InferenceX #4047", href: "https://github.com/SemiAnalysisAI/InferenceX/pull/4047", key: "kimik3-fp4-gb300-dynamo-vllm-agentic-dspark-mooncake-dcp8-disagg" }
+    description: "TP8 + DCP8 prefill and decode workers across two GB300 trays each, served by Dynamo with the DSpark K4 draft: latency-specialized decode at 16-32 clients, KV-balanced decode attention at 112-264."
+    source: { label: "InferenceX #4638", href: "https://github.com/SemiAnalysisAI/InferenceX/pull/4638", key: "kimik3-fp4-gb300-dynamo-vllm-agentic-dspark-mooncake-dcp8-disagg", revision: bc4d49f5df307680be5a1567e01b838b83b31b62 }
     image:
-      vllm: "vllm/vllm-openai:nightly-e4a5b5537344b103ed10de7764378ef09604ef3d"
+      vllm: "vllm/vllm-openai:nightly-3966924d566340e971a6da313e790d05492aa505@sha256:f255c4fad2f30ff5ee80db14d3a03afe9ae0f10d3470339864ff04745fe44cb3"
       dynamo: "a0de541ea8c861d5ac01d15953ca529bea68637b"
     # Served by Dynamo with the Rust frontend; tool-call flags off under Dynamo,
     # and the K4 draft is in args (specialized decode rejects the toggle's K8).
@@ -716,17 +657,19 @@ presets:
       pd_router: dynamo
       frontend: rust
       features: { tool_calling: false, spec_decoding: false }
+      # NIXL P/D plus the shared Mooncake store in host DRAM (embedded mode).
+      kv_offload: kv_store_distributed_mooncake
+    # The lane's connector options on the composed MultiConnector.
+    json_args:
+      "--kv-transfer-config":
+        kv_load_failure_policy: recompute
+        NixlConnector: { kv_buffer_device: cuda, kv_connector_extra_config: { enforce_handshake_compat: false, enable_cross_layers_blocks: false } }
+        MooncakeStoreConnector: { kv_load_failure_policy: recompute, kv_connector_extra_config: { enable_cross_layers_blocks: false, enable_offload: false } }
     # GB300: the InferenceX AgentX xPyD lane (SemiAnalysisAI/InferenceX#4047,
-    # refreshed in #4516), TP8 + DCP8 pools across two trays with the DSpark K4
-    # draft. The pd_cluster
-    # Blackwell block above replaces the recipe Blackwell baseline, so the
-    # shared serving flags are restated here; role blocks below add the rest.
-    # It renders the low-concurrency specialized-decode profile (#4373: c16-c32),
-    # so both roles keep FP32 KV-cache SSM state, pin the DCP KV interleave to 1
-    # and leave FlashInfer autotune off. NIXL P/D would otherwise raise the
-    # interleave to the KV block size after the specialized decode has read it,
-    # and its spec-verify causal mask would use the stale value. The uncapped
-    # c48+ points keep BF16 state, NIXL's interleave and autotune on (see the guide).
+    # refreshed in #4516 and #4638), TP8 + DCP8 pools across two trays with the
+    # DSpark K4 draft. The pd_cluster Blackwell block above replaces the recipe
+    # Blackwell baseline, so the shared serving flags are restated here. Roles
+    # add what every point shares; each point adds its decode profile.
     args:
       - "--decode-context-parallel-size"
       - "8"
@@ -736,19 +679,12 @@ presets:
       - "TOKENSPEED_MLA"
       - "--kv-cache-dtype"
       - "fp8"
-      - "--mamba-ssm-cache-dtype"
-      - "float32"
-      - "--cp-kv-cache-interleave-size"
-      - "1"
       - "--enable-prefix-caching"
       - "--prefix-match-unit"
       - "128"
       - "--gpu-memory-utilization"
       - "0.95"
-      - "--max-cudagraph-capture-size"
-      - "512"
       - "--enable-cumem-allocator"
-      - "--no-enable-flashinfer-autotune"
       - "--load-format"
       - "fastsafetensors"
       - "--stream-interval"
@@ -761,9 +697,9 @@ presets:
       VLLM_USE_DIRECT_DCP_Q_GATHER: "1"
       VLLM_USE_DIRECT_DCP_KV_GATHER: "1"
       VLLM_ALLREDUCE_USE_FLASHINFER: "1"
+      VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT: "0"
       VLLM_USE_NCCL_SYMM_MEM: "0"
       VLLM_USE_V2_MODEL_RUNNER: "1"
-      VLLM_PREFIX_CACHE_RETENTION_INTERVAL: "0"
       NCCL_CUMEM_ENABLE: "1"
       NCCL_MNNVL_ENABLE: "1"
       NCCL_NVLS_ENABLE: "1"
@@ -780,15 +716,16 @@ presets:
       # The qualified workers left UCX on its default transports; this
       # replaces the Blackwell PD block's rc,cuda_copy restriction.
       UCX_TLS: "all"
+      DYN_TCP_REQUEST_TIMEOUT: "120"
+      ETCD_LEASE_TTL: "600"
       # The pd_cluster Blackwell block replaces the recipe baseline, so the
       # startup timeouts are restated: loading and autotune take ~25 min.
       VLLM_ENGINE_READY_TIMEOUT_S: "3600"
       VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS: "1800"
-      VLLM_RPC_TIMEOUT: "600000"
     roles:
       prefill:
         nodes: 2
-        # Prefill captures CUDA graphs up to 8192 tokens (#4516).
+        # TP8 + EP8 with MegaMoE; CUDA graphs up to 8192 tokens.
         args:
           - "--moe-backend"
           - "deep_gemm_mega_moe"
@@ -809,47 +746,156 @@ presets:
         # GB300 decode is a TP8 + DCP8 pool without expert parallelism (#4047).
         parallelism: tp
         nodes: 2
-        # Undo the DEP-only decode flags below for the TP pool. Hybrid LB is
-        # invalid on the --headless follower; 16384 is vLLM's API-server default
-        # on >=160 GB GPUs, which the qualified decode ran with.
-        # FULL_AND_PIECEWISE is vLLM's V1 default; it replaces the strategy's
-        # FULL_DECODE_ONLY, which the qualified decode did not set.
-        # Specialized decode (InferenceX #4373, c16-c32), decode only: it needs no
-        # expert parallelism and FULL decode CUDA graphs. --max-num-seqs is
-        # 64 // (K + 1) = 12 for the K4 draft, and graphs are captured for every
-        # batch size up to 64. The lane also drops tokenspeed_mla_min_split_kv.
+        # Drop the DEP-only decode flags for the TP pool: hybrid LB is invalid
+        # on the --headless follower, and the lane keeps prefix caching on.
+        # 16384 is vLLM's API-server default on >=160 GB GPUs, which the
+        # qualified decode ran with.
+        remove_args: ["--data-parallel-hybrid-lb", "--no-enable-prefix-caching"]
         args:
-          - "--no-data-parallel-hybrid-lb"
           - "--max-num-batched-tokens"
           - "16384"
-          - "--additional-config"
-          - '{"decode_backend":"latency_specialized"}'
-          - "--max-num-seqs"
-          - "12"
-          - "--max-cudagraph-capture-size"
-          - "64"
-          - "--compilation-config"
-          - '{"cudagraph_mode":"FULL_AND_PIECEWISE","cudagraph_capture_sizes":[1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36,37,38,39,40,41,42,43,44,45,46,47,48,49,50,51,52,53,54,55,56,57,58,59,60,61,62,63,64]}'
           - "--attention-config"
           - '{"mla_prefill_backend":"TRTLLM_RAGGED","use_prefill_query_quantization":true}'
         env:
           VLLM_ENABLE_GEMM_RS: "0"
+    # One point per measured operating point. The builder renders one prefill
+    # and one decode worker; the label and description give the lane's worker
+    # counts (each worker is two trays), which the schema cannot express.
+    #
+    # Specialized decode (c16-c32): the latency_specialized backend on decode
+    # needs FULL decode CUDA graphs, FP32 SSM state and no expert parallelism.
+    # --max-num-seqs is 64 // (K + 1) = 12 for the K4 draft, with graphs for
+    # every batch size up to 64. Both roles pin the DCP KV interleave to 1 and
+    # keep FlashInfer autotune off: otherwise NIXL P/D raises the interleave to
+    # the KV block size after the specialized decode has read it, and its
+    # spec-verify causal mask uses the stale value.
+    #
+    # Regular decode (c112-c264): BF16 SSM state and FlashInfer autotune on both
+    # roles; decode runs target and draft attention on KIMI_K3_KV_BALANCED_MLA
+    # (prefill stays on TOKENSPEED_MLA) with 128 sequences, ReplaySSM and graphs
+    # up to 512 tokens in vLLM's default FULL_AND_PIECEWISE mode.
+    points:
+      1p3d_c16:
+        label: "16 (1P3D)"
+        concurrency: 16
+        description: "One prefill and three decode workers (32 GPUs), specialized decode."
+        args: &k3_gb300_spec_args
+          - "--mamba-ssm-cache-dtype"
+          - "float32"
+          - "--cp-kv-cache-interleave-size"
+          - "1"
+          - "--no-enable-flashinfer-autotune"
+        roles: &k3_gb300_spec_roles
+          decode:
+            args:
+              - "--additional-config"
+              - '{"decode_backend":"latency_specialized"}'
+              - "--max-num-seqs"
+              - "12"
+              - "--max-cudagraph-capture-size"
+              - "64"
+              - "--compilation-config"
+              - '{"cudagraph_mode":"FULL_AND_PIECEWISE","cudagraph_capture_sizes":[1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36,37,38,39,40,41,42,43,44,45,46,47,48,49,50,51,52,53,54,55,56,57,58,59,60,61,62,63,64]}'
+      1p2d_c24:
+        label: "24 (1P2D)"
+        concurrency: 24
+        default: true
+        description: "One prefill and two decode workers (24 GPUs), specialized decode."
+        args: *k3_gb300_spec_args
+        roles: *k3_gb300_spec_roles
+      1p2d_c32:
+        label: "32 (1P2D)"
+        concurrency: 32
+        description: "One prefill and two decode workers (24 GPUs), specialized decode."
+        args: *k3_gb300_spec_args
+        roles: *k3_gb300_spec_roles
+      2p4d_c112:
+        label: "112 (2P4D)"
+        concurrency: 112
+        description: "Two prefill and four decode workers (48 GPUs), KV-balanced decode."
+        remove_args: &k3_gb300_reg_remove ["--disable-flashinfer-autotune"]
+        args: &k3_gb300_reg_args
+          - "--mamba-ssm-cache-dtype"
+          - "bfloat16"
+          - "--enable-flashinfer-autotune"
+        roles: &k3_gb300_reg_roles
+          decode:
+            # vLLM's default CUDA graph mode, as the lane ran it.
+            remove_args: ["--compilation-config"]
+            args:
+              - "--attention-backend"
+              - "KIMI_K3_KV_BALANCED_MLA"
+              - "--max-num-seqs"
+              - "128"
+              - "--max-cudagraph-capture-size"
+              - "512"
+              - "--use-replayssm"
+            json_args:
+              "--attention-config": { tokenspeed_mla_min_split_kv: 8 }
+              "--speculative-config": { attention_backend: "KIMI_K3_KV_BALANCED_MLA" }
+      2p4d_c128:
+        label: "128 (2P4D)"
+        concurrency: 128
+        description: "Two prefill and four decode workers (48 GPUs), KV-balanced decode."
+        remove_args: *k3_gb300_reg_remove
+        args: *k3_gb300_reg_args
+        roles: *k3_gb300_reg_roles
+      2p3d_c128:
+        label: "128 (2P3D)"
+        concurrency: 128
+        description: "Two prefill and three decode workers (40 GPUs), KV-balanced decode."
+        remove_args: *k3_gb300_reg_remove
+        args: *k3_gb300_reg_args
+        roles: *k3_gb300_reg_roles
+      2p3d_c152:
+        label: "152 (2P3D)"
+        concurrency: 152
+        description: "Two prefill and three decode workers (40 GPUs), KV-balanced decode."
+        remove_args: *k3_gb300_reg_remove
+        args: *k3_gb300_reg_args
+        roles: *k3_gb300_reg_roles
+      3p4d_c240:
+        label: "240 (3P4D)"
+        concurrency: 240
+        description: "Three prefill and four decode workers (56 GPUs), KV-balanced decode."
+        remove_args: *k3_gb300_reg_remove
+        args: *k3_gb300_reg_args
+        roles: *k3_gb300_reg_roles
+      3p3d_c240:
+        label: "240 (3P3D)"
+        concurrency: 240
+        description: "Three prefill and three decode workers (48 GPUs), KV-balanced decode."
+        remove_args: *k3_gb300_reg_remove
+        args: *k3_gb300_reg_args
+        roles: *k3_gb300_reg_roles
+      3p3d_c264:
+        label: "264 (3P3D)"
+        concurrency: 264
+        description: "Three prefill and three decode workers (48 GPUs), KV-balanced decode."
+        remove_args: *k3_gb300_reg_remove
+        args: *k3_gb300_reg_args
+        roles: *k3_gb300_reg_roles
     runbook: |
-      The GB300 AgentX disaggregated lane in [SemiAnalysis InferenceX](https://inferencex.semianalysis.com/) (`kimik3-fp4-gb300-dynamo-vllm-agentic-dspark-mooncake-dcp8-disagg`, [SemiAnalysisAI/InferenceX#4047](https://github.com/SemiAnalysisAI/InferenceX/pull/4047), points refined in [#4373](https://github.com/SemiAnalysisAI/InferenceX/pull/4373) and [#4516](https://github.com/SemiAnalysisAI/InferenceX/pull/4516)) runs Dynamo + vLLM with several prefill and decode workers. Prefill and decode workers are both TP8 + DCP8 across two 4-GPU GB300 trays; prefill workers also use EP8 with MegaMoE, and decode workers use no expert parallelism. KV moves over NIXL, and all workers share an RDMA Mooncake store. Image `vllm/vllm-openai:nightly-e4a5b5537344b103ed10de7764378ef09604ef3d`, Dynamo `a0de541ea8c861d5ac01d15953ca529bea68637b`.
+      The GB300 AgentX disaggregated lane in [SemiAnalysis InferenceX](https://inferencex.semianalysis.com/) (`kimik3-fp4-gb300-dynamo-vllm-agentic-dspark-mooncake-dcp8-disagg`, [SemiAnalysisAI/InferenceX#4047](https://github.com/SemiAnalysisAI/InferenceX/pull/4047), points refined in [#4373](https://github.com/SemiAnalysisAI/InferenceX/pull/4373), [#4516](https://github.com/SemiAnalysisAI/InferenceX/pull/4516) and [#4638](https://github.com/SemiAnalysisAI/InferenceX/pull/4638)) runs Dynamo + vLLM with several prefill and decode workers. Every worker is TP8 + DCP8 across two 4-GPU GB300 trays; prefill workers also use EP8 with MegaMoE, and decode workers use no expert parallelism. KV moves over NIXL, and all workers share an RDMA Mooncake store in host DRAM. Image `vllm/vllm-openai:nightly-3966924d566340e971a6da313e790d05492aa505` (vLLM `0.33.0rc1.dev229+g3966924d5`), Dynamo built from `a0de541ea8c861d5ac01d15953ca529bea68637b`, Mooncake `0.3.13.post1` from the image.
 
-      | Clients | Prefill workers | Decode workers | GPUs | Decode profile | Decode `--max-num-seqs` |
-      |---:|---:|---:|---:|---|---:|
-      | 16 | 1 | 3 | 32 | [Specialized](#specialized-decode-low-concurrency) | 12 |
-      | 24-32 | 1 | 2 | 24 | [Specialized](#specialized-decode-low-concurrency) | 12 |
-      | 48-64 | 1 | 2 | 24 | Regular | 128 |
-      | 240 | 2 | 3 | 40 | Regular | 128 |
-      | 240 | 3 | 4 | 56 | Regular | 128 |
+      | Clients | Prefill workers | Decode workers | GPUs | Decode profile | Decode attention | Decode `--max-num-seqs` |
+      |---:|---:|---:|---:|---|---|---:|
+      | 16 | 1 | 3 | 32 | Specialized | `TOKENSPEED_MLA` | 12 |
+      | 24, 32 | 1 | 2 | 24 | Specialized | `TOKENSPEED_MLA` | 12 |
+      | 112, 128 | 2 | 4 | 48 | Regular | `KIMI_K3_KV_BALANCED_MLA` | 128 |
+      | 128, 152 | 2 | 3 | 40 | Regular | `KIMI_K3_KV_BALANCED_MLA` | 128 |
+      | 240 | 3 | 4 | 56 | Regular | `KIMI_K3_KV_BALANCED_MLA` | 128 |
+      | 240, 264 | 3 | 3 | 48 | Regular | `KIMI_K3_KV_BALANCED_MLA` | 128 |
 
-      Prefill workers always use `--max-num-seqs 32`. The regular profile is the Workers commands below; the specialized one changes both roles as described under [Specialized decode](#specialized-decode-low-concurrency).
+      Each point renders one prefill worker and one decode worker; start the table's number of workers per role with the same commands, two trays each. Run every worker command on both of its trays with the same `--master-addr` (the first tray's address), `--nnodes 2` and `--node-rank 0` or `1`. Prefill workers always use `--max-num-seqs 32` and target and draft attention on `TOKENSPEED_MLA`.
 
-      The preset renders one prefill pool and one decode pool of the specialized profile (the 1P2D c24-c32 row); its decode uses `--max-num-seqs 12`. For the regular points use the Workers commands below. It turns **Spec decoding** off, because the toggle's generic `RedHatAI/Kimi-K3-speculator.dspark` K8 draft was not qualified and specialized decode allows at most 7 draft tokens, and **Tool calling** off (see the parser note under Workers). This preset's DCP8 pools do not combine with a DEP decode pool, which renders TP1 with `--decode-context-parallel-size 8` that vLLM rejects; the DEP16 decode lane is the [AgentX 3P1D DEP16 preset](#preset-gb300_agentx_pd_dep16).
+      **Specialized decode (16-32 clients).** The decode workers use the `latency_specialized` decode backend from the [GB300 low-latency preset](#preset-gb300_low_latency). Its requirements are listed there: V2 model runner, FULL decode CUDA graphs, BF16 convolution state with FP32 recurrent state, DSpark with at most 7 draft tokens, no pipeline or expert parallelism, and DCP of 1 or TP. Both roles pin `--cp-kv-cache-interleave-size 1`: left unset, NIXL P/D raises the DCP KV interleave to the KV block size after the specialized decode has read it, so its spec-verify causal mask uses the stale value. The prefill workers do not set `--additional-config`, because the backend rejects the expert parallelism they use.
 
-      Documented here only: worker counts beyond one pool per role, the Mooncake store connector (the builder's Mooncake option renders a plain `MooncakeStoreConnector` without the qualified NIXL MultiConnector, recompute policy and load options), the Mooncake services, prefill `--kv-events-config`, and the frontend's router mode and session affinity (the builder renders round-robin).
+      **Regular decode (112-264 clients).** The decode workers run the target and the DSpark draft on `KIMI_K3_KV_BALANCED_MLA`, an FP8 MLA decode backend that plans its split-KV work on the device; it is never selected automatically, and prefill steps use the configured MLA prefill backend. These points cap decode at 128 sequences with ReplaySSM, so TTFT is not traded for ITL by a standing decode queue.
+
+      The preset turns **Spec decoding** off, because the toggle's generic `RedHatAI/Kimi-K3-speculator.dspark` K8 draft was not qualified and specialized decode allows at most 7 draft tokens, and **Tool calling** off: `python3 -m dynamo.vllm` parses the remaining arguments as vLLM engine arguments only, so the vLLM API-server flags `--enable-auto-tool-choice` and `--tool-call-parser` are rejected there. For tool calling and reasoning output, add `--dyn-tool-call-parser kimi_k3 --dyn-reasoning-parser kimi_k3 --dyn-enable-structural-tag` to both roles. This preset's DCP8 pools do not combine with a DEP decode pool, which renders TP1 with `--decode-context-parallel-size 8` that vLLM rejects; the DEP16 decode lane is the [AgentX 3P1D DEP16 preset](#preset-gb300_agentx_pd_dep16).
+
+      The preset selects the distributed (embedded) Mooncake store on the **KV Offload** row, which wraps NIXL and the `MooncakeStoreConnector` in one `MultiConnector` on both roles; the preset sets the lane's connector options on it (recompute on a failed load, offload off). Documented here only: the worker counts, the lane's Mooncake master options and store sizing (the builder's Mooncake option shows its own defaults), and the frontend's router mode and session affinity (the builder renders round-robin).
 
       ### Services
 
@@ -862,96 +908,14 @@ presets:
         --default_kv_lease_ttl=60000 --eviction_high_watermark_ratio=0.95 --eviction_ratio=0.10
       ```
 
-      The workers' Mooncake store uses P2P handshake metadata, RDMA over the node's NICs, a 160GB global segment and a 4GB local buffer, with offload disabled.
-
-      ### Environment (every worker)
+      Every worker reads the Mooncake store configuration from the JSON file that `MOONCAKE_CONFIG_PATH` points to. The lane registered 160GB per rank:
 
-      ```bash
-      export VLLM_USE_DIRECT_DCP_A2A=1 VLLM_USE_DIRECT_DCP_Q_GATHER=1 VLLM_USE_DIRECT_DCP_KV_GATHER=1
-      export VLLM_ALLREDUCE_USE_FLASHINFER=1 VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT=0 VLLM_USE_NCCL_SYMM_MEM=0
-      export NCCL_CUMEM_ENABLE=1 NCCL_MNNVL_ENABLE=1 NCCL_NVLS_ENABLE=1 NCCL_P2P_LEVEL=NVL
-      export VLLM_USE_V2_MODEL_RUNNER=1 VLLM_USE_RUST_FRONTEND=1 VLLM_SSM_CONV_STATE_LAYOUT=DS
-      export VLLM_PREFIX_CACHE_RETENTION_INTERVAL=0
-      export VLLM_MOONCAKE_LOAD_RECV_THREADS=4 MC_SLICE_SIZE=1048576 MC_MAX_MR_SIZE=4294967296 MC_ENABLE_DEST_DEVICE_AFFINITY=1
-      export UCX_MEMTYPE_CACHE=n UCX_MEMTYPE_REG_WHOLE=n UCX_RCACHE_MAX_UNRELEASED=1024 UCX_TCP_AF_PRIO=inet
-      export TOKENSPEED_MLA_FP16_PARTIALS=0
-      export DYN_REQUEST_PLANE=tcp DYN_TCP_REQUEST_TIMEOUT=120 ETCD_LEASE_TTL=600
-      export VLLM_ENGINE_READY_TIMEOUT_S=3600 VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS=1800 VLLM_RPC_TIMEOUT=600000
+      ```json
+      {"master_server_address": "<infra-host>:<master-port>", "metadata_server": "P2PHANDSHAKE",
+       "global_segment_size": "160GB", "local_buffer_size": "4GB", "protocol": "rdma",
+       "device_name": "<the tray's RDMA devices>", "mode": "embedded", "enable_offload": false}
       ```
 
-      Set `VLLM_ENABLE_GEMM_RS=1` on prefill workers and `VLLM_ENABLE_GEMM_RS=0` on decode workers. Prefill workers also set `DYN_BACKEND_ADMISSION_ENGINE_REQUEST_LIMIT=1024` and `DYN_BACKEND_ADMISSION_ENGINE_REQUEST_WAITING_LIMIT=1024`. Loading the 2.8T checkpoint across trays and autotuning takes a while, which is why the engine-ready, execute-model and RPC timeouts are raised.
-
-      ### Workers
-
-      Each worker spans two trays. Run the command on both trays with the same `--master-addr` (the first tray's address), `--nnodes 2` and `--node-rank 0` or `1`.
-
-      Prefill:
-
-      ```bash
-      python3 -m dynamo.vllm --model moonshotai/Kimi-K3 --served-model-name moonshotai/Kimi-K3 \
-        --disaggregation-mode prefill --request-plane tcp \
-        --nnodes 2 --node-rank <0|1> --master-addr <first-tray-ip> --device-ids 0,1,2,3 \
-        --tensor-parallel-size 8 --decode-context-parallel-size 8 --dcp-comm-backend a2a \
-        --enable-expert-parallel --moe-backend deep_gemm_mega_moe \
-        --attention-backend TOKENSPEED_MLA \
-        --attention-config '{"mla_prefill_backend":"TRTLLM_RAGGED","use_prefill_query_quantization":true,"trtllm_ragged_backend":"cute-dsl"}' \
-        --speculative-config '{"model":"Inferact/Kimi-K3-DSpark","attention_backend":"TOKENSPEED_MLA","method":"dspark","num_speculative_tokens":4,"draft_sample_method":"probabilistic","rejection_sample_method":"block"}' \
-        --kv-cache-dtype fp8 --mamba-ssm-cache-dtype bfloat16 \
-        --enable-prefix-caching --prefix-match-unit 128 \
-        --gpu-memory-utilization 0.95 --max-num-seqs 32 --max-cudagraph-capture-size 8192 \
-        --compilation-config '{"cudagraph_mode":"FULL_AND_PIECEWISE","cudagraph_capture_sizes":[4,8,12,16,20,24,32,40,48,64,80,96,128,160,256,384,512,640,768,896,1024,1152,1280,1408,1536,1664,1792,1920,2048,2304,2560,2816,3072,3328,3584,3840,4096,4608,5120,5632,6144,6656,7168,7680,8192]}' \
-        --enable-cumem-allocator --enable-flashinfer-autotune \
-        --load-format fastsafetensors --stream-interval 10 --language-model-only --trust-remote-code \
-        --kv-events-config '{"publisher":"zmq","topic":"kv-events","enable_kv_cache_events":true,"endpoint":"tcp://*:5200"}' \
-        --kv-transfer-config '{"kv_connector":"MultiConnector","kv_role":"kv_both","kv_load_failure_policy":"recompute","kv_connector_extra_config":{"connectors":[{"kv_connector":"NixlConnector","kv_role":"kv_both","kv_load_failure_policy":"fail","kv_buffer_device":"cuda","kv_connector_extra_config":{"enforce_handshake_compat":false,"enable_cross_layers_blocks":false}},{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_both","kv_load_failure_policy":"recompute","kv_connector_extra_config":{"load_async":true,"lookup_async":true,"enable_cross_layers_blocks":false,"enable_offload":false}}]}}'
-      ```
-
-      Decode (regular profile, `--max-num-seqs 128`):
-
-      ```bash
-      python3 -m dynamo.vllm --model moonshotai/Kimi-K3 --served-model-name moonshotai/Kimi-K3 \
-        --disaggregation-mode decode --request-plane tcp \
-        --nnodes 2 --node-rank <0|1> --master-addr <first-tray-ip> --device-ids 0,1,2,3 \
-        --tensor-parallel-size 8 --decode-context-parallel-size 8 --dcp-comm-backend a2a \
-        --attention-backend TOKENSPEED_MLA \
-        --attention-config '{"mla_prefill_backend":"TRTLLM_RAGGED","use_prefill_query_quantization":true,"tokenspeed_mla_min_split_kv":8}' \
-        --speculative-config '{"model":"Inferact/Kimi-K3-DSpark","attention_backend":"TOKENSPEED_MLA","method":"dspark","num_speculative_tokens":4,"draft_sample_method":"probabilistic","rejection_sample_method":"block"}' \
-        --kv-cache-dtype fp8 --mamba-ssm-cache-dtype bfloat16 \
-        --enable-prefix-caching --prefix-match-unit 128 \
-        --gpu-memory-utilization 0.95 --max-num-seqs 128 --max-cudagraph-capture-size 512 --use-replayssm \
-        --enable-cumem-allocator --enable-flashinfer-autotune \
-        --load-format fastsafetensors --stream-interval 10 --language-model-only --trust-remote-code \
-        --kv-transfer-config '{"kv_connector":"MultiConnector","kv_role":"kv_both","kv_load_failure_policy":"recompute","kv_connector_extra_config":{"connectors":[{"kv_connector":"NixlConnector","kv_role":"kv_both","kv_load_failure_policy":"fail","kv_buffer_device":"cuda","kv_connector_extra_config":{"enforce_handshake_compat":false,"enable_cross_layers_blocks":false}},{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_both","kv_load_failure_policy":"recompute","kv_connector_extra_config":{"load_async":true,"lookup_async":true,"enable_cross_layers_blocks":false,"enable_offload":false}}]}}'
-      ```
-
-      For tool calling and reasoning output, add `--reasoning-parser kimi_k3 --dyn-tool-call-parser kimi_k3 --dyn-reasoning-parser kimi_k3 --dyn-enable-structural-tag` to both roles. `python3 -m dynamo.vllm` parses the remaining arguments as vLLM engine arguments only, so the vLLM API-server flags `--enable-auto-tool-choice` and `--tool-call-parser` are rejected there.
-
-      ### Specialized decode (low concurrency)
-
-      The c16-c32 points use the `latency_specialized` decode backend from [GB300 low-latency preset](#preset-gb300_low_latency) on the decode workers (1P3D for c16, 1P2D for c24-c32), with the Workers commands above changed as follows. The backend's requirements are listed there: V2 model runner (set in Environment), FULL decode CUDA graphs, BF16 convolution state with FP32 recurrent state, DSpark with at most 7 draft tokens, no pipeline or expert parallelism, and DCP of 1 or TP. The decode workers already meet the last two.
-
-      Both roles:
-
-      ```bash
-      --mamba-ssm-cache-dtype float32 --cp-kv-cache-interleave-size 1 --no-enable-flashinfer-autotune
-      # replacing --mamba-ssm-cache-dtype bfloat16 and --enable-flashinfer-autotune
-      ```
-
-      Pin `--cp-kv-cache-interleave-size 1` on both roles. Left unset, NIXL P/D raises the DCP KV interleave to the KV block size after the specialized decode has read it, so its spec-verify causal mask uses the stale value.
-
-      Decode only:
-
-      ```bash
-      --additional-config '{"decode_backend":"latency_specialized"}' \
-      --attention-config '{"mla_prefill_backend":"TRTLLM_RAGGED","use_prefill_query_quantization":true}' \
-      --max-num-seqs 12 --max-cudagraph-capture-size 64 \
-      --compilation-config "{\"cudagraph_mode\":\"FULL_AND_PIECEWISE\",\"cudagraph_capture_sizes\":[$(seq -s, 1 64)]}"
-      # replacing the regular decode's --attention-config (it also sets tokenspeed_mla_min_split_kv),
-      # --max-num-seqs 128 and --max-cudagraph-capture-size 512, and without --use-replayssm
-      ```
-
-      `--max-num-seqs` is `64 // (K + 1)` with the K4 draft, so a full verify batch (12 x 5 = 60 tokens) fits the 64-token graph, and CUDA graphs are captured for every batch size up to 64. The prefill workers do not set `--additional-config`: the backend rejects the expert parallelism they use. The uncapped c48+ points keep the regular profile, so TTFT is not traded for ITL by a standing decode queue.
-
       ### Frontend
 
       ```bash
@@ -967,10 +931,10 @@ presets:
     hardware: gb300
     strategy: pd_cluster
     workload: "Agentic (InferenceX AgentX)"
-    description: "Three TP8 + EP8 prefill workers across two GB300 trays each and one DEP16 decode worker across four trays, served by Dynamo with the DSpark K4 draft and online MXFP8 target projections."
-    source: { label: "InferenceX #4516", href: "https://github.com/SemiAnalysisAI/InferenceX/pull/4516", key: "kimik3-fp4-gb300-dynamo-vllm-agentic-dspark-mooncake-dep16-disagg" }
+    description: "Three TP8 + EP8 prefill workers across two GB300 trays each and one DEP16 decode worker across four trays on KV-balanced decode attention, served by Dynamo with the DSpark K4 draft and online MXFP8 target projections."
+    source: { label: "InferenceX #4638", href: "https://github.com/SemiAnalysisAI/InferenceX/pull/4638", key: "kimik3-fp4-gb300-dynamo-vllm-agentic-dspark-mooncake-dep16-disagg", revision: bc4d49f5df307680be5a1567e01b838b83b31b62 }
     image:
-      vllm: "vllm/vllm-openai:nightly-e4a5b5537344b103ed10de7764378ef09604ef3d"
+      vllm: "vllm/vllm-openai:nightly-3966924d566340e971a6da313e790d05492aa505@sha256:f255c4fad2f30ff5ee80db14d3a03afe9ae0f10d3470339864ff04745fe44cb3"
       dynamo: "a0de541ea8c861d5ac01d15953ca529bea68637b"
     # Served by Dynamo with the Rust frontend; tool-call flags off under Dynamo,
     # and the K4 draft is in args instead of the toggle's generic K8 draft.
@@ -978,10 +942,19 @@ presets:
       pd_router: dynamo
       frontend: rust
       features: { tool_calling: false, spec_decoding: false }
-    # GB300: the InferenceX AgentX 3P1D lane (SemiAnalysisAI/InferenceX#4516),
-    # TP8 + EP8 prefill pools (no DCP) and one DEP16 decode pool with MegaMoE.
-    # The pd_cluster Blackwell block replaces the recipe Blackwell baseline, so
-    # the shared serving flags are restated here; role blocks add the rest.
+      # NIXL P/D plus the shared Mooncake store in host DRAM (embedded mode).
+      kv_offload: kv_store_distributed_mooncake
+    # The lane's connector options on the composed MultiConnector.
+    json_args:
+      "--kv-transfer-config":
+        kv_load_failure_policy: recompute
+        NixlConnector: { kv_buffer_device: cuda, kv_connector_extra_config: { enforce_handshake_compat: false, enable_cross_layers_blocks: false } }
+        MooncakeStoreConnector: { kv_load_failure_policy: recompute, kv_connector_extra_config: { enable_cross_layers_blocks: false, enable_offload: false } }
+    # GB300: the InferenceX AgentX 3P1D lane (SemiAnalysisAI/InferenceX#4516,
+    # refreshed in #4638), TP8 + EP8 prefill pools (no DCP) and one DEP16
+    # decode pool with MegaMoE. The pd_cluster Blackwell block replaces the
+    # recipe Blackwell baseline, so the shared serving flags are restated here;
+    # role blocks add the rest.
     args:
       - "--attention-backend"
       - "TOKENSPEED_MLA"
@@ -1034,9 +1007,10 @@ presets:
       # The qualified workers left UCX on its default transports; this
       # replaces the Blackwell PD block's rc,cuda_copy restriction.
       UCX_TLS: "all"
+      DYN_TCP_REQUEST_TIMEOUT: "120"
+      ETCD_LEASE_TTL: "600"
       VLLM_ENGINE_READY_TIMEOUT_S: "3600"
       VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS: "1800"
-      VLLM_RPC_TIMEOUT: "600000"
     roles:
       prefill:
         parallelism: tep
@@ -1063,9 +1037,11 @@ presets:
         # One DEP16 pool (TP1 x DP16 with expert parallelism) across four trays.
         parallelism: dep
         nodes: 4
-        # Decode also quantizes o_proj and the shared experts online to MXFP8.
+        # Decode also quantizes o_proj and the shared experts online to MXFP8,
+        # and runs target and draft attention on KIMI_K3_KV_BALANCED_MLA.
         # FULL_AND_PIECEWISE is vLLM's V1 default, which the qualified decode
-        # ran with; it replaces the DEP strategy's FULL_DECODE_ONLY.
+        # ran with; it replaces the DEP strategy's FULL_DECODE_ONLY. The decode
+        # ranks only add capacity to the Mooncake store (kv_consumer, no lookup).
         args:
           - "--compilation-config"
           - '{"cudagraph_mode":"FULL_AND_PIECEWISE"}'
@@ -1076,81 +1052,44 @@ presets:
           - "--max-cudagraph-capture-size"
           - "512"
           - "--use-replayssm"
+          - "--attention-backend"
+          - "KIMI_K3_KV_BALANCED_MLA"
           - "--attention-config"
           - '{"mla_prefill_backend":"TRTLLM_RAGGED","use_prefill_query_quantization":true,"tokenspeed_mla_min_split_kv":8}'
+        json_args:
+          "--speculative-config": { attention_backend: "KIMI_K3_KV_BALANCED_MLA" }
+          "--kv-transfer-config": { MooncakeStoreConnector: { kv_role: kv_consumer, kv_connector_extra_config: { enable_lookup: false } } }
+    # The lane ran one configuration at three client counts.
+    points:
+      c256:
+        label: "256 (3P1D)"
+        concurrency: 256
+        description: "Three prefill workers and one DEP16 decode worker (40 GPUs)."
+      c288:
+        label: "288 (3P1D)"
+        concurrency: 288
+        description: "Three prefill workers and one DEP16 decode worker (40 GPUs)."
+      c336:
+        label: "336 (3P1D)"
+        concurrency: 336
+        default: true
+        description: "Three prefill workers and one DEP16 decode worker (40 GPUs)."
     runbook: |
-      The GB300 AgentX 3P1D lane in [SemiAnalysis InferenceX](https://inferencex.semianalysis.com/) (`kimik3-fp4-gb300-dynamo-vllm-agentic-dspark-mooncake-dep16-disagg`, [SemiAnalysisAI/InferenceX#4516](https://github.com/SemiAnalysisAI/InferenceX/pull/4516)) runs Dynamo + vLLM with three prefill workers and one decode worker at 336 clients (40 GPUs). Each prefill worker is TP8 + EP8 across two 4-GPU GB300 trays, without DCP. The decode worker is DEP16 (`--tensor-parallel-size 1 --data-parallel-size 16` with expert parallelism) across four trays. Both roles use MegaMoE, `--max-num-seqs 32` (per data-parallel rank on decode) and online MXFP8 for selected target projections. KV moves over NIXL, and all workers share an RDMA Mooncake store. Image `vllm/vllm-openai:nightly-e4a5b5537344b103ed10de7764378ef09604ef3d`, Dynamo `a0de541ea8c861d5ac01d15953ca529bea68637b`.
+      The GB300 AgentX 3P1D lane in [SemiAnalysis InferenceX](https://inferencex.semianalysis.com/) (`kimik3-fp4-gb300-dynamo-vllm-agentic-dspark-mooncake-dep16-disagg`, [SemiAnalysisAI/InferenceX#4516](https://github.com/SemiAnalysisAI/InferenceX/pull/4516), refreshed in [#4638](https://github.com/SemiAnalysisAI/InferenceX/pull/4638)) runs Dynamo + vLLM with three prefill workers and one decode worker at 256, 288 and 336 clients (40 GPUs), with the same configuration at every point. Each prefill worker is TP8 + EP8 across two 4-GPU GB300 trays, without DCP. The decode worker is DEP16 (`--tensor-parallel-size 1 --data-parallel-size 16` with expert parallelism) across four trays and runs the target and the DSpark draft on `KIMI_K3_KV_BALANCED_MLA`; prefill stays on `TOKENSPEED_MLA`. Both roles use MegaMoE, `--max-num-seqs 32` (per data-parallel rank on decode) and online MXFP8 for selected target projections. KV moves over NIXL, and all workers share an RDMA Mooncake store in host DRAM. Image `vllm/vllm-openai:nightly-3966924d566340e971a6da313e790d05492aa505` (vLLM `0.33.0rc1.dev229+g3966924d5`), Dynamo built from `a0de541ea8c861d5ac01d15953ca529bea68637b`, Mooncake `0.3.13.post1` from the image.
 
-      The preset renders one prefill pool and the decode pool. It turns **Spec decoding** off, because the toggle's generic `RedHatAI/Kimi-K3-speculator.dspark` K8 draft was not qualified, and **Tool calling** off (see the parser note under Workers).
-
-      Documented here only: the three prefill workers, the Mooncake store connector (the builder renders a plain NIXL connector, without the qualified MultiConnector, recompute policy and load options), the Mooncake services, prefill `--kv-events-config`, `NVSHMEM_HCA_LIST`, and the frontend's router mode and session affinity (the builder renders round-robin).
-
-      **Target quantization:** `--quantization-config` quantizes these target layers online to MXFP8: the KDA `in_proj_qkvgfab` and MLA `fused_qkv_a_g_proj` projections and the latent-MoE `routed_expert_up_proj`/`routed_expert_down_proj` projections on both roles, plus `o_proj` and the shared experts (`gate_up_proj`, `down_proj`) on decode. The DSpark draft is not quantized: vLLM builds the named draft from its own unquantized config, and the lane's logs show the online quantization pass only for the target layers.
-
-      ### Services
-
-      Start etcd, NATS and a Mooncake master once, as for the [AgentX xPyD preset](#preset-gb300_agentx_pd): `mooncake_master --enable_http_metadata_server=true --default_kv_lease_ttl=60000 --eviction_high_watermark_ratio=0.95 --eviction_ratio=0.10`. The workers' Mooncake store uses P2P handshake metadata, RDMA over the node's NICs, a 160GB global segment and a 4GB local buffer, with offload disabled.
-
-      ### Environment (every worker)
-
-      ```bash
-      export VLLM_ALLREDUCE_USE_FLASHINFER=1 VLLM_ENABLE_GEMM_RS=1 VLLM_USE_NCCL_SYMM_MEM=0
-      export NCCL_CUMEM_ENABLE=1 NCCL_MNNVL_ENABLE=1 NCCL_NVLS_ENABLE=1 NCCL_P2P_LEVEL=NVL
-      export NVSHMEM_REMOTE_TRANSPORT=none NVSHMEM_DISABLE_MNNVL=0 NVSHMEM_IB_ENABLE_IBGDA=0 NVSHMEM_ENABLE_NIC_PE_MAPPING=1
-      export NVSHMEM_HCA_LIST=<the node's RDMA HCAs>
-      export VLLM_USE_V2_MODEL_RUNNER=1 VLLM_USE_RUST_FRONTEND=1 VLLM_SSM_CONV_STATE_LAYOUT=DS
-      export VLLM_MOONCAKE_LOAD_RECV_THREADS=4 MC_SLICE_SIZE=1048576 MC_MAX_MR_SIZE=4294967296 MC_ENABLE_DEST_DEVICE_AFFINITY=1
-      export UCX_MEMTYPE_CACHE=n UCX_MEMTYPE_REG_WHOLE=n UCX_RCACHE_MAX_UNRELEASED=1024 UCX_TCP_AF_PRIO=inet
-      export DYN_REQUEST_PLANE=tcp DYN_TCP_REQUEST_TIMEOUT=120 ETCD_LEASE_TTL=600
-      export VLLM_ENGINE_READY_TIMEOUT_S=3600 VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS=1800 VLLM_RPC_TIMEOUT=600000
-      ```
+      Each point renders one prefill worker and the decode worker; start three prefill workers with the same command. Run each prefill worker on its two trays with the same `--master-addr` (the first tray's address), `--nnodes 2` and `--node-rank 0` or `1`, and the decode worker on each of its four trays with `--data-parallel-start-rank` 0, 4, 8 and 12 and the same `--data-parallel-address`. Set `NVSHMEM_HCA_LIST` to the node's RDMA HCAs on every worker.
 
-      Prefill workers also set `VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT=0`, `DYN_BACKEND_ADMISSION_ENGINE_REQUEST_LIMIT=1024` and `DYN_BACKEND_ADMISSION_ENGINE_REQUEST_WAITING_LIMIT=1024`.
+      `KIMI_K3_KV_BALANCED_MLA` is an FP8 MLA decode backend that plans its split-KV work on the device; it is never selected automatically, and prefill steps use the configured MLA prefill backend.
 
-      ### Workers
+      The preset turns **Spec decoding** off, because the toggle's generic `RedHatAI/Kimi-K3-speculator.dspark` K8 draft was not qualified, and **Tool calling** off; for tool calling and reasoning output under Dynamo, add `--dyn-tool-call-parser kimi_k3 --dyn-reasoning-parser kimi_k3 --dyn-enable-structural-tag` to both roles (`python3 -m dynamo.vllm` rejects the vLLM API-server flags `--enable-auto-tool-choice` and `--tool-call-parser`).
 
-      Prefill (three workers; run each on its two trays with the same `--master-addr`, the first tray's address, `--nnodes 2` and `--node-rank 0` or `1`):
+      The preset selects the distributed (embedded) Mooncake store on the **KV Offload** row with the same connector options as the [AgentX xPyD preset](#preset-gb300_agentx_pd); the decode ranks only add capacity to the store (`kv_consumer` with `enable_lookup: false`). Documented here only: the three prefill workers, the lane's Mooncake master options and store sizing, `NVSHMEM_HCA_LIST`, and the frontend's router mode and session affinity (the builder renders round-robin).
 
-      ```bash
-      python3 -m dynamo.vllm --model moonshotai/Kimi-K3 --served-model-name moonshotai/Kimi-K3 \
-        --disaggregation-mode prefill --request-plane tcp \
-        --nnodes 2 --node-rank <0|1> --master-addr <first-tray-ip> --device-ids 0,1,2,3 \
-        --tensor-parallel-size 8 --enable-expert-parallel --enable-ep-weight-filter --moe-backend deep_gemm_mega_moe \
-        --quantization-config '{"targets":{"*.in_proj_qkvgfab":"mxfp8","*.fused_qkv_a_g_proj":"mxfp8","*.routed_expert_down_proj":"mxfp8","*.routed_expert_up_proj":"mxfp8"}}' \
-        --attention-backend TOKENSPEED_MLA \
-        --attention-config '{"mla_prefill_backend":"TRTLLM_RAGGED","use_prefill_query_quantization":true,"trtllm_ragged_backend":"cute-dsl"}' \
-        --speculative-config '{"model":"Inferact/Kimi-K3-DSpark","attention_backend":"TOKENSPEED_MLA","method":"dspark","num_speculative_tokens":4,"draft_sample_method":"probabilistic","rejection_sample_method":"block"}' \
-        --kv-cache-dtype fp8 --mamba-ssm-cache-dtype bfloat16 \
-        --enable-prefix-caching --prefix-match-unit 128 \
-        --gpu-memory-utilization 0.95 --max-num-seqs 32 --max-num-batched-tokens 16384 --max-cudagraph-capture-size 8192 \
-        --compilation-config '{"cudagraph_mode":"FULL_AND_PIECEWISE","cudagraph_capture_sizes":[4,8,12,16,20,24,32,40,48,64,80,96,128,160,256,384,512,640,768,896,1024,1152,1280,1408,1536,1664,1792,1920,2048,2304,2560,2816,3072,3328,3584,3840,4096,4608,5120,5632,6144,6656,7168,7680,8192]}' \
-        --enable-cumem-allocator --enable-flashinfer-autotune \
-        --load-format fastsafetensors --stream-interval 10 --language-model-only --trust-remote-code \
-        --kv-events-config '{"publisher":"zmq","topic":"kv-events","enable_kv_cache_events":true,"endpoint":"tcp://*:5200"}' \
-        --kv-transfer-config '{"kv_connector":"MultiConnector","kv_role":"kv_both","kv_load_failure_policy":"recompute","kv_connector_extra_config":{"connectors":[{"kv_connector":"NixlConnector","kv_role":"kv_both","kv_load_failure_policy":"fail","kv_buffer_device":"cuda","kv_connector_extra_config":{"enforce_handshake_compat":false,"enable_cross_layers_blocks":false}},{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_both","kv_load_failure_policy":"recompute","kv_connector_extra_config":{"load_async":true,"lookup_async":true,"enable_cross_layers_blocks":false,"enable_offload":false}}]}}'
-      ```
-
-      Decode (one DEP16 worker; run on each of its four trays with `--data-parallel-start-rank` 0, 4, 8 and 12 and the same `--data-parallel-address`, the first tray's address):
+      **Target quantization:** `--quantization-config` quantizes these target layers online to MXFP8: the KDA `in_proj_qkvgfab` and MLA `fused_qkv_a_g_proj` projections and the latent-MoE `routed_expert_up_proj`/`routed_expert_down_proj` projections on both roles, plus `o_proj` and the shared experts (`gate_up_proj`, `down_proj`) on decode. The DSpark draft is not quantized: vLLM builds the named draft from its own unquantized config, and the lane's logs show the online quantization pass only for the target layers.
 
-      ```bash
-      python3 -m dynamo.vllm --model moonshotai/Kimi-K3 --served-model-name moonshotai/Kimi-K3 \
-        --disaggregation-mode decode --request-plane tcp --device-ids 0,1,2,3 \
-        --tensor-parallel-size 1 --data-parallel-size 16 --data-parallel-size-local 4 \
-        --data-parallel-start-rank <0|4|8|12> --data-parallel-address <first-tray-ip> --data-parallel-rpc-port <port> \
-        --data-parallel-hybrid-lb --enable-expert-parallel --enable-ep-weight-filter --moe-backend deep_gemm_mega_moe \
-        --quantization-config '{"targets":{"*.in_proj_qkvgfab":"mxfp8","*.fused_qkv_a_g_proj":"mxfp8","*.routed_expert_down_proj":"mxfp8","*.routed_expert_up_proj":"mxfp8","*.o_proj":"mxfp8","*.shared_experts.gate_up_proj":"mxfp8","*.shared_experts.down_proj":"mxfp8"}}' \
-        --attention-backend TOKENSPEED_MLA \
-        --attention-config '{"mla_prefill_backend":"TRTLLM_RAGGED","use_prefill_query_quantization":true,"tokenspeed_mla_min_split_kv":8}' \
-        --speculative-config '{"model":"Inferact/Kimi-K3-DSpark","attention_backend":"TOKENSPEED_MLA","method":"dspark","num_speculative_tokens":4,"draft_sample_method":"probabilistic","rejection_sample_method":"block"}' \
-        --kv-cache-dtype fp8 --mamba-ssm-cache-dtype bfloat16 \
-        --no-enable-prefix-caching --prefix-match-unit 128 \
-        --gpu-memory-utilization 0.95 --max-num-seqs 32 --max-num-batched-tokens 512 --max-cudagraph-capture-size 512 --use-replayssm \
-        --enable-cumem-allocator --enable-flashinfer-autotune \
-        --load-format fastsafetensors --stream-interval 10 --language-model-only --trust-remote-code \
-        --kv-transfer-config '{"kv_connector":"MultiConnector","kv_role":"kv_both","kv_load_failure_policy":"recompute","kv_connector_extra_config":{"connectors":[{"kv_connector":"NixlConnector","kv_role":"kv_both","kv_load_failure_policy":"fail","kv_buffer_device":"cuda","kv_connector_extra_config":{"enforce_handshake_compat":false,"enable_cross_layers_blocks":false}},{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_consumer","kv_load_failure_policy":"recompute","kv_connector_extra_config":{"load_async":true,"lookup_async":true,"enable_lookup":false,"enable_cross_layers_blocks":false,"enable_offload":false}}]}}'
-      ```
+      ### Services
 
-      For tool calling and reasoning output, add `--reasoning-parser kimi_k3 --dyn-tool-call-parser kimi_k3 --dyn-reasoning-parser kimi_k3 --dyn-enable-structural-tag` to both roles; `python3 -m dynamo.vllm` rejects the vLLM API-server flags `--enable-auto-tool-choice` and `--tool-call-parser`.
+      Start etcd, NATS and a Mooncake master once, and give every worker the Mooncake store configuration through `MOONCAKE_CONFIG_PATH`, as for the [AgentX xPyD preset](#preset-gb300_agentx_pd). The store uses P2P handshake metadata, RDMA over the node's NICs, a 160GB global segment and a 4GB local buffer, with offload disabled.
 
       ### Frontend
 
@@ -1238,6 +1177,129 @@ presets:
 
       Make the same target checkpoint available at the same local path on every node, and make `Inferact/Kimi-K3-DSpark` available through the mounted Hugging Face cache. Run the preset's commands inside the nightly image on every allocated node, with four visible GPUs, host networking, sufficient shared memory, and the devices and network access required by your NVLink/RDMA deployment. Use the reachable address of rank 0 for `MASTER_ADDR` on every node and a distinct `NODE_RANK` from `0` through `NODES - 1`. The lane served the checkpoint from a local path with `--served-model-name moonshotai/Kimi-K3` and `--device-ids 0,1,2,3`; the builder serves the Hugging Face id.
 
+  b300_agentx_agg:
+    label: "AgentX aggregated (Mooncake)"
+    hardware: b300
+    strategy: single_node_tp
+    workload: "Agentic (InferenceX AgentX)"
+    description: "TP8 + DCP8 on one B300 node with the Inferact DSpark draft (K7 at concurrency 1, K5 at 4-8, K3 at 12-28) and an embedded RDMA Mooncake DRAM store."
+    source: { label: "InferenceX #4381", href: "https://github.com/SemiAnalysisAI/InferenceX/pull/4381", key: "kimik3-fp4-b300-vllm-agentic-dspark", revision: 809f7045172ea5effde45cef7cbe02918f25f968 }
+    image: "vllm/vllm-openai@sha256:8c1434bfee7d7626eef78ef70600434ffba4a1b8813cb1d1c8707ad25f47035e"
+    # vLLM's own Python frontend. The draft is in args (TOKENSPEED_MLA attention,
+    # Inferact checkpoint), so the generic Spec decoding toggle stays off.
+    select:
+      frontend: python
+      kv_offload: kv_store_distributed_mooncake
+      features: { spec_decoding: false, text_only: true, tool_calling: true, reasoning: true }
+    # Lane settings that hold at every concurrency. The Blackwell baseline already
+    # supplies fastsafetensors, 1M context, fp8 KV, prefix caching, TOKENSPEED_MLA
+    # decode and --prefix-match-unit 128. The recipe has no DCP selector, so the
+    # lane's DCP8 pairing (a2a, TRTLLM_RAGGED prefill) is written out here.
+    args:
+      - "--decode-context-parallel-size"
+      - "8"
+      - "--dcp-comm-backend"
+      - "a2a"
+      - "--max-num-batched-tokens"
+      - "16384"
+      - "--enable-flashinfer-autotune"
+      - "--enable-cumem-allocator"
+      - "--mamba-ssm-cache-dtype"
+      - "bfloat16"
+      - "--stream-interval"
+      - "10"
+      - "--gpu-memory-utilization"
+      - "0.92"
+      - "--moe-backend"
+      - "auto"
+      - "--speculative-config"
+      - '{"model":"Inferact/Kimi-K3-DSpark","attention_backend":"TOKENSPEED_MLA","method":"dspark","num_speculative_tokens":3,"draft_sample_method":"probabilistic","rejection_sample_method":"block"}'
+    remove_args:
+      - "--no-enable-flashinfer-autotune"
+      - "--disable-flashinfer-autotune"
+    env:
+      VLLM_ENABLE_K3_LATENT_MOE_TAIL_FUSION: "1"
+      VLLM_USE_DIRECT_DCP_A2A: "1"
+      VLLM_USE_DIRECT_DCP_Q_GATHER: "1"
+      VLLM_USE_DIRECT_DCP_KV_GATHER: "1"
+      VLLM_RPC_TIMEOUT: "600000"
+      VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS: "1800"
+      VLLM_PREFIX_CACHE_RETENTION_INTERVAL: "0"
+      VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS: "0"
+      VLLM_HTTP_TIMEOUT_KEEP_ALIVE: "900"
+      VLLM_MOONCAKE_LOAD_RECV_THREADS: "4"
+      TORCH_CUDA_ARCH_LIST: "10.0"
+      PYTHONHASHSEED: "42"
+      MC_STORE_MEMCPY: "1"
+      MC_ENABLE_DEST_DEVICE_AFFINITY: "1"
+      MC_SLICE_SIZE: "1048576"
+      MC_WORKERS_PER_CTX: "4"
+    json_args:
+      "--attention-config": { mla_prefill_backend: "TRTLLM_RAGGED", use_prefill_query_quantization: true }
+      "--compilation-config": { cudagraph_mode: "FULL_AND_PIECEWISE" }
+      "--kv-transfer-config": { kv_load_failure_policy: "recompute", kv_connector_extra_config: { load_async: true, lookup_async: true, enable_offload: false } }
+    # One point per measured concurrency. S is --max-num-seqs (2x concurrency); the
+    # graphs capture (1 + K) x 1..S tokens, then the shared 256..8192 tail.
+    points:
+      c1:
+        label: "1"
+        concurrency: 1
+        args: ["--max-num-seqs", "2"]
+        json_args:
+          "--speculative-config": { num_speculative_tokens: 7 }
+          "--compilation-config": { cudagraph_capture_sizes: [8, 16, 64, 128, 256, 512, 1024, 2048, 4096, 8192] }
+      c4:
+        label: "4"
+        concurrency: 4
+        args: ["--max-num-seqs", "8"]
+        json_args:
+          "--speculative-config": { num_speculative_tokens: 5 }
+          "--compilation-config": { cudagraph_capture_sizes: [6, 12, 18, 24, 30, 36, 42, 48, 54, 60, 66, 72, 78, 84, 90, 96, 128, 256, 512, 1024, 2048, 4096, 8192] }
+      c8:
+        label: "8"
+        concurrency: 8
+        args: ["--max-num-seqs", "16"]
+        json_args:
+          "--speculative-config": { num_speculative_tokens: 5 }
+          "--compilation-config": { cudagraph_capture_sizes: [6, 12, 18, 24, 30, 36, 42, 48, 54, 60, 66, 72, 78, 84, 90, 96, 128, 256, 512, 1024, 2048, 4096, 8192] }
+      c12:
+        label: "12"
+        concurrency: 12
+        args: ["--max-num-seqs", "24"]
+        json_args:
+          "--compilation-config": { cudagraph_capture_sizes: [4, 8, 12, 16, 20, 24, 28, 32, 36, 40, 44, 48, 52, 56, 60, 64, 68, 72, 76, 80, 84, 88, 92, 96, 100, 104, 108, 112, 128, 256, 512, 1024, 2048, 4096, 8192] }
+      c14:
+        label: "14"
+        concurrency: 14
+        args: ["--max-num-seqs", "28"]
+        json_args:
+          "--compilation-config": { cudagraph_capture_sizes: [4, 8, 12, 16, 20, 24, 28, 32, 36, 40, 44, 48, 52, 56, 60, 64, 68, 72, 76, 80, 84, 88, 92, 96, 100, 104, 108, 112, 128, 256, 512, 1024, 2048, 4096, 8192] }
+      c16:
+        label: "16"
+        concurrency: 16
+        default: true
+        args: ["--max-num-seqs", "32"]
+        json_args:
+          "--compilation-config": { cudagraph_capture_sizes: [4, 8, 12, 16, 20, 24, 28, 32, 36, 40, 44, 48, 52, 56, 60, 64, 68, 72, 76, 80, 84, 88, 92, 96, 100, 104, 108, 112, 116, 120, 124, 128, 256, 512, 1024, 2048, 4096, 8192] }
+      c24:
+        label: "24"
+        concurrency: 24
+        args: ["--max-num-seqs", "48"]
+        json_args:
+          "--compilation-config": { cudagraph_capture_sizes: [4, 8, 12, 16, 20, 24, 28, 32, 36, 40, 44, 48, 52, 56, 60, 64, 68, 72, 76, 80, 84, 88, 92, 96, 100, 104, 108, 112, 116, 120, 124, 128, 132, 136, 140, 144, 148, 152, 156, 160, 164, 168, 172, 176, 180, 184, 188, 192, 196, 200, 204, 208, 212, 216, 220, 224, 228, 232, 236, 240, 244, 248, 252, 256, 512, 1024, 2048, 4096, 8192] }
+      c28:
+        label: "28"
+        concurrency: 28
+        args: ["--max-num-seqs", "56"]
+        json_args:
+          "--compilation-config": { cudagraph_capture_sizes: [4, 8, 12, 16, 20, 24, 28, 32, 36, 40, 44, 48, 52, 56, 60, 64, 68, 72, 76, 80, 84, 88, 92, 96, 100, 104, 108, 112, 116, 120, 124, 128, 132, 136, 140, 144, 148, 152, 156, 160, 164, 168, 172, 176, 180, 184, 188, 192, 196, 200, 204, 208, 212, 216, 220, 224, 228, 232, 236, 240, 244, 248, 252, 256, 512, 1024, 2048, 4096, 8192] }
+    runbook: |
+      The B300 AgentX lane in [SemiAnalysis InferenceX](https://inferencex.semianalysis.com/) (`kimik3-fp4-b300-vllm-agentic-dspark`, [SemiAnalysisAI/InferenceX#4381](https://github.com/SemiAnalysisAI/InferenceX/pull/4381), pinned to [`809f704`](https://github.com/SemiAnalysisAI/InferenceX/tree/809f7045172ea5effde45cef7cbe02918f25f968/inferencex-e2e/benchmarks/single_node/srt-slurm-recipes/kimik3/vllm/b300-fp4-mtp)) serves the default `moonshotai/Kimi-K3` checkpoint (MXFP4 weights) on one 8-GPU B300 node with DSpark speculative decoding and an embedded Mooncake DRAM tier. It is a benchmark configuration; the settings in this preset are its production-useful parts. Draft length falls with concurrency (K7 at 1, K5 at 4-8, K3 at 12-28), and each point sets `--max-num-seqs` to twice the concurrency with CUDA graphs covering `(1 + K) x 1..max-num-seqs` tokens. Throughput runs use InferenceX's golden acceptance lengths for DSpark (3.84, 3.62 and 3.00 for K7, K5 and K3).
+
+      **Mooncake store.** The builder's **Mooncake** option renders the master and store config with a 100GB segment, which is the default. The lane gives each of the 8 TP ranks a 256GB `global_segment_size` and a 4GB `local_buffer_size`, 2,080 GiB of DRAM across the node in total, so raise the segment above the default only when the node has that much pinned memory to spare. The lane's master sets `--eviction_high_watermark_ratio=0.95 --eviction_ratio=0.10`, matching the builder, and leaves the lease TTL at the Mooncake default. The store runs embedded over RDMA with `P2PHANDSHAKE` metadata. Pin `device_name` to one ACTIVE InfiniBand port driven by `mlx5_core`: this lane's ranks use a single rail because the rails are isolated from each other, so a port that is down or an EFA `rdmap*` device (no RC queue pairs) fails registration. An InfiniBand port addresses peers with GID index 0 (`MC_GID_INDEX=0`).
+
+      **Kept out of the preset.** `--disable-uvicorn-access-log`, `PYTHONNOUSERSITE=1` and the benchmark client variables are InferenceX harness or logging controls with no serving effect. `--moe-backend auto` is vLLM's default selection and is stated only to match the lane. The lane's `RDMAV_DRIVERS` and `WITH_NVIDIA_PEERMEM=0` are specific to the DSXE cluster (a host-library mount for its `libmlx5` verbs provider, and peermem not loaded) and apply only on nodes with the same setup.
+
   mi355x_agentx:
     label: "AgentX TP8"
     hardware: mi355x
diff --git a/recipes/models/zai-org/GLM-5.3-Flash.yaml b/recipes/models/zai-org/GLM-5.3-Flash.yaml
index 875be29630a5..de4d7213128a 100644
--- a/recipes/models/zai-org/GLM-5.3-Flash.yaml
+++ b/recipes/models/zai-org/GLM-5.3-Flash.yaml
@@ -4,7 +4,7 @@ meta:
   provider: "GLM (Z-AI)"
   description: "GLM-5.3-Flash is a 320B-total / 18B-active multimodal MoE with hybrid KDA and sparse MLA attention, native FP8 weights, MTP, and a 1M-token context window."
   date_added: 2026-08-27
-  date_updated: 2026-10-08
+  date_updated: 2026-10-19
   tasks:
     - text
     - multimodal
@@ -14,23 +14,26 @@ meta:
     b200: verified
     gb200: verified
     mi355x: verified
-    ascend_950pr_8x: verified
 
 model:
   model_id: "zai-org/GLM-5.3-Flash"
   min_vllm_version: "0.29.0"
   nightly_required: true
+  default_frontend: rust
   docker_image:
     nvidia: "vllm/vllm-openai:nightly"
     amd: "vllm/vllm-openai-rocm:nightly"
   install:
     docker:
-      note: "NVIDIA: use vllm/vllm-openai:glm53-flash until support lands in the standard image. AMD: use the official vllm/vllm-openai-rocm:nightly image, which includes merged GLM-5.3-Flash support. Ascend 950PR: pin quay.io/ascend/vllm-ascend:nightly-main-a5 (same as GLM-5.3) with GLM-5.3-Flash overlays vLLM#53906 and vllm-ascend#15127."
+      note: "Native GLM5Next rendering requires an image containing vllm-project/vllm#62424, including its Rust binary. Use NVIDIA/ROCm nightly images only once they include this change. Ascend 950PR also requires the nightly-main-a5 platform support and GLM-5.3-Flash overlays vLLM#53906 and vllm-ascend#15127."
     pip: false
   architecture: moe
   parameter_count: "321B"
   active_parameters: "18B"
   context_length: 1048576
+  base_args:
+    - "--tokenizer-mode"
+    - "glm5_next"
   base_env:
     VLLM_ENGINE_READY_TIMEOUT_S: "3600"
 
@@ -76,13 +79,6 @@ variants:
     vram_minimum_gb: 386
     description: "Default native FP8 checkpoint (zai-org/GLM-5.3-Flash)"
     hardware_overrides:
-      # Exact-GPU pin: computeDockerMeta has no Huawei brand, so a model-level
-      # docker_image string would only affect NVIDIA. Hide pip — official
-      # wheels have no NPU kernels.
-      ascend_950pr_8x:
-        docker_image: "quay.io/ascend/vllm-ascend:nightly-main-a5"
-        install:
-          pip: false
       mi350x: &gfx950_aiter
         extra_env:
           VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS: "1"
@@ -116,22 +112,6 @@ kv_offload_support:
   offloading_cpu: verified
   offloading_fs: verified
 
-strategy_hardware:
-  single_node_tep:
-    npu: unsupported
-  single_node_dep:
-    npu: unsupported
-  multi_node_tp:
-    npu: unsupported
-  multi_node_tp_pp:
-    npu: unsupported
-  multi_node_tep:
-    npu: unsupported
-  multi_node_dep:
-    npu: unsupported
-  pd_cluster:
-    npu: unsupported
-
 hardware_overrides:
   hopper:
     extra_args:
@@ -149,34 +129,12 @@ hardware_overrides:
       - "ROCM_AITER_MLA_SPARSE"
     extra_env:
       VLLM_ROCM_USE_AITER: "1"
-  # Shared by 950PR 8x (8×128 GB). SFA / kpool indexer needs --block-size 512.
-  npu:
-    extra_args:
-      - "--trust-remote-code"
-      - "--block-size"
-      - "512"
-      - "--max-model-len"
-      - "131072"
-      - "--max-num-seqs"
-      - "16"
-      - "--max-num-batched-tokens"
-      - "8192"
-      - "--gpu-memory-utilization"
-      - "0.90"
-    extra_env:
-      PYTORCH_NPU_ALLOC_CONF: "expandable_segments:True"
-      HCCL_BUFFSIZE: "512"
-      HCCL_OP_EXPANSION_MODE: "AIV"
-      HCCL_CONNECT_TIMEOUT: "7200"
 
 hardware_notes:
   mi350x: &cdna4_indexer_topk_notes
     title: "Long-context indexer top-k opt-in"
     body: "When serving long contexts (>=64k), add `--sparse-indexer-topk-backend aiter` for improved decode performance."
   mi355x: *cdna4_indexer_topk_notes
-  ascend_950pr_8x:
-    title: "vLLM Ascend A5 image required"
-    body: "Single-node TP8 on 8×950PR. Pin quay.io/ascend/vllm-ascend:nightly-main-a5 and overlay vLLM#53906 / vllm-ascend#15127. Do not pass --quantization ascend. --block-size 512 is required (kpool indexer). EP is not available on this interconnect."
 
 strategy_overrides:
   pd_cluster:
@@ -219,7 +177,7 @@ guide: |
 
   - **Weights:** about 306 GiB for the default native FP8 checkpoint before runtime
     and KV-cache overhead; the BF16 variant requires roughly twice the weight memory
-  - **vLLM:** use docker before the integration is included in the public repo
+  - **vLLM:** an image containing [native GLM5Next support (#62424)](https://github.com/vllm-project/vllm/pull/62424), including the matching Rust frontend binary. Older images need an updated build for the commands below.
   - **FlashInfer:** 0.6.17 or newer is required for NoPE sparse MLA
 
   The default model ID, `zai-org/GLM-5.3-Flash`, is the FP8 checkpoint. Select the BF16
@@ -227,12 +185,26 @@ guide: |
   `RedHatAI/GLM-5.3-Flash-NVFP4`. Its MoE experts are quantized to 4-bit (NVFP4),
   and it requires NVIDIA Blackwell GPUs.
 
+  ## Frontends
+
+  The command builder defaults to **Rust** and adds `--tokenizer-mode glm5_next`.
+  With `VLLM_USE_RUST_FRONTEND=1`, this selects native GLM5Next chat rendering for
+  the GLM-5.3-Flash template, including tool calls and reasoning options. Image and
+  video inputs are not yet supported on the Rust frontend; use the Python frontend
+  for multimodal requests.
+
+  Select **Python** in the command builder, or set `VLLM_USE_RUST_FRONTEND=0`,
+  to use the Hugging Face renderer and processor on the same build. On older builds
+  that do not recognize `glm5_next`, also replace `--tokenizer-mode glm5_next` with
+  `--tokenizer-mode auto`.
+
   ## Launching the server
 
   ### FP8 with TP4 and MTP on one GB200 tray
 
   ```bash
-  vllm serve zai-org/GLM-5.3-Flash \
+  VLLM_USE_RUST_FRONTEND=1 vllm serve zai-org/GLM-5.3-Flash \
+    --tokenizer-mode glm5_next \
     --tensor-parallel-size 4 \
     --kv-cache-dtype fp8 \
     --speculative-config '{"method":"mtp","num_speculative_tokens":5}' \
@@ -259,7 +231,9 @@ guide: |
   VLLM_KV_CACHE_LAYOUT=HND \
   UCX_NET_DEVICES=all \
   VLLM_NIXL_SIDE_CHANNEL_PORT=5557 \
+  VLLM_USE_RUST_FRONTEND=1 \
   vllm serve zai-org/GLM-5.3-Flash \
+    --tokenizer-mode glm5_next \
     --port 8001 \
     --tensor-parallel-size 4 \
     --kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_producer"}' \
@@ -275,7 +249,9 @@ guide: |
   VLLM_KV_CACHE_LAYOUT=HND \
   UCX_NET_DEVICES=all \
   VLLM_NIXL_SIDE_CHANNEL_PORT=5558 \
+  VLLM_USE_RUST_FRONTEND=1 \
   vllm serve zai-org/GLM-5.3-Flash \
+    --tokenizer-mode glm5_next \
     --port 8002 \
     --tensor-parallel-size 4 \
     --kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_consumer"}' \
@@ -433,15 +409,16 @@ guide: |
   
   ## Running on AMD Instinct
 
-  Use the official `vllm/vllm-openai-rocm:nightly` image. GLM-5.3-Flash support
-  landed in [vLLM #53906](https://github.com/vllm-project/vllm/pull/53906).
+  Use a `vllm/vllm-openai-rocm:nightly` build containing #62424 for native Rust
+  rendering. Model support landed in [vLLM #53906](https://github.com/vllm-project/vllm/pull/53906).
 
   ```bash
   export VLLM_ROCM_USE_AITER=1
   # gfx950 (MI350X/MI355X) only
   export VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=1
   export VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA=1
-  vllm serve zai-org/GLM-5.3-Flash \
+  VLLM_USE_RUST_FRONTEND=1 vllm serve zai-org/GLM-5.3-Flash \
+    --tokenizer-mode glm5_next \
     --served-model-name zai-org/GLM-5.3-Flash \
     --tensor-parallel-size 4 \
     --max-num-seqs 512 \
@@ -466,7 +443,8 @@ guide: |
   (MI350X/MI355X), using the topk indexer backend from AITER is recommended:
 
   ```bash
-  vllm serve zai-org/GLM-5.3-Flash \
+  VLLM_USE_RUST_FRONTEND=1 vllm serve zai-org/GLM-5.3-Flash \
+    --tokenizer-mode glm5_next \
     --attention-backend ROCM_AITER_MLA_SPARSE \
     --sparse-indexer-topk-backend aiter
   ```
diff --git a/recipes/models/zai-org/GLM-5.3.yaml b/recipes/models/zai-org/GLM-5.3.yaml
index 3ac7e3ca2a8e..2570d061215a 100644
--- a/recipes/models/zai-org/GLM-5.3.yaml
+++ b/recipes/models/zai-org/GLM-5.3.yaml
@@ -4,7 +4,7 @@ meta:
   provider: "GLM (Z-AI)"
   description: "GLM-5.3 — Frontier Coding with Emergent Cyber Capabilities"
   date_added: 2026-08-28
-  date_updated: 2026-10-15
+  date_updated: 2026-10-19
   tasks:
     - text
   performance_headline: "Latest GLM-5 series MoE — native FP8 by default, 5-token MTP, 1M context"
@@ -20,9 +20,6 @@ meta:
     mi325x: verified
     # FP8 + 5-token MTP validated 2026-10-03 on 8×MI350X (TP8).
     mi350x: verified
-    # Validated 2026-08-30 on one Atlas A5 node (8×950PR, TP8, graph mode).
-    # Native FP8 path needs vllm-ascend#15346.
-    ascend_950pr_8x: verified
     # NVFP4 + --cpu-offload-gb on 1× GB300 (252 GB HBM + unified host RAM).
     # Do not raise taxonomy vram_gb — host RAM is recipe-specific, not HBM.
     dgx_station_gb300: verified
@@ -37,11 +34,18 @@ model:
   default_frontend: rust
   docker_image:
     amd: "vllm/vllm-openai-rocm:v0.29.0"
+  install:
+    pip:
+      note: "Native GLM5Next rendering requires a build containing vllm-project/vllm#62424, including its Rust binary. The minimum version alone does not include this support."
+    docker:
+      note: "Use an image containing vllm-project/vllm#62424, including its Rust binary. Older hardware-specific image pins require an updated build for --tokenizer-mode glm5_next."
   architecture: moe
   parameter_count: "743B"
   active_parameters: "39B"
   context_length: 1048576
   base_args:
+    - "--tokenizer-mode"
+    - "glm5_next"
     - "--kv-cache-dtype"
     - "fp8_e4m3"
   base_env: {}
@@ -69,10 +73,6 @@ features:
           - "--speculative-config"
           - '{"method":"mtp","num_speculative_tokens":5}'
         hardware_overrides:
-          npu:
-            args:
-              - "--speculative-config"
-              - '{"method":"mtp","num_speculative_tokens":1,"enforce_eager":true}'
           dgx_station_gb300:
             args:
               - "--speculative-config"
@@ -83,7 +83,6 @@ features:
         variants: [default]
         hardware:
           amd: unsupported
-          npu: unsupported
         args:
           - "--speculative-config"
           - '{"method":"dflash","model":"incoai/GLM-5.3-DFlash2","num_speculative_tokens":7}'
@@ -111,13 +110,6 @@ variants:
         extra_args:
           - "--max-model-len"
           - "262144"
-      # Exact-GPU pin: computeDockerMeta has no Huawei brand, so a model-level
-      # docker_image string would only affect NVIDIA. Hide pip — official
-      # wheels have no NPU kernels.
-      ascend_950pr_8x:
-        docker_image: "quay.io/ascend/vllm-ascend:nightly-main-a5"
-        install:
-          pip: false
   nvfp4:
     model_id: "Inferact/GLM-5.3-NVFP4"
     precision: nvfp4
@@ -186,22 +178,6 @@ compatible_strategies:
   - multi_node_dep
   - pd_cluster
 
-# EP/MC2 cannot acquire HCCL links on Atlas 350 (950PR). Multi-node and PD
-# were not validated. Only single-node TP8 is offered on NPU.
-strategy_hardware:
-  single_node_tep:
-    npu: unsupported
-  multi_node_tp:
-    npu: unsupported
-  multi_node_tp_pp:
-    npu: unsupported
-  multi_node_tep:
-    npu: unsupported
-  multi_node_dep:
-    npu: unsupported
-  pd_cluster:
-    npu: unsupported
-
 hardware_overrides:
   hopper:
     # H100/H200 use plain fp8 KV cache; Blackwell keeps base_args' fp8_e4m3.
@@ -218,35 +194,10 @@ hardware_overrides:
     extra_env:
       VLLM_ROCM_USE_AITER: "1"
       VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS: "1"
-  # Shared by 950PR 8x (8×128 GB). last-wins shadows --kv-cache-dtype fp8_e4m3
-  # from base_args; SFA on 950 needs --block-size 512.
-  npu:
-    extra_args:
-      - "--trust-remote-code"
-      - "--kv-cache-dtype"
-      - "auto"
-      - "--block-size"
-      - "512"
-      - "--max-model-len"
-      - "131072"
-      - "--max-num-seqs"
-      - "16"
-      - "--gpu-memory-utilization"
-      - "0.9"
-    extra_env:
-      PYTORCH_NPU_ALLOC_CONF: "expandable_segments:True"
-      OMP_PROC_BIND: "false"
-      OMP_NUM_THREADS: "1"
-      TASK_QUEUE_ENABLE: "1"
-      HCCL_BUFFSIZE: "1024"
-      HCCL_OP_EXPANSION_MODE: "AIV"
 
 strategy_overrides: {}
 
 hardware_notes:
-  ascend_950pr_8x:
-    title: "vLLM Ascend A5 image required"
-    body: "Single-node TP8 on 8×950PR. Pin quay.io/ascend/vllm-ascend:nightly-main-a5 and overlay vllm-ascend#15346 for native FP8. EP is not available on this interconnect. Set HCCL_BUFFSIZE=1024 before launch."
   dgx_station_gb300:
     title: "CPU weight offload (NVFP4)"
     body: "1× GB300 (252 GB HBM) plus unified host RAM. NVFP4 offloads MoE expert weights (w13_weight, w2_weight) with --offload-backend uva and --cpu-offload-gb 222. Set VLLM_WEIGHT_OFFLOADING_DISABLE_PIN_MEMORY=1. Context is capped at 262K / 8 seqs. FP8/BF16 still do not fit this box."
@@ -255,32 +206,41 @@ guide: |
   ## Overview
 
   **GLM-5.3** is the newest model in the GLM-5 series — a ~743B-parameter MoE (39B active) from **Z-AI**.
-  Architecture and serving flags are identical to GLM-5.2; the packaging change is that the
-  **default `zai-org/GLM-5.3` checkpoint is now native FP8** — BF16 weights live under the suffixed `zai-org/GLM-5.3-BF16` repo.
+  The **default `zai-org/GLM-5.3` checkpoint is native FP8** — BF16 weights live under the suffixed `zai-org/GLM-5.3-BF16` repo.
 
   The FP8 checkpoint fits on a single 8xH200 / 8xH20 node and — with FP8 KV cache — reaches the full 1M-token context on 8xB200.
   An NVFP4 checkpoint for Blackwell is published by Inferact (`Inferact/GLM-5.3-NVFP4`).
 
   ## Prerequisites
 
-  - **vLLM 0.28.0** or newer.
+  - **vLLM:** a build containing [native GLM5Next support (#62424)](https://github.com/vllm-project/vllm/pull/62424), including the matching Rust frontend binary.
+  - **Transformers:** 5.15.0 or newer.
   - **GPU:** 8xH200 or 8xH20 (141 GB each) for single-node FP8; 8xB200 (180 GB each) for the full 1M context.
 
   ## Installation
 
-  ```bash
-  uv venv
-  source .venv/bin/activate
-  uv pip install "vllm==0.28.0" --torch-backend=auto
-  uv pip install "transformers>=5.15.0"
-  ```
+  Use a wheel or container built with #62424 for the Rust commands below.
+  The older release and hardware-specific image pins record prior model validation;
+  they need an updated build to use `--tokenizer-mode glm5_next`.
+
+  ## Frontends
+
+  The command builder defaults to **Rust** and adds `--tokenizer-mode glm5_next`.
+  With `VLLM_USE_RUST_FRONTEND=1`, this selects native GLM5Next chat rendering for
+  GLM-5.3, including tool calls and reasoning options.
+
+  Select **Python** in the command builder, or set `VLLM_USE_RUST_FRONTEND=0`,
+  to use the Hugging Face renderer on the same build. On older builds that do not
+  recognize `glm5_next`, also replace `--tokenizer-mode glm5_next` with
+  `--tokenizer-mode auto`.
 
   ## Launching the server
 
   ### FP8 on 8xH200 (standard)
 
   ```bash
-  vllm serve zai-org/GLM-5.3 \
+  VLLM_USE_RUST_FRONTEND=1 vllm serve zai-org/GLM-5.3 \
+    --tokenizer-mode glm5_next \
     --kv-cache-dtype fp8 \
     --tensor-parallel-size 8 \
     --speculative-config.method mtp \
@@ -296,7 +256,9 @@ guide: |
   ```bash
   VLLM_ROCM_USE_AITER=1 \
   VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=1 \
+  VLLM_USE_RUST_FRONTEND=1 \
   vllm serve zai-org/GLM-5.3 \
+    --tokenizer-mode glm5_next \
     --kv-cache-dtype fp8_e4m3 \
     --tensor-parallel-size 8 \
     --speculative-config.method mtp \
@@ -311,7 +273,7 @@ guide: |
     --moe-backend aiter
   ```
 
-  On 8x MI325X with `vllm/vllm-openai-rocm:v0.29.0`, this command serves the full
+  Earlier validation on 8x MI325X with `vllm/vllm-openai-rocm:v0.29.0` served
   `--max-model-len 524288` at `--gpu-memory-utilization 0.80`, reporting a 1,477,168-token
   KV cache and 2.82x maximum concurrency. MTP is healthy there: a mean acceptance length of
   2.77 with 5 tokens drafted per step on a ~500-token greedy generation.
@@ -328,7 +290,8 @@ guide: |
   (`--kv-cache-dtype fp8_e4m3`, already in the base flags) roughly halves that budget.
 
   ```bash
-  vllm serve zai-org/GLM-5.3 \
+  VLLM_USE_RUST_FRONTEND=1 vllm serve zai-org/GLM-5.3 \
+    --tokenizer-mode glm5_next \
     --kv-cache-dtype fp8_e4m3 \
     --tensor-parallel-size 8 \
     --speculative-config.method mtp \
@@ -351,7 +314,8 @@ guide: |
   Select the **NVFP4** variant above (Blackwell-only) or run:
 
   ```bash
-  vllm serve Inferact/GLM-5.3-NVFP4 \
+  VLLM_USE_RUST_FRONTEND=1 vllm serve Inferact/GLM-5.3-NVFP4 \
+    --tokenizer-mode glm5_next \
     --tensor-parallel-size 8 \
     --enable-expert-parallel \
     --reasoning-parser glm47 \
@@ -368,7 +332,8 @@ guide: |
   with eight speculative tokens:
 
   ```bash
-  vllm serve zai-org/GLM-5.3 \
+  VLLM_USE_RUST_FRONTEND=1 vllm serve zai-org/GLM-5.3 \
+    --tokenizer-mode glm5_next \
     --tensor-parallel-size 8 \
     --speculative-config '{"method":"dspark","model":"RedHatAI/GLM-5.3-speculator.dspark","num_speculative_tokens":8}'
   ```
diff --git a/requirements/build/rock.txt b/requirements/build/rock.txt
index 5ca5ada2af06..932057f11737 100644
--- a/requirements/build/rock.txt
+++ b/requirements/build/rock.txt
@@ -15,5 +15,5 @@ setuptools-rust>=1.9.0
 wheel
 jinja2>=3.1.6
 amdsmi==7.0.2
-timm>=1.0.17
+timm>=1.0.30
 tilelang==0.1.10
diff --git a/requirements/build/rocm.txt b/requirements/build/rocm.txt
deleted file mode 100644
index 3c9b6c179613..000000000000
--- a/requirements/build/rocm.txt
+++ /dev/null
@@ -1,19 +0,0 @@
-# Common dependencies
--r ../common.txt
-
-
---extra-index-url https://repo.amd.com/rocm/whl-multi-arch/
-torch==2.12.0+rocm7.14.0
-torchvision==0.27.0+rocm7.14.0
-torchaudio==2.11.0+rocm7.14.0
-triton==3.7.1+git0263a6a6
-cmake>=3.26.1,<4
-packaging>=26.2
-setuptools>=78.1.1,<80.0.0
-setuptools-scm>=8
-setuptools-rust>=1.9.0
-wheel
-jinja2>=3.1.6
-amdsmi==7.0.2
-timm>=1.0.17
-tilelang==0.1.10
diff --git a/requirements/common.txt b/requirements/common.txt
index d235d9e8bd81..5dab173e4c12 100644
--- a/requirements/common.txt
+++ b/requirements/common.txt
@@ -8,9 +8,9 @@ certifi >= 2026.10.19
 tqdm
 blake3
 py-cpuinfo
-transformers >= 5.16.1, < 5.21.0
-huggingface_hub >= 1.31.0  # Required for the `huggingface_hub.utils.httpx` re-export
-tokenizers >= 0.23.2  # Required for fast incremental detokenization.
+transformers >= 5.16.1, < 5.23.0
+huggingface_hub >= 2.0.0
+tokenizers >= 0.23.3  # Required for fast incremental detokenization and huggingface_hub 2.
 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.*; platform_machine != "s390x"
 protobuf == 7.36.1; platform_machine == "s390x"
@@ -36,7 +36,7 @@ partial-json-parser # used for parsing partial JSON outputs
 jsonschema >= 4.23.0 # required for MiniMax M3 tool schema validation
 pyzmq >= 25.0.0
 msgspec
-mistral_common[image] >= 1.11.6
+mistral_common[image] >= 1.13.0
 opencv-python-headless >= 4.13.0    # required for video IO
 pyyaml
 six>=1.16.0; python_version > '3.11' # transitive dependency of pandas that needs to be the latest version for python 3.12
diff --git a/requirements/kv_connectors.txt b/requirements/kv_connectors.txt
index baddf1cd633c..2c5593c3c6a0 100644
--- a/requirements/kv_connectors.txt
+++ b/requirements/kv_connectors.txt
@@ -2,7 +2,7 @@ lmcache >= 0.5.6
 # 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.5.1
+nixl == 1.5.2
 # 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
index 60bb194b32ec..e1c6b58af3c6 100644
--- a/requirements/kv_connectors_cu12.txt
+++ b/requirements/kv_connectors_cu12.txt
@@ -3,5 +3,5 @@
 lmcache == 0.5.6
 cuda-python >= 12, < 13
 cupy-cuda12x != 14.1.0
-nixl == 1.5.1
+nixl == 1.5.2
 mooncake-transfer-engine >= 0.3.12
diff --git a/requirements/rocm.txt b/requirements/rocm.txt
index cd7a11472157..a36e798ab072 100644
--- a/requirements/rocm.txt
+++ b/requirements/rocm.txt
@@ -16,9 +16,9 @@ packaging>=24.2
 setuptools>=78.1.1,<80.0.0
 setuptools-scm>=8
 setuptools-rust>=1.9.0
-runai-model-streamer[s3,gcs,azure]==0.15.7
+runai-model-streamer[s3,gcs,azure]==0.17.1
 conch-triton-kernels==1.2.1
-timm>=1.0.17
+timm>=1.0.30
 # amd-quark: required for Quark quantization on ROCm
 # To be consistent with test_quark.py
 amd-quark==0.12.post1
diff --git a/requirements/rubin-prerelease.txt b/requirements/rubin-prerelease.txt
index 720c08843159..157c8f421a3a 100644
--- a/requirements/rubin-prerelease.txt
+++ b/requirements/rubin-prerelease.txt
@@ -11,4 +11,4 @@ torchaudio==2.11.0.dev20261009+cu134
 
 cuda-python==13.4.1
 cuda-bindings==13.4.1
-nixl==1.4.1
+nixl==1.5.2
diff --git a/requirements/test/cpu.txt b/requirements/test/cpu.txt
index 1fe62a75d014..ee94974a116f 100644
--- a/requirements/test/cpu.txt
+++ b/requirements/test/cpu.txt
@@ -347,7 +347,6 @@ httpx==0.27.2
     #   cohere
     #   fastapi
     #   fastapi-cloud-cli
-    #   huggingface-hub
     #   model-hosting-container-standards
     #   openai
     #   perceptron
@@ -355,8 +354,9 @@ httpx==0.27.2
 httpx2==2.12.0
     # via
     #   -r requirements/test/cuda.in
+    #   huggingface-hub
     #   mcp
-huggingface-hub==1.31.0
+huggingface-hub==2.4.0
     # via
     #   -r requirements/test/../common.txt
     #   accelerate
@@ -508,7 +508,7 @@ mcp-types==2.3.0
     # via mcp
 mdurl==0.1.2
     # via markdown-it-py
-mistral-common==1.11.6
+mistral-common==1.13.0
     # via
     #   -r requirements/test/../common.txt
     #   -r requirements/test/cuda.in
@@ -1010,13 +1010,13 @@ rpds-py==0.20.1
     #   referencing
 rsa==4.9.1
     # via google-auth
-runai-model-streamer==0.15.7
+runai-model-streamer==0.17.1
     # via -r requirements/test/cuda.in
-runai-model-streamer-azure==0.15.7
+runai-model-streamer-azure==0.17.1
     # via runai-model-streamer
-runai-model-streamer-gcs==0.15.7
+runai-model-streamer-gcs==0.17.1
     # via runai-model-streamer
-runai-model-streamer-s3==0.15.7
+runai-model-streamer-s3==0.17.1
     # via runai-model-streamer
 s3transfer==0.10.3
     # via boto3
@@ -1159,12 +1159,12 @@ tiktoken==0.12.0
     #   gpt-oss
     #   lm-eval
     #   mistral-common
-timm==1.0.17
+timm==1.0.30
     # via
     #   -r requirements/test/cuda.in
     #   open-clip-torch
     #   segmentation-models-pytorch
-tokenizers==0.23.2
+tokenizers==0.23.3
     # via
     #   -r requirements/test/../common.txt
     #   cohere
@@ -1227,7 +1227,7 @@ traitlets==5.16.1
     # via
     #   jupyter-client
     #   jupyter-core
-transformers==5.20.2
+transformers==5.22.0
     # via
     #   -r requirements/test/../common.txt
     #   -r requirements/test/cuda.in
@@ -1286,7 +1286,6 @@ typing-extensions==4.16.0
     #   lm-eval
     #   mcp
     #   mcp-types
-    #   mistral-common
     #   mteb
     #   openai
     #   opentelemetry-api
diff --git a/requirements/test/cuda.in b/requirements/test/cuda.in
index 299fee0864b5..0c129abcad8e 100644
--- a/requirements/test/cuda.in
+++ b/requirements/test/cuda.in
@@ -27,20 +27,20 @@ sentence-transformers>=5.6.0 # required for embedding tests
 soundfile # required for audio tests
 jiwer # required for audio tests
 tblib # for pickling test exceptions
-timm >=1.0.17 # required for internvl and gemma3n-mm test
+timm >=1.0.30 # required for internvl and gemma3n-mm test
 torch==2.14.0
 torchaudio==2.11.0
 torchvision==0.29.0
 transformers_stream_generator # required for qwen-vl test
 matplotlib # required for qwen-vl test
-mistral_common[image,audio] >= 1.11.6 # required for voxtral test
+mistral_common[image,audio] >= 1.13.0 # required for voxtral test
 num2words # required for smolvlm test
 open_clip_torch==2.32.0 # Required for nemotron_vl test, Nemotron Parse in test_common.py
 datamodel_code_generator>=0.81.0 # required for minicpm3 test
 nltk>=3.10.3 # pin transitive rouge-score/lm-eval dep for known GHSAs
 lm-eval[api]>=0.4.12 # required for model evaluation test
 mteb[bm25s]>=2, <3 # required for mteb test
-transformers==5.20.2
+transformers==5.22.0
 schemathesis>=4.0.0 # Required for openai schema test.
 # quantization
 buildkite-test-collector==0.1.9
@@ -55,7 +55,7 @@ grpcio-reflection==1.78.0
 
 arctic-inference == 0.1.1; platform_machine == "x86_64" # Required for suffix decoding test
 numba == 0.65.0 # Required for N-gram speculative decoding
-runai-model-streamer[s3,gcs,azure]==0.15.7
+runai-model-streamer[s3,gcs,azure]==0.17.1
 fastsafetensors>=0.3.3
 instanttensor>=0.2.2; platform_machine == "x86_64"
 decord==0.6.0; platform_machine == "x86_64"
diff --git a/requirements/test/cuda.txt b/requirements/test/cuda.txt
index 58706a39d07e..49209ea1adab 100644
--- a/requirements/test/cuda.txt
+++ b/requirements/test/cuda.txt
@@ -367,7 +367,6 @@ httpx==0.27.2
     #   cohere
     #   fastapi
     #   fastapi-cloud-cli
-    #   huggingface-hub
     #   model-hosting-container-standards
     #   openai
     #   perceptron
@@ -375,8 +374,9 @@ httpx==0.27.2
 httpx2==2.12.0
     # via
     #   -r requirements/test/cuda.in
+    #   huggingface-hub
     #   mcp
-huggingface-hub==1.31.0
+huggingface-hub==2.4.0
     # via
     #   -c requirements/common.txt
     #   -r requirements/test/../common.txt
@@ -540,7 +540,7 @@ mcp-types==2.3.0
     # via mcp
 mdurl==0.1.2
     # via markdown-it-py
-mistral-common==1.11.6
+mistral-common==1.13.0
     # via
     #   -c requirements/common.txt
     #   -r requirements/test/../common.txt
@@ -1109,13 +1109,13 @@ rpds-py==0.20.1
     #   referencing
 rsa==4.9.1
     # via google-auth
-runai-model-streamer==0.15.7
+runai-model-streamer==0.17.1
     # via -r requirements/test/cuda.in
-runai-model-streamer-azure==0.15.7
+runai-model-streamer-azure==0.17.1
     # via runai-model-streamer
-runai-model-streamer-gcs==0.15.7
+runai-model-streamer-gcs==0.17.1
     # via runai-model-streamer
-runai-model-streamer-s3==0.15.7
+runai-model-streamer-s3==0.17.1
     # via runai-model-streamer
 s3transfer==0.10.3
     # via boto3
@@ -1263,12 +1263,12 @@ tiktoken==0.12.0
     #   gpt-oss
     #   lm-eval
     #   mistral-common
-timm==1.0.17
+timm==1.0.30
     # via
     #   -r requirements/test/cuda.in
     #   open-clip-torch
     #   segmentation-models-pytorch
-tokenizers==0.23.2
+tokenizers==0.23.3
     # via
     #   -c requirements/common.txt
     #   -r requirements/test/../common.txt
@@ -1337,7 +1337,7 @@ traitlets==5.16.1
     # via
     #   jupyter-client
     #   jupyter-core
-transformers==5.20.2
+transformers==5.22.0
     # via
     #   -c requirements/common.txt
     #   -r requirements/test/../common.txt
@@ -1400,7 +1400,6 @@ typing-extensions==4.16.0
     #   lm-eval
     #   mcp
     #   mcp-types
-    #   mistral-common
     #   mteb
     #   openai
     #   opentelemetry-api
diff --git a/requirements/test/nightly-torch.txt b/requirements/test/nightly-torch.txt
index ed4ad5f011a9..708d8097ae84 100644
--- a/requirements/test/nightly-torch.txt
+++ b/requirements/test/nightly-torch.txt
@@ -23,13 +23,13 @@ jiwer # required for audio tests
 timm # required for internvl test
 transformers_stream_generator # required for qwen-vl test
 matplotlib # required for qwen-vl test
-mistral_common[image,audio] >= 1.11.6 # required for voxtral test
+mistral_common[image,audio] >= 1.13.0 # required for voxtral test
 num2words # required for smolvlm test
 opencv-python-headless >= 4.13.0 # required for video test
 datamodel_code_generator>=0.81.0 # 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.20.2
+transformers==5.22.0
 schemathesis>=4.0.0 # Required for openai schema test.
 # quantization
 buildkite-test-collector==0.1.9
@@ -40,7 +40,7 @@ tritonclient>=2.51.0
 
 numba == 0.65.0 # Required for N-gram speculative decoding
 numpy
-runai-model-streamer[s3,gcs,azure]==0.15.7
+runai-model-streamer[s3,gcs,azure]==0.17.1
 fastsafetensors>=0.3.3
 instanttensor>=0.2.2
 pydantic>=2.12 # 2.11 leads to error on python 3.13
diff --git a/requirements/test/rocm.in b/requirements/test/rocm.in
index 8d090b29b3c5..8e279eba7218 100644
--- a/requirements/test/rocm.in
+++ b/requirements/test/rocm.in
@@ -26,17 +26,17 @@ sentence-transformers>=5.6.0 # required for embedding tests
 soundfile # required for audio tests
 jiwer # required for audio tests
 tblib # for pickling test exceptions
-timm>=1.0.17 # required for internvl and gemma3n-mm test
+timm>=1.0.30 # required for internvl and gemma3n-mm test
 transformers_stream_generator # required for qwen-vl test
 matplotlib # required for qwen-vl test
-mistral_common[image,audio]>=1.11.6 # required for voxtral test
+mistral_common[image,audio]>=1.13.0 # required for voxtral test
 num2words # required for smolvlm test
 open_clip_torch==2.32.0 # Required for nemotron_vl test, Nemotron Parse in test_common.py
 datamodel_code_generator>=0.81.0 # required for minicpm3 test
 nltk>=3.10.3 # pin transitive rouge-score/lm-eval dep for known GHSAs
 lm-eval[api]>=0.4.12 # required for model evaluation test
 mteb[bm25s]>=2, <3 # required for mteb test
-transformers==5.20.2
+transformers==5.22.0
 schemathesis>=4.0.0 # Required for openai schema test
 # quantization
 bitsandbytes==0.50.0
@@ -52,7 +52,7 @@ grpcio-reflection==1.78.0
 
 arctic-inference==0.1.1 # Required for suffix decoding test
 numba==0.65.0 # Required for N-gram speculative decoding
-runai-model-streamer[s3,gcs,azure]==0.15.7
+runai-model-streamer[s3,gcs,azure]==0.17.1
 fastsafetensors>=0.3.3
 instanttensor>=0.1.9
 decord==0.6.0
diff --git a/requirements/test/rocm.txt b/requirements/test/rocm.txt
index ba87f1c3a336..d9db92d4bb6a 100644
--- a/requirements/test/rocm.txt
+++ b/requirements/test/rocm.txt
@@ -1,5 +1,5 @@
 # This file was autogenerated by uv via the following command:
-#    uv pip compile requirements/test/rocm.in -c requirements/rocm.txt -o requirements/test/rocm.txt --index-strategy unsafe-best-match --python-platform x86_64-manylinux_2_28 --python-version 3.12 --no-emit-package torch --no-emit-package torchvision --no-emit-package torchaudio --no-emit-package triton --no-emit-package cuda-bindings --no-emit-package cuda-pathfinder --no-emit-package cuda-toolkit --no-emit-package cupy-cuda12x --no-emit-package nvidia-cublas --no-emit-package nvidia-cuda-cupti --no-emit-package nvidia-cuda-nvrtc --no-emit-package nvidia-cuda-runtime --no-emit-package nvidia-cudnn --no-emit-package nvidia-cufft --no-emit-package nvidia-cufile --no-emit-package nvidia-curand --no-emit-package nvidia-cusolver --no-emit-package nvidia-cusparse --no-emit-package nvidia-cusparselt --no-emit-package nvidia-nccl --no-emit-package nvidia-nvjitlink --no-emit-package nvidia-nvshmem --no-emit-package nvidia-nvtx --no-emit-package nvidia-cublas-cu12 --no-emit-package nvidia-cuda-cupti-cu12 --no-emit-package nvidia-cuda-nvrtc-cu12 --no-emit-package nvidia-cuda-runtime-cu12 --no-emit-package nvidia-cudnn-cu12 --no-emit-package nvidia-cufft-cu12 --no-emit-package nvidia-cufile-cu12 --no-emit-package nvidia-curand-cu12 --no-emit-package nvidia-cusolver-cu12 --no-emit-package nvidia-cusparse-cu12 --no-emit-package nvidia-cusparselt-cu12 --no-emit-package nvidia-nccl-cu12 --no-emit-package nvidia-nvjitlink-cu12 --no-emit-package nvidia-nvshmem-cu12 --no-emit-package nvidia-nvtx-cu12 --no-emit-package nvidia-cublas-cu13 --no-emit-package nvidia-cuda-cupti-cu13 --no-emit-package nvidia-cuda-nvrtc-cu13 --no-emit-package nvidia-cuda-runtime-cu13 --no-emit-package nvidia-cudnn-cu13 --no-emit-package nvidia-cufft-cu13 --no-emit-package nvidia-cufile-cu13 --no-emit-package nvidia-curand-cu13 --no-emit-package nvidia-cusolver-cu13 --no-emit-package nvidia-cusparse-cu13 --no-emit-package nvidia-cusparselt-cu13 --no-emit-package nvidia-nccl-cu13 --no-emit-package nvidia-nvjitlink-cu13 --no-emit-package nvidia-nvshmem-cu13 --no-emit-package nvidia-nvtx-cu13
+#    uv pip compile requirements/test/rocm.in -c requirements/rocm.txt -o requirements/test/rocm.txt --index-strategy unsafe-best-match --python-platform x86_64-manylinux_2_28 --python-version 3.12 --no-emit-package torch --no-emit-package torchvision --no-emit-package torchaudio --no-emit-package triton --no-emit-package cuda-bindings --no-emit-package cuda-pathfinder --no-emit-package cuda-toolkit --no-emit-package cupy-cuda12x --no-emit-package nvidia-cublas --no-emit-package nvidia-cuda-cccl --no-emit-package nvidia-cuda-cupti --no-emit-package nvidia-cuda-nvrtc --no-emit-package nvidia-cuda-runtime --no-emit-package nvidia-cudnn --no-emit-package nvidia-cufft --no-emit-package nvidia-cufile --no-emit-package nvidia-curand --no-emit-package nvidia-cusolver --no-emit-package nvidia-cusparse --no-emit-package nvidia-cusparselt --no-emit-package nvidia-nccl --no-emit-package nvidia-nvjitlink --no-emit-package nvidia-nvshmem --no-emit-package nvidia-nvtx --no-emit-package nvidia-cublas-cu12 --no-emit-package nvidia-cuda-cupti-cu12 --no-emit-package nvidia-cuda-nvrtc-cu12 --no-emit-package nvidia-cuda-runtime-cu12 --no-emit-package nvidia-cudnn-cu12 --no-emit-package nvidia-cufft-cu12 --no-emit-package nvidia-cufile-cu12 --no-emit-package nvidia-curand-cu12 --no-emit-package nvidia-cusolver-cu12 --no-emit-package nvidia-cusparse-cu12 --no-emit-package nvidia-cusparselt-cu12 --no-emit-package nvidia-nccl-cu12 --no-emit-package nvidia-nvjitlink-cu12 --no-emit-package nvidia-nvshmem-cu12 --no-emit-package nvidia-nvtx-cu12 --no-emit-package nvidia-cublas-cu13 --no-emit-package nvidia-cuda-cupti-cu13 --no-emit-package nvidia-cuda-nvrtc-cu13 --no-emit-package nvidia-cuda-runtime-cu13 --no-emit-package nvidia-cudnn-cu13 --no-emit-package nvidia-cufft-cu13 --no-emit-package nvidia-cufile-cu13 --no-emit-package nvidia-curand-cu13 --no-emit-package nvidia-cusolver-cu13 --no-emit-package nvidia-cusparse-cu13 --no-emit-package nvidia-cusparselt-cu13 --no-emit-package nvidia-nccl-cu13 --no-emit-package nvidia-nvjitlink-cu13 --no-emit-package nvidia-nvshmem-cu13 --no-emit-package nvidia-nvtx-cu13
 absl-py==2.4.0
     # via rouge-score
 accelerate==1.13.0
@@ -363,7 +363,6 @@ httpx==0.27.2
     #   cohere
     #   fastapi
     #   fastapi-cloud-cli
-    #   huggingface-hub
     #   model-hosting-container-standards
     #   openai
     #   perceptron
@@ -371,8 +370,9 @@ httpx==0.27.2
 httpx2==2.12.0
     # via
     #   -r requirements/test/rocm.in
+    #   huggingface-hub
     #   mcp
-huggingface-hub==1.31.0
+huggingface-hub==2.4.0
     # via
     #   -c requirements/common.txt
     #   -r requirements/test/../common.txt
@@ -538,7 +538,7 @@ mcp-types==2.3.0
     # via mcp
 mdurl==0.1.2
     # via markdown-it-py
-mistral-common==1.11.6
+mistral-common==1.13.0
     # via
     #   -c requirements/common.txt
     #   -r requirements/test/../common.txt
@@ -1083,15 +1083,15 @@ rpds-py==0.30.0
     # via
     #   jsonschema
     #   referencing
-runai-model-streamer==0.15.7
+runai-model-streamer==0.17.1
     # via
     #   -c requirements/rocm.txt
     #   -r requirements/test/rocm.in
-runai-model-streamer-azure==0.15.7
+runai-model-streamer-azure==0.17.1
     # via runai-model-streamer
-runai-model-streamer-gcs==0.15.7
+runai-model-streamer-gcs==0.17.1
     # via runai-model-streamer
-runai-model-streamer-s3==0.15.7
+runai-model-streamer-s3==0.17.1
     # via runai-model-streamer
 s3transfer==0.16.0
     # via boto3
@@ -1248,13 +1248,13 @@ tilelang==0.1.10
     # via
     #   -c requirements/rocm.txt
     #   -r requirements/test/rocm.in
-timm==1.0.17
+timm==1.0.30
     # via
     #   -c requirements/rocm.txt
     #   -r requirements/test/rocm.in
     #   open-clip-torch
     #   segmentation-models-pytorch
-tokenizers==0.23.2
+tokenizers==0.23.3
     # via
     #   -c requirements/common.txt
     #   -r requirements/test/../common.txt
@@ -1287,7 +1287,7 @@ traitlets==5.16.1
     # via
     #   jupyter-client
     #   jupyter-core
-transformers==5.20.2
+transformers==5.22.0
     # via
     #   -c requirements/common.txt
     #   -r requirements/test/../common.txt
@@ -1348,7 +1348,6 @@ typing-extensions==4.16.0
     #   lm-eval
     #   mcp
     #   mcp-types
-    #   mistral-common
     #   mteb
     #   openai
     #   opentelemetry-api
@@ -1446,12 +1445,12 @@ zipp==3.23.0
 # cuda-pathfinder
 # cuda-toolkit
 # nvidia-cublas
+# nvidia-cuda-cccl
 # nvidia-cuda-cupti
 # nvidia-cuda-nvrtc
 # nvidia-cuda-runtime
 # nvidia-cufft
 # nvidia-cufile
-# nvidia-curand
 # nvidia-cusolver
 # nvidia-cusparse
 # nvidia-nvjitlink
diff --git a/requirements/test/xpu.in b/requirements/test/xpu.in
index 333cfb132d0b..f16f1a5de031 100644
--- a/requirements/test/xpu.in
+++ b/requirements/test/xpu.in
@@ -19,9 +19,9 @@ arctic-inference
 lm_eval[api]>=0.4.12
 nltk>=3.10.3 # pin transitive rouge-score/lm-eval dep for known GHSAs
 modelscope<1.38
-runai-model-streamer[s3,gcs,azure]>=0.15.7
+runai-model-streamer[s3,gcs,azure]>=0.17.1
 tensorizer==2.10.1
-transformers==5.20.2
+transformers==5.22.0
 
 # --- Audio Processing ---
 av
@@ -46,12 +46,12 @@ num2words
 pqdm
 
 # --- Vision & Multimodal ---
-timm>=1.0.17 # required for internvl and gemma3n-mm test
+timm>=1.0.30 # required for internvl and gemma3n-mm test
 albumentations
 ftfy
 matplotlib
 decord==0.6.0; platform_machine == "x86_64"
-mistral-common[image,audio]>=1.11.6 # required for voxtral test
+mistral-common[image,audio]>=1.13.0 # required for voxtral test
 
 open_clip_torch==2.32.0 # Required for nemotron_vl test, Nemotron Parse in test_common.py
 
diff --git a/requirements/test/xpu.txt b/requirements/test/xpu.txt
index 3d026574702b..39c9ab930002 100644
--- a/requirements/test/xpu.txt
+++ b/requirements/test/xpu.txt
@@ -147,7 +147,7 @@ dataproperty==1.1.1
     # via
     #   pytablewriter
     #   tabledata
-datasets==5.0.1
+datasets==5.1.0
     # via
     #   evaluate
     #   lm-eval
@@ -302,17 +302,16 @@ httpx==0.28.1
     # via
     #   anthropic
     #   cohere
-    #   datasets
     #   fastapi
     #   fastapi-cloud-cli
-    #   huggingface-hub
     #   model-hosting-container-standards
 httpx2==2.13.1
     # via
     #   -r requirements/test/xpu.in
+    #   huggingface-hub
     #   mcp
     #   openai
-huggingface-hub==1.31.0
+huggingface-hub==2.4.0
     # via
     #   -c requirements/common.txt
     #   -r requirements/test/../common.txt
@@ -489,7 +488,7 @@ mcp-types==2.3.0
     # via mcp
 mdurl==0.1.2
     # via markdown-it-py
-mistral-common==1.11.7
+mistral-common==1.13.0
     # via
     #   -c requirements/common.txt
     #   -r requirements/test/../common.txt
@@ -930,7 +929,6 @@ requests==2.34.2
     #   -r requirements/test/../common.txt
     #   azure-core
     #   cohere
-    #   datasets
     #   docker
     #   evaluate
     #   google-api-core
@@ -964,13 +962,13 @@ rpds-py==2026.6.3
     # via
     #   jsonschema
     #   referencing
-runai-model-streamer==0.16.1
+runai-model-streamer==0.17.1
     # via -r requirements/test/xpu.in
-runai-model-streamer-azure==0.16.1
+runai-model-streamer-azure==0.17.1
     # via runai-model-streamer
-runai-model-streamer-gcs==0.16.1
+runai-model-streamer-gcs==0.17.1
     # via runai-model-streamer
-runai-model-streamer-s3==0.16.1
+runai-model-streamer-s3==0.17.1
     # via runai-model-streamer
 s3transfer==0.19.2
     # via boto3
@@ -1105,11 +1103,11 @@ tiktoken==0.14.0
     #   gpt-oss
     #   lm-eval
     #   mistral-common
-timm==1.0.28
+timm==1.0.30
     # via
     #   -r requirements/test/xpu.in
     #   open-clip-torch
-tokenizers==0.23.2
+tokenizers==0.23.3
     # via
     #   -c requirements/common.txt
     #   -r requirements/test/../common.txt
@@ -1156,7 +1154,7 @@ traitlets==5.16.1
     # via
     #   jupyter-client
     #   jupyter-core
-transformers==5.20.2
+transformers==5.22.0
     # via
     #   -c requirements/common.txt
     #   -r requirements/test/../common.txt
@@ -1212,7 +1210,6 @@ typing-extensions==4.16.0
     #   lm-eval
     #   mcp
     #   mcp-types
-    #   mistral-common
     #   mteb
     #   openai
     #   opentelemetry-api
diff --git a/requirements/xpu.txt b/requirements/xpu.txt
index 14372071f0b5..14a9cdcc4d1d 100644
--- a/requirements/xpu.txt
+++ b/requirements/xpu.txt
@@ -20,4 +20,5 @@ torchvision
 torchcodec >= 0.14 # Required for the torchcodec video decoding backend
 
 auto_round_lib==0.16.0
-vllm_xpu_kernels==0.1.15.4
+--extra-index-url=https://vllm-project.github.io/vllm-xpu-kernels/whl
+vllm_xpu_kernels==0.1.16
diff --git a/rust/Cargo.lock b/rust/Cargo.lock
index a722d362cb52..da33ca59188a 100644
--- a/rust/Cargo.lock
+++ b/rust/Cargo.lock
@@ -5726,6 +5726,7 @@ dependencies = [
  "serde_repr",
  "serde_tuple",
  "serde_with",
+ "sha2",
  "socket2",
  "task-local",
  "tempfile",
diff --git a/rust/README.md b/rust/README.md
index 42d6c3e16a05..6f94d24d859d 100644
--- a/rust/README.md
+++ b/rust/README.md
@@ -64,6 +64,12 @@ SMG `VllmEngine` gRPC server in place of HTTP. The Rust listener exposes the
 `vllm.Inference` and `vllm.Control` services alongside HTTP. These are different
 gRPC APIs; `--grpc` and `--grpc-port` are mutually exclusive.
 
+### gRPC health checks
+
+The gRPC listener also serves `grpc.health.v1.Health`. Kubernetes probes should
+use the `liveness` service for liveness and `readiness` for readiness; see
+[Rust frontend gRPC health checks](../docs/deployment/k8s.md#rust-frontend-grpc-health-checks).
+
 ### RL weight synchronization
 
 With `VLLM_SERVER_DEV_MODE=1`, the Rust frontend supports the HTTP weight-transfer
@@ -107,6 +113,23 @@ To build the `vllm-rs` in isolation:
 ./tools/build_rust.sh
 ```
 
+### Hybrid and External DP Load Balancing
+
+Like `vllm serve` (see [Data Parallel Deployment](../docs/serving/data_parallel_deployment.md)), `vllm-rs serve` can run one Rust frontend per node that serves only the node's data parallel ranks, with an upstream load balancer spreading requests across nodes. Every node passes the global `--data-parallel-size`, the address of the node that owns rank 0 as `--data-parallel-address`, and the same explicit `--data-parallel-rpc-port`, because all engines join the handshake of the rank-0 frontend. For MoE models, the rank-0 frontend also runs the DP coordinator.
+
+```bash
+# Node 0: ranks 0-1
+vllm-rs serve <moe-model> \
+  --data-parallel-size 4 --data-parallel-size-local 2 --data-parallel-hybrid-lb \
+  --data-parallel-address 10.0.0.1 --data-parallel-rpc-port 13345
+# Node 1: ranks 2-3
+vllm-rs serve <moe-model> \
+  --data-parallel-size 4 --data-parallel-size-local 2 --data-parallel-start-rank 2 \
+  --data-parallel-address 10.0.0.1 --data-parallel-rpc-port 13345
+```
+
+For external load balancing, start one `vllm-rs serve` per rank with `--data-parallel-rank <rank>`. A request can be pinned to one of a frontend's ranks with the `X-data-parallel-rank` header, which takes the global rank. Unlike `vllm serve`, `--data-parallel-rpc-port` has no default here and ranks are not inferred from `--nnodes`/`--node-rank`.
+
 ### Engine-free renderer
 
 `vllm-rs render` serves text-only request preprocessing without starting or
diff --git a/rust/src/bench/src/backends/mod.rs b/rust/src/bench/src/backends/mod.rs
index b5400c5d0f49..da23eb46a0ab 100644
--- a/rust/src/bench/src/backends/mod.rs
+++ b/rust/src/bench/src/backends/mod.rs
@@ -52,6 +52,21 @@ pub struct ChatDelta {
 #[derive(Deserialize)]
 pub struct ChunkUsage {
     pub completion_tokens: Option<u64>,
+    pub prompt_tokens: Option<u64>,
+}
+
+impl ChunkUsage {
+    /// Record the server-reported token counts on `output`. A zero
+    /// `prompt_tokens` keeps the client estimate: some OpenAI-compatible
+    /// servers send placeholder zeros in usage chunks.
+    pub fn apply(&self, output: &mut RequestFuncOutput) {
+        if let Some(ct) = self.completion_tokens {
+            output.output_tokens = ct as usize;
+        }
+        if let Some(pt) = self.prompt_tokens.filter(|&pt| pt > 0) {
+            output.prompt_len = pt as usize;
+        }
+    }
 }
 
 use crate::cli::BackendKind;
@@ -246,6 +261,21 @@ pub fn build_headers(
 mod tests {
     use super::*;
 
+    /// A placeholder zero `prompt_tokens` must not replace the client estimate.
+    #[test]
+    fn chunk_usage_ignores_zero_prompt_tokens() {
+        let mut output = RequestFuncOutput {
+            prompt_len: 100,
+            ..Default::default()
+        };
+        ChunkUsage {
+            completion_tokens: Some(3),
+            prompt_tokens: Some(0),
+        }
+        .apply(&mut output);
+        assert_eq!((output.prompt_len, output.output_tokens), (100, 3));
+    }
+
     #[test]
     fn test_update_server_metrics_converts_ms_to_seconds() {
         let chunk: CompletionChunk = serde_json::from_value(serde_json::json!({
diff --git a/rust/src/bench/src/backends/openai_chat.rs b/rust/src/bench/src/backends/openai_chat.rs
index 30c6093296d0..36ce961abf81 100644
--- a/rust/src/bench/src/backends/openai_chat.rs
+++ b/rust/src/bench/src/backends/openai_chat.rs
@@ -130,10 +130,8 @@ impl OpenAIChatBackend {
                             }
                             // Separate `if` (not `else if`) — Dynamo may send
                             // both choices and usage in the same chunk.
-                            if let Some(ref usage) = data.usage
-                                && let Some(ct) = usage.completion_tokens
-                            {
-                                output.output_tokens = ct as usize;
+                            if let Some(usage) = &data.usage {
+                                usage.apply(&mut output);
                             }
                         }
                     }
@@ -363,7 +361,11 @@ mod tests {
     /// Serve one streaming chat response over a raw socket: `num_tokens`
     /// content chunks back-to-back, then `usage_delay` of silence, then the
     /// choice-less usage chunk and `[DONE]`. Returns the endpoint URL.
-    fn spawn_sse_server(num_tokens: usize, usage_delay: std::time::Duration) -> String {
+    fn spawn_sse_server(
+        num_tokens: usize,
+        usage_delay: std::time::Duration,
+        prompt_tokens: Option<u64>,
+    ) -> String {
         use std::io::{BufRead, BufReader, Read, Write};
 
         let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
@@ -393,8 +395,10 @@ mod tests {
                 send("data: {\"choices\":[{\"delta\":{\"content\":\"a\"}}]}\n\n");
             }
             std::thread::sleep(usage_delay);
+            let prompt_tokens =
+                prompt_tokens.map_or(String::new(), |pt| format!(",\"prompt_tokens\":{pt}"));
             send(&format!(
-                "data: {{\"choices\":[],\"usage\":{{\"completion_tokens\":{num_tokens}}}}}\n\n"
+                "data: {{\"choices\":[],\"usage\":{{\"completion_tokens\":{num_tokens}{prompt_tokens}}}}}\n\n"
             ));
             send("data: [DONE]\n\n");
         });
@@ -411,7 +415,7 @@ mod tests {
         // real delay is needed here; there is no observable event to wait on.
         let usage_delay = std::time::Duration::from_millis(150);
         let input = RequestFuncInput {
-            api_url: spawn_sse_server(3, usage_delay),
+            api_url: spawn_sse_server(3, usage_delay, None),
             model: "test-model".to_string(),
             output_len: 3,
             ..Default::default()
@@ -442,4 +446,42 @@ mod tests {
         // keys already set above must not be overridden by extra_body
         assert_eq!(v["stream"], true);
     }
+
+    /// The server's `usage.prompt_tokens` replaces the client-side estimate,
+    /// so input-token totals count what the server actually processed
+    /// (chat template, multimodal tokens). Mirrors Python #38654.
+    #[tokio::test]
+    async fn test_usage_prompt_tokens_override_prompt_len() {
+        let input = RequestFuncInput {
+            api_url: spawn_sse_server(3, std::time::Duration::ZERO, Some(128)),
+            model: "test-model".to_string(),
+            prompt_len: 100,
+            output_len: 3,
+            ..Default::default()
+        };
+
+        let output = OpenAIChatBackend.send_request(&input, &reqwest::Client::new()).await.unwrap();
+
+        assert!(output.success, "{}", output.error);
+        assert_eq!(output.prompt_len, 128);
+    }
+
+    /// Servers that omit `usage.prompt_tokens` keep the client-side estimate,
+    /// and `completion_tokens` is still recorded.
+    #[tokio::test]
+    async fn test_missing_prompt_tokens_keeps_client_prompt_len() {
+        let input = RequestFuncInput {
+            api_url: spawn_sse_server(3, std::time::Duration::ZERO, None),
+            model: "test-model".to_string(),
+            prompt_len: 100,
+            output_len: 3,
+            ..Default::default()
+        };
+
+        let output = OpenAIChatBackend.send_request(&input, &reqwest::Client::new()).await.unwrap();
+
+        assert!(output.success, "{}", output.error);
+        assert_eq!(output.prompt_len, 100);
+        assert_eq!(output.output_tokens, 3);
+    }
 }
diff --git a/rust/src/bench/src/backends/openai_completions.rs b/rust/src/bench/src/backends/openai_completions.rs
index 0b62d342bcde..049988191318 100644
--- a/rust/src/bench/src/backends/openai_completions.rs
+++ b/rust/src/bench/src/backends/openai_completions.rs
@@ -156,10 +156,8 @@ impl OpenAICompletionsBackend {
                             }
                             // Separate `if` (not `else if`) — Dynamo may send
                             // both choices and usage in the same chunk.
-                            if let Some(ref usage) = data.usage
-                                && let Some(ct) = usage.completion_tokens
-                            {
-                                output.output_tokens = ct as usize;
+                            if let Some(usage) = &data.usage {
+                                usage.apply(&mut output);
                             }
                         }
                     }
@@ -201,3 +199,69 @@ impl OpenAICompletionsBackend {
         Ok(output)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    /// Serve `body` as one streaming completions response over a raw socket.
+    /// Returns the endpoint URL.
+    fn spawn_sse_server(body: &'static str) -> String {
+        use std::io::{BufRead, BufReader, Read, Write};
+
+        let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
+        let addr = listener.local_addr().unwrap();
+        std::thread::spawn(move || {
+            let (mut stream, _) = listener.accept().unwrap();
+            let mut reader = BufReader::new(stream.try_clone().unwrap());
+            let mut content_length = 0;
+            loop {
+                let mut line = String::new();
+                reader.read_line(&mut line).unwrap();
+                if let Some(v) = line.to_ascii_lowercase().strip_prefix("content-length:") {
+                    content_length = v.trim().parse().unwrap();
+                }
+                if line == "\r\n" {
+                    break;
+                }
+            }
+            reader.read_exact(&mut vec![0; content_length]).unwrap();
+            stream
+                .write_all(
+                    format!(
+                        "HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\n\
+                         connection: close\r\n\r\n{body}"
+                    )
+                    .as_bytes(),
+                )
+                .unwrap();
+        });
+        format!("http://{addr}/v1/completions")
+    }
+
+    /// The server's `usage.prompt_tokens` replaces the client-side estimate,
+    /// as in the chat backend.
+    #[tokio::test]
+    async fn test_usage_prompt_tokens_override_prompt_len() {
+        let input = RequestFuncInput {
+            api_url: spawn_sse_server(concat!(
+                "data: {\"choices\":[{\"text\":\"a\"}]}\n\n",
+                "data: {\"choices\":[{\"text\":\"b\"}]}\n\n",
+                "data: {\"choices\":[],\"usage\":{\"completion_tokens\":2,\"prompt_tokens\":128}}\n\n",
+                "data: [DONE]\n\n",
+            )),
+            model: "test-model".to_string(),
+            prompt_len: 100,
+            output_len: 2,
+            ..Default::default()
+        };
+
+        let output = OpenAICompletionsBackend
+            .send_request(&input, &reqwest::Client::new())
+            .await
+            .unwrap();
+
+        assert!(output.success, "{}", output.error);
+        assert_eq!((output.prompt_len, output.output_tokens), (128, 2));
+    }
+}
diff --git a/rust/src/bench/src/backends/openai_responses.rs b/rust/src/bench/src/backends/openai_responses.rs
index 2a9982636137..34d519e3ef95 100644
--- a/rust/src/bench/src/backends/openai_responses.rs
+++ b/rust/src/bench/src/backends/openai_responses.rs
@@ -27,6 +27,7 @@ struct ResponsesEventResponse {
 
 #[derive(serde::Deserialize)]
 struct ResponsesUsage {
+    input_tokens: Option<u64>,
     output_tokens: Option<u64>,
 }
 
@@ -156,9 +157,17 @@ impl ResponsesStream {
         };
 
         match event.event_type.as_deref() {
-            // Reasoning deltas count toward TTFT and ITL because the server is
-            // already decoding. Only output text is collected.
-            Some(kind @ ("response.output_text.delta" | "response.reasoning_text.delta")) => {
+            // Reasoning and tool call deltas count toward TTFT and ITL because
+            // the server is already decoding; a response that only calls a tool
+            // emits nothing else. Only output text is collected.
+            Some(
+                kind @ ("response.output_text.delta"
+                | "response.reasoning_text.delta"
+                | "response.function_call_arguments.delta"
+                | "response.custom_tool_call_input.delta"
+                | "response.mcp_call_arguments.delta"
+                | "response.code_interpreter_call_code.delta"),
+            ) => {
                 if self.first_token_received {
                     self.output
                         .itl
@@ -180,8 +189,14 @@ impl ResponsesStream {
             // with `status: "incomplete"`; the spec also has `response.incomplete`.
             Some("response.completed" | "response.incomplete") => {
                 self.saw_terminal_event = true;
-                if let Some(n) = event.response.and_then(|r| r.usage?.output_tokens) {
-                    self.output.output_tokens = n as usize;
+                if let Some(usage) = event.response.and_then(|r| r.usage) {
+                    if let Some(n) = usage.output_tokens {
+                        self.output.output_tokens = n as usize;
+                    }
+                    // Same rule as `ChunkUsage::apply`: a zero keeps the client estimate.
+                    if let Some(n) = usage.input_tokens.filter(|&n| n > 0) {
+                        self.output.prompt_len = n as usize;
+                    }
                 }
             }
             kind if matches!(kind, Some("error" | "response.failed")) || event.error.is_some() => {
@@ -253,6 +268,21 @@ mod tests {
         stream.finish()
     }
 
+    /// The terminal event's `usage.input_tokens` replaces the client estimate,
+    /// as in Python's `async_request_openai_responses`.
+    #[test]
+    fn test_terminal_usage_input_tokens_override_prompt_len() {
+        let output = run(&[
+            text_delta("Hi"),
+            event(
+                "response.completed",
+                ",\"response\":{\"usage\":{\"input_tokens\":128,\"output_tokens\":1}}",
+            ),
+        ]);
+        assert!(output.success, "{}", output.error);
+        assert_eq!((output.prompt_len, output.output_tokens), (128, 1));
+    }
+
     #[test]
     fn test_reasoning_starts_ttft_and_latency_ends_at_last_token() {
         let output = run(&[
@@ -271,6 +301,48 @@ mod tests {
         assert_eq!(output.output_tokens, 3);
     }
 
+    /// A forced tool call streams only argument deltas. They are generated
+    /// tokens, as `tool_calls` deltas are on `openai-chat`.
+    #[test]
+    fn test_tool_call_only_response_is_timed_and_counted() {
+        let args = |delta: &str| {
+            event(
+                "response.function_call_arguments.delta",
+                &format!(",\"delta\":\"{delta}\""),
+            )
+        };
+        let output = run(&[
+            event("response.output_item.added", ""),
+            args("{\\\"city\\\": "),
+            args("\\\"Paris\\\"}"),
+            event("response.function_call_arguments.done", ""),
+            completed(6),
+        ]);
+
+        assert!(output.success, "{}", output.error);
+        assert_eq!(output.ttft, 0.02);
+        assert_eq!(output.itl, vec![0.01]);
+        assert_eq!(output.latency, 0.03);
+        // Tool arguments are not output text.
+        assert_eq!(output.generated_text, "");
+        assert_eq!(output.output_tokens, 6);
+    }
+
+    #[test]
+    fn test_every_tool_delta_event_starts_the_clock() {
+        for kind in [
+            "response.function_call_arguments.delta",
+            "response.custom_tool_call_input.delta",
+            "response.mcp_call_arguments.delta",
+            "response.code_interpreter_call_code.delta",
+        ] {
+            let delta = event(kind, ",\"delta\":\"x\"");
+            let output = run(&[delta.clone(), delta, completed(2)]);
+            assert!(output.success, "{kind}: {}", output.error);
+            assert_eq!(output.itl.len(), 1, "{kind}");
+        }
+    }
+
     #[test]
     fn test_event_line_split_from_data_line() {
         let delta = text_delta("Hi");
diff --git a/rust/src/bench/src/benchmark.rs b/rust/src/bench/src/benchmark.rs
index 5833a3bdcea5..1f6b0567ac2f 100644
--- a/rust/src/bench/src/benchmark.rs
+++ b/rust/src/bench/src/benchmark.rs
@@ -645,7 +645,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
         let total_output_tokens: usize = input_requests.iter().map(|r| r.expected_output_len).sum();
         println!("Dry run stats:");
         println!("  Total prompts: {}", input_requests.len());
-        println!("  Total input tokens: {total_input_tokens}");
+        println!("  Estimated total input tokens: {total_input_tokens}");
         println!("  Total expected output tokens: {total_output_tokens}");
         println!(
             "  Avg input tokens: {:.1}",
@@ -1118,7 +1118,6 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
         (m, Vec::new())
     } else {
         calculate_metrics(
-            &input_requests,
             &outputs,
             benchmark_duration,
             &config.selected_percentiles,
@@ -1148,7 +1147,6 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
         ) {
             let ss = steady_state::compute(
                 &outputs,
-                &input_requests,
                 &window,
                 &config.selected_percentiles,
                 config.backend.is_pooling(),
@@ -1392,7 +1390,7 @@ pub(crate) async fn send_profile_request(
     Ok(resp.status().is_success())
 }
 
-/// Parse `vllm:num_requests_running` from the server's `/metrics` Prometheus endpoint.
+/// Fetch total `vllm:num_requests_running` across engines from `/metrics`.
 pub(crate) async fn fetch_num_requests_running(
     client: &reqwest::Client,
     base_url: &str,
@@ -1400,18 +1398,35 @@ pub(crate) async fn fetch_num_requests_running(
     let url = format!("{base_url}/metrics");
     let resp = client.get(&url).send().await.ok()?;
     let text = resp.text().await.ok()?;
+
+    parse_num_requests_running(&text)
+}
+
+/// Parse total running requests across the engine series in `/metrics`.
+fn parse_num_requests_running(text: &str) -> Option<usize> {
+    let mut total: Option<usize> = None;
     for line in text.lines() {
-        if let Some(rest) = line.strip_prefix("vllm:num_requests_running") {
-            // Line format: `vllm:num_requests_running{model_name="..."} 42`
-            // The value is after the last space.
-            if let Some(val_str) = rest.rsplit(' ').next()
-                && let Ok(v) = val_str.parse::<f64>()
-            {
-                return Some(v as usize);
-            }
+        let Some(rest) = line.trim().strip_prefix("vllm:num_requests_running") else {
+            continue;
+        };
+        let sample = if rest.trim_start().starts_with('{') {
+            let Some((_, sample)) = rest.rsplit_once('}') else {
+                continue;
+            };
+            sample
+        } else if rest.starts_with(char::is_whitespace) {
+            rest
+        } else {
+            continue;
+        };
+        if let Some(value) = sample.split_whitespace().next().and_then(|s| s.parse::<f64>().ok())
+            && value.is_finite()
+            && value >= 0.0
+        {
+            total = Some(total.unwrap_or(0).saturating_add(value as usize));
         }
     }
-    None
+    total
 }
 
 /// Poll `/metrics` until `num_requests_running >= threshold`, then start the
@@ -1923,6 +1938,44 @@ async fn sample_verify_prompts(
 mod tests {
     use super::*;
 
+    #[test]
+    fn test_parse_num_requests_running_sums_dp_engines() {
+        let metrics: String = (0..8)
+            .map(|engine| {
+                format!("vllm:num_requests_running{{model_name=\"m\",engine=\"{engine}\"}} 16.0\n")
+            })
+            .collect();
+
+        assert_eq!(parse_num_requests_running(&metrics), Some(128));
+    }
+
+    #[test]
+    fn test_parse_num_requests_running_matches_exact_metric_name() {
+        let metrics = "\
+vllm:num_requests_running_total{engine=\"0\"} 1024
+vllm:num_requests_running{engine=\"0\"} 16
+vllm:num_requests_waiting{model_name=\"vllm:num_requests_running\"} 256
+vllm:num_requests_running{engine=\"1\"} 16
+";
+
+        assert_eq!(parse_num_requests_running(metrics), Some(32));
+    }
+
+    #[test]
+    fn test_parse_num_requests_running_distinguishes_zero_and_missing() {
+        for (metrics, expected) in [("vllm:num_requests_running 0\n", Some(0)), ("", None)] {
+            assert_eq!(parse_num_requests_running(metrics), expected);
+        }
+    }
+
+    #[test]
+    fn test_parse_num_requests_running_reads_sample_before_timestamp() {
+        let metrics =
+            "vllm:num_requests_running {model_name=\"model with spaces\"}\t4.2e1 123456\n";
+
+        assert_eq!(parse_num_requests_running(metrics), Some(42));
+    }
+
     fn make_metrics(
         drafts: u64,
         draft_tokens: u64,
diff --git a/rust/src/bench/src/cli.rs b/rust/src/bench/src/cli.rs
index 850fcb35ddfb..693f7f6585ba 100644
--- a/rust/src/bench/src/cli.rs
+++ b/rust/src/bench/src/cli.rs
@@ -406,10 +406,11 @@ pub struct BenchServeArgs {
     #[arg(long, default_value_t = false)]
     pub profile: bool,
 
-    /// Minimum server batch size (num_requests_running) before starting the
-    /// profiler. When set, profiling is deferred until the /metrics endpoint
-    /// reports at least this many running requests, then captures for
-    /// --profile-duration seconds. Requires --profile.
+    /// Minimum total running requests across engines reported by /metrics
+    /// before starting the profiler. When set, profiling is deferred until
+    /// this threshold is reached, then captures for --profile-duration seconds.
+    /// Only engines exposed by the /metrics endpoint at --base-url are counted.
+    /// Requires --profile.
     #[arg(long)]
     pub profile_batch_threshold: Option<usize>,
 
diff --git a/rust/src/bench/src/config.rs b/rust/src/bench/src/config.rs
index 3e180cfb3085..ed8f8f7d36cd 100644
--- a/rust/src/bench/src/config.rs
+++ b/rust/src/bench/src/config.rs
@@ -246,6 +246,17 @@ impl BenchConfig {
             ));
         }
 
+        if args.max_concurrency == Some(0) {
+            return Err(BenchError::Config(
+                "--max-concurrency must be greater than 0".into(),
+            ));
+        }
+        if args.multi_turn_concurrency == Some(0) {
+            return Err(BenchError::Config(
+                "--multi-turn-concurrency must be greater than 0".into(),
+            ));
+        }
+
         let base_url = args.resolve_base_url();
         let api_url = args.resolve_api_url();
 
@@ -1398,6 +1409,60 @@ mod tests {
 
         assert!(BenchConfig::from_args(&args).is_err());
     }
+
+    #[test]
+    fn test_zero_max_concurrency_fails() {
+        let args = parse_args(vec![
+            "vllm-bench",
+            "--model",
+            "test-model",
+            "--max-concurrency",
+            "0",
+        ]);
+        let err = BenchConfig::from_args(&args).expect_err("0 concurrency must be rejected");
+        assert!(
+            err.to_string().contains("--max-concurrency must be greater than 0"),
+            "{}",
+            err
+        );
+    }
+
+    #[test]
+    fn test_positive_max_concurrency_passes() {
+        for value in ["1", "16"] {
+            let args = parse_args(vec![
+                "vllm-bench",
+                "--model",
+                "test-model",
+                "--max-concurrency",
+                value,
+            ]);
+            assert!(
+                BenchConfig::from_args(&args).is_ok(),
+                "value {value} should pass"
+            );
+        }
+    }
+
+    #[test]
+    fn test_zero_multi_turn_concurrency_fails() {
+        let args = parse_args(vec![
+            "vllm-bench",
+            "--model",
+            "test-model",
+            "--multi-turn",
+            "--multi-turn-concurrency",
+            "0",
+        ]);
+        let err =
+            BenchConfig::from_args(&args).expect_err("0 multi-turn concurrency must be rejected");
+        assert!(
+            err.to_string().contains("--multi-turn-concurrency must be greater than 0"),
+            "{}",
+            err
+        );
+    }
+
     #[test]
     fn test_range_ratio_parse_float() {
         let rr = RangeRatio::parse("0.2").unwrap();
diff --git a/rust/src/bench/src/metrics/calculator.rs b/rust/src/bench/src/metrics/calculator.rs
index 7e3fc77565ce..b55ed1cbde6d 100644
--- a/rust/src/bench/src/metrics/calculator.rs
+++ b/rust/src/bench/src/metrics/calculator.rs
@@ -3,7 +3,6 @@
 
 use crate::backends::RequestFuncOutput;
 use crate::config::GoodputConfig;
-use crate::datasets::SampleRequest;
 use crate::metrics::{BenchmarkMetrics, MultiTurnMetrics};
 use crate::multi_turn::ConversationOutput;
 
@@ -27,7 +26,6 @@ fn log_failed_requests(outputs: &[RequestFuncOutput]) {
 ///
 /// Mirrors Python's `calculate_metrics()` from serve.py:392-599.
 pub fn calculate_metrics(
-    input_requests: &[SampleRequest],
     outputs: &[RequestFuncOutput],
     dur_s: f64,
     selected_percentiles: &[f64],
@@ -44,7 +42,7 @@ pub fn calculate_metrics(
     let mut ttfts: Vec<f64> = Vec::new();
     let mut e2els: Vec<f64> = Vec::new();
 
-    for (i, output) in outputs.iter().enumerate() {
+    for output in outputs.iter() {
         if output.success {
             let output_len = if output.output_tokens > 0 {
                 output.output_tokens
@@ -57,7 +55,7 @@ pub fn calculate_metrics(
             };
 
             actual_output_lens.push(output_len);
-            total_input += input_requests[i].prompt_len;
+            total_input += output.prompt_len;
 
             if output_len > 1 {
                 let latency_minus_ttft = output.latency - output.ttft;
@@ -283,7 +281,6 @@ fn percentiles_from_sorted(sorted: &[f64], percentiles: &[f64]) -> Vec<(f64, f64
 ///
 /// Mirrors Python's `calculate_metrics_for_embeddings()` from serve.py.
 /// Key differences from `calculate_metrics()`:
-/// - Uses `output.prompt_len` (server-reported) instead of `input_requests[i].prompt_len`
 /// - Only computes E2EL (no TTFT/TPOT/ITL since pooling is non-streaming)
 /// - `total_output` is 0 (pooling produces no output tokens)
 /// - `total_token_throughput` = `total_input / dur_s` (input tokens only)
@@ -380,7 +377,7 @@ pub fn calculate_embedding_metrics(
 
 /// Calculate multi-turn benchmark metrics with per-turn breakdown.
 ///
-/// 1. Flatten all turns into SampleRequest/RequestFuncOutput pairs → overall metrics
+/// 1. Flatten all turns' outputs → overall metrics
 /// 2. Group by turn_index → per-turn metrics
 /// 3. Conversation-level stats
 pub fn calculate_multi_turn_metrics(
@@ -391,24 +388,15 @@ pub fn calculate_multi_turn_metrics(
     max_turn_count: usize,
 ) -> MultiTurnMetrics {
     // Flatten all turns for overall metrics
-    let mut all_requests: Vec<SampleRequest> = Vec::new();
     let mut all_outputs: Vec<RequestFuncOutput> = Vec::new();
 
     for conv in conversation_outputs {
         for turn in &conv.turns {
-            all_requests.push(SampleRequest {
-                prompt: std::sync::Arc::from(""),
-                prompt_len: turn.request_output.prompt_len,
-                expected_output_len: turn.request_output.output_tokens,
-                request_id: None,
-                ..Default::default()
-            });
             all_outputs.push(turn.request_output.clone());
         }
     }
 
     let (overall, _) = calculate_metrics(
-        &all_requests,
         &all_outputs,
         dur_s,
         selected_percentiles,
@@ -419,25 +407,16 @@ pub fn calculate_multi_turn_metrics(
     // Per-turn breakdown
     let mut per_turn: Vec<BenchmarkMetrics> = Vec::new();
     for turn_idx in 0..max_turn_count {
-        let mut turn_requests: Vec<SampleRequest> = Vec::new();
         let mut turn_outputs: Vec<RequestFuncOutput> = Vec::new();
 
         for conv in conversation_outputs {
             if let Some(turn) = conv.turns.get(turn_idx) {
-                turn_requests.push(SampleRequest {
-                    prompt: std::sync::Arc::from(""),
-                    prompt_len: turn.request_output.prompt_len,
-                    expected_output_len: turn.request_output.output_tokens,
-                    request_id: None,
-                    ..Default::default()
-                });
                 turn_outputs.push(turn.request_output.clone());
             }
         }
 
         if !turn_outputs.is_empty() {
             let (turn_metrics, _) = calculate_metrics(
-                &turn_requests,
                 &turn_outputs,
                 dur_s,
                 selected_percentiles,
@@ -509,6 +488,26 @@ mod tests {
         assert!((p99 - 98.01).abs() < 0.1);
     }
 
+    /// Input-token totals come from each output's `prompt_len` (server-reported
+    /// via `usage.prompt_tokens`), not the client-side request estimate.
+    /// Mirrors Python's `serve.py`.
+    #[test]
+    fn test_total_input_uses_output_prompt_len() {
+        let outputs = vec![
+            RequestFuncOutput {
+                success: true,
+                prompt_len: 128,
+                output_tokens: 3,
+                ..Default::default()
+            };
+            2
+        ];
+
+        let (metrics, _) = calculate_metrics(&outputs, 1.0, &[], false, &GoodputConfig::default());
+
+        assert_eq!(metrics.total_input, 256);
+    }
+
     #[test]
     fn test_multi_turn_metrics_use_request_prompt_lengths() {
         let turn = |turn_index, prompt_len, output_tokens| crate::multi_turn::TurnOutput {
diff --git a/rust/src/bench/src/metrics/steady_state.rs b/rust/src/bench/src/metrics/steady_state.rs
index a403275b7ac6..f108bb7c8501 100644
--- a/rust/src/bench/src/metrics/steady_state.rs
+++ b/rust/src/bench/src/metrics/steady_state.rs
@@ -170,16 +170,12 @@ pub fn detect_window(
     })
 }
 
-use crate::datasets::SampleRequest;
-
-/// Compute steady-state metrics from per-request outputs, the detected window,
-/// and the original request list (for input-token attribution).
+/// Compute steady-state metrics from per-request outputs and the detected window.
 ///
 /// `is_pooling` gates off TTFT because pooling backends write `ttft = latency`
 /// as a placeholder; reporting TTFT over that would silently be E2EL.
 pub fn compute(
     outputs: &[RequestFuncOutput],
-    input_requests: &[SampleRequest],
     window: &SteadyStateWindow,
     percentiles: &[f64],
     is_pooling: bool,
@@ -200,7 +196,7 @@ pub fn compute(
     let mut ttfts_in_window: Vec<f64> = Vec::new();
     let mut tpots_in_window: Vec<f64> = Vec::new();
 
-    for (idx, o) in outputs.iter().enumerate() {
+    for o in outputs.iter() {
         if !o.success {
             continue;
         }
@@ -213,7 +209,7 @@ pub fn compute(
         }
 
         if started_in {
-            input_tokens_in_window += input_requests[idx].prompt_len;
+            input_tokens_in_window += o.prompt_len;
             if !is_pooling {
                 ttfts_in_window.push(o.ttft);
                 // Per-request TPOT: (latency - ttft) / (output_tokens - 1).
@@ -417,10 +413,6 @@ mod tests {
         assert!((w.end_s - 10.0).abs() < 1e-9);
     }
 
-    use std::sync::Arc;
-
-    use crate::datasets::SampleRequest;
-
     fn mk_full(
         start: f64,
         latency: f64,
@@ -440,16 +432,6 @@ mod tests {
         }
     }
 
-    fn mk_sample(prompt_len: usize) -> SampleRequest {
-        SampleRequest {
-            prompt: Arc::from(""),
-            prompt_len,
-            expected_output_len: 0,
-            request_id: None,
-            ..Default::default()
-        }
-    }
-
     #[test]
     fn compute_constant_concurrency_matches_full_run() {
         // 3 identical requests, each 10s, each emitting 5 tokens uniformly via itl=2s.
@@ -459,13 +441,10 @@ mod tests {
             mk_full(0.0, 10.0, 2.0, itl.clone(), 5),
             mk_full(0.0, 10.0, 2.0, itl.clone(), 5),
         ];
-        let requests: Vec<SampleRequest> =
-            outputs.iter().map(|o| mk_sample(o.prompt_len)).collect();
 
         let w = detect_window(&outputs, Some(3), 1.0, 0.1, 10.0).unwrap();
         let m = compute(
             &outputs,
-            &requests,
             &w,
             &[99.0],
             false, // is_pooling
@@ -502,12 +481,10 @@ mod tests {
         let late_itl = vec![0.02; 999]; // 1000 tokens; emission starts after end_s
         let late = mk_full(9.9, 30.0, 0.05, late_itl, 1000);
         outputs.push(late);
-        let requests: Vec<SampleRequest> =
-            outputs.iter().map(|o| mk_sample(o.prompt_len)).collect();
 
         // Target = 2 so threshold_abs = 2; window opens when both are in flight.
         let w = detect_window(&outputs, Some(2), 1.0, 0.1, 40.0).unwrap();
-        let m = compute(&outputs, &requests, &w, &[99.0], false);
+        let m = compute(&outputs, &w, &[99.0], false);
 
         // Late request started in window but most of its tokens are emitted AFTER end_s.
         // Only a handful (≤3) should count.
@@ -534,10 +511,8 @@ mod tests {
             mk_full(0.0, 5.0, 5.0, Vec::new(), 0),
             mk_full(0.0, 5.0, 5.0, Vec::new(), 0),
         ];
-        let requests: Vec<SampleRequest> =
-            outputs.iter().map(|o| mk_sample(o.prompt_len)).collect();
         let w = detect_window(&outputs, Some(2), 1.0, 0.1, 10.0).unwrap();
-        let m = compute(&outputs, &requests, &w, &[99.0], /* is_pooling */ true);
+        let m = compute(&outputs, &w, &[99.0], /* is_pooling */ true);
         assert_eq!(m.mean_ttft_ms, 0.0);
         assert_eq!(m.median_ttft_ms, 0.0);
         assert!(m.percentiles_ttft_ms.is_empty());
@@ -560,10 +535,8 @@ mod tests {
             mk_full(0.0, 10.0, 2.0, itl.clone(), 5),
             mk_full(0.0, 10.0, 2.0, itl.clone(), 5),
         ];
-        let requests: Vec<SampleRequest> =
-            outputs.iter().map(|o| mk_sample(o.prompt_len)).collect();
         let w = detect_window(&outputs, Some(3), 1.0, 0.1, 10.0).unwrap();
-        let m = compute(&outputs, &requests, &w, &[99.0], false);
+        let m = compute(&outputs, &w, &[99.0], false);
 
         assert!((m.mean_tpot_ms - 2000.0).abs() < 1e-6);
         assert!((m.median_tpot_ms - 2000.0).abs() < 1e-6);
@@ -574,10 +547,8 @@ mod tests {
     #[test]
     fn json_roundtrip_some() {
         let outputs = vec![mk_full(0.0, 10.0, 2.0, vec![2.0; 4], 5); 3];
-        let requests: Vec<SampleRequest> =
-            outputs.iter().map(|o| mk_sample(o.prompt_len)).collect();
         let w = detect_window(&outputs, Some(3), 1.0, 0.1, 10.0).unwrap();
-        let m = compute(&outputs, &requests, &w, &[99.0], false);
+        let m = compute(&outputs, &w, &[99.0], false);
         let s = serde_json::to_string(&m).unwrap();
         let parsed: SteadyStateMetrics = serde_json::from_str(&s).unwrap();
         assert!((parsed.output_throughput - m.output_throughput).abs() < 1e-12);
@@ -605,14 +576,13 @@ mod tests {
         let itl = vec![1.0; 9];
         let outputs: Vec<RequestFuncOutput> =
             (0..5).map(|i| mk_full(i as f64, 20.0, 1.0, itl.clone(), 10)).collect();
-        let requests: Vec<SampleRequest> = (0..5).map(|_| mk_sample(100)).collect();
 
         let w = detect_window(&outputs, Some(5), 1.0, 1.0, 25.0).unwrap();
         assert_eq!(w.threshold_abs, 5);
         assert!((w.start_s - 4.0).abs() < 1e-9);
         assert!((w.end_s - 20.0).abs() < 1e-9);
 
-        let m = compute(&outputs, &requests, &w, &[50.0, 99.0], false);
+        let m = compute(&outputs, &w, &[50.0, 99.0], false);
 
         // Round-trip through JSON.
         let s = serde_json::to_string(&m).unwrap();
diff --git a/rust/src/bench/src/mm_processor.rs b/rust/src/bench/src/mm_processor.rs
index 3454dc7993c7..1f7668a0e9e5 100644
--- a/rust/src/bench/src/mm_processor.rs
+++ b/rust/src/bench/src/mm_processor.rs
@@ -135,6 +135,9 @@ pub struct MmProcessorArgs {
 
 /// Entry point for `vllm-bench mm-processor`.
 pub async fn run_mm_processor(args: MmProcessorArgs, timing: RequestTimingStats) -> Result<()> {
+    if args.max_concurrency == 0 {
+        anyhow::bail!("--max-concurrency must be greater than 0");
+    }
     if args.engine.data_parallel_size_local == Some(0) {
         anyhow::bail!(
             "`--data-parallel-size-local 0` is not supported; the benchmark requires \
@@ -157,6 +160,7 @@ pub async fn run_mm_processor(args: MmProcessorArgs, timing: RequestTimingStats)
         handshake_host: args.engine.handshake_host.clone(),
         handshake_port,
         data_parallel_size: args.engine.data_parallel_size,
+        client_handshake_address: None,
         python_args,
     };
     eprintln!(
@@ -215,14 +219,15 @@ async fn run_mm_processor_with_engine(
     }
 
     // Connect to the engine over the same transport the serving frontend uses.
+    let engine_count =
+        args.engine.data_parallel_size_local.unwrap_or(args.engine.data_parallel_size);
     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),
+            engine_start_index: 0,
+            engine_count,
+            data_parallel_size: engine_count,
             ready_timeout: ENGINE_READY_TIMEOUT,
             local_input_address: None,
             local_output_address: None,
diff --git a/rust/src/bench/src/sweep.rs b/rust/src/bench/src/sweep.rs
index 367e7adfc34e..7fb6a6129e00 100644
--- a/rust/src/bench/src/sweep.rs
+++ b/rust/src/bench/src/sweep.rs
@@ -376,9 +376,15 @@ impl SummaryColumn {
 pub fn parse_concurrency_values(s: &str) -> Result<Vec<usize>> {
     s.split(',')
         .map(|v| {
-            v.trim().parse::<usize>().map_err(|_| {
+            let value = v.trim().parse::<usize>().map_err(|_| {
                 BenchError::Config(format!("Invalid concurrency value: '{}'", v.trim()))
-            })
+            })?;
+            if value == 0 {
+                return Err(BenchError::Config(
+                    "--sweep-max-concurrency values must be greater than 0, got 0".into(),
+                ));
+            }
+            Ok(value)
         })
         .collect()
 }
@@ -561,4 +567,25 @@ mod tests {
             "expected at least 2 dashes, got {dash_count}: {row}"
         );
     }
+
+    #[test]
+    fn parse_concurrency_values_rejects_zero() {
+        for input in ["0", "0,1", "1,0", "1,0,16"] {
+            let err = parse_concurrency_values(input).expect_err("0 must be rejected");
+            let msg = err.to_string();
+            assert!(
+                msg.contains("--sweep-max-concurrency") && msg.contains("greater than 0"),
+                "{input}: {msg}"
+            );
+        }
+    }
+
+    #[test]
+    fn parse_concurrency_values_accepts_positive() {
+        assert_eq!(
+            parse_concurrency_values("1,10,50").unwrap(),
+            vec![1, 10, 50]
+        );
+        assert_eq!(parse_concurrency_values(" 8 , 16 ").unwrap(), vec![8, 16]);
+    }
 }
diff --git a/rust/src/chat/examples/external_engine_chat_qwen.rs b/rust/src/chat/examples/external_engine_chat_qwen.rs
index 1ba4d8a71a48..d210580d1da7 100644
--- a/rust/src/chat/examples/external_engine_chat_qwen.rs
+++ b/rust/src/chat/examples/external_engine_chat_qwen.rs
@@ -62,7 +62,9 @@ async fn main() -> Result<()> {
         transport_mode: TransportMode::HandshakeOwner {
             handshake_address: args.handshake_address.clone(),
             advertised_host: args.host.clone(),
+            engine_start_index: 0,
             engine_count: args.engine_count,
+            data_parallel_size: args.engine_count,
             ready_timeout,
             local_input_address: None,
             local_output_address: None,
diff --git a/rust/src/chat/src/backend/hf.rs b/rust/src/chat/src/backend/hf.rs
index f0da863534c7..a0b3d3451b5b 100644
--- a/rust/src/chat/src/backend/hf.rs
+++ b/rust/src/chat/src/backend/hf.rs
@@ -16,7 +16,7 @@ use crate::backend::{
     ChatBackend, DynChatBackend, LoadModelBackendsOptions, LoadedModelBackends,
     NewChatOutputProcessorOptions,
 };
-use crate::error::Result;
+use crate::error::{Error, Result};
 use crate::multimodal::{MultimodalConfigFiles, MultimodalModelInfo};
 use crate::output::{
     DefaultChatOutputProcessor, HarmonyChatOutputProcessor, validate_harmony_parser_overrides,
@@ -25,7 +25,7 @@ use crate::parser::unified::{HfTemplateError, ResponseTemplate};
 use crate::renderer::hf::{HfChatRenderer, MultimodalRenderInfo};
 use crate::renderer::{
     DeepSeekV4ChatRenderer, DeepSeekV32ChatRenderer, DeepSeekV41ChatRenderer, DynChatRenderer,
-    HarmonyChatRenderer, InklingChatRenderer, KimiK3ChatRenderer,
+    Glm5NextChatRenderer, HarmonyChatRenderer, InklingChatRenderer, KimiK3ChatRenderer,
 };
 use crate::request::ChatRequest;
 use crate::{DynChatOutputProcessor, RendererSelection};
@@ -72,6 +72,10 @@ impl HfChatBackend {
         let response_template = load_response_template(&files)?;
 
         let renderer = options.renderer.resolve(model_type);
+        // Only the HF renderer renders Jinja templates.
+        if renderer != RendererSelection::Hf && options.chat_template.is_some() {
+            return Err(Error::UnsupportedChatTemplateOption { renderer });
+        }
         let chat_renderer: DynChatRenderer = match renderer {
             RendererSelection::Auto => unreachable!("renderer auto should be resolved above"),
             RendererSelection::Hf => Arc::new(HfChatRenderer::load(
@@ -99,6 +103,11 @@ impl HfChatBackend {
                 tokenizer.clone(),
                 options.default_chat_template_kwargs,
             )),
+            RendererSelection::Glm5Next => Arc::new(
+                Glm5NextChatRenderer::new(options.default_chat_template_kwargs)
+                    .with_multimodal(model_type == "glm5_next")
+                    .with_tokenizer(tokenizer.clone())?,
+            ),
         };
 
         info!(
@@ -237,8 +246,8 @@ mod tests {
 
     use super::HfChatBackend;
     use crate::backend::{ChatBackend, LoadModelBackendsOptions, NewChatOutputProcessorOptions};
-    use crate::request::{ChatContent, ChatMessage, ChatRequest};
-    use crate::{ParserSelection, RendererSelection};
+    use crate::request::{ChatContent, ChatMessage, ChatRequest, ResolvedToolContext};
+    use crate::{ChatTool, ParserSelection, RendererSelection};
 
     fn request_with_user_text(text: &str) -> ChatRequest {
         ChatRequest {
@@ -282,6 +291,15 @@ mod tests {
         config_json: &str,
         tokenizer_config_json: &str,
     ) -> HfChatBackend {
+        load_backend(renderer, config_json, tokenizer_config_json, None).unwrap()
+    }
+
+    fn load_backend(
+        renderer: RendererSelection,
+        config_json: &str,
+        tokenizer_config_json: &str,
+        chat_template: Option<&str>,
+    ) -> crate::Result<HfChatBackend> {
         HfChatBackend::from_resolved_model_files(
             resolved_files(config_json, tokenizer_config_json),
             "test-model".to_string(),
@@ -294,14 +312,13 @@ mod tests {
                 renderer,
                 language_model_only: false,
                 chat_template_content_format: Default::default(),
-                chat_template: None,
+                chat_template: chat_template.map(str::to_string),
                 trust_request_chat_template: false,
                 default_chat_template_kwargs: HashMap::new(),
                 limit_mm_per_prompt: HashMap::new(),
             },
             test_tokenizer(),
         )
-        .unwrap()
     }
 
     fn render_prompt(
@@ -343,6 +360,72 @@ mod tests {
         assert_eq!(prompt, "hello");
     }
 
+    #[test]
+    fn hf_renderer_selects_named_tokenizer_config_templates_like_transformers() {
+        let config = r#"{"model_type":"llama"}"#;
+        let named_templates = r#"{"chat_template":[
+            {"name":"default","template":"default:{{ messages[0].content }}"},
+            {"name":"tool_use","template":"tool_use:{{ tools[0].function.name }}"}
+        ]}"#;
+        let load = |files, chat_template: Option<&str>| {
+            HfChatBackend::from_resolved_model_files(
+                files,
+                "test-model".to_string(),
+                LoadModelBackendsOptions {
+                    chat_template: chat_template.map(str::to_owned),
+                    ..Default::default()
+                },
+                test_tokenizer(),
+            )
+            .unwrap()
+        };
+
+        let tokenizer_config_only = load(resolved_files(config, named_templates), None);
+        // A configured or dedicated template replaces every named template,
+        // `tool_use` included.
+        let configured = load(
+            resolved_files(config, named_templates),
+            Some("configured:{{ messages[0].content }}"),
+        );
+        let mut files = resolved_files(config, named_templates);
+        let dedicated_path =
+            files.config_path.as_ref().unwrap().with_file_name("chat_template.jinja");
+        std::fs::write(&dedicated_path, "dedicated:{{ messages[0].content }}").unwrap();
+        files.chat_template_path = Some(dedicated_path);
+        let dedicated = load(files, None);
+
+        let request = request_with_user_text("hello");
+        let mut with_tools = request.clone();
+        with_tools.tool_context = ResolvedToolContext::new(
+            &request.messages,
+            vec![ChatTool {
+                name: "lookup".to_string(),
+                description: None,
+                parameters: serde_json::json!({"type": "object"}),
+                strict: None,
+                defer_loading: None,
+            }],
+            None,
+            true,
+        )
+        .unwrap();
+        let rendered = [tokenizer_config_only, configured, dedicated].map(|backend| {
+            [&request, &with_tools].map(|request| {
+                backend.chat_renderer().render(request).unwrap().prompt.into_text().unwrap()
+            })
+        });
+
+        assert_eq!(
+            rendered,
+            [
+                ["default:hello", "tool_use:lookup"],
+                ["configured:hello", "configured:hello"],
+                ["dedicated:hello", "dedicated:hello"],
+            ]
+            .map(|row| row.map(str::to_owned))
+        );
+    }
+
     #[test]
     fn native_renderer_inherits_deployment_reasoning_below_request_controls() {
         let backend = HfChatBackend::from_resolved_model_files(
@@ -404,6 +487,43 @@ mod tests {
         );
     }
 
+    #[tokio::test]
+    async fn only_hf_renderer_accepts_chat_template() {
+        use RendererSelection::Auto;
+
+        const HF_CONFIG: &str = r#"{"model_type":"llama"}"#;
+        const NATIVE_CONFIG: &str = r#"{"model_type":"deepseek_v4"}"#;
+        const TOKENIZER_CONFIG: &str = r#"{"chat_template":"{{ messages[0].content }}"}"#;
+        const TEMPLATE: &str = "{{ messages[0].content }}!";
+
+        load_backend(Auto, HF_CONFIG, TOKENIZER_CONFIG, Some(TEMPLATE)).unwrap();
+        let Err(error) = load_backend(Auto, NATIVE_CONFIG, TOKENIZER_CONFIG, Some(TEMPLATE)) else {
+            panic!("native renderers should reject --chat-template");
+        };
+        expect_test::expect!["--chat-template is not supported by the `deepseek_v4` renderer; use --tokenizer-mode hf to render a Jinja chat template"]
+            .assert_eq(&error.to_report_string());
+
+        let prepare = async |config_json| {
+            let backend = backend_for_selection(Auto, config_json, TOKENIZER_CONFIG);
+            let mut request = request_with_user_text("hello");
+            request.chat_options.chat_template = Some(TEMPLATE.to_string());
+            crate::ChatRequestProcessor::render_only(Arc::new(backend))
+                .prepare_for_tokenization(request)
+                .await
+                .unwrap_err()
+        };
+        // The HF renderer goes on to its trust check.
+        assert!(matches!(
+            prepare(HF_CONFIG).await,
+            crate::Error::UntrustedRequestChatTemplate
+        ));
+        let error = prepare(NATIVE_CONFIG).await;
+        assert!(error.is_request_validation_error());
+        assert_eq!(error.param(), Some("chat_template"));
+        expect_test::expect!["chat_template is not supported by this chat renderer"]
+            .assert_eq(&error.to_report_string());
+    }
+
     #[test]
     fn auto_uses_harmony_renderer_and_output_processor_for_gpt_oss_model_type() {
         let backend = backend_for_selection(
diff --git a/rust/src/chat/src/error.rs b/rust/src/chat/src/error.rs
index 7245a07b1339..3e48d242ee89 100644
--- a/rust/src/chat/src/error.rs
+++ b/rust/src/chat/src/error.rs
@@ -4,6 +4,8 @@
 use thiserror::Error;
 use thiserror_ext::{AsReport as _, Macro};
 
+use crate::RendererSelection;
+
 type BoxedError = Box<dyn std::error::Error + Send + Sync>;
 
 #[derive(Debug, Error, Macro)]
@@ -30,10 +32,21 @@ pub enum Error {
          Refused request with untrusted chat template."
     )]
     UntrustedRequestChatTemplate,
+    #[error("chat_template is not supported by this chat renderer")]
+    UnsupportedRequestChatTemplate,
+    #[error(
+        "--chat-template is not supported by the `{renderer}` renderer; use --tokenizer-mode hf to render a Jinja chat template"
+    )]
+    UnsupportedChatTemplateOption { renderer: RendererSelection },
     #[error("multimodal input is not supported by this chat renderer")]
     UnsupportedMultimodalRenderer,
     #[error("unsupported multimodal content: {0}")]
     UnsupportedMultimodalContent(&'static str),
+    #[error(
+        "multimodal content in a non-leading system or developer message is not supported \
+         when the chat template does not support the `developer` role"
+    )]
+    UnsupportedMultimodalSystemFallback,
     #[error("`{modality}` input is not supported by this model")]
     UnsupportedModality { modality: String },
     #[error("At most {limit} {modality}(s) may be provided in one prompt.")]
@@ -120,10 +133,12 @@ impl Error {
             | Self::ToolChoiceRequiresTools
             | Self::ToolChoiceFunctionNotFound { .. }
             | Self::UnsupportedChatRole { .. }
-            | Self::UntrustedRequestChatTemplate => true,
+            | Self::UntrustedRequestChatTemplate
+            | Self::UnsupportedRequestChatTemplate => true,
             Self::Text(error) => error.is_request_validation_error(),
             Self::UnsupportedMultimodalRenderer
             | Self::UnsupportedMultimodalContent(_)
+            | Self::UnsupportedMultimodalSystemFallback
             | Self::UnsupportedModality { .. }
             | Self::InvalidPreprocessedMultimodal { .. }
             | Self::MmLimitExceeded { .. } => true,
@@ -137,6 +152,7 @@ impl Error {
     pub fn param(&self) -> Option<&'static str> {
         match self {
             Self::Text(error) => error.param(),
+            Self::UnsupportedRequestChatTemplate => Some("chat_template"),
             _ => None,
         }
     }
diff --git a/rust/src/chat/src/event.rs b/rust/src/chat/src/event.rs
index bbcec8121c79..0302684a5e03 100644
--- a/rust/src/chat/src/event.rs
+++ b/rust/src/chat/src/event.rs
@@ -6,7 +6,7 @@ use std::sync::Arc;
 
 use serde::{Deserialize, Serialize};
 use vllm_engine_core_client::protocol::output::RequestSpecDecodeMetrics;
-use vllm_llm::{RequestTimestamps, TokenUsage};
+use vllm_llm::{RequestTimestamps, TokenUsage, WeightVersionSpan};
 use vllm_text::{DecodedLogprobs, DecodedPromptLogprobs};
 
 use crate::FinishReason;
@@ -248,5 +248,7 @@ pub enum ChatEvent {
         /// Connector-specific encoder cache transfer parameters for
         /// disaggregated serving.
         ec_transfer_params: Option<Box<serde_json::Value>>,
+        /// Weight-version spans over all generated token positions.
+        weight_versions: Option<Vec<WeightVersionSpan>>,
     },
 }
diff --git a/rust/src/chat/src/lib.rs b/rust/src/chat/src/lib.rs
index 47d813aae66a..a8f67c4ff680 100644
--- a/rust/src/chat/src/lib.rs
+++ b/rust/src/chat/src/lib.rs
@@ -221,21 +221,34 @@ impl ChatRequestProcessor {
         Ok(Some(features))
     }
 
-    async fn prepare_text_request(&self, mut request: ChatRequest) -> Result<TextRequest> {
+    /// Render one request into a text request, also returning the rendered
+    /// prompt text when the request asked for it and the renderer produced
+    /// text.
+    async fn prepare_text_request(
+        &self,
+        mut request: ChatRequest,
+    ) -> Result<(TextRequest, Option<String>)> {
         // Stamp before rendering so render and tokenize count toward TTFT/e2e.
         let arrival_time = vllm_llm::current_unix_timestamp_secs();
         if self.exclude_tools_when_tool_choice_none {
             request.tool_context.exclude_initial_tools_if_none();
         }
         let renderer = self.backend.chat_renderer();
+        if request.chat_options.chat_template.is_some() && !renderer.renders_request_chat_template()
+        {
+            return Err(Error::UnsupportedRequestChatTemplate);
+        }
         let (mut rendered, request) = tokio::task::spawn_blocking(move || {
             renderer.render(&request).map(|rendered| (rendered, request))
         })
         .await
         .map_err(|e| Error::ChatTemplate(format!("render task join error: {e}")))??;
         let chat_template_kwargs = std::mem::take(&mut rendered.effective_template_kwargs);
+        // Taken before media placeholders are expanded, like Python.
+        let prompt_text =
+            request.return_prompt_text.then(|| rendered.prompt.as_text().cloned()).flatten();
         let (prompt, mm_features) = self.finalize_rendered_prompt(&request, rendered).await?;
-        Ok(TextRequest {
+        let text_request = TextRequest {
             request_id: request.request_id,
             prompt,
             mm_features,
@@ -255,14 +268,16 @@ impl ChatRequestProcessor {
             reasoning_ended: None,
             lora_request: request.lora_request,
             arrival_time: Some(arrival_time),
-        })
+        };
+        Ok((text_request, prompt_text))
     }
 
     /// Prepare one chat request for tokenization without constructing an output processor.
     pub async fn prepare_for_tokenization(&self, mut request: ChatRequest) -> Result<TextRequest> {
         request.validate()?;
         self.map_reasoning_effort(&mut request);
-        self.prepare_text_request(request).await
+        let (text_request, _) = self.prepare_text_request(request).await?;
+        Ok(text_request)
     }
 
     /// Prepare one chat request without submitting it to an engine.
@@ -275,7 +290,7 @@ impl ChatRequestProcessor {
         &self,
         mut request: ChatRequest,
         text_processor: &TextRequestProcessor,
-    ) -> Result<(TextRequest, DynChatOutputProcessor)> {
+    ) -> Result<PreparedChatRequest> {
         request.validate()?;
         self.map_reasoning_effort(&mut request);
 
@@ -287,7 +302,7 @@ impl ChatRequestProcessor {
                 tool_strict_level: self.tool_strict_level,
             },
         )?;
-        let mut text_request = self.prepare_text_request(request).await?;
+        let (mut text_request, prompt_text) = 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.
@@ -311,10 +326,25 @@ impl ChatRequestProcessor {
             }
         }
 
-        Ok((text_request, output_processor))
+        Ok(PreparedChatRequest {
+            text_request,
+            output_processor,
+            prompt_text,
+        })
     }
 }
 
+/// One chat request rendered, tokenized, and ready for submission.
+pub struct PreparedChatRequest {
+    /// The tokenized text request, with any output grammar applied.
+    pub text_request: TextRequest,
+    /// Request-scoped processor for the decoded output.
+    pub output_processor: DynChatOutputProcessor,
+    /// The rendered prompt text, when the request asked for it and the
+    /// renderer produced text.
+    pub prompt_text: Option<String>,
+}
+
 /// Structured chat facade above [`TextLlm`].
 ///
 /// This layer stays above raw text semantics: it takes care of chat-template
@@ -446,14 +476,17 @@ 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, self.text.request_processor()).await?;
+        let PreparedChatRequest {
+            text_request,
+            output_processor,
+            prompt_text,
+        } = 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();
 
         let structured_stream = output_processor.process(decoded_stream)?;
 
-        Ok(ChatEventStream::new(request_id, structured_stream))
+        Ok(ChatEventStream::new(request_id, structured_stream).with_prompt_text(prompt_text))
     }
 
     /// Abort in-flight requests by their external (user-supplied) request ids.
diff --git a/rust/src/chat/src/output/default/mod.rs b/rust/src/chat/src/output/default/mod.rs
index 8e125b0679d1..3b4bb0617a33 100644
--- a/rust/src/chat/src/output/default/mod.rs
+++ b/rust/src/chat/src/output/default/mod.rs
@@ -410,6 +410,7 @@ mod tests {
                     ec_transfer_params: None,
                     sampling_mask: None,
                     spec_decode_metrics: None,
+                    weight_versions: None,
                     timestamps: Default::default(),
                 })),
             })]);
diff --git a/rust/src/chat/src/output/default/unified.rs b/rust/src/chat/src/output/default/unified.rs
index da4f16efc075..cd1208888e05 100644
--- a/rust/src/chat/src/output/default/unified.rs
+++ b/rust/src/chat/src/output/default/unified.rs
@@ -282,6 +282,7 @@ pub(crate) async fn unified_event_stream(
                         finish_reason: finished.finish_reason,
                         kv_transfer_params: finished.kv_transfer_params,
                         ec_transfer_params: finished.ec_transfer_params,
+                        weight_versions: finished.weight_versions,
                     })
                     .await;
                 }
@@ -422,6 +423,7 @@ mod tests {
                 ec_transfer_params: None,
                 sampling_mask: None,
                 spec_decode_metrics: None,
+                weight_versions: None,
                 timestamps: Default::default(),
             })),
         }
@@ -751,6 +753,7 @@ mod tests {
                     finish_reason: crate::FinishReason::Stop(None),
                     kv_transfer_params: None,
                     ec_transfer_params: None,
+                    weight_versions: None,
                 },
             ]
         );
@@ -793,6 +796,7 @@ mod tests {
                     finish_reason: crate::FinishReason::Stop(None),
                     kv_transfer_params: None,
                     ec_transfer_params: None,
+                    weight_versions: None,
                 },
             ]
         );
@@ -827,6 +831,7 @@ mod tests {
                     finish_reason: crate::FinishReason::Stop(None),
                     kv_transfer_params: None,
                     ec_transfer_params: None,
+                    weight_versions: None,
                 },
             ]
         );
@@ -881,6 +886,7 @@ mod tests {
                     finish_reason: crate::FinishReason::Stop(None),
                     kv_transfer_params: None,
                     ec_transfer_params: None,
+                    weight_versions: None,
                 },
             ]
         );
@@ -929,6 +935,7 @@ mod tests {
                     finish_reason: crate::FinishReason::Stop(None),
                     kv_transfer_params: None,
                     ec_transfer_params: None,
+                    weight_versions: None,
                 },
             ]
         );
diff --git a/rust/src/chat/src/output/harmony/mod.rs b/rust/src/chat/src/output/harmony/mod.rs
index d4227041a7e3..2dc1b2e18739 100644
--- a/rust/src/chat/src/output/harmony/mod.rs
+++ b/rust/src/chat/src/output/harmony/mod.rs
@@ -15,7 +15,7 @@ use vllm_text::output::{DecodedTextEvent, SampledDelta};
 
 use crate::Result as ChatResult;
 use crate::error::{Error, Result};
-use crate::event::AssistantBlockKind;
+use crate::event::{AssistantBlockKind, ChatTokenUsage};
 use crate::output::{
     AssistantEvent, ChatOutputProcessor, DynChatEventStream, DynDecodedTextEventStream,
     generate_tool_call_id,
@@ -47,6 +47,8 @@ struct HarmonyGroupKey {
 struct HarmonyGroup {
     key: HarmonyGroupKey,
     text: String,
+    /// Reasoning tokens attributed to this group's text.
+    reasoning_tokens: usize,
 }
 
 #[derive(Debug)]
@@ -57,6 +59,8 @@ struct OpenHarmonyToolCall {
 struct HarmonyState {
     /// Incremental Harmony parser over assistant token IDs.
     parser: StreamableParser,
+    /// Harmony encoding, used to tell control tokens from content tokens.
+    encoding: &'static HarmonyEncoding,
     /// Whether tool-call content should surface as structured tool events.
     tool_calls_enabled: bool,
     /// Count of completed visible assistant messages for newline insertion.
@@ -67,6 +71,13 @@ struct HarmonyState {
     current_text_group: Option<HarmonyGroupKey>,
     /// The currently open Harmony tool recipient, if any.
     open_tool_call: Option<OpenHarmonyToolCall>,
+    /// Total generated tokens counted as reasoning, reported in the final
+    /// usage.
+    reasoning_tokens: usize,
+    /// Reasoning tokens of the current message not yet attributed to an
+    /// emitted reasoning delta, such as a token whose content delta is still
+    /// empty.
+    unattributed_reasoning_tokens: usize,
 }
 
 impl HarmonyChatOutputProcessor {
@@ -119,18 +130,48 @@ impl ChatOutputProcessor for HarmonyChatOutputProcessor {
 
 impl HarmonyState {
     /// Create one fresh Harmony streaming state for a new assistant response.
-    fn new(encoding: HarmonyEncoding, tool_calls_enabled: bool) -> Result<Self> {
+    fn new(encoding: &'static HarmonyEncoding, tool_calls_enabled: bool) -> Result<Self> {
         Ok(Self {
-            parser: StreamableParser::new(encoding, Some(Role::Assistant))
+            parser: StreamableParser::new(encoding.clone(), Some(Role::Assistant))
                 .map_err(harmony_output_parsing_error)?,
+            encoding,
             tool_calls_enabled,
             completed_visible_messages: 0,
             completed_reasoning_messages: 0,
             current_text_group: None,
             open_tool_call: None,
+            reasoning_tokens: 0,
+            unattributed_reasoning_tokens: 0,
         })
     }
 
+    /// Count the token the parser just consumed as reasoning if it is content
+    /// of an `analysis` message or of a `commentary` message with a recipient,
+    /// excluding Harmony control tokens, like Python's
+    /// `HarmonyParser._is_reasoning_token`.
+    fn count_reasoning_token(&mut self, token_id: u32) {
+        let is_reasoning_channel = match self.parser.current_channel().as_deref() {
+            Some("analysis") => true,
+            Some("commentary") => self.parser.current_recipient().is_some(),
+            _ => false,
+        };
+        if is_reasoning_channel && !self.encoding.tokenizer().is_special_token(token_id) {
+            self.reasoning_tokens += 1;
+            self.unattributed_reasoning_tokens += 1;
+        }
+    }
+
+    /// Take the reasoning tokens not yet attributed if `key` is a reasoning
+    /// text group.
+    fn take_reasoning_tokens_for(&mut self, key: &HarmonyGroupKey) -> usize {
+        match text_block_kind(key.channel.as_deref(), key.recipient.as_deref()) {
+            Some(AssistantBlockKind::Reasoning) => {
+                std::mem::take(&mut self.unattributed_reasoning_tokens)
+            }
+            _ => 0,
+        }
+    }
+
     fn process_token_ids(&mut self, token_ids: &[u32]) -> Result<Vec<AssistantEvent>> {
         let mut events = Vec::new();
         let mut pending_group: Option<HarmonyGroup> = None;
@@ -139,6 +180,7 @@ impl HarmonyState {
             let completed_before = self.parser.messages().len();
             self.parser.process(token_id).map_err(harmony_output_parsing_error)?;
             let completed_after = self.parser.messages().len();
+            self.count_reasoning_token(token_id);
 
             if let Some(delta) = self
                 .parser
@@ -151,14 +193,22 @@ impl HarmonyState {
                     channel: self.parser.current_channel(),
                     recipient: self.parser.current_recipient(),
                 };
+                let reasoning_tokens = self.take_reasoning_tokens_for(&key);
 
                 match pending_group.as_mut() {
-                    Some(group) if group.key == key => group.text.push_str(&delta),
+                    Some(group) if group.key == key => {
+                        group.text.push_str(&delta);
+                        group.reasoning_tokens += reasoning_tokens;
+                    }
                     _ => {
                         if let Some(group) = pending_group.take() {
                             self.emit_group(group, &mut events);
                         }
-                        pending_group = Some(HarmonyGroup { key, text: delta });
+                        pending_group = Some(HarmonyGroup {
+                            key,
+                            text: delta,
+                            reasoning_tokens,
+                        });
                     }
                 }
             }
@@ -214,10 +264,12 @@ impl HarmonyState {
         let final_text = harmony_message_text(final_message);
         let tail = final_text.strip_prefix(&pending_content).unwrap_or(final_text).to_string();
         if !tail.is_empty() {
+            let reasoning_tokens = self.take_reasoning_tokens_for(&pending_key);
             self.emit_group(
                 HarmonyGroup {
                     key: pending_key,
                     text: tail,
+                    reasoning_tokens,
                 },
                 &mut events,
             );
@@ -258,8 +310,6 @@ impl HarmonyState {
                     events.push(AssistantEvent::TextDelta {
                         kind,
                         delta: "\n".to_string(),
-                        // TODO: measure reasoning tokens from native
-                        // analysis-channel token counts.
                         token_count: None,
                     });
                 }
@@ -267,12 +317,12 @@ impl HarmonyState {
                 self.current_text_group = Some(group.key.clone());
             }
 
+            let token_count =
+                (kind == AssistantBlockKind::Reasoning).then_some(group.reasoning_tokens);
             events.push(AssistantEvent::TextDelta {
                 kind,
                 delta: group.text,
-                // TODO: measure reasoning tokens from native
-                // analysis-channel token counts.
-                token_count: None,
+                token_count,
             });
             return;
         }
@@ -307,6 +357,10 @@ impl HarmonyState {
 
     /// Update newline and open-tool state after one Harmony message completes.
     fn handle_completed_message(&mut self, key: HarmonyGroupKey) {
+        // Reasoning tokens of a tool-call message are never attributed to a
+        // delta; they reach the client through the final usage only.
+        self.unattributed_reasoning_tokens = 0;
+
         if self.current_text_group.as_ref() == Some(&key) {
             self.current_text_group = None;
         }
@@ -334,7 +388,7 @@ async fn harmony_assistant_event_stream(
     tool_calls_enabled: bool,
     mut y: TryYielder<AssistantEvent, Error>,
 ) -> Result<()> {
-    let mut state = HarmonyState::new(encoding.clone(), tool_calls_enabled)?;
+    let mut state = HarmonyState::new(encoding, tool_calls_enabled)?;
     futures::pin_mut!(decoded);
 
     while let Some(event) = decoded.next().await.transpose()? {
@@ -378,12 +432,16 @@ async fn harmony_assistant_event_stream(
 
                 if let Some(finished) = finished {
                     y.yield_ok(AssistantEvent::Done {
-                        usage: finished.usage.into(),
+                        usage: ChatTokenUsage {
+                            engine: finished.usage,
+                            reasoning_tokens: state.reasoning_tokens,
+                        },
                         timestamps: finished.timestamps,
                         spec_decode_metrics: finished.spec_decode_metrics,
                         finish_reason: finished.finish_reason,
                         kv_transfer_params: finished.kv_transfer_params,
                         ec_transfer_params: finished.ec_transfer_params,
+                        weight_versions: finished.weight_versions,
                     })
                     .await;
                 }
diff --git a/rust/src/chat/src/output/harmony/tests.rs b/rust/src/chat/src/output/harmony/tests.rs
index e05a15bfd53a..1c8a91c713ef 100644
--- a/rust/src/chat/src/output/harmony/tests.rs
+++ b/rust/src/chat/src/output/harmony/tests.rs
@@ -11,6 +11,7 @@ use vllm_text::output::{
 };
 
 use super::*;
+use crate::event::AssistantBlockKind;
 use crate::output::ChatOutputProcessor;
 use crate::request::{ChatRequest, ChatTool, ChatToolChoice, ResolvedToolContext};
 use crate::{AssistantMessageExt, ChatEvent, ChatTokenUsage, FinishReason};
@@ -60,6 +61,7 @@ fn finished() -> Finished {
         ec_transfer_params: None,
         sampling_mask: None,
         spec_decode_metrics: None,
+        weight_versions: None,
         timestamps: Default::default(),
     }
 }
@@ -139,6 +141,7 @@ fn interrupted_final_message_is_preserved() {
             finish_reason: FinishReason::stop_eos(),
             kv_transfer_params: None,
             ec_transfer_params: None,
+            weight_versions: None,
         })
     );
 }
@@ -182,16 +185,20 @@ fn interrupted_analysis_message_is_preserved() {
                     text: "think".to_string(),
                 }],
             },
-            usage: ChatTokenUsage::from(vllm_llm::TokenUsage {
-                prompt_token_count: 0,
-                output_token_count: 0,
-                prompt_cache_usage: Default::default(),
-            }),
+            usage: ChatTokenUsage {
+                engine: vllm_llm::TokenUsage {
+                    prompt_token_count: 0,
+                    output_token_count: 0,
+                    prompt_cache_usage: Default::default(),
+                },
+                reasoning_tokens: 1,
+            },
             timestamps: Default::default(),
             spec_decode_metrics: None,
             finish_reason: FinishReason::stop_eos(),
             kv_transfer_params: None,
             ec_transfer_params: None,
+            weight_versions: None,
         })
     );
 }
@@ -286,6 +293,124 @@ fn tool_calls_stream_arguments_and_finish_with_local_id_shape() {
     assert_eq!(end.arguments, r#"{"city":"Paris"}"#);
 }
 
+fn content_token_count(text: &str) -> usize {
+    harmony_encoding().unwrap().tokenizer().encode_ordinary(text).len()
+}
+
+fn reasoning_token_counts(events: &[ChatEvent]) -> (Vec<usize>, usize) {
+    let deltas = events
+        .iter()
+        .filter_map(|event| match event {
+            ChatEvent::BlockDelta {
+                kind: AssistantBlockKind::Reasoning,
+                token_count,
+                ..
+            // The newline separating two reasoning messages has no tokens.
+            } => Some(token_count.unwrap_or_default()),
+            _ => None,
+        })
+        .collect();
+    let ChatEvent::Done { usage, .. } = events.last().unwrap() else {
+        panic!("expected done");
+    };
+    (deltas, usage.reasoning_tokens)
+}
+
+#[test]
+fn reasoning_tokens_count_analysis_and_tool_call_content() {
+    let analysis = "Let me check the weather.";
+    let arguments = r#"{"city":"Paris"}"#;
+    let tokens = completion_tokens(&[
+        text_message("analysis", analysis),
+        tool_message("get_weather", arguments, "commentary"),
+    ]);
+    let events = block_on(collect_events(
+        HarmonyChatOutputProcessor::new(&request_with_tools()).unwrap(),
+        vec![
+            decoded_start(),
+            decoded_tokens(tokens, None, Some(finished())),
+        ],
+    ));
+
+    // Like Python, content of `analysis` and of `commentary` with a recipient
+    // counts, Harmony control and header tokens do not. Tool-call tokens are
+    // only in the final usage.
+    let (deltas, total) = reasoning_token_counts(&events);
+    assert_eq!(deltas.iter().sum::<usize>(), content_token_count(analysis));
+    assert_eq!(
+        total,
+        content_token_count(analysis) + content_token_count(arguments)
+    );
+}
+
+#[test]
+fn reasoning_tokens_exclude_final_and_preamble_content() {
+    let tokens = completion_tokens(&[
+        text_message("commentary", "Let me check."),
+        text_message("final", "Hello"),
+    ]);
+    let events = block_on(collect_events(
+        HarmonyChatOutputProcessor::new(&ChatRequest::for_test()).unwrap(),
+        vec![
+            decoded_start(),
+            decoded_tokens(tokens, None, Some(finished())),
+        ],
+    ));
+
+    assert_eq!(reasoning_token_counts(&events), (vec![], 0));
+}
+
+fn stream_one_token_per_update(messages: &[Message]) -> Vec<ChatEvent> {
+    let tokens = completion_tokens(messages);
+    let last = tokens.len() - 1;
+    let mut decoded = vec![decoded_start()];
+    for (i, token) in tokens.into_iter().enumerate() {
+        decoded.push(decoded_tokens(
+            vec![token],
+            None,
+            (i == last).then(finished),
+        ));
+    }
+    block_on(collect_events(
+        HarmonyChatOutputProcessor::new(&ChatRequest::for_test()).unwrap(),
+        decoded,
+    ))
+}
+
+#[test]
+fn reasoning_tokens_are_attributed_to_their_streamed_deltas() {
+    let first = "Let me think about this.";
+    let second = "Then answer.";
+    let events = stream_one_token_per_update(&[
+        text_message("analysis", first),
+        text_message("analysis", second),
+        text_message("final", "Hi"),
+    ]);
+
+    // One delta per content token, and the separating newline carries none.
+    let mut expected = vec![1; content_token_count(first)];
+    expected.push(0);
+    expected.extend(vec![1; content_token_count(second)]);
+    let total = expected.iter().sum();
+    assert_eq!(reasoning_token_counts(&events), (expected, total));
+}
+
+#[test]
+fn reasoning_tokens_of_a_split_character_join_its_delta() {
+    // Each of these characters spans several byte-level tokens, so the first
+    // tokens of a character produce no text delta of their own.
+    let text = "\u{1D518}\u{1D52B}";
+    let tokens = content_token_count(text);
+    assert!(tokens > text.chars().count());
+    let events =
+        stream_one_token_per_update(&[text_message("analysis", text), text_message("final", "Hi")]);
+
+    let (deltas, total) = reasoning_token_counts(&events);
+    assert_eq!(deltas.len(), text.chars().count());
+    assert_eq!(deltas.iter().sum::<usize>(), tokens);
+    assert_eq!(total, tokens);
+}
+
 #[test]
 fn semantic_events_precede_same_update_logprobs() {
     let tokens = completion_tokens(&[text_message("final", "hello")]);
diff --git a/rust/src/chat/src/output/mod.rs b/rust/src/chat/src/output/mod.rs
index 6e6ad4909e2a..e06f70dabaf6 100644
--- a/rust/src/chat/src/output/mod.rs
+++ b/rust/src/chat/src/output/mod.rs
@@ -8,7 +8,7 @@ use futures::Stream;
 use trait_set::trait_set;
 use uuid::Uuid;
 use vllm_engine_core_client::protocol::output::RequestSpecDecodeMetrics;
-use vllm_llm::RequestTimestamps;
+use vllm_llm::{RequestTimestamps, WeightVersionSpan};
 use vllm_parser::output_grammar::BuiltOutputGrammar;
 use vllm_text::output::{DecodedLogprobs, DecodedPromptLogprobs, DecodedTextEvent};
 
@@ -64,6 +64,8 @@ pub(crate) enum AssistantEvent {
         /// Connector-specific encoder cache transfer parameters for
         /// disaggregated serving.
         ec_transfer_params: Option<Box<serde_json::Value>>,
+        /// Weight-version spans over all generated token positions.
+        weight_versions: Option<Vec<WeightVersionSpan>>,
     },
 }
 
diff --git a/rust/src/chat/src/output/structured.rs b/rust/src/chat/src/output/structured.rs
index 583d8e78f487..aca2d7b61423 100644
--- a/rust/src/chat/src/output/structured.rs
+++ b/rust/src/chat/src/output/structured.rs
@@ -11,7 +11,7 @@
 use asynk_strim_attr::{TryYielder, try_stream};
 use futures::{StreamExt as _, pin_mut};
 use vllm_engine_core_client::protocol::output::RequestSpecDecodeMetrics;
-use vllm_llm::RequestTimestamps;
+use vllm_llm::{RequestTimestamps, WeightVersionSpan};
 use vllm_text::DecodedLogprobs;
 
 use super::{AssistantEvent, AssistantEventStream};
@@ -157,6 +157,7 @@ impl StructuredEventState {
         finish_reason: FinishReason,
         kv_transfer_params: Option<Box<serde_json::Value>>,
         ec_transfer_params: Option<Box<serde_json::Value>>,
+        weight_versions: Option<Vec<WeightVersionSpan>>,
     ) -> Result<Vec<ChatEvent>> {
         let mut events = Vec::new();
         self.close_open_text_block(&mut events);
@@ -169,6 +170,7 @@ impl StructuredEventState {
             finish_reason,
             kv_transfer_params,
             ec_transfer_params,
+            weight_versions,
         });
         Ok(events)
     }
@@ -324,6 +326,7 @@ pub(crate) async fn structured_chat_event_stream(
                 finish_reason,
                 kv_transfer_params,
                 ec_transfer_params,
+                weight_versions,
             } => {
                 for next in state.finish(
                     usage,
@@ -332,6 +335,7 @@ pub(crate) async fn structured_chat_event_stream(
                     finish_reason,
                     kv_transfer_params,
                     ec_transfer_params,
+                    weight_versions,
                 )? {
                     y.yield_ok(next).await;
                 }
@@ -372,6 +376,7 @@ mod tests {
                 finish_reason: FinishReason::stop_eos(),
                 kv_transfer_params: None,
                 ec_transfer_params: None,
+                weight_versions: None,
             }),
         ]);
 
@@ -429,6 +434,7 @@ mod tests {
                 finish_reason: FinishReason::stop_eos(),
                 kv_transfer_params: None,
                 ec_transfer_params: None,
+                weight_versions: None,
             }),
         ]);
 
@@ -484,6 +490,7 @@ mod tests {
                 finish_reason: FinishReason::stop_eos(),
                 kv_transfer_params: None,
                 ec_transfer_params: None,
+                weight_versions: None,
             }),
         ]);
 
@@ -539,6 +546,7 @@ mod tests {
                 finish_reason: FinishReason::stop_eos(),
                 kv_transfer_params: None,
                 ec_transfer_params: None,
+                weight_versions: None,
             }),
         ]);
 
@@ -609,6 +617,7 @@ mod tests {
                 finish_reason: FinishReason::stop_eos(),
                 kv_transfer_params: None,
                 ec_transfer_params: None,
+                weight_versions: None,
             }),
         ]);
 
@@ -662,6 +671,7 @@ mod tests {
                 finish_reason: FinishReason::stop_eos(),
                 kv_transfer_params: None,
                 ec_transfer_params: None,
+                weight_versions: None,
             }),
         ]);
 
diff --git a/rust/src/chat/src/renderer/glm5_next/encoding.rs b/rust/src/chat/src/renderer/glm5_next/encoding.rs
new file mode 100644
index 000000000000..3fa988a5fbe0
--- /dev/null
+++ b/rust/src/chat/src/renderer/glm5_next/encoding.rs
@@ -0,0 +1,535 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+use std::collections::{HashMap, HashSet};
+use std::fmt::Write as _;
+use std::io::{self, Write};
+
+use serde::Serialize;
+use serde_json::ser::{Formatter, Serializer};
+use serde_json::{Value, json};
+use thiserror_ext::AsReport as _;
+
+use super::writer::Glm5PromptWriter;
+use crate::error::{Error, Result, bail_chat_template_thrown, chat_template_thrown};
+use crate::renderer::MediaPartSource;
+use crate::request::{ChatContent, ChatContentPart, ChatMessage, ChatTool, GenerationPromptMode};
+use crate::{AssistantContentBlock, AssistantMessageExt};
+
+const TOOLS_PREFIX: &str = "<|system|>\n# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>\n";
+const TOOLS_SUFFIX: &str = "</tools>\n\nFor each function call, output the function name and arguments within the following XML format:\n<tool_call>{function-name}<arg_key>{arg-key-1}</arg_key><arg_value>{arg-value-1}</arg_value><arg_key>{arg-key-2}</arg_key><arg_value>{arg-value-2}</arg_value>...</tool_call>";
+
+pub(super) fn truthy(value: &Value) -> bool {
+    match value {
+        Value::Null => false,
+        Value::Bool(value) => *value,
+        Value::Number(value) => value.as_f64() != Some(0.0),
+        Value::String(value) => !value.is_empty(),
+        Value::Array(value) => !value.is_empty(),
+        Value::Object(value) => !value.is_empty(),
+    }
+}
+
+fn text(value: &Value) -> String {
+    match value {
+        Value::Null => "None".into(),
+        Value::Bool(true) => "True".into(),
+        Value::Bool(false) => "False".into(),
+        Value::String(value) => value.clone(),
+        _ => value.to_string(),
+    }
+}
+
+#[derive(Default)]
+struct MediaContext {
+    enabled: bool,
+    order: Vec<MediaPartSource>,
+}
+
+fn visible_text(
+    out: &mut Glm5PromptWriter<'_>,
+    value: &Value,
+    media: &mut MediaContext,
+) -> Result<()> {
+    match value {
+        Value::Null => {}
+        Value::String(value) => out.ordinary(value),
+        Value::Array(parts) => {
+            for part in parts {
+                if let Some(part) = part.as_str() {
+                    out.ordinary(part);
+                    continue;
+                }
+                match part["type"].as_str().unwrap_or_default() {
+                    "text" => out.ordinary(&part.get("text").map(text).unwrap_or_default()),
+                    kind @ ("image" | "image_url" | "video" | "video_url" | "audio"
+                    | "audio_url" | "input_audio") => {
+                        let kind = kind.replace("_url", "").replace("input_", "");
+                        if media.enabled {
+                            out.control(&format!("<|begin_of_{kind}|>"))?;
+                            if kind != "audio" {
+                                out.control(&format!("<|{kind}|>"))?;
+                            }
+                            out.control(&format!("<|end_of_{kind}|>"))?;
+                            if let (Some(message_index), Some(content_part_index)) = (
+                                part["message_index"].as_u64(),
+                                part["content_part_index"].as_u64(),
+                            ) {
+                                media.order.push(MediaPartSource {
+                                    message_index: message_index as usize,
+                                    content_part_index: content_part_index as usize,
+                                });
+                            }
+                        } else {
+                            out.control(&format!("<reminder>You are unable to process this {kind} because you don't have multi-modal input ability. Try different methods.</reminder>"))?;
+                        }
+                    }
+                    _ => {}
+                }
+            }
+        }
+        _ => out.ordinary(&text(value)),
+    }
+    Ok(())
+}
+
+fn tool_definition(tool: &Value) -> Result<String> {
+    let Some(fields) = tool.as_object() else {
+        bail_chat_template_thrown!("Each tool must be a JSON object");
+    };
+    let mut out = String::from("{");
+    for (key, value) in fields {
+        if key == "strict" || key == "defer_loading" {
+            continue;
+        }
+        if out.len() > 1 {
+            out.push_str(", ");
+        }
+        write!(out, "\"{key}\": {}", json(value)?).unwrap();
+    }
+    out.push('}');
+    Ok(out)
+}
+
+fn unwrap_tool(tool: &Value) -> &Value {
+    tool.get("function").unwrap_or(tool)
+}
+
+fn tool_references(value: &Value) -> bool {
+    value
+        .as_array()
+        .and_then(|parts| parts.first())
+        .is_some_and(|part| part["type"] == "tool_reference")
+}
+
+fn reference_response(out: &mut Glm5PromptWriter<'_>, refs: &Value, tools: &[Value]) -> Result<()> {
+    if tools.is_empty() {
+        bail_chat_template_thrown!("Rendering a tool_reference requires tools");
+    }
+    out.control("<tool_response><tools>\n")?;
+    for reference in refs.as_array().into_iter().flatten() {
+        for tool in tools.iter().map(unwrap_tool) {
+            if tool.get("name") == reference.get("name") {
+                out.ordinary(&tool_definition(tool)?);
+                out.control("\n")?;
+            }
+        }
+    }
+    out.control("</tools></tool_response>")
+}
+
+fn response(
+    out: &mut Glm5PromptWriter<'_>,
+    value: &Value,
+    tools: &[Value],
+    media: &mut MediaContext,
+) -> Result<()> {
+    if tool_references(value) {
+        reference_response(out, value, tools)
+    } else {
+        out.control("<tool_response>")?;
+        visible_text(out, value, media)?;
+        out.control("</tool_response>")
+    }
+}
+
+fn bundled_outputs(value: &Value) -> Result<bool> {
+    if !truthy(value) || value.is_string() {
+        return Ok(false);
+    }
+    let Some(parts) = value.as_array() else {
+        bail_chat_template_thrown!("A tool message content must be a string or a list");
+    };
+    Ok(parts.first().is_some_and(|part| part.get("output").is_some()))
+}
+
+fn tool_response(
+    out: &mut Glm5PromptWriter<'_>,
+    message: &Value,
+    tools: &[Value],
+    media: &mut MediaContext,
+) -> Result<()> {
+    let value = &message["content"];
+    if tool_references(value) {
+        return reference_response(out, value, tools);
+    }
+    if bundled_outputs(value)? {
+        for part in value.as_array().unwrap() {
+            response(out, &part["output"], tools, media)?;
+        }
+        Ok(())
+    } else {
+        response(out, value, tools, media)
+    }
+}
+
+fn id(value: &Value) -> String {
+    ["tool_call_id", "id"]
+        .into_iter()
+        .filter_map(|key| value.get(key))
+        .find(|value| truthy(value))
+        .map(text)
+        .unwrap_or_default()
+}
+
+fn ordered_tool_block(
+    out: &mut Glm5PromptWriter<'_>,
+    block: &[Value],
+    assistant: Option<&Value>,
+    tools: &[Value],
+    media: &mut MediaContext,
+) -> Result<()> {
+    let calls = assistant
+        .filter(|message| message["role"] == "assistant")
+        .and_then(|message| message["tool_calls"].as_array())
+        .filter(|calls| !calls.is_empty());
+    if let Some(calls) = calls {
+        let mut can_sort = true;
+        let mut responses = HashMap::new();
+        for message in block {
+            if bundled_outputs(&message["content"])? {
+                for entry in message["content"].as_array().unwrap() {
+                    let key = id(entry);
+                    if key.is_empty() || responses.insert(key, &entry["output"]).is_some() {
+                        can_sort = false;
+                    }
+                }
+            } else {
+                let key = id(message);
+                if key.is_empty() || responses.insert(key, &message["content"]).is_some() {
+                    can_sort = false;
+                }
+            }
+        }
+        let call_ids: Vec<_> = calls.iter().map(id).collect();
+        let mut seen = HashSet::new();
+        for key in &call_ids {
+            if key.is_empty() || !seen.insert(key) {
+                can_sort = false;
+            }
+        }
+        if can_sort && responses.keys().all(|key| seen.contains(key)) {
+            for value in call_ids.iter().filter_map(|key| responses.get(key)) {
+                response(out, value, tools, media)?;
+            }
+            return Ok(());
+        }
+    }
+    for message in block {
+        tool_response(out, message, tools, media)?;
+    }
+    Ok(())
+}
+
+fn tool_call(out: &mut Glm5PromptWriter<'_>, call: &Value) -> Result<()> {
+    let call = call.get("function").filter(|value| truthy(value)).unwrap_or(call);
+    let Some(name) = call["name"].as_str() else {
+        bail_chat_template_thrown!("Each tool_call needs a string name");
+    };
+    let parsed;
+    let arguments = if let Some(raw) = call["arguments"].as_str() {
+        parsed = serde_json::from_str::<Value>(raw).map_err(|error| {
+            chat_template_thrown!("Invalid tool_call arguments: {}", error.as_report())
+        })?;
+        &parsed
+    } else {
+        &call["arguments"]
+    };
+    let Some(arguments) = arguments.as_object() else {
+        bail_chat_template_thrown!("tool_call arguments must be a JSON object");
+    };
+    out.control("<tool_call>")?;
+    out.ordinary(name);
+    for (key, value) in arguments {
+        out.control("<arg_key>")?;
+        out.ordinary(key);
+        out.control("</arg_key><arg_value>")?;
+        if let Some(value) = value.as_str() {
+            out.ordinary(value);
+        } else {
+            out.ordinary(&json(value)?);
+        }
+        out.control("</arg_value>")?;
+    }
+    out.control("</tool_call>")
+}
+
+#[cfg(test)]
+pub(super) fn render(
+    messages: &[Value],
+    tools: &[Value],
+    mode: GenerationPromptMode,
+    kwargs: &HashMap<String, Value>,
+) -> Result<String> {
+    let mut out = Glm5PromptWriter::new(None);
+    render_into(&mut out, messages, tools, mode, kwargs, false)?;
+    Ok(out.text().to_owned())
+}
+
+/// Render into `out`, returning media parts in placeholder order.
+pub(super) fn render_into(
+    out: &mut Glm5PromptWriter<'_>,
+    messages: &[Value],
+    tools: &[Value],
+    mode: GenerationPromptMode,
+    kwargs: &HashMap<String, Value>,
+    multimodal: bool,
+) -> Result<Vec<MediaPartSource>> {
+    let mut media = MediaContext {
+        enabled: multimodal,
+        order: Vec::new(),
+    };
+    render_inner(out, messages, tools, mode, kwargs, &mut media)?;
+    Ok(media.order)
+}
+
+fn render_inner(
+    out: &mut Glm5PromptWriter<'_>,
+    messages: &[Value],
+    tools: &[Value],
+    mode: GenerationPromptMode,
+    kwargs: &HashMap<String, Value>,
+    media: &mut MediaContext,
+) -> Result<()> {
+    let effort = match kwargs.get("reasoning_effort").and_then(Value::as_str) {
+        Some("low") => "Low",
+        Some("high") => "High",
+        _ => "Max",
+    };
+    let clear_thinking = kwargs.get("clear_thinking").is_some_and(truthy);
+    out.control(&format!("[gMASK]<sop><|system|>Reasoning Effort: {effort}"))?;
+    if !tools.is_empty() {
+        out.control(TOOLS_PREFIX)?;
+        for tool in tools.iter().map(unwrap_tool) {
+            if !truthy(&tool["defer_loading"]) {
+                out.ordinary(&tool_definition(tool)?);
+                out.control("\n")?;
+            }
+        }
+        out.control(TOOLS_SUFFIX)?;
+    }
+    let last_user = messages.iter().rposition(|message| message["role"] == "user");
+    let mut index = 0;
+    while index < messages.len() {
+        let message = &messages[index];
+        match message["role"].as_str().unwrap_or_default() {
+            "user" | "system" => {
+                out.control(&format!("<|{}|>", message["role"].as_str().unwrap()))?;
+                visible_text(out, &message["content"], media)?;
+            }
+            "assistant" => {
+                out.control("<|assistant|><think>")?;
+                let mut content = {
+                    let mut text = Glm5PromptWriter::new(None);
+                    visible_text(&mut text, &message["content"], media)?;
+                    text.text().to_owned()
+                };
+                let original = content.clone();
+                let mut reasoning = message["reasoning_content"].as_str().map(str::to_owned);
+                if reasoning.is_none() && content.contains("</think>") {
+                    reasoning = Some(
+                        content
+                            .split("</think>")
+                            .next()
+                            .unwrap()
+                            .rsplit("<think>")
+                            .next()
+                            .unwrap()
+                            .to_owned(),
+                    );
+                    content = content.rsplit("</think>").next().unwrap().to_owned();
+                }
+                if (!clear_thinking || last_user.is_none_or(|last| index > last))
+                    && let Some(reasoning) = reasoning
+                {
+                    out.ordinary(&reasoning);
+                }
+                out.control("</think>")?;
+                let trim = |ch: char| ch.is_whitespace() || ('\u{1c}'..='\u{1f}').contains(&ch);
+                if mode == GenerationPromptMode::ContinueFinalAssistant
+                    && index + 1 == messages.len()
+                {
+                    // Match HF continuation, which truncates the template output
+                    // at the end of the final content: tool calls are dropped,
+                    // trailing whitespace is trimmed, and content the template
+                    // rewrote is rejected.
+                    out.ordinary(content.trim_matches(trim));
+                    let original = original.trim_matches(trim);
+                    if !out.text().contains(original) {
+                        return Err(Error::ChatTemplate(format!(
+                            "continue_final_message is set but the final message does not appear \
+                             in the prompt after applying the chat template! This can happen if \
+                             the chat template deletes portions of the final message. Final \
+                             message to continue: {original}"
+                        )));
+                    }
+                    break;
+                }
+                out.ordinary(content.trim_matches(trim));
+                for call in message["tool_calls"].as_array().into_iter().flatten() {
+                    tool_call(out, call)?;
+                }
+            }
+            "tool" => {
+                let start = index;
+                while index + 1 < messages.len() && messages[index + 1]["role"] == "tool" {
+                    index += 1;
+                }
+                out.control("<|observation|>")?;
+                ordered_tool_block(
+                    out,
+                    &messages[start..=index],
+                    start.checked_sub(1).map(|i| &messages[i]),
+                    tools,
+                    media,
+                )?;
+            }
+            _ => {}
+        }
+        index += 1;
+    }
+    if mode == GenerationPromptMode::StartNewAssistant {
+        out.control("<|assistant|><think>")?;
+    }
+    Ok(())
+}
+
+// TODO: Expose tool_reference and bundled outputs in ChatContentPart.
+
+fn content(value: &ChatContent, message_index: usize) -> Value {
+    match value {
+        ChatContent::Text(text) => json!(text),
+        ChatContent::Parts(parts) => Value::Array(
+            parts
+                .iter()
+                .enumerate()
+                .map(|(content_part_index, part)| match part {
+                    ChatContentPart::Text { text } => json!({"type": "text", "text": text}),
+                    ChatContentPart::ImageUrl { .. } => json!({"type": "image", "message_index": message_index, "content_part_index": content_part_index}),
+                    ChatContentPart::VideoUrl { .. } => json!({"type": "video", "message_index": message_index, "content_part_index": content_part_index}),
+                    ChatContentPart::InputAudio { .. } | ChatContentPart::AudioUrl { .. } => {
+                        json!({"type": "audio", "message_index": message_index, "content_part_index": content_part_index})
+                    }
+                })
+                .collect(),
+        ),
+    }
+}
+
+pub(super) fn messages(messages: &[ChatMessage]) -> Vec<Value> {
+    messages.iter().enumerate().map(|(index, message)| match message {
+        ChatMessage::System { content: value } => json!({"role": "system", "content": content(value, index)}),
+        ChatMessage::Developer { content: value, .. } => json!({"role": "developer", "content": content(value, index)}),
+        ChatMessage::User { content: value } => json!({"role": "user", "content": content(value, index)}),
+        ChatMessage::Custom { role, content: value } => json!({"role": role, "content": content(value, index)}),
+        ChatMessage::ToolResponse { content: value, tool_call_id } => {
+            json!({"role": "tool", "content": content(value, index), "tool_call_id": tool_call_id})
+        }
+        ChatMessage::Assistant { content } => {
+            let mut message = json!({"role": "assistant", "content": content.text()});
+            if content.iter().any(|block| matches!(block, AssistantContentBlock::Reasoning { .. })) {
+                message["reasoning_content"] = json!(content.iter().filter_map(|block| match block {
+                    AssistantContentBlock::Reasoning { text } => Some(text.as_str()),
+                    _ => None,
+                }).collect::<String>());
+            }
+            if content.has_tool_calls() {
+                message["tool_calls"] = Value::Array(content.tool_calls().map(|call| {
+                    json!({"id": call.id, "function": {"name": call.name, "arguments": call.arguments}})
+                }).collect());
+            }
+            message
+        }
+    }).collect()
+}
+
+pub(super) fn tools(tools: &[ChatTool]) -> Vec<Value> {
+    tools
+        .iter()
+        .map(|tool| {
+            json!({"function": {
+                "name": tool.name,
+                "description": tool.description,
+                "parameters": tool.parameters,
+                "defer_loading": tool.defer_loading,
+            }})
+        })
+        .collect()
+}
+
+struct PythonFormatter;
+
+impl Formatter for PythonFormatter {
+    fn begin_array_value<W: ?Sized + Write>(
+        &mut self,
+        writer: &mut W,
+        first: bool,
+    ) -> io::Result<()> {
+        if !first {
+            writer.write_all(b", ")?;
+        }
+        Ok(())
+    }
+
+    fn begin_object_key<W: ?Sized + Write>(
+        &mut self,
+        writer: &mut W,
+        first: bool,
+    ) -> io::Result<()> {
+        if !first {
+            writer.write_all(b", ")?;
+        }
+        Ok(())
+    }
+
+    fn begin_object_value<W: ?Sized + Write>(&mut self, writer: &mut W) -> io::Result<()> {
+        writer.write_all(b": ")
+    }
+
+    fn write_f64<W: ?Sized + Write>(&mut self, writer: &mut W, value: f64) -> io::Result<()> {
+        let magnitude = value.abs();
+        if magnitude != 0.0 && !(1e-4..1e16).contains(&magnitude) {
+            let scientific = format!("{value:e}");
+            let (mantissa, exponent) = scientific.split_once('e').unwrap();
+            let exponent: i32 = exponent.parse().unwrap();
+            write!(writer, "{mantissa}e{exponent:+03}")
+        } else {
+            let decimal = value.to_string();
+            writer.write_all(decimal.as_bytes())?;
+            if !decimal.contains('.') {
+                writer.write_all(b".0")?;
+            }
+            Ok(())
+        }
+    }
+}
+
+fn json(value: &Value) -> Result<String> {
+    let mut bytes = Vec::new();
+    value
+        .serialize(&mut Serializer::with_formatter(&mut bytes, PythonFormatter))
+        .map_err(|error| {
+            chat_template_thrown!("GLM JSON serialization failed: {}", error.as_report())
+        })?;
+    Ok(String::from_utf8(bytes).expect("JSON serializer must write UTF-8"))
+}
diff --git a/rust/src/chat/src/renderer/glm5_next/fixtures/glm_5_3.jinja b/rust/src/chat/src/renderer/glm5_next/fixtures/glm_5_3.jinja
new file mode 100644
index 000000000000..55fd3945292b
--- /dev/null
+++ b/rust/src/chat/src/renderer/glm5_next/fixtures/glm_5_3.jinja
@@ -0,0 +1,251 @@
+[gMASK]<sop>
+{%- set effective_reasoning_effort = reasoning_effort if reasoning_effort is defined and reasoning_effort in ['low', 'high'] else 'max' -%}
+{%- if effective_reasoning_effort is not none -%}<|system|>Reasoning Effort: {{ effective_reasoning_effort | capitalize }}{%- endif -%}
+{%- set clear_thinking = clear_thinking if clear_thinking is defined else false -%}
+{%- if tools -%}
+{%- macro tool_to_json(tool) -%}
+    {%- set ns_tool = namespace(first=true) -%}
+    {{ '{' -}}
+    {%- for k, v in tool.items() -%}
+        {%- if k != 'defer_loading' and k != 'strict' -%}
+            {%- if not ns_tool.first -%}{{- ', ' -}}{%- endif -%}
+            {%- set ns_tool.first = false -%}
+            "{{ k }}": {{ v | tojson(ensure_ascii=False) }}
+        {%- endif -%}
+    {%- endfor -%}
+    {{- '}' -}}
+{%- endmacro -%}
+{%- macro tool_references_to_response(refs) -%}
+    {{- '<tool_response><tools>\n' -}}
+    {%- for tr in refs -%}
+        {%- for tool in tools -%}
+            {%- if 'function' in tool -%}
+                {%- set tool = tool['function'] -%}
+            {%- endif -%}
+            {%- if tool.name == tr.name -%}
+                {{- tool_to_json(tool) + '\n' -}}
+            {%- endif -%}
+        {%- endfor -%}
+    {%- endfor -%}
+    {{- '</tools></tool_response>' -}}
+{%- endmacro -%}
+<|system|>
+# Tools
+
+You may call one or more functions to assist with the user query.
+
+You are provided with function signatures within <tools></tools> XML tags:
+<tools>
+{% for tool in tools %}
+{%- if 'function' in tool -%}
+    {%- set tool = tool['function'] -%}
+{%- endif -%}
+{% if tool.defer_loading is not defined or not tool.defer_loading %}
+{{ tool_to_json(tool) }}
+{% endif %}
+{% endfor %}
+</tools>
+
+For each function call, output the function name and arguments within the following XML format:
+<tool_call>{function-name}<arg_key>{arg-key-1}</arg_key><arg_value>{arg-value-1}</arg_value><arg_key>{arg-key-2}</arg_key><arg_value>{arg-value-2}</arg_value>...</tool_call>{%- endif -%}
+{%- macro visible_text(content) -%}
+    {%- if content is string -%}
+        {{- content }}
+    {%- elif content is iterable and content is not mapping -%}
+        {%- for item in content -%}
+            {%- if item is mapping and item.type == 'text' -%}
+                {{- item.text }}
+            {%- elif item is string -%}
+                {{- item }}
+            {%- elif item is mapping and item.type in ['image', 'image_url', 'video', 'video_url', 'audio', 'audio_url', 'input_audio'] -%}
+                {%- set media_type = item.type | replace('_url', '') | replace('input_', '') -%}
+                {{- "<reminder>You are unable to process this " ~ media_type ~ " because you don't have multi-modal input ability. Try different methods.</reminder>" }}
+            {%- endif -%}
+        {%- endfor -%}
+    {%- elif content is not none -%}
+        {{- content }}
+    {%- endif -%}
+{%- endmacro -%}
+{%- macro tool_response(text) -%}
+{{- '<tool_response>' + text + '</tool_response>' -}}
+{%- endmacro -%}
+{%- macro render_tool_response(m) -%}
+{%- if m.content is string -%}
+    {{- tool_response(m.content) -}}
+{%- elif m.content and m.content is not mapping and m.content.0.type == "tool_reference" -%}
+    {{- tool_references_to_response(m.content) -}}
+{%- elif is_list_of_outputs(m) -%}
+    {%- for tr in m.content -%}
+        {%- if tr.output is iterable and tr.output is not string and tr.output is not mapping and tr.output and tr.output.0.type == "tool_reference" -%}
+            {{- tool_references_to_response(tr.output) -}}
+        {%- else -%}
+            {{- tool_response(visible_text(tr.output)) -}}
+        {%- endif -%}
+    {%- endfor -%}
+{%- else -%}
+    {{- tool_response(visible_text(m.content)) -}}
+{%- endif -%}
+{%- endmacro -%}
+{%- macro id_of(obj) -%}
+    {%- if obj.tool_call_id -%}
+        {{- obj.tool_call_id -}}
+    {%- elif obj.id -%}
+        {{- obj.id -}}
+    {%- endif -%}
+{%- endmacro -%}
+{%- macro is_list_of_outputs(m) -%}
+    {%- if m.content and m.content.0.output is defined -%}1{%- endif -%}
+{%- endmacro -%}
+{%- macro has_dup_tool_result_id(lo, hi, target) -%}
+    {%- set ns_cnt = namespace(n=0) -%}
+    {%- for k in range(lo, hi + 1) -%}
+        {%- set m = messages[k] -%}
+        {%- if is_list_of_outputs(m) -%}
+            {%- for entry in m.content -%}
+                {%- if id_of(entry) == target -%}
+                    {%- set ns_cnt.n = ns_cnt.n + 1 -%}
+                {%- endif -%}
+            {%- endfor -%}
+        {%- elif id_of(m) == target -%}
+            {%- set ns_cnt.n = ns_cnt.n + 1 -%}
+        {%- endif -%}
+        {%- if ns_cnt.n > 1 -%}{%- break -%}{%- endif -%}
+    {%- endfor -%}
+    {%- if ns_cnt.n > 1 -%}1{%- endif -%}
+{%- endmacro -%}
+{%- macro tc_id_exists(tcs, target) -%}
+    {%- set ns_f = namespace(found=false) -%}
+    {%- for tc in tcs -%}
+        {%- if id_of(tc) == target -%}
+            {%- set ns_f.found = true -%}
+            {%- break -%}
+        {%- endif -%}
+    {%- endfor -%}
+    {%- if ns_f.found -%}1{%- endif -%}
+{%- endmacro -%}
+{%- set ns = namespace(last_user_index=-1) -%}
+{%- for m in messages %}
+    {%- if m.role == 'user' %}
+        {%- set ns.last_user_index = loop.index0 -%}
+    {%- endif %}
+{%- endfor %}
+{%- for m in messages -%}
+{%- if m.role == 'user' -%}<|user|>{{ visible_text(m.content) }}
+{%- elif m.role == 'assistant' -%}
+<|assistant|>
+{%- set content = visible_text(m.content) %}
+{%- if m.reasoning_content is string %}
+    {%- set reasoning_content = m.reasoning_content %}
+{%- elif '</think>' in content %}
+    {%- set reasoning_content = content.split('</think>')[0].split('<think>')[-1] %}
+    {%- set content = content.split('</think>')[-1] %}
+{%- endif %}
+{%- if (not clear_thinking or loop.index0 > ns.last_user_index) and reasoning_content is defined -%}
+{{ '<think>' + reasoning_content +  '</think>'}}
+{%- else -%}
+{{ '<think></think>' }}
+{%- endif -%}
+{%- if content.strip() -%}
+{{ content.strip() }}
+{%- endif -%}
+{% if m.tool_calls %}
+{% for tc in m.tool_calls %}
+{%- if tc.function %}
+    {%- set tc = tc.function %}
+{%- endif %}
+{{- '<tool_call>' + tc.name -}}
+{% set _args = tc.arguments %}{% for k, v in _args.items() %}<arg_key>{{ k }}</arg_key><arg_value>{{ v | tojson(ensure_ascii=False) if v is not string else v }}</arg_value>{% endfor %}</tool_call>{% endfor %}
+{% endif %}
+{%- elif m.role == 'tool' -%}
+{%- if loop.first or (messages[loop.index0 - 1].role != "tool") %}
+    {{- '<|observation|>' -}}
+    {%- set block_start = loop.index0 -%}
+    {%- set ns_blk = namespace(end=block_start) -%}
+    {%- for j in range(block_start, messages|length) -%}
+        {%- if messages[j].role == 'tool' -%}
+            {%- set ns_blk.end = j -%}
+        {%- else -%}
+            {%- break -%}
+        {%- endif -%}
+    {%- endfor -%}
+    {%- set ns_a = namespace(tool_calls=none) -%}
+    {%- if block_start > 0 and messages[block_start - 1].role == 'assistant' and messages[block_start - 1].tool_calls -%}
+        {%- set ns_a.tool_calls = messages[block_start - 1].tool_calls -%}
+    {%- endif -%}
+    {%- set ns_chk = namespace(can_sort=true, res=[], mk=[], ek=[], tcs=[]) -%}
+    {%- for k in range(block_start, ns_blk.end + 1) -%}
+        {%- set m = messages[k] -%}
+        {%- if is_list_of_outputs(m) -%}
+            {%- for entry in m.content -%}
+                {%- set ns_chk.res = ns_chk.res + [id_of(entry)] -%}
+                {%- set ns_chk.mk = ns_chk.mk + [k] -%}
+                {%- set ns_chk.ek = ns_chk.ek + [loop.index0] -%}
+            {%- endfor -%}
+        {%- else -%}
+            {%- set ns_chk.res = ns_chk.res + [id_of(m)] -%}
+            {%- set ns_chk.mk = ns_chk.mk + [k] -%}
+            {%- set ns_chk.ek = ns_chk.ek + [-1] -%}
+        {%- endif -%}
+    {%- endfor -%}
+    {%- if not ns_a.tool_calls -%}
+        {%- set ns_chk.can_sort = false -%}
+    {%- else -%}
+        {%- for tc in ns_a.tool_calls -%}
+            {%- set ns_chk.tcs = ns_chk.tcs + [id_of(tc)] -%}
+        {%- endfor -%}
+        {%- if '' in ns_chk.res or '' in ns_chk.tcs -%}
+            {%- set ns_chk.can_sort = false -%}
+        {%- elif ns_chk.res | unique(case_sensitive=true) | list | length != ns_chk.res | length -%}
+            {%- set ns_chk.can_sort = false -%}
+        {%- elif ns_chk.tcs | unique(case_sensitive=true) | list | length != ns_chk.tcs | length -%}
+            {%- set ns_chk.can_sort = false -%}
+        {%- else -%}
+            {%- for eid in ns_chk.res -%}
+                {%- if eid not in ns_chk.tcs -%}
+                    {%- set ns_chk.can_sort = false -%}
+                    {%- break -%}
+                {%- endif -%}
+            {%- endfor -%}
+        {%- endif -%}
+    {%- endif -%}
+    {%- if ns_chk.can_sort -%}
+        {%- set ns_ord = namespace(res=[], tcs=[], j=0, pairs=[]) -%}
+        {%- for eid in ns_chk.res -%}
+            {%- set ns_ord.res = ns_ord.res + [{'id': eid, 'p': loop.index0}] -%}
+        {%- endfor -%}
+        {%- for tc_id in ns_chk.tcs -%}
+            {%- set ns_ord.tcs = ns_ord.tcs + [{'id': tc_id, 'i': loop.index0}] -%}
+        {%- endfor -%}
+        {%- set ns_ord.res = ns_ord.res | sort(attribute='id', case_sensitive=true) -%}
+        {%- for tc in ns_ord.tcs | sort(attribute='id', case_sensitive=true) -%}
+            {%- if ns_ord.j < ns_ord.res | length and ns_ord.res[ns_ord.j].id == tc.id -%}
+                {%- set ns_ord.pairs = ns_ord.pairs + [{'i': tc.i, 'p': ns_ord.res[ns_ord.j].p}] -%}
+                {%- set ns_ord.j = ns_ord.j + 1 -%}
+            {%- endif -%}
+        {%- endfor -%}
+        {%- for pair in ns_ord.pairs | sort(attribute='i') -%}
+            {%- set m = messages[ns_chk.mk[pair.p]] -%}
+            {%- if ns_chk.ek[pair.p] >= 0 -%}
+                {%- set entry = m.content[ns_chk.ek[pair.p]] -%}
+                {%- if entry.output is iterable and entry.output is not string and entry.output is not mapping and entry.output and entry.output.0.type == "tool_reference" -%}
+                    {{- tool_references_to_response(entry.output) -}}
+                {%- else -%}
+                    {{- tool_response(visible_text(entry.output)) -}}
+                {%- endif -%}
+            {%- else -%}
+                {{- render_tool_response(m) -}}
+            {%- endif -%}
+        {%- endfor -%}
+    {%- else -%}
+        {%- for k in range(block_start, ns_blk.end + 1) -%}
+            {{- render_tool_response(messages[k]) -}}
+        {%- endfor -%}
+    {%- endif -%}
+{% endif -%}
+{%- elif m.role == 'system' -%}
+<|system|>{{ visible_text(m.content) }}
+{%- endif -%}
+{%- endfor -%}
+{%- if add_generation_prompt -%}
+    <|assistant|>{{- '<think>' -}}
+{%- endif -%}
diff --git a/rust/src/chat/src/renderer/glm5_next/fixtures/glm_5_3_flash.jinja b/rust/src/chat/src/renderer/glm5_next/fixtures/glm_5_3_flash.jinja
new file mode 100644
index 000000000000..14180451f166
--- /dev/null
+++ b/rust/src/chat/src/renderer/glm5_next/fixtures/glm_5_3_flash.jinja
@@ -0,0 +1,257 @@
+[gMASK]<sop>
+{%- set effective_reasoning_effort = reasoning_effort if reasoning_effort is defined and reasoning_effort in ['low', 'high'] else 'max' -%}
+{%- if effective_reasoning_effort is not none -%}<|system|>Reasoning Effort: {{ effective_reasoning_effort | capitalize }}{%- endif -%}
+{%- set clear_thinking = clear_thinking if clear_thinking is defined else false -%}
+{%- if tools -%}
+{%- macro tool_to_json(tool) -%}
+    {%- set ns_tool = namespace(first=true) -%}
+    {{ '{' -}}
+    {%- for k, v in tool.items() -%}
+        {%- if k != 'defer_loading' and k != 'strict' -%}
+            {%- if not ns_tool.first -%}{{- ', ' -}}{%- endif -%}
+            {%- set ns_tool.first = false -%}
+            "{{ k }}": {{ v | tojson(ensure_ascii=False) }}
+        {%- endif -%}
+    {%- endfor -%}
+    {{- '}' -}}
+{%- endmacro -%}
+{%- macro tool_references_to_response(refs) -%}
+    {{- '<tool_response><tools>\n' -}}
+    {%- for tr in refs -%}
+        {%- for tool in tools -%}
+            {%- if 'function' in tool -%}
+                {%- set tool = tool['function'] -%}
+            {%- endif -%}
+            {%- if tool.name == tr.name -%}
+                {{- tool_to_json(tool) + '\n' -}}
+            {%- endif -%}
+        {%- endfor -%}
+    {%- endfor -%}
+    {{- '</tools></tool_response>' -}}
+{%- endmacro -%}
+<|system|>
+# Tools
+
+You may call one or more functions to assist with the user query.
+
+You are provided with function signatures within <tools></tools> XML tags:
+<tools>
+{% for tool in tools %}
+{%- if 'function' in tool -%}
+    {%- set tool = tool['function'] -%}
+{%- endif -%}
+{% if tool.defer_loading is not defined or not tool.defer_loading %}
+{{ tool_to_json(tool) }}
+{% endif %}
+{% endfor %}
+</tools>
+
+For each function call, output the function name and arguments within the following XML format:
+<tool_call>{function-name}<arg_key>{arg-key-1}</arg_key><arg_value>{arg-value-1}</arg_value><arg_key>{arg-key-2}</arg_key><arg_value>{arg-value-2}</arg_value>...</tool_call>{%- endif -%}
+{%- macro emit_image() -%}<|begin_of_image|><|image|><|end_of_image|>{%- endmacro -%}
+{%- macro emit_video() -%}<|begin_of_video|><|video|><|end_of_video|>{%- endmacro -%}
+{%- macro emit_audio() -%}<|begin_of_audio|><|end_of_audio|>{%- endmacro -%}
+{%- macro visible_text(content) -%}
+    {%- if content is string -%}
+        {{- content -}}
+    {%- elif content is iterable and content is not mapping -%}
+        {%- for item in content -%}
+            {%- if item is mapping and item.type == 'text' -%}
+                {{- item.text -}}
+            {%- elif item is string -%}
+                {{- item -}}
+            {%- elif item is mapping and item.type in ['image', 'image_url'] -%}
+                {{- emit_image() -}}
+            {%- elif item is mapping and item.type in ['video', 'video_url'] -%}
+                {{- emit_video() -}}
+            {%- elif item is mapping and item.type in ['audio', 'audio_url', 'input_audio'] -%}
+                {{- emit_audio() -}}
+            {%- endif -%}
+        {%- endfor -%}
+    {%- elif content is not none -%}
+        {{- content -}}
+    {%- endif -%}
+{%- endmacro -%}
+{%- macro tool_response(text) -%}
+{{- '<tool_response>' + text + '</tool_response>' -}}
+{%- endmacro -%}
+{%- macro render_tool_response(m) -%}
+{%- if m.content is string -%}
+    {{- tool_response(m.content) -}}
+{%- elif m.content and m.content is not mapping and m.content.0.type == "tool_reference" -%}
+    {{- tool_references_to_response(m.content) -}}
+{%- elif is_list_of_outputs(m) -%}
+    {%- for tr in m.content -%}
+        {%- if tr.output is iterable and tr.output is not string and tr.output is not mapping and tr.output and tr.output.0.type == "tool_reference" -%}
+            {{- tool_references_to_response(tr.output) -}}
+        {%- else -%}
+            {{- tool_response(visible_text(tr.output)) -}}
+        {%- endif -%}
+    {%- endfor -%}
+{%- else -%}
+    {{- tool_response(visible_text(m.content)) -}}
+{%- endif -%}
+{%- endmacro -%}
+{%- macro id_of(obj) -%}
+    {%- if obj.tool_call_id -%}
+        {{- obj.tool_call_id -}}
+    {%- elif obj.id -%}
+        {{- obj.id -}}
+    {%- endif -%}
+{%- endmacro -%}
+{%- macro is_list_of_outputs(m) -%}
+    {%- if m.content and m.content.0.output is defined -%}1{%- endif -%}
+{%- endmacro -%}
+{%- macro has_dup_tool_result_id(lo, hi, target) -%}
+    {%- set ns_cnt = namespace(n=0) -%}
+    {%- for k in range(lo, hi + 1) -%}
+        {%- set m = messages[k] -%}
+        {%- if is_list_of_outputs(m) -%}
+            {%- for entry in m.content -%}
+                {%- if id_of(entry) == target -%}
+                    {%- set ns_cnt.n = ns_cnt.n + 1 -%}
+                {%- endif -%}
+            {%- endfor -%}
+        {%- elif id_of(m) == target -%}
+            {%- set ns_cnt.n = ns_cnt.n + 1 -%}
+        {%- endif -%}
+        {%- if ns_cnt.n > 1 -%}{%- break -%}{%- endif -%}
+    {%- endfor -%}
+    {%- if ns_cnt.n > 1 -%}1{%- endif -%}
+{%- endmacro -%}
+{%- macro tc_id_exists(tcs, target) -%}
+    {%- set ns_f = namespace(found=false) -%}
+    {%- for tc in tcs -%}
+        {%- if id_of(tc) == target -%}
+            {%- set ns_f.found = true -%}
+            {%- break -%}
+        {%- endif -%}
+    {%- endfor -%}
+    {%- if ns_f.found -%}1{%- endif -%}
+{%- endmacro -%}
+{%- set ns = namespace(last_user_index=-1) -%}
+{%- for m in messages %}
+    {%- if m.role == 'user' %}
+        {%- set ns.last_user_index = loop.index0 -%}
+    {%- endif %}
+{%- endfor %}
+{%- for m in messages -%}
+{%- if m.role == 'user' -%}<|user|>{{ visible_text(m.content) }}
+{%- elif m.role == 'assistant' -%}
+<|assistant|>
+{%- set content = visible_text(m.content) %}
+{%- if m.reasoning_content is string %}
+    {%- set reasoning_content = m.reasoning_content %}
+{%- elif '</think>' in content %}
+    {%- set reasoning_content = content.split('</think>')[0].split('<think>')[-1] %}
+    {%- set content = content.split('</think>')[-1] %}
+{%- endif %}
+{%- if (not clear_thinking or loop.index0 > ns.last_user_index) and reasoning_content is defined -%}
+{{ '<think>' + reasoning_content +  '</think>'}}
+{%- else -%}
+{{ '<think></think>' }}
+{%- endif -%}
+{%- if content.strip() -%}
+{{ content.strip() }}
+{%- endif -%}
+{% if m.tool_calls %}
+{% for tc in m.tool_calls %}
+{%- if tc.function %}
+    {%- set tc = tc.function %}
+{%- endif %}
+{{- '<tool_call>' ~ tc.name -}}
+{% set _args = tc.arguments %}{% for k, v in _args.items() %}<arg_key>{{ k }}</arg_key><arg_value>{{ v | tojson(ensure_ascii=False) if v is not string else v }}</arg_value>{% endfor %}</tool_call>{% endfor %}
+{% endif %}
+{%- elif m.role == 'tool' -%}
+{%- if loop.first or (messages[loop.index0 - 1].role != "tool") %}
+    {{- '<|observation|>' -}}
+    {%- set block_start = loop.index0 -%}
+    {%- set ns_blk = namespace(end=block_start) -%}
+    {%- for j in range(block_start, messages|length) -%}
+        {%- if messages[j].role == 'tool' -%}
+            {%- set ns_blk.end = j -%}
+        {%- else -%}
+            {%- break -%}
+        {%- endif -%}
+    {%- endfor -%}
+    {%- set ns_a = namespace(tool_calls=none) -%}
+    {%- if block_start > 0 and messages[block_start - 1].role == 'assistant' and messages[block_start - 1].tool_calls -%}
+        {%- set ns_a.tool_calls = messages[block_start - 1].tool_calls -%}
+    {%- endif -%}
+    {%- set ns_chk = namespace(can_sort=true, res=[], mk=[], ek=[], tcs=[]) -%}
+    {%- for k in range(block_start, ns_blk.end + 1) -%}
+        {%- set m = messages[k] -%}
+        {%- if is_list_of_outputs(m) -%}
+            {%- for entry in m.content -%}
+                {%- set ns_chk.res = ns_chk.res + [id_of(entry)] -%}
+                {%- set ns_chk.mk = ns_chk.mk + [k] -%}
+                {%- set ns_chk.ek = ns_chk.ek + [loop.index0] -%}
+            {%- endfor -%}
+        {%- else -%}
+            {%- set ns_chk.res = ns_chk.res + [id_of(m)] -%}
+            {%- set ns_chk.mk = ns_chk.mk + [k] -%}
+            {%- set ns_chk.ek = ns_chk.ek + [-1] -%}
+        {%- endif -%}
+    {%- endfor -%}
+    {%- if not ns_a.tool_calls -%}
+        {%- set ns_chk.can_sort = false -%}
+    {%- else -%}
+        {%- for tc in ns_a.tool_calls -%}
+            {%- set ns_chk.tcs = ns_chk.tcs + [id_of(tc)] -%}
+        {%- endfor -%}
+        {%- if '' in ns_chk.res or '' in ns_chk.tcs -%}
+            {%- set ns_chk.can_sort = false -%}
+        {%- elif ns_chk.res | unique(case_sensitive=true) | list | length != ns_chk.res | length -%}
+            {%- set ns_chk.can_sort = false -%}
+        {%- elif ns_chk.tcs | unique(case_sensitive=true) | list | length != ns_chk.tcs | length -%}
+            {%- set ns_chk.can_sort = false -%}
+        {%- else -%}
+            {%- for eid in ns_chk.res -%}
+                {%- if eid not in ns_chk.tcs -%}
+                    {%- set ns_chk.can_sort = false -%}
+                    {%- break -%}
+                {%- endif -%}
+            {%- endfor -%}
+        {%- endif -%}
+    {%- endif -%}
+    {%- if ns_chk.can_sort -%}
+        {%- set ns_ord = namespace(res=[], tcs=[], j=0, pairs=[]) -%}
+        {%- for eid in ns_chk.res -%}
+            {%- set ns_ord.res = ns_ord.res + [{'id': eid, 'p': loop.index0}] -%}
+        {%- endfor -%}
+        {%- for tc_id in ns_chk.tcs -%}
+            {%- set ns_ord.tcs = ns_ord.tcs + [{'id': tc_id, 'i': loop.index0}] -%}
+        {%- endfor -%}
+        {%- set ns_ord.res = ns_ord.res | sort(attribute='id', case_sensitive=true) -%}
+        {%- for tc in ns_ord.tcs | sort(attribute='id', case_sensitive=true) -%}
+            {%- if ns_ord.j < ns_ord.res | length and ns_ord.res[ns_ord.j].id == tc.id -%}
+                {%- set ns_ord.pairs = ns_ord.pairs + [{'i': tc.i, 'p': ns_ord.res[ns_ord.j].p}] -%}
+                {%- set ns_ord.j = ns_ord.j + 1 -%}
+            {%- endif -%}
+        {%- endfor -%}
+        {%- for pair in ns_ord.pairs | sort(attribute='i') -%}
+            {%- set m = messages[ns_chk.mk[pair.p]] -%}
+            {%- if ns_chk.ek[pair.p] >= 0 -%}
+                {%- set entry = m.content[ns_chk.ek[pair.p]] -%}
+                {%- if entry.output is iterable and entry.output is not string and entry.output is not mapping and entry.output and entry.output.0.type == "tool_reference" -%}
+                    {{- tool_references_to_response(entry.output) -}}
+                {%- else -%}
+                    {{- tool_response(visible_text(entry.output)) -}}
+                {%- endif -%}
+            {%- else -%}
+                {{- render_tool_response(m) -}}
+            {%- endif -%}
+        {%- endfor -%}
+    {%- else -%}
+        {%- for k in range(block_start, ns_blk.end + 1) -%}
+            {{- render_tool_response(messages[k]) -}}
+        {%- endfor -%}
+    {%- endif -%}
+{% endif -%}
+{%- elif m.role == 'system' -%}
+<|system|>{{ visible_text(m.content) }}
+{%- endif -%}
+{%- endfor -%}
+{%- if add_generation_prompt -%}
+    <|assistant|>{{- '<think>' -}}
+{%- endif -%}
diff --git a/rust/src/chat/src/renderer/glm5_next/fixtures/test_cases.json b/rust/src/chat/src/renderer/glm5_next/fixtures/test_cases.json
new file mode 100644
index 000000000000..ac18c322edad
--- /dev/null
+++ b/rust/src/chat/src/renderer/glm5_next/fixtures/test_cases.json
@@ -0,0 +1,3001 @@
+{
+  "cases": [
+    {
+      "name": "basic",
+      "messages": [
+        {
+          "role": "user",
+          "content": "你好，GLM！"
+        }
+      ],
+      "tools": [],
+      "add_generation_prompt": true,
+      "kwargs": {},
+      "expected": "[gMASK]<sop><|system|>Reasoning Effort: Max<|user|>你好，GLM！<|assistant|><think>"
+    },
+    {
+      "name": "no_generation_prompt",
+      "messages": [
+        {
+          "role": "user",
+          "content": "你好，GLM！"
+        }
+      ],
+      "tools": [],
+      "add_generation_prompt": false,
+      "kwargs": {},
+      "expected": "[gMASK]<sop><|system|>Reasoning Effort: Max<|user|>你好，GLM！"
+    },
+    {
+      "name": "effort_None",
+      "messages": [
+        {
+          "role": "user",
+          "content": "你好，GLM！"
+        }
+      ],
+      "tools": [],
+      "add_generation_prompt": true,
+      "kwargs": {
+        "reasoning_effort": null
+      },
+      "expected": "[gMASK]<sop><|system|>Reasoning Effort: Max<|user|>你好，GLM！<|assistant|><think>"
+    },
+    {
+      "name": "effort_low",
+      "messages": [
+        {
+          "role": "user",
+          "content": "你好，GLM！"
+        }
+      ],
+      "tools": [],
+      "add_generation_prompt": true,
+      "kwargs": {
+        "reasoning_effort": "low"
+      },
+      "expected": "[gMASK]<sop><|system|>Reasoning Effort: Low<|user|>你好，GLM！<|assistant|><think>"
+    },
+    {
+      "name": "effort_high",
+      "messages": [
+        {
+          "role": "user",
+          "content": "你好，GLM！"
+        }
+      ],
+      "tools": [],
+      "add_generation_prompt": true,
+      "kwargs": {
+        "reasoning_effort": "high"
+      },
+      "expected": "[gMASK]<sop><|system|>Reasoning Effort: High<|user|>你好，GLM！<|assistant|><think>"
+    },
+    {
+      "name": "effort_medium",
+      "messages": [
+        {
+          "role": "user",
+          "content": "你好，GLM！"
+        }
+      ],
+      "tools": [],
+      "add_generation_prompt": true,
+      "kwargs": {
+        "reasoning_effort": "medium"
+      },
+      "expected": "[gMASK]<sop><|system|>Reasoning Effort: Max<|user|>你好，GLM！<|assistant|><think>"
+    },
+    {
+      "name": "effort_none",
+      "messages": [
+        {
+          "role": "user",
+          "content": "你好，GLM！"
+        }
+      ],
+      "tools": [],
+      "add_generation_prompt": true,
+      "kwargs": {
+        "reasoning_effort": "none"
+      },
+      "expected": "[gMASK]<sop><|system|>Reasoning Effort: Max<|user|>你好，GLM！<|assistant|><think>"
+    },
+    {
+      "name": "effort_max",
+      "messages": [
+        {
+          "role": "user",
+          "content": "你好，GLM！"
+        }
+      ],
+      "tools": [],
+      "add_generation_prompt": true,
+      "kwargs": {
+        "reasoning_effort": "max"
+      },
+      "expected": "[gMASK]<sop><|system|>Reasoning Effort: Max<|user|>你好，GLM！<|assistant|><think>"
+    },
+    {
+      "name": "effort_42",
+      "messages": [
+        {
+          "role": "user",
+          "content": "你好，GLM！"
+        }
+      ],
+      "tools": [],
+      "add_generation_prompt": true,
+      "kwargs": {
+        "reasoning_effort": 42
+      },
+      "expected": "[gMASK]<sop><|system|>Reasoning Effort: Max<|user|>你好，GLM！<|assistant|><think>"
+    },
+    {
+      "name": "effort_False",
+      "messages": [
+        {
+          "role": "user",
+          "content": "你好，GLM！"
+        }
+      ],
+      "tools": [],
+      "add_generation_prompt": true,
+      "kwargs": {
+        "reasoning_effort": false
+      },
+      "expected": "[gMASK]<sop><|system|>Reasoning Effort: Max<|user|>你好，GLM！<|assistant|><think>"
+    },
+    {
+      "name": "system_developer_custom",
+      "messages": [
+        {
+          "role": "system",
+          "content": "系统规则\n"
+        },
+        {
+          "role": "developer",
+          "content": "Ignored by this reference"
+        },
+        {
+          "role": "custom",
+          "content": "Ignored too"
+        },
+        {
+          "role": "user",
+          "content": "你好，GLM！"
+        }
+      ],
+      "tools": [],
+      "add_generation_prompt": true,
+      "kwargs": {},
+      "expected": "[gMASK]<sop><|system|>Reasoning Effort: Max<|system|>系统规则\n<|user|>你好，GLM！<|assistant|><think>"
+    },
+    {
+      "name": "content_parts",
+      "messages": [
+        {
+          "role": "user",
+          "content": [
+            {
+              "type": "text",
+              "text": "a<&>\n"
+            },
+            "b",
+            {
+              "type": "unknown"
+            },
+            {
+              "type": "image"
+            },
+            {
+              "type": "image_url"
+            },
+            {
+              "type": "video"
+            },
+            {
+              "type": "video_url"
+            },
+            {
+              "type": "audio"
+            },
+            {
+              "type": "audio_url"
+            },
+            {
+              "type": "input_audio"
+            }
+          ]
+        }
+      ],
+      "tools": [],
+      "add_generation_prompt": true,
+      "kwargs": {},
+      "expected": "[gMASK]<sop><|system|>Reasoning Effort: Max<|user|>a<&>\nb<reminder>You are unable to process this image because you don't have multi-modal input ability. Try different methods.</reminder><reminder>You are unable to process this image because you don't have multi-modal input ability. Try different methods.</reminder><reminder>You are unable to process this video because you don't have multi-modal input ability. Try different methods.</reminder><reminder>You are unable to process this video because you don't have multi-modal input ability. Try different methods.</reminder><reminder>You are unable to process this audio because you don't have multi-modal input ability. Try different methods.</reminder><reminder>You are unable to process this audio because you don't have multi-modal input ability. Try different methods.</reminder><reminder>You are unable to process this audio because you don't have multi-modal input ability. Try different methods.</reminder><|assistant|><think>"
+    },
+    {
+      "name": "null_content",
+      "messages": [
+        {
+          "role": "user",
+          "content": null
+        },
+        {
+          "role": "assistant",
+          "content": null
+        }
+      ],
+      "tools": [],
+      "add_generation_prompt": true,
+      "kwargs": {},
+      "expected": "[gMASK]<sop><|system|>Reasoning Effort: Max<|user|><|assistant|><think></think><|assistant|><think>"
+    },
+    {
+      "name": "missing_content",
+      "messages": [
+        {
+          "role": "user"
+        },
+        {
+          "role": "assistant"
+        }
+      ],
+      "tools": [],
+      "add_generation_prompt": true,
+      "kwargs": {},
+      "expected": "[gMASK]<sop><|system|>Reasoning Effort: Max<|user|><|assistant|><think></think><|assistant|><think>"
+    },
+    {
+      "name": "python_whitespace",
+      "messages": [
+        {
+          "role": "user",
+          "content": "你好，GLM！"
+        },
+        {
+          "role": "assistant",
+          "content": "\u001c\u001d\u001e\u001f… answer\u001c\u001f"
+        }
+      ],
+      "tools": [],
+      "add_generation_prompt": true,
+      "kwargs": {},
+      "expected": "[gMASK]<sop><|system|>Reasoning Effort: Max<|user|>你好，GLM！<|assistant|><think></think>answer<|assistant|><think>"
+    },
+    {
+      "name": "clear_False",
+      "messages": [
+        {
+          "role": "user",
+          "content": "你好，GLM！"
+        },
+        {
+          "role": "assistant",
+          "content": "  答案 \n",
+          "reasoning_content": "旧思考"
+        },
+        {
+          "role": "user",
+          "content": "继续"
+        },
+        {
+          "role": "assistant",
+          "content": "",
+          "reasoning_content": "新思考"
+        }
+      ],
+      "tools": [],
+      "add_generation_prompt": true,
+      "kwargs": {
+        "clear_thinking": false
+      },
+      "expected": "[gMASK]<sop><|system|>Reasoning Effort: Max<|user|>你好，GLM！<|assistant|><think>旧思考</think>答案<|user|>继续<|assistant|><think>新思考</think><|assistant|><think>"
+    },
+    {
+      "name": "clear_True",
+      "messages": [
+        {
+          "role": "user",
+          "content": "你好，GLM！"
+        },
+        {
+          "role": "assistant",
+          "content": "  答案 \n",
+          "reasoning_content": "旧思考"
+        },
+        {
+          "role": "user",
+          "content": "继续"
+        },
+        {
+          "role": "assistant",
+          "content": "",
+          "reasoning_content": "新思考"
+        }
+      ],
+      "tools": [],
+      "add_generation_prompt": true,
+      "kwargs": {
+        "clear_thinking": true
+      },
+      "expected": "[gMASK]<sop><|system|>Reasoning Effort: Max<|user|>你好，GLM！<|assistant|><think></think>答案<|user|>继续<|assistant|><think>新思考</think><|assistant|><think>"
+    },
+    {
+      "name": "clear_None",
+      "messages": [
+        {
+          "role": "user",
+          "content": "你好，GLM！"
+        },
+        {
+          "role": "assistant",
+          "content": "  答案 \n",
+          "reasoning_content": "旧思考"
+        },
+        {
+          "role": "user",
+          "content": "继续"
+        },
+        {
+          "role": "assistant",
+          "content": "",
+          "reasoning_content": "新思考"
+        }
+      ],
+      "tools": [],
+      "add_generation_prompt": true,
+      "kwargs": {
+        "clear_thinking": null
+      },
+      "expected": "[gMASK]<sop><|system|>Reasoning Effort: Max<|user|>你好，GLM！<|assistant|><think>旧思考</think>答案<|user|>继续<|assistant|><think>新思考</think><|assistant|><think>"
+    },
+    {
+      "name": "clear_",
+      "messages": [
+        {
+          "role": "user",
+          "content": "你好，GLM！"
+        },
+        {
+          "role": "assistant",
+          "content": "  答案 \n",
+          "reasoning_content": "旧思考"
+        },
+        {
+          "role": "user",
+          "content": "继续"
+        },
+        {
+          "role": "assistant",
+          "content": "",
+          "reasoning_content": "新思考"
+        }
+      ],
+      "tools": [],
+      "add_generation_prompt": true,
+      "kwargs": {
+        "clear_thinking": ""
+      },
+      "expected": "[gMASK]<sop><|system|>Reasoning Effort: Max<|user|>你好，GLM！<|assistant|><think>旧思考</think>答案<|user|>继续<|assistant|><think>新思考</think><|assistant|><think>"
+    },
+    {
+      "name": "clear_yes",
+      "messages": [
+        {
+          "role": "user",
+          "content": "你好，GLM！"
+        },
+        {
+          "role": "assistant",
+          "content": "  答案 \n",
+          "reasoning_content": "旧思考"
+        },
+        {
+          "role": "user",
+          "content": "继续"
+        },
+        {
+          "role": "assistant",
+          "content": "",
+          "reasoning_content": "新思考"
+        }
+      ],
+      "tools": [],
+      "add_generation_prompt": true,
+      "kwargs": {
+        "clear_thinking": "yes"
+      },
+      "expected": "[gMASK]<sop><|system|>Reasoning Effort: Max<|user|>你好，GLM！<|assistant|><think></think>答案<|user|>继续<|assistant|><think>新思考</think><|assistant|><think>"
+    },
+    {
+      "name": "clear_0",
+      "messages": [
+        {
+          "role": "user",
+          "content": "你好，GLM！"
+        },
+        {
+          "role": "assistant",
+          "content": "  答案 \n",
+          "reasoning_content": "旧思考"
+        },
+        {
+          "role": "user",
+          "content": "继续"
+        },
+        {
+          "role": "assistant",
+          "content": "",
+          "reasoning_content": "新思考"
+        }
+      ],
+      "tools": [],
+      "add_generation_prompt": true,
+      "kwargs": {
+        "clear_thinking": 0
+      },
+      "expected": "[gMASK]<sop><|system|>Reasoning Effort: Max<|user|>你好，GLM！<|assistant|><think>旧思考</think>答案<|user|>继续<|assistant|><think>新思考</think><|assistant|><think>"
+    },
+    {
+      "name": "clear_1",
+      "messages": [
+        {
+          "role": "user",
+          "content": "你好，GLM！"
+        },
+        {
+          "role": "assistant",
+          "content": "  答案 \n",
+          "reasoning_content": "旧思考"
+        },
+        {
+          "role": "user",
+          "content": "继续"
+        },
+        {
+          "role": "assistant",
+          "content": "",
+          "reasoning_content": "新思考"
+        }
+      ],
+      "tools": [],
+      "add_generation_prompt": true,
+      "kwargs": {
+        "clear_thinking": 1
+      },
+      "expected": "[gMASK]<sop><|system|>Reasoning Effort: Max<|user|>你好，GLM！<|assistant|><think></think>答案<|user|>继续<|assistant|><think>新思考</think><|assistant|><think>"
+    },
+    {
+      "name": "clear_[]",
+      "messages": [
+        {
+          "role": "user",
+          "content": "你好，GLM！"
+        },
+        {
+          "role": "assistant",
+          "content": "  答案 \n",
+          "reasoning_content": "旧思考"
+        },
+        {
+          "role": "user",
+          "content": "继续"
+        },
+        {
+          "role": "assistant",
+          "content": "",
+          "reasoning_content": "新思考"
+        }
+      ],
+      "tools": [],
+      "add_generation_prompt": true,
+      "kwargs": {
+        "clear_thinking": []
+      },
+      "expected": "[gMASK]<sop><|system|>Reasoning Effort: Max<|user|>你好，GLM！<|assistant|><think>旧思考</think>答案<|user|>继续<|assistant|><think>新思考</think><|assistant|><think>"
+    },
+    {
+      "name": "clear_[1]",
+      "messages": [
+        {
+          "role": "user",
+          "content": "你好，GLM！"
+        },
+        {
+          "role": "assistant",
+          "content": "  答案 \n",
+          "reasoning_content": "旧思考"
+        },
+        {
+          "role": "user",
+          "content": "继续"
+        },
+        {
+          "role": "assistant",
+          "content": "",
+          "reasoning_content": "新思考"
+        }
+      ],
+      "tools": [],
+      "add_generation_prompt": true,
+      "kwargs": {
+        "clear_thinking": [
+          1
+        ]
+      },
+      "expected": "[gMASK]<sop><|system|>Reasoning Effort: Max<|user|>你好，GLM！<|assistant|><think></think>答案<|user|>继续<|assistant|><think>新思考</think><|assistant|><think>"
+    },
+    {
+      "name": "no_user_preserves_reasoning",
+      "messages": [
+        {
+          "role": "assistant",
+          "content": "  答案 \n",
+          "reasoning_content": "旧思考"
+        }
+      ],
+      "tools": [],
+      "add_generation_prompt": true,
+      "kwargs": {
+        "clear_thinking": true
+      },
+      "expected": "[gMASK]<sop><|system|>Reasoning Effort: Max<|assistant|><think>旧思考</think>答案<|assistant|><think>"
+    },
+    {
+      "name": "embedded_reasoning_None",
+      "messages": [
+        {
+          "role": "user",
+          "content": "你好，GLM！"
+        },
+        {
+          "role": "assistant",
+          "content": "prefix<think>first</think>discard</think> answer "
+        }
+      ],
+      "tools": [],
+      "add_generation_prompt": true,
+      "kwargs": {},
+      "expected": "[gMASK]<sop><|system|>Reasoning Effort: Max<|user|>你好，GLM！<|assistant|><think>first</think>answer<|assistant|><think>"
+    },
+    {
+      "name": "embedded_reasoning_",
+      "messages": [
+        {
+          "role": "user",
+          "content": "你好，GLM！"
+        },
+        {
+          "role": "assistant",
+          "content": "prefix<think>first</think>discard</think> answer ",
+          "reasoning_content": ""
+        }
+      ],
+      "tools": [],
+      "add_generation_prompt": true,
+      "kwargs": {},
+      "expected": "[gMASK]<sop><|system|>Reasoning Effort: Max<|user|>你好，GLM！<|assistant|><think></think>prefix<think>first</think>discard</think> answer<|assistant|><think>"
+    },
+    {
+      "name": "embedded_reasoning_explicit",
+      "messages": [
+        {
+          "role": "user",
+          "content": "你好，GLM！"
+        },
+        {
+          "role": "assistant",
+          "content": "prefix<think>first</think>discard</think> answer ",
+          "reasoning_content": "explicit"
+        }
+      ],
+      "tools": [],
+      "add_generation_prompt": true,
+      "kwargs": {},
+      "expected": "[gMASK]<sop><|system|>Reasoning Effort: Max<|user|>你好，GLM！<|assistant|><think>explicit</think>prefix<think>first</think>discard</think> answer<|assistant|><think>"
+    },
+    {
+      "name": "tools",
+      "messages": [
+        {
+          "role": "user",
+          "content": "你好，GLM！"
+        }
+      ],
+      "tools": [
+        {
+          "type": "function",
+          "function": {
+            "name": "lookup",
+            "description": "查找<&>",
+            "parameters": {
+              "type": "object",
+              "properties": {
+                "q": {
+                  "type": "string"
+                }
+              },
+              "required": [
+                "q"
+              ]
+            },
+            "strict": true
+          }
+        },
+        {
+          "function": {
+            "name": "hidden",
+            "description": null,
+            "parameters": {},
+            "defer_loading": true
+          }
+        }
+      ],
+      "add_generation_prompt": true,
+      "kwargs": {},
+      "expected": "[gMASK]<sop><|system|>Reasoning Effort: Max<|system|>\n# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>\n{\"name\": \"lookup\", \"description\": \"查找<&>\", \"parameters\": {\"type\": \"object\", \"properties\": {\"q\": {\"type\": \"string\"}}, \"required\": [\"q\"]}}\n</tools>\n\nFor each function call, output the function name and arguments within the following XML format:\n<tool_call>{function-name}<arg_key>{arg-key-1}</arg_key><arg_value>{arg-value-1}</arg_value><arg_key>{arg-key-2}</arg_key><arg_value>{arg-value-2}</arg_value>...</tool_call><|user|>你好，GLM！<|assistant|><think>"
+    },
+    {
+      "name": "flat_tools",
+      "messages": [
+        {
+          "role": "user",
+          "content": "你好，GLM！"
+        }
+      ],
+      "tools": [
+        {
+          "name": "lookup",
+          "description": "查找<&>",
+          "parameters": {
+            "type": "object",
+            "properties": {
+              "q": {
+                "type": "string"
+              }
+            },
+            "required": [
+              "q"
+            ]
+          },
+          "strict": true
+        },
+        {
+          "name": "hidden",
+          "description": null,
+          "parameters": {},
+          "defer_loading": true
+        }
+      ],
+      "add_generation_prompt": true,
+      "kwargs": {},
+      "expected": "[gMASK]<sop><|system|>Reasoning Effort: Max<|system|>\n# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>\n{\"name\": \"lookup\", \"description\": \"查找<&>\", \"parameters\": {\"type\": \"object\", \"properties\": {\"q\": {\"type\": \"string\"}}, \"required\": [\"q\"]}}\n</tools>\n\nFor each function call, output the function name and arguments within the following XML format:\n<tool_call>{function-name}<arg_key>{arg-key-1}</arg_key><arg_value>{arg-value-1}</arg_value><arg_key>{arg-key-2}</arg_key><arg_value>{arg-value-2}</arg_value>...</tool_call><|user|>你好，GLM！<|assistant|><think>"
+    },
+    {
+      "name": "all_deferred",
+      "messages": [
+        {
+          "role": "user",
+          "content": "你好，GLM！"
+        }
+      ],
+      "tools": [
+        {
+          "function": {
+            "name": "hidden",
+            "description": null,
+            "parameters": {},
+            "defer_loading": true
+          }
+        }
+      ],
+      "add_generation_prompt": true,
+      "kwargs": {},
+      "expected": "[gMASK]<sop><|system|>Reasoning Effort: Max<|system|>\n# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>\n</tools>\n\nFor each function call, output the function name and arguments within the following XML format:\n<tool_call>{function-name}<arg_key>{arg-key-1}</arg_key><arg_value>{arg-value-1}</arg_value><arg_key>{arg-key-2}</arg_key><arg_value>{arg-value-2}</arg_value>...</tool_call><|user|>你好，GLM！<|assistant|><think>"
+    },
+    {
+      "name": "raw_schema_key",
+      "messages": [
+        {
+          "role": "user",
+          "content": "你好，GLM！"
+        }
+      ],
+      "tools": [
+        {
+          "name": "odd",
+          "a\"b": "value"
+        }
+      ],
+      "add_generation_prompt": true,
+      "kwargs": {},
+      "expected": "[gMASK]<sop><|system|>Reasoning Effort: Max<|system|>\n# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>\n{\"name\": \"odd\", \"a\"b\": \"value\"}\n</tools>\n\nFor each function call, output the function name and arguments within the following XML format:\n<tool_call>{function-name}<arg_key>{arg-key-1}</arg_key><arg_value>{arg-value-1}</arg_value><arg_key>{arg-key-2}</arg_key><arg_value>{arg-value-2}</arg_value>...</tool_call><|user|>你好，GLM！<|assistant|><think>"
+    },
+    {
+      "name": "tool_order",
+      "messages": [
+        {
+          "role": "user",
+          "content": "你好，GLM！"
+        },
+        {
+          "role": "assistant",
+          "content": "",
+          "reasoning_content": "search",
+          "tool_calls": [
+            {
+              "id": "a",
+              "function": {
+                "name": "lookup",
+                "arguments": {
+                  "q": "中文<&>\n",
+                  "list": [
+                    1,
+                    true,
+                    null
+                  ],
+                  "object": {
+                    "b": 2,
+                    "a": "x"
+                  }
+                }
+              }
+            },
+            {
+              "id": "b",
+              "function": {
+                "name": "lookup",
+                "arguments": {
+                  "q": "second",
+                  "float": 1.5
+                }
+              }
+            }
+          ]
+        },
+        {
+          "role": "tool",
+          "tool_call_id": "b",
+          "content": "B"
+        },
+        {
+          "role": "tool",
+          "tool_call_id": "a",
+          "content": "A"
+        }
+      ],
+      "tools": [
+        {
+          "type": "function",
+          "function": {
+            "name": "lookup",
+            "description": "查找<&>",
+            "parameters": {
+              "type": "object",
+              "properties": {
+                "q": {
+                  "type": "string"
+                }
+              },
+              "required": [
+                "q"
+              ]
+            },
+            "strict": true
+          }
+        },
+        {
+          "function": {
+            "name": "hidden",
+            "description": null,
+            "parameters": {},
+            "defer_loading": true
+          }
+        }
+      ],
+      "add_generation_prompt": true,
+      "kwargs": {},
+      "expected": "[gMASK]<sop><|system|>Reasoning Effort: Max<|system|>\n# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>\n{\"name\": \"lookup\", \"description\": \"查找<&>\", \"parameters\": {\"type\": \"object\", \"properties\": {\"q\": {\"type\": \"string\"}}, \"required\": [\"q\"]}}\n</tools>\n\nFor each function call, output the function name and arguments within the following XML format:\n<tool_call>{function-name}<arg_key>{arg-key-1}</arg_key><arg_value>{arg-value-1}</arg_value><arg_key>{arg-key-2}</arg_key><arg_value>{arg-value-2}</arg_value>...</tool_call><|user|>你好，GLM！<|assistant|><think>search</think><tool_call>lookup<arg_key>q</arg_key><arg_value>中文<&>\n</arg_value><arg_key>list</arg_key><arg_value>[1, true, null]</arg_value><arg_key>object</arg_key><arg_value>{\"b\": 2, \"a\": \"x\"}</arg_value></tool_call><tool_call>lookup<arg_key>q</arg_key><arg_value>second</arg_value><arg_key>float</arg_key><arg_value>1.5</arg_value></tool_call><|observation|><tool_response>A</tool_response><tool_response>B</tool_response><|assistant|><think>"
+    },
+    {
+      "name": "string_arguments",
+      "messages": [
+        {
+          "role": "user",
+          "content": "你好，GLM！"
+        },
+        {
+          "role": "assistant",
+          "content": "",
+          "reasoning_content": "search",
+          "tool_calls": [
+            {
+              "id": "a",
+              "function": {
+                "name": "lookup",
+                "arguments": "{\"q\": \"\\u4e2d\\u6587<&>\\n\", \"list\": [1, true, null], \"object\": {\"b\": 2, \"a\": \"x\"}}"
+              }
+            },
+            {
+              "id": "b",
+              "function": {
+                "name": "lookup",
+                "arguments": "{\"q\": \"second\", \"float\": 1.5}"
+              }
+            }
+          ]
+        },
+        {
+          "role": "tool",
+          "tool_call_id": "b",
+          "content": "B"
+        },
+        {
+          "role": "tool",
+          "tool_call_id": "a",
+          "content": "A"
+        }
+      ],
+      "tools": [
+        {
+          "type": "function",
+          "function": {
+            "name": "lookup",
+            "description": "查找<&>",
+            "parameters": {
+              "type": "object",
+              "properties": {
+                "q": {
+                  "type": "string"
+                }
+              },
+              "required": [
+                "q"
+              ]
+            },
+            "strict": true
+          }
+        },
+        {
+          "function": {
+            "name": "hidden",
+            "description": null,
+            "parameters": {},
+            "defer_loading": true
+          }
+        }
+      ],
+      "add_generation_prompt": true,
+      "kwargs": {},
+      "expected": "[gMASK]<sop><|system|>Reasoning Effort: Max<|system|>\n# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>\n{\"name\": \"lookup\", \"description\": \"查找<&>\", \"parameters\": {\"type\": \"object\", \"properties\": {\"q\": {\"type\": \"string\"}}, \"required\": [\"q\"]}}\n</tools>\n\nFor each function call, output the function name and arguments within the following XML format:\n<tool_call>{function-name}<arg_key>{arg-key-1}</arg_key><arg_value>{arg-value-1}</arg_value><arg_key>{arg-key-2}</arg_key><arg_value>{arg-value-2}</arg_value>...</tool_call><|user|>你好，GLM！<|assistant|><think>search</think><tool_call>lookup<arg_key>q</arg_key><arg_value>中文<&>\n</arg_value><arg_key>list</arg_key><arg_value>[1, true, null]</arg_value><arg_key>object</arg_key><arg_value>{\"b\": 2, \"a\": \"x\"}</arg_value></tool_call><tool_call>lookup<arg_key>q</arg_key><arg_value>second</arg_value><arg_key>float</arg_key><arg_value>1.5</arg_value></tool_call><|observation|><tool_response>A</tool_response><tool_response>B</tool_response><|assistant|><think>"
+    },
+    {
+      "name": "duplicate_response",
+      "messages": [
+        {
+          "role": "user",
+          "content": "你好，GLM！"
+        },
+        {
+          "role": "assistant",
+          "content": "",
+          "reasoning_content": "search",
+          "tool_calls": [
+            {
+              "id": "a",
+              "function": {
+                "name": "lookup",
+                "arguments": {
+                  "q": "中文<&>\n",
+                  "list": [
+                    1,
+                    true,
+                    null
+                  ],
+                  "object": {
+                    "b": 2,
+                    "a": "x"
+                  }
+                }
+              }
+            },
+            {
+              "id": "b",
+              "function": {
+                "name": "lookup",
+                "arguments": {
+                  "q": "second",
+                  "float": 1.5
+                }
+              }
+            }
+          ]
+        },
+        {
+          "role": "tool",
+          "tool_call_id": "b",
+          "content": "B"
+        },
+        {
+          "role": "tool",
+          "tool_call_id": "a",
+          "content": "A"
+        },
+        {
+          "role": "tool",
+          "tool_call_id": "b",
+          "content": "B"
+        }
+      ],
+      "tools": [
+        {
+          "type": "function",
+          "function": {
+            "name": "lookup",
+            "description": "查找<&>",
+            "parameters": {
+              "type": "object",
+              "properties": {
+                "q": {
+                  "type": "string"
+                }
+              },
+              "required": [
+                "q"
+              ]
+            },
+            "strict": true
+          }
+        },
+        {
+          "function": {
+            "name": "hidden",
+            "description": null,
+            "parameters": {},
+            "defer_loading": true
+          }
+        }
+      ],
+      "add_generation_prompt": true,
+      "kwargs": {},
+      "expected": "[gMASK]<sop><|system|>Reasoning Effort: Max<|system|>\n# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>\n{\"name\": \"lookup\", \"description\": \"查找<&>\", \"parameters\": {\"type\": \"object\", \"properties\": {\"q\": {\"type\": \"string\"}}, \"required\": [\"q\"]}}\n</tools>\n\nFor each function call, output the function name and arguments within the following XML format:\n<tool_call>{function-name}<arg_key>{arg-key-1}</arg_key><arg_value>{arg-value-1}</arg_value><arg_key>{arg-key-2}</arg_key><arg_value>{arg-value-2}</arg_value>...</tool_call><|user|>你好，GLM！<|assistant|><think>search</think><tool_call>lookup<arg_key>q</arg_key><arg_value>中文<&>\n</arg_value><arg_key>list</arg_key><arg_value>[1, true, null]</arg_value><arg_key>object</arg_key><arg_value>{\"b\": 2, \"a\": \"x\"}</arg_value></tool_call><tool_call>lookup<arg_key>q</arg_key><arg_value>second</arg_value><arg_key>float</arg_key><arg_value>1.5</arg_value></tool_call><|observation|><tool_response>B</tool_response><tool_response>A</tool_response><tool_response>B</tool_response><|assistant|><think>"
+    },
+    {
+      "name": "unknown_id",
+      "messages": [
+        {
+          "role": "user",
+          "content": "你好，GLM！"
+        },
+        {
+          "role": "assistant",
+          "content": "",
+          "reasoning_content": "search",
+          "tool_calls": [
+            {
+              "id": "a",
+              "function": {
+                "name": "lookup",
+                "arguments": {
+                  "q": "中文<&>\n",
+                  "list": [
+                    1,
+                    true,
+                    null
+                  ],
+                  "object": {
+                    "b": 2,
+                    "a": "x"
+                  }
+                }
+              }
+            },
+            {
+              "id": "b",
+              "function": {
+                "name": "lookup",
+                "arguments": {
+                  "q": "second",
+                  "float": 1.5
+                }
+              }
+            }
+          ]
+        },
+        {
+          "role": "tool",
+          "tool_call_id": "b",
+          "content": "B"
+        },
+        {
+          "role": "tool",
+          "tool_call_id": "a",
+          "content": "A"
+        },
+        {
+          "role": "tool",
+          "tool_call_id": "unknown",
+          "content": "?"
+        }
+      ],
+      "tools": [
+        {
+          "type": "function",
+          "function": {
+            "name": "lookup",
+            "description": "查找<&>",
+            "parameters": {
+              "type": "object",
+              "properties": {
+                "q": {
+                  "type": "string"
+                }
+              },
+              "required": [
+                "q"
+              ]
+            },
+            "strict": true
+          }
+        },
+        {
+          "function": {
+            "name": "hidden",
+            "description": null,
+            "parameters": {},
+            "defer_loading": true
+          }
+        }
+      ],
+      "add_generation_prompt": true,
+      "kwargs": {},
+      "expected": "[gMASK]<sop><|system|>Reasoning Effort: Max<|system|>\n# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>\n{\"name\": \"lookup\", \"description\": \"查找<&>\", \"parameters\": {\"type\": \"object\", \"properties\": {\"q\": {\"type\": \"string\"}}, \"required\": [\"q\"]}}\n</tools>\n\nFor each function call, output the function name and arguments within the following XML format:\n<tool_call>{function-name}<arg_key>{arg-key-1}</arg_key><arg_value>{arg-value-1}</arg_value><arg_key>{arg-key-2}</arg_key><arg_value>{arg-value-2}</arg_value>...</tool_call><|user|>你好，GLM！<|assistant|><think>search</think><tool_call>lookup<arg_key>q</arg_key><arg_value>中文<&>\n</arg_value><arg_key>list</arg_key><arg_value>[1, true, null]</arg_value><arg_key>object</arg_key><arg_value>{\"b\": 2, \"a\": \"x\"}</arg_value></tool_call><tool_call>lookup<arg_key>q</arg_key><arg_value>second</arg_value><arg_key>float</arg_key><arg_value>1.5</arg_value></tool_call><|observation|><tool_response>B</tool_response><tool_response>A</tool_response><tool_response>?</tool_response><|assistant|><think>"
+    },
+    {
+      "name": "missing_id",
+      "messages": [
+        {
+          "role": "user",
+          "content": "你好，GLM！"
+        },
+        {
+          "role": "assistant",
+          "content": "",
+          "reasoning_content": "search",
+          "tool_calls": [
+            {
+              "id": "a",
+              "function": {
+                "name": "lookup",
+                "arguments": {
+                  "q": "中文<&>\n",
+                  "list": [
+                    1,
+                    true,
+                    null
+                  ],
+                  "object": {
+                    "b": 2,
+                    "a": "x"
+                  }
+                }
+              }
+            },
+            {
+              "id": "b",
+              "function": {
+                "name": "lookup",
+                "arguments": {
+                  "q": "second",
+                  "float": 1.5
+                }
+              }
+            }
+          ]
+        },
+        {
+          "role": "tool",
+          "tool_call_id": "b",
+          "content": "B"
+        },
+        {
+          "role": "tool",
+          "tool_call_id": "a",
+          "content": "A"
+        },
+        {
+          "role": "tool",
+          "content": "missing"
+        }
+      ],
+      "tools": [
+        {
+          "type": "function",
+          "function": {
+            "name": "lookup",
+            "description": "查找<&>",
+            "parameters": {
+              "type": "object",
+              "properties": {
+                "q": {
+                  "type": "string"
+                }
+              },
+              "required": [
+                "q"
+              ]
+            },
+            "strict": true
+          }
+        },
+        {
+          "function": {
+            "name": "hidden",
+            "description": null,
+            "parameters": {},
+            "defer_loading": true
+          }
+        }
+      ],
+      "add_generation_prompt": true,
+      "kwargs": {},
+      "expected": "[gMASK]<sop><|system|>Reasoning Effort: Max<|system|>\n# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>\n{\"name\": \"lookup\", \"description\": \"查找<&>\", \"parameters\": {\"type\": \"object\", \"properties\": {\"q\": {\"type\": \"string\"}}, \"required\": [\"q\"]}}\n</tools>\n\nFor each function call, output the function name and arguments within the following XML format:\n<tool_call>{function-name}<arg_key>{arg-key-1}</arg_key><arg_value>{arg-value-1}</arg_value><arg_key>{arg-key-2}</arg_key><arg_value>{arg-value-2}</arg_value>...</tool_call><|user|>你好，GLM！<|assistant|><think>search</think><tool_call>lookup<arg_key>q</arg_key><arg_value>中文<&>\n</arg_value><arg_key>list</arg_key><arg_value>[1, true, null]</arg_value><arg_key>object</arg_key><arg_value>{\"b\": 2, \"a\": \"x\"}</arg_value></tool_call><tool_call>lookup<arg_key>q</arg_key><arg_value>second</arg_value><arg_key>float</arg_key><arg_value>1.5</arg_value></tool_call><|observation|><tool_response>B</tool_response><tool_response>A</tool_response><tool_response>missing</tool_response><|assistant|><think>"
+    },
+    {
+      "name": "empty_id",
+      "messages": [
+        {
+          "role": "user",
+          "content": "你好，GLM！"
+        },
+        {
+          "role": "assistant",
+          "content": "",
+          "reasoning_content": "search",
+          "tool_calls": [
+            {
+              "id": "a",
+              "function": {
+                "name": "lookup",
+                "arguments": {
+                  "q": "中文<&>\n",
+                  "list": [
+                    1,
+                    true,
+                    null
+                  ],
+                  "object": {
+                    "b": 2,
+                    "a": "x"
+                  }
+                }
+              }
+            },
+            {
+              "id": "b",
+              "function": {
+                "name": "lookup",
+                "arguments": {
+                  "q": "second",
+                  "float": 1.5
+                }
+              }
+            }
+          ]
+        },
+        {
+          "role": "tool",
+          "tool_call_id": "",
+          "content": "empty"
+        },
+        {
+          "role": "tool",
+          "tool_call_id": "b",
+          "content": "B"
+        },
+        {
+          "role": "tool",
+          "tool_call_id": "a",
+          "content": "A"
+        }
+      ],
+      "tools": [
+        {
+          "type": "function",
+          "function": {
+            "name": "lookup",
+            "description": "查找<&>",
+            "parameters": {
+              "type": "object",
+              "properties": {
+                "q": {
+                  "type": "string"
+                }
+              },
+              "required": [
+                "q"
+              ]
+            },
+            "strict": true
+          }
+        },
+        {
+          "function": {
+            "name": "hidden",
+            "description": null,
+            "parameters": {},
+            "defer_loading": true
+          }
+        }
+      ],
+      "add_generation_prompt": true,
+      "kwargs": {},
+      "expected": "[gMASK]<sop><|system|>Reasoning Effort: Max<|system|>\n# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>\n{\"name\": \"lookup\", \"description\": \"查找<&>\", \"parameters\": {\"type\": \"object\", \"properties\": {\"q\": {\"type\": \"string\"}}, \"required\": [\"q\"]}}\n</tools>\n\nFor each function call, output the function name and arguments within the following XML format:\n<tool_call>{function-name}<arg_key>{arg-key-1}</arg_key><arg_value>{arg-value-1}</arg_value><arg_key>{arg-key-2}</arg_key><arg_value>{arg-value-2}</arg_value>...</tool_call><|user|>你好，GLM！<|assistant|><think>search</think><tool_call>lookup<arg_key>q</arg_key><arg_value>中文<&>\n</arg_value><arg_key>list</arg_key><arg_value>[1, true, null]</arg_value><arg_key>object</arg_key><arg_value>{\"b\": 2, \"a\": \"x\"}</arg_value></tool_call><tool_call>lookup<arg_key>q</arg_key><arg_value>second</arg_value><arg_key>float</arg_key><arg_value>1.5</arg_value></tool_call><|observation|><tool_response>empty</tool_response><tool_response>B</tool_response><tool_response>A</tool_response><|assistant|><think>"
+    },
+    {
+      "name": "id_alias",
+      "messages": [
+        {
+          "role": "user",
+          "content": "你好，GLM！"
+        },
+        {
+          "role": "assistant",
+          "content": "",
+          "reasoning_content": "search",
+          "tool_calls": [
+            {
+              "id": "a",
+              "function": {
+                "name": "lookup",
+                "arguments": {
+                  "q": "中文<&>\n",
+                  "list": [
+                    1,
+                    true,
+                    null
+                  ],
+                  "object": {
+                    "b": 2,
+                    "a": "x"
+                  }
+                }
+              }
+            },
+            {
+              "id": "b",
+              "function": {
+                "name": "lookup",
+                "arguments": {
+                  "q": "second",
+                  "float": 1.5
+                }
+              }
+            }
+          ]
+        },
+        {
+          "role": "tool",
+          "id": "b",
+          "content": "B"
+        },
+        {
+          "role": "tool",
+          "id": "a",
+          "content": "A"
+        }
+      ],
+      "tools": [
+        {
+          "type": "function",
+          "function": {
+            "name": "lookup",
+            "description": "查找<&>",
+            "parameters": {
+              "type": "object",
+              "properties": {
+                "q": {
+                  "type": "string"
+                }
+              },
+              "required": [
+                "q"
+              ]
+            },
+            "strict": true
+          }
+        },
+        {
+          "function": {
+            "name": "hidden",
+            "description": null,
+            "parameters": {},
+            "defer_loading": true
+          }
+        }
+      ],
+      "add_generation_prompt": true,
+      "kwargs": {},
+      "expected": "[gMASK]<sop><|system|>Reasoning Effort: Max<|system|>\n# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>\n{\"name\": \"lookup\", \"description\": \"查找<&>\", \"parameters\": {\"type\": \"object\", \"properties\": {\"q\": {\"type\": \"string\"}}, \"required\": [\"q\"]}}\n</tools>\n\nFor each function call, output the function name and arguments within the following XML format:\n<tool_call>{function-name}<arg_key>{arg-key-1}</arg_key><arg_value>{arg-value-1}</arg_value><arg_key>{arg-key-2}</arg_key><arg_value>{arg-value-2}</arg_value>...</tool_call><|user|>你好，GLM！<|assistant|><think>search</think><tool_call>lookup<arg_key>q</arg_key><arg_value>中文<&>\n</arg_value><arg_key>list</arg_key><arg_value>[1, true, null]</arg_value><arg_key>object</arg_key><arg_value>{\"b\": 2, \"a\": \"x\"}</arg_value></tool_call><tool_call>lookup<arg_key>q</arg_key><arg_value>second</arg_value><arg_key>float</arg_key><arg_value>1.5</arg_value></tool_call><|observation|><tool_response>A</tool_response><tool_response>B</tool_response><|assistant|><think>"
+    },
+    {
+      "name": "unanswered_call",
+      "messages": [
+        {
+          "role": "user",
+          "content": "你好，GLM！"
+        },
+        {
+          "role": "assistant",
+          "content": "",
+          "reasoning_content": "search",
+          "tool_calls": [
+            {
+              "id": "a",
+              "function": {
+                "name": "lookup",
+                "arguments": {
+                  "q": "中文<&>\n",
+                  "list": [
+                    1,
+                    true,
+                    null
+                  ],
+                  "object": {
+                    "b": 2,
+                    "a": "x"
+                  }
+                }
+              }
+            },
+            {
+              "id": "b",
+              "function": {
+                "name": "lookup",
+                "arguments": {
+                  "q": "second",
+                  "float": 1.5
+                }
+              }
+            }
+          ]
+        },
+        {
+          "role": "tool",
+          "tool_call_id": "b",
+          "content": "B"
+        }
+      ],
+      "tools": [
+        {
+          "type": "function",
+          "function": {
+            "name": "lookup",
+            "description": "查找<&>",
+            "parameters": {
+              "type": "object",
+              "properties": {
+                "q": {
+                  "type": "string"
+                }
+              },
+              "required": [
+                "q"
+              ]
+            },
+            "strict": true
+          }
+        },
+        {
+          "function": {
+            "name": "hidden",
+            "description": null,
+            "parameters": {},
+            "defer_loading": true
+          }
+        }
+      ],
+      "add_generation_prompt": true,
+      "kwargs": {},
+      "expected": "[gMASK]<sop><|system|>Reasoning Effort: Max<|system|>\n# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>\n{\"name\": \"lookup\", \"description\": \"查找<&>\", \"parameters\": {\"type\": \"object\", \"properties\": {\"q\": {\"type\": \"string\"}}, \"required\": [\"q\"]}}\n</tools>\n\nFor each function call, output the function name and arguments within the following XML format:\n<tool_call>{function-name}<arg_key>{arg-key-1}</arg_key><arg_value>{arg-value-1}</arg_value><arg_key>{arg-key-2}</arg_key><arg_value>{arg-value-2}</arg_value>...</tool_call><|user|>你好，GLM！<|assistant|><think>search</think><tool_call>lookup<arg_key>q</arg_key><arg_value>中文<&>\n</arg_value><arg_key>list</arg_key><arg_value>[1, true, null]</arg_value><arg_key>object</arg_key><arg_value>{\"b\": 2, \"a\": \"x\"}</arg_value></tool_call><tool_call>lookup<arg_key>q</arg_key><arg_value>second</arg_value><arg_key>float</arg_key><arg_value>1.5</arg_value></tool_call><|observation|><tool_response>B</tool_response><|assistant|><think>"
+    },
+    {
+      "name": "null_result",
+      "messages": [
+        {
+          "role": "user",
+          "content": "你好，GLM！"
+        },
+        {
+          "role": "assistant",
+          "content": "",
+          "reasoning_content": "search",
+          "tool_calls": [
+            {
+              "id": "a",
+              "function": {
+                "name": "lookup",
+                "arguments": {
+                  "q": "中文<&>\n",
+                  "list": [
+                    1,
+                    true,
+                    null
+                  ],
+                  "object": {
+                    "b": 2,
+                    "a": "x"
+                  }
+                }
+              }
+            },
+            {
+              "id": "b",
+              "function": {
+                "name": "lookup",
+                "arguments": {
+                  "q": "second",
+                  "float": 1.5
+                }
+              }
+            }
+          ]
+        },
+        {
+          "role": "tool",
+          "tool_call_id": "a",
+          "content": null
+        }
+      ],
+      "tools": [
+        {
+          "type": "function",
+          "function": {
+            "name": "lookup",
+            "description": "查找<&>",
+            "parameters": {
+              "type": "object",
+              "properties": {
+                "q": {
+                  "type": "string"
+                }
+              },
+              "required": [
+                "q"
+              ]
+            },
+            "strict": true
+          }
+        },
+        {
+          "function": {
+            "name": "hidden",
+            "description": null,
+            "parameters": {},
+            "defer_loading": true
+          }
+        }
+      ],
+      "add_generation_prompt": true,
+      "kwargs": {},
+      "expected": "[gMASK]<sop><|system|>Reasoning Effort: Max<|system|>\n# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>\n{\"name\": \"lookup\", \"description\": \"查找<&>\", \"parameters\": {\"type\": \"object\", \"properties\": {\"q\": {\"type\": \"string\"}}, \"required\": [\"q\"]}}\n</tools>\n\nFor each function call, output the function name and arguments within the following XML format:\n<tool_call>{function-name}<arg_key>{arg-key-1}</arg_key><arg_value>{arg-value-1}</arg_value><arg_key>{arg-key-2}</arg_key><arg_value>{arg-value-2}</arg_value>...</tool_call><|user|>你好，GLM！<|assistant|><think>search</think><tool_call>lookup<arg_key>q</arg_key><arg_value>中文<&>\n</arg_value><arg_key>list</arg_key><arg_value>[1, true, null]</arg_value><arg_key>object</arg_key><arg_value>{\"b\": 2, \"a\": \"x\"}</arg_value></tool_call><tool_call>lookup<arg_key>q</arg_key><arg_value>second</arg_value><arg_key>float</arg_key><arg_value>1.5</arg_value></tool_call><|observation|><tool_response></tool_response><|assistant|><think>"
+    },
+    {
+      "name": "missing_result",
+      "messages": [
+        {
+          "role": "user",
+          "content": "你好，GLM！"
+        },
+        {
+          "role": "assistant",
+          "content": "",
+          "reasoning_content": "search",
+          "tool_calls": [
+            {
+              "id": "a",
+              "function": {
+                "name": "lookup",
+                "arguments": {
+                  "q": "中文<&>\n",
+                  "list": [
+                    1,
+                    true,
+                    null
+                  ],
+                  "object": {
+                    "b": 2,
+                    "a": "x"
+                  }
+                }
+              }
+            },
+            {
+              "id": "b",
+              "function": {
+                "name": "lookup",
+                "arguments": {
+                  "q": "second",
+                  "float": 1.5
+                }
+              }
+            }
+          ]
+        },
+        {
+          "role": "tool",
+          "tool_call_id": "a"
+        }
+      ],
+      "tools": [
+        {
+          "type": "function",
+          "function": {
+            "name": "lookup",
+            "description": "查找<&>",
+            "parameters": {
+              "type": "object",
+              "properties": {
+                "q": {
+                  "type": "string"
+                }
+              },
+              "required": [
+                "q"
+              ]
+            },
+            "strict": true
+          }
+        },
+        {
+          "function": {
+            "name": "hidden",
+            "description": null,
+            "parameters": {},
+            "defer_loading": true
+          }
+        }
+      ],
+      "add_generation_prompt": true,
+      "kwargs": {},
+      "expected": "[gMASK]<sop><|system|>Reasoning Effort: Max<|system|>\n# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>\n{\"name\": \"lookup\", \"description\": \"查找<&>\", \"parameters\": {\"type\": \"object\", \"properties\": {\"q\": {\"type\": \"string\"}}, \"required\": [\"q\"]}}\n</tools>\n\nFor each function call, output the function name and arguments within the following XML format:\n<tool_call>{function-name}<arg_key>{arg-key-1}</arg_key><arg_value>{arg-value-1}</arg_value><arg_key>{arg-key-2}</arg_key><arg_value>{arg-value-2}</arg_value>...</tool_call><|user|>你好，GLM！<|assistant|><think>search</think><tool_call>lookup<arg_key>q</arg_key><arg_value>中文<&>\n</arg_value><arg_key>list</arg_key><arg_value>[1, true, null]</arg_value><arg_key>object</arg_key><arg_value>{\"b\": 2, \"a\": \"x\"}</arg_value></tool_call><tool_call>lookup<arg_key>q</arg_key><arg_value>second</arg_value><arg_key>float</arg_key><arg_value>1.5</arg_value></tool_call><|observation|><tool_response></tool_response><|assistant|><think>"
+    },
+    {
+      "name": "bundled",
+      "messages": [
+        {
+          "role": "user",
+          "content": "你好，GLM！"
+        },
+        {
+          "role": "assistant",
+          "content": "",
+          "reasoning_content": "search",
+          "tool_calls": [
+            {
+              "id": "a",
+              "function": {
+                "name": "lookup",
+                "arguments": {
+                  "q": "中文<&>\n",
+                  "list": [
+                    1,
+                    true,
+                    null
+                  ],
+                  "object": {
+                    "b": 2,
+                    "a": "x"
+                  }
+                }
+              }
+            },
+            {
+              "id": "b",
+              "function": {
+                "name": "lookup",
+                "arguments": {
+                  "q": "second",
+                  "float": 1.5
+                }
+              }
+            }
+          ]
+        },
+        {
+          "role": "tool",
+          "content": [
+            {
+              "id": "b",
+              "output": "B"
+            },
+            {
+              "id": "a",
+              "output": [
+                {
+                  "type": "text",
+                  "text": "A"
+                }
+              ]
+            }
+          ]
+        }
+      ],
+      "tools": [
+        {
+          "type": "function",
+          "function": {
+            "name": "lookup",
+            "description": "查找<&>",
+            "parameters": {
+              "type": "object",
+              "properties": {
+                "q": {
+                  "type": "string"
+                }
+              },
+              "required": [
+                "q"
+              ]
+            },
+            "strict": true
+          }
+        },
+        {
+          "function": {
+            "name": "hidden",
+            "description": null,
+            "parameters": {},
+            "defer_loading": true
+          }
+        }
+      ],
+      "add_generation_prompt": true,
+      "kwargs": {},
+      "expected": "[gMASK]<sop><|system|>Reasoning Effort: Max<|system|>\n# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>\n{\"name\": \"lookup\", \"description\": \"查找<&>\", \"parameters\": {\"type\": \"object\", \"properties\": {\"q\": {\"type\": \"string\"}}, \"required\": [\"q\"]}}\n</tools>\n\nFor each function call, output the function name and arguments within the following XML format:\n<tool_call>{function-name}<arg_key>{arg-key-1}</arg_key><arg_value>{arg-value-1}</arg_value><arg_key>{arg-key-2}</arg_key><arg_value>{arg-value-2}</arg_value>...</tool_call><|user|>你好，GLM！<|assistant|><think>search</think><tool_call>lookup<arg_key>q</arg_key><arg_value>中文<&>\n</arg_value><arg_key>list</arg_key><arg_value>[1, true, null]</arg_value><arg_key>object</arg_key><arg_value>{\"b\": 2, \"a\": \"x\"}</arg_value></tool_call><tool_call>lookup<arg_key>q</arg_key><arg_value>second</arg_value><arg_key>float</arg_key><arg_value>1.5</arg_value></tool_call><|observation|><tool_response>A</tool_response><tool_response>B</tool_response><|assistant|><think>"
+    },
+    {
+      "name": "references",
+      "messages": [
+        {
+          "role": "user",
+          "content": "你好，GLM！"
+        },
+        {
+          "role": "assistant",
+          "content": "",
+          "reasoning_content": "search",
+          "tool_calls": [
+            {
+              "id": "a",
+              "function": {
+                "name": "lookup",
+                "arguments": {
+                  "q": "中文<&>\n",
+                  "list": [
+                    1,
+                    true,
+                    null
+                  ],
+                  "object": {
+                    "b": 2,
+                    "a": "x"
+                  }
+                }
+              }
+            },
+            {
+              "id": "b",
+              "function": {
+                "name": "lookup",
+                "arguments": {
+                  "q": "second",
+                  "float": 1.5
+                }
+              }
+            }
+          ]
+        },
+        {
+          "role": "tool",
+          "tool_call_id": "a",
+          "content": [
+            {
+              "type": "tool_reference",
+              "name": "hidden"
+            }
+          ]
+        }
+      ],
+      "tools": [
+        {
+          "type": "function",
+          "function": {
+            "name": "lookup",
+            "description": "查找<&>",
+            "parameters": {
+              "type": "object",
+              "properties": {
+                "q": {
+                  "type": "string"
+                }
+              },
+              "required": [
+                "q"
+              ]
+            },
+            "strict": true
+          }
+        },
+        {
+          "function": {
+            "name": "hidden",
+            "description": null,
+            "parameters": {},
+            "defer_loading": true
+          }
+        }
+      ],
+      "add_generation_prompt": true,
+      "kwargs": {},
+      "expected": "[gMASK]<sop><|system|>Reasoning Effort: Max<|system|>\n# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>\n{\"name\": \"lookup\", \"description\": \"查找<&>\", \"parameters\": {\"type\": \"object\", \"properties\": {\"q\": {\"type\": \"string\"}}, \"required\": [\"q\"]}}\n</tools>\n\nFor each function call, output the function name and arguments within the following XML format:\n<tool_call>{function-name}<arg_key>{arg-key-1}</arg_key><arg_value>{arg-value-1}</arg_value><arg_key>{arg-key-2}</arg_key><arg_value>{arg-value-2}</arg_value>...</tool_call><|user|>你好，GLM！<|assistant|><think>search</think><tool_call>lookup<arg_key>q</arg_key><arg_value>中文<&>\n</arg_value><arg_key>list</arg_key><arg_value>[1, true, null]</arg_value><arg_key>object</arg_key><arg_value>{\"b\": 2, \"a\": \"x\"}</arg_value></tool_call><tool_call>lookup<arg_key>q</arg_key><arg_value>second</arg_value><arg_key>float</arg_key><arg_value>1.5</arg_value></tool_call><|observation|><tool_response><tools>\n{\"name\": \"hidden\", \"description\": null, \"parameters\": {}}\n</tools></tool_response><|assistant|><think>"
+    },
+    {
+      "name": "bundled_references",
+      "messages": [
+        {
+          "role": "user",
+          "content": "你好，GLM！"
+        },
+        {
+          "role": "assistant",
+          "content": "",
+          "reasoning_content": "search",
+          "tool_calls": [
+            {
+              "id": "a",
+              "function": {
+                "name": "lookup",
+                "arguments": {
+                  "q": "中文<&>\n",
+                  "list": [
+                    1,
+                    true,
+                    null
+                  ],
+                  "object": {
+                    "b": 2,
+                    "a": "x"
+                  }
+                }
+              }
+            },
+            {
+              "id": "b",
+              "function": {
+                "name": "lookup",
+                "arguments": {
+                  "q": "second",
+                  "float": 1.5
+                }
+              }
+            }
+          ]
+        },
+        {
+          "role": "tool",
+          "content": [
+            {
+              "id": "a",
+              "output": [
+                {
+                  "type": "tool_reference",
+                  "name": "hidden"
+                }
+              ]
+            }
+          ]
+        }
+      ],
+      "tools": [
+        {
+          "type": "function",
+          "function": {
+            "name": "lookup",
+            "description": "查找<&>",
+            "parameters": {
+              "type": "object",
+              "properties": {
+                "q": {
+                  "type": "string"
+                }
+              },
+              "required": [
+                "q"
+              ]
+            },
+            "strict": true
+          }
+        },
+        {
+          "function": {
+            "name": "hidden",
+            "description": null,
+            "parameters": {},
+            "defer_loading": true
+          }
+        }
+      ],
+      "add_generation_prompt": true,
+      "kwargs": {},
+      "expected": "[gMASK]<sop><|system|>Reasoning Effort: Max<|system|>\n# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>\n{\"name\": \"lookup\", \"description\": \"查找<&>\", \"parameters\": {\"type\": \"object\", \"properties\": {\"q\": {\"type\": \"string\"}}, \"required\": [\"q\"]}}\n</tools>\n\nFor each function call, output the function name and arguments within the following XML format:\n<tool_call>{function-name}<arg_key>{arg-key-1}</arg_key><arg_value>{arg-value-1}</arg_value><arg_key>{arg-key-2}</arg_key><arg_value>{arg-value-2}</arg_value>...</tool_call><|user|>你好，GLM！<|assistant|><think>search</think><tool_call>lookup<arg_key>q</arg_key><arg_value>中文<&>\n</arg_value><arg_key>list</arg_key><arg_value>[1, true, null]</arg_value><arg_key>object</arg_key><arg_value>{\"b\": 2, \"a\": \"x\"}</arg_value></tool_call><tool_call>lookup<arg_key>q</arg_key><arg_value>second</arg_value><arg_key>float</arg_key><arg_value>1.5</arg_value></tool_call><|observation|><tool_response><tools>\n{\"name\": \"hidden\", \"description\": null, \"parameters\": {}}\n</tools></tool_response><|assistant|><think>"
+    },
+    {
+      "name": "orphan_tool_block",
+      "messages": [
+        {
+          "role": "tool",
+          "tool_call_id": "b",
+          "content": "B"
+        },
+        {
+          "role": "tool",
+          "tool_call_id": "a",
+          "content": "A"
+        }
+      ],
+      "tools": [
+        {
+          "type": "function",
+          "function": {
+            "name": "lookup",
+            "description": "查找<&>",
+            "parameters": {
+              "type": "object",
+              "properties": {
+                "q": {
+                  "type": "string"
+                }
+              },
+              "required": [
+                "q"
+              ]
+            },
+            "strict": true
+          }
+        },
+        {
+          "function": {
+            "name": "hidden",
+            "description": null,
+            "parameters": {},
+            "defer_loading": true
+          }
+        }
+      ],
+      "add_generation_prompt": true,
+      "kwargs": {},
+      "expected": "[gMASK]<sop><|system|>Reasoning Effort: Max<|system|>\n# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>\n{\"name\": \"lookup\", \"description\": \"查找<&>\", \"parameters\": {\"type\": \"object\", \"properties\": {\"q\": {\"type\": \"string\"}}, \"required\": [\"q\"]}}\n</tools>\n\nFor each function call, output the function name and arguments within the following XML format:\n<tool_call>{function-name}<arg_key>{arg-key-1}</arg_key><arg_value>{arg-value-1}</arg_value><arg_key>{arg-key-2}</arg_key><arg_value>{arg-value-2}</arg_value>...</tool_call><|observation|><tool_response>B</tool_response><tool_response>A</tool_response><|assistant|><think>"
+    },
+    {
+      "name": "duplicate_call",
+      "messages": [
+        {
+          "role": "user",
+          "content": "你好，GLM！"
+        },
+        {
+          "role": "assistant",
+          "content": "",
+          "reasoning_content": "search",
+          "tool_calls": [
+            {
+              "id": "a",
+              "function": {
+                "name": "lookup",
+                "arguments": {
+                  "q": "中文<&>\n",
+                  "list": [
+                    1,
+                    true,
+                    null
+                  ],
+                  "object": {
+                    "b": 2,
+                    "a": "x"
+                  }
+                }
+              }
+            },
+            {
+              "id": "a",
+              "function": {
+                "name": "lookup",
+                "arguments": {
+                  "q": "second",
+                  "float": 1.5
+                }
+              }
+            }
+          ]
+        },
+        {
+          "role": "tool",
+          "tool_call_id": "b",
+          "content": "B"
+        },
+        {
+          "role": "tool",
+          "tool_call_id": "a",
+          "content": "A"
+        }
+      ],
+      "tools": [
+        {
+          "type": "function",
+          "function": {
+            "name": "lookup",
+            "description": "查找<&>",
+            "parameters": {
+              "type": "object",
+              "properties": {
+                "q": {
+                  "type": "string"
+                }
+              },
+              "required": [
+                "q"
+              ]
+            },
+            "strict": true
+          }
+        },
+        {
+          "function": {
+            "name": "hidden",
+            "description": null,
+            "parameters": {},
+            "defer_loading": true
+          }
+        }
+      ],
+      "add_generation_prompt": true,
+      "kwargs": {},
+      "expected": "[gMASK]<sop><|system|>Reasoning Effort: Max<|system|>\n# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>\n{\"name\": \"lookup\", \"description\": \"查找<&>\", \"parameters\": {\"type\": \"object\", \"properties\": {\"q\": {\"type\": \"string\"}}, \"required\": [\"q\"]}}\n</tools>\n\nFor each function call, output the function name and arguments within the following XML format:\n<tool_call>{function-name}<arg_key>{arg-key-1}</arg_key><arg_value>{arg-value-1}</arg_value><arg_key>{arg-key-2}</arg_key><arg_value>{arg-value-2}</arg_value>...</tool_call><|user|>你好，GLM！<|assistant|><think>search</think><tool_call>lookup<arg_key>q</arg_key><arg_value>中文<&>\n</arg_value><arg_key>list</arg_key><arg_value>[1, true, null]</arg_value><arg_key>object</arg_key><arg_value>{\"b\": 2, \"a\": \"x\"}</arg_value></tool_call><tool_call>lookup<arg_key>q</arg_key><arg_value>second</arg_value><arg_key>float</arg_key><arg_value>1.5</arg_value></tool_call><|observation|><tool_response>B</tool_response><tool_response>A</tool_response><|assistant|><think>"
+    },
+    {
+      "name": "missing_call",
+      "messages": [
+        {
+          "role": "user",
+          "content": "你好，GLM！"
+        },
+        {
+          "role": "assistant",
+          "content": "",
+          "reasoning_content": "search",
+          "tool_calls": [
+            {
+              "id": "a",
+              "function": {
+                "name": "lookup",
+                "arguments": {
+                  "q": "中文<&>\n",
+                  "list": [
+                    1,
+                    true,
+                    null
+                  ],
+                  "object": {
+                    "b": 2,
+                    "a": "x"
+                  }
+                }
+              }
+            },
+            {
+              "id": null,
+              "function": {
+                "name": "lookup",
+                "arguments": {
+                  "q": "second",
+                  "float": 1.5
+                }
+              }
+            }
+          ]
+        },
+        {
+          "role": "tool",
+          "tool_call_id": "b",
+          "content": "B"
+        },
+        {
+          "role": "tool",
+          "tool_call_id": "a",
+          "content": "A"
+        }
+      ],
+      "tools": [
+        {
+          "type": "function",
+          "function": {
+            "name": "lookup",
+            "description": "查找<&>",
+            "parameters": {
+              "type": "object",
+              "properties": {
+                "q": {
+                  "type": "string"
+                }
+              },
+              "required": [
+                "q"
+              ]
+            },
+            "strict": true
+          }
+        },
+        {
+          "function": {
+            "name": "hidden",
+            "description": null,
+            "parameters": {},
+            "defer_loading": true
+          }
+        }
+      ],
+      "add_generation_prompt": true,
+      "kwargs": {},
+      "expected": "[gMASK]<sop><|system|>Reasoning Effort: Max<|system|>\n# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>\n{\"name\": \"lookup\", \"description\": \"查找<&>\", \"parameters\": {\"type\": \"object\", \"properties\": {\"q\": {\"type\": \"string\"}}, \"required\": [\"q\"]}}\n</tools>\n\nFor each function call, output the function name and arguments within the following XML format:\n<tool_call>{function-name}<arg_key>{arg-key-1}</arg_key><arg_value>{arg-value-1}</arg_value><arg_key>{arg-key-2}</arg_key><arg_value>{arg-value-2}</arg_value>...</tool_call><|user|>你好，GLM！<|assistant|><think>search</think><tool_call>lookup<arg_key>q</arg_key><arg_value>中文<&>\n</arg_value><arg_key>list</arg_key><arg_value>[1, true, null]</arg_value><arg_key>object</arg_key><arg_value>{\"b\": 2, \"a\": \"x\"}</arg_value></tool_call><tool_call>lookup<arg_key>q</arg_key><arg_value>second</arg_value><arg_key>float</arg_key><arg_value>1.5</arg_value></tool_call><|observation|><tool_response>B</tool_response><tool_response>A</tool_response><|assistant|><think>"
+    },
+    {
+      "name": "multiple_blocks",
+      "messages": [
+        {
+          "role": "user",
+          "content": "你好，GLM！"
+        },
+        {
+          "role": "assistant",
+          "content": "",
+          "reasoning_content": "search",
+          "tool_calls": [
+            {
+              "id": "a",
+              "function": {
+                "name": "lookup",
+                "arguments": {
+                  "q": "中文<&>\n",
+                  "list": [
+                    1,
+                    true,
+                    null
+                  ],
+                  "object": {
+                    "b": 2,
+                    "a": "x"
+                  }
+                }
+              }
+            },
+            {
+              "id": "b",
+              "function": {
+                "name": "lookup",
+                "arguments": {
+                  "q": "second",
+                  "float": 1.5
+                }
+              }
+            }
+          ]
+        },
+        {
+          "role": "tool",
+          "tool_call_id": "b",
+          "content": "B"
+        },
+        {
+          "role": "tool",
+          "tool_call_id": "a",
+          "content": "A"
+        },
+        {
+          "role": "user",
+          "content": "你好，GLM！"
+        },
+        {
+          "role": "assistant",
+          "content": "",
+          "reasoning_content": "search",
+          "tool_calls": [
+            {
+              "id": "a",
+              "function": {
+                "name": "lookup",
+                "arguments": {
+                  "q": "中文<&>\n",
+                  "list": [
+                    1,
+                    true,
+                    null
+                  ],
+                  "object": {
+                    "b": 2,
+                    "a": "x"
+                  }
+                }
+              }
+            },
+            {
+              "id": "b",
+              "function": {
+                "name": "lookup",
+                "arguments": {
+                  "q": "second",
+                  "float": 1.5
+                }
+              }
+            }
+          ]
+        },
+        {
+          "role": "tool",
+          "tool_call_id": "b",
+          "content": "B"
+        },
+        {
+          "role": "tool",
+          "tool_call_id": "a",
+          "content": "A"
+        }
+      ],
+      "tools": [
+        {
+          "type": "function",
+          "function": {
+            "name": "lookup",
+            "description": "查找<&>",
+            "parameters": {
+              "type": "object",
+              "properties": {
+                "q": {
+                  "type": "string"
+                }
+              },
+              "required": [
+                "q"
+              ]
+            },
+            "strict": true
+          }
+        },
+        {
+          "function": {
+            "name": "hidden",
+            "description": null,
+            "parameters": {},
+            "defer_loading": true
+          }
+        }
+      ],
+      "add_generation_prompt": true,
+      "kwargs": {},
+      "expected": "[gMASK]<sop><|system|>Reasoning Effort: Max<|system|>\n# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>\n{\"name\": \"lookup\", \"description\": \"查找<&>\", \"parameters\": {\"type\": \"object\", \"properties\": {\"q\": {\"type\": \"string\"}}, \"required\": [\"q\"]}}\n</tools>\n\nFor each function call, output the function name and arguments within the following XML format:\n<tool_call>{function-name}<arg_key>{arg-key-1}</arg_key><arg_value>{arg-value-1}</arg_value><arg_key>{arg-key-2}</arg_key><arg_value>{arg-value-2}</arg_value>...</tool_call><|user|>你好，GLM！<|assistant|><think>search</think><tool_call>lookup<arg_key>q</arg_key><arg_value>中文<&>\n</arg_value><arg_key>list</arg_key><arg_value>[1, true, null]</arg_value><arg_key>object</arg_key><arg_value>{\"b\": 2, \"a\": \"x\"}</arg_value></tool_call><tool_call>lookup<arg_key>q</arg_key><arg_value>second</arg_value><arg_key>float</arg_key><arg_value>1.5</arg_value></tool_call><|observation|><tool_response>A</tool_response><tool_response>B</tool_response><|user|>你好，GLM！<|assistant|><think>search</think><tool_call>lookup<arg_key>q</arg_key><arg_value>中文<&>\n</arg_value><arg_key>list</arg_key><arg_value>[1, true, null]</arg_value><arg_key>object</arg_key><arg_value>{\"b\": 2, \"a\": \"x\"}</arg_value></tool_call><tool_call>lookup<arg_key>q</arg_key><arg_value>second</arg_value><arg_key>float</arg_key><arg_value>1.5</arg_value></tool_call><|observation|><tool_response>A</tool_response><tool_response>B</tool_response><|assistant|><think>"
+    },
+    {
+      "name": "numeric_and_string_serialization",
+      "messages": [
+        {
+          "role": "user",
+          "content": "你好，GLM！"
+        },
+        {
+          "role": "assistant",
+          "content": "",
+          "reasoning_content": "search",
+          "tool_calls": [
+            {
+              "id": "a",
+              "function": {
+                "name": "lookup",
+                "arguments": {
+                  "floats": [
+                    1e-07,
+                    1e-06,
+                    1e-05,
+                    1000000000000000.0,
+                    1e+16,
+                    1e+20,
+                    -0.0,
+                    0.0001,
+                    1.2345678901234567,
+                    1.234567890123456e-100,
+                    5e-324,
+                    2.2250738585072014e-308
+                  ],
+                  "unicode": "中文 😀 \\ \" \t\n"
+                }
+              }
+            },
+            {
+              "id": "b",
+              "function": {
+                "name": "lookup",
+                "arguments": {
+                  "q": "second",
+                  "float": 1.5
+                }
+              }
+            }
+          ]
+        }
+      ],
+      "tools": [
+        {
+          "type": "function",
+          "function": {
+            "name": "lookup",
+            "description": "查找<&>",
+            "parameters": {
+              "type": "object",
+              "properties": {
+                "q": {
+                  "type": "string"
+                }
+              },
+              "required": [
+                "q"
+              ]
+            },
+            "strict": true
+          }
+        },
+        {
+          "function": {
+            "name": "hidden",
+            "description": null,
+            "parameters": {},
+            "defer_loading": true
+          }
+        }
+      ],
+      "add_generation_prompt": true,
+      "kwargs": {},
+      "expected": "[gMASK]<sop><|system|>Reasoning Effort: Max<|system|>\n# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>\n{\"name\": \"lookup\", \"description\": \"查找<&>\", \"parameters\": {\"type\": \"object\", \"properties\": {\"q\": {\"type\": \"string\"}}, \"required\": [\"q\"]}}\n</tools>\n\nFor each function call, output the function name and arguments within the following XML format:\n<tool_call>{function-name}<arg_key>{arg-key-1}</arg_key><arg_value>{arg-value-1}</arg_value><arg_key>{arg-key-2}</arg_key><arg_value>{arg-value-2}</arg_value>...</tool_call><|user|>你好，GLM！<|assistant|><think>search</think><tool_call>lookup<arg_key>floats</arg_key><arg_value>[1e-07, 1e-06, 1e-05, 1000000000000000.0, 1e+16, 1e+20, -0.0, 0.0001, 1.2345678901234567, 1.234567890123456e-100, 5e-324, 2.2250738585072014e-308]</arg_value><arg_key>unicode</arg_key><arg_value>中文 😀 \\ \" \t\n</arg_value></tool_call><tool_call>lookup<arg_key>q</arg_key><arg_value>second</arg_value><arg_key>float</arg_key><arg_value>1.5</arg_value></tool_call><|assistant|><think>"
+    },
+    {
+      "name": "bad_arguments_None",
+      "messages": [
+        {
+          "role": "user",
+          "content": "你好，GLM！"
+        },
+        {
+          "role": "assistant",
+          "content": "",
+          "reasoning_content": "search",
+          "tool_calls": [
+            {
+              "id": "a",
+              "function": {
+                "name": "lookup",
+                "arguments": null
+              }
+            },
+            {
+              "id": "b",
+              "function": {
+                "name": "lookup",
+                "arguments": {
+                  "q": "second",
+                  "float": 1.5
+                }
+              }
+            }
+          ]
+        }
+      ],
+      "tools": [
+        {
+          "type": "function",
+          "function": {
+            "name": "lookup",
+            "description": "查找<&>",
+            "parameters": {
+              "type": "object",
+              "properties": {
+                "q": {
+                  "type": "string"
+                }
+              },
+              "required": [
+                "q"
+              ]
+            },
+            "strict": true
+          }
+        },
+        {
+          "function": {
+            "name": "hidden",
+            "description": null,
+            "parameters": {},
+            "defer_loading": true
+          }
+        }
+      ],
+      "add_generation_prompt": true,
+      "kwargs": {},
+      "error": "ValueError"
+    },
+    {
+      "name": "bad_arguments_[]",
+      "messages": [
+        {
+          "role": "user",
+          "content": "你好，GLM！"
+        },
+        {
+          "role": "assistant",
+          "content": "",
+          "reasoning_content": "search",
+          "tool_calls": [
+            {
+              "id": "a",
+              "function": {
+                "name": "lookup",
+                "arguments": []
+              }
+            },
+            {
+              "id": "b",
+              "function": {
+                "name": "lookup",
+                "arguments": {
+                  "q": "second",
+                  "float": 1.5
+                }
+              }
+            }
+          ]
+        }
+      ],
+      "tools": [
+        {
+          "type": "function",
+          "function": {
+            "name": "lookup",
+            "description": "查找<&>",
+            "parameters": {
+              "type": "object",
+              "properties": {
+                "q": {
+                  "type": "string"
+                }
+              },
+              "required": [
+                "q"
+              ]
+            },
+            "strict": true
+          }
+        },
+        {
+          "function": {
+            "name": "hidden",
+            "description": null,
+            "parameters": {},
+            "defer_loading": true
+          }
+        }
+      ],
+      "add_generation_prompt": true,
+      "kwargs": {},
+      "error": "ValueError"
+    },
+    {
+      "name": "bad_arguments_invalid",
+      "messages": [
+        {
+          "role": "user",
+          "content": "你好，GLM！"
+        },
+        {
+          "role": "assistant",
+          "content": "",
+          "reasoning_content": "search",
+          "tool_calls": [
+            {
+              "id": "a",
+              "function": {
+                "name": "lookup",
+                "arguments": "invalid"
+              }
+            },
+            {
+              "id": "b",
+              "function": {
+                "name": "lookup",
+                "arguments": {
+                  "q": "second",
+                  "float": 1.5
+                }
+              }
+            }
+          ]
+        }
+      ],
+      "tools": [
+        {
+          "type": "function",
+          "function": {
+            "name": "lookup",
+            "description": "查找<&>",
+            "parameters": {
+              "type": "object",
+              "properties": {
+                "q": {
+                  "type": "string"
+                }
+              },
+              "required": [
+                "q"
+              ]
+            },
+            "strict": true
+          }
+        },
+        {
+          "function": {
+            "name": "hidden",
+            "description": null,
+            "parameters": {},
+            "defer_loading": true
+          }
+        }
+      ],
+      "add_generation_prompt": true,
+      "kwargs": {},
+      "error": "ValueError"
+    },
+    {
+      "name": "bad_arguments_[]",
+      "messages": [
+        {
+          "role": "user",
+          "content": "你好，GLM！"
+        },
+        {
+          "role": "assistant",
+          "content": "",
+          "reasoning_content": "search",
+          "tool_calls": [
+            {
+              "id": "a",
+              "function": {
+                "name": "lookup",
+                "arguments": "[]"
+              }
+            },
+            {
+              "id": "b",
+              "function": {
+                "name": "lookup",
+                "arguments": {
+                  "q": "second",
+                  "float": 1.5
+                }
+              }
+            }
+          ]
+        }
+      ],
+      "tools": [
+        {
+          "type": "function",
+          "function": {
+            "name": "lookup",
+            "description": "查找<&>",
+            "parameters": {
+              "type": "object",
+              "properties": {
+                "q": {
+                  "type": "string"
+                }
+              },
+              "required": [
+                "q"
+              ]
+            },
+            "strict": true
+          }
+        },
+        {
+          "function": {
+            "name": "hidden",
+            "description": null,
+            "parameters": {},
+            "defer_loading": true
+          }
+        }
+      ],
+      "add_generation_prompt": true,
+      "kwargs": {},
+      "error": "ValueError"
+    },
+    {
+      "name": "bad_arguments_5",
+      "messages": [
+        {
+          "role": "user",
+          "content": "你好，GLM！"
+        },
+        {
+          "role": "assistant",
+          "content": "",
+          "reasoning_content": "search",
+          "tool_calls": [
+            {
+              "id": "a",
+              "function": {
+                "name": "lookup",
+                "arguments": 5
+              }
+            },
+            {
+              "id": "b",
+              "function": {
+                "name": "lookup",
+                "arguments": {
+                  "q": "second",
+                  "float": 1.5
+                }
+              }
+            }
+          ]
+        }
+      ],
+      "tools": [
+        {
+          "type": "function",
+          "function": {
+            "name": "lookup",
+            "description": "查找<&>",
+            "parameters": {
+              "type": "object",
+              "properties": {
+                "q": {
+                  "type": "string"
+                }
+              },
+              "required": [
+                "q"
+              ]
+            },
+            "strict": true
+          }
+        },
+        {
+          "function": {
+            "name": "hidden",
+            "description": null,
+            "parameters": {},
+            "defer_loading": true
+          }
+        }
+      ],
+      "add_generation_prompt": true,
+      "kwargs": {},
+      "error": "ValueError"
+    },
+    {
+      "name": "reference_without_tools",
+      "messages": [
+        {
+          "role": "tool",
+          "content": [
+            {
+              "type": "tool_reference",
+              "name": "hidden"
+            }
+          ]
+        }
+      ],
+      "tools": [],
+      "add_generation_prompt": true,
+      "kwargs": {},
+      "error": "ValueError"
+    },
+    {
+      "name": "bad_tool_content",
+      "messages": [
+        {
+          "role": "tool",
+          "content": {
+            "bad": 1
+          }
+        }
+      ],
+      "tools": [
+        {
+          "type": "function",
+          "function": {
+            "name": "lookup",
+            "description": "查找<&>",
+            "parameters": {
+              "type": "object",
+              "properties": {
+                "q": {
+                  "type": "string"
+                }
+              },
+              "required": [
+                "q"
+              ]
+            },
+            "strict": true
+          }
+        },
+        {
+          "function": {
+            "name": "hidden",
+            "description": null,
+            "parameters": {},
+            "defer_loading": true
+          }
+        }
+      ],
+      "add_generation_prompt": true,
+      "kwargs": {},
+      "error": "ValueError"
+    },
+    {
+      "name": "bad_tool_definition",
+      "messages": [
+        {
+          "role": "user",
+          "content": "你好，GLM！"
+        }
+      ],
+      "tools": [
+        "bad"
+      ],
+      "add_generation_prompt": true,
+      "kwargs": {},
+      "error": "ValueError"
+    },
+    {
+      "name": "continue_plain",
+      "messages": [
+        {
+          "role": "user",
+          "content": "你好，GLM！"
+        },
+        {
+          "role": "assistant",
+          "content": "部分回答"
+        }
+      ],
+      "tools": [],
+      "add_generation_prompt": false,
+      "continue_final_message": true,
+      "kwargs": {},
+      "expected": "[gMASK]<sop><|system|>Reasoning Effort: Max<|user|>你好，GLM！<|assistant|><think></think>部分回答"
+    },
+    {
+      "name": "continue_spaces",
+      "messages": [
+        {
+          "role": "user",
+          "content": "你好，GLM！"
+        },
+        {
+          "role": "assistant",
+          "content": "  partial answer \n"
+        }
+      ],
+      "tools": [],
+      "add_generation_prompt": false,
+      "continue_final_message": true,
+      "kwargs": {},
+      "expected": "[gMASK]<sop><|system|>Reasoning Effort: Max<|user|>你好，GLM！<|assistant|><think></think>partial answer"
+    },
+    {
+      "name": "continue_unicode_spaces",
+      "messages": [
+        {
+          "role": "user",
+          "content": "你好，GLM！"
+        },
+        {
+          "role": "assistant",
+          "content": "　x …"
+        }
+      ],
+      "tools": [],
+      "add_generation_prompt": false,
+      "continue_final_message": true,
+      "kwargs": {},
+      "expected": "[gMASK]<sop><|system|>Reasoning Effort: Max<|user|>你好，GLM！<|assistant|><think></think>x"
+    },
+    {
+      "name": "continue_python_whitespace",
+      "messages": [
+        {
+          "role": "user",
+          "content": "你好，GLM！"
+        },
+        {
+          "role": "assistant",
+          "content": "\u001c\u001d answer\u001e\u001f"
+        }
+      ],
+      "tools": [],
+      "add_generation_prompt": false,
+      "continue_final_message": true,
+      "kwargs": {},
+      "expected": "[gMASK]<sop><|system|>Reasoning Effort: Max<|user|>你好，GLM！<|assistant|><think></think>answer"
+    },
+    {
+      "name": "continue_inline_think",
+      "messages": [
+        {
+          "role": "user",
+          "content": "你好，GLM！"
+        },
+        {
+          "role": "assistant",
+          "content": "<think>inline</think>partial"
+        }
+      ],
+      "tools": [],
+      "add_generation_prompt": false,
+      "continue_final_message": true,
+      "kwargs": {},
+      "expected": "[gMASK]<sop><|system|>Reasoning Effort: Max<|user|>你好，GLM！<|assistant|><think>inline</think>partial"
+    },
+    {
+      "name": "continue_inline_close_only",
+      "messages": [
+        {
+          "role": "user",
+          "content": "你好，GLM！"
+        },
+        {
+          "role": "assistant",
+          "content": "partial</think>more"
+        }
+      ],
+      "tools": [],
+      "add_generation_prompt": false,
+      "continue_final_message": true,
+      "kwargs": {},
+      "expected": "[gMASK]<sop><|system|>Reasoning Effort: Max<|user|>你好，GLM！<|assistant|><think>partial</think>more"
+    },
+    {
+      "name": "continue_inline_think_with_spaces",
+      "messages": [
+        {
+          "role": "user",
+          "content": "你好，GLM！"
+        },
+        {
+          "role": "assistant",
+          "content": "<think>a</think>  b"
+        }
+      ],
+      "tools": [],
+      "add_generation_prompt": false,
+      "continue_final_message": true,
+      "kwargs": {},
+      "error": "ValueError"
+    },
+    {
+      "name": "continue_reasoning_text_tool",
+      "messages": [
+        {
+          "role": "user",
+          "content": "你好，GLM！"
+        },
+        {
+          "role": "assistant",
+          "content": "partial",
+          "reasoning_content": "reasoning",
+          "tool_calls": [
+            {
+              "id": "call_0",
+              "function": {
+                "name": "lookup",
+                "arguments": {
+                  "q": "x",
+                  "n": 1
+                }
+              }
+            }
+          ]
+        }
+      ],
+      "tools": [],
+      "add_generation_prompt": false,
+      "continue_final_message": true,
+      "kwargs": {},
+      "expected": "[gMASK]<sop><|system|>Reasoning Effort: Max<|user|>你好，GLM！<|assistant|><think>reasoning</think>partial"
+    },
+    {
+      "name": "continue_reasoning_only",
+      "messages": [
+        {
+          "role": "user",
+          "content": "你好，GLM！"
+        },
+        {
+          "role": "assistant",
+          "content": "",
+          "reasoning_content": "r"
+        }
+      ],
+      "tools": [],
+      "add_generation_prompt": false,
+      "continue_final_message": true,
+      "kwargs": {},
+      "expected": "[gMASK]<sop><|system|>Reasoning Effort: Max<|user|>你好，GLM！<|assistant|><think>r</think>"
+    },
+    {
+      "name": "continue_empty_reasoning",
+      "messages": [
+        {
+          "role": "user",
+          "content": "你好，GLM！"
+        },
+        {
+          "role": "assistant",
+          "content": "x ",
+          "reasoning_content": ""
+        }
+      ],
+      "tools": [],
+      "add_generation_prompt": false,
+      "continue_final_message": true,
+      "kwargs": {},
+      "expected": "[gMASK]<sop><|system|>Reasoning Effort: Max<|user|>你好，GLM！<|assistant|><think></think>x"
+    },
+    {
+      "name": "continue_tool_call_only",
+      "messages": [
+        {
+          "role": "user",
+          "content": "你好，GLM！"
+        },
+        {
+          "role": "assistant",
+          "content": "",
+          "tool_calls": [
+            {
+              "id": "call_0",
+              "function": {
+                "name": "lookup",
+                "arguments": {
+                  "q": "x",
+                  "n": 1
+                }
+              }
+            }
+          ]
+        }
+      ],
+      "tools": [],
+      "add_generation_prompt": false,
+      "continue_final_message": true,
+      "kwargs": {},
+      "expected": "[gMASK]<sop><|system|>Reasoning Effort: Max<|user|>你好，GLM！<|assistant|><think></think>"
+    },
+    {
+      "name": "continue_whitespace_only",
+      "messages": [
+        {
+          "role": "user",
+          "content": "你好，GLM！"
+        },
+        {
+          "role": "assistant",
+          "content": "   "
+        }
+      ],
+      "tools": [],
+      "add_generation_prompt": false,
+      "continue_final_message": true,
+      "kwargs": {},
+      "expected": "[gMASK]<sop><|system|>Reasoning Effort: Max<|user|>你好，GLM！<|assistant|><think></think>"
+    },
+    {
+      "name": "continue_no_user",
+      "messages": [
+        {
+          "role": "assistant",
+          "content": "partial",
+          "reasoning_content": "r"
+        }
+      ],
+      "tools": [],
+      "add_generation_prompt": false,
+      "continue_final_message": true,
+      "kwargs": {},
+      "expected": "[gMASK]<sop><|system|>Reasoning Effort: Max<|assistant|><think>r</think>partial"
+    },
+    {
+      "name": "continue_clear_thinking",
+      "messages": [
+        {
+          "role": "user",
+          "content": "你好，GLM！"
+        },
+        {
+          "role": "assistant",
+          "content": "old",
+          "reasoning_content": "old r"
+        },
+        {
+          "role": "user",
+          "content": "你好，GLM！"
+        },
+        {
+          "role": "assistant",
+          "content": "new",
+          "reasoning_content": "new r"
+        }
+      ],
+      "tools": [],
+      "add_generation_prompt": false,
+      "continue_final_message": true,
+      "kwargs": {
+        "clear_thinking": true
+      },
+      "expected": "[gMASK]<sop><|system|>Reasoning Effort: Max<|user|>你好，GLM！<|assistant|><think></think>old<|user|>你好，GLM！<|assistant|><think>new r</think>new"
+    },
+    {
+      "name": "continue_after_tool_response",
+      "messages": [
+        {
+          "role": "user",
+          "content": "你好，GLM！"
+        },
+        {
+          "role": "assistant",
+          "content": "",
+          "reasoning_content": "r",
+          "tool_calls": [
+            {
+              "id": "call_0",
+              "function": {
+                "name": "lookup",
+                "arguments": {
+                  "q": "x",
+                  "n": 1
+                }
+              }
+            }
+          ]
+        },
+        {
+          "role": "tool",
+          "content": "结果",
+          "tool_call_id": "call_0"
+        },
+        {
+          "role": "assistant",
+          "content": "根据结果"
+        }
+      ],
+      "tools": [],
+      "add_generation_prompt": false,
+      "continue_final_message": true,
+      "kwargs": {
+        "reasoning_effort": "low"
+      },
+      "expected": "[gMASK]<sop><|system|>Reasoning Effort: Low<|user|>你好，GLM！<|assistant|><think>r</think><tool_call>lookup<arg_key>q</arg_key><arg_value>x</arg_value><arg_key>n</arg_key><arg_value>1</arg_value></tool_call><|observation|><tool_response>结果</tool_response><|assistant|><think></think>根据结果"
+    }
+  ]
+}
diff --git a/rust/src/chat/src/renderer/glm5_next/mod.rs b/rust/src/chat/src/renderer/glm5_next/mod.rs
new file mode 100644
index 000000000000..10b2960db4f1
--- /dev/null
+++ b/rust/src/chat/src/renderer/glm5_next/mod.rs
@@ -0,0 +1,74 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+mod encoding;
+mod writer;
+
+use std::collections::HashMap;
+use std::sync::Arc;
+
+use serde_json::Value;
+use vllm_text::tokenizer::DynTokenizer;
+
+use super::{ChatRenderer, RenderedPrompt};
+use crate::error::Result;
+use crate::request::ChatRequest;
+
+#[derive(Clone, Default)]
+pub struct Glm5NextChatRenderer {
+    default_template_kwargs: HashMap<String, Value>,
+    multimodal: bool,
+    tokenizer: Option<Arc<writer::ControlTokenizer>>,
+}
+
+impl Glm5NextChatRenderer {
+    pub fn new(default_template_kwargs: HashMap<String, Value>) -> Self {
+        Self {
+            default_template_kwargs,
+            multimodal: false,
+            tokenizer: None,
+        }
+    }
+
+    /// Render token IDs: control tokens use added-vocabulary IDs and literal
+    /// text uses the base vocabulary.
+    pub fn with_tokenizer(mut self, tokenizer: DynTokenizer) -> Result<Self> {
+        self.tokenizer = Some(Arc::new(writer::ControlTokenizer::new(tokenizer)?));
+        Ok(self)
+    }
+
+    pub fn with_multimodal(mut self, enabled: bool) -> Self {
+        self.multimodal = enabled;
+        self
+    }
+}
+
+impl ChatRenderer for Glm5NextChatRenderer {
+    fn render(&self, request: &ChatRequest) -> Result<RenderedPrompt> {
+        request.validate()?;
+        let mut kwargs = self.default_template_kwargs.clone();
+        kwargs.extend(request.chat_options.template_kwargs.clone());
+        if let Some(effort) = &request.chat_options.reasoning_effort {
+            kwargs.insert("reasoning_effort".into(), serde_json::json!(effort));
+        }
+        let messages = encoding::messages(&request.messages);
+        let tools = encoding::tools(request.tools());
+        let mut out = writer::Glm5PromptWriter::new(self.tokenizer.as_deref());
+        let media_order = encoding::render_into(
+            &mut out,
+            &messages,
+            &tools,
+            request.chat_options.generation_prompt_mode,
+            &kwargs,
+            self.multimodal,
+        )?;
+        Ok(RenderedPrompt {
+            prompt: out.finish()?,
+            media_order: Some(media_order),
+            effective_template_kwargs: kwargs,
+        })
+    }
+}
+
+#[cfg(test)]
+mod tests;
diff --git a/rust/src/chat/src/renderer/glm5_next/tests.rs b/rust/src/chat/src/renderer/glm5_next/tests.rs
new file mode 100644
index 000000000000..fa87129db685
--- /dev/null
+++ b/rust/src/chat/src/renderer/glm5_next/tests.rs
@@ -0,0 +1,492 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+use std::collections::HashMap;
+use std::sync::Arc;
+
+use serde::Deserialize;
+use serde_json::{Value, json};
+use vllm_text::Prompt;
+use vllm_tokenizer::{DynTokenizer, HuggingFaceTokenizer};
+
+use super::writer::{ControlTokenizer, Glm5PromptWriter};
+use super::{Glm5NextChatRenderer, encoding};
+use crate::renderer::hf::{ChatTemplateContentFormatOption, HfChatRenderer, MultimodalRenderInfo};
+use crate::request::{
+    ChatMessage, ChatRequest, ChatTool, GenerationPromptMode, ResolvedToolContext,
+};
+use crate::{AssistantContentBlock, AssistantToolCall, ChatRenderer, EffortValue};
+
+#[derive(Deserialize)]
+struct TestCase {
+    name: String,
+    messages: Vec<Value>,
+    tools: Vec<Value>,
+    add_generation_prompt: bool,
+    #[serde(default)]
+    continue_final_message: bool,
+    kwargs: HashMap<String, Value>,
+    expected: Option<String>,
+    error: Option<String>,
+}
+
+#[derive(Deserialize)]
+struct TestCases {
+    cases: Vec<TestCase>,
+}
+
+#[test]
+fn renders_expected_prompts_and_rejects_invalid_input() {
+    let fixtures: TestCases =
+        serde_json::from_str(include_str!("fixtures/test_cases.json")).unwrap();
+    for case in fixtures.cases {
+        let result = encoding::render(
+            &case.messages,
+            &case.tools,
+            if case.continue_final_message {
+                GenerationPromptMode::ContinueFinalAssistant
+            } else if case.add_generation_prompt {
+                GenerationPromptMode::StartNewAssistant
+            } else {
+                GenerationPromptMode::NoGenerationPrompt
+            },
+            &case.kwargs,
+        );
+        if case.error.is_some() {
+            assert!(
+                result.is_err(),
+                "{}: expected rejection, got {result:?}",
+                case.name
+            );
+        } else {
+            assert_eq!(result.unwrap(), case.expected.unwrap(), "{}", case.name);
+        }
+    }
+}
+
+fn tool_request(count: usize) -> ChatRequest {
+    let calls = (0..count)
+        .map(|index| {
+            AssistantContentBlock::ToolCall(AssistantToolCall {
+                id: format!("call_{index}"),
+                name: "lookup".into(),
+                arguments: json!({"q": format!("城市 {index}"), "values": [true, null, 1.5]})
+                    .to_string(),
+            })
+        })
+        .collect();
+    let mut messages = vec![
+        ChatMessage::system("助手"),
+        ChatMessage::user("查询天气"),
+        ChatMessage::assistant_blocks(calls),
+    ];
+    messages.extend(
+        (0..count).rev().map(|index| {
+            ChatMessage::tool_response(format!("结果 {index}"), format!("call_{index}"))
+        }),
+    );
+    let tools = vec![ChatTool {
+        name: "lookup".into(),
+        description: Some("查询天气".into()),
+        parameters: json!({"type": "object", "properties": {"q": {"type": "string"}}}),
+        strict: Some(true),
+        defer_loading: None,
+    }];
+    let tool_context = ResolvedToolContext::new(&messages, tools, None, true).unwrap();
+    ChatRequest {
+        messages,
+        tool_context,
+        ..ChatRequest::for_test()
+    }
+}
+
+const CHECKPOINT_TEMPLATES: [(&str, &str); 2] = [
+    ("GLM-5.3", include_str!("fixtures/glm_5_3.jinja")),
+    (
+        "GLM-5.3-Flash",
+        include_str!("fixtures/glm_5_3_flash.jinja"),
+    ),
+];
+
+fn hf_renderer(template: &str) -> HfChatRenderer {
+    HfChatRenderer::new(
+        Some(template.into()),
+        HashMap::new(),
+        ChatTemplateContentFormatOption::OpenAi,
+    )
+    .unwrap()
+}
+
+#[test]
+fn typed_requests_match_existing_hf_renderer() {
+    let native = Glm5NextChatRenderer::default();
+    for (name, template) in CHECKPOINT_TEMPLATES {
+        let hf = hf_renderer(template);
+        for count in [1, 2, 8] {
+            let mut request = tool_request(count);
+            for clear in [false, true] {
+                request
+                    .chat_options
+                    .template_kwargs
+                    .insert("clear_thinking".into(), json!(clear));
+                assert_eq!(
+                    native.render(&request).unwrap().prompt,
+                    hf.render(&request).unwrap().prompt,
+                    "{name}"
+                );
+            }
+            request.tool_context.tool_choice = crate::request::ChatToolChoice::None;
+            assert_eq!(
+                native.render(&request).unwrap().prompt,
+                hf.render(&request).unwrap().prompt
+            );
+        }
+    }
+}
+
+#[test]
+fn flash_media_matches_checkpoint_template() {
+    let hf = hf_renderer(CHECKPOINT_TEMPLATES[1].1).with_multimodal(Some(MultimodalRenderInfo {
+        image_token: Some("<|image|>".into()),
+        video_token: Some("<|video|>".into()),
+        audio_token: Some("<|begin_of_audio|>".into()),
+    }));
+    let mut request = tool_request(2);
+    request.messages.push(ChatMessage::user(vec![
+        crate::request::ChatContentPart::image_url("unused"),
+    ]));
+    let expected = hf.render(&request).unwrap();
+    let renderer = Glm5NextChatRenderer::default().with_multimodal(true);
+    let actual = renderer.render(&request).unwrap();
+    assert_eq!(actual.prompt, expected.prompt);
+    assert_eq!(
+        actual.media_order.as_ref().unwrap(),
+        &[crate::renderer::MediaPartSource {
+            message_index: 5,
+            content_part_index: 0
+        }]
+    );
+    assert!(
+        actual
+            .prompt
+            .into_text()
+            .unwrap()
+            .contains("<|begin_of_image|><|image|><|end_of_image|>")
+    );
+}
+
+#[test]
+fn media_order_tracks_reordered_tool_responses() {
+    use crate::renderer::MediaPartSource;
+    use crate::request::ChatContentPart;
+    let mut request = tool_request(2);
+    request.messages[3] =
+        ChatMessage::tool_response(vec![ChatContentPart::image_url("second")], "call_1");
+    request.messages[4] =
+        ChatMessage::tool_response(vec![ChatContentPart::image_url("first")], "call_0");
+    let native = Glm5NextChatRenderer::default().with_multimodal(true).render(&request).unwrap();
+    let hf = hf_renderer(CHECKPOINT_TEMPLATES[1].1)
+        .with_multimodal(Some(MultimodalRenderInfo {
+            image_token: Some("<|image|>".into()),
+            video_token: Some("<|video|>".into()),
+            audio_token: None,
+        }))
+        .render(&request)
+        .unwrap();
+    assert_eq!(native.prompt, hf.prompt);
+    assert_eq!(
+        native.media_order.unwrap(),
+        vec![
+            MediaPartSource {
+                message_index: 4,
+                content_part_index: 0
+            },
+            MediaPartSource {
+                message_index: 3,
+                content_part_index: 0
+            }
+        ]
+    );
+}
+
+#[test]
+fn defaults_request_kwargs_and_typed_effort_have_defined_precedence() {
+    let renderer = Glm5NextChatRenderer::new(HashMap::from([
+        ("reasoning_effort".into(), json!("low")),
+        ("clear_thinking".into(), json!(true)),
+    ]));
+    let mut request = ChatRequest::for_test();
+    request
+        .chat_options
+        .template_kwargs
+        .insert("reasoning_effort".into(), json!("high"));
+    let rendered = renderer.render(&request).unwrap();
+    assert!(rendered.prompt.into_text().unwrap().contains("Reasoning Effort: High"));
+    request.chat_options.reasoning_effort = Some(EffortValue::from("low"));
+    let rendered = renderer.render(&request).unwrap();
+    assert!(rendered.prompt.into_text().unwrap().contains("Reasoning Effort: Low"));
+    assert_eq!(
+        rendered.effective_template_kwargs["reasoning_effort"],
+        "low"
+    );
+    assert_eq!(rendered.effective_template_kwargs["clear_thinking"], true);
+}
+
+#[test]
+fn empty_explicit_reasoning_suppresses_embedded_tag_parsing() {
+    let request = ChatRequest {
+        messages: vec![ChatMessage::assistant_blocks(vec![
+            AssistantContentBlock::Reasoning {
+                text: String::new(),
+            },
+            AssistantContentBlock::Text {
+                text: "<think>literal</think>answer".into(),
+            },
+        ])],
+        ..ChatRequest::for_test()
+    };
+    let rendered = Glm5NextChatRenderer::default().render(&request).unwrap();
+    assert_eq!(
+        rendered.prompt.into_text().unwrap(),
+        "[gMASK]<sop><|system|>Reasoning Effort: Max<|assistant|><think></think><think>literal</think>answer<|assistant|><think>"
+    );
+}
+
+#[test]
+fn continue_final_message_matches_existing_hf_renderer() {
+    let final_messages = [
+        ChatMessage::assistant_text("  partial answer \n"),
+        ChatMessage::assistant_text("<think>inline</think>partial"),
+        ChatMessage::assistant_blocks(vec![
+            AssistantContentBlock::Reasoning {
+                text: "reasoning".into(),
+            },
+            AssistantContentBlock::Text {
+                text: "partial".into(),
+            },
+            AssistantContentBlock::ToolCall(AssistantToolCall {
+                id: "call_0".into(),
+                name: "lookup".into(),
+                arguments: json!({"q": "x"}).to_string(),
+            }),
+        ]),
+    ];
+    let native = Glm5NextChatRenderer::default();
+    for (name, template) in CHECKPOINT_TEMPLATES {
+        let hf = hf_renderer(template);
+        for final_message in &final_messages {
+            let mut request = tool_request(1);
+            request.messages.push(ChatMessage::user("继续"));
+            request.messages.push(final_message.clone());
+            request.chat_options.generation_prompt_mode =
+                GenerationPromptMode::ContinueFinalAssistant;
+            assert_eq!(
+                native.render(&request).unwrap().prompt,
+                hf.render(&request).unwrap().prompt,
+                "{name}: {final_message:?}"
+            );
+        }
+    }
+}
+
+#[test]
+fn media_parts_render_reminders_without_multimodal_inputs() {
+    let request = ChatRequest {
+        messages: vec![ChatMessage::user(vec![
+            crate::request::ChatContentPart::image_url("unused"),
+        ])],
+        ..ChatRequest::for_test()
+    };
+    let rendered = Glm5NextChatRenderer::default().render(&request).unwrap();
+    assert!(rendered.prompt.into_text().unwrap().contains("unable to process this image"));
+    assert_eq!(rendered.media_order, Some(Vec::new()));
+}
+
+/// Exercise BPE merges, added-token recognition and Unicode without a checkpoint download.
+fn tokenizer_fixture() -> Value {
+    let vocab: serde_json::Map<String, Value> = (0u32..256)
+        .map(|byte| {
+            let ch = match byte {
+                33..=126 | 161..=172 | 174..=255 => char::from_u32(byte).unwrap(),
+                _ => {
+                    let index = (0..byte)
+                        .filter(|b| !matches!(b, 33..=126 | 161..=172 | 174..=255))
+                        .count();
+                    char::from_u32(256 + index as u32).unwrap()
+                }
+            };
+            (ch.to_string(), json!(byte))
+        })
+        .chain([("ab".into(), json!(256)), ("abc".into(), json!(257))])
+        .collect();
+    let tokens = [
+        "[gMASK]",
+        "<sop>",
+        "<|system|>",
+        "<|user|>",
+        "<|assistant|>",
+        "<|observation|>",
+        "<think>",
+        "</think>",
+        "<tool_call>",
+        "</tool_call>",
+        "<tool_response>",
+        "</tool_response>",
+        "<arg_key>",
+        "</arg_key>",
+        "<arg_value>",
+        "</arg_value>",
+        "<|image|>",
+        "<|video|>",
+        "<|begin_of_image|>",
+        "<|end_of_image|>",
+        "<|begin_of_video|>",
+        "<|end_of_video|>",
+        "<|begin_of_audio|>",
+        "<|end_of_audio|>",
+    ];
+    json!({
+        "version": "1.0", "truncation": null, "padding": null, "normalizer": null,
+        "added_tokens": tokens.iter().enumerate().map(|(i, content)| json!({
+            "id": 258 + i, "content": content, "single_word": false,
+            "lstrip": false, "rstrip": false, "normalized": false, "special": i < 6,
+        })).collect::<Vec<_>>(),
+        "pre_tokenizer": {"type": "Sequence", "pretokenizers": [
+            {"type": "Split", "pattern": {"Regex": " ?\\p{L}+| ?[^\\s\\p{L}]+|\\s+"},
+             "behavior": "Isolated", "invert": false},
+            {"type": "ByteLevel", "add_prefix_space": false,
+             "trim_offsets": true, "use_regex": false}
+        ]},
+        "post_processor": {"type": "ByteLevel", "add_prefix_space": true,
+                           "trim_offsets": false, "use_regex": true},
+        "decoder": {"type": "ByteLevel", "add_prefix_space": false,
+                    "trim_offsets": false, "use_regex": false},
+        "model": {"type": "BPE", "dropout": null, "unk_token": null,
+                  "continuing_subword_prefix": null, "end_of_word_suffix": null,
+                  "fuse_unk": false, "byte_fallback": false, "ignore_merges": false,
+                  "vocab": vocab, "merges": [["a", "b"], ["ab", "c"]]}
+    })
+}
+
+/// Load `VLLM_GLM_TOKENIZER_JSON`, or the fixture, with both tokenizer backends.
+fn tokenizers(dir: &std::path::Path) -> [DynTokenizer; 2] {
+    let path = std::env::var_os("VLLM_GLM_TOKENIZER_JSON")
+        .map(std::path::PathBuf::from)
+        .unwrap_or_else(|| {
+            let path = dir.join("tokenizer.json");
+            std::fs::write(&path, serde_json::to_vec(&tokenizer_fixture()).unwrap()).unwrap();
+            path
+        });
+    [
+        HuggingFaceTokenizer::new_hf,
+        HuggingFaceTokenizer::new_fastokens,
+    ]
+    .map(|constructor| Arc::new(constructor(&path).unwrap()) as DynTokenizer)
+}
+
+#[test]
+fn token_ids_match_whole_prompt_encoding() {
+    let dir = tempfile::tempdir().unwrap();
+    for tokenizer in tokenizers(dir.path()) {
+        let mut request = tool_request(2);
+        request.messages.push(ChatMessage::user(vec![
+            crate::request::ChatContentPart::text("abc 中文 e\u{301}🙂\r\n"),
+            crate::request::ChatContentPart::image_url("unused"),
+        ]));
+        for mode in [
+            GenerationPromptMode::StartNewAssistant,
+            GenerationPromptMode::NoGenerationPrompt,
+            GenerationPromptMode::ContinueFinalAssistant,
+        ] {
+            request.chat_options.generation_prompt_mode = mode;
+            if mode == GenerationPromptMode::ContinueFinalAssistant {
+                request.messages.push(ChatMessage::assistant_text("abc  "));
+            }
+            for multimodal in [false, true] {
+                let renderer = Glm5NextChatRenderer::default().with_multimodal(multimodal);
+                let text = renderer.render(&request).unwrap();
+                let tokens =
+                    renderer.with_tokenizer(tokenizer.clone()).unwrap().render(&request).unwrap();
+                let ids = tokenizer.encode(&text.prompt.into_text().unwrap(), false).unwrap();
+                assert_eq!(tokens.prompt, Prompt::TokenIds(ids));
+                assert_eq!(tokens.media_order, text.media_order);
+            }
+        }
+    }
+}
+
+#[test]
+fn literal_control_token_spellings_use_the_base_vocabulary() {
+    let dir = tempfile::tempdir().unwrap();
+    for tokenizer in tokenizers(dir.path()) {
+        let request = ChatRequest {
+            messages: vec![ChatMessage::user("<|user|><think>x</think>")],
+            ..ChatRequest::for_test()
+        };
+        let text = Glm5NextChatRenderer::default().render(&request).unwrap().prompt;
+        let Prompt::TokenIds(ids) = Glm5NextChatRenderer::default()
+            .with_tokenizer(tokenizer.clone())
+            .unwrap()
+            .render(&request)
+            .unwrap()
+            .prompt
+        else {
+            panic!("expected token IDs");
+        };
+        assert_eq!(
+            tokenizer.decode(&ids, false).unwrap(),
+            text.into_text().unwrap()
+        );
+        let user = tokenizer.token_to_id("<|user|>").unwrap();
+        assert_eq!(ids.iter().filter(|&&id| id == user).count(), 1);
+    }
+}
+
+#[test]
+fn fixture_token_ids_decode_to_rendered_prompts() {
+    let dir = tempfile::tempdir().unwrap();
+    let fixtures: TestCases =
+        serde_json::from_str(include_str!("fixtures/test_cases.json")).unwrap();
+    for tokenizer in tokenizers(dir.path()) {
+        let control = ControlTokenizer::new(tokenizer.clone()).unwrap();
+        for case in fixtures.cases.iter().filter(|case| case.expected.is_some()) {
+            let mode = if case.continue_final_message {
+                GenerationPromptMode::ContinueFinalAssistant
+            } else if case.add_generation_prompt {
+                GenerationPromptMode::StartNewAssistant
+            } else {
+                GenerationPromptMode::NoGenerationPrompt
+            };
+            let mut out = Glm5PromptWriter::new(Some(&control));
+            encoding::render_into(
+                &mut out,
+                &case.messages,
+                &case.tools,
+                mode,
+                &case.kwargs,
+                false,
+            )
+            .unwrap();
+            let Prompt::TokenIds(ids) = out.finish().unwrap() else {
+                panic!("expected token IDs");
+            };
+            assert_eq!(
+                &tokenizer.decode(&ids, false).unwrap(),
+                case.expected.as_ref().unwrap(),
+                "{}",
+                case.name
+            );
+        }
+    }
+}
+
+#[test]
+fn tokenizer_without_control_tokens_is_rejected() {
+    let dir = tempfile::tempdir().unwrap();
+    let mut config = tokenizer_fixture();
+    config["added_tokens"].as_array_mut().unwrap().pop();
+    let path = dir.path().join("tokenizer.json");
+    std::fs::write(&path, serde_json::to_vec(&config).unwrap()).unwrap();
+    let tokenizer: DynTokenizer = Arc::new(HuggingFaceTokenizer::new_hf(&path).unwrap());
+    assert!(Glm5NextChatRenderer::default().with_tokenizer(tokenizer).is_err());
+}
diff --git a/rust/src/chat/src/renderer/glm5_next/writer.rs b/rust/src/chat/src/renderer/glm5_next/writer.rs
new file mode 100644
index 000000000000..e81d2e67d91b
--- /dev/null
+++ b/rust/src/chat/src/renderer/glm5_next/writer.rs
@@ -0,0 +1,129 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+use vllm_text::Prompt;
+use vllm_text::tokenizer::DynTokenizer;
+
+use crate::error::{Result, chat_template_thrown};
+
+/// Control tokens emitted by the GLM5Next template.
+const CONTROL_TOKENS: [&str; 24] = [
+    "[gMASK]",
+    "<sop>",
+    "<|system|>",
+    "<|user|>",
+    "<|assistant|>",
+    "<|observation|>",
+    "<think>",
+    "</think>",
+    "<tool_call>",
+    "</tool_call>",
+    "<arg_key>",
+    "</arg_key>",
+    "<arg_value>",
+    "</arg_value>",
+    "<tool_response>",
+    "</tool_response>",
+    "<|begin_of_image|>",
+    "<|image|>",
+    "<|end_of_image|>",
+    "<|begin_of_video|>",
+    "<|video|>",
+    "<|end_of_video|>",
+    "<|begin_of_audio|>",
+    "<|end_of_audio|>",
+];
+
+/// Tokenizer with the added-vocabulary IDs of [`CONTROL_TOKENS`].
+pub(super) struct ControlTokenizer {
+    tokenizer: DynTokenizer,
+    ids: [u32; CONTROL_TOKENS.len()],
+}
+
+impl ControlTokenizer {
+    pub(super) fn new(tokenizer: DynTokenizer) -> Result<Self> {
+        let mut ids = [0; CONTROL_TOKENS.len()];
+        for (id, token) in ids.iter_mut().zip(CONTROL_TOKENS) {
+            *id = tokenizer
+                .added_vocab()
+                .iter()
+                .find_map(|(text, id)| (text == token).then_some(*id))
+                .ok_or_else(|| {
+                    chat_template_thrown!("GLM tokenizer lacks added token {token:?}")
+                })?;
+        }
+        Ok(Self { tokenizer, ids })
+    }
+}
+
+/// GLM prompt writer: control tokens map to added-vocabulary IDs and the text
+/// between them is encoded once with the base vocabulary.
+pub(super) struct Glm5PromptWriter<'a> {
+    text: String,
+    tokens: Option<(&'a ControlTokenizer, Vec<u32>)>,
+    pending: usize,
+}
+
+impl<'a> Glm5PromptWriter<'a> {
+    pub(super) fn new(tokenizer: Option<&'a ControlTokenizer>) -> Self {
+        Self {
+            text: String::new(),
+            tokens: tokenizer.map(|tokenizer| (tokenizer, Vec::new())),
+            pending: 0,
+        }
+    }
+
+    pub(super) fn text(&self) -> &str {
+        &self.text
+    }
+
+    /// Write template text, recognizing the GLM control tokens it contains.
+    pub(super) fn control(&mut self, mut text: &str) -> Result<()> {
+        if self.tokens.is_none() {
+            self.text.push_str(text);
+            return Ok(());
+        }
+        while let Some(start) = text.find(['<', '[']) {
+            let rest = &text[start..];
+            let Some(index) = CONTROL_TOKENS.iter().position(|token| rest.starts_with(token))
+            else {
+                self.text.push_str(&text[..=start]);
+                text = &text[start + 1..];
+                continue;
+            };
+            self.text.push_str(&text[..start]);
+            self.flush()?;
+            let (tokenizer, token_ids) = self.tokens.as_mut().unwrap();
+            token_ids.push(tokenizer.ids[index]);
+            self.text.push_str(CONTROL_TOKENS[index]);
+            self.pending = self.text.len();
+            text = &rest[CONTROL_TOKENS[index].len()..];
+        }
+        self.text.push_str(text);
+        Ok(())
+    }
+
+    /// Write literal text; control-token spellings in it stay plain text.
+    pub(super) fn ordinary(&mut self, text: &str) {
+        self.text.push_str(text);
+    }
+
+    fn flush(&mut self) -> Result<()> {
+        if let Some((tokenizer, token_ids)) = &mut self.tokens
+            && self.pending < self.text.len()
+        {
+            let text = &self.text[self.pending..];
+            token_ids.extend(tokenizer.tokenizer.encode_ordinary(text)?);
+            self.pending = self.text.len();
+        }
+        Ok(())
+    }
+
+    pub(super) fn finish(mut self) -> Result<Prompt> {
+        self.flush()?;
+        Ok(match self.tokens {
+            Some((_, token_ids)) => Prompt::TokenIds(token_ids),
+            None => Prompt::Text(self.text),
+        })
+    }
+}
diff --git a/rust/src/chat/src/renderer/hf/mod.rs b/rust/src/chat/src/renderer/hf/mod.rs
index dfa6a1d2165a..083f7d30f639 100644
--- a/rust/src/chat/src/renderer/hf/mod.rs
+++ b/rust/src/chat/src/renderer/hf/mod.rs
@@ -9,7 +9,7 @@ use thiserror_ext::AsReport as _;
 use tracing::{info, trace, warn};
 use vllm_text::Prompt;
 use vllm_text::backend::hf::{
-    HfSpecialTokens, HfTokenizerConfig, ResolvedModelFiles, load_tokenizer_config,
+    HfChatTemplate, HfSpecialTokens, HfTokenizerConfig, ResolvedModelFiles, load_tokenizer_config,
 };
 
 use self::format::{
@@ -73,6 +73,9 @@ pub struct MultimodalRenderInfo {
 /// state.
 pub struct HfChatRenderer {
     default_template: Option<CompiledChatTemplate>,
+    /// The model's template for requests with tools, from a list-form
+    /// `chat_template` entry named `tool_use`; falls back to the default.
+    tool_use_template: Option<CompiledChatTemplate>,
     default_template_kwargs: HashMap<String, JsonValue>,
     content_format: ContentFormatOption,
     special_tokens: Option<HfSpecialTokens>,
@@ -97,6 +100,7 @@ impl HfChatRenderer {
                         .map_err(|error| Error::ChatTemplate(error.to_report_string()))
                 })
                 .transpose()?,
+            tool_use_template: None,
             default_template_kwargs,
             content_format,
             special_tokens: None,
@@ -109,6 +113,18 @@ impl HfChatRenderer {
         })
     }
 
+    /// Use a separate model template for requests with tools, as Transformers
+    /// selects a `tool_use` template when tools are passed.
+    pub fn with_tool_use_template(mut self, template: Option<String>) -> Result<Self> {
+        self.tool_use_template = template
+            .map(|template| {
+                CompiledChatTemplate::new(template, self.content_format)
+                    .map_err(|error| Error::ChatTemplate(error.to_report_string()))
+            })
+            .transpose()?;
+        Ok(self)
+    }
+
     pub fn with_special_tokens(mut self, special_tokens: Option<HfSpecialTokens>) -> Self {
         self.special_tokens = special_tokens;
         self
@@ -137,7 +153,16 @@ impl HfChatRenderer {
             chat_template,
             ..
         } = load_tokenizer_config(files.tokenizer_config_path.as_deref())?;
-        let mut template = chat_template;
+        let mut template = chat_template
+            .as_ref()
+            .and_then(HfChatTemplate::default_template)
+            .map(str::to_owned);
+        // A configured template or a dedicated template file replaces every
+        // tokenizer_config template, `tool_use` included.
+        let mut tool_use_template = chat_template
+            .as_ref()
+            .and_then(HfChatTemplate::tool_use_template)
+            .map(str::to_owned);
         let special_tokens = (!special_tokens.is_empty()).then_some(special_tokens);
 
         if let Some(configured_template) = options.chat_template.as_deref() {
@@ -145,6 +170,7 @@ impl HfChatRenderer {
                 resolve_chat_template(configured_template)
                     .map_err(|error| Error::ChatTemplate(error.to_report_string()))?,
             );
+            tool_use_template = None;
             info!("using configured chat template override");
         } else if let Some(chat_template_path) = files.chat_template_path.as_deref() {
             // If independent chat template file(s) exist and contain non-empty content,
@@ -158,6 +184,7 @@ impl HfChatRenderer {
                     "loaded dedicated chat template file, overriding tokenizer_config chat_template"
                 );
                 template = file_template;
+                tool_use_template = None;
             } else {
                 warn!(
                     path = %chat_template_path.display(),
@@ -171,6 +198,7 @@ impl HfChatRenderer {
             options.default_chat_template_kwargs,
             options.chat_template_content_format,
         )?
+        .with_tool_use_template(tool_use_template)?
         .with_special_tokens(special_tokens)
         .with_multimodal(multimodal)
         .with_trust_request_chat_template(options.trust_request_chat_template))
@@ -183,7 +211,9 @@ impl HfChatRenderer {
     /// temporary template is compiled from that string and used instead of
     /// the model's default. Overrides are refused unless request templates
     /// are trusted, and render under a fuel budget; the model's default
-    /// template has none, as in the Python frontend.
+    /// template has none, as in the Python frontend. Without an override, a
+    /// request with tools uses the model's `tool_use` template if it has one,
+    /// as Transformers' `get_chat_template` does.
     fn apply_chat_template(&self, request: &ChatRequest) -> Result<RenderedPrompt> {
         if request.chat_options.chat_template.is_some() && !self.trust_request_chat_template {
             return Err(Error::UntrustedRequestChatTemplate);
@@ -198,9 +228,15 @@ impl HfChatRenderer {
                     .map_err(|error| Error::ChatTemplate(error.to_report_string()))
             })
             .transpose()?;
+        let model_template = if request.tools().is_empty() {
+            None
+        } else {
+            self.tool_use_template.as_ref()
+        }
+        .or(self.default_template.as_ref());
         let template = override_template
             .as_ref()
-            .or(self.default_template.as_ref())
+            .or(model_template)
             .ok_or(Error::MissingChatTemplate)?;
 
         self.apply_chat_template_inner(template, request)
@@ -215,7 +251,28 @@ impl HfChatRenderer {
             &request.messages,
             effective_template.content_format(),
             self.multimodal.as_ref(),
+            effective_template.supports_developer_role(),
         )?;
+        let has_developer_messages = request
+            .messages
+            .iter()
+            .any(|message| matches!(message, ChatMessage::Developer { .. }));
+        if has_developer_messages && !effective_template.supports_developer_role() {
+            // Consolidation moves every non-leading system message to the front,
+            // ahead of earlier media. Media data stays in message order, so the
+            // rendered placeholders would no longer match it.
+            let moves_media = request.messages.iter().skip(1).any(|message| {
+                matches!(
+                    message,
+                    ChatMessage::System { content } | ChatMessage::Developer { content, .. }
+                        if content.has_multimodal()
+                )
+            });
+            if moves_media {
+                return Err(Error::UnsupportedMultimodalSystemFallback);
+            }
+            messages = consolidate_system_messages(messages);
+        }
 
         // Handling of `continue_final_message`:
         // Append a sentinel tag to the final message content, render as usual, then
@@ -280,6 +337,10 @@ impl ChatRenderer for HfChatRenderer {
     fn render(&self, request: &ChatRequest) -> Result<RenderedPrompt> {
         self.apply_chat_template(request)
     }
+
+    fn renders_request_chat_template(&self) -> bool {
+        true
+    }
 }
 
 /// Chat message in the JSON shape expected by Jinja chat templates.
@@ -320,6 +381,91 @@ enum TemplateContentPart {
     Audio,
 }
 
+/// Match Python vLLM's fallback for templates that do not support `developer`.
+fn consolidate_system_messages<'a>(messages: Vec<TemplateMessage<'a>>) -> Vec<TemplateMessage<'a>> {
+    let needs_consolidation = messages
+        .iter()
+        .enumerate()
+        .any(|(index, message)| message.role == "system" && index > 0)
+        || messages.iter().filter(|message| message.role == "system").count() > 1;
+    if !needs_consolidation {
+        return messages;
+    }
+
+    let mut system_contents = Vec::new();
+    let mut non_system = Vec::with_capacity(messages.len());
+    for message in messages {
+        if message.role == "system" {
+            system_contents.push(message.content);
+        } else {
+            non_system.push(message);
+        }
+    }
+
+    let mut consolidated = Vec::with_capacity(non_system.len() + 1);
+    consolidated.push(TemplateMessage {
+        role: "system",
+        content: merge_system_content(system_contents),
+        tools: None,
+        reasoning: None,
+        reasoning_content: None,
+        tool_calls: None,
+        tool_call_id: None,
+    });
+    consolidated.extend(non_system);
+    consolidated
+}
+
+/// Merge system/developer content without changing the format selected by the
+/// chat template. OpenAI content arrays retain all text and multimodal parts.
+fn merge_system_content(contents: Vec<TemplateContent>) -> TemplateContent {
+    let mut contents = contents.into_iter();
+    let Some(mut merged) = contents.next() else {
+        return TemplateContent::String(String::new());
+    };
+
+    for content in contents {
+        merged = match (merged, content) {
+            (TemplateContent::String(mut left), TemplateContent::String(right)) => {
+                if !left.is_empty() && !right.is_empty() {
+                    left.push_str("\n\n");
+                }
+                left.push_str(&right);
+                TemplateContent::String(left)
+            }
+            (TemplateContent::OpenAi(mut left), TemplateContent::OpenAi(mut right)) => {
+                if !left.is_empty() && !right.is_empty() {
+                    left.push(TemplateContentPart::Text {
+                        text: "\n\n".to_string(),
+                    });
+                }
+                left.append(&mut right);
+                TemplateContent::OpenAi(left)
+            }
+            (left, right) => {
+                let mut left = into_openai_parts(left);
+                let mut right = into_openai_parts(right);
+                if !left.is_empty() && !right.is_empty() {
+                    left.push(TemplateContentPart::Text {
+                        text: "\n\n".to_string(),
+                    });
+                }
+                left.append(&mut right);
+                TemplateContent::OpenAi(left)
+            }
+        };
+    }
+
+    merged
+}
+
+fn into_openai_parts(content: TemplateContent) -> Vec<TemplateContentPart> {
+    match content {
+        TemplateContent::String(text) => vec![TemplateContentPart::Text { text }],
+        TemplateContent::OpenAi(parts) => parts,
+    }
+}
+
 #[derive(Debug, Serialize)]
 struct TemplateToolCall {
     id: String,
@@ -356,10 +502,13 @@ fn to_template_messages<'a>(
     messages: &'a [ChatMessage],
     content_format: ChatTemplateContentFormat,
     multimodal: Option<&MultimodalRenderInfo>,
+    supports_developer_role: bool,
 ) -> Result<Vec<TemplateMessage<'a>>> {
     messages
         .iter()
-        .map(|message| to_template_message(message, content_format, multimodal))
+        .map(|message| {
+            to_template_message(message, content_format, multimodal, supports_developer_role)
+        })
         .collect()
 }
 
@@ -367,6 +516,7 @@ fn to_template_message<'a>(
     message: &'a ChatMessage,
     content_format: ChatTemplateContentFormat,
     multimodal: Option<&MultimodalRenderInfo>,
+    supports_developer_role: bool,
 ) -> Result<TemplateMessage<'a>> {
     Ok(match message {
         ChatMessage::System { content } => TemplateMessage {
@@ -379,9 +529,15 @@ fn to_template_message<'a>(
             tool_call_id: None,
         },
         ChatMessage::Developer { content, tools } => TemplateMessage {
-            role: "developer",
+            role: if supports_developer_role {
+                "developer"
+            } else {
+                "system"
+            },
             content: to_template_content(content, content_format, multimodal)?,
-            tools: tools.as_deref().map(to_template_tools),
+            tools: supports_developer_role
+                .then(|| tools.as_deref().map(to_template_tools))
+                .flatten(),
             reasoning: None,
             reasoning_content: None,
             tool_calls: None,
@@ -520,31 +676,29 @@ fn to_template_string_content(
     match content {
         ChatContent::Text(text) => Ok(text.clone()),
         ChatContent::Parts(parts) => {
-            let mut out = String::new();
-            for part in parts {
-                match part {
-                    ChatContentPart::Text { text } => out.push_str(text),
-                    ChatContentPart::ImageUrl { .. } => {
-                        let image_token = multimodal
-                            .and_then(|multimodal| multimodal.image_token.as_ref())
-                            .ok_or(Error::UnsupportedMultimodalContent("image_url"))?;
-                        out.push_str(image_token);
-                    }
-                    ChatContentPart::VideoUrl { .. } => {
-                        let video_token = multimodal
-                            .and_then(|multimodal| multimodal.video_token.as_ref())
-                            .ok_or(Error::UnsupportedMultimodalContent("video_url"))?;
-                        out.push_str(video_token);
-                    }
+            // Same as Python with `--interleave-mm-strings`: each part becomes one segment,
+            // media as its modality's placeholder in place, empty text parts are dropped, and
+            // the segments are joined with "\n". Python's default moves placeholders to the
+            // front instead.
+            let segments: Vec<&str> = parts
+                .iter()
+                .filter(|part| !part.is_empty_text())
+                .map(|part| match part {
+                    ChatContentPart::Text { text } => Ok(text.as_str()),
+                    ChatContentPart::ImageUrl { .. } => multimodal
+                        .and_then(|multimodal| multimodal.image_token.as_deref())
+                        .ok_or(Error::UnsupportedMultimodalContent("image_url")),
+                    ChatContentPart::VideoUrl { .. } => multimodal
+                        .and_then(|multimodal| multimodal.video_token.as_deref())
+                        .ok_or(Error::UnsupportedMultimodalContent("video_url")),
                     ChatContentPart::InputAudio { .. } | ChatContentPart::AudioUrl { .. } => {
-                        let audio_token = multimodal
-                            .and_then(|multimodal| multimodal.audio_token.as_ref())
-                            .ok_or(Error::UnsupportedMultimodalContent("audio"))?;
-                        out.push_str(audio_token);
+                        multimodal
+                            .and_then(|multimodal| multimodal.audio_token.as_deref())
+                            .ok_or(Error::UnsupportedMultimodalContent("audio"))
                     }
-                }
-            }
-            Ok(out)
+                })
+                .collect::<Result<_>>()?;
+            Ok(segments.join("\n"))
         }
     }
 }
@@ -690,7 +844,7 @@ mod tests {
         };
 
         let rendered = render(
-            Some("{{ messages|length }}:{{ messages[0].role }}:{{ messages[0].tools[0].function.name }}:{{ tools[0].function.name }}"),
+            Some("{% set developer_role = 'developer' %}{{ messages|length }}:{{ messages[0].role }}:{{ messages[0].tools[0].function.name }}:{{ tools[0].function.name }}"),
             &request,
         )
         .unwrap();
@@ -779,7 +933,7 @@ mod tests {
         )
         .unwrap();
 
-        assert_eq!(rendered.prompt, Prompt::Text("a<image>b".to_string()));
+        assert_eq!(rendered.prompt, Prompt::Text("a\n<image>\nb".to_string()));
     }
 
     #[test]
@@ -791,7 +945,7 @@ mod tests {
         )
         .unwrap();
 
-        assert_eq!(rendered.prompt, Prompt::Text("a<video>b".to_string()));
+        assert_eq!(rendered.prompt, Prompt::Text("a\n<video>\nb".to_string()));
     }
 
     #[test]
@@ -805,7 +959,7 @@ mod tests {
 
         assert_eq!(
             rendered.prompt,
-            Prompt::Text("a<audio><audio>b".to_string())
+            Prompt::Text("a\n<audio>\n<audio>\nb".to_string())
         );
     }
 
@@ -848,6 +1002,56 @@ mod tests {
         );
     }
 
+    #[test]
+    fn developer_fallback_preserves_openai_content_parts() {
+        let request = sample_request(vec![
+            ChatMessage::system(vec![
+                ChatContentPart::text("base policy"),
+                ChatContentPart::image_url("data:image/png;base64,test"),
+            ]),
+            ChatMessage::developer("additional policy", None),
+            ChatMessage::user("question"),
+        ]);
+        let rendered = render_mm(
+            "{% for message in messages %}{{ message.role }}={% for part in message.content %}{% if part.type == 'text' %}{{ part.text }}{% elif part.type == 'image' %}<image>{% endif %}{% endfor %};{% endfor %}",
+            &request,
+            ChatTemplateContentFormatOption::OpenAi,
+        )
+        .unwrap();
+
+        assert_eq!(
+            rendered.prompt,
+            Prompt::Text(
+                "system=base policy<image>\n\nadditional policy;user=question;".to_string()
+            )
+        );
+    }
+
+    #[test]
+    fn developer_fallback_rejects_media_in_moved_messages() {
+        let request = sample_request(vec![
+            ChatMessage::user(vec![
+                ChatContentPart::text("question"),
+                ChatContentPart::image_url("data:image/png;base64,a"),
+            ]),
+            ChatMessage::developer(
+                vec![
+                    ChatContentPart::text("policy"),
+                    ChatContentPart::image_url("data:image/png;base64,b"),
+                ],
+                None,
+            ),
+        ]);
+        let error = render_mm(
+            "{% for message in messages %}{{ message.role }}={{ message.content }};{% endfor %}",
+            &request,
+            ChatTemplateContentFormatOption::String,
+        )
+        .unwrap_err();
+
+        assert!(matches!(error, Error::UnsupportedMultimodalSystemFallback));
+    }
+
     #[test]
     fn video_parts_are_rejected_when_model_lacks_video_support() {
         let error = HfChatRenderer::new(
@@ -1028,14 +1232,19 @@ mod tests {
 
     #[test]
     fn chat_template_flattens_text_parts_for_string_templates() {
-        let request = sample_request(vec![ChatMessage::user(vec![
-            ChatContentPart::text("hello"),
-            ChatContentPart::text(" world"),
-        ])]);
-
-        let rendered = render(Some("{{ messages[0].content }}"), &request).unwrap();
-
-        assert_eq!(rendered, "hello world");
+        for (parts, expected) in [
+            (vec!["hello", " world"], "hello\n world"),
+            (vec!["", "hello", "", "world", ""], "hello\nworld"),
+            (vec!["hello\n", "world"], "hello\n\nworld"),
+            (vec!["hello"], "hello"),
+            (vec![""], ""),
+        ] {
+            let request = sample_request(vec![ChatMessage::user(
+                parts.into_iter().map(ChatContentPart::text).collect::<Vec<_>>(),
+            )]);
+            let rendered = render(Some("{{ messages[0].content }}"), &request).unwrap();
+            assert_eq!(rendered, expected);
+        }
     }
 
     #[test]
@@ -1056,7 +1265,7 @@ mod tests {
         )]);
 
         let rendered = render(
-            Some("{{ messages[0].role }}|{{ messages[0].content }}|{{ messages[0].tools[0].function.name }}|{{ messages[0].tools[0].function.parameters.required[0] }}"),
+            Some("{% set developer_role = 'developer' %}{{ messages[0].role }}|{{ messages[0].content }}|{{ messages[0].tools[0].function.name }}|{{ messages[0].tools[0].function.parameters.required[0] }}"),
             &request,
         )
         .unwrap();
@@ -1064,6 +1273,26 @@ mod tests {
         assert_eq!(rendered, "developer|policy|get_weather|city");
     }
 
+    #[test]
+    fn chat_template_without_developer_role_consolidates_developer_as_system() {
+        let request = sample_request(vec![
+            ChatMessage::system("base policy"),
+            ChatMessage::user("question"),
+            ChatMessage::developer("additional policy", None),
+        ]);
+
+        let rendered = render(
+            Some("{% for m in messages %}{{ m.role }}={{ m.content }};{% endfor %}"),
+            &request,
+        )
+        .unwrap();
+
+        assert_eq!(
+            rendered,
+            "system=base policy\n\nadditional policy;user=question;"
+        );
+    }
+
     #[test]
     fn chat_template_passes_custom_roles_unchanged() {
         let request = sample_request(vec![
@@ -1243,7 +1472,7 @@ mod tests {
         .unwrap()
         .prompt;
 
-        assert_eq!(rendered, Prompt::Text("hello world".to_string()));
+        assert_eq!(rendered, Prompt::Text("hello\n world".to_string()));
     }
 
     #[test]
@@ -1740,7 +1969,7 @@ mod tests {
 
         expect![[r#"
             Text(
-                "<｜begin▁of▁sentence｜>You are a helpful assistant.<|BOT|>user\nDescribe this image:<im_patch>What is in it?<|EOT|><|BOT|>assistant\n<think>\n",
+                "<｜begin▁of▁sentence｜>You are a helpful assistant.<|BOT|>user\nDescribe this image:\n<im_patch>\nWhat is in it?<|EOT|><|BOT|>assistant\n<think>\n",
             )
         "#]]
         .assert_debug_eq(&rendered.prompt);
diff --git a/rust/src/chat/src/renderer/hf/template.rs b/rust/src/chat/src/renderer/hf/template.rs
index 3ded4130a88b..50923c1784bb 100644
--- a/rust/src/chat/src/renderer/hf/template.rs
+++ b/rust/src/chat/src/renderer/hf/template.rs
@@ -49,6 +49,10 @@ pub(super) fn request_template_fuel(value: Option<&OsStr>) -> Result<Option<u64>
     Ok((fuel > 0).then_some(fuel))
 }
 
+fn supports_developer_role(template: &str) -> bool {
+    template.contains("\"developer\"") || template.contains("'developer'")
+}
+
 /// Build a pre-configured environment with the given template string.
 fn build_environment(template: String) -> Result<Environment<'static>> {
     let mut env = Environment::new();
@@ -164,12 +168,14 @@ pub(super) struct CompiledChatTemplate {
     /// Cached, fully-configured environment for one compiled template.
     env: Environment<'static>,
     content_format: ChatTemplateContentFormat,
+    supports_developer_role: bool,
 }
 
 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 supports_developer_role = supports_developer_role(&template);
         let content_format = match content_format {
             ChatTemplateContentFormatOption::Auto => detect_chat_template_content_format(&template),
             ChatTemplateContentFormatOption::String => ChatTemplateContentFormat::String,
@@ -179,6 +185,7 @@ impl CompiledChatTemplate {
         Ok(Self {
             env,
             content_format,
+            supports_developer_role,
         })
     }
 
@@ -202,6 +209,10 @@ impl CompiledChatTemplate {
     pub fn content_format(&self) -> ChatTemplateContentFormat {
         self.content_format
     }
+
+    pub fn supports_developer_role(&self) -> bool {
+        self.supports_developer_role
+    }
 }
 
 #[cfg(test)]
@@ -340,10 +351,10 @@ mod tests {
     fn render_turns(turns: usize, fuel: Option<u64>) -> Result<String> {
         let template =
             "{% for turn in turns %}{% for _ in range(10) %}{{ turn }}{% endfor %}{% endfor %}";
-        let template = CompiledChatTemplate {
-            env: build_environment(template.to_string())?,
-            content_format: ChatTemplateContentFormat::String,
-        }
+        let template = CompiledChatTemplate::new(
+            template.to_string(),
+            ChatTemplateContentFormatOption::String,
+        )?
         .with_fuel(fuel);
         let mut kwargs = HashMap::new();
         kwargs.insert("turns".to_string(), serde_json::json!(vec![0; turns]));
diff --git a/rust/src/chat/src/renderer/mod.rs b/rust/src/chat/src/renderer/mod.rs
index 347ec43d03e1..1ddfaabc97aa 100644
--- a/rust/src/chat/src/renderer/mod.rs
+++ b/rust/src/chat/src/renderer/mod.rs
@@ -14,6 +14,7 @@ mod deepseek;
 pub mod deepseek_v32;
 pub mod deepseek_v4;
 pub mod deepseek_v41;
+mod glm5_next;
 pub mod harmony;
 pub mod hf;
 mod inkling;
@@ -25,6 +26,7 @@ mod test_utils;
 pub use deepseek_v4::DeepSeekV4ChatRenderer;
 pub use deepseek_v32::DeepSeekV32ChatRenderer;
 pub use deepseek_v41::DeepSeekV41ChatRenderer;
+pub use glm5_next::Glm5NextChatRenderer;
 pub use harmony::HarmonyChatRenderer;
 pub use inkling::InklingChatRenderer;
 pub use kimi_k3::KimiK3ChatRenderer;
@@ -58,6 +60,12 @@ pub trait ChatRenderer: Send + Sync {
     /// Render one chat request into the text prompt submitted to the text
     /// backend.
     fn render(&self, request: &ChatRequest) -> Result<RenderedPrompt>;
+
+    /// Whether [`render`](Self::render) renders a request's `chat_template`
+    /// override. Requests carrying one are rejected for other renderers.
+    fn renders_request_chat_template(&self) -> bool {
+        false
+    }
 }
 
 /// Shared trait-object form of [`ChatRenderer`].
diff --git a/rust/src/chat/src/renderer/selection.rs b/rust/src/chat/src/renderer/selection.rs
index 0ece893762f3..9792d0fb0ca3 100644
--- a/rust/src/chat/src/renderer/selection.rs
+++ b/rust/src/chat/src/renderer/selection.rs
@@ -30,6 +30,8 @@ pub enum RendererSelection {
     Inkling,
     /// Force the Kimi K3 XTML renderer.
     KimiK3,
+    /// Force the native GLM-5.3 / GLM-5.3-Flash renderer.
+    Glm5Next,
 }
 
 impl RendererSelection {
@@ -43,6 +45,7 @@ impl RendererSelection {
     pub const INKLING_LITERAL: &str = "inkling";
     pub const INKLING_MODEL_TYPE: &str = "inkling_mm_model";
     pub const KIMI_K3_LITERAL: &str = "kimi_k3";
+    pub const GLM5_NEXT_LITERAL: &str = "glm5_next";
 
     /// Resolve the renderer selection using the given model type string, if
     /// it's `Auto`.
@@ -55,6 +58,7 @@ impl RendererSelection {
                 Self::GPT_OSS_MODEL_TYPE => Self::Harmony,
                 Self::INKLING_MODEL_TYPE => Self::Inkling,
                 Self::KIMI_K3_LITERAL => Self::KimiK3,
+                Self::GLM5_NEXT_LITERAL => Self::Glm5Next,
                 _ => Self::Hf,
             },
             selection => selection,
@@ -82,6 +86,8 @@ impl FromStr for RendererSelection {
             Ok(Self::Inkling)
         } else if value.eq_ignore_ascii_case(Self::KIMI_K3_LITERAL) {
             Ok(Self::KimiK3)
+        } else if value.eq_ignore_ascii_case(Self::GLM5_NEXT_LITERAL) {
+            Ok(Self::Glm5Next)
         } else {
             Err(format!(
                 "unknown renderer `{value}` (expected one of: {})",
@@ -102,6 +108,7 @@ impl fmt::Display for RendererSelection {
             Self::Harmony => f.write_str(Self::HARMONY_LITERAL),
             Self::Inkling => f.write_str(Self::INKLING_LITERAL),
             Self::KimiK3 => f.write_str(Self::KIMI_K3_LITERAL),
+            Self::Glm5Next => f.write_str(Self::GLM5_NEXT_LITERAL),
         }
     }
 }
@@ -124,11 +131,25 @@ mod tests {
         }
     }
 
+    #[test]
+    fn auto_resolves_glm5_next_from_model_type_only() {
+        // GLM-5.3 shares `glm_moe_dsa` with GLM-5 through GLM-5.2, whose
+        // templates differ, so it needs an explicit `glm5_next`.
+        assert_eq!(
+            RendererSelection::Auto.resolve("glm5_next"),
+            RendererSelection::Glm5Next
+        );
+        assert_eq!(
+            RendererSelection::Auto.resolve("glm_moe_dsa"),
+            RendererSelection::Hf
+        );
+    }
+
     #[test]
     fn renderer_selection_expected_error_message() {
         let err = RendererSelection::from_str("unknown").unwrap_err();
         expect_test::expect![
-            "unknown renderer `unknown` (expected one of: auto, hf, deepseek_v32, deepseek_v4, deepseek_v41, harmony, inkling, kimi_k3)"
+            "unknown renderer `unknown` (expected one of: auto, hf, deepseek_v32, deepseek_v4, deepseek_v41, harmony, inkling, kimi_k3, glm5_next)"
         ]
         .assert_eq(&err);
     }
diff --git a/rust/src/chat/src/request.rs b/rust/src/chat/src/request.rs
index 2c54bdd93ec3..14131db90d3b 100644
--- a/rust/src/chat/src/request.rs
+++ b/rust/src/chat/src/request.rs
@@ -609,6 +609,11 @@ pub struct ChatRequest {
     /// LoRA adapter selected for this request.
     #[serde(default)]
     pub lora_request: Option<LoraRequest>,
+    /// Whether to report the rendered prompt text through
+    /// [`ChatEventStream::prompt_text`](crate::ChatEventStream::prompt_text)
+    /// and [`CollectedAssistantMessage::prompt_text`](crate::CollectedAssistantMessage::prompt_text).
+    #[serde(default)]
+    pub return_prompt_text: bool,
 }
 
 impl ChatRequest {
@@ -630,6 +635,7 @@ impl ChatRequest {
             data_parallel_rank: None,
             session_id: None,
             lora_request: None,
+            return_prompt_text: false,
         }
     }
 
diff --git a/rust/src/chat/src/stream.rs b/rust/src/chat/src/stream.rs
index 19af129a1c71..042c99aedda5 100644
--- a/rust/src/chat/src/stream.rs
+++ b/rust/src/chat/src/stream.rs
@@ -8,7 +8,7 @@ use std::task::{Context, Poll};
 use futures::Stream;
 use trait_set::trait_set;
 use vllm_engine_core_client::protocol::output::RequestSpecDecodeMetrics;
-use vllm_llm::RequestTimestamps;
+use vllm_llm::{RequestTimestamps, WeightVersionSpan};
 use vllm_text::{DecodedLogprobs, DecodedPositionLogprobs, DecodedPromptLogprobs};
 
 use crate::FinishReason;
@@ -20,6 +20,8 @@ use crate::event::{AssistantContentBlock, AssistantMessage, ChatEvent, ChatToken
 pub struct CollectedAssistantMessage {
     pub message: AssistantMessage,
     pub prompt_token_ids: Arc<[u32]>,
+    /// See [`ChatEventStream::prompt_text`].
+    pub prompt_text: Option<String>,
     pub prompt_logprobs: Option<DecodedPromptLogprobs>,
     pub logprobs: Option<DecodedLogprobs>,
     pub token_ids: Vec<u32>,
@@ -33,11 +35,14 @@ pub struct CollectedAssistantMessage {
     /// Connector-specific encoder cache transfer parameters for disaggregated
     /// serving.
     pub ec_transfer_params: Option<serde_json::Value>,
+    /// Weight-version spans over all generated token positions.
+    pub weight_versions: Option<Vec<WeightVersionSpan>>,
 }
 
 /// Per-request stream of chat events.
 pub struct ChatEventStream {
     request_id: String,
+    prompt_text: Option<String>,
     inner: Pin<Box<dyn Stream<Item = Result<ChatEvent>> + Send>>,
 }
 
@@ -45,15 +50,36 @@ impl ChatEventStream {
     pub(crate) fn new(request_id: String, inner: impl crate::output::ChatEventStream) -> Self {
         Self {
             request_id,
+            prompt_text: None,
             inner: Box::pin(inner),
         }
     }
 
+    /// Attach the rendered prompt text reported by [`Self::prompt_text`].
+    pub(crate) fn with_prompt_text(mut self, prompt_text: Option<String>) -> Self {
+        self.prompt_text = prompt_text;
+        self
+    }
+
     /// Return the request ID associated with this stream.
     pub fn request_id(&self) -> &str {
         &self.request_id
     }
 
+    /// Return the prompt text rendered from the chat template, when requested
+    /// with [`ChatRequest::return_prompt_text`](crate::ChatRequest::return_prompt_text)
+    /// and the renderer produced text rather than token IDs. With multimodal
+    /// content, placeholders are not expanded yet.
+    pub fn prompt_text(&self) -> Option<&str> {
+        self.prompt_text.as_deref()
+    }
+
+    /// Take the rendered prompt text out of the stream, so a streaming
+    /// consumer can move it into its first chunk without a copy.
+    pub fn take_prompt_text(&mut self) -> Option<String> {
+        self.prompt_text.take()
+    }
+
     /// Collect the stream to completion and return the final assembled
     /// assistant message.
     pub async fn collect_message(mut self) -> Result<CollectedAssistantMessage> {
@@ -62,6 +88,7 @@ impl ChatEventStream {
         let mut message = AssistantMessage::default();
         let mut prompt_logprobs = None;
         let mut prompt_token_ids: Arc<[u32]> = Arc::from([]);
+        let prompt_text = self.prompt_text.take();
         let mut logprob_positions: Vec<DecodedPositionLogprobs> = Vec::new();
         let mut token_ids: Vec<u32> = Vec::new();
         while let Some(event) = self.next().await.transpose()? {
@@ -91,10 +118,12 @@ impl ChatEventStream {
                     finish_reason,
                     kv_transfer_params,
                     ec_transfer_params,
+                    weight_versions,
                 } => {
                     return Ok(CollectedAssistantMessage {
                         message: done,
                         prompt_token_ids,
+                        prompt_text,
                         prompt_logprobs,
                         logprobs: (!logprob_positions.is_empty()).then_some(DecodedLogprobs {
                             positions: logprob_positions,
@@ -106,6 +135,7 @@ impl ChatEventStream {
                         finish_reason,
                         kv_transfer_params: kv_transfer_params.map(|value| *value),
                         ec_transfer_params: ec_transfer_params.map(|value| *value),
+                        weight_versions,
                     });
                 }
                 ChatEvent::ToolCallEnd { call, .. } => {
@@ -220,16 +250,19 @@ mod tests {
                     finish_reason: FinishReason::stop_eos(),
                     kv_transfer_params: None,
                     ec_transfer_params: None,
+                    weight_versions: None,
                 }),
             ]),
         );
 
+        let stream = stream.with_prompt_text(Some("user: hi\nassistant:".to_string()));
         let collected = stream.collect_message().await.unwrap();
         assert_eq!(
             collected,
             CollectedAssistantMessage {
                 message: Default::default(),
                 prompt_token_ids: vec![10, 11].into(),
+                prompt_text: Some("user: hi\nassistant:".to_string()),
                 prompt_logprobs: Some(DecodedPromptLogprobs {
                     first_token_id: 0,
                     first_token: "o".to_string(),
@@ -268,6 +301,7 @@ mod tests {
                 finish_reason: FinishReason::stop_eos(),
                 kv_transfer_params: None,
                 ec_transfer_params: None,
+                weight_versions: None,
             }
         );
     }
diff --git a/rust/src/chat/tests/chat.rs b/rust/src/chat/tests/chat.rs
index b4fc80b399d8..6eb9c034f4f1 100644
--- a/rust/src/chat/tests/chat.rs
+++ b/rust/src/chat/tests/chat.rs
@@ -1346,8 +1346,13 @@ async fn chat_stream_and_collect_preserve_prompt_and_sample_logprobs() {
     let mut request = sample_request("chat-logprobs");
     request.sampling_params.logprobs = Some(1);
     request.sampling_params.prompt_logprobs = Some(1);
+    request.return_prompt_text = true;
 
     let mut stream = chat.chat(request.clone()).await.unwrap();
+    assert_eq!(
+        stream.prompt_text(),
+        Some("system: You are terse.\nuser: Say hi\nassistant:")
+    );
     match next_semantic(&mut stream).await.unwrap().unwrap() {
         ChatEvent::Start {
             prompt_token_ids,
@@ -1432,6 +1437,10 @@ async fn chat_stream_and_collect_preserve_prompt_and_sample_logprobs() {
     request.request_id = "chat-logprobs-collect".to_string();
     let collected = chat.chat(request).await.unwrap().collect_message().await.unwrap();
     assert_eq!(collected.message.text(), "Hi");
+    assert_eq!(
+        collected.prompt_text.as_deref(),
+        Some("system: You are terse.\nuser: Say hi\nassistant:")
+    );
     assert_eq!(
         collected.prompt_logprobs,
         Some(DecodedPromptLogprobs {
diff --git a/rust/src/chat/tests/roundtrip.rs b/rust/src/chat/tests/roundtrip.rs
index 0d98bc26dbe3..04b088816f78 100644
--- a/rust/src/chat/tests/roundtrip.rs
+++ b/rust/src/chat/tests/roundtrip.rs
@@ -1127,6 +1127,7 @@ fn decoded_completion_stream(
                     ec_transfer_params: None,
                     sampling_mask: None,
                     spec_decode_metrics: None,
+                    weight_versions: None,
                     timestamps: Default::default(),
                 })),
             }
@@ -1141,6 +1142,7 @@ fn decoded_completion_stream(
                 ec_transfer_params: None,
                 sampling_mask: None,
                 spec_decode_metrics: None,
+                weight_versions: None,
                 timestamps: Default::default(),
             });
             events.push(DecodedTextEvent::TextDelta {
diff --git a/rust/src/cmd/src/cli.rs b/rust/src/cmd/src/cli.rs
index ae79c3836da1..b2fed917745f 100644
--- a/rust/src/cmd/src/cli.rs
+++ b/rust/src/cmd/src/cli.rs
@@ -33,11 +33,12 @@ use vllm_engine_core_client::TransportMode;
 use vllm_engine_core_client::protocol::structured_outputs::StructuredOutputBackend;
 use vllm_managed_engine::ManagedEngineConfig;
 use vllm_managed_engine::cli::{ManagedEngineArgs, repartition_managed_engine_args};
+use vllm_managed_engine::data_parallel::{DataParallelLbArgs, DataParallelTopology};
 use vllm_server::{
-    ApiServerOptions, ChatTemplateContentFormatOption, Config, CoordinatorMode, CorsConfig,
-    CustomHistogramBuckets, DEFAULT_KEEP_ALIVE_TIMEOUT, DEFAULT_LOG_STATS_INTERVAL,
-    DEFAULT_REQUEST_BODY_LIMIT_BYTES, HttpListenerMode, LoraModulePath, OutputLogging,
-    ParserSelection, ReasoningEffort, RenderConfig, RendererSelection,
+    AccessLogConfig, ApiServerOptions, ChatTemplateContentFormatOption, Config, CoordinatorMode,
+    CorsConfig, CustomHistogramBuckets, DEFAULT_KEEP_ALIVE_TIMEOUT, DEFAULT_LOG_STATS_INTERVAL,
+    DEFAULT_REQUEST_BODY_LIMIT_BYTES, HttpListenerMode, LoraModulePath, OutputLogConfig,
+    OutputLogging, ParserSelection, ReasoningEffort, RenderConfig, RendererSelection,
 };
 use vllm_text::backend::hf::{HfOverrides, ModelFileSources};
 
@@ -131,6 +132,17 @@ pub struct RenderArgs {
     /// Public model names accepted by the API.
     #[arg(long, num_args = 0..)]
     served_model_name: Vec<String>,
+    /// Accept any requested model name. Like Python, it is set only through
+    /// `VLLM_SKIP_MODEL_NAME_VALIDATION`.
+    #[arg(
+        long,
+        hide = true,
+        env = "VLLM_SKIP_MODEL_NAME_VALIDATION",
+        action = ArgAction::Set,
+        value_parser = parse_truthy_flag,
+        default_value = "0"
+    )]
+    skip_model_name_validation: bool,
     /// Select the tool-call parser. Use `auto` to infer from the model or
     /// `none` to disable parsing.
     #[arg(long, default_value_t)]
@@ -148,6 +160,7 @@ pub struct RenderArgs {
     #[arg(long = "tokenizer-mode", default_value_t)]
     renderer: RendererSelection,
     /// Override the model chat template with a file path or inline template.
+    /// Only the `hf` renderer supports it.
     #[arg(long)]
     chat_template: Option<String>,
     /// Default JSON keyword arguments merged into every chat-template render.
@@ -200,6 +213,7 @@ impl RenderArgs {
             hf_config,
             hf_overrides: self.hf_overrides,
             served_model_name: self.served_model_name,
+            skip_model_name_validation: self.skip_model_name_validation,
             host: self.host,
             port: self.port,
             tool_call_parser: self.tool_call_parser,
@@ -319,7 +333,7 @@ pub struct SharedRuntimeArgs {
     pub http_timeout_keep_alive: Option<u64>,
 
     /// The file path to the chat template, or the template in single-line form
-    /// for the specified model.
+    /// for the specified model. Only the `hf` renderer supports it.
     #[arg(long)]
     #[serde(default)]
     pub chat_template: Option<String>,
@@ -425,6 +439,14 @@ pub struct SharedRuntimeArgs {
     #[serde(default = "default_enable_log_deltas")]
     pub enable_log_deltas: bool,
 
+    /// Max number of prompt characters or prompt ID numbers being printed in
+    /// log. The default of None means unlimited.
+    ///
+    /// Also bounds logged output characters and output token IDs.
+    #[arg(long)]
+    #[serde(default)]
+    pub max_log_len: Option<usize>,
+
     /// Include prompt_tokens_details in usage, including cached and newly
     /// cached prompt tokens.
     #[arg(
@@ -539,6 +561,20 @@ pub struct SharedRuntimeArgs {
     #[serde(skip)]
     pub http_max_request_body_size: Option<usize>,
 
+    /// Accept any requested model name, served by the base model unless it
+    /// names a loaded LoRA adapter. Like Python, it is set only through
+    /// `VLLM_SKIP_MODEL_NAME_VALIDATION`.
+    #[arg(
+        long,
+        hide = true,
+        env = "VLLM_SKIP_MODEL_NAME_VALIDATION",
+        action = ArgAction::Set,
+        value_parser = parse_truthy_flag,
+        default_value = "0"
+    )]
+    #[serde(skip)]
+    pub skip_model_name_validation: bool,
+
     /// The model name(s) used in the API. If multiple names are provided, the
     /// server will respond to any of the provided names. The model name in the
     /// model field of a response will be the first name in this list. If not
@@ -574,11 +610,36 @@ pub struct SharedRuntimeArgs {
     #[serde(default)]
     pub allow_credentials: bool,
 
+    /// Disable the per-response HTTP access log.
+    #[arg(
+        long,
+        visible_alias = "disable-uvicorn-access-log",
+        default_missing_value = "true",
+        num_args = 0..=1
+    )]
+    #[serde(default, alias = "disable_uvicorn_access_log")]
+    pub disable_access_log: bool,
+
+    /// Comma-separated request paths whose responses the HTTP access log
+    /// skips, such as "/health,/metrics" for frequent probes. Paths match
+    /// exactly, without the query string.
+    #[arg(long)]
+    #[serde(default)]
+    pub disable_access_log_for_endpoints: Option<String>,
+
     /// TLS options for HTTPS/mTLS.
     #[command(flatten)]
     #[serde(default, flatten)]
     pub ssl: SslArgs,
 
+    /// Path prefix for a server behind a path-based routing proxy.
+    ///
+    /// Every route is served both under this prefix and at its bare path:
+    /// with `/llm`, `/llm/v1/models` and `/v1/models` reach the same route.
+    #[arg(long)]
+    #[serde(default)]
+    pub root_path: Option<String>,
+
     /// Profiler configuration forwarded by the Python supervisor.
     ///
     /// When set with a non-null `profiler` type, the Rust frontend registers
@@ -695,6 +756,9 @@ impl SharedRuntimeArgs {
                 .map_err(|error| format!("invalid VLLM_HTTP_MAX_REQUEST_BODY_SIZE: {error}"))?;
             self.http_max_request_body_size = Some(size);
         }
+        if let Some(value) = lookup("VLLM_SKIP_MODEL_NAME_VALIDATION") {
+            self.skip_model_name_validation = parse_truthy_flag(&value)?;
+        }
         Ok(())
     }
 
@@ -719,6 +783,7 @@ impl SharedRuntimeArgs {
         let log_stats_interval = self.log_stats_interval();
         let api_server_options = self.api_server_options();
         let cors = self.cors_config();
+        let access_log = self.access_log_config();
         let tls = self.ssl.tls_config();
         let profiler = self.profiler();
         let structured_outputs_backend = self.structured_outputs_backend();
@@ -768,7 +833,9 @@ impl SharedRuntimeArgs {
             stream_interval: self.stream_interval,
             api_server_options,
             cors,
+            access_log,
             tls,
+            root_path: self.root_path.filter(|root_path| !root_path.is_empty()),
             api_keys: self.api_key,
             disable_log_stats: self.disable_log_stats,
             custom_histogram_buckets: self.custom_histogram_buckets,
@@ -790,7 +857,7 @@ impl SharedRuntimeArgs {
         grpc_listener_mode: Option<HttpListenerMode>,
         handshake_address: String,
         advertised_host: String,
-        engine_count: usize,
+        topology: &DataParallelTopology,
         local_input_address: Option<String>,
         local_output_address: Option<String>,
         manages_engine: bool,
@@ -801,6 +868,7 @@ impl SharedRuntimeArgs {
         let log_stats_interval = self.log_stats_interval();
         let api_server_options = self.api_server_options();
         let cors = self.cors_config();
+        let access_log = self.access_log_config();
         let tls = self.ssl.tls_config();
         let profiler = self.profiler();
         let structured_outputs_backend = self.structured_outputs_backend();
@@ -814,7 +882,9 @@ impl SharedRuntimeArgs {
             transport_mode: TransportMode::HandshakeOwner {
                 handshake_address,
                 advertised_host,
-                engine_count,
+                engine_start_index: topology.frontend_engine_start_index(),
+                engine_count: topology.frontend_engine_count(),
+                data_parallel_size: topology.data_parallel_size,
                 ready_timeout,
                 local_input_address,
                 local_output_address,
@@ -847,7 +917,9 @@ impl SharedRuntimeArgs {
             stream_interval: self.stream_interval,
             api_server_options,
             cors,
+            access_log,
             tls,
+            root_path: self.root_path.filter(|root_path| !root_path.is_empty()),
             api_keys: self.api_key,
             disable_log_stats: self.disable_log_stats,
             custom_histogram_buckets: self.custom_histogram_buckets,
@@ -863,10 +935,13 @@ impl SharedRuntimeArgs {
     fn api_server_options(&self) -> ApiServerOptions {
         ApiServerOptions {
             enable_log_requests: self.enable_log_requests,
-            log_outputs: match (self.enable_log_outputs, self.enable_log_deltas) {
-                (false, _) => OutputLogging::Disabled,
-                (true, false) => OutputLogging::Complete,
-                (true, true) => OutputLogging::CompleteAndDeltas,
+            log_outputs: OutputLogConfig {
+                outputs: match (self.enable_log_outputs, self.enable_log_deltas) {
+                    (false, _) => OutputLogging::Disabled,
+                    (true, false) => OutputLogging::Complete,
+                    (true, true) => OutputLogging::CompleteAndDeltas,
+                },
+                max_log_len: self.max_log_len,
             },
             enable_prompt_tokens_details: self.enable_prompt_tokens_details,
             default_return_tokens_as_token_ids: self.return_tokens_as_token_ids,
@@ -874,6 +949,7 @@ impl SharedRuntimeArgs {
             enable_request_id_headers: self.enable_request_id_headers,
             enable_scale_out: self.enable_scale_out,
             enable_tokenizer_info_endpoint: self.enable_tokenizer_info_endpoint,
+            skip_model_name_validation: self.skip_model_name_validation,
             sse_keep_alive_interval: (self.sse_keep_alive_interval > 0)
                 .then(|| Duration::from_secs(self.sse_keep_alive_interval)),
             request_body_limit_bytes: match self.http_max_request_body_size {
@@ -884,6 +960,18 @@ impl SharedRuntimeArgs {
         }
     }
 
+    fn access_log_config(&self) -> AccessLogConfig {
+        AccessLogConfig {
+            enabled: !self.disable_access_log,
+            excluded_paths: (self.disable_access_log_for_endpoints.iter())
+                .flat_map(|paths| paths.split(','))
+                .map(str::trim)
+                .filter(|path| !path.is_empty())
+                .map(str::to_owned)
+                .collect(),
+        }
+    }
+
     fn cors_config(&self) -> CorsConfig {
         CorsConfig {
             allow_origins: self.allowed_origins.0.clone(),
@@ -985,6 +1073,12 @@ fn parse_int_flag(value: &str) -> Result<bool, String> {
     Ok(value != 0)
 }
 
+/// Parse a boolean flag environment variable as Python's
+/// `value.strip().lower() in ("1", "true")`; any other value is `false`.
+fn parse_truthy_flag(value: &str) -> Result<bool, String> {
+    Ok(matches!(value.trim().to_lowercase().as_str(), "1" | "true"))
+}
+
 /// Parse a byte count environment variable as a non-negative decimal integer.
 fn parse_byte_count(value: &str) -> Result<usize, String> {
     value.trim().parse().map_err(|_| format!("`{value}` is not a number of bytes"))
@@ -1096,15 +1190,23 @@ pub struct ServeArgs {
     /// Managed Python headless-engine arguments.
     #[command(flatten)]
     pub managed_engine: ManagedEngineArgs,
+
+    /// Data parallel load-balancing arguments.
+    #[command(flatten)]
+    pub data_parallel_lb: DataParallelLbArgs,
 }
 
 impl ServeArgs {
     /// Build the OpenAI-server runtime config used after the managed Python
     /// engine starts.
-    pub fn to_frontend_config(&self, handshake_address: String) -> Config {
+    pub fn to_frontend_config(
+        &self,
+        handshake_address: String,
+        topology: &DataParallelTopology,
+    ) -> Config {
         // Prefer IPC sockets for local engine input/output.
         let (local_input_address, local_output_address) =
-            self.managed_engine.frontend_local_only().then(frontend_ipc_addresses).unzip();
+            topology.frontend_local_only().then(frontend_ipc_addresses).unzip();
         let listener_mode = match &self.uds {
             Some(path) => HttpListenerMode::BindUnix { path: path.clone() },
             None => HttpListenerMode::BindTcp {
@@ -1128,7 +1230,7 @@ impl ServeArgs {
             grpc_listener_mode,
             handshake_address,
             self.managed_engine.handshake_host.clone(),
-            self.managed_engine.data_parallel_size,
+            topology,
             local_input_address,
             local_output_address,
             // `--data-parallel-size-local 0` runs the frontend without a local engine.
@@ -1138,7 +1240,12 @@ impl ServeArgs {
 
     /// Build the managed Python-engine spawn configuration with the given
     /// handshake port.
-    pub fn to_managed_engine_config(&self, handshake_port: u16) -> ManagedEngineConfig {
+    pub fn to_managed_engine_config(
+        &self,
+        handshake_port: u16,
+        topology: &DataParallelTopology,
+        client_handshake_address: Option<String>,
+    ) -> ManagedEngineConfig {
         let reasoning_parser = effective_engine_reasoning_parser(
             &self.runtime.reasoning_parser,
             &self.runtime.model.id,
@@ -1162,6 +1269,8 @@ impl ServeArgs {
             self.runtime.limit_mm_per_prompt_json(),
             hf_overrides,
             self.runtime.structured_outputs_config_json(),
+            topology,
+            client_handshake_address,
         )
     }
 }
@@ -1170,6 +1279,16 @@ fn effective_engine_reasoning_parser(selection: &ParserSelection, model: &str) -
     selection.resolve_reasoning_name(model).map(str::to_owned)
 }
 
+/// Allocate a fresh IPC handshake endpoint for a frontend that serves only its
+/// node's engines (hybrid or external DP LB).
+pub fn frontend_ipc_handshake_address() -> String {
+    let preferred_base_path = std::env::var_os("VLLM_RPC_BASE_PATH")
+        .map(PathBuf::from)
+        .unwrap_or_else(std::env::temp_dir);
+    let handshake = preferred_base_path.join(format!("vllm-rs-h-{}", Uuid::new_v4().simple()));
+    format!("ipc://{}", handshake.to_string_lossy())
+}
+
 /// Allocate fresh IPC endpoints for one managed frontend instance.
 fn frontend_ipc_addresses() -> (String, String) {
     let preferred_base_path = std::env::var_os("VLLM_RPC_BASE_PATH")
diff --git a/rust/src/cmd/src/cli/tests.rs b/rust/src/cmd/src/cli/tests.rs
index a01a485cfdf8..613b595b2475 100644
--- a/rust/src/cmd/src/cli/tests.rs
+++ b/rust/src/cmd/src/cli/tests.rs
@@ -5,16 +5,21 @@ use std::collections::HashMap;
 use std::time::Duration;
 
 use expect_test::expect;
-use vllm_engine_core_client::TransportMode;
 use vllm_engine_core_client::protocol::structured_outputs::StructuredOutputBackend;
+use vllm_engine_core_client::{CoordinatorMode as EngineCoreCoordinatorMode, TransportMode};
 use vllm_server::{
-    Config, GenerationConfigMode, HttpListenerMode, LoraModulePath, OutputLogging, ParserSelection,
-    ReasoningEffort, RendererSelection,
+    Config, GenerationConfigMode, HttpListenerMode, LoraModulePath, OutputLogConfig, OutputLogging,
+    ParserSelection, ReasoningEffort, RendererSelection,
 };
 use vllm_text::backend::hf::ModelSource;
 
 use super::{BenchCommand, Cli, Command, LogStatsInterval, SharedRuntimeArgs, parse_int_flag};
 
+/// Resolve the data parallel topology the way `vllm-rs serve` does at startup.
+fn topology(args: &super::ServeArgs) -> vllm_managed_engine::data_parallel::DataParallelTopology {
+    args.data_parallel_lb.resolve(&args.managed_engine, args.headless).unwrap()
+}
+
 #[test]
 fn bench_serve_args_parse_without_managed_engine_repartition() {
     let cli = Cli::try_parse_from([
@@ -55,6 +60,10 @@ fn render_args_build_config_without_tls() {
         "deepseek_v32",
         "--max-logprobs",
         "-1",
+        // Set through `VLLM_SKIP_MODEL_NAME_VALIDATION`, which binds this
+        // hidden flag.
+        "--skip-model-name-validation",
+        "True",
     ])
     .unwrap();
 
@@ -77,6 +86,7 @@ fn render_args_build_config_without_tls() {
     assert_eq!(config.port, 8080);
     assert_eq!(config.max_model_len, Some(32768));
     assert_eq!(config.served_model_name, ["qwen"]);
+    assert!(config.skip_model_name_validation);
     assert_eq!(config.tool_call_parser, ParserSelection::Auto);
     assert_eq!(config.reasoning_parser, ParserSelection::Auto);
     assert_eq!(config.renderer, RendererSelection::DeepSeekV32);
@@ -94,6 +104,23 @@ fn render_args_allow_omitted_max_model_len() {
     let config = args.into_config();
 
     assert_eq!(config.max_model_len, None);
+    assert!(!config.skip_model_name_validation);
+}
+
+#[test]
+fn render_args_select_native_glm5_next() {
+    let cli = Cli::try_parse_from([
+        "vllm-rs",
+        "render",
+        "local-glm5_next",
+        "--tokenizer-mode",
+        "glm5_next",
+    ])
+    .unwrap();
+    let Command::Render(args) = cli.command else {
+        panic!("expected render args");
+    };
+    assert_eq!(args.into_config().renderer, RendererSelection::Glm5Next);
 }
 
 #[test]
@@ -225,6 +252,7 @@ fn serve_args_forward_python_flags_with_separator() {
                         enable_log_requests: false,
                         enable_log_outputs: false,
                         enable_log_deltas: true,
+                        max_log_len: None,
                         enable_prompt_tokens_details: false,
                         enable_per_request_metrics: false,
                         enable_request_id_headers: false,
@@ -238,6 +266,7 @@ fn serve_args_forward_python_flags_with_separator() {
                         log_stats_interval: None,
                         disable_request_id_randomization: false,
                         http_max_request_body_size: None,
+                        skip_model_name_validation: false,
                         served_model_name: [],
                         allowed_origins: JsonStringList(
                             [],
@@ -253,6 +282,8 @@ fn serve_args_forward_python_flags_with_separator() {
                             ],
                         ),
                         allow_credentials: false,
+                        disable_access_log: false,
+                        disable_access_log_for_endpoints: None,
                         ssl: SslArgs {
                             ssl_keyfile: None,
                             ssl_certfile: None,
@@ -260,6 +291,7 @@ fn serve_args_forward_python_flags_with_separator() {
                             ssl_cert_reqs: 0,
                             ssl_ciphers: None,
                         },
+                        root_path: None,
                         profiler_config: None,
                     },
                     managed_engine: ManagedEngineArgs {
@@ -276,6 +308,12 @@ fn serve_args_forward_python_flags_with_separator() {
                             "float16",
                         ],
                     },
+                    data_parallel_lb: DataParallelLbArgs {
+                        data_parallel_rank: None,
+                        data_parallel_start_rank: None,
+                        data_parallel_hybrid_lb: false,
+                        data_parallel_external_lb: false,
+                    },
                 },
             ),
         }
@@ -340,7 +378,7 @@ fn serve_args_forward_shutdown_timeout_to_managed_engine() {
     };
     assert_eq!(args.runtime.shutdown_timeout, 60);
 
-    let config = args.to_managed_engine_config(5555);
+    let config = args.to_managed_engine_config(5555, &topology(&args), None);
     expect![[r#"
         [
             "--reasoning-parser",
@@ -362,7 +400,7 @@ fn serve_args_forward_disable_log_stats_to_managed_engine() {
     };
     assert!(args.runtime.disable_log_stats);
 
-    let config = args.to_managed_engine_config(5555);
+    let config = args.to_managed_engine_config(5555, &topology(&args), None);
     expect![[r#"
         [
             "--reasoning-parser",
@@ -389,7 +427,7 @@ fn serve_args_forward_profiler_config_to_managed_engine() {
     };
     assert_eq!(args.runtime.profiler().as_deref(), Some("torch"));
 
-    let config = args.to_managed_engine_config(5555);
+    let config = args.to_managed_engine_config(5555, &topology(&args), None);
     let profiler_flag_index = config
         .python_args
         .iter()
@@ -415,10 +453,10 @@ fn serve_args_select_structured_outputs_backend_and_forward_config() {
         let Command::Serve(args) = cli.command else {
             panic!("expected serve args");
         };
-        let frontend = args.to_frontend_config("tcp://127.0.0.1:0".into());
+        let frontend = args.to_frontend_config("tcp://127.0.0.1:0".into(), &topology(&args));
         (
             frontend.structured_outputs_backend,
-            args.to_managed_engine_config(5555).python_args,
+            args.to_managed_engine_config(5555, &topology(&args), None).python_args,
         )
     };
 
@@ -516,7 +554,7 @@ fn serve_args_tool_strict_level_requires_xgrammar() {
         let Command::Serve(args) = cli.command else {
             panic!("expected serve args");
         };
-        args.to_frontend_config("tcp://127.0.0.1:0".into()).validate()
+        args.to_frontend_config("tcp://127.0.0.1:0".into(), &topology(&args)).validate()
     };
 
     assert!(validate("function", "xgrammar").is_ok());
@@ -542,7 +580,8 @@ fn serve_args_forward_stream_interval_to_frontend_config() {
     let Command::Serve(args) = cli.command else {
         panic!("expected serve args");
     };
-    let frontend_config = args.to_frontend_config("tcp://127.0.0.1:62100".to_string());
+    let frontend_config =
+        args.to_frontend_config("tcp://127.0.0.1:62100".to_string(), &topology(&args));
     assert_eq!(frontend_config.stream_interval.get(), 4);
 
     let error = Cli::try_parse_from([
@@ -565,20 +604,22 @@ fn serve_args_forward_env_only_settings_to_frontend_config() {
         let Command::Serve(args) = cli.command else {
             panic!("expected serve args");
         };
-        args.to_frontend_config("tcp://127.0.0.1:62100".to_string())
+        args.to_frontend_config("tcp://127.0.0.1:62100".to_string(), &topology(&args))
     };
 
     let config = parse(&[]);
     assert_eq!(config.log_stats_interval, Some(Duration::from_secs(10)));
     assert!(!config.disable_request_id_randomization);
+    assert!(!config.api_server_options.skip_model_name_validation);
     assert_eq!(
         config.api_server_options.request_body_limit_bytes,
         Some(32 * 1024 * 1024)
     );
 
     // Set through `VLLM_LOG_STATS_INTERVAL`,
-    // `VLLM_DISABLE_REQUEST_ID_RANDOMIZATION`, and
-    // `VLLM_HTTP_MAX_REQUEST_BODY_SIZE`, which bind these hidden flags.
+    // `VLLM_DISABLE_REQUEST_ID_RANDOMIZATION`,
+    // `VLLM_HTTP_MAX_REQUEST_BODY_SIZE`, and
+    // `VLLM_SKIP_MODEL_NAME_VALIDATION`, which bind these hidden flags.
     let config = parse(&[
         "--log-stats-interval",
         "2.5",
@@ -586,13 +627,19 @@ fn serve_args_forward_env_only_settings_to_frontend_config() {
         "1",
         "--http-max-request-body-size",
         "1048576",
+        "--skip-model-name-validation",
+        "1",
     ]);
     assert_eq!(config.log_stats_interval, Some(Duration::from_millis(2500)));
     assert!(config.disable_request_id_randomization);
+    assert!(config.api_server_options.skip_model_name_validation);
     assert_eq!(
         config.api_server_options.request_body_limit_bytes,
         Some(1024 * 1024)
     );
+    // Like Python, any value other than `1` or `true` leaves it off.
+    let config = parse(&["--skip-model-name-validation", "yes"]);
+    assert!(!config.api_server_options.skip_model_name_validation);
     let config = parse(&["--http-max-request-body-size", "0"]);
     assert_eq!(config.api_server_options.request_body_limit_bytes, None);
 
@@ -620,6 +667,7 @@ fn args_json_reads_env_only_settings_from_env() {
     let args = apply(&[]).unwrap();
     assert_eq!(args.log_stats_interval(), Some(Duration::from_secs(10)));
     assert!(!args.disable_request_id_randomization);
+    assert!(!args.api_server_options().skip_model_name_validation);
     assert_eq!(
         args.api_server_options().request_body_limit_bytes,
         Some(32 * 1024 * 1024)
@@ -629,10 +677,12 @@ fn args_json_reads_env_only_settings_from_env() {
         ("VLLM_LOG_STATS_INTERVAL", "0.5"),
         ("VLLM_DISABLE_REQUEST_ID_RANDOMIZATION", "1"),
         ("VLLM_HTTP_MAX_REQUEST_BODY_SIZE", "1048576"),
+        ("VLLM_SKIP_MODEL_NAME_VALIDATION", " TRUE "),
     ])
     .unwrap();
     assert_eq!(args.log_stats_interval(), Some(Duration::from_millis(500)));
     assert!(args.disable_request_id_randomization);
+    assert!(args.api_server_options().skip_model_name_validation);
     assert_eq!(
         args.api_server_options().request_body_limit_bytes,
         Some(1024 * 1024)
@@ -711,10 +761,11 @@ fn serve_args_forward_max_logprobs_to_frontend_and_managed_engine() {
     };
     assert_eq!(args.runtime.max_logprobs, Some(-1));
 
-    let frontend_config = args.to_frontend_config("tcp://127.0.0.1:62100".to_string());
+    let frontend_config =
+        args.to_frontend_config("tcp://127.0.0.1:62100".to_string(), &topology(&args));
     assert_eq!(frontend_config.max_logprobs, Some(-1));
 
-    let engine_config = args.to_managed_engine_config(5555);
+    let engine_config = args.to_managed_engine_config(5555, &topology(&args), None);
     expect![[r#"
         [
             "--max-logprobs",
@@ -735,7 +786,7 @@ fn serve_args_resolve_auto_reasoning_parser_for_managed_engine() {
     };
     assert_eq!(args.runtime.reasoning_parser, ParserSelection::Auto);
 
-    let config = args.to_managed_engine_config(5555);
+    let config = args.to_managed_engine_config(5555, &topology(&args), None);
     expect![[r#"
         [
             "--reasoning-parser",
@@ -757,7 +808,7 @@ fn serve_args_resolve_unified_reasoning_parser_for_managed_engine() {
             "kimi_k3",
         ]
     "#]]
-    .assert_debug_eq(&args.to_managed_engine_config(5555).python_args);
+    .assert_debug_eq(&args.to_managed_engine_config(5555, &topology(&args), None).python_args);
 }
 
 #[test]
@@ -775,7 +826,7 @@ fn serve_args_forward_explicit_reasoning_parser_to_managed_engine() {
         panic!("expected serve args");
     };
 
-    let config = args.to_managed_engine_config(5555);
+    let config = args.to_managed_engine_config(5555, &topology(&args), None);
     expect![[r#"
         [
             "--reasoning-parser",
@@ -800,7 +851,7 @@ fn serve_args_do_not_forward_disabled_reasoning_parser_to_managed_engine() {
         panic!("expected serve args");
     };
 
-    let config = args.to_managed_engine_config(5555);
+    let config = args.to_managed_engine_config(5555, &topology(&args), None);
     assert!(config.python_args.is_empty());
 }
 
@@ -820,7 +871,7 @@ fn serve_args_forward_reasoning_parser_even_with_passthrough_reasoning_parser()
         panic!("expected serve args");
     };
 
-    let config = args.to_managed_engine_config(5555);
+    let config = args.to_managed_engine_config(5555, &topology(&args), None);
     expect![[r#"
         [
             "--reasoning-parser",
@@ -878,7 +929,7 @@ fn serve_passes_enable_request_id_headers_into_config() {
     let Command::Serve(args) = cli.command else {
         panic!("expected serve args");
     };
-    let config = args.to_frontend_config("tcp://127.0.0.1:62100".to_string());
+    let config = args.to_frontend_config("tcp://127.0.0.1:62100".to_string(), &topology(&args));
     assert!(config.api_server_options.enable_request_id_headers);
 }
 
@@ -895,7 +946,7 @@ fn serve_passes_enable_prompt_tokens_details_into_config() {
     let Command::Serve(args) = cli.command else {
         panic!("expected serve args");
     };
-    let config = args.to_frontend_config("tcp://127.0.0.1:62100".to_string());
+    let config = args.to_frontend_config("tcp://127.0.0.1:62100".to_string(), &topology(&args));
     assert!(config.api_server_options.enable_prompt_tokens_details);
 }
 
@@ -912,7 +963,7 @@ fn serve_passes_return_tokens_as_token_ids_into_config() {
     let Command::Serve(args) = cli.command else {
         panic!("expected serve args");
     };
-    let config = args.to_frontend_config("tcp://127.0.0.1:62100".to_string());
+    let config = args.to_frontend_config("tcp://127.0.0.1:62100".to_string(), &topology(&args));
     assert!(config.api_server_options.default_return_tokens_as_token_ids);
 }
 
@@ -924,7 +975,7 @@ fn serve_rejects_log_outputs_without_log_requests() {
         let Command::Serve(args) = cli.command else {
             panic!("expected serve args");
         };
-        args.to_frontend_config("tcp://127.0.0.1:62100".to_string())
+        args.to_frontend_config("tcp://127.0.0.1:62100".to_string(), &topology(&args))
     };
 
     let error = config(&["--enable-log-outputs"]).validate().unwrap_err();
@@ -940,9 +991,10 @@ fn serve_resolves_output_logging_from_log_flags() {
         let Command::Serve(args) = cli.command else {
             panic!("expected serve args");
         };
-        args.to_frontend_config("tcp://127.0.0.1:62100".to_string())
+        args.to_frontend_config("tcp://127.0.0.1:62100".to_string(), &topology(&args))
             .api_server_options
             .log_outputs
+            .outputs
     };
 
     assert_eq!(
@@ -971,7 +1023,7 @@ fn frontend_args_json_passes_output_logging_into_config() {
         "--output-listener-fd",
         "999998",
         "--args-json",
-        r#"{"model_tag":"Qwen/Qwen3-0.6B","enable_log_requests":true,"enable_log_outputs":true,"enable_log_deltas":false}"#,
+        r#"{"model_tag":"Qwen/Qwen3-0.6B","enable_log_requests":true,"enable_log_outputs":true,"enable_log_deltas":false,"max_log_len":5}"#,
     ])
     .unwrap();
 
@@ -981,10 +1033,73 @@ fn frontend_args_json_passes_output_logging_into_config() {
     let config = args.into_config();
     assert_eq!(
         config.api_server_options.log_outputs,
-        OutputLogging::Complete
+        OutputLogConfig {
+            outputs: OutputLogging::Complete,
+            max_log_len: Some(5),
+        }
     );
 }
 
+#[test]
+fn frontend_args_json_passes_access_log_flags_into_config() {
+    let config_for = |args_json: &str| {
+        let cli = Cli::try_parse_from([
+            "vllm-rs",
+            "frontend",
+            "--listen-fd",
+            "3",
+            "--input-listener-fd",
+            "999999",
+            "--output-listener-fd",
+            "999998",
+            "--args-json",
+            args_json,
+        ])
+        .unwrap();
+        let Command::Frontend(args) = cli.command else {
+            panic!("expected frontend args");
+        };
+        args.into_config().access_log
+    };
+
+    expect![[r#"
+        AccessLogConfig {
+            enabled: true,
+            excluded_paths: [],
+        }
+    "#]]
+    .assert_debug_eq(&config_for(r#"{"model_tag":"Qwen/Qwen3-0.6B"}"#));
+    // Like Python, the endpoint list is split on commas and trimmed.
+    expect![[r#"
+        AccessLogConfig {
+            enabled: false,
+            excluded_paths: [
+                "/health",
+                "/metrics",
+            ],
+        }
+    "#]]
+    .assert_debug_eq(&config_for(
+        r#"{"model_tag":"Qwen/Qwen3-0.6B","disable_uvicorn_access_log":true,"disable_access_log_for_endpoints":"/health, /metrics,"}"#,
+    ));
+}
+
+#[test]
+fn serve_disable_access_log_accepts_python_flag_name() {
+    let enabled_for = |flag: &str| {
+        let cli = Cli::try_parse_from(["vllm-rs", "serve", "Qwen/Qwen3-0.6B", flag]).unwrap();
+        let Command::Serve(args) = cli.command else {
+            panic!("expected serve args");
+        };
+        let config = args.to_frontend_config("tcp://127.0.0.1:62100".to_string(), &topology(&args));
+        config.access_log.enabled
+    };
+
+    assert!(!enabled_for("--disable-access-log"));
+    assert!(!enabled_for("--disable-uvicorn-access-log"));
+    assert!(enabled_for("--disable-access-log=false"));
+}
+
 #[test]
 fn serve_passes_tls_into_config() {
     let cli = Cli::try_parse_from([
@@ -1005,7 +1120,7 @@ fn serve_passes_tls_into_config() {
     let Command::Serve(args) = cli.command else {
         panic!("expected serve args");
     };
-    let config = args.to_frontend_config("tcp://127.0.0.1:62100".to_string());
+    let config = args.to_frontend_config("tcp://127.0.0.1:62100".to_string(), &topology(&args));
     let tls = config.tls.expect("tls configured");
     assert_eq!(tls.cert_file.as_deref(), Some("/tmp/cert.pem"));
     assert_eq!(tls.key_file.as_deref(), Some("/tmp/key.pem"));
@@ -1020,7 +1135,7 @@ fn serve_without_ssl_flags_has_no_tls() {
     let Command::Serve(args) = cli.command else {
         panic!("expected serve args");
     };
-    let config = args.to_frontend_config("tcp://127.0.0.1:62100".to_string());
+    let config = args.to_frontend_config("tcp://127.0.0.1:62100".to_string(), &topology(&args));
     assert!(config.tls.is_none());
 }
 
@@ -1038,7 +1153,7 @@ fn serve_ssl_keyfile_without_certfile_fails_validation() {
     let Command::Serve(args) = cli.command else {
         panic!("expected serve args");
     };
-    let config = args.to_frontend_config("tcp://127.0.0.1:62100".to_string());
+    let config = args.to_frontend_config("tcp://127.0.0.1:62100".to_string(), &topology(&args));
     // TLS is requested (a key was given) but there is no certificate, so
     // validation fails loud rather than silently serving plaintext.
     assert_eq!(config.tls.as_ref().expect("tls requested").cert_file, None);
@@ -1062,7 +1177,7 @@ fn serve_mtls_without_ca_certs_fails_validation() {
     let Command::Serve(args) = cli.command else {
         panic!("expected serve args");
     };
-    let config = args.to_frontend_config("tcp://127.0.0.1:62100".to_string());
+    let config = args.to_frontend_config("tcp://127.0.0.1:62100".to_string(), &topology(&args));
     // Client-cert verification without a CA bundle has nothing to verify
     // against, so it fails loud at startup.
     let err = config.validate().unwrap_err().to_string();
@@ -1162,7 +1277,7 @@ fn serve_args_parse_lora_modules_in_both_forms() {
     };
     // Frontend-only flag: the managed Python engine must not see it.
     assert_eq!(args.managed_engine.python_args, vec!["--enable-lora"]);
-    let config = args.to_frontend_config("tcp://127.0.0.1:62100".to_string());
+    let config = args.to_frontend_config("tcp://127.0.0.1:62100".to_string(), &topology(&args));
     expect![[r#"
         [
             LoraModulePath {
@@ -1297,7 +1412,7 @@ fn serve_passes_api_keys_into_config() {
     let Command::Serve(args) = cli.command else {
         panic!("expected serve args");
     };
-    let config = args.to_frontend_config("tcp://127.0.0.1:62100".to_string());
+    let config = args.to_frontend_config("tcp://127.0.0.1:62100".to_string(), &topology(&args));
     assert_eq!(config.api_keys, vec!["secret-a", "secret-b"]);
     let debug = format!("{config:#?}");
     assert!(debug.contains("api_keys: [<redacted>; 2]"));
@@ -1363,7 +1478,7 @@ fn serve_args_reject_unknown_renderer_value() {
     .unwrap_err();
 
     expect![[r#"
-        error: invalid value 'definitely_missing' for '--tokenizer-mode <RENDERER>': unknown renderer `definitely_missing` (expected one of: auto, hf, deepseek_v32, deepseek_v4, deepseek_v41, harmony, inkling, kimi_k3)
+        error: invalid value 'definitely_missing' for '--tokenizer-mode <RENDERER>': unknown renderer `definitely_missing` (expected one of: auto, hf, deepseek_v32, deepseek_v4, deepseek_v41, harmony, inkling, kimi_k3, glm5_next)
 
         For more information, try '--help'.
     "#]]
@@ -1376,13 +1491,13 @@ fn serve_args_reject_unsupported_flag_arg() {
         "vllm-rs",
         "serve",
         "Qwen/Qwen3-0.6B",
-        "--root-path",
-        "/prefix",
+        "--uvicorn-log-level",
+        "debug",
     ])
     .unwrap_err();
 
     expect![[r#"
-        error: invalid value '/prefix' for '--root-path <ROOT_PATH>': argument is not implemented in Rust frontend yet
+        error: invalid value 'debug' for '--uvicorn-log-level <UVICORN_LOG_LEVEL>': argument is not implemented in Rust frontend yet
 
         Remove this unsupported argument to continue.
 
@@ -1405,7 +1520,7 @@ fn serve_args_enable_per_request_metrics_requires_engine_stats() {
     let Command::Serve(args) = cli.command else {
         panic!("expected serve args");
     };
-    let config = args.to_frontend_config("tcp://127.0.0.1:0".into());
+    let config = args.to_frontend_config("tcp://127.0.0.1:0".into(), &topology(&args));
     assert!(config.api_server_options.enable_per_request_metrics);
     assert!(config.validate().is_ok());
 
@@ -1420,7 +1535,10 @@ fn serve_args_enable_per_request_metrics_requires_engine_stats() {
     let Command::Serve(args) = cli.command else {
         panic!("expected serve args");
     };
-    let error = args.to_frontend_config("tcp://127.0.0.1:0".into()).validate().unwrap_err();
+    let error = args
+        .to_frontend_config("tcp://127.0.0.1:0".into(), &topology(&args))
+        .validate()
+        .unwrap_err();
     assert!(
         format!("{error:#}")
             .contains("--enable-per-request-metrics requires engine statistics logging")
@@ -1535,6 +1653,7 @@ fn frontend_args_accept_json() {
                         enable_log_requests: false,
                         enable_log_outputs: false,
                         enable_log_deltas: true,
+                        max_log_len: None,
                         enable_prompt_tokens_details: false,
                         enable_per_request_metrics: false,
                         enable_request_id_headers: false,
@@ -1548,6 +1667,7 @@ fn frontend_args_accept_json() {
                         log_stats_interval: None,
                         disable_request_id_randomization: false,
                         http_max_request_body_size: None,
+                        skip_model_name_validation: false,
                         served_model_name: [],
                         allowed_origins: JsonStringList(
                             [],
@@ -1563,6 +1683,8 @@ fn frontend_args_accept_json() {
                             ],
                         ),
                         allow_credentials: false,
+                        disable_access_log: false,
+                        disable_access_log_for_endpoints: None,
                         ssl: SslArgs {
                             ssl_keyfile: None,
                             ssl_certfile: None,
@@ -1570,6 +1692,7 @@ fn frontend_args_accept_json() {
                             ssl_cert_reqs: 0,
                             ssl_ciphers: None,
                         },
+                        root_path: None,
                         profiler_config: None,
                     },
                 },
@@ -1651,7 +1774,7 @@ fn serve_args_resolve_model_sources_and_forward_them_to_managed_engine() {
         panic!("expected serve args")
     };
 
-    let config = args.to_frontend_config("tcp://127.0.0.1:1234".to_string());
+    let config = args.to_frontend_config("tcp://127.0.0.1:1234".to_string(), &topology(&args));
     expect![[r#"
         (
             Some(
@@ -1683,7 +1806,7 @@ fn serve_args_resolve_model_sources_and_forward_them_to_managed_engine() {
             "org/config",
         ]
     "#]]
-    .assert_debug_eq(&args.to_managed_engine_config(1234).python_args);
+    .assert_debug_eq(&args.to_managed_engine_config(1234, &topology(&args), None).python_args);
 }
 
 #[test]
@@ -1749,14 +1872,14 @@ fn hf_overrides_preserve_merge_patch_in_cli_and_python_bootstrap() {
         serde_json::to_value(&args.runtime.hf_overrides).unwrap(),
         expected
     );
-    let engine = args.to_managed_engine_config(1234);
+    let engine = args.to_managed_engine_config(1234, &topology(&args), None);
     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());
+    let config = args.to_frontend_config("tcp://127.0.0.1:1234".to_string(), &topology(&args));
     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});
@@ -1849,7 +1972,7 @@ fn serve_args_forward_auto_max_model_len_to_managed_engine() {
     };
     assert_eq!(args.managed_engine.max_model_len.as_deref(), Some("auto"));
 
-    let config = args.to_managed_engine_config(5555);
+    let config = args.to_managed_engine_config(5555, &topology(&args), None);
     expect![[r#"
         [
             "--max-model-len",
@@ -2001,14 +2124,14 @@ fn frontend_args_json_rejects_unsupported_fields() {
         "--output-listener-fd",
         "999998",
         "--args-json",
-        r#"{"model_tag":"Qwen/Qwen3-0.6B","root_path":"/prefix"}"#,
+        r#"{"model_tag":"Qwen/Qwen3-0.6B","uvicorn_log_level":"debug"}"#,
     ])
     .unwrap_err();
 
     expect![[r#"
-        error: invalid value '{"model_tag":"Qwen/Qwen3-0.6B","root_path":"/prefix"}' for '--args-json <JSON>': 
+        error: invalid value '{"model_tag":"Qwen/Qwen3-0.6B","uvicorn_log_level":"debug"}' for '--args-json <JSON>': 
         The following arguments are not implemented in Rust frontend yet:
-        - root_path
+        - uvicorn_log_level
 
         Remove these arguments to continue.
 
@@ -2028,16 +2151,16 @@ fn frontend_args_json_aggregates_multiple_unsupported_fields() {
         "--output-listener-fd",
         "999998",
         "--args-json",
-        r#"{"model_tag":"Qwen/Qwen3-0.6B","response_role":"assistant","root_path":"/prefix"}"#,
+        r#"{"model_tag":"Qwen/Qwen3-0.6B","response_role":"assistant","uvicorn_log_level":"debug"}"#,
     ])
     .unwrap_err();
 
     let actual = error.to_string().replace(": \n", ":\n");
     expect![[r#"
-        error: invalid value '{"model_tag":"Qwen/Qwen3-0.6B","response_role":"assistant","root_path":"/prefix"}' for '--args-json <JSON>':
+        error: invalid value '{"model_tag":"Qwen/Qwen3-0.6B","response_role":"assistant","uvicorn_log_level":"debug"}' for '--args-json <JSON>':
         The following arguments are not implemented in Rust frontend yet:
         - response_role
-        - root_path
+        - uvicorn_log_level
 
         Remove these arguments to continue.
 
@@ -2088,7 +2211,8 @@ fn custom_histogram_buckets_reach_frontend_config() {
     let Command::Serve(args) = cli.command else {
         panic!("expected serve args");
     };
-    let serve_config = args.to_frontend_config("tcp://127.0.0.1:62100".to_string());
+    let serve_config =
+        args.to_frontend_config("tcp://127.0.0.1:62100".to_string(), &topology(&args));
 
     let args_json =
         format!(r#"{{"model_tag":"Qwen/Qwen3-0.6B","custom_histogram_buckets":{buckets}}}"#);
@@ -2456,6 +2580,7 @@ fn serve_args_accept_handshake_aliases() {
                         enable_log_requests: false,
                         enable_log_outputs: false,
                         enable_log_deltas: true,
+                        max_log_len: None,
                         enable_prompt_tokens_details: false,
                         enable_per_request_metrics: false,
                         enable_request_id_headers: false,
@@ -2469,6 +2594,7 @@ fn serve_args_accept_handshake_aliases() {
                         log_stats_interval: None,
                         disable_request_id_randomization: false,
                         http_max_request_body_size: None,
+                        skip_model_name_validation: false,
                         served_model_name: [],
                         allowed_origins: JsonStringList(
                             [],
@@ -2484,6 +2610,8 @@ fn serve_args_accept_handshake_aliases() {
                             ],
                         ),
                         allow_credentials: false,
+                        disable_access_log: false,
+                        disable_access_log_for_endpoints: None,
                         ssl: SslArgs {
                             ssl_keyfile: None,
                             ssl_certfile: None,
@@ -2491,6 +2619,7 @@ fn serve_args_accept_handshake_aliases() {
                             ssl_cert_reqs: 0,
                             ssl_ciphers: None,
                         },
+                        root_path: None,
                         profiler_config: None,
                     },
                     managed_engine: ManagedEngineArgs {
@@ -2504,6 +2633,12 @@ fn serve_args_accept_handshake_aliases() {
                         max_model_len: None,
                         python_args: [],
                     },
+                    data_parallel_lb: DataParallelLbArgs {
+                        data_parallel_rank: None,
+                        data_parallel_start_rank: None,
+                        data_parallel_hybrid_lb: false,
+                        data_parallel_external_lb: false,
+                    },
                 },
             ),
         }
@@ -2535,6 +2670,113 @@ fn serve_args_accept_data_parallel_primary_flags() {
     assert_eq!(args.managed_engine.data_parallel_size, 4);
 }
 
+#[test]
+fn serve_rank_zero_frontend_hosts_the_coordinator_for_its_single_rank() {
+    let cli = Cli::try_parse_from([
+        "vllm-rs",
+        "serve",
+        "Qwen/Qwen3-0.6B",
+        "--data-parallel-rpc-port",
+        "13345",
+        "--data-parallel-size",
+        "2",
+        "--data-parallel-size-local",
+        "1",
+        "--data-parallel-hybrid-lb",
+    ])
+    .unwrap();
+    let Command::Serve(args) = cli.command else {
+        panic!("expected serve args");
+    };
+
+    // Coordinated DP spans both ranks even though this frontend connects to one.
+    let config = args.to_frontend_config("tcp://127.0.0.1:13345".to_string(), &topology(&args));
+    assert_eq!(
+        config.effective_coordinator_mode(true),
+        Some(EngineCoreCoordinatorMode::InProc)
+    );
+    assert_eq!(config.effective_coordinator_mode(false), None);
+}
+
+#[test]
+fn serve_hybrid_lb_frontend_connects_only_its_ranks() {
+    let cli = Cli::try_parse_from([
+        "vllm-rs",
+        "serve",
+        "Qwen/Qwen3-0.6B",
+        "--data-parallel-address",
+        "10.99.48.128",
+        "--data-parallel-rpc-port",
+        "13345",
+        "--data-parallel-size",
+        "4",
+        "--data-parallel-size-local",
+        "2",
+        "--data-parallel-start-rank",
+        "2",
+    ])
+    .unwrap();
+    let Command::Serve(args) = cli.command else {
+        panic!("expected serve args");
+    };
+    let topology = topology(&args);
+
+    // A non-zero rank's frontend serves only its own engines' second handshake,
+    // over local IPC, and runs no coordinator.
+    let config = args.to_frontend_config("ipc:///tmp/vllm-rs-h".to_string(), &topology);
+    let TransportMode::HandshakeOwner {
+        handshake_address,
+        engine_start_index,
+        engine_count,
+        data_parallel_size,
+        local_input_address,
+        ..
+    } = &config.transport_mode
+    else {
+        panic!("expected handshake-owned transport");
+    };
+    assert_eq!(handshake_address, "ipc:///tmp/vllm-rs-h");
+    assert_eq!(
+        (*engine_start_index, *engine_count, *data_parallel_size),
+        (2, 2, 4)
+    );
+    assert!(
+        local_input_address
+            .as_deref()
+            .is_some_and(|address| address.starts_with("ipc://"))
+    );
+    assert_eq!(config.effective_coordinator_mode(true), None);
+
+    // The engines dial the rank-0 handshake and this frontend's local one.
+    let engine =
+        args.to_managed_engine_config(13345, &topology, Some("ipc:///tmp/vllm-rs-h".to_string()));
+    let command = engine.to_command();
+    expect![[r#"
+        [
+            "-m",
+            "vllm.entrypoints.cli.main",
+            "serve",
+            "Qwen/Qwen3-0.6B",
+            "--headless",
+            "--data-parallel-address",
+            "10.99.48.128",
+            "--data-parallel-rpc-port",
+            "13345",
+            "--data-parallel-size",
+            "4",
+            "--headless-client-handshake-address",
+            "ipc:///tmp/vllm-rs-h",
+            "--reasoning-parser",
+            "qwen3",
+            "--data-parallel-size-local",
+            "2",
+            "--data-parallel-start-rank",
+            "2",
+        ]
+    "#]]
+    .assert_debug_eq(&command.get_args().collect::<Vec<_>>());
+}
+
 #[test]
 fn serve_frontend_config_uses_dp_address_as_advertised_host() {
     let cli = Cli::try_parse_from([
@@ -2551,12 +2793,14 @@ fn serve_frontend_config_uses_dp_address_as_advertised_host() {
     let Command::Serve(args) = cli.command else {
         panic!("expected serve args");
     };
-    let config = args.to_frontend_config("tcp://10.99.48.128:29550".to_string());
+    let config = args.to_frontend_config("tcp://10.99.48.128:29550".to_string(), &topology(&args));
 
     let TransportMode::HandshakeOwner {
         handshake_address,
         advertised_host,
+        engine_start_index,
         engine_count,
+        data_parallel_size,
         ready_timeout,
         local_input_address,
         local_output_address,
@@ -2567,7 +2811,9 @@ fn serve_frontend_config_uses_dp_address_as_advertised_host() {
 
     assert_eq!(handshake_address, "tcp://10.99.48.128:29550");
     assert_eq!(advertised_host, "10.99.48.128");
+    assert_eq!(*engine_start_index, 0);
     assert_eq!(*engine_count, 4);
+    assert_eq!(*data_parallel_size, 4);
     assert!(
         local_input_address
             .as_deref()
@@ -2585,7 +2831,9 @@ fn serve_frontend_config_uses_dp_address_as_advertised_host() {
             transport_mode: HandshakeOwner {
                 handshake_address: "tcp://10.99.48.128:29550",
                 advertised_host: "10.99.48.128",
+                engine_start_index: 0,
                 engine_count: 4,
+                data_parallel_size: 4,
                 ready_timeout: 600s,
                 local_input_address: Some(
                     "<ipc input>",
@@ -2627,13 +2875,17 @@ fn serve_frontend_config_uses_dp_address_as_advertised_host() {
             stream_interval: 1,
             api_server_options: ApiServerOptions {
                 enable_log_requests: false,
-                log_outputs: Disabled,
+                log_outputs: OutputLogConfig {
+                    outputs: Disabled,
+                    max_log_len: None,
+                },
                 enable_prompt_tokens_details: false,
                 default_return_tokens_as_token_ids: false,
                 enable_per_request_metrics: false,
                 enable_request_id_headers: false,
                 enable_scale_out: false,
                 enable_tokenizer_info_endpoint: false,
+                skip_model_name_validation: false,
                 sse_keep_alive_interval: None,
                 request_body_limit_bytes: Some(
                     33554432,
@@ -2649,7 +2901,12 @@ fn serve_frontend_config_uses_dp_address_as_advertised_host() {
                 ],
                 allow_credentials: false,
             },
+            access_log: AccessLogConfig {
+                enabled: true,
+                excluded_paths: [],
+            },
             tls: None,
+            root_path: None,
             api_keys: [],
             disable_log_stats: false,
             custom_histogram_buckets: CustomHistogramBuckets(
@@ -2669,7 +2926,9 @@ fn serve_frontend_config_uses_dp_address_as_advertised_host() {
         transport_mode: TransportMode::HandshakeOwner {
             handshake_address: handshake_address.clone(),
             advertised_host: advertised_host.clone(),
+            engine_start_index: 0,
             engine_count: *engine_count,
+            data_parallel_size: *engine_count,
             ready_timeout: *ready_timeout,
             local_input_address: Some("<ipc input>".to_string()),
             local_output_address: Some("<ipc output>".to_string()),
@@ -2696,14 +2955,16 @@ fn serve_frontend_config_keeps_tcp_transport_for_non_local_only_topology() {
     let Command::Serve(args) = cli.command else {
         panic!("expected serve args");
     };
-    let config = args.to_frontend_config("tcp://10.99.48.128:29550".to_string());
+    let config = args.to_frontend_config("tcp://10.99.48.128:29550".to_string(), &topology(&args));
 
     expect![[r#"
         Config {
             transport_mode: HandshakeOwner {
                 handshake_address: "tcp://10.99.48.128:29550",
                 advertised_host: "10.99.48.128",
+                engine_start_index: 0,
                 engine_count: 4,
+                data_parallel_size: 4,
                 ready_timeout: 600s,
                 local_input_address: None,
                 local_output_address: None,
@@ -2741,13 +3002,17 @@ fn serve_frontend_config_keeps_tcp_transport_for_non_local_only_topology() {
             stream_interval: 1,
             api_server_options: ApiServerOptions {
                 enable_log_requests: false,
-                log_outputs: Disabled,
+                log_outputs: OutputLogConfig {
+                    outputs: Disabled,
+                    max_log_len: None,
+                },
                 enable_prompt_tokens_details: false,
                 default_return_tokens_as_token_ids: false,
                 enable_per_request_metrics: false,
                 enable_request_id_headers: false,
                 enable_scale_out: false,
                 enable_tokenizer_info_endpoint: false,
+                skip_model_name_validation: false,
                 sse_keep_alive_interval: None,
                 request_body_limit_bytes: Some(
                     33554432,
@@ -2763,7 +3028,12 @@ fn serve_frontend_config_keeps_tcp_transport_for_non_local_only_topology() {
                 ],
                 allow_credentials: false,
             },
+            access_log: AccessLogConfig {
+                enabled: true,
+                excluded_paths: [],
+            },
             tls: None,
+            root_path: None,
             api_keys: [],
             disable_log_stats: false,
             custom_histogram_buckets: CustomHistogramBuckets(
@@ -2821,7 +3091,7 @@ fn serve_frontend_config_does_not_manage_engine_without_local_engines() {
     let Command::Serve(args) = cli.command else {
         panic!("expected serve args");
     };
-    let config = args.to_frontend_config("tcp://10.99.48.128:29550".to_string());
+    let config = args.to_frontend_config("tcp://10.99.48.128:29550".to_string(), &topology(&args));
 
     assert!(!config.manages_engine);
 }
@@ -2848,6 +3118,35 @@ fn frontend_config_trusts_request_chat_template_from_args_json() {
     assert!(args.into_config().trust_request_chat_template);
 }
 
+#[test]
+fn root_path_reaches_config_from_serve_args_and_args_json() {
+    let cli = Cli::try_parse_from(["vllm-rs", "serve", "Qwen/Qwen3-0.6B", "--root-path", "/llm"])
+        .unwrap();
+    let Command::Serve(args) = cli.command else {
+        panic!("expected serve args");
+    };
+    let config = args.to_frontend_config("tcp://127.0.0.1:62100".to_string(), &topology(&args));
+    assert_eq!(config.root_path.as_deref(), Some("/llm"));
+
+    let cli = Cli::try_parse_from([
+        "vllm-rs",
+        "frontend",
+        "--listen-fd",
+        "3",
+        "--input-listener-fd",
+        "999999",
+        "--output-listener-fd",
+        "999998",
+        "--args-json",
+        r#"{"model_tag":"Qwen/Qwen3-0.6B","root_path":"/llm"}"#,
+    ])
+    .unwrap();
+    let Command::Frontend(args) = cli.command else {
+        panic!("expected frontend args");
+    };
+    assert_eq!(args.into_config().root_path.as_deref(), Some("/llm"));
+}
+
 #[test]
 fn frontend_config_excludes_tools_when_tool_choice_none_from_args_json() {
     let cli = Cli::try_parse_from([
@@ -2943,13 +3242,17 @@ fn frontend_config_uses_external_coordinator_when_coordinator_address_is_present
             stream_interval: 1,
             api_server_options: ApiServerOptions {
                 enable_log_requests: false,
-                log_outputs: Disabled,
+                log_outputs: OutputLogConfig {
+                    outputs: Disabled,
+                    max_log_len: None,
+                },
                 enable_prompt_tokens_details: false,
                 default_return_tokens_as_token_ids: false,
                 enable_per_request_metrics: false,
                 enable_request_id_headers: false,
                 enable_scale_out: false,
                 enable_tokenizer_info_endpoint: false,
+                skip_model_name_validation: false,
                 sse_keep_alive_interval: None,
                 request_body_limit_bytes: Some(
                     33554432,
@@ -2965,7 +3268,12 @@ fn frontend_config_uses_external_coordinator_when_coordinator_address_is_present
                 ],
                 allow_credentials: false,
             },
+            access_log: AccessLogConfig {
+                enabled: true,
+                excluded_paths: [],
+            },
             tls: None,
+            root_path: None,
             api_keys: [],
             disable_log_stats: false,
             custom_histogram_buckets: CustomHistogramBuckets(
@@ -3000,7 +3308,7 @@ fn serve_frontend_config_uses_unix_listener_when_uds_is_present() {
     let Command::Serve(args) = cli.command else {
         panic!("expected serve args");
     };
-    let config = args.to_frontend_config("tcp://127.0.0.1:29550".to_string());
+    let config = args.to_frontend_config("tcp://127.0.0.1:29550".to_string(), &topology(&args));
 
     assert_eq!(
         config.listener_mode,
@@ -3033,7 +3341,7 @@ fn serve_frontend_config_binds_grpc_on_http_host() {
     let Command::Serve(args) = cli.command else {
         panic!("expected serve args");
     };
-    let config = args.to_frontend_config("tcp://127.0.0.1:29550".to_string());
+    let config = args.to_frontend_config("tcp://127.0.0.1:29550".to_string(), &topology(&args));
 
     assert_eq!(
         config.grpc_listener_mode,
@@ -3153,9 +3461,11 @@ fn serve_revision_reaches_frontend_and_managed_engine() {
         panic!("expected serve args");
     };
     assert_eq!(
-        args.to_frontend_config("tcp://localhost:1234".into()).revision.as_deref(),
+        args.to_frontend_config("tcp://localhost:1234".into(), &topology(&args))
+            .revision
+            .as_deref(),
         Some("release")
     );
-    let engine = args.to_managed_engine_config(1234);
+    let engine = args.to_managed_engine_config(1234, &topology(&args), None);
     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 6c60d65af430..52207a174999 100644
--- a/rust/src/cmd/src/cli/unsupported.rs
+++ b/rust/src/cmd/src/cli/unsupported.rs
@@ -242,36 +242,6 @@ pub struct EngineUnsupportedArgs {
     #[arg(long)]
     pub reasoning_parser_plugin: Option<Unsupported>,
 
-    /// Rank of the data parallel group.
-    #[arg(long, env = "VLLM_DP_RANK")]
-    pub data_parallel_rank: Option<Unsupported>,
-
-    /// Whether to use "hybrid" DP LB mode. Applies only to online serving
-    /// and when data_parallel_size > 0. Enables running an AsyncLLM
-    /// and API server on a "per-node" basis where vLLM load balances
-    /// between local data parallel ranks, but an external LB balances
-    /// between vLLM nodes/replicas. Set explicitly in conjunction with
-    /// --data-parallel-start-rank.
-    #[arg(
-        long,
-        visible_alias = "no-data-parallel-hybrid-lb",
-        default_missing_value = "true",
-        num_args = 0..=1
-    )]
-    pub data_parallel_hybrid_lb: Option<Unsupported>,
-
-    /// Whether to use "external" DP LB mode. Applies only to online serving
-    /// and when data_parallel_size > 0. This is useful for a "one-pod-per-rank"
-    /// wide-EP setup in Kubernetes. Set implicitly when --data-parallel-rank
-    /// is provided explicitly to vllm serve.
-    #[arg(
-        long,
-        visible_alias = "no-data-parallel-external-lb",
-        default_missing_value = "true",
-        num_args = 0..=1
-    )]
-    pub data_parallel_external_lb: Option<Unsupported>,
-
     /// Enable fault tolerance for detailed error recovery, such as scaling
     /// down fault DPEngineCore.
     #[arg(
@@ -439,11 +409,6 @@ pub struct ServerUnsupportedArgs {
     #[arg(long, /* env = "VLLM_LOGGING_CONFIG_PATH" */)]
     pub log_config_file: Option<Unsupported>,
 
-    /// Max number of prompt characters or prompt ID numbers being printed in
-    /// log. The default of None means unlimited.
-    #[arg(long)]
-    pub max_log_len: Option<Unsupported>,
-
     /// If set to True, enable tracking server_load_metrics in the app state.
     #[arg(
         long,
@@ -495,23 +460,6 @@ pub struct ServerUnsupportedArgs {
     #[arg(long)]
     pub uvicorn_log_level: Option<Unsupported>,
 
-    /// Disable uvicorn access log.
-    #[arg(
-        long,
-        visible_alias = "no-disable-uvicorn-access-log",
-        default_missing_value = "true",
-        num_args = 0..=1
-    )]
-    pub disable_uvicorn_access_log: Option<Noop>,
-
-    /// Comma-separated list of endpoint paths to exclude from uvicorn access
-    /// logs. This is useful to reduce log noise from high-frequency endpoints
-    /// like health checks. Example: "/health,/metrics,/ping".
-    /// When set, access logs for requests to these paths will be suppressed
-    /// while keeping logs for other endpoints.
-    #[arg(long)]
-    pub disable_access_log_for_endpoints: Option<Noop>,
-
     /// Refresh SSL Context when SSL certificate files change
     #[arg(
         long,
@@ -521,10 +469,6 @@ pub struct ServerUnsupportedArgs {
     )]
     pub enable_ssl_refresh: Option<Unsupported>,
 
-    /// FastAPI root_path when app is behind a path based routing proxy.
-    #[arg(long)]
-    pub root_path: Option<Unsupported>,
-
     /// Additional ASGI middleware to apply to the app. We accept multiple
     /// --middleware arguments. The value should be an import path. If a
     /// function is provided, vLLM will add it to the server using
diff --git a/rust/src/cmd/src/main.rs b/rust/src/cmd/src/main.rs
index b9535561e417..da4945f93d18 100644
--- a/rust/src/cmd/src/main.rs
+++ b/rust/src/cmd/src/main.rs
@@ -13,7 +13,7 @@ use tokio_util::sync::CancellationToken;
 use tracing::{error, info, warn};
 use vllm_managed_engine::ManagedEngineHandle;
 
-use crate::cli::{BenchCommand, Cli, Command};
+use crate::cli::{BenchCommand, Cli, Command, frontend_ipc_handshake_address};
 
 #[global_allocator]
 static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
@@ -121,6 +121,7 @@ async fn async_main(cli: Cli) -> Result<()> {
             vllm_bench::run(bench_args).await
         }
         Command::Serve(args) => {
+            let topology = args.data_parallel_lb.resolve(&args.managed_engine, args.headless)?;
             let handshake_port = args.managed_engine.resolve_handshake_port()?;
 
             if args.managed_engine.data_parallel_size_local == Some(0) {
@@ -134,13 +135,24 @@ async fn async_main(cli: Cli) -> Result<()> {
                     engine_count = args.managed_engine.data_parallel_size,
                     "running Rust frontend without a managed local Python engine"
                 );
-                let config = args.to_frontend_config(handshake_address);
+                let config = args.to_frontend_config(handshake_address, &topology);
                 return vllm_server::serve(config, shutdown_signal()).await;
             }
 
             let shutdown_timeout = args.runtime.shutdown_timeout();
-            let engine_config = args.to_managed_engine_config(handshake_port);
-            let handshake_address = engine_config.handshake_address();
+            // The frontend owning rank 0 hosts the deployment's handshake. In hybrid or
+            // external DP LB, another node's frontend serves only its engines' second
+            // handshake, on a local IPC endpoint.
+            let client_handshake_address = (!args.headless
+                && topology.frontend_engine_start_index() != 0)
+                .then(frontend_ipc_handshake_address);
+            let engine_config = args.to_managed_engine_config(
+                handshake_port,
+                &topology,
+                client_handshake_address.clone(),
+            );
+            let handshake_address =
+                client_handshake_address.unwrap_or_else(|| engine_config.handshake_address());
 
             let engine = ManagedEngineHandle::spawn(engine_config)
                 .await
@@ -156,7 +168,7 @@ async fn async_main(cli: Cli) -> Result<()> {
                     Ok(())
                 })
             } else {
-                let config = args.to_frontend_config(handshake_address);
+                let config = args.to_frontend_config(handshake_address, &topology);
                 let shutdown = shutdown.clone();
                 tokio::spawn(async move {
                     let result = vllm_server::serve(config, shutdown).await;
diff --git a/rust/src/engine-core-client/Cargo.toml b/rust/src/engine-core-client/Cargo.toml
index 86ffe4c4cdae..127604fa0dc6 100644
--- a/rust/src/engine-core-client/Cargo.toml
+++ b/rust/src/engine-core-client/Cargo.toml
@@ -28,6 +28,7 @@ serde_json.workspace = true
 serde_repr.workspace = true
 serde_tuple.workspace = true
 serde_with.workspace = true
+sha2.workspace = true
 socket2.workspace = true
 task-local.workspace = true
 tempfile = { workspace = true, optional = true }
diff --git a/rust/src/engine-core-client/examples/external_engine_logprobs.rs b/rust/src/engine-core-client/examples/external_engine_logprobs.rs
index e0da15623ca3..5b8b48004c1a 100644
--- a/rust/src/engine-core-client/examples/external_engine_logprobs.rs
+++ b/rust/src/engine-core-client/examples/external_engine_logprobs.rs
@@ -111,7 +111,9 @@ async fn main() -> Result<()> {
         transport_mode: TransportMode::HandshakeOwner {
             handshake_address: args.handshake_address.clone(),
             advertised_host: args.host.clone(),
+            engine_start_index: 0,
             engine_count: args.engine_count,
+            data_parallel_size: args.engine_count,
             ready_timeout,
             local_input_address: None,
             local_output_address: None,
diff --git a/rust/src/engine-core-client/examples/external_engine_utility_call.rs b/rust/src/engine-core-client/examples/external_engine_utility_call.rs
index 8b73a12ce2f6..bf06aa33030f 100644
--- a/rust/src/engine-core-client/examples/external_engine_utility_call.rs
+++ b/rust/src/engine-core-client/examples/external_engine_utility_call.rs
@@ -6,7 +6,7 @@ use std::time::Duration;
 use anyhow::{Context, Result, bail};
 use clap::Parser;
 use tracing_subscriber::EnvFilter;
-use vllm_engine_core_client::protocol::utility::PauseMode;
+use vllm_engine_core_client::protocol::utility::{PauseMode, SleepLevel};
 use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig, TransportMode};
 
 #[derive(Debug, Parser)]
@@ -60,7 +60,9 @@ async fn main() -> Result<()> {
         transport_mode: TransportMode::HandshakeOwner {
             handshake_address: args.handshake_address.clone(),
             advertised_host: args.host.clone(),
+            engine_start_index: 0,
             engine_count: args.engine_count,
+            data_parallel_size: args.engine_count,
             ready_timeout: Duration::from_secs(args.ready_timeout_secs),
             local_input_address: None,
             local_output_address: None,
@@ -111,7 +113,8 @@ async fn main() -> Result<()> {
     if args.skip_sleep_wake {
         println!("sleep_wake=skipped");
     } else {
-        client.sleep(args.sleep_level, args.sleep_mode).await.with_context(|| {
+        let sleep_level = SleepLevel::try_from(args.sleep_level).map_err(anyhow::Error::msg)?;
+        client.sleep(sleep_level, args.sleep_mode).await.with_context(|| {
             format!(
                 "failed to call sleep utility with level={} mode={}",
                 args.sleep_level, args.sleep_mode
diff --git a/rust/src/engine-core-client/src/client.rs b/rust/src/engine-core-client/src/client.rs
index 6eb9b9197260..8f0528622f94 100644
--- a/rust/src/engine-core-client/src/client.rs
+++ b/rust/src/engine-core-client/src/client.rs
@@ -15,18 +15,22 @@ use tokio_util::task::AbortOnDropHandle;
 use tracing::{debug, info, trace};
 
 use crate::client::imp::{ClientInner, run_abort_loop, run_output_dispatcher_loop};
+use crate::client::readiness::{RecheckGuard, run_readiness_loop};
 use crate::coordinator::CoordinatorHandle;
 use crate::error::{Error, Result, bail_invalid_client_config};
+use crate::metrics::EngineSleepState;
+use crate::prompt_block_hasher::PromptBlockHasher;
 use crate::protocol::dtype::ModelDtype;
 use crate::protocol::handshake::EngineCoreReadyResponse;
 use crate::protocol::lora::LoraRequest;
-use crate::protocol::request::{EngineCoreRequest, EngineCoreRequestType};
-use crate::protocol::utility::{CheckpointState, EngineCoreUtilityRequest, PauseMode};
+use crate::protocol::request::EngineCoreRequest;
+use crate::protocol::utility::{CheckpointState, PauseMode, SleepLevel};
 use crate::runtime::{BackgroundShutdownRuntime, build_zmq_runtime};
 use crate::transport::{self, ConnectedEngine};
 
 pub(crate) mod imp;
 mod output_channel;
+mod readiness;
 mod state;
 mod stream;
 
@@ -45,8 +49,18 @@ pub enum TransportMode {
         /// Host/IP that engines should use to connect back to the frontend
         /// transport sockets.
         advertised_host: String,
-        /// Total number of engines expected to join this transport.
+        /// First data-parallel engine rank that this frontend connects to.
+        ///
+        /// The frontend whose range starts at rank 0 handshakes with every
+        /// rank of the deployment; other frontends handshake only with their
+        /// own engines.
+        engine_start_index: u32,
+        /// Total number of engines this frontend connects to.
         engine_count: usize,
+        /// Deployment-wide data-parallel size. This is larger than
+        /// `engine_count` when frontends partition ranks across nodes (hybrid
+        /// or external DP LB).
+        data_parallel_size: usize,
         /// Maximum time to wait for each startup phase to complete.
         ready_timeout: Duration,
         /// Optional explicit bind address for the input ROUTER socket.
@@ -81,13 +95,27 @@ impl TransportMode {
     /// Return the deployment-wide data-parallel size for this transport.
     pub fn data_parallel_size(&self) -> usize {
         match self {
-            Self::HandshakeOwner { engine_count, .. } => *engine_count,
-            Self::Bootstrapped {
+            Self::HandshakeOwner {
+                data_parallel_size, ..
+            }
+            | Self::Bootstrapped {
                 data_parallel_size, ..
             } => *data_parallel_size,
         }
     }
 
+    /// Return the first data-parallel engine rank this frontend connects to.
+    pub fn engine_start_index(&self) -> u32 {
+        match self {
+            Self::HandshakeOwner {
+                engine_start_index, ..
+            }
+            | Self::Bootstrapped {
+                engine_start_index, ..
+            } => *engine_start_index,
+        }
+    }
+
     /// Validate the transport topology before opening any sockets.
     pub fn validate(&self) -> Result<()> {
         let data_parallel_size = self.data_parallel_size();
@@ -101,8 +129,12 @@ impl TransportMode {
         }
 
         match self {
-            Self::HandshakeOwner { .. } => {}
-            Self::Bootstrapped {
+            Self::HandshakeOwner {
+                engine_start_index,
+                engine_count,
+                ..
+            }
+            | Self::Bootstrapped {
                 engine_start_index,
                 engine_count,
                 ..
@@ -171,7 +203,9 @@ impl EngineCoreClientConfig {
             transport_mode: TransportMode::HandshakeOwner {
                 handshake_address: handshake_address.into(),
                 advertised_host: "127.0.0.1".to_string(),
+                engine_start_index: 0,
                 engine_count: 1,
+                data_parallel_size: 1,
                 ready_timeout: Duration::from_secs(30),
                 local_input_address: None,
                 local_output_address: None,
@@ -284,6 +318,8 @@ pub struct EngineCoreClient {
     inner: Arc<ClientInner>,
     coordinator: Option<CoordinatorHandle>,
     abort_tx: mpsc::UnboundedSender<AbortRequest>,
+    /// Hashes prompt blocks for the engines, if they all accept the same hashes.
+    prompt_block_hasher: Option<PromptBlockHasher>,
     /// Whether a profiling session is awaiting an explicit stop. Held across each
     /// start/stop utility call so a stale stop reply cannot clear a newer session.
     profile_active: Mutex<bool>,
@@ -294,6 +330,7 @@ pub struct EngineCoreClient {
     output_task: AbortOnDropHandle<()>,
     dispatcher_task: AbortOnDropHandle<()>,
     abort_task: AbortOnDropHandle<()>,
+    readiness_task: AbortOnDropHandle<()>,
     coordinator_output_task: Option<AbortOnDropHandle<()>>,
     coordinator_task: Option<AbortOnDropHandle<()>>,
 }
@@ -307,11 +344,24 @@ impl EngineCoreClient {
     /// sockets and waits for the expected engine registration frames.
     pub async fn connect(config: EngineCoreClientConfig) -> Result<Self> {
         config.validate()?;
+        let runtime = build_zmq_runtime();
+        let handle = runtime.handle().clone();
+        let connect_task =
+            AbortOnDropHandle::new(handle.spawn(Self::connect_on_runtime(config, runtime)));
+        connect_task.await?
+    }
+
+    async fn connect_on_runtime(
+        config: EngineCoreClientConfig,
+        runtime: BackgroundShutdownRuntime,
+    ) -> Result<Self> {
         let connected = match &config.transport_mode {
             TransportMode::HandshakeOwner {
                 handshake_address,
                 advertised_host,
+                engine_start_index,
                 engine_count,
+                data_parallel_size,
                 ready_timeout,
                 local_input_address,
                 local_output_address,
@@ -326,7 +376,9 @@ impl EngineCoreClient {
 
                 transport::connect_handshake(
                     handshake_address,
+                    *engine_start_index,
                     *engine_count,
+                    *data_parallel_size,
                     advertised_host,
                     local_input_address.as_deref(),
                     local_output_address.as_deref(),
@@ -359,7 +411,7 @@ impl EngineCoreClient {
             }
         };
 
-        Self::from_connected(config, connected).await
+        Self::from_connected(config, connected, runtime).await
     }
 
     /// Create a new client instance from the connected transport state after
@@ -367,16 +419,18 @@ impl EngineCoreClient {
     async fn from_connected(
         config: EngineCoreClientConfig,
         connected: transport::ConnectedTransport,
+        runtime: BackgroundShutdownRuntime,
     ) -> Result<Self> {
         validate_lora_capabilities(&connected.engines)?;
+        let prompt_block_hasher = prompt_block_hasher(&connected.engines);
         let (output_tx, output_rx) = mpsc::channel(64);
         let (abort_tx, abort_rx) = mpsc::unbounded_channel();
         let engines = connected.engines;
-        let runtime = build_zmq_runtime();
         let inner = Arc::new(ClientInner::new(
             connected.input_send,
             runtime.handle().clone(),
             config.model_name.clone(),
+            config.client_index,
             config.engine_stats_enabled,
             &engines,
         ));
@@ -389,6 +443,8 @@ impl EngineCoreClient {
         );
         let abort_task =
             AbortOnDropHandle::new(runtime.spawn(run_abort_loop(inner.clone(), abort_rx)));
+        let readiness_task =
+            AbortOnDropHandle::new(runtime.spawn(run_readiness_loop(inner.clone())));
 
         // If any engine reported a dp_stats_address in its ready response, use it
         // as the external coordinator address.
@@ -397,8 +453,10 @@ impl EngineCoreClient {
 
         let (coordinator, coordinator_output_task, coordinator_task) =
             if let Some(coordinator_transport) = connected.coordinator {
-                let (handle, runner) =
-                    CoordinatorHandle::new_inproc(coordinator_transport.input_socket);
+                let (handle, runner) = CoordinatorHandle::new_inproc(
+                    coordinator_transport.input_socket,
+                    coordinator_transport.frontend_publish_socket,
+                );
                 let (coordinator_output_tx, coordinator_output_rx) = mpsc::channel(64);
                 let coordinator_output_task =
                     AbortOnDropHandle::new(runtime.spawn(transport::run_output_loop(
@@ -427,6 +485,14 @@ impl EngineCoreClient {
                 (None, None, None)
             };
 
+        // Only engines with sleep mode can start paused, prepared for a
+        // checkpoint. Asking here also fails startup with a clear error when
+        // the engines predate the `is_inference_ready` utility.
+        if engines.iter().all(|engine| engine.ready_response.enable_sleep_mode) {
+            let ready = inner.is_inference_ready().await?;
+            inner.readiness().start(ready);
+        }
+
         Ok(Self {
             config,
             input_address: connected.input_address,
@@ -435,11 +501,13 @@ impl EngineCoreClient {
             inner,
             coordinator,
             abort_tx,
+            prompt_block_hasher,
             profile_active: Mutex::new(false),
             runtime,
             output_task,
             dispatcher_task,
             abort_task,
+            readiness_task,
             coordinator_output_task,
             coordinator_task,
         })
@@ -473,6 +541,11 @@ impl EngineCoreClient {
         self.inner.pending_utility_call_count()
     }
 
+    #[cfg(test)]
+    pub(crate) fn readiness_recheck_pending(&self) -> bool {
+        self.inner.readiness().pending_ask().is_some()
+    }
+
     /// Return the engine-side indices connected to this client, as stamped on
     /// request output streams: an engine identity that does not encode an
     /// index maps to 0, as in [`Self::call`].
@@ -586,6 +659,31 @@ impl EngineCoreClient {
     pub fn health_error(&self) -> Option<Arc<Error>> {
         self.inner.health_error()
     }
+
+    /// Subscribe to whether the engines can accept and process new requests.
+    /// It is `false` while they are paused, asleep, or prepared for a
+    /// checkpoint, or while a weight update has not finished successfully.
+    pub fn subscribe_readiness(&self) -> tokio::sync::watch::Receiver<bool> {
+        self.inner.readiness().subscribe()
+    }
+}
+
+/// Build the prompt block hasher if every engine reports the same supported
+/// prefix hash config. Requests can go to any engine, and an engine that
+/// doesn't hash blocks has no use for the hashes.
+fn prompt_block_hasher(engines: &[ConnectedEngine]) -> Option<PromptBlockHasher> {
+    let config = (engines.iter())
+        .map(|engine| engine.ready_response.prefix_hash_config.as_ref())
+        .all_equal_value()
+        .ok()??;
+    let hasher = PromptBlockHasher::new(config);
+    if hasher.is_none() {
+        debug!(
+            ?config,
+            "engine prefix hash config not supported; engine hashes prompts"
+        );
+    }
+    hasher
 }
 
 fn validate_lora_capabilities(engines: &[ConnectedEngine]) -> Result<()> {
@@ -627,6 +725,11 @@ impl EngineCoreClient {
     pub async fn call(&self, mut req: EngineCoreRequest) -> Result<EngineCoreOutputStream> {
         req.client_index = self.config.client_index;
         req.validate()?;
+        if req.block_hashes.is_none()
+            && let Some(hasher) = &self.prompt_block_hasher
+        {
+            req.block_hashes = hasher.hash_request(&req);
+        }
         trace!(
             request_id = %req.request_id,
             client_index = req.client_index,
@@ -712,59 +815,16 @@ impl EngineCoreClient {
     /// Callers should pass utility arguments using Rust tuple semantics so the
     /// encoded payload matches Python's `(client_index, call_id,
     /// method_name, args)` contract: `()`, `(arg,)`, `(arg1, arg2)`, etc.
+    ///
+    /// Once its requests are prepared, the call is not cancellable: dropping
+    /// the returned future stops waiting, but the requests are still sent and
+    /// their replies collected.
     pub async fn call_utility<T, A>(&self, method: &str, args: A) -> Result<Vec<T>>
     where
         T: serde::de::DeserializeOwned,
         A: serde::Serialize + std::fmt::Debug,
     {
-        trace!(
-            method,
-            client_index = self.config.client_index,
-            engine_count = self.engines.len(),
-            "sending utility request"
-        );
-
-        /// Removes utility waiters if a call future is cancelled before completion.
-        struct UtilityCallGuard<'a> {
-            inner: &'a ClientInner,
-            call_ids: Vec<u64>,
-        }
-
-        impl Drop for UtilityCallGuard<'_> {
-            fn drop(&mut self) {
-                self.inner.unregister_utility_calls(self.call_ids.drain(..));
-            }
-        }
-
-        let mut call_guard = UtilityCallGuard {
-            inner: self.inner.as_ref(),
-            call_ids: Vec::with_capacity(self.engines.len()),
-        };
-
-        let mut prepared_calls = Vec::with_capacity(self.engines.len());
-        for engine in &self.engines {
-            let (call_id, rx) = self.inner.allocate_and_register_utility_call()?;
-            call_guard.call_ids.push(call_id);
-            let request =
-                EngineCoreUtilityRequest::new(self.config.client_index, call_id, method, &args)?;
-            prepared_calls.push((&engine.engine_id, call_id, rx, request));
-        }
-
-        let outcomes = join_all(prepared_calls.into_iter().map(
-            |(engine_id, call_id, rx, request)| async move {
-                self.inner
-                    .send_to_engine(engine_id, EngineCoreRequestType::Utility, &request)
-                    .await?;
-                rx.await
-                    .map_err(|_| Error::UtilityCallClosed {
-                        method: method.to_string(),
-                        call_id,
-                    })??
-                    .into_typed_result(method)
-            },
-        ))
-        .await;
-        outcomes.into_iter().collect()
+        self.inner.call_utility(method, args, None).await
     }
 
     /// Call a utility method on all connected engines and return the shared
@@ -797,12 +857,35 @@ impl EngineCoreClient {
         args: A,
         kwargs: K,
     ) -> Result<Vec<rmpv::Value>>
+    where
+        A: serde::Serialize + std::fmt::Debug,
+        K: serde::Serialize + std::fmt::Debug,
+    {
+        // Worker methods can change readiness, as a weight update does.
+        let recheck = self.inner.readiness().recheck_after_call();
+        self.collective_rpc_with(recheck, method, timeout, args, kwargs).await
+    }
+
+    /// [`Self::collective_rpc`] as a control call that passes `recheck`.
+    async fn collective_rpc_with<A, K>(
+        &self,
+        recheck: RecheckGuard,
+        method: &str,
+        timeout: Option<f64>,
+        args: A,
+        kwargs: K,
+    ) -> Result<Vec<rmpv::Value>>
     where
         A: serde::Serialize + std::fmt::Debug,
         K: serde::Serialize + std::fmt::Debug,
     {
         let results = self
-            .call_utility::<rmpv::Value, _>("collective_rpc", (method, timeout, args, kwargs))
+            .inner
+            .call_utility::<rmpv::Value, _>(
+                "collective_rpc",
+                (method, timeout, args, kwargs),
+                Some(recheck),
+            )
             .await?;
 
         Ok(results
@@ -841,7 +924,8 @@ impl EngineCoreClient {
 
     /// Start a weight update for the base model.
     pub async fn start_weight_update(&self) -> Result<()> {
-        self.collective_rpc(
+        self.collective_rpc_with(
+            self.inner.readiness().withdraw_until_recheck(),
             "start_weight_update",
             None,
             Vec::<JsonValue>::new(),
@@ -853,7 +937,8 @@ impl EngineCoreClient {
 
     /// Start a weight update for the speculative draft model.
     pub async fn start_draft_weight_update(&self) -> Result<()> {
-        self.collective_rpc(
+        self.collective_rpc_with(
+            self.inner.readiness().withdraw_until_recheck(),
             "start_draft_weight_update",
             None,
             Vec::<JsonValue>::new(),
@@ -953,37 +1038,51 @@ impl EngineCoreClient {
     }
 
     /// Put the engine to sleep.
-    pub async fn sleep(&self, level: u32, mode: PauseMode) -> Result<()> {
-        self.call_utility::<(), _>("sleep", (level, mode)).await?;
+    pub async fn sleep(&self, level: SleepLevel, mode: PauseMode) -> Result<()> {
+        let recheck = self.inner.readiness().withdraw_until_recheck();
+        self.inner.call_utility::<(), _>("sleep", (level, mode), Some(recheck)).await?;
+        self.inner.record_sleep_state(EngineSleepState::Asleep(level));
         Ok(())
     }
 
     /// 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?;
+        self.inner.record_sleep_state(EngineSleepState::KvCacheReleased);
         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,))
+        let recheck = self.inner.readiness().recheck_after_call();
+        let fully_awake = self
+            .inner
+            .call_utility::<bool, _>("wake_up", (tags,), Some(recheck))
             .await?
             .into_iter()
-            .all(|fully_awake| fully_awake))
+            .all(|fully_awake| fully_awake);
+        // A partial wake-up keeps the previous sleep state, as in Python.
+        if fully_awake {
+            self.inner.record_sleep_state(EngineSleepState::Awake);
+        }
+        Ok(fully_awake)
     }
 
     /// Pause generation (e.g. for weight updates). In `abort` and `wait` mode
     /// the engine rejects new requests with a retryable error until resume.
     pub async fn pause_generation(&self, mode: PauseMode, clear_cache: bool) -> Result<()> {
-        self.call_utility::<(), _>("pause_generation", (mode, clear_cache)).await?;
+        let recheck = self.inner.readiness().withdraw_until_recheck();
+        self.inner
+            .call_utility::<(), _>("pause_generation", (mode, clear_cache), Some(recheck))
+            .await?;
         Ok(())
     }
 
     /// Resume the scheduler after a pause
     pub async fn resume_scheduler(&self) -> Result<()> {
-        self.call_utility::<(), _>("resume_scheduler", ()).await?;
+        let recheck = self.inner.readiness().recheck_after_call();
+        self.inner.call_utility::<(), _>("resume_scheduler", (), Some(recheck)).await?;
         Ok(())
     }
 
@@ -992,14 +1091,18 @@ impl EngineCoreClient {
     /// Returns the first refusal from any engine. An engine that refuses
     /// changes nothing, but with data parallelism the others may have started.
     pub async fn checkpoint_prepare(&self, mode: PauseMode) -> Result<Option<String>> {
+        let recheck = self.inner.readiness().withdraw_until_recheck();
         let refusals: Vec<Option<String>> =
-            self.call_utility("checkpoint_prepare", (mode,)).await?;
+            self.inner.call_utility("checkpoint_prepare", (mode,), Some(recheck)).await?;
         Ok(refusals.into_iter().flatten().next())
     }
 
     /// Bring back what `checkpoint_prepare` released and serve again.
     pub async fn checkpoint_restore(&self) -> Result<()> {
-        self.call_utility::<(), _>("checkpoint_restore", ()).await?;
+        let recheck = self.inner.readiness().recheck_after_call();
+        self.inner
+            .call_utility::<(), _>("checkpoint_restore", (), Some(recheck))
+            .await?;
         Ok(())
     }
 
@@ -1081,6 +1184,7 @@ impl EngineCoreClient {
             output_task,
             dispatcher_task,
             abort_task,
+            readiness_task,
             coordinator_output_task,
             coordinator_task,
             ..
@@ -1092,7 +1196,7 @@ impl EngineCoreClient {
 
         // Abort all client tasks first, then await them.
         // Note the aborting orders here.
-        let mut tasks = vec![abort_task, dispatcher_task, output_task];
+        let mut tasks = vec![readiness_task, abort_task, dispatcher_task, output_task];
         tasks.extend(coordinator_task);
         tasks.extend(coordinator_output_task);
 
diff --git a/rust/src/engine-core-client/src/client/imp.rs b/rust/src/engine-core-client/src/client/imp.rs
index fa82d9591aa8..4f9edb584a7d 100644
--- a/rust/src/engine-core-client/src/client/imp.rs
+++ b/rust/src/engine-core-client/src/client/imp.rs
@@ -8,6 +8,7 @@ use std::time::{SystemTime, UNIX_EPOCH};
 
 use arc_swap::ArcSwapOption;
 use bytes::Bytes;
+use futures::future::join_all;
 use parking_lot::Mutex;
 use thiserror_ext::AsReport as _;
 use tokio::runtime::Handle;
@@ -16,18 +17,21 @@ use tracing::{debug, info, trace, warn};
 use vllm_metrics::METRICS;
 use zeromq::RouterSendHalf;
 
+use crate::client::readiness::{Readiness, RecheckGuard};
 use crate::client::state::{
     OutputReceiver, OutputSender, RequestRegistry, UtilityReceiver, UtilityRegistry,
 };
 use crate::client::stream::EngineCoreStreamOutput;
 use crate::client::{AbortCause, AbortRequest};
 use crate::error::{client_closed, dispatcher_closed, unexpected_dispatcher_output};
-use crate::metrics::{IterationMetricHandles, LoraInfoExporter, SchedulerStatsRecorder};
+use crate::metrics::{
+    EngineSleepState, IterationMetricHandles, LoraInfoExporter, SchedulerStatsRecorder,
+};
 use crate::protocol::encode_msgpack;
 use crate::protocol::output::{EngineCoreOutput, EngineCoreOutputs};
 use crate::protocol::request::{EngineCoreRequest, EngineCoreRequestType};
 use crate::protocol::stats::SchedulerStats;
-use crate::protocol::utility::UtilityOutput;
+use crate::protocol::utility::{EngineCoreUtilityRequest, UtilityOutput};
 use crate::transport::{ConnectedEngine, EngineId};
 use crate::{Error, Result, transport};
 
@@ -46,6 +50,9 @@ pub(crate) struct ClientInner {
     /// The runtime handle used for sending messages to the engine.
     handle: Handle,
     model_name: String,
+    /// Engines that utility calls go to, and this client's index on them.
+    engine_ids: Vec<EngineId>,
+    client_index: u32,
     /// Per-tensor byte threshold loaded from env variable
     /// `VLLM_MSGPACK_ZERO_COPY_THRESHOLD` when this inner client is created.
     msgpack_zero_copy_threshold: usize,
@@ -60,6 +67,9 @@ pub(crate) struct ClientInner {
     utility_reg: Mutex<UtilityRegistry>,
     health_error: ArcSwapOption<Error>,
     health_tx: watch::Sender<bool>,
+    /// Whether the engines can accept and process new requests, maintained
+    /// around the control calls this client sends.
+    readiness: Arc<Readiness>,
 }
 
 impl ClientInner {
@@ -69,6 +79,7 @@ impl ClientInner {
         input_send: RouterSendHalf,
         handle: Handle,
         model_name: String,
+        client_index: u32,
         engine_stats_enabled: bool,
         engines: &[ConnectedEngine],
     ) -> Self {
@@ -86,6 +97,8 @@ impl ClientInner {
             input_send,
             handle,
             model_name,
+            engine_ids: engines.iter().map(|engine| engine.engine_id.clone()).collect(),
+            client_index,
             msgpack_zero_copy_threshold: msgpack_zero_copy_threshold(),
             engine_stats_enabled,
             scheduler_stats_recorder,
@@ -94,6 +107,7 @@ impl ClientInner {
             utility_reg: Mutex::new(UtilityRegistry::default()),
             health_error: ArcSwapOption::empty(),
             health_tx: watch::Sender::new(true),
+            readiness: Arc::new(Readiness::new()),
         }
     }
 
@@ -103,6 +117,14 @@ impl ClientInner {
         &self.model_name
     }
 
+    /// Record a completed sleep or wake transition in
+    /// `vllm:engine_sleep_state`, exported only when engines record stats.
+    pub fn record_sleep_state(&self, state: EngineSleepState) {
+        if let Some(recorder) = &self.scheduler_stats_recorder {
+            recorder.record_sleep_state(state);
+        }
+    }
+
     /// Register a newly added request. Return the selected engine id and the
     /// per-request output channel bound to its `request_id`.
     ///
@@ -243,6 +265,95 @@ impl ClientInner {
         self.health_tx.subscribe()
     }
 
+    pub fn readiness(&self) -> &Arc<Readiness> {
+        &self.readiness
+    }
+
+    /// Call a typed utility method on every engine. See
+    /// [`crate::EngineCoreClient::call_utility`].
+    ///
+    /// A control call passes its `recheck` guard, which the call holds until
+    /// the engines reply.
+    pub async fn call_utility<T, A>(
+        self: &Arc<Self>,
+        method: &str,
+        args: A,
+        recheck: Option<RecheckGuard>,
+    ) -> Result<Vec<T>>
+    where
+        T: serde::de::DeserializeOwned,
+        A: serde::Serialize + std::fmt::Debug,
+    {
+        trace!(
+            method,
+            client_index = self.client_index,
+            engine_count = self.engine_ids.len(),
+            "sending utility request"
+        );
+
+        let mut call_ids = Vec::with_capacity(self.engine_ids.len());
+        let mut calls = Vec::with_capacity(self.engine_ids.len());
+        let prepared = self.engine_ids.iter().try_for_each(|engine_id| {
+            let (call_id, rx) = self.allocate_and_register_utility_call()?;
+            call_ids.push(call_id);
+            let request = EngineCoreUtilityRequest::new(self.client_index, call_id, method, &args)?;
+            let payload = Bytes::from(encode_msgpack(&request)?);
+            calls.push((engine_id.clone(), call_id, rx, payload));
+            Ok(())
+        });
+        if let Err(error) = prepared {
+            self.unregister_utility_calls(call_ids);
+            return Err(error);
+        }
+
+        // The engines execute a sent request whatever happens to the caller,
+        // so the call runs to completion in a task of its own on this client's
+        // runtime, and dropping this future only stops waiting for it.
+        let inner = self.clone();
+        let method_name = method.to_string();
+        let call = self.handle.spawn(async move {
+            let _recheck = recheck;
+            let outputs = join_all(calls.into_iter().map(|(engine_id, call_id, rx, payload)| {
+                let mut input_send = inner.input_send.clone();
+                let method_name = &method_name;
+                async move {
+                    transport::send_message(
+                        &mut input_send,
+                        &engine_id,
+                        EngineCoreRequestType::Utility.to_frame(),
+                        payload,
+                        Vec::new(),
+                    )
+                    .await?;
+                    rx.await.map_err(|_| Error::UtilityCallClosed {
+                        method: method_name.clone(),
+                        call_id,
+                    })?
+                }
+            }))
+            .await;
+            // A request that failed to send leaves its waiter without a reply.
+            inner.unregister_utility_calls(call_ids);
+            outputs
+        });
+        let outputs = match call.await {
+            Ok(outputs) => outputs,
+            Err(error) if error.is_panic() => std::panic::resume_unwind(error.into_panic()),
+            Err(_) => {
+                return Err(client_closed!(
+                    "engine-core client shut down during utility call `{method}`"
+                ));
+            }
+        };
+        outputs.into_iter().map(|output| output?.into_typed_result(method)).collect()
+    }
+
+    /// Ask every engine whether it can accept and process new requests.
+    pub async fn is_inference_ready(self: &Arc<Self>) -> Result<bool> {
+        let ready: Vec<bool> = self.call_utility("is_inference_ready", (), None).await?;
+        Ok(ready.into_iter().all(|ready| ready))
+    }
+
     /// Resolve one utility output to the waiting caller. Returns `true` if a
     /// waiting caller existed.
     pub fn resolve_utility_output(&self, output: UtilityOutput) -> bool {
@@ -552,6 +663,7 @@ mod tests {
             send,
             Handle::current(),
             "test-model".to_string(),
+            0,
             true,
             &[ConnectedEngine {
                 engine_id: EngineId::from(b"engine-0"),
diff --git a/rust/src/engine-core-client/src/client/readiness.rs b/rust/src/engine-core-client/src/client/readiness.rs
new file mode 100644
index 000000000000..c47d1d46adc9
--- /dev/null
+++ b/rust/src/engine-core-client/src/client/readiness.rs
@@ -0,0 +1,247 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+//! Whether the connected engines can accept and process new requests.
+//!
+//! The engines own this state and report it through the `is_inference_ready`
+//! utility. The client withdraws readiness while a call that stops serving is
+//! in flight, and rechecks after every control call ends, asking the engines
+//! again since only they know what the call left behind. Readiness therefore
+//! becomes `true` only from an engine answer.
+//!
+//! Withdrawing readiness publishes `false`, which the gRPC `readiness` and
+//! `vllm.Inference` health services report as `NOT_SERVING`, the same
+//! withdrawal the server publishes on shutdown.
+//!
+//! Asking runs an RPC on every worker, so the client asks only while someone
+//! subscribes; a new subscriber triggers an ask if the answer is unknown.
+
+use std::sync::Arc;
+use std::time::Duration;
+
+use parking_lot::Mutex;
+use thiserror_ext::AsReport as _;
+use tokio::sync::{Notify, watch};
+use tracing::warn;
+
+use crate::client::imp::ClientInner;
+
+/// How long to wait before asking the engines again after a failed attempt.
+const ASK_RETRY_DELAY: Duration = Duration::from_secs(1);
+
+#[derive(Debug)]
+struct ReadinessState {
+    /// Calls that stop serving and have not ended yet; readiness stays
+    /// withdrawn while any is in flight.
+    withdrawals: usize,
+    /// The engines' latest answer, or `None` from the end of a control call
+    /// until they answer again.
+    engines_ready: Option<bool>,
+    /// Advanced whenever a control call ends, so an answer that may predate
+    /// the call is discarded.
+    epoch: u64,
+}
+
+impl ReadinessState {
+    /// Readiness to publish, or `None` to keep the published value until the
+    /// engines answer.
+    fn ready(&self) -> Option<bool> {
+        if self.withdrawals > 0 {
+            Some(false)
+        } else {
+            self.engines_ready
+        }
+    }
+}
+
+/// Whether the connected engines can accept and process new requests, as
+/// published to subscribers. Control calls report through [`RecheckGuard`]
+/// guards, and [`run_readiness_loop`] records the engines' answers.
+pub(crate) struct Readiness {
+    state: Mutex<ReadinessState>,
+    tx: watch::Sender<bool>,
+    /// Wakes the ask loop when an answer may be needed.
+    ask: Notify,
+}
+
+impl Readiness {
+    /// Start ready: engines without sleep mode cannot start paused, and
+    /// [`Readiness::start`] records the engines' answer for those with it.
+    pub fn new() -> Self {
+        Self {
+            state: Mutex::new(ReadinessState {
+                withdrawals: 0,
+                engines_ready: Some(true),
+                epoch: 0,
+            }),
+            tx: watch::Sender::new(true),
+            ask: Notify::new(),
+        }
+    }
+
+    /// Subscribe to the published readiness, asking the engines if their
+    /// answer is unknown.
+    pub fn subscribe(&self) -> watch::Receiver<bool> {
+        let rx = self.tx.subscribe();
+        self.ask.notify_one();
+        rx
+    }
+
+    /// For a control call that may change readiness, such as a resume: leave
+    /// readiness unchanged during the call and recheck after it ends. Pass the
+    /// guard to the call's [`ClientInner::call_utility`].
+    pub fn recheck_after_call(self: &Arc<Self>) -> RecheckGuard {
+        RecheckGuard {
+            readiness: self.clone(),
+            withdraws: false,
+        }
+    }
+
+    /// For a control call that stops serving, such as a pause: withdraw
+    /// readiness from before the call is sent until the recheck after it ends.
+    /// Pass the guard to the call's [`ClientInner::call_utility`].
+    pub fn withdraw_until_recheck(self: &Arc<Self>) -> RecheckGuard {
+        let mut state = self.state.lock();
+        state.withdrawals += 1;
+        self.publish(&state);
+        RecheckGuard {
+            readiness: self.clone(),
+            withdraws: true,
+        }
+    }
+
+    /// Record the engines' answer at startup, before any control call.
+    pub fn start(&self, ready: bool) {
+        let mut state = self.state.lock();
+        state.engines_ready = Some(ready);
+        self.publish(&state);
+    }
+
+    /// Record the engines' answer to an ask made at `epoch`, unless a control
+    /// call ended since.
+    fn answer(&self, epoch: u64, ready: bool) {
+        let mut state = self.state.lock();
+        if state.epoch == epoch {
+            state.engines_ready = Some(ready);
+            self.publish(&state);
+        }
+    }
+
+    /// Return the epoch to ask at while the engines' answer is unknown and
+    /// someone subscribes.
+    pub(super) fn pending_ask(&self) -> Option<u64> {
+        let state = self.state.lock();
+        (state.engines_ready.is_none() && self.tx.receiver_count() > 0).then_some(state.epoch)
+    }
+
+    fn end(&self, withdraws: bool) {
+        let mut state = self.state.lock();
+        if withdraws {
+            state.withdrawals -= 1;
+        }
+        state.epoch += 1;
+        state.engines_ready = None;
+        self.ask.notify_one();
+        self.publish(&state);
+    }
+
+    fn publish(&self, state: &ReadinessState) {
+        if let Some(ready) = state.ready() {
+            self.tx.send_if_modified(|current| std::mem::replace(current, ready) != ready);
+        }
+    }
+}
+
+/// A control call in flight. Dropping it, whether the call succeeded or failed,
+/// rechecks: the client asks the engines again. [`ClientInner::call_utility`]
+/// holds it until the engines reply, even if the caller is cancelled.
+pub(crate) struct RecheckGuard {
+    readiness: Arc<Readiness>,
+    /// Whether readiness is withdrawn until the recheck.
+    withdraws: bool,
+}
+
+impl Drop for RecheckGuard {
+    fn drop(&mut self) {
+        self.readiness.end(self.withdraws);
+    }
+}
+
+/// Background loop that asks the engines whenever their answer is unknown,
+/// until an answer arrives with no control call ending in between.
+pub(crate) async fn run_readiness_loop(inner: Arc<ClientInner>) {
+    let readiness = inner.readiness();
+    loop {
+        let Some(epoch) = readiness.pending_ask() else {
+            readiness.ask.notified().await;
+            continue;
+        };
+        match inner.is_inference_ready().await {
+            Ok(ready) => readiness.answer(epoch, ready),
+            // Engine failure is permanent and fails every health service.
+            Err(_) if !inner.is_healthy() => return,
+            Err(error) => {
+                warn!(
+                    error = %error.as_report(),
+                    "failed to ask the engines whether they are ready; keeping the published readiness"
+                );
+                tokio::time::sleep(ASK_RETRY_DELAY).await;
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn answer_that_predates_a_control_call_is_discarded() {
+        let readiness = Arc::new(Readiness::new());
+        let published = readiness.subscribe();
+
+        drop(readiness.recheck_after_call()); // resume
+        let before_pause = readiness.pending_ask().expect("ask after the resume");
+        drop(readiness.withdraw_until_recheck()); // pause
+        let after_pause = readiness.pending_ask().expect("ask after the pause");
+
+        // Answered before the pause reached the engines.
+        readiness.answer(before_pause, true);
+        assert!(!*published.borrow());
+        assert_eq!(readiness.pending_ask(), Some(after_pause));
+
+        readiness.answer(after_pause, false);
+        assert!(!*published.borrow());
+        assert_eq!(readiness.pending_ask(), None);
+    }
+
+    #[test]
+    fn withdrawal_in_flight_hides_engine_answers() {
+        let readiness = Arc::new(Readiness::new());
+        let published = readiness.subscribe();
+
+        let pause = readiness.withdraw_until_recheck();
+        assert!(!*published.borrow());
+        // A resume that overtakes the pause.
+        drop(readiness.recheck_after_call());
+        let epoch = readiness.pending_ask().unwrap();
+        readiness.answer(epoch, true);
+        assert!(!*published.borrow());
+
+        drop(pause);
+        assert!(!*published.borrow());
+        let epoch = readiness.pending_ask().unwrap();
+        readiness.answer(epoch, true);
+        assert!(*published.borrow());
+    }
+
+    #[test]
+    fn asks_only_while_subscribed() {
+        let readiness = Arc::new(Readiness::new());
+        drop(readiness.recheck_after_call());
+        assert_eq!(readiness.pending_ask(), None);
+
+        let _published = readiness.subscribe();
+        assert!(readiness.pending_ask().is_some());
+    }
+}
diff --git a/rust/src/engine-core-client/src/coordinator/bootstrap.rs b/rust/src/engine-core-client/src/coordinator/bootstrap.rs
index 17c346c933e3..b4bce935c1f0 100644
--- a/rust/src/engine-core-client/src/coordinator/bootstrap.rs
+++ b/rust/src/engine-core-client/src/coordinator/bootstrap.rs
@@ -9,16 +9,21 @@ use zeromq::{PullSocket, XPubSocket, ZmqMessage};
 
 use crate::error::{Error, Result, bail_unexpected_handshake_message};
 
-/// Engine-facing sockets owned by the in-process coordinator.
+/// Sockets owned by the in-process coordinator: the engine-facing broadcast
+/// and control sockets, and the frontend-facing state publish socket.
 pub(crate) struct CoordinatorBootstrap {
     pub input_address: String,
     pub output_address: String,
+    /// Sent to engines in INIT, which report it back to every frontend as
+    /// READY `dp_stats_address`.
+    pub frontend_publish_address: String,
     pub input_socket: XPubSocket,
     pub output_socket: PullSocket,
+    pub frontend_publish_socket: XPubSocket,
 }
 
 impl CoordinatorBootstrap {
-    /// Bind the engine-facing coordinator sockets on the given host.
+    /// Bind the coordinator sockets on the given host.
     pub(crate) async fn bind(local_host: &str) -> Result<Self> {
         let mut input_socket = XPubSocket::new();
         let input_address = input_socket.bind(&format!("tcp://{local_host}:0")).await?.to_string();
@@ -27,11 +32,19 @@ impl CoordinatorBootstrap {
         let output_address =
             output_socket.bind(&format!("tcp://{local_host}:0")).await?.to_string();
 
+        let mut frontend_publish_socket = XPubSocket::new();
+        let frontend_publish_address = frontend_publish_socket
+            .bind(&format!("tcp://{local_host}:0"))
+            .await?
+            .to_string();
+
         Ok(Self {
             input_address,
             output_address,
+            frontend_publish_address,
             input_socket,
             output_socket,
+            frontend_publish_socket,
         })
     }
 
@@ -63,7 +76,7 @@ async fn wait_for_engine_subscriptions(
                     timeout: ready_timeout,
                 }
             })??;
-        if !is_engine_subscription(&message) {
+        if !is_subscribe_all(&message) {
             bail_unexpected_handshake_message!(
                 "expected coordinator subscription [0x01], got {message:?}"
             );
@@ -75,9 +88,9 @@ async fn wait_for_engine_subscriptions(
 }
 
 /// Whether `message` is the single-frame subscribe-all notification (`[0x01]`)
-/// that an engine's XSUB sends whenever it connects or reconnects to the
-/// coordinator broadcast socket.
-pub(super) fn is_engine_subscription(message: &ZmqMessage) -> bool {
+/// that an XSUB peer sends whenever it connects or reconnects to a coordinator
+/// XPUB socket.
+pub(super) fn is_subscribe_all(message: &ZmqMessage) -> bool {
     message.len() == 1 && message.get(0).is_some_and(|frame| frame.as_ref() == [0x01])
 }
 
diff --git a/rust/src/engine-core-client/src/coordinator/external.rs b/rust/src/engine-core-client/src/coordinator/external.rs
index d0418352a6d8..ae49acc1503a 100644
--- a/rust/src/engine-core-client/src/coordinator/external.rs
+++ b/rust/src/engine-core-client/src/coordinator/external.rs
@@ -3,40 +3,18 @@
 
 use std::sync::Arc;
 
-use serde_tuple::Deserialize_tuple;
 use thiserror_ext::AsReport;
 use tracing::{debug, warn};
 use zeromq::prelude::SocketRecv;
 use zeromq::{XSubSocket, ZmqMessage};
 
 use crate::client::imp::ClientInner;
-use crate::coordinator::handle::CoordinatorState;
+use crate::coordinator::handle::{CoordinatorState, CoordinatorStateUpdate};
 use crate::error::{Error, Result, bail_unexpected_coordinator_output};
-use crate::protocol::{OpaqueValue, decode_msgpack};
+use crate::protocol::decode_msgpack;
 
-/// Coordinator-to-frontend state publish received on the front-side coordinator
-/// socket.
-///
-/// This matches the msgpack tuple periodically published by Python
-/// `DPCoordinatorProc.run_coordinator` to all connected frontends.
-///
-/// Original Python definitions:
-/// <https://github.com/vllm-project/vllm/blob/694449050f8dac3d9853e97e518b4a43ec52106a/vllm/v1/engine/coordinator.py#L282-L283>
-/// <https://github.com/vllm-project/vllm/blob/694449050f8dac3d9853e97e518b4a43ec52106a/vllm/v1/engine/coordinator.py#L445-L447>
-#[derive(Debug, Clone, PartialEq, Deserialize_tuple)]
-struct CoordinatorStateUpdate {
-    /// Global per-engine request counts published by the coordinator.
-    ///
-    /// The Rust bootstrapped external-coordinator path preserves this field for
-    /// wire compatibility but intentionally ignores it for routing decisions.
-    counts: OpaqueValue,
-    /// Current global DP wave number stamped onto newly admitted requests.
-    wave: u32,
-    /// Whether engines are currently running (`true`) or paused (`false`).
-    engines_running: bool,
-}
-
-/// Background half of an external Python-owned coordinator connection.
+/// Background half of a connection to a coordinator in another process: a
+/// Python coordinator, or the in-process coordinator of another Rust frontend.
 ///
 /// This owns one frontend-facing XSUB socket. It mirrors the subset of Python's
 /// coordinator protocol needed by the Rust bootstrapped frontend: receive
diff --git a/rust/src/engine-core-client/src/coordinator/handle.rs b/rust/src/engine-core-client/src/coordinator/handle.rs
index 9ebc1717e30c..2a9cc0f75a7b 100644
--- a/rust/src/engine-core-client/src/coordinator/handle.rs
+++ b/rust/src/engine-core-client/src/coordinator/handle.rs
@@ -4,12 +4,42 @@
 use std::sync::Arc;
 
 use parking_lot::Mutex;
+use serde_tuple::{Deserialize_tuple, Serialize_tuple};
 use zeromq::prelude::Socket;
 use zeromq::{XPubSocket, XSubSocket};
 
 use crate::coordinator::external::ExternalCoordinatorService;
 use crate::coordinator::inproc::InProcCoordinatorRunner;
 use crate::error::Result;
+use crate::protocol::OpaqueValue;
+
+/// Coordinator-to-frontend state publish received on the front-side coordinator
+/// socket.
+///
+/// This matches the msgpack tuple periodically published by Python
+/// `DPCoordinatorProc.run_coordinator` to all connected frontends.
+///
+/// Original Python definitions:
+/// <https://github.com/vllm-project/vllm/blob/694449050f8dac3d9853e97e518b4a43ec52106a/vllm/v1/engine/coordinator.py#L282-L283>
+/// <https://github.com/vllm-project/vllm/blob/694449050f8dac3d9853e97e518b4a43ec52106a/vllm/v1/engine/coordinator.py#L445-L447>
+///
+/// The in-process coordinator publishes the same tuple to frontends of other
+/// DP ranks.
+#[derive(Debug, Clone, PartialEq, Serialize_tuple, Deserialize_tuple)]
+pub(super) struct CoordinatorStateUpdate {
+    /// Global per-engine request counts published by the coordinator.
+    ///
+    /// The Rust bootstrapped external-coordinator path preserves this field for
+    /// wire compatibility but intentionally ignores it for routing decisions.
+    ///
+    /// `None` marks a wave-only update. The in-process coordinator does not
+    /// track request counts, so it always publishes `None`.
+    pub counts: Option<OpaqueValue>,
+    /// Current global DP wave number stamped onto newly admitted requests.
+    pub wave: u32,
+    /// Whether engines are currently running (`true`) or paused (`false`).
+    pub engines_running: bool,
+}
 
 /// Snapshot to the coordinator state for request routing and stamping.
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -47,18 +77,22 @@ impl CoordinatorHandle {
         )
     }
 
-    /// Build the paired frontend handle and background runner around one
-    /// engine-facing coordinator broadcast socket.
-    pub(crate) fn new_inproc(coordinator_input: XPubSocket) -> (Self, InProcCoordinatorRunner) {
+    /// Build the paired frontend handle and background runner around the
+    /// engine-facing coordinator broadcast socket and the frontend-facing state
+    /// publish socket.
+    pub(crate) fn new_inproc(
+        coordinator_input: XPubSocket,
+        frontend_publish: XPubSocket,
+    ) -> (Self, InProcCoordinatorRunner) {
         let (handle, state) = Self::new_parts();
         (
             handle,
-            InProcCoordinatorRunner::new(state, coordinator_input),
+            InProcCoordinatorRunner::new(state, coordinator_input, frontend_publish),
         )
     }
 
-    /// Build the paired frontend handle and background service around an
-    /// external Python-owned frontend-side coordinator socket.
+    /// Build the paired frontend handle and background service around the
+    /// frontend-side socket of a coordinator in another process.
     pub(crate) async fn connect_external(
         coordinator_address: &str,
     ) -> Result<(Self, ExternalCoordinatorService)> {
diff --git a/rust/src/engine-core-client/src/coordinator/inproc.rs b/rust/src/engine-core-client/src/coordinator/inproc.rs
index d780b64e99fc..65b3f4885f48 100644
--- a/rust/src/engine-core-client/src/coordinator/inproc.rs
+++ b/rust/src/engine-core-client/src/coordinator/inproc.rs
@@ -5,14 +5,14 @@ use std::sync::Arc;
 
 use serde_tuple::{Deserialize_tuple, Serialize_tuple};
 use thiserror_ext::AsReport;
-use tokio::sync::mpsc;
+use tokio::sync::{Notify, mpsc};
 use tracing::{debug, warn};
 use zeromq::prelude::{SocketRecv, SocketSend};
 use zeromq::{XPubSocket, ZmqError, ZmqMessage};
 
 use crate::client::imp::ClientInner;
-use crate::coordinator::bootstrap::is_engine_subscription;
-use crate::coordinator::handle::CoordinatorState;
+use crate::coordinator::bootstrap::is_subscribe_all;
+use crate::coordinator::handle::{CoordinatorState, CoordinatorStateUpdate};
 use crate::error::{Error, Result, bail_unexpected_coordinator_output};
 use crate::protocol::encode_msgpack;
 use crate::protocol::output::{DpControlMessage, DpControlOutput, EngineCoreOutputs};
@@ -43,13 +43,83 @@ struct StartDpWaveMessage {
 pub(crate) struct InProcCoordinatorRunner {
     state: Arc<CoordinatorState>,
     coordinator_input: XPubSocket,
+    /// Wakes the frontend publisher after each wave transition.
+    state_changed: Arc<Notify>,
+    /// Moved out by `run` to run beside the engine-facing loop.
+    frontend_publisher: Option<FrontendStatePublisher>,
+}
+
+/// Frontend-facing half of the in-process coordinator, through which frontends
+/// of other DP ranks follow the wave state that this process reads in memory.
+///
+/// It runs beside the engine-facing loop, so a frontend that stops reading can
+/// delay frontend updates but never engine wave control. Updates coalesce to
+/// the latest state, which is all a frontend reads.
+struct FrontendStatePublisher {
+    state: Arc<CoordinatorState>,
+    socket: XPubSocket,
+    state_changed: Arc<Notify>,
+}
+
+impl FrontendStatePublisher {
+    /// Publish the current wave state to subscribed frontends.
+    async fn publish_state(&mut self) -> Result<()> {
+        let state = *self.state.lock();
+        let payload = encode_msgpack(&CoordinatorStateUpdate {
+            counts: None,
+            wave: state.current_wave,
+            engines_running: state.engines_running,
+        })?;
+        self.socket.send(ZmqMessage::from(payload)).await?;
+        Ok(())
+    }
+
+    /// Publish on every wave transition and to every new subscriber until a
+    /// fatal socket error.
+    async fn run(mut self) -> Result<()> {
+        loop {
+            tokio::select! {
+                () = self.state_changed.notified() => self.publish_state().await?,
+                // As on the engine-facing socket, zmq.rs registers a subscription only
+                // when `recv()` yields it. Answering it with the current state covers a
+                // frontend that subscribes after the last wave transition.
+                subscription = self.socket.recv() => {
+                    match subscription {
+                        Ok(message) => {
+                            if is_subscribe_all(&message) {
+                                self.publish_state().await?;
+                            }
+                        }
+                        Err(error @ (ZmqError::Codec(_) | ZmqError::Network(_))) => {
+                            warn!(
+                                error = %error,
+                                "coordinator frontend subscription peer disconnected"
+                            );
+                        }
+                        Err(error) => return Err(error.into()),
+                    }
+                }
+            }
+        }
+    }
 }
 
 impl InProcCoordinatorRunner {
-    pub(super) fn new(state: Arc<CoordinatorState>, coordinator_input: XPubSocket) -> Self {
+    pub(super) fn new(
+        state: Arc<CoordinatorState>,
+        coordinator_input: XPubSocket,
+        frontend_publish: XPubSocket,
+    ) -> Self {
+        let state_changed = Arc::new(Notify::new());
         Self {
+            frontend_publisher: Some(FrontendStatePublisher {
+                state: state.clone(),
+                socket: frontend_publish,
+                state_changed: state_changed.clone(),
+            }),
             state,
             coordinator_input,
+            state_changed,
         }
     }
 
@@ -94,16 +164,24 @@ impl InProcCoordinatorRunner {
                 // The engines signals they completed the current wave and are now paused.
                 // Advance the current wave and mark the state as paused.
                 DpControlMessage::WaveComplete(wave) => {
-                    let mut state = self.state.lock();
-                    if wave >= state.current_wave {
-                        let next_wave = wave + 1;
-                        debug!(
-                            wave,
-                            next_wave,
-                            "DP wave finished; pausing engines and advancing coordinator state"
-                        );
-                        state.current_wave = wave + 1;
-                        state.engines_running = false;
+                    let advanced = {
+                        let mut state = self.state.lock();
+                        if wave >= state.current_wave {
+                            let next_wave = wave + 1;
+                            debug!(
+                                wave,
+                                next_wave,
+                                "DP wave finished; pausing engines and advancing coordinator state"
+                            );
+                            state.current_wave = wave + 1;
+                            state.engines_running = false;
+                            true
+                        } else {
+                            false
+                        }
+                    };
+                    if advanced {
+                        self.state_changed.notify_one();
                     }
                 }
                 // An engine requests to start the wave.
@@ -128,6 +206,7 @@ impl InProcCoordinatorRunner {
                             "starting DP wave on notification from engine"
                         );
                         self.broadcast_start_wave(wave, Some(engine_index)).await?;
+                        self.state_changed.notify_one();
                     }
                 }
             },
@@ -150,7 +229,9 @@ impl InProcCoordinatorRunner {
         mut output_rx: mpsc::Receiver<Result<EngineCoreOutputs>>,
         inner: Arc<ClientInner>,
     ) {
-        let result: Result<()> = async {
+        let frontend_publisher =
+            self.frontend_publisher.take().expect("the coordinator runner runs once");
+        let engine_loop = async {
             loop {
                 tokio::select! {
                     // Keep receiving engine subscriptions: zmq.rs XPUB registers a
@@ -161,7 +242,7 @@ impl InProcCoordinatorRunner {
                     subscription = self.coordinator_input.recv() => {
                         match subscription {
                             Ok(message) => {
-                                if is_engine_subscription(&message) {
+                                if is_subscribe_all(&message) {
                                     let state = *self.state.lock();
                                     if state.engines_running {
                                         debug!(
@@ -191,8 +272,11 @@ impl InProcCoordinatorRunner {
                     }
                 }
             }
-        }
-        .await;
+        };
+        let result: Result<()> = tokio::select! {
+            result = engine_loop => result,
+            result = frontend_publisher.run() => result,
+        };
         let Err(error) = result else { return };
 
         warn!(error = %error.as_report(), "coordinator runner exiting with error");
diff --git a/rust/src/engine-core-client/src/lib.rs b/rust/src/engine-core-client/src/lib.rs
index 8f79e227de11..3fc08abebf6c 100644
--- a/rust/src/engine-core-client/src/lib.rs
+++ b/rust/src/engine-core-client/src/lib.rs
@@ -6,6 +6,7 @@ mod coordinator;
 mod error;
 mod metrics;
 pub mod mock_engine;
+mod prompt_block_hasher;
 pub mod protocol;
 pub mod runtime;
 #[cfg(any(test, feature = "test-util"))]
diff --git a/rust/src/engine-core-client/src/metrics.rs b/rust/src/engine-core-client/src/metrics.rs
index dd2835972fdc..66f8e532131f 100644
--- a/rust/src/engine-core-client/src/metrics.rs
+++ b/rust/src/engine-core-client/src/metrics.rs
@@ -6,21 +6,50 @@ use std::collections::BTreeSet;
 use std::time::{SystemTime, UNIX_EPOCH};
 
 use vllm_metrics::{
-    CacheHitSourceLabels, EngineLabels, EnginePositionLabels, F64Gauge, Family, HistogramMetric,
-    LoraAdapterNames, LoraInfoLabels, MooncakeOperationCounterFamily,
-    MooncakeOperationHistogramFamily, MooncakeOperationLabels, RequestMetrics,
-    SchedulerLogStatsAccumulator, SchedulerMetrics, U64Counter, U64Gauge, WaitingReasonLabels,
+    CacheHitSourceLabels, EngineLabels, EnginePositionLabels, EngineSleepStateLabels, F64Gauge,
+    Family, HistogramMetric, KvFetchStageLabels, LoraAdapterNames, LoraInfoLabels,
+    MooncakeOperationCounterFamily, MooncakeOperationHistogramFamily, MooncakeOperationLabels,
+    RequestMetrics, SchedulerLogStatsAccumulator, SchedulerMetrics, U64Counter, U64Gauge,
+    UmbpOperationCounterFamily, UmbpOperationLabels, WaitingReasonLabels,
 };
 
 use crate::protocol::stats::{
-    CachedTokensBySource, KvConnectorStats, MooncakeStats, MooncakeTransferStats,
-    MultiConnectorStats, NixlStats, SchedulerStats,
+    CachedTokensBySource, KvConnectorStats, KvFetchReqsByStage, MooncakeStats,
+    MooncakeTransferStats, MultiConnectorStats, NixlStats, SchedulerStats, UmbpStats,
 };
+use crate::protocol::utility::SleepLevel;
 use crate::transport::ConnectedEngine;
 
 const WAITING_REASON_CAPACITY: &str = "capacity";
 const WAITING_REASON_DEFERRED: &str = "deferred";
 
+/// `sleep_state` label values of `vllm:engine_sleep_state`, in the order of
+/// [`EngineSleepState::gauge_values`].
+const SLEEP_STATES: [&str; 3] = ["awake", "weights_offloaded", "discard_all"];
+
+/// Engine sleep state exported as `vllm:engine_sleep_state`.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub(crate) enum EngineSleepState {
+    Awake,
+    Asleep(SleepLevel),
+    /// After `release_kv_cache_memory`, which Python records like a level-0
+    /// sleep.
+    KvCacheReleased,
+}
+
+impl EngineSleepState {
+    /// Gauge values for the `awake`, `weights_offloaded` and `discard_all`
+    /// series, as Python's `PrometheusStatLogger.record_sleep_state` sets them.
+    fn gauge_values(self) -> [u64; 3] {
+        match self {
+            Self::Awake => [1, 0, 0],
+            Self::Asleep(SleepLevel::PauseOnly) | Self::KvCacheReleased => [0, 0, 0],
+            Self::Asleep(SleepLevel::OffloadWeights) => [0, 1, 0],
+            Self::Asleep(SleepLevel::DiscardAll) => [0, 0, 1],
+        }
+    }
+}
+
 /// Cached output-batch metric handles for one model and engine index.
 pub(crate) struct IterationMetricHandles {
     pub iteration_tokens_total: HistogramMetric,
@@ -54,7 +83,11 @@ struct SchedulerStatsHandles {
     scheduler_waiting: U64Gauge,
     scheduler_waiting_capacity: U64Gauge,
     scheduler_waiting_deferred: U64Gauge,
+    /// One handle per `KvFetchReqsByStage::STAGES` entry.
+    scheduler_kv_fetch_by_stage: [U64Gauge; 3],
     kv_cache_usage: F64Gauge,
+    /// One handle per [`SLEEP_STATES`] entry.
+    engine_sleep_state: [U64Gauge; 3],
 
     // Prefix-cache counters, including the connector-backed external cache path.
     prefix_cache_queries: U64Counter,
@@ -97,6 +130,14 @@ struct SchedulerStatsHandles {
     mooncake_num_failed_recvs: U64Counter,
     mooncake_num_kv_expired_reqs: U64Counter,
 
+    // UMBP store connector telemetry, decoded from `kv_connector_stats`.
+    // Kept as `Family` like the Mooncake store telemetry: an operation's
+    // series appears once the connector first reports it, as in Python.
+    umbp_transfer_submitted_total: UmbpOperationCounterFamily,
+    umbp_transfer_completed_total: UmbpOperationCounterFamily,
+    umbp_transfer_failed_total: UmbpOperationCounterFamily,
+    umbp_transfer_bytes_total: UmbpOperationCounterFamily,
+
     // NIXL connector telemetry, decoded from `kv_connector_stats`.
     nixl_xfer_time_seconds: HistogramMetric,
     nixl_post_time_seconds: HistogramMetric,
@@ -129,7 +170,21 @@ impl SchedulerStatsRecorder {
             })
             .collect();
 
-        Self { engines }
+        let recorder = Self { engines };
+        // Engines start awake, as Python sets when it creates the stat logger.
+        recorder.record_sleep_state(EngineSleepState::Awake);
+        recorder
+    }
+
+    /// Record the sleep state of every connected engine. Sleep and wake are
+    /// utility calls on all engines, so the state is not per engine.
+    pub(crate) fn record_sleep_state(&self, state: EngineSleepState) {
+        let values = state.gauge_values();
+        for handles in self.engines.values() {
+            for (gauge, value) in handles.engine_sleep_state.iter().zip(values) {
+                gauge.set(value);
+            }
+        }
     }
 
     /// Record one scheduler-stats payload for the given engine index.
@@ -170,7 +225,21 @@ fn resolve_scheduler_stats_handles(
         scheduler_waiting_deferred: metrics
             .scheduler_waiting_by_reason
             .get_or_create_owned(&deferred),
+        scheduler_kv_fetch_by_stage: KvFetchReqsByStage::STAGES.map(|stage| {
+            metrics.scheduler_kv_fetch_by_stage.get_or_create_owned(&KvFetchStageLabels {
+                model_name: model_name.to_string(),
+                engine,
+                stage,
+            })
+        }),
         kv_cache_usage: metrics.kv_cache_usage.get_or_create_owned(&labels),
+        engine_sleep_state: SLEEP_STATES.map(|sleep_state| {
+            metrics.engine_sleep_state.get_or_create_owned(&EngineSleepStateLabels {
+                model_name: model_name.to_string(),
+                engine,
+                sleep_state,
+            })
+        }),
         prefix_cache_queries: metrics.prefix_cache_queries.get_or_create_owned(&labels),
         prefix_cache_hits: metrics.prefix_cache_hits.get_or_create_owned(&labels),
         external_prefix_cache_queries: metrics
@@ -224,6 +293,10 @@ fn resolve_scheduler_stats_handles(
         mooncake_num_kv_expired_reqs: metrics
             .mooncake_num_kv_expired_reqs
             .get_or_create_owned(&labels),
+        umbp_transfer_submitted_total: metrics.umbp_transfer_submitted_total.clone(),
+        umbp_transfer_completed_total: metrics.umbp_transfer_completed_total.clone(),
+        umbp_transfer_failed_total: metrics.umbp_transfer_failed_total.clone(),
+        umbp_transfer_bytes_total: metrics.umbp_transfer_bytes_total.clone(),
         nixl_xfer_time_seconds: metrics.nixl_xfer_time_seconds.get_or_create_owned(&labels),
         nixl_post_time_seconds: metrics.nixl_post_time_seconds.get_or_create_owned(&labels),
         nixl_bytes_transferred: metrics.nixl_bytes_transferred.get_or_create_owned(&labels),
@@ -249,6 +322,13 @@ fn record_scheduler_stats_with_handles(handles: &SchedulerStatsHandles, stats: &
         .set(stats.num_waiting_reqs + stats.num_skipped_waiting_reqs);
     handles.scheduler_waiting_capacity.set(stats.num_waiting_reqs);
     handles.scheduler_waiting_deferred.set(stats.num_skipped_waiting_reqs);
+    for (handle, count) in handles
+        .scheduler_kv_fetch_by_stage
+        .iter()
+        .zip(stats.num_kv_fetch_reqs_by_stage.counts())
+    {
+        handle.set(count);
+    }
     handles.kv_cache_usage.set(stats.kv_cache_usage);
 
     // Prefix-cache counters, including the connector-backed external cache path.
@@ -339,6 +419,7 @@ fn record_scheduler_stats_with_handles(handles: &SchedulerStatsHandles, stats: &
                 record_mooncake_transfer_stats(handles, stats)
             }
             KvConnectorStats::Mooncake(stats) => record_mooncake_stats(handles, stats),
+            KvConnectorStats::Umbp(stats) => record_umbp_stats(handles, stats),
             KvConnectorStats::Multi(stats) => record_multi_connector_stats(handles, stats),
             KvConnectorStats::Other(_) => {}
         }
@@ -355,6 +436,9 @@ fn record_multi_connector_stats(handles: &SchedulerStatsHandles, stats: &MultiCo
     if let Some(mooncake) = &stats.mooncake {
         record_mooncake_stats(handles, mooncake);
     }
+    if let Some(umbp) = &stats.umbp {
+        record_umbp_stats(handles, umbp);
+    }
 }
 
 fn record_mooncake_transfer_stats(handles: &SchedulerStatsHandles, stats: &MooncakeTransferStats) {
@@ -406,6 +490,29 @@ fn record_mooncake_stats(handles: &SchedulerStatsHandles, stats: &MooncakeStats)
     }
 }
 
+fn record_umbp_stats(handles: &SchedulerStatsHandles, stats: &UmbpStats) {
+    for (operation, counts) in &stats.0 {
+        let labels = UmbpOperationLabels {
+            model_name: handles.labels.model_name.clone(),
+            engine: handles.labels.engine,
+            operation: operation.as_str(),
+        };
+        handles
+            .umbp_transfer_submitted_total
+            .get_or_create(&labels)
+            .inc_by(counts.submitted);
+        handles
+            .umbp_transfer_completed_total
+            .get_or_create(&labels)
+            .inc_by(counts.completed);
+        handles.umbp_transfer_failed_total.get_or_create(&labels).inc_by(counts.failed);
+        handles
+            .umbp_transfer_bytes_total
+            .get_or_create(&labels)
+            .inc_by(counts.num_bytes);
+    }
+}
+
 fn record_nixl_stats(handles: &SchedulerStatsHandles, stats: &NixlStats) {
     for value in &stats.transfer_duration {
         handles.nixl_xfer_time_seconds.observe(*value);
@@ -487,12 +594,13 @@ mod tests {
     use std::collections::{BTreeMap, BTreeSet};
 
     use expect_test::expect;
+    use itertools::Itertools as _;
     use vllm_metrics::Metrics;
 
     use crate::metrics::LoraInfoExporter;
     use crate::protocol::stats::{
-        KvConnectorStats, MooncakeOperation, MooncakeRecord, MooncakeStats, MooncakeStatus,
-        MooncakeTransferStats, MultiConnectorStats, NixlStats, SchedulerStats,
+        KvConnectorStats, KvFetchReqsByStage, MooncakeOperation, MooncakeRecord, MooncakeStats,
+        MooncakeStatus, MooncakeTransferStats, MultiConnectorStats, NixlStats, SchedulerStats,
     };
 
     fn names(values: &[&str]) -> BTreeSet<String> {
@@ -622,6 +730,91 @@ mod tests {
         .assert_eq(&mooncake_transfer_series);
     }
 
+    /// Records a bare `UMBPStoreConnector` payload, in the shape Python's
+    /// `UMBPStoreConnectorStats` sends, into the `vllm:umbp_transfer_*`
+    /// counters, one series per reported operation.
+    #[test]
+    fn umbp_stats_are_recorded_into_umbp_transfer_metrics() {
+        let metrics = Metrics::new();
+        let handles = super::resolve_scheduler_stats_handles(&metrics.scheduler, "model", 0);
+        // Series of one family render in hash order, so sort them.
+        let umbp_series = || {
+            metrics
+                .render()
+                .unwrap()
+                .lines()
+                .filter(|line| line.starts_with("vllm:umbp_transfer_"))
+                .sorted()
+                .join("\n")
+        };
+
+        // Like Python, no series before the connector reports an operation.
+        expect![[""]].assert_eq(&umbp_series());
+
+        let kv_connector_stats: KvConnectorStats = serde_json::from_value(serde_json::json!({
+            "load": {"submitted": 3, "completed": 2, "failed": 1, "num_bytes": 8192},
+            "store": {"submitted": 1, "completed": 0, "failed": 0, "num_bytes": 0},
+        }))
+        .unwrap();
+        let stats = SchedulerStats {
+            kv_connector_stats: Some(kv_connector_stats),
+            ..Default::default()
+        };
+        super::record_scheduler_stats_with_handles(&handles, &stats);
+        super::record_scheduler_stats_with_handles(&handles, &stats);
+
+        expect![[r#"
+            vllm:umbp_transfer_bytes_total{model_name="model",engine="0",operation="load"} 16384
+            vllm:umbp_transfer_bytes_total{model_name="model",engine="0",operation="store"} 0
+            vllm:umbp_transfer_completed_total{model_name="model",engine="0",operation="load"} 4
+            vllm:umbp_transfer_completed_total{model_name="model",engine="0",operation="store"} 0
+            vllm:umbp_transfer_failed_total{model_name="model",engine="0",operation="load"} 2
+            vllm:umbp_transfer_failed_total{model_name="model",engine="0",operation="store"} 0
+            vllm:umbp_transfer_submitted_total{model_name="model",engine="0",operation="load"} 6
+            vllm:umbp_transfer_submitted_total{model_name="model",engine="0",operation="store"} 2"#]]
+        .assert_eq(&umbp_series());
+    }
+
+    /// Exports every engine's KV fetch stage gauges from startup and sets them
+    /// from each `SchedulerStats` update.
+    #[test]
+    fn kv_fetch_reqs_by_stage_are_recorded_into_stage_gauges() {
+        let metrics = Metrics::new();
+        let handles = super::resolve_scheduler_stats_handles(&metrics.scheduler, "model", 0);
+        // Series of one family render in hash order, so sort them.
+        let kv_fetch_series = || {
+            metrics
+                .render()
+                .unwrap()
+                .lines()
+                .filter(|line| line.starts_with("vllm:num_requests_kv_fetch_by_stage{"))
+                .sorted()
+                .join("\n")
+        };
+
+        expect![[r#"
+            vllm:num_requests_kv_fetch_by_stage{model_name="model",engine="0",stage="completed_waiting"} 0
+            vllm:num_requests_kv_fetch_by_stage{model_name="model",engine="0",stage="in_progress"} 0
+            vllm:num_requests_kv_fetch_by_stage{model_name="model",engine="0",stage="waiting_to_start"} 0"#]]
+        .assert_eq(&kv_fetch_series());
+
+        let stats = SchedulerStats {
+            num_kv_fetch_reqs_by_stage: KvFetchReqsByStage {
+                waiting_to_start: 3,
+                in_progress: 1,
+                completed_waiting: 2,
+            },
+            ..Default::default()
+        };
+        super::record_scheduler_stats_with_handles(&handles, &stats);
+
+        expect![[r#"
+            vllm:num_requests_kv_fetch_by_stage{model_name="model",engine="0",stage="completed_waiting"} 2
+            vllm:num_requests_kv_fetch_by_stage{model_name="model",engine="0",stage="in_progress"} 1
+            vllm:num_requests_kv_fetch_by_stage{model_name="model",engine="0",stage="waiting_to_start"} 3"#]]
+        .assert_eq(&kv_fetch_series());
+    }
+
     /// Records one MultiConnector payload into Mooncake and NIXL metrics.
     #[test]
     fn kv_connector_stats_are_recorded_into_mooncake_and_nixl_metrics() {
diff --git a/rust/src/engine-core-client/src/mock_engine.rs b/rust/src/engine-core-client/src/mock_engine.rs
index 72e05b613c97..ac553fab8487 100644
--- a/rust/src/engine-core-client/src/mock_engine.rs
+++ b/rust/src/engine-core-client/src/mock_engine.rs
@@ -28,6 +28,8 @@ pub struct MockEngineConfig {
     pub local: bool,
     /// Whether the engine should advertise itself as headless.
     pub headless: bool,
+    /// `parallel_config_hash` reported in handshake READY.
+    pub parallel_config_hash: Option<String>,
     /// Engine-ready payload reported after INIT, including max model length,
     /// KV block count, and dtype.
     pub ready_response: EngineCoreReadyResponse,
@@ -40,6 +42,7 @@ impl Default for MockEngineConfig {
         Self {
             local: false,
             headless: true,
+            parallel_config_hash: None,
             ready_response: default_ready_response(),
             connect_timeout: Duration::from_secs(5),
         }
@@ -74,6 +77,7 @@ pub fn default_ready_response() -> EngineCoreReadyResponse {
         supports_draft_weight_updates: false,
         effective_attention_block_size: None,
         vllm_config: None,
+        prefix_hash_config: None,
     }
 }
 
@@ -120,7 +124,7 @@ fn ready_message(status: &str, config: &MockEngineConfig) -> ReadyMessage {
         status: Some(status.to_string()),
         local: Some(config.local),
         headless: Some(config.headless),
-        parallel_config_hash: None,
+        parallel_config_hash: config.parallel_config_hash.clone(),
     }
 }
 
@@ -175,17 +179,77 @@ pub async fn connect_to_frontend(
     engine_id: impl Into<EngineId>,
     config: MockEngineConfig,
 ) -> Result<MockEngineSockets> {
-    let engine_handshake = engine_handshake.as_ref();
-    wait_for_endpoint(engine_handshake, config.connect_timeout).await?;
+    let peer_identity = peer_identity(engine_id)?;
+    let (mut handshake, init) =
+        start_handshake(engine_handshake.as_ref(), &peer_identity, &config).await?;
+    let data_sockets = connect_data_sockets(&init, &peer_identity, &config).await?;
+    let coordinator = connect_coordinator(&init).await?;
+    send_handshake_ready(&mut handshake, &config).await?;
+
+    Ok(MockEngineSockets {
+        init,
+        data_sockets,
+        coordinator,
+    })
+}
 
+/// Join a frontend owning DP rank 0 and this engine's own frontend, as a Python
+/// engine does on a node in hybrid or external DP LB: take the coordinator
+/// addresses and DP rendezvous from the rank-0 INIT and the input/output
+/// addresses from the own frontend's INIT, then send READY to both.
+pub async fn connect_to_hybrid_frontends(
+    head_handshake: impl AsRef<str>,
+    client_handshake: impl AsRef<str>,
+    engine_id: impl Into<EngineId>,
+    mut config: MockEngineConfig,
+) -> Result<MockEngineSockets> {
     let peer_identity = peer_identity(engine_id)?;
+    // Neither handshake is headless, and only the own frontend sees the engine
+    // as local.
+    let head_config = MockEngineConfig {
+        local: false,
+        headless: false,
+        ..config.clone()
+    };
+    config.local = true;
+    config.headless = false;
+    let (mut head, mut init) =
+        start_handshake(head_handshake.as_ref(), &peer_identity, &head_config).await?;
+    let (mut client, client_init) =
+        start_handshake(client_handshake.as_ref(), &peer_identity, &config).await?;
+    init.addresses.inputs = client_init.addresses.inputs;
+    init.addresses.outputs = client_init.addresses.outputs;
+    // The engine reports the rank-0 coordinator to its own frontend.
+    config.ready_response.dp_stats_address = init.addresses.frontend_stats_publish_address.clone();
+
+    let data_sockets = connect_data_sockets(&init, &peer_identity, &config).await?;
+    let coordinator = connect_coordinator(&init).await?;
+    // The engine sends READY to its own frontend first.
+    send_handshake_ready(&mut client, &config).await?;
+    send_handshake_ready(&mut head, &head_config).await?;
+
+    Ok(MockEngineSockets {
+        init,
+        data_sockets,
+        coordinator,
+    })
+}
+
+/// Send HELLO to one frontend handshake endpoint and receive its INIT.
+async fn start_handshake(
+    engine_handshake: &str,
+    peer_identity: &PeerIdentity,
+    config: &MockEngineConfig,
+) -> Result<(DealerSocket, HandshakeInitMessage)> {
+    wait_for_endpoint(engine_handshake, config.connect_timeout).await?;
+
     let mut options = SocketOptions::default();
     options.peer_identity(peer_identity.clone());
     let mut handshake = DealerSocket::with_options(options);
     handshake.connect(engine_handshake).await?;
     handshake
         .send(ZmqMessage::from(encode_msgpack(&ready_message(
-            "HELLO", &config,
+            "HELLO", config,
         ))?))
         .await?;
 
@@ -197,7 +261,15 @@ pub async fn connect_to_frontend(
         );
     }
     let init: HandshakeInitMessage = decode_msgpack(init_frames[0].as_ref())?;
+    Ok((handshake, init))
+}
 
+/// Register on every input/output socket pair named in INIT.
+async fn connect_data_sockets(
+    init: &HandshakeInitMessage,
+    peer_identity: &PeerIdentity,
+    config: &MockEngineConfig,
+) -> Result<Vec<MockEngineDataSockets>> {
     if init.addresses.inputs.is_empty() {
         return Err(Error::UnexpectedHandshakeMessage {
             message: "frontend INIT did not include an input address".to_string(),
@@ -224,15 +296,21 @@ pub async fn connect_to_frontend(
         input_options.peer_identity(peer_identity.clone());
         let mut dealer = DealerSocket::with_options(input_options);
         dealer.connect(input_address).await?;
-        dealer.send(ZmqMessage::from(ready_response_payload(&config)?)).await?;
+        dealer.send(ZmqMessage::from(ready_response_payload(config)?)).await?;
 
         let mut push = PushSocket::new();
         push.connect(output_address).await?;
 
         data_sockets.push(MockEngineDataSockets { dealer, push });
     }
+    Ok(data_sockets)
+}
 
-    let coordinator = match (
+/// Join the coordinator named in INIT, if any, and wait for its READY marker.
+async fn connect_coordinator(
+    init: &HandshakeInitMessage,
+) -> Result<Option<MockCoordinatorSockets>> {
+    match (
         init.addresses.coordinator_input.as_deref(),
         init.addresses.coordinator_output.as_deref(),
     ) {
@@ -252,28 +330,29 @@ pub async fn connect_to_frontend(
                 );
             }
 
-            Some(MockCoordinatorSockets {
+            Ok(Some(MockCoordinatorSockets {
                 input_sub,
                 output_push,
-            })
+            }))
         }
-        (None, None) => None,
+        (None, None) => Ok(None),
         _ => bail_unexpected_handshake_message!(
             "coordinator handshake addresses must be both present or both absent"
         ),
-    };
+    }
+}
 
+/// Send the handshake READY status.
+async fn send_handshake_ready(
+    handshake: &mut DealerSocket,
+    config: &MockEngineConfig,
+) -> Result<()> {
     handshake
         .send(ZmqMessage::from(encode_msgpack(&ready_message(
-            "READY", &config,
+            "READY", config,
         ))?))
         .await?;
-
-    Ok(MockEngineSockets {
-        init,
-        data_sockets,
-        coordinator,
-    })
+    Ok(())
 }
 
 /// Join already-bootstrapped frontend input/output sockets directly.
diff --git a/rust/src/engine-core-client/src/prompt_block_hasher.rs b/rust/src/engine-core-client/src/prompt_block_hasher.rs
new file mode 100644
index 000000000000..2ab8992043ce
--- /dev/null
+++ b/rust/src/engine-core-client/src/prompt_block_hasher.rs
@@ -0,0 +1,538 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+//! Prompt block hashes computed in the frontend, so that the engine doesn't
+//! hash long prompts on its input thread.
+//!
+//! Each hash must be byte-identical to the engine's own, because the engine
+//! chains the hashes of blocks completed by outputs onto them, and because KV
+//! connectors and KV-event consumers compare them across instances.
+//!
+//! Original Python definitions:
+//! - `hash_block_tokens`, `generate_block_hash_extra_keys` and
+//!   `get_request_block_hasher`:
+//!   <https://github.com/vllm-project/vllm/blob/main/vllm/v1/core/kv_cache_utils.py>
+//! - `sha256` and `sha256_cbor`:
+//!   <https://github.com/vllm-project/vllm/blob/main/vllm/utils/hashing.py>
+
+use sha2::{Digest, Sha256};
+
+use crate::protocol::handshake::PrefixHashConfig;
+use crate::protocol::request::{BlockHash, EngineCoreRequest};
+
+/// Python's C pickler starts a new frame once one reaches 64 KiB, which this
+/// encoder doesn't reproduce, so longer keys are left to the engine.
+const PICKLE_FRAME_SIZE_TARGET: usize = 64 * 1024;
+
+/// How the engine serializes `(parent_hash, tuple(block_tokens), extra_keys)`
+/// before hashing it with SHA-256.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+enum KeyEncoding {
+    /// `sha256`: `pickle.dumps(key, protocol=5)`.
+    Pickle,
+    /// `sha256_cbor`: `cbor2.dumps(key, canonical=True)`.
+    Cbor,
+}
+
+/// A Python value in a block's extra keys, limited to the kinds this frontend
+/// encodes.
+#[derive(Debug, Clone, PartialEq, Eq)]
+enum ExtraKey<'a> {
+    None,
+    Int(u32),
+    Str(&'a str),
+    Tuple(Vec<ExtraKey<'a>>),
+}
+
+/// The request-level inputs to `generate_block_hash_extra_keys` that this
+/// frontend encodes.
+struct RequestExtraKeys<'a> {
+    /// `("lora", lora_name, lora_path)`, on every block.
+    lora: Option<ExtraKey<'a>>,
+    /// A non-empty cache salt, keyed on the first block only.
+    cache_salt: Option<&'a str>,
+}
+
+impl<'a> RequestExtraKeys<'a> {
+    /// The extra keys of the block starting at `start`: a tuple of the LoRA
+    /// and cache-salt keys, or `None` without either.
+    fn for_block(&self, start: usize) -> ExtraKey<'a> {
+        let cache_salt = self
+            .cache_salt
+            .filter(|_| start == 0)
+            .map(|salt| ExtraKey::Tuple(vec![ExtraKey::Str("cache_salt"), ExtraKey::Str(salt)]));
+        let keys: Vec<_> = self.lora.iter().cloned().chain(cache_salt).collect();
+        if keys.is_empty() {
+            ExtraKey::None
+        } else {
+            ExtraKey::Tuple(keys)
+        }
+    }
+}
+
+/// Computes the hashes of a prompt's full blocks the way the engine would.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub(crate) struct PromptBlockHasher {
+    encoding: KeyEncoding,
+    block_size: usize,
+    none_hash: Vec<u8>,
+    num_lookahead_tokens: usize,
+}
+
+impl PromptBlockHasher {
+    /// Return a hasher for the engine's config, or `None` when this frontend
+    /// can't reproduce its hashes, in which case the engine hashes prompts.
+    pub(crate) fn new(config: &PrefixHashConfig) -> Option<Self> {
+        // TODO: support `xxhash` and `xxhash_cbor` (XXH3-128 over the same
+        // encodings).
+        let encoding = match config.hash_algo.as_str() {
+            "sha256" => KeyEncoding::Pickle,
+            "sha256_cbor" => KeyEncoding::Cbor,
+            _ => return None,
+        };
+        let block_size = usize::try_from(config.hash_block_size).ok()?;
+        let num_lookahead_tokens = usize::try_from(config.num_lookahead_tokens).ok()?;
+        let none_hash = config.none_hash.0.clone();
+        // Both encodings below write the parent hash with a 1-byte length.
+        let supported = block_size > 0 && none_hash.len() < 256;
+        supported.then_some(Self {
+            encoding,
+            block_size,
+            none_hash,
+            num_lookahead_tokens,
+        })
+    }
+
+    /// Return the request's prompt block hashes, or `None` if the engine has to
+    /// compute them because they need inputs this frontend doesn't encode.
+    pub(crate) fn hash_request(&self, request: &EngineCoreRequest) -> Option<Vec<BlockHash>> {
+        // TODO: encode multimodal extra keys, `("mm", identifier, offset)` per
+        // item overlapping the block, and the prompt-embeds keys of mixed-mode
+        // requests once this frontend supports prompt embeds.
+        let unsupported = request.mm_features.as_ref().is_some_and(|features| !features.is_empty())
+            || request.prompt_embeds.is_some()
+            || request.prompt_is_token_ids.is_some();
+        if unsupported {
+            return None;
+        }
+        let lora = request.lora_request.as_ref();
+        // Python adds no cache-salt key for an empty salt.
+        let cache_salt = request.cache_salt.as_deref().filter(|salt| !salt.is_empty());
+        let strings = (lora.into_iter())
+            .flat_map(|lora| [lora.lora_name.as_str(), lora.lora_path.as_str()])
+            .chain(cache_salt);
+        if self.encoding == KeyEncoding::Pickle && may_share_pickled_objects(strings) {
+            return None;
+        }
+        let keys = RequestExtraKeys {
+            lora: lora.map(|lora| {
+                ExtraKey::Tuple(vec![
+                    ExtraKey::Str("lora"),
+                    ExtraKey::Str(&lora.lora_name),
+                    ExtraKey::Str(&lora.lora_path),
+                ])
+            }),
+            cache_salt,
+        };
+        self.hash_prompt(request.prompt_token_ids.as_deref()?, &keys)
+    }
+
+    /// Hash each full block of `token_ids` whose lookahead tokens are also in
+    /// the prompt, chained from the engine's `NONE_HASH`. Returns `None` for a
+    /// key this frontend can't encode exactly.
+    fn hash_prompt(
+        &self,
+        token_ids: &[u32],
+        keys: &RequestExtraKeys<'_>,
+    ) -> Option<Vec<BlockHash>> {
+        let lookahead = self.num_lookahead_tokens;
+        let num_blocks = token_ids.len().saturating_sub(lookahead) / self.block_size;
+        let mut key = Vec::new();
+        let mut parent = self.none_hash.clone();
+        (0..num_blocks)
+            .map(|index| {
+                let start = index * self.block_size;
+                let end = start + self.block_size;
+                let mut extra_keys = keys.for_block(start);
+                // With lookahead, a block's extra keys become `(E, *tokens)`
+                // with the lookahead tokens after it.
+                if lookahead > 0 {
+                    let tokens = token_ids[end..end + lookahead].iter().map(|&t| ExtraKey::Int(t));
+                    extra_keys = ExtraKey::Tuple([extra_keys].into_iter().chain(tokens).collect());
+                }
+                key.clear();
+                match self.encoding {
+                    KeyEncoding::Pickle => {
+                        encode_pickle_key(&mut key, &parent, &token_ids[start..end], &extra_keys)?
+                    }
+                    KeyEncoding::Cbor => {
+                        encode_cbor_key(&mut key, &parent, &token_ids[start..end], &extra_keys)
+                    }
+                }
+                parent = Sha256::digest(&key).to_vec();
+                Some(BlockHash(parent.clone()))
+            })
+            .collect()
+    }
+}
+
+/// Whether pickle may write one of these strings as a reference to an
+/// identical object earlier in the key instead of writing it again: pickle's
+/// memo is by object identity, CPython shares the empty and single-Latin-1-
+/// character strings, and a string could also be the very object of a key
+/// literal. Conservative; requests where this matters are left to the engine.
+fn may_share_pickled_objects<'a>(strings: impl Iterator<Item = &'a str>) -> bool {
+    let shared_by_cpython = |s: &str| {
+        let mut chars = s.chars();
+        match (chars.next(), chars.next()) {
+            (None, _) => true,
+            (Some(c), None) => u32::from(c) < 0x100,
+            _ => false,
+        }
+    };
+    let strings: Vec<_> = strings.collect();
+    strings.iter().enumerate().any(|(i, s)| {
+        ["lora", "cache_salt"].contains(s) || (shared_by_cpython(s) && strings[i + 1..].contains(s))
+    })
+}
+
+const MEMOIZE: u8 = 0x94;
+
+/// Write `pickle.dumps((parent, tuple(block), extra_keys), protocol=5)`, as
+/// CPython's C pickler writes it.
+fn encode_pickle_key(
+    out: &mut Vec<u8>,
+    parent: &[u8],
+    block: &[u32],
+    extra_keys: &ExtraKey<'_>,
+) -> Option<()> {
+    // PROTO 5, then FRAME with the length of everything after it.
+    out.extend_from_slice(&[0x80, 0x05, 0x95]);
+    let frame_len_at = out.len();
+    out.extend_from_slice(&[0; 8]);
+    // SHORT_BINBYTES parent.
+    out.push(b'C');
+    out.push(u8::try_from(parent.len()).ok()?);
+    out.extend_from_slice(parent);
+    out.push(MEMOIZE);
+    write_pickle_tuple(out, block.len(), |out| {
+        block.iter().try_for_each(|&token| write_pickle_int(out, token))
+    })?;
+    write_pickle_value(out, extra_keys)?;
+    // TUPLE3 for the key itself, then STOP.
+    out.extend_from_slice(&[0x87, MEMOIZE, b'.']);
+    let frame_len = out.len() - frame_len_at - 8;
+    if frame_len >= PICKLE_FRAME_SIZE_TARGET {
+        return None;
+    }
+    out[frame_len_at..frame_len_at + 8].copy_from_slice(&(frame_len as u64).to_le_bytes());
+    Some(())
+}
+
+/// Write a non-empty tuple whose items `write_items` writes.
+fn write_pickle_tuple(
+    out: &mut Vec<u8>,
+    len: usize,
+    write_items: impl FnOnce(&mut Vec<u8>) -> Option<()>,
+) -> Option<()> {
+    if len == 0 {
+        // EMPTY_TUPLE, which pickle doesn't memoize.
+        out.push(b')');
+        return Some(());
+    }
+    // Tuples of up to three items have dedicated opcodes; longer ones use MARK.
+    if len > 3 {
+        out.push(b'(');
+    }
+    write_items(out)?;
+    out.push(match len {
+        1 => 0x85, // TUPLE1
+        2 => 0x86, // TUPLE2
+        3 => 0x87, // TUPLE3
+        _ => b't', // TUPLE
+    });
+    out.push(MEMOIZE);
+    Some(())
+}
+
+fn write_pickle_int(out: &mut Vec<u8>, value: u32) -> Option<()> {
+    match value {
+        // BININT1, BININT2, BININT (signed).
+        0..=0xff => out.extend_from_slice(&[b'K', value as u8]),
+        0x100..=0xffff => {
+            out.push(b'M');
+            out.extend_from_slice(&(value as u16).to_le_bytes());
+        }
+        0x1_0000..=0x7fff_ffff => {
+            out.push(b'J');
+            out.extend_from_slice(&value.to_le_bytes());
+        }
+        // Larger ints are pickled with LONG1, which no real token id needs.
+        _ => return None,
+    }
+    Some(())
+}
+
+fn write_pickle_value(out: &mut Vec<u8>, value: &ExtraKey<'_>) -> Option<()> {
+    match value {
+        ExtraKey::None => out.push(b'N'),
+        ExtraKey::Int(value) => write_pickle_int(out, *value)?,
+        ExtraKey::Str(s) => {
+            // SHORT_BINUNICODE or BINUNICODE, of the UTF-8 bytes.
+            match u8::try_from(s.len()) {
+                Ok(len) => out.extend_from_slice(&[0x8c, len]),
+                Err(_) => {
+                    out.push(b'X');
+                    out.extend_from_slice(&u32::try_from(s.len()).ok()?.to_le_bytes());
+                }
+            }
+            out.extend_from_slice(s.as_bytes());
+            out.push(MEMOIZE);
+        }
+        ExtraKey::Tuple(items) => write_pickle_tuple(out, items.len(), |out| {
+            items.iter().try_for_each(|item| write_pickle_value(out, item))
+        })?,
+    }
+    Some(())
+}
+
+const CBOR_BYTES: u8 = 0x40;
+const CBOR_TEXT: u8 = 0x60;
+const CBOR_ARRAY: u8 = 0x80;
+const CBOR_NULL: u8 = 0xf6;
+
+/// Write `cbor2.dumps((parent, tuple(block), extra_keys), canonical=True)`.
+fn encode_cbor_key(out: &mut Vec<u8>, parent: &[u8], block: &[u32], extra_keys: &ExtraKey<'_>) {
+    out.push(CBOR_ARRAY | 3);
+    write_cbor_head(out, CBOR_BYTES, parent.len() as u64);
+    out.extend_from_slice(parent);
+    write_cbor_head(out, CBOR_ARRAY, block.len() as u64);
+    for &token in block {
+        write_cbor_head(out, 0x00, token.into());
+    }
+    write_cbor_value(out, extra_keys);
+}
+
+fn write_cbor_value(out: &mut Vec<u8>, value: &ExtraKey<'_>) {
+    match value {
+        ExtraKey::None => out.push(CBOR_NULL),
+        ExtraKey::Int(value) => write_cbor_head(out, 0x00, (*value).into()),
+        ExtraKey::Str(s) => {
+            write_cbor_head(out, CBOR_TEXT, s.len() as u64);
+            out.extend_from_slice(s.as_bytes());
+        }
+        ExtraKey::Tuple(items) => {
+            write_cbor_head(out, CBOR_ARRAY, items.len() as u64);
+            items.iter().for_each(|item| write_cbor_value(out, item));
+        }
+    }
+}
+
+/// Write a CBOR item head in its shortest form, as canonical encoding requires.
+fn write_cbor_head(out: &mut Vec<u8>, major: u8, value: u64) {
+    match value {
+        0..=23 => out.push(major | value as u8),
+        24..=0xff => out.extend_from_slice(&[major | 24, value as u8]),
+        0x100..=0xffff => {
+            out.push(major | 25);
+            out.extend_from_slice(&(value as u16).to_be_bytes());
+        }
+        0x1_0000..=0xffff_ffff => {
+            out.push(major | 26);
+            out.extend_from_slice(&(value as u32).to_be_bytes());
+        }
+        _ => {
+            out.push(major | 27);
+            out.extend_from_slice(&value.to_be_bytes());
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::collections::HashMap;
+
+    use serde::Deserialize;
+
+    use super::*;
+    use crate::protocol::lora::LoraRequest;
+
+    const SHA256_NONE_HASH: &str =
+        "066b1e8987e027138e98701c46128929378935f9b6cd1549105ae6b4f0d71198";
+    const SHA256_CBOR_NONE_HASH: &str =
+        "9bd96a485ad84efdafb72ee48a1d7a69bcead0f8f0433173941b276b9581eef0";
+
+    fn hasher(hash_algo: &str, block_size: u64, lookahead: u64) -> PromptBlockHasher {
+        let none_hash = match hash_algo {
+            "sha256" => SHA256_NONE_HASH,
+            _ => SHA256_CBOR_NONE_HASH,
+        };
+        PromptBlockHasher::new(&PrefixHashConfig {
+            hash_algo: hash_algo.to_string(),
+            hash_block_size: block_size,
+            none_hash: BlockHash(hex::decode(none_hash).unwrap()),
+            num_lookahead_tokens: lookahead,
+        })
+        .unwrap()
+    }
+
+    fn request(prompt_token_ids: Vec<u32>) -> EngineCoreRequest {
+        EngineCoreRequest {
+            prompt_token_ids: Some(prompt_token_ids),
+            ..EngineCoreRequest::default()
+        }
+    }
+
+    fn lora_request(lora_name: &str, lora_path: &str) -> LoraRequest {
+        LoraRequest {
+            lora_name: lora_name.to_string(),
+            lora_int_id: 1,
+            lora_path: lora_path.to_string(),
+            base_model_name: None,
+            load_inplace: false,
+            is_3d_lora_weight: false,
+        }
+    }
+
+    /// Prompt block hashes from the Python engine, kept current by
+    /// `tests/v1/core/test_prompt_block_hash_vectors.py` in the repo root, which
+    /// fails if the engine's hashing changes.
+    const PYTHON_VECTORS: &str = include_str!("tests/prompt_block_hash_vectors.json");
+
+    #[derive(Deserialize)]
+    struct Vectors {
+        none_hashes: HashMap<String, String>,
+        cases: Vec<VectorCase>,
+    }
+
+    #[derive(Deserialize)]
+    struct VectorCase {
+        name: String,
+        hash_algo: String,
+        block_size: u64,
+        num_lookahead_tokens: u64,
+        lora: Option<VectorLora>,
+        cache_salt: Option<String>,
+        prompt_token_ids: Vec<u32>,
+        block_hashes: Vec<String>,
+    }
+
+    #[derive(Deserialize)]
+    struct VectorLora {
+        name: String,
+        path: String,
+    }
+
+    /// Plain, LoRA, cache-salt and lookahead requests at several block sizes,
+    /// against `Request.block_hashes` from the real `get_request_block_hasher`
+    /// with `sha256` and `sha256_cbor`.
+    #[test]
+    fn prompt_block_hashes_match_python_engine() {
+        let vectors: Vectors = serde_json::from_str(PYTHON_VECTORS).unwrap();
+        // The vectors chain from the engine's default-seed `NONE_HASH`.
+        assert_eq!(vectors.none_hashes["sha256"], SHA256_NONE_HASH);
+        assert_eq!(vectors.none_hashes["sha256_cbor"], SHA256_CBOR_NONE_HASH);
+        assert!(!vectors.cases.is_empty());
+        for case in vectors.cases {
+            let request = EngineCoreRequest {
+                lora_request: case.lora.map(|lora| lora_request(&lora.name, &lora.path)),
+                cache_salt: case.cache_salt,
+                ..request(case.prompt_token_ids)
+            };
+            let hashes = hasher(&case.hash_algo, case.block_size, case.num_lookahead_tokens)
+                .hash_request(&request)
+                .unwrap_or_else(|| panic!("{} {}: not hashed", case.hash_algo, case.name));
+            let hashes: Vec<_> = hashes.iter().map(|hash| hex::encode(&hash.0)).collect();
+            assert_eq!(
+                hashes, case.block_hashes,
+                "{} {}",
+                case.hash_algo, case.name
+            );
+        }
+    }
+
+    #[test]
+    fn prompt_without_a_full_block_has_no_hashes() {
+        assert_eq!(
+            hasher("sha256", 16, 0).hash_request(&request(vec![1; 15])),
+            Some(vec![])
+        );
+        // The lookahead tokens of the only full block are missing.
+        assert_eq!(
+            hasher("sha256", 16, 2).hash_request(&request(vec![1; 17])),
+            Some(vec![])
+        );
+    }
+
+    #[test]
+    fn token_id_beyond_i32_is_left_to_the_engine() {
+        let tokens = vec![1, 0x8000_0000];
+        assert_eq!(
+            hasher("sha256", 2, 0).hash_request(&request(tokens.clone())),
+            None
+        );
+        // CBOR encodes every u32 directly.
+        assert!(hasher("sha256_cbor", 2, 0).hash_request(&request(tokens)).is_some());
+    }
+
+    #[test]
+    fn key_longer_than_a_pickle_frame_is_left_to_the_engine() {
+        let tokens = vec![70_000; 16 * 1024];
+        assert_eq!(
+            hasher("sha256", 16 * 1024, 0).hash_request(&request(tokens.clone())),
+            None
+        );
+        assert!(hasher("sha256_cbor", 16 * 1024, 0).hash_request(&request(tokens)).is_some());
+    }
+
+    #[test]
+    fn unsupported_configs_are_left_to_the_engine() {
+        let config = |hash_algo: &str, hash_block_size| PrefixHashConfig {
+            hash_algo: hash_algo.to_string(),
+            hash_block_size,
+            none_hash: BlockHash(vec![0; 32]),
+            num_lookahead_tokens: 0,
+        };
+        assert!(PromptBlockHasher::new(&config("xxhash", 16)).is_none());
+        assert!(PromptBlockHasher::new(&config("sha256", 0)).is_none());
+    }
+
+    #[test]
+    fn unsupported_requests_are_left_to_the_engine() {
+        let pickle = hasher("sha256", 2, 0);
+        let base = request(vec![1, 2, 3, 4]);
+        // Python adds no cache-salt key for an empty salt.
+        let empty_salt = EngineCoreRequest {
+            cache_salt: Some(String::new()),
+            ..base.clone()
+        };
+        assert_eq!(pickle.hash_request(&empty_salt), pickle.hash_request(&base));
+        let mixed = EngineCoreRequest {
+            prompt_is_token_ids: Some(vec![true; 4]),
+            ..base.clone()
+        };
+        assert_eq!(pickle.hash_request(&mixed), None);
+        // Pickle might write these strings as references to an earlier object:
+        // repeated shared strings, or one that could be a key literal.
+        let lora = |lora_name: &str, lora_path: &str, cache_salt: Option<&str>| EngineCoreRequest {
+            lora_request: Some(lora_request(lora_name, lora_path)),
+            cache_salt: cache_salt.map(str::to_string),
+            ..base.clone()
+        };
+        for request in [
+            lora("a", "a", None),
+            lora("", "", None),
+            lora("x", "/p", Some("x")),
+            lora("lora", "/p", None),
+            lora("adapter", "/p", Some("cache_salt")),
+        ] {
+            assert_eq!(pickle.hash_request(&request), None, "{request:?}");
+        }
+        // Equal strings that CPython doesn't share are distinct objects.
+        assert!(pickle.hash_request(&lora("adapter", "/p", Some("adapter"))).is_some());
+        assert!(pickle.hash_request(&lora("a", "b", None)).is_some());
+        // CBOR has no memo.
+        let cbor = hasher("sha256_cbor", 2, 0);
+        assert!(cbor.hash_request(&lora("a", "a", None)).is_some());
+    }
+}
diff --git a/rust/src/engine-core-client/src/protocol/handshake.rs b/rust/src/engine-core-client/src/protocol/handshake.rs
index 142ef4729528..73696b52a421 100644
--- a/rust/src/engine-core-client/src/protocol/handshake.rs
+++ b/rust/src/engine-core-client/src/protocol/handshake.rs
@@ -1,12 +1,10 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-use std::collections::BTreeMap;
-
 use serde::{Deserialize, Serialize};
 
-use crate::protocol::OpaqueValue;
 use crate::protocol::dtype::ModelDtype;
+use crate::protocol::request::BlockHash;
 
 /// Decoded engine startup-handshake payload sent on the handshake socket.
 ///
@@ -37,6 +35,25 @@ pub struct KvEventsConfig {
     pub topic: String,
 }
 
+/// What the frontend needs to compute prompt block hashes itself, reported
+/// only when the engine hashes request blocks.
+///
+/// Original Python definition (`PrefixHashConfig`):
+/// <https://github.com/vllm-project/vllm/blob/main/vllm/v1/engine/__init__.py>
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+pub struct PrefixHashConfig {
+    /// The engine's `--prefix-caching-hash-algo`, e.g. `sha256`.
+    pub hash_algo: String,
+    /// Tokens per hashed block.
+    pub hash_block_size: u64,
+    /// Parent hash of the first block.
+    pub none_hash: BlockHash,
+    /// Tokens after each block that its hash also covers, so a block is
+    /// hashed only once they exist.
+    #[serde(default)]
+    pub num_lookahead_tokens: u64,
+}
+
 /// Post-initialization configuration sent from each engine on the input socket
 /// registration message, after the handshake completes.
 ///
@@ -111,6 +128,10 @@ pub struct EngineCoreReadyResponse {
     /// engine runs with `VLLM_SERVER_DEV_MODE` (for `/server_info`).
     #[serde(default)]
     pub vllm_config: Option<serde_json::Value>,
+    /// Set when the engine hashes request blocks and accepts prompt block
+    /// hashes from the frontend.
+    #[serde(default)]
+    pub prefix_hash_config: Option<PrefixHashConfig>,
 }
 
 /// Frontend-owned ZMQ addresses that are sent to the engine during startup
@@ -135,5 +156,56 @@ pub struct HandshakeAddresses {
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct HandshakeInitMessage {
     pub addresses: HandshakeAddresses,
-    pub parallel_config: BTreeMap<String, OpaqueValue>,
+    pub parallel_config: HandshakeParallelConfig,
+}
+
+/// `ParallelConfig` fields that the engine overwrites from INIT, encoded as a
+/// map.
+#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
+pub struct HandshakeParallelConfig {
+    /// Sent for coordinated (MoE) DP only; otherwise the map is empty.
+    #[serde(flatten)]
+    pub data_parallel_rendezvous: Option<DataParallelRendezvous>,
+}
+
+/// Where all DP ranks rendezvous to form their cross-rank torch process
+/// groups.
+///
+/// `_coord_store_port` is not sent: no coordination store runs, so engines
+/// take the listed ports.
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+pub struct DataParallelRendezvous {
+    pub data_parallel_master_ip: String,
+    pub data_parallel_master_port: u16,
+    /// Free ports that the engine takes one per process group before falling
+    /// back to incrementing `data_parallel_master_port`.
+    #[serde(rename = "_data_parallel_master_port_list")]
+    pub data_parallel_master_port_list: Vec<u16>,
+    pub data_parallel_size: usize,
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::protocol::{decode_value, encode_msgpack};
+
+    #[test]
+    fn handshake_parallel_config_encodes_rendezvous_fields_or_empty_map() {
+        let encode = |parallel_config: HandshakeParallelConfig| {
+            decode_value(&encode_msgpack(&parallel_config).unwrap()).unwrap().to_string()
+        };
+
+        let rendezvous = DataParallelRendezvous {
+            data_parallel_master_ip: "10.0.0.1".to_string(),
+            data_parallel_master_port: 29500,
+            data_parallel_master_port_list: vec![29501, 29502, 29503, 29504],
+            data_parallel_size: 4,
+        };
+
+        expect_test::expect!["{}"].assert_eq(&encode(HandshakeParallelConfig::default()));
+        expect_test::expect![[r#"{"data_parallel_master_ip": "10.0.0.1", "data_parallel_master_port": 29500, "_data_parallel_master_port_list": [29501, 29502, 29503, 29504], "data_parallel_size": 4}"#]]
+            .assert_eq(&encode(HandshakeParallelConfig {
+                data_parallel_rendezvous: Some(rendezvous),
+            }));
+    }
 }
diff --git a/rust/src/engine-core-client/src/protocol/request.rs b/rust/src/engine-core-client/src/protocol/request.rs
index 5390eed16f2c..c0ea5a132859 100644
--- a/rust/src/engine-core-client/src/protocol/request.rs
+++ b/rust/src/engine-core-client/src/protocol/request.rs
@@ -7,6 +7,7 @@ use bytes::Bytes;
 use serde::{Deserialize, Serialize};
 use serde_default::DefaultFromSerde;
 use serde_tuple::{Deserialize_tuple, Serialize_tuple};
+use serde_with::serde_as;
 
 use crate::protocol::kv_hints::KvHintsEnvelope;
 use crate::protocol::multimodal::MmFeatures;
@@ -131,8 +132,20 @@ pub struct EngineCoreRequest {
     pub session_id: Option<String>,
     #[serde(default)]
     pub kv_hints: Option<KvHintsEnvelope>,
+    /// Hashes of the prompt's full blocks, computed with the engine's
+    /// [`PrefixHashConfig`](crate::protocol::handshake::PrefixHashConfig) so
+    /// the engine doesn't hash the prompt itself. They must include any extra
+    /// keys the engine would add (LoRA, multimodal, cache salt).
+    #[serde(default)]
+    pub block_hashes: Option<Vec<BlockHash>>,
 }
 
+/// A prefix-cache block hash, carried as msgpack bin like Python `bytes`.
+#[serde_as]
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+#[serde(transparent)]
+pub struct BlockHash(#[serde_as(as = "serde_with::Bytes")] pub Vec<u8>);
+
 impl EngineCoreRequest {
     /// Validate fields intentionally not supported in the first-stage client.
     pub fn validate(&self) -> Result<()> {
@@ -211,7 +224,7 @@ mod tests {
             other => panic!("expected array, got {other:?}"),
         };
 
-        assert_eq!(array.len(), 22);
+        assert_eq!(array.len(), 23);
         assert_eq!(array[0], Value::from("req-1"));
         assert_eq!(array[2], Value::Nil);
         assert_eq!(array[4], Value::Nil);
@@ -219,6 +232,8 @@ mod tests {
         assert_eq!(array[11], Value::from(7));
         assert_eq!(array[20], Value::from("session-1"));
         assert!(matches!(&array[21], Value::Map(_)));
+        // `block_hashes`, filled in by the client.
+        assert_eq!(array[22], Value::Nil);
     }
 
     #[test]
diff --git a/rust/src/engine-core-client/src/protocol/stats.rs b/rust/src/engine-core-client/src/protocol/stats.rs
index c56f32eb244e..16cf220ced25 100644
--- a/rust/src/engine-core-client/src/protocol/stats.rs
+++ b/rust/src/engine-core-client/src/protocol/stats.rs
@@ -49,6 +49,32 @@ impl CachedTokensBySource {
     }
 }
 
+/// Waiting requests with an async KV load, per load stage. Python sends a
+/// `dict[str, int]` keyed by stage, so absent stages default to zero.
+#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
+#[serde(default)]
+pub struct KvFetchReqsByStage {
+    /// Needs an async KV load that has not started.
+    pub waiting_to_start: u64,
+    /// Load started and not yet reported finished.
+    pub in_progress: u64,
+    /// Load finished, request not running yet.
+    pub completed_waiting: u64,
+}
+
+impl KvFetchReqsByStage {
+    /// Label values, in the same order as [`Self::counts`].
+    pub const STAGES: [&'static str; 3] = ["waiting_to_start", "in_progress", "completed_waiting"];
+
+    pub fn counts(&self) -> [u64; 3] {
+        [
+            self.waiting_to_start,
+            self.in_progress,
+            self.completed_waiting,
+        ]
+    }
+}
+
 /// Stores prefix cache hit statistics.
 /// - `reset`: Whether `reset_prefix_cache` was invoked.
 /// - `queries`: Refers to the number of tokens that were queried.
@@ -295,6 +321,44 @@ pub mod kv_connector {
     #[serde(transparent)]
     pub struct MooncakeStats(pub BTreeMap<MooncakeOperation, Vec<MooncakeRecord>>);
 
+    /// UMBP store connector transfer operation name.
+    #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
+    #[serde(rename_all = "snake_case")]
+    pub enum UmbpOperation {
+        Load,
+        Store,
+    }
+
+    impl UmbpOperation {
+        pub(crate) const fn as_str(self) -> &'static str {
+            match self {
+                Self::Load => "load",
+                Self::Store => "store",
+            }
+        }
+    }
+
+    /// Interval counters of one UMBP store transfer operation.
+    #[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
+    pub struct UmbpOperationCounts {
+        /// Objects submitted.
+        pub submitted: u64,
+        /// Objects completed.
+        pub completed: u64,
+        /// Objects failed.
+        pub failed: u64,
+        /// Bytes in successfully completed transfer ranges.
+        pub num_bytes: u64,
+    }
+
+    /// UMBP store connector telemetry grouped by operation.
+    ///
+    /// Original Python definition: `UMBPStoreConnectorStats` in
+    /// `vllm/distributed/kv_transfer/kv_connector/v1/umbp/stats.py`.
+    #[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
+    #[serde(transparent)]
+    pub struct UmbpStats(pub BTreeMap<UmbpOperation, UmbpOperationCounts>);
+
     /// Telemetry emitted by a `MultiConnector` with a flat child payload map.
     #[serde_with::skip_serializing_none]
     #[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
@@ -309,6 +373,8 @@ pub mod kv_connector {
         pub mooncake_transfer: Option<MooncakeTransferStats>,
         #[serde(rename = "MooncakeStoreConnector")]
         pub mooncake: Option<MooncakeStats>,
+        #[serde(rename = "UMBPStoreConnector")]
+        pub umbp: Option<UmbpStats>,
         /// Child connector payloads without a Rust telemetry implementation.
         #[serde(flatten)]
         pub other: BTreeMap<String, OpaqueValue>,
@@ -329,6 +395,7 @@ pub mod kv_connector {
         Nixl(NixlStats),
         MooncakeTransfer(MooncakeTransferStats),
         Mooncake(MooncakeStats),
+        Umbp(UmbpStats),
         Multi(Box<MultiConnectorStats>),
         Other(BTreeMap<String, OpaqueValue>),
     }
@@ -349,6 +416,9 @@ pub struct SchedulerStats {
     /// Number of waiting requests deferred by transient constraints or a blocked status.
     #[serde(default)]
     pub num_skipped_waiting_reqs: u64,
+    /// Waiting requests with an async KV load, by load stage.
+    #[serde(default)]
+    pub num_kv_fetch_reqs_by_stage: KvFetchReqsByStage,
     /// Internal DP load-balancing step counter.
     pub step_counter: u64,
     /// Internal DP load-balancing wave number.
@@ -373,7 +443,10 @@ pub struct SchedulerStats {
 
 #[cfg(test)]
 mod tests {
-    use super::{CachedTokensBySource, KvConnectorStats, PrefixCacheStats};
+    use super::{
+        CachedTokensBySource, KvConnectorStats, KvFetchReqsByStage, PrefixCacheStats,
+        SchedulerStats, UmbpOperation,
+    };
 
     #[test]
     fn prefix_cache_stats_hits_by_source_defaults_and_decodes() {
@@ -422,14 +495,52 @@ mod tests {
             decode(mooncake_transfer.clone()),
             KvConnectorStats::MooncakeTransfer(stats) if stats.num_failed_transfers == [1]
         ));
+        let umbp = serde_json::json!({
+            "load": {"submitted": 2, "completed": 1, "failed": 1, "num_bytes": 4096,
+                     "future_counter": 1}
+        });
+        assert!(matches!(
+            decode(umbp.clone()),
+            KvConnectorStats::Umbp(stats)
+                if stats.0[&UmbpOperation::Load].num_bytes == 4096
+        ));
         // A newer child payload must not drop its siblings.
         assert!(matches!(
             decode(serde_json::json!({
                 "NixlConnector": nixl,
                 "MooncakeConnector": mooncake_transfer,
+                "UMBPStoreConnector": umbp,
             })),
             KvConnectorStats::Multi(stats)
                 if stats.nixl.is_some() && stats.mooncake_transfer.is_some()
+                    && stats.umbp.is_some()
         ));
     }
+
+    #[test]
+    fn scheduler_stats_kv_fetch_reqs_by_stage_defaults_and_decodes() {
+        let mut payload = serde_json::json!({
+            "num_running_reqs": 1, "num_waiting_reqs": 2,
+            "step_counter": 0, "current_wave": 0, "kv_cache_usage": 0.5,
+            "prefix_cache_stats": {
+                "reset": false, "requests": 0, "queries": 0, "hits": 0,
+                "preempted_requests": 0, "preempted_queries": 0, "preempted_hits": 0
+            },
+            "kv_cache_eviction_events": []
+        });
+
+        // Engines before vllm-project/vllm#58874 omit the field entirely.
+        let wire = rmp_serde::to_vec_named(&payload).unwrap();
+        let stats: SchedulerStats = rmp_serde::from_slice(&wire).unwrap();
+        assert_eq!(
+            stats.num_kv_fetch_reqs_by_stage,
+            KvFetchReqsByStage::default()
+        );
+
+        payload["num_kv_fetch_reqs_by_stage"] =
+            serde_json::json!({"waiting_to_start": 3, "in_progress": 1});
+        let wire = rmp_serde::to_vec_named(&payload).unwrap();
+        let stats: SchedulerStats = rmp_serde::from_slice(&wire).unwrap();
+        assert_eq!(stats.num_kv_fetch_reqs_by_stage.counts(), [3, 1, 0]);
+    }
 }
diff --git a/rust/src/engine-core-client/src/protocol/utility.rs b/rust/src/engine-core-client/src/protocol/utility.rs
index 5379ff49e9e8..97a880c3acf2 100644
--- a/rust/src/engine-core-client/src/protocol/utility.rs
+++ b/rust/src/engine-core-client/src/protocol/utility.rs
@@ -118,6 +118,43 @@ impl FromStr for CheckpointState {
     }
 }
 
+/// Level of the `sleep` utility call, sent as Python's integer level.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
+#[serde(into = "u32", try_from = "u32")]
+pub enum SleepLevel {
+    /// Level 0: pause scheduling only, without changing GPU memory.
+    PauseOnly,
+    /// Level 1: offload model weights to CPU and discard the KV cache.
+    OffloadWeights,
+    /// Level 2: discard all GPU memory.
+    DiscardAll,
+}
+
+impl From<SleepLevel> for u32 {
+    fn from(level: SleepLevel) -> Self {
+        match level {
+            SleepLevel::PauseOnly => 0,
+            SleepLevel::OffloadWeights => 1,
+            SleepLevel::DiscardAll => 2,
+        }
+    }
+}
+
+impl TryFrom<u32> for SleepLevel {
+    type Error = String;
+
+    fn try_from(level: u32) -> std::result::Result<Self, Self::Error> {
+        match level {
+            0 => Ok(Self::PauseOnly),
+            1 => Ok(Self::OffloadWeights),
+            2 => Ok(Self::DiscardAll),
+            other => Err(format!(
+                "invalid sleep level `{other}`; expected one of: 0, 1, 2"
+            )),
+        }
+    }
+}
+
 /// Utility call id as carried on the engine-core MessagePack wire.
 ///
 /// Python emits utility ids as MessagePack integers, including values that may
diff --git a/rust/src/engine-core-client/src/test_utils.rs b/rust/src/engine-core-client/src/test_utils.rs
index 13d473f59790..a7d00cedbcff 100644
--- a/rust/src/engine-core-client/src/test_utils.rs
+++ b/rust/src/engine-core-client/src/test_utils.rs
@@ -13,7 +13,7 @@ use crate::EngineId;
 pub use crate::mock_engine::{MockCoordinatorSockets, MockEngineSockets};
 use crate::mock_engine::{
     MockEngineConfig, MockEngineDataSockets, connect_to_bootstrapped_frontend, connect_to_frontend,
-    default_ready_response,
+    connect_to_hybrid_frontends, default_ready_response,
 };
 use crate::protocol::handshake::{EngineCoreReadyResponse, HandshakeInitMessage};
 
@@ -85,6 +85,23 @@ pub async fn setup_mock_engine_sockets(
         .expect("connect mock engine")
 }
 
+/// Join a rank-0 frontend and the engine's own frontend like an engine on a
+/// node in hybrid or external DP LB.
+pub async fn setup_hybrid_mock_engine_sockets(
+    head_handshake: String,
+    client_handshake: String,
+    engine_id: impl Into<EngineId>,
+) -> MockEngineSockets {
+    connect_to_hybrid_frontends(
+        head_handshake,
+        client_handshake,
+        engine_id,
+        test_mock_engine_config(),
+    )
+    .await
+    .expect("connect hybrid mock engine")
+}
+
 /// Connect one mock engine directly to already-bootstrapped frontend
 /// input/output sockets.
 pub async fn setup_bootstrapped_mock_engine(
diff --git a/rust/src/engine-core-client/src/tests/client.rs b/rust/src/engine-core-client/src/tests/client.rs
index f483892d4c94..b2517fcd6e6a 100644
--- a/rust/src/engine-core-client/src/tests/client.rs
+++ b/rust/src/engine-core-client/src/tests/client.rs
@@ -14,6 +14,7 @@ use std::time::Duration;
 use futures::StreamExt;
 use rmpv::Value;
 use thiserror_ext::AsReport as _;
+use tokio::runtime::Builder;
 use tokio::sync::{mpsc, oneshot};
 use tokio::time::timeout;
 use tracing_subscriber::EnvFilter;
@@ -22,7 +23,11 @@ use zeromq::util::PeerIdentity;
 use zeromq::{DealerSocket, PushSocket, SocketOptions, SubSocket, XPubSocket, ZmqMessage};
 
 use crate::metrics::SchedulerStatsRecorder;
-use crate::protocol::handshake::{EngineCoreReadyResponse, HandshakeInitMessage, ReadyMessage};
+use crate::mock_engine::{MockEngineConfig, default_ready_response};
+use crate::protocol::handshake::{
+    EngineCoreReadyResponse, HandshakeInitMessage, HandshakeParallelConfig, PrefixHashConfig,
+    ReadyMessage,
+};
 use crate::protocol::kv_hints::{KvHintAction, KvHintsEnvelope};
 use crate::protocol::logprobs::MaybeWireLogprobs;
 use crate::protocol::lora::LoraRequest;
@@ -34,14 +39,15 @@ use crate::protocol::output::{
     DpControlMessage, DpControlOutput, EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs,
     RequestBatchOutputs, UtilityCallOutput, decode_engine_core_outputs,
 };
-use crate::protocol::request::{EngineCoreRequest, EngineCoreRequestType};
+use crate::protocol::request::{BlockHash, EngineCoreRequest, EngineCoreRequestType};
 use crate::protocol::sampling::EngineCoreSamplingParams;
-use crate::protocol::stats::{KvConnectorStats, MooncakeOperation, SchedulerStats};
+use crate::protocol::stats::{KvConnectorStats, MooncakeOperation, SchedulerStats, UmbpOperation};
 use crate::protocol::tensor::{WireArrayData, WireTensor};
-use crate::protocol::utility::{UtilityOutput, UtilityResultEnvelope};
+use crate::protocol::utility::{PauseMode, SleepLevel, UtilityOutput, UtilityResultEnvelope};
 use crate::test_utils::{
-    IpcNamespace, setup_bootstrapped_mock_engine, setup_mock_engine_sockets,
-    setup_mock_engine_with_init, spawn_mock_engine_task,
+    IpcNamespace, setup_bootstrapped_mock_engine, setup_hybrid_mock_engine_sockets,
+    setup_mock_engine_sockets, setup_mock_engine_with_init, spawn_mock_engine_task,
+    spawn_mock_engine_task_with_ready,
 };
 use crate::transport::ConnectedEngine;
 use crate::{
@@ -181,6 +187,7 @@ fn sample_request_with_id(request_id: &str) -> EngineCoreRequest {
                 payload: BTreeMap::from([("key".to_string(), serde_json::json!("value"))]),
             }],
         }),
+        block_hashes: Some(vec![BlockHash(vec![1; 32]), BlockHash(vec![2; 32])]),
         ..EngineCoreRequest::default()
     }
 }
@@ -270,6 +277,17 @@ async fn recv_start_dp_wave(sub: &mut SubSocket) -> (u32, Option<u32>) {
     rmp_serde::from_slice(&frames[1]).expect("decode START_DP_WAVE payload")
 }
 
+/// Receive one coordinator state publish and return its msgpack bytes as hex.
+async fn recv_coordinator_publish(sub: &mut SubSocket) -> String {
+    let frames = timeout(Duration::from_secs(2), sub.recv())
+        .await
+        .expect("coordinator state publish should arrive")
+        .unwrap()
+        .into_vec();
+    assert_eq!(frames.len(), 1);
+    hex::encode(&frames[0])
+}
+
 async fn send_start_wave(push: &mut PushSocket, engine_index: u32, wave: u32) {
     send_outputs(
         push,
@@ -309,7 +327,9 @@ fn handshake_test_config(
         transport_mode: TransportMode::HandshakeOwner {
             handshake_address,
             advertised_host: "127.0.0.1".to_string(),
+            engine_start_index: 0,
             engine_count,
+            data_parallel_size: engine_count,
             ready_timeout,
             local_input_address: None,
             local_output_address: None,
@@ -419,6 +439,67 @@ fn client_config_validates_bootstrapped_dp_range() {
     .assert_eq(&error.to_string());
 }
 
+#[test]
+fn client_transport_outlives_connecting_runtime() {
+    init_tracing();
+    let caller = Builder::new_current_thread().enable_all().build().unwrap();
+    let independent = Builder::new_multi_thread().worker_threads(2).enable_all().build().unwrap();
+    let ipc = IpcNamespace::new().unwrap();
+    let handshake_address = ipc.handshake_endpoint();
+    let (shutdown_tx, engine_task) = independent.block_on(async {
+        spawn_mock_engine_task(
+            handshake_address.clone(),
+            vec![0x00, 0x00],
+            |dealer, push| {
+                Box::pin(async move {
+                    let _request = recv_engine_message(dealer).await;
+                    send_outputs(
+                        push,
+                        RequestBatchOutputs {
+                            outputs: vec![request_output(
+                                "req-1",
+                                vec![42],
+                                Some(EngineCoreFinishReason::Length),
+                            )],
+                            ..Default::default()
+                        }
+                        .into(),
+                    )
+                    .await;
+                })
+            },
+        )
+    });
+    let client = caller.block_on(connect_client_with_ipc(
+        handshake_test_config(
+            handshake_address,
+            1,
+            "test-model",
+            Duration::from_secs(2),
+            0,
+            None,
+        ),
+        &ipc,
+    ));
+
+    caller.shutdown_timeout(Duration::from_secs(1));
+
+    let mut stream = independent
+        .block_on(client.call(sample_request()))
+        .expect("the client owns its transport runtime")
+        .into_outputs();
+    let output = independent
+        .block_on(async { timeout(Duration::from_secs(2), stream.next()).await })
+        .expect("engine output timed out")
+        .expect("engine output stream ended")
+        .expect("engine output failed");
+    assert_eq!(output.new_token_ids, [42]);
+
+    let _ = shutdown_tx.send(());
+    independent.block_on(engine_task).unwrap();
+    independent.block_on(client.shutdown()).unwrap();
+}
+
 async fn recv_xpub_message(xpub: &mut XPubSocket) -> Vec<bytes::Bytes> {
     xpub.recv().await.unwrap().into_vec()
 }
@@ -572,21 +653,23 @@ where
 async fn coordinator_handshake_includes_engine_control_addresses() {
     init_tracing();
     let ipc = IpcNamespace::new().unwrap();
-    let handshake_address = ipc.handshake_endpoint();
-    let engine_id = [0x00, 0x00];
-
-    let (init_tx, init_rx) = oneshot::channel();
-    let (shutdown_tx, shutdown_rx) = oneshot::channel();
-    let engine_task = tokio::spawn(async move {
-        let sockets = setup_mock_engine_sockets(handshake_address, &engine_id).await;
-        let _ = init_tx.send(sockets.init.clone());
-        let _ = shutdown_rx.await;
-    });
+    let mut engines = Vec::new();
+    for engine_id in [[0x00, 0x00], [0x01, 0x00]] {
+        let handshake_address = ipc.handshake_endpoint();
+        let (init_tx, init_rx) = oneshot::channel();
+        let (shutdown_tx, shutdown_rx) = oneshot::channel();
+        let engine_task = tokio::spawn(async move {
+            let sockets = setup_mock_engine_sockets(handshake_address, &engine_id).await;
+            let _ = init_tx.send(sockets.init.clone());
+            let _ = shutdown_rx.await;
+        });
+        engines.push((init_rx, shutdown_tx, engine_task));
+    }
 
     let client = connect_client_with_ipc(
         handshake_test_config(
             ipc.handshake_endpoint(),
-            1,
+            2,
             "test-model",
             Duration::from_secs(2),
             0,
@@ -596,13 +679,35 @@ async fn coordinator_handshake_includes_engine_control_addresses() {
     )
     .await;
 
-    let init = init_rx.await.unwrap();
+    let mut inits = Vec::new();
+    for (init_rx, shutdown_tx, engine_task) in engines {
+        inits.push(init_rx.await.unwrap());
+        let _ = shutdown_tx.send(());
+        engine_task.await.unwrap();
+    }
+    let init = &inits[0];
+    assert_eq!(init.parallel_config, inits[1].parallel_config);
     assert!(init.addresses.coordinator_input.is_some());
     assert!(init.addresses.coordinator_output.is_some());
-    assert!(init.addresses.frontend_stats_publish_address.is_none());
+    assert!(init.addresses.frontend_stats_publish_address.is_some());
+
+    // Coordinated DP sends one DP rendezvous to every engine: the handshake
+    // host plus five distinct ports.
+    let rendezvous = init
+        .parallel_config
+        .data_parallel_rendezvous
+        .as_ref()
+        .expect("coordinated DP sends the DP rendezvous");
+    assert_eq!(rendezvous.data_parallel_master_ip, "127.0.0.1");
+    assert_eq!(rendezvous.data_parallel_size, 2);
+    let mut ports = rendezvous.data_parallel_master_port_list.clone();
+    assert_eq!(ports.len(), 4);
+    ports.push(rendezvous.data_parallel_master_port);
+    ports.sort_unstable();
+    ports.dedup();
+    assert_eq!(ports.len(), 5);
+    assert!(!ports.contains(&0));
 
-    let _ = shutdown_tx.send(());
-    engine_task.await.unwrap();
     client.shutdown().await.unwrap();
 }
 
@@ -896,6 +1001,395 @@ async fn coordinator_rebroadcasts_engine_start_wave_control() {
     client.shutdown().await.unwrap();
 }
 
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn coordinator_publishes_wave_state_to_frontend_subscribers() {
+    init_tracing();
+    let ipc = IpcNamespace::new().unwrap();
+    let handshake_address = ipc.handshake_endpoint();
+
+    let (init_tx, init_rx) = oneshot::channel();
+    let (subscribed_tx, subscribed_rx) = oneshot::channel();
+    let (shutdown_tx, shutdown_rx) = oneshot::channel();
+    let engine_task = tokio::spawn(async move {
+        let mut engine = setup_mock_engine_sockets(handshake_address, &[0x00, 0x00]).await;
+        let mut coordinator =
+            engine.coordinator.take().expect("coordinator sockets should be present");
+
+        // Finish wave 0 and start wave 1 before any frontend subscribes. The
+        // `START_DP_WAVE` broadcast shows the coordinator processed both.
+        send_start_wave(&mut coordinator.output_push, 0, 0).await;
+        assert_eq!(
+            recv_start_dp_wave(&mut coordinator.input_sub).await,
+            (0, Some(0))
+        );
+        send_outputs(
+            &mut coordinator.output_push,
+            DpControlOutput {
+                engine_index: 0,
+                timestamp: 0.0,
+                control: DpControlMessage::WaveComplete(0),
+            }
+            .into(),
+        )
+        .await;
+        send_start_wave(&mut coordinator.output_push, 0, 1).await;
+        assert_eq!(
+            recv_start_dp_wave(&mut coordinator.input_sub).await,
+            (1, Some(0))
+        );
+        let _ = init_tx.send(engine.init.clone());
+
+        subscribed_rx.await.unwrap();
+        send_outputs(
+            &mut coordinator.output_push,
+            DpControlOutput {
+                engine_index: 0,
+                timestamp: 0.0,
+                control: DpControlMessage::WaveComplete(1),
+            }
+            .into(),
+        )
+        .await;
+
+        let _ = shutdown_rx.await;
+    });
+
+    let client = connect_client_with_ipc(
+        handshake_test_config(
+            ipc.handshake_endpoint(),
+            1,
+            "test-model",
+            Duration::from_secs(2),
+            0,
+            Some(CoordinatorMode::InProc),
+        ),
+        &ipc,
+    )
+    .await;
+
+    let init = init_rx.await.unwrap();
+    let publish_address = init
+        .addresses
+        .frontend_stats_publish_address
+        .expect("frontend publish address should be present");
+    let mut frontend = SubSocket::new();
+    frontend.connect(&publish_address).await.unwrap();
+    frontend.subscribe("").await.unwrap();
+
+    // The wave-only publish `(None, wave, engines_running)` as a msgpack
+    // array. A frontend subscribing mid-wave first receives the current state.
+    assert_eq!(recv_coordinator_publish(&mut frontend).await, "93c001c3");
+
+    // A wave transition still pending when the frontend subscribed may publish
+    // the same state again; the next distinct state is the one that follows.
+    subscribed_tx.send(()).unwrap();
+    let next_state = loop {
+        let state = recv_coordinator_publish(&mut frontend).await;
+        if state != "93c001c3" {
+            break state;
+        }
+    };
+    assert_eq!(next_state, "93c002c2");
+
+    let _ = shutdown_tx.send(());
+    engine_task.await.unwrap();
+    client.shutdown().await.unwrap();
+}
+
+/// Override the DP ranks a handshake-owned test config connects to.
+fn with_engine_range(
+    mut config: EngineCoreClientConfig,
+    start_index: u32,
+    dp_size: usize,
+) -> EngineCoreClientConfig {
+    let TransportMode::HandshakeOwner {
+        engine_start_index,
+        data_parallel_size,
+        ..
+    } = &mut config.transport_mode
+    else {
+        unreachable!("expected handshake-owned transport")
+    };
+    *engine_start_index = start_index;
+    *data_parallel_size = dp_size;
+    config
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
+async fn hybrid_lb_frontends_split_ranks_and_follow_rank_zero_coordinator() {
+    init_tracing();
+    let head_ipc = IpcNamespace::new().unwrap();
+    let node_ipc = IpcNamespace::new().unwrap();
+
+    // Engine 0 runs on the rank-0 node and joins only that node's frontend.
+    let (wave_done_tx, wave_done_rx) = oneshot::channel();
+    let (head_shutdown_tx, head_shutdown_rx) = oneshot::channel();
+    let engine0_task = tokio::spawn({
+        let head_handshake = head_ipc.handshake_endpoint();
+        async move {
+            let mut engine = setup_mock_engine_sockets(head_handshake, &[0x00, 0x00]).await;
+            let mut coordinator =
+                engine.coordinator.take().expect("coordinator sockets should be present");
+
+            // Woken by the wave that engine 1 starts, engine 0 later ends it.
+            assert_eq!(
+                recv_start_dp_wave(&mut coordinator.input_sub).await,
+                (0, Some(1))
+            );
+            send_outputs(
+                &mut coordinator.output_push,
+                DpControlOutput {
+                    engine_index: 0,
+                    timestamp: 0.0,
+                    control: DpControlMessage::WaveComplete(0),
+                }
+                .into(),
+            )
+            .await;
+            let _ = wave_done_tx.send(());
+            let _ = head_shutdown_rx.await;
+        }
+    });
+
+    // Engine 1 runs on the other node and joins both frontends.
+    let (init1_tx, init1_rx) = oneshot::channel();
+    let (wave_tx, mut wave_rx) = mpsc::unbounded_channel();
+    let (node_shutdown_tx, mut node_shutdown_rx) = oneshot::channel();
+    let engine1_task = tokio::spawn({
+        let head_handshake = head_ipc.handshake_endpoint();
+        let node_handshake = node_ipc.handshake_endpoint();
+        async move {
+            let mut engine =
+                setup_hybrid_mock_engine_sockets(head_handshake, node_handshake, &[0x01, 0x00])
+                    .await;
+            let _ = init1_tx.send(engine.init.clone());
+            let mut coordinator =
+                engine.coordinator.take().expect("coordinator sockets should be present");
+            send_start_wave(&mut coordinator.output_push, 1, 0).await;
+            assert_eq!(
+                recv_start_dp_wave(&mut coordinator.input_sub).await,
+                (0, Some(1))
+            );
+
+            // Report the wave stamped on every request until shutdown.
+            let data_socket = engine.data_sockets.first_mut().expect("data socket");
+            loop {
+                tokio::select! {
+                    message = recv_engine_message(&mut data_socket.dealer) => {
+                        if message[0].as_ref() == [0x00] {
+                            let request: EngineCoreRequest =
+                                rmp_serde::from_slice(&message[1]).unwrap();
+                            let _ = wave_tx.send(request.current_wave);
+                        }
+                    }
+                    _ = &mut node_shutdown_rx => break,
+                }
+            }
+        }
+    });
+
+    let head_config = with_engine_range(
+        handshake_test_config(
+            head_ipc.handshake_endpoint(),
+            1,
+            "test-model",
+            Duration::from_secs(2),
+            0,
+            Some(CoordinatorMode::InProc),
+        ),
+        0,
+        2,
+    );
+    let node_config = with_engine_range(
+        handshake_test_config(
+            node_ipc.handshake_endpoint(),
+            1,
+            "test-model",
+            Duration::from_secs(2),
+            0,
+            None,
+        ),
+        1,
+        2,
+    );
+    let (head, node) = tokio::join!(
+        connect_client_with_ipc(head_config, &head_ipc),
+        connect_client_with_ipc(node_config, &node_ipc),
+    );
+
+    // Rank 0 handshakes with both engines but connects only to its own.
+    assert_eq!(head.engine_identities(), vec![[0x00, 0x00].as_slice()]);
+    assert_eq!(node.engine_identities(), vec![[0x01, 0x00].as_slice()]);
+    assert_eq!(
+        (head.data_parallel_size(), node.data_parallel_size()),
+        (2, 2)
+    );
+
+    // Engine 1 takes its I/O from its own frontend and the DP rendezvous from
+    // rank 0.
+    let init1 = init1_rx.await.unwrap();
+    assert_eq!(init1.addresses.inputs, vec![node_ipc.input_endpoint()]);
+    assert_eq!(
+        init1
+            .parallel_config
+            .data_parallel_rendezvous
+            .map(|rendezvous| rendezvous.data_parallel_size),
+        Some(2)
+    );
+
+    // The rank-1 frontend comes to stamp its requests with the wave that rank
+    // 0's coordinator advanced to.
+    wave_done_rx.await.unwrap();
+    let mut streams = Vec::new();
+    timeout(Duration::from_secs(2), async {
+        for i in 0.. {
+            streams.push(node.call(sample_request_with_id(&format!("req-{i}"))).await.unwrap());
+            if wave_rx.recv().await.unwrap() == 1 {
+                break;
+            }
+            tokio::task::yield_now().await;
+        }
+    })
+    .await
+    .expect("the rank-1 frontend should follow the rank-0 wave");
+    drop(streams);
+
+    let _ = head_shutdown_tx.send(());
+    let _ = node_shutdown_tx.send(());
+    engine0_task.await.unwrap();
+    engine1_task.await.unwrap();
+    head.shutdown().await.unwrap();
+    node.shutdown().await.unwrap();
+}
+
+/// Connect a rank-0 frontend serving DP ranks [0, 1) of 2 while the given
+/// engines join it, and return the startup error.
+async fn rank_zero_startup_error(engines: Vec<([u8; 2], MockEngineConfig)>) -> String {
+    let ipc = IpcNamespace::new().unwrap();
+    let engine_tasks = engines
+        .into_iter()
+        .map(|(engine_id, config)| {
+            tokio::spawn(crate::mock_engine::connect_to_frontend(
+                ipc.handshake_endpoint(),
+                EngineId::from(&engine_id),
+                config,
+            ))
+        })
+        .collect::<Vec<_>>();
+    let config = with_engine_range(
+        handshake_test_config(
+            ipc.handshake_endpoint(),
+            1,
+            "test-model",
+            Duration::from_secs(2),
+            0,
+            None,
+        ),
+        0,
+        2,
+    );
+    let Err(error) = EngineCoreClient::connect(config.with_local_input_output_addresses(
+        Some(ipc.input_endpoint()),
+        Some(ipc.output_endpoint()),
+    ))
+    .await
+    else {
+        panic!("rank-0 startup should fail");
+    };
+    for task in engine_tasks {
+        task.abort();
+    }
+    error.to_report_string()
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn partitioned_handshake_rejects_unexpected_ranks_and_headless_remote_engines() {
+    init_tracing();
+    let headless = MockEngineConfig::default();
+    let hybrid = MockEngineConfig {
+        headless: false,
+        ..MockEngineConfig::default()
+    };
+
+    expect_test::expect!["unexpected startup handshake message: engine EngineId(0500) is not one of the expected DP ranks 0..2"].assert_eq(
+        &rank_zero_startup_error(vec![
+            ([0x00, 0x00], headless.clone()),
+            ([0x05, 0x00], hybrid),
+        ])
+        .await,
+    );
+    expect_test::expect!["unexpected startup handshake message: the engine of DP rank 1 joined headless; in hybrid or external DP LB every node runs its own frontend"].assert_eq(
+        &rank_zero_startup_error(vec![
+            ([0x00, 0x00], headless.clone()),
+            ([0x01, 0x00], headless),
+        ])
+        .await,
+    );
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn handshake_checks_parallel_config_hashes_only_for_coordinated_dp() {
+    init_tracing();
+
+    async fn connect_engines_with_different_hashes(
+        ipc: &IpcNamespace,
+        coordinator_mode: Option<CoordinatorMode>,
+    ) -> crate::Result<EngineCoreClient> {
+        let engine_tasks = ["a", "b"]
+            .into_iter()
+            .enumerate()
+            .map(|(index, hash)| {
+                tokio::spawn(crate::mock_engine::connect_to_frontend(
+                    ipc.handshake_endpoint(),
+                    EngineId::from(&[index as u8, 0x00]),
+                    MockEngineConfig {
+                        parallel_config_hash: Some(hash.to_string()),
+                        ..MockEngineConfig::default()
+                    },
+                ))
+            })
+            .collect::<Vec<_>>();
+        let result = EngineCoreClient::connect(
+            handshake_test_config(
+                ipc.handshake_endpoint(),
+                2,
+                "test-model",
+                Duration::from_secs(2),
+                0,
+                coordinator_mode,
+            )
+            .with_local_input_output_addresses(
+                Some(ipc.input_endpoint()),
+                Some(ipc.output_endpoint()),
+            ),
+        )
+        .await;
+        for task in engine_tasks {
+            task.abort();
+        }
+        result
+    }
+
+    // Dense DP ranks share no collectives, so their configs may differ.
+    let ipc = IpcNamespace::new().unwrap();
+    let client = connect_engines_with_different_hashes(&ipc, None).await.unwrap();
+    client.shutdown().await.unwrap();
+
+    let ipc = IpcNamespace::new().unwrap();
+    let Err(error) =
+        connect_engines_with_different_hashes(&ipc, Some(CoordinatorMode::InProc)).await
+    else {
+        panic!("coordinated DP startup should fail");
+    };
+    assert!(
+        error
+            .to_report_string()
+            .contains("launch every DP rank with the same parallel config"),
+        "{}",
+        error.to_report_string()
+    );
+}
+
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 async fn coordinator_accepts_stats_only_outputs() {
     init_tracing();
@@ -1602,18 +2096,38 @@ async fn call_utility_waits_for_all_engines_before_returning_error() {
 }
 
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
-async fn cancelled_utility_call_unregisters_waiter() {
+async fn cancelled_utility_call_still_collects_the_reply() {
     init_tracing();
     let ipc = IpcNamespace::new().unwrap();
     let handshake_address = ipc.handshake_endpoint();
     let (received_tx, received_rx) = oneshot::channel();
+    let (cancelled_tx, cancelled_rx) = oneshot::channel::<()>();
     let (_shutdown, engine_task) = spawn_mock_engine_task(
         handshake_address.clone(),
         vec![0x00, 0x00],
-        |dealer, _push| {
+        |dealer, push| {
             Box::pin(async move {
-                let _utility = recv_engine_message(dealer).await;
+                let utility = recv_engine_message(dealer).await;
+                let payload = decode_value(&utility[1]);
+                let call_id = payload.as_array().expect("utility payload array")[1]
+                    .as_u64()
+                    .expect("call_id");
                 let _ = received_tx.send(());
+                cancelled_rx.await.unwrap();
+                send_outputs(
+                    push,
+                    UtilityCallOutput {
+                        engine_index: 0,
+                        timestamp: 0.0,
+                        output: UtilityOutput {
+                            call_id: call_id.into(),
+                            failure_message: None,
+                            result: Some(utility_result_value(true)),
+                        },
+                    }
+                    .into(),
+                )
+                .await;
                 std::future::pending::<()>().await;
             })
         },
@@ -1637,10 +2151,18 @@ async fn cancelled_utility_call_unregisters_waiter() {
         tokio::spawn(async move { call_client.call_utility::<bool, _>("add_lora", ()).await });
 
     received_rx.await.unwrap();
-    assert_eq!(client.pending_utility_call_count(), 1);
     call.abort();
     assert!(call.await.unwrap_err().is_cancelled());
-    assert_eq!(client.pending_utility_call_count(), 0);
+    // The call keeps waiting for the engine after its caller is gone.
+    assert_eq!(client.pending_utility_call_count(), 1);
+    cancelled_tx.send(()).unwrap();
+    timeout(Duration::from_secs(2), async {
+        while client.pending_utility_call_count() > 0 {
+            tokio::task::yield_now().await;
+        }
+    })
+    .await
+    .expect("the reply resolves the waiter");
 
     engine_task.abort();
     let client = std::sync::Arc::try_unwrap(client)
@@ -1906,6 +2428,90 @@ async fn client_decodes_multipart_logprob_outputs() {
     client.shutdown().await.unwrap();
 }
 
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn client_sends_prompt_block_hashes_for_engine_prefix_hash_config() {
+    init_tracing();
+    let ipc = IpcNamespace::new().unwrap();
+    let handshake_address = ipc.handshake_endpoint();
+    // The engine's NONE_HASH for `sha256` with the default seed, and the
+    // hashes the Python engine computes for this prompt.
+    let ready_response = EngineCoreReadyResponse {
+        prefix_hash_config: Some(PrefixHashConfig {
+            hash_algo: "sha256".to_string(),
+            hash_block_size: 16,
+            none_hash: BlockHash(
+                hex::decode("066b1e8987e027138e98701c46128929378935f9b6cd1549105ae6b4f0d71198")
+                    .unwrap(),
+            ),
+            num_lookahead_tokens: 0,
+        }),
+        ..default_ready_response()
+    };
+    let edge_tokens = [0, 1, 23, 24, 255, 256, 65535, 65536, 151_645, 0x7fff_ffff];
+    let request = EngineCoreRequest {
+        prompt_token_ids: Some([edge_tokens.repeat(3), vec![7, 8]].concat()),
+        block_hashes: None,
+        ..sample_request()
+    };
+
+    let (shutdown_tx, engine_task) = spawn_mock_engine_task_with_ready(
+        handshake_address.clone(),
+        b"engine-block-hashes".to_vec(),
+        ready_response,
+        move |dealer, push| {
+            Box::pin(async move {
+                let add = recv_engine_message(dealer).await;
+                let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap();
+                let hashes = request.block_hashes.expect("prompt block hashes");
+                let hashes = hashes.iter().map(|hash| hex::encode(&hash.0)).collect::<Vec<_>>();
+                assert_eq!(
+                    hashes,
+                    [
+                        "00e88e4087ec88315eb29afe9c001411b274ead7ffbc56e35b824d2c8817b280",
+                        "70de07ec6eed1e78ee32aa5f2c17e7eddda8068507c29e9562b43e0f0a872fe0",
+                    ]
+                );
+
+                send_outputs(
+                    push,
+                    RequestBatchOutputs {
+                        outputs: vec![request_output(
+                            "req-1",
+                            vec![],
+                            Some(EngineCoreFinishReason::Length),
+                        )],
+                        finished_requests: Some(BTreeSet::from(["req-1".to_string()])),
+                        ..Default::default()
+                    }
+                    .into(),
+                )
+                .await;
+            })
+        },
+    );
+
+    let client = connect_client_with_ipc(
+        handshake_test_config(
+            handshake_address,
+            1,
+            "test-model",
+            Duration::from_secs(2),
+            0,
+            None,
+        ),
+        &ipc,
+    )
+    .await;
+
+    let outputs = client.call(request).await.unwrap().into_outputs().collect::<Vec<_>>().await;
+    assert_eq!(outputs.len(), 1);
+    assert!(outputs[0].is_ok());
+
+    let _ = shutdown_tx.send(());
+    engine_task.await.unwrap();
+    client.shutdown().await.unwrap();
+}
+
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 async fn client_sends_large_multimodal_tensor_as_aux_frame() {
     init_tracing();
@@ -2095,6 +2701,9 @@ async fn multi_engine_client_shares_transport_and_routes_by_inflight_count() {
     assert_eq!(init_1.addresses.inputs, vec![ipc.input_endpoint()]);
     assert_eq!(init_0.addresses.outputs, vec![ipc.output_endpoint()]);
     assert_eq!(init_1.addresses.outputs, vec![ipc.output_endpoint()]);
+    // Without coordinated DP, engines keep their own rendezvous config.
+    assert_eq!(init_0.parallel_config, HandshakeParallelConfig::default());
+    assert_eq!(init_1.parallel_config, HandshakeParallelConfig::default());
 
     assert_eq!(client.input_address(), ipc.input_endpoint());
     assert_eq!(client.output_address(), ipc.output_endpoint());
@@ -2416,6 +3025,383 @@ async fn wake_up_requires_every_engine_to_be_fully_awake() {
     client.shutdown().await.unwrap();
 }
 
+/// Serve utility calls as an engine that is ready unless paused or in a weight
+/// update, in whatever order the client sends them, until `stop` fires.
+async fn serve_readiness_engine(
+    dealer: &mut DealerSocket,
+    push: &mut PushSocket,
+    mut paused: bool,
+    mut stop: oneshot::Receiver<()>,
+) {
+    let mut updating_weights = false;
+    loop {
+        let utility = tokio::select! {
+            utility = recv_engine_message(dealer) => utility,
+            _ = &mut stop => return,
+        };
+        assert_eq!(utility[0].as_ref(), &[0x03]);
+        let payload = decode_value(&utility[1]);
+        let array = payload.as_array().expect("utility payload array");
+        let call_id = array[1].as_u64().expect("call_id");
+        let result = match array[2].as_str().expect("utility method") {
+            "is_inference_ready" => utility_result_value(!paused && !updating_weights),
+            "pause_generation" => {
+                paused = true;
+                utility_result_value(())
+            }
+            "resume_scheduler" | "checkpoint_restore" => {
+                paused = false;
+                utility_result_value(())
+            }
+            "collective_rpc" => {
+                match array[3][0].as_str().expect("worker method") {
+                    "start_weight_update" => updating_weights = true,
+                    "finish_weight_update" => updating_weights = false,
+                    method => panic!("unexpected worker method {method}"),
+                }
+                utility_result_value([()])
+            }
+            method => panic!("unexpected utility call {method}"),
+        };
+        send_outputs(
+            push,
+            UtilityCallOutput {
+                engine_index: 0,
+                timestamp: 0.0,
+                output: UtilityOutput {
+                    call_id: call_id.into(),
+                    failure_message: None,
+                    result: Some(result),
+                },
+            }
+            .into(),
+        )
+        .await;
+    }
+}
+
+/// Connect a client to [`serve_readiness_engine`], with `enable_sleep_mode`
+/// set on the engine when it starts paused.
+async fn connect_readiness_engine(
+    engine_id: &'static [u8],
+    start_paused: bool,
+) -> (
+    EngineCoreClient,
+    IpcNamespace,
+    [oneshot::Sender<()>; 2],
+    tokio::task::JoinHandle<()>,
+) {
+    init_tracing();
+    let ipc = IpcNamespace::new().unwrap();
+    let handshake_address = ipc.handshake_endpoint();
+    let mut ready_response = crate::mock_engine::default_ready_response();
+    ready_response.enable_sleep_mode = start_paused;
+    let (stop_tx, stop_rx) = oneshot::channel();
+    let (shutdown_tx, engine_task) = spawn_mock_engine_task_with_ready(
+        handshake_address.clone(),
+        engine_id.to_vec(),
+        ready_response,
+        move |dealer, push| Box::pin(serve_readiness_engine(dealer, push, start_paused, stop_rx)),
+    );
+    let client = connect_client_with_ipc(
+        handshake_test_config(
+            handshake_address,
+            1,
+            "test-model",
+            Duration::from_secs(2),
+            0,
+            None,
+        ),
+        &ipc,
+    )
+    .await;
+    (client, ipc, [stop_tx, shutdown_tx], engine_task)
+}
+
+async fn shutdown_readiness_engine(
+    client: EngineCoreClient,
+    stop: [oneshot::Sender<()>; 2],
+    engine_task: tokio::task::JoinHandle<()>,
+) {
+    for tx in stop {
+        let _ = tx.send(());
+    }
+    engine_task.await.unwrap();
+    client.shutdown().await.unwrap();
+}
+
+async fn wait_until_ready(readiness: &mut tokio::sync::watch::Receiver<bool>) {
+    timeout(Duration::from_secs(2), readiness.wait_for(|ready| *ready))
+        .await
+        .expect("the engine reports ready")
+        .unwrap();
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn readiness_follows_engine_answers_around_pause() {
+    let (client, _ipc, stop, engine_task) =
+        connect_readiness_engine(b"engine-readiness-pause", false).await;
+    let mut readiness = client.subscribe_readiness();
+    assert!(*readiness.borrow());
+
+    client.pause_generation(PauseMode::Abort, false).await.unwrap();
+    assert!(!*readiness.borrow());
+    client.resume_scheduler().await.unwrap();
+    wait_until_ready(&mut readiness).await;
+
+    shutdown_readiness_engine(client, stop, engine_task).await;
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn readiness_follows_engine_answers_around_weight_update() {
+    let (client, _ipc, stop, engine_task) =
+        connect_readiness_engine(b"engine-readiness-weights", false).await;
+    let mut readiness = client.subscribe_readiness();
+
+    client.start_weight_update().await.unwrap();
+    assert!(!*readiness.borrow());
+    client.finish_weight_update().await.unwrap();
+    wait_until_ready(&mut readiness).await;
+
+    shutdown_readiness_engine(client, stop, engine_task).await;
+}
+
+/// Receive the next utility call, which must be `method`, and return its id.
+async fn recv_utility_call(dealer: &mut DealerSocket, method: &str) -> u64 {
+    let utility = recv_engine_message(dealer).await;
+    assert_eq!(utility[0].as_ref(), &[0x03]);
+    let payload = decode_value(&utility[1]);
+    let array = payload.as_array().expect("utility payload array");
+    assert_eq!(array[2], Value::from(method));
+    array[1].as_u64().expect("call_id")
+}
+
+async fn reply_utility(push: &mut PushSocket, call_id: u64, result: UtilityResultEnvelope) {
+    send_outputs(
+        push,
+        UtilityCallOutput {
+            engine_index: 0,
+            timestamp: 0.0,
+            output: UtilityOutput {
+                call_id: call_id.into(),
+                failure_message: None,
+                result: Some(result),
+            },
+        }
+        .into(),
+    )
+    .await;
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn cancelled_pause_rechecks_only_after_the_engine_replies() {
+    init_tracing();
+    let ipc = IpcNamespace::new().unwrap();
+    let handshake_address = ipc.handshake_endpoint();
+    let (pause_seen_tx, pause_seen_rx) = oneshot::channel();
+    let (reply_tx, reply_rx) = oneshot::channel::<()>();
+    let (shutdown_tx, engine_task) = spawn_mock_engine_task(
+        handshake_address.clone(),
+        b"engine-readiness-cancel".to_vec(),
+        move |dealer, push| {
+            Box::pin(async move {
+                let pause = recv_utility_call(dealer, "pause_generation").await;
+                pause_seen_tx.send(()).unwrap();
+                reply_rx.await.unwrap();
+                reply_utility(push, pause, utility_result_value(())).await;
+                let ask = recv_utility_call(dealer, "is_inference_ready").await;
+                reply_utility(push, ask, utility_result_value(false)).await;
+            })
+        },
+    );
+    let client = connect_client_with_ipc(
+        handshake_test_config(
+            handshake_address,
+            1,
+            "test-model",
+            Duration::from_secs(2),
+            0,
+            None,
+        ),
+        &ipc,
+    )
+    .await;
+    let readiness = client.subscribe_readiness();
+
+    tokio::select! {
+        _ = client.pause_generation(PauseMode::Wait, false) => {
+            panic!("the pause returned before the engine replied")
+        }
+        _ = pause_seen_rx => {}
+    }
+    // The caller is gone, but the call still waits for the engine.
+    assert!(!client.readiness_recheck_pending());
+    assert!(!*readiness.borrow());
+    reply_tx.send(()).unwrap();
+
+    let _ = shutdown_tx.send(());
+    engine_task.await.unwrap();
+    assert!(!*readiness.borrow());
+    client.shutdown().await.unwrap();
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn readiness_starts_from_engine_answer_with_sleep_mode() {
+    // Prepared for a checkpoint on startup.
+    let (client, _ipc, stop, engine_task) =
+        connect_readiness_engine(b"engine-readiness-startup", true).await;
+    let mut readiness = client.subscribe_readiness();
+    assert!(!*readiness.borrow());
+
+    client.checkpoint_restore().await.unwrap();
+    wait_until_ready(&mut readiness).await;
+
+    shutdown_readiness_engine(client, stop, engine_task).await;
+}
+
+/// Render the `vllm:engine_sleep_state` series of one model as
+/// `awake,weights_offloaded,discard_all` for engines 0 and 1.
+fn engine_sleep_state_series(model_name: &str) -> [String; 2] {
+    let rendered = vllm_metrics::METRICS.render().unwrap();
+    [0, 1].map(|engine| {
+        ["awake", "weights_offloaded", "discard_all"]
+            .map(|state| {
+                let prefix = format!(
+                    "vllm:engine_sleep_state{{model_name=\"{model_name}\",engine=\"{engine}\",sleep_state=\"{state}\"}} "
+                );
+                rendered
+                    .lines()
+                    .find_map(|line| line.strip_prefix(&prefix))
+                    .unwrap_or("missing")
+                    .to_owned()
+            })
+            .join(",")
+    })
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn sleep_and_wake_calls_record_engine_sleep_state() {
+    init_tracing();
+    // A model name of its own keeps other tests' clients off these series.
+    const MODEL: &str = "engine-sleep-state-model";
+    let ipc = IpcNamespace::new().unwrap();
+    let handshake_address = ipc.handshake_endpoint();
+    let mut engines = Vec::new();
+    for engine_index in 0..2_u16 {
+        // Each engine answers the calls in order; `None` stands for a `()`
+        // result. On the partial wake-up only engine 1 reports fully awake.
+        let answers = [
+            ("sleep", None),
+            ("wake_up", Some(engine_index == 1)),
+            ("wake_up", Some(true)),
+            ("release_kv_cache_memory", None),
+            ("wake_up", Some(true)),
+            ("sleep", None),
+        ];
+        engines.push(spawn_mock_engine_task(
+            handshake_address.clone(),
+            EngineId::from_engine_index(engine_index).into_frame().to_vec(),
+            move |dealer, push| {
+                Box::pin(async move {
+                    for (method, fully_awake) in answers {
+                        let utility = recv_engine_message(dealer).await;
+                        let payload = decode_value(&utility[1]);
+                        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(method));
+                        send_outputs(
+                            push,
+                            UtilityCallOutput {
+                                engine_index: u32::from(engine_index),
+                                timestamp: 0.0,
+                                output: UtilityOutput {
+                                    call_id: call_id.into(),
+                                    failure_message: None,
+                                    result: fully_awake.map(utility_result_value),
+                                },
+                            }
+                            .into(),
+                        )
+                        .await;
+                    }
+                })
+            },
+        ));
+    }
+    let client = connect_client_with_ipc(
+        handshake_test_config(handshake_address, 2, MODEL, Duration::from_secs(2), 5, None),
+        &ipc,
+    )
+    .await;
+
+    let mut series = vec![("connected", engine_sleep_state_series(MODEL))];
+    client.sleep(SleepLevel::OffloadWeights, PauseMode::default()).await.unwrap();
+    series.push(("sleep level 1", engine_sleep_state_series(MODEL)));
+    assert!(!client.wake_up(Some(vec!["weights".to_owned()])).await.unwrap());
+    series.push(("partial wake", engine_sleep_state_series(MODEL)));
+    assert!(client.wake_up(None).await.unwrap());
+    series.push(("full wake", engine_sleep_state_series(MODEL)));
+    client.release_kv_cache_memory().await.unwrap();
+    series.push(("release kv cache", engine_sleep_state_series(MODEL)));
+    assert!(client.wake_up(None).await.unwrap());
+    client.sleep(SleepLevel::DiscardAll, PauseMode::default()).await.unwrap();
+    series.push(("sleep level 2", engine_sleep_state_series(MODEL)));
+
+    expect_test::expect![[r#"
+        [
+            (
+                "connected",
+                [
+                    "1,0,0",
+                    "1,0,0",
+                ],
+            ),
+            (
+                "sleep level 1",
+                [
+                    "0,1,0",
+                    "0,1,0",
+                ],
+            ),
+            (
+                "partial wake",
+                [
+                    "0,1,0",
+                    "0,1,0",
+                ],
+            ),
+            (
+                "full wake",
+                [
+                    "1,0,0",
+                    "1,0,0",
+                ],
+            ),
+            (
+                "release kv cache",
+                [
+                    "0,0,0",
+                    "0,0,0",
+                ],
+            ),
+            (
+                "sleep level 2",
+                [
+                    "0,0,1",
+                    "0,0,1",
+                ],
+            ),
+        ]
+    "#]]
+    .assert_debug_eq(&series);
+
+    for (shutdown_tx, engine_task) in engines {
+        let _ = shutdown_tx.send(());
+        engine_task.await.unwrap();
+    }
+    client.shutdown().await.unwrap();
+}
+
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 async fn collective_rpc_flattens_results_from_all_engines() {
     init_tracing();
@@ -2779,6 +3765,7 @@ fn python_msgpack_fixtures_match_rust_encoding_and_record_nixl_metrics() {
     let mooncake_stats_hex = lines.next().expect("missing Mooncake stats fixture line");
     let mooncake_transfer_stats_hex =
         lines.next().expect("missing Mooncake transfer stats fixture line");
+    let umbp_stats_hex = lines.next().expect("missing UMBP stats fixture line");
     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");
@@ -3006,6 +3993,13 @@ fn python_msgpack_fixtures_match_rust_encoding_and_record_nixl_metrics() {
         KvConnectorStats::MooncakeTransfer(stats) if stats.num_failed_transfers == [1]
     ));
 
+    let umbp_stats: KvConnectorStats =
+        rmp_serde::from_slice(&hex::decode(umbp_stats_hex).unwrap()).unwrap();
+    assert!(matches!(
+        umbp_stats,
+        KvConnectorStats::Umbp(stats) if stats.0[&UmbpOperation::Store].num_bytes == 4096
+    ));
+
     let multi_connector_stats: KvConnectorStats =
         rmp_serde::from_slice(&hex::decode(multi_connector_stats_hex).unwrap()).unwrap();
     let KvConnectorStats::Multi(multi_connector_stats) = multi_connector_stats else {
@@ -3014,6 +4008,7 @@ fn python_msgpack_fixtures_match_rust_encoding_and_record_nixl_metrics() {
     assert!(multi_connector_stats.nixl.is_some());
     assert!(multi_connector_stats.mooncake_transfer.is_some());
     assert!(multi_connector_stats.mooncake.is_some());
+    assert!(multi_connector_stats.umbp.is_some());
     assert!(multi_connector_stats.other.contains_key("UnsupportedConnector"));
     let nixl_stats = multi_connector_stats.nixl.as_ref().expect("NIXL stats decode");
     assert_eq!(nixl_stats.num_failed_handshakes, [1]);
@@ -3060,6 +4055,10 @@ fn python_msgpack_fixtures_match_rust_encoding_and_record_nixl_metrics() {
     assert!(rendered_metrics.contains(
         "vllm:nixl_num_notifications_after_expiry_total{model_name=\"test-model\",engine=\"0\"} 2"
     ));
+    assert!(rendered_metrics.contains(
+        "vllm:umbp_transfer_bytes_total{model_name=\"test-model\",engine=\"0\",\
+         operation=\"load\"} 8192"
+    ));
 
     let map_keys = |bytes: &[u8]| -> BTreeSet<String> {
         match decode_value(bytes) {
@@ -3103,6 +4102,15 @@ fn python_msgpack_fixtures_match_rust_encoding_and_record_nixl_metrics() {
             "speculative_config": null,
         }))
     );
+    assert_eq!(
+        ready_response.prefix_hash_config,
+        Some(PrefixHashConfig {
+            hash_algo: "sha256".to_string(),
+            hash_block_size: 16,
+            none_hash: BlockHash((0..32).collect()),
+            num_lookahead_tokens: 1,
+        })
+    );
     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/prompt_block_hash_vectors.json b/rust/src/engine-core-client/src/tests/prompt_block_hash_vectors.json
new file mode 100644
index 000000000000..132d727f5073
--- /dev/null
+++ b/rust/src/engine-core-client/src/tests/prompt_block_hash_vectors.json
@@ -0,0 +1,578 @@
+{
+  "none_hashes": {
+    "sha256": "066b1e8987e027138e98701c46128929378935f9b6cd1549105ae6b4f0d71198",
+    "sha256_cbor": "9bd96a485ad84efdafb72ee48a1d7a69bcead0f8f0433173941b276b9581eef0"
+  },
+  "cases": [
+    {
+      "name": "plain_block_size_16",
+      "hash_algo": "sha256",
+      "block_size": 16,
+      "num_lookahead_tokens": 0,
+      "lora": null,
+      "cache_salt": null,
+      "prompt_token_ids": [0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 7, 8],
+      "block_hashes": [
+        "00e88e4087ec88315eb29afe9c001411b274ead7ffbc56e35b824d2c8817b280",
+        "70de07ec6eed1e78ee32aa5f2c17e7eddda8068507c29e9562b43e0f0a872fe0"
+      ]
+    },
+    {
+      "name": "plain_block_size_1",
+      "hash_algo": "sha256",
+      "block_size": 1,
+      "num_lookahead_tokens": 0,
+      "lora": null,
+      "cache_salt": null,
+      "prompt_token_ids": [0, 300, 70000],
+      "block_hashes": [
+        "1876a790824e1c62cc68ce0472a1542ccc7bf2d1d5c9e53827fc516bb9f88601",
+        "e66ffb905fd6547fc758c2df71343e839aa7a51f71ad3c7b0770016a493ed7cd",
+        "e4b0a447a032eac68f97ed77776d9f506446b12d46ec513801b12717ecaf3c8c"
+      ]
+    },
+    {
+      "name": "plain_block_size_3",
+      "hash_algo": "sha256",
+      "block_size": 3,
+      "num_lookahead_tokens": 0,
+      "lora": null,
+      "cache_salt": null,
+      "prompt_token_ids": [0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647],
+      "block_hashes": [
+        "588c1268ece56d895e2730c39f59ae5f370f7ea121e4fdaa9569c7bca0ea6a00",
+        "27f3c01741e55cd9536e7db68dfba985c5d39f267f66081f50f4de486588e20c",
+        "bcfc515ddd274ab841f7c2a32d5b90a6dc4efb881d1c237b89437edafaedb691"
+      ]
+    },
+    {
+      "name": "plain_block_size_24",
+      "hash_algo": "sha256",
+      "block_size": 24,
+      "num_lookahead_tokens": 0,
+      "lora": null,
+      "cache_salt": null,
+      "prompt_token_ids": [100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111, 112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 127, 128, 129, 130, 131, 132, 133, 134, 135, 136, 137, 138, 139, 140, 141, 142, 143, 144, 145, 146, 147, 148, 149],
+      "block_hashes": [
+        "160d21b10f24f083b8b07babc68f413df4dbedaaf45c1c171057b12b5f6599c7",
+        "823d4823d3b4dcc76fe5735ec2cf07ab4d3030d095050db13e783bfb69d7a742"
+      ]
+    },
+    {
+      "name": "plain_block_size_300",
+      "hash_algo": "sha256",
+      "block_size": 300,
+      "num_lookahead_tokens": 0,
+      "lora": null,
+      "cache_salt": null,
+      "prompt_token_ids": [0, 7919, 15838, 23757, 31676, 39595, 47514, 55433, 63352, 71271, 79190, 87109, 95028, 102947, 110866, 118785, 126704, 134623, 142542, 150461, 158380, 166299, 174218, 182137, 190056, 197975, 5891, 13810, 21729, 29648, 37567, 45486, 53405, 61324, 69243, 77162, 85081, 93000, 100919, 108838, 116757, 124676, 132595, 140514, 148433, 156352, 164271, 172190, 180109, 188028, 195947, 3863, 11782, 19701, 27620, 35539, 43458, 51377, 59296, 67215, 75134, 83053, 90972, 98891, 106810, 114729, 122648, 130567, 138486, 146405, 154324, 162243, 170162, 178081, 186000, 193919, 1835, 9754, 17673, 25592, 33511, 41430, 49349, 57268, 65187, 73106, 81025, 88944, 96863, 104782, 112701, 120620, 128539, 136458, 144377, 152296, 160215, 168134, 176053, 183972, 191891, 199810, 7726, 15645, 23564, 31483, 39402, 47321, 55240, 63159, 71078, 78997, 86916, 94835, 102754, 110673, 118592, 126511, 134430, 142349, 150268, 158187, 166106, 174025, 181944, 189863, 197782, 5698, 13617, 21536, 29455, 37374, 45293, 53212, 61131, 69050, 76969, 84888, 92807, 100726, 108645, 116564, 124483, 132402, 140321, 148240, 156159, 164078, 171997, 179916, 187835, 195754, 3670, 11589, 19508, 27427, 35346, 43265, 51184, 59103, 67022, 74941, 82860, 90779, 98698, 106617, 114536, 122455, 130374, 138293, 146212, 154131, 162050, 169969, 177888, 185807, 193726, 1642, 9561, 17480, 25399, 33318, 41237, 49156, 57075, 64994, 72913, 80832, 88751, 96670, 104589, 112508, 120427, 128346, 136265, 144184, 152103, 160022, 167941, 175860, 183779, 191698, 199617, 7533, 15452, 23371, 31290, 39209, 47128, 55047, 62966, 70885, 78804, 86723, 94642, 102561, 110480, 118399, 126318, 134237, 142156, 150075, 157994, 165913, 173832, 181751, 189670, 197589, 5505, 13424, 21343, 29262, 37181, 45100, 53019, 60938, 68857, 76776, 84695, 92614, 100533, 108452, 116371, 124290, 132209, 140128, 148047, 155966, 163885, 171804, 179723, 187642, 195561, 3477, 11396, 19315, 27234, 35153, 43072, 50991, 58910, 66829, 74748, 82667, 90586, 98505, 106424, 114343, 122262, 130181, 138100, 146019, 153938, 161857, 169776, 177695, 185614, 193533, 1449, 9368, 17287, 25206, 33125, 41044, 48963, 56882, 64801, 72720, 80639, 88558, 96477, 104396, 112315, 120234, 128153, 136072, 143991, 151910, 159829, 167748, 175667, 183586, 191505, 199424, 7340, 15259, 23178, 31097, 39016, 46935, 54854, 62773, 70692, 78611, 86530, 94449, 102368, 110287, 118206, 126125, 134044, 141963, 149882, 157801, 165720, 173639, 181558, 189477, 197396, 5312, 13231, 21150, 29069, 36988, 44907, 52826, 60745, 68664, 76583, 84502, 92421, 100340, 108259, 116178, 124097, 132016, 139935, 147854, 155773, 163692, 171611, 179530, 187449, 195368, 3284, 11203, 19122, 27041, 34960, 42879, 50798, 58717, 66636, 74555, 82474, 90393, 98312, 106231, 114150, 122069, 129988, 137907, 145826, 153745, 161664, 169583, 177502, 185421, 193340, 1256, 9175, 17094, 25013, 32932, 40851, 48770, 56689, 64608, 72527, 80446, 88365, 96284, 104203, 112122, 120041, 127960, 135879, 143798, 151717, 159636, 167555, 175474, 183393, 191312, 199231, 7147, 15066, 22985, 30904, 38823, 46742, 54661, 62580, 70499, 78418, 86337, 94256, 102175, 110094, 118013, 125932, 133851, 141770, 149689, 157608, 165527, 173446, 181365, 189284, 197203, 5119, 13038, 20957, 28876, 36795, 44714, 52633, 60552, 68471, 76390, 84309, 92228, 100147, 108066, 115985, 123904, 131823, 139742, 147661, 155580, 163499, 171418, 179337, 187256, 195175, 3091, 11010, 18929, 26848, 34767, 42686, 50605, 58524, 66443, 74362, 82281, 90200, 98119, 106038, 113957, 121876, 129795, 137714, 145633, 153552, 161471, 169390, 177309, 185228, 193147, 1063, 8982, 16901, 24820, 32739, 40658, 48577, 56496, 64415, 72334, 80253, 88172, 96091, 104010, 111929, 119848, 127767, 135686, 143605, 151524, 159443, 167362, 175281, 183200, 191119, 199038, 6954, 14873, 22792, 30711, 38630, 46549, 54468, 62387, 70306, 78225, 86144, 94063, 101982, 109901, 117820, 125739, 133658, 141577, 149496, 157415, 165334, 173253, 181172, 189091, 197010, 4926, 12845, 20764, 28683, 36602, 44521, 52440, 60359, 68278, 76197, 84116, 92035, 99954, 107873, 115792, 123711, 131630, 139549, 147468, 155387, 163306, 171225, 179144, 187063, 194982, 2898, 10817, 18736, 26655, 34574, 42493, 50412, 58331, 66250, 74169, 82088, 90007, 97926, 105845, 113764, 121683, 129602, 137521, 145440, 153359, 161278, 169197, 177116, 185035, 192954, 870, 8789, 16708, 24627, 32546, 40465, 48384, 56303, 64222, 72141, 80060, 87979, 95898, 103817, 111736, 119655, 127574, 135493, 143412, 151331, 159250, 167169, 175088, 183007, 190926, 198845, 6761, 14680, 22599, 30518, 38437, 46356, 54275, 62194, 70113, 78032, 85951, 93870, 101789, 109708, 117627, 125546, 133465, 141384, 149303, 157222, 165141, 173060, 180979, 188898, 196817, 4733, 12652, 20571, 28490, 36409, 44328, 52247, 60166, 68085, 76004, 83923, 91842, 99761, 107680, 115599, 123518, 131437, 139356],
+      "block_hashes": [
+        "60fc9719c922de9484e11daebe846969dde25dbbfc8975ad9d5a5f52f5be7cd2",
+        "a4fb5f1cf7802dcd1ed08aa20ff6f20ce522547ed9931aca73bd94fcb5f2aa2f"
+      ]
+    },
+    {
+      "name": "plain_lookahead_0",
+      "hash_algo": "sha256",
+      "block_size": 5,
+      "num_lookahead_tokens": 0,
+      "lora": null,
+      "cache_salt": null,
+      "prompt_token_ids": [0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 7, 8],
+      "block_hashes": [
+        "56b141f4a969e34bcc1da4ad6cb476ebcb2b7ae507e856a475f0b53f87859e20",
+        "6a9284439d12381064bae6a95ee7d5cff7101c5d6674fbf49ba3833ef734a750",
+        "70dfd10d0f29ac5d07fd98de759c7cca739f915ca164e87eedbb6531f8ccccd0",
+        "edf42ff3730eeeeb310cb85d85779e5e5470500cc3dc7f9656718cca82d655e7",
+        "e7c2c67f6e954bd1d0ee70cbab18fe2731ca0841a7de21ec03dd0f24e6316f07",
+        "975f6dfd455769cfba2621b0645c8ba660460ceba9162c2cb5a14bf0d25e3e02"
+      ]
+    },
+    {
+      "name": "lora_lookahead_0",
+      "hash_algo": "sha256",
+      "block_size": 5,
+      "num_lookahead_tokens": 0,
+      "lora": {
+        "name": "adapter-a",
+        "path": "/models/adapter-a"
+      },
+      "cache_salt": null,
+      "prompt_token_ids": [0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 7, 8],
+      "block_hashes": [
+        "354f46da231d7fbbe8387f35c2d3b314d827e063c365ebdcd5366d621de464e1",
+        "c07d4a59996cc567b05a1cfd3738d7c1f7de49f293ee0dc69179a38d0e3dcf34",
+        "9449b4123df032bd6b8a709b706f7b9f06c784bb0520d791a746f1136a03aca9",
+        "162f7a3f81986d782a629a8a4ecf718bd28828023744d12457f464888ff4509b",
+        "2dd4479ba655a901a1f4e9e719c1d41ae7f29bc0d45ede4265a3d1bf8195939a",
+        "82aa728271e8aaf4228f737222281aefc7e8bb96dfe2507f760295cc646b4a92"
+      ]
+    },
+    {
+      "name": "salt_lookahead_0",
+      "hash_algo": "sha256",
+      "block_size": 5,
+      "num_lookahead_tokens": 0,
+      "lora": null,
+      "cache_salt": "tenant-1",
+      "prompt_token_ids": [0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 7, 8],
+      "block_hashes": [
+        "06241d30c0c0700c890353e7487fbad2355e3a3f4fc34c8bb54d75e8e3005a7f",
+        "b2b9513d3e227a005a6bcacecf127ca8aded0ac9848aebe59b8fe6b23addae57",
+        "e30ad49081e476993cd27585cab926c337f6148b271eeba88266660ffcf9ed0e",
+        "ed73f2afc31b85ea1b745bcd37622a6e898905da9a21ce2d3c22471724ce311e",
+        "f227e68c3740320b632eab7694f1adf994f84cc3175ee7b312112c390ff85419",
+        "f46b061b047585f338cc8ea37c19f3fe634ea9b23d35985d759ba143ff19bc91"
+      ]
+    },
+    {
+      "name": "lora_salt_lookahead_0",
+      "hash_algo": "sha256",
+      "block_size": 5,
+      "num_lookahead_tokens": 0,
+      "lora": {
+        "name": "adapter-é",
+        "path": "/m/é"
+      },
+      "cache_salt": "ssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssss",
+      "prompt_token_ids": [0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 7, 8],
+      "block_hashes": [
+        "a4c783bc6d449adcdafdbfd6d59b15f9c23b92eaac1efad5abe94e6f48f97dd5",
+        "cc6648abcd50d135d1d1736a531654bfdd25143bd2035ab372afca88b299afee",
+        "b308d734cceccd07a939943ee90da391ef53665415da9bfec1201e528bce05f7",
+        "83ec7658100ca0bcf28d71c4666f562f62424c5465d2f59ce0ebef549067fc65",
+        "1b22570a84168a6c977cdcb9300e939b1a58b157492570313f99cc049d97b2dd",
+        "d94930748344c58edfd4a63ebff6c2515a93dd9a952f6ed8d0f0e210e8652093"
+      ]
+    },
+    {
+      "name": "plain_lookahead_1",
+      "hash_algo": "sha256",
+      "block_size": 5,
+      "num_lookahead_tokens": 1,
+      "lora": null,
+      "cache_salt": null,
+      "prompt_token_ids": [0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 7, 8],
+      "block_hashes": [
+        "3c7c5bb68aa46add1384d9ea79ddb81fc222eada29c779ff6a3a50485fc44e28",
+        "c7faac878edccc0bf592933d5c6e0397feb8e38365f2fb5a096f2b1a34d9efe1",
+        "0982d2579c5b61796e6011d4fa3378e8d42a4a7e685b964589e53872a48ee4d2",
+        "09f4c42834d5d79e6cb6cdbbee35595091a99922a94bb63f7798e9a2875306f7",
+        "2ae57b2fcab6e3b75e95c747f62b146c5705b2eeeaddc5baf1d9c0652858e423",
+        "158d6ea13b7352c92eb0cd1fc2ce1683b2b3ab280351d2901a4f500ee4c032e8"
+      ]
+    },
+    {
+      "name": "lora_lookahead_1",
+      "hash_algo": "sha256",
+      "block_size": 5,
+      "num_lookahead_tokens": 1,
+      "lora": {
+        "name": "adapter-a",
+        "path": "/models/adapter-a"
+      },
+      "cache_salt": null,
+      "prompt_token_ids": [0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 7, 8],
+      "block_hashes": [
+        "462d89fbe1de8e64755b339cd9b1f28d6d8fe37fd46f2c7cb500e55b967439f8",
+        "78cfc368b6cf8de5d4be1af33f8d810fd70abb727c4656ee76678c5572f7daef",
+        "1d40e74c91bc2a282bf662c32febee626ad2055d06147ca09c4c58a7753e806f",
+        "5899a58c6775cad2c5254dc4b3da428c90236f17670947efbad316df503e4ad5",
+        "0406c605fc0736ba6b3d5f7461bf721949b91a43ee7f16431203639da1dd1039",
+        "1809ea41c760099ad91a5b028e2fd78bce9207f882df3a8e52616ab247db4baf"
+      ]
+    },
+    {
+      "name": "salt_lookahead_1",
+      "hash_algo": "sha256",
+      "block_size": 5,
+      "num_lookahead_tokens": 1,
+      "lora": null,
+      "cache_salt": "tenant-1",
+      "prompt_token_ids": [0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 7, 8],
+      "block_hashes": [
+        "b75ac67378c80b6c465aa2a1a8819d594035819511971dedf5f3edfb5347d4e9",
+        "35a17d302ccfb3065bd3c29f816e8e477c8c22486f34e7c21806db04c14319c9",
+        "9ac6646ea2f23ea2df92f58d2804ef0399d7bd4c8c6e6bfa7019b26846736ab3",
+        "8ac8982774427c5f8bde9b7b4b8508a3b8e17d77ecd6de0d6752c201ff8afab0",
+        "a04eacb3f55cfb17e044f47c3bbb9225f30ece6e7c490663939d3f90c7df3b19",
+        "e4fa8639fc4ed687dbca1b7b398621403d4945efb58802f1bcb49825d339cd53"
+      ]
+    },
+    {
+      "name": "lora_salt_lookahead_1",
+      "hash_algo": "sha256",
+      "block_size": 5,
+      "num_lookahead_tokens": 1,
+      "lora": {
+        "name": "adapter-é",
+        "path": "/m/é"
+      },
+      "cache_salt": "ssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssss",
+      "prompt_token_ids": [0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 7, 8],
+      "block_hashes": [
+        "281a977700036131e70708dc04067eb00031476da6fc270eeb91ef3bec4af9fd",
+        "89a854fb90b44c0945c71d7257699cdbfd12921d7757f2a819503fe985dd5676",
+        "fb9e93e12a8af10382409506b90c16beefaba61c4d168a581ceeb202138e689b",
+        "61dbb5cd7782148f4e0bed242ad56ff82ea2283279ce7dfa1fef87b991b7133e",
+        "b93c4c36d6b6edd65db43841b0ca7fc1f7ee2aed3c1ce0d8afe0081475fa3bcb",
+        "3109559a165cba4f5adda0d5d9f54c4472ca23df2a354d19fe687df14d6c9be7"
+      ]
+    },
+    {
+      "name": "plain_lookahead_3",
+      "hash_algo": "sha256",
+      "block_size": 5,
+      "num_lookahead_tokens": 3,
+      "lora": null,
+      "cache_salt": null,
+      "prompt_token_ids": [0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 7, 8],
+      "block_hashes": [
+        "befa92d3904be5c3845960b68edee89434b3f3bd4098c3ec2a40f320203e94cc",
+        "ee46c224e6ae5cdbb0595306876642b8c322d00f5fcfdfc5087790b790810c2d",
+        "6909f6c3c43b89f16f3562d8716bd3689309fe7f06843a874e7ed0119f56a6d0",
+        "7fc9cb8fa98603defe464f51934f6d77bb66e8be9d6f3d9902fb75d4de9c3be6",
+        "158e0f4eb727ddc8eded92a45b77e328220b892f296fe8c6d304b3bd3cadb67b"
+      ]
+    },
+    {
+      "name": "lora_lookahead_3",
+      "hash_algo": "sha256",
+      "block_size": 5,
+      "num_lookahead_tokens": 3,
+      "lora": {
+        "name": "adapter-a",
+        "path": "/models/adapter-a"
+      },
+      "cache_salt": null,
+      "prompt_token_ids": [0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 7, 8],
+      "block_hashes": [
+        "51725f059c14065470452d9f3027f992f601bbf622fa83864c78a0efbd063ba2",
+        "5980cd32fcbf79d2ee8d5b78ccd7bba91e9713e91dc0a173cfb9244f3b531753",
+        "3236351da2cacef60579aed7d3864517af8c087f30a1f5b3e978f46448be91ba",
+        "a1ae73f215cf4870326e6b6d56e8ed3ee8fe36623b2a231c8c2d2536df2a2fc3",
+        "e644ef52f092cc1f4b27d340e1a114df1d9d65ebd54be016303b61f6d6c65df4"
+      ]
+    },
+    {
+      "name": "salt_lookahead_3",
+      "hash_algo": "sha256",
+      "block_size": 5,
+      "num_lookahead_tokens": 3,
+      "lora": null,
+      "cache_salt": "tenant-1",
+      "prompt_token_ids": [0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 7, 8],
+      "block_hashes": [
+        "070f8fe284fa36a9dbe70492aa9ee5f8bbfcf185de439f3826f0f59a06d221fc",
+        "f7f9a663ce12d9bb02822c275c0c06c5969383acfbaf58db6a2ae8645cc08f97",
+        "d223c14405bf7eede019a0eeed5a656bff4bfd2a03bc37e265e0362c644a783b",
+        "948bb51f39172864dcc38c94e4b62a4580ac190bf00aa9f3cd173efc90617db3",
+        "08d9d549413cdfe0a3f552adddb1ca025928b2ec4eca1d76365d872370506ea0"
+      ]
+    },
+    {
+      "name": "lora_salt_lookahead_3",
+      "hash_algo": "sha256",
+      "block_size": 5,
+      "num_lookahead_tokens": 3,
+      "lora": {
+        "name": "adapter-é",
+        "path": "/m/é"
+      },
+      "cache_salt": "ssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssss",
+      "prompt_token_ids": [0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 7, 8],
+      "block_hashes": [
+        "d8ee74ffd008d0c7065a877ad2455f910511261a79be770ca9bb444c856126c0",
+        "de21a65cb5fc93aec3e2032f5eaffeb2857a706b2a1d71662df5c6665a16e730",
+        "7d944508506bd46ef2e2418352849886e2396d3b458ea57f7c7580a5d54a4885",
+        "bc54bb8fd64a0b2cb7162b92b9eda3cdb72eadfc675ff66f5f462520cd2495e2",
+        "db3c933bf37e992760e0ce1083e9944ea25e218491bf06a75c2dd284e541fe3c"
+      ]
+    },
+    {
+      "name": "plain_block_size_16",
+      "hash_algo": "sha256_cbor",
+      "block_size": 16,
+      "num_lookahead_tokens": 0,
+      "lora": null,
+      "cache_salt": null,
+      "prompt_token_ids": [0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 7, 8],
+      "block_hashes": [
+        "b4e3d7f45cf989fa34b8f87e41e8da20296ed03ac04537e363d47ede8f6910aa",
+        "ab9339ba27613a3dbdd8ad445cd43a3f8175baca60c547911cc500c3ec5b409d"
+      ]
+    },
+    {
+      "name": "plain_block_size_1",
+      "hash_algo": "sha256_cbor",
+      "block_size": 1,
+      "num_lookahead_tokens": 0,
+      "lora": null,
+      "cache_salt": null,
+      "prompt_token_ids": [0, 300, 70000],
+      "block_hashes": [
+        "493f37dcda994a276ed9f9f27341b870a0da995cb22a36dc0e43fbf55d0cd662",
+        "6edeb0b0002590ff50d03ea30340c737c56ec7670d25b43e727800e1b2b37514",
+        "5b6d403bb7875016220be9de56c3271affb85e0a77f6e4b6ac442523277bad78"
+      ]
+    },
+    {
+      "name": "plain_block_size_3",
+      "hash_algo": "sha256_cbor",
+      "block_size": 3,
+      "num_lookahead_tokens": 0,
+      "lora": null,
+      "cache_salt": null,
+      "prompt_token_ids": [0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647],
+      "block_hashes": [
+        "a8419a131ade760fe7db6f193dd8ea7b0d898acd29aa35ab89176221482e24ac",
+        "081c7b465759791d54990230b455781b5cc69bc5f2454b055ac2f49c4e7baf87",
+        "ddc37216453bcfcd78ceb2903a6a2fe1e9694b11892ecff48cdad33f9a3aeccd"
+      ]
+    },
+    {
+      "name": "plain_block_size_24",
+      "hash_algo": "sha256_cbor",
+      "block_size": 24,
+      "num_lookahead_tokens": 0,
+      "lora": null,
+      "cache_salt": null,
+      "prompt_token_ids": [100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111, 112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 127, 128, 129, 130, 131, 132, 133, 134, 135, 136, 137, 138, 139, 140, 141, 142, 143, 144, 145, 146, 147, 148, 149],
+      "block_hashes": [
+        "4100327c81de908bc85378ee3ed53eb3a25961ff3cad70b6f45a8a2973426f63",
+        "4962a3803c42d379a6976da0ccf67f1f0e17e9cc72d57c5080804866e5e6c2e1"
+      ]
+    },
+    {
+      "name": "plain_block_size_300",
+      "hash_algo": "sha256_cbor",
+      "block_size": 300,
+      "num_lookahead_tokens": 0,
+      "lora": null,
+      "cache_salt": null,
+      "prompt_token_ids": [0, 7919, 15838, 23757, 31676, 39595, 47514, 55433, 63352, 71271, 79190, 87109, 95028, 102947, 110866, 118785, 126704, 134623, 142542, 150461, 158380, 166299, 174218, 182137, 190056, 197975, 5891, 13810, 21729, 29648, 37567, 45486, 53405, 61324, 69243, 77162, 85081, 93000, 100919, 108838, 116757, 124676, 132595, 140514, 148433, 156352, 164271, 172190, 180109, 188028, 195947, 3863, 11782, 19701, 27620, 35539, 43458, 51377, 59296, 67215, 75134, 83053, 90972, 98891, 106810, 114729, 122648, 130567, 138486, 146405, 154324, 162243, 170162, 178081, 186000, 193919, 1835, 9754, 17673, 25592, 33511, 41430, 49349, 57268, 65187, 73106, 81025, 88944, 96863, 104782, 112701, 120620, 128539, 136458, 144377, 152296, 160215, 168134, 176053, 183972, 191891, 199810, 7726, 15645, 23564, 31483, 39402, 47321, 55240, 63159, 71078, 78997, 86916, 94835, 102754, 110673, 118592, 126511, 134430, 142349, 150268, 158187, 166106, 174025, 181944, 189863, 197782, 5698, 13617, 21536, 29455, 37374, 45293, 53212, 61131, 69050, 76969, 84888, 92807, 100726, 108645, 116564, 124483, 132402, 140321, 148240, 156159, 164078, 171997, 179916, 187835, 195754, 3670, 11589, 19508, 27427, 35346, 43265, 51184, 59103, 67022, 74941, 82860, 90779, 98698, 106617, 114536, 122455, 130374, 138293, 146212, 154131, 162050, 169969, 177888, 185807, 193726, 1642, 9561, 17480, 25399, 33318, 41237, 49156, 57075, 64994, 72913, 80832, 88751, 96670, 104589, 112508, 120427, 128346, 136265, 144184, 152103, 160022, 167941, 175860, 183779, 191698, 199617, 7533, 15452, 23371, 31290, 39209, 47128, 55047, 62966, 70885, 78804, 86723, 94642, 102561, 110480, 118399, 126318, 134237, 142156, 150075, 157994, 165913, 173832, 181751, 189670, 197589, 5505, 13424, 21343, 29262, 37181, 45100, 53019, 60938, 68857, 76776, 84695, 92614, 100533, 108452, 116371, 124290, 132209, 140128, 148047, 155966, 163885, 171804, 179723, 187642, 195561, 3477, 11396, 19315, 27234, 35153, 43072, 50991, 58910, 66829, 74748, 82667, 90586, 98505, 106424, 114343, 122262, 130181, 138100, 146019, 153938, 161857, 169776, 177695, 185614, 193533, 1449, 9368, 17287, 25206, 33125, 41044, 48963, 56882, 64801, 72720, 80639, 88558, 96477, 104396, 112315, 120234, 128153, 136072, 143991, 151910, 159829, 167748, 175667, 183586, 191505, 199424, 7340, 15259, 23178, 31097, 39016, 46935, 54854, 62773, 70692, 78611, 86530, 94449, 102368, 110287, 118206, 126125, 134044, 141963, 149882, 157801, 165720, 173639, 181558, 189477, 197396, 5312, 13231, 21150, 29069, 36988, 44907, 52826, 60745, 68664, 76583, 84502, 92421, 100340, 108259, 116178, 124097, 132016, 139935, 147854, 155773, 163692, 171611, 179530, 187449, 195368, 3284, 11203, 19122, 27041, 34960, 42879, 50798, 58717, 66636, 74555, 82474, 90393, 98312, 106231, 114150, 122069, 129988, 137907, 145826, 153745, 161664, 169583, 177502, 185421, 193340, 1256, 9175, 17094, 25013, 32932, 40851, 48770, 56689, 64608, 72527, 80446, 88365, 96284, 104203, 112122, 120041, 127960, 135879, 143798, 151717, 159636, 167555, 175474, 183393, 191312, 199231, 7147, 15066, 22985, 30904, 38823, 46742, 54661, 62580, 70499, 78418, 86337, 94256, 102175, 110094, 118013, 125932, 133851, 141770, 149689, 157608, 165527, 173446, 181365, 189284, 197203, 5119, 13038, 20957, 28876, 36795, 44714, 52633, 60552, 68471, 76390, 84309, 92228, 100147, 108066, 115985, 123904, 131823, 139742, 147661, 155580, 163499, 171418, 179337, 187256, 195175, 3091, 11010, 18929, 26848, 34767, 42686, 50605, 58524, 66443, 74362, 82281, 90200, 98119, 106038, 113957, 121876, 129795, 137714, 145633, 153552, 161471, 169390, 177309, 185228, 193147, 1063, 8982, 16901, 24820, 32739, 40658, 48577, 56496, 64415, 72334, 80253, 88172, 96091, 104010, 111929, 119848, 127767, 135686, 143605, 151524, 159443, 167362, 175281, 183200, 191119, 199038, 6954, 14873, 22792, 30711, 38630, 46549, 54468, 62387, 70306, 78225, 86144, 94063, 101982, 109901, 117820, 125739, 133658, 141577, 149496, 157415, 165334, 173253, 181172, 189091, 197010, 4926, 12845, 20764, 28683, 36602, 44521, 52440, 60359, 68278, 76197, 84116, 92035, 99954, 107873, 115792, 123711, 131630, 139549, 147468, 155387, 163306, 171225, 179144, 187063, 194982, 2898, 10817, 18736, 26655, 34574, 42493, 50412, 58331, 66250, 74169, 82088, 90007, 97926, 105845, 113764, 121683, 129602, 137521, 145440, 153359, 161278, 169197, 177116, 185035, 192954, 870, 8789, 16708, 24627, 32546, 40465, 48384, 56303, 64222, 72141, 80060, 87979, 95898, 103817, 111736, 119655, 127574, 135493, 143412, 151331, 159250, 167169, 175088, 183007, 190926, 198845, 6761, 14680, 22599, 30518, 38437, 46356, 54275, 62194, 70113, 78032, 85951, 93870, 101789, 109708, 117627, 125546, 133465, 141384, 149303, 157222, 165141, 173060, 180979, 188898, 196817, 4733, 12652, 20571, 28490, 36409, 44328, 52247, 60166, 68085, 76004, 83923, 91842, 99761, 107680, 115599, 123518, 131437, 139356],
+      "block_hashes": [
+        "f8d904340e7699f0e56074debe137e5c25c91fae8e2c20b77d0a34f9dd0aee23",
+        "5c579410ecd94807db1dd87e335a1c9467b20bafbbc50f5d981e8abd9bfc1eaf"
+      ]
+    },
+    {
+      "name": "plain_lookahead_0",
+      "hash_algo": "sha256_cbor",
+      "block_size": 5,
+      "num_lookahead_tokens": 0,
+      "lora": null,
+      "cache_salt": null,
+      "prompt_token_ids": [0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 7, 8],
+      "block_hashes": [
+        "339f6de80342500080b21b7e3c61c46e5efdecfe290287e0bd565a9d42fdbcbd",
+        "823175a0e66a40f5b0d974e2b9f9b3c2894e4555c53dbd6088602a23b41eb90e",
+        "ad8b056538560fecc39da6e6350a012c7cf0b8a6eb489a8079c16161fab14852",
+        "351f10c51c199624768e861dab09dde15546a40eeaf80add7275c0501c033c8c",
+        "8493ac9bc37e64cb8430441f79c882114a34c341d19a8c9656417102a3387bb5",
+        "f60a875b5ea76b5d48f036138ee14ecd08c723e23756794bd57295cb229fe5a8"
+      ]
+    },
+    {
+      "name": "lora_lookahead_0",
+      "hash_algo": "sha256_cbor",
+      "block_size": 5,
+      "num_lookahead_tokens": 0,
+      "lora": {
+        "name": "adapter-a",
+        "path": "/models/adapter-a"
+      },
+      "cache_salt": null,
+      "prompt_token_ids": [0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 7, 8],
+      "block_hashes": [
+        "613539868e6bca2d39892c001f93b6883795c683f171e7108e1044215318a23a",
+        "2a6b9a8ce9937b57db421677b7ed6fa110438660da9c06aceb166444e3cf8da0",
+        "caa09bb5d77923985e0e4305b23184cdc83586ee415043e2b7bf9f29c37231ee",
+        "0cd5ba89d6850a7b0efe0b6615fa5dcab7777dc69614290430c31ee708142e3c",
+        "90ec760108fb12b6c2be881ecae3a3a74ab4e8fb6d93e382f6aa1b839e402ca0",
+        "5663159167f01a066ce2482a2155247542adaebe7eab0ed01f29d71317176693"
+      ]
+    },
+    {
+      "name": "salt_lookahead_0",
+      "hash_algo": "sha256_cbor",
+      "block_size": 5,
+      "num_lookahead_tokens": 0,
+      "lora": null,
+      "cache_salt": "tenant-1",
+      "prompt_token_ids": [0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 7, 8],
+      "block_hashes": [
+        "f578cca646674a49a675fe07f79f5126c7cf9f103e9e2676b484e5ce8ea6e48f",
+        "c99bd81e931ab2ffd35e2a9990cc0aef3c237591cc90820441b89faf1755a918",
+        "fed8e890c0ebc0027881d50f0b751ba5b231f79ea73f48b43c06acc4acf37027",
+        "dd26988b757ea027a78a6ea74403c5926f126864301a30fdbe4c4d59c266ea7f",
+        "f21321f2b5d3b43cdc43b10c86d7cb05053e2041987489f3b2aee320462d4873",
+        "7acf3ac9e258df746a17015fc5b8aeac00511f599028d2b670fe20b2b92eaaa7"
+      ]
+    },
+    {
+      "name": "lora_salt_lookahead_0",
+      "hash_algo": "sha256_cbor",
+      "block_size": 5,
+      "num_lookahead_tokens": 0,
+      "lora": {
+        "name": "adapter-é",
+        "path": "/m/é"
+      },
+      "cache_salt": "ssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssss",
+      "prompt_token_ids": [0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 7, 8],
+      "block_hashes": [
+        "5d8fc62a8373282ad491aa6c56943ed885da3f29e5fff16640f797f013ad0746",
+        "c5e2dedb410e2a7661063e9f5e3027734dca589e7f27fd3bfbce7616ac56d04c",
+        "6dbeb08efe56b2aa5b62dc86423db72a2101b4132bed6ce39ed4576e2bee763e",
+        "5473c762b8ea057df489bd23f3f8678e13799566ac80985e7132a888c73d63bd",
+        "49e9bd622f10fd84bf318d1f21fd2e952503c50fb5be1655cc9a93c394a90482",
+        "a01f56a47079cb79fa8d8edf9acf11aafbea3f417baf708b0fd06d5c64b75da9"
+      ]
+    },
+    {
+      "name": "plain_lookahead_1",
+      "hash_algo": "sha256_cbor",
+      "block_size": 5,
+      "num_lookahead_tokens": 1,
+      "lora": null,
+      "cache_salt": null,
+      "prompt_token_ids": [0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 7, 8],
+      "block_hashes": [
+        "2cd1c535592d7e8f81d1153412b0c7fb319d247d1c27098daee7261bd53392c1",
+        "773c3ba2cfb414d212b313cee38840e6cdb0e8f4d28aed50194dfb99e5b7b573",
+        "12739b963b318d97e3736473f4a1a24a9dce2e47900849ab3cb3ceb7df580e5d",
+        "b4139498ecdcb5acd651e0fda8abd3cd3815379c8f77263f1f755d7a5065c9e8",
+        "5c54ad3e1bcdb973d01ff5ca4a7a8a55d07e83cc0baa3011150af49bd8be3767",
+        "a7f254fe782288d4a742c4fb9057e94edb970e8584c3e3b958baae0b259614ad"
+      ]
+    },
+    {
+      "name": "lora_lookahead_1",
+      "hash_algo": "sha256_cbor",
+      "block_size": 5,
+      "num_lookahead_tokens": 1,
+      "lora": {
+        "name": "adapter-a",
+        "path": "/models/adapter-a"
+      },
+      "cache_salt": null,
+      "prompt_token_ids": [0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 7, 8],
+      "block_hashes": [
+        "05788a34017c1104e7f55f0ea41670eb8011d0aa1e5884f86462382adc7c3c1f",
+        "c2f8cab4b889d484aa6a4c5c5fbd235222b0dd1054e13f9fe7a65510df5910f2",
+        "76b651b5cd1332cdef01605dbf111784890de9d59a4cd9b9a01ce7415473b301",
+        "3d1bc40af524f03d0d597715cc3fe9822b720a969837dd7e90d401837a59dc1c",
+        "5f7da2ce93b3132b29675cc045d322435a5f0f78631836a8f98dff930c68de1e",
+        "1f4d67a4d69c9d6b9ba2f374c6a84db0a81893195fdf527d9e7715c4b073aebf"
+      ]
+    },
+    {
+      "name": "salt_lookahead_1",
+      "hash_algo": "sha256_cbor",
+      "block_size": 5,
+      "num_lookahead_tokens": 1,
+      "lora": null,
+      "cache_salt": "tenant-1",
+      "prompt_token_ids": [0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 7, 8],
+      "block_hashes": [
+        "d0932ff71fb9dd44ee41671687167b8c0fdc842cc9c32805fe30fdc2423c7418",
+        "a78e0618d987f8a9d0e28027611a726890df0d759c626d0ae2a1dd4244feb83c",
+        "462faa38d09910f8f228f4c4d6683c59e6d8fcbc156daf5745a32fbaf3e659c8",
+        "c4ae8b9bc29e6d4a23dbe7ded0b3bdc5291193e56020c43803ccdf075fe3d04c",
+        "2d551d04128267ea64a02e1d8d77d5df7e0b6a681c0d7b475b835aed5a06512e",
+        "d172f715798145183e406ab09bb7aaaaa484aec2099653e192cc6ed94d14dbeb"
+      ]
+    },
+    {
+      "name": "lora_salt_lookahead_1",
+      "hash_algo": "sha256_cbor",
+      "block_size": 5,
+      "num_lookahead_tokens": 1,
+      "lora": {
+        "name": "adapter-é",
+        "path": "/m/é"
+      },
+      "cache_salt": "ssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssss",
+      "prompt_token_ids": [0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 7, 8],
+      "block_hashes": [
+        "5322c6f9418a576fbb8888cc4b19bf0c1364e70f819e94b9e68d979f3c837568",
+        "3ca41ac5336e878f02eeaf2d470d72e6b5291e03a1c8563e6e79d32880ab7583",
+        "8f0529119de1c34e4f20ce39d3bdb6beddb387a4f6f8a1bf2815b2becdebea43",
+        "550b2d3f46e9b89e01c8f3eaf3f9637fb86bf2ed9641d24cd35937934b32c0fc",
+        "bd8707578b7be2ee143be3e2165eafb6cc6e092577ef2ab5216b7b45b1155007",
+        "43ec7a88990016de65d464768e3265259c41c550b5af2a56360413db52e1894b"
+      ]
+    },
+    {
+      "name": "plain_lookahead_3",
+      "hash_algo": "sha256_cbor",
+      "block_size": 5,
+      "num_lookahead_tokens": 3,
+      "lora": null,
+      "cache_salt": null,
+      "prompt_token_ids": [0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 7, 8],
+      "block_hashes": [
+        "a108c7ac9e75b566b45c0fbc271641e9796a1c4d519594e4a7eb6956903ffbba",
+        "f1d81991a4dbba087a7f051e64dfe029aac3f8b39f0d42cc45dc39c8e2a91619",
+        "70362d9102cb050b4accf4f0558d076d60eebb2d1367a1c32ade73598f629578",
+        "28290effb9f23c528780b3430bf3f68da8c517e7086364774d3d8ab0b51dbb50",
+        "d5d56edf84af5784f508a9157f1979425c3e491f1c1824abf0a9986c4307041e"
+      ]
+    },
+    {
+      "name": "lora_lookahead_3",
+      "hash_algo": "sha256_cbor",
+      "block_size": 5,
+      "num_lookahead_tokens": 3,
+      "lora": {
+        "name": "adapter-a",
+        "path": "/models/adapter-a"
+      },
+      "cache_salt": null,
+      "prompt_token_ids": [0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 7, 8],
+      "block_hashes": [
+        "41ad34e840c0acd2077d256363a6bc6be245458955ab89587911ced23931cc2b",
+        "5e7fe7c71803c5d4c5c9af4cb45dd40b39f84f11f12030c271b6e7dbe5c3b7bc",
+        "957dc13d51b2489fc56ce3744e5ba45bfb24eab4ba02f8665cc3841c184311b1",
+        "e4b5d8d814a6146da674e40579dc2aa60329d540b47d5435dfc108094bf89944",
+        "9d5d87932ca2c41207ecadc63d39c9335d48b84e6be33a64aa02453768624cb5"
+      ]
+    },
+    {
+      "name": "salt_lookahead_3",
+      "hash_algo": "sha256_cbor",
+      "block_size": 5,
+      "num_lookahead_tokens": 3,
+      "lora": null,
+      "cache_salt": "tenant-1",
+      "prompt_token_ids": [0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 7, 8],
+      "block_hashes": [
+        "ec7caf3a4a17c4680fa13870361419627219bf4b90f6c678370584de10016ec8",
+        "65f8d7ae0df13ce22ffe42343e0376664363cb41382f553263c79d3ab3aff18a",
+        "f2a9ffd77b13a777998504e2ebcda27a964ee857964bfbce9f10aa3abdb28d25",
+        "65541a79ee789b6f09e34a1619430402e6884bf2b1cfc04acb0e346baafb7e26",
+        "798b35d9c87d9e14182d410b87757a5c90fd3b3054204af3b0a0e014646deb01"
+      ]
+    },
+    {
+      "name": "lora_salt_lookahead_3",
+      "hash_algo": "sha256_cbor",
+      "block_size": 5,
+      "num_lookahead_tokens": 3,
+      "lora": {
+        "name": "adapter-é",
+        "path": "/m/é"
+      },
+      "cache_salt": "ssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssss",
+      "prompt_token_ids": [0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 0, 1, 23, 24, 255, 256, 65535, 65536, 151645, 2147483647, 7, 8],
+      "block_hashes": [
+        "8e7c16de0fa89e4d029c9cf08f9d1924298da3e483afc2d7e1ef584b707db4ab",
+        "e0f7590a34be191bbfd8f1ab89150a19508053dff39ee1d9544e5435604c1de3",
+        "64c300bd2dd2c0eb8739dbc15aee77888bf2791fdd717b13592ff0f6dc762e92",
+        "4a43f99d50833c3d1beacd5f5fd0e2de0c246f8ec4f6d5cbd86b35e306f3bbbf",
+        "04ef0e5cc97bb4382dba452903f6af733972ff16582a1e6ca446269b52122d14"
+      ]
+    }
+  ]
+}
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 e2c12da61f78..c06716cd6301 100755
--- a/rust/src/engine-core-client/src/tests/python_compat.py
+++ b/rust/src/engine-core-client/src/tests/python_compat.py
@@ -105,6 +105,7 @@ class EngineCoreRequest(
     abort_immediately: bool = False
     session_id: str | None = None
     kv_hints: KvHintsEnvelope | None = None
+    block_hashes: list[bytes] | None = None
 
 
 @dataclass(frozen=True)
@@ -202,6 +203,7 @@ class EngineCoreOutputs(
             )
         ],
     ),
+    block_hashes=[b"\x01" * 32, b"\x02" * 32],
 )
 
 # All defaults -> empty map. Regression guard for the sparse-map decode.
@@ -498,6 +500,14 @@ class KVEventsConfig:
     topic: str
 
 
+@dataclass
+class PrefixHashConfig:
+    hash_algo: str
+    hash_block_size: int
+    none_hash: bytes
+    num_lookahead_tokens: int = 0
+
+
 @dataclass
 class EngineCoreReadyResponse:
     max_model_len: int
@@ -525,6 +535,7 @@ class EngineCoreReadyResponse:
     supports_draft_weight_updates: bool = False
     effective_attention_block_size: int | None = None
     vllm_config: dict | None = None
+    prefix_hash_config: PrefixHashConfig | None = None
 
 
 ready_response = EngineCoreReadyResponse(
@@ -549,6 +560,12 @@ class EngineCoreReadyResponse:
     enable_sleep_mode=True,
     supports_draft_weight_updates=True,
     effective_attention_block_size=64,
+    prefix_hash_config=PrefixHashConfig(
+        hash_algo="sha256",
+        hash_block_size=16,
+        none_hash=bytes(range(32)),
+        num_lookahead_tokens=1,
+    ),
     kv_events_config=KVEventsConfig(
         enable_kv_cache_events=True,
         publisher="zmq",
@@ -597,10 +614,16 @@ class EngineCoreReadyResponse:
     "num_failed_recvs": [],
     "num_kv_expired_reqs": [1],
 }
+# `UMBPStoreConnectorStats.data`: interval counters keyed by operation.
+umbp_stats = {
+    "load": {"submitted": 3, "completed": 2, "failed": 1, "num_bytes": 8192},
+    "store": {"submitted": 1, "completed": 1, "failed": 0, "num_bytes": 4096},
+}
 multi_connector_stats = {
     "NixlConnector": nixl_stats,
     "MooncakeConnector": mooncake_transfer_stats,
     "MooncakeStoreConnector": mooncake_stats,
+    "UMBPStoreConnector": umbp_stats,
     "UnsupportedConnector": {"sample_count": 1},
 }
 
@@ -626,6 +649,7 @@ class EngineCoreReadyResponse:
 print(msgspec.msgpack.encode(nixl_stats).hex())
 print(msgspec.msgpack.encode(mooncake_stats).hex())
 print(msgspec.msgpack.encode(mooncake_transfer_stats).hex())
+print(msgspec.msgpack.encode(umbp_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/engine-core-client/src/transport.rs b/rust/src/engine-core-client/src/transport.rs
index dc1c743ede6c..2768256918fd 100644
--- a/rust/src/engine-core-client/src/transport.rs
+++ b/rust/src/engine-core-client/src/transport.rs
@@ -3,7 +3,7 @@
 
 use std::collections::{BTreeMap, BTreeSet};
 use std::fmt::Debug;
-use std::ops::Deref;
+use std::ops::{Deref, Range};
 use std::os::fd::{FromRawFd, OwnedFd, RawFd};
 use std::os::unix::net::UnixListener as StdUnixListener;
 use std::time::Duration;
@@ -21,11 +21,14 @@ use zeromq::{PullSocket, RouterSendHalf, RouterSocket, ZmqError, ZmqMessage};
 use crate::coordinator::CoordinatorBootstrap;
 use crate::error::{Error, Result, bail_unexpected_handshake_message};
 use crate::protocol::handshake::{
-    EngineCoreReadyResponse, HandshakeAddresses, HandshakeInitMessage, ReadyMessage,
+    DataParallelRendezvous, EngineCoreReadyResponse, HandshakeAddresses, HandshakeInitMessage,
+    HandshakeParallelConfig, ReadyMessage,
 };
 use crate::protocol::output::{EngineCoreOutputs, decode_engine_core_outputs};
 use crate::protocol::{decode_msgpack, encode_msgpack};
 
+mod dp_rendezvous;
+
 /// Dedicated single-frame sentinel emitted by Python `EngineCoreProc` when the
 /// engine dies.
 pub const ENGINE_CORE_DEAD_SENTINEL: &[u8] = b"ENGINE_CORE_DEAD";
@@ -141,9 +144,18 @@ enum EngineStartupState {
 
 /// Connect to one or more engines through the startup handshake protocol,
 /// returning the shared data-plane transport plus the registered engines.
+///
+/// The frontend connects to the engines of DP ranks `[engine_start_index,
+/// engine_start_index + engine_count)`. The frontend whose range starts at
+/// rank 0 handshakes with all `data_parallel_size` engines: in hybrid or external DP LB, engines of other
+/// ranks fetch the coordinator addresses and the DP rendezvous here and take
+/// their input/output addresses from their own frontend. Other frontends
+/// handshake only with their own engines.
 pub async fn connect_handshake(
     handshake_address: &str,
+    engine_start_index: u32,
     engine_count: usize,
+    data_parallel_size: usize,
     local_host: &str,
     local_input_address: Option<&str>,
     local_output_address: Option<&str>,
@@ -153,10 +165,15 @@ pub async fn connect_handshake(
     if engine_count == 0 {
         bail_unexpected_handshake_message!("expected engine_count >= 1");
     }
+    let handshake_count = if engine_start_index == 0 {
+        data_parallel_size
+    } else {
+        engine_count
+    };
 
     info!(
-        engine_count,
-        handshake_address, "waiting for engines to connect"
+        engine_start_index,
+        engine_count, handshake_count, handshake_address, "waiting for engines to connect"
     );
 
     // 1. Bind shared local input/output sockets first so every engine receives the same data-plane
@@ -176,6 +193,13 @@ pub async fn connect_handshake(
     } else {
         None
     };
+    // Coordinated (MoE) DP ranks form cross-rank process groups and need one
+    // rendezvous; the in-process coordinator runs exactly in that case.
+    let parallel_config = HandshakeParallelConfig {
+        data_parallel_rendezvous: enable_inproc_coordinator
+            .then(|| DataParallelRendezvous::allocate(local_host, data_parallel_size))
+            .transpose()?,
+    };
 
     // 2. Bind the shared handshake socket once. All engines connect to this socket with their own
     //    identities, and startup order does not matter.
@@ -183,14 +207,22 @@ pub async fn connect_handshake(
     handshake_socket.bind(handshake_address).await?;
 
     let mut engines = BTreeMap::new();
+    // A frontend serving a subset of the DP ranks knows which engines may join:
+    // every rank for rank 0, its own ranks otherwise.
+    let own_ranks = engine_start_index..engine_start_index + engine_count as u32;
+    let handshake_ranks = if engine_start_index == 0 {
+        0..data_parallel_size as u32
+    } else {
+        own_ranks.clone()
+    };
 
     // 3. Receive HELLO from every engine and send a matching INIT. When coordinator mode is
     //    enabled, the engines will not emit READY until the coordinator barrier below completes.
-    while engines.len() < engine_count {
+    while engines.len() < handshake_count {
         debug!(
             handshake_address,
             connected = engines.len(),
-            waiting_for = engine_count,
+            waiting_for = handshake_count,
             "waiting for engine HELLO"
         );
         let message = timeout(ready_timeout, handshake_socket.recv()).await.map_err(|_| {
@@ -208,6 +240,14 @@ pub async fn connect_handshake(
                     );
                 }
                 debug!(handshake_address, ?engine_id, "received HELLO from engine");
+                if engine_count < data_parallel_size {
+                    validate_partitioned_engine(
+                        &engine_id,
+                        &handshake_message,
+                        &handshake_ranks,
+                        &own_ranks,
+                    )?;
+                }
 
                 send_init_message(
                     &mut handshake_socket,
@@ -215,6 +255,7 @@ pub async fn connect_handshake(
                     &input_address,
                     &output_address,
                     coordinator.as_ref(),
+                    &parallel_config,
                 )
                 .await?;
                 debug!(handshake_address, ?engine_id, "sent INIT to engine");
@@ -251,16 +292,18 @@ pub async fn connect_handshake(
 
     // 4. Optional coordinator startup gate. Without coordinator there is nothing to do.
     if let Some(coordinator) = coordinator.as_mut() {
-        coordinator.wait_for_startup_gate(engine_count, ready_timeout).await?;
+        coordinator.wait_for_startup_gate(handshake_count, ready_timeout).await?;
     }
 
-    // 5. After the optional gate has opened, every engine may now send READY.
+    // 5. After the optional gate has opened, every engine may now send READY. Only coordinated
+    //    (MoE) DP ranks share collectives, so only they must agree on the parallel config.
+    let mut parallel_config_hash = None;
     while engines.values().any(|state| !state.is_ready_received()) {
         debug!(
             handshake_address,
             connected = engines.len(),
             ready = engines.values().filter(|state| state.is_ready_received()).count(),
-            waiting_for = engine_count,
+            waiting_for = handshake_count,
             "waiting for engine READY"
         );
         let message = timeout(ready_timeout, handshake_socket.recv()).await.map_err(|_| {
@@ -286,6 +329,13 @@ pub async fn connect_handshake(
                     ?handshake_message,
                     "received READY from engine"
                 );
+                if enable_inproc_coordinator {
+                    check_parallel_config_hash(
+                        &mut parallel_config_hash,
+                        &engine_id,
+                        &handshake_message,
+                    )?;
+                }
                 *state = EngineStartupState::ReadyReceived;
             }
             Some("HELLO") => {
@@ -299,9 +349,26 @@ pub async fn connect_handshake(
         }
     }
 
-    // 6. Wait for every engine to connect to the shared input socket and register itself.
+    // 6. Wait for this frontend's engines to connect to the shared input socket and register
+    //    themselves. A frontend owning a subset of the DP ranks selects them by engine index.
+    let engine_ids: Vec<_> = if engine_count == data_parallel_size {
+        engines.into_keys().collect()
+    } else {
+        let engine_range = engine_start_index..engine_start_index + engine_count as u32;
+        let engine_ids: Vec<_> = engines
+            .into_keys()
+            .filter(|engine_id| engine_id.engine_index().is_some_and(|i| engine_range.contains(&i)))
+            .collect();
+        if engine_ids.len() != engine_count {
+            bail_unexpected_handshake_message!(
+                "expected {engine_count} engines of DP ranks {engine_range:?} to handshake, got {}",
+                engine_ids.len()
+            );
+        }
+        engine_ids
+    };
     let engines =
-        wait_for_input_registrations(&mut input_socket, engines.into_keys(), ready_timeout).await?;
+        wait_for_input_registrations(&mut input_socket, engine_ids, ready_timeout).await?;
     debug!(
         engine_count = engines.len(),
         "all engines registered on shared input socket"
@@ -468,6 +535,50 @@ fn decode_handshake_message(
     Ok((actual_id, handshake_message))
 }
 
+/// Check that an engine joining a frontend that serves a subset of the DP ranks
+/// is one of the expected ranks, and that an engine of another node's rank
+/// joins through its own frontend, not headless, as hybrid and external DP LB
+/// require.
+fn validate_partitioned_engine(
+    engine_id: &EngineId,
+    hello: &ReadyMessage,
+    handshake_ranks: &Range<u32>,
+    own_ranks: &Range<u32>,
+) -> Result<()> {
+    let Some(rank) = engine_id.engine_index().filter(|rank| handshake_ranks.contains(rank)) else {
+        bail_unexpected_handshake_message!(
+            "engine {engine_id:?} is not one of the expected DP ranks {handshake_ranks:?}"
+        );
+    };
+    if !own_ranks.contains(&rank) && hello.headless == Some(true) {
+        bail_unexpected_handshake_message!(
+            "the engine of DP rank {rank} joined headless; in hybrid or external DP LB every node runs its own frontend"
+        );
+    }
+    Ok(())
+}
+
+/// Check that every engine reports the same `parallel_config_hash`. Engines
+/// from launchers that disagree would otherwise pass startup and hang in a
+/// cross-rank collective.
+fn check_parallel_config_hash(
+    expected: &mut Option<String>,
+    engine_id: &EngineId,
+    ready: &ReadyMessage,
+) -> Result<()> {
+    let Some(hash) = &ready.parallel_config_hash else {
+        return Ok(());
+    };
+    match expected {
+        None => *expected = Some(hash.clone()),
+        Some(expected) if expected != hash => bail_unexpected_handshake_message!(
+            "engine {engine_id:?} reports parallel config hash {hash}, but other engines report {expected}; launch every DP rank with the same parallel config"
+        ),
+        Some(_) => {}
+    }
+    Ok(())
+}
+
 /// Send an INIT message to the engine with the local socket addresses for the
 /// engine to connect to, using the handshake socket.
 async fn send_init_message(
@@ -476,6 +587,7 @@ async fn send_init_message(
     input_address: &str,
     output_address: &str,
     coordinator: Option<&CoordinatorBootstrap>,
+    parallel_config: &HandshakeParallelConfig,
 ) -> Result<()> {
     let init_message = HandshakeInitMessage {
         addresses: HandshakeAddresses {
@@ -483,9 +595,9 @@ async fn send_init_message(
             outputs: vec![output_address.to_string()],
             coordinator_input: coordinator.map(|c| c.input_address.clone()),
             coordinator_output: coordinator.map(|c| c.output_address.clone()),
-            frontend_stats_publish_address: None,
+            frontend_stats_publish_address: coordinator.map(|c| c.frontend_publish_address.clone()),
         },
-        parallel_config: Default::default(),
+        parallel_config: parallel_config.clone(),
     };
     let payload = encode_msgpack(&init_message)?;
     let message = ZmqMessage::try_from(vec![engine_id.to_frame(), Bytes::from(payload)])
diff --git a/rust/src/engine-core-client/src/transport/dp_rendezvous.rs b/rust/src/engine-core-client/src/transport/dp_rendezvous.rs
new file mode 100644
index 000000000000..483d7ee5cc4f
--- /dev/null
+++ b/rust/src/engine-core-client/src/transport/dp_rendezvous.rs
@@ -0,0 +1,65 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+use std::net::TcpListener;
+
+use crate::error::Result;
+use crate::protocol::handshake::DataParallelRendezvous;
+
+/// Ports listed besides the master port, as many as Python's `ParallelConfig`
+/// draws.
+const LISTED_PORT_COUNT: usize = 4;
+
+impl DataParallelRendezvous {
+    /// Allocate the rendezvous of all `data_parallel_size` ranks on the
+    /// handshake host.
+    ///
+    /// Each launcher's `ParallelConfig` draws its own ports, so engines spawned
+    /// by different launchers (another node, or a second launcher on the same
+    /// node) only agree on them when the handshake owner sends one set to all
+    /// of them. The ports are bound and released here, on the host the engines
+    /// dial.
+    pub(super) fn allocate(host: &str, data_parallel_size: usize) -> Result<Self> {
+        // ZMQ endpoints spell IPv6 hosts in brackets; sockets and torch take
+        // the bare address.
+        let master_ip =
+            host.strip_prefix('[').and_then(|host| host.strip_suffix(']')).unwrap_or(host);
+
+        // Keep every listener bound until all ports are read so that none
+        // repeats.
+        let master = TcpListener::bind((master_ip, 0))?;
+        let listed = (0..LISTED_PORT_COUNT)
+            .map(|_| TcpListener::bind((master_ip, 0)))
+            .collect::<std::io::Result<Vec<_>>>()?;
+        let port = |listener: &TcpListener| listener.local_addr().map(|addr| addr.port());
+
+        Ok(Self {
+            data_parallel_master_ip: master_ip.to_string(),
+            data_parallel_master_port: port(&master)?,
+            data_parallel_master_port_list: listed
+                .iter()
+                .map(port)
+                .collect::<std::io::Result<_>>()?,
+            data_parallel_size,
+        })
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::net::TcpListener;
+
+    use crate::protocol::handshake::DataParallelRendezvous;
+
+    #[test]
+    fn allocate_sends_bare_ipv6_master_ip() {
+        if TcpListener::bind(("::1", 0)).is_err() {
+            eprintln!("skipping: IPv6 loopback is unavailable");
+            return;
+        }
+
+        let rendezvous = DataParallelRendezvous::allocate("[::1]", 2).unwrap();
+
+        assert_eq!(rendezvous.data_parallel_master_ip, "::1");
+    }
+}
diff --git a/rust/src/llm/examples/external_engine_smoke.rs b/rust/src/llm/examples/external_engine_smoke.rs
index c965ea303853..32d119a472c5 100644
--- a/rust/src/llm/examples/external_engine_smoke.rs
+++ b/rust/src/llm/examples/external_engine_smoke.rs
@@ -114,7 +114,9 @@ async fn main() -> Result<()> {
         transport_mode: TransportMode::HandshakeOwner {
             handshake_address: args.handshake_address.clone(),
             advertised_host: args.host.clone(),
+            engine_start_index: 0,
             engine_count: args.engine_count,
+            data_parallel_size: args.engine_count,
             ready_timeout,
             local_input_address: None,
             local_output_address: None,
diff --git a/rust/src/llm/src/log_stats.rs b/rust/src/llm/src/log_stats.rs
index d16c8bcbe3a3..70e575706479 100644
--- a/rust/src/llm/src/log_stats.rs
+++ b/rust/src/llm/src/log_stats.rs
@@ -6,9 +6,11 @@ use std::time::{Duration, Instant};
 
 use tokio_util::task::AbortOnDropHandle;
 use tracing::{debug, info};
+use vllm_engine_core_client::protocol::stats::KvFetchReqsByStage;
 use vllm_metrics::{
-    EngineLabels, F64Gauge, METRICS, PromptTokenSourceLabels, SchedulerLogStatsAccumulator,
-    SchedulerLogStatsInterval, U64Counter, U64Gauge, WaitingReasonLabels,
+    EngineLabels, F64Gauge, KvFetchStageLabels, METRICS, PromptTokenSourceLabels,
+    SchedulerLogStatsAccumulator, SchedulerLogStatsInterval, U64Counter, U64Gauge,
+    WaitingReasonLabels,
 };
 
 const WAITING_REASON_DEFERRED: &str = "deferred";
@@ -37,6 +39,8 @@ struct EngineMetrics {
     scheduler_running: U64Gauge,
     scheduler_waiting: U64Gauge,
     scheduler_deferred: U64Gauge,
+    /// One handle per `KvFetchReqsByStage::STAGES` entry.
+    scheduler_kv_fetch_by_stage: [U64Gauge; 3],
     kv_cache_usage: F64Gauge,
 }
 
@@ -162,6 +166,15 @@ fn resolve_engine_metrics(model_name: &str, engine_indices: &[u32]) -> Vec<Engin
                     .scheduler
                     .scheduler_waiting_by_reason
                     .get_or_create_owned(&deferred),
+                scheduler_kv_fetch_by_stage: KvFetchReqsByStage::STAGES.map(|stage| {
+                    m.scheduler.scheduler_kv_fetch_by_stage.get_or_create_owned(
+                        &KvFetchStageLabels {
+                            model_name: model_name.to_string(),
+                            engine,
+                            stage,
+                        },
+                    )
+                }),
                 kv_cache_usage: m.scheduler.kv_cache_usage.get_or_create_owned(&el),
             }
         })
@@ -221,6 +234,11 @@ async fn run_stats_logger(model_name: String, engine_indices: Vec<u32>, interval
         // Read scheduler gauges (aggregate across engines).
         let (num_running, num_waiting, kv_cache_usage) = read_scheduler_gauges(&engines);
         let num_deferred = read_deferred_waiting(&engines);
+        let [
+            kv_fetch_waiting_to_start,
+            kv_fetch_in_progress,
+            kv_fetch_completed_waiting,
+        ] = read_kv_fetch_by_stage(&engines);
         let delta_preemptions = curr.num_preemptions.wrapping_sub(prev.num_preemptions);
 
         // Compute prefix cache hit rate over this interval.
@@ -252,6 +270,15 @@ async fn run_stats_logger(model_name: String, engine_indices: Vec<u32>, interval
         if num_deferred > 0 {
             write!(msg, ", Deferred: {num_deferred} reqs").unwrap();
         }
+        if kv_fetch_waiting_to_start + kv_fetch_in_progress + kv_fetch_completed_waiting > 0 {
+            write!(
+                msg,
+                ", KV fetch: {kv_fetch_waiting_to_start} waiting to start, \
+                 {kv_fetch_in_progress} in progress, \
+                 {kv_fetch_completed_waiting} completed waiting"
+            )
+            .unwrap();
+        }
         if delta_preemptions > 0 {
             write!(msg, ", Preemptions: {delta_preemptions}").unwrap();
         }
@@ -366,6 +393,18 @@ fn read_deferred_waiting(engines: &[EngineMetrics]) -> u64 {
     engines.iter().map(|e| e.scheduler_deferred.get()).sum()
 }
 
+/// Read waiting requests per async KV load stage across all engines, in
+/// `KvFetchReqsByStage::STAGES` order.
+fn read_kv_fetch_by_stage(engines: &[EngineMetrics]) -> [u64; 3] {
+    let mut counts = [0u64; 3];
+    for e in engines {
+        for (count, gauge) in counts.iter_mut().zip(&e.scheduler_kv_fetch_by_stage) {
+            *count += gauge.get();
+        }
+    }
+    counts
+}
+
 /// Return the cache hit rate as a percentage for a counter delta.
 fn cache_hit_rate(hits: u64, queries: u64) -> f64 {
     if queries > 0 {
diff --git a/rust/src/llm/src/output.rs b/rust/src/llm/src/output.rs
index d74a2b4cc73e..78d7625014d3 100644
--- a/rust/src/llm/src/output.rs
+++ b/rust/src/llm/src/output.rs
@@ -208,7 +208,11 @@ pub struct GenerateOutput {
     pub sampling_mask: Option<SamplingMask>,
     /// Per-request speculative-decoding metrics, present on terminal outputs.
     pub spec_decode_metrics: Option<Box<RequestSpecDecodeMetrics>>,
-    /// Weight-version spans over all output tokens, set on terminal outputs only.
+    /// Weight-version spans over all output tokens so far; `None` if tokens
+    /// arrived but no label was reported. Set on every output, not only
+    /// terminal ones, because text decoding may finish the request locally
+    /// after matching a stop string and then cuts the spans at the truncated
+    /// token count.
     pub weight_versions: Option<Vec<WeightVersionSpan>>,
     /// Routed-expert rows for the whole request, set on terminal outputs only.
     pub routed_experts: Option<Box<WireNdArray>>,
@@ -444,11 +448,7 @@ impl GenerateOutputStream {
         {
             self.request_metrics.record_finished(received_at, finish_reason.clone());
         }
-        let weight_versions = if finish_reason.is_some() {
-            self.weight_version_spans()
-        } else {
-            None
-        };
+        let weight_versions = self.weight_version_spans();
         let routed_experts = if finish_reason.is_some() {
             self.take_routed_experts(&raw.request_id)?
         } else {
diff --git a/rust/src/llm/src/request.rs b/rust/src/llm/src/request.rs
index 0481851ea70f..f1eae64f75db 100644
--- a/rust/src/llm/src/request.rs
+++ b/rust/src/llm/src/request.rs
@@ -125,6 +125,9 @@ impl GenerateRequest {
                 reasoning_ended,
                 reasoning_parser_kwargs,
                 abort_immediately: false,
+                // Computed by the engine-core client, which knows the engine's
+                // prefix hash config.
+                block_hashes: None,
             },
         })
     }
diff --git a/rust/src/llm/tests/generate.rs b/rust/src/llm/tests/generate.rs
index 837c73c1151b..7efbbb8038e0 100644
--- a/rust/src/llm/tests/generate.rs
+++ b/rust/src/llm/tests/generate.rs
@@ -438,8 +438,8 @@ async fn generate_merges_outputs_by_stream_interval() {
     assert_eq!(
         weight_versions,
         [
-            None,
-            None,
+            Some(vec![span("v7", 0, 1)]),
+            Some(vec![span("v7", 0, 2), span("v8", 2, 5)]),
             Some(vec![span("v7", 0, 2), span("v8", 2, 5), span("v7", 5, 7)])
         ]
     );
diff --git a/rust/src/managed-engine/src/cli.rs b/rust/src/managed-engine/src/cli.rs
index dc6c55d40052..ef51a8a01d5c 100644
--- a/rust/src/managed-engine/src/cli.rs
+++ b/rust/src/managed-engine/src/cli.rs
@@ -7,6 +7,7 @@ use std::ffi::OsString;
 use clap::error::ErrorKind;
 use clap::{Args, CommandFactory};
 
+use crate::data_parallel::DataParallelTopology;
 use crate::{ManagedEngineConfig, allocate_handshake_port};
 
 /// Managed Python headless-engine CLI arguments.
@@ -94,6 +95,8 @@ impl ManagedEngineArgs {
         limit_mm_per_prompt: Option<String>,
         hf_overrides: Option<String>,
         structured_outputs_config: Option<String>,
+        topology: &DataParallelTopology,
+        client_handshake_address: Option<String>,
     ) -> ManagedEngineConfig {
         let mut python_args = self.python_args;
         // Manually forward some args to the Python engine.
@@ -125,9 +128,16 @@ impl ManagedEngineArgs {
             python_args.push("--shutdown-timeout".to_string());
             python_args.push(shutdown_timeout.to_string());
         }
-        if let Some(data_parallel_size_local) = self.data_parallel_size_local {
+        // This node's engines launch as a headless internal-LB node; hybrid
+        // and external LB differ only in which frontend they handshake with,
+        // which `client_handshake_address` carries.
+        if self.data_parallel_size_local.is_some() || topology.local_engines_only {
             python_args.push("--data-parallel-size-local".to_string());
-            python_args.push(data_parallel_size_local.to_string());
+            python_args.push(topology.local_size.to_string());
+        }
+        if topology.start_rank != 0 {
+            python_args.push("--data-parallel-start-rank".to_string());
+            python_args.push(topology.start_rank.to_string());
         }
         if let Some(limit_mm_per_prompt) = limit_mm_per_prompt {
             python_args.push("--limit-mm-per-prompt".to_string());
@@ -148,22 +158,10 @@ impl ManagedEngineArgs {
             handshake_host: self.handshake_host,
             handshake_port,
             data_parallel_size: self.data_parallel_size,
+            client_handshake_address,
             python_args,
         }
     }
-
-    /// Return the number of engines that the Rust frontend should expect to
-    /// coordinate with.
-    fn local_engine_count(&self) -> usize {
-        self.data_parallel_size_local.unwrap_or(self.data_parallel_size)
-    }
-
-    /// Return whether the managed Rust frontend only needs to communicate with
-    /// colocated engines.
-    pub fn frontend_local_only(&self) -> bool {
-        self.data_parallel_size_local != Some(0)
-            && self.local_engine_count() == self.data_parallel_size
-    }
 }
 
 /// Python `argparse` accepts these multi-character single-dash aliases, but
diff --git a/rust/src/managed-engine/src/data_parallel.rs b/rust/src/managed-engine/src/data_parallel.rs
new file mode 100644
index 000000000000..8e0fee0b61de
--- /dev/null
+++ b/rust/src/managed-engine/src/data_parallel.rs
@@ -0,0 +1,311 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+use anyhow::bail;
+use clap::Args;
+
+use crate::cli::ManagedEngineArgs;
+
+/// Data parallel ranks whose engines one `vllm-rs serve` process launches,
+/// and which of them its frontend connects to.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub struct DataParallelTopology {
+    /// First data parallel rank of this node.
+    pub start_rank: u32,
+    /// Number of data parallel ranks whose engines this process launches.
+    pub local_size: usize,
+    /// Number of data parallel ranks across the whole deployment.
+    pub data_parallel_size: usize,
+    /// Whether the frontend connects only to this node's ranks (hybrid or
+    /// external DP LB) while frontends on other nodes serve the rest. Otherwise
+    /// one frontend connects to every rank (internal DP LB).
+    pub local_engines_only: bool,
+}
+
+impl DataParallelTopology {
+    /// Number of engines the frontend connects to.
+    pub fn frontend_engine_count(&self) -> usize {
+        if self.local_engines_only {
+            self.local_size
+        } else {
+            self.data_parallel_size
+        }
+    }
+
+    /// Whether every engine the frontend connects to runs on this node.
+    pub fn frontend_local_only(&self) -> bool {
+        self.local_size != 0
+            && (self.local_engines_only || self.local_size == self.data_parallel_size)
+    }
+
+    /// First data parallel rank the frontend connects to. The frontend whose
+    /// range starts at rank 0 owns the deployment's handshake.
+    pub fn frontend_engine_start_index(&self) -> u32 {
+        if self.local_engines_only {
+            self.start_rank
+        } else {
+            0
+        }
+    }
+}
+
+/// Data parallel load-balancing arguments of `vllm-rs serve`: which ranks this
+/// node serves and whether its frontend serves only them.
+#[derive(Debug, Clone, Default, Args, PartialEq, Eq)]
+pub struct DataParallelLbArgs {
+    /// Rank of the data parallel group served by this process. Implies
+    /// `--data-parallel-external-lb`.
+    #[arg(long, env = "VLLM_DP_RANK")]
+    pub data_parallel_rank: Option<u32>,
+    /// First data parallel rank on this node. Implies
+    /// `--data-parallel-hybrid-lb` unless `--headless` is set.
+    #[arg(long)]
+    pub data_parallel_start_rank: Option<u32>,
+    /// Whether to use "hybrid" DP LB mode: this node's frontend load
+    /// balances between its local data parallel ranks, while an external LB
+    /// balances between nodes. Set explicitly in conjunction with
+    /// `--data-parallel-size-local` and `--data-parallel-start-rank`.
+    #[arg(long)]
+    pub data_parallel_hybrid_lb: bool,
+    /// Whether to use "external" DP LB mode: every data parallel rank runs its
+    /// own frontend, as in a "one-pod-per-rank" wide-EP setup in Kubernetes.
+    /// Set implicitly when `--data-parallel-rank` is provided.
+    #[arg(long)]
+    pub data_parallel_external_lb: bool,
+}
+
+impl DataParallelLbArgs {
+    /// Resolve the data parallel ranks of this process from the DP LB flags,
+    /// which keep their `vllm serve` meaning.
+    pub fn resolve(
+        &self,
+        engine: &ManagedEngineArgs,
+        headless: bool,
+    ) -> anyhow::Result<DataParallelTopology> {
+        let data_parallel_size = engine.data_parallel_size;
+        if self.data_parallel_hybrid_lb && self.data_parallel_external_lb {
+            bail!(
+                "`--data-parallel-hybrid-lb` and `--data-parallel-external-lb` cannot be enabled together"
+            );
+        }
+        if headless && self.data_parallel_hybrid_lb {
+            bail!("`--data-parallel-hybrid-lb` is not applicable with `--headless`");
+        }
+        if self.data_parallel_rank.is_some() && self.data_parallel_start_rank.is_some() {
+            bail!("`--data-parallel-rank` and `--data-parallel-start-rank` cannot be set together");
+        }
+
+        let topology = if self.data_parallel_external_lb || self.data_parallel_rank.is_some() {
+            let Some(rank) = self.data_parallel_rank else {
+                bail!("`--data-parallel-external-lb` requires `--data-parallel-rank`");
+            };
+            if !matches!(engine.data_parallel_size_local, None | Some(1)) {
+                bail!("an external data parallel rank requires `--data-parallel-size-local 1`");
+            }
+            DataParallelTopology {
+                start_rank: rank,
+                local_size: 1,
+                data_parallel_size,
+                local_engines_only: true,
+            }
+        } else if let Some(local_size) = engine.data_parallel_size_local {
+            // A start rank outside `--headless` implies hybrid LB, and hybrid LB
+            // with every rank on this node is internal LB. Hybrid LB with one
+            // local rank is external LB, which connects the same way.
+            let hybrid_lb = (self.data_parallel_hybrid_lb
+                || (self.data_parallel_start_rank.is_some() && !headless))
+                && local_size != data_parallel_size;
+            DataParallelTopology {
+                start_rank: self.data_parallel_start_rank.unwrap_or(0),
+                local_size,
+                data_parallel_size,
+                local_engines_only: hybrid_lb,
+            }
+        } else {
+            if self.data_parallel_hybrid_lb {
+                bail!("`--data-parallel-hybrid-lb` requires `--data-parallel-size-local`");
+            }
+            // Without a local size, the start rank is ignored.
+            DataParallelTopology {
+                start_rank: 0,
+                local_size: data_parallel_size,
+                data_parallel_size,
+                local_engines_only: false,
+            }
+        };
+
+        if topology.start_rank as usize + topology.local_size > data_parallel_size {
+            bail!(
+                "data parallel ranks [{}, {}) of this node exceed `--data-parallel-size {data_parallel_size}`",
+                topology.start_rank,
+                topology.start_rank as usize + topology.local_size
+            );
+        }
+        if topology.local_engines_only {
+            if topology.local_size == 0 {
+                bail!("hybrid and external DP LB require local data parallel ranks");
+            }
+            // Every node dials the handshake of the frontend owning rank 0.
+            if engine.handshake_port.is_none() {
+                bail!("hybrid and external DP LB require `--data-parallel-rpc-port`");
+            }
+        }
+        Ok(topology)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use clap::Parser;
+
+    use super::DataParallelLbArgs;
+    use crate::cli::ManagedEngineArgs;
+
+    /// The `vllm-rs serve` flags that the topology is resolved from.
+    #[derive(Parser)]
+    struct ServeFlags {
+        #[command(flatten)]
+        engine: ManagedEngineArgs,
+        #[command(flatten)]
+        data_parallel_lb: DataParallelLbArgs,
+        #[arg(long)]
+        headless: bool,
+    }
+
+    #[test]
+    fn resolve_follows_vllm_serve_dp_lb_modes() {
+        let cases: &[&[&str]] = &[
+            &["--data-parallel-size", "4"],
+            &[
+                "--data-parallel-size",
+                "4",
+                "--data-parallel-size-local",
+                "2",
+            ],
+            &[
+                "--data-parallel-size",
+                "4",
+                "--data-parallel-size-local",
+                "2",
+                "--data-parallel-start-rank",
+                "2",
+                "--data-parallel-rpc-port",
+                "13345",
+            ],
+            &[
+                "--data-parallel-size",
+                "4",
+                "--data-parallel-size-local",
+                "2",
+                "--data-parallel-hybrid-lb",
+                "--data-parallel-rpc-port",
+                "13345",
+            ],
+            &[
+                "--headless",
+                "--data-parallel-size",
+                "4",
+                "--data-parallel-size-local",
+                "2",
+                "--data-parallel-start-rank",
+                "2",
+            ],
+            &[
+                "--data-parallel-size",
+                "4",
+                "--data-parallel-rank",
+                "3",
+                "--data-parallel-rpc-port",
+                "1",
+            ],
+            &[
+                "--data-parallel-size",
+                "4",
+                "--data-parallel-start-rank",
+                "2",
+            ],
+            &["--data-parallel-size", "4", "--data-parallel-external-lb"],
+            &[
+                "--data-parallel-size",
+                "4",
+                "--data-parallel-rank",
+                "1",
+                "--data-parallel-start-rank",
+                "1",
+            ],
+            &[
+                "--data-parallel-size",
+                "4",
+                "--data-parallel-hybrid-lb",
+                "--data-parallel-external-lb",
+            ],
+            &[
+                "--data-parallel-size",
+                "4",
+                "--data-parallel-size-local",
+                "2",
+                "--data-parallel-start-rank",
+                "2",
+            ],
+            &[
+                "--data-parallel-size",
+                "4",
+                "--data-parallel-size-local",
+                "2",
+                "--data-parallel-rank",
+                "3",
+            ],
+            &[
+                "--data-parallel-size",
+                "4",
+                "--data-parallel-size-local",
+                "2",
+                "--data-parallel-start-rank",
+                "3",
+                "--data-parallel-rpc-port",
+                "1",
+            ],
+        ];
+        let resolved = cases
+            .iter()
+            .map(|flags| {
+                let parsed = ServeFlags::try_parse_from(["vllm-rs"].iter().chain(*flags)).unwrap();
+                let result = match parsed.data_parallel_lb.resolve(&parsed.engine, parsed.headless)
+                {
+                    Ok(topology) => format!("{topology:?}"),
+                    Err(error) => format!("error: {error}"),
+                };
+                format!("{}\n  {result}", flags.join(" "))
+            })
+            .collect::<Vec<_>>()
+            .join("\n");
+
+        expect_test::expect![[r#"
+        --data-parallel-size 4
+          DataParallelTopology { start_rank: 0, local_size: 4, data_parallel_size: 4, local_engines_only: false }
+        --data-parallel-size 4 --data-parallel-size-local 2
+          DataParallelTopology { start_rank: 0, local_size: 2, data_parallel_size: 4, local_engines_only: false }
+        --data-parallel-size 4 --data-parallel-size-local 2 --data-parallel-start-rank 2 --data-parallel-rpc-port 13345
+          DataParallelTopology { start_rank: 2, local_size: 2, data_parallel_size: 4, local_engines_only: true }
+        --data-parallel-size 4 --data-parallel-size-local 2 --data-parallel-hybrid-lb --data-parallel-rpc-port 13345
+          DataParallelTopology { start_rank: 0, local_size: 2, data_parallel_size: 4, local_engines_only: true }
+        --headless --data-parallel-size 4 --data-parallel-size-local 2 --data-parallel-start-rank 2
+          DataParallelTopology { start_rank: 2, local_size: 2, data_parallel_size: 4, local_engines_only: false }
+        --data-parallel-size 4 --data-parallel-rank 3 --data-parallel-rpc-port 1
+          DataParallelTopology { start_rank: 3, local_size: 1, data_parallel_size: 4, local_engines_only: true }
+        --data-parallel-size 4 --data-parallel-start-rank 2
+          DataParallelTopology { start_rank: 0, local_size: 4, data_parallel_size: 4, local_engines_only: false }
+        --data-parallel-size 4 --data-parallel-external-lb
+          error: `--data-parallel-external-lb` requires `--data-parallel-rank`
+        --data-parallel-size 4 --data-parallel-rank 1 --data-parallel-start-rank 1
+          error: `--data-parallel-rank` and `--data-parallel-start-rank` cannot be set together
+        --data-parallel-size 4 --data-parallel-hybrid-lb --data-parallel-external-lb
+          error: `--data-parallel-hybrid-lb` and `--data-parallel-external-lb` cannot be enabled together
+        --data-parallel-size 4 --data-parallel-size-local 2 --data-parallel-start-rank 2
+          error: hybrid and external DP LB require `--data-parallel-rpc-port`
+        --data-parallel-size 4 --data-parallel-size-local 2 --data-parallel-rank 3
+          error: an external data parallel rank requires `--data-parallel-size-local 1`
+        --data-parallel-size 4 --data-parallel-size-local 2 --data-parallel-start-rank 3 --data-parallel-rpc-port 1
+          error: data parallel ranks [3, 5) of this node exceed `--data-parallel-size 4`"#]].assert_eq(&resolved);
+    }
+}
diff --git a/rust/src/managed-engine/src/lib.rs b/rust/src/managed-engine/src/lib.rs
index 895a9f18c786..2e78153d9d41 100644
--- a/rust/src/managed-engine/src/lib.rs
+++ b/rust/src/managed-engine/src/lib.rs
@@ -2,6 +2,7 @@
 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 pub mod cli;
+pub mod data_parallel;
 mod process;
 
 pub use process::{ManagedEngineConfig, ManagedEngineHandle, allocate_handshake_port};
diff --git a/rust/src/managed-engine/src/process.rs b/rust/src/managed-engine/src/process.rs
index 1f0b9f86b2ce..6c0a699604a0 100644
--- a/rust/src/managed-engine/src/process.rs
+++ b/rust/src/managed-engine/src/process.rs
@@ -44,6 +44,11 @@ pub struct ManagedEngineConfig {
     /// The per-node replica count is forwarded separately in `python_args` as
     /// `--data-parallel-size-local`.
     pub data_parallel_size: usize,
+    /// Handshake endpoint of the colocated frontend when it serves only this
+    /// node's engines (hybrid or external DP LB). The engines then also
+    /// handshake there for their input/output addresses, after fetching the
+    /// coordinator and DP rendezvous from the frontend owning rank 0.
+    pub client_handshake_address: Option<String>,
     /// Extra CLI arguments forwarded verbatim to Python vLLM.
     pub python_args: Vec<String>,
 }
@@ -76,8 +81,11 @@ impl ManagedEngineConfig {
             .arg("--data-parallel-rpc-port")
             .arg(self.handshake_port.to_string())
             .arg("--data-parallel-size")
-            .arg(self.data_parallel_size.to_string())
-            .args(&self.python_args);
+            .arg(self.data_parallel_size.to_string());
+        if let Some(client_handshake_address) = &self.client_handshake_address {
+            command.arg("--headless-client-handshake-address").arg(client_handshake_address);
+        }
+        command.args(&self.python_args);
         command.envs(Self::rust_frontend_envs());
         command
     }
@@ -227,6 +235,7 @@ mod tests {
             handshake_host: "127.0.0.1".to_string(),
             handshake_port: 62100,
             data_parallel_size: 4,
+            client_handshake_address: None,
             python_args: vec![
                 "--data-parallel-size-local".to_string(),
                 "2".to_string(),
diff --git a/rust/src/metrics/src/scheduler.rs b/rust/src/metrics/src/scheduler.rs
index 05392a2cf56e..4383a32bf2c4 100644
--- a/rust/src/metrics/src/scheduler.rs
+++ b/rust/src/metrics/src/scheduler.rs
@@ -91,6 +91,22 @@ pub struct WaitingReasonLabels {
     pub reason: &'static str,
 }
 
+/// Labels for `vllm:engine_sleep_state`: one series per engine and state,
+/// `awake`, `weights_offloaded` (sleep level 1) or `discard_all` (level 2).
+#[derive(Clone, Debug, Hash, PartialEq, Eq, EncodeLabelSet)]
+pub struct EngineSleepStateLabels {
+    pub model_name: String,
+    pub engine: u32,
+    pub sleep_state: &'static str,
+}
+
+#[derive(Clone, Debug, Hash, PartialEq, Eq, EncodeLabelSet)]
+pub struct KvFetchStageLabels {
+    pub model_name: String,
+    pub engine: u32,
+    pub stage: &'static str,
+}
+
 /// Labels for per-operation Mooncake store connector telemetry. `operation`
 /// and `status` are dynamic (not known ahead of time), unlike the other
 /// fixed-label families in this module.
@@ -109,6 +125,19 @@ pub type MooncakeOperationHistogramFamily =
     Family<MooncakeOperationLabels, Histogram, fn() -> Histogram>;
 pub type MooncakeOperationCounterFamily = Family<MooncakeOperationLabels, U64Counter>;
 
+/// Labels for per-operation UMBP store connector telemetry.
+#[derive(Clone, Debug, Hash, PartialEq, Eq, EncodeLabelSet)]
+pub struct UmbpOperationLabels {
+    pub model_name: String,
+    pub engine: u32,
+    pub operation: &'static str,
+}
+
+/// Family type alias for the UMBP connector telemetry, exposed so
+/// `engine-core-client` can hold onto the `Family` directly, since an
+/// operation's series appears once the connector first reports it.
+pub type UmbpOperationCounterFamily = Family<UmbpOperationLabels, U64Counter>;
+
 /// Adapter names encoded as a deterministic comma-joined Prometheus label value.
 #[derive(Clone, Debug, Hash, PartialEq, Eq)]
 pub struct LoraAdapterNames(pub BTreeSet<String>);
@@ -216,8 +245,13 @@ pub struct SchedulerMetrics {
     pub scheduler_running: Family<EngineLabels, U64Gauge>,
     pub scheduler_waiting: Family<EngineLabels, U64Gauge>,
     pub scheduler_waiting_by_reason: Family<WaitingReasonLabels, U64Gauge>,
+    pub scheduler_kv_fetch_by_stage: Family<KvFetchStageLabels, U64Gauge>,
     pub kv_cache_usage: Family<EngineLabels, F64Gauge>,
 
+    /// `vllm:engine_sleep_state`, set by the engine client's sleep and wake
+    /// calls rather than from `SchedulerStats`.
+    pub engine_sleep_state: Family<EngineSleepStateLabels, U64Gauge>,
+
     /// `vllm:lora_requests_info`. Value is the emit-time unix timestamp in
     /// seconds.
     pub lora_info: Family<LoraInfoLabels, F64Gauge>,
@@ -270,6 +304,12 @@ pub struct SchedulerMetrics {
     pub nixl_num_kv_expired_reqs: Family<EngineLabels, U64Counter>,
     pub nixl_num_notifications_after_expiry: Family<EngineLabels, U64Counter>,
 
+    // UMBP store connector telemetry. Mirrors `UMBPStorePromMetrics`.
+    pub umbp_transfer_submitted_total: UmbpOperationCounterFamily,
+    pub umbp_transfer_completed_total: UmbpOperationCounterFamily,
+    pub umbp_transfer_failed_total: UmbpOperationCounterFamily,
+    pub umbp_transfer_bytes_total: UmbpOperationCounterFamily,
+
     /// Non-Prometheus interval accumulators for periodic text-log helpers.
     pub log_stats: Family<EngineLabels, SchedulerLogStatsAccumulator>,
 }
@@ -303,6 +343,16 @@ impl SchedulerMetrics {
             scheduler_waiting_by_reason.clone(),
         );
 
+        let scheduler_kv_fetch_by_stage = Family::default();
+        registry.register(
+            "vllm:num_requests_kv_fetch_by_stage",
+            "Number of waiting requests by async KV load stage. \
+             Stage labels: 'waiting_to_start' = needs an async KV load that has not started; \
+             'in_progress' = load started and not yet reported finished; \
+             'completed_waiting' = load finished, request not running yet.",
+            scheduler_kv_fetch_by_stage.clone(),
+        );
+
         let kv_cache_usage = Family::default();
         registry.register(
             "vllm:kv_cache_usage_perc",
@@ -310,6 +360,16 @@ impl SchedulerMetrics {
             kv_cache_usage.clone(),
         );
 
+        let engine_sleep_state = Family::default();
+        registry.register(
+            "vllm:engine_sleep_state",
+            "Engine sleep state; awake = 0 means engine is sleeping; \
+             awake = 1 means engine is awake; \
+             weights_offloaded = 1 means sleep level 1; \
+             discard_all = 1 means sleep level 2.",
+            engine_sleep_state.clone(),
+        );
+
         let lora_info = Family::default();
         registry.register(
             "vllm:lora_requests_info",
@@ -591,11 +651,42 @@ impl SchedulerMetrics {
             nixl_num_notifications_after_expiry.clone(),
         );
 
+        // UMBP store connector telemetry.
+        let umbp_transfer_submitted_total = Family::default();
+        registry.register(
+            "vllm:umbp_transfer_submitted",
+            "Number of UMBP objects submitted",
+            umbp_transfer_submitted_total.clone(),
+        );
+
+        let umbp_transfer_completed_total = Family::default();
+        registry.register(
+            "vllm:umbp_transfer_completed",
+            "Number of UMBP objects completed",
+            umbp_transfer_completed_total.clone(),
+        );
+
+        let umbp_transfer_failed_total = Family::default();
+        registry.register(
+            "vllm:umbp_transfer_failed",
+            "Number of UMBP objects failed",
+            umbp_transfer_failed_total.clone(),
+        );
+
+        let umbp_transfer_bytes_total = Family::default();
+        registry.register(
+            "vllm:umbp_transfer_bytes",
+            "Bytes in successfully completed UMBP transfer ranges",
+            umbp_transfer_bytes_total.clone(),
+        );
+
         Self {
             scheduler_running,
             scheduler_waiting,
             scheduler_waiting_by_reason,
+            scheduler_kv_fetch_by_stage,
             kv_cache_usage,
+            engine_sleep_state,
             lora_info,
             prefix_cache_queries,
             prefix_cache_hits,
@@ -631,6 +722,10 @@ impl SchedulerMetrics {
             nixl_num_failed_notifications,
             nixl_num_kv_expired_reqs,
             nixl_num_notifications_after_expiry,
+            umbp_transfer_submitted_total,
+            umbp_transfer_completed_total,
+            umbp_transfer_failed_total,
+            umbp_transfer_bytes_total,
             log_stats: Family::default(),
         }
     }
diff --git a/rust/src/mock-engine/src/engine.rs b/rust/src/mock-engine/src/engine.rs
index 8685357b1896..1d75d65c28b1 100644
--- a/rust/src/mock-engine/src/engine.rs
+++ b/rust/src/mock-engine/src/engine.rs
@@ -94,6 +94,7 @@ fn utility_response(
     let result = match request.method_name.as_str() {
         "get_supported_tasks" => utility_envelope(vec!["generate"]),
         "is_sleeping" => utility_envelope(false),
+        "is_inference_ready" => utility_envelope(true),
         "reset_prefix_cache" => utility_envelope(true),
         "wake_up" => utility_envelope(true),
         "reset_mm_cache" | "reset_encoder_cache" | "profile" | "sleep" | "execute_dummy_batch" => {
diff --git a/rust/src/mock-engine/src/tests.rs b/rust/src/mock-engine/src/tests.rs
index 0b053cbbb16a..56ae8724d268 100644
--- a/rust/src/mock-engine/src/tests.rs
+++ b/rust/src/mock-engine/src/tests.rs
@@ -21,7 +21,9 @@ fn client_config(handshake_address: String, engine_count: usize) -> EngineCoreCl
         transport_mode: TransportMode::HandshakeOwner {
             handshake_address,
             advertised_host: "127.0.0.1".to_string(),
+            engine_start_index: 0,
             engine_count,
+            data_parallel_size: engine_count,
             ready_timeout: Duration::from_secs(5),
             local_input_address: None,
             local_output_address: None,
diff --git a/rust/src/server/examples/external_engine_openai_qwen.rs b/rust/src/server/examples/external_engine_openai_qwen.rs
index 3ed107f6590a..e8fff2dfdcca 100644
--- a/rust/src/server/examples/external_engine_openai_qwen.rs
+++ b/rust/src/server/examples/external_engine_openai_qwen.rs
@@ -20,8 +20,8 @@ use tracing_subscriber::EnvFilter;
 use vllm_engine_core_client::TransportMode;
 use vllm_engine_core_client::protocol::structured_outputs::StructuredOutputBackend;
 use vllm_server::{
-    ApiServerOptions, ChatTemplateContentFormatOption, Config, CoordinatorMode, CorsConfig,
-    HttpListenerMode, ParserSelection, RendererSelection, serve,
+    AccessLogConfig, ApiServerOptions, ChatTemplateContentFormatOption, Config, CoordinatorMode,
+    CorsConfig, HttpListenerMode, ParserSelection, RendererSelection, serve,
 };
 
 #[derive(Debug, Parser)]
@@ -55,7 +55,9 @@ async fn main() -> Result<()> {
         transport_mode: TransportMode::HandshakeOwner {
             handshake_address: args.handshake_address,
             advertised_host: args.host,
+            engine_start_index: 0,
             engine_count: args.engine_count,
+            data_parallel_size: args.engine_count,
             ready_timeout: Duration::from_secs(args.ready_timeout_secs),
             local_input_address: None,
             local_output_address: None,
@@ -91,7 +93,9 @@ async fn main() -> Result<()> {
         stream_interval: NonZeroU32::MIN,
         api_server_options: ApiServerOptions::default(),
         cors: CorsConfig::default(),
+        access_log: AccessLogConfig::default(),
         tls: None,
+        root_path: None,
         api_keys: Vec::new(),
         disable_log_stats: false,
         custom_histogram_buckets: Default::default(),
diff --git a/rust/src/server/src/config.rs b/rust/src/server/src/config.rs
index 845b1baa3efb..afa3fb639b47 100644
--- a/rust/src/server/src/config.rs
+++ b/rust/src/server/src/config.rs
@@ -65,8 +65,9 @@ pub enum CoordinatorMode {
 pub struct ApiServerOptions {
     /// Log a summary line for each completed request.
     pub enable_log_requests: bool,
-    /// Which generated outputs to log (`--enable-log-outputs`).
-    pub log_outputs: OutputLogging,
+    /// Which generated outputs to log (`--enable-log-outputs`), and how much
+    /// of each (`--max-log-len`).
+    pub log_outputs: OutputLogConfig,
     /// When `true`, include prompt token cache details in response usage.
     pub enable_prompt_tokens_details: bool,
     /// Format logprob tokens as `token_id:{id}` unless the request sets
@@ -80,6 +81,9 @@ pub struct ApiServerOptions {
     pub enable_scale_out: bool,
     /// When `true`, register `/tokenizer_info`.
     pub enable_tokenizer_info_endpoint: bool,
+    /// Accept any requested model name, served by the base model unless it
+    /// names a loaded LoRA adapter (`VLLM_SKIP_MODEL_NAME_VALIDATION`).
+    pub skip_model_name_validation: 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>,
@@ -115,6 +119,49 @@ impl OutputLogging {
     }
 }
 
+/// Generated-output logging, like Python's `RequestLogger` used for
+/// `--enable-log-outputs`: which outputs to log and how much of each.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Default)]
+pub struct OutputLogConfig {
+    /// Which generated outputs to log.
+    pub outputs: OutputLogging,
+    /// Maximum number of characters or token IDs of each logged output
+    /// (`--max-log-len`); `None` logs them in full.
+    pub max_log_len: Option<usize>,
+}
+
+impl OutputLogConfig {
+    /// Whether the complete output of each response is logged.
+    pub fn logs_complete(self) -> bool {
+        self.outputs.logs_complete()
+    }
+
+    /// Whether each streamed delta is logged.
+    pub fn logs_deltas(self) -> bool {
+        self.outputs.logs_deltas()
+    }
+}
+
+/// HTTP access log settings.
+#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
+pub struct AccessLogConfig {
+    /// Whether to log one line per HTTP response (`--disable-access-log`
+    /// turns it off).
+    pub enabled: bool,
+    /// Request paths whose responses are not logged, matched exactly against
+    /// the path without its query string (`--disable-access-log-for-endpoints`).
+    pub excluded_paths: Vec<String>,
+}
+
+impl Default for AccessLogConfig {
+    fn default() -> Self {
+        Self {
+            enabled: true,
+            excluded_paths: Vec::new(),
+        }
+    }
+}
+
 /// CORS settings mirroring Python's `CORSMiddleware`; origins are denied by default.
 #[derive(Debug, Clone, PartialEq, Eq, Serialize)]
 pub struct CorsConfig {
@@ -333,9 +380,14 @@ pub struct Config {
     pub api_server_options: ApiServerOptions,
     /// CORS settings applied to every HTTP response.
     pub cors: CorsConfig,
+    /// HTTP access log settings.
+    pub access_log: AccessLogConfig,
     /// TLS settings. `None` serves plaintext HTTP; `Some` terminates TLS at the
     /// listener.
     pub tls: Option<TlsConfig>,
+    /// Path prefix (`--root-path`) under which every route is also served,
+    /// for a server behind a path-based routing proxy. Never empty.
+    pub root_path: Option<String>,
     /// API keys accepted as bearer tokens for guarded routes.
     #[serde(skip_serializing)]
     #[educe(Debug(method(fmt_redacted_api_keys)))]
@@ -424,7 +476,13 @@ impl Config {
         match &self.coordinator_mode {
             CoordinatorMode::None => None,
             CoordinatorMode::MaybeInProc => {
-                if model_is_moe && self.engine_count() > 1 {
+                // Coordinated DP runs one coordinator, hosted by the frontend
+                // owning rank 0; other frontends learn its address from their
+                // engines' READY `dp_stats_address`.
+                if model_is_moe
+                    && self.transport_mode.data_parallel_size() > 1
+                    && self.transport_mode.engine_start_index() == 0
+                {
                     Some(EngineCoreCoordinatorMode::InProc)
                 } else {
                     None
diff --git a/rust/src/server/src/error.rs b/rust/src/server/src/error.rs
index b53ca55a66e4..802741ed7af2 100644
--- a/rust/src/server/src/error.rs
+++ b/rust/src/server/src/error.rs
@@ -21,6 +21,8 @@ pub enum ApiError {
     },
     /// The requested model name does not match the single configured model.
     ModelNotFound { model: String },
+    /// A resource other than the model, such as a LoRA adapter, does not exist.
+    NotFound { message: String },
     /// The request body could not be parsed as valid JSON.
     JsonParseError { message: String },
     /// The requested operation conflicts with current server state.
@@ -36,7 +38,7 @@ impl ApiError {
     pub fn status_code(&self) -> StatusCode {
         match self {
             Self::InvalidRequest { .. } => StatusCode::BAD_REQUEST,
-            Self::ModelNotFound { .. } => StatusCode::NOT_FOUND,
+            Self::ModelNotFound { .. } | Self::NotFound { .. } => StatusCode::NOT_FOUND,
             Self::Conflict { .. } => StatusCode::CONFLICT,
             Self::ServerError { .. } => StatusCode::INTERNAL_SERVER_ERROR,
             Self::JsonParseError { .. } => StatusCode::BAD_REQUEST,
@@ -58,6 +60,7 @@ impl ApiError {
                 format!("The model `{model}` does not exist."),
                 Some("model".to_string()),
             ),
+            Self::NotFound { message } => ("NotFoundError", message.clone(), None),
             // Python reports a body that fails to parse into the request
             // model through FastAPI's validation handler, which names the
             // HTTP status instead of an error class.
@@ -141,6 +144,7 @@ mod tests {
         let errors = [
             ApiError::invalid_request("top_p must be in (0, 1].", Some("top_p")),
             ApiError::model_not_found("missing"),
+            ApiError::not_found("No adapter found for /not/an/adapter"),
             ApiError::json_parse_error("expected value at line 1 column 1"),
             ApiError::conflict("A profiling session is already active."),
             ApiError::server_error("Internal server error"),
@@ -153,6 +157,7 @@ mod tests {
         expect![[r#"
             {"error":{"message":"top_p must be in (0, 1].","type":"BadRequestError","param":"top_p","code":400}}
             {"error":{"message":"The model `missing` does not exist.","type":"NotFoundError","param":"model","code":404}}
+            {"error":{"message":"No adapter found for /not/an/adapter","type":"NotFoundError","param":null,"code":404}}
             {"error":{"message":"expected value at line 1 column 1","type":"Bad Request","param":null,"code":400}}
             {"error":{"message":"A profiling session is already active.","type":"Conflict","param":null,"code":409}}
             {"error":{"message":"Internal server error","type":"InternalServerError","param":null,"code":500}}
diff --git a/rust/src/server/src/grpc/control.rs b/rust/src/server/src/grpc/control.rs
index 482a7d1d6301..d522c7cc8e35 100644
--- a/rust/src/server/src/grpc/control.rs
+++ b/rust/src/server/src/grpc/control.rs
@@ -10,7 +10,7 @@ use tokio_util::sync::CancellationToken;
 use tonic::{Code, Request, Response, Status};
 use vllm_engine_core_client::EngineCoreClient;
 use vllm_engine_core_client::protocol::handshake::EngineCoreReadyResponse;
-use vllm_engine_core_client::protocol::utility::PauseMode as EnginePauseMode;
+use vllm_engine_core_client::protocol::utility::{PauseMode as EnginePauseMode, SleepLevel};
 
 use super::{ControlServer, pb};
 use crate::config::LoraModulePath;
@@ -159,6 +159,7 @@ fn load_lora_status(error: LoadLoraError) -> Status {
     let code = match &error {
         LoadLoraError::Disabled(_) => Code::FailedPrecondition,
         LoadLoraError::InvalidAdapter { .. } => Code::InvalidArgument,
+        LoadLoraError::PathAccess(LoraPathAccessError::NotFound { .. }) => Code::NotFound,
         LoadLoraError::PathAccess(LoraPathAccessError::InvalidPath { .. }) => Code::InvalidArgument,
         LoadLoraError::PathAccess(LoraPathAccessError::InvalidConfiguration { .. }) => {
             Code::Internal
@@ -377,7 +378,10 @@ impl pb::control_server::Control for ControlServiceImpl {
         self.require_sleep_mode()?;
         let request = request.into_inner();
         let mode = pause_mode(request.mode)?;
-        let level = request.level.unwrap_or(1);
+        let level = match request.level {
+            Some(level) => SleepLevel::try_from(level).map_err(Status::invalid_argument)?,
+            None => SleepLevel::OffloadWeights,
+        };
         let _guard = self.rl_lock.lock().await;
         self.client()
             .sleep(level, mode)
@@ -539,3 +543,19 @@ pub(super) fn kv_event_source(response: &EngineCoreReadyResponse) -> Option<pb::
         max_queue_size: config.max_queue_size,
     })
 }
+
+#[cfg(test)]
+mod tests {
+    use tonic::Code;
+
+    use super::load_lora_status;
+    use crate::lora::{LoadLoraError, LoraPathAccessError};
+
+    #[test]
+    fn load_lora_reports_missing_adapter_path_as_not_found() {
+        let status = load_lora_status(LoadLoraError::PathAccess(LoraPathAccessError::NotFound {
+            lora_path: "/not/an/adapter".to_string(),
+        }));
+        assert_eq!(status.code(), Code::NotFound);
+    }
+}
diff --git a/rust/src/server/src/grpc/convert.rs b/rust/src/server/src/grpc/convert.rs
index f43af34fc551..dd7c6cc772b2 100644
--- a/rust/src/server/src/grpc/convert.rs
+++ b/rust/src/server/src/grpc/convert.rs
@@ -16,6 +16,7 @@ use vllm_text::{
 };
 
 use super::pb;
+use crate::model_names::ServedModelNames;
 
 fn kv_hints_from_proto(hints: pb::KvHintsEnvelope) -> KvHintsEnvelope {
     KvHintsEnvelope {
@@ -49,15 +50,15 @@ fn kv_hints_from_proto(hints: pb::KvHintsEnvelope) -> KvHintsEnvelope {
 
 /// Convert a gRPC `GenerateRequest` into the internal `TextRequest`.
 ///
-/// If `req.model` is non-empty, it must match one of `served_model_names`;
-/// otherwise the request is rejected with `NotFound`. An empty string is
-/// treated as "unset" (proto3 default) and accepted.
+/// `req.model` must be accepted by `served_model_names`; otherwise the request
+/// is rejected with `NotFound`. An empty string is treated as "unset" (proto3
+/// default) and accepted.
 pub fn to_text_request(
     req: pb::GenerateRequest,
     stream: bool,
-    served_model_names: &[String],
+    served_model_names: &ServedModelNames,
 ) -> Result<TextRequest, Status> {
-    if !req.model.is_empty() && !served_model_names.iter().any(|n| n == &req.model) {
+    if !served_model_names.accepts(Some(&req.model)) {
         return Err(Status::not_found(format!(
             "model `{}` not found",
             req.model
@@ -580,6 +581,11 @@ mod tests {
     use super::{
         ResponseOpts, json_to_proto_struct, pb, to_finish_info, to_sequence_output, to_text_request,
     };
+    use crate::model_names::ServedModelNames;
+
+    fn served() -> ServedModelNames {
+        ServedModelNames::new(vec!["test-model".to_string()])
+    }
 
     fn base_request() -> pb::GenerateRequest {
         pb::GenerateRequest {
@@ -605,7 +611,7 @@ mod tests {
         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()]).unwrap();
+            let text = to_text_request(decoded, stream, &served()).unwrap();
             let engine = vllm_text::lower_text_request(
                 text,
                 vec![1],
@@ -669,7 +675,7 @@ mod tests {
             let encoded = request.encode_to_vec();
             for stream in [false, true] {
                 let request = pb::GenerateRequest::decode(encoded.as_slice()).unwrap();
-                let text = to_text_request(request, stream, &["test-model".to_string()])
+                let text = to_text_request(request, stream, &served())
                     .expect("reasoning controls accepted");
                 let engine = vllm_text::lower_text_request(
                     text,
@@ -708,8 +714,7 @@ mod tests {
             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");
+                let text = to_text_request(decoded, stream, &served()).expect("convert request");
                 assert_eq!(text.sampling_params.watermarking, watermarking);
             }
         }
@@ -721,7 +726,7 @@ mod tests {
             temperature: Some(0.7),
             ..base_request()
         };
-        let text = to_text_request(req, false, &["test-model".to_string()]).expect("convert ok");
+        let text = to_text_request(req, false, &served()).expect("convert ok");
         assert_eq!(text.sampling_params.temperature, Some(0.7));
     }
 
@@ -741,7 +746,7 @@ mod tests {
             ..base_request()
         };
 
-        let text = to_text_request(req, false, &["test-model".to_string()]).expect("convert ok");
+        let text = to_text_request(req, false, &served()).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");
@@ -756,8 +761,7 @@ mod tests {
 
     #[test]
     fn unset_temperature_defaults_to_greedy() {
-        let text = to_text_request(base_request(), false, &["test-model".to_string()])
-            .expect("convert ok");
+        let text = to_text_request(base_request(), false, &served()).expect("convert ok");
         // The gRPC API defaults to greedy (0.0) when temperature is not specified.
         assert_eq!(text.sampling_params.temperature, Some(0.0));
     }
@@ -771,7 +775,7 @@ mod tests {
             }),
             ..base_request()
         };
-        let text = to_text_request(req, false, &["test-model".to_string()]).expect("convert ok");
+        let text = to_text_request(req, false, &served()).expect("convert ok");
         assert_eq!(text.sampling_params.seed, None);
     }
 
@@ -784,7 +788,7 @@ mod tests {
             }),
             ..base_request()
         };
-        let text = to_text_request(req, false, &["test-model".to_string()]).expect("convert ok");
+        let text = to_text_request(req, false, &served()).expect("convert ok");
         assert_eq!(text.sampling_params.seed, Some(0));
     }
 
@@ -828,8 +832,7 @@ mod tests {
                 }),
                 ..base_request()
             };
-            let text =
-                to_text_request(req, false, &["test-model".to_string()]).expect("convert ok");
+            let text = to_text_request(req, false, &served()).expect("convert ok");
             let params = lower_sampling_params(
                 text.sampling_params,
                 SamplingHints::default(),
@@ -872,7 +875,7 @@ mod tests {
             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()]).unwrap();
+                let text = to_text_request(decoded, stream, &served()).unwrap();
                 let engine = vllm_text::lower_text_request(
                     text,
                     vec![1],
@@ -908,7 +911,7 @@ mod tests {
             }),
             ..base_request()
         };
-        let text = to_text_request(req, false, &["test-model".to_string()]).expect("convert ok");
+        let text = to_text_request(req, false, &served()).expect("convert ok");
         assert_eq!(text.sampling_params.skip_reading_prefix_cache, Some(true));
     }
 
@@ -921,7 +924,7 @@ mod tests {
             }),
             ..base_request()
         };
-        let text = to_text_request(req, false, &["test-model".to_string()]).expect("convert ok");
+        let text = to_text_request(req, false, &served()).expect("convert ok");
         assert_eq!(text.sampling_params.skip_reading_prefix_cache, None);
         // Prompt conversion still succeeds and reaches the expected variant.
         assert!(matches!(text.prompt, Prompt::Text(s) if s == "hi"));
@@ -939,6 +942,7 @@ mod tests {
             ec_transfer_params: None,
             sampling_mask: None,
             spec_decode_metrics: None,
+            weight_versions: None,
             timestamps: Default::default(),
         }
     }
diff --git a/rust/src/server/src/grpc/health.rs b/rust/src/server/src/grpc/health.rs
index cd35aec972e8..07f4dc41b01a 100644
--- a/rust/src/server/src/grpc/health.rs
+++ b/rust/src/server/src/grpc/health.rs
@@ -2,64 +2,267 @@
 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 use tokio::sync::watch;
-use tokio_util::sync::CancellationToken;
+use tokio_util::sync::{CancellationToken, WaitForCancellationFutureOwned};
 use tonic::server::NamedService;
 use tonic_health::ServingStatus;
 use tonic_health::server::HealthReporter;
-use tracing::{info, warn};
+use tracing::{debug, info, warn};
+use vllm_engine_core_client::EngineCoreClient;
 
 use super::{ControlGrpcService, InferenceGrpcService};
 
-pub(crate) async fn monitor_health(
-    mut health_reporter: HealthReporter,
-    mut engine_health: watch::Receiver<bool>,
-    shutdown: CancellationToken,
-) {
-    let inference_service = InferenceGrpcService::NAME;
-    let control_service = ControlGrpcService::NAME;
-    let status = ServingStatus::NotServing;
-    let health_event_first = tokio::select! {
-        result = engine_health.wait_for(|healthy| !*healthy) => {
-            match result {
-                Ok(_) => warn!(
-                    status = ?status,
-                    reason = "engine_unhealthy",
-                    "marking gRPC health services as not serving"
-                ),
-                Err(error) => warn!(
-                    %error,
-                    status = ?status,
-                    reason = "health_channel_closed",
-                    "engine health channel closed; marking gRPC health services as not serving"
-                ),
-            }
-            true
+/// Health service for liveness probes. It reports `NOT_SERVING` only after an
+/// engine failure: pause, sleep, weight updates, and orderly shutdown leave it
+/// `SERVING` until the listener closes.
+pub(crate) const LIVENESS_SERVICE: &str = "liveness";
+
+/// Health service for readiness probes. Like `vllm.Inference`, it reports
+/// whether the engines can accept and process new requests, and
+/// `NOT_SERVING` on shutdown or engine failure.
+pub(crate) const READINESS_SERVICE: &str = "readiness";
+
+/// The engine client signals that the gRPC health services follow.
+pub(crate) struct EngineHealth {
+    /// `false` once the engines fail, which is permanent.
+    pub liveness: watch::Receiver<bool>,
+    /// Whether the engines can accept and process new requests.
+    pub readiness: watch::Receiver<bool>,
+}
+
+impl EngineHealth {
+    pub fn subscribe(client: &EngineCoreClient) -> Self {
+        Self {
+            liveness: client.subscribe_health(),
+            readiness: client.subscribe_readiness(),
+        }
+    }
+}
+
+/// Publishes the gRPC health services from [`EngineHealth`].
+pub(crate) struct HealthMonitor {
+    reporter: HealthReporter,
+    engine: EngineHealth,
+    /// Readiness as last published to `readiness` and `vllm.Inference`.
+    ready: bool,
+}
+
+impl HealthMonitor {
+    /// Register every named health service with the engines' current
+    /// readiness. Call before the listener accepts connections so clients
+    /// never see an unregistered service.
+    pub async fn register(reporter: HealthReporter, mut engine: EngineHealth) -> Self {
+        let ready = *engine.readiness.borrow_and_update();
+        reporter.set_serving::<ControlGrpcService>().await;
+        reporter.set_service_status(LIVENESS_SERVICE, ServingStatus::Serving).await;
+        set_readiness(&reporter, ready).await;
+        Self {
+            reporter,
+            engine,
+            ready,
         }
-        _ = shutdown.cancelled() => {
+    }
+
+    /// Run a gRPC server until shutdown, keeping the health services in line
+    /// with the engines.
+    ///
+    /// `server` builds the server from the signal that stops its listener. On
+    /// shutdown, every service except `liveness` reports `NOT_SERVING` before
+    /// that signal fires and before `force` can abort the server, so watchers
+    /// see the withdrawal before their connection closes. Returns `None` if
+    /// `force` aborted the server.
+    pub async fn serve<S: Future>(
+        self,
+        shutdown: CancellationToken,
+        force: impl Future<Output = ()>,
+        server: impl FnOnce(WaitForCancellationFutureOwned) -> S,
+    ) -> Option<S::Output> {
+        // Cancelled by `run` once the withdrawal is published.
+        let withdrawn = CancellationToken::new();
+        let server = server(withdrawn.clone().cancelled_owned());
+        let server = async {
+            tokio::select! {
+                output = server => Some(output),
+                // `force` fires with `shutdown` under `--shutdown-timeout 0`.
+                () = async { force.await; withdrawn.cancelled().await } => None,
+            }
+        };
+        // The graceful drain ends once `run` removes the services, ending
+        // their watches.
+        let (output, ()) = tokio::join!(server, self.run(shutdown, &withdrawn));
+        output
+    }
+
+    /// Keep the health services in line with the engines until shutdown,
+    /// cancel `withdrawn`, then remove the services, which ends their watches.
+    ///
+    /// `readiness` and `vllm.Inference` follow the engines' readiness. On
+    /// shutdown, every service except `liveness` reports `NOT_SERVING`; on
+    /// engine failure, `liveness` does too.
+    async fn run(mut self, shutdown: CancellationToken, withdrawn: &CancellationToken) {
+        let inference_service = InferenceGrpcService::NAME;
+        let control_service = ControlGrpcService::NAME;
+        let status = ServingStatus::NotServing;
+        let mut readiness_open = true;
+        let engine_failed = loop {
+            let ready = tokio::select! {
+                result = self.engine.liveness.wait_for(|healthy| !*healthy) => {
+                    match result {
+                        Ok(_) => warn!(
+                            status = ?status,
+                            reason = "engine_unhealthy",
+                            "marking gRPC health services as not serving"
+                        ),
+                        Err(error) => warn!(
+                            %error,
+                            status = ?status,
+                            reason = "health_channel_closed",
+                            "engine health channel closed; marking gRPC health services as not serving"
+                        ),
+                    }
+                    break true;
+                }
+                _ = shutdown.cancelled() => {
+                    info!(
+                        status = ?status,
+                        reason = "server_shutdown",
+                        "server shutting down; marking gRPC health services except liveness as not serving"
+                    );
+                    break false;
+                }
+                changed = self.engine.readiness.changed(), if readiness_open => {
+                    readiness_open = changed.is_ok();
+                    *self.engine.readiness.borrow_and_update()
+                }
+            };
+            // Changes can coalesce back to the published value.
+            if ready != self.ready {
+                self.ready = ready;
+                debug!(ready, "engine readiness changed");
+                set_readiness(&self.reporter, ready).await;
+            }
+        };
+
+        let reporter = &mut self.reporter;
+        set_readiness(reporter, false).await;
+        reporter.set_not_serving::<ControlGrpcService>().await;
+        // Both gRPC services use the same engine client, so overall server health
+        // mirrors their shared engine health.
+        reporter.set_service_status("", status).await;
+
+        if engine_failed {
+            reporter.set_service_status(LIVENESS_SERVICE, status).await;
+            shutdown.cancelled().await;
             info!(
-                status = ?status,
                 reason = "server_shutdown",
-                "server shutting down; marking gRPC health services as not serving"
+                "server shutting down; closing gRPC health watches"
             );
-            false
         }
+        // Published, so the server may stop.
+        withdrawn.cancel();
+
+        reporter.clear_service_status(inference_service).await;
+        reporter.clear_service_status(control_service).await;
+        reporter.clear_service_status("").await;
+        reporter.clear_service_status(LIVENESS_SERVICE).await;
+        reporter.clear_service_status(READINESS_SERVICE).await;
+    }
+}
+
+async fn set_readiness(reporter: &HealthReporter, ready: bool) {
+    let status = if ready {
+        ServingStatus::Serving
+    } else {
+        ServingStatus::NotServing
     };
+    for service in [READINESS_SERVICE, InferenceGrpcService::NAME] {
+        reporter.set_service_status(service, status).await;
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::sync::{Arc, Mutex};
+
+    use futures::FutureExt as _;
+    use tonic::Request;
+    use tonic_health::pb::HealthCheckRequest;
+    use tonic_health::pb::health_server::Health as _;
+    use tonic_health::server::HealthService;
+
+    use super::*;
+
+    /// A monitor for ready, live engines, a way to read its statuses, and the
+    /// engine signals, which must outlive the test: a closed channel reads as
+    /// an engine failure.
+    async fn monitor() -> (HealthMonitor, HealthService, [watch::Sender<bool>; 2]) {
+        let (reporter, _) = tonic_health::server::health_reporter();
+        let health = HealthService::from_health_reporter(reporter.clone());
+        let (liveness_tx, liveness) = watch::channel(true);
+        let (readiness_tx, readiness) = watch::channel(true);
+        let engine = EngineHealth {
+            liveness,
+            readiness,
+        };
+        let monitor = HealthMonitor::register(reporter, engine).await;
+        (monitor, health, [liveness_tx, readiness_tx])
+    }
+
+    /// Whether `readiness` reports `SERVING`, read without yielding. By the
+    /// time the listener stops, the service is `NOT_SERVING` or removed.
+    fn serving(health: &HealthService) -> bool {
+        let request = Request::new(HealthCheckRequest {
+            service: READINESS_SERVICE.to_string(),
+        });
+        let response = health.check(request).now_or_never().expect("uncontended read");
+        response.is_ok_and(|response| response.get_ref().status == ServingStatus::Serving as i32)
+    }
+
+    #[tokio::test]
+    async fn listener_stops_after_readiness_is_withdrawn() {
+        let (monitor, health, _engine) = monitor().await;
+        let shutdown = CancellationToken::new();
+        shutdown.cancel();
+
+        let serving_at_stop = monitor
+            .serve(shutdown, std::future::pending(), |stop| async move {
+                stop.await;
+                serving(&health)
+            })
+            .await;
 
-    health_reporter.set_not_serving::<InferenceGrpcService>().await;
-    health_reporter.set_not_serving::<ControlGrpcService>().await;
-    // Both gRPC services use the same engine client, so overall server health
-    // mirrors their shared engine health.
-    health_reporter.set_service_status("", status).await;
-
-    if health_event_first {
-        shutdown.cancelled().await;
-        info!(
-            reason = "server_shutdown",
-            "server shutting down; closing gRPC health watches"
-        );
+        assert_eq!(serving_at_stop, Some(false));
     }
 
-    health_reporter.clear_service_status(inference_service).await;
-    health_reporter.clear_service_status(control_service).await;
-    health_reporter.clear_service_status("").await;
+    #[tokio::test]
+    async fn forced_abort_waits_until_readiness_is_withdrawn() {
+        /// Records readiness when the server is dropped, as an abort does.
+        struct Server {
+            health: HealthService,
+            serving_at_drop: Arc<Mutex<Option<bool>>>,
+        }
+        impl Drop for Server {
+            fn drop(&mut self) {
+                *self.serving_at_drop.lock().unwrap() = Some(serving(&self.health));
+            }
+        }
+
+        let (monitor, health, _engine) = monitor().await;
+        let serving_at_drop = Arc::new(Mutex::new(None));
+        let server = Server {
+            health,
+            serving_at_drop: serving_at_drop.clone(),
+        };
+        let shutdown = CancellationToken::new();
+        shutdown.cancel();
+
+        let output = monitor
+            .serve(shutdown, async {}, |_stop| async move {
+                let _server = server;
+                std::future::pending::<()>().await
+            })
+            .await;
+
+        assert_eq!(output, None);
+        assert_eq!(*serving_at_drop.lock().unwrap(), Some(false));
+    }
 }
diff --git a/rust/src/server/src/grpc/inference.rs b/rust/src/server/src/grpc/inference.rs
index 2ca1a650f4d3..0c418dff9f50 100644
--- a/rust/src/server/src/grpc/inference.rs
+++ b/rust/src/server/src/grpc/inference.rs
@@ -134,7 +134,7 @@ impl InferenceServiceImpl {
         let result = async {
             let media = std::mem::take(&mut proto_request.media);
             let mut text_request =
-                convert::to_text_request(proto_request, stream, self.state.served_model_names())?;
+                convert::to_text_request(proto_request, stream, &self.state.base_model_names())?;
             text_request.arrival_time = Some(arrival_time);
             text_request.data_parallel_rank = data_parallel_rank;
 
@@ -260,6 +260,7 @@ impl pb::inference_server::Inference for InferenceServiceImpl {
             ec_transfer_params: collected.ec_transfer_params.map(Box::new),
             sampling_mask: collected.sampling_mask,
             spec_decode_metrics: collected.spec_decode_metrics,
+            weight_versions: collected.weight_versions,
             timestamps: collected.timestamps,
         };
 
diff --git a/rust/src/server/src/grpc/mod.rs b/rust/src/server/src/grpc/mod.rs
index c61c1e5e1419..6a15475865ca 100644
--- a/rust/src/server/src/grpc/mod.rs
+++ b/rust/src/server/src/grpc/mod.rs
@@ -14,7 +14,7 @@ pub use vllm_proto as pb;
 
 pub(crate) use control::ControlGrpcService;
 pub use control::ControlServiceImpl;
-pub(crate) use health::monitor_health;
+pub(crate) use health::{EngineHealth, HealthMonitor};
 pub(crate) use inference::InferenceGrpcService;
 pub use inference::InferenceServiceImpl;
 pub use pb::control_server::ControlServer;
diff --git a/rust/src/server/src/grpc/tests.rs b/rust/src/server/src/grpc/tests.rs
index d1590c2129e6..521b3bafa8c8 100644
--- a/rust/src/server/src/grpc/tests.rs
+++ b/rust/src/server/src/grpc/tests.rs
@@ -59,6 +59,7 @@ use zeromq::{DealerSocket, PushSocket, ZmqMessage};
 
 use super::control::kv_event_source;
 use super::convert::json_to_proto_struct;
+use super::health::{EngineHealth, HealthMonitor};
 use super::pb::control_client::ControlClient;
 use super::pb::inference_client::InferenceClient;
 use super::{ControlServer, ControlServiceImpl, InferenceServer, InferenceServiceImpl, pb};
@@ -353,7 +354,7 @@ async fn setup_grpc_service(
 ) -> (
     InferenceServer<InferenceServiceImpl>,
     ControlServer<ControlServiceImpl>,
-    tokio::sync::watch::Receiver<bool>,
+    EngineHealth,
     MockEngineTask,
 ) {
     setup_grpc_service_with_backend(
@@ -373,7 +374,7 @@ async fn setup_grpc_service_with_backend<F>(
 ) -> (
     InferenceServer<InferenceServiceImpl>,
     ControlServer<ControlServiceImpl>,
-    tokio::sync::watch::Receiver<bool>,
+    EngineHealth,
     MockEngineTask,
 )
 where
@@ -410,7 +411,7 @@ async fn setup_grpc_service_with_engine_script<F>(
 ) -> (
     InferenceServer<InferenceServiceImpl>,
     ControlServer<ControlServiceImpl>,
-    tokio::sync::watch::Receiver<bool>,
+    EngineHealth,
     MockEngineTask,
 )
 where
@@ -437,7 +438,7 @@ where
     )
     .await
     .expect("connect client");
-    let engine_health = client.subscribe_health();
+    let engine_health = EngineHealth::subscribe(&client);
 
     let chat = ChatLlm::from_shared_backend(Llm::new(client), backend);
     let state = Arc::new(AppState::new(vec!["test-model".to_string()], chat));
@@ -475,31 +476,33 @@ async fn grpc_test_server(
 async fn start_grpc_test_server(
     inference_service: InferenceServer<InferenceServiceImpl>,
     control_service: ControlServer<ControlServiceImpl>,
-    engine_health: tokio::sync::watch::Receiver<bool>,
+    engine_health: EngineHealth,
     shutdown: tokio_util::sync::CancellationToken,
 ) -> (Channel, tokio::task::JoinHandle<()>) {
     let (health_reporter, health_service) = health_reporter();
-    health_reporter.set_serving::<InferenceServer<InferenceServiceImpl>>().await;
-    health_reporter.set_serving::<ControlServer<ControlServiceImpl>>().await;
+    let health_monitor = HealthMonitor::register(health_reporter, engine_health).await;
 
     let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.expect("bind grpc listener");
     let addr = listener.local_addr().expect("local addr");
 
     let server_task = tokio::spawn(async move {
         let incoming = MaybeTlsListener::plain(Listener::Tcp(listener));
-        let server = TonicServer::builder()
-            .add_service(health_service)
-            .add_service(control_service)
-            .add_service(inference_service)
-            .serve_with_incoming_shutdown(incoming, shutdown.clone().cancelled_owned());
-        let health_monitor =
-            super::monitor_health(health_reporter, engine_health, shutdown.clone());
-        let server = async move {
-            let result = server.await;
-            shutdown.cancel();
-            result
-        };
-        let (server_result, ()) = tokio::join!(server, health_monitor);
+        let server_shutdown = shutdown.clone();
+        let server_result = health_monitor
+            .serve(shutdown, std::future::pending(), |stop| {
+                let server = TonicServer::builder()
+                    .add_service(health_service)
+                    .add_service(control_service)
+                    .add_service(inference_service)
+                    .serve_with_incoming_shutdown(incoming, stop);
+                async move {
+                    let result = server.await;
+                    server_shutdown.cancel();
+                    result
+                }
+            })
+            .await
+            .expect("grpc server is never forced to abort");
         server_result.expect("grpc server");
     });
 
@@ -2392,6 +2395,48 @@ async fn control_list_loras_requires_lora_enabled_engine() {
     server_task.abort();
 }
 
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+#[serial]
+async fn control_sleep_rejects_unknown_level() {
+    let mut ready = default_ready_response();
+    ready.enable_sleep_mode = true;
+    let (inference_service, control_service, engine_health, engine_task) =
+        setup_grpc_service_with_engine_script(
+            b"engine-grpc-sleep-level".to_vec(),
+            ready,
+            Arc::new(FakeTextBackend::default()),
+            None,
+            // With sleep mode, the client asks for readiness at connect.
+            |dealer, push| {
+                boxed_test_future(reply_utility_bool(dealer, push, "is_inference_ready", true))
+            },
+        )
+        .await;
+    let (channel, server_task) = start_grpc_test_server(
+        inference_service,
+        control_service,
+        engine_health,
+        tokio_util::sync::CancellationToken::new(),
+    )
+    .await;
+    let mut control_client = ControlClient::new(channel);
+    let status = control_client
+        .sleep(pb::SleepRequest {
+            level: Some(3),
+            mode: pb::PauseMode::Unspecified.into(),
+        })
+        .await
+        .expect_err("level 3 should be rejected");
+    assert_eq!(status.code(), tonic::Code::InvalidArgument);
+    assert_eq!(
+        status.message(),
+        "invalid sleep level `3`; expected one of: 0, 1, 2"
+    );
+
+    engine_task.await.expect("mock engine task");
+    server_task.abort();
+}
+
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 #[serial]
 async fn control_forwards_weight_update_without_pause_guard() {
@@ -2589,20 +2634,29 @@ fn kv_event_source_filters_and_exposes_zmq_publisher() {
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 #[serial]
 async fn grpc_health_transitions_to_not_serving_when_engine_becomes_unhealthy() {
-    let (inference_service, control_service, _connected_engine_health, _engine_task) =
+    let (inference_service, control_service, connected_engine_health, _engine_task) =
         setup_grpc_service(b"engine-grpc-health-failure", default_stream_output_specs()).await;
-    let (engine_health_tx, engine_health) = tokio::sync::watch::channel(true);
+    let (engine_health_tx, liveness) = tokio::sync::watch::channel(true);
     let (channel, server_task) = start_grpc_test_server(
         inference_service,
         control_service,
-        engine_health,
+        EngineHealth {
+            liveness,
+            ..connected_engine_health
+        },
         tokio_util::sync::CancellationToken::new(),
     )
     .await;
     let mut health_client = HealthClient::new(channel);
 
     let mut health_streams = Vec::new();
-    for service in ["vllm.Inference", "vllm.Control", ""] {
+    for service in [
+        "vllm.Inference",
+        "vllm.Control",
+        "",
+        "liveness",
+        "readiness",
+    ] {
         let service_label = if service.is_empty() {
             "overall"
         } else {
@@ -2730,6 +2784,17 @@ async fn grpc_shutdown_closes_health_watch_and_server() {
         HealthServingStatus::Serving as i32,
         "unexpected initial health status for vllm.Inference"
     );
+    let mut liveness = health_client
+        .watch(HealthCheckRequest {
+            service: "liveness".to_string(),
+        })
+        .await
+        .expect("start health watch for liveness")
+        .into_inner();
+    assert_eq!(
+        next_health_status(&mut liveness).await,
+        HealthServingStatus::Serving
+    );
 
     control_client
         .shutdown(pb::ShutdownRequest {})
@@ -2759,9 +2824,143 @@ async fn grpc_shutdown_closes_health_watch_and_server() {
         stream_end.is_none(),
         "vllm.Inference health watch remained open"
     );
+    let liveness_end = tokio::time::timeout(Duration::from_secs(2), liveness.message())
+        .await
+        .expect("timed out waiting for liveness health watch to close")
+        .expect("failed while closing liveness health watch");
+    assert!(
+        liveness_end.is_none(),
+        "orderly shutdown must not fail liveness: {liveness_end:?}"
+    );
 
     tokio::time::timeout(Duration::from_secs(2), server_task)
         .await
         .expect("timed out waiting for gRPC server shutdown")
         .expect("gRPC server task failed");
 }
+
+async fn next_health_status(
+    stream: &mut tonic::Streaming<tonic_health::pb::HealthCheckResponse>,
+) -> HealthServingStatus {
+    let response = tokio::time::timeout(Duration::from_secs(2), stream.message())
+        .await
+        .expect("timed out waiting for a health update")
+        .expect("failed to read health update")
+        .expect("health watch ended");
+    HealthServingStatus::try_from(response.status).expect("known health status")
+}
+
+/// Serve pause, resume, and readiness asks as an engine that is ready unless
+/// paused, in whatever order the client sends them.
+async fn serve_pausable_engine(dealer: &mut DealerSocket, push: &mut PushSocket) {
+    let mut paused = false;
+    loop {
+        let frames = recv_engine_message(dealer).await;
+        assert_eq!(frames[0].as_ref(), &[0x03]);
+        let payload = decode_value(&frames[1]).expect("decode utility payload");
+        let fields = payload.as_array().expect("utility payload array");
+        let call_id = fields[1].as_u64().expect("utility call id");
+        let result = match fields[2].as_str().expect("utility method") {
+            "is_inference_ready" => Value::Boolean(!paused),
+            "pause_generation" => {
+                paused = true;
+                Value::Nil
+            }
+            "resume_scheduler" => {
+                paused = false;
+                Value::Nil
+            }
+            method => panic!("unexpected utility call {method}"),
+        };
+        send_outputs(
+            push,
+            UtilityCallOutput {
+                output: UtilityOutput {
+                    call_id: call_id.into(),
+                    failure_message: None,
+                    result: Some(UtilityResultEnvelope::without_type_info(result)),
+                },
+                ..Default::default()
+            }
+            .into(),
+        )
+        .await;
+    }
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+#[serial]
+async fn grpc_pause_withdraws_readiness_but_not_liveness() {
+    let (inference_service, control_service, engine_health, _engine_task) =
+        setup_grpc_service_with_engine_script(
+            b"engine-grpc-pause-readiness",
+            default_ready_response(),
+            Arc::new(FakeTextBackend::default()),
+            None,
+            |dealer, push| boxed_test_future(serve_pausable_engine(dealer, push)),
+        )
+        .await;
+    let (channel, server_task) = start_grpc_test_server(
+        inference_service,
+        control_service,
+        engine_health,
+        tokio_util::sync::CancellationToken::new(),
+    )
+    .await;
+    let mut control_client = ControlClient::new(channel.clone());
+    let mut health_client = HealthClient::new(channel);
+    let mut readiness_streams = Vec::new();
+    for service in ["readiness", "vllm.Inference"] {
+        let mut stream = health_client
+            .watch(HealthCheckRequest {
+                service: service.to_string(),
+            })
+            .await
+            .expect("start health watch")
+            .into_inner();
+        assert_eq!(
+            next_health_status(&mut stream).await,
+            HealthServingStatus::Serving
+        );
+        readiness_streams.push(stream);
+    }
+
+    control_client
+        .pause_generation(pb::PauseGenerationRequest::default())
+        .await
+        .expect("pause generation");
+    for stream in &mut readiness_streams {
+        assert_eq!(
+            next_health_status(stream).await,
+            HealthServingStatus::NotServing
+        );
+    }
+    for service in ["liveness", "vllm.Control", ""] {
+        let status = health_client
+            .check(HealthCheckRequest {
+                service: service.to_string(),
+            })
+            .await
+            .expect("check health")
+            .into_inner()
+            .status;
+        assert_eq!(
+            status,
+            HealthServingStatus::Serving as i32,
+            "a pause must not fail {service:?}"
+        );
+    }
+
+    control_client
+        .resume_generation(pb::ResumeGenerationRequest {})
+        .await
+        .expect("resume generation");
+    for stream in &mut readiness_streams {
+        assert_eq!(
+            next_health_status(stream).await,
+            HealthServingStatus::Serving
+        );
+    }
+
+    server_task.abort();
+}
diff --git a/rust/src/server/src/lib.rs b/rust/src/server/src/lib.rs
index 4298ebfaead7..ed556ecc5c24 100644
--- a/rust/src/server/src/lib.rs
+++ b/rust/src/server/src/lib.rs
@@ -9,6 +9,7 @@ mod grpc;
 mod listener;
 mod lora;
 mod middleware;
+mod model_names;
 mod render;
 mod routes;
 mod runtime;
@@ -26,11 +27,12 @@ use std::time::Duration;
 use anyhow::{Context as _, Result};
 use axum::Router;
 use axum::body::Body;
+use axum::extract::ConnectInfo;
 use axum::http::Request;
 pub use config::{
-    ApiServerOptions, Config, CoordinatorMode, CorsConfig, DEFAULT_KEEP_ALIVE_TIMEOUT,
-    DEFAULT_LOG_STATS_INTERVAL, DEFAULT_REQUEST_BODY_LIMIT_BYTES, HttpListenerMode, LoraModulePath,
-    OutputLogging, TlsConfig,
+    AccessLogConfig, ApiServerOptions, Config, CoordinatorMode, CorsConfig,
+    DEFAULT_KEEP_ALIVE_TIMEOUT, DEFAULT_LOG_STATS_INTERVAL, DEFAULT_REQUEST_BODY_LIMIT_BYTES,
+    HttpListenerMode, LoraModulePath, OutputLogConfig, OutputLogging, TlsConfig,
 };
 use hyper::body::Incoming;
 use hyper::server::conn::http1;
@@ -55,7 +57,7 @@ pub use vllm_metrics::CustomHistogramBuckets;
 use vllm_metrics::METRICS;
 use vllm_text::TextLlm;
 
-use crate::listener::{Listener, MaybeTlsListener};
+use crate::listener::{Listener, ListenerAddr, MaybeTlsListener};
 use crate::routes::build_router;
 use crate::server_info::ServerInfoSnapshot;
 use crate::state::AppState;
@@ -161,6 +163,7 @@ async fn build_state(config: &Config) -> Result<Arc<AppState>> {
             .with_server_info(server_info)
             .with_api_keys(config.api_keys.clone())
             .with_cors(config.cors.clone())
+            .with_access_log(config.access_log.clone())
             .with_profiler(config.profiler.clone()),
     );
 
@@ -228,7 +231,10 @@ where
         _ = shutdown.cancelled() => return Ok(()),
     };
     let model = state.primary_model_name().to_owned();
-    let app = extend_router(build_router(state.clone()));
+    let mut app = extend_router(build_router(state.clone()));
+    if let Some(root_path) = config.root_path.clone() {
+        app = routes::with_root_path(app, root_path);
+    }
 
     info!(model, "starting vLLM server");
 
@@ -253,9 +259,8 @@ where
             .transpose()
             .context("invalid gRPC TLS configuration")?;
         let (health_reporter, health_service) = health_reporter();
-        let engine_health = state.engine_core_client().subscribe_health();
-        health_reporter.set_serving::<grpc::InferenceGrpcService>().await;
-        health_reporter.set_serving::<grpc::ControlGrpcService>().await;
+        let engine_health = grpc::EngineHealth::subscribe(state.engine_core_client());
+        let health_monitor = grpc::HealthMonitor::register(health_reporter, engine_health).await;
         let control_service = grpc::ControlGrpcService::new(
             grpc::ControlServiceImpl::new(state.clone())
                 .with_engine_shutdown(config.manages_engine.then(|| shutdown.clone())),
@@ -271,14 +276,7 @@ where
             .add_service(health_service)
             .add_service(control_service)
             .add_service(inference_service);
-        Some((
-            addr,
-            grpc_listener,
-            svc,
-            grpc_tls,
-            health_reporter,
-            engine_health,
-        ))
+        Some((addr, grpc_listener, svc, grpc_tls, health_monitor))
     } else {
         None
     };
@@ -367,9 +365,7 @@ where
         let server_shutdown = server_shutdown.clone();
         let force_shutdown = force_shutdown.clone();
         async move {
-            let Some((addr, grpc_listener, svc, grpc_tls, health_reporter, engine_health)) =
-                grpc_setup
-            else {
+            let Some((addr, grpc_listener, svc, grpc_tls, health_monitor)) = grpc_setup else {
                 // No gRPC configured: just wait for shutdown so we do not race the
                 // join! by resolving early and tripping the cancellation token.
                 shutdown.cancelled().await;
@@ -380,29 +376,21 @@ where
                 Some(context) => MaybeTlsListener::tls(grpc_listener, context),
                 None => MaybeTlsListener::plain(grpc_listener),
             };
-            let server =
-                svc.serve_with_incoming_shutdown(incoming, shutdown.clone().cancelled_owned());
-            let health_monitor = grpc::monitor_health(health_reporter, engine_health, shutdown);
+            let server = health_monitor.serve(shutdown, force_shutdown.cancelled(), |stop| {
+                let server = svc.serve_with_incoming_shutdown(incoming, stop);
+                async move {
+                    let result = server.await.context("gRPC server failed");
+                    server_shutdown.cancel();
+                    result
+                }
+            });
 
             info!(%addr, tls, model, "gRPC server is ready to accept requests");
 
-            let server = async move {
-                let result = tokio::select! {
-                    result = server => {
-                        result.context("gRPC server failed")
-                    }
-                    _ = force_shutdown.cancelled() => {
-                        warn!("gRPC graceful shutdown deadline elapsed; aborting server");
-                        Ok(())
-                    }
-                };
-
-                server_shutdown.cancel();
-                result
-            };
-
-            let (result, ()) = tokio::join!(server, health_monitor);
-            result
+            server.await.unwrap_or_else(|| {
+                warn!("gRPC graceful shutdown deadline elapsed; aborting server");
+                Ok(())
+            })
         }
     };
 
@@ -434,19 +422,29 @@ async fn serve_connections<L>(
     timeouts: ConnectionTimeouts,
 ) -> Result<()>
 where
-    L: axum::serve::Listener,
+    L: axum::serve::Listener<Addr = ListenerAddr>,
 {
     let graceful = GracefulShutdown::new();
     let mut shutdown = std::pin::pin!(shutdown);
     loop {
-        let (io, _addr) = tokio::select! {
+        let (io, addr) = tokio::select! {
             conn = listener.accept() => conn,
             () = &mut shutdown => break,
         };
 
-        let service = TowerToHyperService::new(
-            app.clone().map_request(|req: Request<Incoming>| req.map(Body::new)),
-        );
+        // Expose the TCP peer to handlers, as `ConnectInfo` does for `axum::serve`.
+        let peer = match addr {
+            ListenerAddr::Tcp(addr) => Some(ConnectInfo(addr)),
+            ListenerAddr::Unix(_) => None,
+        };
+        let service =
+            TowerToHyperService::new(app.clone().map_request(move |req: Request<Incoming>| {
+                let mut req = req.map(Body::new);
+                if let Some(peer) = peer {
+                    req.extensions_mut().insert(peer);
+                }
+                req
+            }));
         let mut builder = http1::Builder::new();
         builder.timer(TokioTimer::new()).header_read_timeout(timeouts.header_read);
         if !timeouts.keep_alive_enabled {
diff --git a/rust/src/server/src/lora.rs b/rust/src/server/src/lora.rs
index 9c3c491cc135..459579784738 100644
--- a/rust/src/server/src/lora.rs
+++ b/rust/src/server/src/lora.rs
@@ -1,6 +1,7 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+use std::io;
 use std::path::{Component, Path, PathBuf};
 use std::sync::atomic::{AtomicU64, Ordering};
 
@@ -12,11 +13,14 @@ use vllm_engine_core_client::EngineCoreClient;
 use vllm_engine_core_client::protocol::lora::LoraRequest;
 
 use crate::config::LoraModulePath;
+use crate::model_names::ServedModelNames;
 
 const RUNTIME_LORA_ALLOWED_PATH_PREFIXES_ENV: &str = "VLLM_RUNTIME_LORA_ALLOWED_PATH_PREFIXES";
 
 #[derive(Debug, Error, Macro)]
 pub(crate) enum LoraPathAccessError {
+    #[error("No adapter found for {lora_path}")]
+    NotFound { lora_path: String },
     #[error("{message}")]
     InvalidPath { message: String },
     #[error("{message}")]
@@ -60,9 +64,12 @@ fn validate_lora_path_access(
         ));
     }
 
-    let canonical_path = path
-        .canonicalize()
-        .map_err(|_| invalid_path!("Local LoRA adapter path must exist and be accessible."))?;
+    let canonical_path = path.canonicalize().map_err(|error| match error.kind() {
+        io::ErrorKind::NotFound => LoraPathAccessError::NotFound {
+            lora_path: lora_path.to_string(),
+        },
+        _ => invalid_path!("Local LoRA adapter path must exist and be accessible."),
+    })?;
     let canonical_prefixes = allowed_prefixes
         .iter()
         .map(|prefix| {
@@ -115,10 +122,21 @@ impl LoadedLora {
 /// the model name resolves to a dynamic adapter.
 #[derive(Debug, Clone)]
 pub(crate) struct LoraModelResolution {
-    pub model_names: Vec<String>,
+    pub model_names: ServedModelNames,
     pub lora_request: Option<LoraRequest>,
 }
 
+impl LoraModelResolution {
+    /// The `model` reported in responses: the LoRA adapter name when the
+    /// request resolved to one, otherwise the first served model name.
+    pub fn response_model(&self) -> String {
+        self.lora_request
+            .as_ref()
+            .map(|request| request.lora_name.clone())
+            .unwrap_or_else(|| self.model_names.primary().to_string())
+    }
+}
+
 /// Runtime registry for dynamically loaded LoRA adapters.
 pub(crate) struct LoraManager {
     /// Dynamically loaded LoRA adapters keyed by public model name, in load order.
@@ -197,12 +215,11 @@ impl LoraManager {
     /// snapshot.
     pub async fn resolve_model(
         &self,
-        base_model_names: &[String],
+        base_model_names: ServedModelNames,
         model_name: Option<&str>,
     ) -> LoraModelResolution {
         let adapters = self.adapters.read().await;
-        let mut model_names = base_model_names.to_vec();
-        model_names.extend(adapters.keys().cloned());
+        let model_names = base_model_names.with_names(adapters.keys().cloned());
         let lora_request =
             model_name.and_then(|name| adapters.get(name)).map(LoadedLora::to_request);
 
@@ -347,7 +364,8 @@ mod tests {
     use std::path::PathBuf;
     use std::time::{SystemTime, UNIX_EPOCH};
 
-    use super::validate_lora_path_access;
+    use super::{LoadedLora, LoraManager, LoraPathAccessError, validate_lora_path_access};
+    use crate::model_names::ServedModelNames;
 
     fn temp_lora_dir(test_name: &str) -> PathBuf {
         let suffix = SystemTime::now()
@@ -362,6 +380,33 @@ mod tests {
         path
     }
 
+    #[tokio::test]
+    async fn resolve_model_selects_adapter_when_any_model_name_is_accepted() {
+        let manager = LoraManager::new();
+        manager.adapters.write().await.insert(
+            "adapter-a".to_string(),
+            LoadedLora {
+                lora_name: "adapter-a".to_string(),
+                lora_int_id: 1,
+                lora_path: "org/adapter-a".to_string(),
+                base_model_name: None,
+                is_3d_lora_weight: false,
+            },
+        );
+        let base = ServedModelNames::new(vec!["base".to_string()]).with_accept_any(true);
+
+        let resolution = manager.resolve_model(base.clone(), Some("adapter-a")).await;
+        assert_eq!(
+            resolution.lora_request.expect("adapter selected").lora_name,
+            "adapter-a"
+        );
+
+        let resolution = manager.resolve_model(base, Some("unknown")).await;
+        assert!(resolution.lora_request.is_none());
+        assert!(resolution.model_names.accepts(Some("unknown")));
+        assert_eq!(resolution.model_names.primary(), "base");
+    }
+
     #[test]
     fn lora_path_allows_hf_repo_ids_without_prefixes() {
         assert_eq!(
@@ -417,6 +462,25 @@ mod tests {
         fs::remove_dir_all(root).ok();
     }
 
+    /// Python answers a missing adapter with 404 `NotFoundError`, so the
+    /// guard reports a missing path apart from an inaccessible one.
+    #[test]
+    fn lora_path_reports_missing_paths_as_not_found() {
+        let root = temp_lora_dir("missing-path");
+        let missing = root.join("missing-adapter");
+
+        let prefixes = [root.clone()];
+        let error =
+            validate_lora_path_access(missing.to_str().expect("utf-8 temp path"), Some(&prefixes))
+                .expect_err("missing adapter path should be rejected");
+        assert!(
+            matches!(&error, LoraPathAccessError::NotFound { lora_path } if *lora_path == missing.to_str().unwrap()),
+            "{error:?}"
+        );
+
+        fs::remove_dir_all(root).ok();
+    }
+
     #[test]
     fn lora_path_rejects_parent_escape_from_configured_prefixes() {
         let root = temp_lora_dir("parent-escape");
diff --git a/rust/src/server/src/middleware/access_log.rs b/rust/src/server/src/middleware/access_log.rs
new file mode 100644
index 000000000000..8a264873b0db
--- /dev/null
+++ b/rust/src/server/src/middleware/access_log.rs
@@ -0,0 +1,57 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+//! Per-response HTTP access log.
+
+use std::net::SocketAddr;
+use std::sync::Arc;
+use std::time::Instant;
+
+use axum::extract::{ConnectInfo, Request, State};
+use axum::middleware::Next;
+use axum::response::Response;
+use tracing::info;
+
+use crate::state::AppState;
+
+/// Log one INFO event per response, unless the request path is excluded: the
+/// client address, method, request target with its query string, HTTP
+/// version, status, and the latency until the response head is ready.
+///
+/// A streaming response is logged when its stream starts, so its latency
+/// does not cover the body.
+pub async fn log_access(
+    State(state): State<Arc<AppState>>,
+    request: Request,
+    next: Next,
+) -> Response {
+    let path = request.uri().path();
+    if (state.access_log.excluded_paths.iter()).any(|excluded| excluded == path) {
+        return next.run(request).await;
+    }
+
+    // Absent for Unix-domain sockets.
+    let client = request
+        .extensions()
+        .get::<ConnectInfo<SocketAddr>>()
+        .map(|ConnectInfo(addr)| *addr);
+    let method = request.method().clone();
+    let uri = request
+        .uri()
+        .path_and_query()
+        .map_or_else(|| path.to_owned(), |target| target.as_str().to_owned());
+    let version = request.version();
+
+    let start = Instant::now();
+    let response = next.run(request).await;
+    info!(
+        client = client.map(display),
+        %method,
+        uri,
+        ?version,
+        status = response.status().as_u16(),
+        latency = ?start.elapsed(),
+        "HTTP response"
+    );
+    response
+}
diff --git a/rust/src/server/src/middleware/mod.rs b/rust/src/server/src/middleware/mod.rs
index 444099175d32..9e869b1ac52c 100644
--- a/rust/src/server/src/middleware/mod.rs
+++ b/rust/src/server/src/middleware/mod.rs
@@ -1,6 +1,7 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+mod access_log;
 mod auth;
 mod cors;
 mod load;
@@ -8,6 +9,7 @@ mod metrics;
 mod offload;
 mod request_id;
 
+pub use access_log::log_access;
 pub use auth::authenticate_api_key;
 pub use cors::{cors_layer, strip_cors_on_no_origin};
 pub use load::track_server_load;
diff --git a/rust/src/server/src/model_names.rs b/rust/src/server/src/model_names.rs
new file mode 100644
index 000000000000..0c1fbd76aea2
--- /dev/null
+++ b/rust/src/server/src/model_names.rs
@@ -0,0 +1,96 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+use crate::error::ApiError;
+
+/// The public model names a request's `model` field may select.
+#[derive(Debug, Clone)]
+pub(crate) struct ServedModelNames {
+    /// Served base model names, then any loaded LoRA adapters. The first
+    /// entry is the primary name.
+    names: Vec<String>,
+    /// Accept any other name too, served by the base model, like Python with
+    /// `VLLM_SKIP_MODEL_NAME_VALIDATION` (for proxies and gateways that
+    /// rewrite or ignore model names).
+    accept_any: bool,
+}
+
+impl ServedModelNames {
+    /// `names` must be non-empty; the first entry is the primary name.
+    pub fn new(names: Vec<String>) -> Self {
+        assert!(!names.is_empty(), "served model names must not be empty");
+        Self {
+            names,
+            accept_any: false,
+        }
+    }
+
+    /// Also accept any other name, served by the base model
+    /// (`VLLM_SKIP_MODEL_NAME_VALIDATION`).
+    pub fn with_accept_any(mut self, accept_any: bool) -> Self {
+        self.accept_any = accept_any;
+        self
+    }
+
+    /// Extend the names with more selectable names, such as loaded LoRA
+    /// adapters.
+    pub fn with_names(mut self, names: impl IntoIterator<Item = String>) -> Self {
+        self.names.extend(names);
+        self
+    }
+
+    /// The name reported as `model` in responses that select no LoRA adapter.
+    pub fn primary(&self) -> &str {
+        &self.names[0]
+    }
+
+    /// All selectable names, primary first.
+    pub fn names(&self) -> &[String] {
+        &self.names
+    }
+
+    /// Whether a request may name `model`. Like Python, an absent or empty
+    /// name selects the primary model.
+    pub fn accepts(&self, model: Option<&str>) -> bool {
+        match model {
+            None | Some("") => true,
+            Some(model) => self.accept_any || self.names.iter().any(|name| name == model),
+        }
+    }
+
+    /// Reject a model name this server does not serve with a 404.
+    pub fn check(&self, model: Option<&str>) -> Result<(), ApiError> {
+        match model {
+            Some(model) if !self.accepts(Some(model)) => {
+                Err(ApiError::model_not_found(model.to_string()))
+            }
+            _ => Ok(()),
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::ServedModelNames;
+
+    fn served(accept_any: bool) -> ServedModelNames {
+        ServedModelNames::new(vec!["base".to_string(), "alias".to_string()])
+            .with_accept_any(accept_any)
+    }
+
+    #[test]
+    fn strict_names_accept_only_served_names() {
+        let served = served(false);
+        for model in [None, Some(""), Some("base"), Some("alias")] {
+            assert!(served.check(model).is_ok(), "{model:?}");
+        }
+        assert!(served.check(Some("other")).is_err());
+    }
+
+    #[test]
+    fn skip_validation_accepts_any_name() {
+        let served = served(true);
+        assert!(served.check(Some("other")).is_ok());
+        assert_eq!(served.primary(), "base");
+    }
+}
diff --git a/rust/src/server/src/render.rs b/rust/src/server/src/render.rs
index b596c263109c..b55c8ccdbfdf 100644
--- a/rust/src/server/src/render.rs
+++ b/rust/src/server/src/render.rs
@@ -15,6 +15,7 @@ use vllm_chat::{
 use vllm_text::TextRequestProcessor;
 use vllm_text::backend::hf::{HfOverrides, ModelSource};
 
+use crate::model_names::ServedModelNames;
 use crate::{
     HttpListenerMode, ReasoningEffort, TlsConfig,
     listener::{Listener, MaybeTlsListener},
@@ -32,6 +33,8 @@ pub struct RenderConfig {
     pub hf_config: Option<ModelSource>,
     pub hf_overrides: HfOverrides,
     pub served_model_name: Vec<String>,
+    /// Accept any requested model name (`VLLM_SKIP_MODEL_NAME_VALIDATION`).
+    pub skip_model_name_validation: bool,
     pub host: String,
     pub port: u16,
     pub tool_call_parser: ParserSelection,
@@ -66,7 +69,7 @@ impl RenderConfig {
 
 pub(crate) struct RenderState {
     pub(crate) model: String,
-    pub(crate) served_model_names: Vec<String>,
+    pub(crate) served_model_names: ServedModelNames,
     /// Unlike `text.max_model_len()` this stays `None` when unset,
     /// so model cards advertise `null` instead of the `u32::MAX`.
     pub(crate) max_model_len: Option<u32>,
@@ -94,8 +97,11 @@ async fn build_state(config: &RenderConfig) -> Result<Arc<RenderState>> {
     )
     .await
     .context("failed to load renderer/tokenizer backends")?;
-    let served_model_names =
-        crate::effective_served_model_names(&config.model, &config.served_model_name);
+    let served_model_names = ServedModelNames::new(crate::effective_served_model_names(
+        &config.model,
+        &config.served_model_name,
+    ))
+    .with_accept_any(config.skip_model_name_validation);
     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);
@@ -182,6 +188,7 @@ mod tests {
                 hf_config: None,
                 hf_overrides: Default::default(),
                 served_model_name: Vec::new(),
+                skip_model_name_validation: false,
                 host: "127.0.0.1".to_string(),
                 port: 8000,
                 tool_call_parser: ParserSelection::Auto,
diff --git a/rust/src/server/src/routes.rs b/rust/src/server/src/routes.rs
index ff4a70c5f848..182745d54ac8 100644
--- a/rust/src/server/src/routes.rs
+++ b/rust/src/server/src/routes.rs
@@ -26,12 +26,12 @@ mod world_size;
 
 use std::sync::Arc;
 
-use axum::extract::DefaultBodyLimit;
-use axum::extract::MatchedPath;
-use axum::http::{Method, StatusCode};
+use axum::extract::{DefaultBodyLimit, MatchedPath, Request};
+use axum::http::{Method, StatusCode, Uri};
 use axum::middleware::{from_fn, from_fn_with_state};
 use axum::routing::{get, post};
 use axum::{Json, Router};
+use tower::ServiceExt as _;
 use tower_http::trace::TraceLayer;
 use tracing::{info, warn};
 
@@ -267,6 +267,7 @@ fn build_router_with_options(
 
     let enable_request_id_headers = state.api_server_options.enable_request_id_headers;
     let enable_api_key_auth = state.has_api_keys();
+    let enable_access_log = state.access_log.enabled;
     let mut router = router
         .with_state(state.clone())
         .layer(match state.api_server_options.request_body_limit_bytes {
@@ -293,6 +294,12 @@ fn build_router_with_options(
     // requests are visible, while metrics/load only see authenticated traffic.
     router = router.layer(TraceLayer::new_for_http());
 
+    // Log every response, including rejected and unrouted ones, like Python's
+    // uvicorn access log.
+    if enable_access_log {
+        router = router.layer(from_fn_with_state(state.clone(), middleware::log_access));
+    }
+
     if enable_request_id_headers {
         router = router.layer(from_fn(middleware::set_request_id_header));
     }
@@ -300,6 +307,37 @@ fn build_router_with_options(
     router
 }
 
+/// Also serve every route of `router` under `root_path`, like FastAPI's
+/// `root_path`: with `/llm`, both `/llm/v1/models` and `/v1/models` reach
+/// `/v1/models`.
+///
+/// The prefix is stripped before routing, so every layer sees the bare path.
+/// Like Starlette's `get_route_path`, it only matches up to a path segment
+/// boundary: `/llmx/v1/models` keeps its path.
+pub(crate) fn with_root_path(router: Router, root_path: String) -> Router {
+    Router::new().fallback_service(
+        router.map_request(move |request: Request| strip_root_path(request, &root_path)),
+    )
+}
+
+fn strip_root_path(mut request: Request, root_path: &str) -> Request {
+    let uri = request.uri();
+    let path = match uri.path().strip_prefix(root_path) {
+        Some("") => "/",
+        Some(path) if path.starts_with('/') => path,
+        _ => return request,
+    };
+    let path_and_query = match uri.query() {
+        Some(query) => format!("{path}?{query}"),
+        None => path.to_owned(),
+    };
+    let mut parts = uri.clone().into_parts();
+    parts.path_and_query =
+        Some(path_and_query.parse().expect("a suffix of a valid path is a valid path"));
+    *request.uri_mut() = Uri::from_parts(parts).expect("only the path changed");
+    request
+}
+
 #[cfg(test)]
 mod tests;
 
diff --git a/rust/src/server/src/routes/inference/generate.rs b/rust/src/server/src/routes/inference/generate.rs
index 2591a7dd2a36..2e9c53b40422 100644
--- a/rust/src/server/src/routes/inference/generate.rs
+++ b/rust/src/server/src/routes/inference/generate.rs
@@ -31,9 +31,9 @@ use vllm_llm::{
 
 use self::convert::{ResponseOptions, prepare_generate_request};
 use self::types::{
-    GenerateLogProb, GenerateLogProbs, GenerateLogProbsContent, GenerateLogprob, GenerateResponse,
-    GenerateResponseChoice, GenerateResponseStreamChoice, GenerateStreamResponse,
-    MultiModalPlaceholders, PerRequestMetrics, PlaceholderRangeInfo,
+    GenerateLogProb, GenerateLogProbs, GenerateLogProbsContent, GenerateLogprob,
+    GenerateOutputMode, GenerateResponse, GenerateResponseChoice, GenerateResponseStreamChoice,
+    GenerateStreamResponse, MultiModalPlaceholders, PerRequestMetrics, PlaceholderRangeInfo,
 };
 pub(crate) use self::types::{GenerateRequest, GenerateSamplingParams};
 pub(crate) use self::validate::validate_request_compat;
@@ -45,9 +45,9 @@ use crate::routes::openai::utils::metrics::{
 };
 use crate::routes::openai::utils::types::Usage;
 use crate::routes::openai::utils::validated_json::ValidatedJson;
-use crate::routes::output_log::{LoggedOutput, log_output};
+use crate::routes::output_log::LoggedOutput;
 use crate::state::AppState;
-use crate::utils::{resolve_request_context, sse_response};
+use crate::utils::{resolve_request_context, sse_response, unix_timestamp};
 
 /// Validate one token-in/token-out request and proxy it into the shared
 /// `vllm-text` stack.
@@ -97,6 +97,7 @@ pub async fn generate(
 
     let api_server_options = state.api_server_options;
     let stream = prepared.stream;
+    let created = unix_timestamp();
     let raw_stream = match state
         .chat
         .text()
@@ -140,6 +141,8 @@ pub async fn generate(
         collect_generate(
             collected,
             prepared.request_id,
+            prepared.response_model,
+            created,
             api_server_options,
             prepared.options,
             mm_placeholders,
@@ -232,9 +235,12 @@ async fn generate_chunk_stream(
 
                 let stop_reason =
                     finish_reason.as_ref().and_then(|reason| reason.as_stop_reason().cloned());
+                // Like Python, the spans cover the whole output on the final chunk only.
+                let weight_versions = output.weight_versions.filter(|_| finish_reason.is_some());
                 let prompt_token_ids = prompt_token_ids.take();
                 y.yield_ok(GenerateStreamResponse {
                     request_id: request_id.clone(),
+                    output_mode: GenerateOutputMode::Tokens,
                     choices: vec![GenerateResponseStreamChoice {
                         index: 0,
                         logprobs,
@@ -247,10 +253,11 @@ async fn generate_chunk_stream(
                             .map(npy_base64)
                             .transpose()?,
                         sampling_mask: output.sampling_mask.map(|mask| mask.rows),
-                        weight_versions: output.weight_versions,
+                        weight_versions,
                     }],
-                    usage: include_continuous_usage
-                        .then(|| Usage::from_token_usage(usage, enable_prompt_tokens_details)),
+                    usage: include_continuous_usage.then(|| {
+                        Usage::from_token_usage(usage, enable_prompt_tokens_details, None)
+                    }),
                     mm_placeholders: prompt_token_ids.as_ref().and_then(|_| mm_placeholders.take()),
                     prompt_token_ids,
                     metrics: None,
@@ -270,8 +277,13 @@ async fn generate_chunk_stream(
     if include_usage {
         y.yield_ok(GenerateStreamResponse {
             request_id,
+            output_mode: GenerateOutputMode::Tokens,
             choices: Vec::new(),
-            usage: Some(Usage::from_token_usage(usage, enable_prompt_tokens_details)),
+            usage: Some(Usage::from_token_usage(
+                usage,
+                enable_prompt_tokens_details,
+                None,
+            )),
             prompt_token_ids: None,
             mm_placeholders: None,
             metrics: spec_decode_metrics.map(|speculative_decoding| PerRequestMetrics {
@@ -289,13 +301,16 @@ async fn generate_chunk_stream(
 fn collect_generate(
     collected: CollectedGenerateOutput,
     request_id: String,
+    response_model: String,
+    created: u64,
     ApiServerOptions {
         enable_log_requests,
+        enable_prompt_tokens_details,
         log_outputs,
         ..
     }: ApiServerOptions,
     ResponseOptions {
-        // Ignored: non-streaming raw generate responses do not include usage.
+        // Ignored: like Python, non-streaming responses always include usage.
         include_usage: _,
         // Ignored: continuous usage is a streaming-only option.
         include_continuous_usage: _,
@@ -352,7 +367,7 @@ fn collect_generate(
     }
     // Like Python, raw generate logs only the output token IDs.
     if log_outputs.logs_complete() && !collected.token_ids.is_empty() {
-        log_output(
+        log_outputs.log(
             LoggedOutput::Response,
             "",
             Some(&collected.token_ids),
@@ -362,6 +377,10 @@ fn collect_generate(
 
     Ok(GenerateResponse {
         request_id,
+        model: response_model,
+        created,
+        usage: Usage::from_token_usage(collected.usage, enable_prompt_tokens_details, None),
+        output_mode: GenerateOutputMode::Tokens,
         choices: vec![GenerateResponseChoice {
             index: 0,
             logprobs,
@@ -585,7 +604,7 @@ mod tests {
     use vllm_engine_core_client::protocol::multimodal::{MmModality, PlaceholderRange};
     use vllm_engine_core_client::protocol::output::{RequestSpecDecodeMetrics, StopReason};
     use vllm_engine_core_client::protocol::sampling_mask::SamplingMask;
-    use vllm_llm::{GeneratePromptInfo, PromptCacheUsage};
+    use vllm_llm::{GeneratePromptInfo, PromptCacheUsage, WeightVersionSpan};
 
     use super::*;
 
@@ -833,6 +852,25 @@ mod tests {
         assert_eq!(chunks[0].choices[0].sampling_mask, Some(vec![vec![33, 44]]));
     }
 
+    #[tokio::test]
+    async fn generate_chunk_stream_returns_weight_versions_on_final_chunk_only() {
+        let span = |end| WeightVersionSpan {
+            version: "v1".into(),
+            start: 0,
+            end,
+        };
+        let mut first = stream_output(Some(&[11]), vec![33], None);
+        first.weight_versions = Some(vec![span(1)]);
+        let mut last = stream_output(None, vec![44], Some(FinishReason::Length));
+        last.weight_versions = Some(vec![span(2)]);
+
+        let chunks = collect_chunks(vec![first, last], false, None).await;
+
+        let weight_versions: Vec<_> =
+            chunks.iter().map(|chunk| chunk.choices[0].weight_versions.clone()).collect();
+        assert_eq!(weight_versions, [None, Some(vec![span(2)])]);
+    }
+
     #[tokio::test]
     async fn generate_chunk_stream_omits_prompt_metadata_by_default() {
         let chunks = collect_chunks(
@@ -979,6 +1017,8 @@ mod tests {
         let response = collect_generate(
             output,
             "raw-1".to_string(),
+            "model".to_string(),
+            0,
             ApiServerOptions::default(),
             ResponseOptions::default(),
             Some(MultiModalPlaceholders::from([(
@@ -1085,6 +1125,8 @@ mod tests {
         let response = collect_generate(
             output,
             "raw-metrics".to_string(),
+            "model".to_string(),
+            0,
             ApiServerOptions::default(),
             ResponseOptions::default(),
             None,
@@ -1179,6 +1221,8 @@ mod tests {
         let response = collect_generate(
             output,
             "raw-1".to_string(),
+            "model".to_string(),
+            0,
             ApiServerOptions::default(),
             ResponseOptions {
                 return_token_ids: true,
@@ -1217,6 +1261,8 @@ mod tests {
         let response = collect_generate(
             output_without_payload(vec![9707]),
             "raw-1".to_string(),
+            "model".to_string(),
+            0,
             ApiServerOptions::default(),
             ResponseOptions {
                 include_prompt_logprobs: true,
@@ -1232,6 +1278,8 @@ mod tests {
         collect_generate(
             output_without_payload(vec![9707, 11]),
             "raw-2".to_string(),
+            "model".to_string(),
+            0,
             ApiServerOptions::default(),
             ResponseOptions {
                 include_prompt_logprobs: true,
@@ -1309,6 +1357,8 @@ mod tests {
         let response = collect_generate(
             output,
             "raw-stop".to_string(),
+            "model".to_string(),
+            0,
             ApiServerOptions::default(),
             ResponseOptions::default(),
             None,
@@ -1321,4 +1371,54 @@ mod tests {
             Some(StopReason::TokenId(7))
         );
     }
+
+    /// Like Python, the non-streaming envelope reports cached prompt tokens
+    /// under `--enable-prompt-tokens-details`, which is what disaggregated
+    /// serving reads from it.
+    #[test]
+    fn collect_generate_reports_cached_prompt_tokens_in_usage() {
+        let output = CollectedGenerateOutput {
+            request_id: "raw-usage".to_string(),
+            prompt_logprobs: None,
+            prompt_token_id_logprobs: None,
+            token_ids: vec![30],
+            logprobs: None,
+            finish_reason: FinishReason::stop_eos(),
+            usage: TokenUsage {
+                prompt_token_count: 4,
+                output_token_count: 1,
+                prompt_cache_usage: PromptCacheUsage {
+                    cached_token_count: 3,
+                    cache_creation_token_count: 1,
+                },
+            },
+            kv_transfer_params: None,
+            ec_transfer_params: None,
+            prompt_token_ids: vec![10, 20, 30, 40],
+            sampling_mask: None,
+            spec_decode_metrics: None,
+            weight_versions: None,
+            routed_experts: None,
+        };
+
+        let response = collect_generate(
+            output,
+            "raw-usage".to_string(),
+            "model".to_string(),
+            0,
+            ApiServerOptions {
+                enable_prompt_tokens_details: true,
+                ..Default::default()
+            },
+            ResponseOptions::default(),
+            None,
+        )
+        .expect("response");
+
+        let usage = serde_json::to_value(&response.usage).expect("usage serializes");
+        assert_eq!(usage["prompt_tokens"], 4);
+        assert_eq!(usage["completion_tokens"], 1);
+        assert_eq!(usage["total_tokens"], 5);
+        assert_eq!(usage["prompt_tokens_details"]["cached_tokens"], 3);
+    }
 }
diff --git a/rust/src/server/src/routes/inference/generate/convert.rs b/rust/src/server/src/routes/inference/generate/convert.rs
index 1c4a24ad3db6..f4f9a7b4a0aa 100644
--- a/rust/src/server/src/routes/inference/generate/convert.rs
+++ b/rust/src/server/src/routes/inference/generate/convert.rs
@@ -14,6 +14,8 @@ use crate::utils::{ResolvedRequestContext, merge_ec_transfer_params, merge_kv_tr
 #[derive(Debug, Clone, PartialEq)]
 pub(super) struct PreparedRequest {
     pub request_id: String,
+    /// The `model` reported in non-streaming responses.
+    pub response_model: String,
     pub text_request: TextRequest,
     pub stream: bool,
     /// Public response rendering options for route-layer helpers.
@@ -91,6 +93,7 @@ pub(super) fn prepare_generate_request(
 
     Ok(PreparedRequest {
         request_id: ctx.request_id,
+        response_model: lora_resolution.response_model(),
         text_request,
         stream,
         options: ResponseOptions {
@@ -110,12 +113,13 @@ mod tests {
 
     use super::prepare_generate_request;
     use crate::lora::LoraModelResolution;
+    use crate::model_names::ServedModelNames;
     use crate::routes::inference::generate::types::GenerateRequest;
     use crate::utils::ResolvedRequestContext;
 
     fn served(names: &[&str]) -> LoraModelResolution {
         LoraModelResolution {
-            model_names: names.iter().map(|s| s.to_string()).collect(),
+            model_names: ServedModelNames::new(names.iter().map(|s| s.to_string()).collect()),
             lora_request: None,
         }
     }
diff --git a/rust/src/server/src/routes/inference/generate/types.rs b/rust/src/server/src/routes/inference/generate/types.rs
index 68ca5eee60d1..fbe3a878b67e 100644
--- a/rust/src/server/src/routes/inference/generate/types.rs
+++ b/rust/src/server/src/routes/inference/generate/types.rs
@@ -132,6 +132,9 @@ pub(super) struct GenerateResponseStreamChoice {
 #[derive(Debug, Clone, Serialize)]
 pub(super) struct GenerateStreamResponse {
     pub request_id: String,
+    /// Always `tokens`: the only supported level. Python serializes it on
+    /// every chunk.
+    pub output_mode: GenerateOutputMode,
     pub choices: Vec<GenerateResponseStreamChoice>,
     pub usage: Option<Usage>,
     pub prompt_token_ids: Option<Vec<u32>>,
@@ -143,6 +146,13 @@ pub(super) struct GenerateStreamResponse {
 #[derive(Debug, Clone, Serialize)]
 pub(super) struct GenerateResponse {
     pub request_id: String,
+    pub model: String,
+    pub created: u64,
+    /// Always present, as in Python, which reports token counts and cached
+    /// prompt tokens here for disaggregated serving.
+    pub usage: Usage,
+    /// Always `tokens`: the only supported level.
+    pub output_mode: GenerateOutputMode,
     pub choices: Vec<GenerateResponseChoice>,
     pub prompt_logprobs: Option<Vec<Option<HashMap<u32, GenerateLogprob>>>>,
     /// Base64 `.npy` float32 array, as Python `numpy2base64` encodes it.
diff --git a/rust/src/server/src/routes/inference/generate/validate.rs b/rust/src/server/src/routes/inference/generate/validate.rs
index b53b424f1104..3036225285ed 100644
--- a/rust/src/server/src/routes/inference/generate/validate.rs
+++ b/rust/src/server/src/routes/inference/generate/validate.rs
@@ -3,18 +3,15 @@
 
 use super::types::{GenerateOutputMode, GenerateRequest};
 use crate::error::{ApiError, bail_invalid_request};
+use crate::model_names::ServedModelNames;
 
 /// Enforce the minimal compatibility contract for the Rust token generate
 /// route.
 pub(crate) fn validate_request_compat(
     request: &GenerateRequest,
-    served_model_names: &[String],
+    served_model_names: &ServedModelNames,
 ) -> Result<(), ApiError> {
-    if let Some(model) = request.model.as_ref()
-        && !served_model_names.iter().any(|n| n == model)
-    {
-        return Err(ApiError::model_not_found(model.clone()));
-    }
+    served_model_names.check(request.model.as_deref())?;
 
     if request.sampling_params.n.unwrap_or(1) != 1 {
         bail_invalid_request!(param = "n", "Only n=1 is supported.");
@@ -96,6 +93,7 @@ mod tests {
     use serde_json::json;
 
     use super::validate_request_compat;
+    use crate::model_names::ServedModelNames;
     use crate::routes::inference::generate::types::GenerateRequest;
 
     fn base_request() -> GenerateRequest {
@@ -107,8 +105,8 @@ mod tests {
         .expect("parse request")
     }
 
-    fn served(names: &[&str]) -> Vec<String> {
-        names.iter().map(|s| s.to_string()).collect()
+    fn served(names: &[&str]) -> ServedModelNames {
+        ServedModelNames::new(names.iter().map(|s| s.to_string()).collect())
     }
 
     #[test]
diff --git a/rust/src/server/src/routes/lora.rs b/rust/src/server/src/routes/lora.rs
index ac0cb11e1a77..ec7b37a138b0 100644
--- a/rust/src/server/src/routes/lora.rs
+++ b/rust/src/server/src/routes/lora.rs
@@ -41,6 +41,9 @@ fn load_lora_api_error(error: LoadLoraError) -> ApiError {
     match error {
         LoadLoraError::Disabled(_) => ApiError::invalid_request(message, None),
         LoadLoraError::InvalidAdapter { .. } => ApiError::invalid_request(message, None),
+        LoadLoraError::PathAccess(LoraPathAccessError::NotFound { .. }) => {
+            ApiError::not_found(message)
+        }
         LoadLoraError::PathAccess(LoraPathAccessError::InvalidPath { .. }) => {
             ApiError::invalid_request(message, Some("lora_path"))
         }
@@ -120,3 +123,19 @@ pub async fn unload_lora_adapter(
         lora_request.lora_name
     ))
 }
+
+#[cfg(test)]
+mod tests {
+    use axum::http::StatusCode;
+
+    use super::load_lora_api_error;
+    use crate::lora::{LoadLoraError, LoraPathAccessError};
+
+    #[test]
+    fn load_lora_reports_missing_adapter_path_as_not_found() {
+        let error = load_lora_api_error(LoadLoraError::PathAccess(LoraPathAccessError::NotFound {
+            lora_path: "/not/an/adapter".to_string(),
+        }));
+        assert_eq!(error.status_code(), StatusCode::NOT_FOUND);
+    }
+}
diff --git a/rust/src/server/src/routes/openai/chat_completions.rs b/rust/src/server/src/routes/openai/chat_completions.rs
index d243ffdfd1af..52ba4328bef7 100644
--- a/rust/src/server/src/routes/openai/chat_completions.rs
+++ b/rust/src/server/src/routes/openai/chat_completions.rs
@@ -25,10 +25,11 @@ use vllm_chat::{
     ChatRequest, CollectedAssistantMessage, FinishReason,
 };
 use vllm_engine_core_client::protocol::output::StopReason;
+use vllm_llm::WeightVersionSpan;
 
 use self::convert::{ResponseOptions, prepare_chat_request};
 pub(crate) use self::types::ChatCompletionRequest;
-use crate::config::ApiServerOptions;
+use crate::config::{ApiServerOptions, OutputLogConfig};
 use crate::error::{ApiError, bail_server_error, chat_submit_error, engine_error, server_error};
 use crate::lora::LoraModelResolution;
 use crate::routes::openai::chat_completions::types::{
@@ -45,8 +46,9 @@ use crate::routes::openai::utils::types::{
     ToolCallDelta, Usage,
 };
 use crate::routes::openai::utils::usage::ContinuousUsage;
+use crate::routes::openai::utils::validate_max_output_tokens;
 use crate::routes::openai::utils::validated_json::ValidatedJson;
-use crate::routes::output_log::{LoggedOutput, chat_output_text, log_output};
+use crate::routes::output_log::{LoggedOutput, chat_output_text};
 use crate::state::AppState;
 use crate::utils::{ResolvedRequestContext, resolve_request_context, sse_response, unix_timestamp};
 
@@ -74,6 +76,13 @@ pub async fn chat_completions(
         Ok(prepared) => prepared,
         Err(error) => return error.into_response(),
     };
+    if let Err(error) = validate_max_output_tokens(
+        "max_completion_tokens",
+        prepared.chat_request.sampling_params.max_tokens,
+        state.chat.engine_core_client().max_model_len(),
+    ) {
+        return error.into_response();
+    }
     let request_span = tracing::info_span!(
         "chat_completions",
         request_id = %prepared.request_id,
@@ -82,8 +91,11 @@ pub async fn chat_completions(
 
     let created = unix_timestamp();
     let api_server_options = state.api_server_options;
+    // Like Python, report reasoning-token details only when a reasoning parser
+    // is configured to count them.
+    let count_reasoning_tokens = state.chat.reasoning_parser_name().is_some();
 
-    let chat_stream =
+    let mut chat_stream =
         match state.chat.chat(prepared.chat_request).instrument(request_span.clone()).await {
             Ok(stream) => stream,
             Err(error) => {
@@ -92,12 +104,15 @@ pub async fn chat_completions(
         };
 
     if stream {
+        let prompt_text = chat_stream.take_prompt_text();
         let chunk_stream = chat_completion_chunk_stream(
             chat_stream,
+            prompt_text,
             prepared.request_id,
             prepared.response_model,
             created,
             api_server_options,
+            count_reasoning_tokens,
             prepared.options,
         );
         let sse_stream = chat_completion_sse_stream(chunk_stream).instrument(request_span);
@@ -110,6 +125,7 @@ pub async fn chat_completions(
             prepared.response_model,
             created,
             api_server_options,
+            count_reasoning_tokens,
             prepared.options,
         )
         .instrument(request_span.clone())
@@ -136,6 +152,7 @@ async fn collect_chat_completion(
         enable_per_request_metrics,
         ..
     }: ApiServerOptions,
+    count_reasoning_tokens: bool,
     ResponseOptions {
         // Ignored: non-streaming responses always include usage.
         include_usage: _,
@@ -148,6 +165,7 @@ async fn collect_chat_completion(
         echo,
         return_token_ids,
         return_tokens_as_token_ids,
+        return_weight_versions,
         is_named_tool_choice,
     }: ResponseOptions,
 ) -> Result<ChatCompletionResponse, ApiError> {
@@ -162,6 +180,7 @@ async fn collect_chat_completion(
     let CollectedAssistantMessage {
         message,
         prompt_token_ids,
+        prompt_text,
         prompt_logprobs,
         logprobs,
         token_ids,
@@ -171,6 +190,7 @@ async fn collect_chat_completion(
         finish_reason,
         kv_transfer_params,
         ec_transfer_params,
+        weight_versions,
     } = collected;
     let stop_reason = finish_reason.as_stop_reason().map(stop_reason_to_json);
     let saw_tool_calls = message.tool_calls().next().is_some();
@@ -218,7 +238,11 @@ async fn collect_chat_completion(
         usage.output_token_count,
         spec_decode_metrics,
     );
-    let usage = Usage::from_token_usage(usage, enable_prompt_tokens_details);
+    let usage = Usage::from_token_usage(
+        usage.engine,
+        enable_prompt_tokens_details,
+        count_reasoning_tokens.then_some(usage.reasoning_tokens),
+    );
 
     if enable_log_requests {
         info!(
@@ -245,7 +269,7 @@ async fn collect_chat_completion(
             }),
         );
         if !output.is_empty() {
-            log_output(
+            log_outputs.log(
                 LoggedOutput::Response,
                 &output,
                 Some(&token_ids),
@@ -276,7 +300,7 @@ async fn collect_chat_completion(
             stop_reason,
             token_ids: (return_token_ids && include_output_metadata).then_some(token_ids),
             routed_experts: None,
-            weight_versions: None,
+            weight_versions: weight_versions.filter(|_| return_weight_versions),
         }],
         service_tier: None,
         usage: Some(usage),
@@ -284,7 +308,7 @@ async fn collect_chat_completion(
         system_fingerprint: None,
         prompt_logprobs,
         prompt_token_ids: return_token_ids.then(|| prompt_token_ids.to_vec()),
-        prompt_text: None,
+        prompt_text,
         kv_transfer_params,
         ec_transfer_params,
     })
@@ -294,6 +318,7 @@ async fn collect_chat_completion(
 #[try_stream]
 async fn chat_completion_chunk_stream(
     mut stream: impl ChatEventStreamTrait + Unpin,
+    mut prompt_text: Option<String>,
     request_id: String,
     response_model: String,
     created: u64,
@@ -305,6 +330,7 @@ async fn chat_completion_chunk_stream(
         enable_per_request_metrics,
         ..
     }: ApiServerOptions,
+    count_reasoning_tokens: bool,
     ResponseOptions {
         include_usage,
         include_continuous_usage,
@@ -316,6 +342,7 @@ async fn chat_completion_chunk_stream(
         echo,
         return_token_ids,
         return_tokens_as_token_ids,
+        return_weight_versions,
         is_named_tool_choice,
     }: ResponseOptions,
     mut y: TryYielder<ChatCompletionStreamResponse, ApiError>,
@@ -334,8 +361,8 @@ async fn chat_completion_chunk_stream(
     // starts or ends, omit its token metadata as well as its visible delta.
     let mut inside_hidden_reasoning = false;
     let mut suppress_current_update_metadata = false;
-    let mut continuous_usage = ContinuousUsage::default();
-    let mut delta_log = log_outputs.logs_deltas().then(StreamDeltaLog::default);
+    let mut continuous_usage = ContinuousUsage::new(count_reasoning_tokens);
+    let mut delta_log = log_outputs.logs_deltas().then(|| StreamDeltaLog::new(log_outputs));
 
     /// Yield a chunk with optional continuous usage attached.
     macro_rules! yield_chunk {
@@ -365,7 +392,7 @@ async fn chat_completion_chunk_stream(
                 let mut chunk = start_chunk(&envelope);
                 chunk.prompt = Some(ChatCompletionStreamPrompt {
                     prompt_token_ids: return_token_ids.then(|| prompt_token_ids.to_vec()),
-                    prompt_text: None,
+                    prompt_text: prompt_text.take(),
                 });
                 yield_chunk!(chunk);
                 // When echo=true, emit the last assistant message content as a delta chunk.
@@ -495,6 +522,7 @@ async fn chat_completion_chunk_stream(
                 spec_decode_metrics,
                 finish_reason,
                 kv_transfer_params,
+                weight_versions,
                 ..
             }) => {
                 if enable_log_requests {
@@ -527,6 +555,7 @@ async fn chat_completion_chunk_stream(
                     finish_reason,
                     saw_tool_calls && !is_named_tool_choice,
                     kv_transfer_params.map(|value| *value),
+                    weight_versions.filter(|_| return_weight_versions),
                 ) {
                     Ok(chunk) => {
                         let finish_reason = chunk.choices[0].finish_reason.as_deref();
@@ -534,7 +563,7 @@ async fn chat_completion_chunk_stream(
                             delta_log.finish(finish_reason);
                         }
                         if log_outputs.logs_complete() {
-                            log_output(
+                            log_outputs.log(
                                 LoggedOutput::StreamComplete,
                                 &chat_output_text(
                                     &message.text(),
@@ -566,7 +595,11 @@ async fn chat_completion_chunk_stream(
                     );
                     y.yield_ok(usage_chunk(
                         &envelope,
-                        Usage::from_token_usage(final_usage, enable_prompt_tokens_details),
+                        Usage::from_token_usage(
+                            final_usage.engine,
+                            enable_prompt_tokens_details,
+                            count_reasoning_tokens.then_some(final_usage.reasoning_tokens),
+                        ),
                         metrics,
                     ))
                     .await;
@@ -649,8 +682,10 @@ impl UpdateOutput {
 /// last delta line carries the finish reason: `vllm-chat` emits `Done` after
 /// the last update's `LogprobsDelta`, and a parser may release held-back output
 /// in between, which belongs to that last update.
-#[derive(Debug, Default)]
+#[derive(Debug)]
 struct StreamDeltaLog {
+    /// Where the deltas are logged.
+    log_outputs: OutputLogConfig,
     /// Output of the update in progress.
     current: UpdateOutput,
     /// The last ended update and its output token IDs.
@@ -658,13 +693,21 @@ struct StreamDeltaLog {
 }
 
 impl StreamDeltaLog {
+    fn new(log_outputs: OutputLogConfig) -> Self {
+        Self {
+            log_outputs,
+            current: UpdateOutput::default(),
+            ended: None,
+        }
+    }
+
     /// End the update in progress, whose output token IDs are `token_ids`,
     /// and log the update before it.
     fn end_update(&mut self, token_ids: &[u32]) {
         if let Some((output, token_ids)) = self.ended.take()
             && let Some(delta) = output.describe()
         {
-            log_output(LoggedOutput::StreamDelta, &delta, Some(&token_ids), None);
+            self.log_outputs.log(LoggedOutput::StreamDelta, &delta, Some(&token_ids), None);
         }
         self.ended = Some((std::mem::take(&mut self.current), token_ids.to_vec()));
     }
@@ -677,7 +720,7 @@ impl StreamDeltaLog {
         };
         output.append(std::mem::take(&mut self.current));
         if let Some(delta) = output.describe() {
-            log_output(
+            self.log_outputs.log(
                 LoggedOutput::StreamDelta,
                 &delta,
                 token_ids.as_deref(),
@@ -969,6 +1012,7 @@ fn final_chunk(
     finish_reason: FinishReason,
     use_tool_calls_finish_reason: bool,
     kv_transfer_params: Option<Value>,
+    weight_versions: Option<Vec<WeightVersionSpan>>,
 ) -> Result<ChatCompletionStreamResponse, ApiError> {
     let stop_reason = finish_reason.as_stop_reason().map(stop_reason_to_json);
     let finish_reason = chat_finish_reason_to_openai(&finish_reason, use_tool_calls_finish_reason)?;
@@ -983,6 +1027,7 @@ fn final_chunk(
     chunk.choices.push(ChatCompletionStreamChoice {
         finish_reason: Some(finish_reason.to_string()),
         stop_reason,
+        weight_versions,
         ..Default::default()
     });
     chunk.kv_transfer_params = kv_transfer_params;
@@ -1108,6 +1153,7 @@ mod tests {
             FinishReason::Stop(Some(StopReason::Text("stop".to_string()))),
             false,
             None,
+            None,
         )
         .expect("finish reason is valid");
 
@@ -1117,7 +1163,7 @@ mod tests {
 
     #[test]
     fn final_chunk_maps_length_finish_reason() {
-        let chunk = final_chunk(&stream_envelope(), FinishReason::Length, false, None)
+        let chunk = final_chunk(&stream_envelope(), FinishReason::Length, false, None, None)
             .expect("finish reason is valid");
 
         assert_eq!(chunk.choices[0].finish_reason.as_deref(), Some("length"));
@@ -1126,7 +1172,7 @@ mod tests {
 
     #[test]
     fn final_chunk_maps_abort_finish_reason() {
-        let chunk = final_chunk(&stream_envelope(), FinishReason::Abort, false, None)
+        let chunk = final_chunk(&stream_envelope(), FinishReason::Abort, false, None, None)
             .expect("abort is a valid finish reason");
 
         assert_eq!(chunk.choices[0].finish_reason.as_deref(), Some("abort"));
@@ -1135,13 +1181,28 @@ mod tests {
 
     #[test]
     fn final_chunk_rejects_error_finish_reason() {
-        assert!(final_chunk(&stream_envelope(), FinishReason::Error(None), false, None).is_err());
+        assert!(
+            final_chunk(
+                &stream_envelope(),
+                FinishReason::Error(None),
+                false,
+                None,
+                None
+            )
+            .is_err()
+        );
     }
 
     #[test]
     fn final_chunk_maps_stop_to_tool_calls_when_tool_calls_were_streamed() {
-        let chunk = final_chunk(&stream_envelope(), FinishReason::stop_eos(), true, None)
-            .expect("finish reason is valid");
+        let chunk = final_chunk(
+            &stream_envelope(),
+            FinishReason::stop_eos(),
+            true,
+            None,
+            None,
+        )
+        .expect("finish reason is valid");
 
         assert_eq!(
             chunk.choices[0].finish_reason.as_deref(),
@@ -1197,11 +1258,13 @@ mod tests {
                 finish_reason: FinishReason::stop_eos(),
                 kv_transfer_params: None,
                 ec_transfer_params: None,
+                weight_versions: None,
             }),
         ]);
 
         let chunks = chat_completion_chunk_stream(
             stream,
+            None,
             "chatcmpl-1".to_string(),
             "model".to_string(),
             1,
@@ -1210,6 +1273,7 @@ mod tests {
                 enable_per_request_metrics: true,
                 ..Default::default()
             },
+            false,
             ResponseOptions {
                 include_usage: true,
                 requested_logprobs: true,
@@ -1291,15 +1355,18 @@ mod tests {
                 finish_reason: FinishReason::stop_eos(),
                 kv_transfer_params: None,
                 ec_transfer_params: None,
+                weight_versions: None,
             }),
         ]);
 
         let chunks = chat_completion_chunk_stream(
             stream,
+            None,
             "chatcmpl-1".to_string(),
             "model".to_string(),
             1,
             ApiServerOptions::default(),
+            false,
             ResponseOptions {
                 requested_logprobs: true,
                 include_reasoning: true,
@@ -1338,6 +1405,7 @@ mod tests {
                 usage: done_usage(1, 1, 0),
                 timestamps: Default::default(),
                 spec_decode_metrics: None,
+                weight_versions: None,
                 finish_reason: FinishReason::stop_eos(),
                 kv_transfer_params: Some(Box::new(json!({"remote_block_ids": [1, 2]}))),
                 ec_transfer_params: None,
@@ -1346,10 +1414,12 @@ mod tests {
 
         let chunks = chat_completion_chunk_stream(
             stream,
+            None,
             "chatcmpl-1".to_string(),
             "model".to_string(),
             1,
             ApiServerOptions::default(),
+            false,
             ResponseOptions {
                 include_usage: true,
                 ..Default::default()
@@ -1404,15 +1474,18 @@ mod tests {
                 finish_reason: FinishReason::stop_eos(),
                 kv_transfer_params: None,
                 ec_transfer_params: None,
+                weight_versions: None,
             }),
         ]);
 
         let chunks = chat_completion_chunk_stream(
             stream,
+            None,
             "chatcmpl-1".to_string(),
             "model".to_string(),
             1,
             ApiServerOptions::default(),
+            false,
             ResponseOptions::default(),
         )
         .collect::<Vec<_>>()
@@ -1486,15 +1559,18 @@ mod tests {
                 finish_reason: FinishReason::stop_eos(),
                 kv_transfer_params: None,
                 ec_transfer_params: None,
+                weight_versions: None,
             }),
         ]);
 
         let chunks = chat_completion_chunk_stream(
             stream,
+            None,
             "chatcmpl-1".to_string(),
             "model".to_string(),
             1,
             ApiServerOptions::default(),
+            false,
             ResponseOptions {
                 requested_logprobs: true,
                 return_token_ids: true,
@@ -1620,15 +1696,18 @@ mod tests {
                 finish_reason: FinishReason::stop_eos(),
                 kv_transfer_params: None,
                 ec_transfer_params: None,
+                weight_versions: None,
             }),
         ]);
 
         let chunks = chat_completion_chunk_stream(
             stream,
+            None,
             "chatcmpl-1".to_string(),
             "model".to_string(),
             1,
             ApiServerOptions::default(),
+            false,
             ResponseOptions {
                 requested_logprobs: true,
                 return_token_ids: true,
@@ -1700,15 +1779,18 @@ mod tests {
                 finish_reason: FinishReason::stop_eos(),
                 kv_transfer_params: None,
                 ec_transfer_params: None,
+                weight_versions: None,
             }),
         ]);
 
         let chunks = chat_completion_chunk_stream(
             stream,
+            None,
             "chatcmpl-1".to_string(),
             "model".to_string(),
             1,
             ApiServerOptions::default(),
+            false,
             ResponseOptions {
                 include_reasoning: true,
                 ..Default::default()
@@ -1763,15 +1845,18 @@ mod tests {
                 finish_reason: FinishReason::stop_eos(),
                 kv_transfer_params: None,
                 ec_transfer_params: None,
+                weight_versions: None,
             }),
         ]);
 
         let chunks = chat_completion_chunk_stream(
             stream,
+            None,
             "chatcmpl-1".to_string(),
             "model".to_string(),
             1,
             ApiServerOptions::default(),
+            false,
             ResponseOptions {
                 return_token_ids: true,
                 ..Default::default()
@@ -1801,7 +1886,8 @@ mod tests {
         chunk
             .usage
             .as_ref()
-            .map(|usage| usage.completion_tokens_details.reasoning_tokens)
+            .and_then(|usage| usage.completion_tokens_details.as_ref())
+            .map(|details| details.reasoning_tokens)
     }
 
     #[tokio::test]
@@ -1843,15 +1929,18 @@ mod tests {
                 finish_reason: FinishReason::stop_eos(),
                 kv_transfer_params: None,
                 ec_transfer_params: None,
+                weight_versions: None,
             }),
         ]);
 
         let chunks = chat_completion_chunk_stream(
             stream,
+            None,
             "chatcmpl-1".to_string(),
             "model".to_string(),
             1,
             ApiServerOptions::default(),
+            true,
             ResponseOptions {
                 include_usage: true,
                 include_continuous_usage: true,
@@ -1916,15 +2005,18 @@ mod tests {
                 finish_reason: FinishReason::stop_eos(),
                 kv_transfer_params: None,
                 ec_transfer_params: None,
+                weight_versions: None,
             }),
         ]);
 
         let chunks = chat_completion_chunk_stream(
             stream,
+            None,
             "chatcmpl-1".to_string(),
             "model".to_string(),
             1,
             ApiServerOptions::default(),
+            true,
             ResponseOptions {
                 include_usage: true,
                 include_continuous_usage: true,
@@ -1948,58 +2040,65 @@ mod tests {
         assert_eq!(chunk_reasoning_tokens(&chunks[0]), Some(0));
         let final_usage = chunks.last().unwrap().usage.as_ref().expect("final usage chunk");
         assert_eq!(final_usage.completion_tokens, Some(2));
-        assert_eq!(final_usage.completion_tokens_details.reasoning_tokens, 2);
+        assert_eq!(chunk_reasoning_tokens(chunks.last().unwrap()), Some(2));
     }
 
+    /// Like Python, a configured reasoning parser reports zero reasoning tokens
+    /// for output without reasoning, while a server without one reports the
+    /// details as `null`.
     #[tokio::test]
-    async fn chunk_stream_reports_zero_reasoning_tokens_without_reasoning() {
-        let stream = stream::iter(vec![
-            Ok(ChatEvent::Start {
-                prompt_token_ids: vec![].into(),
-                prompt_logprobs: None,
-            }),
-            Ok(ChatEvent::BlockDelta {
-                index: 0,
-                kind: AssistantBlockKind::Text,
-                delta: "answer".to_string(),
-                token_count: None,
-            }),
-            Ok(ChatEvent::Done {
-                message: Default::default(),
-                usage: done_usage(1, 1, 0),
-                timestamps: Default::default(),
-                spec_decode_metrics: None,
-                finish_reason: FinishReason::stop_eos(),
-                kv_transfer_params: None,
-                ec_transfer_params: None,
-            }),
-        ]);
-
-        let chunks = chat_completion_chunk_stream(
-            stream,
-            "chatcmpl-1".to_string(),
-            "model".to_string(),
-            1,
-            ApiServerOptions::default(),
-            ResponseOptions {
-                include_usage: true,
-                include_continuous_usage: true,
-                ..Default::default()
-            },
-        )
-        .collect::<Vec<_>>()
-        .await
-        .into_iter()
-        .collect::<Result<Vec<_>, _>>()
-        .expect("stream chunks");
-
-        // Reasoning-token details are always reported; a stream without
-        // reasoning simply reports 0 on every usage payload.
-        for chunk in &chunks {
-            let usage = chunk.usage.as_ref().expect("usage");
-            assert_eq!(usage.completion_tokens_details.reasoning_tokens, 0);
-            let json = serde_json::to_value(usage).expect("usage serializes");
-            assert_eq!(json["completion_tokens_details"]["reasoning_tokens"], 0);
+    async fn chunk_stream_reports_reasoning_token_details_only_with_reasoning_parser() {
+        for (count_reasoning_tokens, expected) in
+            [(true, json!({"reasoning_tokens": 0})), (false, json!(null))]
+        {
+            let stream = stream::iter(vec![
+                Ok(ChatEvent::Start {
+                    prompt_token_ids: vec![].into(),
+                    prompt_logprobs: None,
+                }),
+                Ok(ChatEvent::BlockDelta {
+                    index: 0,
+                    kind: AssistantBlockKind::Text,
+                    delta: "answer".to_string(),
+                    token_count: None,
+                }),
+                Ok(ChatEvent::Done {
+                    message: Default::default(),
+                    usage: done_usage(1, 1, 0),
+                    timestamps: Default::default(),
+                    spec_decode_metrics: None,
+                    finish_reason: FinishReason::stop_eos(),
+                    kv_transfer_params: None,
+                    ec_transfer_params: None,
+                    weight_versions: None,
+                }),
+            ]);
+
+            let chunks = chat_completion_chunk_stream(
+                stream,
+                None,
+                "chatcmpl-1".to_string(),
+                "model".to_string(),
+                1,
+                ApiServerOptions::default(),
+                count_reasoning_tokens,
+                ResponseOptions {
+                    include_usage: true,
+                    include_continuous_usage: true,
+                    ..Default::default()
+                },
+            )
+            .collect::<Vec<_>>()
+            .await
+            .into_iter()
+            .collect::<Result<Vec<_>, _>>()
+            .expect("stream chunks");
+
+            for chunk in &chunks {
+                let usage = chunk.usage.as_ref().expect("usage");
+                let json = serde_json::to_value(usage).expect("usage serializes");
+                assert_eq!(json["completion_tokens_details"], expected);
+            }
         }
     }
 }
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 da8b01ac2f33..0761140315a0 100644
--- a/rust/src/server/src/routes/openai/chat_completions/convert.rs
+++ b/rust/src/server/src/routes/openai/chat_completions/convert.rs
@@ -57,6 +57,8 @@ pub(super) struct ResponseOptions {
     /// The request's `return_tokens_as_token_ids`; `None` falls back to the
     /// server's `--return-tokens-as-token-ids`.
     pub return_tokens_as_token_ids: Option<bool>,
+    /// Whether to return the weight-version spans on each finished choice.
+    pub return_weight_versions: bool,
     /// Whether the request forces one named function tool.
     pub is_named_tool_choice: bool,
 }
@@ -80,11 +82,7 @@ pub(super) fn prepare_chat_request(
     .map_err(|error| text_submit_error("invalid prompt truncation", error))?;
 
     let request_id = format!("chatcmpl-{}", ctx.request_id);
-    let response_model = lora_resolution
-        .lora_request
-        .as_ref()
-        .map(|request| request.lora_name.clone())
-        .unwrap_or_else(|| lora_resolution.model_names.first().cloned().unwrap_or_default());
+    let response_model = lora_resolution.response_model();
     let include_reasoning = request.include_reasoning;
     let echo = request
         .echo
@@ -214,6 +212,7 @@ pub(super) fn prepare_chat_request(
         data_parallel_rank: ctx.data_parallel_rank,
         session_id,
         lora_request: lora_resolution.lora_request.clone(),
+        return_prompt_text: request.return_prompt_text.unwrap_or(false),
     };
 
     Ok(PreparedRequest {
@@ -229,6 +228,7 @@ pub(super) fn prepare_chat_request(
             echo,
             return_token_ids: request.return_token_ids.unwrap_or(false),
             return_tokens_as_token_ids: request.return_tokens_as_token_ids,
+            return_weight_versions: request.return_weight_versions.unwrap_or(false),
             is_named_tool_choice,
         },
         chat_request,
@@ -294,9 +294,7 @@ pub(crate) fn convert_message(message: ChatMessage) -> Result<VllmChatMessage, A
             name: _,
         } => {
             let mut blocks = Vec::new();
-            if let Some(reasoning) = reasoning
-                && !reasoning.is_empty()
-            {
+            if let Some(reasoning) = reasoning {
                 blocks.push(AssistantContentBlock::Reasoning { text: reasoning });
             }
             if let Some(content) = content {
@@ -305,7 +303,9 @@ pub(crate) fn convert_message(message: ChatMessage) -> Result<VllmChatMessage, A
             if let Some(tool_calls) = tool_calls {
                 blocks.extend(convert_assistant_tool_calls(tool_calls)?);
             }
-            if blocks.is_empty() {
+            if blocks.iter().all(|block| {
+                matches!(block, AssistantContentBlock::Reasoning { text } if text.is_empty())
+            }) {
                 bail_invalid_request!(
                     "Assistant messages must contain text, reasoning content, or tool_calls."
                 );
@@ -482,6 +482,7 @@ mod tests {
 
     use super::prepare_chat_request;
     use crate::lora::LoraModelResolution;
+    use crate::model_names::ServedModelNames;
     use crate::routes::openai::chat_completions::types::{
         AssistantRole, ChatCompletionMessage, ChatCompletionRequest,
     };
@@ -498,7 +499,7 @@ mod tests {
 
     fn served(names: &[&str]) -> LoraModelResolution {
         LoraModelResolution {
-            model_names: names.iter().map(|s| s.to_string()).collect(),
+            model_names: ServedModelNames::new(names.iter().map(|s| s.to_string()).collect()),
             lora_request: None,
         }
     }
@@ -1428,36 +1429,38 @@ mod tests {
 
     #[test]
     fn prepare_chat_request_accepts_legacy_reasoning_content_alias() {
-        let request = ChatCompletionRequest {
-            messages: vec![
-                serde_json::from_value(json!({
-                    "role": "assistant",
-                    "content": "answer",
-                    "reasoning_content": "inner",
-                }))
-                .expect("legacy reasoning_content alias is accepted"),
-            ],
-            add_generation_prompt: Some(false),
-            ..base_request()
-        };
+        for reasoning in ["inner", ""] {
+            let request = ChatCompletionRequest {
+                messages: vec![
+                    serde_json::from_value(json!({
+                        "role": "assistant",
+                        "content": "answer",
+                        "reasoning_content": reasoning,
+                    }))
+                    .expect("legacy reasoning_content alias is accepted"),
+                ],
+                add_generation_prompt: Some(false),
+                ..base_request()
+            };
 
-        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::assistant_blocks(vec![
-                AssistantContentBlock::Reasoning {
-                    text: "inner".to_string(),
-                },
-                AssistantContentBlock::Text {
-                    text: "answer".to_string(),
-                },
-            ])]
-        );
+            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::assistant_blocks(vec![
+                    AssistantContentBlock::Reasoning {
+                        text: reasoning.to_string(),
+                    },
+                    AssistantContentBlock::Text {
+                        text: "answer".to_string(),
+                    },
+                ])]
+            );
+        }
     }
 
     #[test]
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 4603de95595a..1b9b748a4c32 100644
--- a/rust/src/server/src/routes/openai/chat_completions/types.rs
+++ b/rust/src/server/src/routes/openai/chat_completions/types.rs
@@ -256,7 +256,7 @@ pub struct ChatCompletionRequest {
     pub return_prompt_text: Option<bool>,
 
     /// Include the weight-version spans that sampled each choice's output
-    /// tokens. Not supported yet.
+    /// tokens.
     pub return_weight_versions: Option<bool>,
 
     /// Salt for prefix cache isolation in multi-user environments
@@ -386,7 +386,7 @@ pub(super) struct ChatCompletionResponse {
     pub system_fingerprint: Option<String>,
     pub prompt_logprobs: Option<PromptLogprobs>,
     pub prompt_token_ids: Option<Vec<u32>>,
-    /// Always `None` until `return_prompt_text` is supported.
+    /// The rendered prompt text when the request sets `return_prompt_text`.
     pub prompt_text: Option<String>,
     pub kv_transfer_params: Option<Value>,
     pub ec_transfer_params: Option<Value>,
@@ -404,7 +404,7 @@ pub(super) struct ChatCompletionChoice {
     /// Always `None`.
     /// TODO: return routed experts when `--enable-return-routed-experts` is set.
     pub routed_experts: Option<String>,
-    /// Always `None` until `return_weight_versions` is supported.
+    /// Contiguous weight-version spans over output-token indexes.
     pub weight_versions: Option<Vec<WeightVersionSpan>>,
 }
 
@@ -470,7 +470,7 @@ impl ChatCompletionStreamResponse {
 #[derive(Debug, Clone, Serialize)]
 pub(super) struct ChatCompletionStreamPrompt {
     pub prompt_token_ids: Option<Vec<u32>>,
-    /// Always `None` until `return_prompt_text` is supported.
+    /// The rendered prompt text when the request sets `return_prompt_text`.
     pub prompt_text: Option<String>,
 }
 
@@ -484,6 +484,8 @@ pub(super) struct ChatCompletionStreamChoice {
     pub finish_reason: Option<String>,
     pub stop_reason: Option<Value>,
     pub token_ids: Option<Vec<u32>>,
+    /// Set on the final chunk only; covers the whole output.
+    pub weight_versions: Option<Vec<WeightVersionSpan>>,
 }
 
 /// Mirrors the Python vLLM `DeltaMessage` class.
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 fe98bac3753d..310ff51efa8e 100644
--- a/rust/src/server/src/routes/openai/chat_completions/validate.rs
+++ b/rust/src/server/src/routes/openai/chat_completions/validate.rs
@@ -3,19 +3,16 @@
 
 use super::types::ChatCompletionRequest;
 use crate::error::{ApiError, bail_invalid_request};
+use crate::model_names::ServedModelNames;
 use crate::routes::openai::utils::types::{ChatMessage, Tool};
 use crate::routes::openai::utils::validate_generation_prompt_truncation;
 
 /// Enforce the minimal compatibility contract for the Rust OpenAI server.
 pub(super) fn validate_request_compat(
     request: &ChatCompletionRequest,
-    served_model_names: &[String],
+    served_model_names: &ServedModelNames,
 ) -> Result<(), ApiError> {
-    if let Some(model) = request.model.as_ref().filter(|model| !model.is_empty())
-        && !served_model_names.iter().any(|name| name == model)
-    {
-        return Err(ApiError::model_not_found(model.clone()));
-    }
+    served_model_names.check(request.model.as_deref())?;
 
     if request.stream_options.as_ref().is_some_and(|options| !options.is_empty()) && !request.stream
     {
@@ -101,20 +98,6 @@ pub(super) fn validate_request_compat(
             "spaces_between_special_tokens is not supported."
         );
     }
-    // TODO: return the chat-templated prompt as `prompt_text`.
-    if request.return_prompt_text == Some(true) {
-        bail_invalid_request!(
-            param = "return_prompt_text",
-            "return_prompt_text is not supported."
-        );
-    }
-    // TODO: return the per-choice `weight_versions` spans.
-    if request.return_weight_versions == Some(true) {
-        bail_invalid_request!(
-            param = "return_weight_versions",
-            "return_weight_versions is not supported."
-        );
-    }
     validate_generation_prompt_truncation(request.truncate_prompt_tokens, request.echo)?;
     reject_non_default(
         request.media_io_kwargs.as_ref(),
@@ -160,14 +143,15 @@ mod tests {
     use validator::Validate as _;
 
     use super::validate_request_compat;
+    use crate::model_names::ServedModelNames;
     use crate::routes::openai::chat_completions::types::ChatCompletionRequest;
     use crate::routes::openai::utils::structured_outputs::ResponseFormat;
     use crate::routes::openai::utils::types::{
         ChatMessage, Function, MessageContent, StringOrArray, Tool, ToolChoice, ToolChoiceValue,
     };
 
-    fn served(names: &[&str]) -> Vec<String> {
-        names.iter().map(|s| s.to_string()).collect()
+    fn served(names: &[&str]) -> ServedModelNames {
+        ServedModelNames::new(names.iter().map(|s| s.to_string()).collect())
     }
 
     fn non_streaming_request(stream_options: serde_json::Value) -> ChatCompletionRequest {
@@ -331,42 +315,6 @@ mod tests {
         assert!(validate_request_compat(&request, &served(&["Qwen/Qwen1.5-0.5B-Chat"])).is_err());
     }
 
-    #[test]
-    fn validate_request_compat_rejects_return_prompt_text() {
-        for return_prompt_text in [None, Some(false)] {
-            let request = ChatCompletionRequest {
-                return_prompt_text,
-                ..base_request()
-            };
-            validate_request_compat(&request, &served(&["Qwen/Qwen1.5-0.5B-Chat"]))
-                .expect("an unset or false return_prompt_text should be accepted");
-        }
-
-        let request = ChatCompletionRequest {
-            return_prompt_text: Some(true),
-            ..base_request()
-        };
-        assert!(validate_request_compat(&request, &served(&["Qwen/Qwen1.5-0.5B-Chat"])).is_err());
-    }
-
-    #[test]
-    fn validate_request_compat_rejects_return_weight_versions() {
-        for return_weight_versions in [None, Some(false)] {
-            let request = ChatCompletionRequest {
-                return_weight_versions,
-                ..base_request()
-            };
-            validate_request_compat(&request, &served(&["Qwen/Qwen1.5-0.5B-Chat"]))
-                .expect("an unset or false return_weight_versions should be accepted");
-        }
-
-        let request = ChatCompletionRequest {
-            return_weight_versions: Some(true),
-            ..base_request()
-        };
-        assert!(validate_request_compat(&request, &served(&["Qwen/Qwen1.5-0.5B-Chat"])).is_err());
-    }
-
     #[test]
     fn validate_request_compat_accepts_output_logprobs() {
         let request = ChatCompletionRequest {
diff --git a/rust/src/server/src/routes/openai/completions.rs b/rust/src/server/src/routes/openai/completions.rs
index 6400a9ced88a..c4ecd1ebeaa4 100644
--- a/rust/src/server/src/routes/openai/completions.rs
+++ b/rust/src/server/src/routes/openai/completions.rs
@@ -21,6 +21,7 @@ use thiserror_ext::AsReport as _;
 use tracing::{debug, error, info, trace};
 use tracing_futures::Instrument as _;
 use vllm_engine_core_client::protocol::output::StopReason;
+use vllm_llm::WeightVersionSpan;
 use vllm_text::tokenizer::Tokenizer;
 use vllm_text::{
     DecodedPromptLogprobs, DecodedTextEvent, FinishReason, SampledDelta, TextOutputStream,
@@ -44,8 +45,9 @@ use crate::routes::openai::completions::types::{
 };
 use crate::routes::openai::utils::types::LogProbs;
 use crate::routes::openai::utils::usage::ContinuousUsage;
+use crate::routes::openai::utils::validate_max_output_tokens;
 use crate::routes::openai::utils::validated_json::ValidatedJson;
-use crate::routes::output_log::{LoggedOutput, log_output};
+use crate::routes::output_log::LoggedOutput;
 use crate::state::AppState;
 use crate::utils::{ResolvedRequestContext, resolve_request_context, sse_response, unix_timestamp};
 
@@ -81,6 +83,13 @@ pub async fn completions(
         Ok(prepared) => prepared,
         Err(error) => return error.into_response(),
     };
+    if let Err(error) = validate_max_output_tokens(
+        "max_tokens",
+        prepared.text_request.sampling_params.max_tokens,
+        state.chat.engine_core_client().max_model_len(),
+    ) {
+        return error.into_response();
+    }
     let request_span = tracing::info_span!(
         "completions",
         request_id = %prepared.request_id,
@@ -158,6 +167,7 @@ async fn collect_completion(
         include_prompt_logprobs,
         return_token_ids,
         return_tokens_as_token_ids,
+        return_weight_versions,
     }: ResponseOptions,
 ) -> Result<CompletionResponse, ApiError> {
     let return_tokens_as_token_ids =
@@ -175,7 +185,7 @@ async fn collect_completion(
         let prompt = echo.as_deref().ok_or_else(|| {
             server_error!("prompt-only completion response missing echoed prompt")
         })?;
-        Some(prompt_only_logprobs_to_openai(
+        Some(echoed_prompt_logprobs_to_openai(
             collected.prompt_logprobs.as_ref(),
             prompt,
             collected.prompt_token_ids.as_ref(),
@@ -207,7 +217,7 @@ async fn collect_completion(
         } else {
             (collected.text.as_str(), collected.token_ids.as_slice())
         };
-        log_output(
+        log_outputs.log(
             LoggedOutput::Response,
             output,
             Some(token_ids),
@@ -225,7 +235,7 @@ async fn collect_completion(
         collected.usage.output_token_count,
         collected.spec_decode_metrics,
     );
-    let usage = Usage::from_token_usage(collected.usage, enable_prompt_tokens_details);
+    let usage = Usage::from_token_usage(collected.usage, enable_prompt_tokens_details, None);
 
     if enable_log_requests {
         info!(
@@ -252,7 +262,7 @@ async fn collect_completion(
             token_ids: return_token_ids.then(|| collected.token_ids.clone()),
             prompt_token_ids: return_token_ids.then(|| collected.prompt_token_ids.to_vec()),
             routed_experts: None,
-            weight_versions: None,
+            weight_versions: collected.weight_versions.filter(|_| return_weight_versions),
         }],
         service_tier: None,
         usage: Some(usage),
@@ -288,6 +298,7 @@ async fn completion_chunk_stream(
         include_prompt_logprobs: _,
         return_token_ids,
         return_tokens_as_token_ids,
+        return_weight_versions,
     }: ResponseOptions,
     mut y: TryYielder<CompletionSseChunk, ApiError>,
 ) -> Result<(), ApiError> {
@@ -337,9 +348,10 @@ async fn completion_chunk_stream(
                 debug!("completion stream started");
                 continuous_usage.set_prompt_tokens(prompt_token_ids.len());
                 if let Some(prompt) = echo.as_ref() {
-                    next_text_offset = text_len(prompt);
-                    let logprobs = if prompt_only && requested_logprobs.is_some() {
-                        Some(prompt_only_logprobs_to_openai(
+                    // Like Python, the echoed prompt carries its own logprobs, and
+                    // generated tokens continue from the offsets they report.
+                    let logprobs = if requested_logprobs.is_some() {
+                        Some(echoed_prompt_logprobs_to_openai(
                             prompt_logprobs.as_ref(),
                             prompt,
                             prompt_token_ids.as_ref(),
@@ -348,6 +360,10 @@ async fn completion_chunk_stream(
                     } else {
                         None
                     };
+                    next_text_offset = logprobs
+                        .as_ref()
+                        .and_then(text_end_offset)
+                        .unwrap_or_else(|| text_len(prompt));
                     let mut chunk = delta_chunk(&envelope, prompt.clone(), logprobs);
                     if return_token_ids && first_chunk {
                         if let Some(choice) = chunk.choices.first_mut() {
@@ -399,9 +415,10 @@ async fn completion_chunk_stream(
                             &envelope,
                             finished.finish_reason,
                             finished.kv_transfer_params.map(|value| *value),
+                            finished.weight_versions.filter(|_| return_weight_versions),
                         )?;
                         if let Some(output_log) = output_log.take() {
-                            log_output(
+                            log_outputs.log(
                                 LoggedOutput::StreamComplete,
                                 &output_log,
                                 None,
@@ -422,6 +439,7 @@ async fn completion_chunk_stream(
                                 Usage::from_token_usage(
                                     finished.usage,
                                     enable_prompt_tokens_details,
+                                    None,
                                 ),
                                 metrics,
                             )))
@@ -451,7 +469,7 @@ async fn completion_chunk_stream(
                 if let Some(output_log) = output_log.as_mut() {
                     output_log.push_str(&delta);
                     if log_outputs.logs_deltas() && !delta.is_empty() {
-                        log_output(
+                        log_outputs.log(
                             LoggedOutput::StreamDelta,
                             &delta,
                             Some(&token_ids),
@@ -488,9 +506,10 @@ async fn completion_chunk_stream(
                         &envelope,
                         finished.finish_reason,
                         finished.kv_transfer_params.map(|value| *value),
+                        finished.weight_versions.filter(|_| return_weight_versions),
                     )?;
                     if let Some(output_log) = output_log.take() {
-                        log_output(
+                        log_outputs.log(
                             LoggedOutput::StreamComplete,
                             &output_log,
                             None,
@@ -508,7 +527,11 @@ async fn completion_chunk_stream(
                         );
                         y.yield_ok(CompletionSseChunk::Usage(usage_chunk(
                             &envelope,
-                            Usage::from_token_usage(finished.usage, enable_prompt_tokens_details),
+                            Usage::from_token_usage(
+                                finished.usage,
+                                enable_prompt_tokens_details,
+                                None,
+                            ),
                             metrics,
                         )))
                         .await;
@@ -547,6 +570,7 @@ fn final_chunk(
     envelope: &Arc<StreamResponseEnvelope>,
     finish_reason: FinishReason,
     kv_transfer_params: Option<Value>,
+    weight_versions: Option<Vec<WeightVersionSpan>>,
 ) -> Result<CompletionStreamResponse, ApiError> {
     let stop_reason = finish_reason.as_stop_reason().map(stop_reason_to_json);
     let finish_reason = completion_finish_reason_to_openai(&finish_reason)?;
@@ -555,6 +579,7 @@ fn final_chunk(
     chunk.choices.push(CompletionStreamChoice {
         finish_reason: Some(finish_reason.to_string()),
         stop_reason,
+        weight_versions,
         ..Default::default()
     });
     chunk.kv_transfer_params = kv_transfer_params;
@@ -577,7 +602,7 @@ fn stop_reason_to_json(stop_reason: &StopReason) -> Value {
     serde_json::to_value(stop_reason).expect("StopReason must serialize to JSON")
 }
 
-fn prompt_only_logprobs_to_openai(
+fn echoed_prompt_logprobs_to_openai(
     prompt_logprobs: Option<&DecodedPromptLogprobs>,
     prompt: &str,
     prompt_token_ids: &[u32],
@@ -603,7 +628,7 @@ fn prompt_only_logprobs_to_openai(
     }
 
     Err(server_error!(
-        "prompt-only completion requested logprobs but generation returned none"
+        "echoed completion prompt requested logprobs but generation returned none"
     ))
 }
 
@@ -717,7 +742,7 @@ mod tests {
 
     #[test]
     fn final_chunk_maps_stop_finish_reason() {
-        let chunk = final_chunk(&stream_envelope(), FinishReason::stop_eos(), None)
+        let chunk = final_chunk(&stream_envelope(), FinishReason::stop_eos(), None, None)
             .expect("finish reason valid");
         assert_eq!(chunk.choices[0].finish_reason.as_deref(), Some("stop"));
         assert_eq!(chunk.choices[0].text, "");
@@ -725,21 +750,21 @@ mod tests {
 
     #[test]
     fn final_chunk_maps_length_finish_reason() {
-        let chunk = final_chunk(&stream_envelope(), FinishReason::Length, None)
+        let chunk = final_chunk(&stream_envelope(), FinishReason::Length, None, None)
             .expect("finish reason valid");
         assert_eq!(chunk.choices[0].finish_reason.as_deref(), Some("length"));
     }
 
     #[test]
     fn final_chunk_maps_abort_finish_reason() {
-        let chunk = final_chunk(&stream_envelope(), FinishReason::Abort, None)
+        let chunk = final_chunk(&stream_envelope(), FinishReason::Abort, None, None)
             .expect("finish reason valid");
         assert_eq!(chunk.choices[0].finish_reason.as_deref(), Some("abort"));
     }
 
     #[test]
     fn final_chunk_rejects_error_finish_reason() {
-        assert!(final_chunk(&stream_envelope(), FinishReason::Error(None), None).is_err());
+        assert!(final_chunk(&stream_envelope(), FinishReason::Error(None), None, None).is_err());
     }
 
     #[tokio::test]
@@ -771,6 +796,7 @@ mod tests {
                             first_token_ts: 10.5,
                             last_token_ts: 10.5,
                         },
+                        weight_versions: None,
                     }),
                 )),
             ])
@@ -906,6 +932,7 @@ mod tests {
                         first_token_ts: 10.5,
                         last_token_ts: 11.0,
                     },
+                    weight_versions: None,
                 }),
             )),
         ]);
@@ -1060,6 +1087,7 @@ mod tests {
                         ec_transfer_params: None,
                         sampling_mask: None,
                         spec_decode_metrics: None,
+                        weight_versions: None,
                         timestamps: vllm_llm::RequestTimestamps::default(),
                     }),
                 ))
@@ -1115,6 +1143,7 @@ mod tests {
                     ec_transfer_params: None,
                     sampling_mask: None,
                     spec_decode_metrics: None,
+                    weight_versions: None,
                     timestamps: Default::default(),
                 }),
             )),
@@ -1177,6 +1206,7 @@ mod tests {
                     ec_transfer_params: None,
                     sampling_mask: None,
                     spec_decode_metrics: None,
+                    weight_versions: None,
                     timestamps: Default::default(),
                 }),
             )),
@@ -1243,6 +1273,7 @@ mod tests {
                     ec_transfer_params: None,
                     sampling_mask: None,
                     spec_decode_metrics: None,
+                    weight_versions: None,
                     timestamps: Default::default(),
                 }),
             )),
@@ -1298,6 +1329,7 @@ mod tests {
                     ec_transfer_params: None,
                     sampling_mask: None,
                     spec_decode_metrics: None,
+                    weight_versions: None,
                     timestamps: Default::default(),
                 }),
             )),
@@ -1358,6 +1390,7 @@ mod tests {
                     ec_transfer_params: None,
                     sampling_mask: None,
                     spec_decode_metrics: None,
+                    weight_versions: None,
                     timestamps: Default::default(),
                 }),
             )),
@@ -1437,6 +1470,7 @@ mod tests {
                     ec_transfer_params: None,
                     sampling_mask: None,
                     spec_decode_metrics: None,
+                    weight_versions: None,
                     timestamps: Default::default(),
                 }),
             )),
@@ -1514,6 +1548,7 @@ mod tests {
                     ec_transfer_params: None,
                     sampling_mask: None,
                     spec_decode_metrics: None,
+                    weight_versions: None,
                     timestamps: Default::default(),
                 }),
             )),
diff --git a/rust/src/server/src/routes/openai/completions/convert.rs b/rust/src/server/src/routes/openai/completions/convert.rs
index 790f006f3983..ba0fa740b53e 100644
--- a/rust/src/server/src/routes/openai/completions/convert.rs
+++ b/rust/src/server/src/routes/openai/completions/convert.rs
@@ -49,6 +49,8 @@ pub(super) struct ResponseOptions {
     /// The request's `return_tokens_as_token_ids`; `None` falls back to the
     /// server's `--return-tokens-as-token-ids`.
     pub return_tokens_as_token_ids: Option<bool>,
+    /// Whether to return the weight-version spans on each finished choice.
+    pub return_weight_versions: bool,
 }
 
 /// Validate and lower one OpenAI completions request into the internal
@@ -71,20 +73,11 @@ pub(super) fn prepare_completion_request(
     .map_err(|error| text_submit_error("invalid prompt truncation", error))?;
 
     let request_id = format!("cmpl-{}", ctx.request_id);
-    let response_model = lora_resolution
-        .lora_request
-        .as_ref()
-        .map(|request| request.lora_name.clone())
-        .unwrap_or_else(|| lora_resolution.model_names.first().cloned().unwrap_or_default());
+    let response_model = lora_resolution.response_model();
 
     let logprobs = request.logprobs;
     let prompt_only = request.echo && request.max_tokens == Some(0);
-    let prompt_logprobs =
-        request.prompt_logprobs.or(if request.echo && (!request.stream || prompt_only) {
-            logprobs
-        } else {
-            None
-        });
+    let prompt_logprobs = request.prompt_logprobs.or(if request.echo { logprobs } else { None });
     let include_usage = (request.stream_options.as_ref())
         .and_then(|options| options.include_usage)
         .unwrap_or(false);
@@ -142,7 +135,7 @@ pub(super) fn prepare_completion_request(
             ignore_eos: request.ignore_eos,
             logit_bias: convert_logit_bias(request.logit_bias)?,
             allowed_token_ids: request.allowed_token_ids,
-            bad_words: None,
+            bad_words: request.bad_words,
             bad_words_token_ids: None,
             logprob_token_ids,
             structured_outputs,
@@ -187,6 +180,7 @@ pub(super) fn prepare_completion_request(
             include_prompt_logprobs,
             return_token_ids: request.return_token_ids.unwrap_or(false),
             return_tokens_as_token_ids: request.return_tokens_as_token_ids,
+            return_weight_versions: request.return_weight_versions.unwrap_or(false),
         },
         text_request,
     })
@@ -231,6 +225,7 @@ mod tests {
 
     use super::prepare_completion_request;
     use crate::lora::LoraModelResolution;
+    use crate::model_names::ServedModelNames;
     use crate::routes::openai::completions::types::CompletionRequest;
     use crate::routes::openai::utils::types::Normalizable;
     use crate::utils::{ResolvedRequestContext, resolve_request_context};
@@ -241,7 +236,7 @@ mod tests {
 
     fn served(names: &[&str]) -> LoraModelResolution {
         LoraModelResolution {
-            model_names: names.iter().map(|s| s.to_string()).collect(),
+            model_names: ServedModelNames::new(names.iter().map(|s| s.to_string()).collect()),
             lora_request: None,
         }
     }
@@ -411,6 +406,7 @@ mod tests {
             "frequency_penalty": 0.2,
             "presence_penalty": 0.3,
             "repetition_penalty": 1.1,
+            "bad_words": ["hello"],
             "ignore_eos": true,
             "skip_special_tokens": false,
             "stream_interval": 4
@@ -447,6 +443,10 @@ mod tests {
             prepared.text_request.sampling_params.repetition_penalty,
             Some(1.1)
         );
+        assert_eq!(
+            prepared.text_request.sampling_params.bad_words,
+            Some(vec!["hello".to_string()])
+        );
         assert!(prepared.text_request.sampling_params.ignore_eos);
         assert!(!prepared.text_request.decode_options.skip_special_tokens);
         assert_eq!(
diff --git a/rust/src/server/src/routes/openai/completions/types.rs b/rust/src/server/src/routes/openai/completions/types.rs
index 18ec43c51f11..c1969010c051 100644
--- a/rust/src/server/src/routes/openai/completions/types.rs
+++ b/rust/src/server/src/routes/openai/completions/types.rs
@@ -164,6 +164,9 @@ pub struct CompletionRequest {
     /// Restrict output to these token IDs only
     pub allowed_token_ids: Option<Vec<u32>>,
 
+    /// List of bad words to avoid during generation
+    pub bad_words: Option<Vec<String>>,
+
     /// Number of prompt logprobs to return
     pub prompt_logprobs: Option<i32>,
 
@@ -206,7 +209,7 @@ pub struct CompletionRequest {
     pub return_token_ids: Option<bool>,
 
     /// Include the weight-version spans that sampled each choice's output
-    /// tokens. Not supported yet.
+    /// tokens.
     pub return_weight_versions: Option<bool>,
 
     /// Salt for prefix cache isolation in multi-user environments
@@ -274,7 +277,7 @@ pub(super) struct CompletionChoice {
     /// Always `None`.
     /// TODO: return routed experts when `--enable-return-routed-experts` is set.
     pub routed_experts: Option<String>,
-    /// Always `None` until `return_weight_versions` is supported.
+    /// Contiguous weight-version spans over output-token indexes.
     pub weight_versions: Option<Vec<WeightVersionSpan>>,
 }
 
@@ -315,6 +318,8 @@ pub(super) struct CompletionStreamChoice {
     pub stop_reason: Option<Value>,
     pub token_ids: Option<Vec<u32>>,
     pub prompt_token_ids: Option<Vec<u32>>,
+    /// Set on the final chunk only; covers the whole output.
+    pub weight_versions: Option<Vec<WeightVersionSpan>>,
 }
 
 #[derive(Debug, Clone, Serialize)]
diff --git a/rust/src/server/src/routes/openai/completions/validate.rs b/rust/src/server/src/routes/openai/completions/validate.rs
index e08e35bbe866..7edd66f3ae0b 100644
--- a/rust/src/server/src/routes/openai/completions/validate.rs
+++ b/rust/src/server/src/routes/openai/completions/validate.rs
@@ -3,21 +3,18 @@
 
 use super::types::CompletionRequest;
 use crate::error::{ApiError, bail_invalid_request};
+use crate::model_names::ServedModelNames;
 use crate::routes::openai::utils::validate_generation_prompt_truncation;
 
 /// Enforce the minimal compatibility contract for the Rust OpenAI server.
 pub(super) fn validate_request_compat(
     request: &CompletionRequest,
-    served_model_names: &[String],
+    served_model_names: &ServedModelNames,
 ) -> Result<(), ApiError> {
     // This path is intentionally scoped to the minimum surface needed by
     // `vllm-bench` random workload compatibility, so unsupported legacy
     // completions features fail early here.
-    if let Some(model) = request.model.as_ref().filter(|model| !model.is_empty())
-        && !served_model_names.iter().any(|name| name == model)
-    {
-        return Err(ApiError::model_not_found(model.clone()));
-    }
+    served_model_names.check(request.model.as_deref())?;
 
     if request.stream_options.as_ref().is_some_and(|options| !options.is_empty()) && !request.stream
     {
@@ -46,14 +43,6 @@ pub(super) fn validate_request_compat(
         bail_invalid_request!(param = "prompt_embeds", "prompt_embeds is not supported.");
     }
 
-    // TODO: return the per-choice `weight_versions` spans.
-    if request.return_weight_versions == Some(true) {
-        bail_invalid_request!(
-            param = "return_weight_versions",
-            "return_weight_versions is not supported."
-        );
-    }
-
     if let Some(logprobs) = request.logprobs
         && logprobs < 0
         && logprobs != -1
@@ -126,6 +115,7 @@ mod tests {
     use serde_json::json;
 
     use super::validate_request_compat;
+    use crate::model_names::ServedModelNames;
     use crate::routes::openai::completions::types::CompletionRequest;
 
     fn base_request() -> CompletionRequest {
@@ -137,8 +127,8 @@ mod tests {
         .expect("parse request")
     }
 
-    fn served_names(names: &[&str]) -> Vec<String> {
-        names.iter().map(|s| s.to_string()).collect()
+    fn served_names(names: &[&str]) -> ServedModelNames {
+        ServedModelNames::new(names.iter().map(|s| s.to_string()).collect())
     }
 
     fn non_streaming_request(stream_options: serde_json::Value) -> CompletionRequest {
@@ -247,30 +237,6 @@ mod tests {
         );
     }
 
-    #[test]
-    fn validate_request_compat_rejects_return_weight_versions() {
-        for return_weight_versions in [json!(null), json!(false)] {
-            let request: CompletionRequest = serde_json::from_value(json!({
-                "prompt": "hello",
-                "return_weight_versions": return_weight_versions,
-            }))
-            .expect("parse request");
-            assert!(
-                validate_request_compat(&request, &served_names(&["Qwen/Qwen1.5-0.5B-Chat"]))
-                    .is_ok()
-            );
-        }
-
-        let request: CompletionRequest = serde_json::from_value(json!({
-            "prompt": "hello",
-            "return_weight_versions": true,
-        }))
-        .expect("parse request");
-        assert!(
-            validate_request_compat(&request, &served_names(&["Qwen/Qwen1.5-0.5B-Chat"])).is_err()
-        );
-    }
-
     #[test]
     fn validate_request_compat_accepts_any_served_name() {
         let request = base_request();
diff --git a/rust/src/server/src/routes/openai/utils/mod.rs b/rust/src/server/src/routes/openai/utils/mod.rs
index a5f00ae0eadf..f0159d6cddea 100644
--- a/rust/src/server/src/routes/openai/utils/mod.rs
+++ b/rust/src/server/src/routes/openai/utils/mod.rs
@@ -33,6 +33,27 @@ pub(crate) fn validate_generation_prompt_truncation(
     Ok(())
 }
 
+/// Reject an output budget larger than the whole context window, like
+/// Python's `TokenizeParams` check for chat and completions. A budget that
+/// fits the window is still clamped to the room the prompt leaves when the
+/// request is lowered; `/inference/v1/generate` only clamps, as in Python.
+pub(crate) fn validate_max_output_tokens(
+    param: &'static str,
+    max_tokens: Option<u32>,
+    max_model_len: u32,
+) -> Result<(), ApiError> {
+    if let Some(max_tokens) = max_tokens
+        && max_tokens > max_model_len
+    {
+        bail_invalid_request!(
+            param = param,
+            "{param}={max_tokens} cannot be greater than max_model_len={max_model_len}. \
+             Please request fewer output tokens."
+        );
+    }
+    Ok(())
+}
+
 pub(crate) fn resolve_generation_prompt_truncation(
     limit: Option<i64>,
     side: Option<TruncationSide>,
diff --git a/rust/src/server/src/routes/openai/utils/types.rs b/rust/src/server/src/routes/openai/utils/types.rs
index 8970711c5e73..d9dd47f766f5 100644
--- a/rust/src/server/src/routes/openai/utils/types.rs
+++ b/rust/src/server/src/routes/openai/utils/types.rs
@@ -7,7 +7,6 @@ use indexmap::IndexMap;
 use llm_multimodal::ImageDetail;
 use serde::{Deserialize, Deserializer, Serialize};
 use serde_json::Value;
-use vllm_chat::ChatTokenUsage;
 use vllm_llm::{PromptCacheUsage, TokenUsage};
 
 // ============================================================================
@@ -125,7 +124,7 @@ pub fn validate_stop(stop: &StringOrArray) -> Result<(), validator::ValidationEr
 
 #[cfg(test)]
 mod tests {
-    use super::{StringOrArray, validate_stop};
+    use super::{ChatMessage, StringOrArray, validate_stop};
 
     #[test]
     fn validate_stop_accepts_four_stop_strings() {
@@ -152,6 +151,47 @@ mod tests {
             "stop strings must contain at most 4 items"
         );
     }
+
+    #[test]
+    fn assistant_reasoning_aliases_fold_into_reasoning() {
+        for (fields, expected) in [
+            (
+                r#""reasoning":"canonical","reasoning_content":"legacy","thinking":"thought""#,
+                Some("canonical"),
+            ),
+            (
+                r#""reasoning_content":"legacy","reasoning":"canonical""#,
+                Some("canonical"),
+            ),
+            (
+                r#""reasoning":null,"reasoning_content":"legacy""#,
+                Some("legacy"),
+            ),
+            (r#""reasoning":"","reasoning_content":"legacy""#, Some("")),
+            (
+                r#""reasoning":"canonical","reasoning_content":42"#,
+                Some("canonical"),
+            ),
+            (
+                r#""reasoning_content":"legacy","thinking":"thought""#,
+                Some("legacy"),
+            ),
+            (
+                r#""reasoning_content":null,"thinking":"thought""#,
+                Some("thought"),
+            ),
+            (r#""reasoning_content":null"#, None),
+        ] {
+            let message = serde_json::from_str(&format!(
+                r#"{{"role":"assistant","content":"answer",{fields}}}"#
+            ))
+            .unwrap_or_else(|error| panic!("{fields}: {error}"));
+            let ChatMessage::Assistant { reasoning, .. } = message else {
+                panic!("{fields}: expected an assistant message, got {message:?}");
+            };
+            assert_eq!(reasoning.as_deref(), expected, "{fields}");
+        }
+    }
 }
 
 // ============================================================================
@@ -459,9 +499,8 @@ pub enum ChatMessage {
         content: Option<MessageContent>,
         name: Option<String>,
         tool_calls: Option<Vec<ToolCall>>,
-        /// Reasoning content for reasoning-capable models.
-        #[serde(alias = "reasoning_content")]
-        #[serde(alias = "thinking")]
+        /// Reasoning content for reasoning-capable models. Its aliases are
+        /// folded in by `normalize_reasoning_aliases`.
         reasoning: Option<String>,
     },
     #[serde(rename = "tool")]
@@ -517,14 +556,39 @@ impl Serialize for ChatMessage {
     }
 }
 
+/// Fold the `reasoning_content` and `thinking` aliases of a message into
+/// `reasoning`, so clients replaying history may send several spellings.
+///
+/// Mirrors the Python frontend's `_normalize_messages_before`: aliases are
+/// always removed, and the first non-null one fills `reasoning` only when it is
+/// absent or null. `thinking` is a Rust-only alias with the lowest precedence.
+fn normalize_reasoning_aliases(message: &mut Value) {
+    let Some(message) = message.as_object_mut() else {
+        return;
+    };
+    for alias in ["reasoning_content", "thinking"] {
+        let Some(alias_value) = message.remove(alias) else {
+            continue;
+        };
+        if alias_value.is_null() {
+            continue;
+        }
+        if message.get("reasoning").is_none_or(Value::is_null) {
+            message.insert("reasoning".to_string(), alias_value);
+        }
+    }
+}
+
 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`].
+    /// Normalize the reasoning aliases, then 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 mut value = Value::deserialize(deserializer)?;
+        normalize_reasoning_aliases(&mut value);
         let is_custom = value
             .get("role")
             .and_then(Value::as_str)
@@ -561,9 +625,10 @@ pub struct Usage {
     pub total_tokens: usize,
     pub completion_tokens: Option<usize>,
     pub prompt_tokens_details: Option<PromptTokenUsageInfo>,
-    /// Reasoning-token breakdown.
-    /// Always present, 0 when the configured parser has no reasoning channel.
-    pub completion_tokens_details: CompletionTokenUsageInfo,
+    /// Reasoning-token breakdown. `None` when the frontend cannot count
+    /// reasoning tokens: on chat without a reasoning parser, and always on
+    /// completions and generate.
+    pub completion_tokens_details: Option<CompletionTokenUsageInfo>,
 }
 
 impl Usage {
@@ -573,23 +638,22 @@ impl Usage {
     /// `usage`, and `enable_prompt_tokens_details` only decides whether
     /// prompt-token details are reported. `usage` is destructured exhaustively,
     /// so a new usage field fails to compile until it is mapped here.
+    /// `reasoning_tokens` is `None` when the route does not count reasoning
+    /// tokens.
     pub fn from_token_usage(
-        usage: impl Into<ChatTokenUsage>,
+        usage: TokenUsage,
         enable_prompt_tokens_details: bool,
+        reasoning_tokens: Option<usize>,
     ) -> Self {
-        let ChatTokenUsage {
-            engine:
-                TokenUsage {
-                    prompt_token_count: prompt_tokens,
-                    output_token_count: completion_tokens,
-                    prompt_cache_usage:
-                        PromptCacheUsage {
-                            cached_token_count,
-                            cache_creation_token_count,
-                        },
+        let TokenUsage {
+            prompt_token_count: prompt_tokens,
+            output_token_count: completion_tokens,
+            prompt_cache_usage:
+                PromptCacheUsage {
+                    cached_token_count,
+                    cache_creation_token_count,
                 },
-            reasoning_tokens,
-        } = usage.into();
+        } = usage;
         Self {
             prompt_tokens,
             total_tokens: prompt_tokens + completion_tokens,
@@ -598,7 +662,8 @@ impl Usage {
                 cached_tokens: cached_token_count,
                 created_cache_tokens: cache_creation_token_count,
             }),
-            completion_tokens_details: CompletionTokenUsageInfo { reasoning_tokens },
+            completion_tokens_details: reasoning_tokens
+                .map(|reasoning_tokens| CompletionTokenUsageInfo { reasoning_tokens }),
         }
     }
 }
@@ -620,24 +685,23 @@ pub struct CompletionTokenUsageInfo {
 
 #[cfg(test)]
 mod usage_tests {
-    use vllm_chat::ChatTokenUsage;
+    use serde_json::json;
     use vllm_llm::{PromptCacheUsage, TokenUsage};
 
     use super::Usage;
 
+    const USAGE: TokenUsage = TokenUsage {
+        prompt_token_count: 5,
+        output_token_count: 2,
+        prompt_cache_usage: PromptCacheUsage {
+            cached_token_count: 3,
+            cache_creation_token_count: 1,
+        },
+    };
+
     #[test]
     fn token_usage_hides_prompt_token_details_by_default() {
-        let usage = Usage::from_token_usage(
-            TokenUsage {
-                prompt_token_count: 5,
-                output_token_count: 2,
-                prompt_cache_usage: PromptCacheUsage {
-                    cached_token_count: 3,
-                    cache_creation_token_count: 1,
-                },
-            },
-            false,
-        );
+        let usage = Usage::from_token_usage(USAGE, false, None);
 
         assert_eq!(usage.prompt_tokens, 5);
         assert_eq!(usage.completion_tokens, Some(2));
@@ -646,22 +710,12 @@ mod usage_tests {
 
     #[test]
     fn token_usage_includes_prompt_token_details_when_enabled() {
-        let usage = Usage::from_token_usage(
-            TokenUsage {
-                prompt_token_count: 5,
-                output_token_count: 2,
-                prompt_cache_usage: PromptCacheUsage {
-                    cached_token_count: 3,
-                    cache_creation_token_count: 1,
-                },
-            },
-            true,
-        );
+        let usage = Usage::from_token_usage(USAGE, true, None);
 
         let json = serde_json::to_value(&usage).expect("usage serializes");
         assert_eq!(
             json["prompt_tokens_details"],
-            serde_json::json!({"cached_tokens": 3, "created_cache_tokens": 1})
+            json!({"cached_tokens": 3, "created_cache_tokens": 1})
         );
     }
 
@@ -674,62 +728,25 @@ mod usage_tests {
                 prompt_cache_usage: Default::default(),
             },
             true,
+            None,
         );
         let json = serde_json::to_value(&usage).expect("usage serializes");
         assert_eq!(json["prompt_tokens_details"]["cached_tokens"], 0);
     }
 
+    /// Like Python's `UsageInfo`, reasoning-token details are reported only
+    /// when counted (including a zero count), and `null` otherwise.
     #[test]
-    fn token_usage_includes_reasoning_token_details() {
-        let usage = Usage::from_token_usage(
-            ChatTokenUsage {
-                engine: TokenUsage {
-                    prompt_token_count: 5,
-                    output_token_count: 4,
-                    prompt_cache_usage: Default::default(),
-                },
-                reasoning_tokens: 3,
-            },
-            false,
-        );
-
-        assert_eq!(usage.completion_tokens_details.reasoning_tokens, 3);
-        let json = serde_json::to_value(&usage).expect("usage serializes");
-        assert_eq!(json["completion_tokens_details"]["reasoning_tokens"], 3);
-    }
-
-    #[test]
-    fn token_usage_serializes_zero_reasoning_tokens() {
-        let usage = Usage::from_token_usage(
-            ChatTokenUsage {
-                engine: TokenUsage {
-                    prompt_token_count: 5,
-                    output_token_count: 0,
-                    prompt_cache_usage: Default::default(),
-                },
-                reasoning_tokens: 0,
-            },
-            false,
-        );
-
-        let json = serde_json::to_value(&usage).expect("usage serializes");
-        assert_eq!(json["completion_tokens_details"]["reasoning_tokens"], 0);
-    }
-
-    #[test]
-    fn token_usage_reports_zero_reasoning_tokens_without_reasoning_parser() {
-        let usage = Usage::from_token_usage(
-            TokenUsage {
-                prompt_token_count: 5,
-                output_token_count: 2,
-                prompt_cache_usage: Default::default(),
-            },
-            false,
-        );
-
-        assert_eq!(usage.completion_tokens_details.reasoning_tokens, 0);
-        let json = serde_json::to_value(&usage).expect("usage serializes");
-        assert_eq!(json["completion_tokens_details"]["reasoning_tokens"], 0);
+    fn token_usage_serializes_reasoning_token_details_only_when_counted() {
+        let json = |reasoning_tokens| {
+            serde_json::to_value(Usage::from_token_usage(USAGE, false, reasoning_tokens))
+                .expect("usage serializes")["completion_tokens_details"]
+                .clone()
+        };
+
+        assert_eq!(json(Some(3)), json!({"reasoning_tokens": 3}));
+        assert_eq!(json(Some(0)), json!({"reasoning_tokens": 0}));
+        assert_eq!(json(None), json!(null));
     }
 }
 
diff --git a/rust/src/server/src/routes/openai/utils/usage.rs b/rust/src/server/src/routes/openai/utils/usage.rs
index fee7c426ae35..046be07f948a 100644
--- a/rust/src/server/src/routes/openai/utils/usage.rs
+++ b/rust/src/server/src/routes/openai/utils/usage.rs
@@ -1,7 +1,6 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-use vllm_chat::ChatTokenUsage;
 use vllm_llm::TokenUsage;
 
 use super::types::Usage;
@@ -15,11 +14,21 @@ use super::types::Usage;
 pub(crate) struct ContinuousUsage {
     prompt_tokens: usize,
     output_tokens: usize,
-    /// Running reasoning-token count, following the decoder clock.
-    reasoning_tokens: usize,
+    /// Running reasoning-token count, following the decoder clock. `None`
+    /// when the route does not count reasoning tokens.
+    reasoning_tokens: Option<usize>,
 }
 
 impl ContinuousUsage {
+    /// Create a tracker; with `count_reasoning_tokens`, reasoning tokens are
+    /// counted from zero.
+    pub(crate) fn new(count_reasoning_tokens: bool) -> Self {
+        Self {
+            reasoning_tokens: count_reasoning_tokens.then_some(0),
+            ..Self::default()
+        }
+    }
+
     /// Record the prompt-token count reported when a stream starts.
     pub(crate) fn set_prompt_tokens(&mut self, prompt_tokens: usize) {
         self.prompt_tokens = prompt_tokens;
@@ -32,13 +41,17 @@ impl ContinuousUsage {
 
     /// Add the tokens attributed to one reasoning delta to the running count.
     pub(crate) fn add_reasoning_tokens(&mut self, token_count: usize) {
-        self.reasoning_tokens = self.reasoning_tokens.saturating_add(token_count);
+        if let Some(reasoning_tokens) = &mut self.reasoning_tokens {
+            *reasoning_tokens = reasoning_tokens.saturating_add(token_count);
+        }
     }
 
     /// Converge the running reasoning count onto the authoritative terminal
     /// value reported at stream end.
     pub(crate) fn set_reasoning_tokens(&mut self, reasoning_tokens: usize) {
-        self.reasoning_tokens = reasoning_tokens;
+        if let Some(count) = &mut self.reasoning_tokens {
+            *count = reasoning_tokens;
+        }
     }
 
     /// Replace the running counts with the final counts reported by generation.
@@ -49,14 +62,11 @@ impl ContinuousUsage {
 
     /// Build a streaming usage snapshot without prompt cache details.
     pub(crate) fn to_usage(&self) -> Usage {
-        let usage = ChatTokenUsage {
-            engine: TokenUsage {
-                prompt_token_count: self.prompt_tokens,
-                output_token_count: self.output_tokens,
-                prompt_cache_usage: Default::default(),
-            },
-            reasoning_tokens: self.reasoning_tokens,
+        let usage = TokenUsage {
+            prompt_token_count: self.prompt_tokens,
+            output_token_count: self.output_tokens,
+            prompt_cache_usage: Default::default(),
         };
-        Usage::from_token_usage(usage, false)
+        Usage::from_token_usage(usage, false, self.reasoning_tokens)
     }
 }
diff --git a/rust/src/server/src/routes/output_log.rs b/rust/src/server/src/routes/output_log.rs
index 7138a641b9d3..ac9d044d29be 100644
--- a/rust/src/server/src/routes/output_log.rs
+++ b/rust/src/server/src/routes/output_log.rs
@@ -6,6 +6,8 @@
 
 use tracing::{debug, info};
 
+use crate::config::OutputLogConfig;
+
 /// Which part of a response one output log line covers.
 #[derive(Debug, Clone, Copy)]
 pub(crate) enum LoggedOutput {
@@ -17,24 +19,38 @@ pub(crate) enum LoggedOutput {
     StreamComplete,
 }
 
-/// Log one generated output like Python's `RequestLogger.log_outputs`: the
-/// output text and finish reason at INFO, the output token IDs at DEBUG. The
-/// request ID comes from the enclosing request span.
-pub(crate) fn log_output(
-    kind: LoggedOutput,
-    output: &str,
-    token_ids: Option<&[u32]>,
-    finish_reason: Option<&str>,
-) {
-    let stream = match kind {
-        LoggedOutput::Response => None,
-        LoggedOutput::StreamDelta => Some("delta"),
-        LoggedOutput::StreamComplete => Some("complete"),
-    };
-    if let Some(token_ids) = token_ids {
-        debug!(stream, ?token_ids, "generated response details");
+impl OutputLogConfig {
+    /// Log one generated output like Python's `RequestLogger.log_outputs`: the
+    /// output text and finish reason at INFO, the output token IDs at DEBUG.
+    /// The request ID comes from the enclosing request span. With
+    /// `max_log_len`, only that many leading characters and token IDs are
+    /// logged.
+    pub(crate) fn log(
+        self,
+        kind: LoggedOutput,
+        output: &str,
+        token_ids: Option<&[u32]>,
+        finish_reason: Option<&str>,
+    ) {
+        let stream = match kind {
+            LoggedOutput::Response => None,
+            LoggedOutput::StreamDelta => Some("delta"),
+            LoggedOutput::StreamComplete => Some("complete"),
+        };
+        if let Some(token_ids) = token_ids {
+            let token_ids = match self.max_log_len {
+                Some(max_log_len) => &token_ids[..token_ids.len().min(max_log_len)],
+                None => token_ids,
+            };
+            debug!(stream, ?token_ids, "generated response details");
+        }
+        let output =
+            match self.max_log_len.and_then(|max_log_len| output.char_indices().nth(max_log_len)) {
+                Some((end, _)) => &output[..end],
+                None => output,
+            };
+        info!(stream, ?output, finish_reason, "generated response");
     }
-    info!(stream, ?output, finish_reason, "generated response");
 }
 
 /// Describe a chat output for the log the way Python does: the content, or
diff --git a/rust/src/server/src/routes/render.rs b/rust/src/server/src/routes/render.rs
index a2de402dac2d..efbe12e4af48 100644
--- a/rust/src/server/src/routes/render.rs
+++ b/rust/src/server/src/routes/render.rs
@@ -23,6 +23,7 @@ use crate::routes::inference::generate::{
     GenerateRequest, GenerateSamplingParams, validate_request_compat as validate_generate_request,
 };
 use crate::routes::openai::utils::types::{ListModelsResponse, ModelObject, StreamOptions};
+use crate::routes::openai::utils::validate_max_output_tokens;
 use crate::routes::openai::utils::validated_json::ValidatedJson;
 use crate::routes::openai::{
     ChatCompletionRequest, CompletionRequest, lower_chat_request, lower_completion_request,
@@ -50,6 +51,7 @@ async fn list_models(State(state): State<Arc<RenderState>>) -> Json<ListModelsRe
         object: "list".to_string(),
         data: state
             .served_model_names
+            .names()
             .iter()
             .map(|model| ModelObject {
                 id: model.clone(),
@@ -74,8 +76,7 @@ fn model_resolution(state: &RenderState) -> LoraModelResolution {
 fn response_model(state: &RenderState, requested_model: Option<&str>) -> String {
     requested_model
         .filter(|model| !model.is_empty())
-        .or_else(|| state.served_model_names.first().map(String::as_str))
-        .unwrap_or_default()
+        .unwrap_or(state.served_model_names.primary())
         .to_string()
 }
 
@@ -164,11 +165,17 @@ async fn render_chat(
     let return_token_ids = body.return_token_ids;
     let request_context = resolve_request_context(&headers, body.request_id.as_deref());
     let chat_request = lower_chat_request(body, &model_resolution(&state), request_context)?;
-    let (text_request, _) = state
+    validate_max_output_tokens(
+        "max_completion_tokens",
+        chat_request.sampling_params.max_tokens,
+        state.text.max_model_len(),
+    )?;
+    let text_request = state
         .chat
         .prepare(chat_request, &state.text)
         .await
-        .map_err(|error| ApiError::invalid_request(error.to_report_string(), None))?;
+        .map_err(|error| ApiError::invalid_request(error.to_report_string(), None))?
+        .text_request;
     Ok(Json(lower_render_request(
         &state,
         text_request,
@@ -196,6 +203,11 @@ async fn render_completion(
         request_context,
         tokenizer.as_ref(),
     )?;
+    validate_max_output_tokens(
+        "max_tokens",
+        text_request.sampling_params.max_tokens,
+        state.text.max_model_len(),
+    )?;
     Ok(Json(vec![lower_render_request(
         &state,
         text_request,
diff --git a/rust/src/server/src/routes/sleep.rs b/rust/src/server/src/routes/sleep.rs
index b6239dc2d4b5..33b78169f55b 100644
--- a/rust/src/server/src/routes/sleep.rs
+++ b/rust/src/server/src/routes/sleep.rs
@@ -7,7 +7,7 @@ use axum::Json;
 use axum::extract::rejection::QueryRejection;
 use axum::extract::{Query, RawQuery, State};
 use serde::{Deserialize, Serialize};
-use vllm_engine_core_client::protocol::utility::PauseMode;
+use vllm_engine_core_client::protocol::utility::{PauseMode, SleepLevel};
 use vllm_metrics::METRICS;
 
 use crate::error::ApiError;
@@ -22,7 +22,7 @@ pub(crate) struct IsSleepingResponse {
 #[derive(Serialize)]
 pub struct SleepResponse {
     status: &'static str,
-    level: u32,
+    level: SleepLevel,
 }
 
 #[derive(Serialize)]
@@ -40,13 +40,13 @@ pub struct WakeUpResponse {
 #[derive(Debug, Deserialize)]
 pub(crate) struct SleepParams {
     #[serde(default = "default_sleep_level")]
-    level: u32,
+    level: SleepLevel,
     #[serde(default)]
     mode: PauseMode,
 }
 
-const fn default_sleep_level() -> u32 {
-    1
+const fn default_sleep_level() -> SleepLevel {
+    SleepLevel::OffloadWeights
 }
 
 fn invalid_query(error: QueryRejection) -> ApiError {
@@ -59,12 +59,6 @@ pub async fn sleep(
     params: Result<Query<SleepParams>, QueryRejection>,
 ) -> Result<Json<SleepResponse>, ApiError> {
     let Query(SleepParams { level, mode }) = params.map_err(invalid_query)?;
-    if level > 2 {
-        return Err(ApiError::invalid_request(
-            "level must be between 0 and 2",
-            Some("level"),
-        ));
-    }
     let _recorder = METRICS.api_server.record_sleep_mode_operation("sleep");
 
     state
diff --git a/rust/src/server/src/routes/tests.rs b/rust/src/server/src/routes/tests.rs
index 587810226290..1e3f2e144ea4 100644
--- a/rust/src/server/src/routes/tests.rs
+++ b/rust/src/server/src/routes/tests.rs
@@ -26,6 +26,7 @@ use vllm_chat::{
     ChatRequestProcessor, ChatTextBackend, DefaultChatOutputProcessor, DynChatOutputProcessor,
     DynChatRenderer, NewChatOutputProcessorOptions, ParserSelection,
 };
+use vllm_engine_core_client::mock_engine::DEFAULT_MOCK_MAX_MODEL_LEN;
 use vllm_engine_core_client::mock_engine::default_ready_response;
 use vllm_engine_core_client::protocol::decode_value;
 use vllm_engine_core_client::protocol::dtype::TensorDtype;
@@ -60,11 +61,13 @@ 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, render::build_router as build_render_router,
+    strip_root_path, with_root_path,
 };
 use crate::config::{
-    ApiServerOptions, CorsConfig, DEFAULT_REQUEST_BODY_LIMIT_BYTES, LoraModulePath,
+    AccessLogConfig, ApiServerOptions, CorsConfig, DEFAULT_REQUEST_BODY_LIMIT_BYTES, LoraModulePath,
 };
 use crate::lora::LoadLoraError;
+use crate::model_names::ServedModelNames;
 use crate::render::RenderState;
 use crate::state::AppState;
 
@@ -776,7 +779,7 @@ fn test_render_app_with(
     let backend = Arc::new(FakeChatBackend::new());
     build_render_router(Arc::new(RenderState {
         model: "backend-model".to_string(),
-        served_model_names: vec!["render-model".to_string()],
+        served_model_names: ServedModelNames::new(vec!["render-model".to_string()]),
         max_model_len,
         text: TextRequestProcessor::new(backend.clone(), max_model_len.unwrap_or(u32::MAX)),
         chat: ChatRequestProcessor::render_only(backend)
@@ -1721,6 +1724,113 @@ async fn render_list_models_reports_configured_max_model_len() {
     }
 }
 
+/// Like Python's `TokenizeParams` check, a requested output budget above the
+/// context window is a 400, while one that fits the window is accepted.
+#[tokio::test]
+async fn render_rejects_max_tokens_above_max_model_len() {
+    let mut app = test_render_app_with_max_model_len(Some(64));
+
+    let mut errors = Vec::new();
+    for (uri, body) in [
+        (
+            "/v1/chat/completions/render",
+            json!({
+                "model": "render-model",
+                "messages": [{"role": "user", "content": "hello"}],
+                "max_tokens": 65,
+            }),
+        ),
+        (
+            "/v1/completions/render",
+            json!({"model": "render-model", "prompt": "hello", "max_tokens": 65}),
+        ),
+    ] {
+        let (status, json) = post_json(&mut app, uri, body).await;
+        assert_eq!(status, StatusCode::BAD_REQUEST, "{uri}");
+        errors.push(json.to_string());
+    }
+    expect_test::expect![[r#"
+        {"error":{"message":"max_completion_tokens=65 cannot be greater than max_model_len=64. Please request fewer output tokens.","type":"BadRequestError","param":"max_completion_tokens","code":400}}
+        {"error":{"message":"max_tokens=65 cannot be greater than max_model_len=64. Please request fewer output tokens.","type":"BadRequestError","param":"max_tokens","code":400}}"#]].assert_eq(&errors.join("\n"));
+
+    let (status, json) = post_json(
+        &mut app,
+        "/v1/completions/render",
+        json!({"model": "render-model", "prompt": "hello", "max_tokens": 64}),
+    )
+    .await;
+    assert_eq!(status, StatusCode::OK, "{json}");
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+#[serial]
+async fn render_glm_tools_for_none_follow_server_exclusion_setting() {
+    let dir = tempfile::tempdir().unwrap();
+    fs::write(
+        dir.path().join("config.json"),
+        r#"{"model_type":"glm_moe_dsa"}"#,
+    )
+    .unwrap();
+    fs::write(dir.path().join("tokenizer.json"), "{}").unwrap();
+    let files = vllm_text::backend::hf::ResolvedModelFiles::new(dir.path().to_str().unwrap(), None)
+        .await
+        .unwrap();
+    let mut tokenizer = TestTokenizer::new();
+    for (index, token) in "[gMASK] <sop> <|system|> <|user|> <|assistant|> <|observation|> \
+        <think> </think> <tool_call> </tool_call> <arg_key> </arg_key> <arg_value> </arg_value> \
+        <tool_response> </tool_response> <|begin_of_image|> <|image|> <|end_of_image|> \
+        <|begin_of_video|> <|video|> <|end_of_video|> <|begin_of_audio|> <|end_of_audio|>"
+        .split_whitespace()
+        .enumerate()
+    {
+        tokenizer = tokenizer.with_regular_token(token, 256 + index as u32);
+    }
+    let backend = Arc::new(
+        vllm_chat::HfChatBackend::from_resolved_model_files(
+            files,
+            "glm-test".into(),
+            vllm_chat::LoadModelBackendsOptions {
+                renderer: vllm_chat::RendererSelection::Glm5Next,
+                language_model_only: true,
+                ..Default::default()
+            },
+            Arc::new(tokenizer.clone()),
+        )
+        .unwrap(),
+    );
+    for exclude in [false, true] {
+        let mut app = build_render_router(Arc::new(RenderState {
+            model: "glm-test".into(),
+            served_model_names: ServedModelNames::new(vec!["glm-test".into()]),
+            max_model_len: Some(4096),
+            text: TextRequestProcessor::new(Arc::new(FakeChatBackend::new()), 4096),
+            chat: ChatRequestProcessor::render_only(backend.clone())
+                .with_exclude_tools_when_tool_choice_none(exclude),
+        }));
+        let (status, json) = post_json(
+            &mut app,
+            "/v1/chat/completions/render",
+            json!({
+                "model": "glm-test",
+                "messages": [{"role": "user", "content": "hello"}],
+                "tools": [{"type": "function", "function": {
+                    "name": "get_weather", "parameters": {"type": "object"},
+                }}],
+                "tool_choice": "none",
+            }),
+        )
+        .await;
+        assert_eq!(status, StatusCode::OK, "{json}");
+        let ids: Vec<u32> = serde_json::from_value(json["token_ids"].clone()).unwrap();
+        let prompt = tokenizer.decode(&ids, false).unwrap();
+        assert_eq!(
+            prompt.contains("get_weather"),
+            !exclude,
+            "exclude={exclude}"
+        );
+    }
+}
+
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 #[serial]
 async fn render_chat_response_can_be_submitted_to_generate_unchanged() {
@@ -2005,6 +2115,65 @@ async fn api_key_auth_rejects_missing_token_on_guarded_route() {
     assert_eq!(json, json!({ "error": "Unauthorized" }));
 }
 
+/// The root path is stripped up to a path segment boundary, like Starlette's
+/// `get_route_path`, keeping the query.
+#[test]
+fn strip_root_path_keeps_query_and_segment_boundary() {
+    let rewritten = [
+        "/llm/v1/models?limit=1",
+        "/llm/v1/models",
+        "/llm",
+        "/llmx/v1/models",
+        "/v1/models",
+    ]
+    .map(|uri| {
+        let request = Request::builder().uri(uri).body(Body::empty()).expect("build request");
+        strip_root_path(request, "/llm").uri().to_string()
+    });
+    expect_test::expect![[r#"
+        [
+            "/v1/models?limit=1",
+            "/v1/models",
+            "/",
+            "/llmx/v1/models",
+            "/v1/models",
+        ]
+    "#]]
+    .assert_debug_eq(&rewritten);
+}
+
+/// Like FastAPI's `root_path`, routes are served both under the prefix and at
+/// their bare paths, and API-key auth guards the prefixed paths too. The
+/// prefix only matches up to a path segment boundary.
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+#[serial]
+async fn root_path_serves_routes_under_prefix_and_at_bare_paths() {
+    let (app, _engine_task) = test_app_with_api_keys(vec!["secret".to_string()]).await;
+    let mut app = with_root_path(app, "/llm".to_string());
+
+    for (uri, authorization, expected) in [
+        ("/v1/models", Some("Bearer secret"), StatusCode::OK),
+        ("/llm/v1/models", Some("Bearer secret"), StatusCode::OK),
+        ("/llm/v1/models", None, StatusCode::UNAUTHORIZED),
+        (
+            "/llmx/v1/models",
+            Some("Bearer secret"),
+            StatusCode::NOT_FOUND,
+        ),
+        ("/llm", Some("Bearer secret"), StatusCode::NOT_FOUND),
+    ] {
+        let mut request = Request::builder().uri(uri);
+        if let Some(authorization) = authorization {
+            request = request.header("authorization", authorization);
+        }
+        let response = app
+            .call(request.body(Body::empty()).expect("build request"))
+            .await
+            .expect("call app");
+        assert_eq!(response.status(), expected, "{uri} {authorization:?}");
+    }
+}
+
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 #[serial]
 async fn api_key_auth_rejects_wrong_token_on_guarded_route() {
@@ -3069,6 +3238,48 @@ async fn wrong_model_returns_not_found() {
     assert_eq!(response.status(), StatusCode::NOT_FOUND);
 }
 
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+#[serial]
+async fn skip_model_name_validation_serves_unknown_model_with_base_model() {
+    let (chat, engine_task) = test_models_with_engine_outputs_and_backend(
+        b"engine-openai-skip-model-name",
+        default_stream_output_specs(),
+        Arc::new(FakeChatBackend::new()),
+    )
+    .await;
+    let mut app = build_router(Arc::new(
+        AppState::new(vec!["Qwen/Qwen1.5-0.5B-Chat".to_string()], chat).with_api_server_options(
+            ApiServerOptions {
+                skip_model_name_validation: true,
+                ..Default::default()
+            },
+        ),
+    ));
+
+    let (status, json) = post_json(
+        &mut app,
+        "/v1/chat/completions",
+        json!({
+            "model": "wrong-model",
+            "stream": false,
+            "messages": [{"role": "user", "content": "hello"}]
+        }),
+    )
+    .await;
+    assert_eq!(status, StatusCode::OK, "{json}");
+    assert_eq!(json["model"], "Qwen/Qwen1.5-0.5B-Chat");
+
+    let (status, json) = post_json(
+        &mut app,
+        "/tokenize",
+        json!({"model": "wrong-model", "prompt": "hello"}),
+    )
+    .await;
+    assert_eq!(status, StatusCode::OK, "{json}");
+
+    engine_task.await.expect("mock engine task");
+}
+
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 #[serial]
 async fn invalid_request_returns_openai_error() {
@@ -3185,6 +3396,7 @@ async fn non_stream_chat_returns_json_response() {
     let response_object = json.as_object().expect("response object");
     let choice = json["choices"][0].as_object().expect("choice object");
     let message = choice["message"].as_object().expect("message object");
+    let usage = json["usage"].as_object().expect("usage object");
     for (object, key) in [
         (response_object, "service_tier"),
         (response_object, "system_fingerprint"),
@@ -3201,6 +3413,8 @@ async fn non_stream_chat_returns_json_response() {
         (message, "audio"),
         (message, "function_call"),
         (message, "reasoning"),
+        // No reasoning parser is configured, so reasoning tokens are unknown.
+        (usage, "completion_tokens_details"),
     ] {
         assert!(
             object.contains_key(key) && object[key].is_null(),
@@ -3211,6 +3425,61 @@ 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 chat_return_prompt_text_reports_the_rendered_prompt() {
+    for stream in [false, true] {
+        let (mut app, engine_task) = test_app_with_backend_and_engine_request_check(
+            Arc::new(FakeChatBackend::new()),
+            |_| {},
+        )
+        .await;
+        let response = app
+            .call(
+                Request::builder()
+                    .method("POST")
+                    .uri("/v1/chat/completions")
+                    .header("content-type", "application/json")
+                    .body(Body::from(
+                        json!({
+                            "stream": stream,
+                            "messages": [{"role": "user", "content": "hello"}],
+                            "return_prompt_text": true,
+                        })
+                        .to_string(),
+                    ))
+                    .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");
+        engine_task.await.expect("mock engine task");
+
+        let prompt_texts = if stream {
+            let text = String::from_utf8(body.to_vec()).expect("utf8 body");
+            sse_json_payloads(&text)
+                .into_iter()
+                .map(|chunk| chunk.get("prompt_text").cloned())
+                .collect::<Vec<_>>()
+        } else {
+            let json: serde_json::Value = serde_json::from_slice(&body).expect("decode json");
+            vec![json.get("prompt_text").cloned()]
+        };
+        // Streaming sends it on the first chunk only.
+        let (first, rest) = prompt_texts.split_first().expect("response");
+        assert_eq!(
+            first,
+            &Some(json!("user: hello\nassistant:")),
+            "stream={stream}"
+        );
+        assert!(
+            rest.iter().all(Option::is_none),
+            "stream={stream}: {prompt_texts:?}"
+        );
+    }
+}
+
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 #[serial]
 async fn openai_watermarking_defaults_and_opt_out_reach_engine_core() {
@@ -3260,6 +3529,44 @@ async fn openai_watermarking_defaults_and_opt_out_reach_engine_core() {
     }
 }
 
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+#[serial]
+async fn openai_bad_words_reach_engine_core_as_token_ids() {
+    for endpoint in ["/v1/chat/completions", "/v1/completions"] {
+        let (mut app, engine_task) = test_app_with_backend_and_engine_request_check(
+            Arc::new(FakeChatBackend::new()),
+            |request| {
+                let params = request.sampling_params.as_ref().expect("sampling params");
+                let bad_words_token_ids =
+                    params.bad_words_token_ids.as_ref().expect("bad words token ids");
+                assert!(!bad_words_token_ids.is_empty());
+            },
+        )
+        .await;
+
+        let mut body = json!({"max_tokens": 2, "bad_words": ["blocked"]});
+        if endpoint == "/v1/chat/completions" {
+            body["messages"] = json!([{"role": "user", "content": "hi"}]);
+        } else {
+            body["prompt"] = json!("hi");
+        }
+        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, "{endpoint}");
+        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 chat_structured_outputs_use_configured_backend() {
@@ -3331,6 +3638,121 @@ async fn chat_structured_outputs_use_configured_backend() {
     }
 }
 
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+#[serial]
+async fn openai_return_weight_versions_reports_spans_on_final_choice() {
+    for endpoint in ["/v1/chat/completions", "/v1/completions"] {
+        for stream in [false, true] {
+            for opt_in in [false, true] {
+                let ipc = IpcNamespace::new().expect("create ipc namespace");
+                let handshake_address = ipc.handshake_endpoint();
+                let engine_task = MockEngineTask::new(spawn_mock_engine_task(
+                    handshake_address.clone(),
+                    b"engine-openai-weight-versions".to_vec(),
+                    |dealer, push| {
+                        boxed_test_future(async move {
+                            let add = recv_engine_message(dealer).await;
+                            let request: EngineCoreRequest =
+                                rmp_serde::from_slice(&add[1]).expect("decode request");
+                            // The weights change from `v1` to `v2` before the last token.
+                            for (token_id, version, finish_reason) in [
+                                (b'h', "v1", None),
+                                (b'i', "v1", None),
+                                (b'!', "v2", Some(EngineCoreFinishReason::Length)),
+                            ] {
+                                send_outputs(
+                                    push,
+                                    RequestBatchOutputs {
+                                        outputs: vec![request_output(
+                                            &request.request_id,
+                                            vec![token_id as u32],
+                                            finish_reason,
+                                        )],
+                                        weight_version: Some(version.into()),
+                                        ..Default::default()
+                                    }
+                                    .into(),
+                                )
+                                .await;
+                            }
+                        })
+                    },
+                ));
+                let client = EngineCoreClient::connect(
+                    EngineCoreClientConfig::new_single(handshake_address)
+                        .with_model_name("test-model")
+                        .with_local_input_output_addresses(
+                            Some(ipc.input_endpoint()),
+                            Some(ipc.output_endpoint()),
+                        ),
+                )
+                .await
+                .expect("connect client");
+                let chat = ChatLlm::from_shared_backend(
+                    test_llm(client),
+                    Arc::new(FakeChatBackend::new()),
+                );
+                let mut app = build_router(Arc::new(AppState::new(
+                    vec!["Qwen/Qwen1.5-0.5B-Chat".to_string()],
+                    chat,
+                )));
+
+                let mut body = json!({
+                    "stream": stream,
+                    "max_tokens": 3,
+                    "return_weight_versions": opt_in,
+                });
+                if endpoint == "/v1/chat/completions" {
+                    body["messages"] = json!([{"role": "user", "content": "hi"}]);
+                } else {
+                    body["prompt"] = json!("hi");
+                }
+                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);
+                let bytes = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
+                engine_task.await.expect("mock engine task");
+                let text = String::from_utf8(bytes.to_vec()).expect("utf8 body");
+
+                let expected = json!([
+                    {"version": "v1", "start": 0, "end": 2},
+                    {"version": "v2", "start": 2, "end": 3},
+                ]);
+                let case = format!("{endpoint} stream={stream} opt_in={opt_in}: {text}");
+                if stream {
+                    // Only the final chunk carries the spans, covering the whole output.
+                    let choices: Vec<_> = sse_json_payloads(&text)
+                        .into_iter()
+                        .flat_map(|chunk| chunk["choices"].as_array().cloned().unwrap_or_default())
+                        .collect();
+                    let (finished, unfinished): (Vec<_>, Vec<_>) =
+                        choices.iter().partition(|choice| !choice["finish_reason"].is_null());
+                    assert_eq!(finished.len(), 1, "{case}");
+                    let expected = opt_in.then_some(&expected);
+                    assert_eq!(finished[0].get("weight_versions"), expected, "{case}");
+                    for choice in unfinished {
+                        assert_eq!(choice.get("weight_versions"), None, "{case}");
+                    }
+                } else {
+                    let json: serde_json::Value = serde_json::from_str(&text).expect("json body");
+                    let spans = &json["choices"][0]["weight_versions"];
+                    let expected = if opt_in { expected } else { json!(null) };
+                    assert_eq!(spans, &expected, "{case}");
+                }
+            }
+        }
+    }
+}
+
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 #[serial]
 async fn non_stream_chat_image_url_reaches_engine_mm_features() {
@@ -3367,6 +3789,7 @@ async fn non_stream_chat_image_url_reaches_engine_mm_features() {
                     json!({
                         "model": "Qwen/Qwen1.5-0.5B-Chat",
                         "stream": false,
+                        "return_prompt_text": true,
                         "messages": [{
                             "role": "user",
                             "content": [
@@ -3394,6 +3817,11 @@ async fn non_stream_chat_image_url_reaches_engine_mm_features() {
     let json: serde_json::Value = serde_json::from_slice(&body).expect("decode json");
 
     assert_eq!(json["object"], "chat.completion");
+    // The rendered prompt, before the image placeholder is expanded.
+    assert_eq!(
+        json["prompt_text"],
+        json!("user: describe <|image_pad|>\nassistant:")
+    );
     assert_eq!(json["choices"][0]["message"]["content"], "hi");
 }
 
@@ -4700,6 +5128,7 @@ async fn non_stream_completions_return_json_response() {
     // non-streaming responses.
     let response_object = json.as_object().expect("response object");
     let choice = json["choices"][0].as_object().expect("choice object");
+    let usage = json["usage"].as_object().expect("usage object");
     for (object, key) in [
         (response_object, "service_tier"),
         (response_object, "system_fingerprint"),
@@ -4711,6 +5140,8 @@ async fn non_stream_completions_return_json_response() {
         (choice, "prompt_token_ids"),
         (choice, "routed_experts"),
         (choice, "weight_versions"),
+        // Completions never count reasoning tokens.
+        (usage, "completion_tokens_details"),
     ] {
         assert!(
             object.contains_key(key) && object[key].is_null(),
@@ -5750,6 +6181,14 @@ async fn non_stream_raw_generate_returns_token_output_envelope() {
     let json: serde_json::Value = serde_json::from_slice(&body).expect("decode json");
 
     assert_eq!(json["request_id"], "raw-req");
+    // Like Python's `model_dump()`, the envelope always carries the response
+    // level, the served model name, the creation time and usage.
+    assert_eq!(json["output_mode"], "tokens");
+    assert_eq!(json["model"], "Qwen/Qwen1.5-0.5B-Chat");
+    assert!(json["created"].as_u64().is_some_and(|created| created > 0));
+    assert_eq!(json["usage"]["prompt_tokens"], 2);
+    assert_eq!(json["usage"]["completion_tokens"], 2);
+    assert_eq!(json["usage"]["total_tokens"], 4);
     assert_eq!(json["choices"][0]["index"], 0);
     assert_eq!(json["choices"][0]["token_ids"], json!([33, 44]));
     assert_eq!(json["choices"][0]["finish_reason"], "stop");
@@ -5902,6 +6341,14 @@ async fn stream_raw_generate_returns_sse_chunks_and_usage() {
     assert_eq!(usage["usage"]["completion_tokens"], 2);
     assert_eq!(usage["usage"]["total_tokens"], 4);
     assert_eq!(payloads[3], "[DONE]");
+
+    // Like Python, every chunk names its response level; the non-streaming
+    // `model` and `created` stay off the chunks.
+    for chunk in [&first, &second, &usage] {
+        assert_eq!(chunk["output_mode"], "tokens", "{chunk}");
+        assert!(chunk.get("model").is_none(), "{chunk}");
+        assert!(chunk.get("created").is_none(), "{chunk}");
+    }
 }
 
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
@@ -6426,6 +6873,147 @@ async fn completions_echo_stream_decodes_token_id_prompt_chunk() {
     assert_eq!(usage_chunk["usage"]["completion_tokens"], 3);
 }
 
+/// Like Python, a streamed echo with `logprobs` reports the prompt tokens'
+/// logprobs with the echoed prompt, and the generated tokens' offsets continue
+/// after the reported prompt tokens rather than after the prompt text.
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+#[serial]
+async fn completions_echo_stream_reports_prompt_logprobs() {
+    for return_tokens_as_token_ids in [false, true] {
+        let ipc = IpcNamespace::new().expect("create ipc namespace");
+        let handshake_address = ipc.handshake_endpoint();
+        let engine_task = MockEngineTask::new(spawn_mock_engine_task(
+            handshake_address.clone(),
+            b"engine-openai-completion-echo-stream-logprobs".to_vec(),
+            |dealer, push| {
+                boxed_test_future(async move {
+                    let add = recv_engine_message(dealer).await;
+                    let request: EngineCoreRequest =
+                        rmp_serde::from_slice(&add[1]).expect("decode request");
+                    let params = request.sampling_params.as_ref().expect("sampling params");
+                    assert_eq!(params.prompt_logprobs, Some(1));
+                    for (token_id, alternate_token_id, prompt_logprobs, finish_reason) in [
+                        (b'h', b'H', Some(prompt_logprobs_for_hello()), None),
+                        (b'i', b'I', None, Some(EngineCoreFinishReason::Length)),
+                    ] {
+                        send_outputs(
+                            push,
+                            RequestBatchOutputs {
+                                outputs: vec![request_output_with_logprobs(
+                                    &request.request_id,
+                                    vec![token_id as u32],
+                                    finish_reason,
+                                    None,
+                                    Some(sample_logprobs_for_token(
+                                        token_id as u32,
+                                        alternate_token_id as u32,
+                                    )),
+                                    prompt_logprobs,
+                                )],
+                                ..Default::default()
+                            }
+                            .into(),
+                        )
+                        .await;
+                    }
+                })
+            },
+        ));
+        let client = EngineCoreClient::connect(
+            EngineCoreClientConfig::new_single(handshake_address)
+                .with_model_name("test-model")
+                .with_local_input_output_addresses(
+                    Some(ipc.input_endpoint()),
+                    Some(ipc.output_endpoint()),
+                ),
+        )
+        .await
+        .expect("connect client");
+        let chat = ChatLlm::from_shared_backend(test_llm(client), Arc::new(FakeChatBackend::new()));
+        let mut app = build_router(Arc::new(AppState::new(
+            vec!["Qwen/Qwen1.5-0.5B-Chat".to_string()],
+            chat,
+        )));
+
+        let response = app
+            .call(
+                Request::builder()
+                    .method("POST")
+                    .uri("/v1/completions")
+                    .header("content-type", "application/json")
+                    .body(Body::from(
+                        json!({
+                            "model": "Qwen/Qwen1.5-0.5B-Chat",
+                            "prompt": "hello",
+                            "stream": true,
+                            "echo": true,
+                            "logprobs": 1,
+                            "return_tokens_as_token_ids": return_tokens_as_token_ids,
+                            "add_special_tokens": false
+                        })
+                        .to_string(),
+                    ))
+                    .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");
+        engine_task.await.expect("mock engine task");
+        let text = String::from_utf8(body.to_vec()).expect("utf8 body");
+
+        let choices: Vec<_> = sse_json_payloads(&text)
+            .into_iter()
+            .map(|chunk| chunk["choices"][0].clone())
+            .filter(|choice| choice["logprobs"].is_object())
+            .map(|choice| {
+                let logprobs = &choice["logprobs"];
+                json!([
+                    choice["text"],
+                    logprobs["tokens"],
+                    logprobs["token_logprobs"],
+                    logprobs["top_logprobs"][0],
+                    logprobs["text_offset"],
+                ])
+            })
+            .collect();
+        // As token IDs, each `token_id:NNN` string is 12 characters, so the
+        // generated offsets start after 60 characters, not after `hello`.
+        let expected = if return_tokens_as_token_ids {
+            [
+                json!([
+                    "hello",
+                    [
+                        "token_id:104",
+                        "token_id:101",
+                        "token_id:108",
+                        "token_id:108",
+                        "token_id:111"
+                    ],
+                    [null, -0.3, -0.4, -0.45, -0.5],
+                    null,
+                    [0, 12, 24, 36, 48],
+                ]),
+                json!(["h", ["token_id:104"], [-0.1], {"token_id:104": -0.1, "token_id:72": -0.2}, [60]]),
+                json!(["i", ["token_id:105"], [-0.1], {"token_id:105": -0.1, "token_id:73": -0.2}, [72]]),
+            ]
+        } else {
+            [
+                json!([
+                    "hello",
+                    ["h", "e", "l", "l", "o"],
+                    [null, -0.3, -0.4, -0.45, -0.5],
+                    null,
+                    [0, 1, 2, 3, 4],
+                ]),
+                json!(["h", ["h"], [-0.1], {"h": -0.1, "H": -0.2}, [5]]),
+                json!(["i", ["i"], [-0.1], {"i": -0.1, "I": -0.2}, [6]]),
+            ]
+        };
+        assert_eq!(choices, expected, "{text}");
+    }
+}
+
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 #[serial]
 async fn reasoning_blocks_are_mapped_to_reasoning_sse_chunks() {
@@ -6521,6 +7109,11 @@ async fn include_reasoning_false_suppresses_reasoning_in_non_stream_chat() {
     assert!(
         json["choices"][0]["message"]["reasoning"].is_null(),
         "{text}"
+    ); // The auto-selected reasoning parser still counts the hidden reasoning.
+    assert_eq!(
+        json["usage"]["completion_tokens_details"],
+        json!({"reasoning_tokens": 5}),
+        "{text}"
     );
 }
 
@@ -8881,6 +9474,48 @@ async fn tokenize_allows_prompts_at_or_above_max_model_len() {
     assert_eq!(chat_json["max_model_len"], 4);
 }
 
+/// Like Python's `TokenizeParams` check, chat and completions reject a
+/// requested output budget above the context window before submitting. The
+/// mock engine answers the first request, so a submitted one returns 200.
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+#[serial]
+async fn chat_and_completions_reject_max_tokens_above_max_model_len() {
+    let (mut app, _engine_task) = test_app_with_engine_handle().await;
+    let max_tokens = DEFAULT_MOCK_MAX_MODEL_LEN + 1;
+
+    let mut errors = Vec::new();
+    for (uri, body) in [
+        (
+            "/v1/chat/completions",
+            json!({
+                "model": "Qwen/Qwen1.5-0.5B-Chat",
+                "messages": [{"role": "user", "content": "hello"}],
+                "max_tokens": max_tokens,
+            }),
+        ),
+        (
+            "/v1/chat/completions",
+            json!({
+                "model": "Qwen/Qwen1.5-0.5B-Chat",
+                "messages": [{"role": "user", "content": "hello"}],
+                "max_completion_tokens": max_tokens,
+            }),
+        ),
+        (
+            "/v1/completions",
+            json!({"model": "Qwen/Qwen1.5-0.5B-Chat", "prompt": "hello", "max_tokens": max_tokens}),
+        ),
+    ] {
+        let (status, json) = post_json(&mut app, uri, body).await;
+        assert_eq!(status, StatusCode::BAD_REQUEST, "{uri}: {json}");
+        errors.push(json.to_string());
+    }
+    expect_test::expect![[r#"
+        {"error":{"message":"max_completion_tokens=1048577 cannot be greater than max_model_len=1048576. Please request fewer output tokens.","type":"BadRequestError","param":"max_completion_tokens","code":400}}
+        {"error":{"message":"max_completion_tokens=1048577 cannot be greater than max_model_len=1048576. Please request fewer output tokens.","type":"BadRequestError","param":"max_completion_tokens","code":400}}
+        {"error":{"message":"max_tokens=1048577 cannot be greater than max_model_len=1048576. Please request fewer output tokens.","type":"BadRequestError","param":"max_tokens","code":400}}"#]].assert_eq(&errors.join("\n"));
+}
+
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 #[serial]
 async fn tokenize_chat_conflicting_generation_flags_returns_400() {
@@ -9512,6 +10147,127 @@ async fn profile_routes_are_hidden_when_profiling_is_disabled() {
     engine_task.abort_and_join().await;
 }
 
+/// Captures the access-log events as `LEVEL field=value ...`, with the
+/// nondeterministic `latency` shown only as present.
+#[derive(Clone, Default)]
+struct AccessLogCapture(Arc<std::sync::Mutex<Vec<String>>>);
+
+impl<S: tracing::Subscriber> tracing_subscriber::Layer<S> for AccessLogCapture {
+    fn on_event(
+        &self,
+        event: &tracing::Event<'_>,
+        _ctx: tracing_subscriber::layer::Context<'_, S>,
+    ) {
+        if event.metadata().target() != "vllm_server::middleware::access_log" {
+            return;
+        }
+        let mut line = event.metadata().level().to_string();
+        event.record(
+            &mut |field: &tracing::field::Field, value: &dyn fmt::Debug| match field.name() {
+                "latency" => line.push_str(" latency=<elapsed>"),
+                name => line.push_str(&format!(" {name}={value:?}")),
+            },
+        );
+        self.0.lock().expect("access log lock").push(line);
+    }
+}
+
+/// Send each request to an app with an API key and `access_log`, and return
+/// the access-log lines. The middleware runs on the test's own task, so a
+/// thread-local subscriber sees its events.
+async fn access_log_lines(
+    access_log: AccessLogConfig,
+    requests: Vec<Request<Body>>,
+) -> Vec<String> {
+    use tracing_subscriber::layer::SubscriberExt as _;
+
+    let (chat, engine_task) = test_chat_with_engine_handle().await;
+    let app = build_router(Arc::new(
+        AppState::new(vec!["test-model".to_string()], chat)
+            .with_api_keys(vec!["secret".to_string()])
+            .with_access_log(access_log),
+    ));
+    let capture = AccessLogCapture::default();
+    let _guard =
+        tracing::subscriber::set_default(tracing_subscriber::registry().with(capture.clone()));
+
+    for request in requests {
+        let response = app.clone().oneshot(request).await.expect("call app");
+        to_bytes(response.into_body(), usize::MAX).await.expect("read body");
+    }
+
+    engine_task.abort_and_join().await;
+    capture.0.lock().expect("access log lock").clone()
+}
+
+/// One INFO event per response, like Python's uvicorn access log, with the
+/// client address (absent without one), request target including the query
+/// string, status and latency, including rejected and unrouted requests;
+/// excluded paths match without the query string.
+#[tokio::test]
+#[serial]
+async fn access_log_records_each_response() {
+    let get = |uri: &str, authorized: bool| {
+        let mut request = Request::builder().method("GET").uri(uri);
+        if authorized {
+            request = request.header("authorization", "Bearer secret");
+        }
+        request.body(Body::empty()).expect("build request")
+    };
+    let mut with_client = get("/v1/models?limit=1", true);
+    with_client
+        .extensions_mut()
+        .insert(axum::extract::ConnectInfo(std::net::SocketAddr::from((
+            [127, 0, 0, 1],
+            52344,
+        ))));
+    let access_log = AccessLogConfig {
+        enabled: true,
+        excluded_paths: vec!["/version".to_string()],
+    };
+
+    let lines = access_log_lines(
+        access_log,
+        vec![
+            with_client,
+            get("/v1/models", false),
+            get("/missing", true),
+            get("/version?verbose=1", false),
+        ],
+    )
+    .await;
+
+    expect_test::expect![[r#"
+        [
+            "INFO message=HTTP response client=127.0.0.1:52344 method=GET uri=\"/v1/models?limit=1\" version=HTTP/1.1 status=200 latency=<elapsed>",
+            "INFO message=HTTP response method=GET uri=\"/v1/models\" version=HTTP/1.1 status=401 latency=<elapsed>",
+            "INFO message=HTTP response method=GET uri=\"/missing\" version=HTTP/1.1 status=404 latency=<elapsed>",
+        ]
+    "#]].assert_debug_eq(&lines);
+}
+
+/// `--disable-access-log` removes the access log.
+#[tokio::test]
+#[serial]
+async fn access_log_disabled_logs_nothing() {
+    let lines = access_log_lines(
+        AccessLogConfig {
+            enabled: false,
+            excluded_paths: Vec::new(),
+        },
+        vec![
+            Request::builder()
+                .method("GET")
+                .uri("/version")
+                .body(Body::empty())
+                .expect("build request"),
+        ],
+    )
+    .await;
+
+    assert_eq!(lines, Vec::<String>::new());
+}
+
 /// Captures the events `--enable-log-outputs` logs, one `LEVEL field=value ...`
 /// line per event.
 #[derive(Clone, Default)]
@@ -9587,7 +10343,7 @@ where
 /// Send one request to a fresh app whose engine replies with `output_specs`,
 /// and return the logged outputs.
 async fn output_logs_for_request(
-    log_outputs: crate::config::OutputLogging,
+    api_server_options: ApiServerOptions,
     output_specs: Vec<(Vec<u32>, Option<EngineCoreFinishReason>)>,
     uri: &str,
     body: serde_json::Value,
@@ -9601,9 +10357,8 @@ async fn output_logs_for_request(
     let mut app = build_router(Arc::new(
         AppState::new(vec!["Qwen/Qwen1.5-0.5B-Chat".to_string()], chat).with_api_server_options(
             ApiServerOptions {
-                log_outputs,
                 enable_scale_out: true,
-                ..Default::default()
+                ..api_server_options
             },
         ),
     ));
@@ -9626,6 +10381,16 @@ async fn output_logs_for_request(
     capture.lines()
 }
 
+fn log_outputs(outputs: crate::config::OutputLogging) -> ApiServerOptions {
+    ApiServerOptions {
+        log_outputs: crate::config::OutputLogConfig {
+            outputs,
+            max_log_len: None,
+        },
+        ..Default::default()
+    }
+}
+
 fn chat_body(stream: bool) -> serde_json::Value {
     json!({
         "model": "Qwen/Qwen1.5-0.5B-Chat",
@@ -9648,7 +10413,7 @@ fn reasoning_then_content_output_specs() -> Vec<(Vec<u32>, Option<EngineCoreFini
 #[serial]
 async fn log_outputs_logs_chat_stream_deltas_and_complete_output() {
     let logs = output_logs_for_request(
-        crate::config::OutputLogging::CompleteAndDeltas,
+        log_outputs(crate::config::OutputLogging::CompleteAndDeltas),
         reasoning_then_content_output_specs(),
         "/v1/chat/completions",
         chat_body(true),
@@ -9676,7 +10441,7 @@ async fn log_outputs_logs_chat_text_released_at_finish_with_its_update() {
         "function": {"name": "get_weather", "parameters": {"type": "object"}}
     }]);
     let logs = output_logs_for_request(
-        crate::config::OutputLogging::CompleteAndDeltas,
+        log_outputs(crate::config::OutputLogging::CompleteAndDeltas),
         vec![(
             token_ids("Hello<tool_ca"),
             Some(EngineCoreFinishReason::Length),
@@ -9696,7 +10461,7 @@ async fn log_outputs_logs_chat_text_released_at_finish_with_its_update() {
 #[serial]
 async fn log_outputs_logs_chat_deltas_sent_before_a_stream_error() {
     let logs = output_logs_for_request(
-        crate::config::OutputLogging::CompleteAndDeltas,
+        log_outputs(crate::config::OutputLogging::CompleteAndDeltas),
         vec![
             (token_ids("Hello"), None),
             (token_ids("!"), Some(EngineCoreFinishReason::Error)),
@@ -9718,7 +10483,7 @@ async fn log_outputs_logs_chat_deltas_sent_before_a_stream_error() {
 #[serial]
 async fn log_outputs_logs_non_stream_chat_content() {
     let logs = output_logs_for_request(
-        crate::config::OutputLogging::CompleteAndDeltas,
+        log_outputs(crate::config::OutputLogging::CompleteAndDeltas),
         reasoning_then_content_output_specs(),
         "/v1/chat/completions",
         chat_body(false),
@@ -9734,7 +10499,7 @@ async fn log_outputs_logs_non_stream_chat_content() {
 #[serial]
 async fn log_outputs_logs_completion_stream_deltas_and_complete_output() {
     let logs = output_logs_for_request(
-        crate::config::OutputLogging::CompleteAndDeltas,
+        log_outputs(crate::config::OutputLogging::CompleteAndDeltas),
         default_stream_output_specs(),
         "/v1/completions",
         json!({"model": "Qwen/Qwen1.5-0.5B-Chat", "prompt": "hello", "stream": true}),
@@ -9754,7 +10519,7 @@ async fn log_outputs_logs_completion_stream_deltas_and_complete_output() {
 #[serial]
 async fn log_outputs_without_deltas_logs_only_complete_completion_stream() {
     let logs = output_logs_for_request(
-        crate::config::OutputLogging::Complete,
+        log_outputs(crate::config::OutputLogging::Complete),
         default_stream_output_specs(),
         "/v1/completions",
         json!({"model": "Qwen/Qwen1.5-0.5B-Chat", "prompt": "hello", "stream": true}),
@@ -9771,14 +10536,14 @@ async fn log_outputs_without_deltas_logs_only_complete_completion_stream() {
 #[serial]
 async fn log_outputs_logs_non_stream_completion_and_generate_token_ids() {
     let completion = output_logs_for_request(
-        crate::config::OutputLogging::CompleteAndDeltas,
+        log_outputs(crate::config::OutputLogging::CompleteAndDeltas),
         default_stream_output_specs(),
         "/v1/completions",
         json!({"model": "Qwen/Qwen1.5-0.5B-Chat", "prompt": "hello"}),
     )
     .await;
     let generate = output_logs_for_request(
-        crate::config::OutputLogging::CompleteAndDeltas,
+        log_outputs(crate::config::OutputLogging::CompleteAndDeltas),
         default_stream_output_specs(),
         "/inference/v1/generate",
         json!({
@@ -9798,11 +10563,74 @@ async fn log_outputs_logs_non_stream_completion_and_generate_token_ids() {
     .assert_eq(&format!("{completion}\n---\n{generate}"));
 }
 
+/// Like Python's `RequestLogger`, `--max-log-len` bounds the logged output
+/// characters and output token IDs.
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+#[serial]
+async fn max_log_len_truncates_logged_outputs_and_token_ids() {
+    let logs = output_logs_for_request(
+        ApiServerOptions {
+            log_outputs: crate::config::OutputLogConfig {
+                outputs: crate::config::OutputLogging::CompleteAndDeltas,
+                max_log_len: Some(1),
+            },
+            ..Default::default()
+        },
+        reasoning_then_content_output_specs(),
+        "/v1/chat/completions",
+        chat_body(true),
+    )
+    .await;
+
+    expect_test::expect![[r#"
+        DEBUG message=generated response details stream="delta" token_ids=[61441]
+        INFO message=generated response stream="delta" output="["
+        DEBUG message=generated response details stream="delta" token_ids=[72]
+        INFO message=generated response stream="delta" output="H"
+        DEBUG message=generated response details stream="delta" token_ids=[33]
+        INFO message=generated response stream="delta" output="!" finish_reason="length"
+        INFO message=generated response stream="complete" output="H" finish_reason="length""#]]
+    .assert_eq(&logs);
+}
+
+/// `--max-log-len` counts characters, not bytes, like Python's `str`
+/// slicing, and 0 logs empty values.
+#[test]
+#[serial]
+fn max_log_len_counts_characters() {
+    use super::output_log::LoggedOutput;
+
+    let capture = OutputLogCapture::start();
+    tracing::info_span!("request", request_id = "test").in_scope(|| {
+        for max_log_len in [Some(2), Some(0), Some(100)] {
+            let log_outputs = crate::config::OutputLogConfig {
+                outputs: crate::config::OutputLogging::Complete,
+                max_log_len,
+            };
+            log_outputs.log(
+                LoggedOutput::Response,
+                "中😀a",
+                Some(&[1, 2, 3]),
+                Some("stop"),
+            );
+        }
+    });
+
+    expect_test::expect![[r#"
+        DEBUG message=generated response details token_ids=[1, 2]
+        INFO message=generated response output="中😀" finish_reason="stop"
+        DEBUG message=generated response details token_ids=[]
+        INFO message=generated response output="" finish_reason="stop"
+        DEBUG message=generated response details token_ids=[1, 2, 3]
+        INFO message=generated response output="中😀a" finish_reason="stop""#]]
+    .assert_eq(&capture.lines());
+}
+
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 #[serial]
 async fn log_outputs_disabled_logs_nothing() {
     let logs = output_logs_for_request(
-        crate::config::OutputLogging::Disabled,
+        log_outputs(crate::config::OutputLogging::Disabled),
         default_stream_output_specs(),
         "/v1/completions",
         json!({"model": "Qwen/Qwen1.5-0.5B-Chat", "prompt": "hello", "stream": true}),
diff --git a/rust/src/server/src/routes/tokenize.rs b/rust/src/server/src/routes/tokenize.rs
index 87da11f42291..3fd5e544a98c 100644
--- a/rust/src/server/src/routes/tokenize.rs
+++ b/rust/src/server/src/routes/tokenize.rs
@@ -34,16 +34,6 @@ fn tokenize_request_id(headers: &HeaderMap) -> String {
     format!("tokenize-{base}")
 }
 
-/// Reject an unknown model name, matching the other handlers.
-fn check_model(state: &AppState, model: Option<&str>) -> Result<(), ApiError> {
-    if let Some(model) = model
-        && !state.served_model_names().iter().any(|n| n == model)
-    {
-        return Err(ApiError::model_not_found(model.to_string()));
-    }
-    Ok(())
-}
-
 /// Build the `token_strs` vector when requested, via the tokenizer vocab.
 fn token_strs(tokenizer: &vllm_text::tokenizer::DynTokenizer, ids: &[u32]) -> Vec<String> {
     // Unknown IDs yield "" — intentional; matches Python's convert_ids_to_tokens behaviour.
@@ -84,7 +74,7 @@ fn tokenize_completion(
     request_id: &str,
     req: TokenizeCompletionRequest,
 ) -> Result<(Vec<u32>, bool), ApiError> {
-    check_model(state, req.model.as_deref())?;
+    state.base_model_names().check(req.model.as_deref())?;
     let return_token_strs = req.return_token_strs;
     let tokens = state
         .chat
@@ -104,7 +94,7 @@ async fn tokenize_chat(
     request_id: &str,
     req: TokenizeChatRequest,
 ) -> Result<(Vec<u32>, bool), ApiError> {
-    check_model(state, req.model.as_deref())?;
+    state.base_model_names().check(req.model.as_deref())?;
     let return_token_strs = req.return_token_strs;
     // `continue_final_message` / `add_generation_prompt` mutual exclusion is
     // enforced in `normalize_generation_prompt_mode` inside `into_chat_request`.
@@ -127,7 +117,7 @@ pub async fn detokenize(
     State(state): State<Arc<AppState>>,
     ValidatedJson(body): ValidatedJson<DetokenizeRequest>,
 ) -> Response {
-    if let Err(error) = check_model(&state, body.model.as_deref()) {
+    if let Err(error) = state.base_model_names().check(body.model.as_deref()) {
         return error.into_response();
     }
     let tokenizer = state.chat.text().tokenizer();
diff --git a/rust/src/server/src/routes/tokenize/types.rs b/rust/src/server/src/routes/tokenize/types.rs
index a449304101b7..208abe1cd056 100644
--- a/rust/src/server/src/routes/tokenize/types.rs
+++ b/rust/src/server/src/routes/tokenize/types.rs
@@ -132,6 +132,7 @@ impl TokenizeChatRequest {
             data_parallel_rank: None,
             session_id: None,
             lora_request: None,
+            return_prompt_text: false,
         })
     }
 }
diff --git a/rust/src/server/src/state.rs b/rust/src/server/src/state.rs
index a23fa3668dbe..ea370b3e863c 100644
--- a/rust/src/server/src/state.rs
+++ b/rust/src/server/src/state.rs
@@ -13,10 +13,11 @@ use vllm_chat::ChatLlm;
 use vllm_engine_core_client::EngineCoreClient;
 use vllm_engine_core_client::runtime::BackgroundShutdownRuntime;
 
-use crate::config::{ApiServerOptions, CorsConfig, LoraModulePath};
+use crate::config::{AccessLogConfig, ApiServerOptions, CorsConfig, LoraModulePath};
 use crate::lora::{
     LoadLoraError, LoadedLora, LoraDisabledError, LoraManager, LoraModelResolution, UnloadLoraError,
 };
+use crate::model_names::ServedModelNames;
 use crate::runtime::build_request_runtime;
 use crate::server_info::{ServerInfoConfigFormat, ServerInfoSnapshot};
 
@@ -39,6 +40,8 @@ pub struct AppState {
     pub api_server_options: ApiServerOptions,
     /// CORS settings applied to every HTTP response.
     pub cors: CorsConfig,
+    /// HTTP access log settings.
+    pub access_log: AccessLogConfig,
     /// Runtime server information returned by `/server_info`, when available.
     server_info: Option<ServerInfoSnapshot>,
     /// SHA-256 hashes of API keys accepted as bearer tokens for guarded routes.
@@ -75,6 +78,7 @@ impl AppState {
             chat,
             api_server_options: ApiServerOptions::default(),
             cors: CorsConfig::default(),
+            access_log: AccessLogConfig::default(),
             server_info: None,
             api_key_hashes: Vec::new(),
             server_load: AtomicU64::new(0),
@@ -97,6 +101,12 @@ impl AppState {
         self
     }
 
+    /// Set the HTTP access log settings.
+    pub fn with_access_log(mut self, access_log: AccessLogConfig) -> Self {
+        self.access_log = access_log;
+        self
+    }
+
     /// Set the backend model path reported as `root` for base-model cards.
     pub fn with_model_path(mut self, model_path: String) -> Self {
         self.model_path = Some(model_path);
@@ -170,9 +180,15 @@ impl AppState {
         Ok(self.lora_manager.served_loras().await)
     }
 
+    /// The base model names a request may select, before LoRA adapters.
+    pub(crate) fn base_model_names(&self) -> ServedModelNames {
+        ServedModelNames::new(self.served_model_names.clone())
+            .with_accept_any(self.api_server_options.skip_model_name_validation)
+    }
+
     /// Resolve the requested model against one dynamic LoRA registry snapshot.
     pub async fn resolve_model_with_loras(&self, model_name: Option<&str>) -> LoraModelResolution {
-        self.lora_manager.resolve_model(&self.served_model_names, model_name).await
+        self.lora_manager.resolve_model(self.base_model_names(), model_name).await
     }
 
     /// Load one dynamic LoRA adapter and register it as a public model name.
diff --git a/rust/src/server/src/tls_tests.rs b/rust/src/server/src/tls_tests.rs
index f553d4b5c345..c70ebeff3825 100644
--- a/rust/src/server/src/tls_tests.rs
+++ b/rust/src/server/src/tls_tests.rs
@@ -5,10 +5,12 @@
 //! through the production listener/connection path, with a trivial router since TLS
 //! terminates below the app.
 
+use std::net::SocketAddr;
 use std::pin::Pin;
 use std::time::Duration;
 
 use axum::Router;
+use axum::extract::ConnectInfo;
 use axum::routing::get;
 use openssl::asn1::Asn1Time;
 use openssl::bn::{BigNum, MsbOption};
@@ -279,7 +281,10 @@ async fn spawn_server_with_timeouts(
 
     let server_config =
         tls_config.map(|cfg| tls::build_server_config(&cfg).expect("build server config"));
-    let app = Router::new().route("/health", get(|| async { "ok" }));
+    let app = Router::new().route("/health", get(|| async { "ok" })).route(
+        "/peer",
+        get(|ConnectInfo(peer): ConnectInfo<SocketAddr>| async move { peer.to_string() }),
+    );
     let shutdown = CancellationToken::new();
     let server_shutdown = shutdown.clone();
     tokio::spawn(async move {
@@ -359,10 +364,12 @@ async fn legacy_suite_handshake(certs: &TestCerts, addr: &str) -> std::io::Resul
     stream.as_mut().connect().await.map_err(std::io::Error::other)
 }
 
-async fn plain_get(addr: &str) -> std::io::Result<String> {
+async fn plain_get(addr: &str, path: &str) -> std::io::Result<String> {
     let mut tcp = TcpStream::connect(addr).await?;
-    tcp.write_all(b"GET /health HTTP/1.1\r\nHost: 127.0.0.1\r\nConnection: close\r\n\r\n")
-        .await?;
+    tcp.write_all(
+        format!("GET {path} HTTP/1.1\r\nHost: 127.0.0.1\r\nConnection: close\r\n\r\n").as_bytes(),
+    )
+    .await?;
     let mut response = String::new();
     tcp.read_to_string(&mut response).await?;
     Ok(response)
@@ -514,8 +521,18 @@ async fn mtls_rejects_untrusted_client_certificate() {
 #[tokio::test]
 async fn plain_http_serves_when_tls_is_disabled() {
     let (addr, shutdown) = spawn_server(None).await;
-    let response = plain_get(&addr).await.expect("http request");
+    let response = plain_get(&addr, "/health").await.expect("http request");
+    assert!(response.starts_with("HTTP/1.1 200"), "{response}");
+    shutdown.cancel();
+}
+
+/// Handlers, such as the access log, see the TCP peer as `ConnectInfo`.
+#[tokio::test]
+async fn plain_http_exposes_tcp_peer_as_connect_info() {
+    let (addr, shutdown) = spawn_server(None).await;
+    let response = plain_get(&addr, "/peer").await.expect("http request");
     assert!(response.starts_with("HTTP/1.1 200"), "{response}");
+    assert!(response.contains("\r\n\r\n127.0.0.1:"), "{response}");
     shutdown.cancel();
 }
 
diff --git a/rust/src/text/src/backend/hf/config.rs b/rust/src/text/src/backend/hf/config.rs
index 9f22e7989329..d204757fa0ce 100644
--- a/rust/src/text/src/backend/hf/config.rs
+++ b/rust/src/text/src/backend/hf/config.rs
@@ -17,7 +17,7 @@ use crate::normalize_top_k;
 pub struct HfTokenizerConfig {
     #[serde(flatten)]
     pub special_tokens: HfSpecialTokens,
-    pub chat_template: Option<String>,
+    pub chat_template: Option<HfChatTemplate>,
     /// The `tokenizer_class` field from HuggingFace tokenizer configs. Some
     /// tiktoken-based models (e.g. DeepSeek, Kimi K2) set this to a value
     /// containing "Tiktoken" which can be used as a hint for backend
@@ -28,6 +28,53 @@ pub struct HfTokenizerConfig {
     pub response_template: Option<serde_json::Value>,
 }
 
+/// The `chat_template` field of `tokenizer_config.json`: one template, or the
+/// named templates Transformers serializes as a list of `{name, template}`
+/// (e.g. `default` and `tool_use`).
+#[derive(Debug, Clone, Deserialize)]
+#[serde(untagged)]
+pub enum HfChatTemplate {
+    Single(String),
+    Named(Vec<HfNamedChatTemplate>),
+}
+
+/// One entry of a list-form `chat_template`.
+#[derive(Debug, Clone, Deserialize)]
+pub struct HfNamedChatTemplate {
+    pub name: String,
+    pub template: String,
+}
+
+impl HfChatTemplate {
+    /// The template Transformers uses when no tools are passed: the single
+    /// template, or the one named `default`.
+    pub fn default_template(&self) -> Option<&str> {
+        match self {
+            Self::Single(template) => Some(template),
+            Self::Named(_) => self.named("default"),
+        }
+    }
+
+    /// The template Transformers prefers when tools are passed: the one named
+    /// `tool_use`, if any.
+    pub fn tool_use_template(&self) -> Option<&str> {
+        self.named("tool_use")
+    }
+
+    /// Look up a named template. Transformers collects the list into a dict,
+    /// so a later entry wins over an earlier one with the same name.
+    fn named(&self, name: &str) -> Option<&str> {
+        match self {
+            Self::Single(_) => None,
+            Self::Named(templates) => templates
+                .iter()
+                .rev()
+                .find(|template| template.name == name)
+                .map(|template| template.template.as_str()),
+        }
+    }
+}
+
 /// Hugging Face named special tokens may be serialized as a string or an
 /// object carrying the token content.
 #[derive(Debug, Clone, Deserialize)]
@@ -321,7 +368,39 @@ where
 
 #[cfg(test)]
 mod tests {
-    use super::{GenerationConfig, ModelConfig};
+    use super::{GenerationConfig, HfTokenizerConfig, ModelConfig};
+
+    #[test]
+    fn tokenizer_config_reads_single_and_named_chat_templates() {
+        let single: HfTokenizerConfig =
+            serde_json::from_str(r#"{"chat_template":"single"}"#).unwrap();
+        let named: HfTokenizerConfig = serde_json::from_str(
+            r#"{"chat_template":[
+                {"name":"default","template":"old default"},
+                {"name":"tool_use","template":"tool use"},
+                {"name":"default","template":"default"}
+            ]}"#,
+        )
+        .unwrap();
+        let without_default: HfTokenizerConfig =
+            serde_json::from_str(r#"{"chat_template":[{"name":"rag","template":"rag"}]}"#).unwrap();
+
+        let selected = [single, named, without_default].map(|config| {
+            let template = config.chat_template.unwrap();
+            (
+                template.default_template().map(str::to_owned),
+                template.tool_use_template().map(str::to_owned),
+            )
+        });
+        assert_eq!(
+            selected,
+            [
+                (Some("single".to_string()), None),
+                (Some("default".to_string()), Some("tool use".to_string())),
+                (None, None),
+            ]
+        );
+    }
 
     #[test]
     fn generation_config_normalizes_vllm_top_k_values() {
diff --git a/rust/src/text/src/backend/hf/mod.rs b/rust/src/text/src/backend/hf/mod.rs
index 77b47f62a31c..910fbf35ba1f 100644
--- a/rust/src/text/src/backend/hf/mod.rs
+++ b/rust/src/text/src/backend/hf/mod.rs
@@ -15,8 +15,8 @@ use vllm_tokenizer::{
 
 use self::config::{GenerationConfig, load_generation_config};
 pub use self::config::{
-    HfSpecialTokens, HfTokenizerConfig, ModelConfig, NamedSpecialToken, load_model_config,
-    load_tokenizer_config,
+    HfChatTemplate, HfNamedChatTemplate, HfSpecialTokens, HfTokenizerConfig, ModelConfig,
+    NamedSpecialToken, load_model_config, load_tokenizer_config,
 };
 pub use self::model_files::{ModelFileSources, ModelSource, ResolvedModelFiles, TokenizerSource};
 pub use self::overrides::HfOverrides;
diff --git a/rust/src/text/src/lower.rs b/rust/src/text/src/lower.rs
index 12b4cf254f39..5cbf58e7361b 100644
--- a/rust/src/text/src/lower.rs
+++ b/rust/src/text/src/lower.rs
@@ -1693,6 +1693,15 @@ mod tests {
         assert_eq!(result.unwrap(), 100);
     }
 
+    /// Shared lowering only clamps, even a budget above the whole context
+    /// window: `/inference/v1/generate` accepts it like Python, while chat and
+    /// completions reject it in their HTTP handlers.
+    #[test]
+    fn resolve_max_tokens_caps_budget_above_model_len() {
+        let result = resolve_max_tokens(Some(1000), None, 200, 100);
+        assert_eq!(result.unwrap(), 100);
+    }
+
     #[test]
     fn lower_text_request_preserves_non_streaming_request_metadata() {
         let mut request = sample_request();
diff --git a/rust/src/text/src/output/decoded.rs b/rust/src/text/src/output/decoded.rs
index 638e8556261e..e49a26060a4a 100644
--- a/rust/src/text/src/output/decoded.rs
+++ b/rust/src/text/src/output/decoded.rs
@@ -10,7 +10,10 @@ use tracing::{Level, debug, trace};
 use vllm_engine_core_client::AbortCause;
 use vllm_engine_core_client::protocol::output::{RequestSpecDecodeMetrics, StopReason};
 use vllm_engine_core_client::protocol::sampling_mask::SamplingMask;
-use vllm_llm::{FinishReason, GenerateOutput, PromptCacheUsage, RequestTimestamps, TokenUsage};
+use vllm_llm::{
+    FinishReason, GenerateOutput, PromptCacheUsage, RequestTimestamps, TokenUsage,
+    WeightVersionSpan,
+};
 use vllm_tokenizer::{DecodedText, DynTokenizer, IncrementalDecoder};
 
 use super::logprobs::{
@@ -79,6 +82,8 @@ pub struct Finished {
     pub sampling_mask: Option<SamplingMask>,
     /// Per-request speculative-decoding statistics emitted by the engine.
     pub spec_decode_metrics: Option<RequestSpecDecodeMetrics>,
+    /// Weight-version spans over all generated token positions.
+    pub weight_versions: Option<Vec<WeightVersionSpan>>,
     pub timestamps: RequestTimestamps,
 }
 
@@ -219,6 +224,7 @@ pub async fn decoded_text_event_stream(
         let kv_transfer_params = output.kv_transfer_params;
         let ec_transfer_params = output.ec_transfer_params;
         let spec_decode_metrics = output.spec_decode_metrics.map(|metrics| *metrics);
+        let mut weight_versions = output.weight_versions;
         let mut finish_reason = output.finish_reason;
         let mut stop_str_matched = false;
         let suppress_terminal_stop_token = finish_reason.as_ref().is_some_and(|r| r.is_stop())
@@ -309,6 +315,14 @@ pub async fn decoded_text_event_stream(
                 }
                 .into());
             }
+            // The spans cover every token received; cut them at the stop-string
+            // truncated count, like Python's spans over its truncated output.
+            if let Some(spans) = weight_versions.as_mut() {
+                spans.retain(|span| span.start < output_token_count);
+                if let Some(last) = spans.last_mut() {
+                    last.end = last.end.min(output_token_count);
+                }
+            }
             // Flush any remaining buffered text.
             let (last_chunk, full_decoded) = decoder.flush(truncate_output_to)?;
             let text_len = full_decoded.text.len();
@@ -358,6 +372,7 @@ pub async fn decoded_text_event_stream(
                     ec_transfer_params,
                     sampling_mask,
                     spec_decode_metrics,
+                    weight_versions,
                     timestamps: output.timestamps,
                 })),
             })
@@ -638,6 +653,7 @@ mod tests {
                                 first_token_ts: 10.0,
                                 last_token_ts: 11.5,
                             },
+                            weight_versions: None,
                         }),
                     ),
                 ],
@@ -671,6 +687,7 @@ mod tests {
                             ec_transfer_params: None,
                             sampling_mask: None,
                             spec_decode_metrics: None,
+                            weight_versions: None,
                             timestamps: Default::default(),
                         }),
                     ),
@@ -892,6 +909,39 @@ mod tests {
         assert_eq!(output.usage.output_token_count, 4);
     }
 
+    #[tokio::test]
+    async fn stream_stop_string_cuts_weight_version_spans_at_truncated_count() {
+        let span = |version: &str, start, end| WeightVersionSpan {
+            version: version.into(),
+            start,
+            end,
+        };
+        let mut first = GenerateOutput::for_test(Some(Arc::from([])), ascii_tokens("ab"), None);
+        first.weight_versions = Some(vec![span("v1", 0, 2)]);
+        // One merged output whose tokens span two more weight versions; the
+        // stop string ends inside the `v2` span.
+        let mut second = GenerateOutput::for_test(None, ascii_tokens("cxyz"), None);
+        second.weight_versions = Some(vec![span("v1", 0, 2), span("v2", 2, 5), span("v3", 5, 6)]);
+        let tokenizer: DynTokenizer = Arc::new(TestTokenizer::new());
+
+        let output = decoded_text_event_stream(
+            "test-weight-versions-stop".into(),
+            tokenizer,
+            stream::iter([Ok(first), Ok(second)]),
+            opts(&["cx"], 0),
+            false,
+        )
+        .collect_output()
+        .await
+        .unwrap();
+
+        assert_eq!(output.token_ids, ascii_tokens("abcx"));
+        assert_eq!(
+            output.weight_versions,
+            Some(vec![span("v1", 0, 2), span("v2", 2, 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 eebc2dd69623..4a812cd54033 100644
--- a/rust/src/text/src/output/mod.rs
+++ b/rust/src/text/src/output/mod.rs
@@ -19,6 +19,7 @@ use std::sync::Arc;
 use futures::{StreamExt as _, pin_mut};
 use vllm_engine_core_client::protocol::output::RequestSpecDecodeMetrics;
 use vllm_engine_core_client::protocol::sampling_mask::SamplingMask;
+use vllm_llm::WeightVersionSpan;
 
 use crate::{Error, FinishReason, Result, TextOutputStream};
 
@@ -42,6 +43,8 @@ pub struct CollectedTextOutput {
     pub sampling_mask: Option<SamplingMask>,
     /// Per-request speculative-decoding statistics emitted by the engine.
     pub spec_decode_metrics: Option<RequestSpecDecodeMetrics>,
+    /// Weight-version spans over generated token positions.
+    pub weight_versions: Option<Vec<WeightVersionSpan>>,
 }
 
 #[allow(clippy::manual_async_fn, reason = "specify `Send` bound")]
@@ -100,6 +103,7 @@ impl<T: TextOutputStream> T {
                                 ec_transfer_params: None,
                                 sampling_mask: None,
                                 spec_decode_metrics: None,
+                                weight_versions: None,
                             })
                         };
 
@@ -114,6 +118,7 @@ impl<T: TextOutputStream> T {
                                 finished.ec_transfer_params.map(|value| *value);
                             collected.sampling_mask = finished.sampling_mask;
                             collected.spec_decode_metrics = finished.spec_decode_metrics;
+                            collected.weight_versions = finished.weight_versions;
                             return Ok(collected);
                         }
                     }
@@ -190,6 +195,7 @@ mod tests {
                     ec_transfer_params: None,
                     sampling_mask: None,
                     spec_decode_metrics: None,
+                    weight_versions: None,
                     timestamps: Default::default(),
                 })),
             }),
@@ -315,6 +321,7 @@ mod tests {
                     ec_transfer_params: None,
                     sampling_mask: None,
                     spec_decode_metrics: None,
+                    weight_versions: None,
                     timestamps: Default::default(),
                 })),
             }),
diff --git a/setup.py b/setup.py
index 8859d0a93671..0a31c10f221d 100644
--- a/setup.py
+++ b/setup.py
@@ -1624,6 +1624,12 @@ def get_rocm_device_requirements() -> list[str]:
         "py.typed",
         "libs/*.so*",
         "model_executor/layers/fused_moe/configs/*.json",
+        # Tuned kernel configs must be explicit: the setuptools-scm git
+        # file-finder is inactive when building from a .git-less tree
+        # (e.g. tarball/sdist builds), which silently drops them from
+        # the wheel and forces fallback to default launch configs.
+        "model_executor/layers/mamba/ops/configs/selective_state_update/*.json",
+        "kernels/helion/configs/*/*.json",
         "model_executor/layers/quantization/utils/configs/*.json",
         "entrypoints/serve/instrumentator/static/*.js",
         "entrypoints/serve/instrumentator/static/*.css",
@@ -1716,7 +1722,7 @@ def add_vllm_package_data(filename: str) -> None:
         "tensorizer": ["tensorizer==2.10.1"],
         "fastsafetensors": ["fastsafetensors >= 0.3.3"],
         "instanttensor": ["instanttensor >= 0.2.2"],
-        "runai": ["runai-model-streamer[s3,gcs,azure] >= 0.15.7"],
+        "runai": ["runai-model-streamer[s3,gcs,azure] >= 0.17.1"],
         "audio": [
             "av",
             "scipy",
diff --git a/tests/basic_correctness/cpu_offload/test_cpu_offload.py b/tests/basic_correctness/cpu_offload/test_cpu_offload.py
index 46bdb7218bd2..39a162b1ee6f 100644
--- a/tests/basic_correctness/cpu_offload/test_cpu_offload.py
+++ b/tests/basic_correctness/cpu_offload/test_cpu_offload.py
@@ -66,7 +66,7 @@ def test_mrv2_weight_offloading(
         with vllm_runner(
             "hmellor/tiny-random-LlamaForCausalLM",
             enforce_eager=True,
-            gpu_memory_utilization=0.02,
+            gpu_memory_utilization=0.03,
             max_model_len=128,
             max_num_seqs=1,
             **offload_kwargs,
diff --git a/tests/basic_correctness/memory/cumem/test_cumem.py b/tests/basic_correctness/memory/cumem/test_cumem.py
index 956373f4d37e..4fb94f0c4e8f 100644
--- a/tests/basic_correctness/memory/cumem/test_cumem.py
+++ b/tests/basic_correctness/memory/cumem/test_cumem.py
@@ -10,10 +10,12 @@
 from vllm import LLM, SamplingParams
 from vllm.device_allocator import get_mem_allocator_instance
 from vllm.platforms import current_platform
+from vllm.utils.torch_utils import PIN_MEMORY
 
 from ....utils import create_new_process_for_each_test
 
 DEVICE_TYPE = current_platform.device_type
+OFFLOAD_TAGS = {1: {"weights", "runtime", "cudagraph"}, 2: {"runtime", "cudagraph"}}
 
 
 def mapped_usage(allocator) -> int:
@@ -31,6 +33,69 @@ def mapped_usage(allocator) -> int:
     )
 
 
+def host_cache_bytes() -> int:
+    if not current_platform.is_cuda_alike():
+        return 0
+    return torch.cuda.host_memory_stats()["allocated_bytes.allocated"]
+
+
+def cpu_backups(allocator, tags, pinned=PIN_MEMORY) -> dict[int, torch.Tensor]:
+    backups = {
+        ptr: data.cpu_backup_tensor
+        for ptr, data in allocator.pointer_to_data.items()
+        if data.cpu_backup_tensor is not None
+    }
+    assert backups.keys() == {
+        ptr for ptr, data in allocator.pointer_to_data.items() if data.tag in tags
+    }
+    for ptr, backup in backups.items():
+        assert backup.nbytes == allocator.pointer_to_data[ptr].handle[1]
+        assert backup.is_pinned() == pinned
+    return backups
+
+
+def data_ptrs(backups: dict[int, torch.Tensor]) -> dict[int, int]:
+    return {ptr: backup.data_ptr() for ptr, backup in backups.items()}
+
+
+def assert_released(backups: dict[int, torch.Tensor]) -> None:
+    assert backups
+    if current_platform.is_cuda_alike():
+        assert not any(backup.is_pinned() for backup in backups.values())
+
+
+def backup_state(worker) -> tuple[dict[int, tuple], int]:
+    allocator = get_mem_allocator_instance()
+    state = {}
+    for ptr, data in allocator.pointer_to_data.items():
+        backup = data.cpu_backup_tensor
+        state[ptr] = (data.tag, data.handle[1], None, 0, False)
+        if backup is not None:
+            state[ptr] = (
+                data.tag,
+                data.handle[1],
+                backup.data_ptr(),
+                backup.nbytes,
+                backup.is_pinned(),
+            )
+    return state, host_cache_bytes()
+
+
+def sleep_and_check_backups(llm, level: int) -> dict[int, tuple]:
+    host_before = llm.collective_rpc(backup_state)[0][1]
+    llm.sleep(level=level)
+    state, host_after = llm.collective_rpc(backup_state)[0]
+    assert host_after == host_before
+    for tag, size, backup_ptr, backup_bytes, pinned in state.values():
+        if tag in OFFLOAD_TAGS[level]:
+            assert backup_ptr is not None
+            assert (backup_bytes, pinned) == (size, PIN_MEMORY)
+        else:
+            assert backup_ptr is None
+    assert any(tag == "weights" for tag, *_ in state.values())
+    return {ptr: s for ptr, s in state.items() if s[2] is not None}
+
+
 def _wake_up_with_poisoned_mappings(allocator, byte_value: int = 0xA5) -> None:
     """Wake discarded allocations with deterministic nonzero contents."""
     original_create_and_map = cumem.create_and_map
@@ -72,8 +137,9 @@ def test_python_error():
         allocator.wake_up()
 
 
+@pytest.mark.parametrize("register_error", [0, 1], ids=["pinned", "register-fails"])
 @create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn")
-def test_basic_cumem(monkeypatch):
+def test_basic_cumem(register_error, monkeypatch):
     # some tensors from default memory pool
     shape = (1024, 1024)
     x = torch.empty(shape, device=DEVICE_TYPE)
@@ -106,15 +172,43 @@ def record_empty_cache():
         return original_empty_cache()
 
     monkeypatch.setattr(torch.accelerator, "empty_cache", record_empty_cache)
-    allocator.sleep()
-    assert len(empty_cache_calls) == 1
-    assert mapped_usage(allocator) == 0
-    allocator.wake_up()
-    assert mapped_usage(allocator) == mapped_bytes
+    if register_error:
+        monkeypatch.setattr(
+            cumem.libcudart, "cudaHostRegister", lambda *args: register_error
+        )
+    unregistered = []
+    original_unregister = cumem.libcudart.cudaHostUnregister
 
-    # they can be used together
-    output = x + y + z
-    assert torch.allclose(output, torch.ones_like(output) * 3)
+    def record_unregister(ptr):
+        unregistered.append(ptr)
+        return original_unregister(ptr)
+
+    monkeypatch.setattr(cumem.libcudart, "cudaHostUnregister", record_unregister)
+    warnings = []
+    monkeypatch.setattr(
+        cumem.logger, "warning_once", lambda msg, *args: warnings.append(msg % args)
+    )
+    pinned = PIN_MEMORY and not register_error
+    backup_ptrs = None
+    for _ in range(3):
+        host_before = host_cache_bytes()
+        empty_cache_calls.clear()
+        allocator.sleep()
+        assert len(empty_cache_calls) == 1
+        assert mapped_usage(allocator) == 0
+        assert host_cache_bytes() == host_before
+        backups = cpu_backups(allocator, {"default"}, pinned)
+        assert backup_ptrs in (None, data_ptrs(backups))
+        backup_ptrs = data_ptrs(backups)
+        allocator.wake_up()
+        assert mapped_usage(allocator) == mapped_bytes
+        assert data_ptrs(cpu_backups(allocator, {"default"}, pinned)) == backup_ptrs
+
+        # they can be used together
+        output = x + y + z
+        assert torch.allclose(output, torch.ones_like(output) * 3)
+        y.add_(1)
+        z.sub_(1)
 
     # Freeing a cuMem allocation synchronizes its device through the accelerator.
     sync_devices = []
@@ -126,7 +220,12 @@ def record_synchronize(device=None):
 
     monkeypatch.setattr(torch.accelerator, "synchronize", record_synchronize)
     with allocator.use_memory_pool("free-test"):
-        temp = torch.empty(4 << 20, dtype=torch.uint8, device=DEVICE_TYPE)
+        temp = torch.full((4 << 20,), 7, dtype=torch.uint8, device=DEVICE_TYPE)
+    allocator.sleep(offload_tags="free-test")
+    assert_released(backups)
+    temp_backups = cpu_backups(allocator, {"free-test"}, pinned)
+    allocator.wake_up()
+    assert torch.equal(temp, torch.full_like(temp, 7))
     device_index = torch.accelerator.current_device_index()
     del temp
     gc.collect()
@@ -135,6 +234,13 @@ def record_synchronize(device=None):
         pass
     assert device_index is not None
     assert device_index in sync_devices
+    assert_released(temp_backups)
+    assert not cpu_backups(allocator, set())
+    assert backup_ptrs is not None
+    released = [*backup_ptrs.values(), *data_ptrs(temp_backups).values()]
+    assert sorted(unregistered) == (sorted(released) if pinned else [])
+    assert bool(warnings) == (PIN_MEMORY and bool(register_error))
+    assert all("pageable" in warning for warning in warnings)
 
 
 @pytest.mark.parametrize("full_sleep", [0, 1, 2], ids=["kv-only", "sleep-1", "sleep-2"])
@@ -174,25 +280,33 @@ def get_mapped_bytes_by_tag(worker):
     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
+    backups = None
+    for _ in range(3):
+        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:
+            cycle_backups = sleep_and_check_backups(llm, full_sleep)
+            assert backups in (None, cycle_backups)
+            backups = cycle_backups
+            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")
+        if full_sleep:
+            assert backups is not None
+            state = llm.collective_rpc(backup_state)[0][0]
+            assert {ptr: state[ptr] for ptr in backups} == backups
+        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])
@@ -217,18 +331,26 @@ def test_sleep_with_only_weights_asleep(first_level, second_level, monkeypatch):
     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
+    backups_by_level: dict[int, dict[int, tuple]] = {}
+    for level, other_level in (
+        (first_level, second_level),
+        (second_level, first_level),
+    ):
+        backups = sleep_and_check_backups(llm, level)
+        assert backups_by_level.setdefault(level, backups) == backups
+        llm.wake_up(tags=["kv_cache"])
+        mapped_before = llm.collective_rpc(get_mapped_bytes)[0]
+        assert mapped_before > 0
+        llm.sleep(level=other_level)
+        assert llm.collective_rpc(get_mapped_bytes)[0] == mapped_before
+        state = llm.collective_rpc(backup_state)[0][0]
+        assert {ptr: s for ptr, s in state.items() if s[2] is not None} == backups
+        llm.wake_up(tags=["weights"])
+        if 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")
@@ -244,9 +366,21 @@ def test_discard_tags():
         kv = torch.ones(512, 512, device=DEVICE_TYPE)
 
     mapped_bytes = mapped_usage(allocator)
+    allocator.sleep(offload_tags=("weights", "kv_cache"))
+    allocator.wake_up()
+    assert torch.equal(kv, torch.ones_like(kv))
+    backups = cpu_backups(allocator, {"weights", "kv_cache"})
+    kv_backups = {
+        ptr: backup
+        for ptr, backup in backups.items()
+        if allocator.pointer_to_data[ptr].tag == "kv_cache"
+    }
 
     # Discard kv_cache only — weights should remain valid
     allocator.discard("kv_cache")
+    assert_released(kv_backups)
+    weight_ptrs = data_ptrs(cpu_backups(allocator, {"weights"}))
+    assert weight_ptrs.items() <= data_ptrs(backups).items()
 
     mapped_bytes_after_discard = mapped_usage(allocator)
     assert mapped_bytes - mapped_bytes_after_discard >= kv.nbytes
@@ -258,11 +392,14 @@ def test_discard_tags():
     # Wake up and verify kv_cache is remapped; discarded contents are undefined.
     allocator.wake_up()
     assert kv.shape == (512, 512)
+    assert data_ptrs(cpu_backups(allocator, {"weights"})) == weight_ptrs
 
     # Full sleep/wake cycle still works after discard
+    weights.fill_(2)
     allocator.sleep(offload_tags="weights")
+    assert data_ptrs(cpu_backups(allocator, {"weights"})) == weight_ptrs
     allocator.wake_up()
-    assert torch.allclose(weights, torch.ones_like(weights))
+    assert torch.equal(weights, torch.full_like(weights, 2))
 
 
 @pytest.mark.parametrize("tag", ["workspace", None], ids=["workspace", "default"])
@@ -354,6 +491,13 @@ def test_reentered_tag(free_x_early):
     allocator.sleep(offload_tags="t")
     allocator.wake_up()
     assert int(y.sum()) == y.numel()
+    backups = cpu_backups(allocator, {"t"})
+    del y
+    gc.collect()
+    with allocator.use_memory_pool("t"):
+        pass
+    assert not allocator.pointer_to_data
+    assert_released(backups)
     allocator.release_pools()
 
 
@@ -376,10 +520,14 @@ def test_runtime_state_survives_sleep(level):
     }
 
     backend = CuMemBackend()
+    host_before = host_cache_bytes()
     backend.suspend(level=level)
     assert mapped_usage(allocator) == 0
+    assert host_cache_bytes() == host_before
+    backups = cpu_backups(allocator, {"runtime"})
 
     backend.resume(tags=["weights"])
+    assert data_ptrs(cpu_backups(allocator, {"runtime"})) == data_ptrs(backups)
     assert torch.equal(state, torch.arange(1 << 20, device=DEVICE_TYPE))
     backend.resume(tags=["kv_cache"])
     kv.fill_(1)
@@ -399,7 +547,13 @@ def test_level2_discards_ordinary_tensor_with_weights_tag():
         )
 
     pointers = (fake_weight.data_ptr(), ordinary_tensor.data_ptr())
+    allocator.sleep(offload_tags="weights")
+    allocator.wake_up()
+    assert torch.all(fake_weight == 0x44)
+    backups = cpu_backups(allocator, {"weights"})
     allocator.sleep(offload_tags=())
+    assert_released(backups)
+    assert not cpu_backups(allocator, set())
     _wake_up_with_poisoned_mappings(allocator)
     torch.accelerator.synchronize()
 
diff --git a/tests/basic_correctness/memory/sleep_mode/test_process_checkpoint.py b/tests/basic_correctness/memory/sleep_mode/test_process_checkpoint.py
new file mode 100644
index 000000000000..b60394353bf2
--- /dev/null
+++ b/tests/basic_correctness/memory/sleep_mode/test_process_checkpoint.py
@@ -0,0 +1,133 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Opt-in Linux/NVIDIA driver integration test, with exclusive GPU access.
+
+Run with VLLM_TEST_CUDA_PROCESS_CHECKPOINT=1. The driver must expose CUDA
+process checkpoint APIs. Override VLLM_TEST_MODEL and VLLM_TEST_SLEEP_CYCLES
+for larger-model validation; this test requires enough RAM for weight backups.
+The selected model must support batch invariance. Hybrid GDN models require
+separate paired accuracy evaluation because that mode is unsupported for GDN.
+"""
+
+import os
+
+import pytest
+
+from vllm import LLM, SamplingParams
+
+
+def _graph_identity(worker):
+    import torch
+
+    if not hasattr(worker, "_checkpoint_test_replays"):
+        worker._checkpoint_test_replays = 0
+        worker._checkpoint_test_captures = 0
+        replay = torch.cuda.CUDAGraph.replay
+        capture = torch.cuda.CUDAGraph.capture_begin
+
+        def counted_replay(graph):
+            worker._checkpoint_test_replays += 1
+            return replay(graph)
+
+        def counted_capture(graph, *args, **kwargs):
+            worker._checkpoint_test_captures += 1
+            return capture(graph, *args, **kwargs)
+
+        torch.cuda.CUDAGraph.replay = counted_replay
+        torch.cuda.CUDAGraph.capture_begin = counted_capture
+    torch.accelerator.synchronize()
+    graphs = worker.model_runner.cudagraph_manager.graphs
+    return {
+        "pid": os.getpid(),
+        "graphs": sorted(g.raw_cuda_graph_exec() for g in graphs.values()),
+        "replays": worker._checkpoint_test_replays,
+        "captures": worker._checkpoint_test_captures,
+    }
+
+
+@pytest.mark.skipif(
+    os.environ.get("VLLM_TEST_CUDA_PROCESS_CHECKPOINT") != "1",
+    reason="requires exclusive GPUs and a checkpoint-capable NVIDIA driver",
+)
+@pytest.mark.parametrize(
+    "tp,transport", [(1, "p2p"), (2, "p2p"), (4, "p2p"), (2, "shm"), (4, "shm")]
+)
+def test_process_checkpoint_preserves_outputs_and_graphs(tp, transport, monkeypatch):
+    import torch
+
+    if torch.accelerator.device_count() < tp:
+        pytest.skip(f"requires {tp} GPUs")
+    config = {
+        "VLLM_USE_V2_MODEL_RUNNER": "1",
+        "VLLM_WORKER_MULTIPROC_METHOD": "spawn",
+        "VLLM_ALLOW_INSECURE_SERIALIZATION": "1",  # Test observation RPC only.
+        "VLLM_BATCH_INVARIANT": "1",
+        "NCCL_P2P_DISABLE": "1" if transport == "shm" else "0",
+        "NCCL_SHM_DISABLE": "0" if transport == "shm" else "1",
+        "NCCL_IB_DISABLE": "1",
+        "NCCL_CUMEM_ENABLE": "1",
+        "NCCL_CUMEM_HOST_ENABLE": "0",
+        "NCCL_NVLS_ENABLE": "0",
+        "VLLM_ALLREDUCE_USE_FLASHINFER": "0",
+        "VLLM_ALLREDUCE_USE_SYMM_MEM": "0",
+        "VLLM_ALLREDUCE_USE_FLASHINFER_PCIE_IPC": "0",
+        "VLLM_USE_NCCL_SYMM_MEM": "0",
+    }
+    for key, value in config.items():
+        monkeypatch.setenv(key, value)
+    llm = LLM(
+        os.environ.get("VLLM_TEST_MODEL", "Qwen/Qwen3-0.6B"),
+        dtype="bfloat16",
+        tensor_parallel_size=tp,
+        distributed_executor_backend=(
+            "vllm.v1.executor.cuda_checkpoint_executor.CUDACheckpointExecutor"
+        ),
+        disable_custom_all_reduce=tp > 1,
+        enable_nccl_comm_suspend=tp > 1,
+        enable_sleep_mode=True,
+        enable_prefix_caching=False,
+        kv_cache_memory_bytes=1024**3,
+        max_model_len=256,
+        max_num_seqs=4,
+        compilation_config={
+            "cudagraph_capture_sizes": [1, 2, 4],
+            "pass_config": {"fuse_allreduce_rms": False},
+        },
+    )
+    try:
+        params = SamplingParams(temperature=0, max_tokens=32, ignore_eos=True)
+
+        def generate():
+            outputs = llm.generate(
+                ["The capital of France is", "Compute 7 times 8:"],
+                params,
+                use_tqdm=False,
+            )
+            return [output.outputs[0].token_ids for output in outputs]
+
+        reference = generate()
+        graphs = llm.collective_rpc(_graph_identity)
+        assert all(rank["graphs"] for rank in graphs)
+        for cycle in range(int(os.environ.get("VLLM_TEST_SLEEP_CYCLES", "3"))):
+            assert generate() == reference
+            llm.sleep(level=1)
+            llm.sleep(level=1)
+            order = (("weights", "kv_cache"), ("kv_cache", "weights"), (None,))[
+                cycle % 3
+            ]
+            for tag in order:
+                llm.wake_up(tags=[tag] if tag else None)
+            awake = llm.collective_rpc(_graph_identity)
+            assert generate() == reference
+            after = llm.collective_rpc(_graph_identity)
+            for original, before_generate, current in zip(
+                graphs, awake, after, strict=True
+            ):
+                assert current["pid"] == original["pid"]
+                assert current["graphs"] == original["graphs"]
+                assert current["replays"] > before_generate["replays"]
+                assert current["captures"] == 0
+            graphs = after
+        llm.sleep(level=1)
+    finally:
+        llm.llm_engine.engine_core.shutdown()
diff --git a/tests/benchmarks/test_custom_dataset_seed.py b/tests/benchmarks/test_custom_dataset_seed.py
index d23ce40b53e7..e38b1451655c 100644
--- a/tests/benchmarks/test_custom_dataset_seed.py
+++ b/tests/benchmarks/test_custom_dataset_seed.py
@@ -5,9 +5,11 @@
 from pathlib import Path
 
 import pytest
-from transformers import AutoTokenizer, PreTrainedTokenizerBase
+from tokenizers import Tokenizer
+from tokenizers.models import WordLevel
+from transformers import AutoTokenizer, PreTrainedTokenizerBase, TokenizersBackend
 
-from vllm.benchmarks.datasets import get_samples
+from vllm.benchmarks.datasets import CustomDataset, get_samples
 
 
 @pytest.fixture(scope="session")
@@ -75,3 +77,22 @@ def test_custom_dataset_same_seed_is_deterministic(
     prompts_b = [s.prompt for s in samples_b]
 
     assert prompts_a == prompts_b
+
+
+@pytest.mark.benchmark
+def test_custom_dataset_keeps_numeric_string_prompts(tmp_path: Path) -> None:
+    """All-numeric string prompts must not be inferred as numbers by pandas."""
+    prompts = ["00123", "00456"]
+    jsonl = tmp_path / "data.jsonl"
+    jsonl.write_text(
+        "".join(json.dumps({"prompt": p, "output_tokens": 4}) + "\n" for p in prompts)
+    )
+    tokenizer = TokenizersBackend(
+        tokenizer_object=Tokenizer(WordLevel({"[UNK]": 0}, unk_token="[UNK]")),
+        unk_token="[UNK]",
+    )
+
+    dataset = CustomDataset(dataset_path=str(jsonl), disable_shuffle=True)
+    requests = dataset.sample(tokenizer, num_requests=2, skip_chat_template=True)
+
+    assert [r.prompt for r in requests] == prompts
diff --git a/tests/benchmarks/test_forward_cli.py b/tests/benchmarks/test_forward_cli.py
index 9b66086b76e0..9f1012414870 100644
--- a/tests/benchmarks/test_forward_cli.py
+++ b/tests/benchmarks/test_forward_cli.py
@@ -16,6 +16,7 @@
 from vllm.benchmarks.forward.kv_init import touched_kv_views, zero_kv_blocks
 from vllm.benchmarks.forward.main import run_case
 from vllm.benchmarks.forward.profile import summarize
+from vllm.benchmarks.forward.routing import RoutingFixture, zipf_topk_ids
 from vllm.benchmarks.forward.shapes import jitter_kv_lens, parse_case
 from vllm.platforms import current_platform
 from vllm.v1.attention.backends.mla.cpu_mla import CPUMLABackend
@@ -27,11 +28,146 @@
 MODEL_NAME = "meta-llama/Llama-3.2-1B-Instruct"
 
 
+@pytest.mark.skipif(
+    not current_platform.is_device_capability_family(100),
+    reason="Native Muse TRTLLM validation requires Blackwell.",
+)
+def test_bench_forward_muse_balanced_routing(tmp_path):
+    """Muse's native router supports finite, reproducible balanced graph replay."""
+    config = dict(
+        architectures=["MuseSpark1xForCausalLM"],
+        model_type="muse_spark_text",
+        vocab_size=128,
+        hidden_size=512,
+        intermediate_size=512,
+        moe_hidden_size=512,
+        moe_intermediate_size=512,
+        num_hidden_layers=1,
+        num_attention_heads=8,
+        num_key_value_heads=8,
+        head_dim=64,
+        num_local_experts=8,
+        num_experts_per_tok=2,
+        num_shared_experts=0,
+        max_position_embeddings=64,
+        hidden_act="silu",
+        attention_bias=False,
+        layer_types=["full_attention"],
+        mlp_layer_types=["sparse"],
+        layer_rope_theta=[10000.0],
+        rope_parameters={"rope_type": "default", "rope_theta": 10000.0},
+        rms_norm_eps=1e-5,
+        post_norm_eps=1e-8,
+        residual_gate_temperature=0.3,
+        final_logit_softcapping=None,
+        output_multiplier=0.5,
+        routed_scaling_factor=1.0,
+        qk_scale_factor=1.0,
+        tie_word_embeddings=False,
+        torch_dtype="bfloat16",
+    )
+    (tmp_path / "config.json").write_text(json.dumps(config))
+    entrypoint = (
+        "from vllm.model_executor.models import ModelRegistry; "
+        "ModelRegistry.register_model('MuseSpark1xForCausalLM', "
+        "'vllm.models.muse_spark.nvidia.model:MuseSpark1xForCausalLM'); "
+        "from vllm.entrypoints.cli.main import main; main()"
+    )
+    output = tmp_path / "results.json"
+    command = [
+        sys.executable,
+        "-c",
+        entrypoint,
+        "bench",
+        "forward",
+        "--model",
+        str(tmp_path),
+        "--skip-tokenizer-init",
+        "--moe-backend",
+        "flashinfer_trtllm",
+        "--routing",
+        "balanced",
+        "--max-model-len",
+        "64",
+        "--max-num-seqs",
+        "4",
+        "--max-num-batched-tokens",
+        "64",
+        "--kv-cache-memory-bytes",
+        "16777216",
+        "--kv-len-jitter",
+        "0",
+        "--kernel-config.enable_jit_warmup",
+        "false",
+        "--kernel-config.enable_cutedsl_warmup",
+        "false",
+        "--disable-flashinfer-autotune",
+        "--case",
+        "prefill=32",
+        "--case",
+        "decode=2*1@32",
+        "--iters",
+        "2",
+        "--samples",
+        "2",
+        "--output-json",
+        str(output),
+    ]
+    result = subprocess.run(command, capture_output=True, text=True, timeout=600)
+    assert result.returncode == 0, result.stdout + result.stderr
+    rows = json.loads(output.read_text())
+    assert len(rows) == 2
+    for row in rows:
+        (rank,) = row["ranks"]
+        assert rank["finite"] and rank["graph_deterministic"]
+        assert rank["graph_vs_eager"]["relative_rms"] == 0
+
+
+def test_balanced_routing_restores_expert_bias(dist_init):
+    """Fixture routes stay balanced despite selection bias, then restore it."""
+    from vllm.config import get_current_vllm_config
+    from vllm.forward_context import set_forward_context
+    from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear
+    from vllm.models.muse_spark.common.moe import MuseSpark1xRouter
+
+    device = torch.device(current_platform.device_type)
+    layer = torch.nn.Module()
+    gate = GateLinear(4, 8, params_dtype=torch.float32).to(device)
+    layer.add_module("router", gate)
+    layer.top_k = 2
+    bias = torch.tensor([100.0, 100.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], device=device)
+    layer.register_buffer("expert_bias", bias.clone())
+    selector = MuseSpark1xRouter(
+        top_k=2,
+        num_experts=8,
+        expert_bias=layer.expert_bias,
+        route_scale=1.0,
+    )
+    hidden = torch.zeros(64, 4, device=device)
+    with torch.no_grad():
+        gate.weight.zero_()
+    fixture = RoutingFixture(layer, 64, 0, 1, 4, 0, device, mode="balanced")
+    try:
+        with set_forward_context(None, get_current_vllm_config(), num_tokens=64):
+            logits, _ = gate(hidden)
+            _, ids = selector.select_experts(hidden, logits)
+            fixture.check_applied()
+            counts = torch.bincount((ids // 2).flatten(), minlength=4)
+            assert counts.min() == counts.max()
+    finally:
+        fixture.remove()
+    torch.testing.assert_close(layer.expert_bias, bias, rtol=0, atol=0)
+    logits, _ = gate(hidden)
+    _, ids = selector.select_experts(hidden, logits)
+    assert set(ids.flatten().tolist()) == {0, 1}
+
+
 @pytest.mark.benchmark
 @pytest.mark.skipif(
     not current_platform.is_cuda(), reason="Forward graph timing requires CUDA."
 )
-def test_bench_forward_hybrid_verify(tmp_path):
+@pytest.mark.parametrize("routing", ["auto", "balanced"])
+def test_bench_forward_hybrid_verify(tmp_path, routing):
     """Repeated target verification starts from the same recurrent state."""
     config = dict(
         architectures=["Qwen3NextForCausalLM"],
@@ -72,6 +208,8 @@ def test_bench_forward_hybrid_verify(tmp_path):
             "--model",
             str(tmp_path),
             "--skip-tokenizer-init",
+            "--routing",
+            routing,
             "--exclude-drafter",
             "--speculative-config",
             '{"method":"ngram","num_speculative_tokens":5}',
@@ -107,11 +245,27 @@ def test_bench_forward_hybrid_verify(tmp_path):
     assert [row["case"] for row in rows] == ["verify", "mixed"]
     for row in rows:
         (rank,) = row["ranks"]
+        assert rank["routing"] == ("zipf" if routing == "auto" else routing)
         assert rank["finite"] and rank["graph_deterministic"]
         assert rank["timing_scope"] == "target_only"
         assert row["median_us"] > 0
 
 
+def test_zipf_expert_selection():
+    """Popularity follows Zipf, with distinct experts and reproducible rows."""
+    rng = np.random.default_rng(0)
+    ids = zipf_topk_ids(50000, 32, 1, rng)
+    frequencies = np.sort(np.bincount(ids.ravel(), minlength=32))[::-1] / len(ids)
+    expected = np.arange(1, 33, dtype=float) ** -1.25
+    expected /= expected.sum()
+    np.testing.assert_allclose(frequencies, expected, atol=0.005, rtol=0)
+
+    ids = zipf_topk_ids(2048, 384, 6, np.random.default_rng(0))
+    assert ((ids >= 0) & (ids < 384)).all()
+    assert (np.diff(np.sort(ids, axis=1), axis=1) > 0).all()
+    assert np.array_equal(ids[:6], zipf_topk_ids(6, 384, 6, np.random.default_rng(0)))
+
+
 @pytest.mark.benchmark
 @pytest.mark.skipif(
     not current_platform.is_cuda(), reason="Forward graph timing requires CUDA."
diff --git a/tests/benchmarks/test_responses_request_func.py b/tests/benchmarks/test_responses_request_func.py
index e7d3f60c66f5..93ea754c2851 100644
--- a/tests/benchmarks/test_responses_request_func.py
+++ b/tests/benchmarks/test_responses_request_func.py
@@ -277,6 +277,82 @@ def test_usage_is_taken_from_the_terminal_event():
     assert output.output_tokens == 4
 
 
+def tool_delta(delta: str, event_type: str = "response.function_call_arguments.delta"):
+    return sse(event_type, {"output_index": 0, "item_id": "fc_1", "delta": delta})
+
+
+def test_tool_call_only_response_is_timed_and_counted():
+    """A forced tool call streams only argument deltas, no output text.
+
+    These are generated tokens, as `tool_calls` deltas are on `openai-chat`.
+    Counting only text deltas would fail every such request.
+    """
+    output = run(
+        [
+            sse("response.created", {"response": {"status": "in_progress"}}),
+            sse("response.output_item.added", {"item": {"type": "function_call"}}),
+            GAP_PREFILL,
+            tool_delta('{"city": '),
+            GAP_TOKEN,
+            tool_delta('"Paris"'),
+            tool_delta("}"),
+            sse("response.function_call_arguments.done", {"arguments": "{}"}),
+            sse("response.output_item.done", {"output_index": 0}),
+            GAP_FINALIZE,
+            usage_event(20, 6),
+        ]
+    )
+
+    assert output.success, output.error
+    # Tool arguments are not output text, as on `openai-chat`.
+    assert output.generated_text == ""
+    assert output.output_tokens == 6
+    assert GAP_PREFILL * 0.8 <= output.ttft < GAP_PREFILL + GAP_TOKEN * 0.5
+    assert len(output.itl) == 2
+    assert output.itl[0] >= GAP_TOKEN * 0.5
+    assert output.latency - output.ttft == pytest.approx(sum(output.itl), abs=1e-9)
+
+
+@pytest.mark.parametrize(
+    "event_type",
+    [
+        "response.function_call_arguments.delta",
+        "response.custom_tool_call_input.delta",
+        "response.mcp_call_arguments.delta",
+        "response.code_interpreter_call_code.delta",
+    ],
+)
+def test_every_tool_delta_event_starts_the_clock(event_type: str):
+    """Each tool delta event the server can emit carries generated tokens."""
+    output = run(
+        [
+            PREAMBLE,
+            tool_delta("a", event_type),
+            tool_delta("b", event_type),
+            usage_event(5, 2),
+        ]
+    )
+
+    assert output.success, output.error
+    assert len(output.itl) == 1
+
+
+def test_text_then_tool_call_times_both():
+    """Text before a tool call starts the clock; argument deltas add ITL."""
+    output = run(
+        [
+            PREAMBLE,
+            text_delta("Checking."),
+            tool_delta('{"city": "Paris"}'),
+            usage_event(5, 9),
+        ]
+    )
+
+    assert output.success, output.error
+    assert output.generated_text == "Checking."
+    assert len(output.itl) == 1
+
+
 @pytest.mark.parametrize("event_type", ["response.completed", "response.incomplete"])
 def test_incomplete_response_is_a_successful_request(event_type: str):
     """Hitting max_output_tokens is the Responses form of finish_reason=length.
diff --git a/tests/benchmarks/test_throughput_cli.py b/tests/benchmarks/test_throughput_cli.py
index 37ced49838c3..c994d9ed0bf7 100644
--- a/tests/benchmarks/test_throughput_cli.py
+++ b/tests/benchmarks/test_throughput_cli.py
@@ -9,6 +9,7 @@
 """
 
 import subprocess
+import warnings
 from argparse import Namespace
 from types import SimpleNamespace
 
@@ -22,6 +23,7 @@
     add_cli_args,
     assign_loras,
     get_requests,
+    validate_args,
 )
 from vllm.utils.argparse_utils import FlexibleArgumentParser
 
@@ -322,6 +324,68 @@ def test_to_serve_args_keeps_legacy_flag_fallback() -> None:
     assert not serve_args.enable_multimodal_chat
 
 
+def _parse_throughput_args(*argv: str) -> Namespace:
+    parser = FlexibleArgumentParser()
+    add_cli_args(parser)
+    return parser.parse_args(["--model", MODEL_NAME, *argv])
+
+
+@pytest.mark.parametrize(
+    ("argv", "expected"),
+    [
+        # Neither flag: the RandomDataset defaults.
+        ([], (1024, 128, 0)),
+        # Legacy flags alone are honored.
+        (
+            ["--input-len", "2048", "--output-len", "512", "--prefix-len", "64"],
+            (2048, 512, 64),
+        ),
+        # Explicit --random-* wins, including an explicit 0.
+        (
+            [
+                "--input-len",
+                "2048",
+                "--random-input-len",
+                "300",
+                "--prefix-len",
+                "64",
+                "--random-prefix-len",
+                "0",
+            ],
+            (300, 128, 0),
+        ),
+    ],
+)
+def test_parsed_lengths_follow_flag_precedence(argv, expected) -> None:
+    """Lengths from the real parser: --random-* > legacy flags > defaults."""
+    serve_args = _to_serve_args(_parse_throughput_args(*argv))
+
+    assert (
+        serve_args.random_input_len,
+        serve_args.random_output_len,
+        serve_args.random_prefix_len,
+    ) == expected
+
+
+@pytest.mark.parametrize(
+    "argv",
+    [
+        ["--input-len", "32", "--output-len", "1", "--prefix-len", "8"],
+        ["--random-input-len", "32", "--random-output-len", "1"],
+        ["--random-prefix-len", "64"],
+        ["--random-prefix-len", "0"],
+    ],
+)
+def test_one_flag_per_length_does_not_warn(argv) -> None:
+    """Only the legacy or only the --random-* flag: no "Both ... are specified"."""
+    args = _parse_throughput_args(*argv)
+    with warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+        validate_args(args)
+
+    assert not [str(w.message) for w in caught if "are specified" in str(w.message)]
+
+
 @pytest.mark.benchmark
 def test_get_requests_random_mm_with_mm_processor_args(
     hf_tokenizer: PreTrainedTokenizerBase,
diff --git a/tests/compile/backend.py b/tests/compile/backend.py
index 167de6ad68a3..641f2318d91f 100644
--- a/tests/compile/backend.py
+++ b/tests/compile/backend.py
@@ -6,11 +6,11 @@
 from contextlib import nullcontext
 from copy import deepcopy
 
-import depyf
 from torch import fx
 from torch._ops import OpOverload, OpOverloadPacket
 from torch.fx._utils import lazy_format_graph_code
 
+from vllm.compilation import depyf_compat
 from vllm.compilation.passes.fx_utils import find_op_nodes
 from vllm.compilation.passes.inductor_pass import (
     InductorPass,
@@ -71,7 +71,7 @@ def __init__(self, *passes: InductorPass | Callable[[fx.Graph], None]):
 
         if debug_dump_path := vllm_config.compile_debug_dump_path():
             logger.debug("Dumping depyf output to %s", debug_dump_path)
-            self.debug_ctx = depyf.prepare_debug(debug_dump_path.as_posix())
+            self.debug_ctx = depyf_compat.prepare_debug(debug_dump_path.as_posix())
         else:
             self.debug_ctx = nullcontext()
 
diff --git a/tests/compile/fullgraph/test_full_graph.py b/tests/compile/fullgraph/test_full_graph.py
index fceed6346d1e..61953e8dbe7a 100644
--- a/tests/compile/fullgraph/test_full_graph.py
+++ b/tests/compile/fullgraph/test_full_graph.py
@@ -59,7 +59,7 @@ def models_list(*, all: bool = True, keywords: list[str] | None = None):
     [CompilationMode.DYNAMO_TRACE_ONCE, CompilationMode.VLLM_COMPILE],
 )
 @pytest.mark.parametrize("model, model_kwargs", models_list(all=True))
-@create_new_process_for_each_test()
+@create_new_process_for_each_test("spawn")
 def test_full_graph(
     monkeypatch: pytest.MonkeyPatch,
     model: str,
@@ -97,17 +97,6 @@ def test_full_graph(
         )
         for model_info in models_list(keywords=["FP8-dynamic"])
     ]
-    + [
-        # Test depyf integration works
-        (
-            CompilationConfig(
-                mode=CompilationMode.VLLM_COMPILE,
-                debug_dump_path=Path(tempfile.gettempdir()),
-            ),
-            "facebook/opt-125m",
-            {},
-        ),
-    ]
     + [
         # graph inductor partition
         (
@@ -140,7 +129,7 @@ def test_full_graph(
     ],
 )
 # only test some of the models
-@create_new_process_for_each_test()
+@create_new_process_for_each_test("spawn")
 def test_custom_compile_config(
     compilation_config: CompilationConfig,
     model: str,
@@ -155,6 +144,23 @@ def test_custom_compile_config(
     run_model(compilation_config, model, **model_kwargs)
 
 
+@create_new_process_for_each_test()
+def test_depyf_integration(disable_vllm_compile_cache):
+    """Compile with depyf's debugging hooks installed.
+
+    ``disable_vllm_compile_cache`` is what makes this a test rather than a
+    formality: depyf only ever sees a compilation that actually runs, so served
+    from a warm cache this exercises none of its hooks.
+    """
+    run_model(
+        CompilationConfig(
+            mode=CompilationMode.VLLM_COMPILE,
+            debug_dump_path=Path(tempfile.gettempdir()),
+        ),
+        "facebook/opt-125m",
+    )
+
+
 def run_model(
     compile_config: CompilationMode | CompilationConfig, model: str, **model_kwargs
 ):
@@ -184,7 +190,10 @@ def run_model(
     # VllmRunner guards this; building LLM directly does not. No-op off ROCm.
     wait_for_memory_to_settle(
         threshold_ratio=1.0
-        - model_kwargs.get("gpu_memory_utilization", CacheConfig.gpu_memory_utilization)
+        - (
+            model_kwargs.get("gpu_memory_utilization")
+            or CacheConfig.DEFAULT_GPU_MEMORY_UTILIZATION
+        )
     )
 
     llm = LLM(
diff --git a/tests/compile/fullgraph/test_multimodal_compile.py b/tests/compile/fullgraph/test_multimodal_compile.py
index e10f5379f13a..54b8d6daa1ea 100644
--- a/tests/compile/fullgraph/test_multimodal_compile.py
+++ b/tests/compile/fullgraph/test_multimodal_compile.py
@@ -5,6 +5,7 @@
 import pytest
 
 from tests.models.utils import dummy_hf_overrides
+from tests.utils import create_new_process_for_each_test
 from vllm.compilation.counter import compilation_counter
 from vllm.config import VllmConfig
 from vllm.config.compilation import CompilationMode, CUDAGraphMode
@@ -17,7 +18,6 @@ def test_compile():
     assert not vllm_config.compilation_config.compile_mm_encoder
 
 
-@pytest.mark.forked
 @pytest.mark.core_model
 @pytest.mark.skipif(not current_platform.is_cuda(), reason="Skip if not cuda")
 @pytest.mark.parametrize(
@@ -50,6 +50,7 @@ def test_compile():
         pytest.param(True, 3, id="enabled"),
     ],
 )
+@create_new_process_for_each_test("spawn")
 def test_pixtral_compilation(
     vllm_runner,
     monkeypatch,
@@ -88,9 +89,8 @@ def test_pixtral_compilation(
         pass
 
 
-# forked needed to workaround https://github.com/vllm-project/vllm/issues/21073
-@pytest.mark.forked
 @pytest.mark.skipif(not current_platform.is_cuda(), reason="Skip if not cuda")
+@create_new_process_for_each_test("spawn")
 def test_qwen2_5_vl_compilation(vllm_runner, monkeypatch):
     """Test that Qwen2.5-VL vision submodules are compiled.
 
@@ -122,9 +122,8 @@ def test_qwen2_5_vl_compilation(vllm_runner, monkeypatch):
         pass
 
 
-# forked needed to workaround https://github.com/vllm-project/vllm/issues/21073
-@pytest.mark.forked
 @pytest.mark.skipif(not current_platform.is_cuda(), reason="Skip if not cuda")
+@create_new_process_for_each_test("spawn")
 def test_qwen2_5_vl_no_vit_compilation(vllm_runner, monkeypatch):
     """Test that Qwen2.5-VL vision submodules are not compiled when the
     config is passed off
@@ -147,8 +146,6 @@ def test_qwen2_5_vl_no_vit_compilation(vllm_runner, monkeypatch):
         pass
 
 
-# forked needed to workaround https://github.com/vllm-project/vllm/issues/21073
-@pytest.mark.forked
 @pytest.mark.skipif(not current_platform.is_cuda(), reason="Skip if not cuda")
 @pytest.mark.parametrize(
     ("compile_mm_encoder", "expected_models"),
@@ -159,6 +156,7 @@ def test_qwen2_5_vl_no_vit_compilation(vllm_runner, monkeypatch):
         pytest.param(True, 29, id="enabled"),
     ],
 )
+@create_new_process_for_each_test("spawn")
 def test_paddleocr_vl_compilation(
     vllm_runner, monkeypatch, compile_mm_encoder: bool, expected_models: int
 ):
@@ -183,10 +181,9 @@ def test_paddleocr_vl_compilation(
         pass
 
 
-# forked needed to workaround https://github.com/vllm-project/vllm/issues/21073
 # Requires Cuda and 8 gpus as well
-@pytest.mark.forked
 @pytest.mark.skip(reason="Skipping due to CI resource constraints")
+@create_new_process_for_each_test("spawn")
 def test_mllama4_vit_compilation(vllm_runner, monkeypatch):
     """Test that Mllama4 vision submodules are compiled.
 
diff --git a/tests/compile/fusions_e2e/test_tp1_quant.py b/tests/compile/fusions_e2e/test_tp1_quant.py
index 6e675b3b660e..12cd841a09fe 100644
--- a/tests/compile/fusions_e2e/test_tp1_quant.py
+++ b/tests/compile/fusions_e2e/test_tp1_quant.py
@@ -133,6 +133,13 @@ def test_tp1_fp8_fusions(
         # enabling AITER.
         matches_check.remove("norm_rope_fusion")
 
+    if current_platform.is_xpu():
+        # XPU defaults to the weight-only W8A16 FP8 path, which leaves no
+        # activation quant op for rms+quant fusion; opt into W8A8 instead.
+        model_kwargs["linear_backend"] = "xpu"
+        # Attention + quant fusion is not supported on XPU.
+        matches = matches._replace(attn_quant_fusion=0)
+
     run_e2e_fusion_test(
         model_name,
         matches,
diff --git a/tests/compile/passes/test_mla_attn_quant_fusion.py b/tests/compile/passes/test_mla_attn_quant_fusion.py
index 05bfda483a98..6dcb44b79c42 100644
--- a/tests/compile/passes/test_mla_attn_quant_fusion.py
+++ b/tests/compile/passes/test_mla_attn_quant_fusion.py
@@ -114,7 +114,7 @@ def __init__(
         self.mla_attn._v_scale = self.mla_attn._v_scale.to(device)
 
         # Initialize W_UK_T and W_UV from kv_b_proj weights
-        self.mla_attn.process_weights_after_loading(torch.get_default_dtype())
+        self.mla_attn.process_weights_after_loading()
         self.kv_b_proj_weight = kv_b_proj.weight.data.clone()
 
         self.block_size = 16
diff --git a/tests/compile/test_aot_compile.py b/tests/compile/test_aot_compile.py
index bbb9edb3710b..78d63f997e8e 100644
--- a/tests/compile/test_aot_compile.py
+++ b/tests/compile/test_aot_compile.py
@@ -196,6 +196,32 @@ def test_save_and_load(monkeypatch: pytest.MonkeyPatch):
             assert torch.allclose(ret, expected)
 
 
+@pytest.mark.skipif(
+    current_platform.device_count() < 2 or not is_torch_equal_or_newer("2.14.0.dev"),
+    reason="requires 2 GPUs and torch 2.14",
+)
+def test_load_on_another_device(monkeypatch: pytest.MonkeyPatch, vllm_tmp_cache: Path):
+    # With --device-ids, a rank can load a cache compiled on another device.
+    monkeypatch.setenv("VLLM_USE_AOT_COMPILE", "1")
+    monkeypatch.setenv("VLLM_USE_MEGA_AOT_ARTIFACT", "1")
+    monkeypatch.setenv("VLLM_USE_STANDALONE_COMPILE", "1")
+    monkeypatch.setenv("VLLM_COMPILE_ON_ONE_RANK", "1")
+    for device_index, loaded in ((0, 0), (1, 1)):
+        monkeypatch.setenv("VLLM_FORCE_AOT_LOAD", str(loaded))
+        disable_envs_cache()
+        x = torch.randn(10, 10, device=f"cuda:{device_index}")
+        vllm_config = make_vllm_config()
+        with (
+            torch.accelerator.device_index(device_index),
+            use_vllm_config(vllm_config),
+            compilation_counter.expect(
+                num_aot_compiles=1 - loaded, num_aot_artifacts_loaded=loaded
+            ),
+        ):
+            out = CompiledMod(vllm_config=vllm_config)(x)
+        torch.testing.assert_close(out, reference_fn(x))
+
+
 @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",
diff --git a/tests/compile/test_depyf_compat.py b/tests/compile/test_depyf_compat.py
new file mode 100644
index 000000000000..6047144a1712
--- /dev/null
+++ b/tests/compile/test_depyf_compat.py
@@ -0,0 +1,84 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for the depyf ``load_by_key_path`` patch in compilation/depyf_compat.py.
+
+depyf replaces ``torch._inductor.codecache.PyCodeCache.load_by_key_path`` with
+a hook that dumps each generated kernel before loading it. That hook's
+signature is frozen at ``(key, path, linemap, attrs)`` while Inductor also
+passes ``set_sys_modules``, so the patch re-wraps it to forward keywords depyf
+does not know about.
+
+End-to-end coverage lives in
+``tests/compile/fullgraph/test_full_graph.py::test_depyf_integration``, which
+needs a GPU and a cold compile cache. These tests pin the same behaviour by
+driving ``PyCodeCache`` directly, so they catch the incompatibility in seconds
+on CPU -- which is where it should have been caught in the first place.
+"""
+
+import sys
+
+import depyf.explain.enable_debugging as enable_debugging
+import torch._inductor.codecache as codecache
+
+from vllm.compilation import depyf_compat
+
+# Captured at import, which pytest does before any test can install the patch.
+PRISTINE_HOOK = enable_debugging.patched_load_by_key_path
+
+
+class TestPatchedHook:
+    """The hook has to survive the call *and* keep both sides' behaviour."""
+
+    def test_load_by_key_path_accepts_set_sys_modules(self, tmp_path):
+        with depyf_compat.prepare_debug(str(tmp_path)):
+            key, path = codecache.PyCodeCache.write("value = 1\n")
+            module = codecache.PyCodeCache.load_by_key_path(
+                key, path, set_sys_modules=False
+            )
+
+        assert module.value == 1
+
+    def test_set_sys_modules_is_forwarded_not_swallowed(self, tmp_path):
+        # Dropping the keyword would also make the call succeed, but it would
+        # re-register benchmark modules in sys.modules and bring back the
+        # exhaustion that pytorch#184285 fixed.
+        with depyf_compat.prepare_debug(str(tmp_path)):
+            key, path = codecache.PyCodeCache.write("value = 2\n")
+            module = codecache.PyCodeCache.load_by_key_path(
+                key, path, set_sys_modules=False
+            )
+            registered = module.__name__ in sys.modules
+
+        sys.modules.pop(module.__name__, None)
+        assert not registered
+
+    def test_depyf_still_dumps_the_generated_kernel(self, tmp_path):
+        # The patch reimplements depyf's hook body rather than wrapping it,
+        # so this is what catches us drifting away from what depyf does.
+        with depyf_compat.prepare_debug(str(tmp_path)):
+            key, path = codecache.PyCodeCache.write("value = 3\n")
+            codecache.PyCodeCache.load_by_key_path(key, path, set_sys_modules=False)
+
+        dumped = list(tmp_path.rglob("*.kernel_*.py"))
+        assert dumped, f"depyf wrote no kernel dump into {tmp_path}"
+        assert "value = 3" in dumped[0].read_text()
+
+
+class TestPatchApplication:
+    def test_patch_is_still_needed(self):
+        assert not depyf_compat._hook_accepts_inductor_keywords(PRISTINE_HOOK), (
+            "depyf's load_by_key_path now accepts everything Inductor passes; "
+            "drop vllm/compilation/depyf_compat.py and this test"
+        )
+
+    def test_patch_is_skipped_when_unnecessary(self):
+        def accepts_anything(*args, **kwargs):
+            pass
+
+        assert depyf_compat._hook_accepts_inductor_keywords(accepts_anything)
+
+    def test_patch_is_idempotent(self, tmp_path):
+        with depyf_compat.prepare_debug(str(tmp_path)):
+            first = enable_debugging.patched_load_by_key_path
+        with depyf_compat.prepare_debug(str(tmp_path)):
+            assert enable_debugging.patched_load_by_key_path is first
diff --git a/tests/compile/test_inductor_fallback_allow_list_patch.py b/tests/compile/test_inductor_fallback_allow_list_patch.py
deleted file mode 100644
index 29fe9962e345..000000000000
--- a/tests/compile/test_inductor_fallback_allow_list_patch.py
+++ /dev/null
@@ -1,250 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Tests for the Inductor FALLBACK_ALLOW_LIST patch in env_override.py.
-
-The patch wraps ``torch._inductor.lowering.FALLBACK_ALLOW_LIST`` in a thin
-proxy that auto-allows any custom op in the ``vllm::`` or ``vllm_aiter::``
-namespaces. This routes those ops through Inductor's fast-path
-``make_fallback(target, warn=False, override_decomp=True)`` and avoids the
-expensive ``error.operator_str(target, args, kwargs)`` formatting that
-recursively stringifies every input ``TensorBox``.
-
-The slow path is what made ``torch.compile`` effectively hang on Kimi-K2.6
-TP=8 (deep MoE/TP IR provenance trees). These tests cover both the proxy's
-semantics in isolation and the membership-check fast-path that Inductor's
-``GraphLowering.call_function`` actually performs, so we can validate the
-optimization without needing a full GPU compile.
-"""
-
-import time
-
-import pytest
-
-from vllm.env_override import (
-    _patch_inductor_fallback_allow_list,
-    _VllmFallbackAllowList,
-)
-
-
-class TestVllmFallbackAllowListProxy:
-    """Unit tests for the membership-proxy semantics."""
-
-    def test_vllm_namespace_auto_allowed(self):
-        proxy = _VllmFallbackAllowList(set())
-        assert "vllm::all_reduce" in proxy
-        assert "vllm::fused_add_rms_norm" in proxy
-        assert "vllm::all_reduce.default" in proxy
-
-    def test_vllm_aiter_namespace_auto_allowed(self):
-        proxy = _VllmFallbackAllowList(set())
-        assert "vllm_aiter::fused_add_rms_norm" in proxy
-        assert "vllm_aiter::rocm_aiter_fused_moe" in proxy
-
-    def test_unknown_namespace_falls_through(self):
-        proxy = _VllmFallbackAllowList({"torchvision::roi_align"})
-        assert "torchvision::roi_align" in proxy
-        assert "made_up_ns::nonexistent_op" not in proxy
-
-    def test_non_string_falls_through_to_inner(self):
-        sentinel = object()
-        inner = {sentinel}
-        proxy = _VllmFallbackAllowList(inner)
-        assert sentinel in proxy
-        assert object() not in proxy
-
-    def test_prefix_only_match_not_substring(self):
-        proxy = _VllmFallbackAllowList(set())
-        assert "not_vllm::something" not in proxy
-        assert "  vllm::space_prefixed" not in proxy
-
-    def test_standard_entries_preserved(self):
-        base = {"torchvision::roi_align", "aten::index_add"}
-        proxy = _VllmFallbackAllowList(base)
-        assert "torchvision::roi_align" in proxy
-        assert "aten::index_add" in proxy
-        assert "aten::__not_present__" not in proxy
-
-    def test_add_and_discard_delegate_to_inner(self):
-        inner: set[str] = set()
-        proxy = _VllmFallbackAllowList(inner)
-        proxy.add("custom::op")
-        assert "custom::op" in inner
-        proxy.discard("custom::op")
-        assert "custom::op" not in inner
-
-    def test_iter_len_repr(self):
-        base = {"torchvision::roi_align", "aten::index_add"}
-        proxy = _VllmFallbackAllowList(base)
-        assert set(iter(proxy)) == base
-        assert len(proxy) == len(base)
-        assert "torchvision::roi_align" in repr(proxy)
-
-    def test_getattr_delegates_to_inner(self):
-        class _Inner:
-            sentinel = "i_am_inner"
-
-            def some_method(self):
-                return 42
-
-        inner = _Inner()
-        proxy = _VllmFallbackAllowList(inner)
-        assert proxy.sentinel == "i_am_inner"
-        assert proxy.some_method() == 42
-
-    def test_sentinel_attribute(self):
-        proxy = _VllmFallbackAllowList(set())
-        assert proxy._vllm_patched is True
-
-
-class TestPatchApplication:
-    """Integration tests verifying the patch reaches ``torch._inductor``."""
-
-    def test_patch_applied_to_lowering(self):
-        import torch._inductor.lowering as _lowering
-
-        assert getattr(_lowering.FALLBACK_ALLOW_LIST, "_vllm_patched", False), (
-            "env_override._patch_inductor_fallback_allow_list did not run"
-        )
-
-    def test_graph_module_local_binding_rebound(self):
-        # ``torch/_inductor/graph.py`` does:
-        #   from torch._inductor.lowering import FALLBACK_ALLOW_LIST
-        # so the patch has to overwrite the graph module's local binding too,
-        # otherwise the fast-path check in GraphLowering.call_function still
-        # sees the original (unwrapped) OrderedSet.
-        import torch._inductor.graph as _graph
-        import torch._inductor.lowering as _lowering
-
-        if not hasattr(_graph, "FALLBACK_ALLOW_LIST"):
-            pytest.skip(
-                "torch._inductor.graph no longer imports FALLBACK_ALLOW_LIST "
-                "as a module-level symbol; nothing to rebind."
-            )
-
-        assert _graph.FALLBACK_ALLOW_LIST is _lowering.FALLBACK_ALLOW_LIST
-
-    def test_patch_is_idempotent(self):
-        import torch._inductor.lowering as _lowering
-
-        first = _lowering.FALLBACK_ALLOW_LIST
-        _patch_inductor_fallback_allow_list()
-        _patch_inductor_fallback_allow_list()
-        assert _lowering.FALLBACK_ALLOW_LIST is first
-
-    def test_real_vllm_ops_in_real_allow_list(self):
-        # End-to-end membership check using the live (already-patched) object.
-        import torch._inductor.lowering as _lowering
-
-        allow_list = _lowering.FALLBACK_ALLOW_LIST
-        assert "vllm::all_reduce" in allow_list
-        assert "vllm::fused_add_rms_norm" in allow_list
-        assert "vllm_aiter::fused_add_rms_norm" in allow_list
-
-
-class TestInductorFallbackFastPath:
-    """Emulates ``GraphLowering.call_function``'s FALLBACK_ALLOW_LIST check.
-
-    The relevant snippet in ``torch/_inductor/graph.py`` is roughly::
-
-        base_name = target.name()
-        if base_name not in FALLBACK_ALLOW_LIST:
-            log.info(
-                "Creating implicit fallback for:\\n%s",
-                error.operator_str(target, args, kwargs),
-            )
-        out = make_fallback(target, ...)
-
-    On a deep MoE/TP graph (Kimi-K2.6 at TP=4/8) ``operator_str`` recurses
-    through every input ``TensorBox.__str__`` and ends up taking many minutes
-    of CPU per encountered op. The patch ensures the membership test
-    short-circuits for ``vllm::*``/``vllm_aiter::*`` ops so the slow path is
-    never entered. These tests pin that behaviour without needing a real
-    GPU compile.
-    """
-
-    def _simulate_graph_lowering(self, target_names: list[str]):
-        """Returns the set of target names that would have hit the slow
-        operator_str() path under the patched FALLBACK_ALLOW_LIST.
-        """
-        import torch._inductor.lowering as _lowering
-
-        allow_list = _lowering.FALLBACK_ALLOW_LIST
-        slow_path_hits: list[str] = []
-        for name in target_names:
-            if name not in allow_list:
-                slow_path_hits.append(name)
-        return slow_path_hits
-
-    def test_vllm_ops_skip_slow_path(self):
-        slow = self._simulate_graph_lowering(
-            [
-                "vllm::all_reduce",
-                "vllm::fused_add_rms_norm",
-                "vllm_aiter::rocm_aiter_fused_moe",
-                "vllm_aiter::asm_moe",
-            ]
-        )
-        assert slow == [], (
-            "Patched FALLBACK_ALLOW_LIST must short-circuit for all "
-            f"vllm::*/vllm_aiter::* ops; got slow-path hits: {slow}"
-        )
-
-    def test_non_vllm_ops_still_hit_slow_path(self):
-        # Without the patch this is also what would happen; with the patch
-        # the behaviour for non-vllm namespaces must be unchanged.
-        slow = self._simulate_graph_lowering(
-            ["my_user_ns::custom_op", "fancy_ns::something_else"]
-        )
-        assert "my_user_ns::custom_op" in slow
-        assert "fancy_ns::something_else" in slow
-
-    def test_kimi_k2_6_style_op_stream(self):
-        """Emulates one decoder layer's worth of fallback hits.
-
-        Kimi-K2.6 at TP=4 lowers a stream of ``vllm::all_reduce`` +
-        ``vllm_aiter::fused_add_rms_norm`` calls (one per residual block)
-        plus a handful of fused-MoE ops. Pre-patch every one of these would
-        invoke ``operator_str`` and stringify a hundreds-deep IR provenance
-        tree; post-patch they must all short-circuit.
-        """
-        n_layers = 64  # Kimi-K2.6 has ~64 decoder layers per replica
-        op_stream: list[str] = []
-        for _ in range(n_layers):
-            op_stream.extend(
-                [
-                    "vllm::all_reduce",
-                    "vllm_aiter::fused_add_rms_norm",
-                    "vllm_aiter::rocm_aiter_fused_moe",
-                ]
-            )
-
-        start = time.perf_counter()
-        slow = self._simulate_graph_lowering(op_stream)
-        elapsed_s = time.perf_counter() - start
-
-        assert slow == [], (
-            f"Expected all {len(op_stream)} vllm/vllm_aiter ops to take "
-            f"the fast path; got {len(slow)} slow-path hits."
-        )
-        # ``__contains__`` is O(1) per call, so a Kimi-sized stream should
-        # complete in well under a second even on a slow runner. The
-        # pre-patch slow path took many minutes per op on Kimi-K2.6 TP=8.
-        assert elapsed_s < 1.0, (
-            f"FALLBACK_ALLOW_LIST membership check is unexpectedly slow: "
-            f"{elapsed_s:.3f}s for {len(op_stream)} ops"
-        )
-
-    def test_inner_set_membership_still_works_for_standard_ops(self):
-        """The patch must not break Inductor's existing fallback decisions
-        for non-vllm ops such as ``torchvision::roi_align``."""
-        import torch._inductor.lowering as _lowering
-
-        allow_list = _lowering.FALLBACK_ALLOW_LIST
-        # ``torchvision::roi_align`` has been a member of the upstream
-        # FALLBACK_ALLOW_LIST since the original Inductor implementation.
-        # If the proxy ever broke pass-through, this would regress.
-        if "torchvision::roi_align" not in allow_list:
-            pytest.skip(
-                "Upstream FALLBACK_ALLOW_LIST no longer ships "
-                "torchvision::roi_align; nothing to verify."
-            )
diff --git a/tests/compile/test_inductor_operator_str_patch.py b/tests/compile/test_inductor_operator_str_patch.py
new file mode 100644
index 000000000000..5595f189f36a
--- /dev/null
+++ b/tests/compile/test_inductor_operator_str_patch.py
@@ -0,0 +1,43 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for the Inductor implicit-fallback message patch in env_override.py."""
+
+import pytest
+import torch
+from torch._inductor.exc import OperatorIssue
+
+from vllm.env_override import _apply_inductor_operator_str_patch
+from vllm.utils.torch_utils import direct_register_custom_op
+
+
+def _identity(x: torch.Tensor) -> torch.Tensor:
+    return x.clone()
+
+
+direct_register_custom_op(
+    op_name="test_operator_str_identity",
+    op_func=_identity,
+    fake_impl=_identity,
+)
+
+
+class _DeepIRNode:
+    """Stands in for an IR node whose __str__ would recurse for minutes."""
+
+    def __str__(self):
+        raise AssertionError("IR node must not be stringified")
+
+    __repr__ = __str__
+
+
+@pytest.fixture
+def patched_operator_str(monkeypatch):
+    monkeypatch.setattr(OperatorIssue, "operator_str", OperatorIssue.operator_str)
+    _apply_inductor_operator_str_patch()
+    return OperatorIssue.operator_str
+
+
+def test_vllm_op_message_is_target_only(patched_operator_str):
+    op = torch.ops.vllm.test_operator_str_identity.default
+    msg = patched_operator_str(op, [_DeepIRNode()], {"x": _DeepIRNode()})
+    assert msg == f"  target: {op}"
diff --git a/tests/config/test_speculative_draft_hf_overrides.py b/tests/config/test_speculative_draft_hf_overrides.py
index 39ebd7924bf2..923c3d187b09 100644
--- a/tests/config/test_speculative_draft_hf_overrides.py
+++ b/tests/config/test_speculative_draft_hf_overrides.py
@@ -15,8 +15,9 @@
 from unittest.mock import MagicMock, patch
 
 import pytest
-from transformers import Glm5NextConfig, PreTrainedConfig
+from transformers import AutoConfig, Glm5NextConfig, PreTrainedConfig
 
+from vllm.config.model import ModelConfig
 from vllm.config.parallel import ParallelConfig
 from vllm.config.speculative import SpeculativeConfig
 
@@ -38,29 +39,51 @@ def test_dict_overrides_are_not_forwarded_to_draft():
     composed = SpeculativeConfig.compose_draft_hf_overrides(
         {"max_position_embeddings": 1234}
     )
-    assert composed is SpeculativeConfig.hf_config_override
+    hf_config = composed(_make_hf_config(model_type="deepseek_v3"))
+    assert hf_config.architectures == ["DeepSeekMTPModel"]
+    assert getattr(hf_config, "max_position_embeddings", None) != 1234
 
 
 @pytest.mark.cpu_test
 def test_none_overrides_fall_back_to_arch_mapping():
     composed = SpeculativeConfig.compose_draft_hf_overrides(None)
-    assert composed is SpeculativeConfig.hf_config_override
+    hf_config = composed(_make_hf_config(model_type="deepseek_v3"))
+    assert hf_config.architectures == ["DeepSeekMTPModel"]
+
+
+_DSPARK_CONFIG = "vllm.v1.worker.gpu.spec_decode.dspark.config.DSparkSpeculatorConfig"
+_SELF_CONTAINED_DRAFT_SPECULATORS = {
+    "Gemma4DSparkModel": _DSPARK_CONFIG,
+    "DSparkMiniMaxDraftModel": _DSPARK_CONFIG,
+}
 
 
 def _make_speculative_config(
-    hf_kwargs: dict, architecture: str, **speculative_kwargs
+    hf_kwargs: dict,
+    architecture: str,
+    target_model: str = "target",
+    **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
+        hf_config=hf_config,
+        architectures=hf_config.architectures,
+        max_model_len=128,
+        # Self-contained DSpark draft models declare their speculator.
+        speculator_config_class=_SELF_CONTAINED_DRAFT_SPECULATORS.get(
+            hf_config.architectures[0]
+        ),
     )
     draft.registry.inspect_model_cls.return_value = (None, architecture)
-    target = MagicMock(max_model_len=128, quantization=None, hf_overrides={})
+    target = MagicMock(
+        model=target_model, max_model_len=128, quantization=None, hf_overrides={}
+    )
     target.hf_config.model_type = hf_kwargs["model_type"]
+    target.using_transformers_backend.return_value = False
+    speculative_kwargs.setdefault("model", "draft")
 
-    with patch("vllm.config.speculative.ModelConfig", return_value=draft):
+    with patch("vllm.config.ModelConfig", return_value=draft):
         return SpeculativeConfig(
-            model="draft",
             **speculative_kwargs,
             target_model_config=target,
             target_parallel_config=ParallelConfig(),
@@ -113,6 +136,75 @@ def test_dspark_width_is_independent_of_mtp_stages(
     assert config.draft_model_config.hf_config.architectures == [architecture]
 
 
+def _make_minimax_dspark_hf_kwargs() -> dict:
+    return dict(
+        model_type="minimax_m3_vl",
+        architectures=["DSparkMiniMaxDraftModel"],
+        quantization_config={
+            "quant_method": "mxfp8",
+            "ignored_layers": ["language_model.model.dspark.fc"],
+        },
+        text_config=_make_hf_config(
+            architectures=["DSparkMiniMaxDraftModel"],
+            model_type="",
+            num_hidden_layers=5,
+            dspark_block_size=7,
+            dspark_markov_rank=256,
+            dflash_config={"block_size": 7},
+        ),
+    )
+
+
+@pytest.mark.cpu_test
+def test_minimax_dspark_draft_promotes_text_config():
+    """The MiniMax DSpark draft ships a VL wrapper config. It must keep its own
+    draft class (not MiniMax MTP or DeepSeek DSpark), with the DSpark keys and
+    the wrapper's quantization config on the promoted text config."""
+    hf_kwargs = _make_minimax_dspark_hf_kwargs()
+    config = _make_speculative_config(
+        hf_kwargs, "DSparkMiniMaxDraftModel", method="dspark"
+    )
+
+    hf_config = config.draft_model_config.hf_config
+    assert hf_config is hf_kwargs["text_config"]
+    assert hf_config.architectures == ["DSparkMiniMaxDraftModel"]
+    assert config.num_speculative_tokens == 7
+    assert hf_config.quantization_config == hf_kwargs["quantization_config"]
+    assert hf_config.markov_rank == 256
+    assert hf_config.dflash_config["causal"] is False
+
+
+@pytest.mark.cpu_test
+def test_minimax_dspark_defaults_to_dspark_subfolder(tmp_path):
+    """With only method="dspark", MiniMax-M3 loads the draft checkpoint shipped
+    in the target's dspark/ subfolder, with its own quantization config."""
+    (tmp_path / "dspark").mkdir()
+    (tmp_path / "dspark" / "config.json").write_text("{}")
+    config = _make_speculative_config(
+        _make_minimax_dspark_hf_kwargs(),
+        "DSparkMiniMaxDraftModel",
+        target_model=str(tmp_path),
+        model=None,
+        method="dspark",
+    )
+
+    assert config.model == str(tmp_path / "dspark")
+    assert config.quantization is None
+    assert config.num_speculative_tokens == 7
+
+
+@pytest.mark.cpu_test
+def test_minimax_dspark_without_subfolder_raises(tmp_path):
+    with pytest.raises(ValueError, match="no draft checkpoint in the 'dspark/'"):
+        _make_speculative_config(
+            _make_minimax_dspark_hf_kwargs(),
+            "DSparkMiniMaxDraftModel",
+            target_model=str(tmp_path),
+            model=None,
+            method="dspark",
+        )
+
+
 @pytest.mark.cpu_test
 def test_mtp_stages_are_independent_of_dspark_width():
     hf_kwargs = dict(
@@ -238,6 +330,7 @@ def _make_mtp_speculative_config(
         hf_config=draft_hf_config,
         architectures=draft_hf_config.architectures,
         max_model_len=128,
+        speculator_config_class=None,
     )
     target_model_config = MagicMock(
         model="target",
@@ -245,8 +338,9 @@ def _make_mtp_speculative_config(
         quantization=None,
         hf_overrides={},
     )
+    target_model_config.using_transformers_backend.return_value = False
 
-    with patch("vllm.config.speculative.ModelConfig", return_value=draft_model_config):
+    with patch("vllm.config.ModelConfig", return_value=draft_model_config):
         return SpeculativeConfig(
             model="draft",
             method="mtp",
@@ -288,3 +382,37 @@ def test_glm5_next_mtp_override_promotes_index_share():
     assert draft_config.architectures == ["Glm5NextMTPModel"]
     assert draft_config.n_predict == 1
     assert draft_config.index_share_for_mtp_iteration is True
+
+
+@pytest.mark.cpu_test
+@pytest.mark.parametrize(
+    ("model_type", "model_impl", "architecture"),
+    [
+        ("glm4_moe", "transformers", "TransformersMoEForMTP"),
+        ("glm4_moe", "auto", "Glm4MoeMTPModel"),
+        # Does not follow the Transformers MTP standard (no `num_mtp_layers`)
+        ("qwen3_next", "transformers", "Qwen3NextMTP"),
+    ],
+)
+def test_mtp_draft_architecture_follows_target_model_impl(
+    tmp_path, model_type: str, model_impl: str, architecture: str
+):
+    """The MTP draft of a Transformers backend target uses the Transformers backend
+    too if the model follows the Transformers MTP standard, while other targets keep
+    their native MTP draft."""
+    config = AutoConfig.for_model(model_type, num_nextn_predict_layers=1)
+    config.architectures = [config.__class__.__name__.replace("Config", "ForCausalLM")]
+    config.save_pretrained(tmp_path)
+    target = ModelConfig(str(tmp_path), model_impl=model_impl, skip_tokenizer_init=True)
+
+    speculative_config = SpeculativeConfig(
+        method="mtp",
+        num_speculative_tokens=1,
+        target_model_config=target,
+        target_parallel_config=ParallelConfig(),
+    )
+
+    assert speculative_config.method == "mtp"
+    draft = speculative_config.draft_model_config
+    assert draft.architecture == architecture
+    assert draft.hf_config.n_predict == 1
diff --git a/tests/config/test_step5_mtp_config.py b/tests/config/test_step5_mtp_config.py
new file mode 100644
index 000000000000..7124dd659531
--- /dev/null
+++ b/tests/config/test_step5_mtp_config.py
@@ -0,0 +1,32 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+from transformers import PreTrainedConfig
+
+from vllm.config.speculative import SpeculativeConfig
+
+
+@pytest.mark.parametrize("model_type", ["step5", "step5_vl"])
+def test_step5_mtp_hf_config_override(model_type: str):
+    text = PreTrainedConfig(
+        architectures=["Step5ForCausalLM"],
+        num_hidden_layers=92,
+        num_nextn_predict_layers=3,
+    )
+    text.model_type = "step5"
+    if model_type == "step5_vl":
+        config = PreTrainedConfig(
+            architectures=["Step5ForConditionalGeneration"], text_config=text
+        )
+        config.model_type = "step5_vl"
+    else:
+        config = text
+    config.quantization_config = {"quant_method": "fp8"}
+
+    overridden = SpeculativeConfig.hf_config_override(config)
+
+    assert overridden.model_type == "step5_mtp"
+    assert overridden.architectures == ["Step5MTP"]
+    assert overridden.n_predict == 3
+    assert overridden.quantization_config == {"quant_method": "fp8"}
diff --git a/tests/conftest.py b/tests/conftest.py
index a979968c1621..c59bc964161d 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -251,12 +251,16 @@ def dist_init():
 
 
 @pytest.fixture
-def default_vllm_config():
+def default_vllm_config(monkeypatch):
     """Set a default VllmConfig for tests that directly test CustomOps or pathways
     that use get_current_vllm_config() outside of a full engine context.
     """
     from vllm.config import VllmConfig, set_current_vllm_config
 
+    # Creating a breakable-by-default model config sets
+    # VLLM_USE_BREAKABLE_CUDAGRAPH process-wide; shield the default config
+    # from such leaks (it would flip custom_ops to "all").
+    monkeypatch.setenv("VLLM_USE_BREAKABLE_CUDAGRAPH", "0")
     config = VllmConfig()
     with set_current_vllm_config(config):
         yield config
@@ -1012,9 +1016,12 @@ def __init__(
         from vllm.platforms import current_platform
 
         if current_platform.is_rocm():
-            gpu_memory_utilization = kwargs.get(
-                "gpu_memory_utilization",
-                CacheConfig.gpu_memory_utilization,
+            # An unset utilization resolves to 1.0 with the extensible KV cache,
+            # which no longer needs the memory free up front; wait for the
+            # standard fraction in that case.
+            gpu_memory_utilization = (
+                kwargs.get("gpu_memory_utilization")
+                or CacheConfig.DEFAULT_GPU_MEMORY_UTILIZATION
             )
             # V1 startup requires free_memory >= total * gpu_memory_utilization.
             # ROCm CI can hand a test a device that is still lazily releasing
@@ -1378,6 +1385,7 @@ def __exit__(self, exc_type, exc_value, traceback):
         # because when the next test starts some GPU memory is still in use.
         gpu_memory_utilization = (
             self.llm.llm_engine.vllm_config.cache_config.gpu_memory_utilization
+            or CacheConfig.DEFAULT_GPU_MEMORY_UTILIZATION
         )
         from vllm.platforms import current_platform
 
diff --git a/tests/distributed/test_custom_all_reduce.py b/tests/distributed/test_custom_all_reduce.py
index dd462c380d19..57d22a47ce2f 100644
--- a/tests/distributed/test_custom_all_reduce.py
+++ b/tests/distributed/test_custom_all_reduce.py
@@ -182,6 +182,152 @@ def test_all_reduce_mhc_matches_unfused_path(monkeypatch):
     multi_process_parallel(monkeypatch, 4, 1, _all_reduce_mhc)
 
 
+def _swizzled_scale_rows(scales: torch.Tensor, rows: int, cols: int) -> torch.Tensor:
+    """[rows, cols] UE8M0 scales read out of the 128x4 swizzled layout."""
+    row = torch.arange(rows, device=scales.device)[:, None]
+    col = torch.arange(cols, device=scales.device)[None, :]
+    offset = (
+        col % 4
+        + col // 4 * 512
+        + row % 32 * 16
+        + row % 128 // 32 * 4
+        + row // 128 * 128 * cols
+    )
+    return scales[offset]
+
+
+@ray.remote(num_gpus=1, max_calls=1)
+def _finalize_all_reduce_gemma_norm_mxfp8(
+    monkeypatch, tp_size, pp_size, rank, distributed_init_port
+):
+    from vllm.model_executor.layers.fused_moe.moe_output import (
+        MoEOutput,
+        UnfinalizedMoEOutput,
+    )
+    from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
+        mxfp8_e4m3_quantize,
+    )
+    from vllm.models.minimax_m3.nvidia.ops.finalize_allreduce_norm import (
+        FinalizeAllReduceGemmaRMSNormMxfp8,
+    )
+
+    hidden, top_k, eps = 6144, 4, 1e-6
+    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)
+        op = FinalizeAllReduceGemmaRMSNormMxfp8(
+            hidden_size=hidden,
+            max_num_tokens=32,
+            top_k=top_k,
+            fp32_weights=False,
+            device=device,
+        )
+
+        def run(n):
+            torch.manual_seed(7 + rank)
+            # A padded permuted GEMM2 buffer, like the MoE's.
+            rows = n * top_k + 5
+            gemm2 = torch.randn(rows, hidden, device=device, dtype=torch.bfloat16)
+            permuted = torch.randperm(rows, device=device)[: n * top_k]
+            permuted = permuted.view(n, top_k).int()
+            # A route to an expert this rank does not hold.
+            permuted[0, -1] = -1
+            weights = torch.rand(n, top_k, device=device).bfloat16()
+            shared = torch.randn(n, hidden, device=device, dtype=torch.bfloat16)
+            torch.manual_seed(123)
+            residual = torch.randn(n, hidden, device=device, dtype=torch.bfloat16)
+            weight = torch.randn(hidden, device=device, dtype=torch.bfloat16) * 0.1
+            moe_output = MoEOutput(
+                routed=UnfinalizedMoEOutput(gemm2, weights, permuted),
+                shared_output=shared,
+            )
+
+            def fused():
+                return op(moe_output, residual, weight, eps)
+
+            def check(outputs):
+                quantized, residual_out = outputs
+                # One FP32 FMA per route in route order, the shared add, one
+                # rounding: the partial each rank publishes.
+                acc = torch.zeros(n, hidden, device=device)
+                for k in range(top_k):
+                    valid = (permuted[:, k] >= 0).unsqueeze(-1)
+                    rows_k = gemm2[permuted[:, k].clamp_min(0).long()].float()
+                    torch.addcmul(
+                        acc,
+                        torch.where(valid, rows_k, 0.0),
+                        torch.where(valid, weights[:, k : k + 1].float(), 0.0),
+                        out=acc,
+                    )
+                partial = (acc + shared.float()).bfloat16()
+                peers = get_tp_group().all_gather(partial, dim=0).view(tp_size, n, -1)
+                reduced = peers[0].float()
+                for peer in peers[1:]:
+                    reduced = reduced + peer.float()
+                prenorm = (reduced + residual.float()).bfloat16()
+                assert torch.equal(residual_out, prenorm)
+
+                inv_rms = torch.rsqrt(
+                    prenorm.float().square().mean(-1, keepdim=True) + eps
+                )
+                normed = (prenorm.float() * inv_rms * (weight.float() + 1)).bfloat16()
+                ref_values, ref_scales = mxfp8_e4m3_quantize(
+                    normed, is_sf_swizzled_layout=True
+                )
+                # The norm sums the squares in another order with an approximate
+                # rsqrt, so a value can move by one BF16 ULP, and with it rarely
+                # a block's scale. Padding rows' scales are zero in both.
+                assert quantized.scale.shape == ref_scales.shape
+                assert (quantized.scale == ref_scales).float().mean() >= 0.99
+                codes = quantized.data.view(torch.uint8)
+                assert (codes == ref_values.view(torch.uint8)).float().mean() >= 0.99
+                cols = hidden // 32
+                scales = _swizzled_scale_rows(
+                    quantized.scale, quantized.scale.numel() // cols, cols
+                )
+                assert not scales[n:].any()
+                scales = scales[:n]
+                step = torch.exp2(scales.float() - 127).repeat_interleave(32, dim=1)
+                dequantized = quantized.data.float() * step
+                expected = normed.float()
+                # E4M3 rounding (3 mantissa bits) plus one BF16 ULP, and the
+                # smallest E4M3 subnormal step of the block.
+                tolerance = 0.07 * expected.abs() + step * 2.0**-9
+                assert torch.all((dequantized - expected).abs() <= tolerance)
+
+            check(fused())
+            graph = torch.cuda.CUDAGraph()
+            with torch.cuda.graph(graph):
+                for _ in range(5):
+                    fused()
+                captured = fused()
+            for _ in range(20):
+                graph.replay()
+            check(captured)
+            # Replays must pick up in-place input changes.
+            gemm2.mul_(0.5)
+            residual.mul_(2)
+            graph.replay()
+            check(captured)
+
+        # Changing shapes cover shrinking and growing batches in one mailbox.
+        for n in (1, 6, 32, 3, 17, 2, 1):
+            run(n)
+
+
+@pytest.mark.skipif(
+    not current_platform.is_device_capability_family(100), reason="Requires SM100"
+)
+def test_finalize_all_reduce_gemma_norm_mxfp8_matches_unfused_path(monkeypatch):
+    """MiniMax-M3's fused MoE finalize + all-reduce + Gemma RMSNorm + MXFP8."""
+    if torch.accelerator.device_count() < 4:
+        pytest.skip("Requires four GPUs with NVLink multicast")
+    multi_process_parallel(monkeypatch, 4, 1, _finalize_all_reduce_gemma_norm_mxfp8)
+
+
 @pytest.mark.parametrize(
     ("dtype", "expected"),
     [
diff --git a/tests/distributed/test_engram_dp_shard.py b/tests/distributed/test_engram_dp_shard.py
index 4a175e60098a..caa1fa2556f6 100644
--- a/tests/distributed/test_engram_dp_shard.py
+++ b/tests/distributed/test_engram_dp_shard.py
@@ -214,6 +214,8 @@ def _worker(rank: int, tp_size: int, port: int) -> None:
 
         for failure in ("create", "open", "mmap", "register", "pinned", None):
             _check_shared_storage_lifetime_and_failures(failure)
+        for failure in (None, "create", "import"):
+            _check_cumem_host_table(failure)
 
         # GPU peer reads need NVLink P2P, which CI's L4s lack.
         p2p = current_platform.is_fully_connected(
@@ -356,7 +358,11 @@ def unregister(pointer):
         assert result.value == 0
         return result
 
+    from vllm.models.deepseek_v41.nvidia.ops import engram_host_table
+
     with pytest.MonkeyPatch.context() as patch:
+        # Exercise the /dev/shm path even where cuMem tables are used.
+        patch.setattr(engram_host_table, "share_host_table", lambda *args: None)
         patch.setattr(
             engram_ops,
             "tempfile",
@@ -405,6 +411,41 @@ def unregister(pointer):
         assert all(not path.exists() for path in paths)
 
 
+def _check_cumem_host_table(failure):
+    """Chunked cuMem tables read back across ranks; a failed allocation or
+    import on any rank falls back to /dev/shm on all of them."""
+    from vllm.models.deepseek_v41.nvidia.ops import engram_host_table
+
+    if not engram_host_table.fabric_host_memory_supported():
+        return
+
+    group = get_engram_dp_group()
+    # Three 2 MiB chunks, the last one partial.
+    rows = 5 * (1 << 20) // DIM
+    with pytest.MonkeyPatch.context() as patch:
+        patch.setattr(engram_host_table, "CHUNK_BYTES", 2 << 20)
+        if failure == "create" and group.rank_in_group == 0:
+            patch.setattr(engram_host_table, "NUMA_RESERVE_FRACTION", 1.0)
+        if failure == "import" and group.rank_in_group == 1:
+            patch.setattr(
+                engram_host_table.cu,
+                "cuMemImportFromShareableHandle",
+                lambda *args: (
+                    engram_host_table.cu.CUresult.CUDA_ERROR_INVALID_VALUE,
+                    None,
+                ),
+            )
+        storage = engram_ops.DPSharedEngramStorage(rows, DIM, BLOCK, group)
+    assert storage.huge_pages == (failure is None)
+    table = torch.arange(rows * DIM).remainder(251).to(torch.uint8).view(rows, DIM)
+    if group.rank_in_group == 0:
+        storage.weight.view(torch.uint8).copy_(table)
+    group.barrier()
+    weight, _ = storage.get_views(storage.weight, storage.weight_scale_inv)
+    assert torch.equal(weight.view(torch.uint8).cpu(), table)
+    group.barrier()
+
+
 def _check_table(
     vllm_config, cpu_offload, dp_shared_memory, n_heads, sequence_parallel, p2p
 ):
@@ -576,7 +617,8 @@ def unexpected_collective(*args, **kwargs):
 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."""
+    """A metadata-less replica looks up dead rows, so it joins DP sharding and
+    the memory profile sees the Engram activations."""
     from vllm.models.deepseek_v41.common.engram import NgramHashState
 
     if current_platform.is_rocm():
@@ -603,17 +645,11 @@ class Decoder(SimpleNamespace):
         # **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])
+            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 the NVIDIA layer's previous_aux; this stub
             # captures no aux states.
             outputs = (output, hidden, hidden, hidden, hidden, None)
diff --git a/tests/distributed/test_eplb_execute.py b/tests/distributed/test_eplb_execute.py
index 6d0f89248d47..734551c37683 100644
--- a/tests/distributed/test_eplb_execute.py
+++ b/tests/distributed/test_eplb_execute.py
@@ -51,6 +51,16 @@ def test_set_eplb_state_registers_byte_dtypes_as_uint8_views(dtype):
     assert torch.equal(scale.view(torch.uint8), registered_scale)
 
 
+def _skip_unless_symm_mem(world_size: int) -> None:
+    if not current_platform.is_cuda() or not current_platform.is_fully_connected(
+        [
+            current_platform.visible_device_id_to_physical_device_id(i)
+            for i in range(world_size)
+        ]
+    ):
+        pytest.skip("symm_mem requires NVLink-connected CUDA GPUs")
+
+
 @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."""
@@ -436,6 +446,7 @@ def _test_async_transfer_layer_without_mtp_worker(
             expert_weights=expert_weights,
             expert_buffer=expert_buffer,
         )
+        expert_buffer = communicator.owned_expert_buffer or expert_buffer
         communicator.set_stream(stream)
 
         for layer_idx in range(num_layers):
@@ -457,6 +468,8 @@ def _test_async_transfer_layer_without_mtp_worker(
                 new_indices=new_indices_cpu[layer_idx].numpy(),
                 ep_rank=ep_rank,
             )
+            # As in the async worker: expert_buffer writes are not stream-ordered.
+            torch.accelerator.current_stream().synchronize()
 
     local_ok = verify_expert_weights_after_shuffle(
         expert_weights,
@@ -551,6 +564,7 @@ def _test_rearrange_expert_weights_with_redundancy(
             expert_weights=expert_weights,
             expert_buffer=expert_buffer,
         )
+        expert_buffer = communicator.owned_expert_buffer or expert_buffer
 
         # Execute weight rearrangement
         rearrange_expert_weights_inplace(
@@ -614,7 +628,8 @@ def _test_rearrange_expert_weights_with_redundancy(
     ],
 )
 @pytest.mark.parametrize(
-    "eplb_communicator", ["torch_nccl", "torch_gloo", "torch_xccl", "pynccl", "nixl"]
+    "eplb_communicator",
+    ["torch_nccl", "torch_gloo", "torch_xccl", "pynccl", "nixl", "symm_mem"],
 )
 def test_rearrange_expert_weights_with_redundancy(
     world_size,
@@ -634,6 +649,8 @@ def test_rearrange_expert_weights_with_redundancy(
         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")
+    if eplb_communicator == "symm_mem":
+        _skip_unless_symm_mem(world_size)
     distributed_run(
         _test_rearrange_expert_weights_with_redundancy,
         world_size,
@@ -734,7 +751,9 @@ def _test_rearrange_expert_weights_no_change(env, world_size) -> None:
         (2, 2, 2, 3),
     ],
 )
-@pytest.mark.parametrize("eplb_communicator", ["torch_gloo", "torch_xccl", "nixl"])
+@pytest.mark.parametrize(
+    "eplb_communicator", ["torch_gloo", "torch_xccl", "nixl", "symm_mem"]
+)
 def test_async_transfer_layer_without_mtp(
     world_size: int,
     num_layers: int,
@@ -751,6 +770,8 @@ def test_async_transfer_layer_without_mtp(
         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")
+    if eplb_communicator == "symm_mem":
+        _skip_unless_symm_mem(world_size)
 
     distributed_run(
         _test_async_transfer_layer_without_mtp_worker,
diff --git a/tests/distributed/test_eplb_migration_scheduler.py b/tests/distributed/test_eplb_migration_scheduler.py
index 1d63afb2681d..c35d2a3bbf8b 100644
--- a/tests/distributed/test_eplb_migration_scheduler.py
+++ b/tests/distributed/test_eplb_migration_scheduler.py
@@ -1,6 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from collections.abc import Sequence
+
 import numpy as np
 import pytest
 import torch
@@ -32,7 +34,12 @@ def add_recv(
     ) -> None:
         self.recv_calls.append((tensors, src_rank, expert_id))
 
-    def set_transfer_context(self, old_indices: np.ndarray, layer_idx: int) -> None:
+    def set_transfer_context(
+        self,
+        old_indices: np.ndarray,
+        layer_idx: int,
+        recv_rows: Sequence[dict[int, int]] | None = None,
+    ) -> None:
         self.context_calls += 1
 
     def execute(self) -> None:
@@ -276,3 +283,85 @@ def _placements_for_transfers(
         old_counts[src_rank] += 1
         new_counts[dst_rank] += 1
     return num_local_experts, old_indices, new_indices
+
+
+class _PushRecordingCommunicator(_MockEplbCommunicator):
+    """Records each push's target row and each receive's buffer row."""
+
+    def __init__(self, buffer: torch.Tensor) -> None:
+        super().__init__()
+        self.buffer = buffer
+        self.recv_rows: Sequence[dict[int, int]] | None = None
+        self.pushes: list[tuple[int, int, int]] = []
+        self.landings: list[tuple[int, int]] = []
+
+    @property
+    def needs_recv_rows(self) -> bool:
+        return True
+
+    def set_transfer_context(
+        self,
+        old_indices: np.ndarray,
+        layer_idx: int,
+        recv_rows: Sequence[dict[int, int]] | None = None,
+    ) -> None:
+        assert recv_rows is not None
+        self.recv_rows = recv_rows
+
+    def add_send(
+        self, tensors: list[torch.Tensor], dst_rank: int, expert_id: int
+    ) -> None:
+        assert self.recv_rows is not None
+        row = self.recv_rows[dst_rank][expert_id]
+        self.pushes.append((dst_rank, expert_id, row))
+
+    def add_recv(
+        self, tensors: list[torch.Tensor], src_rank: int, expert_id: int
+    ) -> None:
+        row_bytes = self.buffer[0].nbytes
+        row = (tensors[0].data_ptr() - self.buffer.data_ptr()) // row_bytes
+        self.landings.append((expert_id, row))
+
+
+@pytest.mark.parametrize("enable_migration_batching", [False, True])
+@pytest.mark.parametrize("seed", range(8))
+def test_push_recv_rows_match_receiver_rows(
+    enable_migration_batching: bool, seed: int
+) -> None:
+    """A push backend writes each expert into exactly the buffer row its
+    receiver reads it from, for both the default and the batched path."""
+    rng = np.random.default_rng(seed)
+    ep_size, num_local, num_logical = 4, 5, 12
+    num_physical = ep_size * num_local
+
+    def placement() -> np.ndarray:
+        # Every logical expert has a replica; the rest are random redundants.
+        experts = np.concatenate(
+            [
+                np.arange(num_logical),
+                rng.integers(0, num_logical, num_physical - num_logical),
+            ]
+        )
+        return rng.permutation(experts).astype(np.int64)
+
+    old_indices, new_indices = placement(), placement()
+    pushes: set[tuple[int, int, int]] = set()
+    landings: set[tuple[int, int, int]] = set()
+    for rank in range(ep_size):
+        buffer = torch.zeros(num_local, 3)
+        communicator = _PushRecordingCommunicator(buffer)
+        move_to_buffer(
+            num_local_experts=num_local,
+            old_indices=old_indices,
+            new_indices=new_indices,
+            expert_weights=[torch.zeros(num_local, 3)],
+            expert_weights_buffers=[buffer],
+            stream=None,
+            ep_rank=rank,
+            communicator=communicator,
+            enable_migration_batching=enable_migration_batching,
+        )
+        pushes.update(communicator.pushes)
+        landings.update((rank, expert, row) for expert, row in communicator.landings)
+
+    assert pushes == landings
diff --git a/tests/distributed/test_eplb_utils.py b/tests/distributed/test_eplb_utils.py
index dd2bb2760b1e..f54d2cd08742 100644
--- a/tests/distributed/test_eplb_utils.py
+++ b/tests/distributed/test_eplb_utils.py
@@ -157,3 +157,28 @@ def test_commit_eplb_maps_for_layer():
 
     # Layer 1 untouched
     assert torch.equal(model_state.physical_to_logical_map[1], original_phy2log[1])
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
+def test_commit_eplb_maps_for_layer_does_not_sync_host():
+    """Async EPLB commits a layer from the main thread every step, so the commit
+    must not synchronize the host with the GPU."""
+    num_layers, num_logical, max_replicas = 2, 3, 2
+    model_state = _make_model_state(
+        phy2log=torch.zeros(num_layers, 4, dtype=torch.long, device="cuda"),
+        log2phy=torch.full(
+            (num_layers, num_logical, max_replicas),
+            -1,
+            dtype=torch.long,
+            device="cuda",
+        ),
+        logcnt=torch.zeros(num_layers, num_logical, dtype=torch.long, device="cuda"),
+    )
+    new_phy2log = torch.tensor([0, 1, 2, 0], dtype=torch.long)
+
+    prev_mode = torch.cuda.get_sync_debug_mode()
+    torch.cuda.set_sync_debug_mode("error")
+    try:
+        _commit_eplb_maps_for_layer(model_state, new_phy2log, layer=0)
+    finally:
+        torch.cuda.set_sync_debug_mode(prev_mode)
diff --git a/tests/distributed/test_pp_cudagraph.py b/tests/distributed/test_pp_cudagraph.py
index dd2fe227b727..e8df9d3ed553 100644
--- a/tests/distributed/test_pp_cudagraph.py
+++ b/tests/distributed/test_pp_cudagraph.py
@@ -32,7 +32,7 @@ def test_pp_cudagraph(
         # ROCm CI can hand this test a device that is still lazily
         # releasing VRAM from a previous process; wait before starting.
         wait_for_memory_to_settle(
-            threshold_ratio=1.0 - CacheConfig.gpu_memory_utilization
+            threshold_ratio=1.0 - CacheConfig.DEFAULT_GPU_MEMORY_UTILIZATION
         )
 
     cudagraph_args = [
diff --git a/tests/distributed/test_pynccl.py b/tests/distributed/test_pynccl.py
index 924dfb73aba0..5279a0524987 100644
--- a/tests/distributed/test_pynccl.py
+++ b/tests/distributed/test_pynccl.py
@@ -246,6 +246,52 @@ def test_cuda_communicator_all_gather_dim_not_zero():
     distributed_run(cuda_communicator_all_gather_dim_worker_fn, 2)
 
 
+@worker_fn_wrapper
+def all_to_all_single_worker_fn():
+    pynccl_comm = PyNcclCommunicator(
+        get_world_group().cpu_group, device=get_world_group().device
+    )
+    rank = pynccl_comm.rank
+    world_size = pynccl_comm.world_size
+    device = f"cuda:{rank}"
+
+    # Rank src sends (src + dst) % 3 rows to rank dst, so some splits are
+    # empty. Row i of the src -> dst chunk carries src * 100 + dst * 10 + i.
+    def num_rows(src: int, dst: int) -> int:
+        return (src + dst) % 3
+
+    def chunk(src: int, dst: int) -> torch.Tensor:
+        values = src * 100 + dst * 10 + torch.arange(num_rows(src, dst))
+        return values.to(torch.float32).view(-1, 1).expand(-1, 8)
+
+    send = torch.cat([chunk(rank, dst) for dst in range(world_size)]).to(device)
+    expected = torch.cat([chunk(src, rank) for src in range(world_size)]).to(device)
+    recv = torch.empty_like(expected)
+    output_split_sizes = [num_rows(src, rank) for src in range(world_size)]
+    input_split_sizes = [num_rows(rank, dst) for dst in range(world_size)]
+
+    pynccl_comm.all_to_all_single(recv, send, output_split_sizes, input_split_sizes)
+    torch.accelerator.synchronize()
+    torch.testing.assert_close(recv, expected, rtol=0, atol=0)
+
+    # Split sizes fixed at capture, so the collective replays from a CUDA
+    # graph with fresh inputs.
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        pynccl_comm.all_to_all_single(recv, send, output_split_sizes, input_split_sizes)
+    send.add_(1000)
+    graph.replay()
+    torch.accelerator.synchronize()
+    torch.testing.assert_close(recv, expected + 1000, rtol=0, atol=0)
+
+
+@pytest.mark.skipif(
+    torch.accelerator.device_count() < 2, reason="Need at least 2 GPUs to run the test."
+)
+def test_pynccl_all_to_all_single():
+    distributed_run(all_to_all_single_worker_fn, 2)
+
+
 @worker_fn_wrapper
 def all_gatherv_worker_fn():
     pynccl_comm = PyNcclCommunicator(
diff --git a/tests/distributed/test_weight_transfer.py b/tests/distributed/test_weight_transfer.py
index fefd9059ee01..cca39d5f1a88 100644
--- a/tests/distributed/test_weight_transfer.py
+++ b/tests/distributed/test_weight_transfer.py
@@ -8,11 +8,12 @@
 
 import builtins
 import importlib.util
+import math
 import pickle
 import runpy
 import threading
 import time
-from contextlib import nullcontext
+from contextlib import ExitStack, nullcontext
 from functools import cached_property
 from types import SimpleNamespace
 from unittest.mock import MagicMock
@@ -50,7 +51,10 @@
     IPCWeightTransferInitInfo,
     IPCWeightTransferUpdateInfo,
 )
-from vllm.distributed.weight_transfer.nccl_common import run_trainer_rendezvous
+from vllm.distributed.weight_transfer.nccl_common import (
+    run_trainer_rendezvous,
+    uid_init_process_group,
+)
 from vllm.distributed.weight_transfer.nccl_engine import (
     NCCLTrainerInitInfo,
     NCCLTrainerWeightTransferEngine,
@@ -281,12 +285,14 @@ def test_parse_update_info_valid(self):
                 "names": ["w1", "w2"],
                 "dtype_names": ["float32", "bfloat16"],
                 "shapes": [[100, 100], [50]],
+                "src": 1,
             }
         )
         assert isinstance(update_info, NCCLWeightTransferUpdateInfo)
         assert update_info.names == ["w1", "w2"]
         assert update_info.dtype_names == ["float32", "bfloat16"]
         assert update_info.shapes == [[100, 100], [50]]
+        assert update_info.src == 1
 
 
 # --- Unit Tests: Engine Registry ---
@@ -396,6 +402,13 @@ def test_nccl_receive_weights_without_init_raises():
         engine.receive_weights(update_info)
 
 
+def test_nccl_disabled_communicator_raises(monkeypatch):
+    """A disabled PyNccl communicator would turn every broadcast into a no-op."""
+    monkeypatch.setenv("VLLM_DISABLE_PYNCCL", "1")
+    with pytest.raises(RuntimeError, match="would be a no-op"):
+        uid_init_process_group(b"\x00" * 128, rank=0, world_size=2, device=0)
+
+
 def test_sparse_nccl_receive_weights_without_init_raises():
     """Test that sparse receive raises if init_transfer_engine wasn't called."""
     if torch.accelerator.device_count() < 1:
@@ -627,6 +640,197 @@ def test_nccl_weight_transfer_between_processes(batch_invariant):
     )
 
 
+@ray.remote(num_gpus=1)
+class MultiSourceInferenceWorker:
+    def __init__(self, worker_rank: int, num_workers: int):
+        import contextlib
+
+        device = _set_ray_assigned_device()
+
+        from vllm.config.weight_transfer import WeightTransferConfig
+        from vllm.distributed.weight_transfer.nccl_engine import (
+            NCCLWeightTransferEngine,
+        )
+
+        class Recorder(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.received: dict[str, float] = {}
+                self.load_counts: dict[str, int] = {}
+
+            def load_weights(self, weights):
+                for name, tensor in weights:
+                    self.received[name] = tensor.sum().item()
+                    self.load_counts[name] = self.load_counts.get(name, 0) + 1
+
+        self.recorder = Recorder()
+        self.engine = NCCLWeightTransferEngine(
+            WeightTransferConfig(backend="nccl"),
+            create_mock_vllm_config(rank=worker_rank, world_size=num_workers),
+            device,
+            self.recorder,
+        )
+        import vllm.config as vllm_config_mod
+
+        vllm_config_mod.set_current_vllm_config = lambda cfg: contextlib.nullcontext()
+
+    def ready(self):
+        return True
+
+    def initialize(self, init_info):
+        self.engine.init_transfer_engine(self.engine.parse_init_info(init_info))
+
+    def update(self, update_info):
+        self.engine.receive_weights(self.engine.parse_update_info(update_info))
+
+    def results(self):
+        torch.accelerator.synchronize()
+        results = {
+            "received": self.recorder.received,
+            "load_counts": self.recorder.load_counts,
+        }
+        self.engine.shutdown()
+        return results
+
+
+@ray.remote(num_gpus=1)
+def trainer_send_from_source(
+    receivers,
+    master_address: str,
+    master_port: int,
+    world_size: int,
+    num_sources: int,
+    rank: int,
+    tensor_shape: list[int],
+    num_tensors_per_source: int,
+    packed: bool,
+) -> bool:
+    from functools import cached_property
+
+    device = _set_ray_assigned_device()
+
+    from vllm.distributed.weight_transfer import (
+        ParamMeta,
+        WeightSource,
+        WeightTransferTrainerFactory,
+    )
+    from vllm.distributed.weight_transfer.nccl_engine import NCCLTrainerInitInfo
+
+    class Source(WeightSource):
+        """Rank r owns (and broadcasts) `source_r.*`, like a pipeline stage."""
+
+        @cached_property
+        def metadata(self):
+            return [
+                ParamMeta(
+                    f"source_{r}.weight_{i}",
+                    torch.float32,
+                    tuple(tensor_shape),
+                    owners=(r,),
+                )
+                for r in range(num_sources)
+                for i in range(num_tensors_per_source)
+            ]
+
+        def iter_full(self):
+            for i in range(num_tensors_per_source):
+                value = rank * num_tensors_per_source + i + 1
+                yield (
+                    f"source_{rank}.weight_{i}",
+                    torch.full(tensor_shape, value, dtype=torch.float32, device=device),
+                )
+
+    class Client:
+        def call(self, method, *args):
+            return ray.get(
+                [getattr(receiver, method).remote(*args) for receiver in receivers]
+            )
+
+        def init_weight_transfer_engine(self, init_info):
+            self.call("initialize", init_info)
+
+        def start_weight_update(self):
+            pass
+
+        def update_weights(self, update_info):
+            self.call("update", update_info)
+
+        def finish_weight_update(self, weight_version=None):
+            pass
+
+    engine = WeightTransferTrainerFactory.trainer_init(
+        NCCLTrainerInitInfo(
+            master_address=master_address,
+            master_port=master_port,
+            world_size=world_size,
+            rank=rank,
+            packed=packed,
+            packed_buffer_size_bytes=1024,
+        ),
+        client=Client(),
+        source=Source(),
+    )
+    for _ in range(2):
+        engine.send_weights()
+    engine.shutdown()
+    return True
+
+
+@pytest.mark.skipif(
+    torch.accelerator.device_count() < 4,
+    reason="Need 4 GPUs for two NCCL sources and two inference workers.",
+)
+@pytest.mark.parametrize("packed", [False, True], ids=["unpacked", "packed"])
+def test_nccl_weight_transfer_from_multiple_sources(packed):
+    """Two trainer ranks each own and broadcast half the parameters (like
+    pipeline stages); both workers receive all of them, twice (the second
+    round reusing metadata)."""
+    _init_ray_for_weight_transfer()
+
+    master_address = "127.0.0.1"
+    master_port = get_open_port()
+    num_sources = 2
+    num_workers = 2
+    world_size = num_sources + num_workers
+    tensor_shape = [8, 8]
+    num_tensors_per_source = 10
+
+    receivers = [
+        MultiSourceInferenceWorker.remote(  # type: ignore[attr-defined]
+            worker_rank, num_workers
+        )
+        for worker_rank in range(num_workers)
+    ]
+    ray.get([receiver.ready.remote() for receiver in receivers])
+    futures = [
+        trainer_send_from_source.remote(
+            receivers,
+            master_address,
+            master_port,
+            world_size,
+            num_sources,
+            rank,
+            tensor_shape,
+            num_tensors_per_source,
+            packed,
+        )
+        for rank in range(num_sources)
+    ]
+    assert ray.get(futures, timeout=300) == [True] * num_sources
+    expected_numel = math.prod(tensor_shape)
+    expected = {
+        f"source_{r}.weight_{i}": (r * num_tensors_per_source + i + 1) * expected_numel
+        for r in range(num_sources)
+        for i in range(num_tensors_per_source)
+    }
+    results = ray.get([receiver.results.remote() for receiver in receivers])
+    for result in results:
+        assert result["received"] == expected
+        assert result["load_counts"] == dict.fromkeys(expected, 2)
+    for receiver in receivers:
+        ray.kill(receiver)
+
+
 def test_sparse_nccl_checkpoint_chunks_to_ep_local_experts_cpu(monkeypatch):
     """Replay global expert patches through two EP-local loaders on CPU."""
 
@@ -1702,7 +1906,9 @@ def test_nccl_trainer_init_ships_worker_init_info(monkeypatch, packed):
 
     # Bypass the real NCCL rendezvous.
     group = object()
-    monkeypatch.setattr(nccl_engine_mod, "open_trainer_endpoint", lambda info: group)
+    monkeypatch.setattr(
+        nccl_engine_mod, "open_trainer_endpoint", lambda info, rank: group
+    )
 
     client = RecordingClient()
     engine = WeightTransferTrainerFactory.trainer_init(
@@ -1849,7 +2055,7 @@ def test_nccl_trainer_send_weights_drives_client_in_order():
         packed=False,
     )
     # Bypass the real NCCL rendezvous; broadcast is a no-op.
-    engine.model_update_group = MagicMock()
+    engine.model_update_group = MagicMock(rank=0)
 
     engine.send_weights()
 
@@ -1873,7 +2079,7 @@ def test_nccl_trainer_sends_metadata_only_when_changed():
         source=ModuleSource(_module_with(("w", torch.zeros(4, device="cuda")))),
         packed=False,
     )
-    engine.model_update_group = MagicMock()
+    engine.model_update_group = MagicMock(rank=0)
 
     engine.send_weights()
     assert client.last_update_info is not None
@@ -1886,8 +2092,9 @@ def test_nccl_trainer_sends_metadata_only_when_changed():
 
 
 def test_nccl_worker_reuses_metadata_per_model(monkeypatch):
-    """Empty update info reuses the previous call's metadata for the same model;
-    target and draft keep separate copies, and re-init clears them."""
+    """Empty update info reuses the previous call's metadata for the same
+    (model, src) and receives from that src; re-init clears every entry, and
+    src must be a trainer rank."""
     import vllm.distributed.weight_transfer.nccl_engine as nccl_engine_mod
 
     monkeypatch.setattr(
@@ -1896,7 +2103,9 @@ def test_nccl_worker_reuses_metadata_per_model(monkeypatch):
     target, draft = MagicMock(spec=torch.nn.Module), MagicMock(spec=torch.nn.Module)
     received: list = []
     for model in (target, draft):
-        model.load_weights = lambda w, m=model: received.append((m, w[0][0]))
+        model.load_weights = lambda w, m=model: received.append(
+            (m, w[0][0], tuple(w[0][1].shape))
+        )
     engine = NCCLWeightTransferEngine(
         WeightTransferConfig(backend="nccl"),
         create_mock_vllm_config(),
@@ -1904,29 +2113,53 @@ def test_nccl_worker_reuses_metadata_per_model(monkeypatch):
         target,
     )
     init_info = NCCLWeightTransferInitInfo(
-        master_address="127.0.0.1", master_port=29500, rank_offset=1, world_size=2
+        master_address="127.0.0.1", master_port=29500, rank_offset=2, world_size=3
     )
     engine.init_transfer_engine(init_info)
     monkeypatch.setattr(torch.cuda, "current_stream", MagicMock())
 
-    def update(model, *names):
+    def update(model, src, name=None, shape=None):
         engine.set_weight_update_target(model, None)
         engine.receive_weights(
             NCCLWeightTransferUpdateInfo(
-                names=list(names),
-                dtype_names=["float32"] * len(names),
-                shapes=[[2]] * len(names),
+                names=[name] if name else [],
+                dtype_names=["float32"] if name else [],
+                shapes=[shape] if name else [],
+                src=src,
             )
         )
 
-    update(target, "t")
-    update(draft, "d")
-    update(target)
-    update(draft)
+    seeded = {
+        (target, 0): ("t0", [2]),
+        (target, 1): ("t1", [3]),
+        (draft, 0): ("d0", [4]),
+        (draft, 1): ("d1", [5]),
+    }
+    for (model, src), (name, shape) in seeded.items():
+        update(model, src, name, shape)
+    reused = [(draft, 1), (target, 0), (draft, 0), (target, 1)]
+    for model, src in reused:
+        update(model, src)
+
+    order = [*seeded, *reused]
+    assert received == [
+        (model, seeded[model, src][0], tuple(seeded[model, src][1]))
+        for model, src in order
+    ]
+    group = engine.model_update_group
+    assert [c.kwargs["src"] for c in group.broadcast.call_args_list] == [
+        src for _, src in order
+    ]
+
     engine.init_transfer_engine(init_info)
-    update(target)
+    for model, src in seeded:
+        update(model, src)
+    assert len(received) == len(order)
+    engine.model_update_group.broadcast.assert_not_called()
 
-    assert received == [(target, "t"), (draft, "d"), (target, "t"), (draft, "d")]
+    for bad in (-1, 2):
+        with pytest.raises(ValueError, match=r"sender rank in \[0, 2\)"):
+            update(target, bad, "c", [2])
 
 
 class _ScriptedSource(WeightSource):
@@ -1950,7 +2183,7 @@ def _mock_group_engine(source, monkeypatch, **kwargs):
     engine = NCCLTrainerWeightTransferEngine(
         client=RecordingClient(), source=source, packed=False, **kwargs
     )
-    engine.model_update_group = MagicMock()
+    engine.model_update_group = MagicMock(rank=0)
     monkeypatch.setattr(torch.cuda, "current_stream", MagicMock())
     return engine
 
@@ -2007,29 +2240,62 @@ def test_nccl_trainer_send_weights_rejects_truncated_source(monkeypatch):
         engine.send_weights()
 
 
-def test_nccl_trainer_rejects_a_second_sender():
-    """Each parameter is sent by its first owner and broadcast has one sender,
-    so a pipeline stage that does not start at rank 0 cannot be sent; every
-    rank sees the same metadata and raises alike."""
-    meta = [
-        ParamMeta("w", torch.float32, (4,), owners=(0, 1)),
-        ParamMeta("b", torch.float32, (2,), owners=(2, 3)),
-    ]
-    source = _ScriptedSource(meta, [])
+_TWO_STAGES = [
+    ParamMeta("w", torch.float32, (4,), owners=(0, 1)),
+    ParamMeta("b", torch.float32, (2,), owners=(2, 3)),
+]
+
+
+@pytest.mark.parametrize(
+    ("meta", "match"),
+    [
+        (
+            [ParamMeta("w", torch.float32, (4,)), _TWO_STAGES[1]],
+            "'w' has no owners",
+        ),
+        (
+            [*_TWO_STAGES, ParamMeta("x", torch.float32, (2,), owners=(1, 2))],
+            "rank 1 owns parameters broadcast by ranks 0 and 1",
+        ),
+        (
+            [ParamMeta("w", torch.float32, (4,), owners=(1, 0))],
+            r"'w' has owners \(1, 0\)",
+        ),
+        (
+            [ParamMeta("w", torch.float32, (4,), owners=(1, 0)), _TWO_STAGES[1]],
+            r"'w' has owners \(1, 0\)",
+        ),
+    ],
+)
+def test_nccl_trainer_rejects_ranks_spanning_broadcasters(meta, match):
+    """Each broadcaster sends in its own turn, so a rank owning parameters of
+    two broadcasters would gather out of turn; every rank sees the same
+    metadata and raises alike."""
     for rank in (0, 2):
-        with pytest.raises(ValueError, match="'b' is sent by its first owner, rank 2"):
+        with pytest.raises(ValueError, match=match):
             NCCLTrainerWeightTransferEngine.trainer_init(
                 NCCLTrainerInitInfo(
                     master_address="127.0.0.1",
                     master_port=29500,
-                    world_size=2,
+                    world_size=3,
                     rank=rank,
                 ),
                 client=RecordingClient(),
-                source=source,
+                source=_ScriptedSource(meta, []),
             )
 
 
+def test_nccl_trainer_rejects_world_size_without_workers():
+    with pytest.raises(ValueError, match="leaves no inference workers"):
+        NCCLTrainerWeightTransferEngine.trainer_init(
+            NCCLTrainerInitInfo(
+                master_address="127.0.0.1", master_port=29500, world_size=2, rank=0
+            ),
+            client=RecordingClient(),
+            source=_ScriptedSource(_TWO_STAGES, []),
+        )
+
+
 def test_nccl_trainer_non_sender_owner_drives_its_gathers(monkeypatch):
     """A non-sender owner (an FSDP peer) still iterates what it owns, since
     the sender's gathers need it; it never touches the client or the group."""
@@ -2098,7 +2364,7 @@ def iter_full(self):
     engine = NCCLTrainerWeightTransferEngine(
         client=_BlockingClient(), source=_FailingSource(), packed=False
     )
-    engine.model_update_group = MagicMock()
+    engine.model_update_group = MagicMock(rank=0)
     monkeypatch.setattr(torch.cuda, "current_stream", MagicMock())
 
     started = time.perf_counter()
@@ -2354,8 +2620,11 @@ def mx_worker(mx_client, monkeypatch):
     worker.weight_transfer_engine = _make_modelexpress_engine()
     worker.vllm_config = worker.weight_transfer_engine.vllm_config
     worker._weight_update_active = False
+    worker._incomplete_weight_updates = set()
     worker._weight_update_is_draft = False
-    worker.model_runner = MagicMock()
+    worker._kv_cache_reclaimable = False
+    worker._weight_update_kv_cache_loan = ExitStack()
+    worker.model_runner = MagicMock(extensible_kv_cache=None)
     return worker
 
 
diff --git a/tests/engine/test_arg_utils.py b/tests/engine/test_arg_utils.py
index 938cf9b3999b..c183b32a8d32 100644
--- a/tests/engine/test_arg_utils.py
+++ b/tests/engine/test_arg_utils.py
@@ -804,6 +804,11 @@ def test_numa_bind_args():
     assert engine_args.numa_bind_nodes == [0, 0, 1, 1]
     assert engine_args.numa_bind_cpus == ["0-3", "4-7", "8-11", "12-15"]
 
+    # Unset means auto; --no-numa-bind disables explicitly.
+    assert EngineArgs.from_cli_args(args=parser.parse_args([])).numa_bind is None
+    args = parser.parse_args(["--no-numa-bind"])
+    assert EngineArgs.from_cli_args(args=args).numa_bind is False
+
 
 def test_ir_op_priority():
     from vllm.config.kernel import IrOpPriorityConfig, KernelConfig
diff --git a/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py b/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py
index aa50a16b2365..e28e3a8e8b33 100644
--- a/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py
+++ b/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py
@@ -405,6 +405,37 @@ def test_tool_result_with_multiple_images(self):
             "https://example.com/img2.jpg",
         ]
 
+    def test_tool_result_with_tool_references(self):
+        """A tool_result with tool_reference items stays one tool message, with
+        the references first, so templates pair one result with one call."""
+        request = self._make_tool_result_request(
+            [
+                {"type": "text", "text": "2 tools found"},
+                {"type": "tool_reference", "tool_name": "get_weather"},
+                {"type": "tool_reference", "tool_name": "get_time"},
+                {
+                    "type": "image",
+                    "source": {
+                        "type": "base64",
+                        "media_type": "image/png",
+                        "data": "AAAA",
+                    },
+                },
+            ]
+        )
+        result = _convert(request)
+
+        assert [m["role"] for m in result.messages] == ["assistant", "tool", "user"]
+        assert result.messages[1] == {
+            "role": "tool",
+            "tool_call_id": "call_001",
+            "content": [
+                {"type": "tool_reference", "name": "get_weather"},
+                {"type": "tool_reference", "name": "get_time"},
+                {"type": "text", "text": "2 tools found"},
+            ],
+        }
+
     def test_tool_result_none_content(self):
         request = self._make_tool_result_request(None)
         result = _convert(request)
@@ -1454,6 +1485,27 @@ def test_no_template_defaults_merge(self):
         """No chat_template → conservative default: merge."""
         assert AnthropicServingMessages._detect_merge_inline_system(None) is True
 
+    def test_tools_none_check_template_merges(self):
+        """OLMo 3-style ``tools | tojson`` on undefined tools falls back to merge."""
+        template = (
+            "{% if tools is none %}{% else %}{{ tools | tojson }}{% endif %}"
+            "{% for m in messages %}{{ m.content }}{% endfor %}"
+        )
+        assert AnthropicServingMessages._detect_merge_inline_system(template) is True
+
+    def test_tools_loop_template_no_merge(self):
+        """Hermes-style ``for tool in tools`` must not raise in the probe."""
+        template = (
+            "{% for tool in tools %}{{ tool | tojson }}{% endfor %}"
+            "{% for m in messages %}{{ m.content }}{% endfor %}"
+        )
+        assert AnthropicServingMessages._detect_merge_inline_system(template) is False
+
+    def test_non_template_error_defaults_merge(self):
+        """A non-TemplateError render failure must not abort startup."""
+        template = "{{ undefined_var | tojson }}"
+        assert AnthropicServingMessages._detect_merge_inline_system(template) is True
+
     @pytest.mark.parametrize(
         ("resolved", "expected"),
         [(PERMISSIVE_TEMPLATE, False), (Q35_TEMPLATE, True)],
@@ -1486,6 +1538,73 @@ def test_probes_resolved_template_without_cli_override(
         )
         assert mock_logger.warning_once.called is expected
 
+    @pytest.mark.parametrize(
+        ("model", "expected"),
+        [
+            # Renders inline system messages as system turns; merging them
+            # would break prefix caching for Claude Code's per-turn reminders.
+            ("deepseek-ai/DeepSeek-V4.1-Flash", False),
+            # Appends inline system text to the preceding user turn (#46710).
+            ("deepseek-ai/DeepSeek-V4-Flash", True),
+            # Puts a trailing system message after the generation prompt.
+            ("deepseek-ai/DeepSeek-V3.2-Exp", True),
+        ],
+    )
+    def test_probes_non_jinja_renderer(self, model, expected):
+        """Encoder-based renderers are probed by rendering, not registered."""
+        from vllm.config import ModelConfig, VllmConfig
+        from vllm.renderers import renderer_from_config
+        from vllm.renderers.online_renderer import OnlineRenderer
+
+        model_config = ModelConfig(model=model)
+        online_renderer = OnlineRenderer(
+            model_config=model_config,
+            renderer=renderer_from_config(VllmConfig(model_config=model_config)),
+            request_logger=None,
+            chat_template=None,
+            chat_template_content_format="auto",
+        )
+        assert (
+            AnthropicServingMessages._should_merge_inline_system(online_renderer)
+            is expected
+        )
+
+    def test_non_jinja_probe_uses_server_chat_template(self):
+        """Renderers that honor --chat-template (e.g. Cohere) probe with it."""
+        renderer = MagicMock()
+        renderer.render_messages.side_effect = ValueError
+        online_renderer = SimpleNamespace(
+            renderer=renderer,
+            chat_template="custom",
+            chat_template_content_format="string",
+            default_chat_template_kwargs={"thinking": False},
+        )
+        AnthropicServingMessages._should_merge_inline_system(online_renderer)
+        params = renderer.render_messages.call_args.args[1]
+        assert params.chat_template == "custom"
+        assert params.chat_template_content_format == "string"
+        assert params.chat_template_kwargs == {
+            "thinking": False,
+            "add_generation_prompt": True,
+        }
+
+    def test_non_jinja_merge_warning_omits_chat_template_hint(self, monkeypatch):
+        """--chat-template may not apply to non-jinja renderers (e.g. V4)."""
+        import vllm.entrypoints.anthropic.serving as serving_mod
+
+        mock_logger = MagicMock()
+        monkeypatch.setattr(serving_mod, "logger", mock_logger)
+        renderer = MagicMock()
+        renderer.render_messages.side_effect = ValueError
+        online_renderer = SimpleNamespace(
+            renderer=renderer,
+            chat_template=None,
+            chat_template_content_format="auto",
+            default_chat_template_kwargs={},
+        )
+        assert AnthropicServingMessages._should_merge_inline_system(online_renderer)
+        assert "--chat-template" not in mock_logger.warning_once.call_args.args[0]
+
 
 # ======================================================================
 # Full (non-streaming) response conversion: messages_full_converter
@@ -1871,6 +1990,44 @@ def test_count_tokens_validation_error_returns_bad_request(self):
         assert response.json()["error"]["type"] == "invalid_request_error"
 
 
+# ======================================================================
+# tool definitions
+# ======================================================================
+
+
+class TestToolDefinitions:
+    @pytest.mark.parametrize(
+        "tool",
+        [
+            {"type": "web_search_20250305", "name": "web_search", "max_uses": 8},
+            {"type": "bash_20250124", "name": "bash"},
+        ],
+    )
+    def test_anthropic_defined_tool_is_rejected_by_type(self, tool):
+        with pytest.raises(ValidationError) as exc_info:
+            _make_request([{"role": "user", "content": "Hi"}], tools=[tool])
+
+        errors = exc_info.value.errors()
+        assert len(errors) == 1
+        assert f"{tool['type']!r} is not supported" in errors[0]["msg"]
+
+    @pytest.mark.parametrize("extra", [{}, {"type": "custom"}])
+    def test_custom_tool_is_accepted(self, extra):
+        request = _make_request(
+            [{"role": "user", "content": "Hi"}],
+            tools=[{"name": "get_weather", "input_schema": {}, **extra}],
+        )
+        assert request.tools is not None
+        assert request.tools[0].input_schema == {"type": "object"}
+
+    def test_custom_tool_still_requires_input_schema(self):
+        with pytest.raises(ValidationError, match="input_schema"):
+            _make_request(
+                [{"role": "user", "content": "Hi"}],
+                tools=[{"type": "custom", "name": "get_weather"}],
+            )
+
+
 # ======================================================================
 # thinking configuration pass-through
 # ======================================================================
diff --git a/tests/entrypoints/anthropic/test_messages.py b/tests/entrypoints/anthropic/test_messages.py
index 9d7fa9210426..dafcd03adfeb 100644
--- a/tests/entrypoints/anthropic/test_messages.py
+++ b/tests/entrypoints/anthropic/test_messages.py
@@ -243,3 +243,20 @@ async def get_stream_usage(resp):
     warm_usage = await get_stream_usage(await client.messages.create(**request))
     assert warm_usage.cache_read_input_tokens is not None
     assert warm_usage.cache_read_input_tokens > 0
+
+
+@pytest.mark.asyncio
+async def test_anthropic_server_tool_is_rejected(client: anthropic.AsyncAnthropic):
+    with pytest.raises(anthropic.BadRequestError) as exc_info:
+        await client.messages.create(
+            model="claude-3-7-sonnet-latest",
+            max_tokens=16,
+            messages=[{"role": "user", "content": "What's new in vLLM?"}],
+            tools=[
+                {"type": "web_search_20250305", "name": "web_search", "max_uses": 8}
+            ],
+        )
+
+    message = exc_info.value.body["error"]["message"]
+    assert "'web_search_20250305' is not supported" in message
+    assert "Field required" not in message
diff --git a/tests/entrypoints/generate/structured_decisions/test_strategies.py b/tests/entrypoints/generate/structured_decisions/test_strategies.py
index 1f43128ffb30..90d1a5c71974 100644
--- a/tests/entrypoints/generate/structured_decisions/test_strategies.py
+++ b/tests/entrypoints/generate/structured_decisions/test_strategies.py
@@ -1,6 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import asyncio
 import json
 import logging
 import math
@@ -64,7 +63,7 @@
     init_exception_handler,
 )
 from vllm.entrypoints.serve.utils.request_logger import RequestLogger
-from vllm.inputs import mm_input, tokens_input
+from vllm.inputs import mm_input
 from vllm.logprobs import Logprob
 from vllm.multimodal.inputs import PlaceholderRange
 from vllm.multimodal.media import MediaConnector, MediaWithBytes
@@ -279,7 +278,7 @@ def test_gemma4_labels_start_the_reply(model_name, reply_start):
 
 
 @pytest.fixture
-def engine_and_renderer(qwen):
+def engine_and_renderer(qwen, request):
     tokenizer, _ = qwen
     model_config = ModelConfig(model="Qwen/Qwen3-0.6B", dtype="float32")
     renderer = HfRenderer(
@@ -300,6 +299,7 @@ def engine_and_renderer(qwen):
         request_logger=None,
         chat_template=None,
         chat_template_content_format="auto",
+        trust_request_chat_template=getattr(request, "param", False),
     )
     yield engine, online_renderer
     renderer.shutdown()
@@ -988,8 +988,9 @@ def test_liquid_caps_images_like_the_checkpoint():
 
 
 @pytest.fixture
-def decision_server(qwen):
-    tokenizer, _ = qwen
+def decision_server(engine_and_renderer):
+    _, renderer = engine_and_renderer
+    tokenizer = renderer.renderer.get_tokenizer()
 
     class Engine:
         model_config = SimpleNamespace(
@@ -1035,21 +1036,7 @@ async def generate(self, prompt, params, request_id, **kwargs):
                 finished=True,
             )
 
-    class Renderer:
-        renderer = SimpleNamespace(get_tokenizer=lambda: tokenizer)
-
-        async def preprocess_chat(self, request, messages, **kwargs):
-            await asyncio.sleep(0)
-            ids = tokenizer.apply_chat_template(
-                messages,
-                add_generation_prompt=True,
-                enable_thinking=False,
-                tokenize=True,
-                return_dict=False,
-            )
-            return [], [tokens_input(ids)]
-
-    engine, renderer = Engine(), Renderer()
+    engine = Engine()
     strategy = NextTokenStrategy(ReadContext(engine, renderer, None, "auto", {}))
     models = OpenAIServingModels(
         engine, [BaseModelPath(name="test-model", model_path="test-model")]
@@ -1095,6 +1082,43 @@ def decision_body(endpoint, count=1):
     )
 
 
+@pytest.mark.parametrize("engine_and_renderer", [False, True], indirect=True)
+@pytest.mark.parametrize(
+    "template_kwargs",
+    [
+        None,
+        {},
+        {"enable_thinking": False},
+        {"chat_template": None},
+        {"chat_template": "{{ 'MARKER-' ~ (6 * 7) ~ '-EVAL' }}"},
+        {"chat_template": ""},
+    ],
+)
+def test_systemone_enforces_request_template_trust(decision_server, template_kwargs):
+    server = decision_server
+    body = decision_body("systemone") | {"chat_template_kwargs": template_kwargs}
+    untrusted = (
+        not server.renderer.trust_request_chat_template
+        and (template_kwargs or {}).get("chat_template") is not None
+    )
+    # Retain the reply framing to exercise a usable template override.
+    if (template_kwargs or {}).get("chat_template"):
+        body["chat_template_kwargs"] = {
+            "chat_template": "{{ 'MARKER-' ~ (6 * 7) ~ '-EVAL' }}"
+            + server.tokenizer.chat_template,
+        }
+    response = server.client.post("/v1/systemone", json=body)
+    invalid = untrusted or template_kwargs == {"chat_template": ""}
+    assert response.status_code == (400 if invalid else 200), response.text
+    if untrusted:
+        assert "--trust-request-chat-template" in response.json()["error"]["message"]
+        assert not server.engine.admitted
+        assert not server.engine.generated
+    elif response.status_code == 200 and (template_kwargs or {}).get("chat_template"):
+        prompt = server.tokenizer.decode(server.engine.generated[0]["prompt_token_ids"])
+        assert prompt.startswith("MARKER-42-EVAL")
+
+
 @pytest.mark.parametrize("endpoint", ["decisions", "systemone"])
 @pytest.mark.parametrize("count,status", [(64, 200), (65, 400)])
 def test_model_question_limit_is_enforced_before_admission(
diff --git a/tests/entrypoints/generate/structured_decisions/test_structured_decisions_e2e.py b/tests/entrypoints/generate/structured_decisions/test_structured_decisions_e2e.py
index a206c56ccc06..73280b7362b9 100644
--- a/tests/entrypoints/generate/structured_decisions/test_structured_decisions_e2e.py
+++ b/tests/entrypoints/generate/structured_decisions/test_structured_decisions_e2e.py
@@ -44,6 +44,22 @@ def choice(instructions, *names):
 }
 
 
+def test_request_template_is_rejected_by_default(server):
+    response = post(
+        server,
+        {
+            "model": MODEL_NAME,
+            "state": "Team: billing.",
+            "questions": {"team": choice("Which team?", "billing")},
+            "chat_template_kwargs": {
+                "chat_template": "{{ 'MARKER-' ~ (6 * 7) ~ '-EVAL' }}"
+            },
+        },
+    )
+    assert response.status_code == 400, response.text
+    assert "--trust-request-chat-template" in response.json()["error"]["message"]
+
+
 def test_choice_decision(server):
     response = post(
         server,
diff --git a/tests/entrypoints/launchers/test_grpc_shutdown.py b/tests/entrypoints/launchers/test_grpc_shutdown.py
new file mode 100644
index 000000000000..df252eadadc5
--- /dev/null
+++ b/tests/entrypoints/launchers/test_grpc_shutdown.py
@@ -0,0 +1,219 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Shutdown sequencing of the gRPC launcher (``_drain_and_stop``)."""
+
+import asyncio
+import contextlib
+import threading
+import time
+
+import pytest
+
+grpc = pytest.importorskip("grpc")
+pytest.importorskip("smg_grpc_servicer")
+
+from vllm.entrypoints.launchers import grpc_server  # noqa: E402
+
+
+class FakeAsyncLLM:
+    """``shutdown(timeout)`` blocks like EngineCore's drain mode: it returns
+    once no request is in flight, or counts as killed when the timeout passes
+    first. Like EngineCore, it rejects a request it reads after shutdown has
+    started, even one that was added before. Like AsyncLLM, it only counts a
+    request once preprocessing is done."""
+
+    def __init__(self):
+        self.lock = threading.Lock()
+        self.inflight = 0
+        self.rejected = 0
+        self.started = None
+        self.timeout = "not called"
+        self.outcome = None
+
+    def get_num_unfinished_requests(self):
+        with self.lock:
+            return self.inflight
+
+    def shutdown(self, timeout=None):
+        self.started = time.monotonic()
+        self.timeout = timeout
+        deadline = self.started + (5.0 if timeout is None else timeout)
+        while time.monotonic() < deadline:
+            with self.lock:
+                if self.inflight == 0:
+                    self.outcome = "drained"
+                    return
+            time.sleep(0.01)
+        self.outcome = "killed"
+
+
+def make_server(llm):
+    async def generate(request, context):
+        # "<tokens>[:<s queued in EngineCore>[:<s of preprocessing>]]"
+        tokens, queued, preprocess = (request.decode().split(":") + ["0", "0"])[:3]
+        # Tokenization or media processing, before AsyncLLM knows the request.
+        await asyncio.sleep(float(preprocess))
+        with llm.lock:
+            llm.inflight += 1
+        try:
+            # The request can wait in EngineCore's input queue, e.g. behind a
+            # long prefill step, before EngineCore reads it.
+            await asyncio.sleep(float(queued))
+            with llm.lock:
+                if llm.started is not None:
+                    llm.rejected += 1
+                    return
+            for _ in range(int(tokens)):
+                await asyncio.sleep(0.01)
+                yield b"t"
+        finally:  # completion, or the abort a cancelled RPC sends
+            with llm.lock:
+                llm.inflight -= 1
+
+    async def subscribe(request, context):
+        # Open until the client or the server ends it.
+        yield b"e"
+        await asyncio.Event().wait()
+
+    inflight = grpc_server._InflightRpcs("test.Svc")
+    server = grpc.aio.server(interceptors=[inflight])
+    handlers = {
+        "Generate": grpc.unary_stream_rpc_method_handler(generate),
+        "SubscribeKvEvents": grpc.unary_stream_rpc_method_handler(subscribe),
+    }
+    server.add_generic_rpc_handlers(
+        (grpc.method_handlers_generic_handler("test.Svc", handlers),)
+    )
+    port = server.add_insecure_port("127.0.0.1:0")
+    return server, port, inflight
+
+
+async def call(port, method, arg=b""):
+    """Run one streaming RPC; return the number of messages and the status."""
+    async with grpc.aio.insecure_channel(f"127.0.0.1:{port}") as channel:
+        n = 0
+        try:
+            async for _ in channel.unary_stream(f"/test.Svc/{method}")(arg):
+                n += 1
+            return n, grpc.StatusCode.OK
+        except grpc.aio.AioRpcError as e:
+            return n, e.code()
+
+
+async def shut_down(rpcs, shutdown_timeout):
+    """Start the RPCs, then time ``_drain_and_stop`` as SIGTERM would run it."""
+    llm = FakeAsyncLLM()
+    server, port, inflight = make_server(llm)
+    await server.start()
+    clients = [asyncio.create_task(call(port, *rpc)) for rpc in rpcs]
+    await asyncio.sleep(0.3)
+    t0 = time.monotonic()
+    await grpc_server._drain_and_stop(server, llm, shutdown_timeout, inflight)
+    elapsed = time.monotonic() - t0
+    return llm, t0, elapsed, await asyncio.gather(*clients)
+
+
+@pytest.mark.asyncio
+async def test_inflight_count_follows_rpcs():
+    llm = FakeAsyncLLM()
+    server, port, inflight = make_server(llm)
+    await server.start()
+    try:
+        subscription = asyncio.create_task(call(port, "SubscribeKvEvents"))
+        rpc = asyncio.create_task(call(port, "Generate", b"50"))
+        await asyncio.sleep(0.2)
+        assert inflight.count == 1  # the subscription does not count
+        assert await rpc == (50, grpc.StatusCode.OK)
+        await asyncio.sleep(0.1)
+        assert inflight.count == 0
+        # A cancelled RPC stops counting too.
+        rpc = asyncio.create_task(call(port, "Generate", b"1000"))
+        await asyncio.sleep(0.2)
+        assert inflight.count == 1
+        rpc.cancel()
+        await asyncio.sleep(0.2)
+        assert inflight.count == 0
+        subscription.cancel()
+        for task in (rpc, subscription):
+            with contextlib.suppress(asyncio.CancelledError):
+                await task
+    finally:
+        await server.stop(None)
+
+
+@pytest.mark.asyncio
+async def test_engine_keeps_serving_until_requests_finish():
+    llm, t0, elapsed, results = await shut_down([("Generate", b"100")], 10)
+    # EngineCore is told to stop only once the request is done, and gets the
+    # rest of the budget plus the teardown floor.
+    assert llm.started - t0 > 0.5
+    remaining = 10 - (llm.started - t0)
+    assert llm.timeout == pytest.approx(
+        remaining + grpc_server._ENGINE_TEARDOWN_S, abs=0.1
+    )
+    assert llm.outcome == "drained"
+    assert results == [(100, grpc.StatusCode.OK)]
+    assert elapsed < 3
+
+
+@pytest.mark.asyncio
+async def test_requests_queued_at_sigterm_are_served():
+    # Added before SIGTERM, read by EngineCore 0.2 s after it.
+    rpcs = [("Generate", b"50:0.5") for _ in range(8)]
+    llm, _, _, results = await shut_down(rpcs, 10)
+    assert llm.rejected == 0
+    assert results == [(50, grpc.StatusCode.OK)] * 8
+
+
+@pytest.mark.asyncio
+async def test_request_in_preprocessing_at_sigterm_is_served():
+    # Accepted before SIGTERM, still preprocessing 0.3 s after it: AsyncLLM
+    # has no request then, but the RPC is in flight.
+    llm, _, _, results = await shut_down([("Generate", b"50:0:0.6")], 10)
+    assert llm.rejected == 0
+    assert results == [(50, grpc.StatusCode.OK)]
+
+
+@pytest.mark.asyncio
+async def test_open_rpc_without_request_does_not_hold_shutdown(monkeypatch):
+    monkeypatch.setattr(grpc_server, "_RPC_FLUSH_S", 0.5)
+    llm, t0, elapsed, results = await shut_down([("SubscribeKvEvents",)], 30)
+    # EngineCore stops right away, and the open RPC is cancelled after the
+    # flush instead of holding shutdown for the whole grace period.
+    assert llm.started - t0 < 0.5
+    assert llm.outcome == "drained"
+    assert 0.5 <= elapsed < 2
+    assert results[0][1] != grpc.StatusCode.OK
+
+
+@pytest.mark.asyncio
+async def test_rpc_over_budget_aborts_before_engine_deadline(monkeypatch):
+    monkeypatch.setattr(grpc_server, "_RPC_GRACE_S", 0.5)
+    llm, _, elapsed, results = await shut_down([("Generate", b"1000")], 1)
+    # The RPC is cancelled at the deadline and aborts its request; EngineCore
+    # still gets the teardown floor.
+    n, code = results[0]
+    assert n < 1000 and code != grpc.StatusCode.OK
+    assert llm.timeout == pytest.approx(grpc_server._ENGINE_TEARDOWN_S, abs=0.2)
+    assert llm.outcome == "drained"
+    assert 1 <= elapsed < 1 + grpc_server._ENGINE_TEARDOWN_S
+
+
+@pytest.mark.asyncio
+async def test_short_timeout_keeps_default_rpc_grace(monkeypatch):
+    monkeypatch.setattr(grpc_server, "_RPC_GRACE_S", 2.0)
+    llm, _, elapsed, results = await shut_down([("Generate", b"1000")], 1)
+    # A 1 s timeout still gives the RPC the 2 s it gets with 0, instead of
+    # cutting it at 1 s (about 130 tokens in).
+    n, code = results[0]
+    assert 200 < n < 1000 and code != grpc.StatusCode.OK
+    assert 2 <= elapsed < 2 + grpc_server._ENGINE_TEARDOWN_S
+    assert llm.timeout == pytest.approx(grpc_server._ENGINE_TEARDOWN_S, abs=0.2)
+
+
+@pytest.mark.asyncio
+async def test_zero_timeout_keeps_abort_behavior():
+    llm, _, _, _ = await shut_down([], 0)
+    # RPCs get the fixed grace and EngineCore its default join, as before.
+    assert llm.timeout is None
+    assert llm.outcome == "drained"
diff --git a/tests/entrypoints/launchers/test_run_batch.py b/tests/entrypoints/launchers/test_run_batch.py
index 04e6e71e70fc..0484ffb326db 100644
--- a/tests/entrypoints/launchers/test_run_batch.py
+++ b/tests/entrypoints/launchers/test_run_batch.py
@@ -13,6 +13,8 @@
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
+from fastapi.responses import JSONResponse
+from pydantic import TypeAdapter
 
 import vllm.envs as envs
 from vllm.assets.audio import AudioAsset
@@ -21,6 +23,7 @@
 from vllm.entrypoints.launchers import run_batch as run_batch_module
 from vllm.entrypoints.launchers.run_batch import (
     BatchProgressTracker,
+    BatchRequestInput,
     BatchRequestOutput,
     BatchTranscriptionRequest,
     batch_output_writer,
@@ -31,13 +34,16 @@
     make_arg_parser,
     make_transcription_wrapper,
     open_batch_input,
+    run_request,
     upload_data,
     url_matches,
     validate_batch,
     validate_run_batch_args,
 )
+from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionResponse
+from vllm.entrypoints.openai.completion.protocol import CompletionResponse
 from vllm.entrypoints.serve.engine.protocol import ErrorResponse
-from vllm.exceptions import VLLMValidationError
+from vllm.exceptions import RetryableRequestError, VLLMValidationError
 from vllm.utils.argparse_utils import FlexibleArgumentParser
 from vllm.utils.mem_constants import MiB_bytes
 from vllm.v1.engine.exceptions import EngineDeadError
@@ -136,6 +142,32 @@
     ]
 )
 
+INPUT_COMPLETIONS_BATCH = "\n".join(
+    json.dumps(req)
+    for req in [
+        {
+            "custom_id": "request-1",
+            "method": "POST",
+            "url": "/v1/completions",
+            "body": {
+                "model": CHAT_MODEL_NAME,
+                "prompt": "Hello world!",
+                "max_tokens": 10,
+            },
+        },
+        {
+            "custom_id": "request-2",
+            "method": "POST",
+            "url": "/v1/completions",
+            "body": {
+                "model": CHAT_MODEL_NAME,
+                "prompt": ["The capital of France is", "One plus one is"],
+                "max_tokens": 10,
+            },
+        },
+    ]
+)
+
 INVALID_INPUT_BATCH = "\n".join(
     json.dumps(req)
     for req in [
@@ -525,6 +557,52 @@ def test_completions():
             BatchRequestOutput.model_validate_json(line)
 
 
+def test_completions_endpoint():
+    with (
+        tempfile.NamedTemporaryFile("w") as input_file,
+        tempfile.NamedTemporaryFile("r") as output_file,
+    ):
+        # A chat line rides along: /v1/chat/completions and /v1/completions
+        # share their last URL segment, so each must reach its own handler.
+        chat_line = json.dumps(
+            json.loads(INPUT_BATCH.strip().split("\n")[0]) | {"custom_id": "chat-1"}
+        )
+        input_file.write(INPUT_COMPLETIONS_BATCH + "\n" + chat_line)
+        input_file.flush()
+        proc = subprocess.Popen(
+            [
+                "vllm",
+                "run-batch",
+                "-i",
+                input_file.name,
+                "-o",
+                output_file.name,
+                "--model",
+                CHAT_MODEL_NAME,
+            ],
+        )
+        proc.communicate()
+        proc.wait()
+        assert proc.returncode == 0, f"{proc=}"
+
+        contents = output_file.read()
+        outputs = {
+            output.custom_id: output
+            for output in map(
+                BatchRequestOutput.model_validate_json, contents.strip().split("\n")
+            )
+        }
+        assert len(outputs) == 3
+        for output in outputs.values():
+            assert output.error is None
+        assert isinstance(outputs["request-1"].response.body, CompletionResponse)
+        # A list prompt returns one choice per prompt.
+        completion = outputs["request-2"].response.body
+        assert isinstance(completion, CompletionResponse)
+        assert len(completion.choices) == 2
+        assert isinstance(outputs["chat-1"].response.body, ChatCompletionResponse)
+
+
 def test_completions_invalid_input():
     """Ensure that we fail when the input doesn't conform to the openai api."""
     with (
@@ -991,6 +1069,42 @@ async def test_download_bytes_backslash_bypass():
         )
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("is_translation", [False, True])
+@pytest.mark.parametrize(
+    "error, status_code, error_type",
+    [
+        (EngineDeadError(), 500, "InternalServerError"),
+        (RetryableRequestError("Engine is busy"), 503, "ServiceUnavailableError"),
+        (VLLMValidationError("Invalid audio"), 400, "BadRequestError"),
+    ],
+)
+async def test_batch_audio_preserves_handler_error_status(
+    is_translation, error, status_code, error_type
+):
+    operation = "translations" if is_translation else "transcriptions"
+    request = BatchRequestInput(
+        custom_id="audio-error",
+        method="POST",
+        url=f"/v1/audio/{operation}",
+        body={
+            "model": SPEECH_LARGE_MODEL_NAME,
+            "file_url": "data:audio/wav;base64,AA==",
+        },
+    )
+    handler = AsyncMock(side_effect=error)
+    wrapped_handler = make_transcription_wrapper(is_translation)(handler)
+
+    output = await run_request(wrapped_handler, request)
+
+    assert output.response is not None
+    assert output.response.status_code == status_code
+    assert output.error is not None
+    assert output.error.error.code == status_code
+    assert output.error.error.type == error_type
+    handler.assert_awaited_once()
+
+
 @pytest.mark.asyncio
 async def test_transcription_wrapper_rejects_oversized_data_url_before_decode(
     monkeypatch: pytest.MonkeyPatch,
@@ -1177,6 +1291,43 @@ def test_validate_batch_rejects_malformed_request(tmp_path):
         validate_batch(f)
 
 
+@pytest.mark.asyncio
+async def test_requests_and_responses_reuse_type_adapters(monkeypatch):
+    """Building a TypeAdapter rebuilds its schema, too slow to do per request."""
+    type_adapter = MagicMock(wraps=TypeAdapter)
+    monkeypatch.setattr(run_batch_module, "TypeAdapter", type_adapter)
+
+    fallback_line = json.dumps(
+        {
+            "custom_id": "request-1",
+            "method": "POST",
+            "url": "/v1/unknown",
+            "body": {"model": EMBEDDING_MODEL_NAME, "input": "Hello world!"},
+        }
+    )
+    lines = [
+        *INPUT_EMBEDDING_BATCH.split("\n"),
+        *INPUT_SCORE_BATCH.split("\n"),
+        fallback_line,
+    ]
+    requests = [BatchRequestInput.model_validate_json(line) for line in lines]
+
+    async def handler(_body):
+        error = {"message": "bad", "type": "BadRequestError", "code": 400}
+        return JSONResponse(content={"error": error}, status_code=400)
+
+    output = await run_request(handler, requests[0])
+
+    assert [type(r.body).__name__ for r in requests] == [
+        *["EmbeddingCompletionRequest"] * 4,
+        *["ScoreQueriesDocumentsRequest"] * 2,
+        "EmbeddingCompletionRequest",
+    ]
+    assert isinstance(output.error, ErrorResponse)
+    assert output.response.status_code == 400
+    type_adapter.assert_not_called()
+
+
 @pytest.mark.asyncio
 async def test_open_batch_input_keeps_a_carriage_return_inside_a_line(tmp_path):
     """A bare carriage return is JSON whitespace, not a line break."""
diff --git a/tests/entrypoints/multimodal/openai/chat_completion/test_vision.py b/tests/entrypoints/multimodal/openai/chat_completion/test_vision.py
index 6438d50a7808..287def72af7f 100644
--- a/tests/entrypoints/multimodal/openai/chat_completion/test_vision.py
+++ b/tests/entrypoints/multimodal/openai/chat_completion/test_vision.py
@@ -3,13 +3,18 @@
 
 import json
 
+import httpx
 import openai
 import pytest
 import pytest_asyncio
+from PIL import Image
 from transformers import AutoProcessor
 
 from tests.entrypoints.multimodal.conftest import TEST_IMAGE_ASSETS
 from tests.utils import ROCM_EXTRA_ARGS, RemoteOpenAIServer
+from vllm.config import ModelConfig
+from vllm.multimodal import MULTIMODAL_REGISTRY
+from vllm.multimodal.cache.base import MultiModalCache
 from vllm.multimodal.media import MediaWithBytes
 from vllm.multimodal.utils import encode_image_url, fetch_image
 from vllm.platforms import current_platform
@@ -17,6 +22,78 @@
 MODEL_NAME = "microsoft/Phi-3.5-vision-instruct"
 MAXIMUM_IMAGES = 2
 
+
+def test_processor_only_cache_eviction_preserves_chat_logprobs():
+    """A later image hit must survive inserts on the same API worker."""
+    model = "llava-hf/llava-1.5-7b-hf"
+    images = [Image.new("RGB", (112, 112), c) for c in ("red", "green", "blue")]
+    processor = MULTIMODAL_REGISTRY.create_processor(
+        ModelConfig(model, limit_mm_per_prompt={"image": 3})
+    )
+    processed = processor(
+        "<image>" * 3,
+        mm_items=processor.info.parse_mm_data({"image": images}),
+    )
+    items = processed["mm_kwargs"]["image"]
+    capacity = MultiModalCache.get_item_size(items[0]) + MultiModalCache.get_item_size(
+        items[1]
+    )
+    urls = [encode_image_url(image) for image in images]
+    args = [
+        "--api-server-count",
+        "2",
+        "--mm-processor-cache-gb",
+        str(capacity / (1 << 30)),
+        "--limit-mm-per-prompt",
+        '{"image": 3}',
+        "--max-model-len",
+        "4096",
+        "--max-num-seqs",
+        "1",
+        "--no-enable-prefix-caching",
+        "--enforce-eager",
+    ]
+    # Keep all requests on one API worker, whose cache holds exactly two images.
+    limits = httpx.Limits(max_connections=1, max_keepalive_connections=1)
+    with (
+        RemoteOpenAIServer(model, args) as server,
+        server.get_client(http_client=httpx.Client(limits=limits)) as client,
+    ):
+
+        def complete(order):
+            return client.chat.completions.create(
+                model=model,
+                temperature=0,
+                max_tokens=1,
+                logprobs=True,
+                top_logprobs=5,
+                messages=[
+                    {
+                        "role": "user",
+                        "content": [
+                            {"type": "text", "text": "Describe these images."},
+                            *[
+                                {"type": "image_url", "image_url": {"url": urls[i]}}
+                                for i in order
+                            ],
+                        ],
+                    }
+                ],
+            )
+
+        expected = complete((2, 0, 1))
+        complete((0, 1))
+        actual = complete((2, 0, 1))
+        assert actual.usage.prompt_tokens == expected.usage.prompt_tokens
+        assert actual.choices[0].message.content
+        expected_probs = expected.choices[0].logprobs.content[0].top_logprobs
+        actual_probs = actual.choices[0].logprobs.content[0].top_logprobs
+        assert [p.token for p in actual_probs] == [p.token for p in expected_probs]
+        assert [p.logprob for p in actual_probs] == pytest.approx(
+            [p.logprob for p in expected_probs], abs=1e-3
+        )
+
+
 # Required terms for beam search validation
 # Each entry is a list of term groups - ALL groups must match
 # Each group is a list of alternatives - at least ONE term in the group must appear
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 8cc33c4939f2..8fc91225ba80 100644
--- a/tests/entrypoints/openai/chat_completion/test_batched_chat_completions.py
+++ b/tests/entrypoints/openai/chat_completion/test_batched_chat_completions.py
@@ -48,6 +48,10 @@ def default_server_args():
         "2048",
         "--max-num-seqs",
         "128",
+        "--enable-prompt-tokens-details",
+        "--enable-prefix-caching",
+        "--block-size",
+        "16",
         "--enforce-eager",
     ]
 
@@ -251,6 +255,39 @@ async def test_batched_chat_completions_logprob_token_ids(
     }
 
 
+@pytest.mark.asyncio
+async def test_batched_chat_completions_prompt_tokens_details(
+    server: RemoteOpenAIServer,
+) -> None:
+    """Batch cache counts include both conversations."""
+    conversation = [{"role": "user", "content": "Repeat the word alpha. " * 20}]
+    body = {"model": MODEL_NAME, "max_tokens": 1, "temperature": 0}
+    usages = []
+
+    async with httpx.AsyncClient() as http_client:
+        # Warm the prefix, measure one conversation, then measure a batch of two.
+        for endpoint, messages in (
+            ("v1/chat/completions", conversation),
+            ("v1/chat/completions", conversation),
+            ("v1/chat/completions/batch", [conversation] * 2),
+        ):
+            response = await http_client.post(
+                server.url_for(endpoint),
+                json={**body, "messages": messages},
+                timeout=60,
+            )
+            assert response.status_code == 200, response.text
+            usages.append(response.json()["usage"])
+
+    single_details = usages[1]["prompt_tokens_details"]
+    batch_details = usages[2]["prompt_tokens_details"]
+    assert single_details is not None
+    assert single_details["cached_tokens"] > 0
+    assert batch_details is not None
+    for field in ("cached_tokens", "created_cache_tokens"):
+        assert batch_details[field] == 2 * single_details[field]
+
+
 def _make_request_output(
     prompt_idx: int,
     text: str,
@@ -295,6 +332,7 @@ async def test_batched_echo_does_not_prepend_user_prompt() -> None:
     serving = OpenAIServingChatBatch.__new__(OpenAIServingChatBatch)
     serving.response_role = "assistant"
     serving.system_fingerprint = None
+    serving.enable_prompt_tokens_details = False
 
     conversations = [[{"role": "user", "content": "USER PROMPT"}]]
     request = BatchChatCompletionRequest(
@@ -324,6 +362,7 @@ async def test_batched_echo_prepends_matching_assistant_prefix() -> None:
     serving = OpenAIServingChatBatch.__new__(OpenAIServingChatBatch)
     serving.response_role = "assistant"
     serving.system_fingerprint = None
+    serving.enable_prompt_tokens_details = False
 
     conversations = [
         [
@@ -438,6 +477,7 @@ def parse(self, model_output, **kwargs):
     serving = OpenAIServingChatBatch.__new__(OpenAIServingChatBatch)
     serving.response_role = "assistant"
     serving.system_fingerprint = None
+    serving.enable_prompt_tokens_details = False
     serving.chat_template = None
     serving.chat_template_content_format = "auto"  # type: ignore[misc]  # Final attr
     serving.default_chat_template_kwargs = {}
diff --git a/tests/entrypoints/openai/chat_completion/test_chat_completion_with_prompt_embeds.py b/tests/entrypoints/openai/chat_completion/test_chat_completion_with_prompt_embeds.py
index 80b512699a33..0696b3ac70d5 100644
--- a/tests/entrypoints/openai/chat_completion/test_chat_completion_with_prompt_embeds.py
+++ b/tests/entrypoints/openai/chat_completion/test_chat_completion_with_prompt_embeds.py
@@ -12,8 +12,11 @@
 import pytest_asyncio
 import torch
 from openai import BadRequestError
+from PIL import Image
+from transformers import AutoConfig
 
 from tests.utils import VLLM_PATH, RemoteOpenAIServer
+from vllm.multimodal.utils import encode_image_url
 from vllm.platforms import current_platform
 
 MODEL_NAME = "facebook/opt-125m"
@@ -71,6 +74,59 @@ def _encode_embeds(embeds: torch.Tensor) -> str:
     return base64.b64encode(buf.getvalue()).decode("utf-8")
 
 
+@pytest.mark.asyncio
+async def test_prompt_embeds_with_image() -> None:
+    """An unpadded vocabulary must accept mixed embeddings and image inputs."""
+    model = "HuggingFaceTB/SmolVLM-256M-Instruct"
+    config = AutoConfig.from_pretrained(model)
+    embeds = _encode_embeds(
+        torch.zeros(1, config.text_config.hidden_size, dtype=SERVER_DTYPE)
+    )
+    with RemoteOpenAIServer(
+        model,
+        [
+            "--enable-prompt-embeds",
+            "--dtype",
+            "bfloat16",
+            "--max-model-len",
+            "2048",
+            "--gpu-memory-utilization",
+            "0.2",
+            "--enforce-eager",
+        ],
+    ) as remote_server:
+        async with remote_server.get_async_client() as client:
+            response = await client.chat.completions.create(
+                model=model,
+                temperature=0,
+                max_tokens=1,
+                messages=[
+                    {
+                        "role": "user",
+                        "content": [
+                            {"type": "prompt_embeds", "data": embeds},
+                            {
+                                "type": "image_url",
+                                "image_url": {
+                                    "url": encode_image_url(Image.new("RGB", (32, 32)))
+                                },
+                            },
+                            {"type": "text", "text": "Describe the image."},
+                        ],
+                    }
+                ],
+            )
+            assert response.choices[0].finish_reason in ("length", "stop")
+            assert response.usage is not None
+            assert response.usage.completion_tokens == 1
+            with pytest.raises(BadRequestError, match="out of vocabulary"):
+                await client.completions.create(
+                    model=model,
+                    prompt=[config.text_config.vocab_size + 1],
+                    max_tokens=1,
+                )
+
+
 @pytest.fixture(scope="module")
 def prompt_embeds_b64(hf_runner) -> list[str]:
     """Pre-compute embeddings for two short prompts and return as base64."""
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 ecff3f8ee6be..e0aafae13edf 100644
--- a/tests/entrypoints/openai/chat_completion/test_thinking_token_budget.py
+++ b/tests/entrypoints/openai/chat_completion/test_thinking_token_budget.py
@@ -14,6 +14,7 @@
 import pytest
 import pytest_asyncio
 
+import vllm.envs as envs
 from tests.utils import RemoteOpenAIServer, multi_gpu_only, requires_fp8
 from vllm.platforms import current_platform
 from vllm.tokenizers import get_tokenizer
@@ -212,6 +213,39 @@ async def test_thinking_token_budget_limits_reasoning(client: openai.AsyncOpenAI
     )
 
 
+@pytest.mark.asyncio
+@pytest.mark.skipif(
+    envs.VLLM_USE_RUST_FRONTEND,
+    reason="the Rust frontend does not serve /v1/responses",
+)
+@pytest.mark.parametrize("client", ["default"], indirect=True)
+async def test_thinking_token_budget_responses_api(client: openai.AsyncOpenAI):
+    """/v1/responses applies thinking_token_budget like /v1/chat/completions."""
+    extra_body = {"thinking_token_budget": THINK_BUDGET}
+    chat = await client.chat.completions.create(
+        model=MODEL_NAME,
+        messages=MESSAGES,
+        max_tokens=100,
+        temperature=0.0,
+        extra_body=extra_body,
+    )
+    response = await client.responses.create(
+        model=MODEL_NAME,
+        input=MESSAGES,
+        max_output_tokens=100,
+        temperature=0.0,
+        extra_body=extra_body,
+    )
+
+    reasoning = "".join(
+        part.text
+        for item in response.output
+        if item.type == "reasoning"
+        for part in item.content or []
+    )
+    assert reasoning == chat.choices[0].message.reasoning
+
+
 @pytest.mark.asyncio
 @multi_gpu_only(num_gpus=2)
 @requires_fp8
diff --git a/tests/entrypoints/openai/completion/test_completion_text_offset.py b/tests/entrypoints/openai/completion/test_completion_text_offset.py
index 5c66b81ff7fc..a1205b442fd1 100644
--- a/tests/entrypoints/openai/completion/test_completion_text_offset.py
+++ b/tests/entrypoints/openai/completion/test_completion_text_offset.py
@@ -97,3 +97,19 @@ async def results():
     assert response.choices[0].logprobs is not None
     assert response.choices[0].logprobs.text_offset == [0, 3, 6]
     assert stream_offsets == [0, 3, 6]
+
+
+def test_non_stream_response_omits_null_prompt_tokens_details():
+    """The OpenAI spec types ``usage.prompt_tokens_details`` as an object, not
+    nullable, so it is omitted when ``--enable-prompt-tokens-details`` is off."""
+    response = _serving().request_output_to_completion_response(
+        [_output("abc", [0], finished=True)],
+        _request(stream=False),
+        "r",
+        0,
+        "m",
+        None,
+        RequestResponseMetadata(request_id="r"),
+    )
+
+    assert "prompt_tokens_details" not in response.model_dump()["usage"]
diff --git a/tests/entrypoints/openai/parser/test_harmony_utils.py b/tests/entrypoints/openai/parser/test_harmony_utils.py
index 4c1dfe52e033..3fddd9488515 100644
--- a/tests/entrypoints/openai/parser/test_harmony_utils.py
+++ b/tests/entrypoints/openai/parser/test_harmony_utils.py
@@ -938,6 +938,9 @@ def test_drops_non_assistant_analysis_messages(self) -> None:
         assert cleaned_messages == messages[1:]
 
 
+pytestmark = pytest.mark.skip_global_cleanup
+
+
 def test_has_custom_tools() -> None:
     assert not has_custom_tools(set())
     assert not has_custom_tools({"web_search_preview", "code_interpreter", "container"})
@@ -1029,8 +1032,8 @@ def test_reasoning_with_multiple_content_items(self):
         )
         assert msg.channel == "analysis"
 
-    def test_reasoning_without_content_returns_none(self):
-        """Test reasoning item without content field returns None."""
+    def test_reasoning_without_content_returns_empty_message(self):
+        """Test reasoning item without content field returns empty message."""
         item = {
             "type": "reasoning",
             "id": "rs_123",
@@ -1039,10 +1042,11 @@ def test_reasoning_without_content_returns_none(self):
 
         msg = response_input_to_harmony(item, prev_responses=[])
 
-        assert msg is None
+        assert msg is not None
+        assert msg.content[0].text == ""
 
-    def test_reasoning_with_none_content_returns_none(self):
-        """Test reasoning item with content=None returns None."""
+    def test_reasoning_with_none_content_returns_empty_message(self):
+        """Test reasoning item with content=None returns empty message."""
         item = {
             "type": "reasoning",
             "id": "rs_123",
@@ -1052,10 +1056,11 @@ def test_reasoning_with_none_content_returns_none(self):
 
         msg = response_input_to_harmony(item, prev_responses=[])
 
-        assert msg is None
+        assert msg is not None
+        assert msg.content[0].text == ""
 
-    def test_reasoning_with_empty_content_returns_none(self):
-        """Test reasoning item with empty content list returns None."""
+    def test_reasoning_with_empty_content_returns_empty_message(self):
+        """Test reasoning item with empty content list returns empty message."""
         item = {
             "type": "reasoning",
             "id": "rs_123",
@@ -1064,4 +1069,5 @@ def test_reasoning_with_empty_content_returns_none(self):
 
         msg = response_input_to_harmony(item, prev_responses=[])
 
-        assert msg is None
+        assert msg is not None
+        assert msg.content[0].text == ""
diff --git a/tests/entrypoints/openai/responses/test_function_call_parsing.py b/tests/entrypoints/openai/responses/test_function_call_parsing.py
index 5592901dd30e..2757dae2f8ee 100644
--- a/tests/entrypoints/openai/responses/test_function_call_parsing.py
+++ b/tests/entrypoints/openai/responses/test_function_call_parsing.py
@@ -9,6 +9,8 @@
 
 from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
 
+pytestmark = pytest.mark.skip_global_cleanup
+
 
 def test_function_call_dict_converted_to_object():
     """Test that function_call dictionaries are correctly parsed into
@@ -375,3 +377,80 @@ def test_assistant_output_style_content_coerced():
     assert item.content[0].annotations == []
     assert item.status == "completed"
     assert item.id.startswith("msg_")
+
+
+def _item_attr(item, key):
+    if isinstance(item, dict):
+        return item.get(key)
+    return getattr(item, key)
+
+
+def _as_list(value):
+    return value if isinstance(value, list) else list(value)
+
+
+def test_assistant_and_reasoning_items_without_id_status_are_accepted():
+    """Assistant/reasoning history should accept missing id/status fields."""
+    request_data = {
+        "model": "gpt-oss",
+        "input": [
+            {
+                "type": "message",
+                "role": "user",
+                "content": [{"type": "input_text", "text": "hello"}],
+            },
+            {
+                "type": "message",
+                "role": "assistant",
+                "content": [{"type": "output_text", "text": "hi"}],
+            },
+            {
+                "type": "reasoning",
+                "summary": [],
+                "encrypted_content": None,
+            },
+        ],
+    }
+
+    request = ResponsesRequest(**request_data)
+
+    assert len(request.input) == 3
+    assistant_item = request.input[1]
+    reasoning_item = request.input[2]
+
+    assert _item_attr(assistant_item, "id")
+    assert _item_attr(assistant_item, "status") == "completed"
+    assert _item_attr(reasoning_item, "id")
+    assert _item_attr(reasoning_item, "status") == "completed"
+    assert _as_list(_item_attr(reasoning_item, "summary")) == []
+
+
+def test_assistant_and_reasoning_keep_existing_id_status():
+    """Existing id/status should not be overwritten by normalization."""
+    request_data = {
+        "model": "gpt-oss",
+        "input": [
+            {
+                "type": "message",
+                "role": "assistant",
+                "id": "msg_existing",
+                "status": "in_progress",
+                "content": [{"type": "output_text", "text": "partial"}],
+            },
+            {
+                "type": "reasoning",
+                "id": "rs_existing",
+                "status": "incomplete",
+                "summary": [{"type": "summary_text", "text": "thinking"}],
+            },
+        ],
+    }
+
+    request = ResponsesRequest(**request_data)
+
+    assistant_item = request.input[0]
+    reasoning_item = request.input[1]
+    assert _item_attr(assistant_item, "id") == "msg_existing"
+    assert _item_attr(assistant_item, "status") == "in_progress"
+    assert _item_attr(reasoning_item, "id") == "rs_existing"
+    assert _item_attr(reasoning_item, "status") == "incomplete"
diff --git a/tests/entrypoints/openai/responses/test_protocol.py b/tests/entrypoints/openai/responses/test_protocol.py
index 04b07992528b..3207bb9417b2 100644
--- a/tests/entrypoints/openai/responses/test_protocol.py
+++ b/tests/entrypoints/openai/responses/test_protocol.py
@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import pytest
+from openai.types.responses.response_output_message import ResponseOutputMessage
 from openai_harmony import (
     Message,
 )
@@ -13,6 +14,8 @@
 )
 from vllm.exceptions import VLLMValidationError
 
+pytestmark = pytest.mark.skip_global_cleanup
+
 
 def test_serialize_message() -> None:
     dict_value = {"a": 1, "b": "2"}
@@ -43,6 +46,71 @@ def test_serialize_messages() -> None:
     assert serialize_messages([msg, dict_value]) == [msg_value, dict_value]
 
 
+def test_normalize_openresponses_input_item_stamps_type() -> None:
+    """A typeless assistant item must keep its role and content.
+
+    ``type`` is optional in the OpenResponses schema, but filling in
+    ``id``/``status`` without it makes the item match ``ItemReference`` in the
+    input union, which collapses it to ``{"id": ...}`` and silently drops the
+    assistant turn from the conversation. Regression test for that.
+    """
+    item = {
+        "role": "assistant",
+        "content": [{"type": "output_text", "text": "hi"}],
+    }
+
+    normalized = ResponsesRequest._normalize_openresponses_input_item(item)
+
+    assert normalized["type"] == "message"
+    assert normalized["role"] == "assistant"
+    assert normalized["id"].startswith("msg_")
+    assert normalized["status"] == "completed"
+    # annotations are filled in for output_text content
+    assert normalized["content"][0]["annotations"] == []
+    # the caller's dict is not mutated
+    assert "type" not in item
+
+
+def test_normalize_openresponses_input_item_reasoning_type() -> None:
+    """An explicit reasoning item keeps its type and gets a summary default."""
+    normalized = ResponsesRequest._normalize_openresponses_input_item(
+        {"type": "reasoning", "content": [{"type": "reasoning_text", "text": "why"}]}
+    )
+
+    assert normalized["type"] == "reasoning"
+    assert normalized["id"].startswith("rs_")
+    assert normalized["status"] == "completed"
+    assert normalized["summary"] == []
+
+
+def test_input_item_parsing_typeless_assistant_message() -> None:
+    """End to end: a typeless assistant item survives as a real message."""
+    request = ResponsesRequest.model_validate(
+        {
+            "input": [
+                {
+                    "role": "assistant",
+                    "content": [{"type": "output_text", "text": "hi"}],
+                }
+            ]
+        }
+    )
+
+    assert isinstance(request.input, list)
+    (parsed,) = request.input
+    assert isinstance(parsed, ResponseOutputMessage)
+    assert parsed.role == "assistant"
+    assert parsed.content[0].text == "hi"
+
+
+def test_input_item_parsing_leaves_non_assistant_items_alone() -> None:
+    """User messages are untouched — normalization only targets assistant items."""
+    item = {"role": "user", "content": "hello"}
+    request = ResponsesRequest.model_validate({"input": [item]})
+
+    assert request.input == [item]
+
+
 def test_custom_tool_grammar_format_rejected() -> None:
     grammar = {"type": "grammar", "syntax": "lark", "definition": 'start: "pwd"'}
     with pytest.raises(VLLMValidationError) as exc_info:
diff --git a/tests/entrypoints/openai/responses/test_responses_utils.py b/tests/entrypoints/openai/responses/test_responses_utils.py
index 12e6c96ecb0a..c5a44ee62cd4 100644
--- a/tests/entrypoints/openai/responses/test_responses_utils.py
+++ b/tests/entrypoints/openai/responses/test_responses_utils.py
@@ -29,6 +29,8 @@
 )
 from vllm.exceptions import VLLMValidationError
 
+pytestmark = pytest.mark.skip_global_cleanup
+
 
 def _single_chat_message(item):
     message = _construct_message_from_response_item(item)
@@ -253,6 +255,19 @@ def test_construct_chat_messages_preserves_single_item_conversions(self):
         assert formatted_item["role"] == "assistant"
         assert formatted_item["content"] == "dongyi"
 
+        # Assistant messages with empty content (e.g. tool-only turns) should
+        # not crash; they map to empty string content.
+        empty_output_item = ResponseOutputMessage(
+            id="msg_empty_content",
+            content=[],
+            role="assistant",
+            status="completed",
+            type="message",
+        )
+        formatted_item = _single_chat_message(empty_output_item)
+        assert formatted_item["role"] == "assistant"
+        assert formatted_item["content"] == ""
+
 
 class TestCustomTools:
     def test_custom_tool_call_output_item(self):
diff --git a/tests/entrypoints/openai/responses/test_sampling_params.py b/tests/entrypoints/openai/responses/test_sampling_params.py
index eaa3b4cedacf..e502d3db58bf 100644
--- a/tests/entrypoints/openai/responses/test_sampling_params.py
+++ b/tests/entrypoints/openai/responses/test_sampling_params.py
@@ -63,6 +63,7 @@ def test_extra_sampling_params(self):
             stop_token_ids=[42, 43],
             ignore_eos=True,
             spaces_between_special_tokens=False,
+            thinking_token_budget=128,
             vllm_xargs={"custom": "value"},
         )
 
@@ -74,6 +75,7 @@ def test_extra_sampling_params(self):
         assert sampling_params.stop_token_ids == [42, 43]
         assert sampling_params.ignore_eos is True
         assert sampling_params.spaces_between_special_tokens is False
+        assert sampling_params.thinking_token_budget == 128
         assert sampling_params.extra_args == {"custom": "value"}
 
     def test_stop_string_conversion(self):
@@ -99,6 +101,7 @@ def test_default_values(self):
 
         assert sampling_params.repetition_penalty == 1.0  # None → 1.0
         assert sampling_params.stop == []  # Empty list
+        assert sampling_params.thinking_token_budget is None
         assert sampling_params.extra_args == {}  # Empty dict
 
     def test_null_top_logprobs_is_same_as_omitted(self):
diff --git a/tests/entrypoints/openai/responses/test_simple.py b/tests/entrypoints/openai/responses/test_simple.py
index ba55a58e7d3c..b284aa982e27 100644
--- a/tests/entrypoints/openai/responses/test_simple.py
+++ b/tests/entrypoints/openai/responses/test_simple.py
@@ -12,6 +12,49 @@
 MODEL_NAME = "Qwen/Qwen3-8B"
 
 
+def _input_text_message(text: str) -> dict:
+    return {
+        "type": "message",
+        "role": "user",
+        "content": [{"type": "input_text", "text": text}],
+    }
+
+
+def _assistant_output_text_message(
+    text: str,
+    *,
+    item_id: str | None = None,
+    status: str | None = None,
+) -> dict:
+    message = {
+        "type": "message",
+        "role": "assistant",
+        "content": [{"type": "output_text", "text": text}],
+    }
+    if item_id is not None:
+        message["id"] = item_id
+    if status is not None:
+        message["status"] = status
+    return message
+
+
+def _reasoning_item(
+    *,
+    item_id: str | None = None,
+    status: str | None = None,
+) -> dict:
+    item: dict[str, object] = {
+        "type": "reasoning",
+        "summary": [],
+        "encrypted_content": None,
+    }
+    if item_id is not None:
+        item["id"] = item_id
+    if status is not None:
+        item["status"] = status
+    return item
+
+
 @pytest.fixture(scope="module")
 def server():
     from .conftest import BASE_TEST_ENV
@@ -94,6 +137,83 @@ async def test_reasoning_item(client: OpenAI, model_name: str):
     assert type(response.output[1].content[0].text) is str
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("model_name", [MODEL_NAME])
+async def test_input_history_openai_shape_e2e(client: OpenAI, model_name: str):
+    """E2E: OpenAI-style history with assistant/reasoning id+status is accepted."""
+    response = await client.responses.create(
+        model=model_name,
+        input=[
+            _input_text_message("You are a concise assistant."),
+            _assistant_output_text_message(
+                "Acknowledged.",
+                item_id="msg_hist_openai_1",
+                status="completed",
+            ),
+            _reasoning_item(item_id="rs_hist_openai_1", status="completed"),
+            _input_text_message("Reply with one short sentence."),
+        ],
+        temperature=0.0,
+    )
+
+    assert response is not None
+    assert response.status == "completed"
+    assert response.output_text is not None
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("model_name", [MODEL_NAME])
+async def test_input_history_openresponses_shape_e2e(client: OpenAI, model_name: str):
+    """E2E: OpenResponses-style history allows optional assistant/reasoning ids."""
+    response = await client.responses.create(
+        model=model_name,
+        input=[
+            _input_text_message("You are a concise assistant."),
+            _assistant_output_text_message("Acknowledged."),
+            _reasoning_item(),
+            _input_text_message("Reply with one short sentence."),
+        ],
+        temperature=0.0,
+    )
+
+    assert response is not None
+    assert response.status == "completed"
+    assert response.output_text is not None
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("model_name", [MODEL_NAME])
+async def test_multiturn_history_replay_without_ids_e2e(
+    client: OpenAI, model_name: str
+):
+    """E2E: second turn succeeds when replayed assistant/reasoning omit id/status."""
+    first_turn_prompt = "Remember this marker for the next turn: MARKER-729."
+    first_response = await client.responses.create(
+        model=model_name,
+        input=first_turn_prompt,
+        temperature=0.0,
+    )
+
+    assert first_response is not None
+    assert first_response.status == "completed"
+    assert first_response.output_text is not None
+
+    second_response = await client.responses.create(
+        model=model_name,
+        input=[
+            _input_text_message(first_turn_prompt),
+            _assistant_output_text_message(first_response.output_text),
+            _reasoning_item(),
+            _input_text_message("What marker did I ask you to remember?"),
+        ],
+        temperature=0.0,
+    )
+
+    assert second_response is not None
+    assert second_response.status == "completed"
+    assert second_response.output_text is not None
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize("model_name", [MODEL_NAME])
 async def test_streaming_output_consistency(client: OpenAI, model_name: str):
diff --git a/tests/entrypoints/openai/responses/test_streaming_events.py b/tests/entrypoints/openai/responses/test_streaming_events.py
index a37109f85134..d3c7cda85d4c 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
 
+import json
+import random
+
 import pytest
 from openai.types.responses import (
     CustomTool,
@@ -16,6 +19,7 @@
     DeltaMessage,
     DeltaToolCall,
 )
+from vllm.entrypoints.openai.responses import utils as responses_utils
 from vllm.entrypoints.openai.responses.streaming_events import (
     SimpleStreamingEventProcessor,
     StreamingState,
@@ -25,7 +29,10 @@
     emit_reasoning_done_events,
     split_delta,
 )
-from vllm.entrypoints.openai.responses.utils import decode_custom_tool_input
+from vllm.entrypoints.openai.responses.utils import (
+    decode_custom_tool_input,
+    decode_custom_tool_input_prefix,
+)
 
 
 def test_browser_find_uses_responses_action_type():
@@ -282,6 +289,62 @@ def test_unrelated_keys_are_not_streamed_as_payload(self):
         assert item.input == "hello"
         assert decode_custom_tool_input('{"a": "xy", "input": "hello"}') == "hello"
 
+    @staticmethod
+    def _stream_custom_tool_input(arguments: list[str]):
+        processor = SimpleStreamingEventProcessor(
+            tools=[CustomTool(type="custom", name="emit_command")]
+        )
+        _run_through_processor(
+            processor,
+            DeltaMessage(tool_calls=[_make_tool_call(0, name="emit_command")]),
+        )
+        for chunk in arguments:
+            _run_through_processor(
+                processor,
+                DeltaMessage(tool_calls=[_make_tool_call(0, arguments=chunk)]),
+            )
+            yield processor.state
+
+    @pytest.mark.parametrize("ensure_ascii", [False, True])
+    def test_streamed_payload_matches_prefix_decode(self, ensure_ascii):
+        rng = random.Random(0)
+        alphabet = ["a", " ", '"', "\\", "/", "\n", "\t", "\x01", "\x1f", "\x85"]
+        alphabet += ["\u00e9", "\u3000", "\U0001f600"]
+        for _ in range(200):
+            payload = "".join(rng.choices(alphabet, k=rng.randint(0, 40)))
+            arguments = json.dumps({"input": payload}, ensure_ascii=ensure_ascii)
+            cuts = sorted(rng.sample(range(1, len(arguments)), k=len(arguments) // 3))
+            chunks = [arguments[i:j] for i, j in zip([0, *cuts], [*cuts, None])]
+            for state in self._stream_custom_tool_input(chunks):
+                assert state.tool_call_payload == decode_custom_tool_input_prefix(
+                    state.accumulated_text
+                )
+
+    def test_payload_deltas_do_not_reparse_the_input(self, monkeypatch):
+        """Streaming cost must not grow with the input already streamed."""
+        loads = responses_utils.partial_json_parser.loads
+        calls = 0
+
+        def counting_loads(*args, **kwargs):
+            nonlocal calls
+            calls += 1
+            return loads(*args, **kwargs)
+
+        monkeypatch.setattr(
+            responses_utils.partial_json_parser, "loads", counting_loads
+        )
+
+        def full_decodes(lines: int) -> int:
+            nonlocal calls
+            calls = 0
+            arguments = json.dumps({"input": 'echo "a > b"\n' * lines})
+            chunks = [arguments[i : i + 4] for i in range(0, len(arguments), 4)]
+            for _ in self._stream_custom_tool_input(chunks):
+                pass
+            return calls
+
+        assert full_decodes(lines=10) == full_decodes(lines=100)
+
     def test_done_reported_authoritative_payload_on_mismatch(self):
         """If the final decode is not an extension of what was streamed, the
         done event and item carry the authoritative value and a correcting
diff --git a/tests/entrypoints/openai/test_tool_choice_content_none.py b/tests/entrypoints/openai/test_tool_choice_content_none.py
index f08915a93dca..3b37321d8a6c 100644
--- a/tests/entrypoints/openai/test_tool_choice_content_none.py
+++ b/tests/entrypoints/openai/test_tool_choice_content_none.py
@@ -23,7 +23,7 @@
     ChatMessage,
 )
 from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
-from vllm.entrypoints.serve.engine.protocol import UsageInfo
+from vllm.entrypoints.serve.engine.protocol import PromptTokenUsageInfo, UsageInfo
 from vllm.parser.abstract_parser import DelegatingParser
 
 pytestmark = pytest.mark.skip_global_cleanup
@@ -152,6 +152,31 @@ def test_chat_completion_response_keeps_non_empty_tool_calls_payload():
     assert message["tool_calls"][0]["function"]["name"] == "get_weather"
 
 
+def test_chat_completion_response_omits_null_annotations_and_prompt_tokens_details():
+    """The OpenAI spec types both as array/object, not nullable."""
+    payload = _chat_response(ChatMessage(role="assistant", content="done")).model_dump()
+
+    assert "annotations" not in payload["choices"][0]["message"]
+    assert "prompt_tokens_details" not in payload["usage"]
+
+
+def test_chat_completion_response_keeps_annotations_and_prompt_tokens_details():
+    citation = {"start_index": 0, "end_index": 4, "title": "t", "url": "https://x"}
+    response = _chat_response(
+        ChatMessage(
+            role="assistant",
+            content="done",
+            annotations=[{"type": "url_citation", "url_citation": citation}],
+        )
+    )
+    response.usage.prompt_tokens_details = PromptTokenUsageInfo(cached_tokens=1)
+
+    parsed = OpenAIChatCompletion.model_validate(response.model_dump())
+
+    assert parsed.choices[0].message.annotations[0].url_citation.url == "https://x"
+    assert parsed.usage.prompt_tokens_details.cached_tokens == 1
+
+
 def _stream_response(delta: DeltaMessage) -> ChatCompletionStreamResponse:
     return ChatCompletionStreamResponse(
         id="chatcmpl-test",
diff --git a/tests/entrypoints/pooling/embed/test_io_processor.py b/tests/entrypoints/pooling/embed/test_io_processor.py
index 8856759ef459..6f90f2c95210 100644
--- a/tests/entrypoints/pooling/embed/test_io_processor.py
+++ b/tests/entrypoints/pooling/embed/test_io_processor.py
@@ -2,6 +2,8 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Unit tests for EmbedIOProcessor."""
 
+from typing import get_args
+
 import pytest
 import torch
 from pydantic import TypeAdapter, ValidationError
@@ -110,6 +112,65 @@ def test_batched_messages_parses_as_batch_chat_request(self):
         ]
         assert request.chat_template_kwargs == {"instruction": "Represent the query: "}
 
+    @pytest.mark.parametrize(
+        "body",
+        [
+            {"input": "hello"},
+            {"input": [[1, 2], [3]]},
+            {"input": []},
+            {"input": [{"role": "user", "content": "hi"}]},
+            {"input": [[{"role": "user", "content": "hi"}]]},
+            {"messages": [{"role": "user", "content": "hi"}]},
+            {"messages": [[{"role": "user", "content": "hi"}]]},
+            {"messages": []},
+            {"messages": [{"author": {"role": "user"}, "content": []}]},
+            {"input": "hello", "messages": [{"role": "user", "content": "hi"}]},
+            {
+                "input": [{"role": "user", "content": "hi"}],
+                "messages": [{"role": "user", "content": "hi"}],
+            },
+        ],
+    )
+    def test_parses_like_plain_union(self, body):
+        """The completion fast path must parse like the plain union."""
+        plain_union = TypeAdapter(get_args(EmbeddingRequest)[0])
+        expected = plain_union.validate_python(body)
+        request: EmbeddingRequest = TypeAdapter(EmbeddingRequest).validate_python(body)
+
+        assert type(request) is type(expected)
+        assert request.model_dump(exclude={"request_id"}) == expected.model_dump(
+            exclude={"request_id"}
+        )
+
+    @pytest.mark.parametrize(
+        "body",
+        [
+            {"input": [1, "a"]},
+            {"input": [[1], "x"]},
+            {"input": "x", "dimensions": "abc"},
+            {
+                "input": "x",
+                "continue_final_message": True,
+                "add_generation_prompt": True,
+            },
+        ],
+    )
+    def test_rejects_like_plain_union(self, body):
+        """Invalid bodies must fail exactly as the plain union fails them."""
+        plain_union = TypeAdapter(get_args(EmbeddingRequest)[0])
+        with pytest.raises(Exception) as expected:
+            plain_union.validate_python(body)
+        with pytest.raises(Exception) as actual:
+            TypeAdapter(EmbeddingRequest).validate_python(body)
+
+        assert type(actual.value) is type(expected.value)
+        if isinstance(expected.value, ValidationError):
+            assert actual.value.errors(include_url=False) == expected.value.errors(
+                include_url=False
+            )
+        else:
+            assert str(actual.value) == str(expected.value)
+
 
 class TestCohereEmbedRequestParsing:
     """Unit tests for Cohere embed request parsing."""
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 e69dd6d45a5b..7fc2ab226c2b 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
@@ -3,16 +3,13 @@
 """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
@@ -20,57 +17,6 @@
 pytestmark = pytest.mark.skip_global_cleanup
 
 
-def test_online_forwards_truncate_prompt_tokens_to_proxy(monkeypatch):
-    """The proxy request handed to the base factory must carry
-    truncate_prompt_tokens/truncation_side from the real request.
-
-    JinaRankingIOProcessor swaps ctx.request for a proxy
-    PoolingCompletionRequest before delegating to the base factory, which
-    reads truncation off ctx.request. Dropping the fields on the proxy
-    silently disables truncate_prompt_tokens for Jina rerank/score.
-    """
-    proc = JinaRankingIOProcessor.__new__(JinaRankingIOProcessor)
-    monkeypatch.setattr(
-        proc,
-        "valid_inputs_online",
-        MagicMock(return_value=ScoringData(data_1=["query"], data_2=["doc"])),
-    )
-    monkeypatch.setattr(proc, "_get_token_limits", MagicMock(return_value=(0, 0)))
-    monkeypatch.setattr(
-        proc, "ensure_str", MagicMock(side_effect=lambda data: list(data))
-    )
-    monkeypatch.setattr(
-        proc, "format_docs_prompts_func", MagicMock(return_value="formatted prompt")
-    )
-
-    captured: dict[str, object] = {}
-
-    def _spy_base(self, ctx):
-        captured["truncate_prompt_tokens"] = ctx.request.truncate_prompt_tokens
-        captured["truncation_side"] = ctx.request.truncation_side
-        return []
-
-    monkeypatch.setattr(PoolingIOProcessor, "get_request_factory_online", _spy_base)
-
-    request = RerankRequest(
-        model="m",
-        query="query",
-        documents=["doc"],
-        truncate_prompt_tokens=512,
-        truncation_side="left",
-    )
-    ctx = MagicMock()
-    ctx.request = request
-    ctx.prompt_extras = None
-
-    proc.get_request_factory_online(ctx)
-
-    assert captured["truncate_prompt_tokens"] == 512
-    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]"))
diff --git a/tests/entrypoints/scale_out/derender/test_derender.py b/tests/entrypoints/scale_out/derender/test_derender.py
index dd312af88bdf..7f3f3cf7bb62 100644
--- a/tests/entrypoints/scale_out/derender/test_derender.py
+++ b/tests/entrypoints/scale_out/derender/test_derender.py
@@ -19,6 +19,7 @@
 )
 from vllm.renderers.online_derenderer import OnlineDerenderer
 from vllm.tokenizers import get_tokenizer
+from vllm.tokenizers.detokenizer_utils import convert_ids_list_to_tokens
 
 MODEL_NAME = "hmellor/tiny-random-LlamaForCausalLM"
 
@@ -522,6 +523,7 @@ def _make_completion_generate_response(
     request_id: str,
     kv_transfer_params: dict | None = None,
     logprobs: dict | None = None,
+    prompt_logprobs: list | None = None,
 ) -> dict:
     return {
         "request_id": request_id,
@@ -533,7 +535,7 @@ def _make_completion_generate_response(
                 "logprobs": logprobs,
             }
         ],
-        "prompt_logprobs": None,
+        "prompt_logprobs": prompt_logprobs,
         "kv_transfer_params": kv_transfer_params,
     }
 
@@ -737,6 +739,81 @@ async def test_derender_completion_logprobs(client):
     assert logprobs["text_offset"][0] == 0
 
 
+@pytest.mark.asyncio
+async def test_derender_completion_prompt_logprobs_passthrough(client):
+    """Each GenerateResponse's prompt_logprobs reaches that prompt's choices.
+
+    Regression test: /v1/completions/derender used to drop the prompt
+    logprobs the engine had already computed and returned, while the direct
+    /v1/completions endpoint attaches them to every choice of the prompt.
+    """
+    gr1 = await _render_completion(client, "Hello")
+    gr2 = await _render_completion(client, "World")
+    tok = get_tokenizer(MODEL_NAME)
+    # A --tokens-only generate server has no tokenizer: decoded_token is null.
+    plp1 = [None, {str(gr1["token_ids"][0]): {"logprob": -0.5, "rank": 1}}]
+    plp2 = [None, {str(gr2["token_ids"][0]): {"logprob": -0.25, "rank": 2}}]
+
+    response = await client.post(
+        "/v1/completions/derender",
+        json={
+            "model": MODEL_NAME,
+            "generate_responses": [
+                _make_completion_generate_response(
+                    gr1["token_ids"][:3], gr1["request_id"], prompt_logprobs=plp1
+                ),
+                _make_completion_generate_response(
+                    gr2["token_ids"][:3], gr2["request_id"], prompt_logprobs=plp2
+                ),
+            ],
+        },
+    )
+    assert response.status_code == 200
+    choices = response.json()["choices"]
+    assert len(choices) == 2
+    for choice, gr, plp in zip(choices, (gr1, gr2), (plp1, plp2)):
+        tid = gr["token_ids"][0]
+        (text,) = convert_ids_list_to_tokens(tok, [tid])
+        entry = choice["prompt_logprobs"][1][str(tid)]
+        assert plp[1] is not None
+        assert choice["prompt_logprobs"][0] is None
+        assert entry["decoded_token"] == text
+        assert entry["logprob"] == plp[1][str(tid)]["logprob"]
+        assert entry["rank"] == plp[1][str(tid)]["rank"]
+
+
+@pytest.mark.asyncio
+async def test_derender_chat_prompt_logprobs_decoded_token(client):
+    """Null decoded_token is filled (unknown ids as ""); given ones and empty
+    entries are kept."""
+    gen_req = await _render_chat(client)
+    ids = gen_req["token_ids"][:3]
+    tok = get_tokenizer(MODEL_NAME)
+    plp = [
+        None,
+        {str(ids[0]): {"logprob": -0.5, "rank": 1}},
+        {str(ids[1]): {"logprob": -0.1, "rank": 1, "decoded_token": "kept"}},
+        {},
+        {str(len(tok) + 1000): {"logprob": -1.0, "rank": 1}},
+    ]
+    response = await client.post(
+        "/v1/chat/completions/derender",
+        json={
+            "model": MODEL_NAME,
+            "generate_response": _make_generate_response(ids, prompt_logprobs=plp),
+        },
+    )
+    assert response.status_code == 200
+    out = response.json()["prompt_logprobs"]
+    assert (
+        out[1][str(ids[0])]["decoded_token"]
+        == (convert_ids_list_to_tokens(tok, [ids[0]])[0])
+    )
+    assert out[2][str(ids[1])]["decoded_token"] == "kept"
+    assert out[3] == {}
+    assert out[4][str(len(tok) + 1000)]["decoded_token"] == ""
+
+
 @pytest.mark.asyncio
 async def test_derender_completion_kv_transfer_params_passthrough(client):
     """kv_transfer_params passes through to CompletionResponse."""
@@ -1618,7 +1695,7 @@ def _derender_hermes_chat(
         )
     )
 
-    (choice,) = derenderer._derender_chat(response, request)
+    (choice,), _ = derenderer._derender_chat(response, request)
     return choice
 
 
@@ -1777,7 +1854,7 @@ def test_derender_batch_trims_stop_like_engine(
             messages=[{"role": "user", "content": "Count"}],
             **request_kwargs,
         )
-        (choice,) = derenderer._derender_chat(response, request)
+        (choice,), _ = derenderer._derender_chat(response, request)
         output_text = choice.message.content
     else:
         request = CompletionRequest(model=MODEL_NAME, prompt="Count", **request_kwargs)
diff --git a/tests/entrypoints/scale_out/derender/test_derender_parity.py b/tests/entrypoints/scale_out/derender/test_derender_parity.py
index 32288a7155d2..d5b925567db9 100644
--- a/tests/entrypoints/scale_out/derender/test_derender_parity.py
+++ b/tests/entrypoints/scale_out/derender/test_derender_parity.py
@@ -1021,6 +1021,20 @@ def _top_tokens(entry: dict) -> list[str]:
     return [top["token"] for top in entry["top_logprobs"]]
 
 
+def _assert_logprobs_parity(i_content: list, c_content: list, extra: dict) -> None:
+    assert [(e["token"], e["bytes"]) for e in i_content] == [
+        (e["token"], e["bytes"]) for e in c_content
+    ]
+    if "top_logprobs" in extra:
+        # Near ties can reorder candidates across the two generations, so
+        # compare the cut, not the candidates.
+        assert all(len(e["top_logprobs"]) == 3 for e in i_content)
+    else:
+        assert [sorted(_top_tokens(e)) for e in i_content] == [
+            sorted(_top_tokens(e)) for e in c_content
+        ]
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize(
     ("include_reasoning", "extra"),
@@ -1043,16 +1057,155 @@ async def test_inline_derender_logprobs(client, include_reasoning, extra):
         return
 
     coupled = (await _coupled(client, messages, **extra))["choices"][0]
-    i_content = inline["logprobs"]["content"]
-    c_content = coupled["logprobs"]["content"]
-    assert [(e["token"], e["bytes"]) for e in i_content] == [
-        (e["token"], e["bytes"]) for e in c_content
-    ]
-    if "top_logprobs" in extra:
-        # Near ties can reorder candidates across the two generations, so
-        # compare the cut, not the candidates.
-        assert all(len(e["top_logprobs"]) == 3 for e in i_content)
-    else:
-        assert [sorted(_top_tokens(e)) for e in i_content] == [
-            sorted(_top_tokens(e)) for e in c_content
+    _assert_logprobs_parity(
+        inline["logprobs"]["content"], coupled["logprobs"]["content"], extra
+    )
+
+
+async def _inline_derender_stream(
+    client: httpx.AsyncClient, messages: list[dict], **extra
+) -> list[dict]:
+    """Render a chat request, then stream it at the derender level. Returns
+    every streamed choice."""
+    render = await client.post(
+        "/v1/chat/completions/render",
+        json={
+            "model": MODEL,
+            "messages": messages,
+            "temperature": 0,
+            "max_tokens": 128,
+            **extra,
+        },
+    )
+    assert render.status_code == 200, render.text
+    choices: list[dict] = []
+    async with client.stream(
+        "POST",
+        "/inference/v1/generate",
+        json={**render.json(), "output_mode": "derender", "stream": True},
+    ) 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
+            assert chunk["output_mode"] == "derender"
+            choices.extend(chunk["choices"])
+    return choices
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("content", "extra"),
+    [
+        pytest.param(TEXT_MESSAGES[0]["content"], {}, id="plain"),
+        pytest.param(
+            "What is 17 times 23? Think it through.",
+            {"include_reasoning": True, "max_tokens": 256},
+            id="reasoning",
+        ),
+        pytest.param(
+            "What is 17 times 23? Think it through.",
+            {"include_reasoning": False, "max_tokens": 256},
+            id="reasoning-hidden",
+        ),
+        pytest.param(
+            "What's the weather in Paris?",
+            {"tools": TOOLS, "tool_choice": "auto", "max_tokens": 1024},
+            id="tool-auto",
+        ),
+        pytest.param(
+            "What's the weather in Paris?",
+            {"tools": TOOLS, "tool_choice": "required", "max_tokens": 1024},
+            id="tool-required",
+        ),
+        pytest.param(
+            "What's the weather in Paris?",
+            {"tools": TOOLS, "tool_choice": FORCE_WEATHER_TOOL, "max_tokens": 1024},
+            id="tool-named",
+        ),
+        pytest.param(
+            "What's the weather in Paris and in Rome?",
+            {
+                "tools": TOOLS,
+                "tool_choice": "required",
+                "parallel_tool_calls": False,
+                "max_tokens": 1024,
+            },
+            id="tool-no-parallel",
+        ),
+        pytest.param(
+            "Write two short lines about the sea.", {"stop": ["\n"]}, id="stop"
+        ),
+        pytest.param(
+            TEXT_MESSAGES[0]["content"],
+            {"n": 2, "temperature": 0.8, "seed": 42},
+            id="n2",
+        ),
+    ],
+)
+async def test_inline_derender_stream_parity(client, content, extra):
+    """Streaming at the derender level matches coupled chat streaming."""
+    messages = [{"role": "user", "content": content}]
+    hidden = extra.get("include_reasoning") is False
+    coupled_choices, _ = await _coupled_stream(
+        client, messages, require_token_ids=not hidden, **extra
+    )
+    inline_choices = await _inline_derender_stream(client, messages, **extra)
+
+    for index in range(extra.get("n", 1)):
+        coupled_index = [c for c in coupled_choices if c["index"] == index]
+        inline_index = [c for c in inline_choices if c["index"] == index]
+        roles = [c["delta"] for c in inline_index if c["delta"].get("role")]
+        assert len(roles) == 1
+        if hidden:
+            assert all(c["token_ids"] is None for c in inline_index)
+            assert assemble_stream(inline_index)["reasoning"] is None
+            assert assemble_stream(coupled_index)["reasoning"] is None
+            continue
+        assert [t for c in inline_index for t in c["token_ids"]] == [
+            t for c in coupled_index for t in c["token_ids"]
         ]
+        coupled = assemble_stream(coupled_index)
+        inline = assemble_stream(inline_index)
+        for message in (coupled, inline):
+            for tc in message["tool_calls"]:
+                tc["id"] = _tool_call_id_shape(tc["id"])
+        assert inline == coupled
+        if extra.get("parallel_tool_calls") is False:
+            assert len(inline["tool_calls"]) <= 1
+
+
+def _streamed_logprobs(choices: list[dict]) -> list[dict]:
+    return [e for c in choices if c.get("logprobs") for e in c["logprobs"]["content"]]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("include_reasoning", "extra"),
+    [
+        pytest.param(True, {"top_logprobs": 3}, id="top3"),
+        pytest.param(False, {"top_logprobs": 3}, id="top3-reasoning-hidden"),
+        pytest.param(True, {"logprob_token_ids": [15, 16, 17]}, id="token-ids"),
+    ],
+)
+async def test_inline_derender_stream_logprobs(client, include_reasoning, extra):
+    """Streamed logprobs match coupled chat, including the `top_logprobs` cut
+    and `logprob_token_ids`, and are dropped when reasoning is hidden."""
+    messages = [{"role": "user", "content": "What is 2+2?"}]
+    extra = {"logprobs": True, "include_reasoning": include_reasoning, **extra}
+    inline_choices = await _inline_derender_stream(client, messages, **extra)
+
+    if not include_reasoning:
+        assert all(c.get("logprobs") is None for c in inline_choices)
+        assert all(c["token_ids"] is None for c in inline_choices)
+        return
+
+    coupled_choices, _ = await _coupled_stream(client, messages, **extra)
+    _assert_logprobs_parity(
+        _streamed_logprobs(inline_choices), _streamed_logprobs(coupled_choices), extra
+    )
diff --git a/tests/entrypoints/scale_out/derender/test_derender_stream.py b/tests/entrypoints/scale_out/derender/test_derender_stream.py
index c843b4e89f95..6c0b6ce76222 100644
--- a/tests/entrypoints/scale_out/derender/test_derender_stream.py
+++ b/tests/entrypoints/scale_out/derender/test_derender_stream.py
@@ -357,7 +357,7 @@ def test_non_streaming_derender_initializes_parser_prefix(
         tools=_WEATHER_TOOLS,
     )
 
-    choices = derenderer._derender_chat(response, request)
+    choices, _ = derenderer._derender_chat(response, request)
 
     parser.set_prompt_token_ids.assert_called_once_with([11, 12])
     assert choices[0].message.tool_calls[0].function.name == "tool"
@@ -421,7 +421,7 @@ def batch_content(token_ids: list[int], finish_reason: str) -> str:
             response = GenerateTokensResponse(
                 request_id="test", prompt_token_ids=[11, 12], choices=[choice]
             )
-            return derenderer._derender_chat(response, request)[0].message.content
+            return derenderer._derender_chat(response, request)[0][0].message.content
 
         # Same decode window and seeding; "length" never drops a token.
         expected = batch_content(
@@ -755,7 +755,7 @@ async def test_chat_batch_seeded_via_explicit_prompt_token_ids(
         self, derenderer, leading_space_ids
     ):
         prompt_ids, output_ids, expected = leading_space_ids
-        choices = await derenderer.derender_chat(
+        choices, _ = await derenderer.derender_chat(
             GenerateTokensResponse(
                 request_id="t",
                 choices=[
@@ -773,7 +773,7 @@ async def test_chat_batch_seeded_via_generate_response_fallback(
         self, derenderer, leading_space_ids
     ):
         prompt_ids, output_ids, expected = leading_space_ids
-        choices = await derenderer.derender_chat(
+        choices, _ = await derenderer.derender_chat(
             GenerateTokensResponse(
                 request_id="t",
                 choices=[
@@ -792,7 +792,7 @@ async def test_chat_batch_explicit_takes_precedence_over_fallback(
     ):
         prompt_ids, output_ids, expected = leading_space_ids
         wrong_ids = tokenizer.encode("a different prompt", add_special_tokens=False)
-        choices = await derenderer.derender_chat(
+        choices, _ = await derenderer.derender_chat(
             GenerateTokensResponse(
                 request_id="t",
                 choices=[
@@ -811,7 +811,7 @@ async def test_chat_batch_without_prompt_context_stays_unseeded(
         self, derenderer, leading_space_ids
     ):
         _, output_ids, _ = leading_space_ids
-        choices = await derenderer.derender_chat(
+        choices, _ = await derenderer.derender_chat(
             GenerateTokensResponse(
                 request_id="t",
                 choices=[
@@ -1902,7 +1902,7 @@ async def test_stream_matches_batch(self, harmony_derenderer, harmony_encode):
         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(
+        batch_choices, _ = await harmony_derenderer.derender_chat(
             GenerateTokensResponse(
                 request_id="test-harmony-batch",
                 choices=[
diff --git a/tests/entrypoints/scale_out/render/test_render.py b/tests/entrypoints/scale_out/render/test_render.py
index 862a89dc8e46..665df4771cdf 100644
--- a/tests/entrypoints/scale_out/render/test_render.py
+++ b/tests/entrypoints/scale_out/render/test_render.py
@@ -40,6 +40,7 @@
 
 MODEL_NAME = "hmellor/tiny-random-LlamaForCausalLM"
 GPT_OSS_MODEL_NAME = "openai/gpt-oss-20b"
+DEEPSEEK_V41_MODEL_NAME = "deepseek-ai/DeepSeek-V4.1-Flash"
 
 
 @pytest.mark.parametrize("multiple", [False, True])
@@ -1081,3 +1082,88 @@ async def test_tokenize_matches_chat_render_for_gpt_oss(gpt_oss_server):
 
     assert render.status_code == tokenize.status_code == 200
     assert tokenize.json()["tokens"] == render.json()["token_ids"]
+
+
+@pytest.fixture(scope="module")
+def deepseek_v41_server():
+    args = [
+        "--tokenizer-mode",
+        "deepseek_v41",
+        "--tool-call-parser",
+        "deepseek_v41",
+        "--enable-auto-tool-choice",
+        "--language-model-only",
+    ]
+    with RemoteLaunchRenderServer(DEEPSEEK_V41_MODEL_NAME, args) as remote_server:
+        yield remote_server
+
+
+@pytest.mark.asyncio
+async def test_messages_render_tool_search_result_for_deepseek_v41(
+    deepseek_v41_server,
+):
+    """A tool search result renders as one result naming the tools it found."""
+    request = {
+        "model": DEEPSEEK_V41_MODEL_NAME,
+        "max_tokens": 16,
+        "tools": [
+            {"name": "tool_search", "input_schema": {"type": "object"}},
+            {"name": "read_file", "input_schema": {"type": "object"}},
+            {
+                "name": "get_weather",
+                "input_schema": {"type": "object"},
+                "defer_loading": True,
+            },
+        ],
+        "messages": [
+            {"role": "user", "content": "Read notes.txt and find a weather tool."},
+            {
+                "role": "assistant",
+                "content": [
+                    {
+                        "type": "tool_use",
+                        "id": "toolu_1",
+                        "name": "tool_search",
+                        "input": {"query": "weather"},
+                    },
+                    {
+                        "type": "tool_use",
+                        "id": "toolu_2",
+                        "name": "read_file",
+                        "input": {"path": "notes.txt"},
+                    },
+                ],
+            },
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "tool_result",
+                        "tool_use_id": "toolu_1",
+                        "content": [
+                            {"type": "tool_reference", "tool_name": "get_weather"}
+                        ],
+                    },
+                    {
+                        "type": "tool_result",
+                        "tool_use_id": "toolu_2",
+                        "content": "notes",
+                    },
+                ],
+            },
+        ],
+    }
+    async with httpx.AsyncClient(
+        base_url=deepseek_v41_server.url_for(""), timeout=30.0
+    ) as client:
+        render = await client.post("/v1/messages/render", json=request)
+        assert render.status_code == 200, render.text
+        detokenize = await client.post(
+            "/detokenize", json={"tokens": render.json()["token_ids"]}
+        )
+
+    assert detokenize.status_code == 200
+    assert (
+        "<｜User｜><tool_result>get_weather</tool_result>\n\n"
+        "<tool_result>notes</tool_result><｜Assistant｜>"
+    ) in detokenize.json()["prompt"]
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 a85355b29611..1994540bce29 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
@@ -1,15 +1,17 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import asyncio
+import functools
 import json
 from dataclasses import dataclass, field
 from typing import Any
-from unittest.mock import AsyncMock, MagicMock
+from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
 
 from vllm.config.multimodal import MultiModalConfig
-from vllm.entrypoints.generate.base.protocol import StreamOptions
+from vllm.entrypoints.generate.base.protocol import DeltaMessage, 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 (
@@ -19,10 +21,12 @@
     encode_mm_kwargs_item,
 )
 from vllm.entrypoints.scale_out.token_in_token_out.protocol import (
+    GenerateDerenderStreamResponse,
     GenerateRequest,
     GenerateTextResponse,
     GenerateTokensResponse,
     MultiModalFeatures,
+    ParseContext,
     PlaceholderRangeInfo,
 )
 from vllm.entrypoints.scale_out.token_in_token_out.serving import ServingTokens
@@ -1563,3 +1567,313 @@ async def mock_generate(*args, **kwargs):
     assert usage_chunk["output_mode"] == "text"
     assert usage_chunk["choices"] == []
     assert usage_chunk["usage"]["completion_tokens"] == 1
+
+
+class _StubParser:
+    """Echoes text as content. From the first "~" on it holds text back until
+    the finished step, like a parser buffering a partial tool call."""
+
+    def __init__(self, *args, **kwargs):
+        self.held = ""
+        self.history_tool_call_cnt = 0
+
+    def parse_delta(self, delta_text, finished=False, **kwargs):
+        if self.held or "~" in delta_text:
+            self.held += delta_text
+            if not finished:
+                return None
+            delta_text, self.held = self.held, ""
+        return DeltaMessage(content=delta_text)
+
+    def set_history_tool_call_count(self, count):
+        self.history_tool_call_cnt = count
+
+
+def _derender_serving(engine, parser_cls=_StubParser):
+    """Serving whose derenderer builds real stream processors around
+    `parser_cls`, or plain ones when it is None."""
+    serving = _build_serving_tokens(engine)
+    derenderer = serving.online_derenderer
+    derenderer.parser = parser_cls
+    derenderer.enable_auto_tools = False
+    derenderer.renderer = MagicMock()
+    derenderer.model_config = engine.model_config
+    derenderer.make_stream_processor = functools.partial(
+        OnlineDerenderer.make_stream_processor, derenderer
+    )
+    return serving
+
+
+def _derender_request(parse_context=None, n=1, **kwargs):
+    return GenerateRequest(
+        token_ids=[1, 2, 3],
+        sampling_params=SamplingParams(max_tokens=10, n=n),
+        model=MODEL_NAME,
+        stream=True,
+        output_mode="derender",
+        parse_context=parse_context or ParseContext(),
+        **kwargs,
+    )
+
+
+@pytest.mark.asyncio
+async def test_stream_derender_flushes_on_finish_only_output():
+    engine = _mock_engine()
+
+    async def mock_generate(*args, **kwargs):
+        yield _make_request_output("req-1", token_ids=[10], text="a")
+        yield _make_request_output("req-1", token_ids=[11], text="~")
+        yield _make_request_output(
+            "req-1", token_ids=[], text="b", finish_reason="stop", finished=True
+        )
+
+    engine.generate = MagicMock(side_effect=mock_generate)
+    serving = _derender_serving(engine)
+
+    parsed = await _collect_sse(await serving.serve_tokens(_derender_request()))
+    choices = [c["choices"][0] for c in parsed[:-1]]
+
+    assert {c["output_mode"] for c in parsed[:-1]} == {"derender"}
+    assert [c["delta"] for c in choices] == [
+        {"role": "assistant", "content": "a"},
+        {},
+        {"content": "~b"},
+    ]
+    assert [c["token_ids"] for c in choices] == [[10], [11], []]
+    assert [c["finish_reason"] for c in choices] == [None, None, "stop"]
+
+
+@pytest.mark.asyncio
+async def test_stream_derender_abort_without_tokens_sends_empty_delta():
+    engine = _mock_engine()
+
+    async def mock_generate(*args, **kwargs):
+        yield _make_request_output(
+            "req-1", token_ids=[], finish_reason="abort", finished=True
+        )
+
+    engine.generate = MagicMock(side_effect=mock_generate)
+    serving = _derender_serving(engine)
+
+    parsed = await _collect_sse(await serving.serve_tokens(_derender_request()))
+
+    assert len(parsed) == 2
+    assert parsed[0]["choices"][0]["delta"] == {"role": "assistant"}
+    assert parsed[0]["choices"][0]["finish_reason"] == "abort"
+
+
+@pytest.mark.asyncio
+async def test_stream_derender_hidden_reasoning_drops_metadata():
+    engine = _mock_engine()
+
+    async def mock_generate(*args, **kwargs):
+        yield _make_request_output(
+            "req-1",
+            token_ids=[10],
+            text="a",
+            logprobs=[{10: Logprob(logprob=-0.1)}],
+        )
+        yield _make_request_output(
+            "req-1",
+            token_ids=[11],
+            text="b",
+            finish_reason="stop",
+            finished=True,
+            logprobs=[{11: Logprob(logprob=-0.2)}],
+        )
+
+    engine.generate = MagicMock(side_effect=mock_generate)
+    serving = _derender_serving(engine)
+    request = _derender_request(ParseContext(include_reasoning=False))
+    request.sampling_params.logprobs = 1
+
+    parsed = await _collect_sse(await serving.serve_tokens(request))
+
+    for chunk in parsed[:-1]:
+        choice = chunk["choices"][0]
+        assert choice["logprobs"] is None
+        assert choice["token_ids"] is None
+        assert choice["sampling_mask"] is None
+    assert parsed[0]["choices"][0]["delta"]["content"] == "a"
+
+
+@pytest.mark.asyncio
+async def test_stream_derender_n2_has_independent_state():
+    engine = _mock_engine()
+
+    async def mock_generate(*args, **kwargs):
+        yield _make_request_output("req-1", token_ids=[10], text="~a", index=0)
+        yield _make_request_output("req-1", token_ids=[20], text="x", index=1)
+        yield _make_request_output(
+            "req-1", token_ids=[11], text="b", index=0, finish_reason="stop"
+        )
+        yield _make_request_output(
+            "req-1", token_ids=[21], text="y", index=1, finish_reason="stop"
+        )
+
+    engine.generate = MagicMock(side_effect=mock_generate)
+    serving = _derender_serving(engine)
+
+    parsed = await _collect_sse(await serving.serve_tokens(_derender_request(n=2)))
+    choices = [c["choices"][0] for c in parsed[:-1]]
+
+    # Choice 0 buffers, choice 1 doesn't, so each has its own parser.
+    assert [(c["index"], c["delta"]) for c in choices] == [
+        (0, {"role": "assistant"}),
+        (1, {"role": "assistant", "content": "x"}),
+        (0, {"content": "~ab"}),
+        (1, {"content": "y"}),
+    ]
+
+
+@pytest.mark.asyncio
+async def test_stream_derender_without_parser_streams_text_as_content():
+    engine = _mock_engine()
+
+    async def mock_generate(*args, **kwargs):
+        yield _make_request_output("req-1", token_ids=[10], text="a~")
+        yield _make_request_output(
+            "req-1", token_ids=[11], text="b", finish_reason="stop", finished=True
+        )
+
+    engine.generate = MagicMock(side_effect=mock_generate)
+    serving = _derender_serving(engine, parser_cls=None)
+
+    parsed = await _collect_sse(await serving.serve_tokens(_derender_request()))
+    choices = [c["choices"][0] for c in parsed[:-1]]
+
+    assert [c["delta"] for c in choices] == [
+        {"role": "assistant", "content": "a~"},
+        {"content": "b"},
+    ]
+    assert [c["finish_reason"] for c in choices] == [None, "stop"]
+
+
+def test_make_stream_processor_seeds_history_tool_call_count():
+    """Without the seed, deterministic tool call IDs restart at 0."""
+    serving = _derender_serving(_mock_engine())
+
+    processor = serving.online_derenderer.make_stream_processor(
+        ParseContext().to_chat_request(), 3
+    )
+
+    assert processor.parser.history_tool_call_cnt == 3
+
+
+class _RejectingParser(_StubParser):
+    def __init__(self, *args, **kwargs):
+        raise ValueError("unsupported tools")
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("parser_cls", "parse_context"),
+    [
+        pytest.param(
+            _StubParser,
+            ParseContext(
+                tool_choice={"type": "function", "function": {"name": "missing"}}
+            ),
+            id="named-tool-missing",
+        ),
+        pytest.param(_RejectingParser, ParseContext(), id="parser-rejects"),
+    ],
+)
+async def test_stream_derender_invalid_request_fails_before_engine(
+    parser_cls, parse_context
+):
+    engine = _mock_engine()
+    serving = _derender_serving(engine, parser_cls=parser_cls)
+
+    response = await serving.serve_tokens(_derender_request(parse_context))
+
+    assert isinstance(response, ErrorResponse)
+    assert response.error.code == 400
+    engine.generate.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_stream_derender_requires_parse_context_with_parser():
+    engine = _mock_engine()
+    serving = _derender_serving(engine)
+    request = _derender_request()
+    request.parse_context = None
+
+    response = await serving.serve_tokens(request)
+
+    assert isinstance(response, ErrorResponse)
+    assert response.error.code == 400
+    engine.generate.assert_not_called()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "abort_error",
+    [None, RuntimeError("engine dead")],
+    ids=["abort-ok", "abort-fails"],
+)
+async def test_stream_derender_parser_failure_closes_generate(abort_error):
+    """`generate()` is closed in the background, so its abort neither delays
+    the error event and [DONE] nor loses them if it fails. Closing it aborts by
+    internal ID, sparing other requests that share the external ID."""
+    engine = _mock_engine()
+    release = asyncio.Event()
+    closed = asyncio.Event()
+
+    class _RaisingParser(_StubParser):
+        def parse_delta(self, delta_text, finished=False, **kwargs):
+            raise RuntimeError("boom")
+
+    async def mock_generate(*args, **kwargs):
+        try:
+            yield _make_request_output("req-1", token_ids=[10], text="a")
+        except GeneratorExit:
+            await release.wait()
+            closed.set()
+            if abort_error is not None:
+                raise abort_error from None
+            raise
+
+    async def collect_chunks():
+        return [c async for c in await serving.serve_tokens(_derender_request())]
+
+    engine.generate = MagicMock(side_effect=mock_generate)
+    serving = _derender_serving(engine, parser_cls=_RaisingParser)
+
+    with patch(
+        "vllm.entrypoints.scale_out.token_in_token_out.serving.logger"
+    ) as logger:
+        chunks = await asyncio.wait_for(collect_chunks(), timeout=5)
+        assert any('"code": 500' in c for c in chunks)
+        assert chunks[-1] == "data: [DONE]\n\n"
+
+        release.set()
+        await asyncio.wait_for(closed.wait(), timeout=5)
+
+    logged = [call.args[0] for call in logger.exception.call_args_list]
+    assert ("Failed to abort request %s." in logged) == (abort_error is not None)
+    engine.abort.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_stream_derender_chunk_has_no_null_delta_keys():
+    engine = _mock_engine()
+
+    async def mock_generate(*args, **kwargs):
+        yield _make_request_output(
+            "req-1", token_ids=[10], text="a", finish_reason="stop", finished=True
+        )
+
+    engine.generate = MagicMock(side_effect=mock_generate)
+    serving = _derender_serving(engine)
+
+    response = await serving.serve_tokens(_derender_request())
+    chunks = [c async for c in response]
+
+    assert '"delta":{"role":"assistant","content":"a"}' in chunks[0]
+    assert (
+        GenerateDerenderStreamResponse.model_validate_json(chunks[0][len("data: ") :])
+        .choices[0]
+        .delta.content
+        == "a"
+    )
diff --git a/tests/entrypoints/scale_out/token_in_token_out/test_protocol.py b/tests/entrypoints/scale_out/token_in_token_out/test_protocol.py
index 159125115a69..586f19ab06a9 100644
--- a/tests/entrypoints/scale_out/token_in_token_out/test_protocol.py
+++ b/tests/entrypoints/scale_out/token_in_token_out/test_protocol.py
@@ -18,6 +18,7 @@
     DerenderCompletionRequest,
     DerenderCompletionStreamRequest,
     GenerateDerenderResponse,
+    GenerateDerenderStreamResponse,
     GenerateRequest,
     GenerateResponse,
     GenerateStreamResponse,
@@ -276,6 +277,19 @@ def test_text_response_and_chunk_dispatch_on_output_mode():
     assert isinstance(chunk, GenerateTextStreamResponse)
 
 
+def test_derender_chunk_dispatches_and_derender_stream_request_rejects_it():
+    chunk = {
+        "output_mode": "derender",
+        "choices": [{"index": 0, "delta": {"content": "hi"}}],
+    }
+    parsed = TypeAdapter(GenerateStreamResponse).validate_python(chunk)
+    assert isinstance(parsed, GenerateDerenderStreamResponse)
+    with pytest.raises(ValidationError, match="output_mode"):
+        DerenderChatStreamRequest.model_validate(
+            {"stream": True, "generate_chunk": chunk}
+        )
+
+
 def test_text_response_requires_text_on_every_choice():
     with pytest.raises(ValidationError, match="text"):
         TypeAdapter(GenerateResponse).validate_python(
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 70c3abcb7df2..b682fa79e42e 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
@@ -744,9 +744,6 @@ async def test_text_mode_stream_matches_non_stream(client, tokenizer, messages):
             {"sampling_params": {"detokenize": False}}, "detokenize", id="no-detok"
         ),
         pytest.param({"output_mode": "bogus"}, "output_mode", id="unsupported"),
-        pytest.param(
-            {"output_mode": "derender", "stream": True}, "stream", id="derender-stream"
-        ),
     ],
 )
 async def test_text_mode_invalid_requests_return_400(client, overrides, expected):
diff --git a/tests/entrypoints/serve/exception_handling/test_error_sanitization.py b/tests/entrypoints/serve/exception_handling/test_error_sanitization.py
index c783fb7db67c..0f06c3cf9bed 100644
--- a/tests/entrypoints/serve/exception_handling/test_error_sanitization.py
+++ b/tests/entrypoints/serve/exception_handling/test_error_sanitization.py
@@ -153,3 +153,57 @@ def test_engine_dead_error_response_hides_log_oriented_message():
     assert "stack trace" not in err.error.message
     assert "log" not in err.error.message
     assert "stack trace (above)" in str(exc)
+
+
+class TestSanitizeMessageReDoSResistance:
+    """Regression tests for GHSA-f2g9-pmwr-xwc7.
+
+    Verifies that sanitize_message performs bounded near-linear work
+    on adversarial slash-delimited input with no filename extension.
+    """
+
+    def test_slash_segments_no_extension_completes_quickly(self):
+        """50k slash-segments without dot-extension must complete in < 0.1s."""
+        import time
+
+        message = "/a" * 50000
+        t0 = time.perf_counter()
+        result = sanitize_message(message)
+        elapsed = time.perf_counter() - t0
+        assert elapsed < 0.1, f"Took {elapsed:.2f}s (expected < 0.1s)"
+        assert result == message
+
+    def test_alphanumeric_control_completes(self):
+        """100k alphanumeric chars must complete near-instantly."""
+        import time
+
+        message = "a" * 100000
+        t0 = time.perf_counter()
+        result = sanitize_message(message)
+        elapsed = time.perf_counter() - t0
+        assert elapsed < 0.1, f"Took {elapsed:.2f}s"
+        assert result == message
+
+    def test_linear_scaling(self):
+        """Verify near-linear scaling across input sizes."""
+        import time
+
+        times = []
+        for n in [1000, 10000, 50000]:
+            message = "/a" * n
+            t0 = time.perf_counter()
+            sanitize_message(message)
+            times.append(time.perf_counter() - t0)
+
+        ratio = times[-1] / max(times[0], 1e-9)
+        assert ratio < 100, (
+            f"50k/1k time ratio is {ratio:.0f}x (expected < 100x for linear)"
+        )
+
+    def test_path_redaction_still_works(self):
+        """Legitimate paths are still redacted after the fix."""
+        assert "<path>" in sanitize_message("Error in /app/server.py")
+        assert "<path>" in sanitize_message("at /workspace/vllm/engine.py")
+        assert "<path>" in sanitize_message(
+            "error at /usr/lib/python3.12/dist-packages/vllm/core.py"
+        )
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 d985d4586605..abd9633c75da 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
@@ -318,6 +318,7 @@ def run_longform_evaluation(
         ("openai/whisper-large-v3", 12.744980),
         # CohereASR is used to test the variable encoder length code paths
         ("CohereLabs/cohere-transcribe-03-2026", 11.92),
+        ("nvidia/canary-1b-v2", 12.34),
     ],
 )
 # Original dataset is 20GB+ in size, hence we use a pre-filtered slice.
diff --git a/tests/entrypoints/speech_to_text/transcription/test_transcription_validation.py b/tests/entrypoints/speech_to_text/transcription/test_transcription_validation.py
index 2237a091ce64..6a14635ca47d 100644
--- a/tests/entrypoints/speech_to_text/transcription/test_transcription_validation.py
+++ b/tests/entrypoints/speech_to_text/transcription/test_transcription_validation.py
@@ -1,12 +1,20 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-# imports for structured outputs tests
 import json
+from io import BytesIO
 
+import httpx
+import numpy as np
 import pytest
+import soundfile as sf
+from fastapi import UploadFile
+from transformers import WhisperConfig
 
 from tests.utils import ROCM_EXTRA_ARGS, RemoteOpenAIServer
+from vllm.config import ModelConfig, SpeechToTextConfig
+from vllm.entrypoints.speech_to_text.transcription.protocol import TranscriptionRequest
+from vllm.multimodal.audio import AudioResampler
 
 MISTRAL_FORMAT_ARGS = [
     "--tokenizer_mode",
@@ -16,6 +24,69 @@
     "--load_format",
     "mistral",
 ]
+MING_ASR_ARGS = [
+    "--trust-remote-code",
+    "--max-model-len",
+    "8192",
+    *(ROCM_EXTRA_ARGS or ["--max-num-seqs", "4"]),
+    "--max-num-batched-tokens",
+    "1024",
+    "--logits-processors",
+    "vllm.model_executor.models.ming_v3_asr_logits_processor:NoRepeatNGramLogitsProcessor",
+]
+
+
+@pytest.fixture
+def ming_server(model_name, request):
+    args = ["--enforce-eager", *MING_ASR_ARGS]
+    options = getattr(request, "param", {})
+    if options.get("enable_prefix_caching", False):
+        args = [arg for arg in args if arg != "--no-enable-prefix-caching"]
+        args.append("--enable-prefix-caching")
+    with RemoteOpenAIServer(model_name, args) as remote:
+        yield remote
+
+
+@pytest.fixture(scope="module")
+def ming_audio():
+    base_url = (
+        "https://raw.githubusercontent.com/inclusionAI/Ming-3.0-asr/"
+        "6141e8c552815db724d8df0f72c9d865c547a1ea/examples/cases/audio/"
+    )
+    with httpx.Client() as client:
+        clips = {}
+        for name in (
+            "chinese_common_voice.wav",
+            "context_same_sentence_pcm16.wav",
+        ):
+            response = client.get(base_url + name)
+            response.raise_for_status()
+            clips[name] = response.content
+        return clips
+
+
+@pytest.mark.parametrize("encoding", ["absent", "dict", "json"])
+def test_request_extensions_reach_prompt_and_sampling_params(tmp_path, encoding):
+    """Both request consumers must receive extensions, including form JSON."""
+    WhisperConfig().save_pretrained(tmp_path)
+    config = ModelConfig(model=str(tmp_path))
+    extensions = {"frontend_window": 128, "custom_tags": ["speech", 1]}
+    options = {}
+    if encoding != "absent":
+        options["vllm_xargs"] = (
+            json.dumps(extensions) if encoding == "json" else extensions
+        )
+    request = TranscriptionRequest(file=UploadFile(BytesIO()), **options)
+    params = request.build_stt_params(
+        audio=np.zeros(160, dtype=np.float32),
+        stt_config=SpeechToTextConfig(),
+        model_config=config,
+        task_type="transcribe",
+    )
+
+    expected = None if encoding == "absent" else extensions
+    assert params.extra_args == expected
+    assert (request.to_sampling_params(16).extra_args or {}) == (expected or {})
 
 
 async def transcribe_and_check(
@@ -23,10 +94,12 @@ async def transcribe_and_check(
     model_name: str,
     file,
     *,
-    language: str,
+    language: str | None,
     expected_text: str,
     expected_seconds: int | None = None,
     case_sensitive: bool = False,
+    response_format: str = "text",
+    extra_body: dict | None = None,
 ):
     """Run a transcription request and assert the output contains
     *expected_text* and optionally that usage reports *expected_seconds*.
@@ -37,12 +110,16 @@ async def transcribe_and_check(
         model=model_name,
         file=file,
         language=language,
-        response_format="text",
+        response_format=response_format,
         temperature=0.0,
+        extra_body=extra_body,
+    )
+    out = (
+        json.loads(transcription)
+        if isinstance(transcription, str)
+        else transcription.model_dump()
     )
-    out = json.loads(transcription)
     out_text = out["text"]
-    out_usage = out["usage"]
 
     if case_sensitive:
         assert expected_text in out_text, (
@@ -55,21 +132,35 @@ async def transcribe_and_check(
         )
 
     if expected_seconds is not None:
+        out_usage = out["usage"]
         assert out_usage["seconds"] == expected_seconds, (
             f"Expected {expected_seconds}s of audio, "
             f"got {out_usage['seconds']}s. Full usage: {out_usage!r}"
         )
 
+    return out_text
+
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize(
-    "model_name", ["mistralai/Voxtral-Mini-3B-2507", "Qwen/Qwen3-ASR-0.6B"]
+    "model_name",
+    [
+        "mistralai/Voxtral-Mini-3B-2507",
+        "Qwen/Qwen3-ASR-0.6B",
+        "inclusionAI/Ming-3.0-asr-tiny",
+    ],
 )
 async def test_basic_audio(mary_had_lamb, model_name):
     server_args = ["--enforce-eager", *ROCM_EXTRA_ARGS]
-
+    extra_body = None
     if model_name.startswith("mistralai"):
         server_args += MISTRAL_FORMAT_ARGS
+    elif model_name == "inclusionAI/Ming-3.0-asr-tiny":
+        server_args = ["--enforce-eager", *MING_ASR_ARGS]
+        extra_body = {
+            "max_completion_tokens": 256,
+            "vllm_xargs": json.dumps({"no_repeat_ngram_size": 10}),
+        }
 
     # Based on https://github.com/openai/openai-cookbook/blob/main/examples/Whisper_prompting_guide.ipynb.
     with RemoteOpenAIServer(model_name, server_args) as remote_server:
@@ -81,6 +172,7 @@ async def test_basic_audio(mary_had_lamb, model_name):
             language="en",
             expected_text="Mary had a little lamb",
             expected_seconds=16,
+            extra_body=extra_body,
         )
 
 
@@ -142,3 +234,76 @@ async def test_basic_audio_foscolo(foscolo, model_name):
             language="it",
             expected_text="ove il mio corpo fanciulletto",
         )
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("model_name", ["inclusionAI/Ming-3.0-asr-tiny"])
+async def test_ming_context(model_name, ming_server, ming_audio):
+    """Topic and hotword hints must reach the native ASR prompt over HTTP."""
+    async with ming_server.get_async_client() as client:
+        text = await transcribe_and_check(
+            client,
+            model_name,
+            ("context.wav", ming_audio["context_same_sentence_pcm16.wav"], "audio/wav"),
+            language=None,
+            expected_text="同囚",
+            response_format="json",
+            extra_body={
+                "max_completion_tokens": 256,
+                "vllm_xargs": json.dumps({"no_repeat_ngram_size": 10}),
+                "prompt": "This is an audio about 同囚.",
+                "hotwords": (
+                    "同囚,黄国权,游学修,关楚耀,赵永洪,赵咏瑶,狱警,长镜头,"
+                    "吞玻璃戏,糖化玻璃,真实案例数据."
+                ),
+            },
+        )
+    assert "\t" not in text
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("model_name", ["inclusionAI/Ming-3.0-asr-tiny"])
+@pytest.mark.parametrize(
+    "ming_server", [{"enable_prefix_caching": True}], indirect=True
+)
+async def test_ming_chunk_policy_prefix_cache(model_name, ming_server, ming_audio):
+    """Restoring a long request's chunk policy must reuse its cached prefix."""
+    waveform, rate = sf.read(
+        BytesIO(ming_audio["chinese_common_voice.wav"]), dtype="float32"
+    )
+    waveform = AudioResampler(16000).resample(waveform, orig_sr=rate)
+    waveform = np.pad(waveform, (0, 64 * 16000 - len(waveform)))
+    clip = BytesIO()
+    sf.write(clip, waveform, 16000, format="WAV", subtype="PCM_16")
+    async with httpx.AsyncClient() as metrics, ming_server.get_async_client() as client:
+
+        async def prefix_hits():
+            response = await metrics.get(ming_server.url_for("metrics"))
+            response.raise_for_status()
+            return sum(
+                float(line.rsplit(" ", 1)[1])
+                for line in response.text.splitlines()
+                if line.startswith("vllm:prefix_cache_hits_total{")
+            )
+
+        for frames in (3000, 6000, 3000):
+            before = await prefix_hits()
+            text = await transcribe_and_check(
+                client,
+                model_name,
+                ("audio.wav", clip.getvalue(), "audio/wav"),
+                language="zh",
+                expected_text="漓江和湘江有秦代开凿的灵渠沟通",
+                response_format="json",
+                extra_body={
+                    "max_completion_tokens": 256,
+                    "vllm_xargs": json.dumps(
+                        {"audio_chunk_frames": frames, "no_repeat_ngram_size": 10}
+                    ),
+                },
+            )
+            assert "\t" not in text
+        after = await prefix_hits()
+        assert after > before, (
+            f"Restoring the chunk policy added no prefix hits: {before=} {after=}"
+        )
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 d5658e7cf0b4..17b15563fb3a 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
@@ -89,6 +89,7 @@ async def test_basic_audio(whisper_client, mary_had_lamb):
         language="en",
         response_format="text",
         temperature=0.0,
+        extra_body={"vllm_xargs": json.dumps({"frontend_window": 128})},
     )
     out = json.loads(transcription)
     out_text = out["text"]
diff --git a/tests/entrypoints/tool_parsers/test_functiongemma_tool_parser.py b/tests/entrypoints/tool_parsers/test_functiongemma_tool_parser.py
new file mode 100644
index 000000000000..048ddbbaaea5
--- /dev/null
+++ b/tests/entrypoints/tool_parsers/test_functiongemma_tool_parser.py
@@ -0,0 +1,101 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import json
+
+import pytest
+import regex as re
+
+from tests.utils import RemoteOpenAIServer
+
+MODEL = "unsloth/functiongemma-270m-it"
+START = "<start_function_call>"
+END = "<end_function_call>"
+TOOLS = [
+    {
+        "type": "function",
+        "function": {
+            "name": name,
+            "description": name,
+            "parameters": {"type": "object", "properties": {}},
+        },
+    }
+    for name in ("set_alarm", "get_weather", "get_time")
+]
+
+
+@pytest.fixture(scope="module")
+def server():
+    args = [
+        "--enforce-eager",
+        "--enable-auto-tool-choice",
+        "--tool-call-parser",
+        "functiongemma",
+        "--max-model-len",
+        "2048",
+        "--max-num-seqs",
+        "2",
+        "--gpu-memory-utilization",
+        "0.05",
+    ]
+    with RemoteOpenAIServer(MODEL, args, max_wait_seconds=480) as server:
+        yield server
+
+
+def _tool_calls(server, output, stream, stream_interval=None):
+    """Pin the model output with a regex, so only serving and parsing are tested."""
+    extra_body = {"structured_outputs": {"regex": re.escape(output)}}
+    if stream_interval is not None:
+        extra_body["stream_interval"] = stream_interval
+    with server.get_client() as client:
+        response = client.chat.completions.create(
+            model=MODEL,
+            messages=[{"role": "user", "content": "Call the tools."}],
+            tools=TOOLS,
+            tool_choice="auto",
+            temperature=0,
+            max_tokens=128,
+            stream=stream,
+            extra_body=extra_body,
+        )
+        if not stream:
+            calls = response.choices[0].message.tool_calls or []
+            return [(c.function.name, json.loads(c.function.arguments)) for c in calls]
+        streamed: dict[int, dict] = {}
+        for chunk in response:
+            for call in chunk.choices[0].delta.tool_calls or []:
+                entry = streamed.setdefault(call.index, {"name": None, "args": ""})
+                entry["name"] = entry["name"] or call.function.name
+                entry["args"] += call.function.arguments or ""
+    return [(e["name"], json.loads(e["args"])) for _, e in sorted(streamed.items())]
+
+
+@pytest.mark.parametrize("stream", [False, True])
+@pytest.mark.parametrize(
+    ("args", "expected"),
+    [
+        (
+            "days:[<escape>mon<escape>],enabled:true,hour:7",
+            {"days": ["mon"], "enabled": True, "hour": 7},
+        ),
+        ("", {}),
+    ],
+    ids=["non_string_arguments", "no_arguments"],
+)
+def test_arguments_reach_the_client(server, stream, args, expected):
+    output = f"{START}call:set_alarm{{{args}}}{END}"
+    assert _tool_calls(server, output, stream) == [("set_alarm", expected)]
+
+
+@pytest.mark.parametrize("stream_interval", [2, 3, 4])
+def test_streamed_call_closes_when_next_call_starts_in_same_delta(
+    server, stream_interval
+):
+    """These intervals merge the first call's end tag and the second's start tag."""
+    output = (
+        f"{START}call:get_weather{{city:<escape>Paris<escape>}}{END}"
+        f"{START}call:get_time{{tz:<escape>UTC<escape>}}{END}"
+    )
+    assert _tool_calls(server, output, True, stream_interval) == [
+        ("get_weather", {"city": "Paris"}),
+        ("get_time", {"tz": "UTC"}),
+    ]
diff --git a/tests/evals/gsm8k/configs/DeepSeek-V3.2-DP_MI355.yaml b/tests/evals/gsm8k/configs/DeepSeek-V3.2-DP_MI355.yaml
index afb4bed27447..ca7e0e7ef055 100644
--- a/tests/evals/gsm8k/configs/DeepSeek-V3.2-DP_MI355.yaml
+++ b/tests/evals/gsm8k/configs/DeepSeek-V3.2-DP_MI355.yaml
@@ -2,7 +2,7 @@ model_name: "deepseek-ai/DeepSeek-V3.2"
 accuracy_threshold: 0.95
 num_questions: 1319
 num_fewshot: 5
-startup_max_wait_seconds: 1200
+startup_max_wait_seconds: 1800
 env:
   VLLM_ROCM_USE_AITER: "1"
   VLLM_ROCM_USE_AITER_MLA: "1"
diff --git a/tests/evals/gsm8k/configs/DeepSeek-V3.2-TP_MI355.yaml b/tests/evals/gsm8k/configs/DeepSeek-V3.2-TP_MI355.yaml
index 5bc5fa9af08a..bda7f1d6a2c6 100644
--- a/tests/evals/gsm8k/configs/DeepSeek-V3.2-TP_MI355.yaml
+++ b/tests/evals/gsm8k/configs/DeepSeek-V3.2-TP_MI355.yaml
@@ -2,7 +2,7 @@ model_name: "deepseek-ai/DeepSeek-V3.2"
 accuracy_threshold: 0.95
 num_questions: 1319
 num_fewshot: 5
-startup_max_wait_seconds: 1200
+startup_max_wait_seconds: 1800
 env:
   VLLM_ROCM_USE_AITER: "1"
   VLLM_ROCM_USE_AITER_MLA: "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
deleted file mode 100644
index 0852d14633c2..000000000000
--- a/tests/evals/gsm8k/configs/humming/config-a100-act-int8-shard-0.txt
+++ /dev/null
@@ -1,4 +0,0 @@
-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
deleted file mode 100644
index 512cbe99ff34..000000000000
--- a/tests/evals/gsm8k/configs/humming/config-a100-act-int8-shard-1.txt
+++ /dev/null
@@ -1,4 +0,0 @@
-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
deleted file mode 100644
index bdfa9033df0c..000000000000
--- a/tests/evals/gsm8k/configs/humming/config-a100-shard-0.txt
+++ /dev/null
@@ -1,5 +0,0 @@
-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-a100-shard-1.txt b/tests/evals/gsm8k/configs/humming/config-a100-shard-1.txt
deleted file mode 100644
index e0260c504e75..000000000000
--- a/tests/evals/gsm8k/configs/humming/config-a100-shard-1.txt
+++ /dev/null
@@ -1,5 +0,0 @@
-Qwen3.6-35B-A3B-NVFP4-humming.yaml
-Qwen3-30B-A3B-Instruct-2507-quantized.w8a8-humming.yaml
-Qwen3-30B-A3B-Fp8-v1-humming.yaml
-Qwen3-30B-A3B-AWQ-humming.yaml
-Qwen3-0.6B-MXFP8-humming.yaml
diff --git a/tests/evals/gsm8k/configs/humming/config-a100-shard-2.txt b/tests/evals/gsm8k/configs/humming/config-a100-shard-2.txt
deleted file mode 100644
index 8453e5d73cab..000000000000
--- a/tests/evals/gsm8k/configs/humming/config-a100-shard-2.txt
+++ /dev/null
@@ -1,5 +0,0 @@
-Qwen3.5-35B-A3B-FP8-humming.yaml
-NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4-humming.yaml
-Qwen3-30B-A3B-MXFP4A16-humming.yaml
-Qwen3-30B-A3B-NVFP4-humming.yaml
-Qwen2-1.5B-Instruct-FP8W8-humming.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
deleted file mode 100644
index 737a92cf48ae..000000000000
--- a/tests/evals/gsm8k/configs/humming/config-h100-shard-0.txt
+++ /dev/null
@@ -1,6 +0,0 @@
-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
-Qwen2-1.5B-Instruct-FP8W8-humming.yaml
diff --git a/tests/evals/gsm8k/configs/humming/config-h100-shard-1.txt b/tests/evals/gsm8k/configs/humming/config-h100-shard-1.txt
deleted file mode 100644
index 0676dc7623f5..000000000000
--- a/tests/evals/gsm8k/configs/humming/config-h100-shard-1.txt
+++ /dev/null
@@ -1,5 +0,0 @@
-Qwen3.6-35B-A3B-NVFP4-humming.yaml
-Qwen3-30B-A3B-AWQ-humming.yaml
-Qwen3-30B-A3B-GPTQ-Int4-humming.yaml
-Qwen3-30B-A3B-FP8-block-humming.yaml
-Qwen3-0.6B-MXFP8-humming.yaml
diff --git a/tests/evals/gsm8k/configs/humming/config-h100-shard-2.txt b/tests/evals/gsm8k/configs/humming/config-h100-shard-2.txt
deleted file mode 100644
index 9d8fe128d0ff..000000000000
--- a/tests/evals/gsm8k/configs/humming/config-h100-shard-2.txt
+++ /dev/null
@@ -1,4 +0,0 @@
-NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4-humming.yaml
-Qwen3.5-35B-A3B-FP8-humming.yaml
-Qwen3-30B-A3B-MXFP4A16-humming.yaml
-gpt-oss-20b-humming.yaml
diff --git a/tests/evals/gsm8k/configs/humming/config-h100.txt b/tests/evals/gsm8k/configs/humming/config-h100.txt
new file mode 100644
index 000000000000..ba6a1e81da52
--- /dev/null
+++ b/tests/evals/gsm8k/configs/humming/config-h100.txt
@@ -0,0 +1,15 @@
+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
+Qwen2-1.5B-Instruct-FP8W8-humming.yaml
+Qwen3.6-35B-A3B-NVFP4-humming.yaml
+Qwen3-30B-A3B-AWQ-humming.yaml
+Qwen3-30B-A3B-GPTQ-Int4-humming.yaml
+Qwen3-30B-A3B-FP8-block-humming.yaml
+Qwen3-0.6B-MXFP8-humming.yaml
+NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4-humming.yaml
+Qwen3.5-35B-A3B-FP8-humming.yaml
+Qwen3-30B-A3B-MXFP4A16-humming.yaml
+gpt-oss-20b-humming.yaml
diff --git a/tests/evals/gsm8k/test_gsm8k_offloading.py b/tests/evals/gsm8k/test_gsm8k_offloading.py
index 9df8498d45a1..d1fa93484251 100644
--- a/tests/evals/gsm8k/test_gsm8k_offloading.py
+++ b/tests/evals/gsm8k/test_gsm8k_offloading.py
@@ -11,11 +11,13 @@
 it until success also waits for offloading to complete.
 
 Correctness is enforced via GSM8K accuracy on both runs: corrupt reloaded
-KV drops accuracy below threshold.  The reload itself is not asserted, so
-a silently skipped reload (e.g. offloading disabled) would not be flagged.
+KV drops accuracy below threshold.  The reload itself is only asserted for
+configs with ``assert_reload``, so elsewhere a silently skipped reload
+(e.g. offloading disabled) would not be flagged.
 
 Covers both KV offloading connectors (OffloadingConnector and
-SimpleCPUOffloadConnector) across five architecture families:
+SimpleCPUOffloadConnector) across five architecture families, plus
+UMBPStoreConnector on ROCm for Gemma 4:
   - Hybrid Mamba (NemotronH: attention + Mamba)
   - Heterogeneous head dim (Gemma 4)
   - Hybrid GDN (Qwen 3.5: attention + GatedDeltaNet)
@@ -35,6 +37,7 @@
 import requests
 
 from tests.utils import RemoteOpenAIServer
+from tests.v1.kv_connector.umbp_test_utils import mori_has_umbp
 from vllm.platforms import current_platform
 
 from .gsm8k_eval import evaluate_gsm8k
@@ -63,6 +66,17 @@ def _kv_transfer_config(
                 },
             }
         )
+    elif connector == "UMBPStoreConnector":
+        return json.dumps(
+            {
+                "kv_connector": "UMBPStoreConnector",
+                "kv_role": "kv_both",
+                "kv_connector_extra_config": {
+                    "mode": "embedded",
+                    "capacity_bytes": cpu_gib << 30,
+                },
+            }
+        )
     elif connector == "SimpleCPUOffloadConnector":
         return json.dumps(
             {
@@ -110,6 +124,17 @@ def _reset_gpu_prefix_cache(base_url: str) -> None:
         _force_engine_step(base_url)
 
 
+def _external_kv_tokens(base_url: str) -> float:
+    metrics = requests.get(f"{base_url}/metrics", timeout=30)
+    metrics.raise_for_status()
+    return sum(
+        float(line.rsplit(" ", 1)[1])
+        for line in metrics.text.splitlines()
+        if line.startswith("vllm:prompt_tokens_by_source_total{")
+        and 'source="external_kv_transfer"' in line
+    )
+
+
 @dataclass
 class OffloadingModelConfig:
     id: str
@@ -122,6 +147,7 @@ class OffloadingModelConfig:
     cpu_offload_gib: int = 4
     startup_timeout: int = 600
     spec_name: str = "CPUOffloadingSpec"
+    assert_reload: bool = False
 
 
 MODELS = [
@@ -228,6 +254,21 @@ class OffloadingModelConfig:
         cpu_offload_gib=16,
         startup_timeout=1200,
     ),
+    # ── UMBPStoreConnector (embedded MORI DRAM store) ────────────────
+    pytest.param(
+        OffloadingModelConfig(
+            id="umbp-gemma-4-e4b-it",
+            model="google/gemma-4-E4B-it",
+            connector="UMBPStoreConnector",
+            accuracy_threshold=0.55,
+            assert_reload=True,
+        ),
+        id="umbp-gemma-4-e4b-it",
+        marks=pytest.mark.skipif(
+            not (current_platform.is_rocm() and mori_has_umbp()),
+            reason="requires ROCm and MORI built with UMBP",
+        ),
+    ),
 ]
 
 
@@ -288,3 +329,8 @@ def test_gsm8k_offloading_correctness(cfg: OffloadingModelConfig):
 
             if run_idx == 1:
                 _reset_gpu_prefix_cache(base_url)
+                reloaded_before = _external_kv_tokens(base_url)
+            elif cfg.assert_reload:
+                assert _external_kv_tokens(base_url) > reloaded_before, (
+                    f"GSM8K run 2/2 did not reload KV through {cfg.connector}"
+                )
diff --git a/tests/fusion/test_quant_activation_contract.py b/tests/fusion/test_quant_activation_contract.py
index da507e5caf50..81f26753aa65 100644
--- a/tests/fusion/test_quant_activation_contract.py
+++ b/tests/fusion/test_quant_activation_contract.py
@@ -16,6 +16,10 @@
     FlashInferCutedslMxfp8LinearKernel,
     FlashInferCutlassMxfp8LinearKernel,
 )
+from vllm.model_executor.kernels.linear.mxfp8.pytorch import TorchMxFp8LinearKernel
+from vllm.model_executor.kernels.linear.mxfp8.rocm_native import (
+    RocmDotScaledMxfp8LinearKernel,
+)
 from vllm.model_executor.kernels.linear.nvfp4.base import (
     NvFp4LinearKernel,
     NvFp4LinearLayerConfig,
@@ -68,7 +72,10 @@
     AiterPerTokenFp8ScaledMMLinearKernel,
     FlashInferCutedslMxfp8LinearKernel,
     FlashInferCutlassMxfp8LinearKernel,
+    RocmDotScaledMxfp8LinearKernel,
 }
+if current_platform.is_cuda():
+    SUPPORTING.add(TorchMxFp8LinearKernel)
 
 
 def _all_kernel_classes() -> list[type]:
diff --git a/tests/jit_monitor/conftest.py b/tests/jit_monitor/conftest.py
index 3f869995a999..caf5107308c5 100644
--- a/tests/jit_monitor/conftest.py
+++ b/tests/jit_monitor/conftest.py
@@ -19,6 +19,7 @@ def reset():
         jit_monitor._mode = "warn"
         jit_monitor._verbose = False
         jit_monitor._cutedsl_hook_installed = False
+        jit_monitor._flydsl_hook_installed = False
         jit_monitor._tilelang_hook_installed = False
         jit_monitor._tilelang_jitimpl_compile_depth = 0
 
diff --git a/tests/jit_monitor/test_flydsl_contract.py b/tests/jit_monitor/test_flydsl_contract.py
new file mode 100644
index 000000000000..51d7c3d2b970
--- /dev/null
+++ b/tests/jit_monitor/test_flydsl_contract.py
@@ -0,0 +1,19 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Pin the FlyDSL internals the JIT monitor hooks."""
+
+import pytest
+
+pytestmark = pytest.mark.cpu_test
+
+jit_function = pytest.importorskip("flydsl.compiler.jit_function")
+
+_ASM = """  gpu.module @kernels {
+    gpu.func @gdn_h_s4096(%arg0: memref<f32>) kernel {
+  func.func @launch(%arg0: memref<f32>) {
+"""
+
+
+def test_kernel_name_inference_still_selects_gpu_func_symbols():
+    """The monitor names the kernel, not the ``launch`` wrapper, from this."""
+    assert jit_function._infer_kernel_names_from_asm(_ASM) == ["gdn_h_s4096"]
diff --git a/tests/jit_monitor/test_hooks.py b/tests/jit_monitor/test_hooks.py
index 6d4f25314c08..f9151f0f0e12 100644
--- a/tests/jit_monitor/test_hooks.py
+++ b/tests/jit_monitor/test_hooks.py
@@ -89,6 +89,43 @@ def compile(self, *args, **kwargs):
     }
 
 
+def _fake_flydsl_import_modules():
+    """Fake Python's FlyDSL modules touched by ``jit_monitor.activate``."""
+
+    class FakeMlirCompiler:
+        compiled: list[str] = []
+
+        @classmethod
+        def compile(cls, module, *, arch="", func_name="", link_libs=None):
+            cls.compiled.append(func_name)
+            return f"compiled:{func_name}"
+
+    def _fake_infer_kernel_names_from_asm(asm):
+        return [
+            line.split("gpu.func @", 1)[1].split("(", 1)[0].strip()
+            for line in asm.splitlines()
+            if "gpu.func @" in line and " kernel" in line
+        ]
+
+    fake_jit_function = cast(Any, ModuleType("flydsl.compiler.jit_function"))
+    fake_jit_function.MlirCompiler = FakeMlirCompiler
+    fake_jit_function._infer_kernel_names_from_asm = _fake_infer_kernel_names_from_asm
+
+    fake_compiler = cast(Any, ModuleType("flydsl.compiler"))
+    fake_compiler.__path__ = []
+    fake_compiler.jit_function = fake_jit_function
+
+    fake_flydsl = cast(Any, ModuleType("flydsl"))
+    fake_flydsl.__path__ = []
+    fake_flydsl.compiler = fake_compiler
+
+    return {
+        "flydsl": fake_flydsl,
+        "flydsl.compiler": fake_compiler,
+        "flydsl.compiler.jit_function": fake_jit_function,
+    }
+
+
 @contextmanager
 def _patch_jit_modules(fake_knobs, *, cute_compile=_fake_cute_compile):
     """Patch the Triton and CuTeDSL imports touched by ``jit_monitor.activate``."""
@@ -100,6 +137,7 @@ def _patch_jit_modules(fake_knobs, *, cute_compile=_fake_cute_compile):
             {
                 "triton": fake_triton,
                 **_fake_cute_import_modules(cute_compile),
+                **_fake_flydsl_import_modules(),
                 **_fake_tilelang_import_modules(),
             },
         ),
@@ -321,6 +359,88 @@ def __call__(self, *args, **kwargs):
     warning_once.assert_called_once()
 
 
+# ------------------------------------------------------------------
+# FlyDSL hook
+# ------------------------------------------------------------------
+
+
+_KERNEL_IR = (
+    "  gpu.module @kernels {{\n"
+    "    gpu.func @gdn_h_s{width}(%arg0: memref<f32>) kernel {{\n"
+)
+
+
+def test_flydsl_compile_logs_warning_and_passes_through():
+    with _patch_jit_modules(_make_fake_knobs()):
+        from flydsl.compiler.jit_function import MlirCompiler
+
+        jit_monitor.activate()
+        with mock.patch.object(jit_monitor.logger, "warning_once") as warning_once:
+            result = MlirCompiler.compile(
+                _KERNEL_IR.format(width=4096), arch="gfx950", func_name="launch"
+            )
+
+    assert result == "compiled:launch"
+    assert MlirCompiler.compiled == ["launch"]
+    warning_once.assert_called_once()
+    msg = warning_once.call_args[0][0] % warning_once.call_args[0][1:]
+    assert "FlyDSL JIT compilation during inference" in msg
+    assert "gdn_h_s4096" in msg
+
+
+def test_flydsl_reports_each_width_sharing_a_launcher_name():
+    """Launchers often share the name ``launch``, so the width carries it."""
+    with _patch_jit_modules(_make_fake_knobs()):
+        from flydsl.compiler.jit_function import MlirCompiler
+
+        jit_monitor.activate()
+        with mock.patch.object(jit_monitor.logger, "warning_once") as warning_once:
+            for width in (4096, 8192):
+                MlirCompiler.compile(_KERNEL_IR.format(width=width), func_name="launch")
+
+    details = [call[0][-1] for call in warning_once.call_args_list]
+    assert details == [
+        " (kernels=gdn_h_s4096)",
+        " (kernels=gdn_h_s8192)",
+    ]
+
+
+def test_flydsl_unreadable_module_costs_the_detail_not_the_compile():
+    class _ExplodingModule:
+        def __str__(self):
+            raise ValueError("unreadable IR")
+
+    with _patch_jit_modules(_make_fake_knobs()):
+        from flydsl.compiler.jit_function import MlirCompiler
+
+        jit_monitor.activate()
+        with mock.patch.object(jit_monitor.logger, "warning_once") as warning_once:
+            result = MlirCompiler.compile(_ExplodingModule(), func_name="launch")
+
+    assert result == "compiled:launch"
+    warning_once.assert_called_once()
+    assert warning_once.call_args[0][-1] == ""
+
+
+def test_flydsl_error_mode_raises():
+    with _patch_jit_modules(_make_fake_knobs()):
+        from flydsl.compiler.jit_function import MlirCompiler
+
+        jit_monitor.activate(mode="error")
+        with pytest.raises(RuntimeError, match="FlyDSL JIT compilation"):
+            MlirCompiler.compile("module", func_name="fly_error_kernel")
+
+
+def test_flydsl_hook_is_not_double_wrapped():
+    with _patch_jit_modules(_make_fake_knobs()):
+        from flydsl.compiler.jit_function import MlirCompiler
+
+        jit_monitor.activate()
+        first = MlirCompiler.__dict__["compile"]
+        jit_monitor._setup_flydsl_jit_hook()
+        assert MlirCompiler.__dict__["compile"] is first
+
+
 # ------------------------------------------------------------------
 # TileLang hook
 # ------------------------------------------------------------------
diff --git a/tests/kernels/attention/minimax_m3_nvfp4_ref.py b/tests/kernels/attention/minimax_m3_nvfp4_ref.py
new file mode 100644
index 000000000000..05f264eeb882
--- /dev/null
+++ b/tests/kernels/attention/minimax_m3_nvfp4_ref.py
@@ -0,0 +1,148 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Torch reference of MiniMax-M3's NVFP4 K/V and index-K pages on ROCm.
+
+A 128-token page is ``[K side | V side]``; a side holds every KV head's e2m1
+data (8192 B per head), then every head's e4m3 block scales (1024 B per head).
+Within a head: K data ``[token][64 B]``, K scales in the FlyDSL prefill's lane
+order, V data ``[token octet][dim octet][8 dims][4 B]`` (a byte holds two
+consecutive tokens of one dim), V scales ``[group][token]``. An index page is
+``[128 x 64 B data][128 x 8 B scales]``.
+"""
+
+import torch
+
+BLOCK = 128
+D = 128
+E2M1_GRID = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0)
+
+
+def quantize_nvfp4(x: torch.Tensor, global_scale: float):
+    """E2M1 codes ``[T, 128]`` and e4m3 scale bytes ``[T, 8]`` of fp32 rows.
+
+    The CUDA encoder's scale math and a nearest-grid E2M1 rounding with ties to
+    the even code, written independently of the kernels' bit arithmetic.
+    """
+    groups = x.view(-1, 8, 16)
+    one_sixth = torch.tensor(1 / 6, dtype=torch.float32)
+    # fp32 throughout, like the kernels: 1 / 0.002 is not 500 in fp32.
+    gs = torch.tensor(global_scale, dtype=torch.float32, device=x.device)
+    scale = (groups.abs().amax(-1) * one_sixth * (1.0 / gs)).clamp(max=448)
+    scale_e4m3 = scale.to(torch.float8_e4m3fn)
+    scale_value = scale_e4m3.float()
+    quant = torch.where(scale_value != 0, 1.0 / (scale_value * gs), 0.0)
+    y = groups * quant[..., None]
+    grid = torch.tensor(E2M1_GRID, device=x.device)
+    dist = (y.abs()[..., None] - grid).abs()
+    is_min = dist == dist.amin(-1, keepdim=True)
+    ids = torch.arange(8, device=x.device)
+    # Neighbouring codes alternate parity, so every tie has an even candidate.
+    even_min = torch.where(is_min & (ids % 2 == 0), ids, 8).amin(-1)
+    mag = torch.where(even_min < 8, even_min, dist.argmin(-1))
+    codes = mag | (torch.signbit(y).long() << 3)
+    return codes.view(-1, D), scale_e4m3.view(torch.uint8).view(-1, 8)
+
+
+def dequant(data: torch.Tensor, scales: torch.Tensor) -> torch.Tensor:
+    """Token-major NVFP4 rows -> fp32 e2m1 x block scale, ``[..., 128]``."""
+    codes = torch.stack((data & 15, data >> 4), -1).flatten(-2).long()
+    grid = torch.tensor(E2M1_GRID, device=data.device)
+    mag = grid[codes & 7] * torch.where(codes >= 8, -1.0, 1.0)
+    s = scales.view(torch.float8_e4m3fn).float()
+    return (mag.unflatten(-1, (8, 16)) * s[..., None]).flatten(-2)
+
+
+def head_offsets(side: int, device):
+    """Byte offsets inside one head's regions of side 0 (K) or 1 (V).
+
+    Returns the data byte and nibble shift of each (token, dim) ``[128, 128]``
+    and the scale byte of each (token, group) ``[128, 8]``.
+    """
+    t = torch.arange(BLOCK, device=device)[:, None]
+    d = torch.arange(D, device=device)[None, :]
+    g = torch.arange(8, device=device)[None, :]
+    if side == 0:
+        r = t % 32
+        n = 4 * (r // 8) + r % 4
+        data = t * 64 + d // 2
+        shift = (d % 2) * 4 + 0 * t
+        scale = ((((t // 32) * 4 + g // 2) * 16 + n) * 2 + (r % 8) // 4) * 2 + g % 2
+    else:
+        data = ((t // 8 * 16 + d // 8) * 8 + d % 8) * 4 + (t % 8) // 2
+        shift = (t % 2) * 4 + 0 * d
+        scale = g * BLOCK + t
+    return data, shift, scale
+
+
+def _region_bases(side: int, heads: int, device):
+    """``[H, 1, 1]`` page offsets of each head's data and scale regions."""
+    h = torch.arange(heads, device=device)[:, None, None]
+    side0 = side * heads * BLOCK * 72
+    return side0 + h * BLOCK * 64, side0 + heads * BLOCK * 64 + h * BLOCK * 8
+
+
+def _flat(cache: torch.Tensor) -> torch.Tensor:
+    """Writable 1-D view of a cache whose block stride may exceed one page."""
+    n = (cache.shape[0] - 1) * cache.stride(0) + cache[0].numel()
+    return cache.as_strided((n,), (1,))
+
+
+def _insert_side(cache, rows, slots, side, global_scale):
+    tokens, heads, _ = rows.shape
+    codes, scales = quantize_nvfp4(rows.reshape(-1, D).float(), global_scale)
+    slots = slots.long()
+    keep = slots >= 0
+    codes = codes.view(tokens, heads, D)[keep].to(torch.uint8)
+    scales = scales.view(tokens, heads, 8)[keep]
+    blk, pos = slots[keep] // BLOCK, slots[keep] % BLOCK
+    data_off, shift, scale_off = head_offsets(side, rows.device)
+    data_base, scale_base = _region_bases(side, heads, rows.device)
+    page = (blk * cache.stride(0))[:, None, None]
+    byte = page + data_base[:, 0] + data_off[pos][:, None]  # [n, H, 128]
+    sh = shift[pos][:, None].expand_as(byte)
+    flat = _flat(cache)
+    for s in (0, 4):  # each byte takes one nibble per pass
+        m = sh == s
+        flat[byte[m]] = (flat[byte[m]] & (0xF0 >> s)) | (codes[m] << s)
+    flat[page + scale_base[:, 0] + scale_off[pos][:, None]] = scales
+
+
+def insert_nvfp4_kv(k, v, kv_cache, slots, k_scale, v_scale):
+    """Quantize K / V ``[T, Hkv, 128]`` into their pages; slots < 0 skip."""
+    _insert_side(kv_cache, k, slots, 0, float(k_scale))
+    _insert_side(kv_cache, v, slots, 1, float(v_scale))
+
+
+def insert_nvfp4_index(index_cache, keys, slots):
+    """Quantize index keys ``[T, 128]`` into their pages (no global scale)."""
+    codes, scales = quantize_nvfp4(keys.float(), 1.0)
+    slots = slots.long()
+    keep = slots >= 0
+    codes, scales = codes[keep].to(torch.uint8), scales[keep]
+    blk, pos = slots[keep] // BLOCK, slots[keep] % BLOCK
+    base = (blk * index_cache.stride(0))[:, None]
+    flat = _flat(index_cache)
+    flat[base + pos[:, None] * 64 + torch.arange(64, device=keys.device)] = codes[
+        :, 0::2
+    ] | (codes[:, 1::2] << 4)
+    flat[base + BLOCK * 64 + pos[:, None] * 8 + torch.arange(8, device=keys.device)] = (
+        scales
+    )
+
+
+def nvfp4_kv_cache_token_major(kv_cache: torch.Tensor):
+    """``(k_data, k_scale, v_data, v_scale)`` of a main cache, token-major.
+
+    Data ``[blocks, Hkv, 128, 64]`` (the low nibble the even dim), scales
+    ``[blocks, Hkv, 128, 8]``, all uint8.
+    """
+    blocks, heads = kv_cache.shape[0], kv_cache.shape[1] // 2
+    pages = kv_cache.reshape(blocks, -1)
+    out = []
+    for side in (0, 1):
+        data_off, shift, scale_off = head_offsets(side, kv_cache.device)
+        data_base, scale_base = _region_bases(side, heads, kv_cache.device)
+        codes = (pages[:, data_base + data_off] >> shift.to(torch.uint8)) & 0xF
+        out.append(codes[..., 0::2] | (codes[..., 1::2] << 4))
+        out.append(pages[:, scale_base + scale_off])
+    return tuple(out)
diff --git a/tests/kernels/attention/test_cache.py b/tests/kernels/attention/test_cache.py
index 149aedd30fc1..fded1554d92f 100644
--- a/tests/kernels/attention/test_cache.py
+++ b/tests/kernels/attention/test_cache.py
@@ -508,6 +508,192 @@ def quantize_and_dequantize(cache_dtype: str) -> torch.Tensor:
     assert four_over_six_mse < default_mse
 
 
+# MiMo-V2.6 Flash: K 192 / V 128, 4 KV heads, block sizes 16 and 64 in use.
+ASYM_HEAD_SIZE = 192
+ASYM_HEAD_SIZE_V = 128
+ASYM_NUM_KV_HEADS = 4
+ASYM_BLOCK_SIZES = [16, 64]
+ASYM_NUM_TOKENS = [6, 512]
+# "unified": one page holds K's head group then V's, addressed the way ROCm
+# AITER unified attention hands the kernel (transpose(1, 2).split(head_size)).
+# "grouped": the same page content laid out per token, which gives the other
+# (page-contiguous, head-strided) stride pattern. No in-tree backend produces it
+# today; it keeps both stride shapes covered for asymmetric head dims.
+ASYM_LAYOUTS = ["unified", "grouped"]
+
+
+def test_reshape_and_cache_flash_asymmetric_tiny_head() -> None:
+    """A single small head still fills the cache correctly.
+
+    num_heads * max(head_size, v_head_size) is below one warp here, so a launch
+    that sized the block from the element count would schedule zero warps and
+    the per-head loop would never advance.
+    """
+    set_random_seed(0)
+    device = "cuda:0"
+    torch.set_default_device(device)
+    torch.accelerator.set_device_index(device)
+
+    num_tokens, num_heads = 4, 1
+    head_size, v_head_size = 16, 8
+    block_size, num_blocks = 16, 2
+    content = head_size + v_head_size
+
+    packed = torch.zeros(
+        num_blocks, num_heads, block_size, content, dtype=torch.bfloat16
+    )
+    key_cache, value_cache = packed.transpose(1, 2).split(
+        [head_size, v_head_size], dim=-1
+    )
+    slot_mapping = torch.tensor([0, 5, 16, 31], dtype=torch.long)
+    key = torch.randn(num_tokens, num_heads, head_size, dtype=torch.bfloat16)
+    value = torch.randn(num_tokens, num_heads, v_head_size, dtype=torch.bfloat16)
+    k_scale = torch.tensor(1.0, dtype=torch.float32)
+    v_scale = torch.tensor(1.0, dtype=torch.float32)
+
+    ops.reshape_and_cache_flash(
+        key,
+        value,
+        key_cache,
+        value_cache,
+        slot_mapping,
+        "auto",
+        k_scale,
+        v_scale,
+    )
+
+    expected_key = torch.zeros_like(key_cache)
+    expected_value = torch.zeros_like(value_cache)
+    for token, slot in enumerate(slot_mapping.tolist()):
+        block, offset = slot // block_size, slot % block_size
+        expected_key[block, offset] = key[token]
+        expected_value[block, offset] = value[token]
+    torch.testing.assert_close(key_cache, expected_key)
+    torch.testing.assert_close(value_cache, expected_value)
+
+
+@pytest.mark.parametrize("kv_cache_layout", ASYM_LAYOUTS)
+@pytest.mark.parametrize("kv_scale_type", KV_SCALE_TYPES)
+@pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPE)
+@pytest.mark.parametrize("block_size", ASYM_BLOCK_SIZES)
+@pytest.mark.parametrize("num_tokens", ASYM_NUM_TOKENS)
+@torch.inference_mode()
+def test_reshape_and_cache_flash_asymmetric_head_sizes(
+    kv_cache_layout: str,
+    kv_scale_type: str,
+    kv_cache_dtype: str,
+    block_size: int,
+    num_tokens: int,
+) -> None:
+    """V is written with its own head width and head stride.
+
+    Both layouts keep K and V in one layer buffer, so only per-head geometry
+    differs between the two sides.
+    """
+    set_random_seed(0)
+    device = "cuda:0"
+    torch.set_default_device(device)
+    torch.accelerator.set_device_index(device)
+
+    num_heads = ASYM_NUM_KV_HEADS
+    num_blocks = num_tokens // block_size + 2
+    content = ASYM_HEAD_SIZE + ASYM_HEAD_SIZE_V
+    # A kernel that copies K's width into V runs past the page, so leave enough
+    # slack to turn that into a failed assertion instead of an out-of-bounds
+    # write that takes the whole test process down.
+    guard = num_heads * content
+    cache_dtype = torch.uint8 if kv_cache_dtype == "fp8" else torch.bfloat16
+    # The last slot is always written: a too-wide copy would run past the end.
+    num_slots = block_size * num_blocks
+    slot_mapping = torch.randperm(num_slots - 1, device=device)[: num_tokens - 1]
+    slot_mapping = torch.cat(
+        [slot_mapping, torch.tensor([num_slots - 1], device=device)]
+    ).to(torch.long)
+
+    key = torch.randn(
+        num_tokens, num_heads, ASYM_HEAD_SIZE, dtype=torch.bfloat16, device=device
+    )
+    value = torch.randn(
+        num_tokens, num_heads, ASYM_HEAD_SIZE_V, dtype=torch.bfloat16, device=device
+    )
+
+    # K and V are views of one packed layer buffer, as in vLLM.
+    content = ASYM_HEAD_SIZE + ASYM_HEAD_SIZE_V
+    if kv_cache_layout == "unified":
+        # [num_blocks, num_heads, block_size, content], split per side.
+        num_bytes = num_blocks * num_heads * block_size * content
+        buf = torch.zeros(num_bytes + guard, dtype=cache_dtype, device=device)
+        packed = buf[:num_bytes].view(num_blocks, num_heads, block_size, content)
+        key_cache, value_cache = packed.transpose(1, 2).split(
+            [ASYM_HEAD_SIZE, ASYM_HEAD_SIZE_V], dim=-1
+        )
+    else:
+        # [num_blocks, block_size, num_heads, content], K group then V group.
+        num_bytes = num_blocks * block_size * num_heads * content
+        buf = torch.zeros(num_bytes + guard, dtype=cache_dtype, device=device)
+        packed = buf[:num_bytes].view(num_blocks, block_size, num_heads, content)
+        key_cache, value_cache = packed.split(
+            [ASYM_HEAD_SIZE, ASYM_HEAD_SIZE_V], dim=-1
+        )
+
+    if kv_cache_dtype == "fp8":
+        # Compare in fp16 after dequantizing, as the symmetric case does.
+        expected_key = torch.zeros(key_cache.shape, dtype=torch.float16, device=device)
+        expected_value = torch.zeros(
+            value_cache.shape, dtype=torch.float16, device=device
+        )
+    else:
+        expected_key = torch.zeros_like(key_cache)
+        expected_value = torch.zeros_like(value_cache)
+    for token, slot in enumerate(slot_mapping.tolist()):
+        block, offset = slot // block_size, slot % block_size
+        expected_key[block, offset] = key[token]
+        expected_value[block, offset] = value[token]
+
+    scale_shape = (1,) if kv_scale_type == "tensor" else (num_heads,)
+    k_scale = torch.ones(scale_shape, dtype=torch.float32, device=device)
+    v_scale = torch.ones(scale_shape, dtype=torch.float32, device=device)
+
+    ops.reshape_and_cache_flash(
+        key,
+        value,
+        key_cache,
+        value_cache,
+        slot_mapping,
+        kv_cache_dtype,
+        k_scale,
+        v_scale,
+    )
+    torch.accelerator.synchronize()
+
+    if kv_cache_dtype == "fp8":
+
+        def _dequant(cache, scale):
+            # gfx94x stores fp8 as e4m3fnuz; reinterpret with the platform dtype
+            # so the bytes are decoded the way the kernel wrote them.
+            values = cache.view(current_platform.fp8_dtype()).to(torch.float16)
+            scale = scale.to(torch.float16)
+            if scale.numel() == 1:
+                return values * scale
+            # per-attn-head scales broadcast over the head dimension
+            return values * scale.view(1, 1, -1, 1)
+
+        result_key = _dequant(key_cache, k_scale)
+        result_value = _dequant(value_cache, v_scale)
+        torch.testing.assert_close(result_key, expected_key, atol=0.001, rtol=0.1)
+        torch.testing.assert_close(result_value, expected_value, atol=0.001, rtol=0.1)
+    else:
+        torch.testing.assert_close(key_cache, expected_key)
+        torch.testing.assert_close(value_cache, expected_value)
+        # Untouched slots stay zero, and nothing is written behind the cache.
+        assert torch.count_nonzero(key_cache) == num_tokens * num_heads * ASYM_HEAD_SIZE
+        assert (
+            torch.count_nonzero(value_cache)
+            == num_tokens * num_heads * ASYM_HEAD_SIZE_V
+        )
+    assert torch.count_nonzero(buf[num_bytes:]) == 0
+
+
 @pytest.mark.parametrize("dtype", DTYPES)
 @pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPE)
 @pytest.mark.parametrize("kv_cache_layout", CACHE_LAYOUTS)
@@ -1014,9 +1200,15 @@ def test_concat_and_cache_ds_mla(
 
 @pytest.mark.parametrize("device", CUDA_DEVICES)
 @pytest.mark.parametrize("block_size", [64, 256])
+@pytest.mark.parametrize(
+    "entry_size,row_padding,page_padding",
+    [(528, 0, 0), (656, 0, 0), (528, 128, 1), (656, 16, 1)],
+)
 @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."""
+def test_concat_and_cache_ds_mla_nope(
+    device: str, block_size: int, entry_size: int, row_padding: int, page_padding: int
+) -> None:
+    """NoPE matches zero RoPE and preserves unused slots and stride padding."""
     dtype = torch.bfloat16
     if current_platform.is_rocm():
         pytest.skip("concat_and_cache_mla doesn't support fp8_ds_mla on ROCm")
@@ -1032,28 +1224,51 @@ def test_concat_and_cache_ds_mla_nope(device: str, block_size: int) -> None:
     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(
+    storage = torch.full(
+        (num_blocks, block_size + page_padding, entry_size + row_padding),
+        0xFF,
+        dtype=torch.uint8,
+        device=device,
+    )
+    kv_cache = storage[:, :block_size, :entry_size]
+    ref_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)
+    storage.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()
+    torch.testing.assert_close(kv_cache, ref_cache[..., :entry_size], atol=0, rtol=0)
+    rows = kv_cache.reshape(num_blocks * block_size, entry_size)
+    if entry_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()
+    assert (storage[..., entry_size:] == 0xFF).all()
+    assert (storage[:, block_size:] == 0xFF).all()
+
+    if entry_size == 528:
+        with pytest.raises(RuntimeError, match="528-byte rows with pe_dim == 0"):
+            ops.concat_and_cache_mla(
+                kv_c, zero_k_pe, kv_cache, slot_mapping, "fp8_ds_mla", scale
+            )
+        with pytest.raises(RuntimeError, match="at least 656 bytes"):
+            ops.cp_gather_and_upconvert_fp8_kv_cache(
+                kv_cache,
+                torch.empty(1, 576, dtype=dtype, device=device),
+                torch.zeros(1, 1, dtype=torch.int32, device=device),
+                torch.zeros(1, dtype=torch.int32, device=device),
+                1,
+            )
 
 
 # Bytes per token for the nvfp4_ds_mla cache layout (see flashmla_sparse.py).
diff --git a/tests/kernels/attention/test_flashinfer_trtllm_attention.py b/tests/kernels/attention/test_flashinfer_trtllm_attention.py
index 074e53d014c9..957bd2fdc5eb 100644
--- a/tests/kernels/attention/test_flashinfer_trtllm_attention.py
+++ b/tests/kernels/attention/test_flashinfer_trtllm_attention.py
@@ -342,7 +342,7 @@ def test_flashinfer_trtllm_decode_with_baseline(
     if is_nvfp4_kv:
         rtol, atol = 1.0, 1.0  # nvfp4 has higher quantization error
     elif q_quant_dtype == FP8_DTYPE and o_quant_dtype == FP4_DTYPE:
-        rtol, atol = 7e-2, 9e-2
+        rtol, atol = 7e-2, 1e-1
     elif q_quant_dtype == FP8_DTYPE and o_quant_dtype == FP8_DTYPE:
         rtol, atol = 3e-2, 4e-2
     elif q_quant_dtype == FP8_DTYPE and o_quant_dtype == dtype:
@@ -637,7 +637,7 @@ def test_flashinfer_trtllm_prefill_with_baseline(
     elif q_quant_dtype == FP8_DTYPE and o_quant_dtype == FP8_DTYPE:
         rtol, atol = 4e-2, 6e-2
     elif q_quant_dtype == FP8_DTYPE and o_quant_dtype == dtype:
-        rtol, atol = 2e-2, 3e-2
+        rtol, atol = 2e-2, 6e-2
     else:
         rtol, atol = 1e-2, 1e-2
 
diff --git a/tests/kernels/attention/test_fused_mqa_topk.py b/tests/kernels/attention/test_fused_mqa_topk.py
new file mode 100644
index 000000000000..7042cedb6849
--- /dev/null
+++ b/tests/kernels/attention/test_fused_mqa_topk.py
@@ -0,0 +1,517 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Single-launch fused MQA-logits + top-k prefill vs the dense indexer path."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+import vllm.envs as envs
+from vllm import _custom_ops as ops
+from vllm.model_executor.kernels.attention.dsa.fused_mqa_topk import (
+    fused_mqa_topk_available,
+    fused_mqa_topk_min_keys,
+    fused_mqa_topk_prefill,
+    get_fused_mqa_topk_workspaces,
+)
+from vllm.model_executor.layers.sparse_attn_indexer import as_fp32_logits
+from vllm.utils.deep_gemm import (
+    fp8_fp4_mqa_logits,
+    mqa_logits_mx_inputs,
+    mqa_logits_use_mx_scales,
+)
+from vllm.v1.attention.backends.mla.indexer import (
+    DeepseekV32IndexerMetadata,
+    DeepseekV32IndexerPrefillChunkMetadata,
+    DeepseekV32IndexerPrefillMetadata,
+)
+
+requires_fused_mqa_topk = pytest.mark.skipif(
+    not torch.cuda.is_available() or not fused_mqa_topk_available(),
+    reason="Requires the SM100 fused MQA+top-k extension",
+)
+
+TOPKS = (512, 1024, 2048)
+TOPK = 2048
+LOGITS_BUDGET = 512 * 1024**2
+
+
+def _dense(q, k, scales, weights, starts, ends, topk=TOPK):
+    """Upstream dense path: budgeted DeepGEMM logits + top_k_per_row_prefill."""
+    k_scale = scales
+    if mqa_logits_use_mx_scales():
+        # Power-of-two scales as four equal UE8M0 bytes, like the K cache.
+        exp = torch.log2(scales).round().to(torch.int64) + 127
+        k_scale = (exp * 0x01010101).to(torch.int32)
+    rows = q.shape[0]
+    out = torch.full((rows, topk), -1, device=q.device, dtype=torch.int32)
+    lens = ends - starts
+    max_len = max(1, int(lens.max()))
+    step = max(1, LOGITS_BUDGET // (4 * max_len))
+    for i in range(0, rows, step):
+        j = min(rows, i + step)
+        q_scale, w = mqa_logits_mx_inputs(q[i:j], None, weights[i:j])
+        logits = as_fp32_logits(
+            fp8_fp4_mqa_logits(
+                (q[i:j], q_scale), (k, k_scale), w, starts[i:j], ends[i:j], max_len
+            ),
+            lens[i:j],
+        )
+        # Compressed logits: row r holds keys [starts[r], ends[r]) from column 0.
+        ops.top_k_per_row_prefill(
+            logits,
+            torch.zeros_like(lens[i:j]),
+            lens[i:j],
+            out[i:j],
+            j - i,
+            logits.stride(0),
+            logits.stride(1),
+            topk,
+        )
+    return out
+
+
+def _logits(q, k, scales, weights, starts, ends):
+    """fp32 reference scores: sum_h w * relu(q . k) * scale, -inf off-range."""
+    kf = k.float()
+    out = torch.zeros(q.shape[0], k.shape[0], device=q.device)
+    for h in range(q.shape[1]):
+        out += weights[:, h, None] * torch.relu(q[:, h].float() @ kf.T)
+    out *= scales[None]
+    n = torch.arange(k.shape[0], device=q.device)
+    off = (n[None] < starts[:, None]) | (n[None] >= ends[:, None])
+    return out.masked_fill(off, -torch.inf)
+
+
+def _check(rel, q, k, scales, weights, starts, ends, scores=True):
+    """Indices relative to the row start: count, bounds, uniqueness, scores.
+
+    ``scores`` compares the selected scores with the exact fp32 top-k; the
+    SM100 dense path rounds its logits to bf16, so it is checked without.
+    """
+    topk = rel.shape[1]
+    visible = (ends - starts).clamp(min=0)
+    valid = rel >= 0
+    assert (valid.sum(1) == visible.clamp(max=topk)).all(), "valid count"
+    assert (rel[~valid] == -1).all(), "padding must be -1"
+    idx = rel.long() + starts[:, None].long()
+    assert ((~valid) | ((idx >= starts[:, None]) & (idx < ends[:, None]))).all()
+    srt = idx.masked_fill(~valid, -1).sort().values
+    assert ((srt[:, 1:] < 0) | (srt.diff(dim=1) > 0)).all(), "duplicates"
+    if not scores:
+        return
+    # Check scores in row blocks so the reference logits stay small.
+    step = max(1, LOGITS_BUDGET // (4 * k.shape[0]))
+    for i in range(0, q.shape[0], step):
+        j = min(q.shape[0], i + step)
+        logits = _logits(q[i:j], k, scales, weights[i:j], starts[i:j], ends[i:j])
+        ref = logits.topk(topk, dim=1).values
+        chosen = logits.gather(1, idx[i:j].clamp(min=0)).masked_fill(
+            ~valid[i:j], -torch.inf
+        )
+        torch.testing.assert_close(
+            chosen.sort(descending=True).values, ref, rtol=2e-4, atol=1e-4
+        )
+
+
+def _inputs(rows, keys, seed=913):
+    torch.manual_seed(seed)
+    q = torch.randn(rows, 32, 128, device="cuda").to(torch.float8_e4m3fn)
+    k = torch.randn(keys, 128, device="cuda").to(torch.float8_e4m3fn)
+    scales = torch.ones(keys, device="cuda")
+    weights = torch.randn(rows, 32, device="cuda") / 32
+    return q, k, scales, weights
+
+
+@requires_fused_mqa_topk
+@pytest.mark.parametrize("topk", TOPKS)
+@pytest.mark.parametrize(
+    "keys,rows,case",
+    [
+        (4096, 64, "causal"),
+        (16384, 256, "causal"),
+        (16384, 16, "early"),
+        (32768, 16, "bounds"),
+        (65536, 16, "ties"),
+        (131072, 16, "hot"),
+        (262145, 9, "padding"),
+        (1048576, 16, "causal"),
+    ],
+)
+def test_single_request_matches_dense(keys, rows, case, topk):
+    q, k, scales, weights = _inputs(rows, keys)
+    starts = torch.zeros(rows, device="cuda", dtype=torch.int32)
+    ends = torch.arange(keys - rows + 1, keys + 1, device="cuda", dtype=torch.int32)
+    if case == "early":  # rows around each top-k's visible-key boundary
+        ends = torch.tensor(
+            [
+                0,
+                1,
+                511,
+                512,
+                513,
+                1023,
+                1024,
+                1025,
+                2047,
+                2048,
+                2049,
+                4096,
+                8192,
+                12000,
+                16000,
+                keys,
+            ],
+            device="cuda",
+            dtype=torch.int32,
+        )
+    elif case == "bounds":
+        starts.fill_(333)
+        ends[::2] = keys // 2
+    elif case == "ties":
+        weights.zero_()
+    elif case == "hot":
+        weights = weights.abs()
+        scales[:12288] = 4
+    out = torch.empty(rows, topk, device="cuda", dtype=torch.int32)
+    fused_mqa_topk_prefill(q, k, scales, weights, starts, ends, out)
+    _check(out, q, k, scales, weights, starts, ends)
+    # Same convention and valid counts as the dense path (key sets may differ
+    # only among equal scores at the top-k boundary).
+    dense = _dense(q, k, scales, weights, starts, ends, topk)
+    _check(dense, q, k, scales, weights, starts, ends, scores=False)
+    assert ((out >= 0).sum(1) == (dense >= 0).sum(1)).all()
+
+
+@requires_fused_mqa_topk
+@pytest.mark.parametrize("topk", TOPKS)
+@pytest.mark.parametrize(
+    "lens", [[3000, 5000, 1000], [40000, 1500, 70000, 9], [600000, 200000]]
+)
+def test_multi_request_chunk_matches_dense(lens, topk):
+    """Packed requests: rows of request r see [cu[r], cu[r] + causal_end)."""
+    q_lens = [min(n, 1024) for n in lens]
+    rows, keys = sum(q_lens), sum(lens)
+    q, k, scales, weights = _inputs(rows, keys, seed=917)
+    starts, ends, base = [], [], 0
+    for n, m in zip(lens, q_lens):
+        starts += [base] * m
+        ends += list(range(base + n - m + 1, base + n + 1))
+        base += n
+    starts = torch.tensor(starts, device="cuda", dtype=torch.int32)
+    ends = torch.tensor(ends, device="cuda", dtype=torch.int32)
+    out = torch.empty(rows, topk, device="cuda", dtype=torch.int32)
+    fused_mqa_topk_prefill(q, k, scales, weights, starts, ends, out)
+    _check(out, q, k, scales, weights, starts, ends)
+    dense = _dense(q, k, scales, weights, starts, ends, topk)
+    _check(dense, q, k, scales, weights, starts, ends, scores=False)
+    assert ((out >= 0).sum(1) == (dense >= 0).sum(1)).all()
+    # Both paths agree on (nearly) every selected key. The dense path's bf16
+    # logits blur the k-th boundary, the more so for small k at long contexts
+    # (~98.6% overlap for top-512 over 600K keys; the fused_mqa_topk keys are
+    # exact).
+    overlap = sum(
+        len(set(a[a >= 0].tolist()) & set(b[b >= 0].tolist()))
+        for a, b in zip(out.cpu(), dense.cpu())
+    )
+    assert overlap >= 0.98 * int((dense >= 0).sum())
+
+
+@requires_fused_mqa_topk
+@pytest.mark.parametrize("rows", [32768, 32770])
+def test_large_chunk_single_call(rows):
+    """Prefill chunks above 16384 rows (e.g. --max-num-batched-tokens 32768)
+    run in one launch."""
+    keys = 65536
+    q, k, scales, weights = _inputs(rows, keys, seed=921)
+    starts = torch.zeros(rows, device="cuda", dtype=torch.int32)
+    ends = torch.arange(keys - rows + 1, keys + 1, device="cuda", dtype=torch.int32)
+    out = torch.empty(rows, TOPK, device="cuda", dtype=torch.int32)
+    fused_mqa_topk_prefill(q, k, scales, weights, starts, ends, out)
+    _check(out, q, k, scales, weights, starts, ends)
+
+
+@requires_fused_mqa_topk
+def test_rescan_scores_only_its_block():
+    """The exact rescan of a block touches only that block's 8 rows.
+
+    Regression test: the rescan's scheduler used a grid stride of 4096 blocks,
+    so with more than 4096 blocks (rows > 32768) rescanning block v also scored
+    block v + 4096, with row bounds read outside the rescanned block and output
+    written without the relative-index conversion. Block 0 is forced into the
+    rescan (a subnormal weight takes the exact path); block 4096 belongs to a
+    second request (nonzero starts, so its indices are made relative) and must
+    stay exact.
+    """
+    rows, keys_per_request = 33000, 8192  # 4125 blocks of 8 rows
+    split = rows // 2
+    q, k, scales, weights = _inputs(rows, 2 * keys_per_request, seed=927)
+    weights[0, 0] = 1e-40  # subnormal fp32: block 0 takes the exact rescan
+    starts = torch.zeros(rows, device="cuda", dtype=torch.int32)
+    starts[split:] = keys_per_request
+    ends = starts + keys_per_request
+    out = torch.full((rows, TOPK), -7, device="cuda", dtype=torch.int32)
+    fused_mqa_topk_prefill(q, k, scales, weights, starts, ends, out)
+    _check(out, q, k, scales, weights, starts, ends, scores=False)
+    # Scores of the rescanned block, the block 4096 blocks after it, and the
+    # last block.
+    check_rows = torch.cat(
+        [
+            torch.arange(0, 8),
+            torch.arange(4096 * 8, 4097 * 8),
+            torch.arange(rows - 8, rows),
+        ]
+    ).cuda()
+    _check(
+        out[check_rows],
+        q[check_rows],
+        k,
+        scales,
+        weights[check_rows],
+        starts[check_rows],
+        ends[check_rows],
+    )
+
+
+@requires_fused_mqa_topk
+def test_gate_with_common_score_offset():
+    """Regression test: scores sharing a large common offset with a narrow top
+    tail (top ranks within ~1 of each other at ~+68) must still get a tight gate.
+    With absolute fp16 sample bins (~2 wide there) the gate fell a whole bin too
+    low, rows overflowed the candidate slab and most blocks took the exact
+    rescan (4.4x slower at 256K keys)."""
+    rows, keys, topk = 2048, 65536, TOPK
+    q, k, scales, weights = _inputs(rows, keys, seed=5)
+    # A shared component in dimension 0 adds 8 * 8 to every q.k.
+    q = q.float()
+    k = k.float()
+    q[:, :, 0] = 8.0
+    k[:, 0] = 8.0
+    q, k = q.to(torch.float8_e4m3fn), k.to(torch.float8_e4m3fn)
+    weights = torch.full_like(weights, 1.0 / 32)
+    starts = torch.zeros(rows, device="cuda", dtype=torch.int32)
+    ends = torch.arange(keys - rows + 1, keys + 1, device="cuda", dtype=torch.int32)
+    out = torch.empty(rows, topk, device="cuda", dtype=torch.int32)
+    workspaces = get_fused_mqa_topk_workspaces(q.device, rows)
+    torch.ops._C.fused_mqa_topk_prefill(
+        q, k, scales, weights, starts, ends, out, *workspaces
+    )
+    _check(out, q, k, scales, weights, starts, ends)
+    # Per-row candidate counts the kernel left in its workspace: about the gate
+    # target (3072 for top-2048 at <= 128K keys), not a multiple of it.
+    counts = workspaces[-2][:rows]
+    assert (counts >= topk).all() and (counts <= 16384).all(), "rows rescanned"
+    assert counts.float().median() < 2 * 3072
+
+
+@requires_fused_mqa_topk
+def test_cuda_graph_replay():
+    rows, keys = 64, 65536
+    q, k, scales, weights = _inputs(rows, keys)
+    starts = torch.zeros(rows, device="cuda", dtype=torch.int32)
+    ends = torch.arange(keys - rows + 1, keys + 1, device="cuda", dtype=torch.int32)
+    out = torch.empty(rows, TOPK, device="cuda", dtype=torch.int32)
+    fn = lambda: fused_mqa_topk_prefill(  # noqa: E731
+        q, k, scales, weights, starts, ends, out
+    )
+    fn()
+    torch.accelerator.synchronize()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        fn()
+    out.fill_(-7)
+    for _ in range(3):
+        graph.replay()
+    torch.accelerator.synchronize()
+    _check(out, q, k, scales, weights, starts, ends)
+
+
+@requires_fused_mqa_topk
+def test_op_rejects_bad_inputs():
+    """The launcher does not re-check; the op enforces its contract."""
+    q, k, scales, weights = _inputs(16, 4096)
+    starts = torch.zeros(16, device="cuda", dtype=torch.int32)
+    ends = torch.full((16,), 4096, device="cuda", dtype=torch.int32)
+    out768 = torch.empty(16, 768, device="cuda", dtype=torch.int32)
+    with pytest.raises(RuntimeError, match="out must be"):
+        fused_mqa_topk_prefill(q, k, scales, weights, starts, ends, out768)
+    wide = torch.empty(16, 2 * TOPK, device="cuda", dtype=torch.int32)
+    with pytest.raises(RuntimeError, match="must be contiguous"):
+        fused_mqa_topk_prefill(q, k, scales, weights, starts, ends, wide[:, :TOPK])
+
+
+KPOOL = 4
+
+
+def _kpool_inputs(pool_lens=(20000, 70000), q_rows=96, tail=2, seed=919):
+    """Two packed GLM-5.3-Flash-style requests: K holds 4-token pools; the
+    last ``q_rows`` tokens of each request score the pools they can see."""
+    q, k, _, weights = _inputs(q_rows * len(pool_lens), sum(pool_lens), seed)
+    starts, ends, positions, base = [], [], [], 0
+    for n in pool_lens:
+        tokens = KPOOL * n + tail
+        pos = torch.arange(tokens - q_rows, tokens)
+        positions.append(pos)
+        starts += [base] * q_rows
+        ends += (base + (pos + 1) // KPOOL).tolist()
+        base += n
+    i32 = dict(device="cuda", dtype=torch.int32)
+    # Unit power-of-two scales as four UE8M0 bytes, like the SM100 K cache.
+    scales = torch.full((k.shape[0], 4), 127, device="cuda", dtype=torch.uint8)
+    return (
+        q,
+        k,
+        scales,
+        weights,
+        torch.tensor(starts, **i32),
+        torch.tensor(ends, **i32),
+        torch.cat(positions).cuda(),
+    )
+
+
+def _run_kpool_indexer(monkeypatch, inputs, fused_mqa_topk, budget_mb=512):
+    """Run the GLM-5.3-Flash kpool indexer prefill on one packed chunk."""
+    import vllm.models.glm5next.nvidia.sparse_indexer as kpool_indexer
+
+    q, k, scales, weights, starts, ends, positions = inputs
+    rows, keys = q.shape[0], k.shape[0]
+    chunk = DeepseekV32IndexerPrefillChunkMetadata(
+        block_table=torch.zeros(2, 1, device="cuda", dtype=torch.int32),
+        cu_seqlen_ks=starts,
+        cu_seqlen_ke=ends,
+        cu_seq_lens=torch.zeros(3, device="cuda", dtype=torch.int32),
+        token_to_seq=torch.zeros(keys, device="cuda", dtype=torch.int32),
+        total_seq_lens=keys,
+        max_seqlen_k=int((ends - starts).max()),
+        token_start=0,
+        token_end=rows,
+        num_reqs=2,
+        skip_kv_gather=True,  # K is already in the gather workspace
+        local_cu_seq_lens=torch.zeros(3, device="cuda", dtype=torch.int32),
+        local_total_seq_lens=keys,
+        max_local_total_seq_lens=keys,
+    )
+    metadata = DeepseekV32IndexerMetadata(
+        seq_lens=torch.zeros(2, device="cuda", dtype=torch.int32),
+        max_seq_len=int(positions.max()) + 1,
+        slot_mapping=torch.zeros(rows, device="cuda", dtype=torch.int64),
+        num_decodes=0,
+        num_decode_tokens=0,
+        num_prefills=2,
+        num_prefill_tokens=rows,
+        prefill=DeepseekV32IndexerPrefillMetadata(
+            chunks=[chunk], max_prefill_seq_len=int(positions.max()) + 1
+        ),
+    )
+    workspace = SimpleNamespace(get_simultaneous=lambda *specs: (k, scales))
+    monkeypatch.setattr(
+        kpool_indexer,
+        "get_forward_context",
+        lambda: SimpleNamespace(attn_metadata={"indexer": metadata}),
+    )
+    monkeypatch.setattr(kpool_indexer, "current_workspace_manager", lambda: workspace)
+    monkeypatch.setattr(envs, "VLLM_USE_FUSED_MQA_TOPK", fused_mqa_topk)
+    monkeypatch.setattr(envs, "VLLM_SPARSE_INDEXER_MAX_LOGITS_MB", budget_mb)
+    calls: list[torch.Size] = []
+    fused_fn = kpool_indexer.fused_mqa_topk_prefill
+
+    def spy(*a, **kw):
+        calls.append(a[-1].shape)  # the output buffer: [rows, top-k pools]
+        fused_fn(*a, **kw)
+
+    monkeypatch.setattr(kpool_indexer, "fused_mqa_topk_prefill", spy)
+    out = torch.empty(rows, TOPK + KPOOL - 1, device="cuda", dtype=torch.int32)
+    kpool_indexer.sparse_attn_indexer_kpool(
+        hidden_states=torch.empty(rows, 1, device="cuda"),
+        k_cache_prefix="indexer",
+        kv_cache=torch.empty(0, device="cuda"),
+        q_quant=q,
+        q_scale=None,
+        k=None,
+        weights=weights,
+        quant_block_size=128,
+        scale_fmt="ue8m0",
+        topk_tokens=TOPK,
+        head_dim=128,
+        max_pool_len=keys,
+        total_seq_lens=keys,
+        topk_indices_buffer=out,
+        skip_k_cache_insert=True,
+        index_kpool=KPOOL,
+        positions=positions,
+    )
+    return out, calls
+
+
+def _row_sets(out):
+    return [set(r[r >= 0].tolist()) for r in out.cpu()]
+
+
+@requires_fused_mqa_topk
+def test_kpool_indexer_fused_matches_dense(monkeypatch):
+    """GLM-5.3-Flash kpool indexer: the fused_mqa_topk kernel ranks the top-512 pools of
+    a packed multi-request chunk like the dense logits + top-k path."""
+    inputs = _kpool_inputs()
+    # 192 rows take the kernel only when dense needs more than one logits
+    # call, so shrink the budget below one call.
+    fused_mqa_topk_out, calls = _run_kpool_indexer(
+        monkeypatch, inputs, fused_mqa_topk=True, budget_mb=16
+    )
+    assert calls == [torch.Size([inputs[0].shape[0], TOPK // KPOOL])]
+    dense, calls = _run_kpool_indexer(monkeypatch, inputs, fused_mqa_topk=False)
+    assert calls == []
+    assert ((fused_mqa_topk_out >= 0).sum(1) == (dense >= 0).sum(1)).all()
+    # Same pools up to the dense path's bf16 boundary blur (see above); every
+    # row also keeps its request's incomplete tail tokens.
+    overlap = sum(
+        len(a & b) for a, b in zip(_row_sets(fused_mqa_topk_out), _row_sets(dense))
+    )
+    assert overlap >= 0.98 * int((dense >= 0).sum())
+
+
+@requires_fused_mqa_topk
+def test_kpool_indexer_short_chunk_stays_dense(monkeypatch):
+    """Chunks whose rows see fewer than fused_mqa_topk_min_keys pools keep the
+    dense path (the fused_mqa_topk kernel is not faster there)."""
+    inputs = _kpool_inputs(pool_lens=(9000, 20000))
+    assert int((inputs[5] - inputs[4]).max()) < fused_mqa_topk_min_keys(TOPK // KPOOL)
+    _, calls = _run_kpool_indexer(monkeypatch, inputs, fused_mqa_topk=True)
+    assert calls == []
+
+
+@requires_fused_mqa_topk
+def test_kpool_indexer_dense_fallback_splits_rows(monkeypatch):
+    """Chunks planned for fused_mqa_topk skip the logits budget; when such a
+    chunk takes the dense path, its rows are split to stay within the budget
+    and select the
+    same pools as one unbudgeted dense pass."""
+    inputs = _kpool_inputs()
+    whole, _ = _run_kpool_indexer(monkeypatch, inputs, fused_mqa_topk=False)
+    split, calls = _run_kpool_indexer(
+        monkeypatch, inputs, fused_mqa_topk=False, budget_mb=1
+    )
+    assert calls == []
+    # top_k_per_row_prefill breaks bf16 ties at the k-th score arbitrarily (two
+    # identical dense runs differ too), so compare the selected scores.
+    q, k, scales, weights, starts, ends, _ = inputs
+    lens = ends - starts
+    q_scale, w = mqa_logits_mx_inputs(q, None, weights)
+    logits = as_fp32_logits(
+        fp8_fp4_mqa_logits(
+            (q, q_scale),
+            (k, scales.view(torch.int32).squeeze(-1)),
+            w,
+            starts,
+            ends,
+            int(lens.max()),
+        ),
+        lens,
+    )
+    for r, (a, b) in enumerate(zip(_row_sets(split), _row_sets(whole))):
+        # Request-relative token ids -> pools; the appended tail is not scored.
+        pools = [
+            sorted({t // KPOOL for t in s if t // KPOOL < int(lens[r])}) for s in (a, b)
+        ]
+        scores = [logits[r, p].sort().values for p in pools]
+        assert torch.equal(scores[0], scores[1]), r
diff --git a/tests/kernels/attention/test_minimax_m3.py b/tests/kernels/attention/test_minimax_m3.py
index 8f0f38be5d6c..29c0e3d24aa3 100644
--- a/tests/kernels/attention/test_minimax_m3.py
+++ b/tests/kernels/attention/test_minimax_m3.py
@@ -627,13 +627,18 @@ def _reference_decode_index_score(
 
 # 3600: 29 blocks over 12 splits of 3, so the last busy split has 2 blocks.
 @pytest.mark.parametrize("prefix_len", [1024, 3600])
-@pytest.mark.parametrize("force_split_k", [False, True])
-def test_prefill_index_topk_correctness(monkeypatch, prefix_len, force_split_k):
-    if force_split_k:
-        # Take the SM12.0-only split-K launch on any CUDA GPU.
-        monkeypatch.setattr(
-            current_platform, "is_device_capability", lambda cap, *_: cap == (12, 0)
-        )
+@pytest.mark.parametrize("device_capability", [80, 90, 120])
+def test_prefill_index_topk_correctness(monkeypatch, prefix_len, device_capability):
+    monkeypatch.setattr(
+        current_platform,
+        "is_device_capability",
+        lambda cap, *_: device_capability == cap[0] * 10 + cap[1],
+    )
+    monkeypatch.setattr(
+        current_platform,
+        "is_device_capability_family",
+        lambda cap, *_: device_capability // 10 == cap // 10,
+    )
     topk = 6
     init_blocks = 0
     local_blocks = 1
@@ -1174,6 +1179,407 @@ def test_aiter_indexer_impl_matches_reference(
     assert torch.equal(sparse_ctx.cpu(), expected_ctx)
 
 
+# NVFP4 index cache (SM100/SM103): MSA Q8KV4 decode on the packed pages, FP8
+# gather + Q8KV8 prefill. Both rank pages from e4m3_rn_satfinite(e2m1 * scale).
+_E2M1_VALUES = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0)
+_NVFP4_PAGE_WIDTH = HEAD_DIM // 2 + HEAD_DIM // 16
+_requires_q8_indexer = pytest.mark.skipif(
+    not (
+        current_platform.is_device_capability(100)
+        or current_platform.is_device_capability(103)
+        or current_platform.is_device_capability(107)
+    ),
+    reason="The MSA Q8KV4/Q8KV8 indexers require SM100, SM103 or SM107.",
+)
+
+
+def _random_nvfp4_index_pages(
+    num_pages: int, scale_max: float = 0.5, page_pad: int = 0
+) -> torch.Tensor:
+    """``[num_pages, 128, 72]`` uint8 pages: random E2M1 codes, then E4M3 scales."""
+    data_bytes = BLOCK_SIZE * HEAD_DIM // 2
+    page_bytes = BLOCK_SIZE * _NVFP4_PAGE_WIDTH
+    raw = torch.zeros(
+        num_pages, page_bytes + page_pad, dtype=torch.uint8, device="cuda"
+    )
+    raw[:, :data_bytes] = torch.randint(
+        0, 256, (num_pages, data_bytes), dtype=torch.uint8, device="cuda"
+    )
+    scale = torch.rand(num_pages, page_bytes - data_bytes, device="cuda") * scale_max
+    raw[:, data_bytes:page_bytes] = scale.to(torch.float8_e4m3fn).view(torch.uint8)
+    return raw[:, :page_bytes].view(num_pages, BLOCK_SIZE, _NVFP4_PAGE_WIDTH)
+
+
+def _dequant_nvfp4_index_pages(pages: torch.Tensor) -> torch.Tensor:
+    """Reference ``[P, 128, 128]`` e4m3: e2m1 * scale, clamped, rounded once."""
+    num_pages = pages.shape[0]
+    flat = pages.reshape(num_pages, -1)
+    data = flat[:, : BLOCK_SIZE * HEAD_DIM // 2].view(num_pages, BLOCK_SIZE, -1)
+    scale = flat[:, BLOCK_SIZE * HEAD_DIM // 2 :].view(torch.float8_e4m3fn)
+    codes = torch.stack((data & 0xF, data >> 4), dim=-1).reshape(
+        num_pages, BLOCK_SIZE, HEAD_DIM
+    )
+    magnitude = torch.tensor(_E2M1_VALUES, device="cuda")[(codes & 7).long()]
+    value = torch.where(codes >= 8, -magnitude, magnitude)
+    value = value * scale.float().view(num_pages, BLOCK_SIZE, -1).repeat_interleave(
+        16, dim=-1
+    )
+    return value.clamp_(-448.0, 448.0).to(torch.float8_e4m3fn)
+
+
+@_requires_q8_indexer
+@pytest.mark.parametrize("use_qmul4", [False, True])
+@pytest.mark.parametrize("scale_max", [0.5, 448.0])
+def test_gather_nvfp4_index_k_fp8_bitwise(
+    scale_max: float, use_qmul4: bool, monkeypatch
+):
+    """Each request's leading pages land at its compact offset, bit-identical
+    to the reference dequant (incl. saturation), and nothing else is written,
+    on both the QMUL4 and the fp32 dequant paths."""
+    import vllm.models.minimax_m3.nvidia.ops.gather_nvfp4_index_k as gather_mod
+    from vllm.models.minimax_m3.nvidia.ops.gather_nvfp4_index_k import (
+        minimax_m3_gather_nvfp4_index_k_fp8,
+    )
+
+    if use_qmul4 and not gather_mod._qmul4_supported(0):
+        pytest.skip("Triton's Blackwell ptxas emits PTX older than 9.4 (no QMUL4).")
+    monkeypatch.setattr(gather_mod, "_qmul4_supported", lambda _: use_qmul4)
+
+    torch.manual_seed(0)
+    num_physical = 40
+    kv = _random_nvfp4_index_pages(num_physical, scale_max, page_pad=32)
+    block_table = torch.randperm(num_physical, device="cuda", dtype=torch.int32)[
+        :36
+    ].view(3, 12)
+    num_pages = torch.tensor([5, 0, 12], dtype=torch.int32, device="cuda")
+    page_offsets = torch.cumsum(num_pages, 0, dtype=torch.int32) - num_pages
+    total = int(num_pages.sum())
+    out = torch.full((total + 1, BLOCK_SIZE, HEAD_DIM), 1.0, device="cuda").to(
+        torch.float8_e4m3fn
+    )
+
+    minimax_m3_gather_nvfp4_index_k_fp8(
+        kv, block_table, num_pages, page_offsets, 12, out
+    )
+
+    expected = _dequant_nvfp4_index_pages(kv)
+    for req, (count, offset) in enumerate(
+        zip(num_pages.tolist(), page_offsets.tolist())
+    ):
+        pages = block_table[req, :count].long()
+        assert torch.equal(
+            out[offset : offset + count].view(torch.uint8),
+            expected[pages].view(torch.uint8),
+        )
+    assert torch.all(out[total].float() == 1.0)
+
+
+def _assert_forced_tail_topk(
+    actual: torch.Tensor, scores: torch.Tensor, lengths: torch.Tensor
+) -> None:
+    """Check the Q8 top-k contract on ``[rows, 16]`` output.
+
+    ``scores[row, :lengths[row] - 1]`` are the reference history-page scores.
+    Rows with <= 16 candidates are ``0 .. length - 1`` then ``-1``. Longer rows
+    hold 15 distinct history pages, then the local page. The ranked pages must
+    be the top 15 up to the kernel's 16-bit score-key quantization.
+    """
+    actual = actual.long()
+    rows, max_pages = scores.shape
+    cols = torch.arange(TOPK, device=actual.device)
+    short = lengths <= TOPK
+    expected_short = torch.where(cols < lengths[:, None], cols, -1)
+    assert torch.equal(actual[short], expected_short[short])
+
+    long_rows = (~short).nonzero().flatten()
+    if long_rows.numel() == 0:
+        return
+    actual, scores, lengths = actual[long_rows], scores[long_rows], lengths[long_rows]
+    assert torch.equal(actual[:, -1], lengths - 1)
+    ranked = actual[:, :-1]
+    assert torch.all((ranked >= 0) & (ranked < (lengths - 1)[:, None]))
+    assert torch.all(ranked.sort(dim=1).values.diff(dim=1) > 0)
+    valid = torch.arange(max_pages, device=scores.device) < (lengths - 1)[:, None]
+    history = scores.masked_fill(~valid, float("-inf"))
+    high = history.amax(dim=1)
+    low = scores.masked_fill(~valid, float("inf")).amin(dim=1)
+    tol = 1.5 * (high - low) / 65534.0 + 1e-5 * high.abs()
+    threshold = history.topk(TOPK - 1, dim=1).values[:, -1]
+    assert torch.all(history.gather(1, ranked) >= (threshold - tol)[:, None])
+    must_pick = history > (threshold + tol)[:, None]
+    picked = torch.zeros_like(must_pick).scatter_(1, ranked, True)
+    assert not torch.any(must_pick & ~picked)
+
+
+@_requires_q8_indexer
+@pytest.mark.parametrize("num_idx_heads", [1, 2, 4])
+@pytest.mark.parametrize("decode_query_len", [1, 3, 8])
+def test_nvfp4_indexer_impl_matches_reference(
+    num_idx_heads, decode_query_len, monkeypatch, workspace_init
+):
+    """Drive MiniMaxM3IndexerNvfp4Impl through its metadata builder on a mixed
+    batch (short and full 8-query decodes, fresh and chunked prefills, a request
+    without history pages, several gather chunks) and check every row against
+    a reference computed from the dequantized NVFP4 keys."""
+    import vllm.models.minimax_m3.common.indexer as indexer_mod
+    from tests.v1.attention.utils import (
+        BatchSpec,
+        create_common_attn_metadata,
+        create_vllm_config,
+    )
+    from vllm.config import set_current_vllm_config
+    from vllm.forward_context import set_forward_context
+    from vllm.models.minimax_m3.nvidia.indexer_nvfp4 import (
+        MiniMaxM3IndexerNvfp4Impl,
+        MiniMaxM3IndexerNvfp4MetadataBuilder,
+    )
+
+    torch.manual_seed(0)
+    device = torch.device("cuda")
+    monkeypatch.setattr(indexer_mod, "get_tensor_model_parallel_world_size", lambda: 1)
+    vllm_config = create_vllm_config(
+        block_size=BLOCK_SIZE, max_model_len=8192, max_num_batched_tokens=8192
+    )
+    vllm_config.model_config.hf_config.sparse_attention_config = {
+        "sparse_num_index_heads": num_idx_heads
+    }
+
+    decode_seq_lens = [2305, 2561, 9]
+    prefill = [(2624, 64), (2720, 96), (100, 100), (1500, 1500)]
+    batch = BatchSpec(
+        seq_lens=decode_seq_lens + [s for s, _ in prefill],
+        query_lens=[decode_query_len] * len(decode_seq_lens) + [q for _, q in prefill],
+    )
+    common = create_common_attn_metadata(batch, BLOCK_SIZE, device)
+    num_blocks = common.block_table_tensor.shape[1]
+    num_physical = batch.batch_size * num_blocks
+    block_table = torch.randperm(num_physical, device=device, dtype=torch.int32)
+    common.block_table_tensor = block_table.view(batch.batch_size, num_blocks)
+    num_tokens = batch.compute_num_tokens()
+
+    index_cache = _random_nvfp4_index_pages(num_physical)
+    index_q = (
+        torch.randn(num_tokens, num_idx_heads * HEAD_DIM, device=device) * 0.5
+    ).to(torch.float8_e4m3fn)
+
+    impl_kwargs = dict(
+        num_kv_heads=num_idx_heads,
+        scale=HEAD_DIM**-0.5,
+        topk_blocks=TOPK,
+        sparse_block_size=BLOCK_SIZE,
+        num_index_heads=num_idx_heads,
+        index_head_dim=HEAD_DIM,
+        init_blocks=0,
+        local_blocks=1,
+        indexer_kv_dtype="nvfp4",
+    )
+    spec = MLAAttentionSpec(
+        block_size=BLOCK_SIZE,
+        num_kv_heads=1,
+        head_size=_NVFP4_PAGE_WIDTH,
+        dtype=torch.uint8,
+    )
+    with set_current_vllm_config(vllm_config):
+        impl = MiniMaxM3IndexerNvfp4Impl(prefix="idx_nvfp4", **impl_kwargs)
+        builder = MiniMaxM3IndexerNvfp4MetadataBuilder(
+            spec, [impl.index_cache.prefix], vllm_config, device
+        )
+    index_spec = impl.index_cache.get_kv_cache_spec(vllm_config)
+    assert isinstance(index_spec, MLAAttentionSpec)
+    assert index_spec.page_size_bytes == BLOCK_SIZE * _NVFP4_PAGE_WIDTH
+    assert index_spec.dtype == torch.uint8
+    # Spec decode would raise the threshold to 1 + num_speculative_tokens.
+    builder.reorder_batch_threshold = decode_query_len
+    # Fit the longest request's 21 history pages but not two such requests.
+    builder.max_gather_pages = 30
+    impl.index_cache.kv_cache = index_cache
+    impl.topk_indices_buffer = torch.full(
+        (num_tokens, num_idx_heads, TOPK), -2, dtype=torch.int32, device=device
+    )
+
+    # Planning reads only host-side metadata: any device sync is an error.
+    torch.cuda.set_sync_debug_mode("error")
+    try:
+        md = builder.build(0, common)
+    finally:
+        torch.cuda.set_sync_debug_mode("default")
+    assert md.num_decode_tokens == decode_query_len * len(decode_seq_lens)
+    assert len(md.prefill_chunks) == 3
+    with set_forward_context({impl.index_cache.prefix: md}, vllm_config):
+        assert impl(index_q) == (None, None)
+
+    # Reference: max over each history page's 128 tokens of q . k.
+    keys = _dequant_nvfp4_index_pages(index_cache).float()
+    q = index_q.float().view(num_tokens, num_idx_heads, HEAD_DIM)
+    max_pages = num_blocks
+    scores = torch.full(
+        (num_tokens, num_idx_heads, max_pages), float("-inf"), device=device
+    )
+    lengths = torch.empty(num_tokens, dtype=torch.long, device=device)
+    token = 0
+    for req, (seq_len, q_len) in enumerate(zip(batch.seq_lens, batch.query_lens)):
+        pages = (seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE
+        k = keys[common.block_table_tensor[req, :pages].long()]
+        rows = slice(token, token + q_len)
+        scores[rows, :, :pages] = torch.einsum("nhd,pkd->nhpk", q[rows], k).amax(dim=-1)
+        positions = torch.arange(seq_len - q_len, seq_len, device=device)
+        lengths[rows] = positions // BLOCK_SIZE + 1
+        token += q_len
+    _assert_forced_tail_topk(
+        impl.topk_indices_buffer.view(-1, TOPK),
+        scores.view(-1, max_pages),
+        lengths.repeat_interleave(num_idx_heads),
+    )
+
+
+@_requires_q8_indexer
+@pytest.mark.parametrize("decode_query_len", [1, 3, 8])
+@pytest.mark.parametrize("num_idx_heads", [1, 2, 4])
+def test_nvfp4_indexer_decode_cuda_graph_replay(
+    decode_query_len, num_idx_heads, monkeypatch, workspace_init
+):
+    """A full CUDA graph captures the decode indexer once per padded batch
+    size; re-planning a later step (new lengths, pages and keys, one
+    zero-length padding request) outside capture and replaying must give that
+    step's top-k, which fails if any captured tensor was re-allocated. The
+    decode reads index_q and writes the top-k in place, copying nothing."""
+    import vllm.models.minimax_m3.common.indexer as indexer_mod
+    from tests.v1.attention.utils import (
+        BatchSpec,
+        create_common_attn_metadata,
+        create_vllm_config,
+    )
+    from vllm.config import set_current_vllm_config
+    from vllm.forward_context import set_forward_context
+    from vllm.models.minimax_m3.nvidia.indexer_nvfp4 import (
+        MiniMaxM3IndexerNvfp4Impl,
+        MiniMaxM3IndexerNvfp4MetadataBuilder,
+    )
+
+    torch.manual_seed(0)
+    device = torch.device("cuda")
+    monkeypatch.setattr(indexer_mod, "get_tensor_model_parallel_world_size", lambda: 1)
+    vllm_config = create_vllm_config(
+        block_size=BLOCK_SIZE, max_model_len=8192, max_num_batched_tokens=8192
+    )
+    vllm_config.model_config.hf_config.sparse_attention_config = {
+        "sparse_num_index_heads": num_idx_heads
+    }
+    num_reqs, max_blocks = 4, 24
+    num_tokens = num_reqs * decode_query_len
+    num_physical = num_reqs * max_blocks
+    index_cache = _random_nvfp4_index_pages(num_physical)
+
+    with set_current_vllm_config(vllm_config):
+        impl = MiniMaxM3IndexerNvfp4Impl(
+            prefix="idx_nvfp4_graph",
+            num_kv_heads=num_idx_heads,
+            scale=HEAD_DIM**-0.5,
+            topk_blocks=TOPK,
+            sparse_block_size=BLOCK_SIZE,
+            num_index_heads=num_idx_heads,
+            index_head_dim=HEAD_DIM,
+            init_blocks=0,
+            local_blocks=1,
+            indexer_kv_dtype="nvfp4",
+        )
+        builder = MiniMaxM3IndexerNvfp4MetadataBuilder(
+            MLAAttentionSpec(
+                block_size=BLOCK_SIZE,
+                num_kv_heads=1,
+                head_size=_NVFP4_PAGE_WIDTH,
+                dtype=torch.uint8,
+            ),
+            [impl.index_cache.prefix],
+            vllm_config,
+            device,
+        )
+    builder.reorder_batch_threshold = decode_query_len
+    impl.index_cache.kv_cache = index_cache
+    impl.topk_indices_buffer = torch.full(
+        (num_tokens, num_idx_heads, TOPK), -2, dtype=torch.int32, device=device
+    )
+    index_q = torch.empty(num_tokens, num_idx_heads * HEAD_DIM, device=device).to(
+        torch.float8_e4m3fn
+    )
+
+    def step(seq_lens, query_lens):
+        common = create_common_attn_metadata(
+            BatchSpec(seq_lens=seq_lens, query_lens=query_lens), BLOCK_SIZE, device
+        )
+        common.block_table_tensor = torch.randperm(
+            num_physical, device=device, dtype=torch.int32
+        )[: num_reqs * max_blocks].view(num_reqs, max_blocks)
+        index_q.copy_(
+            (torch.randn(num_tokens, num_idx_heads * HEAD_DIM, device=device) * 0.5).to(
+                torch.float8_e4m3fn
+            )
+        )
+        return common, builder.build(0, common)
+
+    # Warm up (compiles), then capture the forward of a full batch.
+    common, md = step([2305, 2561, 2944, 700], [decode_query_len] * num_reqs)
+    with set_forward_context({impl.index_cache.prefix: md}, vllm_config):
+        impl(index_q)
+        with torch.profiler.profile(
+            activities=[
+                torch.profiler.ProfilerActivity.CPU,
+                torch.profiler.ProfilerActivity.CUDA,
+            ]
+        ) as prof:
+            impl(index_q)
+            torch.accelerator.synchronize()
+        copies = [
+            e.name for e in prof.events() if "Memcpy" in e.name or "copy_" in e.name
+        ]
+        assert not copies, copies
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            impl(index_q)
+
+    # A later step: the last request pads the graph batch with no tokens.
+    seq_lens = [2800, 1500, 3000, 0]
+    query_lens = [decode_query_len] * 3 + [0]
+    common, md = step(seq_lens, query_lens)
+    assert md.decode_q8kv4.num_reqs == num_reqs
+    # An eager step of another batch size runs before the replay; every decode
+    # wrapper shares one page-score scratch, which this one overwrites.
+    other = create_common_attn_metadata(
+        BatchSpec(seq_lens=[1800, 900], query_lens=[decode_query_len] * 2),
+        BLOCK_SIZE,
+        device,
+    )
+    other.block_table_tensor = common.block_table_tensor[:2].clone()
+    other_md = builder.build(0, other)
+    with set_forward_context({impl.index_cache.prefix: other_md}, vllm_config):
+        impl(index_q[: 2 * decode_query_len])
+    assert len(builder.decode_wrappers) == 2
+    impl.topk_indices_buffer.fill_(-2)
+    graph.replay()
+
+    num_real = 3 * decode_query_len
+    keys = _dequant_nvfp4_index_pages(index_cache).float()
+    q = index_q[:num_real].float().view(num_real, num_idx_heads, HEAD_DIM)
+    scores = torch.full(
+        (num_real, num_idx_heads, max_blocks), float("-inf"), device=device
+    )
+    lengths = torch.empty(num_real, dtype=torch.long, device=device)
+    for req in range(3):
+        pages = (seq_lens[req] + BLOCK_SIZE - 1) // BLOCK_SIZE
+        k = keys[common.block_table_tensor[req, :pages].long()]
+        rows = slice(req * decode_query_len, (req + 1) * decode_query_len)
+        scores[rows, :, :pages] = torch.einsum("nhd,pkd->nhpk", q[rows], k).amax(dim=-1)
+        positions = torch.arange(
+            seq_lens[req] - decode_query_len, seq_lens[req], device=device
+        )
+        lengths[rows] = positions // BLOCK_SIZE + 1
+    _assert_forced_tail_topk(
+        impl.topk_indices_buffer[:num_real].view(-1, TOPK),
+        scores.view(-1, max_blocks),
+        lengths.repeat_interleave(num_idx_heads),
+    )
+
+
 @pytest.mark.parametrize(
     ("decode_query_len", "max_decode_query_len"),
     [
diff --git a/tests/kernels/attention/test_minimax_m3_index_topk_decode.py b/tests/kernels/attention/test_minimax_m3_index_topk_decode.py
new file mode 100644
index 000000000000..afbd90c009f3
--- /dev/null
+++ b/tests/kernels/attention/test_minimax_m3_index_topk_decode.py
@@ -0,0 +1,228 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""MiniMax-M3 index top-k (FlyDSL) against AITER's ``pa_sparse_block_topk``
+block ids, bit for bit, on realistic and adversarial decode and prefill rows."""
+
+import math
+
+import pytest
+import torch
+
+from vllm.platforms import current_platform
+
+if not current_platform.is_rocm():
+    pytest.skip("ROCm (gfx950) FlyDSL kernels", allow_module_level=True)
+
+from vllm.platforms.rocm import on_gfx950  # noqa: E402
+
+if not on_gfx950():
+    pytest.skip("gfx950 instructions", allow_module_level=True)
+
+from vllm.models.minimax_m3.amd.ops.index_topk_decode.host import (  # noqa: E402
+    minimax_m3_index_topk_decode_indices,
+    minimax_m3_index_topk_prefill_indices,
+    supports_index_topk,
+)
+
+BLOCK = 128
+
+
+def _width(max_seq_len: int) -> int:
+    n = math.ceil(math.ceil(max_seq_len / BLOCK) / 64)
+    return (1 << (n - 1).bit_length()) * 64
+
+
+def _scores(kind, heads, rows, width, nblk, gen):
+    s = torch.randn(heads, rows, width, generator=gen)
+    if kind == "real":  # a few outliers, the init block 1e30, the last two 1e29
+        s = torch.where(torch.rand(s.shape, generator=gen) < 4.0 / width, s + 6, s)
+        for r in range(rows):
+            b = min(int(nblk[r]), width)
+            if b > 0:
+                s[:, r, max(b - 2, 0) : b] = 1e29
+                s[:, r, 0] = 1e30
+    elif kind == "ties":
+        s = torch.randint(-3, 4, s.shape, generator=gen).float()
+    elif kind == "equal":
+        s = torch.full_like(s, 0.5)
+    elif kind == "neginf":
+        s = torch.full_like(s, float("-inf"))
+    elif kind == "nan":
+        s[torch.rand(s.shape, generator=gen) < 0.3] = float("nan")
+        neg_nan = torch.tensor([-4194303], dtype=torch.int32).view(torch.float32)
+        s[torch.rand(s.shape, generator=gen) < 0.1] = neg_nan.item()
+    elif kind == "zeros":  # +0 / -0 / denormals
+        vals = torch.tensor([0.0, -0.0, 1e-40, -1e-40, 1e-45, -1e-45])
+        s = vals[torch.randint(0, 6, s.shape, generator=gen)]
+    elif kind == "sentinels":
+        s = torch.where(torch.rand(s.shape, generator=gen) < 0.2, 1e29, s)
+        s = torch.where(torch.rand(s.shape, generator=gen) < 0.05, 1e30, s)
+    elif kind == "inf":
+        r = torch.rand(s.shape, generator=gen)
+        s = torch.where(r < 0.1, float("inf"), s)
+        s = torch.where((r > 0.1) & (r < 0.2), float("-inf"), s)
+    elif kind == "rising":  # increasing with position
+        s = torch.arange(width, dtype=torch.float32).expand(heads, rows, width).clone()
+    # blocks past a row's length hold garbage the kernel must not read as scores
+    tail = torch.randn(s.shape, generator=gen) * 1e6
+    tail[torch.rand(s.shape, generator=gen) < 0.1] = float("nan")
+    live = torch.arange(width)[None, None, :] < nblk[None, :, None]
+    return torch.where(live, s, tail)
+
+
+def _topk_out(heads, rows, pad):
+    """``topk_idx`` and the tensor compared: with ``pad`` a view (non-zero row
+    start, larger head stride) of a larger persistent buffer, compared whole
+    (nothing outside the view is written)."""
+    if not pad:
+        t = torch.full((heads, rows, 16), -7, dtype=torch.int32, device="cuda")
+        return t, t
+    full = torch.full((heads, rows + 37, 16), -7, dtype=torch.int32, device="cuda")
+    return full[:, 5 : 5 + rows], full
+
+
+def _aiter_topk(score, bt, seq, msl, nkv, pad, **kw):
+    from aiter.ops.msa_attention import pa_sparse_block_topk
+
+    heads, rows, _ = score.shape
+    topk, compared = _topk_out(heads, rows, pad)
+    sbt = torch.empty(rows * nkv, 128, dtype=torch.int32, device="cuda")
+    sctx = torch.empty(rows * nkv, dtype=torch.int32, device="cuda")
+    pa_sparse_block_topk(
+        score, topk, bt, seq, sbt, sctx, max_seq_len=msl, block_size=BLOCK,
+        num_kv_heads=nkv, pages_per_block=8, **kw,
+    )  # fmt: skip
+    return compared
+
+
+SHAPES = [
+    # reqs, query_len, max_seq_len, (lo, hi) seq_len, num_kv_heads, index heads
+    (40, 4, 1 << 20, (30000, 210000), 1, 1),  # the QPS decode batch
+    (8, 4, 262144, (1, 262144), 1, 1),
+    (5, 4, 1 << 20, (1 << 19, 1 << 20), 1, 1),  # 4 windows of 2048 blocks
+    (16, 1, 1 << 20, (2040 * BLOCK, 2100 * BLOCK), 1, 1),  # the window edge
+    (12, 1, 4000, (1, 4000), 1, 1),
+    (4, 2, 300000, (1, 300000), 2, 2),
+    (7, 4, 20000, (20000, 40000), 1, 1),  # rows past max_seq_len
+    (256, 4, 1 << 20, (1, 3000), 1, 1),  # 1024 short rows
+    (32, 6, 204800, (190000, 204800), 4, 4),  # 6 rows, 4 heads
+    # rows past max_seq_len inside the padded width (2048 blocks): the window
+    # sized by max_seq_len must go on to the next one, not cut them
+    (8, 1, 200000, (200000, 262000), 4, 1),
+    (4, 2, 0, (0, 0), 1, 1),  # a zero max_seq_len (dummy batch)
+    (1, 1, 1, (0, 1), 1, 1),
+]
+KINDS = [
+    "real",
+    "ties",
+    "equal",
+    "neginf",
+    "nan",
+    "zeros",
+    "sentinels",
+    "inf",
+    "rising",
+]
+
+
+def _check_decode(kind, shape, pad=False):
+    reqs, ql, msl, (lo, hi), nkv, heads = shape
+    gen = torch.Generator().manual_seed(100 * KINDS.index(kind) + SHAPES.index(shape))
+    seq = torch.randint(lo, hi + 1, (reqs,), generator=gen)
+    if reqs > 1:
+        seq[-1] = 0  # a cudagraph padding request
+    tok = torch.arange(reqs * ql) % ql
+    kv = (seq.repeat_interleave(ql) - ql + tok + 1).clamp(min=0)
+    score = _scores(kind, heads, reqs * ql, _width(msl), (kv + BLOCK - 1) // BLOCK, gen)
+    # the table covers every row's blocks (rows may run past max_seq_len)
+    cols = max(1, math.ceil(max(msl, int(seq.max())) / BLOCK))
+    bt = torch.randint(0, 1 << 22, (reqs, cols), generator=gen).int().cuda()
+    score, seq = score.cuda(), seq.int().cuda()
+    ref = _aiter_topk(score, bt, seq, msl, nkv, pad, query_len=ql)
+    topk, got = _topk_out(heads, reqs * ql, pad)
+    minimax_m3_index_topk_decode_indices(
+        score, topk, seq, max_seq_len=msl, query_len=ql
+    )
+    torch.testing.assert_close(got, ref, rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("kind", KINDS)
+@pytest.mark.parametrize("shape", SHAPES)
+def test_index_topk_decode_matches_aiter(kind, shape):
+    _check_decode(kind, shape)
+
+
+@pytest.mark.parametrize("kind", ["real", "nan", "ties"])
+@pytest.mark.parametrize("shape", [SHAPES[0], SHAPES[5]])
+def test_index_topk_decode_padded_output(kind, shape):
+    _check_decode(kind, shape, pad=True)
+
+
+def test_index_topk_decode_all_padding():
+    """A captured dummy batch: every request padding (seq_len 0)."""
+    reqs, ql, msl = 48, 4, 1 << 20
+    score = torch.randn(1, reqs * ql, _width(msl), device="cuda")
+    bt = torch.zeros(reqs, msl // BLOCK, dtype=torch.int32, device="cuda")
+    seq = torch.zeros(reqs, dtype=torch.int32, device="cuda")
+    ref = _aiter_topk(score, bt, seq, msl, 1, False, query_len=ql)
+    topk, _ = _topk_out(1, reqs * ql, False)
+    minimax_m3_index_topk_decode_indices(
+        score, topk, seq, max_seq_len=msl, query_len=ql
+    )
+    torch.testing.assert_close(topk, ref, rtol=0, atol=0)
+
+
+PREFILL_SHAPES = [
+    # (q_len, seq_len) per request, max_seq_len, num_kv_heads, index heads
+    ([(1000, 1000)], 1000, 1, 1),  # a fresh prompt: every row diagonal
+    ([(700, 30000), (1500, 1500)], 30000, 4, 4),  # cached prefix + fresh
+    ([(2000, 9000), (513, 513), (3000, 20000)], 20000, 4, 1),  # ragged
+    ([(300, 262144), (64, 200000)], 262144, 4, 4),  # one window, 2048 blocks
+    ([(129, 2049 * BLOCK), (40, 1 << 20)], 1 << 20, 2, 2),  # multi-window rows
+    ([(1, 1), (5, 5), (128, 128), (129, 129)], 4096, 1, 1),  # tiny rows
+    # 3K new on a 200K prefix beside a fresh 3K prompt
+    ([(3072, 203072), (3072, 3072), (3072, 207872)], 207872, 4, 1),
+    ([(64, 262000), (3072, 203072)], 200000, 4, 1),  # rows past max_seq_len
+]
+
+
+def _check_prefill(kind, shape, pad=False):
+    reqs, msl, nkv, heads = shape
+    gen = torch.Generator().manual_seed(
+        100 * KINDS.index(kind) + PREFILL_SHAPES.index(shape)
+    )
+    # every request's chunk ends its sequence: row positions seq - q .. seq - 1
+    pos = torch.cat([torch.arange(s - q, s) for q, s in reqs])
+    rid = torch.cat([torch.full((q,), r) for r, (q, _) in enumerate(reqs)])
+    nvp = pos // BLOCK + 1
+    seq = torch.tensor([s for _, s in reqs])
+    score = _scores(kind, heads, pos.numel(), _width(msl), nvp, gen).cuda()
+    cols = max(1, math.ceil(max(msl, int(seq.max())) / BLOCK))
+    bt = torch.randint(0, 1 << 22, (len(reqs), cols), generator=gen).int().cuda()
+    nvp, rid, kvl = (x.int().cuda() for x in (nvp, rid, pos + 1))
+    ref = _aiter_topk(
+        score, bt, seq.int().cuda(), msl, nkv, pad,
+        num_valid_pages=nvp, row_req_id=rid, kv_lens=kvl,
+    )  # fmt: skip
+    topk, got = _topk_out(heads, pos.numel(), pad)
+    minimax_m3_index_topk_prefill_indices(
+        score, topk, max_seq_len=msl, num_valid_pages=nvp
+    )
+    torch.testing.assert_close(got, ref, rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("kind", KINDS)
+@pytest.mark.parametrize("shape", PREFILL_SHAPES)
+def test_index_topk_prefill_matches_aiter(kind, shape):
+    _check_prefill(kind, shape)
+
+
+@pytest.mark.parametrize("kind", ["real", "nan", "ties"])
+@pytest.mark.parametrize("shape", [PREFILL_SHAPES[2], PREFILL_SHAPES[4]])
+def test_index_topk_prefill_padded_output(kind, shape):
+    _check_prefill(kind, shape, pad=True)
+
+
+def test_supports_index_topk():
+    assert supports_index_topk(16, BLOCK)
+    assert not supports_index_topk(8, BLOCK) and not supports_index_topk(16, 64)
diff --git a/tests/kernels/attention/test_minimax_m3_msa_cutlass_sparse_decode.py b/tests/kernels/attention/test_minimax_m3_msa_cutlass_sparse_decode.py
index c7a4916737c0..47ae7854b227 100644
--- a/tests/kernels/attention/test_minimax_m3_msa_cutlass_sparse_decode.py
+++ b/tests/kernels/attention/test_minimax_m3_msa_cutlass_sparse_decode.py
@@ -4,6 +4,7 @@
 NVFP4 KV cache)."""
 
 import math
+from dataclasses import fields, replace
 from types import SimpleNamespace
 
 import pytest
@@ -29,11 +30,13 @@
 )
 from vllm.models.minimax_m3.nvidia.model import MiniMaxM3SparseAttention
 from vllm.models.minimax_m3.nvidia.msa_cutlass_sparse_decode import (
+    NVFP4_BLOCK_SCALE_SHIFT,
     MSACutlassDecodePlanCache,
     msa_cutlass_sparse_decode,
     nvfp4_kv_cache_views,
     prepare_decode_metadata,
     should_prepare_decode_metadata,
+    uses_q8kv4_decode,
 )
 from vllm.models.minimax_m3.nvidia.sparse_attention_msa import (
     MiniMaxM3SparseMSABackend,
@@ -41,6 +44,8 @@
     MiniMaxM3SparseMSAImpl,
     MiniMaxM3SparseMSAMetadataBuilder,
     MiniMaxM3SparseMSANvfp4Backend,
+    MiniMaxM3SparseMSANvfp4Impl,
+    MiniMaxM3SparseMSAPrefillMetadata,
 )
 from vllm.platforms import current_platform
 from vllm.v1.attention.backends.registry import AttentionBackendEnum
@@ -179,6 +184,26 @@ def test_msa_cutlass_decode_static_dispatch_requires_sm100(
     )
 
 
+def test_msa_cutlass_decode_static_dispatch_query_head_limits() -> None:
+    """FP8 plans fit the kv_mode 3 work list (65536 query-head rows); NVFP4
+    plans run on the Q8KV4 kernel, which takes 65535 query tokens."""
+
+    def planned(batch_size: int, kv_cache_dtype: str) -> bool:
+        return should_prepare_decode_metadata(
+            batch_size,
+            8,
+            decode_backend="cutlass",
+            num_q_heads=16,
+            num_kv_heads=1,
+            kv_cache_dtype=kv_cache_dtype,
+            page_size=BLOCK_SIZE,
+            topk_blocks=TOPK,
+        )
+
+    assert planned(512, "fp8") and not planned(513, "fp8")
+    assert planned(8191, "nvfp4") and not planned(8192, "nvfp4")
+
+
 @pytest.mark.parametrize("decode_backend", ["cutlass", "triton"])
 def test_msa_cutlass_decode_static_dispatch_nvfp4_has_no_fallback(
     decode_backend: str,
@@ -205,7 +230,7 @@ def test_nvfp4_selects_msa_nvfp4_backend_and_requires_sm100(
         num_kv_heads=1,
     )
     assert backend_cls is MiniMaxM3SparseMSANvfp4Backend
-    assert impl_cls is MiniMaxM3SparseMSAImpl
+    assert impl_cls is MiniMaxM3SparseMSANvfp4Impl
     with pytest.raises(ValueError, match="nvfp4_4over6"):
         select_main_backend_and_impl_cls(
             topk_blocks=TOPK,
@@ -229,12 +254,13 @@ def test_nvfp4_selects_msa_nvfp4_backend_and_requires_sm100(
 @pytest.mark.parametrize(
     ("backend", "expected"),
     [
+        (None, "cutlass"),
         (AttentionBackendEnum.CUTLASS_MSA, "cutlass"),
         (AttentionBackendEnum.TRITON_MSA, "triton"),
     ],
 )
 def test_msa_attention_backend_alias(
-    backend: AttentionBackendEnum,
+    backend: AttentionBackendEnum | None,
     expected: str,
 ) -> None:
     config = AttentionConfig(backend=backend)
@@ -275,6 +301,7 @@ def test_msa_metadata_builder_prepares_cutlass_for_regular_decode(
     seq_lens = torch.full((batch,), 257, dtype=torch.int32, device="cuda")
     base_metadata = SimpleNamespace(
         num_decodes=batch,
+        prefill=None,
         decode=SimpleNamespace(
             block_table=block_table,
             seq_lens=seq_lens,
@@ -300,6 +327,7 @@ def test_msa_metadata_builder_prepares_cutlass_for_regular_decode(
     builder.kv_cache_spec = SimpleNamespace(num_kv_heads=4)
     builder.kv_cache_dtype = "fp8_e4m3"
     builder.decode_backend = "cutlass"
+    builder.q8kv4_decode = False
     builder.msa_cutlass_plan_cache = object()
     builder.msa_cutlass_min_batch_size = None
 
@@ -315,6 +343,39 @@ def test_msa_metadata_builder_prepares_cutlass_for_regular_decode(
     assert metadata.decode.msa_cutlass is expected_metadata
 
 
+@pytest.mark.parametrize("is_sm107", [True, False])
+def test_msa_metadata_builder_reserves_prefill_workspace_only_on_sm107(
+    monkeypatch: pytest.MonkeyPatch,
+    is_sm107: bool,
+) -> None:
+    monkeypatch.setattr(
+        MiniMaxM3SparseMetadataBuilder, "__init__", lambda *args, **kwargs: None
+    )
+    monkeypatch.setattr(
+        current_platform,
+        "is_device_capability",
+        lambda capability: is_sm107 and capability == 107,
+    )
+    vllm_config = SimpleNamespace(
+        model_config=SimpleNamespace(
+            hf_text_config=SimpleNamespace(
+                num_attention_heads=16,
+                sparse_attention_config={"sparse_topk_blocks": TOPK},
+            )
+        ),
+        parallel_config=SimpleNamespace(tensor_parallel_size=1),
+        cache_config=SimpleNamespace(cache_dtype="nvfp4"),
+        attention_config=SimpleNamespace(minimax_m3_msa_decode_backend="cutlass"),
+        scheduler_config=SimpleNamespace(max_num_batched_tokens=256),
+    )
+
+    builder = MiniMaxM3SparseMSAMetadataBuilder(
+        SimpleNamespace(num_kv_heads=1), [], vllm_config, torch.device("cuda")
+    )
+
+    assert (builder.prefill_workspace is not None) == is_sm107
+
+
 def test_msa_cutlass_plan_cache_keys_query_len(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
@@ -382,7 +443,7 @@ def fake_build_plan(**kwargs):
     assert built_query_lens == [1, 2]
 
 
-def test_query_fp8_stays_valid_when_cutlass_plan_appears_on_replay(
+def test_query_stays_valid_when_cutlass_plan_appears_on_replay(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
     import vllm.envs as envs
@@ -430,20 +491,21 @@ def test_query_fp8_stays_valid_when_cutlass_plan_appears_on_replay(
     observed_ptrs: list[int] = []
 
     @eager_break_during_capture
-    def run_attention(query_fp8: torch.Tensor | None) -> None:
+    def run_attention(query: torch.Tensor) -> None:
         if impl.should_use_msa_decode(layer_name):
-            assert query_fp8 is not None
-            observed_ptrs.append(query_fp8.data_ptr())
+            assert query.dtype == torch.float8_e4m3fn
+            observed_ptrs.append(query.data_ptr())
 
     stream = torch.cuda.Stream()
 
     with torch.cuda.stream(stream), override_forward_context(forward_context):
         capture = BreakableCUDAGraphCapture()
         with capture:
-            query_fp8 = MiniMaxM3SparseAttention._allocate_query_fp8(attention, qkv)
-            if query_fp8 is not None:
-                query_fp8.zero_()
-            run_attention(query_fp8)
+            query = MiniMaxM3SparseAttention._allocate_query(
+                attention, num_tokens, qkv.dtype, qkv.device
+            )
+            query.zero_()
+            run_attention(query)
             qkv.zero_()
 
         assert capture.num_graphs == 2
@@ -797,18 +859,18 @@ def _make_nvfp4_kv_cache(
     return kv_cache, torch.cat(dequantized, dim=-1).to(torch.bfloat16)
 
 
-def _e4m3_staged_kv_reference(
+def _q8kv4_kv_reference(
     kv_reference: torch.Tensor, k_scale: torch.Tensor, v_scale: torch.Tensor
 ) -> torch.Tensor:
-    """The dequantized NVFP4 cache as an E4M3-query kernel reads it: each
-    code * block_scale rounded to E4M3, then the global scale."""
+    """The dequantized cache as the Q8KV4 kernels read it: block scales staged
+    by 2**NVFP4_BLOCK_SCALE_SHIFT, and each code * block_scale rounded to E4M3
+    for the FP8 tensor core."""
     sides = []
     for side, global_scale in zip(
         kv_reference.float().split(HEAD_DIM, dim=-1), (k_scale, v_scale)
     ):
-        sides.append(
-            (side / global_scale).to(torch.float8_e4m3fn).float() * global_scale
-        )
+        unit = global_scale * 2.0**NVFP4_BLOCK_SCALE_SHIFT
+        sides.append((side / unit).to(torch.float8_e4m3fn).float() * unit)
     return torch.cat(sides, dim=-1).to(torch.bfloat16)
 
 
@@ -862,6 +924,8 @@ def test_nvfp4_kv_cache_views_address_head_slots() -> None:
         pytest.param(64, 4, 1, 1, True, id="tp1-small-batch-query-len-1"),
         pytest.param(64, 4, 8, 8, False, id="tp1-query-len-8"),
         pytest.param(16, 1, 16, 4, True, id="tp4-query-len-4"),
+        # 1024 x 8 x 16 = 131072 query-head rows, past the kv_mode 3 work list.
+        pytest.param(16, 1, 512, 8, True, id="tp4-query-len-8-1024-requests"),
     ],
 )
 def test_msa_cutlass_decode_nvfp4_matches_triton_on_dequantized_cache(
@@ -871,8 +935,9 @@ def test_msa_cutlass_decode_nvfp4_matches_triton_on_dequantized_cache(
     query_len: int,
     capture: bool,
 ) -> None:
-    """CUTLASS decode reads vLLM's packed NVFP4 pages, including swizzled V
-    scales and device global scales, as their exact dequantized values."""
+    """Q8KV4 decode reads vLLM's packed NVFP4 pages, including swizzled V
+    scales and device global scales, as their exact dequantized values, and a
+    reused plan follows new lengths and zero-length padding requests."""
     torch.manual_seed(0)
     seq_lens_list = [257, 513] * num_request_pairs
     seq_lens_cpu = torch.tensor(seq_lens_list, dtype=torch.int32)
@@ -890,18 +955,27 @@ def test_msa_cutlass_decode_nvfp4_matches_triton_on_dequantized_cache(
     )
     query_fp8 = query.to(torch.float8_e4m3fn)
     topk = _make_topk(seq_lens_list, num_kv_heads, query_len)
-    expected = torch.empty_like(query)
-    minimax_m3_sparse_attn_decode(
-        query_fp8.to(torch.bfloat16),
-        kv_reference,
-        topk.transpose(0, 1),
-        block_table,
-        seq_lens,
-        num_kv_heads,
-        SM_SCALE,
-        expected,
-        query_len,
-    )
+    kv_reference = _q8kv4_kv_reference(kv_reference, k_scale, v_scale)
+
+    def reference(lens: list[int]) -> torch.Tensor:
+        # Padding requests (length 0) trail the batch and are not compared.
+        num_reqs = sum(1 for n in lens if n > 0)
+        rows = num_reqs * query_len
+        expected = torch.empty_like(query[:rows])
+        minimax_m3_sparse_attn_decode(
+            query_fp8[:rows].to(torch.bfloat16),
+            kv_reference,
+            topk[:rows].transpose(0, 1),
+            block_table[:num_reqs],
+            seq_lens[:num_reqs],
+            num_kv_heads,
+            SM_SCALE,
+            expected,
+            query_len,
+        )
+        return expected
+
+    expected = reference(seq_lens_list)
 
     assert should_prepare_decode_metadata(
         len(seq_lens_list),
@@ -913,16 +987,24 @@ def test_msa_cutlass_decode_nvfp4_matches_triton_on_dequantized_cache(
         page_size=BLOCK_SIZE,
         topk_blocks=TOPK,
     )
-    metadata = prepare_decode_metadata(
-        block_table,
-        seq_lens,
-        seq_lens_cpu,
-        query_len,
-        num_q_heads=num_q_heads,
-        num_kv_heads=num_kv_heads,
-        page_size=BLOCK_SIZE,
-        topk_blocks=TOPK,
-    )
+    plan_cache = MSACutlassDecodePlanCache()
+
+    def prepare(lens_cpu: torch.Tensor):
+        return prepare_decode_metadata(
+            block_table,
+            seq_lens,
+            lens_cpu,
+            query_len,
+            num_q_heads=num_q_heads,
+            num_kv_heads=num_kv_heads,
+            page_size=BLOCK_SIZE,
+            topk_blocks=TOPK,
+            q8kv4=True,
+            plan_cache=plan_cache,
+        )
+
+    metadata = prepare(seq_lens_cpu)
+    assert "q8kv4" in metadata.plan[3]
     actual = torch.zeros_like(query)
 
     def run() -> None:
@@ -940,6 +1022,7 @@ def run() -> None:
             v_scale=v_scale,
         )
 
+    graph = None
     if capture:
         run()
         graph = torch.cuda.CUDAGraph()
@@ -952,6 +1035,25 @@ def run() -> None:
     current_platform.synchronize()
     torch.testing.assert_close(actual, expected, atol=0.05, rtol=0.05)
 
+    # Next step on the same plan: shorter requests and, as a padded full CUDA
+    # graph batch has, a trailing zero-length request.
+    updated_seq_lens_list = [129, 385] * num_request_pairs
+    updated_seq_lens_list[-1] = 0
+    seq_lens.copy_(torch.tensor(updated_seq_lens_list, dtype=torch.int32))
+    topk.copy_(_make_topk(updated_seq_lens_list, num_kv_heads, query_len))
+    updated_metadata = prepare(torch.tensor(updated_seq_lens_list, dtype=torch.int32))
+    assert updated_metadata.plan is metadata.plan
+    expected = reference(updated_seq_lens_list)
+    actual.zero_()
+    if graph is not None:
+        graph.replay()
+    else:
+        run()
+    current_platform.synchronize()
+    torch.testing.assert_close(
+        actual[: expected.shape[0]], expected, atol=0.05, rtol=0.05
+    )
+
 
 def _make_prefill_metadata(
     q_lens: list[int],
@@ -1003,18 +1105,27 @@ def _make_prefill_metadata(
     return metadata, topk
 
 
+def _skip_without_q8kv4_prefill(num_q_heads: int, num_kv_heads: int) -> None:
+    blocker = sparse_attention_msa_module._q8kv4_prefill_blocker(
+        num_q_heads, num_kv_heads, BLOCK_SIZE, TOPK
+    )
+    if blocker is not None:
+        pytest.skip(blocker)
+
+
 @pytest.mark.parametrize(
     ("num_q_heads", "num_kv_heads"),
     [pytest.param(64, 4, id="tp1"), pytest.param(16, 1, id="tp4")],
 )
 def test_msa_prefill_nvfp4_matches_triton_on_dequantized_cache(
-    monkeypatch: pytest.MonkeyPatch,
     num_q_heads: int,
     num_kv_heads: int,
 ) -> None:
-    """MSA prefill reads NVFP4 pages as their dequantized values, staged to E4M3
-    for the E4M3 q, including chunked prefill over cached context; it attends
-    with the fp8 q (the only one the fused insert writes on this path)."""
+    """MSA prefill reads NVFP4 pages as their exact dequantized values,
+    including chunked prefill over cached context, with the fp8 q the fused
+    insert writes (the only q it writes on this path), on the CUTLASS Q8KV4
+    kernel."""
+    _skip_without_q8kv4_prefill(num_q_heads, num_kv_heads)
     torch.manual_seed(0)
     q_lens = [300, 128, 77]
     seq_lens_list = [300, 640, 1100]
@@ -1025,10 +1136,10 @@ def test_msa_prefill_nvfp4_matches_triton_on_dequantized_cache(
         num_pages, num_kv_heads, "nvfp4", k_scale, v_scale
     )
 
+    total_q = sum(q_lens)
     metadata, topk = _make_prefill_metadata(
         q_lens, seq_lens_list, block_table, num_kv_heads
     )
-    total_q = sum(q_lens)
 
     layer_name = "model.layers.0.self_attn.attn"
     q_scale = 0.5
@@ -1046,7 +1157,7 @@ def test_msa_prefill_nvfp4_matches_triton_on_dequantized_cache(
         / q_scale
     ).to(torch.float8_e4m3fn)
     query = query_fp8.to(torch.bfloat16) * q_scale
-    kv_reference = _e4m3_staged_kv_reference(kv_reference, k_scale, v_scale)
+    kv_reference = _q8kv4_kv_reference(kv_reference, k_scale, v_scale)
     common = dict(
         num_heads=num_q_heads,
         head_size=HEAD_DIM,
@@ -1055,7 +1166,6 @@ def test_msa_prefill_nvfp4_matches_triton_on_dequantized_cache(
         topk_blocks=TOPK,
         sparse_block_size=BLOCK_SIZE,
     )
-    query_dtypes = _record_query_dtypes(monkeypatch, "sparse_atten_nvfp4_kv_func")
     forward_context = ForwardContext(
         no_compile_layers={},
         attn_metadata={layer_name: metadata},
@@ -1065,100 +1175,98 @@ def test_msa_prefill_nvfp4_matches_triton_on_dequantized_cache(
         expected = MiniMaxM3SparseTritonImpl(**common).forward(
             layer, query, kv_reference, torch.empty_like(query)
         )
-        actual = MiniMaxM3SparseMSAImpl(**common, kv_cache_dtype="nvfp4").forward(
-            layer,
-            torch.full_like(query, float("nan")),
-            kv_cache,
-            torch.zeros_like(query),
-            query_fp8=query_fp8,
+        # As the model calls it for NVFP4: the fp8 q is the only query.
+        actual = MiniMaxM3SparseMSANvfp4Impl(**common).forward(
+            layer, query_fp8, kv_cache, torch.zeros_like(query)
         )
     current_platform.synchronize()
-    assert query_dtypes == [torch.float8_e4m3fn]
     _assert_close_fp8_prefill(actual, expected)
 
 
-def test_msa_triton_decode_fallback_reads_dequantized_query_fp8() -> None:
-    """With CUTLASS decode on but not planned for the batch, the Triton decode
-    fallback attends with the dequantized fp8 q, the only q written."""
+@pytest.mark.parametrize(
+    ("num_q_heads", "num_kv_heads"),
+    [pytest.param(64, 4, id="tp1"), pytest.param(16, 1, id="tp4")],
+)
+def test_msa_prefill_nvfp4_uses_the_builders_workspace(
+    num_q_heads: int,
+    num_kv_heads: int,
+) -> None:
+    """The NVFP4 prefill hands the builder's split workspace to ``run_prefill``:
+    a workspace reserved for the step's token budget is neither replaced nor
+    grown by the call, and the output equals the per-call-allocation path's."""
+    from vllm.third_party.fmha_sm100.prefill_q8kv4 import PrefillWorkspace
+
+    _skip_without_q8kv4_prefill(num_q_heads, num_kv_heads)
     torch.manual_seed(0)
-    num_q_heads, num_kv_heads, q_scale = 16, 1, 0.5
-    seq_lens_list = [257, 513, 129]
-    seq_lens = torch.tensor(seq_lens_list, dtype=torch.int32, device="cuda")
+    q_lens = [300, 128, 77]
+    seq_lens_list = [300, 640, 1100]
     block_table, num_pages = _make_paged_layout(seq_lens_list)
-    kv_cache = (
-        torch.randn(
-            num_pages,
-            num_kv_heads,
-            BLOCK_SIZE,
-            2 * HEAD_DIM,
-            dtype=torch.bfloat16,
-            device="cuda",
-        )
-        * 0.25
-    ).to(torch.float8_e4m3fn)
-    num_tokens = len(seq_lens_list)
-    metadata = MiniMaxM3SparseMetadata(
-        seq_lens=seq_lens,
-        max_seq_len=max(seq_lens_list),
-        slot_mapping=torch.empty(num_tokens, dtype=torch.int64, device="cuda"),
-        num_actual_tokens=num_tokens,
-        num_decodes=num_tokens,
-        num_decode_tokens=num_tokens,
-        num_prefills=0,
-        num_prefill_tokens=0,
-        decode=MiniMaxM3SparseMSADecodeMetadata(
-            seq_lens=seq_lens,
-            block_table=block_table,
-            decode_query_len=1,
-            msa_cutlass=None,
+    k_scale = torch.tensor(1.0, dtype=torch.float32, device="cuda")
+    v_scale = torch.tensor(1.0, dtype=torch.float32, device="cuda")
+    kv_cache, _ = _make_nvfp4_kv_cache(
+        num_pages, num_kv_heads, "nvfp4", k_scale, v_scale
+    )
+    metadata, topk = _make_prefill_metadata(
+        q_lens, seq_lens_list, block_table, num_kv_heads
+    )
+    total_q = sum(q_lens)
+    max_num_batched_tokens = 2048
+    workspace = PrefillWorkspace(torch.device("cuda"), block_scale_shift=0)
+    workspace.split_buffers(TOPK, max_num_batched_tokens, num_q_heads)
+    reserved = workspace.split_buffers(TOPK, total_q, num_q_heads)[0].data_ptr()
+    assert isinstance(metadata.prefill, MiniMaxM3SparsePrefillMetadata)
+    with_workspace = replace(
+        metadata,
+        prefill=MiniMaxM3SparseMSAPrefillMetadata(
+            **{
+                f.name: getattr(metadata.prefill, f.name)
+                for f in fields(metadata.prefill)
+            },
+            workspace=workspace,
         ),
     )
+
     layer_name = "model.layers.0.self_attn.attn"
     layer = SimpleNamespace(
         layer_name=layer_name,
-        topk_indices_buffer=_make_topk(seq_lens_list, num_kv_heads, 1),
-        _q_scale_float=q_scale,
+        topk_indices_buffer=topk,
+        _q_scale_float=1.0,
+        _k_scale=k_scale,
+        _v_scale=v_scale,
     )
     query_fp8 = torch.randn(
-        num_tokens, num_q_heads * HEAD_DIM, dtype=torch.bfloat16, device="cuda"
+        total_q, num_q_heads * HEAD_DIM, dtype=torch.bfloat16, device="cuda"
     ).to(torch.float8_e4m3fn)
-    query = query_fp8.to(torch.bfloat16) * q_scale
-    common = dict(
+    impl = MiniMaxM3SparseMSANvfp4Impl(
         num_heads=num_q_heads,
         head_size=HEAD_DIM,
         scale=SM_SCALE,
         num_kv_heads=num_kv_heads,
-        kv_cache_dtype="fp8",
         topk_blocks=TOPK,
         sparse_block_size=BLOCK_SIZE,
     )
-    impl = MiniMaxM3SparseMSAImpl(**common, msa_decode_backend="cutlass")
-    assert impl.use_cutlass_decode
-    forward_context = ForwardContext(
-        no_compile_layers={},
-        attn_metadata={layer_name: metadata},
-        slot_mapping={},
-    )
-    with override_forward_context(forward_context):
-        expected = MiniMaxM3SparseTritonImpl(**common).forward(
-            layer, query, kv_cache, torch.empty_like(query)
+    outputs = []
+    for md in (metadata, with_workspace):
+        forward_context = ForwardContext(
+            no_compile_layers={},
+            attn_metadata={layer_name: md},
+            slot_mapping={},
         )
-        actual = impl.forward(
-            layer,
-            torch.full_like(query, float("nan")),
-            kv_cache,
-            torch.zeros_like(query),
-            query_fp8=query_fp8,
+        output = torch.zeros(
+            total_q, num_q_heads * HEAD_DIM, dtype=torch.bfloat16, device="cuda"
         )
+        with override_forward_context(forward_context):
+            outputs.append(impl.forward(layer, query_fp8, kv_cache, output))
     current_platform.synchronize()
-    torch.testing.assert_close(actual, expected, atol=0.02, rtol=0.02)
+    assert torch.equal(outputs[0], outputs[1])
+    assert workspace.split_buffers(TOPK, total_q, num_q_heads)[0].data_ptr() == reserved
 
 
 @pytest.mark.parametrize(
     ("num_q_heads", "num_kv_heads"),
     [pytest.param(64, 4, id="tp1"), pytest.param(16, 1, id="tp4")],
 )
-def test_msa_prefill_reads_query_fp8_over_fp8_cache(
+def test_msa_prefill_reads_the_e4m3_query_over_an_fp8_cache(
     monkeypatch: pytest.MonkeyPatch,
     num_q_heads: int,
     num_kv_heads: int,
@@ -1219,13 +1327,126 @@ def test_msa_prefill_reads_query_fp8_over_fp8_cache(
         expected = MiniMaxM3SparseTritonImpl(**common).forward(
             layer, query, kv_cache, torch.empty_like(query)
         )
-        actual = impl.forward(
-            layer,
-            torch.full_like(query, float("nan")),
-            kv_cache,
-            torch.zeros_like(query),
-            query_fp8=query_fp8,
-        )
+        # As the model calls it: the fp8 q is the only query the insert wrote.
+        actual = impl.forward(layer, query_fp8, kv_cache, torch.zeros_like(query))
     current_platform.synchronize()
     assert query_dtypes == [torch.float8_e4m3fn]
     _assert_close_fp8_prefill(actual, expected)
+
+
+def test_nvfp4_kv_cache_never_takes_a_bf16_query() -> None:
+    """NVFP4 attention always reads an E4M3 q on the CUTLASS Q8KV4 kernels, and
+    a head layout either kernel cannot serve fails at init."""
+    assert uses_q8kv4_decode(kv_cache_dtype="nvfp4", num_q_heads=16, num_kv_heads=1)
+    assert not uses_q8kv4_decode(
+        kv_cache_dtype="fp8_e4m3", num_q_heads=16, num_kv_heads=1
+    )
+    with pytest.raises(ValueError, match="Q8KV4"):
+        uses_q8kv4_decode(kv_cache_dtype="nvfp4", num_q_heads=4, num_kv_heads=1)
+    # 8 Q heads per KV head (TP8) decodes on Q8KV4, but Q8KV4 cannot prefill it.
+    with pytest.raises(ValueError, match="CUTLASS Q8KV4 prefill"):
+        MiniMaxM3SparseMSANvfp4Impl(
+            num_heads=8,
+            head_size=HEAD_DIM,
+            scale=SM_SCALE,
+            num_kv_heads=1,
+            topk_blocks=TOPK,
+            sparse_block_size=BLOCK_SIZE,
+        )
+
+
+@pytest.mark.parametrize(
+    "architecture",
+    ["MiniMaxM3SparseForCausalLM", "MiniMaxM3SparseForConditionalGeneration"],
+)
+def test_nvfp4_kv_cache_with_tp8_fails_at_config_time(architecture: str) -> None:
+    """An NVFP4 KV cache needs 16 query heads per KV head on each rank: the config
+    check rejects TP8 (8 per KV head) before the workers start; TP1/2/4 and other
+    KV cache dtypes pass."""
+    from vllm.model_executor.models.config import MODELS_CONFIG_MAP
+
+    def vllm_config(tensor_parallel_size: int, cache_dtype: str) -> SimpleNamespace:
+        text_config = SimpleNamespace(num_attention_heads=64, num_key_value_heads=4)
+        return SimpleNamespace(
+            model_config=SimpleNamespace(hf_text_config=text_config),
+            cache_config=SimpleNamespace(cache_dtype=cache_dtype),
+            parallel_config=SimpleNamespace(tensor_parallel_size=tensor_parallel_size),
+        )
+
+    verify = MODELS_CONFIG_MAP[architecture].verify_and_update_config
+    for tensor_parallel_size in (1, 2, 4):
+        verify(vllm_config(tensor_parallel_size, "nvfp4"))
+    verify(vllm_config(8, "fp8"))
+    with pytest.raises(ValueError, match="NVFP4 KV cache with tensor_parallel_size=8"):
+        verify(vllm_config(8, "nvfp4"))
+
+
+def test_msa_triton_decode_fallback_reads_dequantized_query_fp8() -> None:
+    """With CUTLASS decode on but not planned for the batch, the Triton decode
+    fallback attends with the dequantized E4M3 query."""
+    torch.manual_seed(0)
+    num_q_heads, num_kv_heads, q_scale = 16, 1, 0.5
+    seq_lens_list = [257, 513, 129]
+    seq_lens = torch.tensor(seq_lens_list, dtype=torch.int32, device="cuda")
+    block_table, num_pages = _make_paged_layout(seq_lens_list)
+    kv_cache = (
+        torch.randn(
+            num_pages,
+            num_kv_heads,
+            BLOCK_SIZE,
+            2 * HEAD_DIM,
+            dtype=torch.bfloat16,
+            device="cuda",
+        )
+        * 0.25
+    ).to(torch.float8_e4m3fn)
+    num_tokens = len(seq_lens_list)
+    metadata = MiniMaxM3SparseMetadata(
+        seq_lens=seq_lens,
+        max_seq_len=max(seq_lens_list),
+        slot_mapping=torch.empty(num_tokens, dtype=torch.int64, device="cuda"),
+        num_actual_tokens=num_tokens,
+        num_decodes=num_tokens,
+        num_decode_tokens=num_tokens,
+        num_prefills=0,
+        num_prefill_tokens=0,
+        decode=MiniMaxM3SparseMSADecodeMetadata(
+            seq_lens=seq_lens,
+            block_table=block_table,
+            decode_query_len=1,
+            msa_cutlass=None,
+        ),
+    )
+    layer_name = "model.layers.0.self_attn.attn"
+    layer = SimpleNamespace(
+        layer_name=layer_name,
+        topk_indices_buffer=_make_topk(seq_lens_list, num_kv_heads, 1),
+        _q_scale_float=q_scale,
+    )
+    query_fp8 = torch.randn(
+        num_tokens, num_q_heads * HEAD_DIM, dtype=torch.bfloat16, device="cuda"
+    ).to(torch.float8_e4m3fn)
+    query = query_fp8.to(torch.bfloat16) * q_scale
+    common = dict(
+        num_heads=num_q_heads,
+        head_size=HEAD_DIM,
+        scale=SM_SCALE,
+        num_kv_heads=num_kv_heads,
+        kv_cache_dtype="fp8",
+        topk_blocks=TOPK,
+        sparse_block_size=BLOCK_SIZE,
+    )
+    impl = MiniMaxM3SparseMSAImpl(**common, msa_decode_backend="cutlass")
+    assert impl.use_cutlass_decode
+    forward_context = ForwardContext(
+        no_compile_layers={},
+        attn_metadata={layer_name: metadata},
+        slot_mapping={},
+    )
+    with override_forward_context(forward_context):
+        expected = MiniMaxM3SparseTritonImpl(**common).forward(
+            layer, query, kv_cache, torch.empty_like(query)
+        )
+        actual = impl.forward(layer, query_fp8, kv_cache, torch.zeros_like(query))
+    current_platform.synchronize()
+    torch.testing.assert_close(actual, expected, atol=0.02, rtol=0.02)
diff --git a/tests/kernels/attention/test_minimax_m3_norm_rope_nvfp4_insert.py b/tests/kernels/attention/test_minimax_m3_norm_rope_nvfp4_insert.py
new file mode 100644
index 000000000000..dd60ed7c9129
--- /dev/null
+++ b/tests/kernels/attention/test_minimax_m3_norm_rope_nvfp4_insert.py
@@ -0,0 +1,210 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""The fused FlyDSL norm + RoPE + NVFP4 insert against the HIP norm + RoPE
+kernel followed by the reference NVFP4 inserts: every output byte, stale page
+bytes included, and qkv left untouched."""
+
+import pytest
+import torch
+
+from vllm.platforms import current_platform
+
+if not current_platform.is_rocm():
+    pytest.skip("ROCm (gfx950) FlyDSL kernel", allow_module_level=True)
+
+from vllm.platforms.rocm import on_gfx950  # noqa: E402
+
+if not on_gfx950():
+    pytest.skip("gfx950 instructions", allow_module_level=True)
+
+from tests.kernels.attention.minimax_m3_nvfp4_ref import (  # noqa: E402
+    insert_nvfp4_index,
+    insert_nvfp4_kv,
+)
+from vllm import _custom_ops as ops  # noqa: E402
+from vllm.models.minimax_m3.amd.ops.norm_rope_nvfp4_insert.host import (  # noqa: E402
+    minimax_m3_norm_rope_nvfp4_insert,
+)
+
+D = 128
+BLK = 128
+EPS = 1e-6
+ROT = 64
+
+
+def _request_slots(starts, lens, num_pages, gen):
+    """VLLM's slot mapping: request r's tokens at positions starts[r].. of its
+    own random pages."""
+    perm = torch.randperm(num_pages, generator=gen)
+    out, p = [], 0
+    for st, n in zip(starts, lens):
+        npg = (st + n + BLK - 1) // BLK
+        pos = torch.arange(st, st + n)
+        out.append(perm[p : p + npg][pos // BLK] * BLK + pos % BLK)
+        p += npg
+    return torch.cat(out) if out else torch.empty(0, dtype=torch.int64)
+
+
+class _Batch:
+    def __init__(self, starts, lens, pad, heads, scales, seed=0):
+        g = torch.Generator().manual_seed(seed)
+        self.nq, self.nkv, self.niq = heads
+        n = self.n = sum(lens) + pad
+        row = (self.nq + 2 * self.nkv + self.niq + 1) * D
+        # magnitudes 1e-3 .. 3e3 per 16-value group, some all-zero groups
+        x = torch.randn(n, row // 16, 16, generator=g)
+        mags = torch.tensor([1e-3, 0.05, 1.0, 1.0, 30.0, 3e3])
+        x *= mags[torch.randint(0, len(mags), (n, row // 16), generator=g)][..., None]
+        x[torch.rand(n, row // 16, generator=g) < 0.02] = 0.0
+        self.qkv = x.view(n, row).to(torch.bfloat16).cuda()
+        self.w = [
+            (torch.randn(D, generator=g) * 0.5).to(torch.bfloat16).cuda()
+            for _ in range(4)
+        ]
+        self.cs = (torch.rand(4096, ROT, generator=g) * 2 - 1).to(torch.bfloat16).cuda()
+        self.pos = torch.randint(0, 4096, (n,), generator=g).cuda()
+        pages = sum((s + m + BLK - 1) // BLK for s, m in zip(starts, lens)) + 2
+        pad_slots = torch.full((pad,), -1, dtype=torch.int64)
+        self.slots = torch.cat(
+            [_request_slots(starts, lens, pages, g), pad_slots]
+        ).cuda()
+        self.islots = torch.cat(
+            [_request_slots(starts, lens, pages, g), pad_slots]
+        ).cuda()
+        # stale bytes; the layer views have a block stride past one page
+        main = 2 * self.nkv * BLK * 72
+        self.kv_full = torch.randint(0, 256, (pages, main + 256), generator=g)
+        self.kv_full = self.kv_full.to(torch.uint8).cuda()
+        self.idx_full = torch.randint(0, 256, (pages, BLK * 72 + 64), generator=g)
+        self.idx_full = self.idx_full.to(torch.uint8).cuda()
+        self.ks = torch.tensor(scales[0], dtype=torch.float32, device="cuda")
+        self.vs = torch.tensor(scales[1], dtype=torch.float32, device="cuda")
+
+    def kv(self, full):
+        return full[:, : 2 * self.nkv * BLK * 72].view(-1, 2 * self.nkv, BLK, 72)
+
+    def idx(self, full):
+        return full[:, : BLK * 72].view(-1, BLK, 72)
+
+    def reference(self):
+        qkv = self.qkv.clone()
+        q = torch.empty(self.n, self.nq * D, dtype=torch.bfloat16, device="cuda")
+        iq = torch.empty(self.n, self.niq * D, dtype=torch.float8_e4m3fn, device="cuda")
+        kv_full, idx_full = self.kv_full.clone(), self.idx_full.clone()
+        ops.fused_minimax_m3_qknorm_rope_kv_insert(
+            qkv, self.w[0], self.w[1], self.cs, self.pos, self.nq, self.nkv, ROT,
+            EPS, self.w[2], self.w[3], self.niq, q_out=q, index_q_out=iq,
+        )  # fmt: skip
+        k0, v0 = self.nq * D, (self.nq + self.nkv) * D
+        ik0 = (self.nq + 2 * self.nkv + self.niq) * D
+        insert_nvfp4_kv(
+            qkv[:, k0:v0].view(self.n, self.nkv, D),
+            qkv[:, v0 : v0 + self.nkv * D].view(self.n, self.nkv, D),
+            self.kv(kv_full),
+            self.slots,
+            self.ks,
+            self.vs,
+        )
+        insert_nvfp4_index(self.idx(idx_full), qkv[:, ik0 : ik0 + D], self.islots)
+        return q, iq, kv_full, idx_full
+
+    def run(self, qkv, pos, slots, islots, kv_full, idx_full, q, iq):
+        minimax_m3_norm_rope_nvfp4_insert(
+            qkv, self.w[0], self.w[1], self.w[2], self.w[3], self.cs, pos,
+            slots, islots, self.kv(kv_full), self.idx(idx_full), self.ks,
+            self.vs, q, iq, num_heads=self.nq, num_kv_heads=self.nkv,
+            num_index_heads=self.niq, eps=EPS,
+        )  # fmt: skip
+
+    def fused(self, iq):
+        qkv = self.qkv.clone()
+        q = torch.empty(self.n, self.nq * D, dtype=torch.bfloat16, device="cuda")
+        kv_full, idx_full = self.kv_full.clone(), self.idx_full.clone()
+        self.run(qkv, self.pos, self.slots, self.islots, kv_full, idx_full, q, iq)
+        assert torch.equal(qkv.view(torch.int16), self.qkv.view(torch.int16))
+        return q, kv_full, idx_full
+
+
+def _check(batch, iq_mode="contiguous"):
+    q_ref, iq_ref, kv_ref, idx_ref = batch.reference()
+    n, w = batch.n, batch.niq * D
+    if iq_mode == "rank":  # the context-parallel indexer's [4, n, 128] buffer
+        buf = torch.full((4, n, w), 0x7F, dtype=torch.uint8, device="cuda")
+        iq = buf[2]
+    elif iq_mode == "cols":
+        buf = torch.full((n, 3 * w), 0x7F, dtype=torch.uint8, device="cuda")
+        iq = buf[:, w : 2 * w]
+    else:
+        buf = iq = torch.full((n, w), 0x7F, dtype=torch.uint8, device="cuda")
+    q, kv, idx = batch.fused(iq.view(torch.float8_e4m3fn))
+    assert torch.equal(q.view(torch.int16), q_ref.view(torch.int16))
+    assert torch.equal(iq, iq_ref.view(torch.uint8))
+    assert torch.equal(kv, kv_ref)
+    assert torch.equal(idx, idx_ref)
+    iq.fill_(0x7F)
+    assert bool((buf == 0x7F).all()), "bytes outside index_q_out written"
+
+
+TP4, TP1 = (16, 1, 1), (64, 4, 4)
+DSPARK = ([3 + 9 * i for i in range(43)], [4] * 43)  # 172 decode rows
+
+
+@pytest.mark.parametrize("heads", [TP4, TP1])
+@pytest.mark.parametrize(
+    ("starts", "lens", "pad"),
+    [
+        ([5], [1], 0),  # one decode token
+        ([0, 13], [7, 4], 1),  # a run of 7, runs across an octet edge
+        ([125, 254, 4091], [6, 6, 6], 12),  # Q6 crossing page boundaries
+        (*DSPARK, 20),  # DSpark decode rows, graph padding
+        ([0, 3, 250, 128], [131, 61, 9, 99], 0),  # unaligned prefill chunks
+        ([0, 100], [4096, 4000], 3),  # the prefill tier
+    ],
+)
+@pytest.mark.parametrize("scales", [(1.0, 1.0), (0.002, 3.0)])
+def test_matches_two_kernel_path(heads, starts, lens, pad, scales):
+    _check(_Batch(starts, lens, pad, heads, scales))
+
+
+@pytest.mark.parametrize("iq_mode", ["rank", "cols"])
+@pytest.mark.parametrize(("starts", "lens"), [DSPARK, ([7], [600])])
+def test_index_q_out_view(iq_mode, starts, lens):
+    _check(_Batch(starts, lens, 5, TP4, (1.0, 1.0)), iq_mode)
+
+
+def test_64k_tokens():
+    _check(_Batch([0, 777], [32768, 32768], 0, TP4, (1.0, 1.0)))
+
+
+def test_all_padding():
+    """A graph-capture dummy batch: every slot -1, the pages stay as they were."""
+    batch = _Batch([], [], 64, TP4, (1.0, 1.0))
+    _check(batch)
+    iq = torch.empty(batch.n, D, dtype=torch.float8_e4m3fn, device="cuda")
+    _, kv, idx = batch.fused(iq)
+    assert torch.equal(kv, batch.kv_full) and torch.equal(idx, batch.idx_full)
+
+
+def test_cuda_graph_replay():
+    """Captured once, replayed on new qkv / positions / slots in place."""
+    a = _Batch(*DSPARK, 4, TP4, (1.0, 1.0), seed=1)
+    b = _Batch(*DSPARK, 4, TP4, (1.0, 1.0), seed=2)
+    b.w, b.cs = a.w, a.cs  # captured with the layer's weights and RoPE cache
+    qkv, pos = a.qkv.clone(), a.pos.clone()
+    slots, islots = a.slots.clone(), a.islots.clone()
+    kv_full, idx_full = b.kv_full.clone(), b.idx_full.clone()
+    q = torch.empty(a.n, a.nq * D, dtype=torch.bfloat16, device="cuda")
+    iq = torch.empty(a.n, D, dtype=torch.float8_e4m3fn, device="cuda")
+    args = (qkv, pos, slots, islots, kv_full, idx_full, q, iq)
+    a.run(*args)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        a.run(*args)
+    new = (b.qkv, b.pos, b.slots, b.islots, b.kv_full, b.idx_full)
+    for dst, src in zip(args[:6], new):
+        dst.copy_(src)
+    graph.replay()
+    q_ref, iq_ref, kv_ref, idx_ref = b.reference()
+    assert torch.equal(q.view(torch.int16), q_ref.view(torch.int16))
+    assert torch.equal(iq.view(torch.uint8), iq_ref.view(torch.uint8))
+    assert torch.equal(kv_full, kv_ref) and torch.equal(idx_full, idx_ref)
diff --git a/tests/kernels/attention/test_minimax_m3_nvfp4_index.py b/tests/kernels/attention/test_minimax_m3_nvfp4_index.py
new file mode 100644
index 000000000000..d1331087ee61
--- /dev/null
+++ b/tests/kernels/attention/test_minimax_m3_nvfp4_index.py
@@ -0,0 +1,504 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""MiniMax-M3 NVFP4 index scorers (FlyDSL) against AITER's FP8 scorers and an
+fp32 reference over the same e4m3 keys; the context-parallel push scorer
+against the single-head scorer, bit for bit, on one GPU."""
+
+import pytest
+import torch
+
+from vllm.platforms import current_platform
+
+if not current_platform.is_rocm():
+    pytest.skip("ROCm (gfx950) FlyDSL kernels", allow_module_level=True)
+
+from vllm.platforms.rocm import on_gfx950  # noqa: E402
+
+if not on_gfx950():
+    pytest.skip("gfx950 instructions", allow_module_level=True)
+
+from tests.kernels.attention.minimax_m3_nvfp4_ref import dequant  # noqa: E402
+from vllm.models.minimax_m3.amd.ops.nvfp4_index_decode import (  # noqa: E402
+    host as decode_host,
+)
+from vllm.models.minimax_m3.amd.ops.nvfp4_index_decode import (  # noqa: E402
+    kernel_cp_push as cp,
+)
+from vllm.models.minimax_m3.amd.ops.nvfp4_index_prefill import (  # noqa: E402
+    host as prefill_host,
+)
+
+minimax_m3_nvfp4_index_score_decode = decode_host.minimax_m3_nvfp4_index_score_decode
+minimax_m3_nvfp4_index_score_prefill = prefill_host.minimax_m3_nvfp4_index_score_prefill
+
+BLOCK = 128
+D = 128
+SPECIAL_LENS = [0, 1, 2, 3, 127, 128, 129, 1024]  # padding, short and edge rows
+
+
+def _index_cache(num_blocks: int, page_step: int = 1) -> torch.Tensor:
+    """Random NVFP4 index pages ``page_step`` pages apart (a view of a larger
+    cache)."""
+    n = num_blocks * page_step
+    cache = torch.randint(0, 256, (n, BLOCK * 72), device="cuda").to(torch.uint8)
+    s = torch.rand(n, BLOCK * 8, device="cuda") * 1.75 + 0.25
+    cache[:, BLOCK * 64 :] = s.to(torch.float8_e4m3fn).view(torch.uint8)
+    return cache.view(n, BLOCK, 72)[::page_step]
+
+
+def _fp8_keys(cache: torch.Tensor) -> torch.Tensor:
+    """Every page's keys as e4m3_rn_satfinite(e2m1 x scale), [nb, 128, 128]."""
+    flat = cache.reshape(cache.shape[0], -1)
+    keys = dequant(
+        flat[:, : BLOCK * 64].view(-1, BLOCK, 64),
+        flat[:, BLOCK * 64 :].view(-1, BLOCK, 8),
+    )
+    return keys.clamp(-448, 448).to(torch.float8_e4m3fn)
+
+
+def _block_table(seq_lens, share=False):
+    """Shuffled pages per request; ``share``: some requests start with another
+    request's pages (a shared prefix)."""
+    widths = [(s + BLOCK - 1) // BLOCK for s in seq_lens]
+    num_blocks = sum(widths) + 2
+    perm = torch.randperm(num_blocks, device="cuda", dtype=torch.int32)
+    table = torch.zeros(len(seq_lens), max(widths), dtype=torch.int32, device="cuda")
+    start = 0
+    for r, w in enumerate(widths):
+        table[r, :w] = perm[start : start + w]
+        start += w
+    for r in range(len(seq_lens) if share else 0):
+        src = 0 if r % 2 == 0 else 3 * (r // 6)
+        n = min(widths[r], widths[src], 10 if r % 2 == 0 else 6)
+        if r != src:
+            table[r, :n] = table[src, :n]
+    return table, num_blocks
+
+
+def _lens(num_reqs: int, hi: int) -> list[int]:
+    lens = torch.randint(1, hi, (num_reqs,)).tolist()
+    for i, v in enumerate(SPECIAL_LENS[: num_reqs // 2]):
+        lens[2 * i + 1] = v
+    return lens
+
+
+def _ref_block_scores(q_rows, keys, kv_lens, init_blocks, local_blocks):
+    """[H, R, nb] per-block max of q . k over each row's visible tokens, -inf
+    where none is visible, AITER's init (1e30) / local (1e29) sentinels."""
+    nb = keys.shape[0] // BLOCK
+    dots = torch.einsum("rhd,td->hrt", q_rows.float(), keys)
+    tok = torch.arange(nb * BLOCK, device=keys.device)
+    dots = dots.masked_fill(tok[None, None, :] >= kv_lens[None, :, None], -torch.inf)
+    mx = dots.view(dots.shape[0], dots.shape[1], nb, BLOCK).amax(-1)
+    rb = (kv_lens.clamp(min=0) + BLOCK - 1) // BLOCK  # visible blocks per row
+    blk = torch.arange(nb, device=keys.device)[None, :]
+    vis = blk < rb[:, None]
+    init = vis & (blk < init_blocks)
+    local = vis & (blk >= (rb - local_blocks).clamp(min=0)[:, None])
+    mx = torch.where(init[None], torch.full_like(mx, 1e30), mx)
+    return torch.where(local[None], torch.full_like(mx, 1e29), mx)
+
+
+def _assert_scores(got, want, msg=None):
+    """Sentinels and -inf identical; finite scores close."""
+    exact = ~torch.isfinite(want) | (want >= 1e28)
+    assert torch.equal(got[exact], want[exact]), msg
+    torch.testing.assert_close(got[~exact], want[~exact], rtol=2e-4, atol=1e-3, msg=msg)
+
+
+def _check_decode(actual, q, keys, table, lens, query_len, init_blocks, local_blocks):
+    """Each request's rows against the reference; nothing past its blocks
+    written."""
+    for r, s in enumerate(lens):
+        nb = (s + BLOCK - 1) // BLOCK
+        sl = slice(r * query_len, (r + 1) * query_len)
+        assert actual[:, sl, nb:].isnan().all(), (r, s)
+        if nb == 0:
+            continue
+        kv_lens = (s - query_len + 1 + torch.arange(query_len, device="cuda")).clamp(
+            min=0
+        )
+        page_keys = keys[table[r, :nb].long()].reshape(-1, D)
+        want = _ref_block_scores(q[sl], page_keys, kv_lens, init_blocks, local_blocks)
+        _assert_scores(actual[:, sl, :nb], want, (r, s))
+
+
+@pytest.mark.parametrize(
+    ("num_heads", "query_len"),
+    [(1, 1), (1, 4), (1, 16), (2, 4), (4, 1), (4, 4), (4, 5), (4, 8)],
+)
+@pytest.mark.parametrize("page_step", [1, 2])
+def test_index_decode_matches_aiter_fp8(num_heads, query_len, page_step):
+    """Every block a row can see against AITER's FP8 decode scorer."""
+    from aiter.ops.msa_attention import pa_sparse_block_score_decode
+
+    torch.manual_seed(0)
+    lens = [5, 300, 2100, 9000]
+    table, num_blocks = _block_table(lens)
+    seq_lens = torch.tensor(lens, dtype=torch.int32, device="cuda")
+    cache = _index_cache(num_blocks, page_step)
+    rows = len(lens) * query_len
+    q = (torch.randn(rows, num_heads, D, device="cuda") * 0.5).to(torch.float8_e4m3fn)
+    width = (max(lens) + BLOCK - 1) // BLOCK
+    args = (table, seq_lens, 1, 2, query_len, max(lens))
+    actual = torch.full((num_heads, rows, width), float("nan"), device="cuda")
+    minimax_m3_nvfp4_index_score_decode(q, cache, actual, *args)
+    expected = torch.full_like(actual, float("nan"))
+    pages = _fp8_keys(cache)
+    for h0 in range(0, num_heads, 2):  # AITER builds one- and two-head scorers
+        pa_sparse_block_score_decode(
+            q[:, h0 : h0 + 2].contiguous(), pages, expected[h0 : h0 + 2], *args
+        )
+    torch.accelerator.synchronize()
+    for r, s in enumerate(lens):
+        for t in range(query_len):
+            row = r * query_len + t
+            seen = (max(s - query_len + t + 1, 0) + BLOCK - 1) // BLOCK
+            got, want = actual[:, row, :seen], expected[:, row, :seen]
+            finite = want < 1e28
+            assert torch.equal(got[~finite], want[~finite])
+            torch.testing.assert_close(got[finite], want[finite], rtol=1e-5, atol=1e-4)
+
+
+@pytest.mark.parametrize(
+    ("num_heads", "query_len"), [(1, 1), (1, 4), (2, 3), (4, 1), (4, 4), (4, 6), (4, 8)]
+)
+@pytest.mark.parametrize("num_reqs", [1, 40, 130])
+@pytest.mark.parametrize(("init_blocks", "local_blocks"), [(1, 2), (0, 1)])
+@pytest.mark.parametrize("buffered", [False, True])
+def test_index_decode_matches_reference(
+    num_heads, query_len, num_reqs, init_blocks, local_blocks, buffered, monkeypatch
+):
+    """Request chunks, shared prefixes, padding requests, sequences shorter
+    than the query; ``buffered`` forces the LDS-buffered variant."""
+    if buffered:
+        monkeypatch.setattr(decode_host, "BUFFER_MIN_BYTES", 0)
+        monkeypatch.setattr(decode_host, "UNITS_PER_SLOT", 1)
+    torch.manual_seed(num_reqs)
+    lens = _lens(num_reqs, 6000)
+    table, num_blocks = _block_table(lens, share=True)
+    seq_lens = torch.tensor(lens, dtype=torch.int32, device="cuda")
+    cache = _index_cache(num_blocks, 2 if num_reqs == 40 else 1)
+    rows = num_reqs * query_len
+    q = (torch.randn(rows, num_heads, D, device="cuda") * 0.5).to(torch.float8_e4m3fn)
+    width = (max(lens) + BLOCK - 1) // BLOCK + 2
+    actual = torch.full((num_heads, rows, width), float("nan"), device="cuda")
+    minimax_m3_nvfp4_index_score_decode(
+        q,
+        cache,
+        actual,
+        table,
+        seq_lens,
+        init_blocks,
+        local_blocks,
+        query_len,
+        max(lens),
+    )
+    keys = _fp8_keys(cache).float()
+    torch.accelerator.synchronize()
+    _check_decode(actual, q, keys, table, lens, query_len, init_blocks, local_blocks)
+
+
+@pytest.mark.parametrize(("num_heads", "query_len"), [(1, 4), (4, 4)])
+def test_index_decode_graph_replays_new_lengths(num_heads, query_len):
+    """Captured once at a large max_seq_len, replayed with new lengths and
+    queries."""
+    torch.manual_seed(5)
+    sets = [
+        [140000, 6, 70001, 129, 7, 5000, 0, 128],
+        [6, 127, 2048, 140000, 40000, 300, 129, 0],
+    ]
+    table, num_blocks = _block_table([max(x) for x in zip(*sets)])
+    cache = _index_cache(num_blocks)
+    seq_lens = torch.tensor(sets[0], dtype=torch.int32, device="cuda")
+    rows = len(sets[0]) * query_len
+    q = torch.zeros(rows, num_heads, D, device="cuda", dtype=torch.float8_e4m3fn)
+    max_seq_len = 1 << 18
+    out = torch.empty(num_heads, rows, max_seq_len // BLOCK, device="cuda")
+    args = (q, cache, out, table, seq_lens, 1, 2, query_len, max_seq_len)
+    minimax_m3_nvfp4_index_score_decode(*args)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        minimax_m3_nvfp4_index_score_decode(*args)
+    keys = _fp8_keys(cache).float()
+    for lens in sets + sets[:1]:
+        seq_lens.copy_(torch.tensor(lens, dtype=torch.int32))
+        q.copy_((torch.randn(q.shape, device="cuda") * 0.5).to(q.dtype))
+        out.fill_(float("nan"))
+        graph.replay()
+        torch.accelerator.synchronize()
+        _check_decode(out, q, keys, table, lens, query_len, 1, 2)
+
+
+def _check_cp_push(lens, page_step, init_blocks, local_blocks, grid):
+    """The context-parallel scorer on one GPU, the other ranks' index-q rows
+    and flags pre-set locally: per rank, the single-head scorer's bits on its
+    blocks, -inf past the length inside its last run, nothing elsewhere; the
+    four ranks cover every block once."""
+    num_reqs = len(lens)
+    table, num_blocks = _block_table(lens, share=True)
+    seq_lens = torch.tensor(lens, dtype=torch.int32, device="cuda")
+    cache = _index_cache(num_blocks, page_step)
+    query_len, heads = cp.QL, cp.HEADS
+    rows = num_reqs * query_len
+    q = (torch.randn(rows, heads, D, device="cuda") * 0.5).to(torch.float8_e4m3fn)
+    max_seq_len = max(lens)
+    width = max(64, ((max_seq_len + BLOCK - 1) // BLOCK + 63) // 64 * 64)
+    expected = torch.full((heads, rows, width), float("nan"), device="cuda")
+    for h in range(heads):
+        decode_host.minimax_m3_nvfp4_index_score_decode(
+            q[:, h : h + 1], cache, expected[h : h + 1], table, seq_lens,
+            init_blocks, local_blocks, query_len, max_seq_len,
+        )  # fmt: skip
+    blocks = torch.arange(width, device="cuda")
+    nblk = ((seq_lens + BLOCK - 1) // BLOCK).repeat_interleave(query_len)[:, None]
+    seen = torch.zeros(heads, rows, width, dtype=torch.int32, device="cuda")
+    for rank in range(cp.WORLD):
+        counts = torch.zeros(cp.CNT_WORDS, dtype=torch.int32, device="cuda")
+        flags = torch.zeros(cp.FLAG_BYTES // 4, dtype=torch.int32, device="cuda")
+        flags[cp.SC_FLAGS : cp.SC_FLAGS + cp.WORLD] = 1 << 30
+        flags[cp.SC_FLAGS + rank] = 0
+        iq_buf = torch.full((heads, rows, D), 0x7F, dtype=torch.uint8, device="cuda")
+        iq_buf = iq_buf.view(torch.float8_e4m3fn)
+        for p in range(heads):
+            if p != rank:
+                iq_buf[p] = q[:, p]
+                w = cp.SC_IQ + p * cp.SC_MAX_REQS
+                flags[w : w + cp.SC_MAX_REQS] = 1 << 30
+        outs = []
+        for _ in range(2):
+            actual = torch.full_like(expected, float("nan"))
+            dests = [actual[h].data_ptr() for h in range(heads)]
+            decode_host.minimax_m3_nvfp4_index_score_push_cp(
+                q[:, rank].contiguous(), cache, dests, width, rows * width * 4,
+                table, seq_lens, init_blocks, local_blocks, rank, counts,
+                [flags.data_ptr()] * 4, [iq_buf.data_ptr()] * 4, rows, grid=grid,
+            )  # fmt: skip
+            outs.append(actual)
+        torch.accelerator.synchronize()
+        actual = outs[0]
+        assert torch.equal(actual.view(torch.int32), outs[1].view(torch.int32))
+        mine = (blocks // cp.CP_CH) % cp.WORLD == rank
+        in_runs = mine & (blocks // cp.CP_CH * cp.CP_CH < nblk)
+        inside = in_runs & (blocks < nblk)
+        assert actual[:, ~in_runs].isnan().all(), rank
+        assert (actual[:, in_runs & ~inside] == float("-inf")).all(), rank
+        assert torch.equal(
+            actual[:, inside].view(torch.int32), expected[:, inside].view(torch.int32)
+        ), rank
+        seen += inside.int()
+    assert torch.equal(seen, (~expected.isnan()).int())
+
+
+@pytest.mark.parametrize("num_reqs", [1, 5, 17, 45])
+@pytest.mark.parametrize("page_step", [1, 2])
+@pytest.mark.parametrize(("init_blocks", "local_blocks"), [(1, 2), (0, 1)])
+@pytest.mark.parametrize("grid", [3, 768])
+def test_index_decode_cp_push_matches_single_head(
+    num_reqs, page_step, init_blocks, local_blocks, grid
+):
+    """Grid 3: many runs and requests per workgroup."""
+    torch.manual_seed(num_reqs)
+    _check_cp_push(_lens(num_reqs, 9000), page_step, init_blocks, local_blocks, grid)
+
+
+@pytest.mark.parametrize("grid", [3, 768])
+def test_index_decode_cp_push_all_padding(grid):
+    """Padding requests only (graph capture): nothing written."""
+    _check_cp_push([0] * 5, 1, 1, 2, grid)
+
+
+@pytest.mark.parametrize("num_reqs", [1, 5, 17, 45])
+@pytest.mark.parametrize("grid", [3, 768])
+def test_index_decode_cp_push_split_requests(num_reqs, grid):
+    """DSpark 7 (8 rows a request) scored as two 4-row CP requests
+    (``indexer_nvfp4._cp_requests``): the four ranks' blocks of every row's
+    visible context match the single-head scorer bit for bit, and so do the
+    top-k ids."""
+    from vllm.models.minimax_m3.amd.indexer_msa import _score_width
+    from vllm.models.minimax_m3.amd.indexer_nvfp4 import _cp_requests
+    from vllm.models.minimax_m3.amd.ops.index_topk_decode.host import (
+        minimax_m3_index_topk_decode_indices,
+    )
+    from vllm.models.minimax_m3.common.indexer import MiniMaxM3IndexerDecodeMetadata
+
+    torch.manual_seed(num_reqs)
+    query_len, heads, init_blocks, local_blocks = 8, cp.HEADS, 1, 2
+    lens = [s and max(s, query_len) for s in _lens(num_reqs, 9000)]  # 0: padding
+    lens[0] = 16 * cp.CP_CH * BLOCK + 2  # the first rows' last block opens a run
+    table, num_blocks = _block_table(lens, share=True)
+    seq_lens = torch.tensor(lens, dtype=torch.int32, device="cuda")
+    cache = _index_cache(num_blocks)
+    rows = num_reqs * query_len
+    q = (torch.randn(rows, heads, D, device="cuda") * 0.5).to(torch.float8_e4m3fn)
+    max_seq_len = max(lens)
+    width = _score_width(max_seq_len, BLOCK)  # the top-k's slots
+    expected = torch.full((heads, rows, width), float("nan"), device="cuda")
+    for h in range(heads):
+        decode_host.minimax_m3_nvfp4_index_score_decode(
+            q[:, h : h + 1], cache, expected[h : h + 1], table, seq_lens,
+            init_blocks, local_blocks, query_len, max_seq_len,
+        )  # fmt: skip
+    meta = MiniMaxM3IndexerDecodeMetadata(
+        seq_lens, table, max_seq_len, query_len, query_len
+    )
+    cp_table, cp_lens = _cp_requests(meta)
+    blocks = torch.arange(width, device="cuda")
+    gathered = torch.full_like(expected, float("nan"))
+    for rank in range(cp.WORLD):
+        counts = torch.zeros(cp.CNT_WORDS, dtype=torch.int32, device="cuda")
+        flags = torch.zeros(cp.FLAG_BYTES // 4, dtype=torch.int32, device="cuda")
+        flags[cp.SC_FLAGS : cp.SC_FLAGS + cp.WORLD] = 1 << 30
+        flags[cp.SC_FLAGS + rank] = 0
+        iq_buf = torch.full((heads, rows, D), 0x7F, dtype=torch.uint8, device="cuda")
+        iq_buf = iq_buf.view(torch.float8_e4m3fn)
+        for p in range(heads):
+            if p != rank:
+                iq_buf[p] = q[:, p]
+                w = cp.SC_IQ + p * cp.SC_MAX_REQS
+                flags[w : w + cp.SC_MAX_REQS] = 1 << 30
+        actual = torch.full_like(expected, float("nan"))
+        dests = [actual[h].data_ptr() for h in range(heads)]
+        decode_host.minimax_m3_nvfp4_index_score_push_cp(
+            q[:, rank].contiguous(), cache, dests, width, rows * width * 4,
+            cp_table, cp_lens, init_blocks, local_blocks,
+            rank, counts, [flags.data_ptr()] * 4, [iq_buf.data_ptr()] * 4, rows,
+            grid=grid,
+        )  # fmt: skip
+        mine = (blocks // cp.CP_CH) % cp.WORLD == rank
+        gathered[:, :, mine] = actual[:, :, mine]
+    torch.accelerator.synchronize()
+    row = torch.arange(rows, device="cuda")
+    visible = seq_lens.repeat_interleave(query_len) - query_len + 1 + row % query_len
+    seen = blocks[None, :] < ((visible + BLOCK - 1) // BLOCK)[:, None]
+    seen = seen[None].expand(heads, -1, -1)
+    assert torch.equal(
+        gathered[seen].view(torch.int32), expected[seen].view(torch.int32)
+    )
+    ids = [torch.empty(heads, rows, 16, dtype=torch.int32, device="cuda") for _ in "ab"]
+    for score, out in zip((expected, gathered), ids):
+        minimax_m3_index_topk_decode_indices(
+            score, out, seq_lens, max_seq_len=max_seq_len, query_len=query_len
+        )
+    assert torch.equal(ids[0], ids[1])
+
+
+def _prefill_case(q_lens, seq_lens_list, num_heads, page_step=1):
+    torch.manual_seed(0)
+    table, num_blocks = _block_table(seq_lens_list)
+    seq_lens = torch.tensor(seq_lens_list, dtype=torch.int32, device="cuda")
+    cu = torch.zeros(len(q_lens) + 1, dtype=torch.int32, device="cuda")
+    cu[1:] = torch.tensor(q_lens, dtype=torch.int32, device="cuda").cumsum(0)
+    cache = _index_cache(num_blocks, page_step)
+    q = torch.randn(sum(q_lens), num_heads, D, device="cuda") * 0.5
+    return table, seq_lens, cu, cache, q.to(torch.float8_e4m3fn)
+
+
+def _causal_blocks(q_len, s, num_heads):
+    """[H, q_len, nb] blocks each row of a chunk ending its sequence sees."""
+    nb = (s + BLOCK - 1) // BLOCK
+    row_pos = s - q_len + torch.arange(q_len, device="cuda")
+    seen = torch.arange(nb, device="cuda")[None, :] <= (row_pos // BLOCK)[:, None]
+    return row_pos, seen[None].expand(num_heads, -1, -1)
+
+
+@pytest.mark.parametrize("num_heads", [1, 2, 4])
+@pytest.mark.parametrize(
+    ("q_lens", "seq_lens_list"),
+    [
+        ([129, 300, 5], [129, 2500, 2049]),  # 128-row tiles, ragged
+        ([1000, 700], [4300, 700]),  # 256-row tiles, unaligned prefix
+    ],
+)
+@pytest.mark.parametrize("page_step", [1, 2])
+def test_index_prefill_matches_aiter_fp8(num_heads, q_lens, seq_lens_list, page_step):
+    """Every block a row can see against AITER's FP8 prefill scorer."""
+    from aiter.ops.msa_attention import pa_sparse_block_score_prefill
+
+    table, seq_lens, cu, cache, q = _prefill_case(
+        q_lens, seq_lens_list, num_heads, page_step
+    )
+    msl = max(seq_lens_list)
+    args = (table, cu, seq_lens, 1, 2, max(q_lens), msl)
+    width = (msl + BLOCK - 1) // BLOCK
+    actual = torch.full((num_heads, sum(q_lens), width), float("nan"), device="cuda")
+    minimax_m3_nvfp4_index_score_prefill(q, cache, actual, *args)
+    expected = torch.full_like(actual, float("nan"))
+    pages = _fp8_keys(cache)
+    for h in range(num_heads):
+        pa_sparse_block_score_prefill(
+            q[:, h : h + 1].contiguous(), pages, expected[h : h + 1], *args
+        )
+    torch.accelerator.synchronize()
+    q_start = 0
+    for q_len, s in zip(q_lens, seq_lens_list):
+        _, seen = _causal_blocks(q_len, s, num_heads)
+        rows = slice(q_start, q_start + q_len)
+        got = actual[:, rows, : seen.shape[2]][seen]
+        want = expected[:, rows, : seen.shape[2]][seen]
+        finite = want < 1e28
+        assert torch.equal(got[~finite], want[~finite])
+        torch.testing.assert_close(got[finite], want[finite], rtol=1e-5, atol=1e-4)
+        q_start += q_len
+
+
+@pytest.mark.parametrize(
+    ("q_lens", "seq_lens_list"),
+    [
+        ([4096], [4096]),  # fresh prompt: every tile diagonal
+        ([2000, 513, 3000], [9000, 513, 20000]),  # ragged, unaligned prefixes
+        ([700, 1500], [30000, 1500]),  # long cached prefix + fresh
+    ],
+)
+@pytest.mark.parametrize(("init_blocks", "local_blocks"), [(0, 1), (2, 2)])
+@pytest.mark.parametrize("num_heads", [1, 2, 4])
+def test_index_prefill_matches_reference(
+    q_lens, seq_lens_list, init_blocks, local_blocks, num_heads, monkeypatch
+):
+    """Every kernel the host can pick (``kernel_shared.py`` with 128- and
+    256-row tiles, ``kernel.py``) bit for bit alike, and against an fp32
+    reference."""
+    table, seq_lens, cu, cache, q = _prefill_case(
+        q_lens, seq_lens_list, num_heads, page_step=2 if num_heads == 2 else 1
+    )
+    width = (max(seq_lens_list) + BLOCK - 1) // BLOCK + 3
+    args = (
+        table,
+        cu,
+        seq_lens,
+        init_blocks,
+        local_blocks,
+        max(q_lens),
+        max(seq_lens_list),
+    )
+
+    def run(shared, mq):
+        monkeypatch.setattr(prefill_host, "PIN_SHARED", shared)
+        monkeypatch.setattr(prefill_host, "PIN_MQ", mq)
+        out = torch.full((num_heads, sum(q_lens), width), float("nan"), device="cuda")
+        minimax_m3_nvfp4_index_score_prefill(q, cache, out, *args)
+        return out
+
+    variants = [(True, 1), (True, 2), (False, None)]
+    outs = [run(*v) for v in variants]
+    keys = _fp8_keys(cache).float()
+    torch.accelerator.synchronize()
+    q_start = 0
+    for req, (q_len, s) in enumerate(zip(q_lens, seq_lens_list)):
+        row_pos, seen = _causal_blocks(q_len, s, num_heads)
+        nb = seen.shape[2]
+        rows = slice(q_start, q_start + q_len)
+        for v, out in zip(variants[1:], outs[1:]):
+            got, ref = out[:, rows, :nb][seen], outs[0][:, rows, :nb][seen]
+            assert torch.equal(got, ref), (v, q_len, s)
+        page_keys = keys[table[req, :nb].long()].reshape(-1, D)
+        for r0 in range(0, q_len, 512):  # [H, rows, tokens] dots in slices
+            r1 = min(q_len, r0 + 512)
+            want = _ref_block_scores(
+                q[q_start + r0 : q_start + r1], page_keys, row_pos[r0:r1] + 1,
+                init_blocks, local_blocks,
+            )  # fmt: skip
+            got = outs[0][:, q_start + r0 : q_start + r1, :nb]
+            sl = seen[:, r0:r1]
+            _assert_scores(got[sl], want[sl], (q_len, s))
+        q_start += q_len
diff --git a/tests/kernels/attention/test_minimax_m3_nvfp4_indexer_decode_cp.py b/tests/kernels/attention/test_minimax_m3_nvfp4_indexer_decode_cp.py
new file mode 100644
index 000000000000..1e658fa72407
--- /dev/null
+++ b/tests/kernels/attention/test_minimax_m3_nvfp4_indexer_decode_cp.py
@@ -0,0 +1,130 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""NVFP4 indexer decode CP at TP4: each rank's top-k must equal an exact fp32
+top-15 over the full-context scores, across repeated runs (sentinel reset)."""
+
+import os
+
+import pytest
+import torch
+import torch.multiprocessing as mp
+
+from vllm.platforms import current_platform
+
+if not current_platform.is_cuda():
+    pytest.skip("MiniMax M3 NVFP4 indexer requires CUDA.", allow_module_level=True)
+
+PAGE = 128
+SLOTS = 8
+H = 4
+
+
+def _cache(num_pages: int, gen: torch.Generator) -> torch.Tensor:
+    kv = torch.randint(
+        0, 256, (num_pages, PAGE, 72), dtype=torch.uint8, device="cuda", generator=gen
+    )
+    # Positive normal E4M3 block scales keep every score finite.
+    s = torch.rand(num_pages, PAGE * 8, device="cuda", generator=gen) * 1.75 + 0.25
+    kv.view(num_pages, -1)[:, 8192:] = s.to(torch.float8_e4m3fn).view(torch.uint8)
+    return kv
+
+
+def _reference(scores, seq_lens, query_len, head) -> torch.Tensor:
+    num_reqs = seq_lens.shape[0]
+    out = torch.full((num_reqs * query_len, 16), -1, dtype=torch.int32)
+    for b in range(num_reqs):
+        for i in range(query_len):
+            t = b * query_len + i
+            own = (int(seq_lens[b]) - query_len + i) // PAGE
+            if own + 1 <= 16:
+                out[t, : own + 1] = torch.arange(own + 1)
+                continue
+            row = scores[b, (SLOTS - query_len + i) * H + head, :own]
+            out[t, :15] = torch.sort(-row, stable=True).indices[:15].to(torch.int32)
+            out[t, 15] = own
+    return out
+
+
+def _run(rank: int, port: int, query_len: int, seq_lens_list: list[int]) -> None:
+    from vllm.distributed import (
+        init_distributed_environment,
+        initialize_model_parallel,
+    )
+    from vllm.models.minimax_m3.nvidia.indexer_nvfp4_decode_cp import (
+        Nvfp4IndexerDecodeCP,
+    )
+
+    torch.accelerator.set_device_index(rank)
+    init_distributed_environment(
+        world_size=4,
+        rank=rank,
+        local_rank=rank,
+        distributed_init_method=f"tcp://127.0.0.1:{port}",
+    )
+    initialize_model_parallel(tensor_model_parallel_size=4)
+
+    gen = torch.Generator(device="cuda").manual_seed(0)  # same on every rank
+    seq_lens = torch.tensor(seq_lens_list, dtype=torch.int32, device="cuda")
+    num_reqs = seq_lens.shape[0]
+    pages = (int(seq_lens.max()) + PAGE - 1) // PAGE
+    kv = _cache(num_reqs * pages + 8, gen)
+    # Sized to the longest context, as vLLM's block table is: the longest
+    # request then owns the last page of its rank's compact table.
+    block_table = (
+        torch.randperm(num_reqs * pages, device="cuda", generator=gen)
+        .to(torch.int32)
+        .view(num_reqs, pages)
+    )
+    rows = num_reqs * query_len
+
+    cp = Nvfp4IndexerDecodeCP(max_reqs=num_reqs, device=torch.device("cuda", rank))
+    layers = [cp.new_layer() for _ in range(2)]
+    plan = cp.plan(block_table, seq_lens)
+
+    from vllm.third_party.fmha_sm100.sparse import BatchDecodeIndexerQ8KV4Wrapper
+
+    full = BatchDecodeIndexerQ8KV4Wrapper(num_heads=H)
+    full.plan(block_table, seq_lens)
+    for step in range(3):
+        for layer in layers:
+            q = (torch.randn(rows, H, 128, device="cuda", generator=gen) * 0.5).to(
+                torch.float8_e4m3fn
+            )
+            out = torch.full((rows, 1, 16), -7, dtype=torch.int32, device="cuda")
+            cp.forward(
+                layer, plan, q[:, rank : rank + 1].contiguous(), kv, out, query_len
+            )
+            q_tile = torch.zeros(
+                num_reqs, SLOTS, H, 128, dtype=torch.float8_e4m3fn, device="cuda"
+            )
+            q_tile[:, SLOTS - query_len :] = q.view(num_reqs, query_len, H, 128)
+            # NaN in the columns the kernel leaves alone, so a selection
+            # that reads one cannot match the reference.
+            full._scores.fill_(float("nan"))
+            scores = full._run_scores(q_tile.view(-1, H, 128), kv)
+            expected = _reference(scores.cpu(), seq_lens.cpu(), query_len, rank)
+            torch.testing.assert_close(out[:, 0].cpu(), expected, rtol=0, atol=0)
+
+
+def _worker(rank: int, port: int, query_len: int, seq_lens_list: list[int]) -> None:
+    from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config
+
+    config = VllmConfig(parallel_config=ParallelConfig(tensor_parallel_size=4))
+    with set_current_vllm_config(config):
+        _run(rank, port, query_len, seq_lens_list)
+
+
+@pytest.mark.skipif(
+    torch.accelerator.device_count() < 4
+    or not (
+        current_platform.is_device_capability(100)
+        or current_platform.is_device_capability(103)
+    ),
+    reason="Needs 4 SM100/SM103 GPUs.",
+)
+@pytest.mark.parametrize("query_len", [1, 4, 8])
+def test_nvfp4_indexer_decode_cp_matches_exact_topk(query_len: int):
+    # Trivial rows (<= 16 pages), page boundaries, and long contexts.
+    seq_lens = [9, 2049, 2048 + query_len, 20000, 131072, 300001]
+    port = 29500 + (os.getpid() % 1000)
+    mp.spawn(_worker, args=(port, query_len, seq_lens), nprocs=4, join=True)
diff --git a/tests/kernels/attention/test_minimax_m3_nvfp4_prefill_flydsl.py b/tests/kernels/attention/test_minimax_m3_nvfp4_prefill_flydsl.py
new file mode 100644
index 000000000000..8e3b368dba38
--- /dev/null
+++ b/tests/kernels/attention/test_minimax_m3_nvfp4_prefill_flydsl.py
@@ -0,0 +1,260 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""MiniMax-M3 NVFP4 sparse prefill (FlyDSL) against fp64 attention over the
+exactly dequantized K / V, and bitwise stable across runs and call splits."""
+
+import pytest
+import torch
+
+from vllm.platforms import current_platform
+
+if not current_platform.is_rocm():
+    pytest.skip("ROCm (gfx950) FlyDSL kernel", allow_module_level=True)
+
+from vllm.platforms.rocm import on_gfx950  # noqa: E402
+
+if not on_gfx950():
+    pytest.skip("gfx950 instructions", allow_module_level=True)
+
+from tests.kernels.attention.minimax_m3_nvfp4_ref import (  # noqa: E402
+    dequant,
+    insert_nvfp4_kv,
+    nvfp4_kv_cache_token_major,
+)
+from tests.kernels.mfma_spinner import mfma_spinner  # noqa: E402
+from vllm.models.minimax_m3.amd.ops.nvfp4_prefill.host import (  # noqa: E402
+    minimax_m3_nvfp4_sparse_prefill,
+    supports_nvfp4_sparse_prefill,
+)
+
+BLOCK = 128
+D = 128
+GROUP = 16  # query heads per KV head
+TOPK = 16
+SM_SCALE = D**-0.5
+
+
+def _case(q_lens, ctx_lens, hkv, k_scale, v_scale, seed=0):
+    """The last q_lens[i] of ctx_lens[i] tokens are queries; shuffled pages;
+    each token picks its own block and 15 random earlier ones."""
+    gen = torch.Generator().manual_seed(seed)
+    nbs = [(c + BLOCK - 1) // BLOCK for c in ctx_lens]
+    n_pages = sum(nbs) + 2
+    cache = torch.zeros(n_pages, 2 * hkv, BLOCK, 72, dtype=torch.uint8, device="cuda")
+    slots = torch.arange(n_pages * BLOCK, device="cuda")
+    k = (torch.randn(slots.numel(), hkv, D, generator=gen) * 2).bfloat16().cuda()
+    v = (torch.randn(slots.numel(), hkv, D, generator=gen) * 2).bfloat16().cuda()
+    insert_nvfp4_kv(k, v, cache, slots, k_scale, v_scale)
+    perm = torch.randperm(n_pages, generator=gen).int()
+    table = torch.zeros(len(q_lens), max(nbs), dtype=torch.int32)
+    off = 0
+    for r, nb in enumerate(nbs):
+        table[r, :nb] = perm[off : off + nb]
+        off += nb
+    pos = torch.cat([torch.arange(c - q, c) for q, c in zip(q_lens, ctx_lens)]).int()
+    req = torch.cat([torch.full((q,), r) for r, q in enumerate(q_lens)]).int()
+    rows = []
+    for _ in range(hkv):
+        for p in pos.tolist():
+            cur = p // BLOCK
+            sel = torch.cat(
+                [torch.tensor([cur]), torch.randperm(cur, generator=gen)[: TOPK - 1]]
+            )
+            rows.append(torch.cat([sel, torch.full((TOPK - sel.numel(),), -1)]))
+    topk = torch.stack(rows).int().view(hkv, pos.numel(), TOPK)
+    q = (torch.randn(pos.numel(), hkv * GROUP, D, generator=gen) * 2).bfloat16()
+    return cache, table.cuda(), topk.cuda(), q.cuda(), pos.cuda(), req.cuda()
+
+
+def _prefill(q, cache, topk, table, req, pos, hkv, k_scale=None, v_scale=None):
+    out = torch.empty_like(q)
+    minimax_m3_nvfp4_sparse_prefill(
+        q, cache, topk, table, req, pos, hkv, SM_SCALE, out, k_scale=k_scale,
+        v_scale=v_scale,
+    )  # fmt: skip
+    return out
+
+
+def _fp64(q, cache, topk, table, req, pos, hkv, k_scale, v_scale, rows):
+    """fp64 attention of the token rows ``rows`` over the dequantized K / V."""
+    k_data, k_sf, v_data, v_sf = nvfp4_kv_cache_token_major(cache)
+    kv = [
+        (dequant(data, sf) * s).to(torch.bfloat16).double()  # [pages, hkv, 128, 128]
+        for data, sf, s in ((k_data, k_sf, k_scale), (v_data, v_sf, v_scale))
+    ]
+    ref = torch.empty(len(rows), hkv * GROUP, D, dtype=torch.float64, device="cuda")
+    for i, r in enumerate(rows):
+        for h in range(hkv):
+            blocks = [b for b in topk[h, r].tolist() if b >= 0]
+            pages = [int(table[int(req[r]), b]) for b in blocks]
+            kk = torch.cat([kv[0][p, h] for p in pages])
+            vv = torch.cat([kv[1][p, h] for p in pages])
+            key_pos = torch.cat(
+                [torch.arange(b * BLOCK, (b + 1) * BLOCK) for b in blocks]
+            ).cuda()
+            s = (q[r, h * GROUP : (h + 1) * GROUP].double() @ kk.T) * SM_SCALE
+            s[:, key_pos > int(pos[r])] = float("-inf")
+            ref[i, h * GROUP : (h + 1) * GROUP] = torch.softmax(s, -1) @ vv
+    return ref
+
+
+def _rel(out, ref, rows):
+    o = out[rows].double()
+    return float((o - ref).norm() / ref.norm())
+
+
+MAX_REL = 3e-3  # fp16 P and the bf16 output: ~2e-3
+
+
+@pytest.mark.parametrize("scales", [(1.0, 1.0), (0.5, 2.0)])
+def test_matches_fp64(scales):
+    k_scale, v_scale = (torch.tensor(s, device="cuda") for s in scales)
+    cache, table, topk, q, pos, req = _case([2048], [2048], 1, k_scale, v_scale)
+    out = _prefill(q, cache, topk, table, req, pos, 1, k_scale, v_scale)
+    assert torch.isfinite(out).all()
+    rows = list(range(0, 2048, 37)) + list(range(2040, 2048))
+    ref = _fp64(q, cache, topk, table, req, pos, 1, k_scale, v_scale, rows)
+    assert _rel(out, ref, rows) < MAX_REL
+
+
+@pytest.mark.parametrize("hkv", [1, 2, 4])
+def test_batch_with_prefix_and_tiles_across_requests(hkv):
+    one = torch.tensor(1.0, device="cuda")
+    q_lens, ctx_lens = [1027, 514, 3, 262], [3000, 514, 1100, 900]
+    cache, table, topk, q, pos, req = _case(q_lens, ctx_lens, hkv, one, one, seed=1)
+    out = _prefill(q, cache, topk, table, req, pos, hkv, one, one)
+    assert torch.isfinite(out).all()
+    rows = [0, 1, 2, 3, 1026, 1027, 1028, 1540, 1541, 1542, 1543, 1544, 1805]
+    ref = _fp64(q, cache, topk, table, req, pos, hkv, one, one, rows)
+    assert _rel(out, ref, rows) < MAX_REL
+
+
+@pytest.mark.parametrize("n_tok", [1, 3, 5])
+def test_few_tokens(n_tok):
+    one = torch.tensor(1.0, device="cuda")
+    cache, table, topk, q, pos, req = _case([n_tok], [700], 1, one, one, seed=2)
+    out = _prefill(q, cache, topk, table, req, pos, 1)  # unit scales by default
+    rows = list(range(n_tok))
+    ref = _fp64(q, cache, topk, table, req, pos, 1, one, one, rows)
+    assert _rel(out, ref, rows) < MAX_REL
+
+
+@pytest.mark.parametrize("hkv", [1, 4])
+def test_block_outer_pages_bitwise(hkv):
+    """A page inside a wider block (stride(0) past one page): the dense
+    cache's output."""
+    one = torch.tensor(1.0, device="cuda")
+    q_lens, ctx_lens = [1027, 3, 262], [3000, 1100, 900]
+    cache, table, topk, q, pos, req = _case(q_lens, ctx_lens, hkv, one, one, seed=3)
+    page = cache[0].numel()
+    wide = torch.full(
+        (cache.shape[0], 3 * page), 0x7F, dtype=torch.uint8, device="cuda"
+    )
+    wide = wide[:, page : 2 * page].view(cache.shape)
+    wide.copy_(cache)
+    out = _prefill(q, cache, topk, table, req, pos, hkv)
+    assert torch.isfinite(out).all()
+    assert torch.equal(out, _prefill(q, wide, topk, table, req, pos, hkv))
+
+
+def _shared_picks(pos, req, hkv, seed):
+    """Indexer-like top-k: the own block, then a per-request base of 15
+    earlier blocks, each swapped for a random one with probability 0.1."""
+    gen = torch.Generator().manual_seed(seed)
+    first_blk: dict[int, int] = {}
+    for p, r in zip(pos.tolist(), req.tolist()):
+        first_blk.setdefault(r, p // BLOCK)
+    base = {
+        r: torch.randperm(max(b, TOPK), generator=gen)[: TOPK - 1].tolist()
+        for r, b in first_blk.items()
+    }
+    rows = []
+    for _ in range(hkv):
+        for p, r in zip(pos.tolist(), req.tolist()):
+            cur = p // BLOCK
+            sel = [cur]
+            for b in base[r]:
+                if torch.rand(1, generator=gen).item() < 0.1:
+                    b = int(torch.randint(0, max(cur, 1), (1,), generator=gen))
+                if b < cur and b not in sel:
+                    sel.append(b)
+            rows.append(sel + [-1] * (TOPK - len(sel)))
+    return torch.tensor(rows, dtype=torch.int32).view(hkv, pos.numel(), TOPK).cuda()
+
+
+# Grids above / below the CU count: the two- and one-round combine.
+SHARED_CASES = [
+    (
+        [1501, 6, 2050],
+        [9000, 2100, 6200],
+        [0, 1, 2, 3, 1499, 1500, 1501, 1502, 1506, 3556],
+    ),  # fmt: skip
+    (
+        [201, 6, 150],
+        [5000, 2100, 3000],
+        [0, 1, 2, 3, 199, 200, 201, 202, 206, 207, 356],
+    ),  # fmt: skip
+]
+
+
+@pytest.mark.parametrize("hkv", [1, 2])
+@pytest.mark.parametrize("case", SHARED_CASES)
+def test_shared_picks_deterministic(hkv, case):
+    """Pages shared within a tile, tiles across requests: matches fp64,
+    bitwise repeatable."""
+    one = torch.tensor(1.0, device="cuda")
+    q_lens, ctx_lens, rows = case
+    cache, table, _, q, pos, req = _case(q_lens, ctx_lens, hkv, one, one, seed=3)
+    topk = _shared_picks(pos, req, hkv, seed=4)
+    out = _prefill(q, cache, topk, table, req, pos, hkv)
+    ref = _fp64(q, cache, topk, table, req, pos, hkv, one, one, rows)
+    assert _rel(out, ref, rows) < MAX_REL
+    assert torch.equal(out, _prefill(q, cache, topk, table, req, pos, hkv))
+
+
+def test_fresh_prompt_calls_bitwise():
+    """A fresh 4096-token prompt in one call equals the same rows in calls of
+    1024 / 2048 / 1024 tokens, bitwise (DBO microbatches)."""
+    one = torch.tensor(1.0, device="cuda")
+    n = 4096
+    cache, table, _, q, pos, req = _case([n], [n], 1, one, one, seed=5)
+    topk = _shared_picks(pos, req, 1, seed=6)
+    assert (topk[0, : 15 * BLOCK] == -1).any(1).all()
+
+    def run(a, b):
+        return _prefill(q[a:b], cache, topk[:, a:b], table, req[a:b], pos[a:b], 1)
+
+    full = run(0, n)
+    assert torch.equal(full, torch.cat([run(0, 1024), run(1024, 3072), run(3072, n)]))
+
+
+@pytest.mark.parametrize("n_tok, ctx", [(1024, 31024), (4096, 64096)])
+def test_bitwise_next_to_mfma_waves(n_tok, ctx):
+    """Another kernel's MFMA waves on every CU leave the output unchanged
+    (v_pk_mul_f32 gives wrong products there on MI355X)."""
+    one = torch.tensor(1.0, device="cuda")
+    cache, table, topk, q, pos, req = _case([n_tok], [ctx], 1, one, one, seed=7)
+
+    def run(stream):
+        with torch.cuda.stream(stream):
+            return _prefill(q, cache, topk, table, req, pos, 1)
+
+    # A high-priority queue, so the spinner never waits behind the attention.
+    main, side = torch.cuda.Stream(), torch.cuda.Stream(priority=-1)
+    ref = run(main)
+    torch.accelerator.synchronize()
+    spin = mfma_spinner(torch.cuda.get_device_properties(0).multi_processor_count)
+    for _ in range(8):
+        spin(200_000, side)  # 2 ms
+        out = run(main)
+        torch.accelerator.synchronize()
+        assert torch.equal(out, ref)
+
+
+def test_supported_shapes():
+    """16 query heads per KV head, 16 selected 128-token blocks only."""
+    assert supports_nvfp4_sparse_prefill(16, 1, D, TOPK, BLOCK)
+    assert supports_nvfp4_sparse_prefill(32, 2, D, TOPK, BLOCK)
+    assert not supports_nvfp4_sparse_prefill(8, 1, D, TOPK, BLOCK)
+    assert not supports_nvfp4_sparse_prefill(16, 1, D, 8, BLOCK)
+    assert not supports_nvfp4_sparse_prefill(16, 1, 64, TOPK, BLOCK)
diff --git a/tests/kernels/attention/test_minimax_m3_nvfp4_sparse_decode.py b/tests/kernels/attention/test_minimax_m3_nvfp4_sparse_decode.py
new file mode 100644
index 000000000000..8b368f36f3f8
--- /dev/null
+++ b/tests/kernels/attention/test_minimax_m3_nvfp4_sparse_decode.py
@@ -0,0 +1,243 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""MiniMax-M3 NVFP4 sparse decode (FlyDSL) against fp64 attention over
+K / V = e4m3_rn_satfinite(e2m1 x block scale) x global scale."""
+
+import pytest
+import torch
+
+from vllm.platforms import current_platform
+
+if not current_platform.is_rocm():
+    pytest.skip("ROCm (gfx950) FlyDSL kernel", allow_module_level=True)
+
+from vllm.platforms.rocm import on_gfx950  # noqa: E402
+
+if not on_gfx950():
+    pytest.skip("gfx950 instructions", allow_module_level=True)
+
+from tests.kernels.attention.minimax_m3_nvfp4_ref import (  # noqa: E402
+    dequant,
+    head_offsets,
+    insert_nvfp4_kv,
+    nvfp4_kv_cache_token_major,
+)
+from vllm.models.minimax_m3.amd.ops.nvfp4_sparse_decode.host import (  # noqa: E402
+    minimax_m3_nvfp4_sparse_decode_fly,
+)
+
+BLOCK = 128
+D = 128
+GROUP = 16  # query heads per KV head
+TOPK = 16
+SM_SCALE = D**-0.5
+
+
+def _nvfp4_cache(num_blocks, scales, num_kv_heads=1):
+    """Every slot of the cache filled with random K / V."""
+    k_scale, v_scale = scales
+    cache = torch.zeros(
+        num_blocks, 2 * num_kv_heads, BLOCK, 72, dtype=torch.uint8, device="cuda"
+    )
+    slots = torch.arange(num_blocks * BLOCK, device="cuda")
+    shape = (slots.numel(), num_kv_heads, D)
+    k = (torch.randn(shape, device="cuda") * 2).bfloat16()
+    v = (torch.randn(shape, device="cuda") * 2).bfloat16()
+    insert_nvfp4_kv(k, v, cache, slots, k_scale, v_scale)
+    return cache
+
+
+def _block_table(seq_lens):
+    widths = [(s + BLOCK - 1) // BLOCK for s in seq_lens]
+    num_blocks = sum(widths) + 3
+    perm = torch.randperm(num_blocks, device="cuda", dtype=torch.int32)
+    table = torch.zeros(len(seq_lens), max(widths), dtype=torch.int32, device="cuda")
+    start = 0
+    for r, w in enumerate(widths):
+        table[r, :w] = perm[start : start + w]
+        start += w
+    return table, num_blocks
+
+
+def _scales():
+    return (
+        torch.tensor(0.5, dtype=torch.float32, device="cuda"),
+        torch.tensor(2.0, dtype=torch.float32, device="cuda"),
+    )
+
+
+def _selections(seq_lens, qlen, num_kv_heads=1, seed=0):
+    """[Hkv, rows, 16] picks, mostly shared within a request, with -1 padding
+    and an out-of-range entry; head h's picks rotated by h."""
+    g = torch.Generator(device="cuda").manual_seed(seed)
+    topk = torch.full(
+        (1, len(seq_lens) * qlen, TOPK), -1, dtype=torch.int32, device="cuda"
+    )
+    for r, s in enumerate(seq_lens):
+        nb = (s + BLOCK - 1) // BLOCK
+        base = torch.randperm(nb, device="cuda", generator=g).to(torch.int32)
+        for q in range(qlen):
+            cur = max(s - qlen + q, 0) // BLOCK
+            older = base[(base != cur) & (base <= cur)]
+            pick = torch.cat([older[:10], older[10 + 3 * q : 13 + 3 * q]])
+            sel = torch.cat([torch.tensor([cur], device="cuda"), pick])
+            sel = torch.cat([sel, torch.tensor([nb + 5], device="cuda")])[:TOPK]
+            topk[0, r * qlen + q, : sel.numel()] = sel.to(torch.int32)
+    return torch.cat([topk.roll(h, -1) for h in range(num_kv_heads)])
+
+
+def _reference(q, cache, topk, table, seq_lens, scales, qlen):
+    """fp64 attention over the dequantized K / V; rows that see nothing give 0."""
+    k_data, k_sf, v_data, v_sf = nvfp4_kv_cache_token_major(cache)
+    k_all, v_all = (
+        dequant(data, sf).clamp(-448, 448).to(torch.float8_e4m3fn).double() * s.item()
+        for data, sf, s in ((k_data, k_sf, scales[0]), (v_data, v_sf, scales[1]))
+    )
+    hkv, rows, _ = topk.shape
+    req = torch.arange(rows, device="cuda") // qlen
+    pos = seq_lens.long()[req] - qlen + torch.arange(rows, device="cuda") % qlen
+    offs = torch.arange(BLOCK, device="cuda")
+    out = torch.zeros(q.shape, dtype=torch.float64, device="cuda")
+    for h in range(hkv):
+        blk = topk[h, :rows].long()
+        ok = (pos[:, None] >= 0) & (blk >= 0) & (blk <= (pos // BLOCK)[:, None])
+        blk = torch.where(ok, blk, 0)
+        tok = blk[..., None] * BLOCK + offs  # [rows, TOPK, BLOCK]
+        ok = (ok[..., None] & (tok <= pos[:, None, None])).flatten(1)
+        page = table[req[:, None], blk].long()[..., None].expand_as(tok)
+        k = k_all[page, h, tok % BLOCK].flatten(1, 2)
+        v = v_all[page, h, tok % BLOCK].flatten(1, 2).masked_fill(~ok[..., None], 0)
+        heads = slice(h * GROUP, (h + 1) * GROUP)
+        s = torch.einsum("nhd,ntd->nht", q[:, heads].double(), k) * SM_SCALE
+        s = s.masked_fill(~ok[:, None, :], float("-inf"))
+        p = torch.softmax(s, -1).nan_to_num(0.0)
+        out[:, heads] = torch.einsum("nht,ntd->nhd", p, v)
+    return out
+
+
+def _rel(a, b):
+    a, b = a.double(), b.double()
+    return float((a - b).norm() / b.norm())
+
+
+def _decode(q, cache, topk, table, seq_lens, qlen, out=None, **kw):
+    k_scale, v_scale = _scales()
+    out = torch.full_like(q, float("nan")) if out is None else out
+    minimax_m3_nvfp4_sparse_decode_fly(
+        q, cache, topk, table, seq_lens, SM_SCALE, out, k_scale=k_scale,
+        v_scale=v_scale, decode_query_len=qlen, **kw,
+    )  # fmt: skip
+    return out
+
+
+@pytest.mark.parametrize("decode_query_len", [1, 4, 8])
+@pytest.mark.parametrize("splits", [None, 1, 2, 4, 8])
+@pytest.mark.parametrize("random_bytes", [False, True])
+@pytest.mark.parametrize("xcd_count", [0, 8])
+def test_decode_matches_reference(decode_query_len, splits, random_bytes, xcd_count):
+    """Realistic and random-byte caches: close to fp64, bitwise repeatable."""
+    torch.manual_seed(0)
+    lens = [9, 300, 2100, 20000, 131, 4096]
+    table, num_blocks = _block_table(lens)
+    seq_lens = torch.tensor(lens, dtype=torch.int32, device="cuda")
+    topk = _selections(lens, decode_query_len)
+    rows = topk.shape[1]
+    if random_bytes:
+        cache = torch.randint(
+            0, 256, (num_blocks, 2, BLOCK, 72), dtype=torch.uint8, device="cuda"
+        )
+        for side in range(2):
+            sf = cache.view(num_blocks, -1)[:, side * 9216 + 8192 : (side + 1) * 9216]
+            sf.masked_fill_((sf & 0x7F) == 0x7F, 0x38)  # no NaN scales
+        q = (torch.randn(rows, GROUP, D, device="cuda") * 0.01).bfloat16()
+    else:
+        cache = _nvfp4_cache(num_blocks, _scales())
+        q = torch.randn(rows, GROUP, D, device="cuda").bfloat16()
+    kw = dict(splits=splits, xcd_count=xcd_count)
+    out = _decode(q, cache, topk, table, seq_lens, decode_query_len, **kw)
+    again = _decode(q, cache, topk, table, seq_lens, decode_query_len, **kw)
+    ref = _reference(q, cache, topk, table, seq_lens, _scales(), decode_query_len)
+    assert torch.isfinite(out).all()
+    assert torch.equal(out, again)
+    assert _rel(out, ref) < 5e-3
+
+
+@pytest.mark.parametrize("num_kv_heads", [1, 4])
+@pytest.mark.parametrize("decode_query_len", [1, 4, 6])
+@pytest.mark.parametrize("splits", [1, 4])
+def test_decode_padding_rows_and_kv_heads(num_kv_heads, decode_query_len, splits):
+    """Padding requests (seq_len 0) give zero rows; nothing past the rows is
+    written."""
+    torch.manual_seed(1)
+    lens = [700, 0, 3000, 0, 0, 129]
+    table, num_blocks = _block_table([max(s, 1) for s in lens])
+    seq_lens = torch.tensor(lens, dtype=torch.int32, device="cuda")
+    topk = _selections([max(s, 1) for s in lens], decode_query_len, num_kv_heads)
+    cache = _nvfp4_cache(num_blocks, _scales(), num_kv_heads)
+    rows = topk.shape[1]
+    q = torch.randn(rows, GROUP * num_kv_heads, D, device="cuda").bfloat16()
+    out = torch.full((rows + 3, *q.shape[1:]), float("nan"), device="cuda").bfloat16()
+    _decode(
+        q, cache, topk, table, seq_lens, decode_query_len, out[:rows], splits=splits
+    )
+    ref = _reference(q, cache, topk, table, seq_lens, _scales(), decode_query_len)
+    real = (seq_lens.repeat_interleave(decode_query_len) > 0).cpu()
+    assert _rel(out[:rows][real], ref[real]) < 5e-3
+    assert (out[:rows][~real] == 0).all()
+    assert out[rows:].isnan().all()
+    zeros = torch.zeros_like(seq_lens)
+    pad = _decode(q, cache, topk, table, zeros, decode_query_len, splits=splits)
+    assert (pad == 0).all()
+
+
+@pytest.mark.parametrize(
+    "num_kv_heads,qlen,splits,xcd_count",
+    [(1, 1, 1, 0), (1, 4, 4, 8), (4, 6, 1, 8), (4, 8, 4, 0)],
+)
+def test_garbage_past_the_length_is_ignored(num_kv_heads, qlen, splits, xcd_count):
+    """NaN scale bytes past each request's length leave the output unchanged."""
+    torch.manual_seed(42)
+    lens = [131, 4097, 2100, 777, 65]
+    table, blocks = _block_table(lens)
+    seq_lens = torch.tensor(lens, dtype=torch.int32, device="cuda")
+    topk = _selections(lens, qlen, num_kv_heads)
+    cache = _nvfp4_cache(blocks, _scales(), num_kv_heads)
+    q = torch.randn(len(lens) * qlen, num_kv_heads * GROUP, D, device="cuda").bfloat16()
+    kw = dict(splits=splits, xcd_count=xcd_count)
+    out = _decode(q, cache, topk, table, seq_lens, qlen, **kw)
+    dirty = cache.clone().view(blocks, -1)
+    for side in range(2):
+        start = side * num_kv_heads * 9216 + num_kv_heads * 8192
+        region = dirty[:, start : start + num_kv_heads * 1024]
+        region = region.view(blocks, num_kv_heads, 1024)
+        _, _, offsets = head_offsets(side, cache.device)
+        for req, length in enumerate(lens):
+            page = table[req, length // BLOCK]
+            region[page, :, offsets[length % BLOCK :].reshape(-1)] = 0x7F
+    dirty_out = _decode(q, dirty.view_as(cache), topk, table, seq_lens, qlen, **kw)
+    assert torch.isfinite(out).all()
+    assert torch.equal(out, dirty_out)
+
+
+@pytest.mark.parametrize("num_kv_heads", [1, 4])
+@pytest.mark.parametrize("decode_query_len", [1, 4])
+def test_decode_block_outer_pages_bitwise(num_kv_heads, decode_query_len):
+    """A page inside a wider block (stride(0) past one page): the dense
+    cache's output."""
+    torch.manual_seed(2)
+    lens = [300, 2100, 20000, 131]
+    table, num_blocks = _block_table(lens)
+    seq_lens = torch.tensor(lens, dtype=torch.int32, device="cuda")
+    topk = _selections(lens, decode_query_len, num_kv_heads)
+    cache = _nvfp4_cache(num_blocks, _scales(), num_kv_heads)
+    page = cache[0].numel()
+    wide = torch.full((num_blocks, 3 * page), 0x7F, dtype=torch.uint8, device="cuda")
+    wide_cache = wide[:, page : 2 * page].view(cache.shape)
+    wide_cache.copy_(cache)
+    q = torch.randn(topk.shape[1], GROUP * num_kv_heads, D, device="cuda").bfloat16()
+    outs = [
+        _decode(q, c, topk, table, seq_lens, decode_query_len)
+        for c in (cache, wide_cache)
+    ]
+    assert torch.isfinite(outs[0]).all()
+    assert torch.equal(outs[0], outs[1])
diff --git a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py
index 66414f723ee2..0c37a127c5a1 100644
--- a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py
+++ b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py
@@ -915,6 +915,7 @@ def fake_dense_fallback(*args, out, **kwargs):
 
     monkeypatch.setattr(mod, "_can_use_aiter_sparse_prefill_opus", lambda *args: True)
     monkeypatch.setattr(mod, "_get_aiter_sparse_prefill_opus", lambda: None)
+    monkeypatch.setattr(mod, "_can_use_aiter_pa_prefill_sparse", lambda *args: False)
     monkeypatch.setattr(mod, "_rocm_sparse_attn_prefill_triton", fake_dense_fallback)
     monkeypatch.setattr(
         mod,
@@ -939,6 +940,112 @@ def fake_dense_fallback(*args, out, **kwargs):
     assert torch.count_nonzero(output) == 0
 
 
+@pytest.mark.parametrize(
+    ("num_queries", "cdna_version", "mla_enabled", "dtype", "kv_dtype", "expected"),
+    [
+        (512, 3, True, torch.bfloat16, torch.bfloat16, True),
+        (512, 3, True, torch.float16, torch.float16, True),
+        (511, 3, True, torch.bfloat16, torch.bfloat16, False),
+        (512, 4, True, torch.bfloat16, torch.bfloat16, False),
+        (512, 0, True, torch.bfloat16, torch.bfloat16, False),
+        (512, 3, False, torch.bfloat16, torch.bfloat16, False),
+        (512, 3, True, torch.float32, torch.float32, False),
+        (512, 3, True, torch.bfloat16, torch.float16, False),
+    ],
+)
+def test_can_use_aiter_pa_prefill_sparse(
+    monkeypatch,
+    num_queries: int,
+    cdna_version: int,
+    mla_enabled: bool,
+    dtype: torch.dtype,
+    kv_dtype: torch.dtype,
+    expected: bool,
+) -> None:
+    """Literal gate. Expected values are not derived from the cutoff constant."""
+    from vllm._aiter_ops import rocm_aiter_ops
+    from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod
+
+    monkeypatch.setattr(mod, "get_cdna_version", lambda: cdna_version)
+    monkeypatch.setattr(rocm_aiter_ops, "is_mla_enabled", lambda: mla_enabled)
+    q = torch.empty(num_queries, 1, HEAD_DIM, dtype=dtype)
+    kv = torch.empty(4, 1, HEAD_DIM, dtype=kv_dtype)
+    assert mod._can_use_aiter_pa_prefill_sparse(q, kv) is expected
+
+
+@pytest.mark.parametrize(
+    ("num_queries", "mla_enabled", "expect_aiter"),
+    [(512, True, True), (511, True, False), (512, False, False)],
+)
+def test_sparse_attn_prefill_routes_gfx942_pa_prefill_sparse(
+    monkeypatch, num_queries: int, mla_enabled: bool, expect_aiter: bool
+) -> None:
+    """Call site launches only when the gate passes."""
+    from unittest.mock import Mock
+
+    from vllm._aiter_ops import rocm_aiter_ops
+    from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod
+
+    kernel = None
+    if expect_aiter:
+        pytest.importorskip("aiter.ops.triton.attention.pa_prefill_sparse")
+        from aiter.ops.triton.attention.pa_prefill_sparse import pa_prefill_sparse
+
+        kernel = pa_prefill_sparse
+
+    monkeypatch.setattr(rocm_aiter_ops, "is_mla_enabled", lambda: mla_enabled)
+    monkeypatch.setattr(mod, "get_cdna_version", lambda: 3)
+    monkeypatch.setattr(mod, "_can_use_aiter_sparse_prefill_opus", lambda *args: False)
+
+    def zero_output(*args, **kwargs):
+        kwargs["output"].zero_()
+
+    def zero_out(*args, out, **kwargs):
+        out.zero_()
+        return out
+
+    fake_aiter = Mock(side_effect=zero_output)
+    fake_triton = Mock(side_effect=zero_out)
+    monkeypatch.setattr(mod, "_rocm_sparse_attn_prefill_aiter", fake_aiter)
+    monkeypatch.setattr(mod, "_rocm_sparse_attn_prefill_triton", fake_triton)
+    monkeypatch.setattr(
+        mod,
+        "_rocm_sparse_attn_prefill_ragged_triton",
+        lambda *args, **kwargs: pytest.fail("unexpected ragged Triton fallback"),
+    )
+    monkeypatch.setattr(
+        mod,
+        "build_ragged_indices_from_dense",
+        lambda *args, **kwargs: (
+            torch.zeros(1, dtype=torch.int32),
+            torch.tensor([0, 1], dtype=torch.int32),
+        ),
+    )
+    q = torch.empty(num_queries, 1, HEAD_DIM, dtype=torch.bfloat16)
+    kv = torch.empty(4, 1, HEAD_DIM, dtype=torch.bfloat16)
+    output = torch.empty_like(q)
+    mod.rocm_sparse_attn_prefill(
+        q=q,
+        kv=kv,
+        indices=torch.zeros(num_queries, 1, dtype=torch.int32),
+        topk_length=torch.ones(num_queries, dtype=torch.int32),
+        scale=HEAD_DIM**-0.5,
+        head_dim=HEAD_DIM,
+        nope_head_dim=NOPE_HEAD_DIM,
+        rope_head_dim=ROPE_HEAD_DIM,
+        attn_sink=torch.empty(1, dtype=torch.float32),
+        output=output,
+    )
+
+    if expect_aiter:
+        fake_aiter.assert_called_once()
+        assert fake_aiter.call_args.args[0] is kernel
+        fake_triton.assert_not_called()
+    else:
+        fake_aiter.assert_not_called()
+        fake_triton.assert_called_once()
+
+
 @requires_gfx950
 @torch.inference_mode()
 def test_sparse_attn_prefill_ragged_aiter_opus(monkeypatch) -> None:
diff --git a/tests/kernels/core/test_fused_pre_fc_norms.py b/tests/kernels/core/test_fused_pre_fc_norms.py
new file mode 100644
index 000000000000..5cffb6d1bb73
--- /dev/null
+++ b/tests/kernels/core/test_fused_pre_fc_norms.py
@@ -0,0 +1,66 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for the fused dual-GemmaRMSNorm + concat kernel used by the
+Qwen3.5 MTP drafter.
+
+The kernel must match the eager 3-op reference: two Gemma-style RMSNorms
+(scale by ``1 + weight``, fp32 accumulation) followed by a last-dim concat.
+"""
+
+import pytest
+import torch
+
+from vllm.kernels.triton.fused_pre_fc_norms import fused_pre_fc_norms
+from vllm.platforms import current_platform
+from vllm.utils.torch_utils import set_random_seed
+
+EPS = 1e-6
+DTYPES = [torch.bfloat16, torch.float16, torch.float32]
+NUM_TOKENS = [1, 129]
+HIDDEN_SIZES = [2048, 4096]
+# (rtol, atol) per dtype; the kernel accumulates in fp32 like the reference,
+# so the gap is output rounding and reduction order only.
+TOLS = {
+    torch.bfloat16: (1e-2, 1e-3),
+    torch.float16: (1e-3, 1e-4),
+    torch.float32: (1e-5, 1e-5),
+}
+
+requires_cuda = pytest.mark.skipif(
+    not current_platform.is_cuda(),
+    reason="fused pre-fc norms Triton kernel requires a CUDA device",
+)
+
+
+def _gemma_rmsnorm(x: torch.Tensor, w: torch.Tensor, eps: float) -> torch.Tensor:
+    dtype = x.dtype
+    x = x.float()
+    var = x.pow(2).mean(dim=-1, keepdim=True)
+    return (x * torch.rsqrt(var + eps) * (w.float() + 1.0)).to(dtype)
+
+
+@requires_cuda
+@torch.inference_mode()
+@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
+@pytest.mark.parametrize("hidden_size", HIDDEN_SIZES)
+@pytest.mark.parametrize("dtype", DTYPES)
+def test_fused_pre_fc_norms_correctness(
+    num_tokens: int, hidden_size: int, dtype: torch.dtype
+):
+    """Fused kernel == cat([norm_e(embeds), norm_h(hidden)]) for Gemma-style
+    (1 + w) norms with non-zero weights."""
+    set_random_seed(0)
+    embeds = torch.randn(num_tokens, hidden_size, dtype=dtype, device="cuda")
+    hidden = torch.randn(num_tokens, hidden_size, dtype=dtype, device="cuda")
+    emb_w = torch.randn(hidden_size, dtype=dtype, device="cuda") * 0.2
+    hid_w = torch.randn(hidden_size, dtype=dtype, device="cuda") * 0.2
+
+    fused = fused_pre_fc_norms(embeds, hidden, emb_w, hid_w, EPS)
+    ref = torch.cat(
+        [_gemma_rmsnorm(embeds, emb_w, EPS), _gemma_rmsnorm(hidden, hid_w, EPS)],
+        dim=-1,
+    )
+
+    assert fused.shape == (num_tokens, 2 * hidden_size)
+    rtol, atol = TOLS[dtype]
+    torch.testing.assert_close(fused.float(), ref.float(), rtol=rtol, atol=atol)
diff --git a/tests/kernels/core/test_fused_sigmoid_gate_mul.py b/tests/kernels/core/test_fused_sigmoid_gate_mul.py
new file mode 100644
index 000000000000..3338e8a4f713
--- /dev/null
+++ b/tests/kernels/core/test_fused_sigmoid_gate_mul.py
@@ -0,0 +1,140 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+
+from vllm.kernels.triton.fused_sigmoid_gate_mul import fused_sigmoid_gate_mul
+from vllm.platforms import current_platform
+from vllm.utils.torch_utils import set_random_seed
+
+DTYPES = [torch.bfloat16, torch.float16, torch.float32]
+# Odd token counts, non-power-of-2 and large hidden, plus an empty batch.
+SHAPES = [(1, 512), (3, 1792), (17, 4096), (128, 8192), (1023, 2048), (0, 1024)]
+requires_cuda = pytest.mark.skipif(
+    not current_platform.is_cuda_alike(),
+    reason="fused_sigmoid_gate_mul Triton kernel requires a CUDA/ROCm device",
+)
+
+
+_TOLS = {
+    torch.float32: (1e-6, 1e-5),
+    torch.float16: (2e-3, 2e-3),
+    torch.bfloat16: (1e-2, 1e-2),
+}
+
+
+@requires_cuda
+@pytest.mark.parametrize("shape", SHAPES)
+@pytest.mark.parametrize("dtype", DTYPES)
+@torch.inference_mode()
+def test_fused_sigmoid_gate_mul_correctness(shape: tuple[int, int], dtype: torch.dtype):
+    """The fused kernel must match ``x * torch.sigmoid(gate)`` and must not
+    write in place."""
+    set_random_seed(13)
+    num_tokens, hidden = shape
+    x = torch.randn(num_tokens, hidden, dtype=dtype, device="cuda")
+    # Scale the gate to exercise both tails of the sigmoid.
+    gate = torch.randn(num_tokens, hidden, dtype=dtype, device="cuda") * 4
+
+    result = fused_sigmoid_gate_mul(x, gate)
+    reference = x * torch.sigmoid(gate)
+
+    if result.numel() > 0:
+        assert result.data_ptr() != x.data_ptr()
+    atol, rtol = _TOLS[dtype]
+    torch.testing.assert_close(result, reference, atol=atol, rtol=rtol)
+
+
+@requires_cuda
+@torch.inference_mode()
+def test_fused_sigmoid_gate_mul_offset_row_view():
+    """A row-offset (contiguous but misaligned-pointer) view must compile and
+    match eager -- this exercises the unaligned pointer specialization."""
+    set_random_seed(13)
+    base_x = torch.randn(65, 1023, dtype=torch.bfloat16, device="cuda")
+    base_g = torch.randn(65, 1023, dtype=torch.bfloat16, device="cuda")
+    x, gate = base_x[1:], base_g[1:]
+
+    result = fused_sigmoid_gate_mul(x, gate)
+    reference = x * torch.sigmoid(gate)
+
+    torch.testing.assert_close(result, reference, atol=1e-2, rtol=1e-2)
+
+
+@torch.inference_mode()
+def test_fused_sigmoid_gate_mul_fallback_noncontiguous():
+    """Non-CUDA or non-contiguous inputs take the eager fallback (CPU-safe)."""
+    x = torch.randn(64, 256, dtype=torch.float32).t()
+    gate = torch.randn(64, 256, dtype=torch.float32).t()
+    assert not x.is_contiguous()
+
+    result = fused_sigmoid_gate_mul(x, gate)
+    expected = x * torch.sigmoid(gate)
+
+    torch.testing.assert_close(result, expected, atol=0.0, rtol=0.0)
+
+
+def _fp8_codes(q: torch.Tensor) -> torch.Tensor:
+    """Sign-magnitude e4m3 bytes -> ordered ints (+0 == -0)."""
+    b = q.view(torch.uint8).int()
+    return torch.where(b >= 128, -(b & 0x7F), b)
+
+
+@requires_cuda
+@pytest.mark.skipif(
+    not current_platform.is_cuda() or not current_platform.has_device_capability(89),
+    reason="FP8 (e4m3fn) output requires NVIDIA SM89+",
+)
+# o_proj input widths at TP1 (6144) and TP2 (3072) of a 24x256-head model.
+@pytest.mark.parametrize("shape", [(1, 6144), (37, 6144), (5, 3072), (300, 4096)])
+@torch.inference_mode()
+def test_fused_sigmoid_gate_mul_fp8_matches_chain(shape: tuple[int, int]):
+    """The FP8 output (fused o_proj input quant) equals the bf16 kernel
+    followed by static scaled_fp8_quant."""
+    from vllm import _custom_ops as ops
+
+    g = torch.Generator(device="cuda").manual_seed(shape[0])
+    x = torch.randn(shape, generator=g, dtype=torch.bfloat16, device="cuda") * 8
+    gate = torch.randn(shape, generator=g, dtype=torch.bfloat16, device="cuda") * 4
+    scale = torch.tensor([0.02], device="cuda")
+
+    ref_q, _ = ops.scaled_fp8_quant(
+        fused_sigmoid_gate_mul(x, gate), scale, group_shape=(-1, -1)
+    )
+    q = fused_sigmoid_gate_mul(x, gate, scale)
+
+    assert q.dtype == torch.float8_e4m3fn
+    assert ref_q.float().abs().max().item() == 448.0
+    # Arch-dependent compilation flips occasional bf16 rounding boundaries;
+    # codes must be within one FP8 step.
+    dist = (_fp8_codes(q) - _fp8_codes(ref_q)).abs()
+    assert dist.max().item() <= 1
+    assert (dist > 0).float().mean().item() <= 1e-4
+
+
+@requires_cuda
+@pytest.mark.skipif(
+    not current_platform.is_cuda() or not current_platform.has_device_capability(89),
+    reason="FP8 (e4m3fn) output requires NVIDIA SM89+",
+)
+@torch.inference_mode()
+def test_fused_sigmoid_gate_mul_fp8_strided_gate_falls_back():
+    """A strided gate view (num_heads == 1 per TP rank keeps the packed qkv
+    row stride through reshape) cannot run the fused kernel; the FP8 path
+    must quantize the eager result instead of asserting."""
+    from vllm import _custom_ops as ops
+
+    x = torch.randn(5, 128, dtype=torch.bfloat16, device="cuda")
+    packed = torch.randn(5, 256, dtype=torch.bfloat16, device="cuda")
+    gate = packed[:, :128].view(5, 1, 128).reshape(5, 128)
+    assert not gate.is_contiguous()
+    scale = torch.tensor([0.02], device="cuda")
+
+    ref_q, _ = ops.scaled_fp8_quant(
+        x * torch.sigmoid(gate), scale, group_shape=(-1, -1)
+    )
+    q = fused_sigmoid_gate_mul(x, gate, scale)
+
+    assert q.dtype == torch.float8_e4m3fn
+    assert torch.equal(q.view(torch.uint8), ref_q.view(torch.uint8))
diff --git a/tests/kernels/core/test_vit_fp8_scaling.py b/tests/kernels/core/test_vit_fp8_scaling.py
index d756bfedbb22..af67f9252b9d 100644
--- a/tests/kernels/core/test_vit_fp8_scaling.py
+++ b/tests/kernels/core/test_vit_fp8_scaling.py
@@ -79,7 +79,7 @@ def _build_attention(mm_config, attn_backend=None):
             head_size=HEAD_DIM,
             prefix=LAYER_0,
         )
-        attn.process_weights_after_loading(torch.bfloat16)
+        attn.process_weights_after_loading()
         yield attn
 
 
@@ -247,7 +247,7 @@ def test_static_scales_missing_layer(tmp_path) -> None:
             prefix=LAYER_0,
         )
         with pytest.raises(ValueError, match="scales not found for layer"):
-            attn.process_weights_after_loading(torch.bfloat16)
+            attn.process_weights_after_loading()
 
 
 def test_dynamic_scales_auto_save(tmp_path) -> None:
diff --git a/tests/kernels/fla/test_chunk_int64_offsets.py b/tests/kernels/fla/test_chunk_int64_offsets.py
new file mode 100644
index 000000000000..06168988c9e1
--- /dev/null
+++ b/tests/kernels/fla/test_chunk_int64_offsets.py
@@ -0,0 +1,116 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""FLA chunk kernels on batches whose offsets do not fit in int32."""
+
+import pytest
+import torch
+
+from vllm.platforms import current_platform
+from vllm.third_party.flash_linear_attention.ops.chunk_delta_h import (
+    chunk_gated_delta_rule_fwd_h,
+)
+from vllm.third_party.flash_linear_attention.ops.solve_tril import solve_tril
+
+pytestmark = pytest.mark.skipif(
+    not current_platform.is_cuda_alike(), reason="needs a GPU"
+)
+
+BT = 64
+
+
+def _skip_if_low_memory(need: int) -> None:
+    free, _ = current_platform.mem_get_info()
+    if free < need + (1 << 30):
+        pytest.skip("not enough free device memory")
+
+
+@torch.inference_mode()
+def test_chunk_delta_h_int64_seq_offset():
+    """A later sequence's start offsets need int64.
+
+    The second sequence starts at token 2^31 // (H * V), so its v, w and v_new
+    offsets (bos * H * V) and its h offset (chunk * H * V * K) pass 2^31.
+    """
+    H, Hg, K, V = 64, 16, 128, 128
+    L1, L2 = 2**31 // (H * V), 2 * BT
+    T = L1 + L2
+    # k, w, u, v_new and h, all bf16.
+    _skip_if_low_memory(T * (Hg * K + 3 * H * V) * 2 + T // BT * H * V * K * 2)
+
+    device = torch.device("cuda")
+    gen = torch.Generator(device=device).manual_seed(0)
+
+    def randn(*shape, scale=1.0):
+        x = torch.randn(*shape, device=device, generator=gen, dtype=torch.bfloat16)
+        return x.mul_(scale)
+
+    k = randn(1, T, Hg, K, scale=0.09)
+    w = randn(1, T, H, K, scale=0.05)
+    u = randn(1, T, H, V)
+    # Chunk-local cumulative log decay, as chunk_local_cumsum produces it.
+    g = torch.rand(1, T // BT, BT, H, device=device, generator=gen) * -0.05
+    g = g.cumsum(2).view(1, T, H)
+    h0 = torch.randn(2, H, V, K, device=device, generator=gen) * 0.1
+    cu_seqlens = torch.tensor([0, L1, T], dtype=torch.int32, device=device)
+
+    h, v_new, ht = chunk_gated_delta_rule_fwd_h(
+        k,
+        w,
+        u,
+        g=g,
+        initial_state=h0,
+        output_final_state=True,
+        cu_seqlens=cu_seqlens,
+    )
+
+    # Reference: the second sequence on its own, where its offsets are 0.
+    h_ref, v_new_ref, ht_ref = chunk_gated_delta_rule_fwd_h(
+        k[:, L1:].contiguous(),
+        w[:, L1:].contiguous(),
+        u[:, L1:].contiguous(),
+        g=g[:, L1:].contiguous(),
+        initial_state=h0[1:].contiguous(),
+        output_final_state=True,
+        cu_seqlens=torch.tensor([0, L2], dtype=torch.int32, device=device),
+    )
+    assert v_new is not None and ht is not None
+    assert v_new_ref is not None and ht_ref is not None
+    torch.testing.assert_close(h[:, L1 // BT :], h_ref, atol=0, rtol=0)
+    torch.testing.assert_close(v_new[:, L1:], v_new_ref, atol=0, rtol=0)
+    torch.testing.assert_close(ht[1:], ht_ref, atol=0, rtol=0)
+
+
+@torch.inference_mode()
+def test_solve_tril_int64_offsets():
+    """Rows past 2^31 // (H * BT) need int64 offsets.
+
+    The first sequence runs past row 2^31 // (H * BT), which the per-row loads
+    reach, and the second sequence starts past it, which the base offset
+    (bos * H * BT) reaches.
+    """
+    H = 64
+    L1, L2 = 2**31 // (H * BT) + 2 * BT, 2 * BT
+    T = L1 + L2
+    # A in fp32 and its inverse in bf16.
+    _skip_if_low_memory(T * H * BT * (4 + 2))
+
+    device = torch.device("cuda")
+    gen = torch.Generator(device=device).manual_seed(0)
+    # Strictly lower triangular inside each chunk of BT rows.
+    A = torch.randn(1, T, H, BT, device=device, generator=gen).mul_(0.1)
+    mask = torch.ones(BT, BT, device=device).tril(-1)
+    A.view(1, T // BT, BT, H, BT).mul_(mask[:, None, :])
+    cu_seqlens = torch.tensor([0, L1, T], dtype=torch.int32, device=device)
+
+    Ai = solve_tril(A, cu_seqlens=cu_seqlens, output_dtype=torch.bfloat16)
+
+    # Chunks are independent, so the last two chunks of each sequence on
+    # their own are the reference.
+    for end in (L1, T):
+        start = end - 2 * BT
+        ref = solve_tril(
+            A[:, start:end].contiguous(),
+            cu_seqlens=torch.tensor([0, 2 * BT], dtype=torch.int32, device=device),
+            output_dtype=torch.bfloat16,
+        )
+        torch.testing.assert_close(Ai[:, start:end], ref, atol=0, rtol=0)
diff --git a/tests/kernels/fla/test_fla_layernorm_guard.py b/tests/kernels/fla/test_fla_layernorm_guard.py
index 9223fb7110ff..5ba3e36723ae 100644
--- a/tests/kernels/fla/test_fla_layernorm_guard.py
+++ b/tests/kernels/fla/test_fla_layernorm_guard.py
@@ -545,6 +545,39 @@ def test_rmsnorm_gated_3d_matches_flattened_2d(
     assert cuda_3d.shape == (num_tokens, num_heads, head_dim)
 
 
+@torch.inference_mode()
+def test_rmsnorm_gated_int64_row_offset() -> None:
+    """Rows past 2^31 // row stride need int64 offsets.
+
+    Mirrors RMSNormGated on the GDN output: [num_tokens, num_heads, 128]
+    flattened to rows of 128, so tokens from 2^31 // (num_heads * 128) on
+    (262,144 for 64 heads) wrap in int32.
+    """
+    num_heads, head_dim = 64, 128
+    num_tokens = 2**31 // (num_heads * head_dim) + 4
+    # x, z and the output, all bf16.
+    need = num_tokens * num_heads * head_dim * 2 * 3
+    free, _ = current_platform.mem_get_info()
+    if free < need + (1 << 30):
+        pytest.skip("not enough free device memory")
+
+    set_random_seed(0)
+    device = torch.device(DEVICE)
+    shape = (num_tokens, num_heads, head_dim)
+    x = torch.randn(shape, dtype=torch.bfloat16, device=device)
+    z = torch.randn(shape, dtype=torch.bfloat16, device=device)
+    weight = torch.randn(head_dim, dtype=torch.bfloat16, device=device)
+
+    out = layernorm_fn(
+        x, weight, None, z=z, eps=1e-6, norm_before_gate=True, is_rms_norm=True
+    )
+
+    # Reference on the first token and the last tokens, past the int32 boundary.
+    rows = torch.tensor([0, *range(num_tokens - 4, num_tokens)], device=device)
+    ref = rms_norm_ref(x[rows], weight, None, z=z[rows], eps=1e-6)
+    torch.testing.assert_close(out[rows], ref, atol=1e-2, rtol=1e-2)
+
+
 if __name__ == "__main__":
     # Run a quick smoke test
     test_layer_norm_fwd_basic(128, 1024, torch.float16, 42, False)
diff --git a/tests/kernels/fla/test_fused_gdn_post_conv.py b/tests/kernels/fla/test_fused_gdn_post_conv.py
index a29395edac9b..df5526bc3676 100644
--- a/tests/kernels/fla/test_fused_gdn_post_conv.py
+++ b/tests/kernels/fla/test_fused_gdn_post_conv.py
@@ -220,6 +220,45 @@ def test_fused_post_conv_l0():
     assert g.shape == (0, HV)
 
 
+def test_fused_post_conv_int64_token_offset():
+    """Rows past 2^31 // qkv_dim need int64 offsets."""
+    device = "cuda"
+    H, HV, K, V = 16, 32, 128, 128
+    qkv_dim = 2 * H * K + HV * V
+    L = 2**31 // qkv_dim + 16
+    # conv_output plus the q, k and v outputs, all bf16.
+    need = L * (qkv_dim + 2 * H * K + HV * V) * 2
+    free, _ = current_platform.mem_get_info()
+    if free < need + (1 << 30):
+        pytest.skip("not enough free device memory")
+
+    torch.manual_seed(0)
+    conv_output = torch.randn(L, qkv_dim, dtype=torch.bfloat16, device=device)
+    a = torch.randn(L, HV, dtype=torch.bfloat16, device=device)
+    b = torch.randn(L, HV, dtype=torch.bfloat16, device=device)
+    A_log = torch.randn(HV, dtype=torch.float32, device=device) - 2.0
+    dt_bias = torch.randn(HV, dtype=torch.float32, device=device) * 0.1
+
+    fused = fused_post_conv_prep(
+        conv_output,
+        a,
+        b,
+        A_log,
+        dt_bias,
+        num_k_heads=H,
+        head_k_dim=K,
+        head_v_dim=V,
+    )
+
+    # Reference on the first row and the last rows, past the int32 boundary.
+    rows = torch.tensor([0, *range(L - 16, L)], device=device)
+    ref = reference_post_conv(
+        conv_output[rows], a[rows], b[rows], A_log, dt_bias, H, K, V
+    )
+    for out, expected in zip(fused, ref):
+        torch.testing.assert_close(out[rows], expected, atol=1e-3, rtol=1e-3)
+
+
 @pytest.mark.parametrize(
     "head_ratio,tp_size,query_lengths,state_dtype,norm_dtype",
     [
@@ -343,6 +382,10 @@ def test_fused_gdn_decode_post_conv_mtp_head_ratios(
     if query_lengths[-1] == 0:
         state_indices[-1].zero_()
 
+    output_rtol = 5e-4
+    if impl == "cuda" and state_dtype == torch.float32 and norm_dtype == torch.bfloat16:
+        output_rtol = 1e-3
+
     for step, accepted_tokens in enumerate((1, min(2, state_width), state_width)):
         num_accepted_tokens.fill_(accepted_tokens)
         if query_lengths[-1] == 0:
@@ -392,9 +435,54 @@ def test_fused_gdn_decode_post_conv_mtp_head_ratios(
 
         output_error = (actual.float() - expected.float()).norm()
         output_relative_l2 = output_error / expected.float().norm().clamp_min(1e-20)
-        assert output_relative_l2 < 5e-4, (
+        assert output_relative_l2 < output_rtol, (
             f"MTP output relative L2 mismatch at step {step}: "
             f"{output_relative_l2.item():.6g}"
         )
 
     torch.testing.assert_close(state_actual, state_ref, atol=3e-2, rtol=3e-2)
+
+
+@torch.inference_mode()
+def test_fused_post_conv_prep_strided_ba_bitwise() -> None:
+    """Packed b/a chunk views (row stride 2*HV) must match contiguous inputs
+    bitwise."""
+    torch.manual_seed(0)
+    dtype = torch.bfloat16
+    device = "cuda"
+    num_tokens = 33
+    H, HV, K, V = 2, 4, 128, 128
+    qkv_dim = 2 * H * K + HV * V
+
+    conv_output = torch.randn(num_tokens, qkv_dim, dtype=dtype, device=device)
+    ba = torch.randn(num_tokens, 2 * HV, dtype=dtype, device=device)
+    b_view, a_view = ba.chunk(2, dim=-1)
+    A_log = torch.randn(HV, dtype=torch.float32, device=device)
+    dt_bias = torch.randn(HV, dtype=torch.float32, device=device)
+
+    ref = fused_post_conv_prep(
+        conv_output=conv_output,
+        a=a_view.contiguous(),
+        b=b_view.contiguous(),
+        A_log=A_log,
+        dt_bias=dt_bias,
+        num_k_heads=H,
+        head_k_dim=K,
+        head_v_dim=V,
+        apply_l2norm=True,
+        output_g_exp=False,
+    )
+    got = fused_post_conv_prep(
+        conv_output=conv_output,
+        a=a_view,
+        b=b_view,
+        A_log=A_log,
+        dt_bias=dt_bias,
+        num_k_heads=H,
+        head_k_dim=K,
+        head_v_dim=V,
+        apply_l2norm=True,
+        output_g_exp=False,
+    )
+    for r, g in zip(ref, got, strict=True):
+        torch.testing.assert_close(g, r, rtol=0, atol=0)
diff --git a/tests/kernels/fla/test_fused_recurrent_packed_decode.py b/tests/kernels/fla/test_fused_recurrent_packed_decode.py
index 5ea66ced2390..e43734a49f16 100644
--- a/tests/kernels/fla/test_fused_recurrent_packed_decode.py
+++ b/tests/kernels/fla/test_fused_recurrent_packed_decode.py
@@ -186,3 +186,57 @@ def test_packed_decode_supports_large_batch_head_grid():
     )
 
     assert torch.count_nonzero(out).item() == 0
+
+
+@pytest.mark.parametrize("num_tokens", [1, 9])
+@torch.inference_mode()
+def test_packed_decode_strided_ba_bitwise(num_tokens: int) -> None:
+    """Packed b/a chunk views (row stride 2*HV) must match contiguous inputs
+    bitwise."""
+    torch.manual_seed(0)
+    dtype = torch.bfloat16
+    device = torch.device(DEVICE)
+    H, HV, K, V = 2, 4, 128, 128
+    qkv_dim = 2 * H * K + HV * V
+
+    qkvz = torch.randn(num_tokens, qkv_dim + HV * V, dtype=dtype, device=device)
+    mixed_qkv = qkvz[:, :qkv_dim]
+    ba = torch.randn(num_tokens, 2 * HV, dtype=dtype, device=device)
+    b_view, a_view = ba.chunk(2, dim=-1)
+    A_log = torch.randn(HV, dtype=torch.float32, device=device)
+    dt_bias = torch.randn(HV, dtype=torch.float32, device=device)
+    state = torch.randn(num_tokens + 1, HV, V, K, dtype=torch.float32, device=device)
+    indices = torch.arange(1, num_tokens + 1, dtype=torch.int32, device=device)
+
+    out_strided = torch.empty(num_tokens, 1, HV, V, dtype=dtype, device=device)
+    state_strided = state.clone()
+    fused_recurrent_gated_delta_rule_packed_decode(
+        mixed_qkv=mixed_qkv,
+        a=a_view,
+        b=b_view,
+        A_log=A_log,
+        dt_bias=dt_bias,
+        scale=K**-0.5,
+        initial_state=state_strided,
+        out=out_strided,
+        ssm_state_indices=indices,
+        use_qk_l2norm_in_kernel=True,
+    )
+
+    out_contig = torch.empty(num_tokens, 1, HV, V, dtype=dtype, device=device)
+    state_contig = state.clone()
+    fused_recurrent_gated_delta_rule_packed_decode(
+        mixed_qkv=mixed_qkv.contiguous(),
+        a=a_view.contiguous(),
+        b=b_view.contiguous(),
+        A_log=A_log,
+        dt_bias=dt_bias,
+        scale=K**-0.5,
+        initial_state=state_contig,
+        out=out_contig,
+        ssm_state_indices=indices,
+        use_qk_l2norm_in_kernel=True,
+    )
+
+    torch.testing.assert_close(out_strided, out_contig, rtol=0, atol=0)
+    torch.testing.assert_close(state_strided, state_contig, rtol=0, atol=0)
diff --git a/tests/kernels/mamba/test_checkpoint_exporter.py b/tests/kernels/mamba/test_checkpoint_exporter.py
index fdd4dbe6f97d..6bca64be13ab 100644
--- a/tests/kernels/mamba/test_checkpoint_exporter.py
+++ b/tests/kernels/mamba/test_checkpoint_exporter.py
@@ -172,3 +172,41 @@ def test_exporter_stores_each_callers_checkpoints(caller: str, state_layout: str
 
     assert torch.equal(conv_state, expected_conv)
     assert torch.equal(recurrent_state, expected_recurrent)
+
+
+@pytest.mark.skipif(not current_platform.is_cuda_alike(), reason="needs a GPU")
+def test_exporter_int64_token_offset():
+    """Conv windows of rows past 2^31 // row stride need int64 input offsets.
+
+    Qwen3.5 passes ``raw_qkv`` as the qkv slice of the packed qkvz projection,
+    so its row stride is wider than its width.
+    """
+    width, row_stride, state_len, slot = 8192, 12288, 3, 5
+    num_tokens = 2**31 // row_stride + 64
+    free, _ = current_platform.mem_get_info()
+    if free < num_tokens * row_stride * 2 + (1 << 30):
+        pytest.skip("not enough free device memory")
+
+    device = torch.device("cuda")
+    gen = torch.Generator(device=device).manual_seed(0)
+    qkvz = torch.randn(
+        num_tokens, row_stride, device=device, dtype=torch.bfloat16, generator=gen
+    )
+    raw_qkv = qkvz[:, :width]
+    cu_seqlens = torch.tensor([0, num_tokens], dtype=torch.int32, device=device)
+    conv_state = torch.zeros(
+        NUM_SLOTS, width, state_len, device=device, dtype=torch.bfloat16
+    )
+    offset = num_tokens - 16
+    expected_conv = conv_state.clone()
+    expected_conv[slot] = raw_qkv[offset - state_len : offset].T
+
+    checkpoint = MambaPrefillCheckpointMetadata(
+        torch.tensor([offset], dtype=torch.int32, device=device),
+        torch.tensor([slot], dtype=torch.int32, device=device),
+    )
+    MambaPrefillCheckpointExporter(state_len=state_len).export(
+        checkpoint, raw_qkv=raw_qkv, conv_state=conv_state, cu_seqlens=cu_seqlens
+    )
+
+    assert torch.equal(conv_state, expected_conv)
diff --git a/tests/kernels/mamba/test_gdn_forward_core_split.py b/tests/kernels/mamba/test_gdn_forward_core_split.py
index 8a99e223be3b..7e8b9ef238cb 100644
--- a/tests/kernels/mamba/test_gdn_forward_core_split.py
+++ b/tests/kernels/mamba/test_gdn_forward_core_split.py
@@ -123,6 +123,7 @@ def _build_layer(
     layer.kv_cache = (conv_state, ssm_state)
     with set_current_vllm_config(vllm_config):
         layer.chunk_gated_delta_rule = ChunkGatedDeltaRule()
+    layer.gdn_prefill_backend = layer.chunk_gated_delta_rule.gdn_prefill_backend
     for name in (
         "rearrange_mixed_qkv",
         "_forward_core",
diff --git a/tests/kernels/mamba/test_gdn_fused_mtp.py b/tests/kernels/mamba/test_gdn_fused_mtp.py
index 24ea20325ef8..f2a6f1698e09 100644
--- a/tests/kernels/mamba/test_gdn_fused_mtp.py
+++ b/tests/kernels/mamba/test_gdn_fused_mtp.py
@@ -128,6 +128,7 @@ def _build_layer(
     )
     with set_current_vllm_config(vllm_config):
         layer.chunk_gated_delta_rule = ChunkGatedDeltaRule()
+    layer.gdn_prefill_backend = layer.chunk_gated_delta_rule.gdn_prefill_backend
     for name in (
         "rearrange_mixed_qkv",
         "_forward_core",
@@ -200,6 +201,7 @@ def test_fused_forward_uses_packed_entrypoint() -> None:
     layer = types.SimpleNamespace(
         prefix=PREFIX,
         enable_fused_gdn_spec_decode=True,
+        _ba_stream=None,
         norm=types.SimpleNamespace(
             weight=torch.empty(V, dtype=torch.bfloat16, device=device)
         ),
@@ -242,18 +244,35 @@ def packed_op(
 
 
 @pytest.mark.parametrize(
-    "seq_lens,query_lens,draft_tokens,expected_fused_calls",
+    "seq_lens,query_lens,device_query_lens,draft_tokens,expected_fused_calls",
     [
-        pytest.param([128], [SPEC_TOKENS], [NUM_SPEC], 1, id="pure-mtp"),
+        pytest.param([128], [SPEC_TOKENS], None, [NUM_SPEC], 1, id="pure-mtp"),
+        pytest.param(
+            [128, 128, 128],
+            [3, 3, 3],
+            [1, 4, 4],
+            [2, 2, 2],
+            1,
+            id="adaptive-ragged-mtp",
+        ),
+        pytest.param(
+            [128] * 8,
+            [1] * 8,
+            [4, 4, 0, 0, 0, 0, 0, 0],
+            [NUM_SPEC] * 8,
+            1,
+            id="adaptive-padded-ragged-mtp",
+        ),
         pytest.param(
             [128, 96],
             [SPEC_TOKENS, 64],
+            None,
             [NUM_SPEC, -1],
             0,
             id="mixed-mtp-falls-back",
         ),
-        pytest.param([96], [64], [-1], 0, id="pure-prefill"),
-        pytest.param([128], [1], [-1], 0, id="pure-decode"),
+        pytest.param([96], [64], None, [-1], 0, id="pure-prefill"),
+        pytest.param([128], [1], None, [-1], 0, id="pure-decode"),
     ],
 )
 @pytest.mark.parametrize("output_gate_activation", ["silu", "sigmoid"])
@@ -261,6 +280,7 @@ def packed_op(
 def test_fused_model_path_matches_reference(
     seq_lens: list[int],
     query_lens: list[int],
+    device_query_lens: list[int] | None,
     draft_tokens: list[int],
     expected_fused_calls: int,
     output_gate_activation: str,
@@ -284,6 +304,18 @@ def test_fused_model_path_matches_reference(
     common = create_common_attn_metadata(
         batch, BLOCK_SIZE, device, arange_block_indices=True
     )
+    if device_query_lens is not None:
+        # Simulate adaptive verification token reallocation
+        assert sum(device_query_lens) == sum(query_lens)
+        device_query_start_loc = torch.zeros(
+            len(device_query_lens) + 1, dtype=torch.int32, device=device
+        )
+        torch.cumsum(
+            torch.tensor(device_query_lens, dtype=torch.int32, device=device),
+            dim=0,
+            out=device_query_start_loc[1:],
+        )
+        common = common.replace(query_start_loc=device_query_start_loc)
     common.block_table_tensor.add_(1)
     with set_current_vllm_config(vllm_config):
         metadata = builder.build(
diff --git a/tests/kernels/mamba/test_gdn_output_projection.py b/tests/kernels/mamba/test_gdn_output_projection.py
index f9abfb4ac666..1094b2920db9 100644
--- a/tests/kernels/mamba/test_gdn_output_projection.py
+++ b/tests/kernels/mamba/test_gdn_output_projection.py
@@ -18,6 +18,7 @@
 from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import (
     QwenGatedDeltaNetAttention,
 )
+from vllm.platforms import current_platform
 
 
 @pytest.mark.parametrize("dtype", [torch.float32])
@@ -51,3 +52,122 @@ def test_output_projection_norms_per_head_and_flattens(
     assert z.shape == (num_tokens, num_heads, head_dim)
     assert out.shape == (num_tokens, hidden)
     assert out.dtype == dtype
+
+
+@pytest.mark.parametrize("num_tokens", [1, 37])
+@pytest.mark.parametrize("norm_before_gate", [True, False])
+@pytest.mark.parametrize("activation", ["silu", "sigmoid"])
+@pytest.mark.skipif(
+    not current_platform.is_cuda(),
+    reason="the fused gated RMS norm Triton kernel requires CUDA",
+)
+@torch.inference_mode()
+def test_fused_gated_rms_norm_strided_z_bitwise(
+    num_tokens: int, norm_before_gate: bool, activation: str
+) -> None:
+    """The Qwen3.5 fused gated RMS norm reads z from its packed strided view
+    of the qkvz projection and matches ``rmsnorm_fn`` on the copied/reshaped
+    z bitwise."""
+    from vllm.models.qwen3_5.nvidia.ops.fused_gdn_gated_norm import (
+        fused_gdn_gated_rms_norm,
+    )
+    from vllm.third_party.flash_linear_attention.ops.layernorm_guard import (
+        rmsnorm_fn,
+    )
+
+    torch.manual_seed(0)
+    dtype = torch.bfloat16
+    num_heads, head_dim = 4, 128
+    qkvz_row = 2 * 2 * 128 + 2 * num_heads * head_dim
+    qkvz = torch.randn(num_tokens, qkvz_row, dtype=dtype, device="cuda")
+    z_view = qkvz[:, qkvz_row - num_heads * head_dim :]
+    x = torch.randn(num_tokens, num_heads, head_dim, dtype=dtype, device="cuda")
+    weight = torch.randn(head_dim, dtype=torch.float32, device="cuda")
+    eps = 1e-6
+
+    ref = rmsnorm_fn(
+        x,
+        weight,
+        None,
+        z=z_view.reshape(num_tokens, num_heads, head_dim),
+        eps=eps,
+        group_size=None,
+        norm_before_gate=norm_before_gate,
+        activation=activation,
+    )
+
+    out = torch.empty(num_tokens, num_heads * head_dim, dtype=dtype, device="cuda")
+    fused_gdn_gated_rms_norm(
+        x=x.view(num_tokens, -1),
+        z=z_view,
+        weight=weight,
+        out=out,
+        eps=eps,
+        num_v_heads=num_heads,
+        head_v_dim=head_dim,
+        norm_before_gate=norm_before_gate,
+        activation=activation,
+    )
+
+    torch.testing.assert_close(
+        out.view(num_tokens, num_heads, head_dim), ref, rtol=0, atol=0
+    )
+
+
+# 48/24/12 value heads per rank: TP1/2/4 of a 48-head model.
+@pytest.mark.parametrize("num_heads", [48, 24, 12])
+@pytest.mark.parametrize("num_tokens", [1, 37, 300])
+@pytest.mark.skipif(
+    not current_platform.is_cuda() or not current_platform.has_device_capability(89),
+    reason="the FP8 (e4m3fn) output of the fused gated RMS norm requires NVIDIA SM89+",
+)
+@torch.inference_mode()
+def test_fused_gated_rms_norm_fp8_matches_chain(
+    num_tokens: int, num_heads: int
+) -> None:
+    """The FP8 output (fused out_proj input quant) equals the bf16 kernel
+    followed by static scaled_fp8_quant."""
+    from vllm import _custom_ops as ops
+    from vllm.models.qwen3_5.nvidia.ops.fused_gdn_gated_norm import (
+        fused_gdn_gated_rms_norm,
+    )
+
+    g = torch.Generator(device="cuda").manual_seed(num_tokens)
+    dtype = torch.bfloat16
+    head_dim = 128
+    qkvz_row = 2 * (num_heads // 3) * 128 + 2 * num_heads * head_dim
+    qkvz = torch.randn(num_tokens, qkvz_row, generator=g, dtype=dtype, device="cuda")
+    z_view = qkvz[:, qkvz_row - num_heads * head_dim :]
+    x = torch.randn(
+        num_tokens, num_heads * head_dim, generator=g, dtype=dtype, device="cuda"
+    )
+    weight = (torch.rand(head_dim, generator=g, device="cuda") + 0.5).to(dtype)
+    # Small enough that the largest outputs saturate at +-448.
+    scale = torch.tensor([0.004], device="cuda")
+    kwargs = dict(
+        z=z_view,
+        weight=weight,
+        eps=1e-6,
+        num_v_heads=num_heads,
+        head_v_dim=head_dim,
+        norm_before_gate=True,
+        activation="silu",
+    )
+
+    ref = torch.empty_like(x)
+    fused_gdn_gated_rms_norm(x=x, out=ref, **kwargs)
+    ref_q, _ = ops.scaled_fp8_quant(ref, scale, group_shape=(-1, -1))
+    q = torch.empty_like(x, dtype=torch.float8_e4m3fn)
+    fused_gdn_gated_rms_norm(x=x, out=q, scale=scale, **kwargs)
+
+    assert ref_q.float().abs().max().item() == 448.0
+
+    # Arch-dependent compilation flips occasional bf16 rounding boundaries;
+    # codes must be within one FP8 step (+0 == -0).
+    def codes(t: torch.Tensor) -> torch.Tensor:
+        b = t.view(torch.uint8).int()
+        return torch.where(b >= 128, -(b & 0x7F), b)
+
+    dist = (codes(q) - codes(ref_q)).abs()
+    assert dist.max().item() <= 1
+    assert (dist > 0).float().mean().item() <= 1e-4
diff --git a/tests/kernels/mamba/test_gdn_prefill_flashinfer.py b/tests/kernels/mamba/test_gdn_prefill_flashinfer.py
index be6ab89b3c43..9c694573a027 100644
--- a/tests/kernels/mamba/test_gdn_prefill_flashinfer.py
+++ b/tests/kernels/mamba/test_gdn_prefill_flashinfer.py
@@ -28,7 +28,8 @@
 from vllm.v1.attention.backends.utils import NULL_BLOCK_ID  # noqa: E402
 
 
-def test_flashinfer_gdn_prefill_uses_int64_cu_seqlens(monkeypatch):
+def test_flashinfer_gdn_prefill_passes_cu_seqlens_through(monkeypatch):
+    """FlashInfer accepts int32 cu_seqlens; the wrapper must not cast them."""
     captured_cu_seqlens = None
 
     def fake_chunk_gated_delta_rule(**kwargs):
@@ -57,7 +58,7 @@ def fake_chunk_gated_delta_rule(**kwargs):
     )
 
     assert captured_cu_seqlens is not None
-    assert captured_cu_seqlens.dtype == torch.int64
+    assert captured_cu_seqlens is cu_seqlens
     assert output.shape == q.shape
     assert final_state is None
 
@@ -99,7 +100,7 @@ def test_flashinfer_checkpoints_export_the_state_at_each_offset(use_cp, ssm_dtyp
         torch.randn(1, num_tokens, num_qk_heads, head_dim, device=device), dim=-1
     ).bfloat16()
     v = torch.randn(1, num_tokens, num_v_heads, head_dim, device=device).bfloat16()
-    g = F.logsigmoid(torch.randn(1, num_tokens, num_v_heads, device=device))
+    g = F.logsigmoid(torch.randn(1, num_tokens, num_v_heads, device=device)).exp()
     beta = torch.sigmoid(torch.randn(1, num_tokens, num_v_heads, device=device))
     initial_state = 0.1 * torch.randn(
         len(query_lens), num_v_heads, head_dim, head_dim, device=device
diff --git a/tests/kernels/mfma_spinner.py b/tests/kernels/mfma_spinner.py
new file mode 100644
index 000000000000..09b7e07b20ac
--- /dev/null
+++ b/tests/kernels/mfma_spinner.py
@@ -0,0 +1,64 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""An MFMA co-runner for gfx950 kernel tests: on MI355X, v_pk_mul_f32 can
+return wrong products while another kernel's MFMA waves share the CU."""
+
+import torch
+
+
+def mfma_spinner(n_wg):
+    """``n_wg`` workgroups looping v_mfma_f32_16x16x32_bf16 for ``ticks``
+    (s_memrealtime, 100 per us).
+
+    Launch it on a high-priority stream: on a normal one it can share a hardware
+    queue with the kernel under test and never overlap it.
+    """
+    import flydsl.compiler as flyc
+    import flydsl.expr as fx
+    from flydsl._mlir.dialects import llvm
+    from flydsl.expr import range_constexpr
+    from flydsl.expr.typing import T
+    from flydsl.expr.typing import Vector as Vec
+
+    @fx.struct
+    class Shared:
+        a: fx.Array[fx.Int32, 4096, 16]
+
+    def now():
+        return fx.Int64(
+            llvm.call_intrinsic(T.i64, "llvm.amdgcn.s.memrealtime", [], [], [])
+        )
+
+    @flyc.kernel(name="test_mfma_spinner", known_block_size=[256, 1, 1])
+    def spinner(arg_out: fx.Int64, i64_ticks: fx.Int64):
+        sm = fx.SharedAllocator().allocate(Shared).peek()
+        tx = fx.thread_idx.x
+        sm.a.ptr[tx] = tx
+        a4 = Vec.from_elements([tx, tx + 1, tx + 2, tx + 3], fx.Int32)
+        acc = Vec.filled(4, 0.0, fx.Float32)
+        t0 = now()
+        t = now()
+        while t - t0 < i64_ticks:
+            for _ in range_constexpr(16):
+                ab = Vec(a4).bitcast(fx.BFloat16)
+                acc = Vec(
+                    fx.rocdl.mfma_f32_16x16x32_bf16(
+                        Vec.make_type(4, fx.Float32), [ab, ab, acc, 0, 0, 0]
+                    )
+                )
+            t = now()
+        f32p = fx.PointerType.get(
+            T.f32, address_space=fx.AddressSpace.Global, alignment=4
+        )
+        out = fx.inttoptr(f32p, arg_out)
+        out[fx.block_idx.x * 256 + tx] = acc[0] + fx.Float32(sm.a.ptr[255 - tx])
+
+    @flyc.jit
+    def launch(arg_out: fx.Int64, i64_ticks: fx.Int64, stream: fx.Stream):
+        spinner(arg_out, i64_ticks).launch(
+            grid=(n_wg, 1, 1), block=(256, 1, 1), stream=stream
+        )
+
+    buf = torch.zeros(n_wg * 256, device="cuda")
+    compiled = flyc.compile(launch, buf.data_ptr(), 1, torch.cuda.current_stream())
+    return lambda ticks, stream: compiled(buf.data_ptr(), ticks, stream)
diff --git a/tests/kernels/mhc/test_mhc_kernels.py b/tests/kernels/mhc/test_mhc_kernels.py
index 40695f8a800e..9aa09aebb4de 100644
--- a/tests/kernels/mhc/test_mhc_kernels.py
+++ b/tests/kernels/mhc/test_mhc_kernels.py
@@ -46,6 +46,7 @@
     is_mega_mhc_supported,
     mhc_shifted_post_pre_deep_gemm,
 )
+from vllm.models.deepseek_v41.nvidia.ops.mhc import hc_expand
 from vllm.platforms import current_platform
 from vllm.utils.torch_utils import set_random_seed
 
@@ -1889,3 +1890,36 @@ def test_mhc_fused_post_pre_delayed_falls_back_for_large_batches():
 
     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)
+
+
+@pytest.mark.parametrize(
+    "num_tokens,hidden_size,hc_mult",
+    [(0, 17, 4), (1, 513, 1), (7, 17, 8), (33, 5137, 4), (8192, 5120, 4)],
+)
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
+@pytest.mark.parametrize("strided", [False, True])
+def test_hc_expand_preserves_bits_in_eager_and_graph(
+    num_tokens, hidden_size, hc_mult, dtype, strided
+):
+    """Expansion preserves every input bit, including NaNs and signed zero."""
+    set_random_seed(42)
+    storage = torch.randn(num_tokens, 2 * hidden_size + 1, dtype=dtype, device=DEVICE)
+    x = storage[:, 1::2] if strided else storage[:, :hidden_size].contiguous()
+    if num_tokens:
+        x[0, 0] = float("nan")
+        x[0, -1] = -0.0
+
+    def check(actual):
+        expected = x.unsqueeze(1).expand(-1, hc_mult, -1).contiguous()
+        assert actual.shape == expected.shape and actual.dtype == x.dtype
+        assert actual.is_contiguous()
+        if actual.numel():
+            assert torch.equal(actual.view(torch.uint8), expected.view(torch.uint8))
+
+    check(hc_expand(x, hc_mult))
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        captured = hc_expand(x, hc_mult)
+    x.copy_(torch.randn_like(x))
+    graph.replay()
+    check(captured)
diff --git a/tests/kernels/models/amd/test_dsv41_gfx942_draft_lm_head.py b/tests/kernels/models/amd/test_dsv41_gfx942_draft_lm_head.py
index 2b3f0a8556aa..752292c124fc 100644
--- a/tests/kernels/models/amd/test_dsv41_gfx942_draft_lm_head.py
+++ b/tests/kernels/models/amd/test_dsv41_gfx942_draft_lm_head.py
@@ -6,9 +6,11 @@
 row of zeros, so that vLLM's dispatch (rocm_unquantized_gemm) takes the torch
 GEMM instead of wvSplitK, and the logits of the first 5 rows come back. Every
 other row count, and every call without --monokernel, keeps vLLM's dispatch,
-which sends up to 5 rows to wvSplitK. The tests count the wvSplitK calls of
-vLLM's dispatch on the draft's lm_head shard at TP 4 and compare the logits
-with an fp32 GEMM.
+which sends up to 5 rows to wvSplitK. This shard is 315 MiB, over the
+memory-side cache, so the skinny-GEMM guard would skip wvSplitK for n > 1.
+The counting test lifts that cap to the shard. The guard has its own tests.
+The calls are counted on the draft's lm_head shard at TP 4 and the logits
+are compared with an fp32 GEMM.
 """
 
 import pytest
@@ -75,6 +77,29 @@ def counted(weight, x, *args, **kwargs):
     return rows
 
 
+@pytest.fixture
+def _shard_under_skinny_cap(monkeypatch):
+    """Let this shard take wvSplitK.
+
+    32320 x 5120 bf16 is 315 MiB, over the 256 MiB cap, so the guard
+    skips wvSplitK for n > 1. The cap here is the shard itself. The
+    accessor caches its one evaluation, so clear it on the way in and out.
+    """
+    import vllm.envs as envs
+    from vllm.model_executor.layers.utils import _skinny_gemm_max_weight_bytes
+
+    monkeypatch.setenv(
+        "VLLM_SKINNY_GEMM_MAX_WEIGHT_BYTES",
+        str(VOCAB_SHARD * HIDDEN * torch.bfloat16.itemsize),
+    )
+    envs.disable_envs_cache()
+    _skinny_gemm_max_weight_bytes.cache_clear()
+    yield
+    envs.disable_envs_cache()
+    _skinny_gemm_max_weight_bytes.cache_clear()
+
+
+@pytest.mark.usefixtures("_shard_under_skinny_cap")
 @pytest.mark.parametrize("monokernel", [False, True])
 @pytest.mark.parametrize("rows", [1, 4, 5, 6, 10])
 def test_draft_lm_head_takes_5_rows_off_wvsplitk(
diff --git a/tests/kernels/models/amd/test_dsv41_gfx942_fused_markov.py b/tests/kernels/models/amd/test_dsv41_gfx942_fused_markov.py
index 42bac7e4e7fc..892987dfd7f3 100644
--- a/tests/kernels/models/amd/test_dsv41_gfx942_fused_markov.py
+++ b/tests/kernels/models/amd/test_dsv41_gfx942_fused_markov.py
@@ -80,6 +80,7 @@ def map_draft_to_target(self, ids):
 
 
 def _speculator(temp: float, seed: int):
+    from vllm.v1.worker.gpu.spec_decode.draft_logits import DraftLogitsCache
     from vllm.v1.worker.gpu.spec_decode.dspark.speculator import DSparkSpeculator
 
     dev = torch.device("cuda")
@@ -98,9 +99,9 @@ def _speculator(temp: float, seed: int):
     s._d2t_scatter_index = None
     s._draft_scatter_buf = None
     s.use_fp64_gumbel = False
-    s.draft_logits = torch.zeros(MAX_REQS, N_SPEC, VOCAB, dtype=bf16, device=dev)
-    s.draft_logits_num_shards = 1
-    s.draft_logits_shard_idx = 0
+    s.draft_logits = DraftLogitsCache(
+        MAX_REQS, N_SPEC, VOCAB, dtype=bf16, fill=0.0, device=dev
+    )
     s.temperature = torch.full((MAX_REQS,), temp, dtype=torch.float32, device=dev)
     s.seeds = torch.randint(0, 2**62, (MAX_REQS,), generator=g, device=dev)
     s._step_cols = torch.arange(N_SPEC, dtype=torch.int32, device=dev)
@@ -132,17 +133,17 @@ def test_fused_markov_matches_vllm_loop(monkeypatch, num_reqs, temp):
     _set_gate(monkeypatch, on=False)
     DSparkSpeculator._sample_sequential(s, num_reqs, hidden)
     torch.accelerator.synchronize()
-    ref_tok, ref_cache = s.draft_tokens.clone(), s.draft_logits.clone()
+    ref_tok, ref_cache = s.draft_tokens.clone(), s.draft_logits.values.clone()
 
     s.draft_tokens.zero_()
-    s.draft_logits.zero_()
+    s.draft_logits.values.zero_()
     _set_gate(monkeypatch, on=True)
     try:
         assert fused_markov.sample_sequential(s, num_reqs, hidden)
         torch.accelerator.synchronize()
     finally:
         _set_gate(monkeypatch, on=False)
-    got_tok, got_cache = s.draft_tokens.clone(), s.draft_logits.clone()
+    got_tok, got_cache = s.draft_tokens.clone(), s.draft_logits.values.clone()
 
     # rc and gc are [num_reqs, N_SPEC, VOCAB]: the cached logits of each
     # request and draft step. The differences are counted for each request
diff --git a/tests/kernels/models/amd/test_dsv41_mono_gfx942.py b/tests/kernels/models/amd/test_dsv41_mono_gfx942.py
index 2609e56495ad..13116c109edb 100644
--- a/tests/kernels/models/amd/test_dsv41_mono_gfx942.py
+++ b/tests/kernels/models/amd/test_dsv41_mono_gfx942.py
@@ -624,3 +624,108 @@ def test_contrib_942_is_bf16_of_product(routing):
     torch.accelerator.synchronize()
     ref = (v * w.view(rs.picks, 1)).to(torch.bfloat16).float()
     assert torch.equal(out.view(rs.picks, 16), ref)
+
+
+def _build_nan_route(rs):
+    """stage_route then load_route, then a flat load of each ordered pick.
+
+    The draft's all-NaN scores pick one expert twice. load_route counts picks
+    with a token bitmap, so the order table is short and the unread entries stay
+    at the sentinel. The flat load of that index is the gather that faults."""
+    from vllm.models.deepseek_v41.amd.mono.common.sync import Addr, Mailbox, Space
+
+    words = rs.tab_words
+    tokens, topk = rs.tokens, rs.topk
+    n = tokens * topk
+
+    @fx.struct
+    class Lds:
+        tab: fx.Array[fx.Int32, words, 16]
+        red: fx.Array[fx.Float32, 16, 16]
+
+    @fx.union
+    class U:
+        t: Lds
+
+    @flyc.kernel(
+        name=f"test_nan_route_{tokens}_{topk}", known_block_size=[THREADS, 1, 1]
+    )
+    def kern(logit: Int64, bias: Int64, route: Int64, x: Int64, ids: Int64):
+        tid = fx.thread_idx.x
+        lds = fx.SharedAllocator().allocate(U).t.peek()
+        for i in range_constexpr(-(-words // THREADS)):
+            u = tid + THREADS * i
+            if u < words:
+                fx.ptr_store(fx.Int32(0x7FFFFFFF), lds.tab.ptr + u)
+        gpu.barrier()
+        mb = Mailbox(fx.Int32(1))
+        c = {
+            "S": tokens,
+            "tid": tid,
+            "bid": fx.block_idx.x,
+            "lane": tid % 64,
+            "wave": tid // 64,
+            "rs": rs,
+            "red": lds.red.ptr,
+            "tab": lds.tab.ptr,
+            "args": {"bias": bias},
+            "put": mb.put,
+            "poll": mb.poll,
+            "logit": Addr("logit", Space.SCRATCH, logit),
+            "route": Addr("route", Space.SCRATCH, route),
+        }
+        for t in range_constexpr(tokens):
+            M.stage_route(c, t)
+        gpu.barrier()
+        M.load_route(c, lds.tab.ptr)
+        gpu.barrier()
+        if tid < n:
+            t = tid // topk
+            eid = fx.ptr_load(lds.tab.ptr + (t * 2 * topk + tid % topk))
+            _st(eid, ids, tid)
+            pick = rs.pick_at(lds.tab.ptr, tid)
+            at = x + (fx.Int64(pick) // topk) * 4
+            ptr = fx.inttoptr(fx.PointerType.get(T.f32, fx.AddressSpace.Global, 4), at)
+            # Kept: an unread load is deleted, and so is the fault it raises.
+            _st(fx.ptr_load(ptr).bitcast(fx.Int32), ids, n + tid)
+
+    @flyc.jit
+    def launch(logit: Int64, bias: Int64, route: Int64, x: Int64, ids: Int64):
+        kern(logit, bias, route, x, ids).launch(grid=(1,), block=(THREADS,))
+
+    return launch
+
+
+def test_nan_draft_route_stays_in_range():
+    """10 draft tokens, 128 experts, top-3, every logit NaN.
+
+    This is the gfx942 multi-tile shape from the capture fault. Each token's
+    three ids are unique and in range, and the ordered-pick load finishes."""
+    from vllm.models.deepseek_v41.amd.mono.stages.moe_shape import RouteShape
+
+    tokens, experts, topk = 10, 128, 3
+    rs = RouteShape(experts, topk, tokens, 4)
+    parts = M.ROUTER_PARTS
+    n_pairs = tokens * experts * parts
+    nan_bits = torch.tensor(float("nan")).view(torch.int32)
+    logit = torch.empty(n_pairs * 2, device="cuda", dtype=torch.int32)
+    logit[0::2] = nan_bits
+    logit[1::2] = 1
+    route = torch.zeros(
+        M.ROUTE_COPIES * tokens * 2 * topk * 2, device="cuda", dtype=torch.int32
+    )
+    bias = torch.zeros(experts, device="cuda", dtype=torch.float32)
+    x = torch.zeros(tokens, device="cuda", dtype=torch.float32)
+    ids = torch.full((2 * tokens * topk,), -1, device="cuda", dtype=torch.int32)
+    _build_nan_route(rs)(
+        logit.data_ptr(),
+        bias.data_ptr(),
+        route.data_ptr(),
+        x.data_ptr(),
+        ids.data_ptr(),
+    )
+    torch.accelerator.synchronize()
+    got = ids[: tokens * topk].view(tokens, topk).cpu()
+    for row in got.tolist():
+        assert all(0 <= e < experts for e in row)
+        assert len(set(row)) == topk
diff --git a/tests/kernels/moe/conftest.py b/tests/kernels/moe/conftest.py
index af99dc756b96..db1b089f5258 100644
--- a/tests/kernels/moe/conftest.py
+++ b/tests/kernels/moe/conftest.py
@@ -54,10 +54,22 @@ def should_do_global_cleanup_after_test() -> bool:
     return False
 
 
+_last_path = None
+
+
 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 or "dist_init" in item.fixturenames:
+    global _last_path
+    # Runtime sharding reorders files between builds, so every file also
+    # starts clean: no file inherits the GPU memory another file left cached.
+    new_file = item.path != _last_path
+    _last_path = item.path
+    if (
+        new_file
+        or 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 1673907d278b..4b26efce70a4 100644
--- a/tests/kernels/moe/modular_kernel_tools/common.py
+++ b/tests/kernels/moe/modular_kernel_tools/common.py
@@ -106,6 +106,11 @@ class Config:
     hidden_dim_unpadded: int | None = None
     intermediate_size_per_partition_unpadded: int | None = None
 
+    # Opt into AITER's DeepSeek-V4.1 a4w4 (DSV4) kernel path (SEPARATED
+    # gate/up layout instead of INTERLEAVE); see
+    # VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4 in config.py.
+    use_mxfp4_w4a4_dsv4: bool = False
+
     def __post_init__(self):
         if self.quant_config is None:
             self.quant_config = TestMoEQuantConfig(None, False, False, None)
@@ -233,7 +238,18 @@ def make_env_data(self) -> tuple[VllmConfig, dict[Any, Any]]:
             "VLLM_USE_DEEP_GEMM": str(int(self.needs_deep_gemm())),
         }
 
+        if self.use_mxfp4_w4a4_dsv4:
+            # model_type must pass FusedMoEConfig.__post_init__'s validation
+            # gate for this env var.
+            vllm_config.model_config.hf_config = SimpleNamespace(
+                model_type="deepseek_v41"
+            )
+            env_dict["VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4"] = "1"
+
         vllm_config.parallel_config.all2all_backend = self.all2all_backend()
+        # "mori" labels the MORI EPv2 prepare/finalize here.
+        if self.all2all_backend() == "mori":
+            vllm_config.parallel_config.mori_ep_v2 = True
 
         return vllm_config, env_dict
 
@@ -353,7 +369,11 @@ def needs_aiter(self):
 
     def needs_mori(self):
         info = prepare_finalize_info(self.prepare_finalize_type)
-        return info.backend in ("mori_high_throughput", "mori_low_latency")
+        return info.backend in (
+            "mori_high_throughput",
+            "mori_low_latency",
+            "mori",
+        )
 
     def all2all_backend(self):
         info = prepare_finalize_info(self.prepare_finalize_type)
@@ -410,6 +430,12 @@ def is_valid(self) -> tuple[bool, str | None]:
                 f"block={self.quant_block_shape})"
             )
 
+        # use_mxfp4_w4a4_dsv4 is only reached through the W4A16 mxfp4 quant
+        # key pair (see mxfp4_quant_key_pair); catch it desyncing from
+        # is_w4a16_mxfp4 here instead of failing deep in weight conversion.
+        if self.use_mxfp4_w4a4_dsv4 and not self.is_w4a16_mxfp4:
+            return False, "use_mxfp4_w4a4_dsv4 requires the W4A16 mxfp4 quant scheme."
+
         # Check activation support; NotImplementedError means no opinion.
         try:
             if not self.fused_experts_type._supports_activation(self.activation):
@@ -716,17 +742,22 @@ def reference_moe_impl(
         dtype = config.dtype
         a = rank_tensors.hidden_states
         w1_q, w1_scale_q = weights.w1, weights.w1_scale
-        if config.is_w4a16_mxfp4:
-            assert config.activation in (
-                MoEActivation.SWIGLUOAI,
-                MoEActivation.SWIGLUOAI_UNINTERLEAVE,
-            )
-            # rocm_aiter_moe.py forces gate_mode=INTERLEAVE and the same
-            # "swiglu" activation_method for any W4A16 mxfp4 call regardless
-            # of activation, so SWIGLUOAI_UNINTERLEAVE's kernel math here is
-            # identical to SWIGLUOAI's. Reference must use the matching
-            # interleaved-reading op (SwigluOAIAndMul), not
-            # SiluAndMulWithClamp's packed-halves reading.
+        if not config.use_mxfp4_w4a4_dsv4 and config.activation in (
+            MoEActivation.SWIGLUOAI,
+            MoEActivation.SWIGLUOAI_UNINTERLEAVE,
+        ):
+            # rocm_aiter_moe.py checks use_mxfp4_w4a4_dsv4 before activation:
+            # DSV4 always forces gate_mode=SEPARATED, and only when DSV4 is
+            # off does gate_mode depend on activation (INTERLEAVE, same for
+            # both SwigluOAI variants since rocm_aiter_moe.py forces the same
+            # "swiglu" activation_method for either one). SwigluOAIAndMul
+            # expects gate/up interleaved row order (x[..., ::2], x[...,
+            # 1::2]), not the contiguous [gate; up] halves that `weights.w1`
+            # is generated in -- mirror the same reorder
+            # `_maybe_convert_weights_for_experts()` applies before the real
+            # MK-path AITER kernel call. DSV4's SEPARATED layout (default
+            # SiluAndMul activation) keeps the contiguous halves as-is and
+            # needs no reorder.
             w1_q = _interleave_gate_up_rows(w1_q)
             w1_scale_q = _interleave_gate_up_rows(w1_scale_q)
             reference_activation = MoEActivation.SWIGLUOAI
@@ -887,6 +918,27 @@ def _maybe_convert_weights_for_experts(
         # The converter mutates its w13/w2 weight and scale args in place, and
         # rank_weights is a view into the shared master WeightTensors, so
         # everything passed in must be a fresh tensor.
+        if config.use_mxfp4_w4a4_dsv4:
+            # DSV4's SEPARATED layout needs the real production converter
+            # (not the gpt-oss one below), on un-interleaved weights.
+            from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
+                convert_weight_to_mxfp4_moe_kernel_format,
+            )
+
+            w1, w2, w1_scale, w2_scale, _, _ = (
+                convert_weight_to_mxfp4_moe_kernel_format(
+                    mxfp4_backend=config.mxfp4_backend,
+                    layer=torch.nn.Module(),
+                    w13_weight=rank_weights.w1.clone(),
+                    w2_weight=rank_weights.w2.clone(),
+                    w13_weight_scale=rank_weights.w1_scale.clone(),
+                    w2_weight_scale=rank_weights.w2_scale.clone(),
+                    activation=config.activation,
+                    use_separated_a4w4=True,
+                )
+            )
+            return WeightTensors(w1=w1, w2=w2, w1_scale=w1_scale, w2_scale=w2_scale)
+
         if config.is_w4a16_mxfp4:
             w1 = _interleave_gate_up_rows(rank_weights.w1)
             w1_scale = _interleave_gate_up_rows(rank_weights.w1_scale)
diff --git a/tests/kernels/moe/modular_kernel_tools/mk_objects.py b/tests/kernels/moe/modular_kernel_tools/mk_objects.py
index 0a3071951213..74e773da6c6d 100644
--- a/tests/kernels/moe/modular_kernel_tools/mk_objects.py
+++ b/tests/kernels/moe/modular_kernel_tools/mk_objects.py
@@ -44,6 +44,7 @@
     has_deep_ep_v2,
     has_deep_gemm,
     has_mori,
+    mori_epv2_unavailable_reason,
 )
 
 
@@ -247,6 +248,7 @@ def expert_info(kind) -> ExpertInfo:
     from vllm.model_executor.layers.fused_moe.prepare_finalize.mori import (
         MoriPrepareAndFinalize,
     )
+    from vllm.platforms.rocm import get_cdna_version
 
     register_prepare_and_finalize(
         MoriPrepareAndFinalize,
@@ -257,6 +259,22 @@ def expert_info(kind) -> ExpertInfo:
         supports_apply_weight_on_input=False,
     )
 
+    # Config.make_env_data pins mori_ep_v2=True, which skips the EPv1 fallback
+    # in RocmPlatform.check_and_update_config, so apply its EPv2 gate here.
+    if get_cdna_version() >= 4 and mori_epv2_unavailable_reason() is None:
+        from vllm.model_executor.layers.fused_moe.prepare_finalize.mori_epv2 import (
+            MoriEPv2PrepareAndFinalize,
+        )
+
+        register_prepare_and_finalize(
+            MoriEPv2PrepareAndFinalize,
+            standard_format,
+            [torch.bfloat16],
+            blocked_quantization_support=False,
+            backend="mori",
+            supports_apply_weight_on_input=False,
+        )
+
 if has_flashinfer_cutlass_fused_moe() and current_platform.has_device_capability(100):
     from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutlass_moe import (
         FlashInferExperts,
diff --git a/tests/kernels/moe/test_block_fp8.py b/tests/kernels/moe/test_block_fp8.py
index 77b638a76427..0a4fd90f093a 100644
--- a/tests/kernels/moe/test_block_fp8.py
+++ b/tests/kernels/moe/test_block_fp8.py
@@ -17,6 +17,7 @@
 from vllm.model_executor.layers.fused_moe import (
     fused_experts,
     fused_topk,
+    override_config,
 )
 from vllm.model_executor.layers.fused_moe.activation import MoEActivation
 from vllm.model_executor.layers.fused_moe.all2all_utils import (
@@ -32,6 +33,9 @@
 from vllm.model_executor.layers.fused_moe.experts.triton_deep_gemm_moe import (
     TritonOrDeepGemmExperts,
 )
+from vllm.model_executor.layers.fused_moe.fused_moe import (
+    try_get_optimal_moe_config,
+)
 from vllm.model_executor.layers.quantization.utils.fp8_utils import (
     per_token_group_quant_fp8,
 )
@@ -260,6 +264,57 @@ def test_w8a8_block_fp8_fused_moe(
     torch.testing.assert_close(m_out, ref_out_m, atol=tol, rtol=tol)
 
 
+@pytest.mark.parametrize("M", [24, 1024])
+@pytest.mark.parametrize("modular", [False, True])
+@torch.inference_mode()
+def test_h200_block_fp8_up_and_down_tma_matches_pointer(M, modular, workspace_init):
+    if "H200" not in torch.cuda.get_device_name():
+        pytest.skip("This block FP8 TMA test targets H200")
+
+    E, N, K, topk = 256, 256, 128, 8
+    torch.manual_seed(123)
+    a = torch.randn((M, K), dtype=torch.bfloat16) / 10
+    score = torch.randn((M, E), dtype=torch.bfloat16)
+    w1, w2, quant_config = make_test_quant_config(
+        E, N, K, torch.bfloat16, torch.float8_e4m3fn, block_shape=[128, 128]
+    )
+    topk_weights, topk_ids, _ = fused_topk(a, score.float(), topk, False)
+
+    config_args = (w1.shape, w2.shape, topk, "fp8_w8a8", M)
+    config = try_get_optimal_moe_config(*config_args, block_shape=[128, 128])
+    assert w2.is_contiguous()
+
+    if modular:
+        moe = modular_triton_fused_moe(make_dummy_moe_config(), quant_config)
+
+        def run():
+            return 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,
+            )
+
+    else:
+
+        def run():
+            return fused_experts(
+                a, w1, w2, topk_weights, topk_ids, quant_config=quant_config
+            )
+
+    with set_current_vllm_config(vllm_config):
+        with override_config({**config, "USE_TMA": False}):
+            pointer_output = run()
+        tma_output = run()
+
+    torch.testing.assert_close(tma_output, pointer_output, atol=0, rtol=0)
+
+
 @pytest.mark.parametrize("clamp_limit", [None, 10.0])
 @torch.inference_mode()
 def test_w8a8_block_fp8_clamp_skips_fused_silu_quant(
diff --git a/tests/kernels/moe/test_deep_gemm_mega_moe.py b/tests/kernels/moe/test_deep_gemm_mega_moe.py
index 0f707e16ea9f..cc54b6f8947c 100644
--- a/tests/kernels/moe/test_deep_gemm_mega_moe.py
+++ b/tests/kernels/moe/test_deep_gemm_mega_moe.py
@@ -63,12 +63,26 @@ def _make_experts(**kwargs) -> DeepGemmMegaMoEExperts:
 
 
 def test_deep_gemm_mega_moe_capture_precedes_eplb(monkeypatch):
+    """Routed-experts capture sees logical ids; staging maps them with EPLB."""
     experts = DeepGemmMegaMoEExperts.__new__(DeepGemmMegaMoEExperts)
     torch.nn.Module.__init__(experts)
     monkeypatch.setattr(experts, "_synchronize_first_launch", lambda: None)
     experts.prefix = "model.layers.3.ffn.experts"
     experts.max_num_tokens = 4
     experts.num_shared_experts = 0
+    experts.shared_spec = None
+    experts.weight_format = MEGA_MOE_MXFP4
+    experts._kernel_kwargs = {}
+    monkeypatch.setattr(
+        experts,
+        "get_symm_buffer",
+        lambda: SimpleNamespace(
+            x=torch.empty(4, 8),
+            x_sf=torch.empty(4, 1),
+            topk_idx=torch.empty(4, 2),
+            topk_weights=torch.empty(4, 2),
+        ),
+    )
     experts.capture_fn = None
     experts.eplb_state = SimpleNamespace(
         logical_to_physical_map=torch.empty(1),
@@ -85,19 +99,20 @@ def test_deep_gemm_mega_moe_capture_precedes_eplb(monkeypatch):
         SimpleNamespace(capture=lambda layer_id, ids: captured.append((layer_id, ids))),
     )
 
-    class MappingReached(Exception):
+    class StagingReached(Exception):
         pass
 
-    def map_ids(**kwargs):
+    def stage(hidden_states, topk_weights, ids, *args, eplb_state, **kwargs):
         assert captured == [(3, topk_ids)]
-        raise MappingReached
+        assert ids is topk_ids
+        assert eplb_state is experts.eplb_state
+        raise StagingReached
 
     monkeypatch.setattr(
-        f"{DeepGemmMegaMoEExperts.__module__}.eplb_map_to_physical_and_record",
-        map_ids,
+        f"{DeepGemmMegaMoEExperts.__module__}.prepare_megamoe_inputs", stage
     )
 
-    with pytest.raises(MappingReached):
+    with pytest.raises(StagingReached):
         experts(torch.empty(2, 8), torch.empty(2, 2), topk_ids)
 
 
@@ -688,7 +703,7 @@ def shared_experts(x):
 
 
 @requires_sm100
-@pytest.mark.parametrize("num_tokens", [7, 63, 64, 65, 127, 128, 135, 434, 16384])
+@pytest.mark.parametrize("num_tokens", [7, 63, 64, 65, 127, 128, 135, 434, 1025, 16384])
 @pytest.mark.parametrize("hidden_size", [256, 5120])
 def test_prepare_megamoe_inputs_is_bitwise_exact(num_tokens, hidden_size):
     from vllm.third_party.deep_gemm.utils import per_token_cast_to_fp8
@@ -762,6 +777,121 @@ def test_prepare_megamoe_inputs_is_bitwise_exact(num_tokens, hidden_size):
     )
 
 
+@requires_sm100
+@pytest.mark.parametrize("nvfp4", [False, True])
+@pytest.mark.parametrize("num_tokens", [7, 100])
+def test_prepare_megamoe_eplb_maps_and_records_under_graph_replay(
+    monkeypatch, nvfp4, num_tokens
+):
+    """Fused EPLB staging maps each routed id to a replica of its logical expert
+    and records unpadded load, and a replayed graph sees map and record-flag
+    updates."""
+    from vllm.distributed.eplb.eplb_state import EplbLayerState, compute_logical_maps
+
+    device = torch.device("cuda")
+    hidden_size, top_k = 256, 6
+    hidden = torch.randn(num_tokens, hidden_size, device=device, dtype=torch.bfloat16)
+    ids = torch.arange(num_tokens * top_k, device=device).reshape(-1, top_k) % 8
+    ids = ids.int()
+    ids[1, 0], ids[2, 0] = -1, 8
+    weights = torch.rand(num_tokens, top_k, device=device)
+    padding = torch.arange(num_tokens, device=device) % 11 == 0
+    replicas = torch.tensor([1, 2, 3, 1, 1, 1, 1, 1])
+    num_physical = int(replicas.sum())
+    mapping, counts = compute_logical_maps(
+        torch.arange(8).repeat_interleave(replicas), 8
+    )
+    # Spare slots make the map stride differ from the max replica count.
+    mapping = F.pad(mapping, (0, 2), value=-1)
+    num_unpadded = [
+        torch.tensor(0, dtype=torch.int32, device=device),
+        torch.tensor(num_tokens - 3, dtype=torch.int32, device=device),
+    ]
+    state = EplbLayerState(
+        logical_to_physical_map=mapping.to(device),
+        logical_replica_count=counts.to(device),
+        expert_load_view=torch.zeros(num_physical, dtype=torch.int32, device=device),
+        should_record_tensor=torch.tensor(True, device=device),
+        num_unpadded_tokens_tensors=num_unpadded,
+    )
+    # Staging must read this ubatch's count; ubatch 0 would record nothing.
+    monkeypatch.setattr(
+        "vllm.model_executor.layers.fused_moe.prepare_megamoe.dbo_current_ubatch_id",
+        lambda: 1,
+    )
+    topk_idx = torch.empty(num_tokens, top_k, dtype=torch.int64, device=device)
+    topk_weights = torch.empty_like(weights)
+    if nvfp4:
+        x = torch.empty(num_tokens, hidden_size // 2, dtype=torch.int8, device=device)
+        x_sf = torch.empty(
+            num_tokens, hidden_size // 64, dtype=torch.int32, device=device
+        )
+        x_scales = torch.empty(num_tokens, device=device)
+
+        def stage():
+            prepare_megamoe_nvfp4_inputs(
+                hidden,
+                weights,
+                ids,
+                x,
+                x_sf,
+                x_scales,
+                topk_idx,
+                topk_weights,
+                is_padding=padding,
+                eplb_state=state,
+            )
+    else:
+        x = torch.empty(
+            num_tokens, hidden_size, dtype=torch.float8_e4m3fn, device=device
+        )
+        x_sf = torch.empty(
+            num_tokens, hidden_size // 128, dtype=torch.int32, device=device
+        )
+
+        def stage():
+            prepare_megamoe_inputs(
+                hidden,
+                weights,
+                ids,
+                x,
+                x_sf,
+                topk_idx,
+                topk_weights,
+                is_padding=padding,
+                eplb_state=state,
+            )
+
+    stage()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        stage()
+    physical = state.logical_to_physical_map
+    record_flag, load = state.should_record_tensor, state.expert_load_view
+    assert physical is not None and record_flag is not None and load is not None
+    logical = torch.where(padding[:, None], -1, ids).long().cpu()
+    valid = (logical >= 0) & (logical < 8)
+    unpadded = torch.arange(num_tokens)[:, None] < num_tokens - 3
+    for record in (False, True):
+        physical.copy_(
+            torch.where(physical >= 0, num_physical - 1 - physical, physical)
+        )
+        record_flag.fill_(record)
+        load.zero_()
+        graph.replay()
+        physical_to_logical = torch.full((num_physical,), -1, dtype=torch.long)
+        for expert, slots in enumerate(physical.cpu()):
+            physical_to_logical[slots[slots >= 0]] = expert
+        mapped = topk_idx.cpu()
+        assert (mapped[~valid] == -1).all()
+        assert torch.equal(physical_to_logical[mapped[valid]], logical[valid])
+        recorded = mapped[valid & unpadded] if record else mapped.new_empty(0)
+        expected_load = torch.bincount(recorded, minlength=num_physical)
+        assert torch.equal(load.cpu(), expected_load.to(load.dtype))
+    # Tokens spread over the replicas of a replicated expert.
+    assert mapped[logical == 2].unique().numel() > 1
+
+
 def _shared_sf_tma_layout(sf: torch.Tensor, block_m: int) -> torch.Tensor:
     """Scatter per-token scales into DeepGEMM's padded MN-major UTCCP pages."""
     aligned_block_m = (block_m + 127) // 128 * 128
@@ -940,7 +1070,7 @@ def test_prepare_megamoe_inputs_masks_padding(nonfinite_padding, num_tokens, qua
 
 @requires_sm100
 @pytest.mark.parametrize("shared", [None, "mxfp8", "bf16"])
-@pytest.mark.parametrize("num_tokens", [0, 7, 65, 135])
+@pytest.mark.parametrize("num_tokens", [0, 7, 65, 135, 600])
 def test_prepare_megamoe_nvfp4_inputs_is_bitwise_exact(num_tokens, shared):
     """Per-token NVFP4 dispatch, shared-expert inputs, unrouted padding and the
     L2 alphas match DeepGEMM's reference quantizers bit for bit."""
diff --git a/tests/kernels/moe/test_fp8_tp_loading.py b/tests/kernels/moe/test_fp8_tp_loading.py
index 680ccb812cca..0ac056919693 100644
--- a/tests/kernels/moe/test_fp8_tp_loading.py
+++ b/tests/kernels/moe/test_fp8_tp_loading.py
@@ -68,6 +68,7 @@ def _make_fp8_tp_experts(
         or Fp8Config(is_checkpoint_fp8_serialized=True, weight_block_size=[128, 128]),
         expert_map_manager=SimpleNamespace(
             local_num_experts=num_experts,
+            num_fused_shared_experts=0,
             placement_strategy="linear",
             expert_map=None,
             expert_mask=None,
diff --git a/tests/kernels/moe/test_grouped_topk.py b/tests/kernels/moe/test_grouped_topk.py
index 2004feae6595..05dcf044bcfc 100644
--- a/tests/kernels/moe/test_grouped_topk.py
+++ b/tests/kernels/moe/test_grouped_topk.py
@@ -76,7 +76,13 @@ def _single_group_reference(
     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):
+@pytest.mark.parametrize("use_fused", [False, True])
+def test_kimi_grouped_topk_masks_padding_when_enabled(
+    skip_padding: bool, use_fused: bool, monkeypatch
+):
+    set_random_seed(0)
+    monkeypatch.setenv("VLLM_USE_FUSED_MOE_GROUPED_TOPK", str(int(use_fused)))
+    monkeypatch.setenv("VLLM_MOE_SKIP_PADDING", "1")
     hidden_states = torch.randn((2, 16), device="cuda")
     gating_output = torch.randn((2, 896), device="cuda")
     is_padding = torch.tensor([False, True], device="cuda")
@@ -86,14 +92,19 @@ def test_kimi_grouped_topk_masks_padding_when_enabled(skip_padding: bool):
         num_expert_group=1,
         topk_group=1,
         skip_padding=skip_padding,
+        e_score_correction_bias=torch.randn((896,), device="cuda"),
     )
 
+    baseline_weights, baseline_ids = router._compute_routing(
+        hidden_states, gating_output, torch.int32
+    )
     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)
+    torch.testing.assert_close(topk_weights[0], baseline_weights[0], rtol=0, atol=0)
+    assert torch.equal(topk_ids[0], baseline_ids[0])
     if skip_padding:
         assert torch.all(topk_ids[1] == -1)
         assert torch.all(topk_weights[1] == 0)
@@ -102,6 +113,80 @@ def test_kimi_grouped_topk_masks_padding_when_enabled(skip_padding: bool):
         assert torch.all(topk_weights[1] > 0)
 
 
+@pytest.mark.skipif(
+    not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform."
+)
+@pytest.mark.parametrize(
+    "num_experts,num_groups,topk_group,topk",
+    [(128, 1, 1, 8), (896, 1, 1, 16), (256, 8, 4, 8), (512, 8, 4, 8)],
+)
+@pytest.mark.parametrize("scoring_func", ["softmax", "sigmoid"])
+def test_grouped_topk_padding_changes_on_graph_replay(
+    num_experts: int, num_groups: int, topk_group: int, topk: int, scoring_func: str
+):
+    """Every kernel family must read the current mask on each graph replay."""
+    set_random_seed(0)
+    logits = torch.randn((33, num_experts), device="cuda")
+    bias = torch.randn((num_experts,), device="cuda")
+    hidden = torch.empty((33, 0), device="cuda")
+    padding = torch.zeros(33, dtype=torch.bool, device="cuda")
+
+    def route(mask):
+        return fused_grouped_topk(
+            hidden,
+            logits,
+            topk,
+            True,
+            bias,
+            num_groups,
+            topk_group,
+            scoring_func,
+            2.5,
+            is_padding=mask,
+        )
+
+    baseline_weights, baseline_ids = route(None)
+    stream = torch.cuda.Stream()
+    stream.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(stream):
+        route(padding)
+    torch.cuda.current_stream().wait_stream(stream)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph, stream=stream):
+        weights, ids = route(padding)
+    for pattern in (
+        torch.arange(33, device="cuda") % 2 == 0,
+        torch.ones_like(padding),
+        torch.zeros_like(padding),
+    ):
+        padding.copy_(pattern)
+        graph.replay()
+        torch.testing.assert_close(
+            weights, baseline_weights.masked_fill(pattern[:, None], 0), rtol=0, atol=0
+        )
+        assert torch.equal(ids, baseline_ids.masked_fill(pattern[:, None], -1))
+
+
+@pytest.mark.skipif(
+    not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform."
+)
+def test_grouped_topk_padding_opcheck():
+    torch.library.opcheck(
+        torch.ops._moe_C.grouped_topk,
+        (
+            torch.randn((4, 256), device="cuda"),
+            8,
+            4,
+            8,
+            True,
+            2.5,
+            torch.randn(256, device="cuda"),
+            1,
+            torch.tensor([False, True, False, True], device="cuda"),
+        ),
+    )
+
+
 @pytest.mark.skipif(
     not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform."
 )
diff --git a/tests/kernels/moe/test_modular_kernel_combinations.py b/tests/kernels/moe/test_modular_kernel_combinations.py
index 61e95aca6669..dbae1e50e084 100644
--- a/tests/kernels/moe/test_modular_kernel_combinations.py
+++ b/tests/kernels/moe/test_modular_kernel_combinations.py
@@ -21,6 +21,7 @@
     AiterFp8BlockExperts,
     AiterFp8ChannelExperts,
     AiterFp8TensorExperts,
+    AiterMxfp4Experts,
     AiterUnquantizedExperts,
 )
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
@@ -518,8 +519,8 @@ def test_modular_kernel_combinations_singlegpu(
 # read once at import time, so each combo needs a fresh child process (see
 # parallel_launch_with_config) rather than in-process monkeypatching.
 #
-# AITER_MOE_SORT_BACKEND is not covered: it only matters when output_aux=True,
-# which vLLM only sets on the MXFP4 path -- not yet wired into AiterExperts.
+# AITER_MOE_SORT_BACKEND is covered separately below (3b): it only matters
+# when output_aux=True, which vLLM only sets on the MXFP4 DSV4 path.
 
 
 def _aiter_sort_backend_hooks_available() -> bool:
@@ -774,6 +775,178 @@ def test_aiter_moe_dispatch_policy_forwarded_through_apply(dispatch_policy: int)
     )
 
 
+# --- 3b: AITER_MOE_SORT_BACKEND (MXFP4 adaptive-vs-Opus aux sort) ------------
+#
+# Only matters when aiter's internal `output_aux` flag is set. aiter sets it
+# unconditionally for a handful of MXFP4 a4w4 kernel families it ships,
+# including DSV4 and Kimi-K3's SiTU path (the latter is in fact the shape
+# aiter's own adaptive/Opus crossover heuristic is tuned against -- see
+# `_aux_uses_opus()`'s docstring in the installed aiter). DSV4 is just the
+# only one of these this test file can currently drive end-to-end: it's
+# wired to a vLLM flag (VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4=1,
+# FusedMoEConfig.__post_init__) and the modular-kernel harness already
+# converts weights for it. Kimi-K3's SiTU activation is real and
+# default-enabled in production (VLLM_ROCM_USE_AITER_MOE_SITUV2="auto") but
+# is blanket-marked NYI by `is_nyi_config()` below pending harness support
+# for its weight-shuffle conversion -- a separate, tracked follow-up, not
+# this PR's scope. Despite the "a4w4" name, DSV4 runs through vLLM's W4A16
+# quant pair: activations aren't pre-quantized here, AITER quantizes them to
+# 4-bit internally.
+#
+# Shape (E=384, K=5120, N=640, topk=6, M=16) is DeepSeek-V4.1-tuned
+# (dsv41_fp4_tuned_fmoe.csv) and selects block_m=16 -- the only block_m
+# where aiter's _aux_uses_opus() crossover ever prefers adaptive over Opus.
+
+MXFP4_SORT_BACKEND_ENV_CASES = ["opus", "ck", "adaptive", "auto"]
+
+
+def _aiter_mxfp4_sort_backend_worker(
+    pgi: ProcessGroupInfo,
+    vllm_config: VllmConfig,
+    cpu_group,
+    config: Config,
+    weights: WeightTensors,
+    verbose: bool,
+    backend: str,
+):
+    import aiter.fused_moe as aiter_fused_moe
+
+    # Captured before patching: a patched attribute's signature would be the
+    # mock's (*args, **kwargs), not the real one.
+    original_moe_sorting_impl = aiter_fused_moe._moe_sorting_impl
+    original_adaptive_sort = aiter_fused_moe._adaptive_moe_sort
+
+    with (
+        mock.patch.object(
+            aiter_fused_moe,
+            "_moe_sorting_impl",
+            wraps=original_moe_sorting_impl,
+        ) as sorting_impl_mock,
+        mock.patch.object(
+            aiter_fused_moe,
+            "_adaptive_moe_sort",
+            wraps=original_adaptive_sort,
+        ) as adaptive_mock,
+    ):
+        # Reuses the real correctness/accuracy checking rank_worker already
+        # does for every other AiterExperts combo, instead of duplicating it.
+        rank_worker(pgi, vllm_config, cpu_group, config, weights, verbose)
+
+    assert sorting_impl_mock.call_count > 0, (
+        "Expected AiterMxfp4Experts.apply() to reach aiter's _moe_sorting_impl."
+    )
+    call = sorting_impl_mock.call_args_list[0]
+    bound = inspect.signature(original_moe_sorting_impl).bind(*call.args, **call.kwargs)
+    bound.apply_defaults()
+    output_aux = bound.arguments["output_aux"]
+    assert output_aux, (
+        "Expected this MXFP4 DSV4 shape to set aiter's internal output_aux "
+        "flag (set automatically by aiter's _make_mxfp4_metadata for the "
+        "DSV4 kernel config) -- without it, AITER_MOE_SORT_BACKEND has no "
+        "effect, and this test can no longer exercise its target branch."
+    )
+
+    backend_forces_opus = backend in ("opus", "ck")
+    block_size = bound.arguments["block_size"]
+    num_experts = bound.arguments["num_experts"]
+    routed_rows = bound.arguments["topk_ids"].numel()
+    aux_uses_opus = aiter_fused_moe._aux_uses_opus(
+        output_aux, block_size, routed_rows, num_experts
+    )
+    # Assert (don't derive expect_adaptive from) this predicate: otherwise a
+    # tuning change that flips it would silently relax the assertion below
+    # instead of failing it.
+    assert not aux_uses_opus, (
+        f"block_size={block_size} was chosen as adaptive-eligible, but the "
+        "installed aiter build's crossover now prefers Opus for it -- pick "
+        "a new adaptive-eligible shape."
+    )
+
+    if backend_forces_opus:
+        assert adaptive_mock.call_count == 0, (
+            f"AITER_MOE_SORT_BACKEND={backend!r} unconditionally disables "
+            "the adaptive sort, but the adaptive sort fired anyway."
+        )
+    else:
+        assert adaptive_mock.call_count > 0, (
+            f"AITER_MOE_SORT_BACKEND={backend!r} was requested for this "
+            "adaptive-eligible shape, but the adaptive sort never fired "
+            "-- a silent fallback to Opus."
+        )
+
+
+@require_aiter_moe
+@pytest.mark.parametrize("backend", MXFP4_SORT_BACKEND_ENV_CASES)
+def test_aiter_moe_sort_backend_dispatch_env_matrix_mxfp4(backend: str):
+    """3b follow-up to test_aiter_moe_sorting_backend_dispatch_env_matrix: covers
+    AITER_MOE_SORT_BACKEND (#54966), left uncovered by #58393. DSV4 is the only
+    aiter output_aux path this harness can currently drive end-to-end (see the
+    module comment above for the other known path, Kimi-K3 SiTU, and why it's
+    out of scope here)."""
+    if not _aiter_sort_backend_hooks_available():
+        pytest.skip(
+            "Installed aiter build predates the sort-backend hooks "
+            "(_moe_sorting_impl, _flydsl_moe_sorting, _USE_CK_MOE_SORTING, "
+            "_USE_FLYDSL_MOE_SORTING)."
+        )
+    import aiter.fused_moe as aiter_fused_moe
+
+    if not hasattr(aiter_fused_moe, "_adaptive_moe_sort") or not hasattr(
+        aiter_fused_moe, "_aux_uses_opus"
+    ):
+        pytest.skip(
+            "Installed aiter build lacks the MXFP4 adaptive-sort hooks "
+            "(_adaptive_moe_sort, _aux_uses_opus)."
+        )
+
+    from vllm.model_executor.layers.fused_moe.prepare_finalize import (
+        MoEPrepareAndFinalizeNoDPEPModular,
+    )
+
+    # W4A16: weights are mxfp4, activations are unquantized -- the DSV4 gate
+    # in rocm_aiter_fused_experts() only fires for this quant scheme.
+    quant_config = TestMoEQuantConfig(
+        quant_dtype=None,
+        weight_dtype="mxfp4",
+        per_out_ch_quant=False,
+        per_act_token_quant=False,
+        block_shape=None,
+    )
+    config = Config(
+        Ms=[16],
+        K=5120,
+        N=640,
+        E=384,
+        topks=[6],
+        dtype=torch.bfloat16,
+        quant_config=quant_config,
+        prepare_finalize_type=MoEPrepareAndFinalizeNoDPEPModular,
+        fused_experts_type=AiterMxfp4Experts,
+        world_size=1,
+        use_mxfp4_w4a4_dsv4=True,
+    )
+    valid, reason = config.is_valid()
+    if not valid:
+        # is_valid() also covers AiterMxfp4Experts's gfx950-only check; skip
+        # (not fail) since e.g. the MI300 CI lane running this file isn't gfx950.
+        pytest.skip(reason)
+
+    weights = WeightTensors.make(config)
+    vllm_config, env_dict = config.make_env_data()
+    env_dict = {**env_dict, "AITER_MOE_SORT_BACKEND": backend}
+
+    parallel_launch_with_config(
+        config.world_size,
+        _aiter_mxfp4_sort_backend_worker,
+        vllm_config,
+        env_dict,
+        config,
+        weights,
+        False,
+        backend,
+    )
+
+
 # --- 4b: hidden_dim_unpadded/intermediate_size_per_partition_unpadded matrix -
 
 # K/N intentionally unaligned to AITER's 64/128 granularity (unlike the
@@ -791,6 +964,26 @@ def test_aiter_moe_dispatch_policy_forwarded_through_apply(dispatch_policy: int)
 #   // 64 * 64 * (2 if tp_size == 1 else 1) (rocm_aiter_moe.py); tp_size==1 here.
 _PADDING_INTERMEDIATE_PAD_EXPECTED = 256
 
+# AITER's W4A16 stage2 kernel raises NotImplementedError below a 256-wide
+# physical intermediate (128 isn't enough), so MXFP4 needs its own N pair.
+_PADDING_N_UNPADDED_MXFP4 = 512  # 2 * 256
+_PADDING_N_PADDED_MXFP4 = 768  # +256
+_PADDING_INTERMEDIATE_PAD_EXPECTED_MXFP4_W4A16 = 512
+
+# Both MXFP4 schemes also need a 256-wide physical hidden dim: W4A4's stage1
+# GEMM silently produces NaN/Inf below it (confirmed: K=896/1152 -> NaN,
+# K=768/1024 -> finite); W4A16 silently produces wrong output instead (K=896
+# fails the cos_sim check below). Production enforces the same 256 floor via
+# mxfp4_round_up_hidden_size_and_intermediate_size() before building weights,
+# so this isn't just a test-only requirement.
+_PADDING_K_UNPADDED_MXFP4 = 1024  # 4 * 256
+_PADDING_K_PADDED_MXFP4 = 1280  # +256
+# hidden_pad = (K_padded_mxfp4 - hidden_dim_unpadded) // 128 * 128.
+_PADDING_HIDDEN_PAD_EXPECTED_MXFP4 = 256
+# W4A4's SiLU path doubles intermediate_pad inside AITER, so this constant
+# (undoubled) differs from the W4A16 one above despite the same N pair.
+_PADDING_INTERMEDIATE_PAD_EXPECTED_MXFP4_W4A4 = 256
+
 # {no padding, hidden-only, intermediate-only, both} -> (pad_hidden, pad_intermediate)
 _PADDING_MODES: dict[str, tuple[bool, bool]] = {
     "none": (False, False),
@@ -800,7 +993,8 @@ def test_aiter_moe_dispatch_policy_forwarded_through_apply(dispatch_policy: int)
 }
 
 # unquantized + the 4 fp8 configs AiterExperts supports (excludes MXFP4,
-# stubbed below pending MI350/gfx950, and MK_QUANT_CONFIGS[1], unsupported).
+# covered separately by _PADDING_MXFP4_QUANT_CONFIGS below, and
+# MK_QUANT_CONFIGS[1], unsupported).
 _PADDING_QUANT_CONFIGS = [
     MK_QUANT_CONFIGS[0],  # unquantized
     MK_QUANT_CONFIGS[2],  # fp8 channel weights / per-token activations
@@ -820,10 +1014,13 @@ def test_aiter_moe_dispatch_policy_forwarded_through_apply(dispatch_policy: int)
 def _padding_matrix_fe_cls(
     quant_config: TestMoEQuantConfig | None,
 ) -> type[AiterExperts]:
-    """Map a _PADDING_QUANT_CONFIGS entry to the AiterExperts leaf class that
-    supports it (mirrors Config.fp8_quant_key_pair's scheme derivation)."""
+    """Map a _PADDING_QUANT_CONFIGS/_PADDING_MXFP4_QUANT_CONFIGS entry to the
+    AiterExperts leaf class that supports it (mirrors
+    Config.fp8_quant_key_pair/mxfp4_quant_key_pair's scheme derivation)."""
     if quant_config is None:
         return AiterUnquantizedExperts
+    if quant_config.weight_dtype == "mxfp4" or quant_config.quant_dtype == "mxfp4":
+        return AiterMxfp4Experts
     if quant_config.block_shape is not None:
         return AiterFp8BlockExperts
     if quant_config.per_out_ch_quant:
@@ -894,7 +1091,9 @@ def _make_padding_matrix_weights(config: Config) -> WeightTensors:
         n=config.N,
         k=config.K,
         in_dtype=config.dtype,
-        quant_dtype=config.quant_dtype,
+        # weight_dtype, not quant_dtype: quant_dtype is None for MXFP4 W4A16
+        # (unquantized activation) even though the weight itself is mxfp4.
+        quant_dtype=config.weight_dtype,
         block_shape=config.quant_block_shape,
         per_out_ch_quant=config.is_per_out_ch_quant,
     )
@@ -915,6 +1114,36 @@ def _slice_unpadded_rank_tensors(
     )
 
 
+def _check_mxfp4_accuracy(
+    ref_out: torch.Tensor, mk_out: torch.Tensor, fe_name: str
+) -> None:
+    """Shared mxfp4 accuracy check for _aiter_padding_matrix_worker and
+    _aiter_token_padding_worker (callers pre-slice to the real region).
+
+    Cosine similarity catches a sign-flipped mk_out that the magnitude
+    guard can't (gfx950: 0.79-0.99 correct vs. <=0.01 corrupted); being
+    scale-invariant, it in turn misses a uniform-scale error like
+    rank_worker's ~5x-small AITER_MXFP4_BF16 bug, which the magnitude
+    guard catches instead.
+    """
+    cos_sim = torch.nn.functional.cosine_similarity(
+        mk_out.flatten().float(), ref_out.flatten().float(), dim=0
+    )
+    assert cos_sim > 0.5, (
+        f"{fe_name} output (cos_sim={cos_sim:.4f} vs. reference) looks "
+        "degenerate/decorrelated, not just numerically imprecise."
+    )
+    ref_scale = ref_out.abs().mean()
+    mk_scale = mk_out.abs().mean()
+    assert mk_scale > 0.5 * ref_scale, (
+        f"{fe_name} output magnitude (mean |mk_out|={mk_scale:.6f}) looks "
+        f"degenerate/zeroed vs. reference (mean |ref_out|={ref_scale:.6f})."
+    )
+    # Looser tolerance than the fp8/unquantized case: matches rank_worker's
+    # main mxfp4 sweep's atol/rtol.
+    check_accuracy(ref_out, mk_out, atol=1e-1, rtol=1e-1, percent=0.9)
+
+
 def _aiter_padding_matrix_worker(
     pgi: ProcessGroupInfo,
     vllm_config: VllmConfig,
@@ -941,54 +1170,77 @@ def _aiter_padding_matrix_worker(
             pgi, vllm_config, padded_config, weights, rank_tensors
         )
 
+    fe_name = padded_config.fused_experts_type.__name__
     assert fused_moe_mock.call_count > 0, (
-        "Expected AiterExperts.apply() to call rocm_aiter_ops.fused_moe."
+        f"Expected {fe_name}.apply() to call rocm_aiter_ops.fused_moe."
     )
     for call in fused_moe_mock.call_args_list:
         assert call.kwargs["hidden_pad"] == hidden_pad_expected, (
-            f"AiterExperts.apply() forwarded hidden_pad="
+            f"{fe_name}.apply() forwarded hidden_pad="
             f"{call.kwargs['hidden_pad']}, expected {hidden_pad_expected}."
         )
         assert call.kwargs["intermediate_pad"] == intermediate_pad_expected, (
-            f"AiterExperts.apply() forwarded intermediate_pad="
+            f"{fe_name}.apply() forwarded intermediate_pad="
             f"{call.kwargs['intermediate_pad']}, expected "
             f"{intermediate_pad_expected}."
         )
 
+    k_unpadded = padded_config.hidden_dim_unpadded or padded_config.K
+    n_unpadded = (
+        padded_config.intermediate_size_per_partition_unpadded or padded_config.N
+    )
+
     # AiterExperts.apply() never slices its own output -- the buffer stays
     # raw (padded) width; only the caller-sliced unpadded prefix is checked.
     assert mk_out.shape[-1] == padded_config.K, (
-        f"AiterExperts output width {mk_out.shape[-1]} != raw hidden_dim "
+        f"{fe_name} output width {mk_out.shape[-1]} != raw hidden_dim "
         f"{padded_config.K}."
     )
-    mk_out = mk_out[..., :_PADDING_K_UNPADDED]
+    mk_out = mk_out[..., :k_unpadded]
     # Only the real slice is checked: beyond it is unspecified/reused memory.
     assert torch.isfinite(mk_out).all(), (
-        "AiterExperts output must not contain NaN/Inf in the real "
+        f"{fe_name} output must not contain NaN/Inf in the real "
         "hidden_dim_unpadded-wide output slice."
     )
 
-    unpadded_weights = _slice_unpadded_weights(
-        weights,
-        padded_config.quant_config,
-        _PADDING_K_UNPADDED,
-        _PADDING_N_UNPADDED,
-    )
-    unpadded_rank_tensors = _slice_unpadded_rank_tensors(
-        rank_tensors, _PADDING_K_UNPADDED
-    )
+    is_mxfp4 = padded_config.weight_dtype == "mxfp4"
+    if is_mxfp4:
+        unpadded_weights = _slice_unpadded_mxfp4_weights(
+            weights, k_unpadded, n_unpadded
+        )
+    else:
+        unpadded_weights = _slice_unpadded_weights(
+            weights, padded_config.quant_config, k_unpadded, n_unpadded
+        )
+    unpadded_rank_tensors = _slice_unpadded_rank_tensors(rank_tensors, k_unpadded)
 
     with set_current_vllm_config(vllm_config):
         ref_out = reference_moe_impl(
             padded_config, unpadded_weights, unpadded_rank_tensors
         )
 
+    if (
+        padded_config.is_w4a16_mxfp4
+        and padded_config.activation == MoEActivation.SWIGLUOAI
+    ):
+        # AITER_MXFP4_BF16 silently miscomputes SWIGLUOAI; smoke test only
+        # until #59766 (see rank_worker). SWIGLUOAI_UNINTERLEAVE isn't
+        # exempted -- same kernel path, gets the full check below instead.
+        assert mk_out.shape == ref_out.shape, (
+            f"{fe_name} output shape {mk_out.shape} != reference shape {ref_out.shape}."
+        )
+        return
+
+    if is_mxfp4:
+        _check_mxfp4_accuracy(ref_out, mk_out, fe_name)
+        return
+
     # ref_out's magnitude here (~1e-3-1e-2) is below atol=3e-2, so a zeroed-out
     # mk_out would still pass check_accuracy below -- guard against that.
     ref_scale = ref_out.abs().mean()
     mk_scale = mk_out.abs().mean()
     assert mk_scale > 0.5 * ref_scale, (
-        f"AiterExperts output magnitude (mean |mk_out|={mk_scale:.6f}) looks "
+        f"{fe_name} output magnitude (mean |mk_out|={mk_scale:.6f}) looks "
         f"degenerate/zeroed vs. reference (mean |ref_out|={ref_scale:.6f})."
     )
 
@@ -1092,13 +1344,121 @@ def test_aiter_moe_padding_matrix(mode: str, quant_config: TestMoEQuantConfig |
     )
 
 
-@pytest.mark.skip(
-    reason="MXFP4 AiterExperts padding requires MI350/gfx950 hardware. "
-    'See https://github.com/vllm-project/vllm/issues/54966 ("Test padding").'
+# mxfp4 weight_dtype/quant_dtype entries conditionally appended to
+# MK_QUANT_CONFIGS when has_aiter() (see mk_objects.py): W4A16 (mxfp4 weight,
+# unquantized activation) and W4A4 (both mxfp4).
+_PADDING_MXFP4_QUANT_CONFIGS = [
+    c
+    for c in MK_QUANT_CONFIGS
+    if c is not None and (c.weight_dtype == "mxfp4" or c.quant_dtype == "mxfp4")
+]
+_PADDING_MXFP4_QUANT_IDS = [
+    "mxfp4_w4a4" if c.quant_dtype == "mxfp4" else "mxfp4_w4a16"
+    for c in _PADDING_MXFP4_QUANT_CONFIGS
+]
+
+
+def _slice_unpadded_mxfp4_weights(
+    weights: WeightTensors, k_unpadded: int, n_unpadded: int
+) -> WeightTensors:
+    """MXFP4 counterpart to _slice_unpadded_weights: slices the packed
+    uint8/block-scale tensors directly so the result can feed
+    reference_moe_impl like the fp8 path does.
+
+    downcast_to_mxfp() packs 2 elements/byte along the quantized axis with
+    32-wide scale blocks (MXFP_BLOCK_SIZE); this file's MXFP4 padding
+    constants are multiples of both (lcm=64), so every slice lands on a
+    byte/scale-block boundary. Gate/up interleave is left to
+    reference_moe_impl, which expects pre-interleave row order.
+    """
+    assert weights.w1_scale is not None and weights.w2_scale is not None, (
+        "mxfp4 weights must carry block scales"
+    )
+    assert k_unpadded % 64 == 0 and n_unpadded % 64 == 0
+
+    w1 = _slice_gate_up_rows(weights.w1[:, :, : k_unpadded // 2], n_unpadded)
+    w1_scale = _slice_gate_up_rows(
+        weights.w1_scale[:, :, : k_unpadded // 32], n_unpadded
+    )
+    w2 = weights.w2[:, :k_unpadded, : n_unpadded // 2]
+    w2_scale = weights.w2_scale[:, :k_unpadded, : n_unpadded // 32]
+
+    return WeightTensors(w1=w1, w2=w2, w1_scale=w1_scale, w2_scale=w2_scale)
+
+
+@require_aiter_moe
+@pytest.mark.parametrize(
+    "quant_config", _PADDING_MXFP4_QUANT_CONFIGS, ids=_PADDING_MXFP4_QUANT_IDS
 )
-def test_aiter_moe_padding_matrix_mxfp4():
-    """Stub for the {no padding, hidden-only, intermediate-only, both} MXFP4
-    padding matrix once MI350/gfx950 hardware is available in CI."""
+@pytest.mark.parametrize("mode", list(_PADDING_MODES))
+def test_aiter_moe_padding_matrix_mxfp4(mode: str, quant_config: TestMoEQuantConfig):
+    """See https://github.com/vllm-project/vllm/issues/54966 ("Test padding").
+
+    MXFP4 counterpart to test_aiter_moe_padding_matrix: W4A16 (mxfp4 weight,
+    unquantized activation) and W4A4 (both mxfp4). W4A16 uses
+    SWIGLUOAI_UNINTERLEAVE to dodge the AITER_MXFP4_BF16 SWIGLUOAI bug
+    (#59766, see rank_worker) and get the full reference comparison like
+    W4A4 (SILU) -- only possible since #62907 taught the shared
+    reference/NYI infra to treat SWIGLUOAI_UNINTERLEAVE this way.
+    """
+    from vllm.model_executor.layers.fused_moe.prepare_finalize import (
+        MoEPrepareAndFinalizeNoDPEPModular,
+    )
+
+    pad_hidden, pad_intermediate = _PADDING_MODES[mode]
+    is_w4a16 = quant_config.quant_dtype is None
+
+    k = _PADDING_K_PADDED_MXFP4 if pad_hidden else _PADDING_K_UNPADDED_MXFP4
+    n = _PADDING_N_PADDED_MXFP4 if pad_intermediate else _PADDING_N_UNPADDED_MXFP4
+
+    config = Config(
+        Ms=_PADDING_M,
+        K=k,
+        N=n,
+        E=_PADDING_E,
+        topks=_PADDING_TOPK,
+        dtype=torch.bfloat16,
+        quant_config=quant_config,
+        activation=(
+            MoEActivation.SWIGLUOAI_UNINTERLEAVE if is_w4a16 else MoEActivation.SILU
+        ),
+        prepare_finalize_type=MoEPrepareAndFinalizeNoDPEPModular,
+        fused_experts_type=_padding_matrix_fe_cls(quant_config),
+        world_size=1,
+        hidden_dim_unpadded=_PADDING_K_UNPADDED_MXFP4 if pad_hidden else None,
+        intermediate_size_per_partition_unpadded=(
+            _PADDING_N_UNPADDED_MXFP4 if pad_intermediate else None
+        ),
+    )
+    if not config.fe_supports_quant_scheme():
+        pytest.skip(
+            f"AiterMxfp4Experts does not support quant scheme {quant_config} "
+            "on this GPU (gfx950-only, not gfx1250)."
+        )
+    assert config.is_valid()[0]
+
+    weights = _make_padding_matrix_weights(config)
+    vllm_config, env_dict = config.make_env_data()
+
+    hidden_pad_expected = _PADDING_HIDDEN_PAD_EXPECTED_MXFP4 if pad_hidden else 0
+    if not pad_intermediate:
+        intermediate_pad_expected = 0
+    elif is_w4a16:
+        intermediate_pad_expected = _PADDING_INTERMEDIATE_PAD_EXPECTED_MXFP4_W4A16
+    else:
+        intermediate_pad_expected = _PADDING_INTERMEDIATE_PAD_EXPECTED_MXFP4_W4A4
+
+    parallel_launch_with_config(
+        config.world_size,
+        _aiter_padding_matrix_worker,
+        vllm_config,
+        env_dict,
+        config,
+        weights,
+        False,
+        hidden_pad_expected,
+        intermediate_pad_expected,
+    )
 
 
 # --- 4c: HIP-graph token padding -- inf/nan garbage rows (topk_ids=-1) ------
@@ -1161,18 +1521,30 @@ def _aiter_token_padding_worker(
         ref_out = reference_moe_impl(config, weights, rank_tensors)
 
     real_rows = [i for i in range(config.Ms) if i not in garbage_rows]
+    fe_name = config.fused_experts_type.__name__
 
     assert torch.isfinite(mk_out[real_rows]).all(), (
-        "Garbage (padding) token rows leaked NaN/Inf into real rows' output "
-        "through AiterExperts."
+        f"Garbage (padding) token rows leaked NaN/Inf into real rows' output "
+        f"through {fe_name}."
     )
 
+    if config.is_w4a16_mxfp4 and config.activation == MoEActivation.SWIGLUOAI:
+        # AITER_MXFP4_BF16 silently miscomputes SWIGLUOAI; smoke test only
+        # (finite real rows, already checked) until #59766 (see rank_worker).
+        # SWIGLUOAI_UNINTERLEAVE isn't exempted -- same kernel path, gets
+        # the full check below instead.
+        return
+
+    if config.weight_dtype == "mxfp4":
+        _check_mxfp4_accuracy(ref_out[real_rows], mk_out[real_rows], fe_name)
+        return
+
     # ref_out's magnitude here (~1e-3-1e-2) is below atol=3e-2, so a zeroed-out
     # mk_out would still pass the checks below -- guard against that.
     ref_scale = ref_out[real_rows].abs().mean()
     mk_scale = mk_out[real_rows].abs().mean()
     assert mk_scale > 0.5 * ref_scale, (
-        f"AiterExperts output magnitude (mean |mk_out|={mk_scale:.6f}) looks "
+        f"{fe_name} output magnitude (mean |mk_out|={mk_scale:.6f}) looks "
         f"degenerate/zeroed vs. reference (mean |ref_out|={ref_scale:.6f})."
     )
 
@@ -1271,12 +1643,65 @@ def test_aiter_moe_token_padding_garbage_rows(
     )
 
 
-@pytest.mark.skip(
-    reason="MXFP4 AiterExperts token padding requires MI350/gfx950 hardware. "
-    'See https://github.com/vllm-project/vllm/issues/54966 ("Test padding").'
+@require_aiter_moe
+@pytest.mark.parametrize(
+    "quant_config", _PADDING_MXFP4_QUANT_CONFIGS, ids=_PADDING_MXFP4_QUANT_IDS
 )
-def test_aiter_moe_token_padding_garbage_rows_mxfp4():
-    pass
+@pytest.mark.parametrize("garbage_mode", list(_TOKEN_PADDING_GARBAGE_MODES))
+def test_aiter_moe_token_padding_garbage_rows_mxfp4(
+    garbage_mode: str, quant_config: TestMoEQuantConfig
+):
+    """See https://github.com/vllm-project/vllm/issues/54966 ("Test padding").
+
+    MXFP4 counterpart to test_aiter_moe_token_padding_garbage_rows. Uses
+    SWIGLUOAI_UNINTERLEAVE for W4A16 for the same reason as
+    test_aiter_moe_padding_matrix_mxfp4.
+    """
+    from vllm.model_executor.layers.fused_moe.prepare_finalize import (
+        MoEPrepareAndFinalizeNoDPEPModular,
+    )
+
+    garbage_rows = _TOKEN_PADDING_GARBAGE_MODES[garbage_mode]
+    is_w4a16 = quant_config.quant_dtype is None
+
+    # Reuses test_aiter_moe_padding_matrix_mxfp4's 256-aligned K/N (required
+    # by AITER's MXFP4 kernels). Token-padding safety doesn't depend on
+    # hidden_pad/intermediate_pad, so no separate shape is needed here.
+    config = Config(
+        Ms=_PADDING_M,
+        K=_PADDING_K_PADDED_MXFP4,
+        N=_PADDING_N_PADDED_MXFP4,
+        E=_PADDING_E,
+        topks=_PADDING_TOPK,
+        dtype=torch.bfloat16,
+        quant_config=quant_config,
+        activation=(
+            MoEActivation.SWIGLUOAI_UNINTERLEAVE if is_w4a16 else MoEActivation.SILU
+        ),
+        prepare_finalize_type=MoEPrepareAndFinalizeNoDPEPModular,
+        fused_experts_type=_padding_matrix_fe_cls(quant_config),
+        world_size=1,
+    )
+    if not config.fe_supports_quant_scheme():
+        pytest.skip(
+            f"AiterMxfp4Experts does not support quant scheme {quant_config} "
+            "on this GPU (gfx950-only, not gfx1250)."
+        )
+    assert config.is_valid()[0]
+
+    weights = WeightTensors.make(config)
+    vllm_config, env_dict = config.make_env_data()
+
+    parallel_launch_with_config(
+        config.world_size,
+        _aiter_token_padding_worker,
+        vllm_config,
+        env_dict,
+        config,
+        weights,
+        False,
+        garbage_rows,
+    )
 
 
 if __name__ == "__main__":
diff --git a/tests/kernels/moe/test_moe.py b/tests/kernels/moe/test_moe.py
index d279c15c2c79..26e9f87f1014 100644
--- a/tests/kernels/moe/test_moe.py
+++ b/tests/kernels/moe/test_moe.py
@@ -75,7 +75,7 @@ def test_triton_moe_launcher_passes_scalar_scale_as_pointer(monkeypatch) -> None
     class FakeKernel:
         def __getitem__(self, grid):
             def launch(*args, **kwargs) -> None:
-                captured["a_scale"] = args[4]
+                captured["a_scale"] = args[5]
 
             return launch
 
diff --git a/tests/kernels/moe/test_moe_layer.py b/tests/kernels/moe/test_moe_layer.py
index 26533eba385e..3b6ac607a07d 100644
--- a/tests/kernels/moe/test_moe_layer.py
+++ b/tests/kernels/moe/test_moe_layer.py
@@ -112,7 +112,7 @@ def on_gfx950() -> bool:
 BACKENDS = ["allgather_reducescatter"]
 
 if has_mori():
-    BACKENDS += ["mori_high_throughput", "mori_low_latency"]
+    BACKENDS += ["mori_high_throughput", "mori_low_latency", "mori"]
 
 if has_flashinfer_nvlink_two_sided():
     BACKENDS += ["flashinfer_nvlink_two_sided"]
@@ -135,7 +135,7 @@ def on_gfx950() -> bool:
     # DeepEP forces IBGDA on, so also hide the NICs from it.
     "NVSHMEM_HCA_LIST": "none",
 }
-MORI_BACKENDS = {"mori_high_throughput", "mori_low_latency"}
+MORI_BACKENDS = {"mori_high_throughput", "mori_low_latency", "mori"}
 
 QUANT_METHODS = [
     None,
@@ -151,6 +151,8 @@ def on_gfx950() -> bool:
     "allgather_reducescatter":     {None,         "fp8", "modelopt_fp8", "modelopt_fp4"}, # noqa: E501
     "mori_high_throughput":        {None,         "fp8", "modelopt_fp8"},
     "mori_low_latency":            {None,         "fp8", "modelopt_fp8"},
+    # Pinned to EPv2, which rejects per-tensor FP8 activations.
+    "mori":                        {None, "fp8_blocked"},
     "flashinfer_nvlink_two_sided": {None, "fp8_blocked",                 "modelopt_fp4"}, # noqa: E501
     "flashinfer_nvlink_one_sided": {None, "fp8_blocked",                 "modelopt_fp4"}, # noqa: E501
     "deepep_low_latency":          {None, "fp8_blocked",                 "modelopt_fp4"}, # noqa: E501
@@ -163,6 +165,7 @@ def on_gfx950() -> bool:
     "allgather_reducescatter":     (True,  True,  True,  True),
     "mori_high_throughput":        (True, False, False,  True),
     "mori_low_latency":            (True, False, False,  True),
+    "mori":                        (True, False, False,  True),
     "flashinfer_nvlink_two_sided": (False, True, False, False),
     "flashinfer_nvlink_one_sided": (True,  True, False, False),
     "deepep_low_latency":          (True, False, False,  True),
@@ -2189,6 +2192,8 @@ def test_moe_layer(
         tensor_parallel_size=tp_size,
         enable_expert_parallel=use_ep,
         all2all_backend=backend,
+        # Pin "mori" to EPv2; mori_high_throughput/mori_low_latency cover EPv1.
+        mori_ep_v2=True if backend == "mori" else None,
         enable_eplb=enable_eplb,
         eplb_config=eplb_config,
     )
diff --git a/tests/kernels/moe/test_moe_weight_loading_padded.py b/tests/kernels/moe/test_moe_weight_loading_padded.py
index 407c71e1e0dc..189b431333d8 100644
--- a/tests/kernels/moe/test_moe_weight_loading_padded.py
+++ b/tests/kernels/moe/test_moe_weight_loading_padded.py
@@ -480,6 +480,7 @@ class TestLoadWeightsExpertBias:
     def _make_experts(self, has_bias: bool) -> torch.nn.Module:
         experts = torch.nn.Module()
         experts.layer_name = "model.layers.0.mlp.experts"
+        experts.shared_expert_online_loader = None
         experts.moe_config = make_dummy_moe_config(num_experts=self.NUM_EXPERTS)
         mapping = RoutedExperts.build_expert_params_mapping(
             "gate_proj",
@@ -573,6 +574,7 @@ def _load(
     ):
         experts = torch.nn.Module()
         experts.layer_name = layer_name
+        experts.shared_expert_online_loader = None
         experts.moe_config = SimpleNamespace(is_act_and_mul=is_gated)
         experts.get_expert_mapping = lambda **_: mapping
         experts.is_fused_checkpoint_transposed = transposed
diff --git a/tests/kernels/moe/test_mori_epv2_prepare_finalize.py b/tests/kernels/moe/test_mori_epv2_prepare_finalize.py
new file mode 100644
index 000000000000..577ac1a5da26
--- /dev/null
+++ b/tests/kernels/moe/test_mori_epv2_prepare_finalize.py
@@ -0,0 +1,471 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""MORI EPv2 prepare/finalize and all2all manager, with MORI faked out."""
+
+import sys
+from datetime import timedelta
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+
+import vllm.model_executor.layers.fused_moe.prepare_finalize.mori_epv2 as mori_epv2
+import vllm.utils.import_utils as import_utils
+from vllm.config.parallel import resolve_mori_all2all
+from vllm.model_executor.layers.fused_moe.prepare_finalize.mori_epv2 import (
+    MoriEPv2PrepareAndFinalize,
+)
+from vllm.utils.network_utils import get_open_port
+
+
+class FakeOp:
+    """MORI EPv2 op whose receive buffers have `capacity` rows, `recv` valid."""
+
+    def __init__(self, capacity=256, recv=0):
+        self.routing = object()
+        self.recv_a1 = torch.full((capacity, 4), float("nan"), dtype=torch.bfloat16)
+        self.recv_weights = torch.full((capacity, 2), float("inf"))
+        self.recv_ids = torch.full((capacity, 2), 5, dtype=torch.int32)
+        self.recv = recv
+        self.combine_in = torch.empty(capacity, 4, dtype=torch.bfloat16)
+        self.sent_ids: torch.Tensor | None = None
+        self.combine_input: torch.Tensor | None = None
+        self.combine_routing = None
+
+    def combine_in_view(self):
+        return self.combine_in
+
+    def dispatch(self, input, weights, scales, indices, *, return_routing=True):
+        self.sent_ids = indices
+        return (
+            self.recv_a1,
+            self.recv_weights,
+            None,
+            self.recv_ids,
+            torch.tensor([self.recv], dtype=torch.int32),
+            self.routing,
+        )
+
+    def combine(self, input, weights=None, indices=None, *, routing):
+        self.combine_input = input
+        self.combine_routing = routing
+        return torch.zeros(1, 4, dtype=torch.bfloat16), None
+
+
+def _make_pf(op, num_local_experts=2, rank_expert_offset=0):
+    return MoriEPv2PrepareAndFinalize(op, 16, 4, num_local_experts, rank_expert_offset)
+
+
+def _run(monkeypatch, op, num_tokens_across_dp, topk_ids, **pf_kwargs):
+    """One prepare + finalize round trip under the given DP token counts."""
+    if num_tokens_across_dp is None:
+        ctx = SimpleNamespace(dp_metadata=None)
+    else:
+        counts = torch.tensor(num_tokens_across_dp)
+        ctx = SimpleNamespace(
+            dp_metadata=SimpleNamespace(num_tokens_across_dp_cpu=counts)
+        )
+    monkeypatch.setattr(mori_epv2, "is_forward_context_available", lambda: True)
+    monkeypatch.setattr(mori_epv2, "get_forward_context", lambda: ctx)
+    monkeypatch.setattr(
+        mori_epv2, "get_dp_group", lambda: SimpleNamespace(rank_in_group=0)
+    )
+    monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False)
+    prepare_finalize = _make_pf(op, **pf_kwargs)
+    num_tokens = topk_ids.shape[0]
+    prepared = prepare_finalize.prepare(
+        torch.ones(num_tokens, 4, dtype=torch.bfloat16),
+        torch.ones(num_tokens, 2, dtype=torch.bfloat16),
+        topk_ids,
+        num_experts=8,
+        expert_map=None,
+        apply_router_weight_on_input=False,
+        quant_config=SimpleNamespace(is_per_tensor=False, a1_scale=None),
+    )
+    prepare_finalize.finalize(
+        torch.empty(num_tokens, 4, dtype=torch.bfloat16),
+        prepared[0],
+        prepared[4],
+        prepared[3],
+        False,
+        SimpleNamespace(),
+    )
+    return prepared
+
+
+@pytest.mark.parametrize(
+    "num_tokens_across_dp,recv,rows",
+    [
+        # Sum of DP batches (6) rounded up to 32 rows, of 256 capacity rows.
+        ([1, 3, 2, 0], 7, 32),
+        ([33, 0, 0, 0], 20, 64),
+        # TP=2 per DP rank: each TP rank may send its DP batch, 2 * 50 rows.
+        ([30, 20], 90, 128),
+        # Dispatchers not a multiple of DP ranks: dispatchers * largest batch.
+        ([10, 20, 5], 60, 96),
+        # Bound above capacity: keep the full buffer.
+        ([100, 100, 100, 100], 200, 256),
+        # No DP metadata, or stale metadata missing this rank's batch: full buffer.
+        (None, 7, 256),
+        ([0, 3, 2, 0], 5, 256),
+    ],
+)
+def test_receive_rows_are_bounded_by_dp_tokens(
+    monkeypatch, num_tokens_across_dp, recv, rows
+):
+    op = FakeOp(recv=recv)
+    a1, _, meta, ids, weights = _run(
+        monkeypatch, op, num_tokens_across_dp, torch.tensor([[0, 1]])
+    )
+    for t, buf in ((a1, op.recv_a1), (ids, op.recv_ids), (weights, op.recv_weights)):
+        assert t.shape[0] == rows
+        assert t.data_ptr() == buf.data_ptr()  # views, no copy
+    # Rows past the receive count are cleared; rows past the bound are untouched.
+    assert torch.isnan(a1[:recv]).all() and (a1[recv:] == 0).all()
+    assert (ids[:recv] == 5).all() and (ids[recv:] == -1).all()
+    assert (weights[recv:] == 0).all()
+    assert torch.isnan(op.recv_a1[rows:]).all()
+    assert int(meta.expert_num_tokens) == recv
+    # combine sees only the bounded expert output, under dispatch's routing.
+    assert op.combine_input is not None and op.combine_input.shape[0] == rows
+    assert op.combine_routing is op.routing
+
+
+def test_padding_tokens_are_routed_to_local_experts(monkeypatch):
+    op = FakeOp(recv=2)
+    # Rank 1 of 4 with 2 experts each owns experts 2 and 3. The second token is
+    # DP padding, which the router marks with expert id -1.
+    topk_ids = torch.tensor([[0, 7], [-1, -1]])
+    _run(monkeypatch, op, None, topk_ids, num_local_experts=2, rank_expert_offset=2)
+    assert op.sent_ids is not None and op.sent_ids.dtype == torch.int32
+    assert op.sent_ids.tolist() == [[0, 7], [2, 3]]
+
+
+@pytest.mark.parametrize("world_size,lsa_size", [(8, 8), (8, 4)])
+def test_manager_requires_one_lsa_team(monkeypatch, world_size, lsa_size):
+    from vllm.distributed.device_communicators import all2all
+    from vllm.distributed.device_communicators.base_device_communicator import (
+        All2AllManagerBase,
+    )
+
+    monkeypatch.setattr(
+        All2AllManagerBase,
+        "__init__",
+        lambda self, group: self.__dict__.update(
+            rank=0, world_size=world_size, cpu_group=group
+        ),
+    )
+    monkeypatch.setattr(
+        torch.distributed, "broadcast_object_list", lambda *a, **k: None
+    )
+    monkeypatch.setattr(all2all, "CudaRTLibrary", Mock)
+    device_comm = Mock(lsa_size=lsa_size)
+    calls = Mock()  # records the order of op.close and comm.destroy
+    comm = calls.comm
+    comm.create_dev_comm.return_value = device_comm
+    communicator = Mock()
+    communicator.init.return_value = comm
+    monkeypatch.setitem(
+        sys.modules,
+        "mori.cco",
+        SimpleNamespace(
+            Communicator=communicator,
+            CCODevCommRequirements=SimpleNamespace,
+            GDA_CONNECTION_NONE=0,
+        ),
+    )
+    if lsa_size != world_size:
+        with pytest.raises(ValueError, match="one LSA team"):
+            all2all.MoriEPv2All2AllManager(object())
+    else:
+        cfg = dict(
+            scale_dim=0,
+            scale_type_size=4,
+            num_experts_per_token=8,
+            token_nbytes=4096,
+            combine_token_nbytes=4096,
+            effective_max_recv=1024,
+            dispatch_dtype=None,
+            combine_dtype=None,
+            schedule=None,
+        )
+        monkeypatch.setitem(
+            sys.modules,
+            "mori.ops.dispatch_combine_v2",
+            SimpleNamespace(
+                EpDispatchCombineConfig=lambda **kw: SimpleNamespace(**(cfg | kw)),
+                EpDispatchCombineOp=lambda config, comm: calls.op,
+            ),
+        )
+        manager = all2all.MoriEPv2All2AllManager(object())
+        assert manager.get_handle({}) is calls.op
+        # The cache holds ops weakly; the manager keeps them alive itself.
+        manager.handle_cache._cache.clear()
+        manager.destroy()
+        # Ops live in the communicator's window: they must go first.
+        order = [c[0] for c in calls.mock_calls]
+        assert order.index("op.close") < order.index("comm.destroy")
+    device_comm.close.assert_called_once()
+    comm.destroy.assert_called_once()
+
+
+def _unique_id_worker(rank, init_method):
+    from vllm.distributed.device_communicators import all2all
+    from vllm.distributed.device_communicators.base_device_communicator import (
+        All2AllManagerBase,
+    )
+
+    dist.init_process_group(
+        "gloo",
+        init_method=init_method,
+        rank=rank,
+        world_size=4,
+        timeout=timedelta(seconds=60),
+    )
+    # With PP=2 and DP=2 the EP groups are [0, 2] and [1, 3].
+    groups = [dist.new_group([0, 2]), dist.new_group([1, 3])]
+    group = groups[rank % 2]
+
+    def fake_init(self, g):
+        self.__dict__.update(rank=g.rank(), world_size=g.size(), cpu_group=g)
+
+    All2AllManagerBase.__init__ = fake_init  # type: ignore[method-assign]
+    all2all.CudaRTLibrary = Mock  # type: ignore[misc]
+    communicator = Mock()
+    communicator.get_unique_id.return_value = f"uid-of-rank-{rank}"
+    communicator.init.return_value.create_dev_comm.return_value = Mock(lsa_size=2)
+    sys.modules["mori.cco"] = SimpleNamespace(
+        Communicator=communicator,
+        CCODevCommRequirements=SimpleNamespace,
+        GDA_CONNECTION_NONE=0,
+    )
+    all2all.MoriEPv2All2AllManager(group)
+    # Every rank joins the communicator with the ID its group's first rank made.
+    assert communicator.init.call_args.args[2] == f"uid-of-rank-{rank % 2}"
+    dist.destroy_process_group()
+
+
+def test_unique_id_comes_from_the_ep_groups_first_rank():
+    mp.spawn(_unique_id_worker, args=(f"tcp://127.0.0.1:{get_open_port()}",), nprocs=4)
+
+
+def test_experts_write_into_combine_staging_only_when_it_fits():
+    from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
+        TopKWeightAndReduceContiguous,
+        TopKWeightAndReduceNoOP,
+    )
+
+    op = FakeOp()
+    alloc = _make_pf(op).allocate_fused_expert_output
+    device = op.combine_in.device
+    out = alloc((32, 4), torch.bfloat16, device, TopKWeightAndReduceNoOP())
+    assert out is not None and out.shape == (32, 4)
+    assert out.data_ptr() == op.combine_in.data_ptr()
+    # Unreduced expert output, other dtypes or too many rows keep the workspace.
+    for args in (
+        ((32, 4), torch.bfloat16, device, TopKWeightAndReduceContiguous()),
+        ((32, 4), torch.float32, device, TopKWeightAndReduceNoOP()),
+        ((257, 4), torch.bfloat16, device, TopKWeightAndReduceNoOP()),
+    ):
+        assert alloc(*args) is None
+
+
+@pytest.mark.parametrize(
+    "in_dtype,quant_dtype,quantizes_itself,per_tensor,ok",
+    [
+        (torch.bfloat16, None, False, False, True),
+        # e.g. AITER, which quantizes after the all2all.
+        (torch.bfloat16, torch.float8_e4m3fn, True, False, True),
+        # AITER cannot quantize the received rows per tensor.
+        (torch.bfloat16, torch.float8_e4m3fn, True, True, False),
+        # e.g. FP8 Triton experts, which expect the dispatch to quantize.
+        (torch.bfloat16, torch.float8_e4m3fn, False, False, False),
+        (torch.float16, None, False, False, False),
+    ],
+)
+def test_experts_it_cannot_feed_are_rejected(
+    in_dtype, quant_dtype, quantizes_itself, per_tensor, ok
+):
+    experts = SimpleNamespace(
+        moe_config=SimpleNamespace(in_dtype=in_dtype),
+        quant_config=SimpleNamespace(quant_dtype=quant_dtype, is_per_tensor=per_tensor),
+        expects_unquantized_inputs=quantizes_itself,
+    )
+    prepare_finalize = _make_pf(FakeOp(16))
+    if ok:
+        prepare_finalize.post_init_setup(experts)
+    else:
+        with pytest.raises(ValueError, match="VLLM_MORI_EP_V2=0"):
+            prepare_finalize.post_init_setup(experts)
+
+
+def _parallel_config(**overrides):
+    fields = dict(
+        all2all_backend="mori",
+        mori_ep_v2=None,
+        tensor_parallel_size=1,
+        data_parallel_size=8,
+        data_parallel_size_local=8,
+        data_parallel_external_lb=False,
+        data_parallel_backend="mp",
+        nnodes_within_dp=1,
+        prefill_context_parallel_size=1,
+    )
+    return SimpleNamespace(**(fields | overrides))
+
+
+def _model_config(quant_cfg=None, **overrides):
+    fields = {"dtype": torch.bfloat16, "quantization": None}
+    arch = SimpleNamespace(quantization_config=quant_cfg)
+    return SimpleNamespace(model_arch_config=arch, **(fields | overrides))
+
+
+_EXTERNAL_LB = {"data_parallel_size_local": 1, "data_parallel_external_lb": True}
+_FP8_BLOCK = {"quantization": "fp8", "quant_cfg": {"weight_block_size": [128, 128]}}
+_MODELOPT_FP8 = {"quantization": "modelopt", "quant_cfg": {"quant_algo": "FP8"}}
+_MODELOPT_NVFP4 = {"quantization": "modelopt", "quant_cfg": {"quant_algo": "NVFP4"}}
+
+
+def _ct(strategy):
+    act = {"type": "float", "num_bits": 8, "strategy": strategy}
+    cfg = {"config_groups": {"group_0": {"input_activations": act}}}
+    return {"quantization": "compressed-tensors", "quant_cfg": cfg}
+
+
+def _quark(qscheme):
+    inputs = {"dtype": "fp8_e4m3", "qscheme": qscheme}
+    cfg = {"global_quant_config": {"input_tensors": inputs}}
+    return {"quantization": "quark", "quant_cfg": cfg}
+
+
+@pytest.mark.parametrize(
+    "parallel,model,moe_backend,env,ep_v2",
+    [
+        ({}, {}, "aiter", {}, True),
+        # EP across hosts: DP over two nodes, even with ep_size <= 8.
+        ({"data_parallel_size": 16}, {}, "aiter", {}, False),
+        ({"data_parallel_size_local": 4}, {}, "aiter", {}, False),
+        # A DP replica (TP/PP) spanning nodes.
+        ({"nnodes_within_dp": 2}, {}, "aiter", {}, False),
+        # External LB only sees its own rank: fall back to the EP size.
+        (_EXTERNAL_LB, {}, "aiter", {}, True),
+        (_EXTERNAL_LB | {"data_parallel_size": 16}, {}, "aiter", {}, False),
+        # The LSA team may span hosts (gfx1250); the manager checks it.
+        ({"data_parallel_size": 16, "lsa_spans_hosts": True}, {}, "aiter", {}, True),
+        ({}, {}, "aiter", {"MORI_ENABLE_SDMA": "1"}, False),
+        # gfx1250 keeps EPv2 under SDMA.
+        ({"sdma_picks_epv1": False}, {}, "aiter", {"MORI_ENABLE_SDMA": "1"}, True),
+        ({}, {}, "aiter_mega_moe", {}, False),
+        ({}, {"dtype": torch.float16}, "aiter", {}, False),
+        # An explicit VLLM_MORI_EP_V2 wins over the deployment checks.
+        ({}, {"dtype": torch.float16}, "aiter", {"VLLM_MORI_EP_V2": "1"}, True),
+        ({}, {}, "aiter", {"VLLM_MORI_EP_V2": "0"}, False),
+        # Already resolved: kept as is.
+        ({"mori_ep_v2": True}, {"dtype": torch.float16}, "aiter", {}, True),
+        # Other MoRI backends always run EPv1.
+        ({"all2all_backend": "mori_high_throughput"}, {}, "aiter", {}, None),
+        ({"validated_arch": False}, {}, "aiter", {}, False),
+        ({"mori_reason": "too old"}, {}, "aiter", {}, False),
+        # Per-tensor FP8 activations; block FP8 and NVFP4 are fine.
+        ({}, {"quantization": "fp8"}, "aiter", {}, False),
+        ({}, _FP8_BLOCK, "aiter", {}, True),
+        ({}, _MODELOPT_FP8, "aiter", {}, False),
+        ({}, _MODELOPT_NVFP4, "aiter", {}, True),
+        ({}, _ct("tensor"), "aiter", {}, False),
+        ({}, _ct("group"), "aiter", {}, True),
+        ({}, _quark("per_tensor"), "aiter", {}, False),
+        ({}, _quark("per_group"), "aiter", {}, True),
+    ],
+)
+def test_mori_picks_epv2_where_supported(
+    monkeypatch, parallel, model, moe_backend, env, ep_v2
+):
+    monkeypatch.delenv("VLLM_MORI_EP_V2", raising=False)
+    monkeypatch.delenv("MORI_ENABLE_SDMA", raising=False)
+    for name, value in env.items():
+        monkeypatch.setenv(name, value)
+    parallel = dict(parallel)
+    lsa_spans_hosts = parallel.pop("lsa_spans_hosts", False)
+    sdma_picks_epv1 = parallel.pop("sdma_picks_epv1", True)
+    validated_arch = parallel.pop("validated_arch", True)
+    mori_reason = parallel.pop("mori_reason", None)
+    monkeypatch.setattr(
+        import_utils, "mori_epv2_unavailable_reason", lambda: mori_reason
+    )
+    parallel_config = _parallel_config(**parallel)
+    model_config = _model_config(**model)
+    resolve_mori_all2all(
+        parallel_config,
+        model_config,
+        moe_backend,
+        lsa_spans_hosts,
+        sdma_picks_epv1,
+        validated_arch,
+    )
+    assert parallel_config.mori_ep_v2 is ep_v2
+
+
+@pytest.mark.parametrize(
+    "modules,version,usable",
+    [
+        ({"mori", "mori.cco", "mori.ops.dispatch_combine_v2"}, "1.2.4", True),
+        ({"mori", "mori.cco", "mori.ops.dispatch_combine_v2"}, "1.2.5.dev3+g1", True),
+        ({"mori", "mori.cco", "mori.ops.dispatch_combine_v2"}, "1.2.3.post1", False),
+        ({"mori", "mori.cco", "mori.ops.dispatch_combine_v2"}, "unknown", False),
+        ({"mori"}, "1.2.4", False),
+        (set(), "1.2.4", False),
+    ],
+)
+def test_mori_epv2_needs_a_recent_mori(monkeypatch, modules, version, usable):
+    monkeypatch.setattr(import_utils, "_has_module", lambda name: name in modules)
+    monkeypatch.setitem(sys.modules, "mori", SimpleNamespace(__version__=version))
+    assert (import_utils.mori_epv2_unavailable_reason() is None) is usable
+
+
+@pytest.mark.parametrize(
+    "backend,internode,max_tokens,kernel",
+    [
+        ("mori", False, 4096, "IntraNode"),
+        ("mori", True, 256, "InterNodeV1LL"),
+        ("mori", True, 257, "InterNodeV1"),
+        ("mori_high_throughput", True, 128, "InterNodeV1"),
+        ("mori_low_latency", True, 4096, "InterNodeV1LL"),
+    ],
+)
+def test_mori_epv1_kernel_choice(monkeypatch, backend, internode, max_tokens, kernel):
+    from vllm.distributed.device_communicators import all2all
+
+    kernel_types = SimpleNamespace(
+        IntraNode="IntraNode", InterNodeV1="InterNodeV1", InterNodeV1LL="InterNodeV1LL"
+    )
+    fake_mori = SimpleNamespace(
+        ops=SimpleNamespace(EpDispatchCombineKernelType=kernel_types)
+    )
+    monkeypatch.setitem(sys.modules, "mori", fake_mori)
+    monkeypatch.setitem(
+        sys.modules,
+        "vllm.platforms.rocm",
+        SimpleNamespace(on_gfx942=lambda: False, on_gfx950=lambda: True),
+    )
+    monkeypatch.delenv(
+        "VLLM_MORI_DISPATCH_INTER_KERNEL_SWITCH_THRESHOLD", raising=False
+    )
+    manager = object.__new__(all2all.MoriAll2AllManager)
+    manager.internode = internode
+    manager._all2all_backend = backend
+    kwargs = manager._make_all2all_kwargs(
+        rank=0,
+        num_ep_ranks=16,
+        input_dtype=torch.bfloat16,
+        quant_dtype=torch.bfloat16,
+        token_hidden_size=7168,
+        scale_dim=0,
+        scale_type_size=4,
+        max_num_tokens_per_dp_rank=max_tokens,
+        num_local_experts=16,
+        num_experts_per_token=8,
+    )
+    assert kwargs["kernel_type"] == kernel
diff --git a/tests/kernels/moe/test_muse_spark_router.py b/tests/kernels/moe/test_muse_spark_router.py
index 7e3938497eba..59cf6af7be62 100644
--- a/tests/kernels/moe/test_muse_spark_router.py
+++ b/tests/kernels/moe/test_muse_spark_router.py
@@ -189,3 +189,174 @@ def test_fused_routing_writes_into_preallocated_buffers():
     assert id_view.data_ptr() == ids.data_ptr()
     assert (weights[num_tokens:] == -7.0).all()
     assert (ids[num_tokens:] == -7).all()
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
+@pytest.mark.parametrize("num_experts", [32, 256])
+def test_routing_padding_is_zero_with_nonfinite_logits(num_experts):
+    """Padding must not introduce NaNs, or change any active route, on replay."""
+    torch.manual_seed(23)
+    logits = torch.randn(8, num_experts, device="cuda")
+    bias = torch.randn(num_experts, device="cuda")
+    padding = torch.zeros(8, dtype=torch.bool, device="cuda")
+    router = MuseSpark1xRouter(8, num_experts, bias, 1.0)
+    expected_weights, expected_ids = router._compute_routing(None, logits, None)
+
+    def route():
+        return fused_musespark1x_routing(logits, bias, 8, 1.0, is_padding=padding)
+
+    stream = torch.cuda.Stream()
+    stream.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(stream):
+        for _ in range(3):
+            route()
+    torch.cuda.current_stream().wait_stream(stream)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        weights, ids = route()
+    for active in (7, 3, 0):
+        padding[active:] = True
+        logits[active:] = float("nan")
+        graph.replay()
+        torch.testing.assert_close(
+            weights[:active], expected_weights[:active], rtol=0, atol=0
+        )
+        assert torch.equal(ids[:active], expected_ids[:active])
+        assert (weights[active:] == 0).all()
+        assert (ids[active:] == -1).all()
+    routed_weights, routed_ids = router._compute_routing(
+        None, logits, None, is_padding=padding
+    )
+    assert torch.equal(routed_weights, weights)
+    assert torch.equal(routed_ids, ids)
+
+
+def _triton_routing_reference(logits, bias, multiplier, scale, eps, padding):
+    from vllm.models.muse_spark.common.ops.fused_routing import _routing_kernel
+    from vllm.platforms import current_platform
+
+    tokens = logits.shape[0]
+    weights = torch.empty((tokens, 8), device=logits.device)
+    ids = torch.empty((tokens, 8), device=logits.device, dtype=torch.int32)
+    if tokens:
+        pdl = current_platform.is_arch_support_pdl()
+        _routing_kernel[(tokens,)](
+            logits,
+            bias,
+            weights,
+            ids,
+            padding,
+            0 if padding is None else padding.stride(0),
+            tokens,
+            256,
+            *logits.stride(),
+            *weights.stride(),
+            *ids.stride(),
+            multiplier,
+            scale,
+            eps,
+            TOPK=8,
+            TOPK_PADDED=8,
+            NUM_EXPERTS_PADDED=256,
+            NORMALIZATION_MODE=0,
+            HAS_MULTIPLIER=True,
+            HAS_EPS=True,
+            USE_PDL=pdl,
+            num_warps=1,
+            **({"launch_pdl": True} if pdl else {}),
+        )
+    return weights, ids
+
+
+@pytest.mark.skipif(
+    not torch.cuda.is_available() or torch.version.hip is not None,
+    reason="requires the CUDA routing kernel",
+)
+@pytest.mark.parametrize("tokens", [0, 1, 12, 16, 48, 128, 2048, 2049, 16384])
+@pytest.mark.parametrize("case", ["random", "ties", "near_ties", "negative", "strided"])
+def test_cuda_muse_routing_matches_triton(tokens, case):
+    """Preserve expert choices, FP32 rounding, padding, and caller-owned buffers."""
+    from vllm import _custom_ops as ops
+
+    torch.manual_seed(17)
+    width = 260 if case == "strided" else 256
+    logits = torch.randn(tokens, width, device="cuda")[:, :256]
+    bias = torch.randn(256, device="cuda") * 0.2
+    if case == "ties":
+        logits.zero_()
+        bias.zero_()
+    elif case == "near_ties":
+        logits.mul_(1e-6)
+        bias.zero_()
+    elif case == "negative":
+        logits.sub_(80)
+    padding = (torch.arange(tokens * 2, device="cuda") % 3 == 0)[::2]
+    weights = torch.full((tokens + 3, 12), -7.0, device="cuda")
+    ids = torch.full((tokens + 3, 12), -7, device="cuda", dtype=torch.int32)
+    w, i = weights[:tokens, :8], ids[:tokens, :8]
+    ref_w, ref_i = _triton_routing_reference(logits, bias, 0.5, 2.5, 1e-15, padding)
+    ops.muse_spark_routing(w, i, logits, bias, 0.5, 2.5, 1e-15, padding)
+    assert torch.equal(i, ref_i)
+    assert torch.equal(w, ref_w)
+    assert (i[padding] == -1).all()
+    assert (w[padding] == 0).all()
+    out_w, out_i = fused_musespark1x_routing(
+        logits,
+        bias,
+        8,
+        2.5,
+        router_output_multiplier=0.5,
+        router_score_division_eps=1e-15,
+        topk_weights=weights,
+        topk_ids=ids,
+        is_padding=padding,
+    )
+    assert out_w.untyped_storage().data_ptr() == weights.untyped_storage().data_ptr()
+    assert out_i.untyped_storage().data_ptr() == ids.untyped_storage().data_ptr()
+    assert torch.equal(out_w, ref_w) and torch.equal(out_i, ref_i)
+    assert (weights[tokens:] == -7).all() and (weights[:, 8:] == -7).all()
+    assert (ids[tokens:] == -7).all() and (ids[:, 8:] == -7).all()
+
+
+@pytest.mark.skipif(
+    not torch.cuda.is_available() or torch.version.hip is not None,
+    reason="requires the CUDA routing kernel",
+)
+def test_cuda_muse_routing_graph_replay():
+    """Replays must consume updated logits, bias, and arbitrary padding masks."""
+    torch.manual_seed(19)
+    logits = torch.randn(48, 256, device="cuda")
+    bias = torch.randn(256, device="cuda")
+    padding = torch.zeros(48, device="cuda", dtype=torch.bool)
+    weights = torch.empty(48, 8, device="cuda")
+    ids = torch.empty(48, 8, device="cuda", dtype=torch.int32)
+
+    def route():
+        return fused_musespark1x_routing(
+            logits,
+            bias,
+            8,
+            1.0,
+            router_score_division_eps=1e-15,
+            topk_weights=weights,
+            topk_ids=ids,
+            is_padding=padding,
+        )
+
+    stream = torch.cuda.Stream()
+    stream.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(stream):
+        for _ in range(3):
+            route()
+    torch.cuda.current_stream().wait_stream(stream)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        route()
+    for offset in range(3):
+        logits.normal_()
+        bias.normal_()
+        padding.copy_(torch.arange(48, device="cuda") % 3 == offset)
+        graph.replay()
+        ref_w, ref_i = _triton_routing_reference(logits, bias, 1.0, 1.0, 1e-15, padding)
+        assert torch.equal(ids, ref_i)
+        assert torch.equal(weights, ref_w)
diff --git a/tests/kernels/moe/test_muse_spark_trtllm_moe.py b/tests/kernels/moe/test_muse_spark_trtllm_moe.py
index cce1bd2ebc85..829edfdfcab5 100644
--- a/tests/kernels/moe/test_muse_spark_trtllm_moe.py
+++ b/tests/kernels/moe/test_muse_spark_trtllm_moe.py
@@ -22,7 +22,12 @@
     dequantize_nvfp4_to_dtype,
 )
 from vllm import _custom_ops as ops
-from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config
+from vllm.config import (
+    KernelConfig,
+    ParallelConfig,
+    VllmConfig,
+    set_current_vllm_config,
+)
 from vllm.distributed import get_tp_group
 from vllm.model_executor.layers.fused_moe.config import RoutingMethodType
 from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput
@@ -432,3 +437,104 @@ def test_muse_spark_moe_cudagraph(nvfp4, dist_init, workspace_init):
         ckpt, quant_config = _bf16_checkpoint(), None
     _, _, _, vllm_config, runner = build_moe(MuseSpark1xRMSNorm, ckpt, quant_config)
     assert_cudagraph_replays(runner, vllm_config, *_inputs(5), *_inputs(6))
+
+
+@torch.inference_mode()
+def test_megamoe_trtllm_weight_layout_matches_native(dist_init, workspace_init):
+    """Shared TRT storage must preserve native MegaMoE outputs and graph replay."""
+    from vllm.models.muse_spark.nvidia.mega_moe import _get_deep_gemm
+    from vllm.models.muse_spark.nvidia.ops.prepare_megamoe_nvfp4 import (
+        prepare_megamoe_nvfp4_inputs,
+    )
+
+    torch.manual_seed(137)
+    ckpt = _nvfp4_checkpoint()
+    norm = MuseSpark1xRMSNormWithInputScale(
+        hidden_size=K,
+        eps=1e-8,
+        has_weight=True,
+        is_post_norm=True,
+        zero_centered_gamma=True,
+    ).cuda()
+    norm.weight.copy_(torch.linspace(-0.5, 0.5, K, device="cuda"))
+    layers = []
+    for layout, tp_size in (("native", 1), ("trtllm", 4)):
+        config = VllmConfig(
+            kernel_config=KernelConfig(moe_backend="deep_gemm_mega_moe"),
+            quant_config=nvfp4_quant_config(),
+            parallel_config=ParallelConfig(
+                tensor_parallel_size=tp_size,
+                enable_expert_parallel=True,
+                distributed_executor_backend="mp",
+            ),
+        )
+        with set_current_vllm_config(config):
+            layer = MuseSparkMegaMoEExperts(
+                config,
+                num_experts=E,
+                num_local_experts=E,
+                experts_start_idx=0,
+                top_k=TOPK,
+                hidden_size=K,
+                intermediate_size=N,
+                post_expert_norm=norm,
+                num_max_tokens_per_rank=128,
+                prefix=f"shared_layout.{layout}",
+            ).cuda()
+            assert layer.weight_layout == layout
+            assert set(layer.load_weights(ckpt.items())) == set(ckpt)
+            layer.finalize_weights()
+            layers.append(layer)
+    native, shared = layers
+    for view, parameter in (
+        (shared._transformed_l1[0], shared.w13_weight),
+        (shared._transformed_l1[1], shared.w13_weight_scale),
+        (shared._transformed_l2[0], shared.w2_weight),
+        (shared._transformed_l2[1], shared.w2_weight_scale),
+    ):
+        assert view.data_ptr() == parameter.data_ptr()
+        assert view.untyped_storage().nbytes() == parameter.untyped_storage().nbytes()
+    shared.finalize_weights()  # Must not shuffle the already finalized weights again.
+
+    dg = _get_deep_gemm()
+    buffer = dg.SymmBuffer(
+        torch.distributed.group.WORLD, E, 128, TOPK, K, N, mma_type="fp4xfp4"
+    )
+    try:
+        for tokens in (1, 8, 32, 64):
+            x = torch.randn(tokens, K, device="cuda", dtype=torch.bfloat16) / 10
+            weights, ids = torch.randn(tokens, E, device="cuda").topk(TOPK)
+            weights = weights.softmax(-1)
+            ids = ids.int()
+            ids[::3, -1] = -1
+            buffer.topk_idx.fill_(-1)
+            buffer.topk_weights.zero_()
+
+            def stage(x=x, weights=weights, ids=ids, tokens=tokens):
+                prepare_megamoe_nvfp4_inputs(
+                    x,
+                    weights,
+                    ids,
+                    buffer.x[:tokens],
+                    buffer.x_sf[:tokens].view(torch.uint8),
+                    buffer.topk_idx[:tokens],
+                    buffer.topk_weights[:tokens],
+                    None,
+                )
+
+            stage()
+            expected = native.forward_prepared(x, buffer)
+            actual = shared.forward_prepared(x, buffer)
+            torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+            graph = torch.cuda.CUDAGraph()
+            with torch.cuda.graph(graph):
+                actual = shared.forward_prepared(x, buffer)
+            for _ in range(3):
+                x.normal_(std=0.1)
+                ids.copy_(torch.randn(tokens, E, device="cuda").topk(TOPK).indices)
+                stage()
+                expected = native.forward_prepared(x, buffer)
+                graph.replay()
+                torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+    finally:
+        buffer.destroy()
diff --git a/tests/kernels/moe/test_unquantized_backend_selection.py b/tests/kernels/moe/test_unquantized_backend_selection.py
index 7081142b3808..6e6d3ca5f931 100644
--- a/tests/kernels/moe/test_unquantized_backend_selection.py
+++ b/tests/kernels/moe/test_unquantized_backend_selection.py
@@ -380,6 +380,7 @@ def test_select_cuda_flashinfer_trtllm_modular_backend(
     [
         "mori_high_throughput",
         "mori_low_latency",
+        "mori",
         "flashinfer_nvlink_two_sided",
         "flashinfer_nvlink_one_sided",
     ],
diff --git a/tests/kernels/quantization/test_k3_gfx942_int4_requant.py b/tests/kernels/quantization/test_k3_gfx942_int4_requant.py
new file mode 100644
index 000000000000..4f9de6fa7b09
--- /dev/null
+++ b/tests/kernels/quantization/test_k3_gfx942_int4_requant.py
@@ -0,0 +1,150 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Kimi-K3 gfx942 MXFP4-to-int4 load-time requant."""
+
+import pytest
+import torch
+
+from vllm.platforms import current_platform
+
+if not current_platform.is_rocm():
+    pytest.skip("requires ROCm", allow_module_level=True)
+
+from vllm.platforms.rocm import on_gfx942  # noqa: E402
+
+if not on_gfx942():
+    pytest.skip("requires gfx942", allow_module_level=True)
+
+pytest.importorskip("aiter")
+
+from aiter.ops.quant import per_1x32_i4_quant  # noqa: E402
+from aiter.ops.shuffle import (  # noqa: E402
+    pack_int8_to_packed_int4,
+    shuffle_scale_for_int4,
+    shuffle_weight,
+)
+from aiter.utility import fp4_utils  # noqa: E402
+
+from vllm.model_executor.layers.quantization import mxfp4  # noqa: E402
+
+NUM_EXPERTS, ROWS, COLS = 128, 256, 2048
+# Kimi-K3 w2 per rank under TP8. 896 experts do not split evenly into chunks.
+TP8_W2_SHAPE = (896, 3584, 384)
+E2M1_VALUES = (0, 0.5, 1, 1.5, 2, 3, 4, 6, -0.0, -0.5, -1, -1.5, -2, -3, -4, -6)
+
+
+def _mxfp4_experts(
+    num_experts: int = NUM_EXPERTS, rows: int = ROWS, cols: int = COLS
+) -> tuple[torch.Tensor, torch.Tensor]:
+    gen = torch.Generator(device="cuda").manual_seed(0)
+    weight = torch.randint(
+        0,
+        256,
+        (num_experts, rows, cols // 2),
+        dtype=torch.uint8,
+        device="cuda",
+        generator=gen,
+    )
+    # Normal E8M0 block scales; 0 and 255 have special meanings.
+    scale = torch.randint(
+        118,
+        132,
+        (num_experts, rows, cols // 32),
+        dtype=torch.uint8,
+        device="cuda",
+        generator=gen,
+    )
+    return weight, scale
+
+
+def _requant(
+    weight: torch.Tensor, scale: torch.Tensor
+) -> tuple[torch.Tensor, torch.Tensor]:
+    out_weight, out_scale = mxfp4._mxfp4_experts_to_int4_gfx942(
+        torch.nn.Parameter(weight.clone(), requires_grad=False),
+        torch.nn.Parameter(scale.clone(), requires_grad=False),
+    )
+    return out_weight.view(torch.uint8), out_scale.view(torch.uint8)
+
+
+def _reference_requant(
+    weight: torch.Tensor, scale: torch.Tensor
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """All experts at once, decoding MXFP4 as the OCP MX spec defines it."""
+    num_experts, rows, packed_cols = weight.shape
+    cols = packed_cols * 2
+    # Element 2i is in the low nibble and element 2i + 1 in the high nibble.
+    codes = torch.stack((weight & 0xF, weight >> 4), dim=-1).view(
+        num_experts, rows, cols
+    )
+    values = torch.tensor(E2M1_VALUES, device=weight.device)[codes.long()]
+    block_scales = torch.ldexp(
+        torch.ones_like(scale, dtype=torch.float32), scale.int() - 127
+    )
+    weight_bf16 = (
+        (values.view(num_experts, rows, cols // 32, 32) * block_scales.unsqueeze(-1))
+        .view(num_experts, rows, cols)
+        .to(torch.bfloat16)
+    )
+
+    weight_int4, weight_scale = per_1x32_i4_quant(weight_bf16)
+    packed = pack_int8_to_packed_int4(shuffle_weight(weight_int4, (16, 16)))
+    packed_scale = shuffle_scale_for_int4(weight_scale, group_size=32).reshape(-1)
+    return (
+        packed.view(torch.uint8).view(num_experts, rows, packed_cols),
+        packed_scale.view(torch.uint8),
+    )
+
+
+@pytest.fixture
+def load_model_allocator_settings():
+    original = torch._C._accelerator_getAllocatorSettings()
+    prefix = original.strip().rstrip(",")
+    settings = "max_split_size_mb:20"
+    torch._C._accelerator_setAllocatorSettings(
+        f"{prefix},{settings}" if prefix else settings
+    )
+    yield
+    torch._C._accelerator_setAllocatorSettings(original)
+
+
+@pytest.mark.parametrize(
+    "experts_per_chunk",
+    [1, 5, NUM_EXPERTS],
+    ids=["one_expert", "uneven_last_chunk", "single_chunk"],
+)
+def test_requant_matches_reference(monkeypatch, experts_per_chunk):
+    monkeypatch.setattr(
+        mxfp4, "_K3_INT4_CONVERT_CHUNK_VALUES", experts_per_chunk * ROWS * COLS
+    )
+    weight, scale = _mxfp4_experts()
+
+    out_weight, out_scale = _requant(weight, scale)
+    ref_weight, ref_scale = _reference_requant(weight, scale)
+
+    assert torch.equal(out_weight, ref_weight)
+    assert torch.equal(out_scale, ref_scale)
+
+
+def test_requant_allocates_device_memory_only_in_first_chunk(
+    monkeypatch, load_model_allocator_settings
+):
+    """Later chunks, the last one included, must reuse the first chunk's cached
+    blocks under load_model's allocator settings."""
+    weight, scale = _mxfp4_experts(*TP8_W2_SHAPE)
+    _requant(weight, scale)
+    torch.accelerator.empty_cache()
+
+    device_allocs = []
+    mxfp4_to_f32 = fp4_utils.mxfp4_to_f32
+
+    def recording_mxfp4_to_f32(x):
+        device_allocs.append(torch.accelerator.memory_stats()["num_device_alloc"])
+        return mxfp4_to_f32(x)
+
+    monkeypatch.setattr(fp4_utils, "mxfp4_to_f32", recording_mxfp4_to_f32)
+    _requant(weight, scale)
+    device_allocs.append(torch.accelerator.memory_stats()["num_device_alloc"])
+
+    assert len(device_allocs) > 3
+    assert device_allocs[1:] == [device_allocs[1]] * (len(device_allocs) - 1)
diff --git a/tests/kernels/quantization/test_nvfp4_emulation.py b/tests/kernels/quantization/test_nvfp4_emulation.py
index 9fc97d0e09a8..93bb25ad6b27 100644
--- a/tests/kernels/quantization/test_nvfp4_emulation.py
+++ b/tests/kernels/quantization/test_nvfp4_emulation.py
@@ -864,3 +864,60 @@ def test_nvfp4_dense_emulation(
         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)
+
+
+@pytest.mark.skip_global_cleanup
+@pytest.mark.parametrize("invert_global_scales", [True, False])
+@pytest.mark.parametrize("experts_per_chunk", [1, 2, 3, 8])
+def test_dequantize_nvfp4_moe_weights_chunked_matches_full(
+    experts_per_chunk: int, invert_global_scales: bool, monkeypatch
+) -> None:
+    """Chunked dequantization must equal the single-chunk dequantization exactly.
+
+    5 experts cover one expert per chunk, a smaller last chunk (2, 3) and a
+    chunk larger than the number of experts (8).
+    """
+    set_random_seed(0)
+    device = torch.device(current_platform.device_type)
+    num_experts, hidden, inter, group = 5, 64, 32, 16
+
+    def packed(rows: int, cols: int) -> torch.Tensor:
+        return torch.randint(
+            0, 256, (num_experts, rows, cols // 2), dtype=torch.uint8, device=device
+        )
+
+    def block_scale(rows: int, cols: int) -> torch.Tensor:
+        scale = torch.rand(num_experts, rows, cols // group, device=device) * 4 + 0.5
+        return scale.to(torch.float8_e4m3fn)
+
+    w13_packed, w13_scale = packed(2 * inter, hidden), block_scale(2 * inter, hidden)
+    w2_packed, w2_scale = packed(hidden, inter), block_scale(hidden, inter)
+    w13_global = torch.rand(num_experts, 2, device=device) + 0.5
+    w2_global = torch.rand(num_experts, device=device) + 0.5
+
+    args = (w13_packed, w13_scale, w13_global, w2_packed, w2_scale, w2_global)
+    dtype = torch.float32
+    # The default chunk budget holds all experts at once.
+    ((_, expected_w13, expected_w2),) = (
+        nvfp4_emulation_utils.dequantize_nvfp4_moe_weights(
+            *args, dtype, invert_global_scales=invert_global_scales
+        )
+    )
+
+    expert_bytes = dtype.itemsize * 2 * (w13_packed[0].numel() + w2_packed[0].numel())
+    monkeypatch.setattr(
+        nvfp4_emulation_utils,
+        "_DEQUANT_CHUNK_BYTES",
+        experts_per_chunk * expert_bytes,
+    )
+    chunks = list(
+        nvfp4_emulation_utils.dequantize_nvfp4_moe_weights(
+            *args, dtype, invert_global_scales=invert_global_scales
+        )
+    )
+
+    expected_starts = list(range(0, num_experts, experts_per_chunk))
+    assert [c[0].start for c in chunks] == expected_starts
+    assert chunks[-1][0].stop == num_experts
+    assert torch.equal(torch.cat([c[1] for c in chunks]), expected_w13)
+    assert torch.equal(torch.cat([c[2] for c in chunks]), expected_w2)
diff --git a/tests/kernels/quantization/test_rdna_hybrid_w4a16.py b/tests/kernels/quantization/test_rdna_hybrid_w4a16.py
index edbfae158531..8cb515c1320c 100644
--- a/tests/kernels/quantization/test_rdna_hybrid_w4a16.py
+++ b/tests/kernels/quantization/test_rdna_hybrid_w4a16.py
@@ -1227,3 +1227,33 @@ def test_rdna_hybrid_w4a16_medium_skinny_matches_reference(dtype, has_zp):
     reference = _rdna_hybrid_w4a16_reference(x_mk, w_int4_nk, scales, zp, G, bias=None)
 
     torch.testing.assert_close(output, reference, rtol=2e-2, atol=2e-2)
+
+
+# ---------------------------------------------------------------------------
+# Row offsets past int32
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.skipif(not on_gfx1x(), reason="Hybrid path is gfx11/gfx12 only")
+@pytest.mark.parametrize(
+    "K,N",
+    [
+        (256, 34816),  # M * N passes 2**31: output offsets
+        (17408, 256),  # M * K passes 2**31: activation offsets
+    ],
+)
+def test_triton_prefill_gemm_rows_past_int32_offsets(K, N):
+    """Rows whose element offset in A or C does not fit in int32."""
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA/HIP device not available")
+    M = 2**31 // max(K, N) + 128
+    G = 128
+
+    set_random_seed(0)
+    x, w_int4, b_q, scales, _, _ = _make_prefill_case(M, K, N, G, torch.bfloat16, False)
+    out = triton_w4a16_skinny_fmt_gemm(a=x, b_q=b_q, scales=scales, group_size=G)
+
+    # Dequantize in bf16 like the kernel does, so a long K adds no rounding gap.
+    w = (w_int4 - 8).to(torch.bfloat16) * scales.repeat_interleave(G, dim=1)
+    ref = (x[-64:].float() @ w.float().t()).to(torch.bfloat16)
+    torch.testing.assert_close(out[-64:], ref, rtol=1e-2, atol=5e-2)
diff --git a/tests/kernels/quantization/test_rocm_mxfp8_linear.py b/tests/kernels/quantization/test_rocm_mxfp8_linear.py
index d5e48f2aafa2..b784551928a6 100644
--- a/tests/kernels/quantization/test_rocm_mxfp8_linear.py
+++ b/tests/kernels/quantization/test_rocm_mxfp8_linear.py
@@ -8,6 +8,13 @@
 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.
+
+They also cover the tiny-M fused quant+GEMM fastpath of
+``rocm_mxfp8_block32_gemm_quant`` (at M <= 8 the activation quant runs inside
+the packed GEMM's K loop): bit-exactness with the stock two-launch
+``mxfp8_e4m3_quantize`` + ``rocm_mxfp8_block32_gemm`` path at the serving
+shapes and under CUDA-graph capture/replay, agreement with a high-precision
+QDQ fp32 reference, and bit-identical fallbacks for ineligible inputs.
 """
 
 import pytest
@@ -26,6 +33,17 @@
     reason="native MX linear requires CDNA4 (gfx95x)",
 )
 
+from vllm.model_executor.kernels.linear.mxfp8.Mxfp8LinearKernel import (  # noqa: E402
+    Mxfp8LinearLayerConfig,
+)
+from vllm.model_executor.kernels.linear.mxfp8.rocm_block32_gemm import (  # noqa: E402
+    rocm_mxfp8_block32_gemm,
+    rocm_mxfp8_block32_gemm_quant,
+)
+from vllm.model_executor.kernels.linear.mxfp8.rocm_native import (  # noqa: E402
+    RocmDotScaledMxfp8LinearKernel,
+)
+
 # (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).
@@ -36,6 +54,13 @@
     ("wo_b", 5120, 8192 // 4),
 ]
 
+M31_TP4_SHAPES = [
+    ("fused_qkv", 2560, 6144),
+    ("o_proj", 6144, 2048),
+    ("shared_gate_up", 1536, 6144),
+    ("shared_down", 6144, 768),
+]
+
 
 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
@@ -46,8 +71,46 @@ def _make_layer(n: int, k: int, device: str) -> torch.nn.Module:
     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])
+@pytest.mark.parametrize(
+    "n,k,selected,fallback",
+    [
+        (2560, 6144, (16, 16, 1024, 2, 3), (16, 16, 1024, 2, 2)),
+        (1536, 6144, (16, 16, 1024, 2, 3), (16, 16, 1024, 2, 2)),
+        (6144, 768, (16, 16, 256, 2, 3), (16, 32, 256, 4, 3)),
+    ],
+)
+def test_rocm_mxfp8_linear_selects_m31_tp4_tiles(n, k, selected, fallback):
+    """The gfx950 small-M tuning is shape-scoped and ends after M=16."""
+    from vllm.model_executor.kernels.linear.mxfp8 import rocm_native
+    from vllm.platforms.rocm import on_gfx950
+
+    if not on_gfx950():
+        pytest.skip("MiniMax-M3.1 small-M tuning requires gfx950")
+    for num_tokens in (1, 8, 16):
+        assert rocm_native._select_cfg(num_tokens, n, k) == selected
+    assert rocm_native._select_cfg(17, n, k) == fallback
+
+
+@pytest.mark.parametrize(
+    "n,k,fallback",
+    [
+        (2560, 6144, (16, 16, 1024, 2, 2)),
+        (1536, 6144, (16, 16, 1024, 2, 2)),
+        (6144, 768, (16, 32, 256, 4, 3)),
+    ],
+)
+def test_rocm_mxfp8_linear_preserves_non_gfx950_tiles(monkeypatch, n, k, fallback):
+    from vllm.model_executor.kernels.linear.mxfp8 import rocm_native
+
+    monkeypatch.setattr(rocm_native, "_ON_GFX950", False)
+    for num_tokens in (1, 8, 16, 17):
+        assert rocm_native._select_cfg(num_tokens, n, k) == fallback
+
+
+@pytest.mark.parametrize(
+    "name,n,k", V41_ATTN_SHAPES + M31_TP4_SHAPES, ids=lambda v: str(v)
+)
+@pytest.mark.parametrize("num_tokens", [1, 2, 7, 8, 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.
@@ -259,6 +322,63 @@ def test_rocm_mxfp8_block32_split_k_first_use_in_capture(packed):
     assert _rel_err(out, expected) < 5e-3
 
 
+FUSED_QUANT_SHAPES = [(6144, 1024), (6144, 768), (6144, 384), (8192, 512), (4096, 640)]
+FUSED_QUANT_TOKENS = [1, 2, 3, 5, 8, 16]
+FUSED_QUANT_CASES = [
+    (n, k, m) for n, k in FUSED_QUANT_SHAPES for m in FUSED_QUANT_TOKENS
+] + [(6144, 2048, m) for m in (1, 2, 3, 4)]
+
+
+def _spy_launch(monkeypatch) -> list[bool]:
+    """Record, per ``_launch``, whether the activation quant ran in-kernel."""
+    from vllm.model_executor.kernels.linear.mxfp8 import rocm_native
+
+    fused: list[bool] = []
+    real_launch = rocm_native._launch
+
+    def spy(x, x_scale, *args, **kwargs):
+        fused.append(x_scale is None)
+        return real_launch(x, x_scale, *args, **kwargs)
+
+    monkeypatch.setattr(rocm_native, "_launch", spy)
+    return fused
+
+
+def _mxfp8_weight(n: int, k: int):
+    weight = torch.randn(n, k, device="cuda", dtype=torch.bfloat16) * 0.05
+    return _mxfp8_e4m3_quantize_torch(weight)
+
+
+@pytest.mark.parametrize("n,k,num_tokens", FUSED_QUANT_CASES, ids=lambda v: str(v))
+@torch.inference_mode()
+def test_rocm_mxfp8_fused_act_quant_bitwise_eq_unfused(monkeypatch, n, k, num_tokens):
+    """Small-M bf16 GEMMs quantize the activation in-kernel, bit-exact with
+    the unfused path: ``mxfp8_e4m3_quantize`` and the ``_select_cfg`` tile."""
+    from vllm.model_executor.kernels.linear.mxfp8 import rocm_native
+
+    torch.manual_seed(num_tokens)
+    w, w_scale = _mxfp8_weight(n, k)
+    x = torch.randn(num_tokens, k, device="cuda", dtype=torch.bfloat16)
+    x[0, :32] = 0.0  # an all-zero block clamps its scale byte to 0
+    if num_tokens > 1:
+        x[1] *= 3e6  # output beyond the fp16 range, which bf16 must keep
+
+    assert rocm_native._select_fused_quant_cfg(num_tokens, n, k) is not None
+    x_q, x_scale = mxfp8_e4m3_quantize(x)
+    expected = torch.empty((num_tokens, n), dtype=torch.bfloat16, device="cuda")
+    rocm_native._launch(
+        x_q, x_scale, w, w_scale, expected, *rocm_native._select_cfg(num_tokens, n, k)
+    )
+
+    fused = _spy_launch(monkeypatch)
+    out = rocm_native._mxfp8_dot_scaled_linear(x, w, w_scale)
+    assert fused == [True]
+    assert torch.isfinite(out).all()
+    assert torch.equal(out, expected), (
+        f"max abs diff {(out.float() - expected.float()).abs().max().item()}"
+    )
+
+
 @pytest.mark.parametrize("num_rows", [1, 64, 65])
 @torch.inference_mode()
 def test_rocm_mxfp8_quantizer_matches_torch_on_degenerate_blocks(num_rows):
@@ -284,3 +404,205 @@ def test_rocm_mxfp8_quantizer_matches_torch_on_degenerate_blocks(num_rows):
     torch.testing.assert_close(
         q_triton.view(torch.uint8), q_torch.view(torch.uint8), rtol=0, atol=0
     )
+
+
+# --- tiny-M fused quant+GEMM (rocm_mxfp8_block32_gemm_quant) -----------------
+#
+# Serving block32 linears at TP=4 that hit the fused path at decode,
+# covering the K-structure variety: the multi-band K=5120 shape with the
+# in-launch split-K tier (shared_gate_up), a second K=5120 tier, the
+# masked-K-tail K=1280 (EVEN_K=False), a small-K tile, and a draft-model
+# shape absent from the tuned table (_default_config tile).
+FUSED_QUANT_SERVING_SHAPES = [
+    ("shared_gate_up_tp4", 1152, 5120),
+    ("fused_wqa_wkv", 1792, 5120),
+    ("wq_b_tp4", 8192, 1280),
+    ("shared_down_tp4", 5120, 576),
+    ("default_config", 1024, 2048),
+]
+
+# Decode single, the DSpark k=5 hot M, and the M <= 8 tier bound.
+FUSED_QUANT_MS = [1, 5, 8]
+
+
+def _make_block32_weight(n: int, k: int, device: str, seed: int = 0):
+    """e4m3 weight [n, k] with compact [n//32, k//32] E8M0 block scales."""
+    torch.manual_seed(seed)
+    blocks = (torch.randn(n, k, device=device) * 0.05).view(n // 32, 32, k // 32, 32)
+    amax = blocks.abs().amax(dim=(1, 3)).clamp(min=torch.finfo(torch.float32).tiny)
+    bits = (torch.ceil(torch.log2(amax / 448.0)) + 127.0).clamp(0, 254)
+    weight = (blocks * torch.exp2(127.0 - bits)[:, None, :, None]).view(n, k)
+    return (
+        weight.to(torch.float8_e4m3fn).contiguous(),
+        bits.to(torch.uint8).contiguous(),
+    )
+
+
+def _stock(x, weight, weight_scale, out_dtype=torch.bfloat16):
+    """The pre-fusion two-launch path, verbatim."""
+    x_q, x_scale = mxfp8_e4m3_quantize(x)
+    return rocm_mxfp8_block32_gemm(x_q, x_scale, weight, weight_scale, out_dtype)
+
+
+def _fused_quant_hit(m, n, k) -> bool:
+    from vllm.model_executor.kernels.linear.mxfp8 import rocm_block32_gemm as g
+
+    return (m, n, k) in g._FUSED_QUANT_HIT_SHAPES
+
+
+@pytest.mark.parametrize("name,n,k", FUSED_QUANT_SERVING_SHAPES, ids=lambda v: str(v))
+@pytest.mark.parametrize("m", FUSED_QUANT_MS)
+@torch.inference_mode()
+def test_rocm_mxfp8_fused_quant_bit_exact_serving_shapes(name, n, k, m):
+    """The fused output equals the stock two-launch path byte-for-byte at the
+    serving shapes, and a high-precision QDQ fp32 reference within tolerance.
+
+    The reference quantizes the activation with the same quantizer and then
+    dequantizes both operands and contracts in fp32, so it is independent of
+    the block32 GEMM kernel under test.
+    """
+    torch.manual_seed(m)
+    weight, weight_scale = _make_block32_weight(n, k, "cuda")
+    x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16) * 0.5
+
+    expected = _stock(x, weight, weight_scale)
+    out = rocm_mxfp8_block32_gemm_quant(x, weight, weight_scale, torch.bfloat16)
+    out2 = rocm_mxfp8_block32_gemm_quant(x, weight, weight_scale, torch.bfloat16)
+
+    assert _fused_quant_hit(m, n, k), f"{name}: fused fastpath did not engage at M={m}"
+    assert torch.equal(out, expected), f"{name}: fused output differs from stock"
+    assert torch.equal(out, out2), f"{name}: fused output not deterministic"
+
+    x_q, x_scale = mxfp8_e4m3_quantize(x)
+    row_scale = weight_scale.repeat_interleave(32, dim=0)
+    reference = (
+        dequant_mxfp8_to_bf16(x_q, x_scale).float()
+        @ dequant_mxfp8_to_bf16(weight, row_scale).float().T
+    )
+    rel = (out.float() - reference).norm() / reference.norm()
+    assert rel < 5e-3, f"{name}: fused output vs QDQ reference, rel err {rel:.4f}"
+
+
+@torch.inference_mode()
+def test_rocm_mxfp8_fused_quant_bit_exact_degenerate_blocks():
+    """All-zero / denormal / huge / negative blocks quantize identically.
+
+    An all-zero 32-block clamps the E8M0 scale to byte 0 (2**-127,
+    subnormal in fp32): the reciprocal-form rescale must keep it NaN-free,
+    exactly as the separate quantizer does.
+    """
+    n, k, m = 512, 1024, 6
+    weight, weight_scale = _make_block32_weight(n, k, "cuda", seed=7)
+    x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16) * 0.5
+    x[0] = 0.0  # all-zero row
+    x[1, :32] = 0.0  # single zero block
+    x[1, 32:64] = 1e-38  # denormal block
+    x[2] = 3.0e30  # huge: scale byte clamps at 254
+    x[3] = -x[3]  # negative mirror
+    x[4, ::32] = 0.0  # zero at every block start
+
+    expected = _stock(x, weight, weight_scale)
+    out = rocm_mxfp8_block32_gemm_quant(x, weight, weight_scale, torch.bfloat16)
+    assert _fused_quant_hit(m, n, k)
+    assert torch.isfinite(out.float()).all(), "non-finite fused output"
+    assert torch.equal(out, expected), "degenerate blocks differ from stock"
+
+
+@torch.inference_mode()
+def test_rocm_mxfp8_fused_quant_fallbacks_match_stock():
+    """Ineligible inputs take the internal stock path, bit-identical."""
+    n, k = 1024, 512
+    weight, weight_scale = _make_block32_weight(n, k, "cuda", seed=11)
+
+    # Non-contiguous bf16 x (stride-2 view).
+    base = torch.randn(6, 2 * k, device="cuda", dtype=torch.bfloat16) * 0.5
+    x_nc = base[:, ::2]
+    assert not x_nc.is_contiguous()
+    out = rocm_mxfp8_block32_gemm_quant(x_nc, weight, weight_scale, torch.bfloat16)
+    assert not _fused_quant_hit(6, n, k), (
+        "non-contiguous x must not take the fused path"
+    )
+    assert torch.equal(out, _stock(x_nc.contiguous(), weight, weight_scale))
+
+    # fp32 x falls back (quantizer loads any float dtype, fused does not).
+    x_f32 = torch.randn(6, k, device="cuda", dtype=torch.float32) * 0.5
+    out = rocm_mxfp8_block32_gemm_quant(x_f32, weight, weight_scale, torch.float32)
+    assert not _fused_quant_hit(6, n, k), "fp32 x must not take the fused path"
+    assert torch.equal(out, _stock(x_f32, weight, weight_scale, torch.float32))
+
+    # M > 8 (prefill-class) falls back.
+    x_big = torch.randn(16, k, device="cuda", dtype=torch.bfloat16) * 0.5
+    out = rocm_mxfp8_block32_gemm_quant(x_big, weight, weight_scale, torch.bfloat16)
+    assert not _fused_quant_hit(16, n, k), "M=16 must not take the fused path"
+    assert torch.equal(out, _stock(x_big, weight, weight_scale))
+
+    # M = 0 returns an empty output.
+    x0 = torch.randn(0, k, device="cuda", dtype=torch.bfloat16)
+    out = rocm_mxfp8_block32_gemm_quant(x0, weight, weight_scale, torch.bfloat16)
+    assert out.shape == (0, n)
+
+
+@torch.inference_mode()
+def test_rocm_mxfp8_fused_quant_per_row_scale_unaffected():
+    """Regression guard: the per-row-scale branch of
+    RocDotScaledMxfp8LinearKernel.apply_weights (weight_scale [N, K//32],
+    NOT compacted to block32) must keep working when the block32 branch
+    dispatches to the fused op. The numerics of this branch are owned by
+    test_rocm_mxfp8_linear_matches_dequant_reference above; this test pins
+    the dispatch interaction only.
+    """
+    n, k, m = 1152, 5120, 7
+    torch.manual_seed(23)
+    weight_bf16 = torch.randn(n, k, device="cuda", dtype=torch.bfloat16) * 0.05
+    weight, weight_scale = _mxfp8_e4m3_quantize_torch(weight_bf16)
+    assert weight_scale.shape == (n, k // 32), "expected per-row scales"
+
+    layer = torch.nn.Module()
+    layer.weight = torch.nn.Parameter(weight, requires_grad=False)
+    layer.weight_scale = torch.nn.Parameter(weight_scale, requires_grad=False)
+    kernel = RocmDotScaledMxfp8LinearKernel(Mxfp8LinearLayerConfig(weight_shape=(n, k)))
+    kernel.process_weights_after_loading(layer)
+    # Non-uniform per-row scales are not compacted to [N//32, K//32].
+    assert layer.weight_scale.shape[0] == n
+
+    x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16) * 0.5
+    out = kernel.apply_weights(layer, x)
+    assert out.shape == (m, n)
+    assert torch.isfinite(out.float()).all()
+    # This branch never reaches the fused op.
+    assert not _fused_quant_hit(m, n, k), (
+        "per-row-scale layer must not hit the fused path"
+    )
+
+
+@pytest.mark.parametrize("warm_counters", [True, False])
+@torch.inference_mode()
+def test_rocm_mxfp8_fused_quant_cuda_graph_capture(warm_counters):
+    """Capture+replay is bit-exact with and without warm split-K counters.
+
+    shared_gate_up TP4 at M=6 uses the in-launch split-K config
+    (fused_splits=3): with warmup the counters exist and the splits meet in
+    one launch; captured first-use on a fresh stream has no counters and
+    must fall back to the two-launch split-K (partials + reduce).
+    """
+    n, k, m = 1152, 5120, 6
+    weight, weight_scale = _make_block32_weight(n, k, "cuda", seed=17)
+    x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16) * 0.5
+    expected = _stock(x, weight, weight_scale)
+
+    stream = torch.cuda.Stream()
+    if warm_counters:
+        with torch.cuda.stream(stream):
+            warm = rocm_mxfp8_block32_gemm_quant(
+                x, weight, weight_scale, torch.bfloat16
+            )
+        stream.synchronize()
+        assert torch.equal(warm, expected)
+
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph, stream=stream):
+        out = rocm_mxfp8_block32_gemm_quant(x, weight, weight_scale, torch.bfloat16)
+    for _ in range(3):
+        graph.replay()
+        torch.accelerator.synchronize()
+        assert torch.equal(out, expected), "replay differs from stock"
diff --git a/tests/kernels/quantization/test_silu_mul_nvfp4_quant.py b/tests/kernels/quantization/test_silu_mul_nvfp4_quant.py
index 02ab4f009fa8..8c0f931a2c98 100644
--- a/tests/kernels/quantization/test_silu_mul_nvfp4_quant.py
+++ b/tests/kernels/quantization/test_silu_mul_nvfp4_quant.py
@@ -76,3 +76,44 @@ def test_silu_mul_nvfp4_quant(
     torch.testing.assert_close(
         ref_output_dequant, fused_output_dequant, atol=atol, rtol=rtol
     )
+
+
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("shape", [(90, 160), (1, 160), (130, 96), (128, 256)])
+@torch.inference_mode()
+def test_silu_mul_nvfp4_quant_scale_padding_initialized(
+    default_vllm_config,
+    dtype: torch.dtype,
+    shape: tuple[int, int],
+) -> None:
+    """The kernel must overwrite dirty scale padding (rows and columns)
+    without relying on a zero-initialized allocation."""
+    set_random_seed(42)
+    device = "cuda:0"
+    torch.set_default_device(device)
+
+    m, two_d = shape
+    d = two_d // 2
+    round_up = lambda x, y: (x + y - 1) // y * y
+    rounded_m = round_up(m, 128)
+    rounded_scale_n = round_up(d // BLOCK_SIZE, 4)
+
+    x = torch.randn(shape, dtype=dtype)
+    global_scale = torch.tensor(1.0, dtype=torch.float32)
+
+    out = torch.empty((m, d // 2), dtype=FP4_DTYPE)
+    out_scale = torch.full(
+        (rounded_m, rounded_scale_n // 4), 0x7F7F7F7F, dtype=torch.int32
+    )
+    torch.ops._C.silu_and_mul_nvfp4_quant(out, out_scale, x, global_scale)
+
+    tmp = out_scale.view(FP8_DTYPE).reshape(
+        1, rounded_m // 128, rounded_scale_n // 4, 32, 4, 4
+    )
+    full_scale = tmp.permute(0, 1, 4, 3, 2, 5).reshape(rounded_m, rounded_scale_n)
+    full_scale = full_scale.to(torch.float32)
+
+    assert torch.count_nonzero(full_scale[m:, :]) == 0
+    assert torch.count_nonzero(full_scale[:m, d // BLOCK_SIZE :]) == 0
+    # Valid entries must have been written (no leftover 0x7F == NaN fill).
+    assert not torch.isnan(full_scale[:m, : d // BLOCK_SIZE]).any()
diff --git a/tests/kernels/quantization/test_triton_w4a16.py b/tests/kernels/quantization/test_triton_w4a16.py
index bd8bd3cae23a..cca3372adc94 100644
--- a/tests/kernels/quantization/test_triton_w4a16.py
+++ b/tests/kernels/quantization/test_triton_w4a16.py
@@ -319,13 +319,15 @@ class DummyLayer(torch.nn.Module):
         ),
     )
 
-    kernel.process_weights_after_loading(layer)
-
-    # Expected transformed layouts.
+    # Expected transformed layouts, taken before processing: `replace_parameter`
+    # writes a same-size result into the loaded parameter's own storage, which
+    # the checkpoint tensors above still alias.
     expected_w_kn8 = _pack_int4_along_n(w_int4_kn)  # [K, N//8]
     expected_scales_gn = scales_ckpt_nkg.t().contiguous()  # [K//G, N]
     expected_zeros_gn8 = zeros_ckpt_n8kg.t().contiguous()  # [K//G, N//8]
 
+    kernel.process_weights_after_loading(layer)
+
     assert tuple(layer.weight_packed.shape) == (K, N // 8)
     assert tuple(layer.weight_scale.shape) == (K // G, N)
     assert tuple(layer.weight_zero_point.shape) == (K // G, N // 8)
@@ -488,3 +490,34 @@ def fake_gemm(*, a, b_q, scales, qzeros, group_size, zp_bias):
     assert output.shape == (2, N)
     assert captured["qzeros"] is None
     assert captured["zp_bias"] == scalar_types.uint4b8.bias
+
+
+@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm only")
+@pytest.mark.parametrize(
+    "K,N",
+    [
+        (256, 34816),  # M * N passes 2**31: output offsets
+        (17408, 256),  # M * K passes 2**31: activation offsets
+    ],
+)
+def test_triton_w4a16_gemm_rows_past_int32_offsets(K, N):
+    """Rows whose element offset in A or C does not fit in int32."""
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA/HIP device not available")
+    M = 2**31 // max(K, N) + 128
+
+    set_random_seed(0)
+    G = 128
+    a = (0.25 * torch.randn((M, K), device=device)).to(torch.bfloat16)
+    w_int4 = torch.randint(0, 16, (K, N), device=device, dtype=torch.int32)
+    b_packed = _pack_int4_along_n(w_int4)
+    scales = (0.05 * torch.rand((K // G, N), device=device)).to(torch.bfloat16)
+
+    out = triton_w4a16_gemm(
+        a=a, b_q=b_packed, scales=scales, qzeros=None, group_size=G, zp_bias=8
+    )
+    # Dequantize in bfloat16 as the kernel does, so a long K adds no rounding gap.
+    w = (w_int4 - 8).to(torch.bfloat16) * scales.repeat_interleave(G, dim=0)
+    ref = (a[-64:].float() @ w.float()).to(torch.bfloat16)
+
+    torch.testing.assert_close(out[-64:], ref, rtol=1e-2, atol=1e-2)
diff --git a/tests/kernels/quantization/test_zen_da8w4.py b/tests/kernels/quantization/test_zen_da8w4.py
index 924b7a865f4e..c82f1f9dd9b3 100644
--- a/tests/kernels/quantization/test_zen_da8w4.py
+++ b/tests/kernels/quantization/test_zen_da8w4.py
@@ -272,9 +272,9 @@ def test_da8w4_process_weights_layout(mock_zentorch_ops, monkeypatch):
     )
     assert layer._zentorch_da8w4_packed.is_contiguous()
     assert layer._zentorch_da8w4_scale.is_contiguous()
-    # The checkpoint parameters are released once repacked.
-    assert layer.weight_packed.numel() == 0
-    assert layer.weight_scale.numel() == 0
+    # The repacked tensors live in the checkpoint parameters' storage.
+    assert layer._zentorch_da8w4_packed.data_ptr() == layer.weight_packed.data_ptr()
+    assert layer._zentorch_da8w4_scale.data_ptr() == layer.weight_scale.data_ptr()
 
 
 def test_da8w4_apply_weights_matches_dequantized_reference(
diff --git a/tests/kernels/test_bf16_skinny_gemm.py b/tests/kernels/test_bf16_skinny_gemm.py
index 72921830f2ae..9824490147c3 100644
--- a/tests/kernels/test_bf16_skinny_gemm.py
+++ b/tests/kernels/test_bf16_skinny_gemm.py
@@ -11,9 +11,13 @@
 import torch
 from torch import nn
 
+from vllm.model_executor.kernels.linear.cute_dsl import (
+    skinny_gemm_plans as skinny_plans,
+)
 from vllm.model_executor.kernels.linear.cute_dsl.skinny_gemm import (
     SkinnyGemmConfig,
 )
+from vllm.model_executor.layers.utils import maybe_use_skinny_gemm
 from vllm.models.deepseek_v32.nvidia import glm52_low_latency_gemm as glm52_gemm
 from vllm.models.kimi_k3.nvidia import low_latency_gemm as k3_gemm
 from vllm.models.kimi_k3.nvidia.low_latency_gemm import KIMI_K3_PROJECTIONS
@@ -74,6 +78,7 @@
 ]
 
 GLM_GPU_TABLES = (
+    ((10, 0), glm52_gemm.GLM52_PROJECTIONS_SM100),
     ((10, 3), glm52_gemm.GLM52_PROJECTIONS),
     ((10, 7), glm52_gemm.GLM52_PROJECTIONS_SM107),
 )
@@ -480,7 +485,7 @@ def test_glm52_projection_plans_are_separate() -> None:
     [
         ((10, 3), glm52_gemm.GLM52_PROJECTIONS),
         ((10, 7), glm52_gemm.GLM52_PROJECTIONS_SM107),
-        ((10, 0), None),
+        ((10, 0), glm52_gemm.GLM52_PROJECTIONS_SM100),
         ((9, 0), None),
     ],
 )
@@ -592,7 +597,7 @@ def __init__(
         root.attn.fused_qkv_a_proj.quant_method._plan
         == table[(qkv_a.n, qkv_a.k)].build_plan()
     )
-    assert root.attn.q_b_proj.quant_method._plan == q_b.build_plan()
+    assert root.attn.q_b_proj.quant_method._plan == table[(q_b.n, q_b.k)].build_plan()
     if (6144, 4096) in table:
         assert root.attn.o_proj.quant_method._plan == table[(6144, 4096)].build_plan()
     else:
@@ -1030,6 +1035,75 @@ def test_qwen4_exp_sm121_selected_shapes(
     assert cosine > 0.999
 
 
+@pytest.mark.parametrize(
+    "n,k,num_tokens,config",
+    [
+        pytest.param(n, k, num_tokens, config, id=f"{capability}-{n}x{k}-m{num_tokens}")
+        for capability, plans in skinny_plans.SKINNY_GEMM_PLANS_BY_CAPABILITY.items()
+        for (n, k), shape_plans in plans.items()
+        for num_tokens, config in shape_plans.items()
+    ],
+)
+def test_skinny_gemm_plans_are_valid(
+    n: int, k: int, num_tokens: int, config: SkinnyGemmConfig
+) -> None:
+    """Every registered plan must satisfy the wrapper and kernel constraints."""
+    assert config.num_rows == num_tokens
+    assert config.block_size % 32 == 0
+    assert n % config.outputs_per_block == 0
+    k_tile = config.block_size * config.vector_width
+    assert k % k_tile == 0
+    assert config.static_k in (None, k)
+    if config.static_k is not None:
+        assert config.static_k >= 2 * k_tile
+
+
+@pytest.mark.parametrize(
+    "dtype,bias,k,ndim,batch_invariant,expected",
+    [
+        pytest.param(torch.bfloat16, False, 1024, 2, False, True, id="planned"),
+        pytest.param(torch.float16, False, 1024, 2, False, False, id="fp16"),
+        pytest.param(torch.float32, False, 1024, 2, False, False, id="fp32"),
+        pytest.param(torch.bfloat16, True, 1024, 2, False, False, id="bias"),
+        pytest.param(torch.bfloat16, False, 128, 2, False, False, id="unlisted"),
+        pytest.param(torch.bfloat16, False, 1024, 2, True, False, id="batch-invariant"),
+        pytest.param(torch.bfloat16, False, 1024, 3, False, False, id="conv1d"),
+        pytest.param(torch.bfloat16, False, 1024, 4, False, False, id="conv2d"),
+    ],
+)
+def test_maybe_use_skinny_gemm(
+    monkeypatch: pytest.MonkeyPatch,
+    dtype: torch.dtype,
+    bias: bool,
+    k: int,
+    ndim: int,
+    batch_invariant: bool,
+    expected: bool,
+) -> None:
+    """Only eligible layers are marked; CPU tensors suffice for selection."""
+    assert (4096, 1024) in skinny_plans.SM121_SKINNY_GEMM_PLANS
+    assert (4096, 128) not in skinny_plans.SM121_SKINNY_GEMM_PLANS
+    # Isolate hardware/dependency detection, keeping the real plan lookup.
+    monkeypatch.setattr(
+        skinny_plans.current_platform, "get_device_capability", lambda: (12, 1)
+    )
+    monkeypatch.setattr(
+        skinny_plans.shape_dynamic_skinny_gemm, "is_available", lambda: True
+    )
+    monkeypatch.setenv("VLLM_BATCH_INVARIANT", str(int(batch_invariant)))
+    layer: nn.Module
+    if ndim == 2:
+        layer = nn.Linear(k, 4096, bias=bias, dtype=dtype, device="cpu")
+    elif ndim == 3:
+        layer = nn.Conv1d(k, 4096, 1, bias=bias, dtype=dtype, device="cpu")
+    else:
+        layer = nn.Conv2d(k, 4096, 1, bias=bias, dtype=dtype, device="cpu")
+
+    maybe_use_skinny_gemm(layer)
+
+    assert getattr(layer, "use_skinny_gemm", False) is expected
+
+
 def test_glm52_q_b_nonpacked_single_row_falls_back() -> None:
     _require_sm103_and_cute()
     spec = glm52_gemm.GLM52_Q_B_PROJECTION
diff --git a/tests/kernels/test_compressor_kv_cache.py b/tests/kernels/test_compressor_kv_cache.py
index 3cd998a60ec4..3cd8c1a0b578 100644
--- a/tests/kernels/test_compressor_kv_cache.py
+++ b/tests/kernels/test_compressor_kv_cache.py
@@ -31,6 +31,7 @@
 )
 from vllm.models.deepseek_v4.compressor import _get_c128_boundary
 from vllm.platforms import current_platform
+from vllm.triton_utils import tl, triton
 from vllm.v1.attention.backends.mla.compressor_utils import (
     get_dspark_swa_index_width,
 )
@@ -239,6 +240,7 @@ def test_v41_pcp_swa_cache_matches_non_pcp_writer(kv_mxfp8, monkeypatch):
         _flashinfer_fp8_kv_scale=None,
         kv_mxfp8=kv_mxfp8,
         accepts_unnormed_unroped_query=False,
+        _pads_q_by_copy=False,
         n_local_heads=num_heads,
         padded_heads=num_heads,
         eps=1e-20,
@@ -799,6 +801,157 @@ def check(shorter_by):
     check(1)
 
 
+@triton.jit
+def _per_token_nvfp4_gather_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,
+    cache_block_size: tl.constexpr,
+    block_stride: tl.constexpr,
+):
+    """Reference for the tiled NVFP4 gather: one token per iteration, with
+    the e2m1 grid decoded arithmetically."""
+    batch_idx = tl.program_id(0)
+    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, 512)
+    for i in range(tl.program_id(1), gather_len, tl.num_programs(1)):
+        pos = start_pos + i
+        physical_block_idx = tl.load(
+            block_table_ptr + batch_idx * max_blocks_per_seq + pos // cache_block_size
+        )
+        page = k_cache_ptr + physical_block_idx.to(tl.int64) * block_stride
+        token_ptr = page + (pos % cache_block_size) * 288
+        packed = tl.load(token_ptr + tl.arange(0, 256))
+        codes = tl.interleave((packed & 0xF).to(tl.int32), (packed >> 4).to(tl.int32))
+        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(token_ptr + 256 + tl.arange(0, 32))
+        scale = sf.to(tl.float8e4nv, bitcast=True).to(tl.float32)
+        dequant = tl.reshape(tl.reshape(vals, (32, 16)) * scale[:, None], (512,))
+        out_row_ptr = out_ptr + batch_idx * out_stride0 + (offset + i) * out_stride1
+        tl.store(out_row_ptr + d, dequant.to(tl.bfloat16))
+
+
+def _per_token_nvfp4_gather(
+    out, k_cache, seq_lens, gather_lens, block_table, block_size, offset
+):
+    _per_token_nvfp4_gather_kernel[(seq_lens.shape[0], 128)](
+        out,
+        out.stride(0),
+        out.stride(1),
+        k_cache,
+        seq_lens,
+        block_table,
+        offset,
+        gather_lens,
+        max_blocks_per_seq=block_table.shape[-1],
+        cache_block_size=block_size,
+        block_stride=k_cache.stride(0),
+    )
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only")
+@pytest.mark.parametrize(
+    "block_size,num_reqs", [(64, 1), (128, 3), (64, 17), (128, 64)]
+)
+@pytest.mark.parametrize("partial", [False, True])
+def test_v41_nvfp4_gather_ragged_batch(block_size, num_reqs, partial):
+    """The NVFP4 gather writes the same bits as the per-token reference on
+    arbitrary cache bytes (NaN, subnormal and negative scales included) for
+    ragged batches: empty requests, partial last blocks and a long context, a
+    shuffled block table holding -1 in every slot the gather must not read,
+    and an offset into a wider strided workspace."""
+    from vllm.models.deepseek_v41.common.ops import dequantize_and_gather_k_cache
+
+    torch.manual_seed(19)
+    device = "cuda"
+    if num_reqs == 1:
+        seq_lens_list = [65_541]
+    else:
+        edges = [0, 1, block_size - 1, block_size, block_size + 1, 3 * block_size + 5]
+        seq_lens_list = (edges + torch.randint(0, 5000, (num_reqs,)).tolist())[
+            :num_reqs
+        ]
+    gather_lens_list = (
+        [int(torch.randint(0, n + 1, ())) for n in seq_lens_list] if partial else None
+    )
+    offset = 7 if partial else 0
+
+    blocks_per_req = [math.ceil(n / block_size) for n in seq_lens_list]
+    num_blocks = sum(blocks_per_req) + 1
+    # Padded pages, so the block stride is not block_size * 288.
+    page_bytes = block_size * 288 + 512
+    backing = torch.randint(
+        0, 256, (num_blocks, page_bytes), dtype=torch.uint8, device=device
+    )
+    cache = backing.as_strided((num_blocks, block_size, 288), (page_bytes, 288, 1))
+    pages = iter(torch.randperm(num_blocks).tolist())
+    block_table = torch.full((num_reqs, max(blocks_per_req) + 2), -1, dtype=torch.int32)
+    for req, seq_len in enumerate(seq_lens_list):
+        start = 0 if gather_lens_list is None else seq_len - gather_lens_list[req]
+        for block in range(start // block_size, blocks_per_req[req]):
+            block_table[req, block] = next(pages)
+    block_table = block_table.to(device)
+    seq_lens = torch.tensor(seq_lens_list, dtype=torch.int32, device=device)
+    gather_lens = (
+        None
+        if gather_lens_list is None
+        else torch.tensor(gather_lens_list, dtype=torch.int32, device=device)
+    )
+
+    width = max(seq_lens_list) + offset
+    storage = torch.full(
+        (num_reqs, width + 9, 512), -123.0, dtype=torch.bfloat16, device=device
+    )
+    args = (cache, seq_lens, gather_lens, block_table)
+    dequantize_and_gather_k_cache(
+        storage[:, :width], *args, block_size=block_size, offset=offset
+    )
+    expected = torch.full_like(storage, -123.0)
+    _per_token_nvfp4_gather(expected[:, :width], *args, block_size, offset)
+    assert torch.equal(storage.view(torch.int16), expected.view(torch.int16))
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only")
+def test_v41_nvfp4_gather_decodes_every_code_and_scale():
+    """Every (packed byte, e4m3 scale byte) pair: token t holds the bytes
+    0..255, and scale byte t for all of its 32 groups."""
+    from vllm.models.deepseek_v41.common.ops import dequantize_and_gather_k_cache
+
+    device, block_size = "cuda", 64
+    codes = torch.arange(256, device=device).to(torch.uint8)
+    records = torch.empty(256, 288, dtype=torch.uint8, device=device)
+    records[:, :256] = codes
+    records[:, 256:] = codes[:, None]
+    args = (
+        records.view(256 // block_size, block_size, 288),
+        torch.tensor([256], dtype=torch.int32, device=device),
+        None,
+        torch.arange(4, dtype=torch.int32, device=device)[None],
+        block_size,
+        0,
+    )
+    out = torch.empty(1, 256, 512, dtype=torch.bfloat16, device=device)
+    expected = torch.empty_like(out)
+    dequantize_and_gather_k_cache(out, *args)
+    _per_token_nvfp4_gather(expected, *args)
+    assert torch.equal(out.view(torch.int16), expected.view(torch.int16))
+
+
 @pytest.mark.skipif(
     not current_platform.is_device_capability_family(100),
     reason="The V4.1 records are SM100-only",
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 a14c889caae1..b3b856283b76 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
@@ -715,6 +715,62 @@ def run(q_in, is_q_interleaved):
     torch.testing.assert_close(cache_inter, cache_major, rtol=0, atol=0)
 
 
+@pytest.mark.parametrize("num_tokens", [17, 1024, 2049])
+@pytest.mark.parametrize("dp_pad", [0, 5])
+def test_q_interleaved_pad_copy_matches_fused_insert(num_tokens: int, dp_pad: int):
+    """A KV-only launch plus ``_pad_interleaved_q`` equals the fused launch.
+
+    Same cache bytes and same padded Q bits; the copy writes into a NaN-filled
+    buffer, and the last ``dp_pad`` rows have no slot, as under DP padding.
+    """
+    from vllm.models.deepseek_v41.attention import _pad_interleaved_q
+
+    torch.manual_seed(8)
+    device = "cuda"
+    dtype = torch.bfloat16
+    block_size, n_heads, q_head_padded = 64, 32, 64
+    total = num_tokens + dp_pad
+
+    q = torch.randn(total, n_heads, HEAD_DIM, dtype=dtype, device=device)
+    kv = torch.randn(total, HEAD_DIM, dtype=dtype, device=device)
+    positions = torch.arange(total, 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 insert(pad_to: int):
+        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,
+            kv,
+            cache,
+            slot_mapping,
+            positions,
+            cos_sin_cache,
+            pad_to,
+            1e-6,
+            block_size,
+            False,  # apply_q_norm
+            True,  # kv_mxfp8
+            False,  # apply_q_rope
+            True,  # is_q_interleaved
+        )
+        return out, cache
+
+    want, want_cache = insert(q_head_padded)
+    _, cache = insert(0)
+    assert want_cache.any()
+    torch.testing.assert_close(cache, want_cache, rtol=0, atol=0)
+
+    got = torch.full_like(want, float("nan"))
+    _pad_interleaved_q(q, got)
+    torch.testing.assert_close(
+        got.view(torch.int16), want.view(torch.int16), rtol=0, atol=0
+    )
+
+
 # ── Test 2b: DP padding (slot_mapping shorter than q/kv) ─────────────────────
 
 
diff --git a/tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py b/tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py
index d6a96c6c9960..711e4d2bda9e 100644
--- a/tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py
+++ b/tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py
@@ -156,6 +156,74 @@ def test_dense_norm_rope(num_tokens, num_heads, num_kv_heads):
     torch.testing.assert_close(v_out, v_in, rtol=0, atol=0)
 
 
+@pytest.mark.parametrize("num_tokens", [1, 7, 513])
+@pytest.mark.parametrize("block_size", [16, 128])
+@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"])
+def test_dense_insert_matches_reshape_and_cache_flash(
+    num_tokens, block_size, kv_cache_dtype
+):
+    """A dense layer's insert stores exactly what reshape_and_cache_flash
+    writes from the normed/roped k (unit scales), and leaves q/k/v as the
+    insert-free call does."""
+    torch.manual_seed(0)
+    device, dtype, eps = "cuda", torch.bfloat16, 1e-6
+    num_heads, num_kv_heads, max_pos = 8, 2, 4096
+    q_w = torch.randn(HEAD_DIM, dtype=dtype, device=device) * 0.1
+    k_w = torch.randn(HEAD_DIM, dtype=dtype, device=device) * 0.1
+    cos_sin = make_cos_sin_cache(max_pos, ROTARY_DIM, 5_000_000.0, dtype, device)
+    positions = torch.randint(
+        0, max_pos, (num_tokens,), dtype=torch.int64, device=device
+    )
+    qsz, kvsz = num_heads * HEAD_DIM, num_kv_heads * HEAD_DIM
+    qkv = torch.randn(num_tokens, qsz + 2 * kvsz, dtype=dtype, device=device)
+    qkv_ref = qkv.clone()
+    num_blocks = (num_tokens + block_size - 1) // block_size + 2
+    slot_mapping = torch.randperm(num_blocks * block_size, device=device)[
+        :num_tokens
+    ].to(torch.int64)
+    slot_mapping[::4] = -1  # padding: no write
+    cache_dtype = torch.uint8 if kv_cache_dtype == "fp8" else dtype
+    kv_cache = torch.zeros(
+        num_blocks,
+        num_kv_heads,
+        block_size,
+        2 * HEAD_DIM,
+        dtype=cache_dtype,
+        device=device,
+    )
+    ref_cache = torch.zeros_like(kv_cache)
+
+    args = (qkv, q_w, k_w, cos_sin, positions, num_heads, num_kv_heads)
+    ops.fused_minimax_m3_qknorm_rope_kv_insert(
+        *args,
+        ROTARY_DIM,
+        eps,
+        slot_mapping=slot_mapping,
+        kv_cache=kv_cache,
+        block_size=block_size,
+        kv_cache_dtype=kv_cache_dtype,
+    )
+    ops.fused_minimax_m3_qknorm_rope_kv_insert(
+        qkv_ref, *args[1:], ROTARY_DIM, eps, kv_cache_dtype="auto"
+    )
+    assert torch.equal(qkv, qkv_ref)
+
+    _, k, v = qkv_ref.split([qsz, kvsz, kvsz], dim=-1)
+    key_cache, value_cache = ref_cache.transpose(1, 2).split(HEAD_DIM, dim=-1)
+    one = torch.ones((), dtype=torch.float32, device=device)
+    ops.reshape_and_cache_flash(
+        k.reshape(num_tokens, num_kv_heads, HEAD_DIM),
+        v.reshape(num_tokens, num_kv_heads, HEAD_DIM),
+        key_cache,
+        value_cache,
+        slot_mapping,
+        kv_cache_dtype,
+        one,
+        one,
+    )
+    assert torch.equal(kv_cache, ref_cache)
+
+
 # ── Test 2: sparse mode (full: index branch + cache inserts) ─────────────────
 
 
@@ -756,3 +824,249 @@ def test_sparse_full_nvfp4(num_tokens, num_kv_heads, block_size, kv_scales):
     torch.testing.assert_close(
         index_cache.view(-1, HEAD_DIM), expected_index_cache, rtol=0, atol=0
     )
+
+
+@pytest.mark.skipif(
+    not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10,
+    reason="The NVFP4 reshape_and_cache_flash reference needs SM100",
+)
+@pytest.mark.parametrize("num_tokens", [1, 7, 513])
+@pytest.mark.parametrize("kv_cache_dtype", ["auto", "nvfp4"])
+@pytest.mark.parametrize("page_pad", [0, 48])
+def test_sparse_full_nvfp4_index(num_tokens, kv_cache_dtype, page_pad):
+    """The NVFP4 index-K insert writes MiniMax-M3.1 KV4 bytes as a single-head
+    NVFP4 K side with no global scale (E2M1 page, then linear E4M3 scales: the
+    MSA Q8KV4 page), honours a padded page stride and skips padded slots, while
+    index-Q and everything else match the fp8 index path exactly."""
+    torch.manual_seed(6)
+    device, dtype, eps = "cuda", torch.bfloat16, 1e-6
+    base, max_pos = 5_000_000.0, 4096
+    num_heads, num_kv_heads, num_idx_heads, block_size = 16, 4, 4, 128
+    page_width = HEAD_DIM // 2 + HEAD_DIM // 16
+
+    weights = [
+        torch.randn(HEAD_DIM, dtype=dtype, device=device) * 0.1 for _ in range(4)
+    ]
+    cos_sin = make_cos_sin_cache(max_pos, ROTARY_DIM, base, dtype, device)
+    positions = torch.randint(
+        0, max_pos, (num_tokens,), dtype=torch.int64, device=device
+    )
+    qsz, kvsz = num_heads * HEAD_DIM, num_kv_heads * HEAD_DIM
+    iqsz = num_idx_heads * HEAD_DIM
+    qkv0 = torch.randn(
+        num_tokens, qsz + 2 * kvsz + iqsz + HEAD_DIM, dtype=dtype, device=device
+    )
+    num_blocks = (num_tokens + block_size - 1) // block_size + 1
+    slot_mapping = torch.randperm(
+        num_blocks * block_size, dtype=torch.int64, device=device
+    )[:num_tokens]
+    index_slot_mapping = torch.randperm(
+        num_blocks * block_size, dtype=torch.int64, device=device
+    )[:num_tokens]
+    index_slot_mapping[::5] = -1  # padded / unscheduled tokens are not written
+    nvfp4_kwargs = {}
+    if kv_cache_dtype == "nvfp4":
+        nvfp4_kwargs = dict(
+            kv_k_scale=torch.tensor(0.5, dtype=torch.float32, device=device),
+            kv_v_scale=torch.tensor(2.0, dtype=torch.float32, device=device),
+        )
+
+    def run(index_cache):
+        qkv = qkv0.clone()
+        if kv_cache_dtype == "nvfp4":
+            kv_cache = torch.zeros(
+                num_blocks,
+                2 * num_kv_heads,
+                block_size,
+                page_width,
+                dtype=torch.uint8,
+                device=device,
+            )
+        else:
+            kv_cache = torch.zeros(
+                num_blocks,
+                num_kv_heads,
+                block_size,
+                2 * HEAD_DIM,
+                dtype=dtype,
+                device=device,
+            )
+        index_q = torch.empty(
+            num_tokens, iqsz, dtype=torch.float8_e4m3fn, device=device
+        )
+        ops.fused_minimax_m3_qknorm_rope_kv_insert(
+            qkv,
+            weights[0],
+            weights[1],
+            cos_sin,
+            positions,
+            num_heads,
+            num_kv_heads,
+            ROTARY_DIM,
+            eps,
+            weights[2],
+            weights[3],
+            num_idx_heads,
+            slot_mapping,
+            index_slot_mapping,
+            kv_cache,
+            index_cache,
+            block_size,
+            torch.empty(num_tokens, qsz, dtype=dtype, device=device),
+            index_q,
+            kv_cache_dtype,
+            **nvfp4_kwargs,
+        )
+        return qkv, kv_cache, index_q
+
+    fp8_index_cache = torch.zeros(
+        num_blocks, block_size, HEAD_DIM, dtype=torch.float8_e4m3fn, device=device
+    )
+    qkv_fp8, kv_cache_fp8, index_q_fp8 = run(fp8_index_cache)
+    raw = torch.zeros(
+        num_blocks,
+        block_size * page_width + page_pad,
+        dtype=torch.uint8,
+        device=device,
+    )
+    index_cache = raw[:, : block_size * page_width].view(
+        num_blocks, block_size, page_width
+    )
+    qkv, kv_cache, index_q = run(index_cache)
+
+    assert torch.equal(qkv, qkv_fp8)
+    assert torch.equal(kv_cache, kv_cache_fp8)
+    assert torch.equal(index_q.view(torch.uint8), index_q_fp8.view(torch.uint8))
+
+    # A one-head NVFP4 page's K side is [bs, 64] E2M1 then [bs, 8] linear
+    # scales, which is exactly the index page.
+    index_k = qkv.split([qsz, kvsz, kvsz, iqsz, HEAD_DIM], dim=-1)[-1]
+    expected = nvfp4_pages(
+        [(index_k.reshape(num_tokens, 1, HEAD_DIM), 1.0)],
+        index_slot_mapping,
+        num_blocks,
+        block_size,
+    )
+    assert torch.equal(index_cache, expected[:, 0]), (
+        f"{int((index_cache != expected[:, 0]).sum())} NVFP4 index bytes differ"
+    )
+    assert not raw[:, block_size * page_width :].any()
+
+
+# ── Test 7: write_qkv=False, for callers that read only the FP8 q ────────────
+
+
+@pytest.mark.skipif(
+    not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10,
+    reason="NVFP4 KV caches need SM100",
+)
+@pytest.mark.parametrize("num_tokens", [1, 7, 513])
+@pytest.mark.parametrize("num_heads,num_kv_heads", [(64, 4), (32, 2), (16, 1)])
+@pytest.mark.parametrize("index_cache_dtype", ["nvfp4", "fp8"])
+@pytest.mark.parametrize("scales", [(1.0, 1.0, 1.0), (0.5, 2.0, 0.5)])
+def test_write_qkv_false(
+    num_tokens, num_heads, num_kv_heads, index_cache_dtype, scales
+):
+    """write_qkv=False without q_out writes the same q_fp8_out, index_q_out and
+    cache bytes as a default call, and leaves qkv either untouched or as the
+    default call rewrites it (k and index_k need not be rewritten)."""
+    torch.manual_seed(7)
+    device, dtype, eps = "cuda", torch.bfloat16, 1e-6
+    base, max_pos = 5_000_000.0, 4096
+    num_idx_heads, block_size = num_kv_heads, 128
+    page_width = HEAD_DIM // 2 + HEAD_DIM // 16
+    k_scale, v_scale, q_scale = scales
+
+    weights = [
+        torch.randn(HEAD_DIM, dtype=dtype, device=device) * 0.1 for _ in range(4)
+    ]
+    cos_sin = make_cos_sin_cache(max_pos, ROTARY_DIM, base, dtype, device)
+    positions = torch.randint(
+        0, max_pos, (num_tokens,), dtype=torch.int64, device=device
+    )
+    qsz, kvsz = num_heads * HEAD_DIM, num_kv_heads * HEAD_DIM
+    iqsz = num_idx_heads * HEAD_DIM
+    qkv0 = torch.randn(
+        num_tokens, qsz + 2 * kvsz + iqsz + HEAD_DIM, dtype=dtype, device=device
+    )
+    qkv0[:, qsz + kvsz : qsz + 2 * kvsz] = nvfp4_edge_values(
+        num_tokens * num_kv_heads, v_scale
+    ).view(num_tokens, kvsz)
+    num_blocks = (num_tokens + block_size - 1) // block_size + 1
+    slot_mapping = torch.randperm(
+        num_blocks * block_size, dtype=torch.int64, device=device
+    )[:num_tokens]
+    slot_mapping[::5] = -1
+    index_slot_mapping = torch.randperm(
+        num_blocks * block_size, dtype=torch.int64, device=device
+    )[:num_tokens]
+    kv_cache0 = torch.randint(
+        0,
+        256,
+        (num_blocks, 2 * num_kv_heads, block_size, page_width),
+        dtype=torch.uint8,
+        device=device,
+    )
+    if index_cache_dtype == "nvfp4":
+        index_cache0 = torch.randint(
+            0,
+            256,
+            (num_blocks, block_size, page_width),
+            dtype=torch.uint8,
+            device=device,
+        )
+    else:
+        index_cache0 = torch.randn(num_blocks, block_size, HEAD_DIM, device=device).to(
+            torch.float8_e4m3fn
+        )
+
+    def run(write_qkv):
+        qkv = qkv0.clone()
+        kv_cache, index_cache = kv_cache0.clone(), index_cache0.clone()
+        q_fp8 = torch.empty(num_tokens, qsz, dtype=torch.float8_e4m3fn, device=device)
+        index_q = torch.empty(
+            num_tokens, iqsz, dtype=torch.float8_e4m3fn, device=device
+        )
+        ops.fused_minimax_m3_qknorm_rope_kv_insert(
+            qkv,
+            weights[0],
+            weights[1],
+            cos_sin,
+            positions,
+            num_heads,
+            num_kv_heads,
+            ROTARY_DIM,
+            eps,
+            weights[2],
+            weights[3],
+            num_idx_heads,
+            slot_mapping,
+            index_slot_mapping,
+            kv_cache,
+            index_cache,
+            block_size,
+            torch.empty(num_tokens, qsz, dtype=dtype, device=device)
+            if write_qkv
+            else None,
+            index_q,
+            "nvfp4",
+            q_fp8_out=q_fp8,
+            q_fp8_scale=q_scale,
+            kv_k_scale=torch.tensor(k_scale, dtype=torch.float32, device=device),
+            kv_v_scale=torch.tensor(v_scale, dtype=torch.float32, device=device),
+            write_qkv=write_qkv,
+        )
+        return qkv, q_fp8, index_q, kv_cache, index_cache
+
+    expected = run(True)
+    qkv, *got = run(False)
+    assert torch.equal(qkv.view(torch.int16), qkv0.view(torch.int16)) or (
+        torch.equal(qkv.view(torch.int16), expected[0].view(torch.int16))
+    )
+    for name, g, e in zip(
+        ("q_fp8", "index_q", "kv_cache", "index_cache"), got, expected[1:]
+    ):
+        assert torch.equal(g.view(torch.uint8), e.view(torch.uint8)), (
+            f"{name}: {int((g.view(torch.uint8) != e.view(torch.uint8)).sum())} "
+            "bytes differ"
+        )
diff --git a/tests/kernels/test_fused_qk_norm_rope_gate.py b/tests/kernels/test_fused_qk_norm_rope_gate.py
index 52c587264701..1f57bbeaf856 100644
--- a/tests/kernels/test_fused_qk_norm_rope_gate.py
+++ b/tests/kernels/test_fused_qk_norm_rope_gate.py
@@ -118,3 +118,79 @@ def test_fused_qk_norm_rope_gate_matches_reference(
     torch.testing.assert_close(q_out, q_ref, atol=1e-2, rtol=1e-2)
     torch.testing.assert_close(k_out, k_ref, atol=1e-2, rtol=1e-2)
     torch.testing.assert_close(gate_out, gate_ref, atol=0, rtol=0)
+
+
+@pytest.mark.skipif(
+    not (current_platform.is_cuda_alike() or current_platform.is_xpu()),
+    reason="fused_qk_rmsnorm_rope_gate Triton kernel requires CUDA/ROCm/XPU",
+)
+@torch.inference_mode()
+def test_fused_qk_norm_rope_gate_int64_row_offset(default_vllm_config) -> None:
+    # q_gate and k are views into the packed qkv output, so a token's input
+    # offset is token * row_stride. Rows from 2^31 // row_stride + 1 on
+    # (149,797 here, 209,716 for Qwen3.5-9B) need int64 offsets.
+    num_q_heads, num_kv_heads = 24, 4
+    q_size = num_q_heads * HEAD_DIM
+    kv_size = num_kv_heads * HEAD_DIM
+    row_stride = 2 * q_size + 2 * kv_size
+    num_tokens = 2**31 // row_stride + 8
+    # packed qkv plus the q, k and gate outputs, all bf16.
+    need = num_tokens * (row_stride + 2 * q_size + kv_size) * 2
+    free, _ = current_platform.mem_get_info()
+    if free < need + (1 << 30):
+        pytest.skip("not enough free device memory")
+
+    device = torch.device(
+        current_platform.device_type, torch.accelerator.current_device_index()
+    )
+    torch.set_default_device(device)
+    set_random_seed(SEED)
+
+    qkv = torch.randn(num_tokens, row_stride, dtype=DTYPE, device=device)
+    q_gate, k, _ = qkv.split([2 * q_size, kv_size, kv_size], dim=-1)
+    q_norm = GemmaRMSNorm(HEAD_DIM, eps=RMS_NORM_EPS).to(device, dtype=DTYPE)
+    k_norm = GemmaRMSNorm(HEAD_DIM, eps=RMS_NORM_EPS).to(device, dtype=DTYPE)
+    q_norm.weight.normal_(std=0.1)
+    k_norm.weight.normal_(std=0.1)
+    rope = RotaryEmbedding(
+        HEAD_DIM,
+        ROTARY_DIM,
+        MAX_POSITION_EMBEDDINGS,
+        ROPE_THETA,
+        True,
+        DTYPE,
+    ).to(device)
+    positions = torch.arange(num_tokens, dtype=torch.long, device=device)
+
+    q_out, k_out, gate_out = fused_qk_rmsnorm_rope_gate(
+        q_gate,
+        k,
+        q_norm.weight,
+        k_norm.weight,
+        rope.cos_sin_cache,
+        positions,
+        RMS_NORM_EPS,
+        num_q_heads,
+        num_kv_heads,
+        HEAD_DIM,
+        ROTARY_DIM,
+        norm_beta=1.0,
+    )
+
+    # Reference on the first row and the last rows, past the int32 boundary.
+    rows = torch.tensor([0, *range(num_tokens - 8, num_tokens)], device=device)
+    n = rows.numel()
+    q_gate_heads = q_gate[rows].view(n, num_q_heads, 2 * HEAD_DIM)
+    gate_ref = q_gate_heads[..., HEAD_DIM:].reshape(n, q_size)
+    q_ref = q_norm.forward_native(q_gate_heads[..., :HEAD_DIM])
+    k_ref = k_norm.forward_native(k[rows].view(n, num_kv_heads, HEAD_DIM))
+    assert isinstance(q_ref, torch.Tensor)
+    assert isinstance(k_ref, torch.Tensor)
+    q_ref, k_ref = rope.forward_native(
+        positions[rows], q_ref.reshape(n, q_size), k_ref.reshape(n, kv_size)
+    )
+    assert k_ref is not None
+
+    torch.testing.assert_close(q_out[rows], q_ref, atol=1e-2, rtol=1e-2)
+    torch.testing.assert_close(k_out[rows], k_ref, atol=1e-2, rtol=1e-2)
+    torch.testing.assert_close(gate_out[rows], gate_ref, atol=0, rtol=0)
diff --git a/tests/kernels/test_gemm_rs_ar.py b/tests/kernels/test_gemm_rs_ar.py
index b1c79af2b925..7812e613eec1 100644
--- a/tests/kernels/test_gemm_rs_ar.py
+++ b/tests/kernels/test_gemm_rs_ar.py
@@ -165,6 +165,31 @@ def _run_mode(
         torch.accelerator.synchronize(device)
         _assert_valid_rows_close(actual, expected, M, rank, all_reduce)
 
+        if quant_config is not None:
+            # Producers that already emit MXFP8 (e.g. MiniMax-M3's fused dense
+            # GEMM1) hand in a QuantizedActivation, which skips re-quantization.
+            from vllm.model_executor.layers.fusion.quant_activation import (
+                QuantizedActivation,
+            )
+            from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
+                mxfp8_e4m3_quantize,
+            )
+            from vllm.model_executor.layers.quantization.utils.quant_utils import (
+                kMxfp8Dynamic,
+            )
+
+            x_q, x_sf = mxfp8_e4m3_quantize(x, is_sf_swizzled_layout=True)
+            x_qa = QuantizedActivation(
+                data=x_q,
+                scale=x_sf,
+                orig_dtype=x.dtype,
+                orig_shape=x.shape,
+                quant_key=kMxfp8Dynamic,
+            )
+            actual = gemm_rs_ar.apply(x_qa, projections[K])
+            torch.accelerator.synchronize(device)
+            _assert_valid_rows_close(actual, expected, M, rank, all_reduce)
+
     if all_reduce:
         # AR outputs must remain valid after the symmetric workspace is reused.
         lifetime_M, lifetime_K = 257, 768
diff --git a/tests/kernels/test_minimax_m3_amd_ops.py b/tests/kernels/test_minimax_m3_amd_ops.py
index 868c2fcac7be..d5d8ecf00e90 100644
--- a/tests/kernels/test_minimax_m3_amd_ops.py
+++ b/tests/kernels/test_minimax_m3_amd_ops.py
@@ -6,9 +6,13 @@
 tests assert the two agree within tolerance:
 
   * Gemma RMSNorm (plain + fused-add-residual)  -> fp32 PyTorch normalize
+  * Fused-add Gemma RMSNorm + MXFP8 epilogue     -> norm, then mxfp8_e4m3_quantize
+                                                   (bitwise)
   * SwiGLU-OAI (split layout)                    -> fp32 PyTorch elementwise
   * Fused MXFP8 activation quant (Triton)        -> _mxfp8_e4m3_quantize_torch
   * Native MXFP8 linear (dot_scaled)             -> dequant-to-bf16 @ matmul
+  * Native MXFP8 linear on the fused norm output -> norm, then the linear
+                                                   (bitwise)
   * Native MXFP8 MoE (dot_scaled grouped GEMM)   -> dequant-to-bf16 MoE math
   * Strided K/V sparse-PA cache insert           -> contiguous-copy insert
 
@@ -27,6 +31,7 @@
 
 import pytest
 import torch
+from transformers import PreTrainedConfig
 
 from vllm.platforms import current_platform
 
@@ -39,15 +44,24 @@
     _mxfp8_e4m3_quantize_torch,
     _mxfp8_e4m3_quantize_triton,
     dequant_mxfp8_to_bf16,
+    mxfp8_e4m3_quantize,
+)
+from vllm.models.minimax_m3.amd.model import (  # noqa: E402
+    MiniMAXGemmaRMSNorm,
+    MiniMaxM3SparseAttention,
+    _build_rotary_emb,
+    _can_fuse_input_norm_mxfp8,
+    _kv_insert_operand,
 )
-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_fused_add_rmsnorm_mxfp8,
     gemma_rmsnorm,
     swiglu_oai_quantize_mxfp8,
     swiglu_oai_split,
 )
 from vllm.models.minimax_m3.amd.ops.gemma_rmsnorm import _num_warps  # noqa: E402
+from vllm.platforms.rocm import on_gfx942  # noqa: E402
 
 DEVICE = "cuda"
 EPS = 1e-6
@@ -144,6 +158,61 @@ def test_num_warps_monotonic():
     assert _num_warps(128) <= _num_warps(2048) <= _num_warps(8192)
 
 
+def _mxfp8_edge_case_inputs(n: int, dtype: torch.dtype):
+    """(x, residual, weight) whose normed output covers the MXFP8 block cases.
+
+    Through the weight, every row gets a huge block (scale above 1.0), an
+    all-zero block, and a tiny block (``1 + w == 2**-8``). Rows: 32 ordinary
+    rows with a residual, then an all-zero row, a row of tiny values, a row
+    whose first block dwarfs the rest, and a row of partly-zero blocks.
+    """
+    torch.manual_seed(0)
+    w = torch.randn(n, device=DEVICE, dtype=dtype) * 0.1
+    w[:32] = 300.0
+    w[32:64] = -1.0 + 2**-8
+    w[64:96] = -1.0
+
+    ordinary = torch.randn(32, n, device=DEVICE, dtype=dtype)
+    zero = torch.zeros(1, n, device=DEVICE, dtype=dtype)
+    tiny = torch.full((1, n), torch.finfo(dtype).tiny * 16, device=DEVICE, dtype=dtype)
+    dominated = torch.randn(1, n, device=DEVICE, dtype=dtype) * 1e-2
+    dominated[:, :32] = 1e4
+    partly_zero = torch.randn(1, n, device=DEVICE, dtype=dtype)
+    partly_zero[:, ::2] = 0.0
+    partly_zero[:, 128:192] = 0.0
+    x = torch.cat([ordinary, zero, tiny, dominated, partly_zero])
+
+    res = torch.zeros_like(x)
+    res[:32] = torch.randn(32, n, device=DEVICE, dtype=dtype)
+    return x, res, w
+
+
+@pytest.mark.parametrize("n", [2048, 3072, 6144])
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+@torch.inference_mode()
+def test_gemma_fused_add_rmsnorm_mxfp8_bitwise_eq_unfused(n, dtype):
+    """The MXFP8 epilogue writes the same values, E8M0 scales and residual as
+    ``gemma_fused_add_rmsnorm`` followed by ``mxfp8_e4m3_quantize``."""
+    x, res, w = _mxfp8_edge_case_inputs(n, dtype)
+    m = x.shape[0]
+    # A leading batch dim is flattened into the quantized rows.
+    x, res = x.unsqueeze(0), res.unsqueeze(0)
+
+    got_q, got_scale, got_res = gemma_fused_add_rmsnorm_mxfp8(x, res, w, EPS)
+    normed, ref_res = gemma_fused_add_rmsnorm(x, res, w, EPS)
+    ref_q, ref_scale = mxfp8_e4m3_quantize(normed.reshape(m, n))
+
+    assert got_q.shape == (m, n) and got_q.dtype == torch.float8_e4m3fn
+    assert got_scale.shape == (m, n // 32) and got_scale.dtype == torch.uint8
+    assert got_res.shape == x.shape
+    # Both scale clamps are reached: E8M0 0 (all-zero/tiny) and above 1.0.
+    assert (ref_scale == 0).any() and (ref_scale > 127).any()
+    assert not ref_q.float().isnan().any()
+    assert torch.equal(got_q.view(torch.uint8), ref_q.view(torch.uint8))
+    assert torch.equal(got_scale, ref_scale)
+    assert torch.equal(got_res, ref_res)
+
+
 # --------------------------------------------------------------------------- #
 # SwiGLU-OAI (split layout)
 # --------------------------------------------------------------------------- #
@@ -511,6 +580,97 @@ def _spy(*a, **k):
     assert _relerr(out.float(), out_ref.float()) < tol
 
 
+# --------------------------------------------------------------------------- #
+# input_layernorm -> qkv_proj with the MXFP8 activation quant fused into the norm
+# --------------------------------------------------------------------------- #
+def _quark_mxfp8_linear(n: int, k: int, block32: bool = False):
+    """A Quark MXFP8 linear layer (weights loaded, not yet processed)."""
+    from vllm.model_executor.layers.quantization.quark.schemes.quark_ocp_mx import (
+        QuarkOCP_MX,
+    )
+    from vllm.model_executor.layers.quantization.utils.quant_utils import (
+        kMxfp8Dynamic,
+        kMxfp8Static,
+    )
+    from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+
+    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)
+    if block32:
+        w_scale = w_scale[::32].repeat_interleave(32, dim=0)
+    scheme = QuarkOCP_MX(
+        weight_quant_key=kMxfp8Static, activation_quant_key=kMxfp8Dynamic
+    )
+    with torch.device(DEVICE):
+        layer = torch.nn.Module()
+        scheme.create_weights(
+            layer=layer,
+            output_partition_sizes=[n],
+            input_size_per_partition=k,
+            params_dtype=torch.bfloat16,
+            weight_loader=default_weight_loader,
+        )
+    layer.weight.data.copy_(w_fp8)
+    layer.weight_scale.data.copy_(w_scale)
+    # As QuarkLinearMethod.create_weights does.
+    layer.scheme = scheme
+    return layer, scheme
+
+
+@requires_gfx950
+@pytest.mark.parametrize(
+    "n,k,block32",
+    [
+        (2560, 6144, False),  # MiniMax-M3.1 qkv_proj, TP4
+        (256, 2048, True),  # 32x32 block scales: the block32 GEMM
+        (5120, 576, False),  # K % 128 != 0: BF16 weight after load
+    ],
+)
+@pytest.mark.parametrize("m", [1, 37])
+@torch.inference_mode()
+def test_fused_input_norm_mxfp8_qkv_proj(dist_init, n, k, block32, m):
+    """The linear on the fused norm's MXFP8 output equals the linear on the
+    unfused norm output, wherever the producer turns the fusion on."""
+    layer, scheme = _quark_mxfp8_linear(n, k, block32)
+    fused = _can_fuse_input_norm_mxfp8(layer)
+    assert fused == (k % 128 == 0)
+    scheme.process_weights_after_loading(layer)
+
+    from vllm.model_executor.layers.fusion.quant_activation import (
+        QuantizedActivation,
+    )
+    from vllm.model_executor.layers.quantization.utils.quant_utils import (
+        kMxfp8Dynamic,
+    )
+
+    x = torch.randn(m, k, device=DEVICE, dtype=torch.bfloat16)
+    res = torch.randn(m, k, device=DEVICE, dtype=torch.bfloat16)
+    norm_w = torch.randn(k, device=DEVICE, dtype=torch.bfloat16) * 0.1
+    normed, _ = gemma_fused_add_rmsnorm(x, res, norm_w, EPS)
+    x_q, x_scale, _ = gemma_fused_add_rmsnorm_mxfp8(x, res, norm_w, EPS)
+    qa = QuantizedActivation(x_q, x_scale, x.dtype, x.shape, kMxfp8Dynamic)
+
+    got = scheme.apply_weights(layer, qa)
+    ref = scheme.apply_weights(layer, normed)
+    assert got.shape == ref.shape == (m, n) and got.dtype == ref.dtype
+    if fused:
+        assert torch.equal(got, ref)
+    else:
+        # A QuantizedActivation reaching a BF16 weight is dequantized back.
+        assert _relerr(got, ref) < 5e-2
+
+
+@requires_gfx950
+def test_fused_input_norm_mxfp8_off_under_batch_invariance(dist_init, monkeypatch):
+    import vllm.envs as envs
+
+    layer, _ = _quark_mxfp8_linear(2560, 6144)
+    assert _can_fuse_input_norm_mxfp8(layer)
+    monkeypatch.setattr(envs, "VLLM_BATCH_INVARIANT", True)
+    assert not _can_fuse_input_norm_mxfp8(layer)
+
+
 # ── 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
@@ -688,3 +848,149 @@ def test_kv_insert_operand_copies_unsupported_layouts():
 
     flat = k.reshape(32, -1)
     assert _kv_insert_operand(flat).is_contiguous()
+
+
+@pytest.mark.skipif(not on_gfx942(), reason="Requires gfx942.")
+@torch.inference_mode()
+def test_gfx942_native_qknorm_rope_writes_main_and_index_caches():
+    """The gfx942 fallback writes both caches and ignores padded slots."""
+    from types import SimpleNamespace
+
+    torch.manual_seed(0)
+    num_tokens = 3
+    num_heads = 2
+    num_kv_heads = num_idx_heads = 1
+    head_dim = idx_head_dim = 16
+    block_size = 4
+
+    attention = object.__new__(MiniMaxM3SparseAttention)
+    torch.nn.Module.__init__(attention)
+    attention.num_heads = num_heads
+    attention.num_kv_heads = num_kv_heads
+    attention.num_idx_heads = num_idx_heads
+    attention.head_dim = head_dim
+    attention.idx_head_dim = idx_head_dim
+    attention.q_size = num_heads * head_dim
+    attention.kv_size = num_kv_heads * head_dim
+    attention.index_q_size = num_idx_heads * idx_head_dim
+    attention.layer_name = "model.layers.0.self_attn"
+    attention.skip_index_topk = False
+    attention.kv_cache_dtype = "fp8"
+
+    attention.q_norm = MiniMAXGemmaRMSNorm(head_dim).to(
+        device=DEVICE, dtype=torch.bfloat16
+    )
+    attention.k_norm = MiniMAXGemmaRMSNorm(head_dim).to(
+        device=DEVICE, dtype=torch.bfloat16
+    )
+    attention.index_q_norm = MiniMAXGemmaRMSNorm(idx_head_dim).to(
+        device=DEVICE, dtype=torch.bfloat16
+    )
+    attention.index_k_norm = MiniMAXGemmaRMSNorm(idx_head_dim).to(
+        device=DEVICE, dtype=torch.bfloat16
+    )
+    config = PreTrainedConfig(
+        rope_theta=10000.0,
+        partial_rotary_factor=0.5,
+        max_position_embeddings=32,
+    )
+    from vllm.config import VllmConfig, set_current_vllm_config
+
+    with set_current_vllm_config(VllmConfig()):
+        attention.rotary_emb = _build_rotary_emb(config, head_dim).to(DEVICE)
+    attention.index_rotary_emb = attention.rotary_emb
+
+    fp8_dtype = current_platform.fp8_dtype()
+    attention.kv_cache = torch.zeros(
+        2,
+        num_kv_heads,
+        block_size,
+        2 * head_dim,
+        device=DEVICE,
+        dtype=fp8_dtype,
+    )
+    attention._k_scale = torch.ones((), device=DEVICE, dtype=torch.float32)
+    attention._v_scale = torch.ones((), device=DEVICE, dtype=torch.float32)
+    index_cache = torch.zeros(
+        2, block_size, idx_head_dim, device=DEVICE, dtype=torch.bfloat16
+    )
+    index_prefix = "model.layers.0.self_attn.indexer"
+    attention.indexer = SimpleNamespace(
+        index_cache=SimpleNamespace(
+            dtype=torch.bfloat16,
+            kv_cache=index_cache,
+            prefix=index_prefix,
+        )
+    )
+
+    width = (
+        attention.q_size
+        + 2 * attention.kv_size
+        + attention.index_q_size
+        + attention.idx_head_dim
+    )
+    qkv = torch.randn(num_tokens, width, device=DEVICE, dtype=torch.bfloat16).mul_(0.25)
+    positions = torch.tensor([0, 1, 2], device=DEVICE, dtype=torch.long)
+    main_slots = torch.tensor([0, -1, 5], device=DEVICE, dtype=torch.long)
+    index_slots = torch.tensor([2, -1, 6], device=DEVICE, dtype=torch.long)
+
+    q_raw, k_raw, v_raw, index_q_raw, index_k_raw = qkv.clone().split(
+        [
+            attention.q_size,
+            attention.kv_size,
+            attention.kv_size,
+            attention.index_q_size,
+            attention.idx_head_dim,
+        ],
+        dim=-1,
+    )
+    expected_q = attention.q_norm(q_raw.view(num_tokens, num_heads, head_dim))
+    expected_k = attention.k_norm(k_raw.view(num_tokens, num_kv_heads, head_dim))
+    expected_q, expected_k = attention.rotary_emb.forward_native(
+        positions, expected_q, expected_k
+    )
+    expected_index_q = attention.index_q_norm(
+        index_q_raw.view(num_tokens, num_idx_heads, idx_head_dim)
+    )
+    expected_index_k = attention.index_k_norm(
+        index_k_raw.view(num_tokens, idx_head_dim)
+    )
+    expected_index_q, expected_index_k = attention.index_rotary_emb.forward_native(
+        positions, expected_index_q, expected_index_k
+    )
+    assert expected_k is not None
+    assert expected_index_k is not None
+
+    q, index_q = attention._native_gfx942_qknorm_rope_cache(
+        qkv,
+        positions,
+        {
+            attention.layer_name: main_slots,
+            index_prefix: index_slots,
+        },
+    )
+
+    torch.testing.assert_close(q, expected_q.reshape(num_tokens, -1))
+    assert index_q is not None
+    torch.testing.assert_close(index_q, expected_index_q.reshape(num_tokens, -1))
+
+    expected_v = v_raw.view(num_tokens, num_kv_heads, head_dim)
+    for token, slot in ((0, 0), (2, 5)):
+        block, offset = divmod(slot, block_size)
+        assert torch.equal(
+            attention.kv_cache[block, 0, offset, :head_dim],
+            expected_k[token, 0].to(fp8_dtype),
+        )
+        assert torch.equal(
+            attention.kv_cache[block, 0, offset, head_dim:],
+            expected_v[token, 0].to(fp8_dtype),
+        )
+    assert torch.count_nonzero(attention.kv_cache[0, 0, 1]) == 0
+
+    for token, slot in ((0, 2), (2, 6)):
+        block, offset = divmod(slot, block_size)
+        assert torch.equal(
+            index_cache[block, offset],
+            expected_index_k[token].to(index_cache.dtype),
+        )
+    assert torch.count_nonzero(index_cache[0, 1]) == 0
diff --git a/tests/kernels/test_minimax_m3_dspark_draft.py b/tests/kernels/test_minimax_m3_dspark_draft.py
new file mode 100644
index 000000000000..207a72ff8378
--- /dev/null
+++ b/tests/kernels/test_minimax_m3_dspark_draft.py
@@ -0,0 +1,443 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""MiniMax-M3.1 DSpark draft fusions against the unfused ops they replace.
+
+- The context-KV CUDA kernel
+  (vllm/models/minimax_m3/common/ops/dspark_context_kv.py) against flashinfer
+  gemma_rmsnorm per layer, vllm._custom_ops.rotary_embedding (partial NeoX) and
+  vllm._custom_ops.reshape_and_cache_flash into FlashAttention-layout caches
+  [num_blocks, nkv, block_size, 2 * hd] viewed as
+  kv_cache.transpose(1, 2).split(hd, dim=-1).
+- The draft attention's fused q/k Gemma norm + RoPE against DFlash's q-norm,
+  k-norm and RoPE.
+"""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+import vllm._custom_ops as ops
+from vllm.models.minimax_m3.common.ops.dspark_context_kv import (
+    DSparkContextKVWriter,
+)
+from vllm.platforms import current_platform
+from vllm.utils.torch_utils import set_default_torch_dtype
+
+flashinfer_norm = pytest.importorskip("flashinfer.norm")
+
+HEAD_DIM = 128
+ROTARY_DIM = 64
+BLOCK_SIZE = 16
+EPS = 1e-6
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available(), reason="CUDA not available"
+)
+
+
+def make_cos_sin_cache(max_pos, rotary_dim, dtype, device):
+    inv_freq = 1.0 / (
+        5_000_000.0
+        ** (
+            torch.arange(0, rotary_dim, 2, dtype=torch.float32, device=device)
+            / rotary_dim
+        )
+    )
+    t = torch.arange(max_pos, dtype=torch.float32, device=device)
+    freqs = torch.einsum("i,j->ij", t, inv_freq)
+    return torch.cat((freqs.cos(), freqs.sin()), dim=-1).to(dtype)
+
+
+def make_slots(num_ctx, num_blocks, device):
+    """Unique valid slots with ~1/4 padding (-1), at least one valid."""
+    num_slots = num_blocks * BLOCK_SIZE
+    slots = torch.randperm(num_slots, device=device)[:num_ctx].to(torch.int64)
+    num_pad = max(1, num_ctx // 4)
+    slots[:num_pad] = -1
+    return slots
+
+
+def build_layers(kv_cache_dtype, L, nkv, num_blocks, device, seed, unit_scale=False):
+    g = torch.Generator(device="cpu").manual_seed(seed)
+    if kv_cache_dtype == "fp8" and not unit_scale:
+        dtype = torch.uint8
+        scale = lambda: (  # noqa: E731
+            torch.rand((), generator=g).to(device) * 1.5 + 0.25
+        )
+    else:
+        dtype = torch.uint8 if kv_cache_dtype == "fp8" else torch.bfloat16
+        scale = lambda: torch.ones((), device=device)  # noqa: E731
+    layers = []
+    for _ in range(L):
+        layers.append(
+            SimpleNamespace(
+                kv_cache=torch.zeros(
+                    num_blocks,
+                    nkv,
+                    BLOCK_SIZE,
+                    2 * HEAD_DIM,
+                    dtype=dtype,
+                    device=device,
+                ),
+                _k_scale=scale(),
+                _v_scale=scale(),
+                kv_cache_dtype=kv_cache_dtype,
+            )
+        )
+    return layers
+
+
+def reference_update(
+    all_kv, k_norm_w, positions, cos_sin, slot_mapping, layers, kv_cache_dtype
+):
+    """The unfused path, run on `layers`' caches."""
+    num_ctx, L, _, nkv, hd = all_kv.shape
+    all_k = torch.stack(
+        [
+            flashinfer_norm.gemma_rmsnorm(all_kv[:, layer, 0], k_norm_w[layer], EPS)
+            for layer in range(L)
+        ]
+    )  # [L, num_ctx, nkv, hd]
+    all_v = all_kv[:, :, 1]  # [num_ctx, L, nkv, hd]
+    k_flat = all_k.view(L * num_ctx, nkv * hd)
+    ops.rotary_embedding(
+        positions.repeat(L),
+        k_flat,
+        None,
+        hd,
+        cos_sin.to(k_flat.dtype),
+        True,
+    )
+    all_k = k_flat.view(L, num_ctx, nkv, hd)
+    for layer, a in enumerate(layers):
+        kc, vc = a.kv_cache.transpose(1, 2).split(hd, dim=-1)
+        ops.reshape_and_cache_flash(
+            all_k[layer],
+            all_v[:, layer].contiguous(),
+            kc,
+            vc,
+            slot_mapping,
+            kv_cache_dtype,
+            a._k_scale,
+            a._v_scale,
+        )
+
+
+def assert_bf16_close(got, exp):
+    diff = (got.view(torch.int16).int() - exp.view(torch.int16).int()).abs()
+    assert int(diff.max()) <= 1, (
+        f"bf16 cache differs by {int(diff.max())} ulp "
+        f"in {int((diff > 0).sum())} elements"
+    )
+
+
+def assert_fp8_close(got, exp):
+    equal = got == exp
+    frac = equal.float().mean().item()
+    assert frac >= 0.999, f"fp8 cache matches only {frac:.4%} of bytes"
+    bad = ~equal
+    if bad.any():
+        byte_diff = (got[bad].int() - exp[bad].int()).abs()
+        g = got[bad].view(torch.float8_e4m3fn).float()
+        e = exp[bad].view(torch.float8_e4m3fn).float()
+        ok = (byte_diff <= 1) | ((g == 0) & (e == 0))
+        assert bool(ok.all()), (
+            f"fp8 cache differs by more than 1 step in {int((~ok).sum())} bytes"
+        )
+
+
+def run_case(
+    kv_cache_dtype, num_ctx, nkv, L, seed=0, rotary_dim=ROTARY_DIM, unit_scale=False
+):
+    torch.manual_seed(seed)
+    device = "cuda"
+    num_blocks = (num_ctx * 2 + BLOCK_SIZE - 1) // BLOCK_SIZE + 1
+
+    layers, ref_layers = (
+        build_layers(kv_cache_dtype, L, nkv, num_blocks, device, seed, unit_scale)
+        for _ in range(2)
+    )
+
+    all_kv = torch.randn(
+        num_ctx, L, 2, nkv, HEAD_DIM, dtype=torch.bfloat16, device=device
+    )
+    k_norm_w = torch.randn(L, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.2
+    positions = torch.randint(0, 2048, (num_ctx,), dtype=torch.int64, device=device)
+    cos_sin = make_cos_sin_cache(2048, rotary_dim, torch.bfloat16, device)
+    slot_mapping = make_slots(num_ctx, num_blocks, device)
+
+    writer = DSparkContextKVWriter(layers)
+    writer(all_kv, k_norm_w, EPS, positions, cos_sin, rotary_dim, slot_mapping)
+    reference_update(
+        all_kv, k_norm_w, positions, cos_sin, slot_mapping, ref_layers, kv_cache_dtype
+    )
+
+    for a, b in zip(layers, ref_layers):
+        if kv_cache_dtype == "fp8":
+            assert_fp8_close(a.kv_cache, b.kv_cache)
+        else:
+            assert_bf16_close(a.kv_cache, b.kv_cache)
+
+
+@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"])
+@pytest.mark.parametrize("num_ctx", [1, 7, 64, 300])
+@pytest.mark.parametrize("nkv", [1, 2])
+@pytest.mark.parametrize("L", [1, 5])
+def test_dspark_context_kv(kv_cache_dtype, num_ctx, nkv, L):
+    run_case(kv_cache_dtype, num_ctx, nkv, L)
+
+
+@pytest.mark.parametrize("num_ctx", [7, 300])
+def test_dspark_context_kv_fp8_unit_scale(num_ctx):
+    """Unit k/v scales take the kernel's direct bf16 -> e4m3 conversion."""
+    run_case("fp8", num_ctx, nkv=2, L=5, unit_scale=True)
+
+
+@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"])
+@pytest.mark.parametrize("rotary_dim", [16, 32])
+def test_dspark_context_kv_short_rope(kv_cache_dtype, rotary_dim):
+    """Only some lanes rotate when rotary_dim < 64."""
+    run_case(kv_cache_dtype, 64, nkv=2, L=2, rotary_dim=rotary_dim)
+
+
+@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"])
+def test_dspark_context_kv_cuda_graph(kv_cache_dtype):
+    torch.manual_seed(0)
+    device = "cuda"
+    num_ctx, nkv, L = 16, 1, 5
+    num_blocks = 8
+
+    layers = build_layers(kv_cache_dtype, L, nkv, num_blocks, device, seed=0)
+    all_kv = torch.randn(
+        num_ctx, L, 2, nkv, HEAD_DIM, dtype=torch.bfloat16, device=device
+    )
+    k_norm_w = torch.randn(L, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.2
+    positions = torch.randint(0, 2048, (num_ctx,), dtype=torch.int64, device=device)
+    cos_sin = make_cos_sin_cache(2048, ROTARY_DIM, torch.bfloat16, device)
+    slot_mapping = make_slots(num_ctx, num_blocks, device)
+
+    writer = DSparkContextKVWriter(layers)
+    writer._ensure_tables()  # build pointer tables outside the graph
+
+    stream = torch.cuda.Stream()
+    stream.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(stream):
+        writer(all_kv, k_norm_w, EPS, positions, cos_sin, ROTARY_DIM, slot_mapping)
+    torch.cuda.current_stream().wait_stream(stream)
+
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        writer(all_kv, k_norm_w, EPS, positions, cos_sin, ROTARY_DIM, slot_mapping)
+
+    # Change every input in place, replay, and compare against the unfused
+    # reference run on the new inputs.
+    all_kv.copy_(torch.randn_like(all_kv))
+    positions.copy_(
+        torch.randint(0, 2048, (num_ctx,), dtype=torch.int64, device=device)
+    )
+    slot_mapping.copy_(make_slots(num_ctx, num_blocks, device))
+
+    sentinel = 0x5A if kv_cache_dtype == "fp8" else 0x3F80
+    for a in layers:
+        a.kv_cache.fill_(sentinel)
+    graph.replay()
+
+    ref_layers = build_layers(kv_cache_dtype, L, nkv, num_blocks, device, seed=0)
+    for a in ref_layers:
+        a.kv_cache.fill_(sentinel)
+    reference_update(
+        all_kv, k_norm_w, positions, cos_sin, slot_mapping, ref_layers, kv_cache_dtype
+    )
+
+    torch.accelerator.synchronize()
+    for a, b in zip(layers, ref_layers):
+        if kv_cache_dtype == "fp8":
+            assert_fp8_close(a.kv_cache, b.kv_cache)
+        else:
+            assert_bf16_close(a.kv_cache, b.kv_cache)
+
+
+def test_dspark_context_kv_supported_by_cache_layout():
+    """The writer engages on the [nb, nkv, bs, 2 * 128] page that FLASH_ATTN
+    and FLASHINFER_TS share, whatever the backend, and not on a 5-D cache."""
+    layers = build_layers("fp8", 5, 1, 4, "cuda", seed=0, unit_scale=True)
+    assert DSparkContextKVWriter.supported(layers)
+    layers[0].kv_cache = torch.zeros(
+        4, 2, BLOCK_SIZE, 1, HEAD_DIM, dtype=torch.uint8, device="cuda"
+    )
+    assert not DSparkContextKVWriter.supported(layers)
+
+
+class _CaptureAttention(torch.nn.Module):
+    """Stands in for the draft's Attention: records the q/k/v it is given."""
+
+    _q_scale_float = 1.0
+
+    def __init__(self, layer_name):
+        super().__init__()
+        self.layer_name = layer_name
+
+    def forward(self, q, k, v, output_dtype=None):
+        self.qkv = tuple(None if t is None else t.clone() for t in (q, k, v))
+        return torch.zeros_like(q, dtype=output_dtype)
+
+
+@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"])
+def test_draft_attention_fused_qk_norm_rope(
+    dist_init, default_vllm_config, kv_cache_dtype
+):
+    """One fused launch gives the q of DFlash's q-norm and RoPE and stores its
+    k/v in a bf16 or fp8 cache as the layer's reshape_and_cache_flash would;
+    the attention layer gets q alone."""
+    from vllm.forward_context import set_forward_context
+    from vllm.model_executor.models.qwen3_dflash import DFlashQwen3Attention
+    from vllm.models.minimax_m3.common.dspark import MiniMaxM3DSparkAttention
+    from vllm.transformers_utils.configs.minimax_m3 import MiniMaxM3TextConfig
+
+    torch.manual_seed(0)
+    config = MiniMaxM3TextConfig(
+        hidden_size=512,
+        num_attention_heads=8,
+        num_key_value_heads=2,
+        head_dim=HEAD_DIM,
+        max_position_embeddings=8192,
+    )
+    with torch.device("cuda"), set_default_torch_dtype(torch.bfloat16):
+        attn = MiniMaxM3DSparkAttention(
+            config,
+            sliding_window=None,
+            causal=False,
+            cache_config=None,
+            quant_config=None,
+            prefix="model.layers.0.self_attn",
+        )
+    for param in (attn.qkv_proj.weight, attn.q_norm.weight, attn.k_norm.weight):
+        param.data.normal_(std=0.05)
+    layer = attn.attn
+    cache_dtype = torch.uint8 if kv_cache_dtype == "fp8" else torch.bfloat16
+    layer.kv_cache = torch.zeros(
+        4, 2, BLOCK_SIZE, 2 * HEAD_DIM, dtype=cache_dtype, device="cuda"
+    )
+    attn.attn = _CaptureAttention(layer.layer_name)
+
+    hidden = torch.randn(7, config.hidden_size, device="cuda", dtype=torch.bfloat16)
+    positions = torch.tensor([0, 1, 127, 128, 4000, 8190, 9], device="cuda")
+    slots = torch.tensor([5, -1, 17, 18, 40, 63, 0], device="cuda")
+    DFlashQwen3Attention.forward(attn, positions, hidden)
+    q_ref, k_ref, v_ref = attn.attn.qkv
+    assert k_ref is not None and v_ref is not None
+    with set_forward_context(
+        None, default_vllm_config, slot_mapping={layer.layer_name: slots}
+    ):
+        attn(positions, hidden)
+    q, k, v = attn.attn.qkv
+    torch.testing.assert_close(q, q_ref, atol=2e-2, rtol=2e-2)
+    assert k is None and v is None
+    expected = torch.zeros_like(layer.kv_cache)
+    key_cache, value_cache = expected.transpose(1, 2).split(HEAD_DIM, dim=-1)
+    one = torch.ones((), dtype=torch.float32, device="cuda")
+    ops.reshape_and_cache_flash(
+        k_ref.reshape(7, 2, HEAD_DIM),
+        v_ref.reshape(7, 2, HEAD_DIM),
+        key_cache,
+        value_cache,
+        slots,
+        kv_cache_dtype,
+        one,
+        one,
+    )
+    got, rtol = layer.kv_cache, 2e-2
+    if kv_cache_dtype == "fp8":
+        # k/v_ref come from DFlash's unfused kernels: allow one e4m3 step.
+        got = got.view(torch.float8_e4m3fn).float()
+        expected = expected.view(torch.float8_e4m3fn).float()
+        rtol = 0.13
+    torch.testing.assert_close(got, expected, atol=2e-2, rtol=rtol)
+
+
+class _FixedOutputAttention(torch.nn.Module):
+    """Stands in for the draft's Attention: returns a fixed output in the
+    requested dtype and records the dtypes asked for."""
+
+    def __init__(self, layer_name, out):
+        super().__init__()
+        self.layer_name, self.out, self.output_dtypes = layer_name, out, []
+
+    def forward(self, q, k, v, output_dtype):
+        self.output_dtypes.append(output_dtype)
+        return self.out.to(output_dtype)
+
+
+@pytest.mark.skipif(
+    not current_platform.is_device_capability_family(100),
+    reason="Native MXFP8 kernels require SM100-family CUDA",
+)
+@torch.inference_mode()
+def test_draft_fp16_attention_output_feeds_mxfp8_o_proj(dist_init, default_vllm_config):
+    """A backend that writes an FP8 query's output as FP16 hands it to the
+    MXFP8 o_proj without a conversion pass; o_proj quantizes it and still
+    returns bf16, bit-identical to the bf16 path for values both dtypes hold."""
+    from vllm.forward_context import set_forward_context
+    from vllm.model_executor.layers.quantization.modelopt import ModelOptMxFp8Config
+    from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
+        _mxfp8_e4m3_quantize_torch,
+    )
+    from vllm.models.minimax_m3.common.dspark import MiniMaxM3DSparkAttention
+    from vllm.transformers_utils.configs.minimax_m3 import MiniMaxM3TextConfig
+
+    torch.manual_seed(0)
+    config = MiniMaxM3TextConfig(
+        hidden_size=512,
+        num_attention_heads=8,
+        num_key_value_heads=2,
+        head_dim=HEAD_DIM,
+        max_position_embeddings=8192,
+    )
+    quant_config = ModelOptMxFp8Config.from_config(
+        {
+            "quant_method": "mxfp8",
+            "activation_scheme": "dynamic",
+            "weight_block_size": [1, 32],
+            "ignored_layers": [],
+        }
+    )
+    with torch.device("cuda"), set_default_torch_dtype(torch.bfloat16):
+        attn = MiniMaxM3DSparkAttention(
+            config,
+            sliding_window=None,
+            causal=False,
+            cache_config=None,
+            quant_config=quant_config,
+            prefix="model.layers.0.self_attn",
+            query_pre_quantized=True,
+        )
+    for proj in (attn.qkv_proj, attn.o_proj):
+        weight, scale = _mxfp8_e4m3_quantize_torch(
+            torch.randn(proj.weight.shape, device="cuda") / 16
+        )
+        proj.weight.data.copy_(weight)
+        proj.weight_scale.data.copy_(scale)
+        proj.quant_method.process_weights_after_loading(proj)
+    layer = attn.attn
+    layer.kv_cache = torch.zeros(
+        4, 2, BLOCK_SIZE, 2 * HEAD_DIM, dtype=torch.uint8, device="cuda"
+    )
+    num_tokens = 7
+    out = torch.randint(-256, 256, (num_tokens, attn.q_size), device="cuda") / 128
+    attn.attn = _FixedOutputAttention(layer.layer_name, out.to(torch.bfloat16))
+
+    hidden = torch.randn(
+        num_tokens, config.hidden_size, device="cuda", dtype=torch.bfloat16
+    )
+    positions = torch.arange(num_tokens, device="cuda")
+    outputs = []
+    for attn_output_dtype in (torch.float16, None):
+        attn._attn_output_dtype = attn_output_dtype
+        with set_forward_context(
+            None, default_vllm_config, slot_mapping={layer.layer_name: positions}
+        ):
+            outputs.append(attn(positions, hidden))
+    assert attn.attn.output_dtypes == [torch.float16, torch.bfloat16]
+    assert outputs[0].dtype == torch.bfloat16
+    torch.testing.assert_close(outputs[0], outputs[1], atol=0, rtol=0)
diff --git a/tests/kernels/test_top_k_per_row.py b/tests/kernels/test_top_k_per_row.py
index 3937f86d8df3..6b015e54160b 100644
--- a/tests/kernels/test_top_k_per_row.py
+++ b/tests/kernels/test_top_k_per_row.py
@@ -1886,37 +1886,43 @@ def test_cooperative_topk_overflow_nan(num_rows: int, seq_len: int) -> None:
 
 
 NAN_ORDERING_CASES = [
-    # (backend, num_rows): persistent_topk runs its radix kernel at <= 32 rows,
+    # (backend, num_rows, dtype): persistent_topk runs its radix kernel at <= 32 rows,
     # FilteredTopK at 33-64 rows, and sampled_topk above 64 rows once
     # max_seq_len reaches its length gate (FilteredTopK fallback below it).
     pytest.param(
         "cooperative_topk",
         1,
+        torch.float32,
         marks=pytest.mark.skipif(not _has_cooperative_topk(), reason="needs SM90+"),
     ),
     pytest.param(
         "cooperative_topk",
         9,
+        torch.float32,
         marks=pytest.mark.skipif(not _has_cooperative_topk(), reason="needs SM90+"),
     ),
     pytest.param(
         "cooperative_topk",
         64,
+        torch.float32,
         marks=pytest.mark.skipif(not _has_cooperative_topk(), reason="needs SM90+"),
     ),
-    ("persistent_topk", 16),
-    ("persistent_topk", 48),
-    ("persistent_topk", 128),
+    ("persistent_topk", 16, torch.float32),
+    ("persistent_topk", 48, torch.float32),
+    ("persistent_topk", 128, torch.float32),
+    ("persistent_topk", 16, torch.bfloat16),
+    ("persistent_topk", 48, torch.bfloat16),
+    ("persistent_topk", 128, torch.bfloat16),
 ]
 
 
 @pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA")
-@pytest.mark.parametrize("backend,num_rows", NAN_ORDERING_CASES)
+@pytest.mark.parametrize("backend,num_rows,dtype", NAN_ORDERING_CASES)
 @pytest.mark.parametrize("seq_len", [65536, 262144])
 @pytest.mark.parametrize("top_k", [512, 2048])
 @torch.inference_mode()
 def test_topk_nan_ordering_matches_torch(
-    backend: str, num_rows: int, seq_len: int, top_k: int
+    backend: str, num_rows: int, seq_len: int, top_k: int, dtype: torch.dtype
 ) -> None:
     # Fewer NaNs than top_k, mixed with +inf and the largest finite values:
     # every kernel's ordered key must rank NaN above +inf, as torch.topk does.
@@ -1925,11 +1931,12 @@ def test_topk_nan_ordering_matches_torch(
     pos = torch.rand((num_rows, seq_len), generator=gen, device="cuda")
     pos = pos.topk(top_k // 4 + 3 * (top_k // 8), dim=1).indices
     n_nan, n_big = top_k // 4, top_k // 8
-    fmax = torch.finfo(torch.float32).max
+    fmax = torch.finfo(dtype).max
     logits.scatter_(1, pos[:, :n_nan], float("nan"))
     logits.scatter_(1, pos[:, n_nan : n_nan + n_big], float("inf"))
     logits.scatter_(1, pos[:, n_nan + n_big : n_nan + 2 * n_big], fmax)
     logits.scatter_(1, pos[:, n_nan + 2 * n_big :], 1e30)
+    logits = logits.to(dtype)
     lengths = torch.full((num_rows,), seq_len, dtype=torch.int32, device="cuda")
     indices = torch.empty((num_rows, top_k), dtype=torch.int32, device="cuda")
 
@@ -1943,6 +1950,75 @@ def test_topk_nan_ordering_matches_torch(
     torch.testing.assert_close(actual, expected, rtol=0, atol=0, equal_nan=True)
 
 
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA")
+@pytest.mark.parametrize("top_k", [512, 1024, 2048])
+@pytest.mark.parametrize(
+    "rows,width,host_bound",
+    [
+        (1, 8192, 8192),
+        (8, 32768, 32768),
+        (8, 65536, 65536),
+        (32, 131072, 131072),
+        (33, 1025, 1025),
+        (33, 65536, 65536),
+        (128, 8192, 1048576),
+        (128, 262144, 1048576),
+    ],
+)
+@torch.inference_mode()
+def test_persistent_topk_bf16_immediate_cast(
+    rows: int, width: int, host_bound: int, top_k: int
+) -> None:
+    """Read BF16 exactly, including short tails and updated graph inputs."""
+    stride = (width + 3) // 4 * 4
+    logits = torch.empty_strided(
+        (rows, width), (stride, 1), dtype=torch.bfloat16, device="cuda"
+    )
+    gen = torch.Generator(device="cuda").manual_seed(0)
+    logits.copy_(torch.randn(rows, width, generator=gen, device="cuda"))
+    lengths = torch.full((rows,), width, dtype=torch.int32, device="cuda")
+    indices = torch.full((rows, top_k), -2, dtype=torch.int32, device="cuda")
+    reference = torch.empty_like(indices)
+    workspace = torch.empty(RADIX_TOPK_WORKSPACE_SIZE, dtype=torch.uint8, device="cuda")
+
+    def run() -> None:
+        torch.ops._C.persistent_topk(
+            logits, lengths, indices, workspace, top_k, host_bound
+        )
+
+    run()
+    torch.accelerator.synchronize()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        run()
+    for changed in (False, True):
+        if changed:
+            logits.fill_(1)
+            bounds = [width, min(top_k + 1, width), top_k - 1, 0, width - 1]
+            lengths.copy_(
+                torch.tensor(
+                    [min(bounds[i % len(bounds)], width) for i in range(rows)],
+                    dtype=torch.int32,
+                    device="cuda",
+                )
+            )
+        indices.fill_(-2)
+        graph.replay()
+        fp32 = logits.float()
+        assert torch.equal(fp32.view(torch.int32), logits.view(torch.int16).int() << 16)
+        _run_topk_backend(
+            "persistent_topk", fp32, lengths, reference, top_k, host_bound
+        )
+        _assert_exact_topk(fp32, indices, lengths)
+        _assert_exact_topk(fp32, reference, lengths)
+        score_bags = []
+        for selected in (indices, reference):
+            values = fp32.gather(1, selected.clamp_min(0).long())
+            values.masked_fill_(selected < 0, float("-inf"))
+            score_bags.append(values.sort(1).values.view(torch.int32))
+        assert torch.equal(*score_bags)
+
+
 def _make_candidate_overflow_row(
     seq_len: int,
     top_k: int,
@@ -2063,6 +2139,7 @@ def _assert_topk_values_match_reference(
         pytest.param("persistent_topk", 128, id="sampled"),
     ],
 )
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
 @torch.inference_mode()
 def test_workspace_topk_candidate_overflow(
     backend: str,
@@ -2072,6 +2149,7 @@ def test_workspace_topk_candidate_overflow(
     score_kind: str,
     candidate_count: int | None,
     top_k: int,
+    dtype: torch.dtype,
 ) -> None:
     """TopK remains exact when one coarse bin exceeds candidate capacity."""
     torch.set_default_device("cuda:0")
@@ -2086,7 +2164,7 @@ def test_workspace_topk_candidate_overflow(
         padded = torch.full((stride,), float("inf"), dtype=torch.float32)
         padded[:seq_len] = row
         row = padded
-    logits = row.expand(num_rows, -1).contiguous()
+    logits = row.to(dtype).expand(num_rows, -1).contiguous()
     lengths = torch.full((num_rows,), seq_len, dtype=torch.int32)
     indices = torch.empty((num_rows, top_k), dtype=torch.int32)
 
@@ -2778,3 +2856,46 @@ def decode_kernel(logits, next_n, seq_lens, *args, **kwargs) -> None:
 
     assert recorded["next_n"] == 1
     assert recorded["row_ends"] == [99, 100, 199, 200]
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA")
+@pytest.mark.parametrize("backend", ["auto", "persistent"])
+@pytest.mark.parametrize("rows", [8, 64, 65, 128])
+@torch.inference_mode()
+def test_sparse_indexer_resolved_persistent_bf16(
+    backend: str, rows: int, workspace_init
+) -> None:
+    """Keep backend selection and exact scores across the BF16 row boundary."""
+    set_random_seed(0)
+    width, k = 32768, 512
+    logits = torch.empty_strided(
+        (rows, width), (33792, 1), dtype=torch.bfloat16, device="cuda"
+    ).normal_()
+    fp32 = logits.float()
+    ends = torch.full((rows, 1), width, dtype=torch.int32, device="cuda")
+    output = torch.empty((rows, k), dtype=torch.int32, device="cuda")
+    reference = torch.empty_like(output)
+    selector = _make_topk(backend)
+    resolved = selector.resolve_backend(logits, k, rows)
+    assert resolved == (
+        "cooperative"
+        if backend == "auto" and rows <= 64 and _has_cooperative_topk()
+        else "persistent"
+    )
+    assert selector.resolve_backend(fp32, k, rows) == resolved
+    takes_bf16 = selector.takes_bf16_logits(logits, k, rows)
+    assert takes_bf16 == (
+        resolved == "persistent" and current_platform.is_device_capability_family(100)
+    )
+    assert not selector.takes_bf16_logits(fp32, k, rows)
+    assert not selector.takes_bf16_logits(logits[:, ::2], k, rows)
+    assert not selector.takes_bf16_logits(logits, 256, rows)
+    actual = logits if takes_bf16 else fp32
+    selector(actual, ends, 1, output, k, width)
+    selector(fp32, ends, 1, reference, k, width)
+    assert torch.equal(fp32.view(torch.int32), logits.view(torch.int16).int() << 16)
+    for selected in (output, reference):
+        _assert_exact_topk(fp32, selected, ends)
+    actual_scores = fp32.gather(1, output.long()).sort(1).values.view(torch.int32)
+    reference_scores = fp32.gather(1, reference.long()).sort(1).values.view(torch.int32)
+    assert torch.equal(actual_scores, reference_scores)
diff --git a/tests/lora/test_fused_moe_lora_kernel.py b/tests/lora/test_fused_moe_lora_kernel.py
index 340ea4978618..7f0a560e39c6 100644
--- a/tests/lora/test_fused_moe_lora_kernel.py
+++ b/tests/lora/test_fused_moe_lora_kernel.py
@@ -18,6 +18,7 @@
     get_tensor_model_parallel_world_size,
 )
 from vllm.lora.ops.triton_ops import fused_moe_lora
+from vllm.lora.ops.triton_ops.utils import FUSED_MOE_LORA_TWO_KERNEL_MAX_TOKENS
 from vllm.platforms import current_platform
 from vllm.utils.network_utils import get_open_port
 from vllm.utils.torch_utils import set_random_seed
@@ -146,6 +147,8 @@ def use_fused_moe_lora_kernel(
     fully_sharded=False,
     offset=0,
     add_inputs=True,
+    shrink_split_k=1,
+    expert_map=None,
 ):
     max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1)
     max_num_tokens_padded = round_up(max_num_tokens_padded, block_size)
@@ -174,6 +177,7 @@ def use_fused_moe_lora_kernel(
         num_tokens_post_padded,
         adapter_enabled,
         lora_ids,
+        expert_map,
     )
 
     config = {
@@ -216,7 +220,7 @@ def use_fused_moe_lora_kernel(
         config["GROUP_SIZE_M"],
         config["NUM_WARPS"],
         config["NUM_STAGES"],
-        config["SPLIT_K"],
+        shrink_split_k,
         config["BLOCK_SIZE_M"],
         config["BLOCK_SIZE_N"],
         config["BLOCK_SIZE_K"],
@@ -471,6 +475,7 @@ def use_fused_moe_lora_kernel_naive(
     fully_sharded=False,
     offset=0,
     add_inputs=True,
+    shrink_split_k=1,
 ):
     """Test helper for naive_block_assignment path.
     Skips moe_lora_align_block_size and uses flattened topk_ids as expert_ids.
@@ -523,7 +528,7 @@ def use_fused_moe_lora_kernel_naive(
         config["GROUP_SIZE_M"],
         config["NUM_WARPS"],
         config["NUM_STAGES"],
-        config["SPLIT_K"],
+        shrink_split_k,
         config["BLOCK_SIZE_M"],
         config["BLOCK_SIZE_N"],
         config["BLOCK_SIZE_K"],
@@ -935,7 +940,7 @@ def test_fused_moe_lora_kernel_small_rank(
     torch.testing.assert_close(output, output_ref, atol=1e-2, rtol=1e-2)
 
 
-@pytest.mark.parametrize("num_tokens", [16, 64])
+@pytest.mark.parametrize("num_tokens", [288, 320])
 @pytest.mark.parametrize("top_k_num", [2])
 @pytest.mark.parametrize("num_experts", [8])
 @pytest.mark.parametrize("max_loras", [4])
@@ -961,7 +966,7 @@ def test_fused_moe_lora_kernel_npid_path(
     device,
     seed,
 ):
-    """Exercise the small-batch / NPID > 1 branch of the one-shot fast path.
+    """Exercise the NPID > 1 branch of the one-shot fast path.
 
     With these sizes the one-shot wrapper computes NPID_FACTOR > 1 (base CTA count
     < SM count), so each program covers only an outer chunk of N. The
@@ -1018,6 +1023,216 @@ def test_fused_moe_lora_kernel_npid_path(
     torch.testing.assert_close(output, output_ref, atol=2e-2, rtol=2e-2)
 
 
+def _record_fast_paths(monkeypatch):
+    from vllm.lora.ops.triton_ops import fused_moe_lora_op
+
+    calls = []
+
+    def record(name):
+        original = getattr(fused_moe_lora_op, name)
+
+        def wrapper(*args, **kwargs):
+            calls.append(name)
+            return original(*args, **kwargs)
+
+        monkeypatch.setattr(fused_moe_lora_op, name, wrapper)
+
+    record("_run_fused_moe_lora_one_shot")
+    record("_run_fused_moe_lora_small_batch")
+    return calls
+
+
+BF16, FP16 = torch.bfloat16, torch.float16
+
+
+@pytest.mark.parametrize(
+    "num_tokens,top_k_num,num_experts,max_loras,max_lora_rank,K,N,num_slices,dtype,"
+    "shrink_split_k",
+    [
+        (1, 2, 64, 4, 32, 2048, 1024, 1, BF16, 1),
+        (4, 2, 64, 4, 32, 2048, 1024, 2, BF16, 1),
+        (16, 2, 8, 4, 32, 2048, 1024, 1, BF16, 1),
+        (16, 1, 64, 4, 4, 2048, 1024, 1, BF16, 1),
+        (32, 2, 8, 4, 32, 2048, 1024, 2, BF16, 1),
+        (64, 2, 8, 4, 32, 2048, 1024, 2, BF16, 1),
+        (65, 2, 8, 4, 32, 2048, 1024, 2, BF16, 1),
+        (96, 2, 128, 16, 32, 2048, 1024, 1, BF16, 1),
+        (96, 2, 8, 4, 32, 2048, 1024, 2, BF16, 1),
+        (256, 2, 8, 4, 32, 2048, 1024, 2, BF16, 1),
+        (257, 2, 8, 4, 32, 2048, 1024, 2, BF16, 1),
+        (300, 1, 256, 16, 4, 2048, 1024, 1, BF16, 1),
+        (300, 1, 256, 16, 8, 2048, 1024, 2, FP16, 1),
+        (8, 6, 64, 4, 16, 2048, 1024, 2, FP16, 1),
+        (4, 8, 128, 4, 64, 2048, 1024, 1, BF16, 1),
+        (32, 8, 64, 4, 128, 1024, 1024, 2, FP16, 1),
+        (1, 6, 64, 4, 128, 2048, 1024, 1, BF16, 1),
+        (16, 8, 128, 4, 32, 6144, 1536, 2, BF16, 1),
+        (2, 8, 128, 4, 16, 6144, 1536, 2, FP16, 1),
+        (64, 8, 128, 4, 64, 768, 6144, 1, FP16, 1),
+        (4, 8, 128, 4, 128, 768, 6144, 1, BF16, 1),
+        (4, 2, 64, 4, 32, 2048, 1024, 2, BF16, 4),
+        (32, 2, 8, 4, 32, 2048, 1024, 1, BF16, 4),
+    ],
+    ids=[
+        "m1-naive",
+        "m4-naive",
+        "m16-sorted",
+        "m16-topk1-r4-naive",
+        "m32-sorted",
+        "m64-sorted",
+        "m65-sorted",
+        "m96-naive",
+        "m96-sorted",
+        "m256-sorted",
+        "m257-sorted-one-shot",
+        "m300-topk1-r4-naive-one-shot",
+        "m300-topk1-r8-naive-one-shot-fp16",
+        "m8-topk6-r16-fp16",
+        "m4-topk8-r64-naive",
+        "m32-topk8-r128-fp16",
+        "m1-topk6-r128-naive",
+        "w13-m16-r32",
+        "w13-m2-r16-fp16-naive",
+        "w2-m64-r64-fp16",
+        "w2-m4-r128-naive",
+        "splitk4-m4-naive",
+        "splitk4-m32-sorted",
+    ],
+)
+@pytest.mark.parametrize("device", DEVICES)
+def test_fused_moe_lora_token_threshold_dispatch(
+    num_tokens,
+    top_k_num,
+    num_experts,
+    max_loras,
+    max_lora_rank,
+    K,
+    N,
+    num_slices,
+    dtype,
+    shrink_split_k,
+    device,
+    monkeypatch,
+):
+    fast_calls = _record_fast_paths(monkeypatch)
+    torch.set_default_device(device)
+    set_random_seed(0)
+    block_size = 16
+    topk_ids, topk_weights, token_lora_mapping, lora_ids = sample_data(
+        num_tokens, min(num_tokens, 4), max_loras, num_experts, top_k_num
+    )
+    lora_a_stacked = [
+        torch.rand((max_loras, num_experts, max_lora_rank, K), dtype=dtype)
+        for _ in range(num_slices)
+    ]
+    lora_b_stacked = [
+        torch.rand(
+            (max_loras, num_experts, N // num_slices, max_lora_rank), dtype=dtype
+        )
+        for _ in range(num_slices)
+    ]
+    hidden_states = torch.rand((num_tokens, K), dtype=dtype)
+    output = torch.zeros((num_tokens, top_k_num, N), dtype=dtype)
+    args = (
+        topk_ids,
+        topk_weights,
+        token_lora_mapping,
+        max_lora_rank,
+        top_k_num,
+        lora_ids,
+        lora_a_stacked,
+        lora_b_stacked,
+        hidden_states,
+        output,
+        max_loras,
+    )
+    if num_tokens * top_k_num * 8 <= num_experts * max_loras:
+        use_fused_moe_lora_kernel_naive(
+            *args, block_size, shrink_split_k=shrink_split_k
+        )
+    else:
+        use_fused_moe_lora_kernel(
+            *args, num_experts, block_size, shrink_split_k=shrink_split_k
+        )
+    assert bool(fast_calls) == (num_tokens > FUSED_MOE_LORA_TWO_KERNEL_MAX_TOKENS), (
+        fast_calls
+    )
+    output_ref = use_torch(
+        hidden_states,
+        token_lora_mapping,
+        topk_ids,
+        lora_a_stacked,
+        lora_b_stacked,
+        top_k_num,
+        num_slices,
+    )
+    torch.testing.assert_close(output, output_ref, atol=2e-2, rtol=2e-2)
+
+
+@pytest.mark.parametrize("num_tokens", [4, 64, 300], ids=["m4", "m64", "m300"])
+@pytest.mark.parametrize("device", DEVICES)
+def test_fused_moe_lora_expert_parallel(num_tokens, device, monkeypatch):
+    fast_calls = _record_fast_paths(monkeypatch)
+    torch.set_default_device(device)
+    set_random_seed(0)
+    top_k_num, num_experts, max_loras, max_lora_rank = 2, 16, 4, 32
+    K, N, num_slices, block_size, dtype = 2048, 1024, 2, 16, torch.bfloat16
+    topk_ids, topk_weights, token_lora_mapping, lora_ids = sample_data(
+        num_tokens, min(num_tokens, 4), max_loras, num_experts, top_k_num
+    )
+    topk_ids[0] = torch.tensor([1, 3], dtype=torch.int32)
+    local_experts = torch.arange(0, num_experts, 2)
+    expert_map = torch.full((num_experts,), -1, dtype=torch.int32)
+    expert_map[local_experts] = torch.arange(local_experts.numel(), dtype=torch.int32)
+    is_local = expert_map[topk_ids.long()] >= 0
+    lora_a_stacked = [
+        torch.rand((max_loras, num_experts, max_lora_rank, K), dtype=dtype)
+        for _ in range(num_slices)
+    ]
+    lora_b_stacked = [
+        torch.rand(
+            (max_loras, num_experts, N // num_slices, max_lora_rank), dtype=dtype
+        )
+        for _ in range(num_slices)
+    ]
+    for lora_a, lora_b in zip(lora_a_stacked, lora_b_stacked):
+        lora_a[:, expert_map < 0] = 0
+        lora_b[:, expert_map < 0] = 0
+    hidden_states = torch.rand((num_tokens, K), dtype=dtype)
+    residual = torch.randn((num_tokens, top_k_num, N), dtype=dtype)
+    output = residual.clone()
+    use_fused_moe_lora_kernel(
+        topk_ids,
+        topk_weights,
+        token_lora_mapping,
+        max_lora_rank,
+        top_k_num,
+        lora_ids,
+        [lora_a[:, local_experts].contiguous() for lora_a in lora_a_stacked],
+        [lora_b[:, local_experts].contiguous() for lora_b in lora_b_stacked],
+        hidden_states,
+        output,
+        max_loras,
+        num_experts,
+        block_size,
+        expert_map=expert_map,
+    )
+    assert bool(fast_calls) == (num_tokens > FUSED_MOE_LORA_TWO_KERNEL_MAX_TOKENS), (
+        fast_calls
+    )
+    expected = residual + use_torch(
+        hidden_states,
+        token_lora_mapping,
+        topk_ids,
+        lora_a_stacked,
+        lora_b_stacked,
+        top_k_num,
+        num_slices,
+    )
+    torch.testing.assert_close(output, expected, atol=2e-2, rtol=2e-2)
+    torch.testing.assert_close(output[~is_local], residual[~is_local], atol=0, rtol=0)
+
+
 # -- one-shot corner-case coverage ------------------------------------------
 # Each of the following exercises a path where the kernel is launched but
 # every program early-exits, leaving the output unchanged. The contract is
@@ -1133,15 +1348,20 @@ def _call_one_shot(
     "trigger",
     ["sorted_lora_ids_neg", "naive_mapping_neg", "naive_all_disabled"],
 )
+@pytest.mark.parametrize("one_shot", [False, True], ids=["two_kernel", "one_shot"])
 @pytest.mark.parametrize("device", DEVICES)
-def test_fused_moe_lora_kernel_one_shot_early_exit(trigger, device):
-    """one-shot must leave the residual byte-identical when every program
-    must early-exit. Three trigger conditions are covered:
+def test_fused_moe_lora_kernel_one_shot_early_exit(
+    trigger, one_shot, device, monkeypatch
+):
+    """Both the one-shot and two-kernel paths must leave the residual
+    byte-identical when every program must early-exit. Three trigger
+    conditions are covered:
 
       - "sorted_lora_ids_neg":  sorted path, lora_ids all -1 (lora_id<0 check)
       - "naive_mapping_neg":    naive path, token_lora_mapping all -1
       - "naive_all_disabled":   naive path, adapter_enabled all 0
     """
+    fast_calls = _record_fast_paths(monkeypatch)
     torch.set_default_device(device)
     set_random_seed(0)
 
@@ -1150,10 +1370,14 @@ def test_fused_moe_lora_kernel_one_shot_early_exit(trigger, device):
     # larger E/max_loras and smaller num_tokens.
     if trigger == "naive_mapping_neg":
         num_tokens, top_k, E, max_loras, R = 8, 2, 64, 8, 16
+        if one_shot:
+            top_k, E, max_loras = 1, 256, 16
     elif trigger == "naive_all_disabled":
         num_tokens, top_k, E, max_loras, R = 32, 2, 8, 4, 32
     else:  # sorted_lora_ids_neg
         num_tokens, top_k, E, max_loras, R = 32, 2, 8, 4, 16
+    if one_shot:
+        num_tokens = 300
     K, N = 1024, 1024
     block_size, num_slices, dtype = 16, 2, torch.bfloat16
 
@@ -1210,16 +1434,19 @@ def test_fused_moe_lora_kernel_one_shot_early_exit(trigger, device):
         adapter_enabled,
         block_size,
     )
+    assert bool(fast_calls) == one_shot, fast_calls
     torch.testing.assert_close(output, residual, atol=0, rtol=0)
 
 
+@pytest.mark.parametrize("num_tokens", [8, 300], ids=["two_kernel", "one_shot"])
 @pytest.mark.parametrize("device", DEVICES)
-def test_fused_moe_lora_kernel_zero_grid_no_crash(device):
+def test_fused_moe_lora_kernel_zero_grid_no_crash(num_tokens, device, monkeypatch):
     """num_active_loras=0 (or num_slices=0) would otherwise launch a grid
-    with a zero dimension. one-shot wrapper must short-circuit before launch."""
+    with a zero dimension. Both paths must short-circuit before launch."""
+    fast_calls = _record_fast_paths(monkeypatch)
     torch.set_default_device(device)
     set_random_seed(0)
-    num_tokens, top_k, E, max_loras, R, K, N = 8, 2, 8, 4, 16, 1024, 1024
+    top_k, E, max_loras, R, K, N = 2, 8, 4, 16, 1024, 1024
     block_size, num_slices, dtype = 16, 2, torch.bfloat16
 
     (
@@ -1271,6 +1498,9 @@ def test_fused_moe_lora_kernel_zero_grid_no_crash(device):
         adapter_enabled,
         block_size,
     )
+    assert bool(fast_calls) == (num_tokens > FUSED_MOE_LORA_TWO_KERNEL_MAX_TOKENS), (
+        fast_calls
+    )
     torch.testing.assert_close(output, residual, atol=0, rtol=0)
 
 
@@ -1281,7 +1511,7 @@ def test_fused_moe_lora_kernel_rejects_bad_block_size_m(device):
     failure (`arange's range must be a power of 2`)."""
     torch.set_default_device(device)
     set_random_seed(0)
-    num_tokens, top_k, E, max_loras, R, K, N = 32, 2, 8, 4, 16, 1024, 1024
+    num_tokens, top_k, E, max_loras, R, K, N = 300, 2, 8, 4, 16, 1024, 1024
     num_slices, dtype = 2, torch.bfloat16
     block_size = 24  # NOT a power of 2
 
diff --git a/tests/lora/test_gemma4_tp.py b/tests/lora/test_gemma4_tp.py
index efbee00d6249..276bba0b2d6f 100644
--- a/tests/lora/test_gemma4_tp.py
+++ b/tests/lora/test_gemma4_tp.py
@@ -62,6 +62,7 @@ def test_gemma4_lora(gemma4_vision_lora_files):
         limit_mm_per_prompt={"image": 1},
         mm_processor_cache_gb=0,
         enable_tower_connector_lora=True,
+        gpu_memory_utilization=0.8,
     )
 
     generate_and_test(llm, gemma4_vision_lora_files, lora_id=1)
diff --git a/tests/lora/test_lora_functions.py b/tests/lora/test_lora_functions.py
index 5bd3ed6b045b..c959e0d21d33 100644
--- a/tests/lora/test_lora_functions.py
+++ b/tests/lora/test_lora_functions.py
@@ -4,6 +4,7 @@
 
 import pytest
 
+from vllm import LLM, SamplingParams
 from vllm.engine.arg_utils import AsyncEngineArgs, EngineArgs
 from vllm.entrypoints.launchers.api_server.entry import (
     build_async_engine_client_from_engine_args,
@@ -116,3 +117,35 @@ async def run_check(fn, args, expected: list):
         await run_check(llm.remove_lora, 12, [10, 11])
         await run_check(llm.remove_lora, 11, [10])
         await run_check(llm.remove_lora, 10, [])
+
+
+def test_same_lora_name_with_different_ids(
+    qwen3_meowing_lora_files, qwen3_woofing_lora_files
+):
+    """An adapter unloaded and loaded again under the same name gets a new id
+    while requests for the old id may still be running; each request in the
+    shared batch must use its own adapter."""
+    llm = LLM(
+        model=MODEL_PATH,
+        enable_lora=True,
+        max_loras=2,
+        max_lora_rank=LORA_RANK,
+        max_model_len=128,
+        gpu_memory_utilization=0.8,
+        enforce_eager=True,
+    )
+    messages = [
+        {"role": "system", "content": "Follow the instructions to make animal noises"},
+        {"role": "user", "content": "Make your favorite animal noise."},
+    ]
+    outputs = llm.chat(
+        [messages, messages],
+        SamplingParams(temperature=0, max_tokens=10),
+        chat_template_kwargs={"enable_thinking": False},
+        lora_request=[
+            LoRARequest("pet", 1, qwen3_meowing_lora_files),
+            LoRARequest("pet", 2, qwen3_woofing_lora_files),
+        ],
+    )
+    assert "Meow Meow Meow" in outputs[0].outputs[0].text
+    assert "Woof Woof Woof" in outputs[1].outputs[0].text
diff --git a/tests/lora/test_lora_manager.py b/tests/lora/test_lora_manager.py
index 54f9a80b6415..d5fa621e36e4 100644
--- a/tests/lora/test_lora_manager.py
+++ b/tests/lora/test_lora_manager.py
@@ -1192,6 +1192,43 @@ def test_packed_loras(default_vllm_config, dist_init, dummy_model_gate_up, devic
     torch.testing.assert_close(packed_lora1.lora_b[1], up_lora1.lora_b)
 
 
+@pytest.mark.parametrize("device", DEVICES)
+def test_packed_loras_layer_override(
+    default_vllm_config, dist_init, dummy_model_gate_up, device
+):
+    model = dummy_model_gate_up
+    model.gate_up_proj.lora_packed_modules = ["gate_proj", "gate_proj"]
+    model_lora = create_packed_lora(
+        1,
+        model,
+        module_name="gate_up_proj",
+        replaced_module_names=["gate_proj", "up_proj"],
+        device=device,
+        empty_replaced_module_name="up_proj",
+    )
+    gate_lora = model_lora.get_lora("gate_proj")
+    assert gate_lora is not None
+
+    manager = LoRAModelManager(
+        model,
+        2,
+        2,
+        2,
+        LoRAConfig(
+            max_lora_rank=8, max_cpu_loras=2, max_loras=2, lora_dtype=DEFAULT_DTYPE
+        ),
+        device=device,
+        vllm_config=default_vllm_config,
+    )
+    assert manager.add_adapter(model_lora)
+
+    packed_lora = model_lora.get_lora("gate_up_proj")
+    assert isinstance(packed_lora, PackedLoRALayerWeights)
+    for i in range(2):
+        torch.testing.assert_close(packed_lora.lora_a[i], gate_lora.lora_a)
+        torch.testing.assert_close(packed_lora.lora_b[i], gate_lora.lora_b)
+
+
 def _test_target_modules(
     model,
     target_modules: list[str] | None,
diff --git a/tests/lora/test_moe_lora_ep_load.py b/tests/lora/test_moe_lora_ep_load.py
index 60d0c2c9ba7e..158c9375a208 100644
--- a/tests/lora/test_moe_lora_ep_load.py
+++ b/tests/lora/test_moe_lora_ep_load.py
@@ -4,11 +4,18 @@
 load-time slicing
 """
 
+from contextlib import nullcontext
+from dataclasses import replace
+from types import SimpleNamespace
+
 import pytest
 import torch
 
+from vllm.lora.layers.fused_moe import FusedMoEWithLoRA
 from vllm.lora.lora_model import LoRAModel, MoEEPLoadSpec
 from vllm.lora.peft_helper import PEFTHelper
+from vllm.model_executor.layers.fused_moe.config import FusedMoEParallelConfig
+from vllm.model_executor.layers.fused_moe.expert_map_manager import ExpertMapManager
 
 NUM_LAYERS = 48
 GLOBAL_NUM_EXPERTS = 128
@@ -17,6 +24,56 @@
 NON_EXPERT_MODULES = ("q_proj", "k_proj", "v_proj", "o_proj", "gate")
 
 
+@pytest.mark.parametrize(
+    "use_ep,requested_placement,num_expert_group,resolved_placement",
+    [
+        (False, "linear", 2, "linear"),
+        (False, "round_robin", 2, "linear"),
+        (True, "linear", 2, "linear"),
+        (True, "round_robin", 1, "linear"),
+        (True, "round_robin", 2, "round_robin"),
+    ],
+)
+def test_moe_lora_ep_requires_contiguous_expert_placement(
+    use_ep, requested_placement, num_expert_group, resolved_placement
+):
+    """Reject non-contiguous ownership, but allow requests resolved to linear."""
+    parallel_config = replace(
+        FusedMoEParallelConfig.make_no_parallel(),
+        use_ep=use_ep,
+        ep_size=2 if use_ep else 1,
+    )
+    expert_map = ExpertMapManager(
+        max_num_batched_tokens=16,
+        top_k=2,
+        global_num_experts=4,
+        num_redundant_experts=0,
+        num_expert_group=num_expert_group,
+        moe_parallel_config=parallel_config,
+        placement_strategy=requested_placement,
+        enable_eplb=False,
+    )
+    assert expert_map.placement_strategy == resolved_placement
+
+    layer = FusedMoEWithLoRA.__new__(FusedMoEWithLoRA)
+    torch.nn.Module.__init__(layer)
+    layer.moe_config = SimpleNamespace(moe_parallel_config=parallel_config)
+    layer.base_layer = SimpleNamespace(
+        expert_placement_strategy=expert_map.placement_strategy
+    )
+    error = (
+        pytest.raises(
+            ValueError,
+            match="Set expert_placement_strategy='linear' to use MoE LoRA "
+            "with expert parallelism",
+        )
+        if use_ep and resolved_placement != "linear"
+        else nullcontext()
+    )
+    with error:
+        layer._ep_check()
+
+
 def _expected_lora_modules() -> set[str]:
     """Replicate the set ``WorkerLoRAManager._load_adapter`` would build
     from this model's ``packed_modules_mapping``."""
diff --git a/tests/lora/test_olmoe_tp.py b/tests/lora/test_olmoe_tp.py
index c9a4fd6df5b1..dc391a12a2c5 100644
--- a/tests/lora/test_olmoe_tp.py
+++ b/tests/lora/test_olmoe_tp.py
@@ -11,7 +11,6 @@
 
 import vllm
 from vllm.lora.request import LoRARequest
-from vllm.platforms import current_platform
 
 from ..utils import multi_gpu_test
 
@@ -59,7 +58,6 @@ def generate_and_test(
     llm: vllm.LLM,
     lora_path: str,
     lora_id: list[int | None] | int | None,
-    compare_lower: bool = False,
 ) -> None:
     prompts = [
         PROMPT_TEMPLATE.format(context="How many candidates are there?"),
@@ -99,38 +97,11 @@ def generate_and_test(
             if req_lora_id is not None
             else EXPECTED_BASE_MODEL_OUTPUT[i]
         )
-
-        if compare_lower:
-            generated_text = generated_text.lower()
-            if isinstance(expected_output, str):
-                expected_output = (expected_output.lower(),)
-            else:
-                expected_output = tuple(s.lower() for s in expected_output)
         assert _output_matches(generated_text, expected_output), (
             f"Output {i}: {generated_text!r} does not match any of {expected_output!r}"
         )
 
 
-@pytest.mark.skipif(
-    current_platform.is_cuda_alike(), reason="Skipping to avoid redundant model tests"
-)
-def test_olmoe_lora(olmoe_lora_files, maybe_enable_lora_dual_stream):
-    # We enable enforce_eager=True here to reduce VRAM usage for lora-test CI,
-    # Otherwise, the lora-test will fail due to CUDA OOM.
-    llm = vllm.LLM(
-        MODEL_PATH,
-        max_model_len=1024,
-        enable_lora=True,
-        max_loras=4,
-        enforce_eager=True,
-        trust_remote_code=True,
-        enable_chunked_prefill=True,
-    )
-
-    generate_and_test(llm, olmoe_lora_files, lora_id=1)
-    generate_and_test(llm, olmoe_lora_files, lora_id=2)
-
-
 def test_olmoe_lora_mixed(olmoe_lora_files):
     llm = vllm.LLM(
         MODEL_PATH,
@@ -183,9 +154,6 @@ def test_olmoe_lora_mixed_random(
     assert outputs[0].outputs[0].text.strip().startswith(EXPECTED_LORA_OUTPUT[0])
 
 
-@pytest.mark.skipif(
-    current_platform.is_cuda_alike(), reason="Skipping to avoid redundant model tests"
-)
 @pytest.mark.parametrize("fully_sharded_loras", [False, True])
 @multi_gpu_test(num_gpus=2)
 def test_olmoe_lora_tp2(olmoe_lora_files, fully_sharded_loras):
@@ -203,25 +171,3 @@ def test_olmoe_lora_tp2(olmoe_lora_files, fully_sharded_loras):
 
     generate_and_test(llm, olmoe_lora_files, lora_id=1)
     generate_and_test(llm, olmoe_lora_files, lora_id=2)
-
-
-@pytest.mark.parametrize("fully_sharded_loras", [False, True])
-@multi_gpu_test(num_gpus=4)
-def test_olmoe_lora_tp4(olmoe_lora_files, fully_sharded_loras):
-    llm = vllm.LLM(
-        MODEL_PATH,
-        max_model_len=1024,
-        enable_lora=True,
-        max_loras=4,
-        enforce_eager=True,
-        trust_remote_code=True,
-        enable_chunked_prefill=True,
-        tensor_parallel_size=4,
-        fully_sharded_loras=fully_sharded_loras,
-    )
-    generate_and_test(
-        llm, olmoe_lora_files, lora_id=1, compare_lower=fully_sharded_loras
-    )
-    generate_and_test(
-        llm, olmoe_lora_files, lora_id=2, compare_lower=fully_sharded_loras
-    )
diff --git a/tests/lora/test_sequence_classification.py b/tests/lora/test_sequence_classification.py
index 773e5448c85f..3ed3cf3905c4 100644
--- a/tests/lora/test_sequence_classification.py
+++ b/tests/lora/test_sequence_classification.py
@@ -1,7 +1,11 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import json
+
+import pytest
 import torch
+from safetensors.torch import save_file
 
 from tests.utils import multi_gpu_test
 from vllm import LLM, PoolingParams
@@ -10,6 +14,7 @@
 
 MODEL_NAME = "Qwen/Qwen3-0.6B"
 NATIVE_MODEL_NAME = "Skywork/Skywork-Reward-V2-Qwen3-0.6B"
+BERT_MODEL_NAME = "cross-encoder/ms-marco-MiniLM-L6-v2"
 PROMPTS = {
     "star_trek": "Does warp drive appear in Star Trek?",
     "new_zealand": "Wellington is the capital of New Zealand.",
@@ -100,6 +105,54 @@ def test_native_classification_model_with_modules_to_save(
     cleanup_dist_env_and_memory()
 
 
+def _save_dense_adapter(path, hidden_size: int, modules: list[str]) -> str:
+    path.mkdir()
+    config = {"r": 8, "lora_alpha": 16, "target_modules": ["dense"]}
+    (path / "adapter_config.json").write_text(json.dumps(config))
+    tensors = {}
+    for module in modules:
+        tensors[f"base_model.model.{module}.lora_A.weight"] = torch.zeros(
+            8, hidden_size
+        )
+        tensors[f"base_model.model.{module}.lora_B.weight"] = torch.zeros(
+            hidden_size, 8
+        )
+    save_file(tensors, str(path / "adapter_model.safetensors"))
+    return str(path)
+
+
+def test_lora_on_unsupported_module_is_rejected(tmp_path) -> None:
+    """BERT's pooler.dense is an nn.Linear that LoRA does not wrap. Its LoRA
+    weights must fail the load instead of being ignored, while the encoder's
+    dense layers with the same suffix still load."""
+    llm = LLM(
+        model=BERT_MODEL_NAME,
+        runner="pooling",
+        dtype="float16",
+        enable_lora=True,
+        max_lora_rank=8,
+        enforce_eager=True,
+        max_model_len=512,
+        gpu_memory_utilization=0.3,
+    )
+    hidden_size = llm.llm_engine.model_config.get_hidden_size()
+    encoder_dense = "bert.encoder.layer.0.attention.output.dense"
+    pooler_adapter = _save_dense_adapter(
+        tmp_path / "pooler", hidden_size, [encoder_dense, "bert.pooler.dense"]
+    )
+    encoder_adapter = _save_dense_adapter(
+        tmp_path / "encoder", hidden_size, [encoder_dense]
+    )
+
+    with pytest.raises(Exception, match="bert.pooler.dense"):
+        llm.llm_engine.add_lora(LoRARequest("pooler", 1, pooler_adapter))
+    assert llm.llm_engine.add_lora(LoRARequest("encoder", 2, encoder_adapter))
+    assert set(llm.llm_engine.list_loras()) == {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,
diff --git a/tests/model_executor/layers/test_fused_shared_expert.py b/tests/model_executor/layers/test_fused_shared_expert.py
index 30b7e2fbd34a..e23feb9d911d 100644
--- a/tests/model_executor/layers/test_fused_shared_expert.py
+++ b/tests/model_executor/layers/test_fused_shared_expert.py
@@ -17,6 +17,7 @@
 import vllm.config as vllm_config_module
 from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config
 from vllm.config.quantization import QuantizationConfigArgs
+from vllm.logger import _print_warning_once
 from vllm.model_executor.layers.fused_moe import utils as fused_moe_utils
 from vllm.model_executor.layers.fused_moe.layer import determine_expert_counts
 from vllm.model_executor.layers.quantization.experts_int8 import ExpertsInt8Config
@@ -31,7 +32,10 @@
 from vllm.model_executor.models.qwen3_next import (
     _should_replicate_misaligned_shared_expert,
 )
-from vllm.model_executor.models.utils import PPMissingLayer
+from vllm.model_executor.models.utils import (
+    get_fused_shared_expert_layers,
+    maybe_fuse_shared_experts,
+)
 from vllm.models.deepseek_v4 import quant_config as deepseek_v4_quant_config
 from vllm.models.minimax_m3.amd import model as minimax_m3_model
 from vllm.platforms import current_platform
@@ -148,13 +152,15 @@ class _StubGlm5NextAttention(nn.Module):
     def __init__(self, *args: Any, **kwargs: Any) -> None:
         super().__init__()
         self.o_proj = SimpleNamespace(reduce_results=True)
+        # GLM5Next's MTP predictor reads this attribute from MLA layers.
+        self.mla_attn = SimpleNamespace()
 
 
 def get_fse_test_model_config(
     model_type: str,
     quantization_config: dict[str, Any],
-) -> tuple[object, type[nn.Module]]:
-    config: object
+) -> tuple[Any, type[nn.Module]]:
+    config: Any
     if model_type == "minimax_m3":
         config = MiniMaxM3TextConfig(
             hidden_size=128,
@@ -166,24 +172,29 @@ def get_fse_test_model_config(
             head_dim=128,
             num_local_experts=2,
             num_experts_per_tok=1,
-            moe_layer_freq=[0, 1],
+            moe_layer_freq=[1, 1],
             sparse_attention_config=None,
             rotary_dim=64,
             quantization_config=quantization_config,
         )
         return config, minimax_m3_model.MiniMaxM3Model
-    if model_type == "deepseek_v4":
-        from vllm.models.deepseek_v4.amd.model import DeepseekV4Model
+    if model_type in ("deepseek_v4", "deepseek_v4_cpu"):
+        if model_type == "deepseek_v4":
+            from vllm.models.deepseek_v4.amd.model import DeepseekV4Model
+        else:
+            from vllm.models.deepseek_v4.cpu.model import DeepseekV4Model
 
         config = SimpleNamespace(
+            model_type="deepseek_v4",
             vocab_size=256,
             hidden_size=128,
-            num_hidden_layers=1,
+            num_hidden_layers=2,
             num_attention_heads=1,
             head_dim=128,
             max_position_embeddings=128,
             sliding_window=None,
-            layer_types=["sliding_attention"],
+            layer_types=["sliding_attention", "sliding_attention"],
+            mlp_layer_types=[],
             compress_rates={},
             rope_theta=10000.0,
             compress_rope_theta=10000.0,
@@ -214,13 +225,13 @@ def get_fse_test_model_config(
             quantization_config=quantization_config,
         )
         return config, DeepseekV4Model
-    if model_type == "qwen3_5":
-        from vllm.model_executor.models.qwen3_5 import Qwen3_5Model
+    if model_type in ("qwen3_next", "qwen3_5", "alice_ai"):
+        from transformers import Qwen3NextConfig
 
-        config = Qwen3_5MoeTextConfig(
+        common = dict(
             vocab_size=256,
             hidden_size=128,
-            num_hidden_layers=1,
+            num_hidden_layers=2,
             num_attention_heads=1,
             num_key_value_heads=1,
             head_dim=128,
@@ -232,20 +243,113 @@ def get_fse_test_model_config(
             shared_expert_intermediate_size=32,
             num_experts_per_tok=1,
             num_experts=2,
-            layer_types=["full_attention"],
+            layer_types=["full_attention", "full_attention"],
             quantization_config=quantization_config,
         )
-        return config, Qwen3_5Model
-    if model_type == "deepseek_v2":
-        from transformers import DeepseekV2Config
+        if model_type == "qwen3_next":
+            from vllm.model_executor.models.qwen3_next import Qwen3NextModel
 
-        from vllm.model_executor.models.deepseek_v2 import DeepseekV2Model
+            return Qwen3NextConfig(**common), Qwen3NextModel
+        if model_type == "qwen3_5":
+            from vllm.model_executor.models.qwen3_5 import Qwen3_5Model
 
-        config = DeepseekV2Config(
+            return Qwen3_5MoeTextConfig(**common), Qwen3_5Model
+        from vllm.model_executor.models.alice_ai import AliceAIModel
+
+        config = Qwen3NextConfig(**common, linear_conv_kernel_dim=4)
+        # AliceAI requires the AttnRes grouping metadata.
+        config.block_attn_res_block_size = 1
+        return config, AliceAIModel
+    if model_type == "qwen4_exp":
+        from transformers import Qwen4ExpTextConfig
+
+        config = Qwen4ExpTextConfig(
+            vocab_size=256,
+            hidden_size=128,
+            num_hidden_layers=2,
+            num_attention_heads=1,
+            num_key_value_heads=1,
+            head_dim=128,
+            linear_key_head_dim=64,
+            linear_value_head_dim=64,
+            linear_num_key_heads=1,
+            linear_num_value_heads=1,
+            moe_intermediate_size=32,
+            shared_expert_intermediate_size=32,
+            num_experts_per_tok=1,
+            num_experts=2,
+            layer_types=["full_attention", "full_attention"],
+            eos_token_id=1,
+            hc_count=2,
+            hc_lowrank=4,
+            ple_layer_ids=[],
+            quantization_config=quantization_config,
+        )
+        from vllm.models.qwen4_exp.amd.model import Qwen4ExpModel
+
+        return config, Qwen4ExpModel
+    if model_type == "AXK1":
+        from transformers import AXK1Config
+
+        from vllm.model_executor.models.AXK1 import AXK1Model
+
+        config = AXK1Config(
             vocab_size=256,
             hidden_size=128,
             intermediate_size=32,
-            num_hidden_layers=1,
+            moe_intermediate_size=32,
+            num_hidden_layers=2,
+            num_attention_heads=1,
+            num_key_value_heads=1,
+            n_shared_experts=1,
+            n_routed_experts=2,
+            num_experts_per_tok=1,
+            n_group=1,
+            topk_group=1,
+            first_k_dense_replace=0,
+            q_lora_rank=0,
+            kv_lora_rank=0,
+            qk_nope_head_dim=0,
+            qk_rope_head_dim=0,
+            v_head_dim=0,
+            quantization_config=quantization_config,
+        )
+        return config, AXK1Model
+    if model_type == "glm4_moe_lite":
+        from transformers.models.glm4_moe_lite import Glm4MoeLiteConfig
+
+        from vllm.model_executor.models.glm4_moe_lite import Glm4MoeLiteModel
+
+        config = Glm4MoeLiteConfig(
+            vocab_size=256,
+            hidden_size=128,
+            intermediate_size=32,
+            moe_intermediate_size=32,
+            num_hidden_layers=2,
+            num_attention_heads=1,
+            num_key_value_heads=1,
+            n_routed_experts=2,
+            n_shared_experts=1,
+            num_experts_per_tok=1,
+            q_lora_rank=0,
+            kv_lora_rank=32,
+            qk_nope_head_dim=32,
+            qk_rope_head_dim=16,
+            v_head_dim=32,
+            dtype=torch.bfloat16,
+            quantization_config=quantization_config,
+        )
+        config.first_k_dense_replace = 0
+        config.num_nextn_predict_layers = 1
+        return config, Glm4MoeLiteModel
+    if model_type in ("deepseek_v2", "deepseek_v32"):
+        from transformers import DeepseekV2Config
+
+        common = dict(
+            vocab_size=256,
+            hidden_size=128,
+            intermediate_size=32,
+            num_hidden_layers=2,
             num_attention_heads=1,
             num_key_value_heads=1,
             head_dim=128,
@@ -263,7 +367,18 @@ def get_fse_test_model_config(
             v_head_dim=0,
             quantization_config=quantization_config,
         )
-        return config, DeepseekV2Model
+        config = DeepseekV2Config(**common)
+        if model_type == "deepseek_v2":
+            from vllm.model_executor.models.deepseek_v2 import DeepseekV2Model
+
+            # Needed to build the DeepSeek-V2 MTP draft model in this test.
+            config.num_nextn_predict_layers = 1
+            return config, DeepseekV2Model
+        from vllm.models.deepseek_v32.amd.model import DeepseekV32Model
+
+        config.index_topk = 1
+        config.num_nextn_predict_layers = 1
+        return config, DeepseekV32Model
     if model_type == "glm4_moe":
         from transformers.models.glm4_moe import Glm4MoeConfig
 
@@ -273,7 +388,7 @@ def get_fse_test_model_config(
             vocab_size=256,
             hidden_size=128,
             intermediate_size=32,
-            num_hidden_layers=1,
+            num_hidden_layers=2,
             num_attention_heads=1,
             num_key_value_heads=1,
             head_dim=128,
@@ -295,14 +410,14 @@ def get_fse_test_model_config(
             hidden_size=128,
             intermediate_size=32,
             moe_intermediate_size=32,
-            num_hidden_layers=1,
+            num_hidden_layers=2,
             num_attention_heads=1,
             num_key_value_heads=1,
             n_routed_experts=2,
             n_shared_experts=1,
             num_experts_per_tok=1,
-            layer_types=["linear_attention"],
-            mlp_layer_types=["sparse"],
+            layer_types=["linear_attention", "linear_attention"],
+            mlp_layer_types=["sparse", "sparse"],
             index_topk=1,
             index_kpool=1,
             pad_token_id=None,
@@ -311,6 +426,8 @@ def get_fse_test_model_config(
             mhc=False,
             quantization_config=quantization_config,
         )
+        config.mla_use_nope = False
+        config.num_nextn_predict_layers = 1
         return config, Glm5NextModel
     raise ValueError(f"Unsupported FSE test model: {model_type}")
 
@@ -402,9 +519,18 @@ def check_compatibility(
     ]
 
 
-def test_resolve_layer_fused_shared_expert_rejects_incompatible_quantization(
+@pytest.mark.parametrize(
+    ("reason", "enable_ep"),
+    [
+        ("shared experts are excluded", False),
+        ("expert parallelism is enabled", True),
+    ],
+)
+def test_resolve_layer_fused_shared_expert_logs_skip_reason(
     monkeypatch: pytest.MonkeyPatch,
     caplog: pytest.LogCaptureFixture,
+    reason: str,
+    enable_ep: bool,
 ) -> None:
     monkeypatch.setattr(
         fused_moe_utils.rocm_aiter_ops,
@@ -414,13 +540,20 @@ def test_resolve_layer_fused_shared_expert_rejects_incompatible_quantization(
     monkeypatch.setattr(
         fused_moe_utils,
         "is_shared_expert_quant_fse_compatible",
-        lambda *_, **__: (False, "shared experts are excluded"),
+        lambda *_, **__: (False, reason),
+    )
+    monkeypatch.setattr(
+        "vllm.config.get_current_vllm_config_or_none",
+        lambda: SimpleNamespace(
+            parallel_config=SimpleNamespace(enable_expert_parallel=enable_ep)
+        ),
     )
+    _print_warning_once.cache_clear()
 
     assert not fused_moe_utils.resolve_layer_fused_shared_expert(
-        ExpertsInt8Config(), "model.layers.0.mlp"
+        ExpertsInt8Config(), "model.layers.0.mlp", disable_under_ep=enable_ep
     )
-    assert "shared experts are excluded" in caplog.text
+    assert reason in caplog.text
 
 
 def test_deepseek_v4_shared_expert_fse_uses_mtp_quantization_config_prefix(
@@ -884,45 +1017,411 @@ def test_aiter_fused_moe_validates_shared_expert_arguments() -> None:
     assert validate(shared_tensor, shared_tensor, shared_tensor, shared_tensor, 0)
 
 
-def test_is_model_fused_shared_expert_compatible() -> None:
-    class MoE(nn.Module):
-        def __init__(self, enabled: bool) -> None:
-            super().__init__()
-            self.is_fused_shared_expert_enabled = enabled
+def test_deepseek_v4_shared_expert_mapper_fallback() -> None:
+    """Mixed and all-excluded DSV4 checkpoints map shared experts correctly."""
+    from vllm.models.deepseek_v4.amd.model import (
+        _make_deepseek_v4_weights_mapper,
+        _redirect_shared_expert_name,
+    )
+
+    tensor = torch.empty(0)
+    fuse_by_layer = frozenset({6})
+    # Mixed: any fused layer puts the mapper in fused mode, which leaves the
+    # checkpoint .w2 names untouched; the loader redirects per layer.
+    mapper = _make_deepseek_v4_weights_mapper("fp8", fuse_shared_experts=True)
+    for raw, expected in [
+        (
+            "layers.5.ffn.shared_experts.w2.weight",
+            "model.layers.5.ffn.shared_experts.down_proj.weight",
+        ),
+        (
+            "layers.5.ffn.shared_experts.w2.scale",
+            "model.layers.5.ffn.shared_experts.down_proj.weight_scale_inv",
+        ),
+        (
+            "layers.6.ffn.shared_experts.w2.weight",
+            "model.layers.6.ffn.experts.4.w2.weight",
+        ),
+        (
+            "layers.6.ffn.shared_experts.w2.scale",
+            "model.layers.6.ffn.experts.4.w2.weight_scale_inv",
+        ),
+    ]:
+        [(mapped, _)] = list(mapper.apply([(raw, tensor)]))
+        assert _redirect_shared_expert_name(mapped, fuse_by_layer, 4) == expected
+
+    # All layers excluded (fused_shared_expert_layers == frozenset()) selects
+    # the unfused mapper, which renames .w2 -> .down_proj itself.
+    mapper = _make_deepseek_v4_weights_mapper("fp8", fuse_shared_experts=False)
+    for raw, expected in [
+        (
+            "layers.5.ffn.shared_experts.w2.weight",
+            "model.layers.5.ffn.shared_experts.down_proj.weight",
+        ),
+        (
+            "layers.5.ffn.shared_experts.w2.scale",
+            "model.layers.5.ffn.shared_experts.down_proj.weight_scale_inv",
+        ),
+    ]:
+        [(mapped, _)] = list(mapper.apply([(raw, tensor)]))
+        assert _redirect_shared_expert_name(mapped, frozenset(), 4) == expected
+
 
-    class Layer(nn.Module):
-        def __init__(self, enabled: bool) -> None:
-            super().__init__()
-            self.mlp = MoE(enabled)
+class _FakeMoE(nn.Module):
+    def __init__(self, enabled: bool) -> None:
+        super().__init__()
+        self.is_fused_shared_expert_enabled = enabled
 
-    enabled_layers = nn.ModuleList([Layer(True)])
-    disabled_layers = nn.ModuleList([Layer(False)])
-    mixed_layers = nn.ModuleList([Layer(True), Layer(False)])
-    empty_layers = nn.ModuleList()
-    pipeline_layers = nn.ModuleList([Layer(True), PPMissingLayer()])
 
-    assert fused_moe_utils.is_model_fused_shared_expert_compatible(
-        enabled_layers, MoE, "mlp"
+def _make_fake_moe_model(
+    layer_names: list[str], enabled: list[bool], use_dict: bool
+) -> nn.Module:
+    layers = []
+    for flag in enabled:
+        layer = nn.Module()
+        layer.mlp = _FakeMoE(flag)
+        layers.append(layer)
+    model = nn.Module()
+    model.layers = (
+        nn.ModuleDict(dict(zip(layer_names, layers, strict=True)))
+        if use_dict
+        else nn.ModuleList(layers)
     )
-    assert not fused_moe_utils.is_model_fused_shared_expert_compatible(
-        disabled_layers, MoE, "mlp"
+    return model
+
+
+def test_fused_shared_expert_layers_matches_container_index_space() -> None:
+    """The fused-layer set follows the module tree, not a hard-coded offset.
+
+    ``nn.ModuleList`` MTP layers (Qwen3 Next/3.5, Qwen4 Exp) are 0-based,
+    while ``nn.ModuleDict`` MTP layers (GLM4, DeepSeek) are keyed by absolute
+    layer index. Both must redirect without a per-loader offset.
+    """
+    list_model = _make_fake_moe_model(["0", "1"], [True, False], use_dict=False)
+    dict_model = _make_fake_moe_model(["28", "29"], [False, True], use_dict=True)
+    list_layers = get_fused_shared_expert_layers(list_model, _FakeMoE)
+    dict_layers = get_fused_shared_expert_layers(dict_model, _FakeMoE)
+    assert list_layers == frozenset({0})
+    assert dict_layers == frozenset({29})
+
+    weight = torch.arange(8.0).reshape(4, 2)
+
+    def redirected(name: str, fused_shared_expert_layers: frozenset[int]) -> bool:
+        [(out_name, _)] = list(
+            maybe_fuse_shared_experts(
+                [(name, weight)],
+                n_routed_experts=4,
+                n_shared_experts=1,
+                fused_shared_expert_layers=fused_shared_expert_layers,
+                ckpt_prefix="mlp.shared_expert",
+            )
+        )
+        return ".experts." in out_name
+
+    # ModuleList: 0-based names, including the Qwen4 Exp MTP case where the
+    # quant-config prefix uses an absolute index but the weight names do not.
+    assert redirected("layers.0.mlp.shared_expert.down_proj.weight", list_layers)
+    assert not redirected("layers.1.mlp.shared_expert.down_proj.weight", list_layers)
+    # ModuleDict: absolute names, including the MTP layer offset.
+    assert redirected("layers.29.mlp.shared_expert.down_proj.weight", dict_layers)
+    assert not redirected("layers.28.mlp.shared_expert.down_proj.weight", dict_layers)
+
+
+@pytest.mark.parametrize(
+    ("enable_expert_parallel", "match"),
+    [
+        (True, "expert parallelism is not tested"),
+        (False, "nixl_ep"),
+    ],
+)
+def test_mixed_fused_shared_experts_rejected_with_ep_or_nixl_ep(
+    enable_expert_parallel: bool, match: str
+) -> None:
+    """Mixed FSE must be rejected with EP or a shared-workspace all2all backend.
+
+    With EP, per-layer geometry differences are not validated. ``nixl_ep``
+    additionally keeps a single process-wide buffer initialized from the first
+    MoE layer's ``num_experts_per_rank`` and ignores later geometries. A mixed
+    model has fused layers with ``n_routed + n_shared`` experts and non-fused
+    layers with ``n_routed``, so reusing that buffer would silently map experts
+    to the wrong slots. Construction must fail cleanly instead of running the
+    untested mixed path.
+    """
+    model = _make_fake_moe_model(["0", "1"], [True, False], use_dict=False)
+
+    vllm_config = VllmConfig()
+    vllm_config.parallel_config.enable_expert_parallel = enable_expert_parallel
+    vllm_config.parallel_config.data_parallel_size = 2
+    vllm_config.parallel_config.all2all_backend = "nixl_ep"
+
+    with (
+        set_current_vllm_config(vllm_config),
+        pytest.raises(NotImplementedError, match=match),
+    ):
+        get_fused_shared_expert_layers(model, _FakeMoE)
+
+
+def _assert_mtp_fused_shared_expert_redirect(
+    mtp: nn.Module, ckpt_prefix: str, moe: nn.Module, n_routed_experts: int
+) -> None:
+    """Load a fused MTP shared expert and verify it lands in the appended slot.
+
+    The MTP loaders see different name/index spaces (0-based ``ModuleList`` or
+    absolute ``ModuleDict``); using the live module map must redirect either.
+    """
+    packed_down = torch.full((128, 16), 3, dtype=torch.uint8)
+    down_scale = torch.full((128, 1), 127, dtype=torch.uint8)
+    mtp.load_weights(
+        [
+            (f"{ckpt_prefix}.down_proj.weight", packed_down),
+            (f"{ckpt_prefix}.down_proj.weight_scale", down_scale),
+        ]
     )
-    assert not fused_moe_utils.is_model_fused_shared_expert_compatible(
-        empty_layers, MoE, "mlp"
+    torch.testing.assert_close(
+        moe.experts.routed_experts.w2_weight[n_routed_experts][:128, :16].float(),
+        packed_down.float(),
     )
-    assert fused_moe_utils.is_model_fused_shared_expert_compatible(
-        pipeline_layers, MoE, "mlp"
+
+
+def _assert_mtp_shared_expert_not_fused(
+    mtp: nn.Module, ckpt_prefix: str, moe: nn.Module, n_routed_experts: int
+) -> None:
+    """Load an unfused MTP shared expert and verify no fused slot is appended.
+
+    Without AITER the shared expert stays a separate MLP, so the routed tensor must
+    keep exactly ``n_routed_experts`` rows.
+    """
+    packed_down = torch.full((128, 16), 3, dtype=torch.uint8)
+    down_scale = torch.full((128, 1), 127, dtype=torch.uint8)
+    mtp.load_weights(
+        [
+            (f"{ckpt_prefix}.down_proj.weight", packed_down),
+            (f"{ckpt_prefix}.down_proj.weight_scale", down_scale),
+        ]
     )
-    with pytest.raises(NotImplementedError, match="1 enabled and 1 disabled layers"):
-        fused_moe_utils.is_model_fused_shared_expert_compatible(
-            mixed_layers, MoE, "mlp"
+    assert moe.experts.routed_experts.w2_weight.shape[0] == n_routed_experts
+
+
+def _make_deepseek_v4_shared_expert_checkpoint(
+    layer: nn.Module, idx: int, fused: bool, exclude: bool
+) -> tuple[list[tuple[str, torch.Tensor]], list[tuple[torch.Tensor, torch.Tensor]]]:
+    """Build DSV4's FP8 shared-expert / MXFP4 routed-slot checkpoint pairs."""
+    ffn = layer.ffn
+    if not fused:
+        dtype = torch.bfloat16 if exclude else torch.float8_e4m3fn
+        hidden = ffn.shared_experts.gate_up_proj.weight.shape[1]
+        intermediate = ffn.shared_experts.down_proj.weight.shape[1]
+        gate_up = torch.full((2 * intermediate, hidden), idx + 1, dtype=dtype)
+        down = torch.full((hidden, intermediate), idx + 1, dtype=dtype)
+        return (
+            [
+                (f"layers.{idx}.ffn.shared_experts.gate_up_proj.weight", gate_up),
+                (f"layers.{idx}.ffn.shared_experts.down_proj.weight", down),
+            ],
+            [
+                (ffn.shared_experts.gate_up_proj.weight, gate_up),
+                (ffn.shared_experts.down_proj.weight, down),
+            ],
         )
 
+    routed = ffn.experts.routed_experts
+    slot = routed.w13_weight.shape[0] - 1
+    w13_rows, hidden_half = routed.w13_weight[slot].shape
+    intermediate = w13_rows // 2
+    hidden = hidden_half * 2
+    w1 = torch.full((intermediate, hidden_half), idx + 1, dtype=torch.uint8)
+    w3 = torch.full((intermediate, hidden_half), idx + 1, dtype=torch.uint8)
+    w2 = torch.full((hidden, intermediate // 2), idx + 1, dtype=torch.uint8)
+    return (
+        [
+            (f"layers.{idx}.ffn.shared_experts.w1.weight", w1),
+            (f"layers.{idx}.ffn.shared_experts.w3.weight", w3),
+            (f"layers.{idx}.ffn.shared_experts.w2.weight", w2),
+        ],
+        [
+            (routed.w13_weight[slot].chunk(2, dim=0)[0], w1),
+            (routed.w13_weight[slot].chunk(2, dim=0)[1], w3),
+            (routed.w2_weight[slot], w2),
+        ],
+    )
+
+
+def _build_mtp_model(
+    mtp_cls: Any,
+    vllm_config: VllmConfig,
+    config: object,
+    *,
+    with_text_config: bool = False,
+    with_target_model: bool = False,
+) -> nn.Module:
+    draft_model_config: dict[str, object] = {"hf_config": config}
+    if with_text_config:
+        draft_model_config["hf_text_config"] = config
+    speculative_config: dict[str, object] = {
+        "draft_model_config": SimpleNamespace(**draft_model_config),
+        "method": "mtp",
+        "speculator_config": SimpleNamespace(speculator_class="MTP"),
+        "parallel_drafting": False,
+        "enable_adaptive_verification": False,
+        "use_dspark": lambda: False,
+        "use_multi_module_mtp": lambda: False,
+    }
+    if with_target_model:
+        speculative_config["target_model_config"] = vllm_config.model_config
+    vllm_config.speculative_config = SimpleNamespace(**speculative_config)
+    return mtp_cls(vllm_config=vllm_config)
+
+
+_MTP_FACTORIES: dict[str, tuple[str, str, dict[str, bool]]] = {
+    "deepseek_v4": ("vllm.models.deepseek_v4.amd.mtp", "DeepSeekV4MTP", {}),
+    "glm4_moe": (
+        "vllm.model_executor.models.glm4_moe_mtp",
+        "Glm4MoeMTP",
+        {},
+    ),
+    "glm4_moe_lite": (
+        "vllm.model_executor.models.glm4_moe_lite_mtp",
+        "Glm4MoeLiteMTP",
+        {},
+    ),
+    "deepseek_v32": (
+        "vllm.models.deepseek_v32.amd.mtp",
+        "DeepseekV32MTP",
+        {},
+    ),
+    "glm5_next": (
+        "vllm.models.glm5next.common.mtp",
+        "Glm5NextMTP",
+        {"with_text_config": True},
+    ),
+    "deepseek_v2": (
+        "vllm.model_executor.models.deepseek_mtp",
+        "DeepSeekMTP",
+        {},
+    ),
+    "qwen3_next": (
+        "vllm.model_executor.models.qwen3_next_mtp",
+        "Qwen3NextMTP",
+        {"with_text_config": True, "with_target_model": True},
+    ),
+    "qwen3_5": (
+        "vllm.model_executor.models.qwen3_5_mtp",
+        "Qwen3_5MoeMTP",
+        {"with_text_config": True, "with_target_model": True},
+    ),
+}
+
+
+def _build_mtp_for_model(
+    model_type: str,
+    vllm_config: VllmConfig,
+    config: object,
+    enabled: bool,
+) -> nn.Module | None:
+    factory = _MTP_FACTORIES.get(model_type)
+    if not enabled or factory is None:
+        return None
+    module_name, class_name, kwargs = factory
+    module = importlib.import_module(module_name)
+    return _build_mtp_model(getattr(module, class_name), vllm_config, config, **kwargs)
+
+
+# TODO: move to conftest.py.
+class _SingleRankGroup:
+    """No-op single-rank stand-in for a compiled ``GroupCoordinator``."""
+
+    world_size = 1
+    rank_in_group = 0
+    local_rank = 0
+    is_first_rank = True
+    is_last_rank = True
+
+    @property
+    def device_group(self):
+        return self
+
+    @property
+    def cpu_group(self):
+        return self
+
+    @property
+    def device_communicator(self):
+        return self
+
+    def rank(self) -> int:
+        return 0
+
+    def size(self) -> int:
+        return 1
+
+    def all_reduce(self, input_: torch.Tensor, *args: Any, **kwargs: Any):
+        return input_
+
+    def all_gather(self, input_: torch.Tensor, dim: int = 0):
+        return input_
+
+    def all_gatherv(self, inputs: Any, dim: int = 0, sizes: Any = None):
+        return inputs
+
+    def reduce_scatter(self, input_: torch.Tensor, dim: int = 0):
+        return input_
+
+    def reduce_scatterv(self, input_: Any, dim: int = 0, sizes: Any = None):
+        return input_
+
+    def broadcast(self, input_: torch.Tensor, src: int = 0):
+        return input_
+
+    def send(self, *args: Any, **kwargs: Any) -> None:
+        return None
+
+    def recv(self, *args: Any, **kwargs: Any) -> None:
+        return None
+
+    def barrier(self) -> None:
+        return None
+
+    def destroy(self) -> None:
+        return None
+
+
+@pytest.fixture
+def single_rank_parallel_state(monkeypatch: pytest.MonkeyPatch):
+    """Single rank parallel groups without initializing a process group.
+
+    This is faster than the `dist_init` fixture.
+    """
+    for name in ("_WORLD", "_TP", "_ETP", "_DCP", "_PP", "_DP", "_EP", "_PCP"):
+        monkeypatch.setattr(
+            f"vllm.distributed.parallel_state.{name}", _SingleRankGroup()
+        )
+    monkeypatch.setattr(
+        "vllm.model_executor.parameter.get_tensor_model_parallel_rank", lambda: 0
+    )
+    monkeypatch.setattr(
+        "vllm.model_executor.parameter.get_tensor_model_parallel_world_size", lambda: 1
+    )
+
 
 @pytest.mark.parametrize(
     "model_type",
-    ["minimax_m3", "deepseek_v4", "qwen3_5", "glm4_moe", "deepseek_v2", "glm5_next"],
+    [
+        "minimax_m3",
+        "deepseek_v4",
+        "deepseek_v4_cpu",
+        "qwen3_next",
+        "qwen3_5",
+        "alice_ai",
+        "glm4_moe",
+        "deepseek_v2",
+        "AXK1",
+        "glm4_moe_lite",
+        "deepseek_v32",
+        "glm5_next",
+        "qwen4_exp",
+    ],
 )
+@pytest.mark.parametrize("shared_expert_precision", ["all_mxfp4", "mixed"])
 @pytest.mark.parametrize(
     ("use_fse", "exclude"),
     [
@@ -933,17 +1432,82 @@ def __init__(self, enabled: bool) -> None:
 )
 def test_models_fse_init(
     model_type: str,
+    shared_expert_precision: str,
     use_fse: bool,
     exclude: list[str],
-    dist_init: None,
+    single_rank_parallel_state: None,
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
-    """Model construction resolves FSE consistently with Quark quantization."""
+    """Construction and loading preserve each shared expert's precision."""
+    _print_warning_once.cache_clear()
+    per_layer_fse = model_type != "deepseek_v4_cpu"
+    moe_attr = {
+        "minimax_m3": "block_sparse_moe",
+        "deepseek_v4": "ffn",
+    }.get(model_type, "mlp")
+    shared_name = (
+        "shared_experts"
+        if model_type
+        in {
+            "minimax_m3",
+            "deepseek_v4",
+            "glm4_moe",
+            "deepseek_v2",
+            "AXK1",
+            "glm4_moe_lite",
+            "deepseek_v32",
+            "glm5_next",
+        }
+        else "shared_expert"
+    )
+    if shared_expert_precision == "mixed" and not per_layer_fse:
+        pytest.skip(
+            "The mixed Quark two-layer setup does not apply to this architecture"
+        )
     quantization_config: dict[str, Any] = (
-        get_deepseek_v4_quark_config(["layers.0.ffn.shared_experts"] if exclude else [])
+        get_deepseek_v4_quark_config(
+            [f"layers.{idx}.ffn.shared_experts" for idx in range(2)] if exclude else []
+        )
         if model_type == "deepseek_v4"
-        else {**_QUARK_FSE_CONFIG, "exclude": exclude}
-    )
+        else {**deepcopy(_QUARK_FSE_CONFIG), "exclude": exclude}
+    )
+    if shared_expert_precision == "mixed":
+        mxfp8_config = deepcopy(quantization_config["global_quant_config"])
+        mxfp8_config["input_tensors"].update(
+            dtype="fp8_e4m3",
+            symmetric=True,
+            scale_type=None,
+            scale_format=None,
+            scale_calculation_mode=None,
+            observer_cls="PerGroupMinMaxObserver",
+            zero_point_type="int32",
+        )
+        mxfp8_config["weight"].update(
+            dtype="fp8_e4m3",
+            qscheme="per_block",
+            ch_axis=None,
+            group_size=None,
+            block_size=[1, 32],
+            symmetric=True,
+            scale_type="float8_e8m0fnu",
+            scale_format=None,
+            scale_calculation_mode=None,
+            observer_cls="PerBlock2DMinMaxObserver",
+            zero_point_type="int32",
+        )
+        if model_type == "deepseek_v4":
+            quantization_config.setdefault("layer_quant_config", {})[
+                "layers.1.ffn.shared_experts.w1"
+            ] = deepcopy(mxfp8_config)
+        else:
+            quantization_config.setdefault("layer_quant_config", {}).update(
+                {
+                    f"model.layers.1.{moe_attr}.{shared_name}.{projection}": deepcopy(
+                        mxfp8_config
+                    )
+                    for projection in ("gate_proj", "up_proj", "down_proj")
+                }
+            )
 
     config, model_constructor = get_fse_test_model_config(
         model_type, quantization_config
@@ -961,11 +1525,12 @@ def test_models_fse_init(
         quantization_config=None,
         runner_type="generate",
         is_moe=True,
+        use_mla=False,
         logits_processors=None,
         rswa_window=None,
     )
     vllm_config.parallel_config.enable_expert_parallel = False
-    if model_type == "deepseek_v4":
+    if model_type in ("deepseek_v4", "deepseek_v4_cpu"):
         from vllm.models.deepseek_v4.quant_config import DeepseekV4FP8Config
 
         vllm_config.cache_config.cache_dtype = "fp8_ds_mla"
@@ -975,6 +1540,10 @@ def test_models_fse_init(
         )
     else:
         vllm_config.quant_config = QuarkConfig(quantization_config)
+        if per_layer_fse:
+            vllm_config.quant_config.packed_modules_mapping = {
+                "gate_up_proj": ["gate_proj", "up_proj"]
+            }
     if model_type == "glm5_next":
         from vllm.models.glm5next.common import model as glm5_next_model
 
@@ -985,18 +1554,40 @@ def test_models_fse_init(
         monkeypatch.setattr(
             glm5_next_model, "_fused_shared_experts_tuned", lambda _: True
         )
+        from vllm.models.glm5next.common import mtp as glm5_next_mtp_model
+
+        # The MTP predictor asserts on the concrete MLA attention class and
+        # reads ``mla_attn``, so patch its imported name to the same stub.
+        monkeypatch.setattr(
+            glm5_next_mtp_model, "Glm5NextMLAAttention", _StubGlm5NextAttention
+        )
+    if model_type == "glm4_moe_lite":
+        from vllm.model_executor.models import glm4_moe_lite as glm4_moe_lite_model
+
+        # Only the MoE/shared-expert wiring is under test here; the MLA backend
+        # does not accept this test's default dtype.
+        for attention in ("Glm4MoeLiteAttention", "Glm4MoeLiteMLAAttention"):
+            monkeypatch.setattr(glm4_moe_lite_model, attention, _StubGlm5NextAttention)
+    if model_type == "deepseek_v32":
+        from vllm.models.deepseek_v32.amd import model as deepseek_v32_model
+
+        # The sparse-MLA attention needs a real MLA backend; only the MoE
+        # wiring is under test here.
+        monkeypatch.setattr(
+            deepseek_v32_model, "DeepseekV32MLAAttention", _StubGlm5NextAttention
+        )
+    if model_type == "qwen4_exp":
+        from vllm.models.qwen4_exp.amd import model as qwen4_exp_model
+
+        # Only the MoE/shared-expert wiring is under test here.
+        monkeypatch.setattr(
+            qwen4_exp_model, "Qwen3NextAttention", _StubGlm5NextAttention
+        )
 
     import vllm.envs as envs
     from vllm._aiter_ops import rocm_aiter_ops
 
-    if model_type == "deepseek_v4":
-        from vllm.models.deepseek_v4.amd import model as deepseek_v4_model
-
-        warning_logger = deepseek_v4_model.logger
-    elif model_type == "minimax_m3":
-        warning_logger = minimax_m3_model.logger
-    else:
-        warning_logger = fused_moe_utils.logger
+    warning_logger = fused_moe_utils.logger
 
     with monkeypatch.context() as mp:
         mp.setenv("VLLM_ROCM_USE_AITER", str(use_fse))
@@ -1010,33 +1601,196 @@ def test_models_fse_init(
         importlib.reload(envs)
         mp.setattr(envs, "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS", use_fse)
         rocm_aiter_ops.refresh_env_variables()
+        from vllm.model_executor.layers.fused_moe.experts import rocm_aiter_moe
+
+        mp.setattr(rocm_aiter_moe, "aiter_topK_meta_data", None)
+        rocm_aiter_moe.init_aiter_topK_meta_data.cache_clear()
         aiter_fse_enabled = bool(rocm_aiter_ops.is_fusion_moe_shared_experts_enabled())
-        # These AMD-specific models currently use the raw environment flag.
-        # and do not rely on `rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()`.
-        fse_enabled = use_fse and (
-            aiter_fse_enabled or model_type in {"deepseek_v4", "minimax_m3"}
-        )
+        fse_enabled = use_fse and aiter_fse_enabled
 
         with (
             patch.object(warning_logger, "warning") as warning,
             set_current_vllm_config(vllm_config),
         ):
-            model = model_constructor(vllm_config=vllm_config)
-            mtp = None
-            if model_type == "deepseek_v4" and use_fse and not exclude:
-                from vllm.models.deepseek_v4.amd.mtp import (
-                    DeepSeekV4MTP,
+            model = model_constructor(
+                vllm_config=vllm_config,
+                **({"prefix": "model"} if per_layer_fse else {}),
+            )
+            mtps = {
+                model_type: _build_mtp_for_model(
+                    model_type, vllm_config, config, use_fse and not exclude
                 )
-
-                vllm_config.speculative_config = SimpleNamespace(
-                    draft_model_config=SimpleNamespace(hf_config=config),
-                    method="mtp",
-                    parallel_drafting=False,
-                    enable_adaptive_verification=False,
-                    use_dspark=lambda: False,
+            }
+            mtp = mtps.get("deepseek_v4")
+        if per_layer_fse:
+            expected_policy = {
+                f"layers.0.{moe_attr}": bool(fse_enabled and not exclude),
+                f"layers.1.{moe_attr}": bool(
+                    fse_enabled
+                    and not exclude
+                    and shared_expert_precision == "all_mxfp4"
+                ),
+            }
+            assert model.fused_shared_expert_layers == frozenset(
+                idx for idx in range(2) if expected_policy[f"layers.{idx}.{moe_attr}"]
+            )
+            if any(expected_policy.values()) and not all(expected_policy.values()):
+                with pytest.raises(NotImplementedError, match="EPLB"):
+                    get_fused_shared_expert_layers(
+                        model,
+                        type(getattr(model.layers[0], moe_attr)),
+                        enable_eplb=True,
+                    )
+            for idx, layer in enumerate(model.layers):
+                fused = expected_policy[f"layers.{idx}.{moe_attr}"]
+                moe = getattr(layer, moe_attr)
+                assert (getattr(moe, shared_name) is None) is fused
+                assert moe.experts.expert_map_manager.num_fused_shared_experts == int(
+                    fused
+                )
+                assert (
+                    moe.experts.moe_config.is_fused_shared_expert_requested
+                    is aiter_fse_enabled
+                )
+                if not fused and not exclude and model_type != "deepseek_v4":
+                    scheme = getattr(moe, shared_name).gate_up_proj.scheme
+                    expected_dtype = (
+                        "mxfp8_e4m3"
+                        if idx == 1 and shared_expert_precision == "mixed"
+                        else "mxfp4"
+                    )
+                    assert scheme.weight_dtype == expected_dtype
+
+            checkpoint = []
+            expected_weights = []
+            for idx, layer in enumerate(model.layers):
+                fused = expected_policy[f"layers.{idx}.{moe_attr}"]
+                if model_type == "deepseek_v4":
+                    dsv4_checkpoint, dsv4_expected = (
+                        _make_deepseek_v4_shared_expert_checkpoint(
+                            layer, idx, fused, exclude
+                        )
+                    )
+                    checkpoint.extend(dsv4_checkpoint)
+                    expected_weights.extend(dsv4_expected)
+                    continue
+                high_precision = idx == 1 and shared_expert_precision == "mixed"
+                dtype = (
+                    torch.bfloat16
+                    if exclude
+                    else torch.float8_e4m3fn
+                    if high_precision
+                    else torch.uint8
+                )
+                pack_factor = 1 if exclude or high_precision else 2
+                for projection in ("gate_proj", "up_proj", "down_proj"):
+                    shape = (
+                        (128, 32 // pack_factor)
+                        if projection == "down_proj"
+                        else (32, 128 // pack_factor)
+                    )
+                    weight = torch.full(shape, idx + 1).to(dtype)
+                    name = f"layers.{idx}.{moe_attr}.{shared_name}.{projection}"
+                    checkpoint.append((f"{name}.weight", weight))
+                    if not exclude:
+                        scale_shape = (128, 1) if projection == "down_proj" else (32, 4)
+                        checkpoint.append(
+                            (
+                                f"{name}.weight_scale",
+                                torch.full(scale_shape, 127, dtype=torch.uint8),
+                            )
+                        )
+                    if fused:
+                        routed = getattr(layer, moe_attr).experts.routed_experts
+                        target = (
+                            routed.w2_weight[2]
+                            if projection == "down_proj"
+                            else routed.w13_weight[2].chunk(2, dim=0)[
+                                int(projection == "up_proj")
+                            ]
+                        )
+                    else:
+                        shared = getattr(getattr(layer, moe_attr), shared_name)
+                        target = (
+                            shared.down_proj.weight
+                            if projection == "down_proj"
+                            else shared.gate_up_proj.weight.chunk(2, dim=0)[
+                                int(projection == "up_proj")
+                            ]
+                        )
+                    expected_weights.append((target, weight))
+                if shared_name == "shared_expert":
+                    checkpoint.append(
+                        (
+                            f"layers.{idx}.{moe_attr}.shared_expert_gate.weight",
+                            torch.full((1, 128), idx + 1, dtype=torch.bfloat16),
+                        )
+                    )
+            with set_current_vllm_config(vllm_config):
+                loaded = model.load_weights(checkpoint)
+                base = config.num_hidden_layers
+                mtp_cases = [
+                    (
+                        mtp_model,
+                        f"model.layers.{base}.mlp.{shared}",
+                        mtp_model.model.layers[str(base)].mtp_block.mlp,
+                        getattr(config, "n_routed_experts", 0),
+                    )
+                    for mtp_model, shared in (
+                        (mtps.get("glm4_moe"), "shared_experts"),
+                        (mtps.get("glm4_moe_lite"), "shared_experts"),
+                        (mtps.get("deepseek_v2"), "shared_experts"),
+                        (mtps.get("deepseek_v32"), "shared_experts"),
+                        (mtps.get("glm5_next"), "shared_experts"),
+                    )
+                    if mtp_model is not None
+                ]
+                mtp_cases += [
+                    (
+                        mtp_model,
+                        "mtp.layers.0.mlp.shared_expert",
+                        mtp_model.model.layers[0].mlp,
+                        getattr(config, "num_experts", 0),
+                    )
+                    for mtp_model in (
+                        mtps.get("qwen3_next"),
+                        mtps.get("qwen3_5"),
+                    )
+                    if mtp_model is not None
+                ]
+                for mtp_model, ckpt_prefix, moe, n_routed in mtp_cases:
+                    if fse_enabled and not exclude:
+                        _assert_mtp_fused_shared_expert_redirect(
+                            mtp_model, ckpt_prefix, moe, n_routed
+                        )
+                    else:
+                        _assert_mtp_shared_expert_not_fused(
+                            mtp_model, ckpt_prefix, moe, n_routed
+                        )
+            for idx in range(2):
+                prefix = f"layers.{idx}.{moe_attr}"
+                expected_name = (
+                    f"{prefix}.experts.routed_experts.w2_weight"
+                    if expected_policy[prefix]
+                    else f"{prefix}.{shared_name}.down_proj.weight"
+                )
+                assert expected_name in loaded
+                if shared_name == "shared_expert":
+                    assert f"{prefix}.shared_expert_gate.weight" in loaded
+            for target, weight in expected_weights:
+                torch.testing.assert_close(
+                    target[: weight.shape[0], : weight.shape[1]].float(), weight.float()
                 )
-                mtp = DeepSeekV4MTP(vllm_config=vllm_config)
-        assert model.is_fused_shared_expert_enabled is (fse_enabled and not exclude)
+        elif model_type == "deepseek_v4_cpu":
+            # CPU never supports fused shared experts; the separate MLP is
+            # always built and no routed slot is appended.
+            for layer in model.layers:
+                assert layer.ffn.shared_experts is not None
+                assert (
+                    layer.ffn.experts.expert_map_manager.num_fused_shared_experts == 0
+                )
+        else:
+            assert model.is_fused_shared_expert_enabled is (fse_enabled and not exclude)
         if model_type == "glm5_next":
             assert (model.layers[0].mlp.shared_experts is None) is (
                 fse_enabled and not exclude
@@ -1046,25 +1800,94 @@ def test_models_fse_init(
         # so should just raise a warning.
         if mtp is not None:
             assert not mtp.model.layers[
-                "1"
+                str(config.num_hidden_layers)
             ].mtp_block.ffn.is_fused_shared_expert_enabled
+            if aiter_fse_enabled:
+                warning.assert_called_once()
+                assert (
+                    "DeepSeek-V4/V4.1 shared experts at <shared_expert>"
+                    in (warning.call_args.args[1])
+                )
+            else:
+                # Without AITER nothing fuses, so no skip reason is reported.
+                warning.assert_not_called()
+        if aiter_fse_enabled and shared_expert_precision == "mixed" and not exclude:
             warning.assert_called_once()
-            assert (
-                "DeepSeek-V4/V4.1 shared experts at mtp.0.ffn.shared_experts"
-                in (warning.call_args.args[1])
-            )
-        if aiter_fse_enabled and exclude:
+        if aiter_fse_enabled and shared_expert_precision == "mixed" and not exclude:
             warning.assert_called_once()
-            assert (
-                "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS is enabled"
-                in (warning.call_args.args[0])
-            )
-            assert "excludes shared experts" in warning.call_args.args[1]
+            if model_type == "deepseek_v4":
+                assert "are not MXFP4" in warning.call_args.args[1]
+            else:
+                assert (
+                    "different quantization configurations" in warning.call_args.args[1]
+                )
+        for mtp_model in (
+            mtps.get("glm4_moe"),
+            mtps.get("deepseek_v2"),
+            mtps.get("glm4_moe_lite"),
+            mtps.get("deepseek_v32"),
+            mtps.get("glm5_next"),
+        ):
+            if mtp_model is not None:
+                assert mtp_model.fused_shared_expert_layers == (
+                    frozenset({config.num_hidden_layers})
+                    if fse_enabled
+                    else frozenset()
+                )
+        if aiter_fse_enabled and exclude:
+            if model_type == "deepseek_v4_cpu":
+                # CPU has no FSE path, so the resolver never warns.
+                warning.assert_not_called()
+            else:
+                assert warning.call_count == 1
+                assert (
+                    "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS is enabled"
+                    in (warning.call_args.args[0])
+                )
+                assert "excludes shared experts" in warning.call_args.args[1]
 
+    rocm_aiter_moe.init_aiter_topK_meta_data.cache_clear()
     importlib.reload(envs)
     rocm_aiter_ops.refresh_env_variables()
 
 
+def _modelopt_mxfp8_config(exclude_modules=(), moe_quant_algo="NVFP4"):
+    from vllm.model_executor.layers.quantization.modelopt import ModelOptMxFp8Config
+
+    return ModelOptMxFp8Config(
+        is_checkpoint_mxfp8_serialized=True,
+        kv_cache_quant_algo=None,
+        exclude_modules=list(exclude_modules),
+        moe_quant_algo=moe_quant_algo,
+    )
+
+
+@pytest.mark.skip_global_cleanup
+@pytest.mark.parametrize(
+    ("moe_quant_algo", "expected"),
+    [
+        # NVFP4 routed experts cannot share a slot layout with the MXFP8 shared
+        # expert without online requantization of both.
+        ("NVFP4", False),
+        # A plain MXFP8 checkpoint (both sides MXFP8).
+        (None, True),
+    ],
+)
+def test_modelopt_mxfp8_shared_expert_fse_keeps_the_stored_format(
+    moe_quant_algo: str | None, expected: bool
+) -> None:
+    """Without online requantization, the shared expert is fused only in the
+    format it is stored in."""
+    compatible, reason = is_shared_expert_quant_fse_compatible(
+        _modelopt_mxfp8_config(moe_quant_algo=moe_quant_algo),
+        "model.layers.0.mlp.experts",
+        "model.layers.0.mlp.shared_experts",
+    )
+
+    assert compatible is expected
+    assert (reason is None) is expected
+
+
 @pytest.mark.parametrize(
     ("setting", "value", "reason"),
     [
@@ -1194,7 +2017,10 @@ def test_online_targets_off_shared_expert_defer_to_checkpoint(
     )
 
 
-def test_online_targeting_only_one_shared_projection_rejects_fse() -> None:
+@pytest.mark.skip_global_cleanup
+def test_online_only_one_shared_projection_matching_keys() -> None:
+    """An online target on part of the shared expert is fine as long as every
+    projection runs with the routed experts' keys (MXFP4 here)."""
     shared_prefix = "model.layers.0.mlp.shared_expert"
     quant_config = QuarkConfig({**_QUARK_FSE_CONFIG, "exclude": []})
     quant_config.online_quantization_config = OnlineQuantizationConfig(
@@ -1205,11 +2031,7 @@ def test_online_targeting_only_one_shared_projection_rejects_fse() -> None:
         quant_config,
         "model.layers.0.mlp.experts",
         shared_prefix,
-    ) == (
-        False,
-        "online quantization targets only part of the shared expert at "
-        f"{shared_prefix}",
-    )
+    ) == (True, None)
 
 
 def test_online_fp8_shared_expert_rejects_quark_mxfp4_fse() -> None:
@@ -1219,14 +2041,15 @@ def test_online_fp8_shared_expert_rejects_quark_mxfp4_fse() -> None:
         QuantizationConfigArgs(targets={f"{shared_prefix}*": "fp8_per_tensor"})
     )
 
-    assert is_shared_expert_quant_fse_compatible(
+    compatible, reason = is_shared_expert_quant_fse_compatible(
         quant_config,
         "model.layers.0.mlp.experts",
         shared_prefix,
-    ) == (
-        False,
-        "online shared-expert quantization keys do not match the routed expert keys",
     )
+    assert not compatible
+    assert reason is not None
+    assert reason.startswith("shared expert (weight, activation) keys ")
+    assert f"at {shared_prefix} do not match the routed expert keys" in reason
 
 
 def test_online_mxfp4_shared_expert_rejects_quark_weight_only_mxfp4_fse() -> None:
@@ -1246,14 +2069,15 @@ def test_online_mxfp4_shared_expert_rejects_quark_weight_only_mxfp4_fse() -> Non
         QuantizationConfigArgs(targets={f"{shared_prefix}*": "mxfp4"})
     )
 
-    assert is_shared_expert_quant_fse_compatible(
+    compatible, reason = is_shared_expert_quant_fse_compatible(
         quant_config,
         "model.layers.0.mlp.experts",
         shared_prefix,
-    ) == (
-        False,
-        "online shared-expert quantization keys do not match the routed expert keys",
     )
+    assert not compatible
+    assert reason is not None
+    assert reason.startswith("shared expert (weight, activation) keys ")
+    assert f"at {shared_prefix} do not match the routed expert keys" in reason
 
 
 @pytest.mark.parametrize(
diff --git a/tests/model_executor/layers/test_lmhead_skinny_gemm_guard.py b/tests/model_executor/layers/test_lmhead_skinny_gemm_guard.py
new file mode 100644
index 000000000000..21c9e1060796
--- /dev/null
+++ b/tests/model_executor/layers/test_lmhead_skinny_gemm_guard.py
@@ -0,0 +1,266 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""lm_head skinny-GEMM weight-size guard: core contracts.
+
+wvSplitK_hf_sml's cost at n activation rows scales as ~n passes over the
+weight, so past the last-level (memory-side) cache its effective
+bandwidth collapses. At the DSV4.1-Flash TP=4 serving point the draft
+vocab-parallel lm_head shard ([32320, 5120] bf16, ~315 MiB, one call per
+decode step) hit this collapse. The guard keeps every weight above
+VLLM_SKINNY_GEMM_MAX_WEIGHT_BYTES on the generic GEMM fall-through at
+0 < n <= 5; the default cap is the memory-side cache size, read once
+from the hardware (256 MiB on MI355X, 128 MiB on MI300X-class parts),
+with a 256 MiB fallback where that cannot be read. The swap is NOT
+bit-exact (different GEMM kernel, different accumulation order) -- these
+tests check the decision table, dispatch selection, cap tuning and
+hardware-derivation, and numerical agreement within bf16 GEMM tolerance.
+"""
+
+import os
+
+import pytest
+import torch
+
+import vllm.envs as envs
+from vllm.model_executor.layers.utils import (
+    _amd_l3_cache_bytes,
+    _skinny_gemm_max_weight_bytes,
+    _skinny_guard_warned,
+    lmhead_skinny_guard_permits,
+    rocm_unquantized_gemm_impl,
+)
+
+try:
+    from vllm import _custom_ops as ops
+
+    WVSPLITK_AVAILABLE = hasattr(ops, "wvSplitK")
+except ImportError:
+    WVSPLITK_AVAILABLE = False
+
+pytestmark = [
+    pytest.mark.skipif(not torch.cuda.is_available(), reason="GEMM kernels run on GPU"),
+    pytest.mark.skipif(
+        not WVSPLITK_AVAILABLE,
+        reason="requires the ROCm skinny-GEMM custom ops",
+    ),
+]
+
+_MAX_BYTES_VAR = "VLLM_SKINNY_GEMM_MAX_WEIGHT_BYTES"
+
+# The serving-site shape: the DSpark draft vocab-parallel lm_head shard at
+# TP=4 (draft hidden size 5120), above the 256 MiB default cap.
+BIG_SHAPE = (32320, 5120)  # 315.5 MiB bf16
+# A dense-projection-class weight: far below the cap.
+SMALL_SHAPE = (2048, 7168)  # 29 MiB bf16
+
+
+@pytest.fixture(autouse=True)
+def _clean_env(monkeypatch):
+    # These tests assert exact dispatch selection, so pin every env var
+    # that feeds it: state leaked by unrelated tests earlier in a long
+    # suite run must not change the outcome here.
+    monkeypatch.delenv(_MAX_BYTES_VAR, raising=False)
+    monkeypatch.setenv("VLLM_ROCM_USE_SKINNY_GEMM", "1")
+    monkeypatch.delenv("VLLM_BATCH_INVARIANT", raising=False)
+    _skinny_guard_warned.clear()
+    envs.disable_envs_cache()
+    _skinny_gemm_max_weight_bytes.cache_clear()
+    yield
+    _skinny_guard_warned.clear()
+    envs.disable_envs_cache()
+    _skinny_gemm_max_weight_bytes.cache_clear()
+
+
+def _weight(shape, seed=0):
+    torch.manual_seed(seed)
+    return torch.randn(*shape, dtype=torch.bfloat16, device="cuda")
+
+
+def _x(n, k, seed=1):
+    torch.manual_seed(seed)
+    return torch.randn(n, k, dtype=torch.bfloat16, device="cuda")
+
+
+def _set_cap(max_bytes):
+    # None -> the hardware-derived default (memory-side cache size, with
+    # the 256 MiB fallback); an int -> that exact cap. The accessor caches
+    # its one evaluation, so clear it after every change.
+    if max_bytes is None:
+        os.environ.pop(_MAX_BYTES_VAR, None)
+    else:
+        os.environ[_MAX_BYTES_VAR] = str(max_bytes)
+    envs.disable_envs_cache()
+    _skinny_gemm_max_weight_bytes.cache_clear()
+
+
+def _skinny_dispatch(monkeypatch, x, w):
+    """Whether rocm_unquantized_gemm_impl took the wvSplitK path, and its output.
+
+    Spied at the op boundary rather than read from torch.profiler, which
+    records nothing while another ROCProfiler SDK tool (e.g. the CI kernel
+    recorder) is registered in the process.
+    """
+    calls = 0
+    wvsplitk = ops.wvSplitK
+
+    def spy(*args):
+        nonlocal calls
+        calls += 1
+        return wvsplitk(*args)
+
+    with monkeypatch.context() as m:
+        m.setattr(ops, "wvSplitK", spy)
+        out = rocm_unquantized_gemm_impl(x, w, None)
+    return calls == 1, out
+
+
+def test_guard_decision():
+    """Decision table: bytes-gated at 0 < n <= 5; n=1 (no pathology)
+    always permits; the cap is tunable. The default (auto) cap gates the
+    315.5 MiB site shape on every part it can run on: the memory-side
+    cache is 256 MiB (MI355X) / 128 MiB (MI300X-class), and the fallback
+    is 256 MiB."""
+    small = torch.zeros(SMALL_SHAPE, dtype=torch.bfloat16)
+    big = torch.zeros(BIG_SHAPE, dtype=torch.bfloat16)
+    assert lmhead_skinny_guard_permits(small, 5)
+    for n in (2, 3, 4, 5):
+        assert not lmhead_skinny_guard_permits(big, n)
+    assert lmhead_skinny_guard_permits(big, 1), "n=1 carve-out missing"
+
+    _set_cap(1024 * 1024 * 1024)
+    assert lmhead_skinny_guard_permits(big, 5), "raised cap must permit"
+
+    _set_cap(128 * 1024 * 1024)
+    assert not lmhead_skinny_guard_permits(big, 5), "lowered cap must gate"
+
+    _set_cap(None)
+    assert not lmhead_skinny_guard_permits(big, 5), "auto cap must gate"
+
+
+# Canned amdsmi cache lists (amdsmi_get_gpu_cache_info(...)["cache"]),
+# trimmed to the fields the parse keys on; sizes are KB as amdsmi
+# reports them.
+_AMD_SMI_MI355X = [
+    # L1 data / L1 inst / L2: the levels below the boundary.
+    {
+        "cache_properties": ["DATA_CACHE", "SIMD_CACHE"],
+        "cache_level": 1,
+        "cache_size": 32,
+    },
+    {
+        "cache_properties": ["INST_CACHE", "SIMD_CACHE"],
+        "cache_level": 1,
+        "cache_size": 64,
+    },
+    {
+        "cache_properties": ["DATA_CACHE", "SIMD_CACHE"],
+        "cache_level": 2,
+        "cache_size": 4096,
+    },
+    # The memory-side cache (MALL): 256 MiB, the level the guard keys on.
+    {
+        "cache_properties": ["DATA_CACHE", "SIMD_CACHE"],
+        "cache_level": 3,
+        "cache_size": 262144,
+    },
+]
+_AMD_SMI_MI300X_CLASS = [
+    {"cache_properties": ["DATA_CACHE"], "cache_level": 2, "cache_size": 8192},
+    {"cache_properties": ["DATA_CACHE"], "cache_level": 3, "cache_size": 131072},
+]
+
+
+def test_l3_parse_selects_level3_data_cache():
+    """The parse takes the level-3 DATA_CACHE (the memory-side cache),
+    not the L2 below it, in amdsmi's KB units."""
+    assert _amd_l3_cache_bytes(_AMD_SMI_MI355X) == 256 * 1024 * 1024
+    assert _amd_l3_cache_bytes(_AMD_SMI_MI300X_CLASS) == 128 * 1024 * 1024
+    # Empty list, or instruction caches only: no L3 to key on.
+    assert _amd_l3_cache_bytes([]) is None
+    assert (
+        _amd_l3_cache_bytes(
+            [{"cache_properties": ["INST_CACHE"], "cache_level": 3, "cache_size": 64}]
+        )
+        is None
+    )
+
+
+def test_cap_env_hardware_fallback_precedence(monkeypatch):
+    """Unset env: the cap is the L3 size the library reports; an
+    unreadable L3 falls back to the documented 256 MiB; an explicit env
+    value overrides the hardware."""
+    _set_cap(None)
+    monkeypatch.setattr(
+        "vllm.model_executor.layers.utils._amd_l3_cache_bytes",
+        lambda cache_info=None: 128 * 1024 * 1024,
+    )
+    _skinny_gemm_max_weight_bytes.cache_clear()
+    assert _skinny_gemm_max_weight_bytes() == 128 * 1024 * 1024
+
+    monkeypatch.setattr(
+        "vllm.model_executor.layers.utils._amd_l3_cache_bytes",
+        lambda cache_info=None: None,
+    )
+    _skinny_gemm_max_weight_bytes.cache_clear()
+    assert _skinny_gemm_max_weight_bytes() == 256 * 1024 * 1024
+
+    _set_cap(1024 * 1024 * 1024)
+    assert _skinny_gemm_max_weight_bytes() == 1024 * 1024 * 1024
+
+
+def test_census_warning_fires_once_per_site(caplog):
+    w = torch.zeros(BIG_SHAPE, dtype=torch.bfloat16)
+    with caplog.at_level("WARNING", logger="vllm.model_executor.layers.utils"):
+        lmhead_skinny_guard_permits(w, 5)
+        lmhead_skinny_guard_permits(w, 5)  # same (shape, n): once
+        lmhead_skinny_guard_permits(w, 4)  # new n: second line
+    warns = [r for r in caplog.records if "skinny-GEMM guard engaged" in r.message]
+    assert len(warns) == 2
+    assert "32320" in warns[0].message and "n=5" in warns[0].message
+
+
+@pytest.mark.parametrize("cap_gib,expect_skinny", [(None, False), (1, True)])
+def test_n5_big_weight_dispatch(monkeypatch, cap_gib, expect_skinny):
+    """The draft-lm_head shape (n=5, >256 MiB): under the default cap the
+    guard redirects to the generic GEMM (the same path its n=6 verify twin
+    dispatches to); a raised cap keeps the stock wvSplitK path."""
+    _set_cap(cap_gib * 1024 * 1024 * 1024 if cap_gib else None)
+    w = _weight(BIG_SHAPE)
+    x = _x(5, BIG_SHAPE[1])
+    ref = torch.nn.functional.linear(x, w)
+    skinny, out = _skinny_dispatch(monkeypatch, x, w)
+    assert skinny == expect_skinny
+    torch.testing.assert_close(out, ref, rtol=2e-2, atol=2e-1)
+
+
+@pytest.mark.parametrize(
+    "shape,n,expect_skinny",
+    [
+        (BIG_SHAPE, 6, False),  # n=6 never matches the skinny branch
+        (SMALL_SHAPE, 5, True),  # small weight: bytes-gated only
+    ],
+)
+def test_guard_no_op_where_it_must_be(monkeypatch, shape, n, expect_skinny):
+    """The guard must not touch dispatch outside its window, under the
+    default cap AND a raised one: n=6 (above the skinny branch) and small
+    weights (below the byte cap)."""
+    w = _weight(shape)
+    x = _x(n, shape[1])
+    ref = torch.nn.functional.linear(x, w)
+    for cap_gib in (None, 1):
+        _set_cap(cap_gib * 1024 * 1024 * 1024 if cap_gib else None)
+        _skinny_guard_warned.clear()
+        skinny, out = _skinny_dispatch(monkeypatch, x, w)
+        assert skinny == expect_skinny, (
+            f"cap={cap_gib} GiB: skinny dispatch changed for shape={shape} n={n}"
+        )
+        torch.testing.assert_close(out, ref, rtol=2e-2, atol=2e-1)
+
+
+def test_n1_big_weight_dispatch_unchanged(monkeypatch):
+    """n=1 must keep the stock wvSplitK skinny kernel under the guard
+    (default ON) -- the carve-out keeps n=1 bit-identical to stock."""
+    w = _weight(BIG_SHAPE)
+    x = _x(1, BIG_SHAPE[1])
+    skinny, _ = _skinny_dispatch(monkeypatch, x, w)
+    assert skinny, "n=1 was redirected by the guard; the carve-out is missing"
diff --git a/tests/model_executor/layers/test_pooler_methods.py b/tests/model_executor/layers/test_pooler_methods.py
index b7eed0aefb78..7b053d63e31d 100644
--- a/tests/model_executor/layers/test_pooler_methods.py
+++ b/tests/model_executor/layers/test_pooler_methods.py
@@ -677,6 +677,49 @@ def test_unknown_raises(self):
             get_tok_pooling_method("UNKNOWN")
 
 
+@pytest.mark.parametrize(
+    ("overrides", "expected_calls"),
+    [
+        pytest.param({}, 1, id="matching"),
+        pytest.param({"dimensions": 2}, 3, id="mixed-dimensions"),
+        pytest.param({"use_activation": False}, 2, id="mixed-activation"),
+    ],
+)
+def test_token_pooler_batches_projection_with_mixed_params(
+    overrides, expected_calls, monkeypatch
+):
+    monkeypatch.setenv("VLLM_BATCH_INVARIANT", "0")
+    torch.manual_seed(0)
+    lengths = [3, 4, 2]
+    hidden = torch.randn(sum(lengths), 4, dtype=torch.bfloat16)
+    projector = torch.nn.Sequential(torch.nn.Linear(4, 3), torch.nn.Tanh())
+    activation = PoolerNormalize()
+    head = TokenEmbeddingPoolerHead(torch.float32, projector, activation)
+    params = [PoolingParams(task="token_embed", use_activation=True) for _ in lengths]
+    for name, value in overrides.items():
+        setattr(params[1], name, value)
+    pooler = TokenPooler(TestAllPool._make_all_pool(async_scheduling=True), head)
+    metadata = _make_metadata(lengths, pooling_params=params)
+    expected = [
+        head.forward_chunk(chunk, param)
+        for chunk, param in zip(hidden.split(lengths), params, strict=True)
+    ]
+    projections, activations = [], []
+    with (
+        projector.register_forward_hook(
+            lambda _module, _inputs, output: projections.append(output)
+        ),
+        activation.register_forward_hook(
+            lambda _module, _inputs, output: activations.append(output)
+        ),
+    ):
+        outputs = pooler(hidden, metadata)
+    for actual, reference in zip(outputs, expected, strict=True):
+        torch.testing.assert_close(actual, reference)
+    assert len(projections) == 1
+    assert len(activations) == expected_calls
+
+
 @pytest.mark.parametrize(
     ("chunk_size", "dtype", "head_dtype"),
     [
@@ -727,7 +770,10 @@ def test_late_chunk_pool_reduces_before_nonlinear_projector():
     assert not torch.allclose(actual, wrong_order)
 
 
-def test_late_chunk_pool_isolates_mixed_requests_and_owns_finished_storage():
+@pytest.mark.parametrize("late_chunking", [False, True])
+def test_late_chunk_pool_isolates_mixed_requests_and_owns_finished_storage(
+    late_chunking,
+):
     hidden = torch.arange(36, dtype=torch.float32).reshape(9, 4)
     params = [
         PoolingParams(
@@ -738,6 +784,9 @@ def test_late_chunk_pool_isolates_mixed_requests_and_owns_finished_storage():
             task="token_embed", late_chunking_params=LateChunkingParams(chunk_size=1)
         ),
     ]
+    if not late_chunking:
+        for param in params:
+            param.late_chunking_params = None
     pooler = TokenPooler(
         TestAllPool._make_all_pool(async_scheduling=True), TokenEmbeddingPoolerHead()
     )
@@ -747,6 +796,8 @@ def test_late_chunk_pool_isolates_mixed_requests_and_owns_finished_storage():
         hidden[3:7].clone(),
         hidden[7:].clone(),
     ]
+    if not late_chunking:
+        expected[0] = hidden[:3].clone()
     hidden.fill_(-999)
     for actual, reference in zip(outputs, expected):
         torch.testing.assert_close(actual, reference)
diff --git a/tests/model_executor/model_loader/runai_streamer_loader/test_runai_model_streamer_loader.py b/tests/model_executor/model_loader/runai_streamer_loader/test_runai_model_streamer_loader.py
index a33f11b61aa4..f946b5c71450 100644
--- a/tests/model_executor/model_loader/runai_streamer_loader/test_runai_model_streamer_loader.py
+++ b/tests/model_executor/model_loader/runai_streamer_loader/test_runai_model_streamer_loader.py
@@ -90,7 +90,7 @@ def test_runai_model_loader_download_files(vllm_runner):
 @create_new_process_for_each_test("spawn")
 def test_runai_mtp_real_checkpoint(vllm_runner, tmp_path, monkeypatch):
     """Skip main-only shards without changing a real MTP model's generation."""
-    from vllm.model_executor.models.qwen3_5_mtp import Qwen3_5MTP
+    from vllm.models.qwen3_5 import Qwen3_5MTP
     from vllm.platforms import current_platform
     from vllm.transformers_utils.repo_utils import hf_api
     from vllm.v1.metrics.reader import Counter
diff --git a/tests/model_executor/model_loader/runai_streamer_loader/test_weight_utils.py b/tests/model_executor/model_loader/runai_streamer_loader/test_weight_utils.py
index 0d2522284845..b58d683c5b39 100644
--- a/tests/model_executor/model_loader/runai_streamer_loader/test_weight_utils.py
+++ b/tests/model_executor/model_loader/runai_streamer_loader/test_weight_utils.py
@@ -5,9 +5,11 @@
 import tempfile
 
 import huggingface_hub.constants
+import pytest
 import torch
 from safetensors.torch import save_file
 
+from vllm.model_executor.model_loader import weight_utils
 from vllm.model_executor.model_loader.weight_utils import (
     download_weights_from_hf,
     runai_safetensors_weights_iterator,
@@ -15,9 +17,7 @@
 )
 
 
-def test_runai_safetensors_weights_iterator_clones_reused_buffers(
-    tmp_path, monkeypatch
-):
+def test_runai_safetensors_weights_iterator_yields_owned_tensors(tmp_path, monkeypatch):
     monkeypatch.setenv("RUNAI_STREAMER_MEMORY_LIMIT", "0")
     weights_file = tmp_path / "model.safetensors"
     expected_tensors = {
@@ -26,16 +26,80 @@ def test_runai_safetensors_weights_iterator_clones_reused_buffers(
     }
     save_file(expected_tensors, weights_file)
 
+    clone_count = 0
+    original_clone = torch.Tensor.clone
+
+    def counting_clone(self, *args, **kwargs):
+        nonlocal clone_count
+        clone_count += 1
+        return original_clone(self, *args, **kwargs)
+
+    monkeypatch.setattr(torch.Tensor, "clone", counting_clone)
     actual_tensors = dict(
         runai_safetensors_weights_iterator([str(weights_file)], False)
     )
 
+    assert clone_count == 0
     assert actual_tensors.keys() == expected_tensors.keys()
     assert actual_tensors["first"].data_ptr() != actual_tensors["second"].data_ptr()
     for name, expected_tensor in expected_tensors.items():
         assert torch.equal(actual_tensors[name], expected_tensor)
 
 
+class _FakeStreamer:
+    """Streams two tensors through one shared buffer, like runai<0.17.0."""
+
+    def __init__(self):
+        self.files_to_tensors_metadata = {0: ["first", "second"]}
+        self.stream_kwargs = None
+        self.buffer = torch.empty(2)
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        pass
+
+    def get_tensors(self):
+        for name, value in (("first", 1.0), ("second", 3.0)):
+            self.buffer.fill_(value)
+            yield name, self.buffer
+
+
+class _FakeOwnedStreamer(_FakeStreamer):
+    """Yields a new tensor each step, like runai>=0.17.0 with owned=True."""
+
+    def stream_files(self, paths, device, is_distributed, owned=True):
+        self.stream_kwargs = {"owned": owned}
+
+    def get_tensors(self):
+        for name, tensor in super().get_tensors():
+            yield name, tensor.clone()
+
+
+class _FakeLegacyStreamer(_FakeStreamer):
+    def stream_files(self, paths, device, is_distributed):
+        self.stream_kwargs = {}
+
+
+@pytest.mark.parametrize(
+    ("streamer_cls", "expected_kwargs"),
+    [(_FakeOwnedStreamer, {"owned": True}), (_FakeLegacyStreamer, {})],
+)
+def test_runai_safetensors_weights_iterator_streamer_versions(
+    monkeypatch, streamer_cls, expected_kwargs
+):
+    """Weights stay distinct for streamers with and without `owned` support."""
+    streamer = streamer_cls()
+    monkeypatch.setattr(weight_utils, "SafetensorsStreamer", lambda: streamer)
+
+    tensors = dict(runai_safetensors_weights_iterator(["unused"], False))
+
+    assert streamer.stream_kwargs == expected_kwargs
+    assert torch.equal(tensors["first"], torch.tensor([1.0, 1.0]))
+    assert torch.equal(tensors["second"], torch.tensor([3.0, 3.0]))
+
+
 def test_runai_model_loader():
     with tempfile.TemporaryDirectory() as tmpdir:
         huggingface_hub.constants.HF_HUB_OFFLINE = False
diff --git a/tests/model_executor/model_loader/tensorizer_loader/test_tensorizer.py b/tests/model_executor/model_loader/tensorizer_loader/test_tensorizer.py
index c19765c462c4..193968d7b8e7 100644
--- a/tests/model_executor/model_loader/tensorizer_loader/test_tensorizer.py
+++ b/tests/model_executor/model_loader/tensorizer_loader/test_tensorizer.py
@@ -8,6 +8,7 @@
 import pathlib
 import subprocess
 import sys
+from types import SimpleNamespace
 from typing import Any
 from unittest.mock import Mock
 
@@ -17,6 +18,7 @@
 import vllm.model_executor.model_loader.tensorizer
 from tests.utils import VLLM_PATH, RemoteOpenAIServer
 from vllm import SamplingParams
+from vllm.config.load import LoadConfig
 from vllm.engine.arg_utils import EngineArgs
 from vllm.model_executor.model_loader.tensorizer import (
     TensorizerConfig,
@@ -27,6 +29,7 @@
 )
 from vllm.model_executor.model_loader.tensorizer_loader import (
     BLACKLISTED_TENSORIZER_ARGS,
+    TensorizerLoader,
 )
 from vllm.utils.import_utils import PlaceholderModule
 
@@ -303,6 +306,64 @@ def test_vllm_tensorized_model_has_same_outputs(
         assert outputs == deserialized_outputs
 
 
+@torch.inference_mode()
+def test_deserialized_hf_weights_are_ready_for_forward(tmp_path, monkeypatch):
+    """Raw weights need CPU linear post-processing before the first forward."""
+    from vllm.model_executor.layers import linear
+    from vllm.model_executor.layers import utils as layer_utils
+    from vllm.model_executor.model_loader import base_loader
+    from vllm.platforms.cpu import CpuPlatform
+
+    # Exercise the CPU method on accelerator CI hosts too.
+    for module in (linear, layer_utils, base_loader):
+        monkeypatch.setattr(module, "current_platform", CpuPlatform())
+
+    class TinyModel(torch.nn.Sequential):
+        def __init__(self):
+            super().__init__(torch.nn.Linear(3, 2, bias=False))
+            self[0].quant_method = linear.UnquantizedLinearMethod()
+            self[0]._cpu_skip_gemm_dispatch = True
+
+        def load_weights(self, weights):
+            self.load_state_dict(dict(weights))
+
+        def forward(self, inputs):
+            return self[0].quant_method.apply(self[0], inputs)
+
+    source = torch.nn.Sequential(torch.nn.Linear(3, 2, bias=False))
+    source[0].weight.copy_(torch.tensor([[1.0, 2.0, 3.0], [-2.0, 0.5, 4.0]]))
+    inputs = torch.tensor([[1.0, 2.0, 3.0], [-1.0, 0.0, 2.0]])
+    checkpoint = tmp_path / "raw.tensors"
+    with open_stream(checkpoint, "wb+") as stream:
+        serializer = TensorSerializer(stream)
+        serializer.write_module(source)
+        serializer.close()
+
+    load_config = LoadConfig(
+        load_format="tensorizer",
+        model_loader_extra_config={
+            "tensorizer_config": {"tensorizer_uri": str(checkpoint)}
+        },
+    )
+    model_config = SimpleNamespace(
+        dtype=torch.float32,
+        quantization=None,
+        word_embeddings_untied_by_checkpoint=False,
+    )
+    config = SimpleNamespace(
+        model_config=model_config,
+        load_config=load_config,
+        device_config=SimpleNamespace(device=torch.device("cpu")),
+        parallel_config=SimpleNamespace(tensor_parallel_size=1),
+        quant_config=None,
+    )
+    monkeypatch.setattr(base_loader, "initialize_model", lambda **kwargs: TinyModel())
+
+    loaded = TensorizerLoader(load_config).load_model(config, model_config)
+
+    torch.testing.assert_close(loaded(inputs), source(inputs))
+
+
 def test_load_with_just_model_tensors(just_serialize_model_tensors, model_ref):
     # For backwards compatibility, ensure Tensorizer can be still be loaded
     # for inference by passing the model reference name, not a local/S3 dir,
diff --git a/tests/model_executor/model_loader/test_direct_reload.py b/tests/model_executor/model_loader/test_direct_reload.py
new file mode 100644
index 000000000000..dd63d673f643
--- /dev/null
+++ b/tests/model_executor/model_loader/test_direct_reload.py
@@ -0,0 +1,578 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Direct weight reload (``VLLM_RELOAD_MODE=direct``) on a toy model whose
+post-load processing has the hazards found on real models:
+
+(a) a parent derives a separate tensor from a child's weight (MLA W_UK/W_UV
+    from a quantized kv_b_proj);
+(b) post-load processing changes a weight's layout;
+(c) the model's loader writes a weight without its ``weight_loader`` (MiMo);
+(d) post-load processing consumes a checkpoint tensor, which stays
+    registered as a retired parameter (NVFP4 ``weight_scale_2``);
+(e) post-load processing that must never run on converted bytes (packing);
+(f) post-load processing grows a weight (padding): a reload writes it into
+    the existing storage to keep its address.
+
+Invariant: a reload equals a cold start on the new weights, at the startup
+addresses.
+"""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm.config import VllmConfig
+from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase
+from vllm.model_executor.layers.quantization.online.base import OnlineQuantizationConfig
+from vllm.model_executor.layers.quantization.online.fp8 import (
+    Fp8PerTensorOnlineLinearMethod,
+)
+from vllm.model_executor.model_loader.reload import finish_reload, start_reload
+from vllm.model_executor.model_loader.reload.direct import (
+    begin_direct_reload,
+    finish_direct_reload,
+)
+from vllm.model_executor.model_loader.utils import process_weights_after_loading
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+from vllm.model_executor.utils import replace_parameter, retire_parameter
+
+_MODEL_CONFIG = SimpleNamespace(
+    word_embeddings_untied_by_checkpoint=False, quantization=None, dtype=torch.bfloat16
+)
+_DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+class _Method(QuantizeMethodBase):
+    """Post-load processing given as a function of the layer; counts runs."""
+
+    def __init__(self, fn):
+        self.fn = fn
+        self.runs = 0
+
+    def create_weights(self, layer, *weight_args, **extra_weight_attrs):
+        pass
+
+    def apply(self, layer, *args, **kwargs):
+        raise NotImplementedError
+
+    def process_weights_after_loading(self, layer):
+        self.runs += 1
+        self.fn(layer)
+
+
+def _derive_w_uk(layer):  # (a): a separate tensor derived from the child
+    w_uk = (layer.kv_b.weight.float() * layer.k_scale).t().contiguous()
+    if hasattr(layer, "W_UK"):  # refilled in place: graphs keep its address
+        layer.W_UK.copy_(w_uk)
+    else:
+        layer.W_UK = w_uk
+
+
+def _transpose(layer):  # (b): kernel layout differs from the checkpoint's
+    replace_parameter(layer, "weight", layer.weight.t().contiguous())
+
+
+def _requantize(layer):  # (c): fused-FP8-like rescale
+    scale = layer.weight.float().abs().amax()
+    replace_parameter(layer, "weight", (layer.weight / scale).to(layer.weight.dtype))
+    replace_parameter(layer, "weight_scale", scale.reshape(1))
+
+
+def _consume_scale_2(layer):  # (d): NVFP4 weight_scale_2 folded in, retired
+    ws2 = layer.weight_scale_2
+    replace_parameter(layer, "global_scale", (1.0 / ws2).reshape(1))
+    replace_parameter(layer, "weight", (layer.weight * ws2).to(layer.weight.dtype))
+    retire_parameter(layer, "weight_scale_2")
+
+
+def _pack(layer):  # (e): not idempotent, must never run on converted bytes
+    replace_parameter(layer, "weight", (layer.weight * 2).t().contiguous())
+
+
+def _pad(layer):  # (f): kernel layout larger than the checkpoint's
+    replace_parameter(layer, "weight", torch.nn.functional.pad(layer.weight, (0, 4)))
+
+
+_SHAPES = {
+    "attn.kv_b.weight": ((8, 4), torch.bfloat16),
+    "proj.weight": ((4, 6), torch.bfloat16),
+    "mimo.weight": ((4, 4), torch.bfloat16),
+    "nvfp4.weight": ((4, 4), torch.bfloat16),
+    "nvfp4.weight_scale_2": ((), torch.float32),
+    "engram.weight": ((3, 5), torch.bfloat16),
+    "pad.weight": ((4, 4), torch.bfloat16),
+}
+_PWAL = {
+    "attn.kv_b": None,
+    "proj": _transpose,
+    "mimo": _requantize,
+    "nvfp4": _consume_scale_2,
+    "engram": _pack,
+    "pad": _pad,
+}
+
+
+class _Layer(torch.nn.Module):
+    def __init__(self, fn):
+        super().__init__()
+        self.quant_method = _Method(fn) if fn is not None else None
+
+
+class _ToyModel(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.attn = _Layer(_derive_w_uk)
+        self.attn.register_buffer("k_scale", torch.tensor(0.5))  # not in checkpoint
+        for name, (shape, dtype) in _SHAPES.items():
+            *path, local = name.split(".")
+            parent = self
+            for i, part in enumerate(path):
+                if not hasattr(parent, part):
+                    parent.add_module(part, _Layer(_PWAL[".".join(path[: i + 1])]))
+                parent = getattr(parent, part)
+            param = torch.nn.Parameter(torch.empty(shape, dtype=dtype), False)
+            parent.register_parameter(local, param)
+
+    def load_weights(self, weights):
+        params = dict(self.named_parameters())
+        loaded = set()
+        for name, weight in weights:
+            param = params[name]
+            if name.startswith("mimo."):
+                param.data.copy_(weight)  # (c): bypasses param.weight_loader
+            else:
+                getattr(param, "weight_loader", default_weight_loader)(param, weight)
+            loaded.add(name)
+        return loaded
+
+
+def _checkpoint(scale: float = 1.0) -> dict[str, torch.Tensor]:
+    gen = torch.Generator().manual_seed(0)
+    return {
+        name: ((torch.rand(shape, generator=gen) + 0.5) * scale).to(dtype)
+        for name, (shape, dtype) in _SHAPES.items()
+    }
+
+
+def _cold_start(weights: dict[str, torch.Tensor]) -> _ToyModel:
+    """As ``BaseModelLoader.load_model`` does."""
+    with _DEVICE:
+        model = _ToyModel()
+    model.load_weights(weights.items())
+    process_weights_after_loading(model, _MODEL_CONFIG, _DEVICE)
+    return model
+
+
+def _reload(model: _ToyModel, weights: dict[str, torch.Tensor]) -> None:
+    begin_direct_reload(model)
+    model.load_weights(weights.items())
+    finish_direct_reload(model, _MODEL_CONFIG, _DEVICE)
+
+
+def _state(model: torch.nn.Module) -> dict[str, tuple]:
+    """Every parameter, buffer and tensor attribute: layout, bytes, address."""
+    tensors: dict[str, torch.Tensor] = {}
+    for module_name, module in model.named_modules():
+        prefix = f"{module_name}." if module_name else ""
+        for attr, value in vars(module).items():
+            if isinstance(value, torch.Tensor):
+                tensors[prefix + attr] = value
+    tensors.update(model.named_parameters())
+    tensors.update(model.named_buffers())
+    return {
+        name: (
+            (tuple(t.shape), t.stride(), t.dtype),
+            t.detach().contiguous().reshape(-1).view(torch.uint8).cpu().clone(),
+            t.data_ptr(),
+        )
+        for name, t in tensors.items()
+    }
+
+
+def _assert_same(actual: dict, expected: dict, names=None, addresses=False) -> None:
+    """Layout and bytes (and addresses) equal, for ``names`` or every tensor."""
+    if names is None:
+        assert sorted(actual) == sorted(expected)
+        names = sorted(expected)
+    for name in names:
+        (layout, data, ptr), (exp_layout, exp_data, exp_ptr) = (
+            actual[name],
+            expected[name],
+        )
+        assert layout == exp_layout, name
+        assert torch.equal(data, exp_data), f"{name} differs"
+        assert not addresses or ptr == exp_ptr, f"{name} moved"
+
+
+def _ptrs(state: dict, names=None) -> dict[str, int]:
+    return {n: state[n][2] for n in (state if names is None else names)}
+
+
+def _runs(model: _ToyModel) -> dict[str, int]:
+    return {
+        name: m.quant_method.runs
+        for name, m in model.named_modules()
+        if getattr(m, "quant_method", None) is not None
+    }
+
+
+@pytest.fixture(autouse=True)
+def _direct_mode(monkeypatch):
+    monkeypatch.setenv("VLLM_RELOAD_MODE", "direct")
+
+
+class TestReloadEqualsColdStart:
+    def test_same_checkpoint_is_bit_identical(self):
+        model = _cold_start(_checkpoint())
+        before = _state(model)
+        _reload(model, _checkpoint())
+        _assert_same(_state(model), before, addresses=True)
+        assert set(_runs(model).values()) == {2}
+
+    def test_new_checkpoint_matches_fresh_cold_start(self):
+        model = _cold_start(_checkpoint())
+        before = _state(model)
+        _reload(model, _checkpoint(1.01))
+        after = _state(model)
+        _assert_same(after, _state(_cold_start(_checkpoint(1.01))))
+        assert _ptrs(after) == _ptrs(before)
+        assert model.nvfp4.weight_scale_2.retired  # as at cold start
+
+    @pytest.mark.parametrize(
+        "names",
+        [
+            ["attn.W_UK"],
+            ["proj.weight"],
+            ["mimo.weight", "mimo.weight_scale"],
+            ["nvfp4.weight", "nvfp4.global_scale"],
+            ["engram.weight"],
+            ["pad.weight"],
+        ],
+        ids=[
+            "a-parent-derived",
+            "b-layout-change",
+            "c-loader-bypass",
+            "d-consumed-scale",
+            "e-not-idempotent",
+            "f-grows",
+        ],
+    )
+    def test_new_checkpoint_hazard(self, names):
+        model = _cold_start(_checkpoint())
+        before = _state(model)
+        _reload(model, _checkpoint(1.01))
+        after = _state(model)
+        assert any(not torch.equal(after[n][1], before[n][1]) for n in names)
+        _assert_same(after, _state(_cold_start(_checkpoint(1.01))), names=names)
+        assert _ptrs(after, names) == _ptrs(before, names)
+
+
+class TestUnsent:
+    def test_tensor_without_post_processing_keeps_its_bytes(self):
+        # a frozen table the trainer does not send (Engram, an n-gram
+        # embedding): its module has no post-load processing, so a reload
+        # keeps its cold-start bytes; derived tensors read them as cold start
+        model = _cold_start(_checkpoint())
+        old = _checkpoint()
+        new = _checkpoint(1.01)
+        sent = {n: t for n, t in new.items() if n != "attn.kv_b.weight"}
+        before = _state(model)
+        _reload(model, sent)
+        after = _state(model)
+        _assert_same(after, before, names=["attn.kv_b.weight"], addresses=True)
+        expected = dict(new, **{"attn.kv_b.weight": old["attn.kv_b.weight"]})
+        _assert_same(after, _state(_cold_start(expected)))
+
+
+class TestRefusals:
+    def test_refused_begin_leaves_the_model_untouched(self):
+        # proj is viewed back first, in module order; mimo's storage is short
+        model = _cold_start(_checkpoint())
+        owner, (_, _, dtype, offset), attrs = model.mimo.weight._ckpt
+        model.mimo.weight._ckpt = (owner, ((64, 64), (64, 1), dtype, offset), attrs)
+        before = _state(model)
+        with pytest.raises(RuntimeError, match="checkpoint layout needs"):
+            begin_direct_reload(model)
+        _assert_same(_state(model), before, addresses=True)
+
+    def test_online_quantization_refused(self, monkeypatch):
+        # finish clears the guard that keeps online quantization from
+        # quantizing its converted bytes again
+        online = object.__new__(Fp8PerTensorOnlineLinearMethod)
+        init = _ToyModel.__init__
+
+        def with_online_layer(self):
+            init(self)
+            self.proj.quant_method = online
+
+        monkeypatch.setattr(_ToyModel, "__init__", with_online_layer)
+        with pytest.raises(RuntimeError, match="online quantization"):
+            _cold_start(_checkpoint())
+
+    def test_online_shared_experts_refused(self, monkeypatch):
+        # post-load processing deletes the staging parameters a reload loads
+        init = _ToyModel.__init__
+
+        def with_online_shared_experts(self):
+            init(self)
+            self.proj.shared_expert_online_loader = object()
+
+        monkeypatch.setattr(_ToyModel, "__init__", with_online_shared_experts)
+        with pytest.raises(RuntimeError, match="online quantization"):
+            _cold_start(_checkpoint())
+
+    @pytest.mark.parametrize("method_cls", ["HummingLinearMethod", "HummingMoEMethod"])
+    def test_humming_refused(self, monkeypatch, method_cls):
+        # its post-load processing runs only once per quant method
+        from vllm.model_executor.layers.quantization import humming as humming_mod
+
+        humming = object.__new__(getattr(humming_mod, method_cls))
+        init = _ToyModel.__init__
+
+        def with_humming_layer(self):
+            init(self)
+            self.proj.quant_method = humming
+
+        monkeypatch.setattr(_ToyModel, "__init__", with_humming_layer)
+        with pytest.raises(RuntimeError, match="humming"):
+            _cold_start(_checkpoint())
+
+    def test_replaced_param_refused(self):
+        # a new Parameter under a checkpoint name has no checkpoint layout:
+        # the loader would write checkpoint bytes into its runtime layout
+        model = _cold_start(_checkpoint())
+        model.proj.weight = torch.nn.Parameter(model.proj.weight.data, False)
+        before = _state(model)
+        with pytest.raises(RuntimeError, match=r"replaced .*proj\.weight"):
+            begin_direct_reload(model)
+        _assert_same(_state(model), before, addresses=True)
+        assert "_direct_reload_begun" not in model.__dict__
+
+
+def test_retry_after_aborted_update_matches_cold_start():
+    # Direct reload has no rollback: a failed update leaves the model partly
+    # updated and pending (nothing refuses to run it), and the next full
+    # update repairs it.
+    model = _cold_start(_checkpoint())
+    assert "_direct_reload_begun" not in model.__dict__
+    begin_direct_reload(model)
+    model.load_weights(list(_checkpoint(1.01).items())[:3])  # update aborts
+    assert "_direct_reload_begun" in model.__dict__
+    _reload(model, _checkpoint(1.01))
+    assert "_direct_reload_begun" not in model.__dict__
+    _assert_same(_state(model), _state(_cold_start(_checkpoint(1.01))))
+
+
+def test_start_and_finish_reload_run_direct_reload():
+    model = _cold_start(_checkpoint())
+    before = _state(model)
+    start_reload(model)
+    model.load_weights(_checkpoint().items())
+    finish_reload(model, _MODEL_CONFIG, _DEVICE)
+    _assert_same(_state(model), before, addresses=True)
+
+
+def test_finish_keeps_guards_of_modules_without_quant_method():
+    # e.g. MoERunnerInterface sets the guard in __init__ for other reasons
+    model = _cold_start(_checkpoint())
+    model.attn.kv_b._already_called_process_weights_after_loading = True
+    model.proj._already_called_process_weights_after_loading = True
+    _reload(model, _checkpoint())
+    assert model.attn.kv_b._already_called_process_weights_after_loading
+    assert "_already_called_process_weights_after_loading" not in vars(model.proj)
+
+
+def test_finish_retry_after_failed_direct_finish_raises(monkeypatch):
+    # some modules may already be processed: only a new full update repairs it
+    model = _cold_start(_checkpoint())
+    start_reload(model)
+    model.load_weights(_checkpoint(1.01).items())
+
+    def fail(layer):
+        raise ValueError("boom")
+
+    monkeypatch.setattr(model.nvfp4.quant_method, "fn", fail)
+    with pytest.raises(ValueError, match="boom"):
+        finish_reload(model, _MODEL_CONFIG, _DEVICE)
+    with pytest.raises(RuntimeError, match="finish without a begin"):
+        finish_reload(model, _MODEL_CONFIG, _DEVICE)
+
+
+def test_begin_restores_loader_metadata(monkeypatch):
+    # AWQ retags the kept parameter's packed/output dims; loaders shard with them
+    init = _ToyModel.__init__
+
+    def retag(layer):
+        new = layer.weight.t().contiguous()
+        new._output_dim, new._packed_dim = 1, 1
+        replace_parameter(layer, "weight", new)
+
+    def with_loader_metadata(self):
+        init(self)
+        self.proj.weight._output_dim = 0
+        self.proj.quant_method = _Method(retag)
+
+    monkeypatch.setattr(_ToyModel, "__init__", with_loader_metadata)
+    model = _cold_start(_checkpoint())
+    weight = model.proj.weight
+    assert (weight._output_dim, weight._packed_dim) == (1, 1)
+    begin_direct_reload(model)
+    assert weight._output_dim == 0 and not hasattr(weight, "_packed_dim")
+    model.load_weights(_checkpoint(1.01).items())
+    finish_direct_reload(model, _MODEL_CONFIG, _DEVICE)
+    assert (weight._output_dim, weight._packed_dim) == (1, 1)
+    _assert_same(_state(model), _state(_cold_start(_checkpoint(1.01))))
+
+
+def test_shared_tensor_with_changed_layout_refused():
+    # a draft sharing the target's lm_head after its cold start: the draft's
+    # finish would reprocess it although a draft update does not resend it
+    target = _cold_start(_checkpoint())
+    draft = _cold_start(_checkpoint())
+    draft.proj = target.proj  # its post-load processing transposes it
+    before = _state(target)
+    with pytest.raises(RuntimeError, match="shared from another model"):
+        begin_direct_reload(draft)
+    assert "_direct_reload_begun" not in draft.__dict__
+    _assert_same(_state(target), before, addresses=True)
+
+
+@pytest.mark.parametrize(
+    "fn",
+    [
+        # ROCm FP8 padding: a view with gaps, padded rows apart
+        lambda w: torch.nn.functional.pad(w, (0, 4))[..., :-4],
+        # a released weight: its checkpoint bytes stay for the next reload
+        lambda w: w.new_empty(0),
+    ],
+    ids=["padded-view", "emptied"],
+)
+def test_reload_keeps_address_of_padded_or_emptied_weight(monkeypatch, fn):
+    init = _ToyModel.__init__
+
+    def with_fn(self):
+        init(self)
+        self.proj.quant_method = _Method(
+            lambda layer: replace_parameter(layer, "weight", fn(layer.weight))
+        )
+
+    monkeypatch.setattr(_ToyModel, "__init__", with_fn)
+    model = _cold_start(_checkpoint())
+    ptrs = _ptrs(_state(model))
+    _reload(model, _checkpoint(1.01))
+    _assert_same(_state(model), _state(_cold_start(_checkpoint(1.01))))
+    assert _ptrs(_state(model)) == ptrs
+
+
+def test_replace_parameter_keeps_the_target_checkpoint_layout():
+    # post-load processing may pass another checkpoint parameter as new_tensor
+    model = _cold_start(_checkpoint())
+    ckpt = model.proj.weight._ckpt
+    replace_parameter(model.proj, "weight", model.mimo.weight)
+    assert model.proj.weight._ckpt == ckpt
+
+
+def _config(**overrides) -> SimpleNamespace:
+    offload = SimpleNamespace(
+        offload_backend="auto",
+        prefetch=SimpleNamespace(offload_group_size=0),
+        uva=SimpleNamespace(cpu_offload_gb=0),
+    )
+    fields = dict(
+        offload_config=offload,
+        model_config=SimpleNamespace(quantization=None),
+        quant_config=None,
+        lora_config=None,
+        speculative_config=None,
+        weight_transfer_config=None,
+    )
+    return SimpleNamespace(**(fields | overrides))
+
+
+_ONLINE = object.__new__(OnlineQuantizationConfig)
+
+
+@pytest.mark.parametrize(
+    "mode,overrides,match",
+    [
+        ("direct", {}, None),
+        ("layerwise", {"lora_config": object()}, None),
+        (
+            "direct",
+            {
+                "offload_config": SimpleNamespace(
+                    offload_backend="uva",
+                    prefetch=SimpleNamespace(offload_group_size=0),
+                    uva=SimpleNamespace(cpu_offload_gb=0),
+                )
+            },
+            "weight offloading",
+        ),
+        (
+            "direct",
+            {
+                "offload_config": SimpleNamespace(
+                    offload_backend="auto",
+                    prefetch=SimpleNamespace(offload_group_size=2),
+                    uva=SimpleNamespace(cpu_offload_gb=0),
+                )
+            },
+            "weight offloading",
+        ),
+        (
+            "direct",
+            {
+                "offload_config": SimpleNamespace(
+                    offload_backend="auto",
+                    prefetch=SimpleNamespace(offload_group_size=0),
+                    uva=SimpleNamespace(cpu_offload_gb=1),
+                )
+            },
+            "weight offloading",
+        ),
+        (
+            "direct",
+            {"model_config": SimpleNamespace(quantization="torchao")},
+            "torchao",
+        ),
+        ("direct", {"quant_config": _ONLINE}, "online quantization"),
+        (
+            "direct",
+            {"quant_config": SimpleNamespace(online_quantization_config=_ONLINE)},
+            "online quantization",
+        ),
+        ("direct", {"lora_config": object()}, "LoRA"),
+        (
+            "direct",
+            {"speculative_config": SimpleNamespace(draft_lm_head_quantization="fp8")},
+            "draft_lm_head_quantization",
+        ),
+        (
+            "direct",
+            {"weight_transfer_config": SimpleNamespace(backend="sharded_rdt")},
+            "sharded_rdt",
+        ),
+    ],
+    ids=[
+        "defaults",
+        "layerwise-lora",
+        "offload-backend",
+        "offload-prefetch",
+        "offload-uva",
+        "torchao",
+        "online-config",
+        "online-composed",
+        "lora",
+        "draft-lm-head-quant",
+        "sharded-rdt",
+    ],
+)
+def test_verify_direct_reload_compat(monkeypatch, mode, overrides, match):
+    monkeypatch.setenv("VLLM_RELOAD_MODE", mode)
+    config = _config(**overrides)
+    if match is None:
+        VllmConfig._verify_direct_reload_compat(config)
+    else:
+        with pytest.raises(ValueError, match=match):
+            VllmConfig._verify_direct_reload_compat(config)
diff --git a/tests/model_executor/model_loader/test_expert_substitution.py b/tests/model_executor/model_loader/test_expert_substitution.py
index 926720c68629..223df6af7052 100644
--- a/tests/model_executor/model_loader/test_expert_substitution.py
+++ b/tests/model_executor/model_loader/test_expert_substitution.py
@@ -202,7 +202,7 @@ def test_layerwise_reload_is_rejected_and_keeps_weights():
 @pytest.mark.parametrize("missing_value", [False, True])
 def test_tensorizer_finalizes_local_substitution_values(monkeypatch, missing_value):
     """Raw Tensorizer loading must reject an unloaded local constant row."""
-    from vllm.model_executor.model_loader import tensorizer_loader
+    from vllm.model_executor.model_loader import base_loader, tensorizer_loader
 
     model = _cpu_model(local_layers=(0,))
     weights = [(_approx_value(0, 1), torch.zeros(2))]
@@ -210,19 +210,32 @@ def test_tensorizer_finalizes_local_substitution_values(monkeypatch, missing_val
         weights.append((_approx_value(0, 2), torch.ones(2)))
     finalize = Mock(wraps=model.process_weights_after_loading)
     monkeypatch.setattr(model, "process_weights_after_loading", finalize)
-    monkeypatch.setattr(tensorizer_loader, "initialize_model", lambda **kwargs: model)
-    loader = object.__new__(tensorizer_loader.TensorizerLoader)
+    monkeypatch.setattr(base_loader, "initialize_model", lambda **kwargs: model)
+    monkeypatch.setattr(tensorizer_loader, "is_vllm_tensorized", lambda _: False)
+    load_config = LoadConfig(
+        load_format="tensorizer",
+        model_loader_extra_config={"tensorizer_config": {"tensorizer_uri": "raw"}},
+    )
+    loader = tensorizer_loader.TensorizerLoader(load_config)
     monkeypatch.setattr(loader, "_get_weights_iterator", lambda: iter(weights))
+    model_config = SimpleNamespace(
+        dtype=torch.float32,
+        quantization=None,
+        word_embeddings_untied_by_checkpoint=False,
+    )
     config = SimpleNamespace(
-        model_config=SimpleNamespace(dtype=torch.float32),
-        device_config=SimpleNamespace(device="cpu"),
+        model_config=model_config,
+        load_config=load_config,
+        device_config=SimpleNamespace(device=torch.device("cpu")),
+        parallel_config=SimpleNamespace(tensor_parallel_size=1),
+        quant_config=None,
     )
 
     if missing_value:
         with pytest.raises(ValueError, match=r"missing approx_value .* \[2\]"):
-            loader._load_model_serialized_cpu(config)
+            loader.load_model(config, model_config)
     else:
-        assert loader._load_model_serialized_cpu(config) is model
+        assert loader.load_model(config, model_config) is model
         torch.testing.assert_close(
             model.model.layers[0].mlp.experts.values,
             torch.tensor([[0.0, 0.0], [1.0, 1.0]]),
diff --git a/tests/model_executor/model_loader/test_reload.py b/tests/model_executor/model_loader/test_reload.py
index eac5245bd988..645c61d70982 100644
--- a/tests/model_executor/model_loader/test_reload.py
+++ b/tests/model_executor/model_loader/test_reload.py
@@ -17,12 +17,19 @@
 
 import vllm.model_executor.model_loader.reload.layerwise as reload_layerwise
 import vllm.model_executor.model_loader.reload.meta as reload_meta
-from vllm.config import ModelConfig
+from vllm.config import CacheConfig, ModelConfig
+from vllm.distributed.weight_transfer.sparse_nccl_engine import (
+    SparseNCCLWeightTransferEngine,
+)
 from vllm.lora.request import LoRARequest
 from vllm.model_executor.layers.attention import MMEncoderAttention
 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
 from vllm.model_executor.layers.linear import QKVParallelLinear
 from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase
+from vllm.model_executor.model_loader.checkpoint_weight_patch import (
+    CheckpointWeightPatch,
+    load_checkpoint_weight_patches,
+)
 from vllm.model_executor.model_loader.reload.layerwise import (
     finalize_layerwise_reload,
     initialize_layerwise_reload,
@@ -42,6 +49,7 @@
 from vllm.model_executor.model_loader.weight_utils import (
     composed_weight_loader,
     default_weight_loader,
+    finalize_model_weights,
     load_attention_sink,
     safetensors_weights_iterator,
 )
@@ -49,6 +57,23 @@
 from vllm.transformers_utils.repo_utils import hf_api
 
 
+def test_preserve_placement_remapping_until_checkpoint_iteration_finishes():
+    """Deferred reload keeps loaders installed through trailing checkpoint entries."""
+    layer = torch.nn.Linear(2, 2, bias=False)
+    record_metadata_for_reloading(layer)
+    initialize_layerwise_reload(layer, defer_processing=True)
+    checkpoint = torch.arange(4, dtype=torch.float32).reshape(2, 2)
+    loader = layer.weight.weight_loader
+    loader(layer.weight, checkpoint)
+    assert layer.weight.is_meta
+    assert layer.weight.weight_loader is loader
+    loader(layer.weight, checkpoint + 1)
+    assert layer.weight.is_meta
+    finalize_layerwise_reload(layer, model_config=None)
+    assert torch.equal(layer.weight, checkpoint + 1)
+    assert not reload_layerwise.get_layerwise_info(layer).defer_processing
+
+
 @pytest.mark.parametrize("tp_rank", [0, 1])
 @pytest.mark.parametrize(
     "module_name,class_name,sink_name",
@@ -76,6 +101,7 @@ def test_attention_sink_model_loading_shards_once(
     model = cls.__new__(cls)
     torch.nn.Module.__init__(model)
     model.hf_to_vllm_mapper = WeightsMapper()
+    model._context_kv_proj = None  # set by DFlashQwen3Model.__init__
     model.model = torch.nn.Module()
     model.model.mtp = torch.nn.Module()
     sink = torch.nn.Parameter(torch.zeros(4), requires_grad=False)
@@ -328,7 +354,7 @@ def __init__(self):
         self.weight.weight_loader = default_weight_loader
         self.post_load_called = False
 
-    def process_weights_after_loading(self, act_dtype: torch.dtype) -> None:
+    def process_weights_after_loading(self) -> None:
         self.post_load_called = True
 
 
@@ -348,7 +374,7 @@ def get_attn_backend(self):
     def get_kv_cache_spec(self, vllm_config):
         return None
 
-    def process_weights_after_loading(self, act_dtype: torch.dtype) -> None:
+    def process_weights_after_loading(self) -> None:
         self.post_load_called = True
 
 
@@ -720,6 +746,7 @@ def test_padded_moe_reload_releases_each_layer(
         has_bias=has_bias,
         tp_rank=tp_rank,
         tp_shard_with_padding=False,
+        num_experts=experts,
         moe_parallel_config=SimpleNamespace(tp_size=2),
     )
     model = torch.nn.ModuleList()
@@ -735,9 +762,11 @@ def test_padded_moe_reload_releases_each_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.expert_map_manager = SimpleNamespace(
+            map_global_to_local=lambda i: i, num_fused_shared_experts=0
+        )
         layer._loaded_expert_biases = set()
-        layer._fused_shared_expert_quantizer = None
+        layer.shared_expert_online_loader = None
         method.create_weights(
             layer,
             experts,
@@ -1355,6 +1384,8 @@ def test_reload_weights(
         enable_prefix_caching=False,
         max_model_len=16,
         max_num_seqs=1,
+        # Reloading without a pause that empties the KV cache needs headroom.
+        gpu_memory_utilization=CacheConfig.DEFAULT_GPU_MEMORY_UTILIZATION,
     ) as llm:
         # Reloading must update tensors in place: captured graphs keep pointers.
         ptrs = llm.collective_rpc(_device_tensor_ptrs)
@@ -1522,3 +1553,168 @@ def test_online_quantize_reload(
         mul_perp = llm.generate_prompt_perplexity(["3 4 = 12"], mask=["3 4 ="])[0]
         add_perp = llm.generate_prompt_perplexity(["3 4 = 7"], mask=["3 4 ="])[0]
         assert add_perp < mul_perp
+
+
+class _TwoLayerModel(torch.nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.a = torch.nn.Linear(2, 2, bias=False)
+        self.b = torch.nn.Linear(2, 2, bias=False)
+        self.hook_saw: list[tuple[float, float]] = []
+
+    def load_weights(self, weights) -> set[str]:
+        params = dict(self.named_parameters())
+        loaded = set()
+        for name, tensor in weights:
+            param = params[name]
+            getattr(param, "weight_loader", default_weight_loader)(param, tensor)
+            loaded.add(name)
+        return loaded
+
+    def process_weights_after_loading(self) -> None:
+        a, b = self.a.weight.mean().item(), self.b.weight.mean().item()
+        self.hook_saw.append((a, b))
+
+
+def _full(value: float) -> torch.Tensor:
+    return torch.full((2, 2), value)
+
+
+class _LoggedHookModule(torch.nn.Module):
+    def __init__(self, name: str, log: list[str], *children: torch.nn.Module):
+        super().__init__()
+        self.name, self.log = name, log
+        self.kids = torch.nn.ModuleList(children)
+
+    def process_weights_after_loading(self) -> None:
+        self.log.append(self.name)
+
+
+class _QuantMethodModule(torch.nn.Module, QuantizeMethodBase):
+    def create_weights(self, *args, **kwargs) -> None:
+        pass
+
+    def apply(self, *args, **kwargs) -> torch.Tensor:
+        raise NotImplementedError
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        raise AssertionError("quant methods run in the per-layer pass")
+
+
+def test_finalize_runs_every_module_hook_after_its_submodules():
+    log: list[str] = []
+    shared = _LoggedHookModule("shared", log)
+    model = _LoggedHookModule(
+        "root",
+        log,
+        _LoggedHookModule("a", log, _LoggedHookModule("a.x", log), shared),
+        _LoggedHookModule("b", log, shared),
+    )
+    model.attn = _ReloadableAttentionLayer()
+    model.quantized = torch.nn.Module()
+    model.quantized.quant_method = _QuantMethodModule()
+    # owned by the quant method, which processes it in the per-layer pass
+    model.quantized.transform = _LoggedHookModule("transform", log)
+
+    finalize_model_weights(model)
+
+    assert sorted(log) == ["a", "a.x", "b", "root", "shared"]
+    for child, parent in [("a.x", "a"), ("shared", "a"), ("shared", "b")]:
+        assert log.index(child) < log.index(parent)
+    assert log[-1] == "root"
+    assert not model.attn.post_load_called
+
+
+class _TransformReadingAttention(_ReloadableAttentionLayer):
+    def __init__(self, read) -> None:
+        super().__init__()
+        self.read, self.seen = read, None
+
+    def process_weights_after_loading(self) -> None:
+        self.seen = self.read()
+
+
+def test_attention_post_load_reads_normalized_transform(dist_init):
+    """MLA absorbs kv_b_proj, transforms included, in the attention pass, so a
+    quantized layer's Hadamard transform must be normalized before it."""
+    # transform.linear is imported by compressed_tensors and imports it back
+    import vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors  # noqa: F401, E501
+    from vllm.model_executor.layers.quantization.compressed_tensors.transform.linear import (  # noqa: E501
+        CompressedTensorsLinearTransformMethod,
+    )
+    from vllm.model_executor.layers.quantization.compressed_tensors.transform.module import (  # noqa: E501
+        HadamardTransform,
+    )
+    from vllm.model_executor.model_loader.utils import process_weights_after_loading
+    from vllm.model_executor.parameter import SharedWeightParameter
+
+    transform = HadamardTransform.__new__(HadamardTransform)
+    torch.nn.Module.__init__(transform)
+    transform.weight = SharedWeightParameter(weight_loader=default_weight_loader)
+    transform.weight.add_partition(0, object(), (4, 4))
+    transform.weight.partitions[0].data.fill_(1.0)
+
+    method = CompressedTensorsLinearTransformMethod.__new__(
+        CompressedTensorsLinearTransformMethod
+    )
+    method.quant_method = SimpleNamespace(process_weights_after_loading=lambda _: None)
+    linear = torch.nn.Module()
+    linear.quant_method = method
+    linear.transform = transform
+
+    model = torch.nn.Module()
+    model.kv_b_proj = linear
+    model.attn = _TransformReadingAttention(
+        lambda: transform.weight.partitions[0].data[0, 0].item()
+    )
+    config = SimpleNamespace(
+        dtype=torch.float32,
+        quantization=None,
+        word_embeddings_untied_by_checkpoint=False,
+    )
+
+    process_weights_after_loading(model, config, torch.device("cpu"))
+
+    # normalized by 1/sqrt(4) before the attention pass, and only once
+    assert model.attn.seen == 0.5
+    assert transform.weight.partitions[0].data[0, 0].item() == 0.5
+
+
+def _loaded_model() -> _TwoLayerModel:
+    model = _TwoLayerModel()
+    model.load_weights([("a.weight", _full(1.0)), ("b.weight", _full(2.0))])
+    return model
+
+
+class TestModelHookOnSplitLoad:
+    def test_layerwise_reload(self):
+        model = _loaded_model()
+        record_metadata_for_reloading(model)
+        config = SimpleNamespace(dtype=torch.float32)
+
+        initialize_layerwise_reload(model)
+        model.load_weights([("a.weight", _full(3.0))])
+        model.load_weights([("b.weight", _full(4.0))])
+        assert model.hook_saw == []
+        finalize_layerwise_reload(model, config)
+        assert model.hook_saw == [(3.0, 4.0)]
+
+        # Partial update: the untouched layer is restored before the hook runs.
+        initialize_layerwise_reload(model)
+        model.load_weights([("a.weight", _full(5.0))])
+        finalize_layerwise_reload(model, config)
+        assert model.hook_saw == [(3.0, 4.0), (5.0, 4.0)]
+
+    def test_sparse_update(self):
+        model = _loaded_model()
+        engine = object.__new__(SparseNCCLWeightTransferEngine)
+        engine.model = model
+
+        engine.start_weight_update()
+        for name, value in (("a.weight", 5.0), ("b.weight", 6.0)):
+            values, indices = torch.full((4,), value), torch.arange(4)
+            patch = CheckpointWeightPatch(name, (2, 2), torch.float32, values, indices)
+            load_checkpoint_weight_patches(model, [patch])
+        assert model.hook_saw == []
+        engine.finish_weight_update()
+        assert model.hook_saw == [(5.0, 6.0)]
diff --git a/tests/model_executor/model_loader/test_sharded_state_loader.py b/tests/model_executor/model_loader/test_sharded_state_loader.py
index 6ec5f1ce43ab..b20f69c0bd05 100644
--- a/tests/model_executor/model_loader/test_sharded_state_loader.py
+++ b/tests/model_executor/model_loader/test_sharded_state_loader.py
@@ -7,11 +7,14 @@
 import shutil
 from tempfile import TemporaryDirectory
 from types import SimpleNamespace
+from unittest.mock import Mock
 
 import pytest
 import torch
+from safetensors.torch import save_file
 
 from vllm import LLM, SamplingParams
+from vllm.config import LoadConfig
 from vllm.model_executor.model_loader import ShardedStateLoader
 from vllm.transformers_utils.repo_utils import hf_api
 
@@ -77,6 +80,44 @@ def _mock_pp_group(
     )
 
 
+@pytest.mark.parametrize(
+    "model,model_weights,repo_id",
+    [("org/model", None, "org/model"), ("local/dir", "org/weights", "org/weights")],
+)
+def test_load_weights_resolves_hf_repo_id(
+    monkeypatch, tmp_path, model, model_weights, repo_id
+):
+    """HF repo IDs are downloaded with the requested revision before loading."""
+    _mock_pp_group(monkeypatch, 1)
+    weights = torch.arange(4, dtype=torch.float32).reshape(2, 2)
+    save_file({"weight": weights}, tmp_path / "model-rank-0-part-0.safetensors")
+    download = Mock(return_value=str(tmp_path))
+    monkeypatch.setattr(
+        "vllm.model_executor.model_loader.sharded_state_loader.download_weights_from_hf",
+        download,
+    )
+    load_config = LoadConfig(
+        load_format="sharded_state",
+        download_dir=str(tmp_path / "cache"),
+        ignore_patterns=["unused/*"],
+    )
+    model_config = SimpleNamespace(
+        model=model, model_weights=model_weights, revision="checkpoint-revision"
+    )
+    linear = torch.nn.Linear(2, 2, bias=False)
+
+    ShardedStateLoader(load_config).load_weights(linear, model_config)
+
+    torch.testing.assert_close(linear.weight, weights)
+    download.assert_called_once_with(
+        repo_id,
+        load_config.download_dir,
+        ["*.safetensors"],
+        "checkpoint-revision",
+        ignore_patterns=load_config.ignore_patterns,
+    )
+
+
 @pytest.mark.parametrize(
     "world_size,expected",
     [
diff --git a/tests/model_executor/model_loader/test_weight_cache.py b/tests/model_executor/model_loader/test_weight_cache.py
index e9ff7dc17c4e..98501c7b538e 100644
--- a/tests/model_executor/model_loader/test_weight_cache.py
+++ b/tests/model_executor/model_loader/test_weight_cache.py
@@ -9,6 +9,7 @@
 
 import os
 import shutil
+import signal
 import subprocess
 import sys
 import tempfile
@@ -19,6 +20,7 @@
 from dataclasses import dataclass, field, replace
 from typing import Any
 
+import psutil
 import pytest
 import regex as re
 
@@ -133,6 +135,53 @@ def _stop(self) -> None:
         shutil.rmtree(self.socket_dir, ignore_errors=True)
 
 
+@pytest.mark.distributed(num_gpus=2)
+@pytest.mark.parametrize(
+    "after_ready,stop_sibling",
+    [(False, False), (True, False), (True, True)],
+    ids=["startup", "serving", "unresponsive-sibling"],
+)
+def test_preload_rank_failure_stops_launcher_and_siblings(after_ready, stop_sibling):
+    """A failed later rank must not leave preload waiting on an earlier rank."""
+    children: list[psutil.Process] = []
+    with WeightCacheDaemon(
+        "facebook/opt-125m",
+        tp_size=2,
+        extra_args=["--disable-flashinfer-autotune"],
+    ) as daemon:
+        try:
+            assert daemon._proc is not None
+            parent = psutil.Process(daemon._proc.pid)
+            deadline = time.monotonic() + 120
+            while len(children) != 2:
+                assert daemon._proc.poll() is None, "preload exited before spawning"
+                children = []
+                for child in parent.children():
+                    try:
+                        if "spawn_main" in " ".join(child.cmdline()):
+                            children.append(child)
+                    except psutil.NoSuchProcess:
+                        # Short-lived children can exit before cmdline is read.
+                        continue
+                assert time.monotonic() < deadline, "daemon ranks did not spawn"
+                time.sleep(0.1)
+            if after_ready:
+                daemon.wait_until_ready()
+            if stop_sibling:
+                # A stopped process cannot handle SIGTERM; cleanup must escalate.
+                children[0].send_signal(signal.SIGSTOP)
+            children[-1].send_signal(signal.SIGKILL)
+            assert daemon._proc.wait(timeout=30) == 128 + signal.SIGKILL
+            _, alive = psutil.wait_procs(children, timeout=10)
+            assert not alive, "preload left sibling ranks running"
+        finally:
+            # Also clean up on the unfixed launcher, which waits indefinitely.
+            for child in children:
+                if child.is_running():
+                    child.kill()
+            psutil.wait_procs(children, timeout=10)
+
+
 @dataclass
 class ModelCase:
     model: str
diff --git a/tests/model_executor/test_attention_warmup.py b/tests/model_executor/test_attention_warmup.py
index 861276280e24..8ed9910ee674 100644
--- a/tests/model_executor/test_attention_warmup.py
+++ b/tests/model_executor/test_attention_warmup.py
@@ -1,11 +1,18 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import threading
+from concurrent.futures import ThreadPoolExecutor
+from dataclasses import dataclass
 from types import SimpleNamespace
 
 import pytest
+import torch
 
-from vllm.model_executor.warmup import attention_warmup
+from vllm.distributed import parallel_state
+from vllm.model_executor.warmup import attention_warmup, kernel_warmup
 from vllm.model_executor.warmup.attention_warmup import mixed_batch_attention_warmup
+from vllm.v1.kv_cache_interface import FullAttentionSpec
+from vllm.v1.worker.utils import AttentionGroup
 
 
 def _worker(*, use_v2: bool, max_num_batched_tokens: int, max_model_len: int = 4096):
@@ -69,3 +76,106 @@ def fake_mixed_warmup(runner, execute_model, sample_tokens, num_tokens):
     ]
     # No dummy-run fallback when the mixed step is skipped.
     assert dummy_calls == []
+
+
+class _PPGroup:
+    def __init__(self, world_size: int):
+        self.world_size = world_size
+        self._barrier = threading.Barrier(world_size, timeout=5)
+        self._lock = threading.Lock()
+        self._tensors: list[torch.Tensor] = []
+
+    def cpu_all_reduce_(self, tensor: torch.Tensor) -> torch.Tensor:
+        with self._lock:
+            self._tensors.append(tensor)
+
+        if self._barrier.wait() == 0:
+            reduced = torch.stack(self._tensors).sum(dim=0)
+            for other in self._tensors:
+                other.copy_(reduced)
+        self._barrier.wait()
+        return tensor
+
+
+@dataclass
+class _ModelRunner:
+    attn_groups: list[list[AttentionGroup]]
+    is_pooling_model: bool = False
+
+
+@dataclass
+class _Backend:
+    name: str
+
+    def get_name(self) -> str:
+        return self.name
+
+
+@dataclass
+class _Worker:
+    model_runner: _ModelRunner
+    use_v2_model_runner: bool = False
+
+
+def _workers(
+    backend_names_by_rank: tuple[tuple[str, ...], ...], *, use_v2: bool
+) -> tuple[_Worker, ...]:
+    workers = []
+    for backend_names in backend_names_by_rank:
+        attention_groups = []
+        for backend_name in backend_names:
+            attention_groups.append(
+                AttentionGroup(
+                    backend=_Backend(backend_name),
+                    layer_names=["layer"],
+                    kv_cache_spec=FullAttentionSpec(
+                        block_size=16,
+                        num_kv_heads=1,
+                        head_size=128,
+                        dtype=torch.float16,
+                    ),
+                    kv_cache_group_id=0,
+                )
+            )
+        workers.append(
+            _Worker(_ModelRunner([attention_groups]), use_v2_model_runner=use_v2)
+        )
+    return tuple(workers)
+
+
+@pytest.mark.parametrize(
+    ("backend_names_by_rank", "use_v2", "expected_warmup"),
+    [
+        ((("FLASHINFER",),), False, True),
+        ((("FLASHINFER",), ("FLASHINFER",)), False, True),
+        ((("FLASHINFER",), ("FLASHINFER",)), True, True),
+        ((("FLASHINFER", "OTHER_ATTENTION"),), False, False),
+        ((("FLASHINFER", "OTHER_ATTENTION"),), True, True),
+        ((("FLASHINFER", "OTHER_ATTENTION"), ("FLASHINFER",)), False, False),
+        ((("FLASHINFER", "OTHER_ATTENTION"), ("OTHER_ATTENTION",)), True, True),
+    ],
+)
+@pytest.mark.cpu_test
+def test_flashinfer_attention_warmup_gate(
+    monkeypatch, backend_names_by_rank, use_v2, expected_warmup
+):
+    workers = _workers(backend_names_by_rank, use_v2=use_v2)
+    pp_group = _PPGroup(len(workers))
+    monkeypatch.setattr(parallel_state, "get_pp_group", lambda: pp_group)
+    warmup_workers = []
+    monkeypatch.setattr(
+        kernel_warmup,
+        "mixed_batch_attention_warmup",
+        lambda worker: warmup_workers.append(worker),
+    )
+
+    with ThreadPoolExecutor(max_workers=pp_group.world_size) as executor:
+        futures = [
+            executor.submit(kernel_warmup._warmup_flashinfer_attention, worker)
+            for worker in workers
+        ]
+        for future in futures:
+            future.result()
+
+    expected_warmup_calls = pp_group.world_size if expected_warmup else 0
+    assert len(warmup_workers) == expected_warmup_calls
diff --git a/tests/model_executor/test_keye_mrope.py b/tests/model_executor/test_keye_mrope.py
deleted file mode 100644
index 23f497468ed4..000000000000
--- a/tests/model_executor/test_keye_mrope.py
+++ /dev/null
@@ -1,147 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-from dataclasses import dataclass, field
-
-import pytest
-import torch
-
-from vllm.model_executor.models.keye import KeyeForConditionalGeneration
-from vllm.multimodal.inputs import (
-    MultiModalFeatureSpec,
-    MultiModalFieldElem,
-    MultiModalKwargsItem,
-    PlaceholderRange,
-)
-
-pytestmark = pytest.mark.skip_global_cleanup
-
-
-@pytest.fixture(autouse=True, scope="module")
-def _force_cpu_default_device():
-    original = torch.get_default_device()
-    torch.set_default_device("cpu")
-    yield
-    torch.set_default_device(original)
-
-
-@dataclass
-class DummyVisionConfig:
-    spatial_merge_size: int = 2
-
-
-@dataclass
-class DummyConfig:
-    vision_config: DummyVisionConfig = field(default_factory=DummyVisionConfig)
-
-
-def make_model(config: DummyConfig) -> KeyeForConditionalGeneration:
-    model = object.__new__(
-        KeyeForConditionalGeneration  # type: ignore[type-abstract]  # protocol attrs unset
-    )
-    model.config = config
-    return model
-
-
-def make_mm_feature(
-    *,
-    modality: str,
-    offset: int,
-    length: int,
-    grid_thw: tuple[int, int, int] | list[tuple[int, int, int]],
-) -> MultiModalFeatureSpec:
-    field_name = "image_grid_thw" if modality == "image" else "video_grid_thw"
-    return MultiModalFeatureSpec(
-        data=MultiModalKwargsItem(
-            {
-                field_name: MultiModalFieldElem(
-                    data=torch.tensor(grid_thw),
-                    field=None,  # HACK.
-                ),
-            }
-        ),
-        modality=modality,
-        identifier="DUMMY",
-        mm_position=PlaceholderRange(offset=offset, length=length),
-    )
-
-
-def test_get_mrope_input_positions_text_only():
-    model = make_model(DummyConfig())
-
-    positions, delta = model.get_mrope_input_positions(
-        input_tokens=[11, 12, 13, 14, 15],
-        mm_features=[],
-    )
-
-    expected = torch.tensor(
-        [
-            [0, 1, 2, 3, 4],
-            [0, 1, 2, 3, 4],
-            [0, 1, 2, 3, 4],
-        ]
-    )
-
-    assert torch.equal(positions, expected)
-    assert delta == 0
-
-
-def test_get_mrope_input_positions_single_image():
-    model = make_model(DummyConfig())
-    mm_features = [
-        make_mm_feature(
-            modality="image",
-            offset=1,
-            length=4,
-            grid_thw=(1, 4, 4),
-        )
-    ]
-
-    positions, delta = model.get_mrope_input_positions(
-        input_tokens=[10, 20, 21, 22, 23, 30, 31],
-        mm_features=mm_features,
-    )
-
-    expected = torch.tensor(
-        [
-            [0, 1, 1, 1, 1, 3, 4],
-            [0, 1, 1, 2, 2, 3, 4],
-            [0, 1, 2, 1, 2, 3, 4],
-        ]
-    )
-
-    assert torch.equal(positions, expected)
-    assert delta == -2
-
-
-def test_get_mrope_input_positions_interleaved_image_and_video():
-    model = make_model(DummyConfig())
-    mm_features = [
-        make_mm_feature(
-            modality="image",
-            offset=1,
-            length=4,
-            grid_thw=(1, 4, 4),
-        ),
-        make_mm_feature(
-            modality="video",
-            offset=7,
-            length=4,
-            grid_thw=[(2, 4, 2)],
-        ),
-    ]
-
-    positions, delta = model.get_mrope_input_positions(
-        input_tokens=[10, 20, 21, 22, 23, 30, 31, 40, 41, 42, 43, 50, 51],
-        mm_features=mm_features,
-    )
-
-    expected = torch.tensor(
-        [
-            [0, 1, 1, 1, 1, 3, 4, 5, 5, 7, 7, 9, 10],
-            [0, 1, 1, 2, 2, 3, 4, 5, 6, 7, 8, 9, 10],
-            [0, 1, 2, 1, 2, 3, 4, 5, 5, 7, 7, 9, 10],
-        ]
-    )
-
-    assert torch.equal(positions, expected)
-    assert delta == -2
diff --git a/tests/model_executor/test_keye_vl1_5_mrope.py b/tests/model_executor/test_keye_vl1_5_mrope.py
deleted file mode 100644
index 8942bee25585..000000000000
--- a/tests/model_executor/test_keye_vl1_5_mrope.py
+++ /dev/null
@@ -1,147 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-from dataclasses import dataclass, field
-
-import pytest
-import torch
-
-from vllm.model_executor.models.keye_vl1_5 import KeyeVL1_5ForConditionalGeneration
-from vllm.multimodal.inputs import (
-    MultiModalFeatureSpec,
-    MultiModalFieldElem,
-    MultiModalKwargsItem,
-    PlaceholderRange,
-)
-
-pytestmark = pytest.mark.skip_global_cleanup
-
-
-@pytest.fixture(autouse=True, scope="module")
-def _force_cpu_default_device():
-    original = torch.get_default_device()
-    torch.set_default_device("cpu")
-    yield
-    torch.set_default_device(original)
-
-
-@dataclass
-class DummyVisionConfig:
-    spatial_merge_size: int = 2
-
-
-@dataclass
-class DummyConfig:
-    vision_config: DummyVisionConfig = field(default_factory=DummyVisionConfig)
-
-
-def make_model(config: DummyConfig) -> KeyeVL1_5ForConditionalGeneration:
-    model = object.__new__(
-        KeyeVL1_5ForConditionalGeneration  # type: ignore[type-abstract]  # protocol attrs unset
-    )
-    model.config = config
-    return model
-
-
-def make_mm_feature(
-    *,
-    modality: str,
-    offset: int,
-    length: int,
-    grid_thw: tuple[int, int, int] | list[tuple[int, int, int]],
-    is_embed: list[bool] | None = None,
-) -> MultiModalFeatureSpec:
-    field_name = "image_grid_thw" if modality == "image" else "video_grid_thw"
-    return MultiModalFeatureSpec(
-        data=MultiModalKwargsItem(
-            {
-                field_name: MultiModalFieldElem(
-                    data=torch.tensor(grid_thw),
-                    field=None,  # HACK.
-                ),
-            }
-        ),
-        modality=modality,
-        identifier="DUMMY",
-        mm_position=PlaceholderRange(
-            offset=offset,
-            length=length,
-            is_embed=None if is_embed is None else torch.tensor(is_embed),
-        ),
-    )
-
-
-def test_get_mrope_input_positions_text_only():
-    model = make_model(DummyConfig())
-
-    positions, delta = model.get_mrope_input_positions(
-        input_tokens=[11, 12, 13, 14, 15],
-        mm_features=[],
-    )
-
-    expected = torch.tensor(
-        [
-            [0, 1, 2, 3, 4],
-            [0, 1, 2, 3, 4],
-            [0, 1, 2, 3, 4],
-        ]
-    )
-
-    assert torch.equal(positions, expected)
-    assert delta == 0
-
-
-def test_get_mrope_input_positions_single_image():
-    model = make_model(DummyConfig())
-    mm_features = [
-        make_mm_feature(
-            modality="image",
-            offset=1,
-            length=4,
-            grid_thw=(1, 4, 4),
-        )
-    ]
-
-    positions, delta = model.get_mrope_input_positions(
-        input_tokens=[10, 20, 21, 22, 23, 30, 31],
-        mm_features=mm_features,
-    )
-
-    expected = torch.tensor(
-        [
-            [0, 1, 1, 1, 1, 3, 4],
-            [0, 1, 1, 2, 2, 3, 4],
-            [0, 1, 2, 1, 2, 3, 4],
-        ]
-    )
-
-    assert torch.equal(positions, expected)
-    assert delta == -2
-
-
-def test_get_mrope_input_positions_video_uses_embed_ranges():
-    model = make_model(DummyConfig())
-    mm_features = [
-        make_mm_feature(
-            modality="video",
-            offset=1,
-            length=8,
-            grid_thw=[(2, 4, 2)],
-            is_embed=[False, False, True, True, False, False, True, True],
-        )
-    ]
-
-    positions, delta = model.get_mrope_input_positions(
-        input_tokens=[10, 101, 102, 20, 21, 103, 104, 30, 31, 40, 41],
-        mm_features=mm_features,
-    )
-
-    expected = torch.tensor(
-        [
-            [0, 1, 2, 3, 3, 5, 6, 7, 7, 9, 10],
-            [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
-            [0, 1, 2, 3, 3, 5, 6, 7, 7, 9, 10],
-        ]
-    )
-
-    assert torch.equal(positions, expected)
-    assert delta == 0
diff --git a/tests/model_executor/test_qwen3_omni.py b/tests/model_executor/test_qwen3_omni.py
index d8acd8196049..c3658f138781 100644
--- a/tests/model_executor/test_qwen3_omni.py
+++ b/tests/model_executor/test_qwen3_omni.py
@@ -387,6 +387,7 @@ def test_dspark_shares_target_embedding_with_smaller_draft_vocabulary():
             attention_backend=None,
             kv_cache_dtype=None,
             draft_load_config=None,
+            draft_lm_head_quantization=None,
         ),
         parallel_config=ParallelConfig(),
         attention_config=SimpleNamespace(backend=None),
diff --git a/tests/model_executor/test_replayssm_warmup.py b/tests/model_executor/test_replayssm_warmup.py
index 5db0f9a50d0f..48f563936c95 100644
--- a/tests/model_executor/test_replayssm_warmup.py
+++ b/tests/model_executor/test_replayssm_warmup.py
@@ -36,8 +36,15 @@ def _autotune_runner(
     max_num_tokens: int = 100,
     use_replayssm: bool = True,
     backend: MambaBackendEnum = MambaBackendEnum.FLASHINFER,
+    committable_blocks: int | None = None,
 ) -> SimpleNamespace:
+    extensible_kv_cache = None
+    if committable_blocks is not None:
+        extensible_kv_cache = SimpleNamespace(
+            warmup_committable_blocks=committable_blocks, commit=Mock()
+        )
     return SimpleNamespace(
+        extensible_kv_cache=extensible_kv_cache,
         vllm_config=SimpleNamespace(
             cache_config=SimpleNamespace(use_replayssm=use_replayssm),
             mamba_config=SimpleNamespace(backend=backend),
@@ -46,7 +53,9 @@ def _autotune_runner(
         uniform_decode_query_len=query_len,
         decode_query_len=query_len,
         max_num_tokens=max_num_tokens,
-        scheduler_config=SimpleNamespace(max_num_seqs=max_num_seqs),
+        scheduler_config=SimpleNamespace(
+            max_num_seqs=max_num_seqs, max_num_batched_tokens=max_num_tokens
+        ),
         kv_cache_config=SimpleNamespace(num_blocks=num_blocks),
     )
 
@@ -59,8 +68,10 @@ def _autotune_runner(
         (dict(query_len=6, use_v2_model_runner=True), 16),
         # num_blocks - 1 is the binding constraint.
         (dict(query_len=1, max_num_tokens=128, max_num_seqs=64, num_blocks=5), 4),
+        # An extensible cache bounds by what warmup may commit, not capacity.
+        (dict(query_len=1, max_num_tokens=128, num_blocks=64, committable_blocks=4), 3),
     ],
-    ids=["v1", "v2", "clamped_to_state_capacity"],
+    ids=["v1", "v2", "clamped_to_state_capacity", "clamped_to_committable"],
 )
 def test_replayssm_autotune_decode_kwargs(runner_kwargs, expected_num_reqs):
     query_len = runner_kwargs["query_len"]
@@ -117,6 +128,25 @@ def test_replayssm_autotune_kwargs_skipped(runner_kwargs, flashinfer_supported):
     assert result is None
 
 
+def test_replayssm_autotune_warmup_commits_dummy_slots_first():
+    """The dummy slots 1..max_num_reqs must be committed memory before the
+    autotune run and its cleanup touch them."""
+    runner = _autotune_runner(
+        query_len=1, max_num_tokens=128, num_blocks=64, committable_blocks=4
+    )
+    events: list[str] = []
+    runner.extensible_kv_cache.commit = lambda n: events.append(f"commit {n}")
+    runner._dummy_run = lambda **kwargs: events.append("dummy_run")
+    with (
+        patch.object(
+            warmup, "flashinfer_replayssm_autotune_supported", return_value=True
+        ),
+        patch.object(warmup, "_temporary_replayssm_autotune_state"),
+    ):
+        warmup.replayssm_autotune_warmup(runner)
+    assert events == ["commit 4", "dummy_run"]
+
+
 def test_replayssm_autotune_slots_restore_state_and_trackers():
     mixer = MambaMixer2.__new__(MambaMixer2)
     torch.nn.Module.__init__(mixer)
diff --git a/tests/model_executor/test_routed_experts_capture.py b/tests/model_executor/test_routed_experts_capture.py
index 7b96d277efdf..0a0fb3e45c52 100644
--- a/tests/model_executor/test_routed_experts_capture.py
+++ b/tests/model_executor/test_routed_experts_capture.py
@@ -660,15 +660,10 @@ def test_get_aux_output_connector_passes_config(monkeypatch):
         scheduler_config=SimpleNamespace(max_num_batched_tokens=32)
     )
     model = Mock()
-    kv_cache_config = Mock()
 
-    result = aux_output_worker.get_aux_output_connector(model, config, kv_cache_config)
+    result = aux_output_worker.get_aux_output_connector(model, config)
 
-    constructor.assert_called_once_with(
-        model=model,
-        kv_cache_config=kv_cache_config,
-        vllm_config=config,
-    )
+    constructor.assert_called_once_with(model=model, vllm_config=config)
     assert result is connector
 
 
@@ -694,11 +689,13 @@ def test_aux_output_worker_connector_binds_capture_on_non_output_rank(monkeypatc
     )
     model = Mock()
     connector = aux_output_worker.AuxOutputWorkerConnector(
-        vllm_config=config,
-        model=model,
-        kv_cache_config=SimpleNamespace(kv_cache_groups=[_full_attention_kv_group()]),
+        vllm_config=config, model=model
+    )
+    connector.initialize_output_storage(
+        SimpleNamespace(kv_cache_groups=[_full_attention_kv_group()])
     )
 
+    assert connector._store is None and connector._buffer is None
     constructor.assert_called_once_with(
         max_num_batched_tokens=32,
         vllm_config=config,
@@ -730,7 +727,8 @@ def test_aux_output_worker_connector_default_capacity(monkeypatch):
     monkeypatch.setattr(
         aux_output_worker, "BackgroundBlockObjectStore", background_store_constructor
     )
-    monkeypatch.setattr(aux_output_worker, "RoutedExpertsBuffer", Mock())
+    buffer_constructor = Mock()
+    monkeypatch.setattr(aux_output_worker, "RoutedExpertsBuffer", buffer_constructor)
 
     config = SimpleNamespace(
         aux_output_config=SimpleNamespace(max_bytes=None),
@@ -739,16 +737,18 @@ def test_aux_output_worker_connector_default_capacity(monkeypatch):
         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(
+    connector = aux_output_worker.AuxOutputWorkerConnector(
+        vllm_config=config, model=Mock()
+    )
+    store_constructor.assert_not_called()
+    buffer_constructor.assert_not_called()
+
+    connector.initialize_output_storage(
+        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
diff --git a/tests/model_executor/test_utils.py b/tests/model_executor/test_utils.py
index 6c098a95f663..b25b042d3e4a 100644
--- a/tests/model_executor/test_utils.py
+++ b/tests/model_executor/test_utils.py
@@ -13,7 +13,14 @@
     PackedvLLMParameter,
     permute_param_layout_,
 )
-from vllm.model_executor.utils import replace_parameter
+from vllm.model_executor.utils import (
+    get_live_param,
+    kernel_parameters,
+    move_parameter,
+    replace_parameter,
+    retire_parameter,
+    set_kernel_parameter,
+)
 
 
 @pytest.fixture
@@ -44,6 +51,7 @@ def test_replace_parameter_preserves_custom_attribute(wrap_in_parameter: bool) -
         "weight", torch.nn.Parameter(torch.zeros(4, 4), requires_grad=False)
     )
     original = layer.weight
+    original_ptr = original.data_ptr()
 
     new_data: torch.Tensor = torch.ones(4, 4).t()
     if wrap_in_parameter:
@@ -63,18 +71,30 @@ def test_replace_parameter_preserves_custom_attribute(wrap_in_parameter: bool) -
 
     assert layer.weight is original
     assert layer.weight.stride() == new_data.stride()
-    assert layer.weight.data_ptr() == new_data.data_ptr()
+    # same byte size: copied into the loaded storage, address kept
+    assert layer.weight.data_ptr() == original_ptr
 
 
 @pytest.mark.parametrize(
     "new_data,copies",
     [
         (torch.ones(4, 4, dtype=torch.float32, device="cpu"), True),
-        (torch.ones(4, 4, dtype=torch.float32, device="cpu").t(), False),
+        (torch.ones(4, 4, dtype=torch.float32, device="cpu").t(), True),
         (torch.ones(8, 4, dtype=torch.float32, device="cpu"), False),
-        (torch.ones(4, 4, dtype=torch.bfloat16, device="cpu"), False),
+        # A result that fits (equal or smaller) goes into the loaded storage,
+        # which never shrinks.
+        (torch.ones(4, 4, dtype=torch.bfloat16, device="cpu"), True),
+        (torch.ones(4, 2, dtype=torch.float32, device="cpu"), True),
+        (torch.tensor(1.0, device="cpu"), True),
+    ],
+    ids=[
+        "same-layout-copies",
+        "new-strides-copies",
+        "larger-shape",
+        "smaller-dtype-copies",
+        "smaller-shape-copies",
+        "scalar-copies",
     ],
-    ids=["same-layout-copies", "new-strides", "new-shape", "new-dtype"],
 )
 def test_replace_parameter_keeps_object_by_default(
     single_rank_tp: None, new_data: torch.Tensor, copies: bool
@@ -102,6 +122,42 @@ def weight_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None:
     assert torch.equal(layer.weight.data, new_data)
     # Copying in place keeps addresses captured in CUDA graphs valid.
     assert (layer.weight.data_ptr() == old_data_ptr) == copies
+    assert layer.weight.untyped_storage().nbytes() >= 4 * 4 * 4
+
+
+def test_replace_parameter_online_quantization_frees_source() -> None:
+    layer = torch.nn.Module()
+    layer.quant_method = type("Online", (), {"uses_meta_device": True})()
+    layer.register_parameter("weight", torch.nn.Parameter(torch.zeros(4, 4)))
+
+    replace_parameter(layer, "weight", torch.ones(4, 4, dtype=torch.bfloat16))
+
+    assert layer.weight.untyped_storage().nbytes() == 4 * 4 * 2
+
+
+def test_humming_replace_keeps_loaded_params(single_rank_tp: None) -> None:
+    """A new tensor in a leftover's storage, or taking a large leftover's place,
+    is written to that leftover param, registered under both names."""
+    from vllm.model_executor.layers.quantization.utils.humming.linear import (
+        replace_layer_parameters,
+    )
+
+    layer = torch.nn.Module()
+    packed = torch.nn.Parameter(torch.ones(64, 8), requires_grad=False)
+    scale_inv = torch.nn.Parameter(torch.ones(4), requires_grad=False)
+    layer.register_parameter("weight_packed", packed)
+    layer.register_parameter("weight_scale_inv", scale_inv)
+    tensors = {
+        "weight": torch.full((32, 8), 2.0),
+        "weight_scale": scale_inv.data.view(2, 2),
+    }
+
+    replace_layer_parameters(layer, tensors, list(layer._parameters))
+
+    assert layer.weight is packed and layer.weight_scale is scale_inv
+    assert torch.equal(layer.weight, torch.full((32, 8), 2.0))
+    assert packed.untyped_storage().nbytes() == 64 * 8 * packed.element_size()
+    assert set(kernel_parameters(layer)) == {"weight", "weight_scale"}
 
 
 def test_replace_parameter_preserves_weight_loader() -> None:
@@ -240,6 +296,54 @@ def test_replace_parameter_registers_new_param_for_requantization_source() -> No
     assert torch.equal(layer.weight, torch.ones(4, 4))
 
 
+def test_move_parameter_aliases_one_param(single_rank_tp: None) -> None:
+    """Both names refer to the loaded param; a later replace updates both."""
+
+    def weight_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None:
+        pass
+
+    layer = torch.nn.Module()
+    src = ModelWeightParameter(
+        data=torch.ones(4, 4), input_dim=1, output_dim=0, weight_loader=weight_loader
+    )
+    layer.register_parameter("weight_packed", src)
+
+    move_parameter(layer, "weight_packed", "weight")
+    replace_parameter(layer, "weight", torch.full((2, 8), 2.0))
+
+    assert layer.weight_packed is src and layer.weight is src
+    assert torch.equal(src.data, torch.full((2, 8), 2.0))
+    assert src.weight_loader is weight_loader
+    assert src._input_dim == 1 and src._output_dim == 0
+    assert list(kernel_parameters(layer)) == ["weight"]
+
+    # Kernels writing through the checkpoint name make it the runtime name again.
+    set_kernel_parameter(layer, "weight_packed", torch.ones(2, 8))
+    assert list(kernel_parameters(layer)) == ["weight_packed"]
+
+
+def test_move_parameter_rejects_another_param(single_rank_tp: None) -> None:
+    layer = torch.nn.Module()
+    layer.register_parameter("a", torch.nn.Parameter(torch.ones(2)))
+    layer.register_parameter("b", torch.nn.Parameter(torch.ones(2)))
+    with pytest.raises(ValueError):
+        move_parameter(layer, "a", "b")
+
+
+def test_retired_param_keeps_data_but_is_not_live(single_rank_tp: None) -> None:
+    layer = torch.nn.Module()
+    layer.register_parameter("scale", torch.nn.Parameter(torch.ones(2)))
+    layer.register_parameter("weight", torch.nn.Parameter(torch.ones(4)))
+
+    retire_parameter(layer, "scale")
+
+    assert torch.equal(layer.scale, torch.ones(2))
+    assert get_live_param(layer, "scale") is None
+    assert list(kernel_parameters(layer)) == ["weight"]
+    set_kernel_parameter(layer, "scale", torch.zeros(2))
+    assert get_live_param(layer, "scale") is layer.scale
+
+
 def test_replace_parameter_with_own_parameter_keeps_layout(single_rank_tp) -> None:
     """Handing back the registered parameter keeps it unchanged."""
 
diff --git a/tests/models/deepseek_v32/test_pp_topk.py b/tests/models/deepseek_v32/test_pp_topk.py
new file mode 100644
index 000000000000..0015c88aed83
--- /dev/null
+++ b/tests/models/deepseek_v32/test_pp_topk.py
@@ -0,0 +1,131 @@
+# 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.distributed import parallel_state
+from vllm.model_executor.models import utils as model_utils
+from vllm.models.deepseek_v32.nvidia import model as deepseek_v32_model
+from vllm.sequence import IntermediateTensors
+
+# GLM-5.3: each group of 4 layers from layer 3 reuses its first layer's top-k.
+GLM5_TOPK = dict(index_topk_freq=4, index_skip_topk_offset=3)
+
+
+def _set_pp_size(monkeypatch, pp_size: int) -> None:
+    monkeypatch.setattr(
+        parallel_state, "get_pp_group", lambda: SimpleNamespace(world_size=pp_size)
+    )
+
+
+class TestPPSplitsLayerSharing:
+    @pytest.mark.parametrize(
+        "num_layers,users,consumers,pp_size,expected",
+        [
+            (8, range(8), {1, 2, 3, 5, 6, 7}, 2, False),
+            (8, range(8), {1, 2, 4, 5, 6, 7}, 2, True),
+            (8, range(8), {1, 2, 4, 5, 6, 7}, 1, False),
+            # Layers 4-5 do not use the state, so stage 1's first user is 6.
+            (8, {0, 1, 2, 3, 6, 7}, {1, 2, 3, 7}, 2, False),
+            (8, {0, 1, 2, 3, 6, 7}, {1, 2, 3, 6, 7}, 2, True),
+            (12, range(12), {1, 2, 3, 5, 6, 7, 9, 10, 11}, 3, False),
+            (12, range(12), {1, 2, 3, 4, 6, 7, 9, 10, 11}, 3, True),
+        ],
+        ids=[
+            "split-on-producer",
+            "split-on-consumer",
+            "no-pp",
+            "skip-unrelated-layers",
+            "unrelated-then-consumer",
+            "pp3-aligned",
+            "pp3-middle-split",
+        ],
+    )
+    def test_stage_boundaries(
+        self, monkeypatch, num_layers, users, consumers, pp_size, expected
+    ):
+        _set_pp_size(monkeypatch, pp_size)
+        assert (
+            model_utils.pp_splits_layer_sharing(num_layers, users, consumers)
+            == expected
+        )
+
+    @pytest.mark.parametrize(
+        "num_layers,expected",
+        [(78, True), (76, False)],
+        ids=["glm5-default-split", "split-on-group-start"],
+    )
+    def test_topk_buffers(self, monkeypatch, num_layers, expected):
+        _set_pp_size(monkeypatch, 2)
+        config = SimpleNamespace(
+            num_hidden_layers=num_layers, index_topk=8, **GLM5_TOPK
+        )
+        buffer = torch.empty(4, 8, dtype=torch.int32)
+        buffers = model_utils.pp_topk_buffers(config, buffer, range(num_layers))
+        assert buffers == ({"topk_indices": buffer} if expected else {})
+
+
+def test_empty_intermediate_tensors_follow_buffers():
+    buffer = torch.empty(16, 8, dtype=torch.int32)
+    make_empty = model_utils.make_empty_intermediate_tensors_factory(
+        ["hidden_states", "residual"], 4, {"topk_indices": buffer}
+    )
+    tensors = make_empty(3, torch.bfloat16, torch.device("cpu"))
+
+    assert tensors["residual"].shape == (3, 4)
+    assert tensors["topk_indices"].shape == (3, 8)
+    assert tensors["topk_indices"].dtype == torch.int32
+
+
+class _TopkReusingLayer(nn.Module):
+    """Reads the shared top-k like a reusing layer, then indexes like the next."""
+
+    def __init__(self, topk_indices_buffer: torch.Tensor) -> None:
+        super().__init__()
+        self.topk_indices_buffer = topk_indices_buffer
+        self.seen_topk: torch.Tensor | None = None
+
+    def forward(self, positions, hidden_states, residual, attn_in=None):
+        num_tokens = hidden_states.shape[0]
+        self.seen_topk = self.topk_indices_buffer[:num_tokens].clone()
+        self.topk_indices_buffer[:num_tokens] += 10
+        return hidden_states, torch.zeros_like(hidden_states)
+
+
+def test_middle_pp_stage_receives_and_forwards_topk(monkeypatch):
+    model = object.__new__(deepseek_v32_model.DeepseekV32Model)
+    nn.Module.__init__(model)
+    model.topk_indices_buffer = torch.full((8, 2), -1, dtype=torch.int32)
+    layer = _TopkReusingLayer(model.topk_indices_buffer)
+    model.layers = nn.ModuleList([layer])
+    model.start_layer, model.end_layer = 0, 1
+    model.use_sequence_parallel = False
+    model.aux_hidden_state_layers = ()
+    model.pp_buffers = {"topk_indices": model.topk_indices_buffer}
+    monkeypatch.setattr(
+        deepseek_v32_model,
+        "get_pp_group",
+        lambda: SimpleNamespace(is_first_rank=False, is_last_rank=False),
+    )
+    monkeypatch.setattr(
+        deepseek_v32_model, "tensor_model_parallel_all_reduce", lambda t: t
+    )
+
+    received_topk = torch.arange(6, dtype=torch.int32).view(3, 2)
+    output = model(
+        None,
+        torch.arange(3),
+        IntermediateTensors(
+            {
+                "hidden_states": torch.zeros(3, 4),
+                "topk_indices": received_topk,
+            }
+        ),
+    )
+
+    torch.testing.assert_close(layer.seen_topk, received_topk)
+    torch.testing.assert_close(output["topk_indices"], received_topk + 10)
diff --git a/tests/models/deepseek_v32/test_sequence_parallel.py b/tests/models/deepseek_v32/test_sequence_parallel.py
index c0378dd5e3c2..b72c5aec5d68 100644
--- a/tests/models/deepseek_v32/test_sequence_parallel.py
+++ b/tests/models/deepseek_v32/test_sequence_parallel.py
@@ -9,6 +9,7 @@
 
 from vllm.models.deepseek_v32.nvidia import model as deepseek_v32_model
 from vllm.models.deepseek_v32.nvidia import mtp as deepseek_v32_mtp
+from vllm.sequence import IntermediateTensors
 
 
 class _IdentityNorm(nn.Module):
@@ -146,3 +147,40 @@ def test_mtp_projects_sequence_shard_and_restores_full_output(monkeypatch):
     torch.testing.assert_close(hidden_states, expected)
     torch.testing.assert_close(recycled_hidden_states, expected)
     norm.assert_called_once()
+
+
+class _PartialSumLayer(nn.Module):
+    """Returns a per-TP-rank partial MLP output, like DeepseekV32DecoderLayer."""
+
+    def forward(self, positions, hidden_states, residual, attn_in=None):
+        assert residual is None
+        return torch.ones_like(hidden_states), hidden_states + 1
+
+
+def test_pp_stage_sends_tp_reduced_residual_stream(monkeypatch):
+    """PP send/recv needs TP-replicated tensors, so a middle stage must reduce
+    its partial hidden_states and send the full residual stream only."""
+    model = object.__new__(deepseek_v32_model.DeepseekV32Model)
+    nn.Module.__init__(model)
+    model.layers = nn.ModuleList([_PartialSumLayer()])
+    model.start_layer, model.end_layer = 0, 1
+    model.use_sequence_parallel = False
+    model.aux_hidden_state_layers = ()
+    model.pp_buffers = {}
+
+    pp_group = SimpleNamespace(is_first_rank=False, is_last_rank=False)
+    monkeypatch.setattr(deepseek_v32_model, "get_pp_group", lambda: pp_group)
+    monkeypatch.setattr(
+        deepseek_v32_model,
+        "tensor_model_parallel_all_reduce",
+        lambda tensor: tensor * 2,
+    )
+
+    received = torch.full((3, 2), 5.0)
+    output = model(
+        None, torch.arange(3), IntermediateTensors({"hidden_states": received})
+    )
+
+    assert set(output.tensors) == {"hidden_states"}
+    # residual (5 + 1) + partial (1) all-reduced over 2 TP ranks.
+    torch.testing.assert_close(output["hidden_states"], torch.full((3, 2), 8.0))
diff --git a/tests/models/deepseek_v41/nvidia/specialized_decode/test_kernels.py b/tests/models/deepseek_v41/nvidia/specialized_decode/test_kernels.py
index 12ecc14a1303..c3d1deb741b4 100644
--- a/tests/models/deepseek_v41/nvidia/specialized_decode/test_kernels.py
+++ b/tests/models/deepseek_v41/nvidia/specialized_decode/test_kernels.py
@@ -435,7 +435,7 @@ def test_sparse_mla_persistent_kv_counters(dtype, query_lens):
     """The decode on the shared KV counters equals the decode on fresh ones
     bitwise, launch after launch, and leaves the counters zeroed."""
     if sparse_mla.kv_counter_buffer(1) is None:
-        pytest.skip("FlashInfer is not the version the persistent counters need")
+        pytest.skip("FlashInfer's sparse-MLA decode takes no KV counter buffer")
     torch.manual_seed(0)
     heads, head_dim, page, top_k = 64, 512, 64, 1024
     tokens = sum(query_lens)
diff --git a/tests/models/deepseek_v41/nvidia/specialized_decode/test_numerics.py b/tests/models/deepseek_v41/nvidia/specialized_decode/test_numerics.py
index 3442d862e491..cc6445c837f0 100644
--- a/tests/models/deepseek_v41/nvidia/specialized_decode/test_numerics.py
+++ b/tests/models/deepseek_v41/nvidia/specialized_decode/test_numerics.py
@@ -18,13 +18,12 @@
 * one CUDA graph of all of these, replayed ``REPLAYS`` times bitwise.
 
 Needs four SM100-family GPUs with NVLink multicast and flashinfer-python
-0.7.2; about 5 minutes. For a race stress run, run this file as a module
-(``python -m tests.models.deepseek_v41.nvidia.specialized_decode.test_numerics``)
+0.7.2 or newer; about 5 minutes. For a race stress run, run this file as a
+module (``python -m tests.models.deepseek_v41.nvidia.specialized_decode.test_numerics``)
 with ``--replays 100000``.
 """
 
 import argparse
-import importlib.metadata
 import os
 import tempfile
 from types import SimpleNamespace
@@ -809,14 +808,13 @@ def run(replays: int) -> str | None:
 def test_specialized_decode_numerics():
     if torch.accelerator.device_count() < TP:
         pytest.skip("needs four GPUs")
-    from vllm.models.deepseek_v41.nvidia.specialized_decode import FLASHINFER_VERSION
-
-    try:
-        version = importlib.metadata.version("flashinfer-python")
-    except importlib.metadata.PackageNotFoundError:
-        version = None
-    if version != FLASHINFER_VERSION:
-        pytest.skip(f"needs flashinfer-python {FLASHINFER_VERSION}, found {version}")
+    from vllm.models.deepseek_v41.nvidia.specialized_decode import (
+        MIN_FLASHINFER_VERSION,
+        uses_checked_flashinfer,
+    )
+
+    if not uses_checked_flashinfer("the numerics test"):
+        pytest.skip(f"needs flashinfer-python >= {MIN_FLASHINFER_VERSION}")
     skipped = run(REPLAYS)
     if skipped:
         pytest.skip(skipped)
diff --git a/tests/models/glm5next/test_epd_metadata.py b/tests/models/glm5next/test_epd_metadata.py
new file mode 100644
index 000000000000..4d3063b98331
--- /dev/null
+++ b/tests/models/glm5next/test_epd_metadata.py
@@ -0,0 +1,132 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""EPD metadata-only parser contract shared by GLM-4V and GLM-5.3-Flash."""
+
+import pytest
+import torch
+
+from vllm.model_executor.models.glm4_1v import Glm4vMultiModalDataParser
+
+
+@pytest.mark.parametrize("allow_missing", [False, True])
+def test_glm5_image_metadata_requires_ec_consumer(allow_missing):
+    parser = Glm4vMultiModalDataParser(
+        spatial_merge_size=2,
+        allow_missing_mm_embeddings=allow_missing,
+    )
+    data = {"image_grid_thw": torch.tensor([[1, 4, 6], [1, 8, 8]])}
+    if not allow_missing:
+        with pytest.raises(ValueError, match="image_embeds"):
+            parser.parse_mm_data({"image": data})
+        return
+
+    items = parser.parse_mm_data({"image": data})["image"]
+    assert len(items) == 2
+    assert items.get(1)["image_grid_thw"].tolist() == [1, 8, 8]
+    assert items.get_processor_data() == {}
+    assert parser.placeholder_metadata_fields("image") == {"image_grid_thw"}
+    assert parser.placeholder_metadata_fields("video") == {
+        "video_grid_thw",
+        "timestamps",
+    }
+
+
+@pytest.mark.parametrize("allow_missing", [False, True])
+def test_glm5_video_metadata_requires_timestamps(allow_missing):
+    parser = Glm4vMultiModalDataParser(
+        spatial_merge_size=2,
+        video_needs_metadata=True,
+        allow_missing_mm_embeddings=allow_missing,
+    )
+    data = {
+        "video_grid_thw": torch.tensor([[2, 4, 4]]),
+        "timestamps": torch.tensor([[0, 1]]),
+    }
+    if not allow_missing:
+        with pytest.raises(ValueError, match="video_embeds"):
+            parser.parse_mm_data({"video": data})
+        return
+
+    items = parser.parse_mm_data({"video": data})["video"]
+    assert len(items) == 1
+    assert items.get(0)["timestamps"].tolist() == [0, 1]
+    assert items.get_processor_data() == {}
+
+
+def test_glm_dict_embeds_must_match_grid_and_hidden_size():
+    """Dict embeddings are checked here; metadata-only inputs stay valid."""
+    parser = Glm4vMultiModalDataParser(
+        spatial_merge_size=2,
+        expected_hidden_size=16,
+        video_needs_metadata=True,
+    )
+    grid = torch.tensor([[1, 4, 4]])
+    good = torch.zeros(4, 16)
+    image = parser.parse_mm_data(
+        {"image": {"image_embeds": good, "image_grid_thw": grid}}
+    )["image"]
+    assert len(image) == 1
+    video = parser.parse_mm_data(
+        {
+            "video": {
+                "video_embeds": good,
+                "video_grid_thw": grid,
+                "timestamps": torch.tensor([[0]]),
+            }
+        }
+    )["video"]
+    assert len(video) == 1
+
+    bad_cases = [
+        (torch.zeros(4, 1), "hidden dimension"),
+        (torch.zeros(4), "must be 2D"),
+        (torch.zeros(3, 16), "row count"),
+    ]
+    for bad, match in bad_cases:
+        with pytest.raises(ValueError, match=match):
+            parser.parse_mm_data(
+                {"image": {"image_embeds": bad, "image_grid_thw": grid}}
+            )
+        with pytest.raises(ValueError, match=match):
+            parser.parse_mm_data(
+                {
+                    "video": {
+                        "video_embeds": bad,
+                        "video_grid_thw": grid,
+                        "timestamps": torch.tensor([[0]]),
+                    }
+                }
+            )
+
+    metadata_only = Glm4vMultiModalDataParser(
+        spatial_merge_size=2,
+        expected_hidden_size=16,
+        video_needs_metadata=True,
+        allow_missing_mm_embeddings=True,
+    )
+    image_items = metadata_only.parse_mm_data({"image": {"image_grid_thw": grid}})
+    assert len(image_items["image"]) == 1
+    two_grids = torch.tensor([[1, 4, 4], [1, 4, 4]])
+    eight_rows = torch.zeros(8, 16)
+    two_images = parser.parse_mm_data(
+        {"image": {"image_embeds": eight_rows, "image_grid_thw": two_grids}}
+    )["image"]
+    assert len(two_images) == 2
+    with pytest.raises(ValueError, match="row count"):
+        parser.parse_mm_data(
+            {
+                "image": {
+                    "image_embeds": torch.zeros(64, 16),
+                    "image_grid_thw": two_grids,
+                }
+            }
+        )
+    video_items = metadata_only.parse_mm_data(
+        {
+            "video": {
+                "video_grid_thw": grid,
+                "timestamps": torch.tensor([[0]]),
+            }
+        }
+    )
+    assert len(video_items["video"]) == 1
diff --git a/tests/models/inkling/test_contract_validation.py b/tests/models/inkling/test_contract_validation.py
index 4bdb75dcc499..e5c13903f4e8 100644
--- a/tests/models/inkling/test_contract_validation.py
+++ b/tests/models/inkling/test_contract_validation.py
@@ -37,8 +37,10 @@ def test_vision_scale_plan_breaks_assignment_ties_in_order():
 @pytest.mark.parametrize(
     ("config_cls", "kwargs", "missing"),
     [
-        (InklingAudioConfig, {"decoder_dmodel": 16}, "n_mel_bins"),
-        (InklingVisionConfig, {"decoder_dmodel": 16}, "vision_encoder_type"),
+        pytest.param(InklingAudioConfig, {"decoder_dmodel": 16}, "n_mel_bins"),
+        pytest.param(
+            InklingVisionConfig, {"decoder_dmodel": 16}, "vision_encoder_type"
+        ),
     ],
 )
 def test_enabled_tower_requires_architecture_fields(config_cls, kwargs, missing):
diff --git a/tests/models/inkling/test_mtp_input_fusion.py b/tests/models/inkling/test_mtp_input_fusion.py
index 7ab33dd45ae8..0e17eb523413 100644
--- a/tests/models/inkling/test_mtp_input_fusion.py
+++ b/tests/models/inkling/test_mtp_input_fusion.py
@@ -85,6 +85,7 @@ def test_embed_rmsnorm(n: int, t: int, ids_dtype: torch.dtype) -> None:
 
     # Gather + embed_norm (base model / MTP prefill embed path).
     out = embed_rmsnorm(ids, table, w, EPS)
+    assert isinstance(out, torch.Tensor)
     assert out.shape == (t, n)
     assert torch.equal(out, rmsnorm(ref_emb, w, EPS) if t else ref_emb)
 
diff --git a/tests/models/kimi_k3/amd_specialized_decode/test_engine_fault.py b/tests/models/kimi_k3/amd_specialized_decode/test_engine_fault.py
index 485dba75f57c..a7191164f50b 100644
--- a/tests/models/kimi_k3/amd_specialized_decode/test_engine_fault.py
+++ b/tests/models/kimi_k3/amd_specialized_decode/test_engine_fault.py
@@ -121,6 +121,8 @@ def test_a_stalled_rank_stops_the_engine(tmp_path):
         "VLLM_USE_V2_MODEL_RUNNER": "1",
         "VLLM_ROCM_USE_AITER": "1",
         "VLLM_ROCM_AITER_MLA_ASM_PADDING": "asm",
+        # M1 runs W_UK / W_UV as AITER's MXFP4 bmms, off by default
+        "VLLM_ROCM_USE_AITER_FP4BMM": "1",
     }
     body = {
         "model": MODEL,
diff --git a/tests/models/kimi_k3/amd_specialized_decode/test_numerics.py b/tests/models/kimi_k3/amd_specialized_decode/test_numerics.py
index dc396b099026..14a8eea7ab56 100644
--- a/tests/models/kimi_k3/amd_specialized_decode/test_numerics.py
+++ b/tests/models/kimi_k3/amd_specialized_decode/test_numerics.py
@@ -10,6 +10,10 @@
   Then a CUDA graph of K2 replayed 300 times, bit-identical, and a stalled rank
   (the fault hook): every rank's launch finishes and the output rank's flags
   fail the step.
+- M1a (AttnRes -> qkv_a / g_proj -> q / kv RMSNorm + the fp8 KV write -> q_b
+  -> fp4 W_UK -> the DCP query gather) and M1b (the a2a LSE combine -> fp4
+  W_UV -> the gate): one GPU, then TP8 + DCP8 with a graph of two MLA layers
+  (M1a -> M1b -> K2) replayed.
 """
 
 import pytest
@@ -60,7 +64,7 @@ def _randn(g, *shape, scale=1.0, dtype=bf):
 # ---------------------------------------------------------------- K1
 
 
-def _kda_inputs(nb: int, nacc: int, dim_first: bool, S: int, slots: int = 16):
+def _kda_inputs(nb: int, nacc: int, dim_first: bool, S: int, L: int):
     from vllm.models.kimi_k3.amd.specialized_decode.mono.attention.kda import (
         HD,
         HIDDEN,
@@ -71,7 +75,9 @@ def _kda_inputs(nb: int, nacc: int, dim_first: bool, S: int, slots: int = 16):
     )
 
     g = torch.Generator(device="cuda").manual_seed(0)
-    sl = 3 + (S - 1)
+    R = S // L
+    slots = S + 4
+    sl = 3 + (L - 1)
     w = {
         "prefix": _randn(g, S, HIDDEN, scale=2.0),
         "delta": _randn(g, S, HIDDEN, scale=0.5),
@@ -95,12 +101,15 @@ def _kda_inputs(nb: int, nacc: int, dim_first: bool, S: int, slots: int = 16):
         w["conv_state"] = conv.view(slots, QKV, sl)
     else:
         w["conv_state"] = conv.view(slots, sl, QKV).transpose(-1, -2)
-    w["st_idx"] = torch.arange(3, 3 + S, dtype=torch.int32, device="cuda").view(1, S)
-    w["num_acc"] = torch.full((1,), nacc, dtype=torch.int32, device="cuda")
+    w["st_idx"] = torch.arange(3, 3 + S, dtype=torch.int32, device="cuda").view(R, L)
+    # the requests' accepted counts differ, from nacc on
+    w["num_acc"] = torch.tensor(
+        [1 + (nacc - 1 + r) % L for r in range(R)], dtype=torch.int32, device="cuda"
+    )
     return w
 
 
-def _kda_reference(w, nb: int, S: int, write_idx: int):
+def _kda_reference(w, nb: int, S: int, L: int, write_idx: int):
     from vllm.config import VllmConfig, set_current_vllm_config
     from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
         causal_conv1d_update,
@@ -132,7 +141,7 @@ def _kda_reference(w, nb: int, S: int, write_idx: int):
     mixed_qkv, g2, f_a, beta = proj.split([QKV, PROJ, HD, NH + 4], dim=-1)[:4]
     g1 = F.linear(f_a, w["w_fb"])
     conv_state, rstate = w["conv_state"].clone(), w["rstate"].clone()
-    qsl = torch.tensor([0, S], dtype=torch.int32, device="cuda")
+    qsl = torch.arange(0, S + 1, L, dtype=torch.int32, device="cuda")
     conv = causal_conv1d_update(
         mixed_qkv.contiguous(),
         conv_state,
@@ -142,7 +151,7 @@ def _kda_reference(w, nb: int, S: int, write_idx: int):
         conv_state_indices=w["st_idx"][:, 0].contiguous(),
         num_accepted_tokens=w["num_acc"],
         query_start_loc=qsl,
-        max_query_len=S,
+        max_query_len=L,
         validate_data=False,
         out=torch.empty_like(mixed_qkv),
     )
@@ -163,7 +172,7 @@ def _kda_reference(w, nb: int, S: int, write_idx: int):
         cu_seqlens=qsl,
         ssm_state_indices=w["st_idx"],
         num_accepted_tokens=w["num_acc"],
-        uniform_sequence_length=S,
+        uniform_sequence_length=L,
         out=core,
     )
     with set_current_vllm_config(VllmConfig()):
@@ -179,14 +188,22 @@ def _kda_reference(w, nb: int, S: int, write_idx: int):
     }
 
 
+# (rows, tokens a request): the widths of concurrency 1 .. 16 at 1 + 3 or 1 + 7
+# speculative tokens
+K1_SHAPES = [(8, 8), (16, 8), (16, 4), (16, 16), (4, 4), (32, 8), (64, 4), (64, 16)]
+
+
 @needs_gfx950
+@pytest.mark.parametrize("S, L", K1_SHAPES)
 @pytest.mark.parametrize("dim_first", [True, False])
 @pytest.mark.parametrize("nacc", [1, 3])
 @pytest.mark.parametrize("write", [False, True])
-def test_k1_matches_vllm_ops(dim_first: bool, nacc: int, write: bool) -> None:
+def test_k1_matches_vllm_ops(
+    S: int, L: int, dim_first: bool, nacc: int, write: bool
+) -> None:
     """K1 against attn_res -> in_proj -> f_b -> causal_conv1d_update ->
-    fused_recurrent_kda -> FusedRMSNormGated on one spec-verify step (one
-    request, 8 tokens), including the state each step leaves behind."""
+    fused_recurrent_kda -> FusedRMSNormGated on one spec-verify step (S / L
+    requests of L tokens), including the state each step leaves behind."""
     from vllm.models.kimi_k3.amd.specialized_decode.mono.attention.kda import (
         PROJ,
         KdaPreBuild,
@@ -195,15 +212,15 @@ def test_k1_matches_vllm_ops(dim_first: bool, nacc: int, write: bool) -> None:
     )
     from vllm.models.kimi_k3.amd.specialized_decode.mono.common.plan import HEALTH_WORDS
 
-    S, nb = 8, 4
+    nb = 4
     # a block-write layer appends block nb (block_write_idx == prev_valid_blocks)
     write_idx = nb if write else -1
-    w = _kda_inputs(nb, nacc, dim_first, S)
-    ref = _kda_reference(w, nb, S, write_idx)
+    w = _kda_inputs(nb, nacc, dim_first, S, L)
+    ref = _kda_reference(w, nb, S, L, write_idx)
 
     key = KdaPreBuild(
         tokens=S,
-        qlen=S,
+        qlen=L,
         nblocks=nb,
         delta=True,
         state_len=w["conv_state"].size(-1),
@@ -254,6 +271,432 @@ def test_k1_matches_vllm_ops(dim_first: bool, nacc: int, write: bool) -> None:
     assert _cos(core, ref["core"]) > COS_MIN, errs
 
 
+# ---------------------------------------------------------------- M1a (MLA)
+
+
+def _rms(x, w, eps):
+    xf = x.float()
+    return (xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + eps) * w.float()).to(bf)
+
+
+def _mla_inputs(S: int, nb: int, rank: int = 0, ranks: int = 1):
+    from aiter.ops.triton.quant import dynamic_mxfp4_quant
+
+    from vllm.models.kimi_k3.amd.specialized_decode.mono.attention.kda import HIDDEN
+    from vllm.models.kimi_k3.amd.specialized_decode.mono.attention.mla import (
+        ENTRY,
+        GROWS,
+        KVL,
+        NH,
+        NOPE,
+        QB,
+        QKV_A,
+        QL,
+    )
+
+    g = torch.Generator(device="cuda").manual_seed(0)  # replicated
+    gr = torch.Generator(device="cuda").manual_seed(300 + rank)  # TP shards
+    w = {
+        "prefix": _randn(g, S, HIDDEN, scale=2.0),
+        "delta": _randn(g, S, HIDDEN, scale=0.5),
+        "blocks": _randn(g, S, 8, HIDDEN, scale=2.0),
+        "ares_nw": (1 + _randn(g, HIDDEN, scale=0.1, dtype=torch.float32)).to(bf),
+        "ares_qk": _randn(g, HIDDEN, scale=0.05),
+        "in_nw": (1 + _randn(g, HIDDEN, scale=0.1, dtype=torch.float32)).to(bf),
+        "w_qkv": _randn(g, QKV_A, HIDDEN, scale=0.02),
+        "w_g": _randn(gr, GROWS, HIDDEN, scale=0.02),
+        "q_nw": (1 + _randn(g, QL, scale=0.1, dtype=torch.float32)).to(bf),
+        "kv_nw": (1 + _randn(g, KVL, scale=0.1, dtype=torch.float32)).to(bf),
+        "w_qb": _randn(gr, QB, QL, scale=0.03),
+    }
+    # W_UK as vLLM's MLA keeps it: rows (n, h) quantized by aiter, the scales
+    # column-major: W_K (12, 512, 64) strides (64, 768, 1), its scales (1, 12, 6144)
+    uk = _randn(gr, KVL * NH, NOPE, scale=0.1)
+    wq, ws = dynamic_mxfp4_quant(uk)
+    w["w_k"] = wq.view(torch.uint8).unflatten(0, (KVL, NH)).transpose(0, 1)
+    w["w_k_scale"] = ws.view(torch.uint8).unflatten(0, (KVL, NH)).transpose(0, 1)
+    # hybrid page alignment: entries dense, blocks strided; a token another
+    # DCP rank's (slot -1): every other one, or t % ranks != rank
+    bs, nblk = 16, 2 * S // 16 + 2
+    # finite E4M3 bytes (0x7f is NaN)
+    pages = torch.randint(
+        0,
+        0x70,
+        (nblk, bs * ENTRY + 512),
+        dtype=torch.uint8,
+        device="cuda",
+        generator=gr,
+    )
+    w["kv_cache"] = pages[:, : bs * ENTRY].view(nblk, bs, ENTRY)
+    perm = torch.randperm(nblk * bs, device="cuda", generator=gr)[:S]
+    if ranks == 1:
+        perm[1::2] = -1
+    else:
+        perm[torch.arange(S, device="cuda") % ranks != rank] = -1
+    w["slots"] = perm.to(torch.int64)
+    return w
+
+
+def _mla_reference(w, nb: int):
+    from vllm._aiter_ops import rocm_aiter_ops
+    from vllm.models.kimi_k3.amd.ops.attn_res import attn_res
+    from vllm.models.kimi_k3.amd.specialized_decode.mono.attention.mla import (
+        KVL,
+        NH,
+        NOPE,
+        QL,
+        ROPE,
+    )
+    from vllm.platforms import current_platform
+
+    fp8 = current_platform.fp8_dtype()
+    prefix, blocks = w["prefix"].clone(), w["blocks"].clone()
+    x = attn_res(
+        prefix,
+        w["delta"],
+        blocks,
+        w["ares_nw"],
+        w["ares_qk"],
+        w["in_nw"],
+        nb,
+        -1,
+        1e-5,
+        1e-5,
+    )
+    qkv = F.linear(x, w["w_qkv"])
+    q_c, kv_c, k_pe = qkv.split([QL, KVL, ROPE], -1)
+    q_c, kv_c = _rms(q_c, w["q_nw"], 1e-6), _rms(kv_c, w["kv_nw"], 1e-6)
+    q = F.linear(q_c, w["w_qb"]).view(-1, NH, NOPE + ROPE)
+    q_nope, q_pe = q.split([NOPE, ROPE], -1)
+    one = torch.ones(1, dtype=torch.float32, device="cuda")
+    ql = rocm_aiter_ops.batched_gemm_a16wfp4(
+        q_nope.transpose(0, 1),
+        w["w_k"],
+        w["w_k_scale"],
+        transpose_bm=True,
+        prequant=True,
+        y_scale=one,
+    )
+    qb = torch.cat([ql, q_pe], -1)
+    q8 = qb.float().clamp(-448, 448).to(fp8)
+    cache = w["kv_cache"].clone()
+    entry = torch.cat([kv_c, k_pe], -1).float().clamp(-448, 448).to(fp8)
+    flat = cache.view(-1, cache.size(1), cache.size(2))
+    for t, s in enumerate(w["slots"].tolist()):
+        if s >= 0:
+            flat[s // cache.size(1), s % cache.size(1)] = entry[t].view(torch.uint8)
+    gate = F.linear(x, w["w_g"]).sigmoid()
+    return {
+        "prefix": prefix, "blocks": blocks, "q": q8, "qb": qb, "cache": cache,
+        "gate": gate,
+    }  # fmt: skip
+
+
+def _fp8_f(t):
+    from vllm.platforms import current_platform
+
+    return t.view(current_platform.fp8_dtype()).float()
+
+
+@needs_gfx950
+@pytest.mark.parametrize("qbf16", [False, True])
+@pytest.mark.parametrize("S", [4, 8, 16, 32, 64])
+def test_m1a_single_gpu_matches_vllm_ops(S: int, qbf16: bool) -> None:
+    """M1a at DCP1 (ranks = 1: no gather, no peer memory) against attn_res ->
+    qkv_a -> q / kv RMSNorm -> q_b -> aiter's fp4 W_UK -> the fp8 query (bf16
+    with ``qbf16``, M1c's), the fp8 cache rows and the gate."""
+    from vllm.models.kimi_k3.amd.specialized_decode.mono.attention.mla import (
+        ENTRY,
+        GROWS,
+        KVL,
+        NH,
+        MlaPreBuild,
+        mla_pre,
+        scratch_bytes,
+    )
+    from vllm.models.kimi_k3.amd.specialized_decode.mono.common.plan import HEALTH_WORDS
+
+    nb = 4
+    w = _mla_inputs(S, nb)
+    ref = _mla_reference(w, nb)
+    key = MlaPreBuild(tokens=S, nblocks=nb, delta=True, ranks=1, qbf16=qbf16)
+    prefix, blocks, cache = (
+        w["prefix"].clone(),
+        w["blocks"].clone(),
+        w["kv_cache"].clone(),
+    )
+    q = torch.zeros(S, NH, ENTRY, dtype=bf if qbf16 else torch.uint8, device="cuda")
+    gate = torch.zeros(S, GROWS, dtype=bf, device="cuda")
+    health = torch.zeros(HEALTH_WORDS, dtype=torch.int32, device="cuda")
+    flags = torch.zeros(1, dtype=torch.int32, device="cuda")
+    mla_pre(
+        key,
+        prefix=prefix,
+        delta=w["delta"],
+        blocks=blocks,
+        ares_nw=w["ares_nw"],
+        ares_qk=w["ares_qk"],
+        in_nw=w["in_nw"],
+        w_qkv=w["w_qkv"],
+        w_g=w["w_g"],
+        q_nw=w["q_nw"],
+        kv_nw=w["kv_nw"],
+        w_qb=w["w_qb"],
+        w_k=w["w_k"],
+        w_k_scale=w["w_k_scale"],
+        kv_cache=cache,
+        slots=w["slots"],
+        gate_out=gate,
+        q_out=q,
+        scratch=torch.zeros(scratch_bytes(key), dtype=torch.uint8, device="cuda"),
+        peers=None,
+        rank=0,
+        layer=3,
+        epoch=torch.ones(1, dtype=torch.int32, device="cuda"),
+        health=health,
+        alarm=flags.data_ptr(),
+    )
+    torch.accelerator.synchronize()
+    assert not health.any() and not flags.any()
+
+    torch.testing.assert_close(prefix, ref["prefix"], atol=0, rtol=0)
+    torch.testing.assert_close(blocks, ref["blocks"], atol=0, rtol=0)
+    # GEMV sums differ from hipBLASLt's / aiter's in order only: an fp8 or bf16
+    # value now and then lands one step off
+    if qbf16:
+        # bf16 values a bf16 step off now and then, the same elementwise
+        # rounding as the fp8 path's before its quant
+        rq = ref["qb"].float()
+        assert _rel_l2(q[..., :KVL].float(), rq[..., :KVL]) < 1e-2
+        assert _rel_l2(q[..., KVL:].float(), rq[..., KVL:]) < 1e-2
+        assert _cos(q.float(), rq) > 0.9999
+        q = ref["q"].view(torch.uint8)  # the fp8 checks below: vacuous
+    errs = {
+        "gate": _rel_l2(gate, ref["gate"]),
+        "q_nope": _rel_l2(_fp8_f(q[..., :KVL]), ref["q"][..., :KVL].float()),
+        "q_pe": _rel_l2(_fp8_f(q[..., KVL:]), ref["q"][..., KVL:].float()),
+        "cache": _rel_l2(_fp8_f(cache), _fp8_f(ref["cache"])),
+        "q_same": (q == ref["q"].view(torch.uint8)).float().mean().item(),
+        "cache_same": (cache == ref["cache"]).float().mean().item(),
+    }
+    assert errs["gate"] < 1e-2 and errs["cache"] < 2e-2, errs
+    assert errs["q_nope"] < 3e-2 and errs["q_pe"] < 3e-2, errs
+    assert _cos(_fp8_f(q), ref["q"].float()) > 0.999, errs
+    # rows of unstored tokens (slot -1) untouched
+    assert errs["cache_same"] > 0.99, errs
+
+
+def _a2a_combine(o, lse, seq_lens, qsl, q: int):
+    """``_dcp_a2a_pack_send`` + ``_dcp_a2a_unpack_combine`` for one rank's heads,
+    every rank's partials stacked: o [R, S, 12, 512], lse [R, S, 12], each
+    sender's local KV lengths seq_lens [R, requests]."""
+    S = o.size(1)
+    t = torch.arange(S, device=o.device)
+    empty = (seq_lens[:, t // q] == 0) | (t >= qsl[S // q])
+    lse = lse.masked_fill(empty.unsqueeze(-1), float("-inf"))
+    lse = torch.where(lse.isnan() | (lse == float("inf")), float("-inf"), lse)
+    m = lse.max(0).values
+    m = torch.where(m == float("-inf"), 0.0, m)
+    glse = (lse - m).exp().sum(0).log() + m
+    w = (lse - glse).exp().nan_to_num(0.0)
+    p = torch.where(w.unsqueeze(-1) == 0, 0.0, o.float())
+    return (p * w.unsqueeze(-1)).sum(0).to(bf)
+
+
+def _mla_back_inputs(S: int, q: int, R: int, rank: int = 0):
+    from aiter.ops.triton.quant import dynamic_mxfp4_quant
+
+    from vllm.models.kimi_k3.amd.specialized_decode.mono.attention.mla import (
+        GROWS,
+        KVL,
+        NH,
+    )
+
+    g = torch.Generator(device="cuda").manual_seed(1 + rank)
+    nreq = S // q
+    w = {
+        # AITER's padded 128-head outputs, viewed to the gathered heads
+        "o": _randn(g, S, 128, KVL, scale=0.5)[:, : R * NH],
+        "lse": (torch.randn(S, 128, generator=g, device="cuda") * 3)[:, : R * NH],
+        "gate": torch.rand(S, GROWS, generator=g, device="cuda").to(bf),
+        # the first request with no local KV (when there are others)
+        "seq_lens": torch.randint(
+            1, 900, (nreq,), dtype=torch.int32, device="cuda", generator=g
+        ),
+        "qsl": torch.arange(0, S + 1, q, dtype=torch.int32, device="cuda"),
+    }
+    # a request this rank holds none of, and with ranks, one no rank holds
+    if nreq > 1:
+        w["seq_lens"][rank % nreq] = 0
+    if R > 1 and nreq > 2:
+        w["seq_lens"][-1] = 0
+    uv = _randn(g, 128 * NH, KVL, scale=0.05)
+    wq, ws = dynamic_mxfp4_quant(uv)
+    # (12, 128, 256) dense rows a head; the scales column-major as vLLM keeps them
+    w["w_v"] = wq.view(torch.uint8).view(NH, 128, KVL // 2)
+    w["w_v_scale"] = ws.view(torch.uint8).unflatten(0, (NH, 128))
+    return w
+
+
+def _uv_gate(lat, w):
+    """The gate times vLLM's ``_v_up_proj`` (aiter's fp4 W_UV)."""
+    from vllm._aiter_ops import rocm_aiter_ops
+    from vllm.models.kimi_k3.amd.specialized_decode.mono.attention.mla import NH
+
+    S = lat.size(0)
+    out = torch.empty(S, NH, 128, dtype=bf, device=lat.device)
+    rocm_aiter_ops.batched_gemm_a16wfp4(
+        lat.transpose(0, 1),
+        w["w_v"],
+        w["w_v_scale"],
+        out,
+        transpose_bm=True,
+        prequant=True,
+        y_scale=None,
+    )
+    return out.view(S, -1) * w["gate"]
+
+
+def _mla_back_reference(w):
+    """DCP1: AITER's output is the latent (vLLM's ``_v_up_proj`` input)."""
+    return _uv_gate(w["o"], w)
+
+
+@needs_gfx950
+@pytest.mark.parametrize("S, q", [(4, 4), (8, 4), (16, 4), (32, 4), (64, 4), (64, 8)])
+def test_m1b_single_gpu_matches_vllm_ops(S: int, q: int) -> None:
+    """M1b at DCP1 (ranks = 1: no LSE, no peer memory): AITER's strided
+    output -> aiter's fp4 W_UV -> the gate."""
+    from vllm.models.kimi_k3.amd.specialized_decode.mono.attention.mla import GROWS
+    from vllm.models.kimi_k3.amd.specialized_decode.mono.attention.mla_back import (
+        MlaPostBuild,
+        mla_post,
+    )
+    from vllm.models.kimi_k3.amd.specialized_decode.mono.common.plan import HEALTH_WORDS
+
+    w = _mla_back_inputs(S, q, 1)
+    ref = _mla_back_reference(w)
+    health = torch.zeros(HEALTH_WORDS, dtype=torch.int32, device="cuda")
+    flags = torch.zeros(1, dtype=torch.int32, device="cuda")
+    core = torch.zeros(S, GROWS, dtype=bf, device="cuda")
+    mla_post(
+        MlaPostBuild(tokens=S, qlen=q, ranks=1),
+        o=w["o"],
+        lse=None,
+        seq_lens=w["seq_lens"],
+        qsl=w["qsl"],
+        w_v=w["w_v"],
+        w_v_scale=w["w_v_scale"],
+        gate=w["gate"],
+        core=core,
+        peers=None,
+        rank=0,
+        layer=5,
+        epoch=torch.ones(1, dtype=torch.int32, device="cuda"),
+        health=health,
+        alarm=flags.data_ptr(),
+    )
+    torch.accelerator.synchronize()
+    assert not health.any() and not flags.any()
+    err = _rel_l2(core, ref)
+    assert err < 1e-2 and _cos(core, ref) > COS_MIN, err
+
+
+def _mla_attn_reference(qb, kv8, idx, lens, q, scale):
+    """fp32 causal MLA decode a request: [S, 12, 512] (a request with no KV 0)."""
+    from vllm.models.kimi_k3.amd.specialized_decode.mono.attention.mla import KVL
+
+    out, p0 = [], 0
+    for r, n in enumerate(lens):
+        qq = qb[r * q : (r + 1) * q].float()
+        if n == 0:
+            out.append(qq.new_zeros(q, qq.size(1), KVL))
+            continue
+        kv = kv8[idx[p0 : p0 + n].long()].float()
+        p0 += n
+        sc = torch.einsum("thd,kd->thk", qq, kv) * scale
+        lim = n - q + torch.arange(q, device=qq.device)
+        pos = torch.arange(n, device=qq.device)
+        sc = sc.masked_fill(pos[None, None] > lim[:, None, None], -torch.inf)
+        out.append(torch.einsum("thk,kd->thd", sc.softmax(-1), kv[:, :KVL]))
+    return torch.cat(out)
+
+
+@needs_gfx950
+@pytest.mark.parametrize(
+    "q, lens",
+    [(8, [8704]), (8, [3000, 0]), (4, [5000, 17, 9000, 64]), (8, [9]), (8, [40000])],
+)
+def test_m1c_m1b_splits_single_gpu(q: int, lens: list) -> None:
+    """M1c (``mla_attn``: bf16 query, fp8 KV, the splits' partials) against
+    fp32 attention, and M1b's ``splits`` combine -> W_UV -> the gate against
+    the partials merged in torch."""
+    from vllm.models.kimi_k3.amd.specialized_decode.mono.attention import mla_attn as m
+    from vllm.models.kimi_k3.amd.specialized_decode.mono.attention.mla import (
+        ENTRY,
+        GROWS,
+        NH,
+    )
+    from vllm.models.kimi_k3.amd.specialized_decode.mono.attention.mla_back import (
+        MlaPostBuild,
+        mla_post,
+    )
+    from vllm.models.kimi_k3.amd.specialized_decode.mono.common.plan import HEALTH_WORDS
+
+    S = q * len(lens)
+    g = torch.Generator(device="cuda").manual_seed(7)
+    ntok = 200_000
+    kv8 = (torch.randn(ntok, ENTRY, generator=g, device="cuda") * 0.5).to(
+        torch.float8_e4m3fn
+    )
+    idx = torch.randperm(ntok, generator=g, device="cuda")[: sum(lens)].int()
+    indptr = torch.tensor([0, *torch.tensor(lens).cumsum(0).tolist()],
+                          dtype=torch.int32, device="cuda")  # fmt: skip
+    qb = _randn(g, S, NH, ENTRY, scale=0.5)
+    scale, kv_scale = 1.0 / 192**0.5, 0.75
+    key = m.MlaAttnBuild(tokens=S, qlen=q)
+    so, sl = m.partial_shapes(key)
+    po = torch.empty(so, dtype=bf, device="cuda")
+    pl = torch.empty(sl, dtype=torch.float32, device="cuda")
+    m.mla_attn(key, q=qb, kv=kv8, pages=idx, indptr=indptr, part_o=po, part_lse=pl,
+               scale=scale, kv_scale=kv_scale)  # fmt: skip
+    att = _mla_attn_reference(qb, kv8, idx, lens, q, scale * kv_scale) * kv_scale
+    # the partials merged in torch: a request's splits only (the rest unwritten;
+    # an empty request's lse is -inf: 0 weight)
+    per = so[0] // len(lens)
+    lat = []
+    for r in range(len(lens)):
+        lo, ll = (
+            po[r * per : (r + 1) * per, r * q : (r + 1) * q],
+            pl[r * per : (r + 1) * per, r * q : (r + 1) * q],
+        )
+        e = (ll - ll.amax(0).nan_to_num(neginf=0.0)).exp()
+        lat.append(
+            (lo.float() * e[..., None]).sum(0) / e.sum(0).clamp_min(1e-30)[..., None]
+        )
+    lat = torch.cat(lat)
+    assert _rel_l2(lat, att) < 5e-3
+    # W_UV quantizes the latent to mxfp4: compare at the same latent
+    w = _mla_back_inputs(S, q, 1)
+    ref = _uv_gate(lat.to(bf), w)
+    health = torch.zeros(HEALTH_WORDS, dtype=torch.int32, device="cuda")
+    flags = torch.zeros(1, dtype=torch.int32, device="cuda")
+    core = torch.zeros(S, GROWS, dtype=bf, device="cuda")
+    mla_post(
+        MlaPostBuild(tokens=S, qlen=q, ranks=1, splits=m.splits(q)),
+        o=po[0], lse=pl[0], seq_lens=indptr.diff().int(), qsl=w["qsl"],
+        w_v=w["w_v"], w_v_scale=w["w_v_scale"], gate=w["gate"], core=core,
+        peers=None, rank=0, layer=5,
+        epoch=torch.ones(1, dtype=torch.int32, device="cuda"),
+        health=health, alarm=flags.data_ptr(),
+    )  # fmt: skip
+    torch.accelerator.synchronize()
+    assert not health.any() and not flags.any()
+    assert torch.isfinite(core.float()).all()
+    err = _rel_l2(core, ref)
+    assert err < 1e-2 and _cos(core, ref) > COS_MIN, err
+
+
 # ---------------------------------------------------------------- K2 / MoE
 
 
@@ -378,29 +821,63 @@ def _moe_args(mw):
     }
 
 
-def _k2_inputs(rank: int, device):
+def _k2_inputs(rank: int, device, S: int = 8, seed: int = 0):
     from vllm.models.kimi_k3.amd.specialized_decode.mono.attention.kda import (
         HIDDEN,
         PROJ,
     )
 
-    gc = torch.Generator(device=device).manual_seed(11)
-    gr = torch.Generator(device=device).manual_seed(200 + rank)
+    gc = torch.Generator(device=device).manual_seed(11 + 1000 * seed)
+    gr = torch.Generator(device=device).manual_seed(200 + rank + 1000 * seed)
     return {
-        "core": _randn(gr, 8, PROJ),
+        "core": _randn(gr, S, PROJ),
         "w_o": _randn(gr, HIDDEN, PROJ, scale=0.02),
-        "prefix": _randn(gc, 8, HIDDEN, scale=2.0),
-        "blocks": _randn(gc, 8, 8, HIDDEN, scale=2.0),
+        "prefix": _randn(gc, S, HIDDEN, scale=2.0),
+        "blocks": _randn(gc, S, 8, HIDDEN, scale=2.0),
         "ares_nw": (1 + _randn(gc, HIDDEN, scale=0.1, dtype=torch.float32)).to(bf),
         "ares_qk": _randn(gc, HIDDEN, scale=0.05),
         "in_nw": (1 + _randn(gc, HIDDEN, scale=0.1, dtype=torch.float32)).to(bf),
     }
 
 
-def _k2_launch(runner, w, mw, prefix, out, reset: bool, nb: int = 4):
+def _k2_launch(
+    runner,
+    w,
+    mw,
+    prefix,
+    out,
+    reset: bool,
+    nb: int = 4,
+    layer: int = 3,
+    m1b=None,
+    splits: int = 0,
+    m1c_in=None,
+    opt: int | None = None,
+):
+    """``m1b``: (``mla_post``'s inputs, qlen, ranks) to run M1b in K2's launch;
+    ``splits`` its M1c splits, ``m1c_in`` (``mla_attn.host_args``'s) to run M1c
+    in the launch too. ``opt``: ``K2Build.opt`` (None: the default)."""
     from vllm.models.kimi_k3.amd.specialized_decode.mono import layer as k2
 
-    key = k2.K2Build(tokens=8, nblocks=nb + int(reset), reset=reset, **runner.fault)
+    m1b_in, fused = None, {}
+    if m1b is not None:
+        m1b_in, q, ranks = m1b
+        fused = dict(
+            m1b_qlen=q,
+            m1b_ranks=ranks,
+            m1b_off=runner.mla_post_off,
+            m1b_splits=splits,
+            m1c=m1c_in is not None,
+        )
+    if opt is not None:
+        fused["opt"] = opt
+    key = k2.K2Build(
+        tokens=w["core"].size(0),
+        nblocks=nb + int(reset),
+        reset=reset,
+        **fused,
+        **runner.fault,
+    )
     k2.k2_launch(
         key,
         core=w["core"],
@@ -411,7 +888,9 @@ def _k2_launch(runner, w, mw, prefix, out, reset: bool, nb: int = 4):
         ares_qk=w["ares_qk"],
         in_nw=w["in_nw"],
         out=out,
-        layer=3,
+        layer=128 + layer,
+        m1b_in=m1b_in,
+        m1c_in=m1c_in,
         **_moe_args(mw),
         **runner.launch_args(),
         **runner.peer_args(),
@@ -442,12 +921,11 @@ def _k2_worker(
     from vllm.models.kimi_k3.amd.specialized_decode.mono.attention.kda import HIDDEN
     from vllm.models.kimi_k3.amd.specialized_decode.runtime.runner import MonoRunner
 
-    S, nb = 8, 4
+    nb = 4
     runner = MonoRunner(device)
     mw = _moe_weights(rank, device)
-    w = _k2_inputs(rank, device)
-
-    for reset in (False, True):
+    for S, reset in K2_CASES:
+        w = _k2_inputs(rank, device, S)
         runner.step_begin()
         prefix = w["prefix"].clone()
         out = torch.empty(S, HIDDEN, dtype=bf, device=device)
@@ -474,14 +952,18 @@ def _k2_worker(
         assert _rel_l2(prefix, ref_prefix) < 1e-2
         lay = k2.scratch_layout(key)["xrow"][0]
         x = runner.scratch[lay : lay + S * HIDDEN * 2].view(bf).view(S, HIDDEN)
-        assert _cos(x, x_ref) > COS_MIN
-        assert _cos(out, _moe_reference(x, mw, rank)) > COS_MIN
+        cos_x = _cos(x, x_ref)
+        assert cos_x > COS_MIN, (S, reset, cos_x)
+        cos_out = _cos(out, _moe_reference(x, mw, rank))
+        assert cos_out > COS_MIN, (S, reset, cos_out)
         outs = [torch.empty_like(out) for _ in range(tp_size)]
         dist.all_gather(outs, out)
         assert all(torch.equal(outs[0], o) for o in outs), "ranks disagree"
 
     # A graph of one step (the epoch bump captured with it), replayed: a
     # block-write launch writes its prefix, so every replay must match the first.
+    S = K2_CASES[-1][0]
+    w = _k2_inputs(rank, device, S)
     prefix = torch.empty_like(w["prefix"])
     out = torch.empty(S, HIDDEN, dtype=bf, device=device)
     graph = torch.cuda.CUDAGraph()
@@ -498,12 +980,461 @@ def _k2_worker(
     assert not runner.health.any() and not runner.flags.any()
 
 
+# (rows, block-write layer): every width up to the runner's MAX_TOKENS
+K2_CASES = [
+    (8, False),
+    (8, True),
+    (4, False),
+    (16, False),
+    (16, True),
+    (32, False),
+    (64, False),
+    (64, True),
+]
+
+
 @needs_gfx950
 @multi_gpu_test(num_gpus=8)
 def test_k2_tp8_matches_vllm_ops_and_replays(monkeypatch: pytest.MonkeyPatch) -> None:
     multi_process_parallel(monkeypatch, 8, 1, _k2_worker)
 
 
+@ray.remote(num_gpus=1, max_calls=1)
+def _k2_opt_worker(
+    monkeypatch: pytest.MonkeyPatch,
+    tp_size: int,
+    pp_size: int,
+    rank: int,
+    distributed_init_port: str,
+) -> None:
+    device = _init_worker(monkeypatch, tp_size, pp_size, rank, distributed_init_port)
+
+    from vllm.models.kimi_k3.amd.specialized_decode.mono import layer as k2
+    from vllm.models.kimi_k3.amd.specialized_decode.mono.attention.kda import HIDDEN
+    from vllm.models.kimi_k3.amd.specialized_decode.runtime.runner import MonoRunner
+
+    runner = MonoRunner(device)
+    mw = _moe_weights(rank, device)
+
+    def run(w, S, reset, layer, opt):
+        runner.step_begin()
+        prefix = w["prefix"].clone()
+        out = torch.empty(S, HIDDEN, dtype=bf, device=device)
+        _k2_launch(runner, w, mw, prefix, out, reset, layer=layer, opt=opt)
+        torch.accelerator.synchronize()
+        return out, prefix
+
+    for S, reset in K2_CASES:
+        # two input sets with different words: a sentinel region the first
+        # K2_OPT launch left un-re-armed would hand the second launch the
+        # first one's words, which no longer match opt=0's bits
+        ws = [_k2_inputs(rank, device, S, seed=i) for i in range(2)]
+        # both layer parities: the sentinel regions alternate by it
+        for layer in (3, 4):
+            want = [run(w, S, reset, layer, 0) for w in ws]
+            assert not torch.equal(want[0][0], want[1][0]), (S, reset, layer)
+            # back to back on the same regions: the second polls re-armed ones
+            for i, w in enumerate(ws):
+                out, prefix = run(w, S, reset, layer, k2.K2_OPT)
+                assert torch.equal(out, want[i][0]), (S, reset, layer, i)
+                assert torch.equal(prefix, want[i][1]), (S, reset, layer, i)
+    assert not runner.health.any() and not runner.flags.any()
+
+
+@needs_gfx950
+@multi_gpu_test(num_gpus=8)
+def test_k2_tp8_opt_bits_are_bit_exact(monkeypatch: pytest.MonkeyPatch) -> None:
+    """K2's critical-path options (``layer.K2_OPT``, the default) give the
+    plain K2's (``opt=0``) bits at every width, block-write layer or not, on
+    both layer parities, and again after their regions were re-armed."""
+    multi_process_parallel(monkeypatch, 8, 1, _k2_opt_worker)
+
+
+def _gather(t: torch.Tensor) -> torch.Tensor:
+    """Every rank's ``t`` stacked [ranks, ...]."""
+    out = [torch.empty_like(t) for _ in range(dist.get_world_size())]
+    dist.all_gather(out, t.contiguous())
+    return torch.stack(out)
+
+
+@ray.remote(num_gpus=1, max_calls=1)
+def _m1_worker(
+    monkeypatch: pytest.MonkeyPatch,
+    tp_size: int,
+    pp_size: int,
+    rank: int,
+    distributed_init_port: str,
+) -> None:
+    device = _init_worker(monkeypatch, tp_size, pp_size, rank, distributed_init_port)
+
+    from vllm.models.kimi_k3.amd.specialized_decode.mono.attention.kda import HIDDEN
+    from vllm.models.kimi_k3.amd.specialized_decode.mono.attention.mla import (
+        ENTRY,
+        GROWS,
+        NH,
+        MlaPreBuild,
+        mla_pre,
+    )
+    from vllm.models.kimi_k3.amd.specialized_decode.mono.attention.mla_back import (
+        MlaPostBuild,
+        mla_post,
+    )
+    from vllm.models.kimi_k3.amd.specialized_decode.runtime.runner import MonoRunner
+
+    R, nb = tp_size, 4
+    # DCP8 over the TP group; the test's distributed init sets up no DCP group
+    runner = MonoRunner(device, dcp=R)
+
+    def pre(w, prefix, blocks, cache, q8, gate, layer=3):
+        S = prefix.size(0)
+        key = MlaPreBuild(
+            tokens=S,
+            nblocks=nb,
+            delta=True,
+            ranks=R,
+            peer_off=runner.mla_peer_off,
+            qheads=q8.size(1) if q8.size(1) != R * NH else 0,
+        )
+        mla_pre(
+            key,
+            prefix=prefix,
+            delta=w["delta"],
+            blocks=blocks,
+            ares_nw=w["ares_nw"],
+            ares_qk=w["ares_qk"],
+            in_nw=w["in_nw"],
+            w_qkv=w["w_qkv"],
+            w_g=w["w_g"],
+            q_nw=w["q_nw"],
+            kv_nw=w["kv_nw"],
+            w_qb=w["w_qb"],
+            w_k=w["w_k"],
+            w_k_scale=w["w_k_scale"],
+            kv_cache=cache,
+            slots=w["slots"],
+            gate_out=gate,
+            q_out=q8,
+            layer=layer,
+            **runner.launch_args(),
+            **runner.peer_args(),
+        )
+
+    def post(b, q, gate, core, layer=3, ranks=R):
+        key = MlaPostBuild(
+            tokens=core.size(0), qlen=q, ranks=ranks, peer_off=runner.mla_post_off
+        )
+        mla_post(
+            key,
+            o=b["o"],
+            lse=b["lse"],
+            seq_lens=b["seq_lens"],
+            qsl=b["qsl"],
+            w_v=b["w_v"],
+            w_v_scale=b["w_v_scale"],
+            gate=gate,
+            core=core,
+            layer=layer,
+            **runner.launch_args(),
+            **runner.peer_args(),
+        )
+
+    from vllm.v1.attention.backends.mla import rocm_aiter_mla as am
+    from vllm.v1.attention.backends.mla.rocm_aiter_mla import AiterMLAHelper
+
+    # the head count cprr runs the gathered query at (128 for 96)
+    qheads = am._asm_dcp_verify_heads(R * NH)
+    assert qheads >= R * NH
+
+    # M1a: this rank's query rows land at heads 12 * rank.. on every rank
+    for S in (4, 8, 16, 32, 64):
+        w = _mla_inputs(S, nb, rank, R)
+        ref = _mla_reference(w, nb)
+        runner.step_begin()
+        prefix, blocks = w["prefix"].clone(), w["blocks"].clone()
+        cache = w["kv_cache"].clone()
+        q8 = torch.zeros(S, R * NH, ENTRY, dtype=torch.uint8, device=device)
+        gate = torch.zeros(S, GROWS, dtype=bf, device=device)
+        pre(w, prefix, blocks, cache, q8, gate)
+        torch.accelerator.synchronize()
+        assert not runner.health.any() and not runner.flags.any(), S
+        torch.testing.assert_close(prefix, ref["prefix"], atol=0, rtol=0)
+        want = _gather(ref["q"].view(torch.uint8)).transpose(0, 1).reshape_as(q8)
+        errs = {
+            "q": _rel_l2(_fp8_f(q8), _fp8_f(want)),
+            "gate": _rel_l2(gate, ref["gate"]),
+            "cache_same": (cache == ref["cache"]).float().mean().item(),
+        }
+        assert errs["q"] < 3e-2 and errs["gate"] < 1e-2, (S, errs)
+        assert _cos(_fp8_f(q8), _fp8_f(want)) > 0.999, (S, errs)
+        assert errs["cache_same"] > 0.99, (S, errs)
+        # AITER cprr's padded query (96 -> 128 heads) equals the pad
+        # forward_mqa builds itself (from the contiguous q8: a view of qp
+        # would hand qp back). test_m1a_pad_through_aiter_cprr runs it.
+        if qheads == R * NH:
+            continue
+        runner.step_begin()
+        prefix, blocks = w["prefix"].clone(), w["blocks"].clone()
+        qp = torch.zeros(S, qheads, ENTRY, dtype=torch.uint8, device=device)
+        pre(w, prefix, blocks, w["kv_cache"].clone(), qp, gate)
+        torch.accelerator.synchronize()
+        assert not runner.health.any() and not runner.flags.any(), S
+        pad = AiterMLAHelper.get_mla_padded_q(R * NH, q8, target_heads=qheads)
+        assert pad.data_ptr() != qp.data_ptr()
+        assert torch.equal(qp, pad), S
+
+    # M1b: every sender's partials for this rank's heads, each with its own
+    # empty-shard mask, combined in rank order
+    for S, q in ((4, 4), (16, 4), (64, 4), (64, 8)):
+        b = _mla_back_inputs(S, q, R, rank)
+        runner.step_begin()
+        core = torch.zeros(S, GROWS, dtype=bf, device=device)
+        post(b, q, b["gate"], core)
+        torch.accelerator.synchronize()
+        assert not runner.health.any() and not runner.flags.any(), (S, q)
+        heads = slice(NH * rank, NH * (rank + 1))
+        lat = _a2a_combine(
+            _gather(b["o"])[:, :, heads],
+            _gather(b["lse"])[:, :, heads],
+            _gather(b["seq_lens"]),
+            b["qsl"],
+            q,
+        )
+        ref = _uv_gate(lat, b)
+        if S // q > 2:  # the last request no rank holds
+            assert not core[-q:].any(), (S, q)
+        err = _rel_l2(core, ref)
+        assert err < 1e-2 and _cos(core, ref) > COS_MIN, (S, q, err)
+
+    def m1b_in(b, gate):
+        keys = ("o", "lse", "seq_lens", "qsl", "w_v", "w_v_scale")
+        return dict({n: b[n] for n in keys}, gate=gate)
+
+    # K2 with M1b in its launch (``K2Build.m1b_qlen``): core, prefix and the
+    # MoE output bit-identical to M1b's launch then K2's, at DCP8 and DCP1
+    mw = _moe_weights(rank, device)
+    for S, q in ((8, 8), (16, 8), (16, 4), (64, 4)):
+        k = _k2_inputs(rank, device, S)
+        for ranks in (R, 1):
+            b = _mla_back_inputs(S, q, ranks, rank)
+            if ranks == 1:
+                b["lse"] = None
+            got = []
+            for fused in (False, True):
+                runner.step_begin()
+                core = torch.zeros(S, GROWS, dtype=bf, device=device)
+                prefix = k["prefix"].clone()
+                out = torch.empty(S, HIDDEN, dtype=bf, device=device)
+                if not fused:
+                    post(b, q, b["gate"], core, ranks=ranks)
+                f = (m1b_in(b, b["gate"]), q, ranks) if fused else None
+                _k2_launch(runner, dict(k, core=core), mw, prefix, out, False, m1b=f)
+                torch.accelerator.synchronize()
+                assert not runner.health.any() and not runner.flags.any(), (S, q)
+                got.append((core, prefix, out))
+            for x, y in zip(*got):
+                assert torch.equal(x, y), (S, q, ranks)
+
+    # DCP1, M1c's partials (``mla_attn``) merged by K2's M1b (``splits``), then
+    # M1c in K2's launch too (``K2Build.m1c``): bit-identical, in a graph of two
+    # layers replayed
+    from vllm.models.kimi_k3.amd.specialized_decode.mono.attention import mla_attn
+
+    gk = torch.Generator(device=device).manual_seed(11)
+    kv8 = (torch.randn(100_000, ENTRY, generator=gk, device=device) * 0.5).to(
+        torch.float8_e4m3fn
+    )
+    for S, q, lens in (
+        (8, 8, [8704]),
+        (16, 8, [3000, 0]),
+        (16, 4, [900, 17, 5000, 64]),
+    ):
+        k = _k2_inputs(rank, device, S)
+        b = _mla_back_inputs(S, q, 1, rank)
+        idx = torch.randperm(kv8.size(0), generator=gk, device=device)[: sum(lens)]
+        indptr = torch.tensor([0, *torch.tensor(lens).cumsum(0).tolist()],
+                              dtype=torch.int32, device=device)  # fmt: skip
+        b["seq_lens"] = indptr.diff().int()
+        qb = _randn(gk, S, NH, ENTRY, scale=0.5)
+        c_in = dict(q=qb, kv=kv8, pages=idx.int(), indptr=indptr,
+                    scale=1 / 192**0.5, kv_scale=0.75)  # fmt: skip
+        key = mla_attn.MlaAttnBuild(tokens=S, qlen=q)
+        so, sl = mla_attn.partial_shapes(key)
+        n = mla_attn.splits(q)
+        po = torch.empty(so, dtype=bf, device=device)
+        pl = torch.empty(sl, dtype=torch.float32, device=device)
+        b.update(o=po[0], lse=pl[0])
+        core = torch.zeros(S, GROWS, dtype=bf, device=device)
+        prefix = torch.empty_like(k["prefix"])
+        out = torch.empty(S, HIDDEN, dtype=bf, device=device)
+        k2w = dict(k, core=core)
+        got = []
+        for fused in (False, True):
+            graph = torch.cuda.CUDAGraph()
+            with torch.cuda.graph(graph):
+                runner.step_begin()
+                prefix.copy_(k["prefix"])
+                for layer in (3, 5):
+                    if not fused:
+                        mla_attn.mla_attn(key, part_o=po, part_lse=pl, **c_in)
+                    _k2_launch(
+                        runner, k2w, mw, prefix, out, False, layer=layer,
+                        m1b=(m1b_in(b, b["gate"]), q, 1), splits=n,
+                        m1c_in=c_in if fused else None,
+                    )  # fmt: skip
+            graph.replay()
+            torch.accelerator.synchronize()
+            want = [t.clone() for t in (core, prefix, out)]
+            for _ in range(50):
+                graph.replay()
+            torch.accelerator.synchronize()
+            assert not runner.health.any() and not runner.flags.any(), (S, q, fused)
+            for x, y in zip((core, prefix, out), want):
+                assert torch.equal(x, y), (S, q, fused)
+            got.append(want)
+        for x, y in zip(*got):
+            assert torch.equal(x, y), (S, q)
+
+    # A graph of one step, two MLA layers of M1a -> M1b -> K2 (each layer its
+    # own tags, the peer regions reused), replayed bit-identical; then the
+    # same with M1b in K2's launch, equal to it
+    S, q = 64, 4
+    w = _mla_inputs(S, nb, rank, R)
+    b = _mla_back_inputs(S, q, R, rank)
+    k = _k2_inputs(rank, device, S)
+    prefix, blocks = torch.empty_like(w["prefix"]), torch.empty_like(w["blocks"])
+    cache = w["kv_cache"].clone()
+    q8 = torch.zeros(S, R * NH, ENTRY, dtype=torch.uint8, device=device)
+    gate = torch.zeros(S, GROWS, dtype=bf, device=device)
+    core = torch.zeros(S, GROWS, dtype=bf, device=device)
+    out = torch.empty(S, HIDDEN, dtype=bf, device=device)
+    wants = []
+    for fused in (False, True):
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            runner.step_begin()
+            prefix.copy_(w["prefix"])
+            blocks.copy_(w["blocks"])
+            for layer in (3, 5):
+                pre(w, prefix, blocks, cache, q8, gate, layer)
+                if not fused:
+                    post(b, q, gate, core, layer)
+                f = (m1b_in(b, gate), q, R) if fused else None
+                k2w = dict(k, core=core, blocks=blocks)
+                _k2_launch(runner, k2w, mw, prefix, out, False, layer=layer, m1b=f)
+        graph.replay()
+        torch.accelerator.synchronize()
+        want = [t.clone() for t in (prefix, q8, core, out)]
+        for _ in range(300):
+            graph.replay()
+        torch.accelerator.synchronize()
+        for got, ref in zip((prefix, q8, core, out), want):
+            assert torch.equal(got, ref), fused
+        assert not runner.health.any() and not runner.flags.any()
+        wants.append(want)
+    for x, y in zip(*wants):
+        assert torch.equal(x, y)
+
+
+@needs_gfx950
+@pytest.mark.parametrize("qlen", [4, 8])
+@torch.inference_mode()
+def test_m1a_pad_through_aiter_cprr(monkeypatch: pytest.MonkeyPatch, qlen: int) -> None:
+    """AITER's cprr asm decode, at DCP8 (96 -> 128 heads), on the padded query
+    M1a writes (mla.attend hands it over) equals the decode on the pad
+    forward_mqa builds itself, and both match attention in torch. With
+    _m1_worker's M1a == get_mla_padded_q, an AITER change the hand-off misses
+    fails here. One GPU, a DCP rank simulated (the cprr numerics test's
+    harness: real metadata builder, real kernel)."""
+    from types import SimpleNamespace
+
+    import vllm.distributed.parallel_state as ps
+    import vllm.model_executor.layers.attention.mla_attention as mla_attention
+    from tests.v1.attention import test_rocm_aiter_mla_dcp_cprr_numerics as cp
+    from vllm.config import set_current_vllm_config
+    from vllm.models.kimi_k3.amd.specialized_decode.mono.attention.mla import NH
+    from vllm.models.kimi_k3.amd.specialized_decode.runtime import mla
+    from vllm.v1.attention.backends.mla import rocm_aiter_mla as am
+    from vllm.v1.kv_cache_interface import MLAAttentionSpec
+
+    dcp, rank, nr, ctx = 8, 3, 3, 700
+    n = dcp * NH
+    H = am._asm_dcp_verify_heads(n)
+    if n == H:
+        pytest.skip(f"cprr runs {n} heads natively: M1a writes no pad")
+    # _install_fake_groups patches these; monkeypatch restores them
+    for mod, name in (
+        (ps, "get_dcp_group"),
+        (ps, "get_tp_group"),
+        (am, "get_dcp_group"),
+        (mla_attention, "get_dcp_group"),
+    ):
+        monkeypatch.setattr(mod, name, getattr(mod, name))
+    monkeypatch.setattr(cp, "QLEN", qlen)
+
+    device = torch.device("cuda", 0)
+    torch.accelerator.set_device_index(device.index)
+    g = torch.Generator(device=device).manual_seed(0)
+    cp._ensure_workspace_manager(device)
+    cp._install_fake_groups(dcp, rank)
+    config = cp._make_vllm_config(NH)
+    with set_current_vllm_config(config):
+        name = "model.layers.0.self_attn.attn"
+        config.compilation_config.static_forward_context[name] = cp._make_ctx_layer(
+            config, monkeypatch
+        )
+        spec = MLAAttentionSpec(
+            block_size=1,
+            num_kv_heads=1,
+            head_size=cp.HEAD_SIZE,
+            dtype=config.model_config.dtype,
+            non_causal_multi_token_decode=False,
+        )
+        L = ctx + qlen
+        fp8 = current_platform.fp8_dtype()
+        kv = _randn(g, nr * L, 1, cp.HEAD_SIZE, dtype=torch.float32).to(fp8)
+        q8 = _randn(g, nr * qlen, n, cp.HEAD_SIZE, scale=0.5, dtype=torch.float32)
+        q8 = q8.to(fp8)
+        pos = torch.arange(rank, L, dcp, device=device)  # rank's residue class
+        # forward_mqa pads q8 itself
+        md, o, lse, _ = cp._run_shard(
+            config, spec, name, device, rank, rank, q8, kv, nr, qlen, ctx, dcp, pos
+        )
+        assert md.decode.dcp_route is am._DCPDecodeRoute.CPRR
+        assert md.decode.asm_decode_num_heads == H
+        # M1a's buffer (_m1_worker: M1a == get_mla_padded_q of its q8)
+        qp = torch.empty(nr * qlen, H, cp.HEAD_SIZE, dtype=fp8, device=device)
+        qp.copy_(am.AiterMLAHelper.get_mla_padded_q(n, q8, target_heads=H))
+        layer = cp._Layer(device)
+        layer.impl = cp._build_impl(NH, dcp, rank)  # type: ignore[attr-defined]
+        o_m1a, lse_m1a = mla.attend(SimpleNamespace(mla_attn=layer), qp, kv, md)
+        torch.accelerator.synchronize()
+
+    assert torch.equal(o_m1a, o) and torch.equal(lse_m1a, lse)
+    want, want_lse = cp._torch_reference(
+        q8.float().view(nr, qlen, n, -1),
+        kv.float().view(nr, L, -1),
+        ctx,
+        qlen,
+        cp.SCALE,
+        pos,
+    )
+    o = o.float().reshape(want.shape)
+    err = (o - want).abs().max().item() / want.abs().max().item()
+    assert err < cp.SHARD_RTOL, err
+    assert _rel_l2(o, want) < 3e-2
+    torch.testing.assert_close(
+        lse.float().reshape(want_lse.shape), want_lse, atol=5e-2, rtol=0
+    )
+
+
+@needs_gfx950
+@multi_gpu_test(num_gpus=8)
+def test_m1_tp8_dcp8_matches_vllm_ops_and_replays(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    multi_process_parallel(monkeypatch, 8, 1, _m1_worker)
+
+
 STALLED = 3
 
 
@@ -530,7 +1461,7 @@ def _fault_worker(
     runner.fault = {
         "fault_rank": STALLED,
         "fault_epoch": 2,
-        "fault_layer": 3,
+        "fault_layer": 128 + 3,  # K2's slot of layer 3 (_k2_launch's default)
         "fault_ticks": WAIT_TICKS * 3 // 2,
     }
     mw = _moe_weights(rank, device)
@@ -562,14 +1493,25 @@ def test_stalled_rank_fails_the_step_without_a_hang(
 
 @needs_gfx950
 def test_warmup_compiles_without_launching() -> None:
-    """The JIT warm-up compiles K1 and K2 before capture without running them:
+    """The JIT warm-up compiles K1, K2, M1a and M1b before capture without running them:
     a launch on its null arguments would fault."""
     from vllm.models.kimi_k3.amd.specialized_decode.mono.attention.kda import (
         KdaPreBuild,
     )
+    from vllm.models.kimi_k3.amd.specialized_decode.mono.attention.mla import (
+        MlaPreBuild,
+    )
+    from vllm.models.kimi_k3.amd.specialized_decode.mono.attention.mla_back import (
+        MlaPostBuild,
+    )
     from vllm.models.kimi_k3.amd.specialized_decode.mono.layer import K2Build
     from vllm.models.kimi_k3.amd.specialized_decode.runtime.layers import compile_build
 
-    compile_build(KdaPreBuild(tokens=8, qlen=8, nblocks=4, delta=True, state_len=10))
-    compile_build(K2Build(tokens=8, nblocks=4))
+    for s, q in ((8, 8), (16, 4), (64, 16)):
+        compile_build(
+            KdaPreBuild(tokens=s, qlen=q, nblocks=4, delta=True, state_len=3 + q - 1)
+        )
+        compile_build(K2Build(tokens=s, nblocks=4))
+        compile_build(MlaPreBuild(tokens=s, nblocks=4, delta=True))
+        compile_build(MlaPostBuild(tokens=s, qlen=q))
     torch.accelerator.synchronize()
diff --git a/tests/models/kimi_k3/amd_specialized_decode/test_runtime.py b/tests/models/kimi_k3/amd_specialized_decode/test_runtime.py
index 266b947dbe01..ef4084975a50 100644
--- a/tests/models/kimi_k3/amd_specialized_decode/test_runtime.py
+++ b/tests/models/kimi_k3/amd_specialized_decode/test_runtime.py
@@ -32,12 +32,14 @@
 )
 from vllm.models.kimi_k3.amd.specialized_decode.runtime.guards import (
     cache_reasons,
+    mla_weight_reasons,
     weight_reasons,
 )
 
 
 def _vllm_config(**over) -> SimpleNamespace:
-    """A supported config: TP8, DCP8, 7 speculative tokens, FULL graphs, MRV2.
+    """A supported config: TP8, DCP8 (interleave 1), 7 speculative tokens,
+    max_num_seqs 1, an fp8 KV cache, FULL graphs, MRV2.
     ``over`` sets dotted paths, e.g. ``{"parallel_config.data_parallel_size": 2}``."""
     spec = SimpleNamespace(
         num_speculative_tokens=7,
@@ -54,6 +56,7 @@ def _vllm_config(**over) -> SimpleNamespace:
             enable_expert_parallel=False,
             enable_eplb=False,
             decode_context_parallel_size=8,
+            cp_kv_cache_interleave_size=1,
             prefill_context_parallel_size=1,
             use_ubatching=False,
         ),
@@ -72,7 +75,10 @@ def _vllm_config(**over) -> SimpleNamespace:
             ),
         ),
         cache_config=SimpleNamespace(
-            mamba_cache_dtype="auto", use_replayssm=False, use_kda_recoverssm=False
+            cache_dtype="fp8",
+            mamba_cache_dtype="auto",
+            use_replayssm=False,
+            use_kda_recoverssm=False,
         ),
         lora_config=None,
         kv_transfer_config=None,
@@ -83,6 +89,7 @@ def _vllm_config(**over) -> SimpleNamespace:
         ),
         weight_transfer_config=None,
         speculative_config=spec,
+        scheduler_config=SimpleNamespace(max_num_seqs=1),
         use_v2_model_runner=True,
         compilation_config=SimpleNamespace(
             mode=CompilationMode.VLLM_COMPILE,
@@ -98,13 +105,26 @@ def _vllm_config(**over) -> SimpleNamespace:
     return c
 
 
+@pytest.fixture(autouse=True)
+def fp4_bmm(monkeypatch):
+    monkeypatch.setenv("VLLM_ROCM_USE_AITER_FP4BMM", "1")
+
+
 @pytest.fixture
 def on_mi355x(monkeypatch):
     monkeypatch.setattr(runtime, "platform_reasons", lambda: [])
 
 
-@pytest.mark.parametrize("k, want", [(7, (8,)), (3, (4, 8)), (1, (2, 4, 6, 8))])
-def test_widths_are_whole_requests_up_to_eight_rows(k, want):
+@pytest.mark.parametrize(
+    "k, want",
+    [
+        (7, (8, 16)),
+        (3, (4, 8, 12, 16)),
+        (1, tuple(range(2, 17, 2))),
+        (15, (16,)),
+    ],
+)
+def test_widths_are_whole_requests_up_to_max_rows(k, want):
     assert DecodeConfig(q=k + 1).widths == want
 
 
@@ -169,7 +189,7 @@ def supported_vllm_config(tmp_path, monkeypatch, on_mi355x):
             "short_conv_kernel_size": 4,
         },
     ).save_pretrained(tmp_path)
-    return VllmConfig(
+    cfg = VllmConfig(
         model_config=ModelConfig(
             model=str(tmp_path), dtype=torch.bfloat16, max_model_len=1024, seed=0
         ),
@@ -180,6 +200,10 @@ def supported_vllm_config(tmp_path, monkeypatch, on_mi355x):
         additional_config={"decode_backend": "latency_specialized"},
         speculative_config=SpeculativeConfig(method="ngram", num_speculative_tokens=7),
     )
+    # a served DCP1 step: one request, an fp8 KV cache (M1c)
+    cfg.scheduler_config.max_num_seqs = 1
+    cfg.cache_config.cache_dtype = "fp8"
+    return cfg
 
 
 @pytest.mark.parametrize("read_mode", [False, True])
@@ -218,18 +242,43 @@ def test_unsupported_kv_transfer_is_refused(supported_vllm_config, overrides):
         specialized_decode_config(supported_vllm_config)
 
 
+@pytest.mark.parametrize("dcp", [1, 8])
+@pytest.mark.parametrize("interleave", [1, 16])
+def test_any_interleave_is_accepted(on_mi355x, dcp, interleave):
+    """M1a writes the slots vLLM maps, M1b masks a request's empty shards."""
+    cfg = _vllm_config(
+        **{
+            "parallel_config.decode_context_parallel_size": dcp,
+            "parallel_config.cp_kv_cache_interleave_size": interleave,
+        }
+    )
+    assert specialized_decode_config(cfg) == DecodeConfig(q=8, interleave=interleave)
+
+
+@pytest.mark.parametrize(
+    "over, env, why",
+    [
+        ({"cache_config.cache_dtype": "auto"}, "1", "an fp8 KV cache"),
+        ({}, "0", "VLLM_ROCM_USE_AITER_FP4BMM=1"),
+    ],
+)
+def test_mla_without_m1_runs_vllm_ops(on_mi355x, monkeypatch, over, env, why):
+    """What only M1 needs is no refusal: the MLA layers run vLLM's ops, K2."""
+    monkeypatch.setenv("VLLM_ROCM_USE_AITER_FP4BMM", env)
+    cfg = _vllm_config(**over)
+    assert runtime.unsupported(cfg) == []
+    reasons = runtime.m1_reasons(cfg)
+    assert len(reasons) == 1 and why in reasons[0]
+    assert specialized_decode_config(cfg) == DecodeConfig(q=8, m1=False)
+
+
 @pytest.mark.parametrize(
     "over, reason",
     [
-        (
-            {
-                "parallel_config.tensor_parallel_size": 4,
-                "parallel_config.decode_context_parallel_size": 4,
-            },
-            "TP 4",
-        ),
+        ({"parallel_config.tensor_parallel_size": 4}, "TP 4"),
         ({"parallel_config.pipeline_parallel_size": 2}, "pipeline parallelism"),
-        ({"parallel_config.decode_context_parallel_size": 2}, "decode context"),
+        ({"parallel_config.decode_context_parallel_size": 2}, "other than 1 or 8"),
+        ({"parallel_config.decode_context_parallel_size": 4}, "other than 1 or 8"),
         (
             {
                 "kv_transfer_config": KVTransferConfig(
@@ -239,7 +288,7 @@ def test_unsupported_kv_transfer_is_refused(supported_vllm_config, overrides):
             "KV transfer connectors",
         ),
         ({"speculative_config": None}, "without speculative decoding"),
-        ({"speculative_config.num_speculative_tokens": 8}, "more than 7"),
+        ({"speculative_config.num_speculative_tokens": 16}, "more than 15"),
         ({"speculative_config.enable_adaptive_verification": True}, "adaptive"),
         ({"speculative_config.dynamic": True}, "dynamic speculative"),
         ({"cache_config.mamba_cache_dtype": "float32"}, "KDA cache dtypes"),
@@ -257,6 +306,62 @@ def test_refusals(on_mi355x, over, reason):
     assert len(reasons) == 1 and reason in reasons[0], reasons
 
 
+def _step(dcp: int, seqs: int, k: int) -> SimpleNamespace:
+    return _vllm_config(
+        **{
+            "parallel_config.decode_context_parallel_size": dcp,
+            "scheduler_config.max_num_seqs": seqs,
+            "speculative_config.num_speculative_tokens": k,
+        }
+    )
+
+
+@pytest.mark.parametrize(
+    "dcp, seqs, k",
+    [
+        *[(8, s, k) for s, k in [(1, 1), (1, 3), (1, 7), (1, 15), (2, 1), (2, 3)]],
+        *[(8, s, k) for s, k in [(2, 7), (4, 1), (4, 3), (8, 1)]],
+        *[(1, s, k) for s, k in [(1, 3), (1, 7), (2, 3), (2, 7)]],
+    ],
+)
+def test_served_steps_are_accepted(on_mi355x, dcp, seqs, k):
+    assert runtime.unsupported(_step(dcp, seqs, k)) == []
+
+
+@pytest.mark.parametrize(
+    "dcp, seqs, k, reason",
+    [
+        (8, 3, 1, "max_num_seqs 3 at 1"),
+        (8, 16, 1, "max_num_seqs 16 at 1"),
+        (8, 4, 7, "max_num_seqs 4 at 7"),
+        (8, 8, 3, "max_num_seqs 8 at 3"),
+        (8, 256, 7, "max_num_seqs 256 at 7"),
+        (1, 2, 15, "max_num_seqs 2 at 15"),
+        # 3 requests at qlen 4 (12 rows): M1c's 80 splits do not share out
+        (1, 4, 3, "DCP1 runs M1c only"),
+        # no M1c at qlen 2 / 16: AITER's asm decode
+        (1, 1, 1, "DCP1 runs M1c only"),
+        (1, 8, 1, "DCP1 runs M1c only"),
+        (1, 1, 15, "DCP1 runs M1c only"),
+    ],
+)
+def test_unserved_steps_are_refused(on_mi355x, dcp, seqs, k, reason):
+    reasons = runtime.unsupported(_step(dcp, seqs, k))
+    assert len(reasons) == 1 and reason in reasons[0], reasons
+
+
+def test_dcp1_without_m1_is_refused(on_mi355x):
+    """DCP1 runs M1c, which needs M1 (at DCP8 vLLM's ops run instead)."""
+    cfg = _vllm_config(
+        **{
+            "parallel_config.decode_context_parallel_size": 1,
+            "cache_config.cache_dtype": "auto",
+        }
+    )
+    reasons = runtime.unsupported(cfg)
+    assert len(reasons) == 1 and "DCP1 without M1 (an fp8 KV cache)" in reasons[0]
+
+
 def test_every_refusal_is_reported_at_once(monkeypatch):
     monkeypatch.setattr(
         runtime, "platform_reasons", lambda: ["this GPU (gfx950 / MI355X only)"]
@@ -285,6 +390,17 @@ def test_cache_refusal_names_the_cache():
     assert len(reasons) == 2 and all(r.startswith("layers.5.attn:") for r in reasons)
 
 
+def test_cache_refusal_checks_the_mla_cache():
+    pytest.importorskip("flydsl")
+    ok = torch.zeros(4, 16, 576, dtype=torch.uint8)
+    assert cache_reasons({"layers.3.attn": ok}) == []
+    strided = torch.zeros(4, 32, 576, dtype=torch.uint8)[:, ::2]
+    bf16 = torch.zeros(4, 16, 576, dtype=torch.bfloat16)
+    for bad in (strided, bf16):
+        reasons = cache_reasons({"layers.3.attn": bad})
+        assert len(reasons) == 1 and reasons[0].startswith("layers.3.attn:")
+
+
 def _kda_metadata(**md):
     """One request of 4 tokens unless overridden."""
     m = SimpleNamespace(
@@ -320,6 +436,38 @@ def test_steps_need_metadata():
     assert step_reason(None, "kda", 4) == "no metadata"
 
 
+@pytest.mark.parametrize(
+    "mla, why",
+    [
+        ({}, None),
+        (None, "no MLA decode metadata"),
+        ({"decode": None}, "no MLA decode metadata"),
+        ({"num_prefills": 1}, "MLA: not pure decode"),
+        ({"num_decode_tokens": 2}, "MLA: not pure decode"),
+        ({"decode": SimpleNamespace(max_qo_len=2)}, "MLA: not uniform"),
+        (
+            {"decode": SimpleNamespace(max_qo_len=4, use_gluon_verify=True)},
+            "MLA: Gluon decode",
+        ),
+        (
+            {"decode": SimpleNamespace(max_qo_len=4, use_gluon_decode=True)},
+            "MLA: Gluon decode",
+        ),
+    ],
+)
+def test_steps_check_the_mla_decode(mla, why):
+    """One request of 4 tokens: the MLA layers' metadata must agree."""
+    md = _kda_metadata()
+    if mla is not None:
+        a = SimpleNamespace(
+            decode=SimpleNamespace(max_qo_len=4), num_prefills=0, num_decode_tokens=4
+        )
+        for k, v in mla.items():
+            setattr(a, k, v)
+        md["mla"] = a
+    assert step_reason(md, "kda", 4, "mla") == why
+
+
 class _Kernel:
     def launch(self, step, positions, hidden_states, residual, prefix_delta):
         return "kernel", step.width
@@ -330,8 +478,12 @@ def _decode(monkeypatch, mode=CUDAGraphMode.FULL, md=None):
     KV cache, under a forward context of ``mode``."""
     decode = object.__new__(SpecializedDecode)
     decode.config = DecodeConfig(q=4)
-    decode.runner = SimpleNamespace(kda_prefix="kda")
-    decode.kernels = {w: {1: _Kernel(), 2: _Kernel()} for w in (4, 8)}
+    decode.runner = SimpleNamespace(kda_prefix="kda", mla_prefix="")
+    # Launch-only stand-ins for the GPU kernels.
+    decode.kernels = {
+        w: {1: _Kernel(), 2: _Kernel()}  # type: ignore[dict-item]
+        for w in (4, 8)
+    }
     decode._kv = (1, 2)
     monkeypatch.setattr(SpecializedDecode, "_kv_key", lambda self: (1, 2))
     context = SimpleNamespace(
@@ -400,3 +552,90 @@ def test_weight_views_reject_every_mismatch():
     assert "layer 5: routed_scaling_factor 2.0" in reasons
     assert any(r.startswith("layer 5: MoE backend must be") for r in reasons)
     assert moe.E == 896
+
+
+def test_attend_hands_m1a_padded_query_to_forward_mqa():
+    """forward_mqa gets M1a's padded query's real heads, and AITER's pad of
+    them is M1a's buffer itself (no copy)."""
+    from vllm.models.kimi_k3.amd.specialized_decode.runtime import mla
+    from vllm.v1.attention.backends.mla.rocm_aiter_mla import AiterMLAHelper
+
+    pad = AiterMLAHelper.get_mla_padded_q
+    q = torch.randn(4, 96, 8)
+    padded = pad(96, q, target_heads=128)  # what M1a writes
+    seen = {}
+
+    def forward_mqa(q, kv, md, attn):
+        assert q.shape[1] == 96
+        seen["q"] = pad(96, q, target_heads=128)
+        return None, None
+
+    impl = SimpleNamespace(_decode_num_heads=96, forward_mqa=forward_mqa)
+    wrapper = SimpleNamespace(mla_attn=SimpleNamespace(impl=impl))
+    mla.attend(wrapper, padded, torch.zeros(4, dtype=torch.uint8), None)
+    got = seen["q"]
+    assert got.data_ptr() == padded.data_ptr()  # no copy
+    assert got.shape == padded.shape and got.stride() == padded.stride()
+
+
+def test_qkv_gate_weights_split_the_merged_gate():
+    """KimiK3MergedQKVGateLinear (vLLM's) stacks [q_a | kv_a | gate]: M1a
+    reads the two row slices in place."""
+    from vllm.models.kimi_k3.amd.specialized_decode.runtime import mla
+
+    w = torch.randn(2112 + 1536, 8)
+    wrapper = SimpleNamespace(
+        fused_qkv_a_proj=SimpleNamespace(weight=w), g_proj=None,
+        _qkv_gate_fused=True, q_lora_rank=1536, qk_rope_head_dim=64,
+        mla_attn=SimpleNamespace(kv_lora_rank=512),
+    )  # fmt: skip
+    qkv, g = mla.qkv_gate_weights(wrapper)
+    assert qkv.shape == (2112, 8) and g.shape == (1536, 8)
+    assert qkv.is_contiguous() and g.is_contiguous()
+    assert g.data_ptr() == w[2112:].data_ptr()
+    wrapper._qkv_gate_fused = False
+    assert mla.qkv_gate_weights(wrapper) == (w, None)
+    gw = torch.randn(1536, 8)
+    wrapper.g_proj = SimpleNamespace(weight=gw)
+    assert mla.qkv_gate_weights(wrapper) == (w, gw)
+
+
+def test_m1c_falls_back_where_its_splits_dont_share_out():
+    """M1c's KV splits share out evenly over a width's requests or the width
+    runs the asm decode: qlen 4's 80 splits over 12 rows (3 requests) don't."""
+    from vllm.models.kimi_k3.amd.specialized_decode.runtime.layers import (
+        m1c_at,
+        m1c_splits,
+    )
+
+    k3 = SimpleNamespace(m1c=True, num_spec=3)
+    assert [w for w in (4, 8, 12, 16) if m1c_at(k3, w)] == [4, 8, 16]
+    assert m1c_splits(k3, 8) == 80 and m1c_splits(k3, 12) == 0
+    k7 = SimpleNamespace(m1c=True, num_spec=7)
+    assert all(m1c_at(k7, w) for w in (8, 16))
+    off = SimpleNamespace(m1c=False, num_spec=3)
+    assert not m1c_at(off, 4) and m1c_splits(off, 4) == 0
+
+
+def test_mla_weight_views_fall_back_per_model():
+    """An MLA layer M1 cannot read is no weight refusal: mla_weight_reasons
+    names it, and the MLA layers run vLLM's ops."""
+    pytest.importorskip("flydsl")
+    no_latent = SimpleNamespace(use_latent_moe=False, shared_experts=None)
+    config = _k3_text_config()
+    ffn = [i for i, kind in runtime.layer_kinds(config).items() if kind == "ffn"]
+    w = SimpleNamespace(
+        num_heads=16, q_lora_rank=1536, qk_nope_head_dim=128, qk_rope_head_dim=64,
+        v_head_dim=128, mla_attn=SimpleNamespace(kv_lora_rank=512),
+    )  # fmt: skip
+    layers = [
+        SimpleNamespace(mlp=no_latent, self_attn=SimpleNamespace(mla_attn=w))
+    ] * 93
+    model = SimpleNamespace(config=config, layers=layers, start_layer=0, end_layer=93)
+    assert not any("MLA" in r for r in weight_reasons(model))
+    reasons = mla_weight_reasons(model)
+    assert ffn and len(reasons) == len(ffn)
+    assert (
+        reasons[0] == f"layer {ffn[0]}: MLA heads/q_lora/kv_lora/nope/rope/v "
+        "(16, 1536, 512, 128, 64, 128)"
+    )
diff --git a/tests/models/kimi_k3/test_amd_attn_res.py b/tests/models/kimi_k3/test_amd_attn_res.py
index 871d53faae12..622e8f989782 100644
--- a/tests/models/kimi_k3/test_amd_attn_res.py
+++ b/tests/models/kimi_k3/test_amd_attn_res.py
@@ -99,6 +99,7 @@ def test_amd_attn_res_matches_reference(
         0.0,
     )
 
+    assert isinstance(actual, torch.Tensor)
     torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2)
     torch.testing.assert_close(prefix, original_prefix, atol=0, rtol=0)
     torch.testing.assert_close(blocks, original_blocks, atol=0, rtol=0)
@@ -185,6 +186,7 @@ def test_amd_attn_res_fused_contract(
         output_eps,
     )
 
+    assert isinstance(actual, torch.Tensor)
     torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2)
     torch.testing.assert_close(prefix, expected_prefix, atol=0, rtol=0)
     if write_block:
diff --git a/tests/models/kimi_k3/test_amd_kda_checkpoint.py b/tests/models/kimi_k3/test_amd_kda_checkpoint.py
index 2f8bd3f0e4f8..ee8fbd3b303c 100644
--- a/tests/models/kimi_k3/test_amd_kda_checkpoint.py
+++ b/tests/models/kimi_k3/test_amd_kda_checkpoint.py
@@ -7,6 +7,8 @@
 there must be the one a resumed prefill would have left.
 """
 
+from types import SimpleNamespace
+
 import pytest
 import torch
 
@@ -53,6 +55,8 @@
 class _StubSpeculativeConfig:
     method = "dspark"
     parallel_drafting = False
+    enable_adaptive_verification = False
+    speculator_config = SimpleNamespace(speculator_class="DSparkSpeculator")
 
     def __init__(self, num_speculative_tokens: int) -> None:
         self.num_speculative_tokens = num_speculative_tokens
@@ -507,7 +511,6 @@ def test_opt_in_publishes_chunk_walk_alignment(monkeypatch) -> None:
     [
         ("use_prefill_checkpoint", False),
         ("kda_prefill_backend", "triton"),
-        ("kda_prefill_backend", "gluon"),
         ("use_safe_gate", False),
     ],
 )
@@ -518,6 +521,25 @@ def test_paths_that_cannot_export_stay_opted_out(monkeypatch, attr, value) -> No
     assert spec.prefill_checkpoint_alignment is None
 
 
+def test_gluon_backend_keeps_the_checkpoint(monkeypatch) -> None:
+    # AITER's chunk_kda cannot export the snapshot, but the steps that carry
+    # one run on the fused kernel, which this module already requires.
+    spec = _spec_for(monkeypatch, kda_prefill_backend="gluon")
+
+    assert spec.num_prefill_checkpoint_blocks == 1
+    assert spec.prefill_checkpoint_alignment == KDA_CHECKPOINT_ALIGNMENT
+
+
+def test_gluon_backend_opts_out_without_the_fused_kernel(monkeypatch) -> None:
+    import vllm.models.kimi_k3.amd.kda as kda
+
+    monkeypatch.setattr(kda, "is_fused_kda_chunk_supported", lambda: False)
+    spec = _spec_for(monkeypatch, kda_prefill_backend="gluon")
+
+    assert spec.num_prefill_checkpoint_blocks == 0
+    assert spec.prefill_checkpoint_alignment is None
+
+
 def test_speculative_decoding_does_not_opt_out(monkeypatch) -> None:
     # num_spec reaches the kernels as a wider conv row and an interleaved
     # batch, both of which the layer handles.
diff --git a/tests/models/kimi_k3/test_amd_kda_chunk.py b/tests/models/kimi_k3/test_amd_kda_chunk.py
index 6865dea65497..a524b149e1e5 100644
--- a/tests/models/kimi_k3/test_amd_kda_chunk.py
+++ b/tests/models/kimi_k3/test_amd_kda_chunk.py
@@ -537,6 +537,46 @@ def test_checkpoint_leaves_the_output_and_final_state_alone() -> None:
     torch.testing.assert_close(state_got, state_ref, rtol=0, atol=0)
 
 
+def _slice_inputs(inp: dict, start: int, end: int, h0: torch.Tensor) -> dict:
+    """Tokens ``[start, end)`` of a one-sequence fixture, resumed from ``h0``."""
+    part = dict(inp)
+    for key in ("q", "k", "v", "raw_g", "raw_beta"):
+        part[key] = inp[key][:, start:end].contiguous()
+    part["h0"] = h0
+    part["cu"] = torch.tensor([0, end - start], device="cuda", dtype=torch.int32)
+    return part
+
+
+def test_gluon_hands_checkpoint_steps_to_fused() -> None:
+    """A chunked prefill can switch backends on the step that checkpoints.
+
+    The layer runs Gluon by default and the fused kernel only on the steps that
+    carry a checkpoint, since AITER's chunk_kda cannot export one. Each backend
+    resumes from the state the other left, so the output, the final state and
+    the checkpoint must all match a Triton prefill of the whole prompt.
+    """
+    _requires_kernel()
+    from vllm._aiter_ops import rocm_aiter_ops
+
+    if not rocm_aiter_ops.is_kda_gluon_prefill_available():
+        pytest.skip("AITER Gluon chunk_kda is not available here")
+
+    inp = _inputs([1024], seed=31)
+    o_ref, state_ref = _run(inp, use_fused=False)
+    _, ckpt_ref = _run(_prefix_inputs(inp, 0, 512), use_fused=False)
+
+    # Three steps of one prompt: Gluon, fused with a checkpoint at token 512,
+    # then Gluon again.
+    o1, s1 = _run(_slice_inputs(inp, 0, 256, inp["h0"]), False, backend="gluon")
+    ckpt, (o2, s2) = _run_with_checkpoints(_slice_inputs(inp, 256, 768, s1), [256])
+    o3, s3 = _run(_slice_inputs(inp, 768, 1024, s2), False, backend="gluon")
+
+    o = torch.cat([o1, o2, o3], dim=1)
+    torch.testing.assert_close(o.float(), o_ref.float(), rtol=2e-2, atol=2e-2)
+    torch.testing.assert_close(s3, state_ref, rtol=5e-3, atol=5e-3)
+    torch.testing.assert_close(ckpt[0], ckpt_ref[0], rtol=5e-3, atol=5e-3)
+
+
 @pytest.mark.parametrize("offset", [100, 63, 65, 1024])
 @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
 def test_unusable_checkpoint_offsets_write_nothing(
diff --git a/tests/models/kimi_k3/test_amd_latent_moe_runner.py b/tests/models/kimi_k3/test_amd_latent_moe_runner.py
index e47b5fc08e87..237c3e20bc52 100644
--- a/tests/models/kimi_k3/test_amd_latent_moe_runner.py
+++ b/tests/models/kimi_k3/test_amd_latent_moe_runner.py
@@ -131,6 +131,7 @@ def _check_matches_replicated(
     row_sharded: bool = False,
 ) -> None:
     transform, up_weight = _build_transform(device, row_sharded=row_sharded)
+    assert transform.norm is not None
     runner = _tail_runner(transform, tp_world, use_ep)
     group = get_tp_group().device_group
 
@@ -174,6 +175,7 @@ def _check_unfused_tail_gathers_full_hidden(
 ) -> None:
     transform, up_weight = _build_transform(device, row_sharded=True)
     assert transform.up_proj.weight.shape[0] == HIDDEN_SIZE // tp_world
+    assert transform.norm is not None
 
     # all-gather reconstructs the full projection when all ranks feed the same
     # input, so use one seed for all ranks.
@@ -194,6 +196,7 @@ def _check_writes_only_its_own_shard(
     """Two ranks that swapped both weight rows and write offsets still sum to
     the right total, so inspect each slice before the final collective."""
     transform, up_weight = _build_transform(device)
+    assert transform.norm is not None
     runner = _tail_runner(transform, tp_world)
     group = get_tp_group().device_group
 
diff --git a/tests/models/kimi_k3/test_amd_mla_direct_return.py b/tests/models/kimi_k3/test_amd_mla_direct_return.py
index 57df5bad3344..47c5e373bfac 100644
--- a/tests/models/kimi_k3/test_amd_mla_direct_return.py
+++ b/tests/models/kimi_k3/test_amd_mla_direct_return.py
@@ -1,7 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-from types import SimpleNamespace
+import sys
+from types import ModuleType, SimpleNamespace
 
 import pytest
 import torch
@@ -91,9 +92,7 @@ def attention_output(path, output_shape):
         _fused_qk_prep=fused_decode,
         _qkv_gate_fused=False,
         _fold_qb_uk=False,
-        _fused_decode=lambda q, kv, pe, pos, shape, *args: attention_output(
-            "fused", shape
-        ),
+        _fused_decode=lambda q, kv, pe, shape, *args: attention_output("fused", shape),
         g_proj=None,
         o_proj=lambda value: (value, None),
     )
@@ -119,3 +118,67 @@ def attention_output(path, output_shape):
 
     assert output is projected
     assert paths == ["fused" if fused_decode else "regular"]
+
+
+def test_fused_decode_feeds_zero_positions_to_kernel(monkeypatch) -> None:
+    """The fused kernel's identity RoPE must see position 0 for every slot.
+
+    Not every AITER variant clamps positions into the one-row cos/sin cache, and
+    the kernel takes its token count from slot_mapping, padding included.
+    """
+    from vllm.models.kimi_k3.amd.mla import KimiK3MultiHeadLatentAttentionWrapper
+
+    calls = []
+    cache_ops = ModuleType("aiter.ops.cache")
+    monkeypatch.setattr(
+        cache_ops,
+        "fused_qk_rope_concat_and_cache_mla",
+        lambda *args, **kwargs: calls.append(args),
+        raising=False,
+    )
+    monkeypatch.setitem(sys.modules, "aiter", ModuleType("aiter"))
+    monkeypatch.setitem(sys.modules, "aiter.ops", ModuleType("aiter.ops"))
+    monkeypatch.setitem(sys.modules, "aiter.ops.cache", cache_ops)
+
+    num_tokens, num_heads, nope, rope, lora, v_dim = 4, 2, 4, 2, 4, 4
+    head_size = lora + rope
+    zero_positions = torch.zeros(16, dtype=torch.int64)
+    output = torch.randn(num_tokens, num_heads * v_dim, dtype=torch.bfloat16)
+    attn = SimpleNamespace(
+        is_aiter_triton_fp4_bmm_enabled=False,
+        W_UK_T=torch.randn(num_heads, nope, lora, dtype=torch.bfloat16),
+        kv_lora_rank=lora,
+        impl=SimpleNamespace(forward_mqa=lambda q, *args: (q, None)),
+        _v_up_proj=lambda attn_out, out: out.copy_(output),
+    )
+    wrapper = SimpleNamespace(
+        mla_attn=attn,
+        qk_nope_head_dim=nope,
+        qk_rope_head_dim=rope,
+        _identity_rope=None,
+        _fused_q_heads=4,
+        _zero_positions=zero_positions,
+    )
+    # A CUDA-graph padded batch: the last slot is a pad.
+    slot_mapping = torch.tensor([[5], [9], [2], [-1]])
+
+    result = KimiK3MultiHeadLatentAttentionWrapper._fused_decode(
+        wrapper,
+        torch.randn(num_tokens, num_heads, nope + rope, dtype=torch.bfloat16),
+        torch.randn(num_tokens, lora, dtype=torch.bfloat16),
+        torch.randn(num_tokens, 1, rope, dtype=torch.bfloat16),
+        tuple(output.shape),
+        None,
+        SimpleNamespace(_k_scale=torch.ones(1), _q_scale=torch.ones(1)),
+        torch.zeros(2, 4, head_size, dtype=torch.uint8),
+        slot_mapping,
+    )
+
+    assert len(calls) == 1
+    kernel_slots, positions = calls[0][6], calls[0][9]
+    assert torch.equal(kernel_slots, slot_mapping.flatten())
+    assert positions.shape == kernel_slots.shape
+    assert not positions.any()
+    # A view of the persistent buffer, so captured graphs allocate nothing.
+    assert positions.data_ptr() == zero_positions.data_ptr()
+    assert torch.equal(result, output)
diff --git a/tests/models/kimi_k3/test_eagle3.py b/tests/models/kimi_k3/test_eagle3.py
index 858622fa07da..562ade2b8b6a 100644
--- a/tests/models/kimi_k3/test_eagle3.py
+++ b/tests/models/kimi_k3/test_eagle3.py
@@ -36,7 +36,8 @@ def test_kimi_linear_advertises_eagle3_support():
 
 
 def test_kimi_k3_uses_shared_eagle3_layer_configuration():
-    target = object.__new__(KimiK3ForConditionalGeneration)
+    # Skip model initialization; the test supplies the required attributes below.
+    target = object.__new__(KimiK3ForConditionalGeneration)  # type: ignore[type-abstract]
     torch.nn.Module.__init__(target)
     model = _make_kimi_linear_model()
     object.__setattr__(model, "layers", [None] * 93)
@@ -79,12 +80,14 @@ def test_kimi_linear_forward_extracts_standard_aux_hidden_states(monkeypatch):
         lambda: SimpleNamespace(is_first_rank=True, is_last_rank=True),
     )
 
-    output, aux_hidden_states = model.forward(
+    result = model.forward(
         input_ids=None,
         positions=torch.tensor([0]),
         intermediate_tensors=None,
         inputs_embeds=initial_hidden_states,
     )
+    assert isinstance(result, tuple)
+    output, aux_hidden_states = result
 
     expected_layer_output = layer_hidden_states + layer_residual
     torch.testing.assert_close(output, expected_layer_output)
@@ -128,12 +131,14 @@ def test_kimi_linear_forward_extracts_attn_res_aux_hidden_states(monkeypatch):
     final_attn_res = Mock(return_value=final_hidden_states)
     monkeypatch.setattr(kimi_model, "attn_res", final_attn_res)
 
-    output, aux_hidden_states = model.forward(
+    result = model.forward(
         input_ids=None,
         positions=torch.tensor([0]),
         intermediate_tensors=None,
         inputs_embeds=initial_hidden_states,
     )
+    assert isinstance(result, tuple)
+    output, aux_hidden_states = result
 
     torch.testing.assert_close(output, final_hidden_states)
     torch.testing.assert_close(aux_hidden_states[0], initial_hidden_states)
@@ -187,12 +192,14 @@ def test_attn_res_stream_capture_receives_the_layer_outputs_in_order(monkeypatch
     capture = Mock(return_value=captured)
     monkeypatch.setattr(KimiLinearModel, "_capture_aux_hidden_stream", capture)
 
-    _, aux_hidden_states = model.forward(
+    result = model.forward(
         input_ids=None,
         positions=torch.tensor([0]),
         intermediate_tensors=None,
         inputs_embeds=initial_hidden_states,
     )
+    assert isinstance(result, tuple)
+    _, aux_hidden_states = result
 
     layer_idx, got_prefix, got_pending, got_residual = capture.call_args.args
     assert layer_idx == 0
diff --git a/tests/models/kimi_k3/test_kda.py b/tests/models/kimi_k3/test_kda.py
index 2dfc08eb0497..becf78b59445 100644
--- a/tests/models/kimi_k3/test_kda.py
+++ b/tests/models/kimi_k3/test_kda.py
@@ -6,6 +6,7 @@
 Uses torch.rand for q/k/v to match FLA's test pattern.
 """
 
+from collections.abc import Callable
 from types import SimpleNamespace
 
 import pytest
@@ -77,7 +78,7 @@
     "nvidia": fused_recurrent_kda_packed_decode,
     "amd": fused_recurrent_kda_packed_decode_amd,
 }
-SPEC_DECODE_IMPLS = {
+SPEC_DECODE_IMPLS: dict[str, Callable[..., tuple[torch.Tensor, torch.Tensor]]] = {
     "nvidia": fused_recurrent_kda_nvidia,
     "amd": fused_recurrent_kda_amd,
 }
diff --git a/tests/models/kimi_k3/test_kda_metadata.py b/tests/models/kimi_k3/test_kda_metadata.py
index 61c19e8961de..edd4ea771ff2 100644
--- a/tests/models/kimi_k3/test_kda_metadata.py
+++ b/tests/models/kimi_k3/test_kda_metadata.py
@@ -3,6 +3,7 @@
 
 from dataclasses import fields
 from types import SimpleNamespace
+from typing import TypeVar
 from unittest.mock import Mock, patch
 
 import pytest
@@ -26,7 +27,6 @@
 from vllm.models.kimi_k3.nvidia.model import KimiLinearForCausalLM
 from vllm.platforms import current_platform
 from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig
-from vllm.v1.attention.backend import AttentionMetadataBuilder
 from vllm.v1.attention.backends.gdn_attn import (
     GDNAttentionBackend,
     GDNAttentionMetadata,
@@ -58,6 +58,8 @@
     "flashinfer_prefill_seq_order",
 }
 
+_BuilderT = TypeVar("_BuilderT", bound=GDNAttentionMetadataBuilder)
+
 
 def _assert_matches_shared_gdn(
     reference: GDNAttentionMetadata, actual: KimiK3KDAMetadata
@@ -92,7 +94,7 @@ def _assert_matches_shared_gdn(
 
 
 def _make_builder(
-    builder_cls: type[AttentionMetadataBuilder],
+    builder_cls: type[_BuilderT],
     num_speculative_tokens: int,
     full_cuda_graph: bool,
     device: torch.device = DEVICE,
@@ -105,7 +107,7 @@ def _make_builder(
     disable_eagle_block_drop: bool = False,
     cache_hit_alignment_tokens: int | None = None,
     hash_block_size: int | None = None,
-) -> AttentionMetadataBuilder:
+) -> _BuilderT:
     vllm_config = create_vllm_config(
         model_name="Qwen/Qwen3.5-0.8B",
         block_size=BLOCK_SIZE,
@@ -213,6 +215,7 @@ def test_kda_recoverssm_startup_metadata_flow_without_model(monkeypatch):
         speculative_config=SimpleNamespace(
             num_speculative_tokens=2,
             parallel_drafting=False,
+            enable_adaptive_verification=False,
             use_eagle_block_drop=Mock(return_value=False),
         ),
         compilation_config=SimpleNamespace(
@@ -659,6 +662,8 @@ def test_recoverssm_spec_uses_one_state_slot_and_current_window(
         )
     context = builder.recoverssm_context
     assert context is not None
+    commit = context.commit
+    assert isinstance(commit, Mock)
     actual = builder.build(
         0,
         common_attn_metadata,
@@ -696,7 +701,7 @@ def test_recoverssm_spec_uses_one_state_slot_and_current_window(
             is common_attn_metadata.compute_num_computed_tokens()
         )
         assert postprocess.block_size == BLOCK_SIZE
-    args = context.commit.call_args.args
+    args = commit.call_args.args
     assert args[0] is num_accepted_tokens
     torch.testing.assert_close(args[1], commit_metadata.state_indices[:, 0])
     torch.testing.assert_close(args[2], commit_metadata.query_start_loc)
@@ -807,7 +812,7 @@ def test_kimi_k3_kda_cudagraph_capture_matches_shared_gdn():
     ],
 )
 def test_cudagraph_capture_metadata_avoids_device_to_host_copy(
-    builder_cls: type[AttentionMetadataBuilder], monkeypatch: pytest.MonkeyPatch
+    builder_cls: type[GDNAttentionMetadataBuilder], monkeypatch: pytest.MonkeyPatch
 ):
     """`build_for_cudagraph_capture` is not capture-only: a DP rank with nothing
     scheduled re-stages its FULL-graph metadata through it on every dummy step,
diff --git a/tests/models/kimi_k3/test_latent_moe_tail.py b/tests/models/kimi_k3/test_latent_moe_tail.py
index dd3136641121..45235a67a923 100644
--- a/tests/models/kimi_k3/test_latent_moe_tail.py
+++ b/tests/models/kimi_k3/test_latent_moe_tail.py
@@ -101,7 +101,8 @@ def fake_tail_initialize(**kwargs):
 
     latent_moe_runner.LatentMoERunner()
 
-    assert moe_config.use_deferred_moe_finalize
+    # The stand-in deliberately borrows the real config's property.
+    assert moe_config.use_deferred_moe_finalize  # type: ignore[misc]
     assert moe_config.defer_moe_finalize_max_num_tokens == 128
     assert initialized_with["experts_per_token"] == 16
 
diff --git a/tests/models/kimi_k3/test_mla_prefill_context.py b/tests/models/kimi_k3/test_mla_prefill_context.py
index e8cfb014f63c..6c5362e35520 100644
--- a/tests/models/kimi_k3/test_mla_prefill_context.py
+++ b/tests/models/kimi_k3/test_mla_prefill_context.py
@@ -126,7 +126,6 @@ def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, None]:
 class _FusedLayer:
     """Only the attributes K3's context loop reads."""
 
-    _compute_prefill_context = MultiHeadLatentAttention._compute_prefill_context
     _gather_context_latent = MultiHeadLatentAttention._gather_context_latent
     _attn_read_kv_cache = MultiHeadLatentAttention._attn_read_kv_cache
 
@@ -144,7 +143,6 @@ def __init__(self, kv_b_proj, kv_cache, kv_cache_dtype, k_scale) -> None:
 class _ReferenceImpl:
     """Only the attributes the generic context loop reads."""
 
-    _compute_prefill_context = MLACommonBaseImpl._compute_prefill_context
     _concat_k_nope_k_pe = MLACommonBaseImpl._concat_k_nope_k_pe
     _use_flashinfer_concat_mla_k = False
 
@@ -163,7 +161,7 @@ def _build_prefill_metadata(
     workspace_dtype: torch.dtype,
     q_data_type: torch.dtype,
     backend: _RecordingPrefillBackend,
-) -> MLACommonPrefillMetadata:
+) -> tuple[MLACommonPrefillMetadata, int]:
     query_start_loc_cpu = torch.zeros(len(_QUERY_LENS) + 1, dtype=torch.int32)
     query_start_loc_cpu[1:] = torch.tensor(_QUERY_LENS, dtype=torch.int32).cumsum(0)
     workspace = torch.empty(
@@ -253,13 +251,13 @@ def test_fused_context_matches_generic_impl(
     ).to(q_data_type)
 
     layer = _FusedLayer(kv_b_proj, kv_cache, kv_cache_dtype, k_scale)
-    fused_out, fused_lse = layer._compute_prefill_context(
-        q, SimpleNamespace(prefill=prefill_fused)
+    fused_out, fused_lse = MultiHeadLatentAttention._compute_prefill_context(
+        layer, q, SimpleNamespace(prefill=prefill_fused)
     )
 
     impl = _ReferenceImpl(kv_b_proj, kv_cache_dtype)
-    ref_out, ref_lse = impl._compute_prefill_context(
-        q, kv_cache, SimpleNamespace(prefill=prefill_ref), k_scale
+    ref_out, ref_lse = MLACommonBaseImpl._compute_prefill_context(
+        impl, q, kv_cache, SimpleNamespace(prefill=prefill_ref), k_scale
     )
 
     assert len(backend_fused.calls) == len(backend_ref.calls)
@@ -281,6 +279,7 @@ def test_fused_context_matches_generic_impl(
 
     # Every chunk but the continuation should have been written in place, i.e.
     # straight into the returned accumulator, with no intermediate copy.
+    assert prefill_fused.chunked_context is not None
     wrote_in_place = [
         out is not None and out.data_ptr() == fused_out[chunk.token_slice].data_ptr()
         for out, chunk in zip(
@@ -318,7 +317,9 @@ def test_fused_context_rejects_an_unquantized_query() -> None:
         dtype=torch.bfloat16,
     )
     with pytest.raises(AssertionError, match="new-token epilogue"):
-        layer._compute_prefill_context(q, SimpleNamespace(prefill=prefill))
+        MultiHeadLatentAttention._compute_prefill_context(
+            layer, q, SimpleNamespace(prefill=prefill)
+        )
 
 
 # --------------------------------------------------------------------------
diff --git a/tests/models/kimi_k3/test_sequence_parallel.py b/tests/models/kimi_k3/test_sequence_parallel.py
index 3057f85c7e44..09564b4221d3 100644
--- a/tests/models/kimi_k3/test_sequence_parallel.py
+++ b/tests/models/kimi_k3/test_sequence_parallel.py
@@ -142,9 +142,10 @@ def test_kimi_decoder_layer_keeps_moe_states_sequence_sharded(monkeypatch):
     layer.input_layernorm = _IdentityNorm()
     layer.post_attention_layernorm = _IdentityNorm()
     layer.mlp = _RecordingMoE()
-    layer._run_self_attn = MethodType(
-        lambda self, positions, hidden_states: hidden_states,
+    monkeypatch.setattr(
         layer,
+        "_run_self_attn",
+        MethodType(lambda self, positions, hidden_states: hidden_states, layer),
     )
 
     _mock_sequence_parallel_collectives(monkeypatch)
@@ -188,9 +189,10 @@ def test_kimi_attn_residual_states_stay_sequence_sharded(monkeypatch):
     layer.self_attention_res_proj = _Projection()
     layer.mlp_res_proj = _Projection()
     layer.mlp = _RecordingMoE()
-    layer._run_self_attn = MethodType(
-        lambda self, positions, hidden_states: hidden_states,
+    monkeypatch.setattr(
         layer,
+        "_run_self_attn",
+        MethodType(lambda self, positions, hidden_states: hidden_states, layer),
     )
 
     _mock_sequence_parallel_collectives(monkeypatch)
diff --git a/tests/models/language/pooling/test_jina_reranker_v3.py b/tests/models/language/pooling/test_jina_reranker_v3.py
index 3fc5c48c5a59..a53689cfbd31 100644
--- a/tests/models/language/pooling/test_jina_reranker_v3.py
+++ b/tests/models/language/pooling/test_jina_reranker_v3.py
@@ -7,8 +7,10 @@
 import torch.nn.functional as F
 
 from tests.utils import RemoteOpenAIServer
+from vllm import PoolingParams
 from vllm.entrypoints.pooling.pooling.protocol import PoolingResponse
 from vllm.entrypoints.pooling.scoring.protocol import RerankResponse, ScoreResponse
+from vllm.exceptions import VLLMValidationError
 
 model_name = "jinaai/jina-reranker-v3"
 query = "What are the health benefits of green tea?"
@@ -53,6 +55,8 @@ def test_offline(vllm_runner, monkeypatch):
         _test_offline_1_v_1(llm)
         _test_offline_1_v_n(llm)
         _test_offline_n_v_n(llm)
+        _test_offline_truncation(llm)
+        _test_offline_nested_markers(llm)
         _test_offline_token_embed_illegal_inputs(llm)
         assert llm.model_config.embedding_size == EMBEDDING_SIZE
 
@@ -65,6 +69,8 @@ def test_online():
         _test_online_1_v_n(server)
         _test_online_n_v_n(server)
         _test_online_instruction(server)
+        _test_online_truncation(server)
+        _test_online_nested_markers(server)
         _test_online_token_embed_illegal_inputs(server)
 
 
@@ -145,6 +151,81 @@ def _test_offline_token_embed_illegal_inputs(llm):
         llm.encode([1, 2, 3], pooling_task="token_embed")
 
 
+def _test_offline_truncation(llm):
+    prompts = documents[:2] + [query]
+    (output,) = llm.encode(prompts, pooling_task="token_embed")
+    for side in ("left", "right"):
+        for limit in (
+            -1,
+            1,
+            len(output.prompt_token_ids) - 1,
+            len(output.prompt_token_ids) + 1,
+        ):
+            kwargs = {"truncate_prompt_tokens": limit, "truncation_side": side}
+            for call, args, extra in (
+                (llm.encode, (prompts,), {"pooling_task": "token_embed"}),
+                (llm.score, (query, documents[:2]), {}),
+            ):
+                with pytest.raises(
+                    VLLMValidationError,
+                    match="JinaForRanking does not support truncate_prompt_tokens",
+                ) as exc_info:
+                    call(*args, **extra, tokenization_kwargs=kwargs)
+                assert exc_info.value.parameter == "truncate_prompt_tokens"
+
+    limited = llm.score(
+        "one two three four five",
+        ["six seven eight nine ten", "red green blue yellow white"],
+        pooling_params=PoolingParams(
+            extra_kwargs={"max_tokens_per_query": 3, "max_tokens_per_doc": 4}
+        ),
+    )
+    expected = llm.score(
+        "one two three", ["six seven eight nine", "red green blue yellow"]
+    )
+    assert len(limited) == len(expected) == 2
+    for actual, reference in zip(limited, expected):
+        assert actual.prompt_token_ids == reference.prompt_token_ids
+        assert actual.outputs.score == pytest.approx(reference.outputs.score, abs=TOL)
+
+
+def _test_offline_nested_markers(llm):
+    prompts = documents[:2] + [query]
+    (expected,) = llm.encode(prompts, pooling_task="token_embed")
+    instruction = "Rank relevant passages."
+    scores = llm.score(
+        query,
+        documents[:2],
+        pooling_params=PoolingParams(
+            extra_kwargs={"chat_template_kwargs": {"instruction": instruction}}
+        ),
+    )
+    for token in ("embed", "rerank"):
+        nested = f"<|{token}<|{token}_token|>_token|>"
+        (actual,) = llm.encode(
+            [text + nested for text in prompts], pooling_task="token_embed"
+        )
+        assert actual.outputs.data.shape == (len(prompts), EMBEDDING_SIZE), (
+            actual.outputs.data.shape
+        )
+        assert actual.prompt_token_ids == expected.prompt_token_ids
+        actual_scores = llm.score(
+            query + nested,
+            [doc + nested for doc in documents[:2]],
+            pooling_params=PoolingParams(
+                extra_kwargs={
+                    "chat_template_kwargs": {"instruction": instruction + nested}
+                }
+            ),
+        )
+        assert len(actual_scores) == len(scores)
+        for actual_score, reference in zip(actual_scores, scores):
+            assert actual_score.prompt_token_ids == reference.prompt_token_ids
+            assert actual_score.outputs.score == pytest.approx(
+                reference.outputs.score, abs=TOL
+            )
+
+
 def _get_score_response(server, query, document, **extra_body):
     payload = {
         "model": model_name,
@@ -306,6 +387,91 @@ def _test_online_instruction(server):
     )
 
 
+def _test_online_truncation(server):
+    docs = documents[:2]
+    for endpoint, inputs in (
+        ("pooling", {"input": docs + [query], "task": "token_embed"}),
+        ("score", {"queries": query, "documents": docs}),
+        ("rerank", {"query": query, "documents": docs}),
+    ):
+        payload = {"model": model_name, **inputs}
+        response = requests.post(server.url_for(endpoint), json=payload)
+        response.raise_for_status()
+        num_tokens = response.json()["usage"]["prompt_tokens"]
+        for side in ("left", "right"):
+            for limit in (-1, 1, num_tokens - 1, num_tokens + 1):
+                payload.update(truncate_prompt_tokens=limit, truncation_side=side)
+                response = requests.post(server.url_for(endpoint), json=payload)
+                assert response.status_code == 400, response.status_code
+                error = response.json()["error"]
+                assert (
+                    "JinaForRanking does not support truncate_prompt_tokens"
+                    in error["message"]
+                )
+                assert error["param"] == "truncate_prompt_tokens"
+
+    for get_response in (_get_score_response, _get_rerank_response):
+        limited = get_response(
+            server,
+            "one two three four five",
+            ["six seven eight nine ten", "red green blue yellow white"],
+            max_tokens_per_query=3,
+            max_tokens_per_doc=4,
+        )
+        expected = get_response(
+            server, "one two three", ["six seven eight nine", "red green blue yellow"]
+        )
+        assert limited.usage.prompt_tokens == expected.usage.prompt_tokens
+        if isinstance(limited, ScoreResponse):
+            assert len(limited.data) == len(expected.data) == 2
+            assert [d.score for d in limited.data] == pytest.approx(
+                [d.score for d in expected.data], abs=TOL
+            )
+        else:
+            assert [d.index for d in limited.results] == [
+                d.index for d in expected.results
+            ]
+            assert [d.relevance_score for d in limited.results] == pytest.approx(
+                [d.relevance_score for d in expected.results], abs=TOL
+            )
+
+
+def _test_online_nested_markers(server):
+    docs = documents[:2]
+    expected = _get_embeds(server, docs + [query])
+    instruction = "Rank relevant passages."
+    for token in ("embed", "rerank"):
+        nested = f"<|{token}<|{token}_token|>_token|>"
+        actual = _get_embeds(server, [text + nested for text in docs + [query]])
+        assert actual.shape == expected.shape == (len(docs) + 1, EMBEDDING_SIZE), (
+            actual.shape
+        )
+        torch.testing.assert_close(actual, expected)
+        for get_response in (_get_score_response, _get_rerank_response):
+            reference = get_response(server, query, docs, instruction=instruction)
+            actual_scores = get_response(
+                server,
+                query + nested,
+                [doc + nested for doc in docs],
+                instruction=instruction + nested,
+            )
+            assert actual_scores.usage.prompt_tokens == reference.usage.prompt_tokens
+            if isinstance(actual_scores, ScoreResponse):
+                assert len(actual_scores.data) == len(reference.data) == len(docs)
+                assert [d.score for d in actual_scores.data] == pytest.approx(
+                    [d.score for d in reference.data], abs=TOL
+                )
+            else:
+                assert [d.index for d in actual_scores.results] == [
+                    d.index for d in reference.results
+                ]
+                assert [
+                    d.relevance_score for d in actual_scores.results
+                ] == pytest.approx(
+                    [d.relevance_score for d in reference.results], abs=TOL
+                )
+
+
 def _test_online_token_embed_illegal_inputs(server):
     response = requests.post(
         server.url_for("pooling"),
diff --git a/tests/models/minimax_m3/__init__.py b/tests/models/minimax_m3/__init__.py
new file mode 100644
index 000000000000..e69de29bb2d1
diff --git a/tests/models/minimax_m3/amd_specialized_decode/test_ops.py b/tests/models/minimax_m3/amd_specialized_decode/test_ops.py
index c72a1fde334b..27d2aa2b70b4 100644
--- a/tests/models/minimax_m3/amd_specialized_decode/test_ops.py
+++ b/tests/models/minimax_m3/amd_specialized_decode/test_ops.py
@@ -144,12 +144,18 @@ def _pairs(region, n):
     return region[: n * 8].view(torch.int32).view(n, 2)[:, 0].contiguous()
 
 
+class _TestFfnRunner(FfnRunner):
+    peer_bufs: list[torch.Tensor]
+
+
 def _runner(weights, s, tp=1, widths=None, **build):
     dev = torch.device("cuda")
     bufs = [torch.zeros(ffn.peer_bytes(tp), dtype=torch.uint8, device=dev)]
     bufs += [torch.zeros_like(bufs[0]) for _ in range(tp - 1)]
     peers = torch.tensor([b.data_ptr() for b in bufs], dtype=torch.int64, device=dev)
-    runner = FfnRunner(tp, 0, bufs[0].data_ptr(), peers, dev, widths or (s,), **build)
+    runner = _TestFfnRunner(
+        tp, 0, bufs[0].data_ptr(), peers, dev, widths or (s,), **build
+    )
     runner.warm(weights)
     runner.peer_bufs = bufs
     return runner
diff --git a/tests/models/minimax_m3/amd_specialized_decode/test_runtime.py b/tests/models/minimax_m3/amd_specialized_decode/test_runtime.py
index 326ebde2555f..7aa5afcd9b92 100644
--- a/tests/models/minimax_m3/amd_specialized_decode/test_runtime.py
+++ b/tests/models/minimax_m3/amd_specialized_decode/test_runtime.py
@@ -139,9 +139,13 @@ def _vllm_config(layers=60, **parallel):
 
 
 def test_config_refusals_report_every_reason():
-    assert specialized_decode_config(_vllm_config()).last_layer == 59
+    full_config = specialized_decode_config(_vllm_config())
+    assert full_config is not None
+    assert full_config.last_layer == 59
     # a truncated checkpoint: the sparse layers end at its last layer
-    assert specialized_decode_config(_vllm_config(layers=5)).last_layer == 4
+    truncated_config = specialized_decode_config(_vllm_config(layers=5))
+    assert truncated_config is not None
+    assert truncated_config.last_layer == 4
 
     config = _vllm_config(tensor_parallel_size=2, enable_expert_parallel=True)
     config.cache_config.block_size = 64
@@ -185,7 +189,7 @@ def _layer():
     w2.is_shuffled = True
     fused_experts = object.__new__(AiterMxfp4Experts)
     fused_experts.moe_config = SimpleNamespace(
-        use_mori_kernels=False,
+        use_mori_epv1_kernels=False,
         use_mxfp4_w4a4_dsv4=False,
     )
     experts = SimpleNamespace(
@@ -263,7 +267,7 @@ def _experts(layer):
     ),
     "effective BF16": lambda x: setattr(
         _experts(x).quant_method.moe_kernel.fused_experts.moe_config,
-        "use_mori_kernels",
+        "use_mori_epv1_kernels",
         True,
     ),
     "router": lambda x: setattr(
diff --git a/tests/models/minimax_m3/test_dense_gemm1_mxfp8.py b/tests/models/minimax_m3/test_dense_gemm1_mxfp8.py
new file mode 100644
index 000000000000..cd7516a82559
--- /dev/null
+++ b/tests/models/minimax_m3/test_dense_gemm1_mxfp8.py
@@ -0,0 +1,90 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Correctness test for the MiniMax-M3.1 fused dense GEMM1 kernel
+(vllm/models/minimax_m3/nvidia/ops/dense_gemm1_mxfp8.py).
+"""
+
+import pytest
+import torch
+
+from vllm.platforms import current_platform
+
+pytestmark = pytest.mark.skipif(
+    not (current_platform.is_cuda() and current_platform.has_device_capability(100)),
+    reason="dense GEMM1 MXFP8 kernel is SM100+ CUDA only",
+)
+
+
+def _mxfp8_quantize_f32(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    """MXFP8-quantize fp32 [M, N] with the kernel's integer-exact scale math:
+    scale byte = E - 8 + (amax > 1.75 * 2^E) with E the biased exponent of the
+    32-elem group amax (0 for amax=0), clamped to [0, 254]. Returns q and the
+    unswizzled [M, N//32] scale bytes.
+    """
+    M, N = x.shape
+    g = x.float().view(M, N // 32, 32)
+    amax = g.abs().amax(-1)
+    e_bias = amax.view(torch.int32) >> 23
+    thr = torch.ldexp(torch.full_like(amax, 1.75), e_bias - 127)
+    sb = (e_bias - 8 + (amax > thr).to(torch.int32)).clamp(0, 254)
+    q = g * torch.ldexp(torch.ones_like(amax), 127 - sb)[..., None]
+    q = q.clamp(-448.0, 448.0).to(torch.float8_e4m3fn)
+    return q.view(M, N), sb.to(torch.uint8)
+
+
+# Per-rank I: TP4 shared expert, TP4 dense / TP1 shared, TP1 dense. 12288 is
+# the only one whose weight exceeds half of L2 (grouped rasterization).
+@pytest.mark.parametrize("num_tokens", [1, 37, 128, 1000, 4096])
+@pytest.mark.parametrize("intermediate", [768, 3072, 12288])
+def test_dense_gemm1_mxfp8_correctness(num_tokens: int, intermediate: int):
+    from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
+        mxfp8_e4m3_quantize,
+        swizzle_mxfp8_scale,
+    )
+    from vllm.models.minimax_m3.nvidia.ops.dense_gemm1_mxfp8 import run
+
+    limit, alpha, beta = 7.0, 1.702, 1.0
+    hidden = 6144
+    torch.manual_seed(num_tokens)
+    x = torch.randn(num_tokens, hidden, device="cuda", dtype=torch.bfloat16) * 0.5
+    w1 = (
+        torch.randn(intermediate * 2, hidden, device="cuda", dtype=torch.bfloat16)
+        * 0.05
+    )
+
+    # What Mxfp8OnlineLinearMethod.process_weights_after_loading produces.
+    w1q, wsf_2d = mxfp8_e4m3_quantize(w1, is_sf_swizzled_layout=False)
+    w1sf = swizzle_mxfp8_scale(wsf_2d, M=intermediate * 2, K=hidden)
+
+    M, N = num_tokens, intermediate
+    e8 = torch.float8_e8m0fnu
+    q = torch.empty((M, N), dtype=torch.float8_e4m3fn, device="cuda")
+    sf = torch.zeros(
+        ((M + 127) // 128) * 128 * (N // 32), dtype=torch.uint8, device="cuda"
+    )
+    xq, xsf = mxfp8_e4m3_quantize(x, is_sf_swizzled_layout=True)
+    run(xq, w1q, xsf.view(-1).view(e8), w1sf.view(e8), q, sf, limit, alpha, beta)
+
+    # fp32 mirror chain: same MXFP8 GEMM (fp32 out) -> fp32 clamped SwiGLU
+    # -> MXFP8 quantize with the kernel's scale math.
+    gemm = torch._scaled_mm(
+        xq,
+        w1q.t(),
+        scale_a=xsf.view(e8),
+        scale_b=w1sf.view(e8),
+        out_dtype=torch.float32,
+    )
+    gate = gemm[:, :N].clamp(max=limit)
+    up = gemm[:, N:].clamp(-limit, limit)
+    act = gate * torch.sigmoid(gate * alpha) * (up + beta)
+    q_ref, sf_ref_2d = _mxfp8_quantize_f32(act)
+    sf_ref = swizzle_mxfp8_scale(sf_ref_2d, M=M, K=N)
+
+    # Measured envelope: sf bytes identical; q bytes differ on <=0.007% of
+    # elements, always by 1 fp8 ulp. Bounds below carry >10x margin.
+    sf_mism = (sf != sf_ref).sum().item()
+    assert sf_mism <= sf.numel() * 1e-4, f"sf byte mismatches: {sf_mism}"
+    q_mism = q.view(torch.uint8) != q_ref.view(torch.uint8)
+    assert q_mism.sum().item() <= q.numel() * 1e-3, (
+        f"q byte mismatches: {q_mism.sum().item()}"
+    )
diff --git a/tests/models/minimax_m3/test_dspark_draft_compile.py b/tests/models/minimax_m3/test_dspark_draft_compile.py
new file mode 100644
index 000000000000..dc24bfb18da0
--- /dev/null
+++ b/tests/models/minimax_m3/test_dspark_draft_compile.py
@@ -0,0 +1,135 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""The MiniMax-M3.1 DSpark draft attention under vLLM's torch.compile.
+
+vLLM traces the draft model once, in the memory-profiling run, whose forward
+context has no slot mapping, and reuses that graph afterwards. The draft
+tokens' k/v insert must still happen on the later calls.
+"""
+
+import pytest
+import torch
+from torch import nn
+
+from vllm.utils.torch_utils import set_default_torch_dtype
+
+HEAD_DIM = 128
+BLOCK_SIZE = 16
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available(), reason="CUDA not available"
+)
+
+
+class _QueryAttention(nn.Module):
+    """Stands in for the draft's Attention: returns its query, so the layer
+    output shows the q the attention got."""
+
+    def __init__(self, layer_name):
+        super().__init__()
+        self.layer_name = layer_name
+
+    def forward(self, q, k, v, output_dtype=None):
+        return q.to(output_dtype)
+
+
+@pytest.mark.parametrize(
+    "kv_cache_dtype,query_pre_quantized",
+    [("auto", False), ("fp8", False), ("fp8", True)],
+)
+def test_draft_attention_stores_kv_after_compile(
+    dist_init, monkeypatch, kv_cache_dtype, query_pre_quantized
+):
+    from vllm.compilation.decorators import support_torch_compile
+    from vllm.config import (
+        CompilationConfig,
+        CompilationMode,
+        CUDAGraphMode,
+        VllmConfig,
+        set_current_vllm_config,
+    )
+    from vllm.forward_context import set_forward_context
+    from vllm.models.minimax_m3.common.dspark import MiniMaxM3DSparkAttention
+    from vllm.transformers_utils.configs.minimax_m3 import MiniMaxM3TextConfig
+
+    monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", "1")
+    config = MiniMaxM3TextConfig(
+        hidden_size=512,
+        num_attention_heads=8,
+        num_key_value_heads=2,
+        head_dim=HEAD_DIM,
+        max_position_embeddings=8192,
+    )
+
+    @support_torch_compile
+    class DraftAttention(nn.Module):
+        def __init__(self, *, vllm_config: VllmConfig, prefix: str = "", **kwargs):
+            super().__init__()
+            self.self_attn = MiniMaxM3DSparkAttention(
+                config,
+                sliding_window=None,
+                causal=False,
+                cache_config=None,
+                quant_config=None,
+                prefix=f"{prefix}.self_attn",
+                query_pre_quantized=query_pre_quantized,
+            )
+
+        def forward(
+            self, positions: torch.Tensor, hidden_states: torch.Tensor
+        ) -> torch.Tensor:
+            return self.self_attn(positions, hidden_states)
+
+    vllm_config = VllmConfig(
+        compilation_config=CompilationConfig(
+            mode=CompilationMode.VLLM_COMPILE, cudagraph_mode=CUDAGraphMode.NONE
+        )
+    )
+    torch.manual_seed(0)
+    with (
+        set_current_vllm_config(vllm_config),
+        torch.device("cuda"),
+        set_default_torch_dtype(torch.bfloat16),
+    ):
+        model = DraftAttention(vllm_config=vllm_config, prefix="model.layers.0")
+    attn = model.self_attn
+    for param in (
+        attn.qkv_proj.weight,
+        attn.q_norm.weight,
+        attn.k_norm.weight,
+        attn.o_proj.weight,
+    ):
+        param.data.normal_(std=0.05)
+    layer = attn.attn
+    cache_dtype = torch.uint8 if kv_cache_dtype == "fp8" else torch.bfloat16
+    layer.kv_cache = torch.zeros(
+        4, 2, BLOCK_SIZE, 2 * HEAD_DIM, dtype=cache_dtype, device="cuda"
+    )
+    attn.attn = _QueryAttention(layer.layer_name)
+
+    hidden = torch.randn(7, config.hidden_size, device="cuda", dtype=torch.bfloat16)
+    positions = torch.tensor([0, 1, 127, 128, 4000, 8190, 9], device="cuda")
+    slots = {layer.layer_name: torch.tensor([5, -1, 17, 18, 40, 63, 0], device="cuda")}
+
+    with torch.inference_mode():
+        with set_forward_context(None, vllm_config, slot_mapping=slots):
+            expected_out = attn(positions, hidden)
+        expected_cache = layer.kv_cache.clone()
+        assert expected_cache.any()
+        layer.kv_cache.zero_()
+
+        # The memory-profiling run: no slot mapping, nothing is stored.
+        with set_forward_context(None, vllm_config):
+            model(positions, hidden)
+        assert not layer.kv_cache.any()
+
+        with set_forward_context(None, vllm_config, slot_mapping=slots):
+            out = model(positions, hidden)
+    assert layer.kv_cache.any(), "the compiled draft attention stored no k/v"
+    got, expected, rtol = layer.kv_cache, expected_cache, 2e-2
+    if kv_cache_dtype == "fp8":
+        got = got.view(torch.float8_e4m3fn).float()
+        expected = expected.view(torch.float8_e4m3fn).float()
+        rtol = 0.13
+    torch.testing.assert_close(got, expected, atol=2e-2, rtol=rtol)
+    torch.testing.assert_close(out, expected_out, atol=2e-2, rtol=2e-2)
diff --git a/tests/models/minimax_m3/test_mono_host.py b/tests/models/minimax_m3/test_mono_host.py
new file mode 100644
index 000000000000..df845b274d17
--- /dev/null
+++ b/tests/models/minimax_m3/test_mono_host.py
@@ -0,0 +1,356 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Host-side logic of the MiniMax-M3 mono decode path that needs no GPU.
+
+The kernel package imports flydsl at module import, so these tests run where it
+is installed (the ROCm images); elsewhere they are skipped.
+"""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+pytest.importorskip("flydsl")
+
+from vllm.models.minimax_m3.amd.mono import config as mono_config  # noqa: E402
+from vllm.models.minimax_m3.amd.mono import dispatch as mono_dispatch  # noqa: E402
+from vllm.models.minimax_m3.amd.mono import envs as mono_envs  # noqa: E402
+from vllm.models.minimax_m3.amd.mono.config import (  # noqa: E402
+    MAX_CONTEXT,
+    MAX_TOKENS,
+    SPARSE_BLOCK,
+    TP,
+    MonoUnsupported,
+)
+from vllm.models.minimax_m3.amd.mono.dispatch import (  # noqa: E402
+    _config_refusal,
+    _step_refusal,
+)
+from vllm.models.minimax_m3.amd.mono.runner import (  # noqa: E402
+    dense_prefix,
+    token_rows,
+)
+
+
+def _layer(moe: bool, sparse: bool):
+    attn = (
+        SimpleNamespace(
+            layer_name="attn",
+            indexer=SimpleNamespace(index_cache=SimpleNamespace(prefix="idx")),
+        )
+        if sparse
+        else SimpleNamespace()
+    )
+    return SimpleNamespace(is_moe_layer=moe, self_attn=attn)
+
+
+def test_token_rows_single_query_token_passes_tables_through():
+    bt = torch.arange(12, dtype=torch.int32).view(3, 4)
+    seq = torch.tensor([10, 20, 30], dtype=torch.int32)
+    rows_bt, rows_seq = token_rows(bt, seq, q=1, n=3)
+    assert rows_bt is bt and rows_seq is seq
+
+
+def test_token_rows_expands_speculative_verify_rows():
+    bt = torch.arange(8, dtype=torch.int32).view(2, 4)
+    seq = torch.tensor([10, 20], dtype=torch.int32)
+    rows_bt, rows_seq = token_rows(bt, seq, q=3, n=6)
+    # token j of a request sees seq_len - (q - 1 - j) keys
+    assert rows_seq.tolist() == [8, 9, 10, 18, 19, 20]
+    assert rows_bt.shape == (6, 4) and torch.equal(rows_bt[:3], bt[0].expand(3, 4))
+    assert rows_bt.stride(0) == rows_bt.shape[1]
+
+
+def test_dense_prefix_counts_leading_dense_ffn_layers():
+    layers = [_layer(False, True)] * 3 + [_layer(True, True)] * 5
+    assert dense_prefix(layers) == 3
+    assert dense_prefix([_layer(True, True)] * 2) == 0
+
+
+@pytest.mark.parametrize(
+    "layers",
+    [
+        [_layer(False, True)] * 3,  # no MoE layer at all
+        [
+            _layer(True, True),
+            _layer(False, True),
+            _layer(True, True),
+        ],  # dense after MoE
+        [
+            _layer(False, True),
+            _layer(True, False),
+        ],  # MoE layer without sparse attention
+    ],
+)
+def test_dense_prefix_refuses_other_layer_mixes(layers):
+    with pytest.raises(MonoUnsupported):
+        dense_prefix(layers)
+
+
+def test_configure_accepts_only_the_two_kv_page_strides():
+    saved = (mono_config.KV_PAGE_STRIDE, mono_config.STATIC_KV_SCALE)
+    try:
+        mono_config.configure(kv_page_stride=16, static_kv_scale=True)
+        assert (mono_config.KV_PAGE_STRIDE, mono_config.STATIC_KV_SCALE) == (16, True)
+        with pytest.raises(AssertionError):
+            mono_config.configure(kv_page_stride=4, static_kv_scale=True)
+    finally:
+        mono_config.configure(kv_page_stride=saved[0], static_kv_scale=saved[1])
+
+
+def _vllm_config(**overrides):
+    cg = SimpleNamespace(requires_piecewise_compilation=lambda: False)
+    cfg = SimpleNamespace(
+        parallel_config=SimpleNamespace(
+            tensor_parallel_size=TP,
+            data_parallel_size=1,
+            pipeline_parallel_size=1,
+            enable_expert_parallel=False,
+        ),
+        cache_config=SimpleNamespace(cache_dtype="fp8", block_size=SPARSE_BLOCK),
+        compilation_config=SimpleNamespace(cudagraph_mode=cg),
+        model_config=SimpleNamespace(
+            max_model_len=32768, hf_text_config=SimpleNamespace()
+        ),
+        attention_config=SimpleNamespace(
+            resolve_indexer_kv_dtype=lambda default: "fp8"
+        ),
+    )
+    for path, value in overrides.items():
+        obj, _, attr = path.rpartition(".")
+        setattr(eval(f"cfg.{obj}") if obj else cfg, attr, value)
+    return cfg
+
+
+def _model(fse: bool = True, shared_mlp: bool = True, aux=()):
+    moe = _layer(True, True)
+    moe.block_sparse_moe = SimpleNamespace(
+        is_fused_shared_expert_enabled=fse,
+        shared_experts=object() if shared_mlp else None,
+    )
+    return SimpleNamespace(
+        is_fused_shared_expert_enabled=fse,
+        aux_hidden_state_layers=aux,
+        layers=[_layer(False, True)] * 3 + [moe] * 2,
+    )
+
+
+@pytest.mark.parametrize("fse", [True, False])
+@pytest.mark.parametrize("aux", [(), (2, 30, 57)])
+def test_config_refusal_accepts_the_supported_deployment(fse, aux):
+    assert _config_refusal(_vllm_config(), _model(fse=fse, aux=aux)) is None
+
+
+@pytest.mark.parametrize(
+    "cfg, model, why",
+    [
+        (_vllm_config(**{"parallel_config.tensor_parallel_size": 8}), _model(), "TP"),
+        (_vllm_config(**{"cache_config.cache_dtype": "auto"}), _model(), "kv cache"),
+        (_vllm_config(**{"cache_config.block_size": 64}), _model(), "block size"),
+        (
+            _vllm_config(**{"model_config.max_model_len": MAX_CONTEXT + 1}),
+            _model(),
+            "max_model_len",
+        ),
+        (_vllm_config(), _model(fse=False, shared_mlp=False), "neither fused nor"),
+    ],
+)
+def test_config_refusal_names_the_unmet_condition(cfg, model, why):
+    reason = _config_refusal(cfg, model)
+    assert reason is not None and why in reason
+
+
+def _fwd(n: int, q: int = 1, reqs: int | None = None, prefills: int = 0):
+    reqs = n // q if reqs is None else reqs
+    decode = SimpleNamespace(
+        decode_query_len=q,
+        block_table=torch.zeros(reqs, 8, dtype=torch.int32),
+        seq_lens=torch.ones(reqs, dtype=torch.int32),
+    )
+    md = SimpleNamespace(decode=decode, num_prefills=prefills)
+    return SimpleNamespace(
+        ubatch_slices=None,
+        attn_metadata={"attn": md},
+        slot_mapping={"attn": torch.zeros(n, dtype=torch.int64)},
+    )
+
+
+def test_step_refusal_accepts_a_pure_decode():
+    pos = torch.zeros(4, dtype=torch.int64)
+    assert _step_refusal(_fwd(4), "attn", 4, pos) is None
+    assert _step_refusal(_fwd(6, q=3), "attn", 6, pos) is None
+
+
+def _step_reason(fwd, n: int, positions: torch.Tensor) -> str:
+    reason = _step_refusal(fwd, "attn", n, positions)
+    assert reason is not None
+    return reason
+
+
+def test_step_refusal_rejects_what_the_kernels_cannot_serve():
+    pos = torch.zeros(4, dtype=torch.int64)
+    assert "tokens" in _step_reason(_fwd(MAX_TOKENS + 1), MAX_TOKENS + 1, pos)
+    assert "prefill" in _step_reason(_fwd(4, prefills=1), 4, pos)
+    assert "per request" in _step_reason(_fwd(4, q=3, reqs=2), 4, pos)
+    assert "shorter" in _step_reason(_fwd(4, reqs=2), 4, pos)
+    assert "positions" in _step_reason(_fwd(4), 4, pos.to(torch.int32))
+    fwd = _fwd(4)
+    fwd.attn_metadata["attn"].decode.block_table = fwd.attn_metadata[
+        "attn"
+    ].decode.block_table.t()
+    assert "contiguous" in _step_reason(fwd, 4, pos)
+
+
+# ---- MonoDecode host contract: enable, runner life cycle, slot check, routing ----
+# Everything the kernels are not: these pin the dispatcher's behaviour so a refactor
+# of the host side (or of the switch) must keep passing them.
+
+
+class _FakeRunner:
+    built = 0
+
+    def __init__(self, model, config, *, stale=False, fp8=False):
+        _FakeRunner.built += 1
+        self._stale = stale
+        self.shared_fp8 = fp8
+        self.sparse = []
+        self.rank = 0
+        self.closed = False
+        self.calls = []
+
+    def stale(self):
+        return self._stale
+
+    def close(self):
+        self.closed = True
+
+    def forward(self, input_ids, positions, inputs_embeds):
+        self.calls.append(positions.numel())
+        return "mono-out"
+
+
+@pytest.fixture
+def mono(monkeypatch):
+    """An enabled dispatcher on a supported deployment, outside graph capture,
+    with the runner build stubbed."""
+    monkeypatch.setattr(mono_dispatch.envs, "MONO_ENABLE", True)
+    monkeypatch.setattr(mono_dispatch.envs, "MONO_CHECK", False)
+    monkeypatch.setattr(mono_dispatch, "MonoDecodeRunner", _FakeRunner)
+    monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False)
+    _FakeRunner.built = 0
+    return mono_dispatch.MonoDecode(_model(), _vllm_config())
+
+
+def _fwd_shared(n: int, shared: bool = True):
+    fwd = _fwd(n)
+    slots = fwd.slot_mapping["attn"]
+    fwd.slot_mapping["idx"] = slots.clone() if shared else slots + 1
+    return fwd
+
+
+def test_env_flag_parsing(monkeypatch):
+    for raw, want in (("1", True), ("true", True), ("True", True), ("0", False)):
+        monkeypatch.setenv("VLLM_ROCM_MINIMAX_M3_MONO_TESTFLAG", raw)
+        assert mono_envs._flag("VLLM_ROCM_MINIMAX_M3_MONO_TESTFLAG") is want
+    monkeypatch.delenv("VLLM_ROCM_MINIMAX_M3_MONO_TESTFLAG")
+    assert mono_envs._flag("VLLM_ROCM_MINIMAX_M3_MONO_TESTFLAG") is False
+
+
+def test_disabled_without_the_switch(monkeypatch):
+    monkeypatch.setattr(mono_dispatch.envs, "MONO_ENABLE", False)
+    mono = mono_dispatch.MonoDecode(_model(), _vllm_config())
+    assert not mono._enabled
+    assert not mono.supports(None, torch.zeros(1, dtype=torch.int64), None, None)
+
+
+def test_unsupported_deployment_falls_back(monkeypatch):
+    monkeypatch.setattr(mono_dispatch.envs, "MONO_ENABLE", True)
+    bad = _vllm_config(**{"parallel_config.tensor_parallel_size": 8})
+    assert not mono_dispatch.MonoDecode(_model(), bad)._enabled
+
+
+def test_enabled_dispatcher_reads_the_first_sparse_layer(mono):
+    assert mono._enabled
+    assert mono._layer_name == "attn" and mono._index_prefix == "idx"
+    assert mono._runner is None  # built lazily, on the first eligible step
+
+
+def test_runner_is_built_once_outside_capture(mono):
+    assert mono._ensure_runner() and _FakeRunner.built == 1
+    assert mono._ensure_runner() and _FakeRunner.built == 1
+
+
+def test_runner_is_not_built_inside_capture(mono, monkeypatch):
+    monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True)
+    assert not mono._ensure_runner() and _FakeRunner.built == 0
+    assert mono._enabled  # not a refusal: the next uncaptured step builds it
+
+
+def test_stale_runner_is_rebuilt_after_a_kv_rebind(mono, monkeypatch):
+    first = _FakeRunner(None, None, stale=True)
+    mono._runner, mono._slots_verified = first, True
+    assert mono._ensure_runner()
+    assert first.closed and mono._runner is not first
+    assert not mono._slots_verified  # the new cache's slots are checked again
+
+
+def test_stale_runner_is_kept_while_capturing(mono, monkeypatch):
+    first = _FakeRunner(None, None, stale=True)
+    mono._runner = first
+    monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True)
+    assert not mono._ensure_runner() and not first.closed
+
+
+@pytest.mark.parametrize("exc", [MonoUnsupported("no CUs"), RuntimeError("boom")])
+def test_runner_build_failure_disables_mono(mono, monkeypatch, exc):
+    def boom(model, config):
+        raise exc
+
+    monkeypatch.setattr(mono_dispatch, "MonoDecodeRunner", boom)
+    assert not mono._ensure_runner()
+    assert not mono._enabled and mono._runner is None
+
+
+def test_slot_check_runs_once_and_caches(mono):
+    assert mono._slots_shared(_fwd_shared(4), 4) and mono._slots_verified
+    assert mono._slots_shared(_fwd_shared(4, shared=False), 4)  # cached verdict
+
+
+@pytest.mark.parametrize("missing", [True, False])
+def test_unshared_index_slots_disable_mono(mono, missing):
+    fwd = _fwd_shared(4, shared=False)  # mismatched ids ...
+    if missing:
+        del fwd.slot_mapping["idx"]  # ... or no index slots at all
+    assert not mono._slots_shared(fwd, 4) and not mono._enabled
+
+
+def test_supports_routes_only_pure_decodes(mono, monkeypatch):
+    pos = torch.zeros(4, dtype=torch.int64)
+    monkeypatch.setattr(mono_dispatch, "get_forward_context", lambda: _fwd_shared(4))
+    assert mono.supports(None, pos, None, None) and _FakeRunner.built == 1
+    assert not mono.supports(None, pos, object(), None)  # pipeline-parallel input
+    monkeypatch.setattr(
+        mono_dispatch, "get_forward_context", lambda: _fwd_shared(4, shared=True)
+    )
+    fwd = _fwd(4, prefills=1)
+    monkeypatch.setattr(mono_dispatch, "get_forward_context", lambda: fwd)
+    assert not mono.supports(None, pos, None, None) and mono._enabled  # just declined
+
+
+def test_refused_step_does_not_build_the_runner(mono, monkeypatch):
+    fwd = _fwd(MAX_TOKENS + 1)
+    monkeypatch.setattr(mono_dispatch, "get_forward_context", lambda: fwd)
+    pos = torch.zeros(MAX_TOKENS + 1, dtype=torch.int64)
+    assert not mono.supports(None, pos, None, None) and _FakeRunner.built == 0
+
+
+def test_forward_runs_the_runner_or_the_self_check(mono, monkeypatch):
+    mono._ensure_runner()
+    pos = torch.zeros(2, dtype=torch.int64)
+    assert mono.forward(None, pos, None) == "mono-out" and mono._runner.calls == [2]
+    monkeypatch.setattr(mono_dispatch.envs, "MONO_CHECK", True)
+    monkeypatch.setattr(
+        mono_dispatch, "forward_checked", lambda runner, *a: ("checked", runner)
+    )
+    assert mono.forward(None, pos, None) == ("checked", mono._runner)
diff --git a/tests/models/minimax_m3/test_norm_mxfp8.py b/tests/models/minimax_m3/test_norm_mxfp8.py
new file mode 100644
index 000000000000..774badb4bb82
--- /dev/null
+++ b/tests/models/minimax_m3/test_norm_mxfp8.py
@@ -0,0 +1,115 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""The MiniMax-M3 Triton input norm against the FlashInfer norm (plus the
+reference MXFP8 quantize for MXFP8 output): residual and UE8M0 scales bitwise,
+E4M3 codes within one quantization step, BF16 output within one BF16 ulp."""
+
+import pytest
+import torch
+
+from vllm.platforms import current_platform
+
+pytestmark = pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only")
+
+HIDDEN = 6144
+
+
+def _rows(n: int, k: int, gen: torch.Generator) -> torch.Tensor:
+    """Per-32-block magnitudes over 2^-20..2^20, plus special rows."""
+    x = torch.randn(n, k, generator=gen, device="cuda")
+    e = torch.randint(-20, 21, (n, k // 32, 1), generator=gen, device="cuda")
+    x = (x.view(n, -1, 32) * torch.pow(2.0, e.float())).view(n, k)
+    if n >= 4:
+        x[0] = 0.0
+        x[1] = 1e-40
+        x[2] = float("inf")
+        x[3] = 448.0
+    return x.to(torch.bfloat16)
+
+
+def _code_steps(a: torch.Tensor, b: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:
+    """Per-element code distance in value order; NaN codes must match exactly.
+
+    ``dtype`` is the value dtype: float8_e4m3fn (E4M3 codes) or bfloat16."""
+    n = 256 if dtype == torch.float8_e4m3fn else 65536
+    int_dtype = torch.uint8 if dtype == torch.float8_e4m3fn else torch.uint16
+    codes = torch.arange(n, dtype=torch.int32, device=a.device).to(int_dtype)
+    vals = codes.view(dtype).float()
+    rank = torch.empty(n, dtype=torch.int32, device=a.device)
+    rank[vals.argsort()] = torch.arange(n, dtype=torch.int32, device=a.device)
+    a_i = a.contiguous().view(int_dtype).long()
+    b_i = b.contiguous().view(int_dtype).long()
+    a_nan, b_nan = vals.isnan()[a_i], vals.isnan()[b_i]
+    assert torch.equal(a_nan, b_nan), "NaN pattern mismatch"
+    return torch.where(a_nan, 0, (rank[a_i] - rank[b_i]).abs())
+
+
+@pytest.mark.parametrize("mxfp8", [True, False])
+@pytest.mark.parametrize("fused_add", [True, False])
+@pytest.mark.parametrize(
+    ("num_tokens", "strided"), [(1, False), (129, True), (1024, False)]
+)
+def test_input_norm_matches_reference(num_tokens, strided, fused_add, mxfp8):
+    from flashinfer.norm import gemma_fused_add_rmsnorm, gemma_rmsnorm
+
+    from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
+        mxfp8_e4m3_quantize,
+    )
+    from vllm.models.minimax_m3.nvidia.ops.norm_mxfp8 import gemma_rmsnorm_mxfp8
+
+    gen = torch.Generator(device="cuda").manual_seed(num_tokens)
+    x = _rows(num_tokens, HIDDEN, gen)
+    residual = _rows(num_tokens, HIDDEN, gen)
+    weight = (torch.randn(HIDDEN, generator=gen, device="cuda") * 0.2).to(x.dtype)
+    eps = 1e-6
+
+    x_ref, res_ref = x.clone(), residual.clone()
+    if fused_add:
+        gemma_fused_add_rmsnorm(x_ref, res_ref, weight, eps)
+        y_ref = x_ref
+    else:
+        y_ref = gemma_rmsnorm(x_ref, weight, eps)
+
+    if strided:
+        # Row-contiguous views with row stride 2*H.
+        x_buf = torch.empty(num_tokens, 2 * HIDDEN, dtype=x.dtype, device="cuda")
+        res_buf = torch.empty_like(x_buf)
+        x_buf[:, :HIDDEN], res_buf[:, :HIDDEN] = x, residual
+        x_new, res_new = x_buf[:, :HIDDEN], res_buf[:, :HIDDEN]
+    else:
+        x_new, res_new = x.clone(), residual.clone()
+    out, s = gemma_rmsnorm_mxfp8(
+        x_new, weight, eps, residual=res_new if fused_add else None, mxfp8=mxfp8
+    )
+
+    if fused_add:
+        assert torch.equal(res_new.contiguous(), res_ref)
+    assert torch.equal(x_new.contiguous(), x)
+    if mxfp8:
+        q_ref, s_ref = mxfp8_e4m3_quantize(y_ref, is_sf_swizzled_layout=True)
+        assert torch.equal(s, s_ref)
+        steps = _code_steps(out, q_ref, torch.float8_e4m3fn)
+    else:
+        assert s is None
+        steps = _code_steps(out, y_ref, torch.bfloat16)
+    assert steps.max().item() <= 1
+
+
+def test_norm_module_accepts_3d_per_head():
+    """The dspark drafter norms per-head context K with (tokens, heads,
+    head_dim); the module routes non-2-D inputs to the FlashInfer norm."""
+    from flashinfer.norm import gemma_rmsnorm
+
+    from vllm.models.minimax_m3.nvidia.model import MiniMAXGemmaRMSNorm
+
+    gen = torch.Generator(device="cuda").manual_seed(0)
+    norm = MiniMAXGemmaRMSNorm(128)
+    norm.weight.data = (torch.randn(128, generator=gen, device="cuda") * 0.2).to(
+        torch.bfloat16
+    )
+    x = torch.randn(129, 4, 128, generator=gen, device="cuda").to(torch.bfloat16)
+
+    y_ref = gemma_rmsnorm(x.clone(), norm.weight, norm.variance_epsilon)
+    out = norm(x)
+    assert out.shape == x.shape
+    assert torch.equal(out, y_ref)
diff --git a/tests/models/minimax_m3/test_qkv_index_fused_epilogue.py b/tests/models/minimax_m3/test_qkv_index_fused_epilogue.py
new file mode 100644
index 000000000000..4f4689eb2896
--- /dev/null
+++ b/tests/models/minimax_m3/test_qkv_index_fused_epilogue.py
@@ -0,0 +1,259 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""MiniMax-M3 QKV+indexer MXFP8 GEMM with fused epilogue
+(ops/qkv_index_fused_epilogue.py) vs the production chain: mxfp8_e4m3_quantize
+-> flashinfer mm_mxfp8 -> fused_minimax_m3_qknorm_rope_kv_insert.
+
+Each case runs both tile configs (1-SM and 2-SM), which must match bitwise,
+and compares them to the chain: bf16 q within tolerances, the quantized
+outputs (fp8 q / index_q, NVFP4 caches) code by code. The chain rounds the
+GEMM output to bf16 first, so values near a rounding boundary land one code
+apart.
+"""
+
+import pytest
+import torch
+
+import vllm._custom_ops as ops
+from vllm.platforms import current_platform
+
+pytestmark = pytest.mark.skipif(
+    not (current_platform.is_cuda() and current_platform.has_device_capability(100)),
+    reason="QKV+indexer MXFP8 GEMM is SM100+ CUDA only",
+)
+
+HIDDEN = 6144
+HEAD_DIM = 128
+ROTARY_DIM = 64
+EPS = 1e-6
+MAX_POS = 8192
+Q_FP8_SCALE = 2.0
+# per-rank (q, kv, index_q) heads; N = (nq + 2 * nkv + niq + 1) * HEAD_DIM
+TP_HEADS = {1: (64, 4, 4), 2: (32, 2, 2), 4: (16, 1, 1)}
+# Caches start as 0xFF: never a scale byte (E4M3 scales saturate at 0x7E), so
+# a stray or missing write shows up in the code comparison.
+CACHE_FILL = 0xFF
+
+BF16_RTOL = 2**-6
+BF16_ATOL = 0.05
+# fp8 values near zero (RoPE cancellation, where an E4M3 step is 2^-7..2^-9)
+# can land a few codes apart but stay within 2^-6.
+FP8_MAX_MISMATCH_FRAC = 0.05
+FP8_SMALL_ABS = 2**-6
+NVFP4_MAX_MISMATCH_FRAC = 0.02
+NVFP4_SF_MAX_MISMATCH_FRAC = 0.05
+
+
+def _make_inputs(nq, nkv, niq, M, nb, bs, slots):
+    """MXFP8 activations / weight, norm weights, positions, the RoPE cache and
+    the KV / index-cache slots: random unique with ~3% -1 (no insert), the
+    last M of the cache, or all -1."""
+    from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
+        mxfp8_e4m3_quantize,
+        swizzle_mxfp8_scale,
+    )
+
+    # CPU RNG: the CUDA one's output depends on the GPU's SM count
+    gen = torch.Generator().manual_seed(M)
+
+    def randn(*shape):
+        return torch.randn(*shape, generator=gen).cuda()
+
+    n = (nq + 2 * nkv + niq + 1) * HEAD_DIM
+    # activations with a few outlier channels, like real hidden states
+    x = randn(M, HIDDEN) * 0.5
+    x[:, torch.randperm(HIDDEN, generator=gen)[:16].cuda()] *= 20
+    x = x.bfloat16()
+    w = (randn(n, HIDDEN) * 0.05).bfloat16()
+    xq, xsf = mxfp8_e4m3_quantize(x, is_sf_swizzled_layout=True)
+    wq, wsf = mxfp8_e4m3_quantize(w, is_sf_swizzled_layout=False)
+    wsf = swizzle_mxfp8_scale(wsf, M=n, K=HIDDEN)
+    # Gemma norm weights (scale 1 + w) as in the checkpoint: q / k near 0,
+    # index_q down to -1 (scale near 0), index_k up to ~2.5
+    offset_std = torch.tensor(
+        [[0.1, 0.1], [0.2, 0.1], [-0.4, 0.25], [0.9, 0.5]], device="cuda"
+    )
+    norm_weights = offset_std[:, :1] + offset_std[:, 1:] * randn(4, HEAD_DIM)
+    norm_weights = norm_weights.bfloat16()
+    positions = torch.randint(0, MAX_POS, (M,), generator=gen).cuda()
+    inv_freq = 1.0 / (
+        5e6 ** (torch.arange(0, ROTARY_DIM, 2, device="cuda").float() / ROTARY_DIM)
+    )
+    freqs = torch.outer(torch.arange(MAX_POS, device="cuda").float(), inv_freq)
+    cos_sin = torch.cat((freqs.cos(), freqs.sin()), dim=-1).bfloat16()
+
+    if slots == "none":
+        sm = ism = torch.full((M,), -1, device="cuda")
+    elif slots == "last_page":
+        sm = ism = torch.arange(nb * bs - M, nb * bs, device="cuda")
+    else:
+        sm, ism = (torch.randperm(nb * bs, generator=gen)[:M] for _ in range(2))
+        sm[torch.rand(M, generator=gen) < 0.03] = -1
+        ism[torch.rand(M, generator=gen) < 0.03] = -1
+        sm, ism = sm.cuda(), ism.cuda()
+    return xq, xsf, wq, wsf, norm_weights, positions, cos_sin, sm, ism
+
+
+def _ordered(codes, sign_bit):
+    """Sign-magnitude codes -> integers ordered like their values (+0 == -0),
+    so adjacent representable values differ by 1."""
+    codes = codes.to(torch.int16)
+    mag = codes & (sign_bit - 1)
+    return torch.where(codes & sign_bit != 0, -mag, mag)
+
+
+def _check_fp8(name, got, want):
+    """E4M3 codes: NaN masks exact, at most FP8_MAX_MISMATCH_FRAC differ, each
+    within one code or FP8_SMALL_ABS of the chain's value."""
+    a, b = got.float(), want.float()
+    assert (a.isnan() == b.isnan()).all(), f"{name}: NaN masks differ"
+    fin = ~a.isnan()
+    d = _ordered(got.view(torch.uint8), 0x80) - _ordered(want.view(torch.uint8), 0x80)
+    d = d.abs()[fin]
+    mism = (d > 0).float().mean()
+    assert mism <= FP8_MAX_MISMATCH_FRAC, f"{name}: {mism:.2%} E4M3 codes differ"
+    far = (d > 1) & ((a - b).abs()[fin] > FP8_SMALL_ABS)
+    assert not far.any(), f"{name}: {far.sum()} E4M3 codes off by > 1 code"
+
+
+def _check_cache(name, got, want, slots, bs):
+    """NVFP4 cache [nb, (nslots,) bs, 72] vs the chain's: pages without an
+    inserted token are byte-identical; in the others, E2M1 codes and E4M3
+    scale bytes are each within one code, and few differ. A scale one code
+    apart re-codes its whole group, so E2M1 mismatches only count in groups
+    with equal scales."""
+    pages = torch.zeros(got.shape[0], dtype=torch.bool, device=got.device)
+    pages[slots[slots >= 0] // bs] = True
+    changed = (got != want).flatten(1).any(1)
+    assert not (changed & ~pages).any(), f"{name}: writes to pages with no token"
+
+    def split(c):
+        """Each (page, slot) holds bs tokens' E2M1 data (8 groups of 8B), then
+        their E4M3 group scales (8B; V's swizzled in 4-token quads) -> codes
+        [rows, bs, 8, 16] and scales [rows, bs, 8]."""
+        c = c[pages].reshape(-1, bs * 72)
+        data = c[:, : bs * 64].reshape(-1, bs, 8, 8)
+        codes = _ordered(torch.stack((data & 0xF, data >> 4), -1), 0x8)
+        sf = c[:, bs * 64 :].reshape(-1, bs, 8).to(torch.int16)
+        if got.dim() == 4:  # K/V cache: odd slots are V
+            # byte (t // 4 * 4 + g // 2) * 8 + g % 2 * 4 + t % 4 -> [t, g]
+            k, v = sf.reshape(-1, 2, bs, 8).unbind(1)
+            v = v.reshape(-1, bs // 4, 4, 2, 4).permute(0, 1, 4, 2, 3)
+            sf = torch.stack((k, v.reshape(-1, bs, 8)), 1).reshape(-1, bs, 8)
+        return codes.flatten(-2), sf
+
+    (g, g_sf), (w, w_sf) = split(got), split(want)
+    d, d_sf = (g - w).abs(), (g_sf - w_sf).abs()
+    assert (d <= 1).all() and (d_sf <= 1).all(), f"{name}: codes off by > 1 code"
+    written = (g_sf != CACHE_FILL) | (w_sf != CACHE_FILL)
+    for what, mism, frac in (
+        ("E2M1 codes", d[written & (d_sf == 0)] > 0, NVFP4_MAX_MISMATCH_FRAC),
+        ("scale bytes", d_sf[written] > 0, NVFP4_SF_MAX_MISMATCH_FRAC),
+    ):
+        if mism.numel() == 0:
+            continue
+        # 5 sigma of binomial slack: small M writes only a few groups
+        rate = mism.float().mean().item()
+        bound = frac + 5 * (frac / mism.numel()) ** 0.5
+        assert rate <= bound, f"{name}: {rate:.2%} of written {what} differ"
+
+
+@pytest.mark.parametrize(
+    "case",
+    [
+        pytest.param(dict(tp=1, M=1), id="tp1-m1"),
+        pytest.param(dict(tp=1, M=37), id="tp1-m37"),
+        pytest.param(dict(tp=1, M=1000), id="tp1-m1000"),
+        pytest.param(dict(tp=1, M=4096, q="fp8"), id="tp1-m4096-fp8q"),
+        pytest.param(dict(tp=1, M=32768, q="fp8"), id="tp1-m32768-fp8q"),
+        pytest.param(dict(tp=2, M=1024), id="tp2-m1024"),
+        pytest.param(dict(tp=4, M=1024, q="bf16"), id="tp4-m1024-bf16q"),
+        pytest.param(dict(tp=4, M=4100, q="fp8"), id="tp4-m4100-fp8q"),
+        pytest.param(dict(tp=4, M=1024, kv_scales=(0.5, 1.7)), id="tp4-kv-scales"),
+        pytest.param(dict(tp=4, M=1024, slots="none"), id="tp4-no-insert"),
+        pytest.param(dict(tp=1, M=1024, layers=3), id="tp1-cross-layer"),
+        pytest.param(dict(tp=4, M=1024, block_size=64), id="tp4-bs64"),
+        # > 2**31 bytes: writers in the last page exercise int64 byte offsets
+        pytest.param(
+            dict(tp=1, M=128, nb=30000, slots="last_page"), id="tp1-big-cache"
+        ),
+    ],
+)
+def test_qkv_index_fused_epilogue(case):
+    import flashinfer
+
+    from vllm.models.minimax_m3.nvidia.ops.qkv_index_fused_epilogue import (
+        minimax_m3_qkv_index_fused_epilogue,
+    )
+
+    tp, M = case["tp"], case["M"]
+    q_mode = case.get("q", "both")
+    kv_scales = case.get("kv_scales", (1.0, 1.0))
+    bs, layers = case.get("block_size", 128), case.get("layers", 1)
+    nb = case.get("nb", (M + M // 2 + bs - 1) // bs + 1)
+    nq, nkv, niq = TP_HEADS[tp]
+
+    xq, xsf, wq, wsf, norm_weights, positions, cos_sin, sm, ism = _make_inputs(
+        nq, nkv, niq, M, nb, bs, case.get("slots", "random")
+    )
+
+    def caches():
+        # the last layer of [nb, layers, *page] buffers: page-strided if layers > 1
+        kv = torch.full((nb, layers, 2 * nkv, bs, 72), CACHE_FILL, dtype=torch.uint8)
+        ic = torch.full((nb, layers, bs, 72), CACHE_FILL, dtype=torch.uint8)
+        return kv.cuda()[:, -1], ic.cuda()[:, -1]
+
+    # the production chain
+    kv_ref, ic_ref = caches()
+    q_ref = torch.empty(M, nq * HEAD_DIM, dtype=torch.bfloat16, device="cuda")
+    q8_ref = torch.empty(M, nq * HEAD_DIM, dtype=torch.float8_e4m3fn, device="cuda")
+    iq_ref = torch.empty(M, niq * HEAD_DIM, dtype=torch.float8_e4m3fn, device="cuda")
+    qkv = flashinfer.mm_mxfp8(
+        xq, wq.t(), xsf, wsf, out_dtype=torch.bfloat16, backend="cute-dsl"
+    )
+    ops.fused_minimax_m3_qknorm_rope_kv_insert(
+        qkv, norm_weights[0], norm_weights[1], cos_sin, positions, nq, nkv,
+        ROTARY_DIM, EPS, norm_weights[2], norm_weights[3], niq, sm, ism,
+        kv_ref, ic_ref, bs, q_ref, iq_ref, "nvfp4",
+        q_fp8_out=q8_ref, q_fp8_scale=Q_FP8_SCALE,
+        kv_k_scale=torch.tensor([kv_scales[0]], device="cuda"),
+        kv_v_scale=torch.tensor([kv_scales[1]], device="cuda"),
+    )  # fmt: skip
+
+    outs = []
+    for cta_group in (1, 2):
+        kv, ic = caches()
+        q = torch.empty_like(q_ref) if q_mode != "fp8" else None
+        q8 = torch.empty_like(q8_ref) if q_mode != "bf16" else None
+        iq = torch.empty_like(iq_ref)
+        minimax_m3_qkv_index_fused_epilogue(
+            xq, xsf, wq, wsf, positions, cos_sin, norm_weights[0],
+            norm_weights[1], norm_weights[2], norm_weights[3], EPS,
+            slot_mapping=sm, index_slot_mapping=ism, kv_cache=kv, index_cache=ic,
+            q_out=q, index_q_out=iq, q_fp8_out=q8, q_fp8_scale=Q_FP8_SCALE,
+            k_scale=kv_scales[0], v_scale=kv_scales[1], cta_group=cta_group,
+        )  # fmt: skip
+        outs.append((q, q8, iq, kv, ic))
+
+    # both tile configs write identical bytes
+    for a, b in zip(*outs):
+        if a is not None:
+            assert torch.equal(a.view(torch.uint8), b.view(torch.uint8))
+
+    q, q8, iq, kv, ic = outs[0]
+    if q is not None:
+        torch.testing.assert_close(
+            q, q_ref, rtol=BF16_RTOL, atol=BF16_ATOL, equal_nan=True
+        )
+    if q8 is not None:
+        _check_fp8("q_fp8", q8, q8_ref)
+    _check_fp8("index_q", iq, iq_ref)
+    _check_cache("kv_cache", kv, kv_ref, sm, bs)
+    _check_cache("index_cache", ic, ic_ref, ism, bs)
+    # the other layers of a cross-layer buffer stay untouched
+    for c in (kv, ic):
+        storage = torch.empty(0, dtype=torch.uint8, device="cuda")
+        storage.set_(c.untyped_storage())
+        outside = torch.ones_like(storage, dtype=torch.bool)
+        outside.as_strided(c.size(), c.stride(), c.storage_offset()).fill_(False)
+        assert (storage[outside] == CACHE_FILL).all(), "writes outside the view"
diff --git a/tests/models/multimodal/conftest.py b/tests/models/multimodal/conftest.py
index 54e545d6bf04..c00932738f47 100644
--- a/tests/models/multimodal/conftest.py
+++ b/tests/models/multimodal/conftest.py
@@ -10,4 +10,4 @@
 # longer re-exports the typing aliases, so executing that module raises NameError.
 # TODO: Remove once the checkpoint fixes its imports.
 if not hasattr(F, "List"):
-    F.List = list
+    F.List = list  # type: ignore[attr-defined]
diff --git a/tests/models/multimodal/generation/test_common.py b/tests/models/multimodal/generation/test_common.py
index 9ea2cc43cd6e..083eba8b5639 100644
--- a/tests/models/multimodal/generation/test_common.py
+++ b/tests/models/multimodal/generation/test_common.py
@@ -697,7 +697,7 @@ def _granite4_vision_vllm_to_hf_output(vllm_output, model):
         img_idx_to_prompt=lambda idx: "(<image>./</image>)\n",
         max_model_len=4096,
         max_num_seqs=2,
-        get_stop_token_ids=lambda tok: [tok.eos_id, tok.eot_id],
+        get_stop_token_ids=lambda tok: [getattr(tok, "eos_id"), getattr(tok, "eot_id")],  # noqa: B009
         hf_output_post_proc=model_utils.minicpmv_trunc_hf_output,
         patch_hf_runner=model_utils.minicpmv_25_patch_hf_runner,
     ),
@@ -927,7 +927,7 @@ def _granite4_vision_vllm_to_hf_output(vllm_output, model):
         auto_cls=AutoModelForImageTextToText,
         vllm_output_post_proc=model_utils.llava_image_vllm_to_hf_output,
         marks=multi_gpu_marks(num_gpus=2),
-        **COMMON_BROADCAST_SETTINGS,  # type: ignore
+        **COMMON_BROADCAST_SETTINGS,
     ),
     "llava_next-broadcast": VLMTestInfo(
         models=["llava-hf/llava-v1.6-mistral-7b-hf"],
@@ -936,7 +936,7 @@ def _granite4_vision_vllm_to_hf_output(vllm_output, model):
         auto_cls=AutoModelForImageTextToText,
         vllm_output_post_proc=model_utils.llava_image_vllm_to_hf_output,
         marks=multi_gpu_marks(num_gpus=2),
-        **COMMON_BROADCAST_SETTINGS,  # type: ignore
+        **COMMON_BROADCAST_SETTINGS,
     ),
     ### Custom input edge-cases for specific models
     "intern_vl-diff-patches": VLMTestInfo(
diff --git a/tests/models/multimodal/generation/test_deepseek_ocr.py b/tests/models/multimodal/generation/test_deepseek_ocr.py
new file mode 100644
index 000000000000..6dc9ea52434f
--- /dev/null
+++ b/tests/models/multimodal/generation/test_deepseek_ocr.py
@@ -0,0 +1,58 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Guard tests for DeepSeek-OCR _encode_local_features.
+
+_encode_local_features must skip processing only for an empty-element
+patches tensor, not for a zero-sum one.
+
+Regression for #64075: the old guard ``torch.sum(patches).item() == 0``
+returned None for a 1536x768 half-white/half-black image whose two
+normalised crop patches cancel to sum==0, causing a shape mismatch in
+_merge_multimodal_embeddings and a fatal EngineCore crash.
+
+The correct guard is ``patches.numel() == 0``.
+"""
+
+import pytest
+import torch
+
+from vllm.model_executor.models.deepseek_ocr import DeepseekOCRForCausalLM
+from vllm.model_executor.models.deepseek_ocr2 import DeepseekOCR2ForCausalLM
+
+
+class TestDeepseekOCREncodeLocalFeaturesGuard:
+    def test_ocr2_empty_patches_returns_none(self):
+        """Zero-element tensor → skip local encoding (non-tiled image)."""
+        empty = torch.empty(0, 3, 224, 224)
+        assert DeepseekOCR2ForCausalLM._encode_local_features(None, empty) is None
+
+    def test_ocr2_zero_sum_patches_enters_model_body(self):
+        """Non-empty patches that happen to sum to zero must NOT return None.
+
+        The method will raise AttributeError when it tries to call
+        self.sam_model on a None self – that is the expected outcome here
+        because it proves the guard was passed, not that encoding succeeded.
+        The real crash was a shape-mismatch ValueError that killed EngineCore.
+        """
+        # Two patches that cancel: white (1.0) + black (-1.0) after Normalize.
+        wb = torch.cat([torch.ones(1, 3, 4, 4), -torch.ones(1, 3, 4, 4)])
+        assert wb.numel() > 0
+        assert wb.sum().item() == 0.0
+        with pytest.raises(AttributeError):
+            DeepseekOCR2ForCausalLM._encode_local_features(None, wb)
+
+    def test_ocr1_empty_patches_returns_none(self):
+        """OCR-1: zero-element tensor → skip local encoding."""
+        empty = torch.empty(0, 3, 224, 224)
+        crop_shape = torch.tensor([0, 0])
+        assert (
+            DeepseekOCRForCausalLM._encode_local_features(None, empty, crop_shape)
+            is None
+        )
+
+    def test_ocr1_zero_sum_patches_enters_model_body(self):
+        """OCR-1: non-empty zero-sum patches must pass the guard."""
+        wb = torch.cat([torch.ones(1, 3, 4, 4), -torch.ones(1, 3, 4, 4)])
+        crop_shape = torch.tensor([1, 2])
+        with pytest.raises(AttributeError):
+            DeepseekOCRForCausalLM._encode_local_features(None, wb, crop_shape)
diff --git a/tests/models/multimodal/generation/test_keye.py b/tests/models/multimodal/generation/test_keye.py
deleted file mode 100644
index e9d6626a20d6..000000000000
--- a/tests/models/multimodal/generation/test_keye.py
+++ /dev/null
@@ -1,76 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-from typing import NamedTuple
-
-import pytest
-from PIL.Image import Image
-from transformers import AutoProcessor
-
-from vllm import LLM, EngineArgs, SamplingParams
-from vllm.multimodal.utils import encode_image_url
-
-MODEL_NAME = "Kwai-Keye/Keye-VL-8B-Preview"
-
-QUESTION = "What is the content of each image?"
-
-
-class ModelRequestData(NamedTuple):
-    engine_args: EngineArgs
-    prompt: str
-    image_data: list[Image]
-    stop_token_ids: list[int] | None = None
-    chat_template: str | None = None
-    sampling_params: SamplingParams | None = None
-
-
-@pytest.mark.parametrize("question", [QUESTION])
-def test_keye_vl(image_assets, question: str):
-    images = [asset.pil_image for asset in image_assets]
-    image_urls = [encode_image_url(image) for image in images]
-
-    placeholders = [{"type": "image", "image": url} for url in image_urls]
-    messages = [
-        {
-            "role": "user",
-            "content": [
-                *placeholders,
-                {"type": "text", "text": question},
-            ],
-        },
-    ]
-
-    processor = AutoProcessor.from_pretrained(MODEL_NAME, trust_remote_code=True)
-
-    prompt = processor.apply_chat_template(
-        messages, tokenize=False, add_generation_prompt=True
-    )
-
-    llm = LLM(
-        model=MODEL_NAME,
-        trust_remote_code=True,
-        max_model_len=8192,
-        max_num_seqs=5,
-        limit_mm_per_prompt={"image": len(image_urls)},
-        seed=42,
-    )
-
-    sampling_params = SamplingParams(
-        temperature=0.0, max_tokens=256, stop_token_ids=None
-    )
-
-    outputs = llm.generate(
-        {
-            "prompt": prompt,
-            "multi_modal_data": {"image": images},
-        },
-        sampling_params=sampling_params,
-    )
-
-    print("-" * 50)
-    for o in outputs:
-        generated_text = o.outputs[0].text
-        print(generated_text)
-        assert len(generated_text) > 10, (
-            f"Generated text is too short: {generated_text}"
-        )
-        print("-" * 50)
diff --git a/tests/models/multimodal/generation/test_maverick.py b/tests/models/multimodal/generation/test_maverick.py
index 0f6bec107fc2..a44491dfec31 100644
--- a/tests/models/multimodal/generation/test_maverick.py
+++ b/tests/models/multimodal/generation/test_maverick.py
@@ -523,7 +523,7 @@ def save_weights_to_safetensors(
             print(f"Saved shard {i + 1}/{len(shards)}: {filename}")
 
     # Create index file
-    index_data = {
+    index_data: dict[str, Any] = {
         "metadata": {
             "total_size": sum(
                 tensor.numel() * tensor.element_size() for tensor in weights.values()
diff --git a/tests/models/multimodal/generation/test_nemotron_parse_35.py b/tests/models/multimodal/generation/test_nemotron_parse_35.py
index bb8fbf51d937..f5b9884df4c7 100644
--- a/tests/models/multimodal/generation/test_nemotron_parse_35.py
+++ b/tests/models/multimodal/generation/test_nemotron_parse_35.py
@@ -51,7 +51,9 @@ def generate_hf_logprobs_from_logits(
                 }
             )
         token_ids = generated.sequences[0, -len(logprobs) :].tolist()
-        outputs.append((token_ids, hf_model.tokenizer.decode(token_ids), logprobs))
+        decoded = hf_model.tokenizer.decode(token_ids)
+        assert isinstance(decoded, str)
+        outputs.append((token_ids, decoded, logprobs))
     return outputs
 
 
diff --git a/tests/models/multimodal/generation/test_nemotron_parse_35_mtp.py b/tests/models/multimodal/generation/test_nemotron_parse_35_mtp.py
index 918bc51ea657..55d5108ee09b 100644
--- a/tests/models/multimodal/generation/test_nemotron_parse_35_mtp.py
+++ b/tests/models/multimodal/generation/test_nemotron_parse_35_mtp.py
@@ -23,10 +23,11 @@
 
 
 def test_parse_mtp_config_routes_only_causal_prefix_architecture() -> None:
-    config = PreTrainedConfig(
-        architectures=["RadioMBartPrefixForCausalLM"],
-        decoder=PreTrainedConfig(causal_vision_prefix=True, vocab_size=8, d_model=4),
-    )
+    config = PreTrainedConfig(architectures=["RadioMBartPrefixForCausalLM"])
+    config.decoder = PreTrainedConfig()
+    config.decoder.causal_vision_prefix = True
+    config.decoder.vocab_size = 8
+    config.decoder.d_model = 4
 
     draft = SpeculativeConfig.hf_config_override(config)
 
diff --git a/tests/models/multimodal/generation/test_qwen2_vl.py b/tests/models/multimodal/generation/test_qwen2_vl.py
index 4508135b2bb9..ea74d117b302 100644
--- a/tests/models/multimodal/generation/test_qwen2_vl.py
+++ b/tests/models/multimodal/generation/test_qwen2_vl.py
@@ -126,6 +126,7 @@ def batch_make_image_embeddings(
     image_processor = processor.image_processor
 
     mm_config = llm.get_llm().model_config.multimodal_config
+    assert mm_config is not None
     mm_kwargs = mm_config.merge_mm_processor_kwargs({})
     preprocess_result = image_processor.preprocess(
         images=images, return_tensors="pt", **mm_kwargs
@@ -211,6 +212,7 @@ def batch_make_video_embeddings(
     video_processor = processor.video_processor
 
     mm_config = llm.get_llm().model_config.multimodal_config
+    assert mm_config is not None
     mm_kwargs = mm_config.merge_mm_processor_kwargs({})
     preprocess_result = video_processor.preprocess(
         videos=videos, return_tensors="pt", **mm_kwargs
diff --git a/tests/models/multimodal/generation/test_ultravox.py b/tests/models/multimodal/generation/test_ultravox.py
index 72f473835800..8ae5e573b4fb 100644
--- a/tests/models/multimodal/generation/test_ultravox.py
+++ b/tests/models/multimodal/generation/test_ultravox.py
@@ -24,7 +24,7 @@
 
 MULTI_AUDIO_PROMPT = "Describe each of the audios above."
 
-AudioTuple = tuple[np.ndarray, int]
+AudioTuple = tuple[np.ndarray, float]
 
 VLLM_PLACEHOLDER = "<|audio|>"
 HF_PLACEHOLDER = "<|audio|>"
@@ -92,7 +92,7 @@ def _get_prompt(audio_count, question, placeholder):
 
 def run_multi_audio_test(
     vllm_runner: type[VllmRunner],
-    prompts_and_audios: list[tuple[str, list[AudioTuple]]],
+    prompts_and_audios: list[tuple[str | list[int], list[AudioTuple]]],
     model: str,
     *,
     dtype: str,
diff --git a/tests/models/multimodal/generation/test_vit_cudagraph.py b/tests/models/multimodal/generation/test_vit_cudagraph.py
index e80dd9dc8515..8751a94fe690 100644
--- a/tests/models/multimodal/generation/test_vit_cudagraph.py
+++ b/tests/models/multimodal/generation/test_vit_cudagraph.py
@@ -34,7 +34,7 @@ class VitCudagraphTestConfig:
 
 def params_with_marks(
     configs: dict[str, VitCudagraphTestConfig],
-) -> list[pytest.param]:
+) -> list[object]:
     return [
         pytest.param(model_id, marks=cfg.marks) for model_id, cfg in configs.items()
     ]
@@ -520,6 +520,7 @@ def test_vit_cudagraph_video(model_id, vllm_runner, video_assets):
             "baby_reading": config.video_prompt,  # type: ignore[typeddict-item]
         }
     )
+    sampled_vids: list
     if config.needs_video_metadata:
         sampled_vids = [
             sample_frames_with_video_metadata(
diff --git a/tests/models/multimodal/generation/test_voxtral.py b/tests/models/multimodal/generation/test_voxtral.py
index 3c9d2ace8cca..434df130a080 100644
--- a/tests/models/multimodal/generation/test_voxtral.py
+++ b/tests/models/multimodal/generation/test_voxtral.py
@@ -58,7 +58,7 @@ def test_models_with_multiple_audios(
     vllm_prompt = _get_prompt(audio_assets, MULTI_AUDIO_PROMPT)
     run_multi_audio_test(
         vllm_runner,
-        [(vllm_prompt, [a.audio_and_sample_rate for a in audio_assets])],  # type: ignore[list-item]
+        [(vllm_prompt, [a.audio_and_sample_rate for a in audio_assets])],
         MODEL_NAME,
         dtype=dtype,
         max_tokens=max_tokens,
diff --git a/tests/models/multimodal/generation/vlm_utils/builders.py b/tests/models/multimodal/generation/vlm_utils/builders.py
index 1b7e2347be2f..e73c69c13897 100644
--- a/tests/models/multimodal/generation/vlm_utils/builders.py
+++ b/tests/models/multimodal/generation/vlm_utils/builders.py
@@ -224,7 +224,7 @@ def build_embedding_inputs_from_test_info(
     images = [asset.pil_image for asset in image_assets]
     embeds = test_info.convert_assets_to_embeddings(image_assets)
     if test_info.dtype != "auto":
-        dtype = getattr(torch, test_info.dtype)  # type: ignore
+        dtype = getattr(torch, test_info.dtype)
         embeds = [e.to(dtype=dtype) for e in embeds]
     assert len(images) == len(model_prompts)
 
diff --git a/tests/models/multimodal/generation/vlm_utils/model_utils.py b/tests/models/multimodal/generation/vlm_utils/model_utils.py
index 9acbe2c0d6b3..f201d5cc2137 100644
--- a/tests/models/multimodal/generation/vlm_utils/model_utils.py
+++ b/tests/models/multimodal/generation/vlm_utils/model_utils.py
@@ -22,6 +22,7 @@
     GenerationMixin,
 )
 from transformers.masking_utils import create_causal_mask
+from transformers.utils import ModelOutput
 from transformers.video_utils import VideoMetadata
 
 from vllm.logprobs import SampleLogprobs
@@ -105,7 +106,9 @@ def _llava_vllm_to_hf_output(
     hf_output_str = output_str[1:] if output_str[0] == " " else output_str
 
     if hf_output_ids[-1] == eos_token_id:
-        hf_output_str = hf_output_str + tokenizer.decode(eos_token_id)
+        decoded_eos = tokenizer.decode(eos_token_id)
+        assert isinstance(decoded_eos, str)
+        hf_output_str = hf_output_str + decoded_eos
 
     return hf_output_ids, hf_output_str, out_logprobs
 
@@ -137,7 +140,9 @@ def llava_onevision_vllm_to_hf_output(
 
     hf_output_str = output_str
     if hf_output_ids[-1] == eos_token_id:
-        hf_output_str = hf_output_str + tokenizer.decode(eos_token_id)
+        decoded_eos = tokenizer.decode(eos_token_id)
+        assert isinstance(decoded_eos, str)
+        hf_output_str = hf_output_str + decoded_eos
 
     return hf_output_ids, hf_output_str, out_logprobs
 
@@ -179,7 +184,9 @@ def paligemma_vllm_to_hf_output(vllm_output: RunnerOutput, model: str) -> Runner
     hf_output_str = output_str
 
     if hf_output_ids[-1] == eos_token_id:
-        hf_output_str = hf_output_str + tokenizer.decode(eos_token_id)
+        decoded_eos = tokenizer.decode(eos_token_id)
+        assert isinstance(decoded_eos, str)
+        hf_output_str = hf_output_str + decoded_eos
 
     return hf_output_ids, hf_output_str, out_logprobs
 
@@ -286,7 +293,9 @@ def gemma3_vllm_to_hf_output(vllm_output: RunnerOutput, model: str) -> RunnerOut
 
     hf_output_str = output_str
     if hf_output_ids[-1] == eos_token_id:
-        hf_output_str = hf_output_str + tokenizer.decode(eos_token_id)
+        decoded_eos = tokenizer.decode(eos_token_id)
+        assert isinstance(decoded_eos, str)
+        hf_output_str = hf_output_str + decoded_eos
 
     return hf_output_ids, hf_output_str, out_logprobs
 
@@ -572,11 +581,10 @@ def patched_forward(
         # Prepare attention mask
         attention_mask = create_causal_mask(
             config=self.config,
-            input_embeds=inputs_embeds,
+            inputs_embeds=inputs_embeds,
             attention_mask=attention_mask,
             past_key_values=past_key_values,
             position_ids=position_ids,
-            cache_position=cache_position,
         )
 
         # Initialize and collect hidden states
@@ -800,7 +808,7 @@ def _internvl_generate(
     input_ids: torch.FloatTensor,
     attention_mask: torch.LongTensor | None = None,
     **generate_kwargs,
-) -> torch.LongTensor:
+) -> torch.Tensor | ModelOutput:
     """Generate method for InternVL2 model without fixed use_cache."""
     assert self.img_context_token_id is not None
     target_dtype = next(self.parameters()).dtype
@@ -1207,7 +1215,7 @@ def patched_generate(*args, **kwargs):
         return output
 
     hf_model.get_inputs = patched_get_inputs  # type: ignore[method-assign]
-    hf_model.model.generate = patched_generate  # type: ignore[method-assign]
+    hf_model.model.generate = patched_generate
     return hf_model
 
 
diff --git a/tests/models/multimodal/generation_ppl_test/ppl_utils.py b/tests/models/multimodal/generation_ppl_test/ppl_utils.py
index c612eb7141da..d19764740995 100644
--- a/tests/models/multimodal/generation_ppl_test/ppl_utils.py
+++ b/tests/models/multimodal/generation_ppl_test/ppl_utils.py
@@ -17,7 +17,7 @@
 PPL_TOL = 0.01
 
 
-@torch.inference_mode
+@torch.inference_mode()
 def vqa_ppl_test(
     hf_runner,
     vllm_runner,
diff --git a/tests/models/multimodal/pooling/test_colultravbert.py b/tests/models/multimodal/pooling/test_colultravbert.py
index 81deedcb0334..d26d5011c715 100644
--- a/tests/models/multimodal/pooling/test_colultravbert.py
+++ b/tests/models/multimodal/pooling/test_colultravbert.py
@@ -9,6 +9,8 @@
 from vllm.model_executor.models.colultravbert import get_capped_longest_edge
 
 MODEL_NAME = "athrael-soju/ColUltraVBERT"
+# Later revisions declare the UltraVBertModel architecture instead.
+REVISION = "509adbab61b71da0484b20f3baf718e15219d761"
 DTYPE = "bfloat16"
 PAGE_PROMPT = "User:<image>Describe the image.<end_of_utterance>\nAssistant:"
 IMAGE_SEQ_LEN = 64
@@ -53,6 +55,7 @@ def test_colultravbert_tile_cap(width, height, grid):
 def colultravbert_model(vllm_runner):
     with vllm_runner(
         MODEL_NAME,
+        revision=REVISION,
         runner="pooling",
         dtype=DTYPE,
         max_model_len=2048,
diff --git a/tests/models/multimodal/pooling/test_dse_qwen2_vl.py b/tests/models/multimodal/pooling/test_dse_qwen2_vl.py
index 0fff8f40ac74..8b975c0ada48 100644
--- a/tests/models/multimodal/pooling/test_dse_qwen2_vl.py
+++ b/tests/models/multimodal/pooling/test_dse_qwen2_vl.py
@@ -180,7 +180,7 @@ def test_models_text(
         hf_runner,
         vllm_runner,
         input_texts,
-        input_images,  # type: ignore
+        input_images,
         embed_texts,
         model,
         dtype=dtype,
diff --git a/tests/models/multimodal/pooling/test_embedding_gemma2.py b/tests/models/multimodal/pooling/test_embedding_gemma2.py
index 0dcb72406da5..ff4f7e9fb167 100644
--- a/tests/models/multimodal/pooling/test_embedding_gemma2.py
+++ b/tests/models/multimodal/pooling/test_embedding_gemma2.py
@@ -9,6 +9,8 @@
 import pytest
 import torch
 
+from vllm.multimodal.video import EmbeddingGemma2VideoBackend
+
 MODEL = os.environ.get("EG2_MODEL_PATH")
 needs_model = pytest.mark.skipif(not MODEL, reason="EG2_MODEL_PATH not set")
 TEXT_COS, MM_COS = 0.999, 0.999
@@ -348,6 +350,7 @@ def test_video_loader_real_video(vllm_runner, st_model):
 
     video_path = "/projects/gemma4-vllm/videos/test_intel.mp4"
     backend_cls = VIDEO_LOADER_REGISTRY.name2class["embedding_gemma2"]
+    assert issubclass(backend_cls, EmbeddingGemma2VideoBackend)
     with open(video_path, "rb") as f:
         data = f.read()
 
@@ -396,6 +399,7 @@ def test_video_loader_duplicate_expansion_e2e(vllm_runner, st_model):
     from vllm.multimodal.video import VIDEO_LOADER_REGISTRY
 
     backend_cls = VIDEO_LOADER_REGISTRY.name2class["embedding_gemma2"]
+    assert issubclass(backend_cls, EmbeddingGemma2VideoBackend)
     # 4s clip at 0.5 fps = 2 unique frames
     lowfps_path = "/tmp/test_lowfps_05.mp4"
     _ensure_synthetic_video(lowfps_path, duration=4.0, fps=0.5, width=64, height=48)
@@ -436,6 +440,7 @@ def test_video_loader_indices():
     )
 
     backend_cls = VIDEO_LOADER_REGISTRY.name2class["embedding_gemma2"]
+    assert issubclass(backend_cls, EmbeddingGemma2VideoBackend)
     hf_proc = EmbeddingGemma2VideoProcessor()
 
     # Case 1: clip shorter than 32s (e.g. 5s @ 24fps = 120 frames)
@@ -521,6 +526,7 @@ def test_video_loader_duplicate_expansion():
     )
 
     backend_cls = VIDEO_LOADER_REGISTRY.name2class["embedding_gemma2"]
+    assert issubclass(backend_cls, EmbeddingGemma2VideoBackend)
     hf_proc = EmbeddingGemma2VideoProcessor()
 
     # 4s clip @ 0.5fps = 2 frames total. Requested target fps = 1.0 -> 4 frames.
diff --git a/tests/models/multimodal/pooling/test_intern_vit.py b/tests/models/multimodal/pooling/test_intern_vit.py
index 4d5e8728589c..7b5db69c6b9a 100644
--- a/tests/models/multimodal/pooling/test_intern_vit.py
+++ b/tests/models/multimodal/pooling/test_intern_vit.py
@@ -49,14 +49,14 @@ def _patched_init(self, *args, **kwargs):
         _orig_init(self, *args, **kwargs)
         self.all_tied_weights_keys = {}
 
-    PreTrainedModel.__init__ = _patched_init
+    PreTrainedModel.__init__ = _patched_init  # type: ignore[method-assign]
 
     try:
         hf_model = AutoModel.from_pretrained(
             model, dtype=torch_dtype, trust_remote_code=True
         ).to(DEVICE_TYPE)
     finally:
-        PreTrainedModel.__init__ = _orig_init
+        PreTrainedModel.__init__ = _orig_init  # type: ignore[method-assign]
     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_llava_next.py b/tests/models/multimodal/pooling/test_llava_next.py
index 97bdaea4ecbe..a329ce4fece3 100644
--- a/tests/models/multimodal/pooling/test_llava_next.py
+++ b/tests/models/multimodal/pooling/test_llava_next.py
@@ -126,7 +126,7 @@ def test_models_text(
         hf_runner,
         vllm_runner,
         input_texts,
-        input_images,  # type: ignore
+        input_images,
         model,
         dtype=dtype,
     )
diff --git a/tests/models/multimodal/pooling/test_phi3v.py b/tests/models/multimodal/pooling/test_phi3v.py
index 35df42127eca..e7442177c3d0 100644
--- a/tests/models/multimodal/pooling/test_phi3v.py
+++ b/tests/models/multimodal/pooling/test_phi3v.py
@@ -11,7 +11,6 @@
 from vllm.assets.image import VLM_IMAGES_DIR
 from vllm.config import ModelConfig
 from vllm.multimodal import MULTIMODAL_REGISTRY
-from vllm.platforms import current_platform
 
 from ....conftest import IMAGE_ASSETS, HfRunner, PromptImageInput, VllmRunner
 from ....utils import large_gpu_test
@@ -103,15 +102,6 @@ def _run_test(
 
 
 @pytest.mark.core_model
-@pytest.mark.skipif(
-    current_platform.is_xpu(),
-    reason=(
-        "vllm-xpu-kernels paged_decode_default.conf is missing the "
-        "8,96,64,false,true,false tuple needed by VLM2Vec-Full at "
-        "block_size=64. Re-enable once vllm-xpu-kernels ships the updated "
-        "config. See github.com/vllm-project/vllm-xpu-kernels/pull/587."
-    ),
-)
 @pytest.mark.parametrize("model", MODELS)
 @pytest.mark.parametrize("dtype", ["half"])
 def test_models_text(
@@ -129,7 +119,7 @@ def test_models_text(
         hf_runner,
         vllm_runner,
         input_texts,
-        input_images,  # type: ignore
+        input_images,
         model,
         dtype=dtype,
     )
diff --git a/tests/models/multimodal/pooling/test_radio.py b/tests/models/multimodal/pooling/test_radio.py
index 7c039ab8b9c6..da8a6aa903a1 100644
--- a/tests/models/multimodal/pooling/test_radio.py
+++ b/tests/models/multimodal/pooling/test_radio.py
@@ -1,8 +1,14 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import inspect
+import tempfile
+from pathlib import Path
+from typing import Any
+
 import pytest
 import torch
 import torch.nn as nn
+from safetensors.torch import load_file
 from transformers import AutoConfig, AutoModel, CLIPImageProcessor
 
 from vllm.distributed import cleanup_dist_env_and_memory
@@ -33,13 +39,8 @@ def run_radio_test(
 
     img_processor = CLIPImageProcessor.from_pretrained(model)
     images = [asset.pil_image for asset in image_assets]
-    # Input resolution must be a multiple of `self.min_resolution_step`.
-    # Using `self.get_nearest_supported_resolution`, for assets 432x642 the
-    # nearest supported resolution is 432x640.
     pixel_values = [
-        img_processor(image, return_tensors="pt").pixel_values.to(torch_dtype)[
-            :, :, :, :640
-        ]
+        img_processor(image, return_tensors="pt").pixel_values.to(torch_dtype)
         for image in images
     ]
 
@@ -64,6 +65,22 @@ def run_radio_test(
     # affects evaluation benchmarks.
     hf_model.make_preprocessor_external()
 
+    # Input resolution must be a multiple of the tower's ``min_resolution_step``
+    # (patch_size * window_size), which differs across RADIO generations. Crop
+    # down to the nearest supported resolution so one test body fits every
+    # checkpoint; for C-RADIOv2-H this is a no-op (432x640, as before).
+    radio = getattr(hf_model, "radio_model", hf_model)
+    step = radio.min_resolution_step
+    pixel_values = [
+        pv[
+            :,
+            :,
+            : (pv.shape[-2] // step) * step,
+            : (min(pv.shape[-1], 640) // step) * step,
+        ]
+        for pv in pixel_values
+    ]
+
     hf_outputs_per_image = [
         hf_model(pixel_value.to(DEVICE_TYPE)) for pixel_value in pixel_values
     ]
@@ -73,7 +90,17 @@ def run_radio_test(
         **hf_config.args,
     )
     vllm_model = RadioModel(vllm_config)
-    vllm_model.load_weights(hf_model.state_dict())
+    loaded = vllm_model.load_weights(hf_model.state_dict())
+    # Guard the remote-code -> vLLM key remap: every parameter must be loaded
+    # from the checkpoint, except the LayerScale gains (layer_scale{1,2}.lambda1),
+    # which are identity-init and may be absent from the remote-code checkpoint.
+    expected = {
+        name
+        for name, _ in vllm_model.named_parameters()
+        if not name.endswith((".layer_scale1.lambda1", ".layer_scale2.lambda1"))
+    }
+    missing = expected - loaded
+    assert not missing, f"parameters not loaded from checkpoint: {sorted(missing)}"
     vllm_model = vllm_model.to(DEVICE_TYPE, torch_dtype)
 
     vllm_outputs_per_image = [
@@ -85,14 +112,22 @@ def run_radio_test(
 
     cos_similar = nn.CosineSimilarity(dim=-1)
     for vllm_output, hf_output in zip(vllm_outputs_per_image, hf_outputs_per_image):
-        assert cos_similar(vllm_output[0], hf_output[0]).mean() > 0.99
-        assert cos_similar(vllm_output[1], hf_output[1]).mean() > 0.99
+        summary_sim = cos_similar(vllm_output[0], hf_output[0]).mean()
+        feature_sim = cos_similar(vllm_output[1], hf_output[1]).mean()
+        print(
+            f"[radio] {model_id} {dtype} summary={summary_sim:.6f} "
+            f"features={feature_sim:.6f}"
+        )
+        assert summary_sim > 0.99
+        assert feature_sim > 0.99
 
 
 @pytest.mark.parametrize(
     "model_id",
     [
         "nvidia/C-RADIOv2-H",
+        "nvidia/C-RADIOv3-H",
+        "nvidia/C-RADIOv4-H",
     ],
 )
 @pytest.mark.parametrize("dtype", ["half", "bfloat16"])
@@ -104,3 +139,168 @@ def test_radio(
         model_id,
         dtype=dtype,
     )
+
+
+def _radio_summary(*, teachers, cls_token_per_teacher):
+    """Summary tensor from a tiny RadioModel with no encoder layers, so it
+    builds no parallel layers and needs neither a GPU nor a checkpoint."""
+    config = RadioConfig(
+        model_name="vit_small_patch16_224",
+        teachers=teachers,
+        cls_token_per_teacher=cls_token_per_teacher,
+    )
+    model = RadioModel(config, num_hidden_layers_override=0)
+    # Synthetic encoder output: [batch, num_skip + num_patches, hidden].
+    y = torch.randn(2, model.embeddings.num_skip + 4, config.hidden_size)
+    summary, _ = model._extract_final(y)
+    return summary
+
+
+def test_summary_idxs_no_teachers_keeps_class_tokens():
+    # No teachers -> summary_idxs is None -> keep the class-token summary.
+    summary = _radio_summary(teachers=[], cls_token_per_teacher=False)
+    assert summary.shape[1] > 0
+
+
+def test_summary_idxs_all_use_summary_false_is_empty():
+    # Teachers present but none flagged use_summary -> empty selection.
+    summary = _radio_summary(
+        teachers=[{"name": "a", "use_summary": False}],
+        cls_token_per_teacher=True,
+    )
+    assert summary.shape[1] == 0
+
+
+def _native_state_dict_from_model(model: RadioModel) -> dict[str, torch.Tensor]:
+    """Reverse-map a model's parameters onto native Transformers checkpoint keys
+    (``encoder.layer.N.*`` with split ``attention.attention.{query,key,value}``
+    and ``attention.output.dense``), so loading it must round-trip back to the
+    original parameters."""
+    state_dict: dict[str, torch.Tensor] = {}
+    for name, param in model.named_parameters():
+        data = param.detach().clone()
+        # The module tree already mirrors the native ``encoder.layer.N.*`` names
+        # (including ``layer_scale{1,2}.lambda1``); only the fused attention
+        # projections differ from the native split ones.
+        if ".attention.qkv." in name:
+            base, suffix = name.split(".attention.qkv.")
+            query, key, value = data.chunk(3, dim=0)
+            for proj, shard in (("query", query), ("key", key), ("value", value)):
+                state_dict[f"{base}.attention.attention.{proj}.{suffix}"] = (
+                    shard.clone()
+                )
+        elif ".attention.proj." in name:
+            base, suffix = name.split(".attention.proj.")
+            state_dict[f"{base}.attention.output.dense.{suffix}"] = data
+        else:
+            state_dict[name] = data
+    return state_dict
+
+
+def test_native_format_weight_loading(default_vllm_config, dist_init):
+    # Native Transformers checkpoints nest the encoder under ``encoder.layer.N``
+    # with split attention projections; loading must fuse q/k/v into the packed
+    # ``qkv`` in the right order and land every other tensor in place, matching
+    # the legacy layout the integration test covers.
+    config = RadioConfig(model_name="vit_small_patch16_224")
+    model = RadioModel(config, num_hidden_layers_override=2)
+    # Uninitialized parameters may hold NaN, which never compares equal.
+    with torch.no_grad():
+        for param in model.parameters():
+            param.normal_()
+
+    # Every parameter (including layer_scale) must load from the native keys.
+    expected = {
+        name: param.detach().clone() for name, param in model.named_parameters()
+    }
+    native_weights = _native_state_dict_from_model(model)
+    # Keys both layouts intentionally drop must be skipped, not error.
+    native_weights["input_conditioner.norm_mean"] = torch.zeros(3)
+    native_weights["input_conditioner.norm_std"] = torch.ones(3)
+    native_weights["summary_idxs"] = torch.zeros(1, dtype=torch.long)
+
+    # Zero every parameter so a weight that fails to load is visible as zeros.
+    with torch.no_grad():
+        for param in model.parameters():
+            param.zero_()
+
+    loaded = model.load_weights(native_weights)
+
+    missing = set(expected) - loaded
+    assert not missing, f"native weights not loaded: {sorted(missing)}"
+    params = dict(model.named_parameters())
+    for name, original in expected.items():
+        assert torch.equal(params[name], original), (
+            f"parameter mismatch after native load: {name}"
+        )
+
+
+def test_native_format_from_transformers(default_vllm_config, dist_init):
+    # Real-checkpoint native-format round-trip: load a published (legacy)
+    # checkpoint through Transformers' native ``RadioModel`` (which converts it to
+    # the native layout), re-save it with ``save_original_format=False`` so the
+    # on-disk checkpoint stays native -- the default re-applies the reverse
+    # native->legacy mapping for backward compatibility -- then confirm vLLM's
+    # ``RadioModel`` loads that native checkpoint with values matching the
+    # reference. ``nvidia/C-RADIOv4-H`` loads on Transformers' config defaults and
+    # is already downloaded by ``test_radio``.
+    import transformers
+
+    hf_radio_cls: Any = getattr(transformers, "RadioModel", None)
+    if hf_radio_cls is None:
+        pytest.skip("this transformers version has no native RadioModel")
+
+    model_id = "nvidia/C-RADIOv4-H"
+    hf_native = hf_radio_cls.from_pretrained(model_id, dtype=torch.float32)
+    if (
+        "save_original_format"
+        not in inspect.signature(hf_native.save_pretrained).parameters
+    ):
+        pytest.skip("transformers save_pretrained has no save_original_format flag")
+
+    with tempfile.TemporaryDirectory() as tmp:
+        hf_native.save_pretrained(tmp, save_original_format=False)
+        native_weights: dict[str, torch.Tensor] = {}
+        for shard in sorted(Path(tmp).glob("*.safetensors")):
+            native_weights.update(load_file(str(shard)))
+
+    # Guard that the saved checkpoint is genuinely native-layout (split attention
+    # under ``encoder.layer.N``), not the legacy ``radio_model.model.blocks.*``.
+    assert any(".attention.q_proj." in name for name in native_weights), (
+        "expected native-layout keys in the save_original_format=False checkpoint"
+    )
+    assert not any(name.startswith("radio_model.") for name in native_weights)
+
+    hf_config = AutoConfig.from_pretrained(model_id, trust_remote_code=True)
+    vllm_config = RadioConfig(model_name=hf_config.args["model"], **hf_config.args)
+    vllm_model = RadioModel(vllm_config)
+    loaded = vllm_model.load_weights(native_weights)
+
+    expected = {
+        name
+        for name, _ in vllm_model.named_parameters()
+        if not name.endswith((".layer_scale1.lambda1", ".layer_scale2.lambda1"))
+    }
+    missing = expected - loaded
+    assert not missing, f"native weights not loaded: {sorted(missing)}"
+
+    # Values match the reference: the fused qkv is the concatenation of the
+    # native split q/k/v_proj, and proj is the native o_proj.
+    for name, param in vllm_model.named_parameters():
+        if ".attention.qkv." in name:
+            base, suffix = name.split(".attention.qkv.")
+            ref = torch.cat(
+                [
+                    native_weights[f"{base}.attention.{proj}_proj.{suffix}"]
+                    for proj in ("q", "k", "v")
+                ],
+                dim=0,
+            )
+        elif ".attention.proj." in name:
+            base, suffix = name.split(".attention.proj.")
+            ref = native_weights[f"{base}.attention.o_proj.{suffix}"]
+        elif name in native_weights:
+            ref = native_weights[name]
+        else:
+            continue
+        assert torch.equal(param, ref), f"value mismatch after native load: {name}"
diff --git a/tests/models/multimodal/processing/test_audio_in_video.py b/tests/models/multimodal/processing/test_audio_in_video.py
index 6b77e8834e12..a8e164d7f39d 100644
--- a/tests/models/multimodal/processing/test_audio_in_video.py
+++ b/tests/models/multimodal/processing/test_audio_in_video.py
@@ -178,9 +178,11 @@ def test_audio_in_video_metadata_only(
     mm_config.enable_mm_embeds = True
     mm_config.allow_missing_mm_embeddings = True
     processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config)
-    prompt = "".join(
+    video_placeholders = [
         model_cls.get_placeholder_str("video", i) for i in range(num_videos)
-    )
+    ]
+    assert all(isinstance(token, str) for token in video_placeholders)
+    prompt = "".join(token for token in video_placeholders if token is not None)
     kwargs = {"use_audio_in_video": True}
     uuids = {
         modality: [f"source-{i}" for i in range(num_videos)]
@@ -192,7 +194,9 @@ def test_audio_in_video_metadata_only(
 
         mm_data["image"] = [Image.new("RGB", (64, 64), "green")]
         uuids["image"] = ["image-source"]
-        prompt = model_cls.get_placeholder_str("image", 0) + prompt
+        image_placeholder = model_cls.get_placeholder_str("image", 0)
+        assert image_placeholder is not None
+        prompt = image_placeholder + prompt
     raw = processor(
         prompt,
         mm_items=processor.info.parse_mm_data(mm_data),
@@ -237,9 +241,16 @@ def test_audio_in_video_metadata_only(
                 "uuid": "image-source",
             },
         )
-    parse = parse_chat_messages_async if is_async else parse_chat_messages
-    parsed = parse(messages, ctx.model_config, "openai", mm_processor_kwargs=kwargs)
-    conversation, metadata, remote_uuids = asyncio.run(parsed) if is_async else parsed
+    if is_async:
+        conversation, metadata, remote_uuids = asyncio.run(
+            parse_chat_messages_async(
+                messages, ctx.model_config, "openai", mm_processor_kwargs=kwargs
+            )
+        )
+    else:
+        conversation, metadata, remote_uuids = parse_chat_messages(
+            messages, ctx.model_config, "openai", mm_processor_kwargs=kwargs
+        )
     assert conversation[0]["content"] == (
         ([{"type": "image"}] if has_image else []) + [{"type": "video"}] * num_videos
     )
@@ -271,7 +282,7 @@ def unexpected_preprocessing(*args, **kwargs):
         )
 
     model = model_cls.__new__(model_cls)
-    model.config = processor.info.get_hf_config()
+    setattr(model, "config", processor.info.get_hf_config())  # noqa: B010
     remote_features = [
         MultiModalFeatureSpec(
             data=remote["mm_kwargs"][f.modality][int(f.identifier.split("-")[-1])],
@@ -281,10 +292,9 @@ def unexpected_preprocessing(*args, **kwargs):
         )
         for f in features
     ]
-    raw_positions, raw_delta = model.get_mrope_input_positions(
-        raw["prompt_token_ids"], features
-    )
-    remote_positions, remote_delta = model.get_mrope_input_positions(
+    get_positions = getattr(model, "get_mrope_input_positions")  # noqa: B009
+    raw_positions, raw_delta = get_positions(raw["prompt_token_ids"], features)
+    remote_positions, remote_delta = get_positions(
         remote["prompt_token_ids"], remote_features
     )
     assert raw_delta == remote_delta
diff --git a/tests/models/multimodal/processing/test_audioflamingo3.py b/tests/models/multimodal/processing/test_audioflamingo3.py
index c3adccb27d59..16182eac3052 100644
--- a/tests/models/multimodal/processing/test_audioflamingo3.py
+++ b/tests/models/multimodal/processing/test_audioflamingo3.py
@@ -134,7 +134,9 @@ def test_dummy_data_generation(mock_ctx):
     assert len(dummy_data["audio"]) == 2
 
     expected_len = 60 * 16000
-    assert len(dummy_data["audio"][0]) == expected_len
+    audio_data = dummy_data["audio"]
+    assert isinstance(audio_data, list)
+    assert len(audio_data[0]) == expected_len
 
 
 def test_audio_token_count_matches_hf_processor_math():
diff --git a/tests/models/multimodal/processing/test_common.py b/tests/models/multimodal/processing/test_common.py
index 667907c3e55c..3ff61ef7099b 100644
--- a/tests/models/multimodal/processing/test_common.py
+++ b/tests/models/multimodal/processing/test_common.py
@@ -26,6 +26,7 @@
 from vllm.inputs import MultiModalDataDict, MultiModalInput
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.multimodal.cache import MultiModalProcessorOnlyCache
+from vllm.multimodal.cache.base import MultiModalCache
 from vllm.multimodal.inputs import MultiModalFeatureSpec, batched_tensors_equal
 from vllm.multimodal.processing import BaseMultiModalProcessor, InputProcessingContext
 from vllm.platforms import current_platform
@@ -43,6 +44,7 @@
 
 def add_video_metadata(mm_data: MultiModalDataDict) -> MultiModalDataDict:
     """Add metadata to video mm_data."""
+    assert isinstance(mm_data, dict)
 
     def create_metadata(frames: np.ndarray):
         num_frames = len(frames)
@@ -71,6 +73,7 @@ def glmasr_patch_mm_data(mm_data: MultiModalDataDict) -> MultiModalDataDict:
     """Patch the multimodal data for GLM-ASR model.
     GLM-ASR requires text and audio to match 1:1, so we limit audio to 1.
     """
+    assert isinstance(mm_data, dict)
     if "audio" in mm_data:
         audio = mm_data["audio"]
         if isinstance(audio, list) and len(audio) > 1:
@@ -218,7 +221,7 @@ def test_audio_metadata_only_roundtrip(model_id, durations, monkeypatch):
     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
+    sample_rate = getattr(processor.info, "get_feature_extractor")().sampling_rate  # noqa: B009
     audios = [
         np.zeros(int(duration * sample_rate), dtype=np.float32)
         for duration in durations
@@ -296,6 +299,38 @@ def random_vision_chunk(
     return {"type": "video_chunk", "video_chunk": video_array}
 
 
+@pytest.mark.parametrize("order", [(2, 0, 1), (2, 0, 1, 0), (0, 2, 1)])
+def test_processor_cache_preserves_hits_when_request_exceeds_capacity(order):
+    """Within-request eviction must preserve tensors, tokens and placeholders."""
+    model = "llava-hf/llava-1.5-7b-hf"
+    config = ModelConfig(model, limit_mm_per_prompt={"image": 4})
+    processor = MULTIMODAL_REGISTRY.create_processor(config)
+    images = [Image.new("RGB", (112, 112), c) for c in ("red", "green", "blue")]
+
+    def process(selected, cache=None):
+        data = {"image": selected}
+        return processor(
+            get_token_prompt(processor, data),
+            mm_items=processor.info.parse_mm_data(data),
+            cache=cache,
+        )
+
+    reference = process(images)
+    capacity = sum(
+        MultiModalCache.get_item_size(item)
+        for item in reference["mm_kwargs"]["image"][:2]
+    )
+    cache = MultiModalProcessorOnlyCache(
+        ModelConfig(model, mm_processor_cache_gb=capacity / (1 << 30))
+    )
+    for _ in range(2):
+        process(images[:2], cache)
+        selected = [images[idx] for idx in order]
+        actual = process(selected, cache)
+        expected = process(selected)
+        _assert_inputs_equal(actual, expected)
+
+
 def _test_processing_correctness(
     model_id_or_arch: str,
     hit_rate: float,
@@ -333,6 +368,7 @@ def _test_processing_correctness(
     )
     # Ensure that the cache can fit all of the data
     # (set after because ModelConfig would set it to 0 for encoder-decoder models)
+    assert model_config.multimodal_config is not None
     model_config.multimodal_config.mm_processor_cache_gb = 2048
 
     model_cls = MULTIMODAL_REGISTRY._get_model_cls(model_config)
diff --git a/tests/models/multimodal/processing/test_funaudiochat.py b/tests/models/multimodal/processing/test_funaudiochat.py
new file mode 100644
index 000000000000..b090c6f38bdb
--- /dev/null
+++ b/tests/models/multimodal/processing/test_funaudiochat.py
@@ -0,0 +1,129 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from unittest.mock import MagicMock
+
+import numpy as np
+import pytest
+import torch
+
+from vllm.model_executor.models.funaudiochat import (
+    FunAudioChatDummyInputsBuilder,
+    FunAudioChatMultiModalProcessor,
+    FunAudioChatProcessingInfo,
+)
+from vllm.transformers_utils.configs.funaudiochat import FunAudioChatConfig
+
+
+class RecordingFeatureExtractor:
+    sampling_rate = 16000
+    n_fft = 400
+    hop_length = 160
+
+    def __init__(self) -> None:
+        self.last_kwargs: dict[str, object] | None = None
+
+    def __call__(self, wavs, **kwargs):
+        self.last_kwargs = kwargs
+        n_frames = [max(1, int(w.shape[0]) // self.hop_length) for w in wavs]
+        max_frames = max(n_frames)
+        features = torch.zeros((len(wavs), 128, max_frames))
+        mask = torch.zeros((len(wavs), max_frames), dtype=torch.long)
+        for i, frames in enumerate(n_frames):
+            mask[i, :frames] = 1
+        return {"input_features": features, "attention_mask": mask}
+
+
+class MockSpeechTokenizer:
+    pad_token = "<|audio_pad|>"
+
+    def __call__(self, speech_strs, **kwargs):
+        token_len = max(1, len(self.pad_token))
+        lengths = [len(s) // token_len for s in speech_strs]
+        max_len = max(lengths)
+        pad_to = kwargs.get("pad_to_multiple_of")
+        if pad_to:
+            max_len = ((max_len + pad_to - 1) // pad_to) * pad_to
+        ids = torch.zeros((len(speech_strs), max_len), dtype=torch.long)
+        mask = torch.zeros((len(speech_strs), max_len), dtype=torch.long)
+        for i, n in enumerate(lengths):
+            mask[i, :n] = 1
+        return {"input_ids": ids, "attention_mask": mask}
+
+
+class MockTokenizer:
+    def encode(self, text, **kwargs):
+        return [1, 2, 3]
+
+
+@pytest.fixture
+def funaudiochat_processor(monkeypatch):
+    feature_extractor = RecordingFeatureExtractor()
+    speech_tokenizer = MockSpeechTokenizer()
+    tokenizer = MockTokenizer()
+
+    monkeypatch.setattr(
+        "vllm.model_executor.models.funaudiochat.WhisperFeatureExtractor.from_pretrained",
+        lambda *args, **kwargs: feature_extractor,
+    )
+    monkeypatch.setattr(
+        "vllm.model_executor.models.funaudiochat.TokenizersBackend.from_pretrained",
+        lambda *args, **kwargs: speech_tokenizer,
+    )
+
+    ctx = MagicMock()
+    ctx.get_hf_config.return_value = FunAudioChatConfig()
+    ctx.model_config.hf_config = FunAudioChatConfig()
+    ctx.model_config.model = "funaudiochat"
+    ctx.model_config.revision = None
+    ctx.model_config.tokenizer_revision = None
+    mm_config = MagicMock()
+    mm_config.enable_mm_embeds = False
+    mm_config.allow_missing_mm_embeddings = False
+    ctx.model_config.multimodal_config = mm_config
+    ctx.model_config.get_multimodal_config.return_value = mm_config
+    ctx.tokenizer = tokenizer
+    ctx.get_tokenizer.return_value = tokenizer
+
+    info = FunAudioChatProcessingInfo(ctx)
+    processor = FunAudioChatMultiModalProcessor(
+        info, FunAudioChatDummyInputsBuilder(info)
+    )
+    return processor, feature_extractor
+
+
+def test_max_audio_samples_match_profiled_300s_budget(funaudiochat_processor):
+    processor, _ = funaudiochat_processor
+    sampling_rate = 16000
+
+    assert processor.info.get_max_speech_frames() == 7500
+    assert processor.info.get_max_audio_samples(sampling_rate) == 4_800_000
+    assert processor.info.get_max_audio_samples(sampling_rate) / sampling_rate == 300.0
+
+
+def test_over_budget_audio_is_rejected(funaudiochat_processor):
+    processor, feature_extractor = funaudiochat_processor
+    sampling_rate = 16000
+    over_budget = np.zeros(
+        processor.info.get_max_audio_samples(sampling_rate) + 1, dtype=np.float32
+    )
+    mm_items = processor.info.parse_mm_data({"audio": [over_budget]}, validate=False)
+
+    with pytest.raises(ValueError, match="too long for FunAudioChat"):
+        processor._apply_hf_processor_main(mm_items, {})
+
+    assert feature_extractor.last_kwargs is None
+
+
+def test_profiled_max_audio_is_accepted(funaudiochat_processor):
+    processor, feature_extractor = funaudiochat_processor
+    sampling_rate = 16000
+    max_samples = processor.info.get_max_audio_samples(sampling_rate)
+    audio = np.zeros(max_samples, dtype=np.float32)
+    mm_items = processor.info.parse_mm_data({"audio": [audio]}, validate=False)
+
+    processed = processor._apply_hf_processor_main(mm_items, {})
+
+    assert processed["speech_attention_mask"].sum().item() == 7500
+    assert feature_extractor.last_kwargs is not None
+    assert feature_extractor.last_kwargs.get("padding") == "max_length"
diff --git a/tests/models/multimodal/processing/test_gemma3.py b/tests/models/multimodal/processing/test_gemma3.py
index 2b4c213695ee..8b61deb91d6f 100644
--- a/tests/models/multimodal/processing/test_gemma3.py
+++ b/tests/models/multimodal/processing/test_gemma3.py
@@ -4,6 +4,7 @@
 import pytest
 import torch
 
+from vllm.model_executor.models.gemma3_mm import Gemma3ProcessingInfo
 from vllm.model_executor.models.gemma3n_audio_utils import (
     adjust_audio_features_to_expected_length,
 )
@@ -162,6 +163,7 @@ def test_get_image_size_with_most_features(
         limit_mm_per_prompt={"image": 1},
     )
     processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config)
+    assert isinstance(processor.info, Gemma3ProcessingInfo)
 
     hf_processor = processor.info.get_hf_processor(**mm_processor_kwargs)
 
@@ -185,5 +187,6 @@ def test_get_image_size_with_most_features(
         )
         mm_kwargs_data = processed_inputs["mm_kwargs"].get_data()
         num_patches_tensor = mm_kwargs_data["num_patches"]
+        assert isinstance(num_patches_tensor, torch.Tensor)
         tokens = int(num_patches_tensor.item()) * image_seq_length
         assert tokens <= max_tokens
diff --git a/tests/models/multimodal/processing/test_gemma4.py b/tests/models/multimodal/processing/test_gemma4.py
index f9d935991e90..90f5d0e24eba 100644
--- a/tests/models/multimodal/processing/test_gemma4.py
+++ b/tests/models/multimodal/processing/test_gemma4.py
@@ -14,6 +14,7 @@
 from vllm.model_executor.models.gemma4_mm import (
     Gemma4ForConditionalGeneration,
     Gemma4ImagePixelInputs,
+    Gemma4ProcessingInfo,
 )
 from vllm.model_executor.models.interfaces import supports_multimodal
 from vllm.multimodal import MULTIMODAL_REGISTRY
@@ -98,6 +99,7 @@ def test_compute_num_soft_tokens_does_not_exceed_max_soft_tokens(
         limit_mm_per_prompt={"image": 1},
     )
     processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config)
+    assert isinstance(processor.info, Gemma4ProcessingInfo)
 
     num_soft_tokens = processor.info._compute_num_soft_tokens(
         image_width=image_width,
@@ -190,6 +192,7 @@ def test_get_prompt_updates_respects_nested_max_soft_tokens(model_id: str):
         limit_mm_per_prompt={"image": 1},
     )
     processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config)
+    assert isinstance(processor.info, Gemma4ProcessingInfo)
     image = PILImage.new("RGB", (1000, 1000), color="white")
     image_size = image.size
     mm_items = processor.info.parse_mm_data({"image": image})
@@ -397,7 +400,7 @@ def test_encoder_chunk_no_free_memory_falls_back_to_one():
 
 
 def _gemma4_for_encoder_cudagraph() -> Gemma4ForConditionalGeneration:
-    model = object.__new__(Gemma4ForConditionalGeneration)
+    model = object.__new__(Gemma4ForConditionalGeneration)  # type: ignore[type-abstract]
     vision_config = SimpleNamespace(
         patch_size=16, pooling_kernel_size=3, default_output_length=280
     )
diff --git a/tests/models/multimodal/processing/test_gemma4_unified.py b/tests/models/multimodal/processing/test_gemma4_unified.py
index 28c1018c6602..3754e4abc14f 100644
--- a/tests/models/multimodal/processing/test_gemma4_unified.py
+++ b/tests/models/multimodal/processing/test_gemma4_unified.py
@@ -15,6 +15,7 @@
 )
 from vllm.model_executor.models.gemma4_unified import (
     Gemma4UnifiedForConditionalGeneration,
+    Gemma4UnifiedProcessingInfo,
     Gemma4UnifiedVisionEmbedder,
 )
 from vllm.multimodal import MULTIMODAL_REGISTRY
@@ -81,6 +82,7 @@ def test_compute_num_soft_tokens_does_not_exceed_max_soft_tokens(
         limit_mm_per_prompt={"image": 1},
     )
     processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config)
+    assert isinstance(processor.info, Gemma4UnifiedProcessingInfo)
 
     num_soft_tokens = processor.info._compute_num_soft_tokens(
         image_width=image_width,
@@ -171,6 +173,7 @@ def test_get_prompt_updates_respects_nested_max_soft_tokens(model_id: str):
         limit_mm_per_prompt={"image": 1},
     )
     processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config)
+    assert isinstance(processor.info, Gemma4UnifiedProcessingInfo)
     image = PILImage.new("RGB", (1000, 1000), color="white")
     image_size = image.size
     mm_items = processor.info.parse_mm_data({"image": image})
@@ -315,7 +318,8 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         return self.embedding_projection(x)
 
 
-class _MockUnifiedModel:
+class _MockUnifiedModel(Gemma4UnifiedForConditionalGeneration):
+    # Initialize only the modules needed by the borrowed vision methods.
     def __init__(
         self,
         embed_dim: int = 48,
@@ -323,6 +327,7 @@ def __init__(
         dtype: torch.dtype = torch.float32,
         device: str = "cpu",
     ):
+        torch.nn.Module.__init__(self)
         self.vision_embedder = _build_test_embedder(
             embed_dim=embed_dim, max_pos=max_pos, dtype=dtype, device=device
         )
@@ -330,11 +335,6 @@ def __init__(
             embed_dim=embed_dim, dtype=dtype, device=device
         )
 
-    _vision_chunk = staticmethod(Gemma4UnifiedForConditionalGeneration._vision_chunk)
-    _embed_vision_patches = Gemma4UnifiedForConditionalGeneration._embed_vision_patches
-    _process_image_input = Gemma4UnifiedForConditionalGeneration._process_image_input
-    _process_video_input = Gemma4UnifiedForConditionalGeneration._process_video_input
-
 
 @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32])
 @pytest.mark.parametrize("shape", [(16, 2), (2, 8, 2), (1, 1, 2)])
diff --git a/tests/models/multimodal/processing/test_glm5next.py b/tests/models/multimodal/processing/test_glm5next.py
index 111fb982cf7d..e7348f2585ea 100644
--- a/tests/models/multimodal/processing/test_glm5next.py
+++ b/tests/models/multimodal/processing/test_glm5next.py
@@ -9,9 +9,14 @@
 the two disagree they collide in ``_merge_multimodal_embeddings``, which raises
 inside a worker and takes the engine down with it.
 
-The checks below are arithmetic: no weights, no GPU.
+The grid checks below are arithmetic: no weights, no GPU.
+``test_glm5_metadata_only_roundtrip`` loads the processor and compares
+raw media with metadata-only prompts.
 """
 
+import json
+
+import numpy as np
 import pytest
 import torch
 from PIL import Image
@@ -21,13 +26,181 @@
 )
 from transformers.video_utils import VideoMetadata
 
+from vllm.distributed.ec_transfer.ec_connector.utils import (
+    PlaceholderMetadataResolver,
+    collect_ec_item_metadata,
+)
 from vllm.model_executor.layers.fusion.mm_input_norm import build_mm_input_norm
 from vllm.multimodal import MULTIMODAL_REGISTRY
+from vllm.multimodal.inputs import MultiModalFeatureSpec
 from vllm.platforms import current_platform
+from vllm.renderers.chat_utils.item_tracker import _get_embeds_data
 
+from ....multimodal.utils import random_image, random_video
 from ...utils import build_model_context
 
 
+def _video_with_metadata(frames: np.ndarray):
+    num_frames = len(frames)
+    return (
+        frames,
+        {
+            "total_num_frames": num_frames,
+            "fps": 2.0,
+            "duration": num_frames / 2.0,
+            "video_backend": "opencv",
+            "frames_indices": list(range(num_frames)),
+            "do_sample_frames": True,
+        },
+    )
+
+
+@pytest.mark.parametrize("modality", ["image", "video"])
+def test_glm5_metadata_only_roundtrip(modality):
+    """EC metadata reproduces the raw prompt for image and video."""
+    rng = np.random.RandomState(0)
+    if modality == "image":
+        mm_data = {"image": random_image(rng, 64, 65)}
+        limit = {"image": 1, "video": 0}
+        grid_key = "image_grid_thw"
+    else:
+        mm_data = {"video": _video_with_metadata(random_video(rng, 4, 5, 64, 65))}
+        limit = {"image": 0, "video": 1}
+        grid_key = "video_grid_thw"
+
+    ctx = build_model_context(
+        "zai-org/GLM-5.3-Flash",
+        limit_mm_per_prompt=limit,
+    )
+    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,
+        tokenizer=ctx.tokenizer,
+    )
+    prompt = processor.dummy_inputs.get_dummy_text({modality: 1})
+    raw = processor(
+        prompt,
+        mm_items=processor.info.parse_mm_data(mm_data),
+        hf_processor_mm_kwargs={},
+    )
+    features = [
+        MultiModalFeatureSpec(
+            data=item,
+            modality=modality,
+            identifier=str(index),
+            mm_position=position,
+        )
+        for index, (item, position) in enumerate(
+            zip(raw["mm_kwargs"][modality], raw["mm_placeholders"][modality])
+        )
+    ]
+    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) == 1
+    assert grid_key in metadata[0]
+    if modality == "video":
+        assert metadata[0]["timestamps"]
+    mm_data_remote = _get_embeds_data(modality, metadata, processor)
+    remote = processor(
+        prompt,
+        mm_items=processor.info.parse_mm_data({modality: mm_data_remote}),
+        hf_processor_mm_kwargs={},
+    )
+    assert remote["prompt_token_ids"] == raw["prompt_token_ids"]
+    assert remote["mm_placeholders"] == raw["mm_placeholders"]
+
+    mm_config.allow_missing_mm_embeddings = False
+    ordinary_processor = MULTIMODAL_REGISTRY.create_processor(
+        ctx.model_config,
+        tokenizer=ctx.tokenizer,
+    )
+    with pytest.raises(ValueError, match="should contain the fields"):
+        ordinary_processor.info.parse_mm_data({modality: mm_data_remote})
+
+
+def test_glm5_raw_video_timestamps_follow_sampled_frames():
+    """Input frames_indices are the pre-sample range, not the HF sample.
+
+    32 frames at 8 fps with ``frames_indices=range(32)`` and
+    ``do_sample_frames=True`` used to emit 16 timestamps for ``grid_t=4``.
+    """
+    rng = np.random.RandomState(0)
+    num_frames = 32
+    frames = rng.randint(0, 255, size=(num_frames, 64, 64, 3), dtype=np.uint8)
+    metadata = {
+        "total_num_frames": num_frames,
+        "fps": 8.0,
+        "duration": num_frames / 8.0,
+        "frames_indices": list(range(num_frames)),
+        "do_sample_frames": True,
+    }
+    mm_data = {"video": (frames, metadata)}
+    ctx = build_model_context(
+        "zai-org/GLM-5.3-Flash",
+        limit_mm_per_prompt={"image": 0, "video": 1},
+    )
+    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,
+        tokenizer=ctx.tokenizer,
+    )
+    prompt = processor.dummy_inputs.get_dummy_text({"video": 1})
+    raw = processor(
+        prompt,
+        mm_items=processor.info.parse_mm_data(mm_data),
+        hf_processor_mm_kwargs={},
+    )
+    video_item = raw["mm_kwargs"]["video"][0]
+    assert video_item is not None
+    item = video_item.get_data()
+    grid = item["video_grid_thw"]
+    timestamps = item["timestamps"]
+    grid_t = int(grid.reshape(-1)[0] if isinstance(grid, torch.Tensor) else grid[0])
+    if isinstance(timestamps, torch.Tensor):
+        timestamp_count = timestamps.numel()
+    else:
+        flat: list[object] = []
+        for entry in timestamps:
+            flat.extend(entry if isinstance(entry, list) else [entry])
+        timestamp_count = len(flat)
+    assert timestamp_count == grid_t
+    assert grid_t != num_frames // 2
+
+    features = [
+        MultiModalFeatureSpec(
+            data=mm_item,
+            modality="video",
+            identifier=str(index),
+            mm_position=position,
+        )
+        for index, (mm_item, position) in enumerate(
+            zip(raw["mm_kwargs"]["video"], raw["mm_placeholders"]["video"])
+        )
+    ]
+    published = collect_ec_item_metadata(
+        features, PlaceholderMetadataResolver(ctx.model_config)
+    )
+    published_metadata = json.loads(
+        json.dumps([entry["metadata"] for entry in published.values()])
+    )
+    mm_data_remote = _get_embeds_data("video", published_metadata, processor)
+    remote = processor(
+        prompt,
+        mm_items=processor.info.parse_mm_data({"video": mm_data_remote}),
+        hf_processor_mm_kwargs={},
+    )
+    assert remote["prompt_token_ids"] == raw["prompt_token_ids"]
+    assert remote["mm_placeholders"] == raw["mm_placeholders"]
+
+
 @pytest.fixture(scope="module")
 def processor():
     ctx = build_model_context(
@@ -143,6 +316,48 @@ def test_video_placeholders_match_encoder_rows_when_presampled(processor):
     assert len(timestamps) == grid_t
 
 
+def test_placeholder_rejects_timestamp_count_mismatch(processor):
+    """A timestamp list that is not one entry per temporal patch is refused."""
+    grid_thw = torch.tensor([4, 36, 46])
+    with pytest.raises(ValueError, match="video_grid_thw"):
+        processor.info._placeholder_from_timestamps([0, 1], grid_thw)
+
+
+def test_video_timestamps_follow_supplied_frame_indices(processor):
+    """When do_sample_frames is false, supplied indices win over a resample."""
+    info = processor.info
+    video_processor = info.get_video_processor()
+    temporal_patch_size = int(getattr(video_processor, "temporal_patch_size", 1) or 1)
+    fps = 8.0
+    indices = [0, 8, 80, 88]
+    metadata = {
+        "total_num_frames": 256,
+        "fps": fps,
+        "duration": 32.0,
+        "do_sample_frames": False,
+        "frames_indices": indices,
+    }
+
+    timestamps = info._get_video_second_idx_glm46v(metadata, 256)
+    resampled = info._get_video_second_idx_glm46v(
+        {
+            **{k: v for k, v in metadata.items() if k != "frames_indices"},
+            "do_sample_frames": True,
+        },
+        256,
+    )
+
+    num_frames = len(indices)
+    num_padded = num_frames + (-num_frames % temporal_patch_size)
+    grid_t = max(num_padded // temporal_patch_size, 1)
+    last = len(indices) - 1
+    expected = [
+        int(indices[min(i * temporal_patch_size, last)] / fps) for i in range(grid_t)
+    ]
+    assert timestamps == expected
+    assert timestamps != resampled
+
+
 def test_video_shorter_than_one_sampling_interval_is_rejected(processor):
     """A clip the sampler cannot pick a single frame from is a bad request."""
     with pytest.raises(ValueError, match="selected no frames"):
@@ -177,10 +392,12 @@ def test_mm_device_do_normalize():
 
     normalized_inputs = processor(prompt, mm_items=mm_items)
     normalized_values = normalized_inputs["mm_kwargs"].get_data()["pixel_values"]
+    assert isinstance(normalized_values, torch.Tensor)
 
     ctx.model_config.multimodal_config.mm_device_do_normalize = True
     raw_inputs = processor(prompt, mm_items=mm_items)
     raw_values = raw_inputs["mm_kwargs"].get_data()["pixel_values"]
+    assert isinstance(raw_values, torch.Tensor)
     assert raw_values.dtype == torch.uint8
 
     input_norm = build_mm_input_norm(ctx.model_config).to(device)
diff --git a/tests/models/multimodal/processing/test_h2ovl.py b/tests/models/multimodal/processing/test_h2ovl.py
index 5a54484e1bdd..bc43130d04f1 100644
--- a/tests/models/multimodal/processing/test_h2ovl.py
+++ b/tests/models/multimodal/processing/test_h2ovl.py
@@ -120,7 +120,9 @@ def _run_check(
     # Ensure we have the right number of placeholders per num_crops size
     image_token_id = tokenizer.convert_tokens_to_ids("<IMG_CONTEXT>")
     img_tok_count = processed_inputs["prompt_token_ids"].count(image_token_id)
-    pixel_shape = processed_inputs["mm_kwargs"].get_data()["pixel_values_flat"].shape
+    pixel_values = processed_inputs["mm_kwargs"].get_data()["pixel_values_flat"]
+    assert isinstance(pixel_values, torch.Tensor)
+    pixel_shape = pixel_values.shape
 
     assert img_tok_count == 256 * total_expected_num_patches
     assert pixel_shape[0] == total_expected_num_patches
@@ -204,6 +206,7 @@ def test_image_cache_identity_across_image_counts(
     )
     processor = MULTIMODAL_REGISTRY.create_processor(config)
     tokenizer = processor.info.get_tokenizer()
+    images: list[torch.Tensor] | list[Image.Image]
     if use_embeds:
         images = [torch.ones(8, config.get_hidden_size())] * 3
     else:
diff --git a/tests/models/multimodal/processing/test_internvl.py b/tests/models/multimodal/processing/test_internvl.py
index ab5fbe9a0130..de126c2c95fb 100644
--- a/tests/models/multimodal/processing/test_internvl.py
+++ b/tests/models/multimodal/processing/test_internvl.py
@@ -5,6 +5,7 @@
 from collections.abc import Mapping
 
 import pytest
+import torch
 from PIL import Image
 from transformers import PreTrainedConfig
 
@@ -75,7 +76,9 @@ def _run_check(
     # Ensure we have the right number of placeholders per num_crops size
     image_token_id = tokenizer.convert_tokens_to_ids("<IMG_CONTEXT>")
     img_tok_count = processed_inputs["prompt_token_ids"].count(image_token_id)
-    pixel_shape = processed_inputs["mm_kwargs"].get_data()["pixel_values_flat"].shape
+    pixel_values = processed_inputs["mm_kwargs"].get_data()["pixel_values_flat"]
+    assert isinstance(pixel_values, torch.Tensor)
+    pixel_shape = pixel_values.shape
 
     assert img_tok_count == 256 * total_expected_num_patches
     assert pixel_shape[0] == total_expected_num_patches
diff --git a/tests/models/multimodal/processing/test_kimi_k3.py b/tests/models/multimodal/processing/test_kimi_k3.py
index 62c99b24f143..599b0413e0bc 100644
--- a/tests/models/multimodal/processing/test_kimi_k3.py
+++ b/tests/models/multimodal/processing/test_kimi_k3.py
@@ -23,7 +23,7 @@
 def image_processors():
     ctx = build_model_context(_MODEL_ID, limit_mm_per_prompt={"image": 1})
     processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config)
-    k25_fused = processor.info.image_processor
+    k25_fused = getattr(processor.info, "image_processor")  # noqa: B009
     assert isinstance(k25_fused, KimiK25FusedVisionProcessor)
     k3_hf = AutoImageProcessor.from_pretrained(_MODEL_ID, trust_remote_code=True)
     return k25_fused, k3_hf
@@ -68,10 +68,12 @@ def test_mm_device_do_normalize():
 
     normalized_inputs = processor(prompt, mm_items=mm_items)
     normalized_values = normalized_inputs["mm_kwargs"].get_data()["pixel_values"]
+    assert isinstance(normalized_values, torch.Tensor)
 
     ctx.model_config.multimodal_config.mm_device_do_normalize = True
     raw_inputs = processor(prompt, mm_items=mm_items)
     raw_values = raw_inputs["mm_kwargs"].get_data()["pixel_values"]
+    assert isinstance(raw_values, torch.Tensor)
     assert raw_values.dtype == torch.uint8
 
     input_norm = build_mm_input_norm(ctx.model_config).to(device)
diff --git a/tests/models/multimodal/processing/test_lightonocr.py b/tests/models/multimodal/processing/test_lightonocr.py
index 3966bf8cb4dd..c138808c7736 100644
--- a/tests/models/multimodal/processing/test_lightonocr.py
+++ b/tests/models/multimodal/processing/test_lightonocr.py
@@ -47,4 +47,5 @@ def test_lightonocr_image_embeddings_preserve_lengths_and_values(
     )
     assert [p.length for p in processed["mm_placeholders"]["image"]] == feature_sizes
     for item, expected in zip(processed["mm_kwargs"]["image"], embeddings):
+        assert item is not None
         torch.testing.assert_close(item.get_data()["image_embeds"], expected)
diff --git a/tests/models/multimodal/processing/test_llama4.py b/tests/models/multimodal/processing/test_llama4.py
index 4bc2e5909980..e3d4a94796b5 100644
--- a/tests/models/multimodal/processing/test_llama4.py
+++ b/tests/models/multimodal/processing/test_llama4.py
@@ -3,7 +3,9 @@
 """Tests for Llama4's multimodal preprocessing kwargs."""
 
 import pytest
+import torch
 
+from vllm.model_executor.models.mllama4 import Mllama4ProcessingInfo
 from vllm.multimodal import MULTIMODAL_REGISTRY
 
 from ....conftest import ImageTestAssets
@@ -31,6 +33,7 @@ def test_processor_override(
         mm_processor_cache_gb=mm_processor_cache_gb,
     )
     processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config)
+    assert isinstance(processor.info, Mllama4ProcessingInfo)
     config = processor.info.get_hf_config()
     tokenizer = processor.info.get_tokenizer()
     hf_processor = processor.info.get_hf_processor()
@@ -62,6 +65,7 @@ def test_processor_override(
     assert prompt_token_ids.count(config.eoi_token_index) == num_imgs
     assert prompt_token_ids.count(vocab[hf_processor.image_token]) == num_imgs
     aspect_ratios = mm_data["aspect_ratios"]
+    assert isinstance(aspect_ratios, torch.Tensor)
     num_x_separators = num_y_separators = 0
     for tiles_y, tiles_x in aspect_ratios:
         if tiles_x * tiles_y > 1:
diff --git a/tests/models/multimodal/processing/test_ming_v3_asr.py b/tests/models/multimodal/processing/test_ming_v3_asr.py
new file mode 100644
index 000000000000..7363e61a3c5e
--- /dev/null
+++ b/tests/models/multimodal/processing/test_ming_v3_asr.py
@@ -0,0 +1,220 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""The native Ming frontend must preserve the released ASR request contract."""
+
+import os
+from io import BytesIO
+
+import numpy as np
+import pytest
+import torch
+from starlette.datastructures import UploadFile
+from transformers import AutoProcessor
+
+from vllm.config import ModelConfig
+from vllm.entrypoints.speech_to_text.transcription.protocol import TranscriptionRequest
+from vllm.model_executor.models.ming_audio import (
+    MingAudioProjector,
+    get_audio_token_count,
+)
+from vllm.model_executor.models.ming_v3_asr import (
+    ASR_INSTRUCTION,
+    MingASRStreamingPostProcessor,
+    MingV3ASRForConditionalGeneration,
+)
+from vllm.multimodal import MULTIMODAL_REGISTRY
+from vllm.multimodal.cache import MultiModalProcessorOnlyCache
+from vllm.transformers_utils.configs.ming_v3_asr import MingAudioConfig
+
+from .test_common import _assert_inputs_equal
+
+MODEL = os.environ.get("VLLM_TEST_MING_ASR_MODEL", "inclusionAI/Ming-3.0-asr-tiny")
+
+
+@pytest.fixture(scope="module")
+def asr_processors():
+    config = ModelConfig(
+        MODEL, dtype="bfloat16", mm_processor_cache_gb=1, trust_remote_code=True
+    )
+    return (
+        config,
+        MULTIMODAL_REGISTRY.create_processor(config),
+        AutoProcessor.from_pretrained(MODEL, trust_remote_code=True),
+    )
+
+
+@pytest.mark.parametrize(
+    "samples,language,frames",
+    [
+        (201, None, 1),
+        (321, None, 3000),
+        (16001, "en", 3000),
+        (480001, "zh", 3000),
+        (480161, "en", 1001),
+        (960000, "zh", 6000),
+        (960161, "zh", 6000),
+    ],
+)
+def test_transcription_request_matches_native(
+    asr_processors, samples, language, frames
+):
+    """Compare features and token IDs, including language seed placement and tails."""
+    config, processor, native = asr_processors
+    audio = np.random.RandomState(42).uniform(-0.5, 0.5, samples).astype(np.float32)
+    prompt = ASR_INSTRUCTION + "This is an audio about 同囚."
+    request = TranscriptionRequest(
+        file=UploadFile(BytesIO()),
+        prompt=prompt,
+        language=language,
+        vllm_xargs={"audio_chunk_frames": frames},
+    )
+    stt_params = request.build_stt_params(
+        audio,
+        MingV3ASRForConditionalGeneration.get_speech_to_text_config(
+            config, "transcribe"
+        ),
+        config,
+        "transcribe",
+    )
+    inputs = MingV3ASRForConditionalGeneration.get_generation_prompt(stt_params)
+    actual = processor(
+        inputs["prompt_token_ids"],
+        mm_items=processor.info.parse_mm_data(inputs["multi_modal_data"]),
+        hf_processor_mm_kwargs=inputs["mm_processor_kwargs"],
+    )
+    expected = native.apply_transcription_request(
+        (audio, 16000),
+        prompt=prompt,
+        language_hint=language,
+        audio_chunk_frames=frames,
+    )
+    assert actual["prompt_token_ids"] == expected.input_ids[0].tolist()
+    data = actual["mm_kwargs"].get_data()
+    torch.testing.assert_close(
+        data["audio_feats"][0],
+        expected.audio_feats[0],
+        rtol=0,
+        atol=0,
+    )
+
+
+@pytest.mark.parametrize(
+    "kernels,strides,paddings", [([3], [2], [1]), ([5, 3], [3, 2], [2, 0])]
+)
+def test_audio_token_count_matches_chunked_convolutions(kernels, strides, paddings):
+    """Count each chunk's rounding separately, including a nonempty tail."""
+    config = MingAudioConfig(
+        n_audio_state=4,
+        downsampler_kernel_size_list=kernels,
+        downsampler_stride_list=strides,
+        downsampler_padding_list=paddings,
+    )
+    encoder_conv = torch.nn.Conv1d(1, config.n_audio_state, 3, stride=2, padding=1)
+    projector = MingAudioProjector(config, hidden_size=4)
+    num_frames, chunk_frames = 3103, 1001
+    with torch.inference_mode():
+        actual = sum(
+            projector(encoder_conv(chunk).transpose(1, 2)).shape[1]
+            for chunk in torch.zeros(1, 1, num_frames).split(chunk_frames, dim=-1)
+        )
+    assert get_audio_token_count(config, num_frames, chunk_frames) == actual
+
+
+@pytest.mark.parametrize("scoped", [False, True])
+def test_chunk_policy_has_separate_cache_entries(asr_processors, scoped):
+    """The same waveform encoded with different chunks must not share cached data."""
+    config, processor, _ = asr_processors
+    cache = MultiModalProcessorOnlyCache(config)
+    audio = np.random.RandomState(1).randn(480161).astype(np.float32)
+    results = []
+    for frames in (3000, 6000, 3000):
+        options: dict[str, object] = {"audio_chunk_frames": frames}
+        if scoped:
+            options = {"audio_kwargs": options}
+        kwargs = dict(
+            mm_items=processor.info.parse_mm_data({"audio": audio}),
+            hf_processor_mm_kwargs=options,
+        )
+        result = processor("<audio><audioPatch></audio>", **kwargs, cache=cache)
+        _assert_inputs_equal(result, processor("<audio><audioPatch></audio>", **kwargs))
+        results.append(result)
+    assert results[0]["mm_hashes"] != results[1]["mm_hashes"]
+    _assert_inputs_equal(results[0], results[2])
+    assert config.hf_config.audio_config.chunk_frames == 3000
+    for result, frames in zip(results, (3000, 6000, 3000)):
+        data = result["mm_kwargs"].get_data()
+        assert torch.as_tensor(data["audio_chunk_frames"]).item() == frames
+
+
+@pytest.mark.parametrize(
+    "options,samples,expected_frames",
+    [
+        ({}, 321, 6000),
+        ({"audio_chunk_frames": 3000}, 321, 3000),
+        ({"audio_kwargs": {"audio_chunk_frames": 4500}}, 321, 4500),
+        (
+            {"audio_kwargs": {"audio_chunk_frames": None, "context_chunks": True}},
+            480161,
+            6000,
+        ),
+    ],
+)
+def test_processor_defaults_and_request_audio_options(
+    options, samples, expected_frames
+):
+    """Configured defaults and both request encodings must select the chunk policy."""
+    config = ModelConfig(
+        MODEL,
+        dtype="bfloat16",
+        trust_remote_code=True,
+        mm_processor_kwargs={"audio_kwargs": {"audio_chunk_frames": 6000}},
+    )
+    processor = MULTIMODAL_REGISTRY.create_processor(config)
+    actual = processor(
+        "<audio><audioPatch></audio>",
+        mm_items=processor.info.parse_mm_data(
+            {"audio": np.zeros(samples, dtype=np.float32)}
+        ),
+        hf_processor_mm_kwargs=options,
+    )
+    frames = actual["mm_kwargs"].get_data()["audio_chunk_frames"]
+    assert torch.as_tensor(frames).item() == expected_frames
+
+
+@pytest.mark.parametrize(
+    "options,max_model_len",
+    [({"audio_chunk_frames": 1}, 512), ({"context_chunks": True}, 2048)],
+)
+def test_dummy_audio_respects_processor_defaults(options, max_model_len):
+    """Configured chunking must keep profiling within token and duration limits."""
+    config = ModelConfig(
+        MODEL,
+        dtype="bfloat16",
+        trust_remote_code=True,
+        max_model_len=max_model_len,
+        mm_processor_kwargs={"audio_kwargs": options},
+    )
+    processor = MULTIMODAL_REGISTRY.create_processor(config)
+    inputs = processor.get_dummy_mm_inputs({"audio": 1})
+    assert len(inputs["prompt_token_ids"]) == max_model_len
+
+
+@pytest.mark.parametrize(
+    "chunks,expected",
+    [
+        (["Chi", "nese", "\t", " 你好", " 世界 "], "你好 世界"),
+        (["same ", "way ", "you did"], "same way you did"),
+        (["Eng", "lish"], ""),
+        (["Canton", "ese", "\t", "你好"], "你好"),
+        (["你", "好"], "你好"),
+    ],
+)
+def test_streaming_envelope_matches_final_text(chunks, expected):
+    processor = MingASRStreamingPostProcessor()
+    deltas = [processor.process_delta(chunk, False) for chunk in chunks]
+    deltas.append(processor.process_delta("", True))
+    assert "".join(deltas) == expected
+    assert (
+        MingV3ASRForConditionalGeneration.post_process_output("".join(chunks))
+        == expected
+    )
diff --git a/tests/models/multimodal/processing/test_minicpm_ocr.py b/tests/models/multimodal/processing/test_minicpm_ocr.py
index 8544112b4473..6f2a1296f38d 100644
--- a/tests/models/multimodal/processing/test_minicpm_ocr.py
+++ b/tests/models/multimodal/processing/test_minicpm_ocr.py
@@ -50,7 +50,12 @@ class _StubImageProcessor:
     slice_mode = True
 
     def find_best_resize(
-        self, image_size, scale_resolution, patch_size, allow_upscale=False
+        self,
+        image_size,
+        scale_resolution,
+        patch_size,
+        downsample_mode,
+        allow_upscale=False,
     ):
         height, width = image_size
         if (height * width > scale_resolution * scale_resolution) or allow_upscale:
@@ -58,13 +63,19 @@ def find_best_resize(
             height = int(scale_resolution / math.sqrt(aspect_ratio))
             width = int(height * aspect_ratio)
         # The OCR delta: one 2x2 merge in "4x" mode, two in "16x".
-        merge_factor = 2 if self.downsample_mode == "4x" else 4
+        merge_factor = 2 if downsample_mode == "4x" else 4
         best_width = ensure_divide(width, patch_size * merge_factor)
         best_height = ensure_divide(height, patch_size * merge_factor)
         return best_height, best_width
 
     def get_refine_size(
-        self, image_size, grid, scale_resolution, patch_size, allow_upscale=False
+        self,
+        image_size,
+        grid,
+        scale_resolution,
+        patch_size,
+        downsample_mode,
+        allow_upscale=False,
     ):
         height, width = image_size
         grid_y, grid_x = grid
@@ -74,6 +85,7 @@ def get_refine_size(
             (refine_height / grid_y, refine_width / grid_x),
             scale_resolution,
             patch_size,
+            downsample_mode,
             allow_upscale=allow_upscale,
         )
         return best_height * grid_y, best_width * grid_x
@@ -167,10 +179,18 @@ def test_4x_mode_aligns_the_patch_grid_to_two_patches():
     # Upscaled to ~448 on the long side, then aligned: "4x" rounds to 28px and
     # "16x" to 56px, which lands on a different grid for the same image.
     assert info._pinned_image_processor("4x").find_best_resize(
-        (200, 300), _SCALE_RESOLUTION, _PATCH_SIZE, allow_upscale=True
+        (200, 300),
+        _SCALE_RESOLUTION,
+        _PATCH_SIZE,
+        downsample_mode="4x",
+        allow_upscale=True,
     ) == (364, 560)
     assert info._pinned_image_processor("16x").find_best_resize(
-        (200, 300), _SCALE_RESOLUTION, _PATCH_SIZE, allow_upscale=True
+        (200, 300),
+        _SCALE_RESOLUTION,
+        _PATCH_SIZE,
+        downsample_mode="16x",
+        allow_upscale=True,
     ) == (392, 560)
 
     # 364*560 / (14*14*4) = 260 against 392*560 / (14*14*16) = 70: the ratio is
@@ -227,7 +247,12 @@ def test_sliced_image_token_counts_match_the_transformers_divisor():
     assert grids != [0, 0]
     processor = info._pinned_image_processor("4x")
     refine = processor.get_refine_size(
-        (1120, 1680), grids, _SCALE_RESOLUTION, _PATCH_SIZE, allow_upscale=True
+        (1120, 1680),
+        grids,
+        _SCALE_RESOLUTION,
+        _PATCH_SIZE,
+        downsample_mode="4x",
+        allow_upscale=True,
     )
     expected_patch = (
         (refine[0] // grids[0])
@@ -453,12 +478,21 @@ def test_placeholder_matches_the_transformers_formula(width, height, mode):
     # Rebuild `target_sizes` exactly as `_preprocess` records it.
     sliced = grids != [0, 0]
     source_h, source_w = processor.find_best_resize(
-        (height, width), _SCALE_RESOLUTION, _PATCH_SIZE, allow_upscale=not sliced
+        (height, width),
+        _SCALE_RESOLUTION,
+        _PATCH_SIZE,
+        downsample_mode=mode,
+        allow_upscale=not sliced,
     )
     target_sizes = [[source_h // _PATCH_SIZE, source_w // _PATCH_SIZE]]
     if sliced:
         refine_h, refine_w = processor.get_refine_size(
-            (height, width), grids, _SCALE_RESOLUTION, _PATCH_SIZE, allow_upscale=True
+            (height, width),
+            grids,
+            _SCALE_RESOLUTION,
+            _PATCH_SIZE,
+            downsample_mode=mode,
+            allow_upscale=True,
         )
         patch_h, patch_w = refine_h // grids[0], refine_w // grids[1]
         target_sizes += [[patch_h // _PATCH_SIZE, patch_w // _PATCH_SIZE]] * (
@@ -500,6 +534,7 @@ def test_pixel_and_target_size_paths_agree_on_the_default_resize():
             (height, width),
             _SCALE_RESOLUTION,
             _PATCH_SIZE,
+            downsample_mode="4x",
             allow_upscale=not sliced,
         )
         target_sizes = [[source_h // _PATCH_SIZE, source_w // _PATCH_SIZE]]
@@ -509,6 +544,7 @@ def test_pixel_and_target_size_paths_agree_on_the_default_resize():
                 grids,
                 _SCALE_RESOLUTION,
                 _PATCH_SIZE,
+                downsample_mode="4x",
                 allow_upscale=True,
             )
             patch_h, patch_w = refine_h // grids[0], refine_w // grids[1]
diff --git a/tests/models/multimodal/processing/test_minicpmv.py b/tests/models/multimodal/processing/test_minicpmv.py
index 59cb4331f3b0..fbcdd77e3f96 100644
--- a/tests/models/multimodal/processing/test_minicpmv.py
+++ b/tests/models/multimodal/processing/test_minicpmv.py
@@ -6,6 +6,7 @@
 import numpy as np
 import pytest
 
+from vllm.model_executor.models.minicpmv import MiniCPMVProcessingInfo
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.multimodal.parse import ImageSize
 
@@ -21,6 +22,7 @@ def test_get_hf_processor_for_same_model_different_kwargs(model_id: str):
     )
     processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config)
     info = processor.info
+    assert isinstance(info, MiniCPMVProcessingInfo)
     processor_1 = info.get_hf_processor(max_slice_nums=1)
     processor_2 = info.get_hf_processor(max_slice_nums=2)
     assert processor_1.image_processor.max_slice_nums == 1
@@ -42,6 +44,7 @@ def test_image_processor_for_different_models(model_ids):
         limit_mm_per_prompt={"image": 1},
     )
     first_processor = MULTIMODAL_REGISTRY.create_processor(first_ctx.model_config)
+    assert isinstance(first_processor.info, MiniCPMVProcessingInfo)
     first_image_processor = first_processor.info.get_image_processor()
 
     second_ctx = build_model_context(
@@ -49,6 +52,7 @@ def test_image_processor_for_different_models(model_ids):
         limit_mm_per_prompt={"image": 1},
     )
     second_processor = MULTIMODAL_REGISTRY.create_processor(second_ctx.model_config)
+    assert isinstance(second_processor.info, MiniCPMVProcessingInfo)
     second_image_processor = second_processor.info.get_image_processor()
 
     second_processor.info.get_sliced_grid(ImageSize(width=128, height=128))
diff --git a/tests/models/multimodal/processing/test_minimax_m3.py b/tests/models/multimodal/processing/test_minimax_m3.py
index 04d6aa4778c3..b24ea2b803c6 100644
--- a/tests/models/multimodal/processing/test_minimax_m3.py
+++ b/tests/models/multimodal/processing/test_minimax_m3.py
@@ -1,138 +1,160 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Unit tests for MiniMax-M3 VL ``max_long_side_pixel`` resize support.
+"""Unit tests for MiniMax-M3 VL preprocessing (sglang-aligned size spec).
 
-These exercise the vendored processor directly (no checkpoint / GPU needed), so
-they validate the long-side resize spec and the resulting prompt-token counts
-deterministically.
+These exercise the vendored processor directly (no checkpoint / GPU needed) and
+pin the ``compute_aligned_size`` / ``resolve_max_long_side_pixel`` contract
+against the sglang reference algorithm
+(``python/sglang/srt/multimodal/processors/minimax_m3_vl.py``).
 """
 
 import pytest
 import torch
 
 from vllm.transformers_utils.processors.minimax_m3 import (
-    IMAGE_MAX_TOTAL_PIXELS,
     MIN_SHORT_SIDE_PIXEL,
-    VIDEO_MAX_TOTAL_PIXELS,
     MiniMaxM3VLImageProcessor,
     MiniMaxM3VLVideoProcessor,
-    smart_resize,
+    compute_aligned_size,
+    resolve_max_long_side_pixel,
 )
 
-# Long sides are multiples of patch_size*merge_size (28) so the rounding is
-# exact and the expected token counts are unambiguous.
-LONG_SIDES = [252, 504, 1008]
+FACTOR = 28  # patch_size * merge_size
+PATCH_SIZE = 14
 MERGE2 = 2**2  # merge_size ** 2
 
+# detail -> max_long_side_pixel, per the size spec.
+IMAGE_DETAIL_CAPS = {"low": 672, "auto": 2016, "high": 3584}
+VIDEO_DETAIL_CAPS = {"low": 504, "auto": 672, "high": 2188}
 
-def _image_tokens(grid_thw) -> int:
-    g = list(grid_thw)
-    return int(g[0] * g[1] * g[2]) // MERGE2
+
+def _merged_tokens(grid_thw) -> int:
+    return int(grid_thw.prod()) // MERGE2
 
 
 # --------------------------------------------------------------------------- #
-# smart_resize: the long-side spec (a) shrink / (b) enlarge / (c) hard cap
+# compute_aligned_size: floor layers / ceil-align / long-side cap
 # --------------------------------------------------------------------------- #
-def test_smart_resize_long_side_shrink():
-    # (a) long side exceeds the cap -> shrink so the long side equals the cap.
-    h, w = smart_resize(
-        2048, 1024, factor=28, max_long_side_pixel=1008, max_total_pixels=10**9
-    )
-    assert max(h, w) == 1008
-    assert (h, w) == (1008, 504)  # aspect ratio preserved
+def test_compute_aligned_size_enlarges_short_side():
+    # (200, 40) hits both floors: area 8000 < 168**2 scales to (376, 76); the
+    # short-side floor then scales by 168/76, where ceil(76 * 168/76) = 169
+    # (float overshoot) aligns up to 196, and 832 aligns up to 840.
+    h, w = compute_aligned_size(200, 40, FACTOR, max_long_side_pixel=1008)
+    assert (h, w) == (840, 196)
+    assert h % FACTOR == 0 and w % FACTOR == 0
 
 
-def test_smart_resize_short_side_enlarge():
-    # (b) long side within the cap but short side below the floor -> enlarge so
-    # the short side reaches min_short_side_pixel.
-    h, w = smart_resize(
-        200, 40, factor=28, max_long_side_pixel=1008, max_total_pixels=10**9
-    )
-    assert min(h, w) == MIN_SHORT_SIDE_PIXEL  # 112
+def test_compute_aligned_size_shrinks_over_cap_isotropically():
+    # Long side 2048 ceil-aligns to 2072 > cap -> integer aspect-ratio
+    # scale-back to (1008, 504), already factor-aligned.
+    h, w = compute_aligned_size(2048, 1024, FACTOR, max_long_side_pixel=1008)
+    assert (h, w) == (1008, 504)
 
 
-def test_smart_resize_total_pixels_raises():
-    # (c) still over the area cap after resizing -> raise instead of inferring.
-    with pytest.raises(ValueError, match="max_total_pixels"):
-        smart_resize(
-            5000,
-            5000,
-            factor=28,
-            max_long_side_pixel=4000,
-            max_total_pixels=IMAGE_MAX_TOTAL_PIXELS,
-        )
+def test_compute_aligned_size_rejects_extreme_aspect_ratio():
+    with pytest.raises(ValueError, match="aspect ratio"):
+        compute_aligned_size(20001, 100, FACTOR, max_long_side_pixel=1008)
 
 
-def test_smart_resize_backward_compatible_area_bound():
-    # Without max_long_side_pixel the original Qwen-style area bound is used.
-    assert smart_resize(2048, 2048, factor=28, max_pixels=451584) == (672, 672)
+# --------------------------------------------------------------------------- #
+# resolve_max_long_side_pixel: detail tiers and explicit override
+# --------------------------------------------------------------------------- #
+@pytest.mark.parametrize(
+    ("detail", "expected"),
+    [("low", 672), ("auto", 2016), ("high", 3584), (None, 2016)],
+)
+def test_resolve_max_long_side_pixel_image(detail, expected):
+    assert resolve_max_long_side_pixel(detail, None, modality="image") == expected
+
+
+@pytest.mark.parametrize(
+    ("detail", "expected"),
+    [("low", 504), ("auto", 672), ("high", 2188), (None, 672)],
+)
+def test_resolve_max_long_side_pixel_video(detail, expected):
+    assert resolve_max_long_side_pixel(detail, None, modality="video") == expected
+
+
+@pytest.mark.parametrize("modality", ["image", "video"])
+def test_resolve_max_long_side_pixel_explicit_overrides_detail(modality):
+    assert resolve_max_long_side_pixel("high", 1024, modality=modality) == 1024
+    assert resolve_max_long_side_pixel(None, 1024, modality=modality) == 1024
+
+
+def test_processor_default_caps():
+    assert MIN_SHORT_SIDE_PIXEL == 168
+    assert MiniMaxM3VLImageProcessor().max_long_side_pixel == 2016
+    assert MiniMaxM3VLVideoProcessor().max_long_side_pixel == 672
 
 
 # --------------------------------------------------------------------------- #
-# Image processor: monotonic prompt-token counts for 252 < 504 < 1008
+# Image processor end to end: grid/patch-counter consistency and token counts
 # --------------------------------------------------------------------------- #
-def test_image_tokens_increase_with_max_long_side_pixel():
+@pytest.fixture(scope="module")
+def square_image() -> torch.Tensor:
+    # 4096x4096 so detail="high" actually reaches the 3584 cap (a 2048x2048
+    # image only ceil-aligns to 2072 and never upscales).
+    g = torch.Generator().manual_seed(0)
+    return torch.randint(0, 255, (3, 4096, 4096), dtype=torch.uint8, generator=g)
+
+
+@pytest.mark.parametrize("detail", ["low", "auto", "high"])
+def test_image_preprocess_grid_matches_patch_counter(detail, square_image):
     proc = MiniMaxM3VLImageProcessor()
-    counts = []
-    for long_side in LONG_SIDES:
-        patches = proc.get_number_of_image_patches(
-            2048, 2048, images_kwargs={"max_long_side_pixel": long_side}
-        )
-        counts.append(patches // MERGE2)
+    cap = IMAGE_DETAIL_CAPS[detail]
 
-    assert counts == [81, 324, 1296]
-    assert counts[0] < counts[1] < counts[2]
+    out = proc.preprocess([square_image], detail=detail, return_tensors="pt")
+    grid_thw = out["image_grid_thw"][0]
 
+    patches = proc.get_number_of_image_patches(
+        4096, 4096, images_kwargs={"max_long_side_pixel": cap}
+    )
+    assert int(grid_thw.prod()) == patches
 
-def test_image_processor_defaults_match_spec():
-    proc = MiniMaxM3VLImageProcessor()
-    assert proc.max_long_side_pixel is None  # opt-in
-    assert proc.min_short_side_pixel == MIN_SHORT_SIDE_PIXEL
-    assert proc.max_total_pixels == IMAGE_MAX_TOTAL_PIXELS
+    # Cross-check against compute_aligned_size: the square long side lands
+    # exactly on the cap, giving (cap / patch_size) ** 2 patches.
+    assert compute_aligned_size(4096, 4096, FACTOR, max_long_side_pixel=cap) == (
+        cap,
+        cap,
+    )
+    assert patches == (cap // PATCH_SIZE) ** 2
 
 
-def test_image_preprocess_pipeline_monotonic():
+def test_image_tokens_increase_with_detail(square_image):
     proc = MiniMaxM3VLImageProcessor()
-    image = torch.randint(0, 255, (3, 2048, 2048), dtype=torch.uint8)
-    counts = []
-    for long_side in LONG_SIDES:
-        out = proc.preprocess(
-            [image],
-            do_resize=True,
-            max_long_side_pixel=long_side,
-            return_tensors="pt",
+    counts = [
+        _merged_tokens(
+            proc.preprocess([square_image], detail=detail, return_tensors="pt")[
+                "image_grid_thw"
+            ][0]
         )
-        counts.append(_image_tokens(out["image_grid_thw"][0]))
-    assert counts == [81, 324, 1296]
+        for detail in ("low", "auto", "high")
+    ]
+    # (cap / patch_size) ** 2 / merge_size**2 for caps 672 / 2016 / 3584.
+    assert counts == [576, 5184, 16384]
+    assert counts[0] < counts[1] < counts[2]
 
 
-# --------------------------------------------------------------------------- #
-# Video processor: same monotonic behavior + volumetric (w*h*frames) cap
-# --------------------------------------------------------------------------- #
-def test_video_tokens_increase_with_max_long_side_pixel():
-    proc = MiniMaxM3VLVideoProcessor()
-    assert proc.max_total_pixels == VIDEO_MAX_TOTAL_PIXELS
-    video = torch.randint(0, 255, (4, 3, 2048, 2048), dtype=torch.uint8)
-    counts = []
-    for long_side in LONG_SIDES:
-        out = proc.preprocess(
-            videos=[video],
-            do_resize=True,
-            max_long_side_pixel=long_side,
-            return_tensors="pt",
-        )
-        counts.append(_image_tokens(out["video_grid_thw"][0]))
-    assert counts[0] < counts[1] < counts[2]
+def test_image_preprocess_explicit_cap_overrides_detail(square_image):
+    proc = MiniMaxM3VLImageProcessor()
+    out = proc.preprocess(
+        [square_image],
+        detail="high",
+        max_long_side_pixel=672,
+        return_tensors="pt",
+    )
+    # The explicit cap wins over detail="high" -> same grid as detail="low".
+    assert _merged_tokens(out["image_grid_thw"][0]) == 576
 
 
-def test_video_volumetric_cap_raises():
+# --------------------------------------------------------------------------- #
+# Video processor: volumetric (frames * h * w) cap raises unconditionally
+# --------------------------------------------------------------------------- #
+def test_video_volume_cap_raises_by_default():
     proc = MiniMaxM3VLVideoProcessor()
-    # 400 frames at a 1008-long-side square: 1008*1008*400 >> 301,056,000.
-    video = torch.randint(0, 255, (400, 3, 2048, 2048), dtype=torch.uint8)
+    # Default cap 672: 700 frames * 672 * 672 = 316,108,800 > 301,056,000.
+    # The volumetric cap raises even without any detail/max_long_side_pixel
+    # override.
+    video = torch.randint(0, 255, (700, 3, 672, 672), dtype=torch.uint8)
     with pytest.raises(ValueError, match="max_total_pixels"):
-        proc.preprocess(
-            videos=[video],
-            do_resize=True,
-            max_long_side_pixel=1008,
-            return_tensors="pt",
-        )
+        proc.preprocess(videos=[video], return_tensors="pt")
diff --git a/tests/models/multimodal/processing/test_mistral3.py b/tests/models/multimodal/processing/test_mistral3.py
index 87ca8c842bf1..6d411b9b249e 100644
--- a/tests/models/multimodal/processing/test_mistral3.py
+++ b/tests/models/multimodal/processing/test_mistral3.py
@@ -6,7 +6,7 @@
 import torch
 import torch.nn.functional as F
 from PIL import Image
-from transformers import AutoProcessor, BatchFeature
+from transformers import AutoProcessor, BatchFeature, PreTrainedConfig
 from transformers.models.pixtral import PixtralProcessor
 
 from vllm.model_executor.layers.fusion.mm_input_norm import (
@@ -79,6 +79,9 @@ def test_pixtral_patch_embed_preserves_image_order_and_strides(
     reference = []
     for image in images:
         if normalize:
+            assert isinstance(norm, FusedMMInputNorm)
+            assert isinstance(norm.weight, torch.Tensor)
+            assert isinstance(norm.bias, torch.Tensor)
             normalized = (
                 image.float() * norm.weight[:, None, None] + norm.bias[:, None, None]
             )
@@ -250,13 +253,19 @@ def test_mm_device_do_normalize():
 
     normalized_inputs = processor(prompt, mm_items=mm_items)
     normalized_values = normalized_inputs["mm_kwargs"].get_data()["pixel_values"]
+    assert isinstance(normalized_values, (list, torch.Tensor))
 
     ctx.model_config.multimodal_config.mm_device_do_normalize = True
     raw_inputs = processor(prompt, mm_items=mm_items)
     raw_values = raw_inputs["mm_kwargs"].get_data()["pixel_values"]
-    assert all(value.dtype == torch.uint8 for value in raw_values)
+    assert isinstance(raw_values, (list, torch.Tensor))
+    assert all(
+        isinstance(value, torch.Tensor) and value.dtype == torch.uint8
+        for value in raw_values
+    )
 
-    input_norm = build_mm_input_norm(ctx.model_config)
+    device = torch.device(current_platform.device_type)
+    input_norm = build_mm_input_norm(ctx.model_config).to(device)
     patch_size = processor.info.get_hf_config().vision_config.patch_size
 
     def pack_patches(image: torch.Tensor) -> torch.Tensor:
@@ -271,9 +280,11 @@ def pack_patches(image: torch.Tensor) -> torch.Tensor:
         )
 
     for raw, normalized in zip(raw_values, normalized_values):
-        output = input_norm(pack_patches(raw), normalized.dtype)
+        assert isinstance(raw, torch.Tensor)
+        assert isinstance(normalized, torch.Tensor)
+        output = input_norm(pack_patches(raw).to(device), normalized.dtype)
         torch.testing.assert_close(
-            output, pack_patches(normalized), rtol=1e-5, atol=1e-6
+            output, pack_patches(normalized).to(device), rtol=1e-5, atol=1e-6
         )
 
 
@@ -325,6 +336,6 @@ def test_lightonocr_keeps_vision_config_image_size():
     encoder_info = processor.info.get_vision_encoder_info()
     assert isinstance(encoder_info, PixtralHFEncoderInfo)
     assert not isinstance(encoder_info, Mistral3HFEncoderInfo)
-    assert encoder_info.get_image_size() == (
-        processor.info.get_hf_config().vision_config.image_size
-    )
+    vision_config = processor.info.get_hf_config().vision_config
+    assert isinstance(vision_config, PreTrainedConfig)
+    assert encoder_info.get_image_size() == vision_config.image_size
diff --git a/tests/models/multimodal/processing/test_mllama4.py b/tests/models/multimodal/processing/test_mllama4.py
index 6cf946eff279..362b327f53e7 100644
--- a/tests/models/multimodal/processing/test_mllama4.py
+++ b/tests/models/multimodal/processing/test_mllama4.py
@@ -3,6 +3,7 @@
 """Tests for mllama's multimodal preprocessing and profiling."""
 
 import pytest
+import torch
 from torch import prod
 from transformers import Llama4Config
 
@@ -36,11 +37,13 @@ def test_profiling(model_id: str, max_model_len: int):
     tokens_per_patch = ((image_size // patch_size) ** 2) // downsample_ratio
 
     mm_data = mm_inputs["mm_kwargs"].get_data()
-    chunks_per_image = prod(mm_data["patches_per_image"])
+    patches_per_image = mm_data["patches_per_image"]
+    aspect_ratios = mm_data["aspect_ratios"]
+    assert isinstance(patches_per_image, torch.Tensor)
+    assert isinstance(aspect_ratios, torch.Tensor)
+    chunks_per_image = prod(patches_per_image)
     total_num_patches = chunks_per_image * tokens_per_patch
-    num_tiles = (
-        mm_data["aspect_ratios"][0][0] * mm_data["aspect_ratios"][0][1]
-    )  # x-y separator tokens
+    num_tiles = aspect_ratios[0][0] * aspect_ratios[0][1]  # x-y separator tokens
     total_tokens = (
         total_num_patches.item() + num_tiles.item() + 3
     )  # image start, image, image end
diff --git a/tests/models/multimodal/processing/test_molmo2.py b/tests/models/multimodal/processing/test_molmo2.py
index 77c6b8253d4f..1ecefb23a310 100644
--- a/tests/models/multimodal/processing/test_molmo2.py
+++ b/tests/models/multimodal/processing/test_molmo2.py
@@ -91,7 +91,9 @@ def test_dummy_video_num_frames_override_honors_min_of_two(
         ),
     )
 
-    video, _metadata = data["video"][0]
+    videos = data["video"]
+    assert isinstance(videos, list)
+    video, _metadata = videos[0]
     assert video.shape[0] == expected_frames
 
 
diff --git a/tests/models/multimodal/processing/test_moondream3.py b/tests/models/multimodal/processing/test_moondream3.py
index a6284ae45fd1..f2c5d967f66d 100644
--- a/tests/models/multimodal/processing/test_moondream3.py
+++ b/tests/models/multimodal/processing/test_moondream3.py
@@ -98,6 +98,7 @@ def test_processor_pixel_values(
 
     # Verify pixel_values shape
     pixel_values = mm_data_result["pixel_values"]
+    assert isinstance(pixel_values, torch.Tensor)
     assert pixel_values.dim() == 5  # [batch, num_crops, C, H, W]
     assert pixel_values.shape[2] == 3  # RGB channels
     assert pixel_values.shape[3] == 378  # crop height
@@ -413,7 +414,9 @@ def test_processor_apply_accepts_non_pil_image_inputs(
     assert image_placeholders[0].length == EXPECTED_IMAGE_TOKENS
 
     mm_kwargs = processed_inputs["mm_kwargs"].get_data()
-    assert mm_kwargs["pixel_values"].shape[2:] == (3, 378, 378)
+    pixel_values = mm_kwargs["pixel_values"]
+    assert isinstance(pixel_values, torch.Tensor)
+    assert pixel_values.shape[2:] == (3, 378, 378)
 
 
 class TestMoondream3TilingLogic:
diff --git a/tests/models/multimodal/processing/test_moss_audio.py b/tests/models/multimodal/processing/test_moss_audio.py
index 294d53f23895..369f005dcc2d 100644
--- a/tests/models/multimodal/processing/test_moss_audio.py
+++ b/tests/models/multimodal/processing/test_moss_audio.py
@@ -34,6 +34,11 @@
 from vllm.sequence import IntermediateTensors
 
 
+def _uninitialized_model() -> MossAudioModel:
+    # Each test sets only the fields needed by the method under test.
+    return object.__new__(MossAudioModel)  # type: ignore[type-abstract]
+
+
 class _Tokenizer:
     def encode(self, text, **kwargs):
         return [ord(char) for char in text]
@@ -294,7 +299,7 @@ def test_moss_audio_multimodal_processor_handles_token_and_cache_paths():
 def test_moss_audio_supports_language_model_lora_only():
     assert supports_lora(MossAudioModel)
 
-    model = object.__new__(MossAudioModel)
+    model = _uninitialized_model()
     assert isinstance(model, SupportsLoRA)
 
     mapping = model.get_mm_mapping()
@@ -304,7 +309,7 @@ def test_moss_audio_supports_language_model_lora_only():
 
 
 def test_moss_audio_error_paths():
-    model = object.__new__(MossAudioModel)
+    model = _uninitialized_model()
     with pytest.raises(ValueError, match="DeepStack audio token count mismatch"):
         model._cache_deepstack_input_embeds(
             inputs_embeds=torch.zeros(4, 8),
@@ -338,7 +343,7 @@ def test_moss_audio_validates_tp_config():
 
 
 def test_moss_audio_rejects_audio_data_list_seqlen_count_mismatch():
-    model = object.__new__(MossAudioModel)
+    model = _uninitialized_model()
 
     with pytest.raises(ValueError, match="audio_data batch size"):
         model._parse_and_validate_audio_input(
@@ -349,7 +354,7 @@ def test_moss_audio_rejects_audio_data_list_seqlen_count_mismatch():
 
 @pytest.mark.parametrize("deepstack_scales", [(), (7, 11)])
 def test_moss_audio_embed_multimodal_packs_by_audio(deepstack_scales):
-    model = object.__new__(MossAudioModel)
+    model = _uninitialized_model()
     model.audio_encoder = _FakeAudioEncoder(len(deepstack_scales))
     model.audio_adapter = lambda hidden_states: hidden_states * 5
     model.deepstack_audio_merger_list = [
@@ -395,7 +400,7 @@ class _FakeLanguageModel:
         def embed_input_ids(self, input_ids):
             return torch.zeros(input_ids.shape[0], 8)
 
-    model = object.__new__(MossAudioModel)
+    model = _uninitialized_model()
     model.language_model = _FakeLanguageModel()
     model.deepstack_audio_merger_list = [object(), object()]
     model.deepstack_input_embeds = None
@@ -442,7 +447,7 @@ def fake_lm(*args, _calls=calls, **kwargs):
             return torch.ones(1, 1)
 
         _patch_pp_group(monkeypatch, first=first)
-        model = object.__new__(MossAudioModel)
+        model = _uninitialized_model()
         torch.nn.Module.__init__(model)
         model.language_model = fake_lm
         cached = IntermediateTensors({"deepstack_input_embeds_0": torch.ones(3, 8)})
@@ -476,7 +481,7 @@ def fake_lm_non_first_rank(*args, **kwargs):
         return torch.ones(1, 1)
 
     _patch_pp_group(monkeypatch, first=False)
-    model = object.__new__(MossAudioModel)
+    model = _uninitialized_model()
     torch.nn.Module.__init__(model)
     model.language_model = fake_lm_non_first_rank
     model.deepstack_input_embeds = IntermediateTensors({})
@@ -509,7 +514,7 @@ def make_model(num_layers, deepstack_layers=None):
         model.start_layer, model.end_layer = 0, num_layers
         model.layers = torch.nn.ModuleList([AddOne() for _ in range(num_layers)])
         model.norm = lambda hidden_states, residual: (hidden_states, residual)
-        model._maybe_add_hidden_state = lambda aux, *args: aux
+        monkeypatch.setattr(model, "_maybe_add_hidden_state", lambda aux, *args: aux)
         model.deepstack_inject_layer_indices = (
             range(0) if deepstack_layers is None else deepstack_layers
         )
diff --git a/tests/models/multimodal/processing/test_moss_transcribe_diarize.py b/tests/models/multimodal/processing/test_moss_transcribe_diarize.py
index 5a10e746d2a6..7f888b9446ac 100644
--- a/tests/models/multimodal/processing/test_moss_transcribe_diarize.py
+++ b/tests/models/multimodal/processing/test_moss_transcribe_diarize.py
@@ -19,9 +19,12 @@ def _question(request_prompt: str = "", hotwords: str | None = None) -> str:
         hotwords=hotwords,
         request_prompt=request_prompt,
     )
-    prompt = MossTranscribeDiarizeForConditionalGeneration.get_generation_prompt(
+    prompt_data = MossTranscribeDiarizeForConditionalGeneration.get_generation_prompt(
         params
-    )["prompt"]
+    )
+    assert isinstance(prompt_data, dict)
+    prompt = prompt_data.get("prompt")
+    assert isinstance(prompt, str)
     return prompt.split("<|audio_end|>\n", 1)[1].split("<|im_end|>", 1)[0]
 
 
diff --git a/tests/models/multimodal/processing/test_nemotron3_5_transcribe.py b/tests/models/multimodal/processing/test_nemotron3_5_transcribe.py
new file mode 100644
index 000000000000..d527d724a004
--- /dev/null
+++ b/tests/models/multimodal/processing/test_nemotron3_5_transcribe.py
@@ -0,0 +1,531 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+"""Nemotron-3.5-Transcribe config checks, audio limits and the audio tower."""
+
+import json
+from typing import Any
+
+import pytest
+import torch
+
+from vllm.transformers_utils.configs.nemotron3_5_transcribe import (
+    SPEECHLM_EMBED_EXTRA_ROWS,
+    Nemotron3_5TranscribeConfig,
+    _checked_online,
+    _preprocessor_config,
+)
+
+pytestmark = pytest.mark.cpu_test
+
+ONLINE = {
+    "window_frames": 500,
+    "left_context_frames": 50,
+    "right_context_frames": 50,
+    "fifo_len": 40,
+    "spkcache_update_period": 300,
+    "spkcache_len": 188,
+}
+
+AUDIO_CONFIG = {
+    "preprocessor": {"sample_rate": 16000, "n_fft": 512, "hop_length": 160},
+    "asr_encoder": {"subsampling_factor": 8},
+}
+
+
+def test_default_preprocessor_resolves():
+    preprocessor = _preprocessor_config({"normalize": None})
+    assert (preprocessor["n_mels"], preprocessor["win_length"]) == (64, 320)
+    # n_fft defaults to the next power of two at or above the window length.
+    assert preprocessor["n_fft"] == 512
+    longer_window = _preprocessor_config({"normalize": None, "window_size": 0.05})
+    assert longer_window["n_fft"] == 1024
+
+
+@pytest.mark.parametrize(
+    ("setting", "message"),
+    [
+        ({}, "normalize"),  # NeMo normalizes per feature when the key is missing.
+        ({"normalize": "per_feature"}, "normalize"),
+        ({"normalize": None, "log": False}, "log-mel"),
+        ({"normalize": None, "log_zero_guard_type": "clamp"}, "log_zero_guard_type"),
+    ],
+)
+def test_unsupported_mel_settings_are_refused(setting, message):
+    with pytest.raises(ValueError, match=message):
+        _preprocessor_config(setting)
+
+
+def test_window_settings():
+    assert _checked_online(dict(ONLINE)) == ONLINE
+    for key, value in (
+        ("window_frames", 0),
+        ("left_context_frames", -1),
+        ("right_context_frames", -1),
+        ("fifo_len", -1),
+        ("spkcache_update_period", 0),
+        ("spkcache_len", 0),
+    ):
+        with pytest.raises(ValueError, match=key.split("_")[0]):
+            _checked_online(ONLINE | {key: value})
+
+
+def test_transcription_pieces_fit_the_context():
+    from vllm.model_executor.models.nemotron3_5_transcribe import (
+        _transcription_clip_s,
+    )
+
+    assert _transcription_clip_s(AUDIO_CONFIG, 65536) == 3600  # one hour at most
+    # 12.5 audio tokens and 4 transcript tokens per second in 16,384 - 256 tokens.
+    assert _transcription_clip_s(AUDIO_CONFIG, 16384) == 977
+    assert _transcription_clip_s(AUDIO_CONFIG, 300) == 2
+
+
+def test_audio_length_sets_the_largest_clip():
+    from vllm.model_executor.models.nemotron3_5_transcribe import (
+        _max_audio_samples,
+        _num_audio_embeddings,
+    )
+
+    def max_tokens(length, seq_len=65536):
+        samples = _max_audio_samples(AUDIO_CONFIG, seq_len, length)
+        return _num_audio_embeddings(AUDIO_CONFIG, samples)
+
+    assert max_tokens(None) == 45_000  # one hour by default
+    # 69 minutes, as `--limit-mm-per-prompt '{"audio": {"length": 66240000}}'`
+    # sets it, so a 65-minute clip (48,750 embeddings) fits the encoder cache.
+    assert max_tokens(66_240_000) == 51_750
+    # Never more than the context holds beside a short prompt.
+    assert max_tokens(115_200_000) == 65_536 - 256
+
+
+ENCODER = {
+    "d_model": 64,
+    "n_heads": 2,
+    "n_layers": 1,
+    "self_attention_model": "rope",
+    "subsampling_factor": 8,
+}
+
+# The smallest config.json in the released checkpoint's NeMo SpeechLM layout.
+CHECKPOINT_CONFIG: dict[str, Any] = {
+    "architectures": ["Nemotron3_5TranscribeForConditionalGeneration"],
+    "model_type": "nemo_speechlm",
+    "llm_config": {"model_type": "nemotron_h", "hidden_size": 64, "vocab_size": 1000},
+    "perception": {
+        "preprocessor": {"features": 128, "normalize": None, "window_size": 0.025},
+        "modality_adapter": {"_target_": "IdentityConnector", "d_model": 64},
+        "output_dim": 64,
+    },
+    "pe_encoder_config": {
+        "asr_encoder_type": "transformer",
+        "asr_encoder_cfg": ENCODER,
+        "diarization_model_cfg": {
+            "encoder": ENCODER,
+            "high_resolution": True,
+            "sortformer_modules": {"use_learnable_sil_emb": True},
+        },
+        "speaker_feature_mode": "thresholded",
+        "speaker_activity_threshold": 0.5,
+    },
+}
+
+
+def test_released_checkpoint_layout_loads(tmp_path):
+    from vllm.config import ModelConfig
+
+    (tmp_path / "config.json").write_text(json.dumps(CHECKPOINT_CONFIG))
+    config = ModelConfig(str(tmp_path)).hf_config
+    assert isinstance(config, Nemotron3_5TranscribeConfig)
+    preprocessor = config.audio_config["preprocessor"]
+    assert (preprocessor["win_length"], preprocessor["hop_length"]) == (400, 160)
+    assert config.audio_config["online"] == ONLINE
+    # The audio placeholder sits in embedding rows past the trained vocabulary.
+    assert config.output_vocab_size == 1000
+    assert config.text_config.vocab_size == 1000 + SPEECHLM_EMBED_EXTRA_ROWS
+
+
+@pytest.mark.parametrize(
+    ("changes", "hint"),
+    [
+        # This model's structure under NeMo's generic name, as older exports have.
+        (
+            {},
+            'set "architectures": ["Nemotron3_5TranscribeForConditionalGeneration"]',
+        ),
+        # Another NeMo SpeechLM: no Parallel Expert Encoder.
+        ({"pe_encoder_config": None}, "NeMo Speech vLLM plugin"),
+    ],
+)
+def test_other_nemo_speechlm_architectures_are_refused(tmp_path, changes, hint):
+    from vllm.transformers_utils.config import get_config
+
+    other = (
+        CHECKPOINT_CONFIG
+        | changes
+        | {"architectures": ["NeMoSpeechLMForConditionalGeneration"]}
+    )
+    (tmp_path / "config.json").write_text(json.dumps(other))
+    with pytest.raises(ValueError) as excinfo:
+        get_config(str(tmp_path), trust_remote_code=False)
+    assert hint in str(excinfo.value)
+
+
+def _tiny_encoder(d_model: int, rotary_fraction: float, qk_norm: bool) -> dict:
+    return {
+        "feat_in": 16,
+        "d_model": d_model,
+        "n_heads": 4,
+        "n_layers": 2,
+        "ff_expansion": 4.0,
+        "qkv_bias": False,
+        "qk_norm": qk_norm,
+        "pre_block_norm": True,
+        "subsampling_factor": 8,
+        "rope_base": 10000.0,
+        "rotary_fraction": rotary_fraction,
+        "xscaling": False,
+    }
+
+
+# 0.8 s windows, so a few seconds of audio spans several windows and fills and
+# compresses the speaker cache. Continuous speaker features carry the speaker
+# probabilities into the embeddings unthresholded.
+TINY_AUDIO_CONFIG = {
+    "preprocessor": {
+        "sample_rate": 16000,
+        "win_length": 400,
+        "hop_length": 160,
+        "n_fft": 512,
+        "n_mels": 16,
+        "preemph": 0.97,
+        "log_zero_guard_value": 2**-24,
+        "mag_power": 2.0,
+        "pad_value": 0.0,
+    },
+    "asr_encoder": _tiny_encoder(32, 0.5, True),
+    "speaker_encoder": _tiny_encoder(24, 1.0, False),
+    "sortformer": {
+        "num_speakers": 4,
+        "fc_d_model": 24,
+        "tf_d_model": 12,
+        "spkcache_sil_frames_per_spk": 1,
+        "pred_score_threshold": 0.25,
+        "max_index": 99999,
+        "scores_boost_latest": 0.05,
+        "strong_boost_rate": 0.75,
+        "weak_boost_rate": 1.5,
+        "min_pos_scores_rate": 0.5,
+    },
+    "online": {
+        "window_frames": 10,
+        "left_context_frames": 2,
+        "right_context_frames": 2,
+        "fifo_len": 3,
+        "spkcache_update_period": 6,
+        "spkcache_len": 12,
+    },
+    "fusion": {
+        "speaker_feature_mode": "continuous",
+        "speaker_activity_threshold": None,
+        "spk_kernel_scale": 1.0,
+    },
+    "projection": {"input_dim": 32, "output_dim": 40},
+}
+
+
+@pytest.fixture
+def tiny_tower(dist_init):
+    from vllm.model_executor.layers.linear import ReplicatedLinear
+    from vllm.model_executor.models.nemotron3_5_transcribe_audio import (
+        Nemotron3_5TranscribeAudioTower,
+    )
+    from vllm.platforms import current_platform
+
+    torch.manual_seed(0)
+    # vLLM registers the encoder attention ops for the platform's device only.
+    tower = (
+        Nemotron3_5TranscribeAudioTower(TINY_AUDIO_CONFIG)
+        .eval()
+        .to(current_platform.device_type)
+    )
+    with torch.no_grad():
+        # vLLM's linear layers allocate their weights uninitialized.
+        for module in tower.modules():
+            if isinstance(module, ReplicatedLinear):
+                module.weight.normal_(0.0, module.input_size**-0.5)
+                if module.bias is not None:
+                    module.bias.normal_(0.0, 0.1)
+        featurizer = tower.preprocessor.featurizer
+        featurizer.fb.uniform_()
+        window = torch.hann_window(featurizer.win_length, periodic=False)
+        featurizer.window.copy_(window)
+        sortformer = tower.encoder.diarization_model.sortformer_modules
+        sortformer.learnable_sil_emb.normal_()
+        for module in tower.modules():
+            if isinstance(module, torch.nn.LayerNorm):
+                module.weight.normal_(1.0, 0.1)
+                module.bias.normal_(0.0, 0.1)
+        # Weight loading ends with this step; on CPU it selects the GEMM kernel.
+        for module in tower.modules():
+            if isinstance(module, ReplicatedLinear):
+                module.quant_method.process_weights_after_loading(module)
+    return tower
+
+
+@pytest.mark.parametrize("num_samples", [160, 161, 1_281, 12_345, 12_960, 48_001])
+def test_audio_tower_returns_one_embedding_per_audio_token(tiny_tower, num_samples):
+    from vllm.model_executor.models.nemotron3_5_transcribe import (
+        _num_audio_embeddings,
+    )
+
+    clip = torch.randn(num_samples) * 0.1
+    with torch.inference_mode():
+        (embeddings,) = tiny_tower.encode([clip], torch.tensor([num_samples]))
+    num_tokens = _num_audio_embeddings(TINY_AUDIO_CONFIG, num_samples)
+    assert embeddings.shape == (num_tokens, 40)
+
+
+def test_decoder_weights_read_both_checkpoint_layouts():
+    from vllm.model_executor.models.nemotron3_5_transcribe import (
+        Nemotron3_5TranscribeForConditionalGeneration as Model,
+    )
+    from vllm.model_executor.models.nemotron3_5_transcribe import _decoder_weights
+
+    nemo = {
+        "llm.model.embed_tokens.weight": torch.ones(8, 4),
+        "llm.model.layers.0.mixer._fp32_params.A_log": torch.ones(2),
+        "llm.model.layers.1.mixer.experts.down_projs": torch.arange(24.0).view(2, 3, 4),
+        "llm.model.norm.weight": torch.ones(4),
+    }
+    # The FP8 and NVFP4 checkpoints: Hugging Face names, one tensor per expert.
+    hf = {
+        "backbone.embeddings.weight": torch.ones(8, 4),
+        "backbone.layers.1.mixer.experts.0.up_proj.weight_scale": torch.ones(()),
+        "lm_head.weight": torch.ones(8, 4),
+    }
+    mapped = Model.hf_to_vllm_mapper.apply((nemo | hf).items())
+    decoder = ((name.removeprefix("language_model."), t) for name, t in mapped)
+    out = dict(_decoder_weights(decoder, vocab_size=10))
+    assert sorted(out) == [
+        "lm_head.weight",
+        "model.embed_tokens.weight",
+        "model.embeddings.weight",
+        "model.layers.0.mixer.A_log",
+        "model.layers.1.mixer.experts.0.down_proj.weight",
+        "model.layers.1.mixer.experts.0.up_proj.weight_scale",
+        "model.layers.1.mixer.experts.1.down_proj.weight",
+        "model.norm_f.weight",
+    ]
+    for name in (
+        "model.embed_tokens.weight",
+        "model.embeddings.weight",
+        "lm_head.weight",
+    ):
+        assert out[name].shape == (10, 4) and out[name][8:].eq(0).all()
+    assert torch.equal(
+        out["model.layers.1.mixer.experts.1.down_proj.weight"],
+        nemo["llm.model.layers.1.mixer.experts.down_projs"][1].t(),
+    )
+    # Names outside both layouts stay unmapped, so load_weights refuses them.
+    unknown = [("llm.vision.weight", torch.ones(1))]
+    assert [name for name, _ in Model.hf_to_vllm_mapper.apply(unknown)] == [
+        "llm.vision.weight"
+    ]
+
+
+def test_only_audio_has_a_placeholder():
+    from vllm.model_executor.models.nemotron3_5_transcribe import (
+        AUDIO_PLACEHOLDER,
+    )
+    from vllm.model_executor.models.nemotron3_5_transcribe import (
+        Nemotron3_5TranscribeForConditionalGeneration as Model,
+    )
+
+    for modality in ("audio", "audio_embeds"):
+        assert Model.get_placeholder_str(modality, 1) == AUDIO_PLACEHOLDER
+    with pytest.raises(ValueError, match="Only audio modality is supported"):
+        Model.get_placeholder_str("image", 1)
+
+
+def test_quantization_layer_names_map_onto_the_model():
+    from vllm.model_executor.models.nemotron3_5_transcribe import (
+        Nemotron3_5TranscribeForConditionalGeneration as Model,
+    )
+
+    mapper = Model.hf_to_vllm_mapper.get_rename_mapper()
+    assert mapper.apply_list(
+        [
+            "backbone.layers.1.mixer.gate",
+            "lm_head",
+            "perception.encoder.asr_encoder.pre_encode",
+            # Names the NVFP4 recipe already gives in vLLM's form stay as they are.
+            "language_model.model.layers.0.mixer.in_proj",
+            "mtp.layers.1.mixer.gate",
+        ]
+    ) == [
+        "language_model.model.layers.1.mixer.gate",
+        "language_model.lm_head",
+        "audio_tower.encoder.asr_encoder.pre_encode",
+        "language_model.model.layers.0.mixer.in_proj",
+        "mtp.layers.1.mixer.gate",
+    ]
+
+
+@pytest.mark.usefixtures("dist_init")
+def test_fp8_recipe_quantizes_the_asr_encoder_linear_layers():
+    from vllm.model_executor.layers.quantization.modelopt import ModelOptFp8Config
+    from vllm.model_executor.model_loader.utils import configure_quant_config
+    from vllm.model_executor.models.nemotron3_5_transcribe import (
+        Nemotron3_5TranscribeForConditionalGeneration as Model,
+    )
+    from vllm.model_executor.models.nemotron3_5_transcribe_audio import (
+        Nemotron3_5TranscribeAudioTower,
+    )
+
+    # As in the FP8 checkpoint's config.json.
+    quant_config = ModelOptFp8Config.from_config(
+        {
+            "quant_method": "modelopt",
+            "quant_algo": "FP8",
+            "config_groups": {
+                "group_0": {
+                    "targets": ["Linear"],
+                    "weights": {"dynamic": False, "num_bits": 8, "type": "float"},
+                    "input_activations": {
+                        "dynamic": False,
+                        "num_bits": 8,
+                        "type": "float",
+                    },
+                }
+            },
+            "ignore": [
+                "perception.encoder.diarization_model*",
+                "perception.encoder.asr_encoder.pre_encode*",
+                "perception.proj",
+            ],
+        }
+    )
+    configure_quant_config(quant_config, Model)
+    tower = Nemotron3_5TranscribeAudioTower(
+        TINY_AUDIO_CONFIG, quant_config=quant_config, prefix="audio_tower"
+    )
+    fp8_layers = sorted(
+        name
+        for name, module in tower.named_modules()
+        if isinstance(getattr(module, "weight", None), torch.Tensor)
+        and module.weight.dtype == torch.float8_e4m3fn
+    )
+    assert fp8_layers == sorted(
+        f"encoder.asr_encoder.layers.{i}.{layer}"
+        for i in range(2)
+        for layer in ("attn.w_qkv", "attn.out_proj", "ffn.net.0", "ffn.net.3")
+    )
+
+    # The checkpoint stores each scale as a scalar.
+    weights = []
+    for name, target in [*tower.named_parameters(), *tower.named_buffers()]:
+        if name == "encoder.diar_kernel":
+            continue
+        tensor = torch.full(target.shape, 0.5).to(target.dtype)
+        if name.endswith("_scale"):
+            tensor = torch.tensor(0.25)
+        weights.append((name, tensor))
+    tower.load_weights(weights)
+    w_qkv = tower.get_submodule("encoder.asr_encoder.layers.0.attn.w_qkv")
+    assert w_qkv.weight_scale.item() == w_qkv.input_scale.item() == 0.25
+
+
+def test_mtp_draft_config_keeps_the_quantization_recipe():
+    from vllm.config.speculative import SpeculativeConfig
+
+    recipe = {"quant_method": "modelopt", "quant_algo": "FP8", "ignore": ["lm_head"]}
+    llm_config = CHECKPOINT_CONFIG["llm_config"] | {"num_nextn_predict_layers": 1}
+    config = Nemotron3_5TranscribeConfig.from_dict(
+        CHECKPOINT_CONFIG | {"llm_config": llm_config, "quantization_config": recipe}
+    )
+    draft = SpeculativeConfig.hf_config_override(config)
+    assert draft.architectures == ["Nemotron3_5TranscribeMTPModel"]
+    # The recipe lives on the SpeechLM config; the MTP head is quantized too.
+    assert draft.quantization_config == recipe
+
+
+def test_audio_tower_encodes_each_clip_as_if_alone(tiny_tower):
+    torch.manual_seed(1)
+    # The 1.3 s and 1.5 s clips span two windows and share a padded batch; the
+    # 3 s clip spans four windows and is encoded apart, between them in order.
+    clips = [torch.randn(num_samples) * 0.1 for num_samples in (20_800, 48_000, 24_000)]
+    lengths = torch.tensor([clip.numel() for clip in clips])
+    with torch.inference_mode():
+        together = tiny_tower.encode(clips, lengths)
+        alone = [
+            tiny_tower.encode([clip], lengths[index : index + 1])[0]
+            for index, clip in enumerate(clips)
+        ]
+    for got, expected in zip(together, alone):
+        torch.testing.assert_close(got, expected, rtol=1e-4, atol=1e-4)
+
+
+def test_encoder_packs_each_rows_valid_frames(tiny_tower):
+    encoder = tiny_tower.encoder.asr_encoder
+    packed = encoder.pack(torch.tensor([2, 0, 3]), 4, torch.device("cpu"))
+    assert packed.rows.tolist() == [0, 1, 8, 9, 10]
+    assert packed.positions.tolist() == [0, 1, 0, 1, 2]
+    assert packed.cu_seqlens.tolist() == [0, 2, 5]
+
+
+def test_encoder_runs_only_the_valid_frames(tiny_tower):
+    encoder = tiny_tower.encoder.asr_encoder
+    device = encoder.final_norm.weight.device
+    torch.manual_seed(2)
+    x = torch.randn(3, 6, encoder.d_model, device=device)
+    lengths = torch.tensor([6, 0, 3])
+    with torch.inference_mode():
+        states, _ = encoder(
+            x, lengths.to(device), bypass_pre_encode=True, lengths_cpu=lengths
+        )
+        alone = [
+            encoder(
+                x[row : row + 1, :length],
+                lengths[row : row + 1].to(device),
+                bypass_pre_encode=True,
+            )[0][0]
+            for row, length in ((0, 6), (2, 3))
+        ]
+    assert not states[1].any()
+    assert not states[2, :, 3:].any()
+    torch.testing.assert_close(states[0], alone[0], rtol=1e-4, atol=1e-4)
+    torch.testing.assert_close(states[2, :, :3], alone[1], rtol=1e-4, atol=1e-4)
+
+
+@pytest.mark.parametrize("name", ["asr_encoder", "diarization_model.encoder"])
+@pytest.mark.parametrize("num_frames", [0, 6])
+def test_encoder_returns_zeros_when_no_row_has_frames(tiny_tower, name, num_frames):
+    encoder = tiny_tower.encoder.get_submodule(name)
+    device = encoder.final_norm.weight.device
+    x = torch.randn(3, num_frames, encoder.d_model, device=device)
+    lengths = torch.zeros(3, dtype=torch.int64)
+    with torch.inference_mode():
+        states, _ = encoder(
+            x, lengths.to(device), bypass_pre_encode=True, lengths_cpu=lengths
+        )
+    assert states.shape == (3, encoder.d_model, num_frames)
+    assert not states.any()
+
+
+@pytest.mark.usefixtures("dist_init")
+def test_fp8_encoder_attention_is_refused(tmp_path):
+    from vllm.config import ModelConfig, VllmConfig, set_current_vllm_config
+    from vllm.model_executor.models.nemotron3_5_transcribe_audio import (
+        Nemotron3_5TranscribeAudioTower,
+    )
+
+    (tmp_path / "config.json").write_text(json.dumps(CHECKPOINT_CONFIG))
+    model_config = ModelConfig(str(tmp_path), mm_encoder_attn_dtype="fp8")
+    # vLLM's encoder attention refuses FP8 with another message, except on the
+    # FlashInfer cuDNN and ROCm AITER backends; the tower refuses it on all.
+    with (
+        set_current_vllm_config(VllmConfig(model_config=model_config)),
+        pytest.raises(ValueError, match=r"Nemotron-3\.5-Transcribe does not support"),
+    ):
+        Nemotron3_5TranscribeAudioTower(TINY_AUDIO_CONFIG)
diff --git a/tests/models/multimodal/processing/test_nemotron_parse.py b/tests/models/multimodal/processing/test_nemotron_parse.py
index 019c06eb15a2..1acbe2eb84c0 100644
--- a/tests/models/multimodal/processing/test_nemotron_parse.py
+++ b/tests/models/multimodal/processing/test_nemotron_parse.py
@@ -37,12 +37,14 @@ def test_mm_device_do_normalize(width: int, height: int):
     normalized_values = processor(prompt, mm_items=mm_items)["mm_kwargs"].get_data()[
         "pixel_values"
     ]
+    assert isinstance(normalized_values, torch.Tensor)
 
     ctx.model_config.multimodal_config.mm_device_do_normalize = True
     processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config)
     raw_values = processor(prompt, mm_items=mm_items)["mm_kwargs"].get_data()[
         "pixel_values"
     ]
+    assert isinstance(raw_values, torch.Tensor)
     assert raw_values.dtype == torch.uint8
     assert raw_values.shape == normalized_values.shape
 
@@ -71,7 +73,11 @@ def test_array_image_inputs(mm_device_do_normalize: bool, layout: str):
 
     def pixel_values(img: object) -> torch.Tensor:
         mm_items = processor.info.parse_mm_data({"image": [img]})
-        return processor([0], mm_items=mm_items)["mm_kwargs"].get_data()["pixel_values"]
+        values = processor([0], mm_items=mm_items)["mm_kwargs"].get_data()[
+            "pixel_values"
+        ]
+        assert isinstance(values, torch.Tensor)
+        return values
 
     torch.testing.assert_close(pixel_values(item), pixel_values(image), rtol=0, atol=0)
 
@@ -88,5 +94,6 @@ def test_explicit_kwargs_without_device_normalize():
 
     mm_items = processor.info.parse_mm_data({"image": [_noise_image(310, 470)]})
     values = processor([0], mm_items=mm_items)["mm_kwargs"].get_data()["pixel_values"]
+    assert isinstance(values, torch.Tensor)
     # The remote processor ignores these kwargs and still normalises on the host.
     assert values.is_floating_point()
diff --git a/tests/models/multimodal/processing/test_nemotron_parse_35.py b/tests/models/multimodal/processing/test_nemotron_parse_35.py
index f8562c3e5220..11fbcc1ae45c 100644
--- a/tests/models/multimodal/processing/test_nemotron_parse_35.py
+++ b/tests/models/multimodal/processing/test_nemotron_parse_35.py
@@ -45,11 +45,15 @@ def test_causal_prefix_preserves_host_normalized_packed_images(
     mm_items = processor.info.parse_mm_data({"image": [image]})
     actual = processor([1], mm_items=mm_items)["mm_kwargs"].get_data()
 
-    assert actual["pixel_values"].is_floating_point()
-    assert actual["pixel_values"].shape[-1] == 3 * 16**2
+    pixel_values = actual["pixel_values"]
+    assert isinstance(pixel_values, torch.Tensor)
+    assert pixel_values.is_floating_point()
+    assert pixel_values.shape[-1] == 3 * 16**2
     for name in ("pixel_values", "image_sizes", "content_sizes"):
+        value = actual[name]
+        assert isinstance(value, torch.Tensor)
         torch.testing.assert_close(
-            actual[name], expected[name].to(actual[name].dtype), rtol=0, atol=0
+            value, expected[name].to(value.dtype), rtol=0, atol=0
         )
 
 
@@ -98,4 +102,4 @@ def test_causal_prefix_fields_config_handles_empty_cached_hf_inputs() -> None:
 
     fields = processor._get_mm_fields_config(BatchFeature(), {})
 
-    assert fields["pixel_values"].field.slices == []
+    assert getattr(fields["pixel_values"].field, "slices") == []  # noqa: B009
diff --git a/tests/models/multimodal/processing/test_nemotron_vl.py b/tests/models/multimodal/processing/test_nemotron_vl.py
index b1c1e11892c7..bc2e9782c01e 100644
--- a/tests/models/multimodal/processing/test_nemotron_vl.py
+++ b/tests/models/multimodal/processing/test_nemotron_vl.py
@@ -5,6 +5,7 @@
 from collections.abc import Mapping
 
 import pytest
+import torch
 from PIL import Image
 from transformers import PreTrainedConfig
 
@@ -77,7 +78,9 @@ def _run_check(
     # Ensure we have the right number of placeholders per num_crops size
     image_token_id = tokenizer.convert_tokens_to_ids("<image>")
     img_tok_count = processed_inputs["prompt_token_ids"].count(image_token_id)
-    pixel_shape = processed_inputs["mm_kwargs"].get_data()["pixel_values_flat"].shape
+    pixel_values = processed_inputs["mm_kwargs"].get_data()["pixel_values_flat"]
+    assert isinstance(pixel_values, torch.Tensor)
+    pixel_shape = pixel_values.shape
     print("Image token count:", img_tok_count, "Pixel shape:", pixel_shape)
     assert img_tok_count == 256 * total_expected_num_patches
     assert pixel_shape[0] == total_expected_num_patches
diff --git a/tests/models/multimodal/processing/test_openvla.py b/tests/models/multimodal/processing/test_openvla.py
index 36775d53e40f..d88c716035bd 100644
--- a/tests/models/multimodal/processing/test_openvla.py
+++ b/tests/models/multimodal/processing/test_openvla.py
@@ -5,6 +5,7 @@
 import numpy as np
 import pytest
 import torch
+from jinja2.sandbox import ImmutableSandboxedEnvironment
 from PIL import Image
 from transformers import LlamaConfig
 
@@ -14,6 +15,7 @@
     OpenVLAProcessingInfo,
 )
 from vllm.multimodal.parse import ImageProcessorItems, MultiModalDataItems
+from vllm.transformers_utils.chat_templates import get_chat_template_fallback_path
 from vllm.transformers_utils.configs.openvla import OpenVLAConfig
 from vllm.transformers_utils.processors.openvla import (
     IMAGENET_MEAN,
@@ -164,8 +166,10 @@ def test_openvla_image_processor_outputs_pixel_values() -> None:
 
     output = processor(images=image)
 
-    assert output["pixel_values"].shape == (1, 6, 224, 224)
-    assert output["pixel_values"].dtype == torch.float32
+    pixel_values = output["pixel_values"]
+    assert isinstance(pixel_values, torch.Tensor)
+    assert pixel_values.shape == (1, 6, 224, 224)
+    assert pixel_values.dtype == torch.float32
 
 
 def test_openvla_processing_info_token_counts() -> None:
@@ -196,6 +200,7 @@ def test_openvla_prompt_update_inserts_image_tokens_after_bos() -> None:
     ] == [(1, 1)]
     assert content.full == [32000] * 256
 
+    assert callable(content.is_embed)
     is_embed = content.is_embed(content.full)
     assert is_embed.dtype == torch.bool
     assert is_embed.tolist() == [True] * 256
@@ -203,3 +208,41 @@ def test_openvla_prompt_update_inserts_image_tokens_after_bos() -> None:
 
 def test_openvla_placeholder_string() -> None:
     assert OpenVLAForActionPrediction.get_placeholder_str("image", 0) is None
+
+
+# OpenVLA's training prompt; predict_action() appends trailing space when the
+# prompt does not already end with it.
+_OPENVLA_PROMPT = (
+    "<s>In: What action should the robot take to pick up the red block?\nOut: "
+)
+
+
+@pytest.mark.parametrize(
+    "text",
+    [
+        "pick up the red block",
+        " Pick up the red block. ",
+        "In: What action should the robot take to pick up the red block?\nOut:",
+        "In: What action should the robot take to pick up the red block?\nOut: ",
+    ],
+)
+@pytest.mark.parametrize("with_image", [False, True])
+def test_openvla_chat_template(text: str, with_image: bool) -> None:
+    path = get_chat_template_fallback_path("openvla", "openvla/openvla-7b")
+    assert path is not None
+    env = ImmutableSandboxedEnvironment(trim_blocks=True, lstrip_blocks=True)
+    template = env.from_string(path.read_text())
+
+    content: str | list[dict] = text
+    if with_image:
+        content = [
+            {"type": "image", "image": None},
+            {"type": "text", "text": text},
+        ]
+    messages = [
+        {"role": "system", "content": "ignored"},
+        {"role": "user", "content": content},
+    ]
+
+    # No image placeholder: the processor inserts image tokens after BOS.
+    assert template.render(messages=messages, bos_token="<s>") == _OPENVLA_PROMPT
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 fda9fa7aa8f4..e6c15dd6b383 100644
--- a/tests/models/multimodal/processing/test_qwen2_5_omni_embed.py
+++ b/tests/models/multimodal/processing/test_qwen2_5_omni_embed.py
@@ -48,7 +48,7 @@ def test_precomputed_vision_embeddings_preserve_items(model_version, modality):
     torch.nn.Module.__init__(model)
     model.visual = Mock(dtype=torch.float32, spatial_merge_size=2)
     embeds = torch.arange(80, dtype=torch.float64).reshape(10, 8)
-    outputs = model.embed_multimodal(
+    outputs = getattr(model, "embed_multimodal")(  # noqa: B009
         **{
             f"{modality}_embeds": embeds,
             f"{modality}_grid_thw": torch.tensor([[1, 2, 4], [2, 4, 4]]),
diff --git a/tests/models/multimodal/processing/test_qwen2_vl.py b/tests/models/multimodal/processing/test_qwen2_vl.py
index 133880b08167..21f45d1f09dd 100644
--- a/tests/models/multimodal/processing/test_qwen2_vl.py
+++ b/tests/models/multimodal/processing/test_qwen2_vl.py
@@ -343,7 +343,9 @@ def test_processor_override(
     hf_processor = processor.info.get_hf_processor(**hf_processor_mm_kwargs)
     image_token_id = tokenizer.convert_tokens_to_ids(hf_processor.image_token)
     img_tok_count = processed_inputs["prompt_token_ids"].count(image_token_id)
-    pixel_shape = processed_inputs["mm_kwargs"].get_data()["pixel_values"].shape
+    pixel_values = processed_inputs["mm_kwargs"].get_data()["pixel_values"]
+    assert isinstance(pixel_values, torch.Tensor)
+    pixel_shape = pixel_values.shape
 
     assert img_tok_count == expected_toks_per_img * num_imgs
     assert pixel_shape[0] == expected_pixels_shape[0] * num_imgs
@@ -371,6 +373,7 @@ def test_get_image_size_with_most_features(
         limit_mm_per_prompt={"image": 1},
     )
     processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config)
+    assert isinstance(processor.info, Qwen2VLProcessingInfo)
 
     hf_processor = processor.info.get_hf_processor(**mm_processor_kwargs)
     merge_size = processor.info.get_hf_config().vision_config.spatial_merge_size
@@ -391,7 +394,9 @@ def test_get_image_size_with_most_features(
             mm_items=processor.info.parse_mm_data(mm_data),
             hf_processor_mm_kwargs=mm_processor_kwargs,
         )
-        grid_thw = processed_inputs["mm_kwargs"].get_data()["image_grid_thw"].tolist()
+        image_grid_thw = processed_inputs["mm_kwargs"].get_data()["image_grid_thw"]
+        assert isinstance(image_grid_thw, torch.Tensor)
+        grid_thw = image_grid_thw.tolist()
         t, h, w = grid_thw[0]
         tokens = (t * h * w) // (merge_size**2)
         assert tokens < max_tokens
@@ -422,7 +427,9 @@ def test_qwen2_5_vl_video_fps_to_second_per_grid_ts(model_id: str) -> None:
     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"]
+        item = processed["mm_kwargs"]["video"][0]
+        assert item is not None
+        spg = item.get_data()["second_per_grid_ts"]
         assert float(spg) == pytest.approx(temporal_patch_size / fps, rel=2e-2)
 
 
@@ -458,6 +465,7 @@ def test_mm_device_do_normalize(
     pixel_values_with_normalize = processed_inputs_with_normalize[
         "mm_kwargs"
     ].get_data()["pixel_values"]
+    assert isinstance(pixel_values_with_normalize, torch.Tensor)
     dtype = pixel_values_with_normalize.dtype
 
     processed_inputs_without_normalize = processor(
@@ -468,6 +476,7 @@ def test_mm_device_do_normalize(
     pixel_values_without_normalize = processed_inputs_without_normalize[
         "mm_kwargs"
     ].get_data()["pixel_values"]
+    assert isinstance(pixel_values_without_normalize, torch.Tensor)
 
     ctx.model_config.multimodal_config.mm_device_do_normalize = True
     input_norm = build_mm_input_norm(ctx.model_config).to(device)
diff --git a/tests/models/multimodal/processing/test_qwen3_vl.py b/tests/models/multimodal/processing/test_qwen3_vl.py
index 0c74a86ff000..f0c855efa984 100644
--- a/tests/models/multimodal/processing/test_qwen3_vl.py
+++ b/tests/models/multimodal/processing/test_qwen3_vl.py
@@ -87,6 +87,7 @@ def test_processor_partial_scoped_image_size(model_id: str) -> None:
         limit_mm_per_prompt={"image": 1, "video": 0},
     )
     processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config)
+    assert isinstance(processor.info, Qwen3VLProcessingInfo)
 
     image = np.zeros((128, 128, 3), dtype=np.uint8)
     min_pixels = 256 * 32 * 32
@@ -95,9 +96,11 @@ def test_processor_partial_scoped_image_size(model_id: str) -> None:
 
     def process(hf_mm_kwargs: dict[str, Any]) -> torch.Tensor:
         mm_items = processor.info.parse_mm_data({"image": [image]})
-        return processor._apply_hf_processor_main(mm_items, hf_mm_kwargs)[
+        grid = processor._apply_hf_processor_main(mm_items, hf_mm_kwargs)[
             "image_grid_thw"
         ]
+        assert isinstance(grid, torch.Tensor)
+        return grid
 
     stock_grid = process({})
     complete_grid = process({"images_kwargs": {"size": complete_size}})
@@ -127,6 +130,7 @@ def test_processor_partial_scoped_video_size(model_id: str) -> None:
         limit_mm_per_prompt={"image": 0, "video": 1},
     )
     processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config)
+    assert isinstance(processor.info, Qwen3VLProcessingInfo)
 
     min_pixels = 512 * 32 * 32
     default_size = dict(processor.info.get_video_processor().size)
@@ -134,9 +138,11 @@ def test_processor_partial_scoped_video_size(model_id: str) -> None:
 
     def process(hf_mm_kwargs: dict[str, Any]) -> torch.Tensor:
         mm_items = processor.info.parse_mm_data(_build_video_mm_data(num_frames=8))
-        return processor._apply_hf_processor_main(mm_items, hf_mm_kwargs)[
+        grid = processor._apply_hf_processor_main(mm_items, hf_mm_kwargs)[
             "video_grid_thw"
         ]
+        assert isinstance(grid, torch.Tensor)
+        return grid
 
     stock_grid = process({})
     complete_grid = process({"videos_kwargs": {"size": complete_size}})
@@ -302,6 +308,7 @@ def _probe_mm_token_budgets(
         limit_mm_per_prompt={"image": 1, "video": 1},
     )
     info = MULTIMODAL_REGISTRY.create_processor(ctx.model_config).info
+    assert isinstance(info, Qwen3VLProcessingInfo)
     video = info.get_max_video_tokens(
         seq_len=500000, mm_counts={"video": 1, "image": 1}
     )
@@ -381,7 +388,9 @@ def process_single(hf_kwargs: dict[str, Any]) -> torch.Tensor:
             mm_items=processor.info.parse_mm_data({"video": [videos[0]]}),
             hf_processor_mm_kwargs=hf_kwargs,
         )
-        return single["mm_kwargs"].get_data()["video_grid_thw"][0]
+        grids = single["mm_kwargs"].get_data()["video_grid_thw"]
+        assert isinstance(grids, torch.Tensor)
+        return grids[0]
 
     expected_grids = [process_single(kwargs) for kwargs in scalar_kwargs]
     assert not torch.equal(expected_grids[0], expected_grids[1])
@@ -466,7 +475,9 @@ def test_processor_profiles_fewer_videos_than_configured_options(
     profiled = processor.get_dummy_mm_inputs({"video": 1})
 
     assert len(profiled["mm_placeholders"]["video"]) == 1
-    assert profiled["mm_kwargs"].get_data()["video_grid_thw"].shape[0] == 1
+    video_grid_thw = profiled["mm_kwargs"].get_data()["video_grid_thw"]
+    assert isinstance(video_grid_thw, torch.Tensor)
+    assert video_grid_thw.shape[0] == 1
 
 
 def _build_video_embeds_mm_data(
@@ -584,7 +595,9 @@ def test_dummy_video_spreads_budget_when_frame_cap_enabled(model_id: str) -> Non
     capped_dummy = capped.dummy_inputs.get_dummy_mm_data(
         1024, {"video": 1}, MultiModalDummyOptions()
     )
-    capped_frames = capped_dummy["video"][0][0].shape[0]
+    capped_videos = capped_dummy["video"]
+    assert isinstance(capped_videos, list)
+    capped_frames = capped_videos[0][0].shape[0]
     assert capped_frames == 16, (
         f"Expected the dummy to spread the budget over 16 frames, got {capped_frames}"
     )
@@ -598,7 +611,9 @@ def test_dummy_video_spreads_budget_when_frame_cap_enabled(model_id: str) -> Non
     uncapped_dummy = uncapped.dummy_inputs.get_dummy_mm_data(
         1024, {"video": 1}, MultiModalDummyOptions()
     )
-    uncapped_frames = uncapped_dummy["video"][0][0].shape[0]
+    uncapped_videos = uncapped_dummy["video"]
+    assert isinstance(uncapped_videos, list)
+    uncapped_frames = uncapped_videos[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 02181d5d6fb6..f259eb7d4f02 100644
--- a/tests/models/multimodal/processing/test_smolvlm.py
+++ b/tests/models/multimodal/processing/test_smolvlm.py
@@ -5,6 +5,7 @@
 import pytest
 from transformers import SmolVLMConfig
 
+from vllm.model_executor.models.smolvlm import SmolVLMProcessingInfo
 from vllm.multimodal import MULTIMODAL_REGISTRY
 
 from ....conftest import ImageTestAssets
@@ -40,6 +41,7 @@ def test_processor_override(
         limit_mm_per_prompt={"image": num_imgs},
     )
     processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config)
+    assert isinstance(processor.info, SmolVLMProcessingInfo)
     hf_processor_mm_kwargs = {} if kwargs_on_init else mm_processor_kwargs
 
     # Build the image str / prompt based on the number of images we pass
diff --git a/tests/models/multimodal/processing/test_tensor_schema.py b/tests/models/multimodal/processing/test_tensor_schema.py
index 73be3ab76c0b..573bdd4a8d8b 100644
--- a/tests/models/multimodal/processing/test_tensor_schema.py
+++ b/tests/models/multimodal/processing/test_tensor_schema.py
@@ -45,11 +45,10 @@
     list[Image.Image] | list[np.ndarray] | list[tuple[np.ndarray, dict[str, Any]]]
 )
 AudioInput = list[tuple[np.ndarray, int]]
+ResizedData: TypeAlias = Image.Image | np.ndarray | list[Image.Image]
 
 
-def _resize_data(
-    _data: Image.Image | np.ndarray, size_factor: float
-) -> Image.Image | np.ndarray:
+def _resize_data(_data: ResizedData, size_factor: float) -> ResizedData:
     assert size_factor <= 1, "Size factor must be less than 1"
     # Image input
     if isinstance(_data, Image.Image):
@@ -75,7 +74,7 @@ def _resize_data(
 
 def resize_mm_data(
     data: ImageInput | VideoInput | AudioInput, size_factors: tuple[float, ...]
-) -> ImageInput | VideoInput | AudioInput:
+) -> list[ResizedData]:
     size_factors = size_factors[: len(data)]
     if is_list_of(data, (Image.Image, np.ndarray, list)):
         return [_resize_data(d, s) for d, s in zip(data, size_factors)]
diff --git a/tests/models/multimodal/processing/test_transformers_audio.py b/tests/models/multimodal/processing/test_transformers_audio.py
index 1c1204d8e5ce..b0dc9ae0cd57 100644
--- a/tests/models/multimodal/processing/test_transformers_audio.py
+++ b/tests/models/multimodal/processing/test_transformers_audio.py
@@ -4,6 +4,9 @@
 import pytest
 
 from vllm.config import ModelConfig
+from vllm.model_executor.models.transformers.multimodal import (
+    MultiModalProcessor as TransformersMultiModalProcessor,
+)
 from vllm.multimodal import MULTIMODAL_REGISTRY
 
 AUDIO_MODEL_SETTINGS = {
@@ -103,6 +106,7 @@ def test_audio_multimodal_processor(model_id):
 
     audio_items = result.get("mm_kwargs", {}).get("audio", [])
     assert len(audio_items) == 1, f"Expected 1 audio item, got {len(audio_items)}"
+    assert audio_items[0] is not None
     item_keys = list(audio_items[0].keys())
     has_features = "input_features" in item_keys or "input_values" in item_keys
     assert has_features, (
@@ -142,6 +146,7 @@ def test_audio_fields_not_claimed_by_image():
     model_id = "ibm-granite/granite-speech-3.3-2b"
     model_config = ModelConfig(model=model_id, model_impl="transformers")
     mm_processor = MULTIMODAL_REGISTRY.create_processor(model_config)
+    assert isinstance(mm_processor, TransformersMultiModalProcessor)
 
     audio_keys = ["input_features", "input_features_mask"]
     owned = mm_processor._partition_keys_by_modality(audio_keys, ["audio", "image"])
@@ -155,6 +160,7 @@ def test_unclaimed_fields_warn_rather_than_raise():
     model_id = "ibm-granite/granite-speech-3.3-2b"
     model_config = ModelConfig(model=model_id, model_impl="transformers")
     mm_processor = MULTIMODAL_REGISTRY.create_processor(model_config)
+    assert isinstance(mm_processor, TransformersMultiModalProcessor)
 
     owned = mm_processor._partition_keys_by_modality(
         ["input_features", "surprise_field"], ["audio", "image"]
diff --git a/tests/models/multimodal/processing/test_transformers_image.py b/tests/models/multimodal/processing/test_transformers_image.py
index d45912a4b141..433a00942b80 100644
--- a/tests/models/multimodal/processing/test_transformers_image.py
+++ b/tests/models/multimodal/processing/test_transformers_image.py
@@ -3,6 +3,7 @@
 from unittest.mock import patch
 
 import pytest
+import torch
 
 from vllm.assets.image import ImageAsset
 from vllm.config import ModelConfig
@@ -113,8 +114,12 @@ def test_batch_padding_removed_from_image_items():
     items = result["mm_kwargs"]["image"]
     shapes = set()
     for item in items:
-        height, width = item["image_sizes"].data.flatten().tolist()
+        assert item is not None
+        image_sizes = item["image_sizes"].data
+        assert isinstance(image_sizes, torch.Tensor)
+        height, width = image_sizes.flatten().tolist()
         pixel_values = item["pixel_values"].data
+        assert isinstance(pixel_values, torch.Tensor)
         assert tuple(pixel_values.shape[-2:]) == (height, width)
         shapes.add(tuple(pixel_values.shape))
 
@@ -219,6 +224,7 @@ def test_repeated_image_hits_the_processor_cache():
     model_config = ModelConfig(
         model="llava-hf/llava-onevision-qwen2-0.5b-ov-hf", model_impl="transformers"
     )
+    assert model_config.multimodal_config is not None
     model_config.multimodal_config.mm_processor_cache_gb = 4
     mm_processor = MULTIMODAL_REGISTRY.create_processor(model_config)
     cache = MultiModalProcessorOnlyCache(model_config)
@@ -303,5 +309,9 @@ def test_nested_image_fields_split_per_image():
     items = result["mm_kwargs"]["image"]
     assert len(items) == 2
     for item in items:
+        assert item is not None
         pixel_values = item["pixel_values"].data
-        assert pixel_values.shape[1] == int(item["num_image_patches"].data)
+        num_image_patches = item["num_image_patches"].data
+        assert isinstance(pixel_values, torch.Tensor)
+        assert isinstance(num_image_patches, torch.Tensor)
+        assert pixel_values.shape[1] == int(num_image_patches)
diff --git a/tests/models/multimodal/processing/test_transformers_video.py b/tests/models/multimodal/processing/test_transformers_video.py
index 8d06843f06d3..00966495f6e0 100644
--- a/tests/models/multimodal/processing/test_transformers_video.py
+++ b/tests/models/multimodal/processing/test_transformers_video.py
@@ -2,9 +2,11 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import numpy as np
 import pytest
+import torch
 from transformers import Qwen3VLVideoProcessor
 
 from vllm.config import ModelConfig
+from vllm.model_executor.models.transformers.multimodal import MultiModalProcessingInfo
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.multimodal.cache import MultiModalProcessorOnlyCache
 
@@ -103,9 +105,14 @@ def test_video_multimodal_processor(model_id):
     )
 
     (item,) = result["mm_kwargs"]["video"]
-    rows = item["pixel_values_videos"].data.shape[0]
+    assert item is not None
+    pixel_values_videos = item["pixel_values_videos"].data
+    assert isinstance(pixel_values_videos, torch.Tensor)
+    rows = pixel_values_videos.shape[0]
     if "video_grid_thw" in item:
-        assert rows == int(item["num_video_patches"].data)
+        num_video_patches = item["num_video_patches"].data
+        assert isinstance(num_video_patches, torch.Tensor)
+        assert rows == int(num_video_patches)
     else:
         assert rows == len(frames)
 
@@ -168,6 +175,8 @@ def test_video_fields_not_claimed_by_image():
 
     (image_item,) = result["mm_kwargs"]["image"]
     (video_item,) = result["mm_kwargs"]["video"]
+    assert image_item is not None
+    assert video_item is not None
     assert {"pixel_values", "image_grid_thw"} <= set(image_item.keys())
     assert {"pixel_values_videos", "video_grid_thw"} <= set(video_item.keys())
     assert len(result["mm_placeholders"]["image"]) == 1
@@ -179,6 +188,7 @@ def test_repeated_video_hits_the_processor_cache():
     model_config = ModelConfig(
         model="Qwen/Qwen3-VL-2B-Instruct", model_impl="transformers"
     )
+    assert model_config.multimodal_config is not None
     model_config.multimodal_config.mm_processor_cache_gb = 4
     mm_processor = MULTIMODAL_REGISTRY.create_processor(model_config)
     cache = MultiModalProcessorOnlyCache(model_config)
@@ -222,7 +232,13 @@ def test_merged_token_fields_split_per_video():
     )
 
     items = result["mm_kwargs"]["video"]
-    assert [len(item["video_compression_mask"].data) for item in items] == [128, 64]
+    assert all(item is not None for item in items)
+    assert [
+        len(item["video_compression_mask"].data) for item in items if item is not None
+    ] == [
+        128,
+        64,
+    ]
 
 
 def test_video_sampled_by_processor():
@@ -262,6 +278,7 @@ def test_dummy_video_has_the_most_tokens(model_id):
         ModelConfig(model=model_id, model_impl="transformers")
     )
     info = mm_processor.info
+    assert isinstance(info, MultiModalProcessingInfo)
     seq_len = info.ctx.model_config.max_model_len
 
     result = mm_processor.get_dummy_mm_inputs({"video": 1})
diff --git a/tests/models/multimodal/processing/test_unlimited_ocr.py b/tests/models/multimodal/processing/test_unlimited_ocr.py
index 87d2f2763119..4f5c08504725 100644
--- a/tests/models/multimodal/processing/test_unlimited_ocr.py
+++ b/tests/models/multimodal/processing/test_unlimited_ocr.py
@@ -6,6 +6,7 @@
 from PIL import Image
 
 from vllm.config import ModelConfig
+from vllm.model_executor.models.unlimited_ocr import UnlimitedOCRProcessingInfo
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.multimodal.cache import MultiModalProcessorOnlyCache
 from vllm.multimodal.processing.context import TimingContext
@@ -24,6 +25,7 @@ def test_profiling_tokens_cover_all_crop_grids():
         model_config_kwargs={"max_model_len": 8192},
     )
     mm_processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config)
+    assert isinstance(mm_processor.info, UnlimitedOCRProcessingInfo)
     hf_processor = mm_processor.info.get_hf_processor()
     dummy_inputs = mm_processor.get_dummy_mm_inputs({"image": 1})
     budget = dummy_inputs["mm_placeholders"]["image"][0].get_num_embeds()
@@ -48,6 +50,7 @@ def test_image_cache_identity_across_image_counts(use_uuid: bool, use_embeds: bo
     )
     processor = MULTIMODAL_REGISTRY.create_processor(config)
     tokenizer = processor.info.get_tokenizer()
+    images: list[torch.Tensor] | list[Image.Image]
     if use_embeds:
         images = [torch.ones(8, config.get_hidden_size())] * 3
     else:
diff --git a/tests/models/multimodal/processing/test_whisper.py b/tests/models/multimodal/processing/test_whisper.py
index 1627b48914c3..e7271fa92e5e 100644
--- a/tests/models/multimodal/processing/test_whisper.py
+++ b/tests/models/multimodal/processing/test_whisper.py
@@ -4,7 +4,9 @@
 
 import numpy as np
 import pytest
+import torch
 
+from vllm.model_executor.models.whisper import WhisperProcessingInfo
 from vllm.multimodal import MULTIMODAL_REGISTRY
 
 from ...utils import build_model_context
@@ -21,6 +23,7 @@ def test_audio_features_fit_encoder_window(
     """
     ctx = build_model_context(model_id, limit_mm_per_prompt={"audio": 1})
     processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config)
+    assert isinstance(processor.info, WhisperProcessingInfo)
     feature_extractor = processor.info.get_feature_extractor()
 
     sampling_rate = feature_extractor.sampling_rate
@@ -34,4 +37,5 @@ def test_audio_features_fit_encoder_window(
     )
 
     input_features = processed_inputs["mm_kwargs"].get_data()["input_features"]
+    assert isinstance(input_features, torch.Tensor)
     assert input_features.shape[-1] == feature_extractor.nb_max_frames
diff --git a/tests/models/multimodal/test_cohere_asr.py b/tests/models/multimodal/test_cohere_asr.py
index 950542f90019..5ef3bc94e14e 100644
--- a/tests/models/multimodal/test_cohere_asr.py
+++ b/tests/models/multimodal/test_cohere_asr.py
@@ -5,12 +5,15 @@
 
 import pytest
 import torch
+from torch import nn
 
 from vllm.config import SpeechToTextConfig
 from vllm.model_executor.models.cohere_asr import (
     CohereAsrForConditionalGeneration,
+    CohereASRModel,
     RelPositionMultiHeadAttention,
 )
+from vllm.utils.torch_utils import set_default_torch_dtype
 
 
 class _ReferenceRelPositionMultiHeadAttention(RelPositionMultiHeadAttention):
@@ -88,3 +91,42 @@ def test_get_num_audio_tokens_streaming_estimate(audio_duration_s, expected_toke
     )
 
     assert got == expected_tokens
+
+
+class _CohereASRWeightLoadingModel(CohereASRModel):
+    def __init__(self) -> None:
+        nn.Module.__init__(self)
+        self.register_buffer("pos_bias_u", torch.zeros(2, 2))
+        self.register_buffer("pos_bias_v", torch.zeros(2, 2))
+        self.attention = RelPositionMultiHeadAttention(
+            n_head=2,
+            n_feat=4,
+            pos_bias_u=self.pos_bias_u,
+            pos_bias_v=self.pos_bias_v,
+        )
+
+
+@pytest.mark.cpu_test
+def test_load_weights_preserves_runtime_bias_dtype() -> None:
+    with set_default_torch_dtype(torch.float16):
+        model = _CohereASRWeightLoadingModel()
+
+    loaded_weight = torch.ones_like(model.pos_bias_u, dtype=torch.float32)
+    model.load_weights([("pos_bias_u", loaded_weight), ("pos_bias_v", loaded_weight)])
+
+    assert model.pos_bias_u.dtype == torch.float16
+    assert model.pos_bias_v.dtype == torch.float16
+    assert model.attention.pos_bias_u.dtype == torch.float16
+    assert model.attention.pos_bias_v.dtype == torch.float16
+    torch.testing.assert_close(model.pos_bias_u, loaded_weight.half())
+
+    output = model.attention(
+        query=torch.randn(1, 3, 4, dtype=torch.float16),
+        key=torch.randn(1, 3, 4, dtype=torch.float16),
+        value=torch.randn(1, 3, 4, dtype=torch.float16),
+        mask=None,
+        pos_emb=torch.randn(1, 5, 4, dtype=torch.float16),
+    )
+
+    assert output.dtype == torch.float16
+    assert output.shape == (1, 3, 4)
diff --git a/tests/models/multimodal/test_mapping.py b/tests/models/multimodal/test_mapping.py
index a213091e3d09..ec92ff9eb854 100644
--- a/tests/models/multimodal/test_mapping.py
+++ b/tests/models/multimodal/test_mapping.py
@@ -342,7 +342,7 @@ def test_hf_model_weights_mapper(model_arch: str):
     hf_dummy_model = create_dummy_model(model_id, model_arch)
     hf_converted_weights = hf_dummy_model.named_parameters()
     hf_converted_buffers = hf_dummy_model.named_buffers()
-    mapper: WeightsMapper = model_cls.hf_to_vllm_mapper
+    mapper: WeightsMapper = getattr(model_cls, "hf_to_vllm_mapper")  # noqa: B009
 
     mapped_original_weights = mapper.apply(original_weights)
     mapped_hf_converted_weights = mapper.apply(hf_converted_weights)
@@ -363,8 +363,7 @@ def test_hf_model_weights_mapper(model_arch: str):
         config = hf_dummy_model.config
         key = "tie_word_embeddings"
         if getattr(config.get_text_config(), key, False) or getattr(config, key, False):
-            mapped_tied_weights = mapper.apply((k, None) for k in tied)
-            tied_weight_names = set(map(lambda x: x[0], mapped_tied_weights))
+            tied_weight_names = set(mapper.apply_list(tied))
             ref_weight_names -= tied_weight_names
 
     weights_missing = ref_weight_names - weight_names
diff --git a/tests/models/multimodal/test_mimo_v2_omni.py b/tests/models/multimodal/test_mimo_v2_omni.py
index deccace50d4e..609bf90266cf 100644
--- a/tests/models/multimodal/test_mimo_v2_omni.py
+++ b/tests/models/multimodal/test_mimo_v2_omni.py
@@ -5,6 +5,11 @@
 import pytest
 import torch
 
+from vllm.config.multimodal import MultiModalConfig
+from vllm.model_executor.models import mimo_v2_omni
+from vllm.model_executor.models.interfaces import (
+    supports_multimodal_encoder_tp_data,
+)
 from vllm.platforms import current_platform
 
 EMBED_DIM = 256
@@ -58,10 +63,9 @@ def test_window_attention_applies_sinks(vision_attn_env, num_kv_heads):
     attn.sinks.data.normal_()
 
     total = sum(SEQ_LENS)
-    opts = dict(device="cuda", dtype=torch.bfloat16)
-    q = torch.randn(total, NUM_HEADS, HEAD_DIM, **opts)
-    k = torch.randn(total, num_kv_heads, HEAD_DIM, **opts)
-    v = torch.randn(total, num_kv_heads, HEAD_DIM, **opts)
+    q = torch.randn(total, NUM_HEADS, HEAD_DIM, device="cuda", dtype=torch.bfloat16)
+    k = torch.randn(total, num_kv_heads, HEAD_DIM, device="cuda", dtype=torch.bfloat16)
+    v = torch.randn(total, num_kv_heads, HEAD_DIM, device="cuda", dtype=torch.bfloat16)
     cu_seqlens = torch.tensor(
         [0, *torch.tensor(SEQ_LENS).cumsum(0).tolist()],
         device="cuda",
@@ -74,3 +78,99 @@ def test_window_attention_applies_sinks(vision_attn_env, num_kv_heads):
     # bf16 attention lands at ~2e-3 here; dropping the sinks lands at ~1e-1.
     error = ((out.float() - ref).norm() / ref.norm()).item()
     assert error < 1e-2, f"sink-corrected output is off by {error:.2e}"
+
+
+def _build_attention(num_kv_heads: int):
+    from vllm.model_executor.models.mimo_v2_omni import MiMoVisionAttention
+
+    torch.manual_seed(0)
+    attn = MiMoVisionAttention(
+        embed_dim=EMBED_DIM,
+        num_heads=NUM_HEADS,
+        num_kv_heads=num_kv_heads,
+        qk_channels=HEAD_DIM,
+        kv_channels=HEAD_DIM,
+        use_sink=True,
+        visual_token_window_size=WINDOW,
+    ).cuda()
+    attn.sinks.data.normal_()
+    return attn
+
+
+def _stub_encoder_parallelism(
+    monkeypatch, mode: str, world_size: int, rank: int
+) -> None:
+    """Set the encoder mode and the rank it sees.
+
+    The mode goes through a real `MultiModalConfig`, so `is_vit_use_data_parallel`
+    runs its own logic; only the rank is stubbed, because a single-process test
+    cannot be rank 1 of a multi-rank group.
+    """
+    monkeypatch.setattr(
+        "vllm.model_executor.models.vision.get_multimodal_config",
+        lambda: MultiModalConfig(mm_encoder_tp_mode=mode),
+    )
+    monkeypatch.setattr(
+        mimo_v2_omni.parallel_state,
+        "get_tensor_model_parallel_world_size",
+        lambda: world_size,
+    )
+    monkeypatch.setattr(
+        mimo_v2_omni.parallel_state, "get_tensor_model_parallel_rank", lambda: rank
+    )
+
+
+def _attention_inputs(num_heads: int, num_kv_heads: int):
+    total = sum(SEQ_LENS)
+    opts = dict(device="cuda", dtype=torch.bfloat16)
+    q = torch.randn(total, num_heads, HEAD_DIM, **opts)
+    k = torch.randn(total, num_kv_heads, HEAD_DIM, **opts)
+    v = torch.randn(total, num_kv_heads, HEAD_DIM, **opts)
+    cu_seqlens = torch.tensor(
+        [0, *torch.tensor(SEQ_LENS).cumsum(0).tolist()],
+        device="cuda",
+        dtype=torch.int32,
+    )
+    return q, k, v, cu_seqlens
+
+
+def _relative_error(attn, q, k, v, cu_seqlens, sinks):
+    out = attn._forward_window_attn(q, k, v, cu_seqlens, max(SEQ_LENS))
+    ref = _reference(q, k, v, cu_seqlens, sinks, attn.scale)
+    # bf16 attention lands at ~2e-3 here; dropping the sinks lands at ~1e-1.
+    return ((out.float() - ref).norm() / ref.norm()).item()
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires a GPU")
+@pytest.mark.parametrize("tp_rank", [1, 3])
+def test_window_attention_applies_sinks_under_encoder_data_parallel(
+    vision_attn_env, monkeypatch, tp_rank
+):
+    """Every rank runs the whole encoder, so every rank keeps all sinks."""
+    _stub_encoder_parallelism(monkeypatch, "data", world_size=4, rank=tp_rank)
+    attn = _build_attention(NUM_HEADS)
+    q, k, v, cu_seqlens = _attention_inputs(NUM_HEADS, NUM_HEADS)
+
+    error = _relative_error(attn, q, k, v, cu_seqlens, attn.sinks)
+    assert error < 1e-2, (
+        f"rank {tp_rank} of the data-parallel encoder is off by {error:.2e}"
+    )
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires a GPU")
+def test_weights_mode_slices_sinks_per_rank(vision_attn_env, monkeypatch):
+    """The collapse is data-parallel only: a sharded encoder keeps its slices."""
+    _stub_encoder_parallelism(monkeypatch, "weights", world_size=2, rank=1)
+    attn = _build_attention(NUM_HEADS)
+    heads = NUM_HEADS // 2
+    q, k, v, cu_seqlens = _attention_inputs(heads, heads)
+
+    error = _relative_error(attn, q, k, v, cu_seqlens, attn.sinks[heads:])
+    assert error < 1e-2, f"rank 1 of the sharded encoder is off by {error:.2e}"
+
+
+def test_model_advertises_encoder_tp_data():
+    """Without the flag the config rejects the mode and falls back to weights."""
+    from vllm.model_executor.models.mimo_v2_omni import MiMoV2OmniForCausalLM
+
+    assert supports_multimodal_encoder_tp_data(MiMoV2OmniForCausalLM)
diff --git a/tests/models/multimodal/test_nano_nemotron_vl.py b/tests/models/multimodal/test_nano_nemotron_vl.py
index 41f9428e6d73..ee6f36893f9e 100644
--- a/tests/models/multimodal/test_nano_nemotron_vl.py
+++ b/tests/models/multimodal/test_nano_nemotron_vl.py
@@ -1,13 +1,19 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from types import SimpleNamespace
 from unittest.mock import patch
 
 import numpy as np
 import pytest
 import torch
+import torch.nn as nn
+from PIL import Image
 
 from vllm import envs
+from vllm.config import MultiModalConfig
+from vllm.engine.arg_utils import EngineArgs
+from vllm.model_executor.layers.fusion.mm_input_norm import FusedMMInputNorm
 from vllm.model_executor.models.nano_nemotron_vl import (
     NanoNemotronVLMultiModalProcessor,
     NemotronH_Nano_VL_V2,
@@ -17,11 +23,192 @@
     MultiModalDataParser,
     VideoProcessorItems,
 )
+from vllm.platforms import current_platform
+from vllm.transformers_utils.processors.nano_nemotron_vl import (
+    DynamicResolutionImageTiler,
+    NanoNemotronVLProcessor,
+    _bicubic_resize_and_normalize,
+)
+
+# FusedMMInputNorm runs a Triton kernel on GPU platforms, so its inputs must
+# live on the platform's device.
+_DEVICE = torch.device(current_platform.device_type)
+
+
+@pytest.mark.parametrize("do_cpu_normalize", [True, False])
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+def test_nemotron_resized_normalization_matches_reference(
+    dtype: torch.dtype, do_cpu_normalize: bool, default_vllm_config
+):
+    pixels = torch.randint(
+        0,
+        256,
+        (1, 19, 23, 3),
+        dtype=torch.uint8,
+        generator=torch.Generator().manual_seed(0),
+    )
+    mean = torch.tensor([0.48145466, 0.4578275, 0.40821073]).view(1, 3, 1, 1)
+    std = torch.tensor([0.26862954, 0.26130258, 0.27577711]).view(1, 3, 1, 1)
+    output = _bicubic_resize_and_normalize(
+        pixels,
+        size=(16, 16),
+        norm_mean=mean,
+        norm_std=std,
+        do_cpu_normalize=do_cpu_normalize,
+        dtype=dtype,
+    )
+
+    if not do_cpu_normalize:
+        assert output.dtype == torch.uint8
+        model = _uninitialized_model()
+        nn.Module.__init__(model)
+        model.input_norm = FusedMMInputNorm(
+            mean.flatten().tolist(), std.flatten().tolist(), 1.0 / 255.0
+        ).to(_DEVICE)
+        model.llm_dtype = dtype
+        output = model._normalize_pixel_values(output.to(_DEVICE)).cpu()
+
+    resized = torch.nn.functional.interpolate(
+        pixels.permute(0, 3, 1, 2).float(),
+        size=(16, 16),
+        mode="bicubic",
+        align_corners=False,
+        antialias=True,
+    )
+    expected = ((resized.float() / 255.0 - mean) / std).to(dtype)
+    assert output.dtype == dtype
+    torch.testing.assert_close(output, expected, rtol=0.02, atol=0.06)
+    assert (output.float() - expected.float()).abs().mean() < 0.01
+
+
+def test_nemotron_dynamic_device_normalization_matches_cpu_reference(
+    default_vllm_config,
+):
+    image_mean = [0.485, 0.456, 0.406]
+    image_std = [0.229, 0.224, 0.225]
+    shape = (1, 7, 3 * 16 * 16)
+    pixels = torch.randint(0, 256, shape, dtype=torch.uint8)
+
+    model = _uninitialized_model()
+    nn.Module.__init__(model)
+    model.input_norm = FusedMMInputNorm(image_mean, image_std, 1.0 / 255.0).to(_DEVICE)
+    model.llm_dtype = torch.bfloat16
+
+    output = model._normalize_pixel_values(pixels.to(_DEVICE)).cpu()
+    batch, patches, _ = pixels.shape
+    expected = pixels.to(torch.float32).view(batch, patches, 3, -1)
+    expected = (
+        expected / 255.0 - torch.tensor(image_mean).view(1, 1, 3, 1)
+    ) / torch.tensor(image_std).view(1, 1, 3, 1)
+    expected = expected.view(shape)
+
+    torch.testing.assert_close(output, expected.to(torch.bfloat16))
+
+
+def test_nemotron_cpu_normalized_pixels_bypass_input_norm():
+    pixels = torch.randn(1, 3, 16, 16)
+    model = _uninitialized_model()
+    nn.Module.__init__(model)
+    model.input_norm = None
+
+    assert model._normalize_pixel_values(pixels) is pixels
+
+
+@pytest.mark.parametrize("modality", ["image", "video", "dynamic_image"])
+def test_nemotron_processor_defers_normalization_to_device(
+    modality: str, default_vllm_config
+):
+    dynamic_image = modality == "dynamic_image"
+    config = SimpleNamespace(
+        force_image_size=16,
+        patch_size=4,
+        downsample_ratio=0.5,
+        use_thumbnail=False,
+        norm_mean=[0.485, 0.456, 0.406],
+        norm_std=[0.229, 0.224, 0.225],
+        dtype=torch.bfloat16,
+        vision_config=SimpleNamespace(
+            args={"min_num_patches": 4, "max_num_patches": 16} if dynamic_image else {}
+        ),
+        sound_config=None,
+    )
+    tokenizer = SimpleNamespace(encode=lambda *args, **kwargs: [1])
+    common_kwargs = {
+        "config": config,
+        "tokenizer": tokenizer,
+        "max_model_len": 1024,
+    }
+    cpu_processor = NanoNemotronVLProcessor(**common_kwargs)
+    device_kwargs = MultiModalConfig(
+        mm_device_do_normalize=True
+    ).merge_mm_processor_kwargs({})
+    device_processor = NanoNemotronVLProcessor(**common_kwargs, **device_kwargs)
+
+    pixels = np.random.default_rng(0).integers(0, 256, (4, 19, 23, 3), dtype=np.uint8)
+    if modality == "image":
+        image = Image.fromarray(pixels[0])
+        cpu_values = cpu_processor._images_to_pixel_values_lst([image], 1)[0]
+        raw_values = device_processor._images_to_pixel_values_lst([image], 1)[0]
+    elif modality == "video":
+        cpu_values = cpu_processor._videos_to_pixel_values_lst(
+            [pixels], dtype=torch.bfloat16
+        )[0]
+        raw_values = device_processor._videos_to_pixel_values_lst([pixels])[0]
+    else:
+        params = DynamicResolutionImageTiler.DynamicResolutionParams(
+            media=Image.fromarray(pixels[0]),
+            num_tiles=1,
+            num_embeddings=4,
+            patch_size=(4, 4),
+        )
+        assert cpu_processor.dynamic_tiler is not None
+        assert device_processor.dynamic_tiler is not None
+        cpu_values = cpu_processor.dynamic_tiler.apply_params(
+            params, dtype=torch.bfloat16
+        )[0]
+        raw_values = device_processor.dynamic_tiler.apply_params(params)[0]
+
+    assert raw_values.dtype == torch.uint8
+    assert cpu_values.dtype == torch.bfloat16
+    assert raw_values.shape == cpu_values.shape
+
+    # Normalizing on the device hands the encoder the same dtype as the CPU
+    # path. Values differ only where the uint8 resize rounds or clips the
+    # bicubic overshoot: a few elements by several levels, tiny on average.
+    model = _uninitialized_model()
+    nn.Module.__init__(model)
+    model.input_norm = FusedMMInputNorm(
+        config.norm_mean, config.norm_std, 1.0 / 255.0
+    ).to(_DEVICE)
+    model.llm_dtype = config.dtype
+    if dynamic_image:
+        cpu_values = DynamicResolutionImageTiler.stack([cpu_values], config.patch_size)
+        raw_values = DynamicResolutionImageTiler.stack([raw_values], config.patch_size)
+    device_values = model._normalize_pixel_values(raw_values.to(_DEVICE)).cpu()
+    assert device_values.dtype == cpu_values.dtype
+    assert (device_values.float() - cpu_values.float()).abs().mean() < 0.01
+
+
+def test_nemotron_processor_rejects_partial_normalization():
+    with pytest.raises(
+        ValueError, match="requires do_rescale and do_normalize to have the same value"
+    ):
+        NanoNemotronVLProcessor(
+            config=None,
+            tokenizer=None,
+            max_model_len=1024,
+            do_rescale=False,
+            do_normalize=True,
+        )
+
+
+def _uninitialized_model() -> NemotronH_Nano_VL_V2:
+    return object.__new__(NemotronH_Nano_VL_V2)  # type: ignore[type-abstract]
 
 
 @pytest.mark.parametrize("input_key", ["image_embeds", "video_embeds"])
 def test_precomputed_multimodal_embeddings(input_key: str):
-    model = object.__new__(NemotronH_Nano_VL_V2)
+    model = _uninitialized_model()
     embeds = torch.randn(2, 4, 8)
 
     outputs = model.embed_multimodal(**{input_key: embeds})
@@ -77,21 +264,8 @@ def load_weights(self, weights):
         self.loaded_weights = list(weights)
 
 
-class _FakeTensor:
-    """Sentinel stand-in for torch.Tensor in load_weights tests. Supports the
-    .detach().clone() chain used by load_weights for buffered mm weights;
-    both methods return self so identity (and the existing equality
-    assertions) are preserved through cloning."""
-
-    def detach(self):
-        return self
-
-    def clone(self):
-        return self
-
-
 def test_nano_nemotron_vl_skips_multimodal_weights_in_text_only_mode():
-    model = object.__new__(NemotronH_Nano_VL_V2)
+    model = _uninitialized_model()
     language_model = _LanguageModel()
     object.__setattr__(model, "model_config", _ModelConfig())
     object.__setattr__(model, "language_model", language_model)
@@ -100,15 +274,15 @@ def test_nano_nemotron_vl_skips_multimodal_weights_in_text_only_mode():
     object.__setattr__(model, "vision_final_layernorm", _MissingMultiModalModule())
     object.__setattr__(model, "sound_encoder", None)
 
-    language_weight = object()
+    language_weight = torch.tensor(1.0)
     model.load_weights(
         [
             ("language_model.layers.0.weight", language_weight),
-            ("mlp1.0.weight", object()),
-            ("vision_model.radio_model.encoder.weight", object()),
-            ("sound_encoder.encoder.weight", object()),
-            ("vision_projector.vision_final_layernorm.weight", object()),
-            ("vision_projector.vision_final_layernorm.bias", object()),
+            ("mlp1.0.weight", torch.empty(0)),
+            ("vision_model.radio_model.encoder.weight", torch.empty(0)),
+            ("sound_encoder.encoder.weight", torch.empty(0)),
+            ("vision_projector.vision_final_layernorm.weight", torch.empty(0)),
+            ("vision_projector.vision_final_layernorm.bias", torch.empty(0)),
         ]
     )
 
@@ -116,7 +290,7 @@ def test_nano_nemotron_vl_skips_multimodal_weights_in_text_only_mode():
 
 
 def test_nano_nemotron_vl_loads_vision_weights_without_sound_encoder():
-    model = object.__new__(NemotronH_Nano_VL_V2)
+    model = _uninitialized_model()
     language_model = _LanguageModel()
     vision_model = _VisionModel()
     object.__setattr__(model, "model_config", _ImageOnlyModelConfig())
@@ -126,8 +300,8 @@ def test_nano_nemotron_vl_loads_vision_weights_without_sound_encoder():
     object.__setattr__(model, "vision_final_layernorm", None)
     object.__setattr__(model, "sound_encoder", None)
 
-    language_weight = object()
-    vision_weight = _FakeTensor()
+    language_weight = torch.tensor(1.0)
+    vision_weight = torch.tensor([2.0, 3.0])
     model.load_weights(
         [
             ("language_model.layers.0.weight", language_weight),
@@ -136,13 +310,16 @@ def test_nano_nemotron_vl_loads_vision_weights_without_sound_encoder():
     )
 
     assert language_model.loaded_weights == [("layers.0.weight", language_weight)]
-    assert vision_model.loaded_weights == [
-        ("radio_model.encoder.weight", vision_weight)
-    ]
+    assert len(vision_model.loaded_weights) == 1
+    name, loaded_weight = vision_model.loaded_weights[0]
+    assert name == "radio_model.encoder.weight"
+    assert isinstance(loaded_weight, torch.Tensor)
+    torch.testing.assert_close(loaded_weight, vision_weight)
+    assert loaded_weight is not vision_weight
 
 
 def test_nano_nemotron_vl_requires_sound_encoder_for_sound_weights():
-    model = object.__new__(NemotronH_Nano_VL_V2)
+    model = _uninitialized_model()
     language_model = _LanguageModel()
     vision_model = _VisionModel()
     object.__setattr__(model, "model_config", _ImageOnlyModelConfig())
@@ -153,16 +330,17 @@ def test_nano_nemotron_vl_requires_sound_encoder_for_sound_weights():
     object.__setattr__(model, "sound_encoder", None)
 
     with pytest.raises(AssertionError):
-        model.load_weights([("sound_encoder.encoder.weight", object())])
+        model.load_weights([("sound_encoder.encoder.weight", torch.empty(0))])
 
 
 def _vision_norm_model(use_layernorm=True):
-    model = object.__new__(NemotronH_Nano_VL_V2)
+    model = _uninitialized_model()
     torch.nn.Module.__init__(model)
     model.model_config = _ImageOnlyModelConfig()
     model.language_model = _LanguageModel()
     model.vision_model = _VisionModel()
     model.sound_encoder = None
+    model.input_norm = None
     model.mlp1 = torch.nn.Identity()
     model.vision_final_layernorm = (
         torch.nn.LayerNorm(4, eps=1e-6, dtype=torch.bfloat16) if use_layernorm else None
@@ -256,3 +434,25 @@ def test_extract_audio_from_videos_rejects_oversized_audio():
 
     assert audio_items == []
     assert has_audio == [False]
+
+
+@pytest.mark.core_model
+@pytest.mark.parametrize(
+    "model",
+    [
+        "nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16",
+        "nvidia/NVIDIA-Nemotron-3.5-Super-VL-120B-A12B-BF16",
+    ],
+)
+@pytest.mark.parametrize(
+    ("requested", "expected"), [("auto", "float32"), ("bfloat16", "bfloat16")]
+)
+def test_omni_mamba_ssm_cache_dtype(model: str, requested: str, expected: str):
+    """Omni archs get the NemotronH float32 SSM-cache default; explicit wins."""
+    vllm_config = EngineArgs(
+        model=model,
+        trust_remote_code=True,
+        mamba_ssm_cache_dtype=requested,
+        max_model_len=4096,
+    ).create_engine_config()
+    assert vllm_config.cache_config.mamba_ssm_cache_dtype == expected
diff --git a/tests/models/multimodal/test_openpangu_vl.py b/tests/models/multimodal/test_openpangu_vl.py
index e8bab16d5d96..7267e893bcc4 100644
--- a/tests/models/multimodal/test_openpangu_vl.py
+++ b/tests/models/multimodal/test_openpangu_vl.py
@@ -22,14 +22,13 @@ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
 
 
 class _OpenPanguVL(OpenPanguVLForConditionalGeneration):
+    language_model: _LanguageModel
+
     def __init__(self, hidden_size: int) -> None:
         nn.Module.__init__(self)
         self.config = SimpleNamespace(image_token_id=41, video_token_id=42)
         self.language_model = _LanguageModel(hidden_size)
 
-    def get_language_model(self) -> _LanguageModel:
-        return self.language_model
-
 
 @pytest.mark.parametrize("device_type", ["cpu", "cuda"])
 def test_get_input_embeddings_merges_image_and_video_embeddings(
@@ -40,7 +39,7 @@ def test_get_input_embeddings_merges_image_and_video_embeddings(
 
     device = torch.device(device_type)
     hidden_size = 4
-    model = _OpenPanguVL(hidden_size).to(device)
+    model = _OpenPanguVL(hidden_size).to(device)  # type: ignore[abstract]
     input_ids = torch.tensor([7, 41, 41, 8, 42, 9], device=device)
     multimodal_embeddings = (
         torch.full((2, hidden_size), -1.0, device=device),
diff --git a/tests/models/quantization/test_fp8_per_channel.py b/tests/models/quantization/test_fp8_per_channel.py
index ef11b17968c6..c63b01b5084b 100644
--- a/tests/models/quantization/test_fp8_per_channel.py
+++ b/tests/models/quantization/test_fp8_per_channel.py
@@ -19,8 +19,10 @@
 
 from ..utils import check_logprobs_close
 
-# Small MoE model that fits on a single GPU and exercises both linear + MoE.
-MOE_MODEL = "allenai/OLMoE-1B-7B-0125-Instruct"
+# Small MoE models that fit on a single GPU and exercise both linear + MoE,
+# one per modeling backend.
+MOE_MODEL = "ibm-granite/granite-3.0-1b-a400m-base"
+TRANSFORMERS_MOE_MODEL = "allenai/OLMoE-1B-7B-0125-Instruct"
 # Small dense model (no MoE) to validate the linear-only path.
 DENSE_MODEL = "Qwen/Qwen3-0.6B"
 
@@ -34,7 +36,11 @@
     reason="fp8 is not supported on this GPU type.",
 )
 @pytest.mark.quant_model
-@pytest.mark.parametrize("model", [DENSE_MODEL, MOE_MODEL], ids=["dense", "moe"])
+@pytest.mark.parametrize(
+    "model",
+    [DENSE_MODEL, MOE_MODEL, TRANSFORMERS_MOE_MODEL],
+    ids=["dense", "moe", "moe_transformers"],
+)
 def test_fp8_per_channel_logprobs(
     vllm_runner,
     quant_baseline_logprobs,
diff --git a/tests/models/quantization/test_mxfp8.py b/tests/models/quantization/test_mxfp8.py
index e6f709e2dd89..5ee0bb13991c 100644
--- a/tests/models/quantization/test_mxfp8.py
+++ b/tests/models/quantization/test_mxfp8.py
@@ -24,8 +24,10 @@
 
 from ..utils import check_logprobs_close
 
-# A small MoE model that fits on a single GPU and has both linear + MoE layers.
-MOE_MODEL = "allenai/OLMoE-1B-7B-0125-Instruct"
+# Small MoE models that fit on a single GPU and exercise both linear + MoE,
+# one per modeling backend.
+MOE_MODEL = "ibm-granite/granite-3.0-1b-a400m-base"
+TRANSFORMERS_MOE_MODEL = "allenai/OLMoE-1B-7B-0125-Instruct"
 # A small dense model (no MoE) to validate the linear-only path.
 DENSE_MODEL = "Qwen/Qwen3-0.6B"
 
@@ -39,7 +41,11 @@
     reason="mxfp8 is not supported on this GPU type (requires sm_100+).",
 )
 @pytest.mark.quant_model
-@pytest.mark.parametrize("model", [DENSE_MODEL, MOE_MODEL], ids=["dense", "moe"])
+@pytest.mark.parametrize(
+    "model",
+    [DENSE_MODEL, MOE_MODEL, TRANSFORMERS_MOE_MODEL],
+    ids=["dense", "moe", "moe_transformers"],
+)
 def test_mxfp8_logprobs(
     vllm_runner,
     quant_baseline_logprobs,
@@ -241,7 +247,11 @@ def test_mxfp8_aiter_accepts_swigluoai_params(
     reason="mxfp8 is not supported on this GPU type (requires sm_100+).",
 )
 @pytest.mark.quant_model
-@pytest.mark.parametrize("model", [DENSE_MODEL, MOE_MODEL], ids=["dense", "moe"])
+@pytest.mark.parametrize(
+    "model",
+    [DENSE_MODEL, MOE_MODEL, TRANSFORMERS_MOE_MODEL],
+    ids=["dense", "moe", "moe_transformers"],
+)
 def test_mxfp8_generation(vllm_runner, model: str) -> None:
     """Smoke test: verify online MXFP8 model generates coherent text."""
     prompt = "1 2 3 4 5"
diff --git a/tests/models/qwen3_5/__init__.py b/tests/models/qwen3_5/__init__.py
new file mode 100644
index 000000000000..208f01a7cb5e
--- /dev/null
+++ b/tests/models/qwen3_5/__init__.py
@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
diff --git a/tests/models/qwen3_5/test_gdn_ba_overlap.py b/tests/models/qwen3_5/test_gdn_ba_overlap.py
new file mode 100644
index 000000000000..5f66ded27273
--- /dev/null
+++ b/tests/models/qwen3_5/test_gdn_ba_overlap.py
@@ -0,0 +1,35 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+from transformers import Qwen3NextConfig
+
+from tests.utils import init_test_distributed_environment
+from tests.v1.attention.utils import create_vllm_config
+from vllm.config import set_current_vllm_config
+from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import (
+    QwenGatedDeltaNetAttention,
+)
+from vllm.platforms import current_platform
+from vllm.utils.network_utils import get_open_port
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA-only overlap")
+def test_ba_proj_overlaps_on_aux_stream_with_default_linear_backend() -> None:
+    """A GEMM dispatch change must not silently disable the ba overlap."""
+    vllm_config = create_vllm_config(dtype="bfloat16")
+    with set_current_vllm_config(vllm_config):
+        init_test_distributed_environment(1, 1, 0, str(get_open_port()))
+        layer = QwenGatedDeltaNetAttention(
+            Qwen3NextConfig(
+                hidden_size=64,
+                linear_num_key_heads=2,
+                linear_num_value_heads=4,
+                linear_key_head_dim=8,
+                linear_value_head_dim=8,
+            ),
+            vllm_config,
+            prefix="model.layers.0.linear_attn",
+            overlap_ba_proj=True,
+        )
+    assert layer._ba_stream is not None
diff --git a/tests/models/qwen3_5/test_gemma_add_rms_norm_quant.py b/tests/models/qwen3_5/test_gemma_add_rms_norm_quant.py
new file mode 100644
index 000000000000..17377ab73c55
--- /dev/null
+++ b/tests/models/qwen3_5/test_gemma_add_rms_norm_quant.py
@@ -0,0 +1,109 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Qwen3.5 Gemma add-RMSNorm + input quant vs the unfused production chain
+(flashinfer gemma norm, then scaled_fp8_quant / scaled_fp4_quant)."""
+
+import pytest
+import torch
+
+from vllm import _custom_ops as ops
+from vllm.platforms import current_platform
+
+pytestmark = pytest.mark.skipif(
+    not current_platform.has_device_capability(100),
+    reason="fused Gemma norm + FP8/NVFP4 quant targets SM100+",
+)
+
+# Partial and multiple 128-row scale tiles, both sides of the num_warps
+# switches (128 for NVFP4, 384 for FP8); 1 and 7 are decode sizes.
+NUM_TOKENS = [1, 7, 128, 129, 384, 385]
+
+
+def _fp8_codes(q: torch.Tensor) -> torch.Tensor:
+    """Sign-magnitude e4m3 bytes -> ordered ints (+0 == -0)."""
+    b = q.view(torch.uint8).int()
+    return torch.where(b >= 128, -(b & 0x7F), b)
+
+
+def _fp4_codes(packed: torch.Tensor) -> torch.Tensor:
+    """Packed e2m1 nibbles (low first) -> ordered ints (+0 == -0)."""
+    c = torch.stack([packed & 0xF, packed >> 4], dim=-1).flatten(-2).int()
+    return torch.where(c >= 8, -(c & 7), c)
+
+
+def _swizzled_sf_rows(sf: torch.Tensor, m: int, h: int) -> torch.Tensor:
+    """(m, h // 16) view of a 128x4-swizzled NVFP4 scale tensor."""
+    r = torch.arange(m, device=sf.device)[:, None]
+    k = torch.arange(h // 16, device=sf.device)[None, :]
+    num_k_tiles = (h // 16 + 3) // 4
+    off = (
+        ((r >> 7) * num_k_tiles + (k >> 2)) * 512
+        + (r & 31) * 16
+        + ((r >> 5) & 3) * 4
+        + (k & 3)
+    )
+    return sf.view(torch.uint8).flatten()[off]
+
+
+# 5120: a production hidden size; 2112: not a power of 2 (33 64-wide scale tiles).
+@pytest.mark.parametrize("H", [5120, 2112])
+@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
+@pytest.mark.parametrize("nvfp4", [False, True])
+@pytest.mark.parametrize("has_residual", [True, False])
+@torch.inference_mode()
+def test_gemma_add_rms_norm_quant_matches_chain(
+    H: int, num_tokens: int, nvfp4: bool, has_residual: bool
+):
+    from flashinfer.norm import gemma_fused_add_rmsnorm, gemma_rmsnorm
+
+    from vllm.models.qwen3_5.nvidia.ops.fused_norm_quant import (
+        gemma_add_rms_norm_quant,
+    )
+
+    g = torch.Generator(device="cuda").manual_seed(num_tokens)
+    # Checkpoint-like Gemma weights: (1 + w) from ~0 (w near -1) to 3.5.
+    w = (torch.rand(H, generator=g, device="cuda") * 3.5 - 1.0).bfloat16()
+    w[:64] = -0.99609375
+    x = torch.randn(num_tokens, H, generator=g, device="cuda").bfloat16()
+    residual = (torch.randn(num_tokens, H, generator=g, device="cuda") * 4).bfloat16()
+    residual[:, 7] *= 50  # outlier channel
+    # FP8: static input scale. NVFP4: global scale (input_global_scale_inv).
+    scale = torch.tensor([1 / 0.0013 if nvfp4 else 0.11], device="cuda")
+
+    ref_res = residual.clone()
+    if has_residual:
+        ref_y = x.clone()
+        gemma_fused_add_rmsnorm(ref_y, ref_res, w, 1e-6)
+    else:
+        ref_y = gemma_rmsnorm(x, w, 1e-6)
+    if nvfp4:
+        ref_q, ref_sf = ops.scaled_fp4_quant(
+            ref_y, scale, True, backend="flashinfer-cutedsl"
+        )
+    else:
+        ref_q, _ = ops.scaled_fp8_quant(ref_y, scale, group_shape=(-1, -1))
+
+    x_in = x.clone()
+    res = residual.clone()
+    q, sf, y = gemma_add_rms_norm_quant(
+        x_in, res if has_residual else None, w, 1e-6, scale, nvfp4, write_y=True
+    )
+
+    torch.testing.assert_close(x_in, x, rtol=0, atol=0)
+    torch.testing.assert_close(res, ref_res, rtol=0, atol=0)
+    # The fp32 sum of squares is reduced in a different order; rare 1-ulp
+    # differences of the bf16 output are expected.
+    assert (y != ref_y).float().mean().item() <= 1e-4
+    torch.testing.assert_close(y, ref_y, rtol=2**-7, atol=0)
+
+    if nvfp4:
+        dist = (_fp4_codes(q) - _fp4_codes(ref_q.view(torch.uint8))).abs()
+        sf_mismatch = _swizzled_sf_rows(sf, num_tokens, H) != _swizzled_sf_rows(
+            ref_sf, num_tokens, H
+        )
+        assert sf_mismatch.float().mean().item() <= 1e-4
+    else:
+        assert sf is scale
+        dist = (_fp8_codes(q) - _fp8_codes(ref_q)).abs()
+    assert dist.max().item() <= 1
+    assert (dist > 0).float().mean().item() <= 1e-4
diff --git a/tests/models/qwen4_exp/test_config.py b/tests/models/qwen4_exp/test_config.py
index 6d94b1e5607f..bb01acdd8431 100644
--- a/tests/models/qwen4_exp/test_config.py
+++ b/tests/models/qwen4_exp/test_config.py
@@ -197,6 +197,7 @@ def test_qwen4_exp_model_state_prepares_ngram_context() -> None:
     model_state.ngram_context = torch.empty((8, 3), dtype=torch.int32)
     model_state.ngram_context_offsets = torch.arange(-3, 0, dtype=torch.int64)
     model_state.ple_query_start_loc = torch.empty(9, dtype=torch.int32)
+    model_state.disk_offload_modules = []
 
     input_batch = SimpleNamespace(
         num_reqs=2,
diff --git a/tests/models/qwen4_exp/test_ple.py b/tests/models/qwen4_exp/test_ple.py
index b4382ebd15b2..292ffc8c8ed5 100644
--- a/tests/models/qwen4_exp/test_ple.py
+++ b/tests/models/qwen4_exp/test_ple.py
@@ -2,13 +2,17 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import logging
+import os
 from dataclasses import dataclass
 from itertools import accumulate
+from pathlib import Path
 from types import SimpleNamespace
 
 import pytest
 import torch
 from compressed_tensors.quantization import QuantizationArgs
+from safetensors import safe_open
+from safetensors.torch import save_file
 from torch import nn
 from torch.nn import functional as F
 from transformers import Qwen4ExpTextConfig
@@ -16,6 +20,8 @@
 import vllm.model_executor.layers.vocab_parallel_embedding as embedding_module
 import vllm.model_executor.parameter as parameter_module
 import vllm.models.qwen4_exp.common.ngram_embedding as ngram_embedding_module
+import vllm.models.qwen4_exp.nvidia.ngram_embedding as nvidia_ngram_embedding_module
+import vllm.utils.gpu_sync_debug as gpu_sync_debug
 from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata
 from vllm.config import SpeculativeConfig, VllmConfig
 from vllm.config.compilation import CUDAGraphMode
@@ -48,14 +54,17 @@
 from vllm.models.qwen4_exp.nvidia.ngram_embedding import (
     Qwen4ExpNGramEmbedding,
     Qwen4ExpPLEDeviceEmbedding,
+    Qwen4ExpPLEDiskOffloadEmbedding,
     Qwen4ExpPLEEmbeddingMethod,
     Qwen4ExpPLEFp8EmbeddingMethod,
     Qwen4ExpPLENvFp4EmbeddingMethod,
     Qwen4ExpPLEPinnedHostEmbedding,
     Qwen4ExpPLEUnquantizedEmbeddingMethod,
+    stage_checkpoint_rows,
 )
 from vllm.models.qwen4_exp.nvidia.ple_layer import Qwen4ExpPLELayer
 from vllm.platforms import current_platform
+from vllm.utils.gpu_sync_debug import with_gpu_sync_check
 from vllm.utils.torch_utils import weak_ref_tensor
 from vllm.v1.attention.backends.short_conv_attn import (
     PleShortConvAttentionMetadata,
@@ -139,6 +148,9 @@ def check(actual, expected):
     assert metadata.num_decodes == sum(query_lens[i] == 1 for i in non_spec_rows)
     assert metadata.num_prefills == sum(query_lens[i] > 1 for i in non_spec_rows)
     if spec_rows:
+        assert metadata.spec_state_indices_tensor is not None
+        assert metadata.num_accepted_tokens is not None
+        assert metadata.spec_query_start_loc is not None
         torch.testing.assert_close(
             metadata.spec_state_indices_tensor[: len(spec_rows)],
             common.block_table_tensor[spec_rows, 0],
@@ -1111,6 +1123,187 @@ def test_ple_fp8_embedding_supports_mixed_precision_config(
     )
 
 
+def _save_ple_checkpoint(path: Path, fmt: str) -> None:
+    """Save a 24-row, 32-wide PLE table as four checkpoint shards of 6 rows."""
+    tensors: dict[str, torch.Tensor] = {}
+    if fmt == "nvfp4":
+        planes = {
+            "weight": torch.arange(24 * 16).reshape(24, 16).to(torch.uint8),
+            "weight_scale": (torch.arange(48).reshape(24, 2) / 8 + 0.25).to(
+                torch.float8_e4m3fn
+            ),
+        }
+        tensors["ngram_embedding.weight_scale_2"] = torch.tensor(0.37)
+    else:
+        values = torch.linspace(-2, 2, 24 * 32).reshape(24, 32)
+        planes = {"weight": values.to(torch.bfloat16)}
+        if fmt == "fp8":
+            planes = {"weight": values.to(torch.float8_e4m3fn)}
+            tensors["ngram_embedding.weight_scale"] = torch.tensor([0.25])
+    for name, table in planes.items():
+        for index, shard in enumerate(table.split(6)):
+            tensors[f"ngram_embedding.shard_{index}.{name}"] = shard.clone()
+    save_file(tensors, path)
+
+
+def _make_ple_from_checkpoint(
+    monkeypatch: pytest.MonkeyPatch,
+    path: Path,
+    fmt: str,
+    *,
+    disk_offload: bool,
+    device: str = "cpu",
+    load: bool = True,
+) -> Qwen4ExpNGramEmbedding:
+    """Build a two-head PLE embedding and load `path` through file-backed views."""
+    _mock_etp_group(monkeypatch)
+    for module in (embedding_module, parameter_module):
+        monkeypatch.setattr(module, "get_tensor_model_parallel_rank", lambda: 0)
+        monkeypatch.setattr(module, "get_tensor_model_parallel_world_size", lambda: 1)
+    group = SimpleNamespace(world_size=1)
+    vllm_config = SimpleNamespace(
+        load_config=SimpleNamespace(load_format="auto"),
+        engram_config=SimpleNamespace(disk_offload=disk_offload, cpu_offload=False),
+    )
+    for name, value in (
+        ("get_current_vllm_config", vllm_config),
+        ("get_etp_group", group),
+        ("get_tp_group", group),
+    ):
+        monkeypatch.setattr(nvidia_ngram_embedding_module, name, lambda v=value: v)
+    config = SimpleNamespace(
+        ngram_size=2,
+        heads_per_ngram=2,
+        eos_token_id=0,
+        vocab_size=64,
+        split_ngram_parts=4,
+        ngram_vocab_size_base=8,
+        make_ngram_vocab_size_divisible_by=1,
+        ple_embedding_dtype={"nvfp4": "nvfp4", "fp8": "float8_e4m3fn"}.get(fmt),
+    )
+    with torch.device(device):
+        module = Qwen4ExpNGramEmbedding(
+            config,
+            64,
+            0,
+            8,
+            data_parallel_rank=0,
+            prefix="t",
+            params_dtype=torch.bfloat16,
+        )
+    if load:
+        with safe_open(path, framework="pt") as checkpoint:
+            module.load_weights(
+                (name, checkpoint.get_tensor(name))
+                for name in checkpoint.keys()  # noqa: SIM118
+            )
+        embedding = module.ngram_embedding
+        embedding.quant_method.process_weights_after_loading(embedding)
+    return module
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="disk offload is CUDA-only")
+@pytest.mark.parametrize("fmt", ["bf16", "fp8", "nvfp4"])
+def test_disk_offload_matches_resident_table_under_graph_replay(
+    monkeypatch, tmp_path, fmt
+) -> None:
+    """Staged rows equal the resident lookup of the same ids in a replayed graph."""
+    path = tmp_path / "model.safetensors"
+    _save_ple_checkpoint(path, fmt)
+    gather, resident = (
+        _make_ple_from_checkpoint(
+            monkeypatch, path, fmt, disk_offload=disk_offload, device="cuda"
+        )
+        for disk_offload in (True, False)
+    )
+    assert isinstance(gather.ngram_embedding, Qwen4ExpPLEDiskOffloadEmbedding)
+    assert gather.ngram_embedding.weight.numel() == 0
+    hidden_states = torch.zeros(8, 64, device="cuda")
+    query_start_loc = torch.tensor([0, 5, 8], dtype=torch.int32, device="cuda")
+    ngram_context = torch.tensor([[3], [0]], dtype=torch.int32, device="cuda")
+    input_ids = torch.randint(1, 64, (8,), device="cuda")
+    monkeypatch.setattr(gpu_sync_debug, "_SYNC_CHECK_MODE", "error")
+    monkeypatch.setattr(gpu_sync_debug, "_sync_check_enabled", True)
+
+    def stage() -> None:
+        with_gpu_sync_check(stage_checkpoint_rows)(
+            [(gather, gather.ngram_embedding)],
+            input_ids,
+            query_start_loc,
+            ngram_context,
+        )
+
+    def forward(module: Qwen4ExpNGramEmbedding) -> torch.Tensor:
+        return module(hidden_states, input_ids, query_start_loc, ngram_context)
+
+    stage()
+    forward(gather)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        output = forward(gather)
+    for seed in range(2):
+        torch.manual_seed(seed)
+        input_ids.copy_(torch.randint(1, 64, (8,)))
+        stage()
+        graph.replay()
+        expected = forward(resident)
+        assert torch.equal(output.view(torch.uint8), expected.view(torch.uint8))
+
+
+@pytest.mark.parametrize(
+    ("corrupt", "error"),
+    [
+        pytest.param(
+            lambda tensors: {k: v.float() for k, v in tensors.items()},
+            "must match the embedding dtype",
+            id="dtype",
+        ),
+        pytest.param(
+            lambda tensors: {k: v for k, v in tensors.items() if "_2." not in k},
+            "cover 18 of 24 rows",
+            id="missing_shard",
+        ),
+        pytest.param(
+            lambda tensors: {k: v.clone() for k, v in tensors.items()},
+            "not a file-backed view",
+            id="copied",
+        ),
+    ],
+)
+def test_disk_offload_rejects_shards_it_cannot_serve_from_files(
+    monkeypatch, tmp_path, corrupt, error
+) -> None:
+    path = tmp_path / "model.safetensors"
+    _save_ple_checkpoint(path, "bf16")
+    module = _make_ple_from_checkpoint(
+        monkeypatch, path, "bf16", disk_offload=True, load=False
+    )
+    embedding = module.ngram_embedding
+    with safe_open(path, framework="pt") as checkpoint:
+        tensors = {name: checkpoint.get_tensor(name) for name in checkpoint.keys()}  # noqa: SIM118
+        with pytest.raises(ValueError, match=error):
+            module.load_weights(corrupt(tensors).items())
+            embedding.quant_method.process_weights_after_loading(embedding)
+
+
+def test_disk_offload_fails_closed_on_truncation_and_reload(
+    monkeypatch, tmp_path
+) -> None:
+    path = tmp_path / "model.safetensors"
+    _save_ple_checkpoint(path, "bf16")
+    module = _make_ple_from_checkpoint(monkeypatch, path, "bf16", disk_offload=True)
+    embedding = module.ngram_embedding
+    assert isinstance(embedding, Qwen4ExpPLEDiskOffloadEmbedding)
+    embedding.copy_ids_to_host(torch.tensor([[0, 23], [5, 6], [12, 17], [23, 23]]))
+    embedding.stage_rows(4)
+    os.truncate(path, path.stat().st_size - 4)
+
+    with pytest.raises(ValueError, match="truncated at offset"):
+        embedding.stage_rows(4)
+    with pytest.raises(RuntimeError, match="does not support weight reload"):
+        module.load_weights([("ngram_embedding.shard_0.weight", torch.zeros(6, 32))])
+
+
 def test_dilated_ple_spec_state_rolls_back_before_next_forward() -> None:
     conv_state_len = 6
     dilation = 2
diff --git a/tests/models/qwen4_exp/test_qsa_amd.py b/tests/models/qwen4_exp/test_qsa_amd.py
index c8ebf0190432..178b3f60bae7 100644
--- a/tests/models/qwen4_exp/test_qsa_amd.py
+++ b/tests/models/qwen4_exp/test_qsa_amd.py
@@ -10,21 +10,21 @@
 import torch
 
 from vllm.model_executor.layers.layernorm import GemmaRMSNorm
-from vllm.models.qwen4_exp.amd import indexer_qsa
 from vllm.models.qwen4_exp.amd import (
     model as _qwen4_exp_model,  # noqa: F401
 )
 from vllm.models.qwen4_exp.amd import ple_layer as ple_layer_module
-from vllm.models.qwen4_exp.amd.indexer_qsa import (
-    QSAIndexer,
-    apply_qsa_rmsnorm,
-    apply_qsa_rope,
-)
+from vllm.models.qwen4_exp.amd import qsa_indexer
 from vllm.models.qwen4_exp.amd.ops import qsa as qsa_ops
 from vllm.models.qwen4_exp.amd.ops.qsa_prepare import (
     qsa_prepare,
     supports_fused_pre_indexer,
 )
+from vllm.models.qwen4_exp.amd.qsa_indexer import (
+    QSAIndexer,
+    apply_qsa_rmsnorm,
+    apply_qsa_rope,
+)
 from vllm.platforms import current_platform
 from vllm.triton_utils import HAS_TRITON
 
@@ -132,6 +132,7 @@ def _make_attn(*, use_fused: bool) -> SimpleNamespace:
     return SimpleNamespace(
         use_fused_qsa_prepare=use_fused,
         kv_cache=torch.zeros(1, 2, 16, 8),
+        kv_cache_dtype="auto",
         q_norm=GemmaRMSNorm(4, eps=1e-6),
         k_norm=GemmaRMSNorm(4, eps=1e-6),
         _k_scale_float=1.0,
@@ -139,7 +140,9 @@ def _make_attn(*, use_fused: bool) -> SimpleNamespace:
     )
 
 
-def _make_indexer(*, use_fused: bool, num_tokens: int = 2) -> QSAIndexer:
+def _make_indexer(
+    monkeypatch: pytest.MonkeyPatch, *, use_fused: bool, num_tokens: int = 2
+) -> QSAIndexer:
     """A QSAIndexer whose unfused path runs for real, with no projection weights.
 
     Both branches have to be reachable: a fused-path test is only meaningful if
@@ -152,6 +155,7 @@ def _make_indexer(*, use_fused: bool, num_tokens: int = 2) -> QSAIndexer:
     indexer.index_kv_heads = 1
     indexer.index_head_dim = head_dim
     indexer.compress_ratio = 2
+    indexer.indexer_dtype = torch.bfloat16
     indexer.use_fused_pre_indexer = use_fused
     indexer.skip_topk = True
     indexer.index_qk_proj = lambda hidden_states: (
@@ -190,7 +194,9 @@ def _make_indexer(*, use_fused: bool, num_tokens: int = 2) -> QSAIndexer:
         slot_mapping=torch.arange(num_tokens),
         k_work_metadata=torch.zeros(1, 2, dtype=torch.int32),
     )
-    indexer._metadata = lambda: (raw_metadata, compressed_metadata)
+    monkeypatch.setattr(
+        indexer, "_metadata", lambda: (raw_metadata, compressed_metadata)
+    )
     indexer.test_metadata = (raw_metadata, compressed_metadata)
     return indexer
 
@@ -199,7 +205,7 @@ def _make_indexer(*, use_fused: bool, num_tokens: int = 2) -> QSAIndexer:
 def test_qsa_fused_prepare_replaces_the_unfused_chain(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
-    indexer = _make_indexer(use_fused=True)
+    indexer = _make_indexer(monkeypatch, use_fused=True)
     attn = _make_attn(use_fused=True)
     calls = []
     main_outputs = (torch.empty(0), torch.empty(0))
@@ -208,7 +214,7 @@ def fake_prepare(*args, **kwargs):
         calls.append((args, kwargs))
         return main_outputs
 
-    monkeypatch.setattr(indexer_qsa, "qsa_prepare", fake_prepare)
+    monkeypatch.setattr(qsa_indexer, "qsa_prepare", fake_prepare)
     monkeypatch.setattr(
         QSAIndexer,
         "_update_and_compress",
@@ -260,9 +266,9 @@ def fake_prepare(*args, **kwargs):
 def test_qsa_unsupported_config_keeps_the_unfused_chain(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
-    indexer = _make_indexer(use_fused=False)
+    indexer = _make_indexer(monkeypatch, use_fused=False)
     monkeypatch.setattr(
-        indexer_qsa,
+        qsa_indexer,
         "qsa_prepare",
         lambda *args, **kwargs: pytest.fail("unfused path ran the fused kernel"),
     )
@@ -393,8 +399,9 @@ def expand_qsa_block_indices(*args) -> None:
 
 
 @requires_qsa_kernels
+@pytest.mark.parametrize("use_fp8", [False, True])
 @pytest.mark.parametrize(
-    ("num_rows", "num_query_heads", "num_kv_heads", "page_size"),
+    "num_rows,num_query_heads,num_kv_heads,page_size",
     [
         pytest.param(1, 24, 2, 1792, id="tp1_split64"),
         pytest.param(16, 12, 1, 1792, id="tp2_split32"),
@@ -408,6 +415,7 @@ def test_qsa_sparse_paged_attention_matches_reference(
     num_query_heads: int,
     num_kv_heads: int,
     page_size: int,
+    use_fp8: bool,
 ) -> None:
     torch.manual_seed(2)
     head_dim = 256
@@ -429,7 +437,18 @@ def test_qsa_sparse_paged_attention_matches_reference(
         device="cuda",
         dtype=torch.bfloat16,
     )
-    k_cache, v_cache = kv_cache.split(head_dim, dim=-1)
+    if use_fp8:
+        fp8_dtype = current_platform.fp8_dtype()
+        # The kernel's contract is that the cache already holds values in
+        # the attention dtype's numeric space; quantize the stored cache and
+        # have the reference consume the stored (quantized) values upcast
+        # to float, not the original bf16 — the fp8 cast is lossy.
+        kv_stored = kv_cache.to(fp8_dtype)
+        k_cache, v_cache = kv_stored.split(head_dim, dim=-1)
+        ref_k, ref_v = k_cache.float(), v_cache.float()
+    else:
+        k_cache, v_cache = kv_cache.split(head_dim, dim=-1)
+        ref_k, ref_v = k_cache, v_cache
     block_table = (
         torch.randperm(num_cache_blocks, device="cuda")
         .reshape(num_requests, num_pages_per_request)
@@ -482,8 +501,8 @@ def test_qsa_sparse_paged_attention_matches_reference(
     )
     expected = _qsa_sparse_paged_attention_reference(
         q,
-        k_cache,
-        v_cache,
+        ref_k,
+        ref_v,
         logical_indices,
         block_table,
         token_to_req,
@@ -491,3 +510,19 @@ def test_qsa_sparse_paged_attention_matches_reference(
     )
 
     torch.testing.assert_close(actual, expected, rtol=2e-2, atol=2e-2)
+
+
+@requires_qsa_kernels
+@pytest.mark.parametrize("kv_cache_dtype", ["fp8", "fp8_e4m3"])
+def test_qsa_impl_constructs_with_fp8_kv(kv_cache_dtype: str) -> None:
+    """QSA accepts FP8 KV cache dtypes without probing FlashAttention."""
+    from vllm.models.qwen4_exp.amd.qsa import Qwen4ExpQSAImpl
+
+    impl = Qwen4ExpQSAImpl(
+        num_heads=24,
+        head_size=128,
+        scale=0.088,
+        num_kv_heads=2,
+        kv_cache_dtype=kv_cache_dtype,
+    )
+    assert impl.kv_cache_dtype == kv_cache_dtype
diff --git a/tests/models/qwen4_exp/test_qsa_prepare.py b/tests/models/qwen4_exp/test_qsa_prepare.py
index e56a50b6c991..4f3424f19dd4 100644
--- a/tests/models/qwen4_exp/test_qsa_prepare.py
+++ b/tests/models/qwen4_exp/test_qsa_prepare.py
@@ -5,17 +5,17 @@
 import pytest
 import torch
 
-from vllm.models.qwen4_exp.common.qsa_cache import (
+from vllm.models.qwen4_exp.common.qsa_indexer import (
     canonical_qsa_rope_positions,
     circular_qsa_slot_mapping,
     compressed_qsa_slot_mapping,
 )
-from vllm.models.qwen4_exp.nvidia.indexer_qsa import apply_qsa_rope
 from vllm.models.qwen4_exp.nvidia.ops.qsa import (
     qsa_compress_groups_with_ratio,
     qsa_store_cache_rows,
 )
 from vllm.models.qwen4_exp.nvidia.ops.qsa_prepare import qsa_prepare
+from vllm.models.qwen4_exp.nvidia.qsa_indexer import apply_qsa_rope
 from vllm.platforms import current_platform
 from vllm.triton_utils import HAS_TRITON
 
diff --git a/tests/models/qwen4_exp/test_qsa_prepare_amd.py b/tests/models/qwen4_exp/test_qsa_prepare_amd.py
index 718929720955..f9ed23b39052 100644
--- a/tests/models/qwen4_exp/test_qsa_prepare_amd.py
+++ b/tests/models/qwen4_exp/test_qsa_prepare_amd.py
@@ -6,13 +6,13 @@
 import torch
 
 from vllm.model_executor.layers.layernorm import GemmaRMSNorm
-from vllm.models.qwen4_exp.amd.indexer_qsa import apply_qsa_rope
 from vllm.models.qwen4_exp.amd.ops.qsa import (
     qsa_compress_groups_with_ratio,
     qsa_store_cache_rows,
 )
 from vllm.models.qwen4_exp.amd.ops.qsa_prepare import qsa_prepare
-from vllm.models.qwen4_exp.common.qsa_cache import (
+from vllm.models.qwen4_exp.amd.qsa_indexer import apply_qsa_rope
+from vllm.models.qwen4_exp.common.qsa_indexer import (
     canonical_qsa_rope_positions,
     circular_qsa_slot_mapping,
     compressed_qsa_slot_mapping,
@@ -77,15 +77,17 @@ def assert_fp8_within_one_ulp(actual: torch.Tensor, expected: torch.Tensor) -> N
     assert bool(((code_diff <= 1) | (abs_diff <= 2**-8)).all())
 
 
-def _make_main_inputs(num_tokens: int) -> dict:
-    """Random main-attention arguments with a paged BF16 cache in vLLM's layout."""
+def _make_main_inputs(num_tokens: int, main_fp8: bool = False) -> dict:
+    """Random main-attention arguments with a paged cache in vLLM's layout."""
     num_blocks = num_tokens // MAIN_PAGE + 2
+    # FP8 configs store e4m3 codes in the platform's native flavor
+    # (FNUZ on MI300/MI325, OCP elsewhere), matching the production views.
     cache = torch.zeros(
         num_blocks,
         MAIN_HK,
         MAIN_PAGE,
         2 * MAIN_D,
-        dtype=torch.bfloat16,
+        dtype=current_platform.fp8_dtype() if main_fp8 else torch.bfloat16,
         device="cuda",
     ).transpose(1, 2)
     slots = torch.randperm(num_blocks * MAIN_PAGE, device="cuda")[:num_tokens]
@@ -103,8 +105,8 @@ def _make_main_inputs(num_tokens: int) -> dict:
         main_eps=EPS,
         main_kv_cache=cache,
         main_slot_mapping=slots,
-        main_k_scale=1.0,
-        main_v_scale=1.0,
+        main_k_scale=0.5 if main_fp8 else 1.0,
+        main_v_scale=2.0 if main_fp8 else 1.0,
         # The AMD op hands the gate back through a buffer it does not own.
         main_gate_out=torch.empty(
             num_tokens, MAIN_HQ, MAIN_D, dtype=torch.bfloat16, device="cuda"
@@ -129,14 +131,86 @@ def _check_main_outputs(main: dict, q_out, gate_out, rope, positions) -> None:
     ).view(num_tokens, -1)
     q, k = rope(positions, q, k)
     kv_cache = main["main_kv_cache"]
-    cache = torch.zeros_like(kv_cache)
-    key_cache, value_cache = cache.split(MAIN_D, dim=-1)
+    if kv_cache.dtype in (torch.float8_e4m3fnuz, torch.float8_e4m3fn):
+        torch.testing.assert_close(q_out.flatten(1), q, rtol=RTOL, atol=ATOL)
+        assert torch.equal(gate_out, gate)
+
+        # C++ reference: both the fused Triton kernel and the C++
+        # reshape_and_cache_flash write the platform-native fp8 flavor
+        # (OCP on gfx950, FNUZ on gfx942), so the C++ op is the unfused
+        # reference on both arches. The reference buffers are separate
+        # allocations so the reference cannot alias the fused cache.
+        slot = main["main_slot_mapping"]
+        valid = slot >= 0
+        blocks = (slot[valid] // MAIN_PAGE).long()
+        pages = (slot[valid] % MAIN_PAGE).long()
+        stored = kv_cache[blocks, pages]
+
+        key_ref = torch.zeros(
+            kv_cache.shape[:-1] + (MAIN_D,),
+            dtype=torch.uint8,
+            device=kv_cache.device,
+        )
+        value_ref = torch.zeros(
+            kv_cache.shape[:-1] + (MAIN_D,),
+            dtype=torch.uint8,
+            device=kv_cache.device,
+        )
+        k_scale_t = torch.tensor(float(main["main_k_scale"]), device="cuda")
+        v_scale_t = torch.tensor(float(main["main_v_scale"]), device="cuda")
+        reshape_and_cache_flash(
+            k.view(-1, MAIN_HK, MAIN_D),
+            v.view(-1, MAIN_HK, MAIN_D),
+            key_ref,
+            value_ref,
+            main["main_slot_mapping"],
+            "fp8",
+            k_scale_t,
+            v_scale_t,
+        )
+        # View the uint8 reference codes as the cache's fp8 dtype so
+        # _assert_codes_match's float tier compares fp8 values, not raw
+        # uint8 codes (advanced indexing produces contiguous copies, so
+        # the dtype view is valid).
+        ref_k_stored = key_ref[blocks, pages].view(kv_cache.dtype)
+        ref_v_stored = value_ref[blocks, pages].view(kv_cache.dtype)
+
+        # Compare in code space (the uint8 codes): the fused kernel and
+        # the C++ op round in a slightly different order, so a value on a
+        # step boundary may flip to the adjacent code. Allowing 2 codes
+        # covers the boundary flips; a real bug — like a scale that was
+        # never applied — moves normal-sized numbers by many codes.
+        def _assert_codes_match(actual, expected):
+            code_diff = (
+                actual.view(torch.uint8).int() - expected.view(torch.uint8).int()
+            ).abs()
+            abs_diff = (actual.float() - expected.float()).abs()
+            assert bool(
+                (
+                    (code_diff <= 2)
+                    | (abs_diff <= RTOL * expected.float().abs() + 2 * ATOL)
+                ).all()
+            )
+
+        _assert_codes_match(stored[..., :MAIN_D], ref_k_stored)
+        _assert_codes_match(stored[..., MAIN_D:], ref_v_stored)
+        return
+    key_ref = torch.zeros(
+        kv_cache.shape[:-1] + (MAIN_D,),
+        dtype=kv_cache.dtype,
+        device=kv_cache.device,
+    )
+    value_ref = torch.zeros(
+        kv_cache.shape[:-1] + (MAIN_D,),
+        dtype=kv_cache.dtype,
+        device=kv_cache.device,
+    )
     scale = torch.tensor(1.0, device="cuda")
     reshape_and_cache_flash(
         k.view(-1, MAIN_HK, MAIN_D),
         v.view(-1, MAIN_HK, MAIN_D),
-        key_cache,
-        value_cache,
+        key_ref,
+        value_ref,
         main["main_slot_mapping"],
         "auto",
         scale,
@@ -144,12 +218,16 @@ def _check_main_outputs(main: dict, q_out, gate_out, rope, positions) -> None:
     )
     torch.testing.assert_close(q_out.flatten(1), q, rtol=RTOL, atol=ATOL)
     assert torch.equal(gate_out, gate)
-    torch.testing.assert_close(kv_cache, cache, rtol=RTOL, atol=ATOL)
+    torch.testing.assert_close(kv_cache[..., :MAIN_D], key_ref, rtol=RTOL, atol=ATOL)
+    torch.testing.assert_close(kv_cache[..., MAIN_D:], value_ref, rtol=RTOL, atol=ATOL)
 
 
 @requires_qsa_kernels
 @pytest.mark.usefixtures("default_vllm_config")
-@pytest.mark.parametrize("indexer_dtype", [torch.bfloat16, torch.float8_e4m3fn])
+@pytest.mark.parametrize("main_fp8", [False, True])
+@pytest.mark.parametrize(
+    "indexer_dtype", [torch.bfloat16, pytest.param(None, id="fp8_platform")]
+)
 @pytest.mark.parametrize(
     "mrope,is_2d_positions,cache_rope_positions,state_size,seq_lens,query_lens,history_lens",
     [
@@ -173,6 +251,7 @@ def _check_main_outputs(main: dict, q_out, gate_out, rope, positions) -> None:
     ],
 )
 def test_qsa_fused_prepare_matches_unfused(
+    main_fp8,
     indexer_dtype,
     mrope,
     is_2d_positions,
@@ -184,9 +263,8 @@ def test_qsa_fused_prepare_matches_unfused(
 ) -> None:
     from vllm.model_executor.layers.rotary_embedding import get_rope
 
-    if indexer_dtype == torch.float8_e4m3fn:
-        pytest.skip("the AMD QSA indexer cache is bf16-only")
-
+    if indexer_dtype is None:
+        indexer_dtype = current_platform.fp8_dtype()
     # The FP8 comparison below is a one-ulp bound on data this test generates,
     # so an unseeded run decides for itself whether it reproduces a rounding
     # difference. Seed it, or the kernel's bit-exactness is only sampled.
@@ -338,7 +416,7 @@ def test_qsa_fused_prepare_matches_unfused(
     k_weight = torch.randn(D, dtype=torch.bfloat16, device=device) * 0.2
 
     fused_query = torch.empty(num_tokens, HQ, D, dtype=indexer_dtype, device=device)
-    main = _make_main_inputs(num_tokens)
+    main = _make_main_inputs(num_tokens, main_fp8)
     main_q_out, main_gate_out = qsa_prepare(
         projected_qk[:, : HQ * D],
         projected_qk[:, HQ * D :],
@@ -400,7 +478,7 @@ def test_qsa_fused_prepare_matches_unfused(
     if rope_positions is not None:
         qsa_store_cache_rows(rope_positions, raw_slots, position_rows)
 
-    if indexer_dtype == torch.float8_e4m3fn:
+    if indexer_dtype != torch.bfloat16:
         # Both paths round the same ~bf16 intermediates to e4m3. Bitwise
         # equality (the dsv4 indexer test's bar) does not hold here: the 1D
         # reference rope (forward_cuda) differs from _norm_rope by up to 1
@@ -410,7 +488,7 @@ def test_qsa_fused_prepare_matches_unfused(
     else:
         torch.testing.assert_close(fused_query, unfused_query, rtol=RTOL, atol=ATOL)
     assert torch.equal(fused_raw.view(torch.int16), unfused_raw.view(torch.int16))
-    if indexer_dtype == torch.float8_e4m3fn:
+    if indexer_dtype != torch.bfloat16:
         assert_fp8_within_one_ulp(fused_compressed, unfused_compressed)
     else:
         torch.testing.assert_close(
diff --git a/tests/models/qwen4_exp/test_qsa_reference.py b/tests/models/qwen4_exp/test_qsa_reference.py
index a55cc028f9db..20612005f285 100644
--- a/tests/models/qwen4_exp/test_qsa_reference.py
+++ b/tests/models/qwen4_exp/test_qsa_reference.py
@@ -7,15 +7,17 @@
 import pytest
 import torch
 
-from vllm.models.qwen4_exp.common import qsa_cache
-from vllm.models.qwen4_exp.common.qsa_cache import QSAMetadataBuilder
-from vllm.models.qwen4_exp.nvidia import indexer_qsa
+from vllm.models.qwen4_exp.common import qsa_indexer as qsa_common
+from vllm.models.qwen4_exp.common.qsa_indexer import (
+    QSAIndexerMetadataBuilder,
+    qsa_kv_cache_dtype,
+)
 from vllm.models.qwen4_exp.nvidia import (
     model as _qwen4_exp_model,  # noqa: F401
 )
+from vllm.models.qwen4_exp.nvidia import qsa_indexer
 from vllm.models.qwen4_exp.nvidia.ops import qsa as qsa_ops
 from vllm.models.qwen4_exp.nvidia.ops import qsa_indexer as qsa_indexer_ops
-from vllm.models.qwen4_exp.nvidia.qsa import qsa_kv_cache_dtype
 from vllm.platforms import current_platform
 from vllm.triton_utils import HAS_TRITON
 from vllm.v1.worker.utils import clear_layer_kv_caches
@@ -73,7 +75,7 @@ def test_qsa_mtp_index_share_updates_cache_but_skips_selection(
     )
 
     monkeypatch.setattr(
-        indexer_qsa,
+        qsa_indexer,
         "qsa_prepare",
         lambda *args, **kwargs: updates.append((args, kwargs)),
     )
@@ -88,7 +90,7 @@ def test_qsa_mtp_index_share_updates_cache_but_skips_selection(
         lambda *args, **kwargs: selections.append((args, kwargs)),
     )
 
-    actual, _ = indexer_qsa.QSAIndexer.forward(
+    actual, _ = qsa_indexer.QSAIndexer.forward(
         indexer,
         torch.zeros(2, 2),
         torch.tensor([7, 8]),
@@ -248,9 +250,9 @@ def _qsa_sparse_paged_attention_reference(
 
 
 @requires_qsa_kernels
-def test_qsa_side_metadata_marks_cudagraph_padding_inert() -> None:
+def test_qsa_indexer_metadata_marks_cudagraph_padding_inert() -> None:
     device = torch.device("cuda")
-    builder = QSAMetadataBuilder.__new__(QSAMetadataBuilder)
+    builder = QSAIndexerMetadataBuilder.__new__(QSAIndexerMetadataBuilder)
     builder.compress_ratio = 1
     builder.reorder_batch_threshold = 4
     builder.is_circular_buffer = False
@@ -302,7 +304,7 @@ def test_qsa_side_metadata_marks_cudagraph_padding_inert() -> None:
 @requires_qsa_kernels
 def test_qsa_circular_buffer_metadata_keeps_only_each_requests_suffix() -> None:
     device = torch.device("cuda")
-    builder = QSAMetadataBuilder.__new__(QSAMetadataBuilder)
+    builder = QSAIndexerMetadataBuilder.__new__(QSAIndexerMetadataBuilder)
     builder.compress_ratio = 4
     builder.reorder_batch_threshold = 1
     builder.is_circular_buffer = True
@@ -370,7 +372,7 @@ def test_qsa_circular_buffer_survives_one_speculative_step(chunk_start: int) ->
     capacity = compress_ratio * -(-(compress_ratio + num_spec) // compress_ratio)
     query_len = num_spec + 1
 
-    slots = qsa_cache.circular_qsa_slot_mapping(
+    slots = qsa_common.circular_qsa_slot_mapping(
         torch.tensor([[1]], dtype=torch.int32),
         torch.zeros(query_len, dtype=torch.int32),
         torch.arange(chunk_start, chunk_start + query_len),
@@ -384,8 +386,8 @@ def test_qsa_circular_buffer_survives_one_speculative_step(chunk_start: int) ->
 
 def _qsa_key_cache(
     block_size: int, compress_ratio: int, **kwargs
-) -> qsa_cache.QSAKeyStateCache:
-    return qsa_cache.QSAKeyStateCache(
+) -> qsa_common.QSARawKeyCache:
+    return qsa_common.QSARawKeyCache(
         head_size=64,
         dtype=torch.bfloat16,
         cache_config=SimpleNamespace(block_size=block_size),
@@ -400,7 +402,7 @@ def _qsa_key_cache(
 
 def test_qsa_state_caches_adapt_the_unified_logical_layout() -> None:
     raw_cache = _qsa_key_cache(block_size=32, compress_ratio=4)
-    compressed_cache = qsa_cache.QSACompressedKeyCache(
+    compressed_cache = qsa_common.QSACompressedKeyCache(
         head_size=64,
         dtype=torch.bfloat16,
         cache_config=SimpleNamespace(block_size=32),
@@ -431,6 +433,7 @@ def test_clearing_qsa_key_cache_releases_its_storage() -> None:
     references_before_bind = torch._C._storage_Use_Count(kv.untyped_storage()._cdata)
     raw_cache.bind_kv_cache(kv)
     assert raw_cache.key_cache.shape[-1] == 64
+    assert raw_cache.rope_position_cache is not None
     assert raw_cache.rope_position_cache.dtype == torch.int64
 
     clear_layer_kv_caches([raw_cache])
@@ -469,7 +472,7 @@ def test_qsa_kv_cache_dtype_honors_skip_layers() -> None:
 @requires_qsa_kernels
 def test_qsa_compressed_metadata_keeps_dummy_slots_inert() -> None:
     device = torch.device("cuda")
-    builder = QSAMetadataBuilder.__new__(QSAMetadataBuilder)
+    builder = QSAIndexerMetadataBuilder.__new__(QSAIndexerMetadataBuilder)
     builder.compress_ratio = 4
     builder.reorder_batch_threshold = 1
     builder.is_circular_buffer = False
@@ -508,7 +511,7 @@ def test_qsa_compressed_metadata_keeps_dummy_slots_inert() -> None:
 @requires_qsa_kernels
 @pytest.mark.usefixtures("default_vllm_config")
 def test_qsa_unfused_cache_update_ignores_padded_qk() -> None:
-    """Padded projected Q/K rows must not affect either side cache."""
+    """Padded projected Q/K rows must not affect either indexer cache."""
     from vllm.model_executor.layers.rotary_embedding import get_rope
 
     device = torch.device("cuda")
@@ -558,7 +561,7 @@ def test_qsa_unfused_cache_update_ignores_padded_qk() -> None:
     )
     padded_keys = torch.full((8, 64), torch.nan, dtype=torch.bfloat16, device=device)
     padded_keys[:5].copy_(keys)
-    indexer_qsa.QSAIndexer.forward(
+    qsa_indexer.QSAIndexer.forward(
         indexer,
         torch.cat((torch.ones_like(padded_keys), padded_keys), dim=-1),
         torch.zeros(8, dtype=torch.long, device=device),
@@ -633,7 +636,7 @@ def make_buffers() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tens
         else None
     )
     actual_buffers = make_buffers()
-    actual = qsa_cache.build_qsa_metadata_triton(
+    actual = qsa_common.build_qsa_metadata_triton(
         common,
         *actual_buffers,
         storage_block_size=2,
@@ -646,7 +649,7 @@ def make_buffers() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tens
         torch.empty_like(actual_k_work) if actual_k_work is not None else None
     )
     expected_buffers = make_buffers()
-    expected = qsa_cache._build_qsa_metadata_torch(
+    expected = qsa_common._build_qsa_metadata_torch(
         common,
         *expected_buffers,
         storage_block_size=2,
@@ -696,14 +699,14 @@ def make_buffers() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tens
     max_num_work = (num_tokens + 3) // 4
     actual_k_work = torch.empty(max_num_work, 2, dtype=torch.int32, device=device)
     expected_k_work = torch.empty_like(actual_k_work)
-    actual = qsa_cache.build_qsa_metadata_triton(
+    actual = qsa_common.build_qsa_metadata_triton(
         common,
         *make_buffers(),
         storage_block_size=8,
         compress_ratio=4,
         k_work_metadata_buffer=actual_k_work,
     )
-    expected = qsa_cache._build_qsa_metadata_torch(
+    expected = qsa_common._build_qsa_metadata_torch(
         common,
         *make_buffers(),
         storage_block_size=8,
diff --git a/tests/models/qwen4_exp/test_tpsp.py b/tests/models/qwen4_exp/test_tpsp.py
index 33e01ab24f76..1d879470ffbc 100644
--- a/tests/models/qwen4_exp/test_tpsp.py
+++ b/tests/models/qwen4_exp/test_tpsp.py
@@ -48,7 +48,9 @@
     from vllm.models.qwen4_exp.amd import model as qwen4_model
     from vllm.models.qwen4_exp.amd.model import Qwen4ExpDecoderLayer
 else:
-    from vllm.models.qwen4_exp.nvidia import model as qwen4_model
+    from vllm.models.qwen4_exp.nvidia import (  # type: ignore[no-redef]
+        model as qwen4_model,
+    )
     from vllm.models.qwen4_exp.nvidia.model import Qwen4ExpDecoderLayer
 
 
@@ -220,6 +222,7 @@ def _check_moe_sp(rank: int, config: VllmConfig) -> None:
         full_hidden, moe_input, injection = layer.mlp_hyper_connection.combine_and_mix(
             full_hidden, attn_output, injection
         )
+        assert injection is not None
         expected = tuple(
             x.clone() for x in (full_hidden, layer.mlp(moe_input), injection)
         )
diff --git a/tests/models/registry.py b/tests/models/registry.py
index 7b8223ccad97..317db44a57c8 100644
--- a/tests/models/registry.py
+++ b/tests/models/registry.py
@@ -701,7 +701,11 @@ def check_available_online(
     "ColModernVBertForRetrieval": _HfExamplesInfo(
         "ModernVBERT/colmodernvbert-merged",
     ),
-    "ColUltraVBertForRetrieval": _HfExamplesInfo("athrael-soju/ColUltraVBERT"),
+    "ColUltraVBertForRetrieval": _HfExamplesInfo(
+        "athrael-soju/ColUltraVBERT",
+        # Later revisions declare the UltraVBertModel architecture instead.
+        revision="509adbab61b71da0484b20f3baf718e15219d761",
+    ),
     "ColPaliForRetrieval": _HfExamplesInfo("vidore/colpali-v1.3-hf"),
     "ModernVBertModel": _HfExamplesInfo("ModernVBERT/colmodernvbert-merged"),
     "ColQwen3": _HfExamplesInfo(
@@ -1042,14 +1046,6 @@ def check_available_online(
         "kakaocorp/kanana-1.5-v-3b-instruct",
         trust_remote_code=True,
     ),
-    "KeyeForConditionalGeneration": _HfExamplesInfo(
-        "Kwai-Keye/Keye-VL-8B-Preview",
-        trust_remote_code=True,
-    ),
-    "KeyeVL1_5ForConditionalGeneration": _HfExamplesInfo(
-        "Kwai-Keye/Keye-VL-1_5-8B",
-        trust_remote_code=True,
-    ),
     "MoonshotKimiaForCausalLM": _HfExamplesInfo(
         "moonshotai/Kimi-Audio-7B-Instruct",
         tokenizer_mode="kimi_audio",
@@ -1120,6 +1116,17 @@ def check_available_online(
     "MiMoV2OmniForCausalLM": _HfExamplesInfo(
         "XiaomiMiMo/MiMo-V2.5-Omni", trust_remote_code=True, is_available_online=False
     ),
+    "MingV3ASRForConditionalGeneration": _HfExamplesInfo(
+        "inclusionAI/Ming-3.0-asr-tiny",
+        trust_remote_code=True,
+        use_original_num_layers=True,
+        hf_overrides={
+            "text_config": {
+                "num_hidden_layers": 4,
+                "layer_types": ["linear_attention"] * 3 + ["full_attention"],
+            }
+        },
+    ),
     "MiniCPMO": _HfExamplesInfo(
         "openbmb/MiniCPM-o-2_6",
         trust_remote_code=True,
@@ -1257,9 +1264,17 @@ def check_available_online(
     "NemotronH_Super_Omni_Reasoning_V3": _HfExamplesInfo(
         "nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16", is_available_online=False
     ),
-    # TODO: Change repo id once pertinent archs are public.
     "NemotronH_Omni_Reasoning_V3": _HfExamplesInfo(
-        "nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16", is_available_online=False
+        "nvidia/NVIDIA-Nemotron-3.5-Super-VL-120B-A12B-BF16",
+        min_transformers_version="5.18.0",
+        max_model_len=4096,
+        use_original_num_layers=True,
+        hf_overrides={
+            "text_config": {"layers_block_type": ["linear_attention", "full_attention"]}
+        },
+    ),
+    "Nemotron3_5TranscribeForConditionalGeneration": _HfExamplesInfo(
+        "nvidia/NVIDIA-Nemotron-3.5-Transcribe-30B-A4B", is_available_online=False
     ),
     "OpenCUAForConditionalGeneration": _HfExamplesInfo(
         "xlangai/OpenCUA-7B",
@@ -1481,6 +1496,7 @@ def check_available_online(
         tokenizer_mode="mistral",
     ),
     # [Encoder-decoder]
+    "CanaryForConditionalGeneration": _HfExamplesInfo("nvidia/canary-1b-v2"),
     "CohereAsrForConditionalGeneration": _HfExamplesInfo(
         "CohereLabs/cohere-transcribe-03-2026", trust_remote_code=True
     ),
@@ -1580,6 +1596,11 @@ def check_available_online(
         is_available_online=False,
         use_original_num_layers=True,  # DSpark has >1 draft block
     ),
+    "DSparkMiniMaxDraftModel": _HfExamplesInfo(
+        "MiniMaxAI/MiniMax-M3",
+        is_available_online=False,
+        use_original_num_layers=True,
+    ),
     "K3DSparkModel": _HfExamplesInfo(
         "moonshotai/Kimi-K3",
         speculative_model="Inferact/Kimi-K3-DSpark",
@@ -1858,6 +1879,11 @@ def check_available_online(
         speculative_model="nvidia/Nemotron-Super-Placeholder",
         is_available_online=False,
     ),
+    "Nemotron3_5TranscribeMTPModel": _HfExamplesInfo(
+        "nvidia/NVIDIA-Nemotron-3.5-Transcribe-30B-A4B",
+        speculative_model="nvidia/NVIDIA-Nemotron-3.5-Transcribe-30B-A4B",
+        is_available_online=False,
+    ),
     "OpenPanguMTPModel": _HfExamplesInfo(
         "FreedomIntelligence/openPangu-Ultra-MoE-718B-V1.1",
         trust_remote_code=True,
@@ -1905,6 +1931,18 @@ def check_available_online(
     ),
     "TransformersMultiModalForCausalLM": _HfExamplesInfo("BAAI/Emu3-Chat-hf"),
     "TransformersMoEForCausalLM": _HfExamplesInfo("allenai/OLMoE-1B-7B-0924"),
+    "TransformersForMTP": _HfExamplesInfo(
+        "zai-org/GLM-4.5",
+        speculative_model="zai-org/GLM-4.5",
+        speculative_method="mtp",
+        # TODO: Use a dense model once one opts into the Transformers MTP standard
+        is_available_online=False,
+    ),
+    "TransformersMoEForMTP": _HfExamplesInfo(
+        "zai-org/GLM-4.5",
+        speculative_model="zai-org/GLM-4.5",
+        speculative_method="mtp",
+    ),
     "TransformersMultiModalMoEForCausalLM": _HfExamplesInfo(
         "Qwen/Qwen3-VL-30B-A3B-Instruct"
     ),
diff --git a/tests/models/step5/test_amd_csa_metadata.py b/tests/models/step5/test_amd_csa_metadata.py
new file mode 100644
index 000000000000..4ceba08407c9
--- /dev/null
+++ b/tests/models/step5/test_amd_csa_metadata.py
@@ -0,0 +1,80 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Step-5 CSA per-step metadata on ROCm."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata
+from vllm.models.step5.amd.csa_cache import Step5CSAMetadataBuilder
+from vllm.platforms import current_platform
+
+REGION = 8
+
+pytestmark = pytest.mark.skipif(
+    not current_platform.is_rocm(), reason="Step5 CSA metadata is the ROCm path"
+)
+
+
+def _builder(num_speculative_tokens: int | None = None) -> Step5CSAMetadataBuilder:
+    config = SimpleNamespace(
+        model_config=SimpleNamespace(
+            hf_text_config=SimpleNamespace(
+                sparse_config=SimpleNamespace(region_block_size=REGION, topk=512)
+            )
+        ),
+        scheduler_config=SimpleNamespace(max_num_batched_tokens=64),
+        speculative_config=(
+            None
+            if num_speculative_tokens is None
+            else SimpleNamespace(num_speculative_tokens=num_speculative_tokens)
+        ),
+    )
+    return Step5CSAMetadataBuilder(SimpleNamespace(), [], config, torch.device("cuda"))
+
+
+def test_builder_positions_fallback():
+    """Spec-decode metadata carries no positions; the builder must derive them."""
+    spec = BatchSpec(seq_lens=[40, 9, 130], query_lens=[4, 9, 1])
+    m = create_common_attn_metadata(spec, block_size=16, device=torch.device("cuda"))
+    assert m.positions is None
+    builder = _builder()
+    md = builder.build(0, m)
+    assert md.positions is not None and md.token_req is not None
+
+    want_pos = [
+        p for s, q in zip(spec.seq_lens, spec.query_lens) for p in range(s - q, s)
+    ]
+    want_req = [r for r, q in enumerate(spec.query_lens) for _ in range(q)]
+    assert md.positions.tolist() == want_pos
+    assert md.token_req.tolist() == want_req
+
+
+def test_builder_plans_prefill_only_for_decode_or_mtp_verify():
+    """Equal-length prefills are not mistaken for decode or MTP verification."""
+    prefill = BatchSpec(seq_lens=[9, 9], query_lens=[9, 9])
+    md = _builder().build(
+        0, create_common_attn_metadata(prefill, 16, device=torch.device("cuda"))
+    )
+    assert md.prefill_groups is not None and md.prefill_groups.shape == (18, 10)
+
+    decode = BatchSpec(seq_lens=[40, 9, 130], query_lens=[1, 1, 1])
+    md = _builder().build(
+        0, create_common_attn_metadata(decode, 16, device=torch.device("cuda"))
+    )
+    assert md.prefill_groups is None
+
+    mtp = BatchSpec(seq_lens=[40, 9, 130], query_lens=[4, 4, 4])
+    md = _builder(num_speculative_tokens=3).build(
+        0, create_common_attn_metadata(mtp, 16, device=torch.device("cuda"))
+    )
+    assert md.prefill_groups is None
+
+    mixed = BatchSpec(seq_lens=[40, 9, 130], query_lens=[1, 9, 1])
+    md = _builder().build(
+        0, create_common_attn_metadata(mixed, 16, device=torch.device("cuda"))
+    )
+    assert md.prefill_groups is not None and md.prefill_groups.shape == (11, 10)
+    assert md.prefill_indexer is None  # 17 regions fit in top-k
diff --git a/tests/models/step5/test_amd_csa_paged.py b/tests/models/step5/test_amd_csa_paged.py
new file mode 100644
index 000000000000..559d16352d5a
--- /dev/null
+++ b/tests/models/step5/test_amd_csa_paged.py
@@ -0,0 +1,258 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Paged Step-5 CSA kernels vs PyTorch references on scrambled block tables.
+
+Each test mixes requests with their own block-table rows, and prefill, chunked
+prefill and decode rows, since a page- or request-addressing bug reads wrong
+tokens without changing any shape.
+"""
+
+import pytest
+import torch
+
+from vllm.models.step5.amd.ops.csa_paged import sparse_attention_paged
+from vllm.models.step5.common.ops.sparse_decode_attention import (
+    sparse_decode_attention,
+)
+from vllm.models.step5.common.ops.sparse_indexer import select_regions
+from vllm.platforms import current_platform
+
+pytestmark = pytest.mark.skipif(
+    not current_platform.is_rocm(), reason="Step5 CSA kernels are the ROCm path"
+)
+
+DEVICE = "cuda"
+REGION = 8
+PROXY = 256
+
+
+def _paged(
+    seqs: list[torch.Tensor], rows_per_page: int
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Scatter each request's rows into shared pages behind its own row of a
+    shuffled ``[requests, pages]`` block table."""
+    pages = [-(-x.shape[0] // rows_per_page) for x in seqs]
+    perm = torch.randperm(sum(pages) + 3, device=DEVICE).int()
+    table = torch.zeros(len(seqs), max(pages), dtype=torch.int32, device=DEVICE)
+    cache = torch.zeros(
+        (sum(pages) + 3, rows_per_page, *seqs[0].shape[1:]),
+        dtype=seqs[0].dtype,
+        device=DEVICE,
+    )
+    start = 0
+    for r, (x, n) in enumerate(zip(seqs, pages)):
+        table[r, :n] = perm[start : start + n]
+        start += n
+        idx = torch.arange(x.shape[0], device=DEVICE)
+        cache[table[r, idx // rows_per_page].long(), idx % rows_per_page] = x
+    return cache, table
+
+
+def ref_select(scores: torch.Tensor, positions: torch.Tensor, topk: int):
+    tokens, groups, _ = scores.shape
+    out = torch.full((tokens, groups, topk + 1), -1, dtype=torch.int32)
+    for t in range(tokens):
+        own = int(positions[t]) // REGION
+        for g in range(groups):
+            order = torch.argsort(-scores[t, g, :own].cpu(), stable=True)[:topk]
+            row = torch.cat([torch.sort(order).values, torch.tensor([own])])
+            out[t, g, -row.numel() :] = row.to(torch.int32)
+    return out
+
+
+def ref_attention(q, k, v, regions, req, positions, scale):
+    tokens, heads, _ = q.shape
+    groups = k.shape[2]
+    per_group = heads // groups
+    out = torch.zeros(q.shape, dtype=torch.float32)
+    for t in range(tokens):
+        kr, vr = k[int(req[t])], v[int(req[t])]
+        for g in range(groups):
+            ids = regions[t, g]
+            ids = ids[ids >= 0].long().cpu()
+            cols = (ids[:, None] * REGION + torch.arange(REGION)).flatten()
+            cols = cols[cols <= int(positions[t])]
+            hs = slice(g * per_group, (g + 1) * per_group)
+            p = torch.softmax(
+                q[t, hs].float().cpu() @ kr[cols, g].float().cpu().T * scale, -1
+            )
+            out[t, hs] = p @ vr[cols, g].float().cpu()
+    return out
+
+
+def test_select_regions_breaks_ties_by_region_id():
+    """Post-ReLU scores are often exactly zero; ties must pick lower ids."""
+    torch.manual_seed(0)
+    scores = torch.randn(48, 2, 40, device=DEVICE)
+    scores[scores.abs() < 0.5] = 0.0
+    positions = torch.randint(0, 40 * REGION, (48,), device=DEVICE)
+    own = (positions // REGION).repeat_interleave(2)
+    got = select_regions(scores.flatten(0, 1), own, 6, "per_row").view(48, 2, -1)
+    want = ref_select(scores, positions, 6)
+    for g, w in zip(got.cpu(), want):
+        for gg, ww in zip(g, w):
+            assert sorted(gg[gg >= 0].tolist()) == sorted(ww[ww >= 0].tolist())
+
+
+FP8_KV = [None, torch.float8_e4m3fn, torch.float8_e4m3fnuz]
+
+
+def _quant_kv(k, v, kv_dtype):
+    """KV as stored (FP8 with per-tensor scales, or BF16) and as the reference
+    sees it."""
+    if kv_dtype is None:
+        return k, v, k, v, None, None
+    k_s, v_s = 0.5, 2.0
+    fmax = torch.finfo(kv_dtype).max
+    kq = (k.float() / k_s).clamp(-fmax, fmax).to(kv_dtype)
+    vq = (v.float() / v_s).clamp(-fmax, fmax).to(kv_dtype)
+    scales = [torch.tensor([x], dtype=torch.float32, device=DEVICE) for x in (k_s, v_s)]
+    return kq, vq, kq.float() * k_s, vq.float() * v_s, *scales
+
+
+@pytest.mark.parametrize("kv_dtype", FP8_KV)
+@pytest.mark.parametrize("topk", [6, 512])
+def test_sparse_attention_paged(topk: int, kv_dtype):
+    """topk=512 covers the whole history, so it also checks dense equivalence."""
+    torch.manual_seed(0)
+    seq_len, heads, groups, dim = 203, 16, 4, 192
+    shape = (2, seq_len, groups, dim)
+    k = torch.randn(shape, dtype=torch.bfloat16, device=DEVICE)
+    v = torch.randn(shape, dtype=torch.bfloat16, device=DEVICE)
+    kq, vq, k, v, k_scale, v_scale = _quant_kv(k, v, kv_dtype)
+    # Request 0 starts at a nonzero position to exercise chunked-prefill
+    # addressing independently of the packed query row.
+    pos = torch.cat([torch.arange(37, seq_len), torch.tensor([202])]).int().to(DEVICE)
+    req = torch.tensor([0] * (seq_len - 37) + [1], dtype=torch.int32, device=DEVICE)
+    q = torch.randn(len(pos), heads, dim, dtype=torch.bfloat16, device=DEVICE)
+    scores = torch.randn(len(pos), groups, seq_len // REGION + 1, device=DEVICE)
+    regions = ref_select(scores, pos, topk).to(DEVICE)
+
+    kv_cache, bt = _paged([torch.stack([kq[r], vq[r]], 1) for r in range(2)], 64)
+    got = torch.empty_like(q)
+    sparse_attention_paged(
+        q,
+        kv_cache[:, :, 0],
+        kv_cache[:, :, 1],
+        regions,
+        bt,
+        req,
+        pos,
+        got,
+        region_size=REGION,
+        scale=dim**-0.5,
+        k_scale=k_scale,
+        v_scale=v_scale,
+    )
+    want = ref_attention(q, k, v, regions, req, pos, dim**-0.5)
+    torch.testing.assert_close(got.float().cpu(), want, atol=2e-2, rtol=2e-2)
+
+
+@pytest.mark.parametrize("kv_dtype", FP8_KV)
+def test_sparse_decode_attention(kv_dtype):
+    """MTP verify rows use each request's pages, live tail, and headwise gate."""
+    torch.manual_seed(0)
+    seq_len, heads, groups, dim, topk, query_len = 203, 16, 4, 192, 32, 4
+    shape = (2, seq_len, groups, dim)
+    k = torch.randn(shape, dtype=torch.bfloat16, device=DEVICE)
+    v = torch.randn(shape, dtype=torch.bfloat16, device=DEVICE)
+    kq, vq, k, v, k_scale, v_scale = _quant_kv(k, v, kv_dtype)
+    pos = torch.tensor(
+        [31, 32, 33, 34, seq_len - 12, seq_len - 11, seq_len - 10, seq_len - 9],
+        dtype=torch.int32,
+        device=DEVICE,
+    )
+    req = torch.arange(2, dtype=torch.int32, device=DEVICE).repeat_interleave(query_len)
+    q = torch.randn(len(pos), heads, dim, dtype=torch.bfloat16, device=DEVICE)
+    scores = torch.randn(len(pos), groups, seq_len // REGION + 1, device=DEVICE)
+    regions = ref_select(scores, pos, topk).to(DEVICE)
+    gate = torch.randn(len(pos), heads, device=DEVICE)
+
+    kv_cache, bt = _paged([torch.cat([kq[r], vq[r]], -1) for r in range(2)], 16)
+    got = torch.empty_like(q)
+    sparse_decode_attention(
+        q.view(len(pos) * groups, heads // groups, dim),
+        kv_cache.transpose(1, 2).contiguous(),
+        regions.view(len(pos) * groups, -1).long(),
+        bt,
+        pos.long().repeat_interleave(groups),
+        torch.ones(len(pos) * groups, dtype=torch.bool, device=DEVICE),
+        got.view(len(pos) * groups, heads // groups, dim),
+        scale=dim**-0.5,
+        query_len=query_len,
+        splits=4,
+        k_scale=k_scale,
+        v_scale=v_scale,
+        gate=gate,
+    )
+    want = ref_attention(q, k, v, regions, req, pos, dim**-0.5)
+    want *= torch.sigmoid(gate.float().cpu()).unsqueeze(-1)
+    torch.testing.assert_close(got.float().cpu(), want, atol=2e-2, rtol=2e-2)
+
+
+def test_summaries_on_coarse_summary_blocks():
+    """ROCm keeps summaries in their own KV group, whose block (384 tokens,
+    48 regions) differs from the main 16-token block; summaries and decode
+    scores must match the main-block layout."""
+    from vllm.models.step5.common.ops.sparse_indexer_scores import (
+        decode_summary_scores,
+    )
+    from vllm.models.step5.common.ops.summary_update import (
+        build_prefill_groups,
+        update_prefill,
+    )
+
+    torch.manual_seed(0)
+    seq, capacity = 1001, 12
+    regions = seq // REGION
+    key = torch.randn((seq, 1, PROXY), device=DEVICE).to(torch.bfloat16)
+    z = torch.randn((seq, 1, PROXY), device=DEVICE).to(torch.bfloat16)
+    positions = torch.arange(seq, device=DEVICE)
+    query = torch.randn((2, 1, 8, PROXY), device=DEVICE).to(torch.float8_e4m3fn)
+    weights = torch.rand((2, 1, 8), device=DEVICE)
+
+    def run(block: int):
+        blocks = -(-seq // block)
+        table = torch.randperm(blocks + 2, device=DEVICE)[:blocks].int()
+        mean = torch.zeros(
+            (blocks + 2, block // REGION, PROXY), dtype=torch.uint8, device=DEVICE
+        )
+        rows = torch.zeros((capacity, 2 * PROXY), dtype=torch.bfloat16, device=DEVICE)
+        slots = table[positions // block].long() * block + positions % block
+        groups = build_prefill_groups(
+            slots,
+            positions,
+            torch.zeros_like(positions),
+            torch.tensor([0, seq], device=DEVICE),
+        )
+        keep = positions >= seq - capacity
+        update_prefill(
+            key,
+            z,
+            groups,
+            positions,
+            torch.zeros_like(positions),
+            torch.where(keep, positions % capacity, -1),
+            mean,
+            rows,
+            capacity,
+        )
+        r = torch.arange(regions, device=DEVICE)
+        per_region = mean[table[r // mean.shape[1]].long(), r % mean.shape[1]]
+        scores, meta = decode_summary_scores(
+            query,
+            weights,
+            mean,
+            table[None].expand(2, -1).contiguous(),
+            slots[-1:].expand(2).contiguous(),
+            torch.tensor([seq, seq - 9], device=DEVICE),
+            regions,
+            num_sms=4,
+        )
+        return per_region, scores, meta.lengths
+
+    fine, fine_scores, lengths = run(16)
+    coarse, coarse_scores, _ = run(384)
+    assert fine.any() and torch.equal(fine, coarse)
+    for row, length in enumerate(lengths.tolist()):
+        assert torch.equal(fine_scores[row, :length], coarse_scores[row, :length])
diff --git a/tests/models/step5/test_amd_csa_warmup.py b/tests/models/step5/test_amd_csa_warmup.py
new file mode 100644
index 000000000000..4776ccdb3f2c
--- /dev/null
+++ b/tests/models/step5/test_amd_csa_warmup.py
@@ -0,0 +1,48 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""The ROCm CSA layers reach the Step5 scored-prefill warmup."""
+
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from vllm.platforms import current_platform
+
+pytestmark = pytest.mark.skipif(
+    not current_platform.is_rocm(), reason="Step5 CSA is the ROCm path"
+)
+
+
+def _worker(context):
+    return SimpleNamespace(
+        vllm_config=SimpleNamespace(
+            model_config=SimpleNamespace(
+                hf_text_config=SimpleNamespace(model_type="step5")
+            ),
+            compilation_config=SimpleNamespace(static_forward_context=context),
+        )
+    )
+
+
+def test_step5_warmup_dispatches_rocm_csa_topk():
+    from vllm.model_executor.warmup.step5_dsa_warmup import step5_dsa_warmup
+    from vllm.models.step5.amd import warmup
+    from vllm.models.step5.amd.model import Step5Attention
+
+    def layer(indexer):
+        mock = MagicMock(indexer=indexer)
+        mock.__class__ = Step5Attention
+        return mock
+
+    sparse = layer(SimpleNamespace(topk=512))
+    dense = layer(None)
+    worker = _worker({"swa": dense, "other": object(), "csa": sparse})
+
+    with patch.object(warmup, "run_step5_scored_prefill_warmup") as run:
+        step5_dsa_warmup(worker)
+    run.assert_called_once_with(worker, 512)
+
+    with patch.object(warmup, "run_step5_scored_prefill_warmup") as run:
+        step5_dsa_warmup(_worker({"swa": dense}))
+    run.assert_not_called()
diff --git a/tests/models/step5/test_amd_fp8_kv_init.py b/tests/models/step5/test_amd_fp8_kv_init.py
new file mode 100644
index 000000000000..a9dd40d58ab3
--- /dev/null
+++ b/tests/models/step5/test_amd_fp8_kv_init.py
@@ -0,0 +1,48 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Step5 on ROCm initializes with an FP8 E4M3 KV cache at block size 32."""
+
+import pytest
+
+from vllm.platforms import current_platform
+
+from ..registry import HF_EXAMPLE_MODELS
+from ..utils import dummy_hf_overrides
+
+pytestmark = pytest.mark.skipif(
+    not current_platform.is_rocm(), reason="ROCm Step5 FP8 KV path"
+)
+
+ARCH = "Step5ForConditionalGeneration"
+
+
+def test_fp8_kv_block_size_32_initializes(vllm_runner, monkeypatch):
+    from vllm.platforms.rocm import on_mi3xx
+
+    if not on_mi3xx():
+        pytest.skip("Step5 on ROCm requires MI300 or MI350")
+    info = HF_EXAMPLE_MODELS.get_hf_info(ARCH)
+    info.check_available_online(on_fail="skip")
+    monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
+    with vllm_runner(
+        info.default,
+        load_format="dummy",
+        max_model_len=info.max_model_len,
+        hf_overrides=lambda cfg: dummy_hf_overrides(
+            cfg,
+            model_arch=ARCH,
+            exist_overrides=info.hf_overrides,
+            use_original_num_layers=True,
+        ),
+        kv_cache_dtype="fp8",
+        block_size=32,
+        enforce_eager=True,
+        gpu_memory_utilization=0.8,
+        max_num_seqs=4,
+        limit_mm_per_prompt={"image": 0, "video": 0},
+    ) as llm:
+        cache = llm.llm.llm_engine.vllm_config.cache_config
+        assert cache.cache_dtype == "fp8"
+        assert cache.block_size == 32
+        out = llm.generate_greedy(["Hello"], max_tokens=4)
+        assert len(out[0][0]) > 0
diff --git a/tests/models/step5/test_strided_summary_cache.py b/tests/models/step5/test_strided_summary_cache.py
index fddc5a4d0a7f..0e075d62cfcd 100644
--- a/tests/models/step5/test_strided_summary_cache.py
+++ b/tests/models/step5/test_strided_summary_cache.py
@@ -14,7 +14,7 @@
 from vllm.platforms import current_platform
 
 pytestmark = pytest.mark.skipif(
-    not current_platform.is_cuda(), reason="Step5 kernels require CUDA"
+    not current_platform.is_cuda_alike(), reason="Step5 kernels require a GPU"
 )
 
 BLOCK_SIZE = 16
@@ -64,7 +64,7 @@ def _slots(positions):
 
 
 def _run_updates(cache: Cache, key, z):
-    from vllm.models.step5.nvidia.ops.summary_update import (
+    from vllm.models.step5.common.ops.summary_update import (
         build_prefill_groups,
         update_decode,
         update_prefill,
@@ -135,8 +135,8 @@ def test_summary_updates_match_on_strided_blocks():
 
 
 def test_summary_reads_match_on_strided_blocks():
-    from vllm.models.step5.nvidia.ops.sparse_indexer import pack_summary_pages
-    from vllm.models.step5.nvidia.ops.sparse_indexer_scores import (
+    from vllm.models.step5.common.ops.sparse_indexer import pack_summary_pages
+    from vllm.models.step5.common.ops.sparse_indexer_scores import (
         decode_summary_scores,
     )
 
diff --git a/tests/models/step5/test_warmup_dispatch.py b/tests/models/step5/test_warmup_dispatch.py
new file mode 100644
index 000000000000..11c89fa72749
--- /dev/null
+++ b/tests/models/step5/test_warmup_dispatch.py
@@ -0,0 +1,66 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Step5 scored-prefill warmup dispatch on the NVIDIA path."""
+
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from vllm.model_executor.warmup import step5_dsa_warmup as warmup
+
+
+def _worker(model_type: str = "step5"):
+    return SimpleNamespace(
+        vllm_config=SimpleNamespace(
+            model_config=SimpleNamespace(
+                hf_text_config=SimpleNamespace(model_type=model_type)
+            )
+        ),
+        device="cuda",
+    )
+
+
+def _impl(adapter):
+    return SimpleNamespace(
+        adapter=adapter, sparse_topk=512, max_num_regions=4096, indexer_heads=4
+    )
+
+
+@pytest.fixture(autouse=True)
+def _nvidia():
+    with patch.object(warmup.current_platform, "is_rocm", return_value=False):
+        yield
+
+
+def test_skips_other_models():
+    with patch.object(warmup, "_step5_dsa_impl") as find:
+        warmup.step5_dsa_warmup(_worker("deepseek_v4"))
+    find.assert_not_called()
+
+
+@pytest.mark.parametrize("ran", [True, False])
+def test_scores_prefill_only_after_warmup_step(ran):
+    adapter = MagicMock()
+    worker = _worker()
+    with (
+        patch.object(warmup, "_step5_dsa_impl", return_value=_impl(adapter)),
+        patch.object(
+            warmup, "run_step5_scored_prefill_warmup", return_value=ran
+        ) as run,
+    ):
+        warmup.step5_dsa_warmup(worker)
+    run.assert_called_once_with(worker, 512)
+    assert adapter.warmup_prefill_scoring.called is ran
+    if ran:
+        adapter.warmup_prefill_scoring.assert_called_once_with(4096, 512, 4, "cuda")
+
+
+@pytest.mark.parametrize("impl", [None, _impl(None)])
+def test_no_dsa_adapter_is_a_noop(impl):
+    with (
+        patch.object(warmup, "_step5_dsa_impl", return_value=impl),
+        patch.object(warmup, "run_step5_scored_prefill_warmup") as run,
+    ):
+        warmup.step5_dsa_warmup(_worker())
+    run.assert_not_called()
diff --git a/tests/models/test_adapters.py b/tests/models/test_adapters.py
index 62442e332e8c..2e07c778308d 100644
--- a/tests/models/test_adapters.py
+++ b/tests/models/test_adapters.py
@@ -6,6 +6,7 @@
 import torch
 from transformers import Gemma3Config, PreTrainedConfig, Qwen2Config
 
+from vllm.config import LoadConfig, VllmConfig
 from vllm.model_executor.models.adapters import (
     _create_pooling_model_cls,
     _resolve_num_labels,
@@ -97,10 +98,11 @@ def _buffer_reusing_iterator(weight_dict):
         yield name, view
 
 
-def _make_pooling_model(base_cls=SimpleModel):
+def _make_pooling_model(base_cls=SimpleModel, load_format="fastsafetensors"):
     PoolingModel = _create_pooling_model_cls(base_cls)
     model = base_cls()
     model.__class__ = PoolingModel
+    model.vllm_config = VllmConfig(load_config=LoadConfig(load_format=load_format))
     model.lm_head = StageMissingLayer("output", model.lm_head)
     return model
 
@@ -246,6 +248,25 @@ def counting_clone(self, *args, **kwargs):
     }
 
 
+def test_pooling_load_weights_skips_clone_for_owned_tensor_formats(monkeypatch):
+    """Whitelisted load formats yield owned tensors, so probing must not clone."""
+    ref = _make_reference_weights()
+    model = _make_pooling_model(load_format="runai_streamer")
+
+    clone_count = 0
+    original_clone = torch.Tensor.clone
+
+    def counting_clone(self, *args, **kwargs):
+        nonlocal clone_count
+        clone_count += 1
+        return original_clone(self, *args, **kwargs)
+
+    monkeypatch.setattr(torch.Tensor, "clone", counting_clone)
+    model.load_weights(iter(ref.items()))
+
+    assert clone_count == 0
+
+
 def _composite_config(outer_labels=None, inner_labels=None):
     """Build a multimodal-style config whose text_config is a separate object."""
     config = Gemma3Config(text_config={"num_hidden_layers": 1})
diff --git a/tests/models/test_alice_ai_config.py b/tests/models/test_alice_ai_config.py
index ede3b59622d7..8c8bdbb0f598 100644
--- a/tests/models/test_alice_ai_config.py
+++ b/tests/models/test_alice_ai_config.py
@@ -6,6 +6,7 @@
 
 import pytest
 
+from tests.v1.spec_decode.utils import fake_speculative_config
 from vllm.model_executor.models.config import AliceAIForCausalLMConfig
 
 
@@ -82,7 +83,10 @@ def test_runtime_config_rejects_unsupported_alice_modes(
         speculative_config=(
             None
             if speculative_method is None
-            else SimpleNamespace(method=speculative_method)
+            else fake_speculative_config(
+                method=speculative_method,
+                draft_model_config=SimpleNamespace(hf_config=SimpleNamespace()),
+            )
         ),
     )
 
diff --git a/tests/models/test_deepseek_v41_decoder_replay_layers.py b/tests/models/test_deepseek_v41_decoder_replay_layers.py
index b7f5ca29ec2d..fafcebb610f8 100644
--- a/tests/models/test_deepseek_v41_decoder_replay_layers.py
+++ b/tests/models/test_deepseek_v41_decoder_replay_layers.py
@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """DecoderReplayLayers runs the replay layers on the forward's replay batch."""
 
+from contextlib import contextmanager
 from unittest.mock import MagicMock
 
 import pytest
@@ -9,6 +10,7 @@
 
 from vllm.config import CompilationConfig, CUDAGraphMode
 from vllm.config.kernel import FLASHINFER_MOE_EP_BACKENDS, NATIVE_MEGA_MOE_BACKENDS
+from vllm.distributed.parallel_state import GraphCaptureContext
 from vllm.forward_context import (
     BatchDescriptor,
     ForwardContext,
@@ -158,8 +160,16 @@ def test_replay_graph_matches_eager(monkeypatch, moe_backend, ids_dtype):
     """The replay graph of the next captured size pads and runs the rows. Input ids
     arrive as the model passes them: int32, cast to int64 for native MegaMoE only."""
     monkeypatch.setattr(cudagraph_utils, "get_pp_group", MagicMock)
-    capture = torch.cuda.stream(torch.cuda.Stream())
-    monkeypatch.setattr(cudagraph_utils, "graph_capture", lambda device: capture)
+
+    @contextmanager
+    def fake_graph_capture(device, stream=None):
+        if stream is None:
+            stream = torch.cuda.Stream(device=device)
+        stream.wait_stream(torch.cuda.current_stream(device))
+        with torch.cuda.stream(stream):
+            yield GraphCaptureContext(stream)
+
+    monkeypatch.setattr(cudagraph_utils, "graph_capture", fake_graph_capture)
     compilation = CompilationConfig(decoder_replay_cudagraph_capture_sizes=[4])
     cfg = MagicMock(compilation_config=compilation, speculative_config=None)
     # Mocking it true would capture into the cuMem pool, which breaks on ROCm.
diff --git a/tests/models/test_deepseek_v41_pcp.py b/tests/models/test_deepseek_v41_pcp.py
index 5a1a8a5efe68..d1e5dd1bb04e 100644
--- a/tests/models/test_deepseek_v41_pcp.py
+++ b/tests/models/test_deepseek_v41_pcp.py
@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import os
 from types import SimpleNamespace
 from unittest.mock import Mock
 
@@ -13,6 +14,55 @@
     ReplayBatch,
 )
 from vllm.models.deepseek_v41.nvidia import model as model_module
+from vllm.platforms import current_platform
+
+
+@pytest.mark.distributed(num_gpus=8)
+@pytest.mark.skipif(
+    not current_platform.is_cuda() or current_platform.device_count() < 8,
+    reason="Requires eight CUDA GPUs",
+)
+@pytest.mark.parametrize("pp_size,dp_size", [(2, 1), (1, 2)], ids=["pp-pcp", "dp-pcp"])
+def test_pcp_generates_with_pipeline_or_data_parallelism(pp_size, dp_size):
+    """PCP composes with PP or DP through real-weight serving and generation."""
+    from tests.utils import RemoteOpenAIServer
+
+    model = os.getenv("MODEL_NAME", "deepseek-ai/DeepSeek-V4.1-Flash")
+    args = f"""
+        --language-model-only --tokenizer-mode deepseek_v41
+        --tensor-parallel-size 2 --prefill-context-parallel-size 2
+        --pipeline-parallel-size {pp_size} --data-parallel-size {dp_size}
+        --enable-expert-parallel --all2all-backend allgather_reducescatter
+        --distributed-executor-backend mp --attention-backend FLASHMLA_SPARSE_DSV41
+        --moe-backend marlin --engram-config {{"cpu_offload":true}}
+        --kv-cache-dtype fp8_ds_mla --block-size 64
+        --max-model-len 4096 --max-num-batched-tokens 512 --max-num-seqs 8
+        --enable-chunked-prefill --no-enable-prefix-caching --no-swa-bounded-replay
+        --enforce-eager --no-async-scheduling
+    """.split()
+    env = {"VLLM_USE_V2_MODEL_RUNNER": "1"}
+    if pp_size > 1:
+        env["VLLM_PP_LAYER_PARTITION"] = "14,26"
+    with (
+        RemoteOpenAIServer(model, args, env_dict=env, max_wait_seconds=1200) as server,
+        server.get_client(timeout=300) as client,
+    ):
+        for rank in range(dp_size):
+            response = client.completions.create(
+                model=model,
+                prompt="A lighthouse guides ships through the night. " * 129,
+                temperature=0,
+                max_tokens=8,
+                extra_body={"ignore_eos": True},
+                extra_headers=(
+                    {"X-data-parallel-rank": str(rank)} if dp_size > 1 else None
+                ),
+            )
+            assert len(response.choices) == 1
+            assert response.choices[0].finish_reason == "length"
+            assert response.choices[0].text
+            assert response.usage is not None
+            assert response.usage.completion_tokens == 8
 
 
 @pytest.mark.parametrize("pcp_size", [1, 2])
diff --git a/tests/models/test_deepseek_v41_replay_batch.py b/tests/models/test_deepseek_v41_replay_batch.py
index 7e492e7667e3..d47a25e919a2 100644
--- a/tests/models/test_deepseek_v41_replay_batch.py
+++ b/tests/models/test_deepseek_v41_replay_batch.py
@@ -178,7 +178,7 @@ def test_replay_batch_keeps_each_request_window(state):
         "swa_first": first.attn_metadata,
     }
     assert torch.equal(context.slot_mapping["mla"], _slot_mappings(401)[1, rows])
-    assert context.dp_metadata is None and context.is_padding is None
+    assert context.dp_metadata is None and not context.is_padding.any()
 
 
 def test_replay_context_restores_tp_routing_rows(state):
@@ -186,7 +186,7 @@ def test_replay_context_restores_tp_routing_rows(state):
     batch = _input_batch(QUERY_LENS, SEQ_LENS, PREFILLING)
     replay, _ = _prepare(state, batch, CUDAGraphMode.NONE)
     context = replay.forward_context
-    assert context.is_padding is None and context.dp_metadata is None
+    assert context.dp_metadata is None
     count = len(REPLAY_ROWS)
     assert context.batch_descriptor.num_tokens == count
     capturer = RoutedExpertsCapturer.__new__(RoutedExpertsCapturer)
@@ -457,6 +457,7 @@ def test_idle_dp_rank_dummy_trims_with_its_peers(dp_state):
     assert replay is not None and replay.rows.shape[0] == 1
     dp_metadata = replay.forward_context.dp_metadata
     assert dp_metadata.num_tokens_across_dp_cpu.tolist() == [1, 129]
+    assert replay.forward_context.is_padding.tolist() == [True]
 
 
 def test_zero_prefix_window_warms_gather_and_skips_pad(monkeypatch):
@@ -546,5 +547,9 @@ def test_dp_ranks_pad_to_one_replay_graph(dp_state):
     replay, _ = _prepare(dp_state, batch, CUDAGraphMode.NONE)
     assert dp_state.replay_cudagraphs.dispatch.call_args.args[:2] == (3, 300)
     assert replay.forward_context.batch_descriptor.num_tokens == 512
+    is_padding = replay.forward_context.is_padding
+    assert is_padding.shape == (512,)
+    assert not is_padding[: len(REPLAY_ROWS)].any()
+    assert is_padding[len(REPLAY_ROWS) :].all()
     dp_metadata = replay.forward_context.dp_metadata
     assert dp_metadata.num_tokens_across_dp_cpu.tolist() == [512, 512]
diff --git a/tests/models/test_deepseek_v4_vl.py b/tests/models/test_deepseek_v4_vl.py
new file mode 100644
index 000000000000..6f38cb1f55f2
--- /dev/null
+++ b/tests/models/test_deepseek_v4_vl.py
@@ -0,0 +1,70 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""The loader reaches the language model's post-load hook inside VL wrappers."""
+
+import importlib
+
+import pytest
+import torch
+from torch import nn
+
+from vllm.model_executor.model_loader.weight_utils import finalize_model_weights
+from vllm.model_executor.models.utils import WeightsMapper
+
+
+class _FakeLanguageModel(nn.Module):
+    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()))
+
+
+@pytest.mark.parametrize(
+    "module,cls_name",
+    [
+        (
+            "vllm.models.deepseek_v4.common.vl_model",
+            "DeepseekV4ForConditionalGeneration",
+        ),
+        ("vllm.models.deepseek_v41.nvidia.vl_model", "DeepseekV41ForCausalLM"),
+        ("vllm.models.deepseek_v41.amd.vl_model", "DeepseekV41ForCausalLM"),
+    ],
+    ids=["v4", "v41-nvidia", "v41-amd"],
+)
+def test_vl_wrapper_runs_language_model_hook(module, cls_name) -> None:
+    wrapper_cls = getattr(importlib.import_module(module), cls_name)
+
+    model = object.__new__(wrapper_cls)
+    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 == []
+
+    finalize_model_weights(model)
+
+    assert model.language_model.finalized_values == [(1.0, 3.0)]
+
+    # A reload reruns the hook: no one-shot guard may skip it.
+    finalize_model_weights(model)
+    assert model.language_model.finalized_values == [(1.0, 3.0)] * 2
diff --git a/tests/models/test_deepseek_v4_vl_rocm.py b/tests/models/test_deepseek_v4_vl_rocm.py
index 762f1cd6c0e0..5ce8b2dee22d 100644
--- a/tests/models/test_deepseek_v4_vl_rocm.py
+++ b/tests/models/test_deepseek_v4_vl_rocm.py
@@ -8,7 +8,6 @@
 from torch import nn
 
 from vllm.config import KernelConfig, ParallelConfig
-from vllm.model_executor.models.utils import WeightsMapper
 from vllm.platforms import current_platform
 
 pytestmark = pytest.mark.skipif(
@@ -249,19 +248,71 @@ def fake_factory(**kwargs):
 def test_rocm_fse_rejects_data_parallel_deployment(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
-    from vllm.models.deepseek_v4.amd.model import _fuse_shared_experts_enabled
+    from vllm.config import set_current_vllm_config
+    from vllm.model_executor.layers.fused_moe.utils import (
+        resolve_layer_fused_shared_expert,
+    )
+    from vllm.models.deepseek_v4.amd import model as rocm_model
 
-    config = SimpleNamespace(n_shared_experts=1)
+    class FakeGate(nn.Module):
+        def __init__(self, **kwargs) -> None:
+            super().__init__()
 
-    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))
+    monkeypatch.setattr(rocm_model, "GateLinear", FakeGate)
+    monkeypatch.setattr(rocm_model, "FusedMoEFactory", lambda **kwargs: nn.Identity())
+    monkeypatch.setattr(rocm_model, "get_tensor_model_parallel_world_size", lambda: 1)
+    monkeypatch.setattr(rocm_model, "get_tensor_model_parallel_rank", lambda: 0)
+    monkeypatch.setattr(
+        rocm_model.rocm_aiter_ops,
+        "is_fusion_moe_shared_experts_enabled",
+        lambda: True,
+    )
 
-    # 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)
+    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=0,
+        vocab_size=32,
+        topk_method="noaux_tc",
+        vision_n_layers=0,
+        n_shared_experts=1,
     )
-    assert _fuse_shared_experts_enabled(config, ParallelConfig())
+
+    def build_vllm_config(parallel_config: ParallelConfig) -> SimpleNamespace:
+        return SimpleNamespace(
+            model_config=SimpleNamespace(hf_config=config),
+            quant_config=None,
+            parallel_config=parallel_config,
+            kernel_config=KernelConfig(),
+        )
+
+    # Data parallelism without expert parallelism fails fast at MoE construction.
+    with pytest.raises(ValueError, match="not supported with data parallelism"):
+        rocm_model.DeepseekV4MoE(
+            build_vllm_config(ParallelConfig(data_parallel_size=2)),
+            prefix="model.layers.0.ffn",
+        )
+
+    # Under expert parallelism the resolver disables FSE instead of raising.
+    with set_current_vllm_config(
+        build_vllm_config(
+            ParallelConfig(data_parallel_size=2, enable_expert_parallel=True)
+        )
+    ):
+        assert not resolve_layer_fused_shared_expert(
+            None, "model.layers.0.ffn", disable_under_ep=True
+        )
+
+    # Nothing blocks FSE once expert parallelism is off.
+    with set_current_vllm_config(build_vllm_config(ParallelConfig())):
+        assert resolve_layer_fused_shared_expert(
+            None, "model.layers.0.ffn", disable_under_ep=True
+        )
 
 
 def test_rocm_mtp_forwards_input_ids_for_vision_routing(
@@ -356,127 +407,3 @@ def logits_processor(
 
     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_deepseek_v4_vl_stream_weights.py b/tests/models/test_deepseek_v4_vl_stream_weights.py
deleted file mode 100644
index df28b00ddba2..000000000000
--- a/tests/models/test_deepseek_v4_vl_stream_weights.py
+++ /dev/null
@@ -1,78 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""The DeepSeek-V4/V4.1 VL wrappers stream the language model's weights to it
-as one contiguous group."""
-
-import pytest
-import torch
-from torch import nn
-
-from vllm.model_executor.models.utils import WeightsMapper
-from vllm.platforms import current_platform
-
-pytestmark = pytest.mark.skipif(
-    not current_platform.is_cuda_alike(), reason="CUDA or ROCm implementations"
-)
-
-
-class _FakeFinalizingLanguageModel(nn.Module):
-    finalizes_weights_during_load = True
-
-    def __init__(self) -> None:
-        super().__init__()
-        self.tensor_a = nn.Parameter(torch.zeros(1))
-        self.tensor_c = nn.Parameter(torch.zeros(1))
-        self.load_calls = 0
-        self.finalized_values: list[tuple[float, float]] = []
-
-    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.finalized_values.append((self.tensor_a.item(), self.tensor_c.item()))
-        return loaded
-
-
-def _v4_wrapper():
-    from vllm.models.deepseek_v4.common.vl_model import (
-        DeepseekV4ForConditionalGeneration,
-    )
-
-    return DeepseekV4ForConditionalGeneration
-
-
-def _v41_wrapper():
-    from vllm.models.deepseek_v41 import DeepseekV41ForCausalLM
-
-    return DeepseekV41ForCausalLM
-
-
-@pytest.mark.parametrize("wrapper", [_v4_wrapper, _v41_wrapper], ids=["v4", "v41"])
-def test_load_weights_streams_language_model_group(wrapper) -> None:
-    model = object.__new__(wrapper())
-    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()
-
-    def interleaved_weights():
-        yield "language_model.tensor_a", torch.tensor([1.0])
-        # Loaded before the stream is drained, not after.
-        assert model.language_model.tensor_a.item() == 1.0
-        # Splits the language model's weights unless the wrapper regroups them.
-        yield "vision.tensor_b", torch.tensor([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.load_calls == 1
-    assert model.language_model.finalized_values == [(1.0, 3.0)]
-    assert model.vision.tensor_b.item() == 2.0
diff --git a/tests/models/test_dsv41_mono_decode.py b/tests/models/test_dsv41_mono_decode.py
index 07a6e66daa05..9831014f8492 100644
--- a/tests/models/test_dsv41_mono_decode.py
+++ b/tests/models/test_dsv41_mono_decode.py
@@ -8,6 +8,7 @@
 unreduced output."""
 
 import sys
+from collections import defaultdict
 from types import ModuleType, SimpleNamespace
 
 import pytest
@@ -732,6 +733,64 @@ def test_rocm_attention_takes_the_gfx942_indexer_with_the_flag(
     assert rocm.DeepseekV41Gfx942Indexer.attn_cls is Gfx942SparseAttnIndexer
 
 
+@pytest.mark.parametrize("monokernel", [True, False])
+@pytest.mark.parametrize("cdna", [3, 4])
+def test_swa_decode_meta_runs_on_cdna3_and_cdna4_with_the_flag(
+    monkeypatch, cdna, monokernel
+):
+    """swa_decode_meta's kernel is plain Triton, so with --monokernel a decode
+    batch gets its SWA metadata from the one launch on CDNA4 (gfx950) as on
+    CDNA3 (gfx942). Without the flag, the builder keeps vLLM's build."""
+    import vllm.platforms.rocm as rocm_platform
+    from vllm.models.deepseek_v41.amd import ops, rocm
+    from vllm.models.deepseek_v41.amd.ops import swa_decode_meta
+
+    monkeypatch.setattr(ops, "current_platform", SimpleNamespace(is_rocm=lambda: True))
+    monkeypatch.setattr(rocm_platform, "get_cdna_version", lambda: cdna)
+    monkeypatch.setattr(ops, "_monokernel", ops._monokernel)
+    ops.record_monokernel(_config(monokernel=monokernel))
+    launches = []
+    monkeypatch.setattr(
+        swa_decode_meta, "launch", lambda *args, noncausal: launches.append(noncausal)
+    )
+    ints = lambda *shape: torch.zeros(shape, dtype=torch.int32)  # noqa: E731
+    builder = SimpleNamespace(
+        decode_threshold=M,
+        window_size=md.SWA_WIDTH,
+        block_size=128,
+        decode_swa_indices=ints(64, 1, md.SWA_WIDTH),
+        decode_swa_lens=ints(64),
+        is_valid_token=torch.zeros(64, dtype=torch.bool),
+        token_to_req_indices=ints(64),
+        no_replay_start=ints(2),
+        decode_swa_ragged_indices_buffer=ints(64 * md.SWA_WIDTH),
+        decode_swa_ragged_indptr_buffer=ints(65),
+        build_tile_scheduler=lambda rows: defaultdict(lambda: None),
+        _build_deepseek_v4_metadata=lambda *args: {},
+    )
+    query_start_loc = torch.tensor([0, M], dtype=torch.int32)
+    cam = SimpleNamespace(
+        causal=True,
+        max_query_len=M,
+        num_reqs=1,
+        num_actual_tokens=M,
+        query_start_loc=query_start_loc,
+        query_start_loc_cpu=query_start_loc,
+        seq_lens=torch.tensor([1000], dtype=torch.int32),
+        seq_lens_cpu_upper_bound=torch.tensor([1000], dtype=torch.int32),
+        slot_mapping=torch.arange(M),
+        block_table_tensor=ints(1, 8),
+        token_to_req_indices=lambda buffer: buffer[:M],
+    )
+    out = swa_decode_meta.build(
+        builder, 0, cam, False, None, rocm.DeepseekV41ROCMAiterSparseSWAMetadata
+    )
+    # No later test may see this test's cached gate.
+    ops.record_monokernel(_config(monokernel=False))
+    assert launches == ([False] if monokernel else [])
+    assert (out is not None) is monokernel
+
+
 def test_dspark_speculator_asks_the_draft_model_first():
     """A draft model that samples a step itself returns True, and vLLM's
     loop then does not run. The stub speculator has none of the loop's
diff --git a/tests/models/test_dsv41_mono_layers.py b/tests/models/test_dsv41_mono_layers.py
index 6c405d03d145..1e428d83d302 100644
--- a/tests/models/test_dsv41_mono_layers.py
+++ b/tests/models/test_dsv41_mono_layers.py
@@ -70,7 +70,7 @@ def _arch() -> str:
 # The (tensor parallel size, decode rows) cases of each GPU. The gfx942 kernels
 # build for TP4 only. 8 rows are two GEMV passes of 6 rows there, the second
 # one padded, and 42 rows are the most that they take.
-CASES = {"cdna4": [(2, 8), (4, 8)], "gfx942": [(4, 8), (4, 42)]}
+CASES = {"cdna4": [(2, 8), (4, 8), (8, 8), (8, 28)], "gfx942": [(4, 8), (4, 42)]}
 
 
 @torch.no_grad()
@@ -314,7 +314,7 @@ def _rel(a: torch.Tensor, b: torch.Tensor) -> float:
     return ((a - b).norm() / b.norm()).item()
 
 
-@pytest.mark.parametrize(("tp", "rows"), [(2, 8), (4, 8), (4, 42)])
+@pytest.mark.parametrize(("tp", "rows"), [(2, 8), (4, 8), (4, 42), (8, 8), (8, 28)])
 def test_mono_decode_matches_the_model(vllm_runner, monkeypatch, tmp_path, tp, rows):
     arch = _arch()
     if not arch:
diff --git a/tests/models/test_dsv41_mono_numerics.py b/tests/models/test_dsv41_mono_numerics.py
index 04488957d828..095d8deaaf57 100644
--- a/tests/models/test_dsv41_mono_numerics.py
+++ b/tests/models/test_dsv41_mono_numerics.py
@@ -16,6 +16,9 @@
 rows, and replays the steps in one CUDA graph. The real model's numerics on
 gfx942 are checked by ``test_dsv41_mono_layers.py``."""
 
+import shutil
+import subprocess
+
 import pytest
 import ray
 import torch
@@ -562,7 +565,7 @@ def step(mw=mw, x0=x0, res0=res0, post0=post0, comb0=comb0, pre0=pre0,
 
 
 @pytest.mark.skipif(not _on_cdna4(), reason="the mono kernels target CDNA4")
-@pytest.mark.parametrize("tp", [2, 4])
+@pytest.mark.parametrize("tp", [2, 4, 8])
 def test_mono_decode_layer_matches_vllm(monkeypatch: pytest.MonkeyPatch, tp: int):
     if torch.accelerator.device_count() < tp:
         pytest.skip(f"needs {tp} GPUs")
@@ -570,6 +573,23 @@ def test_mono_decode_layer_matches_vllm(monkeypatch: pytest.MonkeyPatch, tp: int
     multi_process_parallel(monkeypatch, tp, 1, _mono_numerics)
 
 
+@pytest.mark.skipif(not _on_cdna4(), reason="the mono kernels target CDNA4")
+@pytest.mark.skipif(not shutil.which("llvm-objdump"), reason="needs llvm-objdump")
+def test_mono_decode_k2_tp4_no_scratch(monkeypatch: pytest.MonkeyPatch):
+    """TP4 K2 at 14 rows, where the MoE push's peer bases spilled: no scratch."""
+    flyc = pytest.importorskip("flydsl.compiler")
+    from vllm.models.deepseek_v41.amd.mono.layer import MonoBuild, build_mono_k2
+    from vllm.models.minimax_m3.amd.specialized_decode.mono.residency import code_object
+
+    monkeypatch.setenv("COMPILE_ONLY", "1")
+    monkeypatch.setenv("FLYDSL_RUNTIME_ENABLE_CACHE", "0")
+    k2 = build_mono_k2(MonoBuild(14, 4, 2))
+    flyc.compile(k2, *[0] * (k2.func.__code__.co_argcount - 1))
+    elf = code_object(k2._last_compiled[1]._ir_text)
+    isa = subprocess.check_output(["llvm-objdump", "-d", "-"], input=elf)
+    assert b"scratch_" not in isa
+
+
 # ---------------------------------------------------------------- gfx942
 
 
diff --git a/tests/models/test_initialization.py b/tests/models/test_initialization.py
index f2ea22044457..6d81b8bb8443 100644
--- a/tests/models/test_initialization.py
+++ b/tests/models/test_initialization.py
@@ -184,7 +184,12 @@ def _initialize_kv_caches_v1(self, vllm_config):
         from vllm.platforms import current_platform
         from vllm.platforms.interface import DeviceCapability
 
-        if (
+        if current_platform.is_rocm():
+            from vllm.platforms.rocm import on_mi3xx
+
+            if not on_mi3xx():
+                pytest.skip("Step5 on ROCm requires MI300 or MI350")
+        elif (
             not current_platform.is_cuda()
             or current_platform.get_device_capability()
             not in (
diff --git a/tests/models/test_mimo_v2_attention_backend.py b/tests/models/test_mimo_v2_attention_backend.py
new file mode 100644
index 000000000000..14279b6ee4aa
--- /dev/null
+++ b/tests/models/test_mimo_v2_attention_backend.py
@@ -0,0 +1,205 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""MiMo-V2.6 has head_dim 192 but v_head_dim 128, so only backends that keep
+V's width out of the K path can serve it. An explicit ``--attention-backend``
+from that set has to be honored; any other request has to fall back to the
+DiffKV auto-pick with a warning rather than mount a backend that cannot serve
+the shape.
+"""
+
+import pytest
+import torch
+
+from vllm.config import AttentionConfig, VllmConfig, set_current_vllm_config
+from vllm.platforms import current_platform
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
+
+HEAD_DIM = 192
+V_HEAD_DIM = 128
+
+DIFFKV_BACKENDS = (
+    AttentionBackendEnum.FLASH_ATTN_DIFFKV,
+    AttentionBackendEnum.TRITON_ATTN_DIFFKV,
+)
+
+
+@pytest.fixture(autouse=True)
+def bf16_default_dtype():
+    # Backend selection is dtype sensitive; use the dtype the model ships in.
+    default_dtype = torch.get_default_dtype()
+    torch.set_default_dtype(torch.bfloat16)
+    yield
+    torch.set_default_dtype(default_dtype)
+
+
+@pytest.fixture(autouse=True)
+def restore_diffkv_head_size_v():
+    """``set_head_size_v`` writes backend class state.
+
+    Leaving it modified would make these tests order dependent, so snapshot and
+    restore it on every backend that can be selected here.
+    """
+    from vllm.v1.attention.backends.flash_attn_diffkv import (
+        FlashAttentionDiffKVBackend,
+    )
+    from vllm.v1.attention.backends.triton_attn_diffkv import (
+        TritonAttentionDiffKVBackend,
+    )
+
+    backends = (FlashAttentionDiffKVBackend, TritonAttentionDiffKVBackend)
+    saved = [backend.head_size_v for backend in backends]
+    yield
+    for backend, head_size_v in zip(backends, saved):
+        backend.head_size_v = head_size_v
+
+
+def _build_attention(
+    backend: AttentionBackendEnum | None, v_head_dim: int = V_HEAD_DIM
+):
+    from vllm.model_executor.models.mimo_v2 import MiMoV2Attention
+
+    config = VllmConfig(attention_config=AttentionConfig(backend=backend))
+    with set_current_vllm_config(config):
+        layer = MiMoV2Attention(
+            hidden_size=512,
+            num_heads=4,
+            num_kv_heads=2,
+            head_dim=HEAD_DIM,
+            v_head_dim=v_head_dim,
+            prefix="model.layers.0.self_attn",
+        )
+    return layer
+
+
+def test_default_backend_is_diffkv(dist_init) -> None:
+    """No explicit backend: the DiffKV auto-pick is unchanged."""
+    layer = _build_attention(None)
+    assert layer.attn.backend in DIFFKV_BACKENDS
+    # V's width has to reach the layer: Attention defaults head_size_v to
+    # head_size (192), which would corrupt the output reshape and KV spec.
+    assert layer.attn.head_size_v == V_HEAD_DIM
+
+
+@pytest.mark.skipif(
+    not current_platform.is_rocm(), reason="ROCM_AITER_UNIFIED_ATTN is ROCm only"
+)
+def test_explicit_unified_attention_is_honored(dist_init) -> None:
+    from vllm._aiter_ops import IS_AITER_FOUND
+
+    if not IS_AITER_FOUND:
+        pytest.skip("aiter is not installed")
+
+    layer = _build_attention(AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN)
+    assert layer.attn.backend is AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN
+    # The enum alone would not catch the bug this guards: the impl has to be
+    # unified attention's, not the DiffKV one the auto-pick installs.
+    assert type(layer.attn.impl).__name__ == "RocmAiterUnifiedAttentionImpl"
+
+
+def test_explicit_diffkv_backend_is_honored(dist_init) -> None:
+    """The accept path also works for the DiffKV family, on any platform."""
+    layer = _build_attention(AttentionBackendEnum.TRITON_ATTN_DIFFKV)
+    assert layer.attn.backend is AttentionBackendEnum.TRITON_ATTN_DIFFKV
+    assert layer.attn.head_size_v == V_HEAD_DIM
+
+
+def test_explicit_symmetric_backend_falls_back_to_diffkv(dist_init) -> None:
+    """A backend that cannot take head_size_v != head_size must not be mounted."""
+    layer = _build_attention(AttentionBackendEnum.TRITON_ATTN)
+    assert layer.attn.backend in DIFFKV_BACKENDS
+
+
+def test_layer_v_head_dim_reaches_the_backend(dist_init) -> None:
+    """The backend must be told V's width before its impl is built.
+
+    The DiffKV backends default to ``head_size_v = 128``, which is this model's
+    own V width, so a missing ``set_head_size_v`` call is invisible at the
+    model's geometry. That attribute is what the metadata builder and forward
+    read, so assert it carries the layer's width instead of the default.
+    """
+    from vllm.v1.attention.backends.triton_attn_diffkv import (
+        TritonAttentionDiffKVBackend,
+    )
+
+    _build_attention(AttentionBackendEnum.TRITON_ATTN_DIFFKV, v_head_dim=64)
+    assert TritonAttentionDiffKVBackend.head_size_v == 64
+
+
+def test_only_backends_that_can_serve_v_head_size_are_admitted() -> None:
+    """`supports_head_size_v` is what the model's gate consults."""
+    from vllm.v1.attention.backends.flash_attn_diffkv import (
+        FlashAttentionDiffKVBackend,
+    )
+    from vllm.v1.attention.backends.rocm_aiter_unified_attn import (
+        RocmAiterUnifiedAttentionBackend,
+    )
+    from vllm.v1.attention.backends.triton_attn_diffkv import (
+        TritonAttentionDiffKVBackend,
+    )
+
+    assert FlashAttentionDiffKVBackend.supports_head_size_v(HEAD_DIM, V_HEAD_DIM)
+    assert TritonAttentionDiffKVBackend.supports_head_size_v(HEAD_DIM, V_HEAD_DIM)
+    # The packed cache is split with `split(head_size)`, so V cannot be wider.
+    assert RocmAiterUnifiedAttentionBackend.supports_head_size_v(HEAD_DIM, V_HEAD_DIM)
+    assert not RocmAiterUnifiedAttentionBackend.supports_head_size_v(
+        V_HEAD_DIM, HEAD_DIM
+    )
+
+
+def test_cache_write_matches_pytorch_reference(dist_init) -> None:
+    """Fill the packed cache through the impl, then compare with torch.
+
+    The layer hands the kernel two views of one layer buffer, and V has its own
+    width, so the check is what lands in that buffer: K in the first
+    ``head_dim`` slots of a row, V in the remaining ``v_head_dim``. Same shape
+    as the aiter FA check in tests/kernels/attention/test_rocm_aiter_fa.py, for
+    the asymmetric case those backends cannot serve.
+    """
+    from vllm._aiter_ops import IS_AITER_FOUND
+
+    if not IS_AITER_FOUND:
+        pytest.skip("aiter is not installed")
+
+    from tests.kernels.allclose_default import get_default_atol, get_default_rtol
+
+    layer = _build_attention(AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN)
+    impl = layer.attn.impl
+    num_kv_heads = layer.num_kv_heads
+    num_tokens, block_size, num_blocks = 6, 16, 2
+
+    # One packed layer buffer per token row: [K head_dim | V v_head_dim]. The
+    # spare trailing block absorbs an over-wide write from a regressed kernel,
+    # so that fails the comparison instead of writing out of bounds.
+    packed = torch.zeros(
+        num_blocks + 1,
+        num_kv_heads,
+        block_size,
+        HEAD_DIM + V_HEAD_DIM,
+        dtype=torch.bfloat16,
+        device="cuda",
+    )
+    key = torch.randn(
+        num_tokens, num_kv_heads, HEAD_DIM, dtype=torch.bfloat16, device="cuda"
+    )
+    value = torch.randn(
+        num_tokens, num_kv_heads, V_HEAD_DIM, dtype=torch.bfloat16, device="cuda"
+    )
+    # Slot 31 is the last slot of the cache.
+    slot_mapping = torch.tensor([0, 5, 16, 31, 7, 17], dtype=torch.long, device="cuda")
+
+    # The impl reads only the scales off `layer`; they are not the subject here.
+    scales = torch.nn.Module()
+    scales._k_scale = torch.ones((), dtype=torch.float32, device="cuda")
+    scales._v_scale = torch.ones((), dtype=torch.float32, device="cuda")
+    impl.do_kv_cache_update(scales, key, value, packed, slot_mapping)
+
+    expected = torch.zeros_like(packed)
+    for token, slot in enumerate(slot_mapping.tolist()):
+        block, offset = divmod(slot, block_size)
+        # packed buffer is [blocks, kv_heads, block_tokens, content]
+        expected[block, :, offset, :HEAD_DIM] = key[token]
+        expected[block, :, offset, HEAD_DIM:] = value[token]
+
+    torch.testing.assert_close(
+        packed, expected, atol=get_default_atol(packed), rtol=get_default_rtol(packed)
+    )
diff --git a/tests/models/test_minimax_m3_mega_moe.py b/tests/models/test_minimax_m3_mega_moe.py
new file mode 100644
index 000000000000..e4b834e3f832
--- /dev/null
+++ b/tests/models/test_minimax_m3_mega_moe.py
@@ -0,0 +1,295 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""MiniMax-M3.1 on the DeepGEMM NVFP4 MegaMoE layer: gating, the checkpoint
+adapter, and the fused block against a dequantized reference."""
+
+import gc
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm.model_executor.layers.fused_moe.deep_gemm_mega_moe import (
+    DeepGemmMegaMoEExperts,
+)
+from vllm.models.minimax_m3.nvidia.model import (
+    MEGA_MOE_FC2_INPUT_SCALE,
+    _mega_moe_expert_params_mapping,
+    _mega_moe_expert_weight,
+    _use_mega_moe,
+)
+from vllm.platforms import current_platform
+
+pytestmark = pytest.mark.skipif(
+    not current_platform.is_cuda(), reason="NVFP4 MegaMoE requires CUDA"
+)
+
+
+def _vllm_config(moe_backend: str, moe_quant_algo: str | None):
+    return SimpleNamespace(
+        kernel_config=SimpleNamespace(moe_backend=moe_backend),
+        quant_config=SimpleNamespace(moe_quant_algo=moe_quant_algo),
+        parallel_config=SimpleNamespace(enable_expert_parallel=True),
+    )
+
+
+@pytest.mark.parametrize(
+    ("moe_backend", "moe_quant_algo", "expected"),
+    [
+        ("auto", "NVFP4", False),
+        ("deep_gemm_mega_moe", "NVFP4", True),
+        ("deep_gemm_mega_moe", None, NotImplementedError),
+    ],
+)
+def test_mega_moe_is_gated_on_nvfp4_checkpoints(moe_backend, moe_quant_algo, expected):
+    """MegaMoE is opt-in and refuses checkpoints whose experts are not NVFP4."""
+    vllm_config = _vllm_config(moe_backend, moe_quant_algo)
+    if expected is NotImplementedError:
+        with pytest.raises(NotImplementedError, match="NVFP4 routed experts"):
+            _use_mega_moe(vllm_config)
+    else:
+        assert _use_mega_moe(vllm_config) is expected
+
+
+def test_mega_moe_expert_weight_follows_the_qat_recipe():
+    """Compressed-tensors quantization scales become ModelOpt dequantization
+    scales; the FC1 input scale is dropped (per-token) and the FC2 one is the
+    fixed QAT scale whatever the checkpoint holds."""
+    prefix = "layers.3.block_sparse_moe.experts.7"
+    packed = torch.randint(0, 255, (4, 8), dtype=torch.uint8)
+    block_scale = torch.rand(4, 1).to(torch.float8_e4m3fn)
+
+    assert _mega_moe_expert_weight(f"{prefix}.w1.weight_packed", packed) == (
+        f"{prefix}.w1.weight",
+        packed,
+    )
+    assert _mega_moe_expert_weight(f"{prefix}.w3.weight_scale", block_scale) == (
+        f"{prefix}.w3.weight_scale",
+        block_scale,
+    )
+    name, scale = _mega_moe_expert_weight(
+        f"{prefix}.w2.weight_global_scale", torch.tensor(448.0)
+    )
+    assert name == f"{prefix}.w2.weight_scale_2"
+    assert torch.equal(scale, torch.tensor(1 / 448, dtype=torch.float32))
+    for shard_id in ("w1", "w3"):
+        assert (
+            _mega_moe_expert_weight(
+                f"{prefix}.{shard_id}.input_global_scale", torch.tensor(1.0)
+            )
+            is None
+        )
+    name, scale = _mega_moe_expert_weight(
+        f"{prefix}.w2.input_global_scale", torch.tensor(1.0)
+    )
+    assert name == f"{prefix}.w2.input_scale"
+    assert scale.item() == MEGA_MOE_FC2_INPUT_SCALE
+
+
+def _load_checkpoint(moe, tensors: dict[str, torch.Tensor]) -> None:
+    """``MiniMaxM3Model.load_weights`` for one MegaMoE block's tensors."""
+    experts = moe.experts
+    mapping = _mega_moe_expert_params_mapping(experts.num_experts)
+    for name, tensor in tensors.items():
+        if ".shared_experts." in name:
+            name = name.replace(".gate_proj.", ".w1.").replace(".up_proj.", ".w3.")
+            experts.load_shared_expert_weight(name, tensor)
+            continue
+        converted = _mega_moe_expert_weight(name, tensor)
+        if converted is None:
+            continue
+        name, tensor = converted
+        param_name, weight_name, expert_id, shard_id = next(
+            entry for entry in mapping if entry[1] in name
+        )
+        name = name.replace(weight_name, param_name)
+        param = getattr(experts, name.rpartition(".")[2])
+        param.weight_loader(param, tensor, name, shard_id=shard_id, expert_id=expert_id)
+
+
+def _mega_moe_block(num_experts=8, top_k=2, hidden_size=512, intermediate_size=256):
+    """An SM100 MiniMax-M3.1 MoE block on MegaMoE with an MXFP8 shared expert."""
+    if not current_platform.is_device_capability_family(100):
+        pytest.skip("DeepGEMM NVFP4 MegaMoE requires SM100")
+    from vllm.config import VllmConfig, set_current_vllm_config
+    from vllm.config.kernel import KernelConfig
+    from vllm.model_executor.layers.quantization.modelopt import ModelOptMxFp8Config
+    from vllm.models.minimax_m3.nvidia.model import MiniMaxM3MoE
+
+    config = SimpleNamespace(
+        hidden_size=hidden_size,
+        intermediate_size=intermediate_size,
+        num_local_experts=num_experts,
+        num_experts_per_tok=top_k,
+        n_shared_experts=1,
+        scoring_func="sigmoid",
+        routed_scaling_factor=2.0,
+        use_routing_bias=False,
+        hidden_act="swigluoai",
+        swiglu_limit=7.0,
+        swiglu_alpha=1.702,
+        swiglu_beta=1.0,
+    )
+    quant_config = ModelOptMxFp8Config(
+        is_checkpoint_mxfp8_serialized=True,
+        kv_cache_quant_algo=None,
+        exclude_modules=[],
+        moe_quant_algo="NVFP4",
+    )
+    vllm_config = VllmConfig(
+        kernel_config=KernelConfig(moe_backend="deep_gemm_mega_moe")
+    )
+    vllm_config.scheduler_config.max_num_batched_tokens = 256
+    with set_current_vllm_config(vllm_config), torch.device("cuda"):
+        moe = MiniMaxM3MoE(
+            config,
+            layer_id=0,
+            quant_config=quant_config,
+            prefix="layers.0.block_sparse_moe",
+            use_mega_moe=True,
+        )
+    assert isinstance(moe.experts, DeepGemmMegaMoEExperts)
+    return moe, config
+
+
+@pytest.mark.parametrize("fused", [True, False])
+def test_mega_moe_allocates_the_shared_expert_once(dist_init, monkeypatch, fused):
+    """The block builds a real shared MLP only when MegaMoE does not fuse it
+    (e.g. an older DeepGEMM); a fused layer owns the only copy, and building
+    the block leaves no discarded weights for the garbage collector."""
+    if not fused:
+        monkeypatch.setattr(
+            DeepGemmMegaMoEExperts,
+            "_kernel_fuses_shared_experts",
+            lambda self, spec: False,
+        )
+    gc.collect()
+    gc.disable()
+    try:
+        moe, _ = _mega_moe_block()
+        allocated = torch.accelerator.memory_allocated()
+        gc.collect()
+    finally:
+        gc.enable()
+    assert torch.accelerator.memory_allocated() == allocated
+    assert moe.experts.has_fused_shared_experts is fused
+    assert (moe.shared_experts is None) is fused
+    assert all(param.is_cuda for param in moe.parameters())
+
+
+@pytest.mark.parametrize("num_tokens", [1, 45, 200])
+def test_mega_moe_block_matches_the_qat_reference(dist_init, num_tokens):
+    """A MiniMax-M3.1 MoE block on MegaMoE fuses its MXFP8 shared expert and,
+    loaded from a compressed-tensors checkpoint, reproduces the QAT recipe:
+    per-token FC1 input scales and the fixed FC2 input scale on the unweighted
+    SwiGLU output. Routing weights scale the BF16 expert outputs in combine."""
+    from vllm.third_party.deep_gemm.utils import (
+        cast_back_from_fp8,
+        per_token_cast_to_fp8,
+    )
+    from vllm.third_party.deep_gemm.utils.nvfp4 import (
+        cast_back_from_nvfp4,
+        per_token_cast_to_nvfp4,
+    )
+
+    torch.manual_seed(0)
+    moe, config = _mega_moe_block()
+    assert moe.experts.has_fused_shared_experts and moe.shared_experts is None
+    num_experts, top_k = config.num_local_experts, config.num_experts_per_tok
+    hidden_size, intermediate_size = config.hidden_size, config.intermediate_size
+    clamp, alpha, beta = config.swiglu_limit, config.swiglu_alpha, config.swiglu_beta
+
+    def nvfp4(n, k):
+        weight = torch.randn(n, k, device="cuda") * k**-0.5
+        global_scale = weight.abs().max().item() / (6 * 448)
+        packed, sf = per_token_cast_to_nvfp4(weight, global_scale)
+        dequant = cast_back_from_nvfp4(packed, sf, global_scale)
+        return (
+            packed.view(torch.uint8),
+            sf.view(torch.float8_e4m3fn),
+            torch.tensor(1 / global_scale),
+            dequant,
+        )
+
+    def mxfp8(n, k):
+        weight = torch.randn(n, k, device="cuda") * k**-0.5
+        fp8, sf = per_token_cast_to_fp8(weight, use_ue8m0=True, gran_k=32)
+        ue8m0 = (sf.view(torch.int32) >> 23).to(torch.uint8)
+        return fp8, ue8m0, cast_back_from_fp8(fp8, sf, gran_k=32)
+
+    prefix = "layers.0.block_sparse_moe"
+    checkpoint: dict[str, torch.Tensor] = {}
+    ref_w13, ref_w2 = [], []
+    for e in range(num_experts):
+        w13 = nvfp4(2 * intermediate_size, hidden_size)
+        w2 = nvfp4(hidden_size, intermediate_size)
+        for shard_id, rows, (packed, sf, global_scale, _) in (
+            ("w1", slice(0, intermediate_size), w13),
+            ("w3", slice(intermediate_size, None), w13),
+            ("w2", slice(None), w2),
+        ):
+            name = f"{prefix}.experts.{e}.{shard_id}"
+            checkpoint[f"{name}.weight_packed"] = packed[rows]
+            checkpoint[f"{name}.weight_scale"] = sf[rows]
+            checkpoint[f"{name}.weight_global_scale"] = global_scale
+            # A placeholder the block must ignore.
+            checkpoint[f"{name}.input_global_scale"] = torch.tensor(0.25)
+        ref_w13.append(w13[3])
+        ref_w2.append(w2[3])
+    shared_ref = {}
+    for proj, (n, k) in (
+        ("gate_proj", (intermediate_size, hidden_size)),
+        ("up_proj", (intermediate_size, hidden_size)),
+        ("down_proj", (hidden_size, intermediate_size)),
+    ):
+        fp8, ue8m0, shared_ref[proj] = mxfp8(n, k)
+        checkpoint[f"{prefix}.shared_experts.{proj}.weight"] = fp8
+        checkpoint[f"{prefix}.shared_experts.{proj}.weight_scale"] = ue8m0
+    _load_checkpoint(moe, checkpoint)
+
+    # Rows span six decades, which one tensor-wide input scale cannot cover.
+    x = (
+        torch.randn(num_tokens, hidden_size, device="cuda")
+        * torch.logspace(-4, 2, num_tokens, device="cuda")[:, None]
+    ).to(torch.bfloat16)
+    scores = torch.randn(num_tokens, num_experts, device="cuda").sigmoid()
+    topk_weights, topk_ids = scores.topk(top_k, dim=-1)
+    topk_weights *= config.routed_scaling_factor / topk_weights.sum(-1, keepdim=True)
+    out = moe.experts(x, topk_weights, topk_ids)
+
+    def activation(h):
+        gate, up = h.to(torch.bfloat16).float().chunk(2, dim=-1)
+        gate = gate.clamp(max=clamp)
+        up = up.clamp(-clamp, clamp)
+        return gate * torch.sigmoid(alpha * gate) * (up + beta)
+
+    def fp8_round_trip(h):
+        return cast_back_from_fp8(
+            *per_token_cast_to_fp8(h, use_ue8m0=True, gran_k=32, use_packed_ue8m0=True),
+            gran_k=32,
+            use_packed_ue8m0=True,
+        )
+
+    x_scales = x.float().abs().amax(-1, keepdim=True).clamp_min(1e-12) / 2688
+    x_ref = cast_back_from_nvfp4(*per_token_cast_to_nvfp4(x.float() / x_scales))
+    x_ref *= x_scales
+    ref = torch.zeros(num_tokens, hidden_size, device="cuda")
+    for e in range(num_experts):
+        rows, slots = torch.where(topk_ids == e)
+        if rows.numel() == 0:
+            continue
+        mid = activation(x_ref[rows] @ ref_w13[e].T)
+        mid = cast_back_from_nvfp4(
+            *per_token_cast_to_nvfp4(mid, MEGA_MOE_FC2_INPUT_SCALE),
+            MEGA_MOE_FC2_INPUT_SCALE,
+        )
+        expert_out = (mid @ ref_w2[e].T).to(torch.bfloat16).float()
+        ref.index_add_(0, rows, topk_weights[rows, slots, None] * expert_out)
+    shared_w13 = torch.cat([shared_ref["gate_proj"], shared_ref["up_proj"]])
+    shared_mid = fp8_round_trip(activation(fp8_round_trip(x) @ shared_w13.T))
+    ref += (shared_mid @ shared_ref["down_proj"].T).to(torch.bfloat16).float()
+
+    assert torch.isfinite(out).all()
+    # Per row: a missing per-token scale would only show in the small rows.
+    row_error = (out.float() - ref).norm(dim=-1) / ref.norm(dim=-1)
+    assert row_error.max() < 0.02, row_error
diff --git a/tests/models/test_muse_spark_logits_local.py b/tests/models/test_muse_spark_logits_local.py
index bb2ff5772e72..c9f2a8bc76ec 100644
--- a/tests/models/test_muse_spark_logits_local.py
+++ b/tests/models/test_muse_spark_logits_local.py
@@ -10,6 +10,8 @@
 def test_compute_logits_local_applies_output_transform() -> None:
     model = MuseSpark1xForCausalLM.__new__(MuseSpark1xForCausalLM)
     nn.Module.__init__(model)
+    model.model = nn.Module()
+    model.model.mono = None
     model.lm_head = nn.Linear(4, 8, bias=False, dtype=torch.bfloat16)
     model.logits_processor = lambda lm_head, h, skip_gather=False: lm_head(h)
     model.metap_output_multiplier = 0.5
diff --git a/tests/models/test_muse_spark_rocm_mono.py b/tests/models/test_muse_spark_rocm_mono.py
new file mode 100644
index 000000000000..4367b22b8169
--- /dev/null
+++ b/tests/models/test_muse_spark_rocm_mono.py
@@ -0,0 +1,882 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Host side of the Muse Spark monokernel (``--monokernel``).
+
+The guards run on real vLLM configs and on a tiny real MuseSpark1x model
+(Quark MXFP4 experts, ROCm AITER FA, a bf16 LBHNC KV cache). The FlyDSL
+launches build only for the TEP8 checkpoint shapes, so the step tests run
+them as the stock layer pieces they replace: from layer i's attention output
+to layer i+1's q, o-gate and KV insert, reading and writing the device
+memory at the pointers the host passes. A mono step must then match the
+stock layer loop: hidden states, aux hidden states and KV caches.
+"""
+
+import dataclasses
+import inspect
+import json
+from typing import Any
+
+import pytest
+import torch
+
+from vllm.config import (
+    AttentionConfig,
+    CacheConfig,
+    CUDAGraphMode,
+    LoRAConfig,
+    ModelConfig,
+    ParallelConfig,
+    SpeculativeConfig,
+    VllmConfig,
+    WeightTransferConfig,
+    set_current_vllm_config,
+)
+from vllm.model_executor.layers.attention.attention import get_attention_context
+from vllm.model_executor.layers.fused_moe.config import FusedMoEParallelConfig
+from vllm.models.muse_spark.amd.model import norm_residual_then_norm
+from vllm.models.muse_spark.amd.mono.decoder import (
+    FIRST_LAYER,
+    LAST_LAYER,
+    MonoDims,
+    MonoKernels,
+    MuseSparkMono,
+    bucket_of,
+    built_buckets,
+    collect_layer_weights,
+    config_refusals,
+    deployment_refusals,
+    fused_attention_layers,
+    kv_cache_refusal,
+    layer_refusals,
+)
+from vllm.platforms import current_platform
+from vllm.transformers_utils.configs.muse_spark import MuseSparkTextConfig
+
+# The GPU tests share one tiny model and its single-rank parallel state.
+pytestmark = pytest.mark.skip_global_cleanup
+
+NUM_LAYERS = 62
+ROWS_PER_REQUEST = 8
+
+# The text config of Muse-Spark-Open-816B-A42B (the shape the kernels build).
+CHECKPOINT_TEXT: dict[str, Any] = {
+    "model_type": "muse_spark_text",
+    "attention_bias": False,
+    "bos_token_id": 0,
+    "eos_token_id": 1,
+    "pad_token_id": 2,
+    "final_logit_softcapping": 20.0,
+    "head_dim": 64,
+    "hidden_act": "silu",
+    "hidden_size": 8192,
+    "intermediate_size": 16384,
+    "max_position_embeddings": 1048576,
+    "moe_hidden_size": 4096,
+    "moe_intermediate_size": 4096,
+    "num_attention_heads": 128,
+    "num_experts_per_tok": 8,
+    "num_hidden_layers": NUM_LAYERS,
+    "num_key_value_heads": 16,
+    "num_local_experts": 256,
+    "num_shared_experts": 0,
+    "output_multiplier": 0.17677669529663687,
+    "post_norm_eps": 1e-08,
+    "qk_scale_factor": 5.4730064863875,
+    "residual_gate_temperature": 0.3,
+    "rms_norm_eps": 1e-05,
+    "rope_parameters": {"rope_theta": 500000.0, "rope_type": "default"},
+    "routed_scaling_factor": 1.0,
+    "sliding_window": 2048,
+    "tie_word_embeddings": False,
+    "vocab_size": 202048,
+    # Full attention (and no RoPE) every 4th layer from layer 1.
+    "layer_types": [
+        "full_attention" if i % 4 == 1 else "sliding_attention"
+        for i in range(NUM_LAYERS)
+    ],
+    "layer_rope_theta": [0 if i % 4 == 1 else 500000.0 for i in range(NUM_LAYERS)],
+    "mlp_layer_types": ["sparse"] * NUM_LAYERS,
+}
+# One TP rank of the checkpoint's layout at TP = EP = 1.
+TINY_TEXT = CHECKPOINT_TEXT | {
+    "hidden_size": 512,
+    "intermediate_size": 512,
+    "moe_hidden_size": 256,
+    "moe_intermediate_size": 256,
+    "num_attention_heads": 16,
+    "num_key_value_heads": 2,
+    "num_local_experts": 8,
+    "vocab_size": 512,
+    "max_position_embeddings": 4096,
+}
+_MX_FP4 = {
+    "dtype": "fp4",
+    "qscheme": "per_group",
+    "ch_axis": -1,
+    "group_size": 32,
+    "round_method": "half_even",
+    "scale_type": "float",
+    "scale_format": "e8m0",
+    "scale_calculation_mode": "even",
+    "observer_cls": "PerBlockMXObserver",
+    "is_scale_quant": False,
+}
+
+
+def _quark_exclude(num_experts: int) -> list[str]:
+    """The checkpoint's exclusions: everything but the experts of layers 1-60."""
+    names = ["lm_head"]
+    for i in range(NUM_LAYERS):
+        layer = f"model.language_model.layers.{i}"
+        names += [f"{layer}.self_attn.{p}_proj" for p in ("q", "k", "v", "o", "gate")]
+        names += [f"{layer}.mlp.{p}" for p in ("gate", "pre_expert_proj")]
+        names.append(f"{layer}.mlp.post_expert_proj")
+        if i in (0, NUM_LAYERS - 1):
+            names += [
+                f"{layer}.mlp.experts.{e}.{p}_proj"
+                for e in range(num_experts)
+                for p in ("gate", "up", "down")
+            ]
+    return names
+
+
+def _checkpoint_config(text: dict[str, Any]) -> dict[str, Any]:
+    return {
+        "architectures": ["MuseSparkForConditionalGeneration"],
+        "model_type": "muse_spark",
+        "dtype": "bfloat16",
+        "image_token_id": 500,
+        "audio_token_id": 501,
+        "video_token_id": 502,
+        "projector_hidden_act": "gelu",
+        "projector_hidden_size": 64,
+        "out_hidden_size": 256,
+        "text_config": text,
+        "vision_config": {
+            "model_type": "muse_spark_vision",
+            "hidden_act": "gelu",
+            "hidden_size": 64,
+            "intermediate_size": 128,
+            "layer_norm_eps": 1e-05,
+            "layer_types": ["full_attention"],
+            "max_position_embeddings": 64,
+            "merge_size": 2,
+            "num_attention_heads": 2,
+            "num_hidden_layers": 1,
+            "patch_size": 14,
+            "patch_temporal": 2,
+            "pos_emb_height": 4,
+            "pos_emb_width": 4,
+            "rope_parameters": {"rope_theta": 10000, "rope_type": "axial"},
+        },
+        "audio_config": {
+            "model_type": "muse_spark_audio",
+            "conv_kernel_size": 7,
+            "depths": [1],
+            "hidden_act": "gelu",
+            "hidden_sizes": [64],
+            "mlp_ratio": 4.0,
+            "num_hidden_layers": 1,
+            "out_hidden_size": text["hidden_size"],
+            "patch_sizes": [16],
+            "rms_norm_eps": 1e-06,
+        },
+        "quantization_config": {
+            "quant_method": "quark",
+            "global_quant_config": {
+                "input_tensors": _MX_FP4 | {"is_dynamic": True},
+                "weight": _MX_FP4 | {"is_dynamic": False},
+                "output_tensors": None,
+                "bias": None,
+            },
+            "exclude": _quark_exclude(text["num_local_experts"]),
+            "layer_quant_config": {},
+            "layer_type_quant_config": {},
+            "kv_cache_quant_config": {},
+            "export": {"kv_cache_group": [], "pack_method": "reorder"},
+        },
+    }
+
+
+# The target_layer_ids of Muse-Spark-Open-DSpark.
+DSPARK_TARGETS = [13, 29, 45, 57, 61]
+DSPARK = {
+    "model_type": "muse_spark_dspark",
+    "num_hidden_layers": 2,
+    "moe": {"is_enabled": False},
+    "global_attn_cfg": [2048],
+    "layer_types": ["sliding_attention"] * 2,
+    "target_layer_ids": DSPARK_TARGETS,
+    "num_target_layers": len(DSPARK_TARGETS),
+    "markov_rank": 8,
+    "markov_head_type": "vanilla",
+    "dflash_config": {"causal": True, "target_layer_ids": DSPARK_TARGETS},
+}
+
+
+def _vllm_config(
+    target, draft=None, tp: int = 1, cache_dtype: str = "auto", **parallel
+) -> VllmConfig:
+    model_config = ModelConfig(
+        model=str(target),
+        tokenizer_mode="skip",
+        dtype="bfloat16",
+        max_model_len=4096,
+        language_model_only=True,
+        enforce_eager=True,
+    )
+    # Config-only: "mp" skips the executor's local GPU count check.
+    parallel_config = ParallelConfig(
+        tensor_parallel_size=tp,
+        enable_expert_parallel=True,
+        distributed_executor_backend="mp",
+        **parallel,
+    )
+    speculative_config = None
+    if draft is not None:
+        speculative_config = SpeculativeConfig(
+            model=str(draft),
+            method="dspark",
+            num_speculative_tokens=ROWS_PER_REQUEST - 1,
+            target_model_config=model_config,
+            target_parallel_config=parallel_config,
+        )
+    return VllmConfig(
+        model_config=model_config,
+        parallel_config=parallel_config,
+        speculative_config=speculative_config,
+        cache_config=CacheConfig(block_size=16, cache_dtype=cache_dtype),
+        attention_config=AttentionConfig(backend="ROCM_AITER_FA"),
+    )
+
+
+@pytest.fixture(scope="module")
+def checkpoints(tmp_path_factory):
+    root = tmp_path_factory.mktemp("muse_spark")
+    for name, config in (("tiny", _checkpoint_config(TINY_TEXT)), ("dspark", DSPARK)):
+        (root / name).mkdir()
+        (root / name / "config.json").write_text(json.dumps(config))
+    return root
+
+
+# ---------------------------------------------------------------- host guards
+
+
+def test_buckets_pad_rows_up_to_the_next_build() -> None:
+    assert [bucket_of(r) for r in (8, 16, 24, 32, 40, 64, 72, 96, 104, 128)] == [
+        8, 16, 32, 32, 64, 64, 96, 96, 128, 128,
+    ]  # fmt: skip
+    assert [bucket_of(r) for r in (0, 4, 12, 136)] == [None] * 4
+    assert built_buckets(1) == [8]
+    assert built_buckets(3) == [8, 16, 32]
+    assert built_buckets(16) == built_buckets(128) == [8, 16, 32, 64, 96, 128]
+
+
+def test_kv_cache_guard_accepts_only_dense_bf16_lbhnc() -> None:
+    lbhnc = torch.empty(10, 2, 16, 128, dtype=torch.bfloat16, device="meta")
+    assert kv_cache_refusal(lbhnc, 2) is None
+    lhbnc = torch.empty(2, 10, 16, 128, dtype=torch.bfloat16, device="meta")
+    assert kv_cache_refusal(lhbnc.transpose(0, 1), 2).startswith("kv:stride")
+    fp8 = torch.empty(10, 2, 16, 128, dtype=torch.float8_e4m3fn, device="meta")
+    assert kv_cache_refusal(fp8, 2).startswith("kv:torch.float8")
+    block32 = torch.empty(10, 2, 32, 128, dtype=torch.bfloat16, device="meta")
+    assert kv_cache_refusal(block32, 2).startswith("kv:shape")
+    assert kv_cache_refusal(torch.empty(0, dtype=torch.bfloat16), 2) == "kv:unbound"
+
+
+def test_config_guard_accepts_only_the_checkpoint_shape() -> None:
+    assert config_refusals(MuseSparkTextConfig(**CHECKPOINT_TEXT)) == []
+    assert config_refusals(MuseSparkTextConfig(**TINY_TEXT)) == [
+        "config.hidden_size=512 (kernel: 8192)",
+        "config.moe_hidden_size=256 (kernel: 4096)",
+        "config.moe_intermediate_size=256 (kernel: 4096)",
+        "config.num_local_experts=8 (kernel: 256)",
+        "config.num_attention_heads=16 (kernel: 128)",
+        "config.num_key_value_heads=2 (kernel: 16)",
+    ]
+
+
+def _moe_parallel(use_ep: bool = True, ep_size: int = 8, ep_rank: int = 3):
+    return FusedMoEParallelConfig(
+        tp_size=1 if use_ep else ep_size,
+        pcp_size=1,
+        dp_size=1,
+        ep_size=ep_size if use_ep else 1,
+        tp_rank=0,
+        pcp_rank=0,
+        dp_rank=0,
+        ep_rank=ep_rank if use_ep else 0,
+        sp_size=1,
+        use_ep=use_ep,
+        all2all_backend="allgather_reducescatter",
+        enable_eplb=False,
+    )
+
+
+def test_deployment_guard_accepts_only_tep8_with_dspark(checkpoints) -> None:
+    tiny, dspark = checkpoints / "tiny", checkpoints / "dspark"
+    tep8 = _vllm_config(tiny, dspark, tp=8)
+    assert deployment_refusals(tep8, 3, 8, _moe_parallel()) == []
+
+    tp4 = _vllm_config(tiny, dspark, tp=4, expert_placement_strategy="round_robin")
+    tp4.parallel_config.enable_expert_parallel = False
+    assert deployment_refusals(tp4, 3, 4, _moe_parallel(use_ep=False)) == [
+        "tensor_parallel_size=4 (kernel: 8)",
+        "expert parallelism off (needs --enable-expert-parallel)",
+        "expert placement round_robin",
+        "MoE layers not expert parallel",
+    ]
+
+    no_draft = _vllm_config(tiny, tp=8, cache_dtype="fp8")
+    no_draft.lora_config = LoRAConfig()
+    no_draft.model_config.enable_sleep_mode = True
+    no_draft.weight_transfer_config = WeightTransferConfig()
+    assert deployment_refusals(no_draft, 3, 8, _moe_parallel()) == [
+        "needs speculative decoding with 7 draft tokens",
+        "LoRA on",
+        "sleep mode on",
+        "weight transfer on",
+        "kv_cache_dtype=fp8 (kernel: bf16)",
+    ]
+
+
+# ------------------------------------------------- the tiny model on one GPU
+
+
+def _on_gfx950() -> bool:
+    if not (current_platform.is_rocm() and torch.cuda.is_available()):
+        return False
+    from vllm.platforms.rocm import on_gfx950
+
+    return on_gfx950()
+
+
+requires_gfx950 = pytest.mark.skipif(
+    not _on_gfx950(), reason="MuseSpark1x on ROCm needs AITER's gfx950 kernels"
+)
+
+
+@dataclasses.dataclass
+class TinyModel:
+    lm: Any
+    vllm_config: VllmConfig
+    dims: MonoDims
+    refusals_before_finalize: list[str]
+    kv_specs: dict[str, Any]
+
+    @property
+    def model(self):
+        return self.lm.model
+
+    @property
+    def layers(self):
+        return self.lm.model.layers
+
+
+@pytest.fixture(scope="module")
+def tiny_model(checkpoints):
+    from vllm.config.load import LoadConfig
+    from vllm.distributed import (
+        cleanup_dist_env_and_memory,
+        ensure_model_parallel_initialized,
+        init_distributed_environment,
+    )
+    from vllm.model_executor.model_loader.dummy_loader import DummyModelLoader
+    from vllm.model_executor.model_loader.utils import process_weights_after_loading
+    from vllm.models.muse_spark.amd.multimodal import MuseSparkForConditionalGeneration
+    from vllm.utils.torch_utils import set_default_torch_dtype
+    from vllm.v1.kv_cache_interface import FullAttentionSpec, SlidingWindowSpec
+    from vllm.v1.worker.workspace import init_workspace_manager
+
+    with pytest.MonkeyPatch.context() as env:
+        env.setenv("VLLM_ROCM_USE_AITER", "1")
+        env.setenv("VLLM_KV_CACHE_LAYOUT", "LBHNC")
+        from vllm._aiter_ops import rocm_aiter_ops
+
+        rocm_aiter_ops.refresh_env_variables()
+        device = torch.device("cuda", 0)
+        torch.accelerator.set_device_index(device)
+        vllm_config = _vllm_config(checkpoints / "tiny", checkpoints / "dspark")
+        store = checkpoints / "dist_store"
+        with set_current_vllm_config(vllm_config):
+            init_distributed_environment(
+                world_size=1,
+                rank=0,
+                local_rank=0,
+                distributed_init_method=f"file://{store}",
+                backend="nccl",
+            )
+            ensure_model_parallel_initialized(1, 1)
+            init_workspace_manager(device)
+            with set_default_torch_dtype(torch.bfloat16), device:
+                lm = MuseSparkForConditionalGeneration(vllm_config=vllm_config)
+            model_config = vllm_config.model_config
+            text_config = model_config.hf_text_config
+            dims = MonoDims.from_config(text_config, 1)
+            before = layer_refusals(lm.model, dims)
+            DummyModelLoader(LoadConfig(load_format="dummy")).load_weights(
+                lm, model_config
+            )
+            lm.load_weights([])
+            process_weights_after_loading(lm, model_config, device)
+            kv = dict(
+                block_size=16,
+                num_kv_heads=dims.local_kv_heads,
+                head_size=dims.head_dim,
+                dtype=torch.bfloat16,
+            )
+            specs = {
+                "full": FullAttentionSpec(**kv),
+                "swa": SlidingWindowSpec(**kv, sliding_window=2048),
+            }
+            yield TinyModel(lm, vllm_config, dims, before, specs)
+        cleanup_dist_env_and_memory()
+
+
+@requires_gfx950
+def test_layer_guards_follow_the_weight_finalize(tiny_model) -> None:
+    """Refused until load_weights and the quant finalize produced the kernel
+    layouts (folded router, materialized gates, shuffled experts)."""
+    before = tiny_model.refusals_before_finalize
+    for reason in (
+        "layer 1: add_attn coefficients not materialized",
+        "layer 1: router multiplier not folded",
+        "layer 1: expert weights not shuffled",
+    ):
+        assert reason in before
+    assert layer_refusals(tiny_model.model, tiny_model.dims) == []
+
+
+@requires_gfx950
+def test_weight_views_alias_the_rank_slices(tiny_model) -> None:
+    """The launches read vLLM's parameters in place; only the norm
+    multipliers are copies. Rank 1 of 2 reads the second row half."""
+    dims = dataclasses.replace(tiny_model.dims, world_size=2)
+    layer, nxt = tiny_model.layers[1], tiny_model.layers[2]
+    w = collect_layer_weights(tiny_model.layers, 1, 1, dims)
+    moe = layer.feed_forward
+    experts = moe.experts.routed_experts
+
+    assert w.w_r.shape == (dims.num_experts // 2, dims.hidden_size)
+    for view, base, row0 in (
+        (w.w_r, moe.router.weight, dims.num_experts // 2),
+        (w.w_pre, moe.pooler.weight, dims.pooled_size // 2),
+        (w.w_pp, moe.unpooler.weight, dims.hidden_size // 2),
+        (w.w_o, layer.self_attn.o_proj.weight, 0),
+        (w.w13, experts.w13_weight, 0),
+        (w.w13s, experts.w13_weight_scale, 0),
+        (w.w2, experts.w2_weight, 0),
+        (w.w2s, experts.w2_weight_scale, 0),
+        (w.bias, moe.expert_bias, 0),
+        (w.alpha_a, layer.add_attn.alpha, 0),
+        (w.beta_a, layer.add_attn.beta, 0),
+        (w.alpha, layer.add_ffn.alpha, 0),
+        (w.beta, layer.add_ffn.beta, 0),
+        (w.w_qkvg, nxt.self_attn.qkv_proj.weight, 0),
+    ):
+        assert view.is_contiguous()
+        offset = row0 * base.stride(0) * base.element_size()
+        assert view.data_ptr() == base.data_ptr() + offset
+    for copy, norm in (
+        (w.w_pffn, layer.pre_ffn_norm),
+        (w.w_pren, moe.pre_expert_norm),
+        (w.w_post, moe.post_expert_norm),
+        (w.w_pff, layer.post_ffn_norm),
+        (w.w_in, nxt.self_attn.pre_qkv_proj_norm),
+    ):
+        assert torch.equal(copy, (norm.weight.float() + norm.gain_center).bfloat16())
+    # Layer 2 is a sliding-window layer (RoPE), layer 5 a full one (none).
+    assert w.next_rope
+    assert not collect_layer_weights(tiny_model.layers, 4, 0, dims).next_rope
+
+
+@requires_gfx950
+def test_fused_attention_layers_follow_the_window(tiny_model) -> None:
+    layers = tiny_model.layers
+    fused = fused_attention_layers(layers, 4096, 1 << 20)
+    assert sorted(fused) == list(range(FIRST_LAYER + 1, LAST_LAYER + 1))
+    assert [i for i, t in fused.items() if t == "full"] == list(range(5, 61, 4))
+    assert fused_attention_layers(layers, (1 << 20) + 1, 1 << 20) == {}
+
+
+# ------------------------------------------- the launches, as the stock pieces
+
+# The launch signatures of the FlyDSL sources, up to the peer pointers.
+_K3_HEAD = (
+    "attn_o attn_g stream_pre w_o alpha_a beta_a w_pffn w_r bias w_pre w_pren "
+    "w13 w13s w2 w2s w_post w_pp w_pff alpha beta w_in w_qkvg stream_out xn_out "
+    "part_out q_out g_out k_cache v_cache slots cos_sin pos scratch epoch health tl"
+)
+_HA_HEAD = (
+    "attn_o attn_g stream_in w_o alpha_a beta_a w_pffn w_r bias w_pre w_pren "
+    "h_out ids_out sc_out scratch epoch health tl"
+)
+_HM_HEAD = "h ids sc w13 w13s w2 w2s w_post part_out scratch epoch health tl"
+_HB_HEAD = (
+    "part w_pp w_pff alpha beta w_in w_qkvg stream_out sfull_out q_out g_out "
+    "k_cache v_cache slots cos_sin pos scratch epoch health tl"
+)
+_REPLICAS = " w13r w13sr w2r w2sr eptab"
+_DESC_WORDS = 10
+
+
+def _params(head: str, world_size: int, tail: str) -> list[str]:
+    """A launch's parameters in order."""
+    return head.split() + [f"peer{r}" for r in range(world_size)] + tail.split()
+
+
+def _device_view(ptr: int, shape, dtype: torch.dtype, strides=None) -> torch.Tensor:
+    """The device memory at ``ptr`` as a tensor, as a launch reads it."""
+    strides = strides or torch.empty(shape, device="meta").stride()
+    extent = 1 + sum((n - 1) * s for n, s in zip(shape, strides))
+    nbytes = extent * torch.empty((), dtype=dtype).element_size()
+
+    class _Memory:
+        __cuda_array_interface__ = {
+            "shape": (nbytes,),
+            "typestr": "|u1",
+            "data": (ptr, False),
+            "version": 3,
+        }
+
+    raw = torch.as_tensor(_Memory(), device="cuda")
+    return raw.view(dtype).as_strided(shape, strides)
+
+
+class StockLaunches:
+    """The FlyDSL monokernel sources, with each launch running the stock layer
+    pieces it replaces on the memory at its arguments' pointers."""
+
+    def __init__(self, model, dims: MonoDims, bnd: int) -> None:
+        self.model, self.dims, self.bnd = model, dims, bnd
+        self.calls: list[tuple[str, dict]] = []
+        self.pending: dict[int, dict] = {}
+        self.kernels = MonoKernels(
+            *(_Kernel(self, name) for name in ("k3", "k5", "k5m")), _AttnFrag()
+        )
+
+    def launcher(self, name: str, params: list[str], run):
+        signature = inspect.Signature(
+            [inspect.Parameter(p, inspect.Parameter.POSITIONAL_ONLY) for p in params]
+            + [inspect.Parameter("stream", inspect.Parameter.KEYWORD_ONLY)]
+        )
+
+        def launch(*args, **kwargs):
+            bound = dict(signature.bind(*args, **kwargs).arguments)
+            self.calls.append((name, bound))
+            run(bound)
+
+        return launch
+
+    def _attention(
+        self, i: int, attn_o: int, rows: int, nreq: int, fused: bool
+    ) -> torch.Tensor:
+        """Layer i's attention output: at ``attn_o``, or, for a fused build,
+        run here from the descriptor row at ``attn_o``."""
+        layer = self.model.layers[i]
+        d = self.dims
+        if not fused:
+            return _device_view(attn_o, (rows, d.q_rows), torch.bfloat16)
+        words = _device_view(attn_o, (_DESC_WORDS,), torch.int64).tolist()
+        md, _, kv_cache, _ = _attention_context(layer)
+        assert words[1:5] == [
+            kv_cache.data_ptr(),
+            kv_cache.stride(0),
+            md.block_table.data_ptr(),
+            md.seq_lens.data_ptr(),
+        ]
+        assert words[8] == nreq
+        q = _device_view(words[0], (rows, d.q_rows), torch.bfloat16)
+        out = _device_view(words[5], (rows, d.q_rows), torch.bfloat16)
+        out.copy_(layer.self_attn.attn(q, None, None))
+        return out
+
+    def run_layer(self, rows: int, a: dict, b: dict) -> None:
+        """Layer i from its attention output (launch ``a``) to layer i+1's q,
+        o-gate and stream (launch ``b``)."""
+        i, nreq = a["layer"], rows // ROWS_PER_REQUEST
+        d, layer, nxt = self.dims, self.model.layers[i], self.model.layers[i + 1]
+        o = self._attention(i, a["attn_o"], rows, nreq, a["fused"])
+        g = _device_view(
+            a["attn_g"], (rows, d.q_rows), torch.bfloat16, (a["g_stride"], 1)
+        )
+        stream = _device_view(a["stream"], (rows, d.hidden_size), torch.float32)
+        positions = _device_view(b["pos"], (rows,), torch.int64)
+        _, _, kv_cache, slots = _attention_context(nxt)
+        assert b["k_cache"] == kv_cache.data_ptr()
+        assert b["v_cache"] == kv_cache.stride(0)
+        assert b["slots"] == slots.data_ptr()
+        rope = nxt.self_attn.rotary_emb
+        if rope is not None:
+            assert b["cos_sin"] == rope.cos_sin_cache.data_ptr()
+
+        y = layer.after_attention(stream, layer.self_attn.back(o, g))
+        attn_input, stream_next = norm_residual_then_norm(
+            y, nxt.self_attn.pre_qkv_proj_norm, out_dtype=nxt.default_dtype
+        )
+        qkv, _ = nxt.self_attn.qkv_proj(attn_input)
+        q, _, _, g_next = nxt.self_attn._qk_norm_rope_kv_insert(qkv, positions)
+        streams = ["stream_out"] + (["sfull_out"] if b.get("gather") else [])
+        for names, t in ((streams, stream_next), (["q_out"], q), (["g_out"], g_next)):
+            for name in names:
+                _device_view(b[name], t.shape, t.dtype).copy_(t)
+
+
+@dataclasses.dataclass(frozen=True)
+class _Key:
+    S: int
+    W: int
+    next_rope: bool = True
+    attn: str = ""
+    bnd: int = 0
+
+
+class _Kernel:
+    """One FlyDSL source module: its layouts and builds."""
+
+    BLOCKS = 4
+    HEALTH_WORDS = 8 + 8 * BLOCKS
+    TL_POINTS = 4
+
+    def __init__(self, owner: StockLaunches, name: str) -> None:
+        self.owner, self.name = owner, name
+
+    def peer_layout(self, S, W, bnd=0):
+        return {}, 4096
+
+    def scratch_layout(self, S, W=1, bnd=0):
+        return {}, 4096
+
+    def build_vllm(self, S, W, next_rope=True, attn=""):
+        return _Key(S, W, next_rope, attn, self.owner.bnd if self.name == "k5" else 0)
+
+    def build(self, key: _Key):
+        o, w = self.owner, key.W
+        if self.name == "k3":
+            params = _params(_K3_HEAD, w, "layer rank g_stride" + _REPLICAS)
+            return o.launcher("k3", params, lambda a: _run_k3(o, key, a))
+        if self.name == "k5m":
+            params = _params(_HM_HEAD, w, "layer rank nrows" + _REPLICAS)
+            return o.launcher("hm", params, lambda a: None)
+        ha = _params(_HA_HEAD, w, "layer rank g_stride s_full nrows")
+        hb = _params(_HB_HEAD, w, "layer rank so_stride so_off gather nrows")
+
+        def run_ha(a):
+            o.pending[a["layer"]] = a | {"fused": bool(key.attn)}
+
+        return (
+            o.launcher("ha", ha, run_ha),
+            o.launcher("hb", hb, lambda a: _run_hb(o, key, a)),
+        )
+
+
+class _AttnFrag:
+    TABLE_ROWS = 64
+    DESC_WORDS = _DESC_WORDS
+    VLLM_MAX_KEYS = 1 << 20
+
+    @staticmethod
+    def part_bytes() -> int:
+        return 4096
+
+    @staticmethod
+    def vllm_desc(q, kv_cache, block_table, seq_lens, out, part, nreq, scale):
+        return [
+            q.data_ptr(),
+            kv_cache.data_ptr(),
+            kv_cache.stride(0),
+            block_table.data_ptr(),
+            seq_lens.data_ptr(),
+            out.data_ptr(),
+            part.data_ptr(),
+            block_table.stride(0),
+            nreq,
+            0,
+        ]
+
+
+def _run_k3(o: StockLaunches, key: _Key, a: dict) -> None:
+    o.run_layer(key.S, a | {"stream": a["stream_pre"], "fused": bool(key.attn)}, a)
+
+
+def _run_hb(o: StockLaunches, key: _Key, b: dict) -> None:
+    a = o.pending.pop(b["layer"])
+    assert a["nrows"] == b["nrows"]
+    o.run_layer(b["nrows"], a | {"stream": a["stream_in"]}, b)
+
+
+def _attention_context(layer):
+    return get_attention_context(layer.self_attn.attn.layer_name)
+
+
+class Step:
+    """One verify step's real attention metadata and KV caches."""
+
+    def __init__(self, tiny: TinyModel, query_lens: list[int], seed: int) -> None:
+        from tests.v1.attention.utils import (
+            BatchSpec,
+            create_common_attn_metadata,
+            dense_kv_cache_views,
+        )
+        from vllm.v1.attention.backends.rocm_aiter_fa import AiterFlashAttentionBackend
+        from vllm.v1.kv_cache_layout import KVCacheLayout
+
+        device = torch.device("cuda", 0)
+        generator = torch.Generator().manual_seed(seed)
+        seq_lens = [
+            int(torch.randint(40, 300, (), generator=generator)) + q for q in query_lens
+        ]
+        self.common = create_common_attn_metadata(
+            BatchSpec(seq_lens=seq_lens, query_lens=query_lens),
+            16,
+            device,
+            arange_block_indices=True,
+        )
+        num_blocks = int(self.common.block_table_tensor.max()) + 1
+        builder_cls = AiterFlashAttentionBackend.get_builder_cls()
+        self.names = [layer.self_attn.attn.layer_name for layer in tiny.layers]
+        self.metadata: dict[str, Any] = {}
+        self.kv_caches: dict[str, torch.Tensor] = {}
+        for kind, spec in tiny.kv_specs.items():
+            names = [
+                n
+                for n, layer in zip(self.names, tiny.layers)
+                if (layer.self_attn.attn.impl.sliding_window == (2047, 0))
+                == (kind == "swa")
+            ]
+            builder = builder_cls(spec, names, tiny.vllm_config, device)
+            metadata = builder.build(0, self.common)
+            for name in names:
+                raw = torch.zeros(
+                    num_blocks * spec.page_size_bytes, dtype=torch.int8, device=device
+                )
+                (view,) = dense_kv_cache_views(
+                    raw, spec, num_blocks, 1, KVCacheLayout.LBHNC
+                )
+                self.metadata[name] = metadata
+                self.kv_caches[name] = view
+        self.tokens = sum(query_lens)
+        self.input_ids = torch.randint(
+            0, TINY_TEXT["vocab_size"], (self.tokens,), generator=generator
+        ).to(device)
+        self.positions = torch.cat(
+            [torch.arange(s - q, s) for s, q in zip(seq_lens, query_lens)]
+        ).to(device)
+
+    def run(self, tiny: TinyModel, mode=CUDAGraphMode.NONE):
+        """MuseSpark1xModel.forward on zeroed KV caches; returns the hidden
+        states, the aux hidden states and the KV caches after the step."""
+        from vllm.forward_context import set_forward_context
+        from vllm.v1.worker.utils import bind_kv_cache
+
+        for cache in self.kv_caches.values():
+            cache.zero_()
+        context = tiny.vllm_config.compilation_config.static_forward_context
+        bind_kv_cache(self.kv_caches, context, [])
+        slots = dict.fromkeys(self.names, self.common.slot_mapping)
+        with (
+            set_current_vllm_config(tiny.vllm_config),
+            set_forward_context(
+                self.metadata,
+                tiny.vllm_config,
+                num_tokens=self.tokens,
+                cudagraph_runtime_mode=mode,
+                slot_mapping=slots,
+            ),
+        ):
+            out = tiny.model(self.input_ids, self.positions, None)
+        hidden, aux = out if isinstance(out, tuple) else (out, [])
+        kv = {name: cache.clone() for name, cache in self.kv_caches.items()}
+        return hidden.clone(), [a.clone() for a in aux], kv
+
+
+def _run_mono(
+    tiny: TinyModel, step: Step, bnd: int, buckets: list[int], mode=CUDAGraphMode.NONE
+):
+    """The step with the monokernel attached, on the stock launches."""
+    from vllm.distributed import get_tp_group
+
+    launches = StockLaunches(tiny.model, tiny.dims, bnd)
+    tiny.model.mono = MuseSparkMono(
+        tiny.model,
+        launches.kernels,
+        rank=0,
+        group=get_tp_group().cpu_group,
+        device=torch.device("cuda", 0),
+        dims=tiny.dims,
+        buckets=buckets,
+        fused_attention=fused_attention_layers(
+            tiny.layers, 4096, _AttnFrag.VLLM_MAX_KEYS
+        ),
+    )
+    try:
+        return launches, step.run(tiny, mode)
+    finally:
+        tiny.model.mono.close()
+        tiny.model.mono = None
+
+
+def _assert_same_step(stock, mono) -> None:
+    (hidden_s, aux_s, kv_s), (hidden_m, aux_m, kv_m) = stock, mono
+    torch.testing.assert_close(hidden_m, hidden_s, rtol=0, atol=0)
+    for name in kv_s:
+        torch.testing.assert_close(kv_m[name], kv_s[name], rtol=0, atol=0)
+    assert len(aux_m) == len(aux_s)
+    # The stock loop blends an aux state on its own, the mono path takes the
+    # blend the next layer's input norm makes: fp32 rounding apart.
+    for a_m, a_s in zip(aux_m, aux_s):
+        torch.testing.assert_close(a_m, a_s, rtol=1e-5, atol=1e-6)
+
+
+@requires_gfx950
+@pytest.mark.parametrize(
+    "aux_layers",
+    # vLLM's aux ids (L = the output of layer L - 1) for the target_layer_ids of
+    # Muse-Spark-Open-DSpark, of Muse-Spark-Open-DFlash2, and of the embedding
+    # and every stock / mono seam (0, 1, 59, 60, 61).
+    [(14, 30, 46, 58, 62), (2, 17, 31, 45, 60), (1, 2, 60, 61, 62)],
+    ids=["dspark", "dflash2", "seams"],
+)
+@pytest.mark.parametrize(
+    ("num_reqs", "bnd"),
+    [(1, 0), (2, 0), (3, 0), (3, 3), (4, 3), (5, 0), (7, 0)],
+)
+def test_mono_step_matches_the_stock_loop(
+    tiny_model, aux_layers, num_reqs, bnd
+) -> None:
+    """8 / 16 rows: K3 with fused attention; 24-32 rows: HA + HM + HB on the
+    32-row build (fused attention, replicated (bnd 3) or token-owner (bnd 0)
+    stream); 40-56 rows: the 64-row build with vLLM's attention op between
+    the launches. Same hidden states, aux states and KV caches as stock."""
+    tiny_model.lm.set_aux_hidden_state_layers(aux_layers)
+    step = Step(tiny_model, [ROWS_PER_REQUEST] * num_reqs, seed=num_reqs)
+    stock = step.run(tiny_model)
+    launches, mono = _run_mono(tiny_model, step, bnd, [8, 16, 32, 64])
+
+    kernels = {name for name, _ in launches.calls}
+    assert kernels == ({"k3"} if num_reqs <= 2 else {"ha", "hm", "hb"})
+    assert [a["layer"] for name, a in launches.calls if name in ("k3", "hb")] == list(
+        range(FIRST_LAYER, LAST_LAYER + 1)
+    )
+    _assert_same_step(stock, mono)
+
+
+@requires_gfx950
+@pytest.mark.parametrize(
+    ("query_lens", "mode"),
+    [
+        ([20], CUDAGraphMode.NONE),  # prefill
+        ([8, 20], CUDAGraphMode.NONE),  # mixed
+        ([4, 4], CUDAGraphMode.NONE),  # shorter verify
+        ([8] * 17, CUDAGraphMode.NONE),  # 136 rows: above the largest build
+        ([8], CUDAGraphMode.PIECEWISE),
+    ],
+    ids=["prefill", "mixed", "query-len-4", "136-rows", "piecewise"],
+)
+def test_other_steps_run_stock(tiny_model, query_lens, mode) -> None:
+    tiny_model.lm.set_aux_hidden_state_layers((14, 30, 46, 58, 62))
+    step = Step(tiny_model, query_lens, seed=len(query_lens))
+    stock = step.run(tiny_model, mode)
+    launches, mono = _run_mono(tiny_model, step, 0, [8, 16, 32, 64, 96, 128], mode)
+    assert launches.calls == []
+    _assert_same_step(stock, mono)
diff --git a/tests/models/test_muse_spark_rocm_shared_rope.py b/tests/models/test_muse_spark_rocm_shared_rope.py
index 3847f5a7b1fa..1dac6c37b2e3 100644
--- a/tests/models/test_muse_spark_rocm_shared_rope.py
+++ b/tests/models/test_muse_spark_rocm_shared_rope.py
@@ -6,10 +6,12 @@
 import torch
 from torch import nn
 
+import vllm.model_executor.layers.rotary_embedding as rotary_embedding
 import vllm.models.muse_spark.amd.model as amd_model
 
 
 def test_target_layers_share_fp32_rope_cache(monkeypatch, default_vllm_config):
+    monkeypatch.setattr(rotary_embedding, "_ROPE_DICT", {})
     monkeypatch.setattr(amd_model, "get_tensor_model_parallel_world_size", lambda: 1)
     for name in ("MuseSpark1xQKVOGateParallelLinear", "RowParallelLinear", "Attention"):
         monkeypatch.setattr(amd_model, name, lambda *args, **kwargs: nn.Identity())
diff --git a/tests/models/test_qwen3_5.py b/tests/models/test_qwen3_5.py
new file mode 100644
index 000000000000..333099635ae5
--- /dev/null
+++ b/tests/models/test_qwen3_5.py
@@ -0,0 +1,47 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+
+def test_qwen3_5_models_use_platform_specific_implementation():
+    from vllm.models.qwen3_5 import Qwen3_5ForCausalLM, Qwen3_5MTP
+    from vllm.platforms import current_platform
+
+    # Only the NVIDIA variant drops torch.compile; ROCm and other platforms
+    # keep the previous torch.compile implementation in
+    # vllm/model_executor/models/qwen3_5{,_mtp}.py.
+    if current_platform.is_cuda():
+        model_module = "vllm.models.qwen3_5.nvidia.model"
+        mtp_module = "vllm.models.qwen3_5.nvidia.mtp"
+    else:
+        model_module = "vllm.model_executor.models.qwen3_5"
+        mtp_module = "vllm.model_executor.models.qwen3_5_mtp"
+    assert Qwen3_5ForCausalLM.__module__ == model_module
+    assert Qwen3_5MTP.__module__ == mtp_module
+
+
+def test_qwen3_5_nvidia_models_do_not_use_torch_compile_wrapper():
+    # Import the NVIDIA modules directly so the assertion is independent of
+    # the platform this test runs on.
+    from vllm.compilation.wrapper import TorchCompileWithNoGuardsWrapper
+    from vllm.models.qwen3_5.nvidia import model as model_module
+    from vllm.models.qwen3_5.nvidia import mtp as mtp_module
+
+    assert not issubclass(model_module.Qwen3_5Model, TorchCompileWithNoGuardsWrapper)
+    assert not issubclass(
+        mtp_module.Qwen3_5MultiTokenPredictor, TorchCompileWithNoGuardsWrapper
+    )
+    assert not issubclass(mtp_module.Qwen3_5MTP, TorchCompileWithNoGuardsWrapper)
+
+
+def test_qwen3_5_legacy_models_keep_torch_compile_wrapper():
+    # The legacy implementation in vllm/model_executor/models/ serves ROCm and
+    # other non-NVIDIA platforms; it must keep the torch.compile wrapper.
+    from vllm.compilation.wrapper import TorchCompileWithNoGuardsWrapper
+    from vllm.model_executor.models import qwen3_5 as model_module
+    from vllm.model_executor.models import qwen3_5_mtp as mtp_module
+
+    assert issubclass(model_module.Qwen3_5Model, TorchCompileWithNoGuardsWrapper)
+    assert issubclass(
+        mtp_module.Qwen3_5MultiTokenPredictor, TorchCompileWithNoGuardsWrapper
+    )
+    assert issubclass(mtp_module.Qwen3_5MTP, TorchCompileWithNoGuardsWrapper)
diff --git a/tests/models/test_qwen3_5_mtp_config.py b/tests/models/test_qwen3_5_mtp_config.py
index 885995b906f4..5b6fe6edfbc6 100644
--- a/tests/models/test_qwen3_5_mtp_config.py
+++ b/tests/models/test_qwen3_5_mtp_config.py
@@ -2,11 +2,18 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """CPU-only tests for Qwen3.5 MTP speculative decoding config overrides."""
 
+from pathlib import Path
 from typing import Any
 
 import pytest
-from transformers import AutoConfig, PreTrainedConfig
+from transformers import (
+    AutoConfig,
+    PreTrainedConfig,
+    Qwen3_5Config,
+    Qwen3_5TextConfig,
+)
 
+from vllm.config import ModelConfig, ParallelConfig
 from vllm.config.speculative import SpeculativeConfig
 
 _CHECKPOINTS = {
@@ -119,3 +126,59 @@ def test_mtp_override_downloads_real_hf_hub_configs(
     assert cfg.model_type == "qwen3_5_mtp"
     assert cfg.architectures == [expected_arch]
     assert cfg.n_predict == 1
+
+
+@pytest.mark.parametrize("multimodal", [False, True], ids=["text", "multimodal"])
+@pytest.mark.parametrize(
+    ("mtp_layers", "num_spec_tokens", "index_share", "multi_module"),
+    [
+        (1, 2, None, False),
+        (2, 1, None, False),
+        (2, 2, None, True),
+        (2, 2, True, False),
+    ],
+)
+def test_mtp_num_hidden_layers_selects_multi_module_mtp(
+    tmp_path: Path,
+    multimodal: bool,
+    mtp_layers: int,
+    num_spec_tokens: int,
+    index_share: bool | None,
+    multi_module: bool,
+) -> None:
+    """Qwen3.5 declares its MTP modules as `mtp_num_hidden_layers` (nested in
+    `text_config` when multimodal), not `num_nextn_predict_layers`. MTP index
+    sharing is only supported by the single-module speculator."""
+    text_config = Qwen3_5TextConfig(
+        hidden_size=64,
+        intermediate_size=128,
+        num_hidden_layers=4,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        head_dim=16,
+        vocab_size=128,
+        mtp_num_hidden_layers=mtp_layers,
+    )
+    hf_config: PreTrainedConfig
+    if multimodal:
+        hf_config = Qwen3_5Config(
+            text_config=text_config.to_dict(),
+            architectures=["Qwen3_5ForConditionalGeneration"],
+        )
+    else:
+        text_config.architectures = ["Qwen3_5ForCausalLM"]
+        hf_config = text_config
+    hf_config.save_pretrained(tmp_path)
+
+    speculative_config = SpeculativeConfig(
+        method="mtp",
+        num_speculative_tokens=num_spec_tokens,
+        index_share_for_mtp_iteration=index_share,
+        target_model_config=ModelConfig(
+            model=str(tmp_path), tokenizer_mode="skip", max_model_len=2048
+        ),
+        target_parallel_config=ParallelConfig(),
+    )
+
+    assert speculative_config.use_multi_module_mtp() is multi_module
+    assert speculative_config.num_mtp_modules() == mtp_layers
diff --git a/tests/models/test_qwen3_dflash_mxfp8_context_kv.py b/tests/models/test_qwen3_dflash_mxfp8_context_kv.py
new file mode 100644
index 000000000000..447482762e3a
--- /dev/null
+++ b/tests/models/test_qwen3_dflash_mxfp8_context_kv.py
@@ -0,0 +1,145 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""DFlash/DSpark context K/V projection with MXFP8 QKV weights."""
+
+import pytest
+import torch
+from torch import nn
+
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.layers.linear import (
+    QKVParallelLinear,
+    UnquantizedLinearMethod,
+)
+from vllm.model_executor.layers.quantization.modelopt import ModelOptMxFp8Config
+from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
+    _mxfp8_e4m3_quantize_torch,
+)
+from vllm.model_executor.models.qwen3_dflash import DFlashQwen3Model
+from vllm.platforms import current_platform
+
+pytestmark = pytest.mark.skipif(
+    not current_platform.is_cuda()
+    or not current_platform.is_device_capability_family(100),
+    reason="Native MXFP8 kernels require SM100-family CUDA",
+)
+
+HIDDEN, HEAD_DIM, NUM_HEADS, NUM_LAYERS, NUM_TOKENS = 256, 128, 2, 2, 7
+Q_SIZE, KV_SIZE = NUM_HEADS * HEAD_DIM, HEAD_DIM  # one KV head
+
+
+@pytest.fixture
+def bf16_on_cuda(default_vllm_config, dist_init):
+    old_dtype = torch.get_default_dtype()
+    torch.set_default_dtype(torch.bfloat16)
+    with torch.device("cuda"):
+        yield
+    torch.set_default_dtype(old_dtype)
+
+
+def _build_model(mode: str, has_bias: bool) -> tuple[DFlashQwen3Model, list]:
+    # MiniMax-style MXFP8 config; "mixed" leaves layer 1 in BF16.
+    quant_config = (
+        None
+        if mode == "bf16"
+        else ModelOptMxFp8Config.from_config(
+            {
+                "quant_method": "mxfp8",
+                "activation_scheme": "dynamic",
+                "weight_block_size": [1, 32],
+                "ignored_layers": (
+                    ["layers.1.self_attn.qkv_proj"] if mode == "mixed" else []
+                ),
+            }
+        )
+    )
+    model = DFlashQwen3Model.__new__(DFlashQwen3Model)
+    nn.Module.__init__(model)
+    model.hidden_norm = RMSNorm(HIDDEN)
+    model._rms_norm_eps = 1e-6
+    attentions = []
+    checkpoint = []
+    for i in range(NUM_LAYERS):
+        attn = nn.Module()
+        attn.q_size, attn.kv_size = Q_SIZE, KV_SIZE
+        attn.k_norm = RMSNorm(HEAD_DIM)
+        attn.qkv_proj = QKVParallelLinear(
+            HIDDEN,
+            HEAD_DIM,
+            NUM_HEADS,
+            1,
+            bias=has_bias,
+            quant_config=quant_config,
+            prefix=f"layers.{i}.self_attn.qkv_proj",
+            disable_tp=True,
+        )
+        weight = torch.randn(Q_SIZE + 2 * KV_SIZE, HIDDEN) / 16
+        proj = attn.qkv_proj
+        tensors = {}
+        if isinstance(proj.quant_method, UnquantizedLinearMethod):
+            proj.weight.data.copy_(weight)
+        else:
+            weight, scale = _mxfp8_e4m3_quantize_torch(weight)
+            proj.weight.data.copy_(weight)
+            proj.weight_scale.data.copy_(scale)
+            tensors = {"weight": weight, "weight_scale": scale}
+        if has_bias:
+            proj.bias.data.normal_(std=0.1)
+            tensors["bias"] = proj.bias.data.clone()
+        attentions.append(attn)
+        checkpoint.append(tensors)
+    model.layers = nn.ModuleList()
+    for attn in attentions:
+        layer = nn.Module()
+        layer.self_attn = attn
+        model.layers.append(layer)
+    # As __init__ and load_weights do: the MXFP8 context projection exists
+    # from construction and loads the checkpoint's K/V rows.
+    model._context_kv_proj = None
+    if mode == "mxfp8":
+        model._context_kv_proj = model._make_mxfp8_context_kv_proj(attentions)
+        for i, tensors in enumerate(checkpoint):
+            for name, tensor in tensors.items():
+                model._load_context_kv(i, "k", name, tensor[Q_SIZE : Q_SIZE + KV_SIZE])
+                model._load_context_kv(i, "v", name, tensor[Q_SIZE + KV_SIZE :])
+    return model, attentions
+
+
+@pytest.mark.parametrize("mode", ["bf16", "mxfp8", "mixed"])
+@pytest.mark.parametrize("has_bias", [False, True])
+@torch.inference_mode()
+def test_context_kv_matches_qkv_projection(bf16_on_cuda, mode, has_bias):
+    """The context K/V must equal the K/V slice of each layer's own QKV output
+    after kernel repacking, eagerly and under CUDA graph replay."""
+    torch.manual_seed(0)
+    model, attentions = _build_model(mode, has_bias)
+    context_kv_proj = model._context_kv_proj
+    assert (context_kv_proj is not None) == (mode == "mxfp8")
+    # Load-time order: every module's post-load step (the context projection's
+    # included), then the model-level step that builds the context buffers.
+    for module in model.modules():
+        method = getattr(module, "quant_method", None)
+        if method is not None:
+            method.process_weights_after_loading(module)
+    model._build_context_kv_buffers(attentions, has_bias=has_bias)
+    assert (model._fused_kv_weight is not None) == (mode == "bf16")
+    assert model._context_kv_proj is context_kv_proj
+
+    x = torch.randn(NUM_TOKENS, HIDDEN)
+    normed = model.hidden_norm(x)
+    expected = torch.cat(
+        [a.qkv_proj(normed)[0][:, Q_SIZE:] for a in attentions], dim=-1
+    )
+
+    def project() -> torch.Tensor:
+        k, v = model._project_context_kv(x, NUM_TOKENS, NUM_LAYERS, 1, HEAD_DIM)
+        return torch.stack([k, v]).permute(2, 1, 0, 3, 4).reshape(NUM_TOKENS, -1)
+
+    torch.testing.assert_close(project(), expected, rtol=0.02, atol=0.02)
+
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        captured = project()
+    x.normal_()
+    graph.replay()
+    torch.testing.assert_close(captured, project(), rtol=0, atol=0)
diff --git a/tests/models/test_utils.py b/tests/models/test_utils.py
index 0dac2dbf2a6d..c27ce40df259 100644
--- a/tests/models/test_utils.py
+++ b/tests/models/test_utils.py
@@ -239,3 +239,121 @@ def test_weights_mapper_stacks_one_weight_into_several_shards():
     # 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()
+
+
+def _double(weights):
+    """Stateless recode: double every tensor."""
+    for name, tensor in weights:
+        yield name, tensor * 2
+
+
+def _pair_halves(weights):
+    """Stateful merge: cat a name's ``.first``/``.second`` halves on arrival."""
+    pending: dict[str, dict[str, torch.Tensor]] = {}
+    for name, tensor in weights:
+        for half in (".first", ".second"):
+            if name.endswith(half):
+                base = name[: -len(half)]
+                parts = pending.setdefault(base, {})
+                parts[half] = tensor
+                if len(parts) == 2:
+                    del pending[base]
+                    yield base, torch.cat([parts[".first"], parts[".second"]])
+                break
+        else:
+            yield name, tensor
+    if pending:
+        raise ValueError(f"unpaired: {sorted(pending)}")
+
+
+def test_weights_mapper_transform_runs_before_name_mapping():
+    """Transform-emitted names flow through the name mappings, and an N->1
+    merge completes however its inputs are ordered."""
+    mapper = WeightsMapper(
+        orig_to_new_suffix={".w": ".w_renamed"},
+        weights_transforms=[_pair_halves],
+    )
+    first, second = torch.ones(2), torch.zeros(2)
+    keep = torch.ones(1)
+    weights = [("keep.w", keep), ("e.w.second", second), ("e.w.first", first)]
+
+    mapped = list(mapper.apply(weights))
+
+    assert [name for name, _ in mapped] == ["keep.w_renamed", "e.w_renamed"]
+    assert mapped[0][1] is keep
+    assert torch.equal(mapped[1][1], torch.cat([first, second]))
+
+
+def test_weights_mapper_transform_can_recode_values():
+    mapper = WeightsMapper(weights_transforms=[_double])
+
+    mapped = list(mapper.apply([("w", torch.ones(2))]))
+
+    assert torch.equal(mapped[0][1], torch.full((2,), 2.0))
+
+
+def test_weights_mapper_transform_raises_at_iteration_not_at_apply():
+    """Incomplete input raises when the stream is exhausted; merely calling
+    apply() (or the transform) does nothing yet."""
+    mapper = WeightsMapper(weights_transforms=[_pair_halves])
+    stream = mapper.apply([("e.w.first", torch.ones(2))])
+
+    with pytest.raises(ValueError, match="unpaired"):
+        list(stream)
+
+
+def test_weights_mapper_transform_state_is_fresh_per_apply():
+    """The mapper is shared and long-lived; each apply() must start clean."""
+    mapper = WeightsMapper(weights_transforms=[_pair_halves])
+    with pytest.raises(ValueError, match="unpaired"):
+        list(mapper.apply([("e.w.first", torch.ones(2))]))
+    # A leaked pending entry would pair with this stream instead of raising.
+    with pytest.raises(ValueError, match="unpaired"):
+        list(mapper.apply([("e.w.second", torch.zeros(2))]))
+
+    first, second = torch.ones(2), torch.zeros(2)
+    mapped = list(mapper.apply([("e.w.second", second), ("e.w.first", first)]))
+    assert torch.equal(mapped[0][1], torch.cat([first, second]))
+
+
+def test_weights_mapper_or_concatenates_transforms_in_order():
+    """Self's transforms run before other's; a mapper without any adds none."""
+    seen = []
+
+    def tag(label):
+        def rename(weights):
+            for name, tensor in weights:
+                seen.append((label, name))
+                yield f"{name}.{label}", tensor
+
+        return rename
+
+    mapper = WeightsMapper(weights_transforms=[tag("x"), tag("y")]) | WeightsMapper(
+        weights_transforms=[tag("z")]
+    )
+    (mapped,) = list(mapper.apply([("w", torch.ones(1))]))
+
+    assert mapped[0] == "w.x.y.z"
+    assert seen == [("x", "w"), ("y", "w.x"), ("z", "w.x.y")]
+    assert (WeightsMapper() | WeightsMapper()).weights_transforms == []
+    one_side = WeightsMapper() | WeightsMapper(weights_transforms=[_pair_halves])
+    assert one_side.weights_transforms == [_pair_halves]
+
+
+def _double_w(weights):
+    """Stateless dispatch by name, the idiomatic inline ``endswith`` form."""
+    for name, tensor in weights:
+        if name.endswith(".w"):
+            tensor = tensor * 2
+        yield name, tensor
+
+
+def test_weights_mapper_transforms_see_earlier_emitted_names():
+    mapper = WeightsMapper(weights_transforms=[_pair_halves, _double_w])
+
+    (mapped,) = list(
+        mapper.apply([("e.w.second", torch.zeros(1)), ("e.w.first", torch.ones(1))])
+    )
+
+    # The halves merge into "e.w", which the suffix mapping then doubles.
+    assert torch.equal(mapped[1], torch.tensor([2.0, 0.0]))
diff --git a/tests/models/transformers/fusers/test_linear.py b/tests/models/transformers/fusers/test_linear.py
index bcad62765a76..ecb7ff0b08c8 100644
--- a/tests/models/transformers/fusers/test_linear.py
+++ b/tests/models/transformers/fusers/test_linear.py
@@ -4,6 +4,7 @@
 
 import ast
 import inspect
+from collections.abc import MutableMapping
 from types import MethodType, SimpleNamespace
 
 import pytest
@@ -838,9 +839,14 @@ def test_glu_identifies_down_projection():
     matches the column-parallel merged gate/up; `None` when there is no such
     projection to force (fusion of gate/up still applies)."""
     with torch.device("meta"):
-        assert get_fuser(GLUMLP(), GLUFuser).down_name == "down_proj"
-        assert get_fuser(ReversedGLUMLP(), GLUFuser).down_name == "down_proj"
-        assert get_fuser(NoDownGLU(), GLUFuser).down_name is None
+        for cls, down_name in [
+            (GLUMLP, "down_proj"),
+            (ReversedGLUMLP, "down_proj"),
+            (NoDownGLU, None),
+        ]:
+            fuser = get_fuser(cls(), GLUFuser)
+            assert fuser is not None
+            assert fuser.down_name == down_name
 
 
 @pytest.mark.parametrize("attn_cls", [FakeAttention, ReversedFakeAttention])
@@ -891,14 +897,19 @@ def test_detects_and_rewrites_qkv(attn_cls, kv_heads):
 
 def test_qkv_identifies_output_projection():
     with torch.device("meta"):
-        assert get_fuser(FakeAttention(), QKVFuser).o_name == "o_proj"
-        assert get_fuser(ReversedFakeAttention(), QKVFuser).o_name == "o_proj"
-        assert get_fuser(ExtraProjAttention(), QKVFuser).o_name == "o_proj"
-        # Norm children (q_norm/k_norm) must not disturb o_proj identification.
-        assert get_fuser(QKNormAttention(), QKVFuser).o_name == "o_proj"
-        assert get_fuser(PerHeadQKNormAttention(), QKVFuser).o_name == "o_proj"
-        # A module between o_proj and the return is transparent.
-        assert get_fuser(ResidDropoutAttention(), QKVFuser).o_name == "o_proj"
+        for cls in [
+            FakeAttention,
+            ReversedFakeAttention,
+            ExtraProjAttention,
+            # Norm children must not disturb o_proj identification.
+            QKNormAttention,
+            PerHeadQKNormAttention,
+            # A module between o_proj and the return is transparent.
+            ResidDropoutAttention,
+        ]:
+            fuser = get_fuser(cls(), QKVFuser)
+            assert fuser is not None
+            assert fuser.o_name == "o_proj"
 
 
 @pytest.mark.parametrize(
@@ -1265,6 +1276,7 @@ def test_fuser_is_cached_per_class_and_structure():
         fuser_a = get_fuser(GLUMLP(), GLUFuser)
         fuser_b = get_fuser(GLUMLP(), GLUFuser)
     assert fuser_a is fuser_b
+    assert get_fusers.cache is not None
     assert any(key[0] is GLUMLP for key in get_fusers.cache)
 
 
@@ -1297,8 +1309,10 @@ def test_weight_mappings_are_scoped_to_fused_prefixes():
     with torch.device("meta"):
         glu_fuser = get_fuser(GLUMLP(), GLUFuser)
         qkv_fuser = get_fuser(FakeAttention(), QKVFuser)
+    assert glu_fuser is not None and qkv_fuser is not None
 
     mapper = WeightsMapper()
+    assert isinstance(mapper.orig_to_new_stacked, MutableMapping)
     for prefix in ("model.layers.0.mlp", "model.layers.1.mlp"):
         mapper.orig_to_new_stacked.update(glu_fuser.orig_to_new_stacked(prefix))
     mapper.orig_to_new_stacked.update(
diff --git a/tests/models/transformers/fusers/test_moe.py b/tests/models/transformers/fusers/test_moe.py
index 284b1d6a094e..7a8d801765c4 100644
--- a/tests/models/transformers/fusers/test_moe.py
+++ b/tests/models/transformers/fusers/test_moe.py
@@ -493,6 +493,7 @@ def test_moe_fuser_reads_router_dtype_from_the_gate():
     with torch.device("meta"):
         fp32 = MoEBlockFuser.match(MoEBlock(Fp32Router), "experts")
         default = MoEBlockFuser.match(MoEBlock(TopKRouter), "experts")
+    assert fp32 is not None and default is not None
     assert fp32.router_dtype == torch.float32
     assert default.router_dtype is None
 
@@ -566,6 +567,7 @@ def test_moe_fuser_unrecognized_router_moves_shared_expert():
     with torch.device("meta"):
         block = MoEBlockSharedBiasedRouter()
     fuser = MoEBlockFuser.match(block, "experts")
+    assert fuser is not None
     assert fuser.scoring_func is None
     assert fuser.shared_name == "shared_expert"
 
@@ -580,6 +582,7 @@ def test_moe_fuser_unrecognized_router_moves_shared_expert():
 def test_moe_fuser_divided_shared_output(block_cls, scoring_func):
     with torch.device("meta"):
         fuser = MoEBlockFuser.match(block_cls(), "experts", divides_shared_output=True)
+    assert fuser is not None
     assert fuser.scoring_func == scoring_func
 
 
@@ -600,6 +603,7 @@ def test_moe_fuser_scaled_output_with_vllm_routing():
     """A scaled block fuses because vLLM routing applies `routed_scaling_factor`."""
     with torch.device("meta"):
         fuser = MoEBlockFuser.match(MoEBlockScaled(), "experts")
+    assert fuser is not None
     assert fuser.scoring_func == "softmax"
     assert fuser.shared_name == "shared_expert"
 
@@ -608,4 +612,5 @@ def test_moe_fuser_scaled_output_with_vllm_routing():
 def test_moe_fuser_looks_through_casts_and_identity(block_cls):
     with torch.device("meta"):
         fuser = MoEBlockFuser.match(block_cls(), "experts")
+    assert fuser is not None
     assert fuser.shared_name == "shared_expert"
diff --git a/tests/models/transformers/fusers/test_rms_norm.py b/tests/models/transformers/fusers/test_rms_norm.py
index 256f83649a48..142c8075f454 100644
--- a/tests/models/transformers/fusers/test_rms_norm.py
+++ b/tests/models/transformers/fusers/test_rms_norm.py
@@ -255,6 +255,7 @@ def test_rms_norm_builds_vllm_class(cls, expected, zero_centered, default_vllm_c
     with torch.device("meta"):
         module = cls()
         fuser = get_fuser(module, RMSNormFuser)
+        assert fuser is not None
         built = fuser.fuse(module, "norm", default_vllm_config)
     from vllm.model_executor.models.transformers.fusers.rms_norm import (
         TPAwareNormMixin,
@@ -283,7 +284,9 @@ def test_pow_spelled_norm_fuses_to_equivalent_math(cls, default_vllm_config):
     module = cls(16)
     with torch.no_grad():
         module.weight.copy_(torch.randn(16))
-    built = get_fuser(module, RMSNormFuser).fuse(module, "norm", default_vllm_config)
+    fuser = get_fuser(module, RMSNormFuser)
+    assert fuser is not None
+    built = fuser.fuse(module, "norm", default_vllm_config)
     with torch.no_grad():
         built.weight.copy_(module.weight)
     x = torch.randn(4, 16)
@@ -298,7 +301,9 @@ def test_weightless_norm_has_no_hidden_size(default_vllm_config):
     matching the unfused norm and vLLM's native `RMSNorm(hidden_size=head_dim)`.
     """
     module = WeightlessRMSNorm(72)
-    built = get_fuser(module, RMSNormFuser).fuse(module, "norm", default_vllm_config)
+    fuser = get_fuser(module, RMSNormFuser)
+    assert fuser is not None
+    built = fuser.fuse(module, "norm", default_vllm_config)
     assert built.hidden_size == 0
 
     built.tp_size = 2  # emulate TP=2 without a real process group
@@ -329,6 +334,7 @@ def test_eps_is_derived_per_instance(default_vllm_config):
         for eps in (1e-5, 1e-6):
             module = RMSNorm(16, eps=eps)
             fuser = get_fuser(module, RMSNormFuser)
+            assert fuser is not None
             built = fuser.fuse(module, "norm", default_vllm_config)
             assert built.variance_epsilon == eps
 
@@ -341,6 +347,7 @@ def test_eps_attr_is_found_by_value_not_name(cls, default_vllm_config):
         for eps in (1e-5, 1e-6):
             module = cls(16, eps=eps)
             fuser = get_fuser(module, RMSNormFuser)
+            assert fuser is not None
             assert fuser.eps_attr == cls.attr
             built = fuser.fuse(module, "norm", default_vllm_config)
             assert built.variance_epsilon == eps
@@ -353,6 +360,7 @@ def test_literal_eps_is_not_mistaken_for_an_attribute(default_vllm_config, caplo
     with caplog.at_level("DEBUG", logger=logger), torch.device("meta"):
         module = LiteralEpsRMSNorm()
         fuser = get_fuser(module, RMSNormFuser)
+        assert fuser is not None
         built = fuser.fuse(module, "norm", default_vllm_config)
     assert fuser.eps_attr is None
     assert built.variance_epsilon == 1e-4
@@ -366,6 +374,7 @@ def test_ambiguous_eps_attrs_are_disambiguated(default_vllm_config):
         module = AmbiguousEpsRMSNorm(16, eps=1e-6)
         before = dict(vars(module))
         fuser = get_fuser(module, RMSNormFuser)
+        assert fuser is not None
         assert fuser.eps_attr == "variance_epsilon"
         assert vars(module) == before
 
diff --git a/tests/models/transformers/test_backend.py b/tests/models/transformers/test_backend.py
index dd7cae111fb1..a038eff81adc 100644
--- a/tests/models/transformers/test_backend.py
+++ b/tests/models/transformers/test_backend.py
@@ -13,9 +13,22 @@
 import pytest
 import torch
 import torch.nn as nn
-from transformers import AutoConfig, AutoModel, PreTrainedConfig
+from safetensors.torch import load_file, save_file
+from transformers import (
+    AutoConfig,
+    AutoModel,
+    AutoModelForCausalLM,
+    PreTrainedConfig,
+)
 
-from vllm.config import ModelConfig, VllmConfig
+from vllm.config import (
+    ModelConfig,
+    ParallelConfig,
+    SpeculativeConfig,
+    VllmConfig,
+    set_current_vllm_config,
+)
+from vllm.model_executor import model_loader
 from vllm.model_executor.models.interfaces import SupportsMultiModal
 from vllm.model_executor.models.transformers.base import Base
 from vllm.model_executor.models.transformers.fusers import AttentionFuser
@@ -148,6 +161,26 @@ def test_hybrid_attention(vllm_runner: type[VllmRunner]) -> None:
     )
 
 
+def test_whole_graph_compile(
+    vllm_runner: type[VllmRunner], example_prompts: list[str]
+) -> None:
+    """Compiling attention and the KV-cache op into one Inductor graph must
+    match eager, so vllm:: ops have to receive eager strides."""
+    kwargs_ref = {"model_impl": "transformers", "max_model_len": 2048}
+    kwargs_test = {
+        **kwargs_ref,
+        "compilation_config": {"mode": 2, "backend": "inductor"},
+    }
+    check_implementation(
+        vllm_runner,
+        vllm_runner,
+        example_prompts,
+        model="Qwen/Qwen3-0.6B",
+        kwargs_ref={**kwargs_ref, "enforce_eager": True},
+        kwargs_test=kwargs_test,
+    )
+
+
 def get_sinks(model) -> dict[int, torch.Tensor]:
     """The sink tensor each attention layer was handed, keyed by layer index."""
     sinks = {}
@@ -694,8 +727,11 @@ def test_attention_layer_index_is_the_modules_own():
     attentions = [
         model.get_submodule(f"layers.{i}.self_attn.{j}") for i in (0, 1) for j in (0, 1)
     ]
-    fusers = [AttentionFuser.match(None, attn) for attn in attentions]
-    assert all(fuser is not None for fuser in fusers)
+    fusers: list[AttentionFuser] = []
+    for attn in attentions:
+        fuser = AttentionFuser.match(None, attn)
+        assert fuser is not None
+        fusers.append(fuser)
     assert [f.layer_index(a) for f, a in zip(fusers, attentions)] == [0, 1, 2, 3]
 
     # Not dispatching to vLLM yet, as a vision tower never would be. The vLLM
@@ -705,17 +741,13 @@ def test_attention_layer_index_is_the_modules_own():
     assert all(f.validate(a, None) for f, a in zip(fusers, attentions))
 
 
-def test_attention_dispatch_is_required(monkeypatch: pytest.MonkeyPatch):
+def test_attention_dispatch_is_required():
     """Models selected by the permissive registry fail clearly without dispatch."""
-    monkeypatch.setattr(
-        "vllm.model_executor.models.transformers.base.get_pp_indices",
-        lambda *_: (0, 1),
-    )
     model = SimpleNamespace(
         text_config=SimpleNamespace(num_hidden_layers=1),
-        pp_group=SimpleNamespace(rank_in_group=0, world_size=1),
         attention_fusers={},
         _get_attn_cls=lambda: None,
+        _pp_indices=lambda: (0, 1),
     )
 
     with pytest.raises(
@@ -784,3 +816,99 @@ def test_get_mrope_input_positions_omits_unsupported_grid_kwargs():
 
     assert positions.shape == (4, 3)
     assert delta == 0
+
+
+# The smallest sizes the fused MoE kernels accept
+MTP_HIDDEN_SIZE = 128
+MTP_MOE_KWARGS = dict(
+    vocab_size=128,
+    hidden_size=MTP_HIDDEN_SIZE,
+    intermediate_size=256,
+    moe_intermediate_size=128,
+    n_routed_experts=4,
+    num_experts_per_tok=2,
+)
+# Real layer counts, since the MTP weight patterns hard-code them
+MTP_MODELS = {
+    "glm4_moe": dict(
+        num_hidden_layers=46, num_attention_heads=2, num_key_value_heads=1, head_dim=64
+    ),
+    "deepseek_v3": dict(
+        num_hidden_layers=61,
+        num_attention_heads=16,
+        q_lora_rank=64,
+        kv_lora_rank=512,
+        qk_nope_head_dim=128,
+        qk_rope_head_dim=64,
+        v_head_dim=128,
+        first_k_dense_replace=1,
+        n_group=1,
+        topk_group=1,
+    ),
+}
+
+
+@pytest.mark.parametrize("model_type", MTP_MODELS)
+def test_mtp_draft_loads_mtp_weights(model_type, tmp_path, tp_init, monkeypatch):
+    """An MTP draft for a model following the Transformers MTP standard loads each
+    MTP layer appended to the target's checkpoint into `MtpModel`'s layout."""
+    kwargs = MTP_MODELS[model_type] | MTP_MOE_KWARGS
+    num_hidden_layers = kwargs.pop("num_hidden_layers")
+    num_mtp_layers = 2
+    hidden_size = MTP_HIDDEN_SIZE
+    config = AutoConfig.for_model(
+        model_type, num_hidden_layers=num_hidden_layers + num_mtp_layers, **kwargs
+    )
+    model_cls = AutoModelForCausalLM._model_mapping[type(config)]
+    config.architectures = [model_cls.__name__]
+    # These have one MTP layer, so register the second like Step-3.7 registers its own
+    monkeypatch.setattr(
+        model_cls,
+        "_keys_to_ignore_on_load_unexpected",
+        [rf"model\.layers\.{num_hidden_layers + 1}.*"]
+        + model_cls._keys_to_ignore_on_load_unexpected,
+    )
+    # Save the MTP layers as the last decoder layers, then add their MTP modules
+    model_cls(config).save_pretrained(tmp_path)
+    weights = load_file(tmp_path / "model.safetensors")
+    for i in range(num_hidden_layers, num_hidden_layers + num_mtp_layers):
+        for name, shape in {
+            "enorm.weight": (hidden_size,),
+            "hnorm.weight": (hidden_size,),
+            "eh_proj.weight": (hidden_size, 2 * hidden_size),
+            "shared_head.norm.weight": (hidden_size,),
+        }.items():
+            weights[f"model.layers.{i}.{name}"] = torch.randn(shape)
+    save_file(weights, tmp_path / "model.safetensors", metadata={"format": "pt"})
+    config.num_hidden_layers = num_hidden_layers
+    config.num_nextn_predict_layers = num_mtp_layers
+    config.save_pretrained(tmp_path)
+
+    target = ModelConfig(str(tmp_path), model_impl="transformers", dtype="bfloat16")
+    speculative_config = SpeculativeConfig(
+        method="mtp",
+        num_speculative_tokens=num_mtp_layers,
+        target_model_config=target,
+        target_parallel_config=ParallelConfig(),
+    )
+    draft = speculative_config.draft_model_config
+    assert draft.architecture == "TransformersMoEForMTP"
+    vllm_config = VllmConfig(model_config=target, speculative_config=speculative_config)
+    with set_current_vllm_config(vllm_config):
+        model = model_loader.get_model(vllm_config=vllm_config, model_config=draft)
+
+    assert len(model.model.layers) == num_mtp_layers
+    for k, layer in enumerate(model.model.layers):
+        mtp = f"model.layers.{num_hidden_layers + k}"
+        for ckpt_name, param_name in {
+            "enorm.weight": "enorm.weight",
+            "hnorm.weight": "hnorm.weight",
+            "eh_proj.weight": "eh_proj.weight",
+            "shared_head.norm.weight": "post_norm.weight",
+            "input_layernorm.weight": "mtp_block.input_layernorm.weight",
+            "self_attn.o_proj.weight": "mtp_block.self_attn.o_proj.weight",
+        }.items():
+            torch.testing.assert_close(
+                layer.get_parameter(param_name).cpu(),
+                weights[f"{mtp}.{ckpt_name}"].to(torch.bfloat16),
+            )
diff --git a/tests/models/transformers/test_layer_registry.py b/tests/models/transformers/test_layer_registry.py
index 555d87e379e1..2c6e9ced0a0e 100644
--- a/tests/models/transformers/test_layer_registry.py
+++ b/tests/models/transformers/test_layer_registry.py
@@ -1,23 +1,26 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Unit tests for the Transformers backend's hw-agnostic layer resolution.
+"""Unit tests for resolving layers between the in-tree and hw-agnostic paths.
 
-`layers._resolve` imports a layer symbol from
+`hw_agnostic.resolve` returns a layer from
 `vllm.model_executor.hw_agnostic.layers.<module>` when `VLLM_USE_HW_AGNOSTIC`
-is set and the symbol exists, and otherwise falls back to
-`vllm.model_executor.layers.<module>`. These tests pin that contract and the
-logging that reports which source was used.
+is set and the layer exists, and otherwise from
+`vllm.model_executor.layers.<module>`. The Transformers backend and out-of-tree
+plugins both use it, so a plugin's override reaches the class the backend
+builds.
 
 CPU-only: no engine is built. The end-to-end counterpart that serves a model
 through `vllm_runner` lives in `test_layer_registry_e2e.py` (GPU job).
 """
 
-import logging
+import importlib
 import sys
 import types
 
 import pytest
+import torch
 
+from vllm.model_executor import hw_agnostic
 from vllm.model_executor.models.transformers import layers
 
 HW_MODULE = "vllm.model_executor.hw_agnostic.layers.layernorm"
@@ -36,35 +39,71 @@ def fake_hw_layernorm(monkeypatch):
     return module
 
 
+@pytest.fixture
+def resolve_logs(monkeypatch):
+    """Messages `resolve` logs, captured directly since the `*_once` methods
+    deduplicate."""
+    from vllm.model_executor.hw_agnostic import _resolve
+
+    logged: list[str] = []
+    for method in ("info_once", "warning_once"):
+        monkeypatch.setattr(
+            _resolve.logger, method, lambda msg, *args, **_: logged.append(msg % args)
+        )
+    return logged
+
+
 def test_falls_back_to_vllm_when_disabled(monkeypatch, fake_hw_layernorm):
-    """Disabled: the vLLM class is used even if a hw-agnostic one exists."""
+    """Disabled: the in-tree class is used even if a hw-agnostic one exists."""
     monkeypatch.setenv("VLLM_USE_HW_AGNOSTIC", "0")
     from vllm.model_executor.layers.layernorm import RMSNorm as VllmRMSNorm
 
-    assert layers._resolve("layernorm", "RMSNorm") is VllmRMSNorm
+    assert hw_agnostic.resolve("layernorm", "RMSNorm") is VllmRMSNorm
 
 
-def test_uses_hw_agnostic_when_enabled(monkeypatch, fake_hw_layernorm, caplog_vllm):
+def test_uses_hw_agnostic_when_enabled(monkeypatch, fake_hw_layernorm, resolve_logs):
     """Enabled and available: the hw-agnostic class is used and logged."""
     monkeypatch.setenv("VLLM_USE_HW_AGNOSTIC", "1")
-    with caplog_vllm.at_level(logging.INFO):
-        resolved = layers._resolve("layernorm", "RMSNorm")
+    resolved = hw_agnostic.resolve("layernorm", "RMSNorm")
     assert resolved is fake_hw_layernorm.RMSNorm
-    assert "Using hw-agnostic layer: RMSNorm" in caplog_vllm.text
+    assert resolve_logs == ["Using hardware agnostic layer layernorm.RMSNorm"]
 
 
-def test_falls_back_when_symbol_missing(monkeypatch, caplog_vllm):
-    """Enabled but the symbol is not ported: fall back to vLLM and warn."""
+def test_falls_back_when_symbol_missing(monkeypatch, resolve_logs):
+    """Enabled but the symbol is not ported: fall back to in-tree and warn."""
     monkeypatch.setenv("VLLM_USE_HW_AGNOSTIC", "1")
     # A hw-agnostic module without the requested attribute triggers fallback.
     empty = types.ModuleType(HW_MODULE)
     monkeypatch.setitem(sys.modules, HW_MODULE, empty)
     from vllm.model_executor.layers.layernorm import RMSNorm as VllmRMSNorm
 
-    with caplog_vllm.at_level(logging.WARNING):
-        resolved = layers._resolve("layernorm", "RMSNorm")
+    resolved = hw_agnostic.resolve("layernorm", "RMSNorm")
     assert resolved is VllmRMSNorm
-    assert "falling back to default" in caplog_vllm.text
+    assert resolve_logs == [
+        "hw-agnostic layer layernorm.RMSNorm is not available; using the in-tree layer"
+    ]
+
+
+def test_falls_back_when_module_missing(monkeypatch):
+    """Enabled but the module is not ported: fall back to in-tree."""
+    monkeypatch.setenv("VLLM_USE_HW_AGNOSTIC", "1")
+    # `None` in `sys.modules` makes the import fail as for a missing module.
+    monkeypatch.setitem(sys.modules, HW_MODULE, None)
+    from vllm.model_executor.layers.layernorm import RMSNorm as VllmRMSNorm
+
+    assert hw_agnostic.resolve("layernorm", "RMSNorm") is VllmRMSNorm
+
+
+def test_broken_hw_agnostic_module_raises(monkeypatch, tmp_path):
+    """A port that fails to import raises instead of falling back silently."""
+    (tmp_path / "broken_port.py").write_text("import not_a_real_dependency\n")
+    hw_pkg = importlib.import_module("vllm.model_executor.hw_agnostic.layers")
+    monkeypatch.setattr(hw_pkg, "__path__", [*hw_pkg.__path__, str(tmp_path)])
+    importlib.invalidate_caches()
+    monkeypatch.setenv("VLLM_USE_HW_AGNOSTIC", "1")
+
+    with pytest.raises(ModuleNotFoundError, match="not_a_real_dependency"):
+        hw_agnostic.resolve("broken_port", "Layer")
 
 
 def test_act_and_mul_falls_back_for_unknown_activation(
@@ -78,3 +117,132 @@ def test_act_and_mul_falls_back_for_unknown_activation(
     from vllm.model_executor.layers.activation import GeluAndMul
 
     assert isinstance(layers.get_act_and_mul_fn("gelu"), GeluAndMul)
+
+
+@pytest.fixture
+def oot_registries():
+    """The in-tree and hw-agnostic `op_registry_oot`s, restored afterwards."""
+    from vllm.model_executor.custom_op import op_registry_oot as in_tree
+    from vllm.model_executor.hw_agnostic.custom_op import op_registry_oot as hw
+
+    saved = [(registry, dict(registry)) for registry in (in_tree, hw)]
+    yield in_tree, hw
+    for registry, before in saved:
+        registry.clear()
+        registry.update(before)
+
+
+@pytest.fixture
+def reload_backend():
+    """Re-import the backend modules that resolve layers at import, under the
+    test's `VLLM_USE_HW_AGNOSTIC`. Their original namespaces are restored
+    afterwards, so classes other tests imported stay current."""
+    from vllm.model_executor.models.transformers.fusers import rms_norm
+
+    saved = [(module, dict(vars(module))) for module in (layers, rms_norm)]
+    yield lambda: [importlib.reload(module) for module, _ in saved]
+    for module, namespace in saved:
+        vars(module).clear()
+        vars(module).update(namespace)
+
+
+def _fused_rms_norm(backend, rms_norm, vllm_config):
+    fuser = rms_norm.RMSNormFuser(
+        zero_centered=False, source_cls="LlamaRMSNorm", eps=1e-6
+    )
+    return fuser.fuse(torch.nn.RMSNorm(8), "norm", vllm_config)
+
+
+@pytest.mark.parametrize(
+    "module,name,build",
+    [
+        pytest.param(
+            "layernorm",
+            "RMSNorm",
+            _fused_rms_norm,
+            id="RMSNorm",
+            marks=pytest.mark.xfail(
+                strict=True,
+                raises=AssertionError,
+                reason="The fuser builds `TPAwareRMSNorm`, a subclass, and "
+                "overrides are keyed on the exact class name.",
+            ),
+        ),
+        pytest.param(
+            "activation",
+            "SiluAndMul",
+            lambda backend, *_: backend.get_act_and_mul_fn("silu"),
+            id="SiluAndMul",
+        ),
+        pytest.param(
+            "activation",
+            "GeluAndMul",
+            lambda backend, *_: backend.get_act_and_mul_fn("gelu"),
+            id="GeluAndMul",
+        ),
+    ],
+)
+@pytest.mark.parametrize("use_hw_agnostic", ["0", "1"])
+def test_plugin_override_reaches_the_backend(
+    monkeypatch,
+    oot_registries,
+    reload_backend,
+    default_vllm_config,
+    use_hw_agnostic,
+    module,
+    name,
+    build,
+):
+    """A plugin that subclasses `resolve`'s result overrides what the backend
+    builds, on both paths and for a layer with no hw-agnostic port (GeluAndMul).
+    """
+    monkeypatch.setenv("VLLM_USE_HW_AGNOSTIC", use_hw_agnostic)
+    backend, rms_norm = reload_backend()
+    for registry in oot_registries:
+        registry.pop(name, None)  # e.g. already owned by a loaded plugin
+    # The in-tree `get_act_and_mul_fn` caches the op it builds first; start empty,
+    # as a process does where plugins load before any model is built.
+    from vllm.model_executor.layers import activation
+
+    monkeypatch.setattr(activation._ACTIVATION_AND_MUL_REGISTRY, "_dict", {})
+
+    base = hw_agnostic.resolve(module, name)
+    override = base.register_oot(type(f"Plugin{name}", (base,), {}))
+
+    assert isinstance(build(backend, rms_norm, default_vllm_config), override)
+
+
+def test_in_tree_override_warns_on_hw_agnostic_path(monkeypatch):
+    """An override registered only for the in-tree class is not used on the
+    hw-agnostic path, and the warning says how to opt in."""
+    from vllm.model_executor.custom_op import op_registry_oot as in_tree_registry
+    from vllm.model_executor.hw_agnostic import custom_op as hw_custom_op
+    from vllm.model_executor.hw_agnostic.layers.layernorm import RMSNorm as HwRMSNorm
+    from vllm.model_executor.layers.layernorm import RMSNorm as VllmRMSNorm
+
+    plugin_cls = type("PluginRMSNorm", (VllmRMSNorm,), {})
+    monkeypatch.setitem(in_tree_registry, "RMSNorm", plugin_cls)
+    monkeypatch.delitem(hw_custom_op.op_registry_oot, "RMSNorm", raising=False)
+    warned: list[str] = []
+    monkeypatch.setattr(
+        hw_custom_op.logger,
+        "warning_once",
+        lambda msg, *args, **_: warned.append(msg % args),
+    )
+
+    assert type(HwRMSNorm.__new__(HwRMSNorm)) is HwRMSNorm
+    assert len(warned) == 1
+    assert "PluginRMSNorm is registered for the in-tree RMSNorm" in warned[0]
+
+
+def test_override_not_derived_from_hw_agnostic_class_raises(monkeypatch):
+    """Python would skip such an override's `__init__`, so it is rejected."""
+    from vllm.model_executor.hw_agnostic import custom_op as hw_custom_op
+    from vllm.model_executor.hw_agnostic.layers.layernorm import RMSNorm as HwRMSNorm
+    from vllm.model_executor.layers.layernorm import RMSNorm as VllmRMSNorm
+
+    plugin_cls = type("PluginRMSNorm", (VllmRMSNorm,), {})
+    monkeypatch.setitem(hw_custom_op.op_registry_oot, "RMSNorm", plugin_cls)
+
+    with pytest.raises(TypeError, match="PluginRMSNorm.*does not derive from"):
+        HwRMSNorm.__new__(HwRMSNorm)
diff --git a/tests/models/transformers/test_layer_registry_e2e.py b/tests/models/transformers/test_layer_registry_e2e.py
index 13d4b80a16e7..8ccf5c448b9e 100644
--- a/tests/models/transformers/test_layer_registry_e2e.py
+++ b/tests/models/transformers/test_layer_registry_e2e.py
@@ -5,7 +5,7 @@
 Serves a tiny Llama through `vllm_runner` with and without
 `VLLM_USE_HW_AGNOSTIC` and checks the layer providers and logprobs. Needs a
 real engine, so it runs in the GPU job. The CPU-only unit tests for the
-resolution logic (`layers._resolve`) are in `test_layer_registry.py`.
+resolution logic (`hw_agnostic.resolve`) are in `test_layer_registry.py`.
 """
 
 import pytest
diff --git a/tests/multimodal/test_video.py b/tests/multimodal/test_video.py
index 16334fdb0dc7..2a5cd18e760c 100644
--- a/tests/multimodal/test_video.py
+++ b/tests/multimodal/test_video.py
@@ -40,8 +40,10 @@
 from vllm.multimodal.video_decoders.pynvvideocodec import (
     PYNVVIDEOCODEC_DECODE_BATCH_SIZE,
     PYNVVIDEOCODEC_DECODER_CACHE_SIZE,
+    PYNVVIDEOCODEC_DECODER_CAPS_CACHE_SIZE,
     PyNvVideoCodecDecoderSlot,
     PyNvVideoCodecVideoBackendMixin,
+    _get_decoder_caps,
     _pynv_decoder_pool,
 )
 from vllm.platforms import current_platform
@@ -876,6 +878,54 @@ def GetDecoderCaps(device_index, codec, chroma_format, bit_depth):
     ) == (4096, 4096)
 
 
+def test_pynvvideocodec_caches_decoder_caps_per_stream_format():
+    calls = []
+
+    class FakeDemuxer:
+        def __init__(self, codec):
+            self.codec = codec
+
+        def Width(self):
+            return 1280
+
+        def Height(self):
+            return 720
+
+        def GetNvCodecId(self):
+            return self.codec
+
+        def ChromaFormat(self):
+            return "420"
+
+        def BitDepth(self):
+            return 8
+
+    class FakeNvc:
+        @staticmethod
+        def CreateDemuxer(file_path):
+            return FakeDemuxer("hevc" if "hevc" in file_path else "h264")
+
+        @staticmethod
+        def GetDecoderCaps(device_index, codec, chroma_format, bit_depth):
+            calls.append(codec)
+            return {
+                "supported": 1,
+                "width_min": 48,
+                "width_max": 8192,
+                "height_min": 16,
+                "height_max": 8192,
+                "mb_num_max": 262144,
+            }
+
+    for path in ("a.mp4", "b.mp4", "c-hevc.mp4", "d-hevc.mp4"):
+        assert PyNvVideoCodecVideoBackendMixin._validate_decoder_caps(
+            path, FakeNvc
+        ) == (1280, 1280)
+    assert calls == ["h264", "hevc"]
+    cache_size = _get_decoder_caps.cache_info().maxsize
+    assert cache_size == PYNVVIDEOCODEC_DECODER_CAPS_CACHE_SIZE
+
+
 @pytest.mark.parametrize(
     ("max_width", "max_height"),
     [(3840, None), (None, 2160), (0, 2160), (3840, -1), (True, 2160)],
@@ -1287,12 +1337,13 @@ def unknown_fps(data) -> VideoSourceMetadata:
 
 
 @pytest.mark.parametrize(
-    "model_repo, expected_loader_cls, hf_sample_kwargs",
+    "model_repo, expected_loader_cls, hf_sample_kwargs, appends_last_frame",
     [
         pytest.param(
             "allenai/Molmo2-4B",
             Molmo2VideoBackend,
             None,
+            False,
             marks=pytest.mark.skip(
                 reason="Video processor not aligned, investigate later.",
             ),
@@ -1302,18 +1353,21 @@ def unknown_fps(data) -> VideoSourceMetadata:
             "zai-org/GLM-4.1V-9B-Thinking",
             DynamicVideoBackend,
             None,
+            False,
             id="glm4v",
         ),
         pytest.param(
             "zai-org/GLM-4.6V-Flash",
             GLM46VVideoBackend,
             None,
+            False,
             id="glm46v",
         ),
         pytest.param(
             "Qwen/Qwen3-VL-4B-Instruct",
             Qwen3VLVideoBackend,
             None,
+            False,
             id="qwen3vl",
         ),
         # Qwen2-VL/Qwen2.5-VL ship no ``video_processor_type`` in their
@@ -1325,18 +1379,23 @@ def unknown_fps(data) -> VideoSourceMetadata:
             "Qwen/Qwen2-VL-7B-Instruct",
             Qwen2VLVideoBackend,
             {"fps": 2},
+            False,
             id="qwen2vl",
         ),
         pytest.param(
             "Qwen/Qwen2.5-VL-7B-Instruct",
             Qwen2VLVideoBackend,
             {"fps": 2},
+            False,
             id="qwen2_5_vl",
         ),
         pytest.param(
             "MiniMaxAI/MiniMax-M3",
             MiniMaxM3VideoBackend,
             None,
+            # The MiniMax-M3 sampler also keeps the final frame, which the
+            # Hub ``sample_frames`` drops.
+            True,
             id="minimax_m3_vl",
         ),
     ],
@@ -1345,6 +1404,7 @@ def test_video_processor_from_model_repo(
     model_repo: str,
     expected_loader_cls: type,
     hf_sample_kwargs: dict[str, int | float] | None,
+    appends_last_frame: bool,
 ):
     """Test that a model repo resolves to the correct video loader backend.
 
@@ -1390,6 +1450,9 @@ def test_video_processor_from_model_repo(
             duration=vllm_meta["duration"],
         )
         hf_indices = processor.sample_frames(hf_metadata, **(hf_sample_kwargs or {}))
+        last_frame = vllm_meta["total_num_frames"] - 1
+        if appends_last_frame and hf_indices[-1] != last_frame:
+            hf_indices = np.append(hf_indices, last_frame)
         vllm_indices = np.array(vllm_meta["frames_indices"])
         np.testing.assert_array_equal(
             hf_indices,
diff --git a/tests/parser/engine/test_deepseek_v4.py b/tests/parser/engine/test_deepseek_v4.py
index f5b66d9dc1cb..3bf67b2a759a 100644
--- a/tests/parser/engine/test_deepseek_v4.py
+++ b/tests/parser/engine/test_deepseek_v4.py
@@ -33,6 +33,7 @@
     DSML_TOOL_START_VARIANTS,
     DeepSeekV4Parser,
     _dsml_arg_converter,
+    _dsml_arg_value_extender,
     _unwrap_wrapper_args,
     deepseek_v4_config,
 )
@@ -1302,6 +1303,55 @@ def test_eos_not_leaked_when_reasoning_never_ends(self, chunk_size):
         assert output.tool_calls == []
 
 
+# ── Open string values stream without re-running the converter (#52846) ──
+
+
+class TestArgValueExtender:
+    def test_stops_at_split_param_close(self):
+        extend = _dsml_arg_value_extender(
+            "<" + _PARAM_OPEN.format(name="k", is_str="true") + "x"
+        )
+        assert extend is not None
+        assert [extend(d) for d in ['a"\n', "</｜DSML｜para", "meter>"]] == [
+            'a\\"\\n',
+            "</｜DSML｜para",
+            None,
+        ]
+
+    @pytest.mark.parametrize(
+        "raw_args",
+        [
+            "<" + _PARAM_OPEN.format(name="k", is_str="true") + "x" + _PARAM_CLOSE,
+            "<" + _PARAM_OPEN.format(name="k", is_str="false") + "1",
+            "<"
+            + _PARAM_OPEN.format(name="k", is_str="true")
+            + "x"
+            + _PARAM_CLOSE
+            + "<"
+            + _PARAM_OPEN.format(name="k", is_str="true")
+            + "y",
+        ],
+        ids=["closed_value", "non_string_value", "repeated_key"],
+    )
+    def test_no_extender_unless_open_string_value(self, raw_args):
+        assert _dsml_arg_value_extender(raw_args) is None
+
+    @pytest.mark.parametrize(
+        "partial_tag", ['<｜DSML｜parameter name="j', "<｜DSML｜parameter na"]
+    )
+    def test_holds_back_partial_tag_after_gt(self, partial_tag):
+        raw_args = (
+            "<" + _PARAM_OPEN.format(name="k", is_str="true") + "a > b" + partial_tag
+        )
+        assert json.loads(_dsml_arg_converter(raw_args, True)) == {"k": "a > b"}
+        assert _dsml_arg_value_extender(raw_args) is None
+
+    def test_arms_on_value_ending_in_newline(self):
+        raw_args = "<" + _PARAM_OPEN.format(name="k", is_str="true") + "def f():\n"
+        assert json.loads(_dsml_arg_converter(raw_args, True)) == {"k": "def f():\n"}
+        assert _dsml_arg_value_extender(raw_args) is not None
+
+
 # ── Malformed DSML markup must not leak into content (#51914) ─────────
 
 _WRAPPER_VARIANTS = list(DSML_TOOL_START_VARIANTS)
diff --git a/tests/parser/engine/test_deepseek_v41.py b/tests/parser/engine/test_deepseek_v41.py
index 0b78c4d61785..75ef0bf044c9 100644
--- a/tests/parser/engine/test_deepseek_v41.py
+++ b/tests/parser/engine/test_deepseek_v41.py
@@ -148,3 +148,52 @@ def test_python_argument_conversion_and_partial_values():
         "bare": 7,
         "text": "  a<b",
     }
+
+
+@pytest.mark.parametrize(
+    "partial_tag", ['<｜DSML｜ parameter name="j', "<｜DSML｜ parameter na"]
+)
+def test_partial_tag_after_gt_is_held_back(partial_tag):
+    config = deepseek_v41_config()
+    assert config.arg_converter is not None
+    assert config.arg_value_extender is not None
+    raw = '<｜DSML｜ parameter name="k" string="true">a > b' + partial_tag
+    assert json.loads(config.arg_converter(raw, True)) == {"k": "a > b"}
+    assert config.arg_value_extender(raw) is None
+
+
+def test_long_string_argument_streams_before_close():
+    # Regression for #52846: a string body has no ">" so it used to be
+    # buffered until the closing parameter tag arrived.
+    body = "A" * 4096
+    text = (
+        '<｜DSML｜ calls>\n<｜DSML｜ invoke name="emit">\n'
+        f'<｜DSML｜ parameter name="text" string="true">{body}'
+        "</｜DSML｜ parameter>\n</｜DSML｜ invoke>\n</｜DSML｜ calls>"
+    )
+    tools = [
+        {
+            "type": "function",
+            "function": {
+                "name": "emit",
+                "parameters": {
+                    "type": "object",
+                    "properties": {"text": {"type": "string"}},
+                },
+            },
+        }
+    ]
+    tokenizer, tokens = tokenizer_for(text, False)
+    parser = parser_for(tokenizer, {"thinking": False})
+    results = replay_streaming(
+        parser, tokens, chunk_size=32, finished_on_last=True, tools=tools
+    )
+    arg_chunks = [
+        tool_call.function.arguments
+        for result in results
+        if result
+        for tool_call in result.tool_calls or []
+        if tool_call.function and tool_call.function.arguments
+    ]
+    assert len(arg_chunks) > len(body) // 64
+    assert json.loads("".join(arg_chunks)) == {"text": body}
diff --git a/tests/parser/engine/test_glm47_moe.py b/tests/parser/engine/test_glm47_moe.py
index 7497c31dc5b7..271725308f99 100644
--- a/tests/parser/engine/test_glm47_moe.py
+++ b/tests/parser/engine/test_glm47_moe.py
@@ -18,6 +18,7 @@
 from vllm.entrypoints.openai.chat_completion.protocol import (
     ChatCompletionToolsParam,
 )
+from vllm.parser import glm47_moe
 from vllm.parser.glm47_moe import (
     THINK_END,
     THINK_START,
@@ -140,6 +141,15 @@ def test_thinking_disabled(self, no_thinking_parser):
 CALL_TAIL = f"</arg_value>\n{TOOL_CALL_END}"
 # Escapes, non-ASCII, and tag-like text that must stay inside the value.
 CONTENT_LINE = 'if a < b > c: print("x\\y")  # </div> </arg_ <arg_key>\t\u00e9\n'
+# A model stuck in a loop repeats closed argument pairs inside one call.
+LOOP_HEAD = (
+    f"{THINK_END}{TOOL_CALL_START}Write\n"
+    "<arg_key>file_path</arg_key><arg_value>/repo/app.py</arg_value>\n"
+)
+LOOP_PAIRS = (
+    "<arg_key>content</arg_key><arg_value>a > b</arg_value>\n"
+    "<arg_key>line_count</arg_key><arg_value>1</arg_value>\n"
+)
 
 
 def _pieces(text: str, size: int) -> list[str]:
@@ -210,6 +220,91 @@ def counting_converter(raw_args: str, partial: bool) -> str:
 
         assert streaming_conversions(lines=40) == streaming_conversions(lines=400)
 
+    def test_replaced_value_conversions_do_not_grow(self, mock_request):
+        """Once a repeated key's new value diverges from the streamed text,
+        deltas that only extend it skip the converter."""
+        mock_request.tools = [WRITE_TOOL]
+        head = CALL_HEAD + "x</arg_value><arg_key>content</arg_key><arg_value>"
+
+        def streaming_conversions(lines: int) -> int:
+            calls = 0
+            base = glm47_moe_config()
+            base_converter = base.arg_converter
+            assert base_converter is not None
+
+            def counting_converter(raw_args: str, partial: bool) -> str:
+                nonlocal calls
+                if partial:
+                    calls += 1
+                return base_converter(raw_args, partial)
+
+            parser = Glm47MoeParser(
+                make_mock_tokenizer(VOCAB),
+                parser_engine_config=dataclasses.replace(
+                    base, arg_converter=counting_converter
+                ),
+            )
+            chunks = [
+                *_pieces(head, 4),
+                *_pieces(CONTENT_LINE * lines, 4),
+                *_pieces(CALL_TAIL, 4),
+            ]
+            _stream_arguments(parser, mock_request, chunks)
+            return calls
+
+        assert streaming_conversions(lines=40) == streaming_conversions(lines=400)
+
+    @pytest.mark.parametrize("size", [1, 7])
+    def test_repeated_closed_args_stream_non_streaming_arguments(
+        self, mock_request, size
+    ):
+        """Covers a duplicate open key and a second call sharing the first pair."""
+        mock_request.tools = [WRITE_TOOL]
+        text = (
+            LOOP_HEAD
+            + LOOP_PAIRS * 20
+            + "<arg_key>content</arg_key><arg_value>a > b"
+            + CALL_TAIL
+            + LOOP_HEAD.removeprefix(THINK_END)
+            + TOOL_CALL_END
+        )
+
+        streamed = _stream_arguments(
+            Glm47MoeParser(make_mock_tokenizer(VOCAB)),
+            mock_request,
+            _pieces(text, size),
+        )
+        result = Glm47MoeParser(make_mock_tokenizer(VOCAB)).extract_tool_calls(
+            text, mock_request
+        )
+
+        assert len(result.tool_calls) == 2
+        assert streamed == "".join(c.function.arguments for c in result.tool_calls)
+
+    def test_closed_pairs_are_parsed_once(self, mock_request, monkeypatch):
+        """Streaming cost must not grow with the closed pairs already seen."""
+        mock_request.tools = [WRITE_TOOL]
+        arg_re = glm47_moe._ARG_RE
+        matches = 0
+
+        class CountingPattern:
+            def __getattr__(self, name):
+                return getattr(arg_re, name)
+
+            def finditer(self, *args):
+                nonlocal matches
+                for match in arg_re.finditer(*args):
+                    matches += 1
+                    yield match
+
+        monkeypatch.setattr(glm47_moe, "_ARG_RE", CountingPattern())
+        text = LOOP_HEAD + LOOP_PAIRS * 200 + TOOL_CALL_END
+        _stream_arguments(
+            Glm47MoeParser(make_mock_tokenizer(VOCAB)), mock_request, _pieces(text, 4)
+        )
+
+        assert matches == len(arg_re.findall(text))
+
     def test_value_extender_stops_at_split_value_end(self):
         extend = _glm47_arg_value_extender("<arg_key>k</arg_key><arg_value>x")
         assert extend is not None
diff --git a/tests/parser/engine/test_qwen3.py b/tests/parser/engine/test_qwen3.py
index be4b33e3861f..902a9f61b4f9 100644
--- a/tests/parser/engine/test_qwen3.py
+++ b/tests/parser/engine/test_qwen3.py
@@ -819,6 +819,14 @@ def test_partial_value_with_angle_bracket_and_complete_param(self):
 # Escapes, non-ASCII, blank lines and tag-like text that must stay inside
 # the value, including prefixes of the closing and next-parameter tags.
 CONTENT_LINE = 'if a < b > c: print("x\\y")  # </p> </param <parameter\té\n\n'
+# A model stuck in a loop repeats closed parameters inside one call.
+LOOP_HEAD = (
+    "<tool_call>\n<function=write_file>\n<parameter=path>\n/repo/app.py\n</parameter>\n"
+)
+LOOP_PARAMS = (
+    "<parameter=content>\na > b\n</parameter>\n<parameter=mode>\n644\n</parameter>\n"
+)
+LOOP_TAIL = "</function>\n</tool_call>"
 
 
 class TestArgValueExtender:
@@ -908,6 +916,92 @@ def counting_converter(raw_args: str, partial: bool) -> str:
 
         assert streaming_conversions(lines=20) == streaming_conversions(lines=200)
 
+    @pytest.mark.parametrize("size", [1, 7])
+    def test_repeated_closed_params_stream_non_streaming_arguments(
+        self, mock_tokenizer, mock_request, write_tool, size
+    ):
+        """Covers a duplicate open key and a second call sharing the first pair."""
+        from vllm.parser.qwen3 import Qwen3Parser
+
+        mock_request.tools = [write_tool]
+        text = (
+            LOOP_HEAD
+            + LOOP_PARAMS * 20
+            + "<parameter=content>\na > b"
+            + WRITE_CALL_TAIL
+            + LOOP_HEAD
+            + LOOP_TAIL
+        )
+
+        def make_parser():
+            return Qwen3Parser(
+                mock_tokenizer, chat_template_kwargs={"enable_thinking": False}
+            )
+
+        results = simulate_tool_streaming(
+            make_parser(),
+            mock_request,
+            [text[i : i + size] for i in range(0, len(text), size)],
+        )
+        result = make_parser().extract_tool_calls(text, mock_request)
+
+        assert len(result.tool_calls) == 2
+        assert collect_tool_arguments(results) == "".join(
+            c.function.arguments for c in result.tool_calls
+        )
+
+    def test_closed_params_are_parsed_once(
+        self, mock_tokenizer, mock_request, write_tool, monkeypatch
+    ):
+        """Streaming cost must not grow with the closed parameters already seen."""
+        from vllm.parser import qwen3
+
+        mock_request.tools = [write_tool]
+        param_re = qwen3._PARAM_RE
+        matches = 0
+
+        class CountingPattern:
+            def __getattr__(self, name):
+                return getattr(param_re, name)
+
+            def finditer(self, *args):
+                nonlocal matches
+                for match in param_re.finditer(*args):
+                    matches += 1
+                    yield match
+
+            def sub(self, repl, *args):
+                def count(match):
+                    nonlocal matches
+                    matches += 1
+                    return repl
+
+                return param_re.sub(count, *args)
+
+        monkeypatch.setattr(qwen3, "_PARAM_RE", CountingPattern())
+        text = LOOP_HEAD + LOOP_PARAMS * 200 + LOOP_TAIL
+        simulate_tool_streaming(
+            qwen3.Qwen3Parser(
+                mock_tokenizer, chat_template_kwargs={"enable_thinking": False}
+            ),
+            mock_request,
+            [text[i : i + 4] for i in range(0, len(text), 4)],
+        )
+
+        assert matches == len(param_re.findall(text))
+
+    def test_value_closed_by_next_parameter_needs_that_parameter(self):
+        from vllm.parser.qwen3 import _ClosedParams, _qwen3_arg_converter
+
+        closed = _ClosedParams()
+        _qwen3_arg_converter(
+            "<parameter=path>\n/a.py\n<parameter=mode>\n6", True, closed
+        )
+        text = "<parameter=path>\n/a.py\n.bak\n</parameter>"
+        assert _qwen3_arg_converter(text, False, closed) == _qwen3_arg_converter(
+            text, False
+        )
+
     def test_extender_follows_newline_trimming_and_split_end(self):
         from vllm.parser.qwen3 import _qwen3_arg_value_extender
 
diff --git a/tests/parser/mistral/test_tool_calls.py b/tests/parser/mistral/test_tool_calls.py
index e33685f0fc8a..9fd5eda323f0 100644
--- a/tests/parser/mistral/test_tool_calls.py
+++ b/tests/parser/mistral/test_tool_calls.py
@@ -10,6 +10,9 @@
 import regex as re
 from mistral_common.protocol.instruct.messages import AssistantMessage
 from mistral_common.protocol.instruct.request import InstructRequest
+from mistral_common.protocol.instruct.request import (
+    ResponseFormats as MistralResponseFormats,
+)
 from mistral_common.protocol.instruct.tool_calls import (
     FunctionCall,
     ToolCall,
@@ -35,10 +38,11 @@
 from vllm.entrypoints.openai.chat_completion.protocol import (
     ChatCompletionRequest,
 )
+from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
 from vllm.exceptions import VLLMValidationError
 from vllm.parser.engine.events import EventType
 from vllm.parser.engine.streaming_parser_engine import StreamingParserEngine
-from vllm.parser.mistral import _DEFAULT_JSON_SCHEMA, MistralParser, mistral_config
+from vllm.parser.mistral import MistralParser, mistral_config
 from vllm.parser.parser_manager import ParserManager
 from vllm.sampling_params import StructuredOutputsParams
 from vllm.tokenizers import TokenizerLike, get_tokenizer
@@ -1814,13 +1818,19 @@ def _make_request(**kwargs) -> ChatCompletionRequest:
     return ChatCompletionRequest(**defaults)
 
 
+def _adjusted_grammar(parser: MistralToolParser, request: ChatCompletionRequest) -> str:
+    structured_outputs = parser.adjust_request(request).structured_outputs
+    assert structured_outputs is not None
+    assert isinstance(structured_outputs.grammar, str)
+    return structured_outputs.grammar
+
+
 @pytest.mark.parametrize(
-    "request_kwargs,expected_mode,expected_parallel",
+    "request_kwargs,expected_tool_choice,expected_parallel",
     [
         ({"tool_choice": "auto"}, MistralToolChoiceEnum.auto, True),
         ({"tool_choice": "none"}, MistralToolChoiceEnum.none, True),
         ({"tool_choice": "required"}, MistralToolChoiceEnum.required, True),
-        ({"tool_choice": None, "tools": None}, MistralToolChoiceEnum.auto, True),
         (
             {
                 "tool_choice": {
@@ -1848,7 +1858,6 @@ def _make_request(**kwargs) -> ChatCompletionRequest:
         "auto",
         "none",
         "required",
-        "null_tool_choice",
         "named_tool_choice",
         "parallel_false",
         "response_format_text",
@@ -1857,7 +1866,7 @@ def _make_request(**kwargs) -> ChatCompletionRequest:
 def test_adjust_request_grammar_factory(
     mistral_tool_parser: MistralToolParser,
     request_kwargs: dict,
-    expected_mode: MistralToolChoice,
+    expected_tool_choice: MistralToolChoice,
     expected_parallel: bool,
 ) -> None:
     request = _make_request(**request_kwargs)
@@ -1865,17 +1874,20 @@ def test_adjust_request_grammar_factory(
 
     with patch.object(
         factory,
-        "get_lark_from_jinja",
-        wraps=factory.get_lark_from_jinja,
-    ) as mock_get_lark:
+        "get_grammar_for_chat_request",
+        wraps=factory.get_grammar_for_chat_request,
+    ) as mock_get_grammar:
         result = mistral_tool_parser.adjust_request(request)
 
-        mock_get_lark.assert_called_once()
-        call_kwargs = mock_get_lark.call_args
+        mock_get_grammar.assert_called_once()
+        mistral_request = mock_get_grammar.call_args.kwargs["request"]
 
-        assert call_kwargs.kwargs["mode"] == expected_mode
-        assert call_kwargs.kwargs["json_schema"] is None
-        assert call_kwargs.kwargs["parallel_tool_calls"] == expected_parallel
+        assert mistral_request.tool_choice == expected_tool_choice
+        assert mistral_request.parallel_tool_calls == expected_parallel
+        assert mistral_request.response_format.type == MistralResponseFormats.text
+        if request.tools is not None:
+            assert mistral_request.tools is not None
+            assert len(mistral_request.tools) == len(request.tools)
 
     assert result.structured_outputs is not None
     assert isinstance(result.structured_outputs.grammar, str)
@@ -1914,6 +1926,28 @@ def test_adjust_request_non_mistral_tokenizer(
     assert result.skip_special_tokens is expected_skip
 
 
+@pytest.mark.parametrize("responses", [True, False])
+def test_adjust_request_mistral_tokenizer_keeps_skip_special_tokens(
+    mistral_tool_parser: MistralToolParser, responses: bool
+) -> None:
+    # MistralTokenizer keeps the parser's markers with skip_special_tokens=True,
+    # and forcing False corrupts byte-fallback text (e.g. Korean).
+    request = (
+        ResponsesRequest(
+            model="m", input="hi", tools=[{"type": "function", "name": "f"}]
+        )
+        if responses
+        else _make_request()
+    )
+    with patch.object(
+        type(mistral_tool_parser.model_tokenizer),
+        "supports_grammar",
+        new_callable=lambda: property(lambda self: False),
+    ):
+        result = mistral_tool_parser.adjust_request(request)
+    assert result.skip_special_tokens is None
+
+
 @pytest.mark.parametrize(
     "so_kwargs",
     [
@@ -1975,12 +2009,17 @@ def test_adjust_request_unsupported_response_format(
 
 
 @pytest.mark.parametrize(
-    "so_kwargs,expected_json_schema",
+    "so_kwargs,expected_response_format_type,expected_custom_schema",
     [
-        ({"json_object": True}, _DEFAULT_JSON_SCHEMA),
-        ({"json": '{"type": "object"}'}, {"type": "object"}),
+        ({"json_object": True}, MistralResponseFormats.json_object, None),
+        (
+            {"json": '{"type": "object"}'},
+            MistralResponseFormats.json_schema,
+            {"type": "object"},
+        ),
         (
             {"json": {"type": "object", "properties": {"x": {"type": "integer"}}}},
+            MistralResponseFormats.json_schema,
             {"type": "object", "properties": {"x": {"type": "integer"}}},
         ),
     ],
@@ -1989,7 +2028,8 @@ def test_adjust_request_unsupported_response_format(
 def test_adjust_request_structured_outputs_generates_grammar(
     mistral_tool_parser: MistralToolParser,
     so_kwargs: dict,
-    expected_json_schema: str,
+    expected_response_format_type: MistralResponseFormats,
+    expected_custom_schema: dict | None,
 ) -> None:
     request = _make_request(
         structured_outputs=StructuredOutputsParams(**so_kwargs),
@@ -1998,13 +2038,21 @@ def test_adjust_request_structured_outputs_generates_grammar(
 
     with patch.object(
         factory,
-        "get_lark_from_jinja",
-        wraps=factory.get_lark_from_jinja,
-    ) as mock_get_lark:
+        "get_grammar_for_chat_request",
+        wraps=factory.get_grammar_for_chat_request,
+    ) as mock_get_grammar:
         result = mistral_tool_parser.adjust_request(request)
 
-        mock_get_lark.assert_called_once()
-        assert mock_get_lark.call_args.kwargs["json_schema"] == expected_json_schema
+        mock_get_grammar.assert_called_once()
+        mistral_request = mock_get_grammar.call_args.kwargs["request"]
+        assert mistral_request.response_format.type == expected_response_format_type
+        if expected_custom_schema is not None:
+            assert mistral_request.response_format.json_schema is not None
+            assert (
+                mistral_request.response_format.json_schema.custom_schema
+                == expected_custom_schema
+            )
+            assert mistral_request.response_format.json_schema.strict is True
 
     assert result.structured_outputs is not None
     assert isinstance(result.structured_outputs.grammar, str)
@@ -2024,62 +2072,127 @@ def test_adjust_request_deeply_nested_json_string_rejected(
         mistral_tool_parser.adjust_request(request)
 
 
+_RESPONSE_SCHEMA = {"type": "object", "properties": {"x": {"type": "integer"}}}
+
+
 @pytest.mark.parametrize(
-    "response_format_kwargs,expected_json_schema",
+    "response_format_kwargs,expected_response_format_type,expected_custom_schema,expected_strict",
     [
-        ({"type": "json_object"}, _DEFAULT_JSON_SCHEMA),
+        ({"type": "json_object"}, MistralResponseFormats.json_object, None, None),
+        (
+            {
+                "type": "json_schema",
+                "json_schema": {"name": "my_schema", "schema": _RESPONSE_SCHEMA},
+            },
+            MistralResponseFormats.json_schema,
+            _RESPONSE_SCHEMA,
+            False,
+        ),
         (
             {
                 "type": "json_schema",
                 "json_schema": {
                     "name": "my_schema",
-                    "schema": {
-                        "type": "object",
-                        "properties": {"x": {"type": "integer"}},
-                    },
+                    "schema": _RESPONSE_SCHEMA,
+                    "strict": True,
                 },
             },
-            {"type": "object", "properties": {"x": {"type": "integer"}}},
+            MistralResponseFormats.json_schema,
+            _RESPONSE_SCHEMA,
+            True,
         ),
     ],
-    ids=["json_object", "json_schema_with_schema"],
+    ids=["json_object", "json_schema", "json_schema_strict"],
 )
 def test_adjust_request_response_format_generates_grammar(
     mistral_tool_parser: MistralToolParser,
     response_format_kwargs: dict,
-    expected_json_schema: str,
+    expected_response_format_type: MistralResponseFormats,
+    expected_custom_schema: dict | None,
+    expected_strict: bool | None,
 ) -> None:
+    """The client's response_format, `strict` included, reaches mistral-common
+    unchanged; the grammar then follows mistral-common's defaults."""
     request = _make_request(response_format=response_format_kwargs)
     factory = mistral_tool_parser.model_tokenizer.grammar_factory
 
     with patch.object(
         factory,
-        "get_lark_from_jinja",
-        wraps=factory.get_lark_from_jinja,
-    ) as mock_get_lark:
+        "get_grammar_for_chat_request",
+        wraps=factory.get_grammar_for_chat_request,
+    ) as mock_get_grammar:
         result = mistral_tool_parser.adjust_request(request)
 
-        mock_get_lark.assert_called_once()
-        assert mock_get_lark.call_args.kwargs["json_schema"] == expected_json_schema
+        mock_get_grammar.assert_called_once()
+        mistral_request = mock_get_grammar.call_args.kwargs["request"]
+        assert mistral_request.response_format.type == expected_response_format_type
+        if expected_custom_schema is not None:
+            json_schema = mistral_request.response_format.json_schema
+            assert json_schema is not None
+            assert json_schema.custom_schema == expected_custom_schema
+            assert json_schema.name == "my_schema"
+            assert json_schema.strict is expected_strict
 
+    assert isinstance(result, ChatCompletionRequest)
+    assert result.response_format is None
     assert result.structured_outputs is not None
     assert isinstance(result.structured_outputs.grammar, str)
     assert len(result.structured_outputs.grammar) > 0
 
 
 @pytest.mark.parametrize(
-    "tool_choice, expected_method, not_called_method",
+    "strict, schema_enforced",
+    [(None, False), (False, False), (True, True)],
+    ids=["strict_unset", "strict_false", "strict_true"],
+)
+def test_adjust_request_response_format_strict_controls_schema(
+    mistral_tool_parser: MistralToolParser,
+    strict: bool | None,
+    schema_enforced: bool,
+) -> None:
+    """mistral-common enforces a response_format schema only when `strict` is
+    true; otherwise the grammar accepts any JSON object or array."""
+    json_schema: dict = {"name": "my_schema", "schema": _RESPONSE_SCHEMA}
+    if strict is not None:
+        json_schema["strict"] = strict
+    request = _make_request(
+        tools=None,
+        tool_choice=None,
+        response_format={"type": "json_schema", "json_schema": json_schema},
+    )
+    grammar = _adjusted_grammar(mistral_tool_parser, request)
+
+    assert (f"%json {json.dumps(_RESPONSE_SCHEMA)}" in grammar) is schema_enforced
+
+
+def test_adjust_request_invalid_response_format_is_validation_error(
+    mistral_tool_parser: MistralToolParser,
+) -> None:
+    """A schema mistral-common cannot use is a client error, not a crash."""
+    request = _make_request(
+        tools=None,
+        tool_choice=None,
+        response_format={
+            "type": "json_schema",
+            "json_schema": {"name": "my_schema", "schema": {"type": 1}},
+        },
+    )
+    with pytest.raises(VLLMValidationError):
+        mistral_tool_parser.adjust_request(request)
+
+
+@pytest.mark.parametrize(
+    "tool_choice, expected_tool_choice",
     [
-        ("none", "get_lark_for_json_schema", None),
-        ("auto", "get_lark_from_jinja", "get_lark_for_json_schema"),
+        ("none", MistralToolChoiceEnum.none),
+        ("auto", MistralToolChoiceEnum.auto),
     ],
-    ids=["none_uses_json_schema_factory", "auto_uses_jinja_factory"],
+    ids=["none_with_json_schema", "auto_with_json_schema"],
 )
 def test_adjust_request_tool_choice_with_json_schema_factory_routing(
     mistral_tool_parser: MistralToolParser,
     tool_choice: str,
-    expected_method: str,
-    not_called_method: str | None,
+    expected_tool_choice: MistralToolChoice,
 ) -> None:
     request = _make_request(
         tool_choice=tool_choice,
@@ -2087,37 +2200,83 @@ def test_adjust_request_tool_choice_with_json_schema_factory_routing(
     )
     factory = mistral_tool_parser.model_tokenizer.grammar_factory
 
-    patches = {
-        expected_method: patch.object(
-            factory,
-            expected_method,
-            wraps=getattr(factory, expected_method),
-        ),
-    }
-    if not_called_method:
-        patches[not_called_method] = patch.object(
-            factory,
-            not_called_method,
-            wraps=getattr(factory, not_called_method),
-        )
-
-    with patches[expected_method] as mock_expected:
-        ctx = patches[not_called_method] if not_called_method else None
-        if ctx:
-            with ctx as mock_not_called:
-                result = mistral_tool_parser.adjust_request(request)
-                mock_not_called.assert_not_called()
-        else:
-            result = mistral_tool_parser.adjust_request(request)
+    with patch.object(
+        factory,
+        "get_grammar_for_chat_request",
+        wraps=factory.get_grammar_for_chat_request,
+    ) as mock_get_grammar:
+        result = mistral_tool_parser.adjust_request(request)
 
-        mock_expected.assert_called_once()
-        assert mock_expected.call_args.kwargs["json_schema"] == {"type": "object"}
+        mock_get_grammar.assert_called_once()
+        mistral_request = mock_get_grammar.call_args.kwargs["request"]
+        assert mistral_request.tool_choice == expected_tool_choice
+        assert (
+            mistral_request.response_format.type == MistralResponseFormats.json_schema
+        )
+        assert mistral_request.response_format.json_schema is not None
+        assert mistral_request.response_format.json_schema.custom_schema == {
+            "type": "object"
+        }
 
     assert result.structured_outputs is not None
     assert isinstance(result.structured_outputs.grammar, str)
     assert len(result.structured_outputs.grammar) > 0
 
 
+@pytest.mark.parametrize(
+    "tool_choice, pinned_name",
+    [
+        ("required", None),
+        ({"type": "function", "function": {"name": "get_weather"}}, "get_weather"),
+    ],
+    ids=["required", "named"],
+)
+def test_adjust_request_forced_tool_call_grammar(
+    mistral_tool_parser: MistralToolParser,
+    tool_choice: str | dict,
+    pinned_name: str | None,
+) -> None:
+    """A forced tool call allows no text before the call, and a named choice
+    pins the tool name."""
+    grammar = _adjusted_grammar(
+        mistral_tool_parser, _make_request(tool_choice=tool_choice)
+    )
+
+    assert "body: fcalls\n" in grammar
+    assert "content? fcalls" not in grammar
+    if pinned_name is not None:
+        assert f'"{pinned_name}"' in grammar
+        assert '"add"' not in grammar
+
+
+def test_adjust_request_json_schema_with_tools_auto_grammar(
+    mistral_tool_parser: MistralToolParser,
+) -> None:
+    """JSON schema + tools + auto: the grammar is the schema or tool calls,
+    never free text."""
+    schema = {"type": "object", "properties": {"answer": {"type": "string"}}}
+    request = _make_request(
+        tool_choice="auto", structured_outputs=StructuredOutputsParams(json=schema)
+    )
+    grammar = _adjusted_grammar(mistral_tool_parser, request)
+
+    assert f"%json {json.dumps(schema)}" in grammar
+    assert "body: fcalls\n" in grammar
+    assert "content? fcalls" not in grammar
+
+
+def test_adjust_request_text_auto_no_grammar(
+    mistral_tool_parser: MistralToolParser,
+) -> None:
+    """Text response with no tools and auto tool_choice needs no grammar, so
+    the engine must not treat the request as parser-constrained."""
+    request = _make_request(tools=None, tool_choice=None)
+    result = mistral_tool_parser.adjust_request(request)
+    assert isinstance(result, ChatCompletionRequest)
+    assert result.structured_outputs is None
+    assert result._grammar_from_parser is False
+
+
 def test_grammar_from_parser_default_false() -> None:
     request = _make_request()
     assert request._grammar_from_parser is False
diff --git a/tests/plugins_tests/bitsandbytes/test_bitsandbytes.py b/tests/plugins_tests/bitsandbytes/test_bitsandbytes.py
index 92803d93151d..084eba351be6 100644
--- a/tests/plugins_tests/bitsandbytes/test_bitsandbytes.py
+++ b/tests/plugins_tests/bitsandbytes/test_bitsandbytes.py
@@ -35,7 +35,8 @@
 ]
 
 models_4bit_to_moe_test = [
-    ("allenai/OLMoE-1B-7B-0125-Instruct", "quantize moe model inflight"),
+    ("ibm-granite/granite-3.0-1b-a400m-base", "quantize moe model inflight"),
+    ("allenai/OLMoE-1B-7B-0125-Instruct", "quantize transformers moe model inflight"),
 ]
 
 models_pre_quant_4bit_to_test = [
diff --git a/tests/quantization/test_auto_gptq.py b/tests/quantization/test_auto_gptq.py
index 4f60a974caed..90512e9d837e 100644
--- a/tests/quantization/test_auto_gptq.py
+++ b/tests/quantization/test_auto_gptq.py
@@ -162,12 +162,14 @@ def test_routed_experts_loads_per_expert_biases():
     class Loader:
         quant_config = None
         quant_method = object()
-        _fused_shared_expert_quantizer = None
+        shared_expert_online_loader = None
         moe_config = SimpleNamespace(
             is_act_and_mul=True,
             tp_rank=0,
+            num_experts=1,
             moe_parallel_config=SimpleNamespace(tp_size=1),
         )
+        expert_map_manager = SimpleNamespace(num_fused_shared_experts=0)
         _get_hidden_dim = staticmethod(RoutedExperts._get_hidden_dim)
         _narrow_expert_data_for_padding = staticmethod(
             RoutedExperts._narrow_expert_data_for_padding
diff --git a/tests/quantization/test_auto_round.py b/tests/quantization/test_auto_round.py
index 8869b8d177ed..51b47eb0f19e 100644
--- a/tests/quantization/test_auto_round.py
+++ b/tests/quantization/test_auto_round.py
@@ -1053,7 +1053,7 @@ def process_weights_after_loading(self, layer) -> None:
     method.process_weights_after_loading(layer)
 
     assert layer.weight.data.data_ptr() == packed_data.data_ptr()
-    assert not hasattr(layer, "weight_packed")
+    assert layer.weight_packed is layer.weight
     assert captured["processed_layer"] is layer
 
 
@@ -1141,8 +1141,8 @@ def test_inc_mxfp4_moe_method_preserves_checkpoint_packing(
 
     assert layer.w13_weight.data.data_ptr() == w13_packed_data.data_ptr()
     assert layer.w2_weight.data.data_ptr() == w2_packed_data.data_ptr()
-    assert not hasattr(layer, "w13_weight_packed")
-    assert not hasattr(layer, "w2_weight_packed")
+    assert layer.w13_weight_packed is layer.w13_weight
+    assert layer.w2_weight_packed is layer.w2_weight
     assert captured["quant_config_kwargs"]["w1_scale"] is layer.w13_weight_scale
     assert captured["quant_config_kwargs"]["w2_scale"] is layer.w2_weight_scale
     assert captured["kernel_kwargs"]["moe_quant_config"] is expected_quant_config
diff --git a/tests/quantization/test_blackwell_moe.py b/tests/quantization/test_blackwell_moe.py
index e249a7ad91ca..15fd10028cb8 100644
--- a/tests/quantization/test_blackwell_moe.py
+++ b/tests/quantization/test_blackwell_moe.py
@@ -387,3 +387,87 @@ def check_scales(model):
             for step in completion.logprobs
             for value in step.values()
         )
+
+
+@pytest.mark.parametrize("nvfp4", [False, True])
+def test_muse_megamoe_dummy_loading(tmp_path, monkeypatch, vllm_runner, nvfp4):
+    """Dummy-loaded Muse experts generate finite outputs in both weight formats."""
+    import math
+
+    from vllm import SamplingParams
+    from vllm.model_executor.models import ModelRegistry
+
+    monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
+    monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
+    ModelRegistry.register_model(
+        "MuseSpark1xForCausalLM",
+        "vllm.models.muse_spark.nvidia.model:MuseSpark1xForCausalLM",
+    )
+    config = dict(
+        architectures=["MuseSpark1xForCausalLM"],
+        model_type="muse_spark_text",
+        vocab_size=128,
+        hidden_size=512,
+        intermediate_size=512,
+        moe_hidden_size=512,
+        moe_intermediate_size=512,
+        num_hidden_layers=1,
+        num_attention_heads=8,
+        num_key_value_heads=8,
+        head_dim=64,
+        num_local_experts=8,
+        num_experts_per_tok=2,
+        num_shared_experts=0,
+        max_position_embeddings=64,
+        hidden_act="silu",
+        attention_bias=False,
+        layer_types=["full_attention"],
+        mlp_layer_types=["sparse"],
+        layer_rope_theta=[10000.0],
+        rope_parameters={"rope_type": "default", "rope_theta": 10000.0},
+        rms_norm_eps=1e-5,
+        post_norm_eps=1e-8,
+        residual_gate_temperature=0.3,
+        final_logit_softcapping=None,
+        output_multiplier=0.5,
+        routed_scaling_factor=1.0,
+        qk_scale_factor=1.0,
+        tie_word_embeddings=False,
+        torch_dtype="bfloat16",
+    )
+    if nvfp4:
+        config["quantization_config"] = {
+            "quant_method": "modelopt",
+            "quant_algo": "NVFP4",
+            "group_size": 16,
+            "kv_cache_quant_algo": None,
+            "exclude_modules": ["lm_head", "model.layers.*.self_attn.*"],
+        }
+    (tmp_path / "config.json").write_text(json.dumps(config))
+
+    with vllm_runner(
+        str(tmp_path),
+        dtype="bfloat16",
+        load_format="dummy",
+        moe_backend="deep_gemm_mega_moe",
+        enable_expert_parallel=True,
+        skip_tokenizer_init=True,
+        enforce_eager=True,
+        enable_flashinfer_autotune=False,
+        kernel_config={"enable_jit_warmup": False, "enable_cutedsl_warmup": False},
+        max_model_len=32,
+        max_num_seqs=1,
+        max_num_batched_tokens=32,
+        kv_cache_memory_bytes=16 << 20,
+    ) as llm:
+        outputs = llm.get_llm().generate(
+            [{"prompt_token_ids": [1, 2, 3]}],
+            SamplingParams(temperature=0, max_tokens=2, ignore_eos=True, logprobs=1),
+        )
+        completion = outputs[0].outputs[0]
+        assert len(completion.token_ids) == 2
+        assert all(
+            math.isfinite(value.logprob)
+            for step in completion.logprobs
+            for value in step.values()
+        )
diff --git a/tests/quantization/test_compressed_tensors.py b/tests/quantization/test_compressed_tensors.py
index 8120cc35f128..a2b7ec315307 100644
--- a/tests/quantization/test_compressed_tensors.py
+++ b/tests/quantization/test_compressed_tensors.py
@@ -1109,7 +1109,7 @@ def test_mxfp4_moe_with_fp8_inputs_shares_native_mxfp4_method(monkeypatch):
     layer.quant_method.process_weights_after_loading(layer)
 
     assert layer.w13_weight is w13
-    assert not hasattr(layer, "w13_weight_packed")
+    assert layer.w13_weight_packed is w13
 
 
 @pytest.mark.skipif(
diff --git a/tests/quantization/test_fp8.py b/tests/quantization/test_fp8.py
index 0b59a068cafa..0638c752c512 100644
--- a/tests/quantization/test_fp8.py
+++ b/tests/quantization/test_fp8.py
@@ -659,14 +659,15 @@ def test_static_fp8_moe_input_scales_remain_scalar() -> None:
 )
 @pytest.mark.parametrize("model_id", MODELS)
 @pytest.mark.parametrize(
-    "force_marlin", [True, False] if current_platform.is_cuda() else [False]
+    "linear_backend",
+    ["auto", "marlin", "humming"] if current_platform.is_cuda() else ["auto"],
 )
 @pytest.mark.parametrize(
     "use_rocm_aiter", [True, False] if current_platform.is_rocm() else [False]
 )
 def test_model_load_and_run(
     model_id: str,
-    force_marlin: bool,
+    linear_backend: str,
     use_rocm_aiter: bool,
     monkeypatch,
     dist_init,
@@ -676,8 +677,8 @@ def test_model_load_and_run(
         monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
 
     kernel_config = KernelConfig(
-        linear_backend="marlin" if force_marlin else "auto",
-        moe_backend="marlin" if force_marlin else "auto",
+        linear_backend=linear_backend,
+        moe_backend="marlin" if linear_backend == "marlin" else "auto",
     )
     model, vllm_config = load_model_without_vllm_runner(
         model_id,
diff --git a/tests/quantization/test_fp8_per_channel.py b/tests/quantization/test_fp8_per_channel.py
index d4f2fe264e6c..c0be76769ee1 100644
--- a/tests/quantization/test_fp8_per_channel.py
+++ b/tests/quantization/test_fp8_per_channel.py
@@ -9,6 +9,9 @@
 `pytest tests/quantization/test_fp8_per_channel.py --forked`.
 """
 
+from types import SimpleNamespace
+from unittest import mock
+
 import pytest
 import torch
 
@@ -28,10 +31,18 @@
     Fp8PtpcOnlineMoEMethod,
 )
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    get_fp8_min_max,
     kFp8StaticChannelSym,
 )
 from vllm.platforms import current_platform
 
+if current_platform.is_rocm():
+    from vllm.platforms.rocm import on_gfx942
+else:
+
+    def on_gfx942() -> bool:
+        return False
+
 
 def test_fp8_per_channel_shorthand_registered() -> None:
     """The `fp8_per_channel` CLI shorthand must resolve to a config that
@@ -56,6 +67,10 @@ def test_fp8_per_channel_shorthand_registered() -> None:
     not is_quant_method_supported("fp8"),
     reason="FP8 is not supported on this GPU type.",
 )
+@pytest.mark.skipif(
+    on_gfx942(),
+    reason="The dynamic per-token FP8 kernel can segfault on gfx942.",
+)
 def test_scaled_fp8_quant_per_channel_shape() -> None:
     """Verify the kernel call per-channel quant depends on: passing a 2D
     weight to `ops.scaled_fp8_quant` with `use_per_token_if_dynamic=True`
@@ -69,3 +84,63 @@ def test_scaled_fp8_quant_per_channel_shape() -> None:
     assert y.dtype == current_platform.fp8_dtype()
     assert s.shape == (96, 1)
     assert s.dtype == torch.float32
+
+
+@pytest.mark.skipif(not on_gfx942(), reason="Requires gfx942.")
+def test_gfx942_moe_weight_quant_avoids_dynamic_per_token_op(monkeypatch) -> None:
+    """gfx942 load-time MoE quantization uses the torch per-channel path."""
+    device = current_platform.device_type
+    torch.manual_seed(0)
+    w13_weight = torch.randn(2, 96, 128, device=device, dtype=torch.bfloat16)
+    w2_weight = torch.randn(2, 128, 48, device=device, dtype=torch.bfloat16)
+
+    method = object.__new__(Fp8PtpcOnlineMoEMethod)
+    method.moe = SimpleNamespace(tp_size=1)
+    layer = SimpleNamespace(local_num_experts=2)
+
+    def fail_dynamic_quant(*args, **kwargs):
+        raise AssertionError("gfx942 must not call dynamic per-token FP8 quant")
+
+    monkeypatch.setattr(ops, "scaled_fp8_quant", fail_dynamic_quant)
+    with (
+        mock.patch.object(
+            method,
+            "get_weights_for_quantization",
+            return_value=(w13_weight, w2_weight),
+        ) as get_weights,
+        mock.patch.object(method, "_setup_kernel") as setup_kernel,
+        mock.patch.object(
+            method, "release_requantization_source_weights"
+        ) as release_weights,
+    ):
+        method.process_weights_after_loading(layer)
+
+    get_weights.assert_called_once_with(layer)
+    setup_kernel.assert_called_once()
+    release_weights.assert_called_once_with(layer)
+    _, got_w13, got_w2, got_w13_scale, got_w2_scale = setup_kernel.call_args.args
+    assert setup_kernel.call_args.kwargs == {
+        "w13_input_scale": None,
+        "w2_input_scale": None,
+    }
+
+    fp8_min, fp8_max = get_fp8_min_max()
+
+    def reference(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        fp8_max_tensor = torch.tensor(
+            fp8_max, device=weight.device, dtype=torch.float32
+        )
+        scale = weight.abs().amax(dim=-1, keepdim=True).float() / fp8_max_tensor
+        scale.clamp_min_(1.0 / (fp8_max * 512))
+        quantized = (
+            (weight / scale).clamp(fp8_min, fp8_max).to(current_platform.fp8_dtype())
+        )
+        return quantized, scale
+
+    expected_w13, expected_w13_scale = reference(w13_weight)
+    expected_w2, expected_w2_scale = reference(w2_weight)
+    assert torch.equal(got_w13_scale, expected_w13_scale)
+    assert torch.equal(got_w13, expected_w13)
+    assert torch.equal(got_w2_scale, expected_w2_scale)
+    assert torch.equal(got_w2, expected_w2)
+    assert layer._already_called_process_weights_after_loading
diff --git a/tests/quantization/test_keep_param_objects.py b/tests/quantization/test_keep_param_objects.py
index 94b8110c3b7a..ffe15cc1f36b 100644
--- a/tests/quantization/test_keep_param_objects.py
+++ b/tests/quantization/test_keep_param_objects.py
@@ -10,7 +10,11 @@
 from vllm.config.model import ModelConfig
 from vllm.model_executor.layers.fused_moe import FusedMoEFactory
 from vllm.model_executor.layers.linear import ReplicatedLinear
+from vllm.model_executor.layers.quantization.auto_gptq import AutoGPTQConfig
 from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
+from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import (  # noqa: E501
+    CompressedTensorsConfig,
+)
 from vllm.model_executor.layers.quantization.fp8 import Fp8Config
 from vllm.model_executor.layers.quantization.modelopt import (
     ModelOptFp8Config,
@@ -25,6 +29,7 @@
 from vllm.model_executor.model_loader.reload.meta import materialize_meta_tensor
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 from vllm.platforms import current_platform
+from vllm.utils.import_utils import has_humming
 
 pytestmark = pytest.mark.skipif(
     not current_platform.is_device_capability_family(100),
@@ -46,14 +51,28 @@ def _load(layer: torch.nn.Module) -> None:
             param.data.copy_(torch.rand(param.shape) + 0.5)
 
 
+def _nbytes(param: torch.Tensor) -> int:
+    return param.numel() * param.element_size()
+
+
 def _dims(param: torch.Tensor) -> tuple:
     return tuple(
         getattr(param, k, None) for k in ("input_dim", "output_dim", "packed_dim")
     )
 
 
-def _linear(quant_config: QuantizationConfig, output_size: int = SIZE):
-    layer = ReplicatedLinear(SIZE, output_size, bias=False, quant_config=quant_config)
+def _linear(
+    quant_config: QuantizationConfig,
+    output_size: int = SIZE,
+    params_dtype: torch.dtype | None = None,
+):
+    layer = ReplicatedLinear(
+        SIZE,
+        output_size,
+        bias=False,
+        quant_config=quant_config,
+        params_dtype=params_dtype,
+    )
     return layer, layer.quant_method
 
 
@@ -102,7 +121,19 @@ def _online_moe():
     return layer, method
 
 
+def _ct(fmt: str, group: str, group_size: int, input_activations: bool):
+    args = dict(num_bits=4, type="float", strategy=group, group_size=group_size)
+    args |= dict(symmetric=True, dynamic=False)
+    acts = args | dict(dynamic="local") if input_activations else None
+    group_0 = dict(targets=["Linear"], weights=args, input_activations=acts)
+    return CompressedTensorsConfig.from_config(
+        dict(format=fmt, config_groups=dict(group_0=group_0), ignore=[])
+    )
+
+
 NVFP4 = dict(is_checkpoint_nvfp4_serialized=True, kv_cache_quant_algo=None)
+CT_NVFP4 = ("nvfp4-pack-quantized", "tensor_group", 16, True)
+CT_MXFP4 = ("mxfp4-pack-quantized", "group", 32, False)
 
 CASES: list[Callable] = [
     lambda: _linear(Fp8Config(weight_block_size=[128, 128])),
@@ -117,11 +148,18 @@ def _online_moe():
         output_size=SIZE + 64,
     ),
     lambda: _linear(ModelOptMxFp8Config.from_config({"quant_method": "mxfp8"})),
+    # The consumed per-tensor scales stay registered with their data.
+    lambda: _linear(ModelOptNvFp4Config(**NVFP4, exclude_modules=[])),
+    # The packed weights are also registered under their runtime names.
+    lambda: _linear(_ct(*CT_NVFP4)),
+    lambda: _linear(_ct(*CT_MXFP4)),
     _online_linear,
     lambda: _moe(None),
     lambda: _moe(Fp8Config(weight_block_size=[128, 128])),
     lambda: _moe(ModelOptMxFp8Config.from_config({"quant_method": "mxfp8"})),
     lambda: _moe(ModelOptNvFp4Config(**NVFP4, exclude_modules=[])),
+    lambda: _moe(_ct(*CT_NVFP4)),
+    lambda: _moe(_ct(*CT_MXFP4)),
     lambda: _moe(Mxfp4Config()),
     _online_moe,
 ]
@@ -129,34 +167,97 @@ def _online_moe():
     "fp8-block-linear",
     "modelopt-fp8-pb-wo-linear",
     "modelopt-mxfp8-linear",
+    "modelopt-nvfp4-linear",
+    "ct-nvfp4-linear",
+    "ct-mxfp4-linear",
     "online-mxfp4-linear",
     "unquantized-moe",
     "fp8-block-moe",
     "modelopt-mxfp8-moe",
     "modelopt-nvfp4-moe",
+    "ct-nvfp4-moe",
+    "ct-mxfp4-moe",
     "mxfp4-moe",
     "online-mxfp4-moe",
 ]
 
 
-@pytest.mark.parametrize("build", CASES, ids=IDS)
-def test_process_weights_after_loading_keeps_param_objects(
-    build: Callable, default_vllm_config, dist_init
-) -> None:
-    default_vllm_config.model_config = ModelConfig()
+def _check_keeps_param_objects(build: Callable, vllm_config) -> None:
+    vllm_config.model_config = ModelConfig()
     with torch.device("cuda"):
         layer, method = build()
         _load(layer)
         before = {
-            name: (param, set(param.__dict__), _dims(param))
+            name: (param, set(param.__dict__), _dims(param), _nbytes(param))
             for name, param in layer.named_parameters(recurse=False)
         }
         method.process_weights_after_loading(layer)
 
     after = dict(layer.named_parameters(recurse=False, remove_duplicate=False))
-    for name, (param, attrs, dims) in before.items():
+    for name, (param, attrs, dims, nbytes) in before.items():
         assert after.get(name) is param, name
         assert attrs <= param.__dict__.keys(), name
         # The loader shards by these on reload, so they keep the checkpoint layout.
         assert _dims(param) == dims, name
         assert callable(getattr(param, "weight_loader", None)), name
+        # Online quantization frees the high-precision source by design.
+        if not method.uses_meta_device:
+            kept = after[name]
+            offset = kept.storage_offset() * kept.element_size()
+            assert kept.untyped_storage().nbytes() - offset >= nbytes, name
+    # Two names share storage only as one param object.
+    objs = {id(p): p for p in after.values() if p.numel()}
+    ptrs = [p.untyped_storage().data_ptr() for p in objs.values()]
+    assert len(ptrs) == len(set(ptrs))
+
+
+@pytest.mark.parametrize("build", CASES, ids=IDS)
+def test_process_weights_after_loading_keeps_param_objects(
+    build: Callable, default_vllm_config, dist_init
+) -> None:
+    _check_keeps_param_objects(build, default_vllm_config)
+
+
+FORCED_CASES: list[Callable] = [
+    lambda: _linear(
+        ModelOptNvFp4Config(**NVFP4, exclude_modules=[]), params_dtype=torch.bfloat16
+    ),
+    lambda: _linear(_ct(*CT_NVFP4), params_dtype=torch.bfloat16),
+    lambda: _moe(ModelOptNvFp4Config(**NVFP4, exclude_modules=[])),
+    lambda: _moe(_ct(*CT_NVFP4)),
+]
+FORCED_IDS = [
+    "modelopt-nvfp4-linear",
+    "ct-nvfp4-linear",
+    "modelopt-nvfp4-moe",
+    "ct-nvfp4-moe",
+]
+
+
+@pytest.mark.parametrize("backend", ["marlin", "humming"])
+@pytest.mark.parametrize("build", FORCED_CASES, ids=FORCED_IDS)
+def test_forced_backend_keeps_param_objects(
+    build: Callable, backend: str, default_vllm_config, dist_init
+) -> None:
+    """Weight-only backends leave the input scales unused but registered."""
+    if backend == "humming" and not has_humming():
+        pytest.skip("humming is not installed")
+    default_vllm_config.kernel_config.linear_backend = backend
+    default_vllm_config.kernel_config.moe_backend = backend
+    _check_keeps_param_objects(build, default_vllm_config)
+
+
+def test_marlin_bias_keeps_its_address(default_vllm_config, dist_init) -> None:
+    """Captured CUDA graphs read the bias at its startup address (direct reload
+    writes the checkpoint bias there before rerunning post-load processing)."""
+    default_vllm_config.model_config = ModelConfig()
+    default_vllm_config.kernel_config.linear_backend = "marlin"
+    gptq = AutoGPTQConfig.from_config(
+        dict(bits=4, group_size=128, desc_act=False, sym=True)
+    )
+    with torch.device("cuda"):
+        layer = ReplicatedLinear(SIZE, SIZE, bias=True, quant_config=gptq)
+        _load(layer)
+        ptr = layer.bias.data_ptr()
+        layer.quant_method.process_weights_after_loading(layer)
+    assert layer.bias.data_ptr() == ptr
diff --git a/tests/quantization/test_modelopt.py b/tests/quantization/test_modelopt.py
index d489cb2b9acf..1cffd78714eb 100644
--- a/tests/quantization/test_modelopt.py
+++ b/tests/quantization/test_modelopt.py
@@ -154,6 +154,80 @@ def test_modelopt_mxfp8_preserves_per_row_checkpoint_scales(dist_init, monkeypat
     assert torch.equal(linear.weight_scale, scales)
 
 
+_M31_QUANT_CONFIG = {
+    "quant_method": "mxfp8",
+    "ignored_layers": ["lm_head", "model.layers.3.block_sparse_moe.gate"],
+    "group_size": 16,
+    "moe_quant_algo": "NVFP4",
+    "moe_quant_format": "nvfp4-pack-quantized",
+}
+
+
+def test_modelopt_mxfp8_routes_nvfp4_experts_to_compressed_tensors(monkeypatch):
+    """MiniMax-M3.1 keeps MXFP8 linears but ships compressed-tensors NVFP4
+    routed experts; loading them with the MXFP8 MoE method drops weights."""
+    ct_method = Mock()
+    monkeypatch.setattr(
+        "vllm.model_executor.layers.quantization.compressed_tensors."
+        "compressed_tensors_moe.compressed_tensors_moe_w4a4_nvfp4."
+        "CompressedTensorsW4A4Nvfp4MoEMethod",
+        ct_method,
+    )
+    config = ModelOptMxFp8Config.from_config(dict(_M31_QUANT_CONFIG))
+    layer = Mock(spec=RoutedExperts)
+    layer.moe_config = Mock()
+    prefix = "model.layers.3.block_sparse_moe.experts"
+
+    # MiniMax-M3.1 declares that it QAT'd the FC1 input per token, leaving a
+    # neutral 1.0 input_global_scale. A model that declares nothing keeps the
+    # checkpoint's per-tensor scale.
+    for declared, expected in (
+        (kNvfp4DynamicToken, kNvfp4DynamicToken),
+        (None, kNvfp4Dynamic),
+    ):
+        ct_method.reset_mock()
+        layer.moe_config.qat_activation_key = declared
+        assert config.get_quant_method(layer, prefix) is ct_method.return_value
+        ct_method.assert_called_once_with(
+            layer.moe_config, layer_name=prefix, activation_key=expected
+        )
+
+    plain = ModelOptMxFp8Config.from_config({"quant_method": "mxfp8"})
+    assert plain.moe_quant_algo is None
+
+
+def test_per_token_nvfp4_moe_backends_match_the_experts_that_implement_it():
+    """``make_nvfp4_moe_kernel`` dispatches unquantized inputs for a per-token
+    scheme and only forwards ``per_token_activation`` to this set, so it has to
+    be exactly the experts that take that argument. ``_supports_quant_scheme``
+    cannot be used for the check: several NVFP4 backends match on the weight
+    key alone and accept any activation key."""
+    import inspect
+
+    from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import (
+        PER_TOKEN_NVFP4_MOE_BACKENDS,
+        NvFp4MoeBackend,
+        backend_to_kernel_cls,
+    )
+
+    implementing = {
+        backend
+        for backend in NvFp4MoeBackend
+        for cls in backend_to_kernel_cls(backend)
+        if "per_token_activation" in inspect.signature(cls.__init__).parameters
+    }
+    assert implementing == PER_TOKEN_NVFP4_MOE_BACKENDS
+
+
+@pytest.mark.parametrize(
+    "override",
+    [{"moe_quant_algo": "MXFP4"}, {"moe_quant_format": "float-quantized"}],
+)
+def test_modelopt_mxfp8_rejects_unknown_expert_format(override):
+    with pytest.raises(ValueError, match="MoE quantization"):
+        ModelOptMxFp8Config.from_config({**_M31_QUANT_CONFIG, **override})
+
+
 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."""
@@ -1546,6 +1620,7 @@ def test_modelopt_block_quant_loads_expert_blocks_with_aligned_tp_slices(
     layer = torch.nn.Module()
     layer.layer_name = "model.layers.0.mlp.experts"
     layer.is_fused_checkpoint_transposed = False
+    layer.shared_expert_online_loader = None
     layer._orient_fused_weight = RoutedExperts._orient_fused_weight
     layer.get_expert_mapping = lambda **kwargs: (
         RoutedExperts.build_expert_params_mapping(
diff --git a/tests/quantization/test_moe_wna16.py b/tests/quantization/test_moe_wna16.py
index c295a195a5f3..c320a7e7e9be 100644
--- a/tests/quantization/test_moe_wna16.py
+++ b/tests/quantization/test_moe_wna16.py
@@ -384,8 +384,9 @@ def test_compressed_tensors_wna16_moe_converts_and_sets_up_humming_kernel():
     assert method.moe_kernel is not None
     assert set(layer.weight_schemas) == {"w13", "w2"}
     assert set(layer.humming_configs) == {"w13", "w2"}
-    assert not hasattr(layer, "w13_weight_packed")
-    assert not hasattr(layer, "w2_weight_packed")
+    # The converted weights are written to the loaded params (one object each).
+    assert layer.w13_weight_packed is layer.w13_weight
+    assert layer.w2_weight_packed is layer.w2_weight
     assert layer.w13_weight.dtype is torch.int32
     assert layer.w2_weight.dtype is torch.int32
 
diff --git a/tests/quantization/test_online.py b/tests/quantization/test_online.py
index b572bb425a6d..791fe5871ad1 100644
--- a/tests/quantization/test_online.py
+++ b/tests/quantization/test_online.py
@@ -881,7 +881,7 @@ def test_weight_loading_requantization_moe(
     for name, expected in tensors.items():
         # Public loading must preserve all packed weights and scale conventions.
         torch.testing.assert_close(getattr(layer, name).cpu(), expected, rtol=0, atol=0)
-    actual_w13, actual_w2 = source.dequantize_weight(layer)
+    ((_, actual_w13, actual_w2),) = source.dequantize_weight(layer)
 
     # Gate/up ordering and distinct global scales must survive source loading.
     torch.testing.assert_close(actual_w13.cpu(), expected_w13, rtol=0, atol=0)
@@ -921,9 +921,12 @@ def test_online_requantized_moe_padding(
         w13_num_shards=num_shards, is_act_and_mul=is_act_and_mul, has_bias=has_bias
     )
     method.requantization_source = Mock()
-    method.requantization_source.dequantize_weight.return_value = (w13, w2)
+    method.requantization_source.dequantize_weight.return_value = iter(
+        [(slice(0, 2), w13, w2)]
+    )
     layer = SimpleNamespace(
         local_num_experts=2,
+        expert_map_manager=SimpleNamespace(num_fused_shared_experts=0),
         moe_config=SimpleNamespace(
             hidden_dim_unpadded=16,
             intermediate_size_per_partition_unpadded=32,
@@ -1075,6 +1078,7 @@ def create_layer() -> tuple[torch.nn.Module, Any]:
         layer = torch.nn.Module()
         layer.local_num_experts = 2
         layer.num_experts = 2
+        layer.expert_map_manager = SimpleNamespace(num_fused_shared_experts=0)
         layer._expert_routing_tables = lambda: None
         layer.moe_config = SimpleNamespace(
             hidden_dim_unpadded=size,
@@ -2477,7 +2481,7 @@ def test_nvfp4_moe_dequantization_and_fp8_requantization_numerics(
         **checkpoint_tensors,
         params_dtype=torch.bfloat16,
     )
-    w13_dequantized, w2_dequantized = source.dequantize_weight(source_layer)
+    ((_, w13_dequantized, w2_dequantized),) = source.dequantize_weight(source_layer)
 
     w13 = torch.cat((w1, w3), dim=1)
     torch.testing.assert_close(w13_dequantized, w13_reference, rtol=0, atol=0)
@@ -2562,8 +2566,8 @@ def test_online_int8_moe_w2_scale_matches_unsharded(monkeypatch) -> None:
 
     def quantize(w2_in, moe_tp_size):
         layer = torch.nn.Module()
-        layer.w13_weight = torch.nn.Parameter(w13, requires_grad=False)
-        layer.w2_weight = torch.nn.Parameter(w2_in, requires_grad=False)
+        layer.w13_weight = torch.nn.Parameter(w13.clone(), requires_grad=False)
+        layer.w2_weight = torch.nn.Parameter(w2_in.clone(), requires_grad=False)
         layer.num_experts = layer.local_num_experts = w13.shape[0]
         method = SimpleNamespace(moe=SimpleNamespace(tp_size=moe_tp_size))
         Int8OnlineMoEMethod._quantize_weights(method, layer)
@@ -2619,13 +2623,29 @@ def test_online_nvfp4_quantizes_original_expert_weights() -> None:
     not is_quant_method_supported("fp8"),
     reason="FP8 is not supported on this GPU type.",
 )
+@pytest.mark.parametrize(
+    "model_name, expected_model_memory_gib",
+    [
+        # MoE models covering both linear and MoE paths, one per modeling backend
+        (GRANITE_MODEL_NAME, 1.35),
+        ("allenai/OLMoE-1B-7B-0125-Instruct", 6.7),
+    ],
+    ids=["moe", "moe_transformers"],
+)
 def test_online_quant_peak_mem(
     vllm_runner,
     caplog_mp_spawn,
     monkeypatch,
+    model_name,
+    expected_model_memory_gib,
 ) -> None:
     _test_online_quant_peak_mem_impl(
-        "fp8_per_tensor", vllm_runner, caplog_mp_spawn, monkeypatch
+        "fp8_per_tensor",
+        model_name,
+        expected_model_memory_gib,
+        vllm_runner,
+        caplog_mp_spawn,
+        monkeypatch,
     )
 
 
diff --git a/tests/quantization/test_online_shared_expert_fusion.py b/tests/quantization/test_online_shared_expert_fusion.py
index 394ba04a1b30..cbe41f2d5a47 100644
--- a/tests/quantization/test_online_shared_expert_fusion.py
+++ b/tests/quantization/test_online_shared_expert_fusion.py
@@ -3,13 +3,16 @@
 """Tests ROCm AITER fused shared experts with online quantization."""
 
 import json
+import logging
 from pathlib import Path
 from types import SimpleNamespace
 from typing import Any
+from unittest.mock import Mock
 
 import pytest
 import torch
 
+import vllm.envs as envs
 from tests.quantization.utils import load_model_without_vllm_runner
 from vllm._aiter_ops import rocm_aiter_ops
 from vllm.config.model import ModelConfig
@@ -17,17 +20,37 @@
     QuantizationConfigArgs,
     resolve_quantization_config,
 )
+from vllm.logger import _print_warning_once
 from vllm.model_executor.layers.fused_moe import FusedMoEFactory
 from vllm.model_executor.layers.fused_moe.experts import rocm_aiter_moe
+from vllm.model_executor.layers.fused_moe.oracle import nvfp4 as nvfp4_oracle
 from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import Mxfp4MoeBackend
+from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import NvFp4MoeBackend
+from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_w4a4_nvfp4 import (  # noqa: E501
+    CompressedTensorsW4A4Nvfp4MoEMethod,
+)
+from vllm.model_executor.layers.quantization.modelopt import ModelOptMxFp8Config
 from vllm.model_executor.layers.quantization.online.base import OnlineQuantizationConfig
+from vllm.model_executor.layers.quantization.online.mxfp4 import Mxfp4OnlineMoEMethod
+from vllm.model_executor.layers.quantization.online.mxfp8 import Mxfp8OnlineMoEMethod
+from vllm.model_executor.layers.quantization.online.requantization import (
+    is_requantization_source,
+)
 from vllm.model_executor.layers.quantization.quark.quark import QuarkConfig
 from vllm.model_executor.layers.quantization.quark.quark_moe import (
+    QuarkNvfp4MoEMethod,
     QuarkOCP_MX_MoEMethod,
 )
 from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
     mxfp4_quantize,
 )
+from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
+    dequant_mxfp8_to_bf16,
+    mxfp8_e4m3_quantize,
+)
+from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import (
+    dequantize_to_dtype,
+)
 from vllm.model_executor.model_loader.dummy_loader import DummyModelLoader
 from vllm.model_executor.model_loader.reload.layerwise import (
     finalize_layerwise_reload,
@@ -123,17 +146,58 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
         self.o_proj = SimpleNamespace(reduce_results=True)
 
 
+class _SingleRankGroup:
+    """Minimal stand-in for a size-1 `GroupCoordinator`."""
+
+    world_size = 1
+    rank_in_group = 0
+    rank = 0
+    local_rank = 0
+    ranks = [0]
+    is_first_rank = True
+    is_last_rank = True
+    # Read at model construction, never used for communication here.
+    device_group = SimpleNamespace(size=lambda: 1, rank=lambda: 0)
+    cpu_group = None
+
+
+@pytest.fixture
+def single_rank_parallel_state(monkeypatch: pytest.MonkeyPatch):
+    """Single rank parallel groups without initializing a process group.
+
+    This is faster than the `dist_init` fixture.
+    """
+    for name in ("_WORLD", "_TP", "_ETP", "_DCP", "_PP", "_DP", "_EP", "_PCP"):
+        monkeypatch.setattr(
+            f"vllm.distributed.parallel_state.{name}", _SingleRankGroup()
+        )
+    monkeypatch.setattr(
+        "vllm.model_executor.parameter.get_tensor_model_parallel_rank", lambda: 0
+    )
+    monkeypatch.setattr(
+        "vllm.model_executor.parameter.get_tensor_model_parallel_world_size", lambda: 1
+    )
+
+
 @pytest.fixture(autouse=True)
 def reset_aiter_shared_expert_topk_metadata(monkeypatch: pytest.MonkeyPatch):
     """Isolate AITER's process-global shared-expert routing buffer."""
+    _print_warning_once.cache_clear()  # `logger.warning_once` dedupes across tests
     monkeypatch.setattr(rocm_aiter_moe, "aiter_topK_meta_data", None)
     rocm_aiter_moe.init_aiter_topK_meta_data.cache_clear()
     yield
     rocm_aiter_moe.init_aiter_topK_meta_data.cache_clear()
 
 
-def _write_minimal_moe_config(model_path: Path, architecture: str) -> None:
-    """Write a compact MoE config with one shared expert."""
+def _write_minimal_moe_config(
+    model_path: Path,
+    architecture: str,
+    quantization_config: dict[str, Any] | None = None,
+) -> None:
+    """Write a compact MoE config with one shared expert.
+
+    `quantization_config` overrides the checkpoint quantization of MiniMax-M3.
+    """
     config = {
         "architectures": [architecture],
         "hidden_size": 256,
@@ -170,13 +234,20 @@ def _write_minimal_moe_config(model_path: Path, architecture: str) -> None:
         text_config = {
             **config,
             "model_type": "minimax_m3_text",
+            "intermediate_size": 256,
             "dense_intermediate_size": 512,
+            "shared_intermediate_size": 256,
             "hidden_act": "swigluoai",
+            "swiglu_alpha": 1.702,
+            "swiglu_limit": 7.0,
+            "scoring_func": "sigmoid",
+            "use_routing_bias": True,
+            "routed_scaling_factor": 2.0,
             "num_local_experts": 4,
-            "n_shared_experts": 1,
             "moe_layer_freq": [1],
             "sparse_attention_config": None,
-            "quantization_config": {
+            "quantization_config": quantization_config
+            or {
                 **_QUARK_MXFP4_CONFIG,
                 "layer_quant_config": {
                     "*block_sparse_moe.experts*": _QUARK_MXFP4_LAYER_CONFIG
@@ -271,7 +342,7 @@ def _unpadded_block(padded: torch.Tensor, expected: torch.Tensor) -> torch.Tenso
 )
 def test_online_shared_expert_loads_bf16_weights_into_mxfp4_slot(
     default_vllm_config,
-    dist_init,
+    single_rank_parallel_state,
     load_path: str,
 ) -> None:
     """A BF16 shared expert is quantized while routed MXFP4 weights are loaded.
@@ -363,7 +434,7 @@ def test_online_shared_expert_loads_bf16_weights_into_mxfp4_slot(
 
 def test_online_shared_expert_reload_compatibility(
     default_vllm_config,
-    dist_init,
+    single_rank_parallel_state,
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
     """Reload a BF16 shared weight into its MXFP4 slot via meta staging."""
@@ -458,13 +529,12 @@ def test_online_quantization(
     architecture: str,
     tmp_path: Path,
     monkeypatch,
-    dist_init,
+    caplog_vllm,
+    single_rank_parallel_state,
     workspace_init,
 ) -> None:
     """Online MXFP4 fuses each architecture's shared expert into its MoE."""
-    monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
-    monkeypatch.setenv("VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS", "1")
-    rocm_aiter_ops.refresh_env_variables()
+    _enable_aiter(monkeypatch, fse=True)
     if architecture == "Glm5NextForCausalLM":
         from vllm.models.glm5next.common import model as glm5_next_model
 
@@ -498,23 +568,6 @@ def test_online_quantization(
     )
     target_pattern = f"*{expected_shared_expert_name}*"
 
-    logged_messages: list[str] = []
-    logged_warnings: list[str] = []
-
-    def record_info(message: str, *args: object, **_kwargs: object) -> None:
-        if "Quantizing " in message and "of types" in message:
-            logged_messages.append(message % args)
-
-    def record_warning(message: str, *args: object, **_kwargs: object) -> None:
-        logged_warnings.append(message % args)
-
-    monkeypatch.setattr(
-        "vllm.model_executor.model_loader.base_loader.logger.info", record_info
-    )
-    monkeypatch.setattr(
-        "vllm.model_executor.layers.fused_moe.utils.logger.warning", record_warning
-    )
-
     model, vllm_config = load_model_without_vllm_runner(
         str(tmp_path),
         dtype="bfloat16",
@@ -551,7 +604,8 @@ def record_warning(message: str, *args: object, **_kwargs: object) -> None:
     shared_expert_prefix = routed_experts.moe_config.shared_expert_prefix
     assert shared_expert_prefix is not None
     assert shared_expert_prefix.endswith(shared_expert_name)
-    assert routed_experts._fused_shared_expert_quantizer is not None
+    assert routed_experts.shared_expert_online_loader is not None
+    assert routed_experts.shared_expert_online_loader.quantizer is not None
     expert_map_manager = routed_experts.expert_map_manager
     assert expert_map_manager.num_fused_shared_experts == 1
     assert expert_map_manager.map_global_to_local(
@@ -571,53 +625,567 @@ def record_warning(message: str, *args: object, **_kwargs: object) -> None:
     assert not any(
         "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS is enabled but "
         "cannot be enabled" in warning
-        for warning in logged_warnings
+        for warning in _fse_warnings(caplog_vllm)
+    )
+    parent_name = (
+        "block_sparse_moe" if architecture == "MiniMaxM3SparseForCausalLM" else "mlp"
+    )
+    assert _online_quantization_logs(caplog_vllm) == [
+        "Quantizing 2 layers of types: "
+        f"{parent_name}.{shared_expert_name}.down_proj: 1 "
+        f"(from targets: {target_pattern}, mxfp4); "
+        f"{parent_name}.{shared_expert_name}.gate_up_proj: 1 "
+        f"(from targets: {target_pattern}, mxfp4)"
+    ]
+
+
+_MINIMAX_NVFP4_ROUTED_MXFP8_SHARED_QUANT_CONFIG: dict[str, Any] = {
+    # Mirrors MiniMax-M3.1's checkpoint: MXFP8 linears and shared experts,
+    # compressed-tensors NVFP4 routed experts.
+    "quant_method": "mxfp8",
+    "activation_scheme": "dynamic",
+    "weight_block_size": [1, 32],
+    "ignored_layers": [
+        "lm_head",
+        "model.embed_tokens",
+        "vision_tower",
+        "multi_modal_projector",
+        "patch_merge_mlp",
+        "model.layers.0.block_sparse_moe.gate",
+    ],
+    "group_size": 16,
+    "moe_quant_algo": "NVFP4",
+    "moe_quant_format": "nvfp4-pack-quantized",
+}
+
+
+def _quark_moe_config(
+    routed: dict[str, Any], shared: dict[str, Any] | None = None
+) -> dict[str, Any]:
+    """Quark MiniMax checkpoint quantizing the routed experts as `routed` and the
+    shared experts as `shared` (excluded if None). Other layers are excluded."""
+    modules = "experts" if shared is None else "(?:experts|shared_experts)"
+    return {
+        **_QUARK_MXFP4_CONFIG,
+        "global_quant_config": routed,
+        "layer_quant_config": (
+            {} if shared in (None, routed) else {"*shared_experts*": shared}
+        ),
+        "exclude": [f"re:^(?!.*\\.block_sparse_moe\\.{modules}(?:\\.|$)).*$"],
+    }
+
+
+# NVFP4: fp4 values with per-16 e4m3 scales, plus a per-tensor fp8_e4m3 scale.
+_NVFP4_VALUES = {
+    "dtype": "fp4",
+    "is_dynamic": False,
+    "qscheme": "per_group",
+    "ch_axis": -1,
+    "group_size": 16,
+    "scale_format": "e4m3",
+    "observer_cls": "PerBlockMXObserver",
+}
+_NVFP4_GLOBAL_SCALE = {
+    "dtype": "fp8_e4m3",
+    "is_dynamic": False,
+    "qscheme": "per_tensor",
+    "observer_cls": "PerTensorMinMaxObserver",
+}
+_QUARK_NVFP4_LAYER_CONFIG: dict[str, Any] = {
+    "input_tensors": [{**_NVFP4_VALUES, "is_dynamic": True}, _NVFP4_GLOBAL_SCALE],
+    "output_tensors": None,
+    "weight": [_NVFP4_VALUES, _NVFP4_GLOBAL_SCALE],
+    "bias": None,
+    "target_device": None,
+}
+_QUARK_MXFP8_LAYER_CONFIG: dict[str, Any] = {
+    **_QUARK_MXFP4_LAYER_CONFIG,
+    **{
+        name: {**_QUARK_MXFP4_LAYER_CONFIG[name], "dtype": "fp8_e4m3"}
+        for name in ("input_tensors", "weight")
+    },
+}
+_QUARK_NVFP4_ROUTED_BF16_SHARED_QUANT_CONFIG = _quark_moe_config(
+    _QUARK_NVFP4_LAYER_CONFIG
+)
+_QUARK_NVFP4_ROUTED_NVFP4_SHARED_QUANT_CONFIG = _quark_moe_config(
+    _QUARK_NVFP4_LAYER_CONFIG, _QUARK_NVFP4_LAYER_CONFIG
+)
+_QUARK_MXFP4_ROUTED_MXFP8_SHARED_QUANT_CONFIG = _quark_moe_config(
+    _QUARK_MXFP4_LAYER_CONFIG, _QUARK_MXFP8_LAYER_CONFIG
+)
+# BF16 (excluded) routed experts, MXFP8 shared experts.
+_QUARK_BF16_ROUTED_MXFP8_SHARED_QUANT_CONFIG = {
+    **_QUARK_MXFP4_CONFIG,
+    "global_quant_config": _QUARK_MXFP8_LAYER_CONFIG,
+    "layer_quant_config": {},
+    "exclude": ["re:^(?!.*\\.block_sparse_moe\\.shared_experts(?:\\.|$)).*$"],
+}
+
+_FSE_ROUTED_SHARED_TARGETS = {
+    "*block_sparse_moe.experts*": "mxfp4",
+    "*shared_experts*": "mxfp4",
+}
+
+_ONLINE_QUANTIZATION_LOGGER = "vllm.model_executor.model_loader.base_loader"
+_FSE_LOGGER = "vllm.model_executor.layers.fused_moe.utils"
+
+
+def _enable_aiter(
+    monkeypatch: pytest.MonkeyPatch, *, moe: bool = True, fse: bool | None = None
+) -> None:
+    """Enable AITER; `fse` sets or unsets fused shared experts, None keeps it."""
+    # Other tests reload `envs`, which can leave these flags pinned as attributes
+    # that shadow the environment variables.
+    for name in (
+        "VLLM_ROCM_USE_AITER",
+        "VLLM_ROCM_USE_AITER_MOE",
+        "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS",
+    ):
+        monkeypatch.delitem(vars(envs), name, raising=False)
+    monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
+    if moe:
+        monkeypatch.setenv("VLLM_ROCM_USE_AITER_MOE", "1")
+    if fse is not None:
+        monkeypatch.setenv("VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS", str(int(fse)))
+    rocm_aiter_ops.refresh_env_variables()
+
+
+def _use_checkpoint_quantization(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Only resolve the user's quantization config, as the CLI does.
+
+    The checkpoint's own "mxfp8" method is detected into `model_config.quantization`,
+    which the helper would otherwise resolve as the online MXFP8 shorthand.
+    """
+    monkeypatch.setattr(
+        "tests.quantization.utils.resolve_quantization_config",
+        lambda _quantization, quantization_config: quantization_config,
+    )
+
+
+def _logged_messages(
+    caplog: pytest.LogCaptureFixture, logger: str, level: int
+) -> list[str]:
+    return [
+        r.getMessage()
+        for r in caplog.records
+        if r.name == logger and r.levelno == level
+    ]
+
+
+def _online_quantization_logs(caplog: pytest.LogCaptureFixture) -> list[str]:
+    return [
+        message
+        for message in _logged_messages(
+            caplog, _ONLINE_QUANTIZATION_LOGGER, logging.INFO
+        )
+        if "Quantizing " in message and "of types" in message
+    ]
+
+
+def _fse_warnings(caplog: pytest.LogCaptureFixture) -> list[str]:
+    return _logged_messages(caplog, _FSE_LOGGER, logging.WARNING)
+
+
+_ROUTED_MXFP4 = {"*block_sparse_moe.experts*": "mxfp4"}
+_ROUTED_SHARED_MXFP4 = {**_ROUTED_MXFP4, "*shared_expert*": "mxfp4"}
+# (targets, fse_requested, routed_method_cls, id) of routed experts requantization.
+_ROUTED_TARGET_CASES = [
+    (_ROUTED_SHARED_MXFP4, True, Mxfp4OnlineMoEMethod, "fse_routed_shared_mxfp4"),
+    (
+        {"*block_sparse_moe.experts*": "mxfp8"},
+        False,
+        Mxfp8OnlineMoEMethod,
+        "routed_mxfp8",
+    ),
+    (_ROUTED_MXFP4, False, Mxfp4OnlineMoEMethod, "routed_mxfp4"),
+    (_ROUTED_SHARED_MXFP4, False, Mxfp4OnlineMoEMethod, "routed_shared_mxfp4"),
+    # FSE refused: the shared expert, kept as stored, does not match.
+    (_ROUTED_MXFP4, True, Mxfp4OnlineMoEMethod, "fse_refused_routed_mxfp4"),
+]
+
+
+@pytest.mark.parametrize(
+    "checkpoint_quant_config, source_cls, targets, fse_requested, routed_method_cls",
+    [
+        *(
+            pytest.param(
+                config,
+                source_cls,
+                targets,
+                fse,
+                method_cls,
+                id=f"{config_id}-{case_id}",
+                marks=pytest.mark.skip(
+                    # TODO: enable once `QuarkLinearMethod` implements
+                    # `dequantize_weight`.
+                    reason="Requantizing Quark NVFP4 shared experts is not supported."
+                )
+                if config is _QUARK_NVFP4_ROUTED_NVFP4_SHARED_QUANT_CONFIG
+                and "*shared_expert*" in targets
+                else (),
+            )
+            for config, source_cls, config_id in [
+                (
+                    _MINIMAX_NVFP4_ROUTED_MXFP8_SHARED_QUANT_CONFIG,
+                    CompressedTensorsW4A4Nvfp4MoEMethod,
+                    "nvfp4_routed_mxfp8_shared",
+                ),
+                (
+                    _QUARK_NVFP4_ROUTED_BF16_SHARED_QUANT_CONFIG,
+                    QuarkNvfp4MoEMethod,
+                    "quark_nvfp4_routed_bf16_shared",
+                ),
+                (
+                    _QUARK_NVFP4_ROUTED_NVFP4_SHARED_QUANT_CONFIG,
+                    QuarkNvfp4MoEMethod,
+                    "quark_nvfp4_routed_nvfp4_shared",
+                ),
+            ]
+            for targets, fse, method_cls, case_id in _ROUTED_TARGET_CASES
+        ),
+        # Routed experts are not requantized. With FSE, the MXFP8 shared expert is
+        # staged and quantized to MXFP4 into its slot after loading. Without, the
+        # shared expert linears are requantized from Quark MXFP8.
+        *(
+            pytest.param(
+                _QUARK_MXFP4_ROUTED_MXFP8_SHARED_QUANT_CONFIG,
+                None,
+                {"*shared_expert*": "mxfp4"},
+                fse,
+                QuarkOCP_MX_MoEMethod,
+                id=f"quark_mxfp4_routed_mxfp8_shared-{case_id}",
+            )
+            for fse, case_id in ((True, "fse_shared_mxfp4"), (False, "shared_mxfp4"))
+        ),
+    ],
+)
+def test_fse_online_quantization_routed_shared(
+    checkpoint_quant_config: dict[str, Any],
+    source_cls: type,
+    targets: dict[str, str],
+    fse_requested: bool,
+    routed_method_cls: type,
+    tmp_path: Path,
+    monkeypatch,
+    caplog_vllm,
+    single_rank_parallel_state,
+    workspace_init,
+) -> None:
+    """Routed and/or shared experts are (re)quantized online. The shared experts
+    are fused into the routed experts only with
+    `VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=1`.
+    """
+    architecture = "MiniMaxM3SparseForCausalLM"
+    _enable_aiter(monkeypatch, fse=fse_requested)
+    _write_minimal_moe_config(tmp_path, architecture, checkpoint_quant_config)
+    _use_checkpoint_quantization(monkeypatch)
+    model, vllm_config = load_model_without_vllm_runner(
+        str(tmp_path),
+        dtype="bfloat16",
+        model_config_kwargs={
+            "quantization_config": QuantizationConfigArgs(targets=targets)
+        },
+        model_loader_cls=DummyModelLoader,
+    )
+    expected_model_cls, resolved_architecture = get_model_architecture(
+        vllm_config.model_config
+    )
+    assert resolved_architecture == architecture
+    assert isinstance(model, expected_model_cls)
+
+    # FSE requested with only the routed experts online quantized: FSE is
+    # rightfully refused, as the shared expert (kept as stored) does not run with
+    # the routed experts' online keys.
+    if fse_requested and "*shared_expert*" not in targets:
+        assert any(
+            "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS is enabled but cannot be enabled"
+            in warning
+            and "do not match the routed expert keys" in warning
+            for warning in _fse_warnings(caplog_vllm)
+        ), _fse_warnings(caplog_vllm)
+        return
+
+    assert not any(
+        "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS is enabled but "
+        "cannot be enabled" in warning
+        for warning in _fse_warnings(caplog_vllm)
+    ), _fse_warnings(caplog_vllm)
+
+    moes = [
+        module
+        for module in model.modules()
+        if hasattr(module, "is_fused_shared_expert_enabled")
+        and hasattr(module, "experts")
+    ]
+    assert len(moes) == 1
+    moe = moes[0]
+    assert moe.is_fused_shared_expert_enabled == fse_requested, architecture
+    assert (moe.shared_experts is None) == fse_requested
+
+    routed_experts = moe.experts.routed_experts
+    assert isinstance(routed_experts.quant_method, routed_method_cls)
+    # Without source, the routed experts are not requantized: the MXFP8 shared
+    # expert is staged, and quantized to MXFP4 into its slot after loading.
+    assert isinstance(
+        getattr(routed_experts.quant_method, "requantization_source", None),
+        source_cls or type(None),
+    )
+    num_fused_shared_experts = 1 if fse_requested else 0
+    assert (
+        routed_experts.expert_map_manager.num_fused_shared_experts
+        == num_fused_shared_experts
     )
-    if architecture == "MiniMaxM3SparseForCausalLM":
-        # TODO: Use resolve_layer_fused_shared_expert in MiniMax M3 and remove
-        # this branch once its virtual fused projections are registered.
-        assert logged_messages == ["Quantizing 0 layers of types: "]
+    num_slots = routed_experts.moe_config.num_logical_experts + num_fused_shared_experts
+    assert routed_experts.w13_weight.shape[0] == num_slots
+    assert routed_experts.w2_weight.shape[0] == num_slots
+    if routed_method_cls is Mxfp8OnlineMoEMethod:
+        # MXFP8 scales are stored flattened as [num_slots * rows, K // 32].
+        assert routed_experts.w13_weight_scale.shape[0] == (
+            num_slots * routed_experts.w13_weight.shape[1]
+        )
+        assert routed_experts.w2_weight_scale.shape[0] == (
+            num_slots * routed_experts.w2_weight.shape[1]
+        )
     else:
-        assert logged_messages == [
-            "Quantizing 2 layers of types: "
-            f"mlp.{shared_expert_name}.down_proj: 1 "
-            f"(from targets: {target_pattern}, mxfp4); "
-            f"mlp.{shared_expert_name}.gate_up_proj: 1 "
-            f"(from targets: {target_pattern}, mxfp4)"
+        assert routed_experts.w13_weight_scale.shape[0] == num_slots
+        assert routed_experts.w2_weight_scale.shape[0] == num_slots
+    # Staging and source parameters are released after requantization.
+    for name in (
+        "online_w13_weight",
+        "online_w2_weight",
+        "online_w13_weight_scale",
+        "online_w2_weight_scale",
+        "w13_weight_packed",
+        "w2_weight_packed",
+        "w13_weight_global_scale",
+        "w2_weight_global_scale",
+    ):
+        assert not hasattr(routed_experts, name), name
+
+    routed_target = "*block_sparse_moe.experts*"
+    shared_target = "*shared_expert*"
+    quantized = []
+    if routed_target in targets:
+        quantized += [
+            f"block_sparse_moe.experts: 1 "
+            f"(from targets: {routed_target}, {targets[routed_target]})"
         ]
+    if shared_target in targets:
+        quantized += [
+            f"block_sparse_moe.shared_experts.{projection}: 1 "
+            f"(from targets: {shared_target}, {targets[shared_target]})"
+            for projection in ("down_proj", "gate_up_proj")
+        ]
+    assert _online_quantization_logs(caplog_vllm) == [
+        f"Quantizing {len(quantized)} layers of types: " + "; ".join(quantized)
+    ]
+
+    if fse_requested:
+        # The staged shared experts were quantized into their slots once.
+        with pytest.raises(NotImplementedError, match="Reloading fused shared"):
+            routed_experts.weight_loader(
+                routed_experts.w2_weight,
+                torch.zeros(1),
+                f"{routed_experts.layer_name}.w2_weight",
+                "w2",
+                routed_experts.moe_config.num_logical_experts,
+            )
+    if fse_requested and source_cls is None:
+        # Requantized routed experts parameters are replaced after processing, and
+        # are restored for reloads by the layerwise reload machinery only.
+        shared_prefix = routed_experts.moe_config.shared_expert_prefix
+        with pytest.raises(NotImplementedError, match="Reloading fused shared"):
+            model.load_weights(
+                [(f"{shared_prefix}.down_proj.weight", routed_experts.w2_weight[0])]
+            )
 
 
-def test_online_gate_up_proj_target_disables_fse(
+def test_fse_online_quantization_bf16_routed_quantized_shared_raises(
     tmp_path: Path,
-    monkeypatch: pytest.MonkeyPatch,
-    dist_init,
+    monkeypatch,
+    single_rank_parallel_state,
     workspace_init,
 ) -> None:
-    """A partial shared-expert target quantizes while disabling FSE."""
-    architecture = "DeepseekForCausalLM"
-    target_pattern = "*shared_experts.gate_up_proj*"
-    monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
-    monkeypatch.setenv("VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS", "1")
-    rocm_aiter_ops.refresh_env_variables()
-    _write_minimal_moe_config(tmp_path, architecture)
+    """Routed experts quantized online from BF16 can not take a staged shared
+    expert stored quantized (MXFP8)."""
+    architecture = "MiniMaxM3SparseForCausalLM"
+    _enable_aiter(monkeypatch, fse=True)
+    _write_minimal_moe_config(
+        tmp_path, architecture, _QUARK_BF16_ROUTED_MXFP8_SHARED_QUANT_CONFIG
+    )
+    _use_checkpoint_quantization(monkeypatch)
+    with pytest.raises(NotImplementedError, match="quantized online from high"):
+        load_model_without_vllm_runner(
+            str(tmp_path),
+            dtype="bfloat16",
+            model_config_kwargs={
+                "quantization_config": QuantizationConfigArgs(
+                    targets=_ROUTED_SHARED_MXFP4
+                )
+            },
+            model_loader_cls=DummyModelLoader,
+        )
 
-    logged_messages: list[str] = []
-    logged_warnings: list[str] = []
 
-    def record_info(message: str, *args: object, **_kwargs: object) -> None:
-        if "Quantizing " in message and "of types" in message:
-            logged_messages.append(message % args)
+def test_fse_online_quantization_routed_shared_loads_checkpoint_weights(
+    default_vllm_config,
+    monkeypatch,
+    single_rank_parallel_state,
+) -> None:
+    """Routed NVFP4 and shared MXFP8 weights land in their MXFP4 expert slots.
 
-    def record_warning(message: str, *args: object, **_kwargs: object) -> None:
-        logged_warnings.append(message % args)
+    The weights are loaded through `MiniMaxM3Model.load_weights` with checkpoint names
+    (`w1/w2/w3`, `shared_experts.*`, `weight_scale_inv`), as the model loader does.
+    """
+    from vllm.models.minimax_m3.amd.model import MiniMaxM3Model
 
-    monkeypatch.setattr(
-        "vllm.model_executor.model_loader.base_loader.logger.info", record_info
+    default_vllm_config.model_config = ModelConfig()
+    hidden_size = intermediate_size = 64
+    num_routed_experts = 2
+    shared_slot = num_routed_experts
+    device = current_platform.device_type
+    quant_config = ModelOptMxFp8Config(
+        is_checkpoint_mxfp8_serialized=True,
+        kv_cache_quant_algo=None,
+        exclude_modules=[],
+        moe_quant_algo="NVFP4",
     )
-    monkeypatch.setattr(
-        "vllm.model_executor.layers.fused_moe.utils.logger.warning", record_warning
+    quant_config.online_quantization_config = OnlineQuantizationConfig(
+        QuantizationConfigArgs(targets=_FSE_ROUTED_SHARED_TARGETS)
     )
 
+    def nvfp4_weight(rows: int, cols: int):
+        packed = torch.randint(0, 256, (rows, cols // 2), dtype=torch.uint8)
+        scale = (torch.rand(rows, cols // 16) * 4 + 0.5).to(torch.float8_e4m3fn)
+        global_scale = torch.rand(1, dtype=torch.float32) + 0.5
+        return packed, scale, global_scale
+
+    def dequantize_nvfp4(packed, scale, global_scale):
+        return dequantize_to_dtype(
+            packed[None],
+            scale[None],
+            (1.0 / global_scale).reshape(1),
+            torch.bfloat16,
+            block_size=16,
+            swizzle=False,
+        )[0]
+
+    with torch.device(device):
+        runner = FusedMoEFactory(
+            num_experts=num_routed_experts,
+            top_k=1,
+            hidden_size=hidden_size,
+            intermediate_size=intermediate_size,
+            params_dtype=torch.bfloat16,
+            n_shared_experts=1,
+            fuse_shared_experts=True,
+            shared_expert_prefix="model.layers.0.block_sparse_moe.shared_experts",
+            prefix="model.layers.0.block_sparse_moe.experts",
+            quant_config=quant_config,
+        )
+        layer = runner.routed_experts
+        assert isinstance(layer.quant_method, Mxfp4OnlineMoEMethod)
+
+        shapes = {
+            "gate_proj": (intermediate_size, hidden_size),
+            "up_proj": (intermediate_size, hidden_size),
+            "down_proj": (hidden_size, intermediate_size),
+        }
+        routed = {
+            (expert, projection): nvfp4_weight(*shape)
+            for expert in range(num_routed_experts)
+            for projection, shape in shapes.items()
+        }
+        shared = {
+            projection: mxfp8_e4m3_quantize(
+                torch.randn(*shape, dtype=torch.bfloat16),
+                is_sf_swizzled_layout=False,
+            )
+            for projection, shape in shapes.items()
+        }
+
+        ckpt_names = {"gate_proj": "w1", "up_proj": "w3", "down_proj": "w2"}
+        weights = []
+        for (expert, projection), (packed, scale, global_scale) in routed.items():
+            prefix = (
+                f"layers.0.block_sparse_moe.experts.{expert}.{ckpt_names[projection]}"
+            )
+            weights += [
+                (f"{prefix}.weight_packed", packed),
+                (f"{prefix}.weight_scale", scale),
+                (f"{prefix}.weight_global_scale", global_scale),
+                (f"{prefix}.input_global_scale", torch.ones(1)),
+            ]
+        for projection, (weight, scale) in shared.items():
+            prefix = f"layers.0.block_sparse_moe.shared_experts.{projection}"
+            weights += [
+                (f"{prefix}.weight", weight),
+                (f"{prefix}.weight_scale_inv", scale),
+            ]
+
+        # Only the parts of `MiniMaxM3Model` that `load_weights` uses.
+        class ModelStub(torch.nn.Module):
+            get_expert_mapping = MiniMaxM3Model.get_expert_mapping
+            load_weights = MiniMaxM3Model.load_weights
+
+        model = ModelStub()
+        model.config = SimpleNamespace(
+            num_local_experts=num_routed_experts, n_shared_experts=1
+        )
+        # Per-layer FSE: the shared expert of layer 0 is fused.
+        model.fused_shared_expert_layers = frozenset({0})
+        decoder_layer = torch.nn.Module()
+        decoder_layer.block_sparse_moe = torch.nn.Module()
+        decoder_layer.block_sparse_moe.experts = runner
+        model.layers = torch.nn.ModuleList([decoder_layer])
+        # `load_weights` only redirects the shared expert of the layers in
+        # `fused_shared_expert_layers` when AITER fused shared experts are
+        # enabled. Enable that gate only for loading, so the expert kernels
+        # selected above stay unchanged.
+        with monkeypatch.context() as m:
+            m.setattr(
+                rocm_aiter_ops, "is_fusion_moe_shared_experts_enabled", lambda: True
+            )
+            model.load_weights(weights)
+        layer.quant_method.process_weights_after_loading(layer)
+
+    def check_slot(slot: int, gate, up, down) -> None:
+        for rows, reference in (
+            (slice(0, intermediate_size), gate),
+            (slice(intermediate_size, 2 * intermediate_size), up),
+        ):
+            expected, expected_scale = mxfp4_quantize(reference)
+            assert torch.equal(layer.w13_weight[slot, rows], expected)
+            assert torch.equal(layer.w13_weight_scale[slot, rows], expected_scale)
+        expected, expected_scale = mxfp4_quantize(down)
+        assert torch.equal(layer.w2_weight[slot], expected)
+        assert torch.equal(layer.w2_weight_scale[slot], expected_scale)
+
+    for expert in range(num_routed_experts):
+        check_slot(
+            expert,
+            *(dequantize_nvfp4(*routed[expert, projection]) for projection in shapes),
+        )
+    check_slot(
+        shared_slot,
+        *(dequant_mxfp8_to_bf16(*shared[projection]) for projection in shapes),
+    )
+
+
+def test_online_gate_up_proj_target_disables_fse(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+    caplog_vllm,
+    single_rank_parallel_state,
+    workspace_init,
+) -> None:
+    """A partial shared-expert target quantizes while disabling FSE."""
+    architecture = "DeepseekForCausalLM"
+    target_pattern = "*shared_experts.gate_up_proj*"
+    _enable_aiter(monkeypatch, fse=True)
+    _write_minimal_moe_config(tmp_path, architecture)
+
     model, _ = load_model_without_vllm_runner(
         str(tmp_path),
         dtype="bfloat16",
@@ -638,13 +1206,232 @@ def record_warning(message: str, *args: object, **_kwargs: object) -> None:
     )
     assert not moe.is_fused_shared_expert_enabled
     assert moe.shared_experts is not None
-    assert logged_messages == [
+    assert _online_quantization_logs(caplog_vllm) == [
         "Quantizing 1 layers of types: "
         "mlp.shared_experts.gate_up_proj: 1 "
         f"(from targets: {target_pattern}, mxfp4)"
     ]
-    assert logged_warnings == [
+    # The untargeted down_proj is unquantized: its keys do not match.
+    [warning] = _fse_warnings(caplog_vllm)
+    assert warning.startswith(
         "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS is enabled but cannot be "
-        "enabled - skipping for this layer: online quantization targets only "
-        "part of the shared expert at model.layers.0.mlp.shared_experts."
+        "enabled - skipping shared-expert fusion: shared expert (weight, activation) "
+        "keys "
+    )
+    assert "at <shared_expert> do not match the routed expert keys" in warning
+
+
+@pytest.mark.skip_global_cleanup
+@pytest.mark.parametrize("num_shared", [0, 1, 2])
+@pytest.mark.parametrize("experts_per_chunk", [1, 3, 4])
+def test_online_moe_quantize_weights_places_routed_then_shared_experts(
+    num_shared: int, experts_per_chunk: int
+):
+    """Chunked routed experts come first and the staged shared experts last."""
+    import types
+
+    from vllm.model_executor.layers.quantization.online.moe_base import (
+        OnlineMoEMethodBase,
+    )
+
+    num_routed, w13_shape, w2_shape = 4, (6, 8), (8, 3)
+    num_experts = num_routed + num_shared
+
+    def expert_values(shape: tuple[int, ...]) -> torch.Tensor:
+        # Every element of expert `e` is `e`, so the slot is recoverable.
+        values = torch.arange(num_experts, dtype=torch.float32)
+        return values.view(-1, 1, 1).expand(num_experts, *shape).contiguous()
+
+    w13_all, w2_all = expert_values(w13_shape), expert_values(w2_shape)
+
+    def dequantize_weight(layer):
+        for start in range(0, num_routed, experts_per_chunk):
+            experts = slice(start, min(start + experts_per_chunk, num_routed))
+            yield experts, w13_all[experts], w2_all[experts]
+
+    method = SimpleNamespace(
+        requantization_source=SimpleNamespace(
+            dequantize_weight=dequantize_weight,
+        ),
+        _pad_requantized_weights=lambda layer, w13, w2, n: (w13, w2),
+    )
+    method.iter_weights_for_quantization = types.MethodType(
+        OnlineMoEMethodBase.iter_weights_for_quantization, method
+    )
+    layer = SimpleNamespace(
+        local_num_experts=num_experts,
+        expert_map_manager=SimpleNamespace(num_fused_shared_experts=num_shared),
+        orig_dtype=torch.float32,
+        shared_expert_online_loader=SimpleNamespace(
+            dequantize_staging=lambda layer: (
+                w13_all[num_routed:],
+                w2_all[num_routed:],
+            ),
+            release_staging=lambda layer: None,
+        ),
+    )
+
+    def quantize(w13, w2):
+        return w13 * 2, w13.mean(-1), w2 * 2, w2.mean(-1)
+
+    w13_q, w13_scale, w2_q, w2_scale = OnlineMoEMethodBase.quantize_weights(
+        method, layer, quantize
+    )
+
+    assert torch.equal(w13_q, w13_all * 2)
+    assert torch.equal(w13_scale, w13_all.mean(-1))
+    assert torch.equal(w2_q, w2_all * 2)
+    assert torch.equal(w2_scale, w2_all.mean(-1))
+
+
+if current_platform.is_rocm():
+    from vllm.platforms.rocm import on_gfx950
+
+    ROCM_GFX950 = on_gfx950()
+else:
+    ROCM_GFX950 = False
+
+
+@pytest.mark.parametrize("requantize", [False, True], ids=["no_requant", "requant"])
+def test_nvfp4_moe_requantization_on_rocm(
+    requantize: bool,
+    tmp_path: Path,
+    default_vllm_config,
+    single_rank_parallel_state,
+    workspace_init,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """NVFP4 routed experts must be requantized online on ROCm.
+
+    NVFP4 -> MXFP4 requantization with `--moe-backend aiter` skips the NVFP4 oracle:
+    `aiter` is not an NVFP4 backend.
+    """
+    _enable_aiter(monkeypatch)
+    if not requantize:
+        _write_minimal_moe_config(
+            tmp_path,
+            "MiniMaxM3SparseForCausalLM",
+            _MINIMAX_NVFP4_ROUTED_MXFP8_SHARED_QUANT_CONFIG,
+        )
+        _use_checkpoint_quantization(monkeypatch)
+        with pytest.raises(
+            NotImplementedError,
+            match="NVFP4 MoE experts .* is not supported on ROCm",
+        ):
+            load_model_without_vllm_runner(
+                str(tmp_path),
+                dtype="bfloat16",
+                model_loader_cls=DummyModelLoader,
+            )
+        return
+
+    if not ROCM_GFX950:
+        pytest.skip("AITER MXFP4 MoE requires gfx950")
+    default_vllm_config.model_config = ModelConfig()
+    default_vllm_config.kernel_config.moe_backend = "aiter"
+
+    nvfp4_oracle_logger = Mock()
+    monkeypatch.setattr(
+        "vllm.model_executor.layers.fused_moe.oracle.nvfp4.logger", nvfp4_oracle_logger
+    )
+
+    hidden_size = intermediate_size = 256
+    num_experts, top_k = 8, 2
+    device = current_platform.device_type
+    quant_config = ModelOptMxFp8Config(
+        is_checkpoint_mxfp8_serialized=True,
+        kv_cache_quant_algo=None,
+        exclude_modules=[],
+        moe_quant_algo="NVFP4",
+    )
+    quant_config.online_quantization_config = OnlineQuantizationConfig(
+        QuantizationConfigArgs(targets={"*block_sparse_moe.experts*": "mxfp4"})
+    )
+
+    with torch.device(device):
+        runner = FusedMoEFactory(
+            num_experts=num_experts,
+            top_k=top_k,
+            hidden_size=hidden_size,
+            intermediate_size=intermediate_size,
+            params_dtype=torch.bfloat16,
+            router_logits_dtype=torch.float32,
+            prefix="model.layers.0.block_sparse_moe.experts",
+            quant_config=quant_config,
+        )
+        layer = runner.routed_experts
+        assert layer.moe_config.moe_backend == "aiter"
+
+    # The requantization target selected the AITER MXFP4 backend, and the NVFP4
+    # oracle did not run for the checkpoint method: no log from it.
+    assert isinstance(layer.quant_method, Mxfp4OnlineMoEMethod)
+    assert "AITER" in layer.quant_method.mxfp4_backend.name
+    assert isinstance(
+        layer.quant_method.requantization_source, CompressedTensorsW4A4Nvfp4MoEMethod
+    )
+    assert nvfp4_oracle_logger.mock_calls == []
+
+
+def test_nvfp4_oracle_only_skipped_for_requantized_layers(
+    default_vllm_config,
+    single_rank_parallel_state,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Only the targeted layer's NVFP4 checkpoint method skips backend selection."""
+    _enable_aiter(monkeypatch)
+    default_vllm_config.model_config = ModelConfig()
+    quant_config = ModelOptMxFp8Config(
+        is_checkpoint_mxfp8_serialized=True,
+        kv_cache_quant_algo=None,
+        exclude_modules=[],
+        moe_quant_algo="NVFP4",
+    )
+    quant_config.online_quantization_config = OnlineQuantizationConfig(
+        QuantizationConfigArgs(targets={"*layers.0.block_sparse_moe.experts*": "mxfp4"})
+    )
+
+    real_select = nvfp4_oracle.select_nvfp4_moe_backend
+    selected: list[tuple[str, NvFp4MoeBackend]] = []
+
+    def select(config, weight_key, activation_key):
+        skipped = is_requantization_source()
+        if skipped:
+            backend, experts_cls = real_select(config, weight_key, activation_key)
+        else:
+            # No native NVFP4 MoE backend on ROCm: pretend one was selected.
+            backend, experts_cls = NvFp4MoeBackend.VLLM_CUTLASS, Mock()
+        selected.append(("skipped" if skipped else "selected", backend))
+        return backend, experts_cls
+
+    monkeypatch.setattr(
+        "vllm.model_executor.layers.quantization.compressed_tensors."
+        "compressed_tensors_moe.compressed_tensors_moe_w4a4_nvfp4."
+        "select_nvfp4_moe_backend",
+        select,
+    )
+
+    layers = []
+    with torch.device(current_platform.device_type):
+        for layer_idx in range(2):
+            runner = FusedMoEFactory(
+                num_experts=4,
+                top_k=1,
+                hidden_size=64,
+                intermediate_size=64,
+                params_dtype=torch.bfloat16,
+                prefix=f"model.layers.{layer_idx}.block_sparse_moe.experts",
+                quant_config=quant_config,
+            )
+            layers.append(runner.routed_experts)
+
+    assert selected == [
+        ("skipped", NvFp4MoeBackend.EMULATION),
+        ("selected", NvFp4MoeBackend.VLLM_CUTLASS),
     ]
+    assert isinstance(layers[0].quant_method, Mxfp4OnlineMoEMethod)
+    assert isinstance(
+        layers[0].quant_method.requantization_source,
+        CompressedTensorsW4A4Nvfp4MoEMethod,
+    )
+    assert isinstance(layers[1].quant_method, CompressedTensorsW4A4Nvfp4MoEMethod)
+    assert not is_requantization_source()
diff --git a/tests/quantization/test_per_token_kv_cache.py b/tests/quantization/test_per_token_kv_cache.py
index e66a577737bf..2647e8b5a8e2 100644
--- a/tests/quantization/test_per_token_kv_cache.py
+++ b/tests/quantization/test_per_token_kv_cache.py
@@ -525,10 +525,8 @@ def test_process_weights_sets_placeholder_scales(kv_cache_dtype: str):
 
     assert layer._k_scale_float == 1.0
     assert layer._v_scale_float == 1.0
-    assert not hasattr(layer, "k_scale")
-    assert not hasattr(layer, "v_scale")
-    assert not hasattr(layer, "q_scale")
-    assert not hasattr(layer, "prob_scale")
+    for name in ("k_scale", "v_scale", "q_scale", "prob_scale"):
+        assert hasattr(layer, name)
 
 
 # ===========================================================================
diff --git a/tests/quantization/utils.py b/tests/quantization/utils.py
index d8446952f0d7..d0d781a96a6a 100644
--- a/tests/quantization/utils.py
+++ b/tests/quantization/utils.py
@@ -131,20 +131,12 @@ def is_quant_method_supported(quant_method: str) -> bool:
 
 def _test_online_quant_peak_mem_impl(
     quantization_arg_value,
+    model_name,
+    expected_model_memory_gib,
     vllm_runner,
     caplog_mp_spawn,
     monkeypatch,
 ) -> None:
-    # Note: `allenai/OLMoE-1B-7B-0125-Instruct` was selected because:
-    # 1. it covers both Linear and MoE paths
-    # 2. it is already used by other tests in CI, so adding it here
-    #    does not increase disk space for CI runners
-    # I really wanted to use `ibm-granite/granite-3.0-1b-a400m-base`
-    # which I think is the smallest MoE model in vLLM (2.5 GiB bf16,
-    # 1.3 GiB fp8), but could not as adding one more model makes CI
-    # run out of disk space.
-    model_name = "allenai/OLMoE-1B-7B-0125-Instruct"
-
     # Force spawn to ensure caplog_mp_spawn works consistently
     # (it relies on VLLM_LOGGING_CONFIG_PATH which spawn reads but fork ignores)
     monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
@@ -182,13 +174,8 @@ def _test_online_quant_peak_mem_impl(
     print(f"GPU memory used after loading weights: {model_memory_gib} GiB")
     print(f"Peak GPU memory usage while loading weights: {peak_memory_gib} GiB")
 
-    expected_model_memory_gib = 6.7
-
-    # for allenai/OLMoE-1B-7B-0125-Instruct the number we see today is 9.06
-    # GiB on CUDA, which is 1.36x above model_memory_gib. A slightly higher
-    # number is expected as when we load and quantize weights in a streaming
-    # fashion we need to have individual weights in bf16 + fp8 alive at the
-    # same time.
+    # Streaming quantization keeps each weight alive in both bf16 and fp8, so
+    # peak memory is slightly above the final model memory
     expected_peak_memory_gib = expected_model_memory_gib * 1.4
 
     assert model_memory_gib < expected_model_memory_gib, (
diff --git a/tests/reasoning/test_minimax_m3_reasoning_parser.py b/tests/reasoning/test_minimax_m3_reasoning_parser.py
index 561e4020cb1f..3769b481f0b5 100644
--- a/tests/reasoning/test_minimax_m3_reasoning_parser.py
+++ b/tests/reasoning/test_minimax_m3_reasoning_parser.py
@@ -16,7 +16,7 @@
 class MiniMaxM3Tokenizer:
     """Small tokenizer with MiniMax M3 reasoning tags as special tokens."""
 
-    special_tokens = ("<mm:think>", "</mm:think>")
+    special_tokens = ("<mm:think>", "</mm:think>", "]~b]")
 
     def __init__(self):
         self._token_to_id: dict[str, int] = {}
@@ -459,3 +459,38 @@ def test_token_id_helpers_enabled_mode():
     assert parser.count_reasoning_tokens(open_reasoning_ids) == len(
         tokenizer.encode("abc")
     )
+
+
+# Rendered like the MiniMax-M3.1 chat template: the system message spells out
+# both markers, and a prior assistant turn without reasoning is prefixed with
+# the closer.
+PROMPT = (
+    "]~b]system\nWrap your reasoning in <mm:think></mm:think> tags, or begin "
+    "directly after the </mm:think> prefix.[e~[\n"
+    "]~b]user\nCapital of France?[e~[\n"
+    "]~b]ai\n</mm:think>Paris.[e~[\n"
+    "]~b]user\nAnd Germany?[e~[\n"
+    "]~b]ai\n"
+)
+
+
+def test_prompt_markers_do_not_end_reasoning():
+    parser, tokenizer = make_parser()
+    prompt_ids = tokenizer.encode(PROMPT)
+
+    assert not parser.is_reasoning_end(prompt_ids)
+    assert not parser.is_reasoning_end(tokenizer.encode(PROMPT + "<mm:think>ab"))
+    assert parser.is_reasoning_end(tokenizer.encode(PROMPT + "<mm:think>a</mm:think>"))
+    assert parser.is_reasoning_end(tokenizer.encode(PROMPT + "</mm:think>Berlin"))
+
+
+def test_streaming_end_check_ignores_prompt_markers():
+    parser, tokenizer = make_parser()
+    prompt_ids = tokenizer.encode(PROMPT)
+    open_ids = tokenizer.encode("<mm:think>ab")
+    closed_ids = tokenizer.encode("</mm:think>")
+
+    assert not parser.is_reasoning_end_streaming(prompt_ids + open_ids, open_ids[-1:])
+    assert parser.is_reasoning_end_streaming(
+        prompt_ids + open_ids + closed_ids, closed_ids
+    )
diff --git a/tests/renderers/test_hf.py b/tests/renderers/test_hf.py
index 491dad801de3..77e114188b0a 100644
--- a/tests/renderers/test_hf.py
+++ b/tests/renderers/test_hf.py
@@ -3,16 +3,22 @@
 
 import jinja2
 import pytest
+import torch
 
-from vllm.config import ModelConfig
+from vllm.config import ModelConfig, VllmConfig
 from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
 from vllm.exceptions import VLLMValidationError
+from vllm.inputs.engine import MultiModalInput
+from vllm.multimodal.inputs import MultiModalKwargsItems, PlaceholderRange
 from vllm.renderers.chat_utils import ConversationMessage, load_chat_template
 from vllm.renderers.hf import (
+    HfRenderer,
+    _build_prompt_embeds_updates,
     _consolidate_system_messages,
     _convert_developer_to_system,
     _detect_content_format,
     _detect_developer_role_support,
+    _ensure_prompt_embeds_placeholder_token,
     _get_hf_base_chat_template_params,
     _template_error_reason,
     _try_extract_ast,
@@ -83,6 +89,43 @@
 ASSISTANT_MESSAGE_TO_CONTINUE = {"role": "assistant", "content": "The capital of"}
 
 
+def test_apply_prompt_embeds_keeps_engine_token_ids_in_vocabulary() -> None:
+    """Consume OOV sentinels without changing text IDs or embedding offsets."""
+    model = "HuggingFaceTB/SmolVLM-256M-Instruct"
+    model_config = ModelConfig(model, enable_prompt_embeds=True)
+    renderer = HfRenderer(VllmConfig(model_config=model_config), get_tokenizer(model))
+    try:
+        assert renderer.tokenizer is not None
+        sentinel = _ensure_prompt_embeds_placeholder_token(renderer.tokenizer)
+        vocab_size = model_config.get_vocab_size()
+        assert sentinel >= vocab_size
+        tensors = [
+            torch.zeros(length, model_config.get_hidden_size()) for length in (2, 1)
+        ]
+        engine_input: MultiModalInput = {
+            "type": "multimodal",
+            "prompt_token_ids": [1, sentinel, sentinel, 2, sentinel, 3],
+            "mm_kwargs": MultiModalKwargsItems(),
+            "mm_hashes": {},
+            "mm_placeholders": {},
+        }
+
+        renderer._apply_prompt_embeds_to_engine_input(
+            engine_input, tensors, _build_prompt_embeds_updates(tensors, sentinel)
+        )
+
+        assert all(
+            0 <= token_id < vocab_size for token_id in engine_input["prompt_token_ids"]
+        )
+        assert engine_input["prompt_token_ids"] == [1, 0, 0, 2, 0, 3]
+        assert engine_input["mm_placeholders"]["prompt_embeds"] == [
+            PlaceholderRange(offset=1, length=2),
+            PlaceholderRange(offset=4, length=1),
+        ]
+    finally:
+        renderer.shutdown()
+
+
 def test_load_chat_template():
     # Testing chatml template
     template_content = load_chat_template(chat_template=chatml_jinja_path)
diff --git a/tests/test_config.py b/tests/test_config.py
index 1f71e1c966d2..72e25126753c 100644
--- a/tests/test_config.py
+++ b/tests/test_config.py
@@ -22,6 +22,7 @@
 
 import vllm.config.vllm as vllm_config_module
 import vllm.envs as envs
+from tests.v1.spec_decode.utils import fake_speculative_config
 from vllm.compilation.backends import VllmBackend
 from vllm.config import (
     AttentionConfig,
@@ -50,7 +51,7 @@
 from vllm.config.speculative import _validate_qwen3_omni_dspark
 from vllm.config.utils import get_field
 from vllm.config.vllm import OPTIMIZATION_LEVEL_TO_CONFIG, OptimizationLevel
-from vllm.platforms import Platform, PlatformEnum, current_platform
+from vllm.platforms import CpuArchEnum, Platform, PlatformEnum, current_platform
 from vllm.sampling_params import BeamSearchParams
 from vllm.transformers_utils.config import (
     _patch_hf_transformers_nested_rope_validation,
@@ -58,6 +59,8 @@
     try_get_dense_modules,
 )
 from vllm.v1.attention.backend import AttentionCGSupport
+from vllm.v1.worker.gpu.spec_decode.dflash2.config import DFlash2SpeculatorConfig
+from vllm.v1.worker.gpu.spec_decode.lilicorr.config import LiLiCorrSpeculatorConfig
 
 DEVICE_TYPE = current_platform.device_type
 
@@ -135,6 +138,7 @@ def test_rocm_mm_prefix_lm_disables_chunked_mm_input(
         parallel_config=SimpleNamespace(
             prefill_context_parallel_size=1,
             worker_cls="test-worker",
+            all2all_backend="naive",
         ),
         model_config=(
             None
@@ -152,6 +156,37 @@ def test_rocm_mm_prefix_lm_disables_chunked_mm_input(
     assert config.scheduler_config.disable_chunked_mm_input is expected
 
 
+@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm-specific test")
+@pytest.mark.parametrize("backend", ["mori", "naive"])
+def test_rocm_resolves_the_mori_all2all_backend(monkeypatch, backend: str) -> None:
+    import vllm.config.parallel as parallel
+    from vllm.platforms.rocm import RocmPlatform
+
+    calls = []
+    monkeypatch.setattr(
+        parallel, "resolve_mori_all2all", lambda *args, **kwargs: calls.append(args)
+    )
+    config = SimpleNamespace(
+        compilation_config=SimpleNamespace(cudagraph_mode=CUDAGraphMode.NONE),
+        parallel_config=SimpleNamespace(
+            prefill_context_parallel_size=1,
+            worker_cls="test-worker",
+            all2all_backend=backend,
+        ),
+        model_config=None,
+        scheduler_config=SimpleNamespace(
+            is_multimodal_model=False, disable_chunked_mm_input=False
+        ),
+        kernel_config=SimpleNamespace(moe_backend="aiter"),
+    )
+
+    RocmPlatform.check_and_update_config(config)
+
+    # Only "mori" picks between EPv1 and EPv2, once, at config time.
+    expected = [(config.parallel_config, None, "aiter")] if backend == "mori" else []
+    assert calls == expected
+
+
 def _sampling_replay_config(
     *,
     return_sampling_mask: bool = True,
@@ -578,7 +613,13 @@ def test_model_runner_default_preserves_backend_fallback(
     monkeypatch.setattr(vllm_config_module, "HAS_TRITON", has_triton)
 
     config = VllmConfig()
-    assert config.use_v2_model_runner == (has_triton or platform.is_unspecified())
+    # non-x86 CPUs support MRv2 without Triton
+    supports_v2_without_triton = (
+        platform.is_cpu() and platform.get_cpu_architecture() != CpuArchEnum.X86
+    )
+    assert config.use_v2_model_runner == (
+        has_triton or platform.is_unspecified() or supports_v2_without_triton
+    )
 
 
 def test_v2_frontend_config_can_be_created_without_triton():
@@ -608,21 +649,40 @@ def test_v2_frontend_config_can_be_created_without_triton():
     )
 
 
-def test_hisparse_requires_v2_model_runner():
-    config = object.__new__(VllmConfig)
-    config.attention_config = AttentionConfig(hisparse_config=HiSparseConfig())
+@pytest.mark.parametrize(
+    ("overrides", "feature"),
+    [
+        (
+            {"attention_config": AttentionConfig(hisparse_config=HiSparseConfig())},
+            "HiSparse",
+        ),
+        ({"watermark_config": WatermarkConfig(key=42)}, "watermarking"),
+    ],
+)
+def test_v1_model_runner_rejects_v2_only_config(overrides, feature):
+    config = SimpleNamespace(
+        parallel_config=ParallelConfig(distributed_executor_backend="mp"),
+        scheduler_config=SchedulerConfig.default_factory(async_scheduling=False),
+        attention_config=AttentionConfig(),
+        watermark_config=None,
+        speculative_config=None,
+        model_config=None,
+    )
+    config._dflash_needs_multi_kv_group = lambda: False
+    config._is_dflash_candidate_draft = lambda: False
+    assert VllmConfig._get_v1_model_runner_unsupported_features(config) == []
 
-    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
+    for name, value in overrides.items():
+        setattr(config, name, value)
+    assert VllmConfig._get_v1_model_runner_unsupported_features(config) == [feature]
 
 
 def test_hisparse_rejects_decode_context_parallelism(monkeypatch):
-    monkeypatch.setattr(current_platform, "is_cuda", lambda: True)
+    monkeypatch.setattr(
+        current_platform,
+        "get_hisparse_backend_cls",
+        lambda: "vllm.v1.hisparse.backend.CUDAHiSparseBackend",
+    )
     monkeypatch.setattr(current_platform, "device_count", lambda: 2)
     with pytest.raises(ValueError, match="decode context parallelism"):
         VllmConfig(
@@ -635,7 +695,11 @@ def test_hisparse_rejects_decode_context_parallelism(monkeypatch):
 
 
 def test_hisparse_rejects_pipeline_parallelism(monkeypatch):
-    monkeypatch.setattr(current_platform, "is_cuda", lambda: True)
+    monkeypatch.setattr(
+        current_platform,
+        "get_hisparse_backend_cls",
+        lambda: "vllm.v1.hisparse.backend.CUDAHiSparseBackend",
+    )
     monkeypatch.setattr(current_platform, "device_count", lambda: 2)
     with pytest.raises(ValueError, match="pipeline parallelism"):
         VllmConfig(
@@ -645,7 +709,11 @@ def test_hisparse_rejects_pipeline_parallelism(monkeypatch):
 
 
 def test_hisparse_rejects_full_cudagraph_mode(monkeypatch):
-    monkeypatch.setattr(current_platform, "is_cuda", lambda: True)
+    monkeypatch.setattr(
+        current_platform,
+        "get_hisparse_backend_cls",
+        lambda: "vllm.v1.hisparse.backend.CUDAHiSparseBackend",
+    )
     monkeypatch.setattr(current_platform, "support_static_graph_mode", lambda: True)
     monkeypatch.setattr("vllm.config.vllm.HAS_TRITON", True)
     with pytest.raises(ValueError, match="does not support cudagraph_mode=FULL"):
@@ -656,7 +724,11 @@ def test_hisparse_rejects_full_cudagraph_mode(monkeypatch):
 
 
 def test_hisparse_rejects_disabled_hybrid_kv_cache_manager(monkeypatch):
-    monkeypatch.setattr(current_platform, "is_cuda", lambda: True)
+    monkeypatch.setattr(
+        current_platform,
+        "get_hisparse_backend_cls",
+        lambda: "vllm.v1.hisparse.backend.CUDAHiSparseBackend",
+    )
     monkeypatch.setattr("vllm.config.vllm.HAS_TRITON", True)
     with pytest.raises(ValueError, match="requires the hybrid KV cache manager"):
         VllmConfig(
@@ -670,7 +742,11 @@ def test_hisparse_rejects_disabled_hybrid_kv_cache_manager(monkeypatch):
 
 
 def test_hisparse_rejects_disabled_full_isl_reservation(monkeypatch):
-    monkeypatch.setattr(current_platform, "is_cuda", lambda: True)
+    monkeypatch.setattr(
+        current_platform,
+        "get_hisparse_backend_cls",
+        lambda: "vllm.v1.hisparse.backend.CUDAHiSparseBackend",
+    )
     with pytest.raises(ValueError, match="requires --scheduler-reserve-full-isl"):
         VllmConfig(
             attention_config=AttentionConfig(hisparse_config=HiSparseConfig()),
@@ -682,9 +758,9 @@ def test_hisparse_rejects_disabled_full_isl_reservation(monkeypatch):
         )
 
 
-def test_hisparse_rejects_non_cuda(monkeypatch):
-    monkeypatch.setattr(current_platform, "is_cuda", lambda: False)
-    with pytest.raises(ValueError, match="requires NVIDIA CUDA"):
+def test_hisparse_rejects_unsupported_platform(monkeypatch):
+    monkeypatch.setattr(current_platform, "get_hisparse_backend_cls", lambda: None)
+    with pytest.raises(ValueError, match="platform does not support HiSparse"):
         VllmConfig(attention_config=AttentionConfig(hisparse_config=HiSparseConfig()))
 
 
@@ -732,7 +808,11 @@ def test_hisparse_rejects_non_cuda(monkeypatch):
     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(
+        current_platform,
+        "get_hisparse_backend_cls",
+        lambda: "vllm.v1.hisparse.backend.CUDAHiSparseBackend",
+    )
     monkeypatch.setattr("vllm.config.vllm.HAS_TRITON", True)
     config = VllmConfig(
         kv_transfer_config=kv_transfer_config,
@@ -748,7 +828,11 @@ def test_hisparse_connector_implies_attention_config(monkeypatch, kv_transfer_co
 
 
 def test_hisparse_connector_preserves_explicit_attention_config(monkeypatch):
-    monkeypatch.setattr(current_platform, "is_cuda", lambda: True)
+    monkeypatch.setattr(
+        current_platform,
+        "get_hisparse_backend_cls",
+        lambda: "vllm.v1.hisparse.backend.CUDAHiSparseBackend",
+    )
     monkeypatch.setattr("vllm.config.vllm.HAS_TRITON", True)
     config = VllmConfig(
         attention_config=AttentionConfig(
@@ -768,9 +852,9 @@ def test_hisparse_connector_preserves_explicit_attention_config(monkeypatch):
     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"):
+def test_hisparse_connector_with_unsupported_platform_rejected(monkeypatch):
+    monkeypatch.setattr(current_platform, "get_hisparse_backend_cls", lambda: None)
+    with pytest.raises(ValueError, match="platform does not support HiSparse"):
         VllmConfig(
             kv_transfer_config=KVTransferConfig(
                 kv_connector="HiSparseConnector",
@@ -848,7 +932,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,
+        speculative_config=fake_speculative_config(method="mtp") if with_mtp else None,
         parallel_config=SimpleNamespace(prefill_context_parallel_size=1),
         compilation_config=CompilationConfig(
             cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE
@@ -1000,7 +1084,7 @@ def test_v2_model_runner_supports_speculative_method(method):
     config = VllmConfig()
     config.speculative_config = cast(
         SpeculativeConfig,
-        SimpleNamespace(
+        fake_speculative_config(
             method=method,
             parallel_drafting=False,
             enable_adaptive_verification=False,
@@ -1020,7 +1104,7 @@ def test_v2_model_runner_heterogeneous_vocab_draft_falls_back_to_v1(
     config = VllmConfig()
     config.speculative_config = cast(
         SpeculativeConfig,
-        SimpleNamespace(
+        fake_speculative_config(
             method="draft_model",
             parallel_drafting=False,
             enable_adaptive_verification=False,
@@ -1476,6 +1560,8 @@ def test_v1_model_runner_rejects_v2_only_features():
             distributed_executor_backend="mp",
         ),
         scheduler_config=SchedulerConfig.default_factory(async_scheduling=False),
+        attention_config=AttentionConfig(),
+        watermark_config=None,
         speculative_config=None,
         model_config=None,
     )
@@ -1663,6 +1749,8 @@ def test_v1_model_runner_rejects_pipeline_parallelism_with_async_scheduling():
             distributed_executor_backend="mp",
         ),
         scheduler_config=SchedulerConfig.default_factory(async_scheduling=False),
+        attention_config=AttentionConfig(),
+        watermark_config=None,
         speculative_config=None,
         model_config=None,
     )
@@ -1712,7 +1800,9 @@ def test_engram_tensor_parallel_size(dp_size: int, across_dp: bool, expected: in
     assert config.get_parallel_size(parallel) == expected
 
 
-@pytest.mark.parametrize("option", ["embedding_across_dp", "dp_shared_memory"])
+@pytest.mark.parametrize(
+    "option", ["embedding_across_dp", "dp_shared_memory", "disk_offload"]
+)
 def test_engram_rejects_elastic_cross_dp(option):
     parallel = ParallelConfig(
         tensor_parallel_size=4,
@@ -1898,6 +1988,23 @@ def test_deepseek_engram_rejects_microbatching(explicit, enable_dbo, ubatch_size
         VllmConfig._resolve_and_verify_engram_config(config)
 
 
+@pytest.mark.parametrize("threshold,expected", [(7, 8), (8, 8), (16, 16)])
+def test_dbo_thresholds_cover_a_request_per_microbatch(threshold, expected):
+    # Verifying 3 drafts: 4 tokens per request, so 2 microbatches need 8.
+    parallel_config = ParallelConfig(
+        enable_dbo=True,
+        dbo_decode_token_threshold=threshold,
+        dbo_prefill_token_threshold=threshold,
+    )
+    config = cast(
+        VllmConfig,
+        SimpleNamespace(uniform_decode_query_len=4, parallel_config=parallel_config),
+    )
+    VllmConfig._raise_dbo_thresholds_for_spec_decode(config)
+    assert parallel_config.dbo_decode_token_threshold == expected
+    assert parallel_config.dbo_prefill_token_threshold == expected
+
+
 def test_engram_explicit_config_requires_supported_model():
     """Explicit all-false settings still opt into model validation."""
     with pytest.raises(ValueError, match="requires a model with supported Engram"):
@@ -1936,12 +2043,46 @@ def test_engram_draft_config_validates_target(monkeypatch, target_has_ple, expli
             assert config.engram_config.cpu_offload is True
 
 
+@pytest.mark.parametrize(
+    ("architecture", "layer_field", "rocm", "expect_error"),
+    [
+        pytest.param(
+            "DeepseekV41ForCausalLM", "engram_layer_ids", False, True, id="deepseek"
+        ),
+        pytest.param(
+            "Qwen4ExpForCausalLM", "ple_layer_ids", True, True, id="qwen4exp_rocm"
+        ),
+        pytest.param(
+            "Qwen4ExpForCausalLM", "ple_layer_ids", False, False, id="qwen4exp_cuda"
+        ),
+    ],
+)
+def test_engram_disk_offload_requires_cuda_qwen4exp(
+    monkeypatch, architecture, layer_field, rocm, expect_error
+):
+    monkeypatch.setattr(current_platform, "is_cuda", lambda: not rocm)
+    model = cast(
+        ModelConfig,
+        SimpleNamespace(
+            architecture=architecture,
+            hf_text_config=SimpleNamespace(**{layer_field: [1]}),
+        ),
+    )
+    config = EngramConfig(disk_offload=True)
+    if expect_error:
+        with pytest.raises(ValueError, match="disk_offload"):
+            config.verify_model_config(model)
+    else:
+        config.verify_model_config(model)
+
+
 def test_engram_hash_tracks_execution_options():
     configs = [
         EngramConfig(cpu_offload=True),
         EngramConfig(cpu_offload=False),
         EngramConfig(cpu_offload=True, embedding_across_dp=True),
         EngramConfig(cpu_offload=True, dp_shared_memory=True),
+        EngramConfig(disk_offload=True),
     ]
     assert len({config.compute_hash() for config in configs}) == len(configs)
 
@@ -3565,6 +3706,138 @@ def test_scheduler_config_init():
         print(SchedulerConfig.default_factory().max_model_len)
 
 
+def test_eplb_num_redundant_experts_default():
+    """Test that num_redundant_experts defaults to None and can be set."""
+    from vllm.config.parallel import EPLBConfig, ParallelConfig
+
+    # Test default is None
+    eplb_config = EPLBConfig()
+    assert eplb_config.num_redundant_experts is None
+
+    # Test explicit value
+    eplb_config_explicit = EPLBConfig(num_redundant_experts=4)
+    assert eplb_config_explicit.num_redundant_experts == 4
+
+    # Models read num_redundant_experts without a ModelConfig, so the
+    # disabled case must resolve to 0 at ParallelConfig construction.
+    parallel_config = ParallelConfig(
+        enable_eplb=False,
+        enable_expert_parallel=False,
+    )
+    assert parallel_config.eplb_config.num_redundant_experts == 0
+
+    # Test ParallelConfig validation - EPLB disabled with non-zero value
+    with pytest.raises(ValidationError, match="EPLB is not enabled"):
+        ParallelConfig(
+            enable_eplb=False,
+            enable_expert_parallel=False,
+            eplb_config=EPLBConfig(num_redundant_experts=4),
+        )
+
+    # Test validation for negative value (validated in ParallelConfig)
+    with pytest.raises(ValidationError, match="greater than or equal to"):
+        ParallelConfig(
+            enable_eplb=False,
+            enable_expert_parallel=False,
+            eplb_config=EPLBConfig(num_redundant_experts=-1),
+        )
+
+
+@pytest.mark.parametrize(
+    "num_experts,tp_size,pcp_size,dp_size,expected",
+    [
+        (8, 8, 1, 1, 0),  # ep_size=8, divisible: 8 % 8 = 0
+        (8, 8, 1, 2, 8),  # ep_size=16, ep_size > experts: need 8 redundant
+        (8, 2, 1, 3, 4),  # ep_size=6, non-divisible: need 4 redundant
+        (16, 4, 1, 2, 0),  # ep_size=8, divisible: 16 % 8 = 0
+        (10, 4, 1, 2, 6),  # ep_size=8, non-divisible: need 6 redundant
+        (7, 2, 1, 2, 1),  # ep_size=4, non-divisible: need 1 redundant
+        (1, 2, 1, 2, 3),  # ep_size=4, single expert: need 3 redundant
+        # The EP group spans TP x PCP x DP, so PCP must be accounted for.
+        (8, 2, 2, 1, 0),  # ep_size=4, divisible: 8 % 4 = 0
+        (10, 2, 2, 1, 2),  # ep_size=4, non-divisible: need 2 redundant
+        (8, 2, 3, 1, 4),  # ep_size=6, non-divisible: need 4 redundant
+    ],
+)
+def test_eplb_num_redundant_experts_auto_computation(
+    num_experts, tp_size, pcp_size, dp_size, expected
+):
+    """Test that num_redundant_experts is correctly computed by ParallelConfig.
+
+    The computation ensures (num_logical_experts + num_redundant_experts)
+    is divisible by ep_size (= tp_size * pcp_size * dp_size).
+    """
+    from vllm.config.parallel import ParallelConfig
+
+    with (
+        patch(
+            "vllm.config.parallel.current_platform.is_cuda_alike",
+            return_value=True,
+        ),
+        patch(
+            "vllm.config.parallel.current_platform.device_count",
+            return_value=tp_size * pcp_size * dp_size,
+        ),
+    ):
+        parallel_config = ParallelConfig(
+            tensor_parallel_size=tp_size,
+            prefill_context_parallel_size=pcp_size,
+            data_parallel_size=dp_size,
+            enable_expert_parallel=True,
+            enable_eplb=True,
+        )
+    # num_redundant_experts should be None before computation
+    assert parallel_config.eplb_config.num_redundant_experts is None
+
+    # Call the computation method
+    parallel_config.compute_eplb_num_redundant_experts(num_experts)
+
+    # Verify the computed value matches expected
+    assert parallel_config.eplb_config.num_redundant_experts == expected, (
+        f"Expected num_redundant_experts={expected} for "
+        f"num_experts={num_experts}, ep_size={tp_size * pcp_size * dp_size}, "
+        f"got {parallel_config.eplb_config.num_redundant_experts}"
+    )
+    # Verify divisibility constraint
+    ep_size = tp_size * pcp_size * dp_size
+    total = num_experts + parallel_config.eplb_config.num_redundant_experts
+    assert total % ep_size == 0, (
+        f"Divisibility check failed: ({num_experts} + "
+        f"{parallel_config.eplb_config.num_redundant_experts}) % {ep_size} != 0"
+    )
+
+
+def test_eplb_num_redundant_experts_explicit_value_preserved():
+    """Test that explicitly set num_redundant_experts is not overwritten."""
+    from vllm.config.parallel import EPLBConfig, ParallelConfig
+
+    with (
+        patch(
+            "vllm.config.parallel.current_platform.is_cuda_alike",
+            return_value=True,
+        ),
+        patch(
+            "vllm.config.parallel.current_platform.device_count",
+            return_value=8,
+        ),
+    ):
+        parallel_config = ParallelConfig(
+            tensor_parallel_size=4,
+            data_parallel_size=2,
+            enable_expert_parallel=True,
+            enable_eplb=True,
+            eplb_config=EPLBConfig(num_redundant_experts=10),
+        )
+    # num_redundant_experts is explicitly set
+    assert parallel_config.eplb_config.num_redundant_experts == 10
+
+    # Call the computation method - should not override
+    parallel_config.compute_eplb_num_redundant_experts(num_logical_experts=8)
+
+    # Should still be the explicit value
+    assert parallel_config.eplb_config.num_redundant_experts == 10
+
+
 def test_scheduler_config_rejects_zero_max_num_scheduled_tokens():
     """A zero token budget never schedules anything, so generation would hang."""
     with pytest.raises(ValidationError, match="max_num_scheduled_tokens"):
@@ -3731,7 +4004,7 @@ def test_target_only_gumbel_allows_speculative_decoding(caplog_vllm, disable_log
     config.watermark_config = WatermarkConfig(
         key=42, allow_target_only_watermarking=True
     )
-    config.speculative_config = SimpleNamespace(
+    config.speculative_config = fake_speculative_config(
         method="mtp",
         draft_sample_method="probabilistic",
         rejection_sample_method="standard",
@@ -3751,7 +4024,7 @@ def test_speculative_watermarking_without_context_dedup_does_not_warn(
     config.watermark_config = WatermarkConfig(
         algorithm="dual_key_gumbel", key=42, deduplicate_contexts="none"
     )
-    config.speculative_config = SimpleNamespace(
+    config.speculative_config = fake_speculative_config(
         method="mtp",
         draft_sample_method="probabilistic",
         rejection_sample_method="standard",
@@ -3772,7 +4045,7 @@ def test_speculative_context_dedup_logs_no_unsupported_warning(
     config.watermark_config = WatermarkConfig(
         algorithm="dual_key_gumbel", key=42, deduplicate_contexts="single_turn"
     )
-    config.speculative_config = SimpleNamespace(
+    config.speculative_config = fake_speculative_config(
         method="mtp",
         draft_sample_method="probabilistic",
         rejection_sample_method="standard",
@@ -3791,7 +4064,7 @@ def test_gumbel_rejects_speculative_decoding_without_target_only():
     config.watermark_config = WatermarkConfig(
         algorithm="gumbel", key=42, allow_target_only_watermarking=False
     )
-    config.speculative_config = SimpleNamespace(
+    config.speculative_config = fake_speculative_config(
         method="mtp",
         draft_sample_method="probabilistic",
         rejection_sample_method="standard",
@@ -3809,7 +4082,7 @@ def test_dual_key_gumbel_warns_that_configured_alpha_is_unused(
     config.watermark_config = WatermarkConfig(
         algorithm="dual_key_gumbel", key=42, alpha=0.25
     )
-    config.speculative_config = SimpleNamespace(
+    config.speculative_config = fake_speculative_config(
         method="mtp",
         draft_sample_method="probabilistic",
         rejection_sample_method="standard",
@@ -3835,7 +4108,7 @@ def test_dual_key_gumbel_alpha_warning_is_not_emitted(
         algorithm="dual_key_gumbel", key=42, alpha=alpha
     )
     if speculative:
-        config.speculative_config = SimpleNamespace(
+        config.speculative_config = fake_speculative_config(
             method="mtp",
             draft_sample_method="probabilistic",
             rejection_sample_method="standard",
@@ -3851,7 +4124,7 @@ def test_dual_key_gumbel_alpha_warning_is_not_emitted(
 def test_dual_key_gumbel_requires_probabilistic_drafting():
     config = _watermarked_vllm_config()
     config.watermark_config = WatermarkConfig(algorithm="dual_key_gumbel", key=42)
-    config.speculative_config = SimpleNamespace(
+    config.speculative_config = fake_speculative_config(
         method="mtp",
         draft_sample_method="greedy",
         rejection_sample_method="standard",
@@ -3877,7 +4150,7 @@ def test_dual_key_gumbel_supports_probabilistic_speculative_decoding(
 ):
     config = _watermarked_vllm_config()
     config.watermark_config = WatermarkConfig(algorithm="dual_key_gumbel", key=42)
-    config.speculative_config = SimpleNamespace(
+    config.speculative_config = fake_speculative_config(
         method=method,
         draft_sample_method="probabilistic",
         rejection_sample_method="standard",
@@ -3887,11 +4160,18 @@ def test_dual_key_gumbel_supports_probabilistic_speculative_decoding(
     config._check_supports_watermarking()
 
 
+_DFLASH_CANDIDATE_SPECULATOR_CONFIGS = {
+    "DFlash2DraftModel": DFlash2SpeculatorConfig,
+    "LiLiCorrDraftModel": LiLiCorrSpeculatorConfig,
+}
+
+
 @pytest.mark.parametrize("architecture", ["DFlash2DraftModel", "LiLiCorrDraftModel"])
 def test_dual_key_gumbel_rejects_dflash_candidate_drafts(architecture):
     config = _watermarked_vllm_config()
     config.watermark_config = WatermarkConfig(algorithm="dual_key_gumbel", key=42)
-    config.speculative_config = SimpleNamespace(
+    config.speculative_config = fake_speculative_config(
+        _DFLASH_CANDIDATE_SPECULATOR_CONFIGS[architecture],
         method="dflash",
         draft_sample_method="probabilistic",
         rejection_sample_method="standard",
@@ -3906,7 +4186,7 @@ def test_dual_key_gumbel_rejects_dflash_candidate_drafts(architecture):
 def test_dual_key_gumbel_rejects_unsupported_speculative_method():
     config = _watermarked_vllm_config()
     config.watermark_config = WatermarkConfig(algorithm="dual_key_gumbel", key=42)
-    config.speculative_config = SimpleNamespace(
+    config.speculative_config = fake_speculative_config(
         method="ngram",
         draft_sample_method="probabilistic",
         rejection_sample_method="standard",
@@ -3935,7 +4215,7 @@ def test_dual_key_gumbel_rejects_incompatible_speculative_modes(overrides, match
         "parallel_drafting": False,
         **overrides,
     }
-    config.speculative_config = SimpleNamespace(**values)
+    config.speculative_config = fake_speculative_config(**values)
 
     with pytest.raises(ValueError, match=match):
         config._check_supports_watermarking()
@@ -4000,19 +4280,7 @@ def test_watermark_key_is_excluded_from_serialization():
     assert "key" not in serialized
 
 
-def test_watermarking_forces_model_runner_v2(monkeypatch):
-    monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "0")
-
-    with patch("vllm.config.vllm.logger.info_once") as info_once:
-        assert _watermarked_vllm_config().use_v2_model_runner
-
-    info_once.assert_called_once_with(
-        "Watermarking requires Model Runner V2 and overrides "
-        "VLLM_USE_V2_MODEL_RUNNER=0."
-    )
-
-
-@patch("vllm.config.speculative.ModelConfig")
+@patch("vllm.config.ModelConfig")
 def test_mtp_draft_uses_model_weights_not_local_cache(mock_model_config_cls):
     """Regression test: MTP + runai_streamer should use model_weights (original
     S3 URL) for the draft model, not model (local cache dir set by
@@ -4026,7 +4294,9 @@ def test_mtp_draft_uses_model_weights_not_local_cache(mock_model_config_cls):
     mock_draft.model = local_cache
     mock_draft.hf_config.model_type = "deepseek_mtp"
     mock_draft.hf_config.n_predict = None
+    mock_draft.hf_config.num_nextn_predict_layers = 1
     mock_draft.hf_config.max_position_embeddings = 4096
+    mock_draft.speculator_config_class = None
     mock_draft.max_model_len = 4096
     mock_model_config_cls.return_value = mock_draft
 
@@ -4036,6 +4306,7 @@ def test_mtp_draft_uses_model_weights_not_local_cache(mock_model_config_cls):
     target_config.hf_text_config.model_type = "deepseek_v3"
     target_config.quantization = None
     target_config.max_model_len = 4096
+    target_config.using_transformers_backend.return_value = False
 
     SpeculativeConfig(
         method="mtp",
diff --git a/tests/test_connections.py b/tests/test_connections.py
index 69df261c86a2..f8e1ea6847c4 100644
--- a/tests/test_connections.py
+++ b/tests/test_connections.py
@@ -292,3 +292,58 @@ async def test_get_bytes_allows_body_within_limit(
     finally:
         if connection._async_client is not None:
             await connection._async_client.close()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("use_async", [False, True], ids=["sync", "async"])
+@pytest.mark.parametrize("always_incomplete", [False, True], ids=["recover", "exhaust"])
+async def test_get_bytes_retries_incomplete_response(
+    monkeypatch: pytest.MonkeyPatch,
+    local_media_server: str,
+    use_async: bool,
+    always_incomplete: bool,
+) -> None:
+    monkeypatch.setattr(envs, "VLLM_MEDIA_FETCH_MAX_RETRIES", 2)
+    monkeypatch.setattr("vllm.connections.time.sleep", lambda _: None)
+    monkeypatch.setattr("vllm.connections.asyncio.sleep", AsyncMock())
+    attempts = 0
+    serve_complete_response = _MediaHandler.do_GET
+
+    def serve_response(handler: _MediaHandler) -> None:
+        nonlocal attempts
+        attempts += 1
+        if attempts == 1 or always_incomplete:
+            handler.send_response(200)
+            handler.send_header("Content-Length", str(len(_SMALL_BODY)))
+            handler.end_headers()
+            handler.wfile.write(_SMALL_BODY[:_CHUNK_SIZE])
+            handler.close_connection = True
+        else:
+            serve_complete_response(handler)
+
+    monkeypatch.setattr(_MediaHandler, "do_GET", serve_response)
+    connection = HTTPConnection()
+    url = f"{local_media_server}/small"
+
+    async def download() -> bytes:
+        if use_async:
+            return await connection.async_get_bytes(url, timeout=2)
+        return connection.get_bytes(url, timeout=2)
+
+    try:
+        if always_incomplete:
+            error_type = (
+                aiohttp.ClientPayloadError
+                if use_async
+                else requests.exceptions.ChunkedEncodingError
+            )
+            with pytest.raises(error_type):
+                await download()
+        else:
+            assert await download() == _SMALL_BODY
+        assert attempts == 2
+    finally:
+        if connection._async_client is not None:
+            await connection._async_client.close()
+        if connection._sync_client is not None:
+            connection._sync_client.close()
diff --git a/tests/tokenizers_/test_deepseek_v41.py b/tests/tokenizers_/test_deepseek_v41.py
index 35302c90361e..7b966d355893 100644
--- a/tests/tokenizers_/test_deepseek_v41.py
+++ b/tests/tokenizers_/test_deepseek_v41.py
@@ -263,6 +263,24 @@ def test_unsupported_media_is_a_request_error():
         render([{"role": "user", "content": [{"type": "input_audio"}]}])
 
 
+def test_tool_references_render_as_tool_names():
+    messages = [
+        {"role": "assistant", "tool_calls": [tool_call("tool_search")]},
+        {
+            "role": "tool",
+            "tool_call_id": "tool_search",
+            "content": [
+                {"type": "tool_reference", "name": "get_weather"},
+                {"type": "tool_reference", "name": "get_time"},
+            ],
+        },
+    ]
+    assert render(messages, thinking=False).endswith(
+        "<｜User｜><tool_result>get_weather\n\nget_time</tool_result>"
+        "<｜Assistant｜></think>"
+    )
+
+
 @pytest.fixture(params=[False, True], ids=["sync", "async"])
 def render_with_media(request, monkeypatch):
     # Keep real message parsing and image tracking; no model download is needed.
diff --git a/tests/tokenizers_/test_mistral.py b/tests/tokenizers_/test_mistral.py
index 2eb986151752..c043b54951bd 100644
--- a/tests/tokenizers_/test_mistral.py
+++ b/tests/tokenizers_/test_mistral.py
@@ -2276,7 +2276,10 @@ def test_convert_ids_to_tokens_pre_args_tekken():
 # ---------------------------------------------------------------------------
 
 
-@pytest.mark.parametrize("reasoning_effort", [None, "none", "high"])
+@pytest.mark.parametrize(
+    "reasoning_effort",
+    [None, "none", "minimal", "low", "medium", "high", "xhigh", "max"],
+)
 def test_validate_request_params_valid_effort(reasoning_effort: str | None) -> None:
     request = ChatCompletionRequest(
         model="test-model",
@@ -2286,17 +2289,6 @@ def test_validate_request_params_valid_effort(reasoning_effort: str | None) -> N
     assert validate_request_params(request) is None
 
 
-def test_validate_request_params_rejects_unsupported_effort() -> None:
-    request = ChatCompletionRequest(
-        model="test-model",
-        messages=[],
-        reasoning_effort="medium",
-    )
-    with pytest.raises(VLLMValidationError, match="reasoning_effort") as exc_info:
-        validate_request_params(request)
-    assert exc_info.value.parameter == "reasoning_effort"
-
-
 def test_validate_request_params_rejects_chat_template() -> None:
     request = ChatCompletionRequest(
         model="test-model",
diff --git a/tests/tokenizers_/test_registry.py b/tests/tokenizers_/test_registry.py
index 73ff4bc5df28..c853dc4d9627 100644
--- a/tests/tokenizers_/test_registry.py
+++ b/tests/tokenizers_/test_registry.py
@@ -169,6 +169,14 @@ def fake_from_pretrained(path_or_repo_id: str, *args, **kwargs):
             CONFIG_MAPPING._extra_content["qwen3_5_moe"] = registered_config
 
 
+def test_hub_id_tokenizer_class_survives_local_snapshot():
+    # Transformers picks TokenizersBackend for this checkpoint by its Hub ID;
+    # the class declared on the Hub (LlamaTokenizerFast) drops every space
+    tokenizer = get_tokenizer("deepseek-ai/deepseek-coder-1.3b-base")
+    text = "key 96727 value 4922"
+    assert tokenizer.decode(tokenizer.encode(text, add_special_tokens=False)) == text
+
+
 @pytest.fixture
 def config_tokenizer_backend(monkeypatch):
     """Keep config helpers and dispatch real; replace Hub/backend loading."""
diff --git a/tests/tool_parsers/test_deepseekv32_tool_parser.py b/tests/tool_parsers/test_deepseekv32_tool_parser.py
index 40ad6033d3bf..3783ca8a78bb 100644
--- a/tests/tool_parsers/test_deepseekv32_tool_parser.py
+++ b/tests/tool_parsers/test_deepseekv32_tool_parser.py
@@ -969,6 +969,45 @@ def test_emits_arguments_before_invoke_completes(self, parser):
         assert combined  # some partial args emitted
         assert combined.startswith('{"k"')
 
+    def test_long_string_argument_streams_before_close(self):
+        """A string body has no ">" and must not be buffered until
+        </parameter> arrives (GitHub #52846)."""
+        tool = ChatCompletionToolsParam(
+            function=FunctionDefinition(
+                name="emit",
+                parameters={
+                    "type": "object",
+                    "properties": {"text": {"type": "string"}},
+                },
+            ),
+        )
+        parser = make_parser(tools=[tool])
+        body = "A" * 4096
+        chunks = [
+            f'{FC_START}\n{INV_START}emit">\n{PARAM_START}text" string="true">',
+            *(body[i : i + 32] for i in range(0, len(body), 32)),
+            f"{PARAM_END}\n{INV_END}\n{FC_END}",
+        ]
+        per_chunk_args = []
+        prev = ""
+        for chunk in chunks:
+            result = parser.extract_tool_calls_streaming(
+                previous_text=prev,
+                current_text=prev + chunk,
+                delta_text=chunk,
+                previous_token_ids=[],
+                current_token_ids=[],
+                delta_token_ids=[1],
+                request=make_request(),
+            )
+            prev += chunk
+            per_chunk_args.append(
+                self._reconstruct_args([result]) if result is not None else ""
+            )
+
+        assert sum(1 for args in per_chunk_args[1:-1] if args) > 1
+        assert json.loads("".join(per_chunk_args)) == {"text": body}
+
     def test_no_marker_leak_chunked(self, parser):
         """Chunked streaming must NOT leak DSML start-marker fragments
         as content (GitHub #40801)."""
diff --git a/tests/tool_parsers/test_deepseekv4_tool_parser.py b/tests/tool_parsers/test_deepseekv4_tool_parser.py
index e79465ed9b32..f730e29b9ad5 100644
--- a/tests/tool_parsers/test_deepseekv4_tool_parser.py
+++ b/tests/tool_parsers/test_deepseekv4_tool_parser.py
@@ -254,6 +254,84 @@ def test_streaming_emits_incremental_argument_chunks():
     }
 
 
+def _stream_param_body(properties: dict, param_open: str, body: str) -> list[str]:
+    """Stream one parameter body in 32-char chunks, one argument list per chunk."""
+    tool = ChatCompletionToolsParam(
+        function=FunctionDefinition(
+            name="emit",
+            parameters={"type": "object", "properties": properties},
+        ),
+    )
+    parser = make_parser(tools=[tool])
+    chunks = [
+        f'{TC_START}\n{INV_START}emit">\n{param_open}',
+        *(body[i : i + 32] for i in range(0, len(body), 32)),
+        f"{PARAM_END}\n{INV_END}\n{TC_END}",
+    ]
+    per_chunk_args = []
+    previous_text = ""
+    for delta_text in chunks:
+        current_text = previous_text + delta_text
+        delta = parser.extract_tool_calls_streaming(
+            previous_text=previous_text,
+            current_text=current_text,
+            delta_text=delta_text,
+            previous_token_ids=[],
+            current_token_ids=[],
+            delta_token_ids=[1],
+            request=make_request(),
+        )
+        previous_text = current_text
+        per_chunk_args.append(
+            "".join(
+                tool_call.function.arguments
+                for tool_call in (delta.tool_calls if delta else [])
+                if tool_call.function and tool_call.function.arguments
+            )
+        )
+    return per_chunk_args
+
+
+def test_streaming_long_string_argument_is_incremental():
+    # Regression for #52846: a string body has no ">" so it used to be
+    # buffered until </parameter> arrived.
+    body = "A" * 4096
+    per_chunk_args = _stream_param_body(
+        {"text": {"type": "string"}},
+        f'{PARAM_START}text" string="true">',
+        body,
+    )
+
+    body_chunks = per_chunk_args[1:-1]
+    assert sum(1 for args in body_chunks if args) > 1
+    assert json.loads("".join(per_chunk_args)) == {"text": body}
+
+
+def test_streaming_non_string_argument_is_withheld_until_close():
+    body = "1" * 128
+    per_chunk_args = _stream_param_body(
+        {"count": {"type": "integer"}},
+        f'{PARAM_START}count" string="false">',
+        body,
+    )
+
+    assert not any(per_chunk_args[1:-1])
+    assert json.loads("".join(per_chunk_args)) == {"count": int(body)}
+
+
+def test_streaming_string_argument_with_markup_is_incremental():
+    body = '<div class="a">x > y</div>\n' * 64
+    per_chunk_args = _stream_param_body(
+        {"html": {"type": "string"}},
+        f'{PARAM_START}html" string="true">',
+        body,
+    )
+
+    body_chunks = per_chunk_args[1:-1]
+    assert sum(1 for args in body_chunks if args) > 1
+    assert json.loads("".join(per_chunk_args)) == {"html": body}
+
+
 def _with_strict(
     tools: list[ChatCompletionToolsParam],
 ) -> list[ChatCompletionToolsParam]:
diff --git a/tests/tool_parsers/test_functiongemma_tool_parser.py b/tests/tool_parsers/test_functiongemma_tool_parser.py
index 10792c16607e..92c6dba9ff86 100644
--- a/tests/tool_parsers/test_functiongemma_tool_parser.py
+++ b/tests/tool_parsers/test_functiongemma_tool_parser.py
@@ -1,9 +1,11 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import json
 from unittest.mock import MagicMock
 
 import pytest
+import regex as re
 
 from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
 from vllm.tool_parsers.functiongemma_tool_parser import FunctionGemmaToolParser
@@ -121,6 +123,171 @@ def test_argument_with_spaces(self, parser):
         assert result == {"message": "Hello World"}
 
 
+# Rendered by the model's chat template (``format_argument``): only strings are
+# ``<escape>``-delimited; booleans, numbers, arrays and objects are bare.
+TEMPLATE_CALL = (
+    "<start_function_call>call:set_alarm{"
+    "days:[<escape>mon<escape>,<escape>tue<escape>],enabled:true,hour:7,"
+    "label:<escape>Wake up<escape>,snooze:{minutes:5}}"
+    "<end_function_call>"
+)
+TEMPLATE_ARGS = {
+    "days": ["mon", "tue"],
+    "enabled": True,
+    "hour": 7,
+    "label": "Wake up",
+    "snooze": {"minutes": 5},
+}
+
+_START = "<start_function_call>"
+_END = "<end_function_call>"
+_TAG_RE = re.compile(f"({_START}|{_END})")
+
+
+def _chunks(text: str, size: int) -> list[str]:
+    """Split ``text`` into ``size``-character chunks, keeping the tags whole.
+
+    The tags are single special tokens, so a delta never splits them.
+    """
+    chunks = []
+    for part in _TAG_RE.split(text):
+        if _TAG_RE.fullmatch(part):
+            chunks.append(part)
+        else:
+            chunks.extend(part[i : i + size] for i in range(0, len(part), size))
+    return [chunk for chunk in chunks if chunk]
+
+
+def _stream(parser, request, text: str, size: int) -> list[tuple[str, dict]]:
+    return _stream_deltas(parser, request, _chunks(text, size))
+
+
+def _stream_deltas(parser, request, deltas: list[str]) -> list[tuple[str, dict]]:
+    """Reassemble streamed tool calls as a client would: by index."""
+    previous = ""
+    calls: dict[int, dict] = {}
+    for i, chunk in enumerate(deltas, start=1):
+        current = previous + chunk
+        delta = parser.extract_tool_calls_streaming(
+            previous, current, chunk, [], [i], [i], request
+        )
+        previous = current
+        for call in (delta.tool_calls if delta else None) or []:
+            entry = calls.setdefault(call.index, {"name": None, "arguments": ""})
+            if call.function.name:
+                entry["name"] = call.function.name
+            entry["arguments"] += call.function.arguments or ""
+    return [
+        (entry["name"], json.loads(entry["arguments"]))
+        for _, entry in sorted(calls.items())
+    ]
+
+
+class TestTemplateArgumentTypes:
+    """Regression tests for #63983: bare (non-string) arguments were dropped."""
+
+    def test_non_streaming_keeps_every_argument(self, parser, mock_request):
+        result = parser.extract_tool_calls(TEMPLATE_CALL, mock_request)
+
+        assert result.tools_called is True
+        assert json.loads(result.tool_calls[0].function.arguments) == TEMPLATE_ARGS
+
+    @pytest.mark.parametrize("size", [1, 3, 1000])
+    def test_streaming_keeps_every_argument(self, parser, mock_request, size):
+        assert _stream(parser, mock_request, TEMPLATE_CALL, size) == [
+            ("set_alarm", TEMPLATE_ARGS)
+        ]
+
+    @pytest.mark.parametrize("size", [1, 1000])
+    def test_streaming_multiple_arguments_reassemble_to_json(
+        self, parser, mock_request, size
+    ):
+        """The closing brace is held back until the call ends, so the
+        concatenated deltas of a multi-argument call form one JSON object."""
+        text = (
+            "<start_function_call>call:get_weather{"
+            "location:<escape>San Francisco<escape>,unit:<escape>celsius<escape>}"
+            "<end_function_call>"
+            "<start_function_call>call:set_alarm{hour:7,snooze:{minutes:5}}"
+            "<end_function_call>"
+        )
+        assert _stream(parser, mock_request, text, size) == [
+            ("get_weather", {"location": "San Francisco", "unit": "celsius"}),
+            ("set_alarm", {"hour": 7, "snooze": {"minutes": 5}}),
+        ]
+
+    @pytest.mark.parametrize(
+        "first_call_deltas",
+        [
+            ["city:<escape>Paris<escape>", "}", _END + _START],
+            ["city:<escape>Paris<escape>", "}" + _END + _START],
+            ["city:<escape>Paris<escape>}" + _END + _START],
+        ],
+    )
+    def test_streaming_closes_call_when_next_starts_in_same_delta(
+        self, parser, mock_request, first_call_deltas
+    ):
+        """Deltas merge when the client falls behind the engine, so one call's
+        end tag can share a delta with the next call's start tag."""
+        deltas = [
+            _START,
+            "call:get_weather{",
+            *first_call_deltas,
+            "call:get_time{",
+            "tz:<escape>UTC<escape>",
+            "}",
+            _END,
+        ]
+        assert _stream_deltas(parser, mock_request, deltas) == [
+            ("get_weather", {"city": "Paris"}),
+            ("get_time", {"tz": "UTC"}),
+        ]
+
+    @pytest.mark.parametrize("size", [1, 1000])
+    def test_streaming_call_without_arguments_sends_empty_object(
+        self, parser, mock_request, size
+    ):
+        text = f"{_START}call:get_time{{}}{_END}"
+        assert _stream(parser, mock_request, text, size) == [("get_time", {})]
+
+    def test_escaped_strings_keep_delimiters(self, parser):
+        args = "note:<escape>a, b: {c} [d]<escape>,count:2"
+        assert parser._parse_arguments(args) == {"note": "a, b: {c} [d]", "count": 2}
+
+    def test_nested_containers_and_literals(self, parser):
+        args = "cfg:{on:false,ids:[1,[2,3],{k:<escape>v<escape>}],none:None},f:1.5"
+        assert parser._parse_arguments(args) == {
+            "cfg": {"on": False, "ids": [1, [2, 3], {"k": "v"}], "none": None},
+            "f": 1.5,
+        }
+
+    @pytest.mark.parametrize(
+        ("prefix", "expected"),
+        [
+            ("hour:7", {}),
+            ("hour:7,", {"hour": 7}),
+            ("hour:7,label:<escape>Wake", {"hour": 7}),
+            ("hour:7,days:[<escape>mon<escape>", {"hour": 7}),
+            ("hour:7,snooze:{minutes:5", {"hour": 7}),
+            ("hour:7,snooze:{minutes:5}", {"hour": 7, "snooze": {"minutes": 5}}),
+        ],
+    )
+    def test_partial_parsing_holds_back_incomplete_values(
+        self, parser, prefix, expected
+    ):
+        assert parser._parse_arguments(prefix, partial=True) == expected
+
+    @pytest.mark.parametrize(
+        "args",
+        ["a:1,items:[1}", "a:1,cfg:{x:1]", "a:1,b:,c:2", "a:1,b:}"],
+    )
+    @pytest.mark.parametrize("partial", [False, True])
+    def test_mismatched_delimiters_stop_parsing(self, parser, args, partial):
+        """A stray closing brace must not spin the parser; the arguments
+        parsed before it are kept and the rest is dropped."""
+        assert parser._parse_arguments(args, partial=partial) == {"a": 1}
+
+
 class TestAdjustRequest:
     def test_skip_special_tokens_disabled(self, parser, mock_request):
         mock_request.tools = [{"type": "function", "function": {"name": "test"}}]
diff --git a/tests/tool_parsers/test_hy_v4_tool_parser.py b/tests/tool_parsers/test_hy_v4_tool_parser.py
new file mode 100644
index 000000000000..785cb7cc43ab
--- /dev/null
+++ b/tests/tool_parsers/test_hy_v4_tool_parser.py
@@ -0,0 +1,189 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import json
+
+import pytest
+import regex as re
+
+from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
+from vllm.sampling_params import StructuredOutputsParams
+from vllm.tool_parsers.hy_v4_tool_parser import HYV4ToolParser
+from vllm.tool_parsers.structural_tag_registry import get_model_structural_tag
+
+
+class _StubTokenizer:
+    def __init__(self, suffix):
+        self.vocab = {
+            f"<{closing}{name}{suffix}>": i
+            for i, (name, closing) in enumerate(
+                (name, closing)
+                for name in ("tool_calls", "tool_call", "arg_key", "arg_value")
+                for closing in ("", "/")
+            )
+        }
+        self.pattern = "|".join(map(re.escape, self.vocab)) + r"|[\s\S]"
+
+    def get_vocab(self):
+        return self.vocab
+
+    def pieces(self, text):
+        return re.findall(self.pattern, text)
+
+    def encode(self, text):
+        return [
+            self.vocab[piece] if piece in self.vocab else 256 + ord(piece)
+            for piece in self.pieces(text)
+        ]
+
+
+@pytest.fixture(params=["", ":6124c78e"], ids=["unsuffixed", "suffixed"])
+def suffix(request):
+    return request.param
+
+
+def _tool_call(arguments, suffix):
+    pairs = "".join(
+        f"<arg_key{suffix}>{key}</arg_key{suffix}>"
+        f"<arg_value{suffix}>{value}</arg_value{suffix}>"
+        for key, value in arguments
+    )
+    return f"<tool_call{suffix}>record{pairs}</tool_call{suffix}>"
+
+
+def _assert_streaming_matches(properties, calls, expected, suffix, chunking):
+    tokenizer = _StubTokenizer(suffix)
+    request = ChatCompletionRequest(
+        model="test-model",
+        messages=[],
+        tools=[
+            {
+                "type": "function",
+                "function": {
+                    "name": "record",
+                    "parameters": {"type": "object", "properties": properties},
+                },
+            }
+        ],
+    )
+    parser = HYV4ToolParser(tokenizer, tools=request.tools)
+    output = (
+        f"<tool_calls{suffix}>"
+        + "".join(_tool_call(call, suffix) for call in calls)
+        + f"</tool_calls{suffix}>"
+    )
+    non_streaming = parser.extract_tool_calls(output, request)
+    assert non_streaming.tools_called
+    assert [json.loads(tc.function.arguments) for tc in non_streaming.tool_calls] == (
+        expected
+    )
+    final_json = [tc.function.arguments for tc in non_streaming.tool_calls]
+
+    if chunking == "whole":
+        chunks = [output]
+    elif chunking == "tokens":
+        chunks = tokenizer.pieces(output)
+    else:
+        # Guided decoding can split even the structural markers across deltas.
+        request.tool_choice = "required"
+        tag = get_model_structural_tag(
+            "hy_v4", request.tools, "required", False, token_suffix=suffix
+        )
+        assert tag is not None
+        request.structured_outputs = StructuredOutputsParams(
+            structural_tag=tag.model_dump_json()
+        )
+        chunks = list(output)
+
+    arguments = [""] * len(calls)
+    names = []
+    previous_text = ""
+    previous_ids: list[int] = []
+    for chunk in chunks:
+        delta_ids = tokenizer.encode(chunk)
+        current_ids = previous_ids + delta_ids
+        current_text = previous_text + chunk
+        delta = parser.extract_tool_calls_streaming(
+            previous_text,
+            current_text,
+            chunk,
+            previous_ids,
+            current_ids,
+            delta_ids,
+            request,
+        )
+        previous_text, previous_ids = current_text, current_ids
+        if delta is None:
+            continue
+        assert not delta.content
+        for tc in delta.tool_calls or []:
+            assert tc.function is not None
+            if tc.function.name is not None:
+                names.append((tc.index, tc.function.name))
+            arguments[tc.index] += tc.function.arguments or ""
+            # A client cannot retract a prefix already sent on the wire.
+            assert final_json[tc.index].startswith(arguments[tc.index])
+
+    assert names == [(i, "record") for i in range(len(calls))]
+    assert arguments == final_json
+    assert [json.loads(args) for args in arguments] == expected
+    assert parser.prev_tool_call_arr == [
+        {"name": "record", "arguments": args} for args in expected
+    ]
+    assert parser.streamed_args_for_tool == final_json
+
+
+@pytest.mark.parametrize("chunking", ["whole", "tokens", "characters"])
+@pytest.mark.parametrize("key", [" count ", "\tcount\n", " \t "])
+@pytest.mark.parametrize(
+    ("value_type", "raw", "value"),
+    [
+        ("string", '  a"b\\c  ', '  a"b\\c  '),
+        ("integer", "7", 7),
+        ("number", "1.5", 1.5),
+        ("boolean", "true", True),
+        ("array", "[1, 2]", [1, 2]),
+        ("object", '{"x": 1}', {"x": 1}),
+        ("null", "null", None),
+    ],
+)
+def test_whitespace_argument_name_preserves_value_type(
+    suffix, chunking, key, value_type, raw, value
+):
+    _assert_streaming_matches(
+        {key: {"type": value_type}}, [[(key, raw)]], [{key: value}], suffix, chunking
+    )
+
+
+@pytest.mark.parametrize("chunking", ["whole", "tokens", "characters"])
+@pytest.mark.parametrize("reverse", [False, True])
+@pytest.mark.parametrize("padded_type", ["integer", "string"])
+def test_whitespace_distinguishes_argument_names(
+    suffix, chunking, reverse, padded_type
+):
+    properties = {
+        "count": {"type": "string" if padded_type == "integer" else "integer"},
+        " count ": {"type": padded_type},
+    }
+    pairs = [("count", "2"), (" count ", "7")]
+    values = ["2", 7] if padded_type == "integer" else [2, "7"]
+    if reverse:
+        pairs.reverse()
+        values.reverse()
+    expected = {key: value for (key, _), value in zip(pairs, values)}
+    _assert_streaming_matches(properties, [pairs], [expected], suffix, chunking)
+
+
+@pytest.mark.parametrize("chunking", ["whole", "tokens", "characters"])
+def test_whitespace_argument_state_resets_between_calls(suffix, chunking):
+    _assert_streaming_matches(
+        {" count ": {"type": "integer"}, " label ": {"type": "string"}},
+        [
+            [(" count ", "7"), (" label ", "first")],
+            [(" count ", "9")],
+            [(" label ", "last")],
+        ],
+        [{" count ": 7, " label ": "first"}, {" count ": 9}, {" label ": "last"}],
+        suffix,
+        chunking,
+    )
diff --git a/tests/tool_parsers/test_structural_tag_registry.py b/tests/tool_parsers/test_structural_tag_registry.py
index c5212934201b..b62abbfa997c 100644
--- a/tests/tool_parsers/test_structural_tag_registry.py
+++ b/tests/tool_parsers/test_structural_tag_registry.py
@@ -1500,6 +1500,36 @@ def test_tool_strict_level_from_name():
         ToolStrictLevel.from_name("off")
 
 
+@pytest.mark.parametrize("required", [False, True])
+@pytest.mark.parametrize("token_suffix", ["", ":6124c78e"])
+def test_hy_v4_preserves_whitespace_in_argument_keys(
+    sample_tools, required, token_suffix
+):
+    """Argument names must match the schema exactly, including whitespace."""
+    sample_tools[0].function.parameters = {
+        "type": "object",
+        "properties": {" count ": {"type": "integer"}},
+        "required": [" count "] if required else [],
+    }
+    tag = get_model_structural_tag(
+        "hy_v4",
+        sample_tools,
+        "required",
+        reasoning=False,
+        token_suffix=token_suffix,
+    )
+    grammar = Grammar.from_structural_tag(tag)
+    call = (
+        f"<tool_calls{token_suffix}><tool_call{token_suffix}>get_weather"
+        f"<arg_key{token_suffix}> count </arg_key{token_suffix}>"
+        f"<arg_value{token_suffix}>2</arg_value{token_suffix}>"
+        f"</tool_call{token_suffix}></tool_calls{token_suffix}>"
+    )
+
+    assert _is_grammar_accept_string(grammar, call)
+    assert not _is_grammar_accept_string(grammar, call.replace(" count ", "count"))
+
+
 # ---------------------------------------------------------------------------
 # GLM-4.7 structural tag: builtin content for strict tools, shallow otherwise
 # ---------------------------------------------------------------------------
@@ -1614,6 +1644,65 @@ def test_glm47_non_strict_auto_rejects_invalid(body: str):
     assert not _is_grammar_accept_string(grammar, body)
 
 
+def _glm47_container_tools() -> list[ChatCompletionToolsParam]:
+    return [
+        ChatCompletionToolsParam(
+            type="function",
+            function={
+                "name": "read_files",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "files": {
+                            "type": "array",
+                            "items": {
+                                "type": "object",
+                                "properties": {"path": {"type": "string"}},
+                            },
+                        },
+                        "options": {"type": "object"},
+                        "limit": {"type": ["integer", "null"]},
+                    },
+                },
+            },
+        )
+    ]
+
+
+@pytest.mark.parametrize(
+    ("key", "value"),
+    [
+        # a Windows path copied with bare backslashes is not valid JSON
+        ("files", r'[{"path": "D:\SVILUPPO\docs\PLAN.md"}]'),
+        # an over-escaped quote leaves a JSON string open; </arg_value> still ends it
+        ("files", r'[{"path\": "docs/PLAN.md"}]'),
+        ("files", '[{"path": "docs/PLAN.md"}]'),
+        ("options", '{"recursive": true}'),
+        ("limit", "20"),
+        ("limit", "null"),
+    ],
+)
+def test_glm47_non_strict_container_values_accept_model_text(key: str, value: str):
+    grammar = _glm47_shallow_grammar("auto", _glm47_container_tools(), auto=True)
+    assert _is_grammar_accept_string(
+        grammar, _glm47_call("read_files", _glm47_arg(key, value))
+    )
+
+
+@pytest.mark.parametrize(
+    ("key", "value"),
+    [
+        ("limit", "twenty"),
+        ("files", "[<arg_key>x</arg_key>]"),
+    ],
+)
+def test_glm47_non_strict_container_tools_reject_invalid(key: str, value: str):
+    grammar = _glm47_shallow_grammar("auto", _glm47_container_tools(), auto=True)
+    assert not _is_grammar_accept_string(
+        grammar, _glm47_call("read_files", _glm47_arg(key, value))
+    )
+
+
 def test_glm47_non_strict_required_requires_at_least_one_call():
     grammar = _glm47_shallow_grammar("required")
 
diff --git a/tests/tools/test_amd_ci_fast_exit.py b/tests/tools/test_amd_ci_fast_exit.py
new file mode 100644
index 000000000000..266981add870
--- /dev/null
+++ b/tests/tools/test_amd_ci_fast_exit.py
@@ -0,0 +1,104 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""The AMD CI pytest plugin skips native teardown without changing pytest's result."""
+
+import os
+import signal
+import subprocess
+import sys
+from pathlib import Path
+
+import pytest
+
+PLUGIN_DIR = Path(__file__).resolve().parents[2] / ".buildkite/scripts/hardware_ci"
+
+# An object that segfaults when the interpreter tears its module down, after the
+# summary, as native library destructors did in vllm/ci #96786.
+CRASH_AT_TEARDOWN = """
+import faulthandler
+
+class _CrashOnTeardown:
+    def __del__(self, _crash=faulthandler._sigsegv):
+        _crash()
+
+_keep = _CrashOnTeardown()
+"""
+
+
+def _run(tmp_path, body, *, plugin=True, env=None, args=()):
+    (tmp_path / "test_case.py").write_text(CRASH_AT_TEARDOWN + body)
+    run_env = {k: v for k, v in os.environ.items() if k != "PYTEST_ADDOPTS"}
+    for name in (
+        "ROCP_TOOL_LIBRARIES",
+        "CUDA_INJECTION64_PATH",
+        "LD_PRELOAD",
+        "VLLM_CI_FAST_EXIT",
+    ):
+        run_env.pop(name, None)
+    run_env["PYTHONPATH"] = str(PLUGIN_DIR)
+    run_env.update(env or {})
+    cmd = [sys.executable, "-m", "pytest", "-q", "-p", "no:cacheprovider", *args]
+    if plugin:
+        cmd += ["-p", "amd_ci_fast_exit"]
+    return subprocess.run(
+        [*cmd, "test_case.py"],
+        cwd=tmp_path,
+        env=run_env,
+        capture_output=True,
+        text=True,
+    )
+
+
+PASSING = "\ndef test_ok():\n    assert True\n"
+
+
+def test_reproduces_the_teardown_crash_without_the_plugin(tmp_path):
+    assert _run(tmp_path, PASSING, plugin=False).returncode == -signal.SIGSEGV
+
+
+def test_a_passing_run_exits_zero_despite_the_teardown_crash(tmp_path):
+    result = _run(tmp_path, PASSING)
+    assert result.returncode == 0, result.stderr
+    assert "leaving without native teardown" in result.stderr
+
+
+def test_a_failing_run_keeps_its_exit_status(tmp_path):
+    result = _run(tmp_path, "\ndef test_bad():\n    assert False\n")
+    assert result.returncode == pytest.ExitCode.TESTS_FAILED
+
+
+def test_python_exit_handlers_still_run(tmp_path):
+    marker = tmp_path / "atexit-ran"
+    body = (
+        "\nimport atexit, pathlib\n"
+        f"atexit.register(pathlib.Path({str(marker)!r}).touch)\n" + PASSING
+    )
+    assert _run(tmp_path, body).returncode == 0
+    assert marker.exists()
+
+
+@pytest.mark.parametrize(
+    "env",
+    [
+        {"VLLM_CI_FAST_EXIT": "0"},
+        {"ROCP_TOOL_LIBRARIES": "/tmp/libkernrec_rocm.so"},
+        {"CUDA_INJECTION64_PATH": "/tmp/libkernrec.so"},
+        {"LD_PRELOAD": "/opt/rocm/lib/librocprofiler-sdk.so"},
+    ],
+)
+def test_off_when_disabled_or_a_kernel_recorder_needs_native_teardown(tmp_path, env):
+    assert _run(tmp_path, PASSING, env=env).returncode == -signal.SIGSEGV
+
+
+def test_a_nested_in_process_session_does_not_end_the_outer_one(tmp_path):
+    body = (
+        "\npytest_plugins = ['pytester']\n"
+        "def test_inner_session(pytester):\n"
+        "    pytester.makepyfile('def test_x():\\n    assert True\\n')\n"
+        "    assert pytester.runpytest_inprocess('-p', 'amd_ci_fast_exit').ret == 0\n"
+        "\ndef test_after_the_inner_session():\n"
+        "    assert True\n"
+    )
+    result = _run(tmp_path, body)
+    assert result.returncode == 0, result.stdout + result.stderr
+    assert "2 passed" in result.stdout
diff --git a/tests/tools/test_docker_build_metadata_args.py b/tests/tools/test_docker_build_metadata_args.py
index 1821c7498ec7..5239fc7adcd1 100644
--- a/tests/tools/test_docker_build_metadata_args.py
+++ b/tests/tools/test_docker_build_metadata_args.py
@@ -6,12 +6,8 @@
 import subprocess
 from pathlib import Path
 
-import pytest
-
 REPO_ROOT = Path(__file__).resolve().parents[2]
 HELPER = REPO_ROOT / ".buildkite" / "scripts" / "docker-build-metadata-args.sh"
-ROCM_CI_BAKE = REPO_ROOT / ".buildkite" / "scripts" / "ci-bake-rocm.sh"
-ROCM_IMAGE_SMOKE = REPO_ROOT / ".buildkite" / "scripts" / "rocm" / "smoke-test-image.sh"
 
 
 def run_helper(
@@ -180,353 +176,3 @@ def test_rust_build_cache_excludes_git_metadata() -> None:
         assert "source=.git,target=.git" in exact_version_stage
         assert 'SETUPTOOLS_SCM_PRETEND_METADATA="{dirty=false}"' in exact_version_stage
         assert "bash tools/build_rust.sh" in exact_version_stage
-
-
-def test_rocm_ci_base_bake_embeds_content_hash_label() -> None:
-    bake_file = (REPO_ROOT / "docker" / "docker-bake-rocm.hcl").read_text()
-
-    for expected in (
-        'variable "CI_BASE_CONTENT_HASH"',
-        'target "ci-base-rocm"',
-        'target   = "ci_base"',
-        '"vllm.ci_base.content_hash" = CI_BASE_CONTENT_HASH',
-    ):
-        assert expected in bake_file
-
-
-def test_rocm_ci_base_metadata_inputs_cover_ci_base_files() -> None:
-    ci_bake = ROCM_CI_BAKE.read_text()
-
-    for expected in (
-        "requirements/common.txt",
-        "requirements/rocm.txt",
-        "requirements/test/rocm.txt",
-        "docker/Dockerfile.rocm",
-    ):
-        assert expected in ci_bake
-
-
-def test_rocm_ci_smoke_runs_in_shared_buildkit_graph() -> None:
-    dockerfile = (REPO_ROOT / "docker" / "Dockerfile.rocm").read_text()
-    ci_hcl = (REPO_ROOT / "docker" / "ci-rocm.hcl").read_text()
-    full_image_group = ci_hcl.split('group "test-rocm-ci-with-wheel"', maxsplit=1)[
-        1
-    ].split("}", maxsplit=1)[0]
-
-    for expected in (
-        "FROM test AS test_smoke",
-        "smoke-test-image.sh --inside",
-        "FROM scratch AS export_test_smoke",
-        'target "smoke-test-rocm-ci"',
-        'target     = "export_test_smoke"',
-        'output     = ["type=local,dest=./build/rocm-smoke-export"]',
-    ):
-        assert expected in dockerfile or expected in ci_hcl
-    assert '"smoke-test-rocm-ci"' in full_image_group
-    assert 'target "smoke-test-rocm-ci"' in ROCM_CI_BAKE.read_text()
-
-
-def prepare_rocm_smoke_test(
-    tmp_path: Path,
-    *,
-    marker_id: str,
-    build_id: str,
-) -> tuple[dict[str, str], Path, Path, Path]:
-    fake_bin = tmp_path / "bin"
-    fake_bin.mkdir()
-    docker_called = tmp_path / "docker-called"
-    docker = fake_bin / "docker"
-    docker.write_text('#!/bin/sh\ntouch "$FAKE_DOCKER_CALLED"\nexit 99\n')
-    docker.chmod(0o755)
-
-    marker = tmp_path / "build" / "rocm-smoke-export" / "vllm-smoke-ok"
-    marker.parent.mkdir(parents=True)
-    marker.write_text(f"{marker_id}\n")
-    env = os.environ.copy()
-    env.update(
-        {
-            "BUILDKITE_BUILD_ID": build_id,
-            "FAKE_DOCKER_CALLED": str(docker_called),
-            "PATH": f"{fake_bin}:{env['PATH']}",
-        }
-    )
-    env.pop("ROCM_CI_ARTIFACT_ONLY", None)
-    env.pop("VLLM_CI_SMOKE_IMAGE", None)
-    return env, marker, docker, docker_called
-
-
-def test_rocm_smoke_marker_avoids_host_image_pull(tmp_path: Path) -> None:
-    env, marker, _, docker_called = prepare_rocm_smoke_test(
-        tmp_path,
-        marker_id="build-123",
-        build_id="build-123",
-    )
-
-    result = subprocess.run(
-        ["bash", str(ROCM_IMAGE_SMOKE)],
-        check=True,
-        cwd=tmp_path,
-        env=env,
-        stdout=subprocess.PIPE,
-        text=True,
-    )
-
-    assert "verified inside BuildKit" in result.stdout
-    assert not docker_called.exists()
-    assert not marker.exists()
-
-
-def test_rocm_smoke_rejects_marker_from_another_build(tmp_path: Path) -> None:
-    env, marker, _, docker_called = prepare_rocm_smoke_test(
-        tmp_path,
-        marker_id="previous-build",
-        build_id="current-build",
-    )
-
-    result = subprocess.run(
-        ["bash", str(ROCM_IMAGE_SMOKE)],
-        check=False,
-        cwd=tmp_path,
-        env=env,
-        stderr=subprocess.PIPE,
-        text=True,
-    )
-
-    assert result.returncode == 1
-    assert "previous-build, not current-build" in result.stderr
-    assert not docker_called.exists()
-    assert marker.exists()
-
-
-def test_rocm_smoke_override_streams_current_checks_to_docker(
-    tmp_path: Path,
-) -> None:
-    env, marker, docker, _ = prepare_rocm_smoke_test(
-        tmp_path,
-        marker_id="build-123",
-        build_id="build-123",
-    )
-    docker_args = tmp_path / "docker-args"
-    docker_stdin = tmp_path / "docker-stdin"
-    docker.write_text(
-        '#!/bin/sh\nprintf "%s\\n" "$@" > "$FAKE_DOCKER_ARGS"\n'
-        'cat > "$FAKE_DOCKER_STDIN"\n'
-    )
-    env.update(
-        {
-            "FAKE_DOCKER_ARGS": str(docker_args),
-            "FAKE_DOCKER_STDIN": str(docker_stdin),
-            "IMAGE_TAG": "rocm/vllm-ci:built",
-            "VLLM_CI_SMOKE_IMAGE": "rocm/vllm-ci:override",
-        }
-    )
-
-    subprocess.run(
-        ["bash", str(ROCM_IMAGE_SMOKE)],
-        check=True,
-        cwd=tmp_path,
-        env=env,
-    )
-
-    assert "rocm/vllm-ci:override" in docker_args.read_text().splitlines()
-    assert docker_args.read_text().splitlines()[-3:] == ["-s", "--", "--inside"]
-    assert "run_smoke_checks()" in docker_stdin.read_text()
-    assert marker.exists()
-
-
-def test_rocm_git_fetch_disables_automatic_maintenance(tmp_path: Path) -> None:
-    fake_bin = tmp_path / "bin"
-    fake_bin.mkdir()
-    git = fake_bin / "git"
-    git.write_text('#!/bin/sh\nprintf "%s\\n" "$@"\n')
-    git.chmod(0o755)
-
-    env = os.environ.copy()
-    env["PATH"] = f"{fake_bin}:{env['PATH']}"
-    result = subprocess.run(
-        [
-            "bash",
-            "-c",
-            'source "$1"; git_fetch_with_timeout --quiet origin HEAD',
-            "bash",
-            str(ROCM_CI_BAKE),
-        ],
-        check=True,
-        env=env,
-        stdout=subprocess.PIPE,
-        text=True,
-    )
-
-    assert result.stdout.splitlines() == [
-        "fetch",
-        "--no-auto-maintenance",
-        "--quiet",
-        "origin",
-        "HEAD",
-    ]
-
-
-def run_rocm_map_build(
-    tmp_path: Path, *, target: str = "test-rocm-ci", **settings: str
-) -> tuple[subprocess.CompletedProcess[str], list[list[str]]]:
-    # Exercise real orchestration/errexit; replace expensive build preparation
-    # and external services. Cached exports launch another copy of this harness.
-    harness = tmp_path / "bake.sh"
-    harness.write_text(
-        r"""
-source "$1"
-for fn in configure_ci_base_write_scope print_header validate_inputs load_ci_hcl \
-    init_bake_files compute_ci_base_hash_if_needed configure_ci_base_image_refs \
-    prepare_git_cache_metadata prepare_ci_build_context configure_custom_rocm_stages \
-    extract_dependency_pins write_rocm_build_arg_override \
-    compute_dependency_cache_keys \
-    write_ci_base_label_override compute_rocm_csrc_content_hash_if_needed \
-    compute_rocm_rust_content_hash_if_needed write_rocm_cache_override \
-    write_build_context_override resolve_ci_base_dependency_targets print_bake_config \
-    seed_dependency_caches_if_needed promote_stable_ci_base_tag \
-    publish_ci_base_handoff_ref
-do
-    eval "$fn() { :; }"
-done
-init_config() {
-    TARGET="$1"
-    BAKE_TARGETS=("$TARGET")
-    BAKE_ALLOW_ARGS=(--allow 'fs.read=/owned context')
-    BAKE_FILES=(-f 'resolved args and caches.hcl' -f 'owned context.hcl')
-    CI_HCL_PATH=ci.hcl VLLM_BAKE_FILE=bake.hcl BUILDER_NAME=builder
-}
-setup_builder() {
-    if [[ "$TARGET" == kernel-symbol-map-rocm && "$FAILURE" == setup ]]; then
-        false
-    fi
-    echo setup >> events
-}
-remote_image_exists() { [[ "$CACHED" == 1 ]]; }
-get_remote_image_label() { echo "$BUILDKITE_COMMIT"; }
-bash() { command bash "$0" "$@"; }
-docker() {
-    printf '%q ' "$@" >> calls
-    printf '\n' >> calls
-    if [[ "${*: -1}" == kernel-symbol-map-rocm ]]; then
-        [[ ! -e kernel-symbol-map-rocm ]] || return 90
-        [[ "$FAILURE" != export ]] || return 41
-        [[ "$FAILURE" != missing ]] || return 0
-        mkdir kernel-symbol-map-rocm
-        if [[ "$FAILURE" == empty ]]; then
-            touch kernel-symbol-map-rocm/kernel_symbol_map.rocm.json.gz
-        else
-            echo current > kernel-symbol-map-rocm/kernel_symbol_map.rocm.json.gz
-        fi
-        [[ "$FAILURE" != partial ]] || return 42
-    else
-        [[ "$FAILURE" != primary ]] || return 43
-    fi
-}
-buildkite-agent() {
-    printf '%q ' "$@" >> "$TEST_ROOT/calls"
-    printf '\n' >> "$TEST_ROOT/calls"
-    [[ "$(cat kernel_symbol_map.rocm.json.gz)" == current ]] || return 91
-    [[ "$FAILURE" != upload ]] || return 44
-}
-main "$2"
-"""
-    )
-    export = tmp_path / "kernel-symbol-map-rocm"
-    export.mkdir()
-    (export / "kernel_symbol_map.rocm.json.gz").write_text("stale")
-    (tmp_path / "calls").touch()
-    env = {
-        "PATH": os.environ["PATH"],
-        "TEST_ROOT": str(tmp_path),
-        "VLLM_KERNEL_SYMBOL_MAP": "1",
-        "BUILDKITE": "true",
-        "BUILDKITE_COMMIT": "current-commit",
-        "IMAGE_TAG": "test:current",
-        "CACHED": "0",
-        "FAILURE": "",
-        **settings,
-    }
-    result = subprocess.run(
-        ["bash", str(harness), str(ROCM_CI_BAKE), target],
-        cwd=tmp_path,
-        env=env,
-        capture_output=True,
-        text=True,
-    )
-    return result, [
-        shlex.split(line) for line in (tmp_path / "calls").read_text().splitlines()
-    ]
-
-
-@pytest.mark.parametrize("cached", ["0", "1"])
-@pytest.mark.parametrize(
-    "failure", ["", "export", "partial", "missing", "empty", "upload"]
-)
-def test_rocm_optional_map_failures_preserve_image_success(
-    tmp_path: Path, cached: str, failure: str
-) -> None:
-    result, calls = run_rocm_map_build(tmp_path, CACHED=cached, FAILURE=failure)
-    assert result.returncode == 0, result.stderr
-    targets = [call[-1] for call in calls if call[:2] == ["buildx", "bake"]]
-    assert targets == (
-        ["test-rocm-ci", "kernel-symbol-map-rocm"]
-        if cached == "0"
-        else ["kernel-symbol-map-rocm"]
-    )
-    if cached == "0":
-        assert calls[0][:-1] == calls[1][:-1]
-    uploads = [call for call in calls if call[:2] == ["artifact", "upload"]]
-    assert uploads == (
-        [["artifact", "upload", "kernel_symbol_map.rocm.json.gz"]]
-        if failure in ("", "upload")
-        else []
-    )
-    assert ("continuing without it" in result.stderr) == bool(failure)
-
-
-def test_rocm_cached_map_setup_stops_on_failure_without_failing_image(
-    tmp_path: Path,
-) -> None:
-    result, calls = run_rocm_map_build(tmp_path, CACHED="1", FAILURE="setup")
-    assert result.returncode == 0, result.stderr
-    assert "continuing without it" in result.stderr
-    assert not calls
-    assert not (tmp_path / "events").exists()
-
-
-@pytest.mark.parametrize(
-    ("target", "failure", "expected_status"),
-    [
-        ("test-rocm-ci", "primary", 43),
-        ("kernel-symbol-map-rocm", "export", 41),
-        ("kernel-symbol-map-rocm", "upload", 44),
-        ("smoke-test-rocm-ci", "", 1),
-        ("export-wheel-rocm", "", 1),
-    ],
-)
-def test_rocm_required_build_failures_remain_fatal(
-    tmp_path: Path, target: str, failure: str, expected_status: int
-) -> None:
-    result, calls = run_rocm_map_build(tmp_path, target=target, FAILURE=failure)
-    assert result.returncode == expected_status, result.stderr
-    assert [call[-1] for call in calls if call[:2] == ["buildx", "bake"]] == [target]
-    assert "continuing without it" not in result.stderr
-
-
-@pytest.mark.parametrize("cached", ["0", "1"])
-@pytest.mark.parametrize("setting", ["VLLM_KERNEL_SYMBOL_MAP", "BAKE_PRINT_ONLY"])
-def test_rocm_disabled_or_print_only_build_skips_map_export(
-    tmp_path: Path, cached: str, setting: str
-) -> None:
-    settings = {setting: "0" if setting == "VLLM_KERNEL_SYMBOL_MAP" else "1"}
-    result, calls = run_rocm_map_build(tmp_path, CACHED=cached, **settings)
-    assert result.returncode == 0, result.stderr
-    expected = (
-        ["test-rocm-ci"]
-        if cached == "0" and setting == "VLLM_KERNEL_SYMBOL_MAP"
-        else []
-    )
-    assert [call[-1] for call in calls] == expected
-    assert (
-        tmp_path / "kernel-symbol-map-rocm/kernel_symbol_map.rocm.json.gz"
-    ).read_text() == "stale"
diff --git a/tests/tools/test_rocm_build.py b/tests/tools/test_rocm_build.py
new file mode 100644
index 000000000000..c24c53d7df5d
--- /dev/null
+++ b/tests/tools/test_rocm_build.py
@@ -0,0 +1,289 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""The ROCm image keys in .buildkite/scripts/rocm/build.sh must cover every
+build-context file the keyed Dockerfile parts read, or editing that file would
+ship a stale image."""
+
+import os
+import subprocess
+from pathlib import Path
+
+import pytest
+import regex as re
+
+ROOT = Path(__file__).parents[2]
+MARKER = "# ==== per-commit stages: nothing below is part of the runtime key ===="
+
+
+def context_sources(dockerfile: str) -> set[str]:
+    text = dockerfile.replace("\\\n", " ")
+    srcs = set()
+    for args in re.findall(r"^\s*(?:COPY|ADD)\s+(.*)$", text, re.M):
+        words = [w for w in args.split() if not w.startswith("--")]
+        if "--from=" not in args and not words[0].startswith("http"):
+            srcs.update(words[:-1])
+    for opts in re.findall(r"--mount=type=bind,(\S*)", text):
+        kv = dict(o.partition("=")[::2] for o in opts.split(","))
+        if "from" not in kv:
+            srcs.add(kv.get("source", kv.get("src", ".")))
+    return srcs
+
+
+def key_paths(name: str, stack: str = "") -> set[str]:
+    script = (ROOT / ".buildkite/scripts/rocm/build.sh").read_text()
+    block_match = re.search(rf"^{name}=\$\(.*?sha256sum", script, re.M | re.S)
+    assert block_match is not None, f"{name}= block not found in build.sh"
+    pattern = r"\b(?:requirements|tools|tests)/[\w./-]+|\brust-toolchain\.toml\b"
+    paths = set(re.findall(pattern, block_match.group(0)))
+    # therock_wheels.py is injected into Ks via a helper variable (rocm stack only)
+    if name == "Ks" and stack == "rocm":
+        m = re.search(r"_[Tt][Rr][Ww]=.*?tools/[\w./-]+", script, re.S)
+        if m:
+            paths.update(re.findall(r"\b(?:requirements|tools)/[\w./-]+", m.group(0)))
+    return paths
+
+
+def test_runtime_key_covers_runtime_inputs():
+    dockerfile = (ROOT / "docker/Dockerfile.rocm").read_text()
+    assert dockerfile.count(MARKER) == 1
+    assert context_sources(dockerfile.split(MARKER)[0]) <= key_paths("Kr")
+
+
+def test_stack_key_covers_stack_inputs():
+    dockerfile = (ROOT / "docker/Dockerfile.rocm_base").read_text()
+    assert context_sources(dockerfile) <= key_paths("Ks", stack="rocm")
+
+
+def run_build(
+    tmp_path: Path, mode: str = "ci", existing: tuple[str, ...] = (), **env: str
+) -> tuple[int, str, list[str]]:
+    """build.sh <mode> against fake git, docker and buildkite-agent; returns (status,
+    stderr, calls). `existing` are the image refs the registry already has.
+    FAILURE=bake|missing|upload breaks the kernel symbol map; FAILURE=registry
+    makes every registry lookup fail, FAILURE=pushed only lookups of pushed refs."""
+    (tmp_path / "docker").symlink_to(ROOT / "docker")
+    fake = tmp_path / "bin"
+    fake.mkdir()
+    calls = tmp_path / "calls"
+    calls.touch()
+    registry = tmp_path / "registry"
+    registry.write_text("".join(f"{ref}\n" for ref in existing))
+    (fake / "git").write_text(
+        """#!/bin/bash
+case $1 in
+  describe) echo v0.0.0-0-g0 ;;
+  rev-parse) [[ $2 == HEAD ]] && echo "${HEAD_REV:-c1}" || echo "${2%%^*}" ;;
+  log) echo 0 ;;
+  archive) tar -ch docker ;;
+esac
+"""
+    )
+    (fake / "docker").write_text(
+        """#!/bin/bash
+[[ $1 == buildx ]] || exit 0
+case $2 in
+  inspect) [[ -z $NO_BUILDER ]] || exit 1 ;;
+  create) echo "builder ${*:3}" >> "$CALLS" ;;
+  imagetools)
+    if [[ $3 == inspect ]]; then
+      [[ $FAILURE == registry ]] && { echo "503 Service Unavailable" >&2; exit 1; }
+      found=0
+      while read -r ref; do [[ $4 == "$ref"* ]] && found=1; done < "$REGISTRY"
+      [[ $FAILURE == pushed ]] && ((found)) && { echo "503" >&2; exit 1; }
+      ((found)) || { echo "$4: not found" >&2; exit 1; }
+      echo "Digest: sha256:$(printf '0%.0s' {1..64})"
+    else
+      echo "create ${@: -1}" >> "$CALLS"
+    fi ;;
+  bake)
+    echo "bake ${*:5}" >> "$CALLS"
+    echo "cache: ${CACHE:-own only}" >&2
+    [[ $5 == --metadata-file ]] && for t in stack-rocm stack-wheels runtime test-deps
+    do
+      printf '  "%s": {\\n    "containerimage.digest": "sha256:%064d"\\n  },\\n' "$t" 1
+    done > "$6"
+    case ${@: -1} in
+      base|runtime|test-deps) printf '%s\n' "$TAG" "$TAG-wheels" >> "$REGISTRY" ;;
+      version) mkdir -p .rocm-version && echo 1.0 > .rocm-version/vllm-version.txt ;;
+      ci) mkdir -p .rocm-wheel && touch .rocm-wheel/vllm-0.whl ;;
+      kernel-symbol-map)
+        [[ $FAILURE == bake ]] && exit 1
+        mkdir -p .rocm-kernel-map && cd .rocm-kernel-map
+        [[ $FAILURE == missing ]] || echo map > kernel_symbol_map.rocm.json.gz ;;
+    esac ;;
+esac
+exit 0
+"""
+    )
+    (fake / "buildkite-agent").write_text(
+        """#!/bin/bash
+[[ $1 == meta-data ]] && { echo "ref@sha256:$(printf '0%.0s' {1..64})"; exit 0; }
+echo "upload ${PWD##*/}/$3" >> "$CALLS"
+[[ $FAILURE == upload && $3 == kernel_symbol_map.rocm.json.gz ]] && exit 1
+exit 0
+"""
+    )
+    (fake / "sleep").write_text("#!/bin/bash\n")  # skip the registry retry delays
+    for tool in fake.iterdir():
+        tool.chmod(0o755)
+    # A map left by an earlier build in this checkout must never be uploaded.
+    (tmp_path / ".rocm-kernel-map").mkdir()
+    (tmp_path / ".rocm-kernel-map/kernel_symbol_map.rocm.json.gz").write_text("stale")
+    run_env = {
+        "PATH": f"{fake}:{os.environ['PATH']}",
+        "ROCM_CI_REPO": "rocm/vllm-ci",
+        "BUILDKITE_BUILD_ID": "b1",
+        "BUILDKITE_COMMIT": "c1",
+        "CALLS": str(calls),
+        "REGISTRY": str(registry),
+        "FAILURE": "",
+        **env,
+    }
+    result = subprocess.run(
+        ["bash", str(ROOT / ".buildkite/scripts/rocm/build.sh"), "rocm", mode],
+        cwd=tmp_path,
+        env=run_env,
+        capture_output=True,
+        text=True,
+    )
+    return result.returncode, result.stderr, calls.read_text().splitlines()
+
+
+run_ci_build = run_build
+
+
+def built_tags(calls: list[str], stderr: str) -> list[str]:
+    return [line.split(" as ", 1)[1] for line in stderr.splitlines() if " as " in line]
+
+
+TRUSTED_KEYED = tuple(
+    f"rocm/vllm-ci:{name}-rocm-" for name in ("stack", "runtime", "test-deps")
+)
+
+
+@pytest.mark.parametrize("mode", ["release", "release-base"])
+def test_release_never_ships_images_from_the_ci_repo(tmp_path, mode):
+    # The trusted CI repo has every key; release must still build its own into ECR.
+    status, stderr, calls = run_build(
+        tmp_path,
+        mode,
+        ROCM_CI_REPO="ecr/release",
+        existing=TRUSTED_KEYED,
+    )
+    assert status == 0, stderr
+    assert not any(
+        "rocm/vllm-ci:" in line
+        for line in stderr.splitlines()
+        if line.startswith("Reusing")
+    )
+    assert [t.split(":")[1].rsplit("-", 1)[0] for t in built_tags(calls, stderr)] == [
+        "stack-rocm",
+        "runtime-rocm",
+    ]
+    assert all(t.startswith("ecr/release:") for t in built_tags(calls, stderr))
+    assert ("bake vllm-openai" in calls) == (mode == "release")
+    # Nor builds from the CI cache: CI agents also run PR code.
+    assert "vllm-ci-cache" not in stderr and "cache: own only" in stderr
+
+
+def test_ci_builds_reuse_the_trusted_images(tmp_path):
+    status, stderr, calls = run_build(
+        tmp_path,
+        "base",
+        ROCM_CI_REPO="rocm/vllm-dev",
+        existing=TRUSTED_KEYED,
+        BUILDKITE_PULL_REQUEST="7",
+    )
+    assert status == 0, stderr
+    assert not built_tags(calls, stderr)
+    assert stderr.count("Reusing rocm/vllm-ci:") == 3
+
+
+@pytest.mark.parametrize(
+    ("pull_request", "branch", "prefix"),
+    [
+        ("7", "feature", "pr7-"),
+        ("false", "main", "br-main-"),
+        ("false", "releases/v1.2", "br-releases-v1.2-"),
+    ],
+)
+def test_untrusted_builds_get_their_own_namespace(
+    tmp_path, pull_request, branch, prefix
+):
+    status, stderr, calls = run_build(
+        tmp_path,
+        "base",
+        ROCM_CI_REPO="rocm/vllm-dev",
+        BUILDKITE_PULL_REQUEST=pull_request,
+        BUILDKITE_BRANCH=branch,
+    )
+    assert status == 0, stderr
+    tags = built_tags(calls, stderr)
+    assert tags and all(t.startswith(f"rocm/vllm-dev:{prefix}") for t in tags)
+
+
+def test_ci_build_uploads_the_kernel_symbol_map_when_asked(tmp_path):
+    status, stderr, calls = run_ci_build(tmp_path, VLLM_KERNEL_SYMBOL_MAP="1")
+    assert status == 0, stderr
+    assert calls == [
+        "bake ci",
+        "upload tmp_path/.rocm-wheel/vllm-*.whl".replace("tmp_path", tmp_path.name),
+        "bake kernel-symbol-map",
+        "upload .rocm-kernel-map/kernel_symbol_map.rocm.json.gz",
+    ]
+
+
+def test_ci_build_skips_the_kernel_symbol_map_by_default(tmp_path):
+    status, stderr, calls = run_ci_build(tmp_path)
+    assert status == 0, stderr
+    assert "bake kernel-symbol-map" not in calls
+
+
+@pytest.mark.parametrize("failure", ["bake", "missing", "upload"])
+def test_a_kernel_symbol_map_failure_never_fails_the_image_build(tmp_path, failure):
+    status, stderr, calls = run_ci_build(
+        tmp_path, VLLM_KERNEL_SYMBOL_MAP="1", FAILURE=failure
+    )
+    assert status == 0, stderr
+    assert "continuing without it" in stderr
+    uploaded = "upload .rocm-kernel-map/kernel_symbol_map.rocm.json.gz" in calls
+    assert uploaded == (failure == "upload")  # the stale map is never uploaded
+
+
+@pytest.mark.parametrize("mode", ["base", "release", "release-base", "keys"])
+def test_a_registry_error_fails_the_step_without_building(tmp_path, mode):
+    # Not a miss: no bake, and nothing left running (run() waits for stderr's EOF).
+    status, stderr, calls = run_build(
+        tmp_path, mode, ROCM_CI_REPO="rocm/vllm-dev", FAILURE="registry"
+    )
+    assert status == 75, stderr
+    assert "registry error" in stderr
+    assert not any(call.startswith("bake") for call in calls)
+
+
+def test_pushed_images_are_pinned_by_the_bake_digest_not_by_tag(tmp_path):
+    # A tag lookup after the push could see a concurrent build's push of that key.
+    status, stderr, _ = run_build(
+        tmp_path, "base", ROCM_CI_REPO="rocm/vllm-dev", FAILURE="pushed"
+    )
+    assert status == 0, stderr
+
+
+def test_build_refuses_a_checkout_that_is_not_the_commit(tmp_path):
+    # The context and the keys come from HEAD: it must be the commit under test.
+    status, stderr, calls = run_build(tmp_path, "base", HEAD_REV="c0")
+    assert status == 1 and "HEAD is not BUILDKITE_COMMIT" in stderr
+    assert not calls
+
+
+@pytest.mark.skipif(
+    Path("/run/buildkit/buildkitd.sock").exists(), reason="uses the host buildkitd"
+)
+@pytest.mark.parametrize("mode", ["base", "release"])
+def test_only_release_builders_share_the_host_network(tmp_path, mode):
+    # Release needs EC2 IMDS (sccache); CI builders also run PR code.
+    _, stderr, calls = run_build(tmp_path, mode, NO_BUILDER="1")
+    (builder,) = [call for call in calls if call.startswith("builder")]
+    name = "vllm-rocm-release" if mode == "release" else "vllm-rocm"
+    assert builder.startswith(f"builder --name {name} ")
+    assert ("network=host" in builder) == (mode == "release"), stderr
diff --git a/tests/transformers_utils/test_speculators_override.py b/tests/transformers_utils/test_speculators_override.py
index b3de532d0fa5..4d7481578286 100644
--- a/tests/transformers_utils/test_speculators_override.py
+++ b/tests/transformers_utils/test_speculators_override.py
@@ -19,6 +19,9 @@
 
 from vllm.engine.arg_utils import AsyncEngineArgs, EngineArgs
 from vllm.transformers_utils.config import maybe_override_with_speculators
+from vllm.v1.worker.gpu.spec_decode import speculator_config as speculator_config_module
+from vllm.v1.worker.gpu.spec_decode.dflash.config import DFlashSpeculatorConfig
+from vllm.v1.worker.gpu.spec_decode.eagle.config import Eagle3SpeculatorConfig
 
 _SPECULATORS_CONFIG_DICT: dict[str, Any] = {
     "architectures": ["DFlashDraftModel"],
@@ -126,6 +129,30 @@ def test_speculators_user_cannot_override_method_field():
     assert spec["method"] == "dflash"
 
 
+def test_speculators_keeps_custom_method_based_on_checkpoint_method(monkeypatch):
+    monkeypatch.setattr(speculator_config_module, "_CUSTOM_CONFIGS", {})
+    register = speculator_config_module.register_speculative_method
+    register("my_dflash", DFlashSpeculatorConfig)
+    register("my_eagle3", Eagle3SpeculatorConfig)
+    with _patched_get_config_dict(_SPECULATORS_CONFIG_DICT):
+        _, _, spec = maybe_override_with_speculators(
+            model="fake-org/fake-dflash-spec",
+            tokenizer=None,
+            trust_remote_code=False,
+            vllm_speculative_config={"method": "my_dflash"},
+        )
+        assert spec is not None
+        assert spec["method"] == "my_dflash"
+
+        with pytest.raises(ValueError, match="based on 'eagle3'"):
+            maybe_override_with_speculators(
+                model="fake-org/fake-dflash-spec",
+                tokenizer=None,
+                trust_remote_code=False,
+                vllm_speculative_config={"method": "my_eagle3"},
+            )
+
+
 def test_speculators_user_can_override_num_speculative_tokens():
     """Runtime knobs (num_speculative_tokens, moe_backend, etc.) take user value."""
     user_overrides = {
diff --git a/tests/utils_/test_cache.py b/tests/utils_/test_cache.py
index 3fce04d268b2..535a606919ea 100644
--- a/tests/utils_/test_cache.py
+++ b/tests/utils_/test_cache.py
@@ -127,6 +127,24 @@ def test_lru_cache():
     assert 6 in cache
 
 
+def test_lru_cache_peek_preserves_eviction_order_and_stats():
+    cache = LRUCache[str, object | None](2)
+    payload = object()
+    cache.put("oldest", payload)
+    cache.put("newest", None)
+    initial_stats = cache.stat()
+
+    assert cache.peek("oldest") is payload
+    assert cache.peek("newest", payload) is None
+    assert cache.peek("missing") is None
+    assert cache.peek("missing", payload) is payload
+    assert cache.stat() == initial_stats
+
+    cache.put("replacement", object())
+    assert "oldest" not in cache
+    assert "newest" in cache
+
+
 def test_lru_cache_put_if_fits():
     cache = LRUCache[str, int](10, getsizeof=lambda x: x)
 
diff --git a/tests/utils_/test_extensible_tensor.py b/tests/utils_/test_extensible_tensor.py
new file mode 100644
index 000000000000..0fa4621c8e84
--- /dev/null
+++ b/tests/utils_/test_extensible_tensor.py
@@ -0,0 +1,306 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import subprocess
+import sys
+
+import pytest
+import torch
+
+from tests.utils import create_new_process_for_each_test
+from vllm.utils.extensible_tensor import ExtensibleTensor
+from vllm.utils.vmm_driver import vmm_unavailable_reason
+
+# Each test runs in its own process: touching the driver here would initialize
+# CUDA in the pytest parent and break later fork-based tests in the shard.
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
+
+
+def _skip_unless_vmm() -> None:
+    if (reason := vmm_unavailable_reason()) is not None:
+        pytest.skip(f"VMM unavailable: {reason}")
+
+
+@create_new_process_for_each_test("spawn")
+def test_extensible_tensor_grows_without_moving() -> None:
+    _skip_unless_vmm()
+    buffer = ExtensibleTensor(4096, device="cuda")
+    try:
+        base_ptr = buffer.base_ptr
+        buffer.resize_per_segment_(1024)
+        first_view = buffer.segment_view(0)
+        assert first_view.data_ptr() == base_ptr
+        assert first_view.untyped_storage().nbytes() == 1024
+        first_view.fill_(7)
+
+        buffer.resize_per_segment_(2048)
+        second_view = buffer.segment_view(0)
+        assert second_view.data_ptr() == base_ptr
+        assert second_view.untyped_storage().nbytes() == 2048
+        assert torch.equal(second_view[:1024], torch.full_like(second_view[:1024], 7))
+
+        full_view = buffer.full_view()
+        assert full_view.data_ptr() == base_ptr
+        assert full_view.numel() == 4096
+    finally:
+        buffer.free()
+
+
+@create_new_process_for_each_test("spawn")
+def test_extensible_tensor_rejects_shrink_and_overflow() -> None:
+    _skip_unless_vmm()
+    buffer = ExtensibleTensor(1024, device="cuda")
+    try:
+        buffer.resize_per_segment_(512)
+        with pytest.raises(ValueError, match="grow-only"):
+            buffer.resize_per_segment_(256)
+        with pytest.raises(ValueError, match="exceeds the segment capacity"):
+            buffer.resize_per_segment_(1025)
+    finally:
+        buffer.free()
+
+
+@create_new_process_for_each_test("spawn")
+def test_segments_grow_in_lockstep_and_zero_new() -> None:
+    _skip_unless_vmm()
+    """Each segment's committed prefix grows in lockstep.
+
+    Data written to a segment's committed prefix survives a grow; the newly
+    committed range of each segment is zeroed with `zero_new=True` while old
+    bytes are preserved.
+    """
+    et = ExtensibleTensor(max_num_bytes=8192, device="cuda", num_segments=2)
+    try:
+        assert et.num_segments == 2
+        assert et.segment_capacity_bytes == 4096
+
+        et.resize_per_segment_(256, zero_new=True)
+        assert et.bytes_per_segment == 256
+        assert et.num_bytes == 512
+        fv = et.full_view()
+        assert fv.shape == (8192,)
+        assert torch.count_nonzero(fv[:256]) == 0
+        assert torch.count_nonzero(fv[4096 : 4096 + 256]) == 0
+
+        pattern_a = torch.arange(256, device="cuda", dtype=torch.uint8)
+        pattern_b = 255 - pattern_a
+        fv[:256].copy_(pattern_a)
+        fv[4096 : 4096 + 256].copy_(pattern_b)
+
+        et.resize_per_segment_(1024, zero_new=True)
+        fv2 = et.full_view()
+        assert fv2.data_ptr() == fv.data_ptr()
+        assert torch.equal(fv2[:256], pattern_a)
+        assert torch.equal(fv2[4096 : 4096 + 256], pattern_b)
+        assert torch.count_nonzero(fv2[256:1024]) == 0
+        assert torch.count_nonzero(fv2[4096 + 256 : 4096 + 1024]) == 0
+
+        seg1 = et.segment_view(1)
+        assert seg1.data_ptr() == fv.data_ptr() + 4096
+        assert seg1.untyped_storage().nbytes() == 1024
+        assert torch.equal(seg1[:256], pattern_b)
+    finally:
+        et.free()
+
+
+@create_new_process_for_each_test("spawn")
+def test_segments_at_granularity_scale() -> None:
+    _skip_unless_vmm()
+    """Segments spanning multiple mapping granules commit correctly.
+
+    Uses a segment capacity that is not a multiple of the allocation
+    granularity, so a granule straddles the segment boundary and is shared by
+    the first commit of one segment and a later commit of the other -- it must
+    be mapped exactly once.
+    """
+    probe = ExtensibleTensor(max_num_bytes=1, device="cuda")
+    granularity = probe.granularity
+    assert probe.capacity_bytes == granularity
+    probe.free()
+    max_num_bytes = 3 * granularity
+    et = ExtensibleTensor(max_num_bytes=max_num_bytes, device="cuda", num_segments=2)
+    try:
+        seg = et.segment_capacity_bytes
+        assert seg == max_num_bytes // 2
+
+        step = granularity // 2
+        et.resize_per_segment_(step, zero_new=True)
+        fv = et.full_view()
+        fv[:step].fill_(1)
+        fv[seg : seg + step].fill_(2)
+
+        et.resize_per_segment_(seg, zero_new=True)
+        assert et.physical_bytes == max_num_bytes
+        fv2 = et.full_view()
+        assert torch.all(fv2[:step] == 1)
+        assert torch.all(fv2[seg : seg + step] == 2)
+        assert torch.count_nonzero(fv2[step:seg]) == 0
+        assert torch.count_nonzero(fv2[seg + step :]) == 0
+    finally:
+        et.free()
+
+
+@create_new_process_for_each_test("spawn")
+def test_unequal_segments_grow_independently() -> None:
+    _skip_unless_vmm()
+    """Segments of different capacities commit different byte counts, each a
+    prefix of its own range, and expose their own views."""
+    et = ExtensibleTensor(
+        max_num_bytes=12288, device="cuda", segment_capacities=[4096, 8192]
+    )
+    try:
+        assert et.num_segments == 2
+        assert et.segment_offsets == [0, 4096]
+        with pytest.raises(ValueError, match="differ"):
+            _ = et.segment_capacity_bytes
+        et.resize_segments_([1024, 2048], zero_new=True)
+        assert et.num_bytes == 3072
+        assert et.segment_view(0).untyped_storage().nbytes() == 1024
+        assert et.segment_view(1).untyped_storage().nbytes() == 2048
+        assert et.segment_view(1).data_ptr() == et.base_ptr + 4096
+        with pytest.raises(ValueError, match="grow-only"):
+            et.resize_segments_([512, 2048])
+        with pytest.raises(ValueError, match="exceeds"):
+            et.resize_segments_([4097, 2048])
+        with pytest.raises(ValueError, match="Expected 2 sizes"):
+            et.resize_segments_([4096])
+    finally:
+        et.free()
+
+    with pytest.raises(ValueError, match="sum to"):
+        ExtensibleTensor(max_num_bytes=100, device="cuda", segment_capacities=[50])
+
+
+@create_new_process_for_each_test("spawn")
+def test_invalid_usage_raises() -> None:
+    _skip_unless_vmm()
+    with pytest.raises(ValueError):
+        ExtensibleTensor(max_num_bytes=100, device="cuda", num_segments=3)
+
+    et = ExtensibleTensor(max_num_bytes=8192, device="cuda", num_segments=2)
+    try:
+        with pytest.raises(IndexError):
+            et.segment_view(2)
+        et.resize_per_segment_(256)
+        with pytest.raises(ValueError):
+            et.resize_per_segment_(128)
+        with pytest.raises(ValueError):
+            et.resize_per_segment_(et.segment_capacity_bytes + 1)
+    finally:
+        et.free()
+
+
+@create_new_process_for_each_test("spawn")
+def test_release_physical_keeps_addresses() -> None:
+    _skip_unless_vmm()
+    """Releasing physical pages keeps the reservation; recommitting maps fresh
+    zeroed pages under the same addresses (the sleep/wake contract)."""
+    et = ExtensibleTensor(max_num_bytes=8192, device="cuda", num_segments=2)
+    try:
+        et.resize_per_segment_(1024, zero_new=True)
+        fv = et.full_view()
+        fv[:1024].fill_(9)
+        base = et.base_ptr
+
+        et.release_physical()
+        assert et.num_bytes == 0
+        assert et.physical_bytes == 0
+        assert et.base_ptr == base
+
+        et.resize_per_segment_(1024, zero_new=True)
+        assert et.full_view().data_ptr() == base
+        assert torch.count_nonzero(et.full_view()[:1024]) == 0
+    finally:
+        et.free()
+
+
+class _FailingDriver:
+    """Delegates to the real driver, recording access grants and failing the
+    n-th physical allocation with an error `_map_chunk_at` does not retry."""
+
+    class Failure(Exception):
+        pass
+
+    def __init__(self, driver, fail_on_create: int) -> None:
+        self._driver = driver
+        self._fail_on_create = fail_on_create
+        self.creates = 0
+        self.access_grants: list[tuple[int, int]] = []
+
+    def __getattr__(self, name):
+        return getattr(self._driver, name)
+
+    def create(self, size: int, device_index: int) -> int:
+        self.creates += 1
+        if self.creates == self._fail_on_create:
+            raise self.Failure("simulated allocation failure")
+        return self._driver.create(size, device_index)
+
+    def set_access(self, ptr: int, size: int, device_index: int) -> None:
+        self.access_grants.append((ptr, size))
+        self._driver.set_access(ptr, size, device_index)
+
+
+@create_new_process_for_each_test("spawn")
+def test_failed_commit_leaves_no_inaccessible_granules() -> None:
+    _skip_unless_vmm()
+    """A commit that maps several runs and fails part-way must have granted
+    access to every run it did map; later commits skip mapped granules, so an
+    unmapped-but-recorded or mapped-but-inaccessible granule would fault."""
+    et = ExtensibleTensor(max_num_bytes=1, device="cuda")
+    g = et.granularity
+    et.free()
+    et = ExtensibleTensor(max_num_bytes=4 * g, device="cuda")
+    try:
+        buffer = et._buffer
+        buffer.ensure_committed_range(g, 2 * g)
+        driver = _FailingDriver(buffer._driver, fail_on_create=2)
+        buffer._driver = driver
+        # Runs [0, g) and [2g, 3g) are needed; the second allocation fails.
+        with pytest.raises(_FailingDriver.Failure):
+            buffer.ensure_committed_range(0, 3 * g)
+        assert buffer._mapped_granules == {0, 1}
+        base = buffer.base_ptr
+        assert any(
+            ptr <= base and base + g <= ptr + size for ptr, size in driver.access_grants
+        ), driver.access_grants
+        # The mapped prefix is usable and a retry completes the range.
+        et.full_view()[: 2 * g].fill_(1)
+        torch.accelerator.synchronize()
+        buffer.ensure_committed_range(0, 3 * g)
+        assert buffer._mapped_granules == {0, 1, 2}
+        et.full_view()[: 3 * g].fill_(2)
+        torch.accelerator.synchronize()
+    finally:
+        et.free()
+
+
+@create_new_process_for_each_test("spawn")
+def test_views_survive_interpreter_teardown() -> None:
+    _skip_unless_vmm()
+    """Views handed to torch through DLPack may be destroyed after this
+    module's globals are cleared at exit; that must not call freed memory."""
+    script = """
+import sys, torch
+from vllm.utils.extensible_tensor import ExtensibleTensor
+et = ExtensibleTensor(4096, device="cuda")
+et.resize_per_segment_(1024)
+# Held by modules imported before the library, so destroyed after it.
+torch._extkv_test_view = et.full_view()
+sys._extkv_test_view = et.segment_view(0)
+"""
+    result = subprocess.run(
+        [sys.executable, "-c", script], capture_output=True, text=True, timeout=300
+    )
+    assert result.returncode == 0, result.stderr[-2000:]
+
+
+@create_new_process_for_each_test("spawn")
+def test_vmm_probe_reports_usable_driver() -> None:
+    reason = vmm_unavailable_reason()
+    print(f"vmm_unavailable_reason reported as: {reason}")
+    if torch.version.hip is None:
+        assert reason is None
+    else:
+        # ROCm runtimes older than the hipMemSetAccess fix are rejected.
+        assert reason is None or "HIP runtime" in reason
diff --git a/tests/utils_/test_numa_utils.py b/tests/utils_/test_numa_utils.py
index 3c09722f82df..014e2f7da7dc 100644
--- a/tests/utils_/test_numa_utils.py
+++ b/tests/utils_/test_numa_utils.py
@@ -2,8 +2,10 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import os
+from copy import deepcopy
 from functools import partial
 from types import SimpleNamespace
+from typing import Any
 
 import pytest
 
@@ -39,20 +41,30 @@ def _no_pct_open(path, *args, **kwargs):
 
 
 def _make_config(**parallel_kwargs):
-    parallel_defaults = dict(
+    parallel_defaults: dict[str, Any] = dict(
         numa_bind=False,
+        _numa_bind_auto=False,
         numa_bind_nodes=None,
         numa_bind_cpus=None,
+        assigned_physical_gpu_ids=None,
         distributed_executor_backend="mp",
         data_parallel_backend="mp",
         nnodes_within_dp=1,
         data_parallel_rank_local=0,
         data_parallel_index=0,
+        data_parallel_size_local=1,
         pipeline_parallel_size=1,
         tensor_parallel_size=1,
         prefill_context_parallel_size=1,
     )
     parallel_defaults.update(parallel_kwargs)
+    parallel_defaults.setdefault(
+        "local_world_size",
+        parallel_defaults["pipeline_parallel_size"]
+        * parallel_defaults["tensor_parallel_size"]
+        * parallel_defaults["prefill_context_parallel_size"]
+        // parallel_defaults["nnodes_within_dp"],
+    )
     parallel_config = SimpleNamespace(**parallel_defaults)
     parallel_config.world_size = (
         parallel_config.tensor_parallel_size
@@ -214,12 +226,18 @@ def test_pct_binding_returns_none_when_node_cpulist_missing(monkeypatch):
     assert numa_utils._maybe_get_pct_cpu_binding([0]) is None
 
 
-def test_get_numactl_args_uses_pct_when_user_did_not_specify_cpus(monkeypatch):
+@pytest.mark.parametrize("automatic", [False, True])
+def test_get_numactl_args_uses_pct_when_user_did_not_specify_cpus(
+    monkeypatch, automatic
+):
     _patch_pct_gates(monkeypatch, model_match=True, highest_perf=46)
-    vllm_config = _make_config(numa_bind=True, numa_bind_nodes=[0, 1])
+    vllm_config = _make_config(
+        numa_bind=True, _numa_bind_auto=automatic, numa_bind_nodes=[0, 1]
+    )
+    memory_policy = "preferred" if automatic else "membind"
     assert (
         numa_utils._get_numactl_worker_args(vllm_config.parallel_config, local_rank=1)
-        == "--physcpubind=0,1,16,17,64,65,80,81 --membind=1"
+        == f"--physcpubind=0,1,16,17,64,65,80,81 --{memory_policy}=1"
     )
 
 
@@ -249,7 +267,10 @@ def test_get_numactl_args_engine_core_baseline_spans_shard_numa_nodes():
     )
 
 
-def test_get_numactl_args_engine_core_pct_spans_shard_numa_nodes(monkeypatch):
+@pytest.mark.parametrize("automatic", [False, True])
+def test_get_numactl_args_engine_core_pct_spans_shard_numa_nodes(
+    monkeypatch, automatic
+):
     """PCT: EngineCore for a multi-NUMA shard binds to the union of priority
     cores across all shard nodes, so worker `--physcpubind` is always a
     subset of EngineCore's `cpus_allowed`."""
@@ -261,15 +282,17 @@ def test_get_numactl_args_engine_core_pct_spans_shard_numa_nodes(monkeypatch):
     )
     vllm_config = _make_config(
         numa_bind=True,
+        _numa_bind_auto=automatic,
         numa_bind_nodes=[0, 0, 1, 1],
         tensor_parallel_size=4,
     )
+    memory_policy = "preferred-many" if automatic else "membind"
     assert numa_utils._get_numactl_enginecore_args(
         vllm_config.parallel_config, local_rank=0
     ) == (
         "--physcpubind="
         "0,1,16,17,64,65,80,81,128,129,144,145,192,193,208,209"
-        " --membind=0,1"
+        f" --{memory_policy}=0,1"
     )
 
 
@@ -543,6 +566,18 @@ def test_parallel_config_validates_numa_bind_nodes():
         ParallelConfig(numa_bind_nodes=[0, -1])
 
 
+@pytest.mark.parametrize(
+    ("nodes", "cpus"), [([0, 1], None), (None, ["0-3"])], ids=["nodes", "cpus"]
+)
+def test_parallel_config_numa_lists_imply_numa_bind(nodes, cpus):
+    assert ParallelConfig().numa_bind is None
+    config = ParallelConfig(numa_bind_nodes=nodes, numa_bind_cpus=cpus)
+    assert config.numa_bind is True
+    assert config._numa_bind_auto is False
+    with pytest.raises(ValueError, match="cannot be combined with numa_bind=False"):
+        ParallelConfig(numa_bind=False, numa_bind_nodes=nodes, numa_bind_cpus=cpus)
+
+
 @pytest.mark.parametrize("cpuset", ["", "abc", "1-", "4-1", "1,,2", "1:2"])
 def test_parallel_config_rejects_invalid_numa_bind_cpus(cpuset):
     with pytest.raises(ValueError, match="numa_bind_cpus"):
@@ -564,7 +599,7 @@ def test_configure_subprocess_numa_fallback(monkeypatch):
     import multiprocessing
 
     monkeypatch.setattr("shutil.which", lambda name: "/usr/bin/numactl")
-    monkeypatch.setattr(numa_utils.envs, "VLLM_WORKER_MULTIPROC_METHOD", "spawn")
+    monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
     node_config = _make_config(numa_bind=True, numa_bind_nodes=[0])
 
     monkeypatch.setattr(numa_utils.subprocess, "run", _fake_numactl_run([]))
@@ -593,3 +628,409 @@ def test_configure_subprocess_numa_fallback(monkeypatch):
     with numa_utils.configure_subprocess(node_config, local_rank=0):
         assert multiprocessing.spawn.get_executable() == before
         assert numa_utils._NUMACTL_ARGS_ENV not in os.environ
+
+
+def _patch_auto_numa(
+    monkeypatch, start_method, detect, visible_gpus=1, node_cpus="0-7"
+):
+    monkeypatch.setattr("shutil.which", lambda name: "/usr/bin/numactl")
+    monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", start_method)
+    monkeypatch.setattr(numa_utils, "_is_auto_numa_available", lambda **kwargs: True)
+    monkeypatch.setattr(numa_utils, "get_auto_numa_nodes", detect)
+    monkeypatch.setattr(
+        "vllm.platforms.current_platform.device_count", lambda: visible_gpus
+    )
+    real_read_text = numa_utils.Path.read_text
+
+    def read_text(path, *args, **kwargs):
+        if str(path).startswith("/sys/devices/system/node/node"):
+            return node_cpus
+        return real_read_text(path, *args, **kwargs)
+
+    monkeypatch.setattr(numa_utils.Path, "read_text", read_text)
+
+
+@pytest.mark.parametrize("nodes", [[0], [0, 1]])
+def test_configure_subprocess_auto_binds_and_records_result(monkeypatch, nodes):
+    """Default mode resolves once in the launcher and stores the result, so
+    spawned children bind with it instead of re-checking eligibility."""
+    _patch_auto_numa(monkeypatch, "spawn", lambda: nodes, visible_gpus=len(nodes))
+    monkeypatch.setattr(numa_utils.subprocess, "run", _fake_numactl_run([]))
+
+    vllm_config = _make_config(numa_bind=None, tensor_parallel_size=len(nodes))
+    node_list = ",".join(map(str, nodes))
+    preference = "preferred" if len(nodes) == 1 else "preferred-many"
+    with numa_utils.configure_subprocess(
+        vllm_config, local_rank=0, process_kind="EngineCore"
+    ):
+        assert os.environ[numa_utils._NUMACTL_ARGS_ENV] == (
+            f"--cpunodebind={node_list} --{preference}={node_list}"
+        )
+        child_config = deepcopy(vllm_config)
+    assert vllm_config.parallel_config.numa_bind is True
+    assert child_config.parallel_config._numa_bind_auto is True
+    assert child_config.parallel_config.numa_bind_nodes == nodes
+
+    monkeypatch.setattr(numa_utils, "get_auto_numa_nodes", _raise_probe_error)
+    monkeypatch.setattr(numa_utils, "_is_auto_numa_available", lambda **kwargs: False)
+    with numa_utils.configure_subprocess(child_config, local_rank=len(nodes) - 1):
+        assert os.environ[numa_utils._NUMACTL_ARGS_ENV] == (
+            f"--cpunodebind={nodes[-1]} --preferred={nodes[-1]}"
+        )
+
+
+def test_parallel_config_auto_policy_survives_copy_without_changing_hash():
+    """Copying the launch decision must not change the DP configuration hash."""
+    from vllm.config.utils import replace
+
+    parallel_config = ParallelConfig()
+    original_hash = parallel_config.compute_hash()
+    parallel_config.numa_bind = True
+    parallel_config.numa_bind_nodes = [0]
+    parallel_config._numa_bind_auto = True
+    for copied in (
+        replace(parallel_config),
+        deepcopy(parallel_config),
+    ):
+        assert copied._numa_bind_auto is True
+        assert copied.compute_hash() == original_hash
+
+
+@pytest.mark.parametrize("dp_local_rank", [0, 1])
+def test_auto_numa_accepts_pcp_dp_topology(monkeypatch, dp_local_rank):
+    """Automatic eligibility must cover every PCP worker in both DP shards."""
+    _patch_auto_numa(monkeypatch, "spawn", lambda: [0, 0, 1, 1], visible_gpus=4)
+    config = ParallelConfig(
+        distributed_executor_backend="mp",
+        prefill_context_parallel_size=2,
+        data_parallel_size=2,
+        data_parallel_size_local=2,
+    )
+    assert numa_utils._can_auto_bind(config, local_rank=0, dp_local_rank=dp_local_rank)
+    assert config.numa_bind_nodes == [0, 0, 1, 1]
+
+
+def test_auto_numa_redetects_when_visible_devices_change(monkeypatch):
+    """Independent launchers must not reuse another visible GPU set's mapping."""
+    _patch_auto_numa(
+        monkeypatch, "spawn", numa_utils.get_auto_numa_nodes, visible_gpus=2
+    )
+    monkeypatch.setattr(numa_utils.subprocess, "run", _fake_numactl_run([]))
+    detections = []
+
+    def detect():
+        devices = os.environ["CUDA_VISIBLE_DEVICES"]
+        detections.append(devices)
+        return {"0,1": [0, 0], "2,3": [1, 1]}[devices]
+
+    monkeypatch.setattr(
+        "vllm.platforms.current_platform.get_all_device_numa_nodes",
+        detect,
+        raising=False,
+    )
+    for devices, node in [("0,1", 0), ("2,3", 1)]:
+        monkeypatch.setenv("CUDA_VISIBLE_DEVICES", devices)
+        vllm_config = _make_config(numa_bind=None, tensor_parallel_size=2)
+        with numa_utils.configure_subprocess(
+            vllm_config, local_rank=0, process_kind="EngineCore"
+        ):
+            numactl_args = os.environ[numa_utils._NUMACTL_ARGS_ENV]
+            assert numactl_args == f"--cpunodebind={node} --preferred={node}"
+        assert vllm_config.parallel_config.numa_bind_nodes == [node, node]
+    assert detections == ["0,1", "2,3"]
+
+
+@pytest.mark.parametrize("explicit", [False, True], ids=["automatic", "explicit"])
+@pytest.mark.parametrize(
+    "policy",
+    [0, 1, 2, 3, 4, 5, None],
+    ids=[
+        "default",
+        "preferred",
+        "bind",
+        "interleave",
+        "local",
+        "preferred-many",
+        "unreadable",
+    ],
+)
+def test_configure_subprocess_respects_existing_mempolicy(
+    monkeypatch, policy, explicit
+):
+    """Only explicit enablement may replace an existing non-default policy."""
+    is_auto_numa_available = numa_utils._is_auto_numa_available
+    _patch_auto_numa(monkeypatch, "spawn", numa_utils.get_auto_numa_nodes)
+    monkeypatch.setattr(numa_utils, "_is_auto_numa_available", is_auto_numa_available)
+    monkeypatch.setattr("vllm.platforms.current_platform.is_cuda_alike", lambda: True)
+    monkeypatch.setattr(numa_utils.os.path, "isdir", lambda path: True)
+    monkeypatch.setattr(numa_utils.psutil, "cpu_count", lambda: 4)
+    monkeypatch.setattr(
+        numa_utils.psutil,
+        "Process",
+        lambda pid: SimpleNamespace(cpu_affinity=lambda: [0, 1, 2, 3]),
+    )
+    monkeypatch.setattr(
+        "vllm.platforms.current_platform.get_all_device_numa_nodes",
+        lambda: [0],
+        raising=False,
+    )
+
+    def get_mempolicy(mode, *_args):
+        if policy is None:
+            return -1
+        mode._obj.value = policy
+        return 0
+
+    monkeypatch.setattr(
+        numa_utils,
+        "get_libnuma",
+        lambda: SimpleNamespace(numa_available=lambda: 0, get_mempolicy=get_mempolicy),
+    )
+    monkeypatch.setattr(numa_utils.subprocess, "run", _fake_numactl_run([]))
+    vllm_config = _make_config(numa_bind=True if explicit else None)
+    if explicit and policy is None:
+        with (
+            pytest.raises(RuntimeError, match="could not detect"),
+            numa_utils.configure_subprocess(vllm_config, local_rank=0),
+        ):
+            pass
+        return
+
+    expected_bind = policy is not None and (policy == 0 or explicit)
+    memory_arg = "--membind=0" if explicit else "--preferred=0"
+    with numa_utils.configure_subprocess(vllm_config, local_rank=0):
+        numactl_args = os.environ.get(numa_utils._NUMACTL_ARGS_ENV)
+        assert numactl_args == (
+            f"--cpunodebind=0 {memory_arg}" if expected_bind else None
+        )
+    assert vllm_config.parallel_config.numa_bind is expected_bind
+
+
+def _raise_probe_error():
+    raise RuntimeError("amdsmi init failed")
+
+
+@pytest.mark.parametrize(
+    ("start_method", "detect", "assigned_gpus"),
+    [
+        ("fork", lambda: [0], None),
+        ("spawn", lambda: None, None),
+        ("spawn", lambda: [], None),
+        ("spawn", _raise_probe_error, None),
+        ("spawn", lambda: [0, 1], [1]),
+    ],
+    ids=["fork", "no-topology", "no-visible-gpus", "probe-raises", "device-ids"],
+)
+def test_configure_subprocess_auto_falls_back_to_unbound(
+    monkeypatch, start_method, detect, assigned_gpus
+):
+    """Default mode never fails startup or forces spawn; when binding cannot
+    engage (or would pick the wrong node) it starts unbound."""
+    import multiprocessing
+
+    _patch_auto_numa(monkeypatch, start_method, detect)
+
+    vllm_config = _make_config(numa_bind=None, assigned_physical_gpu_ids=assigned_gpus)
+    before = multiprocessing.spawn.get_executable()
+    with numa_utils.configure_subprocess(vllm_config, local_rank=0):
+        assert multiprocessing.spawn.get_executable() == before
+    assert vllm_config.parallel_config.numa_bind is False
+
+
+@pytest.mark.parametrize(
+    "parallel_kwargs",
+    [
+        {"distributed_executor_backend": "ray"},
+        {"data_parallel_backend": "ray"},
+        {"distributed_executor_backend": "external_launcher"},
+    ],
+    ids=["ray-executor", "ray-dp", "external-launcher"],
+)
+def test_auto_numa_skips_unsupported_launchers(monkeypatch, parallel_kwargs):
+    detections = []
+
+    def detect():
+        detections.append(True)
+        return [0, 1]
+
+    _patch_auto_numa(monkeypatch, "spawn", detect, visible_gpus=2)
+    vllm_config = _make_config(numa_bind=None, **parallel_kwargs)
+    with numa_utils.configure_subprocess(vllm_config, local_rank=0):
+        assert numa_utils._NUMACTL_ARGS_ENV not in os.environ
+    assert vllm_config.parallel_config.numa_bind is False
+    assert detections == []
+
+
+@pytest.mark.parametrize("nodes", [[0], [0, 1, 2], [0, -1], [0, None]])
+def test_auto_numa_rejects_incomplete_or_invalid_topology(monkeypatch, nodes):
+    _patch_auto_numa(monkeypatch, "spawn", lambda: nodes, visible_gpus=2)
+    vllm_config = _make_config(numa_bind=None, tensor_parallel_size=2)
+    with numa_utils.configure_subprocess(vllm_config, local_rank=0):
+        assert numa_utils._NUMACTL_ARGS_ENV not in os.environ
+    assert vllm_config.parallel_config.numa_bind is False
+    assert vllm_config.parallel_config.numa_bind_nodes is None
+
+
+def test_auto_numa_checks_all_local_dp_workers_before_first_engine(monkeypatch):
+    _patch_auto_numa(monkeypatch, "spawn", lambda: [0, 1], visible_gpus=2)
+    vllm_config = _make_config(
+        numa_bind=None, tensor_parallel_size=2, data_parallel_size_local=2
+    )
+    with numa_utils.configure_subprocess(
+        vllm_config, local_rank=0, dp_local_rank=0, process_kind="EngineCore"
+    ):
+        assert numa_utils._NUMACTL_ARGS_ENV not in os.environ
+    assert vllm_config.parallel_config.numa_bind is False
+
+
+def test_auto_numa_multi_node_topology_is_local(monkeypatch):
+    _patch_auto_numa(monkeypatch, "spawn", lambda: [2, 3], visible_gpus=2)
+    monkeypatch.setattr(numa_utils.subprocess, "run", _fake_numactl_run([]))
+    vllm_config = _make_config(
+        numa_bind=None, tensor_parallel_size=4, nnodes_within_dp=2
+    )
+    with numa_utils.configure_subprocess(
+        vllm_config, local_rank=0, process_kind="EngineCore"
+    ):
+        assert os.environ[numa_utils._NUMACTL_ARGS_ENV] == (
+            "--cpunodebind=2,3 --preferred-many=2,3"
+        )
+        child_config = deepcopy(vllm_config)
+    with numa_utils.configure_subprocess(child_config, local_rank=1):
+        assert (
+            os.environ[numa_utils._NUMACTL_ARGS_ENV] == "--cpunodebind=3 --preferred=3"
+        )
+
+
+@pytest.mark.parametrize("node_cpus", [None, "", "invalid"])
+def test_auto_numa_skips_nodes_without_usable_cpus(monkeypatch, node_cpus):
+    _patch_auto_numa(monkeypatch, "spawn", lambda: [0], node_cpus=node_cpus)
+    vllm_config = _make_config(numa_bind=None)
+    with numa_utils.configure_subprocess(vllm_config, local_rank=0):
+        assert numa_utils._NUMACTL_ARGS_ENV not in os.environ
+    assert vllm_config.parallel_config.numa_bind is False
+
+
+@pytest.mark.parametrize(
+    "nodes, rejected_args, expected_args",
+    [
+        ([0], "--preferred=", "--cpunodebind=0"),
+        ([0], "--cpunodebind=", None),
+        ([0, 1], "--preferred-many=", "--cpunodebind=0,1"),
+        ([0, 1], "--cpunodebind=", None),
+    ],
+)
+def test_auto_numa_records_effective_policy(
+    monkeypatch, caplog, nodes, rejected_args, expected_args
+):
+    caplog.set_level("INFO", logger=numa_utils.logger.name)
+    _patch_auto_numa(monkeypatch, "spawn", lambda: nodes, visible_gpus=len(nodes))
+    probes = []
+    fake_run = _fake_numactl_run([rejected_args])
+
+    def run(cmd, *args, **kwargs):
+        probes.extend(cmd)
+        return fake_run(cmd, *args, **kwargs)
+
+    monkeypatch.setattr(numa_utils.subprocess, "run", run)
+    vllm_config = _make_config(numa_bind=None, tensor_parallel_size=len(nodes))
+    with numa_utils.configure_subprocess(
+        vllm_config, local_rank=0, process_kind="EngineCore"
+    ):
+        assert os.environ.get(numa_utils._NUMACTL_ARGS_ENV) == expected_args
+        assert vllm_config.parallel_config.numa_bind is (expected_args is not None)
+    launch_logs = [r.getMessage() for r in caplog.records if "Launching" in r.message]
+    if expected_args is None:
+        assert launch_logs == []
+        assert vllm_config.parallel_config.numa_bind_nodes is None
+        assert vllm_config.parallel_config._numa_bind_auto is False
+    else:
+        assert len(launch_logs) == 1
+        assert launch_logs[0].endswith(f"numactl {expected_args}")
+    assert not any(arg.startswith("--membind=") for arg in probes)
+
+
+def test_auto_numa_does_not_swallow_launch_errors(monkeypatch):
+    import multiprocessing
+
+    _patch_auto_numa(monkeypatch, "spawn", lambda: [0])
+    monkeypatch.setattr(numa_utils.subprocess, "run", _fake_numactl_run([]))
+    monkeypatch.delenv(numa_utils._NUMACTL_ARGS_ENV, raising=False)
+    monkeypatch.delenv(numa_utils._NUMACTL_PYTHON_EXECUTABLE_ENV, raising=False)
+    original_executable = multiprocessing.spawn.get_executable()
+    with (
+        pytest.raises(RuntimeError, match="worker startup failed"),
+        numa_utils.configure_subprocess(_make_config(numa_bind=None), local_rank=0),
+    ):
+        raise RuntimeError("worker startup failed")
+    assert multiprocessing.spawn.get_executable() == original_executable
+    assert numa_utils._NUMACTL_ARGS_ENV not in os.environ
+    assert numa_utils._NUMACTL_PYTHON_EXECUTABLE_ENV not in os.environ
+
+
+@pytest.mark.parametrize("missing_numactl", [False, True])
+def test_auto_numa_skips_unavailable_environment(monkeypatch, missing_numactl):
+    _patch_auto_numa(monkeypatch, "spawn", lambda: [0])
+    if missing_numactl:
+        monkeypatch.setattr("shutil.which", lambda name: None)
+    else:
+        monkeypatch.setattr(
+            numa_utils, "_is_auto_numa_available", lambda **kwargs: False
+        )
+    vllm_config = _make_config(numa_bind=None)
+    with numa_utils.configure_subprocess(vllm_config, local_rank=0):
+        assert numa_utils._NUMACTL_ARGS_ENV not in os.environ
+    assert vllm_config.parallel_config.numa_bind is False
+
+
+def test_explicit_numa_nodes_bypass_auto_eligibility(monkeypatch):
+    _patch_auto_numa(monkeypatch, "spawn", _raise_probe_error)
+    monkeypatch.setattr(numa_utils, "_is_auto_numa_available", lambda **kwargs: False)
+    monkeypatch.setattr(numa_utils.subprocess, "run", _fake_numactl_run([]))
+    vllm_config = _make_config(numa_bind=True, numa_bind_nodes=[7])
+    with numa_utils.configure_subprocess(vllm_config, local_rank=0):
+        assert os.environ[numa_utils._NUMACTL_ARGS_ENV] == "--cpunodebind=7 --membind=7"
+
+
+def test_explicit_numa_requires_detectable_topology(monkeypatch):
+    _patch_auto_numa(monkeypatch, "spawn", lambda: None)
+    with (
+        pytest.raises(RuntimeError, match="could not detect"),
+        numa_utils.configure_subprocess(_make_config(numa_bind=True), local_rank=0),
+    ):
+        pass
+
+
+@pytest.mark.parametrize("numa_bind", [False, True])
+def test_explicit_numa_handles_missing_numactl(monkeypatch, numa_bind):
+    monkeypatch.setattr("shutil.which", lambda name: None)
+    vllm_config = _make_config(numa_bind=numa_bind, numa_bind_nodes=[0])
+    if numa_bind:
+        with (
+            pytest.raises(RuntimeError, match="numactl is required"),
+            numa_utils.configure_subprocess(vllm_config, local_rank=0),
+        ):
+            pass
+    else:
+        with numa_utils.configure_subprocess(vllm_config, local_rank=0):
+            assert numa_utils._NUMACTL_ARGS_ENV not in os.environ
+
+
+@pytest.mark.parametrize("affinity", [[0, 1], None, [3, 2, 1, 0]])
+def test_auto_numa_checks_cpu_affinity(monkeypatch, affinity):
+    monkeypatch.setattr("vllm.platforms.current_platform.is_cuda_alike", lambda: True)
+    monkeypatch.setattr(numa_utils.os.path, "isdir", lambda path: True)
+    monkeypatch.setattr(numa_utils.psutil, "cpu_count", lambda: 4)
+    process = SimpleNamespace()
+    if affinity is not None:
+        process.cpu_affinity = lambda: affinity
+    monkeypatch.setattr(numa_utils.psutil, "Process", lambda pid: process)
+    monkeypatch.setattr(numa_utils, "_get_mempolicy", lambda: 0)
+    monkeypatch.setattr(
+        "vllm.platforms.current_platform.get_all_device_numa_nodes",
+        lambda: [0, 1],
+        raising=False,
+    )
+    assert numa_utils._is_auto_numa_available(
+        verbose=False, require_unbound_process=True
+    ) is (affinity == [3, 2, 1, 0])
diff --git a/tests/utils_/test_system_utils.py b/tests/utils_/test_system_utils.py
index 592af714a5ba..25d4ef4b5d1d 100644
--- a/tests/utils_/test_system_utils.py
+++ b/tests/utils_/test_system_utils.py
@@ -11,6 +11,7 @@
 from pathlib import Path
 
 import psutil
+import pytest
 
 from vllm.utils import system_utils
 from vllm.utils.system_utils import (
@@ -32,11 +33,39 @@ def test_unique_filepath():
     assert len(list(Path(temp_dir).glob("*.txt"))) == 10
 
 
-def test_numa_bind_forces_spawn(monkeypatch):
-    monkeypatch.delenv("VLLM_WORKER_MULTIPROC_METHOD", raising=False)
-    monkeypatch.setattr("sys.argv", ["vllm", "serve", "--numa-bind"])
+@pytest.mark.parametrize("start_method", [None, "fork"])
+@pytest.mark.parametrize(
+    "numa_args, force_spawn",
+    [
+        (["--numa-bind"], True),
+        (["--numa-bind-nodes", "0"], True),
+        (["--numa-bind-nodes=0"], True),
+        (["--numa-bind-cpus", "0-3"], True),
+        (["--numa-bind-cpus=0-3"], True),
+        (["--numa_bind"], True),
+        (["--numa_bind_nodes", "0"], True),
+        (["--numa_bind_cpus=0-3"], True),
+        ([], False),
+        (["--no-numa-bind"], False),
+    ],
+)
+def test_numa_bind_cli_start_method(monkeypatch, start_method, numa_args, force_spawn):
+    if start_method is None:
+        monkeypatch.delenv("VLLM_WORKER_MULTIPROC_METHOD", raising=False)
+    else:
+        monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", start_method)
+    monkeypatch.setattr(sys, "argv", ["vllm", "serve", *numa_args])
+    for check in (
+        "is_in_ray_actor",
+        "cuda_is_initialized",
+        "xpu_is_initialized",
+        "in_wsl",
+    ):
+        monkeypatch.setattr(system_utils, check, lambda: False)
+
     _maybe_force_spawn()
-    assert os.environ["VLLM_WORKER_MULTIPROC_METHOD"] == "spawn"
+    expected_method = "spawn" if force_spawn else start_method
+    assert os.environ.get("VLLM_WORKER_MULTIPROC_METHOD") == expected_method
 
 
 def test_suppress_stdout_when_sys_stdout_has_no_fd(capfd):
diff --git a/tests/utils_/test_vmm_driver.py b/tests/utils_/test_vmm_driver.py
new file mode 100644
index 000000000000..98fc5a01e4e6
--- /dev/null
+++ b/tests/utils_/test_vmm_driver.py
@@ -0,0 +1,141 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import subprocess
+import sys
+
+import pytest
+import torch
+
+from vllm.utils.vmm_driver import _find_loaded_library
+
+_CUDART = "/usr/lib/python3.12/site-packages/nvidia/cuda_runtime/lib/libcudart.so.12"
+_MAPS = [
+    f"7f00-7f01 r-xp 00000000 08:01 1 {_CUDART}",
+    "7f02-7f03 r-xp 00000000 08:01 2 [heap]",
+    "7f04-7f05 r-xp 00000000 08:01 3 /usr/lib64/libcuda.so.580.65.06",
+    "7f06-7f07 r-xp 00000000 08:01 4 /usr/lib64/libcuda.so.1",
+]
+
+
+def test_find_loaded_library_matches_the_file_name_not_a_substring():
+    """Torch maps libcudart before the driver; the driver lookup must not pick
+    it (its symbols differ), nor match a path component."""
+    assert _find_loaded_library("libcuda", _MAPS) == "/usr/lib64/libcuda.so.580.65.06"
+    assert _find_loaded_library("libcudart", _MAPS) == _CUDART
+    assert _find_loaded_library("libamdhip64", _MAPS) is None
+    assert _find_loaded_library("libcuda", ["garbage"]) is None
+
+
+@pytest.mark.skipif(
+    not torch.cuda.is_available() or torch.version.hip is not None,
+    reason="CUDA driver context handling",
+)
+def test_cuda_driver_context_usable_before_torch_initializes_the_device():
+    """The VMM probe may run before torch touches the device; the driver must
+    initialize itself rather than fail with a bare 'not initialized' error."""
+    script = """
+from vllm.utils.vmm_driver import get_vmm_driver
+driver = get_vmm_driver()
+driver.ensure_context(0)
+driver.ensure_context(0)
+print(driver.granularity(0))
+"""
+    result = subprocess.run(
+        [sys.executable, "-c", script], capture_output=True, text=True, timeout=300
+    )
+    assert result.returncode == 0, result.stderr[-2000:]
+    assert int(result.stdout.strip().splitlines()[-1]) > 0
+
+
+class _SetAccessModel:
+    """Models HIP's set-access size check over a reservation's allocations,
+    kept in list_order (the runtime's internal order, which callers do not
+    control)."""
+
+    granule = 4096
+
+    def __init__(self, fixed: bool, reverse_list_order: bool) -> None:
+        self.fixed = fixed
+        self.reverse = reverse_list_order
+        self.chunks: list[tuple[int, int]] = []
+        self.live: set[str] = set()
+
+    def granularity(self, device_index: int) -> int:
+        return self.granule
+
+    def reserve(self, size: int) -> int:
+        self.live.add("reservation")
+        return 1 << 20
+
+    def free_reserved(self, ptr: int, size: int) -> None:
+        self.live.discard("reservation")
+
+    def create(self, size: int, device_index: int) -> int:
+        self.live.add(f"handle{size}")
+        return size
+
+    def release(self, handle: int) -> None:
+        self.live.discard(f"handle{handle}")
+
+    def map(self, ptr: int, size: int, handle: int) -> None:
+        self.chunks.append((ptr, size))
+
+    def unmap(self, ptr: int, size: int) -> None:
+        self.chunks.remove((ptr, size))
+
+    def set_access(self, ptr: int, size: int, device_index: int) -> None:
+        if self.fixed:
+            start = [p for p, _ in self.chunks].index(ptr)
+            ordered = sorted(self.chunks)[start:]
+        else:
+            ordered = self.chunks[::-1] if self.reverse else list(self.chunks)
+        covered = 0
+        for _, chunk_size in ordered:
+            covered += chunk_size
+            if covered >= size:
+                break
+        if covered != size:
+            raise RuntimeError("GPU driver error 1: invalid argument")
+
+
+@pytest.mark.parametrize("reverse_list_order", [False, True])
+@pytest.mark.parametrize("fixed", [False, True])
+def test_set_access_probe_detects_the_hip_defect(fixed, reverse_list_order):
+    """The probe flags runtimes that count a reservation's allocations from the
+    first in their list (ROCm/rocm-systems#2516) whatever that list's order,
+    passes fixed ones, and frees everything it maps either way."""
+    from vllm.utils.vmm_driver import _set_access_fails_within_reservation
+
+    model = _SetAccessModel(fixed, reverse_list_order)
+    assert _set_access_fails_within_reservation(model, 0) is (not fixed)
+    assert model.chunks == []
+    assert model.live == set()
+
+
+def test_driver_prefers_the_library_already_loaded(monkeypatch):
+    """A ROCm torch wheel bundles its own HIP runtime; loading another copy by
+    name would put two runtimes in the process, so the mapped one wins."""
+    from types import SimpleNamespace
+
+    from vllm.utils import vmm_driver
+
+    opened: list[str] = []
+
+    def fake_cdll(name: str) -> str:
+        opened.append(name)
+        return name
+
+    monkeypatch.setattr(vmm_driver.ctypes, "CDLL", fake_cdll)
+    driver = SimpleNamespace(
+        _lib_search_name="libamdhip64",
+        _lib_candidates=("libamdhip64.so.7", "libamdhip64.so"),
+    )
+    bundled = "/venv/lib/python3.12/site-packages/torch/lib/libamdhip64.so.7"
+    monkeypatch.setattr(vmm_driver, "_find_loaded_library", lambda name: bundled)
+    assert vmm_driver.VmmDriver._load_library(driver) == bundled
+
+    monkeypatch.setattr(vmm_driver, "_find_loaded_library", lambda name: None)
+    opened.clear()
+    assert vmm_driver.VmmDriver._load_library(driver) == "libamdhip64.so.7"
+    assert opened == ["libamdhip64.so.7"]
diff --git a/tests/v1/attention/test_attention_backends.py b/tests/v1/attention/test_attention_backends.py
index 0f911594931b..26eac96e26eb 100644
--- a/tests/v1/attention/test_attention_backends.py
+++ b/tests/v1/attention/test_attention_backends.py
@@ -864,6 +864,79 @@ def causal_mask_mod(
         )
 
 
+@pytest.mark.skipif(
+    not current_platform.is_device_capability_family(100),
+    reason="FA4 takes FP8 Q/K/V with descales on Blackwell.",
+)
+def test_flash_attn_fa4_unit_fp8_descale_is_exact(default_vllm_config):
+    """FlashAttentionImpl passes FA4 no descale for a per-tensor scale of 1.0,
+    which FA4 reads as exactly 1.0; calibrated scales keep their descale."""
+    from vllm.v1.attention.backends.fa_utils import (
+        flash_attn_varlen_func,
+        is_fa_version_supported,
+    )
+    from vllm.v1.attention.backends.flash_attn import FlashAttentionImpl
+
+    if not is_fa_version_supported(4):
+        pytest.skip("FA4 is not supported.")
+    set_random_seed(0)
+    device = torch.device(f"{DEVICE_TYPE}:0")
+    fp8 = current_platform.fp8_dtype()
+    seq_lens, q_len = [1328, 523, 2011], 3
+    num_heads, num_kv_heads, head_size = 64, 4, 128
+    impl = FlashAttentionImpl(
+        num_heads=num_heads,
+        head_size=head_size,
+        scale=head_size**-0.5,
+        num_kv_heads=num_kv_heads,
+        alibi_slopes=None,
+        sliding_window=None,
+        kv_cache_dtype="fp8",
+    )
+    assert impl.vllm_flash_attn_version == 4
+
+    q = torch.randn(len(seq_lens) * q_len, num_heads, head_size, device=device)
+    k_cache, v_cache = (
+        torch.randn(64, 128, num_kv_heads, head_size, device=device).to(fp8)
+        for _ in range(2)
+    )
+    block_table = torch.randperm(64, dtype=torch.int32, device=device)[:48]
+
+    def run(descale: torch.Tensor | None) -> torch.Tensor:
+        return flash_attn_varlen_func(
+            q=q.to(fp8),
+            k=k_cache,
+            v=v_cache,
+            cu_seqlens_q=torch.arange(
+                0, q.shape[0] + 1, q_len, dtype=torch.int32, device=device
+            ),
+            seqused_k=torch.tensor(seq_lens, dtype=torch.int32, device=device),
+            max_seqlen_q=q_len,
+            max_seqlen_k=max(seq_lens),
+            softmax_scale=head_size**-0.5,
+            causal=True,
+            block_table=block_table.view(len(seq_lens), -1),
+            fa_version=4,
+            q_descale=descale,
+            k_descale=descale,
+            v_descale=descale,
+        )
+
+    shape = (len(seq_lens), num_kv_heads)
+    unit = torch.ones(1, device=device)
+    assert impl._descale(unit, 1.0, shape) is None
+    expanded = unit.expand(shape)  # What FA4 got for a unit scale before.
+    output = run(impl._descale(unit, 1.0, shape))
+    torch.testing.assert_close(output, run(expanded), atol=0, rtol=0)
+    # Descales are applied when given, so the comparison above is not vacuous.
+    assert not torch.equal(output, run(2 * expanded))
+    # Calibrated and per-head scales (whose float holds the max) keep theirs.
+    calibrated = impl._descale(0.5 * unit, 0.5, shape)
+    assert calibrated is not None and torch.equal(calibrated, 0.5 * expanded)
+    per_head = torch.tensor([1.0, 0.5, 1.0, 0.25], device=device)
+    assert impl._descale(per_head, 1.0, shape) is not None
+
+
 @pytest.mark.skipif(
     AttentionBackendEnum.FLASHINFER not in BACKENDS_TO_TEST,
     reason="FlashInfer is not available.",
@@ -1026,6 +1099,117 @@ def mock_xqa(**kwargs):
     }
 
 
+@pytest.mark.skipif(
+    AttentionBackendEnum.FLASHINFER not in BACKENDS_TO_TEST,
+    reason="FlashInfer is not available.",
+)
+@pytest.mark.parametrize(
+    "global_dtype, checkpoint_dtype, expected",
+    [
+        ("auto", "fp8_e4m3", "fp8_e4m3"),
+        ("auto", "nvfp4", "nvfp4"),
+        ("fp8_e4m3", "nvfp4", "fp8_e4m3"),
+        ("auto", None, "auto"),
+    ],
+)
+def test_flashinfer_resolves_layerwise_cache_dtype(
+    dist_init, global_dtype, checkpoint_dtype, expected
+):
+    from vllm.model_executor.layers.attention import Attention
+    from vllm.model_executor.layers.quantization.fp8 import Fp8Config
+    from vllm.utils.torch_utils import set_default_torch_dtype
+    from vllm.v1.attention.backends.flashinfer import FlashInferMetadataBuilder
+
+    class LayerwiseKVConfig(Fp8Config):
+        def get_kv_cache_dtype(self, prefix):
+            return {"layer": checkpoint_dtype}.get(prefix)
+
+    config = create_vllm_config(
+        dtype="bfloat16", block_size=64, max_num_seqs=1, max_num_batched_tokens=128
+    )
+    config.cache_config.cache_dtype = global_dtype
+    config.attention_config.backend = AttentionBackendEnum.FLASHINFER
+    quant_config = Fp8Config() if checkpoint_dtype is None else LayerwiseKVConfig()
+    with set_current_vllm_config(config), set_default_torch_dtype(torch.bfloat16):
+        layer = Attention(
+            num_heads=16,
+            head_size=128,
+            scale=128**-0.5,
+            num_kv_heads=8,
+            cache_config=config.cache_config,
+            quant_config=quant_config,
+            prefix="layer",
+        )
+        spec = layer.get_kv_cache_spec(config)
+        assert spec is not None
+        builder = FlashInferMetadataBuilder(
+            spec, ["layer"], config, torch.device(DEVICE_TYPE)
+        )
+
+    assert layer.kv_cache_dtype == expected
+    assert builder.cache_dtype == expected
+    assert config.cache_config.cache_dtype == global_dtype
+
+
+@pytest.mark.skipif(
+    AttentionBackendEnum.FLASHINFER not in BACKENDS_TO_TEST,
+    reason="FlashInfer is not available.",
+)
+def test_flashinfer_rejects_unresolved_quantized_layer(dist_init):
+    from vllm.v1.attention.backends.flashinfer import FlashInferMetadataBuilder
+
+    config = create_vllm_config(dtype="bfloat16", block_size=64)
+    spec = FullAttentionSpec(
+        block_size=64,
+        num_kv_heads=8,
+        head_size=128,
+        dtype=torch.uint8,
+        kv_quant_mode=KVQuantMode.FP8_PER_TENSOR,
+    )
+    with (
+        set_current_vllm_config(config),
+        pytest.raises(ValueError, match="requires a resolved KV cache dtype"),
+    ):
+        FlashInferMetadataBuilder(spec, [], config, torch.device(DEVICE_TYPE))
+
+
+@pytest.mark.parametrize(
+    "cache_dtypes, expected",
+    [
+        ([], None),
+        ([None], None),
+        (["auto"], "auto"),
+        (["fp8_e4m3"], "fp8_e4m3"),
+        (["fp8", "fp8_e4m3"], "fp8_e4m3"),
+        (["nvfp4", "nvfp4_4over6"], "nvfp4"),
+    ],
+)
+def test_runtime_cache_dtype_resolution(monkeypatch, cache_dtypes, expected):
+    from vllm.v1.attention.backends import utils as backend_utils
+
+    layers = {
+        str(index): SimpleNamespace(kv_cache_dtype=dtype)
+        for index, dtype in enumerate(cache_dtypes)
+    }
+    monkeypatch.setattr(backend_utils, "get_layers_from_vllm_config", lambda *_: layers)
+    assert (
+        backend_utils.get_kv_cache_dtype_from_layers(SimpleNamespace(), list(layers))
+        == expected
+    )
+
+
+def test_runtime_cache_dtype_resolution_rejects_incompatible_layers(monkeypatch):
+    from vllm.v1.attention.backends import utils as backend_utils
+
+    layers = {
+        "layer0": SimpleNamespace(kv_cache_dtype="fp8_e4m3"),
+        "layer1": SimpleNamespace(kv_cache_dtype="fp8_e5m2"),
+    }
+    monkeypatch.setattr(backend_utils, "get_layers_from_vllm_config", lambda *_: layers)
+    with pytest.raises(ValueError, match="must share the same KV cache dtype"):
+        backend_utils.get_kv_cache_dtype_from_layers(SimpleNamespace(), list(layers))
+
+
 @pytest.mark.skipif(
     AttentionBackendEnum.FLASHINFER not in BACKENDS_TO_TEST,
     reason="FlashInfer is not available.",
diff --git a/tests/v1/attention/test_deepseek_v4_swa_visible.py b/tests/v1/attention/test_deepseek_v4_swa_visible.py
index a46495f79072..1d5653e81259 100644
--- a/tests/v1/attention/test_deepseek_v4_swa_visible.py
+++ b/tests/v1/attention/test_deepseek_v4_swa_visible.py
@@ -883,3 +883,106 @@ def common_metadata(seq_lens):
     assert md.decode_swa_width == builder.noncausal_index_width
     torch.testing.assert_close(md.decode_swa_lens, expected.decode_swa_lens)
     torch.testing.assert_close(md.decode_swa_indices, expected.decode_swa_indices)
+
+
+class _StubLinear(torch.nn.Module):
+    def __init__(self, *args, **kwargs):
+        super().__init__()
+        self.weight_block_size = None
+
+
+def _prefix_replay_attention_cls():
+    from vllm.models.deepseek_v41.attention import DeepseekV4Attention
+
+    class _Backend:
+        @staticmethod
+        def get_name() -> str:
+            return "STUB"
+
+    class Layer(DeepseekV4Attention):
+        backend_cls = _Backend
+
+        @classmethod
+        def get_padded_num_q_heads(cls, num_heads: int) -> int:
+            return num_heads
+
+        def forward_mqa(
+            self,
+            q: torch.Tensor,
+            kv: torch.Tensor,
+            positions: torch.Tensor,
+            output: torch.Tensor,
+        ) -> None:
+            return None
+
+        def _o_proj(self, attn_out: torch.Tensor, positions: torch.Tensor):
+            return attn_out
+
+    return Layer
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
+@pytest.mark.parametrize(
+    "is_rocm,cdna,flag,replays_window",
+    [
+        (False, 0, True, True),
+        (True, 4, True, True),
+        (True, 3, True, False),
+        (True, 4, False, False),
+    ],
+)
+def test_prefix_cache_replays_the_swa_window(
+    monkeypatch, is_rocm, cdna, flag, replays_window
+):
+    """Prefix caching stays on, and the window is replayed only on CDNA4.
+
+    tests/v1/core/test_prefix_replay.py already checks that a
+    bounded-replay SWA spec makes a prefix hit recompute the window.
+    This is the spec the attention layer actually publishes. Forcing
+    the window back into the prefix cache, on CDNA4 or off it, fails here.
+    The replay_start clamp used by the gather kernels is covered by
+    test_builder_replay_start_bounds_prefill_window_and_gather.
+    """
+    from vllm.config import set_current_vllm_config
+    from vllm.models.deepseek_v41 import attention as attn
+    from vllm.platforms import current_platform
+
+    monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1")
+    monkeypatch.setattr(attn, "MergedColumnParallelLinear", _StubLinear)
+    monkeypatch.setattr(attn, "ColumnParallelLinear", _StubLinear)
+    monkeypatch.setattr(attn, "RowParallelLinear", _StubLinear)
+    monkeypatch.setattr(attn, "RMSNorm", _StubLinear)
+    monkeypatch.setattr(attn, "build_deepseek_v4_rope", lambda *a, **k: None)
+    monkeypatch.setattr(attn, "get_tensor_model_parallel_world_size", lambda: 1)
+    monkeypatch.setattr(torch.cuda, "Event", lambda *a, **k: object())
+    monkeypatch.setattr(current_platform, "is_rocm", lambda: is_rocm)
+    if is_rocm:
+        import vllm.platforms.rocm as rocm_platform
+
+        monkeypatch.setattr(rocm_platform, "get_cdna_version", lambda: cdna)
+
+    window = 128
+    vllm_config = create_vllm_config(
+        max_model_len=256, max_num_seqs=4, max_num_batched_tokens=64
+    )
+    vllm_config.cache_config.swa_bounded_replay = flag
+    assert vllm_config.cache_config.enable_prefix_caching
+    vllm_config.model_config.hf_config = DeepseekV4Config(
+        num_hidden_layers=1,
+        hidden_size=128,
+        num_attention_heads=8,
+        q_lora_rank=32,
+        o_lora_rank=32,
+        head_dim=64,
+        qk_rope_head_dim=16,
+        o_groups=8,
+        sliding_window=window,
+        compress_ratios=[0],
+    )
+    layer_cls = _prefix_replay_attention_cls()
+    with set_current_vllm_config(vllm_config):
+        layer = layer_cls(vllm_config, prefix="model.layers.0.self_attn")
+    spec = layer.swa_cache_layer.get_kv_cache_spec(vllm_config)
+    assert spec.bounded_replay is replays_window
+    assert spec.prefix_cacheable is not replays_window
+    assert spec.prefix_replay_tokens == (window if replays_window else 0)
diff --git a/tests/v1/attention/test_flashinfer_mla_dcp.py b/tests/v1/attention/test_flashinfer_mla_dcp.py
index a72dbc21d145..f14238fb38ec 100644
--- a/tests/v1/attention/test_flashinfer_mla_dcp.py
+++ b/tests/v1/attention/test_flashinfer_mla_dcp.py
@@ -87,7 +87,8 @@ def test_flashinfer_mla_forward_uses_gathered_head_count(monkeypatch):
 
 @requires_flashinfer_mla
 @pytest.mark.parametrize("causal", [True, False], ids=["causal", "noncausal"])
-def test_flashinfer_mla_forward_uses_native_dcp_api(monkeypatch, causal):
+@pytest.mark.parametrize("query_lens", [(3, 3), (3, 1)], ids=["uniform", "ragged"])
+def test_flashinfer_mla_forward_uses_native_dcp_api(monkeypatch, causal, query_lens):
     import vllm.v1.attention.backends.mla.flashinfer_mla as flashinfer_mla
 
     impl = MagicMock()
@@ -106,8 +107,11 @@ def test_flashinfer_mla_forward_uses_native_dcp_api(monkeypatch, causal):
         flashinfer_mla.FlashInferMLAImpl._flattened_decode_metadata, impl
     )
 
-    num_reqs, query_len = 2, 3
-    num_tokens = num_reqs * query_len
+    num_reqs = len(query_lens)
+    num_tokens = sum(query_lens)
+    query_len = max(query_lens)
+    split_rows = not causal or len(set(query_lens)) > 1
+    lens = torch.tensor(query_lens)
     block_table = torch.tensor([[10], [20]], dtype=torch.int32)
     seq_lens = torch.tensor([4, 5], dtype=torch.int32)
     global_causal_seq_lens = torch.tensor([31, 35], dtype=torch.int32)
@@ -120,15 +124,15 @@ def test_flashinfer_mla_forward_uses_native_dcp_api(monkeypatch, causal):
     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_start_loc = torch.tensor(
+        [0, *lens.cumsum(0).tolist()], 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)
-    kernel_batch = num_reqs if causal else num_tokens
-    kernel_query_len = query_len if causal else 1
+    kernel_batch = num_tokens if split_rows else num_reqs
+    kernel_query_len = 1 if split_rows else query_len
     kernel = MagicMock(
         return_value=(
             torch.ones(
@@ -155,15 +159,17 @@ def test_flashinfer_mla_forward_uses_native_dcp_api(monkeypatch, causal):
 
     call = kernel.call_args.kwargs
     assert call["query"].shape == (kernel_batch, kernel_query_len, 24, 576)
-    expected_block_table = (
-        block_table if causal else block_table.repeat_interleave(query_len, dim=0)
-    )
-    expected_seq_lens = seq_lens if causal else seq_lens.repeat_interleave(query_len)
-    expected_global_seq_lens = (
-        global_causal_seq_lens
-        if causal
-        else global_causal_seq_lens.repeat_interleave(query_len)
-    )
+    expected_block_table = block_table
+    expected_seq_lens = seq_lens
+    expected_global_seq_lens = global_causal_seq_lens
+    if split_rows:
+        expected_block_table = block_table.repeat_interleave(lens, dim=0)
+        expected_seq_lens = seq_lens.repeat_interleave(lens)
+        expected_global_seq_lens = global_causal_seq_lens.repeat_interleave(lens)
+    if causal and split_rows:
+        # Each split row sees KV only up to its own position in its request.
+        rows_after = torch.cat([torch.arange(n - 1, -1, -1) for n in query_lens])
+        expected_global_seq_lens = expected_global_seq_lens - rows_after.int()
     torch.testing.assert_close(call["block_tables"], expected_block_table)
     torch.testing.assert_close(call["seq_lens"], expected_seq_lens)
     assert call["backend"] == "cute-dsl"
diff --git a/tests/v1/attention/test_flashinfer_sparse_mla_workspace.py b/tests/v1/attention/test_flashinfer_sparse_mla_workspace.py
index 5b63c2c609f9..ea1b1a232787 100644
--- a/tests/v1/attention/test_flashinfer_sparse_mla_workspace.py
+++ b/tests/v1/attention/test_flashinfer_sparse_mla_workspace.py
@@ -62,6 +62,23 @@ def test_non_dcp_size_is_unchanged():
     assert computed == _DEFAULT_WORKSPACE_BUFFER_SIZE
 
 
+def test_dcp_size_is_capped_at_one_launch():
+    # Launches cover at most 65535 tokens, so larger token budgets reuse the slab.
+    at_limit = compute_trtllm_sparse_mla_workspace_bytes(
+        base_workspace_bytes=0,
+        dcp_world_size=2,
+        num_heads_per_rank=8,
+        max_num_batched_tokens=65535,
+    )
+    above_limit = compute_trtllm_sparse_mla_workspace_bytes(
+        base_workspace_bytes=0,
+        dcp_world_size=2,
+        num_heads_per_rank=8,
+        max_num_batched_tokens=131072,
+    )
+    assert above_limit == at_limit
+
+
 def test_combined_pcp_dcp_rejects_flashinfer_sparse():
     config = SimpleNamespace(
         parallel_config=SimpleNamespace(
diff --git a/tests/v1/attention/test_flashinfer_ts.py b/tests/v1/attention/test_flashinfer_ts.py
new file mode 100644
index 000000000000..f77186d4aa03
--- /dev/null
+++ b/tests/v1/attention/test_flashinfer_ts.py
@@ -0,0 +1,229 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""FLASHINFER_TS: the non-causal draft-block decode path.
+
+Covers the two things that are easy to get wrong and silently produce plausible
+output: the K/V plane layout the kernel validates, and the page-aligned window
+restriction that stands in for a window the kernel cannot take with a dense
+mask.
+"""
+
+import pytest
+import torch
+
+from vllm.platforms import current_platform
+from vllm.v1.attention.backend import AttentionType
+from vllm.v1.attention.backends.flashinfer_ts import (
+    FlashInferTSMetadataBuilder,
+)
+
+HEAD_DIM = 128
+PAGE = 128
+WINDOW = 4096
+
+
+def _sm100_family() -> bool:
+    cap = current_platform.get_device_capability()
+    return cap is not None and (cap.major, cap.minor) in ((10, 0), (10, 3), (10, 7))
+
+
+requires_kernel = pytest.mark.skipif(
+    not current_platform.is_cuda() or not _sm100_family(),
+    reason="prims_ts decode runs on SM100a/SM103a/SM107a only",
+)
+
+
+@pytest.mark.parametrize("num_kv_heads", [1, 4])
+def test_reads_vllms_existing_interleaved_page(num_kv_heads):
+    """No layout change: K/V are the halves of vLLM's [pages, Hkv, page, 2*D]
+    page, whose token stride is 2*D. The kernel takes its TMA strides from the
+    tensor, so this must stay a pure view with head_dim contiguous.
+    """
+    page = torch.zeros(8, num_kv_heads, PAGE, 2 * HEAD_DIM, dtype=torch.uint8)
+    k, v = page[..., :HEAD_DIM], page[..., HEAD_DIM:]
+
+    assert k.shape == (8, num_kv_heads, PAGE, HEAD_DIM)
+    assert k.stride(-1) == 1 and v.stride(-1) == 1  # the one hard requirement
+    assert k.stride(2) == 2 * HEAD_DIM  # interleaved, not compact
+    assert k.data_ptr() == page.data_ptr()  # views, never a repack
+    assert v.data_ptr() == page.data_ptr() + HEAD_DIM
+
+
+@pytest.mark.parametrize("query_len", [1, 3, 7])
+def test_window_restriction_is_a_superset_of_the_true_window(query_len):
+    """The restricted page table must never drop a token the symmetric window
+    can reach, and must overshoot by less than one page.
+    """
+    builder = FlashInferTSMetadataBuilder.__new__(FlashInferTSMetadataBuilder)
+    builder.sliding_window = WINDOW
+    builder.page_size = PAGE
+    builder.query_len = query_len
+    # The builder writes metadata into persistent buffers so a replayed CUDA
+    # graph keeps reading the same addresses; the window math is exercised
+    # through them, not around them.
+    pages = builder._window_pages(query_len)
+    builder._buf_seq_lens = torch.zeros(6, dtype=torch.int32)
+    builder._buf_block_table = torch.zeros((6, pages), dtype=torch.int32)
+
+    seq_lens = torch.tensor(
+        [1, PAGE, WINDOW - 1, WINDOW + query_len, 8192, 120_000], dtype=torch.int32
+    )
+    block_table = torch.arange(
+        seq_lens.numel() * (int(seq_lens.max()) // PAGE + 1), dtype=torch.int32
+    ).view(seq_lens.numel(), -1)
+
+    new_lens, new_bt = FlashInferTSMetadataBuilder._restrict_to_window(
+        builder, seq_lens, block_table, query_len
+    )
+
+    span = WINDOW + query_len - 1
+    reachable = torch.minimum(seq_lens, torch.full_like(seq_lens, span))
+    assert torch.all(new_lens >= reachable), "window restriction dropped live tokens"
+    assert int((new_lens - reachable).max()) < PAGE, "overshoot exceeds one page"
+    # The kept pages must be the tail of each row, so the last token of the
+    # sequence is still addressable.
+    last_page = (seq_lens.long() - 1) // PAGE
+    assert torch.all(new_bt.long().max(dim=1).values >= last_page)
+
+
+def test_impl_accepts_the_arguments_the_attention_layer_passes():
+    """Attention() calls impl_cls positionally and forwards extra_impl_args
+    (``sinks`` among them); a signature that only matches AttentionImpl's
+    abstract __init__ raises TypeError at engine init, not at import.
+    """
+    from vllm.v1.attention.backends.flashinfer_ts import FlashInferTSImpl
+
+    impl = FlashInferTSImpl(
+        16,
+        HEAD_DIM,
+        HEAD_DIM**-0.5,
+        1,
+        None,
+        WINDOW,
+        "fp8",
+        None,
+        AttentionType.DECODER,
+        None,
+        sinks=None,
+    )
+    # FP8 KV reads an FP8 query; the layer's own quant is skipped via
+    # query_pre_quantized because the fused QK-norm/RoPE kernel emits it.
+    assert impl.supports_quant_query_input
+
+    with pytest.raises(NotImplementedError):
+        FlashInferTSImpl(
+            16,
+            HEAD_DIM,
+            HEAD_DIM**-0.5,
+            1,
+            None,
+            WINDOW,
+            "fp8",
+            None,
+            AttentionType.DECODER,
+            None,
+            sinks=torch.zeros(16),
+        )
+
+
+@requires_kernel
+def test_dense_plan_matches_flash_attention_non_causal():
+    """A dense prims_ts plan is the non-causal draft block FA4 produces."""
+    from vllm.vllm_flash_attn.cute.interface import _flash_attn_fwd
+
+    pytest.importorskip("flashinfer.attention.prims_ts")
+    from flashinfer.attention.prims_ts import BatchDecodePagedTSWrapper
+
+    torch.manual_seed(0)
+    b, sq, hq, hkv, kv = 2, 3, 16, 1, PAGE
+    dev = torch.device("cuda")
+    scale = HEAD_DIM**-0.5
+    pages = b * (kv // PAGE) + 1
+
+    page = torch.randn(pages, 2 * hkv, PAGE, HEAD_DIM, dtype=torch.bfloat16, device=dev)
+    k_plane, v_plane = page[:, :hkv].contiguous(), page[:, hkv:].contiguous()
+    interleaved = torch.cat([k_plane, v_plane], dim=-1).contiguous()
+    k_fa, v_fa = interleaved.transpose(1, 2).split(HEAD_DIM, dim=-1)
+
+    q = torch.randn(b * sq, hq, HEAD_DIM, dtype=torch.bfloat16, device=dev)
+    bt = torch.arange(1, pages, dtype=torch.int32, device=dev).view(b, -1).contiguous()
+    lens = torch.full((b,), kv, dtype=torch.int32, device=dev)
+    cu_q = torch.arange(0, (b + 1) * sq, sq, dtype=torch.int32, device=dev)
+
+    ref = torch.empty(b * sq, hq, HEAD_DIM, dtype=torch.bfloat16, device=dev)
+    _flash_attn_fwd(
+        q,
+        k_fa,
+        v_fa,
+        cu_seqlens_q=cu_q,
+        seqused_k=lens,
+        max_seqlen_q=sq,
+        max_seqlen_k=kv,
+        page_table=bt,
+        softmax_scale=scale,
+        causal=False,
+        window_size_left=WINDOW - 1,
+        window_size_right=WINDOW - 1,
+        num_splits=0,
+        out=ref,
+    )
+
+    wrapper = BatchDecodePagedTSWrapper(kv_layout="HND")
+    wrapper.plan(
+        device=dev,
+        batch_size=b,
+        num_qo_heads=hq,
+        num_kv_heads=hkv,
+        head_dim=HEAD_DIM,
+        page_size=PAGE,
+        max_kv_len=kv,
+        max_seq_len_q=sq,
+        packed_query=False,
+        q_data_type=torch.bfloat16,
+        kv_data_type=torch.bfloat16,
+        o_data_type=torch.bfloat16,
+        mask_type="dense",
+        window_left=-1,
+    )
+    out = torch.empty(b, sq, hq, HEAD_DIM, dtype=torch.bfloat16, device=dev)
+    wrapper.run(
+        q=q.view(b, sq, hq, HEAD_DIM),
+        paged_kv_cache=(k_plane, v_plane),
+        seq_lens=lens,
+        block_tables=bt,
+        bmm1_scale=scale,
+        bmm2_scale=1.0,
+        out=out,
+    )
+
+    # Within one bf16 ulp: the kernels round differently on some SM100 parts.
+    torch.testing.assert_close(
+        out.reshape(b * sq, hq, HEAD_DIM), ref, rtol=1e-2, atol=2e-3
+    )
+
+
+def test_metadata_tensors_keep_their_addresses_across_builds():
+    """A captured CUDA graph replays the ADDRESS of seq_lens and the page table.
+
+    Returning freshly allocated tensors per step leaves a replayed graph reading
+    memory the allocator has since handed to someone else, which shows up as an
+    engine that pins every GPU and never retires a step.
+    """
+    builder = FlashInferTSMetadataBuilder.__new__(FlashInferTSMetadataBuilder)
+    builder.sliding_window = WINDOW
+    builder.page_size = PAGE
+    builder.query_len = 4
+    pages = builder._window_pages(4)
+    builder._buf_seq_lens = torch.zeros(4, dtype=torch.int32)
+    builder._buf_block_table = torch.zeros((4, pages), dtype=torch.int32)
+
+    block_table = torch.arange(4 * 1024, dtype=torch.int32).view(4, -1)
+    ptrs = set()
+    for step in range(3):
+        seq_lens = torch.tensor([512, 8192, 40_000 + step, 120_000], dtype=torch.int32)
+        lens, table = FlashInferTSMetadataBuilder._restrict_to_window(
+            builder, seq_lens, block_table, 4
+        )
+        ptrs.add((lens.data_ptr(), table.data_ptr(), tuple(table.shape)))
+
+    assert len(ptrs) == 1, f"metadata moved between builds: {ptrs}"
diff --git a/tests/v1/attention/test_gdn_metadata_builder.py b/tests/v1/attention/test_gdn_metadata_builder.py
index 937bf2f30d7d..e078b366c5a2 100644
--- a/tests/v1/attention/test_gdn_metadata_builder.py
+++ b/tests/v1/attention/test_gdn_metadata_builder.py
@@ -19,6 +19,7 @@
 from vllm.config.compilation import CUDAGraphMode
 from vllm.v1.attention.backend import CommonAttentionMetadata
 from vllm.v1.attention.backends.gdn_attn import (
+    GDNAttentionBackend,
     GDNAttentionMetadata,
     GDNAttentionMetadataBuilder,
     plan_flashinfer_checkpoints,
@@ -150,6 +151,7 @@ class GDNBuildTestCase:
 def _create_gdn_builder(
     num_speculative_tokens: int = 0,
     full_cuda_graph: bool = False,
+    adaptive_verification: bool = False,
 ) -> GDNAttentionMetadataBuilder:
     """Create a GDNAttentionMetadataBuilder with minimal config."""
     vllm_config = create_vllm_config(
@@ -163,6 +165,7 @@ def _create_gdn_builder(
         vllm_config.speculative_config = SpeculativeConfig(
             method="ngram",
             num_speculative_tokens=num_speculative_tokens,
+            enable_adaptive_verification=adaptive_verification,
         )
     mamba_spec = MambaSpec(
         block_size=BLOCK_SIZE,
@@ -302,6 +305,64 @@ def test_full_cudagraph_spec_metadata_uses_request_count():
     assert meta.num_accepted_tokens.shape == (batch.batch_size,)
 
 
+@pytest.mark.parametrize("num_speculative_tokens", [3, 15])
+def test_adaptive_verification_uses_device_query_boundaries(
+    num_speculative_tokens: int,
+):
+    """GDN keeps the CPU batch plan while executing the GPU's ragged plan."""
+    builder = _create_gdn_builder(
+        num_speculative_tokens=num_speculative_tokens,
+        adaptive_verification=True,
+    )
+    batch = BatchSpec(seq_lens=[80, 96, 112], query_lens=[3, 3, 3])
+    common = create_common_attn_metadata(batch, BLOCK_SIZE, DEVICE)
+    device_query_start_loc = torch.tensor([0, 1, 5, 9], dtype=torch.int32)
+    common = common.replace(query_start_loc=device_query_start_loc)
+
+    meta = builder.build(
+        common_prefix_len=0,
+        common_attn_metadata=common,
+        num_accepted_tokens=torch.ones(3, dtype=torch.int32),
+        num_decode_draft_tokens_cpu=torch.tensor([2, 2, 2], dtype=torch.int32),
+    )
+
+    assert GDNAttentionBackend.supports_device_cpu_query_lens_mismatch()
+    assert (
+        builder.get_varlen_cudagraph_max_query_len(
+            builder.vllm_config, builder.kv_cache_spec
+        )
+        == num_speculative_tokens + 1
+    )
+    torch.testing.assert_close(meta.spec_query_start_loc, device_query_start_loc)
+    assert meta.num_spec_decode_tokens == 9
+    assert meta.num_spec_decodes == 3
+    assert meta.uniform_spec_sequence_length is None
+
+
+def test_adaptive_cudagraph_capture_uses_spec_decode_path():
+    """A varlen graph captured as one-token rows must replay multi-token rows."""
+    builder = _create_gdn_builder(
+        num_speculative_tokens=3,
+        full_cuda_graph=True,
+        adaptive_verification=True,
+    )
+    # Varlen graph capture maximizes the request count, so a four-token graph
+    # is captured as four one-token dummy rows. At replay it may instead hold
+    # one four-token speculative row. Capture must therefore freeze the
+    # speculative path, whose padded rows can also process one-token decodes.
+    batch = BatchSpec(seq_lens=[1, 1, 1, 1], query_lens=[1, 1, 1, 1])
+    common = create_common_attn_metadata(batch, BLOCK_SIZE, DEVICE)
+
+    meta = builder.build_for_cudagraph_capture(common)
+
+    assert meta.num_spec_decodes == 4
+    assert meta.num_decodes == 0
+    assert meta.spec_query_start_loc is not None
+    assert meta.spec_sequence_masks_cpu is not None
+    assert meta.spec_sequence_masks_cpu.device.type == "cpu"
+    torch.testing.assert_close(meta.spec_query_start_loc, common.query_start_loc)
+
+
 def _build_non_spec(
     batch: BatchSpec,
     is_prefilling: list[bool] | None,
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 f88c5bc0fb50..709e66072888 100644
--- a/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py
+++ b/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py
@@ -7,6 +7,7 @@
 import torch
 
 from tests.v1.attention.utils import create_vllm_config
+from vllm.config import SpeculativeConfig
 from vllm.model_executor.layers.attention.sparse_mla_attention import (
     SparseMLACommonMetadataBuilder,
 )
@@ -41,6 +42,69 @@
 from vllm.v1.worker.utils import select_common_block_size
 
 
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
+def test_adaptive_indexer_uses_configured_max_decode_len():
+    """GLM adaptive verification may replay wider rows than capture dummy rows."""
+    device = torch.device("cuda")
+    num_speculative_tokens = 5
+    configured_max_decode_len = num_speculative_tokens + 1
+
+    kv_cache_spec = MLAAttentionSpec(
+        block_size=64,
+        num_kv_heads=1,
+        head_size=128,
+        dtype=torch.bfloat16,
+    )
+    vllm_config = create_vllm_config(
+        model_name="hmellor/tiny-random-LlamaForCausalLM",
+        max_model_len=1024,
+    )
+    vllm_config.speculative_config = SpeculativeConfig(
+        method="ngram",
+        num_speculative_tokens=num_speculative_tokens,
+        enable_adaptive_verification=True,
+    )
+    max_num_blocks = kv_cache_spec.max_num_blocks_per_req(vllm_config, 1024)
+    block_table_width = get_block_table_width(max_num_blocks, kv_cache_spec.block_size)
+    builder = DeepseekV32IndexerMetadataBuilder(
+        kv_cache_spec=kv_cache_spec,
+        layer_names=["dummy"],
+        vllm_config=vllm_config,
+        device=device,
+        block_table_width=block_table_width,
+    )
+
+    # The dummy capture rows are at most two tokens wide, but adaptive
+    # verification may assign all k+1 tokens to one request during replay.
+    query_lens = torch.tensor([2, 2, 2, 2, 1, 1, 1, 1], dtype=torch.int32)
+    query_start_loc_cpu = torch.cat(
+        (torch.zeros(1, dtype=torch.int32), query_lens.cumsum(0))
+    )
+    query_start_loc = query_start_loc_cpu.to(device)
+    seq_lens = (query_lens + 100).to(device)
+    common = CommonAttentionMetadata(
+        query_start_loc=query_start_loc,
+        query_start_loc_cpu=query_start_loc_cpu,
+        seq_lens=seq_lens,
+        seq_lens_cpu_upper_bound=seq_lens.cpu(),
+        num_reqs=8,
+        num_actual_tokens=12,
+        max_query_len=2,
+        max_seq_len=102,
+        block_table_tensor=torch.ones(
+            (8, block_table_width), dtype=torch.int32, device=device
+        ),
+        slot_mapping=torch.arange(12, dtype=torch.int64, device=device),
+        causal=True,
+    )
+
+    metadata = builder.build(0, common)
+
+    assert metadata.decode is not None
+    assert metadata.decode.write_max_decode_len == configured_max_decode_len
+    assert not metadata.decode.decode_is_uniform
+
+
 @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])
@@ -449,6 +513,7 @@ def test_compressed_slot_mapping_dcp_uses_rank_local_state_positions():
 
 def test_index_conversion_warmup_uses_physical_block_stride():
     config = SimpleNamespace(
+        attention_config=SimpleNamespace(hisparse_config=None),
         cache_config=SimpleNamespace(block_size=64),
         model_config=SimpleNamespace(
             max_model_len=1024,
@@ -667,3 +732,29 @@ def test_indexer_prefill_budget_matches_compressed_workspace(compress_ratio):
     assert len(chunks) > 1
     for req_slice, _ in chunks:
         assert int(compressed_seq_lens[req_slice].sum()) <= workspace_rows
+
+
+def test_indexer_expanded_block_table_sized_by_decode_tokens():
+    """Allocated after the memory profile, so it must not scale with the token
+    budget: only decode tokens (one per request without drafts) are expanded."""
+    kv_cache_spec = MLAAttentionSpec(
+        block_size=256,
+        num_kv_heads=1,
+        head_size=128,
+        dtype=torch.bfloat16,
+        tokens_per_state=4,
+    )
+    vllm_config = create_vllm_config(
+        max_model_len=1024, max_num_seqs=8, max_num_batched_tokens=8192
+    )
+    max_num_blocks = kv_cache_spec.max_num_blocks_per_req(vllm_config, 1024)
+    block_table_width = get_block_table_width(max_num_blocks, kv_cache_spec.block_size)
+    builder = DeepseekV32IndexerMetadataBuilder(
+        kv_cache_spec=kv_cache_spec,
+        layer_names=["dummy"],
+        vllm_config=vllm_config,
+        device=torch.device("cpu"),
+        block_table_width=block_table_width,
+    )
+
+    assert builder.expanded_block_table_buffer.shape == (8, block_table_width)
diff --git a/tests/v1/attention/test_indexer_tp_row_shard.py b/tests/v1/attention/test_indexer_tp_row_shard.py
index 4bbbf93eb905..ae683df8a1a3 100644
--- a/tests/v1/attention/test_indexer_tp_row_shard.py
+++ b/tests/v1/attention/test_indexer_tp_row_shard.py
@@ -210,6 +210,10 @@ def run_dummy(tp_world: int) -> None:
                         decode_context_parallel_size=1,
                         prefill_context_parallel_size=1,
                     ),
+                    scheduler_config=SimpleNamespace(
+                        max_num_seqs=4, max_num_batched_tokens=512
+                    ),
+                    num_speculative_tokens=0,
                 ),
             )
             m.setattr(sparse_indexer.envs, "VLLM_SPARSE_INDEXER_MAX_LOGITS_MB", 1)
diff --git a/tests/v1/attention/test_mla_backends.py b/tests/v1/attention/test_mla_backends.py
index 6883dbc4b363..cc722c2ac33e 100644
--- a/tests/v1/attention/test_mla_backends.py
+++ b/tests/v1/attention/test_mla_backends.py
@@ -10,7 +10,7 @@
 
 import sys
 from types import MethodType, SimpleNamespace
-from typing import Literal
+from typing import Any, Literal
 
 import pytest
 import torch
@@ -36,7 +36,7 @@
 from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
 from vllm.platforms import current_platform
 from vllm.utils.math_utils import cdiv
-from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
+from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE, set_random_seed
 from vllm.v1.attention.backend import CommonAttentionMetadata, MLAAttentionImpl
 from vllm.v1.attention.backends.fa_utils import flash_attn_supports_mla
 from vllm.v1.attention.backends.mla import flashmla as flashmla_module
@@ -219,13 +219,39 @@ def test_masked_mha_routing_is_dimension_specific():
         ("fp8_ds_mla", KVQuantMode.FP8_PER_TENSOR),
     ],
 )
+@pytest.mark.parametrize(
+    ("sm120", "kv_lora_rank", "rope_dim", "row_bytes"),
+    [
+        (False, 512, 64, 656),
+        (False, 512, 0, 656),
+        (True, 512, 64, 656),
+        (True, 512, 0, 528),
+        (True, 256, 0, 656),
+    ],
+)
 def test_mla_kv_cache_spec_uses_layer_cache_dtype(
-    cache_dtype: str, expected_quant_mode: KVQuantMode
+    cache_dtype: str,
+    expected_quant_mode: KVQuantMode,
+    sm120: bool,
+    kv_lora_rank: int,
+    rope_dim: int,
+    row_bytes: int,
 ):
+    from vllm.v1.attention.backends.mla.flashinfer_mla_sparse_sm120 import (
+        FlashInferMLASparseSM120Impl,
+    )
+    from vllm.v1.attention.backends.mla.flashmla_sparse import FlashMLASparseImpl
+
+    impl: Any = object.__new__(
+        FlashInferMLASparseSM120Impl if sm120 else FlashMLASparseImpl
+    )
+    impl.kv_lora_rank = kv_lora_rank
+    impl.qk_rope_head_dim = rope_dim
     layer = SimpleNamespace(
         attn_backend=flashmla_module.FlashMLABackend,
+        impl=impl,
         kv_cache_dtype=cache_dtype,
-        head_size=576,
+        head_size=kv_lora_rank + rope_dim,
         topk_buffer_width=None,
         indexer=None,
         non_causal_multi_token_decode=False,
@@ -242,7 +268,8 @@ def test_mla_kv_cache_spec_uses_layer_cache_dtype(
     assert spec.cache_dtype_str == cache_dtype
     assert spec.kv_quant_mode == expected_quant_mode
     if cache_dtype == "fp8_ds_mla":
-        assert spec.page_size_bytes == 64 * 656
+        assert spec.state_content_bytes == row_bytes
+        assert spec.page_size_bytes == 64 * row_bytes
 
 
 @pytest.mark.cpu_test
@@ -328,6 +355,7 @@ def test_mla_post_load_preserves_runtime_weight_addresses(monkeypatch):
     layer.dcp_q_replicate = False
     layer.quant_config = None
     layer.layer_name = "test"
+    layer.act_dtype = torch.float32
     layer.impl = SimpleNamespace(process_weights_after_loading=lambda act_dtype: None)
 
     monkeypatch.setattr(
@@ -335,7 +363,7 @@ def test_mla_post_load_preserves_runtime_weight_addresses(monkeypatch):
     )
 
     with torch.no_grad():
-        layer.process_weights_after_loading(torch.float32)
+        layer.process_weights_after_loading()
         assert isinstance(layer.W_UV, torch.nn.Parameter)
         assert isinstance(layer.W_UK_T, torch.nn.Parameter)
         w_uv_ptr = layer.W_UV.data_ptr()
@@ -344,7 +372,7 @@ def test_mla_post_load_preserves_runtime_weight_addresses(monkeypatch):
         old_w_uk_t = layer.W_UK_T.clone()
 
         layer.kv_b_proj.weight.add_(100)
-        layer.process_weights_after_loading(torch.float32)
+        layer.process_weights_after_loading()
 
     assert layer.W_UV.data_ptr() == w_uv_ptr
     assert layer.W_UK_T.data_ptr() == w_uk_t_ptr
@@ -1683,6 +1711,7 @@ def _run_backend_correctness(
     multiple GPUs. This tests that backends work correctly with different
     head counts.
     """
+    set_random_seed(42)
     # Filter backends to those that support the requested kv_cache_dtype
     backends_to_test = [
         b
@@ -2058,7 +2087,7 @@ def _run_backend_correctness(
     # 4. Run vLLM backends and compare
     rtol = 1e-2
     atol = {
-        "auto": 1e-2,
+        "auto": 2e-2,
         "fp8": 1.5e-1,
         "fp8_e4m3": 1.5e-1,
     }[kv_cache_dtype]
diff --git a/tests/v1/attention/test_mla_prefill_flash_attn.py b/tests/v1/attention/test_mla_prefill_flash_attn.py
new file mode 100644
index 000000000000..62d271a1dea1
--- /dev/null
+++ b/tests/v1/attention/test_mla_prefill_flash_attn.py
@@ -0,0 +1,40 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Output layout of the FLASH_ATTN MLA prefill backend."""
+
+from unittest.mock import patch
+
+import torch
+
+from vllm.v1.attention.backends.mla.prefill.flash_attn import (
+    FlashAttnPrefillBackend,
+)
+
+_FA_MODULE = "vllm.v1.attention.backends.mla.prefill.flash_attn"
+
+
+def test_upstream_flash_attn_head_major_output_is_made_token_major():
+    """Upstream flash_attn on ROCm may return a head-major output, while
+    merge_attn_states derives each partial's token stride from its head
+    stride. The backend must hand back a token-major partial."""
+    num_tokens, num_heads, head_size = 5, 3, 8
+    expected = torch.randn(num_tokens, num_heads, head_size)
+    lse = torch.randn(num_heads, num_tokens)
+
+    def head_major_kernel(**kwargs):
+        assert kwargs["return_attn_probs"]
+        return expected.transpose(0, 1).contiguous().transpose(0, 1), lse
+
+    backend = object.__new__(FlashAttnPrefillBackend)
+    backend.requires_v_padding = False
+    backend._is_vllm_fa = False
+    backend.flash_attn_varlen_func = None
+    q = torch.empty(num_tokens, num_heads, head_size)
+    with patch(f"{_FA_MODULE}._FA4_MLA_PREFILL_KERNEL", head_major_kernel):
+        attn_out, attn_lse = backend._flash_attn_varlen_diff_headdims(
+            q=q, k=q, v=q, return_softmax_lse=True
+        )
+
+    assert attn_out.is_contiguous()
+    torch.testing.assert_close(attn_out, expected)
+    assert attn_lse is lse
diff --git a/tests/v1/attention/test_rocm_aiter_mla_dcp_cprr_numerics.py b/tests/v1/attention/test_rocm_aiter_mla_dcp_cprr_numerics.py
index 21626e9a675e..4c48508c7159 100644
--- a/tests/v1/attention/test_rocm_aiter_mla_dcp_cprr_numerics.py
+++ b/tests/v1/attention/test_rocm_aiter_mla_dcp_cprr_numerics.py
@@ -547,3 +547,174 @@ def test_cprr_kernel_is_selected_for_the_k3_shape():
     )
     assert route is m._DCPDecodeRoute.CPRR
     assert m._asm_dcp_verify_heads(96) in m._NATIVE_CPRR_HEADS
+
+
+VARLEN_QO = (1, 5, 3, 2, 4, 1, 5, 2, 0, 0)
+VARLEN_CTX = (256, 3, 100, 0, 37, 1, 251, 2, 0, 0)
+PAD_TOKENS = 6
+
+
+def _varlen_reference(q_f32, kv_f32, seq_lens, qo_lens, positions_per_req):
+    heads = q_f32.shape[1]
+    o = torch.zeros(sum(qo_lens), heads, KV_LORA_RANK, device=q_f32.device)
+    lse = torch.full((sum(qo_lens), heads), float("-inf"), device=q_f32.device)
+    token = 0
+    for req, (seq_len, qo_len) in enumerate(zip(seq_lens, qo_lens)):
+        positions = positions_per_req[req]
+        kv_sel = kv_f32[req].index_select(0, positions)
+        for t in range(qo_len):
+            valid = positions <= seq_len - qo_len + t
+            if bool(valid.any()):
+                k = kv_sel[valid]
+                s = (q_f32[token] @ k.T) * SCALE
+                lse[token] = torch.logsumexp(s, dim=-1)
+                o[token] = torch.softmax(s, dim=-1) @ k[:, :KV_LORA_RANK]
+            token += 1
+    return o, lse
+
+
+@requires_cprr
+@pytest.mark.parametrize(
+    ("asm", "expected"),
+    [
+        (True, rocm_aiter_mla._DCPDecodeRoute.CPRR),
+        (False, rocm_aiter_mla._DCPDecodeRoute.SEGMENTED),
+    ],
+)
+@torch.inference_mode()
+def test_adaptive_varlen_shards_match_reference(monkeypatch, asm, expected):
+    import aiter
+
+    from vllm.config import set_current_vllm_config
+    from vllm.v1.attention.backend import CommonAttentionMetadata
+    from vllm.v1.attention.backends.mla.rocm_aiter_mla import AiterMLAMetadataBuilder
+    from vllm.v1.attention.ops.dcp import (
+        CPTritonContext,
+        correct_attn_out,
+        mask_dcp_empty_shards_,
+    )
+    from vllm.v1.kv_cache_interface import MLAAttentionSpec
+
+    if not asm:
+        monkeypatch.setattr(
+            rocm_aiter_mla, "_asm_dcp_verify_configured", lambda *a, **k: False
+        )
+    heads_per_rank = 12
+    device = torch.device("cuda", 0)
+    torch.accelerator.set_device_index(device.index)
+    torch.manual_seed(0)
+    _ensure_workspace_manager(device)
+    _install_fake_groups(DCP, 0)
+    vllm_config = _make_vllm_config(heads_per_rank)
+    vllm_config.speculative_config.enable_adaptive_verification = True
+
+    nr = len(VARLEN_QO)
+    qo_lens = list(VARLEN_QO)
+    seq_lens = [c + q for c, q in zip(VARLEN_CTX, VARLEN_QO)]
+    num_tokens = sum(qo_lens)
+    max_len = max(seq_lens)
+    gathered_heads = heads_per_rank * DCP
+    real_reqs = sum(q > 0 for q in qo_lens)
+    base, extra = divmod(num_tokens, real_reqs)
+    host_qo = [base + (i < extra) for i in range(real_reqs)] + [0] * (nr - real_reqs)
+
+    def qsl(lens):
+        return torch.tensor([0] + lens, dtype=torch.int32).cumsum(0, dtype=torch.int32)
+
+    with set_current_vllm_config(vllm_config):
+        layer_name = "model.layers.0.self_attn.attn"
+        vllm_config.compilation_config.static_forward_context[layer_name] = (
+            _make_ctx_layer(vllm_config, monkeypatch)
+        )
+        spec = MLAAttentionSpec(
+            block_size=1,
+            num_kv_heads=1,
+            head_size=HEAD_SIZE,
+            dtype=vllm_config.model_config.dtype,
+            non_causal_multi_token_decode=False,
+        )
+        kv_flat = torch.randn(
+            nr * max_len, 1, HEAD_SIZE, dtype=torch.float32, device=device
+        ).to(aiter.dtypes.fp8)
+        q_flat = (
+            torch.randn(
+                num_tokens + PAD_TOKENS,
+                gathered_heads,
+                HEAD_SIZE,
+                dtype=torch.float32,
+                device=device,
+            )
+            .mul_(0.5)
+            .to(aiter.dtypes.fp8)
+        )
+        kv_f32 = kv_flat.float().view(nr, max_len, HEAD_SIZE)
+        q_f32 = q_flat.float()
+        ref_o, _ = _varlen_reference(
+            q_f32,
+            kv_f32,
+            seq_lens,
+            qo_lens,
+            [torch.arange(s, device=device) for s in seq_lens],
+        )
+
+        raw_o, raw_lse = [], []
+        for r in range(DCP):
+            positions = [
+                torch.arange(r, max(r, s), DCP, device=device) for s in seq_lens
+            ]
+            local = torch.tensor([p.numel() for p in positions], dtype=torch.int32)
+            block_table = torch.zeros(
+                nr,
+                vllm_config.model_config.max_model_len // DCP,
+                dtype=torch.int32,
+                device=device,
+            )
+            for req, p in enumerate(positions):
+                block_table[req, : p.numel()] = (req * max_len + p).to(torch.int32)
+            common = CommonAttentionMetadata(
+                query_start_loc=qsl(qo_lens).to(device),
+                query_start_loc_cpu=qsl(host_qo),
+                seq_lens=torch.tensor(seq_lens, dtype=torch.int32, device=device),
+                num_reqs=nr,
+                num_actual_tokens=num_tokens,
+                max_query_len=QLEN,
+                max_seq_len=max_len,
+                block_table_tensor=block_table,
+                slot_mapping=torch.arange(num_tokens, dtype=torch.int64, device=device),
+                causal=True,
+                dcp_local_seq_lens=local.to(device),
+                dcp_local_seq_lens_cpu_upper_bound=local,
+            )
+            _install_fake_groups(DCP, r)
+            builder = AiterMLAMetadataBuilder(spec, [layer_name], vllm_config, device)
+            md = builder.build(0, common)
+            assert md.decode is not None
+            assert md.decode.dcp_route is expected
+            impl = _build_impl(builder.num_heads, DCP, r)
+            o, lse = impl.forward_mqa(q_flat, kv_flat, md, _Layer(device))
+            torch.accelerator.synchronize()
+
+            lse = lse.float().view(-1, gathered_heads).contiguous()
+            mask_dcp_empty_shards_(lse, md.decode.seq_lens, common.query_start_loc)
+            raw_o.append(o.view(-1, gathered_heads, KV_LORA_RANK))
+            raw_lse.append(lse)
+
+            po, plse = _varlen_reference(q_f32, kv_f32, seq_lens, qo_lens, positions)
+            seen = torch.isfinite(plse)
+            assert not bool((~seen & (lse[:num_tokens] > -1e30)).any())
+            shard_o = raw_o[-1][:num_tokens].float()[seen]
+            rel = (shard_o - po[seen]).abs().max() / po.abs().max()
+            assert rel < SHARD_RTOL, f"rank {r} rel err {rel:.3e}"
+            lse_err = (lse[:num_tokens][seen] - plse[seen]).abs().max()
+            assert lse_err < SHARD_RTOL, f"rank {r} lse err {lse_err:.3e}"
+
+        lses = torch.stack(raw_lse)
+        merged = sum(
+            correct_attn_out(raw_o[r].contiguous(), lses, r, CPTritonContext())[
+                0
+            ].float()
+            for r in range(DCP)
+        )
+        assert bool(torch.isfinite(merged).all())
+        merge_rel = (merged[:num_tokens] - ref_o).abs().max() / ref_o.abs().max()
+        assert merge_rel < MERGE_RTOL, f"merged rel err {merge_rel:.3e}"
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 8ed68986b267..b9bdf097a52f 100644
--- a/tests/v1/attention/test_rocm_aiter_mla_mtp_split.py
+++ b/tests/v1/attention/test_rocm_aiter_mla_mtp_split.py
@@ -973,13 +973,14 @@ def init_common_builder(self, *args, **kwargs):
     assert info_calls == [
         {
             "max_batch_size": config.scheduler_config.max_num_seqs,
-            "max_qo_len": builder.reorder_batch_threshold,
+            "max_qo_len": qlen,
             "num_attention_heads": AiterMLAHelper.get_actual_mla_num_heads(num_heads),
             "q_dtype": dtypes.fp8,
             "kv_dtype": dtypes.fp8,
             "is_sparse": False,
             "fast_mode": True,
         }
+        for qlen in range(1, builder.reorder_batch_threshold + 1)
     ]
     assert builder._mla_q_dtype == dtypes.fp8
     assert builder._mla_kv_dtype == dtypes.fp8
@@ -1014,6 +1015,118 @@ def init_common_builder(self, *args, **kwargs):
     )
 
 
+def test_builder_sizes_metadata_for_every_decode_qlen(monkeypatch):
+    """Size the persistent metadata for the largest need of any served qlen.
+
+    The builder serves every qlen up to reorder_batch_threshold (a DSpark draft
+    runs k-token blocks under a k + 1 threshold). On gfx950, aiter tiles a bf16
+    block of fewer than 128 packed rows by 64, so 16 heads x 7 tokens needs
+    twice the buffers of 16 heads x 8 tokens.
+    """
+
+    def get_mla_metadata_info_v1(
+        max_batch_size,
+        max_qo_len,
+        num_attention_heads,
+        q_dtype,
+        kv_dtype,
+        *,
+        is_sparse,
+        fast_mode,
+    ):
+        packed = max_qo_len * num_attention_heads
+        tiles = math.ceil(packed / (64 if 64 <= packed < 128 else 128))
+        work = max_batch_size * tiles
+        return (
+            (2, torch.uint64),
+            (257, torch.int32),
+            ((work + 255, 8), torch.int32),
+            (work + 1, torch.int32),
+            ((work, 2), torch.int32),
+            (2 * work, torch.int32),
+        )
+
+    def init_common_builder(self, *args, **kwargs):
+        self.num_heads = 16
+        self.q_data_type = torch.bfloat16
+        self.mla_dims = SimpleNamespace(
+            kv_lora_rank=512,
+            qk_nope_head_dim=128,
+            qk_rope_head_dim=64,
+            v_head_dim=128,
+        )
+        self.dcp_world_size = 1
+        self.cp_kv_cache_interleave_size = 1
+        self.reorder_batch_threshold = 8
+
+    monkeypatch.setitem(
+        sys.modules,
+        "aiter",
+        SimpleNamespace(
+            dtypes=SimpleNamespace(fp8="fp8", fp16="fp16", bf16="bf16"),
+            get_mla_metadata_info_v1=get_mla_metadata_info_v1,
+        ),
+    )
+    monkeypatch.setattr(
+        rocm_aiter_mla.MLACommonMetadataBuilder, "__init__", init_common_builder
+    )
+    monkeypatch.setattr(rocm_aiter_mla, "_fp8_mla_prefill_supported", lambda: False)
+    monkeypatch.setattr(rocm_aiter_mla, "is_aiter_found_and_supported", lambda: True)
+
+    config = SimpleNamespace(
+        speculative_config=SimpleNamespace(
+            method="deepseek_mtp", num_speculative_tokens=7, parallel_drafting=False
+        ),
+        parallel_config=SimpleNamespace(
+            tensor_parallel_size=4,
+            decode_context_parallel_size=1,
+            cp_kv_cache_interleave_size=1,
+        ),
+        model_config=SimpleNamespace(
+            max_model_len=16,
+            dtype=torch.bfloat16,
+            get_num_attention_heads=lambda parallel_config: 16,
+        ),
+        scheduler_config=SimpleNamespace(max_num_seqs=4),
+        cache_config=SimpleNamespace(cache_dtype="auto", num_gpu_blocks=None),
+        compilation_config=SimpleNamespace(
+            cudagraph_mode=SimpleNamespace(has_full_cudagraphs=lambda: False),
+            static_forward_context={},
+        ),
+    )
+    builder = AiterMLAMetadataBuilder(
+        kv_cache_spec=SimpleNamespace(
+            block_size=1, dtype=torch.bfloat16, cache_dtype_str="auto"
+        ),
+        layer_names=["layer.0"],
+        vllm_config=config,
+        device=torch.device("cpu"),
+    )
+
+    infos = [
+        get_mla_metadata_info_v1(
+            4, qlen, 16, torch.bfloat16, "bf16", is_sparse=False, fast_mode=True
+        )
+        for qlen in range(1, 9)
+    ]
+    buffers = [
+        builder._mla_work_meta_data,
+        builder._mla_work_indptr,
+        builder._mla_work_info_set,
+        builder._mla_reduce_indptr,
+        builder._mla_reduce_final_map,
+        builder._mla_reduce_partial_map,
+    ]
+    for i, buffer in enumerate(buffers):
+        needed = [
+            torch.Size(torch.empty(info[i][0], device="meta").shape) for info in infos
+        ]
+        assert all(
+            all(have >= need for have, need in zip(buffer.shape, shape))
+            for shape in needed
+        ), (i, buffer.shape, needed)
+
+
 def test_mtp_decode_qlen4_keeps_uniform_rows_with_metadata(monkeypatch):
     get_mla_metadata_v1 = mock.MagicMock()
     monkeypatch.setitem(
@@ -1462,7 +1575,11 @@ def test_a_two_token_bf16_block_is_allowed(monkeypatch):
 
 
 def _adaptive_config(
-    *, adaptive: bool = True, dcp: int = 1, num_speculative_tokens: int = 7
+    *,
+    adaptive: bool = True,
+    dcp: int = 1,
+    num_speculative_tokens: int = 7,
+    cp_interleave: int = 1,
 ) -> SimpleNamespace:
     return SimpleNamespace(
         speculative_config=SimpleNamespace(
@@ -1470,18 +1587,117 @@ def _adaptive_config(
             enable_adaptive_verification=adaptive,
             parallel_drafting=False,
         ),
-        parallel_config=SimpleNamespace(decode_context_parallel_size=dcp),
+        parallel_config=SimpleNamespace(
+            decode_context_parallel_size=dcp,
+            cp_kv_cache_interleave_size=cp_interleave,
+        ),
         use_v2_model_runner=True,
     )
 
 
-def test_varlen_cudagraph_bound_follows_adaptive_verification():
+def test_varlen_cudagraph_bound_follows_adaptive_verification(monkeypatch):
+    monkeypatch.setattr(rocm_aiter_mla, "_asm_dcp_verify_configured", lambda *a: True)
+    for route in ("_segmented_dcp_verify_supported", "_triton_dcp_verify_supported"):
+        monkeypatch.setattr(rocm_aiter_mla, route, lambda *a: False)
     bound = AiterMLAMetadataBuilder.get_varlen_cudagraph_max_query_len
     assert bound(_adaptive_config(), SimpleNamespace()) == 8
     assert bound(_adaptive_config(adaptive=False), SimpleNamespace()) is None
-    # DCP routes plan per-row windows from host lengths.
-    assert bound(_adaptive_config(dcp=2), SimpleNamespace()) is None
+    assert bound(_adaptive_config(dcp=8), SimpleNamespace()) == 8
+    assert bound(_adaptive_config(dcp=8, cp_interleave=2), SimpleNamespace()) is None
+    assert (
+        bound(_adaptive_config(dcp=8, num_speculative_tokens=1), SimpleNamespace())
+        is None
+    )
     # Kimi-K3 DSpark draft layers share the target's group; causality is
     # decided per batch, so the merged non-causal spec keeps the bound.
     shared_draft_group = SimpleNamespace(non_causal_multi_token_decode=True)
     assert bound(_adaptive_config(), shared_draft_group) == 8
+
+
+def test_backend_rejects_unsupported_adaptive_dcp(monkeypatch):
+    import vllm.config
+
+    monkeypatch.setattr(
+        vllm.config,
+        "get_current_vllm_config_or_none",
+        lambda: _adaptive_config(dcp=8, cp_interleave=2),
+    )
+
+    def validate(**flags):
+        return rocm_aiter_mla.AiterMLABackend.validate_configuration(
+            head_size=576,
+            dtype=torch.bfloat16,
+            kv_cache_dtype="fp8",
+            block_size=None,
+            use_mla=True,
+            has_sink=False,
+            use_sparse=False,
+            use_mm_prefix=False,
+            use_per_head_quant_scales=False,
+            device_capability=None,
+            attn_type="decoder",
+            **flags,
+        )
+
+    with pytest.raises(ValueError, match="decode context parallelism"):
+        validate(use_adaptive_verification=True, use_dcp=True)
+    validate(use_dcp=True)
+
+
+class _RecordingExpandKernel(_NoOpTritonKernel):
+    def __call__(self, *args, CLAMP_TAIL=False, **kwargs):
+        self.clamp_tail = CLAMP_TAIL
+
+
+def test_dcp_varlen_right_aligns_rows_and_keeps_lengths(monkeypatch):
+    monkeypatch.setitem(
+        sys.modules, "aiter", SimpleNamespace(get_mla_metadata_v1=mock.MagicMock())
+    )
+    expand = _RecordingExpandKernel()
+    monkeypatch.setattr(rocm_aiter_mla, "_expand_page_indices_kernel", expand)
+    width, num_reqs = 5, 4
+    builder = _builder(
+        mtp_decode_qlen=width,
+        dcp_world_size=8,
+        max_decode_rows=num_reqs * width,
+        num_heads=12,
+        asm_dcp_verify=True,
+        supports_segmented_dcp_verify=True,
+    )
+    rows = torch.arange(num_reqs * width, dtype=torch.int64)
+    builder._varlen_decode_qlen = width
+    builder._varlen_right_align = AiterMLAMetadataBuilder.__dict__[
+        "_varlen_right_align"
+    ].__get__(builder)
+    builder._varlen_row_req = rows // width
+    builder._varlen_row_pos = rows % width
+    builder._varlen_row_ids = rows
+    builder._varlen_q_rows = torch.zeros_like(rows)
+    builder._varlen_o_rows = torch.zeros(num_reqs * width + 1, dtype=torch.int64)
+    builder._fill_varlen_rows = AiterMLAMetadataBuilder._fill_varlen_rows.__get__(
+        builder
+    )
+
+    metadata = AiterMLAMetadataBuilder._build_decode(
+        builder,
+        block_table_tensor=torch.arange(num_reqs * 8, dtype=torch.int32).view(-1, 8),
+        seq_lens_device=torch.tensor([5, 3, 2, 0], dtype=torch.int32),
+        max_seq_len=40,
+        query_start_loc_cpu=torch.tensor([0, 3, 5, 7, 9], dtype=torch.int32),
+        query_start_loc_device=torch.tensor([0, 1, 6, 9, 9], dtype=torch.int32),
+        num_decode_tokens=9,
+        max_query_len=width,
+        dcp_tot_seq_lens_device=torch.tensor([40, 21, 9, 0], dtype=torch.int32),
+    )
+
+    assert metadata.dcp_route is rocm_aiter_mla._DCPDecodeRoute.CPRR
+    assert metadata.seq_lens.tolist() == [5, 3, 2, width]
+    assert metadata.g_kv_indptr is not None
+    assert metadata.g_kv_indptr.tolist()[: num_reqs + 1] == [0, 40, 61, 70, 110]
+    assert metadata.qo_indptr is not None
+    assert metadata.qo_indptr.tolist() == [0, 5, 10, 15, 20]
+    assert expand.clamp_tail is False
+    assert metadata.varlen_q_rows is not None and metadata.varlen_o_rows is not None
+    real_rows = [4, 5, 6, 7, 8, 9, 12, 13, 14]
+    assert metadata.varlen_o_rows[:9].tolist() == real_rows
+    assert metadata.varlen_q_rows[real_rows].tolist() == list(range(9))
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 201b569629d6..92cb645b7868 100644
--- a/tests/v1/attention/test_sparse_indexer_decode_seq_lens.py
+++ b/tests/v1/attention/test_sparse_indexer_decode_seq_lens.py
@@ -20,6 +20,7 @@
 
 from types import ModuleType
 
+import pytest
 import torch
 
 # Bootstrap the glm5next package before importing kpool_compress directly: its
@@ -35,8 +36,10 @@
 # isort: on
 
 from vllm.models.glm5next.common.sparse_indexer import (
+    _build_decode_scatter_indices,
     _decode_topk_seq_lens,
     _fill_short_decode_causal_indices,
+    _scatter_decode_tokens_by_request,
 )
 from vllm.platforms import current_platform
 
@@ -55,6 +58,67 @@
 SELECT_K = TOPK_TOKENS // KPOOL
 
 
+def test_decode_scatter_discards_full_graph_token_padding():
+    """Adaptive verification can leave fewer live tokens than the graph size."""
+    decode_lens = torch.tensor([3, 3, 0, 0], dtype=torch.int64)
+    # Six padding tokens is deliberately larger than lmax. They must alias a
+    # valid discard cell instead of indexing columns 0..5 in a width-3 row.
+    tokens = torch.arange(12, dtype=torch.int32)
+    scatter_indices = _build_decode_scatter_indices(decode_lens, 4, 12)
+
+    out = _scatter_decode_tokens_by_request(
+        tokens, -1, num_requests=4, lmax=3, scatter_indices=scatter_indices
+    )
+
+    assert out.tolist() == [
+        [0, 1, 2],
+        [3, 4, 5],
+        [-1, -1, -1],
+        [-1, -1, -1],
+    ]
+
+
+def test_decode_scatter_without_graph_padding_is_unchanged():
+    decode_lens = torch.tensor([1, 3, 2], dtype=torch.int64)
+    tokens = torch.arange(6, dtype=torch.int32)
+    scatter_indices = _build_decode_scatter_indices(decode_lens, 3, 6)
+
+    out = _scatter_decode_tokens_by_request(
+        tokens, -1, num_requests=3, lmax=3, scatter_indices=scatter_indices
+    )
+
+    assert out.tolist() == [[0, -1, -1], [1, 2, 3], [4, 5, -1]]
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="requires CUDA")
+@torch.inference_mode()
+def test_decode_scatter_cuda_graph_replays_with_smaller_live_budget():
+    decode_lens = torch.full((4,), 3, dtype=torch.int64, device="cuda")
+    tokens = torch.arange(12, dtype=torch.int32, device="cuda")
+
+    def run_scatter():
+        scatter_indices = _build_decode_scatter_indices(decode_lens, 4, 12)
+        return _scatter_decode_tokens_by_request(
+            tokens, -1, num_requests=4, lmax=3, scatter_indices=scatter_indices
+        )
+
+    # Warm allocation and repeat_interleave kernels before capture.
+    run_scatter()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        out = run_scatter()
+
+    decode_lens.copy_(torch.tensor([3, 3, 0, 0], dtype=torch.int64, device="cuda"))
+    graph.replay()
+
+    assert out.cpu().tolist() == [
+        [0, 1, 2],
+        [3, 4, 5],
+        [-1, -1, -1],
+        [-1, -1, -1],
+    ]
+
+
 def test_kpool_ops_dispatch_matches_platform():
     expected_backend = ".amd." if current_platform.is_rocm() else ".nvidia."
     assert expected_backend in indexer_mod.kpool_ops.__name__
diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py
index 565a517532be..0b9c6fc34423 100644
--- a/tests/v1/attention/test_sparse_mla_backends.py
+++ b/tests/v1/attention/test_sparse_mla_backends.py
@@ -114,9 +114,11 @@
 )
 from vllm.v1.attention.ops import flashmla
 from vllm.v1.hisparse import runtime as hisparse_runtime
+from vllm.v1.hisparse.backend import get_hisparse_backend
 from vllm.v1.hisparse.runtime import (
     HiSparseCacheHandle,
     HiSparseRuntime,
+    HiSparseSlotInvalidator,
     ResolvedHiSparseConfig,
     _has_hisparse_ops,
     build_hisparse_prefill_staging_plan,
@@ -202,10 +204,19 @@ def test_rocm_long_bf16_prefill_avoids_nonpersistent_aiter(
     if not current_platform.is_rocm():
         pytest.skip("ROCm sparse MLA regression")
     from vllm._aiter_ops import rocm_aiter_ops
+    from vllm.v1.attention.ops import rocm_aiter_mla_sparse
 
-    # Exercise the gfx942 routing decision on other ROCm devices too;
-    # the numerical kernel still compiles for the actual device.
+    # Exercise the gfx942 routing decision on other ROCm devices too.
+    # Pin get_cdna_version() to the physical chip before the on_gfx942()
+    # patch, so an MI350 stays off pa_prefill_sparse and the Triton kernel
+    # that compiles there is the one that runs.
+    physical_cdna_version = rocm_aiter_mla_sparse.get_cdna_version()
     monkeypatch.setattr("vllm.platforms.rocm.on_gfx942", lambda: True)
+    monkeypatch.setattr(
+        rocm_aiter_mla_sparse,
+        "get_cdna_version",
+        lambda: physical_cdna_version,
+    )
 
     def unsupported_aiter(*args, **kwargs):
         pytest.fail("Long gfx942 BF16 prefill reached nonpersistent AITER")
@@ -286,6 +297,94 @@ def fake_flashinfer(**kwargs):
     assert recorded_scale != kv_lora_rank**-0.5
 
 
+@pytest.mark.skipif(
+    not current_platform.is_device_capability_family(100),
+    reason="trtllm-gen sparse MLA requires SM10x",
+)
+@pytest.mark.parametrize(
+    ("qk_nope_head_dim", "qk_rope_head_dim", "v_head_dim"),
+    [(128, 64, 128), (256, 0, 256)],
+)
+def test_flashinfer_sparse_mla_attends_tokens_past_grid_limit(
+    default_vllm_config, dist_init, qk_nope_head_dim, qk_rope_head_dim, v_head_dim
+):
+    """trtllm-gen maps each token to gridDim.z (limit 65535); a larger step,
+    e.g. a 64K-token prefill, must still attend every token."""
+    device = torch.device("cuda")
+    num_tokens, num_heads, kv_lora_rank, topk, block_size = (
+        65536 + 4099,
+        16,
+        512,
+        128,
+        64,
+    )
+    head_size = kv_lora_rank + qk_rope_head_dim
+    impl = FlashInferMLASparseImpl(
+        num_heads=num_heads,
+        head_size=head_size,
+        scale=head_size**-0.5,
+        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=ColumnParallelLinear(
+            kv_lora_rank, num_heads * (qk_nope_head_dim + v_head_dim), bias=False
+        ),
+    )
+    gen = torch.Generator(device=device).manual_seed(0)
+    kv_cache = (
+        torch.rand(64, block_size, head_size, generator=gen, device=device) - 0.5
+    ).to(torch.bfloat16)
+    q = (
+        torch.rand(num_tokens, num_heads, head_size, generator=gen, device=device) - 0.5
+    ).to(torch.bfloat16)
+    # Mixed active lengths, valid indices first and -1 padding after, as the
+    # index converters produce.
+    valid_counts = torch.randint(
+        1, topk + 1, (num_tokens,), generator=gen, device=device, dtype=torch.int32
+    )
+    valid_counts[::3] = topk
+    topk_indices = torch.randint(
+        0,
+        kv_cache.shape[0] * block_size,
+        (num_tokens, topk),
+        generator=gen,
+        device=device,
+        dtype=torch.int32,
+    )
+    slots = torch.arange(topk, device=device)
+    topk_indices.masked_fill_(slots[None, :] >= valid_counts[:, None], -1)
+
+    out, lse = impl._forward_mqa_kernel(
+        q,
+        kv_cache,
+        topk_indices,
+        valid_counts,
+        layer=SimpleNamespace(),
+        block_size=block_size,
+        is_decode=False,
+    )
+
+    assert lse is None
+    kv_rows = kv_cache.view(-1, head_size).float()
+    # Both sides of the 65535 boundary, the ends, and a sweep in between.
+    rows = sorted({0, 65534, 65535, 65536, num_tokens - 1, *range(0, num_tokens, 997)})
+    for t in rows:
+        kv = kv_rows[topk_indices[t, : valid_counts[t]].long()]
+        scores = q[t].float() @ kv.T * impl.scale
+        ref = torch.softmax(scores, dim=-1) @ kv[:, :kv_lora_rank]
+        torch.testing.assert_close(out[t].float(), ref, atol=2e-2, rtol=2e-2)
+
+
 def test_hisparse_routes_prefill_to_sparse_mqa():
     layer = SimpleNamespace(hisparse_cache=object())
 
@@ -2014,6 +2113,27 @@ def test_triton_convert_returns_valid_counts(num_topk_tokens: int, reuse_buffers
         assert torch.all(result[row, num_valid:] == -1)
 
 
+def test_hisparse_generic_stream_orders_ipc_event():
+    """Generic copy streams must record IPC events and restore the compute stream."""
+    from vllm.v1.hisparse.backend import CUDAHiSparseBackend
+
+    backend = CUDAHiSparseBackend()
+    compute_stream = torch.accelerator.current_stream()
+    copy_stream = torch.Stream(device=torch.device(DEVICE_TYPE))
+    values = torch.arange(8, dtype=torch.float32, device=DEVICE_TYPE)
+    input_ready = torch.Event()
+    input_ready.record(compute_stream)
+    copy_stream.wait_event(input_ready)
+    output_ready = backend.create_ipc_event()
+    with copy_stream:
+        values.add_(1)
+        output_ready.record(copy_stream)
+    assert torch.cuda.current_stream().stream_id == compute_stream.stream_id
+    compute_stream.wait_event(output_ready)
+
+    torch.testing.assert_close(values.cpu(), torch.arange(8, dtype=torch.float32) + 1)
+
+
 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)
@@ -2135,6 +2255,36 @@ def test_sparse_mla_index_groups_own_distinct_physical_buffers():
     )
 
 
+@pytest.mark.parametrize(
+    ("kv_cache_dtype", "row_width", "kv_dtype"),
+    [
+        ("fp8_ds_mla", 528, torch.uint8),
+        ("fp8_ds_mla", 656, torch.uint8),
+        ("bfloat16", 512, torch.bfloat16),
+    ],
+)
+def test_hisparse_index_group_preserves_row_width(
+    monkeypatch, kv_cache_dtype, row_width, kv_dtype
+):
+    create_handle = MagicMock(side_effect=lambda *args, **kwargs: SimpleNamespace())
+    monkeypatch.setattr(
+        index_group_module, "create_hisparse_cache_handle", create_handle
+    )
+    config = SimpleNamespace(
+        attention_config=SimpleNamespace(hisparse_config=object()), model_config=None
+    )
+    builder = SparseMLAIndexGroupBuilder(torch.empty((2, 4), dtype=torch.int32))
+    for leader in (True, False):
+        group, layer_index = builder.register_layer(
+            leader, config, row_width=row_width, kv_cache_dtype=kv_cache_dtype
+        )
+        assert isinstance(group, HiSparseMLAIndexGroup)
+        assert layer_index == (0 if leader else 1)
+        assert create_handle.call_args.kwargs["row_width"] == row_width
+        assert create_handle.call_args.kwargs["kv_dtype"] == kv_dtype
+        assert create_handle.call_args.kwargs["is_index_group_leader"] == leader
+
+
 # HiSparse is host-resident-only and kernel-only: runtime construction
 # raises without the compiled CUDA ops.
 requires_hisparse_ops = pytest.mark.skipif(
@@ -2153,6 +2303,7 @@ def _make_hisparse_runtime(
     max_swap_rows: int | None = None,
     index_group: hisparse_runtime.HiSparseIndexGroup | None = None,
     max_union_rows: int | None = None,
+    kv_dtype: torch.dtype = torch.float32,
 ) -> HiSparseRuntime:
     runtime = HiSparseRuntime(
         config=ResolvedHiSparseConfig(
@@ -2162,7 +2313,7 @@ def _make_hisparse_runtime(
         ),
         max_num_reqs=max_num_reqs,
         row_width=row_width,
-        kv_dtype=torch.float32,
+        kv_dtype=kv_dtype,
         device=DEVICE_TYPE,
         max_swap_rows=max_swap_rows,
         index_group=index_group,
@@ -2173,7 +2324,7 @@ def _make_hisparse_runtime(
     num_blocks = max_num_reqs * blocks_per_request + 1
     raw = torch.zeros(
         num_blocks * block_size * row_width,
-        dtype=torch.float32,
+        dtype=kv_dtype,
         device=DEVICE_TYPE,
     ).view(torch.int8)
     block_table = torch.arange(
@@ -2184,7 +2335,7 @@ def _make_hisparse_runtime(
     runtime.bind_hot_cache(
         raw,
         byte_offset=0,
-        block_stride=block_size * row_width * torch.float32.itemsize,
+        block_stride=block_size * row_width * kv_dtype.itemsize,
         num_blocks=num_blocks,
         block_size=block_size,
         block_table=block_table,
@@ -2646,17 +2797,20 @@ def test_hisparse_resident_rows_bypass_hot_lru():
 
 
 @requires_hisparse_ops
-def test_hisparse_swap_in_preserves_rows_across_eviction():
+@pytest.mark.parametrize(
+    ("row_width", "kv_dtype"),
+    [(64, torch.float32), (528, torch.uint8), (656, torch.uint8)],
+)
+def test_hisparse_swap_in_preserves_rows_across_eviction(row_width, kv_dtype):
     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
+    kv_pool = torch.randint(
+        0, 256, (num_blocks + 1, block_size, row_width), dtype=kv_dtype
     ).pin_memory()
     flat_pool = kv_pool.reshape(-1, row_width)
 
@@ -2665,6 +2819,7 @@ def test_hisparse_swap_in_preserves_rows_across_eviction():
         device_buffer_size=buf,
         max_num_reqs=num_reqs,
         row_width=row_width,
+        kv_dtype=kv_dtype,
     )
     runtime.bind_source_cache(kv_pool)
     cache = HiSparseCacheHandle(runtime)
@@ -2678,7 +2833,7 @@ def test_hisparse_swap_in_preserves_rows_across_eviction():
     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):
+    for step in (0, 0, 1, 2, 0):
         cache.runtime.begin_forward()
         topk = torch.stack(
             [(base + step * top_k + row * 17) % seq_len for row in range(num_reqs)]
@@ -2709,7 +2864,13 @@ def test_hisparse_swap_in_preserves_rows_across_eviction():
             hot_indices[valid].to(torch.long)
         ].cpu()
         expected = flat_pool[global_ref[valid].cpu().to(torch.long)]
-        torch.testing.assert_close(gathered, expected)
+        torch.testing.assert_close(gathered, expected, rtol=0, atol=0)
+
+    hits, misses = runtime.index_group.swap_stats.tolist()
+    valid_rows_per_step = num_reqs * (top_k - 1)
+    assert hits >= valid_rows_per_step
+    assert misses > 3 * valid_rows_per_step
+    assert hits + misses == 5 * valid_rows_per_step
 
 
 @requires_hisparse_ops
@@ -2949,6 +3110,7 @@ def test_hisparse_row_dma_copies_discontiguous_spans_across_layers():
         for _ in resident_caches
     )
     worker = object.__new__(HiSparseConnectorWorker)
+    worker.backend = get_hisparse_backend()
     worker.is_host_writer = True
     worker.kernel_block_size = block_size
     worker.resident_caches = resident_caches
@@ -2960,7 +3122,7 @@ def test_hisparse_row_dma_copies_discontiguous_spans_across_layers():
         )
         for _ in resident_caches
     ]
-    worker.dma_stream = torch.cuda.Stream(device=device)
+    worker.dma_stream = torch.Stream(device=device)
     worker.host_write_event = torch.Event()
     worker._dma_free_descriptors = []
     worker._pending_dma_descriptors = deque()
@@ -3167,8 +3329,10 @@ def test_hisparse_newest_write_and_recycled_slot_invalidation():
     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
+    HiSparseSlotInvalidator(request_state_indices)(
+        [cache_handle.runtime],
+        torch.tensor([8], dtype=torch.int64, device=device),
+        req_ids,
     )
     cache_handle.runtime.begin_forward()
     hot_indices = cache_handle.swap_in(
@@ -3185,6 +3349,40 @@ def test_hisparse_newest_write_and_recycled_slot_invalidation():
     torch.testing.assert_close(flat_hot[idx].cpu(), flat_pool[8])
 
 
+@requires_hisparse_ops
+def test_hisparse_invalidation_covers_every_index_group():
+    """One launch drops a rewritten slot from each listed group's request row."""
+    device = torch.device(DEVICE_TYPE)
+    tables = [
+        torch.tensor(
+            [[5, 7, 9, 7, 1, 2, 3, 4], [7, 1, 2, 3, 7, 7, 7, 7]],
+            dtype=torch.int32,
+            device=device,
+        )
+        + offset
+        for offset in (0, 0, 100)
+    ]
+    expected = [table.clone() for table in tables]
+    leaders = [
+        SimpleNamespace(
+            index_group=SimpleNamespace(device_global_indices=table),
+            cache_ops=torch.ops._C_cache_ops,
+        )
+        for table in tables[:2]
+    ]
+    for table in expected[:2]:
+        table[1][table[1] == 7] = -1
+
+    HiSparseSlotInvalidator(torch.tensor([1, 0], dtype=torch.int32, device=device))(
+        leaders,
+        torch.tensor([7, 7], dtype=torch.int64, device=device),
+        torch.tensor([0, -1], dtype=torch.int32, device=device),
+    )
+
+    for table, want in zip(tables, expected):
+        torch.testing.assert_close(table, want)
+
+
 @requires_hisparse_ops
 def test_hisparse_mixed_batch_bf16_row_split(
     default_vllm_config, dist_init, workspace_init, monkeypatch
@@ -3573,12 +3771,16 @@ def test_hisparse_prefill_resident_sources_follow_each_resident_group():
 
 
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
-def test_hisparse_gather_prefill_cache_prefers_resident_rows():
+@pytest.mark.parametrize(
+    ("row_width", "kv_dtype"),
+    [(16, torch.float32), (528, torch.uint8), (656, torch.uint8)],
+)
+def test_hisparse_gather_prefill_cache_prefers_resident_rows(row_width, kv_dtype):
     """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_size, resident_block_size = 4, 2
     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(
@@ -3596,24 +3798,23 @@ def test_hisparse_gather_prefill_cache_prefers_resident_rows():
     )
 
     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)
+    host_cache = torch.randint(
+        0, 128, (num_host_blocks, block_size, row_width), dtype=kv_dtype
+    ).pin_memory()
+    resident_cache = torch.randint(
+        128,
+        256,
+        (num_res_blocks, resident_block_size, row_width),
+        dtype=kv_dtype,
+        device=device,
     )
 
-    staged = HiSparseRuntime.gather_prefill_cache(
-        None, host_cache, plan, resident_cache=resident_cache
+    runtime = _make_hisparse_runtime(
+        row_width=row_width, block_size=block_size, kv_dtype=kv_dtype
+    )
+    runtime.bind_source_cache(host_cache)
+    staged = runtime.gather_prefill_cache(
+        host_cache, plan, resident_cache=resident_cache
     )
 
     staged_flat = staged.view(-1, row_width).cpu()
@@ -3634,6 +3835,35 @@ def test_hisparse_gather_prefill_cache_prefers_resident_rows():
         torch.testing.assert_close(staged_flat[i], expected)
 
 
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
+def test_hisparse_staged_prefill_cache_backs_min_attention_rows():
+    """A staged cache of a few blocks still has 256 allocated rows behind it.
+
+    trtllm-gen sparse MLA reads up to ~250 rows from the start of its KV cache
+    whatever its size, faulting when a smaller cache ends at an unmapped page.
+    """
+    if not _has_hisparse_ops():
+        pytest.skip("hisparse CUDA ops unavailable")
+    device = torch.device("cuda")
+    block_size, row_width = 64, 16
+    block_table = torch.tensor([[3, 1]], dtype=torch.int32, device=device)
+    seq_lens = torch.tensor([100], dtype=torch.int32, device=device)
+    plan = build_hisparse_prefill_staging_plan(
+        block_table, seq_lens, block_size, staging_block_capacity=2
+    )
+    host_cache = torch.zeros(4, block_size, row_width).pin_memory()
+
+    runtime = _make_hisparse_runtime(row_width=row_width, block_size=block_size)
+    runtime.bind_source_cache(host_cache)
+    staged = runtime.gather_prefill_cache(host_cache, plan)
+
+    assert staged.shape[0] * block_size < 256
+    backing_rows = staged.untyped_storage().nbytes() // (
+        row_width * staged.element_size()
+    )
+    assert backing_rows >= 256
+
+
 def test_hisparse_fp8_decode_resolves_rows_once_then_runs_batched_attention():
     """Verification rows resolve in one swap_in; FP8 attention sees them batched."""
     device = torch.device("cpu")
@@ -4036,42 +4266,69 @@ def hook(rows, cache, indices, counts, *, is_decode, **kwargs):
     )
 
 
-@pytest.mark.parametrize("max_swap_rows", [1, 4], ids=["over_capacity", "mixed"])
-def test_hisparse_resident_prefill_uses_attention_block_stride(
-    monkeypatch, max_swap_rows
+@pytest.mark.parametrize(
+    ("num_decode_tokens", "max_swap_rows"),
+    [(3, 1), (1, 4)],
+    ids=["over_capacity", "mixed"],
+)
+def test_hisparse_resident_prefill_reads_residency_by_state_row(
+    num_decode_tokens, max_swap_rows
 ):
-    """Resident batches with prefill rows read resident pages directly, whether
-    they exceed the swap-row capacity or fit it."""
-    expected = torch.tensor([[19]], dtype=torch.int32)
-    cache_handle = _StubHiSparseCacheHandle(
-        all_context_pages_resident=True,
-        num_decode_tokens=1,
-        num_actual_tokens=2,
-        view=SimpleNamespace(block_size=64, attention_block_stride=832),
-        block_table=torch.tensor([[3]], dtype=torch.int32),
-        batch_block_table=lambda: torch.tensor([[3]], dtype=torch.int32),
-        runtime=SimpleNamespace(max_swap_rows=max_swap_rows),
+    """Resident prefill conversion reads each request's residency row through
+    its state slot, at the attention block stride. An unset slot reads slot 0."""
+    topk = torch.zeros((3, 128), dtype=torch.int32, device=DEVICE_TYPE)
+    # [resident groups, state slots, pages]; the leader reads group 1.
+    residency = torch.tensor(
+        [[[7], [8]], [[3], [5]]], dtype=torch.int32, device=DEVICE_TYPE
+    )
+    runtime = SimpleNamespace(
+        is_group_leader=True,
+        begin_forward=lambda: None,
+        eager_host_mirror=True,
+        max_swap_rows=max_swap_rows,
+        request_state_indices=torch.tensor(
+            [1, 0, -1], dtype=torch.int32, device=DEVICE_TYPE
+        ),
     )
-    index_group = object.__new__(HiSparseMLAIndexGroup)
-    index_group.caches = [cache_handle]
-    index_group.physical_topk_indices = torch.empty((2, 1), dtype=torch.int32)
-    convert_once = MagicMock(return_value=expected)
-    monkeypatch.setattr(index_group, "_convert_once", convert_once)
-    topk = torch.zeros((2, 1), dtype=torch.int32)
-    metadata = SimpleNamespace(
-        req_id_per_token=torch.zeros(2, dtype=torch.int32),
+    cache = HiSparseCacheHandle(runtime)
+    cache.view = SimpleNamespace(block_size=64, attention_block_stride=832)
+    cache.block_table = residency[1]
+    group = HiSparseMLAIndexGroup(
+        logical_topk_indices=topk,
+        physical_topk_indices=torch.empty_like(topk),
+        valid_topk_counts=torch.empty(3, dtype=torch.int32, device=DEVICE_TYPE),
+        side_stream=torch.Stream(),
+        logical_topk_ready=torch.Event(),
+        physical_topk_ready=torch.Event(),
+        has_indexer=True,
+        caches=[cache],
     )
-
-    result = index_group.convert_logical_to_physical_topk(
-        0,
-        topk,
-        metadata,
-        block_stride_rows=None,
-        return_valid_counts=False,
+    metadata = SimpleNamespace(
+        num_actual_tokens=3,
+        num_decode_tokens=num_decode_tokens,
+        max_query_len=1 if num_decode_tokens == 3 else 2,
+        num_reqs=3,
+        req_id_per_token=torch.tensor([0, 1, 2], dtype=torch.int32, device=DEVICE_TYPE),
     )
 
-    assert result is expected
-    assert convert_once.call_args.kwargs["block_stride_rows"] == 832
+    # Repeat after compilation so its host delay cannot hide a stream race, and
+    # change the residency so a stale table cannot pass.
+    for step in range(2):
+        residency.add_(step * 10)
+        expected = torch.tensor(
+            [[5 + step * 10], [3 + step * 10], [3 + step * 10]], device=DEVICE_TYPE
+        )
+        expected *= 832
+        torch.accelerator.synchronize()
+        group.prepare_for_batch(0, metadata)
+        group.set_logical_topk_ready(0)
+        # Compute-stream work between the indexer and the layer must not delay
+        # the inputs the side-stream converter reads.
+        torch.cuda._sleep(200_000_000)
+        physical = group.convert_logical_to_physical_topk(
+            0, topk, metadata, block_stride_rows=None, return_valid_counts=False
+        )
+        torch.testing.assert_close(physical, expected.int().expand_as(topk))
 
 
 def test_hisparse_flashmla_reorders_full_speculative_window_as_decode():
@@ -4201,10 +4458,11 @@ def fake_fa4(**kwargs):
     def fake_trtllm(**kwargs):
         trtllm_calls.append(kwargs)
         rows = kwargs["query"].shape[0]
-        out = torch.full(
-            (rows, 1, num_heads, 512), 2.0, dtype=torch.bfloat16, device=device
-        )
-        return (out, lse(rows)) if kwargs.get("return_lse") else out
+        # Like FlashInfer, write into the caller-owned ``out`` / ``lse``.
+        out = kwargs["out"].fill_(2.0)
+        if not kwargs.get("return_lse"):
+            return out
+        return out, kwargs["lse"].copy_(lse(rows))
 
     monkeypatch.setattr(flashattn_sparse_module, "flash_attn_varlen_func", fake_fa4)
     monkeypatch.setattr(
@@ -4772,7 +5030,7 @@ 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(torch.accelerator, "current_stream", lambda: compute_stream)
     monkeypatch.setattr(
         index_group_module.ops, "cp_gather_and_upconvert_fp8_kv_cache", gather
     )
@@ -4783,7 +5041,7 @@ def test_hisparse_fp8_prefill_gather_uses_dedicated_stream(monkeypatch):
         view=SimpleNamespace(cache=resident_cache, block_size=4),
         block_table=torch.tensor([[2, 3]], dtype=torch.int32),
         runtime=SimpleNamespace(
-            request_state_indices=torch.tensor([0], dtype=torch.int32)
+            request_state_indices=torch.tensor([0], dtype=torch.int32),
         ),
     )
     index_group = object.__new__(HiSparseMLAIndexGroup)
diff --git a/tests/v1/attention/test_sparse_mla_kv_cache_layout.py b/tests/v1/attention/test_sparse_mla_kv_cache_layout.py
index b3dd64800c24..7f61e235d567 100644
--- a/tests/v1/attention/test_sparse_mla_kv_cache_layout.py
+++ b/tests/v1/attention/test_sparse_mla_kv_cache_layout.py
@@ -80,6 +80,7 @@ def test_allocation_and_warmup_follow_addressing_mode(
     layer = SimpleNamespace(
         _vllm_config=config,
         attn_backend=backend,
+        impl=object.__new__(backend.get_impl_cls()),
         kv_cache_dtype=cache_dtype,
         head_size=576,
         topk_buffer_width=None,
diff --git a/tests/v1/attention/test_xpu_mla_sparse_backend.py b/tests/v1/attention/test_xpu_mla_sparse_backend.py
new file mode 100644
index 000000000000..43e0227f6381
--- /dev/null
+++ b/tests/v1/attention/test_xpu_mla_sparse_backend.py
@@ -0,0 +1,24 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+
+from vllm.platforms import current_platform
+
+pytestmark = pytest.mark.skipif(
+    not current_platform.is_xpu(),
+    reason="XPU sparse MLA backend",
+)
+
+
+def test_xpu_sparse_impl_implements_logical_topk_hook():
+    """``MLAAttention.forward`` calls ``impl.record_logical_topk_ready()``
+    unconditionally for sparse layers, so every sparse MLA impl must define it.
+    Without it DeepSeek-V3.2 / GLM-5 crash on the first forward on XPU."""
+    from vllm.v1.attention.backends.mla.xpu_mla_sparse import XPUMLASparseImpl
+
+    hook = getattr(XPUMLASparseImpl, "record_logical_topk_ready", None)
+    assert callable(hook)
+    # The XPU impl does not use sparse-MLA index groups: the hook is a no-op.
+    impl = object.__new__(XPUMLASparseImpl)
+    impl.record_logical_topk_ready()
diff --git a/tests/v1/attention/utils.py b/tests/v1/attention/utils.py
index 0b85e4010050..14265e490abf 100644
--- a/tests/v1/attention/utils.py
+++ b/tests/v1/attention/utils.py
@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Utility functions for attention-related v1 tests."""
 
+import os
 from dataclasses import dataclass
 
 import pytest
@@ -232,15 +233,25 @@ def create_vllm_config(
     if hf_config_override:
         model_config.hf_config.update(hf_config_override)
 
-    return VllmConfig(
-        model_config=model_config,
-        cache_config=cache_config,
-        parallel_config=parallel_config,
-        scheduler_config=scheduler_config,
-        device_config=device_config,
-        load_config=load_config,
-        compilation_config=compilation_config,
-    )
+    # VllmConfig() auto-enables breakable cudagraphs process-wide via
+    # VLLM_USE_BREAKABLE_CUDAGRAPH for breakable-by-default models; keep
+    # that from leaking into unrelated tests sharing this process.
+    breakable_env = os.environ.get("VLLM_USE_BREAKABLE_CUDAGRAPH")
+    try:
+        return VllmConfig(
+            model_config=model_config,
+            cache_config=cache_config,
+            parallel_config=parallel_config,
+            scheduler_config=scheduler_config,
+            device_config=device_config,
+            load_config=load_config,
+            compilation_config=compilation_config,
+        )
+    finally:
+        if breakable_env is None:
+            os.environ.pop("VLLM_USE_BREAKABLE_CUDAGRAPH", None)
+        else:
+            os.environ["VLLM_USE_BREAKABLE_CUDAGRAPH"] = breakable_env
 
 
 @dataclass
diff --git a/tests/v1/core/test_async_scheduler.py b/tests/v1/core/test_async_scheduler.py
index 627bba3afb02..6a24baa77a76 100644
--- a/tests/v1/core/test_async_scheduler.py
+++ b/tests/v1/core/test_async_scheduler.py
@@ -509,7 +509,7 @@ def _create_async_pp_scheduler(
     # Emulate PP at the scheduler level; constructing with
     # pipeline_parallel_size>1 requires that many visible GPUs. Drive with
     # queue_size=pp_size+1 (V2 async PP runs pp_size+1 concurrent batches).
-    scheduler.pp_size = pp_size
+    scheduler.decode_stagger = pp_size
     scheduler.use_pp = pp_size > 1
     return scheduler
 
diff --git a/tests/v1/core/test_deferred_block_free.py b/tests/v1/core/test_deferred_block_free.py
index 44cdf567f7c9..aa2244e069ed 100644
--- a/tests/v1/core/test_deferred_block_free.py
+++ b/tests/v1/core/test_deferred_block_free.py
@@ -13,13 +13,14 @@
 
 import os
 import time
-from unittest.mock import PropertyMock, patch
+from unittest.mock import MagicMock, PropertyMock, patch
 
 import pytest
 
 from vllm.config import VllmConfig
 from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy
 from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.metrics.stats import PrefixCacheStats
 from vllm.v1.outputs import ModelRunnerOutput
 from vllm.v1.request import RequestStatus
 
@@ -544,3 +545,105 @@ def test_cow_retentions_deferred_until_copy_step_processed():
     assert not scheduler.deferred_frees
     out2 = scheduler.schedule()
     assert [r.req_id for r in out2.scheduled_new_reqs] == ["late"]
+
+
+def _setup_async_load(defer_async_loads: bool):
+    """Schedule a running request's prefill and one in-flight decode, then
+    admit a request whose KV an async connector load fills.
+
+    Returns (scheduler, importer, outputs, update_state_after_alloc mock).
+    """
+    scheduler = create_scheduler(
+        model=MODEL,
+        async_scheduling=True,
+        use_kv_connector=mock_kv(matched_tokens=0, is_async=False),
+    )
+    scheduler.defer_block_free = True
+    scheduler.defer_async_loads = defer_async_loads
+    _, out0, out1 = _setup_request_with_inflight_step(scheduler)
+
+    importer = create_requests(
+        num_requests=1, num_tokens=NUM_PROMPT_TOKENS, max_tokens=5, req_ids=["imp"]
+    )[0]
+    scheduler.add_request(importer)
+    connector = scheduler.connector
+    assert connector is not None
+    with (
+        patch.object(
+            connector,
+            "get_num_new_matched_tokens",
+            return_value=(2 * scheduler.block_size, True),
+        ),
+        patch.object(connector, "update_state_after_alloc") as registered,
+    ):
+        out2 = scheduler.schedule()
+    assert importer.status == RequestStatus.WAITING_FOR_REMOTE_KVS
+    return scheduler, importer, [out0, out1, out2], registered
+
+
+def test_deferred_async_load_registers_after_inflight_step():
+    """HiSparse keeps reading a released page until the pool reuses it, and an
+    async connector load isn't ordered against the in-flight step, so the load
+    is registered only once that step is processed.
+    """
+    scheduler, importer, outputs, registered = _setup_async_load(defer_async_loads=True)
+    out0, out1, _ = outputs
+    registered.assert_not_called()
+
+    connector = scheduler.connector
+    with patch.object(connector, "update_state_after_alloc") as registered:
+        scheduler.update_from_output(out0, _make_model_runner_output(out0))
+        outputs.append(scheduler.schedule())
+        registered.assert_not_called()
+
+        scheduler.update_from_output(out1, _make_model_runner_output(out1))
+        outputs.append(scheduler.schedule())
+    registered.assert_called_once()
+    assert registered.call_args.args[0] is importer
+
+
+def test_deferred_async_load_records_hit_sources_when_started():
+    """The connector attributes cache-hit tiers only once the load is
+    registered, so a deferred load records its prefix cache stats then."""
+    scheduler, importer, outputs, _ = _setup_async_load(defer_async_loads=True)
+    out0, out1, _ = outputs
+    connector = scheduler.connector
+    scheduler.connector_prefix_cache_stats = PrefixCacheStats()
+    sources = MagicMock()
+    with (
+        patch.object(connector, "update_state_after_alloc"),
+        patch.object(
+            connector, "get_external_cache_hit_sources", return_value=sources
+        ) as hit_sources,
+        patch.object(PrefixCacheStats, "record") as record,
+    ):
+        scheduler.update_from_output(out0, _make_model_runner_output(out0))
+        scheduler.update_from_output(out1, _make_model_runner_output(out1))
+    hit_sources.assert_called_once()
+    assert hit_sources.call_args.args[0] is importer
+    assert record.call_args.kwargs["hits_by_source"] is sources
+
+
+def test_async_load_registers_immediately_without_deferral():
+    scheduler, importer, _, registered = _setup_async_load(defer_async_loads=False)
+    registered.assert_called_once()
+
+
+def test_abort_deferred_async_load_frees_after_inflight_step():
+    """No load was posted for a deferred request, so nothing will report it
+    done; its blocks are freed once the step that may still read them is
+    processed.
+    """
+    scheduler, importer, outputs, _ = _setup_async_load(defer_async_loads=True)
+    pool = scheduler.kv_cache_manager.block_pool
+    num_free_held = pool.get_num_free_blocks()
+
+    scheduler.finish_requests(importer.request_id, RequestStatus.FINISHED_ABORTED)
+    assert importer.request_id not in scheduler.requests
+    assert scheduler.deferred_frees
+    assert pool.get_num_free_blocks() == num_free_held
+
+    for output in outputs:
+        scheduler.update_from_output(output, _make_model_runner_output(output))
+    assert not scheduler.deferred_frees
+    assert pool.get_num_free_blocks() > num_free_held
diff --git a/tests/v1/core/test_kv_cache_planning.py b/tests/v1/core/test_kv_cache_planning.py
index 5d2a192920e7..43221779a506 100644
--- a/tests/v1/core/test_kv_cache_planning.py
+++ b/tests/v1/core/test_kv_cache_planning.py
@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import copy
 import mmap
 import subprocess
 import sys
@@ -12,7 +13,6 @@
 
 import vllm.v1.core.kv_cache_planning as kv_cache_planning
 import vllm.v1.core.kv_cache_utils as kv_cache_utils
-import vllm.v1.hisparse.runtime as hisparse_runtime_module
 from tests.v1.core.utils import (
     get_kv_cache_configs,
     new_kv_cache_spec,
@@ -20,6 +20,7 @@
     new_mla_spec,
     new_sliding_window_spec,
 )
+from tests.v1.spec_decode.utils import fake_speculative_config
 from vllm.config import (
     CacheConfig,
     ModelConfig,
@@ -27,6 +28,7 @@
     VllmConfig,
 )
 from vllm.config.attention import HiSparseConfig
+from vllm.platforms import current_platform
 from vllm.utils.mem_constants import GiB_bytes
 from vllm.v1.core.kv_cache_planning import (
     check_enough_kv_cache_memory,
@@ -70,7 +72,9 @@ 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
+        current_platform,
+        "get_hisparse_backend_cls",
+        lambda: "vllm.v1.hisparse.backend.CUDAHiSparseBackend",
     )
     specs = {
         "model.layers.0.self_attn": MLAAttentionSpec(
@@ -190,7 +194,9 @@ def test_hisparse_pool_must_fit_max_model_len(
     plus the null block. Auto-fit (-1) must pick a length that passes the same
     check."""
     monkeypatch.setattr(
-        hisparse_runtime_module.current_platform, "is_cuda_alike", lambda: True
+        current_platform,
+        "get_hisparse_backend_cls",
+        lambda: "vllm.v1.hisparse.backend.CUDAHiSparseBackend",
     )
     monkeypatch.setattr(
         kv_cache_planning, "get_hisparse_host_pool_bytes", lambda _: 2**30
@@ -1983,12 +1989,13 @@ def test_check_enough_kv_cache_memory_reserves_null_block():
 def _spec_decode_kv_cache_groups(worker_specs, use_eagle_block_drop=True):
     config = VllmConfig(model_config=ModelConfig(max_model_len=64))
     config.cache_config.kv_cache_layout = "BLHNC"
-    config.speculative_config = SimpleNamespace(
+    config.speculative_config = fake_speculative_config(
         method="dflash",
         use_eagle=lambda: True,
         use_eagle_block_drop=lambda: use_eagle_block_drop,
         use_multi_module_mtp=lambda: False,
         parallel_drafting=False,
+        draft_model_config=SimpleNamespace(architectures=[]),
     )
     configs = get_kv_cache_configs(config, worker_specs, [2**30] * len(worker_specs))
     return generate_scheduler_kv_cache_config(configs).kv_cache_groups
@@ -2172,3 +2179,88 @@ def test_kv_cache_groups_tp_replicas(specs, tp_size, dcp_size, expected):
     """Replicas are the gcd of each layer's tp_size // max_tp_shards."""
     groups = [KVCacheGroupSpec([f"l.{i}"], spec) for i, spec in enumerate(specs)]
     assert kv_cache_groups_tp_replicas(groups, tp_size, dcp_size) == expected
+
+
+def test_shrink_kv_cache_configs_keeps_tensors_and_checks_fit():
+    """Shrinking to the measured block count lowers `num_blocks` and the
+    tensors' sizes; their placements keep describing the reserved capacity.
+    A count too small for `max_model_len` is rejected."""
+    spec = new_kv_cache_spec()
+    vllm_config = VllmConfig(
+        model_config=ModelConfig(max_model_len=spec.block_size * 4)
+    )
+    vllm_config.cache_config.kv_cache_layout = "LBNHC"
+    groups = [KVCacheGroupSpec(["layer.0", "layer.1"], spec)]
+    config = kv_cache_planning.get_kv_cache_config_from_groups(
+        vllm_config, groups, available_memory=2 * spec.page_size_bytes * 16
+    )
+    scheduler_config = generate_scheduler_kv_cache_config([config])
+    assert config.num_blocks == 16
+    tensors_before = copy.deepcopy(config.kv_cache_tensors)
+
+    kv_cache_planning.shrink_kv_cache_configs(
+        vllm_config, [config, scheduler_config], 10
+    )
+    assert config.num_blocks == 10
+    assert scheduler_config.num_blocks == 10
+    assert config.kv_cache_tensors == [
+        replace(tensor, size=tensor.size * 10 // 16) for tensor in tensors_before
+    ]
+
+    # 4 blocks are needed for max_model_len plus the null block.
+    with pytest.raises(ValueError, match="max seq len"):
+        kv_cache_planning.shrink_kv_cache_configs(vllm_config, [config], 4)
+    kv_cache_planning.shrink_kv_cache_configs(vllm_config, [config], 5)
+    assert config.num_blocks == 5
+
+
+def test_extensible_kv_cache_granule_alignment_under_connector():
+    """With a KV connector, capacity and the final measured count are rounded
+    down so every layer segment spans whole commit granules (one RDMA
+    registration each); without one the measured count is kept as is."""
+    from vllm.v1.worker.worker_base import CompilationTimes
+
+    spec = new_kv_cache_spec()  # 1 head x 128 x fp16 x 2 x block 16 = 8 KiB
+    granule = 2 << 20
+    blocks_per_granule = granule // spec.page_size_bytes
+    vllm_config = VllmConfig(
+        model_config=ModelConfig(max_model_len=spec.block_size * 4)
+    )
+    vllm_config.cache_config.kv_cache_layout = "LBNHC"
+    groups = [KVCacheGroupSpec(["layer.0", "layer.1"], spec)]
+    capacity = 3 * blocks_per_granule + 7
+    config = kv_cache_planning.get_kv_cache_config_from_groups(
+        vllm_config, groups, available_memory=2 * spec.page_size_bytes * capacity
+    )
+    scheduler_config = generate_scheduler_kv_cache_config([config])
+    assert config.num_blocks == capacity
+
+    kv_cache_planning.align_extensible_kv_cache_capacity(
+        vllm_config, [config], scheduler_config, granule
+    )
+    aligned = 3 * blocks_per_granule
+    assert config.num_blocks == aligned
+    assert scheduler_config.num_blocks == aligned
+    assert vllm_config.cache_config.num_gpu_blocks == aligned
+    # The placements describe the smaller reservation to be allocated: the
+    # one two-layer tensor places its layers `aligned` blocks apart.
+    (tensor,) = config.kv_cache_tensors
+    assert tensor.layer_stride == aligned * spec.page_size_bytes
+    assert tensor.size == 2 * aligned * spec.page_size_bytes
+
+    measured = [CompilationTimes(0.0, 0.0, num_kv_blocks=2 * blocks_per_granule + 5)]
+    assert (
+        kv_cache_planning.finalize_extensible_kv_cache(
+            vllm_config, [config], scheduler_config, measured, commit_granule=granule
+        )
+        == 2 * blocks_per_granule
+    )
+    assert config.num_blocks == 2 * blocks_per_granule
+
+    measured = [CompilationTimes(0.0, 0.0, num_kv_blocks=blocks_per_granule + 5)]
+    assert (
+        kv_cache_planning.finalize_extensible_kv_cache(
+            vllm_config, [config], scheduler_config, measured
+        )
+        == blocks_per_granule + 5
+    )
diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py
index 917ed57c0ae7..16dd5cf57336 100644
--- a/tests/v1/core/test_kv_cache_utils.py
+++ b/tests/v1/core/test_kv_cache_utils.py
@@ -13,7 +13,6 @@
 
 import vllm.v1.core.kv_cache_planning as kv_cache_planning
 import vllm.v1.core.kv_cache_utils as kv_cache_utils
-import vllm.v1.hisparse.runtime as hisparse_runtime_module
 from tests.v1.core.utils import (
     new_kv_cache_spec,
     new_mamba_spec,
@@ -35,6 +34,7 @@
     MultiModalKwargsItem,
     PlaceholderRange,
 )
+from vllm.platforms import current_platform
 from vllm.sampling_params import SamplingParams
 from vllm.utils.hashing import sha256, sha256_cbor, xxhash, xxhash_cbor
 from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager
@@ -138,7 +138,9 @@ def test_hisparse_steady_state_concurrency_excludes_spilled_resident_pages():
 def test_hisparse_host_pool_must_fit_max_model_len(monkeypatch, extra_blocks, ok):
     """Every page needs a host block, so one max_model_len request must fit."""
     monkeypatch.setattr(
-        hisparse_runtime_module.current_platform, "is_cuda_alike", lambda: True
+        current_platform,
+        "get_hisparse_backend_cls",
+        lambda: "vllm.v1.hisparse.backend.CUDAHiSparseBackend",
     )
     specs = {
         "model.layers.0.self_attn": MLAAttentionSpec(
@@ -181,9 +183,6 @@ def test_hisparse_derived_specs_do_not_affect_kv_cache_layout(monkeypatch):
     backing; KV cache layout resolution must skip them."""
     from vllm.v1.attention.backends.utils import resolve_kv_cache_layout
 
-    monkeypatch.setattr(
-        hisparse_runtime_module.current_platform, "is_cuda_alike", lambda: True
-    )
     monkeypatch.delenv("VLLM_KV_CACHE_LAYOUT", raising=False)
     specs: dict[str, KVCacheSpec] = {
         "model.layers.0.self_attn": MLAAttentionSpec(
@@ -1044,6 +1043,64 @@ def test_request_block_hasher(hash_fn):
     )
 
 
+@pytest.mark.parametrize("hash_fn", [sha256, sha256_cbor])
+def test_request_continues_from_frontend_block_hashes(hash_fn):
+    """Frontend-computed prompt block hashes are used as is, and blocks that
+    outputs complete chain onto them."""
+    frontend_hashes = [BlockHash(bytes([i]) * 32) for i in range(1, 4)]
+    sampling_params = SamplingParams(max_tokens=17)
+    sampling_params.update_from_generation_config({}, eos_token_id=100)
+    request = Request(
+        request_id="0",
+        prompt_token_ids=list(range(10)),
+        sampling_params=sampling_params,
+        pooling_params=None,
+        block_hasher=get_request_block_hasher(3, hash_fn),
+        block_hashes=list(frontend_hashes),
+    )
+    assert request.block_hashes == frontend_hashes
+
+    request.append_output_token_ids([10, 11])
+    assert request.block_hashes == [
+        *frontend_hashes,
+        hash_fn((frontend_hashes[-1], (9, 10, 11), None)),
+    ]
+
+
+def test_frontend_block_hashes_count_check():
+    from vllm.v1.engine import EngineCoreRequest, PrefixHashConfig
+
+    def frontend_hashes(config, num_hashes):
+        request = EngineCoreRequest(
+            request_id="0",
+            prompt_token_ids=list(range(10)),
+            mm_features=None,
+            sampling_params=SamplingParams(),
+            pooling_params=None,
+            arrival_time=0.0,
+            lora_request=None,
+            cache_salt=None,
+            data_parallel_rank=None,
+            block_hashes=[bytes([i]) * 32 for i in range(num_hashes)],
+        )
+        return kv_cache_utils.frontend_block_hashes(request, config)
+
+    config = PrefixHashConfig(
+        hash_algo="sha256", hash_block_size=3, none_hash=kv_cache_utils.NONE_HASH
+    )
+    # 10 prompt tokens make 3 full blocks of 3.
+    assert frontend_hashes(config, 3) == [bytes([i]) * 32 for i in range(3)]
+    assert frontend_hashes(config, 2) is None
+    assert frontend_hashes(config, 4) is None
+    # The engine doesn't hash blocks.
+    assert frontend_hashes(None, 3) is None
+    # With 2 lookahead tokens, the last full block isn't hashed until the 2
+    # tokens after it exist: (10 - 2) // 3 = 2.
+    config.num_lookahead_tokens = 2
+    assert frontend_hashes(config, 2) == [bytes([i]) * 32 for i in range(2)]
+    assert frontend_hashes(config, 3) is None
+
+
 @pytest.mark.parametrize("hash_fn", [sha256, sha256_cbor])
 @pytest.mark.parametrize("lookahead", [False, True])
 def test_request_block_hasher_incremental_append_with_multiple_mm_features(
@@ -1376,6 +1433,31 @@ def test_dcp_replicated_kpool_tail_keeps_block_interior_layout(monkeypatch):
     assert layout == KVCacheLayout.LBHNC
 
 
+def test_pipeline_stages_resolve_a_layout_every_stage_supports(monkeypatch):
+    """PP stages hold different layers, so their workers can report different
+    supported layouts (e.g. only one stage has FlashInfer layers on SM100)."""
+    from vllm.v1.attention.backends.utils import (
+        get_supported_kv_cache_layouts,
+        resolve_kv_cache_layout,
+    )
+
+    monkeypatch.delenv("VLLM_KV_CACHE_LAYOUT", raising=False)
+    config = VllmConfig(model_config=ModelConfig(max_model_len=1024))
+    config.cache_config.kv_cache_layout = None
+    any_layout = [layout.name for layout in get_supported_kv_cache_layouts([])]
+    layout = resolve_kv_cache_layout(config, [["LBHNC", "BLHNC"], any_layout])
+    assert layout == KVCacheLayout.LBHNC
+
+    # A stage that declares nothing must not outvote a declared preference.
+    config.cache_config.kv_cache_layout = None
+    layout = resolve_kv_cache_layout(config, [["BLHNC", "LBNHC"], any_layout])
+    assert layout == KVCacheLayout.BLHNC
+
+    config.cache_config.kv_cache_layout = None
+    with pytest.raises(ValueError, match="No KV cache layout satisfies"):
+        resolve_kv_cache_layout(config, [["LBHNC"], ["LBNHC"]])
+
+
 @pytest.mark.parametrize("use_mla", [False, True])
 def test_full_attention_merge_preserves_replicated_cache_geometry(use_mla):
     config = VllmConfig(model_config=ModelConfig(max_model_len=1024))
@@ -2004,6 +2086,26 @@ def test_merge_mla_spec():
         kv_cache_specs[0].merge(kv_cache_specs)
 
 
+@pytest.mark.skip_global_cleanup
+def test_merge_sliding_window_mla_spec_preserves_quantization_mode():
+    specs = [
+        SlidingWindowMLASpec(
+            block_size=16,
+            num_kv_heads=1,
+            head_size=576,
+            dtype=torch.uint8,
+            cache_dtype_str="fp8_ds_mla",
+            kv_quant_mode=KVQuantMode.FP8_PER_TENSOR,
+            sliding_window=128,
+        )
+        for _ in range(2)
+    ]
+
+    merged = SlidingWindowMLASpec.merge(specs)
+
+    assert merged.kv_quant_mode == KVQuantMode.FP8_PER_TENSOR
+
+
 @pytest.mark.parametrize("hash_fn", [sha256, sha256_cbor])
 def test_request_block_hasher_with_prompt_embeds(hash_fn: Callable[[Any], bytes]):
     block_size = 3
diff --git a/tests/v1/core/test_prefix_caching.py b/tests/v1/core/test_prefix_caching.py
index e8425ae2d400..f382ac0f1fc7 100644
--- a/tests/v1/core/test_prefix_caching.py
+++ b/tests/v1/core/test_prefix_caching.py
@@ -14,6 +14,7 @@
 
 import vllm.v1.core.kv_cache_manager as kv_cache_manager
 import vllm.v1.core.kv_cache_utils as kv_cache_utils
+from vllm.config import VllmConfig
 from vllm.distributed.kv_events import (
     MEDIUM_CPU,
     MEDIUM_GPU,
@@ -21,6 +22,7 @@
     BlockRemoved,
     BlockStored,
 )
+from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorRole
 from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.connector import (
     HiSparseConnector,
     HiSparseConnectorScheduler,
@@ -130,11 +132,14 @@ def make_request(
     )
 
 
-def make_kv_cache_manager(kv_cache_config: KVCacheConfig, **kwargs) -> KVCacheManager:
+def make_kv_cache_manager(
+    kv_cache_config: KVCacheConfig, vllm_config: VllmConfig | None = None, **kwargs
+) -> KVCacheManager:
     """Build a ``KVCacheManager``, deriving ``scheduler_block_size`` from the
     config (LCM of group block sizes) unless explicitly provided. This mirrors
     ``resolve_kv_cache_block_sizes`` for the non-context-parallel case used by
-    these tests, so callers don't have to pass it at every site."""
+    these tests, so callers don't have to pass it at every site. With
+    ``vllm_config``, a HiSparseConnector builds the HiSparse coordinator."""
     kwargs.setdefault(
         "scheduler_block_size",
         lcm(*(g.kv_cache_spec.block_size for g in kv_cache_config.kv_cache_groups)),
@@ -146,7 +151,12 @@ def make_kv_cache_manager(kv_cache_config: KVCacheConfig, **kwargs) -> KVCacheMa
         )
     manager = KVCacheManager(kv_cache_config, **kwargs)
     if kv_cache_config.hisparse_host_num_blocks is not None:
-        get_hisparse_coordinator(manager)
+        if vllm_config is not None:
+            HiSparseConnector(
+                vllm_config, KVConnectorRole.SCHEDULER, kv_cache_config
+            ).bind_kv_cache_manager(manager)
+        else:
+            get_hisparse_coordinator(manager)
     return manager
 
 
@@ -242,11 +252,13 @@ def make_hisparse_kv_cache_manager(
     max_in_flight_tokens: int | None = None,
     enable_caching: bool = False,
     num_prefill_lookahead: int = 0,
+    vllm_config: VllmConfig | None = None,
     **config_kwargs,
 ) -> KVCacheManager:
     config = make_hisparse_kv_cache_config(num_blocks, host_num_blocks, **config_kwargs)
     return make_kv_cache_manager(
         config,
+        vllm_config,
         max_model_len=max_model_len,
         max_in_flight_tokens=max_in_flight_tokens,
         enable_caching=enable_caching,
@@ -921,15 +933,11 @@ def test_hisparse_inflight_host_import_reserves_remaining_gpu_pages():
     assert required == 4
 
 
-@pytest.mark.parametrize("host_num_blocks,admitted", [(5, False), (6, True)])
-def test_hisparse_async_admission_requires_full_prompt_host_pages(
-    host_num_blocks, admitted, tmp_path
-):
-    """An async load is admitted only if its whole prompt's host pages fit.
-
-    The in-flight prefill's host page is not free, so the waiting request's
-    four prompt pages need six host blocks including the null block.
-    """
+def _make_hisparse_scheduler(
+    tmp_path, host_num_blocks: int, matched_tokens: int, **kwargs
+) -> tuple[Scheduler, KVCacheManager]:
+    """A scheduler over a HiSparse manager whose async mock connector matches
+    ``matched_tokens``."""
     from .utils import create_scheduler, mock_kv
 
     (tmp_path / "config.json").write_text(
@@ -939,73 +947,121 @@ def test_hisparse_async_admission_requires_full_prompt_host_pages(
         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),
+        use_kv_connector=mock_kv(matched_tokens=matched_tokens, is_async=True),
+        **kwargs,
+    )
+    manager = make_hisparse_kv_cache_manager(
+        32, host_num_blocks, vllm_config=scheduler.vllm_config
     )
-    manager = make_hisparse_kv_cache_manager(32, host_num_blocks)
     scheduler.kv_cache_manager = manager
     scheduler.kv_cache_config = manager.kv_cache_config
     assert scheduler.scheduler_reserve_full_isl
-    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
+    return scheduler, manager
+
 
+def _add_hisparse_request(
+    scheduler: Scheduler, request_id: str, num_pages: int, offset: int = 0
+) -> Request:
     request = make_request(
-        "waiting", list(range(4 * HISPARSE_BLOCK_SIZE)), HISPARSE_BLOCK_SIZE, sha256
+        request_id,
+        list(range(offset, offset + num_pages * HISPARSE_BLOCK_SIZE)),
+        HISPARSE_BLOCK_SIZE,
+        sha256,
     )
     scheduler.add_request(request)
     scheduler.schedule()
+    return request
+
+
+def test_hisparse_async_admission_fits_beside_inflight_import(tmp_path):
+    """The in-flight import's three prompt pages stay reserved, so the waiting
+    request's four prompt pages fit in eight host blocks including the null
+    block."""
+    scheduler, manager = _make_hisparse_scheduler(tmp_path, 8, HISPARSE_BLOCK_SIZE)
+    inflight = _add_hisparse_request(scheduler, "inflight", 3)
+    assert inflight in scheduler._inflight_prefills
+
+    request = _add_hisparse_request(scheduler, "waiting", 4)
+    assert request in scheduler._inflight_prefills
+    assert manager.get_blocks(request.request_id).blocks[0]
 
-    assert (request in scheduler._inflight_prefills) == admitted
-    assert bool(manager.get_blocks(request.request_id).blocks[0]) == admitted
 
+@pytest.mark.parametrize("release,second_pages", [("free", 6), ("resume", 4)])
+def test_hisparse_import_reservation_released(release, second_pages, tmp_path):
+    """An import stops reserving host pages once it allocates them or is freed.
 
-def test_hisparse_admitted_async_loads_can_finish_with_nothing_running(tmp_path):
+    Seven usable host blocks: the first import holds one page and reserves two.
+    Once it is released, the second import fits only if nothing is still
+    reserved for the first.
+    """
+    scheduler, manager = _make_hisparse_scheduler(tmp_path, 8, HISPARSE_BLOCK_SIZE)
+    first = _add_hisparse_request(scheduler, "first", 3)
+    assert first in scheduler._inflight_prefills
+    if release == "free":
+        # What abort and preemption run; the mock connector cannot finish a
+        # request over several cache groups.
+        manager.free(first)
+    else:
+        scheduler.finished_recving_kv_req_ids.add(first.request_id)
+        assert scheduler.schedule().num_scheduled_tokens[first.request_id] > 0
+
+    second = _add_hisparse_request(scheduler, "second", second_pages, offset=1000)
+    assert second in scheduler._inflight_prefills
+
+
+@pytest.mark.parametrize(
+    "host_num_blocks,first_pages,matched_tokens,num_speculative_tokens",
+    [
+        # Once first imports one page, three are free: enough for second's
+        # prompt but not beside first's remaining three.
+        (5, 4, HISPARSE_BLOCK_SIZE, None),
+        # Both prompts are fully imported; first's draft tokens need a fourth
+        # page on resume.
+        (7, 3, 3 * HISPARSE_BLOCK_SIZE - 1, 4),
+    ],
+)
+def test_hisparse_admitted_async_loads_can_finish_with_nothing_running(
+    tmp_path, host_num_blocks, first_pages, matched_tokens, num_speculative_tokens
+):
     """Admitted async loads must not strand each other on host capacity.
 
     Waiting requests keep their blocks and are never preempted, so two imports
-    that each hold a host page while together needing more than the pool has
-    left would never be scheduled again. HiSparse relies on full-ISL admission:
-    a load is admitted only if its whole prompt's host pages fit, so the newest
-    admitted load can always finish and free host blocks for the others.
+    that each hold host pages while together needing more than the pool has
+    left would never be scheduled again. A load is admitted only if its whole
+    sequence, including the spec decode step, fits beside the pages earlier
+    imports still need.
     """
-    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),
-    )
-    # One null block plus three usable host blocks.
-    manager = make_hisparse_kv_cache_manager(32, 4)
-    scheduler.kv_cache_manager = manager
-    scheduler.kv_cache_config = manager.kv_cache_config
-    assert scheduler.scheduler_reserve_full_isl
-    first = make_request(
-        "first", list(range(3 * HISPARSE_BLOCK_SIZE)), HISPARSE_BLOCK_SIZE, sha256
-    )
-    second = make_request(
-        "second",
-        list(range(1000, 1000 + 3 * HISPARSE_BLOCK_SIZE)),
-        HISPARSE_BLOCK_SIZE,
-        sha256,
-    )
-    for request in (first, second):
-        scheduler.add_request(request)
-        scheduler.schedule()
+    scheduler, _ = _make_hisparse_scheduler(
+        tmp_path,
+        host_num_blocks,
+        matched_tokens,
+        num_speculative_tokens=num_speculative_tokens,
+    )
+    _add_hisparse_request(scheduler, "first", first_pages)
+    _add_hisparse_request(scheduler, "second", 3, offset=1000)
     admitted = {r.request_id for r in scheduler._inflight_prefills}
+    assert admitted == {"first"}
 
     scheduler.finished_recving_kv_req_ids.update(admitted)
-    num_scheduled_tokens = 0
-    for _ in range(4):
-        num_scheduled_tokens += scheduler.schedule().total_num_scheduled_tokens
-    assert num_scheduled_tokens > 0
+    output = scheduler.schedule()
+    assert output.num_scheduled_tokens.get("first", 0) > 0
+
+
+def test_hisparse_import_reservation_does_not_gate_local_admissions(tmp_path):
+    """Like the GPU pool's, the host reservation only gates async loads; a
+    local admission runs and can be preempted, so it cannot strand an import."""
+    # Seven usable host blocks: the import holds one and reserves five, leaving
+    # one, so gating the local request's two pages on the reservation fails.
+    scheduler, _ = _make_hisparse_scheduler(tmp_path, 8, HISPARSE_BLOCK_SIZE)
+    _add_hisparse_request(scheduler, "import", 6)
+    match_external = scheduler.connector.get_num_new_matched_tokens
+    scheduler.connector.get_num_new_matched_tokens = lambda request, num_computed: (
+        (0, False)
+        if request.request_id == "local"
+        else match_external(request, num_computed)
+    )
+    local = _add_hisparse_request(scheduler, "local", 2, offset=1000)
+    assert local in scheduler.running
 
 
 @pytest.mark.parametrize(
diff --git a/tests/v1/core/test_prompt_block_hash_vectors.py b/tests/v1/core/test_prompt_block_hash_vectors.py
new file mode 100644
index 000000000000..10ac488bfbd8
--- /dev/null
+++ b/tests/v1/core/test_prompt_block_hash_vectors.py
@@ -0,0 +1,172 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Pin the engine's prompt block hashes that the Rust frontend reproduces.
+
+The Rust frontend computes prompt block hashes itself (see
+`rust/src/engine-core-client/src/prompt_block_hasher.rs`), and they must be
+byte-identical to the engine's. Its unit tests check them against the vectors
+in `VECTORS`; this test checks that the engine still produces those same
+vectors, so a change to the engine's hashing fails here instead of silently
+splitting the two apart.
+
+If the change is intended, update the Rust hasher to match, then regenerate
+the vectors by running this file directly:
+
+    .venv/bin/python tests/v1/core/test_prompt_block_hash_vectors.py
+"""
+
+import json
+import os
+from collections.abc import Callable
+from pathlib import Path
+from typing import Any
+
+import pytest
+import regex as re
+
+import vllm.v1.core.kv_cache_utils as kv_cache_utils
+from vllm.lora.request import LoRARequest
+from vllm.sampling_params import SamplingParams
+from vllm.utils.hashing import sha256, sha256_cbor
+from vllm.v1.core.kv_cache_utils import get_request_block_hasher, init_none_hash
+from vllm.v1.request import Request
+
+pytestmark = pytest.mark.cpu_test
+
+REPO_ROOT = Path(__file__).resolve().parents[3]
+VECTORS = (
+    REPO_ROOT
+    / "rust"
+    / "src"
+    / "engine-core-client"
+    / "src"
+    / "tests"
+    / "prompt_block_hash_vectors.json"
+)
+
+HASH_FNS: dict[str, Callable[[Any], bytes]] = {
+    "sha256": sha256,
+    "sha256_cbor": sha256_cbor,
+}
+
+# Token ids at every encoding-width boundary of both the pickle and CBOR
+# encodings.
+EDGE_TOKENS = [0, 1, 23, 24, 255, 256, 65535, 65536, 151_645, 0x7FFF_FFFF]
+
+
+def _case_inputs() -> list[dict[str, Any]]:
+    cases: list[dict[str, Any]] = []
+    plain_prompts = [
+        (16, EDGE_TOKENS * 3 + [7, 8]),
+        (1, [0, 300, 70000]),
+        (3, EDGE_TOKENS),
+        (24, list(range(100, 150))),
+        (300, [i * 7919 % 200_003 for i in range(650)]),
+    ]
+    extra_keys = [
+        ("plain", None, None),
+        ("lora", ("adapter-a", "/models/adapter-a"), None),
+        ("salt", None, "tenant-1"),
+        # Non-ASCII LoRA strings, and a salt long enough for pickle's 4-byte
+        # string length.
+        ("lora_salt", ("adapter-é", "/m/é"), "s" * 300),
+    ]
+    for hash_algo in HASH_FNS:
+        for block_size, prompt in plain_prompts:
+            cases.append(
+                {
+                    "name": f"plain_block_size_{block_size}",
+                    "hash_algo": hash_algo,
+                    "block_size": block_size,
+                    "num_lookahead_tokens": 0,
+                    "lora": None,
+                    "cache_salt": None,
+                    "prompt_token_ids": prompt,
+                }
+            )
+        for num_lookahead_tokens in (0, 1, 3):
+            for name, lora, cache_salt in extra_keys:
+                cases.append(
+                    {
+                        "name": f"{name}_lookahead_{num_lookahead_tokens}",
+                        "hash_algo": hash_algo,
+                        "block_size": 5,
+                        "num_lookahead_tokens": num_lookahead_tokens,
+                        "lora": (
+                            None if lora is None else {"name": lora[0], "path": lora[1]}
+                        ),
+                        "cache_salt": cache_salt,
+                        "prompt_token_ids": EDGE_TOKENS * 3 + [7, 8],
+                    }
+                )
+    return cases
+
+
+def _engine_block_hashes(case: dict[str, Any]) -> list[str]:
+    """The prompt block hashes of the engine's `Request` for this case."""
+    lora = case["lora"]
+    sampling_params = SamplingParams(max_tokens=17)
+    sampling_params.update_from_generation_config({}, eos_token_id=100)
+    request = Request(
+        request_id="0",
+        prompt_token_ids=case["prompt_token_ids"],
+        sampling_params=sampling_params,
+        pooling_params=None,
+        lora_request=(
+            None
+            if lora is None
+            else LoRARequest(
+                lora_name=lora["name"], lora_int_id=1, lora_path=lora["path"]
+            )
+        ),
+        cache_salt=case["cache_salt"],
+        block_hasher=get_request_block_hasher(
+            case["block_size"],
+            HASH_FNS[case["hash_algo"]],
+            case["num_lookahead_tokens"],
+        ),
+    )
+    return [block_hash.hex() for block_hash in request.block_hashes]
+
+
+def _engine_vectors() -> dict[str, Any]:
+    """Compute every vector with the engine's default NONE_HASH seed, which is
+    what an engine without PYTHONHASHSEED sends the frontend."""
+    saved_seed = os.environ.pop("PYTHONHASHSEED", None)
+    try:
+        none_hashes: dict[str, str] = {}
+        cases = []
+        for hash_algo, hash_fn in HASH_FNS.items():
+            init_none_hash(hash_fn)
+            none_hashes[hash_algo] = kv_cache_utils.NONE_HASH.hex()
+            for case in _case_inputs():
+                if case["hash_algo"] == hash_algo:
+                    case["block_hashes"] = _engine_block_hashes(case)
+                    cases.append(case)
+    finally:
+        if saved_seed is not None:
+            os.environ["PYTHONHASHSEED"] = saved_seed
+    return {"none_hashes": none_hashes, "cases": cases}
+
+
+def _format(vectors: dict[str, Any]) -> str:
+    text = json.dumps(vectors, indent=2, ensure_ascii=False)
+    # Keep each token id list on one line.
+    text = re.sub(r"\[[\d,\s]+\]", lambda m: json.dumps(json.loads(m[0])), text)
+    return text + "\n"
+
+
+def test_engine_matches_frontend_prompt_block_hash_vectors():
+    assert VECTORS.is_file(), f"missing {VECTORS}"
+    assert json.loads(VECTORS.read_text()) == _engine_vectors(), (
+        "The engine's prompt block hashes no longer match the vectors that the "
+        "Rust frontend's hasher is tested against. Update "
+        "rust/src/engine-core-client/src/prompt_block_hasher.rs to match, then "
+        "regenerate the vectors with "
+        "`.venv/bin/python tests/v1/core/test_prompt_block_hash_vectors.py`."
+    )
+
+
+if __name__ == "__main__":
+    VECTORS.write_text(_format(_engine_vectors()))
+    print(f"wrote {VECTORS}")
diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py
index 5fa8e206dd4e..9445e9723a9d 100644
--- a/tests/v1/core/test_scheduler.py
+++ b/tests/v1/core/test_scheduler.py
@@ -2532,6 +2532,41 @@ def test_kv_cache_release_after_keep_pause_preserves_requests():
     assert resumed.num_scheduled_tokens == {running.request_id: running.num_tokens}
 
 
+@pytest.mark.parametrize("extensible", [True, False])
+def test_pause_that_clears_the_kv_cache_lets_weight_updates_borrow_it(extensible):
+    """Workers learn that their extensible KV cache holds nothing to keep from
+    a pause that cleared it until the resume, and only then."""
+    from types import SimpleNamespace
+
+    core = object.__new__(EngineCore)
+    core.scheduler = create_scheduler()
+    core.mm_receiver_cache = None
+    core.batch_queue = None
+    core.model_executor = Mock(
+        is_sleeping=False,
+        vllm_config=SimpleNamespace(
+            cache_config=SimpleNamespace(enable_extensible_kv_cache=extensible)
+        ),
+    )
+    rpc = core.model_executor.collective_rpc
+
+    def reclaimable_calls() -> list[bool]:
+        return [
+            call.kwargs["args"][0]
+            for call in rpc.call_args_list
+            if call.args == ("set_kv_cache_reclaimable",)
+        ]
+
+    core.pause_scheduler(mode="keep", clear_cache=False)
+    core.resume_scheduler()
+    assert reclaimable_calls() == []
+
+    core.pause_scheduler(mode="abort", clear_cache=True)
+    core.pause_scheduler(mode="abort", clear_cache=True)
+    core.resume_scheduler()
+    assert reclaimable_calls() == ([True, False] if extensible else [])
+
+
 @pytest.mark.parametrize(
     "sleeping_tags",
     [{"weights", "kv_cache"}, {"weights"}, {"kv_cache"}],
@@ -2584,7 +2619,7 @@ def test_draft_slots_budgeted_per_scheduled_request(tmp_path, monkeypatch):
     (tmp_path / "config.json").write_text(
         '{"architectures": ["OPTForCausalLM"], "model_type": "opt"}'
     )
-    scheduler = create_scheduler(
+    base = create_scheduler(
         model=str(tmp_path),
         max_num_seqs=16,
         max_num_batched_tokens=20,
@@ -2592,8 +2627,24 @@ def test_draft_slots_budgeted_per_scheduled_request(tmp_path, monkeypatch):
         parallel_drafting=True,
         skip_tokenizer_init=True,
     )
-    speculative_config = scheduler.vllm_config.speculative_config
+    # P-EAGLE needs a resolvable draft model, so switch the ngram config in place
+    # and redo the config step.
+    vllm_config = base.vllm_config
+    speculative_config = vllm_config.speculative_config
     assert speculative_config is not None
+    object.__setattr__(speculative_config, "method", "eagle3")
+    object.__setattr__(
+        speculative_config,
+        "draft_model_config",
+        SimpleNamespace(hf_config=SimpleNamespace(), architectures=[]),
+    )
+    vllm_config._set_max_num_scheduled_tokens()
+    scheduler = Scheduler(
+        vllm_config=vllm_config,
+        kv_cache_config=base.kv_cache_config,
+        structured_output_manager=base.structured_output_manager,
+        block_size=base.block_size,
+    )
     assert scheduler.max_num_scheduled_tokens == 20
     assert speculative_config.max_num_new_slots_for_drafting == 3
 
@@ -2624,7 +2675,7 @@ def test_block_drafters_cap_requests_instead_of_draft_slots(
     object.__setattr__(
         vllm_config.speculative_config,
         "draft_model_config",
-        SimpleNamespace(hf_config=SimpleNamespace()),
+        SimpleNamespace(hf_config=SimpleNamespace(), architectures=[]),
     )
     vllm_config._set_max_num_scheduled_tokens()
     scheduler = Scheduler(
@@ -8312,93 +8363,6 @@ def test_encoder_input_skipped_when_connector_already_has_the_item(ec_role: str)
     assert not output.scheduled_encoder_inputs.get(req_id)
 
 
-class _RecordingGrammar:
-    """Records what reaches ``validate_tokens`` and raises on a negative id.
-
-    llguidance raises ``OverflowError`` on a negative token id, so this mirrors
-    the real failure: the tests below fail on the unpatched scheduler (the -1
-    padding reaches ``validate_tokens``) and pass once it is stripped first.
-    """
-
-    def __init__(self) -> None:
-        self.seen: list[list[int]] = []
-
-    def validate_tokens(self, tokens: list[int]) -> list[int]:
-        if any(t < 0 for t in tokens):
-            raise OverflowError("llguidance rejects negative token ids")
-        self.seen.append(list(tokens))
-        return tokens
-
-
-def _decode_ready_request(scheduler):
-    """Add one request and advance it out of prefill into decode."""
-    request = create_requests(num_requests=1, num_tokens=1)[0]
-    scheduler.add_request(request)
-    output = scheduler.schedule()
-    scheduler.update_from_output(
-        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=[],
-        ),
-    )
-    return request
-
-
-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, reasoning_ended=True
-    )
-
-    scheduler.update_draft_token_ids(
-        DraftTokenIds([request.request_id], [[10, 11, -1, -1]])
-    )
-
-    assert grammar.seen == [[10, 11]]
-    assert request.spec_token_ids == [10, 11]
-
-
-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 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, reasoning_ended=True
-    )
-
-    scheduler_output = SimpleNamespace(
-        scheduled_spec_decode_tokens={request.request_id: [0, 0, 0, 0]}
-    )
-    scheduler.update_draft_token_ids_in_output(
-        DraftTokenIds([request.request_id], [[10, 11, -1, -1]]),
-        scheduler_output,
-    )
-
-    # The grammar only saw the stripped prefix, never a -1.
-    assert grammar.seen == [[10, 11]]
-    # Two drafts were rejected (4 scheduled - 2 valid), padded back with -1.
-    assert scheduler_output.scheduled_spec_decode_tokens[request.request_id] == [
-        10,
-        11,
-        -1,
-        -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))
@@ -8464,13 +8428,10 @@ def test_diffusion_scheduler_narrows_the_canvas_per_request():
     assert narrow.num_output_placeholders == 4
 
 
-@pytest.mark.parametrize("structured", [False, True])
 @pytest.mark.parametrize("async_scheduling", [False, True])
 @pytest.mark.usefixtures("diffusion_model_runner")
-def test_diffusion_scheduler_trims_full_width_worker_drafts(
-    structured, async_scheduling
-):
-    """Padded worker drafts must be narrowed before scheduling or grammar validation."""
+def test_diffusion_scheduler_trims_full_width_worker_drafts(async_scheduling):
+    """Padded worker drafts must be narrowed before scheduling."""
     scheduler = _diffusion_scheduler(async_scheduling=async_scheduling)
     wide = _diffusion_request("wide", {})
     narrow = _diffusion_request("narrow", {"diffusion_canvas_length": 4})
@@ -8479,12 +8440,7 @@ def test_diffusion_scheduler_trims_full_width_worker_drafts(
     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
+    tokens = [-1] * 8
     drafts = DraftTokenIds(["wide", "narrow"], [tokens.copy(), tokens.copy()])
     if async_scheduling:
         output = scheduler.schedule()
@@ -8498,9 +8454,6 @@ def test_diffusion_scheduler_trims_full_width_worker_drafts(
         "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.parametrize(
diff --git a/tests/v1/cudagraph/test_breakable_cudagraph.py b/tests/v1/cudagraph/test_breakable_cudagraph.py
index 4e19912f4fef..6a5970f26e25 100644
--- a/tests/v1/cudagraph/test_breakable_cudagraph.py
+++ b/tests/v1/cudagraph/test_breakable_cudagraph.py
@@ -6,6 +6,7 @@
 
 import threading
 from contextlib import nullcontext
+from types import SimpleNamespace
 from unittest.mock import MagicMock, patch
 
 import pytest
@@ -68,10 +69,11 @@ def forward_fn(cg_mode):
 
         return forward_fn
 
+    capture_stream = object()
     with (
         patch(
             "vllm.v1.worker.gpu.cudagraph_utils.graph_capture",
-            return_value=nullcontext(),
+            return_value=nullcontext(SimpleNamespace(stream=capture_stream)),
         ),
         patch(
             "vllm.v1.worker.gpu.cudagraph_utils.is_global_first_rank",
@@ -80,6 +82,7 @@ def forward_fn(cg_mode):
     ):
         manager.capture(create_forward_fn)
 
+    assert manager.capture_stream is capture_stream
     assert [warmup for warmup, _ in create_calls] == [True, False]
     assert [mode for _, mode, _ in forward_calls] == [
         CUDAGraphMode.NONE,
diff --git a/tests/v1/cudagraph/test_cudagraph_manager.py b/tests/v1/cudagraph/test_cudagraph_manager.py
index cfe77253e29b..a81586605673 100644
--- a/tests/v1/cudagraph/test_cudagraph_manager.py
+++ b/tests/v1/cudagraph/test_cudagraph_manager.py
@@ -19,6 +19,7 @@
 from vllm.distributed.device_communicators import pynccl_allocator
 from vllm.v1.kv_cache_interface import KVCacheConfig
 from vllm.v1.worker.gpu import cudagraph_utils as gpu_cudagraph_utils
+from vllm.v1.worker.gpu import dp_utils
 from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor
 from vllm.v1.worker.gpu.dp_utils import dispatch_cg_and_sync_dp
 from vllm.v1.worker.gpu.input_batch import InputBuffers
@@ -26,6 +27,9 @@
 from vllm.v1.worker.gpu.spec_decode.target_dependent_ar import (
     cudagraph_utils as spec_cudagraph_utils,
 )
+from vllm.v1.worker.gpu.spec_decode.target_dependent_ar import (
+    speculator as spec_module,
+)
 from vllm.v1.worker.gpu.spec_decode.target_dependent_ar.cudagraph_utils import (
     SpeculatorCudaGraphManager,
 )
@@ -454,6 +458,132 @@ def dispatch(max_query_len):
     assert dispatch(_DECODE_QUERY_LEN + 1).cg_mode == CUDAGraphMode.PIECEWISE
 
 
+@pytest.mark.parametrize(
+    "adaptive,num_tokens,has_prefill,decode_eligible,dp_size,peer_query_len",
+    [
+        (False, 16, False, True, 1, None),
+        (True, 16, False, True, 1, None),
+        (True, 8, False, True, 1, None),
+        (True, 16, True, False, 1, None),
+        pytest.param(True, 16, True, True, 1, None, id="decode-eligible-prefill-tail"),
+        pytest.param(True, 16, False, True, 2, None, id="dp-peer-prefill"),
+        pytest.param(True, 8, False, True, 2, 4, id="dp-peer-longer-decode"),
+    ],
+)
+def test_draft_prefill_matches_shared_target_metadata_shape(
+    monkeypatch,
+    adaptive,
+    num_tokens,
+    has_prefill,
+    decode_eligible,
+    dp_size,
+    peer_query_len,
+):
+    """Draft prefill must reuse the target's padded request layout, even after AV."""
+    target = _make_spec_decode_manager(
+        monkeypatch,
+        decode_query_len=4,
+        capture_sizes=[1, 2, 4, 8, 16, 32],
+        varlen_decode=adaptive,
+    )
+    collective_calls = 0
+
+    def all_reduce(tensor):
+        nonlocal collective_calls
+        collective_calls += 1
+        assert collective_calls == 1, "Draft prefill must reuse the target's DP sync"
+        tensor[:, 1] = torch.tensor(
+            [
+                16,
+                CUDAGraphMode.PIECEWISE.value
+                if peer_query_len is None
+                else CUDAGraphMode.FULL.value,
+                peer_query_len or 0,
+                peer_query_len or -1,
+                0,
+                4,
+                0,
+            ]
+        )
+
+    monkeypatch.setattr(
+        dp_utils, "get_dp_group", lambda: SimpleNamespace(cpu_all_reduce_=all_reduce)
+    )
+    max_query_len = num_tokens // 4
+    target_desc, target_sync = dispatch_cg_and_sync_dp(
+        target,
+        4,
+        num_tokens,
+        max_query_len if decode_eligible else None,
+        dp_size=dp_size,
+        dp_rank=0,
+        max_query_len=None if has_prefill else max_query_len,
+    )
+    hidden_states = torch.zeros(target_desc.num_tokens, 1)
+    speculator = SimpleNamespace(
+        vllm_config=target.vllm_config,
+        device=torch.device("cpu"),
+        num_speculative_steps=3,
+        enable_adaptive_verification=adaptive,
+        max_model_len=4096,
+        max_num_reqs=8,
+        dp_size=dp_size,
+        dp_rank=0,
+        pcp_manager=None,
+        hidden_states=torch.empty_like(hidden_states),
+        last_token_indices=None,
+        current_draft_step=None,
+        input_buffers=None,
+        _copy_request_inputs=lambda *a, **kw: None,
+    )
+    spec_module.TargetDependentARSpeculator.init_cudagraph_manager(
+        speculator, CUDAGraphMode.FULL_AND_PIECEWISE
+    )
+    speculator.prefill_cudagraph_manager._graphs_captured = True
+    batch = SimpleNamespace(
+        num_reqs=4,
+        num_tokens=num_tokens,
+        num_tokens_after_padding=target_desc.num_tokens,
+        num_scheduled_tokens=torch.full((4,), max_query_len).numpy(),
+        seq_lens_cpu_upper_bound=torch.full((8,), 128),
+        idx_mapping=None,
+        has_prefill=has_prefill,
+        decode_graph_eligible=decode_eligible,
+    )
+    monkeypatch.setattr(spec_module, "prepare_prefill_inputs", lambda *a: None)
+    dispatch = spec_module.dispatch_cg_and_sync_dp
+
+    class Dispatched(Exception):
+        pass
+
+    def stop_after_dispatch(*args, **kwargs):
+        raise Dispatched(dispatch(*args, **kwargs)[0])
+
+    monkeypatch.setattr(spec_module, "dispatch_cg_and_sync_dp", stop_after_dispatch)
+    with pytest.raises(Dispatched) as result:
+        spec_module.TargetDependentARSpeculator.propose(
+            speculator,
+            batch,
+            {},
+            {},
+            hidden_states,
+            None,
+            num_sampled=None,
+            num_rejected=None,
+            last_sampled=None,
+            next_prefill_tokens=None,
+            temperature=None,
+            seeds=None,
+            dp_sync_state=target_sync,
+        )
+
+    assert result.value.args[0] == target_desc
+    if dp_size > 1:
+        assert collective_calls == 1
+        assert target_sync is not None
+        assert target_sync.max_query_len == peer_query_len
+
+
 def test_mixed_batch_at_decode_only_token_count_still_gets_a_graph(monkeypatch):
     """A mixed batch must not fall to eager where only decode graphs are staged.
 
diff --git a/tests/v1/cudagraph/test_encoder_cudagraph.py b/tests/v1/cudagraph/test_encoder_cudagraph.py
index 8015e9af333f..db40abad44c3 100644
--- a/tests/v1/cudagraph/test_encoder_cudagraph.py
+++ b/tests/v1/cudagraph/test_encoder_cudagraph.py
@@ -1522,3 +1522,95 @@ def test_user_budgets_negative_raises(self):
 
         with pytest.raises(ValueError, match="must be positive"):
             CompilationConfig(encoder_cudagraph_token_budgets=[-1, 64])
+
+
+# ---------------------------------------------------------------------------
+# FlashInfer cu_seqlens replay padding (no GPU required)
+# ---------------------------------------------------------------------------
+
+
+class TestFlashInferCuSeqlensPadding:
+    """A replay pads its cu_seqlens into the buffer captured at a larger batch
+    size. The padded buffer must equal the cu_seqlens computed for the same
+    inputs at the capture batch size. FlashInfer packs the q/k/o and v offsets
+    into one tensor and splits it at its midpoint, so each half has to stay in
+    its own half of the capture buffer."""
+
+    HIDDEN_SIZE = 16
+
+    def _cu_seqlens(self, flashinfer: bool, seq_lens: list[int], pad_to: int):
+        import numpy as np
+
+        from vllm.model_executor.layers.attention.mm_encoder_attention import (
+            MMEncoderAttention,
+        )
+        from vllm.v1.attention.backends.registry import AttentionBackendEnum
+
+        backend = (
+            AttentionBackendEnum.FLASHINFER
+            if flashinfer
+            else AttentionBackendEnum.FLASH_ATTN
+        )
+        cu_seqlens = np.concatenate([[0], np.cumsum(seq_lens)]).astype(np.int32)
+        cu_seqlens = np.concatenate(
+            [cu_seqlens, np.full(pad_to - len(seq_lens), cu_seqlens[-1], np.int32)]
+        )
+        return MMEncoderAttention.maybe_recompute_cu_seqlens(
+            backend,
+            cu_seqlens,
+            self.HIDDEN_SIZE,
+            1,
+            torch.device("cpu"),
+        )
+
+    @staticmethod
+    def _qwen3_vl_padding_fn(flashinfer: bool):
+        from vllm.model_executor.models.qwen3_vl import (
+            Qwen3VLForConditionalGeneration,
+        )
+        from vllm.v1.attention.backends.registry import AttentionBackendEnum
+
+        backend = (
+            AttentionBackendEnum.FLASHINFER
+            if flashinfer
+            else AttentionBackendEnum.FLASH_ATTN
+        )
+        model = SimpleNamespace(
+            is_multimodal_pruning_enabled=False,
+            get_max_frames_per_video=lambda: 16,
+            visual=SimpleNamespace(attn_backend=backend, out_hidden_size=8),
+        )
+        config = Qwen3VLForConditionalGeneration.get_encoder_cudagraph_config(
+            cast(Any, model)
+        )
+        # Same lookup as EncoderCudaGraphManager replay.
+        return config.padding_logics.get(
+            "cu_seqlens", EncoderCudaGraphManager._copy_padded_buffer
+        )
+
+    # (seq_lens, replay batch size, capture batch size). The FlashInfer batch
+    # buckets are 8, 16, 32 and 64.
+    CASES = [
+        ([49], 2, 16),  # bucket 8 replayed into bucket 16
+        ([49, 196], 2, 16),
+        ([64] * 9, 9, 32),  # bucket 16 replayed into bucket 32
+        ([49] * 3, 3, 8),  # same bucket
+    ]
+
+    @pytest.mark.parametrize("flashinfer", [True, False])
+    @pytest.mark.parametrize(("seq_lens", "replay_batch", "capture_batch"), CASES)
+    def test_replay_padding_matches_capture_size_metadata(
+        self,
+        flashinfer: bool,
+        seq_lens: list[int],
+        replay_batch: int,
+        capture_batch: int,
+    ):
+        capture = self._cu_seqlens(flashinfer, [49] * capture_batch, capture_batch)
+        replay = self._cu_seqlens(flashinfer, seq_lens, replay_batch)
+        want = self._cu_seqlens(flashinfer, seq_lens, capture_batch)
+        assert want.shape == capture.shape
+
+        buf = torch.full_like(capture, -1)
+        self._qwen3_vl_padding_fn(flashinfer)(buf, replay)
+        assert torch.equal(buf, want)
diff --git a/tests/v1/determinism/test_cutlass_batch_invariance.py b/tests/v1/determinism/test_cutlass_batch_invariance.py
index b8f83817065c..a223bbb3105b 100644
--- a/tests/v1/determinism/test_cutlass_batch_invariance.py
+++ b/tests/v1/determinism/test_cutlass_batch_invariance.py
@@ -205,6 +205,33 @@ def test_online_fp8_linear_batch_invariant(
         assert torch.equal(layer(torch.cat([filler, needle]))[0][-1:], baseline)
 
 
+@requires_fp8
+@_NVFP4_REQUIRES_SM100
+@torch.inference_mode()
+def test_flashinfer_fp8_mm_batch_invariant_after_autotune(monkeypatch):
+    """Under "auto", autotune may pick a different bmm_fp8 backend per M bucket."""
+    from vllm.utils import flashinfer as fi
+
+    if not fi.has_flashinfer():
+        pytest.skip("FlashInfer is not installed.")
+    monkeypatch.setattr(envs, "VLLM_BATCH_INVARIANT", True)
+    set_random_seed(0)
+    fp8 = current_platform.fp8_dtype()
+    x = (0.5 * torch.randn(2048, 4096)).to(fp8)
+    b = (0.5 * torch.randn(4096, 4096)).to(fp8).t()
+    scale_a, scale_b = torch.tensor(0.02), torch.tensor(0.03)
+
+    def run(m):
+        return fi.flashinfer_scaled_fp8_mm(x[:m], b, scale_a, scale_b, torch.bfloat16)
+
+    buckets = fi.flashinfer_get_hybrid_num_tokens_buckets(2048)
+    with fi.autotune(tuning_buckets=buckets, round_up=True):
+        run(2048)
+    ref = run(2048)
+    for m in (1, 7, 64, 128, 512):
+        assert torch.equal(run(m), ref[:m])
+
+
 @requires_fp8
 @torch.inference_mode()
 def test_block_fp8_batch_invariant_across_tuned_block_size_k(monkeypatch):
diff --git a/tests/v1/distributed/test_dbo.py b/tests/v1/distributed/test_dbo.py
index 632bdfff6755..6f2dfc6a2658 100644
--- a/tests/v1/distributed/test_dbo.py
+++ b/tests/v1/distributed/test_dbo.py
@@ -58,11 +58,51 @@
 )
 def test_dbo_dp_ep_gsm8k(all2all_backend: str, num_gpus_available):
     """Test DBO with DP+EP using GSM8K evaluation."""
+    _run_dbo_gsm8k(all2all_backend, num_gpus_available)
+
+
+# Combinations the V2 runner supports only (DBO picks V1 unless told otherwise).
+# ngram verifies multi-token queries without a draft model; DCP shards the KV
+# cache across each DP rank's TP group.
+V2_FEATURES = {
+    "spec_decode": (
+        [
+            "--speculative-config",
+            '{"method": "ngram", "num_speculative_tokens": 2, '
+            '"prompt_lookup_max": 4, "prompt_lookup_min": 2}',
+        ],
+        1,
+    ),
+    "dcp": (["--tensor-parallel-size", "2", "--decode-context-parallel-size", "2"], 2),
+}
+
+
+@pytest.mark.skipif(not has_deep_ep(), reason="These tests require deep_ep to run")
+@pytest.mark.parametrize("feature", list(V2_FEATURES))
+def test_dbo_v2_features_gsm8k(feature: str, num_gpus_available):
+    """Test V2-runner DBO combined with spec decode or DCP."""
+    extra_args, tp_size = V2_FEATURES[feature]
+    _run_dbo_gsm8k(
+        "deepep_low_latency",
+        num_gpus_available,
+        extra_args=extra_args,
+        tp_size=tp_size,
+        env_dict={"VLLM_USE_V2_MODEL_RUNNER": "1"},
+    )
+
+
+def _run_dbo_gsm8k(
+    all2all_backend: str,
+    num_gpus_available: int,
+    extra_args: list[str] | None = None,
+    tp_size: int = 1,
+    env_dict: dict[str, str] | None = None,
+):
     lm_eval = pytest.importorskip("lm_eval")
-    required_gpus = DP_SIZE
+    required_gpus = DP_SIZE * tp_size
 
     if num_gpus_available < required_gpus:
-        pytest.skip(f"Need at least {required_gpus} GPUs (DP={DP_SIZE})")
+        pytest.skip(f"Need at least {required_gpus} GPUs (DP={DP_SIZE}, TP={tp_size})")
 
     # Server arguments for DBO + DP + EP
     server_args = [
@@ -83,11 +123,13 @@ def test_dbo_dp_ep_gsm8k(all2all_backend: str, num_gpus_available):
         "32",
         "--all2all-backend",
         all2all_backend,
+        *(extra_args or []),
     ]
 
     with RemoteOpenAIServer(
         MODEL_NAME,
         server_args,
+        env_dict=env_dict,
         max_wait_seconds=600,  # Allow time for model loading with DP+EP
     ) as remote_server:
         # Use host and port directly from RemoteOpenAIServer
diff --git a/tests/v1/distributed/test_dense_dp_world_size.py b/tests/v1/distributed/test_dense_dp_world_size.py
index 1fc37936463a..dd17b0234a23 100644
--- a/tests/v1/distributed/test_dense_dp_world_size.py
+++ b/tests/v1/distributed/test_dense_dp_world_size.py
@@ -1,17 +1,80 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import os
+from typing import TYPE_CHECKING
 
+import pytest
 import requests
 
-from tests.utils import RemoteOpenAIServer
+from tests.utils import VLLM_PATH, RemoteOpenAIServer
 from vllm.platforms import current_platform
 
+if TYPE_CHECKING:
+    import torch
+
+    from vllm.config import ParallelConfig
+
 MODEL_NAME = "Qwen/Qwen3-0.6B"
 DP_SIZE = int(os.getenv("DP_SIZE", "2"))
 TP_SIZE = int(os.getenv("TP_SIZE", "1"))
 
 
+class NumaBindingWorkerExtension:
+    parallel_config: "ParallelConfig"
+    local_rank: int
+    device: "torch.device"
+
+    def check_numa_binding(self):
+        import subprocess
+        from pathlib import Path
+
+        import psutil
+        import torch.distributed as dist
+
+        from vllm.distributed import get_pcp_group, get_world_group
+        from vllm.utils.cpu_resource_utils import parse_id_list
+
+        parallel = self.parallel_config
+        assert parallel.numa_bind and parallel._numa_bind_auto
+        assert parallel.numa_bind_nodes is not None
+        gpu_index = self.local_rank
+        node = current_platform.get_device_numa_node(self.device.index)
+        assert parallel.numa_bind_nodes[gpu_index] == node
+        node_cpus = parse_id_list(
+            Path(f"/sys/devices/system/node/node{node}/cpulist").read_text().strip()
+        )
+        cpus = set(psutil.Process().cpu_affinity())
+        assert cpus and cpus <= set(node_cpus)
+        result = subprocess.run(
+            ["numactl", "--show"],
+            check=True,
+            capture_output=True,
+            text=True,
+            timeout=10,
+        )
+        policy = dict(line.split(":", 1) for line in result.stdout.splitlines())
+        assert policy["policy"].strip() == "preferred"
+        assert policy["preferred node"].strip() == str(node)
+        parent = psutil.Process(os.getppid())
+        parent_cpus = set(parent.cpu_affinity())
+        assert cpus <= parent_cpus
+        state = {
+            "pid": os.getpid(),
+            "gpu_index": gpu_index,
+            "gpu_uuid": current_platform.get_device_uuid(self.device.index),
+            "dp_rank": parallel.data_parallel_rank,
+            "pcp_rank": get_pcp_group().rank_in_group,
+            "node": node,
+            "cpus": sorted(cpus),
+            "engine_core_pid": parent.pid,
+            "engine_core_cpus": sorted(parent_cpus),
+        }
+        world = get_world_group()
+        states = [None] * world.world_size
+        dist.all_gather_object(states, state, group=world.cpu_group)
+        return states
+
+
 def test_dense_dp_world_size():
     server_args = [
         "--dtype",
@@ -50,3 +113,79 @@ def test_dense_dp_world_size():
         )
         response.raise_for_status()
         assert response.json() == {"world_size": TP_SIZE}
+
+
+@pytest.mark.slow_test
+def test_pcp_dp_with_auto_numa():
+    """Automatic binding must cover both PCP workers in both local DP shards."""
+    from shutil import which
+
+    from vllm.utils import numa_utils
+
+    if not current_platform.is_cuda() or current_platform.device_count() < 4:
+        pytest.skip("PCP+DP NUMA smoke test requires four NVIDIA GPUs")
+    if which("numactl") is None or not numa_utils._is_auto_numa_available(
+        verbose=False, require_unbound_process=True
+    ):
+        pytest.skip("Host cannot engage automatic NUMA binding")
+
+    server_args = [
+        "--dtype",
+        "bfloat16",
+        "--max-model-len",
+        "2048",
+        "--max-num-seqs",
+        "8",
+        "--max-num-batched-tokens",
+        "2048",
+        "--kv-cache-memory-bytes",
+        str(256 << 20),
+        "--enforce-eager",
+        "--distributed-executor-backend",
+        "mp",
+        "--prefill-context-parallel-size",
+        "2",
+        "--data-parallel-size",
+        "2",
+        "--data-parallel-size-local",
+        "2",
+        "--worker-extension-cls",
+        "tests.v1.distributed.test_dense_dp_world_size.NumaBindingWorkerExtension",
+    ]
+    with RemoteOpenAIServer(
+        "deepseek-ai/DeepSeek-V2-Lite-Chat",
+        server_args,
+        env_dict={
+            "VLLM_SERVER_DEV_MODE": "1",
+            "VLLM_WORKER_MULTIPROC_METHOD": "spawn",
+            "PYTHONPATH": os.pathsep.join(
+                filter(None, [str(VLLM_PATH), os.environ.get("PYTHONPATH")])
+            ),
+        },
+        max_wait_seconds=1200,
+    ) as server:
+        completion = server.get_client().completions.create(
+            model="deepseek-ai/DeepSeek-V2-Lite-Chat",
+            prompt="The first five positive integers are",
+            temperature=0,
+            max_tokens=16,
+        )
+        assert completion.choices[0].text.strip()
+        assert completion.usage.completion_tokens > 0
+        print(f"PCP+DP completion: {completion.choices[0].text!r}")
+        response = requests.post(
+            server.url_for("collective_rpc"),
+            json={"method": "check_numa_binding", "timeout": 30},
+            timeout=60,
+        )
+        response.raise_for_status()
+        states = response.json()["results"][0]
+        print(f"PCP+DP NUMA worker states: {states}")
+        assert sorted(s["gpu_index"] for s in states) == [0, 1, 2, 3]
+        assert len({s["gpu_uuid"] for s in states}) == 4
+        assert sorted((s["dp_rank"], s["pcp_rank"]) for s in states) == [
+            (0, 0),
+            (0, 1),
+            (1, 0),
+            (1, 1),
+        ]
diff --git a/tests/v1/e2e/general/test_async_scheduling.py b/tests/v1/e2e/general/test_async_scheduling.py
index 28f9943b3b59..3378e0ac3557 100644
--- a/tests/v1/e2e/general/test_async_scheduling.py
+++ b/tests/v1/e2e/general/test_async_scheduling.py
@@ -241,15 +241,15 @@ def run_tests(
     print(f"BASELINE: config=[{baseline_config}], accept_rates={baseline_acceptances}")
 
     failure = None
-    for test_config, test_outputs, test_acceptance_rates in outputs[1:]:
-        for (base_outs, base_logprobs), base_acceptance_rate, (
+    for test_config, test_outputs, test_acceptance_lengths in outputs[1:]:
+        for (base_outs, base_logprobs), base_acceptance_length, (
             test_outs,
             test_logprobs,
-        ), test_acceptance_rate, params in zip(
+        ), test_acceptance_length, params in zip(
             baseline_tests,
             baseline_acceptances or repeat(None),
             test_outputs,
-            test_acceptance_rates or repeat(None),
+            test_acceptance_lengths or repeat(None),
             test_sampling_params,
         ):
             reason = None
@@ -272,8 +272,8 @@ def run_tests(
             if reason is None:
                 try:
                     if (
-                        base_acceptance_rate is not None
-                        and test_acceptance_rate is not None
+                        base_acceptance_length is not None
+                        and test_acceptance_length is not None
                     ):
                         if "spec_mml=None" in test_config:
                             # Preemption causes more variance in acceptance rates
@@ -285,14 +285,14 @@ def run_tests(
                             else:
                                 tolerance = 0.05
                             assert (
-                                test_acceptance_rate > base_acceptance_rate
-                                or test_acceptance_rate
-                                == pytest.approx(base_acceptance_rate, rel=tolerance)
+                                test_acceptance_length > base_acceptance_length
+                                or test_acceptance_length
+                                == pytest.approx(base_acceptance_length, rel=tolerance)
                             )
                         else:
                             # Currently the reported acceptance rate is expected to be
                             # lower when we sometimes skip drafting altogether.
-                            assert test_acceptance_rate > 0.1
+                            assert test_acceptance_length > 1.1
                 except AssertionError as e:
                     reason = "accept  ", e
 
@@ -300,14 +300,14 @@ def run_tests(
                 print(
                     f"\033[32mPASSED\033[0m:           "
                     f"config=[{test_config}], params={params}"
-                    f" accept_rate={test_acceptance_rate}"
+                    f" accept_rate={test_acceptance_length}"
                 )
             else:
                 reason_str, _ = reason
                 print(
                     f"\033[31mFAILED\033[0m({reason_str}): "
                     f"config=[{test_config}], params={params}"
-                    f" accept_rate={test_acceptance_rate}"
+                    f" accept_rate={test_acceptance_length}"
                 )
                 if failure is None:
                     _, failure = reason
@@ -366,7 +366,7 @@ def run_test(
         **cache_arg,
     ) as vllm_model:
         results = []
-        acceptance_rates: list[float] | None = [] if spec_decoding else None
+        acceptance_lengths: list[float] | None = [] if spec_decoding else None
         for override_params in sampling_param_tests:
             metrics_before = vllm_model.llm.get_metrics()
             print(f"----------- RUNNING PARAMS: {override_params}")
@@ -378,10 +378,12 @@ def run_test(
                 )
             )
             metrics_after = vllm_model.llm.get_metrics()
-            if acceptance_rates is not None:
-                acceptance_rate = _get_acceptance_rate(metrics_before, metrics_after)
-                acceptance_rates.append(acceptance_rate)
-                print(f"ACCEPTANCE RATE {acceptance_rate}")
+            if acceptance_lengths is not None:
+                acceptance_length = _get_acceptance_length(
+                    metrics_before, metrics_after
+                )
+                acceptance_lengths.append(acceptance_length)
+                print(f"ACCEPTANCE RATE {acceptance_length}")
 
             if test_preemption:
                 preemptions = _get_count(
@@ -404,7 +406,7 @@ def run_test(
                 )
                 assert _all_logprobs_match(results[0][1], other_test_logprobs)
 
-    return test_config, results, acceptance_rates
+    return test_config, results, acceptance_lengths
 
 
 def _all_logprobs_match(req_a, req_b) -> bool:
@@ -438,10 +440,10 @@ def _logprobs_match(
     )
 
 
-def _get_acceptance_rate(before: list[Metric], after: list[Metric]) -> float:
-    draft = _get_count(before, after, "vllm:spec_decode_num_draft_tokens")
+def _get_acceptance_length(before: list[Metric], after: list[Metric]) -> float:
+    draft = _get_count(before, after, "vllm:spec_decode_num_drafts")
     accept = _get_count(before, after, "vllm:spec_decode_num_accepted_tokens")
-    return accept / draft if draft > 0 else 0.0
+    return 1 + accept / draft if draft > 0 else 1.0
 
 
 def _get_count(before: list[Metric], after: list[Metric], name: str) -> int:
diff --git a/tests/v1/e2e/general/test_async_scheduling_accuracy.py b/tests/v1/e2e/general/test_async_scheduling_accuracy.py
index 0fd208b1ea7c..1133a4dbf10b 100644
--- a/tests/v1/e2e/general/test_async_scheduling_accuracy.py
+++ b/tests/v1/e2e/general/test_async_scheduling_accuracy.py
@@ -292,7 +292,7 @@ def run_tests(
                         f"{context}: acceptance={actual}, baseline={baseline}"
                     )
                 else:
-                    assert actual > 0.1, f"{context}: acceptance={actual}"
+                    assert actual > 1.1, f"{context}: acceptance={actual}"
 
 
 def run_test(
@@ -346,7 +346,7 @@ def run_test(
         **cache_arg,
     ) as vllm_model:
         results = []
-        acceptance_rates: list[float] | None = [] if spec_decoding else None
+        acceptance_lengths: list[float] | None = [] if spec_decoding else None
         for override_params in sampling_param_tests:
             metrics_before = vllm_model.llm.get_metrics()
             print(f"----------- RUNNING PARAMS: {override_params}")
@@ -357,10 +357,12 @@ def run_test(
                 )
             )
             metrics_after = vllm_model.llm.get_metrics()
-            if acceptance_rates is not None:
-                acceptance_rate = _get_acceptance_rate(metrics_before, metrics_after)
-                acceptance_rates.append(acceptance_rate)
-                print(f"ACCEPTANCE RATE {acceptance_rate}")
+            if acceptance_lengths is not None:
+                acceptance_length = _get_acceptance_length(
+                    metrics_before, metrics_after
+                )
+                acceptance_lengths.append(acceptance_length)
+                print(f"ACCEPTANCE RATE {acceptance_length}")
 
             if test_preemption:
                 preemptions = _get_count(
@@ -387,7 +389,7 @@ def run_test(
                 f"{test_config}: structured sampling parameters had no effect"
             )
 
-    return test_config, results, acceptance_rates
+    return test_config, results, acceptance_lengths
 
 
 def _result_signature(requests):
@@ -402,12 +404,11 @@ def _result_signature(requests):
     ]
 
 
-def _get_acceptance_rate(before: list[Metric], after: list[Metric]) -> float:
-    draft = _get_count(before, after, "vllm:spec_decode_num_draft_tokens")
+def _get_acceptance_length(before: list[Metric], after: list[Metric]) -> float:
+    draft = _get_count(before, after, "vllm:spec_decode_num_drafts")
     accept = _get_count(before, after, "vllm:spec_decode_num_accepted_tokens")
     assert draft > 0, "speculative test did not draft any tokens"
-    assert 0 <= accept <= draft, f"invalid acceptance counters: {accept=}, {draft=}"
-    return accept / draft
+    return 1 + accept / draft
 
 
 def _get_count(before: list[Metric], after: list[Metric], name: str) -> int:
diff --git a/tests/v1/e2e/general/test_streaming_input.py b/tests/v1/e2e/general/test_streaming_input.py
index 48c7eb2596db..1d9aa6995de9 100644
--- a/tests/v1/e2e/general/test_streaming_input.py
+++ b/tests/v1/e2e/general/test_streaming_input.py
@@ -26,8 +26,10 @@
 from vllm.utils.torch_utils import set_default_torch_num_threads
 from vllm.v1.engine.async_llm import AsyncLLM
 
-if not current_platform.is_cuda():
-    pytest.skip(reason="V1 currently only supported on CUDA.", allow_module_level=True)
+if not (current_platform.is_cuda() or current_platform.is_xpu()):
+    pytest.skip(
+        reason="V1 currently only supported on CUDA and XPU.", allow_module_level=True
+    )
 
 # Use a small model that doesn't require authentication for fast tests
 MODEL = "facebook/opt-125m"
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 37a4a61d5259..177734672a16 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
@@ -1,7 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-from dataclasses import dataclass
+from dataclasses import dataclass, field
+from typing import Any
 
 import pytest
 
@@ -42,6 +43,7 @@ class ArgsTest:
     parallel_drafting: bool = False
     target_tensor_parallel_size: int = 1
     draft_tensor_parallel_size: int = 1
+    pipeline_parallel_size: int = 1
     max_model_len: int = 2048
     gpu_memory_utilization: float = 0.5
     dataset: str = "test_prompts"
@@ -51,6 +53,7 @@ class ArgsTest:
     enable_prefix_caching: bool | None = None
     chat_template_kwargs: dict | None = None
     reject_full_acceptance: bool = False
+    engine_kwargs: dict[str, Any] = field(default_factory=dict)
 
 
 def get_messages(dataset: str, n: int) -> list[Messages]:
@@ -374,29 +377,8 @@ def assert_draft_model_correctness(args: ArgsTest, vllm_runner):
         dataset=args.dataset, n=args.num_prompts
     )
 
-    with vllm_runner(
-        args.target_model,
-        block_size=None,
-        trust_remote_code=False,
-        enable_chunked_prefill=None,
-        compilation_config=CompilationConfig(),
-        speculative_config={
-            "model": args.draft_model,
-            "method": args.method,
-            "num_speculative_tokens": args.num_speculative_tokens,
-            "max_model_len": args.max_model_len,
-            "enforce_eager": args.enforce_eager,
-            "draft_tensor_parallel_size": args.draft_tensor_parallel_size,
-            "parallel_drafting": args.parallel_drafting,
-        },
-        max_num_seqs=args.max_num_seqs,  # limit cudagraph capture runtime
-        max_model_len=args.max_model_len,
-        gpu_memory_utilization=args.gpu_memory_utilization,
-        tensor_parallel_size=args.target_tensor_parallel_size,
-        enforce_eager=args.enforce_eager,
-        disable_log_stats=False,  # enables get_metrics()
-        enable_prefix_caching=args.enable_prefix_caching,
-    ) as spec_runner:
+    kwargs = vllm_runner_kwargs(args)
+    with vllm_runner(args.target_model, **kwargs) as spec_runner:
         spec_llm = spec_runner.llm
 
         # we don't check the outputs, only check the metrics
@@ -443,3 +425,31 @@ def assert_draft_model_correctness(args: ArgsTest, vllm_runner):
         has_async = spec_llm.llm_engine.vllm_config.scheduler_config.async_scheduling
 
     assert has_async, f"Expected async_scheduling=True for {args.method} spec decode"
+
+
+def vllm_runner_kwargs(args: ArgsTest) -> dict[str, Any]:
+    kwargs = dict(
+        block_size=None,
+        trust_remote_code=False,
+        enable_chunked_prefill=None,
+        compilation_config=CompilationConfig(),
+        speculative_config={
+            "model": args.draft_model,
+            "method": args.method,
+            "num_speculative_tokens": args.num_speculative_tokens,
+            "max_model_len": args.max_model_len,
+            "enforce_eager": args.enforce_eager,
+            "draft_tensor_parallel_size": args.draft_tensor_parallel_size,
+            "parallel_drafting": args.parallel_drafting,
+        },
+        max_num_seqs=args.max_num_seqs,
+        max_model_len=args.max_model_len,
+        gpu_memory_utilization=args.gpu_memory_utilization,
+        tensor_parallel_size=args.target_tensor_parallel_size,
+        pipeline_parallel_size=args.pipeline_parallel_size,
+        enforce_eager=args.enforce_eager,
+        disable_log_stats=False,  # enables get_metrics()
+        enable_prefix_caching=args.enable_prefix_caching,
+    )
+    kwargs.update(args.engine_kwargs)
+    return kwargs
diff --git a/tests/v1/e2e/spec_decode/draft_model/test_structured_output_pp.py b/tests/v1/e2e/spec_decode/draft_model/test_structured_output_pp.py
new file mode 100644
index 000000000000..fc7f72d26550
--- /dev/null
+++ b/tests/v1/e2e/spec_decode/draft_model/test_structured_output_pp.py
@@ -0,0 +1,140 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Structured outputs must remain constrained across pipelined draft handoffs."""
+
+import json
+
+import jsonschema
+import pytest
+
+from tests.utils import multi_gpu_only
+from vllm import SamplingParams
+from vllm.outputs import RequestOutput
+from vllm.sampling_params import StructuredOutputsParams
+from vllm.v1.metrics.reader import Metric
+
+from ..utils import compute_acceptance_len, compute_acceptance_rate
+from .test_draft_model import ArgsTest, vllm_runner_kwargs
+
+NUM_REQUEST_CYCLES = 6
+STRUCTURED_OUTPUT_SCHEMA = {
+    "type": "array",
+    "minItems": 16,
+    "maxItems": 16,
+    "items": {"type": "string", "enum": ["red", "blue"]},
+}
+
+
+def case_to_test(pp_size: int, tp_size: int) -> ArgsTest:
+    """Exercise structured draft handoffs across repeated request cohorts."""
+    return ArgsTest(
+        target_model="Qwen/Qwen3-4B",
+        draft_model="z-lab/Qwen3-4B-DFlash-b16",
+        sampling_config=SamplingParams(temperature=0, max_tokens=256),
+        num_speculative_tokens=15,
+        method="dflash",
+        # reference for AL and AR is TP=1 PP=1
+        expected_acceptance_rate=0.0878 * 0.95,  # ref - 5% tolerance
+        expected_acceptance_len=2.3168 * 0.95,  # ref - 5% tolerance
+        pipeline_parallel_size=pp_size,
+        target_tensor_parallel_size=tp_size,
+        max_num_seqs=8,  # force multiple PP batches
+        gpu_memory_utilization=0.8,
+        engine_kwargs={
+            # The handoff race only exists with deferred (async) sampling.
+            "async_scheduling": True,
+            "structured_outputs_config": {
+                "backend": "xgrammar",
+                "disable_any_whitespace": True,
+            },
+        },
+    )
+
+
+@multi_gpu_only(num_gpus=2)
+@pytest.mark.parametrize(
+    "args",
+    [
+        pytest.param(case_to_test(pp_size=2, tp_size=1), id="pp2-tp1"),
+    ],
+)
+def test_specdec_pp_structured_output(args: ArgsTest, vllm_runner) -> None:
+    # Alternate full-size cohorts to turn over request IDs while PP batches are
+    # still in flight.
+    messages, sampling_params = _create_requests(
+        args.sampling_config,
+        num_cycles=NUM_REQUEST_CYCLES,
+        cohort_size=args.max_num_seqs,
+    )
+    kwargs = vllm_runner_kwargs(args)
+    with vllm_runner(args.target_model, **kwargs) as runner:
+        outputs = runner.llm.chat(
+            messages,
+            sampling_params,
+            chat_template_kwargs={"enable_thinking": False},
+        )
+        metrics = runner.llm.get_metrics()
+    _assert_outputs_are_structured(outputs, sampling_params)
+    _assert_specdec_metrics(
+        metrics,
+        expected_acceptance_rate=args.expected_acceptance_rate,
+        expected_acceptance_len=args.expected_acceptance_len,
+    )
+
+
+def _create_requests(
+    sampling: SamplingParams, num_cycles: int, cohort_size: int
+) -> tuple[list, list[SamplingParams]]:
+    messages = []
+    sampling_params: list[SamplingParams] = []
+    request_index = 0
+    for _ in range(num_cycles):
+        for is_structured in (True, False):
+            for _ in range(cohort_size):
+                messages.append(_make_messages(request_index, is_structured))
+                params = sampling.clone()
+                if is_structured:
+                    params.structured_outputs = StructuredOutputsParams(
+                        json=STRUCTURED_OUTPUT_SCHEMA
+                    )
+                sampling_params.append(params)
+                request_index += 1
+    return messages, sampling_params
+
+
+def _make_messages(index: int, is_structured: bool):
+    task = (
+        "list sixteen animals as JSON."
+        if is_structured
+        else "write a detailed travel story."
+    )
+    messages = [{"role": "user", "content": f"Case {index}: {task}"}]
+    return messages
+
+
+def _assert_specdec_metrics(
+    metrics: list[Metric],
+    expected_acceptance_rate: float,
+    expected_acceptance_len: float,
+) -> None:
+    actual_rate = compute_acceptance_rate(metrics)
+    actual_len = compute_acceptance_len(metrics)
+    print(f"acceptance_rate={actual_rate:.4f}, acceptance_len={actual_len:.4f}")
+    assert actual_rate >= expected_acceptance_rate, (
+        actual_rate,
+        expected_acceptance_rate,
+    )
+    assert actual_len >= expected_acceptance_len, (actual_len, expected_acceptance_len)
+
+
+def _assert_outputs_are_structured(
+    outputs: list[RequestOutput], sampling_params: list[SamplingParams]
+) -> None:
+    assert len(outputs) == len(sampling_params)
+    for output, params in zip(outputs, sampling_params, strict=True):
+        completion = output.outputs[0]
+        assert output.finished and completion.text
+        schema = params.structured_outputs
+        if schema is not None and schema.json is not None:
+            assert completion.finish_reason == "stop", completion.text
+            jsonschema.validate(json.loads(completion.text), schema.json)
diff --git a/tests/v1/e2e/spec_decode/mtp/qwen3_5/test_mtp.py b/tests/v1/e2e/spec_decode/mtp/qwen3_5/test_mtp.py
index 0feaa9efc604..5147097ef290 100644
--- a/tests/v1/e2e/spec_decode/mtp/qwen3_5/test_mtp.py
+++ b/tests/v1/e2e/spec_decode/mtp/qwen3_5/test_mtp.py
@@ -1,20 +1,96 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import json
+from pathlib import Path
+
 import pytest
+import torch
+from safetensors import safe_open
+from safetensors.torch import save_file
 
 from tests.utils import single_gpu_only
 from vllm import SamplingParams
+from vllm.model_executor.model_loader.weight_utils import download_weights_from_hf
 from vllm.platforms import current_platform
+from vllm.v1.metrics.reader import Vector
 
 from ...utils import (
     assert_request_outputs_match,
     compute_acceptance_len,
+    get_spec_decode_metric_value,
     get_test_prompts,
 )
 from .._correctness import check_mtp_correctness
 
 
+def _add_random_mtp_module(model_name: str, path: Path) -> Path:
+    """Copy of `model_name` with a second MTP module of random weights."""
+    src = Path(download_weights_from_hf(model_name, None, allow_patterns=["*"]))
+    for file in src.iterdir():
+        if file.name not in ("config.json", "model.safetensors.index.json"):
+            (path / file.name).symlink_to(file)
+    config = json.loads((src / "config.json").read_text())
+    config["text_config"]["mtp_num_hidden_layers"] = 2
+    (path / "config.json").write_text(json.dumps(config))
+
+    index = json.loads((src / "model.safetensors.index.json").read_text())
+    generator = torch.Generator().manual_seed(0)
+    module_1 = {}
+    for name, shard in index["weight_map"].items():
+        if name.startswith("mtp.layers.0."):
+            with safe_open(src / shard, framework="pt") as f:
+                weight = f.get_tensor(name)
+            module_1[name.replace("layers.0.", "layers.1.")] = torch.randn(
+                weight.shape, generator=generator
+            ).to(weight.dtype)
+    save_file(module_1, path / "mtp_module_1.safetensors")
+    index["weight_map"].update(dict.fromkeys(module_1, "mtp_module_1.safetensors"))
+    (path / "model.safetensors.index.json").write_text(json.dumps(index))
+    return path
+
+
+@single_gpu_only
+def test_mtp_draft_steps_use_their_own_module(
+    monkeypatch: pytest.MonkeyPatch,
+    sampling_config: SamplingParams,
+    tmp_path: Path,
+    vllm_runner,
+):
+    """With two MTP modules, draft step 2 must run module 1, also with the
+    default drafter CUDA graphs. Module 1 has random weights, so its drafts are
+    rejected while module 0's are still accepted."""
+    monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1")
+    model = _add_random_mtp_module("Qwen/Qwen3.5-0.8B-Base", tmp_path)
+    with vllm_runner(
+        str(model),
+        block_size=None,
+        enable_chunked_prefill=None,
+        max_model_len=2048,
+        limit_mm_per_prompt={"image": 0, "video": 0},
+        speculative_config={
+            "method": "mtp",
+            "num_speculative_tokens": 2,
+            # Also covers get_top_tokens(..., spec_step_idx) on this path.
+            "use_local_argmax_reduction": True,
+        },
+        disable_log_stats=False,
+    ) as runner:
+        runner.llm.chat(get_test_prompts(mm_enabled=False), sampling_config)
+        metrics = runner.llm.get_metrics()
+
+    num_drafts = get_spec_decode_metric_value(metrics, "vllm:spec_decode_num_drafts")
+    accepted_per_pos = next(
+        m.values
+        for m in metrics
+        if isinstance(m, Vector)
+        and m.name == "vllm:spec_decode_num_accepted_tokens_per_pos"
+    )
+    step_1, step_2 = (accepted / num_drafts for accepted in accepted_per_pos)
+    assert step_1 > 0.3, f"module 0 acceptance {step_1:.3f}"
+    assert step_2 < 0.05, f"module 1 acceptance {step_2:.3f}: module 0 was reused"
+
+
 @pytest.mark.parametrize(
     ["model_setup", "mm_enabled", "expected_accuracy_threshold"],
     [
diff --git a/tests/v1/e2e/spec_decode/test_distribution_shift.py b/tests/v1/e2e/spec_decode/test_distribution_shift.py
new file mode 100644
index 000000000000..7cbe4827611b
--- /dev/null
+++ b/tests/v1/e2e/spec_decode/test_distribution_shift.py
@@ -0,0 +1,76 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from collections import Counter
+
+import pytest
+import regex as re
+from scipy.stats import chi2_contingency
+
+from tests.utils import single_gpu_only
+from vllm import SamplingParams
+from vllm.sampling_params import StructuredOutputsParams
+
+from .utils import compute_acceptance_rate
+
+
+@pytest.mark.parametrize("rejection_sample_method", ["standard", "block"])
+@pytest.mark.parametrize("async_scheduling", [False, True], ids=["sync", "async"])
+@single_gpu_only
+def test_structured_output_preserves_sampling_distribution(
+    rejection_sample_method, async_scheduling, vllm_runner
+):
+    spec_config = {
+        "method": "dflash",
+        "model": "z-lab/Qwen3-4B-DFlash-b16",
+        "num_speculative_tokens": 3,
+        "draft_sample_method": "probabilistic",
+        "rejection_sample_method": rejection_sample_method,
+    }
+    regex = r'\{"city": "(Paris|Berlin|Madrid)"\}'
+    num_samples = 4000
+    cities_counts = []
+    for config in (None, spec_config):
+        with vllm_runner(
+            "Qwen/Qwen3-4B",
+            speculative_config=config,
+            async_scheduling=async_scheduling,
+            enforce_eager=True,
+            enable_prefix_caching=False,
+            gpu_memory_utilization=0.8,
+            disable_log_stats=False,
+        ) as runner:
+            params = [
+                SamplingParams(
+                    temperature=1.0,
+                    max_tokens=32,
+                    seed=i,
+                    structured_outputs=StructuredOutputsParams(regex=regex),
+                )
+                for i in range(num_samples)
+            ]
+            outputs = runner.get_llm().chat(
+                [[{"role": "user", "content": "Name a random European capital."}]]
+                * num_samples,
+                params,
+                use_tqdm=False,
+                chat_template_kwargs={"enable_thinking": False},
+            )
+            if config is not None:
+                assert compute_acceptance_rate(runner.get_llm().get_metrics()) > 0
+            texts = [o.outputs[0].text for o in outputs]
+            assert all(re.fullmatch(regex, text) for text in texts), (
+                "Grammar-invalid output"
+            )
+            cities_counts.append(Counter(texts))
+    assert_grammar_distribution_unchanged(*cities_counts)
+
+
+def assert_grammar_distribution_unchanged(counts1: Counter[str], counts2: Counter[str]):
+    outcomes = sorted(counts1.keys() | counts2.keys())
+    row1 = [counts1[outcome] for outcome in outcomes]
+    row2 = [counts2[outcome] for outcome in outcomes]
+    _, p, _, _ = chi2_contingency([row1, row2], correction=False)
+    assert p > 1e-3, (
+        f"Grammar output distributions differ: {counts1=}, {counts2=}, {p=}"
+    )
diff --git a/tests/v1/e2e/spec_decode/test_sharded_sampling.py b/tests/v1/e2e/spec_decode/test_sharded_sampling.py
index ea323810bd95..ab0a5f59103e 100644
--- a/tests/v1/e2e/spec_decode/test_sharded_sampling.py
+++ b/tests/v1/e2e/spec_decode/test_sharded_sampling.py
@@ -101,3 +101,73 @@ def run(disable_sharding: bool):
                 f"{name}: {len(divergent)}/{len(ref_outputs)} prompts diverged "
                 f"at {divergent}, beyond near-tie noise"
             )
+
+
+def _dspark_sharding_state(worker) -> tuple[bool, bool]:
+    speculator = worker.model_runner.speculator
+    draft_logits = speculator.draft_logits
+    sharded = draft_logits is not None and draft_logits.row_mapping is not None
+    return sharded, speculator._sharded_sequential_comm is not None
+
+
+def test_dspark_sharded_sampling_equivalence(monkeypatch: pytest.MonkeyPatch):
+    """Batch-sharded sampling, including the sharded draft logits cache and
+    the batch-sharded DSpark sequential sampling loop, must match replicated
+    sampling up to near-tie noise."""
+    tp_size = 2
+    _skip_if_insufficient_gpus_for_tp(tp_size)
+    test_prompts = get_test_prompts(mm_enabled=False, num_prompts=32)
+
+    with monkeypatch.context() as m:
+        m.setenv("VLLM_USE_V2_MODEL_RUNNER", "1")
+        m.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
+        m.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
+
+        def run(disable_sharding: bool):
+            llm = LLM(
+                model="Qwen/Qwen3-4B-FP8",
+                tensor_parallel_size=tp_size,
+                enable_batch_sharded_sampling=not disable_sharding,
+                speculative_config={
+                    "method": "dspark",
+                    "model": "deepseek-ai/dspark_qwen3_4b_block7",
+                    "num_speculative_tokens": 7,
+                    "draft_sample_method": "probabilistic",
+                },
+                max_model_len=2048,
+                max_num_seqs=len(test_prompts),
+                enable_flashinfer_autotune=False,
+            )
+            try:
+                states = llm.llm_engine.collective_rpc(_dspark_sharding_state)
+                for state in states:
+                    # (draft logits cache sharded, sequential loop sharded)
+                    assert state == (not disable_sharding, not disable_sharding)
+                greedy = llm.chat(test_prompts, greedy_sampling())
+                seeded = llm.chat(
+                    test_prompts,
+                    SamplingParams(temperature=1.0, seed=33, max_tokens=32),
+                )
+                return greedy, seeded
+            finally:
+                del llm
+                torch.accelerator.empty_cache()
+                cleanup_dist_env_and_memory()
+                wait_for_memory_to_settle()
+
+        ref_greedy, ref_seeded = run(disable_sharding=True)
+        shard_greedy, shard_seeded = run(disable_sharding=False)
+
+        for name, ref_outputs, shard_outputs in (
+            ("greedy", ref_greedy, shard_greedy),
+            ("seeded", ref_seeded, shard_seeded),
+        ):
+            divergent = [
+                i
+                for i, (ref, out) in enumerate(zip(ref_outputs, shard_outputs))
+                if list(ref.outputs[0].token_ids) != list(out.outputs[0].token_ids)
+            ]
+            assert len(divergent) <= MAX_DIVERGENT_PROMPTS, (
+                f"{name}: {len(divergent)}/{len(ref_outputs)} prompts diverged "
+                f"at {divergent}, beyond near-tie noise"
+            )
diff --git a/tests/v1/ec_connector/integration/test_nixl_failure.py b/tests/v1/ec_connector/integration/test_nixl_failure.py
index 9e73d068bcc3..a5ca488040c9 100644
--- a/tests/v1/ec_connector/integration/test_nixl_failure.py
+++ b/tests/v1/ec_connector/integration/test_nixl_failure.py
@@ -30,7 +30,6 @@
 )
 from vllm.multimodal.utils import encode_image_url
 from vllm.platforms import current_platform
-from vllm.utils.network_utils import get_open_port
 
 MODEL = os.environ.get("MODEL", "Qwen/Qwen2.5-VL-3B-Instruct")
 
@@ -77,11 +76,7 @@ def config(role):
                 "--ec-transfer-config",
                 config("ec_producer"),
             ],
-            env_dict={
-                "CUDA_VISIBLE_DEVICES": devices[0],
-                "VLLM_EC_SIDE_CHANNEL_HOST": "127.0.0.1",
-                "VLLM_EC_SIDE_CHANNEL_PORT": str(get_open_port()),
-            },
+            env_dict={"CUDA_VISIBLE_DEVICES": devices[0]},
         ) as encoder,
         RemoteOpenAIServer(
             MODEL,
diff --git a/tests/v1/ec_connector/unit/test_control.py b/tests/v1/ec_connector/unit/test_control.py
index 93701738eb52..64bb077c9067 100644
--- a/tests/v1/ec_connector/unit/test_control.py
+++ b/tests/v1/ec_connector/unit/test_control.py
@@ -70,6 +70,7 @@ def test_close_sets_alive_false():
 from vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq import (  # noqa: E402
     ZmqClientConnection,
     ZmqClientTransport,
+    ZmqServerTransport,
 )
 from vllm.distributed.ec_transfer.ec_connector.cpu.utils import PeerAddr  # noqa: E402
 
@@ -283,3 +284,38 @@ def test_zmq_client_transport_close_closes_all_connections(
     dealer_a.close.assert_called()
     dealer_b.close.assert_called()
     assert transport._connections == {}
+
+
+def test_zmq_servers_allocate_distinct_reachable_ports():
+    """Simultaneous producers must publish usable ports without configuration."""
+    servers = []
+    client = ZmqClientTransport()
+    try:
+        for _ in range(2):
+            servers.append(ZmqServerTransport(host="127.0.0.1"))
+        assert len({server.port for server in servers}) == 2
+        for server in servers:
+            assert 0 < server.port <= 65535
+            conn = client.connect(("127.0.0.1", server.port))
+            conn.send(b"request")
+            messages = server.poll(timeout_ms=5000)
+            assert len(messages) == 1
+            identity, payload = messages[0]
+            assert payload == b"request"
+            server.send(identity, b"reply")
+            assert conn.dealer.poll(timeout=5000)
+            assert conn.recv() == [b"reply"]
+    finally:
+        client.close()
+        for server in servers:
+            server.close()
+
+
+def test_zmq_server_rejects_an_occupied_explicit_port():
+    server = ZmqServerTransport(host="127.0.0.1")
+    try:
+        with pytest.raises(zmq.ZMQError) as exc:
+            ZmqServerTransport(host="127.0.0.1", port=server.port)
+        assert exc.value.errno == zmq.EADDRINUSE
+    finally:
+        server.close()
diff --git a/tests/v1/ec_connector/unit/test_scheduler_nixl_ctor.py b/tests/v1/ec_connector/unit/test_scheduler_nixl_ctor.py
index 8cad9ec65c4d..2f3ea893174d 100644
--- a/tests/v1/ec_connector/unit/test_scheduler_nixl_ctor.py
+++ b/tests/v1/ec_connector/unit/test_scheduler_nixl_ctor.py
@@ -14,6 +14,7 @@
 )
 from vllm.distributed.ec_transfer.ec_connector.cpu.scheduler import ECCPUScheduler
 from vllm.distributed.nixl_utils import NixlWrapper
+from vllm.utils.network_utils import get_open_port
 
 _N = 16
 _BS = 64
@@ -126,13 +127,24 @@ def fail_control_transport():
 
 
 @pytest.mark.skipif(NixlWrapper is None, reason="Requires NIXL package")
-def test_gate_on_wires_data_transport_and_producer_session(monkeypatch):
-    # Port 0 lets the OS pick an ephemeral port so the real ZMQ ROUTER bind
-    # in ProducerSession.start() cannot collide with another test/process.
-    monkeypatch.setenv("VLLM_EC_SIDE_CHANNEL_HOST", "127.0.0.1")
-    monkeypatch.setenv("VLLM_EC_SIDE_CHANNEL_PORT", "0")
+@pytest.mark.parametrize("dp_index", [0, 1])
+@pytest.mark.parametrize("port_mode", ["unset", "zero", "fixed"])
+def test_gate_on_wires_data_transport_and_producer_session(
+    monkeypatch, dp_index, port_mode
+):
+    monkeypatch.setenv("VLLM_HOST_IP", "127.0.0.1")
+    monkeypatch.delenv("VLLM_EC_SIDE_CHANNEL_HOST", raising=False)
+    monkeypatch.delenv("VLLM_EC_SIDE_CHANNEL_PORT", raising=False)
+    port = 0
+    if port_mode == "fixed":
+        port = get_open_port()
+        monkeypatch.setenv("VLLM_EC_SIDE_CHANNEL_PORT", str(port - dp_index))
+        monkeypatch.setenv("VLLM_EC_SIDE_CHANNEL_HOST", "localhost")
+    elif port_mode == "zero":
+        monkeypatch.setenv("VLLM_EC_SIDE_CHANNEL_PORT", "0")
     monkeypatch.setattr(sched_mod, "create_ec_shared_region", lambda cfg: _region())
     cfg = create_ec_vllm_config(ec_role="ec_both")
+    cfg.parallel_config.data_parallel_index = dp_index
     cfg.ec_transfer_config.ec_connector_extra_config["ec_enable_nixl"] = True
     s = ECCPUScheduler(cfg)
     try:
@@ -140,7 +152,48 @@ def test_gate_on_wires_data_transport_and_producer_session(monkeypatch):
         assert s._data is not None
         assert s._compat_hash is not None
         assert s._producer_session is not None
-        assert s._peer_host == "127.0.0.1"
+        assert s._peer_host == ("localhost" if port else "127.0.0.1")
+        assert s._peer_port is not None and 0 < s._peer_port <= 65535
+        if port:
+            assert s._peer_port == port
+        else:
+            assert s._peer_port != dp_index  # Port 0 must not get the DP offset.
         assert s._transport is not None  # is_ec_consumer is also True
     finally:
         s.shutdown()
+
+
+@pytest.mark.skipif(NixlWrapper is None, reason="Requires NIXL package")
+def test_producers_allocate_distinct_ports_and_agent_names(monkeypatch):
+    """DP replicas retain their configured identity without peer collisions."""
+    monkeypatch.setenv("VLLM_HOST_IP", "127.0.0.1")
+    monkeypatch.delenv("VLLM_EC_SIDE_CHANNEL_HOST", raising=False)
+    monkeypatch.delenv("VLLM_EC_SIDE_CHANNEL_PORT", raising=False)
+    monkeypatch.setattr(sched_mod, "create_ec_shared_region", lambda cfg: _region())
+    engine_id = f"ec-test-{uuid.uuid4()}"
+    schedulers = []
+    expected_names = []
+    try:
+        for engine, dp_index in [(engine_id, 0), (engine_id, 1), (None, 0)]:
+            cfg = create_ec_vllm_config(ec_role="ec_producer")
+            if engine is not None:
+                cfg.ec_transfer_config.engine_id = engine
+            cfg.parallel_config.data_parallel_index = dp_index
+            cfg.ec_transfer_config.ec_connector_extra_config["ec_enable_nixl"] = True
+            expected_names.append(f"{cfg.ec_transfer_config.engine_id}_dp{dp_index}")
+            schedulers.append(ECCPUScheduler(cfg))
+        ports = [s._peer_port for s in schedulers]
+        assert all(port is not None and 0 < port <= 65535 for port in ports)
+        assert len(set(ports)) == len(schedulers)
+        for index, scheduler in enumerate(schedulers):
+            peer = schedulers[(index + 1) % len(schedulers)]._data
+            name = peer.add_remote_peer(
+                scheduler._data.get_agent_metadata(),
+                scheduler._data.get_mem_descriptor(),
+            )
+            if isinstance(name, bytes):
+                name = name.decode()
+            assert name == expected_names[index]
+    finally:
+        for scheduler in reversed(schedulers):
+            scheduler.shutdown()
diff --git a/tests/v1/engine/test_admission_control.py b/tests/v1/engine/test_admission_control.py
index dabd42711ca1..1f33bb52cc27 100644
--- a/tests/v1/engine/test_admission_control.py
+++ b/tests/v1/engine/test_admission_control.py
@@ -26,6 +26,7 @@
 from prometheus_client import REGISTRY
 from pydantic import ValidationError
 
+from vllm.config import VllmConfig
 from vllm.config.scheduler import SchedulerConfig
 from vllm.entrypoints.serve.exception_handling.error_response import (
     create_error_response,
@@ -70,6 +71,7 @@ def _make_async_llm(
 ) -> AsyncLLM:
     """Create a bare AsyncLLM with just the attributes needed for scheduling."""
     llm = AsyncLLM.__new__(AsyncLLM)
+    llm.vllm_config = VllmConfig()
     llm.scheduler_config = SimpleNamespace(
         max_num_queued_reqs=max_num_queued_reqs,
         max_num_queued_tokens=max_num_queued_tokens,
@@ -146,9 +148,6 @@ def _make_request_test_llm(
         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()
diff --git a/tests/v1/engine/test_core_engine_actor_manager.py b/tests/v1/engine/test_core_engine_actor_manager.py
index f8e2373d4c58..711db295d5db 100644
--- a/tests/v1/engine/test_core_engine_actor_manager.py
+++ b/tests/v1/engine/test_core_engine_actor_manager.py
@@ -15,7 +15,7 @@
 import zmq
 from torch.distributed import TCPStore
 
-from vllm.utils.network_utils import get_open_port, make_zmq_socket, split_zmq_path
+from vllm.utils.network_utils import make_zmq_socket, split_zmq_path
 from vllm.v1.engine.core import EngineCoreActorMixin
 from vllm.v1.engine.core_client import BackgroundResources
 from vllm.v1.engine.utils import (
@@ -162,8 +162,9 @@ def __init__(self, **kwargs) -> None:
     parallel_config = vllm_config.parallel_config
     for name, value in overrides.items():
         setattr(parallel_config, name, value)
-    # A fresh port per case so a lingering handshake socket cannot fail the next.
-    parallel_config.data_parallel_rpc_port = get_open_port()
+    # Port 0 makes the handshake socket take a free port at bind time; a port
+    # reserved up front can be handed to the coordination store meanwhile.
+    parallel_config.data_parallel_rpc_port = 0
     parallel_config._coord_store_port = 0
     monkeypatch.setattr("vllm.v1.engine.utils.CoreEngineProcManager", FakeManager)
     monkeypatch.setattr("vllm.v1.engine.utils.CoreEngineActorManager", FakeManager)
diff --git a/tests/v1/engine/test_engine_args.py b/tests/v1/engine/test_engine_args.py
index 5516ed840e47..c4c2963412c4 100644
--- a/tests/v1/engine/test_engine_args.py
+++ b/tests/v1/engine/test_engine_args.py
@@ -5,6 +5,7 @@
 from unittest.mock import patch
 
 import pytest
+import torch
 
 from vllm.config import ModelConfig
 from vllm.engine.arg_utils import EngineArgs
@@ -227,3 +228,368 @@ def test_external_lb_infers_rank_for_multinode_replicas():
         vllm_config = engine_args.create_engine_config(UsageContext.OPENAI_API_SERVER)
 
     assert vllm_config.parallel_config.data_parallel_rank == 0
+
+
+def test_extensible_kv_cache_from_cli():
+    parser = EngineArgs.add_cli_args(FlexibleArgumentParser())
+
+    args = parser.parse_args([])
+    engine_args = EngineArgs.from_cli_args(args=args)
+    assert engine_args.enable_extensible_kv_cache is None
+    assert engine_args.gpu_memory_utilization is None
+
+    args = parser.parse_args(["--enable-extensible-kv-cache"])
+    engine_args = EngineArgs.from_cli_args(args=args)
+    assert engine_args.enable_extensible_kv_cache
+
+    args = parser.parse_args(["--no-enable-extensible-kv-cache"])
+    engine_args = EngineArgs.from_cli_args(args=args)
+    assert engine_args.enable_extensible_kv_cache is False
+
+
+# Off CUDA the platform check trips first, before the behavior under test.
+requires_cuda = pytest.mark.skipif(
+    not torch.cuda.is_available(), reason="requires CUDA"
+)
+
+
+@requires_cuda
+def test_extensible_kv_cache_defaults():
+    """Unset, the extensible KV cache is on for the V2 runner on CUDA and the
+    budget is the whole device; turned off, or where unsupported, the standard
+    0.92 default returns. An explicit utilization is always kept."""
+    config = EngineArgs(model="facebook/opt-125m").create_engine_config(
+        UsageContext.OPENAI_API_SERVER
+    )
+    assert config.cache_config.enable_extensible_kv_cache
+    assert config.cache_config.gpu_memory_utilization is None
+    assert config.cache_config.resolved_gpu_memory_utilization == 1.0
+
+    config.cache_config.enable_extensible_kv_cache = False
+    assert config.cache_config.resolved_gpu_memory_utilization == 0.92
+
+    config = EngineArgs(
+        model="facebook/opt-125m", enable_extensible_kv_cache=False
+    ).create_engine_config(UsageContext.OPENAI_API_SERVER)
+    assert config.cache_config.resolved_gpu_memory_utilization == 0.92
+
+    config = EngineArgs(
+        model="facebook/opt-125m", gpu_memory_utilization=0.5
+    ).create_engine_config(UsageContext.OPENAI_API_SERVER)
+    assert config.cache_config.enable_extensible_kv_cache
+    assert config.cache_config.resolved_gpu_memory_utilization == 0.5
+    config.cache_config.enable_extensible_kv_cache = False
+    assert config.cache_config.resolved_gpu_memory_utilization == 0.5
+
+    # Manual KV sizing leaves the feature off by default but rejects an
+    # explicit request.
+    config = EngineArgs(
+        model="facebook/opt-125m", kv_cache_memory_bytes=1 << 30
+    ).create_engine_config(UsageContext.OPENAI_API_SERVER)
+    assert not config.cache_config.enable_extensible_kv_cache
+    assert config.cache_config.resolved_gpu_memory_utilization == 0.92
+    config = EngineArgs(
+        model="facebook/opt-125m", num_gpu_blocks_override=64
+    ).create_engine_config(UsageContext.OPENAI_API_SERVER)
+    assert not config.cache_config.enable_extensible_kv_cache
+
+
+@pytest.mark.parametrize(
+    "manual_size",
+    [dict(kv_cache_memory_bytes=1 << 30), dict(num_gpu_blocks_override=64)],
+)
+@requires_cuda
+def test_extensible_kv_cache_rejects_manual_kv_cache_size(manual_size):
+    """Measured sizing and a manual size conflict: silently clamping the manual
+    one would misreport the cache, so an explicit request is an error."""
+    (arg_name,) = manual_size
+    engine_args = EngineArgs(
+        model="facebook/opt-125m", enable_extensible_kv_cache=True, **manual_size
+    )
+    with pytest.raises(ValueError, match=arg_name):
+        engine_args.create_engine_config(UsageContext.OPENAI_API_SERVER)
+
+
+def _fake_executor(vllm_config, collective_rpc):
+    from types import MethodType, SimpleNamespace
+
+    from vllm.v1.executor.abstract import Executor
+
+    fake = SimpleNamespace(vllm_config=vllm_config, collective_rpc=collective_rpc)
+    for name in (
+        "extensible_kv_cache_unsupported_reason",
+        "disable_extensible_kv_cache",
+    ):
+        setattr(fake, name, MethodType(getattr(Executor, name), fake))
+    return fake
+
+
+@requires_cuda
+def test_extensible_kv_cache_falls_back_when_driver_unsupported():
+    calls: list[str] = []
+
+    def collective_rpc(method: str):
+        calls.append(method)
+        if method == "extensible_kv_cache_unsupported_reason":
+            return [None, "no VMM support"]
+        return [None, None]
+
+    engine_args = EngineArgs(model="facebook/opt-125m", enable_extensible_kv_cache=True)
+    vllm_config = engine_args.create_engine_config(UsageContext.OPENAI_API_SERVER)
+    assert vllm_config.cache_config.enable_extensible_kv_cache
+    assert vllm_config.cache_config.resolved_gpu_memory_utilization == 1.0
+    specs = [{"layer": object()}]
+    fake = _fake_executor(vllm_config, collective_rpc)
+    assert fake.extensible_kv_cache_unsupported_reason(specs) == "no VMM support"
+    fake.disable_extensible_kv_cache("no VMM support")
+    assert not vllm_config.cache_config.enable_extensible_kv_cache
+    assert vllm_config.cache_config.resolved_gpu_memory_utilization == 0.92
+    assert calls == [
+        "extensible_kv_cache_unsupported_reason",
+        "disable_extensible_kv_cache",
+    ]
+
+    vllm_config = engine_args.create_engine_config(UsageContext.OPENAI_API_SERVER)
+    fake = _fake_executor(vllm_config, lambda method: [None, None])
+    assert fake.extensible_kv_cache_unsupported_reason(specs) is None
+
+    # No KV cache at all: nothing to size, so the feature is turned off.
+    assert fake.extensible_kv_cache_unsupported_reason([{}]) is not None
+
+
+def test_all_reduce_int(monkeypatch):
+    """A MIN all-reduce by default; without a group or with a single rank, the
+    value stands."""
+    from types import SimpleNamespace
+
+    import torch.distributed as dist
+
+    from vllm.distributed.utils import all_reduce_int
+
+    group = SimpleNamespace(size=lambda: 2)
+    seen: list[tuple[object, object]] = []
+
+    def all_reduce(tensor, op, group):
+        seen.append((op, group))
+        tensor.fill_(3)
+
+    monkeypatch.setattr(dist, "all_reduce", all_reduce)
+    assert all_reduce_int(None, 7) == 7
+    assert all_reduce_int(SimpleNamespace(size=lambda: 1), 7) == 7
+    assert seen == []
+    assert all_reduce_int(group, 7) == 3
+    assert all_reduce_int(group, 7, dist.ReduceOp.MAX) == 3
+    assert seen == [(dist.ReduceOp.MIN, group), (dist.ReduceOp.MAX, group)]
+
+
+@pytest.mark.parametrize("requested", [None, True])
+def test_extensible_kv_cache_off_with_elastic_ep(monkeypatch, requested):
+    """Elastic EP scale-up skips warmup and reuses the profiled KV cache size,
+    so nothing would measure an extensible cache: it stays off by default and an
+    explicit request is an error."""
+    from types import SimpleNamespace
+
+    import vllm.platforms
+    from vllm.config.vllm import VllmConfig
+
+    monkeypatch.setattr(
+        vllm.platforms, "current_platform", SimpleNamespace(is_cuda_alike=lambda: True)
+    )
+    config = SimpleNamespace(
+        cache_config=SimpleNamespace(
+            _extensible_kv_cache_resolved=False,
+            enable_extensible_kv_cache=requested,
+            kv_cache_memory_bytes=None,
+            num_gpu_blocks_override=None,
+        ),
+        use_v2_model_runner=True,
+        attention_config=SimpleNamespace(hisparse_config=None),
+        parallel_config=SimpleNamespace(enable_elastic_ep=True),
+        kv_transfer_config=None,
+    )
+    if requested:
+        with pytest.raises(ValueError, match="elastic EP"):
+            VllmConfig._resolve_extensible_kv_cache(config)
+    else:
+        VllmConfig._resolve_extensible_kv_cache(config)
+        assert config.cache_config.enable_extensible_kv_cache is False
+
+
+@pytest.mark.parametrize(
+    ("connector", "buffer_device", "registers"),
+    [
+        ("NixlConnector", "cuda", True),
+        ("NixlPushConnector", "cuda", True),
+        ("NixlConnector", "cpu", False),
+        ("MooncakeConnector", "cuda", True),
+        ("MooncakeStoreConnector", "cuda", True),
+        ("MoRIIOConnector", "cuda", True),
+        ("OffloadingConnector", "cuda", False),
+        ("LMCacheConnectorV1", "cuda", False),
+    ],
+)
+@pytest.mark.parametrize("multi", [False, True])
+def test_kv_transfer_registers_kv_cache_memory(
+    connector, buffer_device, registers, multi
+):
+    """Connectors whose workers register the device KV cache with UCX or RDMA,
+    also when nested in MultiConnector."""
+    from vllm.config.kv_transfer import KVTransferConfig
+
+    if multi:
+        transfer = KVTransferConfig(
+            kv_connector="MultiConnector",
+            kv_role="kv_both",
+            kv_buffer_device=buffer_device,
+            kv_connector_extra_config={
+                "connectors": [
+                    {"kv_connector": "OffloadingConnector", "kv_role": "kv_both"},
+                    {"kv_connector": connector, "kv_role": "kv_both"},
+                ]
+            },
+        )
+    else:
+        transfer = KVTransferConfig(
+            kv_connector=connector, kv_role="kv_both", kv_buffer_device=buffer_device
+        )
+    assert transfer.registers_kv_cache_memory() is registers
+
+
+@pytest.mark.parametrize("rocm", [False, True])
+@pytest.mark.parametrize("requested", [None, True])
+def test_extensible_kv_cache_off_for_rocm_registering_connector(
+    monkeypatch, rocm, requested
+):
+    """UCX and RDMA cannot yet register driver-mapped memory on ROCm, so a
+    connector that registers the KV cache keeps the extensible cache off there
+    by default and an explicit request is an error; CUDA is unaffected."""
+    from types import SimpleNamespace
+
+    import vllm.platforms
+    from vllm.config.kv_transfer import KVTransferConfig
+    from vllm.config.vllm import VllmConfig
+
+    monkeypatch.setattr(
+        vllm.platforms,
+        "current_platform",
+        SimpleNamespace(is_cuda_alike=lambda: True, is_rocm=lambda: rocm),
+    )
+    config = SimpleNamespace(
+        cache_config=SimpleNamespace(
+            _extensible_kv_cache_resolved=False,
+            enable_extensible_kv_cache=requested,
+            kv_cache_memory_bytes=None,
+            num_gpu_blocks_override=None,
+        ),
+        use_v2_model_runner=True,
+        attention_config=SimpleNamespace(hisparse_config=None),
+        parallel_config=SimpleNamespace(enable_elastic_ep=False),
+        kv_transfer_config=KVTransferConfig(
+            kv_connector="NixlConnector", kv_role="kv_both", kv_buffer_device="cuda"
+        ),
+    )
+    if rocm and requested:
+        with pytest.raises(ValueError, match="ROCm"):
+            VllmConfig._resolve_extensible_kv_cache(config)
+        return
+    VllmConfig._resolve_extensible_kv_cache(config)
+    assert config.cache_config.enable_extensible_kv_cache is not rocm
+
+
+@pytest.mark.parametrize("multi", [False, True])
+@pytest.mark.parametrize("requested", [None, True])
+def test_extensible_kv_cache_off_with_lmcache_mp(monkeypatch, multi, requested):
+    """LMCacheMPConnector shares the KV cache through legacy CUDA IPC, which
+    cannot export driver-mapped memory on any platform."""
+    from types import SimpleNamespace
+
+    import vllm.platforms
+    from vllm.config.kv_transfer import KVTransferConfig
+    from vllm.config.vllm import VllmConfig
+
+    lmcache = {"kv_connector": "LMCacheMPConnector", "kv_role": "kv_both"}
+    transfer = (
+        KVTransferConfig(
+            kv_connector="MultiConnector",
+            kv_role="kv_both",
+            kv_connector_extra_config={"connectors": [lmcache]},
+        )
+        if multi
+        else KVTransferConfig(**lmcache)
+    )
+    monkeypatch.setattr(
+        vllm.platforms,
+        "current_platform",
+        SimpleNamespace(is_cuda_alike=lambda: True, is_rocm=lambda: False),
+    )
+    config = SimpleNamespace(
+        cache_config=SimpleNamespace(
+            _extensible_kv_cache_resolved=False,
+            enable_extensible_kv_cache=requested,
+            kv_cache_memory_bytes=None,
+            num_gpu_blocks_override=None,
+        ),
+        use_v2_model_runner=True,
+        attention_config=SimpleNamespace(hisparse_config=None),
+        parallel_config=SimpleNamespace(enable_elastic_ep=False),
+        kv_transfer_config=transfer,
+    )
+    if requested:
+        with pytest.raises(ValueError, match="legacy CUDA IPC"):
+            VllmConfig._resolve_extensible_kv_cache(config)
+    else:
+        VllmConfig._resolve_extensible_kv_cache(config)
+        assert config.cache_config.enable_extensible_kv_cache is False
+
+
+@requires_cuda
+@pytest.mark.parametrize("multi", [False, True])
+def test_extensible_kv_cache_rejects_connector_memory_pool(multi):
+    """A connector's custom memory pool (Mooncake NVLink/BAREX) must own the KV
+    allocation, which the driver-mapped extensible cache bypasses: the feature
+    stays off by default and an explicit request is an error."""
+    from vllm.config.kv_transfer import KVTransferConfig
+
+    mooncake = {
+        "kv_connector": "MooncakeStoreConnector",
+        "kv_role": "kv_both",
+        "kv_connector_extra_config": {"custom_mem_pool": "NVLINK"},
+    }
+    if multi:
+        transfer = KVTransferConfig(
+            kv_connector="MultiConnector",
+            kv_role="kv_both",
+            kv_connector_extra_config={"connectors": [mooncake]},
+        )
+    else:
+        transfer = KVTransferConfig(**mooncake)
+
+    def make_config(**kwargs):
+        return EngineArgs(
+            model="facebook/opt-125m", kv_transfer_config=transfer, **kwargs
+        ).create_engine_config(UsageContext.OPENAI_API_SERVER)
+
+    assert not make_config().cache_config.enable_extensible_kv_cache
+    with pytest.raises(ValueError, match="custom_mem_pool"):
+        make_config(enable_extensible_kv_cache=True)
+
+
+@requires_cuda
+def test_extensible_kv_cache_connector_needs_block_compact_layout():
+    from vllm.config.kv_transfer import KVTransferConfig
+    from vllm.v1.attention.backends.utils import resolve_kv_cache_layout
+
+    def make_config():
+        engine_args = EngineArgs(
+            model="facebook/opt-125m",
+            enable_extensible_kv_cache=True,
+            kv_transfer_config=KVTransferConfig(
+                kv_connector="ExampleConnector", kv_role="kv_both"
+            ),
+        )
+        return engine_args.create_engine_config(UsageContext.OPENAI_API_SERVER)
+
+    layout = resolve_kv_cache_layout(make_config(), [["LHBNC", "LBNHC"]])
+    assert layout.name == "LBNHC"
+    with pytest.raises(ValueError, match="block-compact"):
+        resolve_kv_cache_layout(make_config(), [["LHBNC"]])
diff --git a/tests/v1/engine/test_engine_core.py b/tests/v1/engine/test_engine_core.py
index c41d4d1bfe9d..7c52b32e3760 100644
--- a/tests/v1/engine/test_engine_core.py
+++ b/tests/v1/engine/test_engine_core.py
@@ -539,6 +539,7 @@ def test_encoder_instance_zero_kv_cache(
         trust_remote_code=True,
         dtype="float16",
         seed=42,
+        max_model_len=512,
     )
     scheduler_config = SchedulerConfig(
         max_num_seqs=10,
@@ -653,6 +654,7 @@ def _cadenced_dp_engine_core(
     """A bare DPEngineCoreProc whose all-reduce is scripted by `results`;
     returns the core and the step numbers at which the all-reduce ran."""
     core = object.__new__(DPEngineCoreProc)
+    core.enable_fault_tolerance = False
     core.dp_group = object()
     core.vllm_config = MagicMock(parallel_config=ParallelConfig())
     core.dp_sync_interval = dp_sync_interval
@@ -752,6 +754,7 @@ def _pausable_engine_core_proc() -> EngineCoreProc:
     """A bare EngineCoreProc holding just the state pause_scheduler touches."""
     core = object.__new__(EngineCoreProc)
     core.model_executor = MagicMock()
+    core.model_executor.vllm_config.cache_config.enable_extensible_kv_cache = False
     core.scheduler = MagicMock()
     core.scheduler.has_requests.return_value = False
     core.batch_queue = None
diff --git a/tests/v1/engine/test_kv_cache_group_metadata.py b/tests/v1/engine/test_kv_cache_group_metadata.py
index 0c01d4e307cb..245be26bfa70 100644
--- a/tests/v1/engine/test_kv_cache_group_metadata.py
+++ b/tests/v1/engine/test_kv_cache_group_metadata.py
@@ -38,7 +38,7 @@ def test_global_layout_finalizes_worker_interleave_before_profiling(
 ):
     """Workers with different local blocks must use the global interleave."""
     config = VllmConfig(
-        cache_config=CacheConfig(block_size=64),
+        cache_config=CacheConfig(block_size=64, enable_extensible_kv_cache=False),
         device_config=DeviceConfig(device="cpu"),
         parallel_config=ParallelConfig(
             tensor_parallel_size=2,
diff --git a/tests/v1/engine/test_llm_engine.py b/tests/v1/engine/test_llm_engine.py
index 9c831cac4eb5..491db06a0b88 100644
--- a/tests/v1/engine/test_llm_engine.py
+++ b/tests/v1/engine/test_llm_engine.py
@@ -25,6 +25,7 @@
 def test_step_invalidates_mm_processor_cache_misses():
     engine = object.__new__(LLMEngine)
     engine.should_execute_dummy_batch = False
+    engine.dp_group = None
     engine.log_stats = False
     engine.logger_manager = None
 
@@ -53,6 +54,37 @@ def test_step_invalidates_mm_processor_cache_misses():
     ]
 
 
+@pytest.mark.skip_global_cleanup
+@pytest.mark.parametrize("model_executed", [True, False])
+def test_dp_step_runs_dummy_batch_when_no_forward(model_executed: bool):
+    """Each DP rank must join one forward per step (in-process DP / external
+    launcher). A step whose engine-core step ran no forward (e.g. async
+    scheduling skips requests that already reached max_tokens) must run a
+    dummy batch, or the other ranks hang in their forward's DP sync."""
+    engine = object.__new__(LLMEngine)
+    engine.should_execute_dummy_batch = False
+    engine.dp_group = MagicMock()
+    engine.external_launcher_dp = True
+    engine.log_stats = False
+    engine.logger_manager = None
+    engine.renderer = MagicMock(mm_processor_cache=None)
+    engine.engine_core = MagicMock()
+    engine.engine_core.get_output.return_value = MagicMock(
+        outputs=[], timestamp=0.0, scheduler_stats=None
+    )
+    engine.engine_core.last_step_model_executed = model_executed
+    engine.output_processor = MagicMock()
+    engine.output_processor.process_outputs.return_value = MagicMock(
+        reqs_to_abort=[], request_outputs=[]
+    )
+
+    engine.step()
+
+    assert engine.engine_core.execute_dummy_batch.call_count == (
+        0 if model_executed else 1
+    )
+
+
 def test_duplicate_profile_start_rejected_before_engine_core_dispatch():
     engine = object.__new__(LLMEngine)
     engine.engine_core = MagicMock()
diff --git a/tests/v1/executor/test_cuda_checkpoint_executor.py b/tests/v1/executor/test_cuda_checkpoint_executor.py
new file mode 100644
index 000000000000..b5ad8d88948d
--- /dev/null
+++ b/tests/v1/executor/test_cuda_checkpoint_executor.py
@@ -0,0 +1,417 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""CPU lifecycle tests for the opt-in process checkpoint executor."""
+
+import multiprocessing
+import threading
+import time
+from concurrent.futures import ThreadPoolExecutor
+from types import SimpleNamespace
+from typing import Any
+from unittest.mock import Mock, patch
+
+import pytest
+
+from vllm.v1.executor.cuda_checkpoint_executor import CUDACheckpointExecutor
+from vllm.v1.executor.multiproc_executor import MultiprocExecutor
+
+
+@pytest.fixture
+def executor():
+    calls: list[tuple[Any, ...]] = []
+    obj = CUDACheckpointExecutor.__new__(CUDACheckpointExecutor)
+    obj._checkpoint_lock = threading.RLock()
+    obj._checkpoint_active = False
+    obj._use_nccl_suspend = False
+    obj._checkpoint_flashinfer = False
+    obj.is_failed = False
+    obj.sleeping_tags = set()
+    obj.failure_callback = None
+    obj.workers = [SimpleNamespace(proc=SimpleNamespace(pid=p)) for p in (11, 22)]
+    obj._checkpoint = SimpleNamespace(
+        suspend=lambda pids: calls.append(("checkpoint", pids)),
+        resume=lambda: calls.append(("restore",)),
+        close=lambda: calls.append(("close",)),
+    )
+
+    def sleep(self, level=1):
+        if "weights" in self.sleeping_tags:
+            return
+        calls.append(("sleep", level))
+        self.sleeping_tags = {"weights", "kv_cache"}
+
+    def wake(self, tags=None):
+        calls.append(("wake", tags))
+        self.sleeping_tags -= set(tags) if tags else self.sleeping_tags.copy()
+
+    with (
+        patch.object(MultiprocExecutor, "sleep", sleep),
+        patch.object(MultiprocExecutor, "wake_up", wake),
+        patch.object(
+            MultiprocExecutor, "shutdown", lambda _: calls.append(("shutdown",))
+        ),
+    ):
+        yield obj, calls
+    obj._checkpoint = None
+
+
+def test_sleep_then_restore_precedes_worker_wake(executor):
+    obj, calls = executor
+    obj.sleep()
+    obj.wake_up()
+    assert calls == [
+        ("sleep", 1),
+        ("checkpoint", [11, 22]),
+        ("restore",),
+        ("wake", None),
+    ]
+
+
+def test_repeated_sleep_does_not_checkpoint_twice(executor):
+    obj, calls = executor
+    obj.sleep()
+    obj.sleep()
+    assert calls == [("sleep", 1), ("checkpoint", [11, 22])]
+
+
+def test_partial_wake_restores_process_only_once(executor):
+    obj, calls = executor
+    obj.sleep()
+    obj.wake_up(["weights"])
+    obj.wake_up(["kv_cache"])
+    assert calls.count(("restore",)) == 1
+    assert not obj.is_sleeping
+
+
+def test_invalid_tag_does_not_restore(executor):
+    obj, calls = executor
+    obj.sleep()
+    with pytest.raises(ValueError):
+        obj.wake_up(["invalid"])
+    assert ("restore",) not in calls
+
+
+def test_failed_checkpoint_stops_workers_and_blocks_wake(executor):
+    obj, calls = executor
+
+    def fail(_):
+        raise RuntimeError("partial checkpoint")
+
+    obj._checkpoint.suspend = fail
+    with pytest.raises(RuntimeError, match="partial"):
+        obj.sleep()
+    assert obj.is_failed and ("shutdown",) in calls and ("close",) in calls
+    with pytest.raises(RuntimeError, match="failed"):
+        obj.wake_up()
+    assert not any(c[0] == "wake" for c in calls)
+
+
+def test_shutdown_does_not_restore_checkpointed_memory(executor):
+    obj, calls = executor
+    obj.sleep()
+
+    def fail():
+        raise RuntimeError("restore out of memory")
+
+    obj._checkpoint.resume = fail
+    obj.shutdown()
+    assert not obj.is_failed
+    assert ("shutdown",) in calls and ("close",) in calls
+    assert ("wake", None) not in calls
+
+
+def test_level_two_rejected_before_sleep(executor):
+    obj, calls = executor
+    with pytest.raises(ValueError):
+        obj.sleep(2)
+    assert not calls
+
+
+def test_checkpointed_worker_rejects_cuda_rpc_without_blocking_wake(executor):
+    obj, _ = executor
+    obj.sleep()
+    with pytest.raises(RuntimeError, match="checkpointed"):
+        obj.collective_rpc("profile")
+    obj.wake_up()
+
+
+def test_monitor_shutdown_waits_for_inflight_checkpoint(executor):
+    obj, calls = executor
+    entered, release, closed = threading.Event(), threading.Event(), threading.Event()
+
+    def suspend(pids):
+        entered.set()
+        assert release.wait(10)
+
+    lock, contended = obj._checkpoint_lock, threading.Event()
+
+    class ObservedLock:
+        def __enter__(self):
+            if not lock.acquire(blocking=False):
+                contended.set()
+                assert lock.acquire(timeout=10)
+
+        def __exit__(self, *exc):
+            lock.release()
+
+    obj._checkpoint_lock = ObservedLock()
+    obj._checkpoint.suspend = suspend
+    obj._checkpoint.close = lambda: closed.set()
+    with ThreadPoolExecutor(max_workers=2) as pool:
+        sleeper = pool.submit(obj.sleep)
+        try:
+            assert entered.wait(10)
+            closer = pool.submit(obj.shutdown)
+            assert contended.wait(5)
+            assert not closed.is_set()
+        finally:
+            release.set()
+        sleeper.result(timeout=10)
+        closer.result(timeout=10)
+    assert closed.is_set()
+
+
+def test_repeated_pause_acknowledges_already_completed_device_work(executor):
+    obj, _ = executor
+    obj.sleep()
+    for method in ("synchronize_device", "reset_mm_cache", "reset_encoder_cache"):
+        assert obj.collective_rpc(method) == [None, None]
+    obj.wake_up()
+
+
+def test_nonblocking_pause_returns_completed_future(executor):
+    obj, _ = executor
+    obj.sleep()
+    result = obj.collective_rpc("synchronize_device", non_block=True)
+    assert result.done() and result.result() == [None, None]
+
+
+def test_failed_executor_never_acknowledges_pause(executor):
+    obj, _ = executor
+    obj.sleep()
+    obj.is_failed = True
+    with pytest.raises(RuntimeError, match="failed"):
+        obj.collective_rpc("synchronize_device")
+
+
+@pytest.mark.parametrize(
+    "p2p,shm,tp,nvls,suspend,fusion,backend,accepted",
+    [
+        pytest.param("0", "1", 2, "0", True, False, None, True, id="p2p-only"),
+        pytest.param("0", "0", 3, "0", True, False, None, True, id="p2p-and-shm"),
+        pytest.param("1", "0", 8, "0", True, False, None, True, id="shm-only"),
+        pytest.param(None, None, 2, "0", True, False, None, True, id="defaults"),
+        pytest.param("0", "0", 2, "0", True, True, "trtllm", True, id="trtllm"),
+        pytest.param("0", "0", 2, "1", True, False, None, False, id="nvls"),
+        pytest.param("0", "0", 2, "0", False, False, None, False, id="no-suspend"),
+        pytest.param("0", "0", 2, "0", True, True, None, False, id="fusion-default"),
+        pytest.param("0", "0", 2, "0", True, True, "mnnvl", False, id="mnnvl"),
+    ],
+)
+def test_transport_requires_suspension_and_no_nvls(
+    monkeypatch, nvls, suspend, accepted, fusion, backend, p2p, shm, tp
+):
+    # Transport options are independent of the rejection rules; avoid their
+    # Cartesian product while retaining each option and validation branch.
+    monkeypatch.setattr("vllm.platforms.current_platform.is_cuda", lambda: True)
+    env = {
+        "NCCL_P2P_DISABLE": p2p,
+        "NCCL_SHM_DISABLE": shm,
+        "NCCL_IB_DISABLE": "1",
+        "NCCL_CUMEM_ENABLE": "1",
+        "NCCL_CUMEM_HOST_ENABLE": "0",
+        "NCCL_NVLS_ENABLE": nvls,
+        "VLLM_ALLREDUCE_USE_FLASHINFER": "0",
+        "VLLM_ALLREDUCE_USE_SYMM_MEM": "0",
+        "VLLM_ALLREDUCE_USE_FLASHINFER_PCIE_IPC": "0",
+        "VLLM_USE_NCCL_SYMM_MEM": "0",
+    }
+    for key, value in env.items():
+        if value is None:
+            monkeypatch.delenv(key, raising=False)
+        else:
+            monkeypatch.setenv(key, value)
+    monkeypatch.delenv("VLLM_FLASHINFER_ALLREDUCE_BACKEND", raising=False)
+    if backend is not None:
+        monkeypatch.setenv("VLLM_FLASHINFER_ALLREDUCE_BACKEND", backend)
+    config = SimpleNamespace(
+        compilation_config=SimpleNamespace(
+            mode=3 if fusion else 0,
+            pass_config=SimpleNamespace(fuse_allreduce_rms=fusion),
+        ),
+        parallel_config=SimpleNamespace(
+            tensor_parallel_size=tp,
+            nnodes=1,
+            decode_context_parallel_size=1,
+            prefill_context_parallel_size=1,
+            pipeline_parallel_size=1,
+            data_parallel_size=1,
+            world_size=tp,
+            disable_custom_all_reduce=True,
+            enable_expert_parallel=False,
+            worker_cls="vllm.v1.worker.gpu_worker.Worker",
+        ),
+        model_config=SimpleNamespace(
+            enable_sleep_mode=True,
+            enable_nccl_comm_suspend=suspend,
+            sleep_mode_backend="cumem",
+        ),
+        use_v2_model_runner=True,
+        weight_transfer_config=None,
+        kv_transfer_config=None,
+        ec_transfer_config=None,
+    )
+    with (
+        patch.object(MultiprocExecutor, "__init__", return_value=None),
+        patch.object(MultiprocExecutor, "shutdown"),
+    ):
+        if accepted:
+            obj = CUDACheckpointExecutor(config)
+            assert obj._checkpoint_flashinfer == fusion
+            obj.shutdown()
+        else:
+            with pytest.raises(ValueError, match="NCCL checkpoint configuration"):
+                CUDACheckpointExecutor(config)
+
+
+@pytest.mark.parametrize(
+    "has_symbols,actual,expected",
+    [(False, True, True), (True, False, True), (True, True, False)],
+)
+def test_validate_rejects_missing_or_inconsistent_nccl(
+    monkeypatch, has_symbols, actual, expected
+):
+    from vllm.v1.executor.cuda_checkpoint_executor import _validate_nccl_checkpoint
+
+    comm = SimpleNamespace(
+        available=True,
+        disabled=False,
+        _suspended=actual,
+        nccl=SimpleNamespace(has_symbol=lambda _: has_symbols),
+    )
+    group = SimpleNamespace(device_communicator=SimpleNamespace(pynccl_comm=comm))
+    monkeypatch.setattr("vllm.distributed.get_tp_group", lambda: group)
+    with pytest.raises(RuntimeError, match="active suspend/resume"):
+        _validate_nccl_checkpoint(None, expected)
+
+
+def test_flashinfer_hooks_wrap_driver_checkpoint_once(executor):
+    obj, calls = executor
+    obj._checkpoint_flashinfer = True
+    with patch.object(
+        obj, "collective_rpc", side_effect=lambda method, **kw: calls.append((method,))
+    ):
+        obj.sleep()
+        obj.wake_up(["weights"])
+        obj.wake_up(["kv_cache"])
+    assert calls == [
+        ("sleep", 1),
+        ("checkpoint_prepare",),
+        ("checkpoint", [11, 22]),
+        ("restore",),
+        ("checkpoint_restore",),
+        ("wake", ["weights"]),
+        ("wake", ["kv_cache"]),
+    ]
+
+
+def test_sleep_worker_rpc_has_deadline(executor):
+    obj, _ = executor
+    with patch.object(MultiprocExecutor, "collective_rpc", return_value=[]) as rpc:
+        obj.collective_rpc("sleep")
+        assert rpc.call_args.kwargs["timeout"] == 120
+        obj.collective_rpc("wake_up", timeout=7)
+        assert rpc.call_args.kwargs["timeout"] == 7
+
+
+def test_worker_termination_reaps_real_processes():
+    ctx = multiprocessing.get_context("spawn")
+    processes = [ctx.Process(target=time.sleep, args=(60,)) for _ in range(2)]
+    try:
+        for proc in processes:
+            proc.start()
+        CUDACheckpointExecutor._ensure_worker_termination(processes)
+        assert all(proc.exitcode is not None for proc in processes)
+        assert not any(proc.is_alive() for proc in processes)
+    finally:
+        for proc in processes:
+            if proc.is_alive():
+                proc.kill()
+            proc.join(timeout=5)
+            proc.close()
+
+
+def test_kv_first_partial_wake_then_sleep_preserves_parent_noop(executor):
+    obj, calls = executor
+    obj._use_nccl_suspend = True
+    obj.collective_rpc = Mock()
+    obj.sleep()
+    obj.wake_up(["kv_cache"])
+    assert obj.sleeping_tags == {"weights"}
+    before = list(calls)
+    rpc_count = obj.collective_rpc.call_count
+    obj.sleep()
+    assert calls == before
+    assert obj.collective_rpc.call_count == rpc_count
+    assert not obj.is_failed
+    obj.wake_up(["weights"])
+    assert not obj.is_sleeping
+
+
+@pytest.mark.parametrize(
+    "component,field,value",
+    [
+        ("parallel_config", "tensor_parallel_size", 0),
+        ("parallel_config", "tensor_parallel_size", -1),
+        ("parallel_config", "nnodes", 2),
+        ("parallel_config", "enable_expert_parallel", True),
+        ("parallel_config", "worker_cls", "custom.Worker"),
+        ("model_config", "sleep_mode_backend", "custom"),
+        (None, "use_v2_model_runner", False),
+        (None, "weight_transfer_config", object()),
+        ("platform", "is_cuda", False),
+        ("platform", "linux", False),
+    ],
+)
+def test_unsupported_checkpoint_configuration_fails_before_workers(
+    component, field, value, monkeypatch
+):
+    config = SimpleNamespace(
+        parallel_config=SimpleNamespace(
+            tensor_parallel_size=1,
+            world_size=1,
+            nnodes=1,
+            decode_context_parallel_size=1,
+            prefill_context_parallel_size=1,
+            pipeline_parallel_size=1,
+            data_parallel_size=1,
+            enable_expert_parallel=False,
+            worker_cls="vllm.v1.worker.gpu_worker.Worker",
+        ),
+        model_config=SimpleNamespace(
+            enable_sleep_mode=True,
+            enable_nccl_comm_suspend=False,
+            sleep_mode_backend="cumem",
+        ),
+        use_v2_model_runner=True,
+        weight_transfer_config=None,
+        kv_transfer_config=None,
+        ec_transfer_config=None,
+    )
+    monkeypatch.setattr("vllm.platforms.current_platform.is_cuda", lambda: True)
+    if component == "platform":
+        if field == "is_cuda":
+            monkeypatch.setattr(
+                "vllm.platforms.current_platform.is_cuda", lambda: value
+            )
+        else:
+            monkeypatch.setattr("sys.platform", "darwin")
+    else:
+        setattr(getattr(config, component) if component else config, field, value)
+    with (
+        patch.object(MultiprocExecutor, "__init__") as initialize,
+        patch.object(MultiprocExecutor, "shutdown"),
+        pytest.raises(ValueError, match="checkpoint"),
+    ):
+        CUDACheckpointExecutor(config)
+    initialize.assert_not_called()
diff --git a/tests/v1/executor/test_cuda_process_checkpoint.py b/tests/v1/executor/test_cuda_process_checkpoint.py
new file mode 100644
index 000000000000..1e7aafb90ce1
--- /dev/null
+++ b/tests/v1/executor/test_cuda_process_checkpoint.py
@@ -0,0 +1,231 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Controller ordering and helper cleanup without CUDA initialization."""
+
+import io
+import json
+import subprocess
+import sys
+from unittest.mock import Mock, patch
+
+import pytest
+
+from vllm.utils.cuda_process_checkpoint import CudaProcessCheckpoint, Driver
+
+
+@pytest.mark.parametrize("rank_count", [1, 2])
+def test_all_ranks_restore_before_any_unlock(rank_count):
+    events: list[tuple[int, str]] = []
+    with patch("vllm.utils.cuda_process_checkpoint.Driver") as driver:
+        driver.return_value.run_stage.side_effect = lambda pids, op: events.extend(
+            (pid, op) for pid in pids
+        )
+        checkpoint = CudaProcessCheckpoint(lambda **event: None)
+        checkpoint.suspend(list(range(1, rank_count + 1)))
+        checkpoint.resume()
+    assert events == [
+        (pid, op)
+        for op in ("Lock", "Checkpoint", "Restore", "Unlock")
+        for pid in range(1, rank_count + 1)
+    ]
+    assert checkpoint.pids == []
+
+
+def test_partial_failure_prevents_further_rank_operations():
+    events = []
+
+    def invoke(pid, op):
+        events.append((pid, op))
+        if pid == 22 and op == "Restore":
+            raise RuntimeError("injected restore failure")
+
+    with patch("vllm.utils.cuda_process_checkpoint.Driver") as driver:
+        driver.return_value.run_stage.side_effect = lambda pids, op: [
+            invoke(pid, op) for pid in pids
+        ]
+        checkpoint = CudaProcessCheckpoint(lambda **event: None)
+        checkpoint.suspend([11, 22])
+        with pytest.raises(RuntimeError, match="injected"):
+            checkpoint.resume()
+        count = len(events)
+        with pytest.raises(RuntimeError, match="restart required"):
+            checkpoint.resume()
+        assert len(events) == count
+        assert not any(op == "Unlock" for _, op in events)
+        assert checkpoint.failed
+
+
+def test_helper_initialization_failure_reaps_process():
+    proc = Mock()
+    proc.poll.return_value = None
+    with (
+        patch("subprocess.Popen", return_value=proc),
+        patch.object(Driver, "_read", side_effect=RuntimeError("init")),
+        pytest.raises(RuntimeError, match="init"),
+    ):
+        Driver(lambda **event: None)
+    proc.kill.assert_called_once()
+    proc.communicate.assert_called_once_with(timeout=5)
+
+
+def test_helper_read_timeout_kills_process():
+    driver = Driver.__new__(Driver)
+    driver.proc = Mock()
+    selector = Mock()
+    selector.select.return_value = []
+    with (
+        patch("selectors.DefaultSelector", return_value=selector),
+        pytest.raises(TimeoutError, match="deadline"),
+    ):
+        driver._read(0)
+    driver.proc.kill.assert_called_once()
+    selector.close.assert_called_once()
+
+
+def test_helper_close_reaps_real_process():
+    driver = Driver.__new__(Driver)
+    driver.proc = subprocess.Popen(
+        [sys.executable, "-c", "import time; time.sleep(60)"],
+        stdin=subprocess.PIPE,
+        stdout=subprocess.PIPE,
+        text=True,
+    )
+    try:
+        driver.close()
+        assert driver.proc.returncode is not None
+        assert driver.proc.stdin is not None and driver.proc.stdout is not None
+        assert driver.proc.stdin.closed and driver.proc.stdout.closed
+    finally:
+        if driver.proc.poll() is None:
+            driver.proc.kill()
+            driver.proc.wait(timeout=5)
+
+
+def test_driver_stage_runs_ranks_concurrently():
+    import threading
+
+    from vllm.utils.cuda_process_checkpoint import _run_parallel_stage
+
+    pids = [11, 22, 33, 44]
+    barrier = threading.Barrier(len(pids), timeout=2)
+
+    def invoke(pid, op):
+        barrier.wait()
+        return {"pid": pid, "op": op, "rc": 0}
+
+    rows = _run_parallel_stage(pids, "Restore", invoke)
+    assert rows == [{"pid": pid, "op": "Restore", "rc": 0} for pid in pids]
+
+
+def test_driver_stage_collects_rank_failure():
+    from vllm.utils.cuda_process_checkpoint import _run_parallel_stage
+
+    def invoke(pid, op):
+        if pid == 22:
+            raise RuntimeError("injected")
+        return {"pid": pid, "op": op, "rc": 0}
+
+    rows = _run_parallel_stage([11, 22], "Restore", invoke)
+    assert rows[0]["rc"] == 0
+    assert rows[1]["pid"] == 22 and rows[1]["op"] == "Restore"
+    assert "injected" in rows[1]["error"]
+
+
+@pytest.fixture
+def stage_driver(monkeypatch):
+    driver = Driver.__new__(Driver)
+    driver.emit = Mock()
+    driver.proc = Mock(stdin=io.StringIO())
+    response = Mock()
+    monkeypatch.setattr(driver, "_read", response)
+    return driver, response
+
+
+@pytest.mark.parametrize(
+    "fault", ["missing", "duplicate", "wrong_op", "bad_state", "failed"]
+)
+def test_driver_rejects_invalid_stage_response(fault, stage_driver):
+    driver, response = stage_driver
+    rows = [
+        {"pid": pid, "op": "Restore", "rc": 0, "after": {"rc": 0, "value": 1}}
+        for pid in [11, 22]
+    ]
+    if fault == "missing":
+        rows.pop()
+    elif fault == "duplicate":
+        rows[1]["pid"] = 11
+    elif fault == "wrong_op":
+        rows[1]["op"] = "Unlock"
+    elif fault == "bad_state":
+        rows[1]["after"] = {"rc": 0, "value": 2}
+    else:
+        rows[1]["rc"] = 2
+    response.return_value = {"kind": "stage", "op": "Restore", "results": rows}
+    with pytest.raises(RuntimeError):
+        driver.run_stage([11, 22], "Restore")
+
+
+@pytest.mark.parametrize("response", [None, [], "bad", 0])
+def test_driver_rejects_non_object_response(response, stage_driver):
+    driver, read = stage_driver
+    read.return_value = response
+    with pytest.raises(RuntimeError, match="Invalid CUDA checkpoint stage response"):
+        driver.run_stage([11], "Restore")
+
+
+def test_driver_accepts_complete_stage_response(stage_driver):
+    driver, response = stage_driver
+    rows = [
+        {"pid": pid, "op": "Restore", "rc": 0, "after": {"rc": 0, "value": 1}}
+        for pid in [11, 22]
+    ]
+    response.return_value = {"kind": "stage", "op": "Restore", "results": rows}
+    driver.run_stage([11, 22], "Restore")
+    assert json.loads(driver.proc.stdin.getvalue()) == [[11, 22], "Restore"]
+    for row in rows:
+        driver.emit.assert_any_call(**row)
+
+
+def test_rank_failure_with_blocked_peer_times_out_and_reaps_helper():
+    import textwrap
+
+    driver = Driver.__new__(Driver)
+    driver.emit = Mock()
+    driver.proc = subprocess.Popen(
+        [
+            sys.executable,
+            "-u",
+            "-c",
+            textwrap.dedent("""
+                import json
+                import sys
+                import time
+                from vllm.utils.cuda_process_checkpoint import _run_parallel_stage
+
+                def invoke(pid, op):
+                    if pid == 11:
+                        raise RuntimeError("injected rank failure")
+                    time.sleep(60)
+
+                print(json.dumps({"kind": "ready"}), flush=True)
+                pids, op = json.loads(sys.stdin.readline())
+                rows = _run_parallel_stage(pids, op, invoke)
+                print(json.dumps({"kind": "stage", "op": op, "results": rows}))
+            """),
+        ],
+        stdin=subprocess.PIPE,
+        stdout=subprocess.PIPE,
+        text=True,
+    )
+    try:
+        assert driver._read(10) == {"kind": "ready"}
+        read = driver._read
+        with (
+            patch.object(driver, "_read", side_effect=lambda timeout: read(0.2)),
+            pytest.raises(TimeoutError, match="deadline"),
+        ):
+            driver.run_stage([11, 22], "Restore")
+        driver.proc.wait(timeout=5)
+        assert driver.proc.returncode != 0
+    finally:
+        driver.close()
diff --git a/tests/v1/fault_tolerance/test_expert_reload.py b/tests/v1/fault_tolerance/test_expert_reload.py
new file mode 100644
index 000000000000..495405f2fdb4
--- /dev/null
+++ b/tests/v1/fault_tolerance/test_expert_reload.py
@@ -0,0 +1,42 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+
+from vllm.v1.worker.sentinel.expert_reload import _placement_weight_loader
+
+
+def test_reload_every_local_replica_before_deferred_processing():
+    """Reload recovered placement and every local replica using physical IDs."""
+    weights = torch.full((3,), -1.0)
+    pending = []
+
+    def buffer(param, loaded_weight, expert_id, return_success=False):
+        pending.append((param, loaded_weight, expert_id))
+        return True if return_success else None
+
+    load = _placement_weight_loader(buffer, {2: [4, 5], 0: [6]}, 4, 8)
+    assert load(weights, torch.tensor(20.0), 2, True)
+    assert load(weights, torch.tensor(10.0), 0, True)
+    assert not load(weights, torch.tensor(99.0), 6, True)
+    assert not load(weights, torch.tensor(99.0), 1, True)
+    assert torch.equal(weights, torch.full((3,), -1.0))
+    for param, value, physical_id in pending:
+        param[physical_id - 4].copy_(value)
+    assert torch.equal(weights, torch.tensor([20.0, 20.0, 10.0]))
+
+
+def test_preserve_global_activation_scales_for_nonlocal_experts():
+    """Preserve global activation scales even when no weight slots are local."""
+    scales = torch.zeros(8)
+
+    def load_scale(param, loaded_weight, weight_name, expert_id, return_success=False):
+        param[expert_id].copy_(loaded_weight)
+        return True if return_success else None
+
+    load = _placement_weight_loader(load_scale, {}, 4, 8, use_global_sf=True)
+    for expert_id in range(8):
+        assert load(
+            scales, torch.tensor(float(expert_id)), "input_scale", expert_id, True
+        )
+    assert torch.equal(scales, torch.arange(8, dtype=scales.dtype))
diff --git a/tests/v1/fault_tolerance/test_fault_tolerance_e2e.py b/tests/v1/fault_tolerance/test_fault_tolerance_e2e.py
index 27e45f6635f6..150f8fe565e8 100644
--- a/tests/v1/fault_tolerance/test_fault_tolerance_e2e.py
+++ b/tests/v1/fault_tolerance/test_fault_tolerance_e2e.py
@@ -9,6 +9,7 @@
 import os
 import threading
 import time
+import uuid
 from concurrent.futures import ThreadPoolExecutor
 from typing import Any
 
@@ -18,20 +19,29 @@
 
 from tests.utils import RemoteOpenAIServer, multi_gpu_test
 from vllm.utils.import_utils import has_nixl_ep
+from vllm.utils.network_utils import get_open_port
 
-MODEL_NAME = os.getenv("MODEL_NAME", "ibm-research/PowerMoE-3b")
-DP_SIZE = 2
+MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen3-30B-A3B")
 
 # Fault-detection timeout budget:
 # - CPU: Gloo DP allreduce timeout (30s) detects the dead peer.
 # - nixl_ep: kernel masks the dead rank after Buffer's default timeout_ms=30000 (30s).
-# - Deadline (45s): slowest fallback (30s) + margin.
+# - Eager FP8 quantization can wait for the GPU timeout before the next CPU sync.
+# - Deadline (75s): both 30s waits + margin.
 CPU_DISTRIBUTED_TIMEOUT_S = 30
-FAULT_DETECTION_DEADLINE_S = 45
+FAULT_DETECTION_DEADLINE_S = 75
 
 
-# Patches ``gpu.dp_utils.sync_cudagraph_and_dp_padding`` to raise on ``rank`` at
-# a chosen step. Gated on VLLM_FT_TEST_INJECT_FAULT.
+NUM_REDUNDANT_EXPERTS = 64
+
+# Recovery after scale_down (mask + expert redistribution + weight reload)
+# is expected to finish well within this deadline; the busy-loop wrapper
+# itself waits engine_recovery_timeout_sec (120s) before giving up.
+SCALE_DOWN_DEADLINE_S = 20
+
+
+# Patches ``gpu.dp_utils.sync_cudagraph_and_dp_padding`` to raise on ``rank``
+# once armed (via /collective_rpc). Gated on VLLM_FT_TEST_INJECT_FAULT.
 _FAULT_INJECT_SITECUSTOMIZE = """\
 import builtins
 import os
@@ -40,11 +50,15 @@
 _SPEC = os.environ.get("VLLM_FT_TEST_INJECT_FAULT")
 _MODULE = "vllm.v1.worker.gpu.dp_utils"
 _ATTR = "sync_cudagraph_and_dp_padding"
+_WORKER_MODULE = "vllm.v1.worker.gpu_worker"
 
 if _SPEC:
     _f = dict(kv.split("=", 1) for kv in _SPEC.split(","))
-    _RANK, _STEP = int(_f["rank"]), int(_f["step"])
-    _steps = [0]
+    _RANK = int(_f["rank"])
+    _ARMED = [False]
+
+    def _arm():
+        _ARMED[0] = True
 
     def _patch(m):
         import inspect
@@ -56,11 +70,11 @@ def _wrapped(*args, **kwargs):
             bound.apply_defaults()
             dp_rank = bound.arguments.get("dp_rank")
             if dp_rank == _RANK:
-                _steps[0] += 1
-                if _steps[0] == _STEP:
-                    raise RuntimeError(
-                        "FT test fault injection (rank=%d step=%d)" % (_RANK, _STEP)
-                    )
+                # One-shot: disarm before raising so post-retry steps pass.
+                if not _ARMED[0]:
+                    return result
+                _ARMED[0] = False
+                raise RuntimeError("FT test fault injection (rank=%d)" % _RANK)
             return result
 
         setattr(m, _ATTR, _wrapped)
@@ -79,17 +93,26 @@ def _hook(name, *a, **k):
         ):
             m._ft_patched = True
             _patch(m)
+        # Expose arming to /collective_rpc via a Worker method.
+        w = sys.modules.get(_WORKER_MODULE)
+        if (
+            w is not None
+            and hasattr(w, "Worker")
+            and not getattr(w, "_ft_arm_patched", False)
+        ):
+            w._ft_arm_patched = True
+            w.Worker.arm_ft_fault_injection = lambda self, *a, **k: _arm()
         return module
 
     builtins.__import__ = _hook
 """
 
 
-def _install_fault_injection(monkeypatch, tmp_path, rank: int, step: int) -> None:
-    """Arrange for the DP-sync fn to raise on ``rank`` at serving ``step``.
+def _install_fault_injection(monkeypatch, tmp_path, rank: int) -> None:
+    """Install a sitecustomize that raises on ``rank`` once armed.
 
-    Writes a ``sitecustomize.py`` and prepends its dir to PYTHONPATH so every
-    vLLM subprocess picks it up; the fault spec is read from the environment.
+    Disarmed until the test calls /collective_rpc (needs VLLM_SERVER_DEV_MODE,
+    set here).
     """
     site_dir = tmp_path / "ft_inject"
     site_dir.mkdir()
@@ -99,10 +122,11 @@ def _install_fault_injection(monkeypatch, tmp_path, rank: int, step: int) -> Non
         "PYTHONPATH",
         str(site_dir) + (os.pathsep + existing if existing else ""),
     )
-    monkeypatch.setenv("VLLM_FT_TEST_INJECT_FAULT", f"rank={rank},step={step}")
+    monkeypatch.setenv("VLLM_FT_TEST_INJECT_FAULT", f"rank={rank}")
+    monkeypatch.setenv("VLLM_SERVER_DEV_MODE", "1")
 
 
-def _ft_server_args() -> list[str]:
+def _ft_server_args(extra_args: list[str] | None = None) -> list[str]:
     return [
         "--enforce-eager",
         "--dtype",
@@ -119,18 +143,19 @@ def _ft_server_args() -> list[str]:
         str(CPU_DISTRIBUTED_TIMEOUT_S),
         "--fault-tolerance-config",
         '{"engine_recovery_timeout_sec": 120}',
+        *(extra_args or []),
     ]
 
 
-def _ft_manager():
+def _ft_manager(extra_args: list[str] | None = None, dp_size: int = 2):
     """Build the shared DP+EP fault-tolerant server topology (one engine/server)."""
     from tests.v1.distributed.test_external_lb_dp import ExternalLBServerManager
 
     return ExternalLBServerManager(
         MODEL_NAME,
-        DP_SIZE,
+        dp_size,
         api_server_count=1,  # FT requires a single API server per engine
-        base_server_args=_ft_server_args(),
+        base_server_args=_ft_server_args(extra_args),
         tp_size=1,
     )
 
@@ -156,6 +181,38 @@ def _complete(client):
     )
 
 
+# Invalid responses count as wrong answers. Env overrides exist for tuning.
+NUM_GSM8K_QUESTIONS = int(os.getenv("NUM_GSM8K_QUESTIONS", "256"))
+EXPECTED_ACCURACY = float(os.getenv("EXPECTED_ACCURACY", "0.8"))
+ACCURACY_TOL = float(os.getenv("ACCURACY_TOL", "0.08"))
+
+
+def _run_gsm8k_eval(server, stage: str) -> float:
+    """Run the gsm8k probe; assert the absolute bar, return the accuracy."""
+    from tests.evals.gsm8k.gsm8k_eval import evaluate_gsm8k
+
+    host_port = server.url_for("v1").split("//")[1].split("/")[0]
+    host, port_str = host_port.rsplit(":", 1)
+    results = evaluate_gsm8k(
+        num_questions=NUM_GSM8K_QUESTIONS,
+        num_shots=5,
+        host=f"http://{host}",
+        port=int(port_str),
+        temperature=0.0,
+        seed=42,
+    )
+    accuracy = results["accuracy"]
+    print(
+        f"[{stage}] GSM8K accuracy: {accuracy:.3f} "
+        f"invalid_rate={results['invalid_rate']:.3f} "
+        f"({results['num_questions']} questions)"
+    )
+    assert accuracy >= EXPECTED_ACCURACY, (
+        f"[{stage}] gsm8k accuracy {accuracy:.3f} < {EXPECTED_ACCURACY}"
+    )
+    return accuracy
+
+
 def _in_parallel(fn, servers) -> list:
     """Run ``fn(server)`` for all servers concurrently; return results in order."""
     with ThreadPoolExecutor(max_workers=len(servers)) as ex:
@@ -163,7 +220,7 @@ def _in_parallel(fn, servers) -> list:
 
 
 def _get_ft_status(server) -> dict:
-    resp = requests.get(server.url_for("fault_tolerance/status"), timeout=10)
+    resp = requests.get(server.url_for("v1/fault_tolerance/status"), timeout=10)
     resp.raise_for_status()
     return resp.json()
 
@@ -177,17 +234,37 @@ def _assert_serving_and_healthy(servers) -> None:
     _in_parallel(lambda s: _complete(s.get_client()), servers)
 
 
-def _apply_ft(server, instruction: str, params: dict | None = None) -> dict:
+def _apply_ft(
+    server,
+    instruction: str,
+    params: dict | None = None,
+    request_id: str | None = None,
+) -> dict:
     """POST an FT instruction; assert it is accepted (202) and return the body."""
     resp = requests.post(
-        server.url_for("fault_tolerance/apply"),
-        json={"instruction": instruction, "params": params or {}},
+        server.url_for("v1/fault_tolerance/apply"),
+        json={
+            "instruction": instruction,
+            "params": params or {},
+            "request_id": request_id or str(uuid.uuid4()),
+        },
         timeout=10,
     )
     assert resp.status_code == 202, resp.text
     return resp.json()
 
 
+def _arm_ft_injection(servers) -> None:
+    """Arm the fault injection in every engine's workers via /collective_rpc."""
+    for server in servers:
+        resp = requests.post(
+            server.url_for("collective_rpc"),
+            json={"method": "arm_ft_fault_injection"},
+            timeout=10,
+        )
+        assert resp.status_code == 200, resp.text
+
+
 def _kill_worker_process(server) -> None:
     """SIGKILL only the worker proc, leaving EngineCore and API server alive."""
     workers = []
@@ -207,7 +284,7 @@ def _wait_for_engines(
     match_values: set[str],
     deadline_s: int = FAULT_DETECTION_DEADLINE_S,
 ) -> list[dict[str, Any] | None]:
-    """Poll ``/fault_tolerance/status`` until each server's engine status matches.
+    """Poll ``/v1/fault_tolerance/status`` until each server's engine status matches.
 
     A server matches when its engine-status dict has ``match_key`` equal to
     one of ``match_values``. Returns one engine-status dict per server. Servers still
@@ -256,125 +333,201 @@ def _drive(server):
             t.join(timeout=2)
 
 
-def _wait_for_ft_apply_outcome(server, request_id: str, deadline_s: int) -> str | None:
-    """Wait until ``/fault_tolerance/status`` records the FT apply outcome."""
+def _wait_for_ft_failure(server, deadline_s: int) -> str | None:
+    """Wait until ``/v1/fault_tolerance/status`` records the FT failure outcome."""
     engine_status = _wait_for_engines(
         [server],
-        match_key="last_ft_request_id",
-        match_values={request_id},
+        match_key="ft_state",
+        match_values={"failed"},
         deadline_s=deadline_s,
     )[0]
     return engine_status.get("ft_error") if engine_status else None
 
 
+def _servers_by_rank(servers, dp_size: int) -> dict[int, RemoteOpenAIServer]:
+    """Map each DP rank to its server, so tests can index by rank not position."""
+    return {r: _server_for_rank(servers, r) for r in range(dp_size)}
+
+
+def _drive_to_faulted(
+    servers_by_rank: dict[int, RemoteOpenAIServer],
+    match_values: set[str],
+) -> dict[int, dict[str, Any]]:
+    """Drive every engine until each reports a matching fault status.
+
+    Returns a ``{rank: engine_status}`` map. Asserts no rank hung (all matched
+    within the deadline) so callers can assume every value is present.
+    """
+    ranks = sorted(servers_by_rank)
+    all_servers = [servers_by_rank[r] for r in ranks]
+    with _driving(*all_servers):
+        statuses = _wait_for_engines(
+            all_servers, match_key="status", match_values=match_values
+        )
+    faulted: dict[int, dict[str, Any]] = {}
+    for rank, engine_status in zip(ranks, statuses):
+        assert engine_status is not None, (
+            f"rank {rank} did not report fault within "
+            f"{FAULT_DETECTION_DEADLINE_S}s -- it likely hung"
+        )
+        faulted[rank] = engine_status
+    return faulted
+
+
 @pytest.mark.skipif(not has_nixl_ep(), reason="Requires nixl_ep all2all backend")
 @multi_gpu_test(num_gpus=2)
 def test_injected_fault_retry_recovers_all_ranks(monkeypatch, tmp_path):
     """An exception injected into the inference path drives full retry recovery.
 
-    Injecting an exception into ``sync_cudagraph_and_dp_padding`` at a chosen
-    step on rank 1.
+    The fault is patched into ``sync_cudagraph_and_dp_padding`` on rank 1 via
+    a generated ``sitecustomize``; it raises on the engine's first DP-sync
+    step after arming, so rank 1 goes UNHEALTHY and rank 0 times out on the
+    DP allreduce. Both UNHEALTHY is the precondition for ``retry``.
 
-    - Rank 1 raises inside the busy loop and goes UNHEALTHY.
-    - Rank 0 detects the now-absent peer via the communication timeout and also
-      goes UNHEALTHY.
-
-    Both being UNHEALTHY is the precondition for ``retry``. The fault is patched
-    into the DP-sync fn from the test (via a generated ``sitecustomize``).
+    The injection is armed (via /collective_rpc) only after the baseline
+    accuracy eval; accuracy is compared before and after recovery, mirroring
+    test_elastic_ep.py.
     """
-    fault_step = int(os.getenv("FT_FAULT_STEP", "50"))
-    _install_fault_injection(monkeypatch, tmp_path, rank=1, step=fault_step)
-
-    with _ft_manager() as servers:
-        assert len(servers) == DP_SIZE
-        rank0 = _server_for_rank(servers, 0)
-        rank1 = _server_for_rank(servers, 1)
-
-        # 1. Both engines healthy and serving.
-        _assert_serving_and_healthy((rank0, rank1))
-
-        # 2. Drive both ranks so rank 1 accumulates execute_model steps and trips
-        #    the injected fault; rank 0 then times out on the DP allreduce.
-        with _driving(rank0, rank1):
-            faulted = _wait_for_engines(
-                [rank0, rank1], match_key="status", match_values={"unhealthy"}
-            )
+    _install_fault_injection(monkeypatch, tmp_path, rank=1)
+    dp_size = 2
+    faulted_rank = 1
+    with _ft_manager(dp_size=dp_size) as servers:
+        assert len(servers) == dp_size
+        servers_by_rank = _servers_by_rank(servers, dp_size)
+
+        # 1. All engines healthy and serving.
+        _assert_serving_and_healthy(list(servers_by_rank.values()))
+
+        # 2. Baseline accuracy (the injection is still disarmed).
+        initial_accuracy = _run_gsm8k_eval(servers_by_rank[0], "Initial")
+
+        # 3. Arm, then drive: the injected rank faults on its next engine step.
+        _arm_ft_injection(list(servers_by_rank.values()))
+        faulted = _drive_to_faulted(servers_by_rank, match_values={"unhealthy"})
 
-        for rank, engine_status in enumerate(faulted):
-            assert engine_status is not None, (
-                f"rank {rank} did not report UNHEALTHY within "
-                f"{FAULT_DETECTION_DEADLINE_S}s -- it likely hung"
-            )
         # The rank that raised carries the fault info from its own exception.
-        assert faulted[1] is not None
-        assert faulted[1].get("fault_info"), faulted[1]
+        assert faulted[faulted_rank].get("fault_info"), faulted[faulted_rank]
 
-        # 3. retry both engines.
-        for server in (rank0, rank1):
-            _apply_ft(server, "retry")
+        # 4. retry every engine.
+        round_id = str(uuid.uuid4())
+        for server in servers_by_rank.values():
+            _apply_ft(server, "retry", request_id=round_id)
 
-        # 4. Recovery completes: both engines return to healthy and serve again.
-        _assert_serving_and_healthy((rank0, rank1))
+        # 5. Recovery completes: every engine returns to healthy and serves again.
+        _assert_serving_and_healthy(list(servers_by_rank.values()))
+
+        # 6. Post-recovery accuracy vs the pre-fault baseline.
+        recovered_accuracy = _run_gsm8k_eval(servers_by_rank[0], "After retry")
+        assert recovered_accuracy >= initial_accuracy - ACCURACY_TOL, (
+            f"Accuracy dropped after retry: {recovered_accuracy:.3f} < "
+            f"{initial_accuracy:.3f} - {ACCURACY_TOL}"
+        )
+        print(
+            "\nAccuracy Summary:\n"
+            f"  Initial:     {initial_accuracy:.3f}\n"
+            f"  After retry: {recovered_accuracy:.3f} "
+            f"(diff: {recovered_accuracy - initial_accuracy:+.3f})\n"
+            f"  Tolerance:   {ACCURACY_TOL:.3f}"
+        )
 
 
 @pytest.mark.skipif(not has_nixl_ep(), reason="Requires nixl_ep all2all backend")
-@multi_gpu_test(num_gpus=2)
-def test_worker_kill_survivor_unhealthy_and_dead_rejects_retry():
-    """One worker kill surfaces two status transitions at once.
+@multi_gpu_test(num_gpus=4)
+@pytest.mark.parametrize("quantization", [None, "fp8_per_block"])
+def test_scale_down_removes_dead_rank_and_recovers(quantization):
+    """Preserve physical expert IDs and reload recovered placement after scale-down.
 
-    SIGKILLing only rank 1's worker leaves both EngineCores alive, so the same
-    fault is seen two ways:
+    Survivors keep serving with a smaller DP group, including processed FP8
+    weights whose checkpoint layout differs from their runtime layout.
 
-    - Survivor (rank 0): detects the dead peer via Gloo allreduce / nixl_ep
-      kernel timeout. Its own executor is fine, so ``on_fault`` marks it
-      UNHEALTHY with a ``fault_info``.
-    - Victim (rank 1): detects its own executor failure and marks itself DEAD.
+    SIGKILLing rank 0's worker leaves its EngineCore DEAD while the other ranks
+    detect the peer fault and go UNHEALTHY. The orchestrator then sends
+    ``scale_down`` with ``removed_dp_ranks=[0]`` to every survivor.
 
-    Recovery is gated on UNHEALTHY: the DEAD engine accepts ``retry`` at the
-    HTTP layer (202 = background dispatch) but rejects it in the engine,
-    recording the reason as ``ft_error``.
-    """
-    with _ft_manager() as servers:
-        assert len(servers) == DP_SIZE
-        survivor = _server_for_rank(servers, 0)
-        victim = _server_for_rank(servers, 1)
 
-        # 1. Confirm both engines are healthy and serving.
-        _assert_serving_and_healthy((survivor, victim))
+    Also verifies that a DEAD engine rejects ``retry``: recovery is gated on
+    UNHEALTHY, so trying ``retry`` on the victim records a rejection reason.
 
-        # 2. Kill only the victim's worker; both EngineCores stay alive.
+    Accuracy is measured before the fault and after scale_down; the survivor
+    must stay within ACCURACY_TOL of the pre-fault baseline.
+    """
+    eplb_args = [
+        "--enable-eplb",
+        "--eplb-config.num_redundant_experts",
+        str(NUM_REDUNDANT_EXPERTS),
+    ]
+    if quantization is not None:
+        eplb_args.extend(["--quantization", quantization])
+    dp_size = 4
+    victim_rank = 0
+    with _ft_manager(eplb_args, dp_size=dp_size) as servers:
+        assert len(servers) == dp_size
+        servers_by_rank = _servers_by_rank(servers, dp_size)
+        victim = servers_by_rank[victim_rank]
+        survivor_ranks = [r for r in servers_by_rank if r != victim_rank]
+        survivors = [servers_by_rank[r] for r in survivor_ranks]
+
+        # 1. All engines healthy and serving.
+        _assert_serving_and_healthy(list(servers_by_rank.values()))
+
+        # 2. Baseline accuracy while the full DP group is healthy.
+        initial_accuracy = _run_gsm8k_eval(servers_by_rank[0], "Initial")
+
+        # 3. Kill the victim's worker; drive all engines into the fault.
         _kill_worker_process(victim)
+        faulted = _drive_to_faulted(servers_by_rank, match_values={"dead", "unhealthy"})
 
-        # 3. Drive both engines so each keeps stepping into the failed component.
-        with _driving(survivor, victim):
-            survivor_faulted, victim_faulted = _wait_for_engines(
-                [survivor, victim],
-                match_key="status",
-                match_values={"dead", "unhealthy"},
-            )
-
-        assert survivor_faulted is not None, (
-            "survivor did not report the peer fault within "
-            f"{FAULT_DETECTION_DEADLINE_S}s -- it likely hung"
-        )
-        # The survivor's own executor is fine, so it must be UNHEALTHY, not DEAD.
-        assert survivor_faulted["status"] == "unhealthy", survivor_faulted
-        assert survivor_faulted.get("fault_info"), survivor_faulted
-
-        assert victim_faulted is not None, (
-            "victim did not report its worker's death within "
-            f"{FAULT_DETECTION_DEADLINE_S}s"
+        # 4. DEAD engine rejects retry: recovery requires UNHEALTHY.
+        assert faulted[victim_rank]["status"] == "dead", faulted[victim_rank]
+        _apply_ft(victim, "retry")
+        ft_error = _wait_for_ft_failure(victim, FAULT_DETECTION_DEADLINE_S)
+        assert ft_error is not None, (
+            "rejection was never recorded in /v1/fault_tolerance/status"
         )
-        assert victim_faulted["status"] == "dead", victim_faulted
+        assert "status is DEAD" in ft_error, ft_error
 
-        # 4. retry is accepted at the HTTP layer (202 = background dispatch)...
-        request_id = _apply_ft(victim, "retry")["request_id"]
+        # 5. scale_down sent to every survivor: remove the dead rank.
+        round_id = str(uuid.uuid4())
+        store_port = get_open_port()
+        for server in survivors:
+            _apply_ft(
+                server,
+                "scale_down",
+                {
+                    "removed_dp_ranks": [victim_rank],
+                    "dp_store_port": store_port,
+                    "dp_master_ip": "127.0.0.1",
+                },
+                request_id=round_id,
+            )
 
-        # 5. ...but the DEAD engine must reject it: recovery requires UNHEALTHY.
-        ft_error = _wait_for_ft_apply_outcome(
-            victim, request_id, FAULT_DETECTION_DEADLINE_S
+        # 6. Recovery completes: all survivors are healthy and serving.
+        recovered = _wait_for_engines(
+            survivors,
+            match_key="status",
+            match_values={"healthy"},
+            deadline_s=SCALE_DOWN_DEADLINE_S,
         )
-        assert ft_error is not None, (
-            "rejection was never recorded in /fault_tolerance/status"
+        for rank, engine_status in zip(survivor_ranks, recovered):
+            assert engine_status is not None, (
+                f"survivor {rank} did not recover within {SCALE_DOWN_DEADLINE_S}s "
+                "-- expert reload or DP-group reinit likely failed"
+            )
+        # Verify all survivors actually serve.
+        for server in survivors:
+            completion = _complete(server.get_client())
+            assert completion.choices[0].text
+
+        # 7. Post-recovery accuracy vs the pre-fault baseline.
+        recovered_accuracy = _run_gsm8k_eval(survivors[0], "After scale down")
+        assert recovered_accuracy >= initial_accuracy - ACCURACY_TOL, (
+            f"Accuracy dropped after scale_down: {recovered_accuracy:.3f} < "
+            f"{initial_accuracy:.3f} - {ACCURACY_TOL}"
+        )
+        print(
+            "\nAccuracy Summary:\n"
+            f"  Initial:          {initial_accuracy:.3f}\n"
+            f"  After scale down: {recovered_accuracy:.3f} "
+            f"(diff: {recovered_accuracy - initial_accuracy:+.3f})\n"
+            f"  Tolerance:        {ACCURACY_TOL:.3f}"
         )
-        assert "status is DEAD" in ft_error, ft_error
diff --git a/tests/v1/kv_connector/conftest.py b/tests/v1/kv_connector/conftest.py
new file mode 100644
index 000000000000..1938f63c3993
--- /dev/null
+++ b/tests/v1/kv_connector/conftest.py
@@ -0,0 +1,13 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import tempfile
+
+import pytest
+
+
+@pytest.fixture
+def socket_dir():
+    """Keep AF_UNIX paths short regardless of CI's TMPDIR."""
+    with tempfile.TemporaryDirectory(prefix="umbp-", dir="/tmp") as directory:
+        yield directory
diff --git a/tests/v1/kv_connector/umbp_integration/__init__.py b/tests/v1/kv_connector/umbp_integration/__init__.py
new file mode 100644
index 000000000000..e69de29bb2d1
diff --git a/tests/v1/kv_connector/umbp_integration/test_mori_embedded_gpu.py b/tests/v1/kv_connector/umbp_integration/test_mori_embedded_gpu.py
new file mode 100644
index 000000000000..857a2a013836
--- /dev/null
+++ b/tests/v1/kv_connector/umbp_integration/test_mori_embedded_gpu.py
@@ -0,0 +1,115 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""MORI-backed EmbeddedRuntime GPU integration tests."""
+
+import pytest
+import torch
+
+from tests.v1.kv_connector.umbp_test_utils import mori_has_umbp
+from vllm.distributed.kv_transfer.kv_connector.v1.umbp.data import (
+    BlockTransferPlan,
+    KVLayoutDescriptor,
+    KVRange,
+    KVRegion,
+    RankTopology,
+    TransferJobStatus,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.umbp.runtime import (
+    EmbeddedRuntime,
+    UMBPRuntimeConfig,
+)
+from vllm.platforms import current_platform
+
+pytestmark = pytest.mark.skipif(
+    not (current_platform.is_rocm() and mori_has_umbp()),
+    reason="requires ROCm and MORI built with UMBP",
+)
+
+
+@pytest.mark.parametrize("bulk_load", [False, True])
+def test_mori_embedded_dram_eviction_and_restore(socket_dir, bulk_load):
+    """Overflowing DRAM evicts old objects, and the newest one still restores."""
+    object_size = 1 << 20
+    num_objects = 16
+    runtime = EmbeddedRuntime.from_config(
+        UMBPRuntimeConfig(
+            "embedded",
+            {
+                "capacity_bytes": 8 * object_size,
+                "dram_high_watermark": 0.75,
+                "dram_low_watermark": 0.5,
+                "lookup_dir": socket_dir,
+                "timeout_ms": 60000,
+            },
+        )
+    )
+    worker = runtime.create_worker_handle(
+        "mori-eviction",
+        RankTopology(),
+        KVLayoutDescriptor(
+            regions=(KVRegion("layer0", 0, object_size, object_size, 0),),
+            topology=RankTopology(),
+        ),
+    )
+    try:
+        source = torch.empty((1, object_size), dtype=torch.uint8, device="cuda")
+        destination = torch.zeros_like(source)
+        worker.register_buffers({"layer0": source, "destination": destination})
+
+        keys = [f"eviction-{index}" for index in range(num_objects)]
+        for index, key in enumerate(keys):
+            source.fill_(index)
+            plan = BlockTransferPlan(
+                key,
+                index,
+                ranges=(
+                    KVRange(
+                        "layer0",
+                        0,
+                        index,
+                        source.data_ptr(),
+                        object_size,
+                        object_size,
+                        0,
+                    ),
+                ),
+            )
+            stored = worker.wait(worker.store([plan]))
+            if stored.status is TransferJobStatus.COMPLETED:
+                worker.publish(stored)
+
+        hits = list(worker.batch_exists(keys))
+        assert not all(hits), "the stores never exceeded DRAM capacity"
+        latest = max(index for index, hit in enumerate(hits) if hit)
+
+        if bulk_load:
+            worker.register_buffers({"layer0": destination})
+            job = worker.load_blocks([BlockTransferPlan(keys[latest], 0, group_id=0)])
+            assert job is not None
+        else:
+            job = worker.load(
+                [
+                    BlockTransferPlan(
+                        keys[latest],
+                        0,
+                        ranges=(
+                            KVRange(
+                                "layer0",
+                                0,
+                                0,
+                                destination.data_ptr(),
+                                object_size,
+                                object_size,
+                                0,
+                            ),
+                        ),
+                    )
+                ]
+            )
+        loaded = worker.wait(job)
+        assert loaded.status is TransferJobStatus.COMPLETED
+        assert loaded.completed_bytes == object_size
+        torch.accelerator.synchronize()
+        assert torch.all(destination == latest)
+    finally:
+        worker.close()
diff --git a/tests/v1/kv_connector/umbp_integration/test_mori_token_e2e.py b/tests/v1/kv_connector/umbp_integration/test_mori_token_e2e.py
new file mode 100644
index 000000000000..b14467c44f16
--- /dev/null
+++ b/tests/v1/kv_connector/umbp_integration/test_mori_token_e2e.py
@@ -0,0 +1,128 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""End-to-end restore through the MORI-backed UMBP runtime.
+
+Batch-invariant kernels make a restore bitwise comparable with recomputing the
+prompt, so restored KV must reproduce the cold run's tokens and logprobs.
+"""
+
+import time
+
+import pytest
+
+import vllm.envs as envs
+from tests.v1.kv_connector.rocm_pd_accuracy_utils import _REVISIONS, QWEN_MODEL
+from tests.v1.kv_connector.umbp_test_utils import mori_has_umbp
+from vllm import SamplingParams, TokensPrompt
+from vllm.config import KVTransferConfig
+from vllm.platforms import current_platform
+
+pytestmark = pytest.mark.skipif(
+    not (current_platform.is_rocm() and mori_has_umbp()),
+    reason="requires ROCm and MORI built with UMBP",
+)
+
+BLOCK_SIZE = 16
+PARAMS = SamplingParams(temperature=0, max_tokens=16, logprobs=5, ignore_eos=True)
+
+
+@pytest.fixture
+def umbp_llm(vllm_runner, socket_dir, monkeypatch):
+    monkeypatch.setattr(envs, "VLLM_BATCH_INVARIANT", True)
+    monkeypatch.setenv("VLLM_BATCH_INVARIANT", "1")
+    with vllm_runner(
+        QWEN_MODEL,
+        revision=_REVISIONS[QWEN_MODEL],
+        enforce_eager=True,
+        block_size=BLOCK_SIZE,
+        enable_prefix_caching=True,
+        disable_log_stats=False,
+        kv_transfer_config=KVTransferConfig(
+            kv_connector="UMBPStoreConnector",
+            kv_role="kv_both",
+            kv_connector_extra_config={
+                "mode": "embedded",
+                "capacity_bytes": 1 << 30,
+                "lookup_dir": socket_dir,
+            },
+        ),
+    ) as runner:
+        yield runner.llm
+
+
+def _prompt(seed: int, num_tokens: int) -> TokensPrompt:
+    return TokensPrompt(
+        prompt_token_ids=[
+            100 + (seed * 1009 + i * 37) % 5000 for i in range(num_tokens)
+        ]
+    )
+
+
+def _external_tokens(llm) -> int:
+    return sum(
+        metric.value
+        for metric in llm.get_metrics()
+        if metric.name == "vllm:prompt_tokens_by_source"
+        and metric.labels.get("source") == "external_kv_transfer"
+    )
+
+
+def _reset_gpu_prefix_cache(llm) -> None:
+    """Drop the GPU prefix cache once UMBP has released the stored blocks."""
+    deadline = time.monotonic() + 60
+    while not llm.reset_prefix_cache():
+        assert time.monotonic() < deadline, "UMBP stores never released GPU blocks"
+        llm.generate(
+            TokensPrompt(prompt_token_ids=[0]),
+            SamplingParams(max_tokens=1),
+            use_tqdm=False,
+        )
+
+
+def _assert_same_outputs(expected, actual) -> None:
+    for want, got in zip(expected, actual, strict=True):
+        want_out, got_out = want.outputs[0], got.outputs[0]
+        assert got_out.token_ids == want_out.token_ids
+        assert [
+            {token: entry.logprob for token, entry in step.items()}
+            for step in got_out.logprobs
+        ] == [
+            {token: entry.logprob for token, entry in step.items()}
+            for step in want_out.logprobs
+        ]
+
+
+def test_restore_matches_recompute(umbp_llm):
+    """A batch of prompts ending mid-block restores every full block exactly."""
+    prompts = [
+        _prompt(seed, 4 * BLOCK_SIZE + tail)
+        for seed, tail in enumerate((1, 7, BLOCK_SIZE - 1))
+    ]
+    computed = umbp_llm.generate(prompts, PARAMS, use_tqdm=False)
+    _reset_gpu_prefix_cache(umbp_llm)
+
+    before = _external_tokens(umbp_llm)
+    restored = umbp_llm.generate(prompts, PARAMS, use_tqdm=False)
+
+    assert _external_tokens(umbp_llm) - before == 3 * 4 * BLOCK_SIZE
+    _assert_same_outputs(computed, restored)
+
+
+def test_restore_extends_a_local_prefix_hit(umbp_llm):
+    """UMBP fills in only the blocks after a shorter GPU prefix-cache hit."""
+    prompt = _prompt(0, 8 * BLOCK_SIZE + 7)
+    computed = umbp_llm.generate(prompt, PARAMS, use_tqdm=False)
+    _reset_gpu_prefix_cache(umbp_llm)
+    short = PARAMS.clone()
+    short.skip_reading_prefix_cache = True
+    local_prefix = TokensPrompt(
+        prompt_token_ids=prompt["prompt_token_ids"][: 2 * BLOCK_SIZE + 1]
+    )
+    umbp_llm.generate(local_prefix, short, use_tqdm=False)
+
+    before = _external_tokens(umbp_llm)
+    mixed = umbp_llm.generate(prompt, PARAMS, use_tqdm=False)
+
+    assert mixed[0].num_cached_tokens == 8 * BLOCK_SIZE
+    assert _external_tokens(umbp_llm) - before == 6 * BLOCK_SIZE
+    _assert_same_outputs(computed, mixed)
diff --git a/tests/v1/kv_connector/umbp_test_utils.py b/tests/v1/kv_connector/umbp_test_utils.py
index 4f7ca892e68f..0c79f6d63045 100644
--- a/tests/v1/kv_connector/umbp_test_utils.py
+++ b/tests/v1/kv_connector/umbp_test_utils.py
@@ -225,6 +225,17 @@ def create_worker_handle(
         del namespace
         return _MemoryWorkerHandle(self._store, topology, layout)
 
+    def evict_all(self) -> None:
+        self._store.clear()
+
+
+def mori_has_umbp() -> bool:
+    try:
+        from mori.cpp import UMBPClient  # noqa: F401
+    except ImportError:
+        return False
+    return True
+
 
 def install_memory_runtime(monkeypatch):
     """Share one isolated store across the connectors created by a test."""
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 45c5711306d1..58c0c76ca81d 100644
--- a/tests/v1/kv_connector/unit/offloading_connector/test_events.py
+++ b/tests/v1/kv_connector/unit/offloading_connector/test_events.py
@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import gc
 from types import SimpleNamespace
 from unittest.mock import MagicMock
 
@@ -29,7 +30,9 @@
     FullAttentionSpec,
     KVCacheConfig,
     KVCacheGroupSpec,
+    KVCacheSpec,
     KVCacheSpecKind,
+    MambaSpec,
 )
 from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry
 from vllm.v1.kv_offload.base import (
@@ -50,6 +53,10 @@
     kv_cache_spec_kind=KVCacheSpecKind.FULL_ATTENTION.value,
     kv_cache_spec_sliding_window=None,
 )
+_MAMBA_EVENT_SPEC = OffloadingEventGroupSpec(
+    kv_cache_spec_kind=KVCacheSpecKind.MAMBA.value,
+    kv_cache_spec_sliding_window=None,
+)
 
 
 def _tracker(
@@ -89,17 +96,31 @@ def _group_config(
     blocks_per_chunk: int = 1,
     tokens_per_hash: int | None = None,
     sliding_window_size_in_chunks: int | None = None,
+    mamba: bool = False,
 ) -> GroupOffloadConfig:
     if tokens_per_hash is None:
         tokens_per_hash = block_size
     tokens_per_chunk = block_size * blocks_per_chunk
     assert tokens_per_chunk % tokens_per_hash == 0
-    kv_spec = FullAttentionSpec(
-        block_size=block_size,
-        num_kv_heads=1,
-        head_size=1,
-        dtype=torch.float32,
-    )
+    kv_spec: KVCacheSpec
+    if mamba:
+        # A recurrent group: one state per chunk, offloaded in align mode.
+        kv_spec = MambaSpec(
+            block_size=block_size,
+            shapes=((1, 1),),
+            dtypes=(torch.float32,),
+            mamba_cache_mode="align",
+        )
+        sliding_window_size_in_chunks = 1
+        event_spec = _MAMBA_EVENT_SPEC
+    else:
+        kv_spec = FullAttentionSpec(
+            block_size=block_size,
+            num_kv_heads=1,
+            head_size=1,
+            dtype=torch.float32,
+        )
+        event_spec = _FULL_ATTENTION_EVENT_SPEC
     manager_cls = KVCacheSpecRegistry.get_manager_class(kv_spec)
     assert manager_cls is not None
     return GroupOffloadConfig(
@@ -110,7 +131,8 @@ def _group_config(
         sliding_window_size_in_chunks=sliding_window_size_in_chunks,
         kv_cache_spec=kv_spec,
         manager_cls=manager_cls,
-        kv_event_group_spec=_FULL_ATTENTION_EVENT_SPEC,
+        kv_event_group_spec=event_spec,
+        requires_cow_source=mamba,
     )
 
 
@@ -515,12 +537,84 @@ def test_take_events_opt_out_keeps_placeholders():
     assert len(events[3].block_hashes) == 3
 
 
+def test_mamba_chunk_store_emits_state_checkpoint_event():
+    """A recurrent group's chunk is announced as one state checkpoint: the
+    chunk's last hash, chained to its predecessor, with no tokens and
+    kv_cache_spec_kind "mamba", so a consumer aligns it by hash."""
+    tracker = _tracker()
+    group_config = _group_config(
+        group_idx=1, block_size=8, tokens_per_hash=4, mamba=True
+    )
+    req = _request(block_hashes=[_hash(i) for i in range(4)], token_count=16)
+    keys = _record_chunks(tracker, req, group_config, num_chunks=2)
+
+    events = list(tracker.take_events([_stored_event(keys[1:])]))
+
+    assert len(events) == 1
+    event = events[0]
+    assert isinstance(event, BlockStored)
+    assert event.block_hashes == [_wire_hash(_hash(3))]
+    assert event.parent_block_hash == _wire_hash(_hash(1))
+    assert event.token_ids == []
+    assert event.block_size == 8
+    assert event.group_idx == 1
+    assert event.kv_cache_spec_kind == KVCacheSpecKind.MAMBA.value
+    assert event.kv_cache_spec_sliding_window is None
+
+    [removed] = tracker.take_events([_removed_event(keys[1:])])
+    assert isinstance(removed, BlockRemoved)
+    assert removed.block_hashes == [_wire_hash(_hash(3))]
+    assert removed.group_idx == 1
+
+
+def test_mamba_multi_hash_chunk_announces_only_last_hash():
+    """In chunk mode a recurrent chunk spans several hashes, but only the
+    last one is a resumable checkpoint; the others are not announced."""
+    tracker = _tracker()
+    group_config = _group_config(block_size=4, blocks_per_chunk=4, mamba=True)
+    req = _request(block_hashes=[_hash(i) for i in range(8)], token_count=32)
+    keys = _record_chunks(tracker, req, group_config, num_chunks=2)
+
+    [event] = tracker.take_events([_stored_event(keys[1:])])
+    assert isinstance(event, BlockStored)
+    assert event.block_hashes == [_wire_hash(_hash(7))]
+    assert event.parent_block_hash == _wire_hash(_hash(6))
+    assert event.token_ids == []
+    assert event.block_size == 4
+
+    [removed] = tracker.take_events([_removed_event(keys[1:])])
+    assert isinstance(removed, BlockRemoved)
+    assert removed.block_hashes == [_wire_hash(_hash(7))]
+
+
+@pytest.mark.parametrize(
+    "record_method", ["record_partial_store", "record_partial_lookup"]
+)
+def test_mamba_partial_tail_event_is_state_checkpoint(record_method):
+    tracker = _tracker()
+    group_config = _group_config(
+        block_size=16, blocks_per_chunk=1, mamba=True
+    )._replace(hashes_per_chunk=4)
+    req = _request(block_hashes=[_hash(i) for i in range(8)], token_count=32)
+    key = make_offload_key(req.block_hashes[6], group_config.group_idx)
+
+    getattr(tracker, record_method)(req, group_config, 28, key)
+    [event] = tracker.take_events([_stored_event([key])])
+
+    assert isinstance(event, BlockStored)
+    assert event.block_hashes == [_wire_hash(_hash(6))]
+    assert event.parent_block_hash == _wire_hash(_hash(5))
+    assert event.token_ids == []
+    assert event.block_size == 4
+    assert event.kv_cache_spec_kind == KVCacheSpecKind.MAMBA.value
+
+
 @pytest.mark.parametrize(
     "sliding_window_size_in_chunks",
     [1, 2],
-    ids=["ssm", "sliding-window"],
+    ids=["one-chunk-window", "sliding-window"],
 )
-def test_event_metadata_skips_non_full_attention_group(
+def test_event_metadata_skips_windowed_attention_group(
     sliding_window_size_in_chunks: int,
 ):
     tracker = _tracker()
@@ -773,17 +867,16 @@ def test_reoffload_preserves_secondary_residency(record_method, position):
 
     getattr(tracker, record_method)(req, group_config, position, key)
 
-    assert tracker._pending_event_metadata[key].active_residencies == {
-        (Medium.CPU, None),
-        (Medium.STORAGE, "custom"),
-    }
+    assert not gc.is_tracked(tracker._pending_event_metadata[key])
+    list(tracker.take_events([_removed_event([key])]))
+    assert key in tracker._pending_event_metadata
     [removed] = tracker.take_events(
         [_removed_event([key], Medium.STORAGE, ownership="custom")]
     )
 
     assert stored.token_ids == [1, 2, 3, 4]
     assert stored.ownership == removed.ownership == "custom"
-    assert key in tracker._pending_event_metadata
+    assert key not in tracker._pending_event_metadata
 
 
 def test_take_events_groups_removed_hashes_by_kv_group():
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 1573101b244a..bd15c5562755 100644
--- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py
+++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py
@@ -24,6 +24,8 @@
     create_scheduler,
     create_vllm_config,
 )
+from tests.v1.spec_decode.utils import attach_speculator_config
+from vllm import SamplingParams
 from vllm.config import KVEventsConfig, KVTransferConfig
 from vllm.distributed.kv_events import MEDIUM_CPU, BlockRemoved, BlockStored
 from vllm.distributed.kv_transfer.kv_connector.v1.offloading import (
@@ -49,12 +51,15 @@
 from vllm.distributed.kv_transfer.kv_connector.v1.offloading_connector import (
     OffloadingConnector,
 )
+from vllm.utils.hashing import sha256
 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 (
     BlockHash,
     KVCacheBlock,
     KVCacheBlockCopy,
+    get_request_block_hasher,
+    init_none_hash,
     make_block_hash_with_group_id,
 )
 from vllm.v1.core.sched.output import (
@@ -93,10 +98,12 @@
 )
 from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager
 from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion
+from vllm.v1.kv_offload.cpu.spec import CPUOffloadingSpec
 from vllm.v1.kv_offload.tiering.example.manager import ExampleSecondaryTierManager
 from vllm.v1.metrics.cache_hit_source import CacheHitSource
 from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput
-from vllm.v1.request import RequestStatus
+from vllm.v1.request import Request, RequestStatus
+from vllm.v1.worker.gpu.spec_decode.mtp.config import MTPSpeculatorConfig
 
 
 def _as_mock(value: object) -> MagicMock:
@@ -104,6 +111,39 @@ def _as_mock(value: object) -> MagicMock:
     return value
 
 
+def test_update_offload_keys_without_scanning_hash_prefix():
+    """Incremental key updates must not iterate over the earlier request hashes."""
+
+    class IndexOnlyHashes(list[BlockHash]):
+        def __iter__(self):
+            raise AssertionError("Request hashes must be accessed by index")
+
+    config = MagicMock()
+    config.kv_group_configs = [
+        SimpleNamespace(
+            group_idx=group_idx, hashes_per_chunk=stride, tokens_per_chunk=16 * stride
+        )
+        for group_idx, stride in enumerate((1, 2, 4))
+    ]
+    request = MagicMock(kv_transfer_params=None, block_hashes=IndexOnlyHashes())
+    context = ReqContext("test")
+    state = RequestOffloadState(config, request, context, RequestOffloadingContext())
+
+    for num_hashes in (0, 1, 4, 7, 7, 8):
+        request.block_hashes.extend(
+            BlockHash(bytes([i])) for i in range(len(request.block_hashes), num_hashes)
+        )
+        state.update_offload_keys()
+        for group_idx, stride in enumerate((1, 2, 4)):
+            expected = [
+                make_offload_key(bytes([i]), group_idx)
+                for i in range(stride - 1, num_hashes, stride)
+            ]
+            assert state.group_states[group_idx].offload_keys == expected
+            for chunk_idx, key in enumerate(expected, 1):
+                assert context.get_offload_key_position(key) == chunk_idx * stride * 16
+
+
 @pytest.mark.parametrize(
     "window,extra_retained", [(128, 0), (128, 3), (256, 1), (256, 3)]
 )
@@ -222,7 +262,7 @@ def test_swa_offload_window_covers_unaligned_hit(
     assert len(sched._jobs) == 1
     job = next(iter(sched._jobs.values()))
     if needs_left:
-        assert {left_key, swa[end - 1]} <= job.keys
+        assert {left_key, swa[end - 1]} <= set(job.keys)
     manager.complete_load(job.keys, state.req_context)
 
 
@@ -230,6 +270,7 @@ def _make_partial_tail_scheduler(
     recurrent_only: bool = False,
     eagle: bool = False,
     cache_hit_alignment_tokens: int | None = 4,
+    manager: OffloadingManager | None = None,
 ) -> OffloadingConnectorScheduler:
     vllm_config = _make_vllm_config(extra_config={"self_describing_kv_events": True})
     vllm_config.cache_config.prefix_match_unit = None if recurrent_only else 4
@@ -245,6 +286,8 @@ def _make_partial_tail_scheduler(
     if recurrent_only:
         kv_cache_config.kv_cache_groups = kv_cache_config.kv_cache_groups[1:]
     spec = MockOffloadingSpec(build_offloading_config(vllm_config, kv_cache_config))
+    if manager is not None:
+        spec.manager = manager
     return OffloadingConnectorScheduler(spec, vllm_config, kv_cache_config)
 
 
@@ -252,15 +295,16 @@ def _make_partial_tail_request(
     scheduler: OffloadingConnectorScheduler,
     kv_transfer_params: dict[str, Any] | None = None,
     request_id: str = "req",
+    num_tokens: int = 30,
 ) -> MagicMock:
     request = MagicMock()
     request.request_id = request_id
     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.num_prompt_tokens = num_tokens
+    request.num_tokens = num_tokens
+    request.block_hashes = [BlockHash(f"h{i}".encode()) for i in range(num_tokens // 4)]
+    request.all_token_ids = list(range(num_tokens))
     request.lora_request = None
     request.skip_reading_external_kv = False
     request.is_finished.return_value = False
@@ -268,6 +312,21 @@ def _make_partial_tail_request(
     return request
 
 
+def _make_real_partial_tail_request(
+    scheduler: OffloadingConnectorScheduler, request_id: str
+) -> Request:
+    init_none_hash(sha256)
+    request = Request(
+        request_id=request_id,
+        prompt_token_ids=list(range(30)),
+        sampling_params=SamplingParams(max_tokens=1),
+        pooling_params=None,
+        block_hasher=get_request_block_hasher(4, sha256),
+    )
+    scheduler.on_new_request(request)
+    return request
+
+
 def test_create_req_context_preserves_kv_hints():
     kv_hints = KvHintsEnvelope(
         protocol_version="0.1",
@@ -504,6 +563,44 @@ def test_partial_lookup_returns_exact_boundary_and_group_load_keys():
     assert req_status.partial_tail_boundary is None
 
 
+def test_partial_tail_lookup_delays_duplicate_inflight_load():
+    """A second request must wait for an in-flight partial-tail load."""
+    manager = CPUOffloadingManager(num_chunks=8)
+    scheduler = _make_partial_tail_scheduler(manager=manager)
+
+    first = _make_real_partial_tail_request(scheduler, "first")
+    second = _make_real_partial_tail_request(scheduler, "second")
+
+    first_state = scheduler._req_status[first.request_id]
+    first_state.num_locally_computed_tokens = 16
+    first_state.update_offload_keys()
+    boundary_keys = [
+        scheduler._make_boundary_key(
+            first, group_config.group_idx, 28, first_state.req_context
+        )
+        for group_config in scheduler.config.kv_group_configs
+    ]
+    stored = manager.prepare_store(boundary_keys, first_state.req_context)
+    assert stored is not None
+    manager.complete_store(stored.keys_to_store, first_state.req_context)
+
+    assert scheduler.get_num_new_matched_tokens(first, 16) == (12, True)
+    scheduler.update_state_after_alloc(
+        first,
+        KVCacheBlocks(
+            (
+                [KVCacheBlock(11), KVCacheBlock(12)],
+                [KVCacheBlock(0, is_null=True), KVCacheBlock(21)],
+            )
+        ),
+        num_external_tokens=12,
+    )
+    assert len(scheduler._jobs) == 1
+
+    assert scheduler.get_num_new_matched_tokens(second, 16) == (None, False)
+    assert len(scheduler._jobs) == 1
+
+
 def test_sync_partial_tail_load_targets_hashed_blocks():
     """A sync load's blocks are hashed at allocation; the partial-tail load
     must still target them, matching the async case."""
@@ -593,6 +690,79 @@ def test_max_load_tokens_caps_tokens_beyond_gpu_prefix():
     assert scheduler._req_status["req"].partial_tail_boundary == 24
 
 
+def _hybrid_lookup(stored: set[OffloadKey]):
+    """Full attention (group 0) holds every chunk; Mamba (group 1) only
+    the boundary states in ``stored``."""
+
+    def lookup(key: OffloadKey, req_context: ReqContext) -> LookupResult:
+        if get_offload_group_idx(key) == 0:
+            return LookupResult.HIT
+        return LookupResult.HIT if key in stored else LookupResult.MISS
+
+    return lookup
+
+
+def _make_sparse_retention_scheduler() -> OffloadingConnectorScheduler:
+    # prefix_cache_retention_interval=0, the engine default: one Mamba state
+    # per replay boundary.
+    scheduler = _make_partial_tail_scheduler()
+    scheduler.config = scheduler.config._replace(retention_interval=0)
+    return scheduler
+
+
+def test_sparse_group_miss_registers_shared_prefix_junction():
+    """A revisit that diverges before the only stored Mamba state heals.
+
+    The first request stored one Mamba state at the end of its prompt.
+    The revisit shares the first chunk (16 tokens) but not that state,
+    so the whole lookup misses. The connector must pin a junction at 16
+    so the recompute retains and offloads the Mamba state there; the next
+    revisit then hits the shared prefix.
+    """
+    scheduler = _make_sparse_retention_scheduler()
+    request = _make_partial_tail_request(scheduler)
+    request.shared_prefix_boundary = 0
+    stored: set[OffloadKey] = set()
+    scheduler.manager.lookup.side_effect = _hybrid_lookup(stored)
+
+    assert scheduler.get_num_new_matched_tokens(request, 0) == (0, False)
+    assert request.shared_prefix_boundary == 16
+
+    # The recompute's Mamba state at the junction is handed off and stored.
+    scheduler.manager.prepare_store.side_effect = lambda keys, req_context: (
+        generate_store_output(keys)
+    )
+    jobs = scheduler._build_aligned_boundary_store_jobs({"req": [(1, 77, 16)]})
+    [job_id] = jobs
+    stored.update(scheduler._jobs[job_id].keys)
+
+    scheduler._req_status["req"].transfer_jobs.clear()
+    tokens, _ = scheduler.get_num_new_matched_tokens(request, 0)
+    assert tokens == 16
+
+
+def test_shared_prefix_junction_is_not_lowered():
+    scheduler = _make_sparse_retention_scheduler()
+    request = _make_partial_tail_request(scheduler)
+    # e.g. a longer junction from the local prefix cache
+    request.shared_prefix_boundary = 24
+    scheduler.manager.lookup.side_effect = _hybrid_lookup(set())
+
+    assert scheduler.get_num_new_matched_tokens(request, 0) == (0, False)
+    assert request.shared_prefix_boundary == 24
+
+
+def test_no_shared_prefix_junction_when_sparse_group_hits():
+    scheduler = _make_sparse_retention_scheduler()
+    request = _make_partial_tail_request(scheduler)
+    request.shared_prefix_boundary = 0
+    scheduler.manager.lookup.return_value = LookupResult.HIT
+
+    tokens, _ = scheduler.get_num_new_matched_tokens(request, 0)
+    assert tokens == 28
+    assert request.shared_prefix_boundary == 0
+
+
 def test_max_load_tokens_rounds_down_without_partial_tail():
     scheduler = _make_partial_tail_scheduler()
     scheduler.config = scheduler.config._replace(supports_partial_tail=False)
@@ -968,6 +1138,230 @@ def lookup(key, req_context):
     assert req_status.partial_tail_boundary is None
 
 
+def _resident_lookup(hits: dict[int, set[bytes]]):
+    """Lookup side effect: HIT iff the key's hash is resident for its group."""
+
+    def lookup(key, req_context):
+        resident = hits.get(get_offload_group_idx(key), set())
+        return (
+            LookupResult.HIT
+            if get_offload_block_hash(key) in resident
+            else LookupResult.MISS
+        )
+
+    return lookup
+
+
+def _make_anchored_partial_tail_request(scheduler, hits, num_tokens=30):
+    """16-token chunks, hash 4: FA chunk [0,16) is h3, [16,32) is h7; a
+    30-token prompt's tail at 28 is h6. `hits` holds what each group has
+    resident; a recurrent group typically has its state only at the producer's
+    tail, never at a chunk boundary.
+    """
+    request = _make_partial_tail_request(scheduler, num_tokens=num_tokens)
+    req_status = scheduler._req_status["req"]
+    req_status.update_offload_keys()
+    scheduler.manager.lookup.side_effect = _resident_lookup(hits)
+    return request, req_status
+
+
+@pytest.mark.cpu_test
+def test_partial_lookup_anchors_on_full_attention_prefix():
+    # No recurrent state at the chunk boundary 16, the common case: the search
+    # must still find the tail at 28 from the full-attention prefix alone.
+    scheduler = _make_partial_tail_scheduler()
+    _, req_status = _make_anchored_partial_tail_request(
+        scheduler, {0: {b"h3", b"h6"}, 1: {b"h6"}}
+    )
+    assert scheduler._lookup(req_status) == 28
+    assert req_status.partial_tail_boundary == 28
+
+
+@pytest.mark.cpu_test
+def test_partial_lookup_anchor_loads_only_verified_chunks():
+    scheduler = _make_partial_tail_scheduler()
+    request, req_status = _make_anchored_partial_tail_request(
+        scheduler, {0: {b"h3", b"h6"}, 1: {b"h6"}}
+    )
+    assert scheduler._lookup(req_status) == 28
+    scheduler.update_state_after_alloc(
+        request,
+        KVCacheBlocks(
+            (
+                [KVCacheBlock(31), KVCacheBlock(32)],
+                [KVCacheBlock(0, is_null=True), KVCacheBlock(41)],
+            )
+        ),
+        num_external_tokens=28,
+    )
+    [load_job] = scheduler._current_batch_load_jobs.values()
+    assert [
+        (get_offload_group_idx(key), get_offload_block_hash(key))
+        for key in load_job.src_spec.offload_keys
+    ] == [(0, b"h3"), (0, b"h6"), (1, b"h6")]
+    assert isinstance(load_job.dst_spec, GPULoadStoreSpec)
+    assert load_job.dst_spec.block_ids.tolist() == [31, 32, 41]
+    assert load_job.dst_spec.group_sizes == [2, 1]
+    assert load_job.dst_spec.block_indices == [0, 1]
+
+
+@pytest.mark.cpu_test
+def test_partial_lookup_returns_zero_without_stored_tail():
+    scheduler = _make_partial_tail_scheduler()
+    _, req_status = _make_anchored_partial_tail_request(
+        scheduler, {0: {b"h3"}, 1: set()}
+    )
+    assert scheduler._lookup(req_status) == 0
+    assert req_status.partial_tail_boundary is None
+
+
+@pytest.mark.cpu_test
+def test_partial_lookup_requires_resident_full_attention_prefix():
+    # Tail keys resident but the chunk [0,16) was evicted: reporting the tail
+    # would make the load request a missing chunk.
+    scheduler = _make_partial_tail_scheduler()
+    _, req_status = _make_anchored_partial_tail_request(
+        scheduler, {0: {b"h6"}, 1: {b"h6"}}
+    )
+    assert scheduler._lookup(req_status) == 0
+    assert req_status.partial_tail_boundary is None
+
+
+@pytest.mark.cpu_test
+def test_partial_lookup_prefers_tail_over_aligned_recurrent_state():
+    scheduler = _make_partial_tail_scheduler()
+    _, req_status = _make_anchored_partial_tail_request(
+        scheduler, {0: {b"h3", b"h6"}, 1: {b"h3", b"h6"}}
+    )
+    assert scheduler._lookup(req_status) == 28
+
+
+@pytest.mark.cpu_test
+def test_partial_lookup_defers_while_anchor_chunk_is_loading():
+    scheduler = _make_partial_tail_scheduler()
+    _, req_status = _make_anchored_partial_tail_request(
+        scheduler, {0: {b"h3", b"h6"}, 1: {b"h6"}}
+    )
+    scheduler._chunks_being_loaded.add(req_status.group_states[0].offload_keys[0])
+    assert scheduler._lookup(req_status) is None
+    assert req_status.partial_tail_boundary is None
+
+
+@pytest.mark.cpu_test
+def test_partial_lookup_does_not_defer_on_pending_anchor_chunk():
+    # Like the complete-chunk lookup, a full-attention chunk that is still
+    # being stored does not defer the request; the tail is not reported.
+    scheduler = _make_partial_tail_scheduler()
+    _, req_status = _make_anchored_partial_tail_request(
+        scheduler, {0: {b"h6"}, 1: {b"h6"}}
+    )
+    resident = scheduler.manager.lookup.side_effect
+
+    def lookup(key, req_context):
+        if get_offload_group_idx(key) == 0 and get_offload_block_hash(key) == b"h3":
+            return LookupResult.HIT_PENDING
+        return resident(key, req_context)
+
+    scheduler.manager.lookup.side_effect = lookup
+    assert scheduler._lookup(req_status) == 0
+    assert req_status.partial_tail_boundary is None
+
+
+@pytest.mark.cpu_test
+def test_partial_lookup_anchor_spans_multiple_complete_chunks():
+    # 46-token prompt: chunks [0,16) and [16,32) resident, tail at 44, no
+    # recurrent state at 16 or 32. The bug as seen in production: the tail
+    # lies beyond the first chunk and the complete-chunk anchor is 0.
+    scheduler = _make_partial_tail_scheduler()
+    _, req_status = _make_anchored_partial_tail_request(
+        scheduler, {0: {b"h3", b"h7", b"h10"}, 1: {b"h10"}}, num_tokens=46
+    )
+    assert scheduler._lookup(req_status) == 44
+    assert req_status.partial_tail_boundary == 44
+
+
+@pytest.mark.cpu_test
+def test_partial_lookup_anchor_honours_max_num_new_tokens():
+    scheduler = _make_partial_tail_scheduler()
+    request, _ = _make_anchored_partial_tail_request(
+        scheduler, {0: {b"h3", b"h4", b"h6"}, 1: {b"h4", b"h6"}}
+    )
+    # tails at 20 and 28 are both resident; the cap admits only 20
+    assert scheduler.get_num_new_matched_tokens(request, 0, max_num_new_tokens=20) == (
+        20,
+        True,
+    )
+    assert scheduler._req_status["req"].partial_tail_boundary == 20
+
+
+@pytest.mark.cpu_test
+def test_partial_lookup_anchor_after_gpu_prefix_hit():
+    # The first chunk is already on the GPU; the anchor starts there and the
+    # tail is found without any recurrent state at the chunk boundary.
+    scheduler = _make_partial_tail_scheduler()
+    request, _ = _make_anchored_partial_tail_request(
+        scheduler, {0: {b"h6"}, 1: {b"h6"}}
+    )
+    assert scheduler.get_num_new_matched_tokens(request, 16) == (12, True)
+    assert scheduler._req_status["req"].partial_tail_boundary == 28
+
+
+@pytest.mark.parametrize(
+    "num_tokens,max_new_tokens,pending_anchor,eagle,expected_hit",
+    [
+        (46, None, False, False, 28),
+        (32, None, False, False, 28),
+        (46, 28, False, False, 28),
+        (46, None, True, False, 28),
+        (62, None, True, True, 36),
+        (30, None, True, False, 12),
+        (30, None, True, True, 8),
+    ],
+)
+@pytest.mark.cpu_test
+def test_partial_lookup_keeps_tail_below_longer_attention_prefix(
+    num_tokens, max_new_tokens, pending_anchor, eagle, expected_hit
+):
+    """A sibling's longer attention prefix must not hide a reusable tail."""
+    config = create_vllm_config(
+        kv_connector="OffloadingConnector",
+        kv_role="kv_both",
+        kv_connector_extra_config={"cpu_bytes_to_use": 1 << 20},
+    )
+    config.cache_config.prefix_match_unit = 4
+    kv_config = _make_mamba_hybrid_kv_cache_config()
+    kv_config.cache_hit_alignment_tokens = 4
+    kv_config.kv_cache_groups[0].is_eagle_group = eagle
+    spec = CPUOffloadingSpec(build_offloading_config(config, kv_config))
+    scheduler = OffloadingConnectorScheduler(spec, config, kv_config)
+    request = create_request(num_tokens=num_tokens, block_size=4)
+    scheduler.on_new_request(request)
+    context = scheduler._req_status[request.request_id].req_context
+    positions = (
+        [(0, 16), (0, 32), (0, 40), (1, 16), (1, 36)]
+        if eagle
+        else [(0, 16), (0, 32), (0, 28), (1, 16), (1, 28)]
+    )
+    if expected_hit < 16:
+        positions = [(0, expected_hit + (4 if eagle else 0)), (1, expected_hit)]
+    keys = [
+        make_offload_key(request.block_hashes[boundary // 4 - 1], group)
+        for group, boundary in positions
+    ]
+    pending_boundary = 16 if expected_hit < 16 else 48 if eagle else 32
+    pending_key = make_offload_key(request.block_hashes[pending_boundary // 4 - 1], 0)
+    ready_keys = [key for key in keys if not pending_anchor or key != pending_key]
+    store = scheduler.manager.prepare_store(ready_keys, context)
+    assert store is not None
+    scheduler.manager.complete_store(store.keys_to_store, context)
+    if pending_anchor:
+        assert scheduler.manager.prepare_store([pending_key], context) is not None
+
+    assert scheduler.get_num_new_matched_tokens(
+        request, 0, max_num_new_tokens=max_new_tokens
+    ) == (expected_hit, True)
+
+
 def _key_tokens(scheduler: OffloadingConnectorScheduler, keys) -> dict[int, int]:
     req_context = scheduler._req_status["req"].req_context
     return {
@@ -1487,6 +1881,26 @@ def test_cache_reset_without_in_flight_jobs_holds_nothing():
         scheduler.shutdown()
 
 
+def test_cache_reset_twice_before_a_step_still_flushes():
+    """A second reset before the engine steps must not assert: RL loops call
+    reset_prefix_cache(reset_connector=True) every iteration, and EngineCore
+    drains every queued utility call before it steps again."""
+    scheduler, _ = _make_tiering_engine_scheduler([])
+    try:
+        flushed = _store_in_flight(scheduler)
+        assert scheduler.reset_prefix_cache(reset_connector=True)
+        assert scheduler.reset_prefix_cache(reset_connector=True)
+
+        output = scheduler.schedule()
+        meta = output.kv_connector_metadata
+        assert isinstance(meta, OffloadingConnectorMetadata)
+        assert meta.jobs_to_flush == flushed
+        _report_jobs(scheduler, output, flushed)
+        assert not scheduler.has_requests()
+    finally:
+        scheduler.shutdown()
+
+
 def test_cache_reset_fence_held_too_long_warns_on_idle_steps(monkeypatch):
     """The warning must fire on the empty steps an idle engine takes while the
     fence is held, since no lookup runs there."""
@@ -1600,6 +2014,57 @@ def test_abort_queued_request_does_not_build_store_job(
     assert queued_req_id not in runner.connector_scheduler._req_status
 
 
+@pytest.mark.parametrize(
+    "finish_status", [RequestStatus.FINISHED_ERROR, RequestStatus.FINISHED_ABORTED]
+)
+def test_finish_mid_prefill_stores_only_computed_blocks(
+    request_runner, finish_status: RequestStatus
+):
+    """A request that finishes mid-prefill must not store a block it never
+    computed under that block's full hash."""
+    block_size = 4
+    runner = request_runner(
+        block_size=block_size, num_gpu_blocks=16, async_scheduling=False
+    )
+    runner.manager.prepare_store.side_effect = lambda keys, req_context: (
+        generate_store_output(keys)
+    )
+    # Chunked prefill: 6 of the 16 prompt tokens per step.
+    runner.scheduler.max_num_scheduled_tokens = 6
+    runner.scheduler.scheduler_config.max_num_batched_tokens = 6
+    runner.new_request(token_ids=list(range(100, 116)))
+    req_id = str(runner.req_id)
+
+    scheduler_output = runner.scheduler.schedule()
+    assert scheduler_output.num_scheduled_tokens[req_id] == 6
+    runner._update_gpu_blocks()
+    req = runner.scheduler.requests[req_id]
+    model_runner_output = create_model_runner_output(
+        reqs=[req], token_id=0, use_eos=False
+    )
+    # A partial prefill samples no token.
+    model_runner_output.sampled_token_ids = [[]]
+    runner.scheduler.update_from_output(scheduler_output, model_runner_output)
+    assert req.num_computed_tokens == 6
+
+    # The first step stored the one full block. The second block has only
+    # tokens 4-5 computed when the request finishes.
+    runner.scheduler.finish_requests(req_id, finish_status)
+    metadata = runner.scheduler.schedule().kv_connector_metadata
+    assert isinstance(metadata, OffloadingConnectorMetadata)
+    stored_blocks = sorted(
+        int(block_id)
+        for job in metadata.store_jobs.values()
+        if job.req_id == req_id
+        for block_id in job.src_spec.block_ids
+    )
+    assert not stored_blocks, (
+        f"{finish_status.name}: stored {len(stored_blocks)} blocks {stored_blocks} "
+        f"but only {6 // block_size} full block(s) were computed "
+        f"(num_computed_tokens=6, prompt=16)"
+    )
+
+
 def test_store_jobs_wait_for_missing_offload_key(request_runner):
     """Allocated chunks must wait until their block hash becomes available."""
     block_size = 4
@@ -1998,7 +2463,6 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
         decoded_tokens=[0] * (tokens_per_chunk - 1),
         expected_stored=(3, 4, 5),
     )
-    runner.manager.touch.assert_not_called()
 
     # +1 token -> single block, fail prepare_store
     runner.manager.prepare_store.side_effect = lambda keys, req_context: None
@@ -2027,7 +2491,6 @@ 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_not_called()
 
     # terminate request
     runner.run(
@@ -2573,7 +3036,6 @@ def test_two_groups_full_and_sliding_window(request_runner, async_scheduling: bo
         generate_store_output(keys)
     )
     runner.run(decoded_tokens=[0])
-    runner.manager.touch.assert_not_called()
 
     # store 3 more block
     runner.manager.prepare_store.side_effect = lambda keys, req_context: (
@@ -2584,8 +3046,6 @@ def test_two_groups_full_and_sliding_window(request_runner, async_scheduling: bo
         expected_stored=(0, 1, 2, 3, 4, 5),
     )
 
-    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.
     runner.run(decoded_tokens=[EOS_TOKEN_ID])
@@ -2602,7 +3062,6 @@ def test_two_groups_full_and_sliding_window(request_runner, async_scheduling: bo
         #   are within the window → loads blocks 1,2
         expected_loaded=((0, 0), (0, 1), (0, 2), (1, 1), (1, 2)),
     )
-    runner.manager.touch.assert_not_called()
 
     # 3 blocks are hit on GPU [0, 1, 2]
     # 1 block loaded [3,]
@@ -2720,7 +3179,6 @@ def test_two_groups_different_block_sizes(request_runner, async_scheduling: bool
         generate_store_output(keys)
     )
     runner.run(decoded_tokens=[0])
-    runner.manager.touch.assert_not_called()
 
     # Get to 31 tokens
     # No further blocks offloaded
@@ -2730,7 +3188,6 @@ 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])
-    runner.manager.touch.assert_not_called()
 
     # Get to 35 tokens
     # No further blocks offloaded
@@ -2740,7 +3197,6 @@ 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])
-    runner.manager.touch.assert_not_called()
 
     # Get to 47 tokens
     # No further blocks offloaded
@@ -2750,7 +3206,6 @@ 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])
-    runner.manager.touch.assert_not_called()
 
     runner.run(decoded_tokens=[0], expected_stored=((0, 3), (1, 2)))
 
@@ -3448,7 +3903,7 @@ def test_complete_store_called_per_job(request_runner, async_scheduling: bool):
     # First store: fires when block 0 is fully populated.
     runner.run(decoded_tokens=[0, 0], expected_stored=(0, 1, 2))
     assert runner.manager.complete_store.call_count == 1
-    first_call_keys = set(runner.manager.complete_store.call_args.args[0])
+    first_call_keys = runner.manager.complete_store.call_args.args[0]
     assert len(first_call_keys) == 1
     runner.manager.complete_store.reset_mock()
 
@@ -3458,7 +3913,7 @@ def test_complete_store_called_per_job(request_runner, async_scheduling: bool):
         expected_stored=(3, 4, 5),
     )
     assert runner.manager.complete_store.call_count == 1
-    second_call_keys = set(runner.manager.complete_store.call_args.args[0])
+    second_call_keys = runner.manager.complete_store.call_args.args[0]
     assert first_call_keys != second_call_keys
     runner.manager.complete_store.reset_mock()
 
@@ -3467,6 +3922,47 @@ def test_complete_store_called_per_job(request_runner, async_scheduling: bool):
     assert runner.manager.complete_store.call_count == 0
 
 
+@pytest.mark.parametrize("async_scheduling", [True, False])
+def test_store_job_preserves_order_and_deduplicates_keys(
+    request_runner, async_scheduling: bool
+):
+    tokens_per_block = 4
+    blocks_per_chunk = 3
+    tokens_per_chunk = tokens_per_block * blocks_per_chunk
+    runner = request_runner(
+        blocks_per_chunk=blocks_per_chunk,
+        block_size=tokens_per_block,
+        num_gpu_blocks=100,
+        async_scheduling=async_scheduling,
+    )
+    runner.new_request(token_ids=[0] * (tokens_per_chunk * 2))
+    returned_keys = []
+
+    def prepare_store(keys, req_context):
+        keys = list(keys)
+        returned_keys.extend(keys)
+        output = generate_store_output(keys)
+        output.keys_to_store.append(keys[0])
+        return output
+
+    runner.manager.prepare_store.side_effect = prepare_store
+    job_keys = {}
+
+    def capture_job_keys():
+        job_keys.update(runner.connector_scheduler._jobs.items())
+
+    runner.run(
+        decoded_tokens=[EOS_TOKEN_ID],
+        expected_stored=tuple(range(6)),
+        post_step_fn=capture_job_keys,
+    )
+
+    assert len(job_keys) == 1
+    keys = next(iter(job_keys.values())).keys
+    assert keys == tuple(returned_keys)
+    assert runner.manager.complete_store.call_args.args[0] == keys
+
+
 @pytest.mark.parametrize("async_scheduling", [True, False])
 def test_complete_store_waits_for_all_worker_acks(
     request_runner, async_scheduling: bool
@@ -3512,6 +4008,167 @@ def test_complete_store_waits_for_all_worker_acks(
     assert job_id not in runner.connector_scheduler._jobs
 
 
+@pytest.mark.parametrize("async_scheduling", [True, False])
+def test_failed_store_is_not_published_as_cache(request_runner, async_scheduling: bool):
+    """A store that failed on any worker must not become loadable.
+
+    Its destination blocks were never written, so publishing them would serve
+    whatever the recycled blocks last held -- potentially another request's KV.
+    The failure is reported in an earlier step than the one that completes the
+    job, which is the case a per-step failure lookup would miss.
+    """
+    tokens_per_block = 4
+    blocks_per_chunk = 3
+    runner = request_runner(
+        blocks_per_chunk=blocks_per_chunk,
+        block_size=tokens_per_block,
+        num_gpu_blocks=100,
+        async_scheduling=async_scheduling,
+        worker_count=3,
+    )
+    runner.new_request(token_ids=[0] * (tokens_per_block * blocks_per_chunk))
+    runner.manager.prepare_store.side_effect = lambda keys, req_context: (
+        generate_store_output(keys)
+    )
+    runner.run(decoded_tokens=[0, 0], complete_transfers=False)
+    job_id = next(iter(runner.connector_scheduler._jobs))
+    runner.manager.complete_store.reset_mock()
+
+    runner.connector_scheduler.update_connector_output(
+        KVConnectorOutput(
+            kv_connector_worker_meta=OffloadingWorkerMetadata(
+                completed_jobs={job_id: 1}, failed_jobs={job_id}
+            )
+        )
+    )
+    assert runner.manager.complete_store.call_count == 0
+
+    runner.connector_scheduler.update_connector_output(
+        KVConnectorOutput(
+            kv_connector_worker_meta=OffloadingWorkerMetadata(
+                completed_jobs={job_id: 2}
+            )
+        )
+    )
+    assert runner.manager.complete_store.call_count == 1
+    assert runner.manager.complete_store.call_args.kwargs["success"] is False
+
+
+@pytest.mark.parametrize("async_scheduling", [True, False])
+def test_successful_store_is_published(request_runner, async_scheduling: bool):
+    """The success path still publishes, so the guard cannot pass vacuously."""
+    tokens_per_block = 4
+    blocks_per_chunk = 3
+    runner = request_runner(
+        blocks_per_chunk=blocks_per_chunk,
+        block_size=tokens_per_block,
+        num_gpu_blocks=100,
+        async_scheduling=async_scheduling,
+        worker_count=3,
+    )
+    runner.new_request(token_ids=[0] * (tokens_per_block * blocks_per_chunk))
+    runner.manager.prepare_store.side_effect = lambda keys, req_context: (
+        generate_store_output(keys)
+    )
+    runner.run(decoded_tokens=[0, 0], complete_transfers=False)
+    job_id = next(iter(runner.connector_scheduler._jobs))
+    runner.manager.complete_store.reset_mock()
+
+    runner.connector_scheduler.update_connector_output(
+        KVConnectorOutput(
+            kv_connector_worker_meta=OffloadingWorkerMetadata(
+                completed_jobs={job_id: 3}
+            )
+        )
+    )
+    assert runner.manager.complete_store.call_count == 1
+    assert runner.manager.complete_store.call_args.kwargs["success"] is True
+
+
+@pytest.mark.parametrize("async_scheduling", [True, False])
+@pytest.mark.parametrize("recompute", [False, True], ids=["fail", "recompute"])
+@pytest.mark.parametrize("hybrid", [False, True], ids=["one_group", "full_and_swa"])
+def test_failed_load_ends_or_reruns_its_request(
+    request_runner, recompute: bool, hybrid: bool, async_scheduling: bool
+):
+    """A failed load ends or reruns its request, never the engine.
+
+    Its destination blocks hold undefined KV, so the request is reported through
+    failed_recving, which the scheduler handles for any number of KV cache
+    groups. A rerun must not read the offloaded cache again, or a persistent
+    fault would reload the same chunks forever.
+    """
+    block_size = 4
+    kv_cache_groups = None
+    expected_loaded: tuple[int | tuple[int, int], ...] = (0, 1, 2)
+    if hybrid:
+        # The sliding-window group loads only the blocks inside its window.
+        expected_loaded = ((0, 0), (0, 1), (0, 2), (1, 1), (1, 2))
+        kv_cache_groups = [
+            KVCacheGroupSpec(
+                ["layer0"],
+                FullAttentionSpec(
+                    block_size=block_size,
+                    num_kv_heads=1,
+                    head_size=1,
+                    dtype=torch.float32,
+                ),
+            ),
+            KVCacheGroupSpec(
+                ["layer1"],
+                SlidingWindowSpec(
+                    block_size=block_size,
+                    num_kv_heads=1,
+                    head_size=1,
+                    dtype=torch.float32,
+                    sliding_window=2 * block_size,
+                ),
+            ),
+        ]
+    runner = request_runner(
+        block_size=block_size,
+        num_gpu_blocks=100,
+        async_scheduling=async_scheduling,
+        kv_cache_groups=kv_cache_groups,
+    )
+    runner.scheduler.recompute_kv_load_failures = recompute
+    runner.new_request(token_ids=[0] * block_size * 3)
+    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))
+
+    handler = runner.offloading_spec.handler
+
+    def num_loads() -> int:
+        return sum(
+            isinstance(dst_spec, GPULoadStoreSpec)
+            for _, dst_spec in handler.transfer_specs.values()
+        )
+
+    def at_most_one_load() -> None:
+        assert num_loads() <= 1, "the failed load was retried"
+
+    runner.scheduler.reset_prefix_cache()
+    runner.new_request(token_ids=[0] * (block_size * 3 + 1))
+    request = runner.scheduler.requests[str(runner.req_id)]
+    runner.manager.lookup.return_value = LookupResult.HIT
+    runner.manager.prepare_store.side_effect = lambda keys, req_context: (
+        generate_store_output([])
+    )
+    handler.fail_loads = True
+    runner.run(decoded_tokens=[], expected_loaded=expected_loaded)
+
+    if not recompute:
+        assert request.status == RequestStatus.FINISHED_ERROR
+        return
+    assert request.num_computed_tokens == 0
+    runner.run(decoded_tokens=[EOS_TOKEN_ID], post_step_fn=at_most_one_load)
+    assert num_loads() == 1
+    assert request.status == RequestStatus.FINISHED_STOPPED
+    assert request.num_output_tokens == 1
+
+
 @pytest.mark.parametrize("async_scheduling", [True, False])
 def test_max_offload_tokens_validation(request_runner, async_scheduling: bool):
     """Validates max_offload_tokens: type coercion, boundary values, and capping.
@@ -5153,7 +5810,7 @@ def test_dcp_abort_reconsiders_final_swa_tail(self):
         spec = MockOffloadingSpec(build_offloading_config(vllm_config, kv_cache_config))
         scheduler = OffloadingConnectorScheduler(spec, vllm_config, kv_cache_config)
         _as_mock(scheduler.manager.prepare_store).side_effect = (
-            lambda keys, req_context: (generate_store_output(keys))
+            lambda keys, req_context: generate_store_output(keys)
         )
 
         request = MagicMock()
@@ -5777,7 +6434,7 @@ def _shared_kv_mtp_config():
     spec.num_speculative_tokens_per_batch_size = None
     spec.max_num_new_slots_for_drafting = 0
     spec.num_speculative_tokens = 1
-    return spec
+    return attach_speculator_config(spec, MTPSpeculatorConfig)
 
 
 class TestSharedGroupMTPOffload:
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 7e050e01f737..1e800abfd735 100644
--- a/tests/v1/kv_connector/unit/offloading_connector/test_worker.py
+++ b/tests/v1/kv_connector/unit/offloading_connector/test_worker.py
@@ -309,6 +309,71 @@ def test_load_completion(sync: bool):
     assert meta.transfer_stats.store.is_empty()
 
 
+def _kv_cache_config(num_groups: int) -> KVCacheConfig:
+    return KVCacheConfig(
+        num_blocks=0,
+        kv_cache_tensors=[],
+        kv_cache_groups=[MagicMock() for _ in range(num_groups)],
+    )
+
+
+def _failed(job_id: int) -> list[TransferResult]:
+    return [TransferResult(job_id=job_id, success=False)]
+
+
+def test_failed_store_is_dropped_and_still_completes():
+    """A failed store is a future cache miss, not an error.
+
+    It must still be marked completed, or the scheduler's per-job pending
+    count never reaches zero, and marked failed so it is not published.
+    """
+    worker, _ = _make_worker(_kv_cache_config(1))
+    worker.prepare_store_kv(_store_metadata(11))
+    worker.start_kv_transfers(_empty_metadata())
+    worker.worker.get_finished.return_value = _failed(11)
+
+    finished_sending, finished_recving = worker.get_finished(set())
+
+    assert finished_sending == set()
+    assert finished_recving == set()
+    assert worker.take_failed_recving() == set()
+    meta = worker.build_connector_worker_meta()
+    assert meta is not None
+    assert meta.completed_jobs == {11: 1}
+    assert meta.failed_jobs == {11}
+
+
+def test_failed_load_fails_its_request():
+    """A failed load leaves undefined KV in its destination blocks.
+
+    The request is reported as finished, per the base contract, and as failed,
+    which the scheduler handles for any number of KV cache groups.
+    """
+    worker, _ = _make_worker(_kv_cache_config(1))
+    worker.start_kv_transfers(_load_metadata(12))
+    worker.worker.get_finished.return_value = _failed(12)
+
+    _, finished_recving = worker.get_finished(set())
+
+    assert finished_recving == {"req"}
+    assert worker.take_failed_recving() == {"req"}
+    assert worker.take_failed_recving() == set()
+    meta = worker.build_connector_worker_meta()
+    assert meta is not None
+    assert meta.completed_jobs == {12: 1}
+    assert meta.failed_jobs == {12}
+
+
+def test_failed_sync_load_raises():
+    """A sync load's request already ran on the KV it was meant to load."""
+    worker, _ = _make_worker(_kv_cache_config(1))
+    worker.start_kv_transfers(_load_metadata(13, sync=True))
+    worker.worker.get_finished.return_value = _failed(13)
+
+    with pytest.raises(RuntimeError, match="sync load job 13 failed"):
+        worker.get_finished(set())
+
+
 def test_offloading_connector_worker_accepts_plugin_spec_default_layout():
     from vllm.distributed.kv_transfer.kv_connector.v1.offloading.worker import (
         OffloadingConnectorWorker,
diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_worker_metadata.py b/tests/v1/kv_connector/unit/offloading_connector/test_worker_metadata.py
index ad2505abcefc..7afa93313613 100644
--- a/tests/v1/kv_connector/unit/offloading_connector/test_worker_metadata.py
+++ b/tests/v1/kv_connector/unit/offloading_connector/test_worker_metadata.py
@@ -37,6 +37,13 @@ def test_aggregate_multiple_workers():
     assert result.completed_jobs == {42: 3, 43: 2, 7: 2, 8: 2}
 
 
+def test_aggregate_keeps_a_failure_reported_by_any_worker():
+    ok = OffloadingWorkerMetadata(completed_jobs={1: 1})
+    failed = OffloadingWorkerMetadata(completed_jobs={1: 1}, failed_jobs={1})
+    assert ok.aggregate(failed).failed_jobs == {1}
+    assert failed.aggregate(ok).failed_jobs == {1}
+
+
 def test_aggregate_transfer_stats():
     meta1 = OffloadingWorkerMetadata(
         transfer_stats=TransferStats(
diff --git a/tests/v1/kv_connector/unit/offloading_connector/utils.py b/tests/v1/kv_connector/unit/offloading_connector/utils.py
index f8740e0bab87..30437ad3a111 100644
--- a/tests/v1/kv_connector/unit/offloading_connector/utils.py
+++ b/tests/v1/kv_connector/unit/offloading_connector/utils.py
@@ -84,6 +84,9 @@ def __init__(self):
         self.transfer_specs: dict[int, tuple[LoadStoreSpec, LoadStoreSpec]] = {}
         self.completed_transfers: list[TransferResult] = []
         self.waiting_jobs: set[int] = set()
+        # When set, loads submitted from now on complete as failed.
+        self.fail_loads = False
+        self._failing_jobs: set[int] = set()
         self.completed_jobs: list[int] = []
         self.flushed_jobs: set[int] = set()
 
@@ -108,6 +111,8 @@ def submit_load(
     ) -> bool:
         self.transfer_specs[job_id] = (src_spec, dst_spec)
         self.waiting_jobs.add(job_id)
+        if self.fail_loads:
+            self._failing_jobs.add(job_id)
         return True
 
     def complete_jobs(self, job_ids: set[int]) -> None:
@@ -117,7 +122,7 @@ def complete_jobs(self, job_ids: set[int]) -> None:
                 self.completed_jobs.append(job_id)
                 result = TransferResult(
                     job_id=job_id,
-                    success=True,
+                    success=job_id not in self._failing_jobs,
                     transfer_size=None,
                     transfer_time=None,
                 )
@@ -535,7 +540,7 @@ def _run(
                 self.offloading_spec.complete_transfers()
 
             self.worker_connector.wait_for_save()
-            finished_sending, finished_recving = self.worker_connector.get_finished(
+            transfer_results = self.worker_connector.get_transfer_results(
                 scheduler_output.finished_req_ids
             )
             worker_meta = (
@@ -547,10 +552,11 @@ def _run(
 
             model_runner_output = create_model_runner_output(
                 reqs=self.scheduler.running,
-                finished_sending=finished_sending,
-                finished_recving=finished_recving,
+                finished_sending=transfer_results.finished_sending,
+                finished_recving=transfer_results.finished_recving,
                 token_id=token_id or 0,
                 kv_connector_worker_meta=worker_meta,
+                failed_recving=transfer_results.failed_recving,
             )
 
             prev_token_id = token_id
diff --git a/tests/v1/kv_connector/unit/test_config.py b/tests/v1/kv_connector/unit/test_config.py
index b2853195e3a3..f9164dfa9e62 100644
--- a/tests/v1/kv_connector/unit/test_config.py
+++ b/tests/v1/kv_connector/unit/test_config.py
@@ -3,6 +3,8 @@
 
 """Tests for KV cache offloading configuration."""
 
+import sys
+
 import pytest
 
 from vllm.config import (
@@ -185,6 +187,22 @@ def test_sleep_mode_refuses_unsupported_kv_connector(
         )
 
 
+def test_mooncake_store_supports_sleep_with_explicit_rdma(monkeypatch, tmp_path):
+    """Disabling auto-discovery keeps explicit RDMA despite NVLink overrides."""
+    monkeypatch.setitem(sys.modules, "mooncake", None)
+    config_path = tmp_path / "mooncake.json"
+    config_path.write_text('{"protocol": "rdma", "device_name": "mlx5_0"}')
+    monkeypatch.setenv("MOONCAKE_CONFIG_PATH", str(config_path))
+    monkeypatch.setenv("MC_MS_AUTO_DISC", "0")
+    monkeypatch.delenv("MC_USE_TENT", raising=False)
+    monkeypatch.delenv("MC_USE_TEV1", raising=False)
+    monkeypatch.setenv("MC_FORCE_MNNVL", "1")
+    monkeypatch.setenv("MC_INTRANODE_NVLINK", "1")
+    config = KVTransferConfig(kv_connector="MooncakeStoreConnector", kv_role="kv_both")
+    connector_cls = KVConnectorFactory.get_connector_class(config)
+    assert connector_cls.supports_sleep_mode(config)
+
+
 def test_kv_connector_allows_other_alloc_conf(monkeypatch):
     """Other PYTORCH_CUDA_ALLOC_CONF values must not be rejected."""
     monkeypatch.setenv(
@@ -215,3 +233,21 @@ def test_kv_offloading_size_only_uses_native_default():
     assert kv_transfer_config.kv_connector == "OffloadingConnector"
     assert kv_transfer_config.kv_role == "kv_both"
     assert kv_connector_extra_config["cpu_bytes_to_use"] == 4.0 * (1 << 30)
+
+
+@pytest.mark.parametrize(
+    "mode,expected",
+    [(None, "NixlPullConnector"), ("push", "NixlPushConnector"), ("Push", None)],
+)
+def test_nixl_connector_transfer_mode(mode, expected):
+    """Frameworks that only accept "NixlConnector" select push via extra config."""
+    ktc = KVTransferConfig(
+        kv_connector="NixlConnector",
+        kv_role="kv_both",
+        kv_connector_extra_config={} if mode is None else {"transfer_mode": mode},
+    )
+    if expected is None:
+        with pytest.raises(ValueError, match="transfer_mode"):
+            KVConnectorFactory.get_connector_class(ktc)
+    else:
+        assert KVConnectorFactory.get_connector_class(ktc).__name__ == expected
diff --git a/tests/v1/kv_connector/unit/test_hisparse_connector.py b/tests/v1/kv_connector/unit/test_hisparse_connector.py
index bdfbee28d4c0..5a6615d4cfd0 100644
--- a/tests/v1/kv_connector/unit/test_hisparse_connector.py
+++ b/tests/v1/kv_connector/unit/test_hisparse_connector.py
@@ -1,6 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import contextlib
 import ctypes
 from types import SimpleNamespace
 from unittest.mock import MagicMock, patch
@@ -167,7 +166,6 @@ def test_full_graph_step_prepares_host_mirror_outside_model():
         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])
@@ -180,6 +178,7 @@ def test_full_graph_step_prepares_host_mirror_outside_model():
     worker._draft_layers = ()
     worker.is_host_writer = True
     worker._enqueue_row_dma = MagicMock()
+    worker._invalidate_written_slots = MagicMock()
     worker.start_step = MagicMock(
         side_effect=lambda *_args, **_kwargs: worker._clear_forward_mirror_state()
     )
@@ -208,7 +207,8 @@ def test_full_graph_step_prepares_host_mirror_outside_model():
     worker._enqueue_host_mirror()
 
     worker._enqueue_row_dma.assert_called_once_with((0,), ready_event=None)
-    runtime.invalidate_written_slots.assert_called_once()
+    worker._invalidate_written_slots.assert_called_once()
+    assert worker._invalidate_written_slots.call_args.args[0] == [runtime]
 
 
 class _FakeEvent:
@@ -247,11 +247,8 @@ def swap_blocks_batch(src, dst, sizes):
 
     monkeypatch.setattr(torch, "Event", _FakeEvent)
     monkeypatch.setattr(torch, "Stream", lambda *args, **kwargs: MagicMock())
-    monkeypatch.setattr(torch.cuda, "Stream", lambda *args, **kwargs: MagicMock())
-    monkeypatch.setattr(torch.cuda, "stream", lambda _: contextlib.nullcontext())
     compute_stream = MagicMock()
-    monkeypatch.setattr(worker_module, "current_stream", lambda: compute_stream)
-    monkeypatch.setattr(worker_module.ops, "swap_blocks_batch", swap_blocks_batch)
+    monkeypatch.setattr(torch.accelerator, "current_stream", lambda: compute_stream)
     monkeypatch.setattr(
         runtime_module,
         "get_forward_context",
@@ -269,7 +266,6 @@ def swap_blocks_batch(src, dst, sizes):
             host_cache=torch.zeros(4 * block_size, width),
             request_state_indices=request_state_indices,
             begin_forward=MagicMock(),
-            invalidate_written_slots=MagicMock(),
         )
         handle = HiSparseCacheHandle(runtime)
         handle.view = SimpleNamespace(
@@ -290,6 +286,7 @@ def swap_blocks_batch(src, dst, sizes):
         ),
         MagicMock(),
     )
+    monkeypatch.setattr(worker.backend, "swap_blocks_batch", swap_blocks_batch)
     worker.initialize(
         handles,
         ["target", "draft"],
@@ -299,6 +296,7 @@ def swap_blocks_batch(src, dst, sizes):
         device=torch.device("cpu"),
         pinned_host_pools=[],
     )
+    worker._invalidate_written_slots = MagicMock()
     connector = object.__new__(HiSparseConnector)
     connector.connector_worker = worker
     connector._get_connector_metadata = MagicMock(
@@ -412,7 +410,7 @@ def test_zero_token_step_holds_tp_ranks_after_host_write_wait(
     compute_stream.wait_event.side_effect = lambda event: calls.append(("wait", event))
     tp_group = MagicMock()
     tp_group.barrier.side_effect = lambda: calls.append(("barrier", None))
-    monkeypatch.setattr(worker_module, "current_stream", lambda: compute_stream)
+    monkeypatch.setattr(torch.accelerator, "current_stream", lambda: compute_stream)
     monkeypatch.setattr(worker_module, "get_tp_group", lambda: tp_group)
 
     worker = object.__new__(HiSparseConnectorWorker)
@@ -513,7 +511,7 @@ def test_residency_updates_persist_by_state_row():
     An update for a lost page and an appended page must leave the other pages
     and the other request's row intact across steps that reorder the batch.
     """
-    table = torch.zeros((4, 2, 4), dtype=torch.int32)
+    table = torch.zeros((2, 4, 4), dtype=torch.int32)
     update_hisparse_residency(
         table,
         {
@@ -530,5 +528,5 @@ def test_residency_updates_persist_by_state_row():
         torch.tensor([0, 2], dtype=torch.int32),
     )
 
-    assert table[[0, 2], 0].tolist() == [[4, 5, 0, 0], [1, 0, 3, 6]]
-    assert table[[0, 2], 1].tolist() == [[14, 15, 0, 0], [11, 0, 13, 16]]
+    assert table[0, [0, 2]].tolist() == [[4, 5, 0, 0], [1, 0, 3, 6]]
+    assert table[1, [0, 2]].tolist() == [[14, 15, 0, 0], [11, 0, 13, 16]]
diff --git a/tests/v1/kv_connector/unit/test_mooncake_connector.py b/tests/v1/kv_connector/unit/test_mooncake_connector.py
index 4b835895f845..43b648617b00 100644
--- a/tests/v1/kv_connector/unit/test_mooncake_connector.py
+++ b/tests/v1/kv_connector/unit/test_mooncake_connector.py
@@ -27,9 +27,11 @@
 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector import (
     _SHARED_REGION_GROUP_ID,
     KVConnectorRole,
+    LeaseHeartbeat,
     MooncakeConnector,
     MooncakeConnectorMetadata,
     MooncakeConnectorWorker,
+    MooncakeLeaseHeartbeat,
     MooncakeXferMetadata,
     MooncakeXferResponse,
     MooncakeXferResponseStatus,
@@ -62,6 +64,7 @@
 from vllm.v1.request import RequestStatus
 
 from .utils import (
+    create_model_runner_output,
     create_request,
     create_scheduler,
     create_vllm_config,
@@ -180,12 +183,26 @@ def _make_test_kv_cache_config() -> KVCacheConfig:
     )
 
 
-def _make_worker_vllm_config(kv_role: str):
+def _make_worker_vllm_config(
+    kv_role: str, kv_connector_extra_config: dict | None = None
+):
     return create_vllm_config(
         kv_connector="MooncakeConnector",
         kv_role=kv_role,
         block_size=maybe_update_block_size(16),
+        kv_connector_extra_config=kv_connector_extra_config,
+    )
+
+
+def _make_compressed_kv_cache_config() -> KVCacheConfig:
+    spec = MLAAttentionSpec(
+        block_size=384,
+        num_kv_heads=1,
+        head_size=128,
+        dtype=torch.uint8,
+        tokens_per_state=8,
     )
+    return KVCacheConfig(2, [], [KVCacheGroupSpec(["summary"], spec)])
 
 
 def _make_packed_mla_kv_cache_config(
@@ -1328,6 +1345,29 @@ def test_resolve_need_send_accounts_for_remote_tp_fanout():
     assert send_meta.need_send == 2
 
 
+@pytest.mark.parametrize("kernel_block_size,raises", [(384, False), (64, True)])
+def test_mooncake_compressed_kernel_split(kernel_block_size, raises):
+    """Mooncake keeps compressed pages whole or rejects a layer-outer split."""
+    config = _make_worker_vllm_config("kv_both")
+    config.cache_config.block_size = 384
+    config.cache_config.kv_cache_layout = "LBHNC"
+    ctx = pytest.raises(NotImplementedError) if raises else contextlib.nullcontext()
+    with (
+        set_current_vllm_config(config),
+        patch_worker_dependencies(),
+        patch.object(
+            mooncake_connector,
+            "select_common_block_size",
+            return_value=kernel_block_size,
+        ),
+        ctx,
+    ):
+        worker = MooncakeConnector(
+            config, KVConnectorRole.WORKER, _make_compressed_kv_cache_config()
+        ).connector_worker
+        worker.shutdown()
+
+
 @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."""
@@ -1855,6 +1895,320 @@ async def test_worker_get_finished_timeout(monkeypatch):
         assert "tx-active" in prefill_worker.reqs_need_send
 
 
+class _BootstrapClientStub:
+    """httpx.AsyncClient stand-in for the prefill bootstrap server."""
+
+    def __init__(self, workers: dict | None = None, **kwargs):
+        self.workers = workers or {}
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, *exc_info):
+        return False
+
+    async def post(self, url, json=None):
+        return MagicMock()
+
+    async def get(self, url):
+        response = MagicMock()
+        response.json.return_value = self.workers
+        return response
+
+
+def _lease_decode_scheduler(kv_lease_duration: float | None):
+    extra_config = (
+        {} if kv_lease_duration is None else {"kv_lease_duration": kv_lease_duration}
+    )
+    vllm_config = create_vllm_config(
+        kv_connector="MooncakeConnector",
+        kv_role="kv_consumer",
+        kv_connector_extra_config=extra_config,
+    )
+    return create_scheduler(vllm_config), vllm_config.cache_config.block_size
+
+
+def _remote_prefill_request(request_id: int, block_size: int):
+    request = create_request(
+        request_id=request_id,
+        block_size=block_size,
+        num_tokens=block_size * 2,
+        do_remote_prefill=True,
+    )
+    request.kv_transfer_params.update(
+        {
+            "transfer_id": f"xfer-{request_id}",
+            "remote_bootstrap_addr": "http://prefill:8998",
+        }
+    )
+    return request
+
+
+@pytest.mark.parametrize("kv_lease_duration", [None, 30])
+def test_decode_lease_renewal_is_opt_in_and_throttled(kv_lease_duration):
+    scheduler, block_size = _lease_decode_scheduler(kv_lease_duration)
+    scheduler.add_request(_remote_prefill_request(1, block_size))
+
+    output = scheduler.schedule()
+    if kv_lease_duration is None:
+        assert not output.kv_connector_metadata.lease_heartbeats
+    else:
+        assert output.kv_connector_metadata.lease_heartbeats == {
+            "my-engine-id": LeaseHeartbeat(
+                remote_bootstrap_addr="http://prefill:8998", transfer_ids=["xfer-1"]
+            )
+        }
+    scheduler.update_from_output(output, create_model_runner_output(reqs=[]))
+    # At most one renewal per kv_lease_duration / 6 seconds.
+    assert not scheduler.schedule().kv_connector_metadata.lease_heartbeats
+
+
+def _next_lease_heartbeats(scheduler, scheduler_output, timeout: float = 5.0):
+    connector = scheduler.get_kv_connector()
+    deadline = time.monotonic() + timeout
+    while time.monotonic() < deadline:
+        meta = connector.build_connector_meta(scheduler_output)
+        if meta.lease_heartbeats:
+            return meta.lease_heartbeats
+        time.sleep(0.005)
+    return {}
+
+
+@pytest.mark.parametrize("settled_by", ["finished_recving", "abort"])
+def test_decode_renews_lease_until_pull_settles(settled_by):
+    # Renew every 10 ms so the test can wait for the next renewal.
+    scheduler, block_size = _lease_decode_scheduler(0.06)
+    settled = _remote_prefill_request(1, block_size)
+    kept = _remote_prefill_request(2, block_size)
+    scheduler.add_request(settled)
+    scheduler.add_request(kept)
+
+    output = scheduler.schedule()
+    meta = output.kv_connector_metadata
+    assert len(meta.reqs_to_recv["my-engine-id"]) == 2
+    assert sorted(meta.lease_heartbeats["my-engine-id"].transfer_ids) == [
+        "xfer-1",
+        "xfer-2",
+    ]
+    if settled_by == "finished_recving":
+        scheduler.update_from_output(
+            output,
+            create_model_runner_output(reqs=[], finished_recving={settled.request_id}),
+        )
+    else:
+        scheduler.update_from_output(output, create_model_runner_output(reqs=[]))
+        scheduler.finish_requests(settled.request_id, RequestStatus.FINISHED_ABORTED)
+
+    # The pull of `kept` has been issued but not settled, so it is still renewed.
+    assert _next_lease_heartbeats(scheduler, output) == {
+        "my-engine-id": LeaseHeartbeat(
+            remote_bootstrap_addr="http://prefill:8998", transfer_ids=["xfer-2"]
+        )
+    }
+
+
+@pytest.mark.asyncio
+async def test_decode_heartbeat_reaches_every_paired_prefill_worker():
+    """A renewal for a request still queued on D queries P's bootstrap first."""
+    vllm_config = _make_worker_vllm_config(
+        "kv_consumer", kv_connector_extra_config={"kv_lease_duration": 30}
+    )
+    # P TP=2 serves this D TP=1 worker from both of its ranks.
+    workers = {
+        "0": {
+            "engine_id": "p-engine",
+            "worker_addr": {
+                "0": {"0": "tcp://prefill-tp0:1234"},
+                "1": {"0": "tcp://prefill-tp1:1234"},
+            },
+        }
+    }
+    with (
+        set_current_vllm_config(vllm_config),
+        patch_worker_dependencies() as mocks,
+        patch.object(
+            mooncake_connector.httpx,
+            "AsyncClient",
+            lambda **kwargs: _BootstrapClientStub(workers),
+        ),
+    ):
+        worker = MooncakeConnector(
+            vllm_config, KVConnectorRole.WORKER, _make_test_kv_cache_config()
+        ).connector_worker
+        sock = mocks["mock_socket_object"]
+        sock.recv.return_value = msgspec.msgpack.encode(
+            MooncakeXferResponse(status=MooncakeXferResponseStatus.FINISH)
+        )
+        metadata = MooncakeConnectorMetadata()
+        metadata.lease_heartbeats = {
+            "p-engine": LeaseHeartbeat(
+                remote_bootstrap_addr="http://prefill:8998",
+                transfer_ids=["xfer-1", "xfer-2"],
+            )
+        }
+        worker.start_load_kv(metadata)
+
+        deadline = time.monotonic() + 10
+        while sock.send.await_count < 2 and time.monotonic() < deadline:
+            await asyncio.sleep(0.01)
+        addrs = sorted(c.args[1] for c in mocks["mock_make_zmq"].call_args_list)
+        assert addrs == ["tcp://prefill-tp0:1234", "tcp://prefill-tp1:1234"]
+        for call in sock.send.await_args_list:
+            heartbeat = msgspec.msgpack.decode(
+                call.args[0], type=MooncakeLeaseHeartbeat
+            )
+            assert heartbeat.lease_transfer_ids == ["xfer-1", "xfer-2"]
+        worker.shutdown()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("kv_lease_duration", [30, None])
+async def test_prefill_side_channel_renews_leases(kv_lease_duration):
+    extra_config = (
+        {} if kv_lease_duration is None else {"kv_lease_duration": kv_lease_duration}
+    )
+    vllm_config = _make_worker_vllm_config(
+        "kv_producer", kv_connector_extra_config=extra_config
+    )
+    with (
+        set_current_vllm_config(vllm_config),
+        patch_worker_dependencies(),
+        patch.object(mooncake_connector.httpx, "AsyncClient", _BootstrapClientStub),
+    ):
+        producer = MooncakeConnector(
+            vllm_config, KVConnectorRole.WORKER, _make_test_kv_cache_config()
+        ).connector_worker
+
+        # Allocation, then prefill end, as the scheduler reports them.
+        for block_ids in ([], [[1]]):
+            meta = MooncakeConnectorMetadata()
+            meta.reqs_to_send = {f"p-req-{i}": (f"xfer-{i}", block_ids) for i in (1, 2)}
+            before = time.perf_counter()
+            await producer.record_send_reqs(meta)
+        initial = kv_lease_duration or envs.VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT
+        for send_meta in producer.reqs_need_send.values():
+            assert before + initial <= send_meta.expire_time
+            assert send_meta.expire_time <= time.perf_counter() + initial
+            # Both deadlines pass.
+            send_meta.expire_time = before - 1
+
+        listener = asyncio.create_task(
+            producer._mooncake_sender_listener(threading.Event())
+        )
+        ctx = zmq.asyncio.Context()
+        dealer = ctx.socket(zmq.DEALER)
+        try:
+            while not producer.side_channel_port:
+                await asyncio.sleep(0.01)
+            dealer.connect(f"tcp://127.0.0.1:{producer.side_channel_port}")
+
+            async def request(message) -> MooncakeXferResponse:
+                await dealer.send(msgspec.msgpack.encode(message))
+                reply = await asyncio.wait_for(dealer.recv(), timeout=10)
+                return msgspec.msgpack.decode(reply, type=MooncakeXferResponse)
+
+            renewal = MooncakeLeaseHeartbeat(lease_transfer_ids=["xfer-1", "xfer-9"])
+            response = await request(renewal)
+            if kv_lease_duration is None:
+                assert response.status == MooncakeXferResponseStatus.ERROR
+                expired = {"p-req-1", "p-req-2"}
+            else:
+                assert response.status == MooncakeXferResponseStatus.FINISH
+                renewed = producer.reqs_need_send["xfer-1"].expire_time
+                assert renewed >= before + kv_lease_duration * 2 / 3
+                # A renewal never recreates a transfer P does not hold, and
+                # never shortens a deadline.
+                assert "xfer-9" not in producer.reqs_need_send
+                producer.reqs_need_send["xfer-1"].expire_time = renewed + 100
+                await request(renewal)
+                assert producer.reqs_need_send["xfer-1"].expire_time == renewed + 100
+                expired = {"p-req-2"}
+
+            # Pulls on the same channel still reach the sender workers.
+            pull = _xfer_meta(
+                [_region(0x2000)], {"d-req-1": ("xfer-1", [[1]])}, remote_tp_rank=1
+            )
+            response = await request(pull)
+            assert response.status == MooncakeXferResponseStatus.ERROR
+            assert "not paired" in response.err_msg
+
+            assert await producer.fetch_finished_sending_reqs() == expired
+        finally:
+            listener.cancel()
+            with contextlib.suppress(asyncio.CancelledError):
+                await listener
+            dealer.close(linger=0)
+            ctx.term()
+            producer.shutdown()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("with_live_req", [False, True])
+async def test_pull_after_expiry_fails_without_waiting(with_live_req):
+    """A pull for KV that P already freed must not wait for a prefill."""
+    vllm_config = _make_worker_vllm_config("kv_producer")
+    with set_current_vllm_config(vllm_config), patch_worker_dependencies():
+        producer = MooncakeConnector(
+            vllm_config, KVConnectorRole.WORKER, _make_test_kv_cache_config()
+        ).connector_worker
+        producer.kv_caches_base_addr = [0x1000]
+        producer.block_len_per_layer = [256]
+        producer.kv_block_len_per_layer = [256]
+        producer.registered_layer_names = ["model.layers.0.self_attn"]
+        producer.registered_layer_indices = [0]
+        sender_loop = producer.sender_loop
+        producer.sender_loop = asyncio.get_running_loop()
+
+        producer.reqs_need_send["xfer-expired"] = SendBlockMeta(
+            p_req_id="p-req-expired",
+            transfer_id="xfer-expired",
+            local_block_ids=[[1]],
+            ready=asyncio.Event(),
+            expire_time=time.perf_counter() - 1,
+        )
+        assert await producer.fetch_finished_sending_reqs() == {"p-req-expired"}
+
+        req_blocks = {"d-req-expired": ("xfer-expired", [[100]])}
+        if with_live_req:
+            live = SendBlockMeta(
+                p_req_id="p-req-live",
+                transfer_id="xfer-live",
+                local_block_ids=[[2]],
+                ready=asyncio.Event(),
+                expire_time=time.perf_counter() + 100,
+            )
+            live.ready.set()
+            producer.reqs_need_send["xfer-live"] = live
+            req_blocks["d-req-live"] = ("xfer-live", [[200]])
+        metadata = _xfer_meta(
+            [_region(0x2000, row_offset=-1, layer_name="model.layers.0.self_attn")],
+            req_blocks,
+        )
+        sock = AsyncMock(spec=zmq.asyncio.Socket, send_multipart=AsyncMock())
+        with patch.object(producer, "_send_blocks", return_value=0):
+            await producer.send_kv_to_decode(b"consumer", sock, metadata)
+
+        # No placeholder that would wait for a prefill that already ran.
+        assert "xfer-expired" not in producer.reqs_need_send
+        responses = [
+            producer._xfer_resp_decoder.decode(c.args[0][1])
+            for c in sock.send_multipart.call_args_list
+        ]
+        assert responses[0].err_reqs == ["d-req-expired"]
+        if with_live_req:
+            assert responses[0].status == MooncakeXferResponseStatus.CONTINUE
+            assert len(responses) == 2
+            assert responses[1].status == MooncakeXferResponseStatus.FINISH
+            assert responses[1].ok_reqs == ["d-req-live"]
+            assert producer.finished_sending_reqs == {"p-req-live"}
+        else:
+            assert responses[0].status == MooncakeXferResponseStatus.FINISH
+            assert len(responses) == 1
+
+        producer.sender_loop = sender_loop
+        producer.shutdown()
+
+
 @pytest.mark.parametrize(
     ("layout", "separate_kv_head_groups"),
     [
diff --git a/tests/v1/kv_connector/unit/test_mooncake_connector_hybrid_mamba.py b/tests/v1/kv_connector/unit/test_mooncake_connector_hybrid_mamba.py
index d4eaef2b3a00..974f1782bf67 100644
--- a/tests/v1/kv_connector/unit/test_mooncake_connector_hybrid_mamba.py
+++ b/tests/v1/kv_connector/unit/test_mooncake_connector_hybrid_mamba.py
@@ -38,23 +38,31 @@
 )
 
 from .test_mooncake_connector import patch_worker_dependencies
-from .utils import create_request, create_vllm_config, maybe_update_block_size
+from .utils import (
+    create_request,
+    create_scheduler,
+    create_vllm_config,
+    maybe_update_block_size,
+)
 
 
 def noop_shutdown():
     pass
 
 
-def make_hybrid_gdn_kv_cache_config(block_size: int) -> KVCacheConfig:
+def make_hybrid_gdn_kv_cache_config(
+    block_size: int, num_blocks: int = 16, mamba_cache_mode: str = "none"
+) -> KVCacheConfig:
     gdn_spec = MambaSpec(
         block_size=block_size,
         shapes=((6, 3), (1, 2, 2)),
         dtypes=(torch.float16, torch.float16),
         mamba_type=MambaAttentionBackendEnum.GDN_ATTN,
+        mamba_cache_mode=mamba_cache_mode,
     )
     assert gdn_spec.mamba_type == MambaAttentionBackendEnum.GDN_ATTN
     return KVCacheConfig(
-        num_blocks=16,
+        num_blocks=num_blocks,
         kv_cache_tensors=[],
         kv_cache_groups=[
             KVCacheGroupSpec(
@@ -102,6 +110,67 @@ def test_hybrid_gdn_remote_prefill_uses_mamba_n_minus_one():
     assert is_async is True
 
 
+@pytest.mark.cpu_test
+@pytest.mark.parametrize(
+    "kv_role,expected_attn_pull_blocks",
+    [
+        # The local attention prefix is kept; only the suffix and GDN state move.
+        ("kv_consumer", 2),
+        # Roles that prefill locally keep the reconciled lookup: no local hit.
+        ("kv_both", 6),
+    ],
+)
+def test_hybrid_gdn_remote_prefill_keeps_local_attention_prefix(
+    kv_role: str, expected_attn_pull_blocks: int
+):
+    vllm_config = create_vllm_config(kv_connector="MooncakeConnector", kv_role=kv_role)
+    vllm_config.scheduler_config.disable_hybrid_kv_cache_manager = False
+    vllm_config.cache_config.mamba_cache_mode = "align"
+    block_size = vllm_config.cache_config.block_size
+    kv_cache_config = make_hybrid_gdn_kv_cache_config(
+        block_size, num_blocks=64, mamba_cache_mode="align"
+    )
+    scheduler = create_scheduler(
+        vllm_config,
+        num_blocks=kv_cache_config.num_blocks,
+        kv_cache_config=kv_cache_config,
+    )
+    manager = scheduler.kv_cache_manager
+
+    # Cache a 4-block prefix, then evict its GDN states: the attention blocks
+    # stay hit while no block boundary has a local recurrent state.
+    prefix_len = 4 * block_size
+    fill = create_request(num_tokens=prefix_len, common_prefix_len=prefix_len)
+    computed_blocks, num_computed, _ = manager.get_computed_blocks(fill)
+    assert num_computed == 0
+    gdn_ids = []
+    for i in range(4):
+        manager.allocate_slots(
+            fill, block_size, new_computed_blocks=computed_blocks if i == 0 else None
+        )
+        fill.num_computed_tokens += block_size
+        gdn_ids.append(manager.get_block_ids(fill.request_id)[1][i])
+    manager.free(fill)
+    manager.block_pool.evict_blocks(set(gdn_ids))
+
+    replay = create_request(
+        num_tokens=6 * block_size, common_prefix_len=prefix_len, max_tokens=1
+    )
+    replay.kv_transfer_params = dict(
+        do_remote_prefill=True,
+        remote_engine_id="prefill-engine",
+        remote_bootstrap_addr="prefill-host:8998",
+        transfer_id="xfer-0",
+    )
+    scheduler.add_request(replay)
+    output = scheduler.schedule()
+
+    [pull] = output.kv_connector_metadata.reqs_to_recv["prefill-engine"].values()
+    attn_pull, gdn_pull = pull.local_block_ids
+    assert len(attn_pull) == expected_attn_pull_blocks
+    assert len(gdn_pull) == 1
+
+
 @pytest.mark.cpu_test
 def test_hybrid_gdn_remote_decode_truncates_prefill_before_cache_lookup():
     vllm_config = create_vllm_config(
diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py b/tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py
index 5a1a00158691..c360d40ea5e1 100644
--- a/tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py
+++ b/tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py
@@ -453,7 +453,9 @@ def test_store_mask_retention_interval_zero_eagle_round_trip(num_prompt_tokens):
     exists: set[tuple[int, bytes]] = set()
     store_masks = coord.store_mask(96, num_prompt_tokens=num_prompt_tokens)
     for gid, (group, mask) in enumerate(zip(groups, store_masks, strict=True)):
-        group_hashes = coord.block_hashes_for_spec(hashes, group.kv_cache_spec)
+        group_hashes = chunk_hashes_for_block_size(
+            hashes, coord.hash_block_size, group.kv_cache_spec.block_size
+        )
         for chunk_id, block_hash in enumerate(group_hashes):
             if mask is None or mask[chunk_id]:
                 exists.add((gid, bytes(block_hash)))
@@ -485,7 +487,9 @@ def test_load_mask_restores_mtp_replay_tail(extra_retained_tokens):
     exists: set[tuple[int, bytes]] = set()
     store_masks = coord.store_mask(96, num_prompt_tokens=96)
     for gid, (group, mask) in enumerate(zip(groups, store_masks, strict=True)):
-        group_hashes = coord.block_hashes_for_spec(hashes, group.kv_cache_spec)
+        group_hashes = chunk_hashes_for_block_size(
+            hashes, coord.hash_block_size, group.kv_cache_spec.block_size
+        )
         for chunk_id, block_hash in enumerate(group_hashes):
             if mask is None or mask[chunk_id]:
                 exists.add((gid, bytes(block_hash)))
@@ -740,7 +744,9 @@ def test_dsv4_five_group_eagle_store_lookup_round_trip():
     exists: set[tuple[int, bytes]] = set()
     store_masks = coord.store_mask(token_len)
     for gid, (group, mask) in enumerate(zip(groups, store_masks, strict=True)):
-        group_hashes = coord.block_hashes_for_spec(hashes, group.kv_cache_spec)
+        group_hashes = chunk_hashes_for_block_size(
+            hashes, coord.hash_block_size, group.kv_cache_spec.block_size
+        )
         for chunk_id, block_hash in enumerate(group_hashes):
             if mask is None or mask[chunk_id]:
                 exists.add((gid, bytes(block_hash)))
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 0682f9f465f1..cf5bc9073604 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
@@ -3,11 +3,14 @@
 """End-to-end save->lookup test for MooncakeStoreConnector on a hybrid
 (SWA + Full) attention config, using a dict-backed mock store."""
 
+import ctypes
 import sys
 import threading
 import types
+from types import SimpleNamespace
 from unittest.mock import MagicMock, patch
 
+import pytest
 import torch
 
 from tests.v1.attention.utils import dense_kv_cache_views
@@ -20,10 +23,12 @@
 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (
     BoundaryPut,
     ChunkedTokenDatabase,
-    KeyMetadata,
     LoadSpec,
     ReqMeta,
 )
+from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.layout import (  # noqa: E501
+    KeyMetadata,
+)
 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.scheduler import (  # noqa: E501
     _partial_tail_non_mamba_puts,
 )
@@ -31,6 +36,7 @@
     KVCacheStoreRecvingThread,
     KVCacheStoreSendingThread,
 )
+from vllm.v1.core.kv_cache_planning import get_kv_cache_config_from_groups
 from vllm.v1.core.kv_cache_utils import BlockHash
 from vllm.v1.kv_cache_interface import (
     CircularBufferSpec,
@@ -39,9 +45,11 @@
     KVCacheGroupSpec,
     KVCacheTensor,
     MambaSpec,
+    MLAAttentionSpec,
     SlidingWindowSpec,
 )
 from vllm.v1.kv_cache_layout import KVCacheLayout
+from vllm.v1.worker.utils import allocate_kv_cache
 
 
 class _DictStore:
@@ -149,6 +157,33 @@ def _build_worker_with_dict_store(vllm_config, kv_cache_config, store):
     return worker
 
 
+def _register_with_fake_threads(worker, kv_caches):
+    """Register kv_caches without starting the transfer threads."""
+
+    def _fake_thread_init(*args, **kwargs):
+        """Mock thread that sets all threading.Event args so waits don't block."""
+        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)
+
+
 def test_e2e_swa_plus_full_save_then_lookup_hits():
     """E2E: build a SWA+Full hybrid worker, save all blocks via the sending
     thread (synchronously), then verify lookup returns the full hit length.
@@ -201,28 +236,13 @@ def test_e2e_swa_plus_full_save_then_lookup_hits():
         "L1": dense_kv_cache_views(raw_swa, swa, 4, 1, KVCacheLayout.LBNHC)[0],
     }
 
-    def _fake_thread_init(*args, **kwargs):
-        """Mock thread that sets all threading.Event args so waits don't block."""
-        for v in list(args) + list(kwargs.values()):
-            if isinstance(v, threading.Event):
-                v.set()
-        m = MagicMock()
-        m.start = lambda: None
-        return m
+    _register_with_fake_threads(worker, kv_caches)
 
-    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)
+    # Each group's store values hold only its own layer.
+    assert [db.store_layout.kv_caches_base_addr for db in worker.token_dbs] == [
+        [raw_full.data_ptr()],
+        [raw_swa.data_ptr()],
+    ]
 
     # Now build a real sending thread (no .start()) over the worker's token_dbs
     # and the dict-backed store, so _handle_request runs synchronously.
@@ -283,6 +303,245 @@ def _fake_thread_init(*args, **kwargs):
     assert worker.lookup(num_tokens=64, block_hashes=hs).hit_length == 0
 
 
+def test_hybrid_mamba_values_hold_only_group_states():
+    """Hybrid groups overlay one allocation. A Mamba group's values hold only
+    its own layers' states: not the layer slot only attention uses, and not
+    the padding of each state up to the attention page. Restoring the value
+    brings back the states."""
+    attn = FullAttentionSpec(
+        block_size=16, num_kv_heads=1, head_size=16, dtype=torch.float16
+    )
+    page = attn.page_size_bytes
+    mamba = MambaSpec(
+        block_size=16,
+        shapes=((4, 8), (2, 8)),
+        dtypes=(torch.float32, torch.float16),
+        page_size_padded=page,
+        mamba_cache_mode="align",
+    )
+    assert mamba.state_content_size_bytes == 160 < page
+    num_layers, num_blocks = 3, 4
+
+    def tensor(layers):
+        return KVCacheTensor(
+            size=num_layers * num_blocks * page,
+            layers=layers,
+            layer_stride=num_blocks * page,
+            block_stride=page,
+        )
+
+    cfg = KVCacheConfig(
+        num_blocks=num_blocks,
+        kv_cache_tensors=[tensor(["A0", "A1", "A2"]), tensor(["M0", "M1"])],
+        kv_cache_groups=[
+            KVCacheGroupSpec(["A0", "A1", "A2"], attn),
+            KVCacheGroupSpec(["M0", "M1"], mamba),
+        ],
+    )
+    worker = _build_worker_with_dict_store(
+        _minimal_vllm_config(cache_block_size=16), cfg, _DictStore()
+    )
+    kv_caches = allocate_kv_cache(cfg, torch.device("cpu"), KVCacheLayout.LBNHC)
+    # Like a Mamba backend, view each layer's page as its two states, here
+    # with a gap between them.
+    for name in ("M0", "M1"):
+        page_view = kv_caches[name]
+        kv_caches[name] = [
+            page_view[..., :128].view(torch.float32),
+            page_view[..., 136:160].view(torch.float16),
+        ]
+    _register_with_fake_threads(worker, kv_caches)
+
+    states = [state for name in ("M0", "M1") for state in kv_caches[name]]
+    state_lens = [state[0].numel() * state.element_size() for state in states]
+    attn_db, mamba_db = worker.token_dbs
+    attn_layout, mamba_layout = attn_db.store_layout, mamba_db.store_layout
+    assert attn_layout.block_len == [page] * num_layers
+    assert mamba_layout.kv_caches_base_addr == [s.data_ptr() for s in states]
+    assert mamba_layout.block_len == state_lens
+
+    raw = torch.empty(0, dtype=torch.int8).set_(kv_caches["A0"].untyped_storage())
+    raw.copy_(torch.randint(-128, 128, raw.shape, dtype=torch.int8))
+    before = raw.clone()
+    chunk_addrs, chunk_sizes, _ = mamba_layout.prepare_values([(0, 16)], [2], [0])
+    addrs, sizes = chunk_addrs[0], chunk_sizes[0]
+    saved = [ctypes.string_at(addr, size) for addr, size in zip(addrs, sizes)]
+    # The value is each state of block 2 in turn.
+    assert saved == [
+        state[2].contiguous().view(torch.int8).numpy().tobytes() for state in states
+    ]
+
+    # Clobber block 2 of every layer slot, then load the saved value: only the
+    # Mamba states come back; the gap, the padding and slot A2 keep the clobber.
+    def block_2(addr, length):
+        start = addr - raw.data_ptr() + 2 * page
+        return slice(start, start + length)
+
+    for name in ("A0", "A1", "A2"):
+        raw[block_2(kv_caches[name].data_ptr(), page)] = 7
+    expected = raw.clone()
+    for state, length in zip(states, state_lens):
+        expected[block_2(state.data_ptr(), length)] = before[
+            block_2(state.data_ptr(), length)
+        ]
+    for addr, data in zip(addrs, saved):
+        ctypes.memmove(addr, data, len(data))
+    assert torch.equal(raw, expected)
+
+
+_ATTN = FullAttentionSpec(
+    block_size=16, num_kv_heads=2, head_size=16, dtype=torch.float16
+)
+_MAMBA = MambaSpec(
+    block_size=16,
+    shapes=((4, 8), (2, 8)),
+    dtypes=(torch.float32, torch.float16),
+    page_size_padded=_ATTN.page_size_bytes,
+    mamba_cache_mode="align",
+)
+# Model shapes as KV cache groups: each group is (layer names, spec).
+_MODELS = {
+    "dense": [(["A0", "A1", "A2"], _ATTN)],
+    "mla": [
+        (
+            ["A0", "A1", "A2"],
+            MLAAttentionSpec(
+                block_size=16, num_kv_heads=1, head_size=40, dtype=torch.float16
+            ),
+        )
+    ],
+    "full_swa": [
+        (["A0", "A1"], _ATTN),
+        (
+            ["S0", "S1"],
+            SlidingWindowSpec(
+                block_size=16,
+                num_kv_heads=2,
+                head_size=16,
+                dtype=torch.float16,
+                sliding_window=32,
+            ),
+        ),
+    ],
+    # Like Kimi K3: Mamba groups with fewer layers, states padded to the page.
+    "attn_mamba": [
+        (["A0", "A1", "A2"], _ATTN),
+        (["M0", "M1"], _MAMBA),
+        (["M2", "M3"], _MAMBA),
+    ],
+    # Pages of different sizes are packed side by side within a block.
+    "mixed_pages": [
+        (["A0", "A1"], _ATTN),
+        (
+            ["B0", "B1", "B2"],
+            FullAttentionSpec(
+                block_size=16, num_kv_heads=1, head_size=8, dtype=torch.float16
+            ),
+        ),
+    ],
+}
+
+
+def _byte_mask(t: torch.Tensor, nbytes: int) -> torch.Tensor:
+    """The bytes of the backing storage that ``t`` covers."""
+    offsets = torch.tensor([t.storage_offset()])
+    for size, stride in zip(t.shape, t.stride()):
+        offsets = (offsets[:, None] + torch.arange(size) * stride).flatten()
+    elem = t.element_size()
+    mask = torch.zeros(nbytes, dtype=torch.bool)
+    mask[(offsets[:, None] * elem + torch.arange(elem)).flatten()] = True
+    return mask
+
+
+def _span_mask(spans, nbytes: int) -> torch.Tensor:
+    mask = torch.zeros(nbytes, dtype=torch.bool)
+    for start, length in spans:
+        mask[start : start + length] = True
+    return mask
+
+
+@pytest.mark.parametrize("layout", list(KVCacheLayout), ids=lambda x: x.name)
+@pytest.mark.parametrize("model", list(_MODELS))
+def test_store_values_round_trip_every_layout(model, layout):
+    """Through the planner's real layout, each group's value for a block holds
+    every byte of its layers' views of that block, reads and writes only that
+    block, and restores it exactly. Without padding the value is the group's
+    whole block, as before; with Mamba padding it is exactly the states."""
+    num_blocks, block = 4, 2
+    groups = [KVCacheGroupSpec(names, spec) for names, spec in _MODELS[model]]
+    planner_config = SimpleNamespace(
+        cache_config=SimpleNamespace(
+            get_resolved_kv_cache_layout=lambda: layout,
+            num_gpu_blocks_override=num_blocks,
+            prefix_cache_retention_interval=None,
+        ),
+        attention_config=SimpleNamespace(hisparse_config=None),
+    )
+    try:
+        cfg = get_kv_cache_config_from_groups(planner_config, groups, 1 << 30)
+        kv_caches = allocate_kv_cache(cfg, torch.device("cpu"), layout)
+    except (ValueError, AssertionError):
+        pytest.skip(f"{layout.name} cannot express {model}")
+    worker = _build_worker_with_dict_store(
+        _minimal_vllm_config(cache_block_size=16), cfg, _DictStore()
+    )
+    _register_with_fake_threads(worker, kv_caches)
+
+    first = next(iter(kv_caches.values()))
+    raw = torch.empty(0, dtype=torch.int8).set_(first.untyped_storage())
+    nbytes = raw.numel()
+    raw.copy_(torch.randint(-128, 128, raw.shape, dtype=torch.int8))
+    before = raw.clone()
+    base = raw.data_ptr()
+
+    # The block of every byte that some layer's view uses (-1: none, e.g.
+    # padding). A value must not reach another block's bytes.
+    byte_block = torch.full((nbytes,), -1)
+    for views in kv_caches.values():
+        for view in views if isinstance(views, list) else [views]:
+            for b in range(num_blocks):
+                mask = _byte_mask(view[b], nbytes)
+                if ((byte_block[mask] != -1) & (byte_block[mask] != b)).any():
+                    pytest.skip(f"{layout.name} overlaps blocks of {model}")
+                byte_block[mask] = b
+    other_blocks = (byte_block != block) & (byte_block != -1)
+
+    for group, db in zip(groups, worker.token_dbs, strict=True):
+        store_layout = db.store_layout
+        chunk_addrs, chunk_sizes, _ = store_layout.prepare_values(
+            [(0, 16), (16, 32)], [block, block], [0, 0]
+        )
+        addrs, sizes = chunk_addrs[0], chunk_sizes[0]
+        assert chunk_addrs[1] == addrs and chunk_sizes[1] == sizes
+
+        value = _span_mask([(a - base, n) for a, n in zip(addrs, sizes)], nbytes)
+        owned = torch.zeros(nbytes, dtype=torch.bool)
+        for name in group.layer_names:
+            views = kv_caches[name]
+            for view in views if isinstance(views, list) else [views]:
+                owned |= _byte_mask(view[block], nbytes)
+        assert not (value & other_blocks).any(), "value reaches other blocks"
+        assert not (owned & ~value).any(), "value misses bytes of the group"
+        # Each piece ends at the last byte the group uses in it.
+        assert all(owned[a - base + n - 1] for a, n in zip(addrs, sizes))
+        if len(groups) == 1 and layout != KVCacheLayout.BHLNC:
+            # A single group stores whole blocks, as before. (A BHLNC head's
+            # block stride spans every head, so only the head's block is.)
+            assert store_layout.block_len == store_layout.block_stride
+        if not layout.is_block_outermost:
+            # Exactly the group's bytes: no padding, no other group's slots.
+            assert torch.equal(value, owned) and sum(sizes) == owned.sum()
+
+        # Clobber everything but the other blocks, then load the saved value.
+        saved = [ctypes.string_at(a, n) for a, n in zip(addrs, sizes)]
+        raw[~other_blocks] = 7
+        for addr, data in zip(addrs, saved):
+            ctypes.memmove(addr, data, len(data))
+        assert torch.equal(raw[owned], before[owned])
+        assert torch.equal(raw[other_blocks], before[other_blocks])
+        raw.copy_(before)
+
+
 def test_recv_skips_swa_blocks_before_window():
     """Producer stored every block for both groups; consumer must only fetch
     SWA blocks within the sliding window, not the head."""
@@ -301,8 +560,8 @@ def test_recv_skips_swa_blocks_before_window():
     ]
     md0 = KeyMetadata("m", 0, 0, 0, 0, group_id=0)
     md1 = KeyMetadata("m", 0, 0, 0, 0, group_id=1)
-    db_full = ChunkedTokenDatabase(md0, block_size=16, hash_block_size=16)
-    db_swa = ChunkedTokenDatabase(md1, block_size=16, hash_block_size=16)
+    db_full = ChunkedTokenDatabase(md0, chunk_size=16, hash_block_size=16)
+    db_swa = ChunkedTokenDatabase(md1, chunk_size=16, hash_block_size=16)
     db_full.set_kv_caches_base_addr([0])
     db_full.set_block_len([1024])
     db_swa.set_kv_caches_base_addr([1 << 20])
@@ -355,7 +614,7 @@ def test_chunked_token_database_hash_block_size_smaller_than_block_size():
     """DSv4-style: hash_block_size=4, group block_size=16 — process_tokens
     keys each chunk by its ending fine hash, including a partial tail."""
     md = KeyMetadata("m", 0, 0, 0, 0, group_id=3)
-    db = ChunkedTokenDatabase(md, block_size=16, hash_block_size=4)
+    db = ChunkedTokenDatabase(md, chunk_size=16, hash_block_size=4)
     db.set_kv_caches_base_addr([0])
     db.set_block_len([512])
     fine_hashes = [BlockHash(bytes([i + 1]) * 4) for i in range(8)]
@@ -386,7 +645,7 @@ def _resolve_partial_tail(thread, req: ReqMeta) -> ReqMeta:
     """Add the tail's non-Mamba puts as the scheduler does before the worker."""
     req.boundary_puts = [BoundaryPut(*put) for put in req.boundary_puts or []]
     req.boundary_puts[:0] = _partial_tail_non_mamba_puts(
-        thread.coord, req, [db.block_size for db in thread.token_databases]
+        thread.coord, req, [db.store_layout.block_size for db in thread.token_databases]
     )
     return req
 
@@ -430,7 +689,7 @@ def batch_put_from_multi_buffers(self, keys, addrs, sizes, *a, **k):
     for g_idx in range(2):
         db = ChunkedTokenDatabase(
             KeyMetadata("m", 0, 0, 0, 0, group_id=g_idx),
-            block_size=16,
+            chunk_size=16,
             hash_block_size=4,
         )
         db.set_kv_caches_base_addr([g_idx * 10_000])
@@ -469,8 +728,8 @@ def batch_put_from_multi_buffers(self, keys, addrs, sizes, *a, **k):
 
     # boundary = 12 // 4 * 4 = 12 -> keyed by hs[12 // 4 - 1] = hs[2].
     partial_hash = hs[2]
-    fa_key = token_dbs[0].key_for(partial_hash)
-    mamba_key = token_dbs[1].key_for(partial_hash)
+    fa_key = token_dbs[0].store_layout.key_for(0, partial_hash)
+    mamba_key = token_dbs[1].store_layout.key_for(0, partial_hash)
     assert set(store.puts) == {fa_key, mamba_key}
     # FA reads block_ids[0][0] = block 1: addr = base(0) + 1 * 512.
     assert store.puts[fa_key] == [512]
@@ -511,7 +770,7 @@ def batch_put_from_multi_buffers(self, keys, addrs, sizes, *a, **k):
     for g_idx in range(2):
         db = ChunkedTokenDatabase(
             KeyMetadata("m", 0, 0, 0, 0, group_id=g_idx),
-            block_size=16,
+            chunk_size=16,
             hash_block_size=4,
         )
         db.set_kv_caches_base_addr([g_idx * 10_000])
@@ -594,7 +853,7 @@ def batch_put_from_multi_buffers(self, keys, addrs, sizes, *a, **k):
     for g_idx, block_size in enumerate([4, 16]):
         db = ChunkedTokenDatabase(
             KeyMetadata("m", 0, 0, 0, 0, group_id=g_idx),
-            block_size=block_size,
+            chunk_size=block_size,
             hash_block_size=4,
         )
         db.set_kv_caches_base_addr([g_idx * 10_000])
@@ -632,9 +891,9 @@ def batch_put_from_multi_buffers(self, keys, addrs, sizes, *a, **k):
     # block-end hashes; mamba (block 16): the partial boundary block under
     # the boundary sub-hash, read from the CoW target.
     expected = {
-        token_dbs[0].key_for(hs[0]): [1 * 512],
-        token_dbs[0].key_for(hs[1]): [2 * 512],
-        token_dbs[1].key_for(hs[1]): [10_000 + mamba_cow_block * 512],
+        token_dbs[0].store_layout.key_for(0, hs[0]): [1 * 512],
+        token_dbs[0].store_layout.key_for(0, hs[1]): [2 * 512],
+        token_dbs[1].store_layout.key_for(0, hs[1]): [10_000 + mamba_cow_block * 512],
     }
     assert store.puts == expected
 
@@ -813,7 +1072,7 @@ def test_worker_setup_tolerates_finer_scratch_group():
         group_participates=[True, True],
     )
 
-    # Persist the sub-block partial tail at boundary 8 (keyed by hs[8//8-1]).
+    # Persist a hash-aligned tail within the 16-token physical block.
     hs = [BlockHash(bytes([i + 1]) * 8) for i in range(3)]
     req = ReqMeta(
         req_id="r0",
diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_layout.py b/tests/v1/kv_connector/unit/test_mooncake_store_layout.py
index 00ff3eea4d58..e56c81d7b276 100644
--- a/tests/v1/kv_connector/unit/test_mooncake_store_layout.py
+++ b/tests/v1/kv_connector/unit/test_mooncake_store_layout.py
@@ -3,23 +3,242 @@
 """Tests for Mooncake Store payload layouts."""
 
 import ctypes
+import os
 import random
 
 import pytest
 import torch
 
-from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (
+from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (  # noqa: E501
     ChunkedTokenDatabase,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.layout import (  # noqa: E501
+    AttentionStoreLayout,
+    HeadMajorStoreLayout,
     KeyMetadata,
-    LBHNCStoreLayout,
-    LBNHCStoreLayout,
+    MambaStoreLayout,
     RankLocalStoreLayout,
+    TokenMajorStoreLayout,
 )
+from vllm.model_executor.layers.mamba.mamba_utils import get_conv_state_layout
 from vllm.utils.math_utils import cdiv
+from vllm.v1.core.kv_cache_utils import BlockHash
+from vllm.v1.kv_cache_layout import KVCacheLayout
 
 BLOCK_SIZE = 128
 
 
+@pytest.fixture(autouse=True)
+def _reset_conv_state_layout_cache():
+    get_conv_state_layout.cache_clear()
+    yield
+    get_conv_state_layout.cache_clear()
+
+
+def _attention_specs(
+    num_layers: int, block_size: int, local_heads: int, content_size: int = 4
+):
+    from vllm.v1.kv_cache_interface import FullAttentionSpec
+
+    return tuple(
+        FullAttentionSpec(
+            block_size=block_size,
+            num_kv_heads=local_heads,
+            head_size=content_size // 2,
+            dtype=torch.float16,
+        )
+        for _ in range(num_layers)
+    )
+
+
+def _make_state_store_layout(
+    spec, *, local_tp_size: int, tp_rank: int
+) -> tuple[MambaStoreLayout, torch.Tensor]:
+    tensor = torch.empty((2, 1, 1, spec.state_content_size_bytes), dtype=torch.uint8)
+    metadata = KeyMetadata(
+        "test-model",
+        tp_rank,
+        0,
+        0,
+        0,
+        group_id=1,
+        store_namespace=(
+            "@store_tp:4@store_pp:1@store_format:mamba_state@store_schema:test"
+        ),
+    )
+    layout = MambaStoreLayout(
+        metadata,
+        block_size=16,
+        hash_block_size=16,
+        local_tp_size=local_tp_size,
+        store_tp_size=4,
+        tp_rank=tp_rank,
+        layer_specs=(spec,),
+    )
+    layout.register_kv_caches([tensor], 2)
+    return layout, tensor
+
+
+def _make_gdn_store_layout(
+    *, local_tp_size: int, tp_rank: int
+) -> tuple[MambaStoreLayout, torch.Tensor]:
+    from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
+    from vllm.v1.kv_cache_interface import MambaSpec
+
+    local_factor = 4 // local_tp_size
+    conv_shape = (6 * local_factor, 3)
+    if os.environ.get("VLLM_SSM_CONV_STATE_LAYOUT") == "SD":
+        conv_shape = conv_shape[::-1]
+    spec = MambaSpec(
+        block_size=16,
+        shapes=(conv_shape, (local_factor, 2, 2)),
+        dtypes=(torch.uint8, torch.uint8),
+        mamba_type=MambaAttentionBackendEnum.GDN_ATTN,
+    )
+    return _make_state_store_layout(spec, local_tp_size=local_tp_size, tp_rank=tp_rank)
+
+
+def _make_mamba2_store_layout(
+    *, local_tp_size: int, tp_rank: int
+) -> tuple[MambaStoreLayout, torch.Tensor]:
+    from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
+    from vllm.v1.kv_cache_interface import MambaSpec
+
+    local_heads = 4 // local_tp_size
+    effective_groups = local_tp_size
+    conv_shape = ((8 + 2 * effective_groups * 2) // local_tp_size, 3)
+    if os.environ.get("VLLM_SSM_CONV_STATE_LAYOUT") == "SD":
+        conv_shape = conv_shape[::-1]
+    spec = MambaSpec(
+        block_size=16,
+        shapes=(
+            conv_shape,
+            (local_heads, 2, 2),
+        ),
+        dtypes=(torch.uint8, torch.uint8),
+        mamba_type=MambaAttentionBackendEnum.MAMBA2,
+    )
+    return _make_state_store_layout(spec, local_tp_size=local_tp_size, tp_rank=tp_rank)
+
+
+def _descriptors_for_block(
+    layout: MambaStoreLayout, shard_id: int
+) -> tuple[list[int], list[int]]:
+    addrs, sizes, _ = layout.prepare_values([(0, layout.block_size)], [0], [shard_id])
+    return addrs[0], sizes[0]
+
+
+def _read_segments(addrs: list[int], sizes: list[int]) -> bytes:
+    return b"".join(
+        ctypes.string_at(addr, size) for addr, size in zip(addrs, sizes, strict=True)
+    )
+
+
+def test_tp_shared_layout_loads_partial_prefix_from_physical_block():
+    metadata = KeyMetadata("test-model", 0, 0, 0, 0)
+    layout = HeadMajorStoreLayout(
+        metadata,
+        16,
+        16,
+        local_tp_size=2,
+        store_tp_size=4,
+        tp_rank=0,
+        layer_specs=_attention_specs(1, 16, 4),
+    )
+    tensor = torch.empty((2, 4, 16, 4), dtype=torch.float16)
+    layout.register_kv_caches([tensor], 2)
+    shard_ids = layout.local_shard_ids
+
+    addrs, sizes, block_ids = layout.prepare_values(
+        [(0, 8)] * len(shard_ids), [1], shard_ids
+    )
+
+    assert len(addrs) == len(sizes) == len(shard_ids)
+    assert block_ids == [1] * len(shard_ids)
+
+
+def _assert_state_store_round_trip(layout_factory):
+    block_hash = BlockHash(b"h")
+    for producer_tp, consumer_tp in ((4, 2), (2, 4)):
+        stored: dict[int, bytes] = {}
+        keys: dict[int, str] = {}
+        schemas: set[str] = set()
+        for tp_rank in range(producer_tp):
+            layout, cache = layout_factory(local_tp_size=producer_tp, tp_rank=tp_rank)
+            schemas.add(
+                MambaStoreLayout.schema_fingerprint(layout.layer_specs, producer_tp, 4)
+            )
+            cache.zero_()
+            for shard_id in layout.local_shard_ids:
+                addrs, sizes = _descriptors_for_block(layout, shard_id)
+                for segment_index, (addr, size) in enumerate(
+                    zip(addrs, sizes, strict=True)
+                ):
+                    value = 32 + segment_index * 16
+                    if segment_index in (0, 3):
+                        value += shard_id
+                    ctypes.memset(addr, value, size)
+                stored[shard_id] = _read_segments(addrs, sizes)
+                keys[shard_id] = layout.key_for(shard_id, block_hash)
+
+        for tp_rank in range(consumer_tp):
+            layout, cache = layout_factory(local_tp_size=consumer_tp, tp_rank=tp_rank)
+            schemas.add(
+                MambaStoreLayout.schema_fingerprint(layout.layer_specs, consumer_tp, 4)
+            )
+            cache.zero_()
+            for shard_id in layout.local_shard_ids:
+                assert layout.key_for(shard_id, block_hash) == keys[shard_id]
+                addrs, sizes = _descriptors_for_block(layout, shard_id)
+                offset = 0
+                for addr, size in zip(addrs, sizes, strict=True):
+                    ctypes.memmove(addr, stored[shard_id][offset : offset + size], size)
+                    offset += size
+                assert _read_segments(addrs, sizes) == stored[shard_id]
+        assert set(stored) == set(range(4))
+        assert len(schemas) == 1
+
+
+@pytest.mark.parametrize("conv_layout", ["DS", "SD"])
+def test_gdn_store_shards_round_trip_in_both_tp_directions(monkeypatch, conv_layout):
+    monkeypatch.setenv("VLLM_SSM_CONV_STATE_LAYOUT", conv_layout)
+    _assert_state_store_round_trip(_make_gdn_store_layout)
+
+
+@pytest.mark.parametrize("conv_layout", ["DS", "SD"])
+def test_mamba2_replicated_groups_round_trip_in_both_tp_directions(
+    monkeypatch, conv_layout
+):
+    monkeypatch.setenv("VLLM_SSM_CONV_STATE_LAYOUT", conv_layout)
+    _assert_state_store_round_trip(_make_mamba2_store_layout)
+
+
+@pytest.mark.parametrize(
+    ("conv_layout", "sizes", "shard0_offsets", "shard1_offsets"),
+    [
+        ("DS", [6, 6, 6, 4], [0, 12, 24, 36], [6, 18, 30, 40]),
+        (
+            "SD",
+            [2] * 9 + [4],
+            [0, 4, 8, 12, 16, 20, 24, 28, 32, 36],
+            [2, 6, 10, 14, 18, 22, 26, 30, 34, 40],
+        ),
+    ],
+)
+def test_gdn_store_shard_segments_preserve_projection_boundaries(
+    monkeypatch, conv_layout, sizes, shard0_offsets, shard1_offsets
+):
+    monkeypatch.setenv("VLLM_SSM_CONV_STATE_LAYOUT", conv_layout)
+    layout, cache = _make_gdn_store_layout(local_tp_size=2, tp_rank=0)
+
+    shard0_addrs, shard0_sizes = _descriptors_for_block(layout, 0)
+    shard1_addrs, shard1_sizes = _descriptors_for_block(layout, 1)
+
+    assert shard0_sizes == shard1_sizes == sizes
+    assert [addr - cache.data_ptr() for addr in shard0_addrs] == shard0_offsets
+    assert [addr - cache.data_ptr() for addr in shard1_addrs] == shard1_offsets
+
+
 def _rank_local_layout(num_regions: int, num_block_lens: int) -> RankLocalStoreLayout:
     metadata = KeyMetadata("test-model", 1, 0, 0, 0)
     layout = RankLocalStoreLayout(metadata, BLOCK_SIZE, BLOCK_SIZE)
@@ -34,9 +253,9 @@ def _rank_local_reference(
     layout: RankLocalStoreLayout, start: int, end: int, block_ids: list[int]
 ) -> tuple[list[int], list[int], int]:
     block_id = block_ids[start // layout.block_size]
-    length = len(layout.block_len)
+    length = len(layout.block_stride)
     addrs = [
-        base_addr + block_id * layout.block_len[index % length]
+        base_addr + block_id * layout.block_stride[index % length]
         for index, base_addr in enumerate(layout.kv_caches_base_addr)
     ]
     sizes = [
@@ -79,69 +298,145 @@ def test_rank_local_descriptors_handle_empty_and_invalid_chunks():
         layout.prepare_values([(0, BLOCK_SIZE + 1)], [0, 1], [0])
 
 
-@pytest.mark.parametrize(
-    ("layout_cls", "store_format"),
-    [
-        (LBHNCStoreLayout, "tp_shared_lbhnc"),
-        (LBNHCStoreLayout, "tp_shared_lbnhc"),
-    ],
-)
-def test_tp_shared_layout_owns_store_namespace(layout_cls, store_format):
-    assert (
-        layout_cls.shared_namespace(4, 2)
-        == f"@store_tp:4@store_pp:2@store_format:{store_format}"
+def test_rank_local_layout_skips_block_padding():
+    """A state padded up to a larger page, like a Mamba state padded to the
+    attention page size, is stored without the padding."""
+    num_layers, num_blocks, page, state = 2, 3, 64, 40
+    # Layer-outermost buffer; each layer views its state at the start of
+    # every page.
+    raw = torch.zeros(num_layers * num_blocks * page, dtype=torch.int8)
+    views = [
+        torch.as_strided(
+            raw,
+            (num_blocks, 1, 1, state),
+            (page, state, state, 1),
+            layer * num_blocks * page,
+        )
+        for layer in range(num_layers)
+    ]
+    layout = RankLocalStoreLayout(
+        KeyMetadata("test-model", 0, 0, 0, 0), BLOCK_SIZE, BLOCK_SIZE
+    )
+    layout.register_kv_caches(views, num_blocks)
+
+    assert layout.kv_caches_base_addr == [view.data_ptr() for view in views]
+    assert layout.block_stride == [page] * num_layers
+    assert layout.block_len == [state] * num_layers
+    addrs, sizes, _ = layout.prepare_values(
+        [(0, BLOCK_SIZE), (BLOCK_SIZE, 3 * BLOCK_SIZE)], [2, 0, 1], [0, 0]
     )
+    assert addrs == [
+        [view.data_ptr() + 2 * page for view in views],
+        [view.data_ptr() for view in views],
+    ]
+    # A two-block chunk copies the padding between its blocks.
+    assert sizes == [[state] * num_layers, [page + state] * num_layers]
+    addrs, sizes, _ = layout.prepare_values([(0, BLOCK_SIZE)], [1], [0])
+    assert addrs == [[view.data_ptr() + page for view in views]]
+    assert sizes == [[state] * num_layers]
 
 
 def test_rank_local_database_api_rejects_tp_shared_layout():
     metadata = KeyMetadata("test-model", 0, 0, 0, 0)
-    layout = LBHNCStoreLayout(
+    layout = HeadMajorStoreLayout(
         metadata,
         block_size=16,
         hash_block_size=16,
         local_tp_size=2,
         store_tp_size=4,
         tp_rank=0,
-        num_kv_heads=8,
+        layer_specs=_attention_specs(1, 16, 4),
     )
     database = ChunkedTokenDatabase(metadata, 16, store_layout=layout)
 
     with pytest.raises(RuntimeError, match="rank-local"):
-        database.prepare_value_for_block(0)
+        database.set_block_len([256])
+
+
+_SHARED_LAYOUTS = [
+    (KVCacheLayout.LBHNC, HeadMajorStoreLayout),
+    (KVCacheLayout.LBNHC, TokenMajorStoreLayout),
+    (KVCacheLayout.BLHNC, HeadMajorStoreLayout),
+    (KVCacheLayout.BLNHC, TokenMajorStoreLayout),
+    (KVCacheLayout.LHBNC, HeadMajorStoreLayout),
+    (KVCacheLayout.BHLNC, HeadMajorStoreLayout),
+]
+
+
+def _physical_layer_views(
+    layout: KVCacheLayout,
+    num_layers: int,
+    num_blocks: int,
+    num_heads: int,
+    block_size: int,
+    content_size: int,
+) -> tuple[torch.Tensor, list[torch.Tensor]]:
+    logical_shape = (num_layers, num_blocks, num_heads, block_size, content_size)
+    axis_order = layout.value
+    physical = torch.empty(
+        tuple(logical_shape[axis] for axis in axis_order), dtype=torch.float16
+    )
+    logical = physical.permute(tuple(axis_order.index(axis) for axis in range(5)))
+    return physical, list(logical.unbind(0))
 
 
 @pytest.mark.parametrize(
-    ("layout_cls", "producer_strides", "consumer_strides"),
-    [
-        (LBHNCStoreLayout, (128, 64, 4, 1), (256, 64, 4, 1)),
-        (LBNHCStoreLayout, (128, 4, 8, 1), (256, 4, 16, 1)),
+    ("producer_layout", "producer_cls", "consumer_layout", "consumer_cls"),
+    [(layout, layout_cls, layout, layout_cls) for layout, layout_cls in _SHARED_LAYOUTS]
+    + [
+        (
+            KVCacheLayout.LBHNC,
+            HeadMajorStoreLayout,
+            KVCacheLayout.BLHNC,
+            HeadMajorStoreLayout,
+        ),
+        (
+            KVCacheLayout.LHBNC,
+            HeadMajorStoreLayout,
+            KVCacheLayout.BHLNC,
+            HeadMajorStoreLayout,
+        ),
+        (
+            KVCacheLayout.LBNHC,
+            TokenMajorStoreLayout,
+            KVCacheLayout.BLNHC,
+            TokenMajorStoreLayout,
+        ),
     ],
 )
 def test_tp_shared_layout_round_trip_across_tp_sizes(
-    layout_cls, producer_strides, consumer_strides
+    producer_layout: KVCacheLayout,
+    producer_cls,
+    consumer_layout: KVCacheLayout,
+    consumer_cls,
 ):
     block_size = 16
-    shape = (1, 2, block_size, 4)
+    num_layers = 2
     stored: dict[int, bytes] = {}
-    producer_tensors = []
+    producer_layers = []
 
     for tp_rank in range(4):
-        tensor = torch.empty_strided(shape, producer_strides, dtype=torch.float16)
-        tensor.copy_(
-            torch.arange(128, dtype=torch.float16).view(shape) + tp_rank * 1000
+        physical, layers = _physical_layer_views(
+            producer_layout, num_layers, 1, 2, block_size, 4
         )
-        producer_tensors.append(tensor)
+        for layer_index, layer in enumerate(layers):
+            layer.copy_(
+                torch.arange(layer.numel(), dtype=torch.float16).view_as(layer)
+                + tp_rank * 1000
+                + layer_index * 100
+            )
+        producer_layers.append(layers)
         metadata = KeyMetadata("test-model", tp_rank, 0, 0, 0)
-        layout = layout_cls(
+        layout = producer_cls(
             metadata,
             block_size,
             block_size,
             local_tp_size=4,
             store_tp_size=4,
             tp_rank=tp_rank,
-            num_kv_heads=8,
+            layer_specs=_attention_specs(num_layers, block_size, 2),
         )
-        layout.register_kv_caches([tensor], 1)
+        layout.register_kv_caches(layers, 1)
         addrs, sizes, _ = layout.prepare_values([(0, block_size)], [0], [tp_rank])
         stored[tp_rank] = b"".join(
             ctypes.string_at(addr, size)
@@ -149,21 +444,21 @@ def test_tp_shared_layout_round_trip_across_tp_sizes(
         )
 
     for tp_rank in range(2):
-        tensor = torch.empty_strided(
-            (1, 4, block_size, 4), consumer_strides, dtype=torch.float16
+        physical, layers = _physical_layer_views(
+            consumer_layout, num_layers, 1, 4, block_size, 4
         )
-        tensor.zero_()
+        physical.zero_()
         metadata = KeyMetadata("test-model", tp_rank, 0, 0, 0)
-        layout = layout_cls(
+        layout = consumer_cls(
             metadata,
             block_size,
             block_size,
             local_tp_size=2,
             store_tp_size=4,
             tp_rank=tp_rank,
-            num_kv_heads=8,
+            layer_specs=_attention_specs(num_layers, block_size, 4),
         )
-        layout.register_kv_caches([tensor], 1)
+        layout.register_kv_caches(layers, 1)
         shard_ids = layout.local_shard_ids
         addrs, sizes, _ = layout.prepare_values(
             [(0, block_size)] * len(shard_ids), [0], shard_ids
@@ -176,5 +471,184 @@ def test_tp_shared_layout_round_trip_across_tp_sizes(
                 ctypes.memmove(addr, stored[shard_id][offset : offset + size], size)
                 offset += size
 
-        expected = torch.cat(producer_tensors[tp_rank * 2 : tp_rank * 2 + 2], dim=1)
-        torch.testing.assert_close(tensor, expected)
+        for layer_index, layer in enumerate(layers):
+            expected = torch.cat(
+                [
+                    rank_layers[layer_index]
+                    for rank_layers in producer_layers[tp_rank * 2 : tp_rank * 2 + 2]
+                ],
+                dim=1,
+            )
+            torch.testing.assert_close(layer, expected)
+
+
+@pytest.mark.parametrize(
+    ("cache_layout", "layout_cls"),
+    [
+        (KVCacheLayout.LBHNC, HeadMajorStoreLayout),
+        (KVCacheLayout.LBNHC, TokenMajorStoreLayout),
+    ],
+)
+def test_attention_store_layout_round_trip_across_local_block_sizes(
+    cache_layout, layout_cls
+):
+    metadata = KeyMetadata("test-model", 0, 0, 0, 0)
+    chunks = [(0, 400), (400, 800)]
+
+    _, (small_cache,) = _physical_layer_views(cache_layout, 1, 2, 2, 400, 4)
+    small_cache.copy_(
+        torch.arange(small_cache.numel(), dtype=torch.float16).view_as(small_cache)
+    )
+    small_layout = layout_cls(
+        metadata,
+        block_size=400,
+        hash_block_size=16,
+        local_tp_size=1,
+        store_tp_size=1,
+        tp_rank=0,
+        layer_specs=_attention_specs(1, 400, 2),
+        store_chunk_size=400,
+    )
+    small_layout.register_kv_caches([small_cache], num_blocks=2)
+    addrs, sizes, _ = small_layout.prepare_values(chunks, [0, 1], [0, 0])
+    stored = [_read_segments(a, s) for a, s in zip(addrs, sizes, strict=True)]
+
+    _, (large_cache,) = _physical_layer_views(cache_layout, 1, 1, 2, 800, 4)
+    large_cache.zero_()
+    large_layout = layout_cls(
+        metadata,
+        block_size=800,
+        hash_block_size=16,
+        local_tp_size=1,
+        store_tp_size=1,
+        tp_rank=0,
+        layer_specs=_attention_specs(1, 800, 2),
+        store_chunk_size=400,
+    )
+    large_layout.register_kv_caches([large_cache], num_blocks=1)
+    addrs, sizes, _ = large_layout.prepare_values(chunks, [0], [0, 0])
+    for value, value_addrs, value_sizes in zip(stored, addrs, sizes, strict=True):
+        offset = 0
+        for addr, size in zip(value_addrs, value_sizes, strict=True):
+            ctypes.memmove(addr, value[offset : offset + size], size)
+            offset += size
+    expected = torch.cat((small_cache[0], small_cache[1]), dim=1)
+    torch.testing.assert_close(large_cache[0], expected)
+
+    large_values = [_read_segments(a, s) for a, s in zip(addrs, sizes, strict=True)]
+    _, (restored_small,) = _physical_layer_views(cache_layout, 1, 2, 2, 400, 4)
+    restored_small.zero_()
+    restored_layout = layout_cls(
+        metadata,
+        block_size=400,
+        hash_block_size=16,
+        local_tp_size=1,
+        store_tp_size=1,
+        tp_rank=0,
+        layer_specs=_attention_specs(1, 400, 2),
+        store_chunk_size=400,
+    )
+    restored_layout.register_kv_caches([restored_small], num_blocks=2)
+    addrs, sizes, _ = restored_layout.prepare_values(chunks, [0, 1], [0, 0])
+    for value, value_addrs, value_sizes in zip(large_values, addrs, sizes, strict=True):
+        offset = 0
+        for addr, size in zip(value_addrs, value_sizes, strict=True):
+            ctypes.memmove(addr, value[offset : offset + size], size)
+            offset += size
+    torch.testing.assert_close(restored_small, small_cache)
+
+
+def test_attention_store_chunk_is_independent_of_k3_local_pages():
+    from vllm.v1.kv_cache_interface import MLAAttentionSpec
+
+    fingerprints = set()
+    for block_size in (1536, 11776):
+        spec = MLAAttentionSpec(
+            block_size=block_size,
+            num_kv_heads=1,
+            head_size=128,
+            dtype=torch.bfloat16,
+        )
+        chunk_size = AttentionStoreLayout.resolve_store_chunk_size(
+            (spec,),
+            requested_store_chunk_size=512,
+        )
+        assert chunk_size == 512
+        fingerprints.add(
+            AttentionStoreLayout.schema_fingerprint((spec,), 1, 1, chunk_size)
+        )
+
+    assert len(fingerprints) == 1
+
+
+@pytest.mark.parametrize("layout_cls", [HeadMajorStoreLayout, TokenMajorStoreLayout])
+@pytest.mark.parametrize("chunk_size", [3, 4])
+def test_chunk_descriptors_preserve_order_across_kernel_blocks(layout_cls, chunk_size):
+    cache = torch.empty((3, 3, 4, 4, 4), dtype=torch.float16)
+    if layout_cls is TokenMajorStoreLayout:
+        cache = cache.transpose(2, 3)
+    cache.copy_(torch.arange(cache.numel(), dtype=cache.dtype).reshape(cache.shape))
+    layout = layout_cls(
+        KeyMetadata("test-model", 1, 0, 0, 0),
+        block_size=12,
+        hash_block_size=1,
+        local_tp_size=2,
+        store_tp_size=4,
+        tp_rank=1,
+        layer_specs=_attention_specs(1, 12, 4),
+        store_chunk_size=chunk_size,
+    )
+    layout.register_kv_caches([cache], 3)
+    starts = [12 + chunk_size, 0, 24 + 2 * chunk_size, 12 + chunk_size] * 8
+    shards = [3, 2, 3, 3] * 8
+    chunks = [(start, start + chunk_size) for start in starts]
+    addrs, sizes, blocks = layout.prepare_values(chunks, [2, 0, 1], shards)
+
+    assert blocks == [0, 2, 1, 0] * 8
+    assert layout.prepare_values(chunks[:4], [2, 0, 1], shards[:4]) == (
+        addrs[:4],
+        sizes[:4],
+        blocks[:4],
+    )
+    for start, shard, block, addresses, lengths in zip(
+        starts, shards, blocks, addrs, sizes, strict=True
+    ):
+        logical = cache[block].permute(1, 0, 2, 3).reshape(4, 12, 4)
+        head_start = (shard - 2) * 2
+        offset = start % 12
+        expected = logical[head_start : head_start + 2, offset : offset + chunk_size]
+        if layout_cls is TokenMajorStoreLayout:
+            expected = expected.transpose(0, 1)
+        expected_bytes = expected.contiguous().numpy().tobytes()
+        assert _read_segments(addresses, lengths) == expected_bytes
+
+
+def test_tp_shared_layout_handles_kernel_blocked_compressed_states():
+    from vllm.v1.kv_cache_interface import MLAAttentionSpec, group_kernel_blocks
+
+    spec = MLAAttentionSpec(
+        block_size=16,
+        num_kv_heads=1,
+        head_size=8,
+        dtype=torch.uint8,
+        tokens_per_state=4,
+        model_version="deepseek_v4",
+    )
+    cache = torch.arange(2 * 4 * 8, dtype=torch.uint8).view(8, 1, 1, 8)
+    cache = group_kernel_blocks(cache, num_blocks=2)
+    metadata = KeyMetadata("test-model", 0, 0, 0, 0)
+    layout = HeadMajorStoreLayout(
+        metadata,
+        block_size=16,
+        hash_block_size=4,
+        local_tp_size=1,
+        store_tp_size=1,
+        tp_rank=0,
+        layer_specs=(spec,),
+    )
+
+    layout.register_kv_caches([cache], num_blocks=2)
+    addrs, sizes, _ = layout.prepare_values([(0, 16)], [0], [0])
+
+    assert sizes == [[8, 8, 8, 8]]
+    assert _read_segments(addrs[0], sizes[0]) == bytes(range(32))
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 ec92fb910242..1187e6548e52 100644
--- a/tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py
+++ b/tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py
@@ -198,6 +198,25 @@ def test_scheduler_only_tracks_token_ids_for_kv_events():
             assert req_meta.token_ids is None
 
 
+@pytest.mark.parametrize("token_len", [4, 12, 15, 16, 20, 32])
+def test_scheduler_saves_only_at_common_hybrid_boundaries(token_len):
+    scheduler = _make_bare_scheduler(hash_block_size=4, enable_partial_hash_hits=True)
+    out = _make_new_scheduler_output()
+    out.num_scheduled_tokens["req-0"] = token_len
+    scheduler._unfinished_requests["req-0"] = (out.request, out.request.block_ids)
+
+    meta = scheduler.build_connector_meta(out)
+
+    checkpoint_len = token_len // scheduler._block_size * scheduler._block_size
+    assert scheduler._request_trackers["req-0"].num_saved_tokens == checkpoint_len
+    if checkpoint_len:
+        assert len(meta.requests) == 1
+        assert meta.requests[0].token_len_chunk == checkpoint_len
+    else:
+        assert meta.requests == []
+        assert scheduler._pinned_saves == {}
+
+
 def _make_preemption_scheduler_output():
     return SimpleNamespace(
         finished_req_ids=set(),
@@ -1645,6 +1664,23 @@ def test_boundary_state_never_claimed_without_a_send_thread():
     assert scheduler._gpu_block_pool.blocks[7].ref_cnt == 0
 
 
+def test_decode_consumer_claims_only_decode_boundary_state():
+    scheduler = _make_bare_scheduler(
+        hash_block_size=4,
+        enable_partial_hash_hits=True,
+        kv_role="kv_consumer",
+        save_decode_cache=True,
+    )
+    _register_offload_request(scheduler, prefill_end_tokens=12, num_prompt_tokens=12)
+    out = _make_offload_only_output([(1, 7, 8), (1, 9, 16)])
+
+    meta = scheduler.build_connector_meta(out)
+
+    assert meta.requests[0].boundary_puts == [(1, 9, 16)]
+    store_job_id = meta.requests[0].store_job_id
+    assert scheduler._pinned_saves[store_job_id][0] == [9]
+
+
 def test_resumed_partial_tail_uses_exact_boundary():
     scheduler = _make_bare_scheduler(hash_block_size=4, enable_partial_hash_hits=True)
     # Resumption replays prompt + previously generated tokens.
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 09d27eb947a0..903fcdfb0e6d 100644
--- a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py
+++ b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py
@@ -10,6 +10,7 @@
 import sys
 import threading
 import types
+from dataclasses import replace
 from types import SimpleNamespace
 from unittest.mock import MagicMock, patch
 
@@ -28,26 +29,31 @@
 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store import (
     worker as mooncake_store_worker,
 )
-from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (
+from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (  # noqa: E501
     BlobBlockHashes,
     BoundaryPut,
     ChunkedTokenDatabase,
-    KeyMetadata,
-    LBHNCStoreLayout,
-    LBNHCStoreLayout,
     LoadSpec,
     MooncakeLookupResult,
-    PoolKey,
-    RankLocalStoreLayout,
     ReqMeta,
+    RequestTracker,
     TailKeyBoundary,
 )
+from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.layout import (  # noqa: E501
+    HeadMajorStoreLayout,
+    KeyMetadata,
+    MambaStoreLayout,
+    PoolKey,
+    RankLocalStoreLayout,
+    TokenMajorStoreLayout,
+)
 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.metrics import (
     MooncakeStoreConnectorStats,
 )
 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.scheduler import (  # noqa: E501
     _partial_tail_non_mamba_puts,
 )
+from vllm.model_executor.layers.mamba.mamba_utils import get_conv_state_layout
 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 (
@@ -56,12 +62,28 @@
     KVCacheConfig,
     KVCacheGroupSpec,
     MambaSpec,
-    RSWASpec,
 )
 from vllm.v1.kv_cache_layout import KVCacheLayout
 
-_TP_SHARED_NAMESPACE = "@store_tp:4@store_pp:1@store_format:tp_shared_lbhnc"
-_TP_SHARED_LBNHC_NAMESPACE = "@store_tp:4@store_pp:1@store_format:tp_shared_lbnhc"
+
+@pytest.fixture(autouse=True)
+def _reset_conv_state_layout_cache():
+    get_conv_state_layout.cache_clear()
+    yield
+    get_conv_state_layout.cache_clear()
+
+
+_TP_SHARED_NAMESPACE = (
+    "@store_tp:4@store_pp:1@store_format:attention_head_major@store_schema:test"
+)
+_STORE_LAYOUT_CASES = [
+    (KVCacheLayout.LBHNC, HeadMajorStoreLayout, "attention_head_major"),
+    (KVCacheLayout.LBNHC, TokenMajorStoreLayout, "attention_token_major"),
+    (KVCacheLayout.BLHNC, HeadMajorStoreLayout, "attention_head_major"),
+    (KVCacheLayout.BLNHC, TokenMajorStoreLayout, "attention_token_major"),
+    (KVCacheLayout.LHBNC, HeadMajorStoreLayout, "attention_head_major"),
+    (KVCacheLayout.BHLNC, HeadMajorStoreLayout, "attention_head_major"),
+]
 
 
 def _tp_shared_prefix(tp_rank: int) -> str:
@@ -103,6 +125,8 @@ def _default_send_coord() -> mooncake_store_worker.MooncakeStoreCoordinator:
 
 def _make_tp_shared_db(
     tp_rank: int = 0,
+    chunk_size: int = 16,
+    hash_block_size: int = 16,
 ) -> tuple[ChunkedTokenDatabase, torch.Tensor]:
     num_blocks = 2
     local_heads = 4
@@ -127,16 +151,26 @@ def _make_tp_shared_db(
         0,
         store_namespace=_TP_SHARED_NAMESPACE,
     )
-    layout = LBHNCStoreLayout(
+    layout = HeadMajorStoreLayout(
         metadata,
         block_size,
-        block_size,
+        hash_block_size,
         local_tp_size=2,
         store_tp_size=4,
         tp_rank=tp_rank,
-        num_kv_heads=8,
+        store_chunk_size=chunk_size,
+        layer_specs=(
+            FullAttentionSpec(
+                block_size=block_size,
+                num_kv_heads=local_heads,
+                head_size=content_size // 2,
+                dtype=torch.float16,
+            ),
+        ),
+    )
+    db = ChunkedTokenDatabase(
+        metadata, chunk_size, hash_block_size=hash_block_size, store_layout=layout
     )
-    db = ChunkedTokenDatabase(metadata, block_size, store_layout=layout)
     layout.register_kv_caches([tensor], num_blocks)
     return db, tensor
 
@@ -152,11 +186,13 @@ def _make_store_sending_thread(
     replicate_config: object | None = None,
     enable_group_semantics: bool = False,
     supports_group_ids: bool = False,
+    kv_role: str = "kv_producer",
+    enable_kv_event: bool = False,
 ) -> mooncake_store_worker.KVCacheStoreSendingThread:
     if coord is None:
         coord = _default_send_coord()
     if token_databases is None:
-        db = ChunkedTokenDatabase(KeyMetadata("test-model", 0, 0, 0, 0), block_size=16)
+        db = ChunkedTokenDatabase(KeyMetadata("test-model", 0, 0, 0, 0), chunk_size=16)
         db.set_kv_caches_base_addr([0x1000])
         db.set_block_len([256])
         token_databases = [db]
@@ -167,8 +203,9 @@ def _make_store_sending_thread(
         coord=coord,
         tp_rank=tp_rank,
         group_put_steps=[put_step] * len(token_databases),
-        kv_role="kv_producer",
+        kv_role=kv_role,
         ready_event=threading.Event(),
+        enable_kv_event=enable_kv_event,
         replicate_config=replicate_config,
         enable_group_semantics=enable_group_semantics,
         supports_group_ids=supports_group_ids,
@@ -187,7 +224,7 @@ def _make_store_recving_thread(
     from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheGroupSpec
 
     token_database = ChunkedTokenDatabase(
-        KeyMetadata("test-model", 0, 0, 0, 0), block_size=16
+        KeyMetadata("test-model", 0, 0, 0, 0), chunk_size=16
     )
     token_database.set_kv_caches_base_addr([0x1000])
     token_database.set_block_len([256])
@@ -266,6 +303,32 @@ def _make_multi_group_store_req(req_id: str, block_hashes: list[bytes]) -> ReqMe
     )
 
 
+def test_request_tracker_keeps_event_tokens_from_hybrid_checkpoint():
+    tracker = RequestTracker(
+        req_id="r0",
+        token_len=12,
+        allocated_block_ids=([1, 2],),
+        num_saved_tokens=8,
+        token_ids=list(range(16)),
+    )
+
+    first = ReqMeta.from_request_tracker(
+        tracker,
+        block_size=8,
+    )
+    assert first is None
+    assert tracker.num_saved_tokens == 8
+
+    tracker.token_len = 16
+    second = ReqMeta.from_request_tracker(
+        tracker,
+        block_size=8,
+    )
+    assert second is not None
+    assert second.token_ids_start == 8
+    assert second.token_ids == list(range(8, 16))
+
+
 _DISK_OFFLOAD_SINGLE_KEY_BYTES = worker._estimate_disk_offload_staging_bytes([256])
 _DISK_OFFLOAD_USABLE_BUDGET_RATIO = 0.9
 _DISK_OFFLOAD_BUDGET_FOR_THREE_KEYS = 4 * _DISK_OFFLOAD_SINGLE_KEY_BYTES
@@ -311,17 +374,23 @@ def _make_vllm_config(
     decode_context_parallel_size: int = 1,
     kv_role: str = "kv_both",
     pipeline_parallel_size: int = 1,
+    tensor_parallel_size: int = 1,
     kv_cache_layout: KVCacheLayout = KVCacheLayout.LBHNC,
     disable_hybrid_kv_cache_manager: bool = True,
+    prefix_match_unit: int | None = None,
 ) -> SimpleNamespace:
     cache_config = SimpleNamespace(
-        block_size=16, num_gpu_blocks=10, prefix_match_unit=None
+        block_size=16,
+        num_gpu_blocks=10,
+        enable_prefix_caching=True,
+        prefix_match_unit=prefix_match_unit,
     )
     cache_config.get_resolved_kv_cache_layout = lambda: kv_cache_layout
     return SimpleNamespace(
         model_config=_FakeModelConfig(),
         parallel_config=SimpleNamespace(
             pipeline_parallel_size=pipeline_parallel_size,
+            tensor_parallel_size=tensor_parallel_size,
             rank=rank,
             data_parallel_index=0,
             decode_context_parallel_size=decode_context_parallel_size,
@@ -340,17 +409,19 @@ def _make_vllm_config(
 
 
 def _make_kv_cache_config(
-    *, block_size: int = 16, prefix_cache_retention_interval: int | None = 0
+    *,
+    block_size: int = 16,
+    num_kv_heads: int = 8,
+    prefix_cache_retention_interval: int | None = 0,
 ) -> object:
     """Minimal single-group KVCacheConfig for topology tests."""
-    from vllm.v1.kv_cache_interface import (
-        FullAttentionSpec,
-        KVCacheConfig,
-        KVCacheGroupSpec,
-    )
+    from vllm.v1.kv_cache_interface import KVCacheConfig
 
     spec = FullAttentionSpec(
-        block_size=block_size, num_kv_heads=8, head_size=64, dtype=None
+        block_size=block_size,
+        num_kv_heads=num_kv_heads,
+        head_size=64,
+        dtype=None,
     )
     return KVCacheConfig(
         num_blocks=10,
@@ -360,6 +431,39 @@ def _make_kv_cache_config(
     )
 
 
+def _make_hybrid_gdn_kv_cache_config(
+    tp_size: int, conv_layout: str = "DS"
+) -> KVCacheConfig:
+    from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
+    from vllm.v1.kv_cache_interface import MambaSpec
+
+    full = FullAttentionSpec(
+        block_size=1600 // tp_size,
+        num_kv_heads=1,
+        head_size=64,
+        dtype=None,
+    )
+    local_factor = 4 // tp_size
+    conv_shape = (6 * local_factor, 3)
+    if conv_layout == "SD":
+        conv_shape = conv_shape[::-1]
+    gdn = MambaSpec(
+        block_size=16,
+        shapes=(conv_shape, (local_factor, 2, 2)),
+        dtypes=(torch.uint8, torch.uint8),
+        mamba_type=MambaAttentionBackendEnum.GDN_ATTN,
+        mamba_cache_mode="align",
+    )
+    return KVCacheConfig(
+        num_blocks=10,
+        kv_cache_tensors=[],
+        kv_cache_groups=[
+            KVCacheGroupSpec(["full"], full),
+            KVCacheGroupSpec(["gdn"], gdn),
+        ],
+    )
+
+
 def _make_qsa_hybrid_kv_cache_config():
     return KVCacheConfig(
         num_blocks=10,
@@ -763,7 +867,7 @@ def test_store_sending_thread_records_mooncake_metrics():
 def test_process_tokens_uses_mask_num_as_start_chunk():
     db = ChunkedTokenDatabase(
         KeyMetadata("test-model", 0, 0, 0, 0),
-        block_size=32,
+        chunk_size=32,
         hash_block_size=8,
     )
     block_hashes = _RecordingBlockHashes([bytes([i]) for i in range(16)])
@@ -783,7 +887,7 @@ def test_process_tokens_uses_mask_num_as_start_chunk():
 def test_process_tokens_applies_chunk_mask_before_hash_access():
     db = ChunkedTokenDatabase(
         KeyMetadata("test-model", 0, 0, 0, 0),
-        block_size=32,
+        chunk_size=32,
         hash_block_size=8,
     )
     block_hashes = _RecordingBlockHashes([bytes([i]) for i in range(16)])
@@ -804,7 +908,7 @@ def test_process_tokens_applies_chunk_mask_before_hash_access():
 def test_process_tokens_applies_stride_before_hash_access():
     db = ChunkedTokenDatabase(
         KeyMetadata("test-model", 0, 0, 0, 0),
-        block_size=32,
+        chunk_size=32,
         hash_block_size=8,
     )
     block_hashes = _RecordingBlockHashes([bytes([i]) for i in range(16)])
@@ -950,7 +1054,7 @@ def _resolve_partial_tail(thread, req: ReqMeta) -> ReqMeta:
     """Add the tail's non-Mamba puts as the scheduler does before the worker."""
     req.boundary_puts = [BoundaryPut(*put) for put in req.boundary_puts or []]
     req.boundary_puts[:0] = _partial_tail_non_mamba_puts(
-        thread.coord, req, [db.block_size for db in thread.token_databases]
+        thread.coord, req, [db.store_layout.block_size for db in thread.token_databases]
     )
     return req
 
@@ -972,7 +1076,7 @@ def _make_partial_tail_send_thread(
     )
     db = ChunkedTokenDatabase(
         KeyMetadata("test-model", 0, 0, 0, 0),
-        block_size=4,
+        chunk_size=4,
         hash_block_size=4,
     )
     db.set_kv_caches_base_addr([0x1000])
@@ -980,7 +1084,7 @@ def _make_partial_tail_send_thread(
     # Group 1 models the mamba "align" group the hand-offs reference.
     db_mamba = ChunkedTokenDatabase(
         KeyMetadata("test-model", 0, 0, 0, 0, group_id=1),
-        block_size=16,
+        chunk_size=16,
         hash_block_size=4,
     )
     db_mamba.set_kv_caches_base_addr([0x2000])
@@ -1044,8 +1148,8 @@ def put(keys, *args):
     assert thread._maybe_offload_boundary_states(
         _resolve_partial_tail(thread, metadata)
     )
-    mamba_key = thread.token_databases[1].key_for(b"a1")
-    attention_key = thread.token_databases[0].key_for(b"a2")
+    mamba_key = thread.token_databases[1].store_layout.key_for(0, b"a1")
+    attention_key = thread.token_databases[0].store_layout.key_for(0, b"a2")
     assert mamba_key in stored
     assert attention_key not in stored
 
@@ -1088,6 +1192,62 @@ def test_partial_tail_offload_skips_null_source_blocks():
     ]
 
 
+@pytest.mark.parametrize("saved", [0, 16])
+def test_partial_tail_offload_is_independent_of_worker_progress(saved):
+    store = MagicMock()
+    store.batch_is_exist.side_effect = lambda keys: [0] * len(keys)
+    store.batch_put_from_multi_buffers.side_effect = lambda keys, *a: [256] * len(keys)
+    thread = _make_partial_tail_send_thread(store)
+    thread._saved_offset["req-a"] = saved
+
+    assert thread._maybe_offload_boundary_states(
+        _resolve_partial_tail(thread, _make_partial_tail_req([1, 2, 3]))
+    )
+
+    keys = store.batch_is_exist.call_args.args[0]
+    attention_keys = [key for key in keys if "@group:0" in key]
+    assert [key.rsplit("@", 1)[-1] for key in attention_keys] == [
+        "6130",
+        "6131",
+        "6132",
+    ]
+    assert any("@group:1@6132" in key for key in keys)
+
+
+@pytest.mark.parametrize("chunk_size", [4, 16])
+def test_partial_tail_offload_writes_each_local_store_shard(chunk_size):
+    store = MagicMock()
+    store.batch_is_exist.side_effect = lambda keys: [0] * len(keys)
+    store.batch_put_from_multi_buffers.side_effect = lambda keys, *a: [256] * len(keys)
+    db, tensor = _make_tp_shared_db(chunk_size=chunk_size, hash_block_size=4)
+    thread = _make_partial_tail_send_thread(store)
+    thread.token_databases[0] = db
+
+    assert thread._maybe_offload_boundary_states(
+        _resolve_partial_tail(thread, _make_partial_tail_req([1]))
+    )
+
+    keys, addrs, sizes, _ = store.batch_put_from_multi_buffers.call_args.args
+    attention_puts = [
+        (key, addr, size)
+        for key, addr, size in zip(keys, addrs, sizes, strict=True)
+        if "@group:0" in key
+    ]
+    expected = []
+    for start in range(0, 12, chunk_size):
+        end = min(start + chunk_size, 12)
+        for shard in range(2):
+            heads = tensor[1, shard * 2 : (shard + 1) * 2]
+            expected.append(
+                (
+                    db.store_layout.key_for(shard, [b"a0", b"a1", b"a2"][end // 4 - 1]),
+                    [head[start].data_ptr() for head in heads],
+                    [chunk_size * 4 * tensor.element_size()] * 2,
+                )
+            )
+    assert attention_puts == expected
+
+
 def test_store_sending_thread_skips_null_sparse_group_blocks():
     from vllm.v1.kv_cache_interface import (
         FullAttentionSpec,
@@ -1112,13 +1272,13 @@ def test_store_sending_thread_skips_null_sparse_group_blocks():
 
     db_full = ChunkedTokenDatabase(
         KeyMetadata("test-model", 0, 0, 0, 0, group_id=0),
-        block_size=16,
+        chunk_size=16,
     )
     db_full.set_kv_caches_base_addr([0x1000])
     db_full.set_block_len([256])
     db_sparse = ChunkedTokenDatabase(
         KeyMetadata("test-model", 0, 0, 0, 0, group_id=1),
-        block_size=16,
+        chunk_size=16,
     )
     db_sparse.set_kv_caches_base_addr([0x2000])
     db_sparse.set_block_len([512])
@@ -1180,7 +1340,7 @@ def test_partial_tail_offload_skips_cap_omitted_mamba_group():
     def make_db(group_id: int, block_size: int, base_addr: int):
         db = ChunkedTokenDatabase(
             KeyMetadata("test-model", 0, 0, 0, 0, group_id=group_id),
-            block_size=block_size,
+            chunk_size=block_size,
             hash_block_size=4,
         )
         db.set_kv_caches_base_addr([base_addr])
@@ -1214,8 +1374,14 @@ def make_db(group_id: int, block_size: int, base_addr: int):
     )
     puts = list(zip(keys, addrs, strict=True))
     boundary_hash = BlockHash(b"a2")
-    assert (db_mamba_1.key_for(boundary_hash), [0x2000 + 7 * 256]) in puts
-    assert (db_mamba_2.key_for(boundary_hash), [0x3000 + 9 * 256]) not in puts
+    assert (
+        db_mamba_1.store_layout.key_for(0, boundary_hash),
+        [0x2000 + 7 * 256],
+    ) in puts
+    assert (
+        db_mamba_2.store_layout.key_for(0, boundary_hash),
+        [0x3000 + 9 * 256],
+    ) not in puts
 
 
 def test_partial_tail_offload_replaces_stale_group_ids_after_filtering():
@@ -1309,12 +1475,12 @@ def test_normal_save_excludes_mamba_group_and_null_blocks():
         hash_block_size=16,
     )
     db_full = ChunkedTokenDatabase(
-        KeyMetadata("test-model", 0, 0, 0, 0, group_id=0), block_size=16
+        KeyMetadata("test-model", 0, 0, 0, 0, group_id=0), chunk_size=16
     )
     db_full.set_kv_caches_base_addr([0x1000])
     db_full.set_block_len([256])
     db_mamba = ChunkedTokenDatabase(
-        KeyMetadata("test-model", 0, 0, 0, 0, group_id=1), block_size=16
+        KeyMetadata("test-model", 0, 0, 0, 0, group_id=1), chunk_size=16
     )
     db_mamba.set_kv_caches_base_addr([0x2000])
     db_mamba.set_block_len([256])
@@ -1365,7 +1531,7 @@ def test_block_aligned_snapshot_offload_uses_provided_block():
     # boundary 32 is block-aligned for the mamba group (block 16): one key,
     # keyed by hs[32 // 4 - 1], read from the handed-off block 7 — not from
     # block_ids[1] and with no FA gap coverage.
-    assert keys == [thread.token_databases[1].key_for(BlockHash(hs[7]))]
+    assert keys == [thread.token_databases[1].store_layout.key_for(0, BlockHash(hs[7]))]
     assert addrs == [[0x2000 + 7 * 256]]
 
 
@@ -1401,11 +1567,14 @@ def test_mixed_snapshot_and_sub_block_offloads(saved_tokens, use_eagle):
     proof_end = 48 if use_eagle else 44
     assert keys == [
         # The tail reads only the LCM gap, even when normal saves lag.
-        *[db_full.key_for(BlockHash(hs[i])) for i in range(8, proof_end // 4)],
+        *[
+            db_full.store_layout.key_for(0, BlockHash(hs[i]))
+            for i in range(8, proof_end // 4)
+        ],
         # Aligned snapshot: boundary 32 from the handed-off block 9.
-        db_mamba.key_for(BlockHash(hs[7])),
+        db_mamba.store_layout.key_for(0, BlockHash(hs[7])),
         # Mamba boundary block from the CoW hand-off (block 7).
-        db_mamba.key_for(BlockHash(hs[10])),
+        db_mamba.store_layout.key_for(0, BlockHash(hs[10])),
     ]
     assert addrs == [
         *[[0x1000 + (i + 1) * 256] for i in range(8, proof_end // 4)],
@@ -1432,13 +1601,13 @@ def test_partial_tail_with_smaller_mamba_blocks_writes_one_mamba_key(tp_rank):
         eagle_proof_margin_by_group={},
     )
     db_full = ChunkedTokenDatabase(
-        KeyMetadata("test-model", 0, 0, 0, 0), block_size=16, hash_block_size=4
+        KeyMetadata("test-model", 0, 0, 0, 0), chunk_size=16, hash_block_size=4
     )
     db_full.set_kv_caches_base_addr([0x1000])
     db_full.set_block_len([256])
     db_mamba = ChunkedTokenDatabase(
         KeyMetadata("test-model", 0, 0, 0, 0, group_id=1),
-        block_size=8,
+        chunk_size=8,
         hash_block_size=4,
     )
     db_mamba.set_kv_caches_base_addr([0x2000])
@@ -1467,8 +1636,10 @@ def test_partial_tail_with_smaller_mamba_blocks_writes_one_mamba_key(tp_rank):
     keys, addrs, _sizes, _ = store.batch_put_from_multi_buffers.call_args.args
     # Mamba block cdiv(44, 8) - 1 = 5 belongs to put_step_rank 5 % 2 = 1,
     # which is tp_rank 0 for group 1.
-    mamba_puts = [db_mamba.key_for(BlockHash(hs[10]))] if tp_rank == 0 else []
-    assert keys == [db_full.key_for(BlockHash(hs[10])), *mamba_puts]
+    mamba_puts = (
+        [db_mamba.store_layout.key_for(0, BlockHash(hs[10]))] if tp_rank == 0 else []
+    )
+    assert keys == [db_full.store_layout.key_for(0, BlockHash(hs[10])), *mamba_puts]
     assert addrs == [[0x1000 + 3 * 256], *([[0x2000 + 7 * 256]] * len(mamba_puts))]
 
 
@@ -1488,13 +1659,13 @@ def test_shared_prefix_junction_handoff_writes_only_its_mamba_key():
         eagle_proof_margin_by_group={0: 4},
     )
     db_full = ChunkedTokenDatabase(
-        KeyMetadata("test-model", 0, 0, 0, 0), block_size=4, hash_block_size=4
+        KeyMetadata("test-model", 0, 0, 0, 0), chunk_size=4, hash_block_size=4
     )
     db_full.set_kv_caches_base_addr([0x1000])
     db_full.set_block_len([256])
     db_mamba = ChunkedTokenDatabase(
         KeyMetadata("test-model", 0, 0, 0, 0, group_id=1),
-        block_size=16,
+        chunk_size=16,
         hash_block_size=4,
     )
     db_mamba.set_kv_caches_base_addr([0x2000])
@@ -1519,7 +1690,7 @@ def test_shared_prefix_junction_handoff_writes_only_its_mamba_key():
     assert req.boundary_puts == [(1, 7, 40)]
 
     keys, addrs, _sizes, _ = store.batch_put_from_multi_buffers.call_args.args
-    assert keys == [db_mamba.key_for(BlockHash(hs[9]))]
+    assert keys == [db_mamba.store_layout.key_for(0, BlockHash(hs[9]))]
     assert addrs == [[0x2000 + 7 * 256]]
 
 
@@ -1581,38 +1752,46 @@ def test_store_sending_thread_delta_start_rank_saves_second_local_chunk():
     assert store.batch_put_from_multi_buffers.call_args.args[0] == keys
 
 
-def test_store_sending_thread_delta_saves_only_new_masked_chunks():
+@pytest.mark.parametrize("chunk_size", [8, 16])
+def test_store_sending_thread_delta_saves_only_new_masked_chunks(chunk_size):
     store = MagicMock()
     store.batch_is_exist.side_effect = lambda keys: [0] * len(keys)
     store.batch_put_from_multi_buffers.side_effect = (
         lambda keys, addrs, sizes, replicate_config: [256] * len(keys)
     )
-    coord = SimpleNamespace(
-        lcm_block_size=16,
-        store_mask=lambda token_len, start_token, num_prompt_tokens=None: (
+    coord = MagicMock(lcm_block_size=16)
+    coord.store_mask.side_effect = (
+        lambda token_len, start_token, num_prompt_tokens=None: (
             None,
-            [True, False],
-        ),
-    )
-    coord.store_mask = MagicMock(side_effect=coord.store_mask)
-
-    db_full = ChunkedTokenDatabase(
-        KeyMetadata("test-model", 0, 0, 0, 0, group_id=0),
-        block_size=16,
+            [True, False, True, False][start_token // 16 : token_len // 16],
+        )
     )
-    db_full.set_kv_caches_base_addr([0x1000])
-    db_full.set_block_len([256])
-    db_masked = ChunkedTokenDatabase(
-        KeyMetadata("test-model", 0, 0, 0, 0, group_id=1),
-        block_size=16,
+    spec = FullAttentionSpec(
+        block_size=16, num_kv_heads=1, head_size=2, dtype=torch.float16
     )
-    db_masked.set_kv_caches_base_addr([0x2000])
-    db_masked.set_block_len([256])
+    caches = torch.empty((2, 4, 1, 16, 4), dtype=torch.float16)
+    databases = []
+    for group_id in range(2):
+        metadata = KeyMetadata("test-model", 0, 0, 0, 0, group_id=group_id)
+        layout = HeadMajorStoreLayout(
+            metadata,
+            block_size=16,
+            hash_block_size=chunk_size,
+            local_tp_size=1,
+            store_tp_size=1,
+            tp_rank=0,
+            layer_specs=(spec,),
+            store_chunk_size=chunk_size,
+        )
+        layout.register_kv_caches([caches[group_id]], num_blocks=4)
+        databases.append(
+            ChunkedTokenDatabase(metadata, chunk_size, store_layout=layout)
+        )
 
     thread = _make_store_sending_thread(
         store,
         coord=coord,
-        token_databases=[db_full, db_masked],
+        token_databases=databases,
     )
 
     thread._saved_offset["req-a"] = 32
@@ -1622,7 +1801,7 @@ def test_store_sending_thread_delta_saves_only_new_masked_chunks():
             req_id="req-a",
             token_len_chunk=64,
             block_ids=([0, 1, 2, 3], [0, 1, 2, 3]),
-            block_hashes=[b"a0", b"a1", b"a2", b"a3"],
+            block_hashes=[f"a{i}".encode() for i in range(64 // chunk_size)],
             can_save=True,
             num_prompt_tokens=33,
             prefill_end_tokens=64,
@@ -1633,11 +1812,39 @@ def test_store_sending_thread_delta_saves_only_new_masked_chunks():
     full_hashes = [k.rsplit("@", 1)[-1] for k in keys if "@group:0" in k]
     masked_hashes = [k.rsplit("@", 1)[-1] for k in keys if "@group:1" in k]
 
-    assert full_hashes == [b"a2".hex(), b"a3".hex()]
-    assert masked_hashes == [b"a2".hex()]
+    assert full_hashes == [
+        f"a{i}".encode().hex() for i in range(32 // chunk_size, 64 // chunk_size)
+    ]
+    assert masked_hashes == [
+        f"a{i}".encode().hex() for i in range(32 // chunk_size, 48 // chunk_size)
+    ]
     coord.store_mask.assert_called_once_with(64, 32, num_prompt_tokens=64)
 
 
+def test_decode_save_marks_completed_interval_as_replay_boundary():
+    store = MagicMock()
+    store.batch_is_exist.return_value = []
+    coord = MagicMock(lcm_block_size=16)
+    coord.store_mask.return_value = (None,)
+    thread = _make_store_sending_thread(store, coord=coord, kv_role="kv_consumer")
+    thread._saved_offset["req-a"] = 32
+
+    _run_store_req(
+        thread,
+        ReqMeta(
+            req_id="req-a",
+            token_len_chunk=64,
+            block_ids=([0, 1, 2, 3],),
+            block_hashes=[b"a0", b"a1", b"a2", b"a3"],
+            can_save=True,
+            num_prompt_tokens=48,
+            prefill_end_tokens=48,
+        ),
+    )
+
+    coord.store_mask.assert_called_once_with(64, 32, num_prompt_tokens=65)
+
+
 def test_store_sending_thread_prepares_missing_chunks_once_per_group():
     store = MagicMock()
     store.batch_is_exist.return_value = [0, 1, 0, 1, 0, 0]
@@ -1652,7 +1859,7 @@ def test_store_sending_thread_prepares_missing_chunks_once_per_group():
 
     db0 = ChunkedTokenDatabase(
         KeyMetadata("test-model", 0, 0, 0, 0, group_id=0),
-        block_size=16,
+        chunk_size=16,
     )
     db0.set_kv_caches_base_addr([0x1000])
     db0.set_block_len([256])
@@ -1660,7 +1867,7 @@ def test_store_sending_thread_prepares_missing_chunks_once_per_group():
 
     db1 = ChunkedTokenDatabase(
         KeyMetadata("test-model", 0, 0, 0, 0, group_id=1),
-        block_size=16,
+        chunk_size=16,
     )
     db1.set_kv_caches_base_addr([0x2000])
     db1.set_block_len([512])
@@ -1759,7 +1966,7 @@ def test_stale_store_job_cannot_touch_a_reused_request_id():
     stale = _make_store_req("req-a", [b"a0", b"a1"])
     stale.store_job_id = 1
     thread.add_request(stale)
-    thread._record_saved(stale, 32)
+    thread._record_saved_offset(stale, 32)
     thread.delete_finished_stored_request("req-a")
 
     live = _make_store_req("req-a", [b"a0", b"a1"])
@@ -1768,7 +1975,7 @@ def test_stale_store_job_cannot_touch_a_reused_request_id():
     thread._retry_token_ids["req-a"] = (32, list(range(32, 64)))
 
     thread.finish_store_job(stale)
-    thread._record_saved(stale, 64)
+    thread._record_saved_offset(stale, 64)
     thread._mark_request_skipped_for_pressure(stale)
     thread._update_retry_token_ids(stale, False, 0, list(range(32)))
     thread._update_retry_token_ids(stale, True, 0, None)
@@ -1832,6 +2039,24 @@ def test_store_recving_thread_reports_all_attempted_blocks_on_exception():
     assert thread.get_and_clear_block_ids_with_load_errors() == {0, 1, 2}
 
 
+def test_store_recving_thread_finishes_when_group_masks_select_no_keys():
+    store = MagicMock()
+    thread = _make_store_recving_thread(store)
+    thread.coord.load_mask = MagicMock(return_value=[[False]])
+
+    thread._handle_request(
+        _make_load_req(
+            "req-a",
+            [b"a0"],
+            token_len=16,
+        )
+    )
+
+    assert thread.get_and_clear_finished_requests() == {"req-a"}
+    assert thread.get_and_clear_block_ids_with_load_errors() == set()
+    store.batch_get_into_multi_buffers.assert_not_called()
+
+
 def test_store_worker_get_block_ids_with_load_errors_delegates_to_recv_thread():
     recv_thread = MagicMock()
     recv_thread.get_and_clear_block_ids_with_load_errors.return_value = {3, 4}
@@ -2020,7 +2245,7 @@ def test_store_sending_thread_group_id_excludes_physical_sharding():
             pp_rank=0,
             group_id=0,
         ),
-        block_size=16,
+        chunk_size=16,
     )
     db.set_kv_caches_base_addr([0x1000])
     db.set_block_len([256])
@@ -2051,7 +2276,7 @@ def test_store_sending_thread_multiple_segments_share_logical_group_id():
     store.batch_is_exist.return_value = [0, 0]
     store.batch_put_from_multi_buffers.return_value = [512, 512]
     replicate_config = SimpleNamespace(group_ids=None)
-    db = ChunkedTokenDatabase(KeyMetadata("test-model", 0, 0, 0, 0), block_size=16)
+    db = ChunkedTokenDatabase(KeyMetadata("test-model", 0, 0, 0, 0), chunk_size=16)
     db.set_kv_caches_base_addr([0x1000, 0x2000])
     db.set_block_len([256, 256])
     thread = _make_store_sending_thread(
@@ -2097,7 +2322,7 @@ def test_store_sending_thread_group_ids_share_across_kv_cache_groups():
     for group_id, base_addr in enumerate([0x1000, 0x3000]):
         db = ChunkedTokenDatabase(
             KeyMetadata("test-model", 0, 0, 0, 0, group_id=group_id),
-            block_size=16,
+            chunk_size=16,
         )
         db.set_kv_caches_base_addr([base_addr])
         db.set_block_len([256])
@@ -2340,10 +2565,10 @@ def test_recv_thread_splits_disk_offload_loads_by_budget():
     layout = thread.token_databases[0].store_layout
     assert isinstance(layout, RankLocalStoreLayout)
     base_addr = layout.kv_caches_base_addr[0]
-    block_len = layout.block_len[0]
-    assert first_addrs == [[base_addr], [base_addr + block_len]]
-    assert second_addrs == [[base_addr + 2 * block_len]]
-    expected_size = block_len
+    block_stride = layout.block_stride[0]
+    assert first_addrs == [[base_addr], [base_addr + block_stride]]
+    assert second_addrs == [[base_addr + 2 * block_stride]]
+    expected_size = layout.block_len[0]
     assert first_sizes == [[expected_size], [expected_size]]
     assert second_sizes == [[expected_size]]
 
@@ -2481,7 +2706,7 @@ def test_worker_init_excludes_nonprefix_cache_groups(monkeypatch):
     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]
+    assert [db.chunk_size for db in store_worker.token_dbs] == [800, 800]
 
 
 def _make_two_full_attention_groups_kv_cache_config():
@@ -2628,7 +2853,12 @@ def test_requester_worker_init_skips_disk_budget_when_offload_disabled(
     assert w.disk_offload_buffer_budget_bytes is None
 
 
-def test_worker_enables_lbhnc_store_tp_layout(tmp_path, monkeypatch):
+@pytest.mark.parametrize(
+    ("cache_layout", "layout_cls", "store_format"), _STORE_LAYOUT_CASES
+)
+def test_worker_enables_store_tp_layout(
+    tmp_path, monkeypatch, cache_layout, layout_cls, store_format
+):
     store = MagicMock()
     store.setup.return_value = 0
     _install_fake_mooncake(monkeypatch, store)
@@ -2646,24 +2876,38 @@ def test_worker_enables_lbhnc_store_tp_layout(tmp_path, monkeypatch):
         ),
     )
 
-    w = worker.MooncakeStoreWorker(
-        _make_vllm_config(extra_config={"store_tp_size": 4}),
-        _make_kv_cache_config(),
-    )
+    with patch.object(worker.logger, "warning_once") as warning_once:
+        w = worker.MooncakeStoreWorker(
+            _make_vllm_config(
+                extra_config={"store_tp_size": 4}, kv_cache_layout=cache_layout
+            ),
+            _make_kv_cache_config(num_kv_heads=4),
+        )
+
+    if cache_layout in (KVCacheLayout.LBNHC, KVCacheLayout.BLNHC):
+        warning_once.assert_called_once()
+    else:
+        warning_once.assert_not_called()
 
     assert w.store_tp_size == 4
-    assert w.token_dbs[0].metadata.store_namespace == _TP_SHARED_NAMESPACE
-    assert isinstance(w.token_dbs[0].store_layout, LBHNCStoreLayout)
+    assert w.token_dbs[0].metadata.store_namespace.startswith(
+        f"@store_tp:4@store_pp:1@store_format:{store_format}@store_schema:"
+    )
+    assert isinstance(w.token_dbs[0].store_layout, layout_cls)
     assert w.token_dbs[0].store_layout.local_shard_ids == (0, 1)
     assert len(w._lookup_key_prefixes[0]) == 4
 
 
-def test_worker_enables_lbnhc_store_tp_layout(tmp_path, monkeypatch):
+@pytest.mark.parametrize("with_scratch", [False, True])
+@pytest.mark.parametrize("conv_layout", ["DS", "SD"])
+def test_hybrid_gdn_workers_share_group_aware_store_tp_namespace(
+    tmp_path, monkeypatch, conv_layout, with_scratch
+):
     store = MagicMock()
     store.setup.return_value = 0
     _install_fake_mooncake(monkeypatch, store)
-    _patch_worker_runtime(monkeypatch, tp_size=2)
-    monkeypatch.setattr(_FakeModelConfig, "get_total_num_kv_heads", lambda _self: 8)
+    monkeypatch.setattr(_FakeModelConfig, "get_total_num_kv_heads", lambda _self: 2)
+    monkeypatch.setenv("VLLM_SSM_CONV_STATE_LAYOUT", conv_layout)
     monkeypatch.setenv(
         "MOONCAKE_CONFIG_PATH",
         _write_mooncake_config(
@@ -2676,37 +2920,93 @@ def test_worker_enables_lbnhc_store_tp_layout(tmp_path, monkeypatch):
         ),
     )
 
-    with patch.object(worker.logger, "warning_once") as warning_once:
-        w = worker.MooncakeStoreWorker(
-            _make_vllm_config(
-                extra_config={"store_tp_size": 4},
-                kv_cache_layout=KVCacheLayout.LBNHC,
-            ),
-            _make_kv_cache_config(),
+    workers = []
+    for tp_size, tp_rank in ((2, 0), (4, 1)):
+        _patch_worker_runtime(monkeypatch, tp_size=tp_size, tp_rank=tp_rank)
+        config = _make_vllm_config(
+            extra_config={"store_tp_size": 4},
+            tensor_parallel_size=tp_size,
+            prefix_match_unit=16,
         )
+        kv_cache_config = _make_hybrid_gdn_kv_cache_config(tp_size, conv_layout)
+        if with_scratch:
+            from vllm.v1.kv_cache_interface import CircularBufferSpec
+
+            kv_cache_config.kv_cache_groups.append(
+                KVCacheGroupSpec(
+                    ["scratch"],
+                    CircularBufferSpec(
+                        block_size=4,
+                        num_kv_heads=1,
+                        head_size=64,
+                        head_size_v=0,
+                        dtype=torch.bfloat16,
+                    ),
+                )
+            )
+        store_worker = worker.MooncakeStoreWorker(config, kv_cache_config)
+        assert len(store_worker.token_dbs) == 2
+        assert [group.layer_names for group in store_worker._kv_cache_groups] == [
+            ["full"],
+            ["gdn"],
+        ]
+        assert store_worker.store_tp_size == 4
+        assert isinstance(store_worker.token_dbs[0].store_layout, HeadMajorStoreLayout)
+        assert isinstance(store_worker.token_dbs[1].store_layout, MambaStoreLayout)
+        assert store_worker.token_dbs[0].chunk_size == 16
+        assert len(store_worker._lookup_key_prefixes[0]) == 2
+        assert len(store_worker._lookup_key_prefixes[1]) == 4
+        workers.append(store_worker)
 
-    warning_once.assert_called_once()
-    assert w.store_tp_size == 4
-    assert w.token_dbs[0].metadata.store_namespace == _TP_SHARED_LBNHC_NAMESPACE
-    assert isinstance(w.token_dbs[0].store_layout, LBNHCStoreLayout)
-    assert w.token_dbs[0].store_layout.local_shard_ids == (0, 1)
-    assert len(w._lookup_key_prefixes[0]) == 4
-
+    full_key_tp2 = workers[0].token_dbs[0].store_layout.key_for(0, BlockHash(b"h"))
+    full_key_tp4 = workers[1].token_dbs[0].store_layout.key_for(0, BlockHash(b"h"))
+    gdn_key_tp2 = workers[0].token_dbs[1].store_layout.key_for(1, BlockHash(b"h"))
+    gdn_key_tp4 = workers[1].token_dbs[1].store_layout.key_for(1, BlockHash(b"h"))
+    assert full_key_tp2 == full_key_tp4
+    assert gdn_key_tp2 == gdn_key_tp4
+    assert "@group:0" in full_key_tp2
+    assert "@group:1" in gdn_key_tp2
+    assert "@store_format:attention_head_major" in full_key_tp2
+    assert "@store_format:mamba_state" in gdn_key_tp2
+
+
+@pytest.mark.parametrize(("store_tp_size", "prefix_match_unit"), [(6, 16), (4, None)])
+def test_hybrid_gdn_falls_back_from_incompatible_store_config(
+    tmp_path, monkeypatch, store_tp_size, prefix_match_unit
+):
+    store = MagicMock()
+    store.setup.return_value = 0
+    _install_fake_mooncake(monkeypatch, store)
+    monkeypatch.setattr(_FakeModelConfig, "get_total_num_kv_heads", lambda _self: 2)
+    _patch_worker_runtime(monkeypatch, tp_size=2)
+    monkeypatch.setenv("VLLM_SSM_CONV_STATE_LAYOUT", "DS")
+    monkeypatch.setenv(
+        "MOONCAKE_CONFIG_PATH",
+        _write_mooncake_config(
+            tmp_path,
+            {
+                "metadata_server": "http://metadata/endpoint",
+                "protocol": "tcp",
+                "master_server_address": "10.0.0.7:50051",
+            },
+        ),
+    )
 
-def test_tp_sharded_layout_rejects_full_attention_subclasses():
-    spec = RSWASpec(
-        block_size=16,
-        num_kv_heads=8,
-        head_size=64,
-        dtype=None,
-        rswa_window=32,
+    config = _make_vllm_config(
+        extra_config={"store_tp_size": store_tp_size},
+        prefix_match_unit=prefix_match_unit,
     )
-    store_worker = object.__new__(worker.MooncakeStoreWorker)
-    store_worker.pcp_size = 1
-    store_worker.dcp_size = 1
-    store_worker._kv_cache_groups = [SimpleNamespace(kv_cache_spec=spec)]
+    kv_cache_config = _make_hybrid_gdn_kv_cache_config(tp_size=2)
+    store_worker = worker.MooncakeStoreWorker(config, kv_cache_config)
 
-    assert not store_worker._supports_tp_sharded_store_layout(LBHNCStoreLayout, {})
+    assert store_worker.store_tp_size is None
+    assert all(
+        isinstance(db.store_layout, RankLocalStoreLayout)
+        for db in store_worker.token_dbs
+    )
+    assert "@store_format:rank_local_tp2" in store_worker.token_dbs[
+        0
+    ].store_layout.key_for(0, BlockHash(b"h"))
 
 
 @pytest.mark.parametrize(
@@ -2778,10 +3078,10 @@ def test_lcm_store_tp_gives_prefill_and_decode_common_namespace(tmp_path, monkey
             extra_config["save_decode_cache"] = True
         store_worker = worker.MooncakeStoreWorker(
             _make_vllm_config(extra_config=extra_config, kv_role=kv_role),
-            _make_kv_cache_config(),
+            _make_kv_cache_config(num_kv_heads=8 // tp_size),
         )
         assert store_worker.store_tp_size == 4
-        assert isinstance(store_worker.token_dbs[0].store_layout, LBHNCStoreLayout)
+        assert isinstance(store_worker.token_dbs[0].store_layout, HeadMajorStoreLayout)
         assert store_worker.token_dbs[0].store_layout.local_shard_ids == expected_shards
         workers.append(store_worker)
 
@@ -2832,14 +3132,14 @@ def test_lcm_store_tp_falls_back_when_topology_is_incompatible(
                 "prefill_tp_sizes": prefill_tp_sizes,
             }
         ),
-        _make_kv_cache_config(),
+        _make_kv_cache_config(num_kv_heads=max(1, 8 // tp_size)),
     )
 
     assert store_worker.store_tp_size is None
     assert type(store_worker.token_dbs[0]) is ChunkedTokenDatabase
-    assert "@store_format:rank_local_tp" in store_worker.token_dbs[0].key_for(
-        BlockHash(b"h")
-    )
+    assert "@store_format:rank_local_tp" in store_worker.token_dbs[
+        0
+    ].store_layout.key_for(0, BlockHash(b"h"))
 
 
 @pytest.mark.parametrize("store_tp_size", [1, 3, 6])
@@ -2870,12 +3170,12 @@ def test_worker_keeps_rank_local_layout_for_unsupported_store_tp(
             extra_config={"store_tp_size": store_tp_size},
             kv_cache_layout=KVCacheLayout.LBNHC,
         ),
-        _make_kv_cache_config(),
+        _make_kv_cache_config(num_kv_heads=4),
     )
 
     assert w.store_tp_size is None
     assert type(w.token_dbs[0]) is ChunkedTokenDatabase
-    key = w.token_dbs[0].key_for(BlockHash(b"h"))
+    key = w.token_dbs[0].store_layout.key_for(0, BlockHash(b"h"))
     assert "@store_tp:" not in key
     assert "@store_pp:1" in key
     assert "@store_format:rank_local_tp2_layout_LBNHC" in key
@@ -2886,6 +3186,19 @@ def test_store_tp_opt_in_isolates_different_pp_sizes(tmp_path, monkeypatch):
     store.setup.return_value = 0
     _install_fake_mooncake(monkeypatch, store)
     _patch_worker_runtime(monkeypatch, tp_size=2)
+    monkeypatch.setattr(worker, "get_pp_group", MagicMock())
+
+    def gather_prefixes(output, local, **kwargs):
+        for pp_rank in range(len(output)):
+            output[pp_rank] = tuple(
+                tuple(
+                    prefix.replace("@pp_rank:0", f"@pp_rank:{pp_rank}")
+                    for prefix in group
+                )
+                for group in local
+            )
+
+    monkeypatch.setattr(torch.distributed, "all_gather_object", gather_prefixes)
     monkeypatch.setattr(_FakeModelConfig, "get_total_num_kv_heads", lambda _self: 8)
     monkeypatch.setenv(
         "MOONCAKE_CONFIG_PATH",
@@ -2906,7 +3219,7 @@ def test_store_tp_opt_in_isolates_different_pp_sizes(tmp_path, monkeypatch):
                 extra_config={"store_tp_size": 4},
                 pipeline_parallel_size=pp_size,
             ),
-            _make_kv_cache_config(),
+            _make_kv_cache_config(num_kv_heads=4),
         )
         layout = w.token_dbs[0].store_layout
         keys.append(layout.key_for(layout.local_shard_ids[0], BlockHash(b"h")))
@@ -2916,6 +3229,107 @@ def test_store_tp_opt_in_isolates_different_pp_sizes(tmp_path, monkeypatch):
     assert keys[0] != keys[1]
 
 
+@pytest.mark.parametrize("tp_size", [2, 4])
+@pytest.mark.parametrize("missing_group", [None, 0, 1])
+def test_hybrid_lookup_uses_each_pp_stage_schema(
+    tmp_path, monkeypatch, tp_size, missing_group
+):
+    store = MagicMock()
+    store.setup.return_value = 0
+    _install_fake_mooncake(monkeypatch, store)
+    _patch_worker_runtime(monkeypatch, tp_size=tp_size)
+    monkeypatch.setattr(_FakeModelConfig, "get_total_num_kv_heads", lambda _self: 2)
+    monkeypatch.setenv("VLLM_SSM_CONV_STATE_LAYOUT", "DS")
+    monkeypatch.setenv(
+        "MOONCAKE_CONFIG_PATH",
+        _write_mooncake_config(
+            tmp_path,
+            {
+                "metadata_server": "http://metadata/endpoint",
+                "protocol": "tcp",
+                "master_server_address": "10.0.0.7:50051",
+            },
+        ),
+    )
+
+    stages = []
+    # Build both stage-local schemas before simulating the startup collective.
+    with patch.object(worker.MooncakeStoreWorker, "_init_lookup_key_prefixes"):
+        for pp_rank, layer_counts in enumerate(((2, 1), (1, 2))):
+            kv_config = _make_hybrid_gdn_kv_cache_config(tp_size)
+            kv_config.kv_cache_groups = [
+                replace(
+                    group,
+                    layer_names=[f"pp{pp_rank}_g{g}_l{i}" for i in range(count)],
+                    kv_cache_spec=replace(group.kv_cache_spec, block_size=16),
+                )
+                for g, (group, count) in enumerate(
+                    zip(kv_config.kv_cache_groups, layer_counts, strict=True)
+                )
+            ]
+            stages.append(
+                worker.MooncakeStoreWorker(
+                    _make_vllm_config(
+                        extra_config={"store_tp_size": 4},
+                        rank=pp_rank * tp_size,
+                        pipeline_parallel_size=2,
+                        tensor_parallel_size=tp_size,
+                        prefix_match_unit=16,
+                    ),
+                    kv_config,
+                )
+            )
+
+    stage_prefixes = [
+        tuple(
+            db.store_layout.lookup_key_prefixes([(0, 0, 0, stage.pp_rank)])
+            for db in stage.token_dbs
+        )
+        for stage in stages
+    ]
+    for g in range(2):
+        assert (
+            stages[0].token_dbs[g].metadata.store_namespace
+            != stages[1].token_dbs[g].metadata.store_namespace
+        )
+    hashes = [BlockHash(b"h0"), BlockHash(b"h1")]
+    present = {
+        PoolKey.build_key_string(prefix, block_hash.hex())
+        for groups in stage_prefixes
+        for group in groups
+        for prefix in group
+        for block_hash in hashes
+    }
+    if missing_group is not None:
+        for block_hash in hashes:
+            present.remove(
+                PoolKey.build_key_string(
+                    stage_prefixes[1][missing_group][-1], block_hash.hex()
+                )
+            )
+    store.batch_is_exist.side_effect = lambda keys: [int(k in present) for k in keys]
+    pp_group = SimpleNamespace(cpu_group=object())
+    monkeypatch.setattr(worker, "get_pp_group", lambda: pp_group)
+
+    def gather_prefixes(output, local, *, group):
+        assert group is pp_group.cpu_group
+        output[:] = stage_prefixes
+
+    for stage in stages:
+        gather = MagicMock(side_effect=gather_prefixes)
+        monkeypatch.setattr(torch.distributed, "all_gather_object", gather)
+        stage._init_lookup_key_prefixes()
+        assert stage.lookup(33, hashes).hit_length == (
+            32 if missing_group is None else 0
+        )
+        gather.assert_called_once()
+        assert gather.call_args.args[1] == stage_prefixes[stage.pp_rank]
+        assert stage._lookup_key_prefixes == tuple(
+            tuple(prefix for groups in stage_prefixes for prefix in groups[g])
+            for g in range(2)
+        )
+
+
 def test_save_decode_cache_keeps_transfer_path_enabled(tmp_path, monkeypatch):
     store = MagicMock()
     store.setup.return_value = 0
@@ -3032,7 +3446,7 @@ def test_worker_scales_uniform_attention_group_under_dcp(tmp_path, monkeypatch):
     assert scaled_spec.block_size == 64
     assert all(spec.block_size == 64 for spec in scaled_spec.kv_cache_specs.values())
     assert w.coord.attention_groups[0].spec.block_size == 64
-    assert w.token_dbs[0].block_size == 64
+    assert w.token_dbs[0].chunk_size == 64
 
 
 @pytest.mark.parametrize("dcp_size", [1, 4])
@@ -3077,7 +3491,7 @@ def test_worker_put_striding_covers_every_rank_get_namespace(
             _make_kv_cache_config(),
         )
         db = w.token_dbs[0]
-        token_len = len(block_hashes) * db.block_size
+        token_len = len(block_hashes) * db.chunk_size
         keys = [
             PoolKey(db.metadata, block_hash.hex()).to_string()
             for _, _, block_hash in db.process_tokens(token_len, block_hashes)
@@ -3124,12 +3538,13 @@ def test_mqa_p4_to_d2_uses_shared_rank_zero_namespace(tmp_path, monkeypatch):
                     rank=tp_rank,
                     extra_config={"store_tp_size": 4},
                 ),
-                _make_kv_cache_config(),
+                _make_kv_cache_config(num_kv_heads=1),
             )
 
-            assert w.store_tp_size is None
+            assert w.store_tp_size == 4
             assert w.token_dbs[0].metadata.tp_rank == 0
-            key = w.token_dbs[0].key_for(block_hash)
+            layout = w.token_dbs[0].store_layout
+            key = layout.key_for(0, block_hash)
             if tp_size == 4:
                 put_step = w._group_tp_replication_factors[0]
                 if tp_rank % put_step == 0:
@@ -3137,16 +3552,11 @@ def test_mqa_p4_to_d2_uses_shared_rank_zero_namespace(tmp_path, monkeypatch):
             else:
                 target.add(key)
 
-    assert (
-        put_keys
-        == get_keys
-        == {
-            (
-                "test-model@store_pp:1@store_format:tp_shared_mqa"
-                "@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@68617368"
-            )
-        }
-    )
+    assert put_keys == get_keys
+    assert len(put_keys) == 1
+    key = next(iter(put_keys))
+    assert "@store_tp:4@store_pp:1@store_format:attention_head_major" in key
+    assert "@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@68617368" in key
 
 
 def test_requester_worker_group_semantics_falls_back_without_group_ids(
@@ -3304,6 +3714,178 @@ def test_store_sending_thread_clamps_token_len_to_lcm():
     assert len(keys) == 2
 
 
+@pytest.mark.parametrize("failure", ["partial", "exception"])
+def test_store_sending_thread_waits_for_hybrid_checkpoint_and_retries(failure):
+    from vllm.v1.kv_cache_interface import MambaSpec
+
+    store = MagicMock()
+    full = FullAttentionSpec(
+        block_size=8,
+        num_kv_heads=1,
+        head_size=2,
+        dtype=torch.float16,
+    )
+    mamba = MambaSpec(
+        block_size=8,
+        shapes=((1, 1),),
+        dtypes=(torch.float32,),
+        mamba_cache_mode="align",
+    )
+    coord = mooncake_store_worker.MooncakeStoreCoordinator(
+        [
+            KVCacheGroupSpec(["full"], full),
+            KVCacheGroupSpec(["mamba"], mamba),
+        ],
+        scheduler_block_size=8,
+        hash_block_size=4,
+    )
+
+    metadata = KeyMetadata("test-model", 0, 0, 0, 0, group_id=0)
+    layout = HeadMajorStoreLayout(
+        metadata,
+        block_size=8,
+        hash_block_size=4,
+        local_tp_size=2,
+        store_tp_size=2,
+        tp_rank=0,
+        layer_specs=(full,),
+        store_chunk_size=4,
+    )
+    tensor = torch.empty((3, 1, 8, 4), dtype=torch.float16)
+    layout.register_kv_caches([tensor], num_blocks=3)
+    full_db = ChunkedTokenDatabase(
+        metadata,
+        chunk_size=4,
+        hash_block_size=4,
+        store_layout=layout,
+    )
+    mamba_db = ChunkedTokenDatabase(
+        KeyMetadata("test-model", 0, 0, 0, 0, group_id=1),
+        chunk_size=8,
+        hash_block_size=4,
+    )
+    mamba_db.set_kv_caches_base_addr([0x2000])
+    mamba_db.set_block_len([256])
+
+    thread = _make_store_sending_thread(
+        store,
+        coord=coord,
+        token_databases=[full_db, mamba_db],
+        block_size=8,
+        enable_kv_event=True,
+    )
+    thread.update_kv_event = MagicMock()
+    thread._saved_offset["r0"] = 8
+    stored_keys: set[str] = set()
+    fail_put = True
+
+    def put(keys, addrs, sizes, replicate_config):
+        if fail_put and "@group:0" in keys[0]:
+            if failure == "exception":
+                raise RuntimeError("Store temporarily unavailable")
+            stored_keys.update(keys[:-1])
+            return [256] * (len(keys) - 1) + [-1]
+        stored_keys.update(keys)
+        return [256] * len(keys)
+
+    store.batch_is_exist.side_effect = lambda keys: [
+        int(key in stored_keys) for key in keys
+    ]
+    store.batch_put_from_multi_buffers.side_effect = put
+    _run_store_req(
+        thread,
+        ReqMeta(
+            req_id="r0",
+            token_len_chunk=12,
+            block_ids=([1, 2, 3], [4, 5, 6]),
+            block_hashes=[b"a0", b"a1", b"a2"],
+            can_save=True,
+            token_ids=[8, 9, 10, 11],
+            token_ids_start=8,
+        ),
+    )
+
+    store.batch_is_exist.assert_not_called()
+    store.batch_put_from_multi_buffers.assert_not_called()
+    thread.update_kv_event.assert_not_called()
+    assert thread._saved_offset["r0"] == 8
+    _run_store_req(
+        thread,
+        ReqMeta(
+            req_id="r0",
+            token_len_chunk=16,
+            block_ids=([1, 2, 3], [4, 5, 6]),
+            block_hashes=[b"a0", b"a1", b"a2", b"a3"],
+            can_save=True,
+            token_ids=list(range(8, 16)),
+            token_ids_start=8,
+            boundary_puts=[(1, 5, 16)],
+        ),
+    )
+
+    keys = store.batch_put_from_multi_buffers.call_args.args[0]
+    assert len(keys) == 2
+    assert all("@group:0" in key for key in keys)
+    assert any("@group:1" in key for key in stored_keys)
+    assert thread._saved_offset["r0"] == 8
+    assert thread._retry_token_ids["r0"] == (8, list(range(8, 16)))
+    if failure == "partial":
+        events = thread.update_kv_event.call_args.args[0]
+        assert [event.token_ids for event in events] == [list(range(8, 12))]
+    else:
+        thread.update_kv_event.assert_not_called()
+
+    fail_put = False
+    store.batch_put_from_multi_buffers.reset_mock()
+    thread.update_kv_event.reset_mock()
+    _run_store_req(
+        thread,
+        ReqMeta(
+            req_id="r0",
+            token_len_chunk=24,
+            block_ids=([1, 2, 3], [4, 5, 6]),
+            block_hashes=[b"a0", b"a1", b"a2", b"a3", b"a4", b"a5"],
+            can_save=True,
+            token_ids=list(range(16, 24)),
+            token_ids_start=16,
+            boundary_puts=[(1, 6, 24)],
+        ),
+    )
+    keys = store.batch_put_from_multi_buffers.call_args.args[0]
+    retry_start = 12 if failure == "partial" else 8
+    assert len(keys) == (24 - retry_start) // 4
+    assert all("@group:0" in key for key in keys)
+    assert thread._saved_offset["r0"] == 24
+    assert "r0" not in thread._retry_token_ids
+    events = thread.update_kv_event.call_args.args[0]
+    assert [event.token_ids for event in events] == [
+        list(range(start, start + 4)) for start in range(retry_start, 24, 4)
+    ]
+
+
+@pytest.mark.parametrize("block_ids", [[1], [1, NULL_BLOCK_ID]])
+@pytest.mark.parametrize("first_exists", [0, 1])
+def test_store_sending_thread_missing_source_keeps_checkpoint(block_ids, first_exists):
+    store = MagicMock()
+    store.batch_is_exist.return_value = [first_exists]
+    store.batch_put_from_multi_buffers.return_value = [256]
+    thread = _make_store_sending_thread(store)
+
+    req = _make_store_req("req-a", [b"a0", b"a1"])
+    req.block_ids = (block_ids,)
+    _run_store_req(thread, req)
+    assert thread._saved_offset.get("req-a", 0) == 0
+
+    store.batch_is_exist.return_value = [1, 0]
+    store.batch_put_from_multi_buffers.reset_mock()
+    _run_store_req(thread, _make_store_req("req-a", [b"a0", b"a1"]))
+
+    keys = store.batch_put_from_multi_buffers.call_args.args[0]
+    assert len(keys) == 1
+    assert keys[0].endswith("@6131")
+    assert thread._saved_offset["req-a"] == 32
+
+
 def test_store_sending_thread_skips_when_token_len_below_lcm():
     """Requests shorter than lcm_block_size cannot produce any aligned chunk,
     so neither the existence check nor the put should be issued."""
@@ -3319,7 +3901,7 @@ def test_store_sending_thread_skips_when_token_len_below_lcm():
     )
     db = ChunkedTokenDatabase(
         KeyMetadata("test-model", 0, 0, 0, 0, group_id=0),
-        block_size=64,
+        chunk_size=64,
         hash_block_size=64,
     )
     db.set_kv_caches_base_addr([0x1000])
@@ -3385,14 +3967,14 @@ def successful_put(keys, addrs, sizes, *_args):
 
     db_full = ChunkedTokenDatabase(
         KeyMetadata("test-model", 0, 0, 0, 0, group_id=0),
-        block_size=32,
+        chunk_size=32,
         hash_block_size=8,
     )
     db_full.set_kv_caches_base_addr([0x1000])
     db_full.set_block_len([512])
     db_swa = ChunkedTokenDatabase(
         KeyMetadata("test-model", 0, 0, 0, 0, group_id=1),
-        block_size=8,
+        chunk_size=8,
         hash_block_size=8,
     )
     db_swa.set_kv_caches_base_addr([0x2000])
@@ -3460,14 +4042,14 @@ def test_store_sending_thread_delta_saves_only_new_swa_boundary_chunks():
 
     db_full = ChunkedTokenDatabase(
         KeyMetadata("test-model", 0, 0, 0, 0, group_id=0),
-        block_size=32,
+        chunk_size=32,
         hash_block_size=8,
     )
     db_full.set_kv_caches_base_addr([0x1000])
     db_full.set_block_len([512])
     db_swa = ChunkedTokenDatabase(
         KeyMetadata("test-model", 0, 0, 0, 0, group_id=1),
-        block_size=8,
+        chunk_size=8,
         hash_block_size=8,
     )
     db_swa.set_kv_caches_base_addr([0x2000])
@@ -3530,14 +4112,14 @@ def test_store_sending_thread_kv_events_use_group_chunk_metadata():
 
     db_full = ChunkedTokenDatabase(
         KeyMetadata("test-model", 0, 0, 0, 0, group_id=0),
-        block_size=32,
+        chunk_size=32,
         hash_block_size=8,
     )
     db_full.set_kv_caches_base_addr([0x1000])
     db_full.set_block_len([512])
     db_swa = ChunkedTokenDatabase(
         KeyMetadata("test-model", 0, 0, 0, 0, group_id=1),
-        block_size=8,
+        chunk_size=8,
         hash_block_size=8,
     )
     db_swa.set_kv_caches_base_addr([0x2000])
@@ -3720,11 +4302,15 @@ def _make_bare_worker(
     worker.kv_recv_threads = []
     worker.num_recv_threads = 1
     worker.recv_request_queue = queue.Queue()
+    worker._device_regions = {}
+    worker._released_regions = set()
     worker.finished_store_req = set()
     worker.tp_size = 1
     worker.store_tp_size = None
+    worker.requested_store_chunk_size = None
     worker.num_kv_head = 1
     worker.pp_size = 1
+    worker.pp_rank = 0
     # Minimal single-full-attention-group config so the coordinator-based
     # lookup path works (the connector no longer carries a legacy single-group
     # path; everything flows through the coordinator).
@@ -3767,7 +4353,7 @@ def _make_bare_worker(
     worker.token_dbs = [
         ChunkedTokenDatabase(
             KeyMetadata("test-model", 0, 0, 0, 0, group_id=0),
-            block_size=block_size,
+            chunk_size=block_size,
             hash_block_size=block_size,
         )
     ]
@@ -3943,15 +4529,22 @@ def test_lookup_key_prefixes_cover_store_tp_shards():
     worker.token_dbs = [
         ChunkedTokenDatabase(
             metadata,
-            block_size=16,
-            store_layout=LBHNCStoreLayout(
+            chunk_size=16,
+            store_layout=HeadMajorStoreLayout(
                 metadata,
                 block_size=16,
                 hash_block_size=16,
                 local_tp_size=2,
                 store_tp_size=4,
                 tp_rank=0,
-                num_kv_heads=8,
+                layer_specs=(
+                    FullAttentionSpec(
+                        block_size=16,
+                        num_kv_heads=4,
+                        head_size=2,
+                        dtype=torch.float16,
+                    ),
+                ),
             ),
         )
     ]
@@ -4000,10 +4593,10 @@ def test_lookup_key_prefixes_expand_tp_sharded_groups_per_rank():
     ]
     worker.token_dbs = [
         ChunkedTokenDatabase(
-            KeyMetadata("test-model", 0, 0, 0, 0, group_id=0), block_size=16
+            KeyMetadata("test-model", 0, 0, 0, 0, group_id=0), chunk_size=16
         ),
         ChunkedTokenDatabase(
-            KeyMetadata("test-model", 1, 0, 0, 0, group_id=1), block_size=16
+            KeyMetadata("test-model", 1, 0, 0, 0, group_id=1), chunk_size=16
         ),
     ]
     _refresh_group_tp_replication_factors(worker)
@@ -4043,7 +4636,7 @@ def test_group_tp_replication_factors_mixed_mla_gqa_mamba():
     ]
     worker.token_dbs = [
         ChunkedTokenDatabase(
-            KeyMetadata("test-model", 0, 0, 0, 0, group_id=g_idx), block_size=16
+            KeyMetadata("test-model", 0, 0, 0, 0, group_id=g_idx), chunk_size=16
         )
         for g_idx in range(3)
     ]
@@ -4103,6 +4696,115 @@ def test_uniform_group_uses_common_inner_replication_factor(spec_order):
     )
 
 
+def test_store_group_plans_cover_kimi_k3_and_deepseek_v4_specs(monkeypatch):
+    from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
+    from vllm.v1.kv_cache_interface import (
+        KVCacheGroupSpec,
+        MambaSpec,
+        MLAAttentionSpec,
+        SlidingWindowMLASpec,
+        UniformTypeKVCacheSpecs,
+    )
+
+    worker = _make_bare_worker(block_size=64)
+    worker.tp_size = 4
+    worker.tp_rank = 3
+    worker.num_kv_head = 1
+    monkeypatch.setenv("VLLM_SSM_CONV_STATE_LAYOUT", "DS")
+
+    mla_specs = {
+        name: MLAAttentionSpec(
+            block_size=64,
+            num_kv_heads=1,
+            head_size=64,
+            dtype=torch.bfloat16,
+            tokens_per_state=4,
+            model_version="deepseek_v4",
+        )
+        for name in ("mla0", "mla1")
+    }
+    swa_specs = {
+        name: SlidingWindowMLASpec(
+            block_size=64,
+            num_kv_heads=1,
+            head_size=64,
+            dtype=torch.bfloat16,
+            tokens_per_state=4,
+            model_version="deepseek_v4",
+            sliding_window=128,
+        )
+        for name in ("swa0", "swa1")
+    }
+    gdn = MambaSpec(
+        block_size=64,
+        shapes=((6, 3), (1, 2, 2)),
+        dtypes=(torch.uint8, torch.uint8),
+        mamba_type=MambaAttentionBackendEnum.GDN_ATTN,
+    )
+    groups = [
+        KVCacheGroupSpec(
+            list(mla_specs),
+            UniformTypeKVCacheSpecs(block_size=64, kv_cache_specs=mla_specs),
+        ),
+        KVCacheGroupSpec(
+            list(swa_specs),
+            UniformTypeKVCacheSpecs(block_size=64, kv_cache_specs=swa_specs),
+        ),
+        KVCacheGroupSpec(["gdn"], gdn),
+    ]
+
+    plans = [
+        worker._make_store_group_plan(
+            group,
+            requested_store_tp_size=4,
+            attention_layout_cls=HeadMajorStoreLayout,
+            replication_factor=worker._spec_tp_replication_factor(group.kv_cache_spec),
+        )
+        for group in groups
+    ]
+
+    assert all(plan is not None for plan in plans)
+    assert [plan.layout_cls for plan in plans if plan is not None] == [
+        HeadMajorStoreLayout,
+        HeadMajorStoreLayout,
+        MambaStoreLayout,
+    ]
+    assert [plan.store_shard_count for plan in plans if plan is not None] == [1, 1, 4]
+
+
+def test_attention_store_plan_uses_value_schema():
+    from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheGroupSpec
+
+    worker = _make_bare_worker(block_size=16)
+    worker.tp_size = 2
+    worker.tp_rank = 0
+    worker.num_kv_head = 8
+
+    def plan(dtype: torch.dtype, layout_cls):
+        spec = FullAttentionSpec(
+            block_size=16,
+            num_kv_heads=4,
+            head_size=64,
+            dtype=dtype,
+        )
+        return worker._make_store_group_plan(
+            KVCacheGroupSpec(["layer"], spec),
+            requested_store_tp_size=4,
+            attention_layout_cls=layout_cls,
+            replication_factor=1,
+        )
+
+    head_major = plan(torch.bfloat16, HeadMajorStoreLayout)
+    token_major = plan(torch.bfloat16, TokenMajorStoreLayout)
+    fp8 = plan(torch.float8_e4m3fn, HeadMajorStoreLayout)
+
+    assert head_major is not None and token_major is not None and fp8 is not None
+    assert head_major.layout_cls.store_format == "attention_head_major"
+    assert token_major.layout_cls.store_format == "attention_token_major"
+    assert head_major.schema_fingerprint == token_major.schema_fingerprint
+    assert head_major.schema_fingerprint != fp8.schema_fingerprint
+
+
 def test_lookup_rejects_boundary_missing_one_mamba_shard():
     from vllm.v1.kv_cache_interface import (
         FullAttentionSpec,
@@ -4126,10 +4828,10 @@ def test_lookup_rejects_boundary_missing_one_mamba_shard():
     ]
     worker.token_dbs = [
         ChunkedTokenDatabase(
-            KeyMetadata("test-model", 0, 0, 0, 0, group_id=0), block_size=16
+            KeyMetadata("test-model", 0, 0, 0, 0, group_id=0), chunk_size=16
         ),
         ChunkedTokenDatabase(
-            KeyMetadata("test-model", 1, 0, 0, 0, group_id=1), block_size=16
+            KeyMetadata("test-model", 1, 0, 0, 0, group_id=1), chunk_size=16
         ),
     ]
     worker.coord = mooncake_store_worker.MooncakeStoreCoordinator(
@@ -4150,6 +4852,56 @@ def test_lookup_rejects_boundary_missing_one_mamba_shard():
     assert worker.lookup(33, [b"h0", b"h1"]).hit_length == 0
 
 
+@pytest.mark.parametrize("mamba_block_size", [4, 16])
+@pytest.mark.parametrize("missing_chunk, expected_hit", [(None, 32), (1, 0), (5, 16)])
+def test_lookup_requires_all_attention_chunks_in_page(
+    mamba_block_size, missing_chunk, expected_hit
+):
+    """A page-end key cannot hide an absent interior Store chunk."""
+    worker = _make_bare_worker(block_size=16)
+    worker.hash_block_size = 4
+    worker._kv_cache_groups = [
+        KVCacheGroupSpec(
+            ["attention"],
+            FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=64, dtype=None),
+        ),
+        KVCacheGroupSpec(
+            ["mamba"],
+            MambaSpec(
+                block_size=mamba_block_size,
+                shapes=((1, 1),),
+                dtypes=(torch.float32,),
+                mamba_cache_mode="align",
+            ),
+        ),
+    ]
+    worker.token_dbs = [
+        ChunkedTokenDatabase(
+            KeyMetadata("test-model", 0, 0, 0, 0, group_id=group_id),
+            chunk_size=chunk_size,
+            hash_block_size=4,
+        )
+        for group_id, chunk_size in enumerate((4, mamba_block_size))
+    ]
+    worker.coord = mooncake_store_worker.MooncakeStoreCoordinator(
+        worker._kv_cache_groups, scheduler_block_size=16, hash_block_size=4
+    )
+    _refresh_group_tp_replication_factors(worker)
+    hashes = [BlockHash(f"h{i}".encode()) for i in range(8)]
+    present = {
+        f"@group:{group_id}@{block_hash.hex()}"
+        for group_id in range(2)
+        for index, block_hash in enumerate(hashes)
+        if (group_id == 0 and index != missing_chunk)
+        or (group_id == 1 and index in (3, 7))
+    }
+    worker.store.batch_is_exist.side_effect = lambda keys: [
+        int(any(key.endswith(suffix) for suffix in present)) for key in keys
+    ]
+
+    assert worker.lookup(33, hashes).hit_length == expected_hit
+
+
 def test_lookup_requires_all_dcp_rank_namespaces():
     worker = _make_bare_worker(block_size=16)
     worker.tp_size = 4
@@ -4197,7 +4949,7 @@ def test_lookup_partial_tail_uses_hash_alignment():
     worker.token_dbs = [
         ChunkedTokenDatabase(
             KeyMetadata("test-model", 0, 0, 0, 0, group_id=group_id),
-            block_size=16,
+            chunk_size=16,
             hash_block_size=4,
         )
         for group_id in range(2)
@@ -4247,6 +4999,44 @@ def test_lookup_full_hit_with_eagle_pops_once_not_twice():
     assert worker.store.batch_is_exist.call_count == 1
 
 
+def test_lookup_skips_tail_key_for_empty_chunked_local_window():
+    from vllm.v1.kv_cache_interface import ChunkedLocalAttentionSpec
+
+    worker = _make_bare_worker(block_size=16)
+    worker._kv_cache_groups.append(
+        KVCacheGroupSpec(
+            ["local"],
+            ChunkedLocalAttentionSpec(
+                block_size=16,
+                num_kv_heads=8,
+                head_size=64,
+                dtype=None,
+                attention_chunk_size=32,
+            ),
+        )
+    )
+    worker.token_dbs = [
+        ChunkedTokenDatabase(
+            KeyMetadata("test-model", 0, 0, 0, 0, group_id=group_id),
+            chunk_size=16,
+            hash_block_size=16,
+        )
+        for group_id in range(2)
+    ]
+    worker.coord = mooncake_store_worker.MooncakeStoreCoordinator(
+        worker._kv_cache_groups, scheduler_block_size=16, hash_block_size=16
+    )
+    _refresh_group_tp_replication_factors(worker)
+    worker.store.batch_is_exist.side_effect = lambda keys: [
+        int("@group:0@" in key) for key in keys
+    ]
+
+    result = worker.lookup(33, [BlockHash(b"h0"), BlockHash(b"h1")])
+
+    assert result.hit_length == 32
+    assert result.tail_key_boundaries == (TailKeyBoundary(0, 32),)
+
+
 def test_lookup_plan_resolves_group_tail_keys_from_existing_hashes():
     from vllm.v1.kv_cache_interface import (
         FullAttentionSpec,
@@ -4270,7 +5060,7 @@ def test_lookup_plan_resolves_group_tail_keys_from_existing_hashes():
     worker.token_dbs = [
         ChunkedTokenDatabase(
             KeyMetadata("test-model", 0, 0, 0, 0, group_id=group_id),
-            block_size=16,
+            chunk_size=16,
             hash_block_size=4,
         )
         for group_id in range(2)
@@ -4341,7 +5131,7 @@ def test_lookup_plan_recovers_tail_key_after_multi_chunk_convergence():
     worker.token_dbs = [
         ChunkedTokenDatabase(
             KeyMetadata("test-model", 0, 0, 0, 0, group_id=group_id),
-            block_size=16,
+            chunk_size=16,
             hash_block_size=4,
         )
         for group_id in range(2)
@@ -4451,12 +5241,12 @@ def test_lookup_checks_all_potential_swa_hit_boundaries():
     worker.token_dbs = [
         ChunkedTokenDatabase(
             KeyMetadata("test-model", 0, 0, 0, 0, group_id=0),
-            block_size=32,
+            chunk_size=32,
             hash_block_size=8,
         ),
         ChunkedTokenDatabase(
             KeyMetadata("test-model", 0, 0, 0, 0, group_id=1),
-            block_size=8,
+            chunk_size=8,
             hash_block_size=8,
         ),
     ]
@@ -4512,12 +5302,12 @@ def test_lookup_applies_swa_mask_before_accessing_hashes():
     worker.token_dbs = [
         ChunkedTokenDatabase(
             KeyMetadata("test-model", 0, 0, 0, 0, group_id=0),
-            block_size=32,
+            chunk_size=32,
             hash_block_size=8,
         ),
         ChunkedTokenDatabase(
             KeyMetadata("test-model", 0, 0, 0, 0, group_id=1),
-            block_size=8,
+            chunk_size=8,
             hash_block_size=8,
         ),
     ]
@@ -4563,6 +5353,27 @@ def test_consumer_starts_send_thread_only_when_put_is_enabled(save_decode_cache)
     assert (worker.kv_send_thread is not None) == save_decode_cache
 
 
+def test_tp_shared_registration_follows_cache_group_layer_order():
+    worker = _make_bare_worker(num_gpu_blocks=2)
+    worker.store_tp_size = 2
+    spec = FullAttentionSpec(
+        block_size=16, num_kv_heads=2, head_size=8, dtype=torch.float16
+    )
+    worker._kv_cache_groups = [KVCacheGroupSpec(["second", "first"], spec)]
+    db = MagicMock()
+    worker.token_dbs = [db]
+    first = torch.zeros((2, 1, 1, 8), dtype=torch.uint8)
+    second = torch.ones_like(first)
+
+    _register_with_mocked_threads(worker, {"first": first, "second": second})
+
+    registered = db.store_layout.register_kv_caches.call_args.args[0]
+    assert [cache.data_ptr() for cache in registered] == [
+        second.data_ptr(),
+        first.data_ptr(),
+    ]
+
+
 def test_start_load_kv_queues_async_load():
     w = _make_bare_worker()
     req = _make_load_req(
@@ -4625,24 +5436,34 @@ def test_register_kv_caches_shared_storage(layout: KVCacheLayout):
     if not layout.is_block_compact:
         # Each head group is its own region, so a block spans L*H regions.
         head_cache = caches[0][:, 0]
-        assert db.kv_caches_base_addr == [
+        assert db.store_layout.kv_caches_base_addr == [
             cache[:, head_idx].data_ptr()
             for cache in caches
             for head_idx in range(cache.shape[1])
         ]
-        assert db.block_len == [head_cache.stride(0) * head_cache.element_size()] * (
-            num_layers * caches[0].shape[1]
-        )
+        assert db.store_layout.block_stride == [
+            head_cache.stride(0) * head_cache.element_size()
+        ] * (num_layers * caches[0].shape[1])
     elif layout.is_layer_compact:
-        assert db.kv_caches_base_addr == [cache.data_ptr() for cache in caches]
-        assert db.block_len == [spec.page_size_bytes] * num_layers
+        assert db.store_layout.kv_caches_base_addr == [
+            cache.data_ptr() for cache in caches
+        ]
+        assert db.store_layout.block_stride == [spec.page_size_bytes] * num_layers
     else:
-        assert db.kv_caches_base_addr == [raw.data_ptr()]
-        assert db.block_len == [num_layers * spec.page_size_bytes]
+        assert db.store_layout.kv_caches_base_addr == [raw.data_ptr()]
+        assert db.store_layout.block_stride == [num_layers * spec.page_size_bytes]
     worker.store.register_buffer.assert_called_once_with(raw.data_ptr(), raw.nbytes)
 
 
-def test_register_kv_caches_uses_transfer_group_memory_domain():
+@pytest.mark.parametrize(
+    "extra_config",
+    [
+        {},
+        {"store_tp_size": 2},
+        {"enable_store_tp_lcm": True, "prefill_tp_sizes": [1, 2]},
+    ],
+)
+def test_register_kv_caches_uses_transfer_group_memory_domain(extra_config):
     """Derived GPU caches must not change host-source transfer addresses."""
     from vllm.v1.kv_cache_interface import (
         KVCacheConfig,
@@ -4694,9 +5515,19 @@ def test_register_kv_caches_uses_transfer_group_memory_domain():
     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)
-    ]
+    worker.cache_config.prefix_match_unit = None
+    worker.cache_config.get_resolved_kv_cache_layout.return_value = KVCacheLayout.LBHNC
+    worker._group_tp_replication_factors = (
+        worker._compute_group_tp_replication_factors()
+    )
+    worker.store_tp_size, namespace, plans = worker._select_store_layout(extra_config)
+    worker.token_dbs = worker._build_token_databases(
+        KeyMetadata("test-model", 0, 0, 0, 0, store_namespace=namespace), plans
+    )
+    assert worker.store_tp_size is None
+    assert isinstance(worker.token_dbs[0].store_layout, RankLocalStoreLayout)
+    if extra_config:
+        assert "@store_format:rank_local_tp1" in namespace
     source = torch.zeros(host_num_blocks, page_size, dtype=torch.uint8)
     indexer = torch.zeros(gpu_num_blocks, page_size, dtype=torch.uint8)
 
@@ -4706,8 +5537,8 @@ def test_register_kv_caches_uses_transfer_group_memory_domain():
     )
 
     db = worker.token_dbs[0]
-    assert db.kv_caches_base_addr == [source.data_ptr()]
-    assert db.block_len == [page_size]
+    assert db.store_layout.kv_caches_base_addr == [source.data_ptr()]
+    assert db.store_layout.block_stride == [page_size]
     worker.store.register_buffer.assert_called_once_with(
         source.data_ptr(), source.nbytes
     )
@@ -4769,11 +5600,122 @@ def test_register_kv_caches_separate_head_groups():
         for head_idx in range(cache.shape[1])
     ]
     db = worker.token_dbs[0]
-    assert db.kv_caches_base_addr == expected_addrs
-    assert db.block_len == [head_block_bytes] * len(expected_addrs)
+    assert db.store_layout.kv_caches_base_addr == expected_addrs
+    assert db.store_layout.block_stride == [head_block_bytes] * len(expected_addrs)
     worker.store.register_buffer.assert_called_once_with(raw.data_ptr(), raw.nbytes)
 
 
+# ---------------------------------------------------------------------------
+# Sleep mode: release / restore of the KV cache registration
+# ---------------------------------------------------------------------------
+
+
+class _FakeRegistrationStore:
+    """Keeps the registered regions like the store; `fail` fails one region."""
+
+    def __init__(self):
+        self.registered: dict[int, int] = {}
+        self.fail: tuple[str, int] | None = None
+
+    def register_buffer(self, addr: int, length: int) -> int:
+        if self.fail == ("register", addr) or addr in self.registered:
+            return -1
+        self.registered[addr] = length
+        return 0
+
+    def unregister_buffer(self, addr: int) -> int:
+        if self.fail == ("unregister", addr) or addr not in self.registered:
+            return -1
+        del self.registered[addr]
+        return 0
+
+
+def _registered_worker(kind: str = "dense"):
+    """A worker that registered its KV caches with a fake store, plus the device
+    and host regions it registered."""
+    worker = _make_bare_worker(num_gpu_blocks=2)
+    if kind == "dense":  # two device regions
+        worker.token_dbs[0].store_layout = MagicMock()
+        caches = {n: torch.zeros(2, 4) for n in ("layer0", "__cross_layer__")}
+    else:  # one host region and one device region
+        from vllm.v1.kv_cache_interface import KVCacheTensor
+
+        spec = FullAttentionSpec(
+            block_size=16, num_kv_heads=1, head_size=1, dtype=torch.float16
+        )
+        config = KVCacheConfig(
+            num_blocks=2,
+            kv_cache_tensors=[
+                KVCacheTensor(
+                    size=4 * 32,
+                    layers=["source"],
+                    layer_stride=4 * 32,
+                    block_stride=32,
+                    host_resident=True,
+                ),
+                KVCacheTensor(
+                    size=2 * 32, layers=["layer0"], layer_stride=2 * 32, block_stride=32
+                ),
+            ],
+            kv_cache_groups=[
+                KVCacheGroupSpec(["source"], spec, host_resident=True),
+                KVCacheGroupSpec(["layer0"], spec),
+            ],
+            hisparse_host_num_blocks=4,
+        )
+        worker._kv_cache_config = config
+        worker._kv_cache_groups = list(config.transfer_groups)
+        worker.token_dbs = [
+            ChunkedTokenDatabase(KeyMetadata("test-model", 0, 0, 0, 0), chunk_size=16)
+            for _ in worker._kv_cache_groups
+        ]
+        caches = {
+            "source": torch.zeros(4, 32, dtype=torch.uint8),
+            "layer0": torch.zeros(2, 32, dtype=torch.uint8),
+        }
+    worker.store = _FakeRegistrationStore()
+    _register_with_mocked_threads(worker, caches)
+    worker.kv_send_thread = SimpleNamespace(request_queue=queue.Queue())
+    host = {c.data_ptr(): c.nbytes for n, c in caches.items() if n == "source"}
+    device = {c.data_ptr(): c.nbytes for n, c in caches.items() if n != "source"}
+    return worker, device, host
+
+
+@pytest.mark.cpu_test
+@pytest.mark.parametrize("kind", ["dense", "hisparse"])
+def test_sleep_cycle_releases_and_restores_device_regions(kind: str):
+    """Asleep, only host regions stay registered; waking registers the same
+    device addresses and lengths again. Both calls are idempotent."""
+    worker, device, host = _registered_worker(kind)
+    assert device and worker.store.registered == device | host
+    worker.restore_kv_caches()  # a wake-up without a sleep
+    for _ in range(2):
+        worker.unregister_kv_caches()
+        worker.unregister_kv_caches()  # sleeping after a discard
+        assert worker.store.registered == host
+        worker.restore_kv_caches()
+        worker.restore_kv_caches()  # a retried wake-up
+        assert worker.store.registered == device | host
+
+
+@pytest.mark.cpu_test
+@pytest.mark.parametrize("op", ["unregister", "register"])
+def test_failed_store_call_raises_and_retry_converges(op: str):
+    """A failed region keeps its state and is the only one a retry touches."""
+    worker, device, _ = _registered_worker()
+    if op == "register":
+        worker.unregister_kv_caches()
+    step = worker.restore_kv_caches if op == "register" else worker.unregister_kv_caches
+    _, failing = device
+    worker.store.fail = (op, failing)
+    with pytest.raises(RuntimeError, match=f"{op}_buffer failed"):
+        step()
+    assert (failing in worker.store.registered) == (op == "unregister")
+    worker.store.fail = None
+    step()
+    assert worker.store.registered == (device if op == "register" else {})
+
+
 # ---------------------------------------------------------------------------
 # Dual-mode (embedded / standalone-store) config validation tests
 # ---------------------------------------------------------------------------
diff --git a/tests/v1/kv_connector/unit/test_moriio_connector.py b/tests/v1/kv_connector/unit/test_moriio_connector.py
index 438e20bd0809..381182869660 100644
--- a/tests/v1/kv_connector/unit/test_moriio_connector.py
+++ b/tests/v1/kv_connector/unit/test_moriio_connector.py
@@ -1182,19 +1182,13 @@ def test_single_group_path_unchanged():
     assert scheduler.get_exchange_clipped_blocks(blocks) == blocks
 
 
-class _FakeSpecConfig(SimpleNamespace):
-    """Minimal speculative_config stand-in. Any drafter-capability probe
-    (use_dflash/use_eagle/use_dspark/uses_draft_model/...) defaults to False so
-    VllmConfig.num_lookahead_tokens resolves cleanly as upstream adds methods."""
-
-    def __getattr__(self, name: str):
-        if name.startswith(("use_", "uses_")):
-            return lambda *a, **k: False
-        raise AttributeError(name)
-
-
 def _fake_speculative_config(num_speculative_tokens: int = 1) -> SimpleNamespace:
-    return _FakeSpecConfig(num_speculative_tokens=num_speculative_tokens)
+    """Minimal speculative_config stand-in, enough for
+    VllmConfig.num_lookahead_tokens."""
+    return SimpleNamespace(
+        num_speculative_tokens=num_speculative_tokens,
+        speculator_config=SimpleNamespace(num_lookahead_tokens=num_speculative_tokens),
+    )
 
 
 def _spec_vllm_config() -> VllmConfig:
diff --git a/tests/v1/kv_connector/unit/test_moriio_hetero_tp_write.py b/tests/v1/kv_connector/unit/test_moriio_hetero_tp_write.py
index a1aeedf2fff5..0003667dd350 100644
--- a/tests/v1/kv_connector/unit/test_moriio_hetero_tp_write.py
+++ b/tests/v1/kv_connector/unit/test_moriio_hetero_tp_write.py
@@ -225,7 +225,7 @@ def _build_rank(model, tp_rank, tp_size, layout, hma, state, is_producer) -> Ran
     worker.layer_to_group = {
         layer.name: groups.index(layer.group if hma else "full") for layer in layers
     }
-    worker._mamba_payload_index_by_layer = {
+    worker._hybrid_payload_index_by_layer = {
         layer.name: groups.index("mamba") for layer in layers if layer.kind == "mamba"
     }
     worker._conv_decomp = next(
diff --git a/tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py b/tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py
index 1120062e5410..850e005debd5 100644
--- a/tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py
+++ b/tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py
@@ -10,6 +10,8 @@
 """
 
 import importlib
+import threading
+from collections import defaultdict
 from dataclasses import asdict
 from types import SimpleNamespace
 from unittest.mock import Mock
@@ -18,9 +20,19 @@
 import torch
 
 from vllm.v1.core.block_pool import BlockPool
+from vllm.v1.core.kv_cache_planning import get_kv_cache_groups
 from vllm.v1.core.sched.output import KVConnectorBlockState, SchedulerOutput
 from vllm.v1.core.single_type_kv_cache_manager import FullAttentionManager
-from vllm.v1.kv_cache_interface import FullAttentionSpec, UniformTypeKVCacheSpecs
+from vllm.v1.kv_cache_interface import (
+    CircularBufferSpec,
+    FullAttentionSpec,
+    KVCacheConfig,
+    KVCacheSpec,
+    MLAAttentionSpec,
+    UniformTypeKVCacheSpecs,
+)
+from vllm.v1.kv_cache_layout import KVCacheLayout
+from vllm.v1.request import RequestStatus
 
 moriio_connector = importlib.import_module(
     "vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector"
@@ -56,6 +68,7 @@ def __init__(self, **attrs):
         self.mode = MoRIIOMode.READ
         self._mamba_group_ids: list[int] = []
         self._attn_group_ids: list[int] = [0]
+        self._attn_group_is_state: list[bool] = [False]
         self._num_ssm_scratch_blocks = 0
         self._is_hma_required = False
         self.kv_cache_config = SimpleNamespace(
@@ -173,7 +186,7 @@ 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 attn == [[1, 2]]
     assert mamba == [[99]]
 
 
@@ -210,7 +223,7 @@ def test_split_block_groups_ignores_transfer_disabled_group():
         kv_cache_config=config,
     )
     assert sched.split_block_groups(([1, 2], [70, 71], [99])) == (
-        [1, 2],
+        [[1, 2]],
         [[99]],
     )
 
@@ -232,21 +245,212 @@ def test_exchange_blocks_ignore_transfer_disabled_group():
     ]
 
 
-def test_scheduler_rejects_multiple_attention_groups_with_mamba():
-    """A drafter owning its own transfer group is refused, with a hint."""
+def test_exchange_clipped_blocks_match_select_for_single_full_attention():
+    # With blocks_per_sw == 0 the wire payload must match the pre-multigroup one.
+    def select(block_ids):
+        # Group 1 is not transfer-enabled.
+        return (block_ids[0], block_ids[2])
+
+    config = SimpleNamespace(
+        kv_cache_groups=[None, None, None],
+        select_transfer_block_ids=select,
+    )
+    raw = ([1, 2, 3], [70, 71], [90, 91, 92])
+
+    hma = _FakeScheduler(
+        kv_cache_config=config, _is_hma_required=True, blocks_per_sw=[0, 0]
+    )
+    plain = _FakeScheduler(kv_cache_config=config, _is_hma_required=False)
+
+    clipped = hma.get_exchange_clipped_blocks(raw)
+
+    assert clipped == [list(group) for group in select(raw)]
+    assert clipped == plain.get_exchange_clipped_blocks(raw)
+
+
+def test_scheduler_rejects_hybrid_without_attention_group():
     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="separate attention group"):
+    with pytest.raises(moriio_common.MoRIIOError, match="attention group"):
         moriio_connector.MoRIIOConnectorScheduler(_gate_vllm_config(), "engine", config)
 
 
+def _glm5_next_mtp_kv_cache_config() -> KVCacheConfig:
+    """GLM-5.3-Flash-shaped groups from the real grouping, with an MTP layer."""
+    kpool, num_spec = 4, 1
+    mamba = _mamba_spec(num_speculative_blocks=num_spec)
+    specs: dict[str, KVCacheSpec] = {}
+    for i in range(7):
+        if i % 3 != 2 and i != 6:
+            specs[f"layers.{i}.linear_attn"] = mamba
+            continue
+        specs[f"layers.{i}.attn"] = MLAAttentionSpec(
+            block_size=1024, num_kv_heads=1, head_size=576, dtype=torch.bfloat16
+        )
+        specs[f"layers.{i}.indexer"] = MLAAttentionSpec(
+            block_size=1024,
+            num_kv_heads=1,
+            head_size=132,
+            dtype=torch.uint8,
+            tokens_per_state=kpool,
+        )
+        specs[f"layers.{i}.tail"] = CircularBufferSpec(
+            block_size=2 * kpool,
+            num_kv_heads=2,
+            head_size=128,
+            head_size_v=0,
+            dtype=torch.bfloat16,
+            dcp_sharded=False,
+        )
+    vllm_config = SimpleNamespace(
+        scheduler_config=SimpleNamespace(disable_hybrid_kv_cache_manager=False),
+        parallel_config=SimpleNamespace(pipeline_parallel_size=1),
+        speculative_config=None,
+        cache_config=SimpleNamespace(
+            get_resolved_kv_cache_layout=lambda: KVCacheLayout.BLHNC
+        ),
+    )
+    groups = get_kv_cache_groups(vllm_config, specs)
+    return KVCacheConfig(num_blocks=64, kv_cache_tensors=[], kv_cache_groups=groups)
+
+
+def _glm5_next_mtp_scheduler(kv_role: str = "kv_consumer"):
+    vllm_config = _gate_vllm_config(
+        speculative_config=_spec_config("mtp"),
+        block_size=1024,
+        num_lookahead_tokens=1,
+        num_speculative_tokens=1,
+    )
+    vllm_config.kv_transfer_config.kv_role = kv_role
+    config = _glm5_next_mtp_kv_cache_config()
+    return moriio_connector.MoRIIOConnectorScheduler(vllm_config, "engine", config)
+
+
+@pytest.mark.usefixtures("restore_moriio_role")
+def test_scheduler_accepts_glm5_next_groups_with_mtp():
+    """Several attention groups (MLA + indexer, kpool tail) plus KDA, with MTP."""
+    scheduler = _glm5_next_mtp_scheduler()
+    groups = scheduler.kv_cache_config.transfer_groups
+
+    assert all(isinstance(g.kv_cache_spec, UniformTypeKVCacheSpecs) for g in groups)
+    # The MTP layer shares the target's MLA and tail groups.
+    assert {"layers.6.attn", "layers.6.indexer"} <= set(groups[1].layer_names)
+    assert "layers.6.tail" in groups[2].layer_names
+    assert scheduler._attn_group_ids == [1, 2]
+    assert scheduler._mamba_group_ids == [0]
+    assert scheduler._attn_group_is_state == [False, True]
+    assert scheduler.blocks_per_sw == [0, 0, 0]
+    assert scheduler._num_ssm_scratch_blocks == 1
+    # 1 recomputed token + 1 padded spec row + 1 drafter lookahead slot.
+    assert scheduler._max_decode_tail_blocks == 1
+
+
+@pytest.mark.usefixtures("restore_moriio_role")
+@pytest.mark.parametrize("external_tokens", [0, 1500])
+def test_glm5_next_mtp_read_pairs_every_group(external_tokens):
+    """P payload -> D alignment -> zeroing filter across all transfer groups."""
+    producer = _glm5_next_mtp_scheduler("kv_producer")
+    p_request = SimpleNamespace(
+        request_id="req",
+        kv_transfer_params={"do_remote_decode": True, "transfer_id": "tx"},
+        status=RequestStatus.FINISHED_LENGTH_CAPPED,
+    )
+    # KDA group: [running state, MTP scratch slot].
+    p_blocks = ([90, 91], [10, 11], [50])
+    delay_free, p_params = _FakeConnector(producer).request_finished_all_groups(
+        p_request, p_blocks
+    )
+    assert delay_free is True
+    assert p_params["remote_block_ids"] == [[10, 11], [50], [90]]
+
+    consumer = _glm5_next_mtp_scheduler()
+    d_request = _make_read_request(p_params["remote_block_ids"])
+    # The decode's MLA group holds one extra local block for the spec tail.
+    d_blocks = _FakeBlocks(([200, 201], [100, 101, 102], [300]))
+    consumer.update_state_after_alloc(
+        d_request, d_blocks, num_external_tokens=external_tokens
+    )
+
+    local = consumer._reqs_need_recv["req"][1]
+    if external_tokens:
+        assert local == [[100, 101], [300], [200]]
+        assert consumer._req_kv_params["req"]["remote_block_ids"] == [
+            [10, 11],
+            [50],
+            [90],
+        ]
+    else:
+        # A full local MLA hit still reads the kpool tail ring and KDA state.
+        assert local == [[], [300], [200]]
+
+    output = SchedulerOutput.make_empty()
+    output.new_block_ids_to_zero = [100, 101, 102, 300]
+    consumer.build_connector_meta(output)
+    expected_zero = [102] if external_tokens else [100, 101, 102]
+    assert output.new_block_ids_to_zero == expected_zero
+
+
+def test_read_blocks_routes_each_hybrid_layer_to_its_group():
+    offsets: dict[str, tuple[list[int], list[int]]] = {}
+    worker = _FakeWorker(
+        mode=MoRIIOMode.READ,
+        world_size=1,
+        _has_mamba=True,
+        _num_attn_transfer_groups=2,
+        _num_mamba_transfer_groups=1,
+        _hybrid_payload_index_by_layer={"mla.0": 0, "tail.0": 1, "kda.0": 2},
+        layer_name_to_local_kv_cache_metadata={"mla.0": 0, "tail.0": 0, "kda.0": 0},
+        moriio_wrapper=SimpleNamespace(lock=threading.Lock(), shutdown=lambda: None),
+        moriio_config=SimpleNamespace(transfer_timeout=1.0),
+        _recving_transfers=defaultdict(dict),
+        _recving_transfers_start={},
+        _recving_local_blocks={},
+        _reads_issued_this_step=[],
+        _mamba_reads_this_step=[],
+        _recving_transfers_callback_addr={},
+    )
+    worker._remote_tp_rank = lambda remote_tp_size: 0
+    worker.get_engine_name_with_dp = lambda engine_id, dp_rank: engine_id
+    worker._get_built_session = lambda engine_id: ([None, None], None)
+    worker._region_session_indices = lambda layer_name: [0, 1]
+    worker._is_mamba_layer = lambda layer_name: layer_name.startswith("kda")
+
+    def _offsets(layer_name, local, remote, meta, *args):
+        offsets[layer_name] = (local, remote)
+        return [], [], []
+
+    def _mamba_reads(layer_name, sessions, regions, local, remote, *args):
+        offsets[layer_name] = (local, remote)
+        return []
+
+    worker._compute_read_attention_offsets = _offsets
+    worker._post_read_with_backoff = lambda *args: None
+    worker._post_mamba_reads = _mamba_reads
+
+    worker._read_blocks(
+        [[100, 101], [300], [200]],
+        [[10, 11], [50], [90]],
+        dst_engine_id="prefill",
+        request_id="req",
+        transfer_id="tx",
+        remote_host="127.0.0.1",
+        remote_notify_port=6001,
+        remote_tp_size=1,
+    )
+
+    assert offsets == {
+        "mla.0": ([100, 101], [10, 11]),
+        "tail.0": ([300], [50]),
+        "kda.0": ([200], [90]),
+    }
+    assert worker._recving_local_blocks["req"] == [100, 101, 300]
+
+
 def test_split_block_groups_preserves_multiple_mamba_groups():
     sched = _FakeScheduler(
         _has_mamba=True,
@@ -255,7 +459,7 @@ def test_split_block_groups_preserves_multiple_mamba_groups():
     )
 
     assert sched.split_block_groups(([1, 2], [40], [70])) == (
-        [1, 2],
+        [[1, 2]],
         [[40], [70]],
     )
 
@@ -289,8 +493,7 @@ def test_scheduler_rejects_mamba_group_without_two_states():
         moriio_connector.MoRIIOConnectorScheduler(_gate_vllm_config(), "engine", config)
 
 
-@pytest.mark.parametrize("read_mode", [True, False])
-def test_scheduler_rejects_non_dspark_speculative_hybrid(read_mode):
+def test_scheduler_rejects_non_dspark_speculative_hybrid_write():
     config = SimpleNamespace(
         kv_cache_groups=[
             SimpleNamespace(enable_kv_transfer=True, kv_cache_spec=object()),
@@ -302,7 +505,7 @@ def test_scheduler_rejects_non_dspark_speculative_hybrid(read_mode):
     with pytest.raises(moriio_common.MoRIIOError, match="speculative decoding"):
         moriio_connector.MoRIIOConnectorScheduler(
             _gate_vllm_config(
-                read_mode=read_mode, speculative_config=_spec_config("ngram")
+                read_mode=False, speculative_config=_spec_config("ngram")
             ),
             "engine",
             config,
@@ -621,6 +824,59 @@ def test_loaded_kv_cache_groups_are_every_transfer_group():
     assert connector.get_loaded_kv_cache_group_ids(SimpleNamespace()) == (0, 2)
 
 
+@pytest.mark.usefixtures("restore_moriio_role")
+def test_scheduler_rejects_circular_buffer_without_mamba():
+    full = FullAttentionSpec(
+        block_size=16, num_kv_heads=1, head_size=1, dtype=torch.float32
+    )
+    tail = CircularBufferSpec(
+        block_size=8,
+        num_kv_heads=2,
+        head_size=128,
+        head_size_v=0,
+        dtype=torch.bfloat16,
+        dcp_sharded=False,
+    )
+    config = SimpleNamespace(
+        kv_cache_groups=[
+            SimpleNamespace(enable_kv_transfer=True, kv_cache_spec=full),
+            SimpleNamespace(enable_kv_transfer=True, kv_cache_spec=tail),
+        ],
+    )
+    config.transfer_groups = tuple(config.kv_cache_groups)
+
+    with pytest.raises(NotImplementedError, match="CircularBufferSpec"):
+        moriio_connector.MoRIIOConnectorScheduler(_gate_vllm_config(), "engine", config)
+
+
+@pytest.mark.usefixtures("restore_moriio_role")
+def test_scheduler_rejects_circular_buffer_in_hybrid_write():
+    full = FullAttentionSpec(
+        block_size=16, num_kv_heads=1, head_size=1, dtype=torch.float32
+    )
+    tail = CircularBufferSpec(
+        block_size=8,
+        num_kv_heads=2,
+        head_size=128,
+        head_size_v=0,
+        dtype=torch.bfloat16,
+        dcp_sharded=False,
+    )
+    config = SimpleNamespace(
+        kv_cache_groups=[
+            SimpleNamespace(enable_kv_transfer=True, kv_cache_spec=full),
+            SimpleNamespace(enable_kv_transfer=True, kv_cache_spec=tail),
+            SimpleNamespace(enable_kv_transfer=True, kv_cache_spec=_mamba_spec()),
+        ],
+    )
+    config.transfer_groups = tuple(config.kv_cache_groups)
+
+    with pytest.raises(NotImplementedError, match="CircularBufferSpec"):
+        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)
 
@@ -640,7 +896,7 @@ def test_split_block_groups_keeps_only_running_state():
         _attn_group_ids=[0],
         _mamba_group_ids=[1],
     )
-    assert sched.split_block_groups(([1], [40, 41])) == ([1], [[41]])
+    assert sched.split_block_groups(([1], [40, 41])) == ([[1]], [[41]])
 
 
 @pytest.mark.parametrize(
@@ -661,7 +917,7 @@ def test_split_block_groups_strips_dspark_scratch_slots(blocks, expected):
         _num_ssm_scratch_blocks=2,
     )
 
-    assert sched.split_block_groups(([1], blocks)) == ([1], [expected])
+    assert sched.split_block_groups(([1], blocks)) == ([[1]], [expected])
 
 
 # --------------------------------------------------------------------------
@@ -670,12 +926,12 @@ def test_split_block_groups_strips_dspark_scratch_slots(blocks, expected):
 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)
+    def _fake_request_finished(request, attn_block_groups, mamba_block_groups):
+        seen["attn"] = [list(group) for group in attn_block_groups]
         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],
+            "remote_block_ids": [*attn_block_groups, *mamba_block_groups],
         }
 
     sched = _FakeScheduler(_has_mamba=True, _attn_group_ids=[0], _mamba_group_ids=[1])
@@ -690,7 +946,7 @@ def _fake_request_finished(request, attn_block_ids, mamba_block_groups):
     # 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["attn"] == [[10, 11]]
     assert seen["mamba"] == [[42]]
     assert params["remote_block_ids"] == [[10, 11], [42]]
     assert "remote_mamba_block_ids" not in params
@@ -1400,7 +1656,7 @@ def test_mamba_tp_ratio_rejects_heterogeneous_tp(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})
+    worker = _FakeWorker(_hybrid_payload_index_by_layer={"kda.0": 1, "kda.1": 2})
     groups = ([10], [40], [70])
 
     assert worker._mamba_blocks_for_layer("kda.0", groups) == [40]
@@ -1408,7 +1664,7 @@ def test_worker_selects_each_mamba_layers_own_block_group():
 
 
 def test_worker_rejects_missing_mamba_block_group():
-    worker = _FakeWorker(_mamba_payload_index_by_layer={"kda.0": 2})
+    worker = _FakeWorker(_hybrid_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 b95819402bd5..84aa615a41e9 100644
--- a/tests/v1/kv_connector/unit/test_moriio_kv_layout.py
+++ b/tests/v1/kv_connector/unit/test_moriio_kv_layout.py
@@ -458,7 +458,7 @@ def test_hybrid_write_posts_both_regions_before_layer_completion():
         get_engine_name_with_dp=lambda name, rank: name,
         _get_built_session=lambda name: (["conv-session", "ssm-session"], None),
         _is_mamba_layer=lambda name: True,
-        _mamba_payload_index_by_layer={"kda": 2},
+        _hybrid_payload_index_by_layer={"kda": 2},
         _region_session_indices=lambda name: [0, 1],
         _compute_mamba_transfer_offsets=Mock(
             return_value=([10, 20, 30, 40], [50, 60, 70, 80], [2, 4, 8, 16], 3)
diff --git a/tests/v1/kv_connector/unit/test_moriio_read_tp.py b/tests/v1/kv_connector/unit/test_moriio_read_tp.py
index 03891d1bbeb7..6fd2b72b6b32 100644
--- a/tests/v1/kv_connector/unit/test_moriio_read_tp.py
+++ b/tests/v1/kv_connector/unit/test_moriio_read_tp.py
@@ -154,6 +154,11 @@ def send_notify(transfer_id, host, port, message_type=None, message_fields=None)
     worker.is_producer = False
     worker.dp_rank = 0
     worker._num_mamba_transfer_groups = 1 if worker._has_mamba else 0
+    worker._num_attn_transfer_groups = (
+        len(set(worker.layer_to_group.values())) - worker._num_mamba_transfer_groups
+    )
+    # _groups() puts attention groups before mamba, as the READ payload does.
+    worker._hybrid_payload_index_by_layer = dict(worker.layer_to_group)
     worker.moriio_config = SimpleNamespace(
         backend=backend, transfer_timeout=30.0, recv_abort_timeout=600.0
     )
diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py
index 2889a8bf36b8..7bf31637f5f6 100644
--- a/tests/v1/kv_connector/unit/test_nixl_connector.py
+++ b/tests/v1/kv_connector/unit/test_nixl_connector.py
@@ -23,6 +23,7 @@
 import torch
 
 from tests.v1.attention.utils import dense_kv_cache_views
+from tests.v1.spec_decode.utils import fake_speculative_config
 from vllm import LLM
 from vllm.config import KVTransferConfig, set_current_vllm_config
 from vllm.distributed.kv_transfer.kv_connector.utils import (
@@ -605,6 +606,7 @@ def _nixl_handshake(
         self.dst_region_num_blocks[self.engine_id] = self.region_num_blocks
         self.dst_region_group_ids[self.engine_id] = self.region_group_ids
         self.dst_region_mem_types[self.engine_id] = self.region_mem_types
+        self._set_region_layers([["layer0"]])
 
         assert expected_engine_id == self.REMOTE_ENGINE_ID
 
@@ -644,6 +646,7 @@ def _nixl_handshake(
                     region_num_blocks=self.region_num_blocks,
                     region_group_ids=self.region_group_ids,
                     region_mem_types=self.region_mem_types,
+                    region_members=self.region_members,
                 ),
                 remote_tp_rank=remote_tp_rank,
                 remote_tp_size=remote_tp_size,
@@ -653,6 +656,18 @@ def _nixl_handshake(
         return remote_agents, 0.0
 
 
+def _advertise_one_layer_per_region(worker, *metas: NixlAgentMetadata) -> None:
+    """Give the worker and its peers one layer per region, as registration does,
+    so their regions pair in order."""
+    layer_names = list(worker.kv_cache_config.transfer_group_index_by_layer)
+    num_regions = len(metas[0].kv_caches_base_addr)
+    assert num_regions <= len(layer_names)
+    region_layers = [[name] for name in layer_names[:num_regions]]
+    worker._set_region_layers(region_layers)
+    for meta in metas:
+        meta.region_members = region_layers
+
+
 class TestNixlHandshake:
     @pytest.mark.parametrize(
         ("pcp_rank", "pcp_size", "dcp_size", "expected_tracked"),
@@ -1217,8 +1232,9 @@ def test_handshake_fails_on_kv_cache_layout_mismatch(
                 attn_backend_name=worker.backend_name,
                 physical_blocks_per_logical_kv_block=1,
             )
+            _advertise_one_layer_per_region(worker, meta)
 
-            with pytest.raises(RuntimeError):
+            with pytest.raises(RuntimeError, match="same kv_cache_layout"):
                 # mismatched layout is expected to fail
                 worker.add_remote_agent(meta, remote_tp_size=2)
                 worker.add_remote_agent(meta, remote_tp_size=1)
@@ -1281,6 +1297,7 @@ def test_handshake_succeed_on_kv_cache_layout_mismatch_with_experimental(
                 attn_backend_name=worker.backend_name,
                 physical_blocks_per_logical_kv_block=1,
             )
+            _advertise_one_layer_per_region(worker, meta)
 
             # We don't check layout for homogeneous TP and MLA for now, as the
             # whole block is moved.
@@ -1401,6 +1418,7 @@ def test_handshake_mixed_fa_mla_hetero_tp(self, default_vllm_config, dist_init):
                 attn_backend_name=worker.backend_name,
                 physical_blocks_per_logical_kv_block=1,
             )
+            _advertise_one_layer_per_region(worker, meta)
             worker.add_remote_agent(meta, remote_tp_size=1)
             assert (
                 FakeNixlConnectorWorker.REMOTE_ENGINE_ID in worker.dst_xfer_side_handles
@@ -1428,6 +1446,7 @@ def test_handshake_mixed_fa_mla_hetero_tp(self, default_vllm_config, dist_init):
                 attn_backend_name=worker2.backend_name,
                 physical_blocks_per_logical_kv_block=1,
             )
+            _advertise_one_layer_per_region(worker2, bad_meta)
             with pytest.raises(AssertionError):
                 worker2.add_remote_agent(bad_meta, remote_tp_size=1)
 
@@ -1525,6 +1544,7 @@ def _mla_target_with_gqa_draft(tp_size: int, tp_rank: int, remote_tp_size: int):
             attn_backend_name=worker.backend_name,
             physical_blocks_per_logical_kv_block=1,
         )
+        _advertise_one_layer_per_region(worker, meta)
         return worker, meta, mla_len, draft_len
 
     @patch(
@@ -1724,6 +1744,7 @@ def test_handshake_validates_gqa_replicated_block_len(
                 attn_backend_name=worker.backend_name,
                 physical_blocks_per_logical_kv_block=1,
             )
+            _advertise_one_layer_per_region(worker, meta)
 
             # Must validate cleanly (used to raise AssertionError).
             worker.add_remote_agent(meta, remote_tp_size=8)
@@ -1783,6 +1804,7 @@ def test_handshake_rejects_wrong_block_len_without_gqa_replication(
                 attn_backend_name=worker.backend_name,
                 physical_blocks_per_logical_kv_block=1,
             )
+            _advertise_one_layer_per_region(worker, bad_meta)
 
             with pytest.raises(AssertionError):
                 worker.add_remote_agent(bad_meta, remote_tp_size=2)
@@ -1869,6 +1891,7 @@ def test_handshake_rejects_nvfp4_kv_head_split(
             attn_backend_name=worker.backend_name,
             physical_blocks_per_logical_kv_block=1,
         )
+        _advertise_one_layer_per_region(worker, meta)
 
         if rejected:
             with pytest.raises(RuntimeError, match="NVFP4"):
@@ -4564,14 +4587,13 @@ def _set_test_speculative_config(
         code_revision=code_revision,
         hf_config=SimpleNamespace(eagle_aux_hidden_state_layer_ids=auxiliary_layer_ids),
     )
-    vllm_config.speculative_config = SimpleNamespace(
+    vllm_config.speculative_config = fake_speculative_config(
         method=method,
         draft_model_config=draft_model_config,
         num_speculative_tokens=num_speculative_tokens,
         parallel_drafting=parallel_drafting,
         kv_cache_dtype=kv_cache_dtype,
         attention_backend=attention_backend,
-        use_eagle=lambda: True,
     )
 
 
@@ -5132,3 +5154,71 @@ def test_explicit_kv_role_no_deprecation_warning(default_vllm_config, dist_init)
             mock_logger.warning_once.assert_not_called(),
             (f"kv_role={role!r} should not emit deprecation warning"),
         )
+
+
+class _BackendRecordingNixlWrapper(FakeNixlWrapper):
+    def __init__(self, agent_name: str, config=None, *args, **kwargs):
+        super().__init__(agent_name, *args, **kwargs)
+        self.config = config
+        self.created_backends: list[tuple[str, dict[str, str]]] = []
+
+    def create_backend(self, backend: str, init_params: dict[str, str]) -> None:
+        self.created_backends.append((backend, dict(init_params)))
+
+
+@pytest.mark.parametrize(
+    "extra_config, agent_backends, created_backends",
+    [
+        # No UCX params: the agent config creates the backends, as before.
+        ({}, ["UCX"], []),
+        (
+            {"ucx_backend_params": {"ucx_error_handling_mode": "none"}},
+            [],
+            [("UCX", {"num_threads": "4", "ucx_error_handling_mode": "none"})],
+        ),
+        (
+            {
+                "backends": ["UCX", "LIBFABRIC"],
+                "ucx_backend_params": {"ucx_error_handling_mode": "none"},
+            },
+            [],
+            [("UCX", {"ucx_error_handling_mode": "none"}), ("LIBFABRIC", {})],
+        ),
+        # UCX not selected: the params have nothing to apply to.
+        (
+            {
+                "backends": ["LIBFABRIC"],
+                "ucx_backend_params": {"ucx_error_handling_mode": "none"},
+            },
+            ["LIBFABRIC"],
+            [],
+        ),
+    ],
+)
+def test_ucx_backend_params(dist_init, extra_config, agent_backends, created_backends):
+    """kv_connector_extra_config.ucx_backend_params reaches UCX's create_backend."""
+
+    def agent_config(backends=None, num_threads=0, capture_telemetry=False):
+        return SimpleNamespace(
+            backends=["UCX"] if backends is None else backends,
+            num_threads=num_threads,
+        )
+
+    vllm_config = create_vllm_config(kv_connector_extra_config=extra_config)
+    with (
+        patch(
+            "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
+            _BackendRecordingNixlWrapper,
+        ),
+        patch(
+            "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.nixl_agent_config",
+            agent_config,
+        ),
+        set_current_vllm_config(vllm_config),
+    ):
+        connector = NixlConnector(
+            vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16)
+        )
+    wrapper = connector.connector_worker.nixl_wrapper
+    assert wrapper.config.backends == agent_backends
+    assert wrapper.created_backends == created_backends
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 f7051cb5234f..979f4255ef67 100644
--- a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py
+++ b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py
@@ -43,6 +43,7 @@
     make_kv_cache_config,
     make_nixl_scheduler,
     maybe_update_block_size,
+    region_pairs_by_index,
 )
 
 
@@ -61,7 +62,6 @@ def region_pull_worker():
     worker._engine_last_active = {}
     worker._handshake_lock = threading.RLock()
     worker._handshake_futures = {}
-    worker._transfer_layer_group_ids = ()
     worker._bidirectional_kv_xfer_enabled = False
     worker._recving_transfers = {}
     worker._failed_remote_engines = set()
@@ -85,6 +85,9 @@ def region_pull_worker():
     worker._group_spec_types = (MLAAttentionSpec, MLAAttentionSpec)
     worker.region_group_ids = [0, 1]
     worker.dst_region_group_ids = {"P": [0, 0]}
+    worker._dst_pairs = {
+        "P": region_pairs_by_index(2, groups=[0, 1], groups_match=False)
+    }
     worker.num_regions = 2
     worker.block_len_per_layer = [1024, 1024]
     worker.transfer_topo = MagicMock()
@@ -97,7 +100,6 @@ def region_pull_worker():
         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}
@@ -193,6 +195,11 @@ def test_dcp_region_pull(region_pull_worker, num_pages, region_groups):
     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_pairs = {
+        "P": region_pairs_by_index(
+            worker.num_regions, groups=region_groups, groups_match=False
+        )
+    }
     worker.dst_region_num_blocks = {
         engine: [100] * worker.num_regions for engine in ("P", "D")
     }
@@ -718,6 +725,11 @@ 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._dst_pairs = {
+        remote_engine_id: region_pairs_by_index(
+            len(resolved_types), groups=worker.region_group_ids
+        )
+    }
 
     worker.transfer_topo = MagicMock()
     # tp_ratio not exercised (all_source_ranks is empty so no reads run),
@@ -1161,8 +1173,6 @@ 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]
@@ -1191,16 +1201,14 @@ def _make_mock_worker_for_desc_ids(
 
 
 @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."""
+def test_shared_nixl_region_reads_once_per_paired_group():
+    """A region shared by three groups is read once at each group's blocks."""
     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]),
@@ -1208,6 +1216,7 @@ def test_shared_nixl_region_broadcasts_every_group_block_table():
         block_size_ratio=None,
         physical_blocks_per_logical=1,
         region_num_blocks=[10, 10, 10],
+        pairs=([0, 1, 1, 1, 2], [0, 0, 1, 2, 2]),
     )
 
     assert result.tolist() == [1, 11, 12, 13, 23]
@@ -1249,6 +1258,8 @@ def test_get_block_descs_ids_hybrid_ssm():
         dst_num_blocks=100,
         block_size_ratio=None,
         physical_blocks_per_logical=1,
+        # FA group in both regions, SSM group in SSM region 0.
+        pairs=([0, 1, 0], [0, 0, 1]),
     )
 
     expected = [3, 5, 103, 105, 201, 202, 301, 302, 401, 402, 501, 502]
@@ -1266,14 +1277,12 @@ def test_get_block_descs_ids_selects_attention_regions_by_group():
         group_spec_types=(FullAttentionSpec, FullAttentionSpec),
         block_len_per_layer=[100, 100, 100],
     )
-    worker.region_group_ids = [0, 0, 1]
-    worker._uses_region_group_mapping = True
-
     result = worker._compute_desc_ids(
         block_ids=([1, 2], [7]),
         dst_num_blocks=10,
         block_size_ratio=None,
         physical_blocks_per_logical=1,
+        pairs=([0, 1, 2], [0, 0, 1]),
     )
 
     assert result.tolist() == [1, 2, 11, 12, 27]
@@ -1283,7 +1292,7 @@ def test_get_block_descs_ids_selects_attention_regions_by_group():
         dst_num_blocks=10,
         block_size_ratio=None,
         physical_blocks_per_logical=1,
-        region_group_ids=[1, 1, 0],
+        pairs=([0, 1, 2], [1, 1, 0]),
     )
     assert remapped.tolist() == [21, 22, 7, 17]
 
@@ -1299,15 +1308,13 @@ def test_get_block_descs_ids_uses_per_region_pool_capacity():
         group_spec_types=(FullAttentionSpec, FullAttentionSpec),
         block_len_per_layer=[100, 100],
     )
-    worker.region_group_ids = [0, 1]
-    worker._uses_region_group_mapping = True
-
     result = worker._compute_desc_ids(
         block_ids=([4], [8]),
         dst_num_blocks=10,
         block_size_ratio=None,
         physical_blocks_per_logical=1,
         region_num_blocks=[5, 10],
+        pairs=([0, 1], [0, 1]),
     )
 
     assert result.tolist() == [4, 13]
@@ -1337,6 +1344,8 @@ def test_get_block_descs_ids_kernel_block_mismatch():
         dst_num_blocks=num_blocks,
         block_size_ratio=None,
         physical_blocks_per_logical=ratio,
+        # FA group in both regions, SSM group in SSM region 0.
+        pairs=([0, 1, 0], [0, 0, 1]),
     )
 
     expected = [3, 7, 403, 407, 801, 802, 901, 902, 1001, 1002, 1101, 1102]
@@ -1366,6 +1375,8 @@ def test_get_block_descs_ids_hetero_block_size_hybrid():
         dst_num_blocks=100,
         block_size_ratio=ratio,
         physical_blocks_per_logical=1,
+        # FA group in both regions, SSM group in SSM region 0.
+        pairs=([0, 1, 0], [0, 0, 1]),
     )
 
     # FA regions have 100*4 entries each; SSM regions (4 per layer) start at
@@ -2371,6 +2382,7 @@ def test_push_write_hybrid_mla_replicates_attention():
     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._dst_pairs = {engine_id: region_pairs_by_index(2, groups=[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 f0d5ea1ae9bf..10f0aea20e18 100644
--- a/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py
+++ b/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py
@@ -11,6 +11,7 @@
 mid-decode (silent corruption of an unrelated request).
 """
 
+import contextlib
 import math
 from collections import defaultdict
 from threading import Event, Lock
@@ -32,12 +33,17 @@
 from vllm.distributed.kv_transfer.kv_connector.v1.nixl.push_worker import (
     NixlPushConnectorWorker,
 )
+from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
 from vllm.v1.kv_cache_interface import (
+    FullAttentionSpec,
     KVCacheConfig,
     KVCacheGroupSpec,
     KVCacheLayout,
     KVCacheTensor,
+    MambaSpec,
     MLAAttentionSpec,
+    SlidingWindowSpec,
+    SparseCacheRole,
     create_kv_cache_views,
 )
 
@@ -201,7 +207,7 @@ def test_packed_mla_pp1_transfers_whole_rows(num_layers, worker_cls):
     assert worker._registered_descs == [[(raw.data_ptr(), raw.nbytes, 0, "")]]
     # One packed region -> one run of 4 whole rows.
     assert worker.src_blocks_data.tolist() == [[raw.data_ptr(), stride, 0, stride, 4]]
-    assert worker._transfer_layer_group_ids == ()
+    assert worker.region_members == [list(layouts)]
     metadata = msgspec.msgpack.decode(
         worker.xfer_handshake_metadata.agent_metadata_bytes, type=NixlAgentMetadata
     )
@@ -216,12 +222,12 @@ def test_per_layer_peer_addresses_pages_inside_packed_rows():
         remote.xfer_handshake_metadata.agent_metadata_bytes, type=NixlAgentMetadata
     )
     assert metadata.block_lens == [192]
-    local = SimpleNamespace(region_names=["L0"])
-    bw.NixlBaseConnectorWorker._split_remote_packed_rows(local, metadata)
+    bw.NixlBaseConnectorWorker._expand_remote_packed_rows(metadata)
     assert metadata.kv_caches_base_addr == [raw.data_ptr() + 64]
     assert metadata.block_lens == [128]
     assert metadata.block_strides == [192]
     assert metadata.region_names == ["L0"]
+    assert metadata.region_members == [["L0"]]
     assert metadata.packed_member_layouts == {}
 
 
@@ -280,21 +286,8 @@ def test_packed_mla_pp_pairs_asymmetric_strides_and_overlapping_layers(
         num_blocks=d_num_blocks,
     )
     assert producer.block_len_per_layer == [128, 64, 64]
-    assert producer._transfer_layer_group_ids == (0, 1, 2)
+    assert producer.region_members == [["L2"], ["L2.swa"], ["L3"]]
     assert producer.num_descs == 3 * p_num_blocks
-    local_ids = producer._compute_desc_ids(
-        [[1], [2], [3]],
-        producer.num_blocks,
-        None,
-        1,
-        region_num_blocks=producer.dst_region_num_blocks[producer.engine_id],
-    )
-    local_descs = expand_strided_descs(producer.src_blocks_data)
-    assert local_descs[local_ids].tolist() == [
-        [p_raw.data_ptr() + 192, 128, 0],
-        [p_raw.data_ptr() + 2 * 192, 64, 0],
-        [p_raw.data_ptr() + 128 + 3 * 192, 64, 0],
-    ]
     # Each decoder TP rank receives the complete MLA page, without head slicing.
     for rank in (0, 1):
         metadata = msgspec.msgpack.decode(
@@ -304,12 +297,28 @@ def test_packed_mla_pp_pairs_asymmetric_strides_and_overlapping_layers(
         producer.add_remote_agent(metadata, remote_tp_rank=rank, remote_tp_size=2)
         assert metadata.region_num_blocks == [d_num_blocks] * 3
         assert metadata.region_names == ["L2", "L2.swa", "L3"]
+        pairs = producer._dst_pairs[consumer.engine_id]
+        local_ids = producer._compute_desc_ids(
+            [[1], [2], [3]],
+            producer.num_blocks,
+            None,
+            1,
+            region_num_blocks=producer.dst_region_num_blocks[producer.engine_id],
+            pairs=(pairs.local, pairs.group),
+        )
+        local_descs = expand_strided_descs(producer.src_blocks_data)
+        assert local_descs[local_ids].tolist() == [
+            [p_raw.data_ptr() + 192, 128, 0],
+            [p_raw.data_ptr() + 2 * 192, 64, 0],
+            [p_raw.data_ptr() + 128 + 3 * 192, 64, 0],
+        ]
         remote_ids = producer._compute_desc_ids(
             [[3], [1], [2]],
             metadata.num_blocks,
             None,
             1,
             region_num_blocks=producer.dst_region_num_blocks[consumer.engine_id],
+            pairs=(pairs.remote, pairs.group),
         )
         handle = producer.dst_xfer_side_handles[consumer.engine_id][rank]
         remote_descs = producer.nixl_wrapper.dlists[handle]
@@ -319,9 +328,6 @@ def test_packed_mla_pp_pairs_asymmetric_strides_and_overlapping_layers(
             [d_raw.data_ptr() + 256, 64, 0],
             [d_raw.data_ptr() + 64 + 2 * 256, 64, 0],
         ]
-        aligned = msgspec.msgpack.encode(metadata)
-        producer._align_remote_regions_by_layer(metadata)
-        assert msgspec.msgpack.encode(metadata) == aligned
 
 
 @pytest.mark.cpu_test
@@ -356,7 +362,6 @@ 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)
 
@@ -435,7 +440,7 @@ def _register_overlaid_mla_worker(
     worker.dst_num_blocks = {}
     worker.dst_region_num_blocks = {}
     worker.dst_region_group_ids = {}
-    worker.dst_uses_region_group_mapping = {}
+    worker._dst_pairs = {}
     worker.dst_region_mem_types = {}
     worker._desc_is_dram_by_block_size = {}
     worker._desc_pos_by_block_size = {}
@@ -452,9 +457,6 @@ def _register_overlaid_mla_worker(
     worker._mixed_mem_types = False
     worker.region_names = []
     worker.region_num_blocks = []
-    worker._transfer_layer_names = ()
-    worker._transfer_layer_region_indices = ()
-    worker._transfer_layer_group_ids = ()
     worker._region_is_mla = []
     worker._head_sharded_draft_kv_heads = None
     worker.block_len_per_layer = []
@@ -522,14 +524,11 @@ def test_overlaid_transfer_groups_share_region_geometry(push_pp):
     expected_addrs = [
         allocation.data_ptr() + block * block_stride for block in range(num_blocks)
     ]
-    num_desc_regions = 2 if push_pp else 1
+    # One descriptor list per region, whichever layers share it.
     local_descs = expand_strided_descs(worker.src_blocks_data)
-    assert local_descs[:, 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
-    )
+    assert local_descs[:, 0].tolist() == expected_addrs
+    assert worker.num_descs == num_blocks
+    assert worker.dst_region_num_blocks[worker.engine_id] == [num_blocks]
 
     metadata = msgspec.msgpack.decode(
         worker.xfer_handshake_metadata.agent_metadata_bytes,
@@ -537,7 +536,7 @@ def test_overlaid_transfer_groups_share_region_geometry(push_pp):
     )
     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 metadata.region_members == [["layer.0", "layer.1"]]
     assert worker._block_ids_by_region(([0], [2]), worker.region_group_ids) == [[0, 2]]
 
 
@@ -704,6 +703,20 @@ def _make_mla_hybrid_worker(
     return worker
 
 
+def _own_region_pairs(worker) -> bw._RegionPairs:
+    """How a worker pairs its regions with an identical peer's."""
+    metadata = msgspec.msgpack.decode(
+        worker.xfer_handshake_metadata.agent_metadata_bytes, type=NixlAgentMetadata
+    )
+    return worker._pair_regions(metadata)
+
+
+def _own_pairs(worker) -> tuple[list[int], list[int]]:
+    """The (regions, groups) a worker reads when paired with an identical peer."""
+    pairs = _own_region_pairs(worker)
+    return pairs.local, pairs.group
+
+
 @pytest.mark.cpu_test
 @pytest.mark.parametrize("kda_only_tensor", [False, True])
 def test_mla_hybrid_groups_read_only_their_own_regions(kda_only_tensor):
@@ -724,6 +737,7 @@ def test_mla_hybrid_groups_read_only_their_own_regions(kda_only_tensor):
         dst_num_blocks=4,
         block_size_ratio=None,
         physical_blocks_per_logical=1,
+        pairs=_own_pairs(worker),
     )
 
     # 4 blocks per region; each SSM layer has 3 conv + 1 temporal regions.
@@ -782,12 +796,52 @@ def test_nixl_keeps_packed_sparse_mla_manager_block():
     assert worker._physical_blocks_per_logical_kv_block == 1
 
 
+@pytest.mark.parametrize(
+    "layout,kernel_block_size,raises",
+    [
+        (KVCacheLayout.LBHNC, 384, False),
+        (KVCacheLayout.LBHNC, 64, True),
+        (KVCacheLayout.BLHNC, 64, False),
+    ],
+)
+def test_nixl_compressed_kernel_split(layout, kernel_block_size, raises):
+    """NIXL keeps compressed pages whole or rejects a layer-outer split."""
+    from vllm.config import set_current_vllm_config
+    from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
+        NixlConnectorWorker,
+    )
+
+    spec = MLAAttentionSpec(
+        384, num_kv_heads=1, head_size=128, dtype=torch.uint8, tokens_per_state=8
+    )
+    kv_cache_config = KVCacheConfig(2, [], [KVCacheGroupSpec(["summary"], spec)])
+    vllm_config = create_vllm_config(block_size=384)
+    vllm_config.cache_config.kv_cache_layout = layout.name
+    vllm_config.kv_transfer_config.kv_buffer_device = "cpu"
+    backend = MagicMock()
+    backend.get_name.return_value = "TEST"
+    ctx = pytest.raises(NotImplementedError) if raises else contextlib.nullcontext()
+    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=[backend]),
+        patch.object(bw, "select_common_block_size", return_value=kernel_block_size),
+        set_current_vllm_config(vllm_config),
+        ctx,
+    ):
+        worker = NixlConnectorWorker(vllm_config, "local-engine", kv_cache_config)
+        worker.shutdown()
+
+
 def _make_remote_meta(
     worker,
     remote_block_size,
     remote_kernel_block_size,
     remote_num_logical,
     remote_ssm_sizes,
+    region_group_ids=None,
+    layer_group_ids=None,
 ):
     from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
         NixlAgentMetadata,
@@ -799,20 +853,30 @@ def _make_remote_meta(
     kernel_page = (
         worker.block_len_per_layer[0] * remote_kernel_block_size // worker.block_size
     )
+    region_members = worker.region_members
+    num_regions = len(region_members)
+    if layer_group_ids is None:
+        group_by_layer = worker.kv_cache_config.transfer_group_index_by_layer
+        layer_group_ids = {
+            name: group_by_layer[name] for names in region_members for name in names
+        }
     return NixlAgentMetadata(
         engine_id="remote-engine",
         agent_metadata=b"remote-agent-meta",
         device_id=0,
-        kv_caches_base_addr=[0x10_000_000, 0x20_000_000],
+        kv_caches_base_addr=[0x10_000_000 * (i + 1) for i in range(num_regions)],
         num_blocks=remote_num_logical * remote_ppl,
-        block_lens=[kernel_page, kernel_page],
+        block_lens=[kernel_page] * num_regions,
         # Non-interleaved remote: consecutive blocks abut, so stride == page length.
-        block_strides=[kernel_page, kernel_page],
+        block_strides=[kernel_page] * num_regions,
         kv_cache_layout=worker.kv_cache_layout,
         block_size=remote_kernel_block_size,
         ssm_sizes=remote_ssm_sizes,
         attn_backend_name=worker.backend_name,
         physical_blocks_per_logical_kv_block=remote_ppl,
+        region_members=region_members,
+        region_group_ids=region_group_ids,
+        layer_group_ids=layer_group_ids,
     )
 
 
@@ -920,6 +984,36 @@ def test_hetero_ppl_multi_read_writes_stay_within_request_blocks():
     assert total > 0
 
 
+@pytest.mark.cpu_test
+def test_mla_hybrid_peer_with_other_cache_groups_is_refused():
+    """A hybrid layer in another cache group on P would be read at the wrong
+    group's blocks, and region-mapped reads cannot serve SSM state, so the
+    handshake refuses the peer before any transfer."""
+    worker = _make_mla_hybrid_worker(
+        local_block_size=16,
+        kernel_block_size=16,
+        num_logical_blocks=4,
+        kda_only_tensor=True,
+    )
+    layer_groups = dict(worker.kv_cache_config.transfer_group_index_by_layer)
+    remote_groups = list(worker.region_group_ids)
+    # The KDA-only region holds kda_a.2 alone, in group 1; P puts it in 2.
+    remote_groups[remote_groups.index(1)] = 2
+    layer_groups["kda_a.2"] = 2
+    meta_r = _make_remote_meta(
+        worker,
+        remote_block_size=16,
+        remote_kernel_block_size=16,
+        remote_num_logical=4,
+        remote_ssm_sizes=worker._mamba_ssm_size,
+        region_group_ids=remote_groups,
+        layer_group_ids=layer_groups,
+    )
+
+    with pytest.raises(NotImplementedError, match="cache groups"):
+        _register_remote_agents(worker, meta_r, 1)
+
+
 def _resolve(
     desc_arr,
     idx,
@@ -1462,6 +1556,7 @@ def test_csa_linear_scratch_descs_follow_the_pages_they_overlay():
         dst_num_blocks=2,
         block_size_ratio=None,
         physical_blocks_per_logical=1,
+        pairs=_own_pairs(worker),
     )
     # Scratch group (block 0) now maps to regions 0 and 1 -> descs 0 and 2,
     # where the compressed-page placement gave 4 and 6.
@@ -1477,6 +1572,7 @@ def test_csa_linear_selects_compressed_sparse_compressor_state_and_mamba_regions
         dst_num_blocks=2,
         block_size_ratio=None,
         physical_blocks_per_logical=1,
+        pairs=_own_pairs(worker),
     )
     assert desc_ids.tolist() == [
         1,
@@ -1497,6 +1593,7 @@ def test_csa_linear_selects_compressed_sparse_compressor_state_and_mamba_regions
         dst_num_blocks=2,
         block_size_ratio=None,
         physical_blocks_per_logical=1,
+        pairs=_own_pairs(worker),
     )
     assert mamba_only.tolist() == [9, 11, 13, 15, 17]
 
@@ -1547,6 +1644,9 @@ def test_csa_linear_remote_ple_is_copied_whole():
     )
 
     worker = _make_csa_linear_ple_worker()
+    # Every region pairs, so the peer's SSM and PLE regions keep their indices.
+    pairs = _own_region_pairs(worker)
+    assert pairs.remote_regions == list(range(worker.num_regions))
     metadata = NixlAgentMetadata(
         engine_id="remote",
         agent_metadata=b"agent",
@@ -1571,6 +1671,7 @@ def test_csa_linear_remote_ple_is_copied_whole():
         metadata,
         tp_ratio=-2,
         transfer_info=SimpleNamespace(remote_physical_blocks_per_logical=1),
+        pairs=pairs,
     )
     assert descriptors.shape == (5, 5)
     # PLE run: 2 whole 256-byte pages at 0x10000, 0x10100.
@@ -1582,22 +1683,25 @@ def test_csa_linear_remote_ple_is_copied_whole():
             metadata,
             tp_ratio=-2,
             transfer_info=SimpleNamespace(remote_physical_blocks_per_logical=1),
+            pairs=pairs,
         )
 
 
 @pytest.mark.cpu_test
 def test_csa_linear_ple_page_wider_than_the_shared_region_page():
-    """The PLE page is discovered in the region of the main KV page registered
-    first (both start at byte 0 of the block); its descriptors must still cover
-    the whole PLE page, not the region's block_len."""
+    """The PLE page shares its base address with the main KV page (both start
+    at byte 0 of the block) but has a wider page, so it registers as its own
+    region; its descriptors must still cover the whole PLE page."""
     worker = _make_csa_linear_ple_worker(ple_page_size=384)
 
-    assert worker._ple_region_index == 0
+    assert worker._ple_region_index == 1
     assert worker.block_len_per_layer[0] == 256
     assert worker.block_stride_per_layer[0] == 384
+    assert worker.block_len_per_layer[1] == 384
     assert worker._ple_block_len == 384
 
-    bases = [0x10000, 0x20000, 0x30000, 0x40000]
+    # Region 1 (the PLE page) aliases region 0's base with a wider block_len.
+    bases = [0x10000, 0x10000, 0x20000, 0x30000, 0x40000]
     mamba = worker._build_mamba_local(bases)
     # PLE run: 2 whole 384-byte pages at 0x10000, 0x10000 + 384.
     assert mamba[-1].tolist() == [0x10000, 384, 0, 384, 2]
@@ -1651,7 +1755,9 @@ def group(specs):
         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.
+    # address only that region while the paged group spans both. The rings'
+    # distinct page geometry keeps them in their own transfer region rather
+    # than folding into the paged region they alias.
     tensor_regions = (("paged.0",), ("paged.1", "ring.0", "ring.1"))
     region_size, page_size = 512, 256
     kv_cache_config = KVCacheConfig(
@@ -1706,13 +1812,518 @@ 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._group_region_indices[1] == [1]
-
+    assert worker._group_region_indices[1] == [2]
+
+    group_by_layer = worker.kv_cache_config.transfer_group_index_by_layer
+    regions, groups = zip(
+        *dict.fromkeys(
+            (region, group_by_layer[name])
+            for region, names in enumerate(worker.region_members)
+            for name in names
+        )
+    )
     desc_ids = worker._compute_desc_ids(
         block_ids=([1], [0]),
         dst_num_blocks=2,
         block_size_ratio=None,
         physical_blocks_per_logical=1,
+        pairs=(regions, groups),
+    )
+    # Paged block 1 in both paged regions; ring block 0 only in its own
+    # scratch region, which no longer folds into the overlaid paged region.
+    assert sorted(desc_ids.tolist()) == sorted([0 * 2 + 1, 1 * 2 + 1, 2 * 2 + 0])
+
+
+def _make_shared_base_ring_worker():
+    """MLA swa group plus a compressor ring group aliasing the same
+    block-major allocation with a larger page: the DeepSeek-V4.1-Flash
+    HMA shape where a compressor state cache shares its storage base with
+    the swa cache it serves. The ring page is wider than the swa page, so
+    folding it into the swa region truncates every transferred block."""
+    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,
+    )
+
+    num_blocks = 2
+    # Every 256-byte block holds one 128-byte MLA swa page (4 tokens x 32B)
+    # and a 256-byte compressor ring page (8 tokens x 32B) aliased at the
+    # same base address.
+    block_stride, swa_page, ring_page = 256, 128, 256
+    raw = torch.zeros(num_blocks * block_stride, dtype=torch.uint8)
+
+    swa_spec = MLAAttentionSpec(
+        block_size=4, num_kv_heads=1, head_size=32, dtype=torch.uint8
+    )
+    ring_spec = CircularBufferSpec(
+        block_size=8,
+        num_kv_heads=1,
+        head_size=32,
+        head_size_v=0,
+        dtype=torch.uint8,
+    )
+
+    def group(specs):
+        uniform = UniformTypeKVCacheSpecs.from_specs(specs)
+        assert uniform is not None
+        return KVCacheGroupSpec(list(specs), uniform)
+
+    swa_name, ring_name = "layers.0.swa", "layers.2.compressor"
+    kv_cache_config = KVCacheConfig(
+        num_blocks=num_blocks,
+        kv_cache_tensors=[
+            KVCacheTensor(
+                size=raw.nbytes,
+                layers=[name],
+                layer_stride=raw.nbytes,
+                block_stride=block_stride,
+            )
+            for name in (swa_name, ring_name)
+        ],
+        kv_cache_groups=[
+            group({swa_name: swa_spec}),
+            group({ring_name: ring_spec}),
+        ],
+    )
+
+    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)
+        swa_cache = raw.as_strided((num_blocks, swa_page), (block_stride, 1))
+        ring_cache = raw.as_strided(
+            (num_blocks, 1, ring_page), (block_stride, block_stride, 1)
+        )
+        worker.register_kv_caches({swa_name: swa_cache, ring_name: ring_cache})
+    return worker
+
+
+@pytest.mark.cpu_test
+def test_shared_base_regions_with_distinct_block_lens_stay_separate():
+    """A compressor ring aliasing an MLA region's base address must register
+    its own region instead of collapsing into the MLA region and inheriting
+    its block length, which truncated and misaligned its transferred state."""
+    worker = _make_shared_base_ring_worker()
+
+    base = worker.kv_caches_base_addr[worker.engine_id][0]
+    assert base[0] == base[1]
+    assert worker.num_regions == 2
+    assert worker.block_len_per_layer == [128, 256]
+    assert worker.block_stride_per_layer == [256, 256]
+    assert worker._region_is_mla == [True, False]
+    assert worker._group_region_indices[1] == [1]
+    assert worker.region_group_ids == [0, 1]
+    assert worker.num_descs == 4
+    strided_descs = worker.src_blocks_data
+    assert strided_descs.shape == (2, 5)
+    assert strided_descs[:, 1].tolist() == [128, 256]
+    assert strided_descs[:, 3].tolist() == [256, 256]
+    assert strided_descs[:, 4].tolist() == [2, 2]
+    descs = expand_strided_descs(strided_descs)
+    assert descs[:, 1].tolist() == [128, 128, 256, 256]
+    np.testing.assert_array_equal(
+        descs[:, 0], [base[0], base[0] + 256, base[0], base[0] + 256]
+    )
+
+
+# Real model KV layouts, end to end on CPU: vLLM's planner and allocator build
+# P's and D's caches, D reads one request from P, and the recorded descriptors
+# are run as memory copies so every byte can be checked. Each case is the
+# first layers of a model with that model's per-layer KV specs. Hybrid block
+# sizes are the ones vLLM aligns to on CPU, where the attention page must hold
+# the Mamba state.
+_BF16 = torch.bfloat16
+
+
+def _attn(heads, head, head_v=None, window=None):
+    def make(block_size, _mamba_page):
+        kwargs = dict(
+            block_size=block_size,
+            num_kv_heads=heads,
+            head_size=head,
+            head_size_v=head_v or head,
+            dtype=_BF16,
+        )
+        if window is not None:
+            return SlidingWindowSpec(sliding_window=window, **kwargs)
+        return FullAttentionSpec(**kwargs)
+
+    return make
+
+
+def _mla(head=576, dtype=_BF16, **kwargs):
+    def make(block_size, _mamba_page):
+        return MLAAttentionSpec(
+            block_size=block_size,
+            num_kv_heads=1,
+            head_size=head,
+            dtype=dtype,
+            **kwargs,
+        )
+
+    return make
+
+
+def _mamba(shapes, mamba_type):
+    def make(block_size, mamba_page):
+        return MambaSpec(
+            block_size=block_size,
+            shapes=shapes,
+            dtypes=(_BF16, torch.float32),
+            page_size_padded=mamba_page,
+            mamba_type=mamba_type,
+        )
+
+    return make
+
+
+def _layers(n, make):
+    return [entry for i in range(n) for entry in make(i)]
+
+
+_GDN = MambaAttentionBackendEnum.GDN_ATTN
+# (block size, [(layer name, spec factory)]) per model.
+_REAL_MODELS = {
+    # GLM-4.5-Air and MiniMax-M2 have the same attention shape.
+    "Qwen3-8B": (
+        128,
+        _layers(4, lambda i: [(f"model.layers.{i}.self_attn.attn", _attn(8, 128))]),
+    ),
+    "gpt-oss-20b": (
+        128,
+        _layers(
+            4,
+            lambda i: [
+                (
+                    f"model.layers.{i}.attn.attn",
+                    _attn(8, 64, window=None if i % 2 else 128),
+                )
+            ],
+        ),
+    ),
+    "gemma-3-1b": (
+        128,
+        _layers(
+            12,
+            lambda i: [
+                (
+                    f"model.layers.{i}.self_attn.attn",
+                    _attn(1, 256, window=None if i % 6 == 5 else 512),
+                )
+            ],
+        ),
+    ),
+    "MiMo-V2-Flash": (
+        128,
+        _layers(
+            12,
+            lambda i: [
+                (
+                    f"model.layers.{i}.self_attn.attn",
+                    _attn(4, 192, 128)
+                    if i in (0, 5, 11)
+                    else _attn(8, 192, 128, window=128),
+                )
+            ],
+        ),
+    ),
+    # Kimi-K2 has the same MLA cache.
+    "DeepSeek-V3": (
+        128,
+        _layers(4, lambda i: [(f"model.layers.{i}.self_attn.attn", _mla())]),
+    ),
+    "DeepSeek-V3.2": (
+        64,
+        _layers(
+            4,
+            lambda i: [
+                (
+                    f"model.layers.{i}.self_attn.attn",
+                    _mla(is_index_group_leader=True),
+                ),
+                (
+                    f"model.layers.{i}.self_attn.indexer.k_cache",
+                    _mla(132, torch.uint8, cache_role=SparseCacheRole.INDEXER),
+                ),
+            ],
+        ),
+    ),
+    "Qwen3-Next-80B-A3B": (
+        1152,
+        _layers(
+            8,
+            lambda i: [
+                (f"model.layers.{i}.self_attn.attn", _attn(2, 256))
+                if i % 4 == 3
+                else (
+                    f"model.layers.{i}.linear_attn",
+                    _mamba(((8192, 3), (32, 128, 128)), _GDN),
+                )
+            ],
+        ),
+    ),
+    "NVIDIA-Nemotron-Nano-9B-v2": (
+        1408,
+        _layers(
+            22,
+            lambda i: (
+                [
+                    (
+                        f"model.layers.{i}.mixer",
+                        _mamba(
+                            ((12288, 3), (128, 80, 128)),
+                            MambaAttentionBackendEnum.MAMBA2,
+                        ),
+                    )
+                ]
+                if i in (0, 2, 4, 6, 7, 9, 11, 13, 16, 18, 20)
+                else [(f"model.layers.{i}.mixer.attn", _attn(8, 128))]
+                if i in (14, 21)
+                else []
+            ),
+        ),
+    ),
+    # KDA at 1-based layers 1-3, 5-7; MLA at 4 and 8. A 1920-token MLA page
+    # (64-token kernel blocks) is the first to hold the 2,170,880 B KDA state.
+    "Kimi-Linear-48B-A3B": (
+        1920,
+        _layers(
+            8,
+            lambda i: [
+                (f"model.layers.{i}.self_attn.attn", _mla())
+                if i % 4 == 3
+                else (
+                    f"model.layers.{i}.self_attn",
+                    _mamba(((12288, 3), (32, 128, 128)), _GDN),
+                )
+            ],
+        ),
+    ),
+}
+_LAYOUT_PREFERENCE = ["LBNHC", "LBHNC", "BLNHC", "BLHNC", "BHLNC", "LHBNC"]
+
+
+def _real_model_side(model, layout):
+    """Plan and allocate one side's KV caches the way the engine does, with
+    just enough blocks for one request of one or two blocks per group."""
+    from vllm.v1.attention.backends.utils import resolve_kv_cache_layout
+    from vllm.v1.core.kv_cache_planning import (
+        get_global_kv_cache_groups,
+        get_kv_cache_configs,
+    )
+    from vllm.v1.worker.utils import allocate_kv_cache
+
+    block_size, layers = _REAL_MODELS[model]
+    vllm_config = create_vllm_config(
+        block_size=block_size, max_model_len=block_size, kv_role="kv_both"
+    )
+    vllm_config.model_config.max_model_len = block_size
+    vllm_config.cache_config.enable_prefix_caching = False
+    vllm_config.cache_config.kv_cache_layout = layout
+    vllm_config.kv_transfer_config.kv_buffer_device = "cuda"
+    # A Mamba page is padded to the attention page, which holds it.
+    specs = {name: make(block_size, None) for name, make in layers}
+    mamba_page = max(
+        (
+            spec.page_size_bytes
+            for spec in specs.values()
+            if not isinstance(spec, MambaSpec)
+        ),
+        default=None,
+    )
+    specs = {name: make(block_size, mamba_page) for name, make in layers}
+    resolved = resolve_kv_cache_layout(
+        vllm_config, [_LAYOUT_PREFERENCE], specs.values()
     )
-    # 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])
+    groups = get_global_kv_cache_groups(vllm_config, [specs])
+    vllm_config.cache_config.num_gpu_blocks_override = 2 + sum(
+        _blocks_per_request(g.kv_cache_spec) for g in groups
+    )
+    kv_cache_config = get_kv_cache_configs(vllm_config, [specs], [1 << 40], groups)[0]
+    kv_cache_config.kv_cache_layout = vllm_config.cache_config.kv_cache_layout
+    caches = allocate_kv_cache(
+        kv_cache_config,
+        torch.device("cpu"),
+        resolved,
+        [g.kv_cache_spec.block_size for g in kv_cache_config.kv_cache_groups],
+    )
+    return vllm_config, kv_cache_config, caches
+
+
+def _blocks_per_request(spec):
+    return 1 if isinstance(spec, MambaSpec) else 2
+
+
+def _real_model_worker(engine_id, vllm_config, kv_cache_config, caches):
+    from vllm.config import set_current_vllm_config
+    from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
+        NixlConnectorWorker,
+    )
+
+    use_mla = any(
+        isinstance(g.kv_cache_spec, MLAAttentionSpec)
+        for g in kv_cache_config.kv_cache_groups
+    )
+    fake_backend = MagicMock()
+    fake_backend.get_name.return_value = "FLASHMLA" if use_mla else "FLASH_ATTN"
+    fake_backend.full_cls_name.return_value = "fake.backend"
+    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, "select_common_block_size", lambda size, *_: size),
+        patch.object(bw, "current_platform", fake_platform),
+        patch(
+            "vllm.model_executor.layers.mamba.mamba_utils.get_conv_state_layout",
+            return_value="DS",
+        ),
+        set_current_vllm_config(vllm_config),
+    ):
+        worker = NixlConnectorWorker(vllm_config, engine_id, kv_cache_config)
+        worker.use_mla = use_mla
+        worker.register_kv_caches(caches)
+    return worker
+
+
+def _cache_pages(cache, block):
+    """A layer's page(s) for one block: one tensor, or one per Mamba state."""
+    return (
+        [c[block] for c in cache]
+        if isinstance(cache, (list, tuple))
+        else [cache[block]]
+    )
+
+
+def _tag(page, *key):
+    gen = torch.Generator().manual_seed(hash(key) & 0x7FFFFFFF)
+    return torch.randint(1, 120, page.shape, generator=gen).to(page.dtype)
+
+
+def _run_recorded_copies(nixl):
+    """Run every recorded READ as memcpy: remote desc bytes -> local desc."""
+    import ctypes
+
+    for op, local_handle, local_ids, remote_handle, remote_ids in nixl.xfers:
+        assert op == "READ" and len(local_ids) == len(remote_ids)
+        for local_id, remote_id in zip(local_ids, remote_ids):
+            dst, dst_len, _ = (int(x) for x in nixl.dlists[local_handle][local_id])
+            src, src_len, _ = (int(x) for x in nixl.dlists[remote_handle][remote_id])
+            assert dst_len == src_len
+            ctypes.memmove(dst, src, dst_len)
+
+
+@pytest.mark.cpu_test
+@pytest.mark.parametrize(
+    "model,p_layout,d_layout",
+    [(model, None, None) for model in _REAL_MODELS]
+    + [
+        ("DeepSeek-V3", "BLNHC", "BLNHC"),
+        ("DeepSeek-V3", "BLNHC", "LBNHC"),
+        ("Qwen3-Next-80B-A3B", "BLHNC", "BLHNC"),
+        ("NVIDIA-Nemotron-Nano-9B-v2", "BLHNC", "BLHNC"),
+        ("Kimi-Linear-48B-A3B", "BLNHC", "BLNHC"),
+    ],
+)
+def test_real_model_layout_pull_copies_every_page(model, p_layout, d_layout):
+    """D pulls one request from P with each model's real KV plan: every page
+    of every layer arrives at D's block for the same layer and request
+    position, and nothing else in D's caches changes."""
+    from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
+        NixlConnectorMetadata,
+    )
+
+    p_config, p_kv_config, p_caches = _real_model_side(model, p_layout)
+    d_config, d_kv_config, d_caches = _real_model_side(model, d_layout)
+    p = _real_model_worker("P", p_config, p_kv_config, p_caches)
+    d = _real_model_worker("D", d_config, d_kv_config, d_caches)
+
+    # Block ids run up from 1 on P and down from the last block on D.
+    groups = d_kv_config.transfer_groups
+    p_ids: list[list[int]] = []
+    d_ids: list[list[int]] = []
+    for g in groups:
+        n = _blocks_per_request(g.kv_cache_spec)
+        p_ids.append([1 + sum(map(len, p_ids)) + k for k in range(n)])
+        d_ids.append(
+            [d_kv_config.num_blocks - 1 - sum(map(len, d_ids)) - k for k in range(n)]
+        )
+    for g, ids in zip(groups, p_ids):
+        for layer in g.layer_names:
+            for pos, block in enumerate(ids):
+                for part, page in enumerate(_cache_pages(p_caches[layer], block)):
+                    page.copy_(_tag(page, layer, pos, part))
+
+    meta = msgspec.msgpack.Decoder(NixlAgentMetadata).decode(
+        p.xfer_handshake_metadata.agent_metadata_bytes
+    )
+    _register_remote_agents(d, meta, 1)
+    metadata = NixlConnectorMetadata()
+    metadata.add_new_req_to_recv(
+        request_id="req",
+        local_block_ids=tuple(d_ids),
+        kv_transfer_params={
+            "remote_block_ids": p_ids,
+            "remote_engine_id": "P",
+            "remote_request_id": "prefill-req",
+            "remote_host": "remote-host",
+            "remote_port": 1234,
+            "tp_size": 1,
+        },
+    )
+    req = metadata.reqs_to_recv["req"]
+    req.local_physical_block_ids = d._logical_to_kernel_block_ids(
+        req.local_block_ids, d._physical_blocks_per_logical_kv_block
+    )
+    d._recving_metadata["req"] = req
+    d._read_blocks_for_req("req", req)
+    assert d.nixl_wrapper.xfers
+    _run_recorded_copies(d.nixl_wrapper)
+
+    wrong = []
+    for g, ids in zip(groups, d_ids):
+        for layer in g.layer_names:
+            for pos, block in enumerate(ids):
+                for part, page in enumerate(_cache_pages(d_caches[layer], block)):
+                    if not torch.equal(page, _tag(page, layer, pos, part)):
+                        wrong.append(f"{layer}[{pos}].{part}")
+                    page.zero_()
+    assert not wrong, wrong[:8]
+    storages = {
+        t.untyped_storage().data_ptr(): t
+        for cache in d_caches.values()
+        for t in (cache if isinstance(cache, (list, tuple)) else [cache])
+    }
+    for t in storages.values():
+        raw = torch.empty(0, dtype=torch.uint8).set_(t.untyped_storage())
+        assert not raw.any(), "D bytes outside the request's pages changed"
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 a36ceb472dea..9498dbf53c09 100644
--- a/tests/v1/kv_connector/unit/test_nixl_push_connector.py
+++ b/tests/v1/kv_connector/unit/test_nixl_push_connector.py
@@ -60,7 +60,12 @@
 from vllm.v1.kv_cache_interface import FullAttentionSpec
 from vllm.v1.outputs import KVConnectorOutput
 
-from .utils import create_request, expand_strided_descs, make_nixl_push_scheduler
+from .utils import (
+    create_request,
+    expand_strided_descs,
+    make_nixl_push_scheduler,
+    region_pairs_by_index,
+)
 
 # ----------------------------------------------------------------- #
 #  Helpers / fakes                                                   #
@@ -386,15 +391,11 @@ def fresh(cls) -> _StubWriterWorker:
         w._invalid_remote_engines = set()
         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._packed_member_layouts = {}
-        w._transfer_layer_group_ids = ()
         w.dst_region_num_blocks = {}
         w.dst_region_group_ids = {}
-        w.dst_uses_region_group_mapping = {}
+        w._dst_pairs = {}
         w.dst_region_mem_types = {}
         # Single non-hybrid attention group, matching the stub block id lists.
         w._has_mamba = False
@@ -1062,7 +1063,7 @@ def test_push_submission_failure_is_remembered_after_sibling_success(
         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._dst_pairs[w.engine_id] = region_pairs_by_index(1)
         w.nixl_wrapper.make_prepped_xfer.return_value = 101
         w.nixl_wrapper.transfer.side_effect = RuntimeError("submit")
         if release_fails:
@@ -1408,7 +1409,7 @@ def test_get_new_notifs_extends_lease_on_heartbeat(self):
 
 class TestPushPipelineParallel:
     """PP-sharded producer: per-stage completion counting, pp_size plumbing,
-    and remote-region slicing."""
+    and remote-region selection."""
 
     def test_completion_waits_for_one_notif_per_pp_stage(self):
         """Each PP stage WRITEs its own layers and sends one notif; D must
@@ -1450,17 +1451,21 @@ def test_req_meta_reads_pp_size_from_kv_transfer_params(self):
         metadata.add_new_req_to_recv("req-default", ([0],), params)
         assert metadata.reqs_to_recv["req-default"].pp_size == 1
 
-    def test_add_remote_agent_slices_remote_regions_to_local_pp_window(self):
-        """With PP>1 the producer registered regions for the full model, but
-        this worker holds only a contiguous layer slice; add_remote_agent
-        trims the remote region list to [offset : offset + num_local_regions]
-        before building descriptors. We stop right after the slice via a
-        sentinel on the next collaborator call."""
+    def test_add_remote_agent_keeps_the_remote_regions_of_local_layers(self):
+        """With PP>1 the consumer registered regions for the full model, but
+        this worker holds only a layer slice; add_remote_agent keeps the
+        consumer's regions holding those layers before building descriptors.
+        We stop right after via a sentinel on the next collaborator call."""
         block_len = 4096 * 16
         w = _StubWriterWorker.fresh()  # seeds writer-thread state for teardown
         w.pp_size = 2
-        w._remote_region_offset = 2  # this worker owns layers [2, 4)
         w.block_len_per_layer = [block_len, block_len]  # 2 local layers
+        w._region_is_mla = [False, False]
+        w.region_group_ids = [0, 0]
+        w.kv_cache_config = MagicMock(
+            transfer_group_index_by_layer={f"L{i}": 0 for i in range(4)}
+        )
+        w._set_region_layers([["L2"], ["L3"]])
         w.nixl_wrapper = MagicMock()
 
         class _StopAfterSlice(RuntimeError):
@@ -1483,12 +1488,13 @@ class _StopAfterSlice(RuntimeError):
             ssm_sizes=(0, 0),
             attn_backend_name="FLASH_ATTN",
             physical_blocks_per_logical_kv_block=1,
+            region_members=[["L0"], ["L1"], ["L2"], ["L3"]],
         )
 
         with pytest.raises(_StopAfterSlice):
             w.add_remote_agent(meta, remote_tp_rank=0, remote_tp_size=1)
 
-        # Sliced down to this worker's layer window (regions [2:4]).
+        # Down to the regions of this worker's layers.
         assert meta.kv_caches_base_addr == [12, 13]
         assert meta.block_lens == [block_len, block_len]
 
@@ -1531,6 +1537,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_pairs[engine_id] = region_pairs_by_index(1)
         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}}
@@ -1613,6 +1620,7 @@ 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_pairs[engine_id] = region_pairs_by_index(4, groups=[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]}
 
@@ -1698,10 +1706,7 @@ def _worker_driving_xfer(engine_id: str = "decode-engine"):
             w.engine_id: [10_000],
         }
         w.dst_region_group_ids = {engine_id: [0], w.engine_id: [0]}
-        w.dst_uses_region_group_mapping = {
-            engine_id: False,
-            w.engine_id: False,
-        }
+        w._dst_pairs = {engine_id: region_pairs_by_index(1)}
         w.dst_xfer_side_handles = {engine_id: {0: 5000}}
         w.src_xfer_handles_by_block_size = {16: 2000}
 
@@ -1746,16 +1751,16 @@ def test_no_prefix_hit_leaves_blocks_untrimmed(self):
         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."""
+    @pytest.mark.parametrize("groups_match", [False, True])
+    def test_push_requires_each_layer_in_the_same_cache_group(self, groups_match):
+        """A layer in a pooled remote region still receives its blocks; a layer
+        in another cache group on D is refused."""
         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,)
+        w._dst_pairs[engine_id].groups_match = groups_match
         reg = _registration_data("req-groups", local_block_ids=([500, 501],))
 
-        if not layer_name_routing:
+        if not groups_match:
             with pytest.raises(NotImplementedError, match="cache-group layouts"):
                 NixlPushConnectorWorker._do_start_push_kv(
                     w, "req-groups", ([10, 11],), reg
@@ -1797,36 +1802,44 @@ 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._is_hma_required = True
     worker.kv_cache_config = MagicMock(transfer_group_index_by_layer=group_by_layer)
     worker.transfer_topo = MagicMock()
+    worker.region_group_ids = [
+        groups.pop() if len(groups := {group_by_layer[n] for n in layers}) == 1 else -1
+        for layers in region_layers
+    ]
+    worker.block_len_per_layer = [128] * len(region_layers)
+    worker._region_is_mla = [False] * len(region_layers)
     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):
+def test_shared_region_reads_once_per_group_from_one_descriptor_list():
+    """A region shared by two cache groups keeps one descriptor list and is
+    read at each group's blocks."""
     worker = _layer_routing_worker(
         [["a", "a.swa"], ["b"]], {"a": 0, "a.swa": 1, "b": 0}
     )
+    pairs = worker._pair_regions(
+        _agent_metadata([["a", "a.swa"], ["b"]], [0xA000, 0xB000], [128, 128])
+    )
+    assert (pairs.local, pairs.group) == ([0, 0, 1], [0, 1, 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,
+        region_num_blocks=[4, 7],
+        pairs=(pairs.local, pairs.group),
     )
-    assert desc_ids.tolist() == expected
+    assert desc_ids.tolist() == [1, 2, 5, 6, 5]
 
 
 def test_layer_metadata_round_trip():
@@ -1835,31 +1848,32 @@ def test_layer_metadata_round_trip():
     metadata.region_group_ids = [-1]
     metadata.region_names = ["L0"]
     metadata.region_mem_types = ["VRAM"]
+    metadata.layer_group_ids = {"L0": 0, "L1": 1}
 
     encoded = msgspec.msgpack.encode(metadata)
     assert msgspec.msgpack.Decoder(NixlAgentMetadata).decode(encoded) == metadata
 
 
 @pytest.mark.parametrize(
-    "layouts, error",
+    ("l0_layout", "error"),
     [
-        ({"L0": (-1, 64)}, "escapes its block"),
-        ({"L0": (0, 0)}, "escapes its block"),
-        ({"L0": (200, 64)}, "escapes its block"),
-        ({"L1": (0, 64)}, "has no layout"),
-        ({"L0": (0, 32)}, "page sizes must match"),
+        ((-1, 64), "escapes its block"),
+        ((0, 0), "escapes its block"),
+        ((200, 64), "escapes its block"),
+        (None, "without a layout"),
     ],
 )
-def test_packed_layer_alignment_rejects_invalid_page_layout(layouts, error):
+def test_packed_row_expansion_rejects_invalid_page_layout(l0_layout, error):
     worker = _layer_routing_worker([["L0"]], {"L0": 0})
-    worker.block_len_per_layer = [64]
     metadata = _agent_metadata([["L0", "L1"]], [0x10000], [256])
-    metadata.packed_member_layouts = layouts
+    metadata.packed_member_layouts = {"L1": (128, 64)}
+    if l0_layout is not None:
+        metadata.packed_member_layouts["L0"] = l0_layout
     with pytest.raises(AssertionError, match=error):
-        worker._align_remote_regions_by_layer(metadata)
+        worker._pair_regions(metadata)
 
 
-def test_layer_alignment_preserves_nonpacked_regions_alongside_packed_layers():
+def test_pairing_expands_packed_rows_alongside_nonpacked_regions():
     worker = _layer_routing_worker([["a"], ["b"]], {"a": 0, "b": 1})
     worker.block_len_per_layer = [64, 128]
     metadata = _agent_metadata(
@@ -1867,79 +1881,55 @@ def test_layer_alignment_preserves_nonpacked_regions_alongside_packed_layers():
     )
     metadata.packed_member_layouts = {"a": (64, 64), "other": (128, 64)}
 
-    worker._align_remote_regions_by_layer(metadata)
+    pairs = worker._pair_regions(metadata)
+    worker._select_remote_regions(metadata, pairs.remote_regions)
 
+    assert (pairs.remote, pairs.local) == ([0, 1], [0, 1])
     assert metadata.kv_caches_base_addr == [0x1040, 0x2000]
     assert metadata.block_lens == [64, 128]
     assert metadata.block_strides == [256, 128]
 
 
-def test_layer_identity_gate_preserves_the_non_hma_path():
-    assert _layer_routing_worker([["a"]], {"a": 0})._transfer_layer_region_indices == (
-        0,
+def test_pairing_keeps_pooled_regions_once():
+    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]
 
-    # A base worker (pull) never routes by layer name.
-    pull = object.__new__(NixlBaseConnectorWorker)
-    pull.pp_size = 2
-    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
-        == ()
-    )
+    pairs = worker._pair_regions(metadata)
 
-    # 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
-        == ()
+    assert pairs.remote_regions == [0, 1]
+    assert (pairs.remote, pairs.local, pairs.group) == (
+        [0, 0, 1],
+        [0, 0, 1],
+        [0, 1, 0],
     )
+    assert pairs.groups_match
 
 
-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():
+@pytest.mark.parametrize(
+    ("remote_layer_groups", "groups_match"),
+    [({"a": 0, "a.swa": 1, "b": 0}, True), ({"a": 0, "a.swa": 0, "b": 0}, False)],
+)
+def test_pairing_compares_each_layer_s_group_in_a_shared_region(
+    remote_layer_groups, groups_match
+):
+    """A region shared by two groups advertises group -1, so only the
+    per-layer groups show that the peer puts a layer in another group (here P
+    treats its sliding-window layer as full attention)."""
     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)
+    metadata.layer_group_ids = remote_layer_groups
 
-    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"]
+    assert worker._pair_regions(metadata).groups_match is groups_match
 
 
-def test_layer_alignment_filters_and_reorders_a_pp_stage():
+def test_pairing_keeps_a_pp_stage_s_regions():
     worker = _layer_routing_worker([["l2"], ["l3"]], {"l2": 0, "l3": 1})
     metadata = _agent_metadata(
         [["l2"], ["l0"], ["l3"], ["l1"]],
@@ -1952,10 +1942,10 @@ def test_layer_alignment_filters_and_reorders_a_pp_stage():
     metadata.region_names = ["l2", "l0", "l3", "l1"]
     metadata.region_mem_types = ["VRAM"] * 4
 
-    worker._align_remote_regions_by_layer(metadata)
+    pairs = worker._pair_regions(metadata)
+    worker._select_remote_regions(metadata, pairs.remote_regions)
 
-    assert worker._transfer_layer_region_indices == (0, 1)
-    assert worker._transfer_layer_group_ids == (0, 1)
+    assert (pairs.remote, pairs.local, pairs.group) == ([0, 1], [0, 1], [0, 1])
     assert metadata.kv_caches_base_addr == [0xC000, 0xD000]
     assert metadata.block_lens == [65536, 32768]
     assert metadata.block_strides == [131072, 65536]
@@ -1979,9 +1969,7 @@ def test_layer_descriptors_pair_layers_across_asymmetric_pp_split(
     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
@@ -1994,11 +1982,11 @@ def test_layer_descriptors_pair_layers_across_asymmetric_pp_split(
     )
     consumer.num_blocks = 4
     consumer.region_num_blocks = remote_num_blocks
-    worker._align_remote_regions_by_layer(consumer)
+    pairs = worker._pair_regions(consumer)
+    worker._select_remote_regions(consumer, pairs.remote_regions)
 
-    # 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]
+    # L3 -> remote region 1, L4 and L5 -> 2; region 0 holds no local layer.
+    assert consumer.kv_caches_base_addr == [0xB000, 0xC000]
 
     plan = TPMapping(((0,), (0,)), (0,), {0: 0}, 0)
     # Builders emit one strided run per region; expand to per-block descriptors.
@@ -2006,103 +1994,101 @@ def test_layer_descriptors_pair_layers_across_asymmetric_pp_split(
         worker._build_fa_local([0x1000, 0x2000], block_size_ratio=1)
     )
     remote_descs = expand_strided_descs(
-        worker._build_fa_remote(plan, consumer, block_size_ratio=1)
+        worker._build_fa_remote(
+            plan, consumer, block_size_ratio=1, local_region_ids=pairs.partner
+        )
     )
-    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
+        [[1, 2], [3]],
+        4,
+        None,
+        1,
+        region_num_blocks=local_num_blocks,
+        pairs=(pairs.local, pairs.group),
     )
     remote_ids = worker._compute_desc_ids(
-        [[0, 1], [2]], 4, None, 1, region_num_blocks=consumer.region_num_blocks
+        [[0, 1], [2]],
+        4,
+        None,
+        1,
+        region_num_blocks=consumer.region_num_blocks,
+        pairs=(pairs.remote, pairs.group),
     )
-    assert len(local_descs) == sum(local_counts)
+    assert len(local_descs) == sum(local_num_blocks)
     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],
+    assert sorted(
+        zip(
+            local_descs[local_ids].tolist(),
+            remote_descs[remote_ids].tolist(),
+            strict=True,
+        )
+    ) == [
+        ([0x1100, 128, 0], [0xB000, 128, 7]),
+        ([0x1200, 128, 0], [0xB200, 128, 7]),
+        ([0x1300, 128, 0], [0xC800, 128, 7]),
+        ([0x2200, 128, 0], [0xC000, 128, 7]),
+        ([0x2400, 128, 0], [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}
+@pytest.mark.parametrize(
+    ("local_layers", "pp_size", "remote_layers", "error", "match"),
+    [
+        # Pairing by index would silently transfer stale KV, so a peer that
+        # does not advertise its layers must fail the handshake instead.
+        ([["a"]], 2, [], RuntimeError, "did not advertise the layers"),
+        ([["a"], ["b"]], 2, [["a"]], RuntimeError, "does not hold the KV layers"),
+        # Only a PP stage may hold a slice of the peer's layers.
+        ([["a"]], 1, [["a"], ["b"]], RuntimeError, "has no counterpart here"),
+        ([["a"]], 2, [["a"], ["a"]], NotImplementedError, "spans 2 NIXL regions"),
+    ],
+)
+def test_pairing_rejects_unpairable_peers(
+    local_layers, pp_size, remote_layers, error, match
+):
+    worker = _layer_routing_worker(local_layers, {"a": 0, "b": 0}, pp_size=pp_size)
+    num_regions = max(len(remote_layers), 1)
+    metadata = _agent_metadata(
+        remote_layers,
+        [0xA000 + 0x1000 * i for i in range(num_regions)],
+        [128] * num_regions,
     )
-    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
 
+    with pytest.raises(error, match=match):
+        worker._pair_regions(metadata)
 
-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})
+def test_pairing_pairs_split_layers_in_order():
+    """A layer registered as two regions (K and V apart) pairs them in order."""
+    worker = _layer_routing_worker([["a"], ["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)
+    pairs = worker._pair_regions(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)
+    assert worker._regions_by_layer == {"a": [0, 1]}
+    assert (pairs.remote, pairs.local) == ([0, 1], [0, 1])
 
 
 @pytest.mark.parametrize(
     "field",
-    ["region_num_blocks", "region_group_ids", "region_names", "region_mem_types"],
+    [
+        "block_lens",
+        "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})
+def test_pairing_rejects_inconsistent_region_geometry(field):
+    worker = _layer_routing_worker([["a"], ["b"]], {"a": 0, "b": 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
+    with pytest.raises(RuntimeError, match="metadata of different lengths"):
+        worker._pair_regions(metadata)
 
 
 def test_head_sharded_draft_push_rejects_heterogeneous_tp():
@@ -2120,16 +2106,6 @@ def test_head_sharded_draft_push_rejects_heterogeneous_tp():
             w._validate_remote_agent_handshake(meta, remote_tp_size=1)
 
 
-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"),
     [
diff --git a/tests/v1/kv_connector/unit/test_offloading_connector.py b/tests/v1/kv_connector/unit/test_offloading_connector.py
index d87fa002705f..719f17648a3f 100644
--- a/tests/v1/kv_connector/unit/test_offloading_connector.py
+++ b/tests/v1/kv_connector/unit/test_offloading_connector.py
@@ -20,6 +20,7 @@
 from vllm.platforms import current_platform
 from vllm.utils.math_utils import cdiv
 from vllm.v1.kv_offload.tiering.manager import TieringOffloadingManager
+from vllm.v1.metrics.reader import Gauge
 
 # The offloaded block size must stay a whole number of GPU blocks.
 CPU_BLOCK_SIZES: int
@@ -843,3 +844,57 @@ def stepped() -> bool:
         assert output.num_cached_tokens == len(prefix)
     finally:
         del llm
+
+
+def test_p2p_cache_reset_with_absent_decoder(monkeypatch):
+    """A real prefill engine can reset while its decoder never fetches KV."""
+    monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
+    llm = LLM(
+        model="facebook/opt-125m",
+        max_model_len=256,
+        max_num_batched_tokens=256,
+        gpu_memory_utilization=0.2,
+        attention_config={"backend": "TRITON_ATTN"},
+        enforce_eager=True,
+        async_scheduling=False,
+        disable_log_stats=False,
+        kv_transfer_config=KVTransferConfig(
+            kv_connector="OffloadingConnector",
+            kv_role="kv_both",
+            kv_connector_extra_config={
+                "cpu_bytes_to_use": 8 << 20,
+                "block_size": 16,
+                "spec_name": "TieringOffloadingSpec",
+                "secondary_tiers": [{"type": "p2p", "host": "127.0.0.1", "port": 0}],
+            },
+        ),
+    )
+    prompt = TokensPrompt(prompt_token_ids=list(range(100, 165)))
+    params = SamplingParams(
+        temperature=0,
+        max_tokens=1,
+        extra_args={
+            "kv_transfer_params": {"remote_decoder": {"kv_request_id": "absent"}}
+        },
+    )
+
+    def has_parked_cascade():
+        llm.llm_engine.step()
+        return any(
+            isinstance(m, Gauge)
+            and m.name == "vllm:kv_offload_tiering_active_cascade_jobs"
+            and m.value > 0
+            for m in llm.get_metrics()
+        )
+
+    try:
+        before = llm.generate(prompt, params, use_tqdm=False)[0]
+        _poll_until(has_parked_cascade, timeout_s=10, what="a parked P2P store")
+        _wait_for_prefix_cache_reset(llm, reset_connector=True)
+        after = llm.generate(
+            prompt, SamplingParams(temperature=0, max_tokens=1), use_tqdm=False
+        )[0]
+        assert after.num_cached_tokens == 0
+        assert after.outputs[0].token_ids == before.outputs[0].token_ids
+    finally:
+        llm.llm_engine.engine_core.shutdown()
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 100b5e39cde9..aa3184917c3f 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
@@ -163,10 +163,15 @@ def test_ring_scratch_group_is_accepted_and_never_offloaded() -> None:
         ("SimpleCPUOffloadConnector", True, None, None),
     ],
 )
+@pytest.mark.parametrize("extensible", [False, True])
 def test_engine_queries_offload_blocks_only_when_enabled(
-    monkeypatch, connector, prefix_caching, worker_blocks, expected_blocks
+    monkeypatch, connector, prefix_caching, worker_blocks, expected_blocks, extensible
 ) -> None:
-    """Only active offload queries workers, which must return a block count."""
+    """Only active offload queries workers, which must return a block count.
+
+    An extensible KV cache creates the connector in `extend_kv_cache`, so the
+    query must follow it.
+    """
     config = create_vllm_config(
         kv_connector=connector or "SimpleCPUOffloadConnector",
         kv_role="kv_both",
@@ -179,6 +184,7 @@ def test_engine_queries_offload_blocks_only_when_enabled(
     if connector is None:
         config.kv_transfer_config = None
     config.cache_config.enable_prefix_caching = prefix_caching
+    config.cache_config.enable_extensible_kv_cache = extensible
     kv_config = KVCacheConfig(num_blocks=16, kv_cache_tensors=[], kv_cache_groups=[])
     executor = MagicMock()
     executor.get_kv_cache_specs.return_value = [{}]
@@ -199,3 +205,9 @@ def test_engine_queries_offload_blocks_only_when_enabled(
         expected_blocks is not None
     )
     assert config.cache_config.simple_cpu_offload_num_blocks == expected_blocks
+    calls = [name for name, _, _ in executor.mock_calls]
+    assert ("extend_kv_cache" in calls) == extensible
+    if extensible and expected_blocks is not None:
+        assert calls.index("extend_kv_cache") < calls.index(
+            "get_simple_cpu_offload_num_blocks"
+        )
diff --git a/tests/v1/kv_connector/unit/test_umbp_embedded.py b/tests/v1/kv_connector/unit/test_umbp_embedded.py
new file mode 100644
index 000000000000..5054b2582936
--- /dev/null
+++ b/tests/v1/kv_connector/unit/test_umbp_embedded.py
@@ -0,0 +1,726 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import importlib.util
+import sys
+import threading
+from dataclasses import replace
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from tests.v1.kv_connector.umbp_test_utils import (
+    _kv_cache_config,
+    _vllm_config,
+    handshake,
+    install_memory_runtime,
+    make_kv_caches,
+)
+from tests.v1.kv_connector.unit.utils import (
+    create_model_runner_output,
+    create_request,
+    create_scheduler,
+    create_vllm_config,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorRole
+from vllm.distributed.kv_transfer.kv_connector.v1.umbp.connector import (
+    UMBPStoreConnector,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.umbp.data import (
+    BlockIdentityCodec,
+    BlockLoadBatch,
+    BlockTransferPlan,
+    KVLayoutPlanner,
+    KVRange,
+    KVRegion,
+    RankTopology,
+    TransferJobStatus,
+    UMBPConnectorMetadata,
+    UMBPNamespace,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.umbp.runtime import (
+    EmbeddedRuntime,
+    UMBPRuntimeConfig,
+    UMBPRuntimeFactory,
+    embedded,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.umbp.runtime.embedded import (
+    _lookup_socket_path,
+    _MoriLookupServer,
+    _MoriSchedulerHandle,
+    _MoriWorkerHandle,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.umbp.worker import (
+    UMBPStoreConnectorWorker,
+)
+from vllm.platforms import current_platform
+from vllm.v1.kv_cache_interface import (
+    FullAttentionSpec,
+    KVCacheConfig,
+    KVCacheGroupSpec,
+    KVCacheTensor,
+    SlidingWindowSpec,
+)
+from vllm.v1.request import RequestStatus
+
+mori_available = importlib.util.find_spec("mori") is not None
+
+pytestmark = pytest.mark.skipif(
+    not (current_platform.is_rocm() and mori_available),
+    reason="UMBP is only available on ROCm with the mori package installed",
+)
+
+
+@pytest.mark.parametrize("for_store", [False, True])
+def test_mori_range_arguments_preserve_sparse_offsets_and_empty_objects(for_store):
+    plans = [
+        BlockTransferPlan("empty", 0),
+        BlockTransferPlan(
+            "scatter",
+            1,
+            ranges=(
+                KVRange("layer0", 0, 1, 1000, 64, 16, 48),
+                KVRange("layer1", 0, 1, 2000, 64, 8, 0),
+            ),
+        ),
+    ]
+    assert _MoriWorkerHandle._range_args(plans, for_store=for_store) == (
+        ["empty", "scatter"],
+        [0, 64] if for_store else [],
+        [[], [1000, 2000]],
+        [[], [16, 8]],
+        [[], [48, 0]],
+    )
+
+
+@pytest.fixture
+def bulk_worker(monkeypatch, socket_dir):
+    calls: list[
+        tuple[list[str], list[int], list[list[int]], list[list[int]], list[list[int]]]
+    ] = []
+
+    def get(keys, pointers, sizes, offsets):
+        calls.append((keys, [], pointers, sizes, offsets))
+        return [key != "missing" for key in keys]
+
+    def put(keys, object_sizes, pointers, sizes, offsets):
+        calls.append((keys, object_sizes, pointers, sizes, offsets))
+        return [True] * len(keys)
+
+    monkeypatch.setitem(
+        sys.modules,
+        "mori.cpp",
+        SimpleNamespace(MemoryLocationType=SimpleNamespace(CPU=0, GPU=1)),
+    )
+    client = SimpleNamespace(
+        register_memory=lambda *a: True,
+        deregister_memory=lambda *a: True,
+        batch_get_ranges_into_ptr=get,
+        batch_put_ranges_from_ptr=put,
+        flush=lambda: True,
+        close=lambda: None,
+    )
+    planner = KVLayoutPlanner(
+        [
+            KVRegion("a", 0, 64, 32, 0, 16),
+            KVRegion("b", 1, 32, 16, 32, 16),
+            KVRegion("c", 0, 128, 64, 48, 16),
+        ]
+    )
+    handle = _MoriWorkerHandle(
+        client,
+        "bulk",
+        RankTopology(),
+        socket_dir,
+        1,
+        5,
+        planner.describe(RankTopology()),
+    )
+    worker = UMBPStoreConnectorWorker(handle, planner)
+    buffers = {
+        r.layer_name: torch.empty((8, r.block_stride), dtype=torch.uint8)
+        for r in planner.regions
+    }
+    worker.register_kv_caches(buffers)
+    yield worker, buffers, calls
+    handle.close()
+
+
+@pytest.mark.parametrize("batched", [False, True])
+def test_bulk_load_preserves_group_order_failures_and_new_buffers(
+    bulk_worker, batched, monkeypatch
+):
+    worker, buffers, calls = bulk_worker
+    plans = [
+        BlockTransferPlan("group1", 3, group_id=1, request_id="r"),
+        BlockTransferPlan("missing", 4, group_id=0, request_id="r"),
+        BlockTransferPlan("group0", 6, group_id=0, request_id="r"),
+    ]
+    for caches in (buffers, {name: torch.empty_like(t) for name, t in buffers.items()}):
+        worker.register_kv_caches(caches)
+        expected = _MoriWorkerHandle._range_args(
+            [worker.layout.materialize(p) for p in plans], for_store=False
+        )
+        supplied = plans
+        if batched:
+            supplied = BlockLoadBatch(
+                [p.key for p in plans],
+                [p.block_id for p in plans],
+                [p.group_id for p in plans],
+                "r",
+            )
+            monkeypatch.setattr(
+                BlockLoadBatch,
+                "__getitem__",
+                lambda *a: pytest.fail(
+                    "bulk loading must not materialize per-block plans"
+                ),
+            )
+        job = worker.runtime.load_blocks(supplied)
+        assert job is not None
+        result = worker.runtime.wait(job)
+        assert calls[-1] == expected
+        assert result.plans is supplied if batched else result.plans == tuple(plans)
+        assert result.failed_block_ids == {4}
+        assert result.completed_keys == {"group0", "group1"}
+        assert result.completed_bytes == 112
+    assert len(calls) == 2
+    assert calls[0][2] != calls[1][2]
+
+
+def test_bulk_store_matches_materialized_ranges(bulk_worker, monkeypatch):
+    worker, _, calls = bulk_worker
+    plans = [
+        BlockTransferPlan("group1", 3, group_id=1, request_id="r"),
+        BlockTransferPlan("group0", 6, group_id=0, request_id="r"),
+    ]
+    expected = _MoriWorkerHandle._range_args(
+        [worker.layout.materialize(p) for p in plans]
+    )
+    monkeypatch.setattr(
+        worker.layout,
+        "materialize",
+        lambda *a: pytest.fail("bulk stores must not materialize per-block plans"),
+    )
+
+    worker.enqueue_stores(UMBPConnectorMetadata(store_requests={"r": plans}))
+    (job,) = worker._store_jobs.values()
+    result = worker.runtime.wait(job)
+
+    assert calls == [expected]
+    assert result.plans == tuple(plans)
+    assert result.completed_keys == {"group0", "group1"}
+    assert result.completed_bytes == 112
+
+
+@pytest.mark.parametrize("special", ["subblocks", "materialized", "batch-subblocks"])
+def test_bulk_load_falls_back_for_entire_mixed_batch(bulk_worker, special):
+    worker, buffers, calls = bulk_worker
+    plans = [
+        BlockTransferPlan("a", 1, group_id=1),
+        BlockTransferPlan("c", 2, group_id=0),
+    ]
+    if special in ("subblocks", "batch-subblocks"):
+        buffers["c"] = torch.empty((16, 32), dtype=torch.uint8)
+        worker.register_kv_caches(buffers)
+        if special == "batch-subblocks":
+            plans = BlockLoadBatch(
+                [p.key for p in plans],
+                [p.block_id for p in plans],
+                [p.group_id for p in plans],
+                "r",
+            )
+    else:
+        plans[1] = worker.layout.materialize(plans[1])
+    assert worker.runtime.load_blocks(plans) is None
+    assert not calls
+    expected = _MoriWorkerHandle._range_args(
+        [p if p.ranges else worker.layout.materialize(p) for p in plans],
+        for_store=False,
+    )
+    worker.start_load_kv(None, UMBPConnectorMetadata(load_requests={"r": plans}))
+    worker.wait_for_layer_load("")
+    assert calls == [expected]
+
+
+@pytest.mark.parametrize(
+    ("options", "error"),
+    [
+        ({"dram_low_watermark": 0.9, "dram_high_watermark": 0.7}, "must not"),
+        ({"dram_use_hugepages": 1}, "boolean"),
+        ({"dram_numa_node": -2}, ">= -1"),
+        ({"dram_hugepage_size": 0}, "positive integer"),
+        ({"capacity_bytes": 1024, "total_capacity_bytes": 4096}, "mutually exclusive"),
+    ],
+)
+def test_runtime_config_rejects_invalid_embedded_dram_options(options, error):
+    with pytest.raises(ValueError, match=error):
+        UMBPRuntimeFactory.build(
+            UMBPRuntimeConfig.from_vllm(_vllm_config({"mode": "embedded", **options}))
+        )
+
+
+def test_mori_client_is_created_only_for_worker(monkeypatch, socket_dir):
+    """Scheduler lookup must not allocate a second, unused DRAM store."""
+    monkeypatch.setitem(sys.modules, "mori.cpp", None)
+    runtime = EmbeddedRuntime.from_config(
+        UMBPRuntimeConfig(
+            "embedded", {"capacity_bytes": 1024, "lookup_dir": socket_dir}
+        )
+    )
+    topology = RankTopology()
+    layout = KVLayoutPlanner.from_kv_cache_config(_kv_cache_config()).describe(topology)
+    scheduler = runtime.create_scheduler_handle("worker-only", topology, layout)
+    scheduler.close()
+    with pytest.raises(RuntimeError, match="requires MORI"):
+        runtime.create_worker_handle("worker-only", topology, layout)
+
+    configs = []
+
+    def create_client(config):
+        configs.append(config)
+        return SimpleNamespace(close=lambda: None)
+
+    monkeypatch.setitem(
+        sys.modules,
+        "mori.cpp",
+        SimpleNamespace(
+            UMBPClient=create_client,
+            UMBPConfig=lambda: SimpleNamespace(dram=SimpleNamespace()),
+        ),
+    )
+    worker = runtime.create_worker_handle("worker-only", topology, layout)
+    try:
+        assert len(configs) == 1
+        assert configs[0].dram.capacity_bytes == 1024
+    finally:
+        worker.close()
+
+
+def test_runtime_config_resolves_total_embedded_capacity_per_rank():
+    config = UMBPRuntimeConfig.from_vllm(
+        _vllm_config({"mode": "embedded", "total_capacity_bytes": 4096})
+    ).resolve_for_rank_count(4)
+
+    runtime = EmbeddedRuntime.from_config(config)
+    assert runtime.options["capacity_bytes"] == 1024
+    assert config.options["total_capacity_bytes"] == 4096
+    assert "capacity_bytes" not in config.options
+
+
+def test_embedded_connector_rejects_pipeline_parallelism():
+    with pytest.raises(NotImplementedError, match="pipeline parallelism"):
+        UMBPStoreConnector(
+            _vllm_config(
+                {"mode": "embedded"},
+                pipeline_parallel_size=2,
+            ),
+            KVConnectorRole.SCHEDULER,
+            _kv_cache_config(),
+        )
+
+
+def test_lookup_sockets_separate_engines_and_data_parallel_ranks(socket_dir):
+    def path(options, dp_index=0):
+        config = UMBPRuntimeConfig("embedded", {"lookup_dir": socket_dir, **options})
+        handle = EmbeddedRuntime.from_config(
+            replace(config, dp_index=dp_index)
+        ).create_scheduler_handle("same-model", RankTopology(), None)
+        return handle._paths[RankTopology().local_namespace]
+
+    assert path({}) == path({})
+    assert path({}, dp_index=0) != path({}, dp_index=1)
+    assert path({"lookup_instance": "a"}) != path({"lookup_instance": "b"})
+    with pytest.raises(ValueError, match="lookup_instance"):
+        path({"lookup_instance": 1})
+
+
+def test_lookup_server_refuses_a_socket_another_engine_serves(socket_dir):
+    client = SimpleNamespace(batch_exists=lambda keys: [True] * len(keys))
+    path = _lookup_socket_path("taken", (0, 0, 0, 0), socket_dir, "dp0")
+    first = _MoriLookupServer(path, client)
+    second = _MoriLookupServer(path, client)
+    first.start()
+    try:
+        with pytest.raises(RuntimeError, match="another live engine"):
+            second.start()
+        second.close()
+        scheduler = _MoriSchedulerHandle("taken", RankTopology(), socket_dir, "dp0")
+        assert scheduler.lookup(["k"]) == [True]
+    finally:
+        second.close()
+        first.close()
+
+
+def test_lookup_server_replaces_a_stale_socket_file(socket_dir):
+    import socket
+
+    path = _lookup_socket_path("stale", (0, 0, 0, 0), socket_dir, "dp0")
+    dead = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
+    dead.bind(path)
+    dead.close()
+    server = _MoriLookupServer(
+        path, SimpleNamespace(batch_exists=lambda keys: [True] * len(keys))
+    )
+    server.start()
+    try:
+        scheduler = _MoriSchedulerHandle("stale", RankTopology(), socket_dir, "dp0")
+        assert scheduler.lookup(["k"]) == [True]
+    finally:
+        server.close()
+
+
+@pytest.mark.parametrize("close_rank", [None, 0, 1])
+def test_mori_scheduler_routes_lookup_and_clear_to_all_ranks(socket_dir, close_rank):
+    class _Client:
+        def __init__(self, keys):
+            self.keys = set(keys)
+
+        def batch_exists(self, keys):
+            return [key in self.keys for key in keys]
+
+        def clear(self):
+            self.keys.clear()
+            return True
+
+    namespace = "routed-lookup"
+    topology = RankTopology(tp_size=2)
+    codec0 = BlockIdentityCodec(UMBPNamespace(namespace), tp_rank=0)
+    codec1 = BlockIdentityCodec(UMBPNamespace(namespace), tp_rank=1)
+    keys = [codec0.key(b"a"), codec1.key(b"a")]
+    servers = []
+    clients = [_Client([key, "legacy-key"]) for key in keys]
+    for rank, client in zip(topology.all_namespaces(), clients, strict=True):
+        server = _MoriLookupServer(
+            _lookup_socket_path(namespace, rank, socket_dir),
+            client,
+        )
+        server.start()
+        servers.append(server)
+    scheduler = _MoriSchedulerHandle(namespace, topology, socket_dir)
+    try:
+        assert scheduler.lookup(keys) == [True, True]
+        assert scheduler.last_lookup_diagnostics == {
+            "unavailable_ranks": (),
+            "missing_keys": (),
+        }
+        new_keys = [codec0.key(b"new"), codec1.key(b"new")]
+        assert scheduler.lookup(new_keys) == [False, False]
+        for client, key in zip(clients, new_keys, strict=True):
+            client.keys.add(key)
+        assert scheduler.lookup(new_keys) == [True, True]
+        assert scheduler.lookup([keys[1], "legacy-key", keys[0]]) == [True] * 3
+        if close_rank is not None:
+            servers[close_rank].close()
+            assert scheduler.lookup(keys) == [rank != close_rank for rank in range(2)]
+            assert scheduler.last_lookup_diagnostics["unavailable_ranks"] == (
+                topology.all_namespaces()[close_rank],
+            )
+            assert scheduler.last_lookup_diagnostics["missing_keys"] == (
+                keys[close_rank],
+            )
+        assert scheduler.clear() is (close_rank is None)
+        for rank, client in enumerate(clients):
+            if rank != close_rank:
+                assert not client.keys
+        assert scheduler.lookup(keys + ["legacy-key"]) == [False] * 3
+    finally:
+        for server in servers:
+            server.close()
+
+
+def test_single_rank_lookup_returns_fresh_results(socket_dir):
+    """The only rank also answers legacy keys, and hits are never cached."""
+    namespace = "single-route"
+    codec = BlockIdentityCodec(UMBPNamespace(namespace))
+    topology = RankTopology()
+    keys = [codec.key(b"known"), "legacy", codec.key(b"new")]
+    resident = set(keys[:2])
+    server = _MoriLookupServer(
+        _lookup_socket_path(namespace, topology.local_namespace, socket_dir),
+        SimpleNamespace(batch_exists=lambda keys: [key in resident for key in keys]),
+    )
+    scheduler = _MoriSchedulerHandle(namespace, topology, socket_dir)
+    server.start()
+    try:
+        assert scheduler.lookup(keys) == [True, True, False]
+        resident.clear()
+        resident.add(keys[-1])
+        assert scheduler.lookup(keys) == [False, False, True]
+        server.close()
+        assert scheduler.lookup(keys) == [False] * len(keys)
+        assert scheduler.last_lookup_diagnostics["unavailable_ranks"] == (
+            topology.local_namespace,
+        )
+    finally:
+        scheduler.close()
+        server.close()
+
+
+@pytest.mark.parametrize("flush_succeeds", [False, True])
+def test_mori_publication_does_not_flush_on_the_worker_thread(
+    socket_dir, flush_succeeds
+):
+    flush_threads = []
+
+    def flush():
+        flush_threads.append(threading.get_ident())
+        return flush_succeeds
+
+    client = SimpleNamespace(
+        batch_put_ranges_from_ptr=lambda keys, *args: [True] * len(keys),
+        batch_exists=lambda keys: [True] * len(keys),
+        flush=flush,
+        close=lambda: None,
+    )
+    handle = _MoriWorkerHandle(client, "flush", RankTopology(), socket_dir, 1, 5)
+    try:
+        job = handle.wait(handle.store([BlockTransferPlan("key", 0)]))
+        assert len(flush_threads) == 1
+        assert flush_threads[0] != threading.get_ident()
+        if flush_succeeds:
+            handle.publish(job)
+        else:
+            assert job.status is TransferJobStatus.FAILED
+            with pytest.raises(RuntimeError, match="incomplete"):
+                handle.publish(job)
+        assert len(flush_threads) == 1
+    finally:
+        handle.close()
+
+
+def test_backend_timeout_fails_the_finished_job(socket_dir):
+    """A TimeoutError from MORI itself must not read as buffers still in use."""
+
+    def get(*args):
+        raise TimeoutError("backend already stopped")
+
+    handle = _MoriWorkerHandle(
+        SimpleNamespace(batch_get_ranges_into_ptr=get, close=lambda: None),
+        "timeout",
+        RankTopology(),
+        socket_dir,
+        1,
+        5,
+    )
+    try:
+        job = handle.load([BlockTransferPlan("failed", 0)])
+        assert handle.wait(job).status is TransferJobStatus.FAILED
+        assert handle.poll(job) is job
+    finally:
+        handle.close()
+
+
+@pytest.fixture(params=["load", "store", "flush"])
+def blocked_transfer(request, monkeypatch, socket_dir):
+    started = threading.Event()
+    release = threading.Event()
+    clock = SimpleNamespace(now=100.0)
+    monkeypatch.setattr(embedded.time, "monotonic", lambda: clock.now)
+
+    def block(stage):
+        if request.param == stage:
+            started.set()
+            assert release.wait(5)
+
+    def get(keys, *args):
+        block("load")
+        return [True] * len(keys)
+
+    def put(keys, *args):
+        block("store")
+        return [True] * len(keys)
+
+    def flush():
+        block("flush")
+        return True
+
+    client = SimpleNamespace(
+        batch_get_ranges_into_ptr=get,
+        batch_put_ranges_from_ptr=put,
+        flush=flush,
+        close=lambda: None,
+    )
+    handle = _MoriWorkerHandle(client, "blocked", RankTopology(), socket_dir, 1, 0.1)
+    transfer = handle.load if request.param == "load" else handle.store
+    try:
+        job = transfer([BlockTransferPlan("running", 0)])
+        assert started.wait(5)
+        yield handle, job, transfer, release, clock
+    finally:
+        release.set()
+        handle.close()
+
+
+@pytest.mark.parametrize("operation", ["poll", "wait", "cancel"])
+def test_transfer_timeout_keeps_buffers_owned(blocked_transfer, operation):
+    handle, job, _, release, clock = blocked_transfer
+    assert handle.poll(job) is None
+    clock.now += 1
+    with pytest.raises(TimeoutError, match="buffers still in use"):
+        getattr(handle, operation)(job)
+    assert job.status is TransferJobStatus.RUNNING
+
+    release.set()
+    handle.close()
+    assert handle.poll(job).status is TransferJobStatus.COMPLETED
+
+
+def test_timeout_does_not_count_time_waiting_for_a_worker(blocked_transfer):
+    handle, _, transfer, release, clock = blocked_transfer
+    queued = transfer([BlockTransferPlan("queued", 1)])
+    finished = threading.Event()
+    handle._futures[id(queued)].add_done_callback(lambda _: finished.set())
+    clock.now += 1
+    assert handle.poll(queued) is None
+
+    release.set()
+    assert finished.wait(5)
+    assert handle.poll(queued).status is TransferJobStatus.COMPLETED
+
+
+def test_lookup_server_close_does_not_unlink_a_replacement(socket_dir):
+    client = SimpleNamespace(batch_exists=lambda keys: [True] * len(keys))
+    path = _lookup_socket_path("replaced", (0, 0, 0, 0), socket_dir)
+    first = _MoriLookupServer(path, client)
+    first.start()
+    first.close()
+    second = _MoriLookupServer(path, client)
+    second.start()
+    try:
+        first.close()
+        scheduler = _MoriSchedulerHandle("replaced", RankTopology(), socket_dir)
+        assert scheduler.lookup(["k"]) == [True]
+    finally:
+        second.close()
+
+
+def _engine_step(scheduler, worker, before_load=lambda: None):
+    """One engine step with the worker connector driven as the model runner does."""
+    output = scheduler.schedule()
+    before_load()
+    worker.bind_connector_metadata(output.kv_connector_metadata)
+    worker.start_load_kv(None)
+    worker.wait_for_save()
+    results = worker.get_transfer_results(output.finished_req_ids)
+    runner_output = create_model_runner_output(
+        [scheduler.requests[req_id] for req_id in output.num_scheduled_tokens],
+        finished_sending=results.finished_sending,
+        finished_recving=results.finished_recving,
+        failed_recving=results.failed_recving,
+        invalid_block_ids=worker.get_block_ids_with_load_errors(),
+        kv_connector_worker_meta=worker.build_connector_worker_meta(),
+    )
+    worker.clear_connector_metadata()
+    scheduler.update_from_output(output, runner_output)
+    return output
+
+
+def _run_to_completion(scheduler, worker) -> dict[str, int]:
+    scheduled: dict[str, int] = {}
+    for _ in range(10):
+        if not scheduler.has_requests():
+            return scheduled
+        output = _engine_step(scheduler, worker)
+        for req_id, count in output.num_scheduled_tokens.items():
+            scheduled[req_id] = scheduled.get(req_id, 0) + count
+    raise AssertionError("the engine did not drain")
+
+
+def _attention_and_sliding_window_config() -> KVCacheConfig:
+    shape = dict(block_size=16, num_kv_heads=2, head_size=8, dtype=torch.float16)
+    groups = [
+        KVCacheGroupSpec(["full"], FullAttentionSpec(**shape)),
+        KVCacheGroupSpec(["window"], SlidingWindowSpec(**shape, sliding_window=64)),
+    ]
+    page_bytes = groups[0].kv_cache_spec.page_size_bytes
+    return KVCacheConfig(
+        num_blocks=16,
+        kv_cache_groups=groups,
+        kv_cache_tensors=[
+            KVCacheTensor(
+                size=16 * page_bytes,
+                layers=group.layer_names,
+                layer_stride=16 * page_bytes,
+                block_stride=page_bytes,
+            )
+            for group in groups
+        ],
+    )
+
+
+@pytest.mark.parametrize(
+    ("num_groups", "load_async"), [(1, True), (1, False), (2, True)]
+)
+@pytest.mark.parametrize("policy", ["fail", "recompute"])
+def test_object_evicted_after_lookup_follows_failure_policy(
+    monkeypatch, socket_dir, num_groups, load_async, policy
+):
+    """A hit that is gone by load time must fail or recompute, never reuse KV."""
+    runtime = install_memory_runtime(monkeypatch)
+    if num_groups == 1:
+        kv_cache_config = _kv_cache_config()
+        kv_caches = make_kv_caches()
+    else:
+        kv_cache_config = _attention_and_sliding_window_config()
+        kv_caches = {
+            name: torch.zeros(
+                (16, kv_cache_config.kv_cache_tensors[0].block_stride),
+                dtype=torch.uint8,
+            )
+            for name in ("full", "window")
+        }
+    vllm_config = create_vllm_config(
+        kv_connector="UMBPStoreConnector",
+        kv_role="kv_both",
+        kv_connector_extra_config={
+            "mode": "embedded",
+            "lookup_dir": socket_dir,
+            "load_async": load_async,
+        },
+        kv_load_failure_policy=policy,
+    )
+    scheduler = create_scheduler(
+        vllm_config,
+        num_blocks=kv_cache_config.num_blocks,
+        kv_cache_config=kv_cache_config,
+    )
+    worker = UMBPStoreConnector(vllm_config, KVConnectorRole.WORKER, kv_cache_config)
+    worker.register_kv_caches(kv_caches)
+    handshake(scheduler.connector, worker)
+    num_tokens = 3 * scheduler.block_size + 1
+    num_hit_tokens = 3 * scheduler.block_size
+
+    def same_prompt():
+        return create_request(
+            num_tokens=num_tokens, common_prefix_len=num_tokens, max_tokens=1
+        )
+
+    scheduler.add_request(same_prompt())
+    _run_to_completion(scheduler, worker)
+    assert scheduler.reset_prefix_cache()
+
+    request = same_prompt()
+    scheduler.add_request(request)
+    free_blocks = scheduler.kv_cache_manager.block_pool.get_num_free_blocks()
+
+    def evict_after_lookup():
+        assert request.prefill_stats.num_external_cached_tokens == num_hit_tokens
+        runtime.evict_all()
+
+    _engine_step(scheduler, worker, before_load=evict_after_lookup)
+    if policy == "fail":
+        assert request.status == RequestStatus.FINISHED_ERROR
+        _run_to_completion(scheduler, worker)
+        assert scheduler.kv_cache_manager.block_pool.get_num_free_blocks() == (
+            free_blocks
+        )
+    else:
+        scheduled = _run_to_completion(scheduler, worker)
+        assert request.status == RequestStatus.FINISHED_LENGTH_CAPPED
+        assert scheduled[request.request_id] >= num_tokens
+    # Blocks the failed load left unfilled must never serve a prefix hit.
+    _, num_cached, _ = scheduler.kv_cache_manager.get_computed_blocks(same_prompt())
+    assert num_cached == (0 if policy == "fail" else num_hit_tokens)
+    worker.shutdown()
diff --git a/tests/v1/kv_connector/unit/test_umbp_hybrid.py b/tests/v1/kv_connector/unit/test_umbp_hybrid.py
index ddd86b22c517..6d211f2a68a8 100644
--- a/tests/v1/kv_connector/unit/test_umbp_hybrid.py
+++ b/tests/v1/kv_connector/unit/test_umbp_hybrid.py
@@ -10,6 +10,7 @@
 import torch
 
 from tests.v1.kv_connector.umbp_test_utils import (
+    MemoryRuntime,
     _EmbeddedRuntime,
     _hybrid_kv_cache_config,
     _kv_cache_config,
@@ -18,6 +19,7 @@
     _WorkerHandle,
     make_request,
     make_scheduler_output,
+    mori_has_umbp,
 )
 from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorRole
 from vllm.distributed.kv_transfer.kv_connector.v1.umbp import (
@@ -29,6 +31,7 @@
 from vllm.distributed.kv_transfer.kv_connector.v1.umbp.data import (
     BlockIdentityCodec,
     BlockTransferPlan,
+    KVLayoutPlanner,
     LoadSpec,
     RankTopology,
     RequestTracker,
@@ -38,6 +41,10 @@
     UMBPConnectorWorkerMetadata,
     UMBPNamespace,
 )
+from vllm.distributed.kv_transfer.kv_connector.v1.umbp.runtime import (
+    EmbeddedRuntime,
+    UMBPRuntimeConfig,
+)
 from vllm.distributed.kv_transfer.kv_connector.v1.umbp.runtime.factory import (
     UMBPRuntimeFactory,
 )
@@ -54,6 +61,7 @@
 from vllm.v1.kv_cache_interface import (
     ChunkedLocalAttentionSpec,
     KVCacheGroupSpec,
+    KVCacheTensor,
     MLAAttentionSpec,
     SlidingWindowSpec,
     SparseCacheRole,
@@ -889,6 +897,141 @@ def test_finished_hybrid_tail_pins_full_pages_until_all_workers_finish():
     assert not scheduler.has_pending_push_work()
 
 
+@pytest.mark.parametrize(
+    "device",
+    [
+        "cpu",
+        pytest.param(
+            "cuda",
+            marks=pytest.mark.skipif(
+                not mori_has_umbp(), reason="MORI was built without UMBP"
+            ),
+        ),
+    ],
+)
+def test_finished_hybrid_tail_restores_every_byte(socket_dir, device):
+    """A hash-only tail survives request cleanup, then restores byte for byte."""
+    if device == "cpu":
+        runtime = MemoryRuntime()
+    else:
+        runtime = EmbeddedRuntime.from_config(
+            UMBPRuntimeConfig(
+                "embedded", {"capacity_bytes": 16 << 20, "lookup_dir": socket_dir}
+            )
+        )
+    config = _hybrid_kv_cache_config()
+    groups = config.kv_cache_groups
+    config = replace(
+        config,
+        kv_cache_tensors=[
+            KVCacheTensor(
+                size=8 * group.kv_cache_spec.page_size_bytes,
+                layers=group.layer_names,
+                layer_stride=8 * group.kv_cache_spec.page_size_bytes,
+                block_stride=group.kv_cache_spec.page_size_bytes,
+            )
+            for group in groups
+        ],
+    )
+    topology = RankTopology()
+    layout = KVLayoutPlanner.from_kv_cache_config(config)
+    descriptor = layout.describe(topology)
+    scheduler = UMBPStoreConnectorScheduler(
+        _vllm_config(
+            {"enable_partial_hash_hits": True, "load_async": True},
+            prefix_match_unit=4,
+        ),
+        config,
+        runtime.create_scheduler_handle("tail", topology, descriptor),
+        BlockIdentityCodec(UMBPNamespace("tail")),
+    )
+    worker = UMBPStoreConnectorWorker(
+        runtime.create_worker_handle("tail", topology, descriptor), layout
+    )
+    pool = BlockPool(8, True, 4)
+    scheduler.bind_gpu_block_pool(pool)
+    caches = {
+        group.layer_names[0]: torch.zeros(
+            (8, group.kv_cache_spec.page_size_bytes), dtype=torch.uint8, device=device
+        )
+        for group in groups
+    }
+    # Distinct int32 words, so a K/V swap or a shifted range cannot match.
+    pages = {
+        name: torch.arange(
+            group << 16,
+            (group << 16) + cache.shape[1] // 4,
+            dtype=torch.int32,
+            device=device,
+        ).view(torch.uint8)
+        for group, (name, cache) in enumerate(caches.items())
+    }
+    for block_id, (name, cache) in enumerate(caches.items(), 1):
+        cache[block_id] = pages[name]
+    worker.register_kv_caches(caches)
+    request = make_request(
+        "tail",
+        num_tokens=13,
+        num_computed_tokens=12,
+        block_hashes=[b"a", b"b", b"c"],
+    )
+    source_ids = ([1], [2])
+    scheduler.update_state_after_alloc(
+        request, SimpleNamespace(get_block_ids=lambda group_ids: source_ids), 0
+    )
+    sources = [pool.blocks[1], pool.blocks[2]]
+    pool.touch(sources)
+    assert not scheduler.register_finished_partial_tail(
+        request, source_ids, [(1, 2, 12)]
+    )
+    assert scheduler.request_finished(request, source_ids) == (False, None)
+    pool.free_blocks(sources)
+    output = make_scheduler_output(finished={"tail"})
+    try:
+        metadata = scheduler.build_connector_meta(output)
+        assert len(metadata.store_plans) == 2
+        worker.enqueue_stores(metadata)
+        worker.wait_for_save()
+        deadline = time.monotonic() + 10
+        while scheduler.has_pending_push_work():
+            assert time.monotonic() < deadline, "store completion was not reported"
+            worker.get_finished(set())
+            scheduler.update_connector_output(
+                SimpleNamespace(
+                    kv_connector_worker_meta=worker.build_connector_worker_meta()
+                )
+            )
+            time.sleep(0.001)
+        assert [block.ref_cnt for block in sources] == [0, 0]
+        for cache in caches.values():
+            cache.zero_()
+
+        request.request_id = "replay"
+        assert scheduler.get_num_new_matched_tokens(request, 0) == (12, True)
+        scheduler.update_state_after_alloc(
+            request, SimpleNamespace(get_block_ids=lambda group_ids: ([3], [4])), 12
+        )
+        output.finished_req_ids = set()
+        worker.start_load_kv(None, scheduler.build_connector_meta(output))
+        deadline = time.monotonic() + 10
+        while (finished := worker.get_finished(set())) == (None, None):
+            assert time.monotonic() < deadline, "load completion was not reported"
+            time.sleep(0.001)
+        assert finished == (None, {"replay"})
+        assert worker.get_failed_recving() == set()
+        for block_id, (name, cache) in enumerate(caches.items(), 3):
+            expected = torch.zeros_like(cache)
+            expected[block_id] = pages[name]
+            assert torch.equal(cache, expected), name
+        counters = worker.get_kv_connector_stats().reduce()
+        page_bytes = sum(group.kv_cache_spec.page_size_bytes for group in groups)
+        assert counters["load_num_bytes"] == page_bytes
+        assert counters["store_num_bytes"] == page_bytes
+    finally:
+        worker.close()
+        scheduler.close()
+
+
 @pytest.mark.parametrize("failure", [None, "missing-rank", "exception", "short"])
 @pytest.mark.parametrize("restored_tokens", [32, 64])
 @pytest.mark.parametrize("null_block", [False, True])
diff --git a/tests/v1/kv_connector/unit/utils.py b/tests/v1/kv_connector/unit/utils.py
index 903ef4e0e327..fc0b6b382547 100644
--- a/tests/v1/kv_connector/unit/utils.py
+++ b/tests/v1/kv_connector/unit/utils.py
@@ -2,10 +2,10 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import tempfile
 from collections import defaultdict
-from collections.abc import Callable
+from collections.abc import Callable, Sequence
 from dataclasses import dataclass
 from itertools import chain, count
-from typing import Any, Literal
+from typing import TYPE_CHECKING, Any, Literal
 from unittest.mock import Mock
 
 import numpy as np
@@ -50,6 +50,11 @@
 from vllm.v1.request import Request
 from vllm.v1.structured_output import StructuredOutputManager
 
+if TYPE_CHECKING:
+    from vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker import (
+        _RegionPairs,
+    )
+
 EOS_TOKEN_ID = 50256
 
 
@@ -289,6 +294,7 @@ def create_model_runner_output(
     use_eos: bool = False,
     token_id: int = 0,
     kv_connector_worker_meta: KVConnectorWorkerMetadata | None = None,
+    failed_recving: set[str] | None = None,
 ) -> ModelRunnerOutput:
     """Make dummy model runner output for testing."""
     # Make request data.
@@ -306,12 +312,14 @@ def create_model_runner_output(
             and finished_recving is None
             and invalid_block_ids is None
             and kv_connector_worker_meta is None
+            and failed_recving is None
         )
         else KVConnectorOutput(
             finished_sending=finished_sending,
             finished_recving=finished_recving,
             invalid_block_ids=invalid_block_ids or set(),
             kv_connector_worker_meta=kv_connector_worker_meta,
+            failed_recving=failed_recving or set(),
         )
     )
 
@@ -700,3 +708,24 @@ def expand_strided_descs(strided_descs: np.ndarray) -> np.ndarray:
     offsets = np.arange(len(out), dtype=np.uint64) - np.repeat(starts, counts)
     out[:, 0] += offsets * np.repeat(strided_descs[:, 3], counts)
     return out
+
+
+def region_pairs_by_index(
+    num_regions: int,
+    groups: Sequence[int] | None = None,
+    groups_match: bool = True,
+) -> "_RegionPairs":
+    """NIXL region pairs where peer region ``i`` pairs with local region ``i``."""
+    from vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker import (
+        _RegionPairs,
+    )
+
+    regions = list(range(num_regions))
+    return _RegionPairs(
+        remote=regions,
+        local=regions,
+        group=list(groups) if groups is not None else [0] * num_regions,
+        remote_regions=regions,
+        partner=regions,
+        groups_match=groups_match,
+    )
diff --git a/tests/v1/kv_offload/cpu/policies/test_factory.py b/tests/v1/kv_offload/cpu/policies/test_factory.py
index 98196a73e9bc..a37eb2559471 100644
--- a/tests/v1/kv_offload/cpu/policies/test_factory.py
+++ b/tests/v1/kv_offload/cpu/policies/test_factory.py
@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-from collections.abc import Iterable
+from collections.abc import Sequence
 
 import pytest
 
@@ -31,7 +31,13 @@ def insert(
     def remove(self, key: OffloadKey) -> None:
         pass
 
-    def touch(self, keys: Iterable[OffloadKey], req_context: ReqContext) -> None:
+    def on_request_finished(
+        self,
+        key_groups: Sequence[Sequence[OffloadKey]],
+        insertion_only_keys: set[OffloadKey],
+        reused_keys: set[OffloadKey],
+        req_context: ReqContext,
+    ) -> None:
         pass
 
     def evict(
diff --git a/tests/v1/kv_offload/cpu/test_gpu_worker.py b/tests/v1/kv_offload/cpu/test_gpu_worker.py
index aad2c2f4fbfb..b2ad143c6af4 100644
--- a/tests/v1/kv_offload/cpu/test_gpu_worker.py
+++ b/tests/v1/kv_offload/cpu/test_gpu_worker.py
@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import contextlib
 import logging
 import random
 import time
@@ -204,6 +205,202 @@ def test_handler_shutdown_skips_transfers_after_event_sync_failure() -> None:
     assert not handler.dst_tensors
 
 
+def _bare_handler(gpu_to_cpu: bool = True):
+    """Handler with only the state the failure paths touch."""
+    handler = gpu_worker.SingleDirectionOffloadingHandler.__new__(
+        gpu_worker.SingleDirectionOffloadingHandler
+    )
+    handler.gpu_to_cpu = gpu_to_cpu
+    handler._transfers = gpu_worker.deque()
+    handler._transfer_events = {}
+    handler._stream_pool = []
+    handler._event_pool = []
+    handler._buffer_pool = []
+    return handler
+
+
+def _submittable_handler(monkeypatch: pytest.MonkeyPatch, gpu_to_cpu: bool = True):
+    """Bare handler wired up far enough to run transfer_async off-device."""
+    handler = _bare_handler(gpu_to_cpu)
+    handler.layer_refs_per_group = [MagicMock()]
+    handler.src_blocks_per_chunk = 1
+    handler.dst_blocks_per_chunk = 1
+    handler._estimate_max_copy_ops = lambda group_sizes: 4
+    handler._fill_group_ops = lambda g_idx, **kwargs: (1, 128)
+    handler._swap_blocks_batch = MagicMock()
+    handler._buffer_pool = [tuple(torch.zeros(4, dtype=torch.int64) for _ in range(3))]
+
+    platform = MagicMock()
+    platform.stream.return_value = contextlib.nullcontext()
+    platform.is_cuda_alike.return_value = False
+    monkeypatch.setattr(gpu_worker, "current_platform", platform)
+    return handler
+
+
+def _submit(handler, job_id: int) -> bool:
+    spec = GPULoadStoreSpec(block_ids=[0], group_sizes=[1], block_indices=[0])
+    return handler.transfer_async(job_id, spec, spec)
+
+
+def test_submit_failure_is_not_reported_until_stream_drains(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """A submit fault leaves copies in flight, so the job completes only later.
+
+    swap_blocks_batch submits the descriptor list in chunks and the driver's
+    batch memcpy stops at the first failing copy, so earlier copies of the same
+    job may still be touching its blocks. Reporting the job before the stream
+    drains would let the scheduler reuse blocks that are still being written.
+    """
+    handler = _submittable_handler(monkeypatch)
+    stream, start_event, end_event = MagicMock(), MagicMock(), MagicMock()
+    handler._stream_pool = [stream]
+    handler._event_pool = [end_event, start_event]
+    handler._swap_blocks_batch.side_effect = RuntimeError("failed at index 1")
+    clear = MagicMock()
+    monkeypatch.setattr(gpu_worker, "_clear_pending_device_error", clear)
+
+    assert _submit(handler, 7) is True
+
+    clear.assert_called_once_with()
+    end_event.record.assert_called_once_with(stream)
+    (transfer,) = handler._transfers
+    assert (transfer.job_id, transfer.failed) == (7, True)
+    assert handler._transfer_events == {7: end_event}
+    assert not handler._stream_pool
+    assert not handler._event_pool
+    assert not handler._buffer_pool
+
+    end_event.query.return_value = False
+    assert handler.get_finished() == []
+
+    end_event.query.return_value = True
+    results = handler.get_finished()
+
+    assert [(r.job_id, r.success) for r in results] == [(7, False)]
+    assert not handler._transfers
+    assert not handler._transfer_events
+    assert not handler._stream_pool
+    assert not handler._event_pool
+    assert not handler._buffer_pool
+    start_event.elapsed_time.assert_not_called()
+
+
+def test_submit_failure_still_fences_wait_and_next_submission(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """The faulted job keeps ordering the work that follows it.
+
+    wait() is the store-flush fence and every submission chains onto the
+    previous transfer's end event; a faulted job that left neither would let
+    both run ahead of its in-flight copies.
+    """
+    handler = _submittable_handler(monkeypatch)
+    failed_stream, failed_start, failed_end = MagicMock(), MagicMock(), MagicMock()
+    handler._stream_pool = [failed_stream]
+    handler._event_pool = [failed_end, failed_start]
+    handler._swap_blocks_batch.side_effect = RuntimeError("failed at index 1")
+    _submit(handler, 7)
+
+    handler.wait({7})
+    failed_end.synchronize.assert_called_once_with()
+
+    next_stream = MagicMock()
+    handler._stream_pool = [next_stream]
+    handler._event_pool = [MagicMock(), MagicMock()]
+    handler._buffer_pool = [tuple(torch.zeros(4, dtype=torch.int64) for _ in range(3))]
+    handler._swap_blocks_batch.side_effect = None
+    _submit(handler, 8)
+
+    next_stream.wait_event.assert_called_once_with(failed_end)
+
+
+def test_submit_failure_propagates_when_the_end_event_cannot_be_recorded(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Without an end event the in-flight copies cannot be bounded.
+
+    The device is unrecoverable at that point, so failing loudly beats
+    reporting a completion the blocks have not actually reached.
+    """
+    handler = _submittable_handler(monkeypatch)
+    end_event = MagicMock()
+    end_event.record.side_effect = RuntimeError("device lost")
+    handler._stream_pool = [MagicMock()]
+    handler._event_pool = [end_event, MagicMock()]
+    handler._swap_blocks_batch.side_effect = RuntimeError("failed at index 1")
+
+    with pytest.raises(RuntimeError, match="device lost"):
+        _submit(handler, 7)
+
+
+def test_get_finished_reports_failure_when_event_query_raises() -> None:
+    """A device error latched after submit surfaces at poll time, not as a raise."""
+    handler = _bare_handler()
+    failing = MagicMock()
+    failing.end_event.query.side_effect = RuntimeError("device lost")
+    failing.job_id = 3
+    ready = MagicMock()
+    ready.job_id = 4
+    ready.failed = False
+    ready.end_event.query.return_value = True
+    ready.start_event.elapsed_time.return_value = 2.0
+    ready.num_bytes = 512
+    handler._transfers.extend([failing, ready])
+    handler._transfer_events = {3: failing.end_event, 4: ready.end_event}
+
+    results = handler.get_finished()
+
+    assert [(r.job_id, r.success) for r in results] == [(3, False), (4, True)]
+    assert not handler._transfers
+    assert not handler._transfer_events
+    assert handler._stream_pool == [ready.stream]
+    assert handler._buffer_pool == [
+        (ready.batch_src, ready.batch_dst, ready.batch_sizes)
+    ]
+
+
+def test_get_finished_propagates_non_runtime_errors() -> None:
+    """Only RuntimeError, which device faults raise, becomes a failed transfer.
+
+    Anything else, such as a missing op, must stay loud rather than silently
+    degrade the cache.
+    """
+    handler = _bare_handler()
+    broken = MagicMock()
+    broken.job_id = 5
+    broken.end_event.query.side_effect = AttributeError("no such op")
+    handler._transfers.append(broken)
+
+    with pytest.raises(AttributeError, match="no such op"):
+        handler.get_finished()
+
+
+def test_submit_propagates_non_runtime_errors(monkeypatch: pytest.MonkeyPatch) -> None:
+    handler = _submittable_handler(monkeypatch)
+    handler._stream_pool = [MagicMock()]
+    handler._event_pool = [MagicMock(), MagicMock()]
+    handler._swap_blocks_batch.side_effect = AttributeError("no such op")
+    clear = MagicMock()
+    monkeypatch.setattr(gpu_worker, "_clear_pending_device_error", clear)
+
+    with pytest.raises(AttributeError, match="no such op"):
+        _submit(handler, 7)
+    clear.assert_not_called()
+
+
+def test_wait_swallows_synchronize_failure() -> None:
+    """wait() must not raise; the failure is reported by the next poll."""
+    handler = _bare_handler()
+    event = MagicMock()
+    event.synchronize.side_effect = RuntimeError("device lost")
+    handler._transfer_events[9] = event
+
+    handler.wait({9})
+
+    event.synchronize.assert_called_once_with()
+
+
 @pytest.mark.parametrize("device_sync_fails", [False, True])
 def test_worker_syncs_before_cleanup_after_handler_failure(
     caplog_vllm, monkeypatch: pytest.MonkeyPatch, device_sync_fails: bool
@@ -710,3 +907,76 @@ def test_load_waits_for_pending_compute_stream_writes(default_vllm_config) -> No
                 torch.testing.assert_close(gpu_tensor[block_id].cpu(), expected)
     finally:
         worker.shutdown()
+
+
+@pytest.mark.skipif(
+    not current_platform.is_rocm(),
+    reason="faults a real hipMemcpyBatchAsync call",
+)
+@pytest.mark.parametrize("gpu_to_cpu", [True, False])
+@pytest.mark.parametrize(
+    "num_blocks", [16, 9000], ids=["one_batch", "fault_in_second_batch"]
+)
+def test_device_fault_does_not_fail_the_next_kernel_launch(
+    default_vllm_config, gpu_to_cpu: bool, num_blocks: int
+) -> None:
+    """A caught copy fault must not leave its error pending on the thread.
+
+    The runtime keeps a failed call's error until it is read, and the next
+    kernel launch check would raise it, killing the model step the catch was
+    meant to protect.
+    """
+    page_size_bytes = 4096
+    gpu_tensor = torch.zeros(
+        (num_blocks, page_size_bytes), dtype=torch.int8, device=DEVICES[0]
+    )
+    worker = CPUOffloadingWorker(
+        kv_caches=CanonicalKVCaches(
+            tensors=[
+                CanonicalKVCacheTensor(
+                    tensor=gpu_tensor, page_size_bytes=page_size_bytes
+                )
+            ],
+            group_data_refs=[
+                [CanonicalKVCacheRef(tensor_idx=0, page_size_bytes=page_size_bytes)]
+            ],
+        ),
+        blocks_per_chunk=1,
+        num_cpu_chunks=num_blocks,
+    )
+    handler = worker._store_handler if gpu_to_cpu else worker._load_handler
+    swap_blocks_batch = handler._swap_blocks_batch
+    faulted = []
+
+    def zero_one_copy(src, dst, sizes, **kwargs):
+        sizes[len(sizes) - 500 if num_blocks > 8192 else len(sizes) // 2] = 0
+        try:
+            swap_blocks_batch(src, dst, sizes, **kwargs)
+        except RuntimeError:
+            faulted.append(True)
+            raise
+
+    handler._swap_blocks_batch = zero_one_copy
+    gpu_spec = GPULoadStoreSpec(
+        list(range(num_blocks)), group_sizes=(num_blocks,), block_indices=(0,)
+    )
+    cpu_spec = CPULoadStoreSpec(list(range(num_blocks)))
+    try:
+        if gpu_to_cpu:
+            assert worker.submit_store(1, gpu_spec, cpu_spec)
+        else:
+            assert worker.submit_load(1, cpu_spec, gpu_spec)
+        if not faulted:
+            pytest.skip("this runtime accepts zero-size copies")
+
+        torch.arange(4096, device=DEVICES[0]).cumsum(0)
+        torch.accelerator.synchronize()
+
+        deadline = time.time() + 10
+        finished: list[TransferResult] = []
+        while time.time() < deadline and not finished:
+            finished = worker.get_finished()
+        assert [(r.job_id, r.success) for r in finished] == [(1, False)]
+    finally:
+        gpu_worker._clear_pending_device_error()
+        worker.shutdown()
diff --git a/tests/v1/kv_offload/cpu/test_manager.py b/tests/v1/kv_offload/cpu/test_manager.py
index 10c9f5e04d2d..188babea76fe 100644
--- a/tests/v1/kv_offload/cpu/test_manager.py
+++ b/tests/v1/kv_offload/cpu/test_manager.py
@@ -78,6 +78,14 @@ def to_keys(int_hashes: list[int]) -> list[OffloadKey]:
     return [to_key(i) for i in int_hashes]
 
 
+def reuse_keys(manager: CPUOffloadingManager, keys: list[OffloadKey]) -> None:
+    """Load ready keys in one finished request, recording a single access."""
+    ctx = make_req_context("reuse")
+    manager.prepare_load(keys, ctx)
+    manager.complete_load(keys, ctx)
+    manager.on_request_finished(ctx)
+
+
 def verify_store_output(
     prepare_store_output: PrepareStoreOutput | None,
     expected_prepare_store_output: ExpectedPrepareStoreOutput,
@@ -143,6 +151,20 @@ def to_key_sets(
     assert tuple(stores) == to_key_sets(expected_stores)
 
 
+def test_cpu_stored_event_carries_request_context():
+    manager = CPUOffloadingManager(
+        num_chunks=1,
+        enable_events=True,
+    )
+    req_context = make_req_context("writer")
+    keys = to_keys([1])
+    assert manager.prepare_store(keys, req_context) is not None
+    manager.complete_store(keys, req_context)
+
+    [event] = manager.take_events()
+    assert event.req_context is req_context
+
+
 def test_cpu_eviction_removed_precedes_stored():
     """An eviction is announced before the store that reuses its capacity."""
     manager = make_cpu_manager(num_chunks=2, enable_events=True)
@@ -170,7 +192,7 @@ def test_already_stored_chunk_not_evicted_during_prepare_store(eviction_policy):
 
     Scenario:
         - Store chunks [1, 2] and complete.
-        - touch([1]) makes chunk 2 the LRU candidate.
+        - Reusing [1] makes chunk 2 the LRU candidate.
         - prepare_store([2, 3, 4, 5]):
             * chunk 2 is filtered out as "already stored"
             * but without the fix, chunk 2 would be evicted as the LRU
@@ -187,8 +209,8 @@ def test_already_stored_chunk_not_evicted_during_prepare_store(eviction_policy):
     manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
     manager.complete_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
 
-    # touch [1] to make chunk 2 the LRU candidate
-    manager.touch(to_keys([1]), _EMPTY_REQ_CTX)
+    # reuse [1] to make chunk 2 the LRU candidate
+    reuse_keys(manager, to_keys([1]))
 
     # prepare_store([2, 3, 4, 5]):
     #   - chunk 2 is already stored -> filtered out of keys_to_store
@@ -464,10 +486,10 @@ def test_cpu_manager():
     # complete store [6, 7, 8]
     cpu_manager.complete_store(to_keys([6, 7, 8]), _EMPTY_REQ_CTX)
 
-    # touch [2, 6, 7] (move to end of LRU order)
-    cpu_manager.touch(to_keys([2, 6, 7]), _EMPTY_REQ_CTX)
+    # reuse [2, 6, 7] (move to end of LRU order)
+    reuse_keys(cpu_manager, to_keys([2, 6, 7]))
 
-    # prepare store [7, 9] -> evicts [8] (oldest following previous touch)
+    # prepare store [7, 9] -> evicts [8] (oldest following previous reuse)
     prepare_store_output = cpu_manager.prepare_store(to_keys([9]), _EMPTY_REQ_CTX)
     verify_store_output(
         prepare_store_output,
@@ -656,8 +678,8 @@ def test_t1_to_t2_promotion(self):
         assert to_keys([1])[0] in arc_policy.t1
         assert to_keys([1])[0] not in arc_policy.t2
 
-        # touch chunk 1 (simulate second access)
-        cpu_manager.touch(to_keys([1]), _EMPTY_REQ_CTX)
+        # reuse chunk 1 (second access)
+        reuse_keys(cpu_manager, to_keys([1]))
 
         # chunk 1 should now be in T2 (frequent)
         assert to_keys([1])[0] not in arc_policy.t1
@@ -723,9 +745,9 @@ def test_adaptive_target(self):
         # chunk 1 should be in B1 (ghost list)
         assert to_keys([1])[0] in arc_policy.b1
 
-        # touch chunk 1 (cache miss, but in B1)
+        # store chunk 1 again (cache miss, but in B1)
         # this should increase target_t1_size (favor recency)
-        cpu_manager.touch(to_keys([1]), _EMPTY_REQ_CTX)
+        cpu_manager.prepare_store(to_keys([1]), make_req_context("ghost-hit"))
 
         # target should have increased
         assert arc_policy.target_t1_size > initial_target
@@ -740,8 +762,8 @@ def test_t1_t2_eviction_policy(self):
         cpu_manager.prepare_store(to_keys([1, 2, 3, 4]), _EMPTY_REQ_CTX)
         cpu_manager.complete_store(to_keys([1, 2, 3, 4]), _EMPTY_REQ_CTX)
 
-        # promote chunks 3, 4 to T2 by touching them
-        cpu_manager.touch(to_keys([3, 4]), _EMPTY_REQ_CTX)
+        # promote chunks 3, 4 to T2 by reusing them
+        reuse_keys(cpu_manager, to_keys([3, 4]))
 
         # now: T1 = {1, 2}, T2 = {3, 4}
         assert len(arc_policy.t1) == 2
@@ -772,7 +794,7 @@ def test_batch_eviction_scans_t1_and_t2_once(self):
         cpu_manager.prepare_store(keys, _EMPTY_REQ_CTX)
         cpu_manager.complete_store(keys, _EMPTY_REQ_CTX)
 
-        cpu_manager.touch(keys[128:], _EMPTY_REQ_CTX)
+        reuse_keys(cpu_manager, keys[128:])
         arc_policy.target_t1_size = 64
 
         protected = {keys[0], keys[2], keys[255]}
@@ -816,7 +838,7 @@ def test_batch_eviction_falls_back_after_t1_iterator_exhausted(self):
         cpu_manager.prepare_store(keys, _EMPTY_REQ_CTX)
         cpu_manager.complete_store(keys, _EMPTY_REQ_CTX)
 
-        cpu_manager.touch(keys[6:], _EMPTY_REQ_CTX)
+        reuse_keys(cpu_manager, keys[6:])
         arc_policy.target_t1_size = 4
 
         t1_order = list(arc_policy.t1)
@@ -868,9 +890,9 @@ def test_store_falls_back_to_t1_when_t2_cannot_satisfy_eviction(
             manager.complete_store(to_keys(keys), _EMPTY_REQ_CTX)
 
         # Ghost hits raise the target to 3; promotion leaves only 2 entries in T1.
-        manager.touch(to_keys([1, 2]), _EMPTY_REQ_CTX)
-        manager.touch(to_keys([1]), _EMPTY_REQ_CTX)
-        manager.touch(to_keys([5, 6]), _EMPTY_REQ_CTX)
+        policy.on_store_miss(to_keys([1, 2]), _EMPTY_REQ_CTX)
+        policy.on_store_miss(to_keys([1]), _EMPTY_REQ_CTX)
+        reuse_keys(manager, to_keys([5, 6]))
         assert policy.target_t1_size == 3
         assert list(policy.t1) == to_keys([3, 4])
         assert list(policy.t2) == to_keys([6, 5])
@@ -937,8 +959,8 @@ def test_ghost_list_bounds(self):
         assert len(arc_policy.b1) <= arc_policy.cache_capacity
         assert len(arc_policy.b2) <= arc_policy.cache_capacity
 
-    def test_touch_ordering(self):
-        """Tests that touch() correctly updates access patterns.
+    def test_reuse_ordering(self):
+        """Tests that request reuse correctly updates access patterns.
         Similar to LRU test but verifies T1/T2 ordering.
         """
         cpu_manager, arc_policy = self._make_manager()
@@ -948,11 +970,11 @@ def test_touch_ordering(self):
         cpu_manager.complete_store(to_keys([1, 2, 3, 4]), _EMPTY_REQ_CTX)
 
         # promote 3, 4 to T2
-        cpu_manager.touch(to_keys([3, 4]), _EMPTY_REQ_CTX)
+        reuse_keys(cpu_manager, to_keys([3, 4]))
 
         # T1 = {1, 2}, T2 = {3, 4}
-        # touch [1, 3, 4] - should promote 1 to T2, and move 3,4 to end of T2
-        cpu_manager.touch(to_keys([1, 3, 4]), _EMPTY_REQ_CTX)
+        # reuse [1, 3, 4] - should promote 1 to T2, and move 3,4 to end of T2
+        reuse_keys(cpu_manager, to_keys([1, 3, 4]))
 
         # T1 = {2}, T2 = {1, 3, 4} (in that order, with 4 most recent)
         assert len(arc_policy.t1) == 1
@@ -1016,7 +1038,7 @@ def test_full_scenario(self):
         cpu_manager.complete_store(to_keys([3, 4, 5]), _EMPTY_REQ_CTX)
 
         # promote some chunks to T2
-        cpu_manager.touch(to_keys([2, 3]), _EMPTY_REQ_CTX)
+        reuse_keys(cpu_manager, to_keys([2, 3]))
 
         # T1 has {4, 5}, T2 has {2, 3}
         assert len(arc_policy.t1) == 2
@@ -1258,29 +1280,6 @@ def spy_evict(*args, **kwargs):
     assert manager._num_evictable_cache_chunks == 4
 
 
-def test_touch_forwards_req_context_to_policy(monkeypatch):
-    """Regression: CPUOffloadingManager.touch forwards ReqContext to policy."""
-    manager = make_cpu_manager(num_chunks=4, cache_policy="lru")
-    received = []
-
-    def spy_touch(keys: Iterable[OffloadKey], req_context: ReqContext) -> None:
-        received.append((list(keys), req_context))
-
-    monkeypatch.setattr(manager._policy, "touch", spy_touch)
-
-    keys = to_keys([1, 2])
-    ctx = make_req_context(
-        req_id="test-req",
-        kv_transfer_params={"test_param": "test_value"},
-    )
-
-    manager.touch(keys, ctx)
-
-    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,
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 04eeea7d8020..c784178ae2b3 100644
--- a/tests/v1/kv_offload/cpu/test_shared_offload_region.py
+++ b/tests/v1/kv_offload/cpu/test_shared_offload_region.py
@@ -7,6 +7,7 @@
 import errno
 import mmap
 import os
+import tempfile
 import threading
 import time
 import uuid
@@ -660,6 +661,27 @@ def test_segmented_region_maps_files_back_to_back(iid, monkeypatch):
     """Rows split over three shm files mapped back to back: each worker sees
     the other's writes in every file, and the creator removes every file."""
     monkeypatch.setattr(region_module, "_MIN_SEGMENT_BYTES", PAGE_SIZE)
+    # The reservation must not need a writable temp directory.
+    monkeypatch.setattr(tempfile, "tempdir", "/nonexistent")
+    map_segments = SharedOffloadRegion._map_segments
+
+    def check_reservation(region):
+        base = ctypes.addressof(ctypes.c_char.from_buffer(region.mmap_obj))
+        in_reservation = False
+        with open("/proc/self/smaps") as smaps:
+            for line in smaps:
+                fields = line.split()
+                if "-" in fields[0]:
+                    start, end = (int(value, 16) for value in fields[0].split("-"))
+                    in_reservation = start <= base < end
+                elif in_reservation and fields[0] == "VmFlags:":
+                    assert "ac" not in fields[1:]
+                    break
+            else:
+                pytest.fail("Reservation mapping was absent from smaps")
+        map_segments(region)
+
+    monkeypatch.setattr(SharedOffloadRegion, "_map_segments", check_reservation)
     regions = [
         SharedOffloadRegion(
             engine_id=iid,
@@ -691,6 +713,50 @@ def test_segmented_region_maps_files_back_to_back(iid, monkeypatch):
     assert not any(os.path.exists(path) for path, _, _ in segments)
 
 
+@pytest.mark.parametrize("num_segments", [1, 3])
+def test_single_chunk_segments_preserve_shared_canonical_views(
+    iid, monkeypatch, num_segments
+):
+    """Backing boundaries must not change flat HiSparse views or peer writes."""
+    monkeypatch.setattr(region_module, "_MIN_SEGMENT_BYTES", PAGE_SIZE)
+    huge_page = 2 << 20
+    pool_bytes = 5 * huge_page + PAGE_SIZE
+    regions = [
+        SharedOffloadRegion(
+            engine_id=iid,
+            num_chunks=1,
+            rank=0,
+            kv_bytes_per_chunk=pool_bytes,
+            cpu_page_size=pool_bytes,
+            num_segments=num_segments,
+        )
+        for _ in range(2)
+    ]
+    paths = [path for path, _, _ in regions[0]._segments]
+    try:
+        assert len(paths) == num_segments
+        segments = regions[0]._segments
+        assert all(offset % huge_page == 0 for _, offset, _ in segments)
+        assert sum(length for _, _, length in segments) == pool_bytes
+        if num_segments > 1:
+            assert [length for _, _, length in segments] == [
+                2 * huge_page,
+                2 * huge_page,
+                huge_page + PAGE_SIZE,
+            ]
+        a = regions[0].create_next_canonical_view(pool_bytes)
+        b = regions[1].create_next_canonical_view(pool_bytes)
+        a.fill_(17)
+        assert (b == 17).all()
+        b[2 * huge_page - 1 : 2 * huge_page + 1] = 29
+        assert (a[2 * huge_page - 1 : 2 * huge_page + 1] == 29).all()
+        del a, b
+    finally:
+        for region in regions:
+            region.cleanup()
+    assert not any(os.path.exists(path) for path in paths)
+
+
 # ---------------------------------------------------------------------------
 # Multi-worker race — concurrent construction
 # ---------------------------------------------------------------------------
diff --git a/tests/v1/kv_offload/test_factory.py b/tests/v1/kv_offload/test_factory.py
index ae9cbedf42c1..e423236df0fd 100644
--- a/tests/v1/kv_offload/test_factory.py
+++ b/tests/v1/kv_offload/test_factory.py
@@ -765,3 +765,64 @@ def test_build_metric_definitions_returns_counter_at_threshold():
     metrics = spec_cls.build_metric_definitions(extra_config)
 
     assert CPUOffloadingMetrics.STORES_SKIPPED in metrics
+
+
+@pytest.mark.parametrize(
+    ("secondary_tiers", "spanning_group", "rejected"),
+    [
+        ([], "tp", False),
+        ([{"type": "example"}], None, False),
+        ([{"type": "example"}], "tp", True),
+        ([{"type": "example"}], "pp", True),
+        ([{"type": "example"}], "pcp", True),
+    ],
+)
+def test_tiering_spec_rejects_secondary_tiers_across_nodes(
+    monkeypatch, secondary_tiers, spanning_group, rejected
+):
+    import vllm.distributed.parallel_state as parallel_state
+    import vllm.v1.kv_offload.tiering.spec as tiering_spec_module
+
+    spec = _create_spec(
+        spec_name="TieringOffloadingSpec",
+        worker_kv_bytes_per_block=4096,
+        rank=1,
+        world_size=4,
+        tp_size=2,
+        pp_size=2,
+        extra_config={"secondary_tiers": secondary_tiers},
+    )
+    assert isinstance(spec, TieringOffloadingSpec)
+
+    groups = {
+        name: MagicMock(world_size=2, cpu_group=name) for name in ("tp", "pp", "pcp")
+    }
+    monkeypatch.setattr(parallel_state, "get_tp_group", lambda: groups["tp"])
+    monkeypatch.setattr(parallel_state, "get_pp_group", lambda: groups["pp"])
+    monkeypatch.setattr(parallel_state, "get_pcp_group", lambda: groups["pcp"])
+    queried: list[str] = []
+
+    def same_node_as(group, source_rank):
+        queried.append(group)
+        return [True, group != spanning_group]
+
+    monkeypatch.setattr(parallel_state, "in_the_same_node_as", same_node_as)
+    region_ctor = MagicMock()
+    worker_ctor = MagicMock()
+    monkeypatch.setattr(tiering_spec_module, "SharedOffloadRegion", region_ctor)
+    monkeypatch.setattr(tiering_spec_module, "CPUOffloadingWorker", worker_ctor)
+
+    if rejected:
+        with pytest.raises(NotImplementedError, match="span several nodes"):
+            spec.create_worker(MagicMock())
+        # Rejected before any CPU memory is allocated.
+        region_ctor.assert_not_called()
+        worker_ctor.assert_not_called()
+    else:
+        spec.create_worker(MagicMock())
+        region_ctor.assert_called_once()
+        worker_ctor.assert_called_once()
+    # Without secondary tiers each node's region serves its own workers. With
+    # them every group is queried, even after one spans nodes, since each
+    # query is a collective that the group's other ranks also join.
+    assert sorted(queried) == (["pcp", "pp", "tp"] if secondary_tiers else [])
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 46822f369030..f3d60e8079db 100644
--- a/tests/v1/kv_offload/tiering/p2p/test_data_transport.py
+++ b/tests/v1/kv_offload/tiering/p2p/test_data_transport.py
@@ -909,6 +909,77 @@ def test_raising_state_check_warns_once(self, monkeypatch):
             transport.poll(peer_id="peer:1")
         assert sum("check_xfer_state failed" in w for w in warnings) == 1
 
+    def test_remove_remote_peer_retires_running_transfer(self):
+        """A WRITE still running at peer removal keeps its handle and the
+        peer's registration until it stops, and leaves every poll() so a
+        reconnected peer under the same id never sees it."""
+        transport = self._make_transport()
+        agent = transport._agent
+        transport.add_remote_peer("peer:1", b"meta", 0x1000, 8, 1024)
+        tid = transport.write_blocks("peer:1", [0], [1])
+
+        transport.remove_remote_peer("peer:1")
+        agent.release_dlist_handle.assert_not_called()
+        agent.remove_remote_agent.assert_not_called()
+        transport.add_remote_peer("peer:1", b"meta", 0x1000, 8, 1024)
+        agent.check_xfer_state.return_value = "DONE"
+        assert not transport.poll(peer_id="peer:1").done
+        assert not transport.poll().done
+        assert tid in transport._inflight
+
+        assert transport.poll_retired() == [tid]
+        assert tid not in transport._inflight
+        agent.release_xfer_handle.assert_called_once()
+        # The old registration only; the new one stays registered.
+        agent.release_dlist_handle.assert_called_once()
+        agent.remove_remote_agent.assert_called_once()
+        assert "peer:1" in transport._remote_dlists
+
+    @pytest.mark.parametrize("state", ["PROC", "PEND"])
+    def test_poll_retired_keeps_running_transfer(self, state):
+        transport = self._make_transport()
+        transport.add_remote_peer("peer:1", b"meta", 0x1000, 8, 1024)
+        tid = transport.write_blocks("peer:1", [0], [1])
+        transport.remove_remote_peer("peer:1")
+        transport._agent.check_xfer_state.return_value = state
+
+        assert transport.poll_retired() == []
+        assert tid in transport._inflight
+        transport._agent.release_xfer_handle.assert_not_called()
+
+    def test_poll_retired_state_check_raise_is_not_terminal(self):
+        """Only a remote disconnect proves a retired WRITE stopped."""
+        transport = self._make_transport()
+        transport.add_remote_peer("peer:1", b"meta", 0x1000, 8, 1024)
+        tid = transport.write_blocks("peer:1", [0], [1])
+        transport.remove_remote_peer("peer:1")
+        transport._agent.check_xfer_state.side_effect = nixlNotFoundError("gone")
+        assert transport.poll_retired() == []
+        transport._agent.check_xfer_state.side_effect = _remote_disconnect()
+        assert transport.poll_retired() == [tid]
+        assert "peer:1" not in transport._peer_faults
+
+    def test_force_release_retired_releases_registration(self):
+        transport = self._make_transport()
+        agent = transport._agent
+        transport.add_remote_peer("peer:1", b"meta", 0x1000, 8, 1024)
+        tid = transport.write_blocks("peer:1", [0], [1])
+        assert tid is not None
+        transport.remove_remote_peer("peer:1")
+
+        transport.force_release_retired([tid])
+
+        assert tid not in transport._inflight
+        agent.release_xfer_handle.assert_called_once()
+        agent.remove_remote_agent.assert_called_once()
+
+    def test_remove_remote_peer_without_inflight_releases_at_once(self):
+        transport = self._make_transport()
+        transport.add_remote_peer("peer:1", b"meta", 0x1000, 8, 1024)
+        transport.remove_remote_peer("peer:1")
+        transport._agent.release_dlist_handle.assert_called_once()
+        transport._agent.remove_remote_agent.assert_called_once()
+
     def test_remove_remote_peer_clears_the_fault(self):
         transport = self._make_transport()
         for _ in range(3):
diff --git a/tests/v1/kv_offload/tiering/p2p/test_manager.py b/tests/v1/kv_offload/tiering/p2p/test_manager.py
index c47010b70790..238a03d36339 100644
--- a/tests/v1/kv_offload/tiering/p2p/test_manager.py
+++ b/tests/v1/kv_offload/tiering/p2p/test_manager.py
@@ -8,6 +8,8 @@
 
 from __future__ import annotations
 
+import mmap
+import threading
 import time
 import uuid
 from concurrent.futures import Future
@@ -16,6 +18,7 @@
 import numpy as np
 import pytest
 
+from tests.v1.kv_offload.conftest import make_offloading_config
 from tests.v1.kv_offload.tiering.test_tiering_offloading import _mock_mmap_region
 from vllm.utils.hashing import sha256
 from vllm.v1.core.kv_cache_utils import DEFAULT_NONE_HASH_SEED, init_none_hash
@@ -26,6 +29,8 @@
     ReqContext,
     ScheduleEndContext,
 )
+from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion
+from vllm.v1.kv_offload.cpu.spec import CPUOffloadingSpec
 from vllm.v1.kv_offload.tiering.base import (
     JobResult,
     TieringOffloadingMetrics,
@@ -39,6 +44,7 @@
 from vllm.v1.kv_offload.tiering.p2p import manager as manager_module
 from vllm.v1.kv_offload.tiering.p2p.manager import (
     _RECONNECT_MAX_BACKOFF_S,
+    _RETAINED_HOLD_CAP_S,
     _UNBOUND_STORE_TIMEOUT_S,
     P2PSecondaryTierManager,
     _annotate_req_context,
@@ -50,7 +56,12 @@
     StoreResult,
 )
 from vllm.v1.kv_offload.tiering.p2p.session.client import _LOOKUP_TIMEOUT_S
-from vllm.v1.kv_offload.tiering.p2p.session.protocol import TYPE_KEY, DisconnectMsg
+from vllm.v1.kv_offload.tiering.p2p.session.protocol import (
+    TYPE_KEY,
+    ConnectMsg,
+    DisconnectMsg,
+)
+from vllm.v1.kv_offload.tiering.p2p.session.server import StoreDrainResult
 
 # ---------------------------------------------------------------------------
 # Helpers
@@ -139,9 +150,43 @@ def _make_manager() -> P2PSecondaryTierManager:
     mgr._reaped_stores = {}
     mgr._unbound_store_timeout_s = _UNBOUND_STORE_TIMEOUT_S
     mgr._failed_serve_ctxs = []
+    mgr._retained = {}
+    mgr._retained_force_released = 0
     return mgr
 
 
+def _attach_fake_transports(mgr: P2PSecondaryTierManager) -> None:
+    """Give a stubbed manager inert control/data transports.
+
+    The shipped configuration runs the tier-poll thread, which drives
+    _poll_once() -> _control.poll()/session.poll() and shutdown()'s
+    _control.close()/_data.close(). A stub built with __new__ has none.
+    """
+
+    class _Control:
+        def poll(self):
+            return []
+
+        def close(self) -> None:
+            pass
+
+    class _Data:
+        def remove_remote_peer(self, pid):
+            pass
+
+        def cancel(self, ids, mode="immediate"):
+            return []
+
+        def poll(self, peer_id=None):
+            return SimpleNamespace(done=[], failed=[])
+
+        def close(self) -> None:
+            pass
+
+    mgr._control = _Control()
+    mgr._data = _Data()
+
+
 def _init_offloading_spec() -> SimpleNamespace:
     """Minimal offloading_spec for driving the real __init__."""
     return SimpleNamespace(
@@ -602,12 +647,98 @@ def capture_finished_jobs():
         assert len(observations[("1:p2p", "success")]) == 1
 
 
+def test_timed_out_load_holds_chunks_until_abort_ack(monkeypatch):
+    """A timed-out P2P load fails at once but keeps its primary chunk out of
+    the free pool: the peer may still be writing it. The AbortAck frees it."""
+    from tests.v1.kv_offload.tiering.p2p.test_sessions import (
+        FakeConnection,
+        FakeDataTransport,
+        _activate,
+    )
+    from vllm.v1.kv_offload.cpu.common import CPULoadStoreSpec
+    from vllm.v1.kv_offload.tiering.p2p.session import P2PSession
+    from vllm.v1.kv_offload.tiering.p2p.session import client as client_module
+    from vllm.v1.kv_offload.tiering.p2p.session.client import _PD_LOAD_TIMEOUT_S
+    from vllm.v1.kv_offload.tiering.p2p.session.protocol import AbortAckMsg
+
+    now = [1000.0]
+    monkeypatch.setattr(client_module.time, "monotonic", lambda: now[0])
+    peer = "10.0.0.1:8000"
+    p2p_tier = _make_manager()
+    p2p_tier._control = SimpleNamespace(poll=lambda: [])  # type: ignore[assignment]
+    region = _mock_mmap_region(4)
+    primary = CPUPrimaryTierOffloadingManager(num_chunks=4, mmap_region=region)
+    tiering = TieringOffloadingManager(
+        primary_tier=primary,
+        secondary_tiers=[p2p_tier],
+        tier_poll_interval_s=0.0,
+    )
+    conn = FakeConnection(peer_id=peer)
+    session = P2PSession(
+        peer_id=peer,
+        local_id=p2p_tier._local_id,
+        transport=FakeDataTransport(),
+        local_block_len=4096,
+        local_hash_seed="0",
+        conn=conn,
+    )
+    _activate(session, conn, peer_id=peer)
+    p2p_tier._sessions[peer] = session
+
+    key = OffloadKey(b"remote-block")
+    ctx = _req_context(_remote_prefiller_kv_params())
+    prepared = primary.prepare_write([key], ctx)
+    assert prepared is not None
+    assert isinstance(prepared.store_spec, CPULoadStoreSpec)
+    job = TransferJob(
+        job_id=7,
+        keys=[key],
+        chunk_ids=prepared.store_spec.chunk_ids,
+        is_promotion=True,
+        req_context=ctx,
+    )
+    tiering._register_job(job, tier_idx=0)
+    p2p_tier.submit_load(job)
+    free_before = primary._get_num_free_chunks()
+
+    now[0] += _PD_LOAD_TIMEOUT_S + 1.0
+    tiering._process_finished_jobs()
+    assert 7 not in tiering._jobs
+    assert primary.lookup(key, ctx) is LookupResult.MISS
+    assert primary.num_deferred_release_chunks == 1
+    assert primary._get_num_free_chunks() == free_before
+
+    conn.enqueue(
+        {
+            TYPE_KEY: AbortAckMsg.TYPE,
+            AbortAckMsg.KV_REQUEST_ID: "req-1",
+            AbortAckMsg.ROUND_SEQ: 0,
+        }
+    )
+    tiering._process_finished_jobs()
+    assert primary.num_deferred_release_chunks == 0
+    assert primary._get_num_free_chunks() == free_before + 1
+
+
 # ---------------------------------------------------------------------------
 # Tests for on_request_finished
 # ---------------------------------------------------------------------------
 
 
 class TestOnRequestFinished:
+    def test_finished_inflight_load_does_not_leak_failed_id(self):
+        mgr, peer = _build_paired_managers()
+        params = _remote_prefiller_kv_params("B", 2, kv_request_id="req-finish")
+        ctx = _req_context(params)
+        mgr.on_new_request(ctx)
+        job = _job_metadata(7, kv_params=params)
+        mgr.submit_load(job)
+        mgr.on_request_finished(ctx)
+        results = list(mgr.get_finished_jobs())
+        assert len(results) == 1 and results[0].job_id == 7
+        assert not results[0].success
+        assert mgr._failed_req_ids == set()
+
     def _make_with_failed(self) -> P2PSecondaryTierManager:
         mgr = _make_manager()
         mgr._failed_req_ids = {"req-1"}
@@ -727,6 +858,7 @@ def __init__(
         close_stores: list[int] | None = None,
         close_failed_serves: list[ReqContext] | None = None,
         last_activity_at: float | None = None,
+        close_retained: dict[int, frozenset[int]] | None = None,
     ) -> None:
         self.peer_id = peer_id
         self.alive = alive
@@ -741,6 +873,7 @@ def __init__(
         self._close_req_ids = close_req_ids or []
         self._close_stores = close_stores or []
         self._close_failed_serves = close_failed_serves or []
+        self._close_retained = close_retained or {}
         self.requests: list[tuple[int, str]] = []
         self.stores_added: list[tuple[str, list, object, int]] = []
         self.attached: list[object] = []
@@ -769,6 +902,9 @@ def poll(self):
         self._new_fetch_ids = []
         return result
 
+    def drain_stores(self, kv_request_ids=None):
+        return StoreDrainResult([], set())
+
     def request_blocks(self, job_id, kv_request_id, keys, block_ids):
         self.requests.append((job_id, kv_request_id))
 
@@ -788,6 +924,7 @@ def close(self):
             failed_req_ids=self._close_req_ids,
             failed_stores=self._close_stores,
             failed_serves=self._close_failed_serves,
+            retained=self._close_retained,
         )
 
 
@@ -923,25 +1060,175 @@ def test_submit_store_parks_unbound_batch(self):
 
 
 class TestHasPendingWork:
-    """has_pending_work() must always return True so the engine keeps
-    ticking the offload pipeline — that's the only thread driving
-    _control.poll() (incoming peer connects) and session.poll()
-    (incoming fetch messages on existing sessions)."""
-
-    def test_returns_true_unconditionally_to_keep_engine_ticking(self):
-        """Even with no sessions and no jobs, has_pending_work() returns
-        True so the engine keeps calling get_finished_jobs(), which is
-        what drives _control.poll() for inbound peer connects."""
+    """has_pending_work() reports peer-originated work that still needs a step.
+
+    Idle peer polling is left to the tier-poll thread, so a quiescent tier must
+    report False and let pause and graceful shutdown observe an idle engine.
+    The tiering manager adds the keep-alive itself when no thread services the
+    tier.
+    """
+
+    def test_idle_manager_reports_no_pending_work(self):
+        """With no sessions and no jobs, the tier is idle."""
+        mgr = _make_manager()
+        assert mgr.has_pending_work() is False
+
+    def test_connected_idle_session_is_not_pending_work(self):
+        """An idle connected session is not work: polling it is the
+        tier-poll thread's job, not a reason to keep the engine stepping."""
+        mgr = _make_manager()
+        mgr._sessions["peer:1"] = _FakeSession(peer_id="peer:1", connected=True)
+        assert mgr.has_pending_work() is False
+
+    def test_session_with_outstanding_transfer_is_pending_work(self):
+        """A session with an in-flight load keeps the engine stepping until
+        the transfer is drained."""
         mgr = _make_manager()
+        session = _FakeSession(peer_id="peer:1", connected=True)
+        session._client._inbound[1] = object()
+        mgr._sessions["peer:1"] = session
         assert mgr.has_pending_work() is True
 
-    def test_returns_true_even_when_sessions_present(self):
-        """The result is the same regardless of session state — there is
-        no 'idle' branch."""
+    def test_finished_jobs_awaiting_handoff_is_pending_work(self):
+        """Results not yet returned by get_finished_jobs() need a step."""
         mgr = _make_manager()
-        mgr._sessions["peer:1"] = _FakeSession(peer_id="peer:1")
+        mgr._finished_jobs.append(JobResult(job_id=1, success=True))
+        assert mgr.has_pending_work() is True
+
+    def test_parked_unbound_store_is_pending_work(self):
+        """A parked store still pins primary slots until it is fetched or
+        reaped, so it keeps the engine stepping."""
+        from vllm.v1.kv_offload.tiering.p2p.manager import _UnboundStoreBatch
+
+        mgr = _make_manager()
+        mgr._unbound_stores["req-1"] = [
+            _UnboundStoreBatch(job_id=1, keys=[OffloadKey(b"k")], block_ids=[0])
+        ]
+        assert mgr.has_pending_work() is True
+
+    def test_failed_serve_ctx_awaiting_release_is_pending_work(self):
+        """A synthetic lookup ctx still owing parent.on_request_finished must
+        be flushed on the next serve_external_requests()."""
+        mgr = _make_manager()
+        mgr._failed_serve_ctxs.append(_req_context(kv_params=None))
         assert mgr.has_pending_work() is True
 
+    @pytest.mark.cpu_test
+    def test_B2_idle_manager_does_not_keep_engine_busy(self):
+        """A quiescent P2P tier must let pause and shutdown observe idle.
+
+        The shipped configuration runs the tier-poll thread, which services
+        idle peers between steps; the tier must not keep EngineCore.has_work()
+        True just for that.
+        """
+        from unittest.mock import MagicMock
+
+        from vllm.distributed.kv_transfer.kv_connector.v1.offloading.scheduler import (
+            OffloadingConnectorScheduler,
+        )
+        from vllm.distributed.kv_transfer.kv_connector.v1.offloading_connector import (
+            OffloadingConnector,
+        )
+        from vllm.v1.core.sched.interface import PauseState
+        from vllm.v1.core.sched.scheduler import Scheduler
+        from vllm.v1.engine.core import EngineCoreProc
+        from vllm.v1.kv_offload.tiering.manager import (
+            CPUPrimaryTierOffloadingManager,
+            TieringOffloadingManager,
+        )
+
+        p2p = _make_manager()
+        _attach_fake_transports(p2p)
+        mmap_region = MagicMock()
+        view = memoryview(np.zeros((1, 16), dtype=np.int8))
+        mmap_region.create_kv_memoryview.return_value = view
+        manager = TieringOffloadingManager(
+            CPUPrimaryTierOffloadingManager(1, mmap_region),
+            [p2p],
+        )
+        try:
+            connector_scheduler = object.__new__(OffloadingConnectorScheduler)
+            connector_scheduler.manager = manager
+            connector_scheduler._jobs = {}
+            connector_scheduler._reset_fence = {}
+            connector = object.__new__(OffloadingConnector)
+            connector.connector_scheduler = connector_scheduler
+            scheduler = object.__new__(Scheduler)
+            scheduler._pause_state = PauseState.PAUSED_NEW
+            scheduler.running = []
+            scheduler.waiting = []
+            scheduler.kv_holding_waiting = []
+            scheduler.num_waiting_for_streaming_input = 0
+            scheduler.finished_req_ids = set()
+            scheduler.requests = {}
+            scheduler.connector = connector
+            scheduler.ec_connector = None
+            core = object.__new__(EngineCoreProc)
+            core.scheduler = scheduler
+            core.engines_running = False
+            core.batch_queue = None
+
+            assert core._pause_complete() is True
+        finally:
+            manager.shutdown()
+
+    @pytest.mark.cpu_test
+    def test_B2_no_poll_thread_keeps_engine_busy(self):
+        """With the tier-poll thread disabled the engine tick is the only
+        driver, so the tiering manager must keep the engine stepping."""
+        from unittest.mock import MagicMock
+
+        from vllm.distributed.kv_transfer.kv_connector.v1.offloading.scheduler import (
+            OffloadingConnectorScheduler,
+        )
+        from vllm.distributed.kv_transfer.kv_connector.v1.offloading_connector import (
+            OffloadingConnector,
+        )
+        from vllm.v1.core.sched.interface import PauseState
+        from vllm.v1.core.sched.scheduler import Scheduler
+        from vllm.v1.engine.core import EngineCoreProc
+        from vllm.v1.kv_offload.tiering.manager import (
+            CPUPrimaryTierOffloadingManager,
+            TieringOffloadingManager,
+        )
+
+        p2p = _make_manager()
+        _attach_fake_transports(p2p)
+        mmap_region = MagicMock()
+        view = memoryview(np.zeros((1, 16), dtype=np.int8))
+        mmap_region.create_kv_memoryview.return_value = view
+        manager = TieringOffloadingManager(
+            CPUPrimaryTierOffloadingManager(1, mmap_region),
+            [p2p],
+            tier_poll_interval_s=0,
+        )
+        try:
+            assert manager._tier_poll_thread is None
+            connector_scheduler = object.__new__(OffloadingConnectorScheduler)
+            connector_scheduler.manager = manager
+            connector_scheduler._jobs = {}
+            connector_scheduler._reset_fence = {}
+            connector = object.__new__(OffloadingConnector)
+            connector.connector_scheduler = connector_scheduler
+            scheduler = object.__new__(Scheduler)
+            scheduler._pause_state = PauseState.PAUSED_NEW
+            scheduler.running = []
+            scheduler.waiting = []
+            scheduler.kv_holding_waiting = []
+            scheduler.num_waiting_for_streaming_input = 0
+            scheduler.finished_req_ids = set()
+            scheduler.requests = {}
+            scheduler.connector = connector
+            scheduler.ec_connector = None
+            core = object.__new__(EngineCoreProc)
+            core.scheduler = scheduler
+            core.engines_running = False
+            core.batch_queue = None
+
+            assert core._pause_complete() is False
+        finally:
+            manager.shutdown()
+
 
 # ---------------------------------------------------------------------------
 # Shutdown drain
@@ -1279,6 +1566,113 @@ def test_unanswered_lookup_falls_back_to_local_prefill(self, lookup_clock):
         mgr_a.on_request_finished(ctx)
         assert "req-1" not in mgr_a._failed_req_ids
 
+    def test_lookup_before_handshake_ack_falls_back_to_local_prefill(self):
+        """A probe never waits on a pending handshake: a gone peer would hold
+        it until VLLM_P2P_HANDSHAKE_TIMEOUT_S. The peer gets no lookup."""
+        mgr_a, _ = _build_paired_managers()
+        ctx = _req_context(_remote_kv_source_kv_params("B", 2))
+        mgr_a.on_new_request(ctx)
+        list(mgr_a.get_finished_jobs())
+        assert not mgr_a._sessions["B:2"].ready
+
+        assert mgr_a.lookup(b"key", ctx) == LookupResult.MISS
+        mgr_a.on_schedule_end(ScheduleEndContext(new_req_ids=[], preempted_req_ids=[]))
+        sent = mgr_a._control._inbound_outbound["B:2"]._outbox  # type: ignore[attr-defined]  # _LoopbackControl
+        assert [m[TYPE_KEY] for m in sent] == [ConnectMsg.TYPE]
+
+    def test_failed_registration_drops_the_peer_session(self, monkeypatch):
+        """B registered A and waits for A's ack when A fails to register B.
+
+        B must drop that session: a reconnect from A landing in it would be
+        acked without B's ConnectMsg, so A would look B up without having
+        registered it. Real managers on real ZMQ transports; only the NIXL
+        agent is faked.
+        """
+        from unittest.mock import MagicMock
+
+        import zmq
+
+        from tests.v1.kv_offload.conftest import make_offloading_config
+        from vllm.v1.kv_offload.cpu.spec import CPUOffloadingSpec
+        from vllm.v1.kv_offload.tiering.p2p.data import nixl as nixl_module
+
+        from .test_zmq_transport import _free_port
+
+        agents: list[MagicMock] = []
+
+        def make_agent(name, _config):
+            agent = MagicMock()
+            agent.get_agent_metadata.return_value = name.encode()
+            agent.add_remote_agent.side_effect = lambda metadata: metadata.decode()
+            agents.append(agent)
+            return agent
+
+        monkeypatch.setattr(nixl_module, "_NixlAgent", make_agent)
+        monkeypatch.setattr(nixl_module, "_NixlAgentConfig", MagicMock())
+        warnings: list[str] = []
+        monkeypatch.setattr(
+            manager_module.logger,
+            "warning",
+            lambda msg, *args, **_: warnings.append(msg % args if args else msg),
+        )
+        spec = CPUOffloadingSpec(make_offloading_config())
+        managers: list[P2PSecondaryTierManager] = []
+
+        def build() -> tuple[P2PSecondaryTierManager, str]:
+            last_err: Exception | None = None
+            for _ in range(8):
+                port = _free_port()
+                try:
+                    mgr = P2PSecondaryTierManager(
+                        spec,
+                        memoryview(np.zeros((8, 1024), dtype=np.uint8)),
+                        host="127.0.0.1",
+                        port=port,
+                    )
+                except (zmq.ZMQError, OSError) as e:
+                    last_err = e
+                    continue
+                managers.append(mgr)
+                return mgr, f"127.0.0.1:{port}"
+            assert last_err is not None
+            raise last_err
+
+        def poll_until(condition) -> bool:
+            deadline = time.monotonic() + 5.0
+            while time.monotonic() < deadline:
+                for mgr in managers:
+                    list(mgr.get_finished_jobs())
+                if condition():
+                    return True
+                time.sleep(0.005)
+            return False
+
+        try:
+            mgr_a, peer_a = build()
+            agent_a = agents[-1]
+            agent_a.add_remote_agent.side_effect = RuntimeError(
+                "NIXL_ERR_INVALID_PARAM"
+            )
+            mgr_b, peer_b = build()
+            host, port = peer_b.rsplit(":", 1)
+            mgr_a.on_new_request(
+                ReqContext(
+                    req_id="req-1",
+                    kv_transfer_params=_remote_kv_source_kv_params(host, int(port)),
+                )
+            )
+            assert poll_until(lambda: agent_a.add_remote_agent.called)
+            assert poll_until(lambda: peer_b not in mgr_a._sessions)
+            assert poll_until(lambda: peer_a not in mgr_b._sessions), (
+                "B kept its session for A"
+            )
+            assert any(
+                peer_b in w and w.endswith("data-plane fault #1") for w in warnings
+            )
+        finally:
+            for mgr in managers:
+                mgr.shutdown()
+
     def test_lookup_expiry_allows_zmq_reconnect(self, lookup_clock):
         """The peer must retire the timed-out session before a new handshake."""
         import zmq
@@ -1654,6 +2048,36 @@ def remove_remote_peer(self_inner, pid):
         assert peer_dead not in mgr._sessions
         assert peer_alive in mgr._sessions
 
+    def test_failed_load_carries_its_deferred_release(self):
+        """A session's deferred release reaches the engine's JobResult, so the
+        TieringManager holds the load's chunks until the event is set."""
+        mgr = _make_manager()
+        peer = "10.0.0.2:9000"
+        release = threading.Event()
+        mgr._sessions[peer] = _FakeSession(  # type: ignore[assignment]
+            peer_id=peer,
+            loads=[
+                LoadResult(
+                    job_id=11,
+                    kv_request_id="req-11",
+                    success=False,
+                    deferred_release=release,
+                )
+            ],
+        )
+
+        class _Ctrl:
+            def poll(self_inner):
+                return []
+
+        mgr._control = _Ctrl()  # type: ignore[assignment]
+        mgr._poll_once()
+
+        (result,) = mgr._finished_jobs
+        assert (result.job_id, result.success) == (11, False)
+        assert result.deferred_release is release
+        assert "req-11" in mgr._failed_req_ids
+
     def test_new_fetch_id_binds_and_replays_unbound_batches(self):
         """When session.poll() reports a kv_request_id whose FetchMsg
         arrived this tick, the manager binds it to that session and
@@ -1712,16 +2136,24 @@ def poll(self_inner):
         assert mgr._kv_to_session["req-fast"] is sess
         assert sess.stores_added == []
 
-    def test_failed_load_records_kv_request_id(self):
-        """A LoadResult(success=False) from session.poll() must add its
-        kv_request_id to _failed_req_ids so future lookups return MISS."""
+    @pytest.mark.parametrize("request_finished", [False, True])
+    def test_failed_load_records_kv_request_id(self, request_finished):
+        """An unfinished failed load records its id so future lookups return MISS.
+        Finished requests still complete jobs without retaining their ids."""
         mgr = _make_manager()
         peer = "10.0.0.1:8000"
         sess = _FakeSession(
             peer_id=peer,
             alive=True,
             connected=True,
-            loads=[LoadResult(job_id=5, kv_request_id="req-5", success=False)],
+            loads=[
+                LoadResult(
+                    job_id=5,
+                    kv_request_id="req-5",
+                    success=False,
+                    request_finished=request_finished,
+                )
+            ],
         )
         mgr._sessions[peer] = sess
 
@@ -1734,7 +2166,7 @@ def poll(self_inner):
         mgr._poll_once()
 
         assert mgr._finished_jobs == [JobResult(job_id=5, success=False)]
-        assert "req-5" in mgr._failed_req_ids
+        assert ("req-5" in mgr._failed_req_ids) is not request_finished
 
 
 # ---------------------------------------------------------------------------
@@ -2325,6 +2757,171 @@ def test_disabled_when_zero(self, monkeypatch):
 # ---------------------------------------------------------------------------
 
 
+class _RetiringFakeData:
+    """Fake DataTransport whose retired transfers stop on demand."""
+
+    def __init__(self) -> None:
+        self.stopped: list[int] = []
+        self.force_released: list[int] = []
+        self.removed: list[str] = []
+
+    def remove_remote_peer(self, pid) -> None:
+        self.removed.append(pid)
+
+    def poll_retired(self) -> list[int]:
+        stopped, self.stopped = self.stopped, []
+        return stopped
+
+    def force_release_retired(self, transfer_ids) -> None:
+        self.force_released.extend(transfer_ids)
+
+
+class TestRetainedTransfers:
+    """A reaped session's still-running transfers keep their source pins
+    until the transport reports them stopped, or the hold cap expires."""
+
+    PEER = "10.0.0.1:8000"
+
+    def _reap(self, retained, close_stores=None):
+        mgr = _make_manager()
+        data = _RetiringFakeData()
+        mgr._data = data  # type: ignore[assignment]
+        mgr._sessions[self.PEER] = _FakeSession(  # type: ignore[assignment]
+            peer_id=self.PEER,
+            alive=False,
+            connected=True,
+            close_stores=close_stores,
+            close_retained=retained,
+        )
+        mgr._reap_dead_sessions()
+        return mgr, data
+
+    def test_reap_holds_pins_until_transfer_stops(self):
+        mgr, data = self._reap({7: frozenset({1})}, close_stores=[2])
+        assert mgr._finished_jobs == [JobResult(job_id=2, success=False)]
+        assert data.removed == [self.PEER]
+
+        mgr._poll_retained()
+        assert mgr._finished_jobs == [JobResult(job_id=2, success=False)]
+
+        data.stopped = [7]
+        mgr._poll_retained()
+        assert mgr._finished_jobs[-1] == JobResult(job_id=1, success=False)
+        assert mgr._retained == {}
+
+    def test_shared_job_held_until_last_transfer_stops(self):
+        mgr, data = self._reap({7: frozenset({1}), 8: frozenset({1})})
+        data.stopped = [7]
+        mgr._poll_retained()
+        assert mgr._finished_jobs == []
+        data.stopped = [8]
+        mgr._poll_retained()
+        assert mgr._finished_jobs == [JobResult(job_id=1, success=False)]
+
+    def test_hold_cap_force_releases(self, monkeypatch):
+        now = [1000.0]
+        monkeypatch.setattr(manager_module.time, "monotonic", lambda: now[0])
+        mgr, data = self._reap({7: frozenset({1})})
+
+        now[0] += _RETAINED_HOLD_CAP_S - 1
+        mgr._poll_retained()
+        assert data.force_released == []
+
+        now[0] += 1
+        mgr._poll_retained()
+        assert data.force_released == [7]
+        assert mgr._finished_jobs == [JobResult(job_id=1, success=False)]
+        assert mgr._retained_force_released == 1
+        assert mgr._retained == {}
+
+    def test_idle_eviction_retains_too(self, monkeypatch):
+        monkeypatch.setenv("VLLM_P2P_IDLE_TIMEOUT_S", "10")
+        monkeypatch.setenv("VLLM_P2P_MAX_PEERS", "1")
+        import vllm.envs
+
+        vllm.envs.__dict__.pop("VLLM_P2P_IDLE_TIMEOUT_S", None)
+        vllm.envs.__dict__.pop("VLLM_P2P_MAX_PEERS", None)
+        mgr = _make_manager()
+        mgr._data = _RetiringFakeData()  # type: ignore[assignment]
+        mgr._sessions[self.PEER] = _FakeSession(  # type: ignore[assignment]
+            peer_id=self.PEER,
+            last_activity_at=time.monotonic() - 20,
+            close_retained={7: frozenset({1})},
+        )
+        mgr._reap_idle_sessions()
+        assert self.PEER not in mgr._sessions
+        assert mgr._retained[7][:2] == (self.PEER, frozenset({1}))
+        assert mgr._finished_jobs == []
+
+    def test_shutdown_fails_retained_jobs(self):
+        mgr, _ = self._reap({7: frozenset({1})})
+        mgr._data = _ShutdownFakeData()  # type: ignore[assignment]
+        mgr._control = _ShutdownFakeControl()  # type: ignore[assignment]
+        mgr.shutdown()
+        assert JobResult(job_id=1, success=False) in mgr._finished_jobs
+        assert mgr._retained == {}
+
+    def test_outbound_waits_for_retained_transfers(self):
+        """A new session would load the peer's NIXL agent under the same
+        name while the reaped one's registration still holds it."""
+        mgr, data = self._reap({7: frozenset({1})})
+
+        class NoConnect:
+            def connect(self, peer_id):
+                raise AssertionError("connect() while transfers are retained")
+
+        mgr._control = NoConnect()  # type: ignore[assignment]
+        assert mgr._get_or_create_session(self.PEER) is None
+
+        data.stopped = [7]
+        mgr._poll_retained()
+        assert not mgr._retained_peer(self.PEER)
+
+    def test_inbound_is_told_to_disconnect_while_retained(self):
+        mgr, data = self._reap({7: frozenset({1})})
+
+        class Conn(_RecordingConn):
+            def __init__(self, peer_id: str) -> None:
+                super().__init__(peer_id)
+                self.sent: list[dict] = []
+
+            def send(self, msg: dict) -> None:
+                self.sent.append(msg)
+
+        conn = Conn(self.PEER)
+        mgr._accept_new_peers([conn])  # type: ignore[list-item]
+        assert conn.sent == [{TYPE_KEY: DisconnectMsg.TYPE}]
+        assert conn.close_calls == 1
+        assert self.PEER not in mgr._sessions
+
+        data.stopped = [7]
+        mgr._poll_retained()
+        assert not mgr._retained_peer(self.PEER)
+
+    def test_other_peers_are_not_held(self):
+        mgr, _ = self._reap({7: frozenset({1})})
+        assert not mgr._retained_peer("10.0.0.2:8000")
+
+    def test_drain_jobs_waits_for_retained_transfers(self, monkeypatch):
+        """reset_cache() asserts every job finished once drain_jobs()
+        returns, so a retained transfer's store jobs must be settled."""
+        mgr, data = self._reap({7: frozenset({1})})
+        mgr._control = _DrainCtrl()  # type: ignore[assignment]
+        polls = [0]
+
+        def poll_retired():
+            polls[0] += 1
+            return [7] if polls[0] == 3 else []
+
+        data.poll_retired = poll_retired  # type: ignore[method-assign]
+        monkeypatch.setattr(manager_module.time, "sleep", lambda s: None)
+        mgr.drain_jobs()
+
+        assert polls[0] == 3
+        assert mgr._retained == {}
+        assert JobResult(job_id=1, success=False) in mgr._finished_jobs
+
+
 class TestReconnectBackoff:
     """A peer whose session ended on a data-plane fault reconnects at once the
     first time; repeats back off from 1 s up to _RECONNECT_MAX_BACKOFF_S."""
@@ -2496,3 +3093,347 @@ def peer_lookup(kv_request_id: str) -> LookupResult:
         assert [job.keys for job in promotions] == [[key]]
     finally:
         tiering.shutdown()
+
+
+# ---------------------------------------------------------------------------
+# Idleness while serving peers
+# ---------------------------------------------------------------------------
+
+
+def _producer_with_cpu_hit(
+    key: OffloadKey, tier_poll_interval_s: float | None = None
+) -> tuple[P2PSecondaryTierManager, TieringOffloadingManager]:
+    """A consumer paired with a lone-P2P producer whose CPU tier holds key."""
+    consumer, server = _build_paired_managers()
+    kwargs = {}
+    if tier_poll_interval_s is not None:
+        kwargs["tier_poll_interval_s"] = tier_poll_interval_s
+    tiering = TieringOffloadingManager(
+        primary_tier=CPUPrimaryTierOffloadingManager(
+            num_chunks=4, mmap_region=_mock_mmap_region(4)
+        ),
+        secondary_tiers=[server],
+        **kwargs,
+    )
+    seed_ctx = ReqContext(req_id="seed")
+    with tiering.lock:
+        assert tiering.primary_tier.prepare_store([key], seed_ctx) is not None
+        tiering.primary_tier.complete_store([key], seed_ctx)
+    return consumer, tiering
+
+
+def _peer_hit(
+    consumer: P2PSecondaryTierManager,
+    tiering: TieringOffloadingManager,
+    key: OffloadKey,
+    step_producer: bool,
+) -> ReqContext:
+    """Drive a consumer lookup until the producer pins key as a HIT."""
+    step_end = ScheduleEndContext(new_req_ids=[], preempted_req_ids=[])
+    host, port = "B", 2
+    ctx = ReqContext(
+        req_id="req-peer",
+        kv_transfer_params=_remote_kv_source_kv_params(host, port, "req-peer"),
+    )
+    _annotate_req_context(ctx)
+    consumer.on_new_request(ctx)
+    deadline = time.monotonic() + 5.0
+    while time.monotonic() < deadline:
+        list(consumer.get_finished_jobs())
+        if step_producer:
+            with tiering.lock:
+                tiering.on_schedule_end(step_end)
+        # A probe before the handshake ack falls back to local prefill (MISS).
+        session = consumer._sessions.get(f"{host}:{port}")
+        if session is not None and session.ready:
+            result = consumer.lookup(key, ctx)
+            if result is LookupResult.HIT:
+                return ctx
+            assert result is LookupResult.RETRY
+            consumer.on_schedule_end(step_end)
+        time.sleep(0.001)
+    raise AssertionError("peer lookup did not reach HIT")
+
+
+class TestIdleWhileServingPeers:
+    @pytest.mark.cpu_test
+    def test_pinned_peer_hit_does_not_block_idleness(self):
+        """A lone-P2P producer with a live tier-poll thread serves a whole peer
+        request without an engine step, so its pin must not count as pending
+        work, and the fetch must still complete and release the pin."""
+        key = OffloadKey(b"cpu-resident")
+        consumer, tiering = _producer_with_cpu_hit(key)
+        try:
+            assert tiering._tier_poll_thread is not None
+            ctx = _peer_hit(consumer, tiering, key, step_producer=False)
+            with tiering.lock:
+                assert tiering._jobs
+                assert tiering.has_pending_work() is False
+
+            consumer.submit_load(
+                TransferJob(
+                    job_id=7,
+                    keys=[key],
+                    chunk_ids=np.array([0]),
+                    is_promotion=True,
+                    req_context=ctx,
+                )
+            )
+            results: list[JobResult] = []
+            deadline = time.monotonic() + 5.0
+            while time.monotonic() < deadline:
+                results.extend(consumer.get_finished_jobs())
+                with tiering.lock:
+                    if not tiering._jobs and results:
+                        break
+                time.sleep(0.001)
+            assert JobResult(job_id=7, success=True) in results
+            with tiering.lock:
+                assert not tiering._jobs
+        finally:
+            tiering.shutdown()
+
+    @pytest.mark.cpu_test
+    def test_pinned_peer_hit_without_poll_thread_is_pending(self):
+        """With the tier-poll thread disabled only a step completes the
+        fetch, so the engine must keep stepping."""
+        key = OffloadKey(b"cpu-resident")
+        consumer, tiering = _producer_with_cpu_hit(key, tier_poll_interval_s=0)
+        try:
+            assert tiering._tier_poll_thread is None
+            _peer_hit(consumer, tiering, key, step_producer=True)
+            assert tiering._jobs
+            assert tiering.has_pending_work() is True
+        finally:
+            tiering.shutdown()
+
+
+class _ControlledDrainData(_FakeData):
+    """Hold external transfers and inject peer traffic at an exact poll."""
+
+    def __init__(self, local_id):
+        super().__init__(local_id)
+        self.hold = False
+        self.on_poll = None
+
+    def poll(self, peer_id=None):
+        from vllm.v1.kv_offload.tiering.p2p.data.base import PollResult
+
+        if self.on_poll is not None:
+            self.on_poll()
+        if self.hold:
+            return PollResult(done=[], failed=[])
+        return super().poll(peer_id)
+
+
+@pytest.fixture
+def drain_transports():
+    return [_ControlledDrainData("A:1"), _ControlledDrainData("B:2")]
+
+
+@pytest.fixture
+def drain_peers(monkeypatch, drain_transports):
+    """Real managers and sessions; only the wire transports are in memory."""
+    init_none_hash(sha256)
+    a = _LoopbackControl("A:1")
+    b = _LoopbackControl("B:2")
+    a.pair(b)
+    controls = {"A:1": a, "B:2": b}
+    monkeypatch.setattr(
+        manager_module, "ZmqTransport", lambda local_id, *_: controls[local_id]
+    )
+    transports = iter(drain_transports)
+    monkeypatch.setattr(
+        manager_module, "NixlTransport", lambda *_, **__: next(transports)
+    )
+    spec = CPUOffloadingSpec(make_offloading_config())
+    peers = [
+        P2PSecondaryTierManager(
+            spec, memoryview(bytearray(65536)), host=host, port=port
+        )
+        for host, port in [("A", 1), ("B", 2)]
+    ]
+    try:
+        yield peers
+    finally:
+        for peer in peers:
+            peer.shutdown()
+
+
+@pytest.mark.parametrize("bound", [False, True])
+def test_drain_fails_unfetched_stores_and_rejects_late_fetch(drain_peers, bound):
+    """Reset must release every queued pin without offering stale chunk ids."""
+    consumer, producer = drain_peers
+    ctx = _req_context(_remote_prefiller_kv_params("B", 2))
+    consumer.on_new_request(ctx)
+    for _ in range(3):
+        list(consumer.get_finished_jobs())
+        list(producer.get_finished_jobs())
+
+    key = OffloadKey(b"unfetched")
+    if bound:
+        first = OffloadKey(b"first")
+        consumer.submit_load(_job_metadata(10, [first], [0], ctx.kv_transfer_params))
+        list(producer.get_finished_jobs())
+        producer.submit_store(
+            _job_metadata(11, [first], [0], _remote_decoder_kv_params())
+        )
+        assert list(producer.get_finished_jobs()) == [JobResult(11, True)]
+        assert list(consumer.get_finished_jobs()) == [JobResult(10, True)]
+
+    for job_id in (20, 21):
+        producer.submit_store(
+            _job_metadata(job_id, [key], [job_id - 20], _remote_decoder_kv_params())
+        )
+    producer.drain_jobs()
+    assert sorted(producer.get_finished_jobs(), key=lambda result: result.job_id) == [
+        JobResult(20, False),
+        JobResult(21, False),
+    ]
+    producer.drain_jobs()
+    assert list(producer.get_finished_jobs()) == []
+
+    consumer.submit_load(_job_metadata(30, [key], [1], ctx.kv_transfer_params))
+    list(producer.get_finished_jobs())
+    assert list(consumer.get_finished_jobs()) == [JobResult(30, False)]
+    producer.submit_store(_job_metadata(31, [key], [0], _remote_decoder_kv_params()))
+    assert list(producer.get_finished_jobs()) == [JobResult(31, False)]
+
+
+def test_reset_cache_releases_parked_p2p_slots(drain_peers):
+    """The real tiering reset can reclaim every slot held for an absent decoder."""
+    _, producer = drain_peers
+    region = SharedOffloadRegion(
+        engine_id=str(uuid.uuid4()),
+        num_chunks=8,
+        rank=0,
+        kv_bytes_per_chunk=mmap.PAGESIZE,
+        cpu_page_size=mmap.PAGESIZE,
+    )
+    tiering = TieringOffloadingManager(
+        CPUPrimaryTierOffloadingManager(num_chunks=8, mmap_region=region),
+        [producer],
+        tier_poll_interval_s=0,
+    )
+    ctx = _req_context(_remote_decoder_kv_params())
+    keys = [OffloadKey(b"one"), OffloadKey(b"two")]
+    try:
+        tiering.on_new_request(ctx)
+        store = tiering.prepare_store(keys, ctx)
+        assert store is not None
+        tiering.complete_store(store.keys_to_store, ctx)
+        tiering.on_schedule_end(ScheduleEndContext([], ()))
+        tiering.reset_cache()
+        tiering.on_reset_flushed()
+        assert all(tiering.lookup(key, ctx) is LookupResult.MISS for key in keys)
+        replacement = tiering.prepare_store(
+            [OffloadKey(str(i).encode()) for i in range(8)], ctx
+        )
+        assert replacement is not None
+        assert len(replacement.keys_to_store) == 8
+    finally:
+        tiering.shutdown()
+        region.cleanup()
+
+
+def test_drain_rejects_fetch_arriving_mid_poll(drain_peers, drain_transports):
+    """A fetch arriving during drain must not leave a live producer binding."""
+    consumer, producer = drain_peers
+    ctx = _req_context(_remote_prefiller_kv_params("B", 2, "active"))
+    consumer.on_new_request(ctx)
+    for _ in range(3):
+        list(consumer.get_finished_jobs())
+        list(producer.get_finished_jobs())
+    key = OffloadKey(b"active")
+    consumer.submit_load(_job_metadata(1, [key], [0], ctx.kv_transfer_params))
+    list(producer.get_finished_jobs())
+    producer.submit_store(
+        _job_metadata(2, [key], [0], _remote_decoder_kv_params("active"))
+    )
+    data = drain_transports[1]
+    data.hold = True
+    late_ctx = _req_context(_remote_prefiller_kv_params("B", 2, "late"))
+    consumer.on_new_request(late_ctx)
+    polls = 0
+
+    def peer_traffic():
+        nonlocal polls
+        polls += 1
+        if polls == 1:
+            consumer.submit_load(
+                _job_metadata(3, [key], [1], late_ctx.kv_transfer_params)
+            )
+        elif polls == 3:
+            data.hold = False
+
+    data.on_poll = peer_traffic
+    producer.drain_jobs()
+    assert list(producer.get_finished_jobs()) == [JobResult(2, True)]
+    assert sorted(consumer.get_finished_jobs(), key=lambda result: result.job_id) == [
+        JobResult(1, True),
+        JobResult(3, False),
+    ]
+    producer.submit_store(
+        _job_metadata(4, [key], [1], _remote_decoder_kv_params("late"))
+    )
+    assert list(producer.get_finished_jobs()) == [JobResult(4, False)]
+
+
+@pytest.mark.parametrize("fetch_before_reset", [False, True])
+def test_reset_symmetric_lookup_rounds(drain_peers, fetch_before_reset):
+    """Old HITs fail promptly after reset; fresh lookups can reuse a completed id."""
+    consumer, producer = drain_peers
+    region = SharedOffloadRegion(
+        engine_id=str(uuid.uuid4()),
+        num_chunks=8,
+        rank=0,
+        kv_bytes_per_chunk=mmap.PAGESIZE,
+        cpu_page_size=mmap.PAGESIZE,
+    )
+    tiering = TieringOffloadingManager(
+        CPUPrimaryTierOffloadingManager(num_chunks=8, mmap_region=region),
+        [producer],
+        tier_poll_interval_s=0,
+    )
+    key = OffloadKey(b"symmetric")
+    local_ctx = ReqContext(req_id="local")
+    ctx = _req_context(_remote_kv_source_kv_params("B", 2, "symmetric"))
+    step = ScheduleEndContext([], ())
+
+    def store():
+        output = tiering.prepare_store([key], local_ctx)
+        assert output is not None
+        tiering.complete_store(output.keys_to_store, local_ctx)
+        tiering.on_schedule_end(step)
+
+    def lookup():
+        assert consumer.lookup(key, ctx) is LookupResult.RETRY
+        consumer.on_schedule_end(step)
+        tiering.on_schedule_end(step)
+        list(consumer.get_finished_jobs())
+        assert consumer.lookup(key, ctx) is LookupResult.HIT
+
+    try:
+        tiering.on_new_request(local_ctx)
+        consumer.on_new_request(ctx)
+        for _ in range(3):
+            list(consumer.get_finished_jobs())
+            tiering.on_schedule_end(step)
+        store()
+        lookup()
+        if fetch_before_reset:
+            consumer.submit_load(_job_metadata(1, [key], [0], ctx.kv_transfer_params))
+            tiering.on_schedule_end(step)
+            assert list(consumer.get_finished_jobs()) == [JobResult(1, True)]
+
+        tiering.reset_cache()
+        tiering.on_reset_flushed()
+        if fetch_before_reset:
+            store()
+            lookup()
+        consumer.submit_load(_job_metadata(2, [key], [0], ctx.kv_transfer_params))
+        tiering.on_schedule_end(step)
+        assert list(consumer.get_finished_jobs()) == [JobResult(2, fetch_before_reset)]
+    finally:
+        tiering.shutdown()
+        region.cleanup()
diff --git a/tests/v1/kv_offload/tiering/p2p/test_sessions.py b/tests/v1/kv_offload/tiering/p2p/test_sessions.py
index 7856867d6a96..79984d25093c 100644
--- a/tests/v1/kv_offload/tiering/p2p/test_sessions.py
+++ b/tests/v1/kv_offload/tiering/p2p/test_sessions.py
@@ -326,6 +326,11 @@ def on_request_finished(self, ctx: ReqContext) -> None:
         self.calls.append(("on_request_finished", ctx.req_id))
 
 
+def _released(result: LoadResult) -> bool:
+    assert result.deferred_release is not None
+    return result.deferred_release.is_set()
+
+
 def _make_session(
     conn: FakeConnection | None = None,
     transport: FakeDataTransport | None = None,
@@ -492,6 +497,8 @@ def test_failed_registration_rejects_peer(self):
 
         assert not conn.alive
         assert not any(m[TYPE_KEY] == ConnectAckMsg.TYPE for m in conn._sent)
+        assert conn._sent[-1][TYPE_KEY] == DisconnectMsg.TYPE
+        assert session.data_plane_fault == "peer registration failed: nixl exploded"
 
     def test_failed_ack_send_marks_connection_dead(self):
         """A ConnectAckMsg send that raises marks the connection dead."""
@@ -684,12 +691,18 @@ def test_finish_preserves_active_load_until_terminal_result(self):
         assert aborts[0][AbortFetchMsg.ROUND_SEQ] == 0
 
     def test_finish_abort_ack_emits_one_failure(self):
+        """finish() fails the load at once; the AbortAck only releases it."""
         session, conn, _ = _make_session()
         _activate(session, conn)
         session.request_blocks(
             job_id=8, kv_request_id="req-1", keys=[OffloadKey(b"k")], block_ids=[0]
         )
         session.finish_request("req-1")
+        (failure,) = session.poll().loads
+        assert (failure.job_id, failure.success) == (8, False)
+        assert failure.request_finished
+        assert not _released(failure)
+        assert session._client.has_active_loads is True
 
         conn.enqueue(
             {
@@ -698,28 +711,30 @@ def test_finish_abort_ack_emits_one_failure(self):
                 AbortAckMsg.KV_REQUEST_ID: "req-1",
             }
         )
-        loads = session.poll().loads
-        assert loads == [LoadResult(job_id=8, kv_request_id="req-1", success=False)]
+        assert session.poll().loads == []
+        assert _released(failure)
         assert session._client.has_active_loads is False
         assert "req-1" not in session._client._requests
-        assert session.poll().loads == []
 
     def test_finish_abort_timeout_emits_one_failure(self):
+        """A missed AbortAck after finish() emits no second result: it ends
+        the session, and close() releases the held blocks."""
         session, conn, _ = _make_session()
         _activate(session, conn)
         session.request_blocks(
             job_id=9, kv_request_id="req-1", keys=[OffloadKey(b"k")], block_ids=[0]
         )
         session.finish_request("req-1")
+        (failure,) = session.poll().loads
         _client_load(session, "req-1").aborted_at = (
             time.monotonic() - _ABORT_ACK_TIMEOUT_S - 1.0
         )
 
-        loads = session.poll().loads
-        assert loads == [LoadResult(job_id=9, kv_request_id="req-1", success=False)]
-        assert session._client.has_active_loads is False
-        assert "req-1" not in session._client._requests
         assert session.poll().loads == []
+        assert not session.alive
+        assert not _released(failure)
+        assert session.close().failed_jobs == []
+        assert _released(failure)
 
     def test_late_transfer_done_after_abort_ack_is_ignored(self):
         session, conn, _ = _make_session()
@@ -736,9 +751,10 @@ def test_late_transfer_done_after_abort_ack_is_ignored(self):
                 AbortAckMsg.KV_REQUEST_ID: "req-1",
             }
         )
-        assert session.poll().loads == [
-            LoadResult(job_id=10, kv_request_id="req-1", success=False)
-        ]
+        (failure,) = session.poll().loads
+        assert (failure.job_id, failure.success) == (10, False)
+        assert failure.request_finished
+        assert _released(failure)
 
         conn.enqueue(
             {
@@ -823,10 +839,9 @@ def test_load_timeout_depends_on_mode(self):
         assert _client_load(session, "req-pd").aborted_at is None
 
     def test_load_abort_ack_timeout_surfaces_failure(self):
-        """After load timeout sends AbortFetch, if no AbortAck arrives within
-        _ABORT_ACK_TIMEOUT_S the request is surfaced as failed and removed
-        from _requests — the engine cannot wait forever on a peer that won't
-        ack.
+        """A load timeout fails the load at once, with its blocks held: the
+        peer may still be writing them. If no AbortAck arrives within
+        _ABORT_ACK_TIMEOUT_S, no second result is emitted; close() releases.
         """
         session, conn, _ = _make_session()
         _activate(session, conn)
@@ -837,8 +852,9 @@ def test_load_abort_ack_timeout_surfaces_failure(self):
         _client_load(session, "req-7").submitted_at = (
             time.monotonic() - _PD_LOAD_TIMEOUT_S - 1.0
         )
-        loads = session.poll().loads
-        assert loads == []
+        (failure,) = session.poll().loads
+        assert (failure.job_id, failure.success) == (7, False)
+        assert not _released(failure)
         assert any(
             m.get(TYPE_KEY) == AbortFetchMsg.TYPE
             and m[AbortFetchMsg.KV_REQUEST_ID] == "req-7"
@@ -851,9 +867,46 @@ def test_load_abort_ack_timeout_surfaces_failure(self):
         _client_load(session, "req-7").aborted_at = (
             time.monotonic() - _ABORT_ACK_TIMEOUT_S - 1.0
         )
-        loads = session.poll().loads
-        assert loads == [LoadResult(job_id=7, kv_request_id="req-7", success=False)]
-        assert "req-7" not in session._client._requests
+        assert session.poll().loads == []
+        assert not _released(failure)
+        session.close()
+        assert _released(failure)
+
+    def test_load_abort_ack_timeout_disconnects_peer(self):
+        """A missed AbortAck faults the peer, so later fetches skip it.
+
+        The peer has stopped confirming that its WRITEs ended; keeping the
+        session would make every later load wait out both timeouts again.
+        Loads still in flight on the session fail through close().
+        """
+        session, conn, _ = _make_session()
+        _activate(session, conn)
+        session.request_blocks(
+            job_id=7, kv_request_id="req-7", keys=[OffloadKey(b"k")], block_ids=[0]
+        )
+        session.request_blocks(
+            job_id=8, kv_request_id="req-8", keys=[OffloadKey(b"j")], block_ids=[1]
+        )
+        _client_load(session, "req-7").submitted_at = (
+            time.monotonic() - _PD_LOAD_TIMEOUT_S - 1.0
+        )
+        (failure,) = session.poll().loads
+        assert (failure.job_id, failure.success) == (7, False)
+        assert session.alive
+
+        _client_load(session, "req-7").aborted_at = (
+            time.monotonic() - _ABORT_ACK_TIMEOUT_S - 1.0
+        )
+        assert session.poll().loads == []
+        assert not session.alive
+        assert session.death_reason == "abort ack timed out"
+        assert session.data_plane_fault == "abort ack timed out"
+        assert conn._sent[-1][TYPE_KEY] == DisconnectMsg.TYPE
+
+        result = session.close()
+        assert result.failed_jobs == [8]
+        assert sorted(result.failed_req_ids) == ["req-7", "req-8"]
+        assert _released(failure)
 
     def test_load_abort_ack_clears_request(self):
         """After load timeout sends AbortFetch, an arriving AbortAckMsg from
@@ -867,8 +920,9 @@ def test_load_abort_ack_clears_request(self):
         _client_load(session, "req-8").submitted_at = (
             time.monotonic() - _PD_LOAD_TIMEOUT_S - 1.0
         )
-        # First poll: AbortFetch goes out.
-        session.poll()
+        # First poll: AbortFetch goes out and the load is failed at once.
+        (failure,) = session.poll().loads
+        assert (failure.job_id, failure.success) == (8, False)
         assert _client_load(session, "req-8").aborted_at is not None
 
         # Peer acks the abort.
@@ -879,9 +933,11 @@ def test_load_abort_ack_clears_request(self):
                 AbortAckMsg.KV_REQUEST_ID: "req-8",
             }
         )
-        loads = session.poll().loads
-        assert loads == [LoadResult(job_id=8, kv_request_id="req-8", success=False)]
+        assert session.poll().loads == []
+        assert _released(failure)
         assert "req-8" not in session._client._requests
+        assert session.alive
+        assert session.data_plane_fault is None
 
 
 # ---------------------------------------------------------------------------
@@ -2965,6 +3021,71 @@ def test_pending_close_returns_pending_stores(self):
 
 
 class TestDisconnect:
+    @pytest.mark.parametrize(
+        "completion", ["abort_ack", "transfer_done", "timeout", "close"]
+    )
+    @pytest.mark.parametrize("finished", [False, True])
+    def test_late_failure_preserves_request_finish(
+        self, lookup_clock, completion, finished
+    ):
+        session, conn, _ = _make_session()
+        _activate(session, conn)
+        session.request_blocks(7, "req-7", [OffloadKey(b"key")], [0])
+        if finished:
+            session.finish_request("req-7")
+        if completion == "close":
+            result = session.close()
+            assert result.failed_jobs == [7]
+            assert result.failed_req_ids == ([] if finished else ["req-7"])
+            return
+        loads = []
+        if not finished and completion in ("timeout", "abort_ack"):
+            lookup_clock.now = _PD_LOAD_TIMEOUT_S + 1
+            loads.extend(session.poll().loads)
+        if completion == "timeout":
+            if finished:
+                lookup_clock.now = _ABORT_ACK_TIMEOUT_S + 1
+            else:
+                lookup_clock.now += _ABORT_ACK_TIMEOUT_S + 1
+        else:
+            terminal = AbortAckMsg if completion == "abort_ack" else TransferDoneMsg
+            msg = {
+                TYPE_KEY: terminal.TYPE,
+                TransferDoneMsg.KV_REQUEST_ID: "req-7",
+                TransferDoneMsg.ROUND_SEQ: 0,
+            }
+            if terminal is TransferDoneMsg:
+                msg[TransferDoneMsg.SUCCESS] = False
+            conn.enqueue(msg)
+        loads.extend(session.poll().loads)
+        assert len(loads) == 1
+        assert (loads[0].job_id, loads[0].kv_request_id, loads[0].success) == (
+            7,
+            "req-7",
+            False,
+        )
+        assert loads[0].request_finished is finished
+
+    def test_finish_marks_buffered_result_before_request_id_reuse(self):
+        session, conn, _ = _make_session()
+        _activate(session, conn)
+        session.request_blocks(7, "req-7", [OffloadKey(b"old")], [0])
+        session._client.on_transfer_done("req-7", success=False, round_seq=0)
+        session.finish_request("req-7")
+        session.request_blocks(8, "req-7", [OffloadKey(b"new")], [1])
+        loads = session.poll().loads
+        assert len(loads) == 1 and loads[0].job_id == 7 and not loads[0].success
+        assert loads[0].request_finished
+        conn.enqueue(
+            {
+                TYPE_KEY: TransferDoneMsg.TYPE,
+                TransferDoneMsg.KV_REQUEST_ID: "req-7",
+                TransferDoneMsg.ROUND_SEQ: 0,
+                TransferDoneMsg.SUCCESS: False,
+            }
+        )
+        assert session.poll().loads == [LoadResult(8, "req-7", success=False)]
+
     def test_disconnect_marks_session_dead(self):
         session, conn, _ = _make_session()
         _activate(session, conn)
@@ -2984,6 +3105,43 @@ def test_close_returns_pending_loads_and_stores(self):
         assert set(result.failed_stores) == {10}
         assert result.failed_serves == []
 
+    def test_close_retains_running_transfer(self):
+        """close() must not release a WRITE that may still be landing.
+
+        It cancels in wait mode; a transfer still running is reported as
+        retained with its store job, which stays pinned, and only the jobs
+        no transfer reads are returned as failed.
+        """
+        session, conn, data = _make_session()
+        _activate(session, conn)
+        session.add_stored_blocks("req", [OffloadKey(b"key")], [0], job_id=1)
+        session.add_stored_blocks("idle", [OffloadKey(b"other")], [1], job_id=2)
+        _fetch(conn, "req", [b"key"], [5])
+        session.poll()
+        (tid,) = data._transfers
+        data._cancel_still_inflight.add(tid)
+
+        result = session.close()
+
+        assert data._cancel_calls == [([tid], "wait")]
+        assert tid in data._transfers
+        assert result.retained == {tid: frozenset({1})}
+        assert result.failed_stores == [2]
+
+    def test_close_releases_stopped_transfer(self):
+        """A transfer that has already stopped is released, not retained."""
+        session, conn, data = _make_session()
+        _activate(session, conn)
+        session.add_stored_blocks("req", [OffloadKey(b"key")], [0], job_id=1)
+        _fetch(conn, "req", [b"key"], [5])
+        session.poll()
+
+        result = session.close()
+
+        assert data._transfers == {}
+        assert result.retained == {}
+        assert result.failed_stores == [1]
+
     def test_send_failure_marks_connection_dead(self):
         """A raising send must mark the connection dead, not silently drop
         the message — otherwise the session lingers alive, is never reaped,
@@ -3596,6 +3754,9 @@ def _fetch(conn: FakeConnection, kv_request_id: str, keys, idxs) -> None:
     )
 
 
+_STALL_S = 20.0
+
+
 class TestPeerFaultEscalation:
     """A transport peer fault ends the session through the normal reap
     path (see "Failure model" in data/base.py): no new WRITE is submitted,
@@ -3657,6 +3818,73 @@ def test_fault_disconnects_after_drain_deadline(self):
         session.poll()
         assert not session.alive
 
+    def _stalled_session(self):
+        session, conn, transport = _make_session()
+        session._server._stall_timeout_s = _STALL_S
+        _activate(session, conn)
+        session.add_stored_blocks("req-1", [OffloadKey(b"k1")], [0], job_id=1)
+        _fetch(conn, "req-1", [b"k1"], [5])
+        session.poll()
+        (tid,) = transport._transfers
+        transport._cancel_still_inflight.add(tid)
+        return session, conn, transport, tid
+
+    def test_transfer_under_stall_timeout_keeps_session(self):
+        session, _, _, tid = self._stalled_session()
+        xfer = session._server._inflight[tid]
+        xfer.submitted_at = time.monotonic() - _STALL_S + 1.0
+        session.poll()
+        assert session.alive
+        assert session._server.peer_fault is None
+
+    def test_transfer_done_in_same_poll_is_not_stalled(self):
+        """A WRITE past the stall timeout that the transport reports done in
+        this very poll completes; it does not fault the peer."""
+        session, _, transport, tid = self._stalled_session()
+        session._server._inflight[tid].submitted_at = 0.0
+        transport._poll_done.append(tid)
+        stores = session.poll().stores
+        assert session.alive
+        assert session._server.peer_fault is None
+        assert StoreResult(job_id=1, success=True) in stores
+
+    def test_stalled_transfer_faults_peer_without_releasing(self):
+        """A WRITE past the stall timeout faults the peer at once (no drain
+        wait), and close() retains it with its store job still pinned."""
+        session, conn, transport, tid = self._stalled_session()
+        xfer = session._server._inflight[tid]
+        xfer.submitted_at = time.monotonic() - _STALL_S - 1.0
+
+        stores = session.poll().stores
+        assert not session.alive
+        assert session.data_plane_fault.startswith("transfer stalled for")
+        assert conn._sent[-1] == {TYPE_KEY: DisconnectMsg.TYPE}
+        assert StoreResult(job_id=1, success=False) not in stores
+
+        result = session.close()
+        assert result.retained == {tid: frozenset({1})}
+        assert 1 not in result.failed_stores
+        assert tid in transport._transfers
+        assert all(mode == "wait" for _, mode in transport._cancel_calls)
+
+    def test_stalled_peer_gets_no_new_writes(self):
+        """The fault lands in the same poll that would serve a new fetch."""
+        session, conn, transport, tid = self._stalled_session()
+        session._server._inflight[tid].submitted_at = 0.0
+        session._server.collect_results()
+        assert session._server.peer_fault is not None
+        session.add_stored_blocks("req-2", [OffloadKey(b"k2")], [1], job_id=2)
+        _fetch(conn, "req-2", [b"k2"], [6])
+        session.poll()
+        assert list(transport._transfers) == [tid]
+
+    def test_stall_timeout_zero_disables(self):
+        session, _, _, tid = self._stalled_session()
+        session._server._stall_timeout_s = 0
+        session._server._inflight[tid].submitted_at = 0.0
+        session.poll()
+        assert session.alive
+
     def test_death_reasons(self):
         session, conn, _ = _make_session()
         _activate(session, conn)
@@ -3672,3 +3900,116 @@ def test_first_death_reason_wins(self):
         transport._peer_faults["peer:8000"] = "late"
         session.poll()
         assert session.death_reason == "peer disconnected"
+
+
+@pytest.mark.parametrize("success", [False, True])
+def test_drain_stores_keeps_inflight_pins_until_transport_finishes(success):
+    """Drain fails queued jobs once, but cannot release slots read by a WRITE."""
+    session, conn, transport = _make_session()
+    _activate(session, conn)
+    session.add_stored_blocks(
+        "req-1", [OffloadKey(b"active"), OffloadKey(b"extra")], [0, 1], job_id=1
+    )
+    session.add_stored_blocks("req-1", [OffloadKey(b"queued")], [2], job_id=2)
+    _fetch(conn, "req-1", [b"active"], [5])
+    assert session.poll().stores == []
+    (tid,) = transport._transfers
+
+    assert session.drain_stores().stores == [StoreResult(2, False)]
+    assert session.has_pending_work
+    assert session.drain_stores().stores == []
+    assert session.poll().stores == []
+    assert transport._cancel_calls == []
+
+    (transport._poll_done if success else transport._poll_failed).append(tid)
+    assert session.poll().stores == [StoreResult(1, success)]
+    assert not session.has_pending_work
+    assert session.drain_stores().stores == []
+    _fetch(conn, "req-1", [b"extra"], [6])
+    session.poll()
+    assert len(transport._transfers) == 1
+    session.finish_request("req-1")
+    session.poll()
+    done = [m for m in conn._sent if m[TYPE_KEY] == TransferDoneMsg.TYPE]
+    assert done[-1][TransferDoneMsg.SUCCESS] is False
+    assert session.poll().stores == []
+
+
+def test_drain_rejects_unfetched_lookup_round_without_timeout():
+    """Withdrawn lookup supply must fail a late FetchMsg in one poll."""
+    session, conn, transport = _make_session()
+    _activate(session, conn)
+    key = OffloadKey(b"lookup")
+    parent = FakeParent(stored={key: 0})
+    conn.enqueue(
+        {
+            TYPE_KEY: LookupMsg.TYPE,
+            LookupMsg.KV_REQUEST_ID: "sym",
+            LookupMsg.KEYS: [key],
+            LookupMsg.ROUND_SEQ: 0,
+        }
+    )
+    session.poll()
+    _serve(session, parent)
+    assert conn._sent[-1][LookupRespMsg.HITS] == [True]
+    session.drain_stores()
+    _fetch(conn, "sym", [key], [0])
+    session.poll()
+    assert transport._transfers == {}
+    done = [m for m in conn._sent if m[TYPE_KEY] == TransferDoneMsg.TYPE]
+    assert len(done) == 1
+    assert done[0][TransferDoneMsg.SUCCESS] is False
+
+    conn.enqueue(
+        {
+            TYPE_KEY: LookupMsg.TYPE,
+            LookupMsg.KV_REQUEST_ID: "sym",
+            LookupMsg.KEYS: [key],
+            LookupMsg.ROUND_SEQ: 1,
+        }
+    )
+    session.poll()
+    _serve(session, parent)
+    assert conn._sent[-1][LookupRespMsg.HITS] == [True]
+    conn.enqueue(
+        {
+            TYPE_KEY: FetchMsg.TYPE,
+            FetchMsg.KV_REQUEST_ID: "sym",
+            FetchMsg.KEYS: [key],
+            FetchMsg.BLOCK_INDEXES: [1],
+            FetchMsg.ROUND_SEQ: 1,
+        }
+    )
+    session.poll()
+    (tid,) = transport._transfers
+    transport._poll_done.append(tid)
+    assert session.poll().stores == [StoreResult(1001, True)]
+    assert conn._sent[-1][TransferDoneMsg.SUCCESS] is True
+
+
+def test_drain_aborted_round_sends_only_abort_ack():
+    """An abort parked behind a WRITE must get exactly one terminal message."""
+    session, conn, transport = _make_session()
+    _activate(session, conn)
+    key = OffloadKey(b"active")
+    session.add_stored_blocks("req", [key], [0], job_id=1)
+    _fetch(conn, "req", [key], [0])
+    session.poll()
+    (tid,) = transport._transfers
+    transport._cancel_still_inflight.add(tid)
+    conn.enqueue(
+        {
+            TYPE_KEY: AbortFetchMsg.TYPE,
+            AbortFetchMsg.KV_REQUEST_ID: "req",
+            AbortFetchMsg.ROUND_SEQ: 0,
+        }
+    )
+    session.poll()
+    session.drain_stores()
+    transport._poll_done.append(tid)
+    result = session.poll()
+    assert result.stores == [StoreResult(1, False)]
+    terminals = [
+        m for m in conn._sent if m[TYPE_KEY] in (TransferDoneMsg.TYPE, AbortAckMsg.TYPE)
+    ]
+    assert [m[TYPE_KEY] for m in terminals] == [AbortAckMsg.TYPE]
diff --git a/tests/v1/kv_offload/tiering/p2p/test_zmq_transport.py b/tests/v1/kv_offload/tiering/p2p/test_zmq_transport.py
index eebb9c01ba1e..d7bf17fb925d 100644
--- a/tests/v1/kv_offload/tiering/p2p/test_zmq_transport.py
+++ b/tests/v1/kv_offload/tiering/p2p/test_zmq_transport.py
@@ -4,6 +4,7 @@
 
 from __future__ import annotations
 
+import select
 import socket
 import time
 
@@ -17,6 +18,11 @@
     _Sockets,
 )
 
+# Transports whose libzmq I/O thread is blocked. Terminating their context
+# blocks forever, and garbage collection would terminate it, so they stay
+# referenced until the interpreter exits (pyzmq skips the term at exit).
+_WEDGED_TRANSPORTS: list[ZmqTransport] = []
+
 
 def _free_port() -> int:
     """Find a free TCP port."""
@@ -353,6 +359,99 @@ def test_inbound_message_survives_dead_registration(self):
             transport_a.close()
             transport_b.close()
 
+    def test_close_with_unsent_message_keeps_transport_usable(self, monkeypatch):
+        """A connection closed with an unsent message does not wedge the transport.
+
+        A closed DEALER that still holds a message keeps reconnecting to an
+        unresponsive peer, and its monitor events then block libzmq's I/O
+        thread, so no socket in the context can send or receive again.
+        """
+        apply_heartbeat = zmq_module._apply_heartbeat
+
+        def apply_heartbeat_fast_handshake(sock: zmq.Socket) -> None:
+            apply_heartbeat(sock)
+            sock.setsockopt(zmq.HANDSHAKE_IVL, 200)
+
+        monkeypatch.setattr(
+            zmq_module, "_apply_heartbeat", apply_heartbeat_fast_handshake
+        )
+        # Accepts TCP connections but never answers the ZMTP greeting.
+        silent = socket.socket()
+        silent.bind(("127.0.0.1", 0))
+        silent.listen(8)
+        transport_a, port_a = _make_transport()
+        transport_b, _ = _make_transport()
+        wedged = True
+        try:
+            stale = transport_b.connect(f"127.0.0.1:{silent.getsockname()[1]}")
+            stale.send({"type": "stale"})
+            stale.close()
+            # Several ZMTP handshake timeouts on the closed DEALER.
+            time.sleep(1.0)
+
+            conn = transport_b.connect(f"127.0.0.1:{port_a}")
+            conn.send({"seq": 1})
+            inbound = _wait_for_inbound(transport_a)[0]
+            assert _wait_for_messages(transport_a, inbound, 1) == [{"seq": 1}]
+            wedged = False
+        finally:
+            transport_a.close()
+            silent.close()
+            if wedged:
+                _WEDGED_TRANSPORTS.append(transport_b)
+            else:
+                transport_b.close()
+
+    def test_close_while_connecting_keeps_transport_usable(self):
+        """Closing a connection whose TCP connect never completed keeps the
+        transport usable.
+
+        A peer address with no host behind it (a departed pod, a lost node)
+        never answers the SYN. Tearing down that connect raises EVENT_CLOSED
+        on libzmq's I/O thread, and a monitor event sent after its reader is
+        closed blocks that thread, so no socket in the context can send or
+        receive again.
+        """
+        # Linux drops SYNs to a listener whose accept queue is full, so a
+        # connect to it stays in progress like one to an address with no host.
+        full = socket.socket()
+        full.bind(("127.0.0.1", 0))
+        full.listen(0)
+        addr = full.getsockname()
+        filler = socket.create_connection(addr, timeout=1.0)
+        probe = socket.socket()
+        probe.setblocking(False)
+        probe.connect_ex(addr)
+        _, connected, _ = select.select([], [probe], [], 0.5)
+        probe.close()
+        if connected:
+            filler.close()
+            full.close()
+            pytest.skip("kernel completes connects to a full accept queue")
+
+        transport_a, port_a = _make_transport()
+        transport_b, _ = _make_transport()
+        wedged = True
+        try:
+            pending = transport_b.connect(f"127.0.0.1:{addr[1]}")
+            # Let the I/O thread start the TCP connect.
+            time.sleep(0.2)
+            pending.close()
+
+            conn = transport_b.connect(f"127.0.0.1:{port_a}")
+            conn.send({"seq": 1})
+            inbound = _wait_for_inbound(transport_a)[0]
+            assert _wait_for_messages(transport_a, inbound, 1) == [{"seq": 1}]
+            wedged = False
+        finally:
+            transport_a.close()
+            filler.close()
+            full.close()
+            if wedged:
+                _WEDGED_TRANSPORTS.append(transport_b)
+            else:
+                transport_b.close()
+
 
 class TestExpandedMonitorEvents:
     """The monitor mask includes EVENT_DISCONNECTED and EVENT_CLOSED so
diff --git a/tests/v1/kv_offload/tiering/test_async_lookup.py b/tests/v1/kv_offload/tiering/test_async_lookup.py
index d23d837371bd..543257dc9b35 100644
--- a/tests/v1/kv_offload/tiering/test_async_lookup.py
+++ b/tests/v1/kv_offload/tiering/test_async_lookup.py
@@ -403,3 +403,105 @@ def test_enqueue_once_invariant_enforced(self):
         with pytest.raises(AssertionError):
             mgr.drain_results()
         mgr.shutdown()
+
+    def test_store_drops_orphaned_in_flight_state(
+        self, monkeypatch: pytest.MonkeyPatch
+    ):
+        """A store must not retain an in-flight state after its request ended."""
+        key = _key(1)
+        mgr = InMemoryLookupManager()
+        probe_started = threading.Event()
+        release_probe = threading.Event()
+        batch_lookup = mgr.batch_lookup
+
+        def blocking_lookup(keys, req_context):
+            results = list(batch_lookup(keys, req_context))
+            probe_started.set()
+            if not release_probe.wait(timeout=5):
+                raise TimeoutError("Test did not release the backend probe")
+            return results
+
+        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)
+            mgr.cleanup("req_a")
+            assert key in mgr._lookup_state
+            assert not mgr._lookup_state[key].request_ids
+
+            mgr._existing.add(key)
+            mgr.mark_present([key])
+            assert key not in mgr._lookup_state
+            release_probe.set()
+        finally:
+            release_probe.set()
+            mgr.shutdown()
+
+        assert key not in mgr._lookup_state
+
+    def test_mark_present_handles_mixed_phases(self, monkeypatch: pytest.MonkeyPatch):
+        """One call updates pending, in-flight and resolved keys independently."""
+        pending_key, in_flight_key, resolved_key, unknown_key = (
+            _key(1),
+            _key(2),
+            _key(3),
+            _key(99),
+        )
+        mgr = InMemoryLookupManager()
+        probe_started = threading.Event()
+        release_probe = threading.Event()
+        batch_lookup = mgr.batch_lookup
+        calls = 0
+
+        def blocking_lookup(keys, req_context):
+            nonlocal calls
+            calls += 1
+            results = list(batch_lookup(keys, req_context))
+            if calls == 2:
+                probe_started.set()
+                if not release_probe.wait(timeout=5):
+                    raise TimeoutError("Test did not release the backend probe")
+            return results
+
+        monkeypatch.setattr(mgr, "batch_lookup", blocking_lookup)
+        try:
+            # Resolved: a probe delivered a miss.
+            assert mgr.lookup(resolved_key, _ctx("req_c")) is None
+            mgr.flush()
+            assert mgr._results_ready.wait(timeout=5)
+            mgr._results_ready.clear()
+            assert mgr.lookup(resolved_key, _ctx("req_c")) is False
+
+            # In flight: the worker is still inside this probe.
+            assert mgr.lookup(in_flight_key, _ctx("req_b")) is None
+            mgr.flush()
+            assert probe_started.wait(timeout=5)
+
+            # Pending: accumulated this step, not flushed yet.
+            assert mgr.lookup(pending_key, _ctx("req_a")) is None
+
+            mgr.mark_present([pending_key, in_flight_key, resolved_key, unknown_key])
+
+            for key in (pending_key, in_flight_key, resolved_key):
+                state = mgr._lookup_state[key]
+                assert state.result is True
+                assert state.phase is LookupPhase.RESOLVED
+            assert unknown_key not in mgr._lookup_state
+
+            # A stale pending entry is skipped when its batch is flushed.
+            mgr.flush()
+
+            assert mgr.lookup(pending_key, _ctx("req_a")) is True
+            assert mgr.lookup(in_flight_key, _ctx("req_b")) is True
+            assert mgr.lookup(resolved_key, _ctx("req_c")) is True
+            assert mgr.lookup(unknown_key, _ctx("req_d")) is None
+
+            release_probe.set()
+        finally:
+            release_probe.set()
+            mgr.shutdown()
+
+        # Only the two flushed probes ran: the pending batch was dropped, the
+        # in-flight result was ignored, and unknown_key was never probed.
+        assert calls == 2
diff --git a/tests/v1/kv_offload/tiering/test_fs_tier.py b/tests/v1/kv_offload/tiering/test_fs_tier.py
index 767a73f6ee6f..8c8ebf01ff6b 100644
--- a/tests/v1/kv_offload/tiering/test_fs_tier.py
+++ b/tests/v1/kv_offload/tiering/test_fs_tier.py
@@ -231,6 +231,18 @@ def test_store_creates_file_and_lookup_succeeds(fs_tier):
     assert os.path.exists(dest), f"Expected file at {dest}"
 
 
+def test_store_marks_a_cached_miss_present(fs_tier):
+    tier, _ = fs_tier
+    block_key = key(1)
+    assert lookup_and_wait(tier, [block_key]) == [LookupResult.MISS]
+
+    tier.submit_store(make_job(1, [block_key], [0]))
+    results = drain(tier)
+    assert len(results) == 1
+    assert results[0].success
+    assert lookup_and_wait(tier, [block_key]) == [LookupResult.HIT]
+
+
 def test_store_accepts_block_published_by_another_writer(fs_tier, monkeypatch):
     """Losing the race to publish a block must not fail the store job.
 
@@ -824,6 +836,7 @@ def test_successful_store_emits_stored_event(fs_tier_with_events):
     assert events[0].medium == Medium.STORAGE
     assert events[0].locality is Locality.LOCAL
     assert not events[0].removed
+    assert events[0].req_context is _CTX
     # take_events drains the buffer.
     assert list(tier.take_events()) == []
 
@@ -920,7 +933,7 @@ def flaky_batch_store_block(paths, *args, **kwargs):
     assert len(results) == 1
     assert not results[0].success
     assert list(tier.take_events()) == []
-    assert tier._store_job_keys == {}
+    assert tier._store_jobs == {}
 
 
 def test_events_disabled_by_default(fs_tier):
@@ -930,7 +943,7 @@ def test_events_disabled_by_default(fs_tier):
     assert len(results) == 1
     assert results[0].success
     assert tier.events is None
-    assert tier._store_job_keys == {}
+    assert tier._store_jobs == {}
     assert list(tier.take_events()) == []
 
 
@@ -951,7 +964,7 @@ def test_events_require_global_kv_events_flag(tmp_path):
         assert len(results) == 1
         assert results[0].success
         assert list(tier.take_events()) == []
-        assert tier._store_job_keys == {}
+        assert tier._store_jobs == {}
     finally:
         tier.shutdown()
 
@@ -995,6 +1008,7 @@ def test_cascade_store_emits_fs_event_through_tiering_manager(tmp_path):
         assert len(fs_events) == 1
         assert set(fs_events[0].keys) == set(keys)
         assert not fs_events[0].removed
+        assert fs_events[0].req_context is _CTX
     finally:
         tier.shutdown()
 
diff --git a/tests/v1/kv_offload/tiering/test_kvcr_tier.py b/tests/v1/kv_offload/tiering/test_kvcr_tier.py
index a6f67f80f3f4..946e5cfabba2 100644
--- a/tests/v1/kv_offload/tiering/test_kvcr_tier.py
+++ b/tests/v1/kv_offload/tiering/test_kvcr_tier.py
@@ -7,7 +7,6 @@
 from typing import Any
 from unittest.mock import Mock
 
-import msgspec
 import numpy as np
 import pytest
 
@@ -367,9 +366,7 @@ def test_kvcr_tier_adapts_request_and_load(monkeypatch, backpressure):
 
     tier.on_new_request(ctx)
 
-    assert kvcr.submit_hint_calls == [
-        ((), {"request_id": "req", "hints": msgspec.to_builtins(kv_hint)})
-    ]
+    assert kvcr.submit_hint_calls == [((), {"request_id": "req", "hints": kv_hint})]
 
     bindings = kvcr.constructor_bindings
     assert bindings is not None
@@ -390,6 +387,7 @@ def test_kvcr_tier_adapts_request_and_load(monkeypatch, backpressure):
     assert list(blocks) == [key]
     (descriptor,) = blocks[key]
     assert descriptor == MemoryRef(end_point_name=kvcr.nixl_agent_name, element_index=2)
+    assert type(descriptor.element_index) is int
     assert list(tier.get_finished_jobs()) == [JobResult(7, True)]
 
     # Here we verify that request cleanup discards the request-scoped hint in
@@ -479,13 +477,24 @@ def test_kvcr_tier_maps_query_status(monkeypatch, status, expected):
     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."""
+@pytest.mark.parametrize(
+    ("job_keys", "chunk_ids", "success"),
+    [
+        (None, [1, 3], True),
+        ((b"k2", b"k0"), [3, 1], True),
+        (None, [1], False),
+        ((b"k0", b"other"), [1, 3], False),
+    ],
+    ids=["original-keys", "reordered-keys", "short-chunks", "unexpected-key"],
+)
+def test_kvcr_tier_serves_primary_pin_request(
+    monkeypatch, job_keys, chunk_ids, success
+):
+    """Validate pin jobs and hold accepted hits until KVCR releases the 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]: 3}
     lifecycle: list[str] = []
 
     class Parent:
@@ -499,8 +508,8 @@ def lookup(self, key, req_context):
         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]),
+                keys=requested_keys if job_keys is None else job_keys,
+                chunk_ids=np.array(chunk_ids, dtype=np.int32),
                 is_promotion=True,
                 req_context=req_context,
             )
@@ -516,6 +525,11 @@ def on_request_finished(self, req_context):
 
     [(queued_request, result)] = bindings.poll_pin_results()
     assert queued_request == request
+    assert lifecycle == ["new", "finished"]
+    if not success:
+        assert result is None
+        assert list(tier.get_finished_jobs()) == [JobResult(11, False)]
+        return
     assert result is not None
     pin_handle, descriptors = result
     assert descriptors[keys[1]] is None
@@ -523,7 +537,7 @@ def on_request_finished(self, req_context):
     assert block_descriptors is not None
     (descriptor,) = block_descriptors
     assert descriptor == MemoryRef(end_point_name=kvcr.nixl_agent_name, element_index=3)
-    assert lifecycle == ["new", "finished"]
+    assert type(descriptor.element_index) is int
 
     polls = 0
 
@@ -626,6 +640,7 @@ def test_kvcr_tier_stores_and_emits_inventory(monkeypatch):
     assert blocks[key] == [
         MemoryRef(end_point_name=kvcr.nixl_agent_name, element_index=2)
     ]
+    assert type(blocks[key][0].element_index) is int
     assert list(tier.get_finished_jobs()) == [JobResult(11, True)]
 
     assert kvcr.inventory_sink is not None
@@ -655,7 +670,16 @@ def test_kvcr_tier_requires_self_describing_inventory_events(monkeypatch):
         )
 
 
-def test_kvcr_tier_accumulates_block_results(monkeypatch):
+@pytest.mark.parametrize(
+    ("successes", "expected"),
+    [
+        ((True, False), JobResult(13, False, {OffloadKey(b"k0")})),
+        ((False, True), JobResult(13, False, {OffloadKey(b"k1")})),
+        ((True, True), JobResult(13, True)),
+    ],
+    ids=["success-then-failure", "failure-then-success", "all-success"],
+)
+def test_kvcr_tier_accumulates_block_results(monkeypatch, successes, expected):
     """Wait for every block result and preserve partial success while draining."""
     kvcr = RecordingKVCR()
     tier = _make_tier(monkeypatch, kvcr)
@@ -672,13 +696,10 @@ def test_kvcr_tier_accumulates_block_results(monkeypatch):
 
     op_handle, _, _ = kvcr.deliver_calls[0]
     kvcr.completed.clear()
-    kvcr.complete(op_handle, {keys[0]: True})
+    kvcr.complete(op_handle, {keys[0]: successes[0]})
     assert list(tier.get_finished_jobs()) == []
 
-    kvcr.complete(op_handle, {keys[1]: False})
+    kvcr.complete(op_handle, {keys[1]: successes[1]})
     tier.drain_jobs()
 
-    [result] = tier.get_finished_jobs()
-    assert result.job_id == 13
-    assert not result.success
-    assert result.successful_keys == {keys[0]}
+    assert list(tier.get_finished_jobs()) == [expected]
diff --git a/tests/v1/kv_offload/tiering/test_obj_tier.py b/tests/v1/kv_offload/tiering/test_obj_tier.py
index 640848bda6f9..f69755ed2286 100644
--- a/tests/v1/kv_offload/tiering/test_obj_tier.py
+++ b/tests/v1/kv_offload/tiering/test_obj_tier.py
@@ -321,6 +321,28 @@ def test_store_and_lookup(self):
         assert results[0].success
         assert lookup_and_wait(self.tier, [key(1)]) == [LookupResult.HIT]
 
+    def test_store_marks_a_cached_miss_present(self):
+        block_key = key(1)
+        assert lookup_and_wait(self.tier, [block_key]) == [LookupResult.MISS]
+
+        self.tier.submit_store(make_job(1, [block_key], [0]))
+        results = drain(self.tier)
+        assert len(results) == 1
+        assert results[0].success
+        assert lookup_and_wait(self.tier, [block_key]) == [LookupResult.HIT]
+
+    def test_failed_store_does_not_mark_cached_miss_present(self):
+        block_key = key(1)
+        assert lookup_and_wait(self.tier, [block_key]) == [LookupResult.MISS]
+
+        self.agent.check_xfer_state = lambda h: "ERR"
+        self.tier.submit_store(make_job(1, [block_key], [0]))
+        results = drain(self.tier)
+
+        assert len(results) == 1
+        assert not results[0].success
+        assert lookup_and_wait(self.tier, [block_key]) == [LookupResult.MISS]
+
     def test_lookup_unrelated_key_returns_false(self):
         self.tier.submit_store(make_job(1, [key(1)], [0]))
         drain(self.tier)
@@ -547,6 +569,26 @@ def test_post_transfer_cleanup_failure_does_not_lose_result(
         assert not tier._transfers
         assert list(tier.get_finished_jobs()) == []
 
+    def test_telemetry_failure_does_not_lose_successful_transfer(self, monkeypatch):
+        """Telemetry is optional, so a raise must not fail a completed store."""
+        tier, agent = _make_tier(num_blocks=4)
+        monkeypatch.setattr(
+            agent,
+            "get_xfer_telemetry",
+            MagicMock(side_effect=RuntimeError("nixlNoTelemetryError")),
+        )
+
+        tier.submit_store(make_job(1, [key(1)], [0]))
+        results = list(tier.get_finished_jobs())
+
+        assert len(results) == 1
+        assert results[0].job_id == 1
+        assert results[0].success
+        # Only the timing is lost.
+        assert results[0].transfer_time is None
+        assert not tier._transfers
+        assert list(tier.get_finished_jobs()) == []
+
     def test_xfer_cleanup_retry_finalizes_parent_job_and_primary_pin(self, monkeypatch):
         num_blocks = 4
         tensor = torch.zeros((num_blocks, _BLOCK_ELEMENTS), dtype=_DTYPE)
@@ -637,6 +679,7 @@ def test_successful_store_emits_stored_event(self):
         assert events[0].medium == Medium.STORAGE
         assert events[0].locality is Locality.REMOTE
         assert not events[0].removed
+        assert events[0].req_context is _CTX
         # take_events drains the buffer.
         assert list(self.tier.take_events()) == []
 
@@ -677,7 +720,7 @@ def test_mixed_job_results_emit_event_only_for_successful_job(self):
         events = list(self.tier.take_events())
         assert len(events) == 1
         assert events[0].keys == [key(2)]
-        assert self.tier._store_job_keys == {}
+        assert self.tier._store_jobs == {}
 
     def test_load_job_emits_no_event(self):
         self.tier.submit_store(make_job(1, [key(1)], [0]))
@@ -698,7 +741,7 @@ def test_failed_transfer_emits_no_event(self):
         results = drain(self.tier)
         assert not results[0].success
         assert list(self.tier.take_events()) == []
-        assert self.tier._store_job_keys == {}
+        assert self.tier._store_jobs == {}
 
     def test_submission_failure_emits_no_event(self):
         self.agent.make_prepped_xfer = lambda *a, **k: None
@@ -706,7 +749,7 @@ def test_submission_failure_emits_no_event(self):
         results = list(self.tier.get_finished_jobs())
         assert not results[0].success
         assert list(self.tier.take_events()) == []
-        assert self.tier._store_job_keys == {}
+        assert self.tier._store_jobs == {}
 
     def test_events_disabled_by_default(self):
         tier, _ = _make_tier()
@@ -715,7 +758,7 @@ def test_events_disabled_by_default(self):
         assert len(results) == 1
         assert results[0].success
         assert tier.events is None
-        assert tier._store_job_keys == {}
+        assert tier._store_jobs == {}
         assert list(tier.take_events()) == []
 
     def test_events_require_global_kv_events_flag(self):
@@ -729,7 +772,7 @@ def test_events_require_global_kv_events_flag(self):
         assert len(results) == 1
         assert results[0].success
         assert tier.events is None
-        assert tier._store_job_keys == {}
+        assert tier._store_jobs == {}
         assert list(tier.take_events()) == []
 
 
diff --git a/tests/v1/kv_offload/tiering/test_p2p_lifecycle_e2e.py b/tests/v1/kv_offload/tiering/test_p2p_lifecycle_e2e.py
new file mode 100644
index 000000000000..bdcdd93ab4e6
--- /dev/null
+++ b/tests/v1/kv_offload/tiering/test_p2p_lifecycle_e2e.py
@@ -0,0 +1,177 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import time
+
+import pytest
+
+
+@pytest.fixture
+def p2p_control_peer():
+    """A control-only peer: delay its handshake or fail fetches on demand."""
+    import queue
+    import socket
+    import threading
+
+    from vllm.v1.kv_offload.tiering.p2p.control.base import ControlConnection
+    from vllm.v1.kv_offload.tiering.p2p.control.zmq import ZmqTransport
+
+    with socket.socket() as sock:
+        sock.bind(("127.0.0.1", 0))
+        port = sock.getsockname()[1]
+    handshake = threading.Event()
+    reply_to_fetch = threading.Event()
+    fetches: queue.Queue[dict] = queue.Queue()
+    ready = threading.Event()
+    stop = threading.Event()
+    errors = []
+
+    def serve():
+        transport = None
+        try:
+            transport = ZmqTransport(f"127.0.0.1:{port}", "127.0.0.1", port)
+            connections: list[ControlConnection] = []
+            acknowledged = set()
+            ready.set()
+            while not stop.is_set():
+                connections.extend(transport.poll())
+                for conn in connections:
+                    if handshake.is_set() and conn.peer_id not in acknowledged:
+                        conn.send(
+                            {"type": "connect_ack", "peer_id": f"127.0.0.1:{port}"}
+                        )
+                        acknowledged.add(conn.peer_id)
+                    for msg in conn.recv():
+                        kind = msg["type"]
+                        if kind == "fetch" and msg["keys"]:
+                            fetches.put(msg)
+                            if not reply_to_fetch.is_set():
+                                continue
+                            conn.send(
+                                {
+                                    "type": "transfer_done",
+                                    "kv_request_id": msg["kv_request_id"],
+                                    "round_seq": msg["round_seq"],
+                                    "success": False,
+                                }
+                            )
+                        elif kind == "abort_fetch":
+                            conn.send(
+                                {
+                                    "type": "abort_ack",
+                                    "kv_request_id": msg["kv_request_id"],
+                                    "round_seq": msg["round_seq"],
+                                }
+                            )
+                stop.wait(0.001)
+        except Exception as exc:
+            errors.append(exc)
+            ready.set()
+        finally:
+            if transport is not None:
+                transport.close()
+
+    thread = threading.Thread(target=serve, daemon=True)
+    thread.start()
+    assert ready.wait(10)
+    assert not errors
+    try:
+        yield port, handshake, reply_to_fetch, fetches
+    finally:
+        stop.set()
+        thread.join(timeout=10)
+        assert not thread.is_alive()
+        assert not errors
+
+
+def test_inference_reuses_a_finished_kv_request_id(
+    vllm_runner, monkeypatch, p2p_control_peer
+):
+    """A cancellation's late failure must not suppress a later request's fetch."""
+    from vllm import SamplingParams
+    from vllm.config import KVTransferConfig
+
+    monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
+    monkeypatch.setenv("VLLM_BATCH_INVARIANT", "1")
+    monkeypatch.setenv("VLLM_P2P_HANDSHAKE_TIMEOUT_S", "120")
+    peer_port, handshake, reply_to_fetch, fetches = p2p_control_peer
+    import socket
+
+    with socket.socket() as sock:
+        sock.bind(("127.0.0.1", 0))
+        local_port = sock.getsockname()[1]
+    config = KVTransferConfig(
+        kv_connector="OffloadingConnector",
+        kv_role="kv_both",
+        kv_load_failure_policy="recompute",
+        kv_connector_extra_config={
+            "spec_name": "TieringOffloadingSpec",
+            "cpu_bytes_to_use": 16 * 1024 * 1024,
+            "secondary_tiers": [
+                {"type": "p2p", "host": "127.0.0.1", "port": local_port}
+            ],
+        },
+    )
+    with vllm_runner(
+        "HuggingFaceTB/SmolLM2-135M",
+        skip_tokenizer_init=True,
+        dtype="float16",
+        max_model_len=256,
+        enforce_eager=True,
+        gpu_memory_utilization=0.35,
+        max_num_seqs=2,
+        max_num_batched_tokens=512,
+        enable_prefix_caching=True,
+        disable_log_stats=False,
+        kv_transfer_config=config,
+    ) as runner:
+        llm = runner.llm
+        params = SamplingParams(
+            temperature=0,
+            max_tokens=4,
+            ignore_eos=True,
+            extra_args={
+                "kv_transfer_params": {
+                    "remote_prefiller": {
+                        "remote_host": "127.0.0.1",
+                        "remote_port": peer_port,
+                        "kv_request_id": "retired",
+                    }
+                }
+            },
+        )
+        prompt = {"prompt_token_ids": [3 + i % 100 for i in range(192)]}
+
+        handshake.set()
+        request_ids = llm.enqueue([prompt], params, use_tqdm=False)
+        deadline = time.monotonic() + 30
+        while fetches.empty():
+            assert time.monotonic() < deadline, "no initial P2P fetch"
+            llm.llm_engine.step()
+        fetches.get_nowait()
+        llm.llm_engine.abort_request(request_ids, internal=True)
+        from vllm.v1.metrics.reader import Counter
+
+        deadline = time.monotonic() + 30
+        while True:
+            llm.generate(
+                [{"prompt_token_ids": [0]}],
+                SamplingParams(temperature=0, max_tokens=1),
+                use_tqdm=False,
+            )
+            failures = [
+                metric.value
+                for metric in llm.get_metrics()
+                if isinstance(metric, Counter)
+                and "promotion_job_failures" in metric.name
+            ]
+            if sum(failures) > 0:
+                break
+            assert time.monotonic() < deadline, (
+                f"cancelled load was not drained: {failures}"
+            )
+        reply_to_fetch.set()
+        second_prompt = {"prompt_token_ids": [103 + i % 100 for i in range(192)]}
+        result = llm.generate([second_prompt], params, use_tqdm=False)[0]
+        assert len(result.outputs[0].token_ids) == 4
+        assert not fetches.empty(), "late failure poisoned the reused KV request id"
+        assert fetches.get_nowait()["kv_request_id"] == "retired"
diff --git a/tests/v1/kv_offload/tiering/test_tier_poll_thread.py b/tests/v1/kv_offload/tiering/test_tier_poll_thread.py
index 0ae2e27e1dfd..6245379a76f7 100644
--- a/tests/v1/kv_offload/tiering/test_tier_poll_thread.py
+++ b/tests/v1/kv_offload/tiering/test_tier_poll_thread.py
@@ -260,8 +260,8 @@ def test_tier_that_did_not_opt_in_is_not_serviced_off_thread():
         with manager.lock:
             manager.on_schedule_end(_EMPTY_SCHEDULE_END)
 
-        # on_schedule_end still serves every tier, so the thread dying degrades
-        # to per-step servicing rather than to none.
+        # on_schedule_end serves every tier, so the unpolled tier is serviced
+        # once per step.
         assert quiet.serves == 1
     finally:
         manager.shutdown()
@@ -350,6 +350,141 @@ def test_parent_lookup_while_serving_keeps_the_round_contained():
         manager.shutdown()
 
 
+class _HitTier(_FakeTier):
+    """Second tier whose lookup HITs, so the parent queues a promotion."""
+
+    tolerates_off_thread_polling: ClassVar[bool] = False
+
+    def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult:
+        return LookupResult.HIT
+
+
+class _PromotingTier(_FakeTier):
+    """Opted-in tier that serves a peer lookup through its parent.
+
+    The parent lookup excludes this tier and lands on _HitTier, so the
+    tiering manager queues a promotion in _pending_load_submissions, which
+    only on_schedule_end() flushes.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.on_serve = self._promote
+        self.promoted = threading.Event()
+
+    def _promote(self, parent: ParentManager) -> None:
+        parent.lookup(OffloadKey(b"\x07" * 8), ReqContext(req_id="peer"))
+        self.promoted.set()
+
+
+def _multi_tier_manager(primary, tier_cls, second_cls, **kwargs):
+    common = {
+        "offloading_spec": MagicMock(),
+        "primary_kv_view": primary.get_kv_memoryview(),
+    }
+    return TieringOffloadingManager(
+        primary_tier=primary,
+        secondary_tiers=[
+            tier_cls(tier_type="first", **common),
+            second_cls(tier_type="second", **common),
+        ],
+        **kwargs,
+    )
+
+
+def test_multi_tier_manager_keeps_the_engine_stepping_while_idle():
+    """With more than one secondary tier, a polled tier must keep the engine
+    stepping even when idle.
+
+    A tier-poll round can serve a peer whose lookup promotes a chunk out of
+    another secondary tier. poll_tiers() never flushes promotions and only a
+    scheduler step collects their completions, so if the engine parked it would
+    never take that step and the load would strand. has_pending_work() is
+    consulted before the park, so it must stay True here.
+    """
+    num_chunks = 8
+    primary = CPUPrimaryTierOffloadingManager(
+        num_chunks=num_chunks,
+        mmap_region=_mock_mmap_region(num_chunks),
+    )
+    manager = _multi_tier_manager(primary, _FakeTier, _HitTier)
+    try:
+        assert manager._tier_poll_thread is not None
+        assert manager.has_pending_work() is True
+    finally:
+        manager.shutdown()
+
+
+def test_two_polled_tiers_keep_the_engine_stepping_while_idle():
+    """Two polled tiers strand the same way as one polled plus one unpolled.
+
+    A serve on the first can promote out of the second, and poll_tiers() never
+    flushes it, so the keep-alive cannot key on the promoted tier being
+    unpolled.
+    """
+    num_chunks = 8
+    primary = CPUPrimaryTierOffloadingManager(
+        num_chunks=num_chunks,
+        mmap_region=_mock_mmap_region(num_chunks),
+    )
+    manager = _multi_tier_manager(primary, _FakeTier, _FakeTier)
+    try:
+        assert manager._tier_poll_thread is not None
+        assert manager.has_pending_work() is True
+    finally:
+        manager.shutdown()
+
+
+def test_off_thread_promotion_is_flushed_by_the_next_step():
+    """An off-thread serve queues a promotion that only a step submits.
+
+    Shows the queue actually fills from the control thread, which is why a
+    multi-tier manager must keep the engine stepping (see
+    test_multi_tier_manager_keeps_the_engine_stepping_while_idle).
+    """
+    num_chunks = 8
+    primary = CPUPrimaryTierOffloadingManager(
+        num_chunks=num_chunks,
+        mmap_region=_mock_mmap_region(num_chunks),
+    )
+    manager = _multi_tier_manager(primary, _PromotingTier, _HitTier)
+    tier = manager.secondary_tiers[0]
+    try:
+        assert tier.promoted.wait(_SETTLE_S), "control thread never served"
+
+        with manager.lock:
+            assert manager._pending_load_submissions
+            manager.on_schedule_end(_EMPTY_SCHEDULE_END)
+            assert not manager._pending_load_submissions
+    finally:
+        manager.shutdown()
+
+
+def test_dead_poll_thread_makes_the_manager_keep_the_engine_stepping(
+    manager_and_tier,
+):
+    """With no live thread the engine tick is the only driver.
+
+    A failing loop now backs off rather than exiting, so this guards the
+    defensive branch: if the field ever holds a finished thread (for example a
+    join that failed to clear it), the manager must keep the engine stepping.
+    """
+    manager, tier = manager_and_tier
+    assert tier.served.wait(_SETTLE_S)
+    assert manager.has_pending_work() is False
+
+    manager._stop_tier_poll_thread()
+    # _stop_tier_poll_thread() clears the field only when the join succeeds.
+    # Simulate a finished thread left in the field.
+    dead = threading.Thread(target=lambda: None)
+    dead.start()
+    dead.join()
+    manager._tier_poll_thread = dead
+
+    assert not dead.is_alive()
+    assert manager.has_pending_work() is True
+
+
 def test_shutdown_joins_the_control_thread_before_tier_teardown(manager_and_tier):
     """The thread drives tier transports, so it must be stopped first.
 
@@ -408,7 +543,7 @@ def test_shutdown_raises_when_the_caller_holds_the_lock(manager_and_tier, monkey
 
 
 def test_failure_burst_does_not_stop_the_thread(manager_and_tier):
-    """The give-up budget is time, not rounds.
+    """The backoff budget is time, not rounds.
 
     At a 1 ms interval, a few hundred failing rounds take well under a second;
     a fault that short must not turn the tier-poll thread off for good.
@@ -431,8 +566,13 @@ def flaky(parent):
     assert _live_control_threads()
 
 
-def test_sustained_failure_stops_the_thread(monkeypatch):
-    monkeypatch.setattr(tiering_manager, "_TIER_POLL_GIVE_UP_AFTER_S", 0.2)
+def test_sustained_failure_backs_off_but_keeps_serving(monkeypatch):
+    """The backoff path is non-terminal: an idle engine parks, so a thread that
+    exited after sustained failure would leave peers unserved with nothing to
+    wake the engine. It must back off and keep retrying, then resume the normal
+    interval once a round succeeds."""
+    monkeypatch.setattr(tiering_manager, "_TIER_POLL_BACK_OFF_AFTER_S", 0.2)
+    monkeypatch.setattr(tiering_manager, "_TIER_POLL_BACKOFF_S", 0.05)
     manager, tier = _make_manager()
     try:
         assert tier.served.wait(_SETTLE_S)
@@ -441,12 +581,135 @@ def always_fails(parent):
             raise RuntimeError("persistent")
 
         tier.on_serve = always_fails
-        deadline = time.monotonic() + 5
-        while _live_control_threads() and time.monotonic() < deadline:
-            time.sleep(0.01)
-        assert not _live_control_threads()
+        time.sleep(0.5)
+        assert _live_control_threads(), "thread must survive sustained failure"
+
+        tier.served.clear()
+        tier.on_serve = None
+        assert tier.served.wait(_SETTLE_S * 2), "thread did not resume serving"
+    finally:
+        manager.shutdown()
+
+
+class _RecordingStop:
+    """Stand-in for the tier-poll stop Event that records the waits it is given.
+
+    Lets a test drive _tier_poll_loop() on the calling thread and read back the
+    interval each round paused for, without real sleeping.
+    """
+
+    def __init__(self, rounds: int):
+        self.waits: list[float] = []
+        self._rounds = rounds
+        self._set = False
+
+    def is_set(self) -> bool:
+        return self._set or len(self.waits) >= self._rounds
+
+    def set(self) -> None:
+        self._set = True
+
+    def wait(self, timeout: float | None = None) -> bool:
+        self.waits.append(timeout)
+        return False
+
+
+def _run_failing_rounds(manager, rounds: int, monkeypatch, **constants):
+    """Drive _tier_poll_loop() on this thread through `rounds` failing rounds.
+
+    Stops the real control thread first so only this call drives the loop, and
+    returns the interval each round paused for.
+    """
+    for name, value in constants.items():
+        monkeypatch.setattr(tiering_manager, name, value)
+    manager._stop_tier_poll_thread()
+    manager.secondary_tiers[0].on_serve = lambda parent: (_ for _ in ()).throw(
+        RuntimeError("persistent")
+    )
+    stop = _RecordingStop(rounds)
+    manager._tier_poll_stop = stop
+    manager._tier_poll_loop()
+    return stop.waits
+
+
+def test_backoff_waits_at_least_the_configured_interval(monkeypatch):
+    """Backing off must never poll faster than tier_poll_interval_s.
+
+    The wait is max(tier_poll_interval_s, 1 s), so a user who set a slow
+    interval does not get a faster loop after a fault.
+    """
+    manager, _ = _make_manager(tier_poll_interval_s=3.0)
+    try:
+        waits = _run_failing_rounds(
+            manager, 4, monkeypatch, _TIER_POLL_BACK_OFF_AFTER_S=0.0
+        )
+    finally:
+        manager.shutdown()
+    assert waits == [3.0, 3.0, 3.0, 3.0]
+
+
+def test_backoff_floors_at_one_second_for_a_fast_interval(monkeypatch):
+    """A sub-second interval still backs off to the 1 s floor."""
+    manager, _ = _make_manager(tier_poll_interval_s=0.001)
+    try:
+        waits = _run_failing_rounds(
+            manager, 3, monkeypatch, _TIER_POLL_BACK_OFF_AFTER_S=0.0
+        )
+    finally:
+        manager.shutdown()
+    assert waits == [1.0, 1.0, 1.0]
+
+
+def test_sustained_failure_logs_one_traceback_then_summaries(monkeypatch):
+    """A permanent fault must not emit a full traceback every round.
+
+    The first failure of a run logs a traceback; later ones are one-line
+    summaries, so an idle rank with a broken peer cannot flood the log.
+    """
+    manager, _ = _make_manager()
+    fake_logger = MagicMock()
+    monkeypatch.setattr(tiering_manager, "logger", fake_logger)
+    try:
+        _run_failing_rounds(
+            manager,
+            4,
+            monkeypatch,
+            _TIER_POLL_BACK_OFF_AFTER_S=0.0,
+            _TIER_POLL_ERROR_LOG_INTERVAL_S=0.0,
+        )
+    finally:
+        manager.shutdown()
+    assert fake_logger.exception.call_count == 1
+    assert fake_logger.warning.call_count == 3
+
+
+def test_alternating_failure_logs_at_most_one_traceback(monkeypatch):
+    """A fault that alternates with successes must not log a traceback per
+    failure.
+
+    The failure run restarts on every success, so keying the traceback off the
+    run's first round would emit one per failing round at the default interval.
+    """
+    manager, tier = _make_manager()
+    fake_logger = MagicMock()
+    monkeypatch.setattr(tiering_manager, "logger", fake_logger)
+    manager._stop_tier_poll_thread()
+    calls = {"n": 0}
+
+    def alternating(parent):
+        calls["n"] += 1
+        if calls["n"] % 2:
+            raise RuntimeError("intermittent")
+
+    tier.on_serve = alternating
+    stop = _RecordingStop(6)
+    manager._tier_poll_stop = stop
+    try:
+        manager._tier_poll_loop()
     finally:
         manager.shutdown()
+    assert calls["n"] == 6
+    assert fake_logger.exception.call_count == 1
 
 
 def test_lock_releases_on_error(manager_and_tier):
diff --git a/tests/v1/kv_offload/tiering/test_tiering_offloading.py b/tests/v1/kv_offload/tiering/test_tiering_offloading.py
index 1de7850d7972..9795cd8bfbe0 100644
--- a/tests/v1/kv_offload/tiering/test_tiering_offloading.py
+++ b/tests/v1/kv_offload/tiering/test_tiering_offloading.py
@@ -10,6 +10,7 @@
 5. Eviction coordination between tiers
 """
 
+import threading
 from collections.abc import Iterable
 from unittest.mock import MagicMock, patch
 
@@ -40,8 +41,7 @@
     TierMatcher,
     make_offload_key,
 )
-from vllm.v1.kv_offload.cpu.common import CPUOffloadingInfo
-from vllm.v1.kv_offload.cpu.policies.lru import LRUCachePolicy
+from vllm.v1.kv_offload.cpu.common import CPULoadStoreSpec, CPUOffloadingInfo
 from vllm.v1.kv_offload.tiering.base import (
     JobResult,
     SecondaryTierManager,
@@ -187,6 +187,7 @@ def test_tiering_spec_collects_secondary_metric_definitions(monkeypatch):
         TieringOffloadingMetrics.PRIMARY_READ_USAGE_PERC,
         TieringOffloadingMetrics.ACTIVE_PROMOTION_JOBS,
         TieringOffloadingMetrics.ACTIVE_CASCADE_JOBS,
+        TieringOffloadingMetrics.DEFERRED_RELEASE_CHUNKS,
     ):
         assert isinstance(metrics[metric_name], OffloadingGaugeMetadata)
         assert metrics[metric_name].labelnames == ("tier",)
@@ -656,6 +657,154 @@ def submit_partial(job_metadata: TransferJob) -> None:
             expected_source_keys
         )
 
+    def _promote_with_deferred_release(
+        self, monkeypatch, chunks, successful_indices=(), preset=False
+    ) -> list[threading.Event]:
+        """Promote ``chunks`` from tier1 into a failure with a deferred release.
+
+        Returns the events the tier reported; setting one releases its job.
+        """
+        for chunk in chunks:
+            self.secondary_tier1.chunks[chunk] = True
+        events: list[threading.Event] = []
+
+        def submit_deferring(job_metadata: TransferJob) -> None:
+            event = threading.Event()
+            if preset:
+                event.set()
+            events.append(event)
+            self.secondary_tier1.completed_jobs.append(
+                JobResult(
+                    job_id=job_metadata.job_id,
+                    success=False,
+                    successful_keys=tuple(chunks[i] for i in successful_indices)
+                    or None,
+                    deferred_release=event,
+                )
+            )
+
+        monkeypatch.setattr(self.secondary_tier1, "submit_load", submit_deferring)
+        for chunk in chunks:
+            assert self.manager.lookup(chunk, _CTX) is LookupResult.HIT_PENDING
+        self._simulate_on_schedule_end()
+        self._simulate_on_schedule_end()
+        return events
+
+    def _reduced_stats(self) -> dict[str, int | float]:
+        stats = self.manager.get_stats()
+        assert stats is not None
+        return stats.reduce()
+
+    def _release(self, events: list[threading.Event]) -> None:
+        for event in events:
+            event.set()
+        self.manager._process_finished_jobs()
+
+    def test_deferred_release_misses_but_holds_chunks_until_released(
+        self, manager_setup, monkeypatch
+    ):
+        """A deferred release fails the keys at once but keeps their chunks out
+        of the pool until the tier sets the event."""
+        self._start_request()
+        chunks = to_keys(range(2))
+        free_before = self.primary_tier._get_num_free_chunks()
+
+        events = self._promote_with_deferred_release(monkeypatch, chunks)
+
+        assert len(events) == 1
+        assert [self.primary_tier.lookup(c, _CTX) for c in chunks] == [
+            LookupResult.MISS,
+            LookupResult.MISS,
+        ]
+        assert self.primary_tier.num_deferred_release_chunks == 2
+        assert self.primary_tier._get_num_free_chunks() == free_before - 2
+        assert self.manager.has_pending_work()
+        self.manager._process_finished_jobs()
+        assert self.primary_tier.num_deferred_release_chunks == 2
+        reduced = self._reduced_stats()
+        label = self.manager._metrics.tier_label(0)
+        assert (
+            reduced[f"{TieringOffloadingMetrics.DEFERRED_RELEASE_CHUNKS}:{label}"] == 2
+        )
+        hold = TieringOffloadingMetrics.DEFERRED_RELEASE_HOLD
+        assert not has_histogram_count(reduced, hold)
+
+        self._release(events)
+
+        reduced = self._reduced_stats()
+        assert reduced[histogram_count_key(hold, label)] == 1
+
+        assert self.primary_tier.num_deferred_release_chunks == 0
+        assert self.primary_tier._get_num_free_chunks() == free_before
+        assert self.primary_tier._num_write_pending_chunks == 0
+        assert not self.manager._deferred_releases
+
+    def test_deferred_release_restore_gets_a_fresh_chunk(
+        self, manager_setup, monkeypatch
+    ):
+        """A key whose old chunk awaits release can be stored again, into a
+        different chunk the late writer cannot touch."""
+        self._start_request()
+        chunk = to_keys([7])
+        events = self._promote_with_deferred_release(monkeypatch, chunk)
+        ((held,),) = (d.chunk_ids for d in self.manager._deferred_releases.values())
+
+        output = self.primary_tier.prepare_write(chunk, _CTX)
+        assert output is not None
+        assert isinstance(output.store_spec, CPULoadStoreSpec)
+        assert list(output.store_spec.chunk_ids) != [held]
+        self.primary_tier.complete_write(chunk, _CTX, True)
+
+        self._release(events)
+        assert self.primary_tier.lookup(chunk[0], _CTX) is LookupResult.HIT
+
+    def test_deferred_release_partial_failure_holds_only_failed_chunks(
+        self, manager_setup, monkeypatch
+    ):
+        self._start_request()
+        chunks = to_keys(range(3))
+        events = self._promote_with_deferred_release(
+            monkeypatch, chunks, successful_indices=(0,)
+        )
+
+        assert self.primary_tier.lookup(chunks[0], _CTX) is LookupResult.HIT
+        assert self.primary_tier.num_deferred_release_chunks == 2
+        self._release(events)
+        assert self.primary_tier.num_deferred_release_chunks == 0
+
+    def test_deferred_release_already_set_with_result(self, manager_setup, monkeypatch):
+        """An event already set when its result is reported releases in the
+        same poll round."""
+        self._start_request()
+        self._promote_with_deferred_release(monkeypatch, to_keys(range(2)), preset=True)
+        assert self.primary_tier.num_deferred_release_chunks == 0
+        assert not self.manager._deferred_releases
+
+    def test_defer_write_release_asserts_on_non_pending_key(self, manager_setup):
+        """Every failed key must still be pending: a ready or missing key
+        means a chunk the tier may still write to would go unheld."""
+        self._start_request()
+        chunk = to_keys([5])
+        assert self.primary_tier.prepare_write(chunk, _CTX) is not None
+        self.primary_tier.complete_write(chunk, _CTX, True)
+        with pytest.raises(AssertionError, match="is not a pending write"):
+            self.primary_tier.defer_write_release(chunk)
+        with pytest.raises(AssertionError, match="is not a pending write"):
+            self.primary_tier.defer_write_release(to_keys([6]))
+
+    def test_reset_after_deferred_release_succeeds(self, manager_setup, monkeypatch):
+        """drain_jobs() completing the deferred release lets reset_cache proceed."""
+        self._start_request()
+        events = self._promote_with_deferred_release(monkeypatch, to_keys(range(2)))
+
+        def drain_jobs() -> None:
+            for event in events:
+                event.set()
+
+        monkeypatch.setattr(self.secondary_tier1, "drain_jobs", drain_jobs)
+        self.manager.reset_cache()
+        assert self.primary_tier.num_deferred_release_chunks == 0
+
     def test_lookup_reports_sync_delay_for_resolved_lookups(self, manager_setup):
         """Resolved lookups report one sync delay sample per tier and chunk."""
         self._start_request()
@@ -821,36 +970,6 @@ def test_eviction_in_primary_tier(self, manager_setup):
         assert len(result.evicted_keys) == 2
         assert len(result.keys_to_store) == 2
 
-    def test_touch_propagates_to_all_tiers(self, manager_setup, monkeypatch):
-        """Test that touch() propagates to all tiers."""
-        chunks = to_keys(range(3))
-
-        # Store chunks
-        self._start_request()
-        self.manager.prepare_store(chunks, _CTX)
-        self.manager.complete_store(chunks, _CTX, success=True)
-        self._simulate_on_schedule_end()
-        # for secondary tiers to drain jobs, so primary tier's chunks are evictable.
-        self._simulate_on_schedule_end()
-
-        tier1_touch = MagicMock(wraps=self.secondary_tier1.touch)
-        tier2_touch = MagicMock(wraps=self.secondary_tier2.touch)
-        monkeypatch.setattr(self.secondary_tier1, "touch", tier1_touch)
-        monkeypatch.setattr(self.secondary_tier2, "touch", tier2_touch)
-
-        # Touch chunks
-        self.manager.touch(chunks, _CTX)
-
-        # 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
-        tier1_touch.assert_called_once_with(chunks, _CTX)
-        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))
diff --git a/tests/v1/metrics/test_engine_logger_apis.py b/tests/v1/metrics/test_engine_logger_apis.py
index 2d7f077e6c33..b713d649704b 100644
--- a/tests/v1/metrics/test_engine_logger_apis.py
+++ b/tests/v1/metrics/test_engine_logger_apis.py
@@ -8,11 +8,20 @@
     DummyStatLogger,
 )
 from tests.utils import wait_for_memory_to_settle
+from vllm.config import CacheConfig
 from vllm.v1.engine.async_llm import AsyncEngineArgs, AsyncLLM
 from vllm.v1.metrics.loggers import PerEngineStatLoggerAdapter
 from vllm.v1.metrics.ray_wrappers import RayPrometheusStatLogger
 
 
+def _required_utilization(engine_args: AsyncEngineArgs) -> float:
+    # Unset resolves to 1.0 under the extensible KV cache, which does not need
+    # that memory free up front; wait for the standard fraction instead.
+    return (
+        engine_args.gpu_memory_utilization or CacheConfig.DEFAULT_GPU_MEMORY_UTILIZATION
+    )
+
+
 @pytest.fixture
 def log_stats_enabled_engine_args():
     """Shared fixture providing common AsyncEngineArgs configuration
@@ -40,7 +49,7 @@ async def test_async_llm_replace_default_loggers(log_stats_enabled_engine_args):
     finally:
         engine.shutdown()
         wait_for_memory_to_settle(
-            threshold_ratio=1.0 - log_stats_enabled_engine_args.gpu_memory_utilization
+            threshold_ratio=1.0 - _required_utilization(log_stats_enabled_engine_args)
         )
 
 
@@ -75,5 +84,5 @@ async def test_async_llm_add_to_default_loggers(log_stats_enabled_engine_args):
     finally:
         engine.shutdown()
         wait_for_memory_to_settle(
-            threshold_ratio=1.0 - disabled_log_engine_args.gpu_memory_utilization
+            threshold_ratio=1.0 - _required_utilization(disabled_log_engine_args)
         )
diff --git a/tests/v1/metrics/test_ray_metrics.py b/tests/v1/metrics/test_ray_metrics.py
index 8552e896c0ab..4f19e31f27fc 100644
--- a/tests/v1/metrics/test_ray_metrics.py
+++ b/tests/v1/metrics/test_ray_metrics.py
@@ -7,6 +7,7 @@
 import ray
 
 from tests.utils import wait_for_memory_to_settle
+from vllm.config import CacheConfig
 from vllm.config.model import ModelDType
 from vllm.platforms import current_platform
 from vllm.sampling_params import SamplingParams
@@ -77,8 +78,14 @@ async def run(self):
         ray.get(actor.run.remote())
     finally:
         ray.shutdown()
+        # Unset resolves to 1.0 under the extensible KV cache, which does not
+        # need that memory free up front; wait for the standard fraction.
         wait_for_memory_to_settle(
-            threshold_ratio=1.0 - engine_args.gpu_memory_utilization
+            threshold_ratio=1.0
+            - (
+                engine_args.gpu_memory_utilization
+                or CacheConfig.DEFAULT_GPU_MEMORY_UTILIZATION
+            )
         )
 
 
diff --git a/tests/v1/sample/test_logprobs.py b/tests/v1/sample/test_logprobs.py
index f7f51ac47d88..4f8c378589cf 100644
--- a/tests/v1/sample/test_logprobs.py
+++ b/tests/v1/sample/test_logprobs.py
@@ -1628,3 +1628,128 @@ def test_token_logprobs_large_batch_int64_row_offset():
     assert torch.allclose(logprobs[last, 0], ref, atol=1e-2), (
         f"logprob {logprobs[last, 0].item()} != ref {ref.item()}"
     )
+
+
+@pytest.mark.parametrize("split_topk", [False, True])
+@pytest.mark.parametrize("num_prompt_logprobs", [-1, 3])
+def test_chunked_prompt_logprobs_match_unchunked(
+    num_prompt_logprobs: int, split_topk: bool, monkeypatch
+):
+    """Chunks, and top-k blocks within them, written into preallocated outputs
+    must match one unchunked pass."""
+    if not current_platform.is_cuda_alike():
+        pytest.skip("the logprob kernels are Triton")
+    from vllm.v1.worker.gpu.sample import prompt_logprob
+    from vllm.v1.worker.gpu.sample.logprob import compute_topk_scores
+    from vllm.v1.worker.gpu.sample.prompt_logprob import (
+        CHUNK_SIZE,
+        compute_prompt_logprobs_with_chunking,
+    )
+
+    if split_topk:
+        # Blocks of 100 rows at full vocab, straddling the chunk boundaries.
+        monkeypatch.setattr(prompt_logprob, "TOPK_MAX_ELEMENTS", 100 * 513)
+    device = torch.device(current_platform.device_type)
+    vocab_size = 512
+    num_tokens = 2 * CHUNK_SIZE + 3
+    # Distinct values per row, since topk's order among ties can vary with the
+    # number of rows ranked together.
+    logits = torch.rand(num_tokens, vocab_size, device=device).argsort(dim=-1)
+    logits = logits.float() / vocab_size
+    prompt_token_ids = torch.randint(vocab_size, (num_tokens,), device=device)
+
+    token_ids, logprobs, ranks = compute_prompt_logprobs_with_chunking(
+        prompt_token_ids,
+        logits,
+        lambda hidden_states: hidden_states,
+        num_prompt_logprobs,
+    )
+
+    num_logprobs = vocab_size if num_prompt_logprobs == -1 else num_prompt_logprobs
+    expected = compute_topk_scores(logits, num_logprobs, prompt_token_ids)
+    assert token_ids.dtype == torch.int32
+    torch.testing.assert_close(token_ids, expected.logprob_token_ids.to(torch.int32))
+    torch.testing.assert_close(logprobs, expected.logprobs)
+    torch.testing.assert_close(ranks, expected.selected_token_ranks)
+
+
+def test_trace_replay_with_spec_decode(monkeypatch):
+    monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1")
+
+    prompt = "Hello world " * 50
+    llm_kwargs = dict(
+        max_model_len=512,
+        max_logprobs=5,
+        seed=42,
+        enable_trace_replay=True,
+        enable_chunked_prefill=True,
+        max_num_batched_tokens=32,
+        enable_prefix_caching=False,
+        gpu_memory_utilization=0.25,
+        **GPU_DETERMINISM_KWARGS,
+    )
+
+    def replay(runner, traces):
+        outputs = runner.llm.generate(
+            [prompt] * len(traces),
+            [
+                SamplingParams(
+                    trace_decode_token_ids=trace, max_tokens=len(trace), logprobs=3
+                )
+                for trace in traces
+            ],
+        )
+        results = []
+        for output in outputs:
+            completion = output.outputs[0]
+            assert completion.logprobs is not None
+            results.append((list(completion.token_ids), completion.logprobs))
+        return results
+
+    with VllmRunner("Qwen/Qwen3-0.6B", **llm_kwargs) as ref_runner:
+        sampled = ref_runner.llm.generate(
+            [prompt, prompt],
+            [
+                SamplingParams(temperature=0.0, max_tokens=32, ignore_eos=True),
+                SamplingParams(temperature=1.0, seed=7, max_tokens=32, ignore_eos=True),
+            ],
+        )
+        greedy_trace = list(sampled[0].outputs[0].token_ids)
+        sampled_trace = list(sampled[1].outputs[0].token_ids)
+        diverged_trace = greedy_trace[:10] + sampled_trace[10:20] + greedy_trace[20:]
+        traces = [greedy_trace, sampled_trace, diverged_trace]
+        ref_results = replay(ref_runner, traces)
+
+    with VllmRunner(
+        "Qwen/Qwen3-0.6B",
+        speculative_config={
+            "method": "ngram",
+            "prompt_lookup_max": 5,
+            "prompt_lookup_min": 3,
+            "num_speculative_tokens": 3,
+        },
+        disable_log_stats=False,
+        **llm_kwargs,
+    ) as spec_runner:
+        spec_results = replay(spec_runner, traces)
+        num_accepted = sum(
+            m.value
+            for m in spec_runner.llm.get_metrics()
+            if m.name == "vllm:spec_decode_num_accepted_tokens"
+        )
+    assert num_accepted > 0
+
+    for trace, (ref_ids, ref_lps), (spec_ids, spec_lps) in zip(
+        traces, ref_results, spec_results
+    ):
+        assert ref_ids == trace
+        assert spec_ids == trace
+        assert len(ref_lps) == len(trace)
+        assert len(spec_lps) == len(trace)
+        for token_id, ref_lp, spec_lp in zip(trace, ref_lps, spec_lps):
+            assert math.isclose(
+                ref_lp[token_id].logprob,
+                spec_lp[token_id].logprob,
+                rel_tol=5e-2,
+                abs_tol=2.5e-1,
+            )
diff --git a/tests/v1/sample/test_trace_replay_params.py b/tests/v1/sample/test_trace_replay_params.py
index 43f8266dede5..36e0749e556d 100644
--- a/tests/v1/sample/test_trace_replay_params.py
+++ b/tests/v1/sample/test_trace_replay_params.py
@@ -102,13 +102,12 @@ def test_validate_trace_replay_noop_when_unset():
     )
 
 
-def test_trace_decode_token_ids_rejects_speculative_decoding():
+def test_trace_decode_token_ids_allows_speculative_decoding():
     params = SamplingParams(trace_decode_token_ids=[1])
-    with pytest.raises(ValueError, match="not supported with speculative decoding"):
-        params._validate_trace_replay(
-            _make_model_config(vocab_size=100),
-            speculative_config=object(),
-        )
+    params._validate_trace_replay(
+        _make_model_config(vocab_size=100),
+        speculative_config=object(),
+    )
 
 
 def test_trace_decode_token_ids_rejects_structured_outputs():
diff --git a/tests/v1/simple_kv_offload/test_integration.py b/tests/v1/simple_kv_offload/test_integration.py
index c9119b7bff23..d0916a6e4406 100644
--- a/tests/v1/simple_kv_offload/test_integration.py
+++ b/tests/v1/simple_kv_offload/test_integration.py
@@ -10,8 +10,11 @@
 
 from tests.models.utils import check_logprobs_close
 from vllm import LLM, SamplingParams, TokensPrompt, envs
-from vllm.config import KVTransferConfig
+from vllm.config import KVTransferConfig, ParallelConfig
 from vllm.platforms import current_platform
+from vllm.utils import numa_utils
+from vllm.v1.metrics.reader import Counter
+from vllm.v1.simple_kv_offload.metrics import MetricName
 
 if not current_platform.is_cuda_alike():
     pytest.skip("Requires CUDA or ROCm", allow_module_level=True)
@@ -113,7 +116,7 @@ def reset_gpu_cache():
         assert mixed_hit.num_cached_tokens == 16
 
 
-def _make_llm(model: str, lazy: bool, cpu_bytes_to_use: int) -> LLM:
+def _make_llm(model: str, lazy: bool, cpu_bytes_to_use: int, **llm_kwargs) -> LLM:
     kv_transfer_config = KVTransferConfig(
         kv_connector="SimpleCPUOffloadConnector",
         kv_role="kv_both",
@@ -122,12 +125,18 @@ def _make_llm(model: str, lazy: bool, cpu_bytes_to_use: int) -> LLM:
             "lazy_offload": lazy,
         },
     )
+    llm_kwargs = (
+        dict(
+            kv_cache_memory_bytes=40 << 30,  # 40 GiB
+            disable_hybrid_kv_cache_manager=False,
+            enable_prefix_caching=True,
+        )
+        | llm_kwargs
+    )
     return LLM(
         model=model,
-        kv_cache_memory_bytes=40 << 30,  # 40 GiB
-        disable_hybrid_kv_cache_manager=False,
-        enable_prefix_caching=True,
         kv_transfer_config=kv_transfer_config,
+        **llm_kwargs,
     )
 
 
@@ -231,6 +240,90 @@ def _latency_test(llm: LLM, lazy: bool = False):
     )
 
 
+def _assert_worker_auto_numa(worker):
+    import subprocess
+    from pathlib import Path
+
+    import psutil
+
+    from vllm.utils.cpu_resource_utils import parse_id_list
+
+    parallel_config = worker.parallel_config
+    assert parallel_config.numa_bind and parallel_config._numa_bind_auto
+    gpu_index = parallel_config.get_local_gpu_index(worker.local_rank)
+    node = parallel_config.numa_bind_nodes[gpu_index]
+    node_cpus = parse_id_list(
+        Path(f"/sys/devices/system/node/node{node}/cpulist").read_text().strip()
+    )
+    cpus = set(psutil.Process().cpu_affinity())
+    assert cpus and cpus <= set(node_cpus)
+    result = subprocess.run(
+        ["numactl", "--show"], check=True, capture_output=True, text=True, timeout=10
+    )
+    policy = dict(line.split(":", 1) for line in result.stdout.splitlines())
+    assert policy["policy"].strip() == "preferred"
+    assert policy["preferred node"].strip() == str(node)
+    return {"node": node, "cpus": sorted(cpus), "policy": "preferred"}
+
+
+@pytest.mark.slow_test
+def test_simple_cpu_offload_with_auto_numa(monkeypatch, enable_pickle):
+    """Soft-bound NVIDIA workers must store and restore KV, not silently recompute."""
+    if not current_platform.is_cuda():
+        pytest.skip("Automatic NUMA KV-offload smoke test requires NVIDIA")
+    monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
+    monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "1")
+    preflight = ParallelConfig(distributed_executor_backend="mp")
+    preflight._numa_bind_auto = True
+    if not numa_utils._can_auto_bind(preflight, local_rank=0) or not (
+        numa_utils._probe_numactl_args(
+            numa_utils._get_numactl_worker_args(preflight, local_rank=0)
+        )
+    ):
+        pytest.skip("Host cannot engage automatic NUMA memory preference")
+
+    model = "Qwen/Qwen3-0.6B"
+    llm = _make_llm(
+        model,
+        False,
+        256 << 20,
+        kv_cache_memory_bytes=128 << 20,
+        max_model_len=1024,
+        distributed_executor_backend="mp",
+        enforce_eager=True,
+        disable_log_stats=False,
+    )
+
+    def loaded_blocks(llm: LLM):
+        return sum(
+            metric.value
+            for metric in llm.get_metrics()
+            if isinstance(metric, Counter)
+            and metric.name == MetricName.LOAD_BLOCKS.removesuffix("_total")
+            and metric.labels.get("model_name") == model
+        )
+
+    try:
+        states = llm.collective_rpc(_assert_worker_auto_numa, timeout=30)
+        print(f"NUMA worker state: {states}")
+        sampling_params = SamplingParams(max_tokens=8, temperature=0)
+        prompt = "hi " * 256 + "Let's count to ten. One, two, three, "
+        cold = llm.generate(prompt, sampling_params, use_tqdm=False)[0]
+        llm.generate(prompt, sampling_params, use_tqdm=False)
+        deadline = time.monotonic() + 30
+        while not llm.reset_prefix_cache(reset_connector=False):
+            assert time.monotonic() < deadline, "CPU offload did not drain"
+            llm.generate("0", sampling_params, use_tqdm=False)
+        loaded_before = loaded_blocks(llm)
+        restored = llm.generate(prompt, sampling_params, use_tqdm=False)[0]
+        loaded = loaded_blocks(llm) - loaded_before
+        assert loaded > 0
+        assert restored.outputs[0].token_ids == cold.outputs[0].token_ids
+        print(f"KV offload blocks restored={loaded}; tokens match")
+    finally:
+        del llm
+
+
 @pytest.mark.optional
 @pytest.mark.slow_test
 @pytest.mark.parametrize("model", SMALL_MODELS)
diff --git a/tests/v1/simple_kv_offload/test_scheduler.py b/tests/v1/simple_kv_offload/test_scheduler.py
index c63cd614e77d..df18add326aa 100644
--- a/tests/v1/simple_kv_offload/test_scheduler.py
+++ b/tests/v1/simple_kv_offload/test_scheduler.py
@@ -11,6 +11,7 @@
 import pytest
 import torch
 
+from tests.v1.spec_decode.utils import attach_speculator_config
 from vllm import SamplingParams
 from vllm.config import (
     CacheConfig,
@@ -31,6 +32,7 @@
 from vllm.v1.core.block_pool import BlockPool
 from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager
 from vllm.v1.core.kv_cache_utils import (
+    KVCacheBlock,
     get_request_block_hasher,
     init_none_hash,
     make_block_hash_with_group_id,
@@ -53,6 +55,7 @@
     KVCacheSpec,
     KVCacheTensor,
     MambaSpec,
+    SlidingWindowMLASpec,
     SlidingWindowSpec,
 )
 from vllm.v1.metrics.cache_hit_source import CacheHitSource
@@ -60,6 +63,10 @@
 from vllm.v1.request import Request
 from vllm.v1.simple_kv_offload.manager import SimpleCPUOffloadScheduler
 from vllm.v1.simple_kv_offload.metadata import SimpleCPUOffloadWorkerMetadata
+from vllm.v1.worker.gpu.spec_decode.mtp.config import MTPSpeculatorConfig
+from vllm.v1.worker.gpu.spec_decode.multi_module_mtp.config import (
+    MultiModuleMTPSpeculatorConfig,
+)
 
 pytestmark = pytest.mark.skip_global_cleanup
 
@@ -1136,6 +1143,36 @@ def test_lazy_duplicate_store_skipped() -> None:
     assert get_cpu_free_blocks(sched) == cpu_free_after_first
 
 
+# ---------------------------------------------------------------------------
+# Test 2d: Lazy walk covers bounded-replay SWA allocations
+# ---------------------------------------------------------------------------
+def test_lazy_walk_covers_bounded_replay_swa_groups() -> None:
+    """Lazy: the walk covers what every group, offloaded or not, takes per step."""
+    config = _make_kv_cache_config(num_blocks=32, num_groups=3)
+    for group in config.kv_cache_groups[1:]:
+        group.kv_cache_spec = SlidingWindowMLASpec(
+            block_size=BLOCK_SIZE,
+            num_kv_heads=1,
+            head_size=HEAD_SIZE,
+            dtype=DTYPE,
+            sliding_window=BLOCK_SIZE * 4,
+            bounded_replay=True,
+        )
+    fix = make_scheduler(
+        num_cpu_blocks=32, num_gpu_blocks=32, lazy=True, kv_cache_config=config
+    )
+    gpu_pool = fix.gpu_block_pool
+
+    # One step takes cdiv(64, 16) blocks per group.
+    num_blocks = 3 * 4
+    blocks = _allocate_gpu_blocks(gpu_pool, make_request(num_blocks), num_blocks)
+    _flush_old_blocks_to_lru_head(gpu_pool, gpu_pool.get_num_free_blocks())
+    gpu_pool.free_blocks(blocks)
+
+    meta = fix.scheduler.build_connector_meta(make_scheduler_output({}))
+    assert sorted(meta.store_gpu_blocks) == [b.block_id for b in blocks]
+
+
 # ---------------------------------------------------------------------------
 # Test 3: LRU eviction order
 # ---------------------------------------------------------------------------
@@ -1298,6 +1335,68 @@ def test_touched_blocks_survive_eviction() -> None:
         assert cached is not None, f"req_a block {i} should survive (was touched/MRU)"
 
 
+def test_gpu_prefix_hit_refreshes_cpu_lru() -> None:
+    """Eager mode: a prefix that hits on GPU keeps its CPU copy at the MRU end."""
+    fix = make_scheduler(num_cpu_blocks=5, num_gpu_blocks=16, lazy=False)
+    sched = fix.scheduler
+
+    # Fill all 4 usable CPU blocks: req_a is LRU, req_b is MRU.
+    req_a = make_request(num_blocks=2)
+    req_b = make_request(num_blocks=2)
+    kv_a = _alloc_and_register(fix, req_a, 2)
+    kv_b = _alloc_and_register(fix, req_b, 2)
+    sched.update_state_after_alloc(req_a, kv_a, num_external_tokens=0)
+    sched.update_state_after_alloc(req_b, kv_b, num_external_tokens=0)
+    meta = sched.build_connector_meta(
+        make_scheduler_output(
+            {req_a.request_id: 2 * BLOCK_SIZE, req_b.request_id: 2 * BLOCK_SIZE},
+            new_reqs={
+                req_a.request_id: kv_a.get_block_ids(),
+                req_b.request_id: kv_b.get_block_ids(),
+            },
+        )
+    )
+    simulate_store_completion(sched, meta.store_event)
+
+    # req_a's prefix is reused and hits entirely on GPU.
+    req_a2 = Request(
+        request_id="req-a2",
+        prompt_token_ids=req_a.prompt_token_ids,
+        sampling_params=req_a.sampling_params,
+        pooling_params=None,
+        mm_features=None,
+        block_hasher=get_request_block_hasher(BLOCK_SIZE, sha256),
+    )
+    assert sched.get_num_new_matched_tokens(req_a2, 2 * BLOCK_SIZE) == (0, False)
+    sched.update_state_after_alloc(
+        req_a2, KVCacheBlocks(blocks=([],)), num_external_tokens=0
+    )
+
+    # Storing 2 new blocks evicts the 2 LRU blocks, which are now req_b's.
+    req_c = make_request(num_blocks=2)
+    kv_c = _alloc_and_register(fix, req_c, 2)
+    sched.update_state_after_alloc(req_c, kv_c, num_external_tokens=0)
+    meta2 = sched.build_connector_meta(
+        make_scheduler_output(
+            {req_c.request_id: 2 * BLOCK_SIZE},
+            new_reqs={req_c.request_id: kv_c.get_block_ids()},
+        )
+    )
+    simulate_store_completion(sched, meta2.store_event)
+
+    cache_map = sched.cpu_block_pool.cached_block_hash_to_block
+
+    def cpu_cached(req: Request) -> list[bool]:
+        return [
+            cache_map.get_one_block(make_block_hash_with_group_id(h, 0)) is not None
+            for h in req.block_hashes[:2]
+        ]
+
+    assert cpu_cached(req_a) == [True, True]
+    assert cpu_cached(req_b) == [False, False]
+    assert cpu_cached(req_c) == [True, True]
+
+
 # ---------------------------------------------------------------------------
 # Test 5: Preemption no CPU block leak
 # ---------------------------------------------------------------------------
@@ -2935,6 +3034,45 @@ def test_finished_eager_store_does_not_duplicate_completed_partial_tail() -> Non
     assert len(finish_meta.store_cpu_blocks) == len(stored)
 
 
+def test_gpu_prefix_hit_refreshes_cpu_mamba_boundary() -> None:
+    """A hybrid GPU prefix hit refreshes the CPU attention blocks and the Mamba
+    state at the hit boundary, not older Mamba states, once admitted."""
+    block_size = 4 * BLOCK_SIZE
+    fix = _make_hybrid_attention_mamba_scheduler(block_size=block_size)
+    sched = fix.scheduler
+    cpu_pool = sched.cpu_block_pool
+    req = _make_cp_request(num_blocks=2, virtual_block_size=block_size)
+    other = _make_cp_request(num_blocks=2, virtual_block_size=block_size)
+
+    def cache_on_cpu(request: Request, group_id: int) -> list[KVCacheBlock]:
+        blocks = cpu_pool.get_new_blocks(2)
+        cpu_pool.cache_full_blocks(
+            request=request,
+            blocks=blocks,
+            num_cached_blocks=0,
+            num_full_blocks=2,
+            block_size=sched.group_block_sizes[group_id],
+            kv_cache_group_id=group_id,
+        )
+        return blocks
+
+    attention, mamba = cache_on_cpu(req, 0), cache_on_cpu(req, 1)
+    newer = cache_on_cpu(other, 0)
+    cpu_pool.free_blocks([*attention, *mamba])
+    cpu_pool.free_blocks(newer)
+
+    lru_before = cpu_pool.free_block_queue.get_all_free_blocks()
+    assert sched.get_num_new_matched_tokens(req, 2 * block_size) == (0, False)
+    assert cpu_pool.free_block_queue.get_all_free_blocks() == lru_before
+    sched.update_state_after_alloc(
+        req, KVCacheBlocks(blocks=([], [])), num_external_tokens=0
+    )
+
+    lru_order = cpu_pool.free_block_queue.get_all_free_blocks()
+    assert lru_order[-3:] == [mamba[1], attention[1], attention[0]]
+    assert lru_order.index(mamba[0]) < lru_order.index(newer[0])
+
+
 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)
@@ -3361,6 +3499,12 @@ def test_eager_store_matches_gpu_cacheable_prefix(
     speculative_config.use_eagle.return_value = True
     speculative_config.use_eagle_block_drop.return_value = True
     speculative_config.use_multi_module_mtp.return_value = prefill_lookahead > 1
+    attach_speculator_config(
+        speculative_config,
+        MultiModuleMTPSpeculatorConfig
+        if prefill_lookahead > 1
+        else MTPSpeculatorConfig,
+    )
     vllm_config.speculative_config = speculative_config
     scheduler = SimpleCPUOffloadScheduler(
         vllm_config=vllm_config,
diff --git a/tests/v1/spec_decode/test_acceptance_estimator.py b/tests/v1/spec_decode/test_acceptance_estimator.py
index 022ed463879c..41e42a1b197f 100644
--- a/tests/v1/spec_decode/test_acceptance_estimator.py
+++ b/tests/v1/spec_decode/test_acceptance_estimator.py
@@ -15,6 +15,8 @@
 from vllm.platforms import current_platform
 from vllm.v1.worker.gpu.spec_decode.acceptance_estimator import (
     OnlineAcceptanceEstimator,
+    _accumulate_kernel,
+    _refit_kernel,
 )
 
 pytest.importorskip("triton")
@@ -295,3 +297,34 @@ def test_feature_ignores_which_token_was_drawn():
         atol=2e-3,
         rtol=2e-3,
     )
+
+
+def test_warmup_covers_every_batch_size_and_the_refit():
+    """After warmup, no batch size and no refit may JIT-compile during serving."""
+    device = current_platform.device_type
+
+    def compiled_variants() -> set:
+        return {
+            (kernel.fn.__name__, key)
+            for kernel in (_accumulate_kernel, _refit_kernel)
+            for caches in kernel.device_caches.values()
+            for key in caches[0]
+        }
+
+    estimator = OnlineAcceptanceEstimator(REFIT_NUM_REQS, NUM_STEPS, device)
+    estimator.REFIT_INTERVAL = 2
+    slots = torch.arange(REFIT_NUM_REQS, device=device, dtype=torch.int32)
+    num_accepted = torch.ones(REFIT_NUM_REQS, device=device)
+    features = torch.rand(REFIT_NUM_REQS, NUM_STEPS, device=device)
+    _drive_round(estimator, features[:1], num_accepted[:1], slots[:1])
+    estimator.warmup()
+    warmed = compiled_variants()
+
+    for num_reqs in (8, 300):
+        _drive_round(
+            estimator, features[:num_reqs], num_accepted[:num_reqs], slots[:num_reqs]
+        )
+    torch.accelerator.synchronize()
+
+    assert estimator._refits == 1
+    assert compiled_variants() == warmed
diff --git a/tests/v1/spec_decode/test_custom_speculator.py b/tests/v1/spec_decode/test_custom_speculator.py
new file mode 100644
index 000000000000..9405adb431f2
--- /dev/null
+++ b/tests/v1/spec_decode/test_custom_speculator.py
@@ -0,0 +1,364 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for speculator configs and custom V2 speculators registered with
+`register_speculative_method`."""
+
+import subprocess
+import sys
+import textwrap
+from types import SimpleNamespace
+from typing import Any
+from unittest.mock import patch
+
+import pytest
+import torch
+
+import vllm.envs as envs
+from tests.v1.spec_decode.utils import fake_speculative_config
+from vllm.config import SpeculativeConfig, VllmConfig
+from vllm.model_executor.models.registry import ModelRegistry
+from vllm.v1.worker.gpu.spec_decode import init_speculator
+from vllm.v1.worker.gpu.spec_decode import speculator_config as speculator_config_module
+from vllm.v1.worker.gpu.spec_decode.ngram.config import NgramGPUSpeculatorConfig
+from vllm.v1.worker.gpu.spec_decode.speculator import BaseSpeculator
+from vllm.v1.worker.gpu.spec_decode.speculator_config import SpeculatorConfig
+from vllm.v1.worker.gpu.spec_decode.standalone_ar.config import (
+    StandaloneARSpeculatorConfig,
+)
+from vllm.v1.worker.gpu.spec_decode.v1_speculator_configs import (
+    V1CustomProposerConfig,
+)
+
+
+def _qualname(cls: type) -> str:
+    return f"{__name__}.{cls.__qualname__}"
+
+
+class DummySpeculator(BaseSpeculator):
+    def __init__(self, vllm_config: VllmConfig, device: torch.device, req_states):
+        self.vllm_config = vllm_config
+        self.req_states = req_states
+
+    def init_cudagraph_manager(self, cudagraph_mode) -> None:
+        pass
+
+    def capture(self) -> None:
+        pass
+
+    def propose(self, *args: Any, **kwargs: Any) -> torch.Tensor:
+        raise NotImplementedError
+
+
+class DummySpeculatorConfig(SpeculatorConfig):
+    speculator_class = _qualname(DummySpeculator)
+
+
+class DummyNgramGPUSpeculatorConfig(NgramGPUSpeculatorConfig):
+    speculator_class = _qualname(DummySpeculator)
+
+
+class DummyStandaloneARSpeculatorConfig(StandaloneARSpeculatorConfig):
+    speculator_class = _qualname(DummySpeculator)
+
+
+class OptionsSpeculatorConfig(DummySpeculatorConfig):
+    supported_options = frozenset({"speculator_options"})
+
+    @classmethod
+    def apply_options(cls, spec: SpeculativeConfig) -> None:
+        if (spec.speculator_options or {}).keys() - {"seed"}:
+            raise ValueError("unknown option")
+
+
+class NoSpeculatorClassConfig(SpeculatorConfig):
+    pass
+
+
+class NotASpeculatorConfig(SpeculatorConfig):
+    speculator_class = _qualname(NoSpeculatorClassConfig)
+
+
+@pytest.fixture
+def register(monkeypatch):
+    """Register custom speculative methods for a test only."""
+    monkeypatch.setattr(speculator_config_module, "_CUSTOM_CONFIGS", {})
+    return speculator_config_module.register_speculative_method
+
+
+def _custom_spec_config(method: str = "dummy", **kwargs: Any) -> SpeculativeConfig:
+    return SpeculativeConfig(method=method, num_speculative_tokens=3, **kwargs)
+
+
+@pytest.mark.parametrize("lazy", [False, True])
+def test_custom_speculator(register, lazy: bool):
+    register(
+        "dummy", _qualname(DummySpeculatorConfig) if lazy else DummySpeculatorConfig
+    )
+    spec_config = _custom_spec_config()
+
+    # Not based on a built-in method.
+    assert spec_config.custom_method == "dummy"
+    assert spec_config.method is None
+    assert type(spec_config.speculator_config) is DummySpeculatorConfig
+
+    config = VllmConfig()
+    config.speculative_config = spec_config
+    assert config.num_lookahead_tokens == 0
+    with patch.object(envs, "VLLM_USE_V2_MODEL_RUNNER", None):
+        assert config.use_v2_model_runner
+    with pytest.raises(
+        ValueError, match="V1 does not support: custom speculative method 'dummy'"
+    ):
+        config._validate_v1_model_runner()
+
+
+def test_custom_speculator_inherits_builtin_config(register):
+    register("my_ngram", DummyNgramGPUSpeculatorConfig)
+    spec_config = _custom_spec_config("my_ngram")
+
+    # Set up and behaves as the built-in speculator, but builds the custom one.
+    assert spec_config.custom_method == "my_ngram"
+    assert spec_config.method == "ngram_gpu"
+    assert spec_config.use_ngram_gpu()
+    assert spec_config.prompt_lookup_min == spec_config.prompt_lookup_max == 5
+    assert spec_config.speculator_config.prefer_dense_token_history
+    assert type(spec_config.speculator_config) is DummyNgramGPUSpeculatorConfig
+
+    # A draft model speculator needs a draft model.
+    register("my_draft", DummyStandaloneARSpeculatorConfig)
+    with pytest.raises(ValueError, match="without speculative model"):
+        _custom_spec_config("my_draft")
+
+
+def test_speculator_options(register):
+    register("dummy", DummySpeculatorConfig)
+    register("with_options", OptionsSpeculatorConfig)
+
+    spec_config = _custom_spec_config("with_options", speculator_options={"seed": 1})
+    assert spec_config.speculator_options == {"seed": 1}
+    other_seed = _custom_spec_config("with_options", speculator_options={"seed": 2})
+    assert spec_config.compute_hash() != other_seed.compute_hash()
+
+    with pytest.raises(ValueError, match="unknown option"):
+        _custom_spec_config("with_options", speculator_options={"nope": 1})
+    # Only speculators that support options accept them.
+    for method in ("dummy", "ngram"):
+        with pytest.raises(ValueError, match="speculator_options is not supported"):
+            _custom_spec_config(method, speculator_options={"seed": 1})
+
+
+def test_register_speculative_method(register):
+    register("dummy", DummySpeculatorConfig)
+    # Registering the same config again (e.g. by another process) is fine.
+    register("dummy", DummySpeculatorConfig)
+    assert speculator_config_module.registered_speculative_methods() == ["dummy"]
+
+    with pytest.raises(ValueError, match="already registered"):
+        register("dummy", DummyNgramGPUSpeculatorConfig)
+    for name in ("eagle3", "deepseek_mtp"):
+        with pytest.raises(ValueError, match="built-in speculative method"):
+            register(name, DummySpeculatorConfig)
+    with pytest.raises(ValueError, match=r"Unknown speculative method 'nope'.*dummy"):
+        _custom_spec_config("nope")
+
+
+@pytest.mark.parametrize(
+    ("config", "kwargs", "error", "match"),
+    [
+        (_qualname(DummySpeculator), {}, TypeError, "must be a subclass of"),
+        (NoSpeculatorClassConfig, {}, ValueError, "must set speculator_class"),
+        (DummySpeculatorConfig, {"dspark_draft_topk": 4}, ValueError, "not supported"),
+    ],
+)
+def test_custom_speculator_validation(register, config, kwargs, error, match):
+    register("dummy", config)
+    with pytest.raises(error, match=match):
+        _custom_spec_config(**kwargs)
+
+
+def test_v1_custom_proposer_unsupported_by_v2_model_runner():
+    config = VllmConfig()
+    config.speculative_config = SpeculativeConfig(
+        method="custom_class", model="my_module.MyProposer", num_speculative_tokens=3
+    )
+
+    assert type(config.speculative_config.speculator_config) is V1CustomProposerConfig
+    assert config._get_v2_model_runner_unsupported_features() == [
+        "speculative method 'custom_class'"
+    ]
+
+
+def _vllm_config(config_cls: type[SpeculatorConfig], method: str = "dummy") -> Any:
+    return SimpleNamespace(
+        speculative_config=fake_speculative_config(config_cls, method=method)
+    )
+
+
+def test_init_speculator():
+    vllm_config = _vllm_config(DummySpeculatorConfig)
+    req_states = object()
+
+    speculator = init_speculator(vllm_config, torch.device("cpu"), req_states)
+
+    assert type(speculator) is DummySpeculator
+    assert speculator.vllm_config is vllm_config
+    assert speculator.req_states is req_states
+    # Defaults the model runner relies on for any speculator.
+    assert speculator.draft_logits is None
+    assert not speculator.supports_mm_inputs
+    speculator.prepare_watermarking(sampler=None, idx_mapping=torch.empty(0))
+
+
+@pytest.mark.parametrize(
+    ("vllm_config", "error", "match"),
+    [
+        (_vllm_config(NotASpeculatorConfig), TypeError, "must be a subclass of"),
+        (
+            SimpleNamespace(
+                speculative_config=fake_speculative_config(method="suffix")
+            ),
+            NotImplementedError,
+            "suffix is not supported",
+        ),
+    ],
+)
+def test_init_speculator_rejects(vllm_config, error, match):
+    with pytest.raises(error, match=match):
+        init_speculator(vllm_config, torch.device("cpu"), req_states=object())
+
+
+def test_speculator_configs_do_not_import_speculators():
+    """Speculator configs are used in every process (e.g. the scheduler's), so
+    importing them must not import any speculator or the model runner."""
+    script = textwrap.dedent(
+        """
+        import importlib, pkgutil, sys
+
+        import vllm.v1.worker.gpu.spec_decode as pkg
+
+        for module in pkgutil.walk_packages(pkg.__path__, pkg.__name__ + "."):
+            if module.name.endswith(("config", "_configs")):
+                importlib.import_module(module.name)
+
+        suffixes = (".speculator", ".model_runner", ".extract_hidden_states")
+        heavy = [
+            name
+            for name in sys.modules
+            if name.startswith("vllm.v1.worker.gpu.") and name.endswith(suffixes)
+        ]
+        assert not heavy, heavy
+        """
+    )
+    subprocess.run([sys.executable, "-c", script], check=True)
+
+
+_SPEC = "vllm.v1.worker.gpu.spec_decode"
+_DSPARK = f"{_SPEC}.dspark.config.DSparkSpeculatorConfig"
+_LILICORR = f"{_SPEC}.lilicorr.config.LiLiCorrSpeculatorConfig"
+_GEMMA4_MTP = f"{_SPEC}.gemma4.config.Gemma4MTPSpeculatorConfig"
+
+
+def _checkpoint_spec(
+    method: str | None, model: str, declared: str | None = None, model_type="llama"
+) -> Any:
+    return SimpleNamespace(
+        method=method,
+        custom_method=None,
+        draft_model_config=SimpleNamespace(
+            model=model,
+            # The speculator config class declared by the draft model.
+            speculator_config_class=declared,
+            hf_config=SimpleNamespace(model_type=model_type),
+        ),
+        use_gemma4_mtp=lambda: False,
+        use_multi_module_mtp=lambda: False,
+    )
+
+
+@pytest.mark.parametrize(
+    ("method", "model", "declared", "model_type", "expected"),
+    [
+        # No method configured: detected from the draft checkpoint.
+        (None, "org/llama-draft", None, "llama", "draft_model"),
+        (None, "yuhuili/EAGLE-LLaMA3-Instruct-8B", None, "llama", "eagle"),
+        (None, "AngelSlim/Qwen3-8B_eagle3", None, "llama", "eagle3"),
+        (None, "org/mtp", None, "deepseek_mtp", "mtp"),
+        (None, "org/draft", _LILICORR, "x", "dflash"),
+        # An explicitly configured method is never overridden.
+        ("draft_model", "yuhuili/EAGLE-LLaMA3-Instruct-8B", None, "llama", None),
+        ("mtp", "org/draft", _GEMMA4_MTP, "gemma4_mtp", None),
+        # Only checkpoints declaring a speculator are checked.
+        ("mtp", "yuhuili/EAGLE-LLaMA3-Instruct-8B", None, "llama", None),
+        ("eagle", "eagle618/eagle-deepseek-v3-random", None, "deepseek_mtp", None),
+        ("dspark", "org/draft", None, "deepseek_mtp", None),
+    ],
+)
+def test_resolve_checkpoint_config_class(method, model, declared, model_type, expected):
+    spec = _checkpoint_spec(method, model, declared, model_type)
+    configured_cls = speculator_config_module.get_method_config_class(spec)
+
+    resolved_cls = speculator_config_module._resolve_checkpoint_config_class(
+        spec, configured_cls
+    )
+
+    assert resolved_cls.base_method == (expected or method)
+
+
+@pytest.mark.parametrize(
+    ("method", "model", "declared", "model_type", "match"),
+    [
+        ("draft_model", "org/draft", _DSPARK, "qwen3", "is a 'dspark'"),
+        ("eagle", "org/draft", _DSPARK, "llama", "is a 'dspark'"),
+        ("dflash", "org/draft", _GEMMA4_MTP, "gemma4_mtp", "is a 'mtp'"),
+    ],
+)
+def test_resolve_checkpoint_config_class_rejects_mismatch(
+    method, model, declared, model_type, match
+):
+    spec = _checkpoint_spec(method, model, declared, model_type)
+    configured_cls = speculator_config_module.get_method_config_class(spec)
+
+    with pytest.raises(ValueError, match=match):
+        speculator_config_module._resolve_checkpoint_config_class(spec, configured_cls)
+
+
+@pytest.mark.parametrize(
+    ("method", "declared", "multi_module", "expected"),
+    [
+        ("dflash", None, False, "DFlashSpeculatorConfig"),
+        ("dflash", _LILICORR, False, "LiLiCorrSpeculatorConfig"),
+        # A declared config of another speculator than configured is ignored.
+        ("eagle", _DSPARK, False, "EagleSpeculatorConfig"),
+        ("mtp", _GEMMA4_MTP, False, "Gemma4MTPSpeculatorConfig"),
+        ("mtp", None, True, "MultiModuleMTPSpeculatorConfig"),
+    ],
+)
+def test_speculator_config_selection(method, declared, multi_module, expected):
+    spec = _checkpoint_spec(method, "org/draft", declared)
+    spec.use_multi_module_mtp = lambda: multi_module
+
+    config_cls = speculator_config_module.get_speculator_config_class(spec)
+
+    assert config_cls.__name__ == expected
+
+
+@pytest.mark.parametrize(
+    ("architecture", "expected"),
+    [
+        ("MuseGlimmerAssistantModel", "DFlashSpeculatorConfig"),
+        ("LiLiCorrDraftModel", "LiLiCorrSpeculatorConfig"),
+        ("DFlashQwen3XPressModel", "XPressSpeculatorConfig"),
+        ("MuseSparkDSparkModel", "DSparkSpeculatorConfig"),
+        ("DFlash2DraftModel", "DFlash2SpeculatorConfig"),
+        ("Qwen3OmniDSparkModel", "DSparkSpeculatorConfig"),
+        ("DSparkMiniMaxDraftModel", "DSparkSpeculatorConfig"),
+        ("Gemma4MTPModel", "Gemma4MTPSpeculatorConfig"),
+        ("DeepSeekMTPModel", None),
+    ],
+)
+def test_draft_models_declare_speculator(architecture, expected):
+    model_info = ModelRegistry._try_inspect_model_cls(architecture)
+    assert model_info is not None
+
+    declared = model_info.speculator_config_class
+    assert (declared and declared.rsplit(".", 1)[-1]) == expected
diff --git a/tests/v1/spec_decode/test_dflash2.py b/tests/v1/spec_decode/test_dflash2.py
index 94de5f90a369..1d7e3b706ef5 100644
--- a/tests/v1/spec_decode/test_dflash2.py
+++ b/tests/v1/spec_decode/test_dflash2.py
@@ -1,11 +1,14 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-from types import SimpleNamespace
+from types import MethodType, SimpleNamespace
+from unittest import mock
 
 import pytest
 import torch
 
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
+from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
 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
@@ -606,6 +609,7 @@ def test_beam_dflash2_loads_every_checkpoint_tensor(monkeypatch):
     checkpoint = {name: torch.randn(shape) for name, shape in shapes.items()}
 
     loaded = model.load_weights(iter(checkpoint.items()))
+    model.process_weights_after_loading()  # builds the context K/V weights
 
     assert loaded == {name for name, _ in model.named_parameters()}
     context_kv = [
@@ -762,6 +766,13 @@ def forward(self, hidden_states):
     model = SimpleNamespace(
         hidden_norm=SimpleNamespace(weight=nn.Parameter(torch.ones(2))),
         _rms_norm_eps=1e-6,
+        _context_kv_proj=None,  # no MXFP8 context projection
+    )
+    model._normalize_context_states = MethodType(
+        qwen3_dflash.DFlashQwen3Model._normalize_context_states, model
+    )
+    model._project_context_kv_token_major = MethodType(
+        qwen3_dflash.DFlashQwen3Model._project_context_kv_token_major, model
     )
     layers_attn = [
         SimpleNamespace(
@@ -797,3 +808,57 @@ def forward(self, hidden_states):
     torch.testing.assert_close(actual_k, expected_k)
     torch.testing.assert_close(actual_v, expected_v)
     assert [projection.calls for projection in projections] == [1, 1]
+
+
+def test_score_edges_keeps_out_of_vocabulary_anchors_in_range():
+    """A request in chunked prefill anchors on its next prompt token, which can
+    lie past the draft vocabulary (a multimodal placeholder)."""
+    generator = torch.Generator().manual_seed(0)
+    vocab_size, rank, top_k = 8, 4, 3
+    predecessor_table = torch.randn(vocab_size, rank, generator=generator)
+    successor_table = torch.randn(vocab_size, rank, generator=generator)
+    candidate_ids = torch.randint(0, vocab_size, (2, 5, top_k), generator=generator)
+    unary_logits = torch.randn(2, 5, top_k, generator=generator)
+    hidden = torch.randn(2, 5, rank, generator=generator)
+
+    def score(anchor_token_ids: list[int]) -> torch.Tensor:
+        return _score_edges(
+            predecessor_table,
+            successor_table,
+            candidate_ids,
+            unary_logits,
+            hidden,
+            torch.tensor(anchor_token_ids),
+            top_k,
+        )
+
+    torch.testing.assert_close(score([vocab_size, 2]), score([vocab_size - 1, 2]))
+
+
+_EMBEDDING = "vllm.model_executor.layers.vocab_parallel_embedding"
+
+
+def test_top_k_tokens_stay_in_the_processor_vocabulary(default_vllm_config):
+    """A DFlash2 drafter takes candidates from the target's head, which can have
+    rows past the draft vocabulary (a padded multimodal vocabulary). The
+    candidates index the draft's codebooks, so top-k must not return those rows.
+    """
+    head_rows, vocab_size, hidden_size = 140, 128, 16
+    # One vocabulary shard, so the reductions need no process group.
+    with (
+        mock.patch(f"{_EMBEDDING}.get_tensor_model_parallel_rank", return_value=0),
+        mock.patch(
+            f"{_EMBEDDING}.get_tensor_model_parallel_world_size", return_value=1
+        ),
+    ):
+        head = ParallelLMHead(head_rows, hidden_size)
+    with torch.no_grad():
+        head.weight.zero_()
+        head.weight[vocab_size:head_rows] = 1.0
+    head.quant_method.process_weights_after_loading(head)
+    hidden_states = torch.ones(4, hidden_size, dtype=head.weight.dtype)
+    processor = LogitsProcessor(vocab_size)
+
+    _, ids = processor.get_top_k_tokens(head, hidden_states, 4)
+    assert (ids < vocab_size).all()
+    assert (processor.get_top_tokens(head, hidden_states) < vocab_size).all()
diff --git a/tests/v1/spec_decode/test_dflash_prepare_inputs.py b/tests/v1/spec_decode/test_dflash_prepare_inputs.py
index b40c5d2d906e..3290163aaf12 100644
--- a/tests/v1/spec_decode/test_dflash_prepare_inputs.py
+++ b/tests/v1/spec_decode/test_dflash_prepare_inputs.py
@@ -8,10 +8,12 @@
 import pytest
 import torch
 
+from vllm.config import VllmConfig
 from vllm.platforms import current_platform
 from vllm.v1.attention.backends.utils import PAD_SLOT_ID
 from vllm.v1.worker.gpu.spec_decode.dflash.speculator import (
     DFlashSpeculator,
+    _prepare_dflash_inputs_kernel,
 )
 
 pytestmark = pytest.mark.skipif(
@@ -28,6 +30,7 @@ def _run_prepare(
     cp_size: int = 1,
     cp_interleave: int = 1,
     draft_dcp_size: int | None = None,
+    num_scheduled_tokens: int = 4,
 ):
     device = torch.device(current_platform.device_type)
     max_num_reqs = 4
@@ -46,7 +49,7 @@ def _run_prepare(
     input_batch = SimpleNamespace(
         num_reqs=1,
         num_tokens=4,
-        num_scheduled_tokens=np.array([4], dtype=np.int32),
+        num_scheduled_tokens=np.array([num_scheduled_tokens], dtype=np.int32),
         seq_lens_cpu_upper_bound=torch.tensor([target_positions[-1] + 1]),
         positions=torch.tensor(target_positions, dtype=torch.int64, device=device),
         query_start_loc=torch.tensor([0, 4], dtype=torch.int32, device=device),
@@ -85,6 +88,7 @@ class InputsPrepared(Exception):
         pass
 
     draft = SimpleNamespace(
+        vllm_config=VllmConfig(),
         input_buffers=input_buffers,
         context_positions=context_positions,
         _context_slot_mappings=context_slot_mapping.unsqueeze(0),
@@ -220,3 +224,18 @@ def test_prepare_dflash_inputs_never_writes_the_null_block():
         PAD_SLOT_ID,
         PAD_SLOT_ID,
     ]
+
+
+def test_prepare_dflash_inputs_compiles_once_for_all_query_lengths():
+    def compiled_variants() -> set:
+        return {
+            key
+            for caches in _prepare_dflash_inputs_kernel.device_caches.values()
+            for key in caches[0]
+        }
+
+    kwargs = dict(target_positions=[10, 11, 12, 13], block_table_values=[0, 7, 8])
+    _run_prepare(**kwargs)
+    warmed = compiled_variants()
+    _run_prepare(**kwargs, num_scheduled_tokens=300)
+    assert compiled_variants() == warmed
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 26bcd68fea8f..33ec7f7b551c 100644
--- a/tests/v1/spec_decode/test_draft_attention_backend_override.py
+++ b/tests/v1/spec_decode/test_draft_attention_backend_override.py
@@ -12,13 +12,15 @@
 
 import pytest
 
-from vllm.config import LoadConfig, SpeculativeConfig
+from vllm.config import CacheConfig, LoadConfig, SpeculativeConfig
+from vllm.v1.worker.gpu.spec_decode.dflash.speculator import DFlashSpeculator
 from vllm.v1.worker.gpu.spec_decode.eagle.utils import load_eagle_model
 
 
 @dataclass
 class _AttentionConfig:
     backend: str | None = None
+    use_non_causal: bool = False
 
 
 @dataclass
@@ -98,3 +100,32 @@ def test_override_does_not_mutate_the_target_config():
     cfg = _config("FLASHINFER", "TRITON_ATTN")
     _capture_draft_config(cfg)
     assert cfg.attention_config.backend == "FLASHINFER"
+
+
+def test_draft_attn_builders_get_overrides_and_keep_resolved_cache_state():
+    """The draft's metadata builders read the same kernel overrides as its
+    layers, without losing state the engine core resolved on the target, e.g.
+    an fp8 draft next to an NVFP4 target with a resolved KV cache layout."""
+    cache_config = CacheConfig(cache_dtype="nvfp4")
+    cache_config.kv_cache_layout = "HND"
+    cfg = _VllmConfig(
+        attention_config=_AttentionConfig(backend="FLASHINFER"),
+        kernel_config=_KernelConfig(),
+        cache_config=cache_config,
+        speculative_config=_SpeculativeConfig(
+            attention_backend="TRITON_ATTN", kv_cache_dtype="fp8"
+        ),
+    )
+    speculator = object.__new__(DFlashSpeculator)
+    speculator.vllm_config = cfg
+    speculator.speculative_config = cfg.speculative_config
+    speculator.requires_non_causal = True
+
+    used = speculator.attn_vllm_config
+
+    assert used.attention_config.backend == "TRITON_ATTN"
+    assert used.attention_config.use_non_causal
+    assert used.cache_config.cache_dtype == "fp8"
+    assert used.cache_config.kv_cache_layout == "HND"
+    assert cfg.attention_config.backend == "FLASHINFER"
+    assert cfg.cache_config.cache_dtype == "nvfp4"
diff --git a/tests/v1/spec_decode/test_dspark_topk.py b/tests/v1/spec_decode/test_dspark_topk.py
index fe5f5a389bea..941b81236c20 100644
--- a/tests/v1/spec_decode/test_dspark_topk.py
+++ b/tests/v1/spec_decode/test_dspark_topk.py
@@ -16,6 +16,7 @@
 from vllm.model_executor.models.qwen3_dspark import DSparkMarkovHead
 from vllm.platforms import current_platform
 from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample
+from vllm.v1.worker.gpu.spec_decode.draft_logits import DraftLogitsCache
 from vllm.v1.worker.gpu.spec_decode.dspark.speculator import DSparkSpeculator
 
 
@@ -86,9 +87,9 @@ def test_dspark_revisited_position_uses_step_seed(use_fp64):
     spec.temperature = torch.ones(batch, device=device)
     spec.seeds = torch.arange(batch, dtype=torch.int64, device=device)
     spec._step_cols = torch.arange(2, dtype=torch.int32, device=device)
-    spec.draft_logits = torch.empty(batch, 2, 2, device=device)
-    spec.draft_logits_num_shards = 1
-    spec.draft_logits_shard_idx = 0
+    spec.draft_logits = DraftLogitsCache(
+        batch, 2, 2, dtype=torch.float32, fill=0.0, device=device
+    )
     logits = torch.tensor([0.7, 0.3], device=device).log().expand(batch, -1)
     mapping = torch.arange(batch, device=device)
     positions = torch.full((batch,), 22, dtype=torch.int64, device=device)
@@ -105,7 +106,7 @@ def test_dspark_revisited_position_uses_step_seed(use_fp64):
         positions - 1,
         apply_temperature=True,
         is_drafting=True,
-        logits_cache=spec.draft_logits,
+        logits_cache=spec.draft_logits.values,
         logits_cache_col=spec._step_cols[0],
         use_fp64=use_fp64,
     )
diff --git a/tests/v1/spec_decode/test_lilicorr.py b/tests/v1/spec_decode/test_lilicorr.py
index 7ded16e68e88..40d295f750fa 100644
--- a/tests/v1/spec_decode/test_lilicorr.py
+++ b/tests/v1/spec_decode/test_lilicorr.py
@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-from types import SimpleNamespace
+from types import MethodType, SimpleNamespace
 from typing import Any
 
 import pytest
@@ -24,6 +24,7 @@
 from vllm.transformers_utils.configs.eagle import EAGLEConfig
 from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample
 from vllm.v1.worker.gpu.spec_decode.dflash2.speculator import CandidateSampler
+from vllm.v1.worker.gpu.spec_decode.draft_logits import DraftLogitsCache
 from vllm.v1.worker.gpu.spec_decode.lilicorr.speculator import LiLiCorrSpeculator
 
 
@@ -228,8 +229,12 @@ def __call__(self, value, dim=-1):
         fake = SimpleNamespace(
             scale=1.0,
             soft_cap=None,
+            org_vocab_size=3 * tp_size,
             _apply_head=lambda *args, rank=rank: shards[rank].clone(),
         )
+        fake._mask_outside_vocab = MethodType(
+            logits_processor.LogitsProcessor._mask_outside_vocab, fake
+        )
         lm_head = SimpleNamespace(
             tp_size=tp_size,
             shard_indices=SimpleNamespace(
@@ -364,8 +369,6 @@ def init_base(self, config, device):
     spec.seeds = torch.arange(3)
     spec.draft_tokens = torch.empty(3, 2, dtype=torch.long)
     spec.draft_logits = None
-    spec.draft_logits_num_shards = 1
-    spec.draft_logits_shard_idx = 0
     spec.use_fp64_gumbel = False
     spec.enable_adaptive_verification = True
     spec.input_buffers = SimpleNamespace(input_ids=torch.arange(9))
@@ -452,7 +455,9 @@ def test_candidate_walk_preserves_conditional_scores_and_padding(probabilistic):
     scores = torch.randn(batch, steps, top_k, top_k, device=device)
     tokens = torch.full((batch, steps), -1, dtype=torch.long, device=device)
     cache = (
-        torch.full((batch, steps, vocab), -float("inf"), device=device)
+        DraftLogitsCache(
+            batch, steps, vocab, dtype=torch.float32, fill=-float("inf"), device=device
+        )
         if probabilistic
         else None
     )
@@ -494,13 +499,17 @@ def run():
                 assert index == realized.argmax().item()
             if cache is not None:
                 torch.testing.assert_close(
-                    cache[state, step, candidates[row, step]], realized
+                    cache.values[state, step, candidates[row, step]], realized
+                )
+                assert (
+                    cache.values[state, step, old_candidates[row, step]]
+                    .isneginf()
+                    .all()
                 )
-                assert cache[state, step, old_candidates[row, step]].isneginf().all()
             previous = index
     assert (tokens[2] == -1).all()
     if cache is not None:
-        assert cache[1].isneginf().all()
+        assert cache.values[1].isneginf().all()
 
 
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
@@ -518,7 +527,9 @@ def test_candidate_walk_revisited_position_matches_step_seed(use_fp64):
     seeds = mapping.to(torch.int64)
     positions = torch.arange(21, 21 + steps, device=device).repeat(batch)
     tokens = torch.empty(batch, steps, dtype=torch.int64, device=device)
-    cache = torch.empty(batch, steps, top_k, device=device)
+    cache = DraftLogitsCache(
+        batch, steps, top_k, dtype=torch.float32, fill=0.0, device=device
+    )
     first_round = None
     for start in (21, 22):
         positions.view(batch, steps).copy_(
@@ -545,7 +556,7 @@ def test_candidate_walk_revisited_position_matches_step_seed(use_fp64):
                 positions.view(batch, steps)[:, step] - 1,
                 apply_temperature=True,
                 is_drafting=True,
-                logits_cache=cache,
+                logits_cache=cache.values,
                 logits_cache_col=torch.tensor(step, device=device),
                 use_fp64=use_fp64,
             )
@@ -582,6 +593,7 @@ def test_checkpoint_coverage_rejects_incomplete_or_wrong_heads(
     nn.Module.__init__(wrapper)
     wrapper.model = LiLiCorr.__new__(LiLiCorr)
     nn.Module.__init__(wrapper.model)
+    wrapper.model._context_kv_proj = None  # set by __init__ (no MXFP8 projection)
     wrapper.has_own_lm_head = False
     with set_current_vllm_config(
         VllmConfig(
diff --git a/tests/v1/spec_decode/test_rejection_sampler_utils.py b/tests/v1/spec_decode/test_rejection_sampler_utils.py
index fd282f434624..95e61cdd5cb0 100644
--- a/tests/v1/spec_decode/test_rejection_sampler_utils.py
+++ b/tests/v1/spec_decode/test_rejection_sampler_utils.py
@@ -12,6 +12,7 @@
 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.draft_logits import DraftLogitsCache
 from vllm.v1.worker.gpu.spec_decode.rejection_sampler import (
     RejectionSampler,
     gather_draft_sampled,
@@ -28,6 +29,40 @@
     pytest.skip("CUDA required for rejection sampler tests", allow_module_level=True)
 
 
+def _draft_logits(
+    values: torch.Tensor | None, row_mapping: torch.Tensor | None = None
+) -> DraftLogitsCache | None:
+    """Wrap the draft logits tensor a test builds in a DraftLogitsCache."""
+    if values is None:
+        return None
+    draft_logits = DraftLogitsCache(
+        1, 1, 1, dtype=values.dtype, fill=0.0, device=values.device
+    )
+    draft_logits.values = values
+    draft_logits.row_mapping = row_mapping
+    return draft_logits
+
+
+def _rejection_sample(
+    *,
+    draft_logits: torch.Tensor | None,
+    draft_logits_row_mapping: torch.Tensor | None = None,
+    **kwargs,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    return rejection_sample(
+        draft_logits_cache=_draft_logits(draft_logits, draft_logits_row_mapping),
+        **kwargs,
+    )
+
+
+def _watermarked_rejection_sample(
+    *, draft_logits: torch.Tensor | None, **kwargs
+) -> tuple[torch.Tensor, torch.Tensor]:
+    return watermarked_rejection_sample(
+        draft_logits_cache=_draft_logits(draft_logits), **kwargs
+    )
+
+
 def _build_rejection_sample_inputs(
     target_logits_1d: torch.Tensor,
     draft_logits_1d: torch.Tensor,
@@ -138,7 +173,7 @@ def test_fly_requires_entropy_and_native_lookahead(
         [0, *torch.tensor(counts).cumsum(0).tolist()], dtype=torch.int32, device=device
     )
 
-    sampled, num_sampled = rejection_sample(
+    sampled, num_sampled = _rejection_sample(
         **inputs,
         num_speculative_steps=4,
         fly_window_size=2,
@@ -196,10 +231,10 @@ def test_fly_pending_window_preserves_recovery_position(
         if step < rejected_step:
             reference_logits[:, step] = native_logits[step]
     reference_inputs["target_logits"] = reference_logits.flatten(0, 1)
-    expected, expected_counts = rejection_sample(
+    expected, expected_counts = _rejection_sample(
         **reference_inputs, num_speculative_steps=num_steps
     )
-    sampled, counts = rejection_sample(
+    sampled, counts = _rejection_sample(
         **inputs, num_speculative_steps=num_steps, fly_window_size=2
     )
     assert counts.tolist() == [rejected_step + 1] * num_trials
@@ -302,7 +337,7 @@ def test_watermarked_bonus_uses_target_key():
         [[10, 11], [11, draft_token.item()]], dtype=torch.int64, device="cuda"
     )
 
-    sampled, num_sampled = watermarked_rejection_sample(
+    sampled, num_sampled = _watermarked_rejection_sample(
         **inputs,
         num_speculative_steps=1,
         contexts=contexts,
@@ -330,7 +365,7 @@ def test_watermarked_recovery_uses_target_key():
     inputs["draft_sampled"][1] = 0
     contexts = torch.tensor([[10, 11], [11, 0]], dtype=torch.int64, device="cuda")
 
-    sampled, num_sampled = watermarked_rejection_sample(
+    sampled, num_sampled = _watermarked_rejection_sample(
         **inputs,
         num_speculative_steps=1,
         contexts=contexts,
@@ -358,9 +393,11 @@ def test_repeated_draft_context_does_not_change_acceptance():
         temperature=1.0,
         num_trials=num_trials,
     )
-    expected, expected_num_sampled = rejection_sample(**inputs, num_speculative_steps=1)
+    expected, expected_num_sampled = _rejection_sample(
+        **inputs, num_speculative_steps=1
+    )
 
-    sampled, num_sampled = watermarked_rejection_sample(
+    sampled, num_sampled = _watermarked_rejection_sample(
         **inputs,
         num_speculative_steps=1,
         contexts=torch.tensor([[10, 11], [11, 12]], device="cuda").repeat(
@@ -403,14 +440,14 @@ def test_repeated_context_skip_mask_is_read_per_row_in_a_block():
     contexts = torch.tensor(
         [[10, 11], [11, 3], [3, 0], [0, 1]], dtype=torch.int64, device="cuda"
     )
-    expected, expected_num_sampled = rejection_sample(
+    expected, expected_num_sampled = _rejection_sample(
         **inputs, num_speculative_steps=num_speculative_steps
     )
     assert expected_num_sampled.item() == 2
     assert expected[0, 0].item() == 3
 
     watermarker = GumbelWatermarker(key=42, context_width=2)
-    skipped, skipped_num_sampled = watermarked_rejection_sample(
+    skipped, skipped_num_sampled = _watermarked_rejection_sample(
         **inputs,
         num_speculative_steps=num_speculative_steps,
         contexts=contexts,
@@ -422,7 +459,7 @@ def test_repeated_context_skip_mask_is_read_per_row_in_a_block():
     assert skipped_num_sampled.item() == 2
     torch.testing.assert_close(skipped[:, :2], expected[:, :2], rtol=0, atol=0)
 
-    watermarked, watermarked_num_sampled = watermarked_rejection_sample(
+    watermarked, watermarked_num_sampled = _watermarked_rejection_sample(
         **inputs,
         num_speculative_steps=num_speculative_steps,
         contexts=contexts,
@@ -454,9 +491,11 @@ def test_repeated_recovery_context_uses_ordinary_residual_sample():
         num_trials=1,
     )
     inputs["draft_sampled"][1] = 0
-    expected, expected_num_sampled = rejection_sample(**inputs, num_speculative_steps=1)
+    expected, expected_num_sampled = _rejection_sample(
+        **inputs, num_speculative_steps=1
+    )
 
-    sampled, num_sampled = watermarked_rejection_sample(
+    sampled, num_sampled = _watermarked_rejection_sample(
         **inputs,
         num_speculative_steps=1,
         contexts=torch.tensor([[10, 11], [11, 0]], device="cuda"),
@@ -478,9 +517,11 @@ def test_repeated_bonus_context_uses_ordinary_target_sample():
         temperature=1.0,
         num_trials=1,
     )
-    expected, expected_num_sampled = rejection_sample(**inputs, num_speculative_steps=1)
+    expected, expected_num_sampled = _rejection_sample(
+        **inputs, num_speculative_steps=1
+    )
 
-    sampled, num_sampled = watermarked_rejection_sample(
+    sampled, num_sampled = _watermarked_rejection_sample(
         **inputs,
         num_speculative_steps=1,
         contexts=torch.tensor([[10, 11], [11, 12]], device="cuda"),
@@ -503,9 +544,11 @@ def test_repeated_context_does_not_change_unwatermarked_or_greedy_rows():
         num_trials=2,
     )
     inputs["temperature"][1] = 0
-    expected, expected_num_sampled = rejection_sample(**inputs, num_speculative_steps=1)
+    expected, expected_num_sampled = _rejection_sample(
+        **inputs, num_speculative_steps=1
+    )
 
-    sampled, num_sampled = watermarked_rejection_sample(
+    sampled, num_sampled = _watermarked_rejection_sample(
         **inputs,
         num_speculative_steps=1,
         contexts=torch.tensor([[10, 11], [11, 12], [20, 21], [21, 22]], device="cuda"),
@@ -536,7 +579,7 @@ def test_watermarked_recovery_supports_smaller_draft_vocabulary():
     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(
+    sampled, num_sampled = _watermarked_rejection_sample(
         **inputs,
         num_speculative_steps=1,
         contexts=contexts,
@@ -592,7 +635,7 @@ def test_stochastic_rejection_sample(
         num_trials=num_trials,
     )
 
-    sampled, num_sampled = rejection_sample(
+    sampled, num_sampled = _rejection_sample(
         **inputs, num_speculative_steps=num_speculative_steps
     )
 
@@ -684,7 +727,7 @@ def test_gumbel_drafted_rejection_sample_is_unbiased(num_speculative_steps: int)
         inputs, draft_logits_1d, NARROW_NUM_TRIALS, num_speculative_steps
     )
 
-    sampled, num_sampled = rejection_sample(
+    sampled, num_sampled = _rejection_sample(
         **inputs, num_speculative_steps=num_speculative_steps
     )
 
@@ -744,7 +787,7 @@ def test_multiround_rejection_sample_is_unbiased(
                     logits_cache=logits_cache,
                     logits_cache_col=torch.tensor(step, device=device),
                 )
-        sampled, counts = rejection_sample(
+        sampled, counts = _rejection_sample(
             target_logits=target_probs.log().expand(n * (k + 1), -1).contiguous(),
             draft_logits=logits_cache,
             draft_sampled=draft_sampled.flatten(),
@@ -792,7 +835,7 @@ def test_greedy_rejection_sample(num_speculative_steps: int):
         num_trials=num_trials,
     )
 
-    sampled, num_sampled = rejection_sample(
+    sampled, num_sampled = _rejection_sample(
         **inputs, num_speculative_steps=num_speculative_steps
     )
 
@@ -855,7 +898,7 @@ def test_synthetic_rejection_sample(
         conditional_rates, dtype=torch.float32, device=device
     )
 
-    _, num_sampled = rejection_sample(
+    _, num_sampled = _rejection_sample(
         **inputs,
         num_speculative_steps=num_speculative_steps,
         synthetic_conditional_rates=synthetic_conditional_rates,
@@ -903,7 +946,7 @@ def test_all_nan_target_logits_in_range(temperature: float):
         num_trials=num_trials,
     )
 
-    sampled, num_sampled = rejection_sample(**inputs, num_speculative_steps=K)
+    sampled, num_sampled = _rejection_sample(**inputs, num_speculative_steps=K)
 
     assert ((num_sampled >= 1) & (num_sampled <= K + 1)).all()
     steps = torch.arange(K + 1, device=device).unsqueeze(0)
@@ -934,7 +977,7 @@ def test_placeholder_draft_token_rejected():
     )
     inputs["draft_sampled"].view(num_trials, K + 1)[:, 1:] = -1
 
-    sampled, num_sampled = rejection_sample(**inputs, num_speculative_steps=K)
+    sampled, num_sampled = _rejection_sample(**inputs, num_speculative_steps=K)
 
     assert torch.equal(num_sampled, torch.ones_like(num_sampled))
     recovered = sampled[:, 0]
@@ -971,7 +1014,7 @@ def test_block_verification_placeholder_truncates_block(
         inputs["draft_logits"] = None
     inputs["draft_sampled"].view(num_trials, K + 1)[:, K + 1 - num_placeholders :] = -1
 
-    sampled, num_sampled = rejection_sample(
+    sampled, num_sampled = _rejection_sample(
         **inputs, num_speculative_steps=K, use_block_verification=True
     )
 
@@ -1013,7 +1056,7 @@ def test_greedy_placeholder_emits_target_argmax():
     draft_sampled[:, 1] = target_argmax
     draft_sampled[:, 2:] = -1
 
-    sampled, num_sampled = rejection_sample(**inputs, num_speculative_steps=K)
+    sampled, num_sampled = _rejection_sample(**inputs, num_speculative_steps=K)
 
     assert torch.equal(num_sampled, torch.full_like(num_sampled, 2))
     steps = torch.arange(K + 1, device=device).unsqueeze(0)
@@ -1049,7 +1092,7 @@ def test_placeholder_blocks_later_draft_tokens(use_block_verification: bool):
     )
     inputs["draft_sampled"].view(num_trials, K + 1)[:, 2] = -1
 
-    _, num_sampled = rejection_sample(
+    _, num_sampled = _rejection_sample(
         **inputs,
         num_speculative_steps=K,
         use_block_verification=use_block_verification,
@@ -1098,6 +1141,7 @@ def test_verify_rejects_unproposed_drafts():
             seeds=SimpleNamespace(gpu=inputs["seed"]),
         ),
         use_fp64_gumbel=False,
+        trace_replay_state=None,
     )
     rejection_sampler.num_speculative_steps = K
     rejection_sampler.synthetic_conditional_rates = None
@@ -1105,8 +1149,14 @@ def test_verify_rejects_unproposed_drafts():
     rejection_sampler.fly_window_size = 0
     rejection_sampler.fly_entropy_threshold = 0.3
     rejection_sampler.fly_entropy_top_k = 3
-    rejection_sampler.draft_logits_num_shards = 1
     rejection_sampler.watermark_key = None
+    draft_logits = DraftLogitsCache(
+        *inputs["draft_logits"].shape,
+        dtype=inputs["draft_logits"].dtype,
+        fill=0.0,
+        device=device,
+    )
+    draft_logits.values.copy_(inputs["draft_logits"])
 
     def verify() -> tuple[torch.Tensor, torch.Tensor]:
         draft_sampled, pos = gather_draft_sampled(
@@ -1119,7 +1169,7 @@ def verify() -> tuple[torch.Tensor, torch.Tensor]:
         )
         _, sampled, num_sampled = rejection_sampler._verify(
             inputs["target_logits"],
-            inputs["draft_logits"],
+            draft_logits,
             draft_sampled,
             pos,
             inputs["cu_num_logits"],
@@ -1128,7 +1178,6 @@ def verify() -> tuple[torch.Tensor, torch.Tensor]:
             inputs["expanded_idx_mapping"],
             inputs["expanded_local_pos"],
             None,
-            draft_sampled,
         )
         return sampled, num_sampled
 
@@ -1183,7 +1232,7 @@ def test_block_verification_rejection_sample(
     if not has_draft_logits:
         inputs["draft_logits"] = None
 
-    sampled, num_sampled = rejection_sample(
+    sampled, num_sampled = _rejection_sample(
         **inputs,
         num_speculative_steps=num_speculative_steps,
         use_block_verification=True,
@@ -1223,10 +1272,10 @@ def test_block_verification_accepts_at_least_as_many(num_speculative_steps: int)
         num_trials=num_trials,
     )
 
-    _, num_sampled_standard = rejection_sample(
+    _, num_sampled_standard = _rejection_sample(
         **inputs, num_speculative_steps=num_speculative_steps
     )
-    _, num_sampled_block = rejection_sample(
+    _, num_sampled_block = _rejection_sample(
         **inputs,
         num_speculative_steps=num_speculative_steps,
         use_block_verification=True,
@@ -1272,10 +1321,10 @@ def test_chunked_requests_match_full_batch(
         inputs["draft_logits"] = None
 
     if fly_window_size:
-        native, native_counts = rejection_sample(
+        native, native_counts = _rejection_sample(
             **inputs, num_speculative_steps=num_speculative_steps
         )
-        gated, gated_counts = rejection_sample(
+        gated, gated_counts = _rejection_sample(
             **inputs,
             num_speculative_steps=num_speculative_steps,
             fly_window_size=fly_window_size,
@@ -1288,7 +1337,7 @@ def test_chunked_requests_match_full_batch(
         )
         assert torch.equal(gated[valid], native[valid])
 
-    sampled, num_sampled = rejection_sample(
+    sampled, num_sampled = _rejection_sample(
         **inputs,
         num_speculative_steps=num_speculative_steps,
         fly_window_size=fly_window_size,
@@ -1312,7 +1361,7 @@ def test_chunked_requests_match_full_batch(
         chunk_inputs["cu_num_logits"] = inputs["cu_num_logits"][start : end + 1] - lo
         chunk_inputs["idx_mapping"] = inputs["idx_mapping"][start:end]
 
-        chunk_sampled, chunk_num_sampled = rejection_sample(
+        chunk_sampled, chunk_num_sampled = _rejection_sample(
             **chunk_inputs,
             num_speculative_steps=num_speculative_steps,
             fly_window_size=fly_window_size,
@@ -1380,7 +1429,7 @@ def test_adaptive_draft_lengths_match_and_bonus_oracle(
         num_speculative_steps=K,
         use_block_verification=use_block_verification,
     )
-    sampled, num_sampled = rejection_sample(**inputs)
+    sampled, num_sampled = _rejection_sample(**inputs)
 
     for req_idx, num_matches in greedy_num_matches.items():
         lo = int(cu_num_logits[req_idx])
@@ -1404,7 +1453,7 @@ def test_adaptive_draft_lengths_match_and_bonus_oracle(
         req_inputs["cu_num_logits"] = cu_num_logits[req_idx : req_idx + 2] - lo
         req_inputs["idx_mapping"] = idx_mapping[req_idx : req_idx + 1]
 
-        req_sampled, req_num_sampled = rejection_sample(**req_inputs)
+        req_sampled, req_num_sampled = _rejection_sample(**req_inputs)
 
         n = int(req_num_sampled[0])
         assert n <= num_drafts[req_idx] + 1
@@ -1506,7 +1555,7 @@ def test_watermarked_recovery_matches_philox_gumbel_sample(
     )
     watermarker = GumbelWatermarker(key=key, context_width=context_width)
 
-    sampled, num_sampled = watermarked_rejection_sample(
+    sampled, num_sampled = _watermarked_rejection_sample(
         **inputs,
         num_speculative_steps=2,
         contexts=contexts,
@@ -1544,10 +1593,10 @@ def test_watermarked_recovery_leaves_disabled_rows_on_the_stock_draw():
     watermarking[::2] = True
     inputs["temperature"][:8] = 0.0
 
-    stock_sampled, stock_num_sampled = rejection_sample(
+    stock_sampled, stock_num_sampled = _rejection_sample(
         **inputs, num_speculative_steps=3
     )
-    sampled, num_sampled = watermarked_rejection_sample(
+    sampled, num_sampled = _watermarked_rejection_sample(
         **inputs,
         num_speculative_steps=3,
         contexts=contexts,
@@ -1589,8 +1638,8 @@ def test_watermarked_recovery_is_unchanged_by_fp64_gumbel():
         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(
+    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")
@@ -1601,3 +1650,140 @@ def test_watermarked_recovery_is_unchanged_by_fp64_gumbel():
         rtol=0,
         atol=0,
     )
+
+
+def _position_owned_rows(
+    slots: torch.Tensor, num_shards: int, shard_idx: int, max_num_reqs: int
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """The batch positions a shard owns and its DraftLogitsCache.row_mapping table:
+    the request at batch position i is cached by shard i % num_shards at row
+    i // num_shards."""
+    owned = torch.arange(shard_idx, slots.shape[0], num_shards, device=slots.device)
+    rows = torch.full((max_num_reqs,), -1, dtype=torch.int32, device=slots.device)
+    rows[slots[owned].long()] = torch.arange(
+        owned.shape[0], dtype=torch.int32, device=slots.device
+    )
+    return owned, rows
+
+
+@pytest.mark.parametrize("num_shards", [2, 4])
+@pytest.mark.parametrize("use_block_verification", [False, True])
+def test_sharded_draft_logits_verify_like_replicated(
+    num_shards: int, use_block_verification: bool
+):
+    """A TP rank that caches only its owned requests' draft logits, at the
+    rows of draft_logits_row_mapping, verifies those requests exactly like the
+    replicated cache does."""
+    torch.manual_seed(0)
+    device = "cuda"
+    num_reqs = 12
+    num_speculative_steps = 3
+    vocab_size = 1000
+    num_rows = num_speculative_steps + 1
+    inputs = _build_rejection_sample_inputs(
+        torch.zeros(vocab_size, device=device),
+        torch.zeros(vocab_size, device=device),
+        num_speculative_steps,
+        temperature=1.0,
+        num_trials=num_reqs,
+    )
+    # Per-request distributions, close enough that some drafts are accepted.
+    draft_logits = torch.randn(
+        num_reqs, num_speculative_steps, vocab_size, device=device
+    )
+    target_logits = torch.randn(num_reqs, num_rows, vocab_size, device=device)
+    target_logits[:, :num_speculative_steps] = draft_logits + 0.5 * torch.randn_like(
+        draft_logits
+    )
+    inputs["target_logits"] = target_logits.view(-1, vocab_size)
+    draft_tokens = torch.multinomial(
+        torch.softmax(draft_logits.view(-1, vocab_size), dim=-1), 1
+    ).view(num_reqs, num_speculative_steps)
+    inputs["draft_sampled"].view(num_reqs, num_rows)[:, 1:] = draft_tokens
+    # Shuffle request slots so cache rows aren't request slots.
+    slots = torch.randperm(num_reqs, device=device).to(torch.int32)
+    inputs["idx_mapping"] = slots
+    inputs["expanded_idx_mapping"] = slots.repeat_interleave(num_rows)
+    inputs["draft_logits"] = torch.empty_like(draft_logits)
+    inputs["draft_logits"][slots.long()] = draft_logits
+
+    sampled, num_sampled = _rejection_sample(
+        **inputs,
+        num_speculative_steps=num_speculative_steps,
+        use_block_verification=use_block_verification,
+    )
+    assert 1 < num_sampled.float().mean() < num_rows
+
+    for shard_idx in range(num_shards):
+        owned, draft_logits_row_mapping = _position_owned_rows(
+            slots, num_shards, shard_idx, num_reqs
+        )
+        rows = owned.view(-1, 1) * num_rows + torch.arange(num_rows, device=device)
+        rows = rows.view(-1)
+        local_inputs = dict(inputs)
+        for name in (
+            "target_logits",
+            "draft_sampled",
+            "pos",
+            "expanded_idx_mapping",
+            "expanded_local_pos",
+        ):
+            local_inputs[name] = inputs[name][rows]
+        local_inputs["idx_mapping"] = slots[owned]
+        local_inputs["cu_num_logits"] = inputs["cu_num_logits"][: len(owned) + 1]
+        local_inputs["draft_logits"] = draft_logits[owned]
+
+        local_sampled, local_num_sampled = _rejection_sample(
+            **local_inputs,
+            num_speculative_steps=num_speculative_steps,
+            use_block_verification=use_block_verification,
+            draft_logits_row_mapping=draft_logits_row_mapping,
+        )
+        assert torch.equal(local_num_sampled, num_sampled[owned])
+        valid = torch.arange(num_rows, device=device) < local_num_sampled[:, None]
+        assert torch.equal(local_sampled[valid], sampled[owned][valid])
+
+
+@pytest.mark.parametrize("num_shards", [2, 4])
+def test_sharded_draft_logits_cache_writes(num_shards: int):
+    """gumbel_sample caches only the owned rows, at the rows of
+    logits_cache_row_mapping, without changing what it samples."""
+    torch.manual_seed(0)
+    device = "cuda"
+    max_num_reqs = 16
+    num_cols = 3
+    num_tokens = max_num_reqs * num_cols
+    slots = torch.randperm(max_num_reqs, device=device).to(torch.int32)
+    idx_mapping = slots.repeat_interleave(num_cols)
+    # A CUDA graph padding row.
+    idx_mapping[-num_cols:] = -1
+    cols = torch.arange(num_cols, device=device, dtype=torch.int32).repeat(max_num_reqs)
+    logits = torch.randn(num_tokens, VOCAB_SIZE, device=device)
+    kwargs = dict(
+        expanded_idx_mapping=idx_mapping,
+        temperature=torch.ones(max_num_reqs, device=device),
+        seed=torch.arange(max_num_reqs, device=device),
+        pos=torch.arange(num_tokens, device=device),
+        apply_temperature=True,
+        is_drafting=True,
+        logits_cache_col=cols,
+    )
+    cache = torch.zeros(max_num_reqs, num_cols, VOCAB_SIZE, device=device)
+    sampled = gumbel_sample(logits, logits_cache=cache, **kwargs)
+
+    for shard_idx in range(num_shards):
+        owned, logits_cache_row_mapping = _position_owned_rows(
+            slots, num_shards, shard_idx, max_num_reqs
+        )
+        local_cache = torch.zeros(len(owned), num_cols, VOCAB_SIZE, device=device)
+        local_sampled = gumbel_sample(
+            logits,
+            logits_cache=local_cache,
+            logits_cache_row_mapping=logits_cache_row_mapping,
+            **kwargs,
+        )
+        assert torch.equal(local_sampled, sampled)
+        # The padding request caches nothing.
+        is_real = owned < max_num_reqs - 1
+        expected = cache[slots[owned].long()] * is_real.view(-1, 1, 1)
+        assert torch.equal(local_cache, expected)
diff --git a/tests/v1/spec_decode/utils.py b/tests/v1/spec_decode/utils.py
new file mode 100644
index 000000000000..daaa8ac2a240
--- /dev/null
+++ b/tests/v1/spec_decode/utils.py
@@ -0,0 +1,64 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from types import SimpleNamespace
+from typing import Any
+
+from vllm.v1.worker.gpu.spec_decode.speculator_config import (
+    SpeculatorConfig,
+    get_speculator_config_class,
+)
+
+
+class _FakeSpeculativeConfig(SimpleNamespace):
+    def __getattr__(self, name: str) -> Any:
+        # Like SpeculativeConfig's shortcuts to speculator config properties.
+        if name != "speculator_config" and "speculator_config" in vars(self):
+            return getattr(self.speculator_config, name)
+        raise AttributeError(name)
+
+
+def fake_speculative_config(
+    speculator_config_cls: type[SpeculatorConfig] | None = None, **fields: Any
+) -> Any:
+    """A `SimpleNamespace` standing in for a `SpeculativeConfig` with `fields`.
+
+    Its `speculator_config` is of `speculator_config_cls`, or by default of the
+    class a real config with the same `method` would use (ignoring the draft
+    architecture). It reads `fields` lazily, as for a real config. Attributes
+    that aren't fields are taken from the speculator config, like
+    SpeculativeConfig's shortcuts to its properties.
+    """
+    spec = _FakeSpeculativeConfig(**{"custom_method": None, **fields})
+    if speculator_config_cls is None:
+        probe = SimpleNamespace(
+            method=fields["method"],
+            custom_method=None,
+            draft_model_config=SimpleNamespace(architectures=[]),
+            use_gemma4_mtp=lambda: False,
+            use_multi_module_mtp=lambda: False,
+        )
+        speculator_config_cls = get_speculator_config_class(probe)  # type: ignore[arg-type]
+    spec.speculator_config = speculator_config_cls(spec)
+    return spec
+
+
+# SpeculativeConfig's shortcuts to speculator config properties.
+_SPECULATOR_CONFIG_SHORTCUTS = (
+    "uses_target_hidden_states",
+    "uses_aux_hidden_states",
+    "uses_separate_workspace_lane",
+)
+
+
+def attach_speculator_config(
+    spec: Any, speculator_config_cls: type[SpeculatorConfig]
+) -> Any:
+    """Give a mock standing in for a `SpeculativeConfig` a real
+    `speculator_config` of `speculator_config_cls` (and the shortcuts to its
+    properties), which reads the mock's attributes lazily."""
+    spec.speculator_config = speculator_config_cls(spec)
+    spec.method = speculator_config_cls.base_method
+    spec.custom_method = None
+    for name in _SPECULATOR_CONFIG_SHORTCUTS:
+        setattr(spec, name, getattr(spec.speculator_config, name))
+    return spec
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 fc316e336411..4714062b59fc 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
@@ -36,6 +36,7 @@ def mock_model_runner_with_req_states():
     runner.model_state = Mock()
     runner.block_tables = Mock()
     runner.lora_state = Mock()
+    runner.draft_tokens_handler = Mock()
     runner.pp_handler = None
     runner.sampler = None
     runner.prompt_logprobs_worker = None
diff --git a/tests/v1/structured_output/test_backend_xgrammar_lark.py b/tests/v1/structured_output/test_backend_xgrammar_lark.py
index cda9acf311ce..cdce3d9cd417 100644
--- a/tests/v1/structured_output/test_backend_xgrammar_lark.py
+++ b/tests/v1/structured_output/test_backend_xgrammar_lark.py
@@ -95,6 +95,7 @@ def test_auto_backend_selects_guidance_for_mistral_lark(mode: ToolChoiceEnum):
         json_schema=None,
         parallel_tool_calls=True,
         json_only=False,
+        force_tool_names=False,
     )
     assert _validate(grammar, "auto", tokenizer)._backend == "guidance"
 
diff --git a/tests/v1/structured_output/test_grammar_backends.py b/tests/v1/structured_output/test_grammar_backends.py
new file mode 100644
index 000000000000..6ed0756b31ae
--- /dev/null
+++ b/tests/v1/structured_output/test_grammar_backends.py
@@ -0,0 +1,73 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from collections.abc import Callable
+
+import pytest
+from transformers import AutoTokenizer
+
+from vllm.config import DeviceConfig, VllmConfig
+from vllm.v1.structured_output.backend_guidance import GuidanceBackend
+from vllm.v1.structured_output.backend_lm_format_enforcer import (
+    LMFormatEnforcerBackend,
+)
+from vllm.v1.structured_output.backend_outlines import OutlinesBackend
+from vllm.v1.structured_output.backend_types import (
+    StructuredOutputBackend,
+    StructuredOutputOptions,
+)
+from vllm.v1.structured_output.backend_xgrammar import XgrammarBackend
+
+TOKENIZER = "gpt2"
+REGEX = "1a"
+BACKENDS = (
+    pytest.param("xgrammar", id="xgrammar"),
+    pytest.param("guidance", id="guidance"),
+    pytest.param("outlines", id="outlines"),
+    pytest.param(
+        "lm-format-enforcer",
+        marks=pytest.mark.xfail(
+            reason=(
+                "lm-format-enforcer 0.11.3 imports the removed "
+                "vllm.transformers_utils.tokenizer path"
+            ),
+            strict=True,
+        ),
+        id="lm-format-enforcer",
+    ),
+)
+BACKEND_CLASSES: dict[str, Callable[..., StructuredOutputBackend]] = {
+    "xgrammar": XgrammarBackend,
+    "guidance": GuidanceBackend,
+    "outlines": OutlinesBackend,
+    "lm-format-enforcer": LMFormatEnforcerBackend,
+}
+
+
+@pytest.fixture(scope="module")
+def tokenizer():
+    return AutoTokenizer.from_pretrained(TOKENIZER)
+
+
+def _single_token(tokenizer, text: str) -> int:
+    token_ids = tokenizer.encode(text)
+    assert len(token_ids) == 1, (text, token_ids)
+    return token_ids[0]
+
+
+@pytest.mark.parametrize("backend", BACKENDS)
+def test_accept_tokens_rejection_is_atomic(tokenizer, backend: str):
+    structured_output_backend = BACKEND_CLASSES[backend](
+        VllmConfig(device_config=DeviceConfig(device="cpu")),
+        tokenizer,
+        len(tokenizer),
+    )
+    grammar = structured_output_backend.compile_grammar(
+        StructuredOutputOptions.REGEX, REGEX
+    )
+
+    first = _single_token(tokenizer, "1")
+    invalid = _single_token(tokenizer, "z")
+
+    assert not grammar.accept_tokens("request", [first, invalid])
+    assert grammar.accept_tokens("request", [first])
diff --git a/tests/v1/structured_output/test_structured_output_manager.py b/tests/v1/structured_output/test_structured_output_manager.py
index 8d3b98be7a5b..c68f58f9e705 100644
--- a/tests/v1/structured_output/test_structured_output_manager.py
+++ b/tests/v1/structured_output/test_structured_output_manager.py
@@ -47,9 +47,9 @@ def tokenizer():
 
 @dataclass(frozen=True)
 class FlowCase:
-    """One validate -> bitmask -> accept -> post-accept bitmask scenario."""
+    """One bitmask -> accept -> post-accept bitmask scenario."""
 
-    raw_drafts: tuple[str, ...]
+    raw_drafts: tuple[str | int, ...]
     expected_row_pattern: str
     expected_reasoning: bool | None
     expect_terminated: bool
@@ -102,10 +102,14 @@ def _single_token(tokenizer, text: str) -> int:
     return token_ids[0]
 
 
-def _to_token_ids(tokenizer, texts: tuple[str, ...]) -> list[int]:
-    """Convert literal text (or EOS sentinel) to token IDs."""
+def _to_token_ids(tokenizer, texts: tuple[str | int, ...]) -> list[int]:
+    """Convert literal text, EOS, or speculative-padding tokens to IDs."""
     return [
-        tokenizer.eos_token_id if t == EOS else _single_token(tokenizer, t)
+        t
+        if isinstance(t, int)
+        else tokenizer.eos_token_id
+        if t == EOS
+        else _single_token(tokenizer, t)
         for t in texts
     ]
 
@@ -182,6 +186,40 @@ def _build_harness(
     return manager, request
 
 
+@pytest.mark.parametrize("backend", BACKENDS)
+@pytest.mark.parametrize(
+    ("raw_drafts", "expected_length"),
+    [
+        ((THINK_END, "1", "2", -1), 3),
+        ((THINK_END, "1", "z", "2"), 3),
+        (("1", "z", "2", -1), 3),
+        (("z", THINK_END, "1", "z", "2", -1), 4),
+    ],
+)
+def test_trim_spec_tokens_with_grammar(
+    tokenizer,
+    backend: StructuredOutputsBackend,
+    raw_drafts: tuple[str | int, ...],
+    expected_length: int,
+):
+    manager, request = _build_harness(
+        tokenizer,
+        backend,
+        reasoning_ended=False,
+        reasoning_parser_kwargs={
+            "marker": _single_token(tokenizer, THINK_END),
+        },
+        structured_outputs=StructuredOutputsParams(regex=DIGITS_REGEX),
+    )
+
+    spec_tokens = _to_token_ids(tokenizer, raw_drafts)
+
+    # The reasoning marker makes the manager find the grammar boundary from
+    # the real request state instead of hard-coding a boundary in the test.
+    trimmed = manager.trim_spec_tokens(request, spec_tokens)
+    assert trimmed == spec_tokens[:expected_length]
+
+
 def _row_pattern(bitmask, masked_ids=()) -> str:
     """U: unconstrained; R: all but `masked_ids` allowed; C: grammar rows."""
     reasoning_row = np.full(bitmask.shape[1], -1, dtype=np.int32)
@@ -210,23 +248,16 @@ def _run_real_flow(
 ) -> 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},
+        scheduled_spec_decode_tokens={request.request_id: raw_drafts},
     )
     assert bitmask is not None
     assert bitmask.shape[0] == NUM_SPEC_TOKENS + 1
     row_pattern = _row_pattern(bitmask[: len(raw_drafts) + 1], masked_ids or ())
     assert row_pattern == 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
@@ -238,8 +269,7 @@ def _run_real_flow(
     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.
+    # StructuredOutputManager.accept_tokens() advances the grammar.
     request.append_output_token_ids(list(expected_validated))
     assert manager.accept_tokens(request, list(expected_validated))
 
@@ -268,8 +298,9 @@ def _run_real_flow(
     pytest.param(
         FlowCase(
             prefix='{"a"',
-            raw_drafts=(":", ' "', "b", '"}'),
-            expected_row_pattern="CCCCC",
+            raw_drafts=(":", ' "', "b", '"}', -1, -1),
+            expected_validated=(":", ' "', "b", '"}'),
+            expected_row_pattern="CCCCCUU",
             expected_reasoning=True,
             expect_terminated=False,
             reasoning_ended=True,
@@ -490,7 +521,6 @@ def test_reasoning_rows_mask_grammar_invalid_implicit_end(
     )
     manager.reasoner_cls = MockEngineReasoner
     drafts = _to_token_ids(tokenizer, raw_drafts)
-    assert manager.validate_tokens(request, drafts) == drafts
 
     bitmask = manager.grammar_bitmask(
         requests={request.request_id: request},
@@ -562,19 +592,17 @@ def test_initial_constraint_activation(
             if reasoning_ended is not None
             else None
         },
+        num_spec_tokens=1,
     )
 
-    # "{" 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={},
+        scheduled_spec_decode_tokens={request.request_id: [open_brace]},
     )
     assert bitmask is not None
+    assert _row_pattern(bitmask) == "CC"
     request.append_output_token_ids([open_brace])
     assert manager.accept_tokens(request, [open_brace])
 
@@ -625,30 +653,6 @@ def test_regex_flow(
     )
 
 
-@pytest.mark.parametrize("backend", REGEX_BACKENDS)
-def test_rejected_draft_keeps_later_rows_constrained(
-    tokenizer, backend: StructuredOutputsBackend
-):
-    """A draft the grammar rejects during bitmask fill must not unconstrain
-    the remaining rows or the bonus row."""
-    manager, request = _build_harness(
-        tokenizer,
-        backend,
-        use_reasoner=False,
-        structured_outputs=StructuredOutputsParams(regex=DIGITS_REGEX),
-    )
-    drafts = _to_token_ids(tokenizer, ("1", "z", "2"))
-    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 _row_pattern(bitmask[: len(drafts) + 1]) == "CCCC"
-    assert (bitmask[2] == bitmask[1]).all()
-    assert (bitmask[3] == bitmask[1]).all()
-
-
 def test_outlines_termination(tokenizer):
     """outlines_core never advances on EOS; termination must still be
     tracked exactly once per accepted EOS and undone by rollback."""
diff --git a/tests/v1/worker/cpu/test_cpu_attn_utils.py b/tests/v1/worker/cpu/test_cpu_attn_utils.py
new file mode 100644
index 000000000000..69b254bfae9d
--- /dev/null
+++ b/tests/v1/worker/cpu/test_cpu_attn_utils.py
@@ -0,0 +1,76 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import numpy as np
+import pytest
+import torch
+
+from vllm.config import VllmConfig
+from vllm.platforms import CpuArchEnum, current_platform
+from vllm.triton_utils import HAS_TRITON
+from vllm.v1.worker.cpu import attn_utils
+from vllm.v1.worker.cpu.attn_utils import FastPrefillHelper
+
+pytestmark = pytest.mark.skipif(
+    not current_platform.is_cpu()
+    or HAS_TRITON
+    or current_platform.get_cpu_architecture() == CpuArchEnum.X86,
+    reason="requires a non-x86 CPU platform without Triton",
+)
+
+
+@pytest.mark.parametrize(
+    ("layer_names", "shared_layers", "enabled", "expected"),
+    [
+        (["t0", "t1"], {}, True, set()),
+        (
+            ["t0", "t1", "t2", "t3"],
+            {"t2": "t1", "t3": "t1"},
+            True,
+            {"t2", "t3"},
+        ),
+        (
+            ["t0", "t1", "t2", "t3"],
+            {"t1": "t0", "t3": "t0"},
+            True,
+            {"t3"},
+        ),
+        (["t0"], {"t0": "t0"}, False, set()),
+    ],
+)
+def test_fast_prefill_selects_only_trailing_shared_layers(
+    monkeypatch: pytest.MonkeyPatch,
+    layer_names: list[str],
+    shared_layers: dict[str, str],
+    enabled: bool,
+    expected: set[str],
+) -> None:
+    monkeypatch.setattr(
+        attn_utils,
+        "get_layers_from_vllm_config",
+        lambda *args, **kwargs: dict.fromkeys(layer_names),
+    )
+    monkeypatch.setattr(
+        attn_utils,
+        "get_shared_kv_cache_layers",
+        lambda config: shared_layers,
+    )
+    config = VllmConfig()
+    config.cache_config.kv_sharing_fast_prefill = enabled
+
+    assert attn_utils.get_kv_sharing_fast_prefill_eligible_layers(config) == expected
+
+
+def test_fast_prefill_uses_unpadded_cpu_logits_indices() -> None:
+    helper = FastPrefillHelper()
+    logits_indices = torch.tensor([3, 9], dtype=torch.int32)
+    cu_num_logits = np.array([0, 1, 2], dtype=np.int32)
+
+    assert helper.prepare(logits_indices, cu_num_logits, has_prefill=False) is None
+
+    metadata = helper.prepare(logits_indices, cu_num_logits, has_prefill=True)
+
+    assert metadata is not None
+    assert metadata.logits_indices_padded is logits_indices
+    assert metadata.num_logits_indices == 2
+    assert metadata.max_logits_per_req == 1
diff --git a/tests/v1/worker/cpu/test_cpu_block_table.py b/tests/v1/worker/cpu/test_cpu_block_table.py
new file mode 100644
index 000000000000..5f4fc2d0d7e5
--- /dev/null
+++ b/tests/v1/worker/cpu/test_cpu_block_table.py
@@ -0,0 +1,227 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import pytest
+import torch
+
+from vllm.platforms import CpuArchEnum, current_platform
+from vllm.triton_utils import HAS_TRITON
+from vllm.v1.attention.backends.utils import PAD_SLOT_ID
+from vllm.v1.worker.cpu.block_table import CPUBlockTables
+
+pytestmark = pytest.mark.skipif(
+    not current_platform.is_cpu()
+    or HAS_TRITON
+    or current_platform.get_cpu_architecture() == CpuArchEnum.X86,
+    reason="requires a non-x86 CPU platform without Triton",
+)
+
+
+def test_append_block_ids_updates_cpu_tensors_directly() -> None:
+    block_tables = CPUBlockTables(
+        block_sizes=[8, 16],
+        max_num_reqs=2,
+        max_num_batched_tokens=8,
+        max_num_blocks_per_group=[4, 4],
+        device=torch.device("cpu"),
+    )
+
+    block_tables.append_block_ids(1, ([1, 2], [10]), overwrite=True)
+    block_tables.append_block_ids(1, ([3], [11, 12]), overwrite=False)
+
+    assert block_tables.block_tables[0][1].tolist() == [1, 2, 3, 0]
+    assert block_tables.block_tables[1][1].tolist() == [10, 11, 12, 0]
+    assert block_tables.num_blocks[:, 1].tolist() == [3, 3]
+
+
+def test_append_block_ids_rejects_write_past_row_capacity() -> None:
+    block_tables = CPUBlockTables(
+        block_sizes=[8],
+        max_num_reqs=2,
+        max_num_batched_tokens=8,
+        max_num_blocks_per_group=[4],
+        device=torch.device("cpu"),
+    )
+    block_tables.append_block_ids(1, ([1, 2, 3],), overwrite=True)
+
+    with pytest.raises(
+        RuntimeError,
+        match=r"request 1, group 0 exceeds row capacity \(5 > 4\)",
+    ):
+        block_tables.append_block_ids(1, ([4, 5],), overwrite=False)
+
+    assert block_tables.block_tables[0][1].tolist() == [1, 2, 3, 0]
+    assert block_tables.num_blocks[0, 1].item() == 3
+
+
+def test_gather_block_tables_cpu() -> None:
+    block_tables = CPUBlockTables(
+        block_sizes=[8],
+        max_num_reqs=4,
+        max_num_batched_tokens=8,
+        max_num_blocks_per_group=[6],
+        device=torch.device("cpu"),
+    )
+    source_storage = torch.arange(4 * 8, dtype=torch.int32).view(4, 8)
+    block_tables.block_tables[0].copy_(source_storage[:, 1:7])
+    block_tables.num_blocks.copy_(torch.tensor([[3, 1, 4, 2]], dtype=torch.int32))
+    output_storage = torch.full((5, 8), -1, dtype=torch.int32)
+    output = output_storage[:, 1:7]
+
+    gathered = block_tables.gather_block_tables(
+        idx_mapping=torch.tensor([2, 0], dtype=torch.int64),
+        num_reqs_padded=3,
+        out=(output,),
+    )
+
+    assert gathered[0].data_ptr() == output.data_ptr()
+    assert gathered[0].tolist() == [
+        [17, 18, 19, 20, 0, 0],
+        [1, 2, 3, 0, 0, 0],
+        [0, 0, 0, 0, 0, 0],
+    ]
+    assert output[3:].eq(-1).all()
+
+
+def test_compute_slot_mappings_cpu() -> None:
+    block_tables = CPUBlockTables(
+        block_sizes=[8, 8],
+        max_num_reqs=2,
+        max_num_batched_tokens=8,
+        max_num_blocks_per_group=[3, 3],
+        device=torch.device("cpu"),
+        slot_mapping_enabled=[True, False],
+    )
+    block_tables.block_tables[0].copy_(
+        torch.tensor([[2, 3, 4], [10, 11, 12]], dtype=torch.int32)
+    )
+    output_storage = torch.full((2, 10), 99, dtype=torch.int64)
+    output = output_storage[:, 1:9]
+    positions_storage = torch.tensor(
+        [-99, 0, 4, 100, 101, 3, 8, -99], dtype=torch.int64
+    )
+
+    actual = block_tables.compute_slot_mappings(
+        idx_mapping=torch.tensor([1, -1, 0], dtype=torch.int64),
+        query_start_loc=torch.tensor([0, 2, 4, 6], dtype=torch.int32),
+        positions=positions_storage[1:7],
+        num_tokens_padded=7,
+        out=output,
+    )
+
+    assert actual.data_ptr() == output.data_ptr()
+    assert actual.tolist() == [
+        [80, 84, -1, -1, 19, 24, -1],
+        [-1, -1, -1, -1, -1, -1, -1],
+    ]
+    assert output[:, 7:].eq(-1).all()
+
+
+def test_dummy_block_tables_clear_previous_request_rows() -> None:
+    block_tables = CPUBlockTables(
+        block_sizes=[16],
+        max_num_reqs=4,
+        max_num_batched_tokens=64,
+        max_num_blocks_per_group=[8],
+        device=torch.device("cpu"),
+    )
+    block_tables.append_block_ids(0, ([1, 2],), overwrite=True)
+    block_tables.gather_block_tables(torch.tensor([0]), num_reqs_padded=1)
+    assert block_tables.input_block_tables[0][0, 0].item() == 1
+
+    dummy = block_tables.get_dummy_block_tables(num_reqs=1)
+
+    assert torch.count_nonzero(dummy[0]).item() == 0
+    assert dummy[0].data_ptr() == block_tables.input_block_tables[0].data_ptr()
+
+
+def test_dummy_request_slot_mapping_is_pad() -> None:
+    block_tables = CPUBlockTables(
+        block_sizes=[4],
+        max_num_reqs=2,
+        max_num_batched_tokens=16,
+        max_num_blocks_per_group=[4],
+        device=torch.device("cpu"),
+    )
+    block_tables.append_block_ids(0, ([7, 8],), overwrite=True)
+    query_start_loc = torch.tensor([0, 3], dtype=torch.int32)
+    positions = torch.tensor([1, 2, 3], dtype=torch.int64)
+
+    real = block_tables.compute_slot_mappings(
+        torch.tensor([0]), 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]), query_start_loc, positions, num_tokens_padded=3
+    )
+    assert dummy[0].tolist() == [PAD_SLOT_ID] * 3
+
+
+def test_gather_nulls_blocks_below_sliding_window() -> None:
+    block_tables = CPUBlockTables(
+        block_sizes=[16, 16],
+        max_num_reqs=4,
+        max_num_batched_tokens=64,
+        max_num_blocks_per_group=[16, 16],
+        device=torch.device("cpu"),
+        sliding_window_reach=[0, 31],
+    )
+    full_ids = [list(range(1, 13)), list(range(41, 53))]
+    swa_ids = [list(range(21, 33)), list(range(61, 73))]
+    req_slots = [0, 2]
+    for req, slot in enumerate(req_slots):
+        block_tables.append_block_ids(
+            slot, (full_ids[req], swa_ids[req]), overwrite=True
+        )
+
+    idx_mapping = torch.tensor(req_slots)
+    num_computed_tokens = torch.tensor([145, 0, 120, 0], dtype=torch.int32)
+    full, swa = block_tables.gather_block_tables(
+        idx_mapping,
+        num_reqs_padded=3,
+        num_computed_tokens=num_computed_tokens,
+    )
+
+    assert full[0, :12].tolist() == full_ids[0]
+    assert full[1, :12].tolist() == full_ids[1]
+    assert swa[0, :12].tolist() == [0] * 7 + swa_ids[0][7:]
+    assert swa[1, :12].tolist() == [0] * 5 + swa_ids[1][5:]
+    assert torch.count_nonzero(swa[2]).item() == 0
+    assert block_tables.block_tables[1][0, :12].tolist() == swa_ids[0]
+
+    _, swa = block_tables.gather_block_tables(idx_mapping, num_reqs_padded=2)
+    assert swa[0, :12].tolist() == swa_ids[0]
+
+
+def test_shared_block_tables_clear_rows_from_larger_batch() -> None:
+    tables = CPUBlockTables([16], 8, 64, [3], torch.device("cpu"))
+    for req_index in range(8):
+        tables.append_block_ids(req_index, ([req_index + 1] * 3,), overwrite=True)
+    tables.gather_block_tables(torch.arange(8), num_reqs_padded=8)
+    ptr = tables.input_block_tables[0].data_ptr()
+
+    idx = torch.tensor([7, 2, 0, 5, 1])
+    for num_reqs in (3, 5, 0):
+        num_rows = num_reqs + 1
+        out = tables.gather_block_tables(idx[:num_reqs], num_rows)[0]
+        assert out.shape == (num_rows, 3)
+        assert tables.input_block_tables[0].data_ptr() == ptr
+        assert torch.count_nonzero(tables.input_block_tables[0][num_reqs:]).item() == 0
+        expected = (idx[:num_reqs, None] + 1).expand(-1, 3)
+        torch.testing.assert_close(out[:num_reqs], expected.to(torch.int32))
+
+
+def test_external_block_table_output_preserves_shared_rows() -> None:
+    tables = CPUBlockTables([16], 4, 64, [3], torch.device("cpu"))
+    for req_index in range(4):
+        tables.append_block_ids(req_index, ([req_index + 1] * 3,), overwrite=True)
+    tables.gather_block_tables(torch.arange(4), num_reqs_padded=4)
+    previous = tables.input_block_tables[0].clone()
+    external = torch.full_like(previous[:3], -1)
+
+    result = tables.gather_block_tables(torch.tensor([3, 1]), 3, out=(external,))[0]
+
+    torch.testing.assert_close(tables.input_block_tables[0], previous)
+    assert result[:, 0].tolist() == [4, 2, 0]
+    tables.gather_block_tables(torch.tensor([3]), num_reqs_padded=1)
+    assert torch.count_nonzero(tables.input_block_tables[0][1:]).item() == 0
diff --git a/tests/v1/worker/cpu/test_cpu_ec_connector.py b/tests/v1/worker/cpu/test_cpu_ec_connector.py
new file mode 100644
index 000000000000..c72ea6266ca3
--- /dev/null
+++ b/tests/v1/worker/cpu/test_cpu_ec_connector.py
@@ -0,0 +1,184 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from typing import Any
+
+import pytest
+import torch
+
+from vllm.config import VllmConfig
+from vllm.config.ec_transfer import ECTransferConfig
+from vllm.distributed.ec_transfer.ec_connector.base import (
+    ECConnectorBase,
+    ECConnectorMetadata,
+    ECConnectorRole,
+)
+from vllm.platforms import CpuArchEnum, current_platform
+from vllm.triton_utils import HAS_TRITON
+from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT
+from vllm.v1.worker.cpu import ec_connector as cpu_ec_connector
+from vllm.v1.worker.cpu.ec_connector import (
+    NO_OP_EC_CONNECTOR,
+    ActiveECConnector,
+)
+
+pytestmark = pytest.mark.skipif(
+    not current_platform.is_cpu()
+    or HAS_TRITON
+    or current_platform.get_cpu_architecture() == CpuArchEnum.X86,
+    reason="requires a non-x86 CPU platform without Triton",
+)
+
+WORKER_META = object()
+CONNECTOR_STATS = object()
+
+
+class _ConnectorMetadata(ECConnectorMetadata):
+    pass
+
+
+class _RecordingECConnector(ECConnectorBase):
+    def __init__(self, vllm_config: VllmConfig) -> None:
+        super().__init__(vllm_config, ECConnectorRole.WORKER)
+        self.started_save_cache: dict[str, torch.Tensor] | None = None
+        self.loaded_cache: dict[str, torch.Tensor] | None = None
+        self.saved_hashes: list[str] = []
+        self.finished_req_ids: set[str] | None = None
+        self.metadata_cleared = False
+
+    def start_save_caches(self, **kwargs: Any) -> None:
+        self.started_save_cache = kwargs["encoder_cache"]
+
+    def start_load_caches(
+        self, encoder_cache: dict[str, torch.Tensor], **kwargs: Any
+    ) -> None:
+        self.loaded_cache = encoder_cache
+
+    def save_caches(
+        self,
+        encoder_cache: dict[str, torch.Tensor],
+        mm_hash: str,
+        **kwargs: Any,
+    ) -> None:
+        assert encoder_cache is self.started_save_cache
+        self.saved_hashes.append(mm_hash)
+
+    def get_finished(
+        self, finished_req_ids: set[str]
+    ) -> tuple[set[str] | None, set[str] | None]:
+        self.finished_req_ids = finished_req_ids
+        return {"sent"}, {"received"}
+
+    def build_connector_worker_meta(self) -> Any:
+        return WORKER_META
+
+    def get_ec_connector_stats(self) -> Any:
+        return CONNECTOR_STATS
+
+    def clear_connector_metadata(self) -> None:
+        super().clear_connector_metadata()
+        self.metadata_cleared = True
+
+    def has_cache_item(self, identifier: str) -> bool:
+        return False
+
+    def update_state_after_alloc(self, request: Any, index: int) -> None:
+        pass
+
+    def build_connector_meta(
+        self, scheduler_output: SchedulerOutput
+    ) -> ECConnectorMetadata:
+        return _ConnectorMetadata()
+
+
+def _config(ec_role: str) -> VllmConfig:
+    config = VllmConfig()
+    config.ec_transfer_config = ECTransferConfig(
+        ec_connector="RecordingECConnector",
+        ec_role=ec_role,
+    )
+    return config
+
+
+def _scheduler_output() -> SchedulerOutput:
+    output = SchedulerOutput.make_empty()
+    output.ec_connector_metadata = _ConnectorMetadata()
+    output.finished_req_ids = {"finished"}
+    return output
+
+
+def _connector(
+    monkeypatch: pytest.MonkeyPatch,
+    ec_role: str,
+    encoder_cache: dict[str, torch.Tensor] | None = None,
+) -> tuple[ActiveECConnector, _RecordingECConnector]:
+    config = _config(ec_role)
+    recording_connector = _RecordingECConnector(config)
+    monkeypatch.setattr(
+        cpu_ec_connector, "get_ec_transfer", lambda: recording_connector
+    )
+    return (
+        ActiveECConnector(config, encoder_cache if encoder_cache is not None else {}),
+        recording_connector,
+    )
+
+
+@pytest.mark.parametrize(
+    ("ec_role", "loads", "saves"),
+    [
+        ("ec_producer", False, True),
+        ("ec_both", True, True),
+        ("ec_consumer", True, False),
+    ],
+)
+def test_transfers_encoder_caches_for_connector_role(
+    monkeypatch: pytest.MonkeyPatch,
+    ec_role: str,
+    loads: bool,
+    saves: bool,
+) -> None:
+    encoder_cache = {"old": torch.tensor([1.0])}
+    connector, recording = _connector(monkeypatch, ec_role, encoder_cache)
+
+    with connector.maybe_get_output(_scheduler_output()):
+        encoder_cache["new"] = torch.tensor([2.0])
+
+    assert (recording.loaded_cache is encoder_cache) is loads
+    assert (recording.started_save_cache is encoder_cache) is saves
+    assert recording.saved_hashes == (["new"] if saves else [])
+
+
+def test_reports_connector_results_and_clears_step_metadata(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    connector, recording = _connector(monkeypatch, "ec_both")
+    scheduler_output = _scheduler_output()
+
+    with connector.maybe_get_output(scheduler_output) as output:
+        assert output is not None
+        assert output.ec_connector_worker_meta is None
+        assert output.ec_connector_stats is None
+
+    assert output.finished_sending == {"sent"}
+    assert output.finished_recving == {"received"}
+    assert output.ec_connector_worker_meta is WORKER_META
+    assert output.ec_connector_stats is CONNECTOR_STATS
+    assert recording.finished_req_ids == {"finished"}
+    assert recording.metadata_cleared
+
+
+def test_no_forward_completes_connector_step(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    connector, _ = _connector(monkeypatch, "ec_both")
+
+    output = connector.no_forward(_scheduler_output())
+
+    assert output.ec_connector_output is not None
+    assert output.ec_connector_output.finished_sending == {"sent"}
+    assert output.ec_connector_output.finished_recving == {"received"}
+    assert output.ec_connector_output.ec_connector_worker_meta is WORKER_META
+    assert NO_OP_EC_CONNECTOR.no_forward(_scheduler_output()) is (
+        EMPTY_MODEL_RUNNER_OUTPUT
+    )
diff --git a/tests/v1/worker/cpu/test_cpu_encoder_cache.py b/tests/v1/worker/cpu/test_cpu_encoder_cache.py
new file mode 100644
index 000000000000..d3e00cc813cc
--- /dev/null
+++ b/tests/v1/worker/cpu/test_cpu_encoder_cache.py
@@ -0,0 +1,95 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import weakref
+from dataclasses import replace
+
+import pytest
+import torch
+
+from vllm.multimodal.inputs import (
+    MultiModalFeatureSpec,
+    MultiModalFieldElem,
+    MultiModalKwargsItem,
+    MultiModalSharedField,
+    PlaceholderRange,
+)
+from vllm.platforms import CpuArchEnum, current_platform
+from vllm.triton_utils import HAS_TRITON
+from vllm.v1.worker.cpu.encoder_cache import EncoderCache
+from vllm.v1.worker.cpu.encoder_runner import EncoderRunner
+
+pytestmark = pytest.mark.skipif(
+    not current_platform.is_cpu()
+    or HAS_TRITON
+    or current_platform.get_cpu_architecture() == CpuArchEnum.X86,
+    reason="requires a non-x86 CPU platform without Triton",
+)
+
+HIDDEN_SIZE = 4
+
+
+def _embeds_item(embeds: torch.Tensor) -> MultiModalKwargsItem:
+    return MultiModalKwargsItem(
+        {
+            "embedding": MultiModalFieldElem(
+                data=embeds, field=MultiModalSharedField(batch_size=1)
+            )
+        }
+    )
+
+
+def _feature(identifier: str, offset: int, length: int) -> MultiModalFeatureSpec:
+    return MultiModalFeatureSpec(
+        data=None,
+        modality="image",
+        identifier=identifier,
+        mm_position=PlaceholderRange(offset=offset, length=length),
+    )
+
+
+def test_consumed_input_updates_are_request_local_and_replayable() -> None:
+    pixels = torch.ones(8, HIDDEN_SIZE)
+    feature = replace(_feature("same", 0, 8), data=_embeds_item(pixels))
+    later = replace(feature, mm_position=PlaceholderRange(offset=16, length=8))
+    source = [feature, later]
+    cache = EncoderCache()
+    cache.add_request("a", source)
+    cache.add_request("b", source)
+    cache.encoder_outputs["same"] = torch.ones(8, HIDDEN_SIZE)
+
+    cache.update_encoder_inputs("a", {0: None})
+
+    assert cache.mm_features["a"][0].data is None
+    assert cache.mm_features["a"][0].mm_position is feature.mm_position
+    assert cache.mm_features["a"][1].data is feature.data
+    assert cache.mm_features["b"][0].data is feature.data
+    assert source[0].data is not None
+    assert "same" in cache.encoder_outputs
+
+    cache.free_encoder_cache("same")
+    cache.update_encoder_inputs("a", {0: source[0].data})
+    runner = EncoderRunner(
+        model=torch.nn.Module(),
+        max_num_tokens=32,
+        hidden_size=HIDDEN_SIZE,
+        encoder_cache=cache,
+        dtype=torch.float32,
+        device=torch.device("cpu"),
+    )
+    hashes, inputs = runner.prepare_mm_inputs({"a": [0]})
+    assert hashes == ["same"]
+    assert inputs == [(feature.modality, feature.data)]
+
+
+def test_consumed_input_releases_storage_before_request_finishes() -> None:
+    cache = EncoderCache()
+    pixels = torch.ones(8, HIDDEN_SIZE)
+    ref = weakref.ref(pixels)
+    cache.add_request("a", [replace(_feature("h", 0, 8), data=_embeds_item(pixels))])
+    del pixels
+
+    assert ref() is not None
+    cache.update_encoder_inputs("a", {0: None})
+    assert ref() is None
+    assert "a" in cache.mm_features
diff --git a/tests/v1/worker/cpu/test_cpu_encoder_runner.py b/tests/v1/worker/cpu/test_cpu_encoder_runner.py
new file mode 100644
index 000000000000..5b592d1cec55
--- /dev/null
+++ b/tests/v1/worker/cpu/test_cpu_encoder_runner.py
@@ -0,0 +1,337 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import numpy as np
+import pytest
+import torch
+from transformers import LlamaConfig
+
+from vllm.config import ModelConfig, SchedulerConfig, VllmConfig
+from vllm.entrypoints.pooling.scoring.utils import compress_token_type_ids
+from vllm.multimodal.inputs import (
+    MultiModalBatchedField,
+    MultiModalFeatureSpec,
+    MultiModalFieldElem,
+    MultiModalKwargsItem,
+    PlaceholderRange,
+)
+from vllm.platforms import CpuArchEnum, current_platform
+from vllm.pooling_params import PoolingParams
+from vllm.triton_utils import HAS_TRITON
+from vllm.v1.core.sched.output import NewRequestData
+from vllm.v1.worker.cpu.encoder_cache import EncoderCache
+from vllm.v1.worker.cpu.encoder_runner import EncoderRunner
+from vllm.v1.worker.cpu.input_batch import InputBatch, InputBuffers
+from vllm.v1.worker.cpu.model_states.default import DefaultModelState
+from vllm.v1.worker.cpu.model_states.encoder_only import EncoderOnlyModelState
+from vllm.v1.worker.cpu.states import CPURequestState
+
+pytestmark = pytest.mark.skipif(
+    not current_platform.is_cpu()
+    or HAS_TRITON
+    or current_platform.get_cpu_architecture() == CpuArchEnum.X86,
+    reason="requires a non-x86 CPU platform without Triton",
+)
+
+HIDDEN_SIZE = 4
+
+
+class _EncoderModel(torch.nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.batches: list[list[int]] = []
+
+    def embed_multimodal(self, values: torch.Tensor) -> list[torch.Tensor]:
+        self.batches.append(values.flatten().tolist())
+        return [
+            value.reshape(1, 1).expand(1, HIDDEN_SIZE).to(torch.float32)
+            for value in values
+        ]
+
+
+def _item(value: int) -> MultiModalKwargsItem:
+    return MultiModalKwargsItem(
+        {
+            "values": MultiModalFieldElem(
+                data=torch.tensor([value]),
+                field=MultiModalBatchedField(),
+            )
+        }
+    )
+
+
+def _feature(
+    identifier: str,
+    offset: int,
+    length: int,
+    modality: str = "image",
+    data: MultiModalKwargsItem | None = None,
+) -> MultiModalFeatureSpec:
+    return MultiModalFeatureSpec(
+        data=data,
+        modality=modality,
+        identifier=identifier,
+        mm_position=PlaceholderRange(offset=offset, length=length),
+    )
+
+
+def _make_runner(
+    features: list[MultiModalFeatureSpec],
+    cached: list[MultiModalFeatureSpec],
+) -> EncoderRunner:
+    cache = EncoderCache()
+    cache.mm_features["req0"] = features
+    for feature in cached:
+        length = feature.mm_position.length
+        cache.encoder_outputs[feature.identifier] = torch.arange(
+            length * HIDDEN_SIZE, dtype=torch.float32
+        ).reshape(length, HIDDEN_SIZE)
+    return EncoderRunner(
+        model=None,
+        max_num_tokens=64,
+        hidden_size=HIDDEN_SIZE,
+        encoder_cache=cache,
+        dtype=torch.float32,
+        device=torch.device("cpu"),
+    )
+
+
+def _gather(
+    runner: EncoderRunner, num_scheduled_tokens: int
+) -> tuple[list[torch.Tensor], torch.Tensor]:
+    return runner.gather_mm_embeddings(
+        req_ids=["req0"],
+        total_num_scheduled_tokens=num_scheduled_tokens,
+        num_scheduled_tokens=np.array([num_scheduled_tokens]),
+        query_start_loc=np.array([0]),
+        prefill_lens=np.array([1000]),
+        num_computed_tokens=np.array([0]),
+    )
+
+
+def test_target_path_raises_on_encoder_cache_miss() -> None:
+    runner = _make_runner([_feature("missing", 0, 8)], cached=[])
+
+    with pytest.raises(RuntimeError, match="Encoder cache miss"):
+        _gather(runner, num_scheduled_tokens=8)
+
+
+def test_multi_request_batch_gathers_per_request() -> None:
+    first = _feature("first", 0, 8)
+    second = _feature("second", 0, 8)
+    cache = EncoderCache()
+    cache.mm_features = {"req0": [first], "req1": [second]}
+    for feature in (first, second):
+        cache.encoder_outputs[feature.identifier] = torch.arange(
+            feature.mm_position.length * HIDDEN_SIZE, dtype=torch.float32
+        ).reshape(feature.mm_position.length, HIDDEN_SIZE)
+    runner = EncoderRunner(
+        model=None,
+        max_num_tokens=64,
+        hidden_size=HIDDEN_SIZE,
+        encoder_cache=cache,
+        dtype=torch.float32,
+        device=torch.device("cpu"),
+    )
+
+    mm_embeds, is_mm_embed = runner.gather_mm_embeddings(
+        req_ids=["req0", "req1"],
+        total_num_scheduled_tokens=16,
+        num_scheduled_tokens=np.array([8, 8]),
+        query_start_loc=np.array([0, 8]),
+        prefill_lens=np.array([1000, 1000]),
+        num_computed_tokens=np.array([0, 0]),
+    )
+
+    assert len(mm_embeds) == 2
+    assert [embedding.modality for embedding in mm_embeds] == ["image", "image"]
+    assert is_mm_embed.all()
+
+
+def test_gather_preserves_mixed_modalities() -> None:
+    video = _feature("video", 0, 4, modality="video")
+    audio = _feature("audio", 4, 4, modality="audio")
+    runner = _make_runner([video, audio], cached=[video, audio])
+
+    mm_embeds, is_mm_embed = _gather(runner, num_scheduled_tokens=8)
+
+    assert [embedding.modality for embedding in mm_embeds] == ["video", "audio"]
+    assert is_mm_embed.all()
+
+
+def test_execute_mm_encoder_batches_by_modality_and_restores_order() -> None:
+    model = _EncoderModel()
+    runner = EncoderRunner(
+        model=model,
+        max_num_tokens=64,
+        hidden_size=HIDDEN_SIZE,
+        encoder_cache=EncoderCache(),
+        dtype=torch.float32,
+        device=torch.device("cpu"),
+    )
+    items = [
+        ("video", _item(0)),
+        ("image", _item(1)),
+        ("video", _item(2)),
+        ("image", _item(3)),
+    ]
+
+    outputs = runner.execute_mm_encoder(items)
+
+    assert model.batches == [[1, 3], [0, 2]]
+    assert [int(output[0, 0]) for output in outputs] == [0, 1, 2, 3]
+
+
+@pytest.fixture
+def vllm_config(tmp_path) -> VllmConfig:
+    hf_config = LlamaConfig(
+        hidden_size=HIDDEN_SIZE,
+        intermediate_size=8,
+        num_hidden_layers=1,
+        num_attention_heads=1,
+        num_key_value_heads=1,
+        max_position_embeddings=64,
+        vocab_size=32,
+        architectures=["LlamaForCausalLM"],
+    )
+    hf_config.save_pretrained(tmp_path)
+    model_config = ModelConfig(
+        model=str(tmp_path),
+        dtype="float32",
+        max_model_len=64,
+        skip_tokenizer_init=True,
+    )
+    scheduler_config = SchedulerConfig(
+        max_model_len=64,
+        is_encoder_decoder=False,
+        max_num_batched_tokens=64,
+        max_num_seqs=2,
+        is_multimodal_model=True,
+    )
+    return VllmConfig(
+        model_config=model_config,
+        scheduler_config=scheduler_config,
+    )
+
+
+def test_model_state_caches_encoder_output_without_gathering(
+    vllm_config: VllmConfig,
+) -> None:
+    model = _EncoderModel()
+    cache = EncoderCache()
+    cache.add_request(
+        "req0",
+        [
+            _feature("encoded", 0, 1, data=_item(7)),
+            _feature("not-encoded", 1, 1),
+        ],
+    )
+    state = DefaultModelState(
+        vllm_config,
+        model,
+        cache,
+        torch.device("cpu"),
+    )
+
+    state.execute_mm_encoder({"req0": [0, 1]})
+
+    assert model.batches == [[7]]
+    assert set(cache.encoder_outputs) == {"encoded"}
+    torch.testing.assert_close(
+        cache.encoder_outputs["encoded"],
+        torch.full((1, HIDDEN_SIZE), 7.0),
+    )
+
+
+def test_model_state_empty_encoder_schedule_is_a_noop(
+    vllm_config: VllmConfig,
+) -> None:
+    model = _EncoderModel()
+    cache = EncoderCache()
+    cached = torch.ones(1, HIDDEN_SIZE)
+    cache.encoder_outputs["existing"] = cached
+    state = DefaultModelState(
+        vllm_config,
+        model,
+        cache,
+        torch.device("cpu"),
+    )
+
+    state.execute_mm_encoder({})
+
+    assert model.batches == []
+    assert set(cache.encoder_outputs) == {"existing"}
+    assert cache.encoder_outputs["existing"] is cached
+
+
+def test_encoder_only_state_slices_compressed_token_types_for_chunk(
+    vllm_config: VllmConfig,
+) -> None:
+    tokens = [10, 11, 12, 13, 14, 15]
+    request = NewRequestData(
+        req_id="req0",
+        prompt_token_ids=tokens,
+        mm_features=[],
+        sampling_params=None,
+        pooling_params=PoolingParams(
+            extra_kwargs={
+                "compressed_token_type_ids": compress_token_type_ids([0, 0, 1, 1, 1, 0])
+            }
+        ),
+        block_ids=(),
+        num_computed_tokens=1,
+        lora_request=None,
+        prefill_token_ids=tokens,
+    )
+    req_states = CPURequestState(
+        max_num_reqs=2,
+        max_model_len=64,
+        max_num_batched_tokens=64,
+        vocab_size=32,
+        device=torch.device("cpu"),
+    )
+    req_states.add_request(
+        req_id=request.req_id,
+        prompt_len=len(tokens),
+        all_token_ids=tokens,
+        num_computed_tokens=1,
+        max_tokens=1,
+    )
+    req_index = req_states.req_id_to_index[request.req_id]
+    state = EncoderOnlyModelState(
+        vllm_config, torch.nn.Module(), None, torch.device("cpu")
+    )
+    state.add_request(req_index, request)
+
+    input_batch = InputBatch.make_dummy(
+        num_reqs=1,
+        num_tokens=5,
+        input_buffers=InputBuffers(2, 64, torch.device("cpu")),
+        is_padding=False,
+    )
+    input_batch.req_ids[0] = request.req_id
+    input_batch.idx_mapping.fill_(req_index)
+    input_batch.idx_mapping_np.fill(req_index)
+    input_batch.expanded_idx_mapping.fill_(req_index)
+    input_batch.num_computed_tokens_np[0] = 1
+
+    model_inputs = state.prepare_inputs(input_batch, req_states)
+
+    token_type_ids = model_inputs["token_type_ids"]
+    assert token_type_ids is not None
+    assert token_type_ids.tolist() == [0, 1, 1, 1, 0]
+
+
+def test_encoder_timing_stats_registry() -> None:
+    runner = _make_runner([], [])
+    runner.enable_timing = True
+
+    with runner.timed_encoder_operation({"req0"}):
+        pass
+    with runner.timed_encoder_operation({"req0"}):
+        pass
+
+    stats = runner.get_encoder_timing_stats()
+    assert stats["req0"]["num_encoder_calls"] == 2
+    assert stats["req0"]["encoder_forward_secs"] >= 0
+    assert runner.get_encoder_timing_stats() == {}
diff --git a/tests/v1/worker/cpu/test_cpu_gumbel_sample.py b/tests/v1/worker/cpu/test_cpu_gumbel_sample.py
new file mode 100644
index 000000000000..11e85fcc272f
--- /dev/null
+++ b/tests/v1/worker/cpu/test_cpu_gumbel_sample.py
@@ -0,0 +1,423 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import math
+
+import pytest
+import torch
+
+from vllm.platforms import CpuArchEnum, current_platform
+from vllm.triton_utils import HAS_TRITON
+from vllm.v1.worker.cpu.gumbel import gumbel_sample
+from vllm.v1.worker.cpu.sampler import CPUSamplingStates
+
+if current_platform.is_cpu():
+    import vllm._C  # noqa: F401
+
+pytestmark = pytest.mark.skipif(
+    not current_platform.is_cpu()
+    or HAS_TRITON
+    or current_platform.get_cpu_architecture() == CpuArchEnum.X86,
+    reason="requires a non-x86 CPU platform without Triton",
+)
+
+_UINT32_MASK = (1 << 32) - 1
+_UINT64_MASK = (1 << 64) - 1
+VOCAB_SIZE = 200_000
+NUM_SAMPLES = 500_000
+HEAD_LOG_GAP = 18.0
+Z_TOLERANCE = 10.0
+
+
+def _sample(
+    logits: torch.Tensor,
+    expanded_idx_mapping: torch.Tensor,
+    temperature: torch.Tensor,
+    seeds: torch.Tensor,
+    positions: torch.Tensor,
+    apply_temperature: bool = True,
+    use_fp64: bool = False,
+) -> torch.Tensor:
+    return gumbel_sample(
+        logits,
+        expanded_idx_mapping,
+        temperature,
+        seeds,
+        positions,
+        apply_temperature=apply_temperature,
+        is_drafting=False,
+        use_fp64=use_fp64,
+    )
+
+
+def _sample_many(
+    logits_1d: torch.Tensor,
+    num_samples: int,
+    *,
+    use_fp64: bool = False,
+) -> torch.Tensor:
+    vocab_size = logits_1d.shape[0]
+    logits = logits_1d.unsqueeze(0).expand(num_samples, vocab_size)
+    mapping = torch.zeros(num_samples, dtype=torch.int64)
+    temperature = torch.ones(1)
+    seeds = torch.tensor([0xABCD], dtype=torch.int64)
+    positions = torch.arange(num_samples, dtype=torch.int64)
+    return _sample(
+        logits,
+        mapping,
+        temperature,
+        seeds,
+        positions,
+        use_fp64=use_fp64,
+    )
+
+
+def _sample_histogram(
+    logits_1d: torch.Tensor,
+    num_samples: int,
+    *,
+    chunk_size: int = 1_000_000,
+) -> torch.Tensor:
+    histogram = torch.zeros(logits_1d.shape[0], dtype=torch.float64)
+    for start in range(0, num_samples, chunk_size):
+        size = min(chunk_size, num_samples - start)
+        vocab_size = logits_1d.shape[0]
+        logits = logits_1d.unsqueeze(0).expand(size, vocab_size)
+        mapping = torch.zeros(size, dtype=torch.int64)
+        temperature = torch.ones(1)
+        seeds = torch.tensor([0xABCD], dtype=torch.int64)
+        positions = torch.arange(start, start + size, dtype=torch.int64)
+        sampled = _sample(logits, mapping, temperature, seeds, positions)
+        histogram += torch.bincount(sampled, minlength=vocab_size).double()
+    return histogram
+
+
+def _murmur3_rotl32(value: int, shift: int) -> int:
+    return ((value << shift) | (value >> (32 - shift))) & _UINT32_MASK
+
+
+def _murmur3_mix(hash_: int, key: int) -> int:
+    key = (key * 0xCC9E2D51) & _UINT32_MASK
+    key = _murmur3_rotl32(key, 15)
+    key = (key * 0x1B873593) & _UINT32_MASK
+    hash_ ^= key
+    hash_ = _murmur3_rotl32(hash_, 13)
+    return (hash_ * 5 + 0xE6546B64) & _UINT32_MASK
+
+
+def _murmur3_fmix32(hash_: int) -> int:
+    hash_ ^= hash_ >> 16
+    hash_ = (hash_ * 0x85EBCA6B) & _UINT32_MASK
+    hash_ ^= hash_ >> 13
+    hash_ = (hash_ * 0xC2B2AE35) & _UINT32_MASK
+    return (hash_ ^ (hash_ >> 16)) & _UINT32_MASK
+
+
+def _murmur3_hash32(seed: int, position: int, token_id: int, domain: int = 0) -> int:
+    seed_bits = seed & _UINT64_MASK
+    hash_ = domain
+    hash_ = _murmur3_mix(hash_, seed_bits & _UINT32_MASK)
+    hash_ = _murmur3_mix(hash_, seed_bits >> 32)
+    hash_ = _murmur3_mix(hash_, position & _UINT32_MASK)
+    hash_ = _murmur3_mix(hash_, token_id & _UINT32_MASK)
+    return _murmur3_fmix32(hash_ ^ 16)
+
+
+def _murmur3_uniform32(seed: int, position: int, vocab_size: int) -> torch.Tensor:
+    hashes = torch.tensor(
+        [_murmur3_hash32(seed, position, token_id) for token_id in range(vocab_size)],
+        dtype=torch.int64,
+    )
+    high = (hashes >> 16).to(torch.float32)
+    low = (hashes & 0xFFFF).to(torch.float32)
+    return high * 1.52587890625e-05 + (low + 0.5) * 2.3283064365386963e-10
+
+
+def _murmur3_uniform64(seed: int, position: int, vocab_size: int) -> torch.Tensor:
+    random53 = []
+    for token_id in range(vocab_size):
+        low = _murmur3_hash32(seed, position, token_id)
+        high = _murmur3_hash32(seed, position, token_id, 0x9E3779B9)
+        random53.append(((high << 32) | low) >> 11)
+
+    uniform = (torch.tensor(random53, dtype=torch.float64) + 0.5) * 2.0**-53
+    return torch.where(uniform == 1.0, uniform - 2.0**-53, uniform)
+
+
+def _log1p_neg_stable(value: torch.Tensor) -> torch.Tensor:
+    polynomial = torch.full_like(value, 1.0 / 8.0)
+    for denominator in range(7, 1, -1):
+        polynomial = 1.0 / denominator + value * polynomial
+    polynomial = 1.0 + value * polynomial
+    series = -value * polynomial
+    direct = torch.log(torch.clamp(1.0 - value, min=2.0**-24))
+    return torch.where(value < 0.25, series, direct)
+
+
+def _reference_sample(
+    logits: torch.Tensor,
+    expanded_idx_mapping: torch.Tensor,
+    temperature: torch.Tensor,
+    seeds: torch.Tensor,
+    positions: torch.Tensor,
+    apply_temperature: bool = True,
+    use_fp64: bool = False,
+) -> torch.Tensor:
+    output = []
+    for row_idx, row in enumerate(logits.float()):
+        req_state_idx = int(expanded_idx_mapping[row_idx])
+        temp = float(temperature[req_state_idx]) if req_state_idx >= 0 else 0.0
+        if temp == 0.0:
+            output.append(int(row.argmax()))
+            continue
+
+        position = int(positions[row_idx])
+        seed = int(seeds[req_state_idx])
+
+        if use_fp64:
+            uniform = _murmur3_uniform64(seed, position, row.numel())
+            noise = -torch.log(-torch.log(uniform))
+        else:
+            uniform = _murmur3_uniform32(seed, position, row.numel())
+            noise = -torch.log(-_log1p_neg_stable(uniform))
+
+        if apply_temperature:
+            row = row / temp
+        output.append(int((row.to(noise.dtype) + noise).argmax()))
+
+    return torch.tensor(output, dtype=torch.int64)
+
+
+def _make_heavy_tailed_counts(seed: int = 1234) -> torch.Tensor:
+    counts = torch.randint(
+        1,
+        4,
+        (VOCAB_SIZE,),
+        generator=torch.Generator().manual_seed(seed),
+        dtype=torch.int64,
+    )
+    counts[0] = round(math.exp(HEAD_LOG_GAP))
+    return counts
+
+
+def _counts_to_logits(counts: torch.Tensor) -> torch.Tensor:
+    return counts.double().log().float()
+
+
+def _z_score(observed: int, expected: float, num_trials: int) -> float:
+    probability = expected / num_trials
+    return (observed - expected) / math.sqrt(
+        num_trials * probability * (1.0 - probability)
+    )
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.float16])
+@pytest.mark.parametrize("use_fp64", [False, True])
+def test_random_and_mixed_rows_match_reference(
+    dtype: torch.dtype, use_fp64: bool
+) -> None:
+    logits = torch.randn(4, 257, generator=torch.Generator().manual_seed(1234)).to(
+        dtype
+    )
+    mapping = torch.tensor([3, 0, 2, 1], dtype=torch.int64)
+    temperature = torch.tensor([0.7, 0.0, 1.3, 1.0])
+    seeds = torch.tensor([0xABCD, -1, (1 << 63) - 1, 17], dtype=torch.int64)
+    positions = torch.tensor([0, 11, 123456789, 19], dtype=torch.int64)
+
+    expected = _reference_sample(
+        logits,
+        mapping,
+        temperature,
+        seeds,
+        positions,
+        use_fp64=use_fp64,
+    )
+    actual = _sample(
+        logits,
+        mapping,
+        temperature,
+        seeds,
+        positions,
+        use_fp64=use_fp64,
+    )
+
+    assert torch.equal(actual, expected)
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.float16])
+def test_greedy_supports_dtypes_strides_and_invalid_mapping(
+    dtype: torch.dtype,
+) -> None:
+    storage = torch.full((3, 48), -100.0, dtype=dtype)
+    logits = storage[:, 1:46]
+    logits[0, [0, 17, 44]] = torch.tensor([1.0, 5.0, 3.0], dtype=dtype)
+    logits[1, [0, 31, 44]] = torch.tensor([-3.0, -2.0, -1.0], dtype=dtype)
+    logits[2, [3, 40]] = 7.0
+    mapping = torch.tensor([2, -1, 0], dtype=torch.int64)
+    temperature = torch.tensor([0.0, 1.0, 0.0])
+    seeds = torch.tensor([11, 22, 33], dtype=torch.int64)
+    positions = torch.tensor([100, 101, 102], dtype=torch.int64)
+
+    actual = _sample(logits, mapping, temperature, seeds, positions)
+
+    assert torch.equal(actual, logits.float().argmax(dim=-1))
+
+
+def test_pre_scaled_logits_match_temperature_application() -> None:
+    logits = torch.randn(4, 1031, generator=torch.Generator().manual_seed(7))
+    mapping = torch.tensor([2, 0, 3, 1], dtype=torch.int64)
+    temperature = torch.tensor([0.5, 1.5, 0.8, 2.0])
+    seeds = torch.tensor([19, 23, 29, 31], dtype=torch.int64)
+    positions = torch.tensor([31, 37, 41, 43], dtype=torch.int64)
+    scaled_logits = logits / temperature[mapping].unsqueeze(1)
+
+    expected = _sample(
+        scaled_logits,
+        mapping,
+        temperature,
+        seeds,
+        positions,
+        apply_temperature=False,
+    )
+    actual = _sample(
+        logits,
+        mapping,
+        temperature,
+        seeds,
+        positions,
+        apply_temperature=True,
+    )
+
+    assert torch.equal(actual, expected)
+
+
+def test_sampling_states_apply_temperature() -> None:
+    states = CPUSamplingStates(4, 17, torch.device("cpu"))
+    temperatures = torch.tensor([0.5, 0.0, 2.0, 1.0])
+    states.temperature.copy_(temperatures)
+
+    logits = torch.arange(1, 69, dtype=torch.float32).reshape(4, 17)
+    expanded_idx_mapping = torch.tensor([2, 0, 3, 1], dtype=torch.int64)
+    expected = logits / torch.tensor([2.0, 0.5, 1.0, 1.0]).unsqueeze(1)
+
+    states.apply_temperature(logits, expanded_idx_mapping, torch.arange(4).numpy())
+
+    assert torch.equal(logits, expected)
+
+
+def test_sampling_states_apply_min_p() -> None:
+    states = CPUSamplingStates(4, 4, torch.device("cpu"))
+    min_p = torch.tensor([0.0, 0.25, 0.5, 1.0])
+    states.min_p.copy_(min_p)
+
+    logits = torch.tensor(
+        [
+            [2.0, 1.0, 0.0, -1.0],
+            [3.0, 2.0, 1.0, 0.0],
+            [2.0, 2.0, 1.0, 0.0],
+            [4.0, 3.0, 2.0, 1.0],
+        ]
+    )
+    expanded_idx_mapping = torch.tensor([2, 0, 3, 1], dtype=torch.int64)
+    expected = torch.tensor(
+        [
+            [2.0, -torch.inf, -torch.inf, -torch.inf],
+            [3.0, 2.0, 1.0, 0.0],
+            [2.0, 2.0, -torch.inf, -torch.inf],
+            [4.0, 3.0, -torch.inf, -torch.inf],
+        ]
+    )
+
+    states.apply_min_p(logits, expanded_idx_mapping, torch.arange(4).numpy())
+
+    assert torch.equal(logits, expected)
+
+
+@pytest.mark.parametrize("use_fp64", [False, True])
+def test_sampling_matches_target_distribution(use_fp64: bool) -> None:
+    counts = _make_heavy_tailed_counts()
+    total = counts.sum().item()
+    sampled = _sample_many(_counts_to_logits(counts), NUM_SAMPLES, use_fp64=use_fp64)
+
+    assert sampled.min() >= 0 and sampled.max() < VOCAB_SIZE
+    tail_probability = (total - counts[0].item()) / total
+    tail_count = (sampled != 0).sum().item()
+    z_score = _z_score(tail_count, NUM_SAMPLES * tail_probability, NUM_SAMPLES)
+    assert abs(z_score) < Z_TOLERANCE
+
+
+def test_full_vocab_distribution_fidelity() -> None:
+    counts = torch.randint(
+        500,
+        1500,
+        (VOCAB_SIZE,),
+        generator=torch.Generator().manual_seed(2024),
+        dtype=torch.int64,
+    )
+    total = counts.sum().item()
+    num_samples = 4_000_000
+    histogram = _sample_histogram(_counts_to_logits(counts), num_samples)
+
+    coverage = (histogram > 0).sum().item() / VOCAB_SIZE
+    assert coverage > 0.99
+
+    expected = counts.double() / total * num_samples
+    chi_square = (((histogram - expected) ** 2) / expected).sum().item()
+    degrees_of_freedom = VOCAB_SIZE - 1
+    assert chi_square < degrees_of_freedom + 10 * math.sqrt(2 * degrees_of_freedom)
+
+
+def test_greedy_temperature_zero_returns_argmax() -> None:
+    torch.manual_seed(0)
+    num_requests = 128
+    logits = torch.randn(num_requests, VOCAB_SIZE)
+    mapping = torch.arange(num_requests, dtype=torch.int64)
+    temperature = torch.zeros(num_requests)
+    seeds = torch.arange(num_requests, dtype=torch.int64)
+    positions = torch.arange(num_requests, dtype=torch.int64)
+
+    sampled = _sample(logits, mapping, temperature, seeds, positions)
+
+    assert torch.equal(sampled, logits.argmax(dim=-1))
+
+
+def test_zero_count_tokens_are_never_sampled() -> None:
+    counts = _make_heavy_tailed_counts(seed=7)
+    counts[1::2] = 0
+    sampled = _sample_many(_counts_to_logits(counts), NUM_SAMPLES)
+
+    zero_count = torch.zeros(VOCAB_SIZE, dtype=torch.bool)
+    zero_count[1::2] = True
+    assert sampled.min() >= 0 and sampled.max() < VOCAB_SIZE
+    assert not zero_count[sampled].any()
+
+
+def test_single_nonzero_token_is_always_sampled() -> None:
+    counts = torch.zeros(VOCAB_SIZE, dtype=torch.int64)
+    counts[123_456] = 1000
+
+    sampled = _sample_many(_counts_to_logits(counts), 10_000)
+
+    assert (sampled == 123_456).all()
+
+
+@pytest.mark.parametrize("vocab_size", [1, 999, 1024, 4097])
+def test_vocab_size_not_multiple_of_block(vocab_size: int) -> None:
+    counts = torch.randint(
+        20,
+        200,
+        (vocab_size,),
+        generator=torch.Generator().manual_seed(vocab_size),
+        dtype=torch.int64,
+    )
+    total = counts.sum().item()
+    num_samples = max(40 * vocab_size, 50_000)
+    sampled = _sample_many(_counts_to_logits(counts), num_samples)
+
+    assert sampled.min() >= 0 and sampled.max() < vocab_size
+    observed = torch.bincount(sampled, minlength=vocab_size).double()
+    expected = counts.double() / total * num_samples
+    degrees_of_freedom = vocab_size - 1
+    if degrees_of_freedom > 0:
+        chi_square = (((observed - expected) ** 2) / expected).sum().item()
+        assert chi_square < degrees_of_freedom + 10 * math.sqrt(2 * degrees_of_freedom)
diff --git a/tests/v1/worker/cpu/test_cpu_input_batch.py b/tests/v1/worker/cpu/test_cpu_input_batch.py
new file mode 100644
index 000000000000..f715aa99ee44
--- /dev/null
+++ b/tests/v1/worker/cpu/test_cpu_input_batch.py
@@ -0,0 +1,258 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+
+from vllm import _custom_ops as ops
+from vllm.platforms import CpuArchEnum, current_platform
+from vllm.triton_utils import HAS_TRITON
+from vllm.v1.worker.cpu.input_batch import (
+    InputBatch,
+    InputBuffers,
+    combine_sampled_and_draft_tokens,
+    post_update,
+    prepare_pos_seq_lens,
+    prepare_prefill_inputs,
+)
+from vllm.v1.worker.cpu.rope import RopeState
+
+pytestmark = pytest.mark.skipif(
+    not current_platform.is_cpu()
+    or HAS_TRITON
+    or current_platform.get_cpu_architecture() == CpuArchEnum.X86,
+    reason="requires a non-x86 CPU platform without Triton",
+)
+
+
+@pytest.mark.parametrize(
+    "num_reqs,num_tokens",
+    [(256, 496), (128, 512), (3, 8), (1, 7)],
+)
+def test_make_dummy_distributes_remainder(num_reqs: int, num_tokens: int) -> None:
+    buffers = InputBuffers(
+        max_num_reqs=num_reqs,
+        max_num_tokens=num_tokens,
+        device=torch.device("cpu"),
+    )
+    batch = InputBatch.make_dummy(num_reqs, num_tokens, buffers)
+
+    max_per_req = -(-num_tokens // num_reqs)
+    assert batch.num_scheduled_tokens.sum() == num_tokens
+    assert batch.num_scheduled_tokens.max() == max_per_req
+    assert batch.num_scheduled_tokens.min() >= num_tokens // num_reqs
+    assert (batch.num_scheduled_tokens[:-1] <= batch.num_scheduled_tokens[1:]).all()
+
+    query_lens = batch.query_start_loc_np[1:] - batch.query_start_loc_np[:-1]
+    assert (query_lens == batch.num_scheduled_tokens).all()
+    assert torch.equal(batch.seq_lens, torch.from_numpy(batch.num_scheduled_tokens))
+    assert batch.query_start_loc_np[-1] == num_tokens
+    assert torch.equal(
+        batch.query_start_loc, torch.from_numpy(batch.query_start_loc_np)
+    )
+
+
+def test_prepare_prefill_inputs_cpu() -> None:
+    input_ids_storage = torch.full((8,), -1, dtype=torch.int32)
+    input_ids = input_ids_storage[1:7]
+    next_prefill_storage = torch.full((1, 8), -1, dtype=torch.int32)
+    next_prefill_tokens = next_prefill_storage[:, 1:7]
+    all_token_ids_storage = torch.full((6, 12), -1, dtype=torch.int32)
+    all_token_ids = all_token_ids_storage[:, 1:11]
+    for req_state_idx in range(6):
+        all_token_ids[req_state_idx] = req_state_idx * 100 + torch.arange(
+            10, dtype=torch.int32
+        )
+
+    idx_mapping = torch.tensor([4, 1, 3], dtype=torch.int64)
+    query_start_loc = torch.tensor([0, 3, 5, 6], dtype=torch.int32)
+    prefill_len = torch.tensor([0, 4, 0, 8, 7, 0], dtype=torch.int32)
+    num_computed_tokens = torch.tensor([0, 4, 0, 4, 2, 0], dtype=torch.int32)
+
+    prepare_prefill_inputs(
+        input_ids,
+        next_prefill_tokens,
+        idx_mapping,
+        query_start_loc,
+        all_token_ids,
+        prefill_len,
+        num_computed_tokens,
+    )
+
+    assert input_ids.tolist() == [402, 403, 404, -1, -1, 304]
+    assert next_prefill_tokens[:, 4].tolist() == [405]
+    assert next_prefill_tokens[:, 3].tolist() == [305]
+    assert next_prefill_tokens[:, 1].tolist() == [-1]
+    assert torch.all(next_prefill_tokens[:, [0, 2, 5]] == -1)
+
+
+def test_get_prompt_logprobs_token_ids_cpu() -> None:
+    all_token_ids_storage = torch.full((6, 12), -1, dtype=torch.int32)
+    all_token_ids = all_token_ids_storage[:, 1:11]
+    for req_state_idx in range(6):
+        all_token_ids[req_state_idx] = req_state_idx * 100 + torch.arange(
+            10, dtype=torch.int32
+        )
+
+    token_ids = ops.get_prompt_logprobs_token_ids_cpu(
+        num_tokens=6,
+        query_start_loc=torch.tensor([0, 3, 5, 6], dtype=torch.int32),
+        idx_mapping=torch.tensor([4, 1, 3], dtype=torch.int64),
+        num_computed_tokens=torch.tensor([0, 4, 0, 1, 2, 0], dtype=torch.int32),
+        all_token_ids=all_token_ids,
+    )
+
+    assert token_ids.dtype == torch.int64
+    assert token_ids.tolist() == [403, 404, 405, 105, 106, 302]
+
+
+def test_prepare_pos_seq_lens_cpu() -> None:
+    positions_storage = torch.full((8,), -1, dtype=torch.int64)
+    positions = positions_storage[1:]
+    seq_lens_storage = torch.full((8,), -1, dtype=torch.int32)
+    seq_lens = seq_lens_storage[1:7]
+
+    idx_mapping = torch.tensor([3, 1, 4], dtype=torch.int64)
+    query_start_loc = torch.tensor([0, 3, 5, 6], dtype=torch.int32)
+    num_computed_tokens = torch.tensor([0, 7, 0, 2, 10], dtype=torch.int32)
+
+    prepare_pos_seq_lens(
+        idx_mapping,
+        query_start_loc,
+        num_computed_tokens,
+        positions,
+        seq_lens,
+    )
+
+    assert positions.tolist() == [2, 3, 4, 7, 8, 10, -1]
+    assert seq_lens.tolist() == [5, 9, 11, 0, 0, 0]
+    assert positions_storage[0].item() == -1
+    assert seq_lens_storage[[0, 7]].tolist() == [-1, -1]
+
+
+def test_prepare_mrope_positions_cpu() -> None:
+    num_dims = 3
+    positions_storage = torch.full((num_dims, 10), -1, dtype=torch.int64)
+    positions = positions_storage[:, 1:9]
+
+    prefill_positions_storage = torch.full((15, 10), -1, dtype=torch.int32)
+    prefill_positions = prefill_positions_storage[:, 1:9]
+    for req_state_idx in range(5):
+        for dim in range(num_dims):
+            prefill_positions[req_state_idx * num_dims + dim] = (
+                req_state_idx * 1000 + dim * 100 + torch.arange(8, dtype=torch.int32)
+            )
+
+    ops.prepare_mrope_positions_cpu(
+        positions,
+        prefill_positions,
+        prefill_delta=torch.tensor([0, 10, 0, 20, 30], dtype=torch.int32),
+        idx_mapping=torch.tensor([3, 1, 4], dtype=torch.int64),
+        query_start_loc=torch.tensor([0, 3, 6, 8], dtype=torch.int32),
+        prefill_lens=torch.tensor([0, 5, 0, 6, 4], dtype=torch.int32),
+        num_computed_tokens=torch.tensor([0, 5, 0, 2, 3], dtype=torch.int32),
+    )
+
+    assert positions.tolist() == [
+        [3002, 3003, 3004, 15, 16, 17, 4003, 34],
+        [3102, 3103, 3104, 15, 16, 17, 4103, 34],
+        [3202, 3203, 3204, 15, 16, 17, 4203, 34],
+    ]
+    assert torch.all(positions_storage[:, [0, 9]] == -1)
+
+
+def test_cpu_rope_state_prepares_model_positions() -> None:
+    class MRoPEModel(torch.nn.Module):
+        supports_mrope = True
+
+        def get_mrope_input_positions(
+            self, input_tokens: list[int], mm_features: list
+        ) -> tuple[torch.Tensor, int]:
+            return torch.tensor([[100], [200]], dtype=torch.int64), 10
+
+    state = RopeState(
+        num_dims=2,
+        max_num_reqs=2,
+        max_num_tokens=2,
+        max_model_len=2,
+        device=torch.device("cpu"),
+    )
+    state.init_prefill_positions(
+        req_idx=1,
+        model=MRoPEModel(),
+        prefill_token_ids=[42],
+        mm_features=[],
+    )
+
+    state.prepare_positions(
+        idx_mapping=torch.tensor([1], dtype=torch.int64),
+        query_start_loc=torch.tensor([0, 2], dtype=torch.int32),
+        prefill_lens=torch.tensor([0, 1], dtype=torch.int32),
+        num_computed_tokens=torch.tensor([0, 0], dtype=torch.int32),
+    )
+
+    assert state.get_positions(2).tolist() == [[100, 11], [200, 11]]
+
+    state.update_prefill_positions(
+        1, torch.tensor([[300], [400]], dtype=torch.int64), delta=20
+    )
+    state.prepare_positions(
+        idx_mapping=torch.tensor([1], dtype=torch.int64),
+        query_start_loc=torch.tensor([0, 2], dtype=torch.int32),
+        prefill_lens=torch.tensor([0, 1], dtype=torch.int32),
+        num_computed_tokens=torch.tensor([0, 0], dtype=torch.int32),
+    )
+
+    assert state.read_prefill_positions(1, 1).tolist() == [[300], [400]]
+    assert state.get_positions(2).tolist() == [[300, 21], [400, 21]]
+
+
+def test_combine_sampled_and_draft_tokens_cpu_basic() -> None:
+    input_ids = torch.tensor([10, 11, 12, -1], dtype=torch.int32)
+    idx_mapping = torch.tensor([3, 1], dtype=torch.int64)
+    last_sampled_tokens = torch.tensor([[40], [41], [42], [43]])
+    query_start_loc = torch.tensor([0, 3, 4], dtype=torch.int32)
+    seq_lens = torch.tensor([5, 6], dtype=torch.int32)
+    prefill_len = torch.tensor([0, 5, 0, 5], dtype=torch.int32)
+    draft_tokens = torch.empty((4, 0), dtype=torch.int64)
+    cu_num_logits = torch.tensor([0, 1, 2], dtype=torch.int32)
+
+    logits_indices = combine_sampled_and_draft_tokens(
+        input_ids,
+        idx_mapping,
+        last_sampled_tokens,
+        query_start_loc,
+        seq_lens,
+        prefill_len,
+        draft_tokens,
+        cu_num_logits,
+        num_logits=2,
+    )
+
+    assert input_ids.tolist() == [10, 11, 12, 41]
+    assert logits_indices.tolist() == [2, 3]
+
+
+def test_post_update_cpu_without_optional_inputs() -> None:
+    num_computed_tokens = torch.tensor([8, 13], dtype=torch.int32)
+    last_sampled_tokens = torch.tensor([[70], [80]], dtype=torch.int64)
+    all_token_ids = torch.full((2, 6), -1, dtype=torch.int32)
+    total_len = torch.tensor([2, 3], dtype=torch.int32)
+
+    post_update(
+        idx_mapping=torch.tensor([1], dtype=torch.int64),
+        num_computed_tokens=num_computed_tokens,
+        last_sampled_tokens=last_sampled_tokens,
+        output_bin_counts=None,
+        sampled_tokens=torch.tensor([[4]], dtype=torch.int64),
+        num_sampled=torch.tensor([1], dtype=torch.int32),
+        num_rejected=torch.tensor([0], dtype=torch.int32),
+        query_start_loc=None,
+        all_token_ids=all_token_ids,
+        total_len=total_len,
+    )
+
+    assert num_computed_tokens.tolist() == [8, 13]
+    assert last_sampled_tokens.squeeze(1).tolist() == [70, 4]
+    assert total_len.tolist() == [2, 4]
+    assert all_token_ids[1].tolist() == [-1, -1, -1, 4, -1, -1]
diff --git a/tests/v1/worker/cpu/test_cpu_kv_block_zeroer.py b/tests/v1/worker/cpu/test_cpu_kv_block_zeroer.py
new file mode 100644
index 000000000000..6fb7bed5d8e2
--- /dev/null
+++ b/tests/v1/worker/cpu/test_cpu_kv_block_zeroer.py
@@ -0,0 +1,91 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from tests.v1.attention.utils import dense_kv_cache_views
+from vllm.platforms import CpuArchEnum, current_platform
+from vllm.triton_utils import HAS_TRITON
+from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheLayout
+from vllm.v1.worker.cpu.utils import CPUKVBlockZeroer
+from vllm.v1.worker.utils import AttentionGroup
+
+pytestmark = pytest.mark.skipif(
+    not current_platform.is_cpu()
+    or HAS_TRITON
+    or current_platform.get_cpu_architecture() == CpuArchEnum.X86,
+    reason="Requires a non-x86 CPU without Triton",
+)
+
+
+@pytest.mark.parametrize("layout", list(KVCacheLayout))
+def test_cpu_zeroes_exactly_one_block_per_layer(layout: KVCacheLayout):
+    device = torch.device("cpu")
+    num_blocks, num_layers = 4, 2
+    spec = FullAttentionSpec(
+        block_size=4, num_kv_heads=2, head_size=8, dtype=torch.float32
+    )
+    raw = torch.ones(
+        num_blocks * num_layers * spec.page_size_bytes,
+        dtype=torch.int8,
+        device=device,
+    )
+    views = dense_kv_cache_views(raw, spec, num_blocks, num_layers, layout)
+    groups = [
+        AttentionGroup(
+            backend=None,
+            layer_names=[f"layer.{i}" for i in range(num_layers)],
+            kv_cache_spec=spec,
+            kv_cache_group_id=0,
+        )
+    ]
+    zeroer = CPUKVBlockZeroer(
+        device,
+        attn_groups_iter=iter(groups),
+        kernel_block_sizes=[spec.block_size],
+        static_forward_context={
+            f"layer.{i}": SimpleNamespace(kv_cache=views[i]) for i in range(num_layers)
+        },
+        num_blocks=num_blocks,
+    )
+
+    zeroer.zero_block_ids([2])
+
+    for view in views:
+        assert (view[2] == 0).all(), layout
+        assert (view[[0, 1, 3]].view(torch.int8) == 1).all(), layout
+    assert int((raw == 0).sum()) == num_layers * spec.page_size_bytes, layout
+
+
+def test_cpu_zeroes_every_kernel_block_in_a_logical_block():
+    device = torch.device("cpu")
+    num_blocks = 4
+    spec = FullAttentionSpec(
+        block_size=4, num_kv_heads=2, head_size=8, dtype=torch.bfloat16
+    )
+    raw = torch.ones(num_blocks * spec.page_size_bytes, dtype=torch.int8)
+    cache = dense_kv_cache_views(
+        raw,
+        spec,
+        num_blocks,
+        num_layers=1,
+        layout=KVCacheLayout.LBNHC,
+        kernel_block_size=2,
+    )[0]
+    zeroer = CPUKVBlockZeroer(
+        device,
+        attn_groups_iter=iter(
+            [AttentionGroup(None, ["layer"], spec, kv_cache_group_id=0)]
+        ),
+        kernel_block_sizes=[2],
+        static_forward_context={"layer": SimpleNamespace(kv_cache=cache)},
+        num_blocks=num_blocks,
+    )
+
+    zeroer.zero_block_ids([1, 3])
+
+    assert (cache[[2, 3, 6, 7]] == 0).all()
+    assert (cache[[0, 1, 4, 5]].view(torch.int8) == 1).all()
diff --git a/tests/v1/worker/cpu/test_cpu_kv_connector.py b/tests/v1/worker/cpu/test_cpu_kv_connector.py
new file mode 100644
index 000000000000..72e18bcd4699
--- /dev/null
+++ b/tests/v1/worker/cpu/test_cpu_kv_connector.py
@@ -0,0 +1,175 @@
+# 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
+
+import vllm.v1.worker.cpu.kv_connector as kv_connector_module
+from vllm.config import VllmConfig
+from vllm.distributed.kv_transfer.kv_connector.v1.base import (
+    KVConnectorTransferResults,
+)
+from vllm.platforms import CpuArchEnum, current_platform
+from vllm.triton_utils import HAS_TRITON
+from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.worker.cpu.input_batch import InputBatch, InputBuffers
+from vllm.v1.worker.cpu.kv_connector import ActiveKVConnector
+
+pytestmark = pytest.mark.skipif(
+    not current_platform.is_cpu()
+    or HAS_TRITON
+    or current_platform.get_cpu_architecture() == CpuArchEnum.X86,
+    reason="requires a non-x86 CPU platform without Triton",
+)
+
+
+class FakeKVTransferBackend:
+    def __init__(self, events: list[str]):
+        self.events = events
+        self.load_calls: list[tuple[object, dict[str, Any]]] = []
+
+    def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]) -> None:
+        self.kv_caches = kv_caches
+
+    def set_host_xfer_buffer_ops(self, copy_blocks_op: object) -> None:
+        self.copy_blocks_op = copy_blocks_op
+
+    def handle_preemptions(self, metadata: object) -> None:
+        self.events.append("handle_preemptions")
+
+    def bind_connector_metadata(self, metadata: object) -> None:
+        self.events.append("bind_metadata")
+
+    def start_load_kv(self, context: object, **kwargs: Any) -> None:
+        self.events.append("start_load")
+        self.load_calls.append((context, kwargs))
+
+    def finish_forward(self) -> None:
+        self.events.append("finish_forward")
+
+    def wait_for_save(self) -> None:
+        self.events.append("wait_for_save")
+
+    def get_transfer_results(self, finished_req_ids: set[str]):
+        assert finished_req_ids == {"finished"}
+        return KVConnectorTransferResults(
+            finished_sending={"sent"},
+            finished_recving={"received"},
+            failed_recving={"failed"},
+        )
+
+    def get_block_ids_with_load_errors(self) -> set[int]:
+        return {7}
+
+    def get_kv_connector_stats(self) -> None:
+        return None
+
+    def get_kv_connector_kv_cache_events(self) -> None:
+        return None
+
+    def build_connector_worker_meta(self) -> None:
+        return None
+
+    def clear_connector_metadata(self) -> None:
+        self.events.append("clear_metadata")
+
+
+def _make_connector(
+    monkeypatch: pytest.MonkeyPatch, events: list[str]
+) -> tuple[ActiveKVConnector, FakeKVTransferBackend, object]:
+    backend = FakeKVTransferBackend(events)
+    forward_context = object()
+    monkeypatch.setattr(kv_connector_module, "get_kv_transfer_group", lambda: backend)
+    monkeypatch.setattr(
+        kv_connector_module, "is_forward_context_available", lambda: True
+    )
+    monkeypatch.setattr(
+        kv_connector_module, "get_forward_context", lambda: forward_context
+    )
+    connector = ActiveKVConnector(VllmConfig(), {"layer": torch.empty(1)})
+    return connector, backend, forward_context
+
+
+def _scheduler_output(has_sync_kv_loads: bool) -> SchedulerOutput:
+    output = SchedulerOutput.make_empty()
+    output.kv_connector_metadata = object()
+    output.has_sync_kv_loads = has_sync_kv_loads
+    output.finished_req_ids = {"finished"}
+    return output
+
+
+def _input_batch() -> InputBatch:
+    buffers = InputBuffers(4, 5, torch.device("cpu"))
+    batch = InputBatch.make_dummy(2, 5, buffers)
+    batch.req_ids = ["first", "second"]
+    batch.idx_mapping = torch.tensor([3, 1])
+    batch.idx_mapping_np = np.array([3, 1], dtype=np.intp)
+    return batch
+
+
+@pytest.mark.parametrize("sync_load", [False, True])
+def test_cpu_kv_connector_preserves_transfer_phase_order(
+    monkeypatch: pytest.MonkeyPatch, sync_load: bool
+) -> None:
+    events: list[str] = []
+    connector, backend, forward_context = _make_connector(monkeypatch, events)
+    scheduler_output = _scheduler_output(sync_load)
+    input_batch = _input_batch()
+    attn_metadata = {"layer": object()}
+
+    connector.pre_forward(scheduler_output, input_batch, attn_metadata)
+    connector.finish_forward()
+    output = connector.post_forward({"finished"})
+
+    expected = ["handle_preemptions", "bind_metadata"]
+    if sync_load:
+        expected.append("start_load")
+    expected.append("finish_forward")
+    if not sync_load:
+        expected.append("start_load")
+    expected.extend(("wait_for_save", "clear_metadata"))
+    assert events == expected
+
+    assert len(backend.load_calls) == 1
+    context, load_kwargs = backend.load_calls[0]
+    assert context is forward_context
+    assert load_kwargs == {
+        "scheduler_output": scheduler_output,
+        "request_state_indices": input_batch.idx_mapping,
+        "request_ids": input_batch.req_ids,
+        "num_tokens": 5,
+        "attn_metadata": attn_metadata,
+    }
+    assert output is not None
+    assert output.finished_sending == {"sent"}
+    assert output.finished_recving == {"received"}
+    assert output.failed_recving == {"failed"}
+    assert output.invalid_block_ids == {7}
+
+    connector.no_forward(_scheduler_output(False))
+    assert backend.load_calls[-1] == (forward_context, {})
+
+
+def test_cpu_kv_connector_empty_step_does_not_reuse_batch_inputs(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    events: list[str] = []
+    connector, backend, forward_context = _make_connector(monkeypatch, events)
+    scheduler_output = _scheduler_output(False)
+
+    output = connector.no_forward(scheduler_output)
+
+    assert events == [
+        "handle_preemptions",
+        "bind_metadata",
+        "finish_forward",
+        "start_load",
+        "wait_for_save",
+        "clear_metadata",
+    ]
+    assert backend.load_calls == [(forward_context, {})]
+    assert output.kv_connector_output is not None
+    assert output.kv_connector_output.finished_sending == {"sent"}
diff --git a/tests/v1/worker/cpu/test_cpu_logit_bias.py b/tests/v1/worker/cpu/test_cpu_logit_bias.py
new file mode 100644
index 000000000000..2475fb161903
--- /dev/null
+++ b/tests/v1/worker/cpu/test_cpu_logit_bias.py
@@ -0,0 +1,169 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import numpy as np
+import pytest
+import torch
+
+from vllm.platforms import CpuArchEnum, current_platform
+from vllm.sampling_params import SamplingParams, StructuredOutputsParams
+from vllm.triton_utils import HAS_TRITON
+from vllm.v1.worker.cpu.logits_processors import (
+    CPULogitBiasState,
+    LogitsContext,
+    LogitsProcRequestState,
+)
+from vllm.v1.worker.cpu.states import CPURequestState
+
+if current_platform.is_cpu():
+    import vllm._C  # noqa: F401
+
+pytestmark = pytest.mark.skipif(
+    not current_platform.is_cpu()
+    or HAS_TRITON
+    or current_platform.get_cpu_architecture() == CpuArchEnum.X86,
+    reason="requires a non-x86 CPU platform without Triton",
+)
+
+VOCAB_SIZE = 128
+STOP_TOKEN = 3
+OTHER_TOKEN = 4
+PROMPT_LEN = 2
+MIN_TOKENS = 10
+POS = 5
+
+
+def _make_state(
+    params: list[SamplingParams],
+) -> tuple[CPULogitBiasState, list[int]]:
+    req_states = CPURequestState(
+        max_num_reqs=4,
+        max_model_len=64,
+        max_num_batched_tokens=16,
+        vocab_size=VOCAB_SIZE,
+        device=torch.device("cpu"),
+    )
+    state = CPULogitBiasState(
+        None, LogitsProcRequestState.from_request_state(req_states)
+    )
+    req_indices = []
+    for req_idx, sampling_params in enumerate(params):
+        req_id = f"req-{req_idx}"
+        req_states.add_request(
+            req_id,
+            prompt_len=PROMPT_LEN,
+            all_token_ids=[10, 11],
+            num_computed_tokens=PROMPT_LEN,
+            max_tokens=16,
+        )
+        state_idx = req_states.req_id_to_index[req_id]
+        req_indices.append(state_idx)
+        state.add_request(state_idx, sampling_params)
+    return state, req_indices
+
+
+def _apply(
+    logits: torch.Tensor,
+    params: list[SamplingParams],
+    idx_mapping: list[int] | None = None,
+    positions: list[int] | None = None,
+) -> torch.Tensor:
+    state, req_indices = _make_state(params)
+    if idx_mapping is None:
+        idx_mapping = list(range(len(params)))
+    if positions is None:
+        positions = [POS] * len(idx_mapping)
+
+    state_mapping = [req_indices[idx] for idx in idx_mapping]
+    idx_mapping_tensor = torch.tensor(state_mapping, dtype=torch.int64)
+    num_tokens = len(idx_mapping)
+    state.apply(
+        logits,
+        LogitsContext(
+            expanded_idx_mapping=idx_mapping_tensor,
+            idx_mapping=idx_mapping_tensor,
+            idx_mapping_np=np.array(state_mapping, dtype=np.intp),
+            expanded_local_pos=torch.zeros(num_tokens, dtype=torch.int32),
+            input_ids=torch.zeros(num_tokens, dtype=torch.int32),
+            pos=torch.tensor(positions, dtype=torch.int64),
+            seq_lens_upper_bound_np=np.ones(num_tokens, dtype=np.int64),
+        ),
+    )
+    return logits
+
+
+def _min_tokens_params(structured: bool) -> SamplingParams:
+    return SamplingParams(
+        min_tokens=MIN_TOKENS,
+        stop_token_ids=[STOP_TOKEN],
+        structured_outputs=(
+            StructuredOutputsParams(json={"type": "boolean"}) if structured else None
+        ),
+    )
+
+
+def _only_stop_token_left(num_rows: int) -> torch.Tensor:
+    logits = torch.full((num_rows, VOCAB_SIZE), -float("inf"))
+    logits[:, STOP_TOKEN] = 1.0
+    return logits
+
+
+def test_cpu_allowed_tokens_and_additive_bias_follow_request_mapping() -> None:
+    logits = torch.arange(24, dtype=torch.float32).view(3, 8)
+    params = [
+        SamplingParams(logit_bias={1: 2.0, 4: -1.0}),
+        SamplingParams(),
+        SamplingParams(
+            allowed_token_ids=[0, 3, 5],
+            logit_bias={3: 0.5, 4: 100.0},
+        ),
+    ]
+
+    actual = _apply(logits, params, idx_mapping=[2, 0, 2])
+    expected = torch.tensor(
+        [
+            [0.0, -torch.inf, -torch.inf, 3.5, -torch.inf, 5.0, -torch.inf, -torch.inf],
+            [8.0, 11.0, 10.0, 11.0, 11.0, 13.0, 14.0, 15.0],
+            [
+                16.0,
+                -torch.inf,
+                -torch.inf,
+                19.5,
+                -torch.inf,
+                21.0,
+                -torch.inf,
+                -torch.inf,
+            ],
+        ]
+    )
+    torch.testing.assert_close(actual, expected)
+
+
+def test_cpu_min_tokens_restores_terminal_structured_output() -> None:
+    actual = _apply(_only_stop_token_left(1), [_min_tokens_params(True)])
+    assert actual[0, STOP_TOKEN] == 1.0
+
+
+def test_cpu_min_tokens_keeps_empty_unstructured_row() -> None:
+    actual = _apply(_only_stop_token_left(1), [_min_tokens_params(False)])
+    assert torch.isneginf(actual).all()
+
+
+def test_cpu_min_tokens_keeps_masking_when_another_token_remains() -> None:
+    logits = _only_stop_token_left(1)
+    logits[0, OTHER_TOKEN] = 2.0
+
+    actual = _apply(logits, [_min_tokens_params(True)])
+
+    assert torch.isneginf(actual[0, STOP_TOKEN])
+    assert actual[0, OTHER_TOKEN] == 2.0
+
+
+def test_cpu_min_tokens_restores_only_structured_rows_in_mixed_batch() -> None:
+    actual = _apply(
+        _only_stop_token_left(2),
+        [_min_tokens_params(True), _min_tokens_params(False)],
+    )
+
+    assert actual[0, STOP_TOKEN] == 1.0
+    assert torch.isneginf(actual[1]).all()
diff --git a/tests/v1/worker/cpu/test_cpu_logits_processors.py b/tests/v1/worker/cpu/test_cpu_logits_processors.py
new file mode 100644
index 000000000000..a6edb7a786ac
--- /dev/null
+++ b/tests/v1/worker/cpu/test_cpu_logits_processors.py
@@ -0,0 +1,226 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from collections import Counter
+
+import numpy as np
+import pytest
+import torch
+
+from vllm.platforms import CpuArchEnum, current_platform
+from vllm.sampling_params import SamplingParams
+from vllm.triton_utils import HAS_TRITON
+from vllm.v1.worker.cpu.logits_processors import (
+    CPUBadWords,
+    CPUPenalties,
+    LogitsContext,
+    LogitsProcRequestState,
+)
+from vllm.v1.worker.cpu.states import CPURequestState
+
+if current_platform.is_cpu():
+    import vllm._C  # noqa: F401
+
+pytestmark = pytest.mark.skipif(
+    not current_platform.is_cpu()
+    or HAS_TRITON
+    or current_platform.get_cpu_architecture() == CpuArchEnum.X86,
+    reason="requires a non-x86 CPU platform without Triton",
+)
+
+VOCAB_SIZE = 128
+PROMPT = [5, 6, 5]
+COMMITTED_OUTPUT = [7, 7, 8]
+
+
+def _make_req_states() -> tuple[CPURequestState, int]:
+    req_states = CPURequestState(
+        max_num_reqs=4,
+        max_model_len=64,
+        max_num_batched_tokens=16,
+        vocab_size=VOCAB_SIZE,
+        device=torch.device("cpu"),
+    )
+    req_states.add_request(
+        req_id="req",
+        prompt_len=len(PROMPT),
+        all_token_ids=PROMPT + COMMITTED_OUTPUT,
+        num_computed_tokens=len(PROMPT) + len(COMMITTED_OUTPUT),
+        max_tokens=32,
+    )
+    return req_states, req_states.req_id_to_index["req"]
+
+
+def _context(req_idx: int, input_id: int = COMMITTED_OUTPUT[-1]) -> LogitsContext:
+    return LogitsContext(
+        expanded_idx_mapping=torch.tensor([req_idx], dtype=torch.int64),
+        idx_mapping=torch.tensor([req_idx], dtype=torch.int64),
+        idx_mapping_np=np.array([req_idx], dtype=np.intp),
+        expanded_local_pos=torch.zeros(1, dtype=torch.int32),
+        input_ids=torch.tensor([input_id], dtype=torch.int32),
+        pos=torch.zeros(1, dtype=torch.int32),
+        seq_lens_upper_bound_np=np.ones(1, dtype=np.int64),
+    )
+
+
+def _reference_penalties(
+    logits: torch.Tensor,
+    prompt: list[int],
+    output: list[int],
+    repetition: float,
+    frequency: float,
+    presence: float,
+) -> None:
+    for token_id in set(prompt + output):
+        value = logits[token_id]
+        logits[token_id] = value / repetition if value > 0 else value * repetition
+    for token_id, count in Counter(output).items():
+        logits[token_id] -= frequency * count + presence
+
+
+def test_cpu_penalties_match_reference() -> None:
+    req_states, req_idx = _make_req_states()
+    state = CPUPenalties(None, LogitsProcRequestState.from_request_state(req_states))
+    params = SamplingParams(
+        repetition_penalty=1.7,
+        frequency_penalty=0.4,
+        presence_penalty=-0.25,
+    )
+    state.add_request(req_idx, params)
+
+    logits = torch.linspace(-2, 2, VOCAB_SIZE).unsqueeze(0)
+    logits[:, 5] = 2.0
+    logits[:, 6] = -2.0
+    logits[:, 7] = 3.0
+    logits[:, 8] = -3.0
+    expected = logits.clone()
+    _reference_penalties(
+        expected[0],
+        PROMPT,
+        COMMITTED_OUTPUT,
+        params.repetition_penalty,
+        params.frequency_penalty,
+        params.presence_penalty,
+    )
+
+    actual = state.apply(logits, _context(req_idx))
+    torch.testing.assert_close(actual, expected)
+
+
+@pytest.mark.parametrize("placeholder_request", ["first", "second"])
+def test_cpu_penalties_keep_request_rows_independent(
+    placeholder_request: str,
+) -> None:
+    req_states = CPURequestState(
+        max_num_reqs=4,
+        max_model_len=64,
+        max_num_batched_tokens=16,
+        vocab_size=VOCAB_SIZE,
+        device=torch.device("cpu"),
+    )
+    requests = {
+        "first": ([1, 2], [3, 3]),
+        "second": ([4], [5]),
+    }
+    for req_id, (prompt, output) in requests.items():
+        if req_id == placeholder_request:
+            prompt = prompt + [VOCAB_SIZE]
+        tokens = prompt + output
+        req_states.add_request(
+            req_id=req_id,
+            prompt_len=len(prompt),
+            all_token_ids=tokens,
+            num_computed_tokens=len(tokens),
+            max_tokens=32,
+        )
+    first_idx = req_states.req_id_to_index["first"]
+    second_idx = req_states.req_id_to_index["second"]
+    state = CPUPenalties(None, LogitsProcRequestState.from_request_state(req_states))
+    first_params = SamplingParams(
+        repetition_penalty=1.5,
+        frequency_penalty=0.25,
+        presence_penalty=0.1,
+    )
+    second_params = SamplingParams(
+        repetition_penalty=1.2,
+        frequency_penalty=-0.5,
+        presence_penalty=0.3,
+    )
+    state.add_request(first_idx, first_params)
+    state.add_request(second_idx, second_params)
+
+    # Reverse request order so row number and persistent request slot differ.
+    row_mapping = torch.tensor([second_idx, first_idx], dtype=torch.int64)
+    ctx = LogitsContext(
+        expanded_idx_mapping=row_mapping,
+        idx_mapping=row_mapping,
+        idx_mapping_np=row_mapping.numpy(),
+        expanded_local_pos=torch.zeros(2, dtype=torch.int32),
+        input_ids=torch.tensor([5, 3], dtype=torch.int32),
+        pos=torch.zeros(2, dtype=torch.int32),
+        seq_lens_upper_bound_np=np.array([2, 4], dtype=np.int64),
+    )
+    logits = torch.linspace(-2, 2, VOCAB_SIZE).repeat(2, 1)
+    expected = logits.clone()
+    _reference_penalties(
+        expected[0],
+        *requests["second"],
+        second_params.repetition_penalty,
+        second_params.frequency_penalty,
+        second_params.presence_penalty,
+    )
+    _reference_penalties(
+        expected[1],
+        *requests["first"],
+        first_params.repetition_penalty,
+        first_params.frequency_penalty,
+        first_params.presence_penalty,
+    )
+
+    actual = state.apply(logits, ctx)
+    torch.testing.assert_close(actual, expected)
+
+
+def test_cpu_penalties_clear_reused_slot() -> None:
+    req_states, req_idx = _make_req_states()
+    state = CPUPenalties(None, LogitsProcRequestState.from_request_state(req_states))
+    state.add_request(req_idx, SamplingParams(repetition_penalty=1.2))
+    state.add_request(req_idx, SamplingParams())
+
+    logits = torch.randn(1, VOCAB_SIZE)
+    expected = logits.clone()
+    state.apply(logits, _context(req_idx))
+    torch.testing.assert_close(logits, expected)
+
+
+@pytest.mark.parametrize(
+    ("bad_words", "masked"),
+    [
+        ([[7, 8, 60]], {60}),
+        ([[8, 8, 50]], set()),
+        ([[70]], {70}),
+    ],
+)
+def test_cpu_bad_words_match_committed_output_prefix(
+    bad_words: list[list[int]], masked: set[int]
+) -> None:
+    req_states, req_idx = _make_req_states()
+    state = CPUBadWords(None, LogitsProcRequestState.from_request_state(req_states))
+    state.add_request(req_idx, SamplingParams(_bad_words_token_ids=bad_words))
+
+    actual = state.apply(torch.zeros((1, VOCAB_SIZE)), _context(req_idx))
+    expected = torch.zeros_like(actual)
+    for token_id in masked:
+        expected[0, token_id] = -torch.inf
+    torch.testing.assert_close(actual, expected)
+
+
+def test_cpu_bad_words_clear_reused_slot() -> None:
+    req_states, req_idx = _make_req_states()
+    state = CPUBadWords(None, LogitsProcRequestState.from_request_state(req_states))
+    state.add_request(req_idx, SamplingParams(_bad_words_token_ids=[[7, 8, 60]]))
+    state.add_request(req_idx, SamplingParams())
+
+    logits = torch.zeros((1, VOCAB_SIZE))
+    state.apply(logits, _context(req_idx))
+    assert not torch.isinf(logits).any()
diff --git a/tests/v1/worker/cpu/test_cpu_logprob.py b/tests/v1/worker/cpu/test_cpu_logprob.py
new file mode 100644
index 000000000000..67827983a72d
--- /dev/null
+++ b/tests/v1/worker/cpu/test_cpu_logprob.py
@@ -0,0 +1,313 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import math
+
+import numpy as np
+import pytest
+import torch
+
+from vllm.platforms import CpuArchEnum, current_platform
+from vllm.sampling_params import SamplingParams
+from vllm.triton_utils import HAS_TRITON
+from vllm.v1.worker.cpu.input_batch import InputBatch, InputBuffers
+from vllm.v1.worker.cpu.logprob import (
+    CPUPromptLogprobsWorker,
+    LogprobTokenIdsState,
+    compute_prompt_logprobs_with_chunking,
+    compute_token_logprobs,
+    compute_topk_scores,
+    get_prompt_logprobs_token_ids,
+)
+
+if current_platform.is_cpu():
+    import vllm._C  # noqa: F401
+
+pytestmark = pytest.mark.skipif(
+    not current_platform.is_cpu()
+    or HAS_TRITON
+    or current_platform.get_cpu_architecture() == CpuArchEnum.X86,
+    reason="requires a non-x86 CPU platform without Triton",
+)
+
+
+@pytest.mark.parametrize("logits_mode", [False, True])
+def test_cpu_output_logprobs(logits_mode: bool) -> None:
+    logits = torch.tensor(
+        [[4.0, 2.0, 1.0, 0.0], [0.0, 2.0, 1.0, 3.0]],
+        dtype=torch.bfloat16,
+    )
+    sampled_token_ids = torch.tensor([2, 1], dtype=torch.int64)
+
+    output = compute_topk_scores(
+        logits,
+        num_logprobs=2,
+        sampled_token_ids=sampled_token_ids,
+        cu_num_logits=[0, 1, 2],
+        max_per_req_token_ids=0,
+        logits_mode=logits_mode,
+    )
+
+    assert output.logprob_token_ids.tolist() == [
+        [2, 0, 1],
+        [1, 3, 1],
+    ]
+    assert output.logprob_token_ids.dtype == torch.int64
+    expected_scores = torch.tensor([[1.0, 4.0, 2.0], [2.0, 3.0, 2.0]])
+    if not logits_mode:
+        expected_scores -= torch.logsumexp(logits.float(), dim=-1, keepdim=True)
+    torch.testing.assert_close(output.logprobs, expected_scores)
+    assert output.selected_token_ranks.tolist() == [3, 2]
+    assert output.cu_num_generated_tokens == [0, 1, 2]
+
+
+def test_cpu_prompt_logprobs_scores_prompt_targets() -> None:
+    logits = torch.tensor(
+        [[4.0, 2.0, 1.0, 0.0], [0.0, 2.0, 1.0, 3.0]],
+        dtype=torch.bfloat16,
+    )
+
+    prompt_token_ids = get_prompt_logprobs_token_ids(
+        num_tokens=2,
+        query_start_loc=torch.tensor([0, 2], dtype=torch.int32),
+        idx_mapping=torch.tensor([0], dtype=torch.int64),
+        num_computed_tokens=torch.tensor([0], dtype=torch.int32),
+        all_token_ids=torch.tensor([[7, 1, 2]], dtype=torch.int32),
+    )
+    token_ids, logprobs, token_ranks = compute_prompt_logprobs_with_chunking(
+        prompt_token_ids=prompt_token_ids,
+        prompt_hidden_states=logits,
+        logits_fn=lambda hidden_states: hidden_states,
+        num_prompt_logprobs=2,
+    )
+
+    assert token_ids.tolist() == [[1, 0, 1], [2, 3, 1]]
+    expected = torch.tensor([[2.0, 4.0, 2.0], [1.0, 3.0, 2.0]])
+    expected -= torch.logsumexp(logits.float(), dim=-1, keepdim=True)
+    torch.testing.assert_close(logprobs, expected)
+    assert token_ranks.tolist() == [2, 3]
+
+
+def _input_batch(
+    req_id: str,
+    num_computed_tokens: int,
+    num_scheduled_tokens: int,
+    prefill_len: int,
+) -> InputBatch:
+    buffers = InputBuffers(1, num_scheduled_tokens, torch.device("cpu"))
+    batch = InputBatch.make_dummy(1, num_scheduled_tokens, buffers)
+    batch.req_ids = [req_id]
+    batch.num_computed_prefill_tokens_np[:] = num_computed_tokens
+    batch.num_scheduled_tokens[:] = num_scheduled_tokens
+    batch.prefill_len_np[:] = prefill_len
+    return batch
+
+
+def test_cpu_prompt_logprobs_accumulate_chunked_prefill() -> None:
+    worker = CPUPromptLogprobsWorker(max_num_reqs=1, device=torch.device("cpu"))
+    worker.add_request("request", 0, SamplingParams(prompt_logprobs=1))
+
+    logits = torch.tensor(
+        [
+            [4.0, 2.0, 1.0, 0.0],
+            [0.0, 2.0, 1.0, 3.0],
+            [1.0, 4.0, 3.0, 2.0],
+            [3.0, 0.0, 2.0, 1.0],
+        ]
+    )
+    all_token_ids = torch.tensor([[0, 1, 2, 3, 0]], dtype=torch.int32)
+    prompt_lens = np.array([4], dtype=np.int32)
+
+    first = worker.compute_prompt_logprobs(
+        lambda hidden_states: hidden_states,
+        logits[:2],
+        _input_batch("request", 0, 2, 4),
+        all_token_ids,
+        torch.tensor([0], dtype=torch.int32),
+        prompt_lens,
+    )
+    assert first == {}
+
+    second = worker.compute_prompt_logprobs(
+        lambda hidden_states: hidden_states,
+        logits[2:],
+        _input_batch("request", 2, 2, 4),
+        all_token_ids,
+        torch.tensor([2], dtype=torch.int32),
+        prompt_lens,
+    )
+
+    result = second["request"]
+    assert result.logprob_token_ids[:, 0].tolist() == [1, 2, 3]
+    expected = torch.log_softmax(logits[:3], dim=-1)[
+        torch.arange(3), torch.tensor([1, 2, 3])
+    ]
+    torch.testing.assert_close(result.logprobs[:, 0], expected)
+
+
+def test_cpu_token_logprobs_accept_int32_token_ids() -> None:
+    logits = torch.tensor(
+        [[4.0, 2.0, 1.0, 0.0], [0.0, 2.0, 1.0, 3.0]],
+        dtype=torch.bfloat16,
+    )
+    token_ids = torch.tensor([[1, 3], [2, 0]], dtype=torch.int32)
+
+    actual = compute_token_logprobs(logits, token_ids)
+    expected = logits.gather(-1, token_ids).float()
+    expected -= torch.logsumexp(logits.float(), dim=-1, keepdim=True)
+
+    torch.testing.assert_close(actual, expected)
+
+
+def test_cpu_prompt_logprobs_honor_start_position() -> None:
+    worker = CPUPromptLogprobsWorker(max_num_reqs=1, device=torch.device("cpu"))
+    worker.add_request(
+        "request",
+        0,
+        SamplingParams(prompt_logprobs=1, prompt_logprob_start=2),
+    )
+    logits = torch.tensor(
+        [
+            [4.0, 2.0, 1.0, 0.0],
+            [0.0, 2.0, 1.0, 3.0],
+            [1.0, 4.0, 3.0, 2.0],
+            [3.0, 0.0, 2.0, 1.0],
+        ]
+    )
+    all_token_ids = torch.tensor([[0, 1, 2, 3, 0]], dtype=torch.int32)
+
+    result = worker.compute_prompt_logprobs(
+        lambda hidden_states: hidden_states,
+        logits,
+        _input_batch("request", 0, 4, 4),
+        all_token_ids,
+        torch.tensor([0], dtype=torch.int32),
+        np.array([4], dtype=np.int32),
+    )["request"]
+
+    assert result.logprob_token_ids[:, 0].tolist() == [3]
+
+
+def test_cpu_requested_prompt_token_scores_honor_start_position() -> None:
+    worker = CPUPromptLogprobsWorker(max_num_reqs=1, device=torch.device("cpu"))
+    token_ids = np.array([[0, 2], [1, 3]], dtype=np.int64)
+    worker.add_request(
+        "request",
+        0,
+        SamplingParams(
+            prompt_logprob_token_ids=token_ids,
+            prompt_logprob_start=1,
+        ),
+    )
+    logits = torch.tensor(
+        [
+            [4.0, 2.0, 1.0, 0.0],
+            [0.0, 2.0, 1.0, 3.0],
+            [1.0, 4.0, 3.0, 2.0],
+            [3.0, 0.0, 2.0, 1.0],
+        ]
+    )
+
+    result = worker.compute_prompt_token_id_logprobs(
+        lambda hidden_states: hidden_states,
+        logits,
+        _input_batch("request", 0, 4, 4),
+        np.array([4], dtype=np.int32),
+    )["request"]
+    expected = compute_token_logprobs(logits[1:3], torch.from_numpy(token_ids))
+
+    torch.testing.assert_close(result, expected)
+
+
+def _make_logprob_token_ids_state(
+    num_token_ids: list[int], token_ids: list[list[int]]
+) -> LogprobTokenIdsState:
+    state = LogprobTokenIdsState(len(num_token_ids), torch.device("cpu"))
+    for req_idx, num_ids in enumerate(num_token_ids):
+        state.add_request(
+            req_idx,
+            SamplingParams(logprob_token_ids=token_ids[req_idx][:num_ids]),
+        )
+    return state
+
+
+@pytest.mark.parametrize("logits_mode", [False, True])
+def test_cpu_requested_output_logprobs(logits_mode: bool) -> None:
+    logits = torch.tensor(
+        [
+            [6.0, 5.0, 4.0, 3.0, 2.0, 1.0],
+            [0.0, 2.0, 1.0, 5.0, 4.0, 3.0],
+            [1.0, 6.0, 2.0, 5.0, 3.0, 4.0],
+            [4.0, 1.0, 6.0, 0.0, 5.0, 2.0],
+        ],
+        dtype=torch.bfloat16,
+    )
+    sampled_token_ids = torch.tensor([4, 1, 0, 5], dtype=torch.int64)
+    state = _make_logprob_token_ids_state(
+        num_token_ids=[0, 5, 2, 3],
+        token_ids=[
+            [0, 0, 0, 0, 0],
+            [0, 1, 2, 3, 4],
+            [5, 2, 0, 0, 0],
+            [4, 1, 3, 0, 0],
+        ],
+    )
+
+    output = compute_topk_scores(
+        logits,
+        num_logprobs=4,
+        sampled_token_ids=sampled_token_ids,
+        cu_num_logits=[0, 1, 2, 3, 4],
+        num_per_req_token_ids=state.num_token_ids,
+        per_req_token_ids=state.token_ids,
+        expanded_idx_mapping=torch.tensor([2, 0, 3, 1], dtype=torch.int64),
+        max_per_req_token_ids=5,
+        logits_mode=logits_mode,
+    )
+
+    assert output.logprob_token_ids.tolist() == [
+        [4, 5, 2, 0, 0, 0],
+        [1, 3, 4, 5, 1, 0],
+        [0, 4, 1, 3, 0, 0],
+        [5, 0, 1, 2, 3, 4],
+    ]
+    expected_scores = torch.tensor(
+        [
+            [2.0, 1.0, 4.0, -math.inf, -math.inf, -math.inf],
+            [2.0, 5.0, 4.0, 3.0, 2.0, -math.inf],
+            [1.0, 3.0, 6.0, 5.0, -math.inf, -math.inf],
+            [2.0, 4.0, 1.0, 6.0, 0.0, 5.0],
+        ]
+    )
+    if not logits_mode:
+        expected_scores -= torch.logsumexp(logits.float(), dim=-1, keepdim=True)
+    torch.testing.assert_close(output.logprobs, expected_scores)
+    assert output.selected_token_ranks.tolist() == [5, 4, 6, 4]
+    assert output.cu_num_generated_tokens == [0, 1, 2, 3, 4]
+
+
+def test_cpu_requested_output_logprobs_without_topk() -> None:
+    logits = torch.tensor([[4.0, 3.0, 2.0, 1.0], [1.0, 2.0, 3.0, 4.0]])
+    state = _make_logprob_token_ids_state(
+        num_token_ids=[0, 2],
+        token_ids=[[0, 0], [3, 1]],
+    )
+
+    output = compute_topk_scores(
+        logits,
+        num_logprobs=0,
+        sampled_token_ids=torch.tensor([2, 1]),
+        cu_num_logits=None,
+        num_per_req_token_ids=state.num_token_ids,
+        per_req_token_ids=state.token_ids,
+        expanded_idx_mapping=torch.tensor([1, 0], dtype=torch.int64),
+        max_per_req_token_ids=2,
+        logits_mode=True,
+    )
+
+    assert output.logprob_token_ids.tolist() == [[2, 3, 1], [1, 0, 0]]
+    torch.testing.assert_close(
+        output.logprobs,
+        torch.tensor([[2.0, 1.0, 3.0], [2.0, -math.inf, -math.inf]]),
+    )
+    assert output.selected_token_ranks.tolist() == [3, 3]
diff --git a/tests/v1/worker/cpu/test_cpu_lora_utils.py b/tests/v1/worker/cpu/test_cpu_lora_utils.py
new file mode 100644
index 000000000000..3a28565b427d
--- /dev/null
+++ b/tests/v1/worker/cpu/test_cpu_lora_utils.py
@@ -0,0 +1,37 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import numpy as np
+import pytest
+
+from vllm.lora.request import LoRARequest
+from vllm.platforms import CpuArchEnum, current_platform
+from vllm.triton_utils import HAS_TRITON
+from vllm.v1.worker.cpu.lora_utils import LoraState
+
+pytestmark = pytest.mark.skipif(
+    not current_platform.is_cpu()
+    or HAS_TRITON
+    or current_platform.get_cpu_architecture() == CpuArchEnum.X86,
+    reason="requires a non-x86 CPU platform without Triton",
+)
+
+
+def test_cpu_lora_state_maps_each_logits_row_to_its_request() -> None:
+    state = LoraState(max_num_reqs=4)
+    first = LoRARequest("first", 3, "/tmp/first")
+    second = LoRARequest("second", 7, "/tmp/second")
+    state.add_request("first-request", 1, first)
+    state.add_request("base-request", 2, None)
+    state.add_request("second-request", 3, second)
+
+    prompt_mapping, token_mapping, active = state.make_lora_inputs(
+        req_ids=["second-request", "base-request"],
+        idx_mapping=np.array([3, 2], dtype=np.intp),
+        num_scheduled_tokens=np.array([2, 3], dtype=np.int32),
+        num_logits=np.array([1, 1], dtype=np.int32),
+    )
+
+    assert prompt_mapping == (7, 0)
+    assert token_mapping == (7, 7, 0, 0, 0)
+    assert active == {second}
diff --git a/tests/v1/worker/cpu/test_cpu_pp_utils.py b/tests/v1/worker/cpu/test_cpu_pp_utils.py
new file mode 100644
index 000000000000..aef6b3e90801
--- /dev/null
+++ b/tests/v1/worker/cpu/test_cpu_pp_utils.py
@@ -0,0 +1,174 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import os
+
+import numpy as np
+import pytest
+import torch
+
+from vllm.config import VllmConfig, set_current_vllm_config
+from vllm.distributed.parallel_state import (
+    destroy_distributed_environment,
+    destroy_model_parallel,
+    init_distributed_environment,
+    initialize_model_parallel,
+)
+from vllm.platforms import CpuArchEnum, current_platform
+from vllm.triton_utils import HAS_TRITON
+from vllm.utils.network_utils import get_open_port
+from vllm.v1.worker.cpu.input_batch import InputBatch, InputBuffers
+from vllm.v1.worker.cpu.pp_utils import PPHandler, compute_need_sampled_mask
+
+pytestmark = pytest.mark.skipif(
+    not current_platform.is_cpu()
+    or HAS_TRITON
+    or current_platform.get_cpu_architecture() == CpuArchEnum.X86,
+    reason="requires a non-x86 CPU platform without Triton",
+)
+
+
+def _batch(
+    *,
+    num_computed: list[int],
+    prefill_len: list[int],
+    num_scheduled: list[int],
+    is_prefilling: list[bool],
+    max_seq_len: list[int] | None,
+    idx_mapping: list[int] | None = None,
+) -> InputBatch:
+    num_reqs = len(num_computed)
+    idx_mapping = idx_mapping or list(range(num_reqs))
+    idx_mapping_np = np.array(idx_mapping, dtype=np.int64)
+    num_tokens = sum(num_scheduled)
+    batch = InputBatch.make_dummy(
+        num_reqs,
+        num_tokens,
+        InputBuffers(num_reqs, num_tokens, torch.device("cpu")),
+    )
+    batch.idx_mapping = torch.from_numpy(idx_mapping_np)
+    batch.idx_mapping_np = idx_mapping_np
+    batch.num_computed_tokens_np = np.array(num_computed, dtype=np.int32)
+    batch.prefill_len_np = np.array(prefill_len, dtype=np.int32)
+    batch.num_scheduled_tokens = np.array(num_scheduled, dtype=np.int32)
+    batch.is_prefilling_np = np.array(is_prefilling, dtype=np.bool_)
+    batch.max_seq_len_np = (
+        None if max_seq_len is None else np.array(max_seq_len, dtype=np.int32)
+    )
+    return batch
+
+
+@pytest.mark.parametrize(
+    (
+        "num_computed",
+        "prefill_len",
+        "num_scheduled",
+        "is_prefilling",
+        "max_seq_len",
+        "expected",
+    ),
+    [
+        ([512, 1000], [4096, 1004], [448, 4], [True, True], None, [False, True]),
+        ([0, 512], [4096, 4096], [448, 448], [True, True], None, None),
+        ([10, 512], [8, 4096], [1, 448], [False, True], None, [True, False]),
+        ([1000], [1004], [4], [True], [1005], None),
+        ([1000], [1004], [4], [True], [1260], [True]),
+        (
+            [1000, 10],
+            [1004, 8],
+            [4, 1],
+            [True, False],
+            [1005, 1024],
+            [False, True],
+        ),
+    ],
+)
+def test_cpu_pp_sample_mask_matches_supported_gpu_scenarios(
+    num_computed: list[int],
+    prefill_len: list[int],
+    num_scheduled: list[int],
+    is_prefilling: list[bool],
+    max_seq_len: list[int] | None,
+    expected: list[bool] | None,
+) -> None:
+    batch = _batch(
+        num_computed=num_computed,
+        prefill_len=prefill_len,
+        num_scheduled=num_scheduled,
+        is_prefilling=is_prefilling,
+        max_seq_len=max_seq_len,
+    )
+
+    mask = compute_need_sampled_mask(batch)
+
+    assert (None if mask is None else mask.tolist()) == expected
+
+
+def _run_cpu_pp_sample_sync(rank: int, port: int) -> None:
+    try:
+        os.environ["VLLM_DIST_IDENT"] = str(port)
+        with set_current_vllm_config(VllmConfig()):
+            init_distributed_environment(
+                world_size=2,
+                rank=rank,
+                local_rank=rank,
+                backend="gloo",
+                distributed_init_method=f"tcp://127.0.0.1:{port}",
+            )
+            initialize_model_parallel(
+                tensor_model_parallel_size=1,
+                pipeline_model_parallel_size=2,
+                backend="gloo",
+            )
+            handler = PPHandler(max_num_reqs=4, device=torch.device("cpu"))
+            batch = _batch(
+                num_computed=[4, 10],
+                prefill_len=[4, 8],
+                num_scheduled=[1, 1],
+                is_prefilling=[False, False],
+                max_seq_len=None,
+                idx_mapping=[2, 0],
+            )
+
+            assert handler.get_prev_sampled_outputs() is None
+            if handler.is_last_rank:
+                handler.broadcast(
+                    sampled_token_ids=torch.tensor([[71], [82]], dtype=torch.int64),
+                    num_sampled=torch.tensor([1, 1], dtype=torch.int32),
+                    num_rejected=torch.tensor([0, 0], dtype=torch.int32),
+                    input_batch=batch,
+                )
+            else:
+                assert handler.receive(batch)
+                assert handler.get_prev_sampled_outputs() is None
+                handler.on_req_idx_freed(0)
+                outputs = handler.get_prev_sampled_outputs()
+                assert outputs is not None
+                torch.testing.assert_close(
+                    outputs["sampled_tokens"],
+                    torch.tensor([[71], [82]], dtype=torch.int64),
+                )
+                torch.testing.assert_close(
+                    outputs["num_sampled"], torch.tensor([1, 1], dtype=torch.int32)
+                )
+                torch.testing.assert_close(
+                    outputs["num_rejected"], torch.tensor([0, 0], dtype=torch.int32)
+                )
+                torch.testing.assert_close(
+                    outputs["idx_mapping"],
+                    torch.tensor([2, -1], dtype=torch.int64),
+                )
+
+            torch.distributed.barrier()
+    finally:
+        destroy_model_parallel()
+        destroy_distributed_environment()
+
+
+def test_cpu_pp_synchronizes_samples_and_invalidates_reused_slots() -> None:
+    torch.multiprocessing.spawn(
+        _run_cpu_pp_sample_sync,
+        args=(get_open_port(),),
+        nprocs=2,
+        join=True,
+    )
diff --git a/tests/v1/worker/cpu/test_cpu_sampler.py b/tests/v1/worker/cpu/test_cpu_sampler.py
new file mode 100644
index 000000000000..4e578044912f
--- /dev/null
+++ b/tests/v1/worker/cpu/test_cpu_sampler.py
@@ -0,0 +1,747 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import numpy as np
+import pytest
+import torch
+from transformers import LlamaConfig
+
+import vllm.config.reasoning as reasoning_config_module
+from vllm.config import ModelConfig, VllmConfig
+from vllm.config.reasoning import ReasoningConfig as VllmReasoningConfig
+from vllm.platforms import CpuArchEnum, current_platform
+from vllm.sampling_params import SamplingParams
+from vllm.triton_utils import HAS_TRITON
+from vllm.v1.watermarking.gumbel import (
+    DualKeyGumbelWatermarker,
+    GumbelWatermarker,
+)
+from vllm.v1.worker.cpu.logits_processors import LogitsContext
+from vllm.v1.worker.cpu.sampler import (
+    CPUSampler,
+    CPUSamplingMaskTensors,
+    CPUThinkingBudgetState,
+    CPUTraceReplayState,
+    CPUWatermarkSampler,
+    get_num_nans,
+)
+from vllm.v1.worker.cpu.states import CPURequestState
+
+if current_platform.is_cpu():
+    import vllm._C  # noqa: F401
+
+pytestmark = pytest.mark.skipif(
+    not current_platform.is_cpu()
+    or HAS_TRITON
+    or current_platform.get_cpu_architecture() == CpuArchEnum.X86,
+    reason="requires a non-x86 CPU platform without Triton",
+)
+
+START = 90
+END = 91
+END_A = 92
+END_B = 93
+VOCAB_SIZE = 128
+
+
+class ReasoningConfig:
+    reasoning_start_token_ids = [START]
+    reasoning_end_token_ids = [END]
+    natural_reasoning_end_token_ids = [END]
+
+
+class MultiTokenEndReasoningConfig:
+    reasoning_start_token_ids = [START]
+    reasoning_end_token_ids = [END_A, END_B]
+    natural_reasoning_end_token_ids = [END_A, END_B]
+
+
+class DistinctEndReasoningConfig:
+    reasoning_start_token_ids = [START]
+    reasoning_end_token_ids = [END_A, END_B]
+    natural_reasoning_end_token_ids = [END]
+
+
+class _ReasoningTokenizer:
+    def encode(self, text: str, *, add_special_tokens: bool) -> list[int]:
+        del add_special_tokens
+        return {"<reasoning>": [START], "</reasoning>": [END]}[text]
+
+
+def _make_req_states(tokens: list[int], prompt_len: int = 1) -> CPURequestState:
+    req_states = CPURequestState(
+        max_num_reqs=4,
+        max_model_len=max(64, len(tokens) + 2),
+        max_num_batched_tokens=32,
+        vocab_size=VOCAB_SIZE,
+        device=torch.device("cpu"),
+    )
+    req_states.add_request(
+        req_id="req",
+        prompt_len=prompt_len,
+        all_token_ids=tokens,
+        num_computed_tokens=len(tokens),
+        max_tokens=32,
+    )
+    return req_states
+
+
+def _vllm_config(monkeypatch: pytest.MonkeyPatch, tmp_path) -> VllmConfig:
+    LlamaConfig(
+        hidden_size=16,
+        intermediate_size=32,
+        num_hidden_layers=1,
+        num_attention_heads=1,
+        num_key_value_heads=1,
+        max_position_embeddings=64,
+        vocab_size=VOCAB_SIZE,
+        architectures=["LlamaForCausalLM"],
+    ).save_pretrained(tmp_path)
+    model_config = ModelConfig(
+        model=str(tmp_path),
+        dtype="float32",
+        max_model_len=64,
+        skip_tokenizer_init=True,
+    )
+    reasoning_config = VllmReasoningConfig(
+        reasoning_start_str="<reasoning>",
+        reasoning_end_str="</reasoning>",
+    )
+    monkeypatch.setattr(
+        reasoning_config_module,
+        "cached_tokenizer_from_config",
+        lambda **_: _ReasoningTokenizer(),
+    )
+    reasoning_config.initialize_token_ids(model_config)
+    return VllmConfig(
+        model_config=model_config,
+        reasoning_config=reasoning_config,
+    )
+
+
+def _thinking_context(req_idx: int, input_ids: list[int]) -> LogitsContext:
+    num_rows = len(input_ids)
+    return LogitsContext(
+        expanded_idx_mapping=torch.full((num_rows,), req_idx, dtype=torch.int64),
+        idx_mapping=torch.tensor([req_idx], dtype=torch.int64),
+        idx_mapping_np=np.array([req_idx], dtype=np.intp),
+        expanded_local_pos=torch.arange(num_rows, dtype=torch.int32),
+        input_ids=torch.tensor(input_ids, dtype=torch.int32),
+        pos=torch.zeros(num_rows, dtype=torch.int32),
+        seq_lens_upper_bound_np=np.full(1, num_rows, dtype=np.int64),
+    )
+
+
+def test_cpu_thinking_budget_forces_and_continues_end_marker() -> None:
+    tokens = [1, START, 10, 11, 12]
+    req_states = _make_req_states(tokens)
+    req_idx = req_states.req_id_to_index["req"]
+    state = CPUThinkingBudgetState(req_states, MultiTokenEndReasoningConfig())
+    state.add_request(req_idx, SamplingParams(thinking_token_budget=3))
+
+    logits = torch.zeros((1, VOCAB_SIZE))
+    logits[:, [END_A, END_B]] = -torch.inf
+    state.apply(logits, _thinking_context(req_idx, [12]))
+
+    assert logits[0, END_A] == pytest.approx(1.0e9)
+
+    total_len = int(req_states.total_len[req_idx])
+    req_states.all_token_ids[req_idx, total_len] = END_A
+    req_states.total_len[req_idx] = total_len + 1
+    logits = torch.zeros((1, VOCAB_SIZE))
+    state.apply(logits, _thinking_context(req_idx, [END_A]))
+
+    assert logits[0, END_B] == pytest.approx(1.0e9)
+
+
+def test_cpu_thinking_budget_uses_distinct_forced_end_marker() -> None:
+    req_states = _make_req_states([1, START, 10, 11, 12])
+    req_idx = req_states.req_id_to_index["req"]
+    state = CPUThinkingBudgetState(req_states, DistinctEndReasoningConfig())
+    state.add_request(req_idx, SamplingParams(thinking_token_budget=3))
+
+    logits = torch.zeros((1, VOCAB_SIZE))
+    state.apply(logits, _thinking_context(req_idx, [12]))
+
+    assert logits[0, END_A] == pytest.approx(1.0e9)
+
+    total_len = int(req_states.total_len[req_idx])
+    req_states.all_token_ids[req_idx, total_len] = END_A
+    req_states.total_len[req_idx] = total_len + 1
+    logits = torch.zeros((1, VOCAB_SIZE))
+    state.apply(logits, _thinking_context(req_idx, [END_A]))
+
+    assert logits[0, END_B] == pytest.approx(1.0e9)
+
+
+def test_cpu_thinking_budget_stops_after_natural_end() -> None:
+    tokens = [1, START, 10, END, 20, 21, 22]
+    req_states = _make_req_states(tokens)
+    req_idx = req_states.req_id_to_index["req"]
+    state = CPUThinkingBudgetState(req_states, ReasoningConfig())
+    state.add_request(req_idx, SamplingParams(thinking_token_budget=2))
+
+    logits = torch.zeros((1, VOCAB_SIZE))
+    state.apply(logits, _thinking_context(req_idx, [22]))
+    assert torch.count_nonzero(logits) == 0
+
+
+def test_cpu_thinking_budget_ignores_plain_request() -> None:
+    req_states = _make_req_states([1, START, 10, 11, 12])
+    req_idx = req_states.req_id_to_index["req"]
+    state = CPUThinkingBudgetState(req_states, ReasoningConfig())
+    state.add_request(req_idx, SamplingParams())
+
+    logits = torch.zeros((1, VOCAB_SIZE))
+    state.apply(logits, _thinking_context(req_idx, [12]))
+
+    assert torch.count_nonzero(logits) == 0
+
+
+def test_cpu_thinking_budget_does_not_force_before_budget() -> None:
+    req_states = _make_req_states([1, START, 10, 11])
+    req_idx = req_states.req_id_to_index["req"]
+    state = CPUThinkingBudgetState(req_states, ReasoningConfig())
+    state.add_request(req_idx, SamplingParams(thinking_token_budget=3))
+
+    logits = torch.zeros((1, VOCAB_SIZE))
+    state.apply(logits, _thinking_context(req_idx, [11]))
+    assert torch.count_nonzero(logits) == 0
+
+
+def test_cpu_thinking_budget_uses_latest_start_boundary() -> None:
+    tokens = [1, START, 10, 11, 12, END, 13, START, 14, 15, 16]
+    req_states = _make_req_states(tokens)
+    req_idx = req_states.req_id_to_index["req"]
+    state = CPUThinkingBudgetState(req_states, ReasoningConfig())
+    state.add_request(req_idx, SamplingParams(thinking_token_budget=3))
+
+    logits = torch.zeros((1, VOCAB_SIZE))
+    state.apply(logits, _thinking_context(req_idx, [16]))
+    assert logits[0, END] == pytest.approx(1.0e9)
+
+
+def test_cpu_thinking_budget_continues_end_prefix_from_resumed_prompt() -> None:
+    tokens = [1, START, 10, 11, END_A]
+    req_states = _make_req_states(tokens, prompt_len=len(tokens))
+    req_idx = req_states.req_id_to_index["req"]
+    state = CPUThinkingBudgetState(req_states, MultiTokenEndReasoningConfig())
+    state.add_request(req_idx, SamplingParams(thinking_token_budget=3))
+
+    logits = torch.zeros((1, VOCAB_SIZE))
+    state.apply(logits, _thinking_context(req_idx, [END_A]))
+    assert logits[0, END_A] == 0
+    assert logits[0, END_B] == pytest.approx(1.0e9)
+
+
+def test_cpu_thinking_budget_clamps_oversized_budget() -> None:
+    req_states = _make_req_states([1, START, 10, 11, 12])
+    req_idx = req_states.req_id_to_index["req"]
+    state = CPUThinkingBudgetState(req_states, ReasoningConfig())
+    state.add_request(req_idx, SamplingParams(thinking_token_budget=2**40))
+
+    logits = torch.zeros((1, VOCAB_SIZE))
+    state.apply(logits, _thinking_context(req_idx, [12]))
+    assert torch.count_nonzero(logits) == 0
+
+
+def test_cpu_thinking_budget_scans_new_committed_tokens_incrementally() -> None:
+    tokens = [1, START, *([10] * 4094)]
+    req_states = _make_req_states(tokens)
+    req_idx = req_states.req_id_to_index["req"]
+    state = CPUThinkingBudgetState(req_states, ReasoningConfig())
+    state.add_request(req_idx, SamplingParams(thinking_token_budget=8192))
+
+    state.apply(torch.zeros((1, VOCAB_SIZE)), _thinking_context(req_idx, [10]))
+    assert state.cached_scan_pos[req_idx] == len(tokens)
+
+    req_states.all_token_ids[req_idx, len(tokens)] = 10
+    req_states.total_len[req_idx] = len(tokens) + 1
+    state.apply(torch.zeros((1, VOCAB_SIZE)), _thinking_context(req_idx, [10]))
+    assert state.cached_scan_pos[req_idx] == len(tokens) + 1
+
+
+def test_cpu_thinking_budget_clears_reused_slot() -> None:
+    req_states = _make_req_states([1, START, 10, 11, 12])
+    req_idx = req_states.req_id_to_index["req"]
+    state = CPUThinkingBudgetState(req_states, ReasoningConfig())
+    state.add_request(req_idx, SamplingParams(thinking_token_budget=3))
+    state.add_request(req_idx, SamplingParams())
+
+    logits = torch.zeros((1, VOCAB_SIZE))
+    state.apply(logits, _thinking_context(req_idx, [12]))
+    assert torch.count_nonzero(logits) == 0
+
+
+def test_cpu_greedy_sampling_applies_thinking_budget(monkeypatch, tmp_path) -> None:
+    req_states = _make_req_states([1, START, 10, 11, 12])
+    req_idx = req_states.req_id_to_index["req"]
+    sampler = CPUSampler(
+        vllm_config=_vllm_config(monkeypatch, tmp_path),
+        max_num_reqs=4,
+        vocab_size=VOCAB_SIZE,
+        device=torch.device("cpu"),
+        req_states=req_states,
+    )
+    sampler.add_request(
+        req_idx,
+        SamplingParams(temperature=0.0, thinking_token_budget=3),
+    )
+
+    idx_mapping = torch.tensor([req_idx], dtype=torch.int64)
+    logits = sampler.apply_sampling_params(
+        torch.zeros((1, VOCAB_SIZE)),
+        idx_mapping,
+        idx_mapping,
+        idx_mapping.numpy(),
+        torch.tensor([4], dtype=torch.int32),
+        torch.tensor([12], dtype=torch.int32),
+        torch.tensor([0], dtype=torch.int32),
+        np.array([4], dtype=np.int64),
+    )
+
+    assert logits[0, END] == pytest.approx(1.0e9)
+
+
+def test_cpu_trace_replay_uses_current_decode_step() -> None:
+    req_states = _make_req_states([1, 2, 7], prompt_len=2)
+    req_idx = req_states.req_id_to_index["req"]
+    state = CPUTraceReplayState(req_states)
+    state.add_request(req_idx, SamplingParams(trace_decode_token_ids=[11, 12, 13]))
+
+    for step, expected in enumerate([11, 12, 13]):
+        req_states.total_len[req_idx] = 2 + step
+        sampled = torch.tensor([99], dtype=torch.int64)
+        state.apply_trace(sampled, torch.tensor([req_idx], dtype=torch.int64))
+        assert sampled.tolist() == [expected]
+
+
+def test_cpu_trace_replay_maps_requests_and_skips_inactive_rows() -> None:
+    req_states = CPURequestState(
+        max_num_reqs=4,
+        max_model_len=64,
+        max_num_batched_tokens=16,
+        vocab_size=VOCAB_SIZE,
+        device=torch.device("cpu"),
+    )
+    req_states.add_request("plain", 2, [1, 2], 2, 16)
+    req_states.add_request("traced", 2, [3, 4, 7], 3, 16)
+    plain_idx = req_states.req_id_to_index["plain"]
+    traced_idx = req_states.req_id_to_index["traced"]
+    state = CPUTraceReplayState(req_states)
+    state.add_request(plain_idx, SamplingParams())
+    state.add_request(traced_idx, SamplingParams(trace_decode_token_ids=[11, 12]))
+
+    sampled = torch.tensor([20, 21, 22], dtype=torch.int64)
+    state.apply_trace(
+        sampled,
+        torch.tensor([traced_idx, -1, plain_idx], dtype=torch.int64),
+    )
+    assert sampled.tolist() == [12, 21, 22]
+
+    req_states.total_len[traced_idx] = 4
+    sampled = torch.tensor([30], dtype=torch.int64)
+    state.apply_trace(sampled, torch.tensor([traced_idx], dtype=torch.int64))
+    assert sampled.tolist() == [30]
+
+
+def test_cpu_trace_replay_clears_reused_slot() -> None:
+    req_states = _make_req_states([1, 2], prompt_len=2)
+    req_idx = req_states.req_id_to_index["req"]
+    state = CPUTraceReplayState(req_states)
+    state.add_request(req_idx, SamplingParams(trace_decode_token_ids=[11]))
+    state.add_request(req_idx, SamplingParams())
+
+    sampled = torch.tensor([20], dtype=torch.int64)
+    state.apply_trace(sampled, torch.tensor([req_idx], dtype=torch.int64))
+    assert sampled.tolist() == [20]
+
+
+def test_cpu_nan_count() -> None:
+    logits = torch.tensor(
+        [
+            [0.0, -torch.inf, torch.inf, 1.0, torch.nan, -torch.inf, 2.0, -torch.inf],
+            [0.0] * 8,
+        ]
+    )
+    num_nans = get_num_nans(logits)
+    assert num_nans.tolist() == [1, 0]
+
+
+def test_cpu_sampling_mask_tracks_finite_tokens() -> None:
+    logits = torch.tensor(
+        [
+            [0.0, -torch.inf, torch.inf, 1.0, torch.nan, -torch.inf, 2.0, -torch.inf],
+            [0.0] * 8,
+        ]
+    )
+    mask = CPUSamplingMaskTensors.from_logits(
+        logits,
+        torch.tensor([0, 1, 2], dtype=torch.int32),
+        torch.tensor([1, 0], dtype=torch.int32),
+        max_num_kept=2,
+    )
+    assert mask.counts.tolist() == [3, 0]
+    assert mask.token_ids[0].tolist() == [0, 3]
+    assert mask.packed_mask[0].tolist() == [73]
+
+
+@pytest.mark.parametrize("max_num_kept", [512, 20_001])
+def test_cpu_sampling_mask_matches_wide_finite_support(max_num_kept: int) -> None:
+    vocab_size = 20_001
+    sizes = [0, 1, 511, 512, 513, 2053, vocab_size]
+    generator = torch.Generator().manual_seed(0)
+    logits = torch.full((len(sizes), vocab_size), -torch.inf)
+    expected = []
+    for row, size in enumerate(sizes):
+        kept = torch.randperm(vocab_size, generator=generator)[:size].sort().values
+        logits[row, kept] = 0
+        expected.append(kept.tolist())
+    num_sampled = torch.ones(len(sizes), dtype=torch.int32)
+    num_sampled[1] = 0
+    expected[1] = []
+
+    mask = CPUSamplingMaskTensors.from_logits(
+        logits,
+        torch.arange(len(sizes) + 1, dtype=torch.int32),
+        num_sampled,
+        max_num_kept,
+    )
+
+    assert mask.tolists().to_nested_list() == expected
+
+
+def _make_sampler(
+    sampling_params: list[SamplingParams],
+) -> tuple[CPUSampler, torch.Tensor]:
+    req_states = CPURequestState(
+        max_num_reqs=4,
+        max_model_len=64,
+        max_num_batched_tokens=16,
+        vocab_size=VOCAB_SIZE,
+        device=torch.device("cpu"),
+    )
+    sampler = CPUSampler(
+        vllm_config=VllmConfig(),
+        max_num_reqs=4,
+        vocab_size=VOCAB_SIZE,
+        device=torch.device("cpu"),
+        req_states=req_states,
+    )
+    req_indices = []
+    for i, params in enumerate(sampling_params):
+        req_id = f"req-{i}"
+        req_states.add_request(req_id, 1, [i + 1], 1, 16)
+        req_idx = req_states.req_id_to_index[req_id]
+        req_indices.append(req_idx)
+        sampler.add_request(req_idx, params)
+    return sampler, torch.tensor(req_indices, dtype=torch.int64)
+
+
+@pytest.mark.parametrize(
+    ("sampling_params", "expected_finite"),
+    [
+        (SamplingParams(top_k=2), [2, 3]),
+        (SamplingParams(top_p=0.7), [2, 3]),
+        (SamplingParams(top_k=3, top_p=0.7), [2, 3]),
+        (SamplingParams(top_k=VOCAB_SIZE + 1, top_p=0.7), [2, 3]),
+    ],
+)
+def test_cpu_sampler_applies_top_k_top_p(
+    sampling_params: SamplingParams,
+    expected_finite: list[int],
+) -> None:
+    sampler, idx_mapping = _make_sampler([sampling_params])
+
+    logits = sampler.apply_sampling_params(
+        torch.tensor([[0.0, 1.0, 2.0, 3.0]]),
+        idx_mapping,
+        idx_mapping,
+        idx_mapping.numpy(),
+        torch.tensor([1], dtype=torch.int64),
+        torch.tensor([1], dtype=torch.int32),
+        torch.tensor([0], dtype=torch.int32),
+        np.array([1], dtype=np.int64),
+    )
+
+    assert torch.isfinite(logits[0]).nonzero().flatten().tolist() == expected_finite
+
+
+def test_cpu_sampler_applies_temperature() -> None:
+    sampler, idx_mapping = _make_sampler([SamplingParams(temperature=0.5)])
+    logits = torch.tensor([[0.0, 1.0, 2.0, 3.0]])
+
+    actual = sampler.apply_sampling_params(
+        logits,
+        idx_mapping,
+        idx_mapping,
+        idx_mapping.numpy(),
+        torch.tensor([1], dtype=torch.int64),
+        torch.tensor([1], dtype=torch.int32),
+        torch.tensor([0], dtype=torch.int32),
+        np.array([1], dtype=np.int64),
+    )
+
+    torch.testing.assert_close(actual, logits / 0.5)
+
+
+def test_cpu_sampler_only_processes_active_requests() -> None:
+    sampler, idx_mapping = _make_sampler(
+        [SamplingParams(temperature=0.0), SamplingParams(temperature=0.5)]
+    )
+    logits = torch.tensor([[0.0, 1.0, 2.0, 3.0]])
+
+    greedy = sampler.apply_sampling_params(
+        logits,
+        idx_mapping[:1],
+        idx_mapping[:1],
+        idx_mapping[:1].numpy(),
+        torch.tensor([1], dtype=torch.int32),
+        torch.tensor([1], dtype=torch.int32),
+        torch.tensor([0], dtype=torch.int32),
+        np.ones(1, dtype=np.int64),
+    )
+    scaled = sampler.apply_sampling_params(
+        logits,
+        idx_mapping[1:],
+        idx_mapping[1:],
+        idx_mapping[1:].numpy(),
+        torch.tensor([1], dtype=torch.int32),
+        torch.tensor([2], dtype=torch.int32),
+        torch.tensor([0], dtype=torch.int32),
+        np.ones(1, dtype=np.int64),
+    )
+
+    torch.testing.assert_close(greedy, logits)
+    torch.testing.assert_close(scaled, logits / 0.5)
+
+
+def test_cpu_sampler_clears_processing_when_slot_is_reused() -> None:
+    sampler, idx_mapping = _make_sampler([SamplingParams(temperature=0.5)])
+    req_idx = int(idx_mapping[0])
+    sampler.add_request(req_idx, SamplingParams(temperature=0.0))
+    logits = torch.tensor([[0.0, 1.0, 2.0, 3.0]])
+
+    actual = sampler.apply_sampling_params(
+        logits,
+        idx_mapping,
+        idx_mapping,
+        idx_mapping.numpy(),
+        torch.tensor([1], dtype=torch.int32),
+        torch.tensor([1], dtype=torch.int32),
+        torch.tensor([0], dtype=torch.int32),
+        np.ones(1, dtype=np.int64),
+    )
+
+    torch.testing.assert_close(actual, logits)
+
+
+def test_cpu_sampler_greedy_returns_argmax() -> None:
+    sampler, idx_mapping = _make_sampler([SamplingParams(temperature=0.0)])
+    logits = torch.tensor([[0.0, 3.0, 2.0, 1.0]])
+
+    sampled, _ = sampler.sample(
+        logits,
+        idx_mapping,
+        idx_mapping,
+        idx_mapping.numpy(),
+        torch.tensor([1], dtype=torch.int64),
+        torch.tensor([1], dtype=torch.int32),
+        torch.tensor([0], dtype=torch.int32),
+        np.ones(1, dtype=np.int64),
+    )
+
+    assert sampled.tolist() == [1]
+
+
+@pytest.mark.parametrize("algorithm", ["gumbel", "dual-key"])
+@pytest.mark.parametrize("watermarking", [False, True])
+def test_cpu_watermark_sampler_uses_cpu_sampler_pipeline(
+    algorithm: str,
+    watermarking: bool,
+) -> None:
+    req_states = _make_req_states([1])
+    req_idx = req_states.req_id_to_index["req"]
+    watermarker = (
+        GumbelWatermarker(key=42, context_width=1)
+        if algorithm == "gumbel"
+        else DualKeyGumbelWatermarker(key=42, context_width=1)
+    )
+    sampler = CPUWatermarkSampler(
+        watermarker,
+        vllm_config=VllmConfig(),
+        max_num_reqs=4,
+        vocab_size=VOCAB_SIZE,
+        device=torch.device("cpu"),
+        req_states=req_states,
+    )
+    params = SamplingParams(
+        temperature=1.0,
+        top_k=4,
+        top_p=0.9,
+        seed=7,
+        watermarking=watermarking,
+    )
+    sampler.add_request(req_idx, params)
+
+    idx_mapping = torch.tensor([req_idx], dtype=torch.int64)
+    logits = torch.linspace(-1, 1, VOCAB_SIZE).unsqueeze(0)
+    sample_args = (
+        logits,
+        idx_mapping,
+        idx_mapping,
+        idx_mapping.numpy(),
+        torch.tensor([1], dtype=torch.int64),
+        torch.tensor([1], dtype=torch.int32),
+        torch.tensor([0], dtype=torch.int32),
+        np.ones(1, dtype=np.int64),
+    )
+    sampled, output_logits = sampler.sample(*sample_args)
+    repeated, repeated_logits = sampler.sample(*sample_args)
+
+    assert sampled.item() in range(VOCAB_SIZE - 4, VOCAB_SIZE)
+    assert torch.equal(repeated, sampled)
+    torch.testing.assert_close(repeated_logits, output_logits)
+    assert torch.isfinite(output_logits[0, sampled.item()])
+
+    ordinary = CPUSampler(
+        vllm_config=VllmConfig(),
+        max_num_reqs=4,
+        vocab_size=VOCAB_SIZE,
+        device=torch.device("cpu"),
+        req_states=req_states,
+    )
+    ordinary.add_request(req_idx, params)
+    expected, expected_logits = ordinary.sample(*sample_args)
+    if watermarking:
+        assert not torch.equal(sampled, expected)
+    else:
+        assert torch.equal(sampled, expected)
+        torch.testing.assert_close(output_logits, expected_logits)
+
+
+def test_cpu_watermark_sampler_handles_mixed_rows() -> None:
+    req_states = CPURequestState(
+        max_num_reqs=4,
+        max_model_len=64,
+        max_num_batched_tokens=16,
+        vocab_size=VOCAB_SIZE,
+        device=torch.device("cpu"),
+    )
+    params = [
+        SamplingParams(temperature=1.0, top_k=4, seed=7, watermarking=True),
+        SamplingParams(temperature=1.0, top_k=4, seed=11, watermarking=False),
+        SamplingParams(temperature=0.0, top_k=4, watermarking=True),
+    ]
+    req_indices = []
+    for i in range(len(params)):
+        req_id = f"req-{i}"
+        req_states.add_request(req_id, 1, [9, i + 1], 2, 16)
+        req_indices.append(req_states.req_id_to_index[req_id])
+    idx_mapping = torch.tensor(req_indices, dtype=torch.int64)
+
+    sampler = CPUWatermarkSampler(
+        GumbelWatermarker(key=42, context_width=1),
+        vllm_config=VllmConfig(),
+        max_num_reqs=4,
+        vocab_size=VOCAB_SIZE,
+        device=torch.device("cpu"),
+        req_states=req_states,
+        deduplicate_contexts="none",
+    )
+    ordinary = CPUSampler(
+        vllm_config=VllmConfig(),
+        max_num_reqs=4,
+        vocab_size=VOCAB_SIZE,
+        device=torch.device("cpu"),
+        req_states=req_states,
+    )
+    for req_idx, sampling_params in zip(req_indices, params, strict=True):
+        sampler.add_request(req_idx, sampling_params)
+        ordinary.add_request(req_idx, sampling_params)
+
+    logits = torch.linspace(-1, 1, VOCAB_SIZE).repeat(3, 1)
+    sample_args = (
+        logits,
+        idx_mapping,
+        idx_mapping,
+        idx_mapping.numpy(),
+        torch.tensor([2, 2, 2], dtype=torch.int64),
+        torch.tensor([1, 2, 3], dtype=torch.int32),
+        torch.zeros(3, dtype=torch.int32),
+        np.full(3, 2, dtype=np.int64),
+    )
+    sampled, _ = sampler.sample(*sample_args)
+    ordinary_sampled, _ = ordinary.sample(*sample_args)
+
+    assert sampled[0] != ordinary_sampled[0]
+    assert sampled[1] == ordinary_sampled[1]
+    assert sampled[2] == logits[2].argmax()
+    assert torch.all(sampled >= VOCAB_SIZE - 4)
+
+
+def test_cpu_watermark_sampler_detects_repeated_context() -> None:
+    req_states = _make_req_states([9, 1, 2, 3, 1, 2], prompt_len=1)
+    req_idx = req_states.req_id_to_index["req"]
+    sampler = CPUWatermarkSampler(
+        GumbelWatermarker(key=42, context_width=2),
+        vllm_config=VllmConfig(),
+        max_num_reqs=4,
+        vocab_size=VOCAB_SIZE,
+        device=torch.device("cpu"),
+        req_states=req_states,
+    )
+
+    req_indices = torch.tensor([req_idx], dtype=torch.int64)
+    contexts = sampler._get_contexts(req_indices)
+
+    assert contexts.tolist() == [[1, 2]]
+    assert sampler._get_repeated_contexts(req_indices, contexts).tolist() == [True]
+
+
+def test_cpu_watermark_deduplication_scope_and_partial_context() -> None:
+    req_states = _make_req_states([1, 2, 1, 2], prompt_len=2)
+    req_idx = req_states.req_id_to_index["req"]
+    req_indices = torch.tensor([req_idx], dtype=torch.int64)
+
+    single_turn = CPUWatermarkSampler(
+        GumbelWatermarker(key=42, context_width=2),
+        vllm_config=VllmConfig(),
+        max_num_reqs=4,
+        vocab_size=VOCAB_SIZE,
+        device=torch.device("cpu"),
+        req_states=req_states,
+        deduplicate_contexts="single_turn",
+    )
+    all_history = CPUWatermarkSampler(
+        GumbelWatermarker(key=42, context_width=2),
+        vllm_config=VllmConfig(),
+        max_num_reqs=4,
+        vocab_size=VOCAB_SIZE,
+        device=torch.device("cpu"),
+        req_states=req_states,
+        deduplicate_contexts="all",
+    )
+    contexts = single_turn._get_contexts(req_indices)
+
+    assert single_turn._get_repeated_contexts(req_indices, contexts).tolist() == [False]
+    assert all_history._get_repeated_contexts(req_indices, contexts).tolist() == [True]
+
+    partial_states = _make_req_states([9, 1], prompt_len=1)
+    partial_idx = partial_states.req_id_to_index["req"]
+    partial_sampler = CPUWatermarkSampler(
+        GumbelWatermarker(key=42, context_width=2),
+        vllm_config=VllmConfig(),
+        max_num_reqs=4,
+        vocab_size=VOCAB_SIZE,
+        device=torch.device("cpu"),
+        req_states=partial_states,
+        deduplicate_contexts="all",
+    )
+    partial_indices = torch.tensor([partial_idx], dtype=torch.int64)
+    partial_context = partial_sampler._get_contexts(partial_indices)
+
+    assert partial_context.tolist() == [[-1, 1]]
+    assert partial_sampler._get_repeated_contexts(
+        partial_indices, partial_context
+    ).tolist() == [True]
diff --git a/tests/v1/worker/cpu/test_cpu_structured_outputs.py b/tests/v1/worker/cpu/test_cpu_structured_outputs.py
new file mode 100644
index 000000000000..a658939254c7
--- /dev/null
+++ b/tests/v1/worker/cpu/test_cpu_structured_outputs.py
@@ -0,0 +1,61 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import numpy as np
+import pytest
+import torch
+
+from vllm.platforms import CpuArchEnum, current_platform
+from vllm.triton_utils import HAS_TRITON
+from vllm.v1.worker.cpu.input_batch import InputBatch, InputBuffers
+from vllm.v1.worker.cpu.structured_outputs import CPUStructuredOutputsWorker
+
+pytestmark = pytest.mark.skipif(
+    not current_platform.is_cpu()
+    or HAS_TRITON
+    or current_platform.get_cpu_architecture() == CpuArchEnum.X86,
+    reason="requires a non-x86 CPU platform without Triton",
+)
+
+
+def _worker(vocab_size: int) -> CPUStructuredOutputsWorker:
+    return CPUStructuredOutputsWorker(
+        max_num_logits=8,
+        vocab_size=vocab_size,
+        device=torch.device("cpu"),
+        mask_stride=1,
+    )
+
+
+def _input_batch(req_ids: list[str]) -> InputBatch:
+    buffers = InputBuffers(
+        max_num_reqs=len(req_ids),
+        max_num_tokens=len(req_ids),
+        device=torch.device("cpu"),
+    )
+    batch = InputBatch.make_dummy(len(req_ids), len(req_ids), buffers)
+    batch.req_ids = req_ids
+    return batch
+
+
+def test_cpu_grammar_mask_applies_to_mapped_rows_only() -> None:
+    input_batch = _input_batch(["plain", "grammar"])
+    grammar_bitmask = np.array([[(1 << 1) | (1 << 4)]], dtype=np.int32)
+    logits = torch.zeros((2, 10))
+
+    _worker(10).apply_grammar_bitmask(logits, input_batch, ["grammar"], grammar_bitmask)
+
+    assert torch.isfinite(logits[0]).all()
+    assert torch.isfinite(logits[1]).nonzero().flatten().tolist() == [1, 4]
+
+
+def test_cpu_grammar_mask_handles_word_boundaries_and_sign_bit() -> None:
+    input_batch = _input_batch(["grammar"])
+    packed = np.array([[(1 << 0) | (1 << 31), 1 << 0, 1 << 5]], dtype=np.uint32).view(
+        np.int32
+    )
+    logits = torch.zeros((1, 70))
+
+    _worker(70).apply_grammar_bitmask(logits, input_batch, ["grammar"], packed)
+
+    assert torch.isfinite(logits[0]).nonzero().flatten().tolist() == [0, 31, 32, 69]
diff --git a/tests/v1/worker/test_attn_utils.py b/tests/v1/worker/test_attn_utils.py
index 197929731f42..5b284f620626 100644
--- a/tests/v1/worker/test_attn_utils.py
+++ b/tests/v1/worker/test_attn_utils.py
@@ -26,6 +26,7 @@
 from vllm.config.compilation import CompilationConfig, CUDAGraphMode
 from vllm.v1.attention.backend import AttentionBackend, AttentionCGSupport, MultipleOf
 from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy
+from vllm.v1.hisparse.backend import CUDAHiSparseBackend
 from vllm.v1.hisparse.binding import allocate_hisparse_kv_caches
 from vllm.v1.kv_cache_interface import (
     FullAttentionSpec,
@@ -502,7 +503,10 @@ def test_profiling_cleanup_keeps_host_pool_for_real_kv_cache(monkeypatch):
     pool = attn_utils_module.get_hisparse_host_pool(vllm_config, kv_cache_config)
     backing = pool.allocate(8)
     runtime = SimpleNamespace(
-        _host_cache=backing, registered_host_pool=pinned, hot_backing=object()
+        _host_cache=backing,
+        registered_host_pool=pinned,
+        hot_backing=object(),
+        host_pool=pool,
     )
     forward_context = {
         "layer": SimpleNamespace(hisparse_cache=SimpleNamespace(runtime=runtime))
@@ -511,6 +515,7 @@ def test_profiling_cleanup_keeps_host_pool_for_real_kv_cache(monkeypatch):
     attn_utils_module.release_hisparse_profiling_cache(forward_context)
 
     assert not hasattr(runtime, "_host_cache")
+    assert runtime.host_pool is None
     reused = attn_utils_module.get_hisparse_host_pool(vllm_config, kv_cache_config)
     assert reused is pool
     assert reused.allocate(8) is backing
@@ -518,10 +523,21 @@ def test_profiling_cleanup_keeps_host_pool_for_real_kv_cache(monkeypatch):
     hisparse_runtime_module.drop_hisparse_host_pool(vllm_config)
 
 
-@pytest.mark.parametrize("failure_phase", ["allocation", "binding", "buffers"])
+@pytest.mark.parametrize("failure_phase", ["allocation", "views", "binding", "buffers"])
 @pytest.mark.parametrize("shared", [False, True])
 def test_init_hisparse_rolls_back_host_pool(monkeypatch, failure_phase, shared):
     """Failed initialization must release shared or private registrations."""
+
+    class Backend(CUDAHiSparseBackend):
+        def prepare_kv_cache_views(self, *args):
+            if failure_phase == "views":
+                raise RuntimeError("initialization failed")
+
+    backend = Backend()
+    monkeypatch.setattr(attn_utils_module, "get_hisparse_backend", lambda: backend)
+    monkeypatch.setattr(
+        hisparse_runtime_module, "get_hisparse_backend", lambda: backend
+    )
     region = _FakeSharedHostRegion()
     registered = torch.empty(4096, dtype=torch.int8)
     cudart = MagicMock()
@@ -544,6 +560,7 @@ def allocate(*args):
         else:
             args[-1].registered = registered
             args[-1].backing = registered[:8]
+        args[-1].pinned = [registered]
         if failure_phase == "allocation":
             raise RuntimeError("initialization failed")
         return {}
@@ -571,10 +588,7 @@ def buffers(*args, **kwargs):
             SimpleNamespace(),
         )
     assert region.cleanup_calls == int(shared)
-    if shared:
-        cudart.cudaHostUnregister.assert_not_called()
-    else:
-        cudart.cudaHostUnregister.assert_called_once_with(registered.data_ptr())
+    cudart.cudaHostUnregister.assert_called_once_with(registered.data_ptr())
     assert attn_utils_module.get_hisparse_host_pool(vllm_config, None) is not pools[0]
     hisparse_runtime_module.drop_hisparse_host_pool(vllm_config)
 
diff --git a/tests/v1/worker/test_draft_tokens_handler.py b/tests/v1/worker/test_draft_tokens_handler.py
new file mode 100644
index 000000000000..84150c98dcfe
--- /dev/null
+++ b/tests/v1/worker/test_draft_tokens_handler.py
@@ -0,0 +1,84 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from dataclasses import replace
+
+import pytest
+import torch
+
+from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
+from vllm.v1.worker.gpu.spec_decode.utils import DraftTokensHandler
+
+
+def _input_batch(req_ids: list[str], structured: bool) -> InputBatch:
+    num_reqs = len(req_ids)
+    buffers = InputBuffers(num_reqs, num_reqs, torch.device("cuda"))
+    batch = InputBatch.make_dummy(num_reqs, num_reqs, buffers, is_padding=False)
+    return replace(batch, req_ids=req_ids, has_structured_output_reqs=structured)
+
+
+def _set_draft_toks(
+    handler: DraftTokensHandler,
+    req_ids: list[str],
+    tokens: list[list[int]],
+    structured: bool,
+) -> None:
+    batch = _input_batch(req_ids=req_ids, structured=structured)
+    handler.set_draft_tokens(batch, torch.tensor(tokens, device="cuda"))
+
+
+def _get_draft_toks(handler: DraftTokensHandler) -> dict[str, list[int]]:
+    drafts = handler.get_draft_tokens()
+    assert drafts is not None
+    return dict(zip(drafts.req_ids, drafts.draft_token_ids, strict=True))
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
+def test_draft_tokens_follow_request_lifecycle() -> None:
+    handler = DraftTokensHandler(torch.device("cuda"))
+
+    # The worker samples two structured-output requests
+    _set_draft_toks(
+        handler,
+        req_ids=["a", "b"],
+        tokens=[[11, 12], [21, 22]],
+        structured=True,
+    )
+
+    # "b" finishes before this result is handed back to the scheduler
+    handler.discard("b")
+    assert _get_draft_toks(handler) == {"a": [11, 12]}
+
+    # "a" continues and "b" is re-added with a new draft row
+    _set_draft_toks(
+        handler,
+        req_ids=["a", "b"],
+        tokens=[[13, 14], [23, 24]],
+        structured=True,
+    )
+    assert _get_draft_toks(handler) == {
+        "a": [13, 14],
+        "b": [23, 24],
+    }
+
+    # A batch without structured-output must not erase drafts from live requests
+    _set_draft_toks(
+        handler,
+        req_ids=["plain"],
+        tokens=[[31, 32]],
+        structured=False,
+    )
+    assert _get_draft_toks(handler) == {
+        "a": [13, 14],
+        "b": [23, 24],
+        "plain": [-1, -1],
+    }
+
+    # Requests finish independently, and the handler is empty once all are gone
+    handler.discard("a")
+    assert _get_draft_toks(handler) == {
+        "b": [23, 24],
+        "plain": [-1, -1],
+    }
+    handler.discard("b")
+    handler.discard("plain")
+    assert _get_draft_toks(handler) == {}
diff --git a/tests/v1/worker/test_eagle3_aux_hidden_states_pp.py b/tests/v1/worker/test_eagle3_aux_hidden_states_pp.py
index 89bfc6a53d50..021a5b431483 100644
--- a/tests/v1/worker/test_eagle3_aux_hidden_states_pp.py
+++ b/tests/v1/worker/test_eagle3_aux_hidden_states_pp.py
@@ -10,35 +10,56 @@
 from vllm.model_executor.models.mimo import MiMoModel
 from vllm.models.minimax_m3.nvidia import model as minimax_m3
 from vllm.sequence import IntermediateTensors
+from vllm.v1.worker.gpu.spec_decode.dflash.config import DFlashSpeculatorConfig
+from vllm.v1.worker.gpu.spec_decode.dspark.config import DSparkSpeculatorConfig
+from vllm.v1.worker.gpu.spec_decode.eagle.config import Eagle3SpeculatorConfig
 from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import (
-    requires_aux_hidden_states,
     reserve_aux_intermediate_tensor_slots,
     verify_supports_aux_hidden_states_over_pp,
 )
+from vllm.v1.worker.gpu.spec_decode.extract_hidden_states_config import (
+    ExtractHiddenStatesSpeculatorConfig,
+)
+from vllm.v1.worker.gpu.spec_decode.mtp.config import MTPSpeculatorConfig
 
 
 @pytest.mark.parametrize(
-    ("method", "hf_config", "expected"),
+    ("config_cls", "hf_config", "expected"),
     [
-        ("eagle3", {"eagle_config": {"use_aux_hidden_state": False}}, False),
-        ("eagle3", {"eagle_config": {"use_aux_hidden_state": True}}, True),
-        ("eagle3", {"eagle_config": {}}, True),
-        ("eagle3", {"eagle_config": None}, True),
-        ("eagle3", {}, True),
-        ("dflash", {"eagle_config": {"use_aux_hidden_state": False}}, True),
-        ("dspark", {"eagle_config": {"use_aux_hidden_state": False}}, True),
-        ("extract_hidden_states", {}, True),
-        ("mtp", {}, False),
+        (
+            Eagle3SpeculatorConfig,
+            {"eagle_config": {"use_aux_hidden_state": False}},
+            False,
+        ),
+        (
+            Eagle3SpeculatorConfig,
+            {"eagle_config": {"use_aux_hidden_state": True}},
+            True,
+        ),
+        (Eagle3SpeculatorConfig, {"eagle_config": {}}, True),
+        (Eagle3SpeculatorConfig, {"eagle_config": None}, True),
+        (Eagle3SpeculatorConfig, {}, True),
+        (
+            DFlashSpeculatorConfig,
+            {"eagle_config": {"use_aux_hidden_state": False}},
+            True,
+        ),
+        (
+            DSparkSpeculatorConfig,
+            {"eagle_config": {"use_aux_hidden_state": False}},
+            True,
+        ),
+        (ExtractHiddenStatesSpeculatorConfig, {}, True),
+        (MTPSpeculatorConfig, {}, False),
     ],
 )
 def test_aux_hidden_state_requirement_respects_eagle3_config(
-    method, hf_config, expected
+    config_cls, hf_config, expected
 ):
     spec_config = SimpleNamespace(
-        method=method,
         draft_model_config=SimpleNamespace(hf_config=SimpleNamespace(**hf_config)),
     )
-    assert requires_aux_hidden_states(spec_config) is expected
+    assert config_cls(spec_config).uses_aux_hidden_states is expected
 
 
 def test_aux_layers_are_sorted_and_deduplicated():
@@ -57,29 +78,39 @@ def test_mimo_does_not_inherit_aux_hidden_state_pp_support():
 
 
 class _PPAuxLayer(torch.nn.Module):
-    """Aliases and mutates residual in place like MiniMaxM3DecoderLayer."""
+    """Aliases and mutates residual in place like MiniMaxM3DecoderLayer.
 
-    ffn_all_reduce_deferred = False
+    With ``deferred``, the FFN output is this rank's half of a 2-way TP sum, and
+    the next input norm completes the all-reduce.
+    """
 
-    def __init__(self, index):
+    def __init__(self, index, deferred):
         super().__init__()
         self.index = index
+        self.ffn_all_reduce_deferred = deferred
+        self.fuse_input_allreduce = False
+
+    def forward_input_norm(self, hidden_states, residual):
+        if residual is None:
+            return hidden_states, hidden_states
+        if self.fuse_input_allreduce:
+            hidden_states = hidden_states * 2
+        residual.add_(hidden_states)
+        return hidden_states, residual
 
-    def forward(self, positions, hidden_states, residual):
-        total = hidden_states if residual is None else hidden_states + residual
-        residual = hidden_states if residual is None else residual
-        residual.copy_(total + self.index + 1)
-        return torch.full_like(hidden_states, 10 * (self.index + 1)), residual
+    def forward_after_input_norm(self, positions, hidden_states, residual):
+        residual.add_(self.index + 1)
+        output = torch.full_like(hidden_states, 10 * (self.index + 1))
+        return output / 2 if self.ffn_all_reduce_deferred else output, residual
 
 
 class _PPAuxNorm(torch.nn.Module):
-    """Clobbers its inputs so aux taps that alias them would fail."""
+    """Fused add + norm (identity scaling) that overwrites its inputs in place."""
 
     def forward(self, hidden_states, residual):
-        output = hidden_states + residual
-        hidden_states.fill_(-101)
-        residual.fill_(-103)
-        return output, residual
+        residual.add_(hidden_states)
+        hidden_states.copy_(residual)
+        return hidden_states, residual
 
 
 @pytest.fixture
@@ -90,6 +121,9 @@ def minimax_pp_stage(monkeypatch):
     monkeypatch.setattr(minimax_m3, "get_pp_group", lambda: group)
     monkeypatch.setattr(parallel_state, "get_pp_group", lambda: group)
     monkeypatch.setattr(parallel_state, "model_parallel_is_initialized", lambda: True)
+    monkeypatch.setattr(
+        minimax_m3.MiniMaxM3Model, "_init_fused_moe_finalize", lambda *_: None
+    )
     monkeypatch.setattr(
         minimax_m3,
         "VocabParallelEmbedding",
@@ -98,10 +132,18 @@ def minimax_pp_stage(monkeypatch):
     monkeypatch.setattr(
         minimax_m3, "MiniMAXGemmaRMSNorm", lambda *args, **kwargs: _PPAuxNorm()
     )
+    deferred = [False]
     monkeypatch.setattr(
         minimax_m3,
         "MiniMaxM3DecoderLayer",
-        lambda **kwargs: _PPAuxLayer(int(kwargs["prefix"].rsplit(".", 1)[1])),
+        lambda **kwargs: _PPAuxLayer(
+            int(kwargs["prefix"].rsplit(".", 1)[1]), deferred[0]
+        ),
+    )
+    monkeypatch.setattr(
+        minimax_m3,
+        "fused_allreduce_gemma_rms_norm",
+        lambda hidden_states, residual, norm: norm(hidden_states * 2, residual),
     )
 
     def make_layers(num_layers, build, prefix):
@@ -117,8 +159,9 @@ def make_layers(num_layers, build, prefix):
 
     monkeypatch.setattr(minimax_m3, "make_layers", make_layers)
 
-    def make_stage(start, end, taps, world_size=2):
+    def make_stage(start, end, taps, world_size=2, deferred_all_reduce=False):
         bounds[:] = [start, end]
+        deferred[0] = deferred_all_reduce
         group.world_size = world_size
         group.is_first_rank = start == 0
         group.is_last_rank = end == 6
@@ -130,6 +173,7 @@ def make_stage(start, end, taps, world_size=2):
                 model_config=SimpleNamespace(hf_text_config=config),
                 quant_config=None,
                 speculative_config=None,
+                kernel_config=SimpleNamespace(moe_backend="auto"),
                 use_v2_model_runner=True,
             )
         )
@@ -204,3 +248,15 @@ def test_minimax_pp2_without_aux_keeps_hidden_and_residual_transport(minimax_pp_
     ) == {"hidden_states", "residual"}
     output = last(None, None, sent)
     torch.testing.assert_close(output, _pp_aux_state(_expected_aux_value(6)))
+
+
+@pytest.mark.parametrize("taps", [(0, 3, 6), (1, 2, 5)])
+def test_minimax_aux_taps_follow_deferred_tp_all_reduce(minimax_pp_stage, taps):
+    """Under TP each FFN output stays a partial sum until the next input norm or
+    the final norm all-reduces it, so aux taps must read the reduced stream."""
+    model, _ = minimax_pp_stage(0, 6, taps, world_size=1, deferred_all_reduce=True)
+    assert model.fuse_final_norm_allreduce
+    output, aux = model(None, None, None, _pp_aux_state(7))
+    for hidden, tap in zip(aux, taps, strict=True):
+        torch.testing.assert_close(hidden, _pp_aux_state(_expected_aux_value(tap)))
+    torch.testing.assert_close(output, _pp_aux_state(_expected_aux_value(6)))
diff --git a/tests/v1/worker/test_extensible_kv_cache.py b/tests/v1/worker/test_extensible_kv_cache.py
new file mode 100644
index 000000000000..e90a2da79680
--- /dev/null
+++ b/tests/v1/worker/test_extensible_kv_cache.py
@@ -0,0 +1,585 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import copy
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm.config import ModelConfig, VllmConfig
+from vllm.utils.extensible_tensor import (
+    granule_aligned_blocks,
+    granule_block_alignment,
+)
+from vllm.utils.vmm_driver import vmm_unavailable_reason
+from vllm.v1.core.kv_cache_planning import get_kv_cache_config_from_groups
+from vllm.v1.kv_cache_interface import (
+    FullAttentionSpec,
+    KVCacheGroupSpec,
+    KVCacheTensor,
+    UniformTypeKVCacheSpecs,
+)
+from vllm.v1.kv_cache_layout import KVCacheLayout
+from vllm.v1.worker import extensible_kv_cache
+from vllm.v1.worker.extensible_kv_cache import (
+    ExtensibleKVCache,
+    measure_kv_cache_bytes,
+)
+from vllm.v1.worker.utils import allocate_kv_cache
+
+requires_cuda = pytest.mark.skipif(
+    not torch.cuda.is_available(), reason="requires CUDA"
+)
+
+
+@pytest.fixture
+def vmm():
+    # Probed lazily so collection does not initialize CUDA in the pytest parent.
+    if (reason := vmm_unavailable_reason()) is not None:
+        pytest.skip(f"VMM unavailable: {reason}")
+
+
+NUM_BLOCKS = 8
+NUM_LAYERS = 3
+
+
+def _make_config(layout: KVCacheLayout, num_blocks: int = NUM_BLOCKS):
+    spec = FullAttentionSpec(
+        block_size=16, num_kv_heads=2, head_size=64, dtype=torch.float16
+    )
+    vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=16))
+    vllm_config.cache_config.kv_cache_layout = layout.name
+    layers = [f"layer.{i}" for i in range(NUM_LAYERS)]
+    config = get_kv_cache_config_from_groups(
+        vllm_config,
+        [KVCacheGroupSpec(layers, spec)],
+        available_memory=NUM_LAYERS * spec.page_size_bytes * num_blocks,
+    )
+    assert config.num_blocks == num_blocks
+    return config, spec
+
+
+@requires_cuda
+@pytest.mark.parametrize(
+    ("layout", "expected_segments"),
+    [(KVCacheLayout.LBNHC, NUM_LAYERS), (KVCacheLayout.BLNHC, 1)],
+)
+def test_segments_follow_layout(vmm, layout, expected_segments):
+    """Layer-outermost layouts commit one prefix per layer; block-outermost
+    layouts commit a single prefix of the whole allocation."""
+    config, spec = _make_config(layout)
+    kv_cache = ExtensibleKVCache(config, torch.device("cuda"))
+    try:
+        assert kv_cache.buffer.num_segments == expected_segments
+        assert kv_cache.bytes_per_block == NUM_LAYERS * spec.page_size_bytes
+        assert kv_cache.num_committed_blocks == 0
+    finally:
+        kv_cache.free()
+
+
+@requires_cuda
+@pytest.mark.parametrize("layout", [KVCacheLayout.LBNHC, KVCacheLayout.BLNHC])
+def test_views_stay_valid_across_commits(vmm, layout):
+    """Layer views built over the reservation address committed blocks at
+    fixed offsets: data written before a commit survives it and newly
+    committed blocks read as zero."""
+    config, spec = _make_config(layout)
+    device = torch.device("cuda")
+    kv_cache = ExtensibleKVCache(config, device)
+    try:
+        kv_caches = allocate_kv_cache(
+            config, device, layout, allocate=kv_cache.allocate
+        )
+        assert len(kv_caches) == NUM_LAYERS
+        views = list(kv_caches.values())
+        for view in views:
+            assert view.shape[0] == NUM_BLOCKS
+            assert view.untyped_storage().nbytes() == kv_cache.size
+
+        kv_cache.commit(2)
+        assert kv_cache.num_committed_blocks == 2
+        for i, view in enumerate(views):
+            view[:2].fill_(float(i + 1))
+        torch.accelerator.synchronize()
+
+        kv_cache.commit(NUM_BLOCKS)
+        assert kv_cache.num_committed_blocks == NUM_BLOCKS
+        assert kv_cache.physical_bytes >= kv_cache.size
+        for i, view in enumerate(views):
+            assert torch.all(view[:2] == float(i + 1))
+            assert torch.count_nonzero(view[2:]) == 0
+
+        # Grow-only and clamped to capacity.
+        kv_cache.commit(NUM_BLOCKS + 5)
+        assert kv_cache.num_committed_blocks == NUM_BLOCKS
+    finally:
+        kv_cache.free()
+
+
+def test_measure_kv_cache_bytes():
+    gib = 1 << 30
+    # 10 GiB were free at startup; 3 GiB of non-KV memory is now resident and
+    # 1 GiB is committed to the KV cache, leaving 6 GiB free. Budget is 9 GiB.
+    common = dict(
+        init_free_memory=10 * gib,
+        free_memory=6 * gib,
+        committed_bytes=1 * gib,
+        margin_floor_bytes=0,
+        margin_fraction=0.0,
+    )
+    # Budget-bound: 9 - 3 = 6 GiB for the KV cache.
+    assert measure_kv_cache_bytes(requested_memory=9 * gib, **common) == 6 * gib
+    # Headroom-bound: at most what is free plus what is already committed.
+    assert measure_kv_cache_bytes(requested_memory=20 * gib, **common) == 7 * gib
+    # The measured transient peak is kept free in full, on either bound.
+    assert (
+        measure_kv_cache_bytes(
+            requested_memory=20 * gib, transient_peak_bytes=gib, **common
+        )
+        == 6 * gib
+    )
+    assert (
+        measure_kv_cache_bytes(
+            requested_memory=9 * gib, transient_peak_bytes=gib, **common
+        )
+        == 5 * gib
+    )
+    # The margin is a share of that peak with a floor, taken off the headroom
+    # (floor or fraction, whichever is larger), never off an explicit budget
+    # that already leaves it free.
+    assert (
+        measure_kv_cache_bytes(
+            requested_memory=20 * gib,
+            transient_peak_bytes=gib,
+            **{**common, "margin_floor_bytes": gib},
+        )
+        == 5 * gib
+    )
+    assert (
+        measure_kv_cache_bytes(
+            requested_memory=20 * gib,
+            transient_peak_bytes=2 * gib,
+            **{**common, "margin_fraction": 0.5},
+        )
+        == 4 * gib
+    )
+    assert (
+        measure_kv_cache_bytes(
+            requested_memory=9 * gib, **{**common, "margin_floor_bytes": gib}
+        )
+        == 6 * gib
+    )
+    # A small explicit budget (0.5 GiB) survives a floor larger than itself.
+    assert (
+        measure_kv_cache_bytes(
+            requested_memory=3 * gib + gib // 2, **{**common, "margin_floor_bytes": gib}
+        )
+        == gib // 2
+    )
+    # The result never goes negative.
+    assert measure_kv_cache_bytes(requested_memory=2 * gib, **common) == 0
+
+
+def test_measured_sizing_caps_torch_memory(monkeypatch):
+    """Torch's caching allocator may grow into what stays free once the cache
+    is extended, less multimodal IPC memory and the non-torch reserve,
+    rather than caching fragmented segments until other allocations fail."""
+    gib = 1 << 30
+    limits: list[int | None] = []
+    kv_cache = SimpleNamespace(
+        reserved_headroom_bytes=gib,
+        physical_bytes=gib,
+        bytes_per_block=gib,
+        blocks_within=lambda num_bytes, aligned: num_bytes // gib,
+        physical_bytes_for=lambda num_blocks: num_blocks * gib,
+        set_torch_memory_limit=limits.append,
+    )
+    accelerator = extensible_kv_cache.torch.accelerator
+    monkeypatch.setattr(accelerator, "synchronize", lambda *args: None)
+    monkeypatch.setattr(accelerator, "empty_cache", lambda: None)
+    monkeypatch.setattr(accelerator, "get_memory_info", lambda *args: (6 * gib, 0))
+    monkeypatch.setattr(accelerator, "memory_reserved", lambda *args: 2 * gib)
+    monkeypatch.setattr(
+        extensible_kv_cache, "measure_kv_cache_bytes", lambda **kwargs: 5 * gib
+    )
+    # One GiB of the 5 GiB budget is held for multimodal IPC.
+    monkeypatch.setattr(
+        extensible_kv_cache,
+        "reserve_mm_ipc_gpu_memory",
+        lambda available, *args: available - gib,
+    )
+    vllm_config = SimpleNamespace(
+        model_config=SimpleNamespace(multimodal_config=None),
+        parallel_config=SimpleNamespace(),
+        kv_transfer_config=None,
+    )
+
+    num_blocks = extensible_kv_cache.measure_kv_cache_blocks(
+        kv_cache,  # type: ignore[arg-type]
+        vllm_config,  # type: ignore[arg-type]
+        torch.device("cpu"),
+        init_free_memory=10 * gib,
+        requested_memory=10 * gib,
+    )
+
+    assert num_blocks == 4
+    # 6 GiB free, of which extending to 4 blocks takes 3 and IPC keeps 1.
+    left_free = 2 * gib - extensible_kv_cache.KV_CACHE_NON_TORCH_RESERVE_BYTES
+    assert limits == [2 * gib + left_free]
+
+
+def test_freeing_the_cache_lifts_the_torch_memory_cap(monkeypatch):
+    """The process may start another engine, which must not inherit the cap."""
+    fractions: list[float] = []
+    monkeypatch.setattr(
+        torch.accelerator, "get_memory_info", lambda *args: (0, 100 << 30)
+    )
+    monkeypatch.setattr(
+        torch.cuda,
+        "set_per_process_memory_fraction",
+        lambda fraction, device: fractions.append(fraction),
+    )
+    kv_cache = object.__new__(ExtensibleKVCache)
+    kv_cache.buffer = SimpleNamespace(device=torch.device("cpu"), free=lambda: None)
+    kv_cache.num_committed_blocks = 2
+    kv_cache.set_torch_memory_limit(25 << 30)
+    kv_cache.free()
+    assert fractions == [0.25, 1.0]
+    assert kv_cache.torch_memory_limit is None
+
+
+def test_granule_aligned_blocks():
+    """Blocks per stride must span whole granules; with several strides the
+    count is a multiple of every stride's requirement."""
+    granule = 2 << 20
+    # 64 KiB blocks: 32 per granule.
+    assert granule_block_alignment([64 << 10], granule) == 32
+    assert granule_aligned_blocks(1000, [64 << 10], granule) == 992
+    # 96 KiB blocks: 3 blocks cover 288 KiB; the granule is reached at 64.
+    assert granule_block_alignment([96 << 10], granule) == 64
+    # Mixed strides: lcm(32, 64).
+    assert granule_aligned_blocks(1000, [64 << 10, 96 << 10], granule) == 960
+    # A stride that is a granule multiple needs no alignment.
+    assert granule_aligned_blocks(1000, [4 << 20], granule) == 1000
+    assert granule_aligned_blocks(31, [64 << 10], granule) == 0
+
+
+def _make_mixed_config(num_blocks: int = NUM_BLOCKS):
+    """One group whose layers differ in head size, so their block strides differ."""
+    specs = {
+        f"layer.{i}": FullAttentionSpec(
+            block_size=16,
+            num_kv_heads=2,
+            head_size=64 * (1 + i % 2),
+            dtype=torch.float16,
+        )
+        for i in range(NUM_LAYERS)
+    }
+    uniform = UniformTypeKVCacheSpecs.from_specs(specs)
+    assert uniform is not None
+    vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=16))
+    vllm_config.cache_config.kv_cache_layout = "LBNHC"
+    config = get_kv_cache_config_from_groups(
+        vllm_config,
+        [KVCacheGroupSpec(list(specs), uniform)],
+        available_memory=sum(s.page_size_bytes for s in specs.values()) * num_blocks,
+    )
+    assert config.num_blocks == num_blocks
+    assert len({t.block_stride for t in config.kv_cache_tensors}) == 2
+    return config, specs
+
+
+@requires_cuda
+def test_mixed_block_strides_commit_per_layer(vmm):
+    """Layers with different page sizes get segments of their own stride; a
+    commit backs the same block count in each and the views stay valid."""
+    config, specs = _make_mixed_config()
+    device = torch.device("cuda")
+    kv_cache = ExtensibleKVCache(config, device)
+    try:
+        assert len(kv_cache.segment_strides) == NUM_LAYERS
+        assert sorted(kv_cache.segment_strides) == sorted(
+            s.page_size_bytes for s in specs.values()
+        )
+        views = allocate_kv_cache(
+            config, device, KVCacheLayout.LBNHC, allocate=kv_cache.allocate
+        )
+        kv_cache.commit(3)
+        for i, view in enumerate(views.values()):
+            view[:3].fill_(float(i + 1))
+        torch.accelerator.synchronize()
+        assert kv_cache.physical_bytes >= 3 * kv_cache.bytes_per_block
+        for i, view in enumerate(views.values()):
+            assert torch.all(view[:3] == float(i + 1))
+
+        final = copy.copy(config)
+        final.num_blocks = 3
+        committed = kv_cache.committed_views(final, KVCacheLayout.LBNHC)
+        assert committed is not None
+        for name, view in committed.items():
+            assert view.shape[0] == 3
+            assert view.data_ptr() == views[name].data_ptr()
+            assert view.untyped_storage().nbytes() == 3 * specs[name].page_size_bytes
+    finally:
+        kv_cache.free()
+
+
+@requires_cuda
+def test_release_and_recommit_for_sleep(vmm):
+    """Sleep drops the pages; wake maps fresh zeroed pages for the same block
+    count under the same addresses, so existing views keep working."""
+    config, _ = _make_config(KVCacheLayout.LBNHC)
+    device = torch.device("cuda")
+    kv_cache = ExtensibleKVCache(config, device)
+    try:
+        views = list(
+            allocate_kv_cache(
+                config, device, KVCacheLayout.LBNHC, allocate=kv_cache.allocate
+            ).values()
+        )
+        kv_cache.commit(NUM_BLOCKS)
+        views[0].fill_(3.0)
+        torch.accelerator.synchronize()
+
+        kv_cache.release_physical()
+        assert kv_cache.physical_bytes == 0
+        assert kv_cache.num_committed_blocks == NUM_BLOCKS
+
+        kv_cache.recommit()
+        assert kv_cache.physical_bytes >= kv_cache.size
+        assert torch.count_nonzero(views[0]) == 0
+    finally:
+        kv_cache.free()
+
+
+@requires_cuda
+@pytest.mark.parametrize("layout", [KVCacheLayout.LBNHC, KVCacheLayout.BLNHC])
+def test_shrinking_commit_remaps_to_the_smaller_prefix(vmm, layout):
+    """When fewer blocks fit than warmup committed, the final commit remaps
+    only that prefix (contents are warmup garbage), so the physical footprint,
+    the views and the connector-facing placements all describe the same,
+    smaller cache."""
+    config, _ = _make_config(layout)
+    device = torch.device("cuda")
+    kv_cache = ExtensibleKVCache(config, device)
+    try:
+        capacity_views = allocate_kv_cache(
+            config, device, layout, allocate=kv_cache.allocate
+        )
+        kv_cache.commit(6)
+        for view in capacity_views.values():
+            view[:6].fill_(1.0)
+        torch.accelerator.synchronize()
+        # Without `shrink` a smaller request is a no-op.
+        kv_cache.commit(4)
+        assert kv_cache.num_committed_blocks == 6
+
+        kv_cache.commit(4, shrink=True)
+        assert kv_cache.num_committed_blocks == 4
+        granule = kv_cache.buffer.granularity
+        expected_physical = sum(
+            -(-4 * stride // granule) * granule for stride in kv_cache.segment_strides
+        )
+        assert kv_cache.physical_bytes <= expected_physical
+        final = copy.copy(config)
+        final.num_blocks = 4
+        views = kv_cache.committed_views(final, layout)
+        assert views is not None
+        for name, view in views.items():
+            assert view.shape[0] == 4
+            assert view.data_ptr() == capacity_views[name].data_ptr()
+            assert torch.count_nonzero(view) == 0
+        tensors = kv_cache.committed_kv_cache_tensors(final, layout, 4)
+        assert all(t.size == 4 * kv_cache.bytes_per_block for t in tensors)
+    finally:
+        kv_cache.free()
+
+
+@requires_cuda
+@pytest.mark.parametrize("layout", [KVCacheLayout.LBNHC, KVCacheLayout.BLNHC])
+def test_committed_views_cover_exactly_the_committed_bytes(vmm, layout):
+    """After the final commit, per-layer views are rebuilt over storages that
+    span only the committed blocks, at the same addresses as the capacity
+    views, so connectors deriving extents from storage size see backed memory."""
+    config, spec = _make_config(layout)
+    device = torch.device("cuda")
+    kv_cache = ExtensibleKVCache(config, device)
+    try:
+        capacity_views = allocate_kv_cache(
+            config, device, layout, allocate=kv_cache.allocate
+        )
+        committed = 5
+        kv_cache.commit(committed)
+        for i, view in enumerate(capacity_views.values()):
+            view[:committed].fill_(float(i + 1))
+        torch.accelerator.synchronize()
+
+        final_config = copy.copy(config)
+        final_config.num_blocks = committed
+        views = kv_cache.committed_views(final_config, layout)
+        assert views is not None and views.keys() == capacity_views.keys()
+        for i, (name, view) in enumerate(views.items()):
+            assert view.shape[0] == committed
+            assert view.data_ptr() == capacity_views[name].data_ptr()
+            assert torch.all(view == float(i + 1))
+            block_stride = view.stride(0) * view.element_size()
+            if layout.is_block_outermost:
+                assert view.untyped_storage().nbytes() == committed * block_stride
+            else:
+                assert (
+                    view.untyped_storage().nbytes() == committed * spec.page_size_bytes
+                )
+    finally:
+        kv_cache.free()
+
+
+def _segments_backed_by_one_chunk(kv_cache: ExtensibleKVCache) -> list[bool]:
+    """Whether each segment's committed bytes lie within one driver allocation."""
+    return kv_cache.buffer.segments_backed_by_one_chunk(
+        [kv_cache.num_committed_blocks * s for s in kv_cache.segment_strides]
+    )
+
+
+def _granule(layout: KVCacheLayout) -> int:
+    probe = ExtensibleKVCache(_make_config(layout)[0], torch.device("cuda"))
+    try:
+        return probe.buffer.granularity
+    finally:
+        probe.free()
+
+
+@requires_cuda
+def test_defragmenting_commit_backs_each_segment_with_one_allocation(vmm):
+    """Incremental commits map one driver allocation each; a defragmenting
+    commit of a granule-aligned count remaps every segment as a single
+    allocation (RDMA cannot span several) and still ends up zero-filled."""
+    granule = _granule(KVCacheLayout.LBNHC)
+    stride = _make_config(KVCacheLayout.LBNHC)[1].page_size_bytes
+    alignment = granule_block_alignment([stride], granule)
+    # An unaligned capacity puts segment offsets inside granules.
+    capacity = 2 * alignment + alignment // 2
+    config, _ = _make_config(KVCacheLayout.LBNHC, num_blocks=capacity)
+    kv_cache = ExtensibleKVCache(config, torch.device("cuda"))
+    try:
+        assert kv_cache.block_alignment == alignment
+        kv_cache.commit(1)
+        kv_cache.commit(alignment + 1)
+        num_segments = kv_cache.buffer.num_segments
+        assert kv_cache.buffer.num_physical_chunks == 2 * num_segments
+
+        kv_cache.commit(capacity, defragment=True)
+        assert kv_cache.num_committed_blocks == capacity
+        assert torch.count_nonzero(kv_cache.buffer.full_view()) == 0
+        # Segments after the first start mid-granule and share chunks.
+        assert not all(_segments_backed_by_one_chunk(kv_cache))
+    finally:
+        kv_cache.free()
+
+    config, _ = _make_config(KVCacheLayout.LBNHC, num_blocks=2 * alignment)
+    kv_cache = ExtensibleKVCache(config, torch.device("cuda"))
+    try:
+        kv_cache.commit(alignment + 1)
+        kv_cache.commit(
+            kv_cache.blocks_within(kv_cache.size, aligned=True), defragment=True
+        )
+        assert kv_cache.num_committed_blocks == 2 * alignment
+        buffer = kv_cache.buffer
+        assert buffer.num_physical_chunks == buffer.num_segments
+        assert all(_segments_backed_by_one_chunk(kv_cache))
+    finally:
+        kv_cache.free()
+
+
+@requires_cuda
+@pytest.mark.parametrize("layout", [KVCacheLayout.LBNHC, KVCacheLayout.BLNHC])
+def test_aligned_capacity_config_allocates_and_defragments(vmm, layout):
+    """The engine-side alignment yields a config the cache can be built from,
+    with views over it, and whose full defragmenting commit backs each segment
+    with one allocation."""
+    from vllm.v1.core.kv_cache_planning import align_extensible_kv_cache_capacity
+    from vllm.v1.core.kv_cache_utils import generate_scheduler_kv_cache_config
+
+    granule = _granule(layout)
+    config, spec = _make_config(layout, num_blocks=1000)
+    vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=16))
+    vllm_config.cache_config.kv_cache_layout = layout.name
+    scheduler_config = generate_scheduler_kv_cache_config([config])
+    align_extensible_kv_cache_capacity(vllm_config, [config], scheduler_config, granule)
+    assert config.num_blocks < 1000
+    device = torch.device("cuda")
+    kv_cache = ExtensibleKVCache(config, device)
+    try:
+        views = allocate_kv_cache(config, device, layout, allocate=kv_cache.allocate)
+        assert all(v.shape[0] == config.num_blocks for v in views.values())
+        kv_cache.commit(3)
+        kv_cache.commit(config.num_blocks, defragment=True)
+        assert all(_segments_backed_by_one_chunk(kv_cache))
+        for view in views.values():
+            view.fill_(1.0)
+        torch.accelerator.synchronize()
+    finally:
+        kv_cache.free()
+
+
+@requires_cuda
+def test_blocks_within_accounts_for_granule_rounding(vmm):
+    """The count that fits a byte budget is the most whose rounded-up
+    per-segment footprint fits, not the budget over the nominal block size;
+    aligned, it is the largest aligned count under the nominal footprint."""
+    config, spec = _make_config(KVCacheLayout.LBNHC, num_blocks=4096)
+    kv_cache = ExtensibleKVCache(config, torch.device("cuda"))
+    try:
+        granule = kv_cache.buffer.granularity
+        alignment = kv_cache.block_alignment
+        num_segments = kv_cache.buffer.num_segments
+        # One granule per segment holds `alignment` blocks exactly; one more
+        # block per segment needs a whole extra granule each.
+        budget = num_segments * granule
+        assert kv_cache.blocks_within(budget) == alignment
+        assert kv_cache.blocks_within(budget + 1) == alignment
+        assert kv_cache.physical_bytes_for(alignment + 1) == 2 * budget
+        assert kv_cache.blocks_within(2 * budget - 1) == alignment
+        assert kv_cache.blocks_within(2 * budget) == 2 * alignment
+        nominal = (2 * alignment + 5) * kv_cache.bytes_per_block
+        assert kv_cache.blocks_within(nominal, aligned=True) == 2 * alignment
+        assert kv_cache.blocks_within(kv_cache.size * 2) == config.num_blocks
+    finally:
+        kv_cache.free()
+
+
+@requires_cuda
+@pytest.mark.parametrize("layout", [KVCacheLayout.LBNHC, KVCacheLayout.BLNHC])
+def test_committed_kv_cache_tensors_match_views(vmm, layout):
+    """The connector-facing placements locate every block of every layer at
+    the same bytes as the committed views, within storages of their size."""
+    config, _ = _make_config(layout)
+    device = torch.device("cuda")
+    kv_cache = ExtensibleKVCache(config, device)
+    try:
+        allocate_kv_cache(config, device, layout, allocate=kv_cache.allocate)
+        committed = 5
+        kv_cache.commit(committed)
+        final = copy.copy(config)
+        final.num_blocks = committed
+        views = kv_cache.committed_views(final, layout)
+        tensors = kv_cache.committed_kv_cache_tensors(final, layout, committed)
+        assert views is not None
+        assert {name for t in tensors for name in t.layers} == set(views)
+        for tensor in tensors:
+            assert isinstance(tensor, KVCacheTensor)
+            # Consumers read any tensor's size as the whole allocation's.
+            assert tensor.size == committed * kv_cache.bytes_per_block
+            for layer_idx, name in enumerate(tensor.layers):
+                view = views[name]
+                storage = view.untyped_storage()
+                layer_start = tensor.offset + layer_idx * tensor.layer_stride
+                page_bytes = view[0].numel() * view.element_size()
+                last_page_end = (
+                    layer_start + (committed - 1) * tensor.block_stride + page_bytes
+                )
+                assert storage.nbytes() >= last_page_end
+                assert storage.data_ptr() + layer_start == view.data_ptr()
+                assert view.stride(0) * view.element_size() == tensor.block_stride
+    finally:
+        kv_cache.free()
diff --git a/tests/v1/worker/test_gpu_batch_shard.py b/tests/v1/worker/test_gpu_batch_shard.py
index 4ee76818df4f..0548a9c72fe0 100644
--- a/tests/v1/worker/test_gpu_batch_shard.py
+++ b/tests/v1/worker/test_gpu_batch_shard.py
@@ -23,6 +23,7 @@
     BatchSharder,
     _shard_grammar_output,
 )
+from vllm.v1.worker.gpu.spec_decode import draft_logits as draft_logits_module
 
 DEVICE = current_platform.device_type
 
@@ -43,13 +44,10 @@ def _make_batch(
     num_reqs: int,
     max_num_reqs: int,
     max_spec: int,
-    slot_pool: np.ndarray | None = None,
 ) -> InputBatch:
     """Consistent batch with random unique request slots, 1 + k_i logits
     rows per request at the tail of each query segment."""
-    if slot_pool is None:
-        slot_pool = np.arange(max_num_reqs)
-    slots = rng.choice(slot_pool, size=num_reqs, replace=False)
+    slots = rng.choice(max_num_reqs, size=num_reqs, replace=False)
     slots = slots.astype(np.int32)
     k = rng.integers(0, max_spec + 1, size=num_reqs)
     num_logits_per_req = (1 + k).astype(np.int32)
@@ -116,8 +114,21 @@ def _shard_all_ranks(
     max_num_reqs: int,
     tp_size: int,
     adaptive_verification_enabled: bool = False,
+    owners_np: np.ndarray | None = None,
 ):
-    """Simulate every TP rank by faking get_tp_group() per BatchSharder."""
+    """Simulate every TP rank by faking get_tp_group() per BatchSharder.
+
+    ``owners_np`` is a sharded DraftLogitsCache owner table indexed by request
+    slot. Without it, requests are owned by batch position.
+    """
+    draft_logits = None
+    if owners_np is not None:
+        owners_np = owners_np.astype(np.int32)
+        owners = torch.from_numpy(owners_np).to(DEVICE)
+        # BatchSharder reads only the owner table of a sharded cache.
+        draft_logits = SimpleNamespace(
+            owners_of=lambda idx_mapping_np: (owners_np[idx_mapping_np], owners)
+        )
     results = []
     for rank in range(tp_size):
         fake_group = SimpleNamespace(rank_in_group=rank, world_size=tp_size)
@@ -127,6 +138,7 @@ def _shard_all_ranks(
                 max_num_logits_per_req=8,
                 device=torch.device(DEVICE),
                 adaptive_verification_enabled=adaptive_verification_enabled,
+                draft_logits=draft_logits,
             )
         local_batch, sorted_logits_indices, _, metadata = sharder.shard_sampler_inputs(
             batch, None
@@ -215,7 +227,7 @@ def test_local_batch_partition(tp_size: int, seed: int):
     for rank, (local, _, metadata) in enumerate(results):
         owned = _owned_batch_indices(local)
         assert local.num_reqs == len(owned) == metadata.num_local_reqs
-        assert (batch.idx_mapping_np[owned] % tp_size == rank).all()
+        assert (owned % tp_size == rank).all()
         # Stable sort: batch order is preserved within an owner.
         assert (np.diff(owned) > 0).all()
         all_owned.extend(owned.tolist())
@@ -311,7 +323,7 @@ def test_all_to_all_reassembly_simulated(tp_size: int):
             [sends[src][bounds[rank] : bounds[rank + 1]] for src in range(tp_size)]
         )
         # Reassemble source-rank-major chunks into full-vocab rows, mirroring
-        # all_to_all_logits (source-rank order is vocab-shard order).
+        # all_to_all_vocab_shards (source-rank order is vocab-shard order).
         reassembled = (
             recv.view(tp_size, metadata.num_local_logits, shard_w)
             .permute(1, 0, 2)
@@ -398,15 +410,10 @@ def test_gather_src_indices_round_trip(tp_size: int):
     batch order, including when some rank owns zero requests."""
     rng = np.random.default_rng(3)
     max_num_reqs = 40
-    all_slots = np.arange(max_num_reqs)
-    batch = _make_batch(
-        rng,
-        num_reqs=17,
-        max_num_reqs=max_num_reqs,
-        max_spec=2,
-        slot_pool=all_slots[all_slots % tp_size != 0],  # rank 0 owns nothing
-    )
-    results = _shard_all_ranks(batch, max_num_reqs, tp_size)
+    batch = _make_batch(rng, num_reqs=17, max_num_reqs=max_num_reqs, max_spec=2)
+    # An owner table under which rank 0 owns nothing.
+    owners_np = 1 + np.arange(max_num_reqs) % (tp_size - 1)
+    results = _shard_all_ranks(batch, max_num_reqs, tp_size, owners_np=owners_np)
     metadata0 = results[0][2]
     assert metadata0.num_local_reqs == 0
 
@@ -522,7 +529,10 @@ def fake_all_gather(x: torch.Tensor, dim: int = 0) -> torch.Tensor:
 
 
 @requires_gpu
-def test_gather_sampler_output_logprobs_and_nans():
+# (num_logprobs, max_token_ids): top-k only, `logprob_token_ids` only (what a
+# request gets from a speculative-decoding batch), and token ids wider than top-k.
+@pytest.mark.parametrize("logprobs_dims", [(2, 0), (0, 2), (1, 2)])
+def test_gather_sampler_output_logprobs_and_nans(logprobs_dims: tuple[int, int]):
     """num_nans is reduced per request and gathered; LogprobsTensors are
     packed into fixed per-request blocks, width-aligned across ranks (pad
     and truncate paths), and restored to global row order."""
@@ -631,7 +641,7 @@ def fake_all_gather(x: torch.Tensor, dim: int = 0) -> torch.Tensor:
                 global_batch=global_cuda,
                 local_batch=local_batches[rank],
                 gather_num_nans=True,
-                logprobs_dims=(num_cols - 1, 0),
+                logprobs_dims=logprobs_dims,
             )
 
     for rank in range(tp_size):
@@ -670,3 +680,148 @@ def test_finish_requests_frees_slots_in_sorted_order():
     GPUModelRunner.finish_requests(fake_runner, scheduler_output)
     assert removed == sorted(removed)
     assert set(removed) == {"req-a", "req-b", "req-c", "req-d", "req-e"}
+
+
+def _make_sharded_draft_logits(tp_size: int, max_num_reqs: int) -> list:
+    """One sharded DraftLogitsCache per simulated TP rank."""
+    caches = []
+    for rank in range(tp_size):
+        fake_group = SimpleNamespace(rank_in_group=rank, world_size=tp_size)
+        with mock.patch.object(
+            draft_logits_module, "get_tp_group", return_value=fake_group
+        ):
+            caches.append(
+                draft_logits_module.DraftLogitsCache(
+                    max_num_reqs,
+                    1,
+                    1,
+                    dtype=torch.float32,
+                    fill=0.0,
+                    device=torch.device(DEVICE),
+                    sharded=True,
+                )
+            )
+    return caches
+
+
+def _draft_batch(slots: list[int], num_logits_per_req: int) -> SimpleNamespace:
+    """The fields of an InputBatch that DraftLogitsCache reads."""
+    slots_np = np.array(slots, dtype=np.int32)
+    return SimpleNamespace(
+        idx_mapping=torch.from_numpy(slots_np).to(DEVICE),
+        idx_mapping_np=slots_np,
+        expanded_idx_mapping=torch.from_numpy(
+            np.repeat(slots_np, num_logits_per_req)
+        ).to(DEVICE),
+        expanded_local_pos=torch.arange(num_logits_per_req, device=DEVICE).repeat(
+            len(slots)
+        ),
+    )
+
+
+@requires_gpu
+@pytest.mark.parametrize("tp_size", [2, 4])
+def test_sharded_draft_logits_balance_and_stale_rows(tp_size: int):
+    """Owners follow batch position, every rank caches at most
+    cdiv(num_reqs, tp_size) requests, and a request skipped while another
+    request took its row has stale drafts."""
+    max_num_reqs = 64
+    num_logits_per_req = 3
+    owners = _make_sharded_draft_logits(tp_size, max_num_reqs)
+    # Skewed slots: under slot % tp_size, rank 0 would own every request.
+    step0 = [s * tp_size for s in range(2 * tp_size)]
+    batch0 = _draft_batch(step0, num_logits_per_req)
+    for o in owners:
+        o.assign(batch0.idx_mapping, batch0.idx_mapping_np)
+
+    for rank, o in enumerate(owners):
+        assert o.shard_mapping_np[step0].tolist() == [
+            i % tp_size for i in range(len(step0))
+        ]
+        rows = _np(o.row_mapping)[step0]
+        owned = [i for i in range(len(step0)) if i % tp_size == rank]
+        assert rows[owned].tolist() == [i // tp_size for i in owned]
+        assert (np.delete(rows, owned) == -1).all()
+        assert _np(o.inverse_row_mapping)[: len(owned)].tolist() == [
+            step0[i] for i in owned
+        ]
+
+    # Every drafted request verifies against its own row.
+    for rank, o in enumerate(owners):
+        owned = [s for s in step0 if o.shard_mapping_np[s] == rank]
+        stale = _np(o.stale_drafts(_draft_batch(owned, num_logits_per_req)))
+        assert not stale.any()
+
+    # Step 1 skips the request at position 1 and puts a new request there.
+    skipped, new = step0[1], 63
+    step1 = [step0[0], new] + step0[2:]
+    batch1 = _draft_batch(step1, num_logits_per_req)
+    for o in owners:
+        o.assign(batch1.idx_mapping, batch1.idx_mapping_np)
+
+    # The skipped request keeps its owner, whose row now holds `new`.
+    owner_of_skipped = 1 % tp_size
+    o = owners[owner_of_skipped]
+    assert o.shard_mapping_np[skipped] == owner_of_skipped
+    stale = _np(o.stale_drafts(_draft_batch([skipped, new], num_logits_per_req)))
+    stale = stale.reshape(2, num_logits_per_req)
+    # Only draft rows (local position > 0) of the skipped request are stale.
+    assert stale[0].tolist() == [False] + [True] * (num_logits_per_req - 1)
+    assert not stale[1].any()
+
+
+@pytest.mark.parametrize("pcp_size", [1, 2])
+def test_draft_logits_stay_replicated_under_pcp(pcp_size: int):
+    """PCPManager.sync_drafts indexes the draft logits cache by request state,
+    so PCP keeps the cache replicated even under batch-sharded sampling."""
+    from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator
+
+    speculator = SimpleNamespace(
+        vllm_config=SimpleNamespace(
+            parallel_config=SimpleNamespace(prefill_context_parallel_size=pcp_size)
+        ),
+        max_num_reqs=8,
+        num_speculative_steps=2,
+        vocab_size=4,
+        device=torch.device("cpu"),
+        draft_watermarker=None,
+        draft_logits_spec=lambda _: (torch.float32, 0.0),
+        supports_batch_sharded_draft_logits=lambda: True,
+    )
+    fake_group = SimpleNamespace(rank_in_group=0, world_size=2)
+    with (
+        mock.patch(
+            "vllm.v1.worker.gpu.spec_decode.speculator.batch_sharded_sampling_enabled",
+            return_value=True,
+        ),
+        mock.patch.object(draft_logits_module, "get_tp_group", return_value=fake_group),
+    ):
+        DraftModelSpeculator._init_draft_logits(speculator, target_model=None)
+
+    draft_logits = speculator.draft_logits
+    assert (draft_logits.row_mapping is not None) == (pcp_size == 1)
+    assert draft_logits.values.shape[0] == (4 if pcp_size == 1 else 8)
+
+
+def test_owners_of_reads_owner_table_or_batch_position():
+    """A sharded cache owns each request at the rank its owner table records;
+    a replicated cache leaves requests round-robin over batch positions."""
+    fake_group = SimpleNamespace(rank_in_group=1, world_size=4)
+    idx_mapping_np = np.array([6, 1, 3, 4, 0], dtype=np.int32)
+    with mock.patch.object(
+        draft_logits_module, "get_tp_group", return_value=fake_group
+    ):
+        sharded, replicated = (
+            draft_logits_module.DraftLogitsCache(
+                8, 1, 1, dtype=torch.float32, fill=0.0, device="cpu", sharded=sharded
+            )
+            for sharded in (True, False)
+        )
+        sharded_owners_np, sharded_owners = sharded.owners_of(idx_mapping_np)
+        replicated_owners_np, replicated_owners = replicated.owners_of(idx_mapping_np)
+
+    # Before any assign(), a sharded cache owns each request by its slot.
+    assert sharded_owners_np.tolist() == (idx_mapping_np % 4).tolist()
+    assert sharded_owners is sharded.shard_mapping
+    assert replicated_owners_np.tolist() == [0, 1, 2, 3, 0]
+    assert replicated_owners is None
diff --git a/tests/v1/worker/test_gpu_diffusion_sampler.py b/tests/v1/worker/test_gpu_diffusion_sampler.py
new file mode 100644
index 000000000000..17ba78f4c7b2
--- /dev/null
+++ b/tests/v1/worker/test_gpu_diffusion_sampler.py
@@ -0,0 +1,144 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""DiffusionSampler projects logits tile-by-tile (#50699): sampling outputs
+must not depend on the tile size, and peak memory must scale with the tile,
+not the batch."""
+
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+import torch
+
+from vllm.model_executor.models import diffusion_gemma as dg
+from vllm.sampling_params import SamplingParams
+from vllm.v1.worker.gpu.sample.logprob import LogprobTokenIdsState
+from vllm.v1.worker.gpu.sample.states import SamplingStates
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
+
+DEVICE = "cuda"
+
+
+def _make_sampler(num_reqs, canvas_len, vocab, hidden, weight, top_k=0):
+    states = dg.DiffusionGemmaRequestStates(
+        max_num_reqs=num_reqs,
+        canvas_length=canvas_len,
+        vocab_size=vocab,
+        max_denoising_steps=8,
+        device=torch.device(DEVICE),
+        hidden_size=hidden,
+        stability_threshold=2,
+    )
+    base = SimpleNamespace(
+        sampling_states=SamplingStates(num_reqs, vocab),
+        logprob_token_ids_state=LogprobTokenIdsState(num_reqs, torch.device(DEVICE)),
+        req_states=SimpleNamespace(
+            draft_tokens=torch.zeros(
+                num_reqs, canvas_len, dtype=torch.int64, device=DEVICE
+            )
+        ),
+        logprobs_mode="raw_logprobs",
+    )
+    sampler = dg.DiffusionSampler(
+        sampler=base,
+        diffusion_config=SimpleNamespace(canvas_length=canvas_len),
+        vocab_size=vocab,
+        diffusion_states=states,
+        compute_logits=lambda h: h @ weight.t(),
+        confidence_threshold=0.5,
+        t_min=0.0,
+        t_max=0.0,
+        entropy_bound=0.1,
+        embed_weight=weight,
+        normalizer=torch.tensor(1.0, device=DEVICE),
+    )
+    for i in range(num_reqs):
+        sampler.add_request(i, SamplingParams(top_k=top_k))
+        states.is_encoder_phase[i] = False
+    sampler.apply_staged_writes()
+    return sampler, states
+
+
+def _make_batch(num_reqs, valid_lens):
+    cu = np.zeros(num_reqs + 1, dtype=np.int64)
+    cu[1:] = np.cumsum(valid_lens)
+    qsl_np = cu.astype(np.int32)
+    return SimpleNamespace(
+        num_reqs=num_reqs,
+        num_draft_tokens=int(max(valid_lens)),
+        idx_mapping_np=np.arange(num_reqs, dtype=np.intp),
+        idx_mapping=torch.arange(num_reqs, dtype=torch.int64, device=DEVICE),
+        cu_num_logits_np=cu,
+        query_start_loc_np=qsl_np,
+        query_start_loc=torch.from_numpy(qsl_np).to(DEVICE),
+    )
+
+
+@pytest.mark.parametrize("top_k", [0, 4])
+def test_tiled_projection_matches_full(monkeypatch, top_k):
+    """Deterministic outputs (t_min=t_max=0) are identical for group=1 tiling
+    and a single full-batch pass, including a truncated (padded) canvas."""
+    num_reqs, canvas_len, vocab, hidden = 4, 16, 1024, 32
+    torch.manual_seed(0)
+    weight = torch.randn(vocab, hidden, device=DEVICE)
+    valid_lens = [canvas_len, canvas_len, canvas_len, canvas_len - 3]
+    hs = torch.randn(sum(valid_lens), hidden, device=DEVICE)
+
+    results = []
+    for free in (2**60, 0):
+        torch.manual_seed(0)
+        sampler, states = _make_sampler(
+            num_reqs, canvas_len, vocab, hidden, weight, top_k
+        )
+        monkeypatch.setattr(
+            torch.accelerator, "get_memory_info", lambda free=free: (free, 2**60)
+        )
+        out = sampler(hs.clone(), _make_batch(num_reqs, valid_lens))
+        results.append((out, states))
+
+    (out_full, st_full), (out_tiled, st_tiled) = results
+    torch.testing.assert_close(st_full.argmax_canvas, st_tiled.argmax_canvas)
+    torch.testing.assert_close(st_full.is_encoder_phase, st_tiled.is_encoder_phase)
+    torch.testing.assert_close(st_full.confident, st_tiled.confident)
+    torch.testing.assert_close(st_full.step, st_tiled.step)
+    torch.testing.assert_close(out_full.sampled_token_ids, out_tiled.sampled_token_ids)
+    torch.testing.assert_close(out_full.num_sampled, out_tiled.num_sampled)
+
+
+def test_tiling_bounds_peak_memory(monkeypatch):
+    """Forced group=1 tiling must not materialize full-batch logits."""
+    num_reqs, canvas_len, vocab, hidden = 8, 64, 32768, 64
+    torch.manual_seed(0)
+    weight = torch.randn(vocab, hidden, device=DEVICE)
+    valid_lens = [canvas_len] * num_reqs
+    hs = torch.randn(num_reqs * canvas_len, hidden, device=DEVICE)
+
+    peaks = {}
+    for name, free in (("full", 2**60), ("tiled", 0)):
+        torch.manual_seed(0)
+        sampler, _ = _make_sampler(num_reqs, canvas_len, vocab, hidden, weight)
+        monkeypatch.setattr(
+            torch.accelerator, "get_memory_info", lambda free=free: (free, 2**60)
+        )
+        sampler(hs.clone(), _make_batch(num_reqs, valid_lens))
+        torch.accelerator.synchronize()
+        torch.accelerator.reset_peak_memory_stats()
+        base = torch.accelerator.memory_allocated()
+        sampler(hs.clone(), _make_batch(num_reqs, valid_lens))
+        torch.accelerator.synchronize()
+        peaks[name] = torch.accelerator.max_memory_allocated() - base
+
+    assert peaks["tiled"] < peaks["full"] / 2
+
+
+def test_profile_run_resets_state():
+    """profile_run exercises the decode pipeline for KV sizing, then leaves
+    slot 0 freshly initialized."""
+    num_reqs, canvas_len, vocab, hidden = 2, 16, 1024, 32
+    torch.manual_seed(0)
+    weight = torch.randn(vocab, hidden, device=DEVICE)
+    sampler, states = _make_sampler(num_reqs, canvas_len, vocab, hidden, weight)
+    sampler.profile_run(torch.zeros(num_reqs, hidden, device=DEVICE))
+    assert bool(states.is_encoder_phase[0])
+    assert int(states.step[0]) == 0
diff --git a/tests/v1/worker/test_gpu_draft_lm_head.py b/tests/v1/worker/test_gpu_draft_lm_head.py
index f0f0c72aca0d..7aa03babcb87 100644
--- a/tests/v1/worker/test_gpu_draft_lm_head.py
+++ b/tests/v1/worker/test_gpu_draft_lm_head.py
@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import gc
+from contextlib import ExitStack
 from types import SimpleNamespace
 
 import pytest
@@ -119,6 +120,9 @@ def _draft_update_worker(vllm_config, drafter: torch.nn.Module, **engine) -> Wor
     )
     worker.model_runner = SimpleNamespace(get_draft_model=lambda: drafter)
     worker._weight_update_active = False
+    worker._incomplete_weight_updates = set()
+    worker._kv_cache_reclaimable = False
+    worker._weight_update_kv_cache_loan = ExitStack()
     worker._set_draft_weight_update_target = lambda: None  # type: ignore[method-assign]
     return worker
 
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 5fa44d347569..3ab467625fce 100644
--- a/tests/v1/worker/test_gpu_extract_hidden_states_speculator.py
+++ b/tests/v1/worker/test_gpu_extract_hidden_states_speculator.py
@@ -7,6 +7,7 @@
 import pytest
 import torch
 
+from tests.v1.spec_decode.utils import fake_speculative_config
 from vllm.v1.worker.gpu.spec_decode import extract_hidden_states as spec_module
 from vllm.v1.worker.gpu.spec_decode import init_speculator
 from vllm.v1.worker.gpu.spec_decode.extract_hidden_states import (
@@ -35,18 +36,20 @@ def test_init_speculator_dispatches_extract_hidden_states(monkeypatch):
     vllm_config = cast(
         Any,
         SimpleNamespace(
-            speculative_config=SimpleNamespace(method="extract_hidden_states")
+            speculative_config=fake_speculative_config(method="extract_hidden_states")
         ),
     )
     device = torch.device("cpu")
 
-    def fake_speculator(config, target_device):
-        return config, target_device
+    class FakeSpeculator(ExtractHiddenStatesSpeculator):
+        def __init__(self, config, target_device):
+            self.args = (config, target_device)
 
-    monkeypatch.setattr(spec_module, "ExtractHiddenStatesSpeculator", fake_speculator)
+    monkeypatch.setattr(spec_module, "ExtractHiddenStatesSpeculator", FakeSpeculator)
 
     speculator = init_speculator(vllm_config, device, cast(Any, None))
-    assert speculator == (vllm_config, device)
+    assert isinstance(speculator, FakeSpeculator)
+    assert speculator.args == (vllm_config, device)
 
 
 def test_propose_caches_hidden_states_and_returns_sampled_tokens(monkeypatch):
diff --git a/tests/v1/worker/test_gpu_gumbel_sample.py b/tests/v1/worker/test_gpu_gumbel_sample.py
index b3c4128bfc73..2d507d26190d 100644
--- a/tests/v1/worker/test_gpu_gumbel_sample.py
+++ b/tests/v1/worker/test_gpu_gumbel_sample.py
@@ -24,6 +24,7 @@
     pytest.skip("CUDA required for Gumbel sampler tests", allow_module_level=True)
 
 from vllm.triton_utils import tl, triton
+from vllm.v1.worker.gpu.sample.greedy_argmax import greedy_argmax
 from vllm.v1.worker.gpu.sample.gumbel import (
     _uniform64_from_random53,
     gumbel_sample,
@@ -528,3 +529,21 @@ def test_logits_cache_narrower_than_logits_is_rejected():
             logits_cache=cache,
             logits_cache_col=torch.tensor(0, dtype=torch.int32, device=DEVICE),
         )
+
+
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
+@pytest.mark.parametrize("vocab", [4096, 4097, 129280])
+def test_greedy_argmax_matches_torch(dtype, vocab):
+    """Random rows, cross-block ties, NaN/inf, masked tails and strided rows."""
+    logits = torch.randn(64, vocab + 1, device=DEVICE, dtype=dtype)[:, 1:]
+    logits[0].zero_()
+    logits[1].fill_(float("-inf"))
+    logits[2, [5, -1]] = float("inf")
+    logits[3, [5, -1]] = float("nan")
+    logits[4].fill_(float("-inf"))
+    logits[4, -1] = 1
+    logits[5].round_().clamp_(max=2)
+    logits[5, :4096].clamp_(max=1)
+    # NaN is treated as -inf, as in the rejection sampler.
+    expected = logits.masked_fill(logits.isnan(), float("-inf")).argmax(-1)
+    torch.testing.assert_close(greedy_argmax(logits), expected, rtol=0, atol=0)
diff --git a/tests/v1/worker/test_gpu_model_runner.py b/tests/v1/worker/test_gpu_model_runner.py
index a86e3a8e9996..ecb05e72507f 100644
--- a/tests/v1/worker/test_gpu_model_runner.py
+++ b/tests/v1/worker/test_gpu_model_runner.py
@@ -1147,12 +1147,13 @@ def test_reload_weights_path_replaces_object_storage_source(monkeypatch):
     loader = Mock()
     loader.get_all_weights.return_value = iter(())
     monkeypatch.setattr(reload_inplace_module, "get_model_loader", lambda _: loader)
-    monkeypatch.setattr(reload_inplace_module, "initialize_layerwise_reload", Mock())
-    monkeypatch.setattr(reload_inplace_module, "finalize_layerwise_reload", Mock())
+    monkeypatch.setattr(reload_inplace_module, "start_reload", Mock())
+    monkeypatch.setattr(reload_inplace_module, "finish_reload", Mock())
 
     vllm_config = SimpleNamespace(
         lora_config=None,
-        load_config=SimpleNamespace(load_format="runai_streamer"),
+        load_config=SimpleNamespace(load_format="runai_streamer", device=None),
+        device_config=SimpleNamespace(device="cpu"),
         model_config=SimpleNamespace(
             model="/tmp/pulled-config-files",
             model_weights="s3://bucket/original",
diff --git a/tests/v1/worker/test_gpu_model_runner_v2.py b/tests/v1/worker/test_gpu_model_runner_v2.py
index 752107c2f851..2904c7faf46c 100644
--- a/tests/v1/worker/test_gpu_model_runner_v2.py
+++ b/tests/v1/worker/test_gpu_model_runner_v2.py
@@ -4,15 +4,26 @@
 import contextlib
 from collections.abc import Iterator
 from dataclasses import dataclass
+from functools import partial
 from types import SimpleNamespace
 from unittest.mock import Mock
 
 import pytest
 import torch
 
+import vllm.distributed.aux_output_connector.worker as aux_output_worker
 import vllm.v1.worker.gpu.model_runner as model_runner_module
-from vllm.config import CacheConfig, LoRAConfig, ParallelConfig
+from vllm.config import (
+    AttentionConfig,
+    CacheConfig,
+    KernelConfig,
+    LoadConfig,
+    LoRAConfig,
+    ParallelConfig,
+    SpeculativeConfig,
+)
 from vllm.model_executor.warmup.jit_warmup import JitWarmupRegistry
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
 from vllm.v1.kv_cache_interface import (
     CircularBufferSpec,
     FullAttentionSpec,
@@ -65,12 +76,69 @@ def test_non_last_pp_rank_uses_global_batch_for_sample_feedback(monkeypatch):
 
     runner.sample_tokens(None)
 
-    pp_receive.assert_called_once_with(global_batch)
+    pp_receive.assert_called_once_with(global_batch, recv_drafts=False)
     postprocess_num_computed_tokens.assert_called_once_with(global_batch)
     postprocess_state.assert_called_once_with(global_batch.idx_mapping, 0)
     restore_for_sampling.assert_not_called()
 
 
+@pytest.mark.parametrize(
+    "enabled,expected_events",
+    [
+        (True, ["allocate", "bind", "dummy_run"]),
+        (False, ["dummy_run"]),
+    ],
+    ids=["aux-output-enabled", "aux-output-disabled"],
+)
+def test_profile_run_creates_aux_output_connector_before_forward(
+    monkeypatch, enabled, expected_events
+):
+    events = []
+    capturer = object()
+    model = torch.nn.Module()
+    vllm_config = SimpleNamespace(
+        aux_output_config=SimpleNamespace(enabled=enabled),
+        max_concurrent_batches=1,
+        scheduler_config=SimpleNamespace(max_num_batched_tokens=8),
+    )
+
+    def allocate(**kwargs):
+        assert kwargs == {"max_num_batched_tokens": 8, "vllm_config": vllm_config}
+        events.append("allocate")
+        return capturer
+
+    def bind(bind_model, bind_capturer):
+        assert bind_model is model and bind_capturer is capturer
+        events.append("bind")
+
+    def dummy_run(*args, **kwargs):
+        events.append("dummy_run")
+        return None, None
+
+    monkeypatch.setattr(aux_output_worker, "RoutedExpertsCapturer", allocate)
+    monkeypatch.setattr(aux_output_worker, "bind_routed_experts_capturer", bind)
+    monkeypatch.setattr(torch.accelerator, "synchronize", lambda: None)
+    runner = GPUModelRunner.__new__(GPUModelRunner)
+    runner.vllm_config = vllm_config
+    runner.model = model
+    runner.aux_output_connector = None
+    runner.supports_mm_inputs = False
+    runner.is_last_pp_rank = False
+    runner.max_num_tokens = 8
+    runner.encoder_cache = None
+    runner.pooling_runner = None
+    monkeypatch.setattr(runner, "_dummy_run", dummy_run)
+
+    runner.profile_run()
+
+    assert events == expected_events
+    connector = runner.aux_output_connector
+    if enabled:
+        assert connector is not None and connector._capturer is capturer
+    else:
+        assert connector is None
+
+
 def test_qsa_circular_group_uses_custom_slot_mapping(monkeypatch):
     runner = GPUModelRunner.__new__(GPUModelRunner)
     runner.max_model_len = 262144
@@ -283,23 +351,24 @@ def _make_capture_runner(
     cudagraph_manager's needs_capture decision."""
     runner = GPUModelRunner.__new__(GPUModelRunner)
     runner.model_state = SimpleNamespace(
-        supports_mm_inputs=False, capture_inner_cudagraphs=lambda *args: None
+        supports_mm_inputs=False, capture_inner_cudagraphs=Mock()
     )
     runner.cudagraph_manager = SimpleNamespace(
         needs_capture=lambda: captured,
         capture=lambda *args, **kwargs: None,
+        capture_stream=object(),
     )
     runner.lora_config = None
     monkeypatch.setattr(runner, "maybe_setup_dummy_loras", _no_dummy_loras)
     runner.speculator = None
     runner.adaptive_verification = None
     runner.model = None
-    runner.input_buffers = None
+    runner.input_buffers = object()
     runner.pcp_manager = None
     runner.intermediate_tensors = None
-    runner.block_tables = None
-    runner.attn_groups = None
-    runner.kv_cache_config = None
+    runner.block_tables = object()
+    runner.attn_groups = object()
+    runner.kv_cache_config = object()
     runner.use_aux_hidden_state_outputs = False
     runner.kv_connector = model_runner_module.NO_OP_KV_CONNECTOR
     return runner
@@ -325,6 +394,16 @@ def test_capture_model_locks_workspace_after_capture(monkeypatch):
     runner.capture_model()
 
     assert lock_calls == ["lock"]
+    assert runner.cudagraph_manager is not None
+    capture_inner_cudagraphs = runner.model_state.capture_inner_cudagraphs
+    assert isinstance(capture_inner_cudagraphs, Mock)
+    capture_inner_cudagraphs.assert_called_once_with(
+        runner.input_buffers,
+        runner.block_tables,
+        runner.attn_groups,
+        runner.kv_cache_config,
+        runner.cudagraph_manager.capture_stream,
+    )
 
 
 def test_capture_model_skips_lock_when_nothing_captured(monkeypatch):
@@ -435,27 +514,48 @@ def test_async_copy_to_np_does_not_alias_reused_buffer():
     [("auto", "auto"), ("bfloat16", "bfloat16"), (None, "fp8_e4m3")],
 )
 @pytest.mark.parametrize("data_parallel_index", [0, 3])
+@pytest.mark.parametrize(
+    ("draft_moe", "draft_attn"),
+    [(None, None), ("triton", AttentionBackendEnum.TRITON_ATTN)],
+)
 def test_standalone_draft_kv_dtype_override(
-    monkeypatch, draft_dtype, expected_dtype, data_parallel_index
+    monkeypatch, draft_dtype, expected_dtype, data_parallel_index, draft_moe, draft_attn
 ):
     @dataclass
     class Config:
         cache_config: CacheConfig
         parallel_config: ParallelConfig
         model_config: object
+        kernel_config: KernelConfig
+        attention_config: AttentionConfig
         quant_config: object | None = None
 
-    target = Config(CacheConfig(cache_dtype="fp8_e4m3"), ParallelConfig(), object())
+    target = Config(
+        CacheConfig(cache_dtype="fp8_e4m3"),
+        ParallelConfig(),
+        object(),
+        KernelConfig(moe_backend="flashinfer_cutlass"),
+        AttentionConfig(backend=AttentionBackendEnum.FLASHINFER),
+    )
     target.parallel_config.data_parallel_index = data_parallel_index
     speculator = object.__new__(ar_module.StandaloneARSpeculator)
     speculator.vllm_config = target
     speculator.draft_model_config = object()
-    speculator.speculative_config = SimpleNamespace(
-        kv_cache_dtype=draft_dtype, draft_parallel_config=ParallelConfig()
-    )
+    spec = SimpleNamespace(
+        kv_cache_dtype=draft_dtype,
+        moe_backend=draft_moe,
+        attention_backend=draft_attn,
+        draft_parallel_config=ParallelConfig(),
+    )
+    spec.apply_draft_overrides = partial(SpeculativeConfig.apply_draft_overrides, spec)
+    speculator.speculative_config = spec
     get_model = Mock(return_value=object())
     monkeypatch.setattr(ar_module, "get_model", get_model)
     monkeypatch.setattr(ar_module, "set_model_tag", lambda _: contextlib.nullcontext())
+    monkeypatch.setattr(ar_module, "get_draft_load_config", lambda _: LoadConfig())
+    monkeypatch.setattr(
+        ar_module, "get_pp_safe_draft_load_config", lambda load_config: load_config
+    )
 
     speculator.load_draft_model(object(), set())
 
@@ -465,3 +565,9 @@ class Config:
     assert target.cache_config.cache_dtype == "fp8_e4m3"
     assert draft_config.parallel_config.data_parallel_index == data_parallel_index
     assert target.parallel_config.data_parallel_index == data_parallel_index
+    assert draft_config.kernel_config.moe_backend == (draft_moe or "flashinfer_cutlass")
+    assert draft_config.attention_config.backend == (
+        draft_attn or AttentionBackendEnum.FLASHINFER
+    )
+    assert target.kernel_config.moe_backend == "flashinfer_cutlass"
+    assert target.attention_config.backend == AttentionBackendEnum.FLASHINFER
diff --git a/tests/v1/worker/test_gpu_model_runner_v2_cudagraph_profiling.py b/tests/v1/worker/test_gpu_model_runner_v2_cudagraph_profiling.py
index 8a655ff135ac..63a628ba7f65 100644
--- a/tests/v1/worker/test_gpu_model_runner_v2_cudagraph_profiling.py
+++ b/tests/v1/worker/test_gpu_model_runner_v2_cudagraph_profiling.py
@@ -254,13 +254,17 @@ def _capture_model(*, profile_only: bool = False) -> int:
     assert compilation_counter.num_gpu_runner_capture_triggers == triggers_before
 
 
-def test_model_runner_delegates_to_cudagraph_utils(monkeypatch):
+@pytest.mark.parametrize("extensible", [False, True])
+def test_model_runner_delegates_to_cudagraph_utils(monkeypatch, extensible):
+    """The extensible KV cache measures the capture footprint after warmup
+    instead of estimating it."""
     runner = mrv2.GPUModelRunner.__new__(mrv2.GPUModelRunner)
+    runner.cache_config = SimpleNamespace(enable_extensible_kv_cache=extensible)
     runner.eplb = EPLBController(
         SimpleNamespace(enable_eplb=False), torch.device("cpu")
     )
     monkeypatch.setattr(mrv2, "_profile_cudagraph_memory", lambda r: 42)
-    assert runner.profile_cudagraph_memory() == 42
+    assert runner.profile_cudagraph_memory() == (0 if extensible else 42)
 
 
 def test_extrapolate_full_graph_memory():
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 c9dadd2e9f80..4cc000c7989a 100644
--- a/tests/v1/worker/test_gpu_model_runner_v2_eplb.py
+++ b/tests/v1/worker/test_gpu_model_runner_v2_eplb.py
@@ -204,6 +204,7 @@ class FakeDSparkSpeculator:
         def __init__(self):
             self.model = draft_model
             self.eplb_state = None
+            self.draft_logits = None
 
         def load_model(self, target_model):
             return None
diff --git a/tests/v1/worker/test_gpu_pcp_manager.py b/tests/v1/worker/test_gpu_pcp_manager.py
index a4e50a36e512..9fa46d87e8f8 100644
--- a/tests/v1/worker/test_gpu_pcp_manager.py
+++ b/tests/v1/worker/test_gpu_pcp_manager.py
@@ -7,6 +7,7 @@
 import pytest
 import torch
 
+from tests.v1.spec_decode.utils import fake_speculative_config
 from vllm.config import CUDAGraphMode, ParallelConfig, VllmConfig
 from vllm.forward_context import set_forward_context
 from vllm.model_executor.layers.fused_moe.router.fused_topk_router import fused_topk
@@ -17,6 +18,8 @@
 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.spec_decode.draft_logits import DraftLogitsCache
+from vllm.v1.worker.gpu.spec_decode.dspark.config import DSparkSpeculatorConfig
 
 
 def _copy_to_cpu(value, device=None, dtype=None, out=None):
@@ -86,8 +89,8 @@ def test_validate_config_dcp_for_dspark(draft_uses_mla, dcp_size):
     config.parallel_config.prefill_context_parallel_size = 4
     config.parallel_config.decode_context_parallel_size = dcp_size
     config.model_config.hf_text_config.index_topk = 1
-    config.speculative_config = SimpleNamespace(
-        use_dspark=lambda: True,
+    config.speculative_config = fake_speculative_config(
+        DSparkSpeculatorConfig,
         draft_model_config=SimpleNamespace(use_mla=draft_uses_mla),
     )
     if draft_uses_mla and dcp_size == 2:
@@ -165,13 +168,15 @@ def test_sync_drafts_writes_logits_to_request_slots(monkeypatch):
     )
     monkeypatch.setattr(pcp_manager_module, "get_pcp_group", lambda: group)
     draft_tokens = torch.full((1, 3), -1)
-    draft_logits = torch.full((4, 3, 8), -1.0)
+    draft_logits = DraftLogitsCache(
+        4, 3, 8, dtype=torch.float32, fill=-1.0, device="cpu"
+    )
 
     manager.sync_drafts(draft_tokens, torch.tensor([2]), draft_logits)
 
     torch.testing.assert_close(draft_tokens, root_tokens)
-    torch.testing.assert_close(draft_logits[2], root_logits[0])
-    assert (draft_logits[[0, 1, 3]] == -1).all()
+    torch.testing.assert_close(draft_logits.values[2], root_logits[0])
+    assert (draft_logits.values[[0, 1, 3]] == -1).all()
 
 
 @pytest.mark.parametrize(
diff --git a/tests/v1/worker/test_gpu_rejection_sampler_chunking.py b/tests/v1/worker/test_gpu_rejection_sampler_chunking.py
index 13f70139d373..91884303d3c0 100644
--- a/tests/v1/worker/test_gpu_rejection_sampler_chunking.py
+++ b/tests/v1/worker/test_gpu_rejection_sampler_chunking.py
@@ -10,7 +10,12 @@
 
 from vllm.config.model import PROCESSED_LOGPROBS_MODES, LogprobsMode
 from vllm.platforms import current_platform
+from vllm.sampling_params import SamplingParams
+from vllm.v1.worker.gpu.sample.logprob import LogprobTokenIdsState
 from vllm.v1.worker.gpu.sample.output import SamplingMaskTensors
+from vllm.v1.worker.gpu.sample.sampler import Sampler
+from vllm.v1.worker.gpu.sample.states import NO_LOGPROBS
+from vllm.v1.worker.gpu.spec_decode.draft_logits import DraftLogitsCache
 from vllm.v1.worker.gpu.spec_decode.rejection_sampler import (
     RejectionSampler,
     _iter_request_chunks,
@@ -56,7 +61,6 @@ def _make_rejection_sampler(
     rejection_sampler.fly_window_size = 0
     rejection_sampler.fly_entropy_threshold = 0.3
     rejection_sampler.fly_entropy_top_k = 3
-    rejection_sampler.draft_logits_num_shards = 1
     rejection_sampler.watermark_key = None
     return rejection_sampler
 
@@ -85,7 +89,11 @@ def test_chunked_scores_match_full_batch(
     idx_mapping_np = np.array([7, 2, 9, 1], dtype=np.int32)
     input_batch = _make_input_batch(cu_num_logits_np, idx_mapping_np, device)
     rejection_sampler = _make_rejection_sampler(
-        SimpleNamespace(logprobs_mode=logprobs_mode, return_sampling_mask=False),
+        SimpleNamespace(
+            logprobs_mode=logprobs_mode,
+            return_sampling_mask=False,
+            logprob_token_ids_state=None,
+        ),
         num_speculative_steps=3,
     )
 
@@ -118,7 +126,7 @@ def fake_verify(
             draft_sampled=torch.arange(10, device=device),
             pos=torch.arange(10, device=device),
             max_chunk_logits=5,
-            max_num_logprobs=2,
+            logprobs_dims=(2, 0),
         )
     )
     score_logits = logits + 1 if logprobs_mode in PROCESSED_LOGPROBS_MODES else logits
@@ -128,7 +136,8 @@ def fake_verify(
         score_logits,
         input_batch.cu_num_logits,
         input_batch.cu_num_logits_np,
-        max_num_logprobs=2,
+        input_batch.expanded_idx_mapping,
+        logprobs_dims=(2, 0),
     )
 
     assert sampled[:, 0].tolist() == idx_mapping_np.tolist()
@@ -151,6 +160,235 @@ def fake_verify(
     assert sampling_mask_tensors is None
 
 
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA")
+@pytest.mark.parametrize("logprobs_mode", ["raw_logprobs", "raw_logits"])
+def test_logprob_token_ids_without_topk_logprobs(logprobs_mode: str):
+    """Requests that only set `logprob_token_ids` (no `logprobs`) still get the
+    scores of those tokens when their batch is verified with speculative drafts."""
+    device = torch.device("cuda")
+    vocab_size = 17
+    cu_num_logits_np = np.array([0, 3, 4, 8, 10], dtype=np.int32)
+    idx_mapping_np = np.array([3, 0, 2, 1], dtype=np.int32)
+    input_batch = _make_input_batch(cu_num_logits_np, idx_mapping_np, device)
+    wanted = {3: [5, 11], 0: [2], 2: [], 1: [16, 0, 7]}  # request state index -> ids
+    token_ids_state = LogprobTokenIdsState(4, device)
+    for req_idx, ids in wanted.items():
+        token_ids_state.add_request(req_idx, SamplingParams(logprob_token_ids=ids))
+    token_ids_state.apply_staged_writes()
+    max_token_ids = token_ids_state.max_num_token_ids(idx_mapping_np)
+    assert max_token_ids == 3
+    rejection_sampler = _make_rejection_sampler(
+        SimpleNamespace(
+            logprobs_mode=logprobs_mode,
+            return_sampling_mask=False,
+            logprob_token_ids_state=token_ids_state,
+        ),
+        num_speculative_steps=3,
+    )
+
+    def fake_verify(
+        self, logits, _draft_logits, _draft_sampled, _pos, cu_num_logits, *_
+    ):
+        num_sampled = torch.diff(cu_num_logits).to(torch.int32)
+        sampled = (
+            torch.arange(4, device=device).repeat(len(num_sampled), 1) % vocab_size
+        )
+        return logits, sampled, num_sampled
+
+    rejection_sampler._verify = MethodType(fake_verify, rejection_sampler)  # type: ignore[method-assign]
+    torch.manual_seed(0)
+    logits = torch.randn(10, vocab_size, device=device)
+
+    _, _, chunked, _ = rejection_sampler._verify_in_chunks(
+        logits,
+        input_batch,
+        draft_logits=None,
+        draft_sampled=torch.arange(10, device=device),
+        pos=torch.arange(10, device=device),
+        max_chunk_logits=5,
+        logprobs_dims=(0, max_token_ids),
+    )
+    assert chunked is not None
+    assert chunked.logprobs.shape == (10, 1 + max_token_ids)
+
+    expected = logits if logprobs_mode == "raw_logits" else logits.log_softmax(-1)
+    for req, req_idx in enumerate(idx_mapping_np):
+        ids = wanted[int(req_idx)]
+        for row in range(cu_num_logits_np[req], cu_num_logits_np[req + 1]):
+            assert chunked.logprob_token_ids[row, 1 : 1 + len(ids)].tolist() == ids
+            torch.testing.assert_close(
+                chunked.logprobs[row, 1 : 1 + len(ids)], expected[row, ids]
+            )
+            assert torch.isinf(chunked.logprobs[row, 1 + len(ids) :]).all()
+
+    # Without any requested token ids nothing is computed.
+    assert (
+        rejection_sampler._get_logprobs_tensors(
+            torch.zeros(4, 4, dtype=torch.int64, device=device),
+            torch.ones(4, dtype=torch.int32, device=device),
+            logits[:4],
+            input_batch.cu_num_logits[:5],
+            input_batch.cu_num_logits_np[:5],
+            input_batch.expanded_idx_mapping[:4],
+            None,
+        )
+        is None
+    )
+
+
+def _make_sampler_with_token_ids(
+    logprobs_mode: str,
+    wanted: dict[int, list[int]],
+    num_logprobs: int,
+    device: torch.device,
+) -> SimpleNamespace:
+    """Sampler stand-in with the real `Sampler.get_logprobs_dims`: every request
+    state index in `wanted` asks for `logprob_token_ids`, and the batch-wide top-k
+    `logprobs` is `num_logprobs` (NO_LOGPROBS when unset)."""
+    token_ids_state = LogprobTokenIdsState(4, device)
+    for req_idx, ids in wanted.items():
+        token_ids_state.add_request(req_idx, SamplingParams(logprob_token_ids=ids))
+    token_ids_state.apply_staged_writes()
+    sampler = SimpleNamespace(
+        logprobs_mode=logprobs_mode,
+        return_sampling_mask=False,
+        logprob_token_ids_state=token_ids_state,
+        sampling_states=SimpleNamespace(max_num_logprobs=lambda _idx: num_logprobs),
+    )
+    sampler.get_logprobs_dims = MethodType(Sampler.get_logprobs_dims, sampler)
+    return sampler
+
+
+def _fake_verify(device: torch.device, vocab_size: int):
+    def fake_verify(
+        self, logits, _draft_logits, _draft_sampled, _pos, cu_num_logits, *_
+    ):
+        num_sampled = torch.diff(cu_num_logits).to(torch.int32)
+        sampled = (
+            torch.arange(4, device=device).repeat(len(num_sampled), 1) % vocab_size
+        )
+        return logits, sampled, num_sampled
+
+    return fake_verify
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA")
+def test_get_logprobs_dims_with_only_token_ids():
+    """`logprob_token_ids` alone is enough to ask for logprobs, whatever the
+    batch's draft tokens (the sharded gather sizes its buffers from these dims)."""
+    device = torch.device("cuda")
+    idx_mapping_np = np.array([3, 0, 2, 1], dtype=np.int32)
+    wanted = {3: [5, 11], 0: [2], 2: [], 1: [16, 0, 7]}
+    sampler = _make_sampler_with_token_ids("raw_logprobs", wanted, NO_LOGPROBS, device)
+    assert sampler.get_logprobs_dims(idx_mapping_np) == (0, 3)
+    sampler = _make_sampler_with_token_ids("raw_logprobs", wanted, 1, device)
+    assert sampler.get_logprobs_dims(idx_mapping_np) == (1, 3)
+    sampler = _make_sampler_with_token_ids("raw_logprobs", {}, NO_LOGPROBS, device)
+    assert sampler.get_logprobs_dims(idx_mapping_np) is None
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA")
+@pytest.mark.parametrize("logprobs_mode", ["raw_logprobs", "raw_logits"])
+def test_logprob_token_ids_with_topk_logprobs(logprobs_mode: str):
+    """Top-k `logprobs` narrower than the widest `logprob_token_ids`: the rows are
+    as wide as the token ids, requests with token ids get exactly those tokens and
+    the others get the top-k."""
+    device = torch.device("cuda")
+    vocab_size = 17
+    cu_num_logits_np = np.array([0, 3, 4, 8, 10], dtype=np.int32)
+    idx_mapping_np = np.array([3, 0, 2, 1], dtype=np.int32)
+    input_batch = _make_input_batch(cu_num_logits_np, idx_mapping_np, device)
+    wanted = {3: [5, 11], 0: [], 2: [], 1: [16, 0, 7, 9]}
+    sampler = _make_sampler_with_token_ids(logprobs_mode, wanted, 1, device)
+    dims = sampler.get_logprobs_dims(idx_mapping_np)
+    assert dims == (1, 4)
+    rejection_sampler = _make_rejection_sampler(sampler, num_speculative_steps=3)
+    rejection_sampler._verify = MethodType(  # type: ignore[method-assign]
+        _fake_verify(device, vocab_size), rejection_sampler
+    )
+    torch.manual_seed(0)
+    logits = torch.randn(10, vocab_size, device=device)
+
+    _, _, lp, _ = rejection_sampler._verify_in_chunks(
+        logits,
+        input_batch,
+        draft_logits=None,
+        draft_sampled=torch.arange(10, device=device),
+        pos=torch.arange(10, device=device),
+        max_chunk_logits=5,
+        logprobs_dims=dims,
+    )
+    assert lp is not None
+    assert lp.logprobs.shape == (10, 1 + 4)
+
+    expected = logits if logprobs_mode == "raw_logits" else logits.log_softmax(-1)
+    top1 = logits.argmax(-1)
+    for req, req_idx in enumerate(idx_mapping_np):
+        ids = wanted[int(req_idx)] or None
+        for row in range(cu_num_logits_np[req], cu_num_logits_np[req + 1]):
+            if ids is None:
+                assert lp.logprob_token_ids[row, 1] == top1[row]
+                torch.testing.assert_close(
+                    lp.logprobs[row, 1], expected[row, top1[row]]
+                )
+                assert torch.isinf(lp.logprobs[row, 2:]).all()
+            else:
+                assert lp.logprob_token_ids[row, 1 : 1 + len(ids)].tolist() == ids
+                torch.testing.assert_close(
+                    lp.logprobs[row, 1 : 1 + len(ids)], expected[row, ids]
+                )
+                assert torch.isinf(lp.logprobs[row, 1 + len(ids) :]).all()
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA")
+def test_logprob_token_ids_with_adaptive_verification():
+    """Adaptive verification verifies one chunk over compacted logits, with
+    `cu_num_logits` on device and `cu_num_logits_np` an upper bound: the scores of
+    `logprob_token_ids` must match the fixed-boundary layout."""
+    device = torch.device("cuda")
+    vocab_size = 17
+    true_cu = np.array([0, 3, 4, 8, 10], dtype=np.int32)
+    idx_mapping_np = np.array([3, 0, 2, 1], dtype=np.int32)
+    wanted = {3: [5, 11], 0: [2], 2: [], 1: [16, 0, 7]}
+    sampler = _make_sampler_with_token_ids("raw_logprobs", wanted, -1, device)
+    dims = sampler.get_logprobs_dims(idx_mapping_np)
+    torch.manual_seed(0)
+    logits = torch.randn(10, vocab_size, device=device)
+
+    def run(adaptive: bool):
+        input_batch = _make_input_batch(true_cu, idx_mapping_np, device)
+        if adaptive:
+            # Pre-compaction layout on the host (4 rows per request); the device
+            # tensors hold the compacted one.
+            input_batch.cu_num_logits_np = np.arange(0, 17, 4, dtype=np.int32)
+            input_batch.expanded_idx_mapping = torch.cat(
+                [
+                    input_batch.expanded_idx_mapping,
+                    torch.zeros(6, dtype=torch.int64, device=device),
+                ]
+            )
+        rejection_sampler = _make_rejection_sampler(sampler, num_speculative_steps=3)
+        rejection_sampler.enable_adaptive_verification = adaptive
+        rejection_sampler._verify = MethodType(  # type: ignore[method-assign]
+            _fake_verify(device, vocab_size), rejection_sampler
+        )
+        return rejection_sampler._verify_in_chunks(
+            logits,
+            input_batch,
+            draft_logits=None,
+            draft_sampled=torch.arange(10, device=device),
+            pos=torch.arange(10, device=device),
+            max_chunk_logits=16,
+            logprobs_dims=dims,
+        )[2]
+
+    fixed, adaptive = run(False), run(True)
+    assert fixed is not None and adaptive is not None
+    assert torch.equal(fixed.logprob_token_ids, adaptive.logprob_token_ids)
+    assert torch.equal(fixed.logprobs, adaptive.logprobs)
+    assert fixed.logprobs.shape == (10, 1 + 3)
+
+
 @pytest.mark.skipif(
     not (current_platform.is_cuda() or current_platform.is_xpu()),
     reason="Requires CUDA or XPU",
@@ -167,8 +405,13 @@ def test_replay_on_off_preserves_rejection_sampling_and_rng(monkeypatch):
     processed_logits = torch.zeros((12, vocab_size), device=device)
     processed_logits[0, 0] = -float("inf")
     processed_logits[rows_per_request + 1, 1] = -float("inf")
-    draft_logits = torch.zeros(
-        (num_reqs, num_speculative_steps, vocab_size), device=device
+    draft_logits = DraftLogitsCache(
+        num_reqs,
+        num_speculative_steps,
+        vocab_size,
+        dtype=torch.float32,
+        fill=0.0,
+        device=device,
     )
     draft_sampled = torch.zeros(
         (num_reqs, rows_per_request), dtype=torch.int64, device=device
@@ -178,6 +421,7 @@ def test_replay_on_off_preserves_rejection_sampling_and_rng(monkeypatch):
     sampler = SimpleNamespace(
         logprobs_mode="processed_logprobs",
         return_sampling_mask=False,
+        logprob_token_ids_state=None,
         sampling_states=SimpleNamespace(
             top_k=SimpleNamespace(np=np.full(num_reqs, vocab_size)),
             temperature=SimpleNamespace(gpu=torch.ones(num_reqs, device=device)),
@@ -187,6 +431,7 @@ def test_replay_on_off_preserves_rejection_sampling_and_rng(monkeypatch):
         ),
         use_fp64_gumbel=False,
         apply_sampling_params=lambda logits, *_: logits,
+        trace_replay_state=None,
     )
     rejection_sampler = _make_rejection_sampler(sampler, num_speculative_steps)
     pos = torch.arange(12, dtype=torch.int32, device=device)
@@ -233,7 +478,7 @@ def track_pack(
             draft_sampled,
             pos,
             max_chunk_logits=8,
-            max_num_logprobs=0,
+            logprobs_dims=(0, 0),
         )
     )
     rng_after_off = torch.get_device_module(device).get_rng_state()
@@ -247,7 +492,7 @@ def track_pack(
         draft_sampled,
         pos,
         max_chunk_logits=8,
-        max_num_logprobs=0,
+        logprobs_dims=(0, 0),
     )
     rng_after_on = torch.get_device_module(device).get_rng_state()
 
diff --git a/tests/v1/worker/test_gpu_rejection_sampler_i64.py b/tests/v1/worker/test_gpu_rejection_sampler_i64.py
index 769444ed64c6..5a571d03149b 100644
--- a/tests/v1/worker/test_gpu_rejection_sampler_i64.py
+++ b/tests/v1/worker/test_gpu_rejection_sampler_i64.py
@@ -15,6 +15,7 @@
 import torch
 
 from vllm.platforms import current_platform
+from vllm.v1.worker.gpu.spec_decode.draft_logits import DraftLogitsCache
 from vllm.v1.worker.gpu.spec_decode.rejection_sampler_utils import rejection_sample
 
 VOCAB_SIZE = 155264
@@ -42,14 +43,16 @@ def _run(
         num_logits, VOCAB_SIZE, dtype=torch.bfloat16, device=device
     )
     full_target[roi_row : roi_row + LOGITS_PER_REQ] = target_logits
-    full_draft = torch.randn(
+    full_draft = DraftLogitsCache(
         max_num_reqs,
         NUM_SPECULATIVE_STEPS,
         VOCAB_SIZE,
         dtype=torch.bfloat16,
+        fill=0.0,
         device=device,
     )
-    full_draft[req_state_idx] = draft_logits
+    full_draft.values.normal_()
+    full_draft.values[req_state_idx] = draft_logits
 
     # Draft token for step i of request r lives at row r*LOGITS_PER_REQ + i + 1.
     draft_sampled = torch.zeros(num_logits, dtype=torch.int64, device=device)
@@ -80,7 +83,7 @@ def _run(
 
     sampled, num_sampled = rejection_sample(
         target_logits=full_target,
-        draft_logits=full_draft,
+        draft_logits_cache=full_draft,
         draft_sampled=draft_sampled,
         cu_num_logits=cu_num_logits,
         pos=pos,
diff --git a/tests/v1/worker/test_gpu_target_dependent_ar_speculator.py b/tests/v1/worker/test_gpu_target_dependent_ar_speculator.py
index 5e2a3f060dbf..27c841a75754 100644
--- a/tests/v1/worker/test_gpu_target_dependent_ar_speculator.py
+++ b/tests/v1/worker/test_gpu_target_dependent_ar_speculator.py
@@ -200,6 +200,7 @@ def test_speculator_uses_draft_model_hidden_size(monkeypatch, hc_mult, expected)
             data_parallel_size=1,
             data_parallel_rank=0,
             decode_context_parallel_size=1,
+            use_ubatching=False,
         ),
     )
 
@@ -679,3 +680,38 @@ def test_decode_metadata_update_needs_draft_decode_steps(
     speculator._configure_decode_metadata_update()
 
     assert speculator.supports_decode_metadata_update is expected
+
+
+def test_rebuilt_prefill_metadata_uses_target_batch_slots():
+    # Under DBO the target builds only per-microbatch metadata, so the drafter
+    # rebuilds the full batch's on its own builders from the target's slots.
+    speculator = object.__new__(_TestSpeculator)
+    target = ({"draft": "target metadata"}, {"draft": torch.zeros(6)})
+    assert speculator._build_drafter_attn_metadata(None, *target) == target
+
+    speculator.can_reuse_target_attn_metadata = False
+    speculator.attn_groups = [[Mock()]]
+    speculator.model_state = Mock()
+    speculator.model_state.prepare_attn.return_value = {"draft": "metadata"}
+    speculator.kv_cache_config = SimpleNamespace(
+        kv_cache_groups=[SimpleNamespace(layer_names=["target", "draft"])]
+    )
+    speculator.block_tables = SimpleNamespace(
+        input_block_tables=[torch.arange(8, dtype=torch.int32).view(4, 2)],
+        slot_mappings=torch.arange(8).view(1, 8),
+    )
+    input_batch = SimpleNamespace(num_reqs_after_padding=2, num_tokens_after_padding=6)
+
+    attn_metadata, slot_mappings = speculator._build_drafter_attn_metadata(
+        input_batch, *target, CUDAGraphMode.FULL
+    )
+
+    assert attn_metadata == {"draft": "metadata"}
+    args = speculator.model_state.prepare_attn.call_args.args
+    batch, cg_mode, (block_table,), target_slots, attn_groups, _ = args
+    assert batch is input_batch and cg_mode == CUDAGraphMode.FULL
+    assert attn_groups is speculator.attn_groups
+    assert block_table.shape == (2, 2)
+    assert target_slots.data_ptr() == speculator.block_tables.slot_mappings.data_ptr()
+    assert target_slots.shape == (1, 6)
+    assert slot_mappings["draft"].tolist() == list(range(6))
diff --git a/tests/v1/worker/test_gpu_trace_replay.py b/tests/v1/worker/test_gpu_trace_replay.py
index 646842108744..d82ea0b86a9a 100644
--- a/tests/v1/worker/test_gpu_trace_replay.py
+++ b/tests/v1/worker/test_gpu_trace_replay.py
@@ -24,6 +24,7 @@
 from vllm.v1.worker.gpu.sample.trace_replay import (
     TraceReplayState,
     apply_trace_tokens,
+    apply_trace_tokens_to_verified,
 )
 from vllm.v1.worker.gpu.states import RequestState
 
@@ -195,3 +196,96 @@ def test_slot_reuse_clears_trace():
     _set_lens(state, [3, 0], [3, 0])
     state.apply_trace(sampled, idx_mapping)
     assert sampled.tolist() == [888]
+
+
+class TestApplyTraceToVerified:
+    @pytest.mark.parametrize(
+        "drafts,trace,step,expected_tokens,expected_num_sampled",
+        [
+            ([10, 11, 12], [10, 11, 12, 13, 14], 0, [10, 11, 12, 13], 4),
+            ([10, 99, 12], [10, 11, 12, 13, 14], 0, [10, 11], 2),
+            ([99, 11, 12], [10, 11, 12, 13, 14], 0, [10], 1),
+            ([-1, -1, -1], [10, 11, 12, 13, 14], 0, [10], 1),
+            ([11, 12, 13], [10, 11, 12, 13, 14], 1, [11, 12, 13, 14], 4),
+            ([12, 13, 14], [10, 11, 12, 13, 14], 2, [12, 13, 14], 3),
+            ([13, 14, 15], [10, 11, 12, 13, 14], 3, [13, 14], 2),
+            ([], [10, 11, 12], 1, [11], 1),
+        ],
+        ids=[
+            "all-accepted",
+            "reject-middle",
+            "reject-first",
+            "placeholder-drafts",
+            "offset-step",
+            "trace-ends-at-bonus",
+            "trace-ends-mid-step",
+            "no-drafts",
+        ],
+    )
+    def test_accepts_drafts_matching_trace(
+        self, drafts, trace, step, expected_tokens, expected_num_sampled
+    ):
+        num_steps = 3
+        prompt_len = 6
+        trace_token_ids = torch.zeros(
+            1, TEST_MAX_MODEL_LEN, dtype=torch.int32, device=DEVICE
+        )
+        trace_token_ids[0, : len(trace)] = _i32(trace)
+        sampled = torch.full((1, num_steps + 1), -7, dtype=torch.int64, device=DEVICE)
+        num_sampled = _i32([1])
+        draft_sampled = _i64([5, *drafts])
+        cu_num_logits = _i32([0, len(drafts) + 1])
+
+        apply_trace_tokens_to_verified(
+            sampled,
+            num_sampled,
+            draft_sampled,
+            cu_num_logits,
+            _i32([0]),
+            trace_token_ids,
+            _i32([len(trace)]),
+            _i32([prompt_len + step]),
+            _i32([prompt_len]),
+        )
+
+        assert num_sampled.tolist() == [expected_num_sampled]
+        assert sampled[0, :expected_num_sampled].tolist() == expected_tokens
+
+    def test_batch_with_mixed_requests(self):
+        trace_token_ids = torch.zeros(
+            3, TEST_MAX_MODEL_LEN, dtype=torch.int32, device=DEVICE
+        )
+        trace_token_ids[0, :4] = _i32([100, 101, 102, 103])
+        trace_token_ids[2, :3] = _i32([300, 301, 302])
+        trace_len = _i32([4, 0, 3])
+        prompt_len = _i32([4, 4, 4])
+        total_len = _i32([5, 9, 7])
+        sampled = torch.tensor(
+            [[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]],
+            dtype=torch.int64,
+            device=DEVICE,
+        )
+        num_sampled = _i32([1, 2, 3, 1])
+        draft_sampled = _i64([0, 101, 102, 0, 5, 6, 0, 0, 0, 7])
+        cu_num_logits = _i32([0, 3, 6, 9, 10])
+        idx_mapping = _i32([0, 1, -1, 2])
+
+        apply_trace_tokens_to_verified(
+            sampled,
+            num_sampled,
+            draft_sampled,
+            cu_num_logits,
+            idx_mapping,
+            trace_token_ids,
+            trace_len,
+            total_len,
+            prompt_len,
+        )
+
+        assert num_sampled.tolist() == [3, 2, 3, 1]
+        assert sampled.tolist() == [
+            [101, 102, 103],
+            [4, 5, 6],
+            [7, 8, 9],
+            [10, 11, 12],
+        ]
diff --git a/tests/v1/worker/test_gpu_ubatch_slicing.py b/tests/v1/worker/test_gpu_ubatch_slicing.py
index 406b21f7f2a4..1bf8d3d6c057 100644
--- a/tests/v1/worker/test_gpu_ubatch_slicing.py
+++ b/tests/v1/worker/test_gpu_ubatch_slicing.py
@@ -48,6 +48,7 @@
     UBatchState,
     _slice_input_batch,
     create_ubatch_slices,
+    merge_ubatch_outputs,
     slice_model_inputs,
 )
 from vllm.v1.worker.ubatch_utils import (
@@ -662,6 +663,19 @@ def test_ubatch_runner_overlaps_and_matches_single_batch():
     )
 
 
+def test_merge_keeps_aux_hidden_states_per_layer():
+    # EAGLE3 and DFlash targets return (hidden_states, [aux per layer]).
+    hidden = torch.arange(6.0).view(6, 1)
+    aux = [torch.arange(6.0).view(6, 1) + 10 * i for i in range(3)]
+    merged_hidden, merged_aux = merge_ubatch_outputs(
+        [(hidden[:2], [x[:2] for x in aux]), (hidden[2:], [x[2:] for x in aux])]
+    )
+    torch.testing.assert_close(merged_hidden, hidden)
+    assert len(merged_aux) == len(aux)
+    for merged, expected in zip(merged_aux, aux):
+        torch.testing.assert_close(merged, expected)
+
+
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="DBO needs a GPU")
 @pytest.mark.parametrize("failed_ubatch", [0, 1])
 def test_ubatch_runner_names_the_microbatch_that_failed(failed_ubatch):
@@ -1002,3 +1016,36 @@ def test_microbatched_graph_needs_every_rank_to_reach_the_split():
     assert desc.cg_mode == CUDAGraphMode.NONE
     assert desc.num_ubatches == 2
     assert dp_sync_state is not None and dp_sync_state.eager
+
+
+@pytest.mark.parametrize("num_reqs", [3, 4, 5])
+def test_uniform_spec_decode_batches_split_on_request_boundaries(num_reqs: int):
+    # Verifying 2 drafts per request: 3 tokens each. A midpoint split would cut
+    # a request in two whenever num_reqs is odd, leaving neither microbatch a
+    # uniform decode.
+    query_len = 3
+    input_batch = _make_input_batch(
+        [query_len] * num_reqs, [64] * num_reqs, _make_buffers()
+    )
+    slices = create_ubatch_slices(input_batch, 2, uniform_token_count=query_len)
+    first, second = slices
+    assert first.token_slice.stop % query_len == 0
+    assert first.request_slice.stop == second.request_slice.start
+    assert first.num_tokens == query_len * (num_reqs // 2)
+    assert second.token_slice.stop == query_len * num_reqs
+
+
+def test_single_token_or_unknown_uniform_count_keeps_midpoint_split():
+    input_batch = _make_input_batch([3] * 5, [64] * 5, _make_buffers())
+    for uniform_token_count in (None, 1):
+        first, _ = create_ubatch_slices(
+            input_batch, 2, uniform_token_count=uniform_token_count
+        )
+        assert first.token_slice.stop == 15 // 2
+
+
+def test_too_few_requests_to_align_keeps_midpoint_split():
+    # One request cannot be split on a request boundary.
+    input_batch = _make_input_batch([4], [64], _make_buffers())
+    first, _ = create_ubatch_slices(input_batch, 2, uniform_token_count=4)
+    assert first.token_slice.stop == 2
diff --git a/tests/v1/worker/test_gpu_warmup_blocks.py b/tests/v1/worker/test_gpu_warmup_blocks.py
index ddfbdb7fac5c..6232d0a3c508 100644
--- a/tests/v1/worker/test_gpu_warmup_blocks.py
+++ b/tests/v1/worker/test_gpu_warmup_blocks.py
@@ -19,6 +19,7 @@
 from vllm.config.vllm import VllmConfig
 from vllm.platforms import current_platform
 from vllm.utils.math_utils import cdiv
+from vllm.v1.core.sched.output import SchedulerOutput
 from vllm.v1.kv_cache_interface import (
     CircularBufferSpec,
     FullAttentionSpec,
@@ -26,8 +27,10 @@
     KVCacheGroupSpec,
     MambaSpec,
 )
+from vllm.v1.worker.gpu import warmup
 from vllm.v1.worker.gpu.warmup import (
     _reserved_block_count,
+    run_context_prefill_warmup,
     run_mixed_prefill_decode_warmup,
     warmup_kernels,
 )
@@ -90,6 +93,7 @@ def _make_runner(
         num_speculative_steps=num_spec_steps,
         speculative_config=SimpleNamespace() if num_spec_steps else None,
         adaptive_verification=None,
+        speculator=None,
         rejection_sampler=None,
         decode_query_len=num_spec_steps + 1,
         is_pooling_model=False,
@@ -104,10 +108,13 @@ def _make_runner(
             kv_cache_groups=kv_cache_groups, num_blocks=1024
         ),
         vllm_config=SimpleNamespace(
-            num_lookahead_tokens=num_lookahead_tokens, is_mm_encoder_only=False
+            num_lookahead_tokens=num_lookahead_tokens,
+            is_mm_encoder_only=False,
+            parallel_config=SimpleNamespace(decode_context_parallel_size=1),
         ),
         kv_block_zeroer=None,
         kv_connector=SimpleNamespace(set_disabled=lambda disabled: None),
+        extensible_kv_cache=None,
     )
 
 
@@ -197,6 +204,106 @@ def test_mixed_warmup_reserves_lookahead_blocks():
     assert set(recorder.num_spec_tokens_to_schedule) == {NUM_SPEC_STEPS}
 
 
+class _StubExtensibleCache:
+    def __init__(self, warmup_committable_blocks: int) -> None:
+        self.warmup_committable_blocks = warmup_committable_blocks
+        self.commits: list[int] = []
+
+    def commit(self, num_blocks: int) -> None:
+        self.commits.append(num_blocks)
+
+
+def test_warmup_sizes_batches_by_the_rank_agreed_committable_count():
+    """Each rank measures its own free memory; the warmup batch and what it
+    commits must follow the smallest count across ranks, or ranks run
+    different shapes and can skip steps the others wait on."""
+    runner = _make_runner(
+        [_attention_group()], num_lookahead_tokens=0, num_spec_steps=0
+    )
+    # The engine agreed on 3 blocks, whatever this rank alone could commit.
+    cache = _StubExtensibleCache(3)
+    runner.extensible_kv_cache = cache
+    recorder = _StepRecorder()
+
+    warmup_kernels(runner, recorder.execute_model, recorder.sample_tokens)
+
+    # One block per request at this length: 3 blocks less the null block fit
+    # two requests, not the four `max_num_seqs` allows.
+    assert len(recorder._held) == 2
+    assert cache.commits[0] == 3
+
+
+def test_mixed_warmup_skips_when_ranks_agree_too_few_blocks():
+    runner = _make_runner([_attention_group()], num_lookahead_tokens=0)
+    cache = _StubExtensibleCache(9)
+    runner.extensible_kv_cache = cache
+    recorder = _StepRecorder()
+
+    # 128 tokens need 9 blocks (8 prefill + 1 decode) beyond the null block.
+    assert not run_mixed_prefill_decode_warmup(
+        runner, recorder.execute_model, recorder.sample_tokens, num_tokens=128
+    )
+    assert cache.commits == []
+    assert recorder.steps == []
+
+
+def test_context_prefill_warmup_attends_to_the_requested_context():
+    """Full-length requests resume from a prefix hit until their cached
+    context covers the requested tokens, holding what the scheduler would."""
+    num_lookahead_tokens = NUM_SPEC_STEPS + 1
+    runner = _make_runner([_attention_group()], num_lookahead_tokens)
+    runner.extensible_kv_cache = cache = _StubExtensibleCache(1024)
+    recorder = _StepRecorder()
+
+    assert run_context_prefill_warmup(
+        runner,
+        recorder.execute_model,
+        recorder.sample_tokens,
+        num_context_tokens=2 * MAX_MODEL_LEN,
+    )
+
+    _assert_covers_lookahead(recorder.steps, num_lookahead_tokens)
+    # Each request fills the model length, so three cover twice the context.
+    assert len(recorder.steps) == 3
+    assert sum(num_computed for _, num_computed, _ in recorder.steps) >= (
+        2 * MAX_MODEL_LEN
+    )
+    for _, num_computed, num_scheduled in recorder.steps:
+        assert num_computed + num_scheduled < MAX_MODEL_LEN
+        assert num_scheduled == runner.decode_query_len + 1
+    assert cache.commits == [1 + 3 * cdiv(MAX_MODEL_LEN, BLOCK_SIZE)]
+
+
+def test_context_prefill_warmup_stays_within_the_context_and_vocabulary():
+    """A long model length must not inflate the context past what the step
+    needs, nor its prompt ids past the vocabulary they embed from."""
+    runner = _make_runner([_attention_group()], num_lookahead_tokens=0)
+    runner.extensible_kv_cache = _StubExtensibleCache(1024)
+    outputs: list[SchedulerOutput] = []
+
+    assert run_context_prefill_warmup(
+        runner, outputs.append, lambda grammar_output: None, num_context_tokens=100
+    )
+
+    (request,) = outputs[0].scheduled_new_reqs
+    assert request.num_computed_tokens == 100
+    vocab_size = runner.model_config.get_vocab_size()
+    assert max(request.prompt_token_ids) < vocab_size < len(request.prompt_token_ids)
+
+
+def test_context_prefill_warmup_is_capped_by_the_committable_blocks():
+    runner = _make_runner([_attention_group()], num_lookahead_tokens=0)
+    blocks_per_req = cdiv(MAX_MODEL_LEN, BLOCK_SIZE)
+    runner.extensible_kv_cache = _StubExtensibleCache(blocks_per_req)
+    recorder = _StepRecorder()
+
+    # The null block leaves one block short of a single full-length request.
+    assert not run_context_prefill_warmup(
+        runner, recorder.execute_model, recorder.sample_tokens, MAX_MODEL_LEN
+    )
+    assert recorder.steps == []
+
+
 @pytest.mark.parametrize("mamba_cache_mode", ["none", "align"])
 def test_warmup_reserves_mamba_speculative_blocks(mamba_cache_mode):
     """Mamba groups hold the running-state block plus the speculative tail.
@@ -321,6 +428,44 @@ def _step(num_new_tokens: int) -> None:
     manager.free(request)
 
 
+def test_reserved_block_count_scales_dcp_sharded_blocks():
+    """Under DCP the scheduler's block of a sharded group spans every rank, so
+    a long prompt holds `dcp_world_size` times fewer blocks than locally."""
+    dcp_world_size = 4
+    kv_cache_config = KVCacheConfig(
+        num_blocks=256, kv_cache_tensors=[], kv_cache_groups=[_attention_group()]
+    )
+    (spec,) = [g.kv_cache_spec for g in kv_cache_config.kv_cache_groups]
+    manager = make_kv_cache_manager(
+        kv_cache_config,
+        max_model_len=MAX_MODEL_LEN,
+        hash_block_size=BLOCK_SIZE,
+        enable_caching=False,
+        dcp_world_size=dcp_world_size,
+        scheduler_block_size=BLOCK_SIZE * dcp_world_size,
+    )
+    num_tokens = MAX_MODEL_LEN - 1
+    (request,) = create_requests(
+        num_requests=1, num_tokens=num_tokens, block_size=BLOCK_SIZE
+    )
+    computed_blocks, num_new_computed, _ = manager.get_computed_blocks(request)
+    assert manager.allocate_slots(
+        request, num_tokens, num_new_computed, computed_blocks
+    )
+    (held,) = manager.get_block_ids(request.request_id)
+
+    assert len(held) == _reserved_block_count(
+        num_tokens,
+        spec,
+        num_lookahead_tokens=0,
+        max_model_len=MAX_MODEL_LEN,
+        max_encoder_len=0,
+        dcp_world_size=dcp_world_size,
+    )
+    assert len(held) == cdiv(num_tokens, BLOCK_SIZE * dcp_world_size)
+    manager.free(request)
+
+
 def _config_without_post_init(
     speculative_config: SpeculativeConfig | None,
 ) -> VllmConfig:
@@ -369,7 +514,7 @@ def test_num_lookahead_tokens_per_method(
     object.__setattr__(
         speculative_config,
         "draft_model_config",
-        SimpleNamespace(hf_config=hf_config),
+        SimpleNamespace(hf_config=hf_config, architectures=[]),
     )
 
     config = _config_without_post_init(speculative_config)
@@ -381,3 +526,29 @@ def test_num_lookahead_tokens_without_speculation():
     config = _config_without_post_init(None)
 
     assert config.num_lookahead_tokens == 0
+
+
+@pytest.mark.parametrize("max_logprobs", [0, 5, -1])
+def test_sizing_warmup_requests_prompt_logprobs_on_its_prefill(max_logprobs):
+    """Prompt logprobs materialize full-vocab logits a chunk at a time, which
+    no other step reaches. The sizing step requests them on its real prefill,
+    so the runtime path computes them from the model's own hidden states."""
+    runner = _make_runner([_attention_group()], num_lookahead_tokens=0)
+    runner.model_config = SimpleNamespace(
+        get_vocab_size=lambda: 64, max_logprobs=max_logprobs, use_mla=False
+    )
+    recorder = _StepRecorder()
+    requested: dict[str, int | None] = {}
+
+    def execute_model(scheduler_output) -> None:
+        for new_req in scheduler_output.scheduled_new_reqs:
+            requested[new_req.req_id] = new_req.sampling_params.prompt_logprobs
+        recorder.execute_model(scheduler_output)
+
+    warmup.run_kv_cache_sizing_warmup(runner, execute_model, recorder.sample_tokens)
+
+    expected = None if max_logprobs == 0 else max_logprobs
+    assert {
+        "prefill" if "_prefill_" in req_id else "decode": value
+        for req_id, value in requested.items()
+    } == {"prefill": expected, "decode": None}
diff --git a/tests/v1/worker/test_gpu_worker.py b/tests/v1/worker/test_gpu_worker.py
index 3ed2fea99d94..687b57a11fe6 100644
--- a/tests/v1/worker/test_gpu_worker.py
+++ b/tests/v1/worker/test_gpu_worker.py
@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import json
 from contextlib import nullcontext
 from types import SimpleNamespace
 from unittest.mock import patch
@@ -19,8 +20,9 @@
 )
 
 
-def test_v2_worker_requires_triton(monkeypatch):
+def test_v2_gpu_worker_requires_triton(monkeypatch):
     monkeypatch.setattr(gpu_worker, "HAS_TRITON", False)
+    monkeypatch.setattr(gpu_worker.current_platform, "is_cpu", lambda: False)
     config = SimpleNamespace(use_v2_model_runner=True)
     with pytest.raises(ValueError, match="Model Runner V2 requires Triton"):
         gpu_worker.Worker(
@@ -31,46 +33,6 @@ def test_v2_worker_requires_triton(monkeypatch):
         )
 
 
-@pytest.mark.skipif(
-    not current_platform.is_cpu(), reason="CPU worker test; needs the CPU platform"
-)
-def test_cpu_worker_without_triton(monkeypatch):
-    import vllm.config.vllm as config_module
-    from vllm.config import VllmConfig
-    from vllm.utils.cpu_resource_utils import MemoryNodeInfo
-    from vllm.v1.worker import cpu_worker
-
-    monkeypatch.delenv("VLLM_USE_V2_MODEL_RUNNER", raising=False)
-    monkeypatch.setattr(config_module, "HAS_TRITON", False)
-    monkeypatch.setattr(gpu_worker, "HAS_TRITON", False)
-    monkeypatch.setattr(
-        cpu_worker, "get_visible_memory_node", lambda _assigned_node_ids: [0]
-    )
-    monkeypatch.setattr(
-        cpu_worker, "get_allowed_cpu_list", lambda: [SimpleNamespace(numa_node=0)]
-    )
-    monkeypatch.setattr(
-        cpu_worker,
-        "get_memory_node_info",
-        lambda node: MemoryNodeInfo(
-            total_memory=8 * GiB_bytes, available_memory=8 * GiB_bytes
-        ),
-    )
-    monkeypatch.setattr(
-        cpu_worker.torch.ops._C,
-        "init_cpu_memory_env",
-        lambda nodes: None,
-        raising=False,
-    )
-
-    config = VllmConfig()
-    worker = cpu_worker.CPUWorker(
-        config, local_rank=0, rank=0, distributed_init_method="env://"
-    )
-    assert not config.use_v2_model_runner
-    assert not worker.use_v2_model_runner
-
-
 def test_sleep_preservation_restores_private_draft_not_shared_target():
     target = torch.nn.Linear(2, 2, bias=False)
     draft = torch.nn.Module()
@@ -352,6 +314,23 @@ def test_startup_plan_apply_gate(plan_env):
     assert explicit.cache_config.kv_cache_memory_bytes == 7 * GiB_bytes
 
 
+def test_startup_plan_rejects_budget_recorded_before_aux_capture(plan_env):
+    """Ignore an old-schema budget even when the fingerprint matches."""
+    worker = _plan_worker()
+    maybe_save_startup_plan(worker, 1024)
+    fingerprint = startup_plan.compute_plan_fingerprint(worker.vllm_config, 0, 1)
+    path = startup_plan._plan_path(fingerprint)
+    with open(path) as f:
+        plan = json.load(f)
+    plan["schema"] = 1
+    with open(path, "w") as f:
+        json.dump(plan, f)
+
+    maybe_apply_startup_plan(worker)
+
+    assert worker.cache_config.kv_cache_memory_bytes is None
+
+
 # Memory accounting of the profiling run (Worker.determine_available_memory).
 
 # The fallback reads only the sign of the measured drop and this process's torch
@@ -518,3 +497,298 @@ def worker(enable_jit_warmup):
     calls.clear()
     gpu_worker.Worker._maybe_activate_jit_monitor(worker(False))
     assert calls == []
+
+
+def test_extensible_init_defers_the_connector(monkeypatch: pytest.MonkeyPatch):
+    """Under the extensible cache the connector is created at the final extend,
+    not at init, and the config is kept for it."""
+    from vllm.v1.worker import gpu_worker
+    from vllm.v1.worker.gpu_worker import Worker
+
+    def fail(*args, **kwargs):
+        raise AssertionError("connector created at init")
+
+    monkeypatch.setattr(gpu_worker, "ensure_kv_transfer_initialized", fail)
+    log: list[str] = []
+    runner = SimpleNamespace(
+        initialize_kv_cache=lambda *args, **kwargs: log.append("initialize_kv_cache"),
+        extensible_kv_cache=SimpleNamespace(
+            reserved_headroom_bytes=0, committable_blocks=lambda: 42
+        ),
+    )
+    kv_cache_config = SimpleNamespace(
+        num_blocks=8,
+        kv_cache_layout=None,
+        needs_kv_cache_zeroing=False,
+        hash_block_size=16,
+        cache_hit_alignment_tokens=16,
+    )
+    worker = SimpleNamespace(
+        cache_config=SimpleNamespace(enable_extensible_kv_cache=True),
+        model_config=SimpleNamespace(enable_return_routed_experts=False),
+        model_runner=runner,
+        _maybe_get_memory_pool_context=lambda tag: nullcontext(),
+        _v2_model_runner=lambda: runner,
+    )
+    # Warmup commits what every rank can: another can commit two fewer.
+    monkeypatch.setattr(gpu_worker, "_min_across_ranks", lambda value: value - 2)
+    Worker.initialize_from_config(worker, kv_cache_config)
+    assert runner.extensible_kv_cache.warmup_committable_blocks == 40
+    assert log == ["initialize_kv_cache"]
+    assert worker._kv_cache_config is kv_cache_config
+
+
+@pytest.mark.parametrize(
+    ("local_reason", "any_rank_unsupported", "expected"),
+    [
+        (None, False, None),
+        (None, True, "another rank cannot use it"),
+        ("no VMM support", True, "no VMM support"),
+    ],
+)
+def test_ranks_agree_on_extensible_kv_cache(
+    monkeypatch: pytest.MonkeyPatch, local_reason, any_rank_unsupported, expected
+):
+    """Every rank in the world group (this engine's, and with MoE data
+    parallelism or the external launcher, its peers') must decide alike: one
+    that keeps the cache would wait in its block-count all-reduces for one
+    that dropped it."""
+    import torch.distributed as dist
+
+    from vllm.v1.worker.gpu_worker import Worker
+
+    reduced: list[tuple[int, object]] = []
+
+    def all_reduce_int(group, value: int, op) -> int:
+        reduced.append((value, op))
+        return int(any_rank_unsupported)
+
+    monkeypatch.setattr(gpu_worker, "vmm_unavailable_reason", lambda: local_reason)
+    monkeypatch.setattr(gpu_worker, "all_reduce_int", all_reduce_int)
+    monkeypatch.setattr(
+        gpu_worker, "get_world_group", lambda: SimpleNamespace(cpu_group=None)
+    )
+    assert Worker.extensible_kv_cache_unsupported_reason(SimpleNamespace()) == expected
+    assert reduced == [(int(local_reason is not None), dist.ReduceOp.MAX)]
+
+
+@pytest.mark.parametrize("connector_pool", [None, object()])
+def test_extensible_extend_asks_connector_for_its_memory_pool(
+    monkeypatch: pytest.MonkeyPatch, connector_pool
+):
+    """Under the extensible cache the connector is created at the final extend.
+    It is asked for a custom memory pool before the cache is committed, as on
+    the standard path, and one it provides cannot back a driver-mapped cache."""
+    from vllm.v1.worker import extensible_kv_cache
+
+    log: list[str] = []
+
+    class Committed(Exception):
+        pass
+
+    def get_mem_pool_context():
+        log.append("get_mem_pool_context")
+        return connector_pool
+
+    def commit(num_blocks, **kwargs):
+        log.append(f"commit:{num_blocks}")
+        raise Committed
+
+    monkeypatch.setattr(
+        extensible_kv_cache,
+        "ensure_kv_transfer_initialized",
+        lambda config, kv_cache_config: log.append(
+            f"ensure_kv_transfer:{kv_cache_config.num_blocks}"
+        ),
+    )
+    monkeypatch.setattr(extensible_kv_cache, "has_kv_transfer_group", lambda: True)
+    monkeypatch.setattr(
+        extensible_kv_cache,
+        "get_kv_transfer_group",
+        lambda: SimpleNamespace(get_mem_pool_context=get_mem_pool_context),
+    )
+    runner = SimpleNamespace(
+        extensible_kv_cache=SimpleNamespace(
+            committed_kv_cache_tensors=lambda config, layout, num_blocks: [],
+            num_committed_blocks=4,
+            commit=commit,
+        ),
+        cache_config=SimpleNamespace(get_resolved_kv_cache_layout=lambda: None),
+        vllm_config=SimpleNamespace(kv_transfer_config=object()),
+    )
+    kv_cache_config = SimpleNamespace(num_blocks=16, kv_cache_tensors=[])
+
+    if connector_pool is None:
+        with pytest.raises(Committed):
+            extensible_kv_cache.extend_kv_cache(runner, kv_cache_config, 8)
+        assert log == ["ensure_kv_transfer:8", "get_mem_pool_context", "commit:8"]
+    else:
+        with pytest.raises(ValueError, match="own memory pool"):
+            extensible_kv_cache.extend_kv_cache(runner, kv_cache_config, 8)
+        assert log == ["ensure_kv_transfer:8", "get_mem_pool_context"]
+    # The engine's config is left as planned.
+    assert kv_cache_config.num_blocks == 16
+
+
+GiB = GiB_bytes
+
+
+def _weight_update_worker(
+    monkeypatch: pytest.MonkeyPatch,
+    *,
+    free_bytes: int,
+    layer_bytes: int,
+    reclaimable: bool,
+    log: list[str],
+):
+    """A Worker over an extensible KV cache with free_bytes of free memory
+    and a model whose largest layer reloads layer_bytes."""
+    from vllm.v1.worker import extensible_kv_cache
+    from vllm.v1.worker.gpu_worker import Worker
+
+    monkeypatch.setattr(
+        extensible_kv_cache, "max_layer_reload_bytes", lambda model: layer_bytes
+    )
+    monkeypatch.setattr(extensible_kv_cache, "has_kv_transfer_group", lambda: False)
+    monkeypatch.setattr(
+        torch.accelerator, "get_memory_info", lambda *args: (free_bytes, 100 * GiB)
+    )
+    monkeypatch.setattr(torch.accelerator, "memory_reserved", lambda *args: 0)
+    monkeypatch.setattr(torch.accelerator, "memory_allocated", lambda *args: 0)
+    monkeypatch.setattr(torch.accelerator, "empty_cache", lambda: None)
+
+    kv_cache = SimpleNamespace(
+        buffer=SimpleNamespace(physical_bytes=50 * GiB),
+        release_physical=lambda: log.append("release"),
+        recommit=lambda: log.append("recommit"),
+        torch_memory_limit=None,
+    )
+
+    def set_torch_memory_limit(num_bytes):
+        kv_cache.torch_memory_limit = num_bytes
+        log.append(f"torch limit {num_bytes}")
+
+    kv_cache.set_torch_memory_limit = set_torch_memory_limit
+    runner = SimpleNamespace(
+        extensible_kv_cache=kv_cache,
+        get_model=lambda: None,
+        reload_weights=lambda *args: log.append("reload"),
+    )
+    worker = object.__new__(Worker)
+    worker.vllm_config = None
+    worker.model_runner = runner
+    worker._kv_cache_reclaimable = reclaimable
+    return worker
+
+
+@pytest.mark.parametrize(
+    ("free_gib", "reclaimable", "expected"),
+    [
+        pytest.param(4, False, ["reload"], id="room"),
+        pytest.param(1, True, ["release", "reload", "recommit"], id="borrows"),
+        pytest.param(2, False, ["reload"], id="short-warns"),
+    ],
+)
+def test_weight_reload_makes_room_from_an_emptied_kv_cache(
+    monkeypatch: pytest.MonkeyPatch, free_gib, reclaimable, expected
+):
+    """A reload that may not fit borrows the extensible KV cache's memory when
+    a pause emptied it, and backs the cache again afterwards."""
+    log: list[str] = []
+    worker = _weight_update_worker(
+        monkeypatch,
+        free_bytes=free_gib * GiB,
+        layer_bytes=GiB,
+        reclaimable=reclaimable,
+        log=log,
+    )
+    with patch.object(gpu_worker, "set_current_vllm_config", nullcontext):
+        worker.reload_weights()
+    assert log == expected
+
+
+def test_weight_reload_lifts_the_torch_memory_cap(monkeypatch: pytest.MonkeyPatch):
+    """Weight processing allocates through torch, beyond the cap sized for
+    serving; the cap returns afterwards, failed reload or not."""
+    log: list[str] = []
+    worker = _weight_update_worker(
+        monkeypatch, free_bytes=4 * GiB, layer_bytes=GiB, reclaimable=False, log=log
+    )
+    worker.model_runner.extensible_kv_cache.torch_memory_limit = 7
+
+    def fail(*args):
+        log.append("reload")
+        raise ValueError("bad checkpoint")
+
+    worker.model_runner.reload_weights = fail
+    with (
+        patch.object(gpu_worker, "set_current_vllm_config", nullcontext),
+        pytest.raises(ValueError),
+    ):
+        worker.reload_weights()
+    assert log == ["torch limit None", "reload", "torch limit 7"]
+
+
+def test_weight_reload_fails_before_loading_when_a_layer_cannot_fit(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    """Failing midway would leave a mix of old and new weights."""
+    log: list[str] = []
+    worker = _weight_update_worker(
+        monkeypatch, free_bytes=GiB // 2, layer_bytes=GiB, reclaimable=False, log=log
+    )
+    with pytest.raises(RuntimeError, match="largest layer"):
+        worker.reload_weights()
+    assert log == []
+
+
+def test_weight_reload_recommits_the_kv_cache_after_a_failed_reload(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    log: list[str] = []
+    worker = _weight_update_worker(
+        monkeypatch, free_bytes=GiB, layer_bytes=GiB, reclaimable=True, log=log
+    )
+
+    def fail(*args):
+        raise ValueError("bad checkpoint")
+
+    worker.model_runner.reload_weights = fail
+    with (
+        patch.object(gpu_worker, "set_current_vllm_config", nullcontext),
+        pytest.raises(ValueError),
+    ):
+        worker.reload_weights()
+    assert log == ["release", "recommit"]
+
+
+def test_weight_reload_leaves_a_released_kv_cache_released(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    """Whoever released the cache (sleep, release_kv_cache_memory) restores it."""
+    log: list[str] = []
+    worker = _weight_update_worker(
+        monkeypatch, free_bytes=0, layer_bytes=GiB, reclaimable=True, log=log
+    )
+    worker.model_runner.extensible_kv_cache.buffer.physical_bytes = 0
+    with patch.object(gpu_worker, "set_current_vllm_config", nullcontext):
+        worker.reload_weights()
+    assert log == ["reload"]
+
+
+def test_discarding_the_kv_cache_releases_an_extensible_one(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    from vllm.v1.worker.gpu_worker import Worker
+
+    log: list[str] = []
+    monkeypatch.setattr(gpu_worker, "has_kv_transfer_group", lambda: False)
+    kv_cache = SimpleNamespace(release_physical=lambda: log.append("release"))
+    worker = SimpleNamespace(
+        model_runner=SimpleNamespace(extensible_kv_cache=kv_cache),
+        sleep_mode_backend=SimpleNamespace(discard=lambda tags: log.append("pool")),
+    )
+    Worker.discard(worker, ("kv_cache",))
+    assert log == ["release", "pool"]
+    Worker.discard(worker, ("weights",))
+    assert log == ["release", "pool", "pool"]
diff --git a/tests/v1/worker/test_gpu_worker_weight_transfer.py b/tests/v1/worker/test_gpu_worker_weight_transfer.py
index 29c70f2a4078..e97cabac9577 100644
--- a/tests/v1/worker/test_gpu_worker_weight_transfer.py
+++ b/tests/v1/worker/test_gpu_worker_weight_transfer.py
@@ -7,6 +7,8 @@
 session is active. These tests verify that delegation and the session guard.
 """
 
+from contextlib import ExitStack
+
 import pytest
 import torch
 import torch.nn as nn
@@ -14,14 +16,16 @@
 from vllm.config import ParallelConfig, VllmConfig, get_current_vllm_config
 from vllm.lora.layers import BaseLayerWithLoRA
 from vllm.model_executor.model_loader.reload.inplace import get_parameter_for_reload
+from vllm.v1.worker import gpu_worker
 from vllm.v1.worker.gpu_worker import Worker
 
 
 class _RecordingEngine:
     """Minimal stand-in for a weight transfer engine."""
 
-    def __init__(self, raise_on_update: bool = False):
+    def __init__(self, raise_on_update: bool = False, raise_on_finish: bool = False):
         self.raise_on_update = raise_on_update
+        self.raise_on_finish = raise_on_finish
         self.started = False
         self.finished = False
         self.closed = False
@@ -45,6 +49,8 @@ def update_weights(self, update_info: dict) -> None:
 
     def finish_weight_update(self) -> None:
         self._record_config()
+        if self.raise_on_finish:
+            raise ValueError("boom")
         self.finished = True
 
     def reset_weight_update_target(self) -> None:
@@ -59,9 +65,12 @@ def __init__(self) -> None:
         self.seen_config: VllmConfig | None = None
         self.reset_lora_calls = 0
 
-    def reload_weights(self) -> None:
+    def reload_weights(self, *args) -> None:
         self.seen_config = get_current_vllm_config()
 
+    def get_model(self) -> nn.Module:
+        return nn.Module()
+
     def reset_lora_state(self) -> None:
         self.reset_lora_calls += 1
 
@@ -74,7 +83,10 @@ def _make_worker(engine: _RecordingEngine | None) -> Worker:
     worker.vllm_config = VllmConfig()
     worker.weight_transfer_engine = engine
     worker._weight_update_active = False
+    worker._incomplete_weight_updates = set()
     worker._weight_update_is_draft = False
+    worker._kv_cache_reclaimable = False
+    worker._weight_update_kv_cache_loan = ExitStack()
     worker.model_runner = _RecordingModelRunner()
     return worker
 
@@ -140,6 +152,7 @@ def test_start_update_finish_delegates_to_engine():
     assert engine.finished is True
     assert engine.reset_count == 1
     assert worker._weight_update_active is False
+    assert not Worker.has_incomplete_weight_update(worker)
     assert engine.seen_configs == [worker.vllm_config] * 3
     assert worker.model_runner.reset_lora_calls == 1
 
@@ -208,6 +221,57 @@ def test_finish_draft_session_keeps_lora_state(monkeypatch):
     assert worker.model_runner.reset_lora_calls == 0
 
 
+def test_refused_draft_start_keeps_weights_complete(monkeypatch):
+    engine = _RecordingEngine()
+    engine.supports_draft_weight_update = True
+    worker = _make_worker(engine)
+
+    def refuse() -> None:
+        raise RuntimeError("no draft model")
+
+    monkeypatch.setattr(worker, "_set_draft_weight_update_target", refuse)
+    with pytest.raises(RuntimeError, match="no draft model"):
+        Worker.start_draft_weight_update(worker)
+
+    assert not Worker.has_incomplete_weight_update(worker)
+
+
+def test_draft_update_keeps_failed_base_load_incomplete(monkeypatch):
+    engine = _RecordingEngine(raise_on_update=True)
+    engine.supports_draft_weight_update = True
+    worker = _make_worker(engine)
+    monkeypatch.setattr(worker, "_set_draft_weight_update_target", lambda: None)
+
+    Worker.start_weight_update(worker)
+    with pytest.raises(ValueError, match="boom"):
+        Worker.update_weights(worker, {"names": ["w"]})
+    Worker.start_draft_weight_update(worker)
+    Worker.finish_weight_update(worker)
+
+    assert Worker.has_incomplete_weight_update(worker)
+
+
+def test_failed_draft_finish_keeps_weights_incomplete(monkeypatch):
+    engine = _RecordingEngine()
+    engine.supports_draft_weight_update = True
+    worker = _make_worker(engine)
+    draft_model = nn.Linear(1, 1)
+    monkeypatch.setattr(worker, "_set_draft_weight_update_target", lambda: None)
+    monkeypatch.setattr(worker, "get_draft_model", lambda: draft_model)
+    monkeypatch.setattr(gpu_worker, "share_source_lm_heads", lambda model: None)
+
+    def fail_restore(model: nn.Module, refresh: bool = True) -> None:
+        raise RuntimeError("out of memory")
+
+    monkeypatch.setattr(gpu_worker, "restore_quantized_lm_heads", fail_restore)
+
+    Worker.start_draft_weight_update(worker)
+    with pytest.raises(RuntimeError, match="out of memory"):
+        Worker.finish_weight_update(worker)
+
+    assert Worker.has_incomplete_weight_update(worker)
+
+
 def test_double_start_raises():
     worker = _make_worker(_RecordingEngine())
     Worker.start_weight_update(worker)
@@ -240,6 +304,23 @@ def test_update_resets_active_on_error():
     assert worker._weight_update_active is False
 
 
+@pytest.mark.parametrize("mode", ["direct", "layerwise"])
+def test_failed_finish_ends_session_only_in_direct_mode(mode, monkeypatch):
+    monkeypatch.setenv("VLLM_RELOAD_MODE", mode)
+    engine = _RecordingEngine(raise_on_finish=True)
+    worker = _make_worker(engine)
+    Worker.start_weight_update(worker)
+
+    with pytest.raises(ValueError, match="boom"):
+        Worker.finish_weight_update(worker)
+
+    # a failed direct finish cannot be retried, so the next start must work;
+    # a failed layerwise finish stays active and can be retried
+    direct = mode == "direct"
+    assert worker._weight_update_active is not direct
+    assert engine.reset_count == int(direct)
+
+
 def test_missing_engine_raises():
     worker = _make_worker(None)
     with pytest.raises(RuntimeError, match="Weight transfer not configured"):
diff --git a/tests/v1/worker/test_grammar_invalid_drafts.py b/tests/v1/worker/test_grammar_invalid_drafts.py
deleted file mode 100644
index 77e0ef5242d5..000000000000
--- a/tests/v1/worker/test_grammar_invalid_drafts.py
+++ /dev/null
@@ -1,56 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Which speculative drafts a structured-output request may accept.
-
-`StructuredOutputManager.grammar_bitmask` fills row i before it inspects
-`req_tokens[i]`, so with the first -1 placeholder at index j rows 0..j carry a
-real mask and rows j+1.. carry the all-permissive `_full_mask`. Drafts j..K-1
-must therefore be rejected, or the request samples with no grammar constraint
-at a position the model really does sample.
-"""
-
-from types import SimpleNamespace
-
-import numpy as np
-import torch
-
-from vllm.v1.worker.gpu.structured_outputs import grammar_invalid_drafts
-
-
-def _input_batch(req_ids, num_drafts, num_admitted=None):
-    """`num_admitted` is what adaptive verification kept on device."""
-    num_admitted = num_drafts if num_admitted is None else num_admitted
-    rows = [n + 1 for n in num_admitted]
-    return SimpleNamespace(
-        req_ids=list(req_ids),
-        num_reqs=len(req_ids),
-        num_draft_tokens=sum(num_drafts),
-        num_draft_tokens_per_req=np.array(num_drafts, dtype=np.int32),
-        cu_num_logits=torch.tensor([0, *np.cumsum(rows)], dtype=torch.int32),
-        expanded_local_pos=torch.cat(
-            [torch.arange(n, dtype=torch.int32) for n in rows]
-        ),
-    )
-
-
-def _invalid_rows(batch, grammar_req_ids, num_acceptable):
-    mask = grammar_invalid_drafts(batch, grammar_req_ids, num_acceptable)
-    return None if mask is None else mask.nonzero().flatten().tolist()
-
-
-def test_reject_whole_window_when_nothing_was_backfilled():
-    # Request "g" is at rows 0..3 (1 + 3 drafts), so its drafts are rows 1..3.
-    batch = _input_batch(["g", "p"], [3, 3])
-    assert _invalid_rows(batch, ["g"], [0]) == [1, 2, 3]
-
-
-def test_keep_the_drafts_the_bitmask_could_see():
-    batch = _input_batch(["g", "p"], [3, 3])
-    assert _invalid_rows(batch, ["g"], [2]) == [3]
-
-
-def test_follow_the_device_layout_under_adaptive_verification():
-    # Scheduled 3 drafts each, but adaptive verification admitted 2 for "a" and
-    # 3 for "b": the real rows are a=0..2, b=3..6, not the scheduled 0..3, 4..7.
-    batch = _input_batch(["a", "b"], [3, 3], num_admitted=[2, 3])
-    assert _invalid_rows(batch, ["a", "b"], [1, 0]) == [2, 4, 5, 6]
diff --git a/tests/v1/worker/test_kv_cache_allocation_scope.py b/tests/v1/worker/test_kv_cache_allocation_scope.py
index b5b288d9b5e4..3570d64d906e 100644
--- a/tests/v1/worker/test_kv_cache_allocation_scope.py
+++ b/tests/v1/worker/test_kv_cache_allocation_scope.py
@@ -11,6 +11,7 @@
 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.hisparse.backend import CUDAHiSparseBackend
 from vllm.v1.worker.gpu_worker import Worker
 
 
@@ -60,6 +61,9 @@ def bind_hisparse(**kwargs):
                 )
             ]
 
+        monkeypatch.setattr(
+            hisparse_binding, "get_hisparse_backend", CUDAHiSparseBackend
+        )
         monkeypatch.setattr(
             hisparse_binding,
             "get_hisparse_host_pool",
diff --git a/tests/v1/worker/test_mixed_warmup_gate.py b/tests/v1/worker/test_mixed_warmup_gate.py
index 01f303a27597..a872591954c8 100644
--- a/tests/v1/worker/test_mixed_warmup_gate.py
+++ b/tests/v1/worker/test_mixed_warmup_gate.py
@@ -39,12 +39,18 @@ def test_mixed_warmup_disables_watermarking():
         max_num_reqs=2,
         kv_cache_config=SimpleNamespace(
             kv_cache_groups=[
-                SimpleNamespace(kv_cache_spec=SimpleNamespace(block_size=16))
+                SimpleNamespace(
+                    kv_cache_spec=SimpleNamespace(block_size=16, dcp_sharded=False)
+                )
             ],
             num_blocks=128,
         ),
-        vllm_config=SimpleNamespace(num_lookahead_tokens=0),
+        vllm_config=SimpleNamespace(
+            num_lookahead_tokens=0,
+            parallel_config=SimpleNamespace(decode_context_parallel_size=1),
+        ),
         max_model_len=128,
+        model_config=SimpleNamespace(get_vocab_size=lambda: 64),
         num_speculative_steps=0,
         model_state=SimpleNamespace(max_encoder_len=0),
         kv_connector=SimpleNamespace(set_disabled=lambda disabled: None),
@@ -76,6 +82,7 @@ def test_kernel_warmup_restores_uncalibrated_adaptive_manager(monkeypatch, fail_
     rejection_sampler = SimpleNamespace(enable_adaptive_verification=True)
     runner = SimpleNamespace(
         adaptive_verification=manager,
+        speculator=None,
         rejection_sampler=rejection_sampler,
         speculative_config=SimpleNamespace(),
     )
@@ -116,6 +123,7 @@ def test_kernel_warmup_runs_same_rounds_on_every_pp_stage(monkeypatch, speculati
         rounds.clear()
         runner = SimpleNamespace(
             adaptive_verification=None,
+            speculator=None,
             speculative_config=SimpleNamespace() if speculative else None,
             rejection_sampler=(
                 SimpleNamespace(enable_adaptive_verification=False)
diff --git a/tests/v1/worker/test_mrv2_lmhead_lora.py b/tests/v1/worker/test_mrv2_lmhead_lora.py
new file mode 100644
index 000000000000..341854ad9875
--- /dev/null
+++ b/tests/v1/worker/test_mrv2_lmhead_lora.py
@@ -0,0 +1,61 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import numpy as np
+import pytest
+
+from vllm.lora.request import LoRARequest
+from vllm.v1.worker.gpu.lora_utils import LoraState
+
+pytestmark = [pytest.mark.cpu_test, pytest.mark.skip_global_cleanup]
+
+
+def test_prompt_lora_mapping_aligns_with_spec_decode_logits_rows():
+    """lm_head LoRA ids follow logits rows, not requests.
+
+    Under speculative decoding cu_num_logits is one bonus row plus each
+    draft row. A per-request mapping writes the adapter onto a neighbor.
+    """
+    state = LoraState(max_num_reqs=4)
+    lora_b = LoRARequest("adapter-b", 5, "/tmp/adapter-b")
+    # State slots are not batch order: slot 1 is req-a, slot 0 is req-b.
+    state.add_request("req-a", 1, None)
+    state.add_request("req-b", 0, lora_b)
+
+    idx_mapping = np.array([1, 0], dtype=np.int32)
+    # req-a: 1 bonus + 2 drafts. req-b: 1 bonus and no drafts.
+    cu_num_logits = np.array([0, 3, 4], dtype=np.int32)
+    num_logits = np.diff(cu_num_logits).astype(np.int32)
+    # Query length is independent: req-a is still prefilling.
+    num_scheduled = np.array([8, 1], dtype=np.int32)
+
+    prompt, token, requests = state.make_lora_inputs(
+        ["req-a", "req-b"],
+        idx_mapping,
+        num_scheduled,
+        num_logits,
+    )
+
+    assert num_logits.tolist() == [3, 1]
+    assert prompt == (0, 0, 0, 5)
+    assert len(prompt) == int(cu_num_logits[-1])
+    assert token == (0,) * 8 + (5,)
+    assert requests == {lora_b}
+
+
+def test_prompt_lora_mapping_repeats_for_multiple_sampled_tokens():
+    """Multiple sampled tokens per step still get one lm_head id per row."""
+    state = LoraState(max_num_reqs=2)
+    lora = LoRARequest("adapter", 3, "/tmp/adapter")
+    state.add_request("req-a", 0, lora)
+    state.add_request("req-b", 1, None)
+    # Two logits rows per request, no drafts.
+    num_logits = np.array([2, 2], dtype=np.int32)
+    prompt, token, _requests = state.make_lora_inputs(
+        ["req-a", "req-b"],
+        np.array([0, 1], dtype=np.int32),
+        np.array([2, 2], dtype=np.int32),
+        num_logits,
+    )
+    assert prompt == (3, 3, 0, 0)
+    assert token == (3, 3, 0, 0)
diff --git a/tests/v1/worker/test_sleep_mode_backend.py b/tests/v1/worker/test_sleep_mode_backend.py
index a8b8d61ae2b3..34efde606392 100644
--- a/tests/v1/worker/test_sleep_mode_backend.py
+++ b/tests/v1/worker/test_sleep_mode_backend.py
@@ -57,6 +57,7 @@ def resume(self, tags: list[str] | None = None) -> None:
     worker = object.__new__(Worker)
     worker._sleep_mode_backend = Backend()
     worker._sleep_saved_parameters = {}
+    worker.model_runner = SimpleNamespace(extensible_kv_cache=None)
     worker._sleep_saved_buffers = {}
     worker._sleep_saved_draft_buffers = {}
     worker.vllm_config = SimpleNamespace(
diff --git a/tests/v1/worker/test_spec_decode_embed_sharing_pp.py b/tests/v1/worker/test_spec_decode_embed_sharing_pp.py
index 0fef7bb47bdf..ceda7a99b69f 100644
--- a/tests/v1/worker/test_spec_decode_embed_sharing_pp.py
+++ b/tests/v1/worker/test_spec_decode_embed_sharing_pp.py
@@ -7,6 +7,8 @@
 import torch
 import torch.nn as nn
 
+from tests.v1.spec_decode.utils import fake_speculative_config
+from vllm.config import ParallelConfig
 from vllm.model_executor.model_loader.mtp_validation import (
     disable_mtp_completeness_check,
 )
@@ -90,47 +92,44 @@ def test_mtp_style_drafter_is_left_alone_under_pp(monkeypatch):
 
 
 @pytest.mark.parametrize(
-    "method,pp_size,is_last_rank,expected",
+    "method,may_use_target_embed,pp_size,is_last_rank,expected",
     [
-        ("eagle", 2, True, True),
-        ("eagle3", 2, True, True),
-        ("dflash", 2, True, True),
-        ("dspark", 2, True, True),
-        ("eagle3", 1, True, False),
-        ("eagle3", 2, False, False),
-        ("mtp", 2, True, False),
-        (None, 2, True, False),
+        ("eagle", True, 2, True, True),
+        ("eagle3", True, 2, True, True),
+        ("dflash", True, 2, True, True),
+        ("dspark", True, 2, True, True),
+        ("eagle3", True, 1, True, False),
+        ("eagle3", True, 2, False, False),
+        # Depends on the MTP draft model (e.g. MiniMax-M3's may use it).
+        ("mtp", True, 2, True, True),
+        ("mtp", False, 2, True, False),
+        # Not a hidden-state drafter, so never reuses the target's embedding.
+        ("draft_model", True, 2, True, False),
+        (None, True, 2, True, False),
     ],
 )
 def test_target_embedding_provisioning(
-    monkeypatch, method, pp_size, is_last_rank, expected
+    monkeypatch, method, may_use_target_embed, pp_size, is_last_rank, expected
 ):
     monkeypatch.setattr(
         "vllm.distributed.parallel_state.get_pp_group",
         _fake_pp(pp_size, is_last_rank),
         raising=True,
     )
-    speculative_config = None if method is None else SimpleNamespace(method=method)
+    speculative_config = (
+        None
+        if method is None
+        else fake_speculative_config(
+            method=method,
+            draft_model_config=SimpleNamespace(
+                may_use_target_embed=may_use_target_embed
+            ),
+        )
+    )
     config = SimpleNamespace(speculative_config=speculative_config)
     assert spec_decode_needs_target_embed(config) is expected
 
 
-@pytest.mark.parametrize(
-    "pp_size,is_last_rank,expected",
-    [(2, True, True), (2, False, False), (1, True, False)],
-)
-def test_mtp_target_embedding_requires_model_opt_in(
-    monkeypatch, pp_size, is_last_rank, expected
-):
-    monkeypatch.setattr(
-        "vllm.distributed.parallel_state.get_pp_group",
-        _fake_pp(pp_size, is_last_rank),
-    )
-    config = SimpleNamespace(speculative_config=SimpleNamespace(method="mtp"))
-    assert spec_decode_needs_target_embed(config, include_mtp=True) is expected
-    assert not spec_decode_needs_target_embed(config)
-
-
 @pytest.fixture
 def m3_modules():
     if not current_platform.is_cuda_alike():
@@ -155,12 +154,16 @@ def m3_config():
         use_v2_model_runner=True,
         load_config=SimpleNamespace(load_format="auto"),
         model_config=SimpleNamespace(hf_text_config=hf_config),
-        speculative_config=SimpleNamespace(
-            method="mtp", draft_model_config=SimpleNamespace(hf_config=hf_config)
+        speculative_config=fake_speculative_config(
+            method="mtp",
+            draft_model_config=SimpleNamespace(
+                hf_config=hf_config, may_use_target_embed=True
+            ),
         ),
+        parallel_config=ParallelConfig(),
         quant_config=None,
         cache_config=None,
-        kernel_config=SimpleNamespace(monokernel=False),
+        kernel_config=SimpleNamespace(monokernel=False, moe_backend="auto"),
     )
 
 
@@ -245,21 +248,21 @@ def test_m3_preserves_legacy_embedding_sharing(
             ("model.embed_tokens.weight", torch.ones(VOCAB, HIDDEN)),
         ]
     )
-    assert not hasattr(draft, "has_own_embed_tokens")
+    # Not tracked, so the target's embedding is shared.
+    assert draft.has_own_embed_tokens is False
 
 
 @pytest.mark.parametrize(
-    "rank,use_v2,method,expected",
+    "rank,method,expected",
     [
-        (0, True, "mtp", True),
-        (1, True, "mtp", False),
-        (2, True, "mtp", True),
-        (2, False, "mtp", False),
-        (2, True, None, False),
+        (0, "mtp", True),
+        (1, "mtp", False),
+        (2, "mtp", True),
+        (2, None, False),
     ],
 )
 def test_m3_target_allocates_embedding_on_required_stage(
-    monkeypatch, m3_modules, m3_config, rank, use_v2, method, expected
+    monkeypatch, m3_modules, m3_config, rank, method, expected
 ):
     _, model = m3_modules
     pp = lambda: SimpleNamespace(
@@ -267,10 +270,13 @@ def test_m3_target_allocates_embedding_on_required_stage(
     )
     monkeypatch.setattr(model, "get_pp_group", pp)
     monkeypatch.setattr("vllm.distributed.parallel_state.get_pp_group", pp)
+    monkeypatch.setattr(
+        "vllm.distributed.parallel_state.get_tp_group",
+        lambda: SimpleNamespace(world_size=1),
+    )
     monkeypatch.setattr(model, "VocabParallelEmbedding", lambda *_, **kw: _embed())
     monkeypatch.setattr(model, "make_layers", lambda *_, **kw: (0, 0, nn.ModuleList()))
     monkeypatch.setattr(model, "MiniMAXGemmaRMSNorm", lambda *_, **kw: nn.Identity())
-    m3_config.use_v2_model_runner = use_v2
     if method is None:
         m3_config.speculative_config = None
     target = model.MiniMaxM3Model(vllm_config=m3_config)
diff --git a/tests/v1/worker/test_utils.py b/tests/v1/worker/test_utils.py
index ef162e4b0ccb..863900b0d199 100644
--- a/tests/v1/worker/test_utils.py
+++ b/tests/v1/worker/test_utils.py
@@ -28,6 +28,7 @@
 from vllm.utils.mem_constants import GiB_bytes
 from vllm.utils.mem_utils import MemorySnapshot
 from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy
+from vllm.v1.hisparse.backend import CUDAHiSparseBackend, HiSparseBackend
 from vllm.v1.hisparse.runtime import HiSparseCacheHandle
 from vllm.v1.hisparse.types import SparseKVPageTransfer, SparseKVRowMirror
 from vllm.v1.worker.utils import (
@@ -38,8 +39,26 @@
 )
 
 
+@pytest.mark.parametrize("packed_bytes", [False, True])
+def test_hisparse_resident_gather_preserves_stride_and_payload(packed_bytes):
+    # Include a gap between physical pages and select reordered/repeated rows.
+    backing = torch.arange(24, dtype=torch.float32)
+    source = backing.as_strided((2, 2, 4), (12, 4, 1))
+    expected = torch.stack([backing[16:20], backing[0:4], backing[12:16], backing[0:4]])
+    if packed_bytes:
+        source = source.view(torch.uint8)
+    destination = torch.empty((4, 4), dtype=torch.float32)
+
+    HiSparseBackend().gather_resident_rows(
+        source, destination, torch.tensor([3, 0, 2, -1])
+    )
+
+    torch.testing.assert_close(destination, expected, rtol=0, atol=0)
+
+
 def _make_hisparse_worker() -> HiSparseConnectorWorker:
     worker = object.__new__(HiSparseConnectorWorker)
+    worker.backend = CUDAHiSparseBackend()
     worker._slot_mapping_staging = None
     worker._row_mirror_num_rows = 0
     worker._per_layer_mirrored = set()
@@ -58,6 +77,7 @@ def _make_hisparse_worker() -> HiSparseConnectorWorker:
     worker._metrics_event = MagicMock()
     worker._metrics_pending = False
     worker.leader_runtimes = []
+    worker._invalidate_written_slots = MagicMock()
     return worker
 
 
@@ -85,9 +105,7 @@ def test_hisparse_worker_get_kv_connector_stats_reads_completed_snapshot(monkeyp
     )
     worker.leader_runtimes = [SimpleNamespace(index_group=group)]  # type: ignore[list-item]  # runtime stand-in
     monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False)
-    monkeypatch.setattr(
-        hisparse_worker_module, "current_stream", lambda: compute_stream
-    )
+    monkeypatch.setattr(torch.accelerator, "current_stream", lambda: compute_stream)
 
     assert worker.get_kv_connector_stats() is None
     compute_stream.wait_stream.assert_called_once_with(group.copy_stream)
@@ -168,8 +186,7 @@ def test_hisparse_appends_reference_slots_within_a_mirror_phase(monkeypatch):
     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())
+    monkeypatch.setattr(torch.accelerator, "current_stream", lambda: main_stream)
 
     handle.slot_mapping = torch.tensor([7, 8], dtype=torch.int64)
     worker._stage_row_mirror_mapping(2)
@@ -223,7 +240,7 @@ def _hisparse_parallel_config(**overrides):
 
 def test_hisparse_shares_host_pool_only_for_local_tp(monkeypatch):
     monkeypatch.setattr(
-        hisparse_runtime_module.current_platform, "is_cuda_alike", lambda: True
+        hisparse_runtime_module, "get_hisparse_backend", CUDAHiSparseBackend
     )
     config = SimpleNamespace(parallel_config=_hisparse_parallel_config())
 
@@ -297,7 +314,7 @@ def pin_tensor(tensor):
         parallel_config=_hisparse_parallel_config(),
     )
 
-    pools, private_pools, region = hisparse_runtime_module.allocate_hisparse_host_pools(
+    pools, pinned_pools, region = hisparse_runtime_module.allocate_hisparse_host_pools(
         config,
         [24, 40],
         num_blocks=4,
@@ -307,17 +324,7 @@ def pin_tensor(tensor):
 
     assert region is not None
     assert isinstance(region, FakeSharedOffloadRegion)
-    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_shard": (1, 2),
-    }
+    assert [t.data_ptr() for t in pinned_pools] == [t.data_ptr() for t in pinned]
     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)
@@ -325,14 +332,18 @@ def pin_tensor(tensor):
         (tensor.data_ptr() - region.base_tensor.data_ptr(), tensor.nbytes)
         for tensor in pinned
     ] == [(0, page), (page, 3 * page)]
-    assert region.is_pinned
+    assert not 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."""
+    """Only successful registrations may be unpinned on allocation failure."""
+    monkeypatch.setattr(
+        hisparse_runtime_module, "get_hisparse_backend", CUDAHiSparseBackend
+    )
+    monkeypatch.setattr(torch.accelerator, "synchronize", lambda: None)
     page = mmap.PAGESIZE
     tp_group = MagicMock(rank_in_group=0, world_size=2)
     monkeypatch.setattr(hisparse_runtime_module, "get_tp_group", lambda: tp_group)
@@ -350,6 +361,7 @@ def test_shared_host_pool_tracks_successful_registrations(
     cudart.cudaHostRegister.side_effect = [
         SimpleNamespace(value=int(fail_registration == i)) for i in (1, 2)
     ]
+    cudart.cudaHostUnregister.return_value = SimpleNamespace(value=0)
     monkeypatch.setattr(torch.cuda, "cudart", lambda: cudart)
     config = SimpleNamespace(
         instance_id="test", parallel_config=SimpleNamespace(data_parallel_index=0)
@@ -359,16 +371,13 @@ def test_shared_host_pool_tracks_successful_registrations(
         if fail_registration
         else nullcontext()
     )
+    result = None
     with context:
-        hisparse_runtime_module.allocate_hisparse_host_pools(
+        result = 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)
+    pinned_ptrs = [backing.data_ptr() + i * page for i in range(successful)]
     attempts = 2 if fail_registration is None else fail_registration
     assert (
         cudart.cudaHostRegister.call_args_list
@@ -377,6 +386,18 @@ def test_shared_host_pool_tracks_successful_registrations(
             call(backing.data_ptr() + page, 2 * page, 0),
         ][:attempts]
     )
+    assert region.pinned_addresses == []
+    assert not region.is_pinned
+    if fail_registration is None:
+        assert result is not None
+        assert [t.data_ptr() for t in result[1]] == pinned_ptrs
+        cudart.cudaHostUnregister.assert_not_called()
+        region.cleanup.assert_not_called()
+    else:
+        assert cudart.cudaHostUnregister.call_args_list == [
+            call(ptr) for ptr in reversed(pinned_ptrs)
+        ]
+        region.cleanup.assert_called_once_with()
 
 
 def test_shared_host_pool_registers_layer_spans_in_one_backing(monkeypatch):
@@ -600,6 +621,9 @@ def test_hisparse_tail_restore_preserves_imported_rows(
     """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")
+    monkeypatch.setattr(
+        hisparse_runtime_module, "get_hisparse_backend", CUDAHiSparseBackend
+    )
     worker = _make_hisparse_worker()
     worker.is_host_writer = is_host_writer
     worker.kernel_block_size = 4
@@ -636,7 +660,7 @@ def test_hisparse_tail_restore_preserves_imported_rows(
         worker, "_record_transfer_completion", record_transfer_completion
     )
     stream = torch.cuda.current_stream() if device == "cuda" else MagicMock()
-    monkeypatch.setattr(hisparse_worker_module, "current_stream", lambda: stream)
+    monkeypatch.setattr(torch.accelerator, "current_stream", lambda: stream)
     transfer = SparseKVPageTransfer(7, 2, (1, 3), False, restore=True)
 
     restore_pages = worker._restore_pages
@@ -703,14 +727,9 @@ def test_hisparse_dma_row_mirror_builds_descriptors(monkeypatch):
     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.accelerator, "current_stream", lambda: current_stream)
     monkeypatch.setattr(torch, "Event", lambda: event)
-    monkeypatch.setattr(
-        hisparse_worker_module.ops, "swap_blocks_batch", swap_blocks_batch
-    )
+    monkeypatch.setattr(worker.backend, "swap_blocks_batch", swap_blocks_batch)
 
     worker._set_row_mirrors((SparseKVRowMirror((3,), 4, 1),))
     worker._enqueue_row_dma(range(1))
@@ -769,12 +788,10 @@ def test_hisparse_finish_forward_mirrors_all_layers_once(monkeypatch):
     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(
@@ -799,12 +816,12 @@ def test_hisparse_finish_forward_mirrors_all_layers_once(monkeypatch):
     worker._post_forward_transfers = []
     enqueue_row_dma = MagicMock()
     monkeypatch.setattr(worker, "_enqueue_row_dma", enqueue_row_dma)
+    invalidate_written_slots = MagicMock()
+    monkeypatch.setattr(worker, "_invalidate_written_slots", invalidate_written_slots)
     worker.host_write_event = MagicMock()
     worker._forward_ready_event = MagicMock()
     current_stream = MagicMock()
-    monkeypatch.setattr(
-        hisparse_worker_module, "current_stream", lambda: current_stream
-    )
+    monkeypatch.setattr(torch.accelerator, "current_stream", lambda: current_stream)
 
     worker.finish_forward()
 
@@ -812,11 +829,10 @@ def test_hisparse_finish_forward_mirrors_all_layers_once(monkeypatch):
     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()
+    invalidate_written_slots.assert_called_once()
+    leaders, written_slots, _ = invalidate_written_slots.call_args.args
+    assert leaders == [leader]
+    torch.testing.assert_close(written_slots, dst_slots)
     worker.host_write_event.record.assert_called_once_with(current_stream)
 
 
@@ -825,7 +841,6 @@ def test_hisparse_finish_forward_does_not_repeat_per_layer_mirrors(monkeypatch):
     runtime = SimpleNamespace(
         eager_host_mirror=False,
         is_group_leader=False,
-        invalidate_written_slots=MagicMock(),
     )
     handles = [
         SimpleNamespace(
@@ -902,7 +917,6 @@ def test_hisparse_finish_forward_rejects_partial_per_layer_mirror(monkeypatch):
     runtime = SimpleNamespace(
         eager_host_mirror=False,
         is_group_leader=False,
-        invalidate_written_slots=MagicMock(),
     )
     handles = [
         SimpleNamespace(
@@ -945,9 +959,7 @@ def test_hisparse_finish_forward_orders_next_forward_after_dma(monkeypatch):
     worker._dma_free_descriptors = []
     worker._forward_ready_event = MagicMock()
     monkeypatch.setattr(worker, "_finish_mirror_phase", MagicMock())
-    monkeypatch.setattr(
-        hisparse_worker_module, "current_stream", lambda: current_stream
-    )
+    monkeypatch.setattr(torch.accelerator, "current_stream", lambda: current_stream)
 
     worker.finish_forward()
 
@@ -960,7 +972,6 @@ def test_hisparse_finish_forward_mirrors_standalone_mtp_cache(monkeypatch):
     runtime = SimpleNamespace(
         eager_host_mirror=True,
         is_group_leader=False,
-        invalidate_written_slots=MagicMock(),
     )
     inactive = SimpleNamespace(
         runtime=runtime,
@@ -1000,7 +1011,6 @@ def test_hisparse_shared_host_reader_skips_mirror(monkeypatch):
     leader = SimpleNamespace(
         eager_host_mirror=True,
         is_group_leader=True,
-        invalidate_written_slots=MagicMock(),
     )
     handle = SimpleNamespace(
         runtime=leader,
@@ -1018,11 +1028,14 @@ def test_hisparse_shared_host_reader_skips_mirror(monkeypatch):
     worker._per_layer_mirrored = set()
     enqueue_row_dma = MagicMock()
     monkeypatch.setattr(worker, "_enqueue_row_dma", enqueue_row_dma)
+    invalidate_written_slots = MagicMock()
+    monkeypatch.setattr(worker, "_invalidate_written_slots", invalidate_written_slots)
 
     worker._enqueue_host_mirror()
 
     enqueue_row_dma.assert_not_called()
-    leader.invalidate_written_slots.assert_called_once()
+    invalidate_written_slots.assert_called_once()
+    assert invalidate_written_slots.call_args.args[0] == [leader]
 
 
 def test_hisparse_shared_host_reader_skips_transfer_completion():
@@ -1086,7 +1099,7 @@ def test_hisparse_step_waits_for_previous_host_write(monkeypatch, is_host_writer
     worker._per_layer_mirrored = set()
     worker._submitted_mirror_layers = set()
     stream = MagicMock()
-    monkeypatch.setattr(hisparse_worker_module, "current_stream", lambda: stream)
+    monkeypatch.setattr(torch.accelerator, "current_stream", lambda: stream)
 
     worker.start_step(
         SimpleNamespace(
@@ -1172,7 +1185,7 @@ def test_hisparse_empty_step_does_not_replay_stale_host_mirror(monkeypatch):
     enqueue_host_mirror = MagicMock(wraps=worker._enqueue_host_mirror)
     monkeypatch.setattr(worker, "_enqueue_host_mirror", enqueue_host_mirror)
     stream = MagicMock()
-    monkeypatch.setattr(hisparse_worker_module, "current_stream", lambda: stream)
+    monkeypatch.setattr(torch.accelerator, "current_stream", lambda: stream)
 
     worker.start_step(
         SimpleNamespace(
@@ -1195,6 +1208,9 @@ def test_hisparse_empty_step_does_not_replay_stale_host_mirror(monkeypatch):
 
 def test_hisparse_cache_handles_join_index_groups_during_construction(monkeypatch):
     """Followers must not allocate duplicate runtime state before profiling."""
+    monkeypatch.setattr(
+        hisparse_runtime_module, "get_hisparse_backend", CUDAHiSparseBackend
+    )
     config = SimpleNamespace(
         scheduler_config=SimpleNamespace(
             max_num_seqs=2,
@@ -1209,7 +1225,7 @@ def test_hisparse_cache_handles_join_index_groups_during_construction(monkeypatc
         device_buffer_size=8,
         max_union_rows=8,
     )
-    monkeypatch.setattr(hisparse_runtime_module, "_has_hisparse_ops", lambda: True)
+    monkeypatch.setattr(CUDAHiSparseBackend, "check_available", lambda self: None)
     monkeypatch.setattr(
         hisparse_runtime_module.ResolvedHiSparseConfig,
         "from_vllm_config",
@@ -1230,8 +1246,16 @@ def create_stream(_device):
         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())
+    monkeypatch.setattr(
+        torch,
+        "Stream",
+        lambda device: object(),
+    )
+    monkeypatch.setattr(
+        torch,
+        "Event",
+        lambda: object(),
+    )
     index_group_builder = index_group_module.SparseMLAIndexGroupBuilder(
         torch.empty((2, 4), dtype=torch.int32)
     )
@@ -1368,6 +1392,51 @@ def test_bind_kv_cache_shares_replayssm_trackers_by_cache_group(layers_only):
     assert [m._updates_replayssm_trackers for m in mixers] == [False, True, True]
 
 
+def test_rebind_kv_cache_keeps_replayssm_trackers():
+    """A rebind over the same or fewer blocks (the extensible KV cache rebinds
+    after CUDA graph capture) must keep the tracker tensors the captured graphs
+    address; a rebind needing more blocks gets new ones."""
+    mixers = [_TestReplaySSMMixer() for _ in range(2)]
+    layer_names = [f"layers.{i}.mixer" for i in range(2)]
+    ctx = dict(zip(layer_names, mixers))
+    groups = [SimpleNamespace(layer_names=layer_names)]
+    bind_kv_cache_to_layers(
+        {n: _packed_replayssm_cache(8) for n in layer_names},
+        ctx,
+        kv_cache_groups=groups,
+    )
+    ring_start = mixers[0]._replayssm_ring_start
+    prev_accepted = mixers[0]._replayssm_prev_num_accepted
+    prev_query_len = mixers[0]._replayssm_prev_query_len
+    ring_start[3] = 7
+
+    bind_kv_cache_to_layers(
+        {n: _packed_replayssm_cache(5) for n in layer_names},
+        ctx,
+        kv_cache_groups=groups,
+    )
+    # Same storage (captured graphs hold its address), length of the new cache.
+    for m in mixers:
+        assert m._replayssm_ring_start.data_ptr() == ring_start.data_ptr()
+        assert m._replayssm_ring_start.shape == (5,)
+        assert m._replayssm_prev_num_accepted.data_ptr() == prev_accepted.data_ptr()
+        assert m._replayssm_prev_query_len.data_ptr() == prev_query_len.data_ptr()
+        assert m._replayssm_prev_query_len.shape == (5,)
+        assert m._replayssm_ring_start[3] == 7
+    assert mixers[0]._replayssm_ring_start is mixers[1]._replayssm_ring_start
+
+    bind_kv_cache_to_layers(
+        {n: _packed_replayssm_cache(9) for n in layer_names},
+        ctx,
+        kv_cache_groups=groups,
+    )
+    assert mixers[0]._replayssm_ring_start.data_ptr() != ring_start.data_ptr()
+    assert mixers[0]._replayssm_ring_start.shape == (9,)
+    assert all(
+        m._replayssm_ring_start is mixers[0]._replayssm_ring_start for m in mixers
+    )
+
+
 def test_bind_kv_cache(default_vllm_config):
     from vllm.model_executor.layers.attention import Attention
 
diff --git a/tests/v1/worker/test_workspace.py b/tests/v1/worker/test_workspace.py
index 23a489b51975..8f1b80503794 100644
--- a/tests/v1/worker/test_workspace.py
+++ b/tests/v1/worker/test_workspace.py
@@ -19,6 +19,11 @@ def __init__(self, dspark: bool) -> None:
     def use_dspark(self) -> bool:
         return self._dspark
 
+    @property
+    def uses_separate_workspace_lane(self) -> bool:
+        # Only DSpark drafts in a separate workspace lane.
+        return self._dspark
+
 
 class _VllmConfig:
     def __init__(self, spec_config: _SpecConfig | None) -> None:
@@ -187,3 +192,64 @@ def test_persistent_tensor_preserves_contents_and_rejects_changes() -> None:
     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)
+
+
+def test_retained_span_survives_small_scratch_at_the_tail(monkeypatch) -> None:
+    """The retained span sits at the workspace end, so head-sized scratch
+    requests leave its bytes intact and the next same-key call reports so."""
+    monkeypatch.setattr(workspace, "dbo_current_ubatch_id", lambda: 0)
+    manager = workspace.WorkspaceManager(torch.device("cpu"))
+    (full,) = manager.get_simultaneous(((4096,), torch.uint8))
+
+    span, intact = manager.get_retained("kv", (256,), torch.int32)
+    assert not intact
+    assert span.data_ptr() == full.data_ptr() + 4096 - 1024
+    span.copy_(torch.arange(256, dtype=torch.int32))
+
+    (head,) = manager.get_simultaneous(((3072,), torch.uint8))
+    head.fill_(7)
+    again, intact = manager.get_retained("kv", (256,), torch.int32)
+    assert intact
+    assert again.data_ptr() == span.data_ptr()
+    assert torch.equal(again, torch.arange(256, dtype=torch.int32))
+
+
+@pytest.mark.parametrize(
+    "disturb",
+    ["scratch_overlap", "other_key", "other_size", "reallocation"],
+)
+def test_retained_span_reports_disturbance(monkeypatch, disturb: str) -> None:
+    monkeypatch.setattr(workspace, "dbo_current_ubatch_id", lambda: 0)
+    manager = workspace.WorkspaceManager(torch.device("cpu"))
+    manager.get_simultaneous(((4096,), torch.uint8))
+    manager.get_retained("kv", (256,), torch.int32)
+
+    if disturb == "scratch_overlap":
+        # One byte past the head of the span is enough.
+        manager.get_simultaneous(((3073,), torch.uint8))
+    elif disturb == "other_key":
+        manager.get_retained("other", (256,), torch.int32)
+    elif disturb == "other_size":
+        manager.get_retained("kv", (128,), torch.int32)
+    else:
+        manager.get_simultaneous(((8192,), torch.uint8))
+
+    _, intact = manager.get_retained("kv", (256,), torch.int32)
+    assert not intact
+    # The call itself re-arms the span.
+    _, intact = manager.get_retained("kv", (256,), torch.int32)
+    assert intact
+
+
+def test_retained_span_is_per_slot(monkeypatch) -> None:
+    monkeypatch.setattr(workspace, "dbo_current_ubatch_id", lambda: 0)
+    manager = workspace.WorkspaceManager(torch.device("cpu"), num_lanes=2)
+    manager.get_simultaneous(((4096,), torch.uint8))
+    manager.get_retained("kv", (256,), torch.int32)
+    with workspace.use_workspace_lane(1):
+        # Another lane's scratch and span requests do not touch lane 0.
+        manager.get_simultaneous(((4096,), torch.uint8))
+        _, intact = manager.get_retained("kv", (256,), torch.int32)
+        assert not intact
+    _, intact = manager.get_retained("kv", (256,), torch.int32)
+    assert intact
diff --git a/tests/watermarking/test_watermarking.py b/tests/watermarking/test_watermarking.py
index dd1e90926d7a..cf76f772d017 100644
--- a/tests/watermarking/test_watermarking.py
+++ b/tests/watermarking/test_watermarking.py
@@ -37,6 +37,7 @@
     repeated_context_mask,
 )
 from vllm.v1.worker.gpu.spec_decode.dflash.speculator import DFlashSpeculator
+from vllm.v1.worker.gpu.spec_decode.draft_logits import DraftLogitsCache
 from vllm.v1.worker.gpu.spec_decode.dspark.speculator import DSparkSpeculator
 from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator
 
@@ -1387,7 +1388,9 @@ def sample(logits, contexts, random_sampler=None, skip_mask=None):
         temperature=torch.ones(2),
         seeds=torch.zeros(2, dtype=torch.int64),
         draft_step=torch.tensor(0),
-        draft_logits=torch.zeros(2, 1, 8),
+        draft_logits=DraftLogitsCache(
+            2, 1, 8, dtype=torch.float32, fill=0.0, device="cpu"
+        ),
     )
 
     assert torch.equal(sampled, torch.tensor([7, 4]))
@@ -1755,7 +1758,9 @@ def ordinary_sample(logits, idx_mapping, temperature, seed, pos, **kwargs):
     speculator.acceptance_estimator = None
     batch_slots = torch.tensor([2, 0, 1])
     logits = torch.arange(num_reqs * num_steps * vocab_size, dtype=torch.float32)
-    draft_logits = torch.zeros(num_reqs, num_steps, vocab_size)
+    draft_logits = DraftLogitsCache(
+        num_reqs, num_steps, vocab_size, dtype=torch.float32, fill=0.0, device="cpu"
+    )
 
     sampled = speculator.sample_draft(
         hidden_states=logits.view(-1, vocab_size),
@@ -1769,7 +1774,9 @@ def ordinary_sample(logits, idx_mapping, temperature, seed, pos, **kwargs):
     )
 
     assert torch.equal(sampled, torch.tensor([3, 4, 102, 103, 204, 205]))
-    assert torch.equal(draft_logits[batch_slots], logits.view(num_reqs, num_steps, -1))
+    assert torch.equal(
+        draft_logits.values[batch_slots], logits.view(num_reqs, num_steps, -1)
+    )
 
 
 @pytest.mark.skipif(
@@ -1860,9 +1867,9 @@ def run(fused: bool):
 
 def test_dspark_reduced_vocab_draft_sampler_applies_watermarking(monkeypatch):
     speculator = object.__new__(DSparkSpeculator)
-    speculator.draft_logits = torch.zeros(2, 1, 8)
-    speculator.draft_logits_num_shards = 1
-    speculator.draft_logits_shard_idx = 0
+    speculator.draft_logits = DraftLogitsCache(
+        2, 1, 8, dtype=torch.float32, fill=0.0, device="cpu"
+    )
     speculator._d2t_scatter_index = torch.tensor([1, 5])
     speculator._draft_scatter_buf = torch.full((2, 8), float("-inf"))
     speculator.temperature = torch.ones(2)
@@ -1906,9 +1913,9 @@ def sample(
 
 def test_dspark_target_only_watermarking_leaves_drafts_unwatermarked(monkeypatch):
     speculator = object.__new__(DSparkSpeculator)
-    speculator.draft_logits = torch.zeros(2, 1, 8)
-    speculator.draft_logits_num_shards = 1
-    speculator.draft_logits_shard_idx = 0
+    speculator.draft_logits = DraftLogitsCache(
+        2, 1, 8, dtype=torch.float32, fill=0.0, device="cpu"
+    )
     speculator._d2t_scatter_index = None
     speculator.temperature = torch.ones(2)
     speculator.seeds = torch.zeros(2, dtype=torch.int64)
diff --git a/tools/build_nixl_with_torch215_for_rubin.py b/tools/build_nixl_with_torch215_for_rubin.py
deleted file mode 100644
index ee9fa72ae959..000000000000
--- a/tools/build_nixl_with_torch215_for_rubin.py
+++ /dev/null
@@ -1,619 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-#
-# TODO(arpera):
-# This script is temporary. It builds nixl from source code against torch 2.15
-# because currently there is no nixl wheel built against torch 2.15 which
-# vLLM docker image for Rubin requires.
-# Remove this script from vLLM once there is nixl wheel built against torch 2.15.
-
-import argparse
-import glob
-import os
-import shutil
-import subprocess
-import sys
-import urllib.request
-
-# --- Configuration ---
-WHEELS_CACHE_HOME = os.environ.get("WHEELS_CACHE_HOME", "/tmp/wheels_cache")
-ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
-UCX_DIR = os.path.join("/tmp", "ucx_source")
-NIXL_DIR = os.path.join("/tmp", "nixl_source")
-UCX_INSTALL_DIR = os.path.join("/tmp", "ucx_install")
-NIXL_BUILD_DIR = os.path.join("/tmp", "nixl_build")
-BUILD_TOOLS_DIR = os.path.join("/tmp", "nixl_build_tools")
-UCX_REPO_URL = "https://github.com/openucx/ucx.git"
-NIXL_REPO_URL = "https://github.com/ai-dynamo/nixl.git"
-UCX_REF = os.environ.get("UCX_REF", "v1.23.x")
-TEMP_WHEEL_DIR = os.path.join("/tmp", "nixl_temp_wheelhouse")
-NIXL_REQUIREMENTS_FILE = os.environ.get(
-    "NIXL_REQUIREMENTS_FILE",
-    os.path.join(ROOT_DIR, os.pardir, "requirements", "kv_connectors.txt"),
-)
-MAX_JOBS = int(os.environ.get("MAX_JOBS") or os.cpu_count() or 1)
-CUDA_HOME = os.environ.get("CUDA_HOME", "/usr/local/cuda")
-# SM targets NIXL EP can run on (it requires sm_90 or newer).
-NIXL_EP_CUDA_ARCHS = ("90", "100", "103", "107", "110", "120")
-
-# Libraries that must come from the host or from torch, not from the wheel.
-AUDITWHEEL_EXCLUDES = [
-    "libcuda*",
-    "libcufile*",
-    "libcuobjclient*",
-    "libssl*",
-    "libcrypto*",
-    "libefa*",
-    "libhwloc*",
-    "libfabric*",
-    "libtorch*",
-    "libc10*",
-    "libdoca*",
-    "libred_client*",
-    "libred_async*",
-    "liblz4*",
-]
-
-
-# --- Helper Functions ---
-def get_latest_nixl_version():
-    """Helper function to get latest release version of NIXL.
-
-    Raises:
-        RuntimeError: If the latest release cannot be determined.
-
-    """
-    # /releases/latest redirects to the tag page. Unlike the REST API it is not
-    # rate limited for anonymous clients, which are shared on compute nodes.
-    request = urllib.request.Request(
-        f"{NIXL_REPO_URL.removesuffix('.git')}/releases/latest", method="HEAD"
-    )
-    try:
-        with urllib.request.urlopen(request, timeout=30) as response:
-            return response.geturl().rsplit("/", 1)[-1]
-    except OSError as e:
-        raise RuntimeError(
-            "Cannot determine the latest NIXL release, set NIXL_VERSION."
-        ) from e
-
-
-def get_pinned_nixl_version():
-    """Returns the NIXL version pinned in the vLLM requirements, if there is one."""
-    try:
-        with open(NIXL_REQUIREMENTS_FILE) as requirements:
-            for line in requirements:
-                requirement = line.split("#")[0].split(";")[0]
-                name, separator, version = requirement.partition("==")
-                if separator and name.strip() == "nixl" and version.split():
-                    return version.split()[0]
-    except OSError:
-        pass
-    return None
-
-
-def resolve_nixl_version():
-    """Returns the NIXL version to build and where it comes from."""
-    if os.environ.get("NIXL_VERSION"):
-        return os.environ["NIXL_VERSION"], "the NIXL_VERSION environment variable"
-    pinned_version = get_pinned_nixl_version()
-    if pinned_version:
-        return pinned_version, os.path.normpath(NIXL_REQUIREMENTS_FILE)
-    return get_latest_nixl_version(), "the latest NIXL release"
-
-
-NIXL_VERSION, NIXL_VERSION_SOURCE = resolve_nixl_version()
-# Version as it appears in wheel file names (tags may carry a "v" prefix).
-NIXL_WHEEL_VERSION = NIXL_VERSION.removeprefix("v")
-
-
-def run_command(command, cwd=".", env=None):
-    """Helper function to run a shell command and check for errors."""
-    print(f"--> Running command: {' '.join(command)} in '{cwd}'", flush=True)
-    subprocess.check_call(command, cwd=cwd, env=env)
-
-
-def is_pip_package_installed(package_name):
-    """Checks if a package is installed via pip without raising an exception."""
-    result = subprocess.run(
-        [sys.executable, "-m", "pip", "show", package_name],
-        stdout=subprocess.DEVNULL,
-        stderr=subprocess.DEVNULL,
-    )
-    return result.returncode == 0
-
-
-def find_nixl_wheel_in_cache(cache_dir):
-    """Finds the nixl-cuXX wheel file in the specified cache directory."""
-    # The repaired wheel will have a 'manylinux' tag, but this glob still works.
-    search_pattern = os.path.join(cache_dir, f"nixl_*{NIXL_WHEEL_VERSION}*.whl")
-    wheels = glob.glob(search_pattern)
-    if wheels:
-        # Sort to get the most recent/highest version if multiple exist
-        wheels.sort()
-        return wheels[-1]
-    return None
-
-
-def find_nixl_meta_wheel(wheel_dir):
-    """Finds the `nixl` meta package wheel (the nixl_cuXX dispatcher)."""
-    wheels = glob.glob(os.path.join(wheel_dir, f"nixl-{NIXL_WHEEL_VERSION}-*.whl"))
-    return wheels[-1] if wheels else None
-
-
-def get_cuda_version():
-    """Returns the (major, minor) CUDA toolkit version, or None without CUDA."""
-    nvcc = os.path.join(CUDA_HOME, "bin", "nvcc")
-    if not os.path.exists(nvcc):
-        return None
-    output = subprocess.check_output([nvcc, "--version"], text=True)
-    _, separator, release = output.partition("release ")
-    try:
-        if not separator:
-            raise ValueError
-        major, minor = release.split(",")[0].split(".")[:2]
-        return int(major), int(minor)
-    except ValueError:
-        raise RuntimeError(
-            f"Cannot parse the CUDA version from `{nvcc} --version`."
-        ) from None
-
-
-def get_default_cuda_arch_list():
-    """Returns the SM targets of NIXL EP that the installed nvcc can compile."""
-    nvcc = os.path.join(CUDA_HOME, "bin", "nvcc")
-    supported = subprocess.check_output([nvcc, "--list-gpu-code"], text=True).split()
-    return ",".join(sm for sm in NIXL_EP_CUDA_ARCHS if f"sm_{sm}" in supported)
-
-
-def is_torch_available():
-    """Checks whether torch can be imported by the current interpreter."""
-    result = subprocess.run(
-        [sys.executable, "-c", "import torch"],
-        stdout=subprocess.DEVNULL,
-        stderr=subprocess.DEVNULL,
-    )
-    return result.returncode == 0
-
-
-def is_nixl_installed(build_nixl_ep):
-    """Checks that the requested NIXL version, and its EP extension for the
-    installed torch when requested, is already installed."""
-    cuda_version = get_cuda_version()
-    if cuda_version is None:
-        return is_pip_package_installed("nixl")
-    checks = [
-        "import importlib, importlib.metadata as m",
-        f"assert m.version('nixl') == '{NIXL_WHEEL_VERSION}'",
-        f"assert m.version('nixl-cu{cuda_version[0]}') == '{NIXL_WHEEL_VERSION}'",
-    ]
-    if build_nixl_ep:
-        # The package picks its extension by torch version and fails to import
-        # when none was built for the installed torch.
-        checks.append(f"importlib.import_module('nixl_ep_cu{cuda_version[0]}')")
-    result = subprocess.run(
-        [sys.executable, "-c", "\n".join(checks)],
-        stdout=subprocess.DEVNULL,
-        stderr=subprocess.DEVNULL,
-    )
-    return result.returncode == 0
-
-
-def install_system_dependencies(build_nixl_ep=False):
-    """Installs required system packages using apt-get if run as root."""
-    if shutil.which("apt-get") is None:
-        print("--> apt-get not found, relying on the build tools in PATH.", flush=True)
-        return
-    if os.geteuid() != 0:
-        print("\n---", flush=True)
-        print(
-            "WARNING: Not running as root. \
-            Skipping system dependency installation.",
-            flush=True,
-        )
-        print(
-            "Please ensure the listed packages are installed on your system:",
-            flush=True,
-        )
-        print(
-            "  patchelf build-essential git cmake ninja-build \
-            autotools-dev automake libtool libtool-bin \
-            libibverbs-dev",
-            flush=True,
-        )
-        print("---\n", flush=True)
-        return
-
-    print("--- Running as root. Installing system dependencies... ---", flush=True)
-    apt_packages = [
-        "patchelf",  # <-- Add patchelf here
-        "build-essential",
-        "git",
-        "cmake",
-        "ninja-build",
-        "autotools-dev",
-        "automake",
-        "libtool",
-        "libtool-bin",
-        "pkg-config",
-        "libibverbs-dev",
-    ]
-    # -base CUDA images lack the headers UCX (NVML) and the torch extension
-    # (cuSPARSE/cuSOLVER/cuBLAS, included by ATen/cuda/CUDAContext.h) need.
-    cuda_version = get_cuda_version()
-    if cuda_version is not None:
-        major, minor = cuda_version
-        if not os.path.exists(os.path.join(CUDA_HOME, "include", "nvml.h")):
-            apt_packages.append(f"cuda-nvml-dev-{major}-{minor}")
-        if build_nixl_ep:
-            apt_packages += [
-                f"{package}-{major}-{minor}"
-                for package in ("libcusparse-dev", "libcublas-dev", "libcusolver-dev")
-            ]
-    run_command(["apt-get", "update"])
-    run_command(["apt-get", "install", "-y"] + apt_packages)
-    print("--- System dependencies installed successfully. ---\n", flush=True)
-
-
-def install_python_build_dependencies():
-    """Installs the Python tools needed to build and repair the NIXL wheel."""
-    run_command(
-        [
-            sys.executable,
-            "-m",
-            "pip",
-            "install",
-            "auditwheel",
-            "meson-python",
-            "ninja",
-            "patchelf",
-            "pybind11",
-            "pyyaml",
-            "setuptools>=80.9.0",
-            "tomlkit",
-        ]
-    )
-    # An older meson already on the system (e.g. from the distribution) makes
-    # meson-python fail to package NIXL and pip would keep it, so use our own.
-    run_command(
-        [
-            sys.executable,
-            "-m",
-            "pip",
-            "install",
-            "--upgrade",
-            f"--target={BUILD_TOOLS_DIR}",
-            "meson",
-        ]
-    )
-
-
-def build_ucx(ucx_install_path):
-    """Builds UCX from source, with CUDA support if CUDA is installed."""
-    if not os.path.exists(UCX_DIR):
-        run_command(["git", "clone", UCX_REPO_URL, UCX_DIR])
-    ucx_source_path = os.path.abspath(UCX_DIR)
-    run_command(["git", "checkout", UCX_REF], cwd=ucx_source_path)
-    run_command(["./autogen.sh"], cwd=ucx_source_path)
-    configure_command = [
-        "./contrib/configure-release-mt",
-        f"--prefix={ucx_install_path}",
-        "--enable-shared",
-        "--disable-static",
-        "--disable-doxygen-doc",
-        "--enable-experimental-api",
-        "--enable-optimizations",
-        "--without-avx",
-        "--enable-cma",
-        "--enable-devel-headers",
-        *([f"--with-cuda={CUDA_HOME}"] if get_cuda_version() else []),
-        "--with-verbs",
-        "--without-gdrcopy",
-        "--without-dc",
-        "--without-rdmacm",
-        "--without-gga",
-        "--with-ze=no",
-    ]
-    run_command(configure_command, cwd=ucx_source_path)
-    run_command(["make", "-j", str(MAX_JOBS)], cwd=ucx_source_path)
-    run_command(["make", "install"], cwd=ucx_source_path)
-
-
-def build_nixl_wheel(ucx_path, build_env, wheel_dir, args):
-    """Builds the (unrepaired) NIXL wheel against the UCX installed in ucx_path."""
-    if not os.path.exists(NIXL_DIR):
-        run_command(["git", "clone", NIXL_REPO_URL, NIXL_DIR])
-    else:
-        run_command(["git", "fetch", "--tags"], cwd=NIXL_DIR)
-    nixl_source_path = os.path.abspath(NIXL_DIR)
-    # Tags are "v<version>" for current releases and bare for older ones.
-    for tag in (f"v{NIXL_WHEEL_VERSION}", NIXL_WHEEL_VERSION):
-        tag_exists = subprocess.run(
-            ["git", "rev-parse", "--verify", "--quiet", f"refs/tags/{tag}"],
-            cwd=nixl_source_path,
-            stdout=subprocess.DEVNULL,
-        )
-        if tag_exists.returncode == 0:
-            run_command(["git", "checkout", "--force", tag], cwd=nixl_source_path)
-            print(f"--> Checked out NIXL version: {tag}", flush=True)
-            break
-    else:
-        raise RuntimeError(f"NIXL has no tag for version {NIXL_WHEEL_VERSION}.")
-
-    cuda_version = get_cuda_version()
-    if cuda_version is not None:
-        # The wheel name selects which of nixl-cu12 / nixl-cu13 is built.
-        run_command(
-            [
-                sys.executable,
-                "contrib/tomlutil.py",
-                "--wheel-name",
-                f"nixl-cu{cuda_version[0]}",
-                "pyproject.toml",
-            ],
-            cwd=nixl_source_path,
-        )
-
-    meson_options = {
-        "build_tests": "false",
-        "ucx_path": ucx_path,
-    }
-    cuda_arch_list = args.cuda_arch_list
-    if args.build_nixl_ep:
-        meson_options["build_nixl_ep"] = "true"
-        meson_options["build_examples"] = "true"
-        cuda_arch_list = cuda_arch_list or get_default_cuda_arch_list()
-    if cuda_arch_list:
-        meson_options["nixl_cuda_arch_list"] = cuda_arch_list
-    if args.nixl_plugins:
-        meson_options["enable_plugins"] = args.nixl_plugins
-    config_settings = [f"-Csetup-args=-D{k}={v}" for k, v in meson_options.items()]
-    # Keep the build tree: auditwheel needs the subproject libraries it holds.
-    config_settings.append(f"-Cbuild-dir={NIXL_BUILD_DIR}")
-    config_settings.append(f"-Ccompile-args=-j{MAX_JOBS}")
-
-    run_command(
-        [
-            sys.executable,
-            "-m",
-            "pip",
-            "wheel",
-            ".",
-            "--no-deps",
-            # Build against the torch already installed, not a fresh isolated one.
-            "--no-build-isolation",
-            f"--wheel-dir={wheel_dir}",
-            *config_settings,
-        ],
-        cwd=nixl_source_path,
-        env=build_env,
-    )
-
-
-def repair_nixl_wheel(unrepaired_wheel, repaired_dir, ucx_install_path, build_env):
-    """Bundles the shared libraries and UCX modules into a self-contained wheel."""
-    auditwheel_command = ["auditwheel", "repair"]
-    for pattern in AUDITWHEEL_EXCLUDES:
-        auditwheel_command += ["--exclude", pattern]
-    auditwheel_command += [unrepaired_wheel, f"--wheel-dir={repaired_dir}"]
-    # The wheel ships prometheus-cpp as libcore.so.1.3.0 but its users ask for the
-    # soname libcore.so.1.3, which only exists as a symlink in the build tree.
-    repair_env = build_env.copy()
-    prometheus_dir = os.path.join(NIXL_BUILD_DIR, "subprojects", "prometheus-cpp")
-    repair_env["LD_LIBRARY_PATH"] = f"{prometheus_dir}:{build_env['LD_LIBRARY_PATH']}"
-    run_command(auditwheel_command, env=repair_env)
-
-    repaired_wheel = find_nixl_wheel_in_cache(repaired_dir)
-    if not repaired_wheel:
-        raise RuntimeError("Failed to find the repaired NIXL wheel.")
-    # UCX loads its transports (CUDA, IB, ...) with dlopen, so auditwheel does not
-    # see them as dependencies; the NIXL plugin itself is already in the wheel.
-    run_command(
-        [
-            sys.executable,
-            "contrib/wheel_add_ucx_plugins.py",
-            "--skip-nixl-plugins",
-            f"--ucx-plugins-dir={os.path.join(ucx_install_path, 'lib', 'ucx')}",
-            repaired_wheel,
-        ],
-        cwd=os.path.abspath(NIXL_DIR),
-        env=build_env,
-    )
-    return repaired_wheel
-
-
-def install_nixl_wheels(wheels):
-    """Installs the NIXL wheels and drops the nixl-cuXX of other CUDA majors."""
-    # --force-reinstall: the wheels usually have the version of the release that
-    # is already installed, which pip would otherwise keep.
-    # w/o "no-deps", it will install cuda-torch
-    run_command(
-        [
-            sys.executable,
-            "-m",
-            "pip",
-            "install",
-            "--force-reinstall",
-            "--no-deps",
-            *wheels,
-        ]
-    )
-
-    cuda_version = get_cuda_version()
-    if cuda_version is None:
-        return
-    # The dispatcher only loads nixl-cu<installed CUDA major>, and the meta
-    # package built here does not depend on the other variant.
-    for cuda_major in (12, 13):
-        package = f"nixl-cu{cuda_major}"
-        if cuda_major != cuda_version[0] and is_pip_package_installed(package):
-            run_command([sys.executable, "-m", "pip", "uninstall", "-y", package])
-
-
-def build_and_install_prerequisites(args):
-    """Builds UCX and NIXL from source, creating a self-contained wheel."""
-    print(f"--> NIXL version {NIXL_WHEEL_VERSION} (from {NIXL_VERSION_SOURCE})")
-    if args.build_nixl_ep is None:
-        args.build_nixl_ep = get_cuda_version() is not None and is_torch_available()
-    reuse_existing = not (args.force_reinstall or args.build_only)
-    if reuse_existing and is_nixl_installed(args.build_nixl_ep):
-        print("--> NIXL is already installed. Nothing to do.", flush=True)
-        return
-
-    cached_wheel = find_nixl_wheel_in_cache(WHEELS_CACHE_HOME)
-    if reuse_existing and cached_wheel:
-        print(
-            f"\n--> Found self-contained wheel: \
-                {os.path.basename(cached_wheel)}.",
-            flush=True,
-        )
-        print("--> Installing from cache, skipping all source builds.", flush=True)
-        cached_wheels = [cached_wheel]
-        cached_meta_wheel = find_nixl_meta_wheel(WHEELS_CACHE_HOME)
-        if cached_meta_wheel:
-            cached_wheels.append(cached_meta_wheel)
-        install_nixl_wheels(cached_wheels)
-        print("\n--- Installation from cache complete. ---", flush=True)
-        return
-
-    print(
-        "\n--> No installed package or cached wheel found. \
-         Starting full build process...",
-        flush=True,
-    )
-    if not args.skip_system_deps:
-        install_system_dependencies(args.build_nixl_ep)
-    print("\n--> Installing Python build dependencies...", flush=True)
-    install_python_build_dependencies()
-    print(f"--> Using wheel cache directory: {WHEELS_CACHE_HOME}", flush=True)
-    os.makedirs(WHEELS_CACHE_HOME, exist_ok=True)
-
-    # -- Step 1: Build UCX from source, unless an existing install is given --
-    # With an external UCX (e.g. HPC-X) the wheel keeps depending on it at runtime.
-    external_ucx = args.ucx_path is not None
-    if external_ucx:
-        ucx_install_path = os.path.abspath(args.ucx_path)
-        print(f"\n[1/3] Using the existing UCX in {ucx_install_path}", flush=True)
-    else:
-        ucx_install_path = os.path.abspath(UCX_INSTALL_DIR)
-        print("\n[1/3] Configuring and building UCX from source...", flush=True)
-        build_ucx(ucx_install_path)
-        print("--- UCX build and install complete ---", flush=True)
-
-    # -- Step 2: Build NIXL wheel from source --
-    print("\n[2/3] Building NIXL wheel from source...", flush=True)
-    ucx_lib_path = os.path.join(ucx_install_path, "lib")
-    ucx_plugin_path = os.path.join(ucx_lib_path, "ucx")
-    build_env = os.environ.copy()
-    # pip-installed build tools live next to the running interpreter.
-    build_env["PATH"] = os.pathsep.join(
-        [
-            os.path.join(BUILD_TOOLS_DIR, "bin"),
-            os.path.join(CUDA_HOME, "bin"),
-            os.path.dirname(sys.executable),
-            build_env["PATH"],
-        ]
-    )
-    build_env["PYTHONPATH"] = os.pathsep.join(
-        filter(None, [BUILD_TOOLS_DIR, build_env.get("PYTHONPATH")])
-    )
-    build_env["PKG_CONFIG_PATH"] = os.path.join(ucx_lib_path, "pkgconfig")
-    existing_ld_path = os.environ.get("LD_LIBRARY_PATH", "")
-    build_env["LD_LIBRARY_PATH"] = (
-        f"{ucx_lib_path}:{ucx_plugin_path}:{existing_ld_path}".strip(":")
-    )
-    print(f"--> Using LD_LIBRARY_PATH: {build_env['LD_LIBRARY_PATH']}", flush=True)
-
-    build_nixl_wheel(ucx_install_path, build_env, TEMP_WHEEL_DIR, args)
-
-    # -- Step 3: Repair the wheel by copying UCX libraries --
-    unrepaired_wheel = find_nixl_wheel_in_cache(TEMP_WHEEL_DIR)
-    if not unrepaired_wheel:
-        raise RuntimeError("Failed to find the NIXL wheel after building it.")
-    if external_ucx:
-        print("\n[3/3] Skipping the repair step: UCX is provided externally.")
-        built_wheels = [unrepaired_wheel]
-    else:
-        print("\n[3/3] Repairing NIXL wheel to include UCX libraries...", flush=True)
-        built_wheels = [
-            repair_nixl_wheel(
-                unrepaired_wheel,
-                os.path.join(TEMP_WHEEL_DIR, "repaired"),
-                ucx_install_path,
-                build_env,
-            )
-        ]
-
-    meta_wheel = find_nixl_meta_wheel(
-        os.path.join(NIXL_BUILD_DIR, "src", "bindings", "python", "nixl-meta")
-    )
-    if meta_wheel:
-        built_wheels.append(meta_wheel)
-    else:
-        print("WARNING: the nixl meta package was not built.", flush=True)
-
-    new_wheels = []
-    for wheel in built_wheels:
-        run_command(["cp", wheel, WHEELS_CACHE_HOME])
-        new_wheels.append(os.path.join(WHEELS_CACHE_HOME, os.path.basename(wheel)))
-    run_command(["rm", "-rf", TEMP_WHEEL_DIR])
-
-    wheel_names = ", ".join(os.path.basename(w) for w in new_wheels)
-    if args.build_only:
-        print(f"--> Built wheels in {WHEELS_CACHE_HOME}: {wheel_names}", flush=True)
-        return
-    print(
-        f"--> Successfully built wheels: {wheel_names}. Now installing...", flush=True
-    )
-    install_nixl_wheels(new_wheels)
-    print("--- NIXL installation complete ---", flush=True)
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser(
-        description="Build and install UCX and NIXL dependencies."
-    )
-    parser.add_argument(
-        "--force-reinstall",
-        action="store_true",
-        help="Force rebuild and reinstall of UCX and NIXL \
-        even if they are already installed.",
-    )
-    parser.add_argument(
-        "--build-only",
-        action="store_true",
-        help="Build the wheels into WHEELS_CACHE_HOME and stop, without \
-        installing them or looking at what is already installed.",
-    )
-    parser.add_argument(
-        "--skip-system-deps",
-        action="store_true",
-        help="Do not install system packages with apt-get; the build tools \
-        (git, autotools, cmake, patchelf, ...) must already be on PATH.",
-    )
-    parser.add_argument(
-        "--build-nixl-ep",
-        action=argparse.BooleanOptionalAction,
-        default=None,
-        help="Build the NIXL EP extension for the installed torch version. \
-        Defaults to on when CUDA and torch are available.",
-    )
-    parser.add_argument(
-        "--cuda-arch-list",
-        default=os.environ.get("NIXL_CUDA_ARCH_LIST"),
-        help="Comma-separated CUDA SM targets for NIXL, e.g. 100,103,107. \
-        Defaults to the NIXL EP targets the installed nvcc supports.",
-    )
-    parser.add_argument(
-        "--nixl-plugins",
-        default=os.environ.get("NIXL_PLUGINS", "UCX"),
-        help="Comma-separated NIXL plugins to build. Pass an empty string to \
-        build every plugin whose dependencies are available.",
-    )
-    parser.add_argument(
-        "--ucx-path",
-        default=os.environ.get("UCX_PATH"),
-        help="Prefix of an existing UCX to build against instead of building one. \
-        The resulting wheel is not self-contained.",
-    )
-    args = parser.parse_args()
-    build_and_install_prerequisites(args)
diff --git a/tools/pre_commit/mypy.py b/tools/pre_commit/mypy.py
index 5c90eee54963..d68cf06f7cdf 100755
--- a/tools/pre_commit/mypy.py
+++ b/tools/pre_commit/mypy.py
@@ -49,9 +49,18 @@
     "tests/kernels",
     "tests/lora",
     "tests/model_executor",
+    "tests/models/deepseek_v32",
+    "tests/models/glm5next",
+    "tests/models/inkling",
+    "tests/models/kimi_k3",
     "tests/models/language",
+    "tests/models/multimodal",
+    "tests/models/minimax_m3",
     "tests/models/quantization",
+    "tests/models/qwen4_exp",
+    "tests/models/step5",
     "tests/models/test_initialization.py",
+    "tests/models/transformers",
     "tests/multimodal",
     "tests/parser",
     "tests/plugins/bge_m3_sparse_plugin",
@@ -95,7 +104,6 @@
     "tests/compile",
     "tests/distributed",
     "tests/models",
-    "tests/models/multimodal",
     "tests/v1",
     "tests/v1/core",
     "tests/v1/engine",
diff --git a/tools/pre_commit/test_tethering_allowlist.txt b/tools/pre_commit/test_tethering_allowlist.txt
index af3b990c8469..6a900d7f871c 100644
--- a/tools/pre_commit/test_tethering_allowlist.txt
+++ b/tools/pre_commit/test_tethering_allowlist.txt
@@ -104,7 +104,6 @@ tests/rocm/test_moe_weight_replay.py  # CI jobs disabled pending locked torchao
 
 # Temporary gaps from disabling jobs with unpinned install environments.
 tests/device_allocator/test_alloc_conf.py  # CI jobs disabled pending locked cmake installation
-tests/distributed/test_elastic_ep.py  # CI jobs disabled pending locked KV connectors installation
 tests/models/flint/test_decode_attention.py  # CI jobs disabled pending locked gguf installation
 tests/models/flint/test_flint_moe.py  # CI jobs disabled pending locked gguf installation
 tests/models/flint/test_flint_ops.py  # CI jobs disabled pending locked gguf installation
@@ -186,7 +185,6 @@ tests/tokenizers_/test_detokenize.py  # CI jobs disabled pending locked cmake in
 tests/tokenizers_/test_hf.py  # CI jobs disabled pending locked cmake installation
 tests/tokenizers_/test_mistral.py  # CI jobs disabled pending locked cmake installation
 tests/tokenizers_/test_registry.py  # CI jobs disabled pending locked cmake installation
-tests/tracing/test_loading_tracing.py  # CI jobs disabled pending locked tracing packages installation
 tests/transformers_utils/processors/test_pixtral.py  # CI jobs disabled pending locked cmake installation
 tests/transformers_utils/processors/test_voxtral.py  # CI jobs disabled pending locked cmake installation
 tests/transformers_utils/test_bailing_moe_v3_vl_config.py  # CI jobs disabled pending locked cmake installation
@@ -203,23 +201,8 @@ tests/transformers_utils/test_repo_utils.py  # CI jobs disabled pending locked c
 tests/transformers_utils/test_speculators_dspark_config.py  # CI jobs disabled pending locked cmake installation
 tests/transformers_utils/test_speculators_override.py  # CI jobs disabled pending locked cmake installation
 tests/transformers_utils/test_utils.py  # CI jobs disabled pending locked cmake installation
-tests/v1/determinism/test_batch_invariance_vlm.py  # CI jobs disabled pending locked pytest plugins installation
-tests/v1/determinism/test_cutlass_batch_invariance.py  # CI jobs disabled pending locked pytest plugins installation
-tests/v1/determinism/test_lora_shrink_batch_invariant.py  # CI jobs disabled pending locked pytest plugins installation
-tests/v1/determinism/test_matmul_batch_invariant.py  # CI jobs disabled pending locked pytest plugins installation
-tests/v1/determinism/test_mx_linear_batch_invariant.py  # CI jobs disabled pending locked pytest plugins installation
-tests/v1/determinism/test_nvfp4_batch_invariant.py  # CI jobs disabled pending locked pytest plugins installation
-tests/v1/determinism/test_rms_norm_batch_invariant.py  # CI jobs disabled pending locked pytest plugins installation
 tests/v1/fault_tolerance/test_engine_core_sentinel.py  # CI jobs disabled pending locked KV connectors installation
 tests/v1/fault_tolerance/test_fault_tolerance_e2e.py  # CI jobs disabled pending locked KV connectors installation
 tests/v1/kv_connector/mooncake_integration/test_accuracy.py  # CI jobs disabled pending locked KV connectors installation
-tests/v1/kv_connector/mooncake_integration/test_accuracy_rocm.py  # CI jobs disabled pending locked KV connectors installation
 tests/v1/kv_connector/mooncake_integration/test_mooncake_imports.py  # CI jobs disabled pending locked KV connectors installation
 tests/v1/kv_connector/moriio_integration/test_write_dummy_e2e_rocm.py  # CI jobs disabled pending locked KV connectors installation
-tests/v1/kv_connector/nixl_integration/test_accuracy.py  # CI jobs disabled pending locked KV connectors installation
-tests/v1/kv_connector/nixl_integration/test_accuracy_rocm.py  # CI jobs disabled pending locked KV connectors installation
-tests/v1/kv_connector/nixl_integration/test_edge_cases.py  # CI jobs disabled pending locked KV connectors installation
-tests/v1/kv_connector/nixl_integration/test_mamba_prefix_cache.py  # CI jobs disabled pending locked KV connectors installation
-tests/v1/kv_connector/nixl_integration/test_multi_connector_edge_cases.py  # CI jobs disabled pending locked KV connectors installation
-tests/v1/kv_connector/nixl_integration/test_nixl_imports.py  # CI jobs disabled pending locked KV connectors installation
-tests/v1/kv_connector/nixl_integration/test_spec_decode_acceptance.py  # CI jobs disabled pending locked KV connectors installation
diff --git a/tools/vllm-rocm/generate-rocm-wheels-root-index.sh b/tools/vllm-rocm/generate-rocm-wheels-root-index.sh
index 8ccfab17f79f..56c4cf4e4148 100755
--- a/tools/vllm-rocm/generate-rocm-wheels-root-index.sh
+++ b/tools/vllm-rocm/generate-rocm-wheels-root-index.sh
@@ -17,20 +17,14 @@
 #
 # Environment variables:
 #   S3_BUCKET   - Bucket name (default: vllm-wheels)
-#   ROCM_BASE_DOCKERFILE - default ROCm stack whose variant the root index serves
-#                          (default: docker/Dockerfile.rocm_base)
-#   VARIANT     - override the ROCm variant derived from ROCM_BASE_DOCKERFILE
+#   VARIANT     - ROCm variant the root index serves (required, e.g. rocm100)
 #   DRY_RUN     - Set to 1 for preview mode (same as --dry-run)
 
 set -euo pipefail
 
 # ======== Configuration ========
 BUCKET="${S3_BUCKET:-vllm-wheels}"
-if [[ -z "${VARIANT:-}" ]]; then
-    # shellcheck source=.buildkite/scripts/rocm/stack.sh
-    source .buildkite/scripts/rocm/stack.sh "${ROCM_BASE_DOCKERFILE:-docker/Dockerfile.rocm_base}"
-    VARIANT="$ROCM_STACK_VARIANT"
-fi
+VARIANT="${VARIANT:?set VARIANT, e.g. rocm100}"
 DRY_RUN="${DRY_RUN:-0}"
 FORCE_VERSION=""
 
diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py
index c8574a53bccd..c424c31e3552 100644
--- a/vllm/_custom_ops.py
+++ b/vllm/_custom_ops.py
@@ -1987,6 +1987,222 @@ def scaled_int8_quant(
     return output, input_scales, input_azp
 
 
+def gather_block_table_cpu(
+    block_table: torch.Tensor,
+    output: torch.Tensor,
+    idx_mapping: torch.Tensor,
+    num_blocks: torch.Tensor,
+    num_reqs_padded: int,
+) -> None:
+    torch.ops._C.gather_block_table_cpu(
+        block_table,
+        output,
+        idx_mapping,
+        num_blocks,
+        num_reqs_padded,
+    )
+
+
+def compute_slot_mapping_cpu(
+    block_table: torch.Tensor,
+    slot_mapping: torch.Tensor,
+    idx_mapping: torch.Tensor,
+    query_start_loc: torch.Tensor,
+    positions: torch.Tensor,
+    block_size: int,
+    mapping_enabled: bool,
+    pad_slot_id: int,
+) -> None:
+    torch.ops._C.compute_slot_mapping_cpu(
+        block_table,
+        slot_mapping,
+        idx_mapping,
+        query_start_loc,
+        positions,
+        block_size,
+        mapping_enabled,
+        pad_slot_id,
+    )
+
+
+def prepare_prefill_inputs_cpu(
+    input_ids: torch.Tensor,
+    next_prefill_tokens: torch.Tensor,
+    idx_mapping: torch.Tensor,
+    query_start_loc: torch.Tensor,
+    all_token_ids: torch.Tensor,
+    prefill_len: torch.Tensor,
+    num_computed_tokens: torch.Tensor,
+) -> None:
+    torch.ops._C.prepare_prefill_inputs_cpu(
+        input_ids,
+        next_prefill_tokens,
+        idx_mapping,
+        query_start_loc,
+        all_token_ids,
+        prefill_len,
+        num_computed_tokens,
+    )
+
+
+def get_prompt_logprobs_token_ids_cpu(
+    num_tokens: int,
+    query_start_loc: torch.Tensor,
+    idx_mapping: torch.Tensor,
+    num_computed_tokens: torch.Tensor,
+    all_token_ids: torch.Tensor,
+) -> torch.Tensor:
+    return torch.ops._C.get_prompt_logprobs_token_ids_cpu(
+        num_tokens,
+        query_start_loc,
+        idx_mapping,
+        num_computed_tokens,
+        all_token_ids,
+    )
+
+
+def prepare_pos_seq_lens_cpu(
+    idx_mapping: torch.Tensor,
+    query_start_loc: torch.Tensor,
+    num_computed_tokens: torch.Tensor,
+    pos: torch.Tensor,
+    seq_lens: torch.Tensor,
+) -> None:
+    torch.ops._C.prepare_pos_seq_lens_cpu(
+        idx_mapping,
+        query_start_loc,
+        num_computed_tokens,
+        pos,
+        seq_lens,
+    )
+
+
+def prepare_mrope_positions_cpu(
+    positions: torch.Tensor,
+    prefill_positions: torch.Tensor,
+    prefill_delta: torch.Tensor,
+    idx_mapping: torch.Tensor,
+    query_start_loc: torch.Tensor,
+    prefill_lens: torch.Tensor,
+    num_computed_tokens: torch.Tensor,
+) -> None:
+    torch.ops._C.prepare_mrope_positions_cpu(
+        positions,
+        prefill_positions,
+        prefill_delta,
+        idx_mapping,
+        query_start_loc,
+        prefill_lens,
+        num_computed_tokens,
+    )
+
+
+def combine_sampled_and_draft_tokens_cpu(
+    input_ids: torch.Tensor,
+    idx_mapping: torch.Tensor,
+    last_sampled_tokens: torch.Tensor,
+    query_start_loc: torch.Tensor,
+    seq_lens: torch.Tensor,
+    prefill_len: torch.Tensor,
+    draft_tokens: torch.Tensor,
+    cu_num_logits: torch.Tensor,
+    num_logits: int,
+    num_new_sampled_tokens: int,
+) -> torch.Tensor:
+    return torch.ops._C.combine_sampled_and_draft_tokens_cpu(
+        input_ids,
+        idx_mapping,
+        last_sampled_tokens,
+        query_start_loc,
+        seq_lens,
+        prefill_len,
+        draft_tokens,
+        cu_num_logits,
+        num_logits,
+        num_new_sampled_tokens,
+    )
+
+
+def post_update_cpu(
+    idx_mapping: torch.Tensor,
+    num_computed_tokens: torch.Tensor,
+    last_sampled_tokens: torch.Tensor,
+    output_bin_counts: torch.Tensor | None,
+    sampled_tokens: torch.Tensor,
+    num_sampled: torch.Tensor,
+    num_rejected: torch.Tensor,
+    query_start_loc: torch.Tensor | None,
+    all_token_ids: torch.Tensor,
+    total_len: torch.Tensor,
+) -> None:
+    torch.ops._C.post_update_cpu(
+        idx_mapping,
+        num_computed_tokens,
+        last_sampled_tokens,
+        output_bin_counts,
+        sampled_tokens,
+        num_sampled,
+        num_rejected,
+        query_start_loc,
+        all_token_ids,
+        total_len,
+    )
+
+
+def gumbel_sample_cpu(
+    logits: torch.Tensor,
+    expanded_idx_mapping: torch.Tensor,
+    temperature: torch.Tensor,
+    seeds: torch.Tensor,
+    positions: torch.Tensor,
+    apply_temperature: bool,
+    is_drafting: bool,
+    use_fp64: bool,
+) -> torch.Tensor:
+    return torch.ops._C.gumbel_sample_cpu(
+        logits,
+        expanded_idx_mapping,
+        temperature,
+        seeds,
+        positions,
+        apply_temperature,
+        is_drafting,
+        use_fp64,
+    )
+
+
+def apply_logit_bias_cpu(
+    logits: torch.Tensor,
+    expanded_idx_mapping: torch.Tensor,
+    positions: torch.Tensor,
+    num_allowed_token_ids: torch.Tensor,
+    allowed_token_ids: torch.Tensor,
+    num_logit_bias: torch.Tensor,
+    logit_bias_token_ids: torch.Tensor,
+    logit_bias: torch.Tensor,
+    min_lens: torch.Tensor,
+    num_stop_token_ids: torch.Tensor,
+    restore_when_all_masked: torch.Tensor,
+    stop_token_ids: torch.Tensor,
+    check_all_masked_rows: bool,
+) -> None:
+    torch.ops._C.apply_logit_bias_cpu(
+        logits,
+        expanded_idx_mapping,
+        positions,
+        num_allowed_token_ids,
+        allowed_token_ids,
+        num_logit_bias,
+        logit_bias_token_ids,
+        logit_bias,
+        min_lens,
+        num_stop_token_ids,
+        restore_when_all_masked,
+        stop_token_ids,
+        check_all_masked_rows,
+    )
+
+
 # mamba
 def selective_scan_fwd(
     u: torch.Tensor,
@@ -2426,6 +2642,28 @@ def topk_sigmoid(
     )
 
 
+def muse_spark_routing(
+    topk_weights: torch.Tensor,
+    topk_ids: torch.Tensor,
+    logits: torch.Tensor,
+    bias: torch.Tensor,
+    multiplier: float,
+    route_scale: float,
+    division_eps: float,
+    is_padding: torch.Tensor | None = None,
+) -> None:
+    torch.ops._moe_C.muse_spark_routing(
+        topk_weights,
+        topk_ids,
+        logits,
+        bias,
+        multiplier,
+        route_scale,
+        division_eps,
+        is_padding,
+    )
+
+
 def topk_hash_softplus_sqrt(
     topk_weights: torch.Tensor,
     topk_indices: torch.Tensor,
@@ -2465,6 +2703,7 @@ def grouped_topk(
     routed_scaling_factor: float,
     bias: torch.Tensor,
     scoring_func: int = 0,
+    is_padding: torch.Tensor | None = None,
 ):
     """Perform grouped top-k routing for mixture of experts.
 
@@ -2477,6 +2716,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
+        is_padding: Optional CUDA token mask; padded rows get zero weights and -1 IDs.
 
     """
     if not (current_platform.is_cuda() or current_platform.is_xpu()):
@@ -2488,6 +2728,20 @@ def grouped_topk(
             scores.new_empty((0, topk), dtype=torch.float32),
             scores.new_empty((0, topk), dtype=torch.int32),
         )
+    # XPU provides its own operator with the original eight-argument schema.
+    if current_platform.is_xpu():
+        if is_padding is not None:
+            raise NotImplementedError("Fused grouped padding is only supported on CUDA")
+        return torch.ops._moe_C.grouped_topk(
+            scores,
+            num_expert_group,
+            topk_group,
+            topk,
+            renormalize,
+            routed_scaling_factor,
+            bias,
+            scoring_func,
+        )
     return torch.ops._moe_C.grouped_topk(
         scores,
         num_expert_group,
@@ -2497,6 +2751,7 @@ def grouped_topk(
         routed_scaling_factor,
         bias,
         scoring_func,
+        is_padding,
     )
 
 
@@ -2512,6 +2767,7 @@ def _grouped_topk_fake(
         routed_scaling_factor: float,
         bias: torch.Tensor,
         scoring_func: int = 0,
+        is_padding: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         num_tokens = scores.size(0)
         return (
@@ -2782,6 +3038,7 @@ def fused_minimax_m3_qknorm_rope_kv_insert(
     q_fp8_scale: float = 1.0,
     kv_k_scale: torch.Tensor | None = None,
     kv_v_scale: torch.Tensor | None = None,
+    write_qkv: bool = True,
 ) -> None:
     """Fused MiniMax-M3 attention pre-processing (in-place).
 
@@ -2792,11 +3049,13 @@ def fused_minimax_m3_qknorm_rope_kv_insert(
     - sparse layer (``num_index_heads > 0``): ``[q | k | v | index_q |
       index_k]`` — the index branch is read straight out of ``qkv``.
 
-    When ``kv_cache`` is given (sparse serving), also scatter-inserts the
-    normed/roped k & v into the paged KV cache by ``slot_mapping`` and the
-    index key into ``index_cache`` by ``index_slot_mapping``. ``kv_cache_dtype``
-    selects the cache storage/conversion path. If
-    ``index_slot_mapping`` is omitted, ``slot_mapping`` is used for both caches.
+    When ``kv_cache`` is given, also scatter-inserts the normed/roped k & v
+    into the paged KV cache by ``slot_mapping`` and, on sparse layers, the
+    index key into ``index_cache`` by ``index_slot_mapping``.
+    ``kv_cache_dtype`` selects the cache storage/conversion path; a dense
+    layer's insert stores exactly what ``reshape_and_cache_flash`` would with
+    unit k/v scales. If ``index_slot_mapping`` is omitted, ``slot_mapping`` is
+    used for both caches.
 
     If ``q_out`` / ``index_q_out`` (contiguous ``[N, nq*128]`` / ``[N,
     niq*128]``) are given, the normed/roped q / index_q are written there
@@ -2818,6 +3077,11 @@ def fused_minimax_m3_qknorm_rope_kv_insert(
     ``[index_q | index_k]`` tail, but the kernel only processes the main q/k/v
     branches and main KV cache. This is used by MiniMax-M3 index-topk reuse
     layers that consume top-k block ids selected by an earlier sparse layer.
+
+    ``write_qkv=False`` needs ``q_fp8_out`` without ``q_out``, for callers that
+    read only the FP8 q, ``index_q_out`` and the caches: k and index_k need not
+    be rewritten in ``qkv``. The persistent kernel (SM80+, ``rotary_dim`` 64)
+    skips those stores; the fallback kernel may still make them.
     """
     torch.ops._C.fused_minimax_m3_qknorm_rope_kv_insert(
         qkv,
@@ -2845,6 +3109,47 @@ def fused_minimax_m3_qknorm_rope_kv_insert(
         q_fp8_scale,
         kv_k_scale,
         kv_v_scale,
+        write_qkv,
+    )
+
+
+def fused_minimax_m3_dspark_context_kv(
+    all_kv: torch.Tensor,
+    k_norm_weight: torch.Tensor,
+    cos_sin_cache: torch.Tensor,
+    positions: torch.Tensor,
+    slot_mapping: torch.Tensor,
+    kv_cache: torch.Tensor,
+    cache_offsets: torch.Tensor,
+    k_scale: torch.Tensor,
+    scale_offsets: torch.Tensor,
+    rotary_dim: int,
+    eps: float,
+    kv_cache_dtype: str,
+) -> None:
+    """Store the MiniMax-M3 DSpark draft's context K/V for every layer.
+
+    ``all_kv`` ``[N, L, 2, nkv, 128]`` holds each layer's K and V. K gets the
+    layer's Gemma RMSNorm (``k_norm_weight[L, 128]``) and partial NeoX RoPE;
+    both are stored like ``reshape_and_cache_flash`` into the layers' caches.
+    Layer ``l``'s cache is ``kv_cache`` (layer 0, ``[nb, nkv, bs, 2 * 128]``)
+    shifted by ``cache_offsets[l]`` elements; its fp8 k/v scales are
+    ``k_scale`` (layer 0's) shifted by ``scale_offsets[0, l]`` /
+    ``scale_offsets[1, l]`` floats. Slots < 0 are skipped.
+    """
+    torch.ops._C.fused_minimax_m3_dspark_context_kv(
+        all_kv,
+        k_norm_weight,
+        cos_sin_cache,
+        positions,
+        slot_mapping,
+        kv_cache,
+        cache_offsets,
+        k_scale,
+        scale_offsets,
+        rotary_dim,
+        eps,
+        kv_cache_dtype,
     )
 
 
@@ -3050,8 +3355,8 @@ def swap_blocks(
     and their matching destination block number on the destination tensor.
     block_mapping is expected to be a tensor of shape (num_blocks_to_copy, 2)
     where each block_mapping[i] represents a single copy operation, copying
-    block #block_mapping[i][0] from the source tensor
-    to block #block_mapping[i][1] on the destination tensor.
+    block `block_mapping[i][0]` from the source tensor
+    to block `block_mapping[i][1]` on the destination tensor.
     block_mapping should have dtype int64.
 
     The source and the destination tensors can be either on cpu or gpu,
diff --git a/vllm/benchmarks/datasets/datasets.py b/vllm/benchmarks/datasets/datasets.py
index 3746ac727e0f..322e409aac4f 100644
--- a/vllm/benchmarks/datasets/datasets.py
+++ b/vllm/benchmarks/datasets/datasets.py
@@ -2451,7 +2451,10 @@ def load_data(self) -> None:
 
         # Load the JSONL file
         if self.dataset_path.endswith(".jsonl"):
-            jsonl_data = pd.read_json(path_or_buf=self.dataset_path, lines=True)
+            # Keep prompts like "00123" as strings instead of inferring numbers.
+            jsonl_data = pd.read_json(
+                path_or_buf=self.dataset_path, lines=True, dtype={"prompt": str}
+            )
 
             # check if the JSONL file has a 'prompt' column
             if "prompt" not in jsonl_data.columns:
diff --git a/vllm/benchmarks/forward/main.py b/vllm/benchmarks/forward/main.py
index 57d37a7b2c68..22622e6697bf 100644
--- a/vllm/benchmarks/forward/main.py
+++ b/vllm/benchmarks/forward/main.py
@@ -66,10 +66,13 @@ def add_cli_args(parser: FlexibleArgumentParser) -> None:
     )
     parser.add_argument(
         "--routing",
-        choices=["auto", "real", "balanced"],
+        choices=["auto", "real", "balanced", "zipf"],
         default="auto",
         help="MoE routing. balanced: fixed routing with equal load per EP rank. "
-        "auto: balanced with dummy weights, else real.",
+        "zipf: fixed weighted sampling without replacement (alpha=1.25), "
+        "with expert popularity shuffled per layer. "
+        "auto: zipf with dummy weights, else real; falls back to real when "
+        "the synthetic fixture is unsupported.",
     )
     parser.add_argument(
         "--exclude-drafter",
@@ -209,17 +212,17 @@ def run_case(runner, case: BenchCase, args: argparse.Namespace) -> dict[str, Any
         profile_prefix,
         include_drafter=include_drafter,
     )
-    routing, balanced = install_routing(
+    routing, fixture = install_routing(
         runner, args.routing, args.load_format == "dummy", batch.num_tokens, args.seed
     )
     try:
         result = bench.run()
-        if balanced is not None:
-            balanced.check_applied()
+        if fixture is not None:
+            fixture.check_applied()
     finally:
         runner.execute_model(batch.finish())
-        if balanced is not None:
-            balanced.remove()
+        if fixture is not None:
+            fixture.remove()
     return {
         "num_tokens": batch.num_tokens,
         "kv_blocks": batch.num_blocks,
diff --git a/vllm/benchmarks/forward/routing.py b/vllm/benchmarks/forward/routing.py
index 9d48a541e77d..e68c1344a1de 100644
--- a/vllm/benchmarks/forward/routing.py
+++ b/vllm/benchmarks/forward/routing.py
@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Balanced MoE routing for benchmarks with dummy weights.
+"""Synthetic MoE routing for forward benchmarks.
 
 Dummy gate weights give degenerate routing, which skews MoE and all-to-all
 cost. The routing still runs; its result is replaced by the fixture at
@@ -18,10 +18,13 @@
 """
 
 import inspect
+import zlib
 
 import numpy as np
 import torch
 
+from vllm.distributed.parallel_state import get_tp_group
+from vllm.forward_context import get_forward_context
 from vllm.logger import init_logger
 from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear
 
@@ -54,11 +57,41 @@ def balanced_topk_ids(
     return ranks * local + (offset + repeat) % local
 
 
+def zipf_topk_ids(
+    num_tokens: int, num_experts: int, top_k: int, rng: np.random.Generator
+) -> np.ndarray:
+    """Sample distinct experts per token with Zipf(alpha=1.25) popularity."""
+    expert_ids = rng.permutation(num_experts)
+    inverse_weights = np.arange(1, num_experts + 1, dtype=np.float64) ** 1.25
+    ids = np.empty((num_tokens, top_k), dtype=np.int64)
+    for start in range(0, num_tokens, 1024):
+        count = min(1024, num_tokens - start)
+        # Exponential races sample weighted top-k without replacement.
+        scores = rng.exponential(size=(count, num_experts)) * inverse_weights
+        selected = np.argpartition(scores, top_k - 1, axis=1)[:, :top_k]
+        ids[start : start + count] = expert_ids[selected]
+    return ids
+
+
+def _fixture_rows(fixture: torch.Tensor, num_tokens: int) -> torch.Tensor:
+    """Select this rank's rows when the model shards tokens across TP ranks."""
+    descriptor = get_forward_context().batch_descriptor
+    offset = 0
+    if descriptor is not None:
+        tp = get_tp_group()
+        full = descriptor.num_tokens
+        if num_tokens < full and num_tokens == -(-full // tp.world_size):
+            offset = tp.rank_in_group * num_tokens
+    return fixture[offset : offset + num_tokens]
+
+
 def _moe_layers(model: torch.nn.Module) -> list[tuple[torch.nn.Module, GateLinear]]:
     """(MoE module, gate) per layer, in model order."""
     layers: dict[GateLinear, torch.nn.Module] = {}
     for module in model.modules():
         gate = getattr(module, "gate", None)
+        if not isinstance(gate, GateLinear):
+            gate = getattr(module, "router", None)
         if not isinstance(gate, GateLinear) or gate in layers:
             continue
         if _top_k(module) is not None:
@@ -76,8 +109,8 @@ def _top_k(module: torch.nn.Module) -> int | None:
     return None
 
 
-class BalancedRouting:
-    """Installs balanced routing for one case; `remove()` restores the model."""
+class RoutingFixture:
+    """Installs synthetic routing for one case; `remove()` restores the model."""
 
     def __init__(
         self,
@@ -88,6 +121,7 @@ def __init__(
         ep_size: int,
         seed: int,
         device: torch.device,
+        mode: str,
     ):
         layers = _moe_layers(model)
         if not layers:
@@ -101,14 +135,22 @@ def __init__(
         rows = -(-num_tokens // 256) * 256 + 256
         for layer, (module, gate) in enumerate(layers):
             if (getattr(module, "num_expert_group", None) or 1) > 1:
-                raise NotImplementedError("Balanced routing with grouped top-k.")
+                raise NotImplementedError("Synthetic routing with grouped top-k.")
             num_experts = gate.weight.shape[0]
             top_k = _top_k(module)
             assert top_k is not None
-            rng = np.random.default_rng([seed, layer])
-            ids = balanced_topk_ids(rows * dp_size, num_experts, top_k, ep_size, rng)[
-                dp_rank * rows : (dp_rank + 1) * rows
-            ]
+            layer_seed = layer
+            if mode == "zipf" and gate.prefix:
+                # Preserve layer popularity when pipeline stages change.
+                layer_seed = zlib.crc32(gate.prefix.encode())
+            rng = np.random.default_rng([seed, layer_seed])
+            if mode == "zipf":
+                ids = zipf_topk_ids(rows * dp_size, num_experts, top_k, rng)
+            else:
+                ids = balanced_topk_ids(
+                    rows * dp_size, num_experts, top_k, ep_size, rng
+                )
+            ids = ids[dp_rank * rows : (dp_rank + 1) * rows]
             logits = torch.full((rows, num_experts), _OTHER_LOGIT, dtype=torch.float32)
             logits.scatter_(1, torch.from_numpy(ids), _CHOSEN_LOGIT)
             # Distinct values: no ties for any top-k implementation.
@@ -121,6 +163,8 @@ def __init__(
             bias = getattr(gate, "e_score_correction_bias", None)
             if bias is None:
                 bias = getattr(module, "e_score_correction_bias", None)
+            if bias is None:
+                bias = getattr(module, "expert_bias", None)
             if isinstance(bias, torch.Tensor):
                 self._biases.append((bias, bias.detach().clone()))
                 bias.data.zero_()
@@ -149,7 +193,7 @@ def hook(module, inputs, output):
             logits = output[0] if isinstance(output, tuple) else output
             if cast[0].dtype != logits.dtype:
                 cast[0] = fixture.to(logits.dtype)  # Once, at warmup.
-            replaced = cast[0][: logits.shape[0]]
+            replaced = _fixture_rows(cast[0], logits.shape[0])
             return (replaced, *output[1:]) if isinstance(output, tuple) else replaced
 
         return hook
@@ -163,7 +207,7 @@ def hook(module, args, kwargs):
             ids = kwargs["topk_ids"] if "topk_ids" in kwargs else args[index]
             if cast[0].dtype != ids.dtype:
                 cast[0] = fixture.to(ids.dtype)  # Once, at warmup.
-            replaced = cast[0][: ids.shape[0]]
+            replaced = _fixture_rows(cast[0], ids.shape[0])
             if "topk_ids" in kwargs:
                 return args, {**kwargs, "topk_ids": replaced}
             return (*args[:index], replaced, *args[index + 1 :]), kwargs
@@ -179,7 +223,7 @@ def check_applied(self) -> None:
         ]
         if missed:
             raise RuntimeError(
-                f"Balanced routing didn't reach {len(missed)} MoE layers "
+                f"Synthetic routing didn't reach {len(missed)} MoE layers "
                 f"({missed[0]}): their gate is fused away. Use --routing real."
             )
 
@@ -192,15 +236,15 @@ def remove(self) -> None:
 
 def install_routing(
     runner, requested: str, dummy_weights: bool, num_tokens: int, seed: int
-) -> tuple[str, BalancedRouting | None]:
-    """Resolve `--routing` and install balanced routing if selected.
+) -> tuple[str, RoutingFixture | None]:
+    """Resolve `--routing` and install a synthetic fixture if selected.
 
-    `auto` is balanced with dummy weights and real otherwise; it falls back to
+    `auto` is Zipf with dummy weights and real otherwise; it falls back to
     real routing when the model doesn't support the fixture.
     """
     mode = requested
     if mode == "auto":
-        mode = "balanced" if dummy_weights else "real"
+        mode = "zipf" if dummy_weights else "real"
     if mode == "real":
         return mode, None
     parallel = runner.parallel_config
@@ -215,7 +259,9 @@ def install_routing(
     if not parallel.enable_expert_parallel:
         num_tokens *= parallel.prefill_context_parallel_size
     try:
-        routing = BalancedRouting(
+        if mode == "zipf" and parallel.prefill_context_parallel_size > 1:
+            raise NotImplementedError("Zipf routing with prefill context parallelism.")
+        routing = RoutingFixture(
             runner.model,
             num_tokens,
             runner.dp_rank,
@@ -223,9 +269,10 @@ def install_routing(
             ep_size,
             seed,
             runner.device,
+            mode,
         )
     except (ValueError, NotImplementedError) as e:
-        if requested == "balanced":
+        if requested != "auto":
             raise
         logger.warning("Using real routing: %s", e)
         return "real", None
diff --git a/vllm/benchmarks/forward/speculator.py b/vllm/benchmarks/forward/speculator.py
index bbecdce8579d..0e35cc58d2fd 100644
--- a/vllm/benchmarks/forward/speculator.py
+++ b/vllm/benchmarks/forward/speculator.py
@@ -45,6 +45,7 @@ def prepare(self, input_batch, context):
                     dtype=torch.int32,
                 ),
                 uniform_token_count=None,
+                max_query_len=None,
                 eager=True,
                 num_reqs=self.batch.num_reqs,
             )
diff --git a/vllm/benchmarks/lib/endpoint_request_func.py b/vllm/benchmarks/lib/endpoint_request_func.py
index 442b8fa53ae2..041079282c7d 100644
--- a/vllm/benchmarks/lib/endpoint_request_func.py
+++ b/vllm/benchmarks/lib/endpoint_request_func.py
@@ -544,11 +544,17 @@ def _extract_sse_data(message: str) -> str | None:
 
 # Responses events that carry generated tokens. Reasoning deltas are included
 # because the server is already decoding when it emits them, so excluding them
-# would report the whole reasoning phase as time to first token.
+# would report the whole reasoning phase as time to first token. Tool call
+# arguments and code are generated tokens too, as `tool_calls` deltas are on
+# `openai-chat`; a response that only calls a tool emits nothing else.
 _RESPONSES_TOKEN_DELTA_EVENTS = frozenset(
     {
         "response.output_text.delta",
         "response.reasoning_text.delta",
+        "response.function_call_arguments.delta",
+        "response.custom_tool_call_input.delta",
+        "response.mcp_call_arguments.delta",
+        "response.code_interpreter_call_code.delta",
     }
 )
 
diff --git a/vllm/benchmarks/throughput.py b/vllm/benchmarks/throughput.py
index 76f6a72fd98a..b48e5da6895c 100644
--- a/vllm/benchmarks/throughput.py
+++ b/vllm/benchmarks/throughput.py
@@ -17,6 +17,7 @@
 
 from vllm.benchmarks.datasets import (
     BenchmarkDataset,
+    RandomDataset,
     SampleRequest,
     add_random_dataset_base_args,
     add_random_multimodal_dataset_args,
@@ -523,6 +524,14 @@ def save_to_pytorch_benchmark_format(
         write_to_json(pt_file, pt_records)
 
 
+# --random-* length -> (legacy option it overrides, value when neither is set).
+_RANDOM_LEN_FALLBACKS = {
+    "random_input_len": ("input_len", RandomDataset.DEFAULT_INPUT_LEN),
+    "random_output_len": ("output_len", RandomDataset.DEFAULT_OUTPUT_LEN),
+    "random_prefix_len": ("prefix_len", RandomDataset.DEFAULT_PREFIX_LEN),
+}
+
+
 def _to_serve_args(args: argparse.Namespace) -> argparse.Namespace:
     """Translate throughput args into a namespace the shared get_samples reads.
 
@@ -539,16 +548,16 @@ def _to_serve_args(args: argparse.Namespace) -> argparse.Namespace:
 
     """
     d = vars(args).copy()
-    # random_*: prefer --random-* over legacy --input/output/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)
+    # random_*: prefer --random-* over legacy --input/output/prefix-len, then
+    # the RandomDataset defaults. Callers without the legacy flags (e.g. bench
+    # mm-processor) omit them, so read them optionally. Identity checks, not
+    # truthiness: an explicit 0 must not fall through to them.
+    for name, (legacy, default) in _RANDOM_LEN_FALLBACKS.items():
+        value = getattr(args, name, None)
+        if value is None:
+            value = getattr(args, legacy, None)
+        d[name] = default if value is None else value
     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.
@@ -642,12 +651,6 @@ def validate_args(args):
     if not args.dataset_path and args.dataset_name != "prefix_repetition":
         print("When dataset path is not set, it will default to random dataset")
         args.dataset_name = "random"
-        random_input_len = getattr(args, "random_input_len", None)
-        if args.input_len is None and random_input_len is None:
-            raise ValueError(
-                "Either --input-len or --random-input-len must be provided "
-                "for a random dataset"
-            )
 
     # === Dataset Name Specific Checks ===
     # --hf-subset and --hf-split: only used
@@ -820,14 +823,18 @@ def add_cli_args(parser: FlexibleArgumentParser):
         "--input-len",
         type=int,
         default=None,
-        help="Input prompt length for each request",
+        help="Input prompt length for each request. For random datasets, "
+        "--random-input-len takes precedence; if neither is set, "
+        f"{RandomDataset.DEFAULT_INPUT_LEN}.",
     )
     parser.add_argument(
         "--output-len",
         type=int,
         default=None,
         help="Output length for each request. Overrides the "
-        "output length from the dataset.",
+        "output length from the dataset. For random datasets, "
+        "--random-output-len takes precedence; if neither is set, "
+        f"{RandomDataset.DEFAULT_OUTPUT_LEN}.",
     )
     parser.add_argument(
         "--n", type=int, default=1, help="Number of generated sequences per prompt."
@@ -907,9 +914,10 @@ def add_cli_args(parser: FlexibleArgumentParser):
     parser.add_argument(
         "--prefix-len",
         type=int,
-        default=0,
+        default=None,
         help="Number of fixed prefix tokens before the random "
-        "context in a request (default: 0).",
+        "context in a request. --random-prefix-len takes precedence; if "
+        f"neither is set, {RandomDataset.DEFAULT_PREFIX_LEN}.",
     )
 
     # hf dataset
@@ -975,6 +983,16 @@ def add_cli_args(parser: FlexibleArgumentParser):
 
     # (random, random-mm, random-rerank)
     add_random_dataset_base_args(parser)
+    # Unset unless passed, so the legacy --input-len/--output-len/--prefix-len
+    # apply when given; _to_serve_args falls back to the RandomDataset defaults.
+    for action in parser._actions:
+        if action.dest in _RANDOM_LEN_FALLBACKS:
+            legacy, default = _RANDOM_LEN_FALLBACKS[action.dest]
+            action.default = None
+            action.help = (
+                f"{action.help} Takes precedence over "
+                f"--{legacy.replace('_', '-')}; if neither is set, {default}."
+            )
     add_random_multimodal_dataset_args(parser)
 
     # ASR dataset
diff --git a/vllm/compilation/compiler_interface.py b/vllm/compilation/compiler_interface.py
index 7c2e05cc6b3f..1811d6a7d97b 100644
--- a/vllm/compilation/compiler_interface.py
+++ b/vllm/compilation/compiler_interface.py
@@ -394,7 +394,12 @@ def compile(
         else:
             fake_mode_ctx = contextlib.nullcontext()
 
-        with pregrad_ctx, fake_mode_ctx, set_functorch_config(self.use_aot_compile):
+        with (
+            pregrad_ctx,
+            fake_mode_ctx,
+            set_functorch_config(self.use_aot_compile),
+            set_compile_on_one_rank(),
+        ):
             compiled_graph = standalone_compile(graph, example_inputs, **compile_kwargs)
 
         if not self.use_aot_compile:
@@ -576,6 +581,7 @@ def _get_shape_env() -> AlwaysHitShapeEnv:
 
         with ExitStack() as stack:
             stack.enter_context(set_functorch_config(self.use_aot_compile))
+            stack.enter_context(set_compile_on_one_rank())
             # for hijacking the hash of the compiled graph
             stack.enter_context(
                 patch(
@@ -789,6 +795,20 @@ def set_functorch_config(use_aot_compile: bool) -> Generator[None, None, None]:
             yield
 
 
+@contextlib.contextmanager
+def set_compile_on_one_rank() -> Generator[None, None, None]:
+    # Resolve the device at run time so the cached code runs on any device
+    # index (--device-ids). Follow VLLM_COMPILE_ON_ONE_RANK, which is part of
+    # the cache key, rather than torch's own setting.
+    from torch.compiler import config as compiler_config
+
+    if not is_torch_equal_or_newer("2.14.0.dev"):
+        yield
+        return
+    with compiler_config.patch(compile_on_one_rank=envs.VLLM_COMPILE_ON_ONE_RANK):
+        yield
+
+
 def trigger_inductor_lazy_init(device: torch.device | None = None) -> None:
     """Eagerly trigger inductor's once-per-process lazy inits (SFDP pattern
     matcher, pad_mm, misc patterns).
diff --git a/vllm/compilation/depyf_compat.py b/vllm/compilation/depyf_compat.py
new file mode 100644
index 000000000000..4017d99b2b94
--- /dev/null
+++ b/vllm/compilation/depyf_compat.py
@@ -0,0 +1,90 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Keep depyf usable against Inductor versions newer than its last release.
+
+``depyf.prepare_debug`` replaces three torch functions with hooks of its own.
+Two of them tolerate new arguments, but ``load_by_key_path`` has its signature
+frozen at ``(key, path, linemap, attrs)``, and Inductor now also passes
+``set_sys_modules`` (pytorch#184285). Every compilation that reaches an
+autotune code block while depyf is active therefore dies with a ``TypeError``,
+which surfaces as ``InductorError: Failed to run autotuning code block``.
+
+depyf has hit this before -- ``linemap`` and ``attrs`` broke the same hook in
+thuml/depyf#80 -- but 0.20.0 is both the pinned version and the last release
+upstream published, so a version bump cannot resolve it. Patch the hook here
+instead, following what depyf itself does for ``lazy_format_graph_code``:
+accept ``**kwargs`` and forward them.
+
+Route every depyf entry point through ``prepare_debug`` below so the patch is
+never missed.
+"""
+
+import inspect
+import os
+from collections.abc import Callable
+from contextlib import AbstractContextManager
+from typing import Any
+
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+_patched = False
+
+
+def _hook_accepts_inductor_keywords(hook: Callable[..., Any]) -> bool:
+    """Whether depyf's hook already takes everything Inductor passes it.
+
+    Compares against the live ``PyCodeCache.load_by_key_path``, so this only
+    means anything before depyf swaps that out for the hook.
+    """
+    import torch._inductor.codecache as codecache
+
+    accepted = inspect.signature(hook).parameters
+    if any(p.kind is inspect.Parameter.VAR_KEYWORD for p in accepted.values()):
+        return True
+    expected = inspect.signature(codecache.PyCodeCache.load_by_key_path).parameters
+    return expected.keys() <= accepted.keys()
+
+
+def _patch_load_by_key_path() -> None:
+    """Forward unrecognised keywords through depyf's ``load_by_key_path`` hook.
+
+    Mirrors depyf's own implementation -- dump the generated source alongside
+    the rest of the debug output, then defer to the function depyf displaced --
+    and adds ``**kwargs`` so the call survives arguments depyf never knew about.
+    """
+    import depyf.explain.enable_debugging as enable_debugging
+    from depyf.explain.global_variables import data
+    from depyf.explain.utils import (
+        get_current_compiled_fn_name,
+        write_code_to_file_template,
+    )
+
+    def load_by_key_path(key, path, linemap=None, attrs=None, **kwargs):
+        with open(path) as f:
+            src = f.read()
+        func_name = get_current_compiled_fn_name()
+        dumped_path = write_code_to_file_template(
+            src, os.path.join(data["dump_src_dir"], func_name + ".kernel_%s.py")
+        )
+        unpatched = data["unpatched_load_by_key_path"]
+        return unpatched(key, dumped_path, linemap, attrs, **kwargs)
+
+    enable_debugging.patched_load_by_key_path = load_by_key_path
+
+
+def prepare_debug(dump_src_dir: str) -> AbstractContextManager[None]:
+    """``depyf.prepare_debug`` with the compatibility patches applied."""
+    global _patched
+    import depyf
+    import depyf.explain.enable_debugging as enable_debugging
+
+    if not _patched and not _hook_accepts_inductor_keywords(
+        enable_debugging.patched_load_by_key_path
+    ):
+        logger.debug("Patching depyf's load_by_key_path hook for this Inductor")
+        _patch_load_by_key_path()
+        _patched = True
+
+    return depyf.prepare_debug(dump_src_dir)
diff --git a/vllm/compilation/monitor.py b/vllm/compilation/monitor.py
index e38192034b9d..297b1fb8c09b 100644
--- a/vllm/compilation/monitor.py
+++ b/vllm/compilation/monitor.py
@@ -41,11 +41,11 @@ def monitor_torch_compile(
     depyf_cm = None
     path = vllm_config.compile_debug_dump_path()
     if compilation_config.mode == CompilationMode.VLLM_COMPILE and path:
-        import depyf
+        from vllm.compilation import depyf_compat
 
         path.mkdir(parents=True, exist_ok=True)
         logger.debug("Dumping depyf output to %s", path)
-        depyf_cm = depyf.prepare_debug(path.as_posix())
+        depyf_cm = depyf_compat.prepare_debug(path.as_posix())
         depyf_cm.__enter__()
 
     try:
diff --git a/vllm/config/attention.py b/vllm/config/attention.py
index 2ba6fdb6cb1b..a63cc3962631 100644
--- a/vllm/config/attention.py
+++ b/vllm/config/attention.py
@@ -40,8 +40,10 @@ class AttentionConfig:
     backend: AttentionBackendEnum | None = None
     """Attention backend to use. Use "auto" or None for automatic selection."""
 
-    minimax_m3_msa_decode_backend: MiniMaxM3MSADecodeBackend = "triton"
-    """Sparse decode kernel used by the MiniMax M3 MSA backend."""
+    minimax_m3_msa_decode_backend: MiniMaxM3MSADecodeBackend = "cutlass"
+    """Sparse decode kernel used by the MiniMax M3 MSA backend. "cutlass" falls
+    back to Triton where the CUTLASS kernel does not support the platform, KV
+    cache dtype, or head layout."""
 
     backend_per_kind: dict[str, AttentionBackendEnum] = field(default_factory=dict)
     """Per-KV-cache-group attention backend overrides, keyed by
@@ -95,7 +97,20 @@ class AttentionConfig:
     model's default (bf16 for MiniMax M3, fp8 for the DeepSeek sparse
     indexer, mxfp4 for DeepSeek-V4 on ROCm CDNA4 with aiter). Quantized
     formats (fp8, mxfp4, nvfp4) require indexer kernel support in the
-    backend."""
+    backend; MiniMax M3 supports nvfp4 on SM100/SM103/SM107 with at most 8
+    decode queries per request."""
+
+    indexer_decode_cp: bool = False
+    """Split the sparse-attention indexer's decode scoring across the
+    tensor-parallel ranks by context instead of by index head: each rank scores
+    a 1/TP slice of the index-K pages for every head and the ranks exchange
+    their local top-k, so each reads 1/TP of the index-K cache per sparse
+    layer. The index-K cache stays replicated, so this is unrelated to
+    `--decode-context-parallel-size`, which shards the KV cache itself; the two
+    cannot be combined. Implemented for the MiniMax M3 NVFP4 indexer, which
+    needs `indexer_kv_dtype="nvfp4"` and TP 4 (one index head per rank); it is
+    ignored with a warning where it does not apply. It pays off from about 4
+    concurrent decodes at 128K context and grows with batch and context."""
 
     indexer_sparse_logits: bool | None = None
     """DeepSeek V4.1 two-level indexer: score only the candidate blocks with
diff --git a/vllm/config/cache.py b/vllm/config/cache.py
index 626825dba8c7..0a7d268606c0 100644
--- a/vllm/config/cache.py
+++ b/vllm/config/cache.py
@@ -183,6 +183,7 @@ class CacheConfig:
     """Configuration for the KV cache."""
 
     DEFAULT_BLOCK_SIZE: ClassVar[int] = 16
+    DEFAULT_GPU_MEMORY_UTILIZATION: ClassVar[float] = 0.92
 
     block_size: int = Field(default=None, gt=0)  # type: ignore[assignment]
     """Size of a contiguous cache block in number of tokens.
@@ -212,24 +213,26 @@ class CacheConfig:
 
     This equals to the `hash_block_size` used throughout the KV cache code.
     """
-    gpu_memory_utilization: float = Field(default=0.92, gt=0, le=1)
+    gpu_memory_utilization: float | None = Field(default=None, gt=0, le=1)
     """The fraction of GPU memory to be used for the model executor, which can
     range from 0 to 1. For example, a value of 0.5 would imply 50% GPU memory
-    utilization. If unspecified, will use the default value of 0.92. This is a
-    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. On non-GPU
-    installs, this value controls the corresponding device memory utilization.
-    """
+    utilization. This is a 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.
+    On non-GPU installs, this value controls the corresponding device memory
+    utilization.
+
+    If unspecified, 1.0 with the extensible KV cache (sized from measured
+    memory) and 0.92 otherwise; see `resolved_gpu_memory_utilization`."""
 
     @property
-    def device_memory_utilization(self) -> float:
+    def device_memory_utilization(self) -> float | None:
         """Device-neutral alias for ``gpu_memory_utilization``."""
         return self.gpu_memory_utilization
 
     @device_memory_utilization.setter
-    def device_memory_utilization(self, value: float) -> None:
+    def device_memory_utilization(self, value: float | None) -> None:
         self.gpu_memory_utilization = value
 
     cache_dtype: CacheDType | str = "auto"
@@ -380,6 +383,15 @@ def device_memory_utilization(self, value: float) -> None:
     gpu_memory_utilization. Note that kv_cache_memory_bytes
     (when not-None) ignores gpu_memory_utilization"""
 
+    enable_extensible_kv_cache: bool = Field(default=None)  # type: ignore[assignment]
+    """Reserve the KV cache address range with device virtual memory and commit
+    physical pages after warmup, so the cache is sized from the memory actually
+    free after CUDA graph capture rather than from a profiling estimate. If
+    unspecified, enabled on CUDA and ROCm with the V2 model runner unless
+    kv_cache_memory_bytes is set; falls back to estimate-based sizing where the
+    driver lacks VMM support. Explicitly requesting it where unsupported is an
+    error."""
+
     kv_offloading_size: float | None = None
     """Size of the KV cache offloading buffer in GiB. When TP > 1, this is
     the total buffer size summed across all TP ranks. By default, this is set
@@ -418,6 +430,7 @@ def compute_hash(self) -> str:
             "skip_page_size_padded",
             "user_specified_block_size",
             "_block_size_resolved",
+            "_extensible_kv_cache_resolved",
             # Post-init/derived counters
             "num_gpu_blocks",
             "num_cpu_blocks",
@@ -430,6 +443,8 @@ def compute_hash(self) -> str:
             # Feature toggles not impacting compiled graph shape
             "kv_sharing_fast_prefill",
             "swa_bounded_replay",
+            # Runtime memory allocation strategy, not graph structure.
+            "enable_extensible_kv_cache",
         }
 
         from vllm.config.utils import get_hash_factors, hash_factors
@@ -445,7 +460,10 @@ def metrics_info(self):
     _block_size_resolved: bool = field(default=False, init=False)
     """Guard against pydantic re-running _apply_block_size_default."""
 
-    @field_validator("block_size", mode="wrap")
+    _extensible_kv_cache_resolved: bool = field(default=False, init=False)
+    """Set once `enable_extensible_kv_cache` has been settled for the model."""
+
+    @field_validator("block_size", "enable_extensible_kv_cache", mode="wrap")
     @classmethod
     def _skip_none_validation(cls, value: Any, handler: Callable) -> Any:
         if value is None:
@@ -465,6 +483,15 @@ def _apply_block_size_default(self) -> "CacheConfig":
             self.user_specified_block_size = True
         return self
 
+    @property
+    def resolved_gpu_memory_utilization(self) -> float:
+        """`gpu_memory_utilization`, or its default given the sizing mode."""
+        if self.gpu_memory_utilization is not None:
+            return self.gpu_memory_utilization
+        if self.enable_extensible_kv_cache:
+            return 1.0
+        return self.DEFAULT_GPU_MEMORY_UTILIZATION
+
     @field_validator("cache_dtype", mode="after")
     @classmethod
     def _validate_cache_dtype(cls, cache_dtype: CacheDType | str) -> CacheDType | str:
diff --git a/vllm/config/compilation.py b/vllm/config/compilation.py
index 395b6e2c8ff2..bf40417d7c08 100644
--- a/vllm/config/compilation.py
+++ b/vllm/config/compilation.py
@@ -785,6 +785,7 @@ class CompilationConfig:
         "vllm::rocm_aiter_sparse_attn_indexer",
         "vllm::deepseek_v4_attention",
         "vllm::hpc_rope_norm_forward",
+        "vllm::step5_csa_attention",
     ]
 
     def compute_hash(self) -> str:
diff --git a/vllm/config/engram.py b/vllm/config/engram.py
index 156f6fca0514..d017dc841185 100644
--- a/vllm/config/engram.py
+++ b/vllm/config/engram.py
@@ -37,6 +37,11 @@ class EngramConfig:
     cpu_offload: bool = True
     """Store embedding weights in pinned CPU memory for UVA lookup."""
 
+    disk_offload: bool = False
+    """Keep Qwen4Exp PLE tables in the checkpoint files and read each step's
+    rows on the host instead of allocating the table. Takes precedence over
+    cpu_offload."""
+
     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."""
@@ -44,10 +49,11 @@ class EngramConfig:
     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."""
+    memory without per-step Engram DP collectives. Uses host-NUMA cuMem shared
+    by FABRIC handle where supported, else /dev/shm (which then needs 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 the host cannot hold them."""
 
     use_thp: bool = False
     """Back private CPU-offloaded tables with transparent huge pages (best
@@ -81,6 +87,11 @@ def verify_model_config(self, model_config: "ModelConfig | None") -> None:
                 "EngramConfig requires a model with supported Engram "
                 "embeddings and non-empty n-gram layer ids."
             )
+        if self.disk_offload:
+            from vllm.platforms import current_platform
+
+            if field != "ple_layer_ids" or not current_platform.is_cuda():
+                raise ValueError("disk_offload requires a Qwen4Exp model on CUDA")
 
     def resolve_dp_shared_memory(self, parallel_config: "ParallelConfig") -> None:
         """Share host tables by default wherever the configuration permits."""
@@ -107,6 +118,8 @@ def verify_parallel_config(self, parallel_config: "ParallelConfig") -> None:
             raise ValueError(
                 "Engram embedding_across_dp is not supported with elastic EP yet."
             )
+        if self.disk_offload and parallel_config.enable_elastic_ep:
+            raise ValueError("disk_offload is not supported with elastic EP.")
 
     def get_parallel_size(self, parallel_config: "ParallelConfig") -> int:
         """Derive the embedding group size from the parallel configuration."""
diff --git a/vllm/config/fault_tolerance.py b/vllm/config/fault_tolerance.py
index 7ed095d204f6..5465ee8382f2 100644
--- a/vllm/config/fault_tolerance.py
+++ b/vllm/config/fault_tolerance.py
@@ -1,18 +1,34 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-
-from vllm.config.utils import config
-
-
-@config
-class FaultToleranceConfig:
-    """Configuration for fault tolerance."""
-
-    engine_recovery_timeout_sec: int = 120
-    """Timeout (in seconds) to wait for error handling instructions
-    before raising an exception. If the EngineCore encounters an
-    error, it waits up to this many seconds for vLLM to receive 
-    instructions on how to handle the error and then recover from the fault.
-    If vLLM does not recover during this time, the original error is raised.
-    """
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from typing import TYPE_CHECKING
+
+from vllm.config.utils import config
+
+if TYPE_CHECKING:
+    from vllm.config.parallel import ParallelConfig
+
+
+@config
+class FaultToleranceConfig:
+    """Configuration for fault tolerance."""
+
+    engine_recovery_timeout_sec: int = 120
+    """Timeout (in seconds) to wait for error handling instructions
+    before raising an exception. If the EngineCore encounters an
+    error, it waits up to this many seconds for vLLM to receive
+    instructions on how to handle the error and then recover from the fault.
+    If vLLM does not recover during this time, the original error is raised.
+    """
+
+
+def ft_tp_barrier_required(parallel_config: "ParallelConfig") -> bool:
+    """Whether the per-step FT barrier over the TP cpu group is armed.
+
+    The barrier makes TP peers fail on the host within the gloo timeout
+    instead of leaving an orphaned TP collective on the device stream.
+    """
+    return (
+        parallel_config.enable_fault_tolerance
+        and parallel_config.tensor_parallel_size > 1
+    )
diff --git a/vllm/config/kernel.py b/vllm/config/kernel.py
index 9e40c1528f59..f0bfbf3621e2 100644
--- a/vllm/config/kernel.py
+++ b/vllm/config/kernel.py
@@ -303,7 +303,10 @@ class KernelConfig:
     """If True, run the eligible layers of decode steps on the model's
     MonoKernel: a persistent kernel that runs most or all of one decoder layer
     in a single launch (DeepSeek-V4.1 on MI355X, DeepSeek-V4-Pro at TP8 on
-    MI355X). A deployment that cannot run it is refused at startup."""
+    MI355X; Muse Spark MXFP4 on MI355X with tensor parallel size 8, expert
+    parallelism and a 7-token DSpark draft, for speculative verify steps of 1-16
+    requests, with bf16 MoE activations instead of the checkpoint's MXFP4). A
+    deployment that cannot run it is refused at startup."""
 
     moe_backend: MoEBackend = "auto"
     """Backend for MoE expert computation kernels. Available options:
@@ -361,6 +364,10 @@ class KernelConfig:
       requires VLLM_ROCM_USE_AITER)
 
     Explicit values raise RuntimeError when their constraints are not met.
+
+    On SM100 decode without DCP, persistent (including auto selecting it)
+    reads BF16 logits directly and widens scores inside the kernel, avoiding
+    a full FP32 logits copy. Cooperative still requires FP32 logits.
     """
 
     linear_backend: LinearBackend = "auto"
diff --git a/vllm/config/kv_transfer.py b/vllm/config/kv_transfer.py
index 59eb4c80be2d..30482d3a22fe 100644
--- a/vllm/config/kv_transfer.py
+++ b/vllm/config/kv_transfer.py
@@ -47,6 +47,13 @@ def hisparse_host_pool_gib(
     return host_pool_gib
 
 
+# Connectors that register the device KV cache with an RDMA engine whatever
+# their buffer device; NIXL is checked separately as it can stage through host.
+_KV_CACHE_REGISTERING_CONNECTORS = frozenset(
+    {"MooncakeConnector", "MooncakeStoreConnector", "MoRIIOConnector"}
+)
+
+
 def kv_buffer_device_default_factory() -> str:
     from vllm.platforms import current_platform
 
@@ -138,6 +145,33 @@ def is_kv_consumer(self) -> bool:
     def get_from_extra_config(self, key, default) -> Any:
         return self.kv_connector_extra_config.get(key, default)
 
+    def uses_custom_mem_pool(self) -> bool:
+        """Whether a connector allocates the KV cache from its own memory pool
+        (custom_mem_pool), directly or in MultiConnector."""
+        configs = [self.kv_connector_extra_config]
+        if self.kv_connector == "MultiConnector":
+            configs += [
+                child.get("kv_connector_extra_config", {})
+                for child in self.kv_connector_extra_config.get("connectors", [])
+            ]
+        return any(config.get("custom_mem_pool") for config in configs)
+
+    def registers_kv_cache_memory(self) -> bool:
+        """Whether a connector registers the device KV cache itself with UCX or
+        RDMA (NIXL, Mooncake, MoRIIO), directly or in MultiConnector. NIXL with
+        a CPU kv_buffer_device registers only its host staging buffers."""
+        children = [{"kv_connector": self.kv_connector}]
+        if self.kv_connector == "MultiConnector":
+            children = self.kv_connector_extra_config.get("connectors", [])
+        for child in children:
+            name = child.get("kv_connector") or ""
+            if name in _KV_CACHE_REGISTERING_CONNECTORS:
+                return True
+            buffer_device = child.get("kv_buffer_device", self.kv_buffer_device)
+            if name.startswith("Nixl") and buffer_device != "cpu":
+                return True
+        return False
+
     def has_connector(self, connector_name: str) -> bool:
         """Whether ``connector_name`` is configured, directly or in MultiConnector."""
         if self.kv_connector == connector_name:
diff --git a/vllm/config/load.py b/vllm/config/load.py
index d82997be17c3..935436360dda 100644
--- a/vllm/config/load.py
+++ b/vllm/config/load.py
@@ -67,7 +67,7 @@ class LoadConfig:
     Specifies the loading strategy for safetensors weights.
 
     - None (default): Uses memory-mapped (lazy) loading. When an NFS
-      filesystem is detected and the total checkpoint size fits within 90%%
+      filesystem is detected and the total checkpoint size fits within 90%
       of available RAM, prefetching is enabled automatically.
     - "lazy": Weights are memory-mapped from the file. This enables
       on-demand loading and is highly efficient for models on local storage.
diff --git a/vllm/config/model.py b/vllm/config/model.py
index d0deb2c1bae5..382a41486f32 100644
--- a/vllm/config/model.py
+++ b/vllm/config/model.py
@@ -89,6 +89,7 @@
 TokenizerMode = Literal[
     "auto",
     "hf",
+    "glm5_next",
     "alpamayo2_super",
     "mistral",
     "deepseek_v32",
@@ -185,6 +186,10 @@ class ModelConfig:
     - "deepseek_v41" will use the DeepSeek V4.1 prompt encoder.
     - "kimi_k3" will always use the "hf" tokenizer but render chat prompts
       with Kimi K3's Python XTML encoding instead of a Jinja template.
+    - "glm5_next" uses the "hf" tokenizer and renderer on the Python frontend;
+      the Rust frontend renders GLM-5.3 and GLM-5.3-Flash prompts natively,
+      and selects this under "auto" for the `glm5_next` model type
+      (GLM-5.3-Flash).
     - "cohere" uses the standard HF tokenizer but renders the chat template
       via the `cohere_melody` library (cmd3 / cmd4 templates) instead of
       Jinja, and surfaces grounded-citation metadata on responses.
@@ -2067,6 +2072,18 @@ def supports_kv_cache_grouping(self) -> bool:
     def has_inner_state(self):
         return self._model_info.has_inner_state
 
+    @property
+    def speculator_config_class(self) -> str | None:
+        """For a speculative decoding draft model, the config class of the
+        speculator it's for, if it declares one (see `SpeculativeDraftModel`)."""
+        return self._model_info.speculator_config_class
+
+    @property
+    def may_use_target_embed(self) -> bool:
+        """For a speculative decoding draft model, whether it may reuse the
+        target's input embedding under pipeline parallelism."""
+        return self._model_info.may_use_target_embed
+
     @property
     def supports_replayssm(self) -> bool:
         return self._model_info.supports_replayssm
diff --git a/vllm/config/multimodal.py b/vllm/config/multimodal.py
index 5a69787b73af..db3d86909a80 100644
--- a/vllm/config/multimodal.py
+++ b/vllm/config/multimodal.py
@@ -212,31 +212,31 @@ def _merge_mm_processor_kwargs(
 class BaseDummyOptions:
     """Base options for generating dummy data during profiling."""
 
-    count: int = Field(999, ge=0)
+    count: int = Field(default=999, ge=0)
 
 
 @dataclass(config=ConfigDict(extra="forbid"))
 class VideoDummyOptions(BaseDummyOptions):
     """Options for generating dummy video data during profiling."""
 
-    num_frames: int | None = Field(None, gt=0)
-    width: int | None = Field(None, gt=0)
-    height: int | None = Field(None, gt=0)
+    num_frames: int | None = Field(default=None, gt=0)
+    width: int | None = Field(default=None, gt=0)
+    height: int | None = Field(default=None, gt=0)
 
 
 @dataclass(config=ConfigDict(extra="forbid"))
 class ImageDummyOptions(BaseDummyOptions):
     """Options for generating dummy image data during profiling."""
 
-    width: int | None = Field(None, gt=0)
-    height: int | None = Field(None, gt=0)
+    width: int | None = Field(default=None, gt=0)
+    height: int | None = Field(default=None, gt=0)
 
 
 @dataclass(config=ConfigDict(extra="forbid"))
 class AudioDummyOptions(BaseDummyOptions):
     """Options for generating dummy audio data during profiling."""
 
-    length: int | None = Field(None, gt=0)
+    length: int | None = Field(default=None, gt=0)
 
 
 @final
diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py
index f51a35c5c4a9..02b671e2e2aa 100644
--- a/vllm/config/parallel.py
+++ b/vllm/config/parallel.py
@@ -24,7 +24,7 @@
     from ray.runtime_env import RuntimeEnv
     from ray.util.placement_group import PlacementGroup
 
-    from vllm.config.fault_tolerance import FaultToleranceConfig
+    from vllm.config.model import ModelConfig
     from vllm.v1.executor import Executor
 else:
     RuntimeEnv = Any
@@ -40,7 +40,7 @@
 EPLBPolicyOption = Literal["default"]
 DCPCommBackend = Literal["ag_rs", "a2a"]
 EPLBCommunicatorBackend = Literal[
-    "torch_nccl", "torch_gloo", "torch_xccl", "nixl", "pynccl"
+    "torch_nccl", "torch_gloo", "torch_xccl", "nixl", "pynccl", "symm_mem"
 ]
 All2AllBackend = Literal[
     "naive",
@@ -51,6 +51,7 @@
     "mori_high_throughput",
     "mori_low_latency",
     "moonep",
+    "mori",
     "nixl_ep",
     "allgather_reducescatter",
     "flashinfer_all2allv",  # temporary alias for flashinfer_nvlink_two_sided
@@ -74,8 +75,13 @@ class EPLBConfig:
     of the last `lb_window_size` steps will be used for rearranging experts.
     """
 
-    num_redundant_experts: int = Field(default=0, ge=0)
-    """Number of redundant experts to use for expert parallelism."""
+    num_redundant_experts: int = Field(default=None, ge=0)  # type: ignore[assignment]
+    """Number of redundant experts to use for expert parallelism.
+    If None (default), the minimum valid value will be computed automatically
+    based on the number of logical experts and the expert parallel size.
+    The value will be rounded up to the nearest multiple of the expert parallel
+    size.
+    """
 
     log_balancedness: bool = False
     """
@@ -102,6 +108,8 @@ class EPLBConfig:
     - "torch_xccl": Use torch.distributed XCCL device P2P on XPU
     - "nixl": Use NIXL with staged send/recv buffers
     - "pynccl": Use PyNccl send/recv
+    - "symm_mem": Push expert rows over CUDA symmetric memory; needs all EP
+      ranks in one NVLink domain (MNNVL across nodes)
     - None: Auto-select backend ("torch_xccl" on XPU, prefers "nixl" 
       on CUDA, falls back to "torch_gloo")
     """
@@ -224,14 +232,23 @@ class ParallelConfig:
     - "deepep_high_throughput": Use deepep high-throughput kernels
     - "deepep_low_latency": Use deepep low-latency kernels
     - "mori_high_throughput": MoRI EP with InterNodeV1 for multi-node
+      (deprecated; use "mori")
     - "mori_low_latency": MoRI EP with InterNodeV1LL for multi-node
+      (deprecated; use "mori")
     - "moonep": MoonEP balanced EP with dynamic redundant experts (NVLink)
+    - "mori": MoRI EP; uses EPv2 where supported (see `VLLM_MORI_EP_V2`) and
+      EPv1 otherwise, picking the EPv1 inter-node kernel from the token budget
+      (see `VLLM_MORI_DISPATCH_INTER_KERNEL_SWITCH_THRESHOLD`)
     - "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
     - "passthrough": No all2all at all: the MoE backend dispatches and combines
       itself (e.g. the FlashInfer MoE-EP megakernels). Bound automatically for
       such MoE backends; not selectable from the CLI."""
+    mori_ep_v2: bool | None = None
+    """Whether the "mori" all2all backend runs MoRI EPv2. Resolved from
+    `VLLM_MORI_EP_V2` and the deployment by `resolve_mori_all2all`; not meant
+    to be set by users."""
 
     max_parallel_loading_workers: int | None = Field(default=None, ge=1)
     """Maximum number of parallel loading workers when loading model
@@ -326,16 +343,32 @@ class is dynamically inherited by the worker class. This is used to inject
     nnodes: int = Field(default=1, ge=1)
     """num of nodes for multi-node distributed
     inference when distributed_executor_backend is mp."""
-    numa_bind: bool = False
+    numa_bind: bool | None = None
     """Enable NUMA binding for GPU worker subprocesses.
 
-    By default, workers are pinned to their GPU's NUMA-local CPUs and
-    memory; on PCT-capable Xeons they also auto-bind to the SKU's
-    PCT priority cores.
+    Workers are pinned to their GPU's NUMA-local CPUs; on PCT-capable
+    Xeons they also auto-bind to the SKU's PCT priority cores.
+
+    If None (default), vLLM binds only when the environment supports it
+    (multiple NUMA nodes, `numactl` installed, unrestricted CPU affinity,
+    a default NUMA memory policy, and the `spawn` start method with an
+    mp/uni executor and mp DP, including PCP).
+    Ray, external launchers and device-ID remapping are excluded
+    from automatic binding. Automatic mode prefers GPU-local memory but
+    permits allocations on other nodes. Unsupported environments start
+    unbound. With `True`, strict memory binding may override an existing
+    policy, and a missing `numactl` or undetectable GPU-to-NUMA topology is
+    a startup error; `False` disables binding.
     """
+    _numa_bind_auto: bool = Field(default=False, repr=False)
+    """Internal launch decision: use soft memory preference in spawned children."""
     numa_bind_nodes: list[int] | None = None
     """NUMA node to bind each GPU worker to.
 
+    Setting this list implies explicit `numa_bind=True` and requests strict
+    `--membind`, not automatic soft memory preference. It cannot be combined
+    with `numa_bind=False`.
+
     Specify one NUMA node per visible GPU, for example `[0, 0, 1, 1]`
     for a 4-GPU system with GPUs 0-1 on NUMA node 0 and GPUs 2-3 on
     NUMA node 1. If unset and `numa_bind=True`, vLLM auto-detects the
@@ -345,6 +378,10 @@ class is dynamically inherited by the worker class. This is used to inject
     numa_bind_cpus: list[str] | None = None
     """Optional CPU lists to bind each GPU worker to.
 
+    Setting this list implies explicit `numa_bind=True` and requests strict
+    `--membind`, not CPU-only binding or automatic soft memory preference.
+    It cannot be combined with `numa_bind=False`.
+
     Specify one CPU list per visible GPU, for example
     `["0-3", "4-7", "8-11", "12-15"]`. When set, vLLM uses
     `numactl --physcpubind` instead of `--cpunodebind`. This is useful
@@ -465,7 +502,7 @@ class is dynamically inherited by the worker class. This is used to inject
 
     enable_fault_tolerance: bool = False
     """Enable fault tolerance for detailed error recovery,
-    such as scaling down fault DPEngineCore.
+    such as scaling down faulty DPEngineCores.
     """
 
     fault_tolerance_config: FaultToleranceConfig = Field(
@@ -553,12 +590,12 @@ def _validate_parallel_config(self) -> Self:
                 "data_parallel_external_lb can only be set when data_parallel_size > 1"
             )
 
-        if not self.numa_bind and (
-            self.numa_bind_nodes is not None or self.numa_bind_cpus is not None
-        ):
-            raise ValueError(
-                "numa_bind_nodes and numa_bind_cpus require numa_bind=True."
-            )
+        if self.numa_bind_nodes is not None or self.numa_bind_cpus is not None:
+            if self.numa_bind is False:
+                raise ValueError(
+                    "NUMA binding lists cannot be combined with numa_bind=False."
+                )
+            self.numa_bind = True
 
         if self.enable_eplb:
             if not current_platform.is_cuda_alike() and not current_platform.is_xpu():
@@ -582,14 +619,17 @@ def _validate_parallel_config(self) -> Self:
                     f"PCP={self.prefill_context_parallel_size}, "
                     f"DP={self.data_parallel_size}."
                 )
-        else:
-            if self.eplb_config.num_redundant_experts != 0:
+
+        if not self.enable_eplb:
+            if self.eplb_config.num_redundant_experts not in (None, 0):
                 raise ValueError(
                     "num_redundant_experts is set to "
                     f"{self.eplb_config.num_redundant_experts} but EPLB is not "
                     "enabled. Either enable EPLB or unset "
                     "num_redundant_experts."
                 )
+            # Only EPLB needs the model's expert count to pick a default.
+            self.eplb_config.num_redundant_experts = 0
 
         tp = self.tensor_parallel_size
         pcp = self.prefill_context_parallel_size
@@ -608,6 +648,40 @@ def _validate_parallel_config(self) -> Self:
 
         return self
 
+    def compute_eplb_num_redundant_experts(self, num_logical_experts: int) -> None:
+        """Compute and set num_redundant_experts if not explicitly specified.
+
+        This method should be called after ParallelConfig is initialized and
+        when the number of logical experts is known (from ModelConfig).
+
+        Args:
+            num_logical_experts: The number of logical experts from the model.
+
+        """
+        if self.eplb_config.num_redundant_experts is not None:
+            # Already explicitly set, don't override
+            return
+
+        if self.enable_eplb:
+            # The EP group spans the TP x PCP x DP ranks.
+            ep_size = (
+                self.tensor_parallel_size
+                * self.prefill_context_parallel_size
+                * self.data_parallel_size
+            )
+            # Ensure (num_logical_experts + num_redundant_experts) is
+            # divisible by ep_size, supporting non-standard ep_size values
+            min_redundant = (ep_size - num_logical_experts % ep_size) % ep_size
+            self.eplb_config.num_redundant_experts = min_redundant
+            logger.info_once(
+                "EPLB num_redundant_experts not specified, "
+                "defaulting to minimum valid value: %d "
+                "(num_logical_experts=%d, ep_size=%d)",
+                min_redundant,
+                num_logical_experts,
+                ep_size,
+            )
+
     def set_dcp_defaults(
         self,
         comm_backend: DCPCommBackend = "ag_rs",
@@ -661,6 +735,18 @@ def cpu_distributed_timeout(self) -> timedelta | None:
         seconds = self.cpu_distributed_timeout_seconds
         return timedelta(seconds=seconds) if seconds is not None else None
 
+    @property
+    def effective_cpu_distributed_timeout(self) -> timedelta | None:
+        """The timeout cpu groups operate with right now: the configured
+        value, except under FT before steady state, where group timeouts
+        are deferred because uneven rank startup would trip them."""
+        if self.enable_fault_tolerance:
+            from vllm.distributed.utils import steady_state_entered
+
+            if not steady_state_entered():
+                return None
+        return self.cpu_distributed_timeout
+
     def get_next_dp_init_port(self) -> int:
         """We might need to initialize process groups in multiple
         processes that is related to data parallelism,
@@ -728,6 +814,10 @@ def stateless_init_dp_group(
             stateless_init_torch_distributed_process_group,
         )
 
+        # Under FT, startup-created groups get no configured timeout: uneven
+        # rank startup would trip it. The sentinel applies it after startup.
+        timeout = self.effective_cpu_distributed_timeout
+
         max_retries = 5
         last_exc: Exception | None = None
         for _ in range(max_retries):
@@ -742,7 +832,7 @@ def stateless_init_dp_group(
                     backend="gloo",
                     return_store=return_store,
                     listen_socket=listen_socket,
-                    timeout=self.cpu_distributed_timeout,
+                    timeout=timeout,
                 )
             except DistNetworkError as e:
                 # We only want to retry when the root cause is EADDRINUSE.
@@ -777,6 +867,7 @@ def use_sequence_parallel_moe(self) -> bool:
                 "passthrough",
                 "mori_high_throughput",
                 "mori_low_latency",
+                "mori",
                 "nixl_ep",
             )
             and self.enable_expert_parallel
@@ -949,6 +1040,7 @@ def compute_hash(self):
             # NUMA node in numa_bind_nodes (see vllm/utils/numa_utils.py
             # `_get_numa_node`), which would otherwise diverge the DP hash.
             "numa_bind",
+            "_numa_bind_auto",
             "numa_bind_nodes",
             "numa_bind_cpus",
             "assigned_physical_gpu_ids",
@@ -1198,3 +1290,148 @@ def reconfigure_for_independent_dp_rank(self) -> None:
         self.data_parallel_rank = 0
         self.nnodes = nnodes
         self.node_rank = node_rank
+
+
+def resolve_mori_all2all(
+    parallel_config: ParallelConfig,
+    model_config: "ModelConfig | None",
+    moe_backend: str,
+    lsa_spans_hosts: bool = False,
+    sdma_picks_epv1: bool = True,
+    validated_arch: bool = True,
+) -> None:
+    """Resolve whether the "mori" all2all backend runs MoRI EPv2 or EPv1.
+
+    An unset `VLLM_MORI_EP_V2` prefers EPv2 and falls back to EPv1 where EPv2
+    is unsupported, an explicit true keeps EPv2 regardless, and false selects
+    EPv1. `lsa_spans_hosts` says whether a CCO-LSA team can span hosts (e.g.
+    gfx1250's UALink vPOD). `sdma_picks_epv1` says whether `MORI_ENABLE_SDMA=1`
+    selects EPv1; gfx1250 keeps EPv2 under SDMA. `validated_arch` says whether
+    EPv2 is validated on this GPU.
+    """
+    if parallel_config.all2all_backend != "mori":
+        return
+    if parallel_config.mori_ep_v2 is None:
+        parallel_config.mori_ep_v2 = _resolve_mori_ep_v2(
+            parallel_config,
+            model_config,
+            moe_backend,
+            lsa_spans_hosts,
+            sdma_picks_epv1,
+            validated_arch,
+        )
+    logger.info_once("MoRI all2all: EP%s", "v2" if parallel_config.mori_ep_v2 else "v1")
+
+
+def _resolve_mori_ep_v2(
+    parallel_config: ParallelConfig,
+    model_config: "ModelConfig | None",
+    moe_backend: str,
+    lsa_spans_hosts: bool,
+    sdma_picks_epv1: bool,
+    validated_arch: bool,
+) -> bool:
+    if envs.VLLM_MORI_EP_V2 is False:
+        return False
+    reason = _mori_epv2_unsupported_reason(
+        parallel_config,
+        model_config,
+        moe_backend,
+        lsa_spans_hosts,
+        sdma_picks_epv1,
+        validated_arch,
+    )
+    if reason is None:
+        return True
+    if envs.VLLM_MORI_EP_V2:
+        logger.warning(
+            "VLLM_MORI_EP_V2=1 is set explicitly, so MORI EPv2 is used even "
+            "though it does not support %s.",
+            reason,
+        )
+        return True
+    logger.warning(
+        "MORI EPv2 does not support %s; falling back to MORI EPv1. "
+        "Set VLLM_MORI_EP_V2=1 to force EPv2, or 0 to select EPv1.",
+        reason,
+    )
+    return False
+
+
+# Used only when the launch layout cannot tell whether EP spans hosts.
+_MORI_GPUS_PER_HOST = 8
+
+
+def _ep_spans_hosts(parallel_config: ParallelConfig, ep_size: int) -> bool:
+    if parallel_config.nnodes_within_dp > 1:
+        return True
+    if (
+        parallel_config.data_parallel_external_lb
+        or parallel_config.data_parallel_backend == "ray"
+    ):
+        # Each launch only knows its own DP ranks.
+        return ep_size > _MORI_GPUS_PER_HOST
+    return parallel_config.data_parallel_size_local < parallel_config.data_parallel_size
+
+
+def _mori_epv2_unsupported_reason(
+    parallel_config: ParallelConfig,
+    model_config: "ModelConfig | None",
+    moe_backend: str,
+    lsa_spans_hosts: bool,
+    sdma_picks_epv1: bool,
+    validated_arch: bool,
+) -> str | None:
+    from vllm.utils.import_utils import mori_epv2_unavailable_reason
+
+    reason = mori_epv2_unavailable_reason()
+    if reason is not None:
+        return f"this MoRI ({reason}; upgrade MoRI)"
+    if not validated_arch:
+        return "this GPU (EPv2 is validated on gfx950 and gfx1250)"
+    ep_size = (
+        parallel_config.tensor_parallel_size
+        * parallel_config.data_parallel_size
+        * parallel_config.prefill_context_parallel_size
+    )
+    # A guess before startup; MoriEPv2All2AllManager checks the actual
+    # CCO-LSA team.
+    if not lsa_spans_hosts and _ep_spans_hosts(parallel_config, ep_size):
+        return f"ep_size={ep_size} across hosts (EPv2 needs one host)"
+    sdma = os.environ.get("MORI_ENABLE_SDMA", "0").lower() in ("1", "true")
+    if sdma and sdma_picks_epv1:
+        return "MORI_ENABLE_SDMA=1 (SDMA is EPv1-only)"
+    if moe_backend == "aiter_mega_moe":
+        return "moe_backend=aiter_mega_moe (MegaMoE runs on MoRI EPv1 shmem)"
+    if model_config is not None and model_config.dtype != torch.bfloat16:
+        return f"dtype={model_config.dtype} (EPv2 requires bfloat16)"
+    if model_config is not None and _quantizes_activations_per_tensor(model_config):
+        return "per-tensor FP8 activations (AITER cannot quantize them after dispatch)"
+    return None
+
+
+def _quantizes_activations_per_tensor(model_config: "ModelConfig") -> bool:
+    # Under EPv2 AITER quantizes the received rows itself, which its per-tensor
+    # quant does not support. Other formats are caught at kernel set-up.
+    method = getattr(model_config, "quantization", None)
+    arch_config = getattr(model_config, "model_arch_config", None)
+    quant_cfg = getattr(arch_config, "quantization_config", None)
+    if not isinstance(quant_cfg, dict):
+        quant_cfg = {}
+    if method == "fp8":
+        return not quant_cfg.get("weight_block_size")
+    if method == "modelopt":
+        return quant_cfg.get("quantization", quant_cfg).get("quant_algo") == "FP8"
+    if method == "compressed-tensors":
+        groups = (quant_cfg.get("config_groups") or {}).values()
+        acts = [group.get("input_activations") or {} for group in groups]
+        per_tensor = [
+            a.get("strategy") == "tensor" for a in acts if a.get("type") == "float"
+        ]
+        return any(per_tensor)
+    if method == "quark":
+        global_cfg = quant_cfg.get("global_quant_config") or {}
+        inputs = global_cfg.get("input_tensors") or {}
+        dtype = str(inputs.get("dtype", ""))
+        return inputs.get("qscheme") == "per_tensor" and dtype.startswith("fp8")
+    return False
diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py
index 9a6f1a8afc7b..70ecfd1ccf9b 100644
--- a/vllm/config/speculative.py
+++ b/vllm/config/speculative.py
@@ -3,7 +3,9 @@
 
 import copy
 import functools
+import os
 from collections.abc import Callable, Mapping
+from dataclasses import field
 from typing import TYPE_CHECKING, Any, Literal, Self, get_args
 
 from pydantic import Field, SkipValidation, field_validator, model_validator
@@ -16,9 +18,17 @@
 from vllm.config.utils import config, replace
 from vllm.logger import init_logger
 from vllm.transformers_utils.config import get_hf_text_config, iter_rope_parameters
+from vllm.transformers_utils.repo_utils import get_subfolder_path
 from vllm.utils.hashing import safe_hash
 from vllm.utils.import_utils import LazyLoader, has_arctic_inference
 from vllm.v1.attention.backends.registry import AttentionBackendEnum
+from vllm.v1.worker.gpu.spec_decode.speculator_config import (
+    SpeculatorConfig,
+    get_custom_speculator_config_class,
+    get_method_config_class,
+    get_speculator_config_class,
+    registered_speculative_methods,
+)
 
 if TYPE_CHECKING:
     from transformers import PreTrainedConfig
@@ -45,6 +55,7 @@
     "glm_ocr_mtp",
     "ernie_mtp",
     "nemotron_h_mtp",
+    "nemotron3_5_transcribe_mtp",
     "exaone_moe_mtp",
     "exaone4_5_mtp",
     "qwen3_next_mtp",
@@ -90,12 +101,14 @@
     }
 )
 _QWEN3_OMNI_DSPARK_ARCHITECTURE = "Qwen3OmniDSparkModel"
-_MUSE_SPARK_DSPARK_ARCHITECTURE = "MuseSparkDSparkModel"
-
-# Draft architectures that carry an XPress refiner. EAGLEConfig rewrites a dflash
-# draft's architecture to DFlash{arch}, so a checkpoint is seen under either name;
-# the dispatch in v1/worker/gpu/spec_decode reads this same set.
-XPRESS_ARCHITECTURES = frozenset({"Qwen3XPressModel", "DFlashQwen3XPressModel"})
+# See `SpeculatorConfig.supported_options`.
+_METHOD_SPECIFIC_OPTIONS = (
+    "dspark_draft_topk",
+    "draft_lm_head_quantization",
+    "index_share_for_mtp_iteration",
+    "xpress_topc",
+    "speculator_options",
+)
 
 # RoPE types for which get_rope uses max_position_embeddings only to size the
 # cos_sin_cache, so raising it leaves the rotary frequencies unchanged. Others
@@ -113,6 +126,31 @@ def _rope_only_sizes_cache(hf_config: PreTrainedConfig) -> bool:
     )
 
 
+_MINIMAX_M3_MODEL_TYPES = ("minimax_m3_vl", "minimax_m3_text")
+_MINIMAX_M3_DSPARK_SUBFOLDER = "dspark"
+
+
+def _get_minimax_m3_dspark_draft(
+    target_model_config: ModelConfig, cache_dir: str | None
+) -> str:
+    """MiniMax-M3 checkpoints ship the DSpark draft as a standalone checkpoint
+    in their `dspark/` subfolder."""
+    draft = get_subfolder_path(
+        target_model_config.model,
+        _MINIMAX_M3_DSPARK_SUBFOLDER,
+        revision=target_model_config.revision,
+        cache_dir=cache_dir,
+    )
+    if not os.path.isfile(os.path.join(draft, "config.json")):
+        raise ValueError(
+            f"MiniMax-M3 DSpark found no draft checkpoint in the "
+            f"'{_MINIMAX_M3_DSPARK_SUBFOLDER}/' subfolder of "
+            f"{target_model_config.model}. Set speculative_config.model to the "
+            "draft checkpoint."
+        )
+    return draft
+
+
 def _is_qwen3_omni_target(model_config: ModelConfig) -> bool:
     hf_config = model_config.hf_config
     architectures = set(getattr(model_config, "architectures", ()) or ())
@@ -416,14 +454,26 @@ class SpeculativeConfig:
     model: str | None = None
     """The name of the draft model, eagle head, or additional weights, if
     provided."""
-    method: SpeculativeMethod | None = None
+    method: SpeculativeMethod | str | None = None
     """The name of the speculative method to use. If users provide and set the
     `model` param, the speculative method type will be detected automatically
     if possible, if `model` param is not provided, the method name must be
     provided.
 
     If using `ngram` method, the related configuration `prompt_lookup_max` and
-    `prompt_lookup_min` should be considered."""
+    `prompt_lookup_min` should be considered.
+
+    A configured method is never changed based on the draft checkpoint, but it
+    is checked against it where the checkpoint can be recognized (e.g. `mtp`
+    requires an MTP draft, and `draft_model` rejects one).
+
+    Besides the built-in methods, this can be a custom method registered with
+    `vllm.v1.worker.gpu.spec_decode.speculator_config.register_speculative_method`
+    (e.g. by a vLLM plugin), which requires the V2 model runner. It is then
+    moved to `custom_method`, and this is set to the built-in method that its
+    config is based on (None if it isn't based on one)."""
+    custom_method: str | None = field(default=None, init=False)
+    """The configured custom speculative method, if any (see `method`)."""
     draft_tensor_parallel_size: int | None = Field(default=None, ge=1)
     """The degree of the tensor parallelism for the draft model. Can only be 1
     or the same as the target model's tensor parallel size."""
@@ -482,11 +532,11 @@ class SpeculativeConfig:
     sharing a KV store and external consumers of KV cache events must use the
     same setting."""
     draft_lm_head_quantization: Literal["fp8", "nvfp4"] | None = None
-    """Quantize the drafter's copy of a target lm_head it shares (EAGLE and
-    MTP). "fp8" uses per-row weight and per-token activation scales on the
-    platform's FP8 linear kernel; "nvfp4" is weight-only on the platform's
-    W4A16 kernel. The target verifies with its own head, so only the
-    acceptance length can change. Model Runner V2 only."""
+    """Quantize the drafter's copy of a target lm_head it shares (EAGLE, MTP
+    and DeepSeek-V4/V4.1 DSpark). "fp8" uses per-row weight and per-token
+    activation scales on the platform's FP8 linear kernel; "nvfp4" is
+    weight-only on the platform's W4A16 kernel. The target verifies with its
+    own head, so only the acceptance length can change. Model Runner V2 only."""
     use_local_argmax_reduction: bool = False
     """Use vocab-parallel local argmax instead of all-gathering full logits
     for draft token generation. Reduces communication from O(vocab_size) to
@@ -512,7 +562,11 @@ class SpeculativeConfig:
     """Enable parallel drafting, where all speculative tokens are generated
     in parallel rather than sequentially. This can improve performance but
     requires the speculative model be trained to support parallel drafting.
-    Only compatible with EAGLE and draft model methods."""
+    Only compatible with EAGLE and draft model methods.
+
+    Deprecated: V1 model runner only (the V2 model runner falls back to V1 for
+    it), and will be removed with it. V2 speculators that draft in parallel,
+    such as DFlash and DSpark, do so without this option."""
 
     # required configuration params passed from engine
     target_model_config: SkipValidation[ModelConfig] = None  # type: ignore
@@ -649,6 +703,10 @@ def _resolve_synthetic_acceptance_rates(
     vocabulary. Takes precedence over ``xpress_topc`` in the draft
     checkpoint's config, which is used when this is unset."""
 
+    speculator_options: dict[str, Any] | None = None
+    """Options for a custom speculator, interpreted by its `SpeculatorConfig`,
+    which must list `"speculator_options"` in its `supported_options`."""
+
     fly_window_size: int | None = Field(default=None, ge=1)
     """Number of subsequent native acceptance decisions checked by FLy.
     If unset, defaults to ``min(6, num_speculative_tokens - 1)``."""
@@ -673,15 +731,8 @@ def compute_hash(self) -> str:
         factors: list[Any] = []
         # Eagle3 and extract_hidden_states affect the computation graph because
         # they return intermediate hidden states in addition to the final hidden state.
-        uses_aux_hidden_states = self.method in (
-            "eagle3",
-            "extract_hidden_states",
-            "dflash",
-            "dspark",
-        )
+        uses_aux_hidden_states = self.uses_aux_hidden_states
         factors.append(uses_aux_hidden_states)
-        if self.method == "dspark":
-            factors.append(self.enable_adaptive_verification)
 
         if self.draft_model_config is not None:
             factors.append(self.draft_model_config.compute_hash())
@@ -695,21 +746,41 @@ def compute_hash(self) -> str:
                 # Convert to tuple to make it hashable
                 factors.append(tuple(layer_ids))
 
-        if self.method == "mtp" and self.draft_model_config is not None:
-            factors.append(
-                getattr(
-                    self.draft_model_config.hf_config,
-                    "index_share_for_mtp_iteration",
-                    False,
-                )
-            )
+        # A custom speculator may compile a different draft model graph.
+        factors.append(self.custom_method)
+        if self.speculator_options:
+            factors.append(sorted(self.speculator_options.items()))
+        factors.extend(self.speculator_config.hash_factors())
 
         hash_str = safe_hash(str(factors).encode(), usedforsecurity=False).hexdigest()
         return hash_str
 
     @staticmethod
-    def hf_config_override(hf_config: PreTrainedConfig) -> PreTrainedConfig:
+    def hf_config_override(
+        hf_config: PreTrainedConfig,
+        transformers_mtp_overrides: dict[str, Any] | None = None,
+    ) -> PreTrainedConfig:
+        if transformers_mtp_overrides is not None:
+            # The Transformers modeling backend builds the MTP draft from the
+            # target's config, so it needs none of the rewrites below
+            hf_config.update(transformers_mtp_overrides)
+            return hf_config
         initial_architecture = hf_config.architectures[0]
+        if (
+            initial_architecture == "DSparkDraftModel"
+            and hf_config.model_type == "qwen3"
+        ):
+            # Qwen3 DSpark drafts reuse the DeepSeek-V4 DSpark architecture name.
+            hf_config.update({"architectures": ["Qwen3DSparkModel"]})
+            return hf_config
+        if initial_architecture == "Gemma4DSparkModel":
+            # Normalize the Gemma4 DSpark draft's config keys to DSpark's.
+            if (
+                getattr(hf_config, "dspark_target_layer_ids", None) is None
+                and getattr(hf_config, "target_layer_ids", None) is not None
+            ):
+                hf_config.dspark_target_layer_ids = hf_config.target_layer_ids
+            return hf_config
         if initial_architecture == "RadioMBartPrefixForCausalLM" and getattr(
             getattr(hf_config, "decoder", None), "causal_vision_prefix", False
         ):
@@ -884,6 +955,25 @@ def hf_config_override(hf_config: PreTrainedConfig) -> PreTrainedConfig:
             # Promote VLM's text_config so MTP detection below fires correctly
             hf_config = hf_config.text_config
 
+        if (
+            hf_config.architectures[0]
+            == "Nemotron3_5TranscribeForConditionalGeneration"
+        ):
+            # Its MTP head reads NeMo checkpoint names and takes the audio
+            # embeddings, so it has its own draft class.
+            quantization_config = getattr(hf_config, "quantization_config", None)
+            hf_config = hf_config.text_config
+            if quantization_config is not None:
+                # The FP8 and NVFP4 checkpoints quantize the MTP head too.
+                hf_config.quantization_config = copy.deepcopy(quantization_config)
+            hf_config.model_type = "nemotron3_5_transcribe_mtp"
+            hf_config.update(
+                {
+                    "n_predict": hf_config.num_nextn_predict_layers,
+                    "architectures": ["Nemotron3_5TranscribeMTPModel"],
+                }
+            )
+
         if (
             hf_config.model_type in {"nemotron_h", "nemotron_h_puzzle"}
             and hasattr(hf_config, "num_nextn_predict_layers")
@@ -1106,7 +1196,25 @@ def hf_config_override(hf_config: PreTrainedConfig) -> PreTrainedConfig:
                 text_config.num_kv_shared_layers = 0
             hf_config.update({"n_predict": 1, "architectures": ["Gemma4MTPModel"]})
 
-        if (
+        if initial_architecture == "DSparkMiniMaxDraftModel":
+            quantization_config = getattr(hf_config, "quantization_config", None)
+            hf_config = getattr(hf_config, "text_config", hf_config)
+            if (
+                quantization_config is not None
+                and getattr(hf_config, "quantization_config", None) is None
+            ):
+                hf_config.update({"quantization_config": quantization_config})
+            dflash_config = dict(getattr(hf_config, "dflash_config", None) or {})
+            dflash_config.setdefault("causal", False)
+            hf_config.model_type = "minimax_m3_dspark"
+            hf_config.update(
+                {
+                    "architectures": ["DSparkMiniMaxDraftModel"],
+                    "dflash_config": dflash_config,
+                    "markov_rank": hf_config.dspark_markov_rank,
+                }
+            )
+        elif (
             hf_config.model_type == "minimax_m3_vl"
             or initial_architecture == "MiniMaxM3SparseForConditionalGeneration"
         ):
@@ -1164,13 +1272,17 @@ def hf_config_override(hf_config: PreTrainedConfig) -> PreTrainedConfig:
     def _apply_composed_hf_override(
         target_hf_overrides: Callable[[PreTrainedConfig], PreTrainedConfig],
         hf_config: PreTrainedConfig,
+        transformers_mtp_overrides: dict[str, Any] | None = None,
     ) -> PreTrainedConfig:
-        hf_config = SpeculativeConfig.hf_config_override(hf_config)
+        hf_config = SpeculativeConfig.hf_config_override(
+            hf_config, transformers_mtp_overrides
+        )
         return target_hf_overrides(hf_config)
 
     @staticmethod
     def compose_draft_hf_overrides(
         target_hf_overrides: HfOverrides | None,
+        transformers_mtp_overrides: dict[str, Any] | None = None,
     ) -> Callable[[PreTrainedConfig], PreTrainedConfig]:
         """Build the ``hf_overrides`` for the draft ``ModelConfig``.
 
@@ -1188,10 +1300,15 @@ def compose_draft_hf_overrides(
         method instead.
         """
         if not callable(target_hf_overrides):
-            return SpeculativeConfig.hf_config_override
+            return functools.partial(
+                SpeculativeConfig.hf_config_override,
+                transformers_mtp_overrides=transformers_mtp_overrides,
+            )
 
         return functools.partial(
-            SpeculativeConfig._apply_composed_hf_override, target_hf_overrides
+            SpeculativeConfig._apply_composed_hf_override,
+            target_hf_overrides,
+            transformers_mtp_overrides=transformers_mtp_overrides,
         )
 
     @staticmethod
@@ -1216,464 +1333,59 @@ def __post_init__(self):
         # default.
 
         # infer method from user args
-        if self.method is None and SpeculativeConfig._is_custom_proposer_path(
+        if (custom_cls := get_custom_speculator_config_class(self.method)) is not None:
+            logger.warning_once(
+                "Using custom speculative method '%s'. This is an experimental "
+                "feature and the speculator interface is subject to breaking "
+                "changes in future vLLM releases.",
+                self.method,
+            )
+            self.custom_method = self.method
+            self.method = custom_cls.base_method  # type: ignore[assignment]
+        elif self.method is None and SpeculativeConfig._is_custom_proposer_path(
             self.model
         ):
+            # V1 model runner only: a custom proposer's module path.
             self.method = "custom_class"
-        elif self.method is None:
-            if self.model in ("ngram", "[ngram]"):
-                self.method = "ngram"
-            else:
-                self.method = "draft_model"
+        elif self.method is None and self.model in ("ngram", "[ngram]"):
+            self.method = "ngram"
+        elif self.method is None and self.model is None:
+            # No draft checkpoint to detect the method from.
+            self.method = "draft_model"
+        # Otherwise, it stays None until detected from the draft checkpoint.
 
         if self.method in get_args(MTPModelTypes) and self.method != "mtp":
             logger.warning(
                 "method `%s` is deprecated and replaced with mtp.", self.method
             )
             self.method = "mtp"
-
-        if self.model is None and self.num_speculative_tokens is not None:
-            if self.method == "mtp":
-                if self.target_model_config is None:
-                    raise ValueError("target_model_config must be present for mtp")
-                # use the draft model from the same model:
-                # Use model_weights if set (e.g. runai_streamer replaces
-                # model with a local cache dir but keeps the original path
-                # in model_weights), so the draft model can load weights
-                # from the same source as the target model.
-                self.model = (
-                    self.target_model_config.model_weights
-                    or self.target_model_config.model
-                )
-                # Align the quantization of draft model for cases such as
-                # --quantization fp8 with a bf16 checkpoint.
-                if not self.quantization:
-                    self.quantization = self.target_model_config.quantization
-            elif self.method == "dspark":
-                # DeepSeek DSpark can ship the weights inside the target checkpoint
-                if self.target_model_config is None:
-                    raise ValueError("target_model_config must be present for dspark")
-                self.model = self.target_model_config.model
-                if not self.quantization:
-                    self.quantization = self.target_model_config.quantization
-            elif self.method in ("ngram", "[ngram]"):
-                self.model = "ngram"
-            elif self.method == "ngram_gpu":
-                self.model = "ngram_gpu"
-            elif self.method == "suffix":
-                self.model = "suffix"
-            elif self.method == "extract_hidden_states":
-                self.model = "extract_hidden_states"
-            elif self.method == "custom_class":
-                # method was set explicitly, but model should already contain the
-                # custom module path. If not, this is a configuration error.
-                if self.model is None:
-                    raise ValueError(
-                        "method='custom_class' requires 'model' to contain the "
-                        "custom proposer module path (e.g., 'my_module.MyProposer')."
-                    )
-            else:
-                raise ValueError(
-                    "num_speculative_tokens was provided but without speculative model."
-                )
-
-        if self.method in ("ngram", "[ngram]"):
+        elif self.method == "[ngram]":
             self.method = "ngram"
 
-        if self.method in ("ngram", "ngram_gpu"):
-            # Set default values if not provided
-            if self.prompt_lookup_min is None and self.prompt_lookup_max is None:
-                # TODO(woosuk): Tune these values. They are arbitrarily chosen.
-                self.prompt_lookup_min = 5
-                self.prompt_lookup_max = 5
-            elif self.prompt_lookup_min is None:
-                if self.prompt_lookup_max is None:
-                    raise ValueError(
-                        "Either prompt_lookup_max or prompt_lookup_min must be "
-                        "provided when using the ngram method."
-                    )
-                self.prompt_lookup_min = self.prompt_lookup_max
-            elif self.prompt_lookup_max is None:
-                if self.prompt_lookup_min is None:
-                    raise ValueError(
-                        "Either prompt_lookup_max or prompt_lookup_min must be "
-                        "provided when using the ngram method."
-                    )
-                self.prompt_lookup_max = self.prompt_lookup_min
-
-            # Validate values
-            if self.prompt_lookup_min > self.prompt_lookup_max:
-                raise ValueError(
-                    f"prompt_lookup_min={self.prompt_lookup_min} must "
-                    f"be <= prompt_lookup_max={self.prompt_lookup_max}"
-                )
+        config_cls = get_method_config_class(self)
+        # DSpark drafts carry their own width (block size), so the default
+        # draft is resolved even without num_speculative_tokens.
+        if self.model is None and (
+            self.num_speculative_tokens is not None or self.use_dspark()
+        ):
+            self.model = config_cls.default_model(self)
+        config_cls.configure(self)
+        # The draft checkpoint may have determined the speculator.
+        config_cls = get_speculator_config_class(self)
 
-            # TODO: current we still need extract vocab_size from target model
-            # config, in future, we may try refactor it out, and set
-            # draft related config as None here.
-            self.draft_model_config = self.target_model_config
-            self.draft_parallel_config = self.target_parallel_config
-        elif self.method == "suffix":
-            self._validate_suffix_decoding()
-        elif self.method == "custom_class":
-            # Custom class proposer does not need a draft model.
-            # It will dynamically load the user-provided class at runtime.
-            logger.warning_once(
-                "Using a custom class-based proposer backend. This is an "
-                "experimental feature and the proposer interface is subject to "
-                "breaking changes in future vLLM releases."
-            )
-            self.prompt_lookup_max = 0
-            self.prompt_lookup_min = 0
-            self.draft_model_config = self.target_model_config
-            self.draft_parallel_config = self.target_parallel_config
-        elif self.method == "extract_hidden_states":
-            from vllm.transformers_utils.configs.extract_hidden_states import (
-                ExtractHiddenStatesConfig,
-            )
-
-            # ExtractHiddenStatesModel is instantiated manually in load_model()
-            # We just need to store the target model config for KV cache shape info
-            self.model = "extract_hidden_states"
-            self.prompt_lookup_max = 0
-            self.prompt_lookup_min = 0
-
-            if hasattr(self.draft_model_config, "hf_config"):
-                hf_config = self.draft_model_config.hf_config.to_dict()
-            elif (
-                isinstance(self.draft_model_config, dict)
-                and "hf_config" in self.draft_model_config
+        for option in _METHOD_SPECIFIC_OPTIONS:
+            if (
+                getattr(self, option) is not None
+                and option not in config_cls.supported_options
             ):
-                hf_config = self.draft_model_config["hf_config"]
-            else:
-                hf_config = {}
-
-            self.draft_model_config = copy.copy(self.target_model_config)
-            self.draft_model_config.hf_config = ExtractHiddenStatesConfig(
-                self.draft_model_config.hf_config, **hf_config
-            )
-            self.update_arch_()
-            self.draft_parallel_config = self.target_parallel_config
-
-        else:
-            self.prompt_lookup_max = 0
-            self.prompt_lookup_min = 0
-
-            if self.model is not None:
-                # Compose any callable hf_overrides set on the target so the
-                # draft config receives the same transform (e.g. the test
-                # shrink). Dict overrides stay target-only.
-                draft_hf_overrides = SpeculativeConfig.compose_draft_hf_overrides(
-                    self.target_model_config.hf_overrides
-                )
-                self.draft_model_config = ModelConfig(
-                    model=self.model,
-                    runner="draft",
-                    tokenizer=(
-                        self.model
-                        if self.use_heterogeneous_vocab
-                        else self.target_model_config.tokenizer
-                    ),
-                    tokenizer_mode=self.target_model_config.tokenizer_mode,
-                    trust_remote_code=self.target_model_config.trust_remote_code,
-                    allowed_local_media_path=self.target_model_config.allowed_local_media_path,
-                    allowed_media_domains=self.target_model_config.allowed_media_domains,
-                    dtype=self.target_model_config.dtype,
-                    seed=self.target_model_config.seed,
-                    revision=self.revision,
-                    code_revision=self.code_revision,
-                    tokenizer_revision=self.target_model_config.tokenizer_revision,
-                    max_model_len=self.max_model_len,  # type: ignore[arg-type]
-                    spec_target_max_model_len=self.target_model_config.max_model_len,
-                    quantization=self.quantization,
-                    enforce_eager=self.target_model_config.enforce_eager,
-                    max_logprobs=self.target_model_config.max_logprobs,
-                    hf_overrides=draft_hf_overrides,
-                    config_format=self.target_model_config.config_format,
-                )
-
-                # Automatically detect the method
-                if self.method in (
-                    "eagle",
-                    "eagle3",
-                    "dflash",
-                    "dspark",
-                ):
-                    pass
-                # examples:
-                # yuhuili/EAGLE-LLaMA3-Instruct-8B
-                # yuhuili/EAGLE3-LLaMA3.1-Instruct-8B
-                # AngelSlim/Qwen3-8B_eagle3
-                # deepseek-ai/dspark_qwen3_8b_block7
-                elif "eagle-" in self.draft_model_config.model.lower():
-                    self.method = "eagle"
-                elif "eagle3" in self.draft_model_config.model.lower():
-                    self.method = "eagle3"
-                elif (
-                    "dflash" in self.draft_model_config.model.lower()
-                    or "MuseGlimmerAssistantModel"
-                    in self.draft_model_config.architectures
-                ):
-                    self.method = "dflash"
-                elif (
-                    "dspark" in self.draft_model_config.model.lower()
-                    or "Qwen3DSparkModel" in self.draft_model_config.architectures
-                    or _QWEN3_OMNI_DSPARK_ARCHITECTURE
-                    in self.draft_model_config.architectures
-                    or _MUSE_SPARK_DSPARK_ARCHITECTURE
-                    in self.draft_model_config.architectures
-                    or "Gemma4DSparkModel" in self.draft_model_config.architectures
-                    or (
-                        "DSparkDraftModel" in self.draft_model_config.architectures
-                        and self.draft_model_config.hf_config.model_type == "qwen3"
-                    )
-                ):
-                    self.method = "dspark"
-                elif self.draft_model_config.hf_config.model_type in get_args(
-                    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
-                        not in ("step3p5_mtp", "step5_mtp", "inkling_mtp")
-                    ):
-                        logger.warning(
-                            "Enabling num_speculative_tokens > 1 will run "
-                            "multiple times of forward on same MTP layer"
-                            ",which may result in lower acceptance rate"
-                        )
-                elif self.method == "draft_model":
-                    pass
-                else:
-                    raise NotImplementedError(
-                        f"Unsupported speculative method: '{self.method}'"
-                    )
-
-                if self.method != "draft_model":
-                    # Every draft but an independent model shares the target's
-                    # positional space; a draft checkpoint with a smaller
-                    # max_position_embeddings than the target under-sizes its
-                    # rotary cache (#48894, #57941).
-                    SpeculativeConfig._maybe_override_draft_max_position_embeddings(
-                        self.draft_model_config.hf_config,
-                        self.target_model_config.max_model_len,
-                    )
-
-                # Replace hf_config for EAGLE draft_model
-                if self.method in ("eagle", "eagle3", "dflash"):
-                    from vllm.transformers_utils.configs.eagle import EAGLEConfig
-                    from vllm.transformers_utils.configs.speculators import (
-                        SpeculatorsConfig,
-                    )
-
-                    if isinstance(
-                        self.draft_model_config.hf_config,
-                        (EAGLEConfig, SpeculatorsConfig),
-                    ):
-                        pass
-                    else:
-                        eagle_config = EAGLEConfig(
-                            self.draft_model_config.hf_config,
-                            method=self.method,
-                            model_type="eagle",
-                        )
-                        self.draft_model_config.hf_config = eagle_config
-                        self.update_arch_()
-
-                if (
-                    self.method == "dspark"
-                    and "DSparkDraftModel" in self.draft_model_config.architectures
-                    and self.draft_model_config.hf_config.model_type == "qwen3"
-                ):
-                    self.draft_model_config.hf_config.architectures = [
-                        "Qwen3DSparkModel"
-                    ]
-                    self.update_arch_()
-                elif self.method == "dspark" and (
-                    "Qwen3DSparkModel" not in self.draft_model_config.architectures
-                    and _QWEN3_OMNI_DSPARK_ARCHITECTURE
-                    not in self.draft_model_config.architectures
-                    and _MUSE_SPARK_DSPARK_ARCHITECTURE
-                    not in self.draft_model_config.architectures
-                    and "Gemma4DSparkModel" not in self.draft_model_config.architectures
-                    and "K3DSparkModel" not in self.draft_model_config.architectures
-                ):
-                    # DeepSeek-V4(.1) DSpark reuses the full target config
-                    # and its weights ship in the target checkpoint.
-                    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
-                    )
-                    self.update_arch_()
-                elif (
-                    self.method == "dspark"
-                    and "Gemma4DSparkModel" in self.draft_model_config.architectures
-                ):
-                    # Normalize the self-contained Gemma4 draft's config keys to
-                    # the DSpark conventions.
-                    hf = self.draft_model_config.hf_config
-                    if (
-                        getattr(hf, "dspark_target_layer_ids", None) is None
-                        and getattr(hf, "target_layer_ids", None) is not None
-                    ):
-                        hf.dspark_target_layer_ids = hf.target_layer_ids
-
-                if self.method in ("dflash", "dspark"):
-                    self.parallel_drafting = True
-
-                n_predict = getattr(
-                    self.draft_model_config.hf_config, "n_predict", None
-                )
-                if (
-                    "ParseAuxMTPModel" in self.draft_model_config.architectures
-                    and self.num_speculative_tokens not in (None, 1)
-                ):
-                    raise ValueError("Parse 3.5 has only one trained MTP step")
-                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
-                    elif (
-                        self.num_speculative_tokens > n_predict
-                        and self.num_speculative_tokens % n_predict != 0
-                    ):
-                        # Ensure divisibility for MTP module reuse.
-                        raise ValueError(
-                            f"num_speculative_tokens:{self.num_speculative_tokens}"
-                            f" must be divisible by {n_predict=}"
-                        )
-
-                if self.num_speculative_tokens is None:
-                    raise ValueError(
-                        "A speculative model was provided, but "
-                        "`num_speculative_tokens` was not provided"
-                    )
-
-                if self.xpress_topc is not None and not self.is_xpress():
-                    raise ValueError("xpress_topc is only supported by XPress")
-
-                if self.is_xpress():
-                    hf_config = self.draft_model_config.hf_config
-                    xpress_topc = self.xpress_topc
-                    if xpress_topc is None:
-                        xpress_topc = getattr(hf_config, "xpress_topc", None)
-                    if xpress_topc is not None:
-                        draft_vocab_size = (
-                            getattr(hf_config, "draft_vocab_size", None)
-                            or hf_config.vocab_size
-                        )
-                        if not 0 <= xpress_topc <= draft_vocab_size:
-                            raise ValueError(
-                                "xpress_topc must be between 0 and the draft "
-                                f"vocabulary size ({draft_vocab_size})"
-                            )
-                        hf_config.xpress_topc = xpress_topc
-
-                if self.dspark_draft_topk is not None and self.method != "dspark":
-                    raise ValueError("dspark_draft_topk is only supported by DSpark")
-
-                dspark_draft_topk = None
-                if self.method == "dspark":
-                    hf_config = self.draft_model_config.hf_config
-                    dspark_draft_topk = self.dspark_draft_topk
-                    if dspark_draft_topk is None:
-                        dspark_draft_topk = getattr(
-                            hf_config, "dspark_draft_topk", None
-                        )
-                    if dspark_draft_topk is not None:
-                        draft_vocab_size = (
-                            getattr(hf_config, "draft_vocab_size", None)
-                            or hf_config.vocab_size
-                        )
-                        if not 1 <= dspark_draft_topk <= draft_vocab_size:
-                            raise ValueError(
-                                "dspark_draft_topk must be between 1 and the "
-                                f"draft vocabulary size ({draft_vocab_size})"
-                            )
-                        if (
-                            "Qwen3DSparkModel"
-                            not in self.draft_model_config.architectures
-                            and _QWEN3_OMNI_DSPARK_ARCHITECTURE
-                            not in self.draft_model_config.architectures
-                        ):
-                            raise ValueError(
-                                "dspark_draft_topk is only supported by "
-                                "Qwen3DSparkModel and Qwen3OmniDSparkModel"
-                            )
-                        hf_config.dspark_draft_topk = dspark_draft_topk
-
-                    assert self.target_model_config is not None
-                    _validate_qwen3_omni_dspark(
-                        self.target_model_config,
-                        self.draft_model_config,
-                        self.num_speculative_tokens,
-                    )
-
-                self.draft_tensor_parallel_size = (
-                    SpeculativeConfig._verify_and_get_draft_tp(
-                        self.target_parallel_config,
-                        self.draft_tensor_parallel_size,
-                    )
-                )
-                self.draft_model_config.max_model_len = (
-                    SpeculativeConfig._maybe_override_draft_max_model_len(
-                        self.max_model_len,
-                        self.draft_model_config.max_model_len,
-                        self.target_model_config.max_model_len,
-                    )
-                )
-
-                # 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,
-                        draft_model_config=self.draft_model_config,
-                    )
-                )
-
-        if self.index_share_for_mtp_iteration is not None:
-            if self.method != "mtp" or self.draft_model_config is None:
                 raise ValueError(
-                    "index_share_for_mtp_iteration is only supported with method='mtp'"
+                    f"{option} is not supported with speculative method "
+                    f"{self.custom_method or self.method!r}."
                 )
-            self.draft_model_config.hf_config.index_share_for_mtp_iteration = (
-                self.index_share_for_mtp_iteration
-            )
+        config_cls.apply_options(self)
 
         if (
-            self.method != "dspark"
+            config_cls.adaptive_verification_uses_draft_logits
             and self.enable_adaptive_verification
             and self.use_local_argmax_reduction
         ):
@@ -1685,7 +1397,31 @@ def __post_init__(self):
 
         return self
 
+    @property
+    def speculator_config(self) -> SpeculatorConfig:
+        """The configured speculator's config. Derived on each access
+        so that it follows changes to the fields it depends."""
+        return get_speculator_config_class(self)(self)
+
+    # Shortcuts for commonly used speculator config properties.
+
+    @property
+    def uses_target_hidden_states(self) -> bool:
+        """See `SpeculatorConfig.uses_target_hidden_states`."""
+        return self.speculator_config.uses_target_hidden_states
+
+    @property
+    def uses_aux_hidden_states(self) -> bool:
+        """See `SpeculatorConfig.uses_aux_hidden_states`."""
+        return self.speculator_config.uses_aux_hidden_states
+
+    @property
+    def uses_separate_workspace_lane(self) -> bool:
+        """See `SpeculatorConfig.uses_separate_workspace_lane`."""
+        return self.speculator_config.uses_separate_workspace_lane
+
     def _validate_suffix_decoding(self):
+        # V1 model runner only, to be removed with it.
         if not has_arctic_inference():
             raise ImportError(
                 "Arctic Inference is required for suffix decoding. "
@@ -1770,7 +1506,9 @@ def _maybe_override_draft_max_position_embeddings(
         draft_hf_config: PreTrainedConfig,
         target_max_model_len: int,
     ) -> None:
-        """Raise a draft's max_position_embeddings up to the target's.
+        """Raise the max_position_embeddings of a draft that shares the
+        target's positional space (see `DraftModelSpeculatorConfig`'s
+        `shares_target_positions`) up to the target's.
 
         The proposer feeds the draft positions up to the target's
         max_model_len, while max_position_embeddings sizes the draft's
@@ -1888,6 +1626,19 @@ def apply_draft_overrides(self, vllm_config: VllmConfig) -> VllmConfig:
                 vllm_config = replace(vllm_config, **{config_name: sub_config})
         return vllm_config
 
+    @field_validator("method")
+    @classmethod
+    def _validate_method(cls, value: str | None) -> str | None:
+        if value is None or value in get_args(SpeculativeMethod):
+            return value
+        if value not in (registered := registered_speculative_methods()):
+            raise ValueError(
+                f"Unknown speculative method {value!r}. Built-in methods: "
+                f"{sorted(get_args(SpeculativeMethod))}; registered custom "
+                f"methods: {registered}."
+            )
+        return value
+
     @field_validator("attention_backend", mode="before")
     @classmethod
     def _parse_attention_backend(cls, value: Any) -> Any:
@@ -1912,20 +1663,6 @@ def _verify_args(self) -> Self:
                 "n_predict parameter."
             )
 
-        if self.draft_lm_head_quantization is not None:
-            # Gemma-4 MTP scores tied embedding rows directly; no lm_head to share.
-            if self.method not in ("eagle", "eagle3", "mtp") or self.use_gemma4_mtp():
-                raise ValueError(
-                    "draft_lm_head_quantization is not supported for method "
-                    f"{self.method!r}."
-                )
-            draft = self.draft_model_config
-            if draft is not None and draft.head_dtype != draft.dtype:
-                raise ValueError(
-                    "draft_lm_head_quantization needs the lm_head to run in the "
-                    f"model dtype, got head_dtype={draft.head_dtype}."
-                )
-
         if self.num_speculative_tokens <= 0:
             raise ValueError(
                 "Expected num_speculative_tokens to be greater "
@@ -1974,7 +1711,7 @@ def _verify_args(self) -> Self:
             )
 
         if self.rejection_sample_method == "fly":
-            if self.uses_extract_hidden_states():
+            if not self.speculator_config.proposes_draft_tokens:
                 raise ValueError(
                     "FLy is not compatible with method='extract_hidden_states', "
                     "which does not perform speculative decoding."
@@ -2001,7 +1738,7 @@ def _verify_args(self) -> Self:
 
     def verify_equal_vocab_size_if_draft_model(self):
         if (
-            self.method == "draft_model"
+            self.speculator_config.draft_has_own_vocab
             and self.target_model_config is not None
             and self.draft_model_config is not None
         ):
@@ -2039,25 +1776,7 @@ def max_num_new_slots_for_drafting(self) -> int:
         DFlash and DSpark draft in a separate query pass rather than adding
         rows to the target batch.
         """
-        num_draft_tokens = self.num_speculative_tokens
-
-        if self.use_dflash() or self.use_dspark():
-            return 0
-
-        if self.parallel_drafting:
-            if self.uses_draft_model():
-                # PARD does not shift the existing input, so all K query
-                # positions require additional slots.
-                return num_draft_tokens
-
-            # The existing query is reused; only masked queries need new slots.
-            return num_draft_tokens - 1
-
-        if self.uses_draft_model():
-            # The autoregressive draft-model input retains one unsliced token.
-            return 1
-
-        return 0
+        return self.speculator_config.max_num_new_slots_for_drafting
 
     def use_gemma4_mtp(self) -> bool:
         return (
@@ -2090,35 +1809,10 @@ def needs_eagle_block_drop(self) -> bool:
         """
         return self.use_eagle() and not (self.use_dflash() or self.use_dspark())
 
-    def is_xpress(self) -> bool:
-        """Whether the dflash draft carries an XPress refiner.
-
-        XPress serves under method="dflash": it is a refiner on top of the same
-        drafter, and the inheritance mirrors that (XPressSpeculator subclasses
-        DFlashSpeculator).
-        """
-        return (
-            self.method == "dflash"
-            and self.draft_model_config is not None
-            and any(
-                arch in XPRESS_ARCHITECTURES
-                for arch in self.draft_model_config.architectures
-            )
-        )
-
     def use_eagle_block_drop(self) -> bool:
         """Whether the volatile trailing prefix-cache block should be dropped."""
         return self.needs_eagle_block_drop() and not self.disable_eagle_block_drop
 
-    def draft_shares_target_positions(self) -> bool:
-        """Whether the draft's max_position_embeddings is raised to the
-        target's max_model_len, so its length doesn't limit the target's."""
-        return (
-            self.method != "draft_model"
-            and self.draft_model_config is not None
-            and _rope_only_sizes_cache(self.draft_model_config.hf_config)
-        )
-
     def use_dflash(self) -> bool:
         return self.method == "dflash"
 
@@ -2147,19 +1841,32 @@ def use_ngram_gpu(self) -> bool:
     def use_ngram(self) -> bool:
         return self.method in ("ngram", "ngram_gpu")
 
+    def num_mtp_modules(self) -> int:
+        """Number of MTP modules the drafter selects among by draft step."""
+        assert self.draft_model_config is not None
+        return getattr(self.draft_model_config.hf_config, "n_predict", None) or 1
+
     def use_multi_module_mtp(self) -> bool:
         if self.method != "mtp" or self.draft_model_config is None:
             return False
-        num_mtp_layers = getattr(
-            self.draft_model_config.hf_config, "num_nextn_predict_layers", 1
-        )
-        return min(num_mtp_layers, self.num_speculative_tokens) > 1
+        hf_config = self.draft_model_config.hf_config
+        if min(self.num_mtp_modules(), self.num_speculative_tokens) <= 1:
+            return False
+        index_share = self.index_share_for_mtp_iteration
+        if index_share is None:
+            index_share = getattr(hf_config, "index_share_for_mtp_iteration", False)
+        if index_share:
+            # TODO: Support index_share_for_mtp_iteration in
+            # MultiModuleMTPSpeculator.
+            return False
+        # TODO: Verify MiniMax-M3 with MultiModuleMTPSpeculator.
+        return hf_config.model_type != "minimax_m3_mtp"
 
     def __repr__(self) -> str:
-        method = self.method
+        method = self.custom_method or self.method
         model = (
             None
-            if method
+            if self.method
             in (
                 "ngram",
                 "suffix",
diff --git a/vllm/config/speech_to_text.py b/vllm/config/speech_to_text.py
index 6d31713c8667..cc56f91b0086 100644
--- a/vllm/config/speech_to_text.py
+++ b/vllm/config/speech_to_text.py
@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Literal
 
 from vllm.config.utils import config
 
@@ -50,6 +50,9 @@ class SpeechToTextParams:
     to_language: str | None = None
     """Target language for translation (model-dependent)."""
 
+    extra_args: dict[str, str | int | float | list[str | int | float]] | None = None
+    """Request extensions supplied through vllm_xargs."""
+
 
 @config
 class SpeechToTextConfig:
@@ -60,6 +63,10 @@ class SpeechToTextConfig:
     16kHz audio input. The input audio will be automatically resampled to this
     rate before processing."""
 
+    audio_resample_method: Literal["pyav", "scipy", "soxr", "torchaudio"] | None = None
+    """Optional resampler applied after decoding at the native sample rate.
+    None preserves the audio decoder's default resampling behavior."""
+
     max_audio_clip_s: int | None = 30
     """Maximum duration in seconds for a single audio clip without chunking.
     Audio longer than this will be split into smaller chunks if
diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py
index 078da543555d..7f381384dddc 100644
--- a/vllm/config/vllm.py
+++ b/vllm/config/vllm.py
@@ -73,6 +73,17 @@
 
 logger = init_logger(__name__)
 
+_QWEN3_5_ARCHITECTURES = frozenset(
+    {
+        "Qwen3_5ForCausalLM",
+        "Qwen3_5ForConditionalGeneration",
+        "Qwen3_5MTP",
+        "Qwen3_5MoeForCausalLM",
+        "Qwen3_5MoeForConditionalGeneration",
+        "Qwen3_5MoeMTP",
+    }
+)
+
 DEFAULT_BREAKABLE_CUDAGRAPH_ARCHITECTURES = frozenset(
     {
         "AliceAIForCausalLM",
@@ -126,6 +137,10 @@ def default_breakable_cudagraph_architectures() -> frozenset[str]:
         # graphs are on. Enable this architecture so the published AMD
         # recipe can start.
         return frozenset({"DeepseekV41ForCausalLM"})
+    if current_platform.is_cuda():
+        # Qwen3.5 migrates to breakable CUDA graphs on NVIDIA only; other
+        # platforms keep the torch.compile path.
+        return DEFAULT_BREAKABLE_CUDAGRAPH_ARCHITECTURES | _QWEN3_5_ARCHITECTURES
     return DEFAULT_BREAKABLE_CUDAGRAPH_ARCHITECTURES
 
 
@@ -636,20 +651,7 @@ def num_lookahead_tokens(self) -> int:
         speculative_config = self.speculative_config
         if speculative_config is None:
             return 0
-        dspark_fill_in = speculative_config.use_dspark() and not getattr(
-            speculative_config.draft_model_config.hf_config, "sample_from_anchor", True
-        )
-        if speculative_config.use_dflash() or dspark_fill_in:
-            # Fill-in drafting uses a bonus query plus one query per draft token.
-            # DSpark's anchor-sampling layout does not need the extra slot.
-            return self.num_speculative_tokens + 1
-        if speculative_config.use_eagle() or speculative_config.uses_draft_model():
-            # Anchor-sampling DSpark drafts a block of num_speculative_tokens
-            # query tokens in which the anchor itself is the first prediction
-            # position (no separate bonus query), so it needs exactly
-            # num_speculative_tokens lookahead slots.
-            return self.num_speculative_tokens
-        return 0
+        return speculative_config.speculator_config.num_lookahead_tokens
 
     @property
     def num_prefill_lookahead_tokens(self) -> int:
@@ -664,15 +666,9 @@ def num_prefill_lookahead_tokens(self) -> int:
         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():
+        if speculative_config is None:
             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
+        return speculative_config.speculator_config.num_prefill_lookahead_tokens
 
     @property
     def num_hash_lookahead_tokens(self) -> int:
@@ -729,33 +725,22 @@ def use_cumem_cudagraph_pool(self) -> bool:
 
     @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(
-                    "Watermarking requires Model Runner V2 and overrides "
-                    "VLLM_USE_V2_MODEL_RUNNER=0."
-                )
-            return True
-
         use_v2_model_runner = envs.VLLM_USE_V2_MODEL_RUNNER
         if use_v2_model_runner is not None:
             return use_v2_model_runner
 
-        if not HAS_TRITON and not current_platform.is_unspecified():
+        unsupported = self._get_v2_model_runner_unsupported_features()
+
+        if current_platform.uses_standalone_cpu_model_runner():
+            unsupported.extend(
+                self._get_standalone_cpu_model_runner_unsupported_features()
+            )
+        elif not HAS_TRITON and not current_platform.is_unspecified():
             logger.warning_once(
                 "Model Runner V2 requires Triton; using the V1 model runner instead."
             )
             return False
 
-        unsupported = self._get_v2_model_runner_unsupported_features()
         if unsupported:
             logger.warning_once(
                 "Model Runner V2 does not yet support %s; using the V1 model "
@@ -1243,6 +1228,55 @@ def _verify_aux_output_compatibility(self) -> None:
                         f"{connector_name}; PD auxiliary output is not supported."
                     )
 
+    def _verify_direct_reload_compat(self) -> None:
+        """Reject configurations that direct weight reload cannot support."""
+        # not `direct_reload_enabled()`: importing it would load the model
+        # loaders into every process that builds a config
+        if envs.VLLM_RELOAD_MODE != "direct":
+            return
+        from vllm.model_executor.layers.quantization.online.base import (
+            OnlineQuantizationConfig,
+        )
+
+        offload = self.offload_config
+        if (
+            offload.offload_backend != "auto"
+            or offload.prefetch.offload_group_size > 0
+            or offload.uva.cpu_offload_gb > 0
+        ):
+            # CPU/UVA offload gives parameters new storage on every post-load
+            # processing; prefetch offload shares one buffer between layers
+            raise ValueError("Direct reload does not support weight offloading.")
+        model_config = self.model_config
+        if model_config is not None and model_config.quantization == "torchao":
+            # its load_weights runs a layerwise reload, and its post-load
+            # processing replaces the weight parameter
+            raise ValueError("Direct reload does not support torchao quantization.")
+        if isinstance(self.quant_config, OnlineQuantizationConfig) or (
+            getattr(self.quant_config, "online_quantization_config", None) is not None
+        ):
+            # `finish` clears the guard that keeps online quantization from
+            # quantizing its already-quantized bytes again
+            raise ValueError("Direct reload does not support online quantization.")
+        if self.lora_config is not None:
+            # LoRA wraps layers after cold start; their wrappers forward
+            # `quant_method`, so a rerun would process each base layer twice
+            raise ValueError("Direct reload does not support LoRA.")
+        spec = self.speculative_config
+        if spec is not None and spec.draft_lm_head_quantization is not None:
+            # after `begin` the draft's quantized lm_head is a view of the target
+            # lm_head's storage: post-load processing would write it there
+            raise ValueError(
+                "Direct reload does not support draft_lm_head_quantization."
+            )
+        transfer = self.weight_transfer_config
+        if transfer is not None and transfer.backend == "sharded_rdt":
+            # its replay plan is baked by a layerwise (meta) dry run
+            raise ValueError(
+                "Direct reload does not support the sharded_rdt weight-transfer "
+                "backend. Use VLLM_RELOAD_MODE=layerwise or another backend."
+            )
+
     def _verify_kv_transfer_compat(self) -> None:
         """Reject configurations that silently corrupt KV transfers."""
         if (
@@ -1374,20 +1408,10 @@ def _check_supports_watermarking(
                 raise ValueError(
                     "Parallel speculative drafting is not supported with watermarking."
                 )
-            if speculative_config.method not in (
-                "dflash",
-                "dspark",
-                "eagle",
-                "eagle3",
-                "mtp",
-            ):
+            if not speculative_config.speculator_config.supports_watermarking:
                 raise ValueError(
                     "Watermarking supports only dflash, dspark, eagle, eagle3 and "
-                    "mtp speculative decoding."
-                )
-            if self._is_dflash_candidate_draft():
-                raise ValueError(
-                    "Watermarking does not support DFlash2 and LiLiCorr "
+                    "mtp speculative decoding, and not DFlash2 and LiLiCorr "
                     "candidate-head drafters yet."
                 )
             if (
@@ -1549,6 +1573,11 @@ def __post_init__(self):
         self.parallel_config.set_dcp_defaults()
 
         if self.model_config is not None:
+            # Normalize EPLB num_redundant_experts before verification.
+            self.parallel_config.compute_eplb_num_redundant_experts(
+                self.model_config.get_num_experts()
+            )
+
             self.model_config.verify_with_parallel_config(self.parallel_config)
             self.model_config.verify_dual_chunk_attention_config(self.load_config)
 
@@ -1609,6 +1638,7 @@ def __post_init__(self):
             self.quant_config = VllmConfig._get_quantization_config(
                 self.model_config, self.load_config
             )
+        self._verify_direct_reload_compat()
 
         # "dummy" reads no weights at all, and the sharded formats read a vLLM
         # state dict, which stores tied word embeddings under the lm_head only.
@@ -1732,6 +1762,19 @@ def __post_init__(self):
             else:
                 self.parallel_config.disable_nccl_for_dp_synchronization = False
 
+        if (
+            self.parallel_config.enable_fault_tolerance
+            and not self.parallel_config.disable_nccl_for_dp_synchronization
+        ):
+            # FT only supports DP allreduce over the CPU group: recovery
+            # rebuilds only the cpu group, so a device-group (NCCL) sync
+            # would hang after reinit.
+            logger.info_once(
+                "Disabling NCCL for DP synchronization because "
+                "fault tolerance is enabled.",
+            )
+            self.parallel_config.disable_nccl_for_dp_synchronization = True
+
         if (
             self.observability_config.per_request_spec_decode_metrics != "none"
             and self.speculative_config is None
@@ -1882,8 +1925,8 @@ def has_blocked_weights():
             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 current_platform.get_hisparse_backend_cls() is None:
+                raise ValueError("The current platform does not support HiSparse.")
             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:
@@ -2019,7 +2062,7 @@ def has_blocked_weights():
         if self.cache_config.kv_sharing_fast_prefill:
             if (
                 self.speculative_config is not None
-                and self.speculative_config.use_eagle()
+                and self.speculative_config.uses_target_hidden_states
             ):
                 raise ValueError(
                     "Fast prefill optimization for KV sharing is not "
@@ -2095,6 +2138,7 @@ def has_blocked_weights():
             self._validate_v2_model_runner()
         else:
             self._validate_v1_model_runner()
+        self._resolve_extensible_kv_cache()
 
         self._validate_profiler_config()
         self._validate_batch_sharded_sampling()
@@ -2204,6 +2248,8 @@ def has_blocked_weights():
                 " and install the matching kernels."
             )
 
+            self._raise_dbo_thresholds_for_spec_decode()
+
         if not self.instance_id:
             self.instance_id = random_uuid()[:5]
 
@@ -2239,7 +2285,7 @@ def has_blocked_weights():
         ):
             if (
                 self.speculative_config is not None
-                and self.speculative_config.use_eagle()
+                and self.speculative_config.uses_target_hidden_states
             ):
                 # Hybrid KV cache manager is not yet supported with chunked
                 # local attention + eagle.
@@ -2377,19 +2423,21 @@ def _set_max_num_scheduled_tokens(self):
                     f" num_lookahead_tokens={self.num_lookahead_tokens}."
                 )
             spec = self.speculative_config
-            if spec.use_dflash() or spec.use_dspark():
+            max_draft_seqs = spec.speculator_config.max_num_draft_seqs(
+                max_num_batched_tokens
+            )
+            if max_draft_seqs is not None:
                 max_num_active_seqs = (
                     self.scheduler_config.max_num_active_seqs
                     or self.scheduler_config.max_num_seqs
                 )
-                max_draft_seqs = max_num_batched_tokens // self.num_lookahead_tokens
                 if max_draft_seqs < max_num_active_seqs:
                     logger.warning_once(
-                        "The %s draft pass needs %d query tokens per request within "
+                        "The %s draft pass fits at most %d requests within "
                         "max_num_batched_tokens=%d. Overriding max_num_active_seqs "
                         "from %d to %d.",
                         spec.method,
-                        self.num_lookahead_tokens,
+                        max_draft_seqs,
                         max_num_batched_tokens,
                         max_num_active_seqs,
                         max_draft_seqs,
@@ -3030,16 +3078,13 @@ def _limit_max_model_len_to_draft(self) -> None:
         spec_config = self.speculative_config
         if (
             spec_config is None
-            or spec_config.draft_model_config is None
             or self.model_config is None
             or not self.use_v2_model_runner
-            # The draft's positional capacity was raised to match the target.
-            or spec_config.draft_shares_target_positions()
         ):
             return
-        draft_max_model_len = spec_config.draft_model_config.max_model_len
+        draft_max_model_len = spec_config.speculator_config.max_model_len
         target_max_model_len = self.model_config.max_model_len
-        if draft_max_model_len >= target_max_model_len:
+        if draft_max_model_len is None or draft_max_model_len >= target_max_model_len:
             return
         logger.warning(
             "The draft model's max_model_len (%d) is smaller than the target "
@@ -3065,7 +3110,7 @@ def _get_v2_model_runner_unsupported_features(self) -> list[str]:
             unsupported.append("pipeline parallelism with external_launcher")
 
         if speculative_config is not None:
-            if speculative_config.method in ("suffix", "custom_class"):
+            if speculative_config.speculator_config.speculator_class is None:
                 unsupported.append(f"speculative method '{speculative_config.method}'")
 
             # V2 EagleSpeculator does not support parallel_drafting (for P-Eagle).
@@ -3087,8 +3132,122 @@ def _get_v2_model_runner_unsupported_features(self) -> list[str]:
 
         return unsupported
 
+    def _get_standalone_cpu_model_runner_unsupported_features(self) -> list[str]:
+        """Collect features not supported by the standalone CPU model runner."""
+        unsupported: list[str] = []
+
+        if self.is_mm_encoder_only:
+            unsupported.append("multimodal encoder-only instances")
+
+        if self.speculative_config is not None:
+            unsupported.append("speculative decoding")
+
+        if self.model_config is not None:
+            if self.model_config.enable_prompt_embeds:
+                unsupported.append("prompt embeddings")
+            if self.model_config.is_diffusion:
+                unsupported.append("diffusion models")
+            if self.model_config.logits_processors:
+                unsupported.append("custom logits processors")
+
+        if self.parallel_config.decode_context_parallel_size > 1:
+            unsupported.append("decode context parallelism")
+
+        if self.parallel_config.prefill_context_parallel_size > 1:
+            unsupported.append("prefill context parallelism")
+
+        if self.parallel_config.enable_batch_sharded_sampling:
+            unsupported.append("batch-sharded sampling")
+
+        if self.parallel_config.use_ubatching:
+            unsupported.append("microbatching")
+
+        if self.parallel_config.enable_eplb:
+            unsupported.append("expert parallel load balancing")
+
+        if self.aux_output_config.enabled:
+            unsupported.append("routed-expert auxiliary outputs")
+
+        if self.lora_config and self.lora_config.enable_tower_connector_lora:
+            unsupported.append("multimodal tower and connector LoRA")
+
+        return unsupported
+
+    def _resolve_extensible_kv_cache(self) -> None:
+        """Settle `enable_extensible_kv_cache`: on where supported if unset, an
+        error if explicitly requested where unsupported.
+
+        Runs once per model: configs derived for sub-models (`with_hf_config`)
+        may pick a different runner, and must not revisit the decision.
+        """
+        from vllm.platforms import current_platform
+
+        cache_config = self.cache_config
+        if cache_config._extensible_kv_cache_resolved:
+            return
+        cache_config._extensible_kv_cache_resolved = True
+        requested = cache_config.enable_extensible_kv_cache
+        unsupported = None
+        if not current_platform.is_cuda_alike():
+            unsupported = "it is only supported on CUDA and ROCm"
+        elif cache_config.kv_cache_memory_bytes is not None:
+            unsupported = (
+                "it sizes the KV cache from measured memory, which "
+                "kv_cache_memory_bytes overrides"
+            )
+        elif cache_config.num_gpu_blocks_override is not None:
+            unsupported = (
+                "it sizes the KV cache from measured memory, which "
+                "num_gpu_blocks_override overrides"
+            )
+        elif not self.use_v2_model_runner:
+            unsupported = "it requires the V2 model runner"
+        elif self.attention_config.hisparse_config is not None:
+            unsupported = "HiSparse allocates its KV cache itself"
+        elif self.parallel_config.enable_elastic_ep:
+            unsupported = (
+                "elastic EP scale-up skips warmup and reuses the profiled KV "
+                "cache size, so no measurement would size the cache"
+            )
+        elif (
+            self.kv_transfer_config is not None
+            and self.kv_transfer_config.has_connector("LMCacheMPConnector")
+        ):
+            unsupported = (
+                "LMCacheMPConnector shares the KV cache through legacy CUDA IPC, "
+                "which cannot export driver-mapped memory"
+            )
+        elif (
+            self.kv_transfer_config is not None
+            and self.kv_transfer_config.uses_custom_mem_pool()
+        ):
+            unsupported = (
+                "custom_mem_pool allocates the KV cache from a connector's "
+                "memory pool, which a driver-mapped cache cannot use"
+            )
+        elif (
+            self.kv_transfer_config is not None
+            and current_platform.is_rocm()
+            and self.kv_transfer_config.registers_kv_cache_memory()
+        ):
+            # TODO: drop once UCX handles HSA_EXT_POINTER_TYPE_HSA_VMEM and
+            # hsa_amd_portable_export_dmabuf accepts VMM ranges.
+            unsupported = (
+                "on ROCm, UCX and RDMA cannot yet register driver-mapped memory, "
+                "which the NIXL, Mooncake and MoRIIO connectors require"
+            )
+        if requested is None:
+            cache_config.enable_extensible_kv_cache = unsupported is None
+        elif requested and unsupported is not None:
+            raise ValueError(f"enable_extensible_kv_cache: {unsupported}.")
+
     def _get_v1_model_runner_unsupported_features(self) -> list[str]:
         unsupported: list[str] = []
+        if self.attention_config.hisparse_config is not None:
+            unsupported.append("HiSparse")
+
+        if self.watermark_config is not None:
+            unsupported.append("watermarking")
 
         # PCP runtime support is implemented only by the V2 model runner.
         if self.parallel_config.prefill_context_parallel_size > 1:
@@ -3110,9 +3269,21 @@ def _get_v1_model_runner_unsupported_features(self) -> list[str]:
 
         # DSpark is implemented only by the V2 GPU model runner.
         if self.speculative_config:
+            # Imported here since it imports vllm.config.
+            from vllm.v1.worker.gpu.spec_decode.xpress.config import (
+                XPressSpeculatorConfig,
+            )
+
+            if self.speculative_config.custom_method is not None:
+                unsupported.append(
+                    "custom speculative method "
+                    f"'{self.speculative_config.custom_method}'"
+                )
             if self.speculative_config.method == "dspark":
                 unsupported.append("dspark speculative decoding")
-            elif self.speculative_config.is_xpress():
+            elif isinstance(
+                self.speculative_config.speculator_config, XPressSpeculatorConfig
+            ):
                 unsupported.append("xpress speculative decoding")
             if self.speculative_config.enable_adaptive_verification:
                 unsupported.append("adaptive draft verification")
@@ -3220,12 +3391,35 @@ def _validate_batch_sharded_sampling(self) -> None:
                 f"{'; '.join(blockers)}."
             )
 
+    def _raise_dbo_thresholds_for_spec_decode(self) -> None:
+        """The speculative counterpart of ParallelConfig's DBO threshold check.
+
+        A uniform decode batch below one request per microbatch cannot be split
+        on a request boundary, and an idle DP rank's one-request dummy batch
+        must stay below the thresholds too. Such batches gain nothing from
+        microbatching, so raise the thresholds rather than reject the config.
+        """
+        parallel_config = self.parallel_config
+        min_ubatch_tokens = parallel_config.num_ubatches * self.uniform_decode_query_len
+        for name in ("dbo_decode_token_threshold", "dbo_prefill_token_threshold"):
+            threshold = getattr(parallel_config, name)
+            if threshold < min_ubatch_tokens:
+                logger.warning_once(
+                    "Raising %s from %d to %d, the number of microbatches times "
+                    "the decode query length.",
+                    name,
+                    threshold,
+                    min_ubatch_tokens,
+                )
+                setattr(parallel_config, name, min_ubatch_tokens)
+
     def _get_dbo_unsupported_features(self) -> list[str]:
         """Collect what the V2 model runner cannot combine with DBO.
 
-        The V2 runner microbatches a plain decoder forward pass. Anything that
-        slices or replays the batch differently (drafting, adapters, pipeline
-        stages, context parallelism, encoders) is not handled yet.
+        The V2 runner microbatches a decoder forward pass (drafting runs whole,
+        after it). Anything else that slices or replays the batch differently
+        (adapters, pipeline stages, prefill context parallelism, encoders) is
+        not handled yet.
         """
         # TODO: DBO with model runner V2 is under development.
         # It should be enabled with explicit VLLM_USE_V2_MODEL_RUNNER environ.
@@ -3239,15 +3433,10 @@ def _get_dbo_unsupported_features(self) -> list[str]:
 
         if self.lora_config is not None:
             unsupported.append("dual batch overlap with LoRA")
-        if self.speculative_config is not None:
-            unsupported.append("dual batch overlap with speculative decoding")
         if parallel_config.pipeline_parallel_size > 1:
             unsupported.append("dual batch overlap with pipeline parallelism")
-        if (
-            parallel_config.decode_context_parallel_size > 1
-            or parallel_config.prefill_context_parallel_size > 1
-        ):
-            unsupported.append("dual batch overlap with context parallelism")
+        if parallel_config.prefill_context_parallel_size > 1:
+            unsupported.append("dual batch overlap with prefill context parallelism")
         if model_config is not None and (
             model_config.is_multimodal_model or model_config.is_encoder_decoder
         ):
@@ -3329,6 +3518,13 @@ def _validate_v2_model_runner(self) -> None:
             raise ValueError(
                 f"Model Runner V2 does not yet support: {', '.join(unsupported)}"
             )
+        if current_platform.uses_standalone_cpu_model_runner():
+            unsupported = self._get_standalone_cpu_model_runner_unsupported_features()
+            if unsupported:
+                raise ValueError(
+                    "CPU Model Runner V2 without Triton does not yet support: "
+                    f"{', '.join(unsupported)}"
+                )
 
     def _validate_v1_model_runner(self) -> None:
         unsupported = self._get_v1_model_runner_unsupported_features()
diff --git a/vllm/connections.py b/vllm/connections.py
index 2913a640e0cb..12e4e4f5678d 100644
--- a/vllm/connections.py
+++ b/vllm/connections.py
@@ -204,6 +204,7 @@ def _is_retryable(exc: Exception) -> bool:
     Retryable:
       - Timeouts (aiohttp, requests, stdlib)
       - Connection-level failures (refused, reset, DNS)
+      - Response-body transfer failures (incomplete or malformed payloads)
       - Server errors (5xx) -- includes S3 503 SlowDown
     Not retryable:
       - Client errors (4xx) -- bad URL, auth, not-found
@@ -230,6 +231,12 @@ def _is_retryable(exc: Exception) -> bool:
         ),
     ):
         return True
+    # A connection closed during a body read can raise a payload error.
+    if isinstance(
+        exc,
+        (aiohttp.ClientPayloadError, requests.exceptions.ChunkedEncodingError),
+    ):
+        return True
     # aiohttp server-side disconnects
     if isinstance(exc, aiohttp.ServerDisconnectedError):
         return True
diff --git a/vllm/cute_utils/__init__.py b/vllm/cute_utils/__init__.py
index 033533a5c9a0..2cc4564efa26 100644
--- a/vllm/cute_utils/__init__.py
+++ b/vllm/cute_utils/__init__.py
@@ -60,11 +60,15 @@ def to_cta0_smem(ptr: cute.Pointer, *, loc=None, ip=None):
     )
 
 
-def simple_tma_copy(atom, src, dst, mbar=None, cache_policy=None):
+def simple_tma_copy(atom, src, dst, mbar=None, cache_policy=None, mcast_mask=None):
     """A simple helper that wraps group_modes() and tma_partition()
     NOTE: this should be called WITHOUT cute.elect_one()
     """
-    if isinstance(atom.op, cpasync.CopyBulkTensorTileG2SOp):
+    g2s_ops = (
+        cpasync.CopyBulkTensorTileG2SOp,
+        cpasync.CopyBulkTensorTileG2SMulticastOp,
+    )
+    if isinstance(atom.op, g2s_ops):
         gmem = src
         smem = dst
     elif isinstance(atom.op, cpasync.CopyBulkTensorTileS2GOp):
@@ -81,8 +85,13 @@ def simple_tma_copy(atom, src, dst, mbar=None, cache_policy=None):
         cute.group_modes(gmem, 0),
     )
 
-    if isinstance(atom.op, cpasync.CopyBulkTensorTileG2SOp):
-        cute.copy(atom, g_part, s_part, tma_bar_ptr=mbar, cache_policy=cache_policy)
+    if isinstance(atom.op, g2s_ops):
+        kwargs = {}
+        if mcast_mask is not None:
+            kwargs["mcast_mask"] = mcast_mask
+        cute.copy(
+            atom, g_part, s_part, tma_bar_ptr=mbar, cache_policy=cache_policy, **kwargs
+        )
     elif isinstance(atom.op, cpasync.CopyBulkTensorTileS2GOp):
         cute.copy(atom, s_part, g_part, cache_policy=cache_policy)
     else:
diff --git a/vllm/cute_utils/cvt.py b/vllm/cute_utils/cvt.py
index 8bb4ad95198b..ce103ee741b8 100644
--- a/vllm/cute_utils/cvt.py
+++ b/vllm/cute_utils/cvt.py
@@ -38,6 +38,12 @@ def bf16x2_to_fp32x2(data, *, loc=None, ip=None) -> tuple[Float32, Float32]:
         )
 
     elif isinstance(data, (cute.Tensor, cute.TensorSSA)):
+        # TensorSSA in, TensorSSA out (it can't be recast: stage it in rmem)
+        is_ssa = isinstance(data, cute.TensorSSA)
+        if is_ssa:
+            frag = cute.make_rmem_tensor(data.shape, data.dtype)
+            frag.store(data)
+            data = frag
         # recast to 32-bit registers if needed
         if data.element_type == BFloat16:
             data = cute.recast_tensor(data, Uint32)
@@ -48,7 +54,7 @@ def bf16x2_to_fp32x2(data, *, loc=None, ip=None) -> tuple[Float32, Float32]:
         out = cute.make_rmem_tensor(size * 2, Float32)
         for i in range(size):
             out[i * 2], out[i * 2 + 1] = bf16x2_to_fp32x2(data[i])
-        return out
+        return out.load() if is_ssa else out
 
     else:
         raise ValueError(f"Unsupported type {type(data)}")
diff --git a/vllm/device_allocator/cumem.py b/vllm/device_allocator/cumem.py
index be7aba6d6352..61aab422c37d 100644
--- a/vllm/device_allocator/cumem.py
+++ b/vllm/device_allocator/cumem.py
@@ -10,6 +10,7 @@
 # the only successful approach is to call cuda driver API in C.
 import atexit
 import gc
+import mmap
 from collections.abc import Callable, Iterator
 from contextlib import contextmanager
 from typing import Any
@@ -61,6 +62,27 @@ def unmap_and_release(allocation_handle: HandleType) -> None:
     python_unmap_and_release(*allocation_handle)
 
 
+def create_cpu_backup(size_in_bytes: int) -> torch.Tensor:
+    backup = torch.frombuffer(
+        mmap.mmap(-1, size_in_bytes, flags=mmap.MAP_PRIVATE), dtype=torch.uint8
+    )
+    if PIN_MEMORY:
+        error = libcudart.cudaHostRegister(backup.data_ptr(), size_in_bytes)
+        if error:
+            libcudart.cudaGetLastError()
+            logger.warning_once(
+                "cudaHostRegister failed (error %d); sleep mode CPU backups "
+                "fall back to pageable memory.",
+                error,
+            )
+    return backup
+
+
+def release_cpu_backup(backup: torch.Tensor) -> None:
+    if backup.is_pinned() and libcudart.cudaHostUnregister(backup.data_ptr()):
+        libcudart.cudaGetLastError()
+
+
 def get_pluggable_allocator(
     python_malloc_fn: Callable[[HandleType], None],
     python_free_func: Callable[[int], HandleType],
@@ -207,6 +229,7 @@ def _python_free_callback(self, ptr: int) -> HandleType:
         """
         data = self.pointer_to_data.pop(ptr)
         if data.cpu_backup_tensor is not None:
+            release_cpu_backup(data.cpu_backup_tensor)
             data.cpu_backup_tensor = None
         if data.is_asleep and current_platform.is_rocm():
             # On ROCm, sleep() already unmapped and released this allocation's
@@ -260,24 +283,11 @@ def sleep(self, offload_tags: tuple[str, ...] | str | None = None) -> None:
                 if requests_offload != was_offloaded:
                     has_policy_conflict = True
                 continue
-            handle = data.handle
-            total_bytes += handle[1]
-            if data.tag in offload_tags:
-                backup_bytes += handle[1]
-                size_in_bytes = handle[1]
-                cpu_backup_tensor = torch.empty(
-                    size_in_bytes,
-                    dtype=torch.uint8,
-                    device="cpu",
-                    pin_memory=PIN_MEMORY,
-                )
-                cpu_ptr = cpu_backup_tensor.data_ptr()
-                libcudart.cudaMemcpy(cpu_ptr, ptr, size_in_bytes)
-                data.cpu_backup_tensor = cpu_backup_tensor
-            try:
-                unmap_and_release(handle)
-            finally:
-                data.is_asleep = True
+            offload = data.tag in offload_tags
+            total_bytes += data.handle[1]
+            if offload:
+                backup_bytes += data.handle[1]
+            self._put_to_sleep(ptr, data, offload)
 
         logger.info(
             "CuMemAllocator: sleep freed %.2f GiB memory in total, of which "
@@ -297,6 +307,20 @@ def sleep(self, offload_tags: tuple[str, ...] | str | None = None) -> None:
         gc.collect()
         torch.accelerator.empty_cache()
 
+    def _put_to_sleep(self, ptr: int, data: AllocationData, offload: bool) -> None:
+        if offload:
+            if data.cpu_backup_tensor is None:
+                data.cpu_backup_tensor = create_cpu_backup(data.handle[1])
+            cpu_ptr = data.cpu_backup_tensor.data_ptr()
+            libcudart.cudaMemcpy(cpu_ptr, ptr, data.handle[1])
+        elif data.cpu_backup_tensor is not None:
+            release_cpu_backup(data.cpu_backup_tensor)
+            data.cpu_backup_tensor = None
+        try:
+            unmap_and_release(data.handle)
+        finally:
+            data.is_asleep = True
+
     def discard(self, tags: tuple[str, ...] | str) -> None:
         """Discard mapped allocations with the given tags without CPU backup."""
         if isinstance(tags, str):
@@ -304,7 +328,7 @@ def discard(self, tags: tuple[str, ...] | str) -> None:
 
         discarded_bytes = 0
         has_policy_conflict = False
-        for data in self.pointer_to_data.values():
+        for ptr, data in self.pointer_to_data.items():
             if data.tag not in tags:
                 continue
             if data.is_asleep:
@@ -312,8 +336,7 @@ def discard(self, tags: tuple[str, ...] | str) -> None:
                     has_policy_conflict = True
                 continue
             torch.accelerator.synchronize(data.handle[0])
-            unmap_and_release(data.handle)
-            data.is_asleep = True
+            self._put_to_sleep(ptr, data, offload=False)
             discarded_bytes += data.handle[1]
 
         logger.info(
@@ -350,14 +373,8 @@ def wake_up(self, tags: list[str] | None = None) -> None:
                 create_and_map(handle)
                 data.is_asleep = False
                 if data.cpu_backup_tensor is not None:
-                    cpu_backup_tensor = data.cpu_backup_tensor
-                    if cpu_backup_tensor is not None:
-                        size_in_bytes = (
-                            cpu_backup_tensor.numel() * cpu_backup_tensor.element_size()
-                        )
-                        cpu_ptr = cpu_backup_tensor.data_ptr()
-                        libcudart.cudaMemcpy(ptr, cpu_ptr, size_in_bytes)
-                        data.cpu_backup_tensor = None
+                    cpu_ptr = data.cpu_backup_tensor.data_ptr()
+                    libcudart.cudaMemcpy(ptr, cpu_ptr, handle[1])
                 elif current_platform.is_rocm():
                     # amdgpu <6.14 may return stale VRAM (#44972); drop with its support
                     libcudart.cudaMemset(ptr, 0, handle[1])
diff --git a/vllm/device_allocator/xpumem.py b/vllm/device_allocator/xpumem.py
index 5eb836008d68..e3981bec3fa3 100644
--- a/vllm/device_allocator/xpumem.py
+++ b/vllm/device_allocator/xpumem.py
@@ -185,33 +185,11 @@ def sleep(self, offload_tags: tuple[str, ...] | str | None = None) -> None:
                 if requests_offload != was_offloaded:
                     has_policy_conflict = True
                 continue
-            size_in_bytes = data.handle[1]
-            total_bytes += size_in_bytes
-            if data.tag not in offload_tags:
-                unmap_and_release(data.handle)
-                data.is_asleep = True
-                continue
-
-            backup_bytes += size_in_bytes
-            device, _, _, _ = data.handle
-            cpu_backup_tensor = torch.empty(
-                size_in_bytes,
-                dtype=torch.uint8,
-                device="cpu",
-                pin_memory=PIN_MEMORY,
-            )
-            cpu_ptr = cpu_backup_tensor.data_ptr()
-            _xpu_memcpy_sync(
-                cpu_ptr,
-                ptr,
-                size_in_bytes,
-                MEMCPY_DEVICE_TO_HOST,
-                device,
-            )
-            data.cpu_backup_tensor = cpu_backup_tensor
-
-            unmap_and_release(data.handle)
-            data.is_asleep = True
+            offload = data.tag in offload_tags
+            total_bytes += data.handle[1]
+            if offload:
+                backup_bytes += data.handle[1]
+            self._put_to_sleep(ptr, data, offload)
 
         logger.info(
             "XpuMemAllocator: sleep freed %.2f GiB memory in total, of which "
@@ -233,6 +211,23 @@ def sleep(self, offload_tags: tuple[str, ...] | str | None = None) -> None:
         if callable(xpu_empty_cache):
             xpu_empty_cache()
 
+    def _put_to_sleep(self, ptr: int, data: AllocationData, offload: bool) -> None:
+        device, size_in_bytes, _, _ = data.handle
+        if offload:
+            if data.cpu_backup_tensor is None:
+                data.cpu_backup_tensor = torch.empty(
+                    size_in_bytes,
+                    dtype=torch.uint8,
+                    device="cpu",
+                    pin_memory=PIN_MEMORY,
+                )
+            cpu_ptr = data.cpu_backup_tensor.data_ptr()
+            _xpu_memcpy_sync(cpu_ptr, ptr, size_in_bytes, MEMCPY_DEVICE_TO_HOST, device)
+        else:
+            data.cpu_backup_tensor = None
+        unmap_and_release(data.handle)
+        data.is_asleep = True
+
     def discard(self, tags: tuple[str, ...] | str) -> None:
         """Discard mapped allocations with the given tags without CPU backup."""
         if isinstance(tags, str):
@@ -240,7 +235,7 @@ def discard(self, tags: tuple[str, ...] | str) -> None:
 
         discarded_bytes = 0
         has_policy_conflict = False
-        for data in self.pointer_to_data.values():
+        for ptr, data in self.pointer_to_data.items():
             if data.tag not in tags:
                 continue
             if data.is_asleep:
@@ -248,8 +243,7 @@ def discard(self, tags: tuple[str, ...] | str) -> None:
                     has_policy_conflict = True
                 continue
             torch.accelerator.synchronize(data.handle[0])
-            unmap_and_release(data.handle)
-            data.is_asleep = True
+            self._put_to_sleep(ptr, data, offload=False)
             discarded_bytes += data.handle[1]
 
         logger.info(
@@ -289,7 +283,6 @@ def wake_up(self, tags: list[str] | None = None) -> None:
                 MEMCPY_HOST_TO_DEVICE,
                 device,
             )
-            data.cpu_backup_tensor = None
 
     def release_pools(self) -> None:
         """Release pools, retaining allocators until their allocations are freed."""
diff --git a/vllm/distributed/aux_output_connector/worker.py b/vllm/distributed/aux_output_connector/worker.py
index 6044a61cccc3..6c053d17304c 100644
--- a/vllm/distributed/aux_output_connector/worker.py
+++ b/vllm/distributed/aux_output_connector/worker.py
@@ -88,7 +88,6 @@ def __init__(
         *,
         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(
@@ -96,6 +95,7 @@ def __init__(
         )
         bind_routed_experts_capturer(model, capturer)
         self._capturer = capturer
+        self._vllm_config = vllm_config
         self._store: BackgroundBlockObjectStore | None = None
         self._buffer: RoutedExpertsBuffer | None = None
         self._requests: dict[str, _WorkerRequestState] = {}
@@ -107,11 +107,15 @@ def __init__(
         self._pending_outputs: list[PendingAuxOutput] = []
         self._lock = Lock()
         self._max_concurrent_batches = vllm_config.max_concurrent_batches
+
+    def initialize_output_storage(self, kv_cache_config: KVCacheConfig) -> None:
+        """Create the output data plane once the KV cache layout is resolved."""
         # 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
 
+        vllm_config = self._vllm_config
         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(
@@ -131,7 +135,7 @@ def __init__(
             shape_per_token,
             hash_block_size,
             vllm_config.scheduler_config.max_num_seqs,
-            max_num_batched_tokens,
+            vllm_config.scheduler_config.max_num_batched_tokens,
             vllm_config.max_concurrent_batches,
         )
 
@@ -406,10 +410,6 @@ def close(self) -> None:
 
 
 def get_aux_output_connector(
-    model: torch.nn.Module, vllm_config: VllmConfig, kv_cache_config: KVCacheConfig
+    model: torch.nn.Module, vllm_config: VllmConfig
 ) -> AuxOutputWorkerConnector:
-    return AuxOutputWorkerConnector(
-        model=model,
-        kv_cache_config=kv_cache_config,
-        vllm_config=vllm_config,
-    )
+    return AuxOutputWorkerConnector(model=model, vllm_config=vllm_config)
diff --git a/vllm/distributed/device_communicators/all2all.py b/vllm/distributed/device_communicators/all2all.py
index 74ee703a095b..6dc25135cc35 100644
--- a/vllm/distributed/device_communicators/all2all.py
+++ b/vllm/distributed/device_communicators/all2all.py
@@ -13,6 +13,7 @@
 import vllm.envs as envs
 from vllm.config import get_current_vllm_config
 from vllm.distributed import get_dp_group, get_ep_group, get_pcp_group
+from vllm.distributed.device_communicators.cuda_wrapper import CudaRTLibrary
 from vllm.distributed.utils import StatelessProcessGroup
 from vllm.forward_context import get_forward_context
 from vllm.logger import init_logger
@@ -27,6 +28,7 @@
     has_deep_ep_v2,
     has_mori,
 )
+from vllm.utils.math_utils import round_up
 
 from .base_device_communicator import All2AllManagerBase, Cache
 
@@ -412,6 +414,17 @@ def clean_buffers(self) -> None:
         torch.accelerator.synchronize()
         DeepEPLLAll2AllManager._last_mask = None
 
+    def update_mask(self, rank: int, masked: bool = True) -> None:
+        buf = DeepEPLLAll2AllManager._buffer
+        assert buf is not None
+        buf.low_latency_update_mask_buffer(rank, masked)
+        torch.accelerator.synchronize()
+        if DeepEPLLAll2AllManager._last_mask is None:
+            DeepEPLLAll2AllManager._last_mask = torch.zeros(
+                self.world_size, device="cuda", dtype=torch.int32
+            )
+        DeepEPLLAll2AllManager._last_mask[rank] = int(masked)
+
 
 @dataclass
 class _NixlEPBufferState:
@@ -646,6 +659,19 @@ def clean_buffers(self) -> None:
         torch.accelerator.synchronize()
         NixlEPAll2AllManager._last_mask = None
 
+    def update_mask(self, rank: int, masked: bool = True) -> None:
+        state = NixlEPAll2AllManager._buffer
+        assert state is not None
+        state.buffer.update_mask_buffer(rank, mask=masked)
+        torch.accelerator.synchronize()
+        last = NixlEPAll2AllManager._last_mask
+        if last is None or last.shape[0] != state.active_ep_size:
+            NixlEPAll2AllManager._last_mask = torch.zeros(
+                state.active_ep_size, device="cuda", dtype=torch.int32
+            )
+            last = NixlEPAll2AllManager._last_mask
+        last[rank] = int(masked)
+
 
 class FlashInferNVLinkTwoSidedManager(All2AllManagerBase):
     """All2All communication based on flashinfer all2allv/two-sided NVLink kernels."""
@@ -928,6 +954,7 @@ def __init__(self, cpu_group, all2all_backend: str):
         assert all2all_backend in (
             "mori_high_throughput",
             "mori_low_latency",
+            "mori",
         ), f"unsupported MoRI all2all backend: {all2all_backend!r}"
         import mori
 
@@ -982,9 +1009,16 @@ def _make_all2all_kwargs(
             warp_num_per_block = 16
             block_num = 80
         else:
-            # Multi-node: kernel follows --all2all-backend (mirrors deepep_* split).
-            # mori_low_latency → InterNodeV1LL; mori_high_throughput → V1.
-            if self._all2all_backend == "mori_low_latency":
+            # Multi-node: mori_low_latency → InterNodeV1LL; mori_high_throughput
+            # → V1; mori → V1LL for small per-rank token budgets, else V1.
+            if self._all2all_backend == "mori":
+                use_ll = (
+                    max_num_tokens_per_dp_rank
+                    <= envs.VLLM_MORI_DISPATCH_INTER_KERNEL_SWITCH_THRESHOLD
+                )
+            else:
+                use_ll = self._all2all_backend == "mori_low_latency"
+            if use_ll:
                 kernel_type = mori.ops.EpDispatchCombineKernelType.InterNodeV1LL
             else:
                 kernel_type = mori.ops.EpDispatchCombineKernelType.InterNodeV1
@@ -1037,6 +1071,134 @@ def get_handle(self, kwargs):
         return handle
 
 
+def _mori_epv2_arena_bytes(cfg) -> int:
+    """Upper bound of a MORI EPv2 gather-mode op's symmetric arena.
+
+    Mirrors the intranode regions of MORI's hip_backend: per received row
+    disp_out, out_tok, out_scales, out_idx + out_wts (topk * 4 each) and
+    recv_to_src_token (4); per rank recv_num (8 since MORI 1.2.4) and
+    cross_device_barrier (8); and tok_off (4). MORI allocates the arena inside
+    the window itself, so an underestimate fails there rather than overlapping.
+    """
+    # The hip backend pads scale rows to 128 B; FlyDSL uses the unpadded row.
+    scale_row = -(-cfg.scale_dim * cfg.scale_type_size // 128) * 128
+    topk = cfg.num_experts_per_token
+    per_recv_row = cfg.token_nbytes + cfg.combine_token_nbytes + scale_row
+    per_recv_row += topk * 8 + 4
+    per_rank = cfg.world_size * 16
+    # The arena aligns each of its (at most 12) regions to 256 B.
+    return cfg.effective_max_recv * per_recv_row + per_rank + 4 + 12 * 256
+
+
+# cco reserves each rank's window in 4 GiB steps.
+_MORI_EPV2_WINDOW_ALIGN = 4 << 30
+
+
+@dataclass
+class _MoriEPv2Window:
+    comm: Any
+    free_bytes: int
+
+
+class MoriEPv2All2AllManager(All2AllManagerBase):
+    def __init__(self, cpu_group):
+        super().__init__(cpu_group)
+        from mori.cco import GDA_CONNECTION_NONE, CCODevCommRequirements
+
+        self._windows: list[_MoriEPv2Window] = []
+        # The cache holds ops weakly; keep them alive with the windows their
+        # arenas were charged to, so a dropped op is never charged twice.
+        self._ops: list[Any] = []
+        self.handle_cache = Cache()
+        # Open the first window now so a wrong topology fails before the
+        # weights load; ops that need more get their own window later.
+        comm = self._open_window(_MORI_EPV2_WINDOW_ALIGN).comm
+        requirements = CCODevCommRequirements()
+        requirements.gda_connection_type = GDA_CONNECTION_NONE
+        try:
+            device_comm = comm.create_dev_comm(requirements)
+            lsa_size = device_comm.lsa_size
+            device_comm.close()
+            if lsa_size != self.world_size:
+                raise ValueError(
+                    "MORI EPv2 requires all expert-parallel ranks in one LSA team, "
+                    f"got LSA size {lsa_size} for {self.world_size} ranks"
+                )
+        except Exception as e:
+            e.add_note("MORI EPv2: probing the CCO-LSA team of the EP group failed")
+            self.destroy()
+            raise
+        logger.debug("MORI EPv2 all2all over a %d-rank CCO-LSA team", self.world_size)
+
+    def _open_window(self, nbytes: int) -> _MoriEPv2Window:
+        from mori.cco import Communicator
+
+        uid = Communicator.get_unique_id() if self.rank == 0 else None
+        values = [uid]
+        dist.broadcast_object_list(values, group_src=0, group=self.cpu_group)
+        comm = Communicator.init(
+            self.world_size, self.rank, values[0], per_rank_vmm=nbytes
+        )
+        # cco probes fabric handles with hipMemCreate and, where they are
+        # unsupported, leaves hipErrorNotSupported as the HIP last error; the
+        # next torch call that checks it (an allocation) would raise it. MoRI
+        # clears it itself since ROCm/mori#761, which no release has yet.
+        CudaRTLibrary().cudaGetLastError()
+        window = _MoriEPv2Window(comm, nbytes)
+        self._windows.append(window)
+        return window
+
+    def _make_handle(self, **kwargs):
+        from mori.ops.dispatch_combine_v2 import (
+            EpDispatchCombineConfig,
+            EpDispatchCombineOp,
+        )
+
+        config = EpDispatchCombineConfig(
+            rank=self.rank,
+            world_size=self.world_size,
+            data_type=torch.bfloat16,
+            combine_mode="gather",
+            **kwargs,
+        )
+        # Each op carves its arena out of a cco window, which only reserves
+        # address space, so open a big enough one instead of failing. Whole
+        # GiB per op keeps MORI's allocation granularity inside the window.
+        arena = round_up(_mori_epv2_arena_bytes(config), 1 << 30)
+        window = next((w for w in self._windows if w.free_bytes >= arena), None)
+        if window is None:
+            window = self._open_window(round_up(arena, _MORI_EPV2_WINDOW_ALIGN))
+        window.free_bytes -= arena
+        handle = EpDispatchCombineOp(config, window.comm)
+        self._ops.append(handle)
+        window.comm.barrier()
+        logger.debug(
+            "MORI EPv2 op initialized: dispatch=%s combine=%s recv_cap=%d "
+            "schedule=%s arena_gb=%d cco_windows=%d",
+            config.dispatch_dtype,
+            config.combine_dtype,
+            config.effective_max_recv,
+            config.schedule,
+            arena >> 30,
+            len(self._windows),
+        )
+        return handle
+
+    def get_handle(self, kwargs):
+        return self.handle_cache.get_or_create(kwargs, self._make_handle)
+
+    def destroy(self):
+        # Each op's arena lives in a cco window, so free the ops first.
+        with self.handle_cache._lock:
+            self.handle_cache._cache.clear()
+        for handle in self._ops:
+            handle.close()
+        self._ops.clear()
+        for window in self._windows:
+            window.comm.destroy()
+        self._windows.clear()
+
+
 class DeepEPV2All2AllManager(All2AllManagerBase):
     """All2All communication based on DeepEP v2 ElasticBuffer (unified API).
     Uses NCCL Gin backend with analytical SM calculation.
diff --git a/vllm/distributed/device_communicators/base_device_communicator.py b/vllm/distributed/device_communicators/base_device_communicator.py
index 588163b84602..ba39c1cef3ea 100644
--- a/vllm/distributed/device_communicators/base_device_communicator.py
+++ b/vllm/distributed/device_communicators/base_device_communicator.py
@@ -136,6 +136,10 @@ def clean_buffers(self) -> None:
         """
         raise NotImplementedError
 
+    def update_mask(self, rank: int, masked: bool = True) -> None:
+        """Set the mask for a specific EP rank. Used during scale-down."""
+        raise NotImplementedError
+
     def set_num_sms(self, num_sms: int):
         pass
 
diff --git a/vllm/distributed/device_communicators/cuda_communicator.py b/vllm/distributed/device_communicators/cuda_communicator.py
index bba4792fea87..56c9e33ee4f2 100644
--- a/vllm/distributed/device_communicators/cuda_communicator.py
+++ b/vllm/distributed/device_communicators/cuda_communicator.py
@@ -165,9 +165,15 @@ def __init__(
                 register_graph_buffers=register,
             )
 
-        # AITER custom all-gather/reduce-scatter DP-attention dispatch/combine
+        # AITER custom all-gather/reduce-scatter DP-attention dispatch/combine.
+        # AITER requires a torch.distributed-registered group and exchanges IPC
+        # handles via the default store; stateless groups satisfy neither.
+        # Only StatelessGroupCoordinator passes tcp_store_group, so this checks
+        # the property AITER depends on rather than the feature that currently
+        # creates stateless groups (elastic EP).
         if (
             "dp" in unique_name
+            and tcp_store_group is None
             and self.world_size in (2, 4, 8)
             and current_platform.is_rocm()
             and rocm_aiter_ops.is_custom_all_reduce_enabled()
@@ -218,10 +224,21 @@ def __init__(
                 self.all2all_manager = DeepEPLLAll2AllManager(
                     self.cpu_group, tcp_store_group
                 )
+            elif self.all2all_backend == "mori" and self._mori_ep_v2():
+                from .all2all import MoriEPv2All2AllManager
+
+                self.all2all_manager = MoriEPv2All2AllManager(self.cpu_group)
             elif self.all2all_backend in (
                 "mori_high_throughput",
                 "mori_low_latency",
+                "mori",
             ):
+                if self.all2all_backend != "mori":
+                    logger.warning_once(
+                        "'%s' is deprecated and will be removed in a future release. "
+                        "Use 'mori', which picks the MoRI kernels automatically.",
+                        self.all2all_backend,
+                    )
                 from .all2all import MoriAll2AllManager
 
                 self.all2all_manager = MoriAll2AllManager(
@@ -283,6 +300,11 @@ def __init__(
                 scope="global",
             )
 
+    @staticmethod
+    def _mori_ep_v2() -> bool:
+        config = get_current_vllm_config_or_none()
+        return config is not None and bool(config.parallel_config.mori_ep_v2)
+
     def _log_all_reduce_backend_selection(self) -> None:
         """Log the all-reduce backends that are active for this group.
 
diff --git a/vllm/distributed/device_communicators/cuda_wrapper.py b/vllm/distributed/device_communicators/cuda_wrapper.py
index d27d84c03045..a2024d5b996a 100644
--- a/vllm/distributed/device_communicators/cuda_wrapper.py
+++ b/vllm/distributed/device_communicators/cuda_wrapper.py
@@ -78,6 +78,8 @@ class CudaRTLibrary:
             cudaError_t,
             [ctypes.POINTER(ctypes.c_void_p), cudaIpcMemHandle_t, ctypes.c_uint],
         ),
+        # cudaError_t cudaIpcCloseMemHandle ( void* devPtr )
+        Function("cudaIpcCloseMemHandle", cudaError_t, [ctypes.c_void_p]),
         # cudaError_t cudaHostRegister ( void* ptr, size_t size, unsigned int flags )
         Function(
             "cudaHostRegister",
@@ -102,6 +104,7 @@ class CudaRTLibrary:
         "cudaMemcpy": "hipMemcpy",
         "cudaIpcGetMemHandle": "hipIpcGetMemHandle",
         "cudaIpcOpenMemHandle": "hipIpcOpenMemHandle",
+        "cudaIpcCloseMemHandle": "hipIpcCloseMemHandle",
         "cudaHostRegister": "hipHostRegister",
         "cudaHostUnregister": "hipHostUnregister",
         "cudaGetLastError": "hipGetLastError",
@@ -199,6 +202,9 @@ def cudaIpcOpenMemHandle(self, handle: cudaIpcMemHandle_t) -> ctypes.c_void_p:
         )
         return devPtr
 
+    def cudaIpcCloseMemHandle(self, devPtr: ctypes.c_void_p) -> None:
+        self.CUDART_CHECK(self.funcs["cudaIpcCloseMemHandle"](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."""
diff --git a/vllm/distributed/device_communicators/pynccl.py b/vllm/distributed/device_communicators/pynccl.py
index bb508c303927..8981594affda 100644
--- a/vllm/distributed/device_communicators/pynccl.py
+++ b/vllm/distributed/device_communicators/pynccl.py
@@ -486,6 +486,60 @@ def scatter(
             self.recv(output_tensor, root, stream)
         self.nccl.ncclGroupEnd()
 
+    def all_to_all_single(
+        self,
+        output_tensor: torch.Tensor,
+        input_tensor: torch.Tensor,
+        output_split_sizes: list[int],
+        input_split_sizes: list[int],
+        stream=None,
+    ):
+        """All-to-all along dim 0, like `torch.distributed.all_to_all_single`:
+        the next `input_split_sizes[i]` rows of `input_tensor` go to rank i, and
+        the next `output_split_sizes[i]` rows of `output_tensor` come from it.
+        CUDA graph capturable for split sizes fixed at capture.
+        """
+        if self.disabled:
+            return
+        assert input_tensor.device == self.device, (
+            f"this nccl communicator is created to work on {self.device}, "
+            f"but the input tensor is on {input_tensor.device}"
+        )
+        assert input_tensor.is_contiguous() and output_tensor.is_contiguous()
+        if stream is None:
+            stream = current_stream()
+        nccl_stream = cudaStream_t(stream.cuda_stream)
+        nccl_dtype = ncclDataTypeEnum.from_torch(input_tensor.dtype)
+        row_numel = input_tensor[0].numel() if input_tensor.dim() > 1 else 1
+        row_bytes = row_numel * input_tensor.element_size()
+        send_ptr = input_tensor.data_ptr()
+        recv_ptr = output_tensor.data_ptr()
+        self.nccl.ncclGroupStart()
+        for peer in range(self.world_size):
+            send_rows = input_split_sizes[peer]
+            recv_rows = output_split_sizes[peer]
+            if send_rows > 0:
+                self.nccl.ncclSend(
+                    buffer_type(send_ptr),
+                    send_rows * row_numel,
+                    nccl_dtype,
+                    peer,
+                    self.comm,
+                    nccl_stream,
+                )
+            if recv_rows > 0:
+                self.nccl.ncclRecv(
+                    buffer_type(recv_ptr),
+                    recv_rows * row_numel,
+                    nccl_dtype,
+                    peer,
+                    self.comm,
+                    nccl_stream,
+                )
+            send_ptr += send_rows * row_bytes
+            recv_ptr += recv_rows * row_bytes
+        self.nccl.ncclGroupEnd()
+
     def send(self, tensor: torch.Tensor, dst: int, stream=None):
         if self.disabled:
             return
diff --git a/vllm/distributed/ec_transfer/ec_connector/cpu/control/zmq.py b/vllm/distributed/ec_transfer/ec_connector/cpu/control/zmq.py
index d87bdc792057..7d17d069a374 100644
--- a/vllm/distributed/ec_transfer/ec_connector/cpu/control/zmq.py
+++ b/vllm/distributed/ec_transfer/ec_connector/cpu/control/zmq.py
@@ -174,17 +174,27 @@ class ZmqServerTransport:
     routes a reply back to a specific peer by ZMQ identity.
     """
 
-    def __init__(self, host: str, port: int) -> None:
+    def __init__(self, host: str, port: int = 0) -> None:
         self._ctx = zmq.Context()
-        self._router = make_zmq_socket(
-            ctx=self._ctx,
-            path=make_zmq_path(scheme="tcp", host=host, port=port),
-            socket_type=zmq.ROUTER,
-            bind=True,
-        )
+        try:
+            self._router = make_zmq_socket(
+                ctx=self._ctx,
+                path=make_zmq_path(scheme="tcp", host=host, port=port),
+                socket_type=zmq.ROUTER,
+                bind=True,
+            )
+        except Exception:
+            self._ctx.destroy(linger=0)
+            raise
         self._poller = zmq.Poller()
         self._poller.register(self._router, zmq.POLLIN)
 
+    @property
+    def port(self) -> int:
+        """Return the actual bound port, including when port 0 was requested."""
+        endpoint = self._router.getsockopt_string(zmq.LAST_ENDPOINT)
+        return int(endpoint.rsplit(":", 1)[1])
+
     def poll(self, timeout_ms: int = 0) -> list[tuple[bytes, bytes]]:
         """Wait up to timeout_ms for incoming messages, then drain all pending.
 
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 567f2d59c9f3..b487fa7dee40 100644
--- a/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/__init__.py
+++ b/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/__init__.py
@@ -140,6 +140,7 @@ def _setup_nixl(self, vllm_config: "VllmConfig") -> None:
             ProducerSession,
         )
         from vllm.distributed.nixl_utils import NixlWrapper, nixl_agent_config
+        from vllm.utils.network_utils import get_ip
         from vllm.version import __version__ as VLLM_VERSION
 
         if NixlWrapper is None or nixl_agent_config is None:
@@ -149,6 +150,7 @@ def _setup_nixl(self, vllm_config: "VllmConfig") -> None:
             )
         engine_id = self._ec_config.engine_id
         assert engine_id is not None
+        dp_index = vllm_config.parallel_config.data_parallel_index
         # How long a consumer waits for an XferAck. The producer answers from
         # its scheduler step, so its reply latency scales with the encoder's
         # batch size: a deployment whose steps run longer must raise this.
@@ -163,24 +165,31 @@ def _setup_nixl(self, vllm_config: "VllmConfig") -> None:
             dtype=str(self._dtype),
             block_size_bytes=self._region.block_size_bytes,
         )
-        if self._is_producer:
-            self._peer_host = envs.VLLM_EC_SIDE_CHANNEL_HOST
-            self._peer_port = envs.VLLM_EC_SIDE_CHANNEL_PORT
-
         self._data = NixlDataTransport(
-            agent_name=engine_id,
+            agent_name=f"{engine_id}_dp{dp_index}",
             base_ptr=self._region.blocks.data_ptr(),
             num_blocks=self._region.num_blocks,
             block_size_bytes=self._region.block_size_bytes,
             total_size_bytes=self._region.num_blocks * self._region.block_size_bytes,
         )
         if self._is_producer:
-            assert self._peer_host is not None
-            assert self._peer_port is not None
+            self._peer_host = envs.VLLM_EC_SIDE_CHANNEL_HOST or get_ip()
+            if self._peer_host in ("0.0.0.0", "::", "*"):
+                raise ValueError(
+                    "Set VLLM_EC_SIDE_CHANNEL_HOST or VLLM_HOST_IP to a "
+                    "reachable address, not a wildcard."
+                )
+            base_port = envs.VLLM_EC_SIDE_CHANNEL_PORT
+            port = base_port + dp_index if base_port else 0
+            if base_port < 0 or not 0 <= port <= 65535:
+                raise ValueError(
+                    "VLLM_EC_SIDE_CHANNEL_PORT + data_parallel_index "
+                    "must be between 1 and 65535, or use port 0 for auto allocation."
+                )
+            transport = ZmqServerTransport(host=self._peer_host, port=port)
+            self._peer_port = transport.port
             self._producer_session = ProducerSession(
-                transport=ZmqServerTransport(
-                    host=self._peer_host, port=self._peer_port
-                ),
+                transport=transport,
                 data=self._data,
                 cache=self._cache,
                 compat_hash=self._compat_hash,
diff --git a/vllm/distributed/elastic_ep/elastic_execute.py b/vllm/distributed/elastic_ep/elastic_execute.py
index 182a4b9bb4f3..f064e2e86af0 100644
--- a/vllm/distributed/elastic_ep/elastic_execute.py
+++ b/vllm/distributed/elastic_ep/elastic_execute.py
@@ -12,6 +12,7 @@
 import torch.nn as nn
 from torch.distributed import P2POp
 
+from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphWrapper
 from vllm.compilation.cuda_graph import CUDAGraphWrapper
 from vllm.compilation.wrapper import reset_compile_wrapper
 from vllm.config import ParallelConfig, set_current_vllm_config
@@ -460,6 +461,37 @@ def _release_cuda_graphs(self) -> None:
         gc.collect()
         torch.accelerator.synchronize()
         torch.accelerator.empty_cache()
+        self._switch_to_fresh_graph_pool()
+
+    def _switch_to_fresh_graph_pool(self) -> None:
+        # Libraries may keep capture-time tensors alive after their graphs are
+        # released (e.g. FlyDSL retains the last launch's DLPack arguments), so
+        # the old pool cannot be erased. Re-capturing into that pool trips the
+        # allocator's "use_count > 0" assert; capture into a fresh one instead.
+        platform_cls = type(current_platform)
+        if platform_cls._global_graph_pool is None:
+            return
+        new_pool = current_platform.graph_pool_handle()
+        platform_cls._global_graph_pool = new_pool
+
+        model_runner = self.worker.model_runner
+        manager = getattr(model_runner, "cudagraph_manager", None)
+        speculator = getattr(model_runner, "speculator", None)
+        managers = [manager] + [
+            attr
+            for attr in (vars(speculator).values() if speculator else ())
+            if isinstance(attr, CudaGraphManager)
+        ]
+        for m in managers:
+            if m is not None and m.pool is not None:
+                m.pool = new_pool
+
+        wrappers: list[CUDAGraphWrapper | BreakableCUDAGraphWrapper] = [
+            *CUDAGraphWrapper._all_instances,
+            *BreakableCUDAGraphWrapper._all_instances,
+        ]
+        for wrapper in wrappers:
+            wrapper.graph_pool = new_pool
 
     def switch_and_remove(self) -> None:
         # Removing ranks skipped preparation, so wait for prior cleanup here.
diff --git a/vllm/distributed/eplb/async_worker.py b/vllm/distributed/eplb/async_worker.py
index 0f09bb99c493..17e4691e77f1 100644
--- a/vllm/distributed/eplb/async_worker.py
+++ b/vllm/distributed/eplb/async_worker.py
@@ -3,6 +3,7 @@
 """The async worker that transfers experts in the background."""
 
 import threading
+import time
 from typing import TYPE_CHECKING
 
 import torch
@@ -38,7 +39,8 @@ def thread_target() -> None:
                 stream=stream,
                 is_profile=is_profile,
             )
-        except Exception as exc:  # pragma: no cover - diagnostic path
+        except Exception as exc:
+            state.async_worker_error = exc
             logger.exception("async loop error (Rank %d): %s", rank, str(exc))
 
     thread = threading.Thread(target=thread_target, daemon=True)
@@ -77,110 +79,153 @@ def transfer_run_periodically(
     stream: torch.Stream,
     is_profile: bool = False,
 ) -> None:
+    rank = get_eplb_group().device_group.rank()
     while True:
-        state.rearrange_event.wait(stream=stream)
-
-        eplb_group = get_eplb_group().device_group
-        eplb_cpu_group = get_eplb_group().cpu_group
-        ep_rank = eplb_group.rank()
-
-        assert state.is_async
-        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(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 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, stream
+        try:
+            state.rearrange_event.wait(stream=stream)
+            _transfer_cycle(state, stream, is_profile)
+        except Exception:
+            if not state.parallel_config.enable_fault_tolerance:
+                raise
+            # FT keeps the thread alive so EPLB can resume after a retry.
+            logger.exception(
+                "async loop error (Rank %d); continuing under fault tolerance",
+                rank,
             )
 
-            migration_batches_by_layer = None
-            if state.parallel_config.eplb_config.migration_batching_enabled:
-                profile_name = "eplb: schedule migration batches"
-                num_local_experts = model_state.model.expert_weights[0][0].shape[0]
-                # The placement snapshots stay fixed for this rebalance cycle.
-                # Precompute all layer schedules once, but execute transfers one
-                # layer at a time below to preserve the migration lifecycle. This
-                # profiling range therefore covers scheduling for every MoE layer.
-                with torch.profiler.record_function(profile_name):
-                    migration_batches_by_layer = schedule_migration_batches_for_layers(
-                        num_local_experts,
-                        physical_to_logical_map_cpu.numpy(),
-                        new_physical_to_logical_map.numpy(),
-                    )
-
-            # Execute one EPLB layer transfer per model forward pass. Each iteration
-            # of this loop will copy the new set of expert weights into
-            # model_state.expert_buffer, which will be consumed by the main thread in
-            # move_to_workspace.
-            # We sync the rebalanced flag across ranks before each iteration so
-            # all ranks make a coordinated decision to continue or stop.
-            while layer_idx < num_layers:
-                flag = torch.tensor(
-                    [int(model_state.rebalanced)],
-                    dtype=torch.int32,
-                    device="cpu",
-                )
-                torch.distributed.all_reduce(flag, group=eplb_cpu_group)
-                if int(flag.item()) != eplb_cpu_group.size():
-                    logger.warning(
-                        "async worker (rank=%d): layer %d coordinated stop "
-                        "(flag_sum=%d, group_size=%d)",
-                        ep_rank,
-                        layer_idx,
-                        int(flag.item()),
-                        eplb_cpu_group.size(),
-                    )
-                    model_state.rebalanced = False
-                    break
-
-                transfer_metadata = transfer_layer(
-                    old_layer_indices=physical_to_logical_map_cpu[layer_idx],
-                    new_layer_indices=new_physical_to_logical_map[layer_idx],
-                    expert_weights=model_state.model.expert_weights[layer_idx],
-                    expert_weights_buffer=model_state.expert_buffer,
-                    communicator=model_state.communicator,
-                    ep_group=eplb_group,
-                    is_profile=is_profile,
-                    stream=stream,
-                    layer_idx=layer_idx,
-                    enable_migration_batching=(
-                        state.parallel_config.eplb_config.migration_batching_enabled
-                    ),
-                    migration_batches=(
-                        # This is only the current layer's schedule; it does not
-                        # combine transfers from different layers.
-                        migration_batches_by_layer[layer_idx]
-                        if migration_batches_by_layer is not None
-                        else None
-                    ),
-                )
 
-                # Wait until all writes to expert_buffer have finished before making the
-                # AsyncEplbLayerResult visible to the main thread.
-                stream.synchronize()
+def _transfer_cycle(
+    state: "EplbState",
+    stream: torch.Stream,
+    is_profile: bool,
+) -> None:
+    """Run one rebalance pass over all model states and layers."""
+    # Snapshot the FT reset epoch. A change in this value means a fault
+    # recovery happened, and the in-flight rebalance must be abandoned.
+    reset_epoch = state.ft_reset_epoch
 
-                # This event guarantees that expert_buffer will not be overwritten by
-                # subsequent iterations of this loop until the main thread has consumed
-                # it. Record is called by the main thread after move_from_buffer().
-                consumed_event = CpuGpuEvent()
+    eplb_group = get_eplb_group().device_group
+    eplb_cpu_group = get_eplb_group().cpu_group
+    ep_rank = eplb_group.rank()
+
+    assert state.is_async
+    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(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 device_stream(stream):
+            physical_to_logical_map_cpu = model_state.physical_to_logical_map.cpu()
+
+        t_start = time.perf_counter()
+        new_physical_to_logical_map = run_rebalance_experts(
+            model_state, state, physical_to_logical_map_cpu, stream
+        )
+
+        migration_batches_by_layer = None
+        if state.parallel_config.eplb_config.migration_batching_enabled:
+            profile_name = "eplb: schedule migration batches"
+            num_local_experts = model_state.model.expert_weights[0][0].shape[0]
+            # The placement snapshots stay fixed for this rebalance cycle.
+            # Precompute all layer schedules once, but execute transfers one
+            # layer at a time below to preserve the migration lifecycle. This
+            # profiling range therefore covers scheduling for every MoE layer.
+            with torch.profiler.record_function(profile_name):
+                migration_batches_by_layer = schedule_migration_batches_for_layers(
+                    num_local_experts,
+                    physical_to_logical_map_cpu.numpy(),
+                    new_physical_to_logical_map.numpy(),
+                )
 
-                model_state.pending_result = AsyncEplbLayerResult(
-                    layer_idx=layer_idx,
-                    new_physical_to_logical_map=new_physical_to_logical_map[layer_idx],
-                    transfer_metadata=transfer_metadata,
-                    consumed_event=consumed_event,
+        # Execute one EPLB layer transfer per model forward pass. Each iteration
+        # of this loop will copy the new set of expert weights into
+        # model_state.expert_buffer, which will be consumed by the main thread in
+        # move_to_workspace.
+        # We sync the rebalanced flag across ranks before each iteration so
+        # all ranks make a coordinated decision to continue or stop.
+        while layer_idx < num_layers:
+            flag = torch.tensor(
+                [int(model_state.rebalanced)],
+                dtype=torch.int32,
+                device="cpu",
+            )
+            torch.distributed.all_reduce(flag, group=eplb_cpu_group)
+            if int(flag.item()) != eplb_cpu_group.size():
+                logger.warning(
+                    "async worker (rank=%d): layer %d coordinated stop "
+                    "(flag_sum=%d, group_size=%d)",
+                    ep_rank,
+                    layer_idx,
+                    int(flag.item()),
+                    eplb_cpu_group.size(),
                 )
+                model_state.rebalanced = False
+                break
+
+            transfer_metadata = transfer_layer(
+                old_layer_indices=physical_to_logical_map_cpu[layer_idx],
+                new_layer_indices=new_physical_to_logical_map[layer_idx],
+                expert_weights=model_state.model.expert_weights[layer_idx],
+                expert_weights_buffer=model_state.expert_buffer,
+                communicator=model_state.communicator,
+                ep_group=eplb_group,
+                is_profile=is_profile,
+                stream=stream,
+                layer_idx=layer_idx,
+                enable_migration_batching=(
+                    state.parallel_config.eplb_config.migration_batching_enabled
+                ),
+                migration_batches=(
+                    # This is only the current layer's schedule; it does not
+                    # combine transfers from different layers.
+                    migration_batches_by_layer[layer_idx]
+                    if migration_batches_by_layer is not None
+                    else None
+                ),
+            )
 
-                # 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=stream)
-                assert model_state.pending_result is None
-                layer_idx += 1
+            # Wait until all writes to expert_buffer have finished before making the
+            # AsyncEplbLayerResult visible to the main thread.
+            stream.synchronize()
+
+            # This event guarantees that expert_buffer will not be overwritten by
+            # subsequent iterations of this loop until the main thread has consumed
+            # it. Record is called by the main thread after move_from_buffer().
+            consumed_event = CpuGpuEvent()
+
+            model_state.pending_result = AsyncEplbLayerResult(
+                layer_idx=layer_idx,
+                new_physical_to_logical_map=new_physical_to_logical_map[layer_idx],
+                transfer_metadata=transfer_metadata,
+                consumed_event=consumed_event,
+            )
+            # A fault recovery happened during the transfer, so the main
+            # thread will never consume this result. Drop it and abandon
+            # the rebalance.
+            if state.ft_reset_epoch != reset_epoch:
+                model_state.pending_result = None
+                return
+
+            # 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=stream)
+            # Host-issued writes into expert_buffer (e.g. NIXL RDMA reads) are not
+            # stream-ordered, so wait until move_from_buffer has run on the GPU.
+            stream.synchronize()
+            assert model_state.pending_result is None
+            # Woken by an FT reset's record, not a normal consume.
+            if state.ft_reset_epoch != reset_epoch:
+                return
+            layer_idx += 1
+        if ep_rank == 0:
+            logger.info(
+                "Async EPLB rearranged %d/%d layers of %s in %.2f s.",
+                layer_idx,
+                num_layers,
+                model_state.model_name,
+                time.perf_counter() - t_start,
+            )
diff --git a/vllm/distributed/eplb/eplb_communicator.py b/vllm/distributed/eplb/eplb_communicator.py
index 409a031c03dc..61d0f2d94047 100644
--- a/vllm/distributed/eplb/eplb_communicator.py
+++ b/vllm/distributed/eplb/eplb_communicator.py
@@ -3,6 +3,7 @@
 """EPLB communicator implementations and factory."""
 
 import contextlib
+import itertools
 import time
 import uuid
 from abc import ABC, abstractmethod
@@ -26,6 +27,7 @@
 from vllm.distributed.parallel_state import (
     GroupCoordinator,
     get_pp_group,
+    in_the_same_node_as,
     is_local_first_rank,
 )
 from vllm.distributed.stateless_coordinator import StatelessGroupCoordinator
@@ -33,6 +35,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 round_up
 from vllm.utils.torch_utils import PIN_MEMORY
 
 logger = init_logger(__name__)
@@ -74,14 +77,34 @@ def execute(self) -> None:
         """
 
     def set_transfer_context(  # noqa: B027
-        self, old_indices: np.ndarray, layer_idx: int
+        self,
+        old_indices: np.ndarray,
+        layer_idx: int,
+        recv_rows: Sequence[dict[int, int]] | None = None,
     ) -> None:
-        """Pre-set layer context before add_recv calls.
+        """Pre-set layer context before add_send/add_recv calls.
 
         Default is a no-op; overridden by backends (e.g. NIXL) that need
         layer-level context to issue transfers inside add_recv.
+
+        Args:
+            old_indices: Global physical-to-logical map before the move.
+            layer_idx: Index of the MoE layer being transferred.
+            recv_rows: Per EP rank, the ``expert_buffer`` row each expert
+                received by that rank lands in. Only push backends need it.
+
         """
 
+    @property
+    def needs_recv_rows(self) -> bool:
+        """Whether set_transfer_context must be given every rank's recv rows."""
+        return False
+
+    @property
+    def owned_expert_buffer(self) -> list[torch.Tensor] | None:
+        """Backend-allocated receive buffer, or ``None`` to use the caller's."""
+        return None
+
     @property
     def needs_profile_buffer_reservation(self) -> bool:
         """Whether the profile path must run a dummy collective operation to reserve
@@ -416,7 +439,12 @@ def add_send(
         # weights are pre-registered and always readable in-place.
         pass
 
-    def set_transfer_context(self, old_indices: np.ndarray, layer_idx: int) -> None:
+    def set_transfer_context(
+        self,
+        old_indices: np.ndarray,
+        layer_idx: int,
+        recv_rows: Sequence[dict[int, int]] | None = None,
+    ) -> None:
         assert not self._xfer_entries, (
             f"set_transfer_context() called with {len(self._xfer_entries)} "
             f"pending transfers from layer {self._layer_idx}; "
@@ -707,11 +735,124 @@ def execute(self) -> None:
             self._group_started = False
 
 
+class SymmMemEplbCommunicator(EplbCommunicator):
+    """Push expert rows into peers' receive buffers over symmetric memory.
+
+    Senders copy from their live weights into the receiver's ``expert_buffer``
+    row, which every rank derives from the global placements; ``execute`` is
+    collective. CPU barriers only: GPU spin barriers can deadlock MegaMoE
+    cooperative launches.
+    """
+
+    def __init__(
+        self,
+        group_coordinator: GroupCoordinator,
+        layer: Sequence[torch.Tensor],
+    ) -> None:
+        import torch.distributed._symmetric_memory as symm_mem
+
+        from vllm.distributed.device_communicators.custom_all_reduce import (
+            _group_can_attempt_mnnvl,
+        )
+
+        ep_group = group_coordinator.device_group
+        cpu_group = group_coordinator.cpu_group
+        device = layer[0].device
+        if not current_platform.is_cuda():
+            raise ValueError("symm_mem EPLB requires CUDA")
+        if all(in_the_same_node_as(cpu_group, source_rank=0)):
+            physical_ids = [0] * cpu_group.size()
+            torch.distributed.all_gather_object(
+                physical_ids,
+                current_platform.visible_device_id_to_physical_device_id(device.index),
+                group=cpu_group,
+            )
+            in_nvlink_domain = current_platform.is_fully_connected(physical_ids)
+        else:
+            in_nvlink_domain = _group_can_attempt_mnnvl(cpu_group, device)
+        if not in_nvlink_domain:
+            raise ValueError(
+                "symm_mem EPLB requires all EP ranks in one NVLink domain "
+                "(MNNVL across nodes); use nixl or torch_gloo instead."
+            )
+
+        # One buffer and one rendezvous cover all of the layer's tensors.
+        offsets = [0, *itertools.accumulate(round_up(w.nbytes, 256) for w in layer)]
+
+        def as_layer(buffer: torch.Tensor) -> list[torch.Tensor]:
+            return [
+                buffer[o : o + w.nbytes].view(w.dtype).view(w.shape)
+                for o, w in zip(offsets, layer)
+            ]
+
+        buffer = symm_mem.empty(offsets[-1], dtype=torch.uint8, device=device)
+        # For MNNVL, rendezvous is what confirms all ranks share one NVLink domain.
+        self._handle = symm_mem.rendezvous(buffer, ep_group.group_name)
+        self._peer_buffers = [
+            as_layer(self._handle.get_buffer(rank, buffer.shape, torch.uint8))
+            for rank in range(ep_group.size())
+        ]
+        self._buffer = as_layer(buffer)
+
+        self._stream = None
+        self._group_coordinator = group_coordinator
+        self._recv_rows: Sequence[dict[int, int]] = []
+        self._sends: list[tuple[list[torch.Tensor], int, int]] = []
+        self._log_initialized()
+
+    @property
+    def needs_profile_buffer_reservation(self) -> bool:
+        return False
+
+    @property
+    def needs_recv_rows(self) -> bool:
+        return True
+
+    @property
+    def owned_expert_buffer(self) -> list[torch.Tensor]:
+        return self._buffer
+
+    def set_transfer_context(
+        self,
+        old_indices: np.ndarray,
+        layer_idx: int,
+        recv_rows: Sequence[dict[int, int]] | None = None,
+    ) -> None:
+        assert recv_rows is not None
+        self._recv_rows = recv_rows
+
+    def add_send(
+        self, tensors: list[torch.Tensor], dst_rank: int, expert_id: int
+    ) -> None:
+        row = self._recv_rows[dst_rank][expert_id]
+        self._sends.append((tensors, dst_rank, row))
+
+    def add_recv(
+        self, tensors: list[torch.Tensor], src_rank: int, expert_id: int
+    ) -> None:
+        # No-op: senders push into this rank's buffer in execute().
+        pass
+
+    def execute(self) -> None:
+        with device_stream(self._stream):
+            self._barrier()  # Every peer's buffer rows are free to overwrite.
+            for tensors, dst_rank, row in self._sends:
+                for src, peer in zip(tensors, self._peer_buffers[dst_rank]):
+                    peer[row].copy_(src, non_blocking=True)
+            self._barrier()  # Every push has landed in its receiver's buffer.
+        self._sends.clear()
+
+    def _barrier(self) -> None:
+        with gpu_sync_allowed():
+            torch.accelerator.current_stream().synchronize()
+        self._group_coordinator.cpu_all_reduce_(torch.zeros(1, dtype=torch.int32))
+
+
 def create_eplb_communicator(
     group_coordinator: GroupCoordinator,
     backend: str,
     expert_weights: Sequence[Sequence[torch.Tensor]],
-    expert_buffer: Sequence[torch.Tensor],
+    expert_buffer: Sequence[torch.Tensor] | None,
 ) -> EplbCommunicator:
     """Create an EPLB communicator for the given backend.
 
@@ -719,7 +860,7 @@ def create_eplb_communicator(
         group_coordinator: Process-group coordinator that provides the
             device and CPU communication groups.
         backend: Communicator backend name (``"torch_nccl"``,
-            ``"torch_gloo"``, ``"pynccl"``, or ``"nixl"``).
+            ``"torch_gloo"``, ``"pynccl"``, ``"nixl"``, or ``"symm_mem"``).
             Falls back to ``"torch_nccl"`` when *None*.
             Stateless (elastic EP) groups support ``"torch_nccl"``,
             ``"pynccl"``, and ``"nixl"``; ``"torch_nccl"`` is silently
@@ -731,7 +872,8 @@ def create_eplb_communicator(
             NixlEplbCommunicator registers all layers with NIXL for
             zero-copy RDMA reads.
         expert_buffer: Pre-allocated receive buffer tensors (one per
-            weight tensor in a single layer).
+            weight tensor in a single layer). ``symm_mem`` ignores it and
+            uses its own ``owned_expert_buffer``.
 
     """
     first_layer = expert_weights[0] if expert_weights else []
@@ -793,7 +935,9 @@ def _create_pynccl() -> EplbCommunicator:
             backend = "pynccl"
         return _create_pynccl()
 
-    if backend == "nixl":
+    if backend == "symm_mem":
+        return SymmMemEplbCommunicator(group_coordinator, first_layer)
+    elif backend == "nixl":
         if not has_nixl():
             raise RuntimeError(
                 "EPLB communicator 'nixl' requested but NIXL is unavailable."
@@ -803,6 +947,7 @@ def _create_pynccl() -> EplbCommunicator:
                 "EPLB communicator 'nixl' supports only cuda-like devices "
                 f"(got {tensor_device_type})."
             )
+        assert expert_buffer is not None
         try:
             return NixlEplbCommunicator(
                 cpu_group=group_coordinator.cpu_group,
diff --git a/vllm/distributed/eplb/eplb_state.py b/vllm/distributed/eplb/eplb_state.py
index 303d69ba2397..ce6602cc9c32 100644
--- a/vllm/distributed/eplb/eplb_state.py
+++ b/vllm/distributed/eplb/eplb_state.py
@@ -27,6 +27,7 @@
 
 import threading
 import time
+from collections import deque
 from collections.abc import Sequence
 from dataclasses import dataclass
 
@@ -276,6 +277,13 @@ def __init__(self, parallel_config: ParallelConfig, device: torch.device):
         a single ``.fill_()`` updates all layers at once.  Allocated on the
         first call to :meth:`_propagate_shared_tensors`.
         """
+        self._should_record_value: bool | None = None
+        """Last value written to ``should_record_tensor``."""
+        self._pending_balancedness: deque[
+            tuple[int, int, torch.Event, list[tuple[str, torch.Tensor]]]
+        ] = deque()
+        """Balancedness stats whose copies to the host are in flight, oldest
+        first."""
         self.is_async: bool = False
         """
         The flag indicates whether the EPLB is running in async mode.
@@ -284,10 +292,20 @@ def __init__(self, parallel_config: ParallelConfig, device: torch.device):
         """
         Event to signal when a new rearrangement is needed for the async thread.
         """
+        self.ft_reset_epoch: int = 0
+        """
+        Incremented by FT recovery (reset_eplb_async_state). The async worker
+        snapshots it when a transfer cycle starts; a change in the value means
+        a recovery happened, so the cycle is abandoned.
+        """
         self.async_worker: threading.Thread | None = None
         """
         Background thread handling async transfers.
         """
+        self.async_worker_error: Exception | None = None
+        """
+        Exception that stopped the async worker thread, if any.
+        """
         self.device_index: int | None = None
         """
         Device index for the async EPLB worker thread.
@@ -400,7 +418,12 @@ def add_model(
             f"num_redundant_experts {model.num_redundant_experts} "
             f"must be less than or equal to {MAX_EXPERT_REDUNDANCY}"
         )
-        max_slots_per_logical_expert = MAX_EXPERT_REDUNDANCY + 1
+        # Only elastic EP changes the physical expert count at runtime.
+        max_slots_per_logical_expert = (
+            MAX_EXPERT_REDUNDANCY + 1
+            if self.parallel_config.enable_elastic_ep
+            else model.num_redundant_experts + 1
+        )
         logical_to_physical_map = torch.full(
             (model.num_logical_experts, max_slots_per_logical_expert),
             -1,
@@ -489,17 +512,24 @@ def add_model(
             logical_replica_count,
         )
         self._propagate_shared_tensors(model, num_unpadded_tokens_tensors)
-        expert_buffer = [torch.empty_like(w) for w in model.expert_weights[0]]
 
-        assert self.parallel_config.eplb_config.communicator is not None, (
+        backend = self.parallel_config.eplb_config.communicator
+        assert backend is not None, (
             "EPLB communicator backend must be set by ParallelConfig"
         )
+        expert_buffer = (
+            None
+            if backend == "symm_mem"
+            else [torch.empty_like(w) for w in model.expert_weights[0]]
+        )
         communicator = create_eplb_communicator(
             group_coordinator=get_eplb_group(),
-            backend=self.parallel_config.eplb_config.communicator,
+            backend=backend,
             expert_weights=model.expert_weights,
             expert_buffer=expert_buffer,
         )
+        expert_buffer = communicator.owned_expert_buffer or expert_buffer
+        assert expert_buffer is not None
 
         model_state = EplbModelState(
             physical_to_logical_map=physical_to_logical_map,
@@ -593,6 +623,8 @@ def step(
             self.rearrange(is_profile=True)
             return
 
+        if log_stats:
+            self._maybe_log_balancedness()
         if (
             log_stats
             and self.expert_rearrangement_step
@@ -603,46 +635,44 @@ def step(
             # expert_load_pass: (num_moe_layers, num_physical_experts)
             expert_load_pass_list = self._sync_load_pass()
             ep_group = get_ep_group().device_group
-            for expert_load_pass, eplb_model_state in zip(
-                expert_load_pass_list, self.model_states.values()
-            ):
-                # num_tokens_per_rank: (num_moe_layers, num_ranks)
-                num_tokens_per_rank = (
-                    expert_load_pass.reshape(
-                        expert_load_pass.shape[0], ep_group.size(), -1
+            if ep_group.rank() == 0:
+                stats: list[tuple[str, torch.Tensor]] = []
+                for expert_load_pass, eplb_model_state in zip(
+                    expert_load_pass_list, self.model_states.values()
+                ):
+                    # num_tokens_per_rank: (num_moe_layers, num_ranks)
+                    num_tokens_per_rank = (
+                        expert_load_pass.reshape(
+                            expert_load_pass.shape[0], ep_group.size(), -1
+                        )
+                        .sum(dim=-1)
+                        .float()
                     )
-                    .sum(dim=-1)
-                    .float()
-                )
-
-                # Compute balancedness ratio:
-                # for each layer:
-                #   (mean load across ranks) / (max load across ranks)
-                avg_tokens_tensor, max_tokens_tensor = _compute_eplb_load_stats(
-                    num_tokens_per_rank
-                )
-
-                # This gpu/cpu sync only happens with expert stat logging enabled.
-                with gpu_sync_allowed():
-                    tokens_tensors: list[float] = torch.stack(
-                        [avg_tokens_tensor, max_tokens_tensor]
-                    ).tolist()
-                avg_tokens, max_tokens = tokens_tensors
-                balancedness = avg_tokens / max_tokens if max_tokens > 0 else 0.0
-
-                if ep_group.rank() == 0:
-                    logger.info(
-                        "EPLB step: %d for model %s: avg_tokens=%.2f, "
-                        "max_tokens=%d, balancedness=%.4f, "
-                        "steps until the next rearrangement: %d",
+                    # Compute balancedness ratio:
+                    # for each layer:
+                    #   (mean load across ranks) / (max load across ranks)
+                    avg_tokens_tensor, max_tokens_tensor = _compute_eplb_load_stats(
+                        num_tokens_per_rank
+                    )
+                    # A blocking read here would wait for the queued forward.
+                    tokens = torch.stack([avg_tokens_tensor, max_tokens_tensor])
+                    stats.append(
+                        (
+                            eplb_model_state.model_name,
+                            tokens.to("cpu", non_blocking=True),
+                        )
+                    )
+                copied = torch.Event()
+                copied.record()
+                self._pending_balancedness.append(
+                    (
                         self.expert_rearrangement_step,
-                        eplb_model_state.model_name,
-                        avg_tokens,
-                        max_tokens,
-                        balancedness,
                         self.expert_rearrangement_step_interval
                         - self.expert_rearrangement_step,
+                        copied,
+                        stats,
                     )
+                )
 
         # Update the expert load sliding window
         if not is_dummy:
@@ -694,6 +724,28 @@ def step(
 
         self._update_layer_should_record(log_stats=log_stats)
 
+    def _maybe_log_balancedness(self) -> None:
+        pending = self._pending_balancedness
+        while pending:
+            step, steps_left, copied, stats = pending[0]
+            if not copied.query():
+                break
+            pending.popleft()
+            for model_name, tokens in stats:
+                avg_tokens, max_tokens = tokens.tolist()
+                balancedness = avg_tokens / max_tokens if max_tokens > 0 else 0.0
+                logger.info(
+                    "EPLB step: %d for model %s: avg_tokens=%.2f, "
+                    "max_tokens=%d, balancedness=%.4f, "
+                    "steps until the next rearrangement: %d",
+                    step,
+                    model_name,
+                    avg_tokens,
+                    max_tokens,
+                    balancedness,
+                    steps_left,
+                )
+
     def _should_record_current_step(self, log_stats: bool = False) -> bool:
         """Return whether expert-load recording should be enabled this step.
 
@@ -718,10 +770,13 @@ def _should_record_current_step(self, log_stats: bool = False) -> bool:
 
     def _update_layer_should_record(self, log_stats: bool = False) -> None:
         """Update the shared ``should_record_tensor`` for all layers."""
-        if self.should_record_tensor is not None:
-            self.should_record_tensor.fill_(
-                self._should_record_current_step(log_stats=log_stats)
-            )
+        if self.should_record_tensor is None:
+            return
+        should_record = self._should_record_current_step(log_stats=log_stats)
+        # The value only flips at window edges; skip the per-step fill kernel.
+        if should_record != self._should_record_value:
+            self.should_record_tensor.fill_(should_record)
+            self._should_record_value = should_record
 
     def _propagate_shared_tensors(
         self,
@@ -1030,25 +1085,22 @@ def drain_async(self) -> None:
                 )
 
     def _all_ranks_result_ready(self, model_state: EplbModelState) -> bool:
-        parallel_state = get_ep_group()
-        has_result = int(model_state.pending_result is not None)
-
-        cpu_group = getattr(parallel_state, "cpu_group", None)
-        if cpu_group is not None and cpu_group.size() > 1:
-            flag = torch.tensor((has_result,), dtype=torch.int32, device="cpu")
-            all_reduce(flag, group=cpu_group)
-            return int(flag.item()) == cpu_group.size()
-
-        device_group = parallel_state.device_group
-        if device_group.size() <= 1:
-            return bool(has_result)
-
-        device = getattr(
-            parallel_state, "device", model_state.physical_to_logical_map.device
+        flags = torch.tensor(
+            (
+                int(self.async_worker_error is not None),
+                int(model_state.pending_result is None),
+            ),
+            dtype=torch.int32,
+            device="cpu",
         )
-        flag = torch.tensor((has_result,), dtype=torch.int32, device=device)
-        all_reduce(flag, group=device_group)
-        return int(flag.item()) == device_group.size()
+        # cpu_all_reduce_ only sums, so count failed and pending ranks.
+        get_ep_group().cpu_all_reduce_(flags)
+        num_failed, num_pending = flags.tolist()
+        if num_failed:
+            raise RuntimeError(
+                "Async EPLB worker failed on at least one rank; see the worker logs."
+            ) from self.async_worker_error
+        return num_pending == 0
 
     def _allreduce_list(self, tensor_list: list[torch.Tensor]) -> list[torch.Tensor]:
         """All-reduce a list of tensors."""
@@ -1314,24 +1366,22 @@ def compute_logical_maps(
         device=device,
     )
 
-    running_count = torch.zeros_like(logical_replica_count)
-    layer_indices = torch.arange(num_layers, device=device)
-    for phys_idx in range(num_physical):
-        # Logical expert at physical slot phys_idx for each layer
-        logical_expert_ids = physical_to_logical_map_view[:, phys_idx]  # [num_layers]
-
-        # Scale up will set the logical expert ids to -1 for all new physical experts.
-        # Only consider "valid" experts when setting up the logical_to_physical map.
-        valid_expert_mask = logical_expert_ids >= 0
-        if not valid_expert_mask.any():
-            continue
-        valid_layers = layer_indices[valid_expert_mask]
-        valid_experts = logical_expert_ids[valid_expert_mask]
-
-        # Use the current running count as the replica index, then increment it.
-        replica_idx = running_count[valid_layers, valid_experts]
-        logical_to_physical_map_out[valid_layers, valid_experts, replica_idx] = phys_idx
-        running_count[valid_layers, valid_experts] += 1
+    # Stable sorting preserves physical-slot order within each logical expert.
+    logical_ids, physical_slots = torch.sort(
+        physical_to_logical_map_view, dim=1, stable=True
+    )
+    valid_layers, sorted_positions = (logical_ids >= 0).nonzero(as_tuple=True)
+    valid_experts = logical_ids[valid_layers, sorted_positions]
+    expert_offsets = logical_replica_count.cumsum(dim=1) - logical_replica_count
+    num_inactive = (~valid_mask).sum(dim=1)
+    replica_indices = (
+        sorted_positions
+        - num_inactive[valid_layers]
+        - expert_offsets[valid_layers, valid_experts]
+    )
+    logical_to_physical_map_out[valid_layers, valid_experts, replica_indices] = (
+        physical_slots[valid_layers, sorted_positions].to(dtype)
+    )
 
     # If computing maps for a single layer, squeeze out the extra layer dimension
     # before returning
@@ -1341,12 +1391,9 @@ def compute_logical_maps(
 
 
 def _pad_out_tensor(src: torch.Tensor, dst: torch.Tensor) -> None:
-    src_padding = dst.shape[-1] - src.shape[-1]
-    assert src_padding >= 0
-    new_src = torch.nn.functional.pad(src, (0, src_padding), value=-1)
-    # The map is committed from the host once per layer per rearrangement.
-    with gpu_sync_allowed():
-        dst.copy_(new_src)
+    padded = torch.full(dst.shape, -1, dtype=dst.dtype, pin_memory=PIN_MEMORY)
+    padded[..., : src.shape[-1]] = src
+    dst.copy_(padded, non_blocking=True)
 
 
 def _commit_eplb_maps_for_layer(
diff --git a/vllm/distributed/eplb/eplb_utils.py b/vllm/distributed/eplb/eplb_utils.py
index 943f6d9bef06..2f21ed643947 100644
--- a/vllm/distributed/eplb/eplb_utils.py
+++ b/vllm/distributed/eplb/eplb_utils.py
@@ -57,8 +57,12 @@ def wait(self, stream: torch.Stream | None = None):
         Should only be called by the Async Eplb thread.
         """
         self._recorded.wait()
-        self._event.wait(stream)
-        self._recorded.clear()
+        try:
+            self._event.wait(stream)
+        finally:
+            # Clear even if the CUDA wait raises, so a faulted cycle does
+            # not leave the event armed for the next iteration.
+            self._recorded.clear()
 
     def record(self, stream: torch.Stream | None = None):
         """Unblocks the waiting thread after calling event.record().
diff --git a/vllm/distributed/eplb/rebalance_execute.py b/vllm/distributed/eplb/rebalance_execute.py
index 5479ef01df4e..e9d17e0f9626 100644
--- a/vllm/distributed/eplb/rebalance_execute.py
+++ b/vllm/distributed/eplb/rebalance_execute.py
@@ -171,6 +171,35 @@ def get_ep_ranks_with_experts_batch(
     return ranks_to_send_map, ranks_to_recv_map
 
 
+def _first_rows(local_expert_ids: np.ndarray) -> dict[int, int]:
+    """Map each expert to its first local row in ``local_expert_ids``."""
+    rows: dict[int, int] = {}
+    for row, expert_id in enumerate(local_expert_ids.tolist()):
+        if expert_id != -1:
+            rows.setdefault(expert_id, row)
+    return rows
+
+
+def _recv_masks(
+    old_local_expert_ids: np.ndarray, new_local_expert_ids: np.ndarray
+) -> tuple[np.ndarray, np.ndarray]:
+    """Rows keeping their expert, and rows filled without a remote receive."""
+    is_unchanged = old_local_expert_ids == new_local_expert_ids
+    can_recv_local = np.isin(new_local_expert_ids, old_local_expert_ids)
+    is_received_locally = is_unchanged | ((new_local_expert_ids != -1) & can_recv_local)
+    return is_unchanged, is_received_locally
+
+
+def _remote_recv_rows(
+    old_local_expert_ids: np.ndarray, new_local_expert_ids: np.ndarray
+) -> tuple[np.ndarray, np.ndarray]:
+    """Experts a rank receives from peers (sorted), and the first row needing each."""
+    _, is_received_locally = _recv_masks(old_local_expert_ids, new_local_expert_ids)
+    need_recv = ~is_received_locally & (new_local_expert_ids != -1)
+    experts, first = np.unique(new_local_expert_ids[need_recv], return_index=True)
+    return experts, np.flatnonzero(need_recv)[first]
+
+
 def _execute_migration_batch(
     batch: list[MigrationFlow],
     old_indices: np.ndarray,
@@ -181,10 +210,11 @@ def _execute_migration_batch(
     expert_weights_buffers: Sequence[torch.Tensor],
     communicator: EplbCommunicator,
     layer_idx: int,
+    recv_rows: list[dict[int, int]] | None,
 ) -> None:
     """Execute one batch of rank-pair-disjoint expert transfers."""
     # NIXL drains its transfer context in execute(), so each batch restores it.
-    communicator.set_transfer_context(old_indices, layer_idx)
+    communicator.set_transfer_context(old_indices, layer_idx, recv_rows)
     for flow in batch:
         if flow.src_rank == ep_rank:
             for expert_id in flow.expert_ids:
@@ -228,16 +258,15 @@ def _execute_migration_batches(
                 num_local_experts, old_indices, new_indices
             )
     base = ep_rank * num_local_experts
-    old_local = old_indices[base : base + num_local_experts]
-    new_local = new_indices[base : base + num_local_experts]
-    old_rows: dict[int, int] = {}
-    new_rows: dict[int, int] = {}
-    for row, expert_id in enumerate(old_local):
-        if expert_id != -1:
-            old_rows.setdefault(int(expert_id), row)
-    for row, expert_id in enumerate(new_local):
-        if expert_id != -1:
-            new_rows.setdefault(int(expert_id), row)
+    old_rows = _first_rows(old_indices[base : base + num_local_experts])
+    new_rows = _first_rows(new_indices[base : base + num_local_experts])
+    # Batched receives land in each expert's first new row (see new_rows).
+    recv_rows = None
+    if communicator.needs_recv_rows:
+        recv_rows = [
+            _first_rows(new_indices[b : b + num_local_experts])
+            for b in range(0, new_indices.shape[0], num_local_experts)
+        ]
     for batch in batches:
         _execute_migration_batch(
             batch=batch,
@@ -249,6 +278,7 @@ def _execute_migration_batches(
             expert_weights_buffers=expert_weights_buffers,
             communicator=communicator,
             layer_idx=layer_idx,
+            recv_rows=recv_rows,
         )
 
 
@@ -321,16 +351,8 @@ def move_to_buffer(
     old_local_expert_ids = old_indices[local_global]
     new_local_expert_ids = new_indices[local_global]
 
-    # Unchanged mask
-    is_unchanged = old_local_expert_ids == new_local_expert_ids
-
-    # Local receive eligibility
-    new_valid = new_local_expert_ids != -1
-    can_recv_local = np.isin(
-        new_local_expert_ids, old_local_expert_ids, assume_unique=False
-    )
-    is_received_locally = np.logical_or(
-        is_unchanged, np.logical_and(new_valid, can_recv_local)
+    is_unchanged, is_received_locally = _recv_masks(
+        old_local_expert_ids, new_local_expert_ids
     )
 
     # Send map: first src row per unique expert present locally in old mapping
@@ -347,17 +369,13 @@ def move_to_buffer(
         send_src_rows[:send_count] = src_rows
 
     # Recv map: primary dst per unique expert needed remotely
-    recv_count = 0
-    need_recv_mask = np.logical_and(~is_received_locally, new_valid)
-    if np.any(need_recv_mask):
-        desired_experts = new_local_expert_ids[need_recv_mask]
-        desired_dsts = local_rows[need_recv_mask]
-        uniq_recv_experts, uniq_indices = np.unique(desired_experts, return_index=True)
-        dst_rows = desired_dsts[uniq_indices]
-        recv_count = int(uniq_recv_experts.shape[0])
-        recv_expert_ids[:recv_count] = uniq_recv_experts
-        recv_dst_rows[:recv_count] = dst_rows
-        recv_primary_mask[dst_rows] = True
+    uniq_recv_experts, dst_rows = _remote_recv_rows(
+        old_local_expert_ids, new_local_expert_ids
+    )
+    recv_count = int(uniq_recv_experts.shape[0])
+    recv_expert_ids[:recv_count] = uniq_recv_experts
+    recv_dst_rows[:recv_count] = dst_rows
+    recv_primary_mask[dst_rows] = True
 
     eligible_local_buffer_mask = np.logical_and(~is_unchanged, is_received_locally)
 
@@ -398,7 +416,16 @@ def move_to_buffer(
         )
         return transfer_metadata
 
-    communicator.set_transfer_context(old_indices, layer_idx)
+    recv_rows: list[dict[int, int]] | None = None
+    if communicator.needs_recv_rows:
+        recv_rows = []
+        for start in range(0, old_indices.shape[0], num_local_experts):
+            experts, rows = _remote_recv_rows(
+                old_indices[start : start + num_local_experts],
+                new_indices[start : start + num_local_experts],
+            )
+            recv_rows.append(dict(zip(experts.tolist(), rows.tolist())))
+    communicator.set_transfer_context(old_indices, layer_idx, recv_rows)
 
     # 2. Post sends
     if send_count > 0:
diff --git a/vllm/distributed/kv_transfer/kv_connector/factory.py b/vllm/distributed/kv_transfer/kv_connector/factory.py
index 97064e06cdb8..c0f0a4e6600f 100644
--- a/vllm/distributed/kv_transfer/kv_connector/factory.py
+++ b/vllm/distributed/kv_transfer/kv_connector/factory.py
@@ -123,6 +123,11 @@ def get_connector_class(
                 logger.error(msg)
                 raise ValueError(msg)
         elif connector_name in cls._registry:
+            if connector_name == "NixlConnector":
+                mode = kv_transfer_config.get_from_extra_config("transfer_mode", "pull")
+                if mode not in ("pull", "push"):
+                    raise ValueError(f"Invalid NixlConnector transfer_mode: {mode!r}")
+                connector_name = f"Nixl{mode.capitalize()}Connector"
             connector_cls = cls._registry[connector_name]()
         else:
             raise ValueError(f"Unsupported connector type: {connector_name}")
@@ -232,6 +237,11 @@ def supports_hma_config(cls, kv_transfer_config: "KVTransferConfig") -> bool:
     "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.connector",
     "MooncakeStoreConnector",
 )
+KVConnectorFactory.register_connector(
+    "UMBPStoreConnector",
+    "vllm.distributed.kv_transfer.kv_connector.v1.umbp",
+    "UMBPStoreConnector",
+)
 KVConnectorFactory.register_connector(
     "FlexKVConnectorV1",
     "vllm.distributed.kv_transfer.kv_connector.v1.flexkv_connector",
diff --git a/vllm/distributed/kv_transfer/kv_connector/utils.py b/vllm/distributed/kv_transfer/kv_connector/utils.py
index 436666f3dc74..934546e0ac66 100644
--- a/vllm/distributed/kv_transfer/kv_connector/utils.py
+++ b/vllm/distributed/kv_transfer/kv_connector/utils.py
@@ -21,7 +21,7 @@
 from vllm.platforms import current_platform
 from vllm.utils.torch_utils import async_tensor_h2d
 from vllm.v1.attention.backend import AttentionBackend
-from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheConfig
+from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheConfig, iter_layer_specs
 from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput
 
 if TYPE_CHECKING:
@@ -397,8 +397,6 @@ def get_current_attn_backends_and_specs(
         for name, layer in layers.items():
             spec = specs.get(name, group.kv_cache_spec)
             if isinstance(spec, AttentionSpec) and spec.tokens_per_state > 1:
-                layout = vllm_config.cache_config.get_resolved_kv_cache_layout()
-                assert layout.is_block_outermost, "Compressed caches must be packed"
                 continue
             if not isinstance(spec, AttentionSpec):
                 spec = None
@@ -408,6 +406,26 @@ def get_current_attn_backends_and_specs(
     return [backend for backend, _ in pairs], [spec for _, spec in pairs]
 
 
+def check_compressed_kernel_split(
+    vllm_config: VllmConfig, kv_cache_config: KVCacheConfig, kernel_block_size: int
+) -> None:
+    """A layer-outer layout cannot split compressed pages into kernel blocks."""
+    if vllm_config.cache_config.get_resolved_kv_cache_layout().is_block_outermost:
+        return
+    for group in kv_cache_config.transfer_groups:
+        for spec in iter_layer_specs(group.kv_cache_spec):
+            if (
+                isinstance(spec, AttentionSpec)
+                and spec.tokens_per_state > 1
+                and spec.block_size % kernel_block_size == 0
+                and spec.block_size != kernel_block_size
+            ):
+                raise NotImplementedError(
+                    "Compressed KV caches cannot be split into kernel blocks "
+                    "with a layer-outermost KV cache layout."
+                )
+
+
 def get_current_attn_backend(
     vllm_config: VllmConfig, layer_names: list[str] | None = None
 ) -> type[AttentionBackend]:
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/connector.py
index 4f17f1e737e0..04c4e2547af6 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/connector.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/connector.py
@@ -98,10 +98,12 @@ def __init__(
         *,
         async_speculative: bool,
         draft_kv_lookahead: int = 0,
+        import_step_tokens: int = 0,
     ) -> None:
         self.coordinator: HiSparseCoordinator | None = None
         self.async_speculative = async_speculative
         self.draft_kv_lookahead = draft_kv_lookahead
+        self.import_step_tokens = import_step_tokens
         self.requests: dict[str, Request] = {}
 
     def bind_coordinator(self, coordinator: HiSparseCoordinator) -> None:
@@ -244,6 +246,13 @@ def __init__(
                     and speculative_config is not None
                 ),
                 draft_kv_lookahead=vllm_config.num_lookahead_tokens,
+                import_step_tokens=(
+                    1
+                    + vllm_config.num_speculative_tokens
+                    + vllm_config.num_lookahead_tokens
+                    if vllm_config.num_speculative_tokens
+                    else 0
+                ),
             )
             max_model_len = vllm_config.model_config.max_model_len
             steady_concurrency = get_hisparse_steady_state_concurrency(
@@ -265,7 +274,10 @@ def __init__(
     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)
+            get_hisparse_coordinator(
+                kv_cache_manager,
+                import_step_tokens=self.connector_scheduler.import_step_tokens,
+            )
         )
 
     @property
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py
index 6e3d3415c3ad..b1ff69553d58 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py
@@ -13,7 +13,6 @@
 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,
@@ -23,11 +22,13 @@
     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.backend import get_hisparse_backend
 from vllm.v1.hisparse.layout import HISPARSE_HOT_SUFFIX
 from vllm.v1.hisparse.runtime import (
     HiSparseCacheHandle,
+    HiSparseHostPool,
+    HiSparseSlotInvalidator,
     drop_hisparse_host_pool,
     release_pinned_state,
     update_hisparse_residency,
@@ -164,6 +165,7 @@ def _create_hisparse_host_events(
     is_host_writer: bool,
     device: torch.device,
 ) -> tuple[torch.Event, torch.Event]:
+    backend = get_hisparse_backend()
     if shared_host_region is None:
         return torch.Event(), torch.Event()
 
@@ -171,18 +173,18 @@ def _create_hisparse_host_events(
     ipc_handles = None
     if is_host_writer:
         events = (
-            torch.cuda.Event(interprocess=True),
-            torch.cuda.Event(interprocess=True),
+            backend.create_ipc_event(),
+            backend.create_ipc_event(),
         )
-        stream = current_stream()
+        stream = torch.accelerator.current_stream()
         for event in events:
             event.record(stream)
-        ipc_handles = tuple(event.ipc_handle() for event in events)
+        ipc_handles = tuple(backend.export_event(event) 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]),
+            backend.import_event(device, ipc_handles[0]),
+            backend.import_event(device, ipc_handles[1]),
         )
     return events
 
@@ -191,6 +193,7 @@ 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.backend = get_hisparse_backend()
         self.vllm_config = vllm_config
         self.kv_cache_config = kv_cache_config
         self._initialized = False
@@ -211,26 +214,20 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]) -> None:
         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] = {}
+        host_pools: dict[int, HiSparseHostPool] = {}
         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
+            assert cache.runtime.host_pool is not None
+            host_pools[id(cache.runtime.host_pool)] = cache.runtime.host_pool
         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())
-        )
+        if len(host_pools) != 1:
+            raise RuntimeError("HiSparse caches must share one host pool.")
+        host_pool = next(iter(host_pools.values()))
+        shared_host_region = host_pool.shared_region
+        pinned_host_pools = host_pool.pinned
         is_host_writer = _is_hisparse_host_writer(shared_host_region)
 
         resident = cache_handles[0].view
@@ -280,9 +277,7 @@ def initialize(
         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.dma_stream = torch.Stream(device=device) if self.is_host_writer else None
         self._slot_mapping_staging = None
         if self.is_host_writer:
             max_mirror_rows = (
@@ -339,6 +334,9 @@ def initialize(
             raise RuntimeError(
                 "HiSparse request-state mapping does not match max_num_seqs."
             )
+        self._invalidate_written_slots = HiSparseSlotInvalidator(
+            self.request_state_indices
+        )
         self.hot_backing = hot_backing
         self.residency: torch.Tensor | None = next(
             (cache.residency for cache in cache_handles if cache.residency is not None),
@@ -371,7 +369,7 @@ def set_request_state_indices(self, indices: torch.Tensor) -> None:
                 "HiSparse request-state mapping exceeds max_num_seqs: "
                 f"{indices.numel()} > {self.request_state_indices.numel()}."
             )
-        if torch.cuda.is_current_stream_capturing():
+        if self.backend.is_current_stream_capturing():
             return
         # Attention indexes persistent request state by input-batch row. Refresh
         # that indirection after every batch compaction or reorder.
@@ -417,7 +415,7 @@ def start_step(
         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)
+        torch.accelerator.current_stream().wait_event(previous_host_write_event)
         if self.shared_host_region is not None and not num_tokens:
             # Without a forward, no TP collective keeps rank 0 from recording
             # a shared event again before every rank has issued its wait.
@@ -499,9 +497,9 @@ def _stage_row_mirror_mapping(self, num_tokens: int) -> None:
                 "HiSparse row mapping exceeds staging capacity: "
                 f"{end} > {state.slots.shape[0]}."
             )
-        main_stream = current_stream()
+        main_stream = torch.accelerator.current_stream()
         state.stream.wait_stream(main_stream)
-        with torch.cuda.stream(state.stream):
+        with state.stream:
             state.slots[start:end].copy_(
                 slots[:num_tokens],
                 non_blocking=True,
@@ -579,9 +577,9 @@ def get_kv_connector_stats(self) -> HiSparseKVConnectorStats | None:
         if (
             self._metrics_calls % _METRICS_INTERVAL == 0
             and not self._metrics_pending
-            and not torch.cuda.is_current_stream_capturing()
+            and not self.backend.is_current_stream_capturing()
         ):
-            compute_stream = current_stream()
+            compute_stream = torch.accelerator.current_stream()
             for runtime in self.leader_runtimes:
                 group = runtime.index_group
                 compute_stream.wait_stream(group.copy_stream)
@@ -614,12 +612,12 @@ def _submit_dma_descriptors(
         stream = self.dma_stream
         assert stream is not None
         if ready_event is None:
-            stream.wait_stream(current_stream())
+            stream.wait_stream(torch.accelerator.current_stream())
         else:
             stream.wait_event(ready_event)
         completion_event = torch.Event()
-        with torch.cuda.stream(stream):
-            ops.swap_blocks_batch(
+        with stream:
+            self.backend.swap_blocks_batch(
                 descriptors.src[:descriptor_count],
                 descriptors.dst[:descriptor_count],
                 descriptors.sizes[:descriptor_count],
@@ -785,7 +783,9 @@ def _restore_pages(self, transfers: list[SparseKVPageTransfer]) -> None:
                         source[host_start : host_start + self.kernel_block_size],
                         non_blocking=True,
                     )
-        self._record_transfer_completion(transfers, stream=current_stream())
+        self._record_transfer_completion(
+            transfers, stream=torch.accelerator.current_stream()
+        )
 
     def _submit_transfers(self, transfers: list[SparseKVPageTransfer]) -> None:
         if self.cache_handles[0].runtime.eager_host_mirror:
@@ -897,12 +897,11 @@ def mirror_key(handle: HiSparseCacheHandle) -> tuple:
                 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],
-                )
+        self._invalidate_written_slots(
+            [handle.runtime for _, handle in active if handle.runtime.is_group_leader],
+            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
@@ -915,7 +914,7 @@ def _finish_mirror_phase(self, ready_event: torch.Event | None = None) -> None:
         self._clear_forward_mirror_state()
 
     def finish_forward(self) -> None:
-        compute_stream = current_stream()
+        compute_stream = torch.accelerator.current_stream()
         self._forward_ready_event.record()
         self._finish_mirror_phase(self._forward_ready_event)
         if self.is_host_writer:
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 ed887f296bd5..54eba531c9dd 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
@@ -415,7 +415,7 @@ def _calculate_mtp_layers(vllm_config, model_config):
                 model_config.hf_config, "num_nextn_predict_layers", 0
             )
 
-        elif vllm_config.speculative_config.use_eagle():
+        elif vllm_config.speculative_config.uses_target_hidden_states:
             try:
                 draft_model_config = vllm_config.speculative_config.draft_model_config
                 num_mtp_layers = draft_model_config.get_num_layers(
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 277b4832dc10..8fa9b69f55e0 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
@@ -27,6 +27,7 @@
 from vllm.distributed.kv_transfer.kv_connector.utils import (
     EngineId,
     TransferTopology,
+    check_compressed_kernel_split,
     get_current_attn_backends,
     get_current_attn_backends_and_specs,
 )
@@ -77,6 +78,7 @@
     SlidingWindowSpec,
 )
 from vllm.v1.metrics.cache_hit_source import CacheHitSource
+from vllm.v1.outputs import KVConnectorOutput
 from vllm.v1.request import RequestStatus
 from vllm.v1.worker.utils import select_common_block_size
 
@@ -676,6 +678,38 @@ class MooncakeXferResponse(
     layout_id: int = 0
 
 
+class MooncakeLeaseHeartbeat(
+    msgspec.Struct,
+    omit_defaults=True,  # type: ignore[call-arg]
+):
+    """Sent by D to renew P's lease on KV blocks that D has not pulled yet.
+
+    It shares P's side channel with pulls. A pull decodes as a heartbeat with
+    no transfer ids, because msgspec skips the fields it does not know.
+    """
+
+    lease_transfer_ids: list[TransferId] = msgspec.field(default_factory=list)
+
+
+@dataclass
+class LeaseHeartbeat:
+    """Transfers that one P engine holds for this D engine under a lease."""
+
+    remote_bootstrap_addr: str
+    transfer_ids: list[TransferId]
+
+
+def _get_kv_lease_duration(kv_transfer_config: KVTransferConfig) -> float | None:
+    """Return `kv_lease_duration` in seconds, or None when leases are off."""
+    duration = kv_transfer_config.get_from_extra_config("kv_lease_duration", None)
+    if duration is None:
+        return None
+    duration = float(duration)
+    if duration <= 0:
+        raise ValueError(f"kv_lease_duration must be positive, got {duration}.")
+    return duration
+
+
 @dataclass
 class PullReqMeta:
     d_req_id: ReqId
@@ -711,6 +745,8 @@ def __init__(self):
         self.reqs_to_recv: dict[EngineId, dict[ReqId, PullReqMeta]] = defaultdict(dict)
         self.reqs_to_send: dict[ReqId, tuple[TransferId, list[list[int]]]] = {}
         self.reqs_not_processed: set[TransferId] = set()
+        # D side: lease renewals to send this step, by remote P engine.
+        self.lease_heartbeats: dict[EngineId, LeaseHeartbeat] = {}
 
     def add_new_req(
         self,
@@ -736,6 +772,14 @@ def add_new_req(
 class MooncakeConnector(KVConnectorBase_V1, SupportsHMA):
     _cache_hit_source = CacheHitSource.P2P
 
+    @property
+    def supports_divergent_local_hybrid_hits(self) -> bool:
+        # A decode pulls the unhashed blocks of every transfer group, so the
+        # recurrent-state slot arrives even when the attention prefix is local.
+        # Roles that prefill locally keep the lookup that finds the shared
+        # prefix boundary.
+        return self._pulls_mamba_state
+
     def __init__(
         self,
         vllm_config: VllmConfig,
@@ -747,6 +791,13 @@ def __init__(
         assert vllm_config.kv_transfer_config is not None
         assert vllm_config.kv_transfer_config.engine_id is not None
         self.engine_id: EngineId = vllm_config.kv_transfer_config.engine_id
+        self._pulls_mamba_state = (
+            vllm_config.kv_transfer_config.kv_role == "kv_consumer"
+            and any(
+                isinstance(group.kv_cache_spec, MambaSpec)
+                for group in kv_cache_config.transfer_groups
+            )
+        )
 
         if role == KVConnectorRole.SCHEDULER:
             assert kv_cache_config is not None, (
@@ -814,6 +865,10 @@ def on_new_request(self, request: "Request") -> None:
         assert self.connector_scheduler is not None
         self.connector_scheduler.on_new_request(request)
 
+    def update_connector_output(self, connector_output: KVConnectorOutput):
+        assert self.connector_scheduler is not None
+        self.connector_scheduler.update_connector_output(connector_output)
+
     def request_finished(
         self,
         request: "Request",
@@ -980,6 +1035,13 @@ def __init__(
         # remote prefill or aborted.
         self._reqs_not_processed: set[TransferId] = set()
 
+        # D side: requests whose KV P still holds under a lease, from arrival
+        # until the pull settles. A request can wait in the queue for longer
+        # than the lease, so renewal cannot wait until it is scheduled.
+        self._kv_lease_duration = _get_kv_lease_duration(vllm_config.kv_transfer_config)
+        self._lease_reqs: dict[ReqId, tuple[EngineId, str, TransferId]] = {}
+        self._last_heartbeat_time = 0.0
+
         # Compute sliding window block counts per KV cache group.
         sw_sizes_tokens: list[tuple[int, int]] = [
             (g.kv_cache_spec.sliding_window, g.kv_cache_spec.block_size)
@@ -1058,6 +1120,26 @@ def on_new_request(self, request: "Request") -> None:
         ):
             self._truncate_mamba_request_for_prefill(request)
 
+        if (
+            self._kv_lease_duration is not None
+            and params is not None
+            and params.get("do_remote_prefill")
+            and all(
+                params.get(p)
+                for p in ("remote_engine_id", "remote_bootstrap_addr", "transfer_id")
+            )
+        ):
+            self._lease_reqs[request.request_id] = (
+                params["remote_engine_id"],
+                params["remote_bootstrap_addr"],
+                params["transfer_id"],
+            )
+
+    def update_connector_output(self, connector_output: KVConnectorOutput) -> None:
+        # Failed pulls are reported here too.
+        for req_id in connector_output.finished_recving or ():
+            self._lease_reqs.pop(req_id, None)
+
     def get_num_new_matched_tokens(
         self, request: "Request", num_computed_tokens: int
     ) -> tuple[int, bool]:
@@ -1131,6 +1213,10 @@ def update_state_after_alloc(
                 local_block_ids = self.get_sw_clipped_blocks(unhashed_block_ids)
                 # Get unhashed blocks to pull from remote.
                 self._reqs_need_recv[request.request_id] = (request, local_block_ids)
+                if not any(local_block_ids):
+                    # This pull only releases P's blocks and never reaches
+                    # finished_recving, so stop renewing the lease now.
+                    self._lease_reqs.pop(request.request_id, None)
             else:
                 logger.warning(
                     "Got invalid KVTransferParams: %s. This "
@@ -1165,6 +1251,23 @@ def build_connector_meta(
                 )
             self._reqs_need_recv.clear()
 
+            # Renew every lease in one message per P engine, every
+            # kv_lease_duration / 6 seconds (as NixlConnector does).
+            now = time.perf_counter()
+            if (
+                self._lease_reqs
+                and self._kv_lease_duration is not None
+                and now - self._last_heartbeat_time >= self._kv_lease_duration / 6
+            ):
+                self._last_heartbeat_time = now
+                for engine_id, bootstrap_addr, transfer_id in self._lease_reqs.values():
+                    heartbeat = meta.lease_heartbeats.get(engine_id)
+                    if heartbeat is None:
+                        heartbeat = meta.lease_heartbeats[engine_id] = LeaseHeartbeat(
+                            remote_bootstrap_addr=bootstrap_addr, transfer_ids=[]
+                        )
+                    heartbeat.transfer_ids.append(transfer_id)
+
         if not self.is_kv_consumer:
             for req_id, (req, block_ids) in self._reqs_need_send.items():
                 assert req.kv_transfer_params is not None
@@ -1196,6 +1299,7 @@ def request_finished(
             request.status,
             params,
         )
+        self._lease_reqs.pop(request.request_id, None)
         if not params or not params.get("transfer_id"):
             return False, None
 
@@ -1250,6 +1354,19 @@ def __init__(
         logger.info("Initializing Mooncake Transfer Engine worker %s", engine_id)
 
         self.vllm_config = vllm_config
+        # nvidia-peermem cannot pin VMM-backed memory; the transfer engine must
+        # register the extensible KV cache through DMA-BUF. Read at
+        # engine.initialize(), so settle it before constructing the engine.
+        if (
+            vllm_config.cache_config.enable_extensible_kv_cache
+            and os.environ.setdefault("WITH_NVIDIA_PEERMEM", "0") != "0"
+        ):
+            raise RuntimeError(
+                "WITH_NVIDIA_PEERMEM=1 is incompatible with the extensible KV "
+                "cache: nvidia-peermem cannot register VMM-backed memory. "
+                "Unset it to use Mooncake's DMA-BUF registration, or disable "
+                "the extensible KV cache."
+            )
         # Capture device BEFORE TransferEngine init — MNNVL's NVLink allocator
         # may change the current CUDA device during engine.initialize().
         self.device_id = torch.accelerator.current_device_index()
@@ -1335,6 +1452,14 @@ def __init__(
         self.kv_caches_base_addr: list[int] = []
         self.device_kv_caches: dict[str, torch.Tensor] = {}
         self.reqs_need_send: dict[TransferId, SendBlockMeta] = {}
+        # P holds a finished request's blocks for kv_lease_duration seconds,
+        # and each renewal from D moves the deadline to at least 2/3 of that
+        # from now (as NixlConnector does). Unset keeps the fixed
+        # VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT.
+        self._kv_lease_duration = _get_kv_lease_duration(kv_transfer_config)
+        # Expired transfers in expiry order, so a late pull fails at once
+        # instead of waiting for a prefill that already happened.
+        self._expired_transfer_ids: OrderedDict[TransferId, float] = OrderedDict()
 
         # For kv_both, we will act both prefiller and decoder.
         if not self.is_kv_consumer:
@@ -1422,6 +1547,7 @@ def __init__(
         self._encoder = msgspec.msgpack.Encoder()
         self._xfer_meta_decoder = msgspec.msgpack.Decoder(MooncakeXferMetadata)
         self._xfer_resp_decoder = msgspec.msgpack.Decoder(MooncakeXferResponse)
+        self._heartbeat_decoder = msgspec.msgpack.Decoder(MooncakeLeaseHeartbeat)
 
     def _sync_block_size_with_kernel(self) -> None:
         # When speculative decoding (e.g. Eagle) is enabled, the main model
@@ -1433,6 +1559,9 @@ def _sync_block_size_with_kernel(self) -> None:
             self.vllm_config, self.kv_cache_config, backends
         )
         kernel_block_size = select_common_block_size(self.block_size, backends, specs)
+        check_compressed_kernel_split(
+            self.vllm_config, self.kv_cache_config, kernel_block_size
+        )
         if self.block_size != kernel_block_size:
             logger.info_once(
                 "User-specified logical block size (%s) does not match"
@@ -1540,6 +1669,15 @@ async def _mooncake_sender_listener(self, ready_event: threading.Event):
         try:
             while True:
                 identity, metadata_bytes = await sock.recv_multipart()
+                # Renew here, not in a sender worker: those can all be
+                # waiting for prefills, and the renewal must not wait.
+                heartbeat = self._decode_heartbeat(metadata_bytes)
+                if heartbeat is not None:
+                    response = self._renew_kv_leases(heartbeat.lease_transfer_ids)
+                    await sock.send_multipart(
+                        (identity, self._encoder.encode(response))
+                    )
+                    continue
                 await self.sender_worker_queue.put((identity, metadata_bytes))
         except zmq.ContextTerminated:
             logger.debug("ZMQ context terminated, exiting Mooncake sender thread.")
@@ -1552,6 +1690,29 @@ async def _mooncake_sender_listener(self, ready_event: threading.Event):
             await asyncio.gather(*sender_tasks, return_exceptions=True)
             sock.close()
 
+    def _decode_heartbeat(self, message: bytes) -> MooncakeLeaseHeartbeat | None:
+        try:
+            heartbeat = self._heartbeat_decoder.decode(message)
+        except msgspec.DecodeError:
+            # Leave malformed messages to the sender worker, which reports them.
+            return None
+        return heartbeat if heartbeat.lease_transfer_ids else None
+
+    def _renew_kv_leases(self, transfer_ids: list[TransferId]) -> MooncakeXferResponse:
+        if self._kv_lease_duration is None:
+            return MooncakeXferResponse(
+                status=MooncakeXferResponseStatus.ERROR,
+                err_msg="kv_lease_duration is not set on this prefill instance.",
+            )
+        new_expire_time = time.perf_counter() + self._kv_lease_duration * 2 / 3
+        for transfer_id in transfer_ids:
+            # Skip transfers not prefilled yet or already freed; a renewal
+            # must not recreate them.
+            send_meta = self.reqs_need_send.get(transfer_id)
+            if send_meta is not None:
+                send_meta.expire_time = max(send_meta.expire_time, new_expire_time)
+        return MooncakeXferResponse(status=MooncakeXferResponseStatus.FINISH)
+
     async def _sender_worker(self, sock: zmq.asyncio.Socket):
         while True:
             try:
@@ -1656,8 +1817,12 @@ async def send_kv_to_decode(
             )
             await sock.send_multipart((identity, self._encoder.encode(response)))
             return
+        expired_reqs: list[ReqId] = []
         for d_req_id, (transfer_id, _) in meta.req_blocks.items():
             if transfer_id not in self.reqs_need_send:
+                if transfer_id in self._expired_transfer_ids:
+                    expired_reqs.append(d_req_id)
+                    continue
                 # This req is not enqueued in P side yet, create it here.
                 self.reqs_need_send[transfer_id] = SendBlockMeta(
                     p_req_id="",
@@ -1668,6 +1833,21 @@ async def send_kv_to_decode(
             send_meta = self.reqs_need_send[transfer_id]
             pending_reqs[d_req_id] = send_meta
 
+        if expired_reqs:
+            logger.warning(
+                "Requests %s were pulled after their KV expired on P side.",
+                expired_reqs,
+            )
+            response = MooncakeXferResponse(
+                status=MooncakeXferResponseStatus.CONTINUE
+                if pending_reqs
+                else MooncakeXferResponseStatus.FINISH,
+                err_reqs=expired_reqs,
+                err_msg="KV expired on P side before the pull.",
+                layout_id=meta.layout_id,
+            )
+            await sock.send_multipart((identity, self._encoder.encode(response)))
+
         async def wait_and_ret(
             d_req_id: ReqId, send_meta: SendBlockMeta
         ) -> tuple[ReqId, SendBlockMeta]:
@@ -2355,24 +2535,41 @@ async def fetch_finished_sending_reqs(self) -> set[ReqId]:
         now = time.perf_counter()
 
         expired_transfer_id = []
+        # Scan every entry: renewals move deadlines, so the map is not in
+        # expiry order.
         for transfer_id, send_meta in self.reqs_need_send.items():
             if (
                 send_meta.p_req_id
                 and send_meta.expire_time < now
                 and send_meta.sending == 0
             ):
-                logger.warning(
-                    "Request %s timed out after %d seconds without "
-                    "being sent. Freeing its blocks on the producer side.",
-                    send_meta.p_req_id,
-                    envs.VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT,
-                )
+                if self._kv_lease_duration is None:
+                    logger.warning(
+                        "Request %s timed out after %d seconds without "
+                        "being sent. Freeing its blocks on the producer side.",
+                        send_meta.p_req_id,
+                        envs.VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT,
+                    )
+                else:
+                    logger.warning(
+                        "KV lease of request %s expired before the decoder "
+                        "pulled it. Freeing its blocks on the producer side.",
+                        send_meta.p_req_id,
+                    )
                 self.xfer_stats.record_kv_expired_req()
                 finished_sending_reqs.add(send_meta.p_req_id)
                 expired_transfer_id.append(transfer_id)
 
         for transfer_id in expired_transfer_id:
             del self.reqs_need_send[transfer_id]
+            self._expired_transfer_ids[transfer_id] = now
+        # A pull later than this falls back to the ready-wait timeout.
+        horizon = now - envs.VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT
+        while self._expired_transfer_ids:
+            transfer_id, expired_at = next(iter(self._expired_transfer_ids.items()))
+            if expired_at >= horizon:
+                break
+            del self._expired_transfer_ids[transfer_id]
 
         return finished_sending_reqs
 
@@ -2643,31 +2840,46 @@ async def _connect_to_prefiller_bootstrap(self, remote_bootstrap_addr: str):
         self._pending_bootstrap_queries[remote_bootstrap_addr].set()
         del self._pending_bootstrap_queries[remote_bootstrap_addr]
 
-    def receive_kv(
-        self,
-        remote_engine_id: EngineId,
-        pull_metas: dict[ReqId, PullReqMeta],
-    ):
+    def _select_remote_workers(
+        self, remote_engine_id: EngineId
+    ) -> dict[int, list[int]]:
+        """Return the producer PP ranks to pull from, by producer TP rank."""
         remote_tp_ranks = self.transfer_topo.handshake_target_ranks(
             self._tp_size[remote_engine_id]
         )
-        worker_addrs: list[str] = []
         selected_remote_pp: dict[int, list[int]] = {}
         for remote_tp_rank in remote_tp_ranks:
             pp_to_addr = self._remote_agents[remote_engine_id][remote_tp_rank]
             if self.pp_size == len(pp_to_addr) and self.pp_rank in pp_to_addr:
-                pp_ranks = [self.pp_rank]
+                selected_remote_pp[remote_tp_rank] = [self.pp_rank]
             else:
-                pp_ranks = sorted(pp_to_addr)
-            selected_remote_pp[remote_tp_rank] = pp_ranks
-            worker_addrs.extend(pp_to_addr[pp_rank] for pp_rank in pp_ranks)
+                selected_remote_pp[remote_tp_rank] = sorted(pp_to_addr)
+        return selected_remote_pp
+
+    def _remote_worker_addrs(
+        self, remote_engine_id: EngineId, selected_remote_pp: dict[int, list[int]]
+    ) -> list[str]:
+        agents = self._remote_agents[remote_engine_id]
+        return [
+            agents[tp_rank][pp_rank]
+            for tp_rank, pp_ranks in selected_remote_pp.items()
+            for pp_rank in pp_ranks
+        ]
+
+    def receive_kv(
+        self,
+        remote_engine_id: EngineId,
+        pull_metas: dict[ReqId, PullReqMeta],
+    ):
+        selected_remote_pp = self._select_remote_workers(remote_engine_id)
+        worker_addrs = self._remote_worker_addrs(remote_engine_id, selected_remote_pp)
 
         count = len(worker_addrs)
         logger.debug(
             "Receiving Mooncake KV for engine %s from producer TP ranks %s "
             "and PP ranks %s",
             remote_engine_id,
-            remote_tp_ranks,
+            list(selected_remote_pp),
             selected_remote_pp,
         )
         for pull_meta in pull_metas.values():
@@ -2677,17 +2889,76 @@ def receive_kv(
                 self.receive_kv_from_single_worker(worker_addr, pull_metas)
             )
 
+    async def _query_bootstrap(self, remote_bootstrap_addr: str):
+        """Query a bootstrap server, or wait for the query already running."""
+        if remote_bootstrap_addr not in self._pending_bootstrap_queries:
+            self._pending_bootstrap_queries[remote_bootstrap_addr] = asyncio.Event()
+            await self._connect_to_prefiller_bootstrap(remote_bootstrap_addr)
+        else:
+            await self._pending_bootstrap_queries[remote_bootstrap_addr].wait()
+
+    async def _send_lease_heartbeat(
+        self, remote_engine_id: EngineId, heartbeat: LeaseHeartbeat
+    ):
+        """Renew the leases one P engine holds for this worker's requests."""
+        if remote_engine_id not in self._remote_agents:
+            # The requests may still be queued. Query now so this renewal,
+            # and the pull later, can reach P's workers.
+            await self._query_bootstrap(heartbeat.remote_bootstrap_addr)
+            if remote_engine_id not in self._remote_agents:
+                logger.warning(
+                    "Skipping KV lease renewal for %d requests: remote "
+                    "engine_id %s not found from bootstrap server %s.",
+                    len(heartbeat.transfer_ids),
+                    remote_engine_id,
+                    heartbeat.remote_bootstrap_addr,
+                )
+                return
+        payload = self._encoder.encode(
+            MooncakeLeaseHeartbeat(lease_transfer_ids=heartbeat.transfer_ids)
+        )
+        selected_remote_pp = self._select_remote_workers(remote_engine_id)
+        await asyncio.gather(
+            *(
+                self._send_heartbeat_to_worker(addr, payload)
+                for addr in self._remote_worker_addrs(
+                    remote_engine_id, selected_remote_pp
+                )
+            )
+        )
+
+    async def _send_heartbeat_to_worker(self, worker_addr: str, payload: bytes):
+        assert self._kv_lease_duration is not None
+        try:
+            with make_zmq_socket(
+                self.async_zmq_ctx, worker_addr, zmq.DEALER, bind=False, linger=0
+            ) as sock:
+                # Wait for the reply: linger=0 drops a message still unsent.
+                sock.setsockopt(zmq.RCVTIMEO, int(self._kv_lease_duration / 6 * 1000))
+                await sock.send(payload)
+                response = self._xfer_resp_decoder.decode(await sock.recv())
+        except zmq.ContextTerminated:
+            return
+        except Exception as e:
+            # The pull reports a P that is really gone.
+            logger.debug("KV lease renewal to %s failed: %s", worker_addr, e)
+            return
+        if response.status == MooncakeXferResponseStatus.ERROR:
+            logger.warning_once(
+                "Mooncake prefill worker %s rejected a KV lease renewal (%s). "
+                "Set the same kv_lease_duration on prefill and decode "
+                "instances.",
+                worker_addr,
+                response.err_msg,
+            )
+
     async def handle_new_engine_id(
         self,
         remote_engine_id: EngineId,
         pull_metas: dict[ReqId, PullReqMeta],
     ):
         remote_bootstrap_addr = next(iter(pull_metas.values())).remote_bootstrap_addr
-        if remote_bootstrap_addr not in self._pending_bootstrap_queries:
-            self._pending_bootstrap_queries[remote_bootstrap_addr] = asyncio.Event()
-            await self._connect_to_prefiller_bootstrap(remote_bootstrap_addr)
-        else:
-            await self._pending_bootstrap_queries[remote_bootstrap_addr].wait()
+        await self._query_bootstrap(remote_bootstrap_addr)
 
         if remote_engine_id not in self._remote_agents:
             self._handle_failed_recv(
@@ -2718,8 +2989,8 @@ async def record_send_reqs(self, metadata: MooncakeConnectorMetadata):
                 send_meta = self.reqs_need_send[transfer_id]
                 send_meta.p_req_id = p_req_id
                 send_meta.local_block_ids = block_ids
-                send_meta.expire_time = (
-                    time.perf_counter() + envs.VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT
+                send_meta.expire_time = time.perf_counter() + (
+                    self._kv_lease_duration or envs.VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT
                 )
                 send_meta.ready.set()
             else:
@@ -2739,12 +3010,24 @@ async def record_send_reqs(self, metadata: MooncakeConnectorMetadata):
             if send_meta:
                 assert not send_meta.ready.is_set()
 
+    async def _send_lease_heartbeats(
+        self, lease_heartbeats: dict[EngineId, LeaseHeartbeat]
+    ):
+        for remote_engine_id, heartbeat in lease_heartbeats.items():
+            asyncio.create_task(self._send_lease_heartbeat(remote_engine_id, heartbeat))
+
     def start_load_kv(self, metadata: MooncakeConnectorMetadata):
         if not self.is_kv_producer and metadata.reqs_to_recv:
             asyncio.run_coroutine_threadsafe(
                 self._start_load_kv(metadata.reqs_to_recv), self.receiver_loop
             )
 
+        if not self.is_kv_producer and metadata.lease_heartbeats:
+            asyncio.run_coroutine_threadsafe(
+                self._send_lease_heartbeats(metadata.lease_heartbeats),
+                self.receiver_loop,
+            )
+
         if not self.is_kv_consumer and (
             metadata.reqs_to_send or metadata.reqs_not_processed
         ):
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 a1e28b97d06a..a91167234da4 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
@@ -11,6 +11,7 @@
 enabling prefix caching via hash-based deduplication.
 """
 
+import os
 from collections import Counter
 from collections.abc import Iterable, Sequence
 from contextlib import AbstractContextManager
@@ -38,6 +39,7 @@
     PromMetric,
     PromMetricT,
 )
+from vllm.distributed.mooncake_store import MooncakeStoreConfig
 from vllm.forward_context import ForwardContext
 from vllm.logger import init_logger
 from vllm.v1.attention.backend import AttentionMetadata
@@ -129,7 +131,20 @@ class MooncakeStoreConnector(KVConnectorBase_V1, SupportsHMA):
 
     @classmethod
     def supports_sleep_mode(cls, kv_transfer_config: KVTransferConfig) -> bool:
-        return False
+        config = MooncakeStoreConfig.load_from_config()
+        if (
+            config.protocol != "rdma"
+            or not config.device_name
+            or os.getenv("MC_MS_AUTO_DISC", "0") != "0"
+            or any(v in os.environ for v in ("MC_USE_TENT", "MC_USE_TEV1"))
+        ):
+            raise ValueError(
+                "MooncakeStoreConnector sleep requires protocol='rdma', "
+                "device_name set, MC_MS_AUTO_DISC=0 or unset, and "
+                "MC_USE_TENT/MC_USE_TEV1 unset. "
+                "Sleep with NVLink Fabric is not supported yet."
+            )
+        return True
 
     @staticmethod
     def _validate_kv_cache_config(
@@ -342,6 +357,14 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
         assert self.connector_worker is not None
         self.connector_worker.register_kv_caches(kv_caches)
 
+    def unregister_kv_caches(self) -> None:
+        assert self.connector_worker is not None
+        self.connector_worker.unregister_kv_caches()
+
+    def restore_kv_caches(self) -> None:
+        assert self.connector_worker is not None
+        self.connector_worker.restore_kv_caches()
+
     def start_load_kv(self, forward_context: ForwardContext, **kwargs: Any) -> None:
         assert self.connector_worker is not None
         metadata = self._get_connector_metadata()
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 c893478cafbf..a53aaf2f6442 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
@@ -6,9 +6,6 @@
 from dataclasses import replace
 from typing import TYPE_CHECKING, 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 replay_boundaries
@@ -93,10 +90,7 @@ def __init__(
         retention_interval: int | None = None,
         enable_partial_hash_hits: bool = False,
     ) -> None:
-        # Mirrors core's resolve_kv_cache_block_sizes: the hash unit only has
-        # to divide groups that participate in prefix caching. Non-shareable
-        # scratch groups (e.g. GLM-5.3-Flash's kpool tail) are skipped by
-        # _verify_and_split_kv_cache_groups and never probed for hits.
+        # Match core's alignment requirements for prefix-cacheable groups.
         assert all(
             g.kv_cache_spec.block_size % hash_block_size == 0
             for g in kv_cache_groups
@@ -173,15 +167,9 @@ def align_lookup_length(self, length: int) -> int:
         return length // alignment * alignment
 
     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).
-        """
+        """Group compatible specs using their registered cache managers."""
         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
-            # not participate in hit lookup. Mirrors core's
-            # KVCacheCoordinator.verify_and_split_kv_cache_groups.
             if not g.kv_cache_spec.prefix_cacheable:
                 continue
             spec = _coverage_spec(g.kv_cache_spec)
@@ -206,10 +194,7 @@ def _verify_and_split_kv_cache_groups(self) -> None:
         if self.use_eagle and not any(g.use_eagle for g in attention_groups):
             attention_groups = [g._replace(use_eagle=True) for g in attention_groups]
         self.attention_groups = attention_groups
-        # Per-group eagle bits. SpecGroup carries use_eagle for the whole
-        # merged spec group, so the per-group store/lookup masks agree with
-        # the merged-group hit check, which applies the eagle drop to every
-        # group sharing the spec.
+        # Apply the merged spec group's Eagle setting to each member's masks.
         self.eagle_group_ids = {
             gid for g in attention_groups if g.use_eagle for gid in g.group_ids
         }
@@ -234,12 +219,12 @@ def find_longest_cache_hit(
         apply_eagle: bool = True,
     ) -> tuple[tuple[list[bool], ...], int]:
         """Returns ``(load_mask_per_group, hit_length)``. ``mask[g][i]`` is True iff
-        group ``g`` populates chunk ``i`` locally (e.g. SWA and Mamba tail-only);
+        group ``g`` populates physical block ``i`` (e.g. SWA and Mamba tail-only);
         recv-side callers skip False slots.
 
         ``apply_eagle`` controls whether the per-spec ``use_eagle`` last-block
         pop is applied. Lookup callers want it (the drafter requires recomputing
-        the last block); per-chunk mask callers must not, because ``token_len``
+        the last block); load-mask callers must not, because ``token_len``
         already reflects the eagle-pruned hit length and a second pop would
         leave the trailing block unloaded.
         """
@@ -258,15 +243,10 @@ def load_mask(
         token_len: int,
     ) -> tuple[list[bool], ...]:
         """Per-group load masks: ``mask[g][i]`` is True iff group ``g``'s
-        spec would populate chunk ``i`` locally at length ``token_len``
+        spec would populate physical block ``i`` at length ``token_len``
         (e.g. SWA / Mamba tail-only).
         """
-        # ``apply_eagle=False`` because ``token_len`` is already the
-        # eagle-pruned hit length returned by ``client.lookup``. Re-applying
-        # the pop here would shorten the mask by one extra block; the recv
-        # thread would then silently skip the trailing chunk yielded by
-        # ``db.process_tokens`` and leave that block uninitialized in the
-        # local KV pool.
+        # Lookup has already applied Eagle's last-block drop.
         masks, _ = self.find_longest_cache_hit(
             block_hashes,
             token_len,
@@ -283,22 +263,15 @@ def store_mask(
     ) -> tuple[list[bool] | None, ...]:
         """Per-group store masks for the suffix starting at ``start_token``.
 
-        ``mask[g][i]`` is True iff the i-th chunk of group ``g`` *after*
+        ``mask[g][i]`` is True iff the i-th physical block of group ``g`` *after*
         ``start_token`` should be written to the store so a future cache hit
         can consume it. ``None`` is the all-True sentinel for the suffix.
 
         Reuses the engine's ``SingleTypeKVCacheManager.reachable_block_mask``
         so the store retains exactly the blocks the local prefix cache would.
 
-        Mamba groups are always all-False: the normal save resolves blocks
-        positionally from the connector's append-only block-ID snapshot, but
-        an align-mode mamba block table is not append-only (interior state
-        blocks are nulled/freed, and speculative decoding relocates the spec
-        blocks in place), so a positional read may hit a null, freed, or live
-        speculative-state block and persist wrong bytes under a valid prefix
-        hash. Mamba state is persisted only through the connector-pinned exact
-        block hand-off path
-        (``SchedulerOutput.kv_connector_block_state.boundary_state_offloads``).
+        Mamba groups are all-False; checkpoint block IDs arrive through
+        ``SchedulerOutput.kv_connector_block_state.boundary_state_offloads``.
         """
         return self._reachable_masks(
             aligned_token_len,
@@ -314,7 +287,7 @@ def lookup_mask(
     ) -> tuple[list[bool] | None, ...]:
         """Per-group lookup masks.
 
-        ``mask[g][i]`` is True iff chunk ``i`` of group ``g`` should be
+        ``mask[g][i]`` is True iff physical block ``i`` of group ``g`` should be
         looked up as an aligned hit boundary. ``None`` is the all-True
         sentinel.
         """
@@ -355,39 +328,30 @@ def _reachable_masks(
         masks: list[list[bool] | None] = []
         for g_idx, g in enumerate(self.kv_cache_groups):
             spec = _coverage_spec(g.kv_cache_spec)
-            end_chunk = aligned_token_len // spec.block_size
-            start_chunk = min(end_chunk, max(0, cdiv(start_token, spec.block_size)))
+            end_block = aligned_token_len // spec.block_size
+            start_block = min(end_block, max(0, cdiv(start_token, spec.block_size)))
             if exclude_mamba and isinstance(spec, MambaSpec):
-                masks.append([False] * (end_chunk - start_chunk))
+                masks.append([False] * (end_block - start_block))
                 continue
             manager_cls = KVCacheSpecRegistry.get_manager_class(spec)
             assert manager_cls is not None
             use_eagle = g_idx in self.eagle_group_ids
             mask = manager_cls.reachable_block_mask(
-                start_block=start_chunk,
-                end_block=end_chunk,
+                start_block=start_block,
+                end_block=end_block,
                 alignment_tokens=self.lcm_block_size,
                 kv_cache_spec=spec,
                 use_eagle=use_eagle,
                 retention_interval=retention_interval,
                 reachable_boundaries=reachable_boundaries,
-                # ``spec`` is already DCP-resolved (worker.py applies
-                # resolve_dcp_kv_cache_spec) and ``end_chunk`` is indexed in
-                # that scaled block size, so the mask must not scale again.
+                # Worker initialization has already resolved DCP in this spec.
                 dcp_world_size=1,
             )
             if mask is not None:
-                assert len(mask) == end_chunk - start_chunk
+                assert len(mask) == end_block - start_block
             masks.append(mask)
         return tuple(masks)
 
-    def block_hashes_for_spec(
-        self, block_hashes: Sequence[BlockHash], spec: KVCacheSpec
-    ) -> Sequence[BlockHash]:
-        return chunk_hashes_for_block_size(
-            block_hashes, self.hash_block_size, spec.block_size
-        )
-
     def _find_hit_blocks(
         self,
         block_hashes: Sequence[BlockHash],
@@ -396,13 +360,7 @@ def _find_hit_blocks(
         *,
         apply_eagle: bool = True,
     ) -> tuple[tuple[list[KVCacheBlock], ...], int]:
-        """Mirrors HybridKVCacheCoordinator.find_longest_cache_hit but
-        dispatches via spec_manager_map (we don't allocate managers).
-
-        When ``apply_eagle`` is False, ignore each group's ``use_eagle`` —
-        used by ``load_mask`` to avoid popping a second block on top of the
-        one already removed by the lookup.
-        """
+        """Converge registered managers' hits as in HybridKVCacheCoordinator."""
         alignment_tokens = (
             self.hash_block_size
             if self.enable_partial_hash_hits
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 4cd0402b7b4f..47a0e2e1c0bb 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
@@ -9,16 +9,19 @@
 from dataclasses import dataclass, field
 from typing import NamedTuple, cast
 
-import numpy as np
 import torch
 
 from vllm.distributed.kv_transfer.kv_connector.v1.base import (
     KVConnectorMetadata,
     KVConnectorWorkerMetadata,
 )
+from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.layout import (
+    KeyMetadata,
+    RankLocalStoreLayout,
+    StoreLayout,
+)
 from vllm.logger import init_logger
 from vllm.utils.math_utils import cdiv
-from vllm.utils.torch_utils import is_non_overlapping_and_dense
 from vllm.v1.core.kv_cache_utils import (
     BlockHash,
     BlockHashListWithBlockSize,
@@ -52,15 +55,9 @@ def __getitem__(self, idx):
 
 
 class _CompactChunkHashList(BlockHashListWithBlockSize):
-    """View that keys each ``block_size`` chunk by the last constituent
-    ``hash_block_size`` hash instead of concatenating all of them.
-
-    The engine chains block hashes (each hash folds in the previous one), so the
-    final sub-block hash of a chunk already uniquely identifies the whole chunk
-    and its prefix. Using it keeps a Mooncake key at a single hash digest
-    regardless of the ``block_size`` / ``hash_block_size`` ratio, instead of
-    growing the key linearly with it (e.g. 64x for ``block_size=256``,
-    ``hash_block_size=4``).
+    """Key each chunk by its final chained hash, which identifies its full prefix.
+
+    This keeps Store keys one digest long regardless of the chunk/hash size ratio.
     """
 
     def __init__(
@@ -69,8 +66,7 @@ def __init__(
         hash_block_size: int,
         target_block_size: int,
     ):
-        # Accept any indexable sequence (e.g. the lazy ``BlobBlockHashes``), not
-        # just ``list``; the base only indexes/sizes it.
+        # The base class only needs sequence indexing and length.
         assert target_block_size % hash_block_size == 0
         self.block_hashes = block_hashes  # type: ignore[assignment]
         self.scale_factor = target_block_size // hash_block_size
@@ -98,510 +94,39 @@ def chunk_hashes_for_block_size(
     )
 
 
-@dataclass
-class KeyMetadata:
-    """Metadata for constructing pool keys."""
-
-    model_name: str
-    tp_rank: int
-    pcp_rank: int
-    dcp_rank: int
-    pp_rank: int
-    group_id: int = 0
-    # Optional namespace prepended to every key. Lets separate deployments
-    # share one Mooncake master without colliding on identical block hashes.
-    # Empty (the default) keeps keys byte-identical to the unprefixed format.
-    cache_prefix: str = ""
-    # Complete namespace for an opt-in Store payload format. Empty keeps
-    # historical keys byte-identical.
-    store_namespace: str = ""
-
-
-@dataclass(order=True)
-class PoolKey:
-    """Key for addressing KV cache blocks in the distributed store."""
-
-    key_metadata: KeyMetadata
-    chunk_hash: str
-
-    def __hash__(self):
-        return hash(
-            (
-                self.key_metadata.cache_prefix,
-                self.key_metadata.model_name,
-                self.key_metadata.store_namespace,
-                self.key_metadata.tp_rank,
-                self.key_metadata.pcp_rank,
-                self.key_metadata.dcp_rank,
-                self.key_metadata.pp_rank,
-                self.key_metadata.group_id,
-                self.chunk_hash,
-            )
-        )
-
-    @staticmethod
-    def build_prefix(
-        key_metadata: KeyMetadata,
-        *,
-        tp_rank: int | None = None,
-        pcp_rank: int | None = None,
-        dcp_rank: int | None = None,
-        pp_rank: int | None = None,
-    ) -> str:
-        """Return the stable prefix for a Mooncake pool key."""
-        prefix = f"{key_metadata.cache_prefix}@" if key_metadata.cache_prefix else ""
-        return (
-            f"{prefix}"
-            f"{key_metadata.model_name}"
-            f"{key_metadata.store_namespace}"
-            f"@tp_rank:{key_metadata.tp_rank if tp_rank is None else tp_rank}"
-            f"@pcp{key_metadata.pcp_rank if pcp_rank is None else pcp_rank}"
-            f"@dcp{key_metadata.dcp_rank if dcp_rank is None else dcp_rank}"
-            f"@pp_rank:{key_metadata.pp_rank if pp_rank is None else pp_rank}"
-            f"@group:{key_metadata.group_id}"
-        )
-
-    @staticmethod
-    def build_key_string(key_prefix: str, chunk_hash: str) -> str:
-        return f"{key_prefix}@{chunk_hash}"
-
-    def to_string(self) -> str:
-        return self.build_key_string(
-            self.build_prefix(self.key_metadata), self.chunk_hash
-        )
-
-
-StoreShardId = int
-# Physical ranks encoded in a rank-local key: (TP, PCP, DCP, PP).
-RankNamespace = tuple[int, int, int, int]
-
-
-class StoreLayout:
-    """Store payload layout for one KV-cache group on one local TP rank."""
-
-    def __init__(
-        self,
-        metadata: KeyMetadata,
-        block_size: int,
-        hash_block_size: int,
-    ) -> None:
-        self.metadata = metadata
-        self.block_size = block_size
-        self.hash_block_size = hash_block_size
-
-    @property
-    def local_shard_ids(self) -> tuple[StoreShardId, ...]:
-        """Store objects contributed or consumed by this local rank."""
-        raise NotImplementedError
-
-    def key_for(self, shard_id: StoreShardId, chunk_hash: BlockHash) -> str:
-        raise NotImplementedError
-
-    def lookup_key_prefixes(
-        self,
-        rank_namespaces: Sequence[RankNamespace],
-    ) -> tuple[str, ...]:
-        """Prefixes that must exist for a logical block to be reusable."""
-        raise NotImplementedError
-
-    def register_kv_caches(
-        self,
-        kv_caches: Sequence[torch.Tensor],
-        num_blocks: int,
-    ) -> None:
-        """Build descriptors from the local KV cache tensors."""
-        raise NotImplementedError
-
-    def prepare_values(
-        self,
-        chunks: Sequence[tuple[int, int]],
-        block_ids: list[int],
-        shard_ids: Sequence[StoreShardId],
-    ) -> tuple[list[list[int]], list[list[int]], list[int]]:
-        """Build Store multi-buffer descriptors for logical chunks."""
-        raise NotImplementedError
-
-
-class RankLocalStoreLayout(StoreLayout):
-    """Historical rank-local Store payload layout."""
-
-    def __init__(
-        self,
-        metadata: KeyMetadata,
-        block_size: int,
-        hash_block_size: int,
-    ) -> None:
-        super().__init__(metadata, block_size, hash_block_size)
-        self._key_prefix = PoolKey.build_prefix(metadata)
-        self.kv_caches_base_addr: list[int] = []
-        self.block_len: list[int] = []
-
-    @property
-    def local_shard_ids(self) -> tuple[StoreShardId, ...]:
-        return (0,)
-
-    def key_for(self, shard_id: StoreShardId, chunk_hash: BlockHash) -> str:
-        assert shard_id == 0
-        return PoolKey.build_key_string(self._key_prefix, chunk_hash.hex())
-
-    def lookup_key_prefixes(
-        self,
-        rank_namespaces: Sequence[RankNamespace],
-    ) -> tuple[str, ...]:
-        return tuple(
-            PoolKey.build_prefix(
-                self.metadata,
-                tp_rank=tp_rank,
-                pcp_rank=pcp_rank,
-                dcp_rank=dcp_rank,
-                pp_rank=pp_rank,
-            )
-            for tp_rank, pcp_rank, dcp_rank, pp_rank in rank_namespaces
-        )
-
-    def set_kv_caches_base_addr(self, base_addrs: list[int]) -> None:
-        self.kv_caches_base_addr = base_addrs
-
-    def set_block_len(self, block_lens: list[int]) -> None:
-        self.block_len = block_lens
-
-    def register_kv_caches(
-        self,
-        kv_caches: Sequence[torch.Tensor],
-        num_blocks: int,
-    ) -> None:
-        seen_region_ptrs: set[int] = set()
-        base_addrs: list[int] = []
-        block_lens: list[int] = []
-        for cache in kv_caches:
-            storage = cache.untyped_storage()
-            storage_addr = storage.data_ptr()
-            storage_len = storage.nbytes()
-            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])
-                    region_addr = head_cache.data_ptr()
-                    if region_addr in seen_region_ptrs:
-                        continue
-                    seen_region_ptrs.add(region_addr)
-                    base_addrs.append(region_addr)
-                    block_lens.append(head_cache.stride(0) * head_cache.element_size())
-            elif cache.stride(0) * cache.element_size() * num_blocks == storage_len:
-                if storage_addr in seen_region_ptrs:
-                    continue
-                seen_region_ptrs.add(storage_addr)
-                base_addrs.append(storage_addr)
-                block_lens.append(storage_len // num_blocks)
-            else:
-                region_addr = cache.data_ptr()
-                if region_addr in seen_region_ptrs:
-                    continue
-                seen_region_ptrs.add(region_addr)
-                base_addrs.append(region_addr)
-                block_lens.append(cache.stride(0) * cache.element_size())
-        self.kv_caches_base_addr = base_addrs
-        self.block_len = block_lens
-
-    def prepare_values(
-        self,
-        chunks: Sequence[tuple[int, int]],
-        block_ids: list[int],
-        shard_ids: Sequence[StoreShardId],
-    ) -> tuple[list[list[int]], list[list[int]], list[int]]:
-        if not chunks:
-            return [], [], []
-        assert len(chunks) == len(shard_ids) and all(
-            shard_id == 0 for shard_id in shard_ids
-        )
-        base = np.asarray(self.kv_caches_base_addr, dtype=np.int64)
-        length = len(self.block_len)
-        blen = np.asarray(
-            [self.block_len[i % length] for i in range(base.shape[0])],
-            dtype=np.int64,
-        )
-        n = len(chunks)
-        starts = np.fromiter((c[0] for c in chunks), dtype=np.int64, count=n)
-        spans = np.fromiter((c[1] for c in chunks), dtype=np.int64, count=n) - starts
-        assert np.all(spans % self.hash_block_size == 0)
-        bids = np.fromiter(
-            (block_ids[i] for i in (starts // self.block_size).tolist()),
-            dtype=np.int64,
-            count=n,
-        )
-        addrs = base[None, :] + bids[:, None] * blen[None, :]
-        block_counts = (spans + self.block_size - 1) // self.block_size
-        sizes = blen[None, :] * block_counts[:, None]
-        return addrs.tolist(), sizes.tolist(), bids.tolist()
-
-    def prepare_value_for_block(self, block_id: int) -> tuple[list[int], list[int]]:
-        length = len(self.block_len)
-        return (
-            [
-                base_addr + block_id * self.block_len[index % length]
-                for index, base_addr in enumerate(self.kv_caches_base_addr)
-            ],
-            [
-                self.block_len[index % length]
-                for index in range(len(self.kv_caches_base_addr))
-            ],
-        )
-
-
-class TPShardedStoreLayout(StoreLayout):
-    """Store layout shared by divisible TP sizes."""
-
-    store_format: str
-
-    @classmethod
-    def shared_namespace(cls, store_tp_size: int, pp_size: int) -> str:
-        return (
-            f"@store_tp:{store_tp_size}@store_pp:{pp_size}"
-            f"@store_format:{cls.store_format}"
-        )
-
-    def __init__(
-        self,
-        metadata: KeyMetadata,
-        block_size: int,
-        hash_block_size: int,
-        local_tp_size: int,
-        store_tp_size: int,
-        tp_rank: int,
-        num_kv_heads: int,
-    ) -> None:
-        super().__init__(metadata, block_size, hash_block_size)
-        self.shards_per_rank = store_tp_size // local_tp_size
-        first_shard = tp_rank * self.shards_per_rank
-        self.store_shard_ids = tuple(
-            range(first_shard, first_shard + self.shards_per_rank)
-        )
-        self.store_tp_size = store_tp_size
-        self.heads_per_store_shard = num_kv_heads // store_tp_size
-        self.local_num_kv_heads = num_kv_heads // local_tp_size
-        self._shard_key_prefixes = {
-            shard_id: PoolKey.build_prefix(metadata, tp_rank=shard_id)
-            for shard_id in self.store_shard_ids
-        }
-        self._shard_addr_bases = np.empty((self.shards_per_rank, 0), dtype=np.uint64)
-        self._shard_block_strides = np.empty((self.shards_per_rank, 0), dtype=np.uint64)
-        self._shard_sizes = np.empty((self.shards_per_rank, 0), dtype=np.uint64)
-
-    @property
-    def local_shard_ids(self) -> tuple[StoreShardId, ...]:
-        return self.store_shard_ids
-
-    def key_for(self, shard_id: StoreShardId, chunk_hash: BlockHash) -> str:
-        return PoolKey.build_key_string(
-            self._shard_key_prefixes[shard_id], chunk_hash.hex()
-        )
-
-    def lookup_key_prefixes(
-        self,
-        rank_namespaces: Sequence[RankNamespace],
-    ) -> tuple[str, ...]:
-        pp_ranks = sorted({namespace[3] for namespace in rank_namespaces})
-        return tuple(
-            PoolKey.build_prefix(self.metadata, tp_rank=shard_id, pp_rank=pp_rank)
-            for shard_id in range(self.store_tp_size)
-            for pp_rank in pp_ranks
-        )
-
-    def _set_segment_templates(
-        self, templates: Sequence[tuple[list[int], list[int], list[int]]]
-    ) -> None:
-        self._shard_addr_bases = np.asarray(
-            [template[0] for template in templates], dtype=np.uint64
-        )
-        self._shard_block_strides = np.asarray(
-            [template[1] for template in templates], dtype=np.uint64
-        )
-        self._shard_sizes = np.asarray(
-            [template[2] for template in templates], dtype=np.uint64
-        )
-
-    def _has_packed_strides(
-        self, head_stride: int, token_stride: int, content_bytes: int
-    ) -> bool:
-        raise NotImplementedError
-
-    def _shard_segments(
-        self, head_start: int, head_stride: int, token_stride: int, content_bytes: int
-    ) -> Iterable[tuple[int, int]]:
-        raise NotImplementedError
-
-    def register_kv_caches(
-        self,
-        kv_caches: Sequence[torch.Tensor],
-        num_blocks: int,
-    ) -> None:
-        templates: list[tuple[list[int], list[int], list[int]]] = [
-            ([], [], []) for _ in range(self.shards_per_rank)
-        ]
-        for cache in kv_caches:
-            if cache.ndim != 4 or tuple(cache.shape[:3]) != (
-                num_blocks,
-                self.local_num_kv_heads,
-                self.block_size,
-            ):
-                raise ValueError(
-                    "TP-shared Mooncake store requires packed KV caches with "
-                    "logical shape (num_blocks, local_kv_heads, block_size, content)"
-                )
-
-            element_size = cache.element_size()
-            block_stride, head_stride, token_stride, content_stride = (
-                stride * element_size for stride in cache.stride()
-            )
-            content_bytes = cache.shape[3] * element_size
-            if content_stride != element_size or not self._has_packed_strides(
-                head_stride, token_stride, content_bytes
-            ):
-                raise ValueError(
-                    "TP-shared Mooncake store requires packed "
-                    f"{self.store_format} KV layout"
-                )
-
-            for local_shard, (addr_bases, block_strides, sizes) in enumerate(templates):
-                head_start = local_shard * self.heads_per_store_shard
-                for offset, size in self._shard_segments(
-                    head_start, head_stride, token_stride, content_bytes
-                ):
-                    addr_bases.append(cache.data_ptr() + offset)
-                    block_strides.append(block_stride)
-                    sizes.append(size)
-
-        self._set_segment_templates(templates)
-
-    def prepare_values(
-        self,
-        chunks: Sequence[tuple[int, int]],
-        block_ids: list[int],
-        shard_ids: Sequence[StoreShardId],
-    ) -> tuple[list[list[int]], list[list[int]], list[int]]:
-        if not chunks:
-            return [], [], []
-        if len(chunks) != len(shard_ids):
-            raise ValueError("Each Store chunk must have one shard ID")
-
-        count = len(chunks)
-        starts = np.fromiter(
-            (chunk[0] for chunk in chunks), dtype=np.int64, count=count
-        )
-        ends = np.fromiter((chunk[1] for chunk in chunks), dtype=np.int64, count=count)
-        if np.any(ends - starts != self.block_size):
-            raise ValueError("TP-shared Mooncake store requires full KV blocks")
-
-        first_shard = self.store_shard_ids[0]
-        local_shards = np.fromiter(
-            (shard_id - first_shard for shard_id in shard_ids),
-            dtype=np.int64,
-            count=count,
-        )
-        if np.any(local_shards < 0) or np.any(local_shards >= self.shards_per_rank):
-            raise ValueError("Store shard is not owned by this TP rank")
-        chunk_block_ids = np.fromiter(
-            (block_ids[index] for index in (starts // self.block_size).tolist()),
-            dtype=np.uint64,
-            count=count,
-        )
-        addr_bases = self._shard_addr_bases[local_shards]
-        block_strides = self._shard_block_strides[local_shards]
-        addrs = addr_bases + chunk_block_ids[:, None] * block_strides
-        sizes = self._shard_sizes[local_shards]
-        return addrs.tolist(), sizes.tolist(), chunk_block_ids.tolist()
-
-
-class LBHNCStoreLayout(TPShardedStoreLayout):
-    """Native head-major layout shared by divisible TP sizes."""
-
-    store_format = "tp_shared_lbhnc"
-
-    def _has_packed_strides(
-        self, head_stride: int, token_stride: int, content_bytes: int
-    ) -> bool:
-        return (
-            token_stride == content_bytes
-            and head_stride == self.block_size * content_bytes
-        )
-
-    def _shard_segments(
-        self, head_start: int, head_stride: int, token_stride: int, content_bytes: int
-    ) -> Iterable[tuple[int, int]]:
-        yield head_start * head_stride, self.heads_per_store_shard * head_stride
-
-
-class LBNHCStoreLayout(TPShardedStoreLayout):
-    """Native token-major layout shared by divisible TP sizes."""
-
-    store_format = "tp_shared_lbnhc"
-
-    def _has_packed_strides(
-        self, head_stride: int, token_stride: int, content_bytes: int
-    ) -> bool:
-        return (
-            head_stride == content_bytes
-            and token_stride == self.local_num_kv_heads * content_bytes
-        )
-
-    def _shard_segments(
-        self, head_start: int, head_stride: int, token_stride: int, content_bytes: int
-    ) -> Iterable[tuple[int, int]]:
-        for token_idx in range(self.block_size):
-            yield (
-                token_idx * token_stride + head_start * head_stride,
-                self.heads_per_store_shard * content_bytes,
-            )
-
-
 class ChunkedTokenDatabase:
-    """Enumerates logical token chunks and their hashes."""
+    """Enumerates Store chunks and their prefix hashes."""
 
     def __init__(
         self,
         metadata: KeyMetadata,
-        block_size: int,
+        chunk_size: int,
         hash_block_size: int | None = None,
         store_layout: StoreLayout | None = None,
     ):
         self.metadata = metadata
-        self.block_size = block_size
-        self.hash_block_size = hash_block_size or block_size
-        if self.block_size % self.hash_block_size != 0:
+        self.chunk_size = chunk_size
+        self.hash_block_size = hash_block_size or chunk_size
+        if self.chunk_size % self.hash_block_size != 0:
             raise ValueError(
-                f"block_size ({self.block_size}) must be a multiple of "
+                f"chunk_size ({self.chunk_size}) must be a multiple of "
                 f"hash_block_size ({self.hash_block_size})"
             )
         self.store_layout = store_layout or RankLocalStoreLayout(
-            metadata, block_size, self.hash_block_size
+            metadata, chunk_size, self.hash_block_size
         )
 
-    @property
-    def kv_caches_base_addr(self) -> list[int]:
-        return self._rank_local_layout().kv_caches_base_addr
-
-    @property
-    def block_len(self) -> list[int]:
-        return self._rank_local_layout().block_len
-
     def _rank_local_layout(self) -> RankLocalStoreLayout:
         if not isinstance(self.store_layout, RankLocalStoreLayout):
             raise RuntimeError("This operation requires a rank-local Store layout")
         return self.store_layout
 
-    def key_for(self, chunk_hash: BlockHash) -> str:
-        return self._rank_local_layout().key_for(0, chunk_hash)
-
     def set_kv_caches_base_addr(self, kv_caches_base_addr: list[int]):
         self._rank_local_layout().set_kv_caches_base_addr(kv_caches_base_addr)
 
     def set_block_len(self, block_len: list[int]):
         self._rank_local_layout().set_block_len(block_len)
 
-    def prepare_value_for_block(self, block_id: int) -> tuple[list[int], list[int]]:
-        """Return addresses and sizes for one physical block slot."""
-        return self._rank_local_layout().prepare_value_for_block(block_id)
-
     def process_tokens(
         self,
         token_len: int,
@@ -623,11 +148,9 @@ def process_tokens(
             token_len: Total number of tokens. Must be hash-block aligned and
                 covered by ``block_hashes`` when hashes are present.
             block_hashes: Block hashes computed at ``hash_block_size`` granularity.
-                When ``block_size > hash_block_size`` each group's ``block_size`` chunk
-                is keyed by its last sub-hash via ``chunk_hashes_for_block_size``.
+                A chunk spanning multiple hashes is keyed by its last hash.
             mask_num: Number of tokens to skip from the beginning.
-            chunk_mask: Optional mask relative to the first chunk after
-                ``mask_num``. False entries are skipped before hash access.
+            chunk_mask: Optional mask over chunks after ``mask_num``.
             put_step: Stride for distributing chunks across ranks.
             put_step_rank: ``chunk_id % put_step`` value this rank stores.
 
@@ -637,8 +160,8 @@ def process_tokens(
             return
         assert token_len % self.hash_block_size == 0
         assert token_len // self.hash_block_size <= len(block_hashes)
-        start_chunk = max(0, cdiv(mask_num, self.block_size))
-        max_chunks = cdiv(token_len, self.block_size)
+        start_chunk = max(0, cdiv(mask_num, self.chunk_size))
+        max_chunks = cdiv(token_len, self.chunk_size)
         if chunk_mask is not None:
             max_chunks = min(max_chunks, start_chunk + len(chunk_mask))
         for chunk_id in range(start_chunk, max_chunks):
@@ -646,8 +169,8 @@ def process_tokens(
                 continue
             if chunk_id % put_step != put_step_rank:
                 continue
-            start_idx = chunk_id * self.block_size
-            end_idx = min(start_idx + self.block_size, token_len)
+            start_idx = chunk_id * self.chunk_size
+            end_idx = min(start_idx + self.chunk_size, token_len)
             h = block_hashes[end_idx // self.hash_block_size - 1]
             yield start_idx, end_idx, h
 
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/layout.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/layout.py
new file mode 100644
index 000000000000..081ab27d4109
--- /dev/null
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/layout.py
@@ -0,0 +1,848 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+#
+# Adapted from vllm-project/vllm-ascend
+# (vllm_ascend/distributed/kv_transfer/kv_pool/ascend_store/).
+"""Store keys and tensor layouts for MooncakeStoreConnector."""
+
+import hashlib
+from collections.abc import Iterable, Sequence
+from dataclasses import dataclass
+from fractions import Fraction
+from typing import NamedTuple, cast
+
+import numpy as np
+import torch
+
+from vllm.distributed.kv_transfer.kv_connector.v1.ssm_conv_transfer_utils import (
+    derive_mamba_conv_split,
+)
+from vllm.model_executor.layers.mamba.mamba_utils import is_conv_state_dim_first
+from vllm.utils.torch_utils import is_non_overlapping_and_dense
+from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
+from vllm.v1.core.kv_cache_utils import BlockHash
+from vllm.v1.kv_cache_interface import AttentionSpec, MambaSpec
+from vllm.v1.kv_cache_layout import KVCacheLayout
+
+
+def _schema_fingerprint(parts: tuple[object, ...]) -> str:
+    return hashlib.sha256(repr(parts).encode()).hexdigest()[:16]
+
+
+@dataclass
+class KeyMetadata:
+    """Metadata for constructing pool keys."""
+
+    model_name: str
+    tp_rank: int
+    pcp_rank: int
+    dcp_rank: int
+    pp_rank: int
+    group_id: int = 0
+    # Optional namespace prepended to every key. Lets separate deployments
+    # share one Mooncake master without colliding on identical block hashes.
+    # Empty (the default) keeps keys byte-identical to the unprefixed format.
+    cache_prefix: str = ""
+    # Complete namespace for an opt-in Store payload format. Empty keeps
+    # historical keys byte-identical.
+    store_namespace: str = ""
+
+
+@dataclass(order=True)
+class PoolKey:
+    """Key for addressing KV cache blocks in the distributed store."""
+
+    key_metadata: KeyMetadata
+    chunk_hash: str
+
+    def __hash__(self):
+        return hash(
+            (
+                self.key_metadata.cache_prefix,
+                self.key_metadata.model_name,
+                self.key_metadata.store_namespace,
+                self.key_metadata.tp_rank,
+                self.key_metadata.pcp_rank,
+                self.key_metadata.dcp_rank,
+                self.key_metadata.pp_rank,
+                self.key_metadata.group_id,
+                self.chunk_hash,
+            )
+        )
+
+    @staticmethod
+    def build_prefix(
+        key_metadata: KeyMetadata,
+        *,
+        tp_rank: int | None = None,
+        pcp_rank: int | None = None,
+        dcp_rank: int | None = None,
+        pp_rank: int | None = None,
+    ) -> str:
+        """Return the stable prefix for a Mooncake pool key."""
+        prefix = f"{key_metadata.cache_prefix}@" if key_metadata.cache_prefix else ""
+        return (
+            f"{prefix}"
+            f"{key_metadata.model_name}"
+            f"{key_metadata.store_namespace}"
+            f"@tp_rank:{key_metadata.tp_rank if tp_rank is None else tp_rank}"
+            f"@pcp{key_metadata.pcp_rank if pcp_rank is None else pcp_rank}"
+            f"@dcp{key_metadata.dcp_rank if dcp_rank is None else dcp_rank}"
+            f"@pp_rank:{key_metadata.pp_rank if pp_rank is None else pp_rank}"
+            f"@group:{key_metadata.group_id}"
+        )
+
+    @staticmethod
+    def build_key_string(key_prefix: str, chunk_hash: str) -> str:
+        return f"{key_prefix}@{chunk_hash}"
+
+    def to_string(self) -> str:
+        return self.build_key_string(
+            self.build_prefix(self.key_metadata), self.chunk_hash
+        )
+
+
+StoreShardId = int
+# Physical ranks encoded in a rank-local key: (TP, PCP, DCP, PP).
+RankNamespace = tuple[int, int, int, int]
+
+
+class StoreLayout:
+    """Store payload layout for one KV-cache group on one local TP rank."""
+
+    store_format = "rank_local"
+
+    def __init__(
+        self,
+        metadata: KeyMetadata,
+        block_size: int,
+        hash_block_size: int,
+    ) -> None:
+        self.metadata = metadata
+        self.block_size = block_size
+        self.hash_block_size = hash_block_size
+
+    @property
+    def local_shard_ids(self) -> tuple[StoreShardId, ...]:
+        """Store objects contributed or consumed by this local rank."""
+        raise NotImplementedError
+
+    def key_for(self, shard_id: StoreShardId, chunk_hash: BlockHash) -> str:
+        raise NotImplementedError
+
+    def lookup_key_prefixes(
+        self,
+        rank_namespaces: Sequence[RankNamespace],
+    ) -> tuple[str, ...]:
+        """Prefixes that must exist for a logical block to be reusable."""
+        raise NotImplementedError
+
+    def register_kv_caches(
+        self,
+        kv_caches: Sequence[torch.Tensor],
+        num_blocks: int,
+    ) -> None:
+        """Build descriptors from the local KV cache tensors."""
+        raise NotImplementedError
+
+    def prepare_values(
+        self,
+        chunks: Sequence[tuple[int, int]],
+        block_ids: list[int],
+        shard_ids: Sequence[StoreShardId],
+    ) -> tuple[list[list[int]], list[list[int]], list[int]]:
+        """Build Store multi-buffer descriptors for logical chunks."""
+        raise NotImplementedError
+
+
+class RankLocalStoreLayout(StoreLayout):
+    """Historical rank-local Store payload layout."""
+
+    def __init__(
+        self,
+        metadata: KeyMetadata,
+        block_size: int,
+        hash_block_size: int,
+    ) -> None:
+        super().__init__(metadata, block_size, hash_block_size)
+        self._key_prefix = PoolKey.build_prefix(metadata)
+        self.kv_caches_base_addr: list[int] = []
+        self.block_stride: list[int] = []
+        # Bytes at the start of each block that are stored: up to the last
+        # byte of the group's data, skipping e.g. the padding of a Mamba state
+        # up to the attention page size.
+        self.block_len: list[int] = []
+
+    @property
+    def local_shard_ids(self) -> tuple[StoreShardId, ...]:
+        return (0,)
+
+    def key_for(self, shard_id: StoreShardId, chunk_hash: BlockHash) -> str:
+        assert shard_id == 0
+        return PoolKey.build_key_string(self._key_prefix, chunk_hash.hex())
+
+    def lookup_key_prefixes(
+        self,
+        rank_namespaces: Sequence[RankNamespace],
+    ) -> tuple[str, ...]:
+        return tuple(
+            PoolKey.build_prefix(
+                self.metadata,
+                tp_rank=tp_rank,
+                pcp_rank=pcp_rank,
+                dcp_rank=dcp_rank,
+                pp_rank=pp_rank,
+            )
+            for tp_rank, pcp_rank, dcp_rank, pp_rank in rank_namespaces
+        )
+
+    def set_kv_caches_base_addr(self, base_addrs: list[int]) -> None:
+        self.kv_caches_base_addr = base_addrs
+
+    def set_block_len(self, block_lens: list[int]) -> None:
+        self.block_stride = block_lens
+        self.block_len = list(block_lens)
+
+    def register_kv_caches(
+        self,
+        kv_caches: Sequence[torch.Tensor],
+        num_blocks: int,
+    ) -> None:
+        # Map each region's base address to its entry in the parallel lists.
+        region_index: dict[int, int] = {}
+        base_addrs: list[int] = []
+        block_strides: list[int] = []
+        block_lens: list[int] = []
+
+        def add_region(addr: int, block_stride: int, block_len: int) -> None:
+            assert 0 < block_len <= block_stride
+            index = region_index.get(addr)
+            if index is None:
+                region_index[addr] = len(base_addrs)
+                base_addrs.append(addr)
+                block_strides.append(block_stride)
+                block_lens.append(block_len)
+            else:
+                # Several views can share a region (e.g. the layers of a
+                # block-outermost buffer); keep the bytes up to the last one.
+                block_lens[index] = max(block_lens[index], block_len)
+
+        for cache in kv_caches:
+            storage = cache.untyped_storage()
+            storage_addr = storage.data_ptr()
+            storage_len = storage.nbytes()
+            block_bytes = cache[0].numel() * cache.element_size()
+            if not is_non_overlapping_and_dense(cache[0]):
+                # One region per head, copying only that head's block (in a
+                # block-outermost buffer the block stride spans every head).
+                for head_idx in range(cache.shape[1]):
+                    head_cache = cache[:, head_idx]
+                    assert is_non_overlapping_and_dense(head_cache[0])
+                    add_region(
+                        head_cache.data_ptr(),
+                        head_cache.stride(0) * head_cache.element_size(),
+                        head_cache[0].numel() * head_cache.element_size(),
+                    )
+            elif cache.stride(0) * cache.element_size() * num_blocks == storage_len:
+                # Block-outermost buffer: each block holds every layer, so the
+                # whole storage is one region. In block 0 this layer's data
+                # starts at cache.data_ptr() - storage_addr and spans
+                # block_bytes, so the block is used up to that end; the
+                # region keeps the end of its last layer.
+                add_region(
+                    storage_addr,
+                    storage_len // num_blocks,
+                    cache.data_ptr() - storage_addr + block_bytes,
+                )
+            else:
+                add_region(
+                    cache.data_ptr(),
+                    cache.stride(0) * cache.element_size(),
+                    block_bytes,
+                )
+        self.kv_caches_base_addr = base_addrs
+        self.block_stride = block_strides
+        self.block_len = block_lens
+
+    def prepare_values(
+        self,
+        chunks: Sequence[tuple[int, int]],
+        block_ids: list[int],
+        shard_ids: Sequence[StoreShardId],
+    ) -> tuple[list[list[int]], list[list[int]], list[int]]:
+        if not chunks:
+            return [], [], []
+        assert len(chunks) == len(shard_ids) and all(
+            shard_id == 0 for shard_id in shard_ids
+        )
+        base = np.asarray(self.kv_caches_base_addr, dtype=np.int64)
+        length = len(self.block_stride)
+        stride = np.asarray(
+            [self.block_stride[i % length] for i in range(base.shape[0])],
+            dtype=np.int64,
+        )
+        blen = np.asarray(
+            [self.block_len[i % length] for i in range(base.shape[0])],
+            dtype=np.int64,
+        )
+        n = len(chunks)
+        starts = np.fromiter((c[0] for c in chunks), dtype=np.int64, count=n)
+        spans = np.fromiter((c[1] for c in chunks), dtype=np.int64, count=n) - starts
+        assert np.all(spans % self.hash_block_size == 0)
+        bids = np.fromiter(
+            (block_ids[i] for i in (starts // self.block_size).tolist()),
+            dtype=np.int64,
+            count=n,
+        )
+        addrs = base[None, :] + bids[:, None] * stride[None, :]
+        block_counts = (spans + self.block_size - 1) // self.block_size
+        # A chunk spanning several consecutive blocks also copies the
+        # padding between them.
+        sizes = stride[None, :] * (block_counts[:, None] - 1) + blen[None, :]
+        return addrs.tolist(), sizes.tolist(), bids.tolist()
+
+
+class TPShardedStoreLayout(StoreLayout):
+    """Store layout shared by divisible TP sizes."""
+
+    def __init__(
+        self,
+        metadata: KeyMetadata,
+        block_size: int,
+        hash_block_size: int,
+        local_tp_size: int,
+        store_tp_size: int,
+        tp_rank: int,
+        store_chunk_size: int | None = None,
+    ) -> None:
+        super().__init__(metadata, block_size, hash_block_size)
+        self.store_chunk_size = store_chunk_size or block_size
+        if block_size % self.store_chunk_size:
+            raise ValueError("Store chunk size must divide the local block size")
+        self.shards_per_rank = store_tp_size // local_tp_size
+        first_shard = tp_rank * self.shards_per_rank
+        self.store_shard_ids = tuple(
+            range(first_shard, first_shard + self.shards_per_rank)
+        )
+        self.store_tp_size = store_tp_size
+        self._shard_key_prefixes = {
+            shard_id: PoolKey.build_prefix(metadata, tp_rank=shard_id)
+            for shard_id in self.store_shard_ids
+        }
+        self._chunk_addr_bases: np.ndarray | list[np.ndarray] = []
+        self._chunk_block_strides: np.ndarray | list[np.ndarray] = []
+        self._chunk_sizes: np.ndarray | list[np.ndarray] = []
+
+    @property
+    def local_shard_ids(self) -> tuple[StoreShardId, ...]:
+        return self.store_shard_ids
+
+    def key_for(self, shard_id: StoreShardId, chunk_hash: BlockHash) -> str:
+        return PoolKey.build_key_string(
+            self._shard_key_prefixes[shard_id], chunk_hash.hex()
+        )
+
+    def lookup_key_prefixes(
+        self,
+        rank_namespaces: Sequence[RankNamespace],
+    ) -> tuple[str, ...]:
+        pp_ranks = sorted({namespace[3] for namespace in rank_namespaces})
+        return tuple(
+            PoolKey.build_prefix(self.metadata, tp_rank=shard_id, pp_rank=pp_rank)
+            for shard_id in range(self.store_tp_size)
+            for pp_rank in pp_ranks
+        )
+
+    def _set_segment_templates(
+        self,
+        templates: Sequence[Sequence[tuple[list[int], list[int], list[int]]]],
+    ) -> None:
+        self._chunk_addr_bases = []
+        self._chunk_block_strides = []
+        self._chunk_sizes = []
+        for chunk in templates:
+            bases, strides, sizes = zip(*chunk, strict=True)
+            self._chunk_addr_bases.append(np.asarray(bases, dtype=np.uint64))
+            self._chunk_block_strides.append(np.asarray(strides, dtype=np.uint64))
+            self._chunk_sizes.append(np.asarray(sizes, dtype=np.uint64))
+        if len({chunk.shape for chunk in self._chunk_addr_bases}) == 1:
+            self._chunk_addr_bases = np.stack(self._chunk_addr_bases)
+            self._chunk_block_strides = np.stack(self._chunk_block_strides)
+            self._chunk_sizes = np.stack(self._chunk_sizes)
+
+    def prepare_values(
+        self,
+        chunks: Sequence[tuple[int, int]],
+        block_ids: list[int],
+        shard_ids: Sequence[StoreShardId],
+    ) -> tuple[list[list[int]], list[list[int]], list[int]]:
+        if not chunks:
+            return [], [], []
+        if len(chunks) != len(shard_ids):
+            raise ValueError("Each Store chunk must have one shard ID")
+
+        first_shard = self.store_shard_ids[0]
+        # Small batches avoid NumPy indexing setup.
+        if len(chunks) > 16 and isinstance(self._chunk_addr_bases, np.ndarray):
+            starts, ends = np.asarray(chunks, dtype=np.int64).T
+            spans = ends - starts
+            if np.any(starts % self.store_chunk_size) or np.any(
+                (spans <= 0) | (spans > self.store_chunk_size)
+            ):
+                raise ValueError("Invalid Store chunk boundary")
+            block_indices, offsets = np.divmod(starts, self.block_size)
+            chunk_indices = offsets // self.store_chunk_size
+            local_shards = np.asarray(shard_ids, dtype=np.int64) - first_shard
+            if np.any((local_shards < 0) | (local_shards >= self.shards_per_rank)):
+                raise ValueError("Store shard is not owned by this TP rank")
+            if np.any(block_indices >= len(block_ids)):
+                raise ValueError("Store chunk has no local cache block")
+            selected_blocks = np.fromiter(
+                (block_ids[index] for index in block_indices.tolist()),
+                dtype=np.uint64,
+                count=len(chunks),
+            )
+            bases = self._chunk_addr_bases[chunk_indices, local_shards]
+            strides = cast(np.ndarray, self._chunk_block_strides)[
+                chunk_indices, local_shards
+            ]
+            batch_sizes = cast(np.ndarray, self._chunk_sizes)[
+                chunk_indices, local_shards
+            ]
+            batch_addrs = bases + selected_blocks[:, None] * strides
+            return batch_addrs.tolist(), batch_sizes.tolist(), selected_blocks.tolist()
+
+        addrs: list[list[int]] = []
+        sizes: list[list[int]] = []
+        selected_block_ids: list[int] = []
+        for (start, end), shard_id in zip(chunks, shard_ids, strict=True):
+            if (
+                start % self.store_chunk_size
+                or not 0 < end - start <= self.store_chunk_size
+            ):
+                raise ValueError("Invalid Store chunk boundary")
+            block_index, block_offset = divmod(start, self.block_size)
+            chunk_index = block_offset // self.store_chunk_size
+            local_shard = shard_id - first_shard
+            if not 0 <= local_shard < self.shards_per_rank:
+                raise ValueError("Store shard is not owned by this TP rank")
+            if block_index >= len(block_ids):
+                raise ValueError("Store chunk has no local cache block")
+            block_id = block_ids[block_index]
+            bases = self._chunk_addr_bases[chunk_index][local_shard]
+            strides = self._chunk_block_strides[chunk_index][local_shard]
+            addrs.append((bases + block_id * strides).tolist())
+            sizes.append(self._chunk_sizes[chunk_index][local_shard].tolist())
+            selected_block_ids.append(block_id)
+        return addrs, sizes, selected_block_ids
+
+
+class _AttentionCache(NamedTuple):
+    block_stride: int
+    kernel_stride: int
+    head_stride: int
+    state_stride: int
+    kernel_blocks: int
+    num_states: int
+    content_bytes: int
+    num_heads: int
+
+
+class AttentionStoreLayout(TPShardedStoreLayout):
+    store_format: str
+
+    def __init__(
+        self,
+        metadata: KeyMetadata,
+        block_size: int,
+        hash_block_size: int,
+        local_tp_size: int,
+        store_tp_size: int,
+        tp_rank: int,
+        layer_specs: Sequence[AttentionSpec],
+        store_chunk_size: int | None = None,
+    ) -> None:
+        super().__init__(
+            metadata,
+            block_size,
+            hash_block_size,
+            local_tp_size,
+            store_tp_size,
+            tp_rank,
+            store_chunk_size,
+        )
+        self.layer_specs = tuple(layer_specs)
+        self.chunks_per_block = block_size // self.store_chunk_size
+        global_head_slots = {
+            spec.num_heads * local_tp_size for spec in self.layer_specs
+        }
+        if len(global_head_slots) != 1:
+            raise ValueError("Attention layers in a Store group need equal head slots")
+        global_num_head_slots = global_head_slots.pop()
+        if global_num_head_slots % store_tp_size:
+            raise ValueError(
+                "Attention head slots must be divisible by the Store shard count"
+            )
+        self.heads_per_store_shard = global_num_head_slots // store_tp_size
+
+    @staticmethod
+    def resolve_store_shard_count(
+        layer_specs: Sequence[AttentionSpec],
+        local_tp_size: int,
+        requested_store_tp_size: int,
+    ) -> int | None:
+        global_head_slots = {spec.num_heads * local_tp_size for spec in layer_specs}
+        if len(global_head_slots) != 1:
+            return None
+        global_num_head_slots = global_head_slots.pop()
+        smaller = min(requested_store_tp_size, global_num_head_slots)
+        larger = max(requested_store_tp_size, global_num_head_slots)
+        if larger % smaller or smaller < local_tp_size or smaller % local_tp_size:
+            return None
+        return smaller
+
+    @staticmethod
+    def resolve_store_chunk_size(
+        layer_specs: Sequence[AttentionSpec],
+        requested_store_chunk_size: int | None = None,
+    ) -> int | None:
+        if requested_store_chunk_size is not None:
+            for spec in layer_specs:
+                states = Fraction(requested_store_chunk_size, 1) / spec.tokens_per_state
+                if (
+                    spec.block_size % requested_store_chunk_size
+                    or states.denominator != 1
+                ):
+                    return None
+            return requested_store_chunk_size
+
+        block_sizes = {spec.block_size for spec in layer_specs}
+        return block_sizes.pop() if len(block_sizes) == 1 else None
+
+    @staticmethod
+    def schema_fingerprint(
+        layer_specs: Sequence[AttentionSpec],
+        local_tp_size: int,
+        store_shard_count: int,
+        store_chunk_size: int,
+    ) -> str:
+        global_num_head_slots = layer_specs[0].num_heads * local_tp_size
+        heads_per_store_shard = global_num_head_slots // store_shard_count
+        layers = tuple(
+            (
+                type(spec).__name__,
+                store_chunk_size,
+                str(Fraction(store_chunk_size, 1) / spec.tokens_per_state),
+                spec.state_content_size_bytes,
+                heads_per_store_shard,
+                str(spec.dtype),
+                str(spec.kv_quant_mode),
+                str(getattr(spec, "cache_dtype_str", None)),
+                str(spec.tokens_per_state),
+            )
+            for spec in layer_specs
+        )
+        return _schema_fingerprint(layers)
+
+    @staticmethod
+    def _state_range(
+        spec: AttentionSpec, store_chunk_size: int, chunk_index: int
+    ) -> tuple[int, int]:
+        num_states = Fraction(store_chunk_size, 1) / spec.tokens_per_state
+        if num_states.denominator != 1:
+            raise ValueError("Store chunk boundaries must align with attention states")
+        start = chunk_index * num_states.numerator
+        return start, start + num_states.numerator
+
+    @staticmethod
+    def _cache_strides(
+        cache: torch.Tensor, spec: AttentionSpec, num_blocks: int
+    ) -> _AttentionCache:
+        if cache.ndim == 4:
+            cache = cache.unsqueeze(1)
+        if cache.ndim != 5 or cache.shape[0] != num_blocks:
+            raise ValueError(
+                "TP-shared attention cache needs logical shape "
+                "(blocks, kernel_blocks, heads, states, content)"
+            )
+        _, kernel_blocks, num_heads, num_states, content = cache.shape
+        if num_heads != spec.num_heads:
+            raise ValueError("Attention cache head slots do not match its spec")
+        expected_states = spec.get_num_kernel_states(spec.block_size)
+        if kernel_blocks * num_states != expected_states:
+            raise ValueError("Attention cache state count does not match its spec")
+        element_size = cache.element_size()
+        block_stride, kernel_stride, head_stride, state_stride, content_stride = (
+            stride * element_size for stride in cache.stride()
+        )
+        content_bytes = content * element_size
+        if (
+            content_stride != element_size
+            or content_bytes != spec.state_content_size_bytes
+        ):
+            raise ValueError("TP-shared Mooncake store requires packed KV content")
+        return _AttentionCache(
+            block_stride,
+            kernel_stride,
+            head_stride,
+            state_stride,
+            kernel_blocks,
+            num_states,
+            content_bytes,
+            num_heads,
+        )
+
+    def register_kv_caches(
+        self,
+        kv_caches: Sequence[torch.Tensor],
+        num_blocks: int,
+    ) -> None:
+        templates: list[list[tuple[list[int], list[int], list[int]]]] = [
+            [([], [], []) for _ in range(self.shards_per_rank)]
+            for _ in range(self.chunks_per_block)
+        ]
+        for cache, spec in zip(kv_caches, self.layer_specs, strict=True):
+            strides = self._cache_strides(cache, spec, num_blocks)
+            for chunk_index, chunk_templates in enumerate(templates):
+                state_start, state_end = self._state_range(
+                    spec, self.store_chunk_size, chunk_index
+                )
+                if state_end > strides.kernel_blocks * strides.num_states:
+                    raise ValueError("Store chunk exceeds the local attention block")
+                for local_shard, (
+                    addr_bases,
+                    block_strides,
+                    sizes,
+                ) in enumerate(chunk_templates):
+                    head_start = local_shard * self.heads_per_store_shard
+                    for offset, size in self._shard_segments(
+                        strides, head_start, state_start, state_end
+                    ):
+                        addr_bases.append(cache.data_ptr() + offset)
+                        block_strides.append(strides.block_stride)
+                        sizes.append(size)
+
+        self._set_segment_templates(templates)
+
+    def _shard_segments(
+        self,
+        cache: _AttentionCache,
+        head_start: int,
+        state_start: int,
+        state_end: int,
+    ) -> Iterable[tuple[int, int]]:
+        raise NotImplementedError
+
+
+class HeadMajorStoreLayout(AttentionStoreLayout):
+    """Store descriptors with layer/head/state value ordering."""
+
+    store_format = "attention_head_major"
+
+    def _shard_segments(
+        self,
+        cache: _AttentionCache,
+        head_start: int,
+        state_start: int,
+        state_end: int,
+    ) -> Iterable[tuple[int, int]]:
+        if cache.state_stride != cache.content_bytes:
+            raise ValueError(
+                "TP-shared Mooncake store requires a packed head-major layout"
+            )
+        for head_index in range(head_start, head_start + self.heads_per_store_shard):
+            state_index = state_start
+            while state_index < state_end:
+                kernel_index, state_in_kernel = divmod(state_index, cache.num_states)
+                count = min(state_end - state_index, cache.num_states - state_in_kernel)
+                yield (
+                    kernel_index * cache.kernel_stride
+                    + head_index * cache.head_stride
+                    + state_in_kernel * cache.state_stride,
+                    count * cache.content_bytes,
+                )
+                state_index += count
+
+
+class TokenMajorStoreLayout(AttentionStoreLayout):
+    """Store descriptors with layer/state/head value ordering."""
+
+    store_format = "attention_token_major"
+
+    def _shard_segments(
+        self,
+        cache: _AttentionCache,
+        head_start: int,
+        state_start: int,
+        state_end: int,
+    ) -> Iterable[tuple[int, int]]:
+        if not (
+            cache.head_stride == cache.content_bytes
+            and cache.state_stride == cache.num_heads * cache.content_bytes
+        ):
+            raise ValueError(
+                "TP-shared Mooncake store requires a packed token-major layout"
+            )
+        for state_index in range(state_start, state_end):
+            kernel_index, state_in_kernel = divmod(state_index, cache.num_states)
+            yield (
+                kernel_index * cache.kernel_stride
+                + state_in_kernel * cache.state_stride
+                + head_start * cache.head_stride,
+                self.heads_per_store_shard * cache.content_bytes,
+            )
+
+
+ATTENTION_STORE_LAYOUTS: dict[KVCacheLayout, type[AttentionStoreLayout]] = {
+    KVCacheLayout.LBHNC: HeadMajorStoreLayout,
+    KVCacheLayout.LBNHC: TokenMajorStoreLayout,
+    KVCacheLayout.BLHNC: HeadMajorStoreLayout,
+    KVCacheLayout.BLNHC: TokenMajorStoreLayout,
+    KVCacheLayout.LHBNC: HeadMajorStoreLayout,
+    KVCacheLayout.BHLNC: HeadMajorStoreLayout,
+}
+
+
+class MambaStoreLayout(TPShardedStoreLayout):
+    """Store descriptors for cache groups represented by ``MambaSpec``.
+
+    This includes Mamba and GDN convolution/recurrent state caches.
+    """
+
+    store_format = "mamba_state"
+
+    def __init__(
+        self,
+        metadata: KeyMetadata,
+        block_size: int,
+        hash_block_size: int,
+        local_tp_size: int,
+        store_tp_size: int,
+        tp_rank: int,
+        layer_specs: Sequence[MambaSpec],
+    ) -> None:
+        super().__init__(
+            metadata,
+            block_size,
+            hash_block_size,
+            local_tp_size,
+            store_tp_size,
+            tp_rank,
+        )
+        self.layer_specs = tuple(layer_specs)
+        self._segments = tuple(
+            self._state_segments(spec, local_tp_size, self.shards_per_rank)
+            for spec in self.layer_specs
+        )
+
+    @staticmethod
+    def _state_segments(
+        spec: MambaSpec, local_tp_size: int, shards_per_rank: int
+    ) -> tuple[tuple[int, int, int], ...]:
+        conv = derive_mamba_conv_split(spec, local_tp_size, require_dim_first=False)
+        segments = []
+        dim_first = is_conv_state_dim_first()
+        num_rows = 1 if dim_first else conv.conv_rows
+        dim_bytes = conv.conv_dtype_size * (conv.conv_rows if dim_first else 1)
+        for row_index in range(num_rows):
+            offset = row_index * conv.local_conv_dim * conv.conv_dtype_size
+            for proj_index, proj_dim in enumerate(conv.local_proj_dims):
+                size = proj_dim * dim_bytes
+                unit_size = dim_bytes
+                if (
+                    spec.mamba_type == MambaAttentionBackendEnum.MAMBA2
+                    and proj_index > 0
+                ):
+                    unit_size *= spec.shapes[1][-1]
+                segments.append((offset, size, unit_size))
+                offset += size
+        state_offset = conv.ssm_sizes[0]
+        for state_index, shape in enumerate(spec.shapes[1:], start=1):
+            state_dtype: torch.dtype = spec.dtypes[state_index]
+            element_size = torch.empty((), dtype=state_dtype).element_size()
+            state_bytes = torch.Size(shape).numel() * element_size
+            unit_size = torch.Size(shape[1:]).numel() * element_size
+            segments.append((state_offset, state_bytes, unit_size))
+            state_offset += state_bytes
+        for segment in segments:
+            MambaStoreLayout._shard_slice(segment, 0, shards_per_rank)
+        return tuple(segments)
+
+    @staticmethod
+    def _shard_slice(
+        segment: tuple[int, int, int],
+        shard_index: int,
+        shards_per_rank: int,
+    ) -> tuple[int, int]:
+        offset, size, unit_size = segment
+        if size % unit_size:
+            raise ValueError("Mamba state segment contains a partial state unit")
+        unit_count = size // unit_size
+        if unit_count % shards_per_rank == 0:
+            shard_size = size // shards_per_rank
+            return offset + shard_index * shard_size, shard_size
+        if shards_per_rank % unit_count == 0:
+            # A local state unit can serve several finer Store shards.
+            copies_per_unit = shards_per_rank // unit_count
+            return offset + shard_index // copies_per_unit * unit_size, unit_size
+        raise ValueError("Mamba state units cannot be mapped evenly to Store shards")
+
+    @classmethod
+    def schema_fingerprint(
+        cls,
+        layer_specs: Sequence[MambaSpec],
+        local_tp_size: int,
+        store_tp_size: int,
+    ) -> str:
+        shards_per_rank = store_tp_size // local_tp_size
+        layers = tuple(
+            (
+                type(spec).__name__,
+                tuple(
+                    cls._shard_slice(segment, 0, shards_per_rank)[1]
+                    for segment in cls._state_segments(
+                        spec, local_tp_size, shards_per_rank
+                    )
+                ),
+                tuple(str(dtype) for dtype in spec.dtypes),
+                str(spec.mamba_type),
+                "DS" if is_conv_state_dim_first() else "SD",
+            )
+            for spec in layer_specs
+        )
+        return _schema_fingerprint(layers)
+
+    def register_kv_caches(
+        self,
+        kv_caches: Sequence[torch.Tensor],
+        num_blocks: int,
+    ) -> None:
+        templates: list[tuple[list[int], list[int], list[int]]] = [
+            ([], [], []) for _ in range(self.shards_per_rank)
+        ]
+        for cache, spec, segments in zip(
+            kv_caches, self.layer_specs, self._segments, strict=True
+        ):
+            if cache.ndim != 4 or tuple(cache.shape[:3]) != (num_blocks, 1, 1):
+                raise ValueError(
+                    "TP-shared MambaSpec store requires logical cache shape "
+                    "(num_blocks, 1, 1, state_bytes)"
+                )
+            element_size = cache.element_size()
+            if cache.stride(3) * element_size != 1:
+                raise ValueError(
+                    "TP-shared MambaSpec store requires packed state bytes"
+                )
+            block_stride = cache.stride(0) * element_size
+            content_bytes = cache.shape[3] * element_size
+            if content_bytes < spec.state_content_size_bytes:
+                raise ValueError("MambaSpec cache does not contain the complete state")
+            for shard_index, (addr_bases, block_strides, sizes) in enumerate(templates):
+                for segment in segments:
+                    offset, shard_size = self._shard_slice(
+                        segment, shard_index, self.shards_per_rank
+                    )
+                    addr_bases.append(cache.data_ptr() + offset)
+                    block_strides.append(block_stride)
+                    sizes.append(shard_size)
+        self._set_segment_templates([templates])
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 18d961abb286..6830baf46122 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
@@ -173,6 +173,13 @@ def __init__(
             if isinstance(group.kv_cache_spec, MambaSpec)
         }
         self._boundary_state_group_ids = frozenset(mamba_groups)
+        if (
+            self.save_decode_cache
+            and mamba_groups
+            and kv_cache_config.prefix_cache_retention_interval == 0
+        ):
+            # Decode checkpoints become reusable Store boundaries.
+            kv_cache_config.prefix_cache_retention_interval = self._block_size
 
         # Mirrors MooncakeStoreConnector._capacity_only.
         self._capacity_only = (
@@ -505,7 +512,7 @@ def build_connector_meta(
         if (
             block_state is not None
             and block_state.boundary_state_offloads
-            and not is_consumer
+            and can_process_cached
         ):
             self._handle_boundary_state_offloads(
                 block_state.boundary_state_offloads, meta
@@ -576,11 +583,7 @@ def _reference_save_blocks(self, meta: MooncakeStoreConnectorMetadata) -> None:
                 "A null block cannot back a boundary put"
             )
             if req_meta.token_len_chunk:
-                # Normal prefix save: pin all non-Mamba sources for retries.
-                # Every allocated block is referenced, not just the ones covering
-                # this job's token range: a rank resumes from its own last
-                # successful offset, which lags the scheduler's whenever a save was
-                # skipped or failed, so it may read anywhere below the range.
+                # Each rank may retry from an earlier successful save offset.
                 block_ids.extend(
                     block_id
                     for group_id, group in enumerate(req_meta.block_ids)
@@ -693,12 +696,10 @@ def _handle_boundary_state_offloads(
                 continue
             accepted: list[BoundaryPut] = []
             for group_id, block_id, boundary_tokens in entries:
-                # Every other group stops saving at the end of this prefill, so
-                # a mamba-only key past it can never complete a joint hybrid
-                # hit. `prefill_end_tokens` — not the original prompt length —
-                # is the boundary: a resumed request re-prefills and re-saves
-                # its previously generated tokens for every group.
-                if boundary_tokens > tracker.prefill_end_tokens:
+                is_decode_boundary = boundary_tokens > tracker.prefill_end_tokens
+                if is_decode_boundary and not self.save_decode_cache:
+                    continue
+                if not is_decode_boundary and self.kv_role == "kv_consumer":
                     continue
                 if block_id == NULL_BLOCK_ID:
                     continue
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 004b69286ea3..3561b66e66b5 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
@@ -20,7 +20,8 @@
 from collections.abc import Callable, Sequence
 from concurrent.futures import Future, ThreadPoolExecutor
 from contextlib import AbstractContextManager
-from typing import Any, TypeVar
+from dataclasses import dataclass
+from typing import Any, Literal, TypeVar, cast
 
 import torch
 import zmq
@@ -30,6 +31,7 @@
 from vllm.distributed import (
     get_dcp_group,
     get_pcp_group,
+    get_pp_group,
     get_tensor_model_parallel_rank,
     get_tensor_model_parallel_world_size,
 )
@@ -45,17 +47,21 @@
 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (  # noqa: E501
     BlobBlockHashes,
     ChunkedTokenDatabase,
-    KeyMetadata,
-    LBHNCStoreLayout,
-    LBNHCStoreLayout,
     MooncakeLookupResult,
     MooncakeStoreConnectorMetadata,
     MooncakeStoreWorkerMetadata,
-    PoolKey,
     ReqMeta,
-    StoreShardId,
     TailKeyBoundary,
-    TPShardedStoreLayout,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.layout import (  # noqa: E501
+    ATTENTION_STORE_LAYOUTS,
+    AttentionStoreLayout,
+    KeyMetadata,
+    MambaStoreLayout,
+    PoolKey,
+    StoreLayout,
+    StoreShardId,
+    TokenMajorStoreLayout,
 )
 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.protocol import (  # noqa: E501
     LOOKUP_MSG,
@@ -77,8 +83,9 @@
     resolve_kv_cache_block_sizes,
 )
 from vllm.v1.kv_cache_interface import (
-    FullAttentionSpec,
+    AttentionSpec,
     KVCacheConfig,
+    KVCacheGroupSpec,
     KVCacheSpec,
     MambaSpec,
     MLAAttentionSpec,
@@ -86,7 +93,6 @@
     UniformTypeKVCacheSpecs,
     group_kernel_blocks,
 )
-from vllm.v1.kv_cache_layout import KVCacheLayout
 
 from .metrics import MooncakeStoreConnectorStats
 
@@ -96,6 +102,38 @@
 _T = TypeVar("_T")
 
 
+@dataclass(frozen=True)
+class _StoreGroupPlan:
+    layer_specs: tuple[KVCacheSpec, ...]
+    local_block_size: int
+    logical_tp_size: int
+    store_shard_count: int
+    tp_rank: int
+    layout_cls: type[StoreLayout]
+    schema_fingerprint: str
+    store_chunk_size: int | None = None
+
+
+_STORE_SPEC_FAMILIES: dict[type[KVCacheSpec], Literal["attention", "mamba"]] = {
+    AttentionSpec: "attention",
+    MambaSpec: "mamba",
+}
+
+
+def _store_spec_family(spec: KVCacheSpec) -> Literal["attention", "mamba"] | None:
+    for spec_cls in type(spec).__mro__:
+        if family := _STORE_SPEC_FAMILIES.get(spec_cls):
+            return family
+    return None
+
+
+def _group_layer_specs(group: KVCacheGroupSpec) -> tuple[KVCacheSpec, ...]:
+    group_spec = group.kv_cache_spec
+    if isinstance(group_spec, UniformTypeKVCacheSpecs):
+        return tuple(group_spec.kv_cache_specs[name] for name in group.layer_names)
+    return (group_spec,) * len(group.layer_names)
+
+
 def resolve_store_tp_size(extra_config: dict[str, Any]) -> int | None:
     """Resolve the common Store TP requested by connector config."""
     if extra_config.get("enable_store_tp_lcm") is True:
@@ -112,6 +150,52 @@ def resolve_store_tp_size(extra_config: dict[str, Any]) -> int | None:
     return store_tp_size if type(store_tp_size) is int and store_tp_size > 0 else None
 
 
+def resolve_store_chunk_size(extra_config: dict[str, Any]) -> int | None:
+    """Resolve an explicitly shared Attention Store chunk size."""
+    store_chunk_size = extra_config.get("store_chunk_size")
+    return (
+        store_chunk_size
+        if type(store_chunk_size) is int and store_chunk_size > 0
+        else None
+    )
+
+
+def _resolve_attention_store_chunk(
+    group: KVCacheGroupSpec,
+    tp_size: int,
+    replication_factor: int,
+    requested_store_tp_size: int,
+    hash_block_size: int,
+    normalize_hybrid_chunks: bool,
+    requested_store_chunk_size: int | None = None,
+) -> tuple[tuple[AttentionSpec, ...], int, int, int] | None:
+    if tp_size % replication_factor:
+        return None
+    layer_specs = _group_layer_specs(group)
+    if not all(isinstance(spec, AttentionSpec) for spec in layer_specs):
+        return None
+    attention_specs = cast(tuple[AttentionSpec, ...], layer_specs)
+    logical_tp_size = tp_size // replication_factor
+    store_shard_count = AttentionStoreLayout.resolve_store_shard_count(
+        attention_specs, logical_tp_size, requested_store_tp_size
+    )
+    if store_shard_count is None:
+        return None
+    store_chunk_size = AttentionStoreLayout.resolve_store_chunk_size(
+        attention_specs,
+        requested_store_chunk_size=(
+            requested_store_chunk_size
+            if requested_store_chunk_size is not None
+            else hash_block_size
+            if normalize_hybrid_chunks
+            else None
+        ),
+    )
+    if store_chunk_size is None or store_chunk_size % hash_block_size:
+        return None
+    return attention_specs, logical_tp_size, store_shard_count, store_chunk_size
+
+
 def _rotate_list(values: list[_T], offset: int) -> list[_T]:
     return values[offset:] + values[:offset]
 
@@ -453,8 +537,7 @@ def __init__(
         self._store_pressure_active = False
         self._skip_store_requests: set[str] = set()
 
-        # Per-request high-water mark of tokens actually persisted; the next
-        # batch resumes here, so pressure-skipped or failed ranges are retried.
+        # Per-request high-water mark of complete hybrid checkpoints.
         self._saved_offset: dict[str, int] = {}
         # Retained only after a failed store so retry events can recover the
         # token suffix without full snapshots on the normal path.
@@ -508,12 +591,14 @@ def take_completed_saves(self) -> dict[int, int]:
             self._completed_saves = {}
         return completed
 
-    def _record_saved(self, req_meta: ReqMeta, token_len: int) -> None:
+    def _record_saved_offset(self, req_meta: ReqMeta, offset: int) -> None:
         # Guard on job liveness so neither a concurrent finish/preempt pop nor a
         # stale job's offset is written back over the live generation's.
         with self.done_task_lock:
             if req_meta.store_job_id in self.stored_requests.get(req_meta.req_id, ()):
-                self._saved_offset[req_meta.req_id] = token_len
+                self._saved_offset[req_meta.req_id] = max(
+                    self._saved_offset.get(req_meta.req_id, 0), offset
+                )
 
     def _get_retry_token_ids(self, req_meta: ReqMeta) -> tuple[int, list[int]] | None:
         """Return retry state only if this store job is still live."""
@@ -571,17 +656,10 @@ def _clear_store_pressure(self) -> bool:
     def _boundary_puts(
         self, req_meta: ReqMeta
     ) -> list[tuple[str, list[int], list[int], KeyMetadata]]:
-        """Puts for ``req_meta.boundary_puts``: each block under the hash of
-        the prefix ending at its ``num_tokens``.
-
-        The scheduler resolves and pins every block, so none is looked up
-        positionally here. That matters for mamba "align" states: ``store_mask``
-        masks mamba groups out of the normal save, which resolves a chunk's
-        address as ``req_meta.block_ids[g][start // block_size]``. An align-mode
-        table is mutated in place (a superseded state block is freed and
-        nulled; speculative blocks relocate) without telling the connector, so
-        a retried positional read could hit a block that now belongs to another
-        request. The handed-off block *is* the state.
+        """Build Store chunks from scheduler-resolved, pinned blocks.
+
+        Attention pages may contain several chunks. Mamba entries retain the
+        exact handed-off state block under its checkpoint hash.
         """
         hash_block_size = self.coord.hash_block_size
         puts: list[tuple[str, list[int], list[int], KeyMetadata]] = []
@@ -592,7 +670,6 @@ def _boundary_puts(
                 or block_id == NULL_BLOCK_ID
             ):
                 continue
-            # A negative index would silently key the block by the last hash.
             assert num_tokens % hash_block_size == 0, (
                 f"Boundary put at {num_tokens} tokens is not a multiple of the "
                 f"hash block size {hash_block_size}"
@@ -604,12 +681,28 @@ def _boundary_puts(
             # Distribute across ranks by the same rule as normal chunks.
             put_step = self.group_put_steps[group_id]
             put_step_rank = (self.tp_rank + group_id) % put_step
-            if (cdiv(num_tokens, db.block_size) - 1) % put_step != put_step_rank:
-                continue
-            addr, size = db.prepare_value_for_block(block_id)
-            puts.append(
-                (db.key_for(req_meta.block_hashes[hash_idx]), addr, size, db.metadata)
-            )
+            shard_ids = db.store_layout.local_shard_ids
+            block_size = db.store_layout.block_size
+            block_start = (cdiv(num_tokens, block_size) - 1) * block_size
+            for start in range(block_start, num_tokens, db.chunk_size):
+                if start // db.chunk_size % put_step != put_step_rank:
+                    continue
+                end = min(start + db.chunk_size, num_tokens)
+                key_hash = req_meta.block_hashes[end // hash_block_size - 1]
+                addrs, sizes, _ = db.store_layout.prepare_values(
+                    [(start - block_start, end - block_start)] * len(shard_ids),
+                    [block_id],
+                    shard_ids,
+                )
+                for shard_id, addr, size in zip(shard_ids, addrs, sizes, strict=True):
+                    puts.append(
+                        (
+                            db.store_layout.key_for(shard_id, key_hash),
+                            addr,
+                            size,
+                            db.metadata,
+                        )
+                    )
         return puts
 
     def _maybe_offload_boundary_states(self, req_meta: ReqMeta) -> bool:
@@ -728,15 +821,15 @@ def _handle_request(self, req_meta: ReqMeta):
         # The single `finally` is the only way out, so the scheduler releases
         # this job's GPU block references however the job ends.
         save_completed = False
-        token_len = 0
+        available_token_len = 0
+        checkpoint_len = 0
         req_id = req_meta.req_id
         event_token_ids = req_meta.token_ids
         token_ids_start = req_meta.token_ids_start
         try:
-            # Cache hits are always a multiple of ``lcm_block_size`` tokens,
-            # which is also ``store_mask``'s precondition.
             lcm_block_size = self.coord.lcm_block_size
-            token_len = req_meta.token_len_chunk // lcm_block_size * lcm_block_size
+            available_token_len = req_meta.token_len_chunk
+            checkpoint_len = available_token_len // lcm_block_size * lcm_block_size
             block_ids_per_group = req_meta.block_ids
             current_event = req_meta.current_event
 
@@ -764,18 +857,24 @@ def _handle_request(self, req_meta: ReqMeta):
             if not self._maybe_offload_boundary_states(req_meta):
                 return
 
-            if token_len == 0:
-                return
-
-            # Resume from where this rank left off; only the new suffix is saved.
             save_start = self._saved_offset.get(req_id, 0)
+            if checkpoint_len <= save_start:
+                return
 
             # Within each lcm region only per-spec relevant chunks are loaded
             # (e.g., SWA or linear attn), so mask out irrelevant chunks
+            replay_boundary = req_meta.prefill_end_tokens
+            if (
+                self.kv_role == "kv_consumer"
+                and replay_boundary is not None
+                and checkpoint_len > replay_boundary
+            ):
+                replay_boundary = checkpoint_len + 1
+            # Successful saves advance by complete scheduler-aligned intervals.
             store_masks = self.coord.store_mask(
-                token_len,
+                checkpoint_len,
                 save_start,
-                num_prompt_tokens=req_meta.prefill_end_tokens,
+                num_prompt_tokens=replay_boundary,
             )
 
             starts: list[int] = []
@@ -786,26 +885,36 @@ def _handle_request(self, req_meta: ReqMeta):
             )
             group_indices: list[int] = []
             store_shard_ids: list[StoreShardId] = []
+            missing_source_blocks = False
             for g_idx, db in enumerate(self.token_databases):
                 if not self.group_participates[g_idx]:
                     continue
                 # Rotate the stride phase per group to balance load across ranks.
                 put_step = self.group_put_steps[g_idx]
                 put_step_rank = (self.tp_rank + g_idx) % put_step
+                store_mask = store_masks[g_idx]
                 group_blocks = block_ids_per_group[g_idx]
                 for start, end, block_hash in db.process_tokens(
-                    token_len,
+                    checkpoint_len,
                     req_meta.block_hashes,
                     mask_num=save_start,
-                    chunk_mask=store_masks[g_idx],
                     put_step=put_step,
                     put_step_rank=put_step_rank,
                 ):
-                    block_idx = start // db.block_size
-                    group_blocks = block_ids_per_group[g_idx]
+                    block_idx = start // db.store_layout.block_size
+                    if store_mask is not None:
+                        mask_idx = block_idx - cdiv(
+                            save_start, db.store_layout.block_size
+                        )
+                        if (
+                            not 0 <= mask_idx < len(store_mask)
+                            or not store_mask[mask_idx]
+                        ):
+                            continue
                     if block_idx >= len(group_blocks) or (
                         group_blocks[block_idx] == NULL_BLOCK_ID
                     ):
+                        missing_source_blocks = True
                         logger.debug(
                             "Skipping unavailable Mooncake store source block "
                             "(req=%s, group=%d, block=%d)",
@@ -824,8 +933,7 @@ def _handle_request(self, req_meta: ReqMeta):
                             event_specs.append((start, end, g_idx, block_hash))
 
             if not keys:
-                self._record_saved(req_meta, token_len)
-                save_completed = True
+                save_completed = not missing_source_blocks
                 return
 
             # Check which blocks already exist (dedup)
@@ -851,8 +959,7 @@ def _handle_request(self, req_meta: ReqMeta):
             ]
 
             if not missing_indices:
-                self._record_saved(req_meta, token_len)
-                save_completed = True
+                save_completed = not missing_source_blocks
                 return
 
             if len(missing_indices) != len(keys):
@@ -941,6 +1048,7 @@ def _handle_request(self, req_meta: ReqMeta):
                 )
                 logger.error("Failed to put key %s, error: %s", keys, e)
                 put_had_exception = True
+                failed_indices = set(range(len(keys)))
             else:
                 failed_indices = {i for i, value in enumerate(res) if value < 0}
                 self._record_operation(
@@ -973,8 +1081,11 @@ def _handle_request(self, req_meta: ReqMeta):
                         req_id,
                     )
 
-            if not put_had_exception and not failed_indices:
-                self._record_saved(req_meta, token_len)
+            if (
+                not put_had_exception
+                and not failed_indices
+                and not missing_source_blocks
+            ):
                 save_completed = True
                 if self._clear_store_pressure():
                     logger.info(
@@ -1023,7 +1134,7 @@ def _handle_request(self, req_meta: ReqMeta):
                                 else None
                             ),
                             token_ids=token_ids,
-                            block_size=db.block_size,
+                            block_size=db.chunk_size,
                             lora_id=None,
                             medium="cpu",
                             lora_name=None,
@@ -1034,7 +1145,9 @@ def _handle_request(self, req_meta: ReqMeta):
             if self.enable_kv_event and stored_events:
                 self.update_kv_event(stored_events)
         finally:
-            if self.enable_kv_event and token_len:
+            if save_completed:
+                self._record_saved_offset(req_meta, checkpoint_len)
+            if self.enable_kv_event and available_token_len:
                 self._update_retry_token_ids(
                     req_meta,
                     save_completed,
@@ -1150,10 +1263,12 @@ def _handle_request(self, req_meta: ReqMeta):
             chunks: list[tuple[int, int]] = []
             store_shard_ids: list[StoreShardId] = []
             for start, end, block_hash in db.process_tokens(
-                token_len, req_meta.block_hashes, mask_num
+                token_len,
+                req_meta.block_hashes,
+                mask_num,
             ):
-                chunk_idx = start // db.block_size
-                if chunk_idx >= len(mask) or not mask[chunk_idx]:
+                block_index = start // db.store_layout.block_size
+                if block_index >= len(mask) or not mask[block_index]:
                     continue
                 boundary_tokens = (
                     tail_key_boundaries.get(g_idx) if end == token_len else None
@@ -1175,6 +1290,11 @@ def _handle_request(self, req_meta: ReqMeta):
             size_list.extend(g_sizes)
             block_id_list.extend(g_block_ids)
 
+        if not key_list:
+            self.set_finished_request(req_id)
+            self.request_queue.task_done()
+            return
+
         # Rotate aligned lists by tp_rank for load balancing.
         rotation = self.tp_rank % len(key_list)
         key_list_c = _rotate_list(key_list, rotation)
@@ -1391,8 +1511,9 @@ def __init__(
         assert kv_role is not None
         self.kv_role = kv_role
         extra_config = vllm_config.kv_transfer_config.kv_connector_extra_config
-        self.can_put = self.kv_role in ("kv_producer", "kv_both") or (
-            extra_config.get("save_decode_cache", False)
+        self.save_decode_cache = extra_config.get("save_decode_cache", False)
+        self.can_put = (
+            self.kv_role in ("kv_producer", "kv_both") or self.save_decode_cache
         )
         self.load_async = extra_config.get("load_async", True)
         # Mirrors MooncakeStoreConnector._capacity_only.
@@ -1492,6 +1613,10 @@ def __init__(
         self.kv_recv_threads: list[KVCacheStoreRecvingThread] = []
         self.num_recv_threads = max(1, envs.VLLM_MOONCAKE_LOAD_RECV_THREADS)
         self.recv_request_queue: queue.Queue[ReqMeta] = queue.Queue()
+        # Device KV regions (addr -> len) registered with the store.
+        self._device_regions: dict[int, int] = {}
+        # Device regions unregistered for sleep and not yet registered again.
+        self._released_regions: set[int] = set()
         self.finished_store_req: set[str] = set()
         self._kv_connector_stats_lock = threading.Lock()
         self.kv_connector_stats = MooncakeStoreConnectorStats()
@@ -1508,11 +1633,20 @@ def __init__(
             )
             return
 
+        if (
+            self.save_decode_cache
+            and kv_cache_config.has_mamba_layers
+            and kv_cache_config.prefix_cache_retention_interval == 0
+        ):
+            kv_cache_config.prefix_cache_retention_interval = self.block_size
         self.coord = MooncakeStoreCoordinator.from_kv_cache_config(
             kv_cache_config, vllm_config, self.block_size, self.hash_block_size
         )
         self._kv_cache_groups = self.coord.kv_cache_groups
-        self.store_tp_size, store_namespace, store_layout_cls = (
+        self._group_tp_replication_factors = (
+            self._compute_group_tp_replication_factors()
+        )
+        self.store_tp_size, store_namespace, store_group_plans = (
             self._select_store_layout(extra_config)
         )
         metadata = KeyMetadata(
@@ -1528,81 +1662,102 @@ def __init__(
             ),
             store_namespace=store_namespace,
         )
-        self._group_tp_replication_factors: tuple[int, ...] = (
-            self._compute_group_tp_replication_factors()
-        )
-        self.token_dbs = self._build_token_databases(metadata, store_layout_cls)
+        self.token_dbs = self._build_token_databases(metadata, store_group_plans)
         self._init_lookup_key_prefixes()
 
-    def _supports_tp_sharded_store_layout(
-        self,
-        layout_cls: type[TPShardedStoreLayout] | None,
-        extra_config: dict[str, Any],
-    ) -> bool:
-        if (
-            layout_cls is None
-            or self.pcp_size != 1
-            or self.dcp_size != 1
-            or len(self._kv_cache_groups) != 1
-        ):
-            return False
-
-        # Subclasses may use different sharding or cache-lifetime semantics.
-        return (
-            type(self._kv_cache_groups[0].kv_cache_spec) is FullAttentionSpec
-            and str(extra_config.get("enable_cross_layers_blocks", "False")).lower()
-            != "true"
-        )
-
     def _select_store_layout(
         self, extra_config: dict[str, Any]
-    ) -> tuple[int | None, str, type[TPShardedStoreLayout] | None]:
-        """Select the opt-in TP layout and its Store namespace."""
+    ) -> tuple[int | None, str, tuple[_StoreGroupPlan | None, ...] | None]:
+        """Select compatible Store layouts for every KV-cache group."""
         lcm_store_tp_enabled = extra_config.get("enable_store_tp_lcm") is True
         store_tp_requested = (
             lcm_store_tp_enabled or extra_config.get("store_tp_size") is not None
         )
+        self.requested_store_chunk_size = resolve_store_chunk_size(extra_config)
         if not store_tp_requested:
             return None, "", None
 
         requested_store_tp_size = resolve_store_tp_size(extra_config)
-        cache_layout = self.cache_config.get_resolved_kv_cache_layout()
-        layout_cls: type[TPShardedStoreLayout] | None = {
-            KVCacheLayout.LBHNC: LBHNCStoreLayout,
-            KVCacheLayout.LBNHC: LBNHCStoreLayout,
-        }.get(cache_layout)
-
+        store_chunk_requested = extra_config.get("store_chunk_size") is not None
+        hybrid_store_requested = len(self._kv_cache_groups) > 1
+        prefix_match_unit = self.cache_config.prefix_match_unit
         if (
-            requested_store_tp_size is not None
-            and requested_store_tp_size >= self.tp_size
-            and requested_store_tp_size % self.tp_size == 0
-            and self._supports_tp_sharded_store_layout(layout_cls, extra_config)
+            hybrid_store_requested
+            and prefix_match_unit is not None
+            and not store_chunk_requested
         ):
-            assert layout_cls is not None
-            if self.num_kv_head % requested_store_tp_size == 0:
-                if layout_cls is LBNHCStoreLayout:
-                    logger.warning_once(
-                        "Mooncake Store TP sharding is using the LBNHC (NHC) "
-                        "KV cache layout, which creates many transfer segments "
-                        "and may significantly reduce PUT/GET performance. Use "
-                        "LBHNC (HNC) when supported by the attention backend."
-                    )
-                return (
-                    requested_store_tp_size,
-                    layout_cls.shared_namespace(requested_store_tp_size, self.pp_size),
-                    layout_cls,
+            self.requested_store_chunk_size = prefix_match_unit
+
+        cache_layout = self.cache_config.get_resolved_kv_cache_layout()
+        store_layout_cls = ATTENTION_STORE_LAYOUTS.get(cache_layout)
+
+        fallback_reason: str | None = None
+        if requested_store_tp_size is None:
+            fallback_reason = "invalid Store TP configuration"
+        elif self._kv_cache_config.hisparse_host_num_blocks is not None:
+            fallback_reason = (
+                "HiSparse host cache registration requires rank-local layout"
+            )
+        elif hybrid_store_requested and prefix_match_unit is None:
+            fallback_reason = "Hybrid Store TP sharing requires prefix_match_unit"
+        elif store_chunk_requested and self.requested_store_chunk_size is None:
+            fallback_reason = "invalid Store chunk size"
+        elif (
+            hybrid_store_requested
+            and self.requested_store_chunk_size is not None
+            and self.requested_store_chunk_size % self.hash_block_size
+        ):
+            fallback_reason = "Store chunk size is not aligned to prefix_match_unit"
+        elif store_layout_cls is None:
+            fallback_reason = f"unsupported KV cache layout {cache_layout}"
+        elif self.pcp_size != 1 or self.dcp_size != 1:
+            fallback_reason = "PCP or DCP is enabled"
+        elif str(extra_config.get("enable_cross_layers_blocks", "False")).lower() == (
+            "true"
+        ):
+            fallback_reason = "cross-layer KV blocks are enabled"
+        elif (
+            requested_store_tp_size < self.tp_size
+            or requested_store_tp_size % self.tp_size
+        ):
+            fallback_reason = "Store TP is not divisible by local TP"
+
+        plans: list[_StoreGroupPlan | None] = []
+        if fallback_reason is None:
+            assert requested_store_tp_size is not None
+            assert store_layout_cls is not None
+            for group_index, (group, replication_factor) in enumerate(
+                zip(
+                    self._kv_cache_groups,
+                    self._group_tp_replication_factors,
+                    strict=True,
                 )
-            if self.num_kv_head == 1:
-                logger.info(
-                    "Mooncake heterogeneous-TP store sharing uses the replicated "
-                    "MQA layout for store_tp_size=%d",
+            ):
+                if not group.kv_cache_spec.prefix_cacheable:
+                    plans.append(None)
+                    continue
+                plan = self._make_store_group_plan(
+                    group,
                     requested_store_tp_size,
+                    store_layout_cls,
+                    replication_factor,
                 )
-                return (
-                    None,
-                    f"@store_pp:{self.pp_size}@store_format:tp_shared_mqa",
-                    None,
+                if plan is None:
+                    fallback_reason = (
+                        f"cache group {group_index} has an incompatible Store layout"
+                    )
+                    break
+                plans.append(plan)
+
+        if fallback_reason is None:
+            if store_layout_cls is TokenMajorStoreLayout:
+                logger.warning_once(
+                    "Mooncake Store TP sharding is using a token-major KV cache "
+                    "layout, which creates many transfer segments and may "
+                    "significantly reduce PUT/GET performance. Use a head-major "
+                    "layout when supported by the attention backend."
                 )
+            return requested_store_tp_size, "", tuple(plans)
 
         store_namespace = (
             f"@store_pp:{self.pp_size}@store_format:"
@@ -1614,57 +1769,150 @@ def _select_store_layout(
             else extra_config.get("store_tp_size")
         )
         logger.warning(
-            "Mooncake heterogeneous-TP store sharing is disabled for "
-            "Store TP configuration %r with KV layout %s; using a "
-            "compatibility-namespaced rank-local store layout",
+            "Store TP configuration %r with KV layout %s falls back because %s; "
+            "using namespace %s",
             requested_topology,
             cache_layout,
+            fallback_reason,
+            store_namespace,
         )
         return None, store_namespace, None
 
     def _build_token_databases(
         self,
         metadata: KeyMetadata,
-        layout_cls: type[TPShardedStoreLayout] | None,
+        plans: tuple[_StoreGroupPlan | None, ...] | None,
     ) -> list[ChunkedTokenDatabase]:
         """Construct token databases and their Store layouts."""
         token_dbs: list[ChunkedTokenDatabase] = []
         for group_idx, group in enumerate(self._kv_cache_groups):
+            spec = group.kv_cache_spec
             hash_block_size = (
-                self.hash_block_size
-                if group.kv_cache_spec.prefix_cacheable
-                else group.kv_cache_spec.block_size
+                self.hash_block_size if spec.prefix_cacheable else spec.block_size
+            )
+            plan = plans[group_idx] if plans is not None else None
+            if plan is None:
+                token_dbs.append(
+                    ChunkedTokenDatabase(
+                        dataclasses.replace(
+                            metadata,
+                            group_id=group_idx,
+                            tp_rank=(
+                                self.tp_rank
+                                // self._group_tp_replication_factors[group_idx]
+                            ),
+                        ),
+                        spec.block_size,
+                        hash_block_size=hash_block_size,
+                    )
+                )
+                continue
+            assert self.store_tp_size is not None
+            group_namespace = (
+                f"@store_tp:{self.store_tp_size}@store_pp:{self.pp_size}"
+                f"@store_format:{plan.layout_cls.store_format}"
+                f"@store_schema:{plan.schema_fingerprint}"
             )
-            group_tp_rank = self.tp_rank
-            if layout_cls is None:
-                group_tp_rank //= self._group_tp_replication_factors[group_idx]
             group_metadata = dataclasses.replace(
                 metadata,
                 group_id=group_idx,
-                tp_rank=group_tp_rank,
+                tp_rank=plan.tp_rank,
+                store_namespace=group_namespace,
             )
-            store_layout: TPShardedStoreLayout | None = None
-            if layout_cls is not None:
-                assert self.store_tp_size is not None
-                store_layout = layout_cls(
+            if issubclass(plan.layout_cls, AttentionStoreLayout):
+                attention_layout_cls = cast(type[AttentionStoreLayout], plan.layout_cls)
+                store_layout: StoreLayout = attention_layout_cls(
+                    group_metadata,
+                    plan.local_block_size,
+                    hash_block_size,
+                    local_tp_size=plan.logical_tp_size,
+                    store_tp_size=plan.store_shard_count,
+                    tp_rank=plan.tp_rank,
+                    layer_specs=cast(tuple[AttentionSpec, ...], plan.layer_specs),
+                    store_chunk_size=plan.store_chunk_size,
+                )
+            else:
+                store_layout = MambaStoreLayout(
                     group_metadata,
-                    group.kv_cache_spec.block_size,
+                    plan.local_block_size,
                     hash_block_size,
-                    local_tp_size=self.tp_size,
-                    store_tp_size=self.store_tp_size,
-                    tp_rank=self.tp_rank,
-                    num_kv_heads=self.num_kv_head,
+                    local_tp_size=plan.logical_tp_size,
+                    store_tp_size=plan.store_shard_count,
+                    tp_rank=plan.tp_rank,
+                    layer_specs=cast(tuple[MambaSpec, ...], plan.layer_specs),
                 )
             token_dbs.append(
                 ChunkedTokenDatabase(
                     group_metadata,
-                    group.kv_cache_spec.block_size,
+                    plan.store_chunk_size or plan.local_block_size,
                     hash_block_size=hash_block_size,
                     store_layout=store_layout,
                 )
             )
         return token_dbs
 
+    def _make_store_group_plan(
+        self,
+        group: KVCacheGroupSpec,
+        requested_store_tp_size: int,
+        attention_layout_cls: type[AttentionStoreLayout],
+        replication_factor: int,
+    ) -> _StoreGroupPlan | None:
+        layer_specs = _group_layer_specs(group)
+        families = {_store_spec_family(spec) for spec in layer_specs}
+        if len(families) != 1 or None in families:
+            return None
+        family = families.pop()
+        assert family is not None
+
+        if family == "attention":
+            resolved = _resolve_attention_store_chunk(
+                group,
+                self.tp_size,
+                replication_factor,
+                requested_store_tp_size,
+                self.hash_block_size,
+                len(self._kv_cache_groups) > 1,
+                self.requested_store_chunk_size,
+            )
+            if resolved is None:
+                return None
+            attention_specs, logical_tp_size, store_shard_count, store_chunk_size = (
+                resolved
+            )
+            return _StoreGroupPlan(
+                layer_specs=layer_specs,
+                local_block_size=group.kv_cache_spec.block_size,
+                logical_tp_size=logical_tp_size,
+                store_shard_count=store_shard_count,
+                tp_rank=self.tp_rank // replication_factor,
+                layout_cls=attention_layout_cls,
+                schema_fingerprint=AttentionStoreLayout.schema_fingerprint(
+                    attention_specs,
+                    logical_tp_size,
+                    store_shard_count,
+                    store_chunk_size,
+                ),
+                store_chunk_size=store_chunk_size,
+            )
+
+        mamba_specs = cast(tuple[MambaSpec, ...], layer_specs)
+        try:
+            schema_fingerprint = MambaStoreLayout.schema_fingerprint(
+                mamba_specs, self.tp_size, requested_store_tp_size
+            )
+        except (AssertionError, IndexError, NotImplementedError, ValueError):
+            return None
+        return _StoreGroupPlan(
+            layer_specs=layer_specs,
+            local_block_size=group.kv_cache_spec.block_size,
+            logical_tp_size=self.tp_size,
+            store_shard_count=requested_store_tp_size,
+            tp_rank=self.tp_rank,
+            layout_cls=MambaStoreLayout,
+            schema_fingerprint=schema_fingerprint,
+        )
+
     def get_mem_pool_context(self) -> AbstractContextManager | None:
         """Return a context manager for the custom MemPool, or None if
         no custom pool is configured."""
@@ -1680,23 +1928,19 @@ def _spec_tp_replication_factor(self, spec: KVCacheSpec) -> int:
             if isinstance(spec, UniformTypeKVCacheSpecs)
             else (spec,)
         )
-        # Any rank-specific state makes the whole packed value rank-specific.
         if any(isinstance(inner, MambaSpec) for inner in inner_specs):
             return 1
-        # A pure MLA packed value is replicated on every TP rank.
-        if all(
-            isinstance(inner, (MLAAttentionSpec, SlidingWindowMLASpec))
+        factors = [
+            (
+                self.tp_size
+                if isinstance(inner, (MLAAttentionSpec, SlidingWindowMLASpec))
+                else max(1, self.tp_size // self.num_kv_head)
+            )
             for inner in inner_specs
-        ):
-            return self.tp_size
-        return max(1, self.tp_size // self.num_kv_head)
+        ]
+        return min(factors, default=1)
 
     def _compute_group_tp_replication_factors(self) -> tuple[int, ...]:
-        """Return the number of byte-identical TP replicas per cache group.
-
-        DCP and Mamba use 1; MLA uses ``tp_size``; GQA uses
-        ``tp_size // num_kv_head``.
-        """
         return tuple(
             self._spec_tp_replication_factor(group.kv_cache_spec)
             for group in self._kv_cache_groups
@@ -1707,16 +1951,14 @@ def rank_namespaces(factor: int) -> tuple[tuple[int, int, int, int], ...]:
             if self.dcp_size > 1:
                 # DCP is a TP subdivision: dcp_rank == tp_rank % dcp_size.
                 return tuple(
-                    (tp_rank, pcp_rank, tp_rank % self.dcp_size, pp_rank)
+                    (tp_rank, pcp_rank, tp_rank % self.dcp_size, self.pp_rank)
                     for pcp_rank in range(self.pcp_size)
                     for tp_rank in range(self.tp_size)
-                    for pp_rank in range(self.pp_size)
                 )
             return tuple(
-                (shard_rank, pcp_rank, 0, pp_rank)
+                (shard_rank, pcp_rank, 0, self.pp_rank)
                 for pcp_rank in range(self.pcp_size)
                 for shard_rank in range(self.tp_size // factor)
-                for pp_rank in range(self.pp_size)
             )
 
         self._lookup_key_prefixes = tuple(
@@ -1725,6 +1967,18 @@ def rank_namespaces(factor: int) -> tuple[tuple[int, int, int, int], ...]:
             )
             for g_idx, db in enumerate(self.token_dbs)
         )
+        if self.pp_size > 1:
+            # Each stage's schema describes its own layers, not the whole model.
+            stage_prefixes: list[tuple[tuple[str, ...], ...]] = [()] * self.pp_size
+            torch.distributed.all_gather_object(
+                stage_prefixes,
+                self._lookup_key_prefixes,
+                group=get_pp_group().cpu_group,
+            )
+            self._lookup_key_prefixes = tuple(
+                tuple(prefix for stage in stage_prefixes for prefix in stage[g_idx])
+                for g_idx in range(len(self.token_dbs))
+            )
 
     def register_kv_caches(
         self,
@@ -1737,13 +1991,11 @@ 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:
+        # A layer's cache is one tensor, or a list of state tensors that share
+        # one raw storage (Mamba).
+        def _layer_views(v: torch.Tensor | list[torch.Tensor]) -> list[torch.Tensor]:
             assert isinstance(v, torch.Tensor | list)
-            return v if isinstance(v, torch.Tensor) else v[0]
+            return [v] if isinstance(v, torch.Tensor) else v
 
         assert self.cache_config.num_gpu_blocks is not None
         self.num_blocks = self.cache_config.num_gpu_blocks
@@ -1756,10 +2008,16 @@ def _repr_tensor(v: torch.Tensor | list[torch.Tensor]) -> torch.Tensor:
                 for layer_name in group.layer_names
             }
             seen_storage_ptrs: set[int] = set()
-            cache_tensors = [
-                group_kernel_blocks(_repr_tensor(cache), self.num_blocks)
+            cache_by_layer = {
+                layer_name: [
+                    group_kernel_blocks(view, self.num_blocks)
+                    for view in _layer_views(cache)
+                ]
                 for layer_name, cache in kv_caches.items()
                 if layer_name in store_layer_names
+            }
+            cache_tensors = [
+                cache for views in cache_by_layer.values() for cache in views
             ]
             for cache in cache_tensors:
                 cache_storage = cache.untyped_storage()
@@ -1776,8 +2034,28 @@ def _repr_tensor(v: torch.Tensor | list[torch.Tensor]) -> torch.Tensor:
                         region_len,
                         ret,
                     )
-            for db in self.token_dbs:
-                db.store_layout.register_kv_caches(cache_tensors, self.num_blocks)
+                else:
+                    self._device_regions[base_addr] = region_len
+            for group_id, (group, db) in enumerate(
+                zip(self._kv_cache_groups, self.token_dbs, strict=True)
+            ):
+                group_views = [
+                    cache_by_layer[name]
+                    for name in group.layer_names
+                    if name in cache_by_layer
+                ]
+                if self.store_tp_size is not None:
+                    if len(group_views) != len(group.layer_names):
+                        raise ValueError(
+                            "Missing KV cache tensors for TP-shared cache group "
+                            f"{group_id}"
+                        )
+                    group_tensors = [views[0] for views in group_views]
+                else:
+                    group_tensors = [cache for views in group_views for cache in views]
+                db.store_layout.register_kv_caches(
+                    group_tensors or cache_tensors, self.num_blocks
+                )
 
         registered_buffers: dict[int, int] = {}
         num_host_segments = 0
@@ -1803,7 +2081,7 @@ def _repr_tensor(v: torch.Tensor | list[torch.Tensor]) -> torch.Tensor:
             addrs: list[int] = []
             block_lens: list[int] = []
             for layer_name, value in group_caches:
-                cache = _repr_tensor(value)
+                cache = _layer_views(value)[0]
                 is_host_resident = group.host_resident
                 if is_host_resident:
                     num_blocks = self._kv_cache_config.hisparse_host_num_blocks
@@ -1834,6 +2112,8 @@ def _repr_tensor(v: torch.Tensor | list[torch.Tensor]) -> torch.Tensor:
                             f"{base_addr:#x} len {region_len} ({mem_kind}): {ret}"
                         )
                     registered_buffers[base_addr] = region_len
+                    if not is_host_resident:
+                        self._device_regions[base_addr] = region_len
                 elif registered_len != region_len:
                     raise ValueError(
                         f"KV cache views at {base_addr:#x} expose inconsistent "
@@ -1934,6 +2214,20 @@ def _repr_tensor(v: torch.Tensor | list[torch.Tensor]) -> torch.Tensor:
             "Started %d Mooncake KV-load receive thread(s)", self.num_recv_threads
         )
 
+    def unregister_kv_caches(self) -> None:
+        """Unregister device KV buffers after the engine has drained transfers."""
+        for addr in self._device_regions.keys() - self._released_regions:
+            if self.store.unregister_buffer(addr) != 0:
+                raise RuntimeError(f"Mooncake unregister_buffer failed for {addr:#x}")
+            self._released_regions.add(addr)
+
+    def restore_kv_caches(self) -> None:
+        """Register the released regions again: same addresses, new pages."""
+        for addr in list(self._released_regions):
+            if self.store.register_buffer(addr, self._device_regions[addr]) != 0:
+                raise RuntimeError(f"Mooncake register_buffer failed for {addr:#x}")
+            self._released_regions.discard(addr)
+
     def start_load_kv(self, metadata: MooncakeStoreConnectorMetadata):
         """Issue async loads.
 
@@ -2109,44 +2403,35 @@ def lookup(
             return MooncakeLookupResult(0)
 
         # Build per-(group, hash) candidate keys expanded across rank namespaces.
-        # candidate_meta stores the (group, hash_bytes) for key slice.
+        # candidate_meta stores (group, hash_bytes, end_token) for each key slice.
         candidate_keys: list[str] = []
-        candidate_meta: list[tuple[int, bytes]] = []
+        candidate_meta: list[tuple[int, bytes, int]] = []
         fine_grained = self.coord.enable_partial_hash_hits
         lookup_masks = None if fine_grained else self.coord.lookup_mask(token_len)
         for g_idx, db in enumerate(self.token_dbs):
             if not self._kv_cache_groups[g_idx].kv_cache_spec.prefix_cacheable:
                 continue
-            spec_block_size = db.block_size
+            spec_block_size = self._kv_cache_groups[g_idx].kv_cache_spec.block_size
             key_prefixes = self._lookup_key_prefixes[g_idx]
-            if fine_grained:
-                max_units = min(len(block_hashes), token_len // self.hash_block_size)
-                unit_ids: range | list[int] = range(max_units)
-                group_hashes: Sequence[BlockHash] = block_hashes
-            else:
-                lookup_mask = lookup_masks[g_idx]  # type: ignore[index]
-                group_hashes = self.coord.block_hashes_for_spec(
-                    block_hashes, self._kv_cache_groups[g_idx].kv_cache_spec
-                )
-                max_chunks = min(len(group_hashes), cdiv(token_len, spec_block_size))
-                mask_limit = (
-                    max_chunks
-                    if lookup_mask is None
-                    else min(max_chunks, len(lookup_mask))
-                )
-                unit_ids = [
-                    chunk_id
-                    for chunk_id in range(mask_limit)
-                    if lookup_mask is None or lookup_mask[chunk_id]
-                ]
-            for chunk_id in unit_ids:
-                h = group_hashes[chunk_id]
+            lookup_mask = None if lookup_masks is None else lookup_masks[g_idx]
+            unit_size = self.hash_block_size if fine_grained else db.chunk_size
+            max_units = (
+                min(token_len, len(block_hashes) * self.hash_block_size) // unit_size
+            )
+            for unit_id in range(max_units):
+                block_id = unit_id * unit_size // spec_block_size
+                if lookup_mask is not None and (
+                    block_id >= len(lookup_mask) or not lookup_mask[block_id]
+                ):
+                    continue
+                end_token = (unit_id + 1) * unit_size
+                h = block_hashes[end_token // self.hash_block_size - 1]
                 hash_hex = h.hex()
                 for key_prefix in key_prefixes:
                     candidate_keys.append(
                         PoolKey.build_key_string(key_prefix, hash_hex)
                     )
-                candidate_meta.append((g_idx, bytes(h)))
+                candidate_meta.append((g_idx, bytes(h), end_token))
 
         if not candidate_keys:
             return MooncakeLookupResult(0)
@@ -2175,17 +2460,28 @@ def lookup(
         # shard, replicated groups one namespace per unique KV head).
         exists_set = set()
         pos = 0
-        for g_idx, hash_bytes in candidate_meta:
+        previous_page = None
+        preceding_chunks_present = True
+        for g_idx, hash_bytes, end_token in candidate_meta:
+            spec_block_size = self._kv_cache_groups[g_idx].kv_cache_spec.block_size
+            page = (g_idx, (end_token - 1) // spec_block_size)
+            if page != previous_page:
+                preceding_chunks_present = True
+                previous_page = page
             count = len(self._lookup_key_prefixes[g_idx])
-            if all(res[pos + j] == 1 for j in range(count)):
+            present = all(res[pos + j] == 1 for j in range(count))
+            # A page-prefix hit covers every Store chunk before its final chunk.
+            if preceding_chunks_present and present:
                 exists_set.add((g_idx, hash_bytes))
+            if end_token % self.token_dbs[g_idx].chunk_size == 0:
+                preceding_chunks_present &= present
             pos += count
 
         cached_block_pool = ExternalCachedBlockPool(
             self.hash_block_size,
             exists_set,
         )
-        _, hit_length = self.coord.find_longest_cache_hit(
+        load_masks, hit_length = self.coord.find_longest_cache_hit(
             block_hashes,
             token_len,
             cached_block_pool,
@@ -2194,7 +2490,7 @@ def lookup(
             usable_length = self.coord.align_lookup_length(num_tokens - 1)
             if usable_length <= 0:
                 return MooncakeLookupResult(0)
-            _, hit_length = self.coord.find_longest_cache_hit(
+            load_masks, hit_length = self.coord.find_longest_cache_hit(
                 block_hashes,
                 usable_length,
                 cached_block_pool,
@@ -2205,6 +2501,7 @@ def lookup(
                 block_hashes,
                 hit_length,
                 cached_block_pool,
+                load_masks,
             ),
         )
 
@@ -2213,6 +2510,7 @@ def _tail_key_boundaries(
         block_hashes: Sequence[BlockHash],
         hit_length: int,
         cached_block_pool: ExternalCachedBlockPool,
+        load_masks: tuple[list[bool], ...],
     ) -> tuple[TailKeyBoundary, ...]:
         """Return the hash boundary used to store each group's tail block.
 
@@ -2227,10 +2525,9 @@ def _tail_key_boundaries(
         boundaries = []
         hit_boundary_hash_idx = hit_length // self.hash_block_size - 1
         for group_id, db in enumerate(self.token_dbs):
-            if not self._kv_cache_groups[group_id].kv_cache_spec.prefix_cacheable:
-                # Scratch groups are never stored, so they have no tail key.
+            if not load_masks[group_id] or not load_masks[group_id][-1]:
                 continue
-            chunk_id = cdiv(hit_length, db.block_size) - 1
+            chunk_id = cdiv(hit_length, db.chunk_size) - 1
             boundary_tokens = hit_length
             contains_hit_boundary = cached_block_pool.contains(
                 group_id, block_hashes[hit_boundary_hash_idx]
@@ -2239,7 +2536,7 @@ def _tail_key_boundaries(
                 assert contains_hit_boundary
             if not contains_hit_boundary:
                 next_chunk_hash_idx = min(
-                    (chunk_id + 1) * db.block_size // self.hash_block_size,
+                    (chunk_id + 1) * db.chunk_size // self.hash_block_size,
                     len(block_hashes),
                 )
                 for hash_idx in range(hit_boundary_hash_idx + 1, next_chunk_hash_idx):
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 9ae2597a9ebb..4bd75ace0dfe 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
@@ -108,11 +108,16 @@
 )
 from vllm.v1.core.sched.output import SchedulerOutput
 from vllm.v1.kv_cache_interface import (
+    CircularBufferSpec,
+    FullAttentionSpec,
     KVCacheConfig,
+    KVCacheSpec,
     MambaSpec,
     SlidingWindowMLASpec,
     SlidingWindowSpec,
+    get_kv_cache_spec_sliding_window,
     is_full_attention_spec,
+    iter_layer_specs,
 )
 from vllm.v1.metrics.cache_hit_source import CacheHitSource
 from vllm.v1.outputs import KVConnectorOutput
@@ -303,7 +308,7 @@ def __init__(
         )
         self.mode = get_moriio_mode(self.kv_transfer_config)
         self._transfers_mamba_state = any(
-            group.enable_kv_transfer and isinstance(group.kv_cache_spec, MambaSpec)
+            group.enable_kv_transfer and _is_mamba_group_spec(group.kv_cache_spec)
             for group in kv_cache_config.kv_cache_groups
         )
         if role == KVConnectorRole.SCHEDULER:
@@ -495,11 +500,19 @@ def has_connector_metadata(self) -> bool:
             return False
 
 
+def _is_mamba_group_spec(kv_cache_spec: KVCacheSpec) -> bool:
+    """True if every layer spec in the group is a MambaSpec."""
+    layer_specs = iter_layer_specs(kv_cache_spec)
+    return len(layer_specs) > 0 and all(
+        isinstance(spec, MambaSpec) for spec in layer_specs
+    )
+
+
 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.
+    mamba_group_indices) by whether each group's layers are MambaSpecs.
 
     Groups that opt out of external transfer are absent from
     ``KVCacheConfig.transfer_groups``. Returned indices address that filtered
@@ -508,7 +521,7 @@ def _split_kv_cache_group_kinds(
     attn: list[int] = []
     mamba: list[int] = []
     for gi, group in enumerate(kv_cache_config.transfer_groups):
-        if isinstance(group.kv_cache_spec, MambaSpec):
+        if _is_mamba_group_spec(group.kv_cache_spec):
             mamba.append(gi)
         else:
             attn.append(gi)
@@ -522,7 +535,7 @@ def _validate_hybrid_speculation(vllm_config: VllmConfig) -> None:
     # Other hybrid speculative methods need their own transfer validation.
     if not speculative_config.use_dspark():
         raise MoRIIOError(
-            "MoRIIO hybrid transfer supports DSpark speculative decoding only, got "
+            "MoRIIO hybrid WRITE supports DSpark speculative decoding only, got "
             f"method={speculative_config.method!r}"
         )
 
@@ -570,37 +583,41 @@ def __init__(
         self._has_mamba = bool(self._mamba_group_ids)
         mamba_spec = None
         if self._has_mamba:
+            # READ drafters must share the target's attention groups (e.g. MTP).
             if not self._attn_group_ids:
                 raise MoRIIOError("MoRIIO hybrid transfer requires an attention group")
-            if self.mode == MoRIIOMode.READ and len(self._attn_group_ids) != 1:
-                raise MoRIIOError(
-                    "MoRIIO hybrid READ requires exactly one transferable "
-                    "attention group, got "
-                    f"{len(self._attn_group_ids)}; a drafter that owns a "
-                    "separate attention group is not supported, give it the "
-                    "target's group or disable KV transfer for it"
-                )
             mamba_specs = [
-                cast(MambaSpec, kv_cache_config.transfer_groups[group_id].kv_cache_spec)
+                cast(MambaSpec, spec)
                 for group_id in self._mamba_group_ids
+                for spec in iter_layer_specs(
+                    kv_cache_config.transfer_groups[group_id].kv_cache_spec
+                )
             ]
             mamba_spec = _validate_mamba_specs(mamba_specs)
-            _validate_hybrid_speculation(vllm_config)
+            if self.mode == MoRIIOMode.WRITE:
+                _validate_hybrid_speculation(vllm_config)
         self._num_ssm_scratch_blocks = (
             mamba_spec.num_speculative_blocks if mamba_spec is not None else 0
         )
         self.block_size = vllm_config.cache_config.block_size
         attn_block_size = self.block_size
         max_decode_tail_tokens = 1 + vllm_config.num_lookahead_tokens
+        # Per attention group: not prefix-cacheable, so read even on a full hit.
+        self._attn_group_is_state: list[bool] = []
         if self._has_mamba:
-            attn_spec = kv_cache_config.transfer_groups[
-                self._attn_group_ids[0]
-            ].kv_cache_spec
-            attn_block_size = attn_spec.block_size
-            if attn_spec.dcp_sharded:
-                attn_block_size *= (
-                    vllm_config.parallel_config.decode_context_parallel_size
-                )
+            attn_specs = [
+                kv_cache_config.transfer_groups[group_id].kv_cache_spec
+                for group_id in self._attn_group_ids
+            ]
+            self._attn_group_is_state = [
+                not spec.prefix_cacheable for spec in attn_specs
+            ]
+            dcp_size = vllm_config.parallel_config.decode_context_parallel_size
+            # The smallest attention block bounds every group's local tail.
+            attn_block_size = min(
+                spec.block_size * (dcp_size if spec.dcp_sharded else 1)
+                for spec in attn_specs
+            )
             # Admission can pad the final prompt token to 1 + num_spec rows;
             # allocate_slots reserves drafter lookahead in addition to those.
             max_decode_tail_tokens += vllm_config.num_speculative_tokens
@@ -615,14 +632,22 @@ def __init__(
             )
         )
         if self._is_hma_required:
-            # MoRIIO handles sliding-window attention and mamba/KDA groups in
-            # READ mode. Other hybrid cache specs still fail closed.
+            # MoRIIO handles sliding-window and mamba/KDA groups in READ mode, and
+            # circular-buffer groups only alongside mamba/KDA in READ mode. Other
+            # hybrid cache specs still fail closed.
+            supported: tuple[type[KVCacheSpec], ...] = (
+                FullAttentionSpec,
+                SlidingWindowSpec,
+                MambaSpec,
+            )
+            if self._has_mamba and self.mode == MoRIIOMode.READ:
+                supported += (CircularBufferSpec,)
             unsupported = [
                 type(g.kv_cache_spec).__name__
                 for g in kv_cache_config.transfer_groups
-                if not (
-                    is_full_attention_spec(g.kv_cache_spec)
-                    or isinstance(g.kv_cache_spec, (SlidingWindowSpec, MambaSpec))
+                if not all(
+                    isinstance(spec, supported)
+                    for spec in iter_layer_specs(g.kv_cache_spec)
                 )
             ]
             if unsupported:
@@ -646,8 +671,8 @@ def __init__(
                 )
 
         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)
+            (sliding_window, g.kv_cache_spec.block_size)
+            if (sliding_window := get_kv_cache_spec_sliding_window(g.kv_cache_spec))
             else (0, self.block_size)
             for g in kv_cache_config.transfer_groups
         ]
@@ -1053,26 +1078,38 @@ def update_state_after_alloc(
                                 # 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 = (
+                                attn_block_groups, mamba_block_groups = (
                                     self.split_block_groups(blocks.get_block_ids())
                                 )
-                                if len(mamba_block_groups) != len(remote_mamba_groups):
+                                num_attn = len(attn_block_groups)
+                                if len(remote_block_ids) != num_attn + len(
+                                    mamba_block_groups
+                                ):
                                     raise MoRIIOError(
-                                        "MoRIIO hybrid READ Mamba group "
-                                        "count mismatch: "
-                                        f"local={len(mamba_block_groups)}, "
-                                        f"remote={len(remote_mamba_groups)}"
+                                        "MoRIIO hybrid READ block group "
+                                        "count mismatch: local="
+                                        f"{num_attn + len(mamba_block_groups)}, "
+                                        f"remote={len(remote_block_ids)}"
                                     )
-                                local_attn = attn_block_ids
-                                local_attn, remote_attn = self._align_read_blocks(
-                                    local_attn,
-                                    remote_attn,
-                                    self._max_decode_tail_blocks,
-                                )
+                                remote_mamba_groups = [
+                                    list(group) for group in remote_block_ids[num_attn:]
+                                ]
+                                aligned_local_attn: list[list[int]] = []
+                                aligned_remote_attn: list[list[int]] = []
+                                for local_group, remote_group in zip(
+                                    attn_block_groups,
+                                    remote_block_ids[:num_attn],
+                                    strict=True,
+                                ):
+                                    aligned_local, aligned_remote = (
+                                        self._align_read_blocks(
+                                            local_group,
+                                            list(remote_group),
+                                            self._max_decode_tail_blocks,
+                                        )
+                                    )
+                                    aligned_local_attn.append(aligned_local)
+                                    aligned_remote_attn.append(aligned_remote)
                                 aligned_local_mamba: list[list[int]] = []
                                 aligned_remote_mamba: list[list[int]] = []
                                 # Scratch slots were removed by
@@ -1092,9 +1129,12 @@ def update_state_after_alloc(
                                     )
                                     aligned_local_mamba.append(aligned_local)
                                     aligned_remote_mamba.append(aligned_remote)
-                                local_block_ids = [local_attn, *aligned_local_mamba]
+                                local_block_ids = [
+                                    *aligned_local_attn,
+                                    *aligned_local_mamba,
+                                ]
                                 adjusted_remote_block_ids = [
-                                    remote_attn,
+                                    *aligned_remote_attn,
                                     *aligned_remote_mamba,
                                 ]
                             else:
@@ -1125,10 +1165,18 @@ def update_state_after_alloc(
                             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()
+                                attn_block_groups, mamba_block_groups = (
+                                    self.split_block_groups(blocks.get_block_ids())
                                 )
-                                local_block_ids = [[], *mamba_block_groups]
+                                local_block_ids = [
+                                    group if is_state else []
+                                    for group, is_state in zip(
+                                        attn_block_groups,
+                                        self._attn_group_is_state,
+                                        strict=True,
+                                    )
+                                ]
+                                local_block_ids.extend(mamba_block_groups)
                             else:
                                 # Nothing to pull, but retain one empty list per
                                 # cache group so the worker can notify P.
@@ -1416,11 +1464,12 @@ def build_connector_meta(
         ):
             # Hybrid models zero recycled attention pages that held Mamba state.
             # Host-submitted READs overwrite these pages and can race zeroing.
-            # Hybrid READ metadata puts the aligned attention pages first.
+            # Hybrid READ metadata puts the aligned attention groups first.
             read_dst_block_ids = {
                 b
                 for _, block_ids in self._reqs_need_recv.values()
-                for b in block_ids[0]
+                for group in block_ids[: len(self._attn_group_ids)]
+                for b in group
             }
             scheduler_output.new_block_ids_to_zero = [
                 b
@@ -1470,19 +1519,20 @@ def shutdown(self):
 
     def split_block_groups(
         self, block_ids: list[list[int]] | tuple[list[int], ...]
-    ) -> tuple[list[int], list[list[int]]]:
+    ) -> tuple[list[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. Only each group's running-state
-        slot is transferred.
+        The wire payload stores every attention group first, followed by every
+        Mamba group, each in transfer-group order. Sliding-window groups keep
+        only their window; Mamba groups only their running-state slot.
         """
         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]])
+        # Also window-clips sliding-window attention groups.
+        transfer_block_ids = self.get_exchange_clipped_blocks(block_ids)
+        attn = [list(transfer_block_ids[group_id]) for group_id in self._attn_group_ids]
         mamba_groups = [
             clip_ssm_state_blocks(
                 list(transfer_block_ids[group_id]),
@@ -1531,7 +1581,7 @@ def request_finished(
         should be freed now or will be sent asynchronously and freed later.
 
         ``block_ids`` normally contains all cache groups. Hybrid READ callers from
-        ``request_finished_all_groups`` pass a flat attention list and the
+        ``request_finished_all_groups`` pass the attention groups and the
         recurrent-state groups separately in ``mamba_block_groups``.
         """
         request_id = request.request_id
@@ -1571,7 +1621,7 @@ def request_finished(
                 # 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 cast(BlockIds, block_ids)],
                         *[list(group) for group in mamba_block_groups],
                     ]
                 elif block_ids and isinstance(block_ids[0], int):
@@ -1596,7 +1646,7 @@ def request_finished(
         # On producer
         if mamba_block_groups is not None:
             computed_block_ids: BlockIds = [
-                list(cast(list[int], block_ids)),
+                *[list(group) for group in cast(BlockIds, block_ids)],
                 *[list(group) for group in mamba_block_groups],
             ]
         elif self._has_mamba and self.mode == MoRIIOMode.WRITE:
@@ -1766,11 +1816,17 @@ def __init__(
         # 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)
+        attn_transfer_group_ids, mamba_transfer_group_ids = _split_kv_cache_group_kinds(
+            kv_cache_config
+        )
+        self._num_attn_transfer_groups = len(attn_transfer_group_ids)
         self._num_mamba_transfer_groups = len(mamba_transfer_group_ids)
-        self._mamba_payload_index_by_layer = {
+        # Hybrid READ payload order: attention groups, then Mamba groups.
+        self._hybrid_payload_index_by_layer = {
             layer_name: group_id if self.mode == MoRIIOMode.WRITE else payload_index
-            for payload_index, group_id in enumerate(mamba_transfer_group_ids, start=1)
+            for payload_index, group_id in enumerate(
+                [*attn_transfer_group_ids, *mamba_transfer_group_ids]
+            )
             for layer_name in kv_cache_config.transfer_groups[group_id].layer_names
         }
 
@@ -3789,7 +3845,7 @@ def _mamba_blocks_for_layer(
         layer_name: str,
         block_groups: BlockIds,
     ) -> list[int]:
-        payload_index = self._mamba_payload_index_by_layer.get(layer_name)
+        payload_index = self._hybrid_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"
@@ -3990,7 +4046,8 @@ def _read_blocks(
             return
 
         if self._has_mamba:
-            expected_groups = 1 + self._num_mamba_transfer_groups
+            num_attn = self._num_attn_transfer_groups
+            expected_groups = num_attn + self._num_mamba_transfer_groups
             if (
                 len(local_block_ids) != expected_groups
                 or len(remote_block_ids) != expected_groups
@@ -4000,8 +4057,8 @@ def _read_blocks(
                     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])
+            local_attn = [b for group in local_block_ids[:num_attn] for b in group]
+            remote_attn = [b for group in remote_block_ids[:num_attn] for b in group]
             self._mamba_tp_ratio(remote_tp_size)
         else:
             local_attn = remote_attn = []
@@ -4055,13 +4112,13 @@ def _read_blocks(
                     _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]
+                group_idx = (
+                    self._hybrid_payload_index_by_layer
+                    if self._has_mamba
+                    else 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_read_attention_offsets(
                     layer_name,
                     local_layer_blocks,
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 ffd75e2d3e18..4149e28bb21c 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
@@ -12,8 +12,9 @@
 import time
 import uuid
 from collections import defaultdict
-from collections.abc import Iterator
+from collections.abc import Iterator, Sequence
 from concurrent.futures import Future, ThreadPoolExecutor
+from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, cast
 
 import msgspec
@@ -27,6 +28,7 @@
     EngineId,
     EngineTransferInfo,
     TransferTopology,
+    check_compressed_kernel_split,
     get_current_attn_backends,
     get_current_attn_backends_and_specs,
     kv_postprocess_blksize_and_layout_on_receive,
@@ -125,6 +127,34 @@ def _share_storage_and_block_stride(caches: list[torch.Tensor]) -> bool:
     return len(block_strides) == len(storage_ptrs) == 1
 
 
+@dataclass
+class _RegionPairs:
+    """How a peer's KV regions transfer with the local ones.
+
+    Pairs come from the layers both sides hold (``region_members``): pair ``k``
+    moves the peer's region ``remote[k]`` and local region ``local[k]`` at
+    cache group ``group[k]``'s blocks. A region shared by several cache groups
+    pairs once per group it holds a layer of, so which layers share a region
+    may differ between the two sides. ``remote`` indexes ``remote_regions``,
+    the peer's regions that pair (in its order), which are the ones its
+    descriptor list holds.
+    """
+
+    remote: list[int]
+    local: list[int]
+    group: list[int]
+    remote_regions: list[int]
+    # A local region paired with each of ``remote_regions``; they all have the
+    # same block length.
+    partner: list[int]
+    # Whether every paired layer is in the same cache group on both sides.
+    groups_match: bool = True
+    # Hybrid SSM models: the peer's regions holding SSM state, in the order its
+    # SSM descriptors are built, and its PLE region (set by _pair_ssm_regions).
+    remote_ssm_regions: list[int] | None = None
+    remote_ple_region: int | None = None
+
+
 class NixlBaseConnectorWorker:
     """Base implementation of Worker side methods shared by pull and push."""
 
@@ -133,8 +163,7 @@ class NixlBaseConnectorWorker:
     # Overridden by NixlPushConnectorWorker.
     _TRANSFER_MODE: str = "pull"
 
-    # Layer-name routing is supported only by NixlPushConnector for HMA or
-    # packed KV caches under PP.
+    # Only NixlPushConnector supports PP with HMA or packed KV caches.
     _supports_pp_hma = False
 
     def _compute_desc_ids(
@@ -143,11 +172,19 @@ def _compute_desc_ids(
         dst_num_blocks: int,
         block_size_ratio: float | None,
         physical_blocks_per_logical: int,
+        pairs: tuple[Sequence[int], Sequence[int]],
         region_num_blocks: list[int] | None = None,
-        region_group_ids: list[int] | None = None,
-        uses_region_group_mapping: bool | None = None,
+        ssm_regions: Sequence[int] | None = None,
     ) -> np.ndarray:
-        """Compute NIXL descriptor IDs for given block IDs."""
+        """Compute NIXL descriptor IDs for given block IDs.
+
+        ``pairs`` is ``(regions, keys)``: region ``regions[k]`` is read at
+        ``block_ids[keys[k]]``, a cache group's blocks (a region's, for
+        region-mapped transfers).
+
+        ``ssm_regions`` lists this side's SSM regions in the order its SSM
+        descriptors are built; it defaults to the local ones.
+        """
         num_ssm_regions = 0
         if self._has_mamba:
             assert self._conv_decomp is not None
@@ -164,67 +201,28 @@ 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] * num_desc_regions
+            region_num_blocks = [num_blocks] * (
+                self.num_regions or len(self.block_len_per_layer)
+            )
         elif block_size_ratio is not None:
             region_num_blocks = [
                 int(count * block_size_ratio) for count in region_num_blocks
             ]
-        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)
 
-        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
-            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)
-            # NOTE (NickLucche) With HMA, every kv group has the same number of layers
-            # and layers from different groups share the same kv tensor.
-            # eg block_ids=[[1, 2], [3]]->blocks [1, 2] need to be
-            # read across all regions, same for [3], but group0-group1 blocks will
-            # always differ (different areas). Therefore we can just flatten the
-            # block_ids and compute the descs ids for all groups at once.
-            if not uses_region_group_mapping:
-                block_arr = np.concatenate(
-                    [np.asarray(group, dtype=np.int32) for group in block_ids]
-                )[None, :]
-                return (region_offsets[:, None] + block_arr).ravel()
+            regions, keys = pairs
+            offsets = region_offsets[np.asarray(regions, dtype=np.int64)]
+            keys_array = np.asarray(keys, dtype=np.int64)
             desc_ids = []
-            for group_id, group in enumerate(block_ids):
-                if not group:
-                    continue
-                region_ids = np.flatnonzero(
-                    (region_group_ids_array == group_id)
-                    | (region_group_ids_array == _SHARED_REGION_GROUP_ID)
-                )[:, None]
-                assert region_ids.size > 0
-                group_blocks = np.asarray(group, dtype=np.int32)[None, :]
-                desc_ids.append((region_offsets[region_ids] + group_blocks).ravel())
+            for key in np.unique(keys_array).tolist():
+                blocks = np.asarray(block_ids[key], dtype=np.int64)
+                if blocks.size:
+                    key_offsets = offsets[keys_array == key]
+                    desc_ids.append((key_offsets[:, None] + blocks[None, :]).ravel())
             if not desc_ids:
                 return np.array([], dtype=np.int64)
             return np.concatenate(desc_ids)
@@ -235,63 +233,54 @@ def _compute_desc_ids(
         # logical_blocks entries per region (no kernel splitting, and never
         # ratio-expanded since state blocks are indivisible).
         logical_blocks = dst_num_blocks // physical_blocks_per_logical
+        if ssm_regions is None:
+            ssm_regions = self._ssm_region_indices or range(
+                len(self.block_len_per_layer)
+            )
+        ssm_slot = {region: slot for slot, region in enumerate(ssm_regions)}
+        # NOTE (NickLucche) SSM and Attention block regions can be exchanged
+        # arbitrarily by manager. Therefore, descs are laid out as:
+        #   [descs_fa (all regions) | descs_ssm (all SSM regions) | desc_ple].
+        # num_fa_descs offset must be computed per-engine since P and D can
+        # have different num_blocks (and thus different FA desc counts).
+        num_ssm_descs = len(ssm_regions) * ssm_regions_per_layer
+        pair_regions = np.asarray(pairs[0], dtype=np.int64)
+        pair_keys = np.asarray(pairs[1], dtype=np.int64)
+        sub_descs = np.arange(ssm_regions_per_layer, dtype=np.int64)
         all_descs: list[np.ndarray] = []
         for i, group in enumerate(block_ids):
-            group_arr = np.asarray(group, dtype=np.int32)
+            blocks = np.asarray(group, dtype=np.int64)
+            group_regions = pair_regions[pair_keys == i]
+            if not blocks.size or not group_regions.size:
+                continue
             spec_type = self._group_spec_types[i]
-            # Hybrid regions are keyed by base address, so some hold only
-            # another group's layers (e.g. KDA-only or scratch-only tensors).
-            # Reading them for this group would only re-read aliased bytes.
-            group_regions = (
-                self._group_region_indices[i] if self._group_region_indices else None
-            )
             if _is_attention_spec(spec_type):
-                fa_region_ids = (
-                    np.arange(self.num_regions, dtype=np.int32)
-                    if group_regions is None
-                    else np.asarray(group_regions, dtype=np.int32)
-                )[:, None]
-                all_descs.append(
-                    (fa_region_ids * num_blocks + group_arr[None, :]).ravel()
-                )
+                desc_ids = region_offsets[group_regions][:, None] + blocks[None, :]
             elif _is_ssm_spec(spec_type):
-                # NOTE (NickLucche) SSM and Attention block regions can
-                # be exchanged arbitrarily by manager.  Therefore, descs
-                # are laid out as:
-                #   [descs_fa (all regions) | descs_ssm (all regions)].
-                # num_fa_descs offset must be computed per-engine since
-                # P and D can have different num_blocks (and thus
-                # different FA desc counts).
                 if i == self._ple_group_index:
-                    ssm_region_ids = np.arange(
-                        num_ssm_regions, num_ssm_regions + 1, dtype=np.int32
-                    )
-                elif group_regions is None:
-                    ssm_region_ids = np.arange(num_ssm_regions, dtype=np.int32)
+                    ssm_desc_ids = np.array([num_ssm_descs], dtype=np.int64)
                 else:
-                    # Each SSM layer owns ssm_regions_per_layer consecutive
-                    # descs, in _ssm_region_indices order.
-                    ssm_slots = np.flatnonzero(
-                        np.isin(self._ssm_region_indices, group_regions)
-                    ).astype(np.int32)
-                    ssm_region_ids = (
-                        ssm_slots[:, None] * ssm_regions_per_layer
-                        + np.arange(ssm_regions_per_layer, dtype=np.int32)
-                    ).ravel()
-                ssm_region_ids = ssm_region_ids[:, None]
-                all_descs.append(
-                    (
-                        ssm_region_ids * logical_blocks
-                        + group_arr[None, :]
-                        + num_fa_descs
+                    # Each SSM region owns ssm_regions_per_layer consecutive
+                    # descs, in ssm_regions order.
+                    slots = np.array([ssm_slot[r] for r in group_regions.tolist()])
+                    ssm_desc_ids = (
+                        slots[:, None] * ssm_regions_per_layer + sub_descs
                     ).ravel()
+                desc_ids = (
+                    ssm_desc_ids[:, None] * logical_blocks
+                    + blocks[None, :]
+                    + num_fa_descs
                 )
             else:
                 raise ValueError(
                     f"Unknown spec type {self._group_spec_types[i]} at index {i}"
                 )
+            all_descs.append(desc_ids.ravel())
 
-        return np.concatenate(all_descs)
+        if not all_descs:
+            return np.array([], dtype=np.int32)
+        # NIXL takes int32 ids without a conversion copy.
+        return np.concatenate(all_descs).astype(np.int32)
 
     def _build_local_splits_from_plan(
         self,
@@ -382,7 +371,7 @@ def _fa_desc_replicated(
         n_regions = len(self.block_len_per_layer)
         if n_regions == 0 or self.num_regions == 0:
             return np.zeros(len(desc_starts), dtype=bool)
-        region_indices = self._transfer_layer_region_indices or range(n_regions)
+        region_indices = range(n_regions)
         region_ends = np.cumsum(
             [self.region_num_blocks[i] * block_size_ratio for i in region_indices]
         )
@@ -441,138 +430,22 @@ def _get_head_sharded_draft_regions(self) -> set[int]:
         return regions
 
     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 _tracks_region_layers(self) -> bool:
-        """Whether regions advertise their member layer names.
-
-        Push workers do so for HMA and packed layouts, so a PP producer can
-        match layers by name instead of by region index.
-        """
-        return (
-            self._supports_pp_hma
-            and (self._is_hma_required or self._has_packed_cache)
-            and not self._has_mamba
-        )
-
-    def _requires_layer_name_routing(self) -> bool:
-        """Whether PP push must match HMA/packed layers by name, not region index."""
-        return self.pp_size > 1 and self._tracks_region_layers()
-
-    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):
+        regions_by_layer: dict[str, list[int]] = defaultdict(list)
+        for region, layer_names in enumerate(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}"
-
-        packed_regions = {
-            remote_region_by_layer[name]
-            for name in nixl_agent_meta.packed_member_layouts
-        }
-        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)
-        if nixl_agent_meta.packed_member_layouts:
-            for i, layer_name in enumerate(self._transfer_layer_names):
-                if remote_regions[i] not in packed_regions:
-                    continue
-                assert layer_name in nixl_agent_meta.packed_member_layouts, (
-                    f"Remote packed layer {layer_name!r} has no layout"
-                )
-                offset, page_size = nixl_agent_meta.packed_member_layouts[layer_name]
-                assert (
-                    0 <= offset < offset + page_size <= nixl_agent_meta.block_lens[i]
-                ), f"Remote packed layer {layer_name!r} escapes its block"
-                local_region = self._transfer_layer_region_indices[i]
-                assert page_size == self.block_len_per_layer[local_region], (
-                    f"Packed MLA page sizes must match for layer {layer_name!r}"
-                )
-                nixl_agent_meta.kv_caches_base_addr[i] += offset
-                nixl_agent_meta.block_lens[i] = page_size
-            nixl_agent_meta.packed_member_layouts = {}
-        # 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]
+                regions_by_layer[layer_name].append(region)
+        self._regions_by_layer: dict[str, list[int]] = dict(regions_by_layer)
 
-    def _split_remote_packed_rows(self, nixl_agent_meta: NixlAgentMetadata) -> None:
-        """Map the remote's whole-row region onto our per-layer regions.
+    @staticmethod
+    def _expand_remote_packed_rows(nixl_agent_meta: NixlAgentMetadata) -> None:
+        """View each of the remote's packed rows as one region per layer.
 
-        E.g. a BLHNC prefiller registers one region whose block is a full row
-        holding every layer's page, while a HiSparse decoder registers one
-        region per layer. Using the advertised `packed_member_layouts`, the
-        remote row is rewritten in place as one virtual region per local layer:
+        A block-outermost peer may register one region whose block is a full
+        row holding every layer's page, while this worker registers one region
+        per layer (a PP push stage, or a HiSparse decoder on LB*). Using the
+        advertised ``packed_member_layouts``, each row is rewritten in place as
+        one virtual region per layer, so that layers pair one to one:
 
             remote row b:  [ L0 page | L1 page | ... ]   base=B, stride=S
                              ^off0     ^off1
@@ -581,40 +454,209 @@ def _split_remote_packed_rows(self, nixl_agent_meta: NixlAgentMetadata) -> None:
             remote metadata, after:   [(B + off0, page0, S),    # L0
                                        (B + off1, page1, S)]    # L1
 
-        Remote regions outside the row (e.g. a separate indexer) pass through.
+        Remote regions outside the rows (e.g. a separate indexer) pass through.
         """
         layouts = nixl_agent_meta.packed_member_layouts
-        remote_names = nixl_agent_meta.region_names
-        assert remote_names is not None, "Remote advertised no region_names"
-        (row,) = [i for i, name in enumerate(remote_names) if name in layouts]
-        if nixl_agent_meta.region_group_ids is not None and (
-            nixl_agent_meta.region_group_ids[row] == _SHARED_REGION_GROUP_ID
-        ):
-            raise NotImplementedError(
-                "Reading multi-group packed KV rows into per-layer regions is not "
-                "supported; use the same KV cache layout on P and D."
+        regions: list[tuple[int, int, int, list[str]]] = []
+        packed_rows: set[int] = set()
+        for region, layer_names in enumerate(nixl_agent_meta.region_members):
+            packed = [name for name in layer_names if name in layouts]
+            if not packed:
+                block_len = nixl_agent_meta.block_lens[region]
+                regions.append((region, 0, block_len, layer_names))
+                continue
+            assert len(packed) == len(layer_names), (
+                f"Remote packed row {region} holds layers without a layout"
             )
-        remote_index = {name: i for i, name in enumerate(remote_names)}
-        regions = [
-            (row, *layouts[name])
-            if name in layouts
-            else (remote_index[name], 0, nixl_agent_meta.block_lens[remote_index[name]])
-            for name in self.region_names
-        ]
+            packed_rows.add(region)
+            for name in packed:
+                offset, page_size = layouts[name]
+                assert (
+                    0
+                    <= offset
+                    < offset + page_size
+                    <= nixl_agent_meta.block_lens[region]
+                ), f"Remote packed layer {name!r} escapes its block"
+                regions.append((region, offset, page_size, [name]))
         nixl_agent_meta.kv_caches_base_addr = [
-            nixl_agent_meta.kv_caches_base_addr[i] + offset for i, offset, _ in regions
+            nixl_agent_meta.kv_caches_base_addr[region] + offset
+            for region, offset, _, _ in regions
         ]
-        nixl_agent_meta.block_lens = [page_size for _, _, page_size in regions]
+        nixl_agent_meta.block_lens = [page_size for _, _, page_size, _ in regions]
         nixl_agent_meta.block_strides = [
-            nixl_agent_meta.block_strides[i] for i, _, _ in regions
+            nixl_agent_meta.block_strides[region] for region, _, _, _ in regions
         ]
         for attr in ("region_num_blocks", "region_group_ids", "region_mem_types"):
             values = getattr(nixl_agent_meta, attr)
             if values is not None:
-                setattr(nixl_agent_meta, attr, [values[i] for i, _, _ in regions])
-        nixl_agent_meta.region_names = list(self.region_names)
+                setattr(
+                    nixl_agent_meta, attr, [values[region] for region, *_ in regions]
+                )
+        if nixl_agent_meta.region_names is not None:
+            # A row's virtual regions take their layer's name.
+            names = nixl_agent_meta.region_names
+            nixl_agent_meta.region_names = [
+                layer_names[0] if region in packed_rows else names[region]
+                for region, _, _, layer_names in regions
+            ]
+        nixl_agent_meta.region_members = [layer_names for *_, layer_names in regions]
         nixl_agent_meta.packed_member_layouts = {}
 
+    def _pair_regions(self, nixl_agent_meta: NixlAgentMetadata) -> _RegionPairs:
+        """Pair a peer's regions with the local ones through the layers both
+        hold.
+
+        A layer's regions pair in order on both sides (two when K and V
+        register apart). The peer must hold every local layer; layers only
+        the peer holds are skipped on a PP stage, which holds a slice of them.
+        A peer's packed rows are first rewritten as one region per layer when
+        this worker registers its layers apart.
+        """
+        engine_id = nixl_agent_meta.engine_id
+        num_regions = len(nixl_agent_meta.kv_caches_base_addr)
+        if any(
+            values is not None and len(values) != num_regions
+            for values in (
+                nixl_agent_meta.block_lens,
+                nixl_agent_meta.block_strides,
+                nixl_agent_meta.region_num_blocks,
+                nixl_agent_meta.region_group_ids,
+                nixl_agent_meta.region_names,
+                nixl_agent_meta.region_mem_types,
+            )
+        ):
+            raise RuntimeError(
+                f"Engine {engine_id} advertised KV region metadata of different "
+                "lengths."
+            )
+        if len(nixl_agent_meta.region_members) != num_regions:
+            raise RuntimeError(
+                f"Engine {engine_id} did not advertise the layers of its KV regions."
+            )
+        if nixl_agent_meta.packed_member_layouts and not self._packed_member_layouts:
+            self._expand_remote_packed_rows(nixl_agent_meta)
+        remote_members = nixl_agent_meta.region_members
+        remote_groups = nixl_agent_meta.region_group_ids
+        remote_layer_groups = nixl_agent_meta.layer_group_ids
+        group_by_layer = self.kv_cache_config.transfer_group_index_by_layer
+        remote_regions_by_layer: dict[str, list[int]] = defaultdict(list)
+        for region, layer_names in enumerate(remote_members):
+            for layer_name in layer_names:
+                remote_regions_by_layer[layer_name].append(region)
+
+        triples: dict[tuple[int, int, int], None] = {}
+        layers: set[str] = set()
+        groups_match = True
+        for layer_name, remote_regions in remote_regions_by_layer.items():
+            local_regions = self._regions_by_layer.get(layer_name)
+            if local_regions is None:
+                if self.pp_size > 1:
+                    # This PP stage holds a slice of the peer's layers.
+                    continue
+                raise RuntimeError(
+                    f"Engine {engine_id} holds KV layer {layer_name!r}, which "
+                    "has no counterpart here."
+                )
+            if len(local_regions) != len(remote_regions):
+                raise NotImplementedError(
+                    f"KV layer {layer_name!r} spans {len(remote_regions)} NIXL "
+                    f"regions on engine {engine_id} and {len(local_regions)} here."
+                )
+            group = group_by_layer[layer_name]
+            layers.add(layer_name)
+            # Block ids are listed per group, so the group numbers must match;
+            # a region shared by several groups (-1) cannot show that.
+            if remote_layer_groups and remote_layer_groups.get(layer_name) != group:
+                groups_match = False
+            for remote, local in zip(remote_regions, local_regions):
+                remote_group = remote_groups[remote] if remote_groups else group
+                if remote_group not in (group, _SHARED_REGION_GROUP_ID):
+                    groups_match = False
+                triples[(remote, local, group)] = None
+        missing = self._regions_by_layer.keys() - layers
+        if missing:
+            raise RuntimeError(
+                f"Engine {engine_id} does not hold the KV layers "
+                f"{sorted(missing, key=_region_sort_key)}."
+            )
+
+        paired_regions = sorted({remote for remote, _, _ in triples})
+        position = {region: i for i, region in enumerate(paired_regions)}
+        partner = [-1] * len(paired_regions)
+        for remote, local, _ in triples:
+            first = partner[position[remote]]
+            if first < 0:
+                partner[position[remote]] = local
+            elif self.block_len_per_layer[first] != self.block_len_per_layer[
+                local
+            ] or self._is_region_replicated(first) != self._is_region_replicated(local):
+                raise NotImplementedError(
+                    f"KV region {remote} of engine {engine_id} pairs with local "
+                    "regions of different block lengths."
+                )
+        pairs = _RegionPairs(
+            remote=[position[remote] for remote, _, _ in triples],
+            local=[local for _, local, _ in triples],
+            group=[group for _, _, group in triples],
+            remote_regions=paired_regions,
+            partner=partner,
+            groups_match=groups_match,
+        )
+        if self._has_mamba:
+            self._pair_ssm_regions(pairs, engine_id)
+        return pairs
+
+    def _pair_ssm_regions(self, pairs: _RegionPairs, engine_id: EngineId) -> None:
+        """Find the peer's SSM and PLE regions through the pairs.
+
+        SSM descriptors are built per region holding SSM state. The peer's SSM
+        regions are listed in the order of the local SSM regions they pair
+        with, so identical layouts keep identical descriptor lists.
+        """
+        if not pairs.groups_match:
+            # Region-mapped reads, the only fallback, support pure MLA alone.
+            raise NotImplementedError(
+                f"Engine {engine_id} puts Mamba/SSM hybrid KV layers in other "
+                "cache groups than this worker; use the same KV cache plan on P "
+                "and D."
+            )
+        local_ssm = self._ssm_region_indices or range(len(self.block_len_per_layer))
+        ssm_slot = {region: slot for slot, region in enumerate(local_ssm)}
+        ssm_pairs: list[tuple[int, int]] = []
+        ple_regions: set[int] = set()
+        for remote, local, group in zip(pairs.remote, pairs.local, pairs.group):
+            if group == self._ple_group_index:
+                ple_regions.add(remote)
+            elif _is_ssm_spec(self._group_spec_types[group]):
+                ssm_pairs.append((ssm_slot[local], remote))
+        if len(ple_regions) > 1:
+            raise NotImplementedError(
+                f"Engine {engine_id} holds its PLE state in several KV regions."
+            )
+        pairs.remote_ssm_regions = list(
+            dict.fromkeys(remote for _, remote in sorted(ssm_pairs))
+        )
+        pairs.remote_ple_region = next(iter(ple_regions), None)
+
+    @staticmethod
+    def _select_remote_regions(
+        nixl_agent_meta: NixlAgentMetadata, remote_regions: Sequence[int]
+    ) -> None:
+        """Keep the remote regions ``remote_regions``, in that order."""
+        for attr in (
+            "kv_caches_base_addr",
+            "block_lens",
+            "block_strides",
+            "region_num_blocks",
+            "region_group_ids",
+            "region_names",
+            "region_mem_types",
+            "region_members",
+        ):
+            values = getattr(nixl_agent_meta, attr)
+            if values:
+                setattr(nixl_agent_meta, attr, [values[i] for i in remote_regions])
+
     def __init__(
         self,
         vllm_config: "VllmConfig",
@@ -756,8 +798,15 @@ def __init__(
                 else nixl_agent_config(num_threads=num_threads, capture_telemetry=True)
             )
 
+        # Extra UCX backend init params, e.g. {"ucx_error_handling_mode": "none"}.
+        params = self.kv_transfer_config.get_from_extra_config("ucx_backend_params", {})
+        ucx_backend_params = {k: str(v) for k, v in params.items()}
         sleep_mode = vllm_config.model_config.enable_sleep_mode and config is not None
-        if sleep_mode:
+        # With UCX params, create every backend here, in the configured order.
+        explicit_backends = (config is not None and not sleep_mode) and (
+            bool(ucx_backend_params) and "UCX" in self.nixl_backends
+        )
+        if sleep_mode or explicit_backends:
             config = nixl_agent_config(backends=[], capture_telemetry=True)
         self.nixl_wrapper = nixl_wrapper_cls(str(uuid.uuid4()), config)
         if sleep_mode:
@@ -767,8 +816,17 @@ def __init__(
                 {
                     "num_threads": str(num_threads),
                     "engine_config": "RCACHE_ENABLE=n,TLS=^cuda_ipc",
+                    **ucx_backend_params,
                 },
             )
+        elif explicit_backends:
+            for backend in self.nixl_backends:
+                nixl_params: dict[str, str] = {}
+                if backend == "UCX":
+                    if not non_ucx_backends:
+                        nixl_params["num_threads"] = str(num_threads)
+                    nixl_params.update(ucx_backend_params)
+                self.nixl_wrapper.create_backend(backend, nixl_params)
         # Map of engine_id -> {(pp_rank, tp_rank): agent_name, ...}.
         # non-PP remote uses pp_rank 0, i.e. (0, tp_rank).
         self._remote_agents: dict[EngineId, dict[tuple[int, int], str]] = defaultdict(
@@ -869,7 +927,6 @@ def __init__(
         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._mixed_mem_types = False
@@ -881,9 +938,8 @@ def __init__(
         self._dram_src_handles_by_block_size: dict[int, int | None] = {}
 
         # 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.
+        # transfers into the PP=1 remote's regions that hold its layers.
         self.pp_size = vllm_config.parallel_config.pipeline_parallel_size
-        self._remote_region_offset = 0
         if self.pp_size > 1:
             if self._is_hma_required and not self._supports_pp_hma:
                 raise NotImplementedError(
@@ -891,8 +947,7 @@ def __init__(
                     "> 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.
+            # Mamba/SSM hybrid PP stages are not validated yet.
             if self._has_mamba:
                 raise NotImplementedError(
                     "NixlPushConnector does not support pipeline_parallel_size > 1 "
@@ -924,7 +979,8 @@ def __init__(
         self.dst_num_blocks: dict[EngineId, int] = {}
         self.dst_region_num_blocks: dict[EngineId, list[int]] = {}
         self.dst_region_group_ids: dict[EngineId, list[int]] = {}
-        self.dst_uses_region_group_mapping: dict[EngineId, bool] = {}
+        # Engine -> how its regions pair with ours.
+        self._dst_pairs: dict[EngineId, _RegionPairs] = {}
         self.dst_region_mem_types: dict[EngineId, list[str]] = {}
         self._registered_descs: list[Any] = []
         self._kv_unregistered = False
@@ -1037,14 +1093,10 @@ def __init__(
         # Block-outermost layouts (BLHNC/BHLNC) pack every layer's page into one
         # block row of a shared allocation.
         self._has_packed_cache = KVCacheLayout[self.kv_cache_layout].is_block_outermost
-        # Region layer membership is populated for HMA and packed layouts.
+        # The layers of each region, which pair regions with a peer's.
         self.region_members = []
+        self._regions_by_layer = {}
         self._packed_member_layouts: dict[str, tuple[int, int]] = {}
-        # 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] = {}
         # GQA draft under an MLA target: its SPLIT regions use a draft TP mapping.
@@ -1081,6 +1133,9 @@ def _sync_block_size_with_kernel(self) -> None:
             self.vllm_config, self.kv_cache_config, self.attn_backends
         )
         kernel_block_size = select_common_block_size(self.block_size, backends, specs)
+        check_compressed_kernel_split(
+            self.vllm_config, self.kv_cache_config, kernel_block_size
+        )
         # Number of blocks not accounting for kernel block mismatches
         self._logical_num_blocks = self.num_blocks
         if self.block_size != kernel_block_size:
@@ -1496,8 +1551,14 @@ def request_ready(f: Future[Any], entry=(req_id, meta)):
 
     def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
         """Register the KV Cache data in nixl."""
-        use_layer_name_routing = self._requires_layer_name_routing()
-        route_packed_layers = self._has_packed_cache and use_layer_name_routing
+        # A PP push stage's packed rows hold other layers than the consumer's,
+        # so it registers each layer's page as a region of its own.
+        split_packed_rows = (
+            self._has_packed_cache
+            and self._supports_pp_hma
+            and self.pp_size > 1
+            and not self._has_mamba
+        )
         self.transfer_topo = TransferTopology(
             tp_rank=self.transfer_tp_rank,
             tp_size=self.transfer_tp_size,
@@ -1557,6 +1618,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
         registration_ranges: dict[tuple[int, str], tuple[int, int, int]] = {}
         region_mem_types: list[str] = []
         seen_base_addresses: list[int] = []
+        seen_region_keys: list[tuple[int, int]] = []
         self._ssm_region_indices = []
         region_groups: list[set[int]] = []
         self._ple_block_len = None
@@ -1570,13 +1632,12 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
             and self.use_mla
             and not self._has_mamba
         )
-        if route_packed_layers and any(
+        if split_packed_rows and any(
             not isinstance(spec, (MLAAttentionSpec, SlidingWindowMLASpec))
             for spec in self._layer_specs.values()
         ):
             raise NotImplementedError("PP push with packed KV caches requires MLA")
 
-        track_region_layers = self._tracks_region_layers()
         region_layers: list[list[str]] = []
         packed_member_layouts: dict[str, tuple[int, int]] = {}
 
@@ -1663,9 +1724,9 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
                 physical_ratio = self._physical_blocks_per_logical_kv_block
                 block_len = physical_page_size // physical_ratio
                 block_stride = physical_page_size
-                if self._is_csa_linear and cache.ndim != 1:
-                    # CSA-linear packs multiple cache owners in each block, so
-                    # state views advance by the packed block stride.
+                if (self._is_csa_linear or self._has_packed_cache) and cache.ndim != 1:
+                    # CSA-linear and block-outermost layouts pack several pages
+                    # in each block row, so state views advance by the row.
                     block_stride = cache.stride(0) * cache.element_size()
                 assert block_stride % physical_ratio == 0
                 region_specs = [
@@ -1701,7 +1762,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
                     and storage_is_block_major
                     and is_mla_region
                 )
-                if packed_mla and use_layer_name_routing:
+                if packed_mla and split_packed_rows:
                     # PP>1 prefill stage: register only this layer's page of each row.
                     region_specs = [
                         (cache.data_ptr(), physical_page_size, block_stride)
@@ -1758,10 +1819,13 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
                     ]
 
             for base_addr, block_len, block_stride in region_specs:
-                # A packed PP producer keeps one region per layer, so aliases at
-                # one address (a layer and its SWA view) stay distinct.
-                if base_addr in seen_base_addresses and not route_packed_layers:
-                    region_index = seen_base_addresses.index(base_addr)
+                # Overlaid groups can page the same bytes with different page
+                # sizes (DSv4.1 SWA, indexer and compressor caches), so views at
+                # one address share a region only if their block_len matches.
+                # A packed PP producer keeps one region per layer.
+                region_key = (base_addr, block_len)
+                if region_key in seen_region_keys and not split_packed_rows:
+                    region_index = seen_region_keys.index(region_key)
                     assert region_mem_types[region_index] == mem_type
                     self._region_is_mla[region_index] |= is_mla_region
                     if is_mla_region:
@@ -1774,6 +1838,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
                 else:
                     region_index = len(seen_base_addresses)
                     seen_base_addresses.append(base_addr)
+                    seen_region_keys.append(region_key)
                     self.block_len_per_layer.append(block_len)
                     self.block_stride_per_layer.append(block_stride)
                     self.region_group_ids.append(group_id)
@@ -1783,11 +1848,10 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
                     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 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:
@@ -1839,7 +1903,6 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
         self.kv_caches_base_addr[self.engine_id][self.tp_rank] = seen_base_addresses
         self.num_regions = len(seen_base_addresses)
         self.region_mem_types = region_mem_types
-        self._uses_region_group_mapping = len(set(self.region_group_ids)) > 1
         self._mixed_mem_types = len(set(region_mem_types)) > 1
         if self._has_mamba:
             self.region_num_blocks = [self.num_blocks] * self.num_regions
@@ -1847,22 +1910,8 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
         self._set_region_layers(region_layers)
         self._packed_member_layouts = packed_member_layouts
 
-        if self.pp_size > 1 and not use_layer_name_routing:
-            start_layer, end_layer = self.model_config.get_layers_start_end_indices(
-                self.vllm_config.parallel_config
-            )
-            num_local_layers = end_layer - start_layer
-            assert num_local_layers > 0 and self.num_regions % num_local_layers == 0
-            regions_per_layer = self.num_regions // num_local_layers
-            self._remote_region_offset = regions_per_layer * start_layer
-
         # Total local FA descriptors (boundary between FA and mamba descs).
-        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.num_descs = sum(self.region_num_blocks)
 
         self._mixed_mem_types = len(set(region_mem_types)) > 1
         self._skip_dram_xfer = self._leaves_dram_xfer_to_rank0()
@@ -1889,11 +1938,8 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
 
         self.device_kv_caches = kv_caches
         self.dst_num_blocks[self.engine_id] = self.num_blocks
-        self.dst_region_num_blocks[self.engine_id] = xfer_region_num_blocks
+        self.dst_region_num_blocks[self.engine_id] = self.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
-        )
         self.dst_region_mem_types[self.engine_id] = self.region_mem_types
 
         if self._has_mamba:
@@ -1941,6 +1987,11 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
             dcp_size=self.dcp_size,
             pcp_size=self.pcp_size,
             region_members=self.region_members,
+            layer_group_ids={
+                name: self.kv_cache_config.transfer_group_index_by_layer[name]
+                for names in self.region_members
+                for name in names
+            },
             packed_member_layouts=self._packed_member_layouts,
         )
         # Wrap metadata in payload with hash for defensive decoding
@@ -2065,10 +2116,12 @@ def _build_mamba_remote(
         nixl_agent_meta: NixlAgentMetadata,
         tp_ratio: int,
         transfer_info: EngineTransferInfo,
+        pairs: _RegionPairs,
     ) -> np.ndarray:
         """Build remote desc regions (conv sub-projections + ssm) per layer.
         For hetero-TP, each D rank reads only its sub-projection slice from
-        the P rank. Returns Nx5 uint64 strided descs."""
+        the P rank. The remote SSM and PLE regions come from ``pairs``.
+        Returns Nx5 uint64 strided descs."""
         assert nixl_agent_meta.kv_caches_base_addr, (
             "Remote KV cache base addresses must not be empty."
         )
@@ -2092,10 +2145,9 @@ def _build_mamba_remote(
         strided_descs: list[tuple[int, int, int, int, int]] = []
         # NOTE (ZhanqiuHu): use per-layer block_lens[i], not [0], in case
         # block lengths vary across layers (e.g. MLA).
-        region_indices = self._ssm_region_indices or range(
-            len(nixl_agent_meta.kv_caches_base_addr)
-        )
-        for i in region_indices:
+        assert pairs.remote_ssm_regions is not None, "SSM regions were not paired"
+        ple_region = pairs.remote_ple_region
+        for i in pairs.remote_ssm_regions:
             base_addr = nixl_agent_meta.kv_caches_base_addr[i]
             block_stride = (
                 nixl_agent_meta.block_strides[i] * remote_physical_per_logical
@@ -2110,7 +2162,8 @@ def _build_mamba_remote(
                 (ssm_addr, ssm_read_size, device_id, block_stride, num_blocks)
             )
 
-        if (region_index := self._ple_region_index) is not None:
+        if self._ple_region_index is not None:
+            assert ple_region is not None, "Remote holds no PLE region"
             local_block_len = (
                 self._ple_page_len() * self._physical_blocks_per_logical_kv_block
             )
@@ -2125,10 +2178,9 @@ def _build_mamba_remote(
                     f"local={local_block_len}, remote={remote_block_len}."
                 )
             remote_block_stride = (
-                nixl_agent_meta.block_strides[region_index]
-                * remote_physical_per_logical
+                nixl_agent_meta.block_strides[ple_region] * remote_physical_per_logical
             )
-            addr = nixl_agent_meta.kv_caches_base_addr[region_index]
+            addr = nixl_agent_meta.kv_caches_base_addr[ple_region]
             strided_descs.append(
                 (addr, remote_block_len, device_id, remote_block_stride, num_blocks)
             )
@@ -2193,9 +2245,6 @@ def _build_fa_local(
         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 = [
             self._block_strided_descs(
                 base_addresses[i],
@@ -2205,7 +2254,7 @@ def _build_fa_local(
                 block_size_ratio,
                 device_id,
             )
-            for i in region_indices
+            for i in range(len(base_addresses))
         ]
         return np.concatenate(parts)
 
@@ -2214,8 +2263,14 @@ def _build_fa_remote(
         plan: TPMapping,
         nixl_agent_meta: NixlAgentMetadata,
         block_size_ratio: int,
+        local_region_ids: Sequence[int] | None = None,
     ) -> np.ndarray:
-        """Build remote FA descriptors as Nx5 uint64 strided descs."""
+        """Build remote FA descriptors as Nx5 uint64 strided descs.
+
+        ``local_region_ids`` gives a local region each remote region pairs with
+        (``_RegionPairs.partner``); remote region ``i`` pairs with local region
+        ``i`` without it.
+        """
         assert self.transfer_topo is not None
         assert nixl_agent_meta.kv_caches_base_addr, (
             "Remote KV cache base addresses must not be empty."
@@ -2229,11 +2284,8 @@ def _build_fa_remote(
         device_id = nixl_agent_meta.device_id
         block_strides = nixl_agent_meta.block_strides
         strided_descs: list[tuple[int, int, int, int, int]] = []
-        use_layer_name_routing = bool(self._transfer_layer_region_indices)
         for i, base_addr in enumerate(nixl_agent_meta.kv_caches_base_addr):
-            local_region = (
-                self._transfer_layer_region_indices[i] if use_layer_name_routing else i
-            )
+            local_region = local_region_ids[i] if local_region_ids 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[local_region]
@@ -2464,9 +2516,11 @@ def add_remote_agent(
 
         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 self._transfer_layer_region_indices:
+        if (
+            self._supports_pp_hma
+            and self.pp_size > 1
+            and (self._is_hma_required or self._has_packed_cache)
+        ):
             if self.block_size != nixl_agent_meta.block_size:
                 raise NotImplementedError(
                     "Attention-HMA push requires identical P/D block sizes."
@@ -2476,48 +2530,8 @@ def add_remote_agent(
                     "Attention-HMA push does not support decode TP greater than "
                     "prefill TP yet"
                 )
-            self._align_remote_regions_by_layer(nixl_agent_meta)
-        elif nixl_agent_meta.packed_member_layouts and not self._packed_member_layouts:
-            # With HiSparse, P may be on BL* while D is LB*, as only indexer is on HBM.
-            # TODO (NickLucche) we may consider not allowing this case at all (P BL*)
-            # and just clear this function.
-            self._split_remote_packed_rows(nixl_agent_meta)
-        elif (
-            self.pp_size > 1
-            and len(nixl_agent_meta.kv_caches_base_addr) > num_local_regions
-        ):
-            # This worker holds a PP layer-slice; the PP=1 remote registered
-            # the full model. Slice its regions to our layer window so the
-            # logic below sees congruent local/remote lists.
-            start = self._remote_region_offset
-            end = start + num_local_regions
-            assert len(nixl_agent_meta.kv_caches_base_addr) >= end
-            nixl_agent_meta.kv_caches_base_addr = nixl_agent_meta.kv_caches_base_addr[
-                start:end
-            ]
-            nixl_agent_meta.block_lens = nixl_agent_meta.block_lens[start:end]
-            nixl_agent_meta.block_strides = nixl_agent_meta.block_strides[start:end]
-            if nixl_agent_meta.region_num_blocks is not None:
-                nixl_agent_meta.region_num_blocks = nixl_agent_meta.region_num_blocks[
-                    start:end
-                ]
-            if nixl_agent_meta.region_group_ids is not None:
-                nixl_agent_meta.region_group_ids = nixl_agent_meta.region_group_ids[
-                    start:end
-                ]
-            if nixl_agent_meta.region_names is not None:
-                nixl_agent_meta.region_names = nixl_agent_meta.region_names[start:end]
-            if nixl_agent_meta.region_mem_types is not None:
-                nixl_agent_meta.region_mem_types = nixl_agent_meta.region_mem_types[
-                    start:end
-                ]
-        # The push WRITE path is not validated with split remote blocks.
-        if self._has_mamba and self.use_mla and self._TRANSFER_MODE == "pull":
-            self._split_remote_blocks(nixl_agent_meta)
-
         if (
             self._packed_member_layouts
-            and not self._transfer_layer_region_indices
             and self._packed_member_layouts != nixl_agent_meta.packed_member_layouts
         ):
             # Whole rows transfer as one unit, so both sides must pack them alike.
@@ -2525,6 +2539,17 @@ def add_remote_agent(
                 "Packed KV block rows differ between P and D. Use the same "
                 "VLLM_KV_CACHE_LAYOUT and block size on both sides."
             )
+        pairs = self._pair_regions(nixl_agent_meta)
+        if self._dst_pairs.setdefault(engine_id, pairs) != pairs:
+            raise RuntimeError(
+                f"The TP ranks of engine {engine_id} lay out their KV regions "
+                "differently."
+            )
+        # Keep the peer's regions that pair; descriptor ids index them.
+        self._select_remote_regions(nixl_agent_meta, pairs.remote_regions)
+        # The push WRITE path is not validated with split remote blocks.
+        if self._has_mamba and self.use_mla and self._TRANSFER_MODE == "pull":
+            self._split_remote_blocks(nixl_agent_meta)
 
         ### Register remote engine in TransferTopology (idempotent).
         physical_blocks_per_logical = (
@@ -2575,15 +2600,11 @@ def add_remote_agent(
                 nixl_agent_meta.region_num_blocks
                 or [nixl_agent_meta.num_blocks] * num_remote_regions
             )
-            remote_region_group_ids = nixl_agent_meta.region_group_ids or (
+            self.dst_region_group_ids[engine_id] = 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
-            )
             self.dst_region_mem_types[engine_id] = (
                 nixl_agent_meta.region_mem_types
                 or [self.nixl_memory_type] * num_remote_regions
@@ -2594,7 +2615,7 @@ def add_remote_agent(
             nixl_agent_meta.kv_caches_base_addr
         )
         self._validate_remote_agent_handshake(
-            nixl_agent_meta, remote_tp_size, remote_dcp_size
+            nixl_agent_meta, remote_tp_size, remote_dcp_size, pairs.partner
         )
 
         # This is 1 when P and D `--tensor-parallel-size` match. Otherwise,
@@ -2674,6 +2695,7 @@ def add_remote_agent(
             plan,
             nixl_agent_meta,
             block_size_ratio,
+            pairs.partner,
         )
         logger.debug(
             "Created %s blocks for dst engine %s with remote rank %s and local rank %s",
@@ -2688,7 +2710,9 @@ def add_remote_agent(
                 engine_id,
                 remote_tp_rank,
             )
-            mamba = self._build_mamba_remote(nixl_agent_meta, tp_ratio, transfer_info)
+            mamba = self._build_mamba_remote(
+                nixl_agent_meta, tp_ratio, transfer_info, pairs
+            )
             blocks_data = np.concatenate([blocks_data, mamba])
 
         remote_mem_types = set(self.dst_region_mem_types[engine_id])
@@ -2709,9 +2733,13 @@ def _validate_remote_agent_handshake(
         nixl_agent_meta: NixlAgentMetadata,
         remote_tp_size: int,
         remote_dcp_size: int = 1,
+        local_region_ids: Sequence[int] | None = None,
     ):
         """Validate the remote agent handshake metadata ensuring the
         invariants hold true.
+
+        ``local_region_ids`` gives a local region each remote region pairs
+        with; remote region ``i`` pairs with local region ``i`` without it.
         """
         remote_engine_id = nixl_agent_meta.engine_id
 
@@ -2769,8 +2797,7 @@ def _validate_remote_agent_handshake(
         )
         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 not self._dst_pairs[remote_engine_id].groups_match
             and (
                 self._physical_blocks_per_logical_kv_block != 1
                 or remote_physical_per_logical != 1
@@ -2898,9 +2925,7 @@ 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)
-        )
+        local_regions = local_region_ids 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
@@ -2908,13 +2933,13 @@ def _validate_remote_agent_handshake(
             # match up to the kernel block size ratio even under
             # heterogeneous TP (remote kernel blocks may be smaller).
             # SSM geometry is validated via ssm_sizes/conv offsets instead.
-            assert self.block_len_per_layer == [
+            local_lens = [self.block_len_per_layer[r] for r in local_regions]
+            assert local_lens == [
                 block_len * block_size_ratio for block_len in nixl_agent_meta.block_lens
             ], (
                 "Hybrid MLA kernel-granularity block lengths must match "
                 f"between P and D (block_size_ratio={block_size_ratio}): "
-                f"local={self.block_len_per_layer}, "
-                f"remote={nixl_agent_meta.block_lens}."
+                f"local={local_lens}, remote={nixl_agent_meta.block_lens}."
             )
         elif not self._has_mamba:
             assert len(local_regions) == len(nixl_agent_meta.block_lens), (
@@ -2981,16 +3006,6 @@ 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:
-            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"
-            )
         if nixl_agent_meta.region_mem_types is not None:
             assert len(nixl_agent_meta.region_mem_types) == num_remote_regions
 
@@ -4061,7 +4076,7 @@ def _cleanup_remote_engine(
         self.dst_num_blocks.pop(engine_id, None)
         self.dst_region_num_blocks.pop(engine_id, None)
         self.dst_region_group_ids.pop(engine_id, None)
-        self.dst_uses_region_group_mapping.pop(engine_id, None)
+        self._dst_pairs.pop(engine_id, None)
         self.dst_region_mem_types.pop(engine_id, None)
         self.tp_mappings.pop(engine_id, None)
         self.draft_tp_mappings.pop(engine_id, None)
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 ab3e3f83a3a5..c2a1adbcb239 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py
@@ -51,8 +51,10 @@
 #  13: Add packed-member layouts and order-independent packed-push backend hashes
 #  14: Add Qwen3.8-Flash-Next PLE page length - #59997
 #  15: Add remote_registration_generation to kv_transfer_params
+#  16: Pair regions by the layers both sides hold: every peer advertises
+#      region_members and the cache group of each layer
 #
-NIXL_CONNECTOR_VERSION: int = 15
+NIXL_CONNECTOR_VERSION: int = 16
 
 
 @dataclass
@@ -79,6 +81,8 @@ class NixlAgentMetadata:
     ple_block_len: int | None = None
     # Layer names sharing each advertised region, in region order.
     region_members: list[list[str]] = field(default_factory=list)
+    # Cache group of each layer in region_members.
+    layer_group_ids: dict[str, int] = field(default_factory=dict)
     # Packed member -> (byte offset in its region's block, bytes per page).
     packed_member_layouts: dict[str, tuple[int, int]] = field(default_factory=dict)
 
@@ -105,7 +109,7 @@ def _get_speculative_compatibility_factors(
 ) -> dict[str, Any] | None:
     """Return NIXL compatibility factors for hidden-state-based speculators."""
     speculative_config = vllm_config.speculative_config
-    if speculative_config is None or not speculative_config.use_eagle():
+    if speculative_config is None or not speculative_config.uses_target_hidden_states:
         return None
 
     draft_model_config = speculative_config.draft_model_config
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 6c436f892866..a2c9661fda09 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
@@ -187,6 +187,7 @@ def _read_blocks_for_req(self, req_id: str, meta: ReqMeta):
         local_block_ids = meta.local_physical_block_ids
         remote_region_groups = self.dst_region_group_ids[engine_id]
         local_region_groups = self.region_group_ids or remote_region_groups
+        pairs = self._dst_pairs[engine_id]
         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.
@@ -194,7 +195,14 @@ def _read_blocks_for_req(self, req_id: str, meta: ReqMeta):
                 ReadSpec(remote_rank=rank, local_block_ids=[], remote_block_ids=[])
                 for rank in plan.all_source_ranks
             ]
-        elif local_region_groups != remote_region_groups:
+        elif not pairs.groups_match:
+            # The two sides put some layer in different cache groups: read each
+            # region at its own blocks.
+            if pairs.remote != pairs.local or len(pairs.local) != self.num_regions:
+                raise NotImplementedError(
+                    "Different NIXL cache-group layouts need the same KV regions "
+                    "on both sides"
+                )
             if not self.use_mla or self._has_mamba:
                 raise NotImplementedError(
                     "Different NIXL cache-group layouts are only supported for "
@@ -497,23 +505,19 @@ def _read_blocks(
         # corresponding rank. With heterogeneous TP, fixing D>P, the D tp
         # workers will issue xfers to parts of the P worker remote kv caches.
 
-        # Get descs ids.
+        # Get descs ids. Pair k reads remote region pairs.remote[k] into local
+        # region pairs.local[k], at its cache group's blocks (its own regions'
+        # blocks for region-mapped reads).
+        pairs = self._dst_pairs[dst_engine_id]
+        by_region = read_spec.block_ids_by_region
         remote_block_descs_ids = self._compute_desc_ids(
             block_ids=remote_block_ids,
             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.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.get(dst_engine_id) or None)
-            ),
-            uses_region_group_mapping=(
-                self.num_regions > 1
-                if read_spec.block_ids_by_region
-                else self.dst_uses_region_group_mapping[dst_engine_id]
-            ),
+            pairs=(pairs.remote, pairs.remote if by_region else pairs.group),
+            ssm_regions=pairs.remote_ssm_regions,
         )
         local_block_descs_ids = self._compute_desc_ids(
             block_ids=local_block_ids,
@@ -521,16 +525,7 @@ def _read_blocks(
             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.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 or None)
-            ),
-            uses_region_group_mapping=(
-                self.num_regions > 1
-                if read_spec.block_ids_by_region
-                else self._uses_region_group_mapping
-            ),
+            pairs=(pairs.local, pairs.local if by_region else pairs.group),
         )
 
         assert len(local_block_descs_ids) == len(remote_block_descs_ids)
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 00f2b30d7508..6f975a2b5836 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
@@ -35,6 +35,7 @@
 import threading
 import time
 from collections import defaultdict
+from collections.abc import Sequence
 from concurrent.futures import Future
 from typing import TYPE_CHECKING, Any
 
@@ -157,9 +158,10 @@ def _validate_remote_agent_handshake(
         nixl_agent_meta: NixlAgentMetadata,
         remote_tp_size: int,
         remote_dcp_size: int = 1,
+        local_region_ids: Sequence[int] | None = None,
     ) -> None:
         super()._validate_remote_agent_handshake(
-            nixl_agent_meta, remote_tp_size, remote_dcp_size
+            nixl_agent_meta, remote_tp_size, remote_dcp_size, local_region_ids
         )
         assert self.transfer_topo is not None
         if (
@@ -561,10 +563,7 @@ 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
-        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:
+        if not self._dst_pairs[engine_id].groups_match:
             raise NotImplementedError(
                 "NixlPushConnector does not support different producer and "
                 "consumer cache-group layouts"
@@ -707,15 +706,17 @@ def _xfer_blocks(
                 f"{len(remote_block_ids[i])} remote blocks"
             )
 
-        # Get descs ids.
+        # Get descs ids. Pair k writes local region pairs.local[k] into remote
+        # region pairs.remote[k] at its cache group's blocks.
+        pairs = self._dst_pairs[dst_engine_id]
         remote_block_descs_ids = self._compute_desc_ids(
             block_ids=remote_block_ids,
             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_group_ids=self.dst_region_group_ids[dst_engine_id],
-            uses_region_group_mapping=self.dst_uses_region_group_mapping[dst_engine_id],
+            pairs=(pairs.remote, pairs.group),
+            ssm_regions=pairs.remote_ssm_regions,
         )
         local_block_descs_ids = self._compute_desc_ids(
             block_ids=local_block_ids,
@@ -723,8 +724,7 @@ def _xfer_blocks(
             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_group_ids=self.region_group_ids,
-            uses_region_group_mapping=self._uses_region_group_mapping,
+            pairs=(pairs.local, pairs.group),
         )
 
         assert len(local_block_descs_ids) == len(remote_block_descs_ids)
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/common.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/common.py
index 61cbcf99c1a5..afcd787dbaa7 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/common.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/common.py
@@ -85,11 +85,14 @@ class OffloadingWorkerMetadata(KVConnectorWorkerMetadata):
     """
 
     completed_jobs: dict[int, int] = field(default_factory=dict)
+    failed_jobs: set[int] = field(default_factory=set)
     transfer_stats: TransferStats = field(default_factory=TransferStats)
 
-    def mark_completed(self, job_id: int) -> None:
+    def mark_completed(self, job_id: int, success: bool = True) -> None:
         """Record a transfer job completion from this worker."""
         self.completed_jobs[job_id] = 1
+        if not success:
+            self.failed_jobs.add(job_id)
 
     def aggregate(
         self, other: "KVConnectorWorkerMetadata"
@@ -102,5 +105,6 @@ def aggregate(
 
         return OffloadingWorkerMetadata(
             completed_jobs=merged,
+            failed_jobs=self.failed_jobs | other.failed_jobs,
             transfer_stats=self.transfer_stats.aggregate(other.transfer_stats),
         )
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 23fe932451d3..8b53a7dfb50b 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/events.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/events.py
@@ -13,12 +13,21 @@
 of the same hash. Opt-in via
 ``kv_connector_extra_config["self_describing_kv_events"]``; inert unless
 KV cache events are enabled. See the PR description for the full design.
+
+Recurrent (Mamba / linear-attention) state groups announce state checkpoints,
+not token ranges: an offloaded chunk holds the recurrent state at the chunk's
+last hash only. Their ``BlockStored`` therefore carries that single hash, its
+parent, an empty ``token_ids`` and ``kv_cache_spec_kind`` ``"mamba"``. A
+consumer aligns it by hash with the attention group's blocks, and a hybrid
+prefix is reusable only up to a boundary where every group has a residency.
+Windowed attention groups keep the legacy placeholder payload.
 """
 
 from collections.abc import Iterable
-from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, NamedTuple
 
+import msgspec
+
 from vllm.distributed.kv_events import (
     MEDIUM_CPU,
     MEDIUM_STORAGE,
@@ -35,6 +44,7 @@
 )
 from vllm.v1.kv_cache_interface import (
     KVCacheGroupSpec,
+    MambaSpec,
     get_kv_cache_spec_kind,
     get_kv_cache_spec_sliding_window,
 )
@@ -60,8 +70,8 @@
     Medium.STORAGE: MEDIUM_STORAGE,
 }
 
-# The tracker is CPU-primary by construction: the primary tier's residency key.
-_PRIMARY_RESIDENCY: tuple[Medium, str | None] = (Medium.CPU, None)
+# The tracker is CPU-primary by construction: the primary tier's residency bit.
+_PRIMARY_RESIDENCY_BIT = 1
 
 
 class OffloadingEventGroupSpec(NamedTuple):
@@ -79,11 +89,16 @@ def get_offloading_event_group_spec(
     )
 
 
-@dataclass(slots=True)
-class _OffloadEventMetadata:
+class _OffloadEventMetadata(
+    msgspec.Struct,
+    gc=False,  # type: ignore[call-arg]
+):
     """BlockStored payload snapshot for one OffloadKey, captured while the
     Request is available and kept until the final matching removal event.
-    ``medium`` and ``ownership`` are forwarded from the OffloadingEvent."""
+    ``medium`` and ``ownership`` are forwarded from the OffloadingEvent.
+
+    One is kept per offloaded chunk, so it is untracked by the cyclic GC: keep
+    fields free of tracked containers (``active_residencies`` is a bitmask)."""
 
     # The chunk's constituent block hashes; the last one is the OffloadKey.
     block_hashes: tuple[BlockHash, ...]
@@ -96,7 +111,40 @@ class _OffloadEventMetadata:
     extra_keys: tuple[tuple[Any, ...] | None, ...] | None
     group_idx: int
     kv_cache_spec: OffloadingEventGroupSpec
-    active_residencies: set[tuple[Medium, str | None]]
+    active_residencies: int
+
+
+def describes_group(group_config: "GroupOffloadConfig") -> bool:
+    """Whether the tracker builds self-describing payloads for this group.
+
+    Full-attention groups and recurrent state groups (``MambaSpec``) do;
+    windowed attention groups (sliding window, chunked local) keep the legacy
+    placeholder because a window's hashes do not identify a reusable prefix.
+    """
+    if group_config.sliding_window_size_in_chunks is None:
+        return True
+    return _is_recurrent_group(group_config)
+
+
+def _is_recurrent_group(group_config: "GroupOffloadConfig") -> bool:
+    return isinstance(group_config.kv_cache_spec, MambaSpec)
+
+
+def _as_state_checkpoint(meta: _OffloadEventMetadata) -> _OffloadEventMetadata:
+    """Reduce a chunk payload to the recurrent state checkpoint it holds.
+
+    A recurrent group's offloaded chunk is one state, valid at the chunk's
+    last hash. Announcing the whole hash range would let a consumer treat
+    every hash in it as a resumable position, so the payload keeps only the
+    last hash (the OffloadKey), chained to the hash before it, and no tokens:
+    the consumer aligns it with the attention group's blocks by hash.
+    """
+    assert meta.block_hashes
+    if len(meta.block_hashes) > 1:
+        meta.parent_block_hash = meta.block_hashes[-2]
+    meta.block_hashes = meta.block_hashes[-1:]
+    meta.token_ids = ()
+    return meta
 
 
 class OffloadingEventsTracker:
@@ -122,6 +170,10 @@ def __init__(
 
         # OffloadKey -> payload snapshot, kept until final removal or reset.
         self._pending_event_metadata: dict[OffloadKey, _OffloadEventMetadata] = {}
+        # (medium, ownership) -> its bit in active_residencies.
+        self._residency_bits: dict[tuple[Medium, str | None], int] = {
+            (Medium.CPU, None): _PRIMARY_RESIDENCY_BIT
+        }
 
     def record_store(
         self,
@@ -132,16 +184,16 @@ def record_store(
     ) -> None:
         """Snapshot the KV cache event payload for one offloaded chunk.
 
-        No-op when self-describing event capture is disabled or for
-        sliding-window / SSM groups, which keep the legacy placeholder payload.
+        No-op when self-describing event capture is disabled or for windowed
+        attention groups, which keep the legacy placeholder payload.
         """
         if not self.self_describing_enabled:
             return
-        if group_config.sliding_window_size_in_chunks is not None:
+        if not describes_group(group_config):
             return
         meta = self._build_event_metadata(req, group_config, chunk_idx)
         if existing := self._pending_event_metadata.get(offload_key):
-            meta.active_residencies.update(existing.active_residencies)
+            meta.active_residencies |= existing.active_residencies
         self._pending_event_metadata[offload_key] = meta
 
     def record_lookup(
@@ -154,19 +206,17 @@ def record_lookup(
         """Snapshot metadata for a ready or write-pending primary-tier hit."""
         if not self.self_describing_enabled:
             return
-        if group_config.sliding_window_size_in_chunks is not None:
+        if not describes_group(group_config):
             return
-        if offload_key not in self._pending_event_metadata:
-            self._pending_event_metadata[offload_key] = self._build_event_metadata(
-                req, group_config, chunk_idx
-            )
+        meta = self._pending_event_metadata.get(offload_key)
+        if meta is None:
+            meta = self._build_event_metadata(req, group_config, chunk_idx)
+            self._pending_event_metadata[offload_key] = meta
         # Primary residency is registered here rather than by the manager's
         # prepare_write(): the payload needs the Request, which only the
         # scheduler holds at lookup time, and the manager reaches this tracker
         # only through OffloadingEvents.
-        self._pending_event_metadata[offload_key].active_residencies.add(
-            _PRIMARY_RESIDENCY
-        )
+        meta.active_residencies |= _PRIMARY_RESIDENCY_BIT
 
     def record_partial_store(
         self,
@@ -176,7 +226,7 @@ def record_partial_store(
         offload_key: OffloadKey,
     ) -> None:
         """Snapshot metadata for a newly stored partial recurrent tail."""
-        if group_config.sliding_window_size_in_chunks is not None:
+        if not describes_group(group_config):
             return
         self._record_partial_tail(req, group_config, boundary_tokens, offload_key)
 
@@ -188,7 +238,7 @@ def record_partial_lookup(
         offload_key: OffloadKey,
     ) -> None:
         """Backfill metadata for a partial recurrent tail lookup hit."""
-        if group_config.sliding_window_size_in_chunks is not None:
+        if not describes_group(group_config):
             return
         if offload_key not in self._pending_event_metadata:
             self._record_partial_tail(req, group_config, boundary_tokens, offload_key)
@@ -245,18 +295,20 @@ def _record_partial_tail(
             extra_keys=None,
             group_idx=group_config.group_idx,
             kv_cache_spec=group_config.kv_event_group_spec,
-            active_residencies={_PRIMARY_RESIDENCY},
+            active_residencies=_PRIMARY_RESIDENCY_BIT,
         )
+        if _is_recurrent_group(group_config):
+            meta = _as_state_checkpoint(meta)
         if existing := self._pending_event_metadata.get(offload_key):
-            meta.active_residencies.update(existing.active_residencies)
+            meta.active_residencies |= existing.active_residencies
         self._pending_event_metadata[offload_key] = meta
 
     def take_events(self, events: Iterable[OffloadingEvent]) -> Iterable[KVCacheEvent]:
         """Translate raw OffloadingEvents into self-describing KV events.
 
-        Complete metadata is available only for full-attention groups when
-        the tracker is enabled. Other shapes retain the legacy placeholder
-        payload so consumers can ignore them.
+        Complete metadata is available for full-attention and recurrent
+        state groups when the tracker is enabled. Windowed attention groups
+        retain the legacy placeholder payload so consumers can ignore them.
 
         Yields:
             ``BlockStored`` or ``BlockRemoved`` events corresponding to
@@ -292,7 +344,8 @@ def _build_event_metadata(
     ) -> _OffloadEventMetadata:
         """Build the payload snapshot for one offloaded chunk: its
         constituent per-block hashes, the whole chunk's tokens, and the
-        per-block ``block_size``."""
+        per-block ``block_size``. A recurrent group's chunk is reduced to
+        its state checkpoint (see :func:`_as_state_checkpoint`)."""
         hashes_per_chunk = group_config.hashes_per_chunk
         assert hashes_per_chunk > 0
         assert chunk_idx >= 0
@@ -314,9 +367,12 @@ def _build_event_metadata(
             group_config.tokens_per_chunk // group_config.tokens_per_block
         )
 
-        if group_config.sliding_window_size_in_chunks is not None:
+        if not describes_group(group_config):
             # The recording methods filter these out before calling this helper.
-            raise AssertionError("self-describing events only support full attention")
+            raise AssertionError(
+                "self-describing events only support full-attention and "
+                "recurrent state groups"
+            )
 
         parent_block_hash: BlockHash | None
         if first_hash_idx == 0:
@@ -328,7 +384,12 @@ def _build_event_metadata(
         tok_start = chunk_idx * group_config.tokens_per_chunk
         tok_end = tok_start + group_config.tokens_per_chunk
         assert tok_end <= len(req.all_token_ids)
-        token_ids = tuple(req.all_token_ids[tok_start:tok_end])
+        # A recurrent chunk's payload carries no tokens; skip the copy.
+        token_ids: tuple[int, ...] = (
+            ()
+            if _is_recurrent_group(group_config)
+            else tuple(req.all_token_ids[tok_start:tok_end])
+        )
 
         lora_id: int | None = None
         lora_name: str | None = None
@@ -336,7 +397,7 @@ def _build_event_metadata(
             lora_id = req.lora_request.adapter_id
             lora_name = req.lora_request.name
 
-        return _OffloadEventMetadata(
+        meta = _OffloadEventMetadata(
             block_hashes=tuple(chunk_hashes),
             parent_block_hash=parent_block_hash,
             token_ids=token_ids,
@@ -346,8 +407,11 @@ def _build_event_metadata(
             extra_keys=None,
             group_idx=group_config.group_idx,
             kv_cache_spec=group_config.kv_event_group_spec,
-            active_residencies={_PRIMARY_RESIDENCY},
+            active_residencies=_PRIMARY_RESIDENCY_BIT,
         )
+        if _is_recurrent_group(group_config):
+            meta = _as_state_checkpoint(meta)
+        return meta
 
     def _placeholder_stored(
         self,
@@ -386,6 +450,11 @@ def _take_stored_event(self, event: OffloadingEvent) -> Iterable[KVCacheEvent]:
         # eviction event so BlockRemoved can fan out to the same hashes.
         # Events are self-contained (own parent), so key order is free.
         locality = event.locality.value if event.locality is not None else None
+        residency_bit = 0
+        if event.removal_expected:
+            residency_bit = self._residency_bits.setdefault(
+                (event.medium, event.ownership), 1 << len(self._residency_bits)
+            )
         for key in event.keys:
             meta = self._pending_event_metadata.get(key)
             if meta is None:
@@ -394,7 +463,7 @@ def _take_stored_event(self, event: OffloadingEvent) -> Iterable[KVCacheEvent]:
                     logger.warning_once(
                         "OffloadingEventsTracker: no event metadata for "
                         "offload key during BlockStored emission; emitting a "
-                        "placeholder payload. Expected for non-full-attention "
+                        "placeholder payload. Expected for windowed attention "
                         "groups and promotions not observed as a primary-tier "
                         "hit before translation."
                     )
@@ -403,8 +472,7 @@ def _take_stored_event(self, event: OffloadingEvent) -> Iterable[KVCacheEvent]:
                 )
                 continue
 
-            if event.removal_expected:
-                meta.active_residencies.add((event.medium, event.ownership))
+            meta.active_residencies |= residency_bit
             yield BlockStored(
                 block_hashes=list(
                     maybe_convert_block_hash(h) for h in meta.block_hashes
@@ -433,18 +501,19 @@ def _take_removed_event(self, event: OffloadingEvent) -> Iterable[KVCacheEvent]:
         # Keep group_idx unambiguous if a manager batch spans groups.
         locality = event.locality.value if event.locality is not None else None
         by_group: dict[int, list] = {}
+        residency_mask = ~self._residency_bits.get((event.medium, event.ownership), 0)
         for key in event.keys:
             meta = self._pending_event_metadata.get(key)
             if event.metadata_only:
                 if meta is not None:
-                    meta.active_residencies.discard((event.medium, event.ownership))
+                    meta.active_residencies &= residency_mask
                 continue
             if meta is not None:
                 group_idx = meta.group_idx
                 by_group.setdefault(group_idx, []).extend(
                     maybe_convert_block_hash(h) for h in meta.block_hashes
                 )
-                meta.active_residencies.discard((event.medium, event.ownership))
+                meta.active_residencies &= residency_mask
             else:
                 if self.self_describing_enabled:
                     logger.warning_once(
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 5bf8a77d351d..ca0a5e757e3b 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py
@@ -5,7 +5,7 @@
 from collections.abc import Iterable, Iterator, Sequence
 from contextlib import ExitStack, contextmanager
 from dataclasses import dataclass, field, replace
-from itertools import chain, islice
+from itertools import chain
 from typing import Any, NamedTuple
 
 from vllm.config import VllmConfig
@@ -82,11 +82,13 @@ class TransferJobStatus:
     # Number of workers still pending. Starts at num_workers,
     # decremented as each worker reports completion. Job is done at 0.
     pending_count: int
-    # Offload keys this job covers; passed to manager.complete_*().
-    keys: set[OffloadKey]
+    # Offload keys this job covers, in prefix order; passed to
+    # manager.complete_*().
+    keys: tuple[OffloadKey, ...]
     is_store: bool
     # Unlike an async load, store jobs may be created next to it.
     is_sync_load: bool = False
+    failed: bool = False
     # Store source blocks fenced after the request finishes.
     deferred_fence_block_ids: list[int] | None = None
     # Store source blocks fenced when the transfer is created.
@@ -374,6 +376,9 @@ class RequestOffloadState:
     load_key_ranges: list[tuple[int, int, OffloadKey]] = field(default_factory=list)
     # True once on_request_finished has been signaled to the manager.
     finished_signaled: bool = False
+    # A load for this request failed; all its later lookups, including after
+    # preemption, skip the offloaded cache so a faulting load is not retried.
+    load_failed: bool = False
 
     def __post_init__(self) -> None:
         self.group_states = tuple(
@@ -413,14 +418,9 @@ def update_offload_keys(self) -> None:
         for group_config, group_state in zip(
             self.config.kv_group_configs, self.group_states
         ):
-            for req_block_hash in islice(
-                self.req.block_hashes,
-                group_config.hashes_per_chunk * len(group_state.offload_keys)
-                + group_config.hashes_per_chunk
-                - 1,
-                None,
-                group_config.hashes_per_chunk,
-            ):
+            stride = group_config.hashes_per_chunk
+            start = stride * (len(group_state.offload_keys) + 1) - 1
+            for req_block_hash in self.req.block_hashes[start::stride]:
                 key = make_offload_key(req_block_hash, group_config.group_idx)
                 group_state.offload_keys.append(key)
                 self.req_context.set_offload_key_position(
@@ -890,13 +890,17 @@ def _lookup_complete_chunks(
         self,
         req_status: RequestOffloadState,
         max_num_new_tokens: int | None = None,
-    ) -> int | None:
+    ) -> tuple[int | 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
         can invalidate an earlier group's result, so the loop re-runs when that
         happens until num_hit_tokens converges.
+
+        Returns the number of hit tokens (None to defer) and the hit boundary
+        the full-attention groups proved before any sliding-window/Mamba group
+        narrowed it (None if not reached).
         """
         num_computed_tokens = req_status.num_locally_computed_tokens
         max_hit_size_tokens: int = req_status.req.num_tokens
@@ -930,6 +934,7 @@ def _lookup_complete_chunks(
         # in the current convergence iteration. Reset when a non-eagle group
         # tightens the hit boundary, requiring a fresh pop.
         eagle_verified: set[int] = set()
+        full_attention_hit_boundary: int | None = None
         while lookup_groups:
             looked_up_sliding_window: bool = False
             groups_iter = iter(lookup_groups)
@@ -960,11 +965,21 @@ def _lookup_complete_chunks(
                     )
                 if max_hit_size_tokens - num_computed_tokens < tokens_per_chunk:
                     # We can only load less than a chunk, so skip.
-                    return 0
+                    return 0, full_attention_hit_boundary
 
                 sliding_window_size_in_chunks = (
                     group_config.sliding_window_size_in_chunks
                 )
+                if (
+                    sliding_window_size_in_chunks is not None
+                    and full_attention_hit_boundary is None
+                ):
+                    # Groups before the first sparse group are full attention.
+                    full_attention_hit_boundary = (
+                        max_hit_size_tokens
+                        if self._full_attention_groups and not defer_lookup
+                        else 0
+                    )
 
                 # For eagle groups, query one extra chunk that will be popped.
                 # Widening applies to every group type: without it, the pop
@@ -1016,7 +1031,7 @@ def _lookup_complete_chunks(
                         initial_window + int(is_eagle_unverified),
                     )
                 if num_hit_chunks == 0:
-                    return 0
+                    return 0, full_attention_hit_boundary
 
                 if num_hit_chunks is None:
                     defer_lookup = True
@@ -1033,7 +1048,7 @@ def _lookup_complete_chunks(
                 new_num_hit_tokens = max_hit_size_tokens - num_computed_tokens
                 if new_num_hit_tokens < tokens_per_chunk:
                     # We can only load less than a chunk, so skip.
-                    return 0
+                    return 0, full_attention_hit_boundary
 
                 if new_num_hit_tokens < num_hit_tokens:
                     if not group_config.is_eagle_group:
@@ -1056,7 +1071,7 @@ def _lookup_complete_chunks(
                 "Offloading manager delayed request %s as backend requested",
                 req_status.req.request_id,
             )
-            return None
+            return None, full_attention_hit_boundary
 
         # Possibly delay the request if any hit chunk is already being loaded.
         if self._chunks_being_loaded:
@@ -1082,7 +1097,7 @@ def _lookup_complete_chunks(
                         " chunks are already being loaded",
                         req_status.req.request_id,
                     )
-                    return None
+                    return None, full_attention_hit_boundary
 
         logger.debug(
             "Request %s hit %s offloaded tokens after %s GPU hit tokens",
@@ -1091,7 +1106,7 @@ def _lookup_complete_chunks(
             num_computed_tokens,
         )
 
-        return num_hit_tokens
+        return num_hit_tokens, full_attention_hit_boundary
 
     def _make_boundary_key(
         self,
@@ -1110,7 +1125,26 @@ def _lookup(
         req_status: RequestOffloadState,
         max_num_new_tokens: int | None = None,
     ) -> int | None:
-        complete_hit = self._lookup_complete_chunks(req_status, max_num_new_tokens)
+        complete_hit, full_attention_hit_boundary = self._lookup_complete_chunks(
+            req_status, max_num_new_tokens
+        )
+        num_hit_tokens = self._lookup_partial_tail_hit(
+            req_status, complete_hit, full_attention_hit_boundary, max_num_new_tokens
+        )
+        if num_hit_tokens is not None and full_attention_hit_boundary is not None:
+            self._maybe_register_shared_prefix_junction(
+                req_status, full_attention_hit_boundary, num_hit_tokens
+            )
+        return num_hit_tokens
+
+    def _lookup_partial_tail_hit(
+        self,
+        req_status: RequestOffloadState,
+        complete_hit: int | None,
+        full_attention_hit_boundary: int | None,
+        max_num_new_tokens: int | None,
+    ) -> int | None:
+        """Extend a complete-chunk hit with a partial tail, if supported."""
         req_status.partial_tail_boundary = None
         if complete_hit is None or not self.config.supports_partial_tail:
             return complete_hit
@@ -1118,10 +1152,26 @@ def _lookup(
         local_tokens = req_status.num_locally_computed_tokens
         complete_boundary = local_tokens + complete_hit
         tokens_per_hash = self.config.tokens_per_hash
+        block_size = self._partial_tail_block_size
+        # A recurrent group holds state only at a producer's tail, so the
+        # full-attention prefix, not the complete hit, bounds the search.
+        anchor_boundary = max(complete_boundary, full_attention_hit_boundary or 0)
+        anchor_pending = anchor_boundary > complete_boundary and any(
+            key in (self._chunks_being_loaded or ())
+            for group_config, group_state in zip(
+                self.config.kv_group_configs, req_status.group_states
+            )
+            if not group_config.requires_cow_source
+            for key in group_state.offload_keys[
+                complete_boundary // block_size : anchor_boundary // block_size
+            ]
+        )
+        if anchor_pending:
+            anchor_boundary = complete_boundary
         # A partial tail carries its own EAGLE proof, so it may also lie in
         # the block after the chunk that the complete-chunk lookup pops.
         num_search_blocks = 2 if self._max_eagle_proof_margin else 1
-        block_end = complete_boundary + (
+        block_end = anchor_boundary + (
             num_search_blocks * self._partial_tail_block_size
         )
         max_boundary = min(req_status.req.num_prompt_tokens - 1, block_end - 1)
@@ -1143,16 +1193,15 @@ def _lookup(
         if max_boundary <= complete_boundary:
             return complete_hit
 
-        # A boundary past the block of complete_boundary also loads that
+        # A boundary past the block of anchor_boundary also loads that
         # block's complete chunk. Look it up once, and only if it is needed.
         # Partial tails use one block per chunk, so a block index is a chunk
         # index.
-        block_size = self._partial_tail_block_size
-        bridge_chunk_idx = complete_boundary // block_size
+        bridge_chunk_idx = anchor_boundary // block_size
         bridge_boundary = (bridge_chunk_idx + 1) * block_size
         # Bridge keys that the complete-chunk lookup did not confirm.
         unconfirmed_bridge_keys: list[tuple[GroupOffloadConfig, OffloadKey]] = []
-        pending = False
+        pending = anchor_pending
         if max_boundary >= bridge_boundary:
             for group_config, group_state in zip(
                 self.config.kv_group_configs, req_status.group_states
@@ -1171,7 +1220,12 @@ def _lookup(
                     break
                 pending |= result is not LookupResult.HIT
         bridge_pending = pending
+        # A tail of this request may lie below a sibling's longer full-attention
+        # prefix. A block-aligned boundary there is a complete chunk, which the
+        # complete-chunk lookup already missed.
         for boundary in range(max_boundary, complete_boundary, -tokens_per_hash):
+            if boundary <= anchor_boundary and boundary % block_size == 0:
+                continue
             loads_bridge = boundary >= bridge_boundary
             hit, boundary_pending = self._lookup_partial_tail(
                 req_status, boundary, loads_bridge and bridge_pending
@@ -1220,6 +1274,10 @@ def _lookup_partial_tail(
             key = self._make_boundary_key(
                 req, group_config.group_idx, key_tokens, req_context
             )
+            # The offload manager can report HIT as soon as the source chunk
+            # is ready, while its earlier GPU load is still in flight.  Do
+            # not schedule the same partial-tail transfer twice.
+            pending |= key in (self._chunks_being_loaded or ())
             result = self.manager.lookup(key, req_context)
             if result is LookupResult.MISS:
                 return False, pending
@@ -1234,6 +1292,21 @@ def _lookup_partial_tail(
             )
         return True, False
 
+    def _maybe_register_shared_prefix_junction(
+        self, req_status: RequestOffloadState, boundary: int, num_hit_tokens: int
+    ) -> None:
+        """Pin a shared-prefix junction where a sparse group cut the hit, as
+        ``get_computed_blocks`` does for a local hit, so the recompute retains
+        the sliding-window/Mamba state at the full-attention hit boundary."""
+        if self.config.retention_interval is None:
+            return
+        if boundary <= req_status.num_locally_computed_tokens + num_hit_tokens:
+            return
+        request = req_status.req
+        # Never lower a junction the local prefix cache already set.
+        if boundary > request.shared_prefix_boundary:
+            request.shared_prefix_boundary = boundary
+
     def on_new_request(self, request: Request) -> None:
         """Called when a new request is added to the scheduler."""
         req_context = _create_req_context(request, self)
@@ -1299,7 +1372,7 @@ def get_num_new_matched_tokens(
         req_status.num_locally_computed_tokens = num_computed_tokens
 
         num_hit_tokens: int | None
-        if request.skip_reading_external_kv:
+        if request.skip_reading_external_kv or req_status.load_failed:
             num_hit_tokens = 0
         else:
             lookup_start = time.monotonic()
@@ -1533,7 +1606,7 @@ def update_state_after_alloc(
         self._jobs[load_job_id] = TransferJobStatus(
             req_id=request.request_id,
             pending_count=self.config.num_workers,
-            keys=set(keys_to_load),
+            keys=tuple(dict.fromkeys(keys_to_load)),
             is_store=False,
             is_sync_load=sync_load,
         )
@@ -1640,7 +1713,7 @@ def _build_aligned_boundary_store_jobs(
                 self._jobs[job_id] = TransferJobStatus(
                     req_id=req_id,
                     pending_count=self.config.num_workers,
-                    keys={key},
+                    keys=(key,),
                     is_store=True,
                     fenced_block_ids=[block_id],
                 )
@@ -1787,7 +1860,7 @@ def _build_partial_tail_store_jobs(
             self._jobs[job_id] = TransferJobStatus(
                 req_id=req_id,
                 pending_count=self.config.num_workers,
-                keys=set(store_output.keys_to_store),
+                keys=tuple(dict.fromkeys(store_output.keys_to_store)),
                 is_store=True,
                 fenced_block_ids=source_blocks,
             )
@@ -1877,7 +1950,12 @@ def _build_store_jobs(
                 # Clamp to the GPU prefix cache's commit point. The final sampled
                 # token's slot is never committed (under spec decode it holds a
                 # rejected draft's KV), so a block ending there must not be stored.
-                num_tokens_after_batch = max(req.num_prompt_tokens, req.num_tokens - 1)
+                # A request that finished before its prefill completed (e.g. an
+                # error) has computed only num_computed_tokens of its prompt.
+                num_tokens_after_batch = min(
+                    req.num_computed_tokens,
+                    max(req.num_prompt_tokens, req.num_tokens - 1),
+                )
             else:
                 num_scheduled_tokens = scheduler_output.num_scheduled_tokens[req_id]
                 num_tokens_after_batch = req.num_computed_tokens + num_scheduled_tokens
@@ -2043,7 +2121,8 @@ def _build_store_jobs(
                 req_status.advance_stored_idx(num_offloadable_tokens)
                 continue
 
-            keys_to_store = set(store_output.keys_to_store)
+            ordered_keys_to_store = tuple(dict.fromkeys(store_output.keys_to_store))
+            keys_to_store = set(ordered_keys_to_store)
 
             group_sizes: list[int] = []
             block_indices: list[int] = []
@@ -2117,7 +2196,7 @@ def _build_store_jobs(
             self._jobs[job_id] = TransferJobStatus(
                 req_id=req_id,
                 pending_count=self.config.num_workers,
-                keys=set(keys_to_store),
+                keys=ordered_keys_to_store,
                 is_store=True,
                 deferred_fence_block_ids=deferred_fence_block_ids,
                 fenced_block_ids=fenced_block_ids or None,
@@ -2325,6 +2404,8 @@ def _update_connector_output(self, connector_output: KVConnectorOutput):
                 )
                 continue
             job_status = self._jobs[job_id]
+            if job_id in meta.failed_jobs:
+                job_status.failed = True
             job_status.pending_count -= count
             if job_status.pending_count > 0:
                 continue
@@ -2332,8 +2413,27 @@ def _update_connector_output(self, connector_output: KVConnectorOutput):
 
             req_status = self._req_status[job_status.req_id]
             if job_status.is_store:
-                self.manager.complete_store(job_status.keys, req_status.req_context)
+                if job_status.failed:
+                    logger.warning(
+                        "KV offload store job %d failed; dropping %d chunk(s) "
+                        "from the CPU cache rather than publishing them.",
+                        job_id,
+                        len(job_status.keys),
+                    )
+                self.manager.complete_store(
+                    job_status.keys,
+                    req_status.req_context,
+                    success=not job_status.failed,
+                )
             else:
+                if job_status.failed:
+                    logger.warning(
+                        "KV offload load job %d failed; request %s will not "
+                        "load from the offloaded cache again.",
+                        job_id,
+                        job_status.req_id,
+                    )
+                    req_status.load_failed = True
                 self.manager.complete_load(job_status.keys, req_status.req_context)
                 if self._chunks_being_loaded:
                     self._chunks_being_loaded.difference_update(job_status.keys)
@@ -2457,10 +2557,11 @@ def take_events(self) -> Iterable[KVCacheEvent]:
 
     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
+        # reset_cache cannot be called in the middle of a schedule step.
+        # _current_batch_jobs_to_flush may still hold a previous reset's flush
+        # set if no step ran since; the new ids are merged into it.
         assert not self._current_batch_load_jobs
         assert not self._current_batch_sync_loads
-        assert not self._current_batch_jobs_to_flush
         assert not self._current_batch_allocated_block_ids
 
         # Held for the whole reset. manager.reset_cache() drains tiers by
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 616b20e4d085..8698deee8733 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py
@@ -60,6 +60,7 @@ def __init__(
 
         # job_id -> req_id for in-flight loads (None for sync loads).
         self._load_jobs: dict[int, ReqId | None] = {}
+        self._failed_recving: set[ReqId] = set()
         self._unsubmitted_store_jobs: list[
             tuple[int, GPULoadStoreSpec, LoadStoreSpec]
         ] = []
@@ -296,10 +297,19 @@ def get_finished(self, finished_req_ids: set[str]) -> tuple[set[str], set[str]]:
         assert self.worker is not None
         finished_recving: set[str] = set()
         for transfer_result in self.worker.get_finished():
-            # we currently do not support job failures
             job_id = transfer_result.job_id
-            assert transfer_result.success
             is_load = job_id in self._load_jobs
+            if not transfer_result.success:
+                self._connector_worker_meta.mark_completed(job_id, success=False)
+                req_id = self._load_jobs.pop(job_id, None)
+                if is_load and req_id is None:
+                    # Its request already ran this step on the unloaded KV,
+                    # which failed_recving cannot undo.
+                    raise RuntimeError(f"KV offload sync load job {job_id} failed")
+                if req_id is not None:
+                    self._failed_recving.add(req_id)
+                    finished_recving.add(req_id)
+                continue
             if (
                 transfer_result.transfer_time is not None
                 and transfer_result.transfer_size is not None
@@ -320,6 +330,11 @@ def get_finished(self, finished_req_ids: set[str]) -> tuple[set[str], set[str]]:
 
         return set(), finished_recving
 
+    def take_failed_recving(self) -> set[ReqId]:
+        """Drain the requests whose load failed since the last call."""
+        failed_recving, self._failed_recving = self._failed_recving, set()
+        return failed_recving
+
     def build_connector_worker_meta(self) -> OffloadingWorkerMetadata | None:
         """Return completed transfer job IDs since the last call."""
         if not self._connector_worker_meta.completed_jobs:
@@ -331,6 +346,7 @@ def build_connector_worker_meta(self) -> OffloadingWorkerMetadata | None:
     def shutdown(self) -> None:
         self._unsubmitted_store_jobs.clear()
         self._load_jobs.clear()
+        self._failed_recving.clear()
         self._connector_worker_meta = OffloadingWorkerMetadata()
         if self.worker is not None:
             self.worker.shutdown()
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 af3e113b5f09..087d1e88d644 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
@@ -13,7 +13,10 @@
     KVConnectorRole,
     SupportsHMA,
 )
-from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorMetadata
+from vllm.distributed.kv_transfer.kv_connector.v1.base import (
+    KVConnectorMetadata,
+    KVConnectorTransferResults,
+)
 from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
     KVConnectorPromMetrics,
     KVConnectorStats,
@@ -142,6 +145,14 @@ def get_finished(self, finished_req_ids: set[str]) -> tuple[set[str], set[str]]:
         assert isinstance(self._connector_metadata, OffloadingConnectorMetadata)
         return self.connector_worker.get_finished(finished_req_ids)
 
+    def get_transfer_results(
+        self, finished_req_ids: set[str]
+    ) -> KVConnectorTransferResults:
+        results = super().get_transfer_results(finished_req_ids)
+        assert self.connector_worker is not None
+        results.failed_recving = self.connector_worker.take_failed_recving()
+        return results
+
     def build_connector_worker_meta(self) -> OffloadingWorkerMetadata | None:
         if self.connector_worker is not None:
             return self.connector_worker.build_connector_worker_meta()
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 04ce6a0b92c5..462371eddd56 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
@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Mamba conv-state sub-projection decomposition for NIXL transfer.
+"""Mamba conv-state sub-projection decomposition for KV transfer.
 
 With DS conv state layout (dim, state_len), sub-projections are
 contiguous in memory.  Each D rank reads its slices via separate
@@ -27,7 +27,7 @@
 class MambaConvSplitInfo:
     """Per-rank byte sizes of the conv sub-projections.
 
-    Used by both P and D sides for NIXL descriptor registration.
+    Used to build TP-aware transfer descriptors.
     All fields are LOCAL to this engine's TP (already divided by TP size).
 
     DS memory layout within one page (contiguous):
@@ -107,6 +107,8 @@ def remote_conv_offsets(
 def derive_mamba_conv_split(
     mamba_spec: MambaSpec,
     local_tp: int,
+    *,
+    require_dim_first: bool = True,
 ) -> MambaConvSplitInfo:
     """Derive per-rank sub-projection byte sizes from a MambaSpec.
 
@@ -115,9 +117,10 @@ def derive_mamba_conv_split(
 
     Args:
         mamba_spec: MambaSpec whose shapes are:
-            shapes[0] = conv state: (conv_dim_local, conv_rows) in DS layout.
+            shapes[0] = conv state in the configured DS or SD layout.
             shapes[1] = temporal state (model-specific shape).
         local_tp: this engine's tensor-parallel size.
+        require_dim_first: Require the DS layout used by NIXL 3-read transfer.
 
     Returns:
         MambaConvSplitInfo with per-rank sub-projection dims, conv_rows,
@@ -138,11 +141,13 @@ def derive_mamba_conv_split(
     conv_shape = mamba_spec.shapes[0]
     assert len(conv_shape) == 2, f"Expected 2D conv state shape, got {conv_shape}"
 
-    # NOTE (ZhanqiuHu): 3-read requires DS layout, which is already asserted
-    # in nixl worker __init__.  Use it directly instead of heuristic detection.
-    assert is_conv_state_dim_first(), "3-read requires DS conv state layout"
-    local_conv_dim = conv_shape[0]  # DS: (conv_dim_local, conv_rows)
-    conv_rows = conv_shape[1]
+    dim_first = is_conv_state_dim_first()
+    if require_dim_first:
+        assert dim_first, "3-read requires DS conv state layout"
+    if dim_first:
+        local_conv_dim, conv_rows = conv_shape
+    else:
+        conv_rows, local_conv_dim = conv_shape
 
     conv_dtype_size = torch.tensor(
         [],
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/umbp/runtime/__init__.py b/vllm/distributed/kv_transfer/kv_connector/v1/umbp/runtime/__init__.py
index 7b2ce33ebf0d..24ac28628bf3 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/umbp/runtime/__init__.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/umbp/runtime/__init__.py
@@ -2,12 +2,16 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 from .base import IUMBPRuntime, UMBPSchedulerHandle, UMBPWorkerHandle
+from .embedded import EmbeddedRuntime
 from .factory import UMBPRuntimeConfig, UMBPRuntimeFactory
 
+UMBPRuntimeFactory.register("embedded", EmbeddedRuntime.from_config)
+
 __all__ = [
     "IUMBPRuntime",
     "UMBPSchedulerHandle",
     "UMBPWorkerHandle",
     "UMBPRuntimeConfig",
     "UMBPRuntimeFactory",
+    "EmbeddedRuntime",
 ]
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/umbp/runtime/embedded.py b/vllm/distributed/kv_transfer/kv_connector/v1/umbp/runtime/embedded.py
new file mode 100644
index 000000000000..818e9ab37cef
--- /dev/null
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/umbp/runtime/embedded.py
@@ -0,0 +1,880 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""MORI-backed embedded DRAM offloading."""
+
+from __future__ import annotations
+
+import contextlib
+import hashlib
+import os
+import socket
+import threading
+import time
+from collections.abc import Iterable, Sequence
+from concurrent.futures import Future, ThreadPoolExecutor, TimeoutError
+from functools import lru_cache
+from pathlib import Path
+from typing import Any, NamedTuple, Protocol
+
+import msgspec
+import numpy as np
+import regex as re
+import torch
+
+from vllm.logger import init_logger
+
+from ..data import (
+    BlockLoadBatch,
+    BlockStoreBatch,
+    BlockTransferPlan,
+    KVLayoutDescriptor,
+    RankTopology,
+    TransferJobState,
+    TransferJobStatus,
+)
+from .base import IUMBPRuntime, UMBPSchedulerHandle, UMBPWorkerHandle
+from .factory import UMBPRuntimeConfig
+
+logger = init_logger(__name__)
+
+# Bounds one scheduler-to-worker lookup round trip on the local socket.
+_LOOKUP_TIMEOUT_S = 1.0
+
+# MORI evicts asynchronously once the pool crosses its high watermark, so a
+# put into a full pool can fail until eviction catches up.
+_STORE_RETRIES = 2
+_STORE_RETRY_BACKOFF_S = 0.01
+
+_RANK_KEY_PATTERN = re.compile(r":tp(\d+):pcp(\d+):dcp(\d+):pp(\d+):g\d+:")
+
+
+@lru_cache(maxsize=256)
+def _rank_namespace_from_key_prefix(prefix: str) -> tuple[int, int, int, int] | None:
+    match = _RANK_KEY_PATTERN.search(prefix)
+    if match is None:
+        return None
+    tp_rank, pcp_rank, dcp_rank, pp_rank = map(int, match.groups())
+    return tp_rank, pcp_rank, dcp_rank, pp_rank
+
+
+class _LookupClient(Protocol):
+    def batch_exists(self, keys: Sequence[str]) -> Sequence[bool]: ...
+
+    def clear(self) -> bool: ...
+
+
+class _MoriClient(_LookupClient, Protocol):
+    def register_memory(self, *args: Any) -> bool: ...
+
+    def deregister_memory(self, *args: Any) -> bool: ...
+
+    def batch_get_ranges_into_ptr(self, *args: Any) -> Sequence[bool]: ...
+
+    def batch_put_ranges_from_ptr(self, *args: Any) -> list[bool]: ...
+
+    def flush(self) -> bool: ...
+
+
+def _lookup_socket_path(
+    namespace: str,
+    rank_namespace: tuple[int, int, int, int],
+    lookup_dir: str,
+    instance: str = "",
+) -> str:
+    # The namespace identifies the model and layout, not the engine: DP ranks
+    # and independent engines serving one model share it, so the engine's own
+    # identity must separate their private pools.
+    identity = f"{namespace}:{rank_namespace}"
+    if instance:
+        identity = f"{namespace}:{instance}:{rank_namespace}"
+    digest = hashlib.sha256(identity.encode()).hexdigest()[:24]
+    return str(Path(lookup_dir) / f"vllm-umbp-{digest}.sock")
+
+
+def _socket_has_listener(path: str) -> bool:
+    probe = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
+    probe.settimeout(0.5)
+    try:
+        probe.connect(path)
+    except OSError:
+        return False
+    finally:
+        probe.close()
+    return True
+
+
+class _MoriLookupServer:
+    """Small worker-local lookup bridge for the scheduler process."""
+
+    def __init__(self, path: str, client: _LookupClient) -> None:
+        self.path = path
+        self.client = client
+        self._stop = threading.Event()
+        self._thread: threading.Thread | None = None
+        self._socket: socket.socket | None = None
+        self._owns_path = False
+
+    def start(self) -> None:
+        Path(self.path).parent.mkdir(parents=True, exist_ok=True)
+        if os.path.exists(self.path) and _socket_has_listener(self.path):
+            raise RuntimeError(
+                f"UMBP lookup socket {self.path} is served by another live "
+                "engine; set lookup_instance to a distinct value for each "
+                "engine that serves this model on this host"
+            )
+        with contextlib.suppress(FileNotFoundError):
+            os.unlink(self.path)
+        self._socket = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
+        self._socket.bind(self.path)
+        self._owns_path = True
+        self._socket.listen(16)
+        self._socket.settimeout(0.2)
+        self._thread = threading.Thread(
+            target=self._serve,
+            name="umbp-embedded-lookup",
+            daemon=True,
+        )
+        self._thread.start()
+
+    def _serve(self) -> None:
+        assert self._socket is not None
+        while not self._stop.is_set():
+            try:
+                connection, _ = self._socket.accept()
+            except TimeoutError:
+                continue
+            except OSError:
+                break
+            with connection:
+                connection.settimeout(_LOOKUP_TIMEOUT_S)
+                try:
+                    request = b""
+                    while not request.endswith(b"\n"):
+                        chunk = connection.recv(65536)
+                        if not chunk:
+                            break
+                        request += chunk
+                    payload = msgspec.json.decode(request)
+                    if isinstance(payload, dict) and payload.get("op") == "clear":
+                        connection.sendall(
+                            msgspec.json.encode(bool(self.client.clear())) + b"\n"
+                        )
+                        continue
+                    result = self.client.batch_exists(payload)
+                    connection.sendall(
+                        msgspec.json.encode([bool(value) for value in result]) + b"\n"
+                    )
+                except Exception:
+                    with contextlib.suppress(OSError):
+                        connection.sendall(b"[]\n")
+
+    def close(self) -> None:
+        self._stop.set()
+        if self._socket is not None:
+            self._socket.close()
+        if self._thread is not None and self._thread.ident is not None:
+            self._thread.join(timeout=2)
+        if self._owns_path:
+            with contextlib.suppress(FileNotFoundError):
+                os.unlink(self.path)
+            self._owns_path = False
+
+
+class _MoriSchedulerHandle(UMBPSchedulerHandle):
+    def __init__(
+        self,
+        namespace: str,
+        topology: RankTopology,
+        lookup_dir: str,
+        instance: str = "",
+    ) -> None:
+        self._paths = {
+            rank: _lookup_socket_path(namespace, rank, lookup_dir, instance)
+            for rank in topology.all_namespaces()
+        }
+        self.last_lookup_diagnostics: dict[str, tuple[Any, ...]] = {}
+
+    @staticmethod
+    def _request(path: str, payload: Any, timeout: float = _LOOKUP_TIMEOUT_S) -> Any:
+        with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as sock:
+            sock.settimeout(timeout)
+            sock.connect(path)
+            sock.sendall(msgspec.json.encode(payload) + b"\n")
+            response = b""
+            while not response.endswith(b"\n"):
+                chunk = sock.recv(65536)
+                if not chunk:
+                    break
+                response += chunk
+            return msgspec.json.decode(response or b"[]")
+
+    def lookup(self, keys: Sequence[str]) -> Sequence[bool]:
+        unavailable: list[tuple[int, int, int, int]]
+        if len(self._paths) == 1:
+            rank, path = next(iter(self._paths.items()))
+            try:
+                values = self._request(path, keys)
+                if len(values) != len(keys):
+                    raise ValueError("lookup returned an invalid result length")
+                return self._lookup_result(keys, [bool(value) for value in values], [])
+            except (OSError, ValueError, IndexError):
+                unavailable = (
+                    [rank]
+                    if any(
+                        _rank_namespace_from_key_prefix(key.rpartition(":")[0] + ":")
+                        == rank
+                        for key in keys
+                    )
+                    else []
+                )
+                return self._lookup_result(keys, [False] * len(keys), unavailable)
+        result = [False] * len(keys)
+        keys_by_rank: dict[tuple[int, int, int, int], list[tuple[int, str]]] = {}
+        unrouted: list[tuple[int, str]] = []
+        for index, key in enumerate(keys):
+            key_rank = _rank_namespace_from_key_prefix(key.rpartition(":")[0] + ":")
+            if key_rank is None or key_rank not in self._paths:
+                unrouted.append((index, key))
+            else:
+                keys_by_rank.setdefault(key_rank, []).append((index, key))
+        unavailable = []
+        for rank, indexed_keys in keys_by_rank.items():
+            path = self._paths[rank]
+            try:
+                values = self._request(path, [key for _, key in indexed_keys])
+                if len(values) != len(indexed_keys):
+                    raise ValueError("lookup returned an invalid result length")
+                for (index, _), value in zip(indexed_keys, values, strict=True):
+                    result[index] = bool(value)
+            except (OSError, ValueError, IndexError):
+                unavailable.append(rank)
+        for path in self._paths.values():
+            if not unrouted:
+                break
+            try:
+                values = self._request(path, [key for _, key in unrouted])
+                for (index, _), value in zip(unrouted, values, strict=True):
+                    result[index] = result[index] or bool(value)
+            except (OSError, ValueError):
+                continue
+        return self._lookup_result(keys, result, unavailable)
+
+    def _lookup_result(
+        self,
+        keys: Sequence[str],
+        result: list[bool],
+        unavailable: list[tuple[int, int, int, int]],
+    ) -> list[bool]:
+        missing = tuple(
+            key for key, exists in zip(keys, result, strict=True) if not exists
+        )
+        self.last_lookup_diagnostics = {
+            "unavailable_ranks": tuple(unavailable),
+            "missing_keys": missing,
+        }
+        if unavailable:
+            logger.warning("UMBP lookup unavailable for ranks: %s", unavailable)
+        elif missing:
+            logger.debug("UMBP lookup missing %d objects", len(missing))
+        return result
+
+    def close(self) -> None:
+        return
+
+    def clear(self) -> bool:
+        success = True
+        contacted = False
+        for path in self._paths.values():
+            try:
+                cleared = self._request(path, {"op": "clear"}, timeout=2.0)
+                contacted = True
+                success = success and bool(cleared)
+            except (OSError, ValueError):
+                success = False
+        return contacted and success
+
+
+class _BlockLoadLayout(NamedTuple):
+    bases: np.ndarray
+    strides: np.ndarray
+    sizes: list[int]
+    offsets: list[int]
+    num_blocks: int
+    num_bytes: int
+
+
+_LoadArgs = tuple[list[str], list[list[int]], list[list[int]], list[list[int]]]
+_StoreArgs = tuple[
+    list[str], list[int], list[list[int]], list[list[int]], list[list[int]]
+]
+
+
+class _MoriWorkerHandle(UMBPWorkerHandle):
+    def __init__(
+        self,
+        client: _MoriClient,
+        namespace: str,
+        topology: RankTopology,
+        lookup_dir: str,
+        max_workers: int,
+        timeout_s: float,
+        layout: KVLayoutDescriptor | None = None,
+        lookup_instance: str = "",
+    ) -> None:
+        self.client = client
+        self._namespace = namespace
+        self._topology = topology
+        self._lookup_dir = lookup_dir
+        self._lookup_instance = lookup_instance
+        self._lookup_server: _MoriLookupServer | None = None
+        self._registered_storages: set[int] = set()
+        self._gpu_devices: set[int] = set()
+        self._executor = ThreadPoolExecutor(
+            max_workers=max_workers,
+            thread_name_prefix="umbp-embedded-transfer",
+        )
+        self._job_lock = threading.Lock()
+        self._futures: dict[int, Future[TransferJobState]] = {}
+        self._store_deadlines: dict[int, float] = {}
+        self._load_deadlines: dict[int, float] = {}
+        self._timeout_s = timeout_s
+        self._layout = layout
+        self._load_layouts: dict[int, _BlockLoadLayout] = {}
+
+    def register_buffers(self, kv_caches: dict[str, torch.Tensor]) -> None:
+        try:
+            from mori.cpp import MemoryLocationType
+        except ImportError as exc:
+            raise RuntimeError("MORI UMBP Python bindings are unavailable") from exc
+
+        for cache in kv_caches.values():
+            storage = cache.untyped_storage()
+            storage_ptr = storage.data_ptr()
+            if storage_ptr in self._registered_storages:
+                continue
+            location = (
+                MemoryLocationType.GPU
+                if cache.device.type == "cuda"
+                else MemoryLocationType.CPU
+            )
+            device = cache.device.index if cache.device.index is not None else -1
+            if not self.client.register_memory(
+                storage_ptr, storage.nbytes(), location, device
+            ):
+                raise RuntimeError(
+                    f"MORI UMBP failed to register KV storage 0x{storage_ptr:x}"
+                )
+            self._registered_storages.add(storage_ptr)
+            if location == MemoryLocationType.GPU and device >= 0:
+                self._gpu_devices.add(device)
+
+        self._register_load_layouts(kv_caches)
+        if self._lookup_server is None:
+            self._lookup_server = _MoriLookupServer(
+                _lookup_socket_path(
+                    self._namespace,
+                    self._topology.local_namespace,
+                    self._lookup_dir,
+                    self._lookup_instance,
+                ),
+                self,
+            )
+            self._lookup_server.start()
+
+    def _register_load_layouts(self, kv_caches: dict[str, torch.Tensor]) -> None:
+        self._load_layouts = {}
+        if self._layout is None:
+            return
+        for group_id in {r.group_id for r in self._layout.regions}:
+            bases: list[int] = []
+            strides: list[int] = []
+            sizes: list[int] = []
+            offsets: list[int] = []
+            counts: list[int] = []
+            for region in self._layout.regions:
+                if region.group_id != group_id:
+                    continue
+                cache = kv_caches.get(region.layer_name)
+                if cache is None or cache.ndim == 0:
+                    break
+                stride = cache.stride(0) * cache.element_size()
+                if (
+                    stride <= 0
+                    or stride > region.block_stride
+                    or region.block_stride % stride
+                    or region.block_bytes > stride
+                ):
+                    break
+                bases.append(cache.data_ptr())
+                strides.append(stride)
+                offsets.append(sum(sizes))
+                sizes.append(region.block_bytes)
+                counts.append(cache.shape[0])
+            else:
+                self._load_layouts[group_id] = _BlockLoadLayout(
+                    np.asarray(bases, dtype=np.uint64),
+                    np.asarray(strides, dtype=np.uint64),
+                    sizes,
+                    offsets,
+                    min(counts),
+                    sum(sizes),
+                )
+
+    def _bulk_load_args(
+        self, plans: Sequence[BlockTransferPlan]
+    ) -> tuple[_LoadArgs, tuple[int, ...]] | None:
+        groups: dict[int, list[tuple[int, int]]] = {}
+        entries: Iterable[tuple[int | None, int]]
+        if isinstance(plans, (BlockLoadBatch, BlockStoreBatch)):
+            entries = zip(plans.group_ids, plans.block_ids, strict=True)
+            keys = plans.keys
+        else:
+            if any(p.ranges for p in plans):
+                return None
+            entries = ((p.group_id, p.block_id) for p in plans)
+            keys = [p.key for p in plans]
+        for index, (group_id, block_id) in enumerate(entries):
+            if group_id is None:
+                return None
+            layout = self._load_layouts.get(group_id)
+            if layout is None or not 0 <= block_id < layout.num_blocks:
+                return None
+            groups.setdefault(group_id, []).append((index, block_id))
+        pointers: list[list[int]] = [[]] * len(plans)
+        sizes: list[list[int]] = [[]] * len(plans)
+        offsets: list[list[int]] = [[]] * len(plans)
+        plan_bytes = [0] * len(plans)
+        for group_id, group_entries in groups.items():
+            layout = self._load_layouts[group_id]
+            ids = np.asarray([block for _, block in group_entries], dtype=np.uint64)
+            addresses = (
+                layout.bases[None, :] + ids[:, None] * layout.strides[None, :]
+            ).tolist()
+            for (index, _), row in zip(group_entries, addresses, strict=True):
+                pointers[index] = row
+                sizes[index] = layout.sizes
+                offsets[index] = layout.offsets
+                plan_bytes[index] = layout.num_bytes
+        return (
+            (keys, pointers, sizes, offsets),
+            tuple(plan_bytes),
+        )
+
+    def load_blocks(
+        self, plans: Sequence[BlockTransferPlan]
+    ) -> TransferJobState | None:
+        prepared = self._bulk_load_args(plans)
+        if prepared is None:
+            return None
+        args, plan_bytes = prepared
+        return self._submit_load(
+            plans if isinstance(plans, BlockLoadBatch) else tuple(plans),
+            args,
+            plan_bytes,
+        )
+
+    def batch_exists(self, keys: Sequence[str]) -> Sequence[bool]:
+        return [bool(value) for value in self.client.batch_exists(keys)]
+
+    def clear(self) -> bool:
+        return bool(self.client.clear())
+
+    @staticmethod
+    def _range_args(plans: Sequence[BlockTransferPlan], *, for_store: bool = True):
+        keys, object_sizes, pointers, sizes, offsets = [], [], [], [], []
+        for plan in plans:
+            plan_pointers, plan_sizes, plan_offsets = [], [], []
+            object_size = 0
+            for item in plan.ranges:
+                plan_pointers.append(item.base_address)
+                plan_sizes.append(item.length)
+                plan_offsets.append(item.object_offset)
+                if for_store:
+                    end = item.object_offset + item.length
+                    if end > object_size:
+                        object_size = end
+            keys.append(plan.key)
+            if for_store:
+                object_sizes.append(object_size)
+            pointers.append(plan_pointers)
+            sizes.append(plan_sizes)
+            offsets.append(plan_offsets)
+        return keys, object_sizes, pointers, sizes, offsets
+
+    def load(self, plans: Sequence[BlockTransferPlan]) -> TransferJobState:
+        return self._submit_load(tuple(plans))
+
+    def _submit_load(
+        self,
+        plans: tuple[BlockTransferPlan, ...] | BlockLoadBatch,
+        range_args: _LoadArgs | None = None,
+        plan_bytes: tuple[int, ...] | None = None,
+    ) -> TransferJobState:
+        job = TransferJobState(plans, plan_bytes=plan_bytes)
+        job.start()
+        future = self._executor.submit(self._load_sync, job, plans, range_args)
+        with self._job_lock:
+            self._futures[id(job)] = future
+        return job
+
+    def _load_sync(
+        self,
+        job: TransferJobState,
+        plans: tuple[BlockTransferPlan, ...] | BlockLoadBatch,
+        range_args: _LoadArgs | None = None,
+    ) -> TransferJobState:
+        with self._job_lock:
+            self._load_deadlines[id(job)] = time.monotonic() + self._timeout_s
+        if not plans:
+            job.complete()
+            return job
+        if range_args is None:
+            keys, _, pointers, sizes, offsets = self._range_args(plans, for_store=False)
+        else:
+            keys, pointers, sizes, offsets = range_args
+        request_id = (
+            plans.request_id
+            if isinstance(plans, BlockLoadBatch)
+            else plans[0].request_id
+        )
+        started_at = time.monotonic()
+        logger.debug(
+            "MORI UMBP range load started request=%s plans=%d ranges=%d bytes=%d",
+            request_id,
+            len(plans),
+            sum(len(plan_sizes) for plan_sizes in sizes),
+            sum(sum(plan_sizes) for plan_sizes in sizes),
+        )
+        results = self.client.batch_get_ranges_into_ptr(keys, pointers, sizes, offsets)
+        logger.debug(
+            "MORI UMBP range load returned request=%s plans=%d elapsed=%.3fs",
+            request_id,
+            len(plans),
+            time.monotonic() - started_at,
+        )
+        completed = [key for key, ok in zip(keys, results, strict=True) if ok]
+        failed = [key for key, ok in zip(keys, results, strict=True) if not ok]
+        if completed:
+            job.complete(completed)
+        if failed:
+            job.fail(failed, "MORI UMBP range load failed")
+        return job
+
+    def store_blocks(
+        self, plans: Sequence[BlockTransferPlan]
+    ) -> TransferJobState | None:
+        """Store whole blocks from registered layouts; None requests ranges."""
+        prepared = self._bulk_load_args(plans)
+        if prepared is None:
+            return None
+        (keys, pointers, sizes, offsets), plan_bytes = prepared
+        return self._submit_store(
+            plans if isinstance(plans, BlockStoreBatch) else tuple(plans),
+            (keys, list(plan_bytes), pointers, sizes, offsets),
+            plan_bytes,
+        )
+
+    def store(self, plans: Sequence[BlockTransferPlan]) -> TransferJobState:
+        return self._submit_store(tuple(plans))
+
+    def _submit_store(
+        self,
+        plans: tuple[BlockTransferPlan, ...] | BlockStoreBatch,
+        range_args: _StoreArgs | None = None,
+        plan_bytes: tuple[int, ...] | None = None,
+    ) -> TransferJobState:
+        job = TransferJobState(plans, plan_bytes=plan_bytes)
+        job.start()
+        ready_events: list[torch.Event] = []
+        for device in self._gpu_devices:
+            with torch.device(f"cuda:{device}"):
+                event = torch.Event()
+                event.record(torch.accelerator.current_stream())
+                ready_events.append(event)
+        future = self._executor.submit(
+            self._store_sync, job, plans, ready_events, range_args
+        )
+        with self._job_lock:
+            self._futures[id(job)] = future
+        return job
+
+    def _store_sync(
+        self,
+        job: TransferJobState,
+        plans: tuple[BlockTransferPlan, ...] | BlockStoreBatch,
+        ready_events: list[torch.Event],
+        range_args: _StoreArgs | None = None,
+    ) -> TransferJobState:
+        # Timed from here: a store queued behind others has not stalled.
+        with self._job_lock:
+            self._store_deadlines[id(job)] = time.monotonic() + self._timeout_s
+        for event in ready_events:
+            event.synchronize()
+        if not plans:
+            job.complete()
+            return job
+        keys, object_sizes, pointers, sizes, offsets = (
+            self._range_args(plans) if range_args is None else range_args
+        )
+        results = self.client.batch_put_ranges_from_ptr(
+            keys, object_sizes, pointers, sizes, offsets
+        )
+        failed_indices = [index for index, ok in enumerate(results) if not ok]
+        for attempt in range(_STORE_RETRIES):
+            if not failed_indices:
+                break
+            self.client.flush()
+            time.sleep(_STORE_RETRY_BACKOFF_S * (attempt + 1))
+            for index in failed_indices.copy():
+                retry = self.client.batch_put_ranges_from_ptr(
+                    [keys[index]],
+                    [object_sizes[index]],
+                    [pointers[index]],
+                    [sizes[index]],
+                    [offsets[index]],
+                )
+                if retry and retry[0]:
+                    results[index] = True
+                    failed_indices.remove(index)
+            if failed_indices:
+                logger.warning(
+                    "MORI UMBP store retry %d left %d/%d objects pending",
+                    attempt + 1,
+                    len(failed_indices),
+                    len(plans),
+                )
+        if not self.client.flush():
+            job.fail([plan.key for plan in plans], "MORI UMBP flush failed")
+            return job
+        completed = [plan.key for plan, ok in zip(plans, results, strict=True) if ok]
+        failed = [plan.key for plan, ok in zip(plans, results, strict=True) if not ok]
+        if completed:
+            job.complete(completed)
+        if failed:
+            job.fail(failed, "MORI UMBP range store failed")
+        return job
+
+    def wait(self, job: TransferJobState) -> TransferJobState:
+        with self._job_lock:
+            future = self._futures.get(id(job))
+        if future is None:
+            return job
+        try:
+            return future.result(timeout=self._timeout_s)
+        except TimeoutError as exc:
+            if not future.done():
+                # The transfer still owns GPU pointers. Abort this engine step
+                # rather than reporting completion and allowing block reuse.
+                raise TimeoutError(
+                    "MORI UMBP transfer timed out with buffers still in use"
+                ) from exc
+            job.fail([plan.key for plan in job.plans], str(exc))
+            return job
+        except Exception as exc:
+            job.fail([plan.key for plan in job.plans], str(exc))
+            return job
+        finally:
+            if future.done():
+                with self._job_lock:
+                    self._futures.pop(id(job), None)
+                    self._store_deadlines.pop(id(job), None)
+                    self._load_deadlines.pop(id(job), None)
+
+    def poll(self, job: TransferJobState) -> TransferJobState | None:
+        with self._job_lock:
+            future = self._futures.get(id(job))
+        if future is None:
+            if job.status in (
+                TransferJobStatus.COMPLETED,
+                TransferJobStatus.FAILED,
+                TransferJobStatus.CANCELLED,
+            ):
+                return job
+            return None
+        if not future.done():
+            with self._job_lock:
+                deadline = self._store_deadlines.get(
+                    id(job), self._load_deadlines.get(id(job))
+                )
+            if deadline is not None and time.monotonic() >= deadline:
+                raise TimeoutError(
+                    "MORI UMBP transfer timed out with buffers still in use"
+                )
+            return None
+        with self._job_lock:
+            self._futures.pop(id(job), None)
+            self._store_deadlines.pop(id(job), None)
+            self._load_deadlines.pop(id(job), None)
+        try:
+            return future.result()
+        except Exception as exc:
+            job.fail([plan.key for plan in job.plans], str(exc))
+            return job
+
+    def publish(self, job: TransferJobState) -> None:
+        if job.status.value != "completed":
+            raise RuntimeError("cannot publish an incomplete MORI UMBP job")
+
+    def cancel(self, job: TransferJobState) -> TransferJobState:
+        with self._job_lock:
+            future = self._futures.get(id(job))
+        if future is None:
+            return job
+        if future.cancel():
+            with self._job_lock:
+                self._futures.pop(id(job), None)
+                self._store_deadlines.pop(id(job), None)
+                self._load_deadlines.pop(id(job), None)
+            job.cancel("preempted")
+            return job
+        return self.wait(job)
+
+    def take_evicted_keys(self) -> Sequence[str]:
+        return ()
+
+    def close(self) -> None:
+        self._executor.shutdown(wait=True, cancel_futures=True)
+        if self._lookup_server is not None:
+            self._lookup_server.close()
+        for storage_ptr in self._registered_storages:
+            self.client.deregister_memory(storage_ptr)
+        self._registered_storages.clear()
+        close = getattr(self.client, "close", None)
+        if callable(close):
+            close()
+        else:
+            self.client.flush()
+
+
+def _configure_dram(client_config: Any, options: dict[str, Any]) -> None:
+    """Apply embedded-only host-DRAM options to MORI's client config."""
+    dram = client_config.dram
+    dram.capacity_bytes = options.get("capacity_bytes", 64 * 1024**3)
+    option_fields = {
+        "dram_use_shared_memory": "use_shared_memory",
+        "dram_shm_name": "shm_name",
+        "dram_high_watermark": "high_watermark",
+        "dram_low_watermark": "low_watermark",
+        "dram_use_hugepages": "use_hugepages",
+        "dram_hugepage_size": "hugepage_size",
+        "dram_numa_node": "numa_node",
+        "dram_prefault": "prefault",
+    }
+    for option, field in option_fields.items():
+        if option in options:
+            setattr(dram, field, options[option])
+
+
+def _resolve_embedded_options(config: UMBPRuntimeConfig) -> dict[str, Any]:
+    options = dict(config.options)
+    if "capacity_bytes" in options and "total_capacity_bytes" in options:
+        raise ValueError(
+            "capacity_bytes and total_capacity_bytes are mutually exclusive"
+        )
+    for name in ("capacity_bytes", "total_capacity_bytes", "dram_hugepage_size"):
+        if name in options and (type(options[name]) is not int or options[name] <= 0):
+            raise ValueError(f"{name} must be a positive integer")
+    for name in ("dram_use_shared_memory", "dram_use_hugepages", "dram_prefault"):
+        if name in options and type(options[name]) is not bool:
+            raise ValueError(f"{name} must be a boolean")
+    for name in ("dram_high_watermark", "dram_low_watermark"):
+        if name in options and (
+            type(options[name]) not in (int, float) or not 0 < options[name] <= 1
+        ):
+            raise ValueError(f"{name} must be in (0, 1]")
+    low, high = options.get("dram_low_watermark"), options.get("dram_high_watermark")
+    if low is not None and high is not None and low > high:
+        raise ValueError("dram_low_watermark must not exceed dram_high_watermark")
+    if "dram_numa_node" in options and (
+        type(options["dram_numa_node"]) is not int or options["dram_numa_node"] < -1
+    ):
+        raise ValueError("dram_numa_node must be an integer >= -1")
+    if "dram_shm_name" in options and (
+        not isinstance(options["dram_shm_name"], str) or not options["dram_shm_name"]
+    ):
+        raise ValueError("dram_shm_name must be a non-empty string")
+    if {
+        "master_address",
+        "node_address",
+        "io_engine_host",
+        "peer_service_port",
+    } & options.keys():
+        raise ValueError("embedded UMBP cannot configure distributed-only options")
+    total = options.pop("total_capacity_bytes", None)
+    if total is not None:
+        if total < config.rank_count:
+            raise ValueError(
+                "total_capacity_bytes must provide at least one byte per rank"
+            )
+        options["capacity_bytes"] = total // config.rank_count
+        options["_configured_total_capacity_bytes"] = total
+    return options
+
+
+class EmbeddedRuntime(IUMBPRuntime):
+    """MORI-backed embedded DRAM runtime."""
+
+    def __init__(self, config: UMBPRuntimeConfig) -> None:
+        if config.options.get("backend") == "memory":
+            raise ValueError(
+                "The memory backend is test-only; use MORI for embedded mode"
+            )
+        self.options = _resolve_embedded_options(config)
+        self.lookup_dir = self.options.get("lookup_dir", "/tmp")
+        instance = self.options.get("lookup_instance", "")
+        if not isinstance(instance, str):
+            raise ValueError("lookup_instance must be a string")
+        self.lookup_instance = f"{instance}:dp{config.dp_index}"
+
+    def create_scheduler_handle(
+        self,
+        namespace: str,
+        topology: RankTopology,
+        layout: KVLayoutDescriptor,
+    ) -> UMBPSchedulerHandle:
+        del layout
+        return _MoriSchedulerHandle(
+            namespace, topology, self.lookup_dir, self.lookup_instance
+        )
+
+    def create_worker_handle(
+        self,
+        namespace: str,
+        topology: RankTopology,
+        layout: KVLayoutDescriptor,
+    ) -> UMBPWorkerHandle:
+        try:
+            from mori.cpp import UMBPClient, UMBPConfig
+        except ImportError as exc:
+            raise RuntimeError(
+                "Embedded UMBP requires MORI built with BUILD_UMBP=ON"
+            ) from exc
+
+        client_config = UMBPConfig()
+        _configure_dram(client_config, self.options)
+        total = self.options.get("_configured_total_capacity_bytes")
+        if total is None:
+            logger.info(
+                "Embedded UMBP DRAM capacity: %.2f GiB per rank",
+                client_config.dram.capacity_bytes / 1024**3,
+            )
+        else:
+            logger.info(
+                "Embedded UMBP DRAM capacity: %.2f GiB total, %.2f GiB per rank",
+                total / 1024**3,
+                client_config.dram.capacity_bytes / 1024**3,
+            )
+        return _MoriWorkerHandle(
+            UMBPClient(client_config),
+            namespace,
+            topology,
+            self.lookup_dir,
+            int(self.options.get("num_workers", 4)),
+            float(self.options.get("timeout_ms", 30000)) / 1000,
+            layout,
+            lookup_instance=self.lookup_instance,
+        )
+
+    @classmethod
+    def from_config(cls, config: UMBPRuntimeConfig) -> IUMBPRuntime:
+        return cls(config)
diff --git a/vllm/distributed/parallel_state.py b/vllm/distributed/parallel_state.py
index 897518247f13..dc4d23b22192 100644
--- a/vllm/distributed/parallel_state.py
+++ b/vllm/distributed/parallel_state.py
@@ -527,6 +527,10 @@ def __init__(
         self.cpu_group = self_cpu_group
         self.device_group = self_device_group
 
+        # Dead DP slots under fault-tolerance scale-down (empty when FT is
+        # off); read by the DP allreduce to neutralize dead columns.
+        self.dead_dp_ranks: set[int] = set()
+
         from vllm.platforms import current_platform
 
         if current_platform.is_cuda_alike():
@@ -1692,7 +1696,7 @@ def _aiter_registers_graph_buffers(group: GroupCoordinator) -> bool:
 
 
 @contextmanager
-def graph_capture(device: torch.device):
+def graph_capture(device: torch.device, stream: torch.cuda.Stream | None = None):
     """`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
@@ -1704,12 +1708,14 @@ def graph_capture(device: torch.device):
     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 given ``stream`` is reused instead, so graphs captured on it share the
+    caching allocator's free blocks of the graph pool.
     """
     from vllm.device_allocator.alloc_conf import non_expandable_allocations
     from vllm.platforms import current_platform
 
     tp, pp, dp = get_tp_group(), get_pp_group(), get_dp_group()
-    context = GraphCaptureContext(torch.cuda.Stream(device=device))
+    context = GraphCaptureContext(stream or torch.cuda.Stream(device=device))
     # Let AITER all-reduce register captured buffers instead of copying them.
     # Not with DP AITER collectives: they reduce workspaces allocated earlier.
     non_expandable = (
diff --git a/vllm/distributed/stateless_coordinator.py b/vllm/distributed/stateless_coordinator.py
index 53d82efc2a94..c35eb4b4d9b4 100644
--- a/vllm/distributed/stateless_coordinator.py
+++ b/vllm/distributed/stateless_coordinator.py
@@ -1,7 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import socket
-import struct
 from typing import Any, Optional
 
 import torch
@@ -17,6 +15,8 @@
 )
 from vllm.distributed.utils import (
     StatelessProcessGroup,
+    allocate_group_ports,
+    fetch_group_ports,
     get_cpu_distributed_timeout_or_none,
     get_distributed_timeout_or_none,
     stateless_destroy_torch_distributed_process_group,
@@ -27,37 +27,8 @@
 
 logger = init_logger(__name__)
 
-_PORTS_FMT = "!3I"
-
-
-def _allocate_group_ports(
-    key: str,
-    host: str,
-    coord_store: Store,
-) -> tuple[list[int], list[socket.socket]]:
-    """Bind 3 sockets and publish the ports to *coord_store*.
-
-    Called by rank 0 only.  Returns ``(ports, sockets)`` with the
-    sockets still open.
-    """
-    socks: list[socket.socket] = []
-    ports: list[int] = []
-    for _ in range(3):
-        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-        s.bind((host, 0))
-        s.listen()
-        socks.append(s)
-        ports.append(s.getsockname()[1])
-    coord_store.set(key, struct.pack(_PORTS_FMT, *ports))
-    return ports, socks
-
-
-def _fetch_group_ports(key: str, coord_store: Store) -> list[int]:
-    """Read 3 ports published by rank 0 from *coord_store*.
-
-    Blocks until the key is available.
-    """
-    return list(struct.unpack(_PORTS_FMT, coord_store.get(key)))
+# device, cpu and tcp-store groups each get one port
+_NUM_GROUP_PORTS = 3
 
 
 class StatelessGroupCoordinator(GroupCoordinator):
@@ -118,13 +89,11 @@ def __init__(
 
                 key = f"{group_name}_{idx}"
                 if self.rank_in_group == 0:
-                    ports, socks = _allocate_group_ports(
-                        key,
-                        host,
-                        coord_store,
+                    ports, socks = allocate_group_ports(
+                        coord_store, key, host, _NUM_GROUP_PORTS
                     )
                 else:
-                    ports = _fetch_group_ports(key, coord_store)
+                    ports = fetch_group_ports(coord_store, key, _NUM_GROUP_PORTS)
                     socks = []
                 device_port, cpu_port, tcp_store_port = ports
 
diff --git a/vllm/distributed/utils.py b/vllm/distributed/utils.py
index 796da751da13..6afdc5078668 100644
--- a/vllm/distributed/utils.py
+++ b/vllm/distributed/utils.py
@@ -10,6 +10,7 @@
 import os
 import pickle
 import socket
+import struct
 import sys
 import time
 import uuid
@@ -19,7 +20,7 @@
 from typing import Any
 
 import torch
-from torch.distributed import ProcessGroup, Store, TCPStore
+from torch.distributed import ProcessGroup, ReduceOp, Store, TCPStore
 from torch.distributed.distributed_c10d import (
     Backend,
     PrefixStore,
@@ -59,6 +60,18 @@ def ensure_divisibility(numerator, denominator):
 _CPU_ALL_REDUCE_TAG_BASE = 27100
 
 
+def all_reduce_int(
+    group: ProcessGroup | None, value: int, op: ReduceOp = ReduceOp.MIN
+) -> int:
+    """All-reduce one integer over a CPU process group, e.g. for ranks to agree
+    on a size or a decision; `value` itself without a group or with one rank."""
+    if group is None or group.size() == 1:
+        return value
+    tensor = torch.tensor([value], dtype=torch.int64, device="cpu")
+    torch.distributed.all_reduce(tensor, op=op, group=group)
+    return int(tensor.item())
+
+
 def cpu_all_reduce_(
     tensor: torch.Tensor,
     group: ProcessGroup,
@@ -592,13 +605,59 @@ def get_cached_tcp_store_client(host: str, port: int) -> TCPStore:
     return TCPStore(host, port, is_master=False, wait_for_workers=False)
 
 
+def allocate_group_ports(
+    store: Store, key: str, host: str, count: int
+) -> tuple[list[int], list[socket.socket]]:
+    """Bind sockets and publish the ports to *store*.
+
+    Returns ``(ports, sockets)`` with the sockets still open.
+    """
+    socks: list[socket.socket] = []
+    ports: list[int] = []
+    for _ in range(count):
+        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        s.bind((host, 0))
+        s.listen()
+        socks.append(s)
+        ports.append(s.getsockname()[1])
+    store.set(key, struct.pack(f"!{count}I", *ports))
+    return ports, socks
+
+
+def fetch_group_ports(store: Store, key: str, count: int) -> list[int]:
+    """Read ports published under *key* by func:`allocate_group_ports`. Blocks until
+    the key is available.
+    """
+    return list(struct.unpack(f"!{count}I", store.get(key)))
+
+
+# Per-process flag: startup (weight load, KV setup, graph capture) is done and
+# the busy loop is about to start. Under FT, gloo groups created before this
+# point are given no cpu timeout, because uneven rank startup would trip it;
+# groups created later (recovery reinit) get the configured timeout.
+_steady_state_entered = False
+
+
+def enter_steady_state() -> None:
+    global _steady_state_entered
+    _steady_state_entered = True
+
+
+def steady_state_entered() -> bool:
+    return _steady_state_entered
+
+
 def get_cpu_distributed_timeout_or_none() -> timedelta | None:
     from vllm.config import get_current_vllm_config_or_none
 
     vllm_config = get_current_vllm_config_or_none()
     if vllm_config is None:
         return None
-    return vllm_config.parallel_config.cpu_distributed_timeout
+    return vllm_config.parallel_config.effective_cpu_distributed_timeout
+
+
+def set_gloo_backend_timeout(group: ProcessGroup, timeout: timedelta) -> None:
+    group._get_backend(torch.device("cpu"))._set_default_timeout(timeout)
 
 
 def get_distributed_timeout_or_none() -> timedelta | None:
diff --git a/vllm/distributed/weight_transfer/ipc_engine.py b/vllm/distributed/weight_transfer/ipc_engine.py
index d841ecf6f312..565c5ef46b87 100644
--- a/vllm/distributed/weight_transfer/ipc_engine.py
+++ b/vllm/distributed/weight_transfer/ipc_engine.py
@@ -158,20 +158,16 @@ def init_transfer_engine(self, init_info: IPCWeightTransferInitInfo) -> None:
         self.packed = init_info.packed
 
     def start_weight_update(self) -> None:
-        """Initialize layerwise reloading for the incoming checkpoint weights."""
-        from vllm.model_executor.model_loader.reload import (
-            initialize_layerwise_reload,
-        )
+        """Begin reloading the incoming checkpoint weights."""
+        from vllm.model_executor.model_loader.reload import start_reload
 
-        initialize_layerwise_reload(self.model)
+        start_reload(self.model)
 
     def finish_weight_update(self) -> None:
-        """Finalize layerwise reloading after all weights have been received."""
-        from vllm.model_executor.model_loader.reload import (
-            finalize_layerwise_reload,
-        )
+        """Finish reloading after all weights have been received."""
+        from vllm.model_executor.model_loader.reload import finish_reload
 
-        finalize_layerwise_reload(self.model, self.model_config)
+        finish_reload(self.model, self.model_config, self.device)
         # Every reduce_tensor call is a fresh export with its own refcount
         # slot, so releasing once per update always balances this update's
         # export and lets the trainer reclaim its staging buffer. Callers
diff --git a/vllm/distributed/weight_transfer/m2n_engine.py b/vllm/distributed/weight_transfer/m2n_engine.py
index 66dbb6ba2457..07837742c62f 100644
--- a/vllm/distributed/weight_transfer/m2n_engine.py
+++ b/vllm/distributed/weight_transfer/m2n_engine.py
@@ -441,7 +441,7 @@ def _prepare_destination_plan(self, init_info: M2NWeightTransferInitInfo) -> Non
         )
 
     def start_weight_update(self) -> None:
-        """Set up layerwise reloading, but only if some parameter needs it.
+        """Begin reloading, but only if some parameter needs it.
 
         Directly-resharded parameters are written in place and never go through
         `load_weights`, so a plan with no fallback entries has nothing to
@@ -449,19 +449,17 @@ def start_weight_update(self) -> None:
         """
         if not self._uses_load_weights:
             return
-        from vllm.model_executor.model_loader.reload import (
-            initialize_layerwise_reload,
-        )
+        from vllm.model_executor.model_loader.reload import start_reload
 
-        initialize_layerwise_reload(self.model)
+        start_reload(self.model)
 
     def finish_weight_update(self) -> None:
-        """Finalize layerwise reloading when the plan uses fallback entries."""
+        """Finish reloading when the plan uses fallback entries."""
         if not self._uses_load_weights:
             return
-        from vllm.model_executor.model_loader.reload import finalize_layerwise_reload
+        from vllm.model_executor.model_loader.reload import finish_reload
 
-        finalize_layerwise_reload(self.model, self.model_config)
+        finish_reload(self.model, self.model_config, self.device)
 
     def receive_weights(self, update_info: M2NWeightTransferUpdateInfo) -> None:
         """Receive each requested parameter using its initialization-time plan."""
diff --git a/vllm/distributed/weight_transfer/nccl_common.py b/vllm/distributed/weight_transfer/nccl_common.py
index b7de9f8becfc..97d4617fbbed 100644
--- a/vllm/distributed/weight_transfer/nccl_common.py
+++ b/vllm/distributed/weight_transfer/nccl_common.py
@@ -30,6 +30,7 @@
     DEFAULT_PACKED_BUFFER_SIZE_BYTES,
     DEFAULT_PACKED_NUM_BUFFERS,
 )
+from vllm.platforms import current_platform
 from vllm.utils.nccl import unpinned_nccl_env
 
 
@@ -171,6 +172,16 @@ class NCCLRendezvous(Protocol):
     world_size: int
 
 
+def _require_usable_communicator(comm: "PyNcclCommunicator") -> "PyNcclCommunicator":
+    """A disabled communicator makes every broadcast a no-op (stale weights)."""
+    if comm.disabled:
+        raise RuntimeError(
+            f"PyNccl communicator is disabled on {current_platform.device_type}; "
+            "NCCL weight transfer would be a no-op."
+        )
+    return comm
+
+
 def stateless_init_metadata_group(
     master_address: str,
     master_port: int,
@@ -211,7 +222,7 @@ def stateless_init_process_group(
         host=master_address, port=master_port, rank=rank, world_size=world_size
     )
     with unpinned_nccl_env():
-        return PyNcclCommunicator(pg, device=device)
+        return _require_usable_communicator(PyNcclCommunicator(pg, device=device))
 
 
 def uid_init_process_group(
@@ -228,11 +239,13 @@ def uid_init_process_group(
     from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
 
     with unpinned_nccl_env():
-        return PyNcclCommunicator.from_unique_id_bytes(
-            nccl_unique_id_bytes,
-            rank=rank,
-            world_size=world_size,
-            device=device,
+        return _require_usable_communicator(
+            PyNcclCommunicator.from_unique_id_bytes(
+                nccl_unique_id_bytes,
+                rank=rank,
+                world_size=world_size,
+                device=device,
+            )
         )
 
 
@@ -305,9 +318,8 @@ def trainer_init(
             - master_address: str
             - master_port: int
             - world_size: int
-        rank: This trainer process's rank in the group. The broadcast backends
-            have a single trainer at rank 0; multi-rank trainers (m2n) occupy
-            ranks `[0, num_trainer_ranks)` and the workers start after them.
+        rank: This trainer process's rank in the group. Trainer ranks occupy
+            `[0, num_trainer_ranks)` and the workers start after them.
 
     Returns:
         PyNcclCommunicator for weight transfer.
diff --git a/vllm/distributed/weight_transfer/nccl_engine.py b/vllm/distributed/weight_transfer/nccl_engine.py
index 7ec98ef9969c..3d99406746e1 100644
--- a/vllm/distributed/weight_transfer/nccl_engine.py
+++ b/vllm/distributed/weight_transfer/nccl_engine.py
@@ -55,8 +55,7 @@ class NCCLTrainerInitInfo(TrainerInitInfo):
     """Trainer-side init info for the dense NCCL weight transfer backend.
 
     The sender opens its endpoint as NCCL rank 0, so it needs no `rank_offset`.
-    `world_size` is the full trainer+worker NCCL group size. `rank` (from
-    `TrainerInitInfo`) identifies this trainer process; rank 0 is the sender.
+    `rank` (from `TrainerInitInfo`) identifies this trainer process.
 
     The trainer joins over a TCPStore rendezvous (`master_address` +
     `master_port`). Torch-free trainers (e.g. JAX) that cannot join a TCPStore
@@ -76,6 +75,8 @@ class NCCLTrainerInitInfo(TrainerInitInfo):
     master_address: str
     master_port: int
     world_size: int
+    """Transfer group size: rank 0 and each distinct `ParamMeta.owners[0]`, plus
+    the inference workers (`1 + num_workers` for a single sender)."""
     packed: bool = True
     packed_buffer_size_bytes: int = DEFAULT_PACKED_BUFFER_SIZE_BYTES
     packed_num_buffers: int = DEFAULT_PACKED_NUM_BUFFERS
@@ -85,19 +86,16 @@ class NCCLTrainerInitInfo(TrainerInitInfo):
 class NCCLWeightTransferUpdateInfo(WeightTransferUpdateInfo):
     """Per-round update info for the dense NCCL weight transfer backend.
 
-    Whether the transfer is packed (and the buffer geometry) is a must-agree
-    wire param carried on the init info (`NCCLTrainerInitInfo` /
-    `NCCLWeightTransferInitInfo`), not here; this carries only the per-round
-    parameter metadata.
+    Carries the parameter metadata and `src`, the broadcasting group rank; the
+    packed wire params live on the init info.
 
-    Leave all three lists empty to reuse the metadata of the previous
-    `update_weights` call on the same model, so a trainer that sends the same
-    parameters every round ships the names, dtypes and shapes only once.
+    Empty lists reuse the previous update's metadata for the same model and `src`.
     """
 
     names: list[str] = field(default_factory=list)
     dtype_names: list[str] = field(default_factory=list)
     shapes: list[list[int]] = field(default_factory=list)
+    src: int = 0
 
     def __post_init__(self):
         """Validate that all lists have the same length."""
@@ -119,10 +117,8 @@ class NCCLWeightTransferEngine(
 ):
     """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
-    in a process group. Received weights are loaded via the model's
-    `load_weights` using the layerwise reload lifecycle.
+    Receives dense checkpoint-format weights broadcast from the trainer ranks
+    below `rank_offset` and loads them with the layerwise reload lifecycle.
     """
 
     # Define backend-specific dataclass types
@@ -140,11 +136,14 @@ def __init__(
         self.model_update_group: PyNcclCommunicator | None = None
         # Set from the trainer-supplied init info at the handshake; defaults are
         # only for the (unreachable) receive-before-init case.
+        self.rank_offset = 1
         self.packed = False
         self.packed_buffer_size_bytes = DEFAULT_PACKED_BUFFER_SIZE_BYTES
         self.packed_num_buffers = DEFAULT_PACKED_NUM_BUFFERS
-        # Metadata of the last update per model, reused by empty update info.
-        self._last_update_info: dict[torch.nn.Module, NCCLWeightTransferUpdateInfo] = {}
+        # Metadata of the last update per (model, src), reused by empty update info.
+        self._last_update_info: dict[
+            tuple[torch.nn.Module, int], NCCLWeightTransferUpdateInfo
+        ] = {}
 
     def init_transfer_engine(self, init_info: NCCLWeightTransferInitInfo) -> None:
         """Initialize NCCL process group with the trainer and record the
@@ -156,6 +155,7 @@ def init_transfer_engine(self, init_info: NCCLWeightTransferInitInfo) -> None:
                       rank offset, world size, and the packed wire params
 
         """
+        self.rank_offset = init_info.rank_offset
         self.packed = init_info.packed
         self.packed_buffer_size_bytes = init_info.packed_buffer_size_bytes
         self.packed_num_buffers = init_info.packed_num_buffers
@@ -165,20 +165,19 @@ def init_transfer_engine(self, init_info: NCCLWeightTransferInitInfo) -> None:
         )
 
     def start_weight_update(self) -> None:
-        """Initialize layerwise reloading for the incoming checkpoint weights."""
-        from vllm.model_executor.model_loader.reload import (
-            initialize_layerwise_reload,
-        )
+        """Begin reloading the incoming checkpoint weights."""
+        from vllm.model_executor.model_loader.reload import start_reload
 
-        initialize_layerwise_reload(self.model)
+        start_reload(self.model)
 
     def finish_weight_update(self) -> None:
-        """Finalize layerwise reloading after all weights have been received."""
-        from vllm.model_executor.model_loader.reload import (
-            finalize_layerwise_reload,
-        )
+        """Finish reloading after all weights have been received."""
+        from vllm.model_executor.model_loader.reload import finish_reload
 
-        finalize_layerwise_reload(self.model, self.model_config)
+        finish_reload(self.model, self.model_config, self.device)
+        # Release the cached receive buffers and post-load temporaries so memory
+        # allocated outside torch (NCCL, other processes on the GPU) can use it.
+        torch.accelerator.empty_cache()
 
     def receive_weights(self, update_info: NCCLWeightTransferUpdateInfo) -> None:
         """Receive weights from trainer via NCCL broadcast.
@@ -189,8 +188,8 @@ def receive_weights(self, update_info: NCCLWeightTransferUpdateInfo) -> None:
         encoded.
 
         Args:
-            update_info: NCCL update info containing parameter names, dtypes,
-                        and shapes; empty to reuse the previous call's
+            update_info: Parameter names, dtypes, shapes and `src`; empty lists
+                        reuse the previous call's from the same `src`
 
         """
         if self.model_update_group is None:
@@ -198,10 +197,17 @@ def receive_weights(self, update_info: NCCLWeightTransferUpdateInfo) -> None:
                 "NCCL weight transfer not initialized. "
                 "Call init_transfer_engine() first."
             )
+        src = update_info.src
+        if not 0 <= src < self.rank_offset:
+            raise ValueError(
+                f"NCCL update src must be a sender rank in [0, {self.rank_offset}), "
+                f"got {src!r}."
+            )
+        key = (self.model, src)
         if update_info.names:
-            self._last_update_info[self.model] = update_info
+            self._last_update_info[key] = update_info
         else:
-            update_info = self._last_update_info.get(self.model, update_info)
+            update_info = self._last_update_info.get(key, update_info)
 
         from vllm.model_executor.model_loader.mtp_validation import (
             disable_mtp_completeness_check,
@@ -220,7 +226,7 @@ def state_dict_info_iterator():
                 packed_nccl_broadcast_consumer(
                     iterator=state_dict_info_iterator(),
                     group=self.model_update_group,
-                    src=0,
+                    src=src,
                     post_unpack_func=self.model.load_weights,
                     buffer_size_bytes=self.packed_buffer_size_bytes,
                     num_buffers=self.packed_num_buffers,
@@ -234,7 +240,7 @@ def state_dict_info_iterator():
                     dtype = getattr(torch, dtype_name)
                     weight = torch.empty(shape, dtype=dtype, device=self.device)
                     self.model_update_group.broadcast(
-                        weight, src=0, stream=torch.cuda.current_stream()
+                        weight, src=src, stream=torch.cuda.current_stream()
                     )
                     self.model.load_weights([(name, weight)])
                     del weight
@@ -248,17 +254,11 @@ def shutdown(self) -> None:
 class NCCLTrainerWeightTransferEngine(TrainerWeightTransferEngine[NCCLTrainerInitInfo]):
     """Trainer-side NCCL weight transfer engine.
 
-    On the sender (rank 0) holds the NCCL communicator and drives the full
-    update round trip: it runs the inference-side `update_weights` concurrently
-    with the trainer-side broadcast (both rendezvous inside the same NCCL
-    calls), then finishes the update. Non-sender trainer ranks hold no
-    communicator; they only drive `iter_full()` over the names they own, to stay
-    in the trainer-side gathers (e.g. FSDP `full_tensor()`), and skip the client
-    RPCs and the broadcast (all guarded on `is_sender`).
+    Each parameter is broadcast by its first owner (rank 0 when unset); rank 0
+    also drives the client RPCs. See "Broadcasting ranks" in `nccl.md`.
 
-    Each parameter is sent by its first owner, so every parameter's owners must
-    start at rank 0 (or be unset): a partitioned source with several senders
-    (pipeline stages) is rejected.
+    Other trainer ranks hold no communicator and only drive `iter_full()` to
+    stay in the trainer-side gathers (e.g. FSDP `full_tensor()`).
 
     `packed` / buffer sizes come from `NCCLTrainerInitInfo`; the sender
     propagates them to the worker at `trainer_init` (on the worker-side init
@@ -283,6 +283,8 @@ def __init__(
         self.packed_buffer_size_bytes = packed_buffer_size_bytes
         self.packed_num_buffers = packed_num_buffers
         self.model_update_group: PyNcclCommunicator | None = None
+        # This rank's rank in the transfer group, set by trainer_init.
+        self._group_rank = 0
         # Metadata the worker already holds; an unchanged round sends none.
         self._sent_meta: list[ParamMeta] | None = None
 
@@ -296,7 +298,19 @@ def _build_worker_init_payload(cls, init_info: NCCLTrainerInitInfo) -> dict:
         )
 
     def _open_trainer_endpoint(self, init_info: NCCLTrainerInitInfo) -> None:
-        self.model_update_group = open_trainer_endpoint(init_info)
+        self.model_update_group = open_trainer_endpoint(init_info, self._group_rank)
+
+    @staticmethod
+    def _broadcast_plan(
+        meta: list[ParamMeta],
+    ) -> tuple[list[int], list[tuple[int, list[ParamMeta]]]]:
+        """The transfer group's trainer ranks, in group-rank order, and each
+        broadcaster's parameters (in metadata order), in broadcaster order."""
+        slices: dict[int, list[ParamMeta]] = {}
+        for m in meta:
+            slices.setdefault(0 if m.owners is None else m.owners[0], []).append(m)
+        group_ranks = sorted({0, *slices})
+        return group_ranks, sorted(slices.items())
 
     @classmethod
     def trainer_init(
@@ -314,13 +328,39 @@ def trainer_init(
                 "does not implement iter_full()."
             )
         # Identical on every rank, so every rank raises alike.
+        group_ranks, slices = cls._broadcast_plan(source.metadata)
         for m in source.metadata:
-            if m.owners is not None and m.owners[0] != 0:
+            if m.owners is not None and 0 in m.owners[1:]:
                 raise ValueError(
-                    f"parameter {m.name!r} is sent by its first owner, rank "
-                    f"{m.owners[0]}, but NCCL broadcast has a single sender, "
-                    "rank 0. Every parameter's owners must start at rank 0."
+                    f"parameter {m.name!r} has owners {m.owners}: rank 0 only "
+                    "gathers in its own turn, so it must be the first owner of "
+                    "every parameter it owns."
                 )
+        if len(slices) > 1:
+            broadcaster_of: dict[int, int] = {}
+            for broadcaster, metas in slices:
+                for m in metas:
+                    if m.owners is None:
+                        raise ValueError(
+                            f"parameter {m.name!r} has no owners, so every rank "
+                            "owns it, but the source has several broadcasting "
+                            f"ranks {[r for r, _ in slices]}. Every parameter "
+                            "must declare its owners."
+                        )
+                    for r in m.owners:
+                        prev = broadcaster_of.setdefault(r, broadcaster)
+                        if prev != broadcaster:
+                            raise ValueError(
+                                f"rank {r} owns parameters broadcast by ranks "
+                                f"{prev} and {broadcaster}, but each broadcaster "
+                                "sends in its own turn, so every parameter a "
+                                "rank owns must share its first owner."
+                            )
+        if init_info.world_size <= len(group_ranks):
+            raise ValueError(
+                f"world_size {init_info.world_size} leaves no inference workers "
+                f"after the {len(group_ranks)} broadcasting trainer ranks."
+            )
         engine = cls(
             client=client,
             source=source,
@@ -330,13 +370,17 @@ def trainer_init(
             packed_buffer_size_bytes=init_info.packed_buffer_size_bytes,
             packed_num_buffers=init_info.packed_num_buffers,
         )
+        if engine.rank not in group_ranks:
+            # Ranks that broadcast nothing stay out of the transfer group; they
+            # only join the trainer-side gathers in send_weights.
+            return engine
+        engine._group_rank = group_ranks.index(engine.rank)
         if not engine.is_sender:
-            # Non-sender trainer ranks aren't part of the transfer NCCL group and
-            # don't drive the inference side; they only participate in the
-            # trainer-side gather during send_weights.
+            engine._open_trainer_endpoint(init_info)
             return engine
 
         payload = cls._build_worker_init_payload(init_info)
+        payload["rank_offset"] = len(group_ranks)
         run_trainer_rendezvous(
             init_worker=lambda: engine.client.init_weight_transfer_engine(payload),
             init_trainer=lambda: engine._open_trainer_endpoint(init_info),
@@ -351,34 +395,57 @@ def send_weights(self) -> None:
         # Metadata is declared without gathering. For Megatron it is itself a
         # collective, so every rank runs it; only the sender ships it.
         meta = source.metadata
+        group_ranks, slices = self._broadcast_plan(meta)
 
-        if not self.is_sender:
-            # Non-sender ranks only join the trainer-side gather collective.
-            self._broadcast(source, meta)
+        if self.rank not in group_ranks:
+            # Ranks outside the transfer group only join the trainer-side
+            # gather collective.
+            for _ in source.iter_full():
+                pass
             self._post_send_sync()
             return
 
-        if meta == self._sent_meta:
-            update_info = NCCLWeightTransferUpdateInfo()
+        reuse = self.is_sender and meta == self._sent_meta
+        if self.is_sender:
+            self.client.start_weight_update()
+        for broadcaster, metas in slices:
+            src = group_ranks.index(broadcaster)
+            if self.is_sender:
+                self._send_update(metas, src, reuse)
+            else:
+                self._broadcast(metas, src)
+        if self.is_sender:
+            self.client.finish_weight_update()
+            self._sent_meta = meta
+        self._post_send_sync()
+
+    def _send_update(self, meta: list[ParamMeta], src: int, reuse: bool) -> None:
+        """Rank 0's part of one broadcaster's turn: run its `update_weights`
+        concurrently with `_broadcast`."""
+        if reuse:
+            update_info = NCCLWeightTransferUpdateInfo(src=src)
         else:
             update_info = NCCLWeightTransferUpdateInfo(
                 names=[m.name for m in meta],
                 dtype_names=[str(m.dtype).split(".")[-1] for m in meta],
                 shapes=[list(m.shape) for m in meta],
+                src=src,
             )
-
-        self.client.start_weight_update()
+        payload = asdict(update_info)
+        if not src:
+            # Keep single-broadcaster payloads readable by older servers.
+            del payload["src"]
         # update_weights (workers receive) must run concurrently with the
         # trainer-side broadcast — both rendezvous inside the same NCCL calls.
         exe = ThreadPoolExecutor(max_workers=1)
         try:
-            future = exe.submit(self.client.update_weights, asdict(update_info))
+            future = exe.submit(self.client.update_weights, payload)
             # Cheap best-effort: if update_weights already failed (e.g. a bad
             # request rejected before any NCCL call), surface it now instead of
             # hanging in broadcast waiting for a peer that will never arrive.
             if future.done():
                 future.result()
-            self._broadcast(source, meta)
+            self._broadcast(meta, src)
             future.result()  # surface inference-side errors
         finally:
             # Never wait for the RPC thread here. If the broadcast raised, the
@@ -389,28 +456,23 @@ def send_weights(self) -> None:
             # thread is still joined at interpreter exit, so a caller that wants
             # to exit cleanly after such a failure must not wait on it.)
             exe.shutdown(wait=False)
-        self.client.finish_weight_update()
-        self._sent_meta = meta
-        self._post_send_sync()
-
-    def _broadcast(self, source: WeightSource, meta: list[ParamMeta]) -> None:
-        """Drive `iter_full()` (gathering among each parameter's owners) and, on
-        the sender, broadcast from rank 0, packed or one-by-one. Other owners only
-        join the gathers."""
-        if not self.is_sender:
-            for _ in source.iter_full():
-                pass
-            return
 
+    def _broadcast(self, meta: list[ParamMeta], src: int) -> None:
+        """Broadcast `meta` from group rank `src`: that rank sends from
+        `iter_full()`, every other group member receives and discards."""
         assert self.model_update_group is not None, (
             "trainer_init() must be called before _broadcast()."
         )
-        pairs = self._checked_iter(source, meta)
+        if src != self.model_update_group.rank:
+            self._discard(meta, src)
+            return
+        assert self.source is not None
+        pairs = self._checked_iter(self.source, meta)
         if self.packed:
             packed_nccl_broadcast_producer(
                 iterator=pairs,
                 group=self.model_update_group,
-                src=0,
+                src=src,
                 post_iter_func=lambda item: item[1],
                 buffer_size_bytes=self.packed_buffer_size_bytes,
                 num_buffers=self.packed_num_buffers,
@@ -423,7 +485,27 @@ def _broadcast(self, source: WeightSource, meta: list[ParamMeta]) -> None:
                 # pointer. Keep the copy referenced until the broadcast is
                 # enqueued. (The packed path linearizes in `pack_tensors`.)
                 send = tensor if tensor.is_contiguous() else tensor.contiguous()
-                self.model_update_group.broadcast(send, src=0, stream=stream)
+                self.model_update_group.broadcast(send, src=src, stream=stream)
+
+    def _discard(self, meta: list[ParamMeta], src: int) -> None:
+        """Receive another broadcaster's parameters and drop them."""
+        assert self.model_update_group is not None
+        device = self.model_update_group.device
+        if self.packed:
+            packed_nccl_broadcast_consumer(
+                iterator=((m.name, (list(m.shape), m.dtype)) for m in meta),
+                group=self.model_update_group,
+                src=src,
+                post_unpack_func=lambda weights: None,
+                buffer_size_bytes=self.packed_buffer_size_bytes,
+                num_buffers=self.packed_num_buffers,
+                device=device,
+            )
+        else:
+            stream = torch.cuda.current_stream()
+            for m in meta:
+                weight = torch.empty(m.shape, dtype=m.dtype, device=device)
+                self.model_update_group.broadcast(weight, src=src, stream=stream)
 
     @staticmethod
     def _checked_iter(
@@ -437,8 +519,7 @@ def _checked_iter(
         iteration disagrees with it, the two sides split the stream differently
         and the transfer hangs in NCCL or loads garbage. Checking here costs one
         comparison per parameter and turns that into an error naming the first
-        divergent parameter. The sender owns every parameter, so it yields all
-        of `meta`.
+        divergent parameter.
         """
         sent = 0
         for name, tensor in source.iter_full():
diff --git a/vllm/distributed/weight_transfer/sparse_nccl_engine.py b/vllm/distributed/weight_transfer/sparse_nccl_engine.py
index b50a0be2c655..cc81f8b626f7 100644
--- a/vllm/distributed/weight_transfer/sparse_nccl_engine.py
+++ b/vllm/distributed/weight_transfer/sparse_nccl_engine.py
@@ -145,8 +145,12 @@ def start_weight_update(self) -> None:
         pass
 
     def finish_weight_update(self) -> None:
-        """No-op: sparse patches are applied in place, no layerwise reload."""
-        pass
+        """No layerwise reload; rebuild model-level derived state once."""
+        from vllm.model_executor.model_loader.weight_utils import (
+            finalize_model_weights,
+        )
+
+        finalize_model_weights(self.model)
 
     def receive_weights(self, update_info: SparseNCCLWeightTransferUpdateInfo) -> None:
         """Receive sparse flat-index patches from the trainer and apply them."""
diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py
index 423544a2a135..dd5e8c3bbdfa 100644
--- a/vllm/engine/arg_utils.py
+++ b/vllm/engine/arg_utils.py
@@ -282,10 +282,8 @@ def get_type_hints(type_hint: TypeHint) -> set[TypeHint]:
     return type_hints
 
 
-NEEDS_HELP = (
-    any("--help" in arg for arg in sys.argv)  # vllm SUBCOMMAND --help
-    or "mkdocs" in sys.modules  # mkdocs SUBCOMMAND
-)
+IS_ARGPARSE = any("--help" in arg for arg in sys.argv)  # vllm SUBCOMMAND --help
+NEEDS_HELP = IS_ARGPARSE or "mkdocs" in sys.modules  # mkdocs SUBCOMMAND
 
 
 def _maybe_add_docs_url(cls: Any) -> str:
@@ -354,8 +352,9 @@ def _compute_kwargs(cls: ConfigType) -> dict[str, dict[str, Any]]:
         # Get the help text for the field
         name = field.name
         help = cls_docs.get(name, "").strip()
-        # Escape % for argparse
-        help = help.replace("%", "%%")
+        # Escape % for argparse's %-formatting in _expand_help
+        if IS_ARGPARSE:
+            help = help.replace("%", "%%")
 
         # Initialise the kwargs dictionary for the field
         kwargs[name] = {"default": default, "help": help}
@@ -460,7 +459,7 @@ def get_kwargs(cls: ConfigType) -> dict[str, dict[str, Any]]:
 
 
 _MAMBA_CACHE_MODE_DEPRECATION_MESSAGE = (
-    "--mamba-cache-mode is deprecated and will be removed in v0.33.0. "
+    "--mamba-cache-mode is deprecated and will be removed in v0.34.0. "
     "It has no effect: Mamba prefix caching follows --enable-prefix-caching."
 )
 
@@ -541,7 +540,7 @@ class EngineArgs:
     cpu_distributed_timeout_seconds: int | None = (
         ParallelConfig.cpu_distributed_timeout_seconds
     )
-    numa_bind: bool = ParallelConfig.numa_bind
+    numa_bind: bool | None = ParallelConfig.numa_bind
     numa_bind_nodes: list[int] | None = ParallelConfig.numa_bind_nodes
     numa_bind_cpus: list[str] | None = ParallelConfig.numa_bind_cpus
     device_ids: list[int | str] | None = None
@@ -611,8 +610,9 @@ class EngineArgs:
     offload_num_in_group: int = PrefetchOffloadConfig.offload_num_in_group
     offload_prefetch_step: int = PrefetchOffloadConfig.offload_prefetch_step
     offload_params: set[str] = get_field(PrefetchOffloadConfig, "offload_params")
-    gpu_memory_utilization: float = CacheConfig.gpu_memory_utilization
+    gpu_memory_utilization: float | None = CacheConfig.gpu_memory_utilization
     kv_cache_memory_bytes: int | None = CacheConfig.kv_cache_memory_bytes
+    enable_extensible_kv_cache: bool | None = CacheConfig.enable_extensible_kv_cache
     max_num_batched_tokens: int | None = None
     max_num_scheduled_tokens: int | None = None
     long_prefill_token_threshold: int = SchedulerConfig.long_prefill_token_threshold
@@ -1399,6 +1399,10 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
         cache_group.add_argument(
             "--kv-cache-memory-bytes", **cache_kwargs["kv_cache_memory_bytes"]
         )
+        cache_group.add_argument(
+            "--enable-extensible-kv-cache",
+            **cache_kwargs["enable_extensible_kv_cache"],
+        )
         cache_group.add_argument("--kv-cache-dtype", **cache_kwargs["cache_dtype"])
         cache_group.add_argument(
             "--num-gpu-blocks-override", **cache_kwargs["num_gpu_blocks_override"]
@@ -2314,6 +2318,7 @@ def create_engine_config(
             block_size=self.block_size,  # type: ignore[arg-type]
             gpu_memory_utilization=self.gpu_memory_utilization,
             kv_cache_memory_bytes=self.kv_cache_memory_bytes,
+            enable_extensible_kv_cache=self.enable_extensible_kv_cache,  # type: ignore[arg-type]
             cache_dtype=resolved_cache_dtype,  # type: ignore[arg-type]
             is_attention_free=model_config.is_attention_free,
             num_gpu_blocks_override=self.num_gpu_blocks_override,
diff --git a/vllm/entrypoints/anthropic/protocol.py b/vllm/entrypoints/anthropic/protocol.py
index 72bd58e837bf..1559c405205b 100644
--- a/vllm/entrypoints/anthropic/protocol.py
+++ b/vllm/entrypoints/anthropic/protocol.py
@@ -133,12 +133,29 @@ def validate_tool_change_role(self) -> "AnthropicMessage":
 class AnthropicTool(BaseModel):
     """Tool definition."""
 
+    type: Literal["custom"] | None = None
     name: str
     description: str | None = None
     input_schema: dict[str, Any]
     strict: bool | None = None
     defer_loading: bool | None = None
 
+    @model_validator(mode="before")
+    @classmethod
+    def reject_anthropic_defined_tools(cls, data: Any) -> Any:
+        # Anthropic-defined tools (web_search_*, code_execution_*, bash_*,
+        # text_editor_*, ...) carry a versioned ``type`` and no input_schema.
+        # Reject them by name here, before field validation would report a
+        # missing input_schema and point the client at the wrong problem.
+        if isinstance(data, dict):
+            tool_type = data.get("type")
+            if tool_type not in (None, "custom"):
+                raise ValueError(
+                    f"tool type {tool_type!r} is not supported; only custom "
+                    "tools with an input_schema are supported"
+                )
+        return data
+
     @field_validator("input_schema")
     @classmethod
     def validate_input_schema(cls, v):
diff --git a/vllm/entrypoints/anthropic/serving.py b/vllm/entrypoints/anthropic/serving.py
index e1b362ab4136..4ca44cf24f9b 100644
--- a/vllm/entrypoints/anthropic/serving.py
+++ b/vllm/entrypoints/anthropic/serving.py
@@ -56,8 +56,12 @@
 from vllm.logger import init_logger
 from vllm.renderers.chat_utils import ChatTemplateContentFormatOption
 from vllm.renderers.hf import HfRenderer, resolve_chat_template
-from vllm.renderers.inputs.preprocess import extract_prompt_components
+from vllm.renderers.inputs.preprocess import (
+    extract_prompt_components,
+    parse_dec_only_prompt,
+)
 from vllm.renderers.online_renderer import OnlineRenderer
+from vllm.renderers.params import ChatParams
 
 logger = init_logger(__name__)
 
@@ -273,8 +277,10 @@ def _should_merge_inline_system(cls, online_renderer: OnlineRenderer) -> bool:
         """
         renderer = online_renderer.renderer
         tool_variants: tuple[list[dict[str, Any]] | None, ...] = (None, [])
-        if not isinstance(renderer, HfRenderer) or renderer.tokenizer is None:
-            merge = cls._detect_merge_inline_system(online_renderer.chat_template)
+        if renderer.tokenizer is None:
+            merge = True
+        elif not isinstance(renderer, HfRenderer):
+            merge = not cls._renderer_keeps_inline_system(online_renderer)
         else:
             merge = any(
                 cls._detect_merge_inline_system(
@@ -288,13 +294,18 @@ def _should_merge_inline_system(cls, online_renderer: OnlineRenderer) -> bool:
                 for tools in tool_variants
             )
         if merge:
+            hint = (
+                " Pass a --chat-template that accepts non-leading system "
+                "messages to avoid this."
+                if isinstance(renderer, HfRenderer)
+                else ""
+            )
             logger.warning_once(
                 "The chat template requires system-first ordering, so inline "
                 "system messages in /v1/messages requests (e.g. Claude Code's "
                 "per-turn reminders) are merged into the leading system prompt. "
                 "Each new one changes the prompt prefix, so the rest of the "
-                "conversation misses the prefix cache. Pass a --chat-template "
-                "that accepts non-leading system messages to avoid this."
+                "conversation misses the prefix cache." + hint
             )
         return merge
 
@@ -324,9 +335,60 @@ def _detect_merge_inline_system(chat_template: str | None) -> bool:
                 add_generation_prompt=False,
             )
             return False
-        except jinja2.TemplateError:
+        except Exception:
+            # Non-TemplateError failures (e.g. TypeError from filters on
+            # undefined variables) must not abort server startup.
             return True
 
+    @staticmethod
+    def _renderer_keeps_inline_system(online_renderer: OnlineRenderer) -> bool:
+        """Probe a non-jinja renderer (e.g. DeepSeek V4.1's encoder).
+
+        Inline system messages must render in order as distinct system turns,
+        and a trailing one must still be followed by the generation prompt.
+        DeepSeek V4 appends their text to the user turn, V3.2 puts it after the
+        generation prompt, and Mistral moves it into the last user turn.
+        """
+        markers = ("SYSA", "USRA", "SYSB", "USRB")
+        messages: list[Any] = [
+            {"role": role, "content": marker}
+            for role, marker in zip(("system", "user") * 2, markers)
+        ]
+        renderer = online_renderer.renderer
+        params = ChatParams(
+            chat_template=online_renderer.chat_template,
+            chat_template_content_format=online_renderer.chat_template_content_format,
+            chat_template_kwargs={"add_generation_prompt": True},
+        ).with_defaults(online_renderer.default_chat_template_kwargs)
+
+        def render(msgs: list[Any]) -> str:
+            _, prompt = renderer.render_messages(msgs, params)
+            text, token_ids, _ = extract_prompt_components(
+                renderer.model_config, parse_dec_only_prompt(prompt)
+            )
+            if text is None:
+                text = renderer.get_tokenizer().decode(token_ids or [])
+            return text
+
+        def separator(text: str) -> str:
+            return text[text.find(markers[1]) + len(markers[1]) : text.find(markers[2])]
+
+        try:
+            full = render(messages)
+            trailing = render(messages[:3])
+            as_user = render([*messages[:2], {**messages[2], "role": "user"}])
+        except Exception:
+            logger.debug("Inline system probe failed", exc_info=True)
+            return False
+        positions = [full.find(marker) for marker in markers]
+        if -1 in positions or positions != sorted(positions):
+            return False
+        system_separator = separator(full)
+        if not system_separator.strip() or system_separator == separator(as_user):
+            return False
+        generation_prompt = full[positions[-1] + len(markers[-1]) :]
+        return trailing.endswith(markers[2] + generation_prompt)
+
     @staticmethod
     def _convert_image_source_to_url(source: dict[str, Any]) -> str:
         """Convert an Anthropic image source to an OpenAI-compatible URL.
@@ -595,11 +657,18 @@ def _convert_user_tool_result(
                         )
             tool_text = "\n".join(text_parts)
 
+        tool_content: str | list[dict[str, Any]] = tool_text or ""
+        if tool_reference:
+            # One tool message per tool_result. Chat templates that expand
+            # tool references check the first content part.
+            tool_content = tool_reference
+            if tool_text:
+                tool_content.append({"type": "text", "text": tool_text})
         openai_messages.append(
             {
                 "role": "tool",
                 "tool_call_id": block.tool_use_id or "",
-                "content": tool_text or "",
+                "content": tool_content,
             }
         )
 
@@ -614,15 +683,6 @@ def _convert_user_tool_result(
                 }
             )
 
-        if tool_reference:
-            openai_messages.append(
-                {
-                    "role": "tool",
-                    "tool_call_id": block.tool_use_id or "",
-                    "content": tool_reference,  # type: ignore[dict-item]
-                }
-            )
-
     @classmethod
     def _build_base_request(
         cls,
diff --git a/vllm/entrypoints/cli/preload.py b/vllm/entrypoints/cli/preload.py
index ef54163ec5c3..669f6a46db56 100644
--- a/vllm/entrypoints/cli/preload.py
+++ b/vllm/entrypoints/cli/preload.py
@@ -24,9 +24,11 @@
 import signal
 import sys
 import threading
+import time
 import typing
 from collections.abc import Sequence
 from itertools import product
+from multiprocessing.connection import wait
 from multiprocessing.process import BaseProcess
 
 from fastapi import FastAPI
@@ -324,7 +326,7 @@ def _seed_endpoint(
                 return None
             return host, base_port + offset
 
-        procs = []
+        procs: list[BaseProcess] = []
         for group_index, (is_draft, config, init_method) in enumerate(groups):
             # Draft daemons live past the target group's port block so both
             # roles can be reached over a single seed host:base_port.
@@ -371,58 +373,79 @@ def _seed_endpoint(
                 args.weight_cache_health_port,
                 health_state,
             )
-        for proc in procs:
-            proc.start()
+        shutdown_requested = False
+        exit_code = 0
 
         def _shutdown(signum, frame):
-            if health_server is not None:
-                health_server.should_exit = True
-            for proc in procs:
-                proc.terminate()
+            nonlocal shutdown_requested, exit_code
+            shutdown_requested = True
+            exit_code = 128 + signum
 
         signal.signal(signal.SIGINT, _shutdown)
         signal.signal(signal.SIGTERM, _shutdown)
 
-        ready: set[tuple[str, int]] = set()
-        while len(ready) < len(expected_ready):
-            try:
-                ready.add(ready_queue.get(timeout=1.0))
-            except queue.Empty:
+        started: list[BaseProcess] = []
+        try:
+            for proc in procs:
+                if shutdown_requested:
+                    break
+                proc.start()
+                started.append(proc)
+
+            ready: set[tuple[str, int]] = set()
+            while not shutdown_requested:
                 dead = [p for p in procs if p.exitcode is not None]
                 if dead:
                     logger.error(
-                        "Weight cache daemon rank(s) exited during startup "
+                        "Weight cache daemon rank(s) exited during %s "
                         "(exitcodes=%s); shutting down.",
+                        "serving" if ready == expected_ready else "startup",
                         [p.exitcode for p in dead],
                     )
-                    for proc in procs:
-                        proc.terminate()
-                    for proc in procs:
-                        proc.join()
-                    if health_server is not None:
-                        health_server.should_exit = True
-                    if health_thread is not None:
-                        health_thread.join()
-                    sys.exit(max((p.exitcode or 0) for p in procs))
-        health_state.mark_ready()
-        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()
-        if health_server is not None:
-            health_server.should_exit = True
-        if health_thread is not None:
-            health_thread.join()
-        sys.exit(max(proc.exitcode or 0 for proc in procs))
+                    # Preserve the failure before terminating healthy siblings.
+                    # A clean rank exit is also unexpected while serving.
+                    failure = next((code for p in dead if (code := p.exitcode)), 1)
+                    exit_code = 128 - failure if failure < 0 else failure
+                    break
+                if ready == expected_ready:
+                    wait([p.sentinel for p in procs], timeout=1.0)
+                    continue
+                try:
+                    ready.add(ready_queue.get(timeout=1.0))
+                except queue.Empty:
+                    continue
+                if ready == expected_ready:
+                    health_state.mark_ready()
+                    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,
+                    )
+        finally:
+            if health_server is not None:
+                health_server.should_exit = True
+            for proc in started:
+                if proc.is_alive():
+                    proc.terminate()
+            deadline = time.monotonic() + 10.0
+            for proc in started:
+                proc.join(timeout=max(0.0, deadline - time.monotonic()))
+            for proc in started:
+                if proc.is_alive():
+                    proc.kill()
+            for proc in started:
+                proc.join()
+            if health_thread is not None:
+                health_thread.join()
+            ready_queue.close()
+        sys.exit(exit_code)
 
     def subparser_init(
         self, subparsers: argparse._SubParsersAction
diff --git a/vllm/entrypoints/cli/serve.py b/vllm/entrypoints/cli/serve.py
index 40f116514d3a..5d547a582410 100644
--- a/vllm/entrypoints/cli/serve.py
+++ b/vllm/entrypoints/cli/serve.py
@@ -243,16 +243,21 @@ def signal_handler(signum, frame):
         handshake_address,
     )
 
+    # A colocated frontend that serves only these engines (the Rust frontend in
+    # hybrid or external DP LB) also handshakes with them for their I/O.
+    client_handshake_address = args.headless_client_handshake_address
+
     # Create the engines.
     engine_manager = CoreEngineProcManager(
         local_engine_count=local_engine_count,
         start_index=vllm_config.parallel_config.data_parallel_rank,
         local_start_index=0,
         vllm_config=vllm_config,
-        local_client=False,
+        local_client=client_handshake_address is not None,
         handshake_address=handshake_address,
         executor_class=Executor.get_class(vllm_config),
         log_stats=not engine_args.disable_log_stats,
+        client_handshake_address=client_handshake_address,
     )
 
     try:
diff --git a/vllm/entrypoints/common/cli_args.py b/vllm/entrypoints/common/cli_args.py
index b861c9e73fb7..9340a3f945fb 100644
--- a/vllm/entrypoints/common/cli_args.py
+++ b/vllm/entrypoints/common/cli_args.py
@@ -448,6 +448,14 @@ def make_arg_parser(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
         help="Run in headless mode. See multi-node data parallel "
         "documentation for more details.",
     )
+    # Internal: the Rust frontend's managed `serve` passes its local handshake
+    # endpoint so that headless engines also handshake with it.
+    parser.add_argument(
+        "--headless-client-handshake-address",
+        type=str,
+        default=None,
+        help=argparse.SUPPRESS,
+    )
     parser.add_argument(
         "--api-server-count",
         "-asc",
diff --git a/vllm/entrypoints/generate/structured_decisions/serving.py b/vllm/entrypoints/generate/structured_decisions/serving.py
index 80398a450dc5..1ffbe4414df6 100644
--- a/vllm/entrypoints/generate/structured_decisions/serving.py
+++ b/vllm/entrypoints/generate/structured_decisions/serving.py
@@ -73,6 +73,15 @@ async def create_decision(
     ) -> StructuredDecisionResponse | ErrorResponse:
         if (error := await self._check_model(request)) is not None:  # type: ignore[arg-type]
             return error
+        renderer = self.strategy.context.online_renderer
+        if (
+            error := renderer.validate_chat_template(
+                request_chat_template=None,
+                chat_template_kwargs=request.chat_template_kwargs,
+                trust_request_chat_template=renderer.trust_request_chat_template,
+            )
+        ) is not None:
+            return error
         engine_client = self.strategy.context.engine_client
         if engine_client.errored:
             raise engine_client.dead_error
diff --git a/vllm/entrypoints/launchers/grpc_server.py b/vllm/entrypoints/launchers/grpc_server.py
index 9cdb18b6bc7b..d99e0334649a 100644
--- a/vllm/entrypoints/launchers/grpc_server.py
+++ b/vllm/entrypoints/launchers/grpc_server.py
@@ -23,6 +23,7 @@
 import signal
 import sys
 import time
+from functools import partial
 
 try:
     import grpc
@@ -52,6 +53,110 @@
 
 logger = init_logger(__name__)
 
+# Grace in-flight RPCs get with --shutdown-timeout 0. A positive timeout never
+# gives them less, so setting one is never less graceful than the default.
+_RPC_GRACE_S = 5.0
+# Teardown time EngineCore keeps past the drain budget (NIXL, device memory,
+# shm), as the default 5 s join gave it before.
+_ENGINE_TEARDOWN_S = 5.0
+# Time RPCs get to send their last responses once EngineCore has stopped;
+# whatever is still open after that (e.g. KV-event streams) is cancelled.
+_RPC_FLUSH_S = 5.0
+# How often to check whether anything can still send EngineCore a request.
+_IDLE_POLL_S = 0.1
+# RPCs that stay open until the client or the server ends them. They carry no
+# request, so shutdown does not wait for them before stopping EngineCore.
+_SUBSCRIPTION_METHODS = frozenset({"SubscribeKvEvents"})
+
+
+class _InflightRpcs(grpc.aio.ServerInterceptor):
+    """Counts the in-flight RPCs of one service, except subscriptions.
+
+    An RPC counts from the moment the server dispatches it, before its request
+    message is received, so one whose request is still arriving or being
+    preprocessed (tokenization, media fetch and processing), and is not yet
+    known to AsyncLLM, counts too. It stops counting when the RPC's task
+    finishes, whether or not its handler ran.
+    """
+
+    def __init__(self, service: str):
+        self.count = 0
+        self._prefix = f"/{service}/"
+
+    async def intercept_service(self, continuation, handler_call_details):
+        handler = await continuation(handler_call_details)
+        method = handler_call_details.method
+        if (
+            handler is not None
+            and method.startswith(self._prefix)
+            and method[len(self._prefix) :] not in _SUBSCRIPTION_METHODS
+        ):
+            # grpc.aio runs interceptors in the task that serves the RPC.
+            self.count += 1
+            asyncio.current_task().add_done_callback(self._done)
+        return handler
+
+    def _done(self, task: asyncio.Task) -> None:
+        self.count -= 1
+
+
+async def _wait_for_requests(
+    rpcs: asyncio.Future, async_llm: AsyncLLM, inflight: _InflightRpcs
+) -> None:
+    """Return once ``rpcs`` is done or nothing can still send EngineCore a request.
+
+    EngineCore rejects every request it reads once its shutdown has started,
+    even one an RPC forwarded before SIGTERM, so it must keep running while
+    an RPC other than a subscription is in flight or AsyncLLM still has a
+    request. Idle has to hold on two polls in a row, which covers an RPC that
+    gRPC accepted before SIGTERM but has not yet handed to its task.
+    """
+    idle = False
+    while not rpcs.done():
+        was_idle = idle
+        idle = inflight.count == 0 and async_llm.get_num_unfinished_requests() == 0
+        if was_idle and idle:
+            return
+        await asyncio.wait([rpcs], timeout=_IDLE_POLL_S)
+
+
+async def _drain_and_stop(
+    server, async_llm: AsyncLLM, shutdown_timeout: int, inflight: _InflightRpcs
+) -> None:
+    """Stop the gRPC server and AsyncLLM within ``--shutdown-timeout``.
+
+    0 keeps the abort behavior: in-flight RPCs get ``_RPC_GRACE_S``, then
+    EngineCore gets its default join. Otherwise the server refuses new RPCs at
+    once and EngineCore keeps serving until the in-flight RPCs (``inflight``)
+    and their requests are done; RPCs still running at the deadline
+    (``--shutdown-timeout``, but no less than ``_RPC_GRACE_S``) are cancelled
+    and abort their requests. EngineCore then drains work that outlives an
+    RPC, such as NIXL prefill blocks the decode side has not read yet, until
+    the same deadline plus ``_ENGINE_TEARDOWN_S`` for teardown.
+    """
+    loop = asyncio.get_running_loop()
+    if shutdown_timeout == 0:
+        await server.stop(grace=_RPC_GRACE_S)
+        logger.info("gRPC server stopped")
+        await loop.run_in_executor(None, async_llm.shutdown)
+        logger.info("AsyncLLM engine stopped")
+        return
+
+    grace = max(shutdown_timeout, _RPC_GRACE_S)
+    deadline = loop.time() + grace
+    logger.info("Draining in-flight requests (shutdown_timeout=%ds)", shutdown_timeout)
+    rpcs = asyncio.ensure_future(server.stop(grace=grace))
+    await _wait_for_requests(rpcs, async_llm, inflight)
+    remaining = max(deadline - loop.time(), 0.0)
+    await loop.run_in_executor(
+        None, partial(async_llm.shutdown, timeout=remaining + _ENGINE_TEARDOWN_S)
+    )
+    logger.info("AsyncLLM engine stopped")
+    await asyncio.wait([rpcs], timeout=_RPC_FLUSH_S)
+    await server.stop(grace=None)
+    await rpcs
+    logger.info("gRPC server stopped")
+
 
 async def serve_grpc(args: argparse.Namespace):
     """Main gRPC serving function.
@@ -85,7 +190,11 @@ async def serve_grpc(args: argparse.Namespace):
     servicer = VllmEngineServicer(async_llm, start_time)
 
     # Create gRPC server
+    inflight = _InflightRpcs(
+        vllm_engine_pb2.DESCRIPTOR.services_by_name["VllmEngine"].full_name
+    )
     server = grpc.aio.server(
+        interceptors=[inflight],
         options=[
             ("grpc.max_send_message_length", -1),
             ("grpc.max_receive_message_length", -1),
@@ -146,8 +255,15 @@ async def _force_log():
         # Handle shutdown signals
         loop = asyncio.get_running_loop()
         stop_event = asyncio.Event()
+        force_stops = []
 
         def signal_handler():
+            if stop_event.is_set():
+                # A second signal skips the drain: cancelling the RPCs aborts
+                # their requests, so EngineCore runs out of work and exits.
+                logger.info("Received another shutdown signal; cancelling RPCs")
+                force_stops.append(loop.create_task(server.stop(grace=None)))
+                return
             logger.info("Received shutdown signal")
             stop_event.set()
 
@@ -166,10 +282,7 @@ def signal_handler():
             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")
+        await _drain_and_stop(server, async_llm, vllm_config.shutdown_timeout, inflight)
         logger.info("Shutdown complete")
 
 
diff --git a/vllm/entrypoints/launchers/run_batch.py b/vllm/entrypoints/launchers/run_batch.py
index 94c543a8ab44..fcde6fa59222 100644
--- a/vllm/entrypoints/launchers/run_batch.py
+++ b/vllm/entrypoints/launchers/run_batch.py
@@ -40,6 +40,10 @@
     ChatCompletionRequest,
     ChatCompletionResponse,
 )
+from vllm.entrypoints.openai.completion.protocol import (
+    CompletionRequest,
+    CompletionResponse,
+)
 from vllm.entrypoints.pooling.embed.protocol import (
     EmbeddingRequest,
     EmbeddingResponse,
@@ -159,6 +163,15 @@ def validate_no_file(cls, data: Any):
     | BatchTranslationRequest
 )
 
+# Build once: constructing a TypeAdapter rebuilds its core schema (~0.5 ms).
+_EMBEDDING_REQUEST_ADAPTER: TypeAdapter[EmbeddingRequest] = TypeAdapter(
+    EmbeddingRequest
+)
+_SCORE_REQUEST_ADAPTER: TypeAdapter[ScoreRequest] = TypeAdapter(ScoreRequest)
+_BATCH_REQUEST_INPUT_BODY_ADAPTER: TypeAdapter[BatchRequestInputBody] = TypeAdapter(
+    BatchRequestInputBody
+)
+
 
 class BatchRequestInput(OpenAIBaseModel):
     """The per-line object of the batch input file.
@@ -174,12 +187,14 @@ class BatchRequestInput(OpenAIBaseModel):
     # supported.
     method: str
 
-    # The OpenAI API relative URL to be used for the request. Currently
-    # /v1/chat/completions is supported.
+    # The OpenAI API relative URL to be used for the request, such as
+    # /v1/chat/completions or /v1/completions.
     url: str
 
-    # The parameters of the request.
-    body: BatchRequestInputBody
+    # The parameters of the request. CompletionRequest is kept out of
+    # BatchRequestInputBody, which an unknown URL is parsed against: its
+    # validator raises VLLMValidationError, which pydantic does not catch.
+    body: BatchRequestInputBody | CompletionRequest
 
     @field_validator("body", mode="plain")
     @classmethod
@@ -188,21 +203,24 @@ def check_type_for_url(cls, value: Any, info: ValidationInfo):
         url: str = info.data["url"]
         if url == "/v1/chat/completions":
             return ChatCompletionRequest.model_validate(value)
+        if url == "/v1/completions":
+            return CompletionRequest.model_validate(value)
         if url == "/v1/embeddings":
-            return TypeAdapter(EmbeddingRequest).validate_python(value)
+            return _EMBEDDING_REQUEST_ADAPTER.validate_python(value)
         if url.endswith("/score"):
-            return TypeAdapter(ScoreRequest).validate_python(value)
+            return _SCORE_REQUEST_ADAPTER.validate_python(value)
         if url.endswith("/rerank"):
             return RerankRequest.model_validate(value)
         if url == "/v1/audio/transcriptions":
             return BatchTranscriptionRequest.model_validate(value)
         if url == "/v1/audio/translations":
             return BatchTranslationRequest.model_validate(value)
-        return TypeAdapter(BatchRequestInputBody).validate_python(value)
+        return _BATCH_REQUEST_INPUT_BODY_ADAPTER.validate_python(value)
 
 
 AllResponse: TypeAlias = (
     ChatCompletionResponse
+    | CompletionResponse
     | EmbeddingResponse
     | ScoreResponse
     | RerankResponse
@@ -211,6 +229,10 @@ def check_type_for_url(cls, value: Any, info: ValidationInfo):
     | TranslationResponseVerbose
 )
 
+_ALL_RESPONSE_ADAPTER: TypeAdapter[AllResponse | ErrorResponse] = TypeAdapter(
+    AllResponse | ErrorResponse
+)
+
 
 class BatchResponseData(OpenAIBaseModel):
     # HTTP status code of the response.
@@ -617,9 +639,7 @@ async def run_request(
 
     if isinstance(response, JSONResponse):
         with contextlib.suppress(pydantic.ValidationError):
-            response = TypeAdapter(AllResponse | ErrorResponse).validate_python(
-                json.loads(response.body)
-            )
+            response = _ALL_RESPONSE_ADAPTER.validate_python(json.loads(response.body))
 
     if isinstance(response, AllResponse):
         batch_output = BatchRequestOutput(
@@ -743,18 +763,11 @@ async def transcription_wrapper(
                 request_dict = batch_request_body.model_dump(exclude={"file_url"})
                 request_dict["file"] = mock_file
 
+                request: TranslationRequest | TranscriptionRequest
                 if is_translation:
-                    # Create TranslationRequest from BatchTranslationRequest
-                    translation_request = TranslationRequest.model_validate(
-                        request_dict
-                    )
-                    return await handler_fn(audio_data, translation_request)
+                    request = TranslationRequest.model_validate(request_dict)
                 else:
-                    # Create TranscriptionRequest from BatchTranscriptionRequest
-                    transcription_request = TranscriptionRequest.model_validate(
-                        request_dict
-                    )
-                    return await handler_fn(audio_data, transcription_request)
+                    request = TranscriptionRequest.model_validate(request_dict)
             except Exception as e:
                 operation = "translation" if is_translation else "transcription"
                 return ErrorResponse(
@@ -765,6 +778,8 @@ async def transcription_wrapper(
                     )
                 )
 
+            return await handler_fn(audio_data, request)
+
         return transcription_wrapper
 
     return wrapper
@@ -797,6 +812,7 @@ async def build_endpoint_registry(
 
     # Get serving objects from state (defaulting to None if not set)
     openai_serving_chat = getattr(state, "openai_serving_chat", None)
+    openai_serving_completion = getattr(state, "openai_serving_completion", None)
     openai_serving_transcription = getattr(state, "openai_serving_transcription", None)
     openai_serving_translation = getattr(state, "openai_serving_translation", None)
     serving_embedding = getattr(state, "serving_embedding", None)
@@ -806,7 +822,7 @@ async def build_endpoint_registry(
 
     # Registry of endpoint configurations
     endpoint_registry: dict[str, EndpointConfig] = {
-        "completions": {
+        "chat_completions": {
             "url": "/v1/chat/completions",
             "handler_getter": lambda: (
                 openai_serving_chat.create_chat_completion
@@ -815,6 +831,15 @@ async def build_endpoint_registry(
             ),
             "wrapper_fn": None,
         },
+        "completions": {
+            "url": "/v1/completions",
+            "handler_getter": lambda: (
+                openai_serving_completion.create_completion
+                if openai_serving_completion is not None
+                else None
+            ),
+            "wrapper_fn": None,
+        },
         "embeddings": {
             "url": "/v1/embeddings",
             "handler_getter": lambda: (
@@ -941,11 +966,13 @@ async def run_one_request(
     """Route a single line of the batch to its endpoint handler."""
     request = BatchRequestInput.model_validate_json(request_json)
 
-    endpoint_key = request.url.split("/")[-1]
-
+    # Match on the whole URL: /v1/chat/completions and /v1/completions share
+    # their last segment.
     result = None
-    if endpoint_key in endpoint_registry:
-        result = handle_endpoint_request(request, endpoint_registry[endpoint_key])
+    for endpoint in endpoint_registry.values():
+        result = handle_endpoint_request(request, endpoint)
+        if result is not None:
+            break
 
     if result is None:
         result = make_async_error_request_output(
diff --git a/vllm/entrypoints/llm.py b/vllm/entrypoints/llm.py
index c8b9d2325d8b..28cf28d01656 100644
--- a/vllm/entrypoints/llm.py
+++ b/vllm/entrypoints/llm.py
@@ -108,7 +108,8 @@ class LLM(
             reserve for the model weights, activations, and KV cache. Higher
             values will increase the KV cache size and thus improve the model's
             throughput. However, if the value is too high, it may cause out-of-
-            memory (OOM) errors.
+            memory (OOM) errors. Defaults to 1.0 with the extensible KV cache
+            (which sizes the cache from measured memory) and to 0.92 otherwise.
         kv_cache_memory_bytes: 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 gpu_memory_utilization. However, users may want to
@@ -117,6 +118,10 @@ class LLM(
             compared with using gpu_memory_utilization. Note that
             kv_cache_memory_bytes (when not-None) ignores
             gpu_memory_utilization
+        enable_extensible_kv_cache: Reserve the KV cache address range with
+            device virtual memory and commit physical pages after warmup, so
+            the cache is sized from the memory actually free after CUDA graph
+            capture. Defaults to on for CUDA and ROCm with the V2 model runner.
         cpu_offload_gb: The size (GiB) of CPU memory to use for offloading
             the model weights. This virtually increases the GPU memory space
             you can use to hold the model weights, at the cost of CPU-GPU data
@@ -193,7 +198,7 @@ def __init__(
         tokenizer_revision: str | None = None,
         chat_template: Path | str | None = None,
         seed: int = 0,
-        gpu_memory_utilization: float = 0.92,
+        gpu_memory_utilization: float | None = None,
         cpu_offload_gb: float = 0,
         offload_group_size: int = 0,
         offload_num_in_group: int = 1,
@@ -213,6 +218,7 @@ def __init__(
         profiler_config: dict[str, Any] | ProfilerConfig | None = None,
         attention_config: dict[str, Any] | AttentionConfig | None = None,
         kv_cache_memory_bytes: int | None = None,
+        enable_extensible_kv_cache: bool | None = None,
         compilation_config: int | dict[str, Any] | CompilationConfig | None = None,
         quantization_config: dict[str, Any] | QuantizationConfigArgs | None = None,
         logits_processors: list[str | type[LogitsProcessor]] | None = None,
@@ -310,6 +316,7 @@ def _make_config(value: Any, cls: type[_R]) -> _R:
             seed=seed,
             gpu_memory_utilization=gpu_memory_utilization,
             kv_cache_memory_bytes=kv_cache_memory_bytes,
+            enable_extensible_kv_cache=enable_extensible_kv_cache,
             cpu_offload_gb=cpu_offload_gb,
             offload_group_size=offload_group_size,
             offload_num_in_group=offload_num_in_group,
diff --git a/vllm/entrypoints/openai/chat_completion/batch_serving.py b/vllm/entrypoints/openai/chat_completion/batch_serving.py
index 77f83ee58acd..43c07272270c 100644
--- a/vllm/entrypoints/openai/chat_completion/batch_serving.py
+++ b/vllm/entrypoints/openai/chat_completion/batch_serving.py
@@ -3,6 +3,7 @@
 
 import asyncio
 import time
+from collections import Counter
 from collections.abc import AsyncGenerator
 from http import HTTPStatus
 
@@ -16,7 +17,11 @@
     ChatCompletionResponseChoice,
     ChatMessage,
 )
-from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat
+from vllm.entrypoints.openai.chat_completion.serving import (
+    OpenAIServingChat,
+    _get_mm_token_counts,
+    _make_prompt_tokens_details,
+)
 from vllm.entrypoints.serve.engine.protocol import ErrorResponse, UsageInfo
 from vllm.entrypoints.serve.utils.api_utils import get_max_tokens
 from vllm.inputs import EngineInput
@@ -31,6 +36,11 @@
 logger = init_logger(__name__)
 
 
+def _sum_or_none(counts: list[int | None]) -> int | None:
+    """Sum per-prompt counts; ``None`` unless every prompt reported one."""
+    return None if None in counts else sum(count or 0 for count in counts)
+
+
 class OpenAIServingChatBatch(OpenAIServingChat):
     """Extends OpenAIServingChat with the /v1/chat/completions/batch endpoint.
 
@@ -206,6 +216,7 @@ async def create_batch_chat_completion(
             tokenizer,
             request_metadata,
             self.parser_cls,
+            mm_token_counts=[_get_mm_token_counts(p) for p in engine_prompts],
         )
 
     async def chat_completion_full_generator_batch(
@@ -218,6 +229,7 @@ async def chat_completion_full_generator_batch(
         tokenizer: TokenizerLike,
         request_metadata: RequestResponseMetadata,
         parser_cls: type[Parser] | None = None,
+        mm_token_counts: list[dict[str, int]] | None = None,
     ) -> ErrorResponse | ChatCompletionResponse:
         """Handle batched (non-streaming) chat completions.
 
@@ -239,6 +251,8 @@ async def chat_completion_full_generator_batch(
         choices: list[ChatCompletionResponseChoice] = []
         total_prompt_tokens = 0
         total_completion_tokens = 0
+        cached_tokens: list[int | None] = []
+        cache_creation_tokens: list[int | None] = []
 
         for prompt_idx in range(len(generators)):
             final_res = final_results.get(prompt_idx)
@@ -257,6 +271,8 @@ async def chat_completion_full_generator_batch(
             total_completion_tokens += sum(
                 len(output.token_ids) for output in final_res.outputs
             )
+            cached_tokens.append(final_res.num_cached_tokens)
+            cache_creation_tokens.append(final_res.num_cache_creation_tokens)
 
             for output in final_res.outputs:
                 self._raise_if_error(output.finish_reason, request_id)
@@ -334,10 +350,20 @@ async def chat_completion_full_generator_batch(
                 )
                 choices.append(choice_data)
 
+        mm_total: Counter[str] = Counter()
+        for counts in mm_token_counts or ():
+            mm_total.update(counts)
+
         usage = UsageInfo(
             prompt_tokens=total_prompt_tokens,
             completion_tokens=total_completion_tokens,
             total_tokens=total_prompt_tokens + total_completion_tokens,
+            prompt_tokens_details=_make_prompt_tokens_details(
+                self.enable_prompt_tokens_details,
+                _sum_or_none(cached_tokens),
+                _sum_or_none(cache_creation_tokens),
+                mm_total,
+            ),
         )
         request_metadata.final_usage_info = usage
 
diff --git a/vllm/entrypoints/openai/chat_completion/protocol.py b/vllm/entrypoints/openai/chat_completion/protocol.py
index 1341ef2b86fc..ca9e8a01a07b 100644
--- a/vllm/entrypoints/openai/chat_completion/protocol.py
+++ b/vllm/entrypoints/openai/chat_completion/protocol.py
@@ -77,7 +77,9 @@ class ChatMessage(OpenAIBaseModel):
     role: str
     content: str | None = None
     refusal: str | None = None
-    annotations: OpenAIAnnotation | None = None
+    annotations: list[OpenAIAnnotation] | None = Field(
+        default=None, exclude_if=lambda value: value is None
+    )
     audio: OpenAIChatCompletionAudio | None = None
     function_call: FunctionCall | None = None
     tool_calls: list[ToolCall] = Field(default_factory=list)
diff --git a/vllm/entrypoints/openai/responses/harmony.py b/vllm/entrypoints/openai/responses/harmony.py
index 8b2133124c87..bb57d3021f49 100644
--- a/vllm/entrypoints/openai/responses/harmony.py
+++ b/vllm/entrypoints/openai/responses/harmony.py
@@ -154,12 +154,17 @@ def response_input_to_harmony(
 ) -> Message | None:
     """Convert a single ResponseInputOutputItem into a Harmony Message.
 
-    Returns None for reasoning items with empty or absent content so
-    the caller can skip them.
+    Returns None for system/developer messages with empty content so
+    the caller can skip them. Reasoning items with empty or absent
+    content (e.g. placeholders from history replay) are converted to
+    empty assistant messages to preserve the conversation structure.
     """
     if not isinstance(response_msg, dict):
         response_msg = response_msg.model_dump()
-    if "type" not in response_msg or response_msg["type"] == "message":
+
+    item_type = response_msg.get("type")
+
+    if item_type == "message" or (item_type is None and "role" in response_msg):
         role = response_msg["role"]
         content = response_msg["content"]
         if role in ("system", "developer"):
@@ -176,7 +181,7 @@ def response_input_to_harmony(
             msg = Message.from_role_and_contents(role, contents)
         if role == "assistant":
             msg = msg.with_channel("final")
-    elif response_msg["type"] == "function_call_output":
+    elif item_type == "function_call_output":
         call_id = response_msg["call_id"]
         call_response: ResponseFunctionToolCall | None = None
         for prev_response in reversed(prev_responses):
@@ -197,22 +202,25 @@ def response_input_to_harmony(
         )
         msg = msg.with_channel("commentary")
         msg = msg.with_recipient("assistant")
-    elif response_msg["type"] == "reasoning":
+    elif item_type == "reasoning" or (item_type is None and "summary" in response_msg):
         content = response_msg.get("content")
         if content and len(content) >= 1:
             reasoning_text = "\n".join(item["text"] for item in content)
             msg = Message.from_role_and_content(Role.ASSISTANT, reasoning_text)
             msg = msg.with_channel("analysis")
         else:
-            return None
-    elif response_msg["type"] == "function_call":
+            # Reasoning item without content (e.g. placeholder from history
+            # replay) — emit an empty assistant message so the conversation
+            # structure is preserved.
+            msg = Message.from_role_and_content(Role.ASSISTANT, "")
+    elif item_type == "function_call":
         msg = Message.from_role_and_content(Role.ASSISTANT, response_msg["arguments"])
         msg = msg.with_channel("commentary")
         msg = msg.with_recipient(f"functions.{response_msg['name']}")
         msg = msg.with_content_type("json")
     else:
         raise VLLMValidationError(
-            f"Unknown input type: {response_msg['type']}",
+            f"Unknown input type: {item_type}",
             parameter="input",
         )
     return msg
diff --git a/vllm/entrypoints/openai/responses/protocol.py b/vllm/entrypoints/openai/responses/protocol.py
index 80f9bc28d3c1..0d0d5717b6d0 100644
--- a/vllm/entrypoints/openai/responses/protocol.py
+++ b/vllm/entrypoints/openai/responses/protocol.py
@@ -86,6 +86,7 @@
     RequestOutputKind,
     SamplingParams,
     StructuredOutputsParams,
+    ThinkingTokenBudget,
 )
 from vllm.utils import random_uuid
 
@@ -448,6 +449,7 @@ class ResponsesRequest(OpenAIBaseModel):
     stop: StopParam = []
     stop_token_ids: list[int] | None = []
     ignore_eos: bool = False
+    thinking_token_budget: ThinkingTokenBudget = None
     vllm_xargs: dict[str, str | int | float | list[str | int | float]] | None = Field(
         default=None,
         description=(
@@ -645,6 +647,7 @@ def to_sampling_params(
             ),
             structured_outputs=self.extract_structured_outputs(),
             logit_bias=self.logit_bias,
+            thinking_token_budget=self.thinking_token_budget,
             extra_args=extra_args,
             skip_clone=True,  # Created fresh per request, safe to skip clone
             skip_special_tokens=skip_special_tokens,
@@ -731,74 +734,106 @@ def input_item_parsing(cls, data):
 
         processed_input = []
         for item in input_data:
-            if not isinstance(item, dict):
-                processed_input.append(item)
+            normalized_item = cls._normalize_openresponses_input_item(item)
+
+            if not isinstance(normalized_item, dict):
+                processed_input.append(normalized_item)
                 continue
 
-            item_type = item.get("type")
+            item_type = normalized_item.get("type")
 
             if item_type == "function_call":
                 try:
-                    processed_input.append(ResponseFunctionToolCall(**item))
+                    processed_input.append(ResponseFunctionToolCall(**normalized_item))
                 except ValidationError:
                     logger.debug(
                         "Failed to parse function_call to ResponseFunctionToolCall, "
                         "leaving for Pydantic validation"
                     )
-                    processed_input.append(item)
+                    processed_input.append(normalized_item)
 
             elif item_type == "reasoning":
-                if "id" not in item:
-                    item = {**item, "id": f"rs_{random_uuid()}"}
                 try:
-                    processed_input.append(ResponseReasoningItem(**item))
+                    processed_input.append(ResponseReasoningItem(**normalized_item))
                 except ValidationError:
                     logger.debug(
                         "Failed to parse reasoning to ResponseReasoningItem, "
                         "leaving for Pydantic validation"
                     )
-                    processed_input.append(item)
+                    processed_input.append(normalized_item)
 
-            elif item_type == "message" and item.get("role") == "assistant":
-                content = item.get("content")
-                if not isinstance(content, list):
+            elif item_type == "message" and normalized_item.get("role") == "assistant":
+                if not isinstance(normalized_item.get("content"), list):
                     # String content is a valid EasyInputMessageParam,
                     # do not coerce it to ResponseOutputMessage
                     processed_input.append(item)
                     continue
 
-                original_item = item
-                item = dict(item)
-                if "id" not in item:
-                    item["id"] = f"msg_{random_uuid()}"
-                if "status" not in item:
-                    item["status"] = "completed"
-                # ResponseOutputText requires annotations
-                new_content = []
-                for c in content:
-                    if (
-                        isinstance(c, dict)
-                        and c.get("type") == "output_text"
-                        and "annotations" not in c
-                    ):
-                        c = {**c, "annotations": []}
-                    new_content.append(c)
-                item["content"] = new_content
                 try:
-                    processed_input.append(ResponseOutputMessage(**item))
+                    processed_input.append(ResponseOutputMessage(**normalized_item))
                 except ValidationError:
                     logger.debug(
                         "Failed to parse assistant message to ResponseOutputMessage, "
                         "leaving for Pydantic validation"
                     )
-                    processed_input.append(original_item)
+                    processed_input.append(item)
 
             else:
-                processed_input.append(item)
+                processed_input.append(normalized_item)
 
         data["input"] = processed_input
         return data
 
+    @staticmethod
+    def _normalize_openresponses_input_item(item: Any) -> Any:
+        """Normalize input items to bridge OpenResponses schema and openai-python.
+
+        The OpenResponses specification treats several fields as optional
+        (``id``, ``status`` on assistant/reasoning items; ``annotations``
+        on output_text content) that the openai-python SDK marks as
+        required.  Fill sensible defaults so payloads conforming to the
+        OpenResponses spec pass SDK validation.
+        """
+        if not isinstance(item, dict):
+            return item
+
+        item_type = item.get("type", "message")
+        should_fill = item_type == "reasoning" or (
+            item_type == "message" and item.get("role") == "assistant"
+        )
+        if not should_fill:
+            return item
+
+        normalized_item = item.copy()
+        # ``type`` is optional in the OpenResponses schema but load-bearing for
+        # SDK validation: without it a filled-in assistant item matches
+        # ``ItemReference`` in the input union and collapses to ``{"id": ...}``,
+        # silently dropping role/content. Stamp back the type we keyed off.
+        normalized_item["type"] = item_type
+        if normalized_item.get("id") in (None, ""):
+            prefix = "rs" if item_type == "reasoning" else "msg"
+            normalized_item["id"] = f"{prefix}_{random_uuid()}"
+        if normalized_item.get("status") in (None, ""):
+            normalized_item["status"] = "completed"
+        if item_type == "reasoning" and normalized_item.get("summary") is None:
+            normalized_item["summary"] = []
+
+        # OpenResponses schema treats ``annotations`` as optional for
+        # output_text content, but the openai SDK requires it.
+        if item_type == "message" and isinstance(normalized_item.get("content"), list):
+            normalized_content = []
+            for content_item in normalized_item["content"]:
+                if (
+                    isinstance(content_item, dict)
+                    and content_item.get("type") == "output_text"
+                    and "annotations" not in content_item
+                ):
+                    content_item = {**content_item, "annotations": []}
+                normalized_content.append(content_item)
+            normalized_item["content"] = normalized_content
+
+        return normalized_item
+
     @model_validator(mode="before")
     @classmethod
     def check_custom_tool_format(cls, data):
diff --git a/vllm/entrypoints/openai/responses/streaming_events.py b/vllm/entrypoints/openai/responses/streaming_events.py
index b14379cf12f8..614561ae83b3 100644
--- a/vllm/entrypoints/openai/responses/streaming_events.py
+++ b/vllm/entrypoints/openai/responses/streaming_events.py
@@ -79,6 +79,7 @@
     custom_tool_names,
     decode_custom_tool_input,
     decode_custom_tool_input_prefix,
+    extend_custom_tool_input_prefix,
     resolve_responses_tool_call_name,
 )
 from vllm.outputs import CompletionOutput
@@ -829,6 +830,9 @@ class SimpleStreamingState:
     tool_call_index: int | None = None
     tool_call_is_custom: bool = False
     tool_call_payload: str = ""
+    # End of the custom-tool arguments decoded into tool_call_payload, or -1
+    # once they need a full decode.
+    tool_call_payload_end: int = 0
     current_state: _StateType = field(default_factory=lambda: _StateType.NONE)
 
 
@@ -1062,6 +1066,7 @@ def emit_simple_tool_call_open(
     state.tool_call_index = index
     state.tool_call_is_custom = is_custom
     state.tool_call_payload = ""
+    state.tool_call_payload_end = 0
     state.accumulated_text = ""
     item: ResponseOutputItem
     if is_custom:
@@ -1114,15 +1119,27 @@ def _emit_custom_tool_input_delta(
     ]
 
 
+def _custom_tool_input_prefix(state: SimpleStreamingState) -> str:
+    if state.tool_call_payload_end >= 0:
+        result = extend_custom_tool_input_prefix(
+            state.accumulated_text,
+            state.tool_call_payload,
+            state.tool_call_payload_end,
+        )
+        if result is not None:
+            payload, state.tool_call_payload_end = result
+            return payload
+        state.tool_call_payload_end = -1
+    return decode_custom_tool_input_prefix(state.accumulated_text)
+
+
 def emit_simple_tool_call_delta(
     state: SimpleStreamingState,
     delta: str,
 ) -> list[StreamingResponsesResponse]:
     state.accumulated_text += delta
     if state.tool_call_is_custom:
-        return _emit_custom_tool_input_delta(
-            state, decode_custom_tool_input_prefix(state.accumulated_text)
-        )
+        return _emit_custom_tool_input_delta(state, _custom_tool_input_prefix(state))
     return [
         ResponseFunctionCallArgumentsDeltaEvent(
             type="response.function_call_arguments.delta",
diff --git a/vllm/entrypoints/openai/responses/utils.py b/vllm/entrypoints/openai/responses/utils.py
index d3e3d698c4fd..ae04c606d934 100644
--- a/vllm/entrypoints/openai/responses/utils.py
+++ b/vllm/entrypoints/openai/responses/utils.py
@@ -92,6 +92,51 @@ def decode_custom_tool_input_prefix(arguments: str) -> str:
     return payload
 
 
+_CUSTOM_TOOL_INPUT_HEAD = encode_custom_tool_input("")[:-2]
+
+
+def extend_custom_tool_input_prefix(
+    arguments: str, payload: str, end: int
+) -> tuple[str, int] | None:
+    """:func:`decode_custom_tool_input_prefix` of *arguments*, decoding only
+    ``arguments[end:]``. *payload* and *end* come from the previous call on a
+    prefix of *arguments* (``""`` and 0 at first). Returns None when the
+    arguments do not start as :func:`encode_custom_tool_input` writes them,
+    or once the payload string closes; only a full decode is exact then."""
+    if not end:
+        if not arguments.startswith(_CUSTOM_TOOL_INPUT_HEAD):
+            return ("", 0) if _CUSTOM_TOOL_INPUT_HEAD.startswith(arguments) else None
+        end = len(_CUSTOM_TOOL_INPUT_HEAD)
+    # partial_json_parser strips trailing whitespace before it completes the
+    # string. Decode up to the last complete escape, holding back a high
+    # surrogate so that a pair is never split between calls.
+    tail = arguments[end:].rstrip()
+    i = cut = 0
+    high_end = -1
+    while i < len(tail):
+        if tail[i] == "\\":
+            if tail[i + 1 : i + 2] == "u":
+                if "d800" <= tail[i + 2 : i + 6].lower() <= "dbff":
+                    high_end = i + 6
+                i += 6
+            else:
+                i += 2
+            if i > len(tail):
+                break
+        else:
+            i += 1
+        cut = i
+    if cut == high_end:
+        cut -= 6
+    try:
+        chunk = json.loads(f'"{tail[:cut]}"')
+    except ValueError:
+        return None
+    if chunk and "\ud800" <= chunk[-1] <= "\udbff":
+        return None
+    return payload + chunk, end + cut
+
+
 def _item_type(item: ResponseInputOutputItem) -> str | None:
     return item.get("type") if isinstance(item, dict) else getattr(item, "type", None)
 
@@ -411,7 +456,10 @@ def _construct_message_from_response_item(
             "reasoning": reasoning,
         }
     elif isinstance(item, ResponseOutputMessage):
-        output_text = item.content[0].text
+        # Assistant messages may legitimately have empty content (e.g. when the
+        # assistant emitted only a tool call). Fall back to empty string instead
+        # of crashing on item.content[0].text.
+        output_text = item.content[0].text if item.content else ""
         if prev_assistant_msg:
             previous_content = prev_assistant_msg.get("content")
             if previous_content is None:
@@ -434,20 +482,57 @@ def _construct_message_from_response_item(
             content=output.get("output"),
             tool_call_id=output.get("call_id"),
         )
+    elif isinstance(item, dict) and item.get("type") == "reasoning":
+        # Dict-form reasoning item (e.g. from TypedDict union parsing).
+        # Pydantic may leave Iterable fields as ValidatorIterator objects,
+        # so materialise them with list() before indexing.
+        reasoning_content = ""
+        content = item.get("content")
+        summary = item.get("summary")
+        if content:
+            content = list(content)
+            if len(content) >= 1:
+                text = content[0]
+                reasoning_content = (
+                    text.get("text", "") if isinstance(text, dict) else text
+                )
+        elif summary:
+            summary = list(summary)
+            if len(summary) >= 1:
+                s = summary[0]
+                reasoning_content = s.get("text", "") if isinstance(s, dict) else s
+        if prev_assistant_msg:
+            previous_reasoning = prev_assistant_msg.get("reasoning")
+            if previous_reasoning is None:
+                prev_assistant_msg["reasoning"] = reasoning_content
+                return None
+        return {
+            "role": "assistant",
+            "reasoning": reasoning_content,
+        }
     elif isinstance(item, dict) and item.get("role") == "assistant":
+        # Dict-form assistant message (e.g. from TypedDict union parsing).
         content = item.get("content")
-        text: str | None = None
-        if isinstance(content, str):
+        if content and not isinstance(content, str):
+            content = list(content)
+            if len(content) >= 1:
+                first = content[0]
+                text = first.get("text", "") if isinstance(first, dict) else first
+            else:
+                text = ""
+        elif isinstance(content, str):
             text = content
-        elif isinstance(content, list) and content:
-            text = content[0].get("text")
-        if text is not None:
-            if prev_assistant_msg:
-                previous_content = prev_assistant_msg.get("content")
-                if previous_content is None:
-                    prev_assistant_msg["content"] = text
-                    return None
-            return {"role": "assistant", "content": text}
+        else:
+            text = ""
+        if prev_assistant_msg:
+            previous_content = prev_assistant_msg.get("content")
+            if previous_content is None:
+                prev_assistant_msg["content"] = text
+                return None
+        return {
+            "role": "assistant",
+            "content": text,
+        }
     if isinstance(item, dict) and "role" in item:
         return item  # type: ignore[return-value]
     item_type = item.get("type") if isinstance(item, dict) else item.type
diff --git a/vllm/entrypoints/pooling/embed/protocol.py b/vllm/entrypoints/pooling/embed/protocol.py
index d52871d6a6f7..1b7d907d13f9 100644
--- a/vllm/entrypoints/pooling/embed/protocol.py
+++ b/vllm/entrypoints/pooling/embed/protocol.py
@@ -14,7 +14,15 @@
 
 import numpy as np
 import pybase64 as base64
-from pydantic import BaseModel, BeforeValidator, Field, model_validator
+from pydantic import (
+    BaseModel,
+    BeforeValidator,
+    Field,
+    ValidationError,
+    ValidatorFunctionWrapHandler,
+    WrapValidator,
+    model_validator,
+)
 
 from vllm import PoolingParams
 from vllm.entrypoints.serve.engine.protocol import OpenAIBaseModel, UsageInfo
@@ -156,12 +164,36 @@ def normalize_input_messages(cls, data):
         return normalized
 
 
+_CHAT_ONLY_FIELDS = frozenset(EmbeddingChatRequest.model_fields).difference(
+    EmbeddingCompletionRequest.model_fields
+)
+
+
+def _validate_embedding_request(data: Any, handler: ValidatorFunctionWrapHandler):
+    # Smart-mode union validation tries every member, and the chat members are
+    # slow to reject plain inputs, so try the completion request first when no
+    # chat member could accept the body.
+    if (
+        isinstance(data, dict)
+        and "input" in data
+        and _CHAT_ONLY_FIELDS.isdisjoint(data)
+        and not _is_chat_messages(data["input"])
+        and not _is_batched_chat_messages(data["input"])
+    ):
+        try:
+            return EmbeddingCompletionRequest.model_validate(data)
+        except ValidationError:
+            pass
+    return handler(data)
+
+
 EmbeddingRequest: TypeAlias = Annotated[
     EmbeddingCompletionRequest
     | EmbeddingChatRequest
     | EmbeddingBatchChatRequest
     | EmbeddingChatInputRequest
     | EmbeddingBatchChatInputRequest,
+    WrapValidator(_validate_embedding_request),
     BeforeValidator(reject_removed_pooling_parameters),
 ]
 
diff --git a/vllm/entrypoints/pooling/scoring/io_processor.py b/vllm/entrypoints/pooling/scoring/io_processor.py
index cd1b3df61b6a..483f0737aefb 100644
--- a/vllm/entrypoints/pooling/scoring/io_processor.py
+++ b/vllm/entrypoints/pooling/scoring/io_processor.py
@@ -8,7 +8,7 @@
 import torch.nn.functional as F
 
 from vllm import PoolingParams, PoolingRequestOutput, TokensPrompt
-from vllm.exceptions import ChatTemplateResolutionError
+from vllm.exceptions import ChatTemplateResolutionError, VLLMValidationError
 from vllm.logger import init_logger
 from vllm.renderers import TokenizeParams, merge_kwargs
 from vllm.renderers.hf import resolve_chat_template, safe_apply_chat_template
@@ -833,7 +833,17 @@ def default_tokenizer_encode():
         return full_prompt, engine_prompt
 
 
-class JinaRankingIOProcessorMixin:
+class JinaRankingIOProcessorMixin(PoolingIOProcessor):
+    def render(self, render_params: AnyRenderParam) -> PoolingEngineInput:
+        if render_params["tok_params"].truncate_prompt_tokens is not None:
+            raise VLLMValidationError(
+                "JinaForRanking does not support truncate_prompt_tokens. Use "
+                "max_tokens_per_query and max_tokens_per_doc for scoring, or "
+                "truncate each input text on the client.",
+                parameter="truncate_prompt_tokens",
+            )
+        return super().render(render_params)
+
     @staticmethod
     def format_docs_prompts_func(
         query: str,
@@ -852,9 +862,12 @@ def format_docs_prompts_func(
             special_tokens = default_special_tokens
 
         def sanitize_input(text: str) -> str:
-            for token in special_tokens.values():
-                text = text.replace(token, "")
-            return text
+            while True:
+                previous = text
+                for token in special_tokens.values():
+                    text = text.replace(token, "")
+                if text == previous:
+                    return text
 
         query = sanitize_input(query)
         docs = [sanitize_input(doc) for doc in docs]
@@ -950,10 +963,6 @@ def get_request_factory_online(
                 for q, d in zip(queries, docs)
             ]
 
-        # Forward truncation from the real request: the base factory reads
-        # these off ctx.request, so omitting them here silently drops
-        # truncate_prompt_tokens for Jina rerank/score (unlike the embed and
-        # bi/cross-encoder paths, which read them from the real request).
         ctx.request = PoolingCompletionRequest(
             task="token_embed",
             input=prompts,
diff --git a/vllm/entrypoints/scale_out/derender/serving.py b/vllm/entrypoints/scale_out/derender/serving.py
index 28468becd19e..d338d679c754 100644
--- a/vllm/entrypoints/scale_out/derender/serving.py
+++ b/vllm/entrypoints/scale_out/derender/serving.py
@@ -172,7 +172,7 @@ async def derender_chat_response(
             )
 
         try:
-            choices = await self.online_derenderer.derender_chat(
+            choices, prompt_logprobs = await self.online_derenderer.derender_chat(
                 request.generate_response,
                 request.chat_request,
                 request.prompt_token_ids,
@@ -205,7 +205,7 @@ async def derender_chat_response(
             created=int(time.time()),
             choices=choices,
             usage=usage,
-            prompt_logprobs=gen.prompt_logprobs,
+            prompt_logprobs=prompt_logprobs,
             prompt_token_ids=(
                 gen.prompt_token_ids
                 if request.chat_request is not None
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 ee0e8c162319..ebb07ee7cb92 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
@@ -7,7 +7,9 @@
 from http import HTTPStatus
 
 from fastapi import APIRouter, Depends, FastAPI, HTTPException, Request, Response
+from fastapi.exceptions import RequestValidationError
 from fastapi.responses import JSONResponse, StreamingResponse
+from pydantic import ValidationError
 
 from vllm.engine.protocol import EngineClient
 from vllm.entrypoints.serve.tokenize.serving import ServingTokenization
@@ -20,6 +22,7 @@
 
 from ...serve.engine.protocol import ErrorResponse
 from .protocol import (
+    AbortRequestsRequest,
     GenerateRequest,
     GenerateResponseBase,
 )
@@ -87,14 +90,17 @@ async def abort_requests(raw_request: Request):
             status_code=HTTPStatus.BAD_REQUEST.value,
             detail=f"JSON decode error: {e}",
         ) from e
-    request_ids = body.get("request_ids")
-    if request_ids is None:
-        raise HTTPException(
-            status_code=HTTPStatus.BAD_REQUEST.value,
-            detail="Missing 'request_ids' in request body",
-        )
+    # Validate the parsed body instead of declaring it as a FastAPI body param:
+    # FastAPI's strict content-type check would then reject clients that send
+    # no `Content-Type`, which the legacy `/abort_requests` path still allows.
+    try:
+        request = AbortRequestsRequest.model_validate(body)
+    except ValidationError as e:
+        raise RequestValidationError(
+            errors=[{**err, "loc": ("body", *err["loc"])} for err in e.errors()]
+        ) from e
     # Abort requests in background
-    asyncio.create_task(engine_client(raw_request).abort(request_ids))
+    asyncio.create_task(engine_client(raw_request).abort(request.request_ids))
     return Response(status_code=200)
 
 
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 a74a0d20e798..cfa737a214b5 100644
--- a/vllm/entrypoints/scale_out/token_in_token_out/protocol.py
+++ b/vllm/entrypoints/scale_out/token_in_token_out/protocol.py
@@ -11,12 +11,14 @@
     NonNegativeInt,
     PrivateAttr,
     Tag,
+    field_serializer,
     field_validator,
     model_validator,
 )
 
 from vllm.config import ModelConfig
 from vllm.entrypoints.generate.base.protocol import (
+    DeltaMessage,
     PerRequestMetrics,
     StreamOptions,
     validate_cache_salt,
@@ -309,8 +311,8 @@ def _require_ec_for_metadata_only(self) -> "GenerateRequest":
             "strings. 'derender' returns the parsed assistant message and "
             "needs `parse_context` when the server has a parser. 'text' and "
             "'derender' need a server that loads a tokenizer and "
-            "sampling_params.detokenize to be true. 'derender' does not "
-            "support `stream`."
+            "sampling_params.detokenize to be true. With `stream`, "
+            "'derender' returns a parsed chat `delta` per chunk."
         ),
     )
     parse_context: ParseContext | None = Field(
@@ -556,6 +558,19 @@ class GenerateTextStreamChoice(GenerateStreamChoiceBase):
     logprobs: ChatCompletionLogProbs | None = None
 
 
+class GenerateDerenderStreamChoice(GenerateStreamChoiceBase):
+    delta: DeltaMessage
+    """Parsed delta since the previous chunk of this choice."""
+
+    logprobs: ChatCompletionLogProbs | None = None
+    """Logprobs with decoded token strings and `bytes`. None when reasoning
+    is hidden by `include_reasoning=false`."""
+
+    @field_serializer("delta")
+    def _serialize_delta(self, delta: DeltaMessage) -> dict[str, Any]:
+        return delta.model_dump(exclude_none=True)
+
+
 class GenerateStreamResponseBase(BaseModel):
     request_id: str = Field(
         default_factory=lambda: f"{random_uuid()}",
@@ -581,6 +596,11 @@ class GenerateTextStreamResponse(GenerateStreamResponseBase):
     choices: list[GenerateTextStreamChoice]
 
 
+class GenerateDerenderStreamResponse(GenerateStreamResponseBase):
+    output_mode: Literal["derender"] = "derender"
+    choices: list[GenerateDerenderStreamChoice]
+
+
 class GenerateDerenderChoice(GenerateChoiceBase):
     message: ChatMessage
     """Assistant message parsed from the output text."""
@@ -657,11 +677,18 @@ def output_mode_or_tokens(value: Any) -> Any:
 ]
 GenerateStreamResponse: TypeAlias = Annotated[
     Annotated[GenerateTokensStreamResponse, Tag("tokens")]
-    | Annotated[GenerateTextStreamResponse, Tag("text")],
+    | Annotated[GenerateTextStreamResponse, Tag("text")]
+    | Annotated[GenerateDerenderStreamResponse, Tag("derender")],
     Discriminator(output_mode_or_tokens),
 ]
 
 
+class AbortRequestsRequest(BaseModel):
+    """Request for the /inference/v1/abort_requests endpoint."""
+
+    request_ids: str | list[str]
+
+
 class DerenderChatRequest(BaseModel):
     """Request for the /v1/chat/completions/derender endpoint (non-streaming).
 
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 b46720e8a29d..1f4af4ad9b83 100644
--- a/vllm/entrypoints/scale_out/token_in_token_out/serving.py
+++ b/vllm/entrypoints/scale_out/token_in_token_out/serving.py
@@ -6,7 +6,7 @@
 import itertools
 import math
 import time
-from collections.abc import AsyncGenerator
+from collections.abc import AsyncGenerator, Callable
 from collections.abc import Sequence as GenericSequence
 from typing import Any
 
@@ -15,6 +15,7 @@
 
 from vllm.engine.protocol import EngineClient
 from vllm.entrypoints.generate.base.protocol import (
+    DeltaMessage,
     PerRequestMetrics,
     RequestResponseMetadata,
 )
@@ -31,6 +32,9 @@
     ChatCompletionLogProbsContent,
     ChatCompletionRequest,
 )
+from vllm.entrypoints.openai.chat_completion.stream_processor import (
+    ChatStreamProcessor,
+)
 from vllm.entrypoints.openai.models.serving import OpenAIServingModels
 from vllm.entrypoints.serve.engine.protocol import (
     ErrorResponse,
@@ -64,11 +68,15 @@
 from .protocol import (
     GenerateDerenderChoice,
     GenerateDerenderResponse,
+    GenerateDerenderStreamChoice,
+    GenerateDerenderStreamResponse,
     GenerateLogProb,
     GenerateLogProbs,
     GenerateLogProbsContent,
     GenerateRequest,
     GenerateResponse,
+    GenerateStreamChoiceBase,
+    GenerateStreamResponseBase,
     GenerateTextChoice,
     GenerateTextResponse,
     GenerateTextStreamChoice,
@@ -77,10 +85,29 @@
     GenerateTokensResponse,
     GenerateTokensStreamChoice,
     GenerateTokensStreamResponse,
+    OutputMode,
 )
 
 logger = init_logger(__name__)
 
+_STREAM_RESPONSE_CLS: dict[OutputMode, Callable[..., GenerateStreamResponseBase]] = {
+    "tokens": GenerateTokensStreamResponse,
+    "text": GenerateTextStreamResponse,
+    "derender": GenerateDerenderStreamResponse,
+}
+
+_background_tasks: set[asyncio.Task[None]] = set()
+
+
+async def _close_result_generator(
+    result_generator: AsyncGenerator[RequestOutput, None], request_id: str
+) -> None:
+    """Close `generate()` so it aborts this request by its internal ID."""
+    try:
+        await result_generator.aclose()
+    except Exception:
+        logger.exception("Failed to abort request %s.", request_id)
+
 
 def _clamp_logprob(logprob: float) -> float:
     """The OpenAI shapes' floor for a logprob: ``-inf`` and NaN become ``-9999.0``.
@@ -240,13 +267,11 @@ async def serve_tokens(
             )
         # Built before generating so an invalid parse_context fails fast.
         derender_request: ChatCompletionRequest | None = None
+        processors: list[ChatStreamProcessor] | None = None
         if request.output_mode == "derender":
             # TODO(#57442, #60915): reject the sampled-only and top-k logprob
             # flags here once either lands.
-            if request.stream:
-                return self.create_error_response(
-                    "output_mode='derender' does not support stream=true."
-                )
+            history_tool_call_cnt = 0
             if self.online_derenderer.parser is not None:
                 if request.parse_context is None:
                     return self.create_error_response(
@@ -254,10 +279,23 @@ async def serve_tokens(
                         "server with a tool or reasoning parser. Use the one "
                         "returned by /render."
                     )
+                history_tool_call_cnt = request.parse_context.history_tool_call_cnt
                 try:
                     derender_request = request.parse_context.to_chat_request()
                 except (ValueError, VLLMValidationError) as e:
                     return self.create_error_response(e)
+            if request.stream:
+                # One parser per choice, built here so a bad parse_context
+                # fails fast.
+                try:
+                    processors = [
+                        self.online_derenderer.make_stream_processor(
+                            derender_request, history_tool_call_cnt
+                        )
+                        for _ in range(sampling_params.n)
+                    ]
+                except (ValueError, VLLMValidationError) as e:
+                    return self.create_error_response(e)
         try:
             msgspec.msgpack.encode(
                 (
@@ -407,6 +445,7 @@ async def serve_tokens(
                 request_id,
                 model_name,
                 request_metadata,
+                processors=processors,
             )
 
         return await self.serve_tokens_full_generator(
@@ -433,16 +472,7 @@ async def serve_tokens_full_generator(
         derender_mode = request.output_mode == "derender"
         text_mode = request.output_mode == "text" or derender_mode
         tokenizer = self._logprobs_tokenizer(text_mode)
-        # The derender level returns chat logprobs, so it honors
-        # logprob_token_ids and applies the chat top_logprobs cut.
-        if derender_mode:
-            num_logprobs = sampling_params.num_logprobs
-            top_limit = chat_top_logprobs_limit(
-                sampling_params.logprobs, bool(sampling_params.logprob_token_ids)
-            )
-        else:
-            num_logprobs = sampling_params.logprobs
-            top_limit = None
+        num_logprobs, top_limit = self._logprob_settings(sampling_params, derender_mode)
 
         try:
             async for res in result_generator:
@@ -650,6 +680,7 @@ async def serve_tokens_stream_generator(
         request_id: str,
         model_name: str,
         request_metadata: RequestResponseMetadata,
+        processors: list[ChatStreamProcessor] | None = None,
     ) -> AsyncGenerator[str, None]:
         num_prompt_tokens = 0
         first_iteration = True
@@ -659,8 +690,14 @@ async def serve_tokens_stream_generator(
         # With n > 1, the first output may not include every choice yet.
         num_generated_tokens = [0] * sampling_params.n
         last_res: RequestOutput | None = None
-        text_mode = request.output_mode == "text"
+        derender_mode = request.output_mode == "derender"
+        assert derender_mode == (processors is not None)
+        text_mode = request.output_mode == "text" or derender_mode
         tokenizer = self._logprobs_tokenizer(text_mode)
+        num_logprobs, top_limit = self._logprob_settings(sampling_params, derender_mode)
+        response_cls = _STREAM_RESPONSE_CLS[request.output_mode]
+        role_sent = [False] * sampling_params.n
+        finish_reason_sent = [False] * sampling_params.n
 
         include_usage, include_continuous_usage = should_include_usage(
             request.stream_options, False
@@ -691,14 +728,32 @@ async def serve_tokens_stream_generator(
                     # the finish reason, e.g. an abort, which has no new
                     # tokens. Text mode also emits text held back for stop
                     # string matching, which can arrive without token IDs.
+                    # Like coupled chat, derender keeps stepping the parser
+                    # once it has seen tokens.
                     if not (
                         delta_token_ids
                         or finish_reason is not None
                         or (text_mode and output.text)
+                        or (processors is not None and processors[i].num_output_tokens)
                     ):
                         continue
 
-                    if sampling_params.logprobs is not None:
+                    delta: DeltaMessage | None = None
+                    hide_metadata = False
+                    if processors is not None:
+                        if finish_reason_sent[i]:
+                            continue
+                        processor = processors[i]
+                        delta, finish_reason = processor.step(
+                            output.text,
+                            as_list(delta_token_ids),
+                            finish_reason,
+                            res.prompt_token_ids,
+                        )
+                        finish_reason_sent[i] = finish_reason is not None
+                        hide_metadata = processor.hide_metadata
+
+                    if num_logprobs is not None and not hide_metadata:
                         out_logprobs = output.logprobs
                         assert out_logprobs is not None, "Did not output logprobs"
                         logprobs: GenerateLogProbs | ChatCompletionLogProbs | None
@@ -707,6 +762,7 @@ async def serve_tokens_stream_generator(
                                 token_ids=delta_token_ids,
                                 top_logprobs=out_logprobs,
                                 tokenizer=tokenizer,
+                                top_limit=top_limit,
                             )
                         else:
                             logprobs = self._create_tokens_logprobs(
@@ -725,31 +781,45 @@ async def serve_tokens_stream_generator(
                     sampling_mask = None
                     if output.sampling_mask is not None:
                         sampling_mask = output.sampling_mask.token_ids
+                    token_ids_out: list[int] | None = as_list(delta_token_ids)
+                    if hide_metadata:
+                        token_ids_out = None
+                        sampling_mask = None
+
+                    if derender_mode and (
+                        delta is None
+                        and finish_reason is None
+                        and not (logprobs and logprobs.content)
+                        and not token_ids_out
+                    ):
+                        continue
+
                     choice_fields: dict[str, Any] = dict(
                         index=i,
                         logprobs=logprobs,
                         finish_reason=finish_reason,
                         stop_reason=output.stop_reason,
-                        token_ids=as_list(delta_token_ids),
+                        token_ids=token_ids_out,
                         routed_experts=routed_experts_b64,
                         sampling_mask=sampling_mask,
                         weight_versions=output.weight_versions,
                     )
-                    chunk: GenerateTokensStreamResponse | GenerateTextStreamResponse
-                    if text_mode:
-                        chunk = GenerateTextStreamResponse(
-                            request_id=request_id,
-                            choices=[
-                                GenerateTextStreamChoice(
-                                    text=output.text, **choice_fields
-                                )
-                            ],
+                    choice: GenerateStreamChoiceBase
+                    if derender_mode:
+                        delta = delta or DeltaMessage()
+                        if not role_sent[i]:
+                            delta.role = "assistant"
+                            role_sent[i] = True
+                        choice = GenerateDerenderStreamChoice(
+                            delta=delta, **choice_fields
                         )
-                    else:
-                        chunk = GenerateTokensStreamResponse(
-                            request_id=request_id,
-                            choices=[GenerateTokensStreamChoice(**choice_fields)],
+                    elif text_mode:
+                        choice = GenerateTextStreamChoice(
+                            text=output.text, **choice_fields
                         )
+                    else:
+                        choice = GenerateTokensStreamChoice(**choice_fields)
+                    chunk = response_cls(request_id=request_id, choices=[choice])
 
                     if prompt_token_ids is not None:
                         chunk.prompt_token_ids = prompt_token_ids
@@ -790,21 +860,12 @@ async def serve_tokens_stream_generator(
                         per_request_metrics = PerRequestMetrics(
                             speculative_decoding=spec_stats
                         )
-                final_chunk: GenerateTokensStreamResponse | GenerateTextStreamResponse
-                if text_mode:
-                    final_chunk = GenerateTextStreamResponse(
-                        request_id=request_id,
-                        choices=[],
-                        usage=final_usage_info,
-                        metrics=per_request_metrics,
-                    )
-                else:
-                    final_chunk = GenerateTokensStreamResponse(
-                        request_id=request_id,
-                        choices=[],
-                        usage=final_usage_info,
-                        metrics=per_request_metrics,
-                    )
+                final_chunk = response_cls(
+                    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"
 
             request_metadata.final_usage_info = final_usage_info
@@ -815,10 +876,32 @@ async def serve_tokens_stream_generator(
             )
         except Exception as e:
             logger.exception("Error in token generation stream.")
+            # Stop the engine decoding for a stream nobody reads without
+            # delaying the error event or being cancelled by a disconnect.
+            task = asyncio.create_task(
+                _close_result_generator(result_generator, request_id)
+            )
+            _background_tasks.add(task)
+            task.add_done_callback(_background_tasks.discard)
             data = self.create_streaming_error_response(e)
             yield f"data: {data}\n\n"
         yield "data: [DONE]\n\n"
 
+    @staticmethod
+    def _logprob_settings(
+        sampling_params: SamplingParams, derender_mode: bool
+    ) -> tuple[int | None, int | None]:
+        """Logprob count to gate on and the `top_logprobs` cut.
+
+        The derender level returns chat logprobs, so it honors
+        logprob_token_ids and applies the chat top_logprobs cut.
+        """
+        if derender_mode:
+            return sampling_params.num_logprobs, chat_top_logprobs_limit(
+                sampling_params.logprobs, bool(sampling_params.logprob_token_ids)
+            )
+        return sampling_params.logprobs, None
+
     def _logprobs_tokenizer(self, text_mode: bool) -> TokenizerLike | None:
         """Tokenizer that resolves logprob tokens, or None for placeholders.
 
diff --git a/vllm/entrypoints/serve/engine/protocol.py b/vllm/entrypoints/serve/engine/protocol.py
index 704d6b2455a8..ef1592199625 100644
--- a/vllm/entrypoints/serve/engine/protocol.py
+++ b/vllm/entrypoints/serve/engine/protocol.py
@@ -136,5 +136,7 @@ class UsageInfo(OpenAIBaseModel):
     prompt_tokens: int = 0
     total_tokens: int = 0
     completion_tokens: int | None = 0
-    prompt_tokens_details: PromptTokenUsageInfo | None = None
+    prompt_tokens_details: PromptTokenUsageInfo | None = Field(
+        default=None, exclude_if=lambda value: value is None
+    )
     completion_tokens_details: CompletionTokenUsageInfo | None = None
diff --git a/vllm/entrypoints/serve/exception_handling/utils.py b/vllm/entrypoints/serve/exception_handling/utils.py
index 6c1dd205167c..1c26c90cd5c2 100644
--- a/vllm/entrypoints/serve/exception_handling/utils.py
+++ b/vllm/entrypoints/serve/exception_handling/utils.py
@@ -11,5 +11,6 @@ def sanitize_message(message: str) -> str:
     message = re.sub(
         r"/(?:home|usr|opt|var|tmp|root|lib|mnt|srv)(?:/[\w.\-]+)+", "<path>", message
     )
-    message = re.sub(r"(?:/[\w\-]+)+/[\w\-]+\.\w+", "<path>", message)
+    if "." in message:
+        message = re.sub(r"(?>/[\w\-]+)+/[\w\-]+\.\w+", "<path>", message)
     return message.strip()
diff --git a/vllm/entrypoints/serve/fault_tolerance/api_router.py b/vllm/entrypoints/serve/fault_tolerance/api_router.py
index 22dad84a2626..911dffee0416 100644
--- a/vllm/entrypoints/serve/fault_tolerance/api_router.py
+++ b/vllm/entrypoints/serve/fault_tolerance/api_router.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import json
-import uuid
 from http import HTTPStatus
 
 from fastapi import APIRouter, BackgroundTasks, Depends, FastAPI, HTTPException, Request
@@ -11,31 +10,33 @@
 from vllm.entrypoints.serve.engine.protocol import ErrorResponse
 from vllm.entrypoints.serve.utils.api_utils import validate_json_request
 from vllm.logger import init_logger
-from vllm.v1.fault_tolerance.utils import FaultToleranceRequest
+from vllm.v1.fault_tolerance.utils import ALLOWED_FT_INSTRUCTIONS, FaultToleranceRequest
 
 logger = init_logger(__name__)
 
 router = APIRouter()
 
-_ALLOWED_INSTRUCTIONS = {"retry"}
 
-
-def _validate_payload(body: dict) -> tuple[str, dict]:
+def _validate_payload(body: dict) -> tuple[str, dict, str]:
     if not isinstance(body, dict):
         raise HTTPException(400, "Request body must be a JSON object.")
     instruction = body.get("instruction")
     if not instruction:
         raise HTTPException(400, "'instruction' is required.")
-    if instruction not in _ALLOWED_INSTRUCTIONS:
+    if instruction not in ALLOWED_FT_INSTRUCTIONS:
         raise HTTPException(400, f"Invalid instruction: '{instruction}'.")
     params = body.get("params", {})
     if not isinstance(params, dict):
         raise HTTPException(400, "'params' must be an object.")
-    return instruction, params
+
+    request_id = body.get("request_id")
+    if not isinstance(request_id, str) or not request_id:
+        raise HTTPException(400, "'request_id' must be a non-empty string.")
+    return instruction, params, request_id
 
 
 @router.post(
-    "/fault_tolerance/apply",
+    "/v1/fault_tolerance/apply",
     dependencies=[Depends(validate_json_request)],
     responses={
         HTTPStatus.ACCEPTED.value: {"model": dict},
@@ -50,23 +51,26 @@ async def process_fault_tolerance_instruction(
     except json.JSONDecodeError as e:
         raise HTTPException(400, "Invalid JSON format") from e
 
-    instruction, params = _validate_payload(body)
+    instruction, params, request_id = _validate_payload(body)
+    # One recovery round shares one request_id, which namespaces that round's
+    # coordination keys. The orchestrator must send the same request_id to
+    # every engine in a round; the engines verify this during recovery.
     ft_request = FaultToleranceRequest(
         instruction=instruction,
         params=params,
-        request_id=str(uuid.uuid4()),
+        request_id=request_id,
     )
 
     client: EngineClient = raw_request.app.state.engine_client
     # Recovery runs cross-rank collective ops that only complete once every rank
     # has been dispatched. Run it in the background and return immediately so the
     # orchestrator can dispatch to all ranks without blocking; completion is
-    # observed by polling GET /fault_tolerance/status.
+    # observed by polling GET /v1/fault_tolerance/status.
     background_tasks.add_task(_run_fault_recovery, client, ft_request)
     return JSONResponse(
         status_code=HTTPStatus.ACCEPTED.value,
         content={
-            "message": "Request accepted; poll /fault_tolerance/status for updates.",
+            "message": "Request accepted; poll /v1/fault_tolerance/status for updates.",
             "request_id": ft_request.request_id,
         },
         background=background_tasks,
@@ -90,7 +94,7 @@ async def _run_fault_recovery(
         )
 
 
-@router.get("/fault_tolerance/status")
+@router.get("/v1/fault_tolerance/status")
 async def get_status(raw_request: Request):
     client: EngineClient = raw_request.app.state.engine_client
     return JSONResponse(content=await client.get_status())
diff --git a/vllm/entrypoints/speech_to_text/base/serving.py b/vllm/entrypoints/speech_to_text/base/serving.py
index 871f69666f51..11a2ee05b09f 100644
--- a/vllm/entrypoints/speech_to_text/base/serving.py
+++ b/vllm/entrypoints/speech_to_text/base/serving.py
@@ -31,7 +31,7 @@
 from vllm.logger import init_logger
 from vllm.logprobs import FlatLogprobs, Logprob
 from vllm.model_executor.models import SupportsTranscription
-from vllm.multimodal.audio import get_audio_duration, split_audio
+from vllm.multimodal.audio import AudioResampler, get_audio_duration, split_audio
 from vllm.multimodal.media.audio import load_audio
 from vllm.outputs import RequestOutput
 from vllm.renderers.inputs import DictPrompt, EncoderDecoderDictPrompt
@@ -173,14 +173,20 @@ def _decode_and_chunk_speech(
         # NOTE resample to model SR here for efficiency. This is also a
         # pre-requisite for chunking, as it assumes Whisper SR.
         try:
+            resample_method = self.asr_config.audio_resample_method
             with io.BytesIO(audio_data) as buf:
                 y, sr = load_audio(
                     buf,
-                    sr=self.asr_config.sample_rate,
+                    sr=None if resample_method else self.asr_config.sample_rate,
                     mono=True,
                     max_duration_s=self.max_audio_decode_duration_s,
                     max_decode_bytes=self.max_audio_decode_bytes,
                 )
+            if resample_method:
+                y = AudioResampler(
+                    target_sr=self.asr_config.sample_rate, method=resample_method
+                ).resample(y, orig_sr=sr)
+                sr = self.asr_config.sample_rate
         except ValueError:
             raise
         except Exception as exc:
diff --git a/vllm/entrypoints/speech_to_text/transcription/protocol.py b/vllm/entrypoints/speech_to_text/transcription/protocol.py
index 7e66c512c720..a9c6f113b93e 100644
--- a/vllm/entrypoints/speech_to_text/transcription/protocol.py
+++ b/vllm/entrypoints/speech_to_text/transcription/protocol.py
@@ -208,6 +208,7 @@ def build_stt_params(
             request_prompt=self.prompt,
             to_language=self.to_language,
             hotwords=self.hotwords,
+            extra_args=self.vllm_xargs,
         )
 
     def to_beam_search_params(
diff --git a/vllm/env_override.py b/vllm/env_override.py
index bd4c7f6d1e99..f8eb4e824c1d 100644
--- a/vllm/env_override.py
+++ b/vllm/env_override.py
@@ -818,103 +818,74 @@ def _exec_then_patch(module):
 
 _patch_cpp_indirect_assert_if_needed()
 
-# ============================================================
-# Inductor FALLBACK_ALLOW_LIST fast-path for vllm::*/vllm_aiter::* ops
-# ============================================================
-# When Inductor encounters a custom op without a registered lowering or
-# decomposition (e.g. vllm::all_reduce, vllm_aiter::fused_add_rms_norm) it
-# correctly creates an implicit fallback that calls into the eager Python
-# impl. However, unless `base_name` (e.g. "vllm::all_reduce") is in
-# torch._inductor.lowering.FALLBACK_ALLOW_LIST, GraphLowering.call_function
-# (torch/_inductor/graph.py:~1283) takes the slow path that emits
+# ===================================================
+# torch < 2.13 Inductor implicit-fallback log message
+# ===================================================
+# Before torch 2.13, GraphLowering.call_function builds
 #   log.info("Creating implicit fallback for:\n%s",
 #            error.operator_str(target, args, kwargs))
-# `operator_str` eagerly recurses through __str__ on every input TensorBox;
-# for deep MoE/TP graphs (e.g. Kimi-K2.6 at TP=8) the IR provenance tree
-# behind a TP all-reduce input or a residual-fed RMSNorm input is hundreds
-# of layers deep, and stringifying it consumes many minutes of CPU per call,
-# effectively hanging compilation.
-#
-# Patching FALLBACK_ALLOW_LIST membership to also match any "vllm::*" or
-# "vllm_aiter::*" base_name routes our custom ops through the fast path
-# `make_fallback(target, warn=False, override_decomp=True)` instead. This
-# preserves all downstream behaviour (allreduce_rms_fusion still pattern-
-# matches them, partitioning still works, fallback semantics identical) but
-# skips the expensive log formatting on the FIRST encounter of each op.
-#
-# We wrap the OrderedSet in a thin proxy that:
-#   - Returns True from __contains__ for any vllm::*/vllm_aiter::* op
-#   - Otherwise delegates to the underlying set (preserving membership of
-#     the standard entries like "torchvision::roi_align", "aten::index_add")
-#   - Forwards add()/__iter__()/__len__()/etc. so other Inductor code paths
-#     that mutate or iterate the set keep working.
+# eagerly, even when INFO logging is off. operator_str stringifies every IR
+# argument and IRNode.__str__ has no depth or cycle protection, so deep IR
+# with shared nodes can effectively hang compilation (see #42129, #43111).
+# Only print the target for vllm ops. Remove once every platform pins
+# torch >= 2.13.
 
 _VLLM_FALLBACK_NAMESPACE_PREFIXES = ("vllm::", "vllm_aiter::")
 
 
-class _VllmFallbackAllowList:
-    """Membership proxy that auto-allows vllm::*/vllm_aiter::* base_names."""
-
-    _vllm_patched = True
-
-    def __init__(self, inner):
-        self._inner = inner
-
-    def __contains__(self, item):
-        if isinstance(item, str) and item.startswith(_VLLM_FALLBACK_NAMESPACE_PREFIXES):
-            return True
-        return item in self._inner
-
-    def add(self, item):
-        self._inner.add(item)
-
-    def discard(self, item):
-        self._inner.discard(item)
+def _apply_inductor_operator_str_patch() -> None:
+    from torch._inductor.exc import OperatorIssue
 
-    def __iter__(self):
-        return iter(self._inner)
+    original = OperatorIssue.operator_str
 
-    def __len__(self):
-        return len(self._inner)
+    def operator_str(target, args, kwargs) -> str:
+        name = target.name() if hasattr(target, "name") else str(target)
+        if name.startswith(_VLLM_FALLBACK_NAMESPACE_PREFIXES):
+            return f"  target: {target}"
+        return original(target, args, kwargs)
 
-    def __repr__(self):
-        return f"_VllmFallbackAllowList({self._inner!r})"
+    OperatorIssue.operator_str = staticmethod(operator_str)
 
-    def __getattr__(self, name):
-        return getattr(self._inner, name)
 
+def _patch_inductor_operator_str_if_needed():
+    """Shorten the implicit-fallback message for vllm ops on torch < 2.13.
 
-def _patch_inductor_fallback_allow_list() -> None:
-    """Wrap torch._inductor.lowering.FALLBACK_ALLOW_LIST so any custom op in
-    the ``vllm::`` or ``vllm_aiter::`` namespaces is treated as a member.
-
-    Idempotent: a sentinel attribute on the proxy prevents re-wrapping.
+    Defers application until torch._inductor.exc is imported, so that
+    importing vllm does not import torch._inductor.
     """
-    try:
-        from torch._inductor import lowering as _lowering
-    except ImportError:
+    if is_torch_equal_or_newer("2.13.0.dev"):
         return
 
-    base = getattr(_lowering, "FALLBACK_ALLOW_LIST", None)
-    if base is None or getattr(base, "_vllm_patched", False):
+    import sys
+
+    target_name = "torch._inductor.exc"
+    if target_name in sys.modules:
+        _apply_inductor_operator_str_patch()
         return
 
-    _lowering.FALLBACK_ALLOW_LIST = _VllmFallbackAllowList(base)
+    import importlib.abc
 
-    # torch/_inductor/graph.py imports the symbol at module load time:
-    #   from torch._inductor.lowering import FALLBACK_ALLOW_LIST
-    # so we also need to overwrite the local binding in the graph module if
-    # it has already been imported.
-    try:
-        from torch._inductor import graph as _graph
+    class _OperatorStrPatchFinder(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_operator_str_patch()
+
+            spec.loader.exec_module = _exec_then_patch  # type: ignore[method-assign]
+            return spec
 
-        if hasattr(_graph, "FALLBACK_ALLOW_LIST"):
-            _graph.FALLBACK_ALLOW_LIST = _lowering.FALLBACK_ALLOW_LIST
-    except ImportError:
-        pass
+    sys.meta_path.insert(0, _OperatorStrPatchFinder())
 
 
-_patch_inductor_fallback_allow_list()
+_patch_inductor_operator_str_if_needed()
 
 
 def _apply_inductor_pattern_matcher_patch() -> None:
diff --git a/vllm/envs.py b/vllm/envs.py
index a7bde6b336fa..a7a8879fd9ea 100755
--- a/vllm/envs.py
+++ b/vllm/envs.py
@@ -59,6 +59,7 @@
     VLLM_CPU_INT4_W4A8: bool = True
     VLLM_XLA_CACHE_PATH: str = os.path.join(VLLM_CACHE_ROOT, "xla_cache")
     VLLM_XLA_CHECK_RECOMPILATION: bool = False
+    VLLM_USE_FUSED_MQA_TOPK: bool = True
     VLLM_SPARSE_INDEXER_MAX_LOGITS_MB: int = 512
     VLLM_ADAPTIVE_VERIFICATION_PROFILE_CONTEXT_LEN: int = 8192
     VLLM_XPU_PP_MICROBATCH: bool = True
@@ -121,6 +122,7 @@
     VLLM_USE_BYTECODE_HOOK: bool = True
     VLLM_FORCE_AOT_LOAD: bool = False
     VLLM_USE_MEGA_AOT_ARTIFACT: bool = False
+    VLLM_COMPILE_ON_ONE_RANK: bool = False
     VLLM_USE_TRITON_AWQ: bool = False
     VLLM_FASTSAFETENSORS_QUEUE_SIZE: int = 0
     VLLM_TRITON_FORCE_FIRST_CONFIG: bool = False
@@ -240,8 +242,9 @@
     VLLM_P2P_MAX_PEERS: int = 128
     VLLM_P2P_HANDSHAKE_TIMEOUT_S: int = 30
     VLLM_P2P_IDLE_TIMEOUT_S: int = 300
-    VLLM_EC_SIDE_CHANNEL_HOST: str = "localhost"
-    VLLM_EC_SIDE_CHANNEL_PORT: int = 5601
+    VLLM_P2P_TRANSFER_STALL_TIMEOUT_S: float = 20.0
+    VLLM_EC_SIDE_CHANNEL_HOST: str = ""
+    VLLM_EC_SIDE_CHANNEL_PORT: int = 0
     VLLM_MOONCAKE_BOOTSTRAP_PORT: int = 8998
     VLLM_MOONCAKE_STORE_TIER_LOG: bool = False
     VLLM_MOONCAKE_LOAD_RECV_THREADS: int = 1
@@ -294,6 +297,8 @@
     VLLM_OBJECT_STORAGE_SHM_BUFFER_NAME: str = "VLLM_OBJECT_STORAGE_SHM_BUFFER"
     VLLM_DEEPEP_BUFFER_SIZE_MB: int = 1024
     VLLM_DEEPEP_HIGH_THROUGHPUT_FORCE_INTRA_NODE: bool = False
+    VLLM_MORI_EP_V2: bool | None = None
+    VLLM_MORI_DISPATCH_INTER_KERNEL_SWITCH_THRESHOLD: int = 256
     VLLM_DEEPEP_LOW_LATENCY_USE_MNNVL: bool = False
     VLLM_DEEPEP_V2_ALLOW_HYBRID_MODE: bool = True
     VLLM_DEEPEP_V2_PREFER_OVERLAP: bool = False
@@ -309,6 +314,7 @@
     VLLM_DEBUG_WORKSPACE: bool = False
     VLLM_DISABLE_SHARED_EXPERTS_STREAM: bool = False
     VLLM_DISABLE_DSV4_MEGAMOE_SHARED_EXPERT_FUSION: bool = False
+    VLLM_SKINNY_GEMM_MAX_WEIGHT_BYTES: int = 0
     VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD: int = 256
     VLLM_MULTI_STREAM_GEMM_TOKEN_THRESHOLD: int = 1024
     VLLM_COHERE2_ATTN_MOE_DUAL_STREAM_TOKEN_THRESHOLD: int = 256
@@ -318,6 +324,7 @@
     VLLM_DEBUG_MFU_METRICS: bool = False
     VLLM_WEIGHT_OFFLOADING_DISABLE_PIN_MEMORY: bool = False
     VLLM_WEIGHT_OFFLOADING_DISABLE_UVA: bool = False
+    VLLM_RELOAD_MODE: Literal["layerwise", "direct"] = "layerwise"
     VLLM_KV_OFFLOAD_MAX_BATCH_DESCRIPTORS: int = 0
     VLLM_WSL2_ENABLE_PIN_MEMORY: bool = False
     VLLM_DISABLE_LOG_LOGO: bool = False
@@ -402,6 +409,18 @@ def use_mega_aot_artifact():
     return os.environ.get("VLLM_USE_MEGA_AOT_ARTIFACT", default_value) == "1"
 
 
+def compile_on_one_rank() -> bool:
+    from vllm.platforms import current_platform
+    from vllm.utils.torch_utils import is_torch_equal_or_newer
+
+    if not is_torch_equal_or_newer("2.14.0.dev"):
+        return False
+    # Inductor implements the mode for CUDA/ROCm and XPU only.
+    if not (current_platform.is_cuda_alike() or current_platform.is_xpu()):
+        return False
+    return os.environ.get("VLLM_COMPILE_ON_ONE_RANK", "0") == "1"
+
+
 def env_with_choices(
     env_name: str,
     default: str | None,
@@ -807,6 +826,13 @@ def _resolve_rust_cli_path() -> str | None:
     # without re-splitting graph modules. This reduces overhead during model
     # loading by using reconstruct_serializable_fn_from_mega_artifact.
     "VLLM_USE_MEGA_AOT_ARTIFACT": use_mega_aot_artifact,
+    # Compile with torch's compile-on-one-rank mode, which resolves the device
+    # at run time instead of baking the compile-time index into the cached
+    # code, so a compile cache can be reused on any device index (e.g. across
+    # --device-ids). Off by default. Requires torch >= 2.14 on CUDA/ROCm or
+    # XPU; not yet compatible with use_inductor_graph_partition or
+    # debug_dump_path.
+    "VLLM_COMPILE_ON_ONE_RANK": compile_on_one_rank,
     # local rank of the process in the distributed setting, used to determine
     # the GPU device id
     "LOCAL_RANK": lambda: int(os.environ.get("LOCAL_RANK", "0")),
@@ -1093,6 +1119,13 @@ def _resolve_rust_cli_path() -> str | None:
     ),
     # Enable SPMD mode for TPU backend.
     "VLLM_XLA_USE_SPMD": lambda: bool(int(os.getenv("VLLM_XLA_USE_SPMD", "0"))),
+    # Use the single-launch fused FP8 MQA-logits + top-k kernel for the sparse
+    # MLA indexer prefill chunks it supports on SM100-family GPUs (see
+    # use_fused_mqa_topk / fused_mqa_topk_chunk); other chunks take the dense
+    # path. Default: on (set to 0 for the dense path).
+    "VLLM_USE_FUSED_MQA_TOPK": lambda: bool(
+        int(os.getenv("VLLM_USE_FUSED_MQA_TOPK", "1"))
+    ),
     # Maximum size (in MB) for logits tensor in sparse MLA indexer prefill chunks.
     # Bounds the [M, N] float32 logits tensor to prevent CUDA OOM.
     # Default: 512 MB
@@ -1642,7 +1675,8 @@ def _resolve_rust_cli_path() -> str | None:
     ),
     # Skip cudagraph/DP padding tokens in the MoE path by forcing their expert
     # ids to -1 so the dispatch and experts drop them. Requires a MoE kernel that
-    # treats topk_id == -1 as a skip sentinel
+    # treats topk_id == -1 as a skip sentinel. Supported CUDA grouped routers
+    # fuse this into routing; the existing backend opt-in is still required.
     "VLLM_MOE_SKIP_PADDING": lambda: bool(int(os.getenv("VLLM_MOE_SKIP_PADDING", "1"))),
     # Step5 SP dense/shared MLPs: opt into TP-sharded weights with AG/RS
     # instead of replicated weights and local-token computation.
@@ -1775,15 +1809,18 @@ def _resolve_rust_cli_path() -> str | None:
     # only when remaining peer slots are within 8 of VLLM_P2P_MAX_PEERS.
     # Default: 300. Set to 0 to disable.
     "VLLM_P2P_IDLE_TIMEOUT_S": lambda: int(os.getenv("VLLM_P2P_IDLE_TIMEOUT_S", "300")),
-    # IP address used for the EC connector's ZMQ side channel
-    # (producer ROUTER bind, consumer DEALER dial).
-    "VLLM_EC_SIDE_CHANNEL_HOST": lambda: os.getenv(
-        "VLLM_EC_SIDE_CHANNEL_HOST", "localhost"
-    ),
-    # Port for the EC connector's ZMQ side channel; advertised to peers
-    # via `ec_transfer_params.peer_port` on the producer's response.
+    # Seconds a P2P WRITE may stay in progress before the source treats its
+    # peer as faulted: no new WRITEs to it, and the session is torn down.
+    # The stalled transfer and its pins are held until it stops. Default:
+    # 20. Set to 0 to disable.
+    "VLLM_P2P_TRANSFER_STALL_TIMEOUT_S": lambda: float(
+        os.getenv("VLLM_P2P_TRANSFER_STALL_TIMEOUT_S", "20")
+    ),
+    # Optional EC control address override; otherwise use get_ip().
+    "VLLM_EC_SIDE_CHANNEL_HOST": lambda: os.getenv("VLLM_EC_SIDE_CHANNEL_HOST", ""),
+    # EC control base port plus DP index. Zero requests an OS-assigned port.
     "VLLM_EC_SIDE_CHANNEL_PORT": lambda: int(
-        os.getenv("VLLM_EC_SIDE_CHANNEL_PORT", "5601")
+        os.getenv("VLLM_EC_SIDE_CHANNEL_PORT", "0")
     ),
     # Per-request timeout (seconds) when a prefiller worker registers with
     # the Mooncake bootstrap server.
@@ -2067,6 +2104,14 @@ def _resolve_rust_cli_path() -> str | None:
     "VLLM_DEEPEP_HIGH_THROUGHPUT_FORCE_INTRA_NODE": lambda: bool(
         int(os.getenv("VLLM_DEEPEP_HIGH_THROUGHPUT_FORCE_INTRA_NODE", "0"))
     ),
+    # MoRI EPv2 for --all2all-backend mori. Unset prefers EPv2 and falls back to
+    # EPv1 where EPv2 is unsupported; 1 forces EPv2; 0 selects EPv1.
+    "VLLM_MORI_EP_V2": lambda: maybe_convert_bool(os.getenv("VLLM_MORI_EP_V2")),
+    # With --all2all-backend mori on MoRI EPv1 across nodes, use InterNodeV1LL
+    # when the per-rank token budget is at most this, else InterNodeV1.
+    "VLLM_MORI_DISPATCH_INTER_KERNEL_SWITCH_THRESHOLD": lambda: int(
+        os.getenv("VLLM_MORI_DISPATCH_INTER_KERNEL_SWITCH_THRESHOLD", "256")
+    ),
     # Allow DeepEP to use MNNVL (multi-node nvlink) for internode_ll kernel,
     # turn this for better latency on GB200 like system
     "VLLM_DEEPEP_LOW_LATENCY_USE_MNNVL": lambda: bool(
@@ -2119,7 +2164,7 @@ def _resolve_rust_cli_path() -> str | None:
     "VLLM_NCCL_INCLUDE_PATH": lambda: os.environ.get("VLLM_NCCL_INCLUDE_PATH", None),
     # GC debug config
     # - VLLM_GC_DEBUG=0: disable GC debugger
-    # - VLLM_GC_DEBUG=1: enable GC debugger with gc.collect elpased times
+    # - VLLM_GC_DEBUG=1: enable GC debugger with gc.collect elapsed times
     # - VLLM_GC_DEBUG='{"top_objects":5}': enable GC debugger with
     #                                      top 5 collected objects
     "VLLM_GC_DEBUG": lambda: os.getenv("VLLM_GC_DEBUG", ""),
@@ -2136,6 +2181,23 @@ def _resolve_rust_cli_path() -> str | None:
     "VLLM_DISABLE_DSV4_MEGAMOE_SHARED_EXPERT_FUSION": lambda: bool(
         int(os.getenv("VLLM_DISABLE_DSV4_MEGAMOE_SHARED_EXPERT_FUSION", "0"))
     ),
+    # Size cap for the lm_head skinny-GEMM (wvSplitK) weight guard: weights
+    # larger than this skip the skinny path at 1 < n <= 5 (see
+    # vllm.model_executor.layers.utils.lmhead_skinny_guard_permits).
+    # wvSplitK's cost at n activation rows scales as ~n passes over the
+    # weight, so once the weight exceeds the last-level (memory-side)
+    # cache those extra passes are HBM round trips. 0 (the default)
+    # derives the cap from that cache size once at first use (256 MiB on
+    # MI355X, 128 MiB on MI300X-class parts), with a 256 MiB fallback
+    # where the size cannot be read — see the guard for how it's read.
+    # The portable device-properties API cannot be used: it exposes only
+    # the far smaller L2 (4 MiB on gfx950), which would disable the
+    # fastpath on every real lm_head shard. Set a value explicitly to
+    # override the hardware-derived cap; set it very large (e.g. 2**60)
+    # to disable the guard.
+    "VLLM_SKINNY_GEMM_MAX_WEIGHT_BYTES": lambda: int(
+        os.getenv("VLLM_SKINNY_GEMM_MAX_WEIGHT_BYTES", "0")
+    ),
     # Limits when we run shared_experts in a separate stream.
     # We found out that for large batch sizes, the separate stream
     # execution is not beneficial (most likely because of the input clone)
@@ -2192,6 +2254,12 @@ def _resolve_rust_cli_path() -> str | None:
     "VLLM_WEIGHT_OFFLOADING_DISABLE_UVA": lambda: bool(
         int(os.getenv("VLLM_WEIGHT_OFFLOADING_DISABLE_UVA", "0"))
     ),
+    # How weight updates reload the model: "layerwise" rebuilds each layer
+    # from the checkpoint tensors; "direct" writes them into the live
+    # tensors and reruns post-load processing in place.
+    "VLLM_RELOAD_MODE": env_with_choices(
+        "VLLM_RELOAD_MODE", "layerwise", ["layerwise", "direct"]
+    ),
     # Max descriptors per CPU-KV-offload batch-memcpy call. 0 = platform default
     # (ROCm chunks at 8192, since hipMemcpyBatchAsync faults above that on
     # rocm 7.14/7.15 when batch copy is optimized; CUDA uncapped). Set >0 to override.
@@ -2463,6 +2531,7 @@ def compile_factors() -> dict[str, object]:
         "VLLM_CPU_MOE_PREPACK",
         "VLLM_ZENTORCH_WEIGHT_PREPACK",
         "VLLM_TEST_FORCE_LOAD_FORMAT",
+        "VLLM_RELOAD_MODE",
         "VLLM_ENABLE_CUDA_COMPATIBILITY",
         "VLLM_CUDA_COMPATIBILITY_PATH",
         "VLLM_SKIP_MODEL_NAME_VALIDATION",
diff --git a/vllm/exceptions.py b/vllm/exceptions.py
index 85863c2fbf1d..96415f47b4fd 100644
--- a/vllm/exceptions.py
+++ b/vllm/exceptions.py
@@ -198,3 +198,18 @@ def __init__(self):
             "starting another session.",
             HTTPStatus.CONFLICT,
         )
+
+
+class EngineFaultedError(GracefulHTTPError):
+    """Raised when the engine has faulted and is awaiting FT recovery.
+
+    Returns HTTP 503 (Service Unavailable) so that load balancers and
+    client SDKs retry the request on a different instance.
+    """
+
+    def __init__(self):
+        super().__init__(
+            "The engine has faulted and cannot accept new requests. "
+            "Please try again later or on a different instance.",
+            HTTPStatus.SERVICE_UNAVAILABLE,
+        )
diff --git a/vllm/kernels/triton/fused_pre_fc_norms.py b/vllm/kernels/triton/fused_pre_fc_norms.py
new file mode 100644
index 000000000000..453b06f6a45e
--- /dev/null
+++ b/vllm/kernels/triton/fused_pre_fc_norms.py
@@ -0,0 +1,73 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import torch
+
+from vllm.triton_utils import tl, triton
+
+
+@triton.jit
+def _gemma_rms_norm(x, w, eps, HIDDEN_SIZE: tl.constexpr):
+    x = x.to(tl.float32)
+    mean_sq = tl.sum(x * x, axis=0) / HIDDEN_SIZE
+    rrms = tl.rsqrt(mean_sq + eps)
+    w = w.to(tl.float32) + 1.0
+    return (x * rrms) * w
+
+
+@triton.jit
+def _fused_pre_fc_norms_kernel(
+    embeds_ptr,
+    embeds_stride,
+    hidden_ptr,
+    hidden_stride,
+    emb_w_ptr,
+    hid_w_ptr,
+    eps,
+    out_ptr,
+    out_stride,
+    H: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    tok = tl.program_id(0)
+    off = tl.arange(0, BLOCK)
+    mask = off < H
+
+    e = tl.load(embeds_ptr + tok * embeds_stride + off, mask=mask, other=0.0)
+    ew = tl.load(emb_w_ptr + off, mask=mask)
+    e_normed = _gemma_rms_norm(e, ew, eps, H)
+    tl.store(out_ptr + tok * out_stride + off, e_normed, mask=mask)
+
+    h = tl.load(hidden_ptr + tok * hidden_stride + off, mask=mask, other=0.0)
+    hw = tl.load(hid_w_ptr + off, mask=mask)
+    h_normed = _gemma_rms_norm(h, hw, eps, H)
+    tl.store(out_ptr + tok * out_stride + H + off, h_normed, mask=mask)
+
+
+def fused_pre_fc_norms(
+    inputs_embeds: torch.Tensor,
+    hidden_states: torch.Tensor,
+    emb_weight: torch.Tensor,
+    hid_weight: torch.Tensor,
+    eps: float,
+) -> torch.Tensor:
+    """Returns cat([norm_e(inputs_embeds), norm_h(hidden_states)]) -> [N, 2H].
+
+    Both norms are Gemma-style RMSNorm (scale by ``1 + weight``, fp32
+    accumulation, multiply in fp32 before the cast back to the input dtype).
+    """
+    n, h = inputs_embeds.shape
+    out = torch.empty(n, 2 * h, dtype=inputs_embeds.dtype, device=inputs_embeds.device)
+    _fused_pre_fc_norms_kernel[(n,)](
+        inputs_embeds,
+        inputs_embeds.stride(0),
+        hidden_states,
+        hidden_states.stride(0),
+        emb_weight,
+        hid_weight,
+        eps,
+        out,
+        out.stride(0),
+        h,
+        triton.next_power_of_2(h),
+    )
+    return out
diff --git a/vllm/kernels/triton/fused_sigmoid_gate_mul.py b/vllm/kernels/triton/fused_sigmoid_gate_mul.py
new file mode 100644
index 000000000000..def949e0cd3d
--- /dev/null
+++ b/vllm/kernels/triton/fused_sigmoid_gate_mul.py
@@ -0,0 +1,109 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Fused sigmoid-gate-multiply Triton kernel: ``out = x * sigmoid(gate)``."""
+
+import torch
+
+from vllm.platforms import current_platform
+from vllm.triton_utils import tl, triton
+
+
+@triton.jit
+def _fused_sigmoid_gate_mul_kernel(
+    x_ptr,
+    gate_ptr,
+    out_ptr,
+    scale_ptr,
+    hidden,
+    x_stride_t,
+    gate_stride_t,
+    out_stride_t,
+    BLOCK: tl.constexpr,
+    QUANT_FP8: tl.constexpr,
+):
+    row = tl.program_id(0).to(tl.int64)
+    offs = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK)
+    mask = offs < hidden
+
+    x = tl.load(x_ptr + row * x_stride_t + offs, mask=mask, other=0.0)
+    g = tl.load(gate_ptr + row * gate_stride_t + offs, mask=mask, other=0.0)
+    out = x.to(tl.float32) * tl.sigmoid(g.to(tl.float32))
+    if QUANT_FP8:
+        # Static FP8 from the rounded product, like scaled_fp8_quant.
+        out = out.to(x_ptr.dtype.element_ty).to(tl.float32)
+        out = tl.clamp(out * tl.math.div_rn(1.0, tl.load(scale_ptr)), -448.0, 448.0)
+    tl.store(
+        out_ptr + row * out_stride_t + offs,
+        out.to(out_ptr.dtype.element_ty),
+        mask=mask,
+    )
+
+
+def fused_sigmoid_gate_mul(
+    x: torch.Tensor, gate: torch.Tensor, scale: torch.Tensor | None = None
+) -> torch.Tensor:
+    """Compute ``x * sigmoid(gate)`` in a single Triton kernel.
+
+    Falls back to the eager ``x * torch.sigmoid(gate)`` sequence unless both
+    inputs are same-shape, same-dtype, contiguous 2D CUDA tensors.
+
+    Args:
+        x: (num_tokens, hidden) tensor, e.g. attention output.
+        gate: (num_tokens, hidden) pre-sigmoid gate.
+        scale: optional static per-tensor FP8 scale. When given, the result is
+            quantized to float8_e4m3fn; ineligible inputs take the eager
+            sequence followed by scaled_fp8_quant.
+
+    Returns:
+        A fresh (num_tokens, hidden) tensor with x's dtype (or FP8) and device.
+
+    """
+    if (
+        x.ndim != 2
+        or x.shape != gate.shape
+        or x.dtype != gate.dtype
+        or not x.is_contiguous()
+        or not gate.is_contiguous()
+        or not x.is_cuda
+        # The FP8 output is hardcoded to e4m3fn; ROCm GEMMs take fnuz.
+        or (scale is not None and not current_platform.is_cuda())
+    ):
+        out = x * torch.sigmoid(gate)
+        if scale is None:
+            return out
+        # Strided gate view (num_heads == 1 per TP rank): quantize the eager
+        # result, matching the unfused chain.
+        from vllm import _custom_ops as ops
+
+        return ops.scaled_fp8_quant(out, scale, group_shape=(-1, -1))[0]
+
+    num_tokens, hidden = x.shape
+    out = torch.empty_like(x, dtype=x.dtype if scale is None else torch.float8_e4m3fn)
+    if num_tokens == 0 or hidden == 0:
+        return out
+
+    if hidden <= 1024:
+        block = triton.next_power_of_2(hidden)
+        num_warps = 4
+    elif hidden <= 4096:
+        block = 1024
+        num_warps = 4
+    else:
+        block = 2048
+        num_warps = 8
+
+    grid = (num_tokens, triton.cdiv(hidden, block))
+    _fused_sigmoid_gate_mul_kernel[grid](
+        x,
+        gate,
+        out,
+        scale if scale is not None else x,
+        hidden,
+        x.stride(0),
+        gate.stride(0),
+        out.stride(0),
+        BLOCK=block,
+        QUANT_FP8=scale is not None,
+        num_warps=num_warps,
+    )
+    return out
diff --git a/vllm/lora/layers/fused_moe.py b/vllm/lora/layers/fused_moe.py
index 19a1128ce145..aa20290330c2 100644
--- a/vllm/lora/layers/fused_moe.py
+++ b/vllm/lora/layers/fused_moe.py
@@ -233,6 +233,15 @@ def _create_lora_b_weights(self, max_loras: int, lora_config: LoRAConfig):
 
     def _ep_check(self):
         if self.use_ep:
+            placement = self.base_layer.expert_placement_strategy
+            if placement != "linear":
+                raise ValueError(
+                    "Fused MoE LoRA with EP requires linear expert placement: "
+                    "adapter loading assumes a contiguous range of experts "
+                    f"on each rank, got '{placement}'. "
+                    "Set expert_placement_strategy='linear' to use MoE LoRA "
+                    "with expert parallelism."
+                )
             moe_config = self.moe_config
             all2all_backend = moe_config.moe_parallel_config.all2all_backend
             assert all2all_backend == "allgather_reducescatter", (
diff --git a/vllm/lora/lora_model.py b/vllm/lora/lora_model.py
index f4520b12c75b..886d285d9ea7 100644
--- a/vllm/lora/lora_model.py
+++ b/vllm/lora/lora_model.py
@@ -207,6 +207,7 @@ def from_local_checkpoint(
         weights_mapper: WeightsMapper | None = None,
         skip_prefixes: list[str] | None = None,
         moe_ep_spec: MoEEPLoadSpec | None = None,
+        unsupported_lora_modules: set[str] | None = None,
     ) -> "LoRAModel":
         """Create a LoRAModel from a local checkpoint.
 
@@ -231,6 +232,9 @@ def from_local_checkpoint(
                 slicing metadata shared across all MoE layers. Non-local
                 expert weights are skipped at read time instead of being
                 loaded and discarded later.
+            unsupported_lora_modules: Full names of model layers that match
+                expected_lora_modules but are not wrapped by LoRA, such as
+                BERT's pooler.dense. LoRA weights on them are unexpected.
 
         Returns:
             Loaded LoRA Model.
@@ -270,6 +274,9 @@ def check_unexpected_modules(modules: dict):
                 ):
                     unexpected_modules.append(module_name)
 
+                elif module_name in (unsupported_lora_modules or ()):
+                    unexpected_modules.append(f"{module_name} (not a LoRA layer)")
+
             if unexpected_modules:
                 raise ValueError(
                     f"While loading {lora_dir}, expected"
diff --git a/vllm/lora/model_manager.py b/vllm/lora/model_manager.py
index bf270d6d5667..92470952e774 100644
--- a/vllm/lora/model_manager.py
+++ b/vllm/lora/model_manager.py
@@ -482,6 +482,9 @@ def _parent_module(module_name: str) -> str:
             return module_name.rpartition(".")[0]
 
         wrapped_by_id: dict[int, BaseLayerWithLoRA] = {}
+        # Layers that share a suffix with a LoRA layer but are not wrapped,
+        # such as BERT's pooler.dense. Adapters with LoRA on them are refused.
+        self.unsupported_lora_modules: set[str] = set()
 
         for module_name, module in self.model.named_modules(remove_duplicate=False):
             if isinstance(module, PPMissingLayer):
@@ -495,6 +498,11 @@ def _parent_module(module_name: str) -> str:
                     module,
                     self.supported_lora_modules,
                 ):
+                    if (
+                        module_name.rsplit(".", 1)[-1] in self.supported_lora_modules
+                        and self._get_punica_wrapper(module_name) is not None
+                    ):
+                        self.unsupported_lora_modules.add(module_name)
                     continue
             elif not self._match_target_modules(module_name) and not is_classifier_head:
                 # Full classification heads bypass the A/B LoRA target filter.
@@ -618,7 +626,7 @@ def _parent_module(module_name: str) -> str:
                 continue
             self.register_module(module_name, new_module)
 
-            self._register_packed_modules(module_name)
+            self._register_packed_modules(module_name, module)
             # All lora layers share the same punica_wrapper based on reference.
             new_module.set_mapping(punica_wrapper)
 
@@ -833,10 +841,17 @@ def _get_punica_wrapper(self, module_name: str) -> PunicaWrapperBase | None:
 
         return None
 
-    def _register_packed_modules(self, module_full_name: str) -> None:
+    def _register_packed_modules(
+        self, module_full_name: str, module: nn.Module
+    ) -> None:
         parts = module_full_name.split(".")
         module_name = parts[-1]
-        replacements = self.packed_modules_mapping.get(module_name, [])
+        # A base layer may set `lora_packed_modules` to override the leaf-name
+        # mapping for that layer only, e.g. Gemma4 k_eq_v layers pack
+        # [q_proj, k_proj, k_proj].
+        replacements = getattr(
+            module, "lora_packed_modules", None
+        ) or self.packed_modules_mapping.get(module_name, [])
         # When replacements is less than or equal to 1, it indicates that this
         # module is not a packed module.
         if len(replacements) <= 1:
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 717c5edfdf1b..faa4cdea75eb 100644
--- a/vllm/lora/ops/triton_ops/fused_moe_lora_op.py
+++ b/vllm/lora/ops/triton_ops/fused_moe_lora_op.py
@@ -13,7 +13,7 @@
 from vllm.utils.mem_utils import get_max_shared_memory_bytes
 from vllm.utils.torch_utils import async_tensor_h2d, direct_register_custom_op
 
-from .utils import supports_pdl, supports_tma
+from .utils import FUSED_MOE_LORA_TWO_KERNEL_MAX_TOKENS, supports_pdl, supports_tma
 
 
 @triton.jit
@@ -1242,7 +1242,9 @@ def _fused_moe_lora_shrink(
     a_desc = None
     b_desc = None
     if use_tma and num_slices == 1:
-        b_desc = triton.tools.tensor_descriptor.TensorDescriptor.from_tensor(
+        from triton.tools.tensor_descriptor import TensorDescriptor
+
+        b_desc = TensorDescriptor.from_tensor(
             lora_a_stacked[0],
             [1, 1, shrink_config["BLOCK_SIZE_N"], shrink_config["BLOCK_SIZE_K"]],
         )
@@ -1372,13 +1374,15 @@ def _fused_moe_lora_expand(
     a_desc = None
     b_desc = None
     if use_tma:
+        from triton.tools.tensor_descriptor import TensorDescriptor
+
         if sorted_token_ids is not None:
-            a_desc = triton.tools.tensor_descriptor.TensorDescriptor.from_tensor(
+            a_desc = TensorDescriptor.from_tensor(
                 a_intermediate_cache1,
                 [expand_config["BLOCK_SIZE_M"], expand_config["BLOCK_SIZE_K"]],
             )
         if num_slices == 1:
-            b_desc = triton.tools.tensor_descriptor.TensorDescriptor.from_tensor(
+            b_desc = TensorDescriptor.from_tensor(
                 lora_b_stacked[0],
                 [1, 1, expand_config["BLOCK_SIZE_N"], expand_config["BLOCK_SIZE_K"]],
             )
@@ -1492,7 +1496,11 @@ def _fused_moe_lora(
     assert top_k_num == topk_weights.shape[1]
 
     # Fast path: single fused kernel
-    if not fully_sharded and max_lora_rank <= _FUSED_MOE_LORA_ONE_SHOT_MAX_RANK:
+    if (
+        not fully_sharded
+        and max_lora_rank <= _FUSED_MOE_LORA_ONE_SHOT_MAX_RANK
+        and topk_weights.shape[0] > FUSED_MOE_LORA_TWO_KERNEL_MAX_TOKENS
+    ):
         M_pairs = topk_weights.numel()
         if (
             sorted_token_ids is None
@@ -1560,7 +1568,11 @@ def _fused_moe_lora(
 
     # TMA is not currently compatiple with fully_sharded due to the non-determinism
     # of token id sorting across ranks.
-    use_tma = supports_tma(device) and not fully_sharded
+    use_tma = (
+        supports_tma(device)
+        and not fully_sharded
+        and max_lora_rank * output.element_size() % 16 == 0
+    )
 
     intermediate_cache_shape = (
         num_slices,
diff --git a/vllm/lora/ops/triton_ops/utils.py b/vllm/lora/ops/triton_ops/utils.py
index fbcef91002d2..143d502ce5c3 100644
--- a/vllm/lora/ops/triton_ops/utils.py
+++ b/vllm/lora/ops/triton_ops/utils.py
@@ -312,6 +312,9 @@ def get_lora_op_configs(
     return config_data
 
 
+FUSED_MOE_LORA_TWO_KERNEL_MAX_TOKENS = 64
+
+
 @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."""
diff --git a/vllm/lora/peft_helper.py b/vllm/lora/peft_helper.py
index 721b7a039f00..459c36d2d8a5 100644
--- a/vllm/lora/peft_helper.py
+++ b/vllm/lora/peft_helper.py
@@ -31,6 +31,7 @@ class PEFTHelper:
 
     bias: Literal["none"] = field(default="none")
     modules_to_save: list[str] | None = field(default=None)
+    task_type: str | None = field(default=None)
     # True to use Rank-Stabilized LoRA (rsLoRA, see: https://arxiv.org/abs/2312.03732)
     use_rslora: bool = field(default=False)
     # True to use Weight-Decomposed Low-Rank Adaptation (DoRA, see: https://arxiv.org/abs/2402.09353)
@@ -110,6 +111,12 @@ def __post_init__(self):
                     f"LoRA rank for `{key}` in `rank_pattern` must be a positive "
                     f"integer, got {rank}."
                 )
+        # PEFT adds the classification head to modules_to_save whenever it
+        # loads a SEQ_CLS or TOKEN_CLS adapter, even if the config omits it.
+        if self.task_type in ("SEQ_CLS", "TOKEN_CLS"):
+            self.modules_to_save = list(
+                dict.fromkeys([*(self.modules_to_save or []), "classifier", "score"])
+            )
         if self.use_rslora:
             logger.info_once("Loading LoRA weights trained with rsLoRA.")
         self.vllm_lora_scaling_factor = self._scaling(self.r, self.lora_alpha)
diff --git a/vllm/lora/request.py b/vllm/lora/request.py
index 36c6881270cc..ee5288833421 100644
--- a/vllm/lora/request.py
+++ b/vllm/lora/request.py
@@ -54,16 +54,15 @@ def path(self):
         return self.lora_path
 
     def __eq__(self, value: object) -> bool:
-        """Overrides the equality method to compare LoRARequest
-        instances based on lora_name. This allows for identification
-        and comparison lora adapter across engines.
+        """Compares lora_name and lora_int_id. A name loaded again at runtime
+        gets a new lora_int_id while requests for the old one may still be
+        running, and the engine needs both adapters active.
         """
-        return isinstance(value, self.__class__) and self.lora_name == value.lora_name
+        return (
+            isinstance(value, self.__class__)
+            and self.lora_name == value.lora_name
+            and self.lora_int_id == value.lora_int_id
+        )
 
     def __hash__(self) -> int:
-        """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.
-        """
-        return hash(self.lora_name)
+        return hash((self.lora_name, self.lora_int_id))
diff --git a/vllm/lora/worker_manager.py b/vllm/lora/worker_manager.py
index 7e3d485fabf8..a229718fe216 100644
--- a/vllm/lora/worker_manager.py
+++ b/vllm/lora/worker_manager.py
@@ -148,6 +148,7 @@ def _load_adapter(self, lora_request: LoRARequest) -> LoRAModel:
                 weights_mapper=hf_to_vllm_mapper,
                 skip_prefixes=lora_skip_prefixes,
                 moe_ep_spec=self._adapter_manager.moe_ep_load_spec,
+                unsupported_lora_modules=self._adapter_manager.unsupported_lora_modules,
             )
             # Stamp the on-disk MoE layout onto the loaded model so the
             # adapter manager can route 3D-format checkpoints through the
diff --git a/vllm/model_executor/hw_agnostic/__init__.py b/vllm/model_executor/hw_agnostic/__init__.py
index 208f01a7cb5e..a2700e8ce619 100644
--- a/vllm/model_executor/hw_agnostic/__init__.py
+++ b/vllm/model_executor/hw_agnostic/__init__.py
@@ -1,2 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from vllm.model_executor.hw_agnostic._resolve import resolve
+
+__all__ = ["resolve"]
diff --git a/vllm/model_executor/hw_agnostic/_resolve.py b/vllm/model_executor/hw_agnostic/_resolve.py
new file mode 100644
index 000000000000..db963da09090
--- /dev/null
+++ b/vllm/model_executor/hw_agnostic/_resolve.py
@@ -0,0 +1,76 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import importlib
+from types import ModuleType
+from typing import Any
+
+import vllm.envs as envs
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+_HW_AGNOSTIC_PKG = "vllm.model_executor.hw_agnostic.layers"
+_IN_TREE_PKG = "vllm.model_executor.layers"
+
+
+def _import_hw_agnostic(module: str) -> ModuleType | None:
+    """Import `hw_agnostic.layers.<module>`, or return None if there is none.
+
+    Only a missing hw-agnostic module counts as absent; any other import error,
+    e.g. a missing dependency of an existing port, propagates.
+    """
+    name = f"{_HW_AGNOSTIC_PKG}.{module}"
+    try:
+        return importlib.import_module(name)
+    except ModuleNotFoundError as e:
+        if e.name is not None and (name == e.name or name.startswith(f"{e.name}.")):
+            return None
+        raise
+
+
+def resolve(module: str, name: str) -> Any:
+    """Return layer `name` from `module` for the active path.
+
+    Returns `vllm.model_executor.hw_agnostic.layers.<module>.<name>` if
+    `VLLM_USE_HW_AGNOSTIC` is set and that layer exists, else
+    `vllm.model_executor.layers.<module>.<name>`. Falling back to the in-tree
+    layer while `VLLM_USE_HW_AGNOSTIC` is set logs a warning.
+
+    The Transformers backend builds its layers through this function, so an
+    out-of-tree plugin that subclasses the result and registers it with its
+    `register_oot` overrides the class that backend builds on either path.
+    Models with a native vLLM implementation always build the in-tree class, and
+    overrides do not yet reach the backend's norms, which are subclasses of
+    `RMSNorm` and `GemmaRMSNorm`.
+
+    Example:
+        SiluAndMul = hw_agnostic.resolve("activation", "SiluAndMul")
+
+        @SiluAndMul.register_oot
+        class MySiluAndMul(SiluAndMul): ...
+
+    Args:
+        module: Module path relative to `vllm.model_executor.layers`, e.g.
+            `"layernorm"`.
+        name: Name of the layer in that module, e.g. `"RMSNorm"`.
+
+    Returns:
+        The hw-agnostic or the in-tree object named `name`.
+
+    Raises:
+        AttributeError: `name` exists in neither module.
+        ImportError: The hw-agnostic module exists but fails to import, or the
+            in-tree module does not exist.
+
+    """
+    if envs.VLLM_USE_HW_AGNOSTIC:
+        hw_module = _import_hw_agnostic(module)
+        if (layer := getattr(hw_module, name, None)) is not None:
+            logger.info_once("Using hardware agnostic layer %s.%s", module, name)
+            return layer
+        logger.warning_once(
+            "hw-agnostic layer %s.%s is not available; using the in-tree layer",
+            module,
+            name,
+        )
+    return getattr(importlib.import_module(f"{_IN_TREE_PKG}.{module}"), name)
diff --git a/vllm/model_executor/hw_agnostic/custom_op.py b/vllm/model_executor/hw_agnostic/custom_op.py
index ecb37681f267..7468519d370a 100644
--- a/vllm/model_executor/hw_agnostic/custom_op.py
+++ b/vllm/model_executor/hw_agnostic/custom_op.py
@@ -29,6 +29,45 @@ def maybe_get_oot_by_class(class_type: type) -> type:
     return class_type
 
 
+def _class_to_instantiate(cls: type, name: str) -> type:
+    """The out-of-tree override registered for `cls` as `name`, else `cls`.
+
+    The in-tree and the hw-agnostic path keep separate registries, both keyed on
+    the bare class name, so a class can be registered in the hw-agnostic path
+    but not in the in-tree one, and vice versa. This catches the two cases where
+    an override written against the in-tree class would go unnoticed:
+
+    * Registered in the hw-agnostic path, but not derived from `cls`: Python
+      would skip its `__init__`, so this raises.
+    * Registered in the in-tree path, but not in the hw-agnostic one: it would
+      never be used, so this warns.
+    """
+    override = op_registry_oot.get(name)
+    if override is None:
+        from vllm.model_executor.custom_op import op_registry_oot as in_tree_oot
+
+        if name in in_tree_oot:
+            logger.warning_once(
+                "Out-of-tree override %s is registered for the in-tree %s, so the "
+                "hw-agnostic path does not use it. To opt in, subclass "
+                "hw_agnostic.resolve(<module>, %r) and register it with that "
+                "class's register_oot.",
+                in_tree_oot[name].__name__,
+                name,
+                name,
+            )
+        return cls
+    if not issubclass(override, cls):
+        raise TypeError(
+            f"Out-of-tree override {override.__name__!r} for {name!r} does not "
+            f"derive from {cls.__module__}.{name}, the class the hw-agnostic path "
+            f"instantiates. Subclass hw_agnostic.resolve(<module>, {name!r}) "
+            f"instead."
+        )
+    logger.debug("Instantiating %s using %s", name, override)
+    return override
+
+
 class PluggableLayer(nn.Module):
     """Base class for pluggable layers.
 
@@ -53,16 +92,7 @@ def __new__(cls, *args, **kwargs):
                 f"@PluggableLayer.register, or it's the PluggableLayer itself."
             ) from None
 
-        if layer_class_name not in op_registry_oot:
-            layer_cls_to_instantiate = cls
-        else:
-            layer_cls_to_instantiate = op_registry_oot[layer_class_name]
-            logger.debug(
-                "Instantiating pluggable layer: %s using %s",
-                layer_class_name,
-                str(layer_cls_to_instantiate),
-            )
-        return super().__new__(layer_cls_to_instantiate)
+        return super().__new__(_class_to_instantiate(cls, layer_class_name))
 
     # Decorator to register pluggable layers.
     @classmethod
@@ -114,16 +144,7 @@ def __new__(cls, *args, **kwargs):
                 f"@CustomOp.register, or it's the CustomOp base class itself."
             ) from None
 
-        if op_name not in op_registry_oot:
-            op_cls_to_instantiate = cls
-        else:
-            op_cls_to_instantiate = op_registry_oot[op_name]
-            logger.debug(
-                "Instantiating custom op: %s using %s",
-                op_name,
-                str(op_cls_to_instantiate),
-            )
-        return super().__new__(op_cls_to_instantiate)
+        return super().__new__(_class_to_instantiate(cls, op_name))
 
     def __init__(self, *, enforce_enable: bool = False, compile_native: bool = False):
         super().__init__()
diff --git a/vllm/model_executor/kernels/attention/dsa/candidate_blocks.py b/vllm/model_executor/kernels/attention/dsa/candidate_blocks.py
index a822b0ff0bcd..b27fcf0f1bc9 100644
--- a/vllm/model_executor/kernels/attention/dsa/candidate_blocks.py
+++ b/vllm/model_executor/kernels/attention/dsa/candidate_blocks.py
@@ -162,6 +162,7 @@ def select_candidate_blocks(
     block_size: int,
     out: torch.Tensor,
     row_repeat: int = 1,
+    score_dtype: torch.dtype | None = None,
 ) -> None:
     """Select local block IDs by maximum score, pinning each row's newest block.
 
@@ -176,20 +177,21 @@ def select_candidate_blocks(
         out.fill_(-1)
         return
     nblocks = triton.cdiv(width, block_size)
+    score_dtype = score_dtype or logits.dtype
     # Logits are max_model_len wide (CUDA-graph static) but rows are shorter:
     # DeepSelect stops at each row's valid block count. Order does not matter.
     use_deep_select = has_deep_select() and nblocks >= topk_blocks
     if use_deep_select:
         in_align, out_align = get_deep_select_stride_requirement()
-        padded = round_up(nblocks, in_align // logits.element_size())
-        scores = logits.new_empty((rows, padded))[:, :nblocks]
+        padded = round_up(nblocks, in_align // score_dtype.itemsize)
+        scores = logits.new_empty((rows, padded), dtype=score_dtype)[:, :nblocks]
         use_deep_select = (
             out.stride(1) == 1
             and out.stride(0) * out.element_size() % out_align == 0
             and is_deep_select_supported(scores, topk_blocks)
         )
     else:
-        scores = logits.new_empty((rows, nblocks))
+        scores = logits.new_empty((rows, nblocks), dtype=score_dtype)
     valid_blocks = torch.empty(rows, dtype=torch.int32, device=logits.device)
     _block_scores_kernel[(rows, 32)](
         logits,
diff --git a/vllm/model_executor/kernels/attention/dsa/fused_mqa_topk.py b/vllm/model_executor/kernels/attention/dsa/fused_mqa_topk.py
new file mode 100644
index 000000000000..928e222498a7
--- /dev/null
+++ b/vllm/model_executor/kernels/attention/dsa/fused_mqa_topk.py
@@ -0,0 +1,184 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Single-launch fused FP8 MQA-logits + top-k prefill selection (SM100 family)"""
+
+import functools
+
+import torch
+
+import vllm.envs as envs
+from vllm.logger import init_logger
+from vllm.platforms import current_platform
+
+logger = init_logger(__name__)
+
+# Current restrictions of the fused_mqa_topk kernel:
+FUSED_MQA_TOPK_HEADS = 32
+FUSED_MQA_TOPK_HEAD_DIM = 128
+FUSED_MQA_TOPK_TOPKS = (512, 1024, 2048)
+# Most gathered keys (all requests of a chunk) one call accepts: candidates
+# store the key index in 20 bits (kMaxKeys in csrc/fused_mqa_topk/constants.h).
+# Chunks with more keys take the dense path; the planner never packs past it.
+FUSED_MQA_TOPK_MAX_KEYS = 1 << 20
+# Rows per fused_mqa_topk call. The kernel has no row limit, but prefill chunks above
+# 16384 rows (--max-num-batched-tokens > 16384) are not validated end to end yet,
+# so the chunk planner splits fused_mqa_topk chunks at this many rows.
+FUSED_MQA_TOPK_MAX_ROWS = 16384
+
+
+def fused_mqa_topk_min_keys(topk: int) -> int:
+    """Longest-row key count from which fused_mqa_topk beats the dense path
+    on SM100f (shorter chunks stay dense): 32K, or 64K for top-2048."""
+    return 1 << 16 if topk > 1024 else 1 << 15
+
+
+# Mirrors csrc/fused_mqa_topk/constants.h (the op checks the workspace shapes).
+_CANDIDATES_PER_ROW = 16384  # kCandidatesPerRow
+_RETAINED_PER_ROW = 4096  # kRetainedPerRow
+_SLOTS_PER_SM = 2  # candidate slots per persistent CTA
+_BLOCK_ROWS = 8  # kBlockRows: query rows per block, one block per CTA at a time
+
+
+@functools.cache
+def _sm_count(index: int) -> int:
+    return torch.cuda.get_device_properties(index).multi_processor_count
+
+
+@functools.cache
+def fused_mqa_topk_available() -> bool:
+    """True when the fused_mqa_topk op is compiled in and the GPU is SM100-family."""
+    return (
+        current_platform.is_cuda()
+        and current_platform.is_device_capability_family(100)
+        and hasattr(torch.ops._C, "fused_mqa_topk_prefill")
+    )
+
+
+def use_fused_mqa_topk(
+    *,
+    query_shape: tuple[int, ...],
+    topk: int,
+    fp4_cache: bool,
+    full_context: bool,
+    candidate_selection: bool = False,
+) -> bool:
+    """Whether an indexer may use the fused_mqa_topk kernel: the single
+    eligibility rule of the prefill chunk planner and the indexer layers (each
+    prefill chunk then also needs fused_mqa_topk_chunk).
+
+    query_shape: one token's indexer queries, (heads, head_dim) = (32, 128);
+    topk: the per-row top-k (pools for kpool indexers) in {512, 1024, 2048};
+    fp4_cache: the MXFP4 indexer cache (and its FP4 queries) stays dense;
+    full_context: every rank ranks its rows over the whole context (no DCP, or
+    PCP after the KV all-gather); pure DCP needs the per-shard top-k merge;
+    candidate_selection: v4.1 two-level selection needs the logits.
+    VLLM_BATCH_INVARIANT keeps the dense path: the fused gate depends on the
+    chunk's total keys, and ties at the k-th score are taken in atomic order.
+    """
+    return (
+        envs.VLLM_USE_FUSED_MQA_TOPK
+        and not envs.VLLM_BATCH_INVARIANT
+        and not fp4_cache
+        and full_context
+        and not candidate_selection
+        and tuple(query_shape) == (FUSED_MQA_TOPK_HEADS, FUSED_MQA_TOPK_HEAD_DIM)
+        and topk in FUSED_MQA_TOPK_TOPKS
+        and fused_mqa_topk_available()
+    )
+
+
+def fused_mqa_topk_chunk(
+    max_seqlen_k: int, keys: int, topk: int, rows: int, dense_rows: int
+) -> bool:
+    """Whether a prefill chunk of a fused_mqa_topk-eligible layer takes the
+    fused_mqa_topk kernel: its longest row sees at least fused_mqa_topk_min_keys
+    keys (the dense path is as fast below), its gathered keys fit the kernel,
+    and its rows keep the SMs busy.
+
+    Each SM scans every key of one 8-row block at a time, so the blocks run in
+    waves of one per SM. A last wave that fills less than half the SMs costs a
+    whole wave and makes the kernel slower than dense, unless dense needs more
+    than one logits call (rows > dense_rows, its rows per call).
+    """
+    sms = _sm_count(torch.accelerator.current_device_index())
+    last_wave = (-(-rows // _BLOCK_ROWS) - 1) % sms + 1  # blocks, 1..sms
+    return (
+        max_seqlen_k >= fused_mqa_topk_min_keys(topk)
+        and keys <= FUSED_MQA_TOPK_MAX_KEYS
+        and (rows > dense_rows or 2 * last_wave >= sms)
+    )
+
+
+def get_fused_mqa_topk_workspaces(
+    device: torch.device, rows: int
+) -> list[torch.Tensor]:
+    """The workspaces of one fused_mqa_topk call over ``rows`` query rows, in the
+    op's argument order (see fused_mqa_topk.cu), all fresh: the zeroed block
+    counter, then the scratch."""
+    index = (
+        device.index
+        if device.index is not None
+        else torch.accelerator.current_device_index()
+    )
+    blocks = -(-rows // _BLOCK_ROWS)
+    # One persistent CTA per SM, at most one per block; S = 2 slots per CTA.
+    slot_rows = _BLOCK_ROWS * _SLOTS_PER_SM * min(blocks, _sm_count(index))
+    empty = functools.partial(torch.empty, device=device)
+    return [
+        torch.zeros(2, dtype=torch.int32, device=device),  # block counter
+        empty((slot_rows, _CANDIDATES_PER_ROW), dtype=torch.uint16),  # fp24 lows
+        empty((slot_rows, _CANDIDATES_PER_ROW), dtype=torch.int32),  # indices
+        empty((slot_rows, _RETAINED_PER_ROW), dtype=torch.float32),  # retained
+        empty((slot_rows, _RETAINED_PER_ROW), dtype=torch.int32),  # their indices
+        empty(blocks * _BLOCK_ROWS, dtype=torch.int32),  # per-row candidate counts
+        empty(blocks, dtype=torch.int32),  # blocks queued for the exact rescan
+    ]
+
+
+def fused_mqa_topk_key_scales(scales: torch.Tensor) -> torch.Tensor:
+    """Per-key fp32 scales from the gathered indexer K-cache scales.
+
+    scales: [keys, 4] uint8. SM100 caches store each power-of-two scale as four
+    equal UE8M0 bytes; the fp32 value 2**(e - 127) has bit pattern ``e << 23``.
+    """
+    exponents = scales.view(torch.int32).squeeze(-1) & 0xFF
+    return (exponents << 23).view(torch.float32)
+
+
+def fused_mqa_topk_prefill(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    scales: torch.Tensor,
+    weights: torch.Tensor,
+    starts: torch.Tensor,
+    ends: torch.Tensor,
+    output: torch.Tensor,
+) -> None:
+    """Fused replacement of ``fp8_fp4_mqa_logits`` + ``top_k_per_row_prefill``.
+
+    q: [rows, 32, 128] fp8 e4m3; k: [keys, 128] fp8 e4m3 (keys <=
+    FUSED_MQA_TOPK_MAX_KEYS); scales: [keys] fp32; weights: [rows, 32] fp32;
+    starts/ends: [rows] int32, the visible key range ``[start, end)`` of each
+    row; output: [rows, k] int32 with k (the top-k) 512, 1024 or 2048. All
+    contiguous; q, k, scales, weights and output 16-byte aligned (the op checks
+    the contract).
+
+    Writes the same convention as ``ops.top_k_per_row_prefill``: indices
+    relative to the row's ``start`` (i.e. ``key - start``), unordered, ``-1``
+    padded when fewer than k keys are visible.
+    """
+    logger.info_once(
+        "Using the single-launch fused MQA+top-k prefill kernel "
+        "(FP8, 32 heads, top-k %d).",
+        output.shape[1],
+    )
+    torch.ops._C.fused_mqa_topk_prefill(
+        q,
+        k,
+        scales,
+        weights,
+        starts,
+        ends,
+        output,
+        *get_fused_mqa_topk_workspaces(q.device, q.shape[0]),
+    )
diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py
index b4d7bd3ef669..eda81ddbfea3 100644
--- a/vllm/model_executor/kernels/linear/__init__.py
+++ b/vllm/model_executor/kernels/linear/__init__.py
@@ -836,6 +836,9 @@ def fp8_scaled_mm_is_batch_invariant(
     kernel: FP8ScaledMMLinearKernel | Fp8BlockScaledMMLinearKernel,
 ) -> bool:
     """Whether a row's output from the kernel is independent of the other rows."""
+    # Weight-only here, and Humming honors VLLM_BATCH_INVARIANT itself.
+    if isinstance(kernel, HummingFP8ScaledMMLinearKernel):
+        return True
     if not isinstance(
         kernel, (CutlassFP8ScaledMMLinearKernel, TritonFP8ScaledMMLinearKernel)
     ):
@@ -1199,41 +1202,34 @@ def init_nvfp4_linear_kernel(use_a16: bool = False) -> NvFp4LinearKernel:
         HummingNvFp4LinearKernel,
     )
 
-    # VLLM_BATCH_INVARIANT forces deterministic execution. Prefer the
-    # batch-invariant CUTLASS implementation when available, otherwise fall
-    # back to emulation. It overrides --linear-backend.
+    # VLLM_BATCH_INVARIANT forces deterministic execution: Humming for W4A16,
+    # otherwise CUTLASS, or emulation where CUTLASS is unsupported. It
+    # overrides --linear-backend.
     force_kernel: type[NvFp4LinearKernel] | None = None
     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:
-            if linear_backend not in ("auto", "cutlass"):
-                logger.warning_once(
-                    "VLLM_BATCH_INVARIANT overrides --linear-backend=%s; "
-                    "using the CUTLASS backend for deterministic execution.",
-                    linear_backend,
-                )
+        if use_a16:
+            force_kernel, bi_backend = HummingNvFp4LinearKernel, "humming"
+        else:
+            cutlass_supported, reason = CutlassNvFp4LinearKernel.is_supported()
+            if cutlass_supported:
+                force_kernel, bi_backend = CutlassNvFp4LinearKernel, "cutlass"
             else:
                 logger.info_once(
-                    "VLLM_BATCH_INVARIANT forces NVFP4 linear to use the "
-                    "CUTLASS backend for deterministic execution."
-                )
-            force_kernel = CutlassNvFp4LinearKernel
-        else:
-            if linear_backend not in ("auto", "emulation"):
-                logger.warning_once(
-                    "VLLM_BATCH_INVARIANT overrides --linear-backend=%s; "
-                    "using the emulation backend for deterministic execution.",
-                    linear_backend,
+                    "VLLM_BATCH_INVARIANT is set but the batch-invariant NVFP4 "
+                    "kernel is not supported on this platform; falling back to "
+                    "emulation for deterministic execution. Reason: %s",
+                    reason,
                 )
-            logger.info_once(
-                "VLLM_BATCH_INVARIANT is set but the batch-invariant NVFP4 "
-                "kernel is not supported on this platform; falling back to "
-                "emulation for deterministic execution. Reason: %s",
-                reason,
+                force_kernel, bi_backend = EmulationNvFp4LinearKernel, "emulation"
+        if linear_backend not in ("auto", bi_backend):
+            logger.warning_once(
+                "VLLM_BATCH_INVARIANT overrides --linear-backend=%s; using the "
+                "%s backend for deterministic execution.",
+                linear_backend,
+                bi_backend,
             )
-            force_kernel = EmulationNvFp4LinearKernel
     elif linear_backend == "auto" and use_a16:
         _cc = current_platform.get_device_capability()
         compute_capability = _cc.to_int() if _cc is not None else None
diff --git a/vllm/model_executor/kernels/linear/cute_dsl/gemm_rs_ar.py b/vllm/model_executor/kernels/linear/cute_dsl/gemm_rs_ar.py
index ec40fb056a13..68d057a2bb79 100644
--- a/vllm/model_executor/kernels/linear/cute_dsl/gemm_rs_ar.py
+++ b/vllm/model_executor/kernels/linear/cute_dsl/gemm_rs_ar.py
@@ -41,7 +41,12 @@
 from vllm.cute_utils import _tcgen05, mbarrier, simple_tma_copy, to_cta0_smem
 from vllm.distributed import get_tp_group
 from vllm.logger import init_logger
+from vllm.model_executor.layers.fusion.quant_activation import (
+    QuantizedActivation,
+    as_quantized_activation,
+)
 from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
+from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp8Dynamic
 from vllm.platforms import current_platform
 
 if TYPE_CHECKING:
@@ -984,7 +989,9 @@ def should_run(self, x: torch.Tensor) -> bool:
         # supported but faster on the existing LL path.
         return x.shape[0] >= 128
 
-    def apply(self, x: torch.Tensor, linear: LinearBase) -> torch.Tensor:
+    def apply(
+        self, x: torch.Tensor | QuantizedActivation, linear: LinearBase
+    ) -> torch.Tensor:
         from vllm.model_executor.kernels.linear.mxfp8.flashinfer import (
             FlashInferCutedslMxfp8LinearKernel,
         )
@@ -998,16 +1005,23 @@ def apply(self, x: torch.Tensor, linear: LinearBase) -> torch.Tensor:
             w = w.t()
         w_sf = getattr(linear, "weight_scale", None)
 
+        # A pre-quantized activation carries F8_128x4-swizzled MXFP8 scales,
+        # the layout the in-kernel quantization below produces.
+        qa = as_quantized_activation(x, kMxfp8Dynamic)
+        if qa is not None:
+            assert qa.orig_dtype == torch.bfloat16
+            x, x_sf = qa.data, qa.scale
+        else:
+            assert isinstance(x, torch.Tensor) and x.dtype == torch.bfloat16
+            x_sf = None
         assert x.ndim == 2
         M, K = x.shape
         assert 0 < M <= self.max_M
-        assert x.dtype == torch.bfloat16
         assert x.device == self.device
         assert x.is_contiguous()
         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,
@@ -1019,13 +1033,14 @@ def apply(self, x: torch.Tensor, linear: LinearBase) -> torch.Tensor:
             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)
+            if x_sf is None:
+                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 x_sf is None and w_sf is None
             assert w.shape == (N, K) and K % 64 == 0
             assert w.dtype == torch.bfloat16
             assert w.is_contiguous()
diff --git a/vllm/model_executor/kernels/linear/cute_dsl/skinny_gemm_plans.py b/vllm/model_executor/kernels/linear/cute_dsl/skinny_gemm_plans.py
new file mode 100644
index 000000000000..d0a51b70a9cc
--- /dev/null
+++ b/vllm/model_executor/kernels/linear/cute_dsl/skinny_gemm_plans.py
@@ -0,0 +1,78 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Measured skinny GEMM plans for unquantized BF16 linear layers.
+
+Plans are keyed by device capability, local ``(N, K)`` weight shape and token
+count. A layer whose shape is listed runs the planned config for that token
+count; other token counts use ``F.linear``. Entries are kept only where the
+skinny GEMM beats cuBLAS by more than 5% under CUDA graph replay.
+"""
+
+from vllm.model_executor.kernels.linear.cute_dsl.skinny_gemm import (
+    SkinnyGemmConfig,
+    shape_dynamic_skinny_gemm,
+)
+from vllm.platforms import current_platform
+
+SkinnyGemmPlan = dict[int, SkinnyGemmConfig]
+
+# DGX Spark (GB10), measured on GLM-5.3-Flash local shapes at TP=2.
+SM121_SKINNY_GEMM_PLANS: dict[tuple[int, int], SkinnyGemmPlan] = {
+    # Shared-expert fused gate/up and MLA fused q_a/kv_a projections.
+    (2048, 4096): {
+        1: SkinnyGemmConfig(1, 32, 1, k_unroll=4, static_k=4096),
+        2: SkinnyGemmConfig(2, 64, 1, k_unroll=4, vector_width=4, static_k=4096),
+        4: SkinnyGemmConfig(4, 128, 1, vector_width=2, static_k=4096),
+    },
+    # Shared-expert down projection.
+    (4096, 1024): {
+        1: SkinnyGemmConfig(1, 64, 1, vector_width=2, static_k=1024),
+        2: SkinnyGemmConfig(2, 64, 1, vector_width=2, static_k=1024),
+        4: SkinnyGemmConfig(4, 64, 1, vector_width=2, static_k=1024),
+    },
+    # KDA output projection.
+    (4096, 4096): {
+        1: SkinnyGemmConfig(1, 32, 2, k_unroll=4, vector_width=4, static_k=4096),
+        2: SkinnyGemmConfig(2, 64, 1, k_unroll=4, vector_width=2, static_k=4096),
+        4: SkinnyGemmConfig(4, 64, 1, vector_width=2, static_k=4096),
+    },
+    # MLA output projection.
+    (4096, 8192): {
+        2: SkinnyGemmConfig(2, 64, 2, k_unroll=2, vector_width=2, static_k=8192),
+        4: SkinnyGemmConfig(4, 256, 1, k_unroll=2, static_k=8192),
+    },
+    # MLA q_b projection.
+    (8192, 1536): {
+        1: SkinnyGemmConfig(1, 32, 1, k_unroll=4, vector_width=2, static_k=1536),
+        2: SkinnyGemmConfig(2, 32, 1, k_unroll=2, vector_width=2, static_k=1536),
+        4: SkinnyGemmConfig(4, 32, 1, k_unroll=2, vector_width=2, static_k=1536),
+    },
+    # KDA fused QKV/beta/f_a/g_a input projection.
+    (12576, 4096): {
+        1: SkinnyGemmConfig(1, 256, 1, k_unroll=2, vector_width=2),
+        2: SkinnyGemmConfig(2, 256, 1),
+        4: SkinnyGemmConfig(4, 128, 1, k_unroll=4, vector_width=4, static_k=4096),
+    },
+    # LM head.
+    (77440, 4096): {
+        1: SkinnyGemmConfig(1, 256, 1, static_k=4096),
+        2: SkinnyGemmConfig(2, 256, 1, k_unroll=2),
+        4: SkinnyGemmConfig(4, 256, 1, k_unroll=2),
+    },
+}
+
+SKINNY_GEMM_PLANS_BY_CAPABILITY: dict[
+    tuple[int, int] | None, dict[tuple[int, int], SkinnyGemmPlan]
+] = {
+    (12, 1): SM121_SKINNY_GEMM_PLANS,
+}
+
+
+def get_skinny_gemm_plan(n: int, k: int) -> SkinnyGemmPlan | None:
+    """Return the token-count plan for an ``(N, K)`` weight on this device."""
+    plans = SKINNY_GEMM_PLANS_BY_CAPABILITY.get(
+        current_platform.get_device_capability()
+    )
+    if plans is None or not shape_dynamic_skinny_gemm.is_available():
+        return None
+    return plans.get((n, k))
diff --git a/vllm/model_executor/kernels/linear/mixed_precision/MPLinearKernel.py b/vllm/model_executor/kernels/linear/mixed_precision/MPLinearKernel.py
index 2657c2a39cf8..056a89c43348 100644
--- a/vllm/model_executor/kernels/linear/mixed_precision/MPLinearKernel.py
+++ b/vllm/model_executor/kernels/linear/mixed_precision/MPLinearKernel.py
@@ -8,6 +8,7 @@
 import torch
 
 from vllm.model_executor.layers.quantization.utils import replace_parameter
+from vllm.model_executor.utils import retire_parameter
 from vllm.scalar_type import ScalarType
 
 
@@ -69,6 +70,13 @@ def apply_weights(
     ) -> torch.Tensor:
         raise NotImplementedError
 
+    def _retire_unused_zp(self, layer: torch.nn.Module) -> None:
+        name = self.w_zp_name
+        if name is not None and getattr(layer, name, None) is not None:
+            # Loaded but unused: keep the param and its data.
+            retire_parameter(layer, name)
+        self.w_zp_name = None
+
     def _transform_param(
         self, layer: torch.nn.Module, name: str | None, fn: Callable
     ) -> None:
diff --git a/vllm/model_executor/kernels/linear/mixed_precision/conch.py b/vllm/model_executor/kernels/linear/mixed_precision/conch.py
index 7f9910a76a23..efb1473a9945 100644
--- a/vllm/model_executor/kernels/linear/mixed_precision/conch.py
+++ b/vllm/model_executor/kernels/linear/mixed_precision/conch.py
@@ -113,8 +113,8 @@ def transform_w_zp(x):
         self._transform_param(layer, self.w_s_name, transform_w_s)
         if self.config.zero_points:
             self._transform_param(layer, self.w_zp_name, transform_w_zp)
-        elif self.w_zp_name is not None:
-            layer.register_parameter(self.w_zp_name, None)
+        else:
+            self._retire_unused_zp(layer)
 
     def apply_weights(
         self,
diff --git a/vllm/model_executor/kernels/linear/mixed_precision/cpu.py b/vllm/model_executor/kernels/linear/mixed_precision/cpu.py
index e2f8dbac7fd9..b869e5cc85ef 100644
--- a/vllm/model_executor/kernels/linear/mixed_precision/cpu.py
+++ b/vllm/model_executor/kernels/linear/mixed_precision/cpu.py
@@ -101,7 +101,7 @@ def _process_gptq_weights_w4a8(self, layer: torch.nn.Module):
             packed_zp = getattr(layer, self.w_zp_name)
         else:
             # w4a8 kernel always requires zp, allocate a fake zp
-            assert self.w_zp_name
+            self.w_zp_name = "w_zp"
             replace_parameter(
                 layer,
                 self.w_zp_name,
@@ -138,13 +138,13 @@ def _process_gptq_weights_w4a8(self, layer: torch.nn.Module):
         if layer.bias is not None:
             layer.bias.data = layer.bias.float()
 
-        packed_weight.data = blocked_w
-        scales.data = blocked_s
-        packed_zp.data = blocked_zp
+        replace_parameter(layer, self.w_q_name, blocked_w)
+        replace_parameter(layer, self.w_s_name, blocked_s)
+        replace_parameter(layer, self.w_zp_name, blocked_zp)
 
     def process_weights_after_loading(self, layer: torch.nn.Module):
-        if (not self.config.zero_points) and (self.w_zp_name is not None):
-            setattr(layer, self.w_zp_name, None)
+        if not self.config.zero_points:
+            self._retire_unused_zp(layer)
 
         weights = getattr(layer, self.w_q_name)
         # Require GPTQ pack format
diff --git a/vllm/model_executor/kernels/linear/mixed_precision/cutlass.py b/vllm/model_executor/kernels/linear/mixed_precision/cutlass.py
index ea867cafd063..64c920ef53c4 100644
--- a/vllm/model_executor/kernels/linear/mixed_precision/cutlass.py
+++ b/vllm/model_executor/kernels/linear/mixed_precision/cutlass.py
@@ -12,6 +12,7 @@
     convert_packed_uint4b8_to_signed_int4_inplace,
 )
 from vllm.model_executor.parameter import BasevLLMParameter, permute_param_layout_
+from vllm.model_executor.utils import replace_parameter
 from vllm.platforms import current_platform
 from vllm.scalar_type import scalar_types
 
@@ -88,7 +89,7 @@ def transform_w_s(x):
 
         w_s = getattr(layer, self.w_s_name)
         fp8_scales, chan_scales = convert_bf16_scales_to_fp8(self.quant_fp8, w_s.data)
-        w_s.data = fp8_scales
+        replace_parameter(layer, self.w_s_name, fp8_scales)
 
         # register per-channel scales
         layer.register_parameter(
diff --git a/vllm/model_executor/kernels/linear/mixed_precision/dynamic_4bit.py b/vllm/model_executor/kernels/linear/mixed_precision/dynamic_4bit.py
index d0515027628e..0867968360a4 100644
--- a/vllm/model_executor/kernels/linear/mixed_precision/dynamic_4bit.py
+++ b/vllm/model_executor/kernels/linear/mixed_precision/dynamic_4bit.py
@@ -5,6 +5,7 @@
 import torch
 
 from vllm.model_executor.layers.quantization.utils import replace_parameter
+from vllm.model_executor.utils import retire_parameter
 from vllm.platforms import CpuArchEnum, current_platform
 from vllm.scalar_type import scalar_types
 
@@ -111,7 +112,8 @@ def process_weights_after_loading(self, layer: torch.nn.Module):
         replace_parameter(
             layer, self.w_q_name, torch.nn.Parameter(w, requires_grad=False)
         )
-        setattr(layer, self.w_s_name, None)
+        # Scales are packed into the weight; the loaded param keeps its data.
+        retire_parameter(layer, self.w_s_name)
 
     def apply_weights(
         self,
diff --git a/vllm/model_executor/kernels/linear/mixed_precision/marlin.py b/vllm/model_executor/kernels/linear/mixed_precision/marlin.py
index 60687fc8602d..1aabf7abd164 100644
--- a/vllm/model_executor/kernels/linear/mixed_precision/marlin.py
+++ b/vllm/model_executor/kernels/linear/mixed_precision/marlin.py
@@ -21,6 +21,7 @@
     unpack_cols,
 )
 from vllm.model_executor.parameter import BasevLLMParameter, permute_param_layout_
+from vllm.model_executor.utils import replace_parameter
 from vllm.platforms import current_platform
 from vllm.scalar_type import scalar_types
 
@@ -169,13 +170,18 @@ def transform_w_s(x):
                 ),
             )
         else:
-            setattr(layer, self.w_zp_name, marlin_make_empty(device))
+            # Marlin takes an empty zp; a loaded unused one keeps its data.
+            self._retire_unused_zp(layer)
+            self.w_zp_name = "w_zp"
+            replace_parameter(layer, self.w_zp_name, marlin_make_empty(device))
         self._transform_param(layer, self.w_q_name, transform_w_q)
         self._transform_param(layer, self.w_s_name, transform_w_s)
 
         if hasattr(layer, "bias") and layer.bias is not None:
-            layer.bias.data = marlin_permute_bias(
-                marlin_pad_dim(layer.bias, size_n, padded_n)
+            replace_parameter(
+                layer,
+                "bias",
+                marlin_permute_bias(marlin_pad_dim(layer.bias, size_n, padded_n)),
             )
 
     def apply_weights(
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 62c8038d3fba..a5c89238d3cb 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
@@ -84,7 +84,7 @@ def _fits_in_memory_budget(nbytes: int, device: torch.device) -> bool:
     from vllm.v1.worker.gpu.attn_utils import get_kv_cache_spec
 
     vllm_config = get_current_vllm_config()
-    util = vllm_config.cache_config.gpu_memory_utilization
+    util = vllm_config.cache_config.resolved_gpu_memory_utilization
     min_kv_bytes = max_memory_usage_bytes(
         vllm_config, get_kv_cache_spec(vllm_config).values()
     )
@@ -198,7 +198,8 @@ def _triton_w4a16_skinny_fmt_kernel(
     pid_m = tl.program_id(0)
     pid_n = tl.program_id(1)
 
-    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    # int64 rows: offs_m * stride_am and offs_m * N pass 2**31 for large M
+    offs_m = pid_m.to(tl.int64) * BLOCK_M + tl.arange(0, BLOCK_M)
     offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
 
     # ExLlama unshuffle shifts: shift[j] = (j//2)*4 + (j%2)*16
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 ed6c67aa99c5..6b8d7b390b64 100644
--- a/vllm/model_executor/kernels/linear/mixed_precision/triton_w4a16.py
+++ b/vllm/model_executor/kernels/linear/mixed_precision/triton_w4a16.py
@@ -79,7 +79,8 @@ def triton_w4a16_gemm_kernel(
     pid_n = tl.program_id(1)
 
     # Row/col offsets for this tile
-    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    # int64 rows: offs_m * stride_am and offs_m * stride_cm pass 2**31 for large M
+    offs_m = pid_m.to(tl.int64) * BLOCK_M + tl.arange(0, BLOCK_M)
     offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
 
     # b/zeros are stored with N packed: N//8 int32 columns per K row
diff --git a/vllm/model_executor/kernels/linear/mixed_precision/xpu.py b/vllm/model_executor/kernels/linear/mixed_precision/xpu.py
index 3cfbba243ab3..fa8f4afbb90d 100644
--- a/vllm/model_executor/kernels/linear/mixed_precision/xpu.py
+++ b/vllm/model_executor/kernels/linear/mixed_precision/xpu.py
@@ -181,10 +181,6 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
             torch.nn.Parameter(packed, requires_grad=False),
         )
 
-        # Free the original unpacked int8 weight (still registered as "weight")
-        # to avoid double-storing both int8 [N, K] and int32 [N, K/8] in memory.
-        delattr(layer, "weight")
-
     def apply_weights(
         self,
         layer: torch.nn.Module,
diff --git a/vllm/model_executor/kernels/linear/mixed_precision/zentorch.py b/vllm/model_executor/kernels/linear/mixed_precision/zentorch.py
index c62a70507682..5975ddb479db 100644
--- a/vllm/model_executor/kernels/linear/mixed_precision/zentorch.py
+++ b/vllm/model_executor/kernels/linear/mixed_precision/zentorch.py
@@ -8,10 +8,12 @@
 """
 
 import torch
+from torch.nn import Parameter
 
 from vllm import envs
 from vllm.logger import init_logger
 from vllm.model_executor.kernels.linear.zentorch_utils import has_zentorch_op
+from vllm.model_executor.utils import replace_parameter
 from vllm.platforms import current_platform
 from vllm.scalar_type import scalar_types
 
@@ -21,6 +23,17 @@
 logger = init_logger(__name__)
 
 
+def _store_in_param(layer: torch.nn.Module, attr: str, name: str | None) -> None:
+    """Back kernel tensor ``layer.<attr>`` by the storage of loaded param ``name``."""
+    tensor = getattr(layer, attr)
+    if name is None or tensor is None:
+        return
+    if not isinstance(getattr(layer, name, None), Parameter):
+        return
+    replace_parameter(layer, name, tensor)
+    setattr(layer, attr, getattr(layer, name).data)
+
+
 def _import_unpack_from_int32():
     """Import compressed-tensors' ``unpack_from_int32`` across versions."""
     try:
@@ -144,8 +157,7 @@ def _maybe_process_da8w4_weights(self, layer: torch.nn.Module) -> bool:
         if (in_features // num_groups) % 4 != 0 or in_features % 2 != 0:
             return False
 
-        if self.w_zp_name is not None:
-            setattr(layer, self.w_zp_name, None)
+        self._retire_unused_zp(layer)
 
         unpack_from_int32 = _import_unpack_from_int32()
 
@@ -165,10 +177,8 @@ def _maybe_process_da8w4_weights(self, layer: torch.nn.Module) -> bool:
         # CT stores scales as [N, G]; DA8W4 wants per-group {G, N}.
         layer._zentorch_da8w4_scale = weight_scale.t().to(torch.bfloat16).contiguous()
 
-        for param_name in (self.w_q_name, self.w_s_name):
-            param = getattr(layer, param_name, None)
-            if param is not None and hasattr(param, "data"):
-                param.data = torch.empty(0)
+        _store_in_param(layer, "_zentorch_da8w4_packed", self.w_q_name)
+        _store_in_param(layer, "_zentorch_da8w4_scale", self.w_s_name)
 
         layer._zentorch_kind = "compressed_tensors_w4a8_da8w4"
         layer._zentorch_da8w4 = True
@@ -205,8 +215,8 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
             super().process_weights_after_loading(layer)
             return
 
-        if (not self.config.zero_points) and (self.w_zp_name is not None):
-            setattr(layer, self.w_zp_name, None)
+        if not self.config.zero_points:
+            self._retire_unused_zp(layer)
 
         weight_q = getattr(layer, self.w_q_name)
         weight_s = getattr(layer, self.w_s_name)
@@ -255,16 +265,9 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
         layer._zentorch_woq_scale = weight_scale.t().contiguous()
         layer._zentorch_woq_zero_point = zp_tc
 
-        for param_name in (self.w_q_name, self.w_s_name, self.w_zp_name):
-            if param_name is None:
-                continue
-            param = getattr(layer, param_name, None)
-            if param is None:
-                continue
-            if hasattr(param, "data"):
-                param.data = torch.empty(0)
-            else:
-                setattr(layer, param_name, torch.empty(0))
+        _store_in_param(layer, "_zentorch_woq_packed", self.w_q_name)
+        _store_in_param(layer, "_zentorch_woq_scale", self.w_s_name)
+        _store_in_param(layer, "_zentorch_woq_zero_point", self.w_zp_name)
 
         layer._zentorch_kind = "compressed_tensors_w4a16_gptq"
         layer._zentorch_processed_weights = True
diff --git a/vllm/model_executor/kernels/linear/mxfp8/flashinfer.py b/vllm/model_executor/kernels/linear/mxfp8/flashinfer.py
index c3b99fc36ec3..f21d0bf37360 100644
--- a/vllm/model_executor/kernels/linear/mxfp8/flashinfer.py
+++ b/vllm/model_executor/kernels/linear/mxfp8/flashinfer.py
@@ -99,7 +99,7 @@ def apply_weights(
             out_dtype, input_shape = qa.orig_dtype, qa.orig_shape
         else:
             assert isinstance(x, torch.Tensor)
-            out_dtype, input_shape = x.dtype, x.shape
+            out_dtype, input_shape = layer.params_dtype, x.shape
             input_mxfp8, input_scale = mxfp8_e4m3_quantize(
                 x.view(-1, K), is_sf_swizzled_layout=True
             )
@@ -177,7 +177,7 @@ def apply_weights(
             out_dtype, input_shape = qa.orig_dtype, qa.orig_shape
         else:
             assert isinstance(x, torch.Tensor)
-            out_dtype, input_shape = x.dtype, x.shape
+            out_dtype, input_shape = layer.params_dtype, x.shape
             input_mxfp8, input_scale = mxfp8_e4m3_quantize(
                 x.view(-1, K), is_sf_swizzled_layout=True
             )
diff --git a/vllm/model_executor/kernels/linear/mxfp8/pytorch.py b/vllm/model_executor/kernels/linear/mxfp8/pytorch.py
index cc1201d22c48..bc83a7d71b16 100644
--- a/vllm/model_executor/kernels/linear/mxfp8/pytorch.py
+++ b/vllm/model_executor/kernels/linear/mxfp8/pytorch.py
@@ -3,12 +3,20 @@
 
 import torch
 
+from vllm.model_executor.layers.fusion.quant_activation import (
+    QuantizedActivation,
+    as_quantized_activation,
+)
 from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
     MXFP8_BLOCK_SIZE,
     mxfp8_e4m3_quantize,
     swizzle_mxfp8_scale,
     xpu_mxfp8_quantize,
 )
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    QuantKey,
+    kMxfp8Dynamic,
+)
 from vllm.model_executor.utils import replace_parameter
 from vllm.platforms import current_platform
 
@@ -37,6 +45,10 @@ def is_supported(
     def can_implement(cls, c: Mxfp8LinearLayerConfig) -> tuple[bool, str | None]:
         return True, None
 
+    def input_quant_key(self) -> QuantKey | None:
+        # CUDA activations use FlashInfer's F8_128x4 swizzled scale layout.
+        return kMxfp8Dynamic if current_platform.is_cuda() else None
+
     def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
         # weight: [N, K] fp8_e4m3; weight_scale: [N, K//32] (E8M0-as-uint8).
         weight = layer.weight.data  # [N, K]
@@ -60,15 +72,20 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
     def apply_weights(
         self,
         layer: torch.nn.Module,
-        x: torch.Tensor,
+        x: torch.Tensor | QuantizedActivation,
         bias: torch.Tensor | None = None,
     ) -> torch.Tensor:
-        out_dtype = x.dtype
         weight = layer.weight
         weight_scale = layer.weight_scale
-        input_shape = x.shape
+        qa = as_quantized_activation(x, self.input_quant_key())
+        if qa is not None:
+            out_dtype, input_shape = qa.orig_dtype, qa.orig_shape
+        else:
+            assert isinstance(x, torch.Tensor)
+            out_dtype, input_shape = layer.params_dtype, x.shape
 
         if current_platform.is_xpu():
+            assert isinstance(x, torch.Tensor)
             # weight stored column-major [K, N].
             K, N = weight.shape
             mat_b = weight
@@ -79,8 +96,14 @@ def apply_weights(
             # weight stored [N, K]; transpose to column-major for _scaled_mm.
             N, K = weight.shape
             mat_b = weight.t()
-            input_2d = x.reshape(-1, K)
-            x_fp8, x_scale = mxfp8_e4m3_quantize(input_2d, is_sf_swizzled_layout=True)
+            if qa is not None:
+                x_fp8, x_scale = qa.data, qa.scale
+            else:
+                assert isinstance(x, torch.Tensor)
+                input_2d = x.reshape(-1, K)
+                x_fp8, x_scale = mxfp8_e4m3_quantize(
+                    input_2d, is_sf_swizzled_layout=True
+                )
             x_scale = x_scale.view(_E8M0)
 
         out = torch._scaled_mm(
diff --git a/vllm/model_executor/kernels/linear/mxfp8/rocm_block32_gemm.py b/vllm/model_executor/kernels/linear/mxfp8/rocm_block32_gemm.py
index 3d5e9c8357ca..9effa2b50bd2 100644
--- a/vllm/model_executor/kernels/linear/mxfp8/rocm_block32_gemm.py
+++ b/vllm/model_executor/kernels/linear/mxfp8/rocm_block32_gemm.py
@@ -50,9 +50,15 @@
 
 import torch
 
+from vllm.logger import init_logger
+from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
+    mxfp8_e4m3_quantize,
+)
 from vllm.triton_utils import tl, triton
 from vllm.utils.torch_utils import direct_register_custom_op
 
+logger = init_logger(__name__)
+
 BLOCK_ROWS = 32
 
 # gfx950 hands out workgroups to its XCDs round-robin by program id, and 8 is
@@ -272,6 +278,124 @@ def _block32_packed_kernel(
         )
 
 
+# _block32_packed_kernel with the MXFP8 activation quantizer folded into
+# the K loop: x_ptr is the RAW bf16/fp16 [M, K] activation (row-major
+# contiguous), quantized per 32-element block in registers with the exact
+# _mxfp8_quant_triton_kernel arithmetic (amax -> biased E8M0 scale ->
+# reciprocal rescale -> RN cast to e4m3), so the separate quant launch and
+# the x_q/x_scale global round-trip disappear from the decode critical
+# path. Tiling, K-panel packing, split-K reduction and stores are
+# byte-for-byte the packed kernel, so the output is bit-identical to
+# quantize-then-GEMM by construction (same quant arithmetic on the same
+# values, same dot order).
+@triton.jit(do_not_specialize=["M"])
+def _block32_packed_quant_kernel(
+    x_ptr,
+    w_ptr,
+    ws_ptr,
+    out_ptr,
+    slot_ptr,
+    count_ptr,
+    M,
+    N: tl.constexpr,
+    K: tl.constexpr,
+    stride_om,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    K_PACK: tl.constexpr,
+    K_PER_SPLIT: tl.constexpr,
+    EVEN_K: tl.constexpr,
+    FUSED_SPLITS: tl.constexpr,
+    FP8_MAX: tl.constexpr,
+    TINY: tl.constexpr,
+):
+    if FUSED_SPLITS > 1:
+        pid_m, pid_n, tile, pid_k = _split_tile(N, BLOCK_N, FUSED_SPLITS)
+        if pid_m * BLOCK_M >= M:
+            return
+        k_start = pid_k * K_PER_SPLIT
+        k_stop = tl.minimum(k_start + K_PER_SPLIT, K)
+    else:
+        pid_m, pid_n, pid_k = tl.program_id(0), tl.program_id(1), tl.program_id(2)
+        if K_PER_SPLIT >= K:
+            k_start = 0
+            k_stop = K
+        else:
+            k_start = pid_k * K_PER_SPLIT
+            k_stop = tl.minimum(k_start + K_PER_SPLIT, K)
+    # Tile row r is token r // K_PACK reading K panel r % K_PACK (and likewise
+    # for weight rows), so one dot covers K_PACK consecutive BLOCK_K slices.
+    rows = pid_m * BLOCK_M + tl.arange(0, BLOCK_M * K_PACK) // K_PACK
+    x_panel = tl.arange(0, BLOCK_M * K_PACK) % K_PACK
+    cols = pid_n * BLOCK_N + tl.arange(0, BLOCK_N * K_PACK) // K_PACK
+    w_panel = tl.arange(0, BLOCK_N * K_PACK) % K_PACK
+    ks = tl.arange(0, BLOCK_K)
+    gs = tl.arange(0, BLOCK_K // 32)
+    row_ok = rows < M
+    col_ok = cols < N
+    acc = tl.zeros((BLOCK_M * K_PACK, BLOCK_N * K_PACK), dtype=tl.float32)
+    for base in range(k_start, k_stop, BLOCK_K * K_PACK):
+        xk = base + x_panel[:, None] * BLOCK_K + ks[None, :]
+        wk = base + w_panel[:, None] * BLOCK_K + ks[None, :]
+        wg = base // 32 + w_panel[:, None] * (BLOCK_K // 32) + gs[None, :]
+        x_mask = row_ok[:, None]
+        w_mask = col_ok[:, None]
+        ws_mask = col_ok[:, None]
+        if not EVEN_K:
+            x_mask = x_mask & (xk < K)
+            w_mask = w_mask & (wk < K)
+            ws_mask = ws_mask & (wg < K // 32)
+        # Fused MXFP8 quant, verbatim _mxfp8_quant_triton_kernel arithmetic.
+        # K is a multiple of 32 and every 32-block of the tile is 32-aligned,
+        # so each block is either fully inside K or fully masked (loaded 0.0,
+        # scale byte 0) -- the same values the separate quantizer produces
+        # for these rows. Masked rows likewise load 0.0 and quantize to
+        # scale byte 0 / value 0, matching the kernel's other=0 loads.
+        x = tl.load(x_ptr + rows[:, None] * K + xk, mask=x_mask, other=0.0).to(
+            tl.float32
+        )
+        xb = tl.reshape(x, (BLOCK_M * K_PACK, BLOCK_K // 32, 32))
+        amax = tl.maximum(tl.max(tl.abs(xb), axis=2), TINY)
+        sb = tl.ceil(tl.log2(amax / FP8_MAX)) + 127.0
+        sb = tl.minimum(tl.maximum(sb, 0.0), 254.0)
+        rescale = tl.exp2(127.0 - sb)
+        xq = tl.reshape(xb * rescale[:, :, None], (BLOCK_M * K_PACK, BLOCK_K)).to(
+            w_ptr.dtype.element_ty
+        )
+        xs = sb.to(tl.uint8)
+        w = tl.load(w_ptr + cols[:, None] * K + wk, mask=w_mask, other=0.0)
+        ws = tl.load(
+            ws_ptr + (cols[:, None] // 32) * (K // 32) + wg, mask=ws_mask, other=0
+        )
+        acc = tl.dot_scaled(xq, xs, "e4m3", w.T, ws, "e4m3", acc=acc)
+    # Keep only products of matching K panels: the block diagonal.
+    blocks = acc.reshape(BLOCK_M, K_PACK, BLOCK_N, K_PACK).trans(0, 2, 1, 3)
+    pair = tl.arange(0, K_PACK)
+    diagonal = tl.where(
+        pair[None, None, :, None] == pair[None, None, None, :], blocks, 0.0
+    )
+    y = tl.sum(tl.sum(diagonal, 3), 2)
+    out_rows = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    out_cols = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    out_ptrs = out_ptr + out_rows[:, None] * stride_om + out_cols[None, :]
+    out_mask = (out_rows[:, None] < M) & (out_cols[None, :] < N)
+    if FUSED_SPLITS > 1:
+        _sum_splits(
+            y,
+            out_ptrs,
+            out_mask,
+            slot_ptr + tile * (FUSED_SPLITS * BLOCK_M * BLOCK_N),
+            count_ptr + tile,
+            pid_k,
+            FUSED_SPLITS,
+        )
+    else:
+        tl.store(
+            out_ptrs + pid_k * M * N, y.to(out_ptr.dtype.element_ty), mask=out_mask
+        )
+
+
 @triton.jit(do_not_specialize=["M"])
 def _block32_splitk_reduce_kernel(
     partial_ptr,
@@ -331,9 +455,9 @@ def _packed(bm, bn, bk, k_pack, warps, stages, nonkdim, waves, fused=1) -> _Conf
     )
 
 
-# (N, K) -> [(max M, config), ...] for the DeepSeek-V4.1-Flash MXFP8 linears at
-# TP2/TP4, tuned on MI355X under HIP graphs with weights cold in the LLC. A
-# tier serves every M up to its bound; larger M reuse the last tier.
+# (N, K) -> [(max M, config), ...] for the DeepSeek-V4.1-Flash MXFP8
+# linears at TP2/TP4, tuned on MI355X under HIP graphs. A tier serves
+# every M up to its bound; larger M reuse the last tier.
 _TUNED: dict[tuple[int, int], list[tuple[int, _Config]]] = {
     # shared expert gate_up, TP4
     (1152, 5120): [
@@ -563,6 +687,49 @@ def _config(M: int, N: int, K: int) -> _Config:
     return tiers[-1][1]
 
 
+# Retuned tiling/scheduling for the fused-quant kernel, whose per-CTA
+# quant prologue shifts the best (BLOCK_M, BLOCK_N, num_warps, num_stages,
+# waves_per_eu). Entries override ONLY those five axes of the stock tier
+# active at M; the K-structure axes (block_k, k_split/k_pack,
+# fused_splits, matrix_instr_nonkdim) stay at that tier's values so no
+# output element's K-accumulation order changes and the fused output stays
+# bit-identical to the two-launch path at every M. Consulted only on the
+# quant_in_kernel path, never by the stock two-launch path.
+_TUNED_FUSED_QUANT: dict[
+    tuple[int, int], list[tuple[int, tuple[int, int, int, int, int]]]
+] = {
+    # (N, K): [(max_m, (block_m, block_n, warps, stages, waves)), ...]
+    # DeepSeek-V4.1-Flash TP4 decode shapes:
+    (1152, 5120): [(8, (8, 32, 4, 2, 0))],  # shared_gate_up
+    (1792, 5120): [(8, (8, 16, 4, 3, 1))],  # fused_wqa_wkv
+    (4096, 1280): [(8, (8, 16, 4, 2, 0))],  # indexer_wq_b
+    (5120, 576): [(8, (8, 16, 4, 2, 0))],  # shared_down
+    (5120, 2048): [(8, (8, 32, 8, 3, 0))],  # wo_b
+    (8192, 1280): [(8, (8, 32, 2, 3, 1))],  # wq_b
+    # DSpark draft shapes (stock falls to _default_config):
+    (1024, 2048): [(8, (8, 16, 4, 2, 1))],  # draft
+    (768, 1536): [(8, (8, 16, 4, 3, 0))],  # draft
+}
+
+
+def _fused_quant_config(M: int, N: int, K: int) -> _Config:
+    """The stock tier active at M with its tiling/scheduling retuned for the
+    in-kernel quant prologue (see _TUNED_FUSED_QUANT); the stock config
+    where the shape has no fused-tier entry. K-structure axes are always
+    the stock tier's own, so the output stays bit-exact at every M."""
+    base = _config(M, N, K)
+    for max_m, (bm, bn, warps, stages, waves) in _TUNED_FUSED_QUANT.get((N, K), ()):
+        if max_m >= M:
+            return base._replace(
+                block_m=bm,
+                block_n=bn,
+                num_warps=warps,
+                num_stages=stages,
+                waves_per_eu=waves,
+            )
+    return base
+
+
 _split_counters_by_stream: dict[tuple[int, int], torch.Tensor] = {}
 
 
@@ -599,7 +766,11 @@ def _launch(
     weight_scale: torch.Tensor,
     out: torch.Tensor,
     cfg: _Config,
+    quant_in_kernel: bool = False,
 ) -> None:
+    """Launch cfg's kernel. quant_in_kernel=True (packed cfgs only) makes
+    x the raw bf16/fp16 activation and x_scale None: the kernel quantizes
+    x per 32-block in registers (see _block32_packed_quant_kernel)."""
     M, K = x.shape
     N = weight.shape[0]
     opts = dict(
@@ -638,21 +809,46 @@ def _launch(
             grid = (grid_m, grid_n, splits)
         else:
             grid = (grid_n, grid_m, splits)
-    args = (x, x_scale, weight, weight_scale, target, slots, counters, M, N, K)
+    args: tuple
+    if quant_in_kernel:
+        if not cfg.packed:
+            raise ValueError("quant_in_kernel needs a packed tile config")
+        args = (x, weight, weight_scale, target, slots, counters, M, N, K)
+    else:
+        args = (x, x_scale, weight, weight_scale, target, slots, counters, M, N, K)
     stride = out.stride(0) if splits == 1 or fused > 1 else N
     if cfg.packed:
-        _block32_packed_kernel[grid](
-            *args,
-            stride,
-            BLOCK_M=cfg.block_m,
-            BLOCK_N=cfg.block_n,
-            BLOCK_K=cfg.block_k,
-            K_PACK=cfg.k_split,
-            K_PER_SPLIT=k_per_split,
-            EVEN_K=K % step == 0,
-            FUSED_SPLITS=fused,
-            **opts,
-        )
+        if quant_in_kernel:
+            # Same output contract as the packed kernel below (partials at
+            # pid_k * M * N when FUSED_SPLITS == 1 with multiple pid_k), so
+            # the split-K reduce after this branch applies unchanged.
+            _block32_packed_quant_kernel[grid](
+                *args,
+                stride,
+                BLOCK_M=cfg.block_m,
+                BLOCK_N=cfg.block_n,
+                BLOCK_K=cfg.block_k,
+                K_PACK=cfg.k_split,
+                K_PER_SPLIT=k_per_split,
+                EVEN_K=K % step == 0,
+                FUSED_SPLITS=fused,
+                FP8_MAX=_FP8_MAX,
+                TINY=_FP32_TINY,
+                **opts,
+            )
+        else:
+            _block32_packed_kernel[grid](
+                *args,
+                stride,
+                BLOCK_M=cfg.block_m,
+                BLOCK_N=cfg.block_n,
+                BLOCK_K=cfg.block_k,
+                K_PACK=cfg.k_split,
+                K_PER_SPLIT=k_per_split,
+                EVEN_K=K % step == 0,
+                FUSED_SPLITS=fused,
+                **opts,
+            )
     else:
         _block32_tiled_kernel[grid](
             *args,
@@ -735,3 +931,113 @@ def rocm_mxfp8_block32_gemm(
     return torch.ops.vllm.rocm_mxfp8_block32_gemm(
         x, x_scale, weight, weight_scale, out_dtype
     )
+
+
+# --- tiny-M fused quant+GEMM ------------------------------------------------
+# At M <= 8 (decode-sized batches) the fused kernel below computes the
+# activation quant inside the GEMM's K loop, saving one launch and its
+# global round-trip per call.
+
+_TINY_M_FUSED_QUANT_MAX_M = 8
+# Shared with _mxfp8_quant_triton_kernel (the bit-exactness contract).
+_FP8_MAX = float(torch.finfo(torch.float8_e4m3fn).max)
+_FP32_TINY = float(torch.finfo(torch.float32).tiny)
+
+_FUSED_QUANT_HIT_SHAPES: set = set()
+
+
+def _rocm_mxfp8_block32_gemm_quant_impl(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    weight_scale: torch.Tensor,
+    out_dtype: torch.dtype,
+) -> torch.Tensor:
+    """``x @ weight.T`` for raw bf16/fp16 ``x`` and a 32x32 block-scaled
+    MXFP8 weight, quantizing ``x`` inside the GEMM at tiny M.
+
+    M <= 8 with a packed tile config (every tuned M <= 8 tier is packed)
+    runs one launch whose K loop computes the per-32-block E8M0 scale and
+    the e4m3 cast exactly as ``_mxfp8_quant_triton_kernel`` and feeds
+    ``tl.dot_scaled`` from registers -- bit-identical output to
+    ``mxfp8_e4m3_quantize`` + ``rocm_mxfp8_block32_gemm`` with one launch
+    fewer on the decode critical path. Everything else (larger M,
+    non-packed tiles, or non-contiguous or non-bf16/fp16 x) runs the stock
+    two-launch path with identical results. The branch is inside this
+    opaque op so a compiled caller cannot freeze the token-count decision.
+    """
+    M, K = x.shape
+    out = torch.empty((M, weight.shape[0]), dtype=out_dtype, device=x.device)
+    if M <= 0:
+        return out
+    cfg = _config(M, weight.shape[0], K)
+    eligible = (
+        cfg.packed
+        and M <= _TINY_M_FUSED_QUANT_MAX_M
+        and x.dtype in (torch.bfloat16, torch.float16)
+        and x.is_contiguous()
+    )
+    if eligible:
+        # Retuned tiling for the in-kernel quant prologue (same K-structure
+        # as the stock tier at M, so bit-exactness holds); the stock path
+        # below keeps the stock config untouched.
+        fused_cfg = _fused_quant_config(M, weight.shape[0], K)
+        shape = (M, weight.shape[0], K)
+        if shape not in _FUSED_QUANT_HIT_SHAPES:
+            _FUSED_QUANT_HIT_SHAPES.add(shape)
+            # warning_once keeps the engagement log one-shot per shape;
+            # WARNING, not INFO, so it reaches the server log.
+            logger.warning_once(
+                "tiny-M fused MXFP8 quant+GEMM HIT (M=%d, N=%d, K=%d, packed "
+                "tile %dx%dx%d k_pack=%d%s)",
+                M,
+                weight.shape[0],
+                K,
+                fused_cfg.block_m,
+                fused_cfg.block_n,
+                fused_cfg.block_k,
+                fused_cfg.k_split,
+                (
+                    f", fused_splits={fused_cfg.fused_splits}"
+                    if fused_cfg.fused_splits > 1
+                    else ""
+                ),
+            )
+        _launch(x, None, weight, weight_scale, out, fused_cfg, quant_in_kernel=True)
+    else:
+        x_q, x_scale = mxfp8_e4m3_quantize(x)
+        _launch(x_q, x_scale, weight, weight_scale, out, cfg)
+    return out
+
+
+def _rocm_mxfp8_block32_gemm_quant_fake(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    weight_scale: torch.Tensor,
+    out_dtype: torch.dtype,
+) -> torch.Tensor:
+    return x.new_empty((x.shape[0], weight.shape[0]), dtype=out_dtype)
+
+
+# An opaque op for the same reason as rocm_mxfp8_block32_gemm: the fused vs
+# two-launch choice branches on the token count.
+direct_register_custom_op(
+    op_name="rocm_mxfp8_block32_gemm_quant",
+    op_func=_rocm_mxfp8_block32_gemm_quant_impl,
+    fake_impl=_rocm_mxfp8_block32_gemm_quant_fake,
+)
+
+
+def rocm_mxfp8_block32_gemm_quant(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    weight_scale: torch.Tensor,
+    out_dtype: torch.dtype,
+) -> torch.Tensor:
+    """``x @ weight.T`` for a raw bf16/fp16 ``x`` [M, K] and a 32x32
+    block-scaled MXFP8 ``weight`` [N, K] with ``weight_scale``
+    [N / 32, K / 32]. Tiny-M batches quantize inside the packed GEMM (see
+    ``_rocm_mxfp8_block32_gemm_quant_impl``).
+    """
+    return torch.ops.vllm.rocm_mxfp8_block32_gemm_quant(
+        x, weight, weight_scale, out_dtype
+    )
diff --git a/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py b/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py
index e46b61275c66..878a9ee1cd0e 100644
--- a/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py
+++ b/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py
@@ -3,27 +3,61 @@
 """Native MXFP8 linear GEMM for AMD CDNA4 (gfx950) via Triton ``tl.dot_scaled``.
 
 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. ``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.
+activations are MXFP8-quantized per token (inside the GEMM for small M), or
+arrive pre-quantized as a ``QuantizedActivation`` with row-major (not swizzled)
+E8M0 scales. Uses the CDNA4 hardware microscaling 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
 
 import vllm.envs as envs
+from vllm.model_executor.layers.fusion.quant_activation import (
+    QuantizedActivation,
+    as_quantized_activation,
+)
 from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
     MXFP8_BLOCK_SIZE,
     MXFP8_SCALE_DTYPE,
+    MXFP8_VALUE_DTYPE,
     dequant_mxfp8_to_bf16,
     mxfp8_e4m3_quantize,
 )
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    QuantKey,
+    kMxfp8Dynamic,
+)
 from vllm.model_executor.utils import replace_parameter
 from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
 
 from .Mxfp8LinearKernel import Mxfp8LinearKernel, Mxfp8LinearLayerConfig
-from .rocm_block32_gemm import BLOCK_ROWS, rocm_mxfp8_block32_gemm
+from .rocm_block32_gemm import (
+    BLOCK_ROWS,
+    rocm_mxfp8_block32_gemm,
+    rocm_mxfp8_block32_gemm_quant,
+)
+
+_FP8_MAX = float(torch.finfo(MXFP8_VALUE_DTYPE).max)
+_TINY = float(torch.finfo(torch.float32).tiny)
+
+
+@triton.jit
+def _quant_mxfp8_tile(x, FP8_MAX: tl.constexpr, TINY: tl.constexpr):
+    """MXFP8-quantize a [BLOCK_M, BLOCK_K] tile: bit-exact with
+    ``mxfp8_e4m3_quantize`` (see ``_mxfp8_quant_triton_kernel`` for why the
+    scale is applied as a reciprocal)."""
+    BM: tl.constexpr = x.shape[0]
+    BK: tl.constexpr = x.shape[1]
+    xb = tl.reshape(x.to(tl.float32), (BM, BK // 32, 32))
+    amax = tl.maximum(tl.max(tl.abs(xb), axis=2), TINY)
+    sb = tl.ceil(tl.log2(amax / FP8_MAX)) + 127.0
+    sb = tl.minimum(tl.maximum(sb, 0.0), 254.0)
+    rescale = tl.exp2(127.0 - sb)
+    xq = tl.reshape((xb * rescale[:, :, None]).to(tl.float8e4nv), (BM, BK))
+    return xq, sb.to(tl.uint8)
 
 
 @triton.jit
@@ -35,7 +69,7 @@ def _mxfp8_linear_kernel(
     out_ptr,
     M,
     N,
-    K,
+    K: tl.constexpr,
     stride_xm,
     stride_xk,
     stride_xsm,
@@ -49,31 +83,59 @@ def _mxfp8_linear_kernel(
     BLOCK_M: tl.constexpr,
     BLOCK_N: tl.constexpr,
     BLOCK_K: tl.constexpr,
+    QUANT_X: tl.constexpr,
+    QUANT_ROWS: tl.constexpr,
+    FP8_MAX: tl.constexpr,
+    TINY: tl.constexpr,
 ):
     pid_m = tl.program_id(0)
     pid_n = tl.program_id(1)
     offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    # Rows of x that are loaded: all BLOCK_M, or with QUANT_X the first
+    # QUANT_ROWS (a power of 2 >= the rows this tile stores).
+    offs_xm = pid_m * BLOCK_M + tl.arange(0, QUANT_ROWS)
     offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
     offs_k = tl.arange(0, BLOCK_K)
     offs_sk = tl.arange(0, BLOCK_K // 32)
     m_mask = offs_m < M
+    xm_mask = offs_xm < M
     n_mask = offs_n < N
 
-    x_ptrs = x_ptr + offs_m[:, None] * stride_xm + offs_k[None, :] * stride_xk
+    x_ptrs = x_ptr + offs_xm[:, None] * stride_xm + offs_k[None, :] * stride_xk
     xs_ptrs = xs_ptr + offs_m[:, None] * stride_xsm + offs_sk[None, :] * stride_xsk
     w_ptrs = w_ptr + offs_n[:, None] * stride_wn + offs_k[None, :] * stride_wk
     ws_ptrs = ws_ptr + offs_n[:, None] * stride_wsn + offs_sk[None, :] * stride_wsk
 
     acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
     for _ in range(0, tl.cdiv(K, BLOCK_K)):
-        x = tl.load(x_ptrs, mask=m_mask[:, None], other=0.0)
-        w = tl.load(w_ptrs, mask=n_mask[:, None], other=0.0)
-        xs = tl.load(xs_ptrs, mask=m_mask[:, None], other=0)
-        ws = tl.load(ws_ptrs, mask=n_mask[:, None], other=0)
+        if QUANT_X:
+            # x is the unquantized activation; every N-tile quantizes its own
+            # copy, which is cheaper than a separate quant launch at small M.
+            xf = tl.load(x_ptrs, mask=xm_mask[:, None], other=0.0)
+            w = tl.load(w_ptrs, mask=n_mask[:, None], other=0.0)
+            ws = tl.load(ws_ptrs, mask=n_mask[:, None], other=0)
+            x, xs = _quant_mxfp8_tile(xf, FP8_MAX, TINY)
+            if QUANT_ROWS < BLOCK_M:
+                # BLOCK_M is the smallest MFMA tile; repeating the quantized
+                # rows fills it, and rows >= M are never stored.
+                REPS: tl.constexpr = BLOCK_M // QUANT_ROWS
+                x = tl.reshape(
+                    tl.broadcast_to(x[None, :, :], (REPS, QUANT_ROWS, BLOCK_K)),
+                    (BLOCK_M, BLOCK_K),
+                )
+                xs = tl.reshape(
+                    tl.broadcast_to(xs[None, :, :], (REPS, QUANT_ROWS, BLOCK_K // 32)),
+                    (BLOCK_M, BLOCK_K // 32),
+                )
+        else:
+            x = tl.load(x_ptrs, mask=m_mask[:, None], other=0.0)
+            w = tl.load(w_ptrs, mask=n_mask[:, None], other=0.0)
+            xs = tl.load(xs_ptrs, mask=m_mask[:, None], other=0)
+            ws = tl.load(ws_ptrs, mask=n_mask[:, None], other=0)
+            xs_ptrs += (BLOCK_K // 32) * stride_xsk
         acc += tl.dot_scaled(x, xs, "e4m3", w.T, ws, "e4m3")
         x_ptrs += BLOCK_K * stride_xk
         w_ptrs += BLOCK_K * stride_wk
-        xs_ptrs += (BLOCK_K // 32) * stride_xsk
         ws_ptrs += (BLOCK_K // 32) * stride_wsk
 
     o_ptrs = out_ptr + offs_m[:, None] * stride_om + offs_n[None, :] * stride_on
@@ -82,52 +144,125 @@ def _mxfp8_linear_kernel(
     )
 
 
-def _mxfp8_dot_scaled_linear(
-    x: torch.Tensor,  # [M, K] bf16/fp16
-    w: torch.Tensor,  # [N, K] fp8 e4m3
-    w_scale: torch.Tensor,  # [N, K//32] uint8 (E8M0)
-) -> torch.Tensor:
+def _launch(
+    x: torch.Tensor,
+    x_scale: torch.Tensor | None,
+    w: torch.Tensor,
+    w_scale: torch.Tensor,
+    out: torch.Tensor,
+    block_m: int,
+    block_n: int,
+    block_k: int,
+    num_warps: int,
+    num_stages: int,
+) -> None:
+    """``out = x @ w.T``. ``x_scale`` None means ``x`` is the unquantized
+    activation and is MXFP8-quantized inside the kernel."""
     M, K = x.shape
     N = w.shape[0]
-    x_q, x_scale = mxfp8_e4m3_quantize(x)
-    out = torch.empty((M, N), dtype=x.dtype, device=x.device)
-    BLOCK_M, BLOCK_N, BLOCK_K, num_warps, num_stages = _select_cfg(M, N, K)
-    grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(N, BLOCK_N))
+    quant_x = x_scale is None
+    xs = x if x_scale is None else x_scale  # Unused by the kernel if quant_x.
+    # Quantizing only the rows that can be stored keeps the fused quant's cost
+    # (serial per N-tile) from scaling with BLOCK_M at M < BLOCK_M. At
+    # BLOCK_K=128 that cost is small and repeating the rows costs more.
+    if quant_x and block_k >= 256:
+        quant_rows = min(block_m, triton.next_power_of_2(M))
+    else:
+        quant_rows = block_m
+    grid = (triton.cdiv(M, block_m), triton.cdiv(N, block_n))
     _mxfp8_linear_kernel[grid](
-        x_q,
-        x_scale,
+        x,
+        xs,
         w,
         w_scale,
         out,
         M,
         N,
         K,
-        x_q.stride(0),
-        x_q.stride(1),
-        x_scale.stride(0),
-        x_scale.stride(1),
+        x.stride(0),
+        x.stride(1),
+        xs.stride(0),
+        xs.stride(1),
         w.stride(0),
         w.stride(1),
         w_scale.stride(0),
         w_scale.stride(1),
         out.stride(0),
         out.stride(1),
-        BLOCK_M=BLOCK_M,
-        BLOCK_N=BLOCK_N,
-        BLOCK_K=BLOCK_K,
+        BLOCK_M=block_m,
+        BLOCK_N=block_n,
+        BLOCK_K=block_k,
+        QUANT_X=quant_x,
+        QUANT_ROWS=quant_rows,
+        FP8_MAX=_FP8_MAX,
+        TINY=_TINY,
         num_warps=num_warps,
         num_stages=num_stages,
     )
+
+
+# Past this many rows, quantizing the activation once per N-tile inside the
+# GEMM costs more than one separate quant launch.
+_FUSED_QUANT_MAX_M = 16
+
+
+def _select_fused_quant_cfg(M: int, N: int, K: int):
+    """Tile config with the activation quant fused (M <= 16), or None to
+    quantize separately.
+    """
+    if K == 2048 and M <= 4:
+        return 16, 32, 512, 4, 3
+    if K > 1024:
+        return None
+    if K == 1024:
+        return 16, 32, 512, 4, 3
+    return 16, 32, 256 if K % 256 == 0 else 128, 4, 3
+
+
+def _mxfp8_dot_scaled_linear(
+    x: torch.Tensor,  # [M, K] bf16/fp16
+    w: torch.Tensor,  # [N, K] fp8 e4m3
+    w_scale: torch.Tensor,  # [N, K//32] uint8 (E8M0)
+) -> torch.Tensor:
+    M, K = x.shape
+    N = w.shape[0]
+    if (
+        0 < M <= _FUSED_QUANT_MAX_M
+        and x.dtype == torch.bfloat16
+        and not envs.VLLM_BATCH_INVARIANT
+    ):
+        cfg = _select_fused_quant_cfg(M, N, K)
+        if cfg is not None:
+            out = torch.empty((M, N), dtype=x.dtype, device=x.device)
+            _launch(x, None, w, w_scale, out, *cfg)
+            return out
+    x_q, x_scale = mxfp8_e4m3_quantize(x)
+    return _mxfp8_dot_scaled_gemm(x_q, x_scale, w, w_scale, x.dtype)
+
+
+def _mxfp8_dot_scaled_gemm(
+    x_q: torch.Tensor,  # [M, K] fp8 e4m3
+    x_scale: torch.Tensor,  # [M, K//32] uint8 (E8M0), row-major
+    w: torch.Tensor,  # [N, K] fp8 e4m3
+    w_scale: torch.Tensor,  # [N, K//32] uint8 (E8M0)
+    out_dtype: torch.dtype,
+) -> torch.Tensor:
+    M, K = x_q.shape
+    N = w.shape[0]
+    out = torch.empty((M, N), dtype=out_dtype, device=x_q.device)
+    _launch(x_q, x_scale, w, w_scale, out, *_select_cfg(M, N, K))
     return out
 
 
 # gfx950's async-copy LDS (default in the triton 3.8 this branch pins) leaves no
 # room for num_stages=3 at BLOCK_K=256; use 2 there.
 _BK256_STAGES = 3
+_ON_GFX950 = False
 if current_platform.is_rocm():
     from vllm.platforms.rocm import on_gfx950
 
-    if on_gfx950():
+    _ON_GFX950 = on_gfx950()
+    if _ON_GFX950:
         _BK256_STAGES = 2
 
 
@@ -142,6 +277,11 @@ def _select_cfg(M, N, K):
     (the K-loop is unmasked), so every BLOCK_K below is guarded to be K-divisible
     (served K: 384/768/1024/2048/6144).
     """
+    if _ON_GFX950 and M <= 16:
+        if K == 6144 and N in (1536, 2560):
+            return 16, 16, 1024, 2, 3
+        if K == 768 and N == 6144:
+            return 16, 16, 256, 2, 3
     if M <= 64:
         # decode (M in {1,32,64}): tiny-M GEMV is weight-BW + GPU-OCCUPANCY bound. The
         # lever is NARROW BLOCK_N=16 (maximize N-tiles so more CUs stream the weight in
@@ -261,10 +401,24 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
         replace_parameter(layer, "weight", weight.contiguous())
         replace_parameter(layer, "weight_scale", weight_scale)
 
+    def input_quant_key(self) -> QuantKey:
+        # Row-major [M, K // 32] E8M0 scales, unlike FlashInfer's swizzled ones
+        # under the same key: producers must also check the kernel type.
+        return kMxfp8Dynamic
+
+    @staticmethod
+    def supports_quantized_input(layer: torch.nn.Module) -> bool:
+        """Whether ``apply_weights`` reads a ``QuantizedActivation`` for
+        ``layer`` as is (fp8 weights), rather than dequantizing it back."""
+        return layer.weight.element_size() == 1 and (
+            layer.weight_scale.shape[0] != layer.weight.shape[0]
+            or layer.weight.shape[1] % _DOT_SCALED_K_ALIGN == 0
+        )
+
     def apply_weights(
         self,
         layer: torch.nn.Module,
-        x: torch.Tensor,
+        x: torch.Tensor | QuantizedActivation,
         bias: torch.Tensor | None = None,
     ) -> torch.Tensor:
         if layer.weight_scale.dtype != MXFP8_SCALE_DTYPE:
@@ -272,23 +426,50 @@ def apply_weights(
                 f"Expected {MXFP8_SCALE_DTYPE} weight_scale, got "
                 f"{layer.weight_scale.dtype}."
             )
-        out_shape = (*x.shape[:-1], layer.weight.shape[0])
-        x2d = x.reshape(-1, x.shape[-1])
+        qa = as_quantized_activation(x, self.input_quant_key())
+        if qa is not None and not self.supports_quantized_input(layer):
+            # BF16 weight (K % 128 != 0): serve the dequantized activation.
+            x = dequant_mxfp8_to_bf16(qa.data, qa.scale).to(qa.orig_dtype)
+            x, qa = x.reshape(qa.orig_shape), None
+        if qa is not None:
+            out_dtype, in_shape = qa.orig_dtype, qa.orig_shape
+            x_q, x_scale = qa.data.reshape(-1, in_shape[-1]), qa.scale
+        else:
+            assert isinstance(x, torch.Tensor)
+            out_dtype, in_shape = x.dtype, x.shape
+            x2d = x.reshape(-1, in_shape[-1])
+        out_shape = (*in_shape[:-1], layer.weight.shape[0])
         if layer.weight_scale.shape[0] != layer.weight.shape[0]:
             # One scale row per 32 weight rows (see process_weights_after_loading).
-            x_q, x_scale = mxfp8_e4m3_quantize(x2d)
-            out = rocm_mxfp8_block32_gemm(
-                x_q, x_scale, layer.weight, layer.weight_scale, x.dtype
+            if qa is not None:
+                out = rocm_mxfp8_block32_gemm(
+                    x_q, x_scale, layer.weight, layer.weight_scale, out_dtype
+                )
+            else:
+                # The op quantizes the activation itself: at tiny M (the
+                # decode / DSpark shapes) the quant is fused into the packed
+                # GEMM kernel -- one launch, bit-exact with
+                # mxfp8_e4m3_quantize, kill switch
+                # VLLM_DISABLE_FUSED_MXFP8_QUANT_GEMM=1 -- while larger M
+                # keeps the separate launches. The branch lives inside the
+                # opaque op so a compiled caller cannot freeze the
+                # token-count decision.
+                out = rocm_mxfp8_block32_gemm_quant(
+                    x2d, layer.weight, layer.weight_scale, out_dtype
+                )
+        elif qa is not None:
+            out = _mxfp8_dot_scaled_gemm(
+                x_q, x_scale, layer.weight, layer.weight_scale, out_dtype
             )
         elif layer.weight.element_size() >= 2:
-            out = torch.nn.functional.linear(x2d, layer.weight.to(x.dtype))
+            out = torch.nn.functional.linear(x2d, layer.weight.to(out_dtype))
         elif x2d.shape[-1] % _DOT_SCALED_K_ALIGN == 0:
             out = _mxfp8_dot_scaled_linear(x2d, layer.weight, layer.weight_scale)
         else:
             # 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 = torch.nn.functional.linear(x2d, w_bf16).to(out_dtype)
         out = out.reshape(out_shape)
         if bias is not None:
             out = out + bias
diff --git a/vllm/model_executor/kernels/linear/nvfp4/humming.py b/vllm/model_executor/kernels/linear/nvfp4/humming.py
index d6d9cb441b9e..676c3a6c86e7 100644
--- a/vllm/model_executor/kernels/linear/nvfp4/humming.py
+++ b/vllm/model_executor/kernels/linear/nvfp4/humming.py
@@ -11,7 +11,11 @@
     quant_key_to_input_schema,
 )
 from vllm.model_executor.layers.quantization.utils.quant_utils import kNvfp4Dynamic
-from vllm.model_executor.utils import replace_parameter
+from vllm.model_executor.utils import (
+    move_parameter,
+    replace_parameter,
+    set_kernel_parameter,
+)
 from vllm.platforms import current_platform
 from vllm.utils.humming import create_humming_locks
 from vllm.utils.import_utils import has_humming
@@ -53,16 +57,18 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
             "strategy": "group",
             "group_size": 16,
         }
-        # CT pack-quantized reads `weight_packed`; the scheme renamed it to `weight`.
-        if not hasattr(layer, "weight_packed"):
-            layer.weight_packed = layer.weight
-            del layer.weight
+        # CT pack-quantized reads `weight_packed`: make it the weight's runtime name
+        # (a CT checkpoint already has it as an alias of `weight`).
+        move_parameter(layer, "weight", "weight_packed")
         # The CT linear scheme inverts the global scale (1/scale) for
         # marlin/cutlass; humming wants the original.
         replace_parameter(layer, "weight_global_scale", 1.0 / layer.weight_global_scale)
         input_quant_key = None
         if hasattr(layer, "input_global_scale"):
-            replace_parameter(
+            # Quark keeps its loaded (retired) `input_scale_2`: revive that one.
+            if "input_scale_2" not in layer._parameters:
+                move_parameter(layer, "input_global_scale", "input_scale_2")
+            set_kernel_parameter(
                 layer, "input_scale_2", layer.input_global_scale.reshape(-1)
             )
             input_quant_key = kNvfp4Dynamic
diff --git a/vllm/model_executor/kernels/linear/scaled_mm/ScaledMMLinearKernel.py b/vllm/model_executor/kernels/linear/scaled_mm/ScaledMMLinearKernel.py
index e632f6f13f28..20726608b76c 100644
--- a/vllm/model_executor/kernels/linear/scaled_mm/ScaledMMLinearKernel.py
+++ b/vllm/model_executor/kernels/linear/scaled_mm/ScaledMMLinearKernel.py
@@ -200,6 +200,7 @@ def _get_layer_params(self, layer) -> _Int8ParamsT:
             getattr(layer, w_q),
             getattr(layer, w_s),
             getattr(layer, i_s, None),
-            getattr(layer, i_zp, None),
+            # A loaded zero point of a symmetric scheme stays registered, unused.
+            None if self.config.input_symmetric else getattr(layer, i_zp, None),
             getattr(layer, azp_adj, None),
         )
diff --git a/vllm/model_executor/kernels/linear/scaled_mm/cpu.py b/vllm/model_executor/kernels/linear/scaled_mm/cpu.py
index 00a1bf06d617..32ce9296e20d 100644
--- a/vllm/model_executor/kernels/linear/scaled_mm/cpu.py
+++ b/vllm/model_executor/kernels/linear/scaled_mm/cpu.py
@@ -142,8 +142,8 @@ def process_weights_for_onednn(self, layer: torch.nn.Module) -> None:
             32,
         )
         # weight is prepacked and maintained by the dnnl_handler,
-        # release the original weight
-        setattr(layer, w_q_name, None)
+        # release the original weight storage
+        replace_parameter(layer, w_q_name, weight.new_empty(0))
         del weight
 
     def process_weights_for_sgl(self, layer: torch.nn.Module) -> None:
diff --git a/vllm/model_executor/kernels/linear/scaled_mm/cutlass.py b/vllm/model_executor/kernels/linear/scaled_mm/cutlass.py
index 2d4c5fecbdbc..366e62c2d5e2 100644
--- a/vllm/model_executor/kernels/linear/scaled_mm/cutlass.py
+++ b/vllm/model_executor/kernels/linear/scaled_mm/cutlass.py
@@ -79,7 +79,6 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
                     i_s_name,
                     torch.nn.Parameter(input_scale.max(), requires_grad=False),
                 )
-                setattr(layer, i_zp_name, None)
             else:
                 input_zero_point = getattr(layer, i_zp_name)
 
diff --git a/vllm/model_executor/kernels/linear/scaled_mm/marlin.py b/vllm/model_executor/kernels/linear/scaled_mm/marlin.py
index 1df6ad55412d..c1f5ad949a1b 100644
--- a/vllm/model_executor/kernels/linear/scaled_mm/marlin.py
+++ b/vllm/model_executor/kernels/linear/scaled_mm/marlin.py
@@ -73,11 +73,10 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
             replace_parameter(layer, scale_name, weight_scale.data)
         # Non-block: callers must pass weight in (K, N) layout.
 
-        layer.input_scale = None
+        # A loaded input_scale stays registered; weight-only Marlin ignores it.
         prepare_fp8_layer_for_marlin(
             layer, self.size_k_first, input_dtype=self.marlin_input_dtype
         )
-        del layer.input_scale
 
     def apply_weights(
         self,
diff --git a/vllm/model_executor/kernels/linear/scaled_mm/triton.py b/vllm/model_executor/kernels/linear/scaled_mm/triton.py
index b9f7e55d5f07..d641c92f0612 100644
--- a/vllm/model_executor/kernels/linear/scaled_mm/triton.py
+++ b/vllm/model_executor/kernels/linear/scaled_mm/triton.py
@@ -73,7 +73,6 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
                     i_s_name,
                     torch.nn.Parameter(i_s.max(), requires_grad=False),
                 )
-                setattr(layer, i_zp_name, None)
             else:
                 input_zero_point = getattr(layer, i_zp_name)
 
diff --git a/vllm/model_executor/layers/attention/attention.py b/vllm/model_executor/layers/attention/attention.py
index 93958bb12782..fc23e9791c5b 100644
--- a/vllm/model_executor/layers/attention/attention.py
+++ b/vllm/model_executor/layers/attention/attention.py
@@ -256,6 +256,7 @@ def __init__(
         mm_prefix_clamp_sliding_window: bool = False,
         attn_backend: type[AttentionBackend] | None = None,
         head_size_v: int | None = None,
+        query_pre_quantized: bool = False,
         **extra_impl_args,
     ) -> None:
         """The KV cache is stored inside this class and is accessed via
@@ -278,6 +279,11 @@ def __init__(
         else:
             kv_cache_dtype = "auto"
 
+        if quant_config is not None and kv_cache_dtype == "auto":
+            checkpoint_kv_cache_dtype = quant_config.get_kv_cache_dtype(prefix)
+            if checkpoint_kv_cache_dtype is not None:
+                kv_cache_dtype = checkpoint_kv_cache_dtype
+
         # llm-compressor models declare an FP8 KV-cache scheme in their
         # checkpoint config. Honor it only when the user did not explicitly
         # pick a kv_cache_dtype; an explicit choice (e.g. bfloat16) must win.
@@ -346,6 +352,7 @@ def __init__(
         # During model initialization, the default dtype is set as the model
         # weight and activation dtype.
         dtype = torch.get_default_dtype()
+        self.act_dtype = dtype
         if attn_backend is None:
             self.attn_backend = get_attn_backend(
                 head_size,
@@ -469,7 +476,11 @@ def __init__(
         # for attn backends supporting query quantization
         self.query_quant = None
         if (
-            self.impl.supports_quant_query_input
+            # A model whose fused QK-norm/RoPE kernel already emits the
+            # quantized query (MiniMax-M3's DSpark draft) passes
+            # query_pre_quantized and skips this second quantization.
+            not query_pre_quantized
+            and self.impl.supports_quant_query_input
             and (
                 self.kv_cache_dtype.startswith("fp8")
                 or self.kv_cache_dtype.startswith("nvfp4")
@@ -591,8 +602,8 @@ def extra_repr(self) -> str:
         s += f", backend={self.impl.__class__.__name__}"
         return s
 
-    def process_weights_after_loading(self, act_dtype: torch.dtype):
-        self.impl.process_weights_after_loading(act_dtype)
+    def process_weights_after_loading(self) -> None:
+        self.impl.process_weights_after_loading(self.act_dtype)
 
         # If we should not load quant weights, we initialize the scales to 1.0
         # as the default value. See [Note: Register q/k/v/prob scales in state dict]
diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py
index 173c40dc89de..107a3f1c3ce4 100644
--- a/vllm/model_executor/layers/attention/mla_attention.py
+++ b/vllm/model_executor/layers/attention/mla_attention.py
@@ -504,6 +504,7 @@ def __init__(
             )
 
         dtype = torch.get_default_dtype()
+        self.act_dtype = dtype
         if attn_backend is not None:
             assert attn_backend.is_mla(), (
                 f"MLAAttention: attn_backend must be an MLA backend, "
@@ -1239,7 +1240,8 @@ def _use_sparse_mha(self, attn_metadata: "MLACommonMetadata") -> bool:
             self._vllm_config.attention_config.sparse_mla_force_mqa
         )
 
-    def process_weights_after_loading(self, act_dtype: torch.dtype):
+    def process_weights_after_loading(self) -> None:
+        act_dtype = self.act_dtype
         # Let per-backend impls do their own weight packing first (no-op
         # unless overridden), mirroring Attention.process_weights_after_loading.
         self.impl.process_weights_after_loading(act_dtype)
@@ -1382,9 +1384,10 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
             kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype),
             # ds_mla layouts pack NoPE + RoPE + scales into one opaque per-token
             # blob, so the size is not derivable from head_size.
-            # See flashmla_sparse.py.
-            state_content_bytes={"fp8_ds_mla": 656, "nvfp4_ds_mla": 352}.get(
-                self.kv_cache_dtype
+            state_content_bytes=(
+                cast(MLAAttentionImpl, self.impl).get_fp8_ds_mla_row_bytes()
+                if self.kv_cache_dtype == "fp8_ds_mla"
+                else {"nvfp4_ds_mla": 352}.get(self.kv_cache_dtype)
             ),
         )
         if self.sliding_window is not None:
diff --git a/vllm/model_executor/layers/attention/mm_encoder_attention.py b/vllm/model_executor/layers/attention/mm_encoder_attention.py
index 0af2396ed719..fdc4d8b7118b 100644
--- a/vllm/model_executor/layers/attention/mm_encoder_attention.py
+++ b/vllm/model_executor/layers/attention/mm_encoder_attention.py
@@ -456,13 +456,8 @@ def _init_fp8_state(self) -> None:
             _fp8_scale_save_path = mm_cfg.mm_encoder_fp8_scale_save_path
             _fp8_scale_save_margin = mm_cfg.mm_encoder_fp8_scale_save_margin
 
-    def process_weights_after_loading(self, act_dtype: torch.dtype) -> None:
-        """Populate FP8 scale buffers after weights are loaded.
-
-        ``act_dtype`` matches the signature used by :class:`Attention` and
-        :class:`MLAAttention` for the loader auto-scan but is unused:
-        FP8 scales are always float32.
-        """
+    def process_weights_after_loading(self) -> None:
+        """Populate FP8 scale buffers (always float32) after weights are loaded."""
         if not self.fp8_enabled:
             return
 
diff --git a/vllm/model_executor/layers/attention/sparse_mla_attention.py b/vllm/model_executor/layers/attention/sparse_mla_attention.py
index 2f6583f1fcac..cfc3009f52e9 100644
--- a/vllm/model_executor/layers/attention/sparse_mla_attention.py
+++ b/vllm/model_executor/layers/attention/sparse_mla_attention.py
@@ -689,7 +689,11 @@ def __init__(
                 index_group_builder.register_layer(
                     indexer is not None,
                     vllm_config,
-                    head_size=head_size,
+                    row_width=(
+                        self.get_fp8_ds_mla_row_bytes()
+                        if kv_cache_dtype == "fp8_ds_mla"
+                        else head_size
+                    ),
                     kv_cache_dtype=kv_cache_dtype,
                 )
             )
diff --git a/vllm/model_executor/layers/fused_moe/all2all_utils.py b/vllm/model_executor/layers/fused_moe/all2all_utils.py
index 3fbeebaaa88d..2c99cdff4f6d 100644
--- a/vllm/model_executor/layers/fused_moe/all2all_utils.py
+++ b/vllm/model_executor/layers/fused_moe/all2all_utils.py
@@ -119,6 +119,7 @@ def flashinfer_one_sided_dispatch_layout(
         from .prepare_finalize.deepep_v2 import DeepEPV2PrepareAndFinalize
     if has_mori():
         from .prepare_finalize.mori import MoriPrepareAndFinalize
+        from .prepare_finalize.mori_epv2 import MoriEPv2PrepareAndFinalize
     if has_nixl_ep():
         from .prepare_finalize.nixl_ep import (
             NIXL_EP_QUANT_BLOCK_SHAPE,
@@ -334,7 +335,31 @@ def maybe_make_prepare_finalize(
             num_global_experts=moe.num_experts,
         )
 
-    elif moe.use_mori_kernels:
+    elif moe.use_mori_epv2_kernels:
+        num_local_experts = moe.num_experts // all2all_manager.world_size
+        # MORI EPv2 routes expert e to rank e // num_local_experts, so every rank
+        # must hold the same number of experts.
+        if num_local_experts * all2all_manager.world_size != moe.num_experts:
+            raise ValueError(
+                f"MORI EPv2 needs the {moe.num_experts} experts to split evenly "
+                f"over {all2all_manager.world_size} expert-parallel ranks"
+            )
+        all_to_all_args = dict(
+            hidden_dim=moe.hidden_dim,
+            max_num_inp_token_per_rank=moe.max_num_tokens,
+            num_experts_per_rank=num_local_experts,
+            num_experts_per_token=moe.experts_per_token,
+        )
+        handle = all2all_manager.get_handle(all_to_all_args)
+        prepare_finalize = MoriEPv2PrepareAndFinalize(
+            handle,
+            max_tokens_per_rank=moe.max_num_tokens,
+            num_dispatchers=all2all_manager.world_size,
+            num_local_experts=num_local_experts,
+            rank_expert_offset=all2all_manager.rank * num_local_experts,
+        )
+
+    elif moe.use_mori_epv1_kernels:
         assert quant_config is not None
 
         # MXFP4 activations are dispatched as packed FP4 with their e8m0
diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py
index e8bd822becab..8f37fa20206c 100644
--- a/vllm/model_executor/layers/fused_moe/config.py
+++ b/vllm/model_executor/layers/fused_moe/config.py
@@ -16,7 +16,10 @@
     OCP_MX_DTYPES,
     OCP_MX_Scheme,
 )
-from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    GroupShape,
+    QuantKey,
+)
 from vllm.platforms import current_platform
 from vllm.utils.import_utils import get_triton_kernels_version
 from vllm.utils.math_utils import cdiv
@@ -1026,6 +1029,7 @@ class FusedMoEParallelConfig:
     use_ep: bool  # whether to use EP or not
     all2all_backend: str  # all2all backend for MoE communication
     enable_eplb: bool  # whether to enable expert load balancing
+    mori_ep_v2: bool = False  # whether the "mori" all2all backend runs EPv2
 
     @property
     def is_sequence_parallel(self) -> bool:
@@ -1084,12 +1088,22 @@ def use_passthrough_all2all(self):
         return self.all2all_backend == "passthrough"
 
     @property
-    def use_mori_kernels(self):
+    def use_mori_epv1_kernels(self):
+        if self.all2all_backend == "mori":
+            return self.use_all2all_kernels and not self.mori_ep_v2
         return self.use_all2all_kernels and self.all2all_backend in (
             "mori_high_throughput",
             "mori_low_latency",
         )
 
+    @property
+    def use_mori_epv2_kernels(self):
+        return (
+            self.use_all2all_kernels
+            and self.all2all_backend == "mori"
+            and self.mori_ep_v2
+        )
+
     @property
     def use_nixl_ep_kernels(self):
         return self.use_all2all_kernels and self.all2all_backend == "nixl_ep"
@@ -1222,6 +1236,7 @@ def make(
                 use_ep=False,
                 all2all_backend=vllm_parallel_config.all2all_backend,
                 enable_eplb=vllm_parallel_config.enable_eplb,
+                mori_ep_v2=bool(vllm_parallel_config.mori_ep_v2),
             )
         # DP + EP / TP + EP / DP + TP + EP
         assert use_ep
@@ -1242,6 +1257,7 @@ def make(
             use_ep=True,
             all2all_backend=vllm_parallel_config.all2all_backend,
             enable_eplb=vllm_parallel_config.enable_eplb,
+            mori_ep_v2=bool(vllm_parallel_config.mori_ep_v2),
         )
 
     @classmethod
@@ -1338,7 +1354,13 @@ class FusedMoEConfig:
     # 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
+    is_fused_shared_expert_requested: bool = False
+    # Activation scheme the routed experts were trained (QAT) against, when the
+    # checkpoint does not record it. Quant methods use it in place of their
+    # default: MiniMax-M3.1 quantizes the FC1 input per token and ships a 1.0
+    # ``input_global_scale`` placeholder, which is indistinguishable from a
+    # calibrated per-tensor scale of 1.0.
+    qat_activation_key: QuantKey | None = None
     # Whether to force MXFP4 (a4w4) MoE activations for DeepSeek V4.1 on
     # ROCm/AITER. Opt-in via VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4=1; rejected
     # for any other model type. Resolved here, not in the forward path,
@@ -1371,7 +1393,7 @@ def __post_init__(self):
 
         if self.is_act_and_mul:
             self.rocm_aiter_fmoe_enabled = rocm_aiter_ops.is_fused_moe_enabled()
-            self.aiter_fmoe_shared_expert_enabled = (
+            self.is_fused_shared_expert_requested = bool(
                 rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
             )
 
@@ -1398,11 +1420,12 @@ def __post_init__(self):
                 )
             self.use_mxfp4_w4a4_dsv4 = True
 
-        if self.use_mori_kernels:
+        if self.use_mori_epv1_kernels:
             assert self.rocm_aiter_fmoe_enabled, (
-                "Mori needs to be used with aiter fused_moe for now."
+                "MoRI EPv1 needs to be used with aiter fused_moe for now."
             )
-            assert not self.aiter_fmoe_shared_expert_enabled, (
+        if self.use_mori_epv1_kernels or self.use_mori_epv2_kernels:
+            assert not self.is_fused_shared_expert_requested, (
                 "Mori does not support fusion shared expert now. "
                 "Turn it off by setting VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=0"
             )
@@ -1553,8 +1576,12 @@ def use_passthrough_all2all(self):
         return self.moe_parallel_config.use_passthrough_all2all
 
     @property
-    def use_mori_kernels(self):
-        return self.moe_parallel_config.use_mori_kernels
+    def use_mori_epv1_kernels(self):
+        return self.moe_parallel_config.use_mori_epv1_kernels
+
+    @property
+    def use_mori_epv2_kernels(self):
+        return self.moe_parallel_config.use_mori_epv2_kernels
 
     @property
     def use_fi_nvl_two_sided_kernels(self):
diff --git a/vllm/model_executor/layers/fused_moe/deep_gemm_mega_moe.py b/vllm/model_executor/layers/fused_moe/deep_gemm_mega_moe.py
index 986f07238694..99c5c87736b0 100644
--- a/vllm/model_executor/layers/fused_moe/deep_gemm_mega_moe.py
+++ b/vllm/model_executor/layers/fused_moe/deep_gemm_mega_moe.py
@@ -27,9 +27,6 @@
     prepare_megamoe_inputs,
     prepare_megamoe_nvfp4_inputs,
 )
-from vllm.model_executor.layers.fused_moe.router.base_router import (
-    eplb_map_to_physical_and_record,
-)
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
     QuantKey,
     kMxfp8Dynamic,
@@ -38,7 +35,6 @@
 from vllm.model_executor.models.utils import extract_layer_index
 from vllm.model_executor.utils import set_weight_attrs
 from vllm.platforms import current_platform
-from vllm.v1.worker.ubatching import dbo_current_ubatch_id
 
 logger = init_logger(__name__)
 
@@ -768,30 +764,10 @@ def forward(
             if is_padding is not None:
                 is_padding = is_padding[:num_tokens]
 
+        # Capture logical ids; staging maps them to EPLB physical replicas.
         if self.capture_fn is not None:
             self.capture_fn(topk_ids)
 
-        # EPLB: map logical expert IDs to physical replicas and record load.
-        eplb_state = self.eplb_state
-        if eplb_state.logical_to_physical_map is not None:
-            assert eplb_state.expert_load_view is not None
-            assert eplb_state.logical_replica_count is not None
-            assert eplb_state.should_record_tensor is not None
-            if is_padding is not None:
-                topk_ids = torch.where(is_padding.unsqueeze(1), -1, topk_ids)
-            topk_ids = eplb_map_to_physical_and_record(
-                topk_ids=topk_ids,
-                expert_load_view=eplb_state.expert_load_view,
-                logical_to_physical_map=eplb_state.logical_to_physical_map,
-                logical_replica_count=eplb_state.logical_replica_count,
-                record_enabled=eplb_state.should_record_tensor,
-                num_unpadded_tokens=eplb_state.num_unpadded_tokens_tensors[
-                    dbo_current_ubatch_id()
-                ]
-                if eplb_state.num_unpadded_tokens_tensors is not None
-                else None,
-            )
-
         from vllm.utils.deep_gemm import _import_deep_gemm
 
         deep_gemm = _import_deep_gemm()
@@ -823,6 +799,7 @@ def forward(
                 w2_global_scale=self.w2_weight_scale_2,
                 l2_alpha=self._l2_alpha,
                 l2_activation_scale=l2_activation_scale,
+                eplb_state=self.eplb_state,
             )
             kwargs.update(
                 l1_alpha=self.w13_weight_scale_2,
@@ -844,6 +821,7 @@ def forward(
                 shared_x_sf=shared_x_sf,
                 shared_block_m=shared_block_m,
                 hidden_quant=fmt.hidden_quant,
+                eplb_state=self.eplb_state,
             )
 
         if self.has_fused_shared_experts:
diff --git a/vllm/model_executor/layers/fused_moe/eplb_triton.py b/vllm/model_executor/layers/fused_moe/eplb_triton.py
new file mode 100644
index 000000000000..f603bdf1d0a7
--- /dev/null
+++ b/vllm/model_executor/layers/fused_moe/eplb_triton.py
@@ -0,0 +1,47 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""EPLB replica mapping and expert-load recording shared by Triton kernels."""
+
+from vllm.triton_utils import tl, triton
+
+
+@triton.jit
+def eplb_map_and_record(
+    logical_ids,
+    token_idx,
+    mask,
+    logical_to_physical,
+    replica_count,
+    num_logical,
+    map_stride,
+    expert_load,
+    num_physical,
+    record_enabled,
+    num_unpadded_tokens,
+):
+    """Map logical expert ids to physical replicas and record expert load.
+
+    ``token_idx`` broadcasts against ``logical_ids``; masked-off or
+    out-of-range ids map to -1.
+    """
+    valid = mask & (logical_ids >= 0) & (logical_ids < num_logical)
+    safe_ids = tl.where(valid, logical_ids, 0)
+    replicas = tl.load(replica_count + safe_ids, mask=valid, other=1)
+    # floor(2^32 / phi), classic Knuth multiplicative hash multiplier.
+    KNUTH_MULTIPLIER = 2654435769
+    hashed = (token_idx.to(tl.int64) * KNUTH_MULTIPLIER) & 0xFFFFFFFF
+    physical_ids = tl.load(
+        logical_to_physical + safe_ids * map_stride + hashed % replicas,
+        mask=valid,
+        other=-1,
+    )
+
+    record = (
+        (tl.load(record_enabled) != 0)
+        & (physical_ids >= 0)
+        & (physical_ids < num_physical)
+    )
+    if num_unpadded_tokens is not None:
+        record &= token_idx < tl.load(num_unpadded_tokens)
+    tl.atomic_add(expert_load + tl.where(record, physical_ids, 0), 1, mask=record)
+    return physical_ids
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 c1e3231c2048..501dd9ba43ce 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
@@ -488,7 +488,7 @@ def expects_unquantized_inputs(self) -> bool:
         # so we should not defer input quantization.
         # Otherwise, AITER fused MoE kernels handle input quantization
         # internally via a single fused kernel.
-        return not self.moe_config.use_mori_kernels
+        return not self.moe_config.use_mori_epv1_kernels
 
     @staticmethod
     def activation_format() -> mk.FusedMoEActivationFormat:
@@ -554,7 +554,7 @@ def workspace_shapes(
         # MoRI dispatches MXFP4 activations packed two per byte, so K is
         # half the hidden size there.
         packed_fp4_input = (
-            self.moe_config.use_mori_kernels and self.quant_config.use_mxfp4_w4a4
+            self.moe_config.use_mori_epv1_kernels and self.quant_config.use_mxfp4_w4a4
         )
         output = (M, K * 2 if packed_fp4_input else K)
         return (workspace1, workspace2, output)
diff --git a/vllm/model_executor/layers/fused_moe/fused_moe.py b/vllm/model_executor/layers/fused_moe/fused_moe.py
index e06288a3e92b..7e9f0e28f6ac 100644
--- a/vllm/model_executor/layers/fused_moe/fused_moe.py
+++ b/vllm/model_executor/layers/fused_moe/fused_moe.py
@@ -299,6 +299,7 @@ def fused_moe_kernel(
     # Pointers to matrices
     a_ptr,
     b_ptr,
+    b_desc,
     c_ptr,
     b_bias_ptr,
     a_scale_ptr,
@@ -351,6 +352,7 @@ def fused_moe_kernel(
     SWAP_AB: tl.constexpr,
     # Tensor-descriptor path for the A gather and B load in the K-loop.
     USE_TD: tl.constexpr = False,
+    USE_TMA_B: tl.constexpr = False,
 ):
     """Implements the fused computation for a Mixture of Experts (MOE) using
     token and expert matrices.
@@ -441,6 +443,7 @@ def fused_moe_kernel(
     offs_k = tl.arange(0, BLOCK_SIZE_K)
     # TD gather and the SWAP_AB accumulator layout are mutually exclusive.
     tl.static_assert(not (USE_TD and SWAP_AB))
+    tl.static_assert(not (USE_TD and USE_TMA_B))
     if USE_TD:
         # ``tt.descriptor_gather`` requires block_shape[0] == 1 and i32 idx.
         m_td = num_valid_tokens // top_k
@@ -523,16 +526,26 @@ def fused_moe_kernel(
             b = b_desc.load([pid_n * BLOCK_SIZE_N, k * BLOCK_SIZE_K]).T
         elif SWAP_AB:
             a_mask = (offs_k[:, None] < K - k * BLOCK_SIZE_K) & token_mask[None, :]
-            b_mask = offs_k[None, :] < K - k * BLOCK_SIZE_K
             a = tl.load(a_ptrs, mask=a_mask, other=0.0)
-            b = tl.load(b_ptrs, mask=b_mask, other=0.0)
+            if not USE_TMA_B:
+                b_mask = offs_k[None, :] < K - k * BLOCK_SIZE_K
+                b = tl.load(b_ptrs, mask=b_mask, other=0.0)
         else:
             a = tl.load(
                 a_ptrs,
                 mask=token_mask[:, None] & (offs_k[None, :] < K - k * BLOCK_SIZE_K),
                 other=0.0,
             )
-            b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0)
+            if not USE_TMA_B:
+                b = tl.load(
+                    b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0
+                )
+        if USE_TMA_B:
+            b = b_desc.load(
+                [off_experts.to(tl.int32), pid_n * BLOCK_SIZE_N, k * BLOCK_SIZE_K]
+            ).reshape((BLOCK_SIZE_N, BLOCK_SIZE_K))
+            if not SWAP_AB:
+                b = b.T
         # We accumulate along the K dimension.
         if use_int8_w8a16:
             accumulator = tl.dot(a, b.to(compute_type), acc=accumulator)
@@ -842,6 +855,26 @@ def invoke_fused_moe_triton_kernel(
     BLOCK_SIZE_K = config.pop("BLOCK_SIZE_K")
     if block_shape is not None:
         BLOCK_SIZE_K = min(BLOCK_SIZE_K, min(block_shape[0], block_shape[1]))
+    use_tma_b = (
+        bool(config.pop("USE_TMA", True))
+        and use_fp8_w8a8
+        and block_shape is not None
+        and all(block_shape)
+        and current_platform.is_cuda()
+        and current_platform.is_device_capability(90)
+        and B.device.type == "cuda"
+        and B.is_contiguous()
+        and B.size(1) % config["BLOCK_SIZE_N"] == 0
+        and B.size(2) % BLOCK_SIZE_K == 0
+    )
+    b_desc = None
+    if use_tma_b:
+        from triton.tools.tensor_descriptor import TensorDescriptor
+
+        set_triton_allocator(B.device)
+        b_desc = TensorDescriptor.from_tensor(
+            B, [1, config["BLOCK_SIZE_N"], BLOCK_SIZE_K]
+        )
     if use_td and A.size(1) % BLOCK_SIZE_K != 0:
         # TD gather/load feeding tl.dot with a non-block-aligned K
         # miscompiles (~74% of output elements wrong) on real HW;
@@ -864,6 +897,7 @@ def invoke_fused_moe_triton_kernel(
     fused_moe_kernel[grid](
         A,
         B,
+        b_desc,
         C,
         B_bias,
         A_scale,
@@ -906,6 +940,7 @@ def invoke_fused_moe_triton_kernel(
         BLOCK_SIZE_K=BLOCK_SIZE_K,
         SWAP_AB=SWAP_AB,
         USE_TD=use_td,
+        USE_TMA_B=use_tma_b,
         **config,
     )
 
diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py
index 3b3e2aa39973..cbe6d0ce8938 100644
--- a/vllm/model_executor/layers/fused_moe/layer.py
+++ b/vllm/model_executor/layers/fused_moe/layer.py
@@ -40,6 +40,7 @@
 from vllm.model_executor.layers.quantization.base_config import (
     QuantizationConfig,
 )
+from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey
 
 logger = init_logger(__name__)
 
@@ -114,6 +115,7 @@ def FusedMoEFactory(
     swiglu_beta: float | None = None,
     activation_situ_beta: float | None = None,
     activation_situ_linear_beta: float | None = None,
+    qat_activation_key: QuantKey | None = None,
     e_score_correction_bias: torch.Tensor | None = None,
     apply_router_weight_on_input: bool = False,
     activation: str = "silu",
@@ -186,6 +188,8 @@ def FusedMoEFactory(
         swiglu_limit: SwiGLU activation limit
         activation_situ_beta: SituGLU activation beta
         activation_situ_linear_beta: SituGLU linear beta
+        qat_activation_key: Activation quantization scheme the experts were
+            trained against, for checkpoints that do not record it
         e_score_correction_bias: Expert score correction bias tensor
         apply_router_weight_on_input: Whether to apply router weights on input
         activation: Activation function name ("silu", "gelu", etc.)
@@ -387,6 +391,7 @@ def FusedMoEFactory(
         swiglu_beta=swiglu_beta,
         activation_situ_beta=activation_situ_beta,
         activation_situ_linear_beta=activation_situ_linear_beta,
+        qat_activation_key=qat_activation_key,
         max_capture_size=vllm_config.compilation_config.max_cudagraph_capture_size,
         skip_final_all_reduce=skip_final_all_reduce,
     )
diff --git a/vllm/model_executor/layers/fused_moe/oracle/fp8.py b/vllm/model_executor/layers/fused_moe/oracle/fp8.py
index 0aed7b16f2ef..8de9f1a21eb3 100644
--- a/vllm/model_executor/layers/fused_moe/oracle/fp8.py
+++ b/vllm/model_executor/layers/fused_moe/oracle/fp8.py
@@ -39,7 +39,11 @@
     kFp8Dynamic128Sym,
     kFp8Static128BlockSym,
 )
-from vllm.model_executor.utils import replace_parameter
+from vllm.model_executor.utils import (
+    move_parameter,
+    replace_parameter,
+    set_kernel_parameter,
+)
 from vllm.platforms import current_platform
 from vllm.utils.deep_gemm import should_auto_disable_deep_gemm
 from vllm.utils.math_utils import round_up
@@ -639,6 +643,14 @@ def convert_to_fp8_moe_kernel_format(
         convert_to_humming_moe_kernel_format(
             layer, quant_config=_humming_fp8_weight_schema(layer, w13, w13_scale)
         )
+        if scale_name == "weight_scale_inv":
+            # One param under both names: the loaded `_inv` holds the runtime scale.
+            for prefix in ("w13", "w2"):
+                scale = layer._parameters.pop(f"{prefix}_weight_scale").data
+                move_parameter(
+                    layer, f"{prefix}_{scale_name}", f"{prefix}_weight_scale"
+                )
+                set_kernel_parameter(layer, f"{prefix}_weight_scale", scale)
         w13 = layer.w13_weight
         w2 = layer.w2_weight
         w13_scale = layer.w13_weight_scale
diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
index 9de3732f00f1..3e838b4e1d53 100644
--- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
+++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
@@ -82,6 +82,19 @@
         )
 
 
+def store_params_in_triton(layer: torch.nn.Module, **tensors) -> None:
+    """Back each loaded param and its triton_kernels tensor by one storage, so the
+    param keeps its loaded bytes without a second copy of the data."""
+    for name, tensor in tensors.items():
+        replace_parameter(layer, name, tensor.storage.data)
+        tensor.storage.data = getattr(layer, name).data
+
+
+def triton_weight(layer: torch.nn.Module, name: str):
+    """The triton_kernels tensor stored for weight `name`, else the param."""
+    return getattr(layer, f"{name}_triton", getattr(layer, name))
+
+
 def _pack_deepgemm_mxfp4_scales(
     w13_weight: torch.Tensor,
     w2_weight: torch.Tensor,
@@ -1433,8 +1446,7 @@ def swap_every_two_rows(x, axis=-1):
             flex_ctx=FlexCtx(rhs_data=w2_flex, lhs_data=lhs_data2),
         )
 
-        del layer.w13_weight
-        del layer.w2_weight
+        store_params_in_triton(layer, w13_weight=w13_weight, w2_weight=w2_weight)
 
         return (
             w13_weight,
@@ -1467,13 +1479,14 @@ def swap_every_two_rows(x, axis=-1):
             **_mx_scale_kwargs(w2_scale), flex_ctx=FlexCtx(rhs_data=w2_flex)
         )
 
-        # The original mxfp4 block scales have been swizzled into the
-        # precision configs above and are no longer read by the kernel, so
-        # drop the now-dead weight/scale Parameters to free their memory.
-        del layer.w13_weight
-        del layer.w2_weight
-        del layer.w13_weight_scale
-        del layer.w2_weight_scale
+        # The swizzled weights and scales live in the params' storage.
+        store_params_in_triton(
+            layer,
+            w13_weight=w13_weight,
+            w2_weight=w2_weight,
+            w13_weight_scale=w13_scale,
+            w2_weight_scale=w2_scale,
+        )
 
         return (
             w13_weight,
@@ -1929,13 +1942,14 @@ def shuffle_weight(w: torch.Tensor) -> torch.Tensor:
             **_mx_scale_kwargs(w2_scale), flex_ctx=FlexCtx(rhs_data=w2_flex)
         )
 
-        # The original mxfp4 block scales have been swizzled into the
-        # precision configs above and are no longer read by the kernel, so
-        # drop the now-dead weight/scale Parameters to free their memory.
-        del layer.w13_weight
-        del layer.w2_weight
-        del layer.w13_weight_scale
-        del layer.w2_weight_scale
+        # The swizzled weights and scales live in the params' storage.
+        store_params_in_triton(
+            layer,
+            w13_weight=w13_weight,
+            w2_weight=w2_weight,
+            w13_weight_scale=w13_scale,
+            w2_weight_scale=w2_scale,
+        )
 
         return (
             w13_weight,
diff --git a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py
index 80aa9122f447..d265f7629af0 100644
--- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py
+++ b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py
@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from enum import Enum
+from typing import Any
 
 import torch
 
@@ -36,6 +37,7 @@
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
     QuantKey,
 )
+from vllm.model_executor.utils import move_parameter
 
 logger = init_logger(__name__)
 
@@ -54,6 +56,15 @@ class NvFp4MoeBackend(Enum):
     FLASHINFER_MOE_EP_CUTEDSL = "FLASHINFER_MOE_EP_CUTEDSL"
 
 
+# Backends whose experts quantize the MoE input themselves, with one scale per
+# token. ``make_nvfp4_moe_kernel`` only hands ``per_token_activation`` to these.
+PER_TOKEN_NVFP4_MOE_BACKENDS = frozenset(
+    {
+        NvFp4MoeBackend.FLASHINFER_CUTEDSL,
+        NvFp4MoeBackend.FLASHINFER_TRTLLM,
+    }
+)
+
 FLASHINFER_NVFP4_MOE_BACKENDS = [
     NvFp4MoeBackend.FLASHINFER_TRTLLM,
     NvFp4MoeBackend.FLASHINFER_CUTLASS,
@@ -196,6 +207,18 @@ def select_nvfp4_moe_backend(
     """Select the primary NvFP4 MoE backend
     Note: Shape-specific fallbacks may still occur at runtime.
     """
+    from vllm.model_executor.layers.quantization.online.requantization import (
+        is_requantization_source,
+    )
+
+    # TODO: remove and let the oracles have an option to skip init.
+    if is_requantization_source():
+        from vllm.model_executor.layers.fused_moe.experts.nvfp4_emulation_moe import (
+            Nvfp4QuantizationEmulationTritonExperts,
+        )
+
+        return NvFp4MoeBackend.EMULATION, Nvfp4QuantizationEmulationTritonExperts
+
     # 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
@@ -441,10 +464,10 @@ def convert_to_nvfp4_moe_kernel_format(
                 "strategy": "group",
                 "group_size": 16,
             }
-            # CT pack-quantized reads `weight_packed`; the method renamed it to
-            # `weight`. Re-alias (convert replaces all params anyway).
-            layer.w13_weight_packed = layer.w13_weight
-            layer.w2_weight_packed = layer.w2_weight
+            # CT pack-quantized reads `weight_packed`; the method aliased it as
+            # `weight`. Make the loaded `*_weight_packed` name the runtime one.
+            move_parameter(layer, "w13_weight", "w13_weight_packed")
+            move_parameter(layer, "w2_weight", "w2_weight_packed")
         else:
             quant_config = {"quant_method": "modelopt", "quant_algo": "nvfp4"}
 
@@ -615,6 +638,10 @@ def make_nvfp4_moe_kernel(
     routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
     per_token_activation: bool = False,
 ) -> mk.FusedMoEKernel:
+    # Marlin consumes A16 and must not receive per-token A4 metadata.
+    if backend == NvFp4MoeBackend.MARLIN and moe_quant_config.use_nvfp4_w4a16:
+        per_token_activation = False
+
     # Create Prepare/Finalize.
     prepare_finalize = maybe_make_prepare_finalize(
         moe=moe_config,
@@ -628,15 +655,17 @@ def make_nvfp4_moe_kernel(
 
     logger.info_once("Using %s", prepare_finalize.__class__.__name__)
 
-    extra_kwargs = {}
-    if (
-        backend
-        in (
-            NvFp4MoeBackend.FLASHINFER_TRTLLM,
-            NvFp4MoeBackend.FLASHINFER_CUTEDSL,
-        )
-        and per_token_activation
-    ):
+    extra_kwargs: dict[str, Any] = {}
+    if per_token_activation:
+        # Prepare dispatched the input unquantized, so an expert that does
+        # not quantize it itself would read garbage. The backend selection
+        # does not catch this: several NVFP4 experts match on the weight key
+        # alone and accept any activation key.
+        if backend not in PER_TOKEN_NVFP4_MOE_BACKENDS:
+            raise ValueError(
+                f"Per-token NVFP4 activation scales are not supported by "
+                f"{backend.value}."
+            )
         extra_kwargs["per_token_activation"] = True
 
     # Create Experts.
diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/mori_epv2.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/mori_epv2.py
new file mode 100644
index 000000000000..d19b60856729
--- /dev/null
+++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/mori_epv2.py
@@ -0,0 +1,244 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from typing import Any
+
+import torch
+
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
+from vllm.distributed import get_dp_group
+from vllm.forward_context import get_forward_context, is_forward_context_available
+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 (
+    TopKWeightAndReduceNoOP,
+)
+
+logger = init_logger(__name__)
+
+# Receive-row bounds are rounded up to this many rows (as SGLang's MORI path does).
+_RECV_ROW_ALIGN = 32
+
+
+class MoriEPv2PrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
+    def __init__(
+        self,
+        mori_op: Any,
+        max_tokens_per_rank: int,
+        num_dispatchers: int,
+        num_local_experts: int,
+        rank_expert_offset: int,
+    ):
+        super().__init__()
+        self.mori_op = mori_op
+        self.num_dispatchers_ = num_dispatchers
+        self.max_tokens_per_rank = max_tokens_per_rank
+        self.num_local_experts = num_local_experts
+        self.rank_expert_offset = rank_expert_offset
+        self._padding_expert_ids: torch.Tensor | None = None
+        self._routing = None
+        self._experts_read_recv_rows_only = False
+
+    @property
+    def activation_format(self) -> mk.FusedMoEActivationFormat:
+        return mk.FusedMoEActivationFormat.Standard
+
+    def post_init_setup(self, fused_experts: mk.FusedMoEExperts):
+        from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import (
+            AiterExperts,
+        )
+
+        if fused_experts.moe_config.in_dtype != torch.bfloat16:
+            raise ValueError(
+                "MORI EPv2 only dispatches BF16 activations; got "
+                f"{fused_experts.moe_config.in_dtype}. Set VLLM_MORI_EP_V2=0 to "
+                "use MORI EPv1."
+            )
+        quant_dtype = fused_experts.quant_config.quant_dtype
+        if quant_dtype is not None and not fused_experts.expects_unquantized_inputs:
+            raise ValueError(
+                "MORI EPv2 dispatches unquantized BF16 activations, but "
+                f"{type(fused_experts).__name__} expects {quant_dtype} "
+                "activations quantized before dispatch. Use AITER fused MoE experts "
+                "(VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_MOE=1), or set "
+                "VLLM_MORI_EP_V2=0 for MORI EPv1, which also requires AITER fused MoE."
+            )
+
+        if quant_dtype is not None and fused_experts.quant_config.is_per_tensor:
+            raise ValueError(
+                "MORI EPv2 does not support per-tensor activation quantization: "
+                "the experts would quantize the received rows, which AITER's "
+                "per-tensor quant does not support. Set VLLM_MORI_EP_V2=0 to use "
+                "MORI EPv1, which quantizes before dispatch."
+            )
+
+        # AITER computes only the first expert_num_tokens rows.
+        self._experts_read_recv_rows_only = isinstance(fused_experts, AiterExperts)
+
+    def allocate_fused_expert_output(
+        self,
+        shape: tuple[int, ...],
+        dtype: torch.dtype,
+        device: torch.device,
+        weight_and_reduce_impl: mk.TopKWeightAndReduce,
+    ) -> torch.Tensor | None:
+        # Experts that reduce top-k themselves write straight into the arena
+        # rows combine sends, and combine skips its copy-in.
+        if not isinstance(weight_and_reduce_impl, TopKWeightAndReduceNoOP):
+            return None
+        buf = self.mori_op.combine_in_view()
+        if (
+            len(shape) != 2
+            or buf.dtype != dtype
+            or buf.device != device
+            or shape[1] != buf.shape[1]
+            or shape[0] > buf.shape[0]
+        ):
+            return None
+        return buf[: shape[0]]
+
+    def output_is_reduced(self) -> bool:
+        return True
+
+    def num_dispatchers(self):
+        return self.num_dispatchers_
+
+    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
+
+    def _route_padding_locally(self, topk_ids: torch.Tensor) -> torch.Tensor:
+        # DP padding tokens carry expert id -1 (VLLM_MOE_SKIP_PADDING) and zero
+        # weight. MORI EPv2 has no invalid-expert slot and its FlyDSL dispatch
+        # faults on -1, so send them to this rank's own experts instead.
+        if self._padding_expert_ids is None:
+            slots = torch.arange(topk_ids.shape[1], device=topk_ids.device)
+            self._padding_expert_ids = (
+                self.rank_expert_offset + slots % self.num_local_experts
+            ).to(topk_ids.dtype)
+        return torch.where(topk_ids < 0, self._padding_expert_ids, topk_ids)
+
+    def _recv_rows_bound(self, a1: torch.Tensor, capacity: int) -> int:
+        """Host-side upper bound on the rows this rank can receive this step.
+
+        MORI receives each source token at most once per rank, so the receive
+        count is at most the tokens all dispatchers send: the sum of the DP
+        batches, times the dispatchers per DP rank (each TP rank may send the
+        same batch). Slicing the capacity-sized buffers (world_size *
+        max_num_tokens rows) to this bound keeps the row masking, the experts
+        and combine's copy-in off the unused tail. The bound is rounded up to a
+        multiple of _RECV_ROW_ALIGN to limit the distinct expert shapes.
+        """
+        dp_meta = (
+            get_forward_context().dp_metadata
+            if is_forward_context_available()
+            else None
+        )
+        if dp_meta is None:
+            return capacity
+        tokens = dp_meta.num_tokens_across_dp_cpu
+        n_dp = tokens.numel()
+        dp_rank = get_dp_group().rank_in_group
+        if dp_rank >= n_dp or a1.shape[0] > int(tokens[dp_rank]):
+            # Metadata that does not cover this rank's own batch is stale; a
+            # bound derived from it could cut off received rows.
+            logger.debug_once(
+                "MORI EPv2: DP token counts do not cover this rank's batch; "
+                "using the full receive buffer."
+            )
+            return capacity
+        if n_dp > 0 and self.num_dispatchers_ % n_dp == 0:
+            bound = self.num_dispatchers_ // n_dp * int(tokens.sum())
+        else:
+            bound = self.num_dispatchers_ * int(tokens.max())
+        bound = max(bound, 1)
+        bound = -(-bound // _RECV_ROW_ALIGN) * _RECV_ROW_ALIGN
+        return min(capacity, bound)
+
+    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:
+        del num_experts, expert_map, defer_input_quant
+        assert not apply_router_weight_on_input, (
+            "MORI EPv2 does not support apply_router_weight_on_input=True"
+        )
+        if a1.dtype != torch.bfloat16:
+            raise ValueError(f"MORI EPv2 requires BF16 dispatch input, got {a1.dtype}")
+
+        (
+            dispatch_a1,
+            dispatch_weights,
+            dispatch_scale,
+            dispatch_ids,
+            dispatch_recv_token_num,
+            self._routing,
+        ) = self.mori_op.dispatch(
+            a1,
+            topk_weights.to(torch.float32),
+            None,
+            self._route_padding_locally(topk_ids.to(torch.int32)),
+            return_routing=True,
+        )
+        rows = self._recv_rows_bound(a1, dispatch_a1.shape[0])
+        if rows < dispatch_a1.shape[0]:
+            dispatch_a1 = dispatch_a1[:rows]
+            dispatch_weights = dispatch_weights[:rows]
+            dispatch_ids = dispatch_ids[:rows]
+            if dispatch_scale is not None:
+                dispatch_scale = dispatch_scale[:rows]
+        # MORI reuses capacity-sized buffers without clearing unused rows.
+        unused = (
+            torch.arange(dispatch_a1.shape[0], device=dispatch_a1.device)
+            >= dispatch_recv_token_num
+        ).unsqueeze(1)
+        # Experts that read only the received rows can skip clearing the rest,
+        # unless a dynamic per-tensor activation scale is taken over every row.
+        if not self._experts_read_recv_rows_only or (
+            quant_config.is_per_tensor and quant_config.a1_scale is None
+        ):
+            dispatch_a1.masked_fill_(unused, 0)
+        dispatch_weights.masked_fill_(unused, 0)
+        dispatch_ids.masked_fill_(unused, -1)
+        expert_tokens_meta = mk.ExpertTokensMetadata(
+            expert_num_tokens=dispatch_recv_token_num,
+            expert_num_tokens_cpu=None,
+        )
+        return (
+            dispatch_a1,
+            dispatch_scale,
+            expert_tokens_meta,
+            dispatch_ids,
+            dispatch_weights,
+        )
+
+    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:
+        del topk_weights, topk_ids, apply_router_weight_on_input
+        del weight_and_reduce_impl
+        if self._routing is None:
+            raise RuntimeError("MORI EPv2 combine called without a dispatch")
+        # combine takes its token mapping from the routing, not from indices.
+        result = self.mori_op.combine(fused_expert_output, routing=self._routing)[0]
+        output.copy_(result[: output.shape[0]])
+        if not torch.cuda.is_current_stream_capturing():
+            self._routing = None
diff --git a/vllm/model_executor/layers/fused_moe/prepare_megamoe.py b/vllm/model_executor/layers/fused_moe/prepare_megamoe.py
index 5e8d8cde9a62..3d8627824cad 100644
--- a/vllm/model_executor/layers/fused_moe/prepare_megamoe.py
+++ b/vllm/model_executor/layers/fused_moe/prepare_megamoe.py
@@ -4,17 +4,21 @@
 
 Quantize hidden states to fp8 with E8M0 group scales, to NVFP4 with per-token
 global scales, or copy them as BF16, and repack the routing top-k tensors into
-the int64/float32 layout that the DeepGEMM MegaMoE kernels consume.
+the int64/float32 layout that the DeepGEMM MegaMoE kernels consume. Given an
+EPLB layer state, routing ids are mapped to physical replicas on the way.
 """
 
 import torch
 
+from vllm.distributed.eplb.eplb_state import EplbLayerState
+from vllm.model_executor.layers.fused_moe.eplb_triton import eplb_map_and_record
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
     FP8_DTYPE,
     QuantKey,
     kMxfp8Dynamic,
 )
 from vllm.triton_utils import tl, triton
+from vllm.v1.worker.ubatching import dbo_current_ubatch_id
 
 _E8M0_SCALE_DTYPES = (torch.uint8, torch.float8_e8m0fnu)
 # Inverse of the NVFP4 range: E2M1 max (6) times E4M3 max (448).
@@ -81,6 +85,14 @@ def _store_topk(
     is_padding,
     topk_idx_out,
     topk_weights_out,
+    logical_to_physical,
+    replica_count,
+    expert_load,
+    record_enabled,
+    num_unpadded_tokens,
+    num_logical,
+    map_stride,
+    num_physical,
     topk_ids_stride_m: tl.constexpr,
     topk_ids_stride_k: tl.constexpr,
     topk_weights_stride_m: tl.constexpr,
@@ -113,6 +125,20 @@ def _store_topk(
         other=0,
     ).to(tl.int64)
     ids = tl.where(token_is_padding[:, None], -1, ids)
+    if logical_to_physical is not None:
+        ids = eplb_map_and_record(
+            ids,
+            token_id[:, None],
+            topk_mask,
+            logical_to_physical,
+            replica_count,
+            num_logical,
+            map_stride,
+            expert_load,
+            num_physical,
+            record_enabled,
+            num_unpadded_tokens,
+        )
     tl.store(
         topk_idx_out
         + token_id[:, None] * topk_idx_stride_m
@@ -149,6 +175,14 @@ def _prepare_megamoe_inputs_kernel(
     is_padding,
     topk_idx_out,
     topk_weights_out,
+    logical_to_physical,
+    replica_count,
+    expert_load,
+    record_enabled,
+    num_unpadded_tokens,
+    num_logical,
+    map_stride,
+    num_physical,
     hidden_stride_m: tl.constexpr,
     hidden_stride_k: tl.constexpr,
     x_stride_m: tl.constexpr,
@@ -249,6 +283,14 @@ def _prepare_megamoe_inputs_kernel(
             is_padding,
             topk_idx_out,
             topk_weights_out,
+            logical_to_physical,
+            replica_count,
+            expert_load,
+            record_enabled,
+            num_unpadded_tokens,
+            num_logical,
+            map_stride,
+            num_physical,
             topk_ids_stride_m,
             topk_ids_stride_k,
             topk_weights_stride_m,
@@ -277,6 +319,14 @@ def _prepare_megamoe_nvfp4_inputs_kernel(
     is_padding,
     topk_idx_out,
     topk_weights_out,
+    logical_to_physical,
+    replica_count,
+    expert_load,
+    record_enabled,
+    num_unpadded_tokens,
+    num_logical,
+    map_stride,
+    num_physical,
     w2_global_scale,
     l2_alpha,
     hidden_stride_m: tl.constexpr,
@@ -300,52 +350,83 @@ def _prepare_megamoe_nvfp4_inputs_kernel(
     hidden_size: tl.constexpr,
     top_k: tl.constexpr,
     num_local_experts: tl.constexpr,
-    BLOCK_M: tl.constexpr,
+    BLOCK_H: tl.constexpr,
     BLOCK_K: tl.constexpr,
+    NUM_SPLITS: tl.constexpr,
     BLOCK_TOPK: tl.constexpr,
     BLOCK_E: tl.constexpr,
     SHARED_BLOCK_M: tl.constexpr,
     SHARED_MXFP8: tl.constexpr,
 ) -> None:
+    # Grid (token, split): each program quantizes the BLOCK_K chunks split,
+    # split + NUM_SPLITS, ... of one token's row.
+    token = tl.program_id(0)
+    split = tl.program_id(1)
     if l2_alpha is not None:
         experts = tl.arange(0, BLOCK_E)
-        expert_mask = (experts < num_local_experts) & (tl.program_id(0) == 0)
+        expert_mask = (experts < num_local_experts) & (token == 0) & (split == 0)
         alpha = tl.load(w2_global_scale + experts, mask=expert_mask)
         tl.store(l2_alpha + experts, alpha * l2_activation_scale, mask=expert_mask)
+    if token >= num_tokens:
+        return
 
-    token_id = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M)
-    token_mask = token_id < num_tokens
-    k_offsets = tl.arange(0, BLOCK_K)
-    hidden_ptrs = hidden_states + token_id[:, None] * hidden_stride_m
-
-    amax = tl.zeros((BLOCK_M,), tl.float32)
-    for k_block_id in range(hidden_size // BLOCK_K):
-        hidden = tl.load(
-            hidden_ptrs + (k_block_id * BLOCK_K + k_offsets[None, :]) * hidden_stride_k,
-            mask=token_mask[:, None],
-            other=0.0,
-        ).to(tl.float32)
-        amax = tl.maximum(amax, tl.max(tl.abs(hidden), axis=1))
-    x_scale = tl.maximum(amax, 1.0e-12) * _INV_NVFP4_MAX
-    tl.store(x_scales + token_id, x_scale, mask=token_mask)
-
-    for k_block_id in range(hidden_size // BLOCK_K):
-        k_ids = k_block_id * BLOCK_K + k_offsets
-        hidden = tl.load(
-            hidden_ptrs + k_ids[None, :] * hidden_stride_k,
-            mask=token_mask[:, None],
-            other=0.0,
+    if split == 0:
+        token_ids = token + tl.zeros([1], dtype=tl.int32)
+        _store_topk(
+            token_ids,
+            token_ids < num_tokens,
+            topk_ids,
+            topk_weights,
+            is_padding,
+            topk_idx_out,
+            topk_weights_out,
+            logical_to_physical,
+            replica_count,
+            expert_load,
+            record_enabled,
+            num_unpadded_tokens,
+            num_logical,
+            map_stride,
+            num_physical,
+            topk_ids_stride_m,
+            topk_ids_stride_k,
+            topk_weights_stride_m,
+            topk_weights_stride_k,
+            is_padding_stride_m,
+            topk_idx_stride_m,
+            topk_idx_stride_k,
+            topk_weights_out_stride_m,
+            topk_weights_out_stride_k,
+            top_k,
+            1,
+            BLOCK_TOPK,
         )
 
+    # Every program reads the whole row for the token scale.
+    row = hidden_states + token.to(tl.int64) * hidden_stride_m
+    cols = tl.arange(0, BLOCK_H)
+    hidden_row = tl.load(
+        row + cols * hidden_stride_k, mask=cols < hidden_size, other=0.0
+    )
+    amax = tl.max(tl.abs(hidden_row.to(tl.float32)), axis=0)
+    x_scale = tl.maximum(amax, 1.0e-12) * _INV_NVFP4_MAX
+    if split == 0:
+        tl.store(x_scales + token, x_scale)
+
+    shared_sf_row = _shared_sf_row(token, SHARED_BLOCK_M)
+    for k0 in tl.range(split * BLOCK_K, hidden_size, NUM_SPLITS * BLOCK_K):
+        k_ids = k0 + tl.arange(0, BLOCK_K)
+        k_mask = k_ids < hidden_size
+        hidden = tl.load(row + k_ids * hidden_stride_k, mask=k_mask, other=0.0)
+
         # E2M1 values with one E4M3 scale per 16, relative to the token scale.
         values = tl.reshape(
-            tl.div_rn(hidden.to(tl.float32), x_scale[:, None]),
-            [BLOCK_M, BLOCK_K // 16, 16],
+            tl.div_rn(hidden.to(tl.float32), x_scale), [BLOCK_K // 16, 16]
         )
-        block_scale = tl.max(tl.abs(values), axis=2) * (1.0 / 6.0)
+        block_scale = tl.max(tl.abs(values), axis=1) * (1.0 / 6.0)
         block_scale = tl.minimum(tl.maximum(block_scale, 2.0**-9), 448.0)
         block_scale = block_scale.to(tl.float8e4nv)
-        values = tl.div_rn(values, block_scale.to(tl.float32)[:, :, None])
+        values = tl.div_rn(values, block_scale.to(tl.float32)[:, None])
         magnitude = tl.abs(values)
         # Round to nearest even on the grid {0, 0.5, 1, 1.5, 2, 3, 4, 6}.
         code = (
@@ -358,61 +439,45 @@ def _prepare_megamoe_nvfp4_inputs_kernel(
             + (magnitude > 5.0).to(tl.int32)
         )
         code = tl.where((values < 0) & (code != 0), code | 8, code)
-        low, high = tl.split(tl.reshape(code, [BLOCK_M, BLOCK_K // 2, 2]))
-        packed_offsets = k_block_id * (BLOCK_K // 2) + tl.arange(0, BLOCK_K // 2)
+        low, high = tl.split(tl.reshape(code, [BLOCK_K // 2, 2]))
+        packed_ids = k0 // 2 + tl.arange(0, BLOCK_K // 2)
         tl.store(
-            x_out + token_id[:, None] * x_stride_m + packed_offsets[None, :],
+            x_out + token * x_stride_m + packed_ids,
             (low | (high << 4)).to(tl.uint8),
-            mask=token_mask[:, None],
+            mask=packed_ids < hidden_size // 2,
         )
-        sf_offsets = k_block_id * (BLOCK_K // 64) + tl.arange(0, BLOCK_K // 64)
+        # Four E4M3 scales per little-endian int32 word, one word per 64 values.
+        sf_ids = k0 // 64 + tl.arange(0, BLOCK_K // 64)
         sf_bytes = tl.reshape(
-            block_scale.to(tl.uint8, bitcast=True), [BLOCK_M, BLOCK_K // 64, 4]
+            block_scale.to(tl.uint8, bitcast=True), [BLOCK_K // 64, 4]
         )
         tl.store(
-            x_sf + token_id[:, None] * x_sf_stride_m + sf_offsets[None, :],
+            x_sf + token * x_sf_stride_m + sf_ids,
             _pack_bytes(sf_bytes, 4),
-            mask=token_mask[:, None],
+            mask=sf_ids < hidden_size // 64,
         )
 
         if shared_x is not None:
-            shared_ptrs = shared_x + token_id[:, None] * shared_x_stride_m
+            shared_ptrs = shared_x + token * shared_x_stride_m + k_ids
             if SHARED_MXFP8:
                 fp8, _, scale_exp = _quantize_fp8_groups(
-                    hidden.to(tl.float32), BLOCK_M, BLOCK_K, 32, True
+                    tl.reshape(hidden.to(tl.float32), (1, BLOCK_K)),
+                    1,
+                    BLOCK_K,
+                    32,
+                    True,
                 )
-                tl.store(shared_ptrs + k_ids[None, :], fp8, mask=token_mask[:, None])
+                tl.store(shared_ptrs, tl.reshape(fp8, [BLOCK_K]), mask=k_mask)
+                word_ids = k0 // 128 + tl.arange(0, BLOCK_K // 128)
                 tl.store(
                     shared_x_sf
-                    + _shared_sf_row(token_id, SHARED_BLOCK_M) * shared_x_sf_stride_m
-                    + k_block_id * shared_x_sf_stride_k,
-                    _pack_bytes(scale_exp, BLOCK_K // 32),
-                    mask=token_mask,
+                    + shared_sf_row * shared_x_sf_stride_m
+                    + word_ids * shared_x_sf_stride_k,
+                    _pack_bytes(tl.reshape(scale_exp, [BLOCK_K // 128, 4]), 4),
+                    mask=word_ids < hidden_size // 128,
                 )
             else:
-                tl.store(shared_ptrs + k_ids[None, :], hidden, mask=token_mask[:, None])
-
-    _store_topk(
-        token_id,
-        token_mask,
-        topk_ids,
-        topk_weights,
-        is_padding,
-        topk_idx_out,
-        topk_weights_out,
-        topk_ids_stride_m,
-        topk_ids_stride_k,
-        topk_weights_stride_m,
-        topk_weights_stride_k,
-        is_padding_stride_m,
-        topk_idx_stride_m,
-        topk_idx_stride_k,
-        topk_weights_out_stride_m,
-        topk_weights_out_stride_k,
-        top_k,
-        BLOCK_M,
-        BLOCK_TOPK,
-    )
+                tl.store(shared_ptrs, hidden, mask=k_mask)
 
 
 def _resolve_hidden_quant(hidden_quant: QuantKey, block_k: int) -> tuple[int, bool]:
@@ -436,6 +501,29 @@ def _resolve_hidden_quant(hidden_quant: QuantKey, block_k: int) -> tuple[int, bo
     return group_shape.col, scale.dtype in _E8M0_SCALE_DTYPES
 
 
+def _eplb_kernel_args(eplb_state: EplbLayerState | None) -> tuple:
+    """Replica-map and load-recording kernel args; inert without EPLB."""
+    if eplb_state is None or eplb_state.logical_to_physical_map is None:
+        return (None,) * 5 + (0,) * 3
+    mapping = eplb_state.logical_to_physical_map
+    replicas = eplb_state.logical_replica_count
+    load = eplb_state.expert_load_view
+    record = eplb_state.should_record_tensor
+    unpadded = eplb_state.num_unpadded_tokens_tensors
+    assert replicas is not None and load is not None and record is not None
+    assert unpadded is not None
+    return (
+        mapping,
+        replicas,
+        load,
+        record,
+        unpadded[dbo_current_ubatch_id()],
+        mapping.shape[0],
+        mapping.stride(0),
+        load.numel(),
+    )
+
+
 def prepare_megamoe_inputs(
     hidden_states: torch.Tensor,
     topk_weights: torch.Tensor,
@@ -448,6 +536,7 @@ def prepare_megamoe_inputs(
     shared_x_sf: torch.Tensor | None = None,
     shared_block_m: int | None = None,
     hidden_quant: QuantKey | None = kMxfp8Dynamic,
+    eplb_state: EplbLayerState | None = None,
 ) -> None:
     """Quantize hidden states and repack top-k routing for DeepGEMM MegaMoE.
 
@@ -467,6 +556,7 @@ def prepare_megamoe_inputs(
         hidden_quant: Hidden-state quantization scheme. E8M0 scales are packed
             four per int32 into ``x_sf``; fp32 scales are stored directly.
             ``None`` copies the hidden states into ``x`` unquantized.
+        eplb_state: Optional EPLB layer state; maps ids to replicas and records load.
 
     """
     num_tokens, hidden_size = hidden_states.shape
@@ -518,8 +608,9 @@ def prepare_megamoe_inputs(
                 f"{required_rows}, got {shared_x_sf.shape[0]}."
             )
 
-    # On GB200, eight-row tiles win from 64 tokens; keep smaller batches untiled.
-    block_m = 8 if num_tokens >= 64 else 1
+    # On GB200, 8-row tiles win from 64 tokens and 32-row tiles from 1024 (1.3-1.6x
+    # over 8 rows); smaller batches stay untiled.
+    block_m = 32 if num_tokens >= 1024 else 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
@@ -533,6 +624,7 @@ def prepare_megamoe_inputs(
         is_padding,
         topk_idx_out,
         topk_weights_out,
+        *_eplb_kernel_args(eplb_state),
         hidden_states.stride(0),
         hidden_states.stride(1),
         x.stride(0),
@@ -580,6 +672,7 @@ def prepare_megamoe_nvfp4_inputs(
     w2_global_scale: torch.Tensor | None = None,
     l2_alpha: torch.Tensor | None = None,
     l2_activation_scale: float = 1.0,
+    eplb_state: EplbLayerState | None = None,
 ) -> None:
     """Stage DeepGEMM NVFP4 MegaMoE inputs (``kNvfp4DynamicToken``).
 
@@ -606,6 +699,7 @@ def prepare_megamoe_nvfp4_inputs(
             dequantization scale of each local expert. Written even without
             local tokens, as other ranks' tokens still reach these experts.
         l2_activation_scale: Global scale the kernel quantizes L2 inputs with.
+        eplb_state: Optional EPLB layer state; maps ids to replicas and records load.
 
     """
     num_tokens, hidden_size = hidden_states.shape
@@ -632,11 +726,19 @@ def prepare_megamoe_nvfp4_inputs(
         if tensor is not None and tensor.stride(-1) != 1:
             raise ValueError(f"MegaMoE NVFP4 staging requires a row-major {name}.")
 
-    block_m = 8 if num_tokens >= 64 else 1
     top_k = topk_ids.shape[1]
     num_local_experts = 0 if l2_alpha is None else l2_alpha.numel()
+    # Up to 512 tokens each row is split over programs; larger batches fill the
+    # GPU with one program per row looping over wider chunks.
+    block_h = triton.next_power_of_2(hidden_size)
+    if num_tokens <= 512:
+        block_k = min(1024, block_h)
+        num_splits = triton.cdiv(hidden_size, block_k)
+    else:
+        block_k = min(2048, block_h)
+        num_splits = 1
     # One program even without tokens, to write l2_alpha.
-    grid = (max(triton.cdiv(num_tokens, block_m), 1),)
+    grid = (max(num_tokens, 1), num_splits)
     _prepare_megamoe_nvfp4_inputs_kernel[grid](
         hidden_states,
         x.view(torch.uint8),
@@ -649,6 +751,7 @@ def prepare_megamoe_nvfp4_inputs(
         is_padding,
         topk_idx_out,
         topk_weights_out,
+        *_eplb_kernel_args(eplb_state),
         w2_global_scale,
         l2_alpha,
         hidden_states.stride(0),
@@ -672,8 +775,9 @@ def prepare_megamoe_nvfp4_inputs(
         hidden_size,
         top_k,
         num_local_experts,
-        BLOCK_M=block_m,
-        BLOCK_K=128,
+        BLOCK_H=block_h,
+        BLOCK_K=block_k,
+        NUM_SPLITS=num_splits,
         BLOCK_TOPK=triton.next_power_of_2(top_k),
         BLOCK_E=triton.next_power_of_2(max(num_local_experts, 1)),
         SHARED_BLOCK_M=shared_block_m or 1,
diff --git a/vllm/model_executor/layers/fused_moe/routed_experts.py b/vllm/model_executor/layers/fused_moe/routed_experts.py
index 59aea0599b4e..0273f577042e 100644
--- a/vllm/model_executor/layers/fused_moe/routed_experts.py
+++ b/vllm/model_executor/layers/fused_moe/routed_experts.py
@@ -26,11 +26,11 @@
 from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import (
     UnquantizedFusedMoEMethod,
 )
-from vllm.model_executor.layers.linear import LinearBase
 from vllm.model_executor.layers.quantization.base_config import (
     QuantizationConfig,
     resolve_quant_method,
 )
+from vllm.model_executor.utils import kernel_parameters
 from vllm.utils.math_utils import cdiv
 
 if TYPE_CHECKING:
@@ -160,7 +160,6 @@ def __init__(
         object.__setattr__(self, "hash_indices_table", hash_indices_table)
         # End random parameters
         self._loaded_expert_biases: set[str] = set()
-        self._fused_shared_expert_quantizer = self._get_fused_shared_expert_quantizer()
 
         self.quant_method = self._get_quant_method(
             self.layer_name,
@@ -168,6 +167,26 @@ def __init__(
             self.moe_config,
         )
 
+        # Only set when FSE go through online quantization, with pre-quantized routed.
+        self.shared_expert_online_loader = None
+        online_config = getattr(self.quant_config, "online_quantization_config", None)
+        if (
+            online_config is not None
+            and self.expert_map_manager.num_fused_shared_experts > 0
+        ):
+            from vllm.model_executor.layers.quantization.online.shared_expert import (
+                OnlineSharedExpertLoader,
+            )
+
+            self.shared_expert_online_loader = OnlineSharedExpertLoader(
+                self.quant_config,
+                online_config,
+                self.quant_method,
+                self.moe_config.shared_expert_prefix,
+                self.moe_config.num_logical_experts,
+                self.expert_map_manager.num_fused_shared_experts,
+            )
+
         # Round up hidden size and update moe_config.
         # TODO: move roundup to _get_quant_method?
         self.hidden_size, self.intermediate_size_per_partition = (
@@ -212,6 +231,10 @@ def __init__(
 
         self.quant_method.create_weights(layer=self, **moe_quant_params)
 
+        loader = self.shared_expert_online_loader
+        if loader is not None and loader.stages_shared_experts:
+            loader.create_staging(self, params_dtype, weight_loader=self.weight_loader)
+
         self.lora_base_layer_prefix = ""
 
     # TODO(bnell): Temporary hack. Get rid of this.
@@ -638,31 +661,6 @@ def _load_single_value(
         # _to_scalar's reshape(()) would reject the size-2 weight_shape.
         param_data[expert_id] = loaded_weight
 
-    def _get_fused_shared_expert_quantizer(
-        self,
-    ) -> Callable[[torch.Tensor], tuple[torch.Tensor, torch.Tensor]] | None:
-        """Return the quantizer for an online-quantized fused shared expert."""
-        online_config = getattr(self.quant_config, "online_quantization_config", None)
-        prefix = self.moe_config.shared_expert_prefix
-        if (
-            online_config is None
-            or prefix is None
-            or self.expert_map_manager.num_fused_shared_experts == 0
-        ):
-            return None
-        # FSE compatibility checks require every shared-expert projection to
-        # use the routed experts' weight key.
-        resolved = online_config.resolve_quant_method_cls(
-            LinearBase, f"{prefix}.down_proj"
-        )
-        if resolved is None:
-            return None
-        from vllm.model_executor.layers.quantization.online.base import (
-            ONLINE_SHARED_EXPERT_QUANTIZERS,
-        )
-
-        return ONLINE_SHARED_EXPERT_QUANTIZERS[resolved[3].weight]
-
     @overload
     def weight_loader(
         self,
@@ -725,6 +723,24 @@ def weight_loader(
 
         quant_method_name = base_quant_method.__class__.__name__
         global_expert_id = expert_id
+        num_global_experts = (
+            self.moe_config.num_experts
+            + self.expert_map_manager.num_fused_shared_experts
+        )
+        # `num_experts` counts logical and redundant experts; fused shared
+        # slots are appended after them.
+        # Expert substitution checkpoints name retained experts by their
+        # logical ID, which is not bounded by the number of allocated slots.
+        expert_substitution = getattr(self, "expert_substitution", None)
+        if expert_substitution is not None:
+            num_global_experts = (
+                expert_substitution.num_logical_experts
+                + self.expert_map_manager.num_fused_shared_experts
+            )
+        assert global_expert_id < num_global_experts, (
+            f"Expert id {global_expert_id} is out of range for "
+            f"{self.layer_name}; expected < {num_global_experts}"
+        )
         expert_id = self._map_global_expert_id_to_local_expert_id(global_expert_id)
 
         use_global_sf = (
@@ -737,34 +753,18 @@ def weight_loader(
             return False if return_success else None
         # Hereafter, `expert_id` is local physical id
 
-        # A full-precision fused shared expert is quantized into the checkpoint
-        # layout, then loaded like a pre-quantized expert weight and scale.
-        if (
-            self._fused_shared_expert_quantizer is not None
-            and self.moe_config.num_logical_experts
-            <= global_expert_id
-            < self.moe_config.num_logical_experts
-            + self.expert_map_manager.num_fused_shared_experts
-            and weight_name.endswith("_weight")
-            and loaded_weight.is_floating_point()
-            and loaded_weight.dtype != param.dtype
-        ):
-            weight, weight_scale = self._fused_shared_expert_quantizer(
-                loaded_weight.to(self.moe_config.device)
-            )
-            stem = "w2" if shard_id == "w2" else "w13"
-            self.weight_loader(
-                getattr(self, f"{stem}_weight_scale"),
-                weight_scale,
-                f"{weight_name}_scale",
+        if self.shared_expert_online_loader is not None:
+            param, expert_id, loaded = self.shared_expert_online_loader.load(
+                self,
+                param,
+                loaded_weight,
+                weight_name,
                 shard_id,
                 global_expert_id,
-                return_success=True,
-            )
-            loaded = self.weight_loader(
-                param, weight, weight_name, shard_id, global_expert_id, True
+                expert_id,
             )
-            return loaded if return_success else None
+            if loaded is not None:
+                return loaded if return_success else None
 
         # is_transposed: if the dim to shard the weight
         # should be flipped. Required by GPTQ/AWQ (K-first format).
@@ -1042,6 +1042,10 @@ def load_weights(
                 )
                 weight_name = qual_name.replace(weight_name, param_name)
                 param_name = weight_name.removeprefix(f"{self.layer_name}.")
+                if self.shared_expert_online_loader is not None:
+                    param_name = self.shared_expert_online_loader.redirect_param_name(
+                        param_name, shard_id, expert_id
+                    )
                 param = getattr(self, param_name, None)
                 if param is None:
                     if param_name.endswith(("w13_bias", "w2_bias")):
@@ -1147,7 +1151,7 @@ def make_expert_params_mapping(
         has_base_layer = any(".base_layer." in n for n, _ in model.named_parameters())
         prefix = "base_layer." if has_base_layer else ""
         # These loaders index ``params_dict[full_name]``, so both sides get it.
-        return RoutedExperts.build_expert_params_mapping(
+        mapping = RoutedExperts.build_expert_params_mapping(
             ckpt_gate_proj_name,
             ckpt_down_proj_name,
             ckpt_up_proj_name,
@@ -1158,6 +1162,14 @@ def make_expert_params_mapping(
             lora_base_layer_prefix_on_param_name=prefix,
         )
 
+        # All MoE layers of a model share the same FSE and quantization setup.
+        routed_experts = next(
+            (m for m in model.modules() if isinstance(m, RoutedExperts)), None
+        )
+        if routed_experts is None or routed_experts.shared_expert_online_loader is None:
+            return mapping
+        return routed_experts.shared_expert_online_loader.redirect_mapping(mapping)
+
     @staticmethod
     def build_expert_params_mapping(
         ckpt_gate_proj_name: str,
@@ -1329,7 +1341,14 @@ def _maybe_make_contiguous(
                 torch.transpose(p.data, 1, 2), requires_grad=False
             )
 
-        weights = list(self.named_parameters())
+        # Kernel views of the params: live ones only, under their runtime names
+        # (an aliased checkpoint name such as `w13_weight_packed` is reported as
+        # `w13_weight`, which the name filters below expect).
+        weights = [
+            (f"{prefix}.{name}" if prefix else name, p)
+            for prefix, module in self.named_modules()
+            for name, p in kernel_parameters(module).items()
+        ]
 
         # EPLB must use shards when original weight storage has been released.
         # Use explicit metadata: some tensor subclasses lack untyped_storage().
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 9f01324fe92b..90f292f6e3ae 100644
--- a/vllm/model_executor/layers/fused_moe/router/base_router.py
+++ b/vllm/model_executor/layers/fused_moe/router/base_router.py
@@ -8,6 +8,7 @@
 import vllm.envs as envs
 from vllm.distributed.eplb.eplb_state import EplbLayerState
 from vllm.forward_context import get_forward_context, is_forward_context_available
+from vllm.model_executor.layers.fused_moe.eplb_triton import eplb_map_and_record
 from vllm.model_executor.layers.fused_moe.router.fused_moe_router import (
     FusedMoERouter,
 )
@@ -31,7 +32,6 @@ def _eplb_map_and_record_i32_kernel(
         out_size,
         numel,
         num_active_experts,
-        HAS_NUM_UNPADDED: tl.constexpr,
         BLOCK_SIZE: tl.constexpr,
     ):
         pid = tl.program_id(0)
@@ -39,31 +39,6 @@ def _eplb_map_and_record_i32_kernel(
         mask = offs < numel
 
         expert_id = tl.load(topk_ids_ptr + offs, mask=mask, other=0).to(tl.int64)
-        valid_expert = (expert_id >= 0) & (expert_id < num_logical_experts)
-        safe_expert_id = tl.where(valid_expert, expert_id, 0)
-
-        # 1. Convert the logical expert ids to physical expert ids
-        replica_count = tl.load(
-            logical_replica_count_ptr + safe_expert_id,
-            mask=mask & valid_expert,
-            other=1,
-        )
-        # Avoid invalid modulo/div by forcing at least 1.
-        replica_count = tl.maximum(replica_count, 1)
-        # floor(2^32 / phi), classic Knuth multiplicative hash multiplier.
-        KNUTH_MULTIPLIER = 2654435769
-        token_idx = (offs // num_active_experts).to(tl.int64)
-        hashed = (token_idx * KNUTH_MULTIPLIER) & 0xFFFFFFFF
-        replica_idx = hashed % replica_count
-        map_index = safe_expert_id * map_slots + replica_idx
-        physical_id = tl.load(
-            logical_to_physical_ptr + map_index,
-            mask=mask & valid_expert,
-            other=-1,
-        )
-        tl.store(out_ids_ptr + offs, physical_id, mask=mask)
-
-        # 2. Record expert load metrics.
 
         # TODO(bowen): When using `FusedMoEModularKernel`, this
         # can be done in a more unified way, since
@@ -76,23 +51,20 @@ def _eplb_map_and_record_i32_kernel(
         # If later refactor moved all the MoE kernel calls
         # to the modular kernel, we can move this logic there
         # to achieve better efficiency.
-
-        record_enabled = tl.load(record_enabled_ptr) != 0
-        # Skip padded tokens when recording.
-        if HAS_NUM_UNPADDED:
-            num_unpadded_tokens = tl.load(num_unpadded_tokens_ptr)
-            is_unpadded = offs < num_unpadded_tokens * num_active_experts
-        else:
-            is_unpadded = True
-        valid = (
-            mask
-            & record_enabled
-            & is_unpadded
-            & (physical_id >= 0)
-            & (physical_id < out_size)
+        physical_id = eplb_map_and_record(
+            expert_id,
+            offs // num_active_experts,
+            mask,
+            logical_to_physical_ptr,
+            logical_replica_count_ptr,
+            num_logical_experts,
+            map_slots,
+            out_ptr,
+            out_size,
+            record_enabled_ptr,
+            num_unpadded_tokens_ptr,
         )
-        safe_physical_id = tl.where(physical_id >= 0, physical_id, 0)
-        tl.atomic_add(out_ptr + safe_physical_id, 1, mask=valid)
+        tl.store(out_ids_ptr + offs, physical_id, mask=mask)
 
     def _eplb_map_and_record_triton(
         topk_ids: torch.Tensor,
@@ -123,7 +95,6 @@ def _eplb_map_and_record_triton(
             expert_load_view.shape[0],
             numel,
             num_active_experts,
-            HAS_NUM_UNPADDED=num_unpadded_tokens is not None,
             BLOCK_SIZE=256,
         )
         return out_flat.reshape(topk_ids.shape)
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 64a12e57cda5..95b369882b22 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,7 +8,6 @@
 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,
@@ -17,7 +16,10 @@
 from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import (
     rocm_aiter_grouped_topk,
 )
-from vllm.model_executor.layers.fused_moe.router.base_router import BaseRouter
+from vllm.model_executor.layers.fused_moe.router.base_router import (
+    BaseRouter,
+    get_moe_padding_mask,
+)
 from vllm.model_executor.layers.fused_moe.router.fused_topk_bias_router import (
     fused_topk_bias,
 )
@@ -36,9 +38,11 @@ def fused_grouped_topk(
     topk_group: int = 0,
     scoring_func: str = "softmax",
     routed_scaling_factor: float = 1.0,
+    is_padding: torch.Tensor | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     assert hidden_states.size(0) == gating_output.size(0), "Number of tokens mismatch"
 
+    kernel_padding = is_padding if current_platform.is_cuda() else None
     if scoring_func == "sigmoid":
         # Fully fused kernel path for sigmoid
         topk_values, topk_indices = ops.grouped_topk(
@@ -50,6 +54,7 @@ def fused_grouped_topk(
             routed_scaling_factor,
             e_score_correction_bias,
             1,  # scoring_func=1 for sigmoid
+            kernel_padding,
         )
     elif scoring_func == "softmax":
         # Apply softmax in Python, then use fused kernel
@@ -64,10 +69,16 @@ def fused_grouped_topk(
             routed_scaling_factor,
             e_score_correction_bias,
             0,  # scoring_func=0 (no activation, scores already computed)
+            kernel_padding,
         )
     else:
         raise ValueError(f"Unsupported scoring function: {scoring_func}")
 
+    # XPU's external operator retains its original ABI and post-routing mask.
+    if is_padding is not None and not current_platform.is_cuda():
+        topk_values = topk_values.masked_fill(is_padding.unsqueeze(1), 0)
+        topk_indices = topk_indices.masked_fill(is_padding.unsqueeze(1), -1)
+
     # Fused kernel outputs float32 values and int32 indices directly
     return topk_values, topk_indices
 
@@ -88,6 +99,7 @@ def grouped_topk(
     scoring_func: str = "softmax",
     routed_scaling_factor: float = 1.0,
     e_score_correction_bias: torch.Tensor | None = None,
+    is_padding: torch.Tensor | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     if (
         envs.VLLM_USE_FUSED_MOE_GROUPED_TOPK
@@ -106,6 +118,7 @@ def grouped_topk(
             topk_group=topk_group,
             scoring_func=scoring_func,
             routed_scaling_factor=routed_scaling_factor,
+            is_padding=is_padding,
         )
 
     assert hidden_states.size(0) == gating_output.size(0), "Number of tokens mismatch"
@@ -159,6 +172,9 @@ def grouped_topk(
 
     if routed_scaling_factor != 1.0:
         topk_weights = topk_weights * routed_scaling_factor
+    if is_padding is not None:
+        topk_weights = topk_weights.masked_fill(is_padding.unsqueeze(1), 0)
+        topk_ids = topk_ids.masked_fill(is_padding.unsqueeze(1), -1)
     return topk_weights.to(torch.float32), topk_ids.to(torch.int32)
 
 
@@ -325,6 +341,10 @@ def valid_grouping() -> bool:
                 )
             return topk_weights, topk_ids
 
+        is_padding = (
+            get_moe_padding_mask(router_logits.shape[0]) if self.skip_padding else None
+        )
+
         # Select grouped_topk implementation
         if rocm_aiter_ops.is_fused_moe_enabled():
             if not rocm_aiter_ops.is_fusion_moe_shared_experts_enabled():
@@ -334,7 +354,7 @@ def valid_grouping() -> bool:
                 num_fused_shared_experts=self.num_fused_shared_experts,
             )
         else:
-            grouped_topk_impl = grouped_topk
+            grouped_topk_impl = partial(grouped_topk, is_padding=is_padding)
 
         topk_weights, topk_ids = grouped_topk_impl(
             hidden_states=hidden_states,
@@ -348,15 +368,9 @@ 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)
+        # AITER retains its existing interface and masking behavior.
+        if rocm_aiter_ops.is_fused_moe_enabled() and is_padding is not None:
+            topk_weights = topk_weights.masked_fill(is_padding.unsqueeze(1), 0)
+            topk_ids = topk_ids.masked_fill(is_padding.unsqueeze(1), -1)
 
         return topk_weights, topk_ids
diff --git a/vllm/model_executor/layers/fused_moe/utils.py b/vllm/model_executor/layers/fused_moe/utils.py
index 27df12668584..6a4911e8387f 100644
--- a/vllm/model_executor/layers/fused_moe/utils.py
+++ b/vllm/model_executor/layers/fused_moe/utils.py
@@ -1,13 +1,12 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import functools
-from collections.abc import Iterable
 from math import prod
 from typing import TYPE_CHECKING
 
+import regex as re
 import torch
 import torch.nn.functional as F
-from torch import nn
 
 import vllm.envs as envs
 from vllm import _custom_ops as ops
@@ -40,7 +39,6 @@
 from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
     per_tensor_dequantize,
 )
-from vllm.model_executor.models.utils import PPMissingLayer
 from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
 from vllm.utils.math_utils import cdiv
@@ -52,11 +50,23 @@
 logger = init_logger(__name__)
 
 
+# Matches any layer-qualified module path (e.g.
+# ``model.layers.5.mlp.shared_expert`` or
+# ``model.language_model.model.layers.5.mlp.shared_expert``) so skip reasons
+# are reported once per cause instead of once per layer.
+_LAYER_MODULE_PATH_RE = re.compile(r"(?:[\w]+\.)*(?:layers|mtp)\.\d+\.[\w.]*")
+
+
+def _normalize_skip_reason(reason: str) -> str:
+    return _LAYER_MODULE_PATH_RE.sub("<shared_expert>", reason)
+
+
 def resolve_layer_fused_shared_expert(
     quant_config: "QuantizationConfig | None",
     prefix: str,
     shared_expert_name: str = "shared_experts",
     moe_backend: str | None = None,
+    disable_under_ep: bool = False,
 ) -> bool:
     """Resolve whether AITER fused shared-expert execution is enabled.
 
@@ -65,6 +75,9 @@ def resolve_layer_fused_shared_expert(
         prefix: MoE module prefix.
         shared_expert_name: Shared-expert module name under ``prefix``.
         moe_backend: MoE backend that runs the fused shared expert.
+        disable_under_ep: Whether this model family does not support FSE under
+            expert parallelism. Kept as an explicit call-site policy instead of
+            branching on model type here.
 
     Returns:
         Whether AITER fused shared experts are enabled.
@@ -76,16 +89,31 @@ def resolve_layer_fused_shared_expert(
     # NOTE: is_fusion_moe_shared_experts_enabled is decorated with @if_aiter_supported
     # that returns None if AITER is not available.
     fse_requested = bool(rocm_aiter_ops.is_fusion_moe_shared_experts_enabled())
-    fse_compatible, fse_reason = (
-        is_shared_expert_quant_fse_compatible(
+
+    # TODO: should it be disabled for all models?
+    skip_reason: str | None = None
+    if fse_requested and disable_under_ep:
+        from vllm.config import get_current_vllm_config_or_none
+
+        vllm_config = get_current_vllm_config_or_none()
+        parallel_config = (
+            vllm_config.parallel_config if vllm_config is not None else None
+        )
+        if parallel_config is not None and parallel_config.enable_expert_parallel:
+            fse_requested = False
+            skip_reason = "expert parallelism is enabled"
+
+    if fse_requested:
+        fse_compatible, fse_reason = is_shared_expert_quant_fse_compatible(
             quant_config,
             f"{prefix}.experts",
             f"{prefix}.{shared_expert_name}",
             moe_backend=moe_backend,
         )
-        if fse_requested
-        else (True, None)
-    )
+        if not fse_compatible:
+            skip_reason = fse_reason or "unknown reason"
+    else:
+        fse_compatible = True
     is_fused_shared_expert_enabled = fse_requested and fse_compatible
 
     # In online shared_expert quantization is used, register it to
@@ -121,52 +149,18 @@ def resolve_layer_fused_shared_expert(
                     quant_key_str,
                     target_pattern,
                 )
-    if fse_requested and not is_fused_shared_expert_enabled:
-        logger.warning(
+    if skip_reason is not None:
+        # Keep the full per-layer reason at debug level; the warning below is
+        # deduplicated by cause so a long model logs it once.
+        logger.debug("Skipping shared-expert fusion at %s: %s", prefix, skip_reason)
+        logger.warning_once(
             "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS is enabled but "
-            "cannot be enabled - skipping for this layer: %s.",
-            fse_reason,
+            "cannot be enabled - skipping shared-expert fusion: %s.",
+            _normalize_skip_reason(skip_reason),
         )
     return is_fused_shared_expert_enabled
 
 
-def is_model_fused_shared_expert_compatible(
-    layers: nn.ModuleList | Iterable[nn.Module],
-    moe_cls: type[nn.Module],
-    moe_name: str,
-) -> bool:
-    """Resolve one fused-shared-expert state for a model's MoE layers."""
-
-    def get_moe_layer(layer: nn.Module) -> nn.Module | None:
-        for name in moe_name.split("."):
-            layer = getattr(layer, name, None)
-            if layer is None:
-                return None
-        return layer
-
-    moe_layers = (
-        moe_layer
-        for layer in layers
-        if not isinstance(layer, PPMissingLayer)
-        and (moe_layer := get_moe_layer(layer)) is not None
-        and isinstance(moe_layer, moe_cls)
-    )
-
-    enabled = [
-        getattr(layer, "is_fused_shared_expert_enabled", False) for layer in moe_layers
-    ]
-    enabled_count = sum(enabled)
-    disabled_count = len(enabled) - enabled_count
-    if enabled_count > 0 and disabled_count > 0:
-        raise NotImplementedError(
-            "Fused shared experts must be enabled for all MoE layers; found "
-            f"{enabled_count} enabled and {disabled_count} disabled layers. "
-            "Per-layer fused shared experts is not yet supported. Please open "
-            "an issue."
-        )
-    return enabled_count > 0 and disabled_count == 0
-
-
 @triton.jit
 def _count_expert_num_tokens(
     topk_ids_ptr,
diff --git a/vllm/model_executor/layers/fused_qk_norm_rope.py b/vllm/model_executor/layers/fused_qk_norm_rope.py
index a62174b338c4..504a3a4c50f0 100644
--- a/vllm/model_executor/layers/fused_qk_norm_rope.py
+++ b/vllm/model_executor/layers/fused_qk_norm_rope.py
@@ -47,7 +47,8 @@ def _fused_qk_rmsnorm_rope_gate_kernel(
     MROPE_SECTION_H: tl.constexpr,
     MROPE_SECTION_W: tl.constexpr,
 ):
-    token = tl.program_id(0)
+    # int64 so that token * stride does not overflow int32 for large batches.
+    token = tl.program_id(0).to(tl.int64)
     head = tl.program_id(1)
     is_k = head >= num_q_heads
     local_head = tl.where(is_k, head - num_q_heads, head)
diff --git a/vllm/model_executor/layers/indexer_topk.py b/vllm/model_executor/layers/indexer_topk.py
index 02c4280ea0a1..a372057b388d 100644
--- a/vllm/model_executor/layers/indexer_topk.py
+++ b/vllm/model_executor/layers/indexer_topk.py
@@ -233,6 +233,27 @@ def resolve_backend(
             )
         return self._backend
 
+    def takes_bf16_logits(
+        self, logits: torch.Tensor, topk_tokens: int, num_rows: int
+    ) -> bool:
+        """Whether this decode call can consume BF16 without a full FP32 copy."""
+        if logits.dtype != torch.bfloat16:
+            return False
+        if self._backend == "deep_select":
+            return (
+                self._has_deep_select
+                and _HAS_DEEP_SELECT_C
+                and is_deep_select_supported(logits, topk_tokens)
+            )
+        return (
+            self._backend in ("auto", "persistent")
+            and self._is_cuda
+            and current_platform.is_device_capability_family(100)
+            and logits.stride(1) == 1
+            and topk_tokens in (512, 1024, 2048)
+            and self.resolve_backend(logits, topk_tokens, num_rows) == "persistent"
+        )
+
     def _resolve_auto(
         self, logits: torch.Tensor, topk_tokens: int, num_rows: int
     ) -> str:
diff --git a/vllm/model_executor/layers/linear.py b/vllm/model_executor/layers/linear.py
index d3195472d3cd..bd106c0a8c65 100644
--- a/vllm/model_executor/layers/linear.py
+++ b/vllm/model_executor/layers/linear.py
@@ -30,8 +30,10 @@
     QuantizeMethodBase,
 )
 from vllm.model_executor.layers.utils import (
+    cuda_unquantized_gemm,
     dispatch_unquantized_gemm,
     maybe_pad_weight_avoid_cache_aliasing,
+    maybe_use_skinny_gemm,
 )
 from vllm.model_executor.parameter import (
     BasevLLMParameter,
@@ -228,6 +230,8 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
                 layer.weight.data = weight.t().contiguous().t()
         elif current_platform.is_rocm():
             layer.weight.data = maybe_pad_weight_avoid_cache_aliasing(layer.weight.data)
+        elif self._gemm_impl is cuda_unquantized_gemm:
+            maybe_use_skinny_gemm(layer)
 
     def apply(
         self,
diff --git a/vllm/model_executor/layers/logits_processor.py b/vllm/model_executor/layers/logits_processor.py
index 04541212fe0e..1ca1d71381bc 100644
--- a/vllm/model_executor/layers/logits_processor.py
+++ b/vllm/model_executor/layers/logits_processor.py
@@ -198,6 +198,21 @@ def _get_logits(
             logits = logits[..., : self.org_vocab_size]
         return logits
 
+    def _mask_outside_vocab(
+        self, logits: torch.Tensor, lm_head: VocabParallelEmbedding
+    ) -> None:
+        """Mask this shard's padding and any head rows past ``org_vocab_size``.
+
+        A head shared with another model can have rows past this processor's
+        vocabulary, as ``_get_logits`` drops them.
+        """
+        shard_indices = lm_head.shard_indices
+        if shard_indices.num_org_vocab_padding > 0:
+            logits[..., -shard_indices.num_org_vocab_padding :] = -float("inf")
+        cutoff = self.org_vocab_size - shard_indices.org_vocab_start_index
+        if cutoff < logits.shape[-1]:
+            logits[..., max(cutoff, 0) :] = -float("inf")
+
     def get_top_tokens(
         self,
         lm_head: VocabParallelEmbedding,
@@ -223,10 +238,7 @@ def get_top_tokens(
         if self.scale != 1.0:
             logits = logits * self.scale
 
-        # Mask out padding entries beyond org_vocab_size on this shard.
-        num_pad = lm_head.shard_indices.num_org_vocab_padding
-        if num_pad > 0:
-            logits[..., -num_pad:] = -float("inf")
+        self._mask_outside_vocab(logits, lm_head)
 
         local_max_vals, local_max_indices = logits.max(dim=-1)
 
@@ -284,10 +296,7 @@ def get_top_k_tokens(
             if self.scale != 1.0:
                 logits = logits * self.scale
 
-        # Mask out padding entries beyond org_vocab_size on this shard.
-        num_pad = lm_head.shard_indices.num_org_vocab_padding
-        if num_pad > 0:
-            logits[..., -num_pad:] = -float("inf")
+        self._mask_outside_vocab(logits, lm_head)
 
         if return_log_probs:
             log_partition = torch.logsumexp(logits.float(), dim=-1, keepdim=True)
diff --git a/vllm/model_executor/layers/mamba/checkpoint.py b/vllm/model_executor/layers/mamba/checkpoint.py
index da74702a9239..920bad6c8051 100644
--- a/vllm/model_executor/layers/mamba/checkpoint.py
+++ b/vllm/model_executor/layers/mamba/checkpoint.py
@@ -282,7 +282,7 @@ def store_cache_checkpoints_kernel(
     if HAS_TOKEN_MAP:
         token_idx = tl.load(token_indices_ptr + token_idx, mask=valid_conv, other=0)
     values = tl.load(
-        x_ptr + token_idx * x_stride_0 + width_idx * x_stride_1,
+        x_ptr + token_idx.to(tl.int64) * x_stride_0 + width_idx * x_stride_1,
         mask=valid_conv,
     )
     tl.store(
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 b712268a9460..64cc3055c9d5 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
@@ -28,11 +28,13 @@
 from vllm.forward_context import ForwardContext, get_forward_context
 from vllm.logger import init_logger
 from vllm.model_executor.custom_op import CustomOp, PluggableLayer
+from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation
 from vllm.model_executor.layers.layernorm import RMSNormGated
 from vllm.model_executor.layers.linear import (
     ColumnParallelLinear,
     MergedColumnParallelLinear,
     RowParallelLinear,
+    UnquantizedLinearMethod,
 )
 from vllm.model_executor.layers.mamba.abstract import (
     get_mamba_prefill_checkpoint_alignment,
@@ -55,6 +57,7 @@
 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.inc import INCConfig
+from vllm.model_executor.layers.utils import cuda_unquantized_gemm
 from vllm.model_executor.model_loader.weight_utils import (
     sharded_weight_loader,
 )
@@ -74,10 +77,12 @@
 )
 from vllm.third_party.flash_linear_attention.ops.utils import FLA_CHUNK_SIZE
 from vllm.utils.math_utils import cdiv
+from vllm.utils.multi_stream_utils import execute_in_parallel
 from vllm.utils.torch_utils import (
     LayerNameType,
     _encode_layer_name,
     _resolve_layer_name,
+    aux_stream,
     direct_register_custom_op,
 )
 from vllm.v1.attention.backends.gdn_attn import (
@@ -107,6 +112,10 @@
 logger = init_logger(__name__)
 
 MAX_FUSED_GDN_MTP_TOKENS = 8
+# Below this conservative threshold, overlap in_proj_ba with in_proj_qkvz on the aux
+# stream. On GB300, across Qwen3.5/3.8 shapes, quant formats and TP 1-8, the first
+# regression was at 272 tokens (27B BF16 TP8).
+_BA_PROJ_OVERLAP_MAX_TOKENS = 256
 FUSED_GDN_STATE_DTYPES = (torch.float32, torch.bfloat16)
 
 
@@ -324,6 +333,7 @@ def fi_chunk_gated_delta_rule(
     checkpoint_cu_starts: torch.Tensor | None = None,
     checkpoint_every_n_tokens: int = 0,
     use_cp: Literal["auto"] | bool = "auto",
+    output: torch.Tensor | None = None,
 ):
     from flashinfer.gdn_prefill import (
         chunk_gated_delta_rule as chunk_gated_delta_rule_fi,
@@ -341,19 +351,21 @@ def fi_chunk_gated_delta_rule(
     g = g.squeeze(0).contiguous()
     beta = beta.squeeze(0).contiguous()
     fi_state = initial_state.to(torch.float32)
+    # Like FlashInfer's own API, g is the already-exponentiated gate;
+    # fused_post_conv_prep emits it directly (output_g_exp=True).
     fi_g = g.to(torch.float32)
     fi_beta = beta.to(torch.float32)
-    if cu_seqlens is not None:
-        cu_seqlens = cu_seqlens.to(torch.int64)
+    # cu_seqlens goes through as is: FlashInfer accepts int32 and int64.
     result = chunk_gated_delta_rule_fi(
         q=q,
         k=k,
         v=v,
-        g=torch.exp(fi_g),
+        g=fi_g,
         beta=fi_beta,
         initial_state=fi_state,
         output_final_state=output_final_state,
         cu_seqlens=cu_seqlens,
+        output=output,
         state_checkpoints=state_checkpoints,
         checkpoint_cu_starts=checkpoint_cu_starts,
         checkpoint_every_n_tokens=checkpoint_every_n_tokens,
@@ -414,6 +426,9 @@ def forward_cuda(
         core_attn_out: torch.Tensor | None = None,
         aiter_prefill_metadata: object | None = None,
     ):
+        # g is in log space in this dispatch layer (shared with the Triton/FLA
+        # backends); FlashInfer takes exp(g), so exp here.
+        g = torch.exp(g.to(torch.float32))
         o, final_state = fi_chunk_gated_delta_rule(
             q=q,
             k=k,
@@ -562,6 +577,10 @@ def __init__(
         prefix: str = "",
         gqa_interleaved_layout=False,
         reduce_results: bool = True,
+        overlap_ba_proj: bool = False,
+        fused_gdn_gated_rms_norm: (
+            Callable[..., torch.Tensor | QuantizedActivation] | None
+        ) = None,
     ) -> None:
         super().__init__(config, vllm_config, prefix)
 
@@ -625,6 +644,21 @@ def __init__(
             prefix=f"{prefix}.in_proj_ba",
         )
         self.disable_tp_for_ba_proj = self.maybe_disable_tp(self.quant_config)
+        # Opt-in, for callers outside torch.compile (compiled callers would need a
+        # custom op or Dynamo stream tracing; not validated). Plain cuBLAS ba GEMM only:
+        # FlashInfer GEMMs share a process-wide workspace. Off with ubatching and LoRA.
+        self._ba_stream = None
+        ba_method = self.in_proj_ba.quant_method
+        if (
+            overlap_ba_proj
+            and current_platform.is_cuda()
+            and isinstance(ba_method, UnquantizedLinearMethod)
+            and ba_method._gemm_impl is cuda_unquantized_gemm
+            and not vllm_config.parallel_config.use_ubatching
+            and vllm_config.lora_config is None
+        ):
+            self._ba_stream = aux_stream()
+            self._ba_events = (torch.cuda.Event(), torch.cuda.Event())
 
         query_key_settings = (self.key_dim, 0, False)
         value_settings = (self.value_dim, 0, False)
@@ -703,6 +737,10 @@ def __init__(
         self.enable_packed_recurrent_decode = (
             envs.VLLM_ENABLE_FLA_PACKED_RECURRENT_DECODE
         )
+        # Off by default; the Qwen3.5 NVIDIA variant passes a fused gated RMS
+        # norm that reads z strided from the packed qkvz projection and
+        # returns the out_proj input: `out`, or a pre-quantized activation.
+        self._fused_gdn_gated_rms_norm = fused_gdn_gated_rms_norm
         self.gdn_decode_kernel = envs.VLLM_GDN_DECODE_KERNEL.strip().lower()
         if self.gdn_decode_kernel == "cuda" and current_platform.is_cuda_alike():
             reason = self._fused_gdn_decode_unsupported_reason(vllm_config)
@@ -1134,18 +1172,39 @@ def forward_hip(
     def forward_cuda(
         self,
         hidden_states: torch.Tensor,
+        qkvz_input: QuantizedActivation | None = None,
     ) -> torch.Tensor:
         """Forward pass with three parts:
         1. Input projection
         2. Core attention (custom op)
         3. Output projection
+
+        ``qkvz_input``, if given, is ``hidden_states`` pre-quantized for
+        in_proj_qkvz; in_proj_ba still reads ``hidden_states``.
         """
         num_tokens = hidden_states.size(0)
         # ============================================================
         # Part 1: Input Projection
         # ============================================================
-        mixed_qkvz, _ = self.in_proj_qkvz(hidden_states)
-        ba, _ = self.in_proj_ba(hidden_states)
+        qkvz_in = hidden_states if qkvz_input is None else qkvz_input
+        if self._ba_stream is None:
+            mixed_qkvz, _ = self.in_proj_qkvz(qkvz_in)
+            ba, _ = self.in_proj_ba(hidden_states)
+        else:
+            # Hide the tiny ba GEMM on the aux stream. Only under CUDA graph
+            # capture: eagerly the stream switch costs more host time than it
+            # saves. ba is enqueued first (aux-first order); qkvz-first
+            # regressed block-FP8 qkvz from ~100 tokens.
+            start, done = self._ba_events
+            mixed_qkvz, (ba,) = execute_in_parallel(
+                lambda: self.in_proj_qkvz(qkvz_in)[0],
+                [lambda: self.in_proj_ba(hidden_states)[0]],
+                start,
+                [done],
+                [self._ba_stream],
+                enable=num_tokens <= _BA_PROJ_OVERLAP_MAX_TOKENS
+                and torch.cuda.is_current_stream_capturing(),
+            )
 
         use_fused_gdn_spec_decode = (
             self.enable_fused_gdn_spec_decode
@@ -1181,28 +1240,54 @@ def forward_cuda(
             qkv_size = (self.key_dim * 2 + self.value_dim) // self.tp_size
             z_size = self.value_dim // self.tp_size
             mixed_qkv, z = mixed_qkvz.split([qkv_size, z_size], dim=-1)
-            z = z.reshape(z.size(0), -1, self.head_v_dim)
+            if self._fused_gdn_gated_rms_norm is None:
+                z = z.reshape(z.size(0), -1, self.head_v_dim)
             b, a = self.split_ba(ba)
 
         # ============================================================
         # Part 2: Core Attention (Custom Op)
         # ============================================================
-        # Note: we should not use torch.empty here like other attention backends,
-        # see discussions in https://github.com/vllm-project/vllm/pull/28182
-        core_attn_out = torch.zeros(
+        # Uninitialized like the other attention backends: _forward_core
+        # writes every row up to num_actual_tokens. Padding rows behind it
+        # are only read by row-wise ops and never reach sampled outputs.
+        core_attn_out = torch.empty(
             (num_tokens, self.num_v_heads // self.tp_size, self.head_v_dim),
             dtype=hidden_states.dtype,
             device=hidden_states.device,
         )
 
+        # b/a may stay strided chunk views of ba: the packed decode, prefill
+        # (fused_post_conv_prep) and ROCm fallback kernels consume row strides
+        # directly, and fused_sigmoid_gating_delta_rule_update contiguous()
+        # internally where needed.
         torch.ops.vllm.qwen_gdn_attention_core(
             mixed_qkv,
-            b.contiguous(),
-            a.contiguous(),
+            b,
+            a,
             core_attn_out,
             layer_name=_encode_layer_name(self.prefix),
         )
 
+        if (
+            self._fused_gdn_gated_rms_norm is not None
+            and not self.gqa_interleaved_layout
+        ):
+            out2d = core_attn_out.view(num_tokens, -1)
+            out_proj_input = self._fused_gdn_gated_rms_norm(
+                x=out2d,
+                z=z,
+                weight=self.norm.weight,
+                out=out2d,
+                out_proj=self.out_proj,
+                eps=self.norm.eps,
+                num_v_heads=self.num_v_heads // self.tp_size,
+                head_v_dim=self.head_v_dim,
+                norm_before_gate=self.norm.norm_before_gate,
+                activation=self.norm.activation,
+            )
+            output, _ = self.out_proj(out_proj_input)
+            return output
+
         # ============================================================
         # Part 3: Output Projection
         # ============================================================
@@ -1463,6 +1548,9 @@ def _warmup_flashinfer_prefill_variants(
         output. A single sequence under ``use_cp="auto"`` takes the CP route
         wherever the heuristic can pick it; ``use_cp=False`` forces the other.
         """
+        # g arrives in log space (from _warmup_prefill_kernels); FlashInfer
+        # takes exp(g).
+        g = torch.exp(g.to(torch.float32))
         checkpoint_buffers: list[torch.Tensor | None] = [None]
         if self._exports_checkpoints:
             (state_checkpoints,) = current_workspace_manager().get_simultaneous(
@@ -1738,7 +1826,7 @@ def _forward_core(
                 head_k_dim=self.head_k_dim,
                 head_v_dim=self.head_v_dim,
                 apply_l2norm=True,
-                output_g_exp=False,
+                output_g_exp=self.gdn_prefill_backend == "flashinfer",
             )
             query_non_spec = query_non_spec.unsqueeze(0)
             key_non_spec = key_non_spec.unsqueeze(0)
@@ -1804,6 +1892,7 @@ def _forward_core(
             core_attn_out_decode = None
 
         # 2.3: Process the remaining part (prefill chunk, or non-spec decode-only)
+        prefill_in_place = False
         if attn_metadata.num_prefills > 0:
             # State indices, initial-state mask and cu_seqlens for the chunk
             # kernel are precomputed by the metadata builder (the prefill tail
@@ -1813,10 +1902,31 @@ def _forward_core(
             prefill_has_initial_state = attn_metadata.prefill_has_initial_state
             assert prefill_state_indices is not None
             assert prefill_has_initial_state is not None
-            initial_state = ssm_state[prefill_state_indices]
-            initial_state[~prefill_has_initial_state, ...] = 0
-            plan = attn_metadata.flashinfer_checkpoint_plan
-            if plan is None:
+            if self.gdn_prefill_backend == "flashinfer":
+                # Without spec decodes the prefill rows follow the peeled
+                # decodes contiguously, so FlashInfer writes them straight
+                # into core_attn_out. Spec batches scatter by token index.
+                prefill_out = None
+                if spec_sequence_masks is None:
+                    prefill_out = core_attn_out[num_decode_tokens:num_actual_tokens]
+                core_attn_out_non_spec, last_recurrent_state = (
+                    self._forward_flashinfer_prefill(
+                        query_non_spec,
+                        key_non_spec,
+                        value_non_spec,
+                        g_non_spec,
+                        beta_non_spec,
+                        attn_metadata,
+                        ssm_state=ssm_state,
+                        conv_state=conv_state,
+                        conv_input=conv_input_non_spec,
+                        output=prefill_out,
+                    )
+                )
+                prefill_in_place = prefill_out is not None
+            else:
+                initial_state = ssm_state[prefill_state_indices]
+                initial_state[~prefill_has_initial_state, ...] = 0
                 (
                     core_attn_out_non_spec,
                     last_recurrent_state,
@@ -1834,44 +1944,12 @@ def _forward_core(
                     use_qk_l2norm_in_kernel=False,
                     aiter_prefill_metadata=attn_metadata.aiter_prefill_metadata,
                 )
-            else:
-                # Only FlashInfer's prefill emits checkpoints.
-                (state_checkpoints,) = current_workspace_manager().get_simultaneous(
-                    self._checkpoint_buffer_spec
-                )
-                state_checkpoints = state_checkpoints[: plan.num_checkpoints]
-                (
-                    core_attn_out_non_spec,
-                    last_recurrent_state,
-                ) = fi_chunk_gated_delta_rule(
-                    q=query_non_spec,
-                    k=key_non_spec,
-                    v=value_non_spec,
-                    g=g_non_spec,
-                    beta=beta_non_spec,
-                    initial_state=initial_state,
-                    output_final_state=True,
-                    cu_seqlens=attn_metadata.prefill_query_start_loc,
-                    use_qk_l2norm_in_kernel=False,
-                    state_checkpoints=state_checkpoints,
-                    checkpoint_cu_starts=plan.cu_starts,
-                    checkpoint_every_n_tokens=plan.every_n_tokens,
-                )
-                assert attn_metadata.checkpoint is not None
-                assert non_spec_query_start_loc is not None
-                self._checkpoint_exporter.export(
-                    attn_metadata.checkpoint,
-                    raw_qkv=conv_input_non_spec,
-                    conv_state=conv_state,
-                    recurrent_checkpoint=state_checkpoints,
-                    recurrent_state=ssm_state,
-                    cu_seqlens=non_spec_query_start_loc,
-                    recurrent_rows=plan.rows,
-                )
             # Init cache
             ssm_state[prefill_state_indices] = last_recurrent_state.to(ssm_state.dtype)
 
-            if split_non_spec:
+            if split_non_spec and prefill_in_place:
+                core_attn_out[:num_decode_tokens] = core_attn_out_decode.squeeze(0)
+            elif split_non_spec:
                 # Stitch the peeled decode outputs in front of the prefill
                 # outputs (decode-first order).
                 core_attn_out_non_spec = torch.cat(
@@ -1908,9 +1986,81 @@ def _forward_core(
             )
         elif spec_sequence_masks is not None:
             core_attn_out[:num_actual_tokens] = core_attn_out_spec.squeeze(0)
-        else:
+        elif not prefill_in_place:
             core_attn_out[:num_actual_tokens] = core_attn_out_non_spec.squeeze(0)
 
+    def _forward_flashinfer_prefill(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        g: torch.Tensor,
+        beta: torch.Tensor,
+        attn_metadata: GDNAttentionMetadata,
+        *,
+        ssm_state: torch.Tensor,
+        conv_state: torch.Tensor,
+        conv_input: torch.Tensor | None,
+        output: torch.Tensor | None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """FlashInfer chunked prefill of the batch's prefill rows.
+
+        Writes the output into ``output`` when given and exports the
+        batch's internal prefill checkpoints. Returns the output and the
+        rows' final states.
+        """
+        has_initial_state = attn_metadata.prefill_has_initial_state
+        assert has_initial_state is not None
+        initial_state = ssm_state[attn_metadata.prefill_state_indices]
+        initial_state[~has_initial_state, ...] = 0
+        plan = attn_metadata.flashinfer_checkpoint_plan
+        if plan is None:
+            return fi_chunk_gated_delta_rule(
+                q=q,
+                k=k,
+                v=v,
+                g=g,
+                beta=beta,
+                initial_state=initial_state,
+                output_final_state=True,
+                cu_seqlens=attn_metadata.prefill_query_start_loc,
+                use_qk_l2norm_in_kernel=False,
+                output=output,
+            )
+
+        (state_checkpoints,) = current_workspace_manager().get_simultaneous(
+            self._checkpoint_buffer_spec
+        )
+        state_checkpoints = state_checkpoints[: plan.num_checkpoints]
+        core_out, final_state = fi_chunk_gated_delta_rule(
+            q=q,
+            k=k,
+            v=v,
+            g=g,
+            beta=beta,
+            initial_state=initial_state,
+            output_final_state=True,
+            cu_seqlens=attn_metadata.prefill_query_start_loc,
+            use_qk_l2norm_in_kernel=False,
+            output=output,
+            state_checkpoints=state_checkpoints,
+            checkpoint_cu_starts=plan.cu_starts,
+            checkpoint_every_n_tokens=plan.every_n_tokens,
+        )
+        assert attn_metadata.checkpoint is not None
+        assert attn_metadata.non_spec_query_start_loc is not None
+        assert conv_input is not None
+        self._checkpoint_exporter.export(
+            attn_metadata.checkpoint,
+            raw_qkv=conv_input,
+            conv_state=conv_state,
+            recurrent_checkpoint=state_checkpoints,
+            recurrent_state=ssm_state,
+            cu_seqlens=attn_metadata.non_spec_query_start_loc,
+            recurrent_rows=plan.rows,
+        )
+        return core_out, final_state
+
     def _forward_core_decode_aiter(
         self,
         qkvz: torch.Tensor,
diff --git a/vllm/model_executor/layers/mamba/mamba_mixer2.py b/vllm/model_executor/layers/mamba/mamba_mixer2.py
index 76b13c69748d..680e587285c8 100644
--- a/vllm/model_executor/layers/mamba/mamba_mixer2.py
+++ b/vllm/model_executor/layers/mamba/mamba_mixer2.py
@@ -1211,9 +1211,39 @@ def share_replayssm_ring_trackers(
                     "ReplaySSM layers in one cache group must share cache capacity"
                 )
 
-        ring_start = torch.zeros(num_blocks, dtype=torch.int32, device=device)
-        prev_num_accepted = torch.zeros_like(ring_start)
-        prev_query_len = torch.zeros_like(ring_start)
+        # Keep the trackers a previous bind created when they still cover the cache,
+        # as prefix views: CUDA graphs captured since then hold their addresses,
+        # and the kernels want the tracker length to match the cache. An extensible
+        # KV cache rebinds after capture with the same or fewer blocks.
+        trackers = (
+            first_mixer._replayssm_ring_start,
+            first_mixer._replayssm_prev_num_accepted,
+            first_mixer._replayssm_prev_query_len,
+        )
+        if (
+            trackers[0].numel() >= num_blocks
+            and trackers[0].device == device
+            and all(t.shape == trackers[0].shape for t in trackers)
+            and all(
+                mixer_tracker is tracker
+                for mixer in (replayssm_mixers[n] for n in group_layer_names)
+                for mixer_tracker, tracker in zip(
+                    (
+                        mixer._replayssm_ring_start,
+                        mixer._replayssm_prev_num_accepted,
+                        mixer._replayssm_prev_query_len,
+                    ),
+                    trackers,
+                )
+            )
+        ):
+            ring_start, prev_num_accepted, prev_query_len = (
+                t[:num_blocks] for t in trackers
+            )
+        else:
+            ring_start = torch.zeros(num_blocks, dtype=torch.int32, device=device)
+            prev_num_accepted = torch.zeros_like(ring_start)
+            prev_query_len = torch.zeros_like(ring_start)
         for layer_name in group_layer_names:
             mixer = replayssm_mixers[layer_name]
             mixer._replayssm_ring_start = ring_start
diff --git a/vllm/model_executor/layers/pooler/tokwise/heads.py b/vllm/model_executor/layers/pooler/tokwise/heads.py
index 95242ee470a7..e853bf7c89ff 100644
--- a/vllm/model_executor/layers/pooler/tokwise/heads.py
+++ b/vllm/model_executor/layers/pooler/tokwise/heads.py
@@ -1,12 +1,13 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from abc import ABC, abstractmethod
-from collections.abc import Set
+from collections.abc import Callable, Set
 from typing import TypeAlias
 
 import torch
 import torch.nn as nn
 
+import vllm.envs as envs
 from vllm.model_executor.layers.pooler import ActivationFn, ClassifierFn, ProjectorFn
 from vllm.pooling_params import PoolingParams
 from vllm.tasks import PoolingTask
@@ -88,6 +89,57 @@ def extra_repr(self) -> str:
     def get_supported_tasks(self) -> Set[PoolingTask]:
         return {"token_embed"}
 
+    def forward_batched(
+        self,
+        hidden_states: torch.Tensor,
+        pooling_params: list[PoolingParams],
+        split: Callable[[torch.Tensor], list[TokenPoolingMethodOutputItem]],
+    ) -> list[TokenPoolerHeadOutputItem] | None:
+        """Batch projection, splitting only when finishing needs per-request options.
+
+        Keep eligibility checks in sync with PoolingParams fields read by forward_chunk.
+        """
+        if envs.VLLM_BATCH_INVARIANT or any(
+            p.late_chunking_params is not None for p in pooling_params
+        ):
+            return None
+        # [total_tokens, hidden_size] -> [total_tokens, embedding_size]
+        embeddings = self._project(hidden_states)
+        first = pooling_params[0]
+        if all(
+            p.dimensions == first.dimensions
+            and p.use_activation == first.use_activation
+            for p in pooling_params
+        ):
+            embeddings = self._finish(embeddings, first)
+            # Each output has shape [request_tokens, output_size].
+            return split(embeddings)
+        # Each part has shape [request_tokens, embedding_size].
+        pooled_data = split(embeddings)
+        del embeddings
+        return [
+            self._finish(part, param) if part is not None else None
+            for part, param in zip(pooled_data, pooling_params, strict=True)
+        ]
+
+    def _project(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        """Project [num_rows, hidden_size] to [num_rows, embedding_size]."""
+        if self.head_dtype is not None:
+            hidden_states = hidden_states.to(self.head_dtype)
+        if self.projector is not None:
+            return self.projector(hidden_states)
+        return hidden_states
+
+    def _finish(
+        self, embeddings: torch.Tensor, pooling_param: PoolingParams
+    ) -> torch.Tensor:
+        """Slice/normalize [num_rows, embedding_size] to [num_rows, output_size]."""
+        if pooling_param.dimensions is not None:
+            embeddings = embeddings[..., : pooling_param.dimensions]
+        if self.activation is not None and pooling_param.use_activation:
+            embeddings = self.activation(embeddings)
+        return embeddings
+
     def forward_chunk(
         self,
         pooled_data: TokenPoolingMethodOutputItem,
@@ -97,33 +149,16 @@ def forward_chunk(
         if pooled_data is None:
             return None
 
+        # pooled_data: [num_tokens, hidden_size]
         if (late_chunking := pooling_param.late_chunking_params) is not None:
             means = _mean_pool_chunks(pooled_data, late_chunking.chunk_size)
             pooled_data = (
                 means.to(pooled_data.dtype) if self.head_dtype is None else means
             )
+            # pooled_data: [num_chunks, hidden_size]
 
-        if self.head_dtype is not None:
-            pooled_data = pooled_data.to(self.head_dtype)
-        # pooled_data shape: [n_tokens, hidden_size]
-
-        # Apply ST projector
-        if self.projector is not None:
-            embeddings = self.projector(pooled_data)
-        else:
-            embeddings = pooled_data
-        # embeddings shape: [n_tokens, embedding_size]
-
-        # for matryoshka representation
-        if pooling_param.dimensions is not None:
-            embeddings = embeddings[..., : pooling_param.dimensions]
-
-        # for normalize
-        if self.activation is not None and pooling_param.use_activation:
-            embeddings = self.activation(embeddings)
-
-        # embeddings shape: [n_tokens, embedding_size]
-        return embeddings
+        # Output: [num_rows, output_size], one row per token or late chunk.
+        return self._finish(self._project(pooled_data), pooling_param)
 
 
 class TokenClassifierPoolerHead(TokenPoolerHead):
diff --git a/vllm/model_executor/layers/pooler/tokwise/poolers.py b/vllm/model_executor/layers/pooler/tokwise/poolers.py
index 0dab2d5b1d57..6ab77edb7820 100644
--- a/vllm/model_executor/layers/pooler/tokwise/poolers.py
+++ b/vllm/model_executor/layers/pooler/tokwise/poolers.py
@@ -28,6 +28,7 @@
     TokenPoolerHeadOutputItem,
 )
 from .methods import (
+    AllPool,
     TokenPoolingMethod,
     TokenPoolingMethodOutputItem,
     get_tok_pooling_method,
@@ -91,6 +92,22 @@ def forward(
         hidden_states: torch.Tensor,
         pooling_metadata: PoolingMetadata,
     ) -> TokenPoolerOutput:
+        params = pooling_metadata.pooling_params
+        if (
+            params
+            and type(self.pooling) is AllPool
+            and not self.pooling.enable_chunked_prefill
+            and type(self.head) is TokenEmbeddingPoolerHead
+            and (
+                outputs := self.head.forward_batched(
+                    hidden_states,
+                    params,
+                    lambda embeddings: self.pooling(embeddings, pooling_metadata),
+                )
+            )
+            is not None
+        ):
+            return outputs
         pooled_data = self.pooling(hidden_states, pooling_metadata)
         if self.head is not None:
             pooled_data = self.head(pooled_data, pooling_metadata)
diff --git a/vllm/model_executor/layers/quantization/base_config.py b/vllm/model_executor/layers/quantization/base_config.py
index 5bc1e42aa5da..db0d5524f941 100644
--- a/vllm/model_executor/layers/quantization/base_config.py
+++ b/vllm/model_executor/layers/quantization/base_config.py
@@ -3,6 +3,8 @@
 
 import inspect
 from abc import ABC, abstractmethod
+from collections.abc import Iterator
+from contextlib import nullcontext
 from typing import TYPE_CHECKING, Any
 
 import regex as re
@@ -11,6 +13,7 @@
 from transformers import PreTrainedConfig
 
 if TYPE_CHECKING:
+    from vllm.config.cache import CacheDType
     from vllm.model_executor.layers.quantization import QuantizationMethods
     from vllm.model_executor.layers.quantization.online.base import (
         OnlineQuantizationConfig,
@@ -86,8 +89,17 @@ def process_weights_after_loading(self, layer: nn.Module) -> None:
 
     def dequantize_weight(
         self, layer: nn.Module
-    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
-        """Materialize a serialized quantized weight for requantization."""
+    ) -> torch.Tensor | Iterator[tuple[slice, torch.Tensor, torch.Tensor]]:
+        """Materialize a serialized quantized weight for requantization.
+
+        Args:
+            layer: Layer holding the serialized weights.
+
+        Returns:
+            The dequantized weight for linear layers. For MoE layers, an
+            iterator of ``(experts, w13, w2)`` over groups of experts.
+
+        """
         raise NotImplementedError(
             f"The quantization method {type(self)} does not implement "
             "dequantize_weight, so re-quantization from this checkpoint is "
@@ -258,6 +270,10 @@ def get_checkpoint_weight_mapper() -> "WeightsMapper":
 
         return WeightsMapper(orig_to_new_suffix={".g_idx": None})
 
+    def get_kv_cache_dtype(self, prefix: str) -> "CacheDType | None":
+        """Return a checkpoint-selected KV-cache dtype for one layer."""
+        return None
+
     def apply_vllm_mapper(  # noqa: B027
         self, hf_to_vllm_mapper: "WeightsMapper"
     ):
@@ -308,35 +324,44 @@ def resolve_quant_method(
     from vllm.model_executor.layers.quantization.online.moe_base import (
         OnlineMoEMethodBase,
     )
+    from vllm.model_executor.layers.quantization.online.requantization import (
+        requantization_source_scope,
+    )
 
-    base_quant_method = quant_config.get_quant_method(layer, prefix)
-    if quant_config.online_quantization_config is None:
-        # No online configuration: retain the checkpoint method.
-        return base_quant_method
+    online_config = quant_config.online_quantization_config
     # Online quantization currently supports only LinearBase and RoutedExperts.
     # Embeddings and ParallelLMHead retain their checkpoint quantization method.
-    if not isinstance(layer, (LinearBase, RoutedExperts)):
+    is_online_supported_layer = isinstance(layer, (LinearBase, RoutedExperts))
+
+    online_target = None
+    if online_config is not None and is_online_supported_layer:
+        online_config.packed_modules_mapping = quant_config.packed_modules_mapping
+        online_target = online_config.resolve_quant_method_cls(type(layer), prefix)
+
+    # The checkpoint method of requantized experts is only a dequantization source.
+    with (
+        requantization_source_scope()
+        if online_target is not None and isinstance(layer, RoutedExperts)
+        else nullcontext()
+    ):
+        base_quant_method = quant_config.get_quant_method(layer, prefix)
+    if online_config is None:
+        # No online configuration: retain the checkpoint method.
+        return base_quant_method
+    if not is_online_supported_layer:
         # Online quantization only supports linear and routed MoE layers.
         return base_quant_method
 
-    quant_config.online_quantization_config.packed_modules_mapping = (
-        quant_config.packed_modules_mapping
-    )
     checkpoint_is_quantized = base_quant_method is not None and not isinstance(
         base_quant_method, (UnquantizedLinearMethod, UnquantizedFusedMoEMethod)
     )
-    online_target = quant_config.online_quantization_config.resolve_quant_method_cls(
-        type(layer), prefix
-    )
 
     if checkpoint_is_quantized:
         if online_target is None:
             # The checkpoint quant method is applied as there is no online override.
             return base_quant_method
 
-        online_quant_method = quant_config.online_quantization_config.get_quant_method(
-            layer, prefix
-        )
+        online_quant_method = online_config.get_quant_method(layer, prefix)
 
         assert base_quant_method is not None and not isinstance(
             base_quant_method, (UnquantizedLinearMethod, UnquantizedFusedMoEMethod)
@@ -353,4 +378,4 @@ def resolve_quant_method(
         return base_quant_method
 
     # Quantize an unquantized layer with the online method.
-    return quant_config.online_quantization_config.get_quant_method(layer, prefix)
+    return online_config.get_quant_method(layer, prefix)
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 f5b73ade5338..3d17901029bd 100644
--- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py
+++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py
@@ -1271,11 +1271,3 @@ def _to_scalar(tensor: torch.Tensor) -> float:
         # Sync host (cpu) scale copies read by AITER fused kernels.
         layer._k_scale_cpu.fill_(layer._k_scale_float)
         layer._v_scale_cpu.fill_(layer._v_scale_float)
-
-        # Discard all placeholders.
-        del layer.k_scale
-        del layer.v_scale
-        del layer.q_scale
-        del layer.k_zero_point
-        del layer.v_zero_point
-        del layer.q_zero_point
diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4.py
index 227382ff1d03..7bdcda4bf7b1 100644
--- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4.py
+++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4.py
@@ -43,7 +43,11 @@
 from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
     prepare_moe_fp4_layer_for_marlin,
 )
-from vllm.model_executor.utils import replace_parameter, set_weight_attrs
+from vllm.model_executor.utils import (
+    move_parameter,
+    replace_parameter,
+    set_weight_attrs,
+)
 from vllm.platforms import current_platform
 
 logger = init_logger(__name__)
@@ -168,15 +172,8 @@ def get_fused_moe_quant_config(
             )
 
     def process_weights_after_loading(self, layer: RoutedExperts) -> None:
-        layer.w13_weight = torch.nn.Parameter(
-            layer.w13_weight_packed.data, requires_grad=False
-        )
-        delattr(layer, "w13_weight_packed")
-
-        layer.w2_weight = torch.nn.Parameter(
-            layer.w2_weight_packed.data, requires_grad=False
-        )
-        delattr(layer, "w2_weight_packed")
+        move_parameter(layer, "w13_weight_packed", "w13_weight")
+        move_parameter(layer, "w2_weight_packed", "w2_weight")
 
         if self.use_cutlass_mxfp4:
             # Swizzle weight scales from flat checkpoint layout [E, N, K//32]
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 e7fc3f1df365..312c4f65a8b4 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
@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from collections.abc import Iterator
 
 import torch
 
@@ -26,14 +27,23 @@
 from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import (  # noqa E501
     CompressedTensorsMoEMethod,
 )
+from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
+    trtllm_nvfp4_hidden_alignment,
+)
 from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import (
     dequantize_nvfp4_moe_weights,
 )
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    QuantKey,
+    is_dynamic_per_token_outer_scale,
     kNvfp4Dynamic,
     kNvfp4Static,
 )
-from vllm.model_executor.utils import replace_parameter, set_weight_attrs
+from vllm.model_executor.utils import (
+    move_parameter,
+    replace_parameter,
+    set_weight_attrs,
+)
 
 # NVFP4 backends that have been verified to support EPLB for this quantization recipe.
 _EPLB_SUPPORTED_NVFP4_BACKENDS = frozenset(
@@ -54,16 +64,24 @@ def __init__(
         moe: FusedMoEConfig,
         layer_name: str | None = None,
         use_a16: bool = False,
+        activation_key: QuantKey = kNvfp4Dynamic,
     ):
         super().__init__(moe)
         self.group_size = 16
         self.use_a16 = use_a16
+        # FC1 input scheme; a per-token one (``kNvfp4DynamicToken``) replaces
+        # the checkpoint's static input scales. Only the backends that quantize
+        # the input themselves accept it, so selection enforces the invariant.
+        self.activation_key: QuantKey | None = None if use_a16 else activation_key
+        self.per_token_activation = is_dynamic_per_token_outer_scale(
+            self.activation_key
+        )
 
         # Select experts implementation.
         self.nvfp4_backend, self.experts_cls = select_nvfp4_moe_backend(
             config=self.moe,
             weight_key=kNvfp4Static,
-            activation_key=None if use_a16 else kNvfp4Dynamic,
+            activation_key=self.activation_key,
         )
         self.use_global_sf = is_global_sf_supported_for_nvfp4_backend(
             self.nvfp4_backend
@@ -190,8 +208,8 @@ def create_weights(
 
     def dequantize_weight(
         self, layer: torch.nn.Module
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        """Materialize serialized NVFP4 expert weights for requantization."""
+    ) -> Iterator[tuple[slice, torch.Tensor, torch.Tensor]]:
+        """Yield ``(experts, w13, w2)`` dequantized NVFP4 experts."""
         return dequantize_nvfp4_moe_weights(
             layer.w13_weight_packed,
             layer.w13_weight_scale,
@@ -206,19 +224,9 @@ def dequantize_weight(
 
     def process_weights_after_loading(self, layer: RoutedExperts) -> None:
         """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
-        # remove MKM, we should revert this change to ensure compatibility.
-        layer.w13_weight = torch.nn.Parameter(
-            layer.w13_weight_packed.data, requires_grad=False
-        )
-        delattr(layer, "w13_weight_packed")
-
-        layer.w2_weight = torch.nn.Parameter(
-            layer.w2_weight_packed.data, requires_grad=False
-        )
-        delattr(layer, "w2_weight_packed")
+        # MKM requires wN_weight; the loaded wN_weight_packed params stay registered.
+        move_parameter(layer, "w13_weight_packed", "w13_weight")
+        move_parameter(layer, "w2_weight_packed", "w2_weight")
 
         # Use a single gscale for w13.
         if self.moe.is_act_and_mul and not torch.allclose(
@@ -230,6 +238,14 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None:
             )
         w13_weight_global_scale = layer.w13_weight_global_scale[:, 0].contiguous()
 
+        a13_scale = 1.0 / layer.w13_input_global_scale
+        a2_scale = 1.0 / layer.w2_input_global_scale
+        if self.per_token_activation:
+            # Per-token scales cover the FC1 input, and the per-token kernel
+            # requantizes the FC2 input per token: static scales are neutral.
+            a13_scale = torch.ones_like(a13_scale)
+            a2_scale = torch.ones_like(a2_scale)
+
         # Shuffle weights into the NvFp4 kernel format.
         (
             w13,
@@ -246,13 +262,17 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None:
             w13=layer.w13_weight,
             w13_scale=layer.w13_weight_scale,
             w13_scale_2=(1.0 / w13_weight_global_scale),
-            a13_scale=(1.0 / layer.w13_input_global_scale),
+            a13_scale=a13_scale,
             w2=layer.w2_weight,
             w2_scale=layer.w2_weight_scale,
             w2_scale_2=(1.0 / layer.w2_weight_global_scale),
-            a2_scale=(1.0 / layer.w2_input_global_scale),
+            a2_scale=a2_scale,
             is_act_and_mul=self.moe.is_act_and_mul,
             use_a16=self.use_a16,
+            trtllm_hidden_alignment=trtllm_nvfp4_hidden_alignment(
+                per_token_activation=self.per_token_activation,
+                is_act_and_mul=self.moe.is_act_and_mul,
+            ),
         )
 
         replace_parameter(layer, "w13_weight", w13)
@@ -273,6 +293,7 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None:
             experts_cls=self.experts_cls,
             backend=self.nvfp4_backend,
             routing_tables=layer._expert_routing_tables(),
+            per_token_activation=self.per_token_activation,
         )
         self.moe_kernel.fused_experts.process_weights_after_loading(layer)
 
diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a8_fp8.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a8_fp8.py
index 4a2223aab8bb..67feb5b11c34 100644
--- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a8_fp8.py
+++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a8_fp8.py
@@ -25,7 +25,11 @@
 from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import (  # noqa E501
     CompressedTensorsMoEMethod,
 )
-from vllm.model_executor.utils import replace_parameter, set_weight_attrs
+from vllm.model_executor.utils import (
+    move_parameter,
+    replace_parameter,
+    set_weight_attrs,
+)
 
 
 class CompressedTensorsW4A8Fp8MoEMethod(CompressedTensorsMoEMethod):
@@ -149,8 +153,7 @@ def create_weights(
 
     def process_weights_after_loading(self, layer: RoutedExperts) -> None:
         for name in ("w13", "w2"):
-            setattr(layer, f"{name}_weight", getattr(layer, f"{name}_weight_packed"))
-            delattr(layer, f"{name}_weight_packed")
+            move_parameter(layer, f"{name}_weight_packed", f"{name}_weight")
 
         (
             w13_weight,
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 f3952afe0dc1..c071fbc93df0 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
@@ -30,7 +30,12 @@
     QuantKey,
     ScaleDesc,
 )
-from vllm.model_executor.utils import replace_parameter, set_weight_attrs
+from vllm.model_executor.utils import (
+    move_parameter,
+    replace_parameter,
+    retire_parameter,
+    set_weight_attrs,
+)
 
 logger = init_logger(__name__)
 
@@ -203,47 +208,17 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
             )
         )
 
-        # Register packed weights as parameters
-        replace_parameter(
-            layer,
-            "w13_weight_packed",
-            torch.nn.Parameter(w13_packed, requires_grad=False),
-        )
-        replace_parameter(
-            layer,
-            "w2_weight_packed",
-            torch.nn.Parameter(w2_packed, requires_grad=False),
-        )
-
-        replace_parameter(
-            layer, "w13_weight", torch.nn.Parameter(torch.empty(0), requires_grad=False)
-        )
-        replace_parameter(
-            layer, "w2_weight", torch.nn.Parameter(torch.empty(0), requires_grad=False)
-        )
-        replace_parameter(
-            layer,
-            "w13_weight_scale",
-            torch.nn.Parameter(w13_weight_scale, requires_grad=False),
-        )
-        replace_parameter(
-            layer,
-            "w2_weight_scale",
-            torch.nn.Parameter(w2_weight_scale, requires_grad=False),
-        )
-
+        # The packed weights (with scales and biases folded in) go into the loaded
+        # weight storage, one param under both names; the rest keep their data.
+        for name, packed in (("w13", w13_packed), ("w2", w2_packed)):
+            replace_parameter(layer, f"{name}_weight", packed)
+            del layer._parameters[f"{name}_weight_packed"]  # never-loaded placeholder
+            move_parameter(layer, f"{name}_weight", f"{name}_weight_packed")
+        for name in ("w13_weight_scale", "w2_weight_scale"):
+            retire_parameter(layer, name)
         if self.has_bias:
-            replace_parameter(
-                layer,
-                "w13_bias",
-                torch.nn.Parameter(w13_bias, requires_grad=False),
-            )
-        if self.has_bias:
-            replace_parameter(
-                layer,
-                "w2_bias",
-                torch.nn.Parameter(w2_bias, requires_grad=False),
-            )
+            retire_parameter(layer, "w13_bias")
+            retire_parameter(layer, "w2_bias")
 
         quant_config = self.get_fused_moe_quant_config(layer)
         assert quant_config is not None
diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a8_mxfp4.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a8_mxfp4.py
index ae2c88044f1f..d22b7ac1dd2b 100644
--- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a8_mxfp4.py
+++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a8_mxfp4.py
@@ -3,6 +3,7 @@
 
 from vllm.model_executor.layers.fused_moe import RoutedExperts
 from vllm.model_executor.layers.quantization.mxfp4 import Mxfp4MoEMethod
+from vllm.model_executor.utils import move_parameter
 
 _PACKED_WEIGHTS = ("w13_weight", "w2_weight")
 
@@ -28,7 +29,5 @@ def create_weights(self, layer: RoutedExperts, *args, **kwargs):
 
     def process_weights_after_loading(self, layer: RoutedExperts) -> None:
         for name in _PACKED_WEIGHTS:
-            param = getattr(layer, f"{name}_packed")
-            delattr(layer, f"{name}_packed")
-            layer.register_parameter(name, param)
+            move_parameter(layer, f"{name}_packed", name)
         super().process_weights_after_loading(layer)
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 4a410bce6ac8..f0eead077962 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
@@ -3,7 +3,6 @@
 from collections.abc import Callable
 
 import torch
-from torch.nn.parameter import Parameter
 
 from vllm.model_executor.kernels.linear import init_mxfp4_linear_kernel
 from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
@@ -17,6 +16,7 @@
     GroupQuantScaleParameter,
     ModelWeightParameter,
 )
+from vllm.model_executor.utils import move_parameter
 
 __all__ = ["CompressedTensorsW4A4Mxfp4"]
 
@@ -88,8 +88,7 @@ def create_weights(
         layer.register_parameter("weight_scale", weight_scale)
 
     def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
-        layer.weight = Parameter(layer.weight_packed.data, requires_grad=False)
-        del layer.weight_packed
+        move_parameter(layer, "weight_packed", "weight")
         self.kernel.process_weights_after_loading(layer)
 
     def apply_weights(
diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4.py b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4.py
index 78924681e9ae..7a06072bd901 100644
--- a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4.py
+++ b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4.py
@@ -17,7 +17,7 @@
     ModelWeightParameter,
     PerTensorScaleParameter,
 )
-from vllm.model_executor.utils import replace_parameter
+from vllm.model_executor.utils import move_parameter, replace_parameter
 
 logger = init_logger(__name__)
 
@@ -94,8 +94,7 @@ def create_weights(
 
     def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
         # Rename CT checkpoint names to standardized names
-        layer.weight = layer.weight_packed
-        del layer.weight_packed
+        move_parameter(layer, "weight_packed", "weight")
 
         # Check for mismatched weight global scales
         if torch.unique(layer.weight_global_scale).numel() != 1:
diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a16_fp8.py b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a16_fp8.py
index 16a0d9242a18..cdc027fdb367 100644
--- a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a16_fp8.py
+++ b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a16_fp8.py
@@ -121,13 +121,9 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
         if self.strategy == QuantizationStrategy.BLOCK:
             assert self.is_static_input_scheme is False
             # MarlinFP8ScaledMMLinearKernel uses "weight_scale_inv" for block
-            # quant, while CT registers the scale as "weight_scale".
-            # Rename by deleting the old parameter and adding the new one so
-            # that prepare_fp8_layer_for_marlin (which prefers "weight_scale"
-            # over "weight_scale_inv") picks up "weight_scale_inv" correctly.
-            weight_scale_data = layer.weight_scale.data
-            del layer._parameters["weight_scale"]
-            replace_parameter(layer, "weight_scale_inv", weight_scale_data)
+            # quant, while CT registers the scale as "weight_scale". Alias the
+            # loaded parameter (same object) under both names.
+            layer.weight_scale_inv = layer.weight_scale
         else:
             if self.strategy == QuantizationStrategy.TENSOR:
                 # For fused modules with per-tensor scales, expand each scale
diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa8o8.py b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa8o8.py
index dd1b986dc3cd..28e191bf276e 100644
--- a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa8o8.py
+++ b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa8o8.py
@@ -31,6 +31,7 @@
     ModelWeightParameter,
     PackedvLLMParameter,
 )
+from vllm.model_executor.utils import move_parameter, replace_parameter
 from vllm.scalar_type import scalar_types
 
 __all__ = ["CompressedTensorsWNA8O8Int", "fake_quant_static_int8"]
@@ -211,22 +212,14 @@ def _pack_int_quantized_weight(self, layer: torch.nn.Module) -> None:
         weight = layer.weight
         out_features, in_features = weight.shape
         packed = pack_to_int32(weight.data.contiguous(), self.num_bits)
-        delattr(layer, "weight")
+        # One param under both names; the packed data goes into the loaded storage.
+        move_parameter(layer, "weight", "weight_packed")
+        replace_parameter(layer, "weight_packed", packed.contiguous())
+        weight.packed_dim, weight.packed_factor = 1, self.pack_factor
 
         def _noop_loader(*_, **__):
             return None
 
-        layer.register_parameter(
-            "weight_packed",
-            PackedvLLMParameter(
-                data=packed.contiguous(),
-                input_dim=1,
-                output_dim=0,
-                packed_dim=1,
-                packed_factor=self.pack_factor,
-                weight_loader=_noop_loader,
-            ),
-        )
         layer.register_parameter(
             "weight_shape",
             BasevLLMParameter(
@@ -241,7 +234,6 @@ def _take_act_scale(layer, name: str) -> torch.Tensor | None:
         if param is None:
             return None
         scale = param.data.clone()
-        delattr(layer, name)
         return None if float(scale.reshape(-1)[0]) == 0.0 else scale
 
     def apply_weights(
diff --git a/vllm/model_executor/layers/quantization/fbgemm_fp8.py b/vllm/model_executor/layers/quantization/fbgemm_fp8.py
index 06589628215d..5afd493fc3f7 100644
--- a/vllm/model_executor/layers/quantization/fbgemm_fp8.py
+++ b/vllm/model_executor/layers/quantization/fbgemm_fp8.py
@@ -169,9 +169,8 @@ def process_weights_after_loading(self, layer: Module) -> None:
         # The param keeps its checkpoint dims; Humming reads the (K, N) layout.
         layer.processed_weight_dims = (0, 1)
         if self.quant_config.use_marlin:
+            # Activations are not quantized for marlin (input_scale_ub unused).
             prepare_fp8_layer_for_marlin(layer)
-            # Activations not quantized for marlin.
-            del layer.input_scale_ub
 
         self.fp8_linear.process_weights_after_loading(layer)
 
diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py
index 0c4fe9889268..e63efb792fbe 100644
--- a/vllm/model_executor/layers/quantization/fp8.py
+++ b/vllm/model_executor/layers/quantization/fp8.py
@@ -406,6 +406,8 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
 
             # Update layer with new values.
             replace_parameter(layer, "weight", weight.data)
+            # The helper returns (K, N); the param keeps its checkpoint dims.
+            layer.processed_weight_dims = (0, 1)
             replace_parameter(layer, "weight_scale", weight_scale.data)
 
         if input_scale is not None:
diff --git a/vllm/model_executor/layers/quantization/humming.py b/vllm/model_executor/layers/quantization/humming.py
index d0c225cebceb..ea4c9f6c475b 100644
--- a/vllm/model_executor/layers/quantization/humming.py
+++ b/vllm/model_executor/layers/quantization/humming.py
@@ -41,6 +41,11 @@
     select_humming_moe_experts,
     weight_schema_to_quant_key,
 )
+from vllm.model_executor.layers.quantization.utils.humming.linear import (
+    get_live_param,
+    kernel_parameters,
+    replace_layer_parameters,
+)
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 from vllm.model_executor.parameter import (
     BasevLLMParameter,
@@ -542,21 +547,14 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
                 param_dtype=layer.param_dtype,
             )
             tensors.update(input_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)
-
+            replace_layer_parameters(layer, tensors, list(layer._parameters))
             del tensors
 
         # force requant (origin quant setting -> fp16/bf16 -> new_quant setting)
         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()
+            source_tensors = kernel_parameters(layer)
             tensors = self.weight_schema.requant_tensors(
                 tensors=source_tensors,
                 target_weight_schema=self.force_weight_schema,
@@ -568,14 +566,8 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
 
             self.weight_schema = self.force_weight_schema
 
-            for name, _ in list(layer.named_parameters()):
-                if name != "bias":
-                    delattr(layer, name)
-
-            for name, tensor in tensors.items():
-                param = torch.nn.Parameter(tensor, requires_grad=False)
-                setattr(layer, name, param)
-
+            names = [name for name in layer._parameters if name != "bias"]
+            replace_layer_parameters(layer, tensors, names)
             del tensors
 
         self.input_schema = check_and_fallback_input_schema(
@@ -584,7 +576,7 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
             param_dtype=layer.param_dtype,
             allow_fallback=self.quant_config.allow_input_schema_fallback,
         )
-        source_tensors = dict(layer.named_parameters())
+        source_tensors = kernel_parameters(layer)
         self.layer_config, tensors = _hm.prepare_and_transform_humming_tensors(
             tensors=source_tensors,
             shape_n=layer.output_partition_sizes_sum,
@@ -603,11 +595,7 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
         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)
+        replace_layer_parameters(layer, tensors, list(layer._parameters))
 
         self.compute_config = get_humming_linear_compute_config()
         self.locks = _hm.create_humming_locks(layer.weight.device)
@@ -623,12 +611,12 @@ def apply(
             self.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" 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),
+            weight_scale=get_live_param(layer, "weight_scale"),
+            zero_point=get_live_param(layer, "zero_point"),
+            bias=get_live_param(layer, "bias" if bias is not None else "zero_bias"),
+            weight_scale_2=get_live_param(layer, "weight_scale_2"),
+            input_scale=get_live_param(layer, "input_scale"),
+            input_scale_2=get_live_param(layer, "input_scale_2"),
             hadamard_block_size=self.weight_schema.hadamard_block_size,
             locks=self.locks,
             compute_config=self.compute_config,
diff --git a/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_linear.py b/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_linear.py
index a4902fae60c3..06143e73756f 100644
--- a/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_linear.py
+++ b/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_linear.py
@@ -4,7 +4,6 @@
 from typing import TYPE_CHECKING, Any
 
 import torch
-from torch.nn.parameter import Parameter
 
 from vllm.model_executor.kernels.linear import init_mxfp4_linear_kernel
 from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp4Dynamic
@@ -12,6 +11,7 @@
     GroupQuantScaleParameter,
     ModelWeightParameter,
 )
+from vllm.model_executor.utils import move_parameter
 
 from .inc_scheme import INCLinearScheme
 
@@ -84,8 +84,7 @@ def create_weights(
         layer.register_parameter("weight_scale", weight_scale)
 
     def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
-        layer.weight = Parameter(layer.weight_packed.data, requires_grad=False)
-        del layer.weight_packed
+        move_parameter(layer, "weight_packed", "weight")
         self.kernel.process_weights_after_loading(layer)
 
     def apply_weights(
diff --git a/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_moe.py b/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_moe.py
index e51a51aeaa76..ae5fd4fe6b55 100644
--- a/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_moe.py
+++ b/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_moe.py
@@ -35,7 +35,11 @@
 from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
     prepare_moe_fp4_layer_for_marlin,
 )
-from vllm.model_executor.utils import replace_parameter, set_weight_attrs
+from vllm.model_executor.utils import (
+    move_parameter,
+    replace_parameter,
+    set_weight_attrs,
+)
 from vllm.platforms import current_platform
 
 logger = init_logger(__name__)
@@ -165,14 +169,8 @@ def get_fused_moe_quant_config(
         )
 
     def process_weights_after_loading(self, layer: RoutedExperts) -> None:
-        layer.w13_weight = torch.nn.Parameter(
-            layer.w13_weight_packed.data, requires_grad=False
-        )
-        delattr(layer, "w13_weight_packed")
-        layer.w2_weight = torch.nn.Parameter(
-            layer.w2_weight_packed.data, requires_grad=False
-        )
-        delattr(layer, "w2_weight_packed")
+        move_parameter(layer, "w13_weight_packed", "w13_weight")
+        move_parameter(layer, "w2_weight_packed", "w2_weight")
 
         if self.use_cutlass_mxfp4:
             # Swizzle weight scales from flat checkpoint layout [E, N, K//32]
diff --git a/vllm/model_executor/layers/quantization/inc/schemes/inc_wna16_linear.py b/vllm/model_executor/layers/quantization/inc/schemes/inc_wna16_linear.py
index 02a14f51e791..cfd5fdad4581 100644
--- a/vllm/model_executor/layers/quantization/inc/schemes/inc_wna16_linear.py
+++ b/vllm/model_executor/layers/quantization/inc/schemes/inc_wna16_linear.py
@@ -14,7 +14,7 @@
     GroupQuantScaleParameter,
     PackedvLLMParameter,
 )
-from vllm.model_executor.utils import replace_parameter
+from vllm.model_executor.utils import replace_parameter, retire_parameter
 from vllm.scalar_type import scalar_types
 
 from .inc_scheme import INCLinearScheme
@@ -434,9 +434,10 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
         layer.ark_weight_type = ark_linear.wdt
         layer.ark_scale_type = ark_linear.sdt
 
-        if hasattr(layer, "qzeros"):
-            del layer.qzeros
-        del layer.scales
+        # ARK holds its own qzeros/scales; the loaded params keep their data.
+        if getattr(layer, "qzeros", None) is not None:
+            retire_parameter(layer, "qzeros")
+        retire_parameter(layer, "scales")
 
     def apply_weights(
         self,
diff --git a/vllm/model_executor/layers/quantization/kv_cache.py b/vllm/model_executor/layers/quantization/kv_cache.py
index ba503a93f1a5..5428ce17ca86 100644
--- a/vllm/model_executor/layers/quantization/kv_cache.py
+++ b/vllm/model_executor/layers/quantization/kv_cache.py
@@ -77,11 +77,8 @@ def apply(self, layer: torch.nn.Module) -> torch.Tensor:
 
     def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
         if is_weights_pre_processed():
-            # Scale buffers were exported with post-processed values; drop the
-            # consumed scale placeholders and rebuild the host-side copies.
-            for name in ("k_scale", "v_scale", "q_scale", "prob_scale"):
-                if hasattr(layer, name):
-                    delattr(layer, name)
+            # Scale buffers were exported with post-processed values; rebuild the
+            # host-side copies.
             if hasattr(layer, "_k_scale"):
                 layer._k_scale_float = layer._k_scale.item()
                 layer._v_scale_float = layer._v_scale.item()
@@ -90,7 +87,7 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
                 layer._v_scale_cpu = layer._v_scale.detach().to("cpu")
             return
 
-        # skip if there are no weights to process (for example, weight reloading)
+        # skip if there are no weights to process
         if not hasattr(layer, "q_scale"):
             assert not hasattr(layer, "k_scale")
             assert not hasattr(layer, "v_scale")
@@ -104,10 +101,6 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
             layer._v_scale.copy_(1.0)
             layer._k_scale_float = 1.0
             layer._v_scale_float = 1.0
-            del layer.k_scale
-            del layer.v_scale
-            del layer.q_scale
-            del layer.prob_scale
             return
 
         # If the kv-cache is not quantized, we enforce the k/v_scale to be 1.0
@@ -171,13 +164,13 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
         if layer.q_scale > 0.0:
             q_scale = layer.q_scale
             if current_platform.is_fp8_fnuz():
-                q_scale *= 2
+                q_scale = q_scale * 2
         else:
             q_scale = 1.0
         if layer.prob_scale > 0.0:
             prob_scale = layer.prob_scale
             if current_platform.is_fp8_fnuz():
-                prob_scale *= 2
+                prob_scale = prob_scale * 2
         else:
             prob_scale = 1.0
 
@@ -206,8 +199,3 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
                 "issues. Please make sure q/prob scaling factors are "
                 "available in the fp8 checkpoint."
             )
-
-        del layer.k_scale
-        del layer.v_scale
-        del layer.q_scale
-        del layer.prob_scale
diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py
index 29db26053521..a01c5b71c15e 100644
--- a/vllm/model_executor/layers/quantization/modelopt.py
+++ b/vllm/model_executor/layers/quantization/modelopt.py
@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-from collections.abc import Callable
+from collections.abc import Callable, Iterator
 from dataclasses import dataclass
 from enum import Enum
 from fnmatch import fnmatch
@@ -115,6 +115,7 @@
 from vllm.model_executor.utils import (
     is_weights_pre_processed,
     replace_parameter,
+    retire_parameter,
     set_weight_attrs,
 )
 from vllm.platforms import current_platform
@@ -203,7 +204,7 @@ def is_layer_excluded(self, prefix: str) -> bool:
             return True
 
         # TODO: This special hard coded logic is not needed for quantized checkpoints
-        # generated by ModelOpt >= 0.39.0 where they are handled natually by the
+        # generated by ModelOpt >= 0.39.0 where they are handled naturally by the
         # exclude_modules config. But need to keep them for loading quantized
         # checkpoints generated by older versions. Then check substring matching
         # for patterns not caught by exact match
@@ -239,7 +240,7 @@ def get_quant_method(
             return None
 
         # TODO: This special hard coded logic is not needed for quantized checkpoints
-        # generated by ModelOpt >= 0.39.0 where they are handled natually by the
+        # generated by ModelOpt >= 0.39.0 where they are handled naturally by the
         # exclude_modules config. But need to keep them for loading quantized
         # checkpoints generated by older versions. Then check substring matching
         # for patterns not caught by exact match
@@ -855,6 +856,8 @@ def __init__(
         super().__init__(moe_config)
         self.quant_config = quant_config
         self.use_a16 = quant_config.quant_method == "W4A16_NVFP4"
+        # Whether the selected backend reads the activation input scales.
+        self.uses_input_scales = True
         activation_key = None if self.use_a16 else kNvfp4Dynamic
 
         if quant_config.moe_activation_override is not None:
@@ -895,8 +898,8 @@ def uses_weight_scale_2_pattern(self) -> bool:
 
     def dequantize_weight(
         self, layer: torch.nn.Module
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        """Materialize serialized ModelOpt NVFP4 expert weights."""
+    ) -> Iterator[tuple[slice, torch.Tensor, torch.Tensor]]:
+        """Yield ``(experts, w13, w2)`` dequantized ModelOpt NVFP4 experts."""
         return dequantize_nvfp4_moe_weights(
             layer.w13_weight,
             layer.w13_weight_scale,
@@ -1093,11 +1096,18 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None:
         replace_parameter(layer, "w13_weight", w13)
         replace_parameter(layer, "w13_weight_scale", w13_scale)
         replace_parameter(layer, "w13_weight_scale_2", w13_scale_2)
-        replace_parameter(layer, "w13_input_scale", a13_scale)
         replace_parameter(layer, "w2_weight", w2)
         replace_parameter(layer, "w2_weight_scale", w2_scale)
         replace_parameter(layer, "w2_weight_scale_2", w2_scale_2)
-        replace_parameter(layer, "w2_input_scale", a2_scale)
+        # Weight-only backends (Marlin, Humming) return no input scales: the
+        # loaded ones keep their data, retired, and the kernel config gets None.
+        self.uses_input_scales = a13_scale is not None
+        if self.uses_input_scales:
+            replace_parameter(layer, "w13_input_scale", a13_scale)
+            replace_parameter(layer, "w2_input_scale", a2_scale)
+        else:
+            retire_parameter(layer, "w13_input_scale")
+            retire_parameter(layer, "w2_input_scale")
 
         self._build_moe_kernel(layer)
 
@@ -1132,8 +1142,8 @@ def get_fused_moe_quant_config(self, layer: RoutedExperts) -> FusedMoEQuantConfi
             w2_scale=layer.w2_weight_scale,
             w13_scale_2=layer.w13_weight_scale_2,
             w2_scale_2=layer.w2_weight_scale_2,
-            a13_scale=layer.w13_input_scale,
-            a2_scale=layer.w2_input_scale,
+            a13_scale=layer.w13_input_scale if self.uses_input_scales else None,
+            a2_scale=layer.w2_input_scale if self.uses_input_scales else None,
             swiglu_limit=getattr(layer, "swiglu_limit", None),
             swiglu_alpha=getattr(layer, "swiglu_alpha", None),
             swiglu_beta=getattr(layer, "swiglu_beta", None),
@@ -1208,10 +1218,14 @@ def __init__(
         is_checkpoint_mxfp8_serialized: bool,
         kv_cache_quant_algo: str | None,
         exclude_modules: list[str],
+        moe_quant_algo: str | None = None,
     ) -> None:
         super().__init__(exclude_modules)
         self.quant_method = "MXFP8"
         self.is_checkpoint_mxfp8_serialized = is_checkpoint_mxfp8_serialized
+        # MiniMax-M3.1 checkpoints keep MXFP8 linears but store routed experts
+        # as compressed-tensors NVFP4 (`moe_quant_algo: NVFP4`).
+        self.moe_quant_algo = moe_quant_algo
 
         if not is_checkpoint_mxfp8_serialized:
             raise ValueError(
@@ -1247,12 +1261,25 @@ def from_config(cls, config: dict[str, Any]) -> "ModelOptMxFp8Config":
         # (same on-disk format as ModelOpt MXFP8); normalize to the ModelOpt
         # schema and reuse the shared parser.
         if "quantization" not in config and not config.get("quant_algo"):
+            moe_quant_algo = (config.get("moe_quant_algo") or "").upper() or None
+            if moe_quant_algo is not None and (
+                moe_quant_algo != "NVFP4"
+                or config.get("moe_quant_format") != "nvfp4-pack-quantized"
+                or config.get("group_size") != 16
+            ):
+                raise ValueError(
+                    "Unsupported MXFP8 checkpoint MoE quantization: "
+                    f"moe_quant_algo={config.get('moe_quant_algo')!r}, "
+                    f"moe_quant_format={config.get('moe_quant_format')!r}, "
+                    f"group_size={config.get('group_size')!r}."
+                )
             config = {
                 "quant_method": "modelopt",
                 "quantization": {
                     "quant_algo": "MXFP8",
                     "kv_cache_quant_algo": config.get("kv_cache_quant_algo"),
                     "exclude_modules": config.get("ignored_layers", []) or [],
+                    "moe_quant_algo": moe_quant_algo,
                 },
             }
         return cast("ModelOptMxFp8Config", super().from_config(config))
@@ -1286,8 +1313,41 @@ def _from_config(
             is_checkpoint_mxfp8_serialized,
             kv_cache_quant_method,
             exclude_modules,
+            moe_quant_algo=original_config.get("quantization", {}).get(
+                "moe_quant_algo"
+            ),
         )
 
+    def routed_expert_weight_key(self) -> QuantKey:
+        """The weight format the routed experts are stored in.
+
+        MXFP8, or NVFP4 (``moe_quant_algo: NVFP4``).
+        """
+        if self.moe_quant_algo != "NVFP4":
+            return kMxfp8Static
+        return kNvfp4Static
+
+    def get_quant_method(
+        self, layer: torch.nn.Module, prefix: str
+    ) -> "QuantizeMethodBase | None":
+        if (
+            self.moe_quant_algo == "NVFP4"
+            and isinstance(layer, RoutedExperts)
+            and not self.is_layer_excluded(prefix)
+        ):
+            from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_w4a4_nvfp4 import (  # noqa: E501
+                CompressedTensorsW4A4Nvfp4MoEMethod,
+            )
+
+            # Nothing in the checkpoint says how its FC1 input was
+            # quantized during QAT, so the model declares it.
+            return CompressedTensorsW4A4Nvfp4MoEMethod(
+                layer.moe_config,
+                layer_name=prefix,
+                activation_key=(layer.moe_config.qat_activation_key or kNvfp4Dynamic),
+            )
+        return super().get_quant_method(layer, prefix)
+
 
 class ModelOptMxFp8FusedMoE(FusedMoEMethodBase):
     """FlashInfer TRTLLM MXFP8 block-scale MoE for ModelOpt checkpoints."""
@@ -2092,7 +2152,8 @@ def process(self, layer, role) -> None:
         # Raw max, no reciprocation — Marlin/cutlass want ModelOpt's amax/2688.
         weight_global_scale = layer.weight_scale_2.max().to(torch.float32)
         replace_parameter(layer, "weight_global_scale", weight_global_scale)
-        del layer.weight_scale_2
+        # Kernels read weight_global_scale; the loaded param keeps its data.
+        retire_parameter(layer, "weight_scale_2")
 
 
 class KNvfp4Dynamic(QuantKeyScheme):
@@ -2130,7 +2191,8 @@ def process(self, layer, role) -> None:
             "input_global_scale_inv",
             (1.0 / layer.input_global_scale).to(torch.float32),
         )
-        del layer.input_scale
+        # Kernels read input_global_scale; the loaded param keeps its data.
+        retire_parameter(layer, "input_scale")
 
 
 class KFp8StaticTensor(QuantKeyScheme):
@@ -2550,7 +2612,8 @@ def post_process(self, layer) -> None:
                 "Dropping input_scale from a weight-only W4A16 NVFP4 "
                 "checkpoint (deprecated)."
             )
-        del layer.input_scale
+        # Weight-only: no kernel reads it; the loaded param keeps its data.
+        retire_parameter(layer, "input_scale")
 
 
 _DROP_INPUT_SCALE = _DropInputScale()
diff --git a/vllm/model_executor/layers/quantization/mxfp4.py b/vllm/model_executor/layers/quantization/mxfp4.py
index c93977d86c65..6288a47bcde0 100644
--- a/vllm/model_executor/layers/quantization/mxfp4.py
+++ b/vllm/model_executor/layers/quantization/mxfp4.py
@@ -31,6 +31,7 @@
     mxfp4_round_up_hidden_size_and_intermediate_size,
     select_deepseek_v4_mxfp4_moe_backend,
     select_mxfp4_moe_backend,
+    triton_weight,
 )
 from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
 from vllm.model_executor.layers.quantization import QuantizationMethods
@@ -366,8 +367,8 @@ def _setup_kernel(
             replace_parameter(layer, "w13_weight_scale", w13_scale)
             replace_parameter(layer, "w2_weight_scale", w2_scale)
         else:
-            layer.w13_weight = w13
-            layer.w2_weight = w2
+            layer.w13_weight_triton = w13
+            layer.w2_weight_triton = w2
             self.w13_precision_config = w13_scale
             self.w2_precision_config = w2_scale
 
@@ -444,8 +445,8 @@ def apply(
         assert self.moe_kernel is not None
         return self.moe_kernel.apply(
             hidden_states=x,
-            w1=layer.w13_weight,
-            w2=layer.w2_weight,
+            w1=triton_weight(layer, "w13_weight"),
+            w2=triton_weight(layer, "w2_weight"),
             topk_weights=topk_weights,
             topk_ids=topk_ids,
             activation=layer.activation,
@@ -467,8 +468,8 @@ def apply_monolithic(
         assert self.moe_kernel is not None
         return self.moe_kernel.apply_monolithic(
             hidden_states=x,
-            w1=layer.w13_weight,
-            w2=layer.w2_weight,
+            w1=triton_weight(layer, "w13_weight"),
+            w2=triton_weight(layer, "w2_weight"),
             router_logits=router_logits,
             activation=layer.activation,
             global_num_experts=layer.global_num_experts,
@@ -503,6 +504,108 @@ def _use_k3_situ_aiter(moe: FusedMoEConfig) -> bool:
     )
 
 
+# The gfx942 int4 requant unpacks MXFP4 to f32 and peaks above 20 GiB per rank
+# if all experts convert at once, so it converts in chunks. Chunks are sized by
+# unpacked values rather than experts so that pure TP, which holds every expert
+# at 1/tp width, runs as many chunks as EP.
+_K3_INT4_CONVERT_CHUNK_VALUES = 1 << 27
+
+
+def _mxfp4_experts_to_int4_gfx942(
+    weight: torch.nn.Parameter,
+    scale: torch.nn.Parameter,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Requantize packed MXFP4 experts to shuffled groupwise int4 for gfx942.
+
+    Takes ``(E, N, K // 2)`` MXFP4 weights with ``(E, N, K // 32)`` E8M0
+    scales and returns ``(E, N, K // 2)`` packed int4 with flattened
+    per-32 scales.
+    """
+    from aiter import dtypes as aiter_dtypes
+    from aiter.ops.quant import per_1x32_i4_quant
+    from aiter.ops.shuffle import (
+        pack_int8_to_packed_int4,
+        shuffle_scale_for_int4,
+        shuffle_weight,
+    )
+    from aiter.utility import fp4_utils
+
+    fp4_dtype = torch.float4_e2m1fn_x2
+    e8m0_dtype = torch.float8_e8m0fnu
+
+    w_all = weight.data.view(fp4_dtype)
+    s_all = scale.data.view(e8m0_dtype)
+    num_experts, rows, packed_cols = w_all.shape
+    chunk = min(
+        num_experts,
+        max(1, _K3_INT4_CONVERT_CHUNK_VALUES // (rows * packed_cols * 2)),
+    )
+
+    # Write chunks into a preallocated output. torch.cat of a list
+    # would hold the full result twice at the join.
+    out_packed: torch.Tensor | None = None
+    out_scale: torch.Tensor | None = None
+    scale_stride = 0
+    for lo in range(0, num_experts, chunk):
+        # Shift the last chunk back so every chunk has the same shapes. Under
+        # load_model's max_split_size_mb, a shorter chunk cannot reuse the
+        # cached blocks. The overlapping experts are recomputed identically.
+        lo = min(lo, num_experts - chunk)
+        hi = lo + chunk
+        weight_f32 = fp4_utils.mxfp4_to_f32(w_all[lo:hi])
+        scale_f32 = fp4_utils.e8m0_to_f32(s_all[lo:hi])
+        chunk_experts, output_size, input_size = weight_f32.shape
+        weight_bf16 = (
+            (
+                weight_f32.view(chunk_experts, output_size, input_size // 32, 32)
+                * scale_f32.view(chunk_experts, output_size, input_size // 32, 1)
+            )
+            .view(chunk_experts, output_size, input_size)
+            .to(torch.bfloat16)
+        )
+        del weight_f32, scale_f32
+
+        weight_int4, weight_scale = per_1x32_i4_quant(weight_bf16)
+        del weight_bf16
+        weight_int4 = weight_int4.view(aiter_dtypes.i4x2).view(
+            chunk_experts, output_size, input_size
+        )
+        weight_packed = pack_int8_to_packed_int4(
+            shuffle_weight(weight_int4.view(aiter_dtypes.i8), (16, 16))
+        )
+        del weight_int4
+        chunk_packed = weight_packed.view(
+            chunk_experts, output_size, input_size // 2
+        ).view(aiter_dtypes.i4x2)
+        chunk_scale = (
+            shuffle_scale_for_int4(weight_scale, group_size=32).view(-1).contiguous()
+        )
+        if out_packed is None:
+            out_packed = torch.empty(
+                (num_experts,) + tuple(chunk_packed.shape[1:]),
+                dtype=chunk_packed.dtype,
+                device=chunk_packed.device,
+            )
+            scale_stride = chunk_scale.numel() // chunk_experts
+            out_scale = torch.empty(
+                num_experts * scale_stride,
+                dtype=chunk_scale.dtype,
+                device=chunk_scale.device,
+            )
+        assert out_packed is not None
+        assert out_scale is not None
+        out_packed[lo:hi].copy_(chunk_packed)
+        out_scale[lo * scale_stride : hi * scale_stride].copy_(chunk_scale)
+        # No empty_cache() here: the next chunk reuses these cached
+        # blocks, while releasing them costs a hipFree per temporary.
+        del weight_packed, weight_scale, chunk_packed, chunk_scale
+
+    del w_all, s_all
+    assert out_packed is not None
+    assert out_scale is not None
+    return out_packed, out_scale
+
+
 def _use_k3_situ_int4_gfx942(moe: FusedMoEConfig) -> bool:
     """Opt-in gfx942 MXFP4-to-int4 requant for Kimi-K3 SiTU experts."""
     from vllm.platforms import current_platform
@@ -851,8 +954,8 @@ def _setup_kernel(
             replace_parameter(layer, "w13_weight_scale", w13_scale)
             replace_parameter(layer, "w2_weight_scale", w2_scale)
         else:
-            layer.w13_weight = w13
-            layer.w2_weight = w2
+            layer.w13_weight_triton = w13
+            layer.w2_weight_triton = w2
             self.w13_precision_config = w13_scale
             self.w2_precision_config = w2_scale
 
@@ -891,109 +994,17 @@ def _setup_kernel_k3_situ_gfx942(self, layer: RoutedExperts) -> None:
                 "SiLU. Rebuild with the replacement A16W4/int4 pipeline."
             )
 
-        from aiter import dtypes as aiter_dtypes
-        from aiter.ops.quant import per_1x32_i4_quant
-        from aiter.ops.shuffle import (
-            pack_int8_to_packed_int4,
-            shuffle_scale_for_int4,
-            shuffle_weight,
+        w13, w13_scale = _mxfp4_experts_to_int4_gfx942(
+            layer.w13_weight, layer.w13_weight_scale
         )
-        from aiter.utility import fp4_utils
-
-        fp4_dtype = torch.float4_e2m1fn_x2
-        e8m0_dtype = torch.float8_e8m0fnu
-
-        # mxfp4_to_f32 grows the packed weight 8x and peaks above 20 GiB per
-        # rank if all experts convert at once. Convert 8 experts at a time.
-        _CONVERT_CHUNK = 8
-
-        def convert(
-            weight: torch.nn.Parameter,
-            scale: torch.nn.Parameter,
-        ) -> tuple[torch.Tensor, torch.Tensor]:
-            # Drop the source before returning. Packed int4 is the same size as
-            # packed MXFP4, so holding both doubles the footprint. That fits
-            # under expert parallel (1/ep_size experts per rank) and does not
-            # fit under pure tensor parallel.
-            w_all = weight.data.view(fp4_dtype)
-            s_all = scale.data.view(e8m0_dtype)
-            num_experts = w_all.shape[0]
-
-            # Write chunks into a preallocated output. torch.cat of a list
-            # would hold the full result twice at the join.
-            out_packed: torch.Tensor | None = None
-            out_scale: torch.Tensor | None = None
-            scale_stride = 0
-            for lo in range(0, num_experts, _CONVERT_CHUNK):
-                hi = min(lo + _CONVERT_CHUNK, num_experts)
-                weight_f32 = fp4_utils.mxfp4_to_f32(w_all[lo:hi])
-                scale_f32 = fp4_utils.e8m0_to_f32(s_all[lo:hi])
-                chunk_experts, output_size, input_size = weight_f32.shape
-                weight_bf16 = (
-                    (
-                        weight_f32.view(
-                            chunk_experts, output_size, input_size // 32, 32
-                        )
-                        * scale_f32.view(
-                            chunk_experts, output_size, input_size // 32, 1
-                        )
-                    )
-                    .view(chunk_experts, output_size, input_size)
-                    .to(torch.bfloat16)
-                )
-                del weight_f32, scale_f32
-
-                weight_int4, weight_scale = per_1x32_i4_quant(weight_bf16)
-                del weight_bf16
-                weight_int4 = weight_int4.view(aiter_dtypes.i4x2).view(
-                    chunk_experts, output_size, input_size
-                )
-                weight_packed = pack_int8_to_packed_int4(
-                    shuffle_weight(weight_int4.view(aiter_dtypes.i8), (16, 16))
-                )
-                del weight_int4
-                chunk_packed = weight_packed.view(
-                    chunk_experts, output_size, input_size // 2
-                ).view(aiter_dtypes.i4x2)
-                chunk_scale = (
-                    shuffle_scale_for_int4(weight_scale, group_size=32)
-                    .view(-1)
-                    .contiguous()
-                )
-                if out_packed is None:
-                    out_packed = torch.empty(
-                        (num_experts,) + tuple(chunk_packed.shape[1:]),
-                        dtype=chunk_packed.dtype,
-                        device=chunk_packed.device,
-                    )
-                    scale_stride = chunk_scale.numel() // chunk_experts
-                    out_scale = torch.empty(
-                        num_experts * scale_stride,
-                        dtype=chunk_scale.dtype,
-                        device=chunk_scale.device,
-                    )
-                assert out_packed is not None
-                assert out_scale is not None
-                out_packed[lo:hi].copy_(chunk_packed)
-                out_scale[lo * scale_stride : hi * scale_stride].copy_(chunk_scale)
-                del weight_packed, weight_scale, chunk_packed, chunk_scale
-                torch.accelerator.empty_cache()
-
-            del w_all, s_all
-            assert out_packed is not None
-            assert out_scale is not None
-            weight.data = torch.empty(0, dtype=torch.uint8, device=out_packed.device)
-            scale.data = torch.empty(0, dtype=torch.uint8, device=out_packed.device)
-            torch.accelerator.empty_cache()
-            return out_packed, out_scale
-
-        w13, w13_scale = convert(layer.w13_weight, layer.w13_weight_scale)
         replace_parameter(layer, "w13_weight", w13)
         replace_parameter(layer, "w13_weight_scale", w13_scale)
         del w13, w13_scale
         torch.accelerator.empty_cache()
 
-        w2, w2_scale = convert(layer.w2_weight, layer.w2_weight_scale)
+        w2, w2_scale = _mxfp4_experts_to_int4_gfx942(
+            layer.w2_weight, layer.w2_weight_scale
+        )
         replace_parameter(layer, "w2_weight", w2)
         replace_parameter(layer, "w2_weight_scale", w2_scale)
         del w2, w2_scale
@@ -1099,8 +1110,8 @@ def apply(
         assert self.moe_kernel is not None
         return self.moe_kernel.apply(
             hidden_states=x,
-            w1=layer.w13_weight,
-            w2=layer.w2_weight,
+            w1=triton_weight(layer, "w13_weight"),
+            w2=triton_weight(layer, "w2_weight"),
             topk_weights=topk_weights,
             topk_ids=topk_ids,
             activation=layer.activation,
@@ -1122,8 +1133,8 @@ def apply_monolithic(
         assert self.moe_kernel is not None
         return self.moe_kernel.apply_monolithic(
             hidden_states=x,
-            w1=layer.w13_weight,
-            w2=layer.w2_weight,
+            w1=triton_weight(layer, "w13_weight"),
+            w2=triton_weight(layer, "w2_weight"),
             router_logits=router_logits,
             activation=layer.activation,
             global_num_experts=layer.global_num_experts,
diff --git a/vllm/model_executor/layers/quantization/online/base.py b/vllm/model_executor/layers/quantization/online/base.py
index 511714837faa..9ac99e6c04f7 100644
--- a/vllm/model_executor/layers/quantization/online/base.py
+++ b/vllm/model_executor/layers/quantization/online/base.py
@@ -61,6 +61,9 @@
     get_layer_name_after_index,
 )
 from vllm.model_executor.layers.quantization.utils.mxfp4_utils import mxfp4_quantize
+from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
+    dequant_mxfp8_to_bf16,
+)
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
     QuantKey,
     kFp8Static128BlockSym,
@@ -70,6 +73,7 @@
     kInt8StaticChannelSym,
     kMxfp4Static,
     kMxfp8Dynamic,
+    kMxfp8Static,
     kNvfp4Static,
 )
 
@@ -114,6 +118,14 @@ class OnlineQuantizationSource(str, Enum):
     kMxfp4Static: mxfp4_quantize,
 }
 
+# Dequantizers `(weight, weight_scale) -> high precision weight` of quantized fused
+# shared-expert checkpoint weights, keyed by the checkpoint weight QuantKey.
+ONLINE_SHARED_EXPERT_DEQUANTIZERS: dict[
+    QuantKey, Callable[[torch.Tensor, torch.Tensor], torch.Tensor]
+] = {
+    kMxfp8Static: dequant_mxfp8_to_bf16,
+}
+
 
 def _find_matching_targets(
     prefix: str,
diff --git a/vllm/model_executor/layers/quantization/online/fp8.py b/vllm/model_executor/layers/quantization/online/fp8.py
index f4a327192108..c4a85e6b0037 100644
--- a/vllm/model_executor/layers/quantization/online/fp8.py
+++ b/vllm/model_executor/layers/quantization/online/fp8.py
@@ -22,6 +22,9 @@
 from vllm.model_executor.kernels.linear.scaled_mm import (
     MarlinFP8ScaledMMLinearKernel,
 )
+from vllm.model_executor.kernels.linear.scaled_mm.humming import (
+    HummingFP8ScaledMMLinearKernel,
+)
 from vllm.model_executor.layers.fused_moe import RoutedExperts
 from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig
 from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
@@ -249,6 +252,8 @@ def process_weights_after_loading(self, layer: Module) -> None:
 
         # Update layer with new values.
         replace_parameter(layer, "weight", qweight.t().data)
+        # Transposed to (K, N); the param keeps its checkpoint dims.
+        layer.processed_weight_dims = (0, 1)
         replace_parameter(layer, "weight_scale", weight_scale.data)
 
         if self.use_marlin and hasattr(self.fp8_linear, "marlin_input_dtype"):
@@ -415,13 +420,17 @@ def create_weights(
             out_dtype=self.out_dtype,
             module_name=self.__class__.__name__,
         )
-        # PTPC requires per-token activation FP8; MarlinFP8 is W8A16 and
-        # would silently produce a weight-only fp8 model.
-        if isinstance(self.fp8_linear, MarlinFP8ScaledMMLinearKernel):
+        # PTPC requires per-token activation FP8; Marlin and Humming FP8 are
+        # W8A16 and would silently produce a weight-only fp8 model.
+        if isinstance(
+            self.fp8_linear,
+            (MarlinFP8ScaledMMLinearKernel, HummingFP8ScaledMMLinearKernel),
+        ):
             raise ValueError(
                 "FP8 PTPC online quant requires a kernel that honors "
-                "per-token activation quantization; MarlinFP8 is W8A16 "
-                "weight-only. Requires SM89+ for Cutlass FP8 or ROCm MI3xx "
+                "per-token activation quantization; "
+                f"{type(self.fp8_linear).__name__} is W8A16 weight-only. "
+                "Requires SM89+ for Cutlass FP8 or ROCm MI3xx "
                 "for rowwise scaled_mm."
             )
 
@@ -817,26 +826,43 @@ def process_weights_after_loading(self, layer: Module) -> None:
         fp8_dtype = current_platform.fp8_dtype()
         w13 = torch.empty_like(w13_weight, dtype=fp8_dtype)
         w2 = torch.empty_like(w2_weight, dtype=fp8_dtype)
-        # Scale's leading dim is taken from the fp8 weight tensor by
-        # construction, so it cannot drift from the weight's expert count
-        # under EP / padded MoE.
-        n_w13 = w13_weight.shape[1]
-        w13_scale = torch.ones(
-            w13.shape[0], n_w13, 1, device=w13.device, dtype=torch.float32
-        )
         layer.w13_input_scale = None
         layer.w2_input_scale = None
 
+        use_torch_w13_quant = False
+        if current_platform.is_rocm():
+            from vllm.platforms.rocm import on_gfx942
+
+            use_torch_w13_quant = on_gfx942()
+
+        if use_torch_w13_quant:
+            # dynamic_per_token_scaled_fp8_quant can hard-crash on gfx942.
+            w13_amax = weight_amax(w13_weight, dim=-1, keepdim=True)
+            w13_scale = _fp8_channel_scale(w13_amax)
+        else:
+            # Scale's leading dim is taken from the fp8 weight tensor by
+            # construction, so it cannot drift from the weight's expert count
+            # under EP / padded MoE.
+            n_w13 = w13_weight.shape[1]
+            w13_scale = torch.ones(
+                w13.shape[0], n_w13, 1, device=w13.device, dtype=torch.float32
+            )
+
         w2_amax = weight_amax(w2_weight, dim=-1, keepdim=True)
         w2_amax = amax_for_moe_weight_quant(w2_amax, self.moe.tp_size)
         w2_scale = _fp8_channel_scale(w2_amax)
 
         for expert in range(layer.local_num_experts):
-            w13[expert], w13_scale[expert] = ops.scaled_fp8_quant(
-                w13_weight[expert],
-                scale=None,
-                use_per_token_if_dynamic=True,
-            )
+            if use_torch_w13_quant:
+                w13[expert] = _fp8_quant_per_channel(
+                    w13_weight[expert], w13_scale[expert]
+                )
+            else:
+                w13[expert], w13_scale[expert] = ops.scaled_fp8_quant(
+                    w13_weight[expert],
+                    scale=None,
+                    use_per_token_if_dynamic=True,
+                )
             w2[expert] = _fp8_quant_per_channel(w2_weight[expert], w2_scale[expert])
 
         self._setup_kernel(
diff --git a/vllm/model_executor/layers/quantization/online/moe_base.py b/vllm/model_executor/layers/quantization/online/moe_base.py
index eb175ebd33d8..5eb5a6990a73 100644
--- a/vllm/model_executor/layers/quantization/online/moe_base.py
+++ b/vllm/model_executor/layers/quantization/online/moe_base.py
@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 from abc import abstractmethod
+from collections.abc import Callable, Iterator
 
 import torch
 
@@ -12,6 +13,7 @@
 )
 from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig
 from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput
+from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import triton_weight
 from vllm.model_executor.layers.quantization.online.requantization import (
     OnlineRequantizationMixin,
 )
@@ -46,14 +48,17 @@ def create_weights(
         layer.weight_block_size = None
 
         if self.requantization_source is not None:
+            # In case FSE is used, shared experts slots are initialized on RoutedExperts
+            # online quant method, the source only holds the routed experts.
             self.create_requantization_source_weights(
                 layer,
-                num_experts,
+                num_experts - layer.expert_map_manager.num_fused_shared_experts,
                 layer.moe_config.hidden_dim_unpadded,
                 layer.moe_config.intermediate_size_per_partition_unpadded,
                 params_dtype,
                 **extra_weight_attrs,
             )
+            layer.num_experts = num_experts
             return
 
         # Fused gate_up_proj (column parallel) — full precision on meta device
@@ -118,22 +123,95 @@ def create_weights(
     def get_weights_for_quantization(
         self, layer: torch.nn.Module
     ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Return the high precision ``(w13, w2)`` of all local experts."""
         if self.requantization_source is None:
             return layer.w13_weight, layer.w2_weight
+        w13, w2 = self.quantize_weights(layer, lambda w13, w2: (w13, w2))
+        return w13, w2
+
+    def iter_weights_for_quantization(
+        self, layer: torch.nn.Module
+    ) -> Iterator[tuple[slice, torch.Tensor, torch.Tensor]]:
+        """Yield ``(experts, w13, w2)`` for groups of experts to quantize.
+
+        The source dequantizes in chunks to bound staging memory. Without a source,
+        all routed experts are yielded at once. Fused shared experts come last.
+        """
+        num_shared = layer.expert_map_manager.num_fused_shared_experts
+        num_routed = layer.local_num_experts - num_shared
+
+        source = self.requantization_source
+        if source is None:
+            # These include the fused shared experts slots, never staged here.
+            yield slice(0, num_routed), layer.w13_weight, layer.w2_weight
+        else:
+            chunks = source.dequantize_weight(layer)
+            if isinstance(chunks, torch.Tensor):
+                raise TypeError("MoE requantization requires w13 and w2 weights.")
+            for experts, w13, w2 in chunks:
+                yield (
+                    experts,
+                    *self._pad_requantized_weights(
+                        layer, w13, w2, experts.stop - experts.start
+                    ),
+                )
+
+        if num_shared > 0:
+            loader = layer.shared_expert_online_loader
+            shared_w13, shared_w2 = loader.dequantize_staging(layer)
+            loader.release_staging(layer)
+            yield (
+                slice(num_routed, num_routed + num_shared),
+                *self._pad_requantized_weights(
+                    layer,
+                    shared_w13.to(layer.orig_dtype),
+                    shared_w2.to(layer.orig_dtype),
+                    num_shared,
+                ),
+            )
+
+    def quantize_weights(
+        self,
+        layer: torch.nn.Module,
+        quantize: Callable[[torch.Tensor, torch.Tensor], tuple[torch.Tensor, ...]],
+    ) -> tuple[torch.Tensor, ...]:
+        """Quantize the layer's weights one group of experts at a time.
 
-        weights = self.requantization_source.dequantize_weight(layer)
-        if not isinstance(weights, tuple):
-            raise TypeError("MoE requantization requires w13 and w2 weight tensors.")
+        Args:
+            layer: Layer to quantize.
+            quantize: Maps ``(w13, w2)`` of a group of experts to tensors, e.g.
+                ``(w13, w13_scale, w2, w2_scale)``.
 
-        w13, w2 = weights
+        Returns:
+            The ``quantize`` outputs for all local experts.
 
+        """
+        num_experts = layer.local_num_experts
+        outputs: list[torch.Tensor] | None = None
+        for experts, w13, w2 in self.iter_weights_for_quantization(layer):
+            quantized = quantize(w13, w2)
+            if outputs is None:
+                if w13.shape[0] == num_experts:
+                    return quantized
+                outputs = [t.new_empty((num_experts, *t.shape[1:])) for t in quantized]
+            for output, tensor in zip(outputs, quantized):
+                output[experts] = tensor
+        assert outputs is not None
+        return tuple(outputs)
+
+    def _pad_requantized_weights(
+        self,
+        layer: torch.nn.Module,
+        w13: torch.Tensor,
+        w2: torch.Tensor,
+        num_experts: int,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         config = layer.moe_config
         hidden_size = config.hidden_dim_unpadded
         intermediate_size = config.intermediate_size_per_partition_unpadded
         target_hidden_size = config.hidden_dim
         target_intermediate_size = config.intermediate_size_per_partition
 
-        num_experts = layer.local_num_experts
         num_shards = self.moe.w13_num_shards
 
         if w13.shape != (num_experts, num_shards * intermediate_size, hidden_size):
@@ -237,8 +315,8 @@ def apply_monolithic(
         assert self.moe_kernel is not None
         return self.moe_kernel.apply_monolithic(
             x,
-            layer.w13_weight,
-            layer.w2_weight,
+            triton_weight(layer, "w13_weight"),
+            triton_weight(layer, "w2_weight"),
             router_logits,
             activation=layer.activation,
             global_num_experts=layer.global_num_experts,
@@ -264,8 +342,8 @@ def apply(
         assert self.moe_kernel is not None
         return self.moe_kernel.apply(
             x,
-            layer.w13_weight,
-            layer.w2_weight,
+            triton_weight(layer, "w13_weight"),
+            triton_weight(layer, "w2_weight"),
             topk_weights,
             topk_ids,
             activation=layer.activation,
diff --git a/vllm/model_executor/layers/quantization/online/mxfp4.py b/vllm/model_executor/layers/quantization/online/mxfp4.py
index 70a64340f694..828b84fd4057 100644
--- a/vllm/model_executor/layers/quantization/online/mxfp4.py
+++ b/vllm/model_executor/layers/quantization/online/mxfp4.py
@@ -235,8 +235,8 @@ def _setup_kernel(
         ):
             # Triton-based backends: w13/w2 are triton_kernels.tensor.Tensor
             # Store on layer for apply(), scales are PrecisionConfig
-            layer.w13_weight = w13
-            layer.w2_weight = w2
+            layer.w13_weight_triton = w13
+            layer.w2_weight_triton = w2
             self.w13_precision_config = w13_scale
             self.w2_precision_config = w2_scale
         else:
@@ -269,8 +269,8 @@ def get_fused_moe_quant_config(
         # AITER_MXFP4_FP8` to read `self.w13_precision_config`/
         # `self.w2_precision_config` instead of `layer.w13_weight_scale`/
         # `layer.w2_weight_scale`. For those backends,
-        # `convert_weight_to_mxfp4_moe_kernel_format` deletes the latter, so
-        # `getattr` below would raise `AttributeError`. This is currently
+        # `convert_weight_to_mxfp4_moe_kernel_format` empties the latter, so
+        # the scales read below would be empty. This is currently
         # unreachable because `activation_quant_key` is hardcoded to
         # `kMxfp4Dynamic`
         # TODO: When supporting online activation quant key override,
@@ -296,18 +296,23 @@ def process_weights_after_loading(self, layer: Module) -> None:
         if getattr(layer, "_already_called_process_weights_after_loading", False):
             return
 
-        w13_weight, w2_weight = self.get_weights_for_quantization(layer)
-        self._zero_padding(layer, w13_weight, w2_weight)
-
         if self.mxfp4_backend == Mxfp4MoeBackend.NONE:
+            w13_weight, w2_weight = self.get_weights_for_quantization(layer)
+            self._zero_padding(layer, w13_weight, w2_weight)
             layer._already_called_process_weights_after_loading = True
             return
 
         layer.w13_input_scale = None
         layer.w2_input_scale = None
 
-        w13, w13_scale = _quantize_mxfp4_moe_weight(w13_weight)
-        w2, w2_scale = _quantize_mxfp4_moe_weight(w2_weight)
+        def quantize(w13_weight: torch.Tensor, w2_weight: torch.Tensor):
+            self._zero_padding(layer, w13_weight, w2_weight)
+            return (
+                *_quantize_mxfp4_moe_weight(w13_weight),
+                *_quantize_mxfp4_moe_weight(w2_weight),
+            )
+
+        w13, w13_scale, w2, w2_scale = self.quantize_weights(layer, quantize)
 
         self._setup_kernel(
             layer,
diff --git a/vllm/model_executor/layers/quantization/online/mxfp8.py b/vllm/model_executor/layers/quantization/online/mxfp8.py
index f16f2b9fbf24..d405e76c2cde 100644
--- a/vllm/model_executor/layers/quantization/online/mxfp8.py
+++ b/vllm/model_executor/layers/quantization/online/mxfp8.py
@@ -237,12 +237,16 @@ def process_weights_after_loading(self, layer: Module) -> None:
         if getattr(layer, "_already_called_process_weights_after_loading", False):
             return
 
-        w13_weight, w2_weight = self.get_weights_for_quantization(layer)
         layer.w13_input_scale = None
         layer.w2_input_scale = None
 
-        w13, w13_scale = self._quantize_mxfp8_moe_weight(w13_weight)
-        w2, w2_scale = self._quantize_mxfp8_moe_weight(w2_weight)
+        def quantize(w13_weight: torch.Tensor, w2_weight: torch.Tensor):
+            return (
+                *self._quantize_mxfp8_moe_weight(w13_weight),
+                *self._quantize_mxfp8_moe_weight(w2_weight),
+            )
+
+        w13, w13_scale, w2, w2_scale = self.quantize_weights(layer, quantize)
 
         self._setup_kernel(
             layer,
diff --git a/vllm/model_executor/layers/quantization/online/requantization.py b/vllm/model_executor/layers/quantization/online/requantization.py
index 01260342178b..72fb0a2718b6 100644
--- a/vllm/model_executor/layers/quantization/online/requantization.py
+++ b/vllm/model_executor/layers/quantization/online/requantization.py
@@ -1,10 +1,35 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from collections.abc import Iterator
+from contextlib import contextmanager
+from contextvars import ContextVar
+
 import torch
 
 from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase
 
+_requantization_source: ContextVar[bool] = ContextVar(
+    "requantization_source", default=False
+)
+
+
+@contextmanager
+def requantization_source_scope() -> Iterator[None]:
+    """Mark quantization methods created in this scope as requantization sources.
+
+    Their weights are only dequantized, so oracles can skip backend selection.
+    """
+    token = _requantization_source.set(True)
+    try:
+        yield
+    finally:
+        _requantization_source.reset(token)
+
+
+def is_requantization_source() -> bool:
+    return _requantization_source.get()
+
 
 class OnlineRequantizationMixin:
     """Shared lifecycle for converting serialized weights before quantization."""
diff --git a/vllm/model_executor/layers/quantization/online/shared_expert.py b/vllm/model_executor/layers/quantization/online/shared_expert.py
new file mode 100644
index 000000000000..d97638536156
--- /dev/null
+++ b/vllm/model_executor/layers/quantization/online/shared_expert.py
@@ -0,0 +1,344 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from collections.abc import Callable
+from typing import Any
+
+import torch
+
+from vllm.model_executor.layers.fused_moe import FusedMoeWeightScaleSupported
+from vllm.model_executor.layers.linear import LinearBase
+from vllm.model_executor.layers.quantization.online.requantization import (
+    OnlineRequantizationMixin,
+)
+from vllm.model_executor.layers.quantization.utils.config_utils import (
+    get_checkpoint_quant_keys,
+)
+from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
+    MXFP8_BLOCK_SIZE,
+    MXFP8_SCALE_DTYPE,
+    MXFP8_VALUE_DTYPE,
+)
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    QuantKey,
+    kMxfp8Static,
+)
+from vllm.model_executor.utils import set_weight_attrs
+
+SharedExpertQuantizer = Callable[[torch.Tensor], tuple[torch.Tensor, torch.Tensor]]
+
+
+def _param_stem(shard_id: str) -> str:
+    return "w2" if shard_id == "w2" else "w13"
+
+
+class OnlineSharedExpertLoader:
+    """Loads the fused shared experts (FSE) of a ``RoutedExperts`` layer with online
+    quantization.
+
+    - Staged shared experts: the shared experts are loaded as stored in temporary
+      slots (``online_w{13,2}_weight[_scale]``), dequantized after loading, and
+      quantized into the routed experts slots: with the routed experts when they
+      are requantized, else by ``finalize``.
+    - Unquantized shared experts with pre-quantized routed experts: the shared
+      experts are quantized on arrival by ``quantizer``, and loaded in the routed
+      experts slots.
+    """
+
+    def __init__(
+        self,
+        quant_config,
+        online_config,
+        quant_method: Any,
+        shared_expert_prefix: str | None,
+        num_logical_experts: int,
+        num_fused_shared_experts: int,
+    ):
+        self.shared_expert_prefix = shared_expert_prefix
+        self.num_logical_experts = num_logical_experts
+        self.num_fused_shared_experts = num_fused_shared_experts
+        # The routed experts are requantized online from their checkpoint format.
+        self.routed_requantized = (
+            isinstance(quant_method, OnlineRequantizationMixin)
+            and quant_method.requantization_source is not None
+        )
+        # Requantized routed experts quantize the staged shared experts with them.
+        self.stages_shared_experts = self.routed_requantized
+        self.quantizer: SharedExpertQuantizer | None = None
+        target_key = self._get_shared_target_key(online_config)
+        # Format the shared experts are stored in (None: unquantized). FSE
+        # compatibility checks require every projection to share it.
+        self.stored_key: QuantKey | None = None
+        if self.routed_requantized or target_key is not None:
+            self.stored_key = get_checkpoint_quant_keys(
+                quant_config, f"{shared_expert_prefix}.down_proj", LinearBase
+            )[0]
+        if (
+            not self.routed_requantized
+            and target_key is not None
+            and self.stored_key != target_key
+        ):
+            from vllm.model_executor.layers.quantization.online.base import (
+                ONLINE_SHARED_EXPERT_QUANTIZERS,
+            )
+
+            if target_key not in ONLINE_SHARED_EXPERT_QUANTIZERS:
+                raise NotImplementedError(
+                    "Online quantization of fused shared experts is not "
+                    f"supported for weight key {target_key}. Disable FSE or "
+                    "open an issue."
+                )
+            self.quantizer = ONLINE_SHARED_EXPERT_QUANTIZERS[target_key]
+            # Quantized shared experts are staged, as their weight and scale
+            # are needed together to be dequantized.
+            self.stages_shared_experts = self.stored_key is not None
+            from vllm.model_executor.layers.quantization.online.moe_base import (
+                OnlineMoEMethodBase,
+            )
+
+            if self.stages_shared_experts and isinstance(
+                quant_method, OnlineMoEMethodBase
+            ):
+                # `finalize` writes into pre-quantized routed experts slots.
+                raise NotImplementedError(
+                    "Fused shared experts stored quantized are not supported "
+                    "with routed experts quantized online from high "
+                    "precision. Disable FSE or open an issue."
+                )
+        self.staging_parameters: dict[str, torch.nn.Parameter] = {}
+
+    def _get_shared_target_key(self, online_config) -> QuantKey | None:
+        if self.shared_expert_prefix is None:
+            return None
+        # FSE compatibility checks require every shared-expert projection to
+        # use the routed experts' weight key.
+        resolved = online_config.resolve_quant_method_cls(
+            LinearBase, f"{self.shared_expert_prefix}.down_proj"
+        )
+        return None if resolved is None else resolved[3].weight
+
+    def create_staging(
+        self,
+        layer: torch.nn.Module,
+        params_dtype: torch.dtype,
+        **extra_weight_attrs,
+    ) -> None:
+        """Create parameters holding the fused shared experts as stored.
+
+        The checkpoint shared experts are loaded as stored into
+        ``online_w{13,2}_weight`` (plus ``online_w{13,2}_weight_scale`` for
+        quantized shared experts), which are separate from the routed experts'
+        parameters. They are dequantized by ``dequantize_staging``.
+        """
+        from vllm.model_executor.layers.quantization.online.base import (
+            ONLINE_SHARED_EXPERT_DEQUANTIZERS,
+        )
+
+        moe_config = layer.moe_config
+        hidden_size = moe_config.hidden_dim_unpadded
+        intermediate_size = moe_config.intermediate_size_per_partition_unpadded
+        checkpoint_key = self.stored_key
+        if (
+            checkpoint_key is not None
+            and checkpoint_key not in ONLINE_SHARED_EXPERT_DEQUANTIZERS
+        ):
+            raise NotImplementedError(
+                "Fused shared experts (FSE) with online quantization only support "
+                "shared experts stored unquantized or as "
+                f"{list(ONLINE_SHARED_EXPERT_DEQUANTIZERS)}, got {checkpoint_key}. "
+                "Disable FSE or open an issue."
+            )
+        # NOTE: The staging layout is MXFP8's, the only quantized format in
+        # ONLINE_SHARED_EXPERT_DEQUANTIZERS.
+        is_mxfp8 = checkpoint_key == kMxfp8Static
+        if is_mxfp8 and (
+            hidden_size % MXFP8_BLOCK_SIZE or intermediate_size % MXFP8_BLOCK_SIZE
+        ):
+            raise ValueError(
+                "MXFP8 shared experts require hidden/intermediate sizes "
+                f"divisible by {MXFP8_BLOCK_SIZE}."
+            )
+
+        def register(name: str, shape: tuple[int, ...], dtype: torch.dtype, **attrs):
+            param = torch.nn.Parameter(
+                torch.empty(shape, dtype=dtype), requires_grad=False
+            )
+            layer.register_parameter(name, param)
+            set_weight_attrs(param, extra_weight_attrs | attrs)
+            self.staging_parameters[name] = param
+
+        w13_rows = moe_config.w13_num_shards * intermediate_size
+        for stem, rows, cols in (
+            ("w13", w13_rows, hidden_size),
+            ("w2", hidden_size, intermediate_size),
+        ):
+            shape = (self.num_fused_shared_experts, rows)
+            weight_dtype = MXFP8_VALUE_DTYPE if is_mxfp8 else params_dtype
+            register(f"online_{stem}_weight", (*shape, cols), weight_dtype)
+            if is_mxfp8:
+                register(
+                    f"online_{stem}_weight_scale",
+                    (*shape, cols // MXFP8_BLOCK_SIZE),
+                    MXFP8_SCALE_DTYPE,
+                    quant_method=FusedMoeWeightScaleSupported.GROUP.value,
+                )
+
+    def dequantize_staging(
+        self, layer: torch.nn.Module
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Return the staged ``(w13, w2)`` shared experts, dequantized if needed."""
+        w13 = layer.online_w13_weight.data
+        w2 = layer.online_w2_weight.data
+        if self.stored_key is not None:
+            from vllm.model_executor.layers.quantization.online.base import (
+                ONLINE_SHARED_EXPERT_DEQUANTIZERS,
+            )
+
+            dequantize = ONLINE_SHARED_EXPERT_DEQUANTIZERS[self.stored_key]
+            w13 = dequantize(w13, layer.online_w13_weight_scale.data)
+            w2 = dequantize(w2, layer.online_w2_weight_scale.data)
+        return w13, w2
+
+    def release_staging(self, layer: torch.nn.Module) -> None:
+        for name, param in self.staging_parameters.items():
+            if layer._parameters.get(name) is param:
+                delattr(layer, name)
+        self.staging_parameters.clear()
+
+    def finalize(self, layer: torch.nn.Module) -> None:
+        """Quantize the staged shared experts into the pre-quantized routed experts
+        slots. To call after loading, before the routed experts' processing.
+
+        Requantized routed experts quantize the staged shared experts themselves.
+        """
+        if self.routed_requantized or not self.staging_parameters:
+            return
+        assert self.quantizer is not None
+        shared_w13, shared_w2 = self.dequantize_staging(layer)
+        for shared_id in range(self.num_fused_shared_experts):
+            expert_id = layer._map_global_expert_id_to_local_expert_id(
+                self.num_logical_experts + shared_id
+            )
+            if expert_id == -1:
+                continue
+            for stem, weight in (("w13", shared_w13), ("w2", shared_w2)):
+                weight, weight_scale = self.quantizer(weight[shared_id])
+                for name, tensor in (
+                    (f"{stem}_weight", weight),
+                    (f"{stem}_weight_scale", weight_scale),
+                ):
+                    rows, cols = tensor.shape
+                    getattr(layer, name).data[expert_id, :rows, :cols].copy_(tensor)
+        self.release_staging(layer)
+
+    def is_fused_shared_expert(self, global_expert_id: int) -> bool:
+        return (
+            self.num_logical_experts
+            <= global_expert_id
+            < self.num_logical_experts + self.num_fused_shared_experts
+        )
+
+    @staticmethod
+    def _staging_param_name(weight_name: str, shard_id: str) -> str | None:
+        """Staging parameter (``online_w{13,2}_weight[_scale]``) of a shared expert
+        tensor. They are separate from the routed experts' parameters. None for
+        other tensors.
+        """
+        stem = f"online_{_param_stem(shard_id)}"
+        if weight_name.endswith("_weight_scale"):
+            return f"{stem}_weight_scale"
+        if weight_name.endswith("_weight"):
+            return f"{stem}_weight"
+        return None
+
+    def redirect_param_name(
+        self, param_name: str, shard_id: str, expert_id: int
+    ) -> str:
+        """Parameter name a checkpoint tensor of expert `expert_id` loads into."""
+        # Once the staging is released, reloads go to `load`, which raises.
+        if self.staging_parameters and self.is_fused_shared_expert(expert_id):
+            return self._staging_param_name(param_name, shard_id) or param_name
+        return param_name
+
+    def redirect_mapping(
+        self, mapping: list[tuple[str, str, int, str]]
+    ) -> list[tuple[str, str, int, str]]:
+        """Point the expert params mapping of the shared experts to the staging
+        parameters, as e.g. ``w13_`` -> ``online_w13_``."""
+        if not self.staging_parameters:
+            return mapping
+        redirected = []
+        for param_name, weight_name, expert_id, shard_id in mapping:
+            if self.is_fused_shared_expert(expert_id) and weight_name.startswith(
+                f"experts.{expert_id}."
+            ):
+                stem = "w2_" if param_name.endswith("w2_") else "w13_"
+                param_name = f"{param_name.removesuffix(stem)}online_{stem}"
+            redirected.append((param_name, weight_name, expert_id, shard_id))
+        return redirected
+
+    def load(
+        self,
+        layer: Any,
+        param: torch.nn.Parameter,
+        loaded_weight: torch.Tensor,
+        weight_name: str,
+        shard_id: str,
+        global_expert_id: int,
+        expert_id: int,
+    ) -> tuple[torch.nn.Parameter, int, bool | None]:
+        """Prepare the load of a fused shared expert tensor in
+        ``RoutedExperts.weight_loader``.
+
+        Returns:
+            The parameter and local expert id to load into, and the load result if
+            the tensor was already loaded (None to let the caller load it).
+
+        """
+        if not self.is_fused_shared_expert(global_expert_id):
+            return param, expert_id, None
+
+        if self.stages_shared_experts:
+            if not self.staging_parameters:
+                # The staged shared experts were already quantized and released.
+                raise NotImplementedError(
+                    "Reloading fused shared experts staged for online quantization "
+                    "is not supported."
+                )
+            staging_name = self._staging_param_name(weight_name, shard_id)
+            if staging_name is None:
+                raise NotImplementedError(
+                    f"Fused shared expert tensor {weight_name!r} has no online "
+                    "staging parameter: only *_weight and *_weight_scale are "
+                    "supported when staging shared experts."
+                )
+            return (
+                getattr(layer, staging_name),
+                global_expert_id - self.num_logical_experts,
+                None,
+            )
+
+        # Only quantize high-precision shared expert weights, load anything else as is.
+        if (
+            self.quantizer is None
+            or not weight_name.endswith("_weight")
+            or not loaded_weight.is_floating_point()
+            or loaded_weight.dtype == param.dtype
+        ):
+            return param, expert_id, None
+
+        # Full-precision shared expert: quantize it, then load the weight and
+        # scale like a pre-quantized expert.
+        weight, weight_scale = self.quantizer(loaded_weight.to(layer.moe_config.device))
+        layer.weight_loader(
+            getattr(layer, f"{_param_stem(shard_id)}_weight_scale"),
+            weight_scale,
+            f"{weight_name}_scale",
+            shard_id,
+            global_expert_id,
+            return_success=True,
+        )
+        loaded = layer.weight_loader(
+            param, weight, weight_name, shard_id, global_expert_id, True
+        )
+        return param, expert_id, loaded
diff --git a/vllm/model_executor/layers/quantization/quark/quark.py b/vllm/model_executor/layers/quantization/quark/quark.py
index 2951f578b001..01c1172e85ad 100644
--- a/vllm/model_executor/layers/quantization/quark/quark.py
+++ b/vllm/model_executor/layers/quantization/quark/quark.py
@@ -968,6 +968,23 @@ def __init__(self, quantization_config: QuarkConfig):
     def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
         layer.scheme.process_weights_after_loading(layer)
 
+    def dequantize_weight(self, layer: torch.nn.Module) -> torch.Tensor:
+        """Reconstruct serialized weights for online requantization."""
+        from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
+            dequant_mxfp8_to_bf16,
+        )
+        from vllm.model_executor.layers.quantization.utils.quant_utils import (
+            kMxfp8Static,
+        )
+
+        if getattr(layer.scheme, "weight_quant_key", None) != kMxfp8Static:
+            raise NotImplementedError(
+                "Quark weight dequantization is only supported for MXFP8."
+            )
+        return dequant_mxfp8_to_bf16(
+            layer.weight.contiguous(), layer.weight_scale.contiguous()
+        )
+
     def create_weights(
         self,
         layer: torch.nn.Module,
diff --git a/vllm/model_executor/layers/quantization/quark/quark_moe.py b/vllm/model_executor/layers/quantization/quark/quark_moe.py
index a17f7b28da87..189e55d2269c 100644
--- a/vllm/model_executor/layers/quantization/quark/quark_moe.py
+++ b/vllm/model_executor/layers/quantization/quark/quark_moe.py
@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from collections.abc import Iterator
 from typing import TYPE_CHECKING, Any
 
 import torch
@@ -55,6 +56,7 @@
     make_mxfp4_moe_quant_config,
     mxfp4_round_up_hidden_size_and_intermediate_size,
     select_mxfp4_moe_backend,
+    triton_weight,
 )
 from vllm.model_executor.layers.fused_moe.oracle.mxfp8 import (
     select_mxfp8_moe_backend,
@@ -124,7 +126,11 @@
     normalize_e4m3fn_to_e4m3fnuz,
     per_tensor_dequantize,
 )
-from vllm.model_executor.utils import replace_parameter, set_weight_attrs
+from vllm.model_executor.utils import (
+    replace_parameter,
+    retire_parameter,
+    set_weight_attrs,
+)
 from vllm.platforms import current_platform
 
 if TYPE_CHECKING:
@@ -1239,16 +1245,7 @@ def create_weights(
             layer.w13_bias, layer.w2_bias = None, None
 
     def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
-        # Discard zero points (INT8 fused MoE kernel uses symmetric quant)
-        for attr in (
-            "w13_input_zero_point",
-            "w2_input_zero_point",
-            "w13_weight_zero_point",
-            "w2_weight_zero_point",
-        ):
-            if hasattr(layer, attr):
-                delattr(layer, attr)
-
+        # Loaded zero points stay registered; the INT8 fused MoE kernel ignores them.
         # For static input scales, collapse the per-expert scales to a single
         # value (the legacy fused_experts path expects one scale per layer).
         if self.static_input_scales:
@@ -1849,8 +1846,8 @@ def _setup_kernel(self, layer: RoutedExperts):
         ):
             # Triton-based backends: w13/w2 are triton_kernels.tensor.Tensor
             # Store on layer for apply(), scales are PrecisionConfig
-            layer.w13_weight = w13
-            layer.w2_weight = w2
+            layer.w13_weight_triton = w13
+            layer.w2_weight_triton = w2
             self.w13_precision_config = w13_scale
             self.w2_precision_config = w2_scale
         else:
@@ -1970,8 +1967,8 @@ def apply(
         assert self.moe_kernel is not None
         return self.moe_kernel.apply(
             hidden_states=x,
-            w1=layer.w13_weight,
-            w2=layer.w2_weight,
+            w1=triton_weight(layer, "w13_weight"),
+            w2=triton_weight(layer, "w2_weight"),
             topk_weights=topk_weights,
             topk_ids=topk_ids,
             activation=layer.activation,
@@ -1992,8 +1989,8 @@ def apply_monolithic(
         assert self.moe_kernel is not None
         return self.moe_kernel.apply_monolithic(
             hidden_states=x,
-            w1=layer.w13_weight,
-            w2=layer.w2_weight,
+            w1=triton_weight(layer, "w13_weight"),
+            w2=triton_weight(layer, "w2_weight"),
             router_logits=router_logits,
             activation=layer.activation,
             global_num_experts=layer.global_num_experts,
@@ -2204,6 +2201,8 @@ def __init__(
         super().__init__(moe, kNvfp4Static, activation_quant_key)
         self.quant_config = quant_config
         self.group_size = 16
+        # Whether the selected backend reads the activation input scales.
+        self.uses_input_scales = True
 
         # Select experts implementation.
         self.nvfp4_backend, self.experts_cls = select_nvfp4_moe_backend(
@@ -2319,8 +2318,8 @@ def create_weights(
 
     def dequantize_weight(
         self, layer: torch.nn.Module
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        """Dequantize Quark NVFP4 expert weights for requantization."""
+    ) -> Iterator[tuple[slice, torch.Tensor, torch.Tensor]]:
+        """Yield ``(experts, w13, w2)`` dequantized Quark NVFP4 experts."""
         return dequantize_nvfp4_moe_weights(
             layer.w13_weight,
             layer.w13_weight_scale,
@@ -2374,12 +2373,19 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None:
         replace_parameter(layer, "w13_weight", w13)
         replace_parameter(layer, "w13_weight_scale", w13_scale)
         replace_parameter(layer, "w13_weight_scale_2", w13_scale_2)
-        replace_parameter(layer, "w13_input_scale_2", a13_scale)
 
         replace_parameter(layer, "w2_weight", w2)
         replace_parameter(layer, "w2_weight_scale", w2_scale)
         replace_parameter(layer, "w2_weight_scale_2", w2_scale_2)
-        replace_parameter(layer, "w2_input_scale_2", a2_scale)
+        # Weight-only backends (Marlin, Humming) return no input scales: the
+        # loaded ones keep their data, retired, and the kernel config gets None.
+        self.uses_input_scales = a13_scale is not None
+        if self.uses_input_scales:
+            replace_parameter(layer, "w13_input_scale_2", a13_scale)
+            replace_parameter(layer, "w2_input_scale_2", a2_scale)
+        else:
+            retire_parameter(layer, "w13_input_scale_2")
+            retire_parameter(layer, "w2_input_scale_2")
 
         # Setup modular kernel.
         self.moe_quant_config = self.get_fused_moe_quant_config(layer)
@@ -2402,8 +2408,8 @@ def get_fused_moe_quant_config(
             w2_scale=layer.w2_weight_scale,
             w13_scale_2=layer.w13_weight_scale_2,
             w2_scale_2=layer.w2_weight_scale_2,
-            a13_scale=layer.w13_input_scale_2,
-            a2_scale=layer.w2_input_scale_2,
+            a13_scale=layer.w13_input_scale_2 if self.uses_input_scales else None,
+            a2_scale=layer.w2_input_scale_2 if self.uses_input_scales else None,
             swiglu_limit=getattr(layer, "swiglu_limit", None),
             swiglu_alpha=getattr(layer, "swiglu_alpha", None),
             swiglu_beta=getattr(layer, "swiglu_beta", None),
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 4634777c185d..38b373ad7d29 100644
--- a/vllm/model_executor/layers/quantization/quark/schemes/quark_nvfp4.py
+++ b/vllm/model_executor/layers/quantization/quark/schemes/quark_nvfp4.py
@@ -23,7 +23,7 @@
     ModelWeightParameter,
     PerTensorScaleParameter,
 )
-from vllm.model_executor.utils import replace_parameter
+from vllm.model_executor.utils import replace_parameter, retire_parameter
 
 __all__ = ["QuarkNVFP4"]
 
@@ -126,7 +126,8 @@ def create_weights(
     def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
         input_global_scale = layer.input_scale_2.max().to(torch.float32)
         replace_parameter(layer, "input_global_scale", input_global_scale)
-        del layer.input_scale_2
+        # Kernels read input_global_scale; the loaded param keeps its data.
+        retire_parameter(layer, "input_scale_2")
 
         weight_global_scale = layer.weight_scale_2.to(torch.float32)
 
@@ -142,7 +143,7 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
         weight_global_scale = weight_global_scale.max()
 
         replace_parameter(layer, "weight_global_scale", weight_global_scale)
-        del layer.weight_scale_2
+        retire_parameter(layer, "weight_scale_2")
 
         replace_parameter(
             layer, "alpha", layer.input_global_scale * layer.weight_global_scale
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 ee089be321bf..2caf3b7ac8f7 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
@@ -15,6 +15,9 @@
     init_mxfp6_linear_kernel,
     init_mxfp8_linear_kernel,
 )
+from vllm.model_executor.layers.fusion.quant_activation import (
+    expose_input_quant_key,
+)
 from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
     MXFP8_SCALE_DTYPE,
     MXFP8_VALUE_DTYPE,
@@ -246,6 +249,8 @@ def create_weights(
             self.ocp_mx_linear = init_mxfp8_linear_kernel(
                 weight_shape=layer.weight.shape
             )
+            if current_platform.is_rocm():
+                expose_input_quant_key(layer, self.ocp_mx_linear)
         else:
             raise NotImplementedError(
                 "No OCP MX linear kernel for weight_quant_key="
diff --git a/vllm/model_executor/layers/quantization/quark/schemes/quark_w8a8_int8.py b/vllm/model_executor/layers/quantization/quark/schemes/quark_w8a8_int8.py
index 364f5eaf614f..4eac3b91c27c 100644
--- a/vllm/model_executor/layers/quantization/quark/schemes/quark_w8a8_int8.py
+++ b/vllm/model_executor/layers/quantization/quark/schemes/quark_w8a8_int8.py
@@ -152,12 +152,7 @@ def _scale_weight_loader(
     # Checkpoints are serialized in quark format, which is
     # different from the format the kernel may want. Handle repacking here.
     def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
-        layer.register_parameter("weight_zero_point", None)
-        delattr(layer, "weight_zero_point")
-        if self.input_symmetric:
-            layer.register_parameter("input_zero_point", None)
-            delattr(layer, "input_zero_point")
-
+        # The loaded (symmetric) zero points stay registered; kernels ignore them.
         self.kernel.process_weights_after_loading(layer)
 
     def apply_weights(
diff --git a/vllm/model_executor/layers/quantization/utils/config_utils.py b/vllm/model_executor/layers/quantization/utils/config_utils.py
index 0c802e7af0db..c543178bdac1 100644
--- a/vllm/model_executor/layers/quantization/utils/config_utils.py
+++ b/vllm/model_executor/layers/quantization/utils/config_utils.py
@@ -11,7 +11,10 @@
 from vllm.logger import init_logger
 
 if TYPE_CHECKING:
+    import torch
+
     from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
+    from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey
 
 logger = init_logger(__name__)
 
@@ -56,6 +59,72 @@ def get_quark_ocp_mx_group_size(
     return int(group_size) if group_size is not None else None
 
 
+def _get_effective_quant_keys(
+    quant_config: "QuantizationConfig",
+    prefix: str,
+    layer_type: "type[torch.nn.Module]",
+) -> tuple["QuantKey | None", "QuantKey | None"]:
+    """Return the (weight, activation) keys a layer runs with: its online
+    quantization target if any, else its checkpoint format. ``None`` denotes an
+    unquantized weight, or an unknown/unquantized activation.
+
+    Raises:
+        NotImplementedError: If the checkpoint format can not be determined.
+
+    """
+    from vllm.model_executor.layers.linear import UnquantizedLinearMethod
+
+    online_quant_config = quant_config.online_quantization_config
+    target = (
+        online_quant_config.resolve_quant_method_cls(layer_type, prefix)
+        if online_quant_config is not None
+        else None
+    )
+    if target is not None:
+        _, _, _, quant_spec, method_cls = target
+        if method_cls in (None, UnquantizedLinearMethod):
+            return None, None
+        # NOTE: The online shorthands' `QuantSpec` (`_ONLINE_SHORTHANDS` in
+        # quantization.py) do not define the activation key: use the method's.
+        return quant_spec.weight, getattr(method_cls, "activation_quant_key", None)
+    return get_checkpoint_quant_keys(quant_config, prefix, layer_type)
+
+
+def get_checkpoint_quant_keys(
+    quant_config: "QuantizationConfig",
+    prefix: str,
+    layer_type: "type[torch.nn.Module]",
+) -> tuple["QuantKey | None", "QuantKey | None"]:
+    """Return the (weight, activation) keys a layer is stored with in the
+    checkpoint. ``None`` denotes an unquantized weight, or an unknown/unquantized
+    activation.
+
+    Raises:
+        NotImplementedError: If the checkpoint format can not be determined.
+
+    """
+    from vllm.model_executor.layers.fused_moe import RoutedExperts
+    from vllm.model_executor.layers.quantization.modelopt import ModelOptMxFp8Config
+    from vllm.model_executor.layers.quantization.quark.quark import QuarkConfig
+    from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp8Static
+
+    if isinstance(quant_config, QuarkConfig):
+        weight_key, activation_key, _ = quant_config.get_quant_method_target(
+            prefix, layer_type
+        )
+        return weight_key, activation_key
+    if isinstance(quant_config, ModelOptMxFp8Config):
+        # Non-excluded linears are MXFP8. TODO: activation keys are left unknown.
+        if quant_config.is_layer_excluded(prefix):
+            return None, None
+        if issubclass(layer_type, RoutedExperts):
+            return quant_config.routed_expert_weight_key(), None
+        return kMxfp8Static, None
+    raise NotImplementedError(
+        f"Cannot determine the checkpoint format of {type(quant_config).__name__}."
+    )
+
+
 def is_shared_expert_quant_fse_compatible(
     quant_config: "QuantizationConfig | None",
     expert_prefix: str,
@@ -78,11 +147,11 @@ def is_shared_expert_quant_fse_compatible(
     """
     from vllm.model_executor.layers.quantization.fp8 import Fp8Config
     from vllm.model_executor.layers.quantization.modelopt import ModelOptMxFp8Config
-    from vllm.model_executor.layers.quantization.online.fp8 import OnlineLinearBase
     from vllm.model_executor.layers.quantization.quark.quark import QuarkConfig
     from vllm.model_executor.layers.quantization.utils.quant_utils import (
         is_layer_skipped,
         kMxfp4Static,
+        kMxfp8Static,
     )
     from vllm.models.deepseek_v4.quant_config import DeepseekV4FP8Config
     from vllm.models.deepseek_v41.quant_config import (
@@ -95,93 +164,63 @@ def is_shared_expert_quant_fse_compatible(
     if quant_config is None:
         return True, None
 
+    # Re-quantization of a pre-quantized checkpoint.
     online_quant_config = quant_config.online_quantization_config
-    shared_expert_targets = None
     if online_quant_config is not None:
+        from vllm.model_executor.layers.fused_moe import RoutedExperts
         from vllm.model_executor.layers.linear import LinearBase
 
         online_quant_config.packed_modules_mapping = quant_config.packed_modules_mapping
-        shared_expert_targets = [
-            online_quant_config.resolve_quant_method_cls(
-                LinearBase, f"{shared_expert_prefix}.{projection_name}"
-            )
+        shared_prefixes = [
+            f"{shared_expert_prefix}.{projection_name}"
             for projection_name in projection_names
         ]
-        if all(target is None for target in shared_expert_targets):
-            shared_expert_targets = None
-
-    if shared_expert_targets is not None:
-        from vllm.model_executor.layers.fused_moe import RoutedExperts
-        from vllm.model_executor.layers.linear import UnquantizedLinearMethod
-        from vllm.model_executor.layers.quantization.online.base import (
-            ONLINE_SHARED_EXPERT_QUANTIZERS,
+        is_online_quantized = any(
+            online_quant_config.resolve_quant_method_cls(layer_type, prefix) is not None
+            for layer_type, prefix in [
+                (RoutedExperts, expert_prefix),
+                *((LinearBase, prefix) for prefix in shared_prefixes),
+            ]
         )
-
-        # NOTE: online shared experts quantization check is only implemented for quark
-        # quant method at the moment. This can be extended here for other quant methods.
-        # Requirement is:
-        # - Retrieve routed weight quant key
-        # - Retrieve routed activation quant key
-        # and ensure the online quantization set on shared expert match them.
-        if isinstance(quant_config, QuarkConfig):
-            routed_weight_key, routed_activation_key, _ = (
-                quant_config.get_quant_method_target(expert_prefix, RoutedExperts)
-            )
-        elif (
-            isinstance(quant_config, (DeepseekV4FP8Config, DeepseekV41FP8Config))
-            and quant_config.expert_dtype == "fp4"
-            and moe_backend == "aiter_mega_moe"
-        ):
-            # AITER MegaMoE quantizes the activations of every slot itself.
-            routed_weight_key, routed_activation_key = kMxfp4Static, None
-        else:
-            reason = (
-                "online quantization FSE compatibility check is not implemented for "
-                f"{type(quant_config).__name__}"
-            )
-            logger.warning_once(reason)
-            return False, reason
-
-        if routed_weight_key not in ONLINE_SHARED_EXPERT_QUANTIZERS:
-            return (
-                False,
-                "online shared-expert quantization is not supported for routed "
-                f"expert weight key {routed_weight_key}",
-            )
-
-        for shared_expert_target in shared_expert_targets:
-            if shared_expert_target is None:
-                return (
-                    False,
-                    "online quantization targets only part of the shared expert at "
-                    f"{shared_expert_prefix}",
-                )
-
-            _, _, _, shared_quant_spec, shared_method_cls = shared_expert_target
-            if shared_method_cls in (None, UnquantizedLinearMethod):
-                return (
-                    False,
-                    "online quantization excludes FSE weights at "
-                    f"{shared_expert_prefix}",
-                )
-            assert issubclass(shared_method_cls, OnlineLinearBase)
-
-            # NOTE: We can not rely on shared_quant_spec.activation_quant_key
-            # as _ONLINE_SHORTHANDS: dict[str, QuantSpec] in quantization.py
-            # does not define the default activation quant key.
-            # `online_quant_config.resolve_quant_method_cls` should return the
-            # activation quant key as well. Only some online methods define it
-            # on the class, so compare the weight key first.
-            if shared_quant_spec.weight != routed_weight_key or (
-                moe_backend != "aiter_mega_moe"
-                and shared_method_cls.activation_quant_key != routed_activation_key
-            ):
+        if is_online_quantized:
+            # Only check that the routed and shared experts run with the same
+            # keys. Unsupported loading paths (e.g. the format the shared expert
+            # is stored in) raise NotImplementedError at weight loading.
+            try:
+                if (
+                    isinstance(
+                        quant_config, (DeepseekV4FP8Config, DeepseekV41FP8Config)
+                    )
+                    and quant_config.expert_dtype == "fp4"
+                    and moe_backend == "aiter_mega_moe"
+                ):
+                    routed_keys: tuple[QuantKey | None, QuantKey | None] = (
+                        kMxfp4Static,
+                        None,
+                    )
+                else:
+                    routed_keys = _get_effective_quant_keys(
+                        quant_config, expert_prefix, RoutedExperts
+                    )
+                shared_keys = {
+                    _get_effective_quant_keys(quant_config, prefix, LinearBase)
+                    for prefix in shared_prefixes
+                }
+            except NotImplementedError as e:
+                logger.warning_once(str(e))
+                return False, str(e)
+            if moe_backend == "aiter_mega_moe":
+                # AITER MegaMoE quantizes the activations of every slot itself.
+                routed_keys = (routed_keys[0], None)
+                shared_keys = {(weight_key, None) for weight_key, _ in shared_keys}
+            if shared_keys != {routed_keys}:
                 return (
                     False,
-                    "online shared-expert quantization keys do not match the "
-                    "routed expert keys",
+                    f"shared expert (weight, activation) keys {shared_keys} at "
+                    f"{shared_expert_prefix} do not match the routed expert keys "
+                    f"{routed_keys}",
                 )
-        return True, None
+            return True, None
 
     if moe_backend == "aiter_mega_moe":
         return (
@@ -367,6 +406,17 @@ def is_ignored(layer_name: str) -> bool:
                 "MXFP8 excludes routed and shared experts inconsistently at "
                 f"{shared_expert_prefix}",
             )
+        # The shared expert is stored as MXFP8, so it can take an expert slot
+        # only when the routed experts are MXFP8 too (not NVFP4-routed hybrids).
+        if (
+            not expert_excluded
+            and quant_config.routed_expert_weight_key() != kMxfp8Static
+        ):
+            return (
+                False,
+                "routed experts do not run as MXFP8 like the shared expert at "
+                f"{shared_expert_prefix}",
+            )
 
         return True, None
 
diff --git a/vllm/model_executor/layers/quantization/utils/humming/linear.py b/vllm/model_executor/layers/quantization/utils/humming/linear.py
index 821f23aeda85..f8557ccbe4a3 100644
--- a/vllm/model_executor/layers/quantization/utils/humming/linear.py
+++ b/vllm/model_executor/layers/quantization/utils/humming/linear.py
@@ -12,10 +12,80 @@
 from vllm.model_executor.layers.quantization.utils.humming.schema import (
     check_and_fallback_input_schema,
 )
+from vllm.model_executor.utils import (
+    get_live_param,
+    kernel_parameters,
+    move_parameter,
+    replace_parameter,
+    retire_parameter,
+    set_kernel_parameter,
+)
 
 if TYPE_CHECKING:
     from vllm.utils.humming import BaseInputSchema, LayerConfig
 
+# A param the conversion leaves unused keeps its loaded data, retired, when it is
+# at most this fraction of the largest converted tensor (scalars, per-tensor and
+# per-channel scales); a larger one backs a new tensor (its data is in them).
+_KEEP_UNUSED_FRACTION = 1 / 256
+
+
+def _nbytes(t: torch.Tensor) -> int:
+    return t.numel() * t.element_size()
+
+
+def replace_layer_parameters(
+    layer: torch.nn.Module, tensors: dict[str, torch.Tensor], names: list[str]
+) -> None:
+    """Write ``tensors`` into ``layer``, keeping existing parameter objects.
+
+    A param in ``names`` that gets no new tensor stays registered: an alias of a
+    replaced param follows it; one whose storage a new tensor is in, or a large
+    one, is also registered under a new name (see `move_parameter`); any other
+    keeps its data, retired."""
+    params = layer._parameters
+    replaced = {id(p) for name in tensors if (p := params.get(name)) is not None}
+    leftovers = {
+        name: p
+        for name in names
+        if name not in tensors
+        and (p := params.get(name)) is not None
+        and id(p) not in replaced
+    }
+
+    def adopt(src: str, dst: str) -> None:
+        move_parameter(layer, src, dst)
+        for name in [n for n, p in leftovers.items() if p is params[src]]:
+            del leftovers[name]
+
+    new = [name for name in tensors if name not in params]
+    for name in new:
+        ptr = tensors[name].untyped_storage().data_ptr()
+        src = next(
+            (n for n, p in leftovers.items() if p.untyped_storage().data_ptr() == ptr),
+            None,
+        )
+        if ptr and src is not None:
+            adopt(src, name)
+    limit = _KEEP_UNUSED_FRACTION * max(map(_nbytes, tensors.values()), default=0)
+    large = sorted(
+        (n for n, p in leftovers.items() if _nbytes(p) > limit),
+        key=lambda n: _nbytes(leftovers[n]),
+        reverse=True,
+    )
+    free = sorted(
+        (n for n in new if n not in params),
+        key=lambda n: _nbytes(tensors[n]),
+        reverse=True,
+    )
+    for src, dst in zip(large, free):
+        if src in leftovers:
+            adopt(src, dst)
+    for name in leftovers:
+        retire_parameter(layer, name)
+    for name, tensor in tensors.items():
+        set_kernel_parameter(layer, name, tensor)
+
 
 def convert_linear_layer_to_humming_standard(
     layer: LinearBase, name_map: dict[str, str]
@@ -24,15 +94,15 @@ def convert_linear_layer_to_humming_standard(
     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)
+        param = getattr(layer, checkpoint_name)
+        tensor = param.data
 
         if name == "weight":
             # (input_dim, output_dim) after PWAL, if it transposed the weight.
             input_dim, output_dim = getattr(
                 layer,
                 "processed_weight_dims",
-                (getattr(tensor, "input_dim", 1), getattr(tensor, "output_dim", 0)),
+                (getattr(param, "input_dim", 1), getattr(param, "output_dim", 0)),
             )
 
             if input_dim == 0 and output_dim == 1:
@@ -42,19 +112,17 @@ def convert_linear_layer_to_humming_standard(
 
             tensor = tensor.view(tensor.size(0), -1).view(torch.int32)
         elif name in ["weight_scale", "zero_point"]:
-            if getattr(tensor, "output_dim", 0) == 1:
+            if getattr(param, "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)
+        if name != checkpoint_name:
+            # One param under both names: the loader keeps its name.
+            move_parameter(layer, checkpoint_name, name)
+        replace_parameter(layer, name, tensor)
 
 
 def prepare_humming_linear_layer_config(
@@ -91,7 +159,7 @@ def prepare_humming_linear_layer_config(
     shape_n_stacks = layer.output_partition_sizes
 
     # Step 1: convert weight and input schemas to humming standard format
-    source_tensors = dict(layer.named_parameters(recurse=False))
+    source_tensors = kernel_parameters(layer)
     weight_schema, tensors = weight_schema.convert_humming(
         tensors=source_tensors,
         shape_n_stacks=shape_n_stacks,
@@ -148,11 +216,7 @@ def prepare_humming_linear_layer_config(
     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(recurse=False)):
-        delattr(layer, name)
-    for name, tensor in tensors.items():
-        param = torch.nn.Parameter(tensor, requires_grad=False)
-        setattr(layer, name, param)
+    replace_layer_parameters(layer, tensors, list(layer._parameters))
 
     return config
 
@@ -179,19 +243,19 @@ def apply_humming_linear(
     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 = get_live_param(layer, "input_scale") if mode.should_quantize else None
     input_scale_2 = (
-        getattr(layer, "input_scale_2", None) if mode.has_secondary_scale else None
+        get_live_param(layer, "input_scale_2") 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),
+        weight_scale=get_live_param(layer, "weight_scale"),
+        zero_point=get_live_param(layer, "zero_point"),
+        bias=get_live_param(layer, "zero_bias" if skip_bias_add else "bias"),
+        weight_scale_2=get_live_param(layer, "weight_scale_2"),
         input_scale=input_scale,
         input_scale_2=input_scale_2,
         hadamard_block_size=layer.weight_schema.hadamard_block_size,
diff --git a/vllm/model_executor/layers/quantization/utils/humming/moe.py b/vllm/model_executor/layers/quantization/utils/humming/moe.py
index 418e9acbe70a..44b5a9020443 100644
--- a/vllm/model_executor/layers/quantization/utils/humming/moe.py
+++ b/vllm/model_executor/layers/quantization/utils/humming/moe.py
@@ -24,6 +24,11 @@
     HummingGroupedExperts,
     HummingIndexedExperts,
 )
+from vllm.model_executor.layers.quantization.utils.humming.linear import (
+    get_live_param,
+    kernel_parameters,
+    replace_layer_parameters,
+)
 from vllm.model_executor.layers.quantization.utils.humming.schema import (
     humming_is_layer_skipped,
 )
@@ -168,14 +173,14 @@ def get_humming_moe_quant_config(
         weight_dtype=str(weight_schema.b_dtype),
         weight_group_shape=weight_group_shape,
         activation_group_shape=activation_group_shape,
-        w1_scale=getattr(layer, "w13_weight_scale", None),
-        w1_gscale=getattr(layer, "w13_weight_scale_2", None),
-        w1_zp=getattr(layer, "w13_zero_point", None),
-        w1_bias=getattr(layer, "w13_bias", None),
-        w2_scale=getattr(layer, "w2_weight_scale", None),
-        w2_gscale=getattr(layer, "w2_weight_scale_2", None),
-        w2_zp=getattr(layer, "w2_zero_point", None),
-        w2_bias=getattr(layer, "w2_bias", None),
+        w1_scale=get_live_param(layer, "w13_weight_scale"),
+        w1_gscale=get_live_param(layer, "w13_weight_scale_2"),
+        w1_zp=get_live_param(layer, "w13_zero_point"),
+        w1_bias=get_live_param(layer, "w13_bias"),
+        w2_scale=get_live_param(layer, "w2_weight_scale"),
+        w2_gscale=get_live_param(layer, "w2_weight_scale_2"),
+        w2_zp=get_live_param(layer, "w2_zero_point"),
+        w2_bias=get_live_param(layer, "w2_bias"),
         gemm1_alpha=gemm1_alpha,
         gemm1_beta=gemm1_beta,
         gemm1_clamp_limit=gemm1_clamp_limit,
@@ -189,7 +194,7 @@ def get_humming_moe_quant_config(
         )
         for name in ("input_scale", "input_scale_2"):
             key = f"{prefix}_{name}"
-            value = getattr(layer, f"{sublayer}_{name}", None)
+            value = get_live_param(layer, f"{sublayer}_{name}")
             setattr(config, key, value)
     return config
 
@@ -301,12 +306,17 @@ def _extract_sublayer_tensors(
     layer: "RoutedExperts",
     sublayer_name: str,
 ) -> dict[str, torch.Tensor]:
-    """Extract tensors for a specific sublayer from the layer's state dict."""
-    return dict(
-        (key.removeprefix(sublayer_name + "_"), value)
-        for key, value in layer.state_dict().items()
+    """Extract a sublayer's live params and persistent buffers."""
+    buffers = {
+        name: buffer
+        for name, buffer in layer._buffers.items()
+        if buffer is not None and name not in layer._non_persistent_buffers_set
+    }
+    return {
+        key.removeprefix(sublayer_name + "_"): value
+        for key, value in (kernel_parameters(layer) | buffers).items()
         if key.startswith(sublayer_name + "_")
-    )
+    }
 
 
 def _replace_layer_parameters(
@@ -324,19 +334,15 @@ def _replace_layer_parameters(
         preserve_bias: If True, don't delete bias parameters
 
     """
-    # Delete old parameters
-    for name, _ in list(layer.named_parameters()):
-        if not name.startswith(sublayer_name + "_"):
-            continue
-        if preserve_bias and name == sublayer_name + "_bias":
-            continue
-        delattr(layer, name)
-
-    # Set new parameters
-    for name, tensor in tensors.items():
-        param_name = f"{sublayer_name}_{name}"
-        param = torch.nn.Parameter(tensor, requires_grad=False)
-        setattr(layer, param_name, param)
+    names = [
+        name
+        for name in layer._parameters
+        if name.startswith(sublayer_name + "_")
+        and not (preserve_bias and name == sublayer_name + "_bias")
+    ]
+    replace_layer_parameters(
+        layer, {f"{sublayer_name}_{n}": t for n, t in tensors.items()}, names
+    )
 
 
 def _convert_sublayer_to_humming(
diff --git a/vllm/model_executor/layers/quantization/utils/mxfp8_utils.py b/vllm/model_executor/layers/quantization/utils/mxfp8_utils.py
index bfcc69e4b1de..6567eeb86242 100644
--- a/vllm/model_executor/layers/quantization/utils/mxfp8_utils.py
+++ b/vllm/model_executor/layers/quantization/utils/mxfp8_utils.py
@@ -145,6 +145,7 @@ def _kernel(
 
 
 _MXFP8_QUANT_KERNEL = None
+_MXFP8_M1_BLOCK_M = {2048: 4, 3072: 4, 6144: 2}
 
 
 def _mxfp8_e4m3_quantize_triton(
@@ -163,7 +164,7 @@ def _mxfp8_e4m3_quantize_triton(
     scales = torch.empty(
         (M, K // MXFP8_BLOCK_SIZE), dtype=MXFP8_SCALE_DTYPE, device=x.device
     )
-    BLOCK_M = 64
+    BLOCK_M = _MXFP8_M1_BLOCK_M.get(K, 64) if M == 1 else 64
     grid = (triton.cdiv(M, BLOCK_M), K // MXFP8_BLOCK_SIZE)
     _MXFP8_QUANT_KERNEL[grid](
         x,
@@ -180,6 +181,7 @@ def _mxfp8_e4m3_quantize_triton(
         BLOCK_M=BLOCK_M,
         FP8_MAX=float(torch.finfo(MXFP8_VALUE_DTYPE).max),
         TINY=float(torch.finfo(torch.float32).tiny),
+        num_warps=4,
     )
     return xq, scales
 
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 cad141206678..ecede028029f 100644
--- a/vllm/model_executor/layers/quantization/utils/nvfp4_emulation_utils.py
+++ b/vllm/model_executor/layers/quantization/utils/nvfp4_emulation_utils.py
@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from collections.abc import Iterator
 from types import SimpleNamespace
 
 import torch
@@ -17,6 +18,9 @@
     "ref_nvfp4_quant",
 ]
 
+# Target size of the dequantized expert chunks yielded for requantization.
+_DEQUANT_CHUNK_BYTES = 1 << 30
+
 FLOAT4_E2M1_MAX = scalar_types.float4_e2m1f.max()
 FLOAT4_E2M1_MAX_RECIPROCAL = 1 / FLOAT4_E2M1_MAX
 
@@ -630,7 +634,7 @@ def dequantize_to_dtype(
     return out.to(dtype)
 
 
-def dequantize_nvfp4_moe_weights(
+def _dequantize_nvfp4_moe_experts(
     w13_weight_packed: torch.Tensor,
     w13_weight_scale: torch.Tensor,
     w13_weight_global_scale: torch.Tensor,
@@ -638,16 +642,9 @@ def dequantize_nvfp4_moe_weights(
     w2_weight_scale: torch.Tensor,
     w2_weight_global_scale: torch.Tensor,
     dtype: torch.dtype,
-    group_size: int = 16,
-    *,
-    invert_global_scales: bool = True,
+    group_size: int,
+    invert_global_scales: bool,
 ) -> tuple[torch.Tensor, torch.Tensor]:
-    """Materialize serialized NVFP4 MoE weights for requantization.
-
-    Compressed-tensors stores global encode scales, while ModelOpt stores
-    ``weight_scale_2`` as the corresponding decode multiplier. Set
-    ``invert_global_scales`` according to the source checkpoint convention.
-    """
     if invert_global_scales:
         w13_weight_global_scale = 1.0 / w13_weight_global_scale
         w2_weight_global_scale = 1.0 / w2_weight_global_scale
@@ -682,6 +679,54 @@ def dequantize_nvfp4_moe_weights(
     return w13, w2
 
 
+@torch.no_grad()
+def dequantize_nvfp4_moe_weights(
+    w13_weight_packed: torch.Tensor,
+    w13_weight_scale: torch.Tensor,
+    w13_weight_global_scale: torch.Tensor,
+    w2_weight_packed: torch.Tensor,
+    w2_weight_scale: torch.Tensor,
+    w2_weight_global_scale: torch.Tensor,
+    dtype: torch.dtype,
+    group_size: int = 16,
+    *,
+    invert_global_scales: bool = True,
+) -> Iterator[tuple[slice, torch.Tensor, torch.Tensor]]:
+    """Dequantize serialized NVFP4 MoE weights for requantization.
+
+    Compressed-tensors stores global encode scales, while ModelOpt stores
+    ``weight_scale_2`` as the corresponding decode multiplier. Set
+    ``invert_global_scales`` according to the source checkpoint convention.
+
+    Returns:
+        An iterator of ``(experts, w13, w2)`` over groups of experts, so that
+        the dequantized staging tensors stay around ``_DEQUANT_CHUNK_BYTES``.
+
+    """
+    num_experts = w13_weight_packed.shape[0]
+    # Two fp4 values per packed byte.
+    expert_bytes = (
+        dtype.itemsize
+        * 2
+        * (w13_weight_packed[0].numel() + w2_weight_packed[0].numel())
+    )
+    experts_per_chunk = max(1, _DEQUANT_CHUNK_BYTES // expert_bytes)
+    for start in range(0, num_experts, experts_per_chunk):
+        experts = slice(start, min(start + experts_per_chunk, num_experts))
+        w13, w2 = _dequantize_nvfp4_moe_experts(
+            w13_weight_packed[experts],
+            w13_weight_scale[experts],
+            w13_weight_global_scale[experts],
+            w2_weight_packed[experts],
+            w2_weight_scale[experts],
+            w2_weight_global_scale[experts],
+            dtype,
+            group_size,
+            invert_global_scales,
+        )
+        yield experts, w13, w2
+
+
 def get_reciprocal(x):
     if isinstance(x, torch.Tensor):
         # torch.where yields operation not permitted when stream is capturing.
diff --git a/vllm/model_executor/layers/quantization/utils/quant_utils.py b/vllm/model_executor/layers/quantization/utils/quant_utils.py
index 8cdbb3e596ae..557034f3e3d7 100644
--- a/vllm/model_executor/layers/quantization/utils/quant_utils.py
+++ b/vllm/model_executor/layers/quantization/utils/quant_utils.py
@@ -217,6 +217,17 @@ def __str__(self):
     FP4_DTYPE, scale=kNvfp4DynamicGroupScale, scale2=kDynamicTokenScale
 )
 
+
+def is_dynamic_per_token_outer_scale(key: QuantKey | None) -> bool:
+    """Whether ``key``'s second-level scale is computed per token at runtime."""
+    return (
+        key is not None
+        and key.scale2 is not None
+        and not key.scale2.static
+        and key.scale2.group_shape == GroupShape.PER_TOKEN
+    )
+
+
 kNvfp4StaticGroupScale = ScaleDesc(FP8_DTYPE, True, GroupShape(1, 16))
 kNvfp4Static = QuantKey(
     FP4_DTYPE, scale=kNvfp4StaticGroupScale, scale2=kStaticTensorScale
diff --git a/vllm/model_executor/layers/rotary_embedding/mrope.py b/vllm/model_executor/layers/rotary_embedding/mrope.py
index 9bcc5a95e16d..e1dbd1169f09 100644
--- a/vllm/model_executor/layers/rotary_embedding/mrope.py
+++ b/vllm/model_executor/layers/rotary_embedding/mrope.py
@@ -436,38 +436,18 @@ def forward_native(
         key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape)
         return query, key
 
-    def forward_cuda(
+    def _forward_text_rope(
         self,
         positions: torch.Tensor,
         query: torch.Tensor,
-        key: torch.Tensor | None = None,
-        offsets: torch.Tensor | None = None,
-    ) -> tuple[torch.Tensor, torch.Tensor | None]:
-        assert positions.ndim == 1 or positions.ndim == 2
-        assert key is not None
-
+        key: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         cos_sin_cache = self._match_cos_sin_cache_dtype(query)
         num_tokens = positions.shape[-1]
         cos_sin = cos_sin_cache[positions]
         cos, sin = cos_sin.chunk(2, dim=-1)
         query_shape = query.shape
         key_shape = key.shape
-        if positions.ndim == 2:
-            assert self.mrope_section
-
-            q, k = triton_mrope(
-                query,
-                key,
-                cos,
-                sin,
-                self.mrope_section,
-                self.head_size,
-                self.rotary_dim,
-                self.mrope_interleaved,
-                self.is_neox_style,
-            )
-
-            return q.reshape(query_shape), k.reshape(key_shape)
 
         query = query.view(num_tokens, -1, self.head_size)
         query_rot = query[..., : self.rotary_dim]
@@ -490,6 +470,40 @@ def forward_cuda(
         key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape)
         return query, key
 
+    def forward_cuda(
+        self,
+        positions: torch.Tensor,
+        query: torch.Tensor,
+        key: torch.Tensor | None = None,
+        offsets: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        assert positions.ndim == 1 or positions.ndim == 2
+        assert key is not None
+
+        if positions.ndim == 2:
+            assert self.mrope_section
+            cos_sin_cache = self._match_cos_sin_cache_dtype(query)
+            cos_sin = cos_sin_cache[positions]
+            cos, sin = cos_sin.chunk(2, dim=-1)
+            query_shape = query.shape
+            key_shape = key.shape
+
+            q, k = triton_mrope(
+                query,
+                key,
+                cos,
+                sin,
+                self.mrope_section,
+                self.head_size,
+                self.rotary_dim,
+                self.mrope_interleaved,
+                self.is_neox_style,
+            )
+
+            return q.reshape(query_shape), k.reshape(key_shape)
+
+        return self._forward_text_rope(positions, query, key)
+
     def forward_hip(
         self,
         positions: torch.Tensor,
@@ -536,6 +550,26 @@ def forward_xpu(
         offsets: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor | None]:
         return self.forward_cuda(positions, query, key, offsets)
+        if positions.ndim == 2:
+            assert key is not None
+            assert self.mrope_section is not None
+            import vllm._xpu_ops  # noqa: F401
+
+            cos_sin_cache = self._match_cos_sin_cache_dtype(query)
+            torch.ops._xpu_C.multimodal_rotary_embedding(
+                positions,
+                query,
+                key,
+                self.head_size,
+                cos_sin_cache,
+                self.is_neox_style,
+                self.mrope_section,
+                self.mrope_interleaved,
+            )
+            return query, key
+        assert positions.ndim == 1
+        assert key is not None
+        return self._forward_text_rope(positions, query, key)
 
     @staticmethod
     def get_next_input_positions(
diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py
index fb3eb528a856..89358fba8faa 100644
--- a/vllm/model_executor/layers/sparse_attn_indexer.py
+++ b/vllm/model_executor/layers/sparse_attn_indexer.py
@@ -2,6 +2,8 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Custom Sparse Attention Indexer layers."""
 
+import math
+
 import torch
 import torch.distributed as dist
 
@@ -29,9 +31,18 @@
 from vllm.model_executor.kernels.attention.dsa.candidate_blocks import (
     select_candidate_blocks as _select_candidate_blocks,
 )
+from vllm.model_executor.kernels.attention.dsa.fused_mqa_topk import (
+    fused_mqa_topk_chunk,
+    fused_mqa_topk_key_scales,
+    fused_mqa_topk_prefill,
+    get_fused_mqa_topk_workspaces,
+    use_fused_mqa_topk,
+)
 from vllm.model_executor.layers.indexer_topk import (
     RADIX_TOPK_WORKSPACE_SIZE,
+    deep_select_topk,
     get_indexer_topk,
+    is_deep_select_supported,
 )
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
     get_fp8_min_max,
@@ -46,6 +57,7 @@
     mqa_logits_use_mx_scales,
 )
 from vllm.utils.import_utils import has_cutedsl
+from vllm.utils.math_utils import round_up
 from vllm.utils.torch_utils import (
     LayerNameType,
     _encode_layer_name,
@@ -54,6 +66,7 @@
 )
 from vllm.v1.attention.backends.mla.indexer import (
     DeepseekV32IndexerMetadata,
+    dense_logits_rows,
     tp_prefill_row_sharding_supported,
 )
 from vllm.v1.attention.ops.common import pack_seq_triton, unpack_seq_triton
@@ -478,6 +491,18 @@ def sparse_attn_indexer(
         )
         assert candidate_block_size > 0
 
+    fused_mqa_topk = use_fused_mqa_topk(
+        query_shape=q_quant.shape[1:],
+        topk=topk_tokens,
+        fp4_cache=use_fp4_cache,
+        full_context=use_pcp or dcp_world_size == 1,
+        candidate_selection=candidate_blocks is not None,
+    )
+
+    # DeepSelect ranks SM100's BF16 decode logits directly; other top-k
+    # backends and the DCP merge read an FP32 copy.
+    fp32_decode_logits = topk_backend != "deep_select" or dcp_world_size != 1
+
     # assert isinstance(attn_metadata, dict)
     if not isinstance(attn_metadata, dict):
         # Reserve workspace for indexer during profiling run
@@ -497,12 +522,46 @@ def sparse_attn_indexer(
             )
             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
-        max_logits_elems = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024
+        if fused_mqa_topk:
+            # Allocate the largest fused_mqa_topk workspaces so they are
+            # accounted for (freed again, they overlap the dense logits below).
+            get_fused_mqa_topk_workspaces(q_quant.device, q_quant.shape[0])
+
+        # Decode rows per step: each decode is padded to 1 + speculative tokens.
+        # The profiling run sets the vLLM config; other dummy runs may not.
+        max_decode_tokens = 0
+        if (cfg := get_current_vllm_config_or_none()) is not None:
+            max_decode_tokens = min(
+                cfg.scheduler_config.max_num_seqs,
+                cfg.scheduler_config.max_num_batched_tokens,
+            ) * (1 + cfg.num_speculative_tokens)
+
+        # Needed with the fused kernel too: chunks fused_mqa_topk_chunk keeps
+        # dense take this path (row-split to stay within the budget).
+        # Dummy allocation for the peak logits memory during inference: a
+        # prefill chunk's, or the decode logits, which eager and piecewise
+        # steps allocate outside the CUDA graph pools. DeepGEMM emits them in
+        # BF16 on SM100 and FP32 elsewhere, each row holding whole KV splits at
+        # a 1024-byte stride; the FP32 copy keeps that stride. Prefill always
+        # reads the FP32 copy, and the chunk splitter caps it at the budget.
+        if current_platform.is_device_capability_family(100):
+            logits_dtype, split_kv = torch.bfloat16, 384
+        else:
+            logits_dtype, split_kv = torch.float32, 256
+        row_align = math.lcm(split_kv, 1024 // logits_dtype.itemsize)
+        copy_bytes = 0 if logits_dtype == torch.float32 else torch.float32.itemsize
+        decode_copy_bytes = copy_bytes if fp32_decode_logits else 0
+        prefill_budget = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024
+        max_logits_bytes = max(
+            prefill_budget
+            // torch.float32.itemsize
+            * (logits_dtype.itemsize + copy_bytes),
+            max_decode_tokens
+            * round_up(max_model_len, row_align)
+            * (logits_dtype.itemsize + decode_copy_bytes),
+        )
         _ = torch.empty(
-            max_logits_elems, dtype=torch.uint8, device=hidden_states.device
+            max_logits_bytes, dtype=torch.uint8, device=hidden_states.device
         )
 
         # The dummy run returns before exchange_topk executes, so reserve its
@@ -606,13 +665,18 @@ def sparse_attn_indexer(
         prefill_metadata = attn_metadata_narrowed.prefill
         assert prefill_metadata is not None
 
-        # Get the full shared workspace buffers once (will allocate on first use).
+        # Get the shared workspace buffers once, sized for this batch's largest
+        # chunk rather than the profiled bound (total_seq_lens), so the scratch
+        # stays at the workspace head and leaves retained tail spans intact.
         # 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
+            max(c.max_local_total_seq_lens for c in prefill_metadata.chunks),
+            head_dim,
+            fp8_dtype,
+            use_fp4_cache,
         )
         # PCP + DCP needs two more pairs: the rank-major all-gather destination
         # and the de-interleaved result.
@@ -690,9 +754,29 @@ def sparse_attn_indexer(
             # DeepGEMM returns compressed logits: row i starts at K position
             # cu_seqlen_ks[i]. None means every row starts at column zero.
             row_starts: torch.Tensor | None = None
+            # Rows per dense logits call under the logits budget.
+            dense_rows = dense_logits_rows(chunk.max_seqlen_k)
+            use_deep_select = False
             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)
+            elif fused_mqa_topk and fused_mqa_topk_chunk(
+                chunk.max_seqlen_k,
+                k_quant.shape[0],
+                topk_tokens,
+                q_slice.shape[0],
+                dense_rows,
+            ):
+                fused_mqa_topk_prefill(
+                    q_slice,
+                    k_quant,
+                    fused_mqa_topk_key_scales(k_scale),
+                    weights[row_start:row_end],
+                    cu_seqlen_ks,
+                    cu_seqlen_ke,
+                    topk_indices,
+                )
+                logits = q_slice.new_empty((q_slice.shape[0], 0), dtype=torch.float32)
             else:
                 # DeepGEMM scalar-type tags (zero-copy): MXFP4 values → int8
                 # (kPackedFP4), scales → int32 squeezed to 1-D kv_sf / 2-D q_sf.
@@ -706,70 +790,96 @@ def sparse_attn_indexer(
                     k_scale_cast = k_scale.view(
                         torch.int32 if mqa_logits_use_mx_scales() else torch.float32
                     ).squeeze(-1)
-                if current_platform.is_xpu():
-                    if q_scale_slice is not None:
-                        raise RuntimeError("XPU fp8_mqa_logits does not support FP4 Q")
-                    logits = torch.ops.vllm.xpu_fp8_mqa_logits(
-                        q_slice_cast,
-                        k_quant_cast,
-                        k_scale_cast,
-                        weights[row_start:row_end],
-                        cu_seqlen_ks,
-                        cu_seqlen_ke,
-                    )
-                    row_starts = cu_seqlen_ks
-                    row_ends = cu_seqlen_ke
-                else:
-                    q_scale_slice, chunk_weights = mqa_logits_mx_inputs(
-                        q_slice_cast, q_scale_slice, weights[row_start:row_end]
-                    )
-                    row_ends = cu_seqlen_ke - cu_seqlen_ks
-                    logits = as_fp32_logits(
-                        fp8_fp4_mqa_logits(
-                            (q_slice_cast, q_scale_slice),
+                # Chunks planned for fused_mqa_topk are not bounded by the logits
+                # budget: split their rows (budget-planned chunks, the only ones
+                # under DCP, are one window, so the merge below sees all the
+                # rows).
+                for sub in range(0, q_slice.shape[0], dense_rows):
+                    rows = slice(sub, min(sub + dense_rows, q_slice.shape[0]))
+                    tokens = slice(row_start + rows.start, row_start + rows.stop)
+                    if current_platform.is_xpu():
+                        if q_scale_slice is not None:
+                            raise RuntimeError(
+                                "XPU fp8_mqa_logits does not support FP4 Q"
+                            )
+                        logits = torch.ops.vllm.xpu_fp8_mqa_logits(
+                            q_slice_cast[rows],
+                            k_quant_cast,
+                            k_scale_cast,
+                            weights[tokens],
+                            cu_seqlen_ks[rows],
+                            cu_seqlen_ke[rows],
+                        )
+                        row_starts = cu_seqlen_ks[rows]
+                        row_ends = cu_seqlen_ke[rows]
+                    else:
+                        q_scale_rows, chunk_weights = mqa_logits_mx_inputs(
+                            q_slice_cast[rows],
+                            q_scale_slice[rows] if q_scale_slice is not None else None,
+                            weights[tokens],
+                        )
+                        row_ends = cu_seqlen_ke[rows] - cu_seqlen_ks[rows]
+                        logits = fp8_fp4_mqa_logits(
+                            (q_slice_cast[rows], q_scale_rows),
                             (k_quant_cast, k_scale_cast),
                             chunk_weights,
-                            cu_seqlen_ks,
-                            cu_seqlen_ke,
+                            cu_seqlen_ks[rows],
+                            cu_seqlen_ke[rows],
                             chunk.max_seqlen_k,
-                        ),
-                        row_ends,
-                    )
-                num_rows = logits.shape[0]
-                if candidate_blocks is not None:
-                    # Two-level selection (v4.1): the candidate source
-                    # publishes its top blocks; later indexers mask their
-                    # scores to them. Both before the row top-k.
-                    chunk_candidates = candidate_blocks[
-                        chunk.token_start : chunk.token_end
-                    ]
-                    if candidate_write:
-                        _select_candidate_blocks(
+                        )
+                        # DeepSelect ranks the bf16 logits in place of their exact
+                        # fp32 copy; fp32 logits keep the per-row top-k.
+                        use_deep_select = (
+                            topk_backend == "deep_select"
+                            and logits.dtype == torch.bfloat16
+                            and dcp_world_size == 1
+                            and is_deep_select_supported(logits, topk_tokens)
+                        )
+                        if not use_deep_select:
+                            logits = as_fp32_logits(logits, row_ends)
+                    num_rows = logits.shape[0]
+                    if candidate_blocks is not None:
+                        # Two-level selection (v4.1): the candidate source
+                        # publishes its top blocks; later indexers mask their
+                        # scores to them. Both before the row top-k.
+                        chunk_candidates = candidate_blocks[tokens]
+                        if candidate_write:
+                            _select_candidate_blocks(
+                                logits,
+                                row_starts,
+                                row_ends,
+                                chunk_candidates.shape[1],
+                                candidate_block_size,
+                                chunk_candidates,
+                            )
+                        else:
+                            _apply_candidate_mask(
+                                logits,
+                                row_starts,
+                                row_ends,
+                                chunk_candidates,
+                                candidate_block_size,
+                            )
+                    if use_deep_select:
+                        deep_select_topk(
                             logits,
-                            row_starts,
-                            row_ends,
-                            chunk_candidates.shape[1],
-                            candidate_block_size,
-                            chunk_candidates,
+                            topk_tokens,
+                            end=row_ends,
+                            output_idx=topk_indices[rows],
                         )
                     else:
-                        _apply_candidate_mask(
+                        ops.top_k_per_row_prefill(
                             logits,
-                            row_starts,
+                            torch.zeros_like(row_ends)
+                            if row_starts is None
+                            else row_starts,
                             row_ends,
-                            chunk_candidates,
-                            candidate_block_size,
+                            topk_indices[rows],
+                            num_rows,
+                            logits.stride(0),
+                            logits.stride(1),
+                            topk_tokens,
                         )
-                ops.top_k_per_row_prefill(
-                    logits,
-                    torch.zeros_like(row_ends) if row_starts is None else row_starts,
-                    row_ends,
-                    topk_indices,
-                    num_rows,
-                    logits.stride(0),
-                    logits.stride(1),
-                    topk_tokens,
-                )
 
             if deinterleave_idx is None:
                 # Under the PCP path the top-k already ran over the whole
@@ -783,6 +893,8 @@ def sparse_attn_indexer(
                     cp_kv_cache_interleave_size,
                     row_starts=row_starts,
                 )
+            # Free this chunk's logits before the next chunk allocates its own.
+            del logits
 
         if row_shard is not None:
             # Rows are already fully ranked, so this is a layout gather.
@@ -887,9 +999,11 @@ def sparse_attn_indexer(
                 max_model_len=max_model_len,
                 indices=decode_metadata.indices,
             )
-            # Keeping bf16 halves top-k/candidate reads; other top-k backends and
-            # DCP need fp32. SM100 decode is varlen: one context length per row.
-            if topk_backend != "deep_select" or dcp_world_size != 1:
+            # Keeping bf16 halves top-k/candidate reads. SM100 decode is varlen:
+            # one context length per row.
+            if dcp_world_size != 1 or not get_indexer_topk(
+                topk_backend
+            ).takes_bf16_logits(logits, topk_tokens, logits.shape[0]):
                 logits = as_fp32_logits(logits, seq_lens.view(-1))
         num_rows = logits.shape[0]
         if candidate_blocks is not None:
@@ -910,6 +1024,11 @@ def sparse_attn_indexer(
                     candidate_block_size,
                     decode_candidates,
                     row_repeat,
+                    score_dtype=(
+                        torch.float32
+                        if topk_backend in ("auto", "persistent")
+                        else None
+                    ),
                 )
             else:
                 _apply_candidate_mask(
diff --git a/vllm/model_executor/layers/utils.py b/vllm/model_executor/layers/utils.py
index ef4594e11374..88563c743d05 100644
--- a/vllm/model_executor/layers/utils.py
+++ b/vllm/model_executor/layers/utils.py
@@ -328,6 +328,104 @@ def maybe_pad_weight_avoid_cache_aliasing(weight: torch.Tensor) -> torch.Tensor:
     return torch.nn.functional.pad(weight, (0, pad))[..., :-pad]
 
 
+_skinny_guard_warned: set[tuple[tuple[int, ...], int]] = set()
+
+# Cap used when the L3 size cannot be read: the MI355X value, the part
+# this guard was measured on (see envs.py for the cap).
+_SKINNY_GEMM_FALLBACK_WEIGHT_BYTES = 256 * 1024 * 1024
+
+
+def _amd_l3_cache_bytes(cache_info: list | None = None) -> int | None:
+    """The GPU's L3 (memory-side) cache size in bytes, or None.
+
+    The skinny-GEMM weight cap tracks this cache: past it a weight cannot
+    stay resident, and the skinny kernel's cost at n activation rows (~n
+    weight passes) pays HBM round trips instead. The device-properties
+    API exposes only the far smaller L2 (4 MiB on gfx950), so query the
+    L3 through AMD's amdsmi library: 256 MiB on MI355X, 128 MiB on
+    MI300X-class parts. `cache_info` (an amdsmi cache list) lets tests
+    drive the parse; the default queries the library.
+    """
+    if cache_info is None:
+        try:
+            import amdsmi
+
+            amdsmi.amdsmi_init()
+            try:
+                handles = amdsmi.amdsmi_get_processor_handles()
+                cache_info = amdsmi.amdsmi_get_gpu_cache_info(handles[0])["cache"]
+            finally:
+                amdsmi.amdsmi_shut_down()
+        except Exception:  # noqa: BLE001 -- optional lib, any failure -> fallback
+            return None
+    for cache in cache_info:
+        if cache.get("cache_level") == 3 and "DATA_CACHE" in cache.get(
+            "cache_properties", []
+        ):
+            return cache["cache_size"] * 1024  # amdsmi reports KB
+    return None
+
+
+@functools.lru_cache(maxsize=1)
+def _skinny_gemm_max_weight_bytes() -> int:
+    """Effective skinny-GEMM weight cap in bytes, evaluated once.
+
+    VLLM_SKINNY_GEMM_MAX_WEIGHT_BYTES > 0 overrides; 0 (the default)
+    derives the cap from the hardware L3 (memory-side) cache size,
+    falling back to _SKINNY_GEMM_FALLBACK_WEIGHT_BYTES where that cannot
+    be read. lmhead_skinny_guard_permits runs per decode step, so this is
+    computed on the first call only (env access in hot paths: #17067).
+    Tests that change the env var or mock the hardware read must call
+    cache_clear().
+    """
+    env_bytes = envs.VLLM_SKINNY_GEMM_MAX_WEIGHT_BYTES
+    if env_bytes > 0:
+        return env_bytes
+    return _amd_l3_cache_bytes() or _SKINNY_GEMM_FALLBACK_WEIGHT_BYTES
+
+
+def lmhead_skinny_guard_permits(weight: torch.Tensor, n: int) -> bool:
+    """Should this weight still take the wvSplitK skinny path at n?
+
+    wvSplitK_hf_sml's cost at n activation rows scales as ~n passes over
+    the weight (MI355X measurements in the PR description), so past the
+    last-level (memory-side) cache its effective bandwidth collapses:
+    every pass beyond the first is an HBM round trip. Weights above
+    VLLM_SKINNY_GEMM_MAX_WEIGHT_BYTES (default: the memory-side cache
+    size, read from the hardware once at first use, 256 MiB on MI355X;
+    see envs.py for what the cap controls and how to tune it) therefore
+    skip the skinny path at 1 < n <= 5 and fall through to the generic
+    GEMM.
+
+    n == 1 always keeps wvSplitK: the n=1 kernel variant shows no
+    pathology on large weights, and its fall-through would be a numerics
+    change with no speed win, so single-row calls stay on the stock path.
+
+    Returns True (keep wvSplitK) when n <= 1 or the weight is under the
+    cap; otherwise returns False after emitting a census WARNING (once per
+    (shape, n) per process) naming the redirected site.
+    """
+    if n <= 1:
+        return True
+    max_bytes = _skinny_gemm_max_weight_bytes()
+    if weight.nbytes <= max_bytes:
+        return True
+    key = (tuple(weight.shape), n)
+    if key not in _skinny_guard_warned:
+        _skinny_guard_warned.add(key)
+        logger.warning(
+            "lm_head skinny-GEMM guard engaged: weight %s (%d bytes) exceeds "
+            "the weight cap %d (VLLM_SKINNY_GEMM_MAX_WEIGHT_BYTES, or the "
+            "hardware memory-side cache size when unset) at n=%d; skipping "
+            "wvSplitK for the generic GEMM fall-through",
+            tuple(weight.shape),
+            weight.nbytes,
+            max_bytes,
+            n,
+        )
+    return False
+
+
 def rocm_unquantized_gemm_impl(
     x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None = None
 ) -> torch.Tensor:
@@ -396,6 +494,7 @@ def rocm_unquantized_gemm_impl(
         and (bias is None or bias.is_contiguous())
         and (m == 1 or m > 8)
         and 0 < n <= 5
+        and lmhead_skinny_guard_permits(weight, n)
     )
 
     use_llmm1 = (
@@ -420,7 +519,10 @@ def rocm_unquantized_gemm_impl(
         out = ops.LLMM1(weight, x_view, 4)
         return out.reshape(*x.shape[:-1], weight.shape[0])
 
-    if rocm_aiter_ops.is_tgemm_enabled():
+    # AITER's gfx950 FlyDSL A16W16 kernel addresses C with int32 offsets and
+    # reads A/B through 4 GiB buffer descriptors (ROCm/aiter#6853).
+    aiter_hgemm_unsafe = n * m >= 1 << 31 or max(n, m) * k * x.element_size() >= 1 << 32
+    if rocm_aiter_ops.is_tgemm_enabled() and not aiter_hgemm_unsafe:
         from aiter.tuned_gemm import tgemm
 
         return tgemm.mm(x, weight, bias)
@@ -457,6 +559,80 @@ def rocm_unquantized_gemm(
 )
 
 
+def cuda_skinny_unquantized_gemm_impl(
+    x: torch.Tensor, weight: torch.Tensor
+) -> torch.Tensor:
+    from vllm.model_executor.kernels.linear.cute_dsl.skinny_gemm import (
+        row_stride_ok,
+        shape_dynamic_skinny_gemm,
+    )
+    from vllm.model_executor.kernels.linear.cute_dsl.skinny_gemm_plans import (
+        get_skinny_gemm_plan,
+    )
+
+    plan = get_skinny_gemm_plan(weight.shape[0], weight.shape[1])
+    config = None if plan is None else plan.get(x.shape[0])
+    if (
+        config is not None
+        and x.dtype == torch.bfloat16
+        and weight.is_cuda
+        and weight.is_contiguous()
+        and row_stride_ok(x, config)
+    ):
+        return shape_dynamic_skinny_gemm(x, weight, config)
+    return torch.nn.functional.linear(x, weight)
+
+
+def cuda_skinny_unquantized_gemm_fake(
+    x: torch.Tensor, weight: torch.Tensor
+) -> torch.Tensor:
+    return x.new_empty((*x.shape[:-1], weight.shape[0]))
+
+
+direct_register_custom_op(
+    op_name="cuda_skinny_unquantized_gemm",
+    op_func=cuda_skinny_unquantized_gemm_impl,
+    fake_impl=cuda_skinny_unquantized_gemm_fake,
+)
+
+
+def maybe_use_skinny_gemm(layer: torch.nn.Module) -> None:
+    """Mark ``layer`` for the skinny GEMM if its weight shape has a plan."""
+    from vllm.model_executor.kernels.linear.cute_dsl.skinny_gemm_plans import (
+        get_skinny_gemm_plan,
+    )
+
+    weight = layer.weight
+    if weight.ndim != 2:
+        return
+
+    if (
+        not envs.VLLM_BATCH_INVARIANT
+        and weight.dtype == torch.bfloat16
+        and getattr(layer, "bias", None) is None
+        and get_skinny_gemm_plan(*weight.shape) is not None
+    ):
+        layer.use_skinny_gemm = True
+
+
+def cuda_unquantized_gemm(
+    layer: torch.nn.Module,
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor | None = None,
+) -> torch.Tensor:
+    # The op picks the config from the token count at runtime, so the choice
+    # stays out of the compiled graph.
+    if (
+        bias is None
+        and getattr(layer, "use_skinny_gemm", False)
+        and x.is_cuda
+        and x.dim() == 2
+    ):
+        return torch.ops.vllm.cuda_skinny_unquantized_gemm(x, weight)
+    return torch.nn.functional.linear(x, weight, bias)
+
+
 @functools.cache
 def warmup_rocm_skinny_gemm_workspaces(device: torch.device) -> None:
     """Eagerly allocate wvSplitKrc's per-device split-K workspace pool.
@@ -688,7 +864,7 @@ def dispatch_unquantized_gemm(
 
     backend_spec = _FLASHINFER_BF16_BACKENDS.get(linear_backend)
     if backend_spec is None:
-        return default_unquantized_gemm
+        return cuda_unquantized_gemm
 
     if not backend_spec.is_supported():
         logger.warning_once(
@@ -698,7 +874,7 @@ def dispatch_unquantized_gemm(
             linear_backend,
             backend_spec.flashinfer_backend,
         )
-        return default_unquantized_gemm
+        return cuda_unquantized_gemm
 
     logger.info_once(
         "Using FlashInfer %s for eligible unquantized BF16 GEMMs.",
diff --git a/vllm/model_executor/layers/vocab_parallel_embedding.py b/vllm/model_executor/layers/vocab_parallel_embedding.py
index 574aed3090fe..51671d9b0033 100644
--- a/vllm/model_executor/layers/vocab_parallel_embedding.py
+++ b/vllm/model_executor/layers/vocab_parallel_embedding.py
@@ -28,7 +28,10 @@
     resolve_quant_method,
 )
 from vllm.model_executor.layers.quantization.utils.fp8_utils import is_fp8
-from vllm.model_executor.layers.utils import dispatch_unquantized_gemm
+from vllm.model_executor.layers.utils import (
+    dispatch_unquantized_gemm,
+    maybe_use_skinny_gemm,
+)
 from vllm.model_executor.parameter import BasevLLMParameter
 from vllm.model_executor.utils import set_weight_attrs
 from vllm.platforms import current_platform
@@ -69,6 +72,8 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
             from vllm.model_executor.layers.utils import dispatch_cpu_unquantized_gemm
 
             dispatch_cpu_unquantized_gemm(layer, remove_weight=False)
+        elif current_platform.is_cuda() and isinstance(layer, ParallelLMHead):
+            maybe_use_skinny_gemm(layer)
 
     def apply(
         self,
diff --git a/vllm/model_executor/model_loader/base_loader.py b/vllm/model_executor/model_loader/base_loader.py
index 65f64a33267b..b2e5e3344062 100644
--- a/vllm/model_executor/model_loader/base_loader.py
+++ b/vllm/model_executor/model_loader/base_loader.py
@@ -3,6 +3,7 @@
 import gc
 import time
 from abc import ABC, abstractmethod
+from collections.abc import Generator
 
 import torch
 import torch.nn as nn
@@ -59,6 +60,18 @@ def create_model(
         log_model_inspection(model)
         return model
 
+    def get_all_weights(
+        self,
+        model_config: ModelConfig,
+        model: nn.Module,
+    ) -> Generator[tuple[str, torch.Tensor], None, None]:
+        """Iterate over all weights by checkpoint name, for filtered
+        partial reloads (e.g. FT expert reload after scale_down)."""
+        raise NotImplementedError(
+            f"load_format={self.load_config.load_format!r} does not support "
+            "iterating weights by checkpoint name"
+        )
+
     @instrument(span_name="Load model")
     def load_model(
         self, vllm_config: VllmConfig, model_config: ModelConfig, prefix: str = ""
diff --git a/vllm/model_executor/model_loader/reload/__init__.py b/vllm/model_executor/model_loader/reload/__init__.py
index a1ac282d290f..0ab9d3dca6fa 100644
--- a/vllm/model_executor/model_loader/reload/__init__.py
+++ b/vllm/model_executor/model_loader/reload/__init__.py
@@ -21,6 +21,8 @@
     "initialize_layerwise_reload",
     "finalize_layerwise_processing",
     "finalize_layerwise_reload",
+    "finish_reload",
+    "start_reload",
     "set_torchao_reload_attrs",
     "support_quantized_model_reload_from_hp_weights",
 ]
@@ -31,6 +33,7 @@
     initialize_layerwise_reload,
     record_metadata_for_reloading,
 )
+from .lifecycle import finish_reload, start_reload
 from .torchao_decorator import (
     set_torchao_reload_attrs,
     support_quantized_model_reload_from_hp_weights,
diff --git a/vllm/model_executor/model_loader/reload/direct.py b/vllm/model_executor/model_loader/reload/direct.py
new file mode 100644
index 000000000000..5dda849bf0d2
--- /dev/null
+++ b/vllm/model_executor/model_loader/reload/direct.py
@@ -0,0 +1,237 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Direct weight reload (``VLLM_RELOAD_MODE=direct``).
+
+Cold start records each tensor's checkpoint layout before post-load processing
+(PWAL). ``begin`` views re-laid-out tensors back to that layout on their
+storage, the loaders write the checkpoint in place, and ``finish`` reruns PWAL,
+which must keep addresses (captured CUDA graphs) and refill derived tensors.
+Between ``begin`` and a successful ``finish`` the model must not run; there is
+no rollback, the next full update repairs it.
+
+Delivery contract (not checked here): an update delivers every checkpoint
+tensor of each module whose PWAL changes its layout or is not idempotent.
+"""
+
+from collections.abc import Iterator
+
+import torch
+
+import vllm.envs as envs
+from vllm.config import ModelConfig
+from vllm.logger import init_logger
+from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase
+from vllm.model_executor.utils import (
+    is_weights_pre_processed,
+    storage_view,
+)
+
+logger = init_logger(__name__)
+
+__all__ = [
+    "begin_direct_reload",
+    "direct_reload_enabled",
+    "finish_direct_reload",
+    "record_checkpoint_layout",
+]
+
+
+# Loader metadata post-load processing may retag on a kept parameter (AWQ
+# repacks along another dim); the checkpoint loaders shard with it.
+_LOADER_ATTRS = (
+    "_input_dim",
+    "_output_dim",
+    "_packed_dim",
+    "_packed_factor",
+    "_marlin_tile_size",
+    "input_dim",
+    "output_dim",
+    "packed_dim",
+    "packed_factor",
+    "pack_factor",
+)
+
+
+def direct_reload_enabled() -> bool:
+    """Whether weight updates reload in place (``VLLM_RELOAD_MODE=direct``)."""
+    return envs.VLLM_RELOAD_MODE == "direct"
+
+
+def _layout(t: torch.Tensor) -> tuple:
+    return tuple(t.shape), t.stride(), t.dtype, t.storage_offset()
+
+
+def _loadable_tensors(model: torch.nn.Module) -> Iterator[tuple[str, torch.Tensor]]:
+    """Parameters and persistent buffers, each object once."""
+    seen: set[int] = set()
+    for module_name, module in model.named_modules():
+        persistent = (
+            (n, b)
+            for n, b in module._buffers.items()
+            if n not in module._non_persistent_buffers_set
+        )
+        for name, tensor in [*module._parameters.items(), *persistent]:
+            if tensor is None or id(tensor) in seen:
+                continue
+            seen.add(id(tensor))
+            yield (f"{module_name}.{name}" if module_name else name), tensor
+
+
+def _check_supported_modules(model: torch.nn.Module) -> None:
+    # config-level refusals: VllmConfig._verify_direct_reload_compat
+    from vllm.model_executor.layers.quantization.humming import (
+        HummingLinearMethod,
+        HummingMoEMethod,
+    )
+    from vllm.model_executor.layers.quantization.online.fp8 import OnlineLinearBase
+    from vllm.model_executor.layers.quantization.online.moe_base import (
+        OnlineMoEMethodBase,
+    )
+
+    for name, module in model.named_modules():
+        method = getattr(module, "quant_method", None)
+        # also without an online config (experts_int8); fused shared experts
+        # quantize staging parameters that post-load processing deletes
+        if isinstance(method, (OnlineLinearBase, OnlineMoEMethodBase)) or (
+            getattr(module, "shared_expert_online_loader", None) is not None
+        ):
+            raise RuntimeError(
+                f"direct reload: online quantization is not supported "
+                f"({name}: {type(method).__name__})"
+            )
+        if isinstance(method, (HummingLinearMethod, HummingMoEMethod)):
+            # its post-load processing runs once per quant method (a guard on
+            # the method, not the module): a reload would keep checkpoint bytes
+            raise RuntimeError(f"direct reload: humming is not supported ({name})")
+
+
+def record_checkpoint_layout(model: torch.nn.Module) -> None:
+    """Cold start, before post-load processing: remember each tensor's
+    checkpoint layout and loader metadata, which tensors this model owns (those
+    no other model recorded first), and which parameter names hold checkpoint
+    tensors.
+    Runs once; reloads reuse the cold-start record."""
+    if hasattr(model, "_ckpt_param_names") or is_weights_pre_processed():
+        return  # pre-processed weights are already in runtime layout
+    _check_supported_modules(model)
+    # a model owns the tensors it records first; `begin` refuses a shared one
+    # (a target lm_head in a draft) whose post-load processing re-laid it out
+    model._ckpt_owner = object()  # tags owned tensors (ids of freed ones recur)
+    for _, t in _loadable_tensors(model):
+        if not t.is_meta and not hasattr(t, "_ckpt"):
+            attrs = {k: t.__dict__[k] for k in _LOADER_ATTRS if k in t.__dict__}
+            t._ckpt = (model._ckpt_owner, _layout(t), attrs)
+    # every name, aliases included: post-load processing may replace either
+    model._ckpt_param_names = frozenset(
+        name
+        for name, p in model.named_parameters(remove_duplicate=False)
+        if hasattr(p, "_ckpt")
+    )
+
+
+def _check_params_kept(model: torch.nn.Module) -> None:
+    """Raise if post-load processing registered a new parameter under a
+    checkpoint name: `begin` cannot view it back, and the loader would write
+    checkpoint bytes into its runtime layout."""
+    replaced = [
+        name
+        for name, p in model.named_parameters(remove_duplicate=False)
+        if name in model._ckpt_param_names and not hasattr(p, "_ckpt")
+    ]
+    if replaced:
+        raise RuntimeError(
+            f"direct reload: post-load processing replaced the parameter object "
+            f"of {len(replaced)} checkpoint tensors (e.g. {replaced[:3]}); keep "
+            "it with `replace_parameter` / `retire_parameter` instead"
+        )
+
+
+def _owned(model: torch.nn.Module, t: torch.Tensor) -> bool:
+    ckpt = getattr(t, "_ckpt", None)
+    return ckpt is not None and ckpt[0] is model._ckpt_owner
+
+
+def _checkpoint_views(
+    model: torch.nn.Module,
+) -> list[tuple[torch.Tensor, torch.Tensor]]:
+    """(tensor, view of its storage in its checkpoint layout) for every
+    tensor post-load processing re-laid out; raises if a storage cannot
+    hold its checkpoint layout."""
+    views: list[tuple[torch.Tensor, torch.Tensor]] = []
+    for name, t in _loadable_tensors(model):
+        ckpt = getattr(t, "_ckpt", None)
+        if ckpt is None:
+            continue  # no layout: created by post-load processing, derived
+        if _layout(t) == ckpt[1]:
+            continue
+        if ckpt[0] is not model._ckpt_owner:
+            # shared from another model (a target lm_head in a draft): this
+            # finish would reprocess bytes a draft update does not resend
+            raise RuntimeError(
+                f"direct reload: {name} is shared from another model and its "
+                "post-load processing changed its layout"
+            )
+        shape, stride, dtype, offset = ckpt[1]
+        storage = t.untyped_storage()
+        span = 0 if 0 in shape else 1 + sum((n - 1) * s for n, s in zip(shape, stride))
+        need = (offset + span) * dtype.itemsize
+        if need > storage.nbytes():
+            raise RuntimeError(
+                f"direct reload: {name} storage has {storage.nbytes()} bytes, "
+                f"its checkpoint layout needs {need}"
+            )
+        views.append((t, storage_view(storage, offset, shape, stride, dtype)))
+    return views
+
+
+@torch.no_grad()
+def begin_direct_reload(model: torch.nn.Module) -> None:
+    """Make every tensor writable by its checkpoint weight loader.
+
+    Everything that can refuse the reload is checked first and raises with
+    the model untouched. Once this returns, the model holds checkpoint-layout
+    bytes until a finish succeeds."""
+    if not hasattr(model, "_ckpt_param_names"):
+        raise RuntimeError("direct reload: no checkpoint layout was recorded")
+    _check_params_kept(model)
+    views = _checkpoint_views(model)
+
+    # nothing above modified the model
+    model._direct_reload_begun = True  # one finish per begin
+    for t, view in views:
+        t.data = view  # no zeroing: the full sync rewrites the checkpoint region
+    for _, t in _loadable_tensors(model):
+        if not _owned(model, t):
+            continue
+        for k in _LOADER_ATTRS:
+            t.__dict__.pop(k, None)
+        t.__dict__.update(t._ckpt[2])
+    logger.info(
+        "direct reload: %d tensors viewed back to checkpoint layout", len(views)
+    )
+
+
+@torch.no_grad()
+def finish_direct_reload(
+    model: torch.nn.Module, model_config: ModelConfig, device: torch.device
+) -> None:
+    """Rerun cold-start post-load processing."""
+    from vllm.model_executor.model_loader.utils import process_weights_after_loading
+
+    if not model.__dict__.pop("_direct_reload_begun", False):
+        # e.g. a retry after a failed finish, which may have processed some
+        # modules already: only a new full update repairs the model
+        raise RuntimeError(
+            "direct reload: finish without a begin; start a new full update"
+        )
+    if device.type != "cpu":
+        torch.accelerator.synchronize()  # engines may write on other streams
+
+    # One-shot "already processed" guards would make the rerun below a no-op
+    # and leave the freshly written checkpoint layout in place. Only quant
+    # methods' guards: other modules set it for other reasons.
+    for module in model.modules():
+        if isinstance(getattr(module, "quant_method", None), QuantizeMethodBase):
+            module.__dict__.pop("_already_called_process_weights_after_loading", None)
+
+    process_weights_after_loading(model, model_config, device)
diff --git a/vllm/model_executor/model_loader/reload/inplace.py b/vllm/model_executor/model_loader/reload/inplace.py
index dfe11ed4e745..4239e81487fb 100644
--- a/vllm/model_executor/model_loader/reload/inplace.py
+++ b/vllm/model_executor/model_loader/reload/inplace.py
@@ -11,7 +11,7 @@
 from vllm.logger import init_logger
 from vllm.model_executor.model_loader import get_model_loader
 
-from .layerwise import finalize_layerwise_reload, initialize_layerwise_reload
+from .lifecycle import finish_reload, start_reload
 
 logger = init_logger(__name__)
 
@@ -91,9 +91,10 @@ def reload_weights(
     logger.info_once("Reloading weights inplace...")
     if is_checkpoint_format:
         # load weights from checkpoint/ original model format
-        initialize_layerwise_reload(model)
+        device = torch.device(load_config.device or vllm_config.device_config.device)
+        start_reload(model)
         loaded_weights = model.load_weights(weights_iterator)
-        finalize_layerwise_reload(model, model_config)
+        finish_reload(model, model_config, device)
 
     else:
         # load weights from kernel format
diff --git a/vllm/model_executor/model_loader/reload/layerwise.py b/vllm/model_executor/model_loader/reload/layerwise.py
index df91522732fe..a3badca62eda 100644
--- a/vllm/model_executor/model_loader/reload/layerwise.py
+++ b/vllm/model_executor/model_loader/reload/layerwise.py
@@ -12,7 +12,10 @@
 from vllm.logger import init_logger
 from vllm.model_executor.layers.attention import is_deferred_attention_layer
 from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+from vllm.model_executor.model_loader.weight_utils import (
+    default_weight_loader,
+    finalize_model_weights,
+)
 
 from .meta import (
     SKIP_LOAD_TENSORS,
@@ -38,6 +41,7 @@
     "initialize_layerwise_reload",
     "finalize_layerwise_processing",
     "finalize_layerwise_reload",
+    "max_layer_reload_bytes",
 ]
 
 
@@ -79,8 +83,21 @@ def record_metadata_for_reloading(model: torch.nn.Module):
         info.restore_device = torch.get_default_device()
 
 
+def max_layer_reload_bytes(model: torch.nn.Module) -> int:
+    """Bytes of the largest layer in checkpoint format, which a layerwise reload
+    materializes on the device, one layer at a time, before processing it."""
+    sizes = (
+        sum(t.nbytes for tensors in info.restore_metadata for t in tensors.values())
+        for layer in model.modules()
+        if (info := LAYERWISE_INFO.get(layer)) is not None
+    )
+    return max(sizes, default=0)
+
+
 @torch.no_grad()
-def initialize_layerwise_reload(model: torch.nn.Module):
+def initialize_layerwise_reload(
+    model: torch.nn.Module, *, defer_processing: bool = False
+):
     """Set up layerwise weight loading with deferred processing.
 
     Must be called after `record_metadata_for_reloading`. This function:
@@ -93,6 +110,8 @@ def initialize_layerwise_reload(model: torch.nn.Module):
     2. Load all cached weights
     3. Run quantization processing if applicable
     4. Copy processed values back to original tensor storage
+
+    With defer_processing=True, only finalize_layerwise_reload triggers processing.
     """
     # disable torchao reloading to avoid infinite recursion
     model._original_do_torchao_reload = getattr(model, "_do_torchao_reload", False)
@@ -115,6 +134,7 @@ def initialize_layerwise_reload(model: torch.nn.Module):
 
         # Wrap weight loaders to buffer loading
         initialize_online_processing(layer)
+        info.defer_processing = defer_processing
 
 
 def initialize_online_processing(layer: torch.nn.Module):
@@ -194,7 +214,9 @@ def online_process_loader(*args, **kwargs):
         )
 
         # Do not online process attention layers, must wait until finalize
-        if is_deferred_attention_layer(layer):
+        # Preserve placement remapping until checkpoint iteration finishes:
+        # expert reload wrappers must survive trailing redundant mappings.
+        if info.defer_processing or is_deferred_attention_layer(layer):
             return ret
 
         # Log warnings allocating excessive buffers on device
@@ -270,7 +292,7 @@ def finalize_layerwise_processing(model: torch.nn.Module, model_config: ModelCon
         # if the created weight has extra padding elements which are not loaded
         # Having too many of these delayed layers can lead to excess memory usage
         # see Limitations(4)
-        elif info.load_numel > 0 and info.load_numel < info.load_numel_total:  # type: ignore[operator]
+        elif info.defer_processing or info.load_numel < info.load_numel_total:  # type: ignore[operator]
             logger.debug("%s: Delayed processing", layer.__class__.__name__)
             _layerwise_process(layer, info)
 
@@ -284,8 +306,9 @@ def finalize_layerwise_processing(model: torch.nn.Module, model_config: ModelCon
     LOADING_LAYERS.clear()
 
 
-def finalize_layerwise_reload(*args, **kwargs):
-    finalize_layerwise_processing(*args, **kwargs)
+def finalize_layerwise_reload(model: torch.nn.Module, model_config: ModelConfig):
+    finalize_layerwise_processing(model, model_config)
+    finalize_model_weights(model)
 
 
 def _finalize_attention_layer(
@@ -300,7 +323,7 @@ def _finalize_attention_layer(
         _reload_attention_scales(layer, info)
     else:
         _place_kernel_tensors(layer, info)
-    layer.process_weights_after_loading(model_config.dtype)
+    layer.process_weights_after_loading()
 
 
 def _reload_attention_scales(layer: torch.nn.Module, info: LayerReloadingInfo) -> None:
diff --git a/vllm/model_executor/model_loader/reload/lifecycle.py b/vllm/model_executor/model_loader/reload/lifecycle.py
new file mode 100644
index 000000000000..b47b09445983
--- /dev/null
+++ b/vllm/model_executor/model_loader/reload/lifecycle.py
@@ -0,0 +1,31 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Begin and complete a checkpoint-format weight update in the reload mode
+``VLLM_RELOAD_MODE`` selects (fixed for the process)."""
+
+import torch
+
+from vllm.config import ModelConfig
+
+from .direct import begin_direct_reload, direct_reload_enabled, finish_direct_reload
+from .layerwise import finalize_layerwise_reload, initialize_layerwise_reload
+
+__all__ = ["finish_reload", "start_reload"]
+
+
+def start_reload(model: torch.nn.Module) -> None:
+    """Begin a weight update in the configured mode."""
+    if direct_reload_enabled():
+        begin_direct_reload(model)
+    else:
+        initialize_layerwise_reload(model)
+
+
+def finish_reload(
+    model: torch.nn.Module, model_config: ModelConfig, device: torch.device
+) -> None:
+    """Complete the weight update `start_reload` began."""
+    if direct_reload_enabled():
+        finish_direct_reload(model, model_config, device)
+    else:
+        finalize_layerwise_reload(model, model_config)
diff --git a/vllm/model_executor/model_loader/reload/torchao_decorator.py b/vllm/model_executor/model_loader/reload/torchao_decorator.py
index 10ced70d31ce..fb8ba877d143 100644
--- a/vllm/model_executor/model_loader/reload/torchao_decorator.py
+++ b/vllm/model_executor/model_loader/reload/torchao_decorator.py
@@ -10,7 +10,7 @@
 from vllm.config import ModelConfig
 
 from .layerwise import (
-    finalize_layerwise_reload,
+    finalize_layerwise_processing,
     initialize_layerwise_reload,
 )
 
@@ -53,7 +53,7 @@ def patched_model_load_weights(
 
         initialize_layerwise_reload(model)
         loaded_weights = original_load_weights(self, *args, **kwargs)
-        finalize_layerwise_reload(model, model._model_config)
+        finalize_layerwise_processing(model, model._model_config)
 
         return loaded_weights
 
diff --git a/vllm/model_executor/model_loader/reload/types.py b/vllm/model_executor/model_loader/reload/types.py
index 320c8d3602d2..e41c4c2b0eb8 100644
--- a/vllm/model_executor/model_loader/reload/types.py
+++ b/vllm/model_executor/model_loader/reload/types.py
@@ -22,6 +22,7 @@ class LayerReloadingInfo:
     # track how many elements are ready for loading, used by `online_process_loader`
     load_numel: int = 0
     load_numel_total: int | None = None
+    defer_processing: bool = False
 
     # used by `online_process_loader` to buffer args and tensors until ready to load
     loaded_weights: list[tuple[str, BoundArguments]] = field(default_factory=list)
diff --git a/vllm/model_executor/model_loader/sharded_state_loader.py b/vllm/model_executor/model_loader/sharded_state_loader.py
index 9e8e6e9cfc75..0ca9fc10e8bb 100644
--- a/vllm/model_executor/model_loader/sharded_state_loader.py
+++ b/vllm/model_executor/model_loader/sharded_state_loader.py
@@ -191,7 +191,7 @@ def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None:
         model_weights = model_config.model
         if model_weights_override := model_config.model_weights:
             model_weights = model_weights_override
-        local_model_path = model_weights
+        local_model_path = self._prepare_weights(model_weights, model_config.revision)
 
         rendered_pattern = ShardedStateLoader._render_default_pattern(
             self.pattern, part="*"
diff --git a/vllm/model_executor/model_loader/tensorizer_loader.py b/vllm/model_executor/model_loader/tensorizer_loader.py
index ae74beb7ed5f..00ae07fd66f7 100644
--- a/vllm/model_executor/model_loader/tensorizer_loader.py
+++ b/vllm/model_executor/model_loader/tensorizer_loader.py
@@ -21,7 +21,6 @@
 )
 from vllm.model_executor.model_loader.utils import (
     get_model_architecture,
-    initialize_model,
 )
 from vllm.utils.torch_utils import set_default_torch_dtype
 
@@ -65,29 +64,6 @@ def _get_weights_iterator(
         tensorizer_args = self.tensorizer_config._construct_tensorizer_args()
         return tensorizer_weights_iterator(tensorizer_args)
 
-    def _load_model_serialized_cpu(
-        self,
-        vllm_config: VllmConfig,
-        prefix: str = "",
-    ) -> nn.Module:
-        """Load a serialized model with tensorizer to the CPU.
-
-        This is only necessary when the model isn't vLLM-tensorized (see
-        examples/features/tensorize_vllm_model.py) This should still
-        be faster than default HuggingFace loading, but will be slower than
-        loading a vLLM-tensorized model.
-        """
-        device_config = vllm_config.device_config
-        model_config = vllm_config.model_config
-        with set_default_torch_dtype(model_config.dtype):
-            with torch.device(device_config.device):
-                model = initialize_model(vllm_config=vllm_config, prefix=prefix)
-
-            model.load_weights(self._get_weights_iterator())
-            if hasattr(model, "process_weights_after_loading"):
-                model.process_weights_after_loading()
-        return model.eval()
-
     def download_model(self, model_config: ModelConfig) -> None:
         self.tensorizer_config.verify_with_model_config(model_config)
 
@@ -138,7 +114,7 @@ def load_model(
                     )
             self.load_weights(model, model_config)
             return model
-        return self._load_model_serialized_cpu(vllm_config=vllm_config, prefix=prefix)
+        return super().load_model(vllm_config, model_config, prefix)
 
     @staticmethod
     def save_model(
diff --git a/vllm/model_executor/model_loader/utils.py b/vllm/model_executor/model_loader/utils.py
index 8a236941002b..1ab7db8e8ba0 100644
--- a/vllm/model_executor/model_loader/utils.py
+++ b/vllm/model_executor/model_loader/utils.py
@@ -33,6 +33,7 @@
     is_draft_model_cacheable,
 )
 from vllm.model_executor.model_loader.weight_tying import maybe_retie_word_embeddings
+from vllm.model_executor.model_loader.weight_utils import finalize_model_weights
 from vllm.model_executor.models.interfaces import SupportsQuant
 from vllm.model_executor.utils import is_weights_pre_processed
 from vllm.tracing import instrument
@@ -151,9 +152,16 @@ def process_weights_after_loading(
     from vllm.model_executor.layers.quantization.online.moe_base import (
         OnlineMoEMethodBase,
     )
+    from vllm.model_executor.model_loader.reload.direct import (
+        direct_reload_enabled,
+        record_checkpoint_layout,
+    )
 
     maybe_retie_word_embeddings(model, model_config)
 
+    if direct_reload_enabled():
+        record_checkpoint_layout(model)  # cold start only; reloads reuse it
+
     for name, module in model.named_modules():
         quant_method = getattr(module, "quant_method", None)
         if isinstance(quant_method, QuantizeMethodBase):
@@ -183,6 +191,12 @@ def process_weights_after_loading(
                 else nullcontext()
             )
             with loading_context, timing_context:
+                # Quantize the staged fused shared experts into their slots.
+                shared_expert_loader = getattr(
+                    module, "shared_expert_online_loader", None
+                )
+                if shared_expert_loader is not None:
+                    shared_expert_loader.finalize(module)
                 quant_method.process_weights_after_loading(module)
             # process_weights_after_loading may swap in freshly-created
             # Parameters (e.g. FP8 requantization), which are stamped with the
@@ -200,14 +214,11 @@ def process_weights_after_loading(
     # weights.
     for _, module in model.named_modules():
         if is_deferred_attention_layer(module):
-            # TODO(lucas): see if there is a way to unify the signatures
-            # of process_weights_after_loading
             with device_loading_context(module, target_device):
-                module.process_weights_after_loading(model_config.dtype)
+                module.process_weights_after_loading()
 
-    # Model-level post-load hook, after the per-layer quant finalize.
-    if hasattr(model, "process_weights_after_loading"):
-        model.process_weights_after_loading()
+    # Every other module's own post-load hook, after the per-layer passes.
+    finalize_model_weights(model)
 
     # Needed for torchao model reloading via model.reload_weights
     # @kylesayrs @jerryzh168 this can be removed if callers move to `reload_weights`
diff --git a/vllm/model_executor/model_loader/weight_utils.py b/vllm/model_executor/model_loader/weight_utils.py
index c31d0595e015..f7b0fff68c6f 100644
--- a/vllm/model_executor/model_loader/weight_utils.py
+++ b/vllm/model_executor/model_loader/weight_utils.py
@@ -7,6 +7,7 @@
 import fnmatch
 import glob
 import hashlib
+import inspect
 import json
 import os
 import sys
@@ -50,6 +51,7 @@
 from vllm.platforms import current_platform
 from vllm.tracing import instrument
 from vllm.transformers_utils.repo_utils import hf_api, hf_fs
+from vllm.utils.cpu_resource_utils import get_available_memory
 from vllm.utils.import_utils import PlaceholderModule
 
 try:
@@ -966,27 +968,6 @@ def _get_checkpoints_size_bytes(files: list[str]) -> int:
     return sum(os.path.getsize(f) for f in files)
 
 
-def _get_available_ram_bytes() -> int:
-    """Return available RAM, honoring cgroup limits."""
-    import psutil
-
-    host_available = psutil.virtual_memory().available
-
-    from vllm.utils.cpu_resource_utils import (
-        get_cgroup_memory_limit,
-        get_cgroup_memory_usage,
-    )
-
-    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)
-    )
-    return min(host_available, cgroup_available)
-
-
 def _get_fs_type(files: list[str]) -> str:
     """Get the filesystem type of the first file in *files* (Linux only)."""
     if not files:
@@ -1136,7 +1117,7 @@ def safetensors_weights_iterator(
     fs_type = _get_fs_type(sorted_files)
     is_net_fs = fs_type in ("nfs", "nfs4", "lustre")
     total_bytes = _get_checkpoints_size_bytes(sorted_files)
-    avail_bytes = _get_available_ram_bytes()
+    avail_bytes = get_available_memory()
     ram_threshold_pct = 90
     fits_in_ram = total_bytes <= (ram_threshold_pct / 100.0) * avail_bytes
     fs_name = fs_type.upper() if fs_type else "unknown"
@@ -1293,10 +1274,13 @@ def runai_safetensors_weights_iterator(
             else "cpu"
         )
 
+        # runai-model-streamer<0.17.0 has no `owned` and reuses its buffers
+        owned = "owned" in inspect.signature(streamer.stream_files).parameters
         streamer.stream_files(
             hf_weights_files,
             device=device,
             is_distributed=is_distributed,
+            **({"owned": True} if owned else {}),
         )
         total_tensors = sum(
             len(tensors_meta)
@@ -1313,7 +1297,7 @@ def runai_safetensors_weights_iterator(
         )
 
         for name, tensor in tensor_iter:
-            yield name, tensor.clone()
+            yield name, tensor if owned else tensor.clone()
 
 
 def fastsafetensors_weights_iterator(
@@ -1526,6 +1510,42 @@ def _load_file(bin_file: str):
             del state
 
 
+def finalize_model_weights(model: torch.nn.Module) -> None:
+    """Run every module's own post-load hook after a full load or reload,
+    each module after all of its submodules.
+
+    Quantization methods, the submodules of quantized layers (their quant
+    method processes them) and attention layers are skipped: the loader runs
+    them in its earlier per-layer passes, which later passes may read.
+    """
+    from vllm.model_executor.layers.attention import is_deferred_attention_layer
+    from vllm.model_executor.layers.quantization.base_config import (
+        QuantizeMethodBase,
+    )
+
+    seen: set[torch.nn.Module] = set()
+
+    def children_first(module: torch.nn.Module) -> Iterable[torch.nn.Module]:
+        quantized = isinstance(
+            getattr(module, "quant_method", None), QuantizeMethodBase
+        )
+        # a module shared by several parents runs once, before all of them
+        for child in [] if quantized else module.children():
+            if child not in seen:
+                seen.add(child)
+                yield from children_first(child)
+        yield module
+
+    for module in children_first(model):
+        if isinstance(module, QuantizeMethodBase) or is_deferred_attention_layer(
+            module
+        ):
+            continue
+        hook = getattr(module, "process_weights_after_loading", None)
+        if callable(hook):
+            hook()
+
+
 def default_weight_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None:
     """Default weight loader."""
     try:
diff --git a/vllm/model_executor/models/AXK1.py b/vllm/model_executor/models/AXK1.py
index 6a12fc88b255..cf0828e9951d 100644
--- a/vllm/model_executor/models/AXK1.py
+++ b/vllm/model_executor/models/AXK1.py
@@ -49,7 +49,6 @@
     fused_moe_make_expert_params_mapping,
 )
 from vllm.model_executor.layers.fused_moe.utils import (
-    is_model_fused_shared_expert_compatible,
     resolve_layer_fused_shared_expert,
 )
 from vllm.model_executor.layers.layernorm import RMSNorm
@@ -84,6 +83,8 @@
 from .utils import (
     AutoWeightsLoader,
     PPMissingLayer,
+    extract_layer_index,
+    get_fused_shared_expert_layers,
     get_spec_layer_idx_from_weight_name,
     is_pp_missing_parameter,
     make_empty_intermediate_tensors_factory,
@@ -742,10 +743,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
             prefix=f"{prefix}.layers",
         )
 
-        self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible(
-            self.layers,
+        self.fused_shared_expert_layers = get_fused_shared_expert_layers(
+            self,
             AXK1MoE,
-            "mlp",
+            enable_eplb=vllm_config.parallel_config.enable_eplb,
         )
 
         if get_pp_group().is_last_rank:
@@ -839,10 +840,11 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
             ckpt_gate_proj_name="gate_proj",
             ckpt_down_proj_name="down_proj",
             ckpt_up_proj_name="up_proj",
+            # NOTE: global despite per-layer FSE; non-fused layers skip the extras.
             num_experts=self.config.n_routed_experts
             + (
                 (self.config.n_shared_experts or 0)
-                if self.is_fused_shared_expert_enabled
+                if self.fused_shared_expert_layers
                 else 0
             ),
             num_redundant_experts=self.num_redundant_experts,
@@ -858,8 +860,11 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
             if spec_layer is not None:
                 continue  # skip spec decode layers for main model
 
+            # Only parse the layer index when FSE is enabled.
             is_fusion_moe_shared_experts_layer = (
-                self.is_fused_shared_expert_enabled and ("mlp.shared_experts" in name)
+                "mlp.shared_experts" in name
+                and bool(rocm_aiter_ops.is_fusion_moe_shared_experts_enabled())
+                and extract_layer_index(name) in self.fused_shared_expert_layers
             )
 
             for param_name, weight_name, shard_id in stacked_params_mapping:
diff --git a/vllm/model_executor/models/adapters.py b/vllm/model_executor/models/adapters.py
index 03cd55cf1003..cd28af2d60fc 100644
--- a/vllm/model_executor/models/adapters.py
+++ b/vllm/model_executor/models/adapters.py
@@ -29,6 +29,24 @@
 
 logger = init_logger(__name__)
 
+# Load formats whose weight iterators yield tensors that are never reused
+_OWNED_TENSOR_LOAD_FORMATS = frozenset(
+    {
+        "auto",
+        "hf",
+        "safetensors",
+        "mistral",
+        "pt",
+        "npcache",
+        "instanttensor",
+        "tensorizer",
+        "sharded_state",
+        "modelexpress",
+        "runai_streamer",
+        "runai_streamer_sharded",
+    }
+)
+
 _GENERATE_SUFFIXES = [
     "ForCausalLM",
     "ForConditionalGeneration",
@@ -190,9 +208,14 @@ def _load_pooling_model_weights(
             target_prefix = ""
 
             seen_weights = list[tuple[str, torch.Tensor]]()
+            owned = (
+                self.vllm_config.load_config.load_format in _OWNED_TENSOR_LOAD_FORMATS
+            )
             for name, loaded_weight in weights:
                 # Clone because the iterator may reuse the tensor buffer
-                seen_weights.append((name, loaded_weight.clone()))
+                seen_weights.append(
+                    (name, loaded_weight if owned else loaded_weight.clone())
+                )
 
                 lookup_name = name
                 if hf_to_vllm_mapper is not None:
@@ -226,8 +249,8 @@ def _load_pooling_model_weights(
                     self._get_name(),
                 )
 
-            # Lazy chain so buffer-reusing weight iterators (e.g.
-            # runai_streamer) are consumed one tensor at a time.
+            # Lazy chain so buffer-reusing weight iterators are consumed
+            # one tensor at a time.
             mapped_weights = (
                 (target_prefix + name, weight)
                 for name, weight in itertools.chain(seen_weights, weights)
diff --git a/vllm/model_executor/models/alice_ai.py b/vllm/model_executor/models/alice_ai.py
index a0c5f53ecf92..9dfdab545916 100644
--- a/vllm/model_executor/models/alice_ai.py
+++ b/vllm/model_executor/models/alice_ai.py
@@ -20,7 +20,6 @@
 )
 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,
 )
 from vllm.model_executor.layers.layernorm import GemmaRMSNorm
@@ -57,6 +56,7 @@
     AutoWeightsLoader,
     WeightsMapper,
     extract_layer_index,
+    get_fused_shared_expert_layers,
     make_layers,
     maybe_fuse_shared_experts,
     maybe_prefix,
@@ -500,8 +500,10 @@ def get_layer(prefix: str) -> AliceAIDecoderBlock:
             layer_fn=get_layer,
             prefix=f"{prefix}.layers",
         )
-        self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible(
-            layers=self.layers, moe_cls=AliceAISparseMoeBlock, moe_name="mlp"
+        self.fused_shared_expert_layers = get_fused_shared_expert_layers(
+            self,
+            AliceAISparseMoeBlock,
+            enable_eplb=vllm_config.parallel_config.enable_eplb,
         )
         self.norm = GemmaRMSNorm(
             hidden_size=config.hidden_size, eps=config.rms_norm_eps
@@ -567,7 +569,7 @@ def forward(
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         weights = maybe_fuse_shared_experts(
             weights=weights,
-            enabled=self.is_fused_shared_expert_enabled,
+            fused_shared_expert_layers=self.fused_shared_expert_layers,
             n_routed_experts=self.config.num_experts,
             n_shared_experts=1,
             ckpt_prefix="mlp.shared_expert",
diff --git a/vllm/model_executor/models/bailing_moe_v3.py b/vllm/model_executor/models/bailing_moe_v3.py
index 5819c89c9cc2..cbb1e8bc8f3f 100644
--- a/vllm/model_executor/models/bailing_moe_v3.py
+++ b/vllm/model_executor/models/bailing_moe_v3.py
@@ -1100,9 +1100,8 @@ def _forward(
             assert g1_non_spec is not None and beta_non_spec is not None
             assert state_indices is not None
             assert has_initial_state is not None
-            zero_idx = state_indices[~has_initial_state]
-            recurrent_state[zero_idx] = 0
             initial_state = recurrent_state_active[state_indices].contiguous()
+            initial_state.masked_fill_(~has_initial_state[:, None, None, None], 0)
             out, last_state = chunk_kda(
                 q=q,
                 k=k,
diff --git a/vllm/model_executor/models/canary.py b/vllm/model_executor/models/canary.py
new file mode 100644
index 000000000000..681031ebbabc
--- /dev/null
+++ b/vllm/model_executor/models/canary.py
@@ -0,0 +1,652 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Inference-only NVIDIA Canary (FastConformer encoder + Transformer decoder).
+
+Targets the native Transformers checkpoint layout (e.g. ``nvidia/canary-1b-v2``)
+whose ``config.json`` has ``model_type == "canary"``.
+"""
+
+import math
+from collections.abc import Iterable, Mapping, Sequence
+from typing import ClassVar
+
+import regex as re
+import torch
+from torch import nn
+from transformers import BatchFeature, CanaryConfig, ParakeetFeatureExtractor
+from transformers import ParakeetEncoder as HFParakeetEncoder
+
+from vllm.compilation.decorators import support_torch_compile
+from vllm.config import ModelConfig, SpeechToTextConfig, VllmConfig
+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
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
+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.multimodal import MULTIMODAL_REGISTRY
+from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargsItems
+from vllm.multimodal.parse import (
+    AudioProcessorItems,
+    MultiModalDataItems,
+    MultiModalDataParser,
+)
+from vllm.multimodal.processing import (
+    BaseDummyInputsBuilder,
+    BaseProcessingInfo,
+    EncDecMultiModalProcessor,
+    PromptReplacement,
+    PromptUpdate,
+)
+from vllm.multimodal.processing.processor import HFMultiModalInputs
+from vllm.renderers import TokenizeParams
+from vllm.tokenizers import cached_tokenizer_from_config
+from vllm.utils.collection_utils import is_list_of
+from vllm.utils.gpu_sync_debug import gpu_sync_allowed
+from vllm.v1.attention.backend import AttentionType
+
+from .cohere_asr import CohereASRAttention, CohereASRCrossAttention, CohereASRMLP
+from .interfaces import (
+    MultiModalEmbeddings,
+    SupportsMultiModal,
+    SupportsTranscription,
+)
+from .utils import WeightsMapper, make_layers, maybe_prefix
+
+logger = init_logger(__name__)
+
+# See https://huggingface.co/nvidia/canary-1b-v2
+CANARY_SUPPORTED_LANGS = {
+    "bg": "Bulgarian",
+    "hr": "Croatian",
+    "cs": "Czech",
+    "da": "Danish",
+    "nl": "Dutch",
+    "en": "English",
+    "et": "Estonian",
+    "fi": "Finnish",
+    "fr": "French",
+    "de": "German",
+    "el": "Greek",
+    "hu": "Hungarian",
+    "it": "Italian",
+    "lv": "Latvian",
+    "lt": "Lithuanian",
+    "mt": "Maltese",
+    "pl": "Polish",
+    "pt": "Portuguese",
+    "ro": "Romanian",
+    "sk": "Slovak",
+    "sl": "Slovenian",
+    "es": "Spanish",
+    "sv": "Swedish",
+    "ru": "Russian",
+    "uk": "Ukrainian",
+}
+
+
+# ----- Decoder -----
+class CanaryPositionalEmbedding(nn.Module):
+    """Fixed sinusoidal table laid out as ``[sin | cos]`` and scaled by
+    ``1 / sqrt(hidden_size)``; matches ``CanaryPositionalEmbedding`` in
+    Transformers."""
+
+    def __init__(self, max_positions: int, hidden_size: int) -> None:
+        super().__init__()
+        half = hidden_size // 2
+        log_timescale_increment = math.log(10000.0) * 2 / hidden_size
+        inv_timescales = torch.exp(
+            -log_timescale_increment * torch.arange(half).float()
+        )
+        scaled_time = torch.arange(max_positions).float()[:, None] * inv_timescales
+        pos_emb = torch.cat([torch.sin(scaled_time), torch.cos(scaled_time)], dim=1)
+        self.register_buffer(
+            "pos_emb", pos_emb / math.sqrt(hidden_size), persistent=False
+        )
+
+    def forward(self, positions: torch.Tensor) -> torch.Tensor:
+        return self.pos_emb[positions]
+
+
+class CanaryDecoderLayer(nn.Module):
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+        config = vllm_config.model_config.hf_config.decoder_config
+        cache_config = vllm_config.cache_config
+        quant_config = vllm_config.quant_config
+        hidden_size = config.hidden_size
+
+        self.input_layernorm = nn.LayerNorm(hidden_size)
+        self.self_attn = CohereASRAttention(
+            embed_dim=hidden_size,
+            num_heads=config.num_attention_heads,
+            bias=config.attention_bias,
+            attn_type=AttentionType.DECODER,
+            cache_config=cache_config,
+            quant_config=quant_config,
+            prefix=f"{prefix}.self_attn",
+        )
+        self.post_attention_layernorm = nn.LayerNorm(hidden_size)
+        self.encoder_attn = CohereASRCrossAttention(
+            embed_dim=hidden_size,
+            num_heads=config.num_attention_heads,
+            bias=config.attention_bias,
+            cache_config=cache_config,
+            quant_config=quant_config,
+            prefix=f"{prefix}.encoder_attn",
+        )
+        self.final_layernorm = nn.LayerNorm(hidden_size)
+        self.mlp = CohereASRMLP(
+            embed_dim=hidden_size,
+            ffn_dim=config.intermediate_size,
+            act_fn=config.hidden_act,
+            quant_config=quant_config,
+            prefix=f"{prefix}.mlp",
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor | None,
+    ) -> torch.Tensor:
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        hidden_states = residual + self.self_attn(hidden_states=hidden_states)
+
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = residual + self.encoder_attn(
+            hidden_states=hidden_states,
+            encoder_hidden_states=encoder_hidden_states,
+        )
+
+        residual = hidden_states
+        hidden_states = self.final_layernorm(hidden_states)
+        hidden_states = residual + self.mlp(hidden_states)
+        return hidden_states
+
+
+@support_torch_compile(dynamic_arg_dims={"input_ids": 0, "positions": -1})
+class CanaryDecoder(nn.Module):
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+        config = vllm_config.model_config.hf_config.decoder_config
+        self.hidden_size = config.hidden_size
+
+        self.embed_tokens = VocabParallelEmbedding(
+            config.vocab_size, config.hidden_size
+        )
+        self.pos_emb = CanaryPositionalEmbedding(
+            config.max_position_embeddings, config.hidden_size
+        )
+        self.embedding_layernorm = nn.LayerNorm(config.hidden_size)
+        self.start_layer, self.end_layer, self.layers = make_layers(
+            config.num_hidden_layers,
+            lambda prefix: CanaryDecoderLayer(vllm_config=vllm_config, prefix=prefix),
+            prefix=f"{prefix}.layers",
+        )
+        self.norm = nn.LayerNorm(config.hidden_size)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        encoder_hidden_states: torch.Tensor | None,
+    ) -> torch.Tensor:
+        hidden_states = self.get_input_embeddings(input_ids, positions)
+        for layer in self.layers:
+            hidden_states = layer(hidden_states, encoder_hidden_states)
+        return self.norm(hidden_states)
+
+    def get_input_embeddings(
+        self, input_ids: torch.Tensor, positions: torch.Tensor
+    ) -> torch.Tensor:
+        inputs_embeds = self.embed_tokens(input_ids)
+        pos_emb = self.pos_emb(positions).to(inputs_embeds.dtype)
+        return self.embedding_layernorm(inputs_embeds + pos_emb)
+
+
+class CanaryModel(nn.Module):
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+        config: CanaryConfig = vllm_config.model_config.hf_config
+        # Reuses the Transformers FastConformer encoder as-is, like `parakeet.py`.
+        self.encoder = HFParakeetEncoder(config.encoder_config).to(
+            vllm_config.model_config.dtype
+        )
+        self.decoder = CanaryDecoder(
+            vllm_config=vllm_config, prefix=maybe_prefix(prefix, "decoder")
+        )
+        assert config.encoder_config.hidden_size == config.decoder_config.hidden_size
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        encoder_outputs: list[torch.Tensor],
+    ) -> torch.Tensor:
+        enc_states = torch.cat(encoder_outputs, dim=0) if encoder_outputs else None
+        return self.decoder(
+            input_ids=input_ids,
+            positions=positions,
+            encoder_hidden_states=enc_states,
+        )
+
+    def get_encoder_outputs(
+        self,
+        input_features: torch.Tensor,
+        feature_lengths: torch.Tensor,
+    ) -> list[torch.Tensor]:
+        # input_features: (batch, time, num_mel_bins)
+        max_len = input_features.shape[1]
+        attention_mask = (
+            torch.arange(max_len, device=input_features.device)[None, :]
+            < feature_lengths[:, None]
+        )
+        out = self.encoder(
+            input_features=input_features, attention_mask=attention_mask
+        ).last_hidden_state
+        output_lengths = self.encoder._get_subsampling_output_length(feature_lengths)
+        with gpu_sync_allowed():
+            lengths = output_lengths.tolist()
+        return [feat[:length] for feat, length in zip(out, lengths)]
+
+
+# ----- Multi-modal processing -----
+def _subsampled_length(num_frames: int, encoder_config) -> int:
+    """Mirror of ``ParakeetEncoder._get_subsampling_output_length``."""
+    kernel_size = encoder_config.subsampling_conv_kernel_size
+    stride = encoder_config.subsampling_conv_stride
+    add_pad = (kernel_size - 1) // 2 * 2 - kernel_size
+    length = num_frames
+    for _ in range(int(math.log2(encoder_config.subsampling_factor))):
+        length = math.floor((length + add_pad) / stride + 1.0)
+    return length
+
+
+class CanaryProcessingInfo(BaseProcessingInfo):
+    def get_hf_config(self) -> CanaryConfig:
+        return self.ctx.get_hf_config(CanaryConfig)
+
+    def get_default_tok_params(self) -> TokenizeParams:
+        # The decoder prompt is built from control tokens; never add BOS/EOS.
+        return super().get_default_tok_params().with_kwargs(add_special_tokens=False)
+
+    def get_feature_extractor(self, **kwargs: object) -> ParakeetFeatureExtractor:
+        feature_extractor = self.get_hf_processor(**kwargs).feature_extractor
+        assert isinstance(feature_extractor, ParakeetFeatureExtractor)
+        return feature_extractor
+
+    def get_data_parser(self) -> MultiModalDataParser:
+        return MultiModalDataParser(
+            target_sr=self.get_feature_extractor().sampling_rate,
+            target_channels=1,
+        )
+
+    @property
+    def skip_prompt_length_check(self) -> bool:
+        return True  # The encoder prompt is padded
+
+    def get_supported_mm_limits(self) -> Mapping[str, int | None]:
+        return {"audio": 1}
+
+    def get_num_audio_tokens(self, num_samples: int) -> int:
+        fe = self.get_feature_extractor()
+        num_frames = (num_samples + fe.n_fft // 2 * 2 - fe.n_fft) // fe.hop_length
+        return _subsampled_length(num_frames, self.get_hf_config().encoder_config)
+
+
+class CanaryDummyInputsBuilder(BaseDummyInputsBuilder[CanaryProcessingInfo]):
+    def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
+        return "<|startofcontext|>" * mm_counts.get("audio", 0)
+
+    def get_dummy_mm_data(
+        self,
+        seq_len: int,
+        mm_counts: Mapping[str, int],
+        mm_options: MultiModalDummyOptions,
+    ) -> MultiModalDataDict:
+        sampling_rate = self.info.get_feature_extractor().sampling_rate
+        audio_len = CanaryForConditionalGeneration.max_audio_clip_s * sampling_rate
+        return {
+            "audio": self._get_dummy_audios(
+                length=audio_len,
+                num_audios=mm_counts.get("audio", 0),
+                overrides=mm_options.get("audio"),
+            )
+        }
+
+
+class CanaryMultiModalProcessor(EncDecMultiModalProcessor[CanaryProcessingInfo]):
+    skip_decoder_start_token: bool = True
+
+    @property
+    def pad_dummy_encoder_prompt(self) -> bool:
+        return True
+
+    def create_encoder_prompt(
+        self,
+        prompt: list[int],
+        mm_items: MultiModalDataItems,
+    ) -> list[int]:
+        # The encoder only consumes audio features; this placeholder is
+        # expanded to the number of encoder output frames.
+        return [0]
+
+    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)
+        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
+            )
+        )
+
+    def _call_hf_processor(
+        self,
+        hf_data: Mapping[str, object],
+        hf_kwargs: Mapping[str, object],
+    ) -> BatchFeature:
+        # `CanaryProcessor` renames `input_ids` to `decoder_input_ids`, so call
+        # the feature extractor and tokenizer directly instead.
+        feature_extractor = self.info.get_feature_extractor(**hf_kwargs)
+        audios = hf_data.get("audio")
+        if audios is None:
+            result = BatchFeature()
+        else:
+            if not isinstance(audios, list):
+                audios = [audios]
+            # Extract one item at a time so features are never padded to the
+            # longest clip in the batch.
+            input_features = []
+            feature_lengths = []
+            for audio in audios:
+                features = feature_extractor(
+                    audio,
+                    sampling_rate=feature_extractor.sampling_rate,
+                    return_tensors="pt",
+                )
+                input_features.append(features["input_features"][0])
+                feature_lengths.append(features["attention_mask"][0].sum())
+            result = BatchFeature(
+                dict(
+                    input_features=input_features,
+                    feature_lengths=torch.stack(feature_lengths),
+                )
+            )
+
+        text = hf_data.get("text")
+        if text is not None:
+            assert isinstance(text, str)
+            token_ids = self.info.get_tokenizer().encode(text, add_special_tokens=False)
+            result["input_ids"] = torch.tensor([token_ids])
+        return result
+
+    def _get_mm_fields_config(
+        self,
+        hf_inputs: BatchFeature,
+        hf_processor_mm_kwargs: Mapping[str, object],
+    ) -> Mapping[str, MultiModalFieldConfig]:
+        return dict(
+            input_features=MultiModalFieldConfig.batched("audio"),
+            feature_lengths=MultiModalFieldConfig.batched("audio"),
+        )
+
+    def _get_prompt_updates(
+        self,
+        mm_items: MultiModalDataItems,
+        hf_processor_mm_kwargs: Mapping[str, object],
+        out_mm_kwargs: MultiModalKwargsItems,
+    ) -> Sequence[PromptUpdate]:
+        def get_replacement(item_idx: int) -> list[int]:
+            audios = mm_items.get_items("audio", AudioProcessorItems)
+            num_samples = audios.get_audio_length(item_idx)
+            return [0] * self.info.get_num_audio_tokens(num_samples)
+
+        return [
+            PromptReplacement(modality="audio", target=[0], replacement=get_replacement)
+        ]
+
+
+@MULTIMODAL_REGISTRY.register_processor(
+    CanaryMultiModalProcessor,
+    info=CanaryProcessingInfo,
+    dummy_inputs=CanaryDummyInputsBuilder,
+)
+class CanaryForConditionalGeneration(
+    nn.Module, SupportsTranscription, SupportsMultiModal
+):
+    packed_modules_mapping = {
+        "qkv_proj": ["q_proj", "k_proj", "v_proj"],
+        "kv_proj": ["k_proj", "v_proj"],
+    }
+
+    # Only the decoder is built from vLLM layers; encoder weights keep their
+    # Transformers names.
+    hf_to_vllm_mapper = WeightsMapper(
+        orig_to_new_regex={
+            re.compile(r"(decoder\.layers\.\d+\.\w+_attn)\.o_proj\."): (
+                r"\1.out_projection."
+            ),
+            re.compile(r"(decoder\.layers\.\d+\.mlp)\.fc1\."): r"\1.dense_in.",
+            re.compile(r"(decoder\.layers\.\d+\.mlp)\.fc2\."): r"\1.dense_out.",
+            re.compile(r"\.num_batches_tracked$"): None,
+        },
+    )
+
+    supported_languages = CANARY_SUPPORTED_LANGS
+    supports_transcription_only = True
+    skip_warmup_audio_preprocessing = True
+    # TODO: canary-1b-v2 handles long-form audio with dynamic chunking in NeMo;
+    # confirm the right per-chunk limit for vLLM's chunker.
+    max_audio_clip_s: ClassVar[int] = 40
+
+    @classmethod
+    def validate_language(cls, language: str | None) -> str | None:
+        if language is None:
+            logger.warning(
+                "Defaulting to language='en'. If you wish to transcribe audio in "
+                "a different language, pass the `language` field in the "
+                "TranscriptionRequest."
+            )
+            language = "en"
+        # The decoder prompt needs one of the trained language tokens, so
+        # unlike the base class, do not fall back for other languages.
+        if language not in cls.supported_languages:
+            raise ValueError(
+                f"Unsupported language: {language!r}. Canary supports "
+                f"{list(cls.supported_languages)}."
+            )
+        return language
+
+    @classmethod
+    def get_decoder_prompt_tokens(
+        cls, source_language: str, target_language: str, punctuation: bool = True
+    ) -> list[str]:
+        """Same token sequence as the checkpoint's ``chat_template.jinja``."""
+        return [
+            "<|startofcontext|>",
+            "<|startoftranscript|>",
+            "<|emo:undefined|>",
+            f"<|{source_language}|>",
+            f"<|{target_language}|>",
+            "<|pnc|>" if punctuation else "<|nopnc|>",
+            "<|noitn|>",
+            "<|notimestamp|>",
+            "<|nodiarize|>",
+        ]
+
+    @classmethod
+    def get_generation_prompt(cls, stt_params: SpeechToTextParams) -> PromptType:
+        source_language = stt_params.language
+        if source_language is None:
+            raise ValueError("Language must be specified for Canary")
+        target_language = source_language
+        if stt_params.task_type == "translate":
+            target_language = stt_params.to_language or "en"
+            if target_language not in cls.supported_languages:
+                raise ValueError(
+                    f"Unsupported target language {target_language!r}; expected "
+                    f"one of {list(cls.supported_languages)}."
+                )
+            # Canary cannot detect the audio's language, and `validate_language`
+            # defaults a missing one to "en", so a translation without
+            # `language` would silently become English transcription.
+            if source_language == target_language:
+                raise ValueError(
+                    "Canary translation needs `language` (the audio's language) "
+                    f"to differ from `to_language` ({target_language!r}); it "
+                    "cannot detect the source language."
+                )
+            # canary-1b-v2 is trained to translate to or from English only.
+            if "en" not in (source_language, target_language):
+                raise ValueError(
+                    "Canary translates to or from English only; got "
+                    f"{source_language!r} -> {target_language!r}."
+                )
+
+        tokenizer = cached_tokenizer_from_config(stt_params.model_config)
+        tokens = cls.get_decoder_prompt_tokens(source_language, target_language)
+        prompt_token_ids = tokenizer.convert_tokens_to_ids(tokens)
+        unk_token_id = getattr(tokenizer, "unk_token_id", None)
+        if unk_token_id is not None and unk_token_id in prompt_token_ids:
+            raise ValueError("Failed to resolve the Canary decoder control tokens.")
+
+        audio = stt_params.audio
+        return TokensPrompt(
+            prompt_token_ids=prompt_token_ids,
+            multi_modal_data={"audio": (audio, stt_params.stt_config.sample_rate)},
+        )
+
+    @classmethod
+    def get_placeholder_str(cls, modality: str, i: int) -> str | None:
+        if modality.startswith("audio"):
+            return None
+        raise ValueError("Only audio modality is supported")
+
+    @classmethod
+    def get_speech_to_text_config(
+        cls, model_config: ModelConfig, task_type: str
+    ) -> SpeechToTextConfig:
+        return SpeechToTextConfig(
+            max_audio_clip_s=cls.max_audio_clip_s,
+            sample_rate=16_000,
+        )
+
+    @classmethod
+    def get_num_audio_tokens(
+        cls,
+        audio_duration_s: float,
+        stt_config: SpeechToTextConfig,
+        model_config: ModelConfig,
+    ) -> int | None:
+        # Matches ParakeetFeatureExtractor defaults (n_fft=512, hop_length=160).
+        num_frames = int(audio_duration_s * stt_config.sample_rate) // 160
+        return _subsampled_length(num_frames, model_config.hf_config.encoder_config)
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+        config: CanaryConfig = vllm_config.model_config.hf_config
+        quant_config = vllm_config.quant_config
+        self.config = config
+        self.dtype = vllm_config.model_config.dtype
+
+        with self._mark_composite_model(
+            vllm_config,
+            language_targets=CanaryDecoder,
+            tower_targets={"audio": HFParakeetEncoder},
+        ):
+            self.model = CanaryModel(
+                vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
+            )
+
+        decoder_config = config.decoder_config
+        self.proj_out = ParallelLMHead(
+            decoder_config.vocab_size,
+            decoder_config.hidden_size,
+            quant_config=quant_config,
+            bias=True,
+            prefix=maybe_prefix(prefix, "proj_out"),
+        )
+        self.logits_processor = LogitsProcessor(decoder_config.vocab_size)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        encoder_outputs: list[torch.Tensor] | None = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        return self.model(
+            input_ids=input_ids,
+            positions=positions,
+            encoder_outputs=encoder_outputs or [],
+        )
+
+    def get_language_model(self) -> torch.nn.Module:
+        return self.model.decoder
+
+    def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
+        input_features = kwargs.pop("input_features", None)
+        feature_lengths = kwargs.pop("feature_lengths", None)
+        if input_features is None:
+            return []
+        if not isinstance(feature_lengths, torch.Tensor):
+            raise ValueError("Audio feature lengths must be a tensor.")
+
+        if isinstance(input_features, torch.Tensor):
+            features = list(input_features)
+        elif is_list_of(input_features, torch.Tensor):
+            features = input_features
+        else:
+            raise ValueError(
+                f"Incorrect type of audio features: {type(input_features)}"
+            )
+        padded = torch.nn.utils.rnn.pad_sequence(
+            [f.to(self.dtype) for f in features], batch_first=True, padding_value=0.0
+        )
+        return self.model.get_encoder_outputs(padded, feature_lengths.reshape(-1))
+
+    def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.logits_processor(self.proj_out, hidden_states, self.proj_out.bias)
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        stacked_params_mapping = [
+            # (param_name, shard_name, shard_id)
+            (".self_attn.qkv_proj", ".self_attn.q_proj", "q"),
+            (".self_attn.qkv_proj", ".self_attn.k_proj", "k"),
+            (".self_attn.qkv_proj", ".self_attn.v_proj", "v"),
+            (".encoder_attn.kv_proj", ".encoder_attn.k_proj", "k"),
+            (".encoder_attn.kv_proj", ".encoder_attn.v_proj", "v"),
+        ]
+        params_dict = dict(self.named_parameters())
+        params_dict.update(self.named_buffers())
+        loaded_params: set[str] = set()
+
+        for name, loaded_weight in self.hf_to_vllm_mapper.apply(weights):
+            # The encoder is the Transformers module and keeps unfused q/k/v.
+            is_decoder = name.startswith("model.decoder.")
+            for param_name, shard_name, shard_id in stacked_params_mapping:
+                if not is_decoder or shard_name not in name:
+                    continue
+                name = name.replace(shard_name, param_name)
+                param = params_dict[name]
+                param.weight_loader(param, loaded_weight, shard_id)
+                break
+            else:
+                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
diff --git a/vllm/model_executor/models/cohere_asr.py b/vllm/model_executor/models/cohere_asr.py
index 298ccda9522b..a424ef93b4b6 100644
--- a/vllm/model_executor/models/cohere_asr.py
+++ b/vllm/model_executor/models/cohere_asr.py
@@ -1845,16 +1845,6 @@ 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)
 
-                # Convert buffer dtype to match loaded weight for pos_bias tensors
-                if "pos_bias" in name and param.dtype != loaded_weight.dtype:
-                    logger.info(
-                        "Converting buffer %s dtype from %s to %s for loading.",
-                        name,
-                        param.dtype,
-                        loaded_weight.dtype,
-                    )
-                    param.data = param.data.to(loaded_weight.dtype)
-
                 weight_loader(param, loaded_weight)
             loaded_params.add(name)
         return loaded_params
diff --git a/vllm/model_executor/models/cohere_eagle.py b/vllm/model_executor/models/cohere_eagle.py
index a8b540bbcf5c..a984a0bf36ab 100644
--- a/vllm/model_executor/models/cohere_eagle.py
+++ b/vllm/model_executor/models/cohere_eagle.py
@@ -137,15 +137,16 @@ def forward(
 
 
 class EagleCohereForCausalLM(CohereForCausalLM):
+    # Flags checked by the speculative proposer to decide whether to share
+    # embed_tokens / lm_head with the target model. Cohere EAGLE checkpoints
+    # use tied embeddings so these weights are absent from the draft file.
+    has_own_embed_tokens = False
+    has_own_lm_head = False
+
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         nn.Module.__init__(self)
         assert vllm_config.speculative_config is not None
         self.config = vllm_config.speculative_config.draft_model_config.hf_config
-        # Flags checked by the speculative proposer to decide whether to share
-        # embed_tokens / lm_head with the target model. Cohere EAGLE checkpoints
-        # use tied embeddings so these weights are absent from the draft file.
-        self.has_own_embed_tokens = False
-        self.has_own_lm_head = False
         if self.config.tie_word_embeddings:
             self.hf_to_vllm_mapper = self.hf_to_vllm_mapper | WeightsMapper(
                 orig_to_new_prefix={"model.embed_tokens.": None}
diff --git a/vllm/model_executor/models/config.py b/vllm/model_executor/models/config.py
index 40d00345a9ca..57411849e2be 100644
--- a/vllm/model_executor/models/config.py
+++ b/vllm/model_executor/models/config.py
@@ -56,18 +56,11 @@ def verify_and_update_config(vllm_config: "VllmConfig") -> None:
             )
 
         speculative_config = vllm_config.speculative_config
-        speculative_method = (
-            speculative_config.method if speculative_config is not None else None
-        )
-        if speculative_method in {
-            "dflash",
-            "dspark",
-            "eagle3",
-            "extract_hidden_states",
-        }:
+        if speculative_config is not None and speculative_config.uses_aux_hidden_states:
             raise NotImplementedError(
                 "AliceAI does not support speculative method "
-                f"{speculative_method!r} because it requires auxiliary hidden states"
+                f"{speculative_config.method!r} because it requires auxiliary "
+                "hidden states"
             )
 
 
@@ -86,6 +79,30 @@ def verify_and_update_config(cls, vllm_config: "VllmConfig") -> None:
             logger.info("Using bfloat16 kv-cache for DeepSeekV3.2")
 
 
+class MiniMaxM3SparseConfig(VerifyAndUpdateConfig):
+    @staticmethod
+    def verify_and_update_config(vllm_config: "VllmConfig") -> None:
+        from vllm.v1.kv_cache_interface import get_kv_quant_mode
+
+        if not get_kv_quant_mode(vllm_config.cache_config.cache_dtype).is_nvfp4:
+            return
+        # The NVFP4 KV attention kernels (MSA's CUTLASS Q8KV4 prefill) need 16
+        # query heads per KV head on each rank; fail before the workers start.
+        text_config = vllm_config.model_config.hf_text_config
+        tp_size = vllm_config.parallel_config.tensor_parallel_size
+        num_heads = text_config.num_attention_heads // tp_size
+        num_kv_heads = max(1, text_config.num_key_value_heads // tp_size)
+        if num_heads != 16 * num_kv_heads:
+            raise ValueError(
+                "MiniMax-M3 does not support an NVFP4 KV cache with "
+                f"tensor_parallel_size={tp_size} yet: each rank holds {num_heads} "
+                f"query heads over {num_kv_heads} KV head(s), and the NVFP4 "
+                "attention kernels need 16 per KV head. Use a smaller "
+                "tensor_parallel_size (with data/expert parallelism) or another "
+                "kv_cache_dtype."
+            )
+
+
 class GlmMoeDsaForCausalLM(VerifyAndUpdateConfig):
     @staticmethod
     def verify_and_update_config(vllm_config: "VllmConfig") -> None:
@@ -99,12 +116,10 @@ def verify_and_update_config(vllm_config: "VllmConfig") -> None:
         ):
             cache_config.cache_dtype = "fp8_e4m3"
             logger.info_once("Using fp8 kv-cache for GlmMoeDsaForCausalLM on SM10x")
-            # MTP shares the target's KV cache layout; a separate draft model
-            # must not inherit a default it was never validated with.
             spec_config = vllm_config.speculative_config
             if (
                 spec_config is not None
-                and spec_config.method != "mtp"
+                and spec_config.method not in ("mtp", "dspark", "dflash")
                 and spec_config.kv_cache_dtype is None
             ):
                 spec_config.kv_cache_dtype = "auto"
@@ -838,7 +853,7 @@ def verify_and_update_config(cls, vllm_config: "VllmConfig") -> None:
                 cache_config.mamba_block_size = model_config.max_model_len
 
 
-class MiniMaxM3SparseForConditionalGenerationConfig(VerifyAndUpdateConfig):
+class MiniMaxM3SparseForConditionalGenerationConfig(MiniMaxM3SparseConfig):
     INDEX_CACHE_KEYS = ("use_index_cache", "index_topk_freq", "index_skip_topk_offset")
 
     @classmethod
@@ -906,6 +921,15 @@ def verify_and_update_model_config(model_config: "ModelConfig") -> None:
             video_kwargs.setdefault("video_backend", "nemotron_vl")
 
 
+class Nemotron3_5TranscribeForConditionalGenerationConfig(VerifyAndUpdateConfig):
+    @classmethod
+    def verify_and_update_config(cls, vllm_config: "VllmConfig") -> None:
+        NemotronHForCausalLMConfig.update_mamba_ssm_cache_dtype(
+            cache_config=vllm_config.cache_config,
+            hf_config=vllm_config.model_config.hf_config.text_config,
+        )
+
+
 class NomicBertModelConfig(VerifyAndUpdateConfig):
     @staticmethod
     def verify_and_update_model_config(model_config: "ModelConfig") -> None:
@@ -1286,6 +1310,7 @@ def verify_and_update_config(vllm_config: "VllmConfig") -> None:
     "LlamaBidirectionalModel": LlamaBidirectionalConfig,
     "LlamaNemotronVLForSequenceClassification": LlamaNemotronVLConfig,
     "LlamaNemotronVLModel": LlamaNemotronVLConfig,
+    "MiniMaxM3SparseForCausalLM": MiniMaxM3SparseConfig,
     "Mamba2ForCausalLM": MambaModelConfig,
     "MambaForCausalLM": MambaModelConfig,
     "MiniMaxM3SparseForConditionalGeneration": (
@@ -1294,6 +1319,12 @@ def verify_and_update_config(vllm_config: "VllmConfig") -> None:
     "NemotronHForCausalLM": NemotronHForCausalLMConfig,
     "NemotronHPuzzleForCausalLM": NemotronHForCausalLMConfig,
     "NemotronH_Nano_VL_V2": NemotronHNanoVLV2Config,
+    "NemotronH_Nano_Omni_Reasoning_V3": NemotronHNanoVLV2Config,
+    "NemotronH_Super_Omni_Reasoning_V3": NemotronHNanoVLV2Config,
+    "NemotronH_Omni_Reasoning_V3": NemotronHNanoVLV2Config,
+    "Nemotron3_5TranscribeForConditionalGeneration": (
+        Nemotron3_5TranscribeForConditionalGenerationConfig
+    ),
     "NomicBertModel": NomicBertModelConfig,
     "Qwen2ForProcessRewardModel": Qwen2ForProcessRewardModelConfig,
     "Qwen2ForRewardModel": Qwen2ForRewardModelConfig,
diff --git a/vllm/model_executor/models/deepseek_mtp.py b/vllm/model_executor/models/deepseek_mtp.py
index 4f6125a2279b..c870660935bb 100644
--- a/vllm/model_executor/models/deepseek_mtp.py
+++ b/vllm/model_executor/models/deepseek_mtp.py
@@ -7,15 +7,13 @@
 import torch.nn as nn
 from transformers import PreTrainedConfig
 
+from vllm._aiter_ops import rocm_aiter_ops
 from vllm.compilation.decorators import support_torch_compile
 from vllm.config import VllmConfig
 from vllm.distributed import tensor_model_parallel_all_gather
 from vllm.model_executor.layers.fused_moe import (
     fused_moe_make_expert_params_mapping,
 )
-from vllm.model_executor.layers.fused_moe.utils import (
-    is_model_fused_shared_expert_compatible,
-)
 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
@@ -41,6 +39,8 @@
 )
 from .interfaces import SupportsPP
 from .utils import (
+    extract_layer_index,
+    get_fused_shared_expert_layers,
     get_pp_missing_layer_names,
     get_spec_layer_idx_from_weight_name,
     make_empty_intermediate_tensors_factory,
@@ -242,6 +242,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()
         self.config = vllm_config.model_config.hf_config
         self.quant_config = vllm_config.quant_config
+        self.vllm_config = vllm_config
         self.model = DeepSeekMultiTokenPredictor(
             vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
         )
@@ -267,10 +268,10 @@ def set_moe_parameters(self):
                 self.moe_mlp_layers.append(layer.mlp)
                 self.moe_layers.append(layer.mlp.experts)
         self.extract_moe_parameters(example_moe)
-        self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible(
-            self.model.layers.values(),
+        self.fused_shared_expert_layers = get_fused_shared_expert_layers(
+            self,
             DeepseekV2MoE,
-            "mtp_block.mlp",
+            enable_eplb=self.vllm_config.parallel_config.enable_eplb,
         )
 
     def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
@@ -316,17 +317,15 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         ]
         stacked_params_mapping.extend(indexer_fused_mapping)
 
+        fused_shared_expert_layers = self.fused_shared_expert_layers
         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",
+            # NOTE: global despite per-layer FSE; non-fused layers skip the extras.
             num_experts=self.config.n_routed_experts
-            + (
-                self.config.n_shared_experts
-                if self.is_fused_shared_expert_enabled
-                else 0
-            ),
+            + (self.config.n_shared_experts if fused_shared_expert_layers else 0),
         )
 
         pp_missing_layer_names = get_pp_missing_layer_names(self)
@@ -353,8 +352,11 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
                     weight_loader(param, loaded_weight)
                     loaded_params.add(name)
                 continue
+            # Only parse the layer index when FSE is enabled.
             is_fusion_moe_shared_experts_layer = (
-                self.is_fused_shared_expert_enabled and ("mlp.shared_experts" in name)
+                "mlp.shared_experts" in name
+                and bool(rocm_aiter_ops.is_fusion_moe_shared_experts_enabled())
+                and extract_layer_index(name) in fused_shared_expert_layers
             )
             name = self._rewrite_spec_layer_name(spec_layer, name)
 
diff --git a/vllm/model_executor/models/deepseek_ocr.py b/vllm/model_executor/models/deepseek_ocr.py
index 3aef695d0608..d137030a337a 100644
--- a/vllm/model_executor/models/deepseek_ocr.py
+++ b/vllm/model_executor/models/deepseek_ocr.py
@@ -499,7 +499,7 @@ def _encode_global_features(self, image_tensor: torch.Tensor) -> torch.Tensor:
     def _encode_local_features(
         self, patches: torch.Tensor, crop_shape: torch.Tensor
     ) -> torch.Tensor | None:
-        if torch.sum(patches).item() == 0:
+        if patches.numel() == 0:
             return None
 
         local_features_1 = self.sam_model(patches)
diff --git a/vllm/model_executor/models/deepseek_ocr2.py b/vllm/model_executor/models/deepseek_ocr2.py
index 18de7503c721..9139259909e6 100644
--- a/vllm/model_executor/models/deepseek_ocr2.py
+++ b/vllm/model_executor/models/deepseek_ocr2.py
@@ -343,7 +343,7 @@ def _encode_global_features(self, image_tensor: torch.Tensor) -> torch.Tensor:
         return features.view(-1, dim)
 
     def _encode_local_features(self, patches: torch.Tensor) -> torch.Tensor | None:
-        if torch.sum(patches).item() == 0:
+        if patches.numel() == 0:
             return None
 
         local_features = self.sam_model(patches)
diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py
index a1a919fae644..c3befbc1a3a5 100644
--- a/vllm/model_executor/models/deepseek_v2.py
+++ b/vllm/model_executor/models/deepseek_v2.py
@@ -59,7 +59,6 @@
     fused_moe_make_expert_params_mapping,
 )
 from vllm.model_executor.layers.fused_moe.utils import (
-    is_model_fused_shared_expert_compatible,
     resolve_layer_fused_shared_expert,
 )
 from vllm.model_executor.layers.layernorm import LayerNorm, RMSNorm
@@ -97,6 +96,7 @@
 from vllm.model_executor.models.utils import (
     AutoWeightsLoader,
     extract_layer_index,
+    get_fused_shared_expert_layers,
     sequence_parallel_chunk,
 )
 from vllm.platforms import current_platform
@@ -1476,10 +1476,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
             prefix=f"{prefix}.layers",
         )
 
-        self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible(
-            self.layers,
+        self.fused_shared_expert_layers = get_fused_shared_expert_layers(
+            self,
             DeepseekV2MoE,
-            "mlp",
+            enable_eplb=vllm_config.parallel_config.enable_eplb,
         )
 
         if get_pp_group().is_last_rank:
@@ -1629,12 +1629,9 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
             ckpt_gate_proj_name="gate_proj",
             ckpt_down_proj_name="down_proj",
             ckpt_up_proj_name="up_proj",
+            # NOTE: global despite per-layer FSE; non-fused layers skip the extras.
             num_experts=self.config.n_routed_experts
-            + (
-                self.config.n_shared_experts
-                if self.is_fused_shared_expert_enabled
-                else 0
-            ),
+            + (self.config.n_shared_experts if self.fused_shared_expert_layers else 0),
             num_redundant_experts=self.num_redundant_experts,
         )
 
@@ -1659,8 +1656,11 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
             ):
                 continue  # this layer has no indexer; drop its checkpoint weights
 
+            # Only parse the layer index when FSE is enabled.
             is_fusion_moe_shared_experts_layer = (
-                self.is_fused_shared_expert_enabled and ("mlp.shared_experts" in name)
+                "mlp.shared_experts" in name
+                and bool(rocm_aiter_ops.is_fusion_moe_shared_experts_enabled())
+                and extract_layer_index(name) in self.fused_shared_expert_layers
             )
 
             if _try_load_fp8_indexer_wk(
diff --git a/vllm/model_executor/models/diffusion_gemma.py b/vllm/model_executor/models/diffusion_gemma.py
index f2f49cf03297..7feeaf5d51d3 100644
--- a/vllm/model_executor/models/diffusion_gemma.py
+++ b/vllm/model_executor/models/diffusion_gemma.py
@@ -16,7 +16,7 @@
 
 from __future__ import annotations
 
-from collections.abc import Iterable, Sequence
+from collections.abc import Callable, Iterable, Sequence
 from types import SimpleNamespace
 from typing import Any
 
@@ -909,6 +909,13 @@ def __init__(
     def get_supported_generation_tasks(self):
         return ("generate",)
 
+    def compute_sample_logits(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return hidden_states
+
+    def dummy_sampler_run(self, sampler: Any, hidden_states: torch.Tensor) -> bool:
+        sampler.profile_run(hidden_states)
+        return True
+
     def custom_sampler(self, sampler: Any) -> tuple[Any, Any] | None:
         diffusion_config = self.vllm_config.diffusion_config
         gen = self.gen_config
@@ -931,6 +938,7 @@ def custom_sampler(self, sampler: Any) -> tuple[Any, Any] | None:
             diffusion_config=diffusion_config,
             vocab_size=self.model_config.get_vocab_size(),
             diffusion_states=self.diffusion_states,
+            compute_logits=self.model.compute_logits,
             t_min=gen["t_min"],
             t_max=gen["t_max"],
             entropy_bound=entropy_bound,
@@ -1143,6 +1151,7 @@ def __init__(
         vocab_size: int,
         diffusion_states: DiffusionGemmaRequestStates,
         *,
+        compute_logits: Callable[[torch.Tensor], torch.Tensor],
         confidence_threshold: float,
         t_min: float,
         t_max: float,
@@ -1154,6 +1163,7 @@ def __init__(
         tp_size: int = 1,
         tp_group_name: str = "",
     ):
+        self.compute_logits = compute_logits
         self.sampling_states = sampler.sampling_states
         self.logprob_token_ids_state = sampler.logprob_token_ids_state
         self.req_states = sampler.req_states
@@ -1262,6 +1272,24 @@ def penalties_state(self):
         # post_update treats None bin counts as "no penalty bookkeeping".
         return _NO_PENALTIES_STATE
 
+    @torch.inference_mode()
+    def profile_run(self, hidden_states: torch.Tensor) -> None:
+        CL = self.canvas_length
+        device = hidden_states.device
+        bounds = np.array([0, CL], dtype=np.int32)
+        dummy_batch = SimpleNamespace(
+            num_reqs=1,
+            num_draft_tokens=CL,
+            idx_mapping_np=np.zeros(1, dtype=np.int32),
+            idx_mapping=torch.zeros(1, dtype=torch.int64, device=device),
+            cu_num_logits_np=bounds,
+            query_start_loc_np=bounds,
+            query_start_loc=torch.tensor(bounds, device=device),
+        )
+        self.diffusion_states.add_request(0)
+        self(hidden_states.new_zeros((CL, hidden_states.shape[-1])), dummy_batch)
+        self.diffusion_states.add_request(0)
+
     # ------------------------------------------------------------------
     # Prefill
     # ------------------------------------------------------------------
@@ -1359,12 +1387,12 @@ def _build_output(
 
     def __call__(
         self,
-        logits: torch.Tensor,
+        hidden_states: torch.Tensor,
         input_batch: Any,
         draft_logits: torch.Tensor | None = None,
     ) -> SamplerOutput:
         num_reqs = input_batch.num_reqs
-        device = logits.device
+        device = hidden_states.device
 
         if input_batch.num_draft_tokens == 0:
             return self._handle_prefill(input_batch, device)
@@ -1398,26 +1426,7 @@ def __call__(
             valid_canvas_len_np.astype(np.int64), device=device
         )
 
-        # Per-request top_k/top_p, mirroring the AR sampler. Masked tokens
-        # become -inf and survive the temperature scaling in the compiled
-        # step, so Gumbel sampling, probs, and entropy all see the filtered
-        # distribution. The committed argmax (always the top-1 token) is
-        # unaffected; only the canvas exploration is constrained. Applied
-        # before canvas padding so phantom positions stay uniform.
-        if num_decode > 0:
-            top_k, top_p, has_topp_only_rows = self.sampling_states.get_top_k_top_p(
-                decode_slots.repeat_interleave(
-                    valid_canvas_len, output_size=int(valid_canvas_len_np.sum())
-                ),
-                decode_slots_np,
-                logits.shape[-1],
-            )
-            if top_k is not None or top_p is not None:
-                logits = apply_top_k_top_p(
-                    logits.float(), top_k, top_p, has_topp_only_rows=has_topp_only_rows
-                )
-
-        # Where each decode request's rows start in the flat logits. Tiles
+        # Where each decode request's rows start in the flat hidden states. 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])
@@ -1443,23 +1452,18 @@ def __call__(
         embed_weight = self.embed_weight
         vocab_size = self.vocab_size
         sc_vocab_start, sc_vocab_end = self.sc_vocab_start, self.sc_vocab_end
+        per_row = None
         if allowed is not None:
-            assert logits.shape[-1] == allowed.numel()
             embed_weight = self.embed_weight[allowed]
             sc_vocab_start, sc_vocab_end = 0, allowed.numel()
         elif states.constrained and num_decode > 0:
             # The decode slots do not share one set, so this step runs over
             # the full vocabulary. Mask each constrained request's rows to
-            # its own set so it reads the same as on the shared path. The
-            # helper masks a copy, since the runner owns `logits`: one
-            # full-vocab copy per step, as for top_k/top_p above.
+            # its own set so it reads the same as on the shared path.
             per_row = [
                 None if ids is None else states.allowed_tensor(ids)
                 for ids in (states.constrained.get(s) for s in decode_slots_np.tolist())
             ]
-            logits = _mask_rows_to_allowed(
-                logits, row_starts_np, valid_canvas_len_np, per_row
-            )
 
         slots_np = input_batch.idx_mapping_np[:num_reqs]
         max_num_logprobs = self.sampling_states.max_num_logprobs(slots_np)
@@ -1473,8 +1477,9 @@ def __call__(
         # 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.
+        # all slots) is the widest. Each tile projects its own logits, and 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 = torch.accelerator.get_memory_info()[0] if num_decode > 0 else 0
@@ -1484,10 +1489,11 @@ def __call__(
             run_end = run_start
             while run_end < num_decode and widths_np[order[run_end]] == W:
                 run_end += 1
-            # Transient [tile * W, vocab] tensors: the softmax in the embedding
-            # dtype for self-conditioning, plus the fp32 scaled logits when
-            # logprobs are wanted (pad for allocator overhead).
-            budget = max(1, int(free * 0.5) // max(W * self.vocab_size * 4 * 3, 1))
+            # Transient [tile * W, vocab] tensors: the projected fp32 logits and
+            # their top_k/top_p or allowed-set copy, the softmax in the
+            # embedding dtype for self-conditioning, plus the fp32 scaled
+            # logits when logprobs are wanted (pad for allocator overhead).
+            budget = max(1, int(free * 0.5) // max(W * self.vocab_size * 4 * 5, 1))
             for t0 in range(run_start, run_end, budget):
                 sel_np = order[t0 : min(t0 + budget, run_end)]
                 n = len(sel_np)
@@ -1496,21 +1502,54 @@ def __call__(
                 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())
+                phantom = None
                 if contiguous and tile_valid_np.min() == W:
                     r0 = int(row_starts_np[sel_np[0]])
-                    tile_logits = logits[r0 : r0 + n * W]
+                    tile_hidden = hidden_states[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.
+                    # filtering below 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
+                        hidden_states.shape[0] - 1
+                    )
+                    phantom = (ar.unsqueeze(0) >= tile_valid.unsqueeze(1)).reshape(
+                        -1, 1
+                    )
+                    tile_hidden = hidden_states[src.reshape(-1)].masked_fill_(
+                        phantom, 0
+                    )
+                tile_logits = self.compute_logits(tile_hidden)
+                if allowed is not None:
+                    assert tile_logits.shape[-1] == allowed.numel()
+
+                # Per-request top_k/top_p, mirroring the AR sampler. Masked
+                # tokens become -inf and survive the temperature scaling, so
+                # Gumbel sampling, probs, and entropy all see the filtered
+                # distribution. The committed argmax (always the top-1 token)
+                # is unaffected; only the canvas exploration is constrained.
+                top_k, top_p, has_topp_only_rows = self.sampling_states.get_top_k_top_p(
+                    tile_slots.repeat_interleave(W, output_size=n * W),
+                    decode_slots_np[sel_np],
+                    tile_logits.shape[-1],
+                )
+                if top_k is not None or top_p is not None:
+                    tile_logits = apply_top_k_top_p(
+                        tile_logits.float(),
+                        top_k,
+                        top_p,
+                        has_topp_only_rows=has_topp_only_rows,
                     )
-                    valid = ar.unsqueeze(0) < tile_valid.unsqueeze(1)
-                    tile_logits = logits[src.reshape(-1)].masked_fill_(
-                        ~valid.reshape(-1, 1), 0
+                    if phantom is not None:
+                        tile_logits.masked_fill_(phantom, 0)
+                if per_row is not None:
+                    tile_logits = _mask_rows_to_allowed(
+                        tile_logits,
+                        np.arange(n) * W,
+                        tile_valid_np,
+                        [per_row[i] for i in sel_np],
                     )
                 compute_sc = (
                     not states.single_step_slots
diff --git a/vllm/model_executor/models/expert_substitution.py b/vllm/model_executor/models/expert_substitution.py
index 0799bca8b7e0..44767a71d150 100644
--- a/vllm/model_executor/models/expert_substitution.py
+++ b/vllm/model_executor/models/expert_substitution.py
@@ -14,7 +14,7 @@
     decoder_layer_index,
     get_expert_substitution_spec,
 )
-from vllm.model_executor.model_loader.reload import finalize_layerwise_reload
+from vllm.model_executor.model_loader.reload import finalize_layerwise_processing
 
 _T = TypeVar("_T", bound=nn.Module)
 
@@ -69,7 +69,7 @@ def __init__(self, *args, **kwargs):
         def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
             if any(s.values.is_meta for _, s in self._expert_substitutions.values()):
                 # Put the previous weights back so the model stays usable.
-                finalize_layerwise_reload(self, model_config)
+                finalize_layerwise_processing(self, model_config)
                 raise NotImplementedError(
                     "expert substitution does not support layerwise weight reload"
                 )
diff --git a/vllm/model_executor/models/funaudiochat.py b/vllm/model_executor/models/funaudiochat.py
index 8e64fe03e22d..ac31c20d6b4c 100644
--- a/vllm/model_executor/models/funaudiochat.py
+++ b/vllm/model_executor/models/funaudiochat.py
@@ -532,7 +532,7 @@ def speech_tokenizer(self) -> TokenizersBackend:
             revision=self.ctx.model_config.revision,
         )
 
-    def get_feature_extractor(self) -> WhisperFeatureExtractor:
+    def get_feature_extractor(self, **kwargs: object) -> WhisperFeatureExtractor:
         return self.feature_extractor
 
     def get_speech_tokenizer(self) -> TokenizersBackend:
@@ -568,6 +568,26 @@ def get_audio_group_size(self) -> int:
         audio_cfg = getattr(cfg, "audio_config", None)
         return int(getattr(audio_cfg, "group_size", 5))
 
+    def get_max_speech_frames(self) -> int:
+        """Maximum 25 Hz speech frames admitted for one audio item.
+
+        The discrete encoder emits ``max_source_positions`` replacement tokens
+        after grouping frames by ``group_size`` (1500 * 5 = 7500 frames, 300s).
+        """
+        cfg = self.get_hf_config()
+        audio_cfg = getattr(cfg, "audio_config", None)
+        max_audio_tokens = int(getattr(audio_cfg, "max_source_positions", 1500))
+        return max(1, max_audio_tokens) * max(1, self.get_audio_group_size())
+
+    def get_max_audio_samples(self, sampling_rate: int) -> int:
+        """Waveform samples that fill the profiled per-item speech-frame budget."""
+        token_fps = int(self.token_fps)
+        target_num_frames = self.get_max_speech_frames()
+        return max(
+            1,
+            (target_num_frames * sampling_rate + token_fps - 1) // token_fps,
+        )
+
 
 class FunAudioChatDummyInputsBuilder(
     BaseDummyInputsBuilder[FunAudioChatProcessingInfo]
@@ -587,16 +607,7 @@ def get_dummy_mm_data(
 
         # Dummy inputs are used for profiling; construct the worst-case audio
         # length that maximizes the number of encoder tokens.
-        cfg = self.info.get_hf_config()
-        audio_cfg = getattr(cfg, "audio_config", None)
-        max_audio_tokens = int(getattr(audio_cfg, "max_source_positions", 1500))
-        group_size = self.info.get_audio_group_size()
-        token_fps = int(getattr(self.info, "token_fps", 25))
-        target_num_frames = max(1, max_audio_tokens) * max(1, group_size)
-        audio_len = max(
-            1,
-            (target_num_frames * sampling_rate + token_fps - 1) // token_fps,
-        )
+        audio_len = self.info.get_max_audio_samples(sampling_rate)
 
         return {
             "audio": self._get_dummy_audios(
@@ -638,6 +649,8 @@ def _apply_hf_processor_main(
         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)
+        max_samples = self.info.get_max_audio_samples(sr)
+        max_speech_frames = self.info.get_max_speech_frames()
 
         wavs: list[np.ndarray] = []
         speech_strs: list[str] = []
@@ -654,6 +667,15 @@ def _apply_hf_processor_main(
                     audio_np, (0, min_samples - audio_np.shape[0]), mode="constant"
                 )
 
+            if audio_np.shape[0] > max_samples:
+                duration_s = float(audio_np.shape[0]) / float(sr)
+                max_duration_s = float(max_samples) / float(sr)
+                raise ValueError(
+                    "Audio is too long for FunAudioChat: "
+                    f"{duration_s:.1f}s exceeds the profiled maximum of "
+                    f"{max_duration_s:.1f}s ({max_speech_frames} speech frames)."
+                )
+
             wavs.append(audio_np)
             num_frames = int(
                 (float(audio_np.shape[0]) / float(sr)) * float(self.info.token_fps)
diff --git a/vllm/model_executor/models/gemma4.py b/vllm/model_executor/models/gemma4.py
index dc5e3ea81cf9..c67e0197c601 100644
--- a/vllm/model_executor/models/gemma4.py
+++ b/vllm/model_executor/models/gemma4.py
@@ -419,11 +419,14 @@ def __init__(
         self.scaling = 1.0
 
         layer_idx = extract_layer_index(prefix)
+        layer_type = config.layer_types[layer_idx]
         num_kv_shared_layers = getattr(config, "num_kv_shared_layers", 0)
+        attention_k_eq_v = getattr(config, "attention_k_eq_v", False)
         first_kv_shared_layer_idx = config.num_hidden_layers - num_kv_shared_layers
         self.is_kv_shared_layer = (
             num_kv_shared_layers > 0 and layer_idx >= first_kv_shared_layer_idx
         )
+        self.is_sliding = layer_type == "sliding_attention"
 
         if self.is_kv_shared_layer:
             # K/V come from the target layer's cache, so like HF this layer
@@ -448,6 +451,9 @@ def __init__(
                 quant_config=quant_config,
                 prefix=f"{prefix}.qkv_proj",
             )
+            if attention_k_eq_v and not self.is_sliding:
+                # k_proj LoRA must be applied to the V slice
+                self.qkv_proj.lora_packed_modules = ["q_proj", "k_proj", "k_proj"]
             self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
             # V norm: no learnable scale (pure normalization only)
             self.v_norm = RMSNorm(
@@ -464,9 +470,6 @@ def __init__(
         # Q norm: output = norm(x) * weight (learnable per-head scale)
         self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
 
-        # Determine layer type and sliding window
-        layer_type = config.layer_types[layer_idx]
-        self.is_sliding = layer_type == "sliding_attention"
         sliding_window = config.sliding_window if self.is_sliding else None
 
         # Initialize RoPE based on layer type.
diff --git a/vllm/model_executor/models/gemma4_dspark.py b/vllm/model_executor/models/gemma4_dspark.py
index 13c868e59b38..8e6a79af7ec4 100644
--- a/vllm/model_executor/models/gemma4_dspark.py
+++ b/vllm/model_executor/models/gemma4_dspark.py
@@ -19,10 +19,17 @@
     ParallelLMHead,
     VocabParallelEmbedding,
 )
+from vllm.model_executor.model_loader.mtp_validation import (
+    is_mtp_completeness_check_enabled,
+)
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
 from .gemma4_mtp import Gemma4MTPAttention, Gemma4MTPDecoderLayer
-from .qwen3_dflash import DFlashQwen3Model, _dflash_layer_causal
+from .qwen3_dflash import (
+    DFlashQwen3Model,
+    _dflash_layer_causal,
+    _reads_raw_weights,
+)
 from .qwen3_dspark import (
     DSparkConfidenceHead,
     DSparkMarkovHead,
@@ -227,10 +234,13 @@ def _build_context_kv_buffers(
         self, layers_attn: list[nn.Module], has_bias: bool
     ) -> None:
         self._hidden_norm_weight = self.hidden_norm.weight.data
-        self._fused_k_weight = torch.cat([a.k_proj.weight for a in layers_attn], dim=0)
-        self._fused_k_bias: torch.Tensor | None = (
-            torch.cat([a.k_proj.bias for a in layers_attn], dim=0) if has_bias else None
-        )
+        self._context_k_projs = [a.k_proj for a in layers_attn]
+        self._fused_k_weight: torch.Tensor | None = None
+        self._fused_k_bias: torch.Tensor | None = None
+        if _reads_raw_weights(self._context_k_projs):
+            self._fused_k_weight = torch.cat([p.weight for p in self._context_k_projs])
+            if has_bias:
+                self._fused_k_bias = torch.cat([p.bias for p in self._context_k_projs])
         self._k_norm_weights = torch.stack(
             [a.k_norm.weight.data for a in layers_attn], dim=0
         ).contiguous()
@@ -256,7 +266,10 @@ def _project_context_kv(
         ops.rms_norm(
             normed, context_states, self._hidden_norm_weight, self._rms_norm_eps
         )
-        all_k_flat = F.linear(normed, self._fused_k_weight, self._fused_k_bias)
+        if self._fused_k_weight is not None:
+            all_k_flat = F.linear(normed, self._fused_k_weight, self._fused_k_bias)
+        else:
+            all_k_flat = torch.cat([p(normed)[0] for p in self._context_k_projs], -1)
         all_k = (
             all_k_flat.view(num_ctx, num_layers, num_kv_heads, head_dim)
             .permute(1, 0, 2, 3)
@@ -326,7 +339,6 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
                     loaded.add(name)
                     if "confidence_head" in name:
                         includes_confidence_head = True
-        if not includes_confidence_head:
+        if not includes_confidence_head and is_mtp_completeness_check_enabled():
             self.model.confidence_head = None
-        self.model._build_fused_kv_buffers()
         return loaded
diff --git a/vllm/model_executor/models/gemma4_mtp.py b/vllm/model_executor/models/gemma4_mtp.py
index f9931ddafc0a..a0e99deab2d5 100644
--- a/vllm/model_executor/models/gemma4_mtp.py
+++ b/vllm/model_executor/models/gemma4_mtp.py
@@ -45,6 +45,7 @@
     VocabParallelEmbedding,
 )
 from vllm.sequence import IntermediateTensors
+from vllm.v1.worker.gpu.spec_decode.gemma4.config import Gemma4MTPSpeculatorConfig
 
 from .gemma4 import Gemma4MLP, _get_text_config
 from .utils import (
@@ -446,6 +447,9 @@ class Gemma4MTP(nn.Module):
     hidden-state feedback buffer.
     """
 
+    # The speculator this is a draft model of (see SpeculativeDraftModel).
+    speculator_config_class = Gemma4MTPSpeculatorConfig
+
     has_own_lm_head = True
 
     hf_to_vllm_mapper = WeightsMapper(
@@ -516,7 +520,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         draft_cfg = speculative_config.draft_model_config
         gen_cfg = draft_cfg.try_get_generation_config()
         self._suppress_token_ids = gen_cfg.get("suppress_tokens") if gen_cfg else None
-        # Materialized on-device in load_weights: compute_logits runs under CUDA
+        # Materialized on-device after loading: compute_logits runs under CUDA
         # graph capture in the V2 speculator, where indexing with a Python list
         # would issue an unpinned H2D copy (illegal during capture).
         self._suppress_idx: torch.Tensor | None = None
@@ -588,13 +592,14 @@ def get_top_tokens(
         )
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        self._stable_full_lm_head_weight = None
         loader = AutoWeightsLoader(self)
-        loaded = 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._stable_full_lm_head_weight = None
         if self._suppress_token_ids:
             self._suppress_idx = torch.tensor(
                 self._suppress_token_ids,
                 dtype=torch.long,
                 device=next(self.parameters()).device,
             )
-        return loaded
diff --git a/vllm/model_executor/models/glm4_1v.py b/vllm/model_executor/models/glm4_1v.py
index be456a6a5821..9724b4cb4c71 100644
--- a/vllm/model_executor/models/glm4_1v.py
+++ b/vllm/model_executor/models/glm4_1v.py
@@ -89,7 +89,12 @@
     MultiModalKwargsItems,
     VideoItem,
 )
-from vllm.multimodal.parse import ImageSize, MultiModalDataItems, MultiModalDataParser
+from vllm.multimodal.parse import (
+    DictEmbeddingItems,
+    ImageSize,
+    MultiModalDataItems,
+    MultiModalDataParser,
+)
 from vllm.multimodal.processing import (
     BaseDummyInputsBuilder,
     BaseMultiModalProcessor,
@@ -963,6 +968,118 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
 
 
+class Glm4vMultiModalDataParser(MultiModalDataParser):
+    """Placeholder metadata shared by GLM-4V and GLM-5.3-Flash.
+
+    The patch grid sizes image placeholders. Video placeholders also embed one
+    timestamp string per temporal patch, so those timestamps travel with the
+    grid when embeddings arrive through the EC connector.
+    """
+
+    embedding_fields = {
+        "image": {"image_embeds": "values", "image_grid_thw": "metadata"},
+        "video": {
+            "video_embeds": "values",
+            "video_grid_thw": "metadata",
+            "timestamps": "metadata",
+        },
+    }
+
+    def __init__(self, spatial_merge_size: int, *args, **kwargs):
+        self._spatial_merge_size = spatial_merge_size
+        super().__init__(*args, **kwargs)
+
+    def _embed_row_count(self, grid_thw: torch.Tensor) -> int:
+        grid = grid_thw
+        if grid.ndim == 1:
+            grid = grid.unsqueeze(0)
+        if grid.ndim != 2 or grid.shape[-1] != 3:
+            raise ValueError(
+                f"grid_thw must have shape (N, 3), got {tuple(grid_thw.shape)}"
+            )
+        merge = self._spatial_merge_size**2
+        # Sum per item. grid.prod() would multiply the items together.
+        return int((grid.prod(dim=-1) // merge).sum().item())
+
+    def _check_supplied_embeds(
+        self,
+        data: dict,
+        *,
+        embed_key: str,
+        grid_key: str,
+        modality: str,
+    ) -> None:
+        """Reject dict embeddings whose shape cannot match this grid.
+
+        Metadata-only requests omit the embedding tensor and stay valid.
+        """
+        if embed_key not in data:
+            return
+        embeds = data[embed_key]
+        if not isinstance(embeds, torch.Tensor) or embeds.ndim != 2:
+            got = tuple(embeds.shape) if isinstance(embeds, torch.Tensor) else None
+            raise ValueError(
+                f"{modality.capitalize()} embedding must be 2D "
+                f"(num_tokens, hidden_size), got {type(embeds).__name__} "
+                f"with shape {got}"
+            )
+        if (
+            self.expected_hidden_size is not None
+            and embeds.shape[-1] != self.expected_hidden_size
+        ):
+            raise ValueError(
+                f"{modality.capitalize()} embedding hidden dimension mismatch: "
+                f"got {embeds.shape[-1]}, but model expects "
+                f"{self.expected_hidden_size}. "
+                f"Embedding shape: {tuple(embeds.shape)}"
+            )
+        if grid_key not in data:
+            return
+        expected_rows = self._embed_row_count(data[grid_key])
+        if embeds.shape[0] != expected_rows:
+            raise ValueError(
+                f"{modality.capitalize()} embedding row count mismatch: "
+                f"got {embeds.shape[0]}, but grid_thw requires {expected_rows}. "
+                f"Embedding shape: {tuple(embeds.shape)}"
+            )
+
+    def _parse_image_data(self, data):
+        if isinstance(data, dict):
+            self._check_supplied_embeds(
+                data,
+                embed_key="image_embeds",
+                grid_key="image_grid_thw",
+                modality="image",
+            )
+            required, optional = self.embedding_field_sets("image")
+            return DictEmbeddingItems(
+                data,
+                modality="image",
+                required_fields=required,
+                optional_fields=optional,
+                fields_factory=_create_qwen2vl_field_factory(self._spatial_merge_size),
+            )
+        return super()._parse_image_data(data)
+
+    def _parse_video_data(self, data):
+        if isinstance(data, dict):
+            self._check_supplied_embeds(
+                data,
+                embed_key="video_embeds",
+                grid_key="video_grid_thw",
+                modality="video",
+            )
+            required, optional = self.embedding_field_sets("video")
+            return DictEmbeddingItems(
+                data,
+                modality="video",
+                required_fields=required,
+                optional_fields=optional,
+                fields_factory=_create_qwen2vl_field_factory(self._spatial_merge_size),
+            )
+        return super()._parse_video_data(data)
+
+
 class Glm4vProcessingInfo(BaseProcessingInfo):
     def get_supported_mm_limits(self) -> Mapping[str, int | None]:
         return {"image": None, "video": 1}
@@ -1036,9 +1153,11 @@ def get_mm_max_tokens_per_item(
         return result
 
     def get_data_parser(self):
-        return MultiModalDataParser(
+        return Glm4vMultiModalDataParser(
+            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,
         )
 
     def _get_vision_info(
@@ -1391,53 +1510,64 @@ def _get_video_second_idx_glmga(
         timestamps_list = full_second_idxs[::2]
         return list(timestamps_list)
 
-    def _construct_video_placeholder(
+    def _video_timestamps(self, video_array: np.ndarray, metadata: dict[str, Any]):
+        hf_processor = self.get_hf_processor()
+        n_frames = len(video_array)
+        if isinstance(hf_processor, Glm4vProcessor):
+            return self._get_video_second_idx_glm4v(metadata, n_frames)
+        if self._is_glmga_model(hf_processor):
+            return self._get_video_second_idx_glmga(metadata, n_frames)
+        return self._get_video_second_idx_glm46v(metadata, n_frames)
+
+    def _placeholder_from_timestamps(
         self,
-        video_array: np.ndarray,
-        metadata: dict[str, Any],
+        timestamps: torch.Tensor | list,
         grid_thw: torch.Tensor,
     ) -> list[int]:
         hf_processor = self.get_hf_processor()
         tokenizer = self.get_tokenizer()
-        image_processor = hf_processor.image_processor
-
         hf_config = self.get_hf_config()
-        boi_token_id = hf_config.image_start_token_id
-        eoi_token_id = hf_config.image_end_token_id
-        bov_token_id = hf_config.video_start_token_id
-        eov_token_id = hf_config.video_end_token_id
-        merge_length = image_processor.merge_size**2
-
-        assert isinstance(grid_thw, torch.Tensor)
-
-        if isinstance(hf_processor, Glm4vProcessor):
-            timestamps = self._get_video_second_idx_glm4v(metadata, len(video_array))
-        elif self._is_glmga_model(hf_processor):
-            timestamps = self._get_video_second_idx_glmga(metadata, len(video_array))
+        merge_length = hf_processor.image_processor.merge_size**2
+        if isinstance(timestamps, torch.Tensor):
+            timestamp_values = timestamps.flatten().tolist()
         else:
-            timestamps = self._get_video_second_idx_glm46v(metadata, len(video_array))
-
+            timestamp_values = list(timestamps)
+        grid_t = int(grid_thw[0])
+        if len(timestamp_values) != grid_t:
+            raise ValueError(
+                "Video timestamp count does not match video_grid_thw[0]: "
+                f"{len(timestamp_values)} timestamps for grid_t={grid_t}."
+            )
         timestamp_format = (
             "{}" if isinstance(hf_processor, Glm4vProcessor) else "{:.1f} seconds"
         )
         frames_idx_token = [
-            tokenizer.encode(timestamp_format.format(i), add_special_tokens=False)
-            for i in timestamps
+            tokenizer.encode(timestamp_format.format(ts), add_special_tokens=False)
+            for ts in timestamp_values
         ]
-        T, H, W = grid_thw
-        num_tokens_per_frame = int(H * W) // merge_length
-        placeholder = []
-        placeholder.append(bov_token_id)
+        _, height, width = grid_thw
+        num_tokens_per_frame = int(height * width) // merge_length
         frame_embed_token_id = hf_processor.image_token_id
+        placeholder = [hf_config.video_start_token_id]
         for frame_idx in frames_idx_token:
-            placeholder.append(boi_token_id)
+            placeholder.append(hf_config.image_start_token_id)
             placeholder.extend([frame_embed_token_id] * num_tokens_per_frame)
-            placeholder.append(eoi_token_id)
+            placeholder.append(hf_config.image_end_token_id)
             placeholder.extend(frame_idx)
-        placeholder.append(eov_token_id)
-
+        placeholder.append(hf_config.video_end_token_id)
         return placeholder
 
+    def _construct_video_placeholder(
+        self,
+        video_array: np.ndarray,
+        metadata: dict[str, Any],
+        grid_thw: torch.Tensor,
+    ) -> list[int]:
+        return self._placeholder_from_timestamps(
+            self._video_timestamps(video_array, metadata),
+            grid_thw,
+        )
+
 
 class Glm4vDummyInputsBuilder(BaseDummyInputsBuilder[Glm4vProcessingInfo]):
     def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
@@ -1565,6 +1695,46 @@ def _get_direct_path_inputs(
 
         return prepared_data, prepared_kwargs
 
+    @staticmethod
+    def _sampled_frame_metadata_overrides(mm_items, video_metadata):
+        """Build timestamp-helper metadata from HF-sampled frame indices.
+
+        Input ``frames_indices`` can be the pre-sample range. The processor
+        returns a ``VideoMetadata`` with the sample it kept. GLM-5.3's direct
+        path must feed those indices into the timestamp helper and clear
+        ``do_sample_frames`` so the helper keeps them instead of sampling again.
+
+        Overrides are returned and never written back onto ``mm_items``. The
+        same items are reused by ``test_processing_correctness`` (baseline then
+        cache); persisting ``do_sample_frames=False`` while the video tensor
+        stays full-length made pass-2 HF emit ``grid_t != len(timestamps)`` for
+        odd-length videos.
+        """
+        if not video_metadata:
+            return None
+        video_items = mm_items.get("video")
+        if video_items is None:
+            return None
+        overrides: list[dict[str, Any] | None] = [None] * video_items.get_count()
+        for index, meta_obj in enumerate(video_metadata):
+            indices = getattr(meta_obj, "frames_indices", None)
+            if indices is None or index >= video_items.get_count():
+                continue
+            raw = video_items.data[index]
+            if not (
+                isinstance(raw, tuple) and len(raw) == 2 and isinstance(raw[1], dict)
+            ):
+                continue
+            if hasattr(indices, "tolist"):
+                indices = indices.tolist()
+            new_meta = dict(raw[1])
+            new_meta["frames_indices"] = [int(idx) for idx in indices]
+            new_meta["do_sample_frames"] = False
+            overrides[index] = new_meta
+        if not any(override is not None for override in overrides):
+            return None
+        return overrides
+
     def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str:
         return self.dummy_inputs.get_dummy_text(mm_counts)
 
@@ -1578,7 +1748,10 @@ def _apply_hf_processor_main(
         )
 
         if not hf_data:
-            return self._finalize_hf_mm_data(hf_data, hf_kwargs, passthrough_data)
+            return self._with_video_timestamps(
+                self._finalize_hf_mm_data(hf_data, hf_kwargs, passthrough_data),
+                mm_items,
+            )
 
         prompt_text = hf_data.pop("text")
         assert isinstance(prompt_text, str)
@@ -1594,8 +1767,23 @@ def _apply_hf_processor_main(
                 dict(text=prompt_text, **prepared_data),
                 prepared_kwargs,
             )
-            return self._finalize_hf_mm_data(
-                hf_data, hf_kwargs, passthrough_data, processed_data
+            # The processor returns a new VideoMetadata whose frames_indices
+            # are the sample it kept. The object we passed in still holds the
+            # pre-sample range.
+            sampled_metadata = (
+                processed_data.get("video_metadata")
+                if hasattr(processed_data, "get")
+                else None
+            )
+            metadata_overrides = self._sampled_frame_metadata_overrides(
+                mm_items, sampled_metadata
+            )
+            return self._with_video_timestamps(
+                self._finalize_hf_mm_data(
+                    hf_data, hf_kwargs, passthrough_data, processed_data
+                ),
+                mm_items,
+                metadata_overrides=metadata_overrides,
             )
 
         if (
@@ -1619,7 +1807,8 @@ def _apply_hf_processor_main(
 
                 video_mm_data = dict()
                 video_mm_data["videos"] = [[video_array]]
-                video_mm_data["video_metadata"] = [[_to_video_metadata(metadata)]]
+                video_metadata = _to_video_metadata(metadata)
+                video_mm_data["video_metadata"] = [[video_metadata]]
 
                 video_outputs = self.info.ctx.call_hf_processor(
                     self.info.get_hf_processor(**video_mm_kwargs),
@@ -1629,6 +1818,23 @@ def _apply_hf_processor_main(
                     ),
                     video_mm_kwargs,
                 )
+                sampled_indices = None
+                out_meta = (
+                    video_outputs.get("video_metadata")
+                    if hasattr(video_outputs, "get")
+                    else None
+                )
+                if out_meta:
+                    first = (
+                        out_meta[0][0] if isinstance(out_meta[0], list) else out_meta[0]
+                    )
+                    sampled_indices = getattr(first, "frames_indices", None)
+                if sampled_indices is None:
+                    sampled_indices = getattr(video_metadata, "frames_indices", None)
+                if sampled_indices is not None and isinstance(metadata, dict):
+                    if hasattr(sampled_indices, "tolist"):
+                        sampled_indices = sampled_indices.tolist()
+                    metadata["frames_indices"] = [int(idx) for idx in sampled_indices]
                 input_ids = video_outputs.pop("input_ids")
 
                 input_ids[input_ids == processor.image_token_id] = (
@@ -1660,10 +1866,39 @@ def _apply_hf_processor_main(
         processed_data["input_ids"] = input_ids
 
         processed_data.update(video_outputs)
-        return self._finalize_hf_mm_data(
-            hf_data, hf_kwargs, passthrough_data, processed_data
+        return self._with_video_timestamps(
+            self._finalize_hf_mm_data(
+                hf_data, hf_kwargs, passthrough_data, processed_data
+            ),
+            mm_items,
         )
 
+    def _with_video_timestamps(
+        self,
+        processed,
+        mm_items: MultiModalDataItems,
+        metadata_overrides: list[dict[str, Any] | None] | None = None,
+    ):
+        video_items = mm_items.get("video")
+        if (
+            video_items is not None
+            and video_items.get_count() > 0
+            and not isinstance(video_items, DictEmbeddingItems)
+            and "timestamps" not in processed
+        ):
+            timestamps = []
+            for i in range(video_items.get_count()):
+                video_array, metadata = video_items.get(i)
+                if (
+                    metadata_overrides is not None
+                    and i < len(metadata_overrides)
+                    and metadata_overrides[i] is not None
+                ):
+                    metadata = metadata_overrides[i]
+                timestamps.append(self.info._video_timestamps(video_array, metadata))
+            processed["timestamps"] = timestamps
+        return processed
+
     def _get_mm_fields_config(
         self,
         hf_inputs: BatchFeature,
@@ -1698,11 +1933,16 @@ def get_video_replacement(item_idx: int):
             out_item = out_mm_kwargs["video"][item_idx]
             grid_thw = out_item["video_grid_thw"].data
             assert isinstance(grid_thw, torch.Tensor)
-
-            video, metadata = mm_items["video"][item_idx]
-            placeholder = self.info._construct_video_placeholder(
-                video, metadata, grid_thw
-            )
+            timestamps = out_item.get("timestamps")
+            if timestamps is not None:
+                placeholder = self.info._placeholder_from_timestamps(
+                    timestamps.data, grid_thw
+                )
+            else:
+                video, metadata = mm_items["video"][item_idx]
+                placeholder = self.info._construct_video_placeholder(
+                    video, metadata, grid_thw
+                )
             return PromptUpdateDetails.select_token_id(
                 placeholder,
                 embed_token_id=hf_processor.image_token_id,
diff --git a/vllm/model_executor/models/glm4_moe.py b/vllm/model_executor/models/glm4_moe.py
index 4e541a49413e..c43f12e41fad 100644
--- a/vllm/model_executor/models/glm4_moe.py
+++ b/vllm/model_executor/models/glm4_moe.py
@@ -46,7 +46,6 @@
     GateLinear,
 )
 from vllm.model_executor.layers.fused_moe.utils import (
-    is_model_fused_shared_expert_compatible,
     resolve_layer_fused_shared_expert,
 )
 from vllm.model_executor.layers.layernorm import RMSNorm
@@ -69,6 +68,7 @@
     AutoWeightsLoader,
     PPMissingLayer,
     WeightsMapper,
+    get_fused_shared_expert_layers,
     make_empty_intermediate_tensors_factory,
     make_layers,
     maybe_fuse_shared_experts,
@@ -444,10 +444,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
             prefix=f"{prefix}.layers",
         )
 
-        self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible(
-            self.layers,
+        self.fused_shared_expert_layers = get_fused_shared_expert_layers(
+            self,
             Glm4MoE,
-            "mlp",
+            enable_eplb=enable_eplb,
         )
 
         if get_pp_group().is_last_rank:
@@ -495,7 +495,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
             skip_spec_layers(weights, self.config),
             n_routed_experts=self.config.n_routed_experts,
             n_shared_experts=self.config.n_shared_experts or 1,
-            enabled=self.is_fused_shared_expert_enabled,
+            fused_shared_expert_layers=self.fused_shared_expert_layers,
         )
         loader = AutoWeightsLoader(self)
         return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
diff --git a/vllm/model_executor/models/glm4_moe_lite.py b/vllm/model_executor/models/glm4_moe_lite.py
index 019380540d08..367fbbbcd96a 100644
--- a/vllm/model_executor/models/glm4_moe_lite.py
+++ b/vllm/model_executor/models/glm4_moe_lite.py
@@ -34,6 +34,7 @@
 if TYPE_CHECKING:
     from transformers.models.glm4_moe_lite import Glm4MoeLiteConfig
 
+from vllm._aiter_ops import rocm_aiter_ops
 from vllm.compilation.decorators import support_torch_compile
 from vllm.config import VllmConfig
 from vllm.distributed import (
@@ -43,9 +44,6 @@
 from vllm.model_executor.layers.fused_moe import (
     fused_moe_make_expert_params_mapping,
 )
-from vllm.model_executor.layers.fused_moe.utils import (
-    is_model_fused_shared_expert_compatible,
-)
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
 from vllm.model_executor.layers.vocab_parallel_embedding import (
@@ -76,6 +74,8 @@
 from .utils import (
     AutoWeightsLoader,
     PPMissingLayer,
+    extract_layer_index,
+    get_fused_shared_expert_layers,
     get_spec_layer_idx_from_weight_name,
     is_pp_missing_parameter,
     make_layers,
@@ -276,10 +276,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
             prefix=f"{prefix}.layers",
         )
 
-        self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible(
-            self.layers,
+        self.fused_shared_expert_layers = get_fused_shared_expert_layers(
+            self,
             Glm4MoeLite,
-            "mlp",
+            enable_eplb=vllm_config.parallel_config.enable_eplb,
         )
 
         if get_pp_group().is_last_rank:
@@ -364,12 +364,9 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
             ckpt_gate_proj_name="gate_proj",
             ckpt_down_proj_name="down_proj",
             ckpt_up_proj_name="up_proj",
+            # NOTE: global despite per-layer FSE; non-fused layers skip the extras.
             num_experts=self.config.n_routed_experts
-            + (
-                self.config.n_shared_experts
-                if self.is_fused_shared_expert_enabled
-                else 0
-            ),
+            + (self.config.n_shared_experts if self.fused_shared_expert_layers else 0),
         )
 
         params_dict = dict(self.named_parameters())
@@ -382,8 +379,11 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
             if spec_layer is not None:
                 continue  # skip spec decode layers for main model
 
+            # Only parse the layer index when FSE is enabled.
             is_fusion_moe_shared_experts_layer = (
-                self.is_fused_shared_expert_enabled and ("mlp.shared_experts" in name)
+                "mlp.shared_experts" in name
+                and bool(rocm_aiter_ops.is_fusion_moe_shared_experts_enabled())
+                and extract_layer_index(name) in self.fused_shared_expert_layers
             )
 
             for param_name, weight_name, shard_id in stacked_params_mapping:
diff --git a/vllm/model_executor/models/glm4_moe_lite_mtp.py b/vllm/model_executor/models/glm4_moe_lite_mtp.py
index f5a82b72314b..a59eeb8f60fb 100644
--- a/vllm/model_executor/models/glm4_moe_lite_mtp.py
+++ b/vllm/model_executor/models/glm4_moe_lite_mtp.py
@@ -30,14 +30,12 @@
 import torch.nn as nn
 from transformers import PreTrainedConfig
 
+from vllm._aiter_ops import rocm_aiter_ops
 from vllm.config import VllmConfig
 from vllm.model_executor.layers.fused_moe import (
     MoERunner,
     fused_moe_make_expert_params_mapping,
 )
-from vllm.model_executor.layers.fused_moe.utils import (
-    is_model_fused_shared_expert_compatible,
-)
 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
@@ -62,7 +60,12 @@
     Glm4MoeLiteDecoderLayer,
 )
 from .interfaces import SupportsPP
-from .utils import get_spec_layer_idx_from_weight_name, maybe_prefix
+from .utils import (
+    extract_layer_index,
+    get_fused_shared_expert_layers,
+    get_spec_layer_idx_from_weight_name,
+    maybe_prefix,
+)
 
 
 class SharedHead(nn.Module):
@@ -241,10 +244,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
                 self.moe_mlp_layers.append(layer.mlp)
                 self.moe_layers.append(layer.mlp.experts)
         self.extract_moe_parameters(example_moe)
-        self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible(
-            self.model.layers.values(),
+        self.fused_shared_expert_layers = get_fused_shared_expert_layers(
+            self,
             Glm4MoeLite,
-            "mtp_block.mlp",
+            enable_eplb=vllm_config.parallel_config.enable_eplb,
         )
 
     def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
@@ -279,17 +282,15 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
             ("fused_qkv_a_proj", "kv_a_proj_with_mqa", 1),
         ]
 
+        fused_shared_expert_layers = self.fused_shared_expert_layers
         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",
+            # NOTE: global despite per-layer FSE; non-fused layers skip the extras.
             num_experts=self.config.n_routed_experts
-            + (
-                self.config.n_shared_experts
-                if self.is_fused_shared_expert_enabled
-                else 0
-            ),
+            + (self.config.n_shared_experts if fused_shared_expert_layers else 0),
         )
 
         params_dict = dict(self.named_parameters())
@@ -300,8 +301,11 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
             spec_layer = get_spec_layer_idx_from_weight_name(self.config, name)
             if spec_layer is None:
                 continue
+            # Only parse the layer index when FSE is enabled.
             is_fusion_moe_shared_experts_layer = (
-                self.is_fused_shared_expert_enabled and ("mlp.shared_experts" in name)
+                "mlp.shared_experts" in name
+                and bool(rocm_aiter_ops.is_fusion_moe_shared_experts_enabled())
+                and extract_layer_index(name) in fused_shared_expert_layers
             )
             name = self._rewrite_spec_layer_name(spec_layer, name)
             for param_name, weight_name, shard_id in stacked_params_mapping:
diff --git a/vllm/model_executor/models/glm4_moe_mtp.py b/vllm/model_executor/models/glm4_moe_mtp.py
index 330f891d1cf3..0464322cae02 100644
--- a/vllm/model_executor/models/glm4_moe_mtp.py
+++ b/vllm/model_executor/models/glm4_moe_mtp.py
@@ -31,14 +31,12 @@
 import torch.nn as nn
 from transformers import PreTrainedConfig
 
+from vllm._aiter_ops import rocm_aiter_ops
 from vllm.config import CacheConfig, ParallelConfig, VllmConfig
 from vllm.model_executor.layers.fused_moe import (
     MoERunner,
     fused_moe_make_expert_params_mapping,
 )
-from vllm.model_executor.layers.fused_moe.utils import (
-    is_model_fused_shared_expert_compatible,
-)
 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
@@ -54,7 +52,12 @@
     Glm4MoE,
     Glm4MoeDecoderLayer,
 )
-from .utils import get_spec_layer_idx_from_weight_name, maybe_prefix
+from .utils import (
+    extract_layer_index,
+    get_fused_shared_expert_layers,
+    get_spec_layer_idx_from_weight_name,
+    maybe_prefix,
+)
 
 
 class SharedHead(nn.Module):
@@ -215,10 +218,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
                 self.moe_mlp_layers.append(layer.mlp)
                 self.moe_layers.append(layer.mlp.experts)
         self.extract_moe_parameters(example_moe)
-        self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible(
-            self.model.layers.values(),
+        self.fused_shared_expert_layers = get_fused_shared_expert_layers(
+            self,
             Glm4MoE,
-            "mtp_block.mlp",
+            enable_eplb=vllm_config.parallel_config.enable_eplb,
         )
 
     def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
@@ -257,8 +260,10 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
 
         # Params for weights, fp8 weight scales, fp8 activation scales
         # (param_name, weight_name, expert_id, shard_id)
+        fused_shared_expert_layers = self.fused_shared_expert_layers
+        # NOTE: global despite per-layer FSE; non-fused layers skip the extras.
         num_experts = self.config.n_routed_experts
-        if self.is_fused_shared_expert_enabled and self.config.n_shared_experts:
+        if fused_shared_expert_layers and self.config.n_shared_experts:
             num_experts += self.config.n_shared_experts
         expert_params_mapping = fused_moe_make_expert_params_mapping(
             self,
@@ -282,8 +287,11 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
                     continue
                 name = self._rewrite_spec_layer_name(spec_layer, name)
 
+            # Only parse the layer index when FSE is enabled.
             is_fusion_moe_shared_experts_layer = (
-                self.is_fused_shared_expert_enabled and ("mlp.shared_experts" in name)
+                "mlp.shared_experts" in name
+                and bool(rocm_aiter_ops.is_fusion_moe_shared_experts_enabled())
+                and extract_layer_index(name) in fused_shared_expert_layers
             )
 
             for param_name, weight_name, shard_id in stacked_params_mapping:
diff --git a/vllm/model_executor/models/interfaces.py b/vllm/model_executor/models/interfaces.py
index 701c24a072d4..3635a637bc77 100644
--- a/vllm/model_executor/models/interfaces.py
+++ b/vllm/model_executor/models/interfaces.py
@@ -71,6 +71,7 @@
         EncoderCudaGraphReplayBuffers,
         EncoderItemSpec,
     )
+    from vllm.v1.worker.gpu.spec_decode.speculator_config import SpeculatorConfig
 
 logger = init_logger(__name__)
 
@@ -736,7 +737,7 @@ class SupportsLoRA(Protocol):
     packed_modules_mapping: dict[str, list[str]] = {}
     # Module prefixes to skip during LoRA loading (e.g., ["mtp."] for MTP layers)
     lora_skip_prefixes: ClassVar[list[str]] = []
-    lora_manager: "LoRAModelManager | None"
+    lora_manager: "LoRAModelManager | None" = None
 
 
 # We can't use runtime_checkable with ClassVar for issubclass checks
@@ -1537,6 +1538,20 @@ def supports_transcription(
     return getattr(model, "supports_transcription", False)
 
 
+@runtime_checkable
+class SpeculativeDraftModel(Protocol):
+    """The interface of a speculative decoding draft model which declares the
+    speculator it's for. This is read through the model registry, so that it's
+    available without importing the model.
+    """
+
+    speculator_config_class: ClassVar[type["SpeculatorConfig"]]
+    """The config class of the speculator that this is a draft model of.
+    Unless a speculative method is configured that the checkpoint can't
+    override, a draft checkpoint with this model's architecture uses this
+    speculator."""
+
+
 @runtime_checkable
 class SupportsEagleBase(Protocol):
     """Base interface for models that support EAGLE-based speculative decoding."""
diff --git a/vllm/model_executor/models/keye.py b/vllm/model_executor/models/keye.py
deleted file mode 100644
index 83ff14dcf948..000000000000
--- a/vllm/model_executor/models/keye.py
+++ /dev/null
@@ -1,1684 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import math
-from abc import abstractmethod
-from collections.abc import Iterable, Iterator, Mapping, Sequence
-from functools import partial
-from typing import Annotated, Any, Literal, TypeAlias, TypeVar
-
-import numpy as np
-import torch
-import torch.nn as nn
-from einops import rearrange
-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 MultiModalDummyOptions
-from vllm.distributed import get_tensor_model_parallel_world_size
-from vllm.inputs import ModalityData, MultiModalDataDict
-from vllm.logger import init_logger
-from vllm.model_executor.layers.attention import (
-    MMEncoderAttention,
-)
-from vllm.model_executor.layers.conv import Conv2dLayer
-from vllm.model_executor.layers.linear import (
-    ColumnParallelLinear,
-    QKVParallelLinear,
-    RowParallelLinear,
-)
-from vllm.model_executor.layers.quantization import QuantizationConfig
-from vllm.model_executor.layers.rotary_embedding.common import (
-    ApplyRotaryEmb,
-)
-from vllm.model_executor.models.module_mapping import MultiModelKeys
-from vllm.multimodal import MULTIMODAL_REGISTRY
-from vllm.multimodal.inputs import (
-    ImageItem,
-    MultiModalFeatureSpec,
-    MultiModalFieldConfig,
-    MultiModalKwargsItem,
-    MultiModalKwargsItems,
-    VideoItem,
-)
-from vllm.multimodal.parse import (
-    DictEmbeddingItems,
-    ImageSize,
-    ModalityDataItems,
-    MultiModalDataItems,
-    MultiModalDataParser,
-)
-from vllm.multimodal.processing import (
-    BaseDummyInputsBuilder,
-    BaseMultiModalProcessor,
-    BaseProcessingInfo,
-    PromptReplacement,
-    PromptUpdate,
-)
-from vllm.sequence import IntermediateTensors
-from vllm.utils.tensor_schema import TensorSchema, TensorShape
-from vllm.utils.torch_utils import PIN_MEMORY, async_tensor_h2d
-
-from .interfaces import (
-    MultiModalEmbeddings,
-    SupportsLoRA,
-    SupportsMRoPE,
-    SupportsMultiModal,
-    SupportsPP,
-)
-from .siglip import SiglipMLP
-from .utils import (
-    AutoWeightsLoader,
-    WeightsMapper,
-    init_vllm_registered_model,
-    maybe_prefix,
-)
-from .vision import is_vit_use_data_parallel
-
-logger = init_logger(__name__)
-
-
-def smart_resize(
-    height: int,
-    width: int,
-    factor: int,
-    min_pixels: int,
-    max_pixels: int,
-):
-    if height < factor:
-        logger.warning(
-            "smart_resize: height=%s < factor=%s, reset height=factor",
-            height,
-            factor,
-        )
-        width = round((width * factor) / height)
-        height = factor
-
-    if width < factor:
-        logger.warning(
-            "smart_resize: width=%s < factor=%s, reset width=factor",
-            width,
-            factor,
-        )
-        height = round((height * factor) / width)
-        width = factor
-
-    if max(height, width) / min(height, width) > 200:
-        raise ValueError(
-            "absolute aspect ratio must be smaller than 200, got "
-            "{max(height, width) / min(height, width)}"
-        )
-    h_bar = round(height / factor) * factor
-    w_bar = round(width / factor) * factor
-    if h_bar * w_bar > max_pixels:
-        beta = math.sqrt((height * width) / max_pixels)
-        h_bar = math.floor(height / beta / factor) * factor
-        w_bar = math.floor(width / beta / factor) * factor
-    elif h_bar * w_bar < min_pixels:
-        beta = math.sqrt(min_pixels / (height * width))
-        h_bar = math.ceil(height * beta / factor) * factor
-        w_bar = math.ceil(width * beta / factor) * factor
-    return h_bar, w_bar
-
-
-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)
-    """
-
-    type: Literal["pixel_values"]
-    pixel_values: Annotated[
-        torch.Tensor, TensorShape("bnp", 3, "ps", "ps", dynamic_dims={"bnp"})
-    ]
-    image_grid_thw: Annotated[torch.Tensor, TensorShape("ni", 3)]
-
-
-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)
-    """
-
-    type: Literal["image_embeds"]
-    image_embeds: Annotated[torch.Tensor, TensorShape("nf", "hs")]
-    image_grid_thw: Annotated[torch.Tensor, TensorShape("ni", 3)]
-
-
-KeyeImageInputs: TypeAlias = KeyeImagePixelInputs | KeyeImageEmbeddingInputs
-
-
-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)
-    """
-
-    type: Literal["pixel_values_videos"]
-    pixel_values_videos: Annotated[
-        torch.Tensor, TensorShape("bnp", 3, "ps", "ps", dynamic_dims={"bnp"})
-    ]
-    video_grid_thw: Annotated[torch.Tensor, TensorShape("nv", 3)]
-
-
-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)
-    """
-
-    type: Literal["video_embeds"]
-    video_embeds: Annotated[torch.Tensor, TensorShape("nf", "hs")]
-    video_grid_thw: Annotated[torch.Tensor, TensorShape("nv", 3)]
-
-
-KeyeVideoInputs: TypeAlias = KeyeVideoPixelInputs | KeyeVideoEmbeddingInputs
-
-
-class KeyeVisionEmbeddings(nn.Module):
-    def __init__(self, config: PreTrainedConfig):
-        super().__init__()
-        self.config = config
-        self.embed_dim = config.hidden_size
-        self.image_size = config.image_size
-        self.patch_size = config.patch_size
-
-        self.patch_embedding = Conv2dLayer(
-            in_channels=config.num_channels,
-            out_channels=self.embed_dim,
-            kernel_size=self.patch_size,
-            stride=self.patch_size,
-            padding="valid",
-        )
-
-        self.num_patches = (self.image_size // self.patch_size) ** 2
-        self.num_positions = self.num_patches
-        self.cache_position_embedding: dict[tuple[int, int], torch.Tensor] = {}
-        self.cache_position_count: dict[tuple[int, int], int] = {}
-        self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim)
-        self.packing_position_embedding = nn.Embedding(32768, self.embed_dim)
-
-        self.register_buffer(
-            "position_ids",
-            torch.arange(self.num_positions).expand((1, -1)),
-            persistent=False,
-        )
-
-    def interpolate_pos_encoding(
-        self,
-        embeddings: torch.Tensor,
-        height: int,
-        width: int,
-        is_after_patchify: bool = False,
-    ) -> torch.Tensor:
-        num_positions = self.position_embedding.weight.shape[0]
-
-        patch_pos_embed = self.position_embedding.weight.unsqueeze(0)
-
-        dim = embeddings.shape[-1]
-
-        if is_after_patchify:
-            new_height = height
-            new_width = width
-        else:
-            new_height = height // self.patch_size
-            new_width = width // self.patch_size
-
-        sqrt_num_positions = torch_int(num_positions**0.5)
-        patch_pos_embed = patch_pos_embed.reshape(
-            1, sqrt_num_positions, sqrt_num_positions, dim
-        )
-        patch_pos_embed = patch_pos_embed.permute(0, 3, 1, 2)
-
-        patch_pos_embed = nn.functional.interpolate(
-            patch_pos_embed,
-            size=(new_height, new_width),
-            mode="bilinear",
-            align_corners=False,
-        )
-
-        patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
-        return patch_pos_embed
-
-    def forward(
-        self,
-        pixel_values: torch.FloatTensor,
-        position_ids: torch.Tensor | None = None,
-        image_grid_thw: list[tuple[int, int, int]] | None = None,
-        interpolate_pos_encoding=False,
-    ) -> torch.Tensor:
-        if pixel_values.dim() == 4:
-            pixel_values = pixel_values.unsqueeze(0)
-        if pixel_values.dim() == 5:
-            if position_ids is None:
-                raise ValueError(
-                    "position_ids cannot be None when pixel_values.dim() is 5."
-                )
-            (
-                batch_size,
-                sequence_len,
-                channel,
-                height,
-                width,
-            ) = pixel_values.shape
-            target_dtype = self.patch_embedding.weight.dtype
-            pixel_values = rearrange(pixel_values, "b l c h w -> (b l) c h w")
-            patch_embeds = self.patch_embedding(pixel_values.to(dtype=target_dtype))
-            embeddings = patch_embeds.flatten(-2).squeeze(-1)
-
-            if interpolate_pos_encoding and image_grid_thw is not None:
-                start = 0
-                tmp_embeddings = list()
-                for image_grid in image_grid_thw:
-                    t, h, w = image_grid
-                    end = start + t * h * w
-                    image_embeddings = embeddings[start:end, :]
-                    position_embedding = (
-                        self.interpolate_pos_encoding(image_embeddings, h, w, True)
-                        .squeeze(0)
-                        .repeat(t, 1)
-                    )
-                    image_embeddings = image_embeddings + position_embedding
-                    tmp_embeddings.append(image_embeddings)
-                    start = end
-                embeddings = torch.concat(tmp_embeddings, dim=0).unsqueeze(0)
-            else:
-                embeddings = embeddings + self.packing_position_embedding(position_ids)
-            return embeddings
-        else:
-            raise ValueError(
-                "Unsupported pixel_values dimension:"
-                f" {pixel_values.dim()}. Expected 4 or 5."
-            )
-
-
-def apply_rotary_pos_emb_flashatt(
-    q: torch.Tensor,
-    k: torch.Tensor,
-    cos: torch.Tensor,
-    sin: torch.Tensor,
-    apply_rotary_emb: ApplyRotaryEmb,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    cos = cos.chunk(2, dim=-1)[0].contiguous()
-    sin = sin.chunk(2, dim=-1)[0].contiguous()
-
-    q_embed = apply_rotary_emb(q, cos, sin)
-    k_embed = apply_rotary_emb(k, cos, sin)
-
-    return q_embed, k_embed
-
-
-class KeyeSiglipAttention(nn.Module):
-    """Multi-headed attention from 'Attention Is All You
-    Need' paper."""
-
-    def __init__(
-        self,
-        config: PreTrainedConfig,
-        quant_config: QuantizationConfig | None = None,
-        prefix: str = "",
-    ):
-        super().__init__()
-        self.config = config
-
-        hidden_size = config.hidden_size
-        self.hidden_size = config.hidden_size
-        use_data_parallel = is_vit_use_data_parallel()
-        tp_size = 1 if use_data_parallel else get_tensor_model_parallel_world_size()
-        self.total_num_heads = config.num_attention_heads
-        assert self.total_num_heads % tp_size == 0
-        self.num_heads = self.total_num_heads // tp_size
-        self.total_num_kv_heads = config.num_attention_heads
-        if self.total_num_kv_heads >= tp_size:
-            assert self.total_num_kv_heads % tp_size == 0
-        else:
-            assert tp_size % self.total_num_kv_heads == 0
-        self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
-        self.head_dim = config.hidden_size // self.total_num_heads
-        self.q_size = self.num_heads * self.head_dim
-        self.kv_size = self.num_kv_heads * self.head_dim
-        self.scale = self.head_dim**-0.5
-
-        self.qkv_proj = QKVParallelLinear(
-            hidden_size,
-            self.head_dim,
-            self.total_num_heads,
-            self.total_num_kv_heads,
-            bias=True,
-            quant_config=quant_config,
-            prefix=f"{prefix}.qkv_proj",
-        )
-        self.out_proj = RowParallelLinear(
-            input_size=hidden_size,
-            output_size=hidden_size,
-            quant_config=quant_config,
-            prefix=f"{prefix}.out_proj",
-        )
-
-        self.attn = MMEncoderAttention(
-            num_heads=self.num_heads,
-            head_size=self.head_dim,
-            scale=self.scale,
-            num_kv_heads=self.num_kv_heads,
-            prefix=f"{prefix}.attn",
-        )
-
-        self.apply_rotary_emb = ApplyRotaryEmb(
-            enforce_enable=True,
-            enable_fp32_compute=True,
-        )
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        attention_mask: torch.Tensor | None = None,
-        output_attentions: bool | None = False,
-        cu_seqlens: torch.Tensor | None = None,
-        rope_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
-    ) -> torch.Tensor:
-        qkv, _ = self.qkv_proj(hidden_states)
-        q, k, v = qkv.split(
-            [self.q_size, self.kv_size, self.kv_size],
-            dim=-1,
-        )
-
-        if cu_seqlens is None:
-            raise ValueError("cu_seqlens cannot be None.")
-        max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()
-
-        if rope_emb is None:
-            q = q.view(*q.shape[:-1], self.num_heads, self.head_dim)
-            k = k.view(
-                *k.shape[:-1],
-                self.num_kv_heads,
-                self.head_dim,
-            )
-            v = v.view(
-                *v.shape[:-1],
-                self.num_kv_heads,
-                self.head_dim,
-            )
-        else:
-            cos, sin = rope_emb
-            q = q.view(*q.shape[:-1], self.num_heads, self.head_dim)
-            k = k.view(
-                *k.shape[:-1],
-                self.num_kv_heads,
-                self.head_dim,
-            )
-            q, k = apply_rotary_pos_emb_flashatt(q, k, cos, sin, self.apply_rotary_emb)
-            v = v.view(
-                *v.shape[:-1],
-                self.num_kv_heads,
-                self.head_dim,
-            )
-
-        context_layer = self.attn(
-            query=q,
-            key=k,
-            value=v,
-            cu_seqlens=cu_seqlens,
-            max_seqlen=max_seqlen,
-        )
-        context_layer = rearrange(context_layer, "b s h d -> b s (h d)")
-
-        output, _ = self.out_proj(context_layer)
-        return output
-
-
-class SigLIPRotaryEmbedding(nn.Module):
-    def __init__(self, dim: int, theta: float = 10000.0) -> None:
-        super().__init__()
-        self.dim = dim
-        self.theta = theta
-        self.rope_init()
-
-    def rope_init(self):
-        inv_freq = 1.0 / (
-            self.theta ** (torch.arange(0, self.dim, 2, dtype=torch.float) / self.dim)
-        )
-        self.register_buffer("inv_freq", inv_freq, persistent=False)
-
-    def forward(self, seqlen: int) -> torch.Tensor:
-        seq = torch.arange(
-            seqlen,
-            device=self.inv_freq.device,
-            dtype=self.inv_freq.dtype,
-        )
-        freqs = torch.outer(seq, self.inv_freq)
-        return freqs
-
-
-class KeyeSiglipEncoderLayer(nn.Module):
-    def __init__(
-        self,
-        config: PreTrainedConfig,
-        quant_config: QuantizationConfig | None = None,
-        prefix: str = "",
-    ):
-        super().__init__()
-        self.embed_dim = config.hidden_size
-        self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
-        self.self_attn = KeyeSiglipAttention(
-            config,
-            quant_config=quant_config,
-            prefix=f"{prefix}.self_attn",
-        )
-        self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
-        self.mlp = SiglipMLP(
-            config,
-            quant_config=quant_config,
-            prefix=f"{prefix}.mlp",
-        )
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        attention_mask: torch.Tensor,
-        output_attentions: bool | None = False,
-        cu_seqlens: torch.Tensor | None = None,
-        rope_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
-    ) -> tuple[torch.FloatTensor]:
-        residual = hidden_states
-
-        hidden_states = self.layer_norm1(hidden_states)
-        hidden_states = self.self_attn(
-            hidden_states=hidden_states,
-            attention_mask=attention_mask,
-            output_attentions=output_attentions,
-            cu_seqlens=cu_seqlens,
-            rope_emb=rope_emb,
-        )
-
-        hidden_states = residual + hidden_states
-
-        residual = hidden_states
-        hidden_states = self.layer_norm2(hidden_states)
-        hidden_states = self.mlp(hidden_states)
-
-        hidden_states = residual + hidden_states
-
-        return hidden_states
-
-
-class KeyeSiglipEncoder(nn.Module):
-    def __init__(
-        self,
-        config: PreTrainedConfig,
-        quant_config: QuantizationConfig | None = None,
-        prefix: str = "",
-    ):
-        super().__init__()
-        self.config = config
-        embed_dim = config.hidden_size
-        num_heads = config.num_attention_heads
-        head_dim = embed_dim // num_heads
-        self.layers = nn.ModuleList(
-            [
-                KeyeSiglipEncoderLayer(
-                    config,
-                    quant_config=quant_config,
-                    prefix=f"{prefix}.layers.{layer_idx}",
-                )
-                for layer_idx in range(config.num_hidden_layers)
-            ]
-        )
-        self.rotary_pos_emb = SigLIPRotaryEmbedding(head_dim // 2)
-
-    @staticmethod
-    def flatten_list(image_grid_thw):
-        tmp_image_grid_thw = list()
-        for image_grid in image_grid_thw:
-            if isinstance(image_grid, list):
-                tmp_image_grid_thw.extend(image_grid)
-            else:
-                tmp_image_grid_thw.append(image_grid)
-        return tmp_image_grid_thw
-
-    def forward(
-        self,
-        inputs_embeds,
-        attention_mask: torch.Tensor | None = None,
-        output_attentions: bool | None = None,
-        output_hidden_states: bool | None = None,
-        cu_seqlens: torch.Tensor | None = None,
-        image_grid_thw: list[tuple[int, int, int]] | None = None,
-        height_position_ids: torch.Tensor | None = None,
-        width_position_ids: torch.Tensor | None = None,
-        use_rope: bool | None = False,
-        window_size: int = -1,
-        vision_or_text: str = "vision",
-    ) -> BaseModelOutput:
-        device = inputs_embeds.device
-        hidden_states = inputs_embeds
-        if use_rope is True:
-            flatten_image_grid_thw = self.flatten_list(image_grid_thw)
-
-            if width_position_ids is None or height_position_ids is None:
-                split_hids = list()
-                split_wids = list()
-                for t, h, w in flatten_image_grid_thw:
-                    image_pids = torch.arange(t * h * w, device=device) % (h * w)
-                    sample_hids = image_pids // w
-                    sample_wids = image_pids % w
-                    split_hids.append(sample_hids)
-                    split_wids.append(sample_wids)
-                width_position_ids = torch.concat(split_wids, dim=0)
-                height_position_ids = torch.concat(split_hids, dim=0)
-
-            pids = torch.stack(
-                [height_position_ids, width_position_ids],
-                dim=-1,
-            )
-            # The ids are built from the grids above, so `h`/`w` bound them
-            # and the table size is known on the host.
-            max_grid_size = max(max(h, w) for _, h, w in flatten_image_grid_thw)
-            rope_emb_max_grid = self.rotary_pos_emb(max_grid_size)
-            rope_emb = rope_emb_max_grid[pids].flatten(1)
-            rope_emb = rope_emb.repeat(1, 2)
-            rope_emb = (rope_emb.cos(), rope_emb.sin())
-        else:
-            rope_emb = None
-
-        attn_cu_seqlens = cu_seqlens
-        hidden_states = inputs_embeds
-        assert attention_mask is None
-
-        for encoder_layer in self.layers:
-            hidden_states = encoder_layer(
-                hidden_states,
-                attention_mask,
-                output_attentions=output_attentions,
-                cu_seqlens=attn_cu_seqlens,
-                rope_emb=rope_emb,
-            )
-        return hidden_states
-
-
-class KeyeSiglipVisionTransformer(nn.Module):
-    def __init__(
-        self,
-        config: PreTrainedConfig,
-        quant_config: QuantizationConfig | None = None,
-        prefix: str = "",
-    ):
-        super().__init__()
-        self.config = config
-        embed_dim = config.hidden_size
-
-        self.embeddings = KeyeVisionEmbeddings(config)
-        self.encoder = KeyeSiglipEncoder(
-            config,
-            quant_config=quant_config,
-            prefix=f"{prefix}.encoder",
-        )
-        self.post_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
-
-    def forward(
-        self,
-        pixel_values,
-        output_attentions: bool | None = None,
-        output_hidden_states: bool | None = None,
-        interpolate_pos_encoding: bool | None = False,
-        attention_mask: torch.Tensor | None = None,
-        sample_indices: torch.Tensor | None = None,
-        image_indices: torch.Tensor | None = None,
-        position_ids: torch.Tensor | None = None,
-        height_position_ids: torch.Tensor | None = None,
-        width_position_ids: torch.Tensor | None = None,
-        cu_seqlens: torch.Tensor | None = None,
-        padding_mask: torch.Tensor | None = None,
-        vision_return_embed_list: bool | None = False,
-        image_grid_thw: list[tuple[int, int, int]] | None = None,
-        return_pooler_output: bool | None = True,
-        use_rope: bool | None = False,
-        window_size: int = -1,
-    ) -> BaseModelOutputWithPooling:
-        hidden_states = self.embeddings(
-            pixel_values,
-            interpolate_pos_encoding=interpolate_pos_encoding,
-            position_ids=position_ids,
-            image_grid_thw=image_grid_thw,
-        )
-
-        last_hidden_state = self.encoder(
-            inputs_embeds=hidden_states,
-            output_attentions=output_attentions,
-            output_hidden_states=output_hidden_states,
-            attention_mask=attention_mask,
-            cu_seqlens=cu_seqlens,
-            image_grid_thw=image_grid_thw,
-            use_rope=use_rope,
-            height_position_ids=height_position_ids,
-            width_position_ids=width_position_ids,
-            window_size=window_size,
-            vision_or_text="vision",
-        )
-
-        last_hidden_state = self.post_layernorm(last_hidden_state)
-
-        if cu_seqlens is None:
-            raise ValueError(
-                "cu_seqlens cannot be None for "
-                "SiglipVisionTransformer output processing."
-            )
-        # `cu_seqlens` is the running sum of the per-image `t * h * w`, so the
-        # split sizes are known on the host and slicing needs no device read.
-        split_sizes = [
-            int(np.prod(thw)) for thw in self.encoder.flatten_list(image_grid_thw)
-        ]
-        return [
-            tensor.squeeze(0) for tensor in last_hidden_state.split(split_sizes, dim=1)
-        ]
-
-
-class KeyeSiglipVisionModel(nn.Module):
-    config_class = PreTrainedConfig
-    main_input_name = "pixel_values"
-
-    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={"vision_model.head.": None},
-    )
-
-    def __init__(
-        self,
-        config: PreTrainedConfig,
-        quant_config: QuantizationConfig | None = None,
-        prefix: str = "",
-    ):
-        super().__init__()
-
-        self.vision_model = KeyeSiglipVisionTransformer(
-            config,
-            quant_config=quant_config,
-            prefix=maybe_prefix(prefix, "vision_model"),
-        )
-        self.quant_config = quant_config
-
-    @property
-    def dtype(self) -> torch.dtype:
-        return self.vision_model.embeddings.patch_embedding.weight.dtype
-
-    @property
-    def device(self) -> torch.device:
-        return self.vision_model.embeddings.patch_embedding.weight.device
-
-    def get_input_embeddings(self) -> nn.Module:
-        return self.vision_model.embeddings.patch_embedding
-
-    def forward(
-        self,
-        pixel_values,
-        sample_indices: torch.Tensor | None = None,
-        output_attentions: bool | None = None,
-        output_hidden_states: bool | None = None,
-        interpolate_pos_encoding: bool = False,
-        position_ids: torch.Tensor | None = None,
-        vision_return_embed_list: bool | None = False,
-        image_grid_thw: list[tuple[int, int, int]] | None = None,
-        cu_seqlens: torch.Tensor | None = None,
-        return_pooler_output: bool | None = True,
-        use_rope: bool | None = False,
-        window_size: int = -1,
-    ) -> BaseModelOutputWithPooling:
-        return self.vision_model(
-            pixel_values=pixel_values,
-            output_attentions=output_attentions,
-            output_hidden_states=output_hidden_states,
-            interpolate_pos_encoding=interpolate_pos_encoding,
-            position_ids=position_ids,
-            vision_return_embed_list=vision_return_embed_list,
-            image_grid_thw=image_grid_thw,
-            sample_indices=sample_indices,
-            cu_seqlens=cu_seqlens,
-            return_pooler_output=return_pooler_output,
-            use_rope=use_rope,
-            window_size=window_size,
-        )
-
-    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 Projector(nn.Module):
-    def __init__(
-        self,
-        text_config: PreTrainedConfig,
-        vision_config: PreTrainedConfig,
-        quant_config: QuantizationConfig | None = None,
-        prefix: str = "",
-    ):
-        super().__init__()
-        self.text_config = text_config
-        self.vision_config = vision_config
-        self.merge_kernel_size = (2, 2)
-
-        self.hidden_size = (
-            self.vision_config.hidden_size
-            * self.merge_kernel_size[0]
-            * self.merge_kernel_size[1]
-        )
-
-        self.pre_norm = torch.nn.LayerNorm(self.vision_config.hidden_size, eps=1e-05)
-        self.act = GELUActivation()
-
-        self.linear_1 = ColumnParallelLinear(
-            self.hidden_size,
-            self.hidden_size,
-            bias=True,
-            quant_config=quant_config,
-            prefix=f"{prefix}.linear_1",
-        )
-        self.linear_2 = RowParallelLinear(
-            self.hidden_size,
-            self.text_config.hidden_size,
-            bias=True,
-            quant_config=quant_config,
-            prefix=f"{prefix}.linear_2",
-        )
-
-    def forward(
-        self,
-        image_features: torch.Tensor | list[torch.Tensor],
-        image_grid_thw: list[tuple[int, int, int]],
-    ) -> torch.Tensor | list[torch.Tensor]:
-        m1, m2 = self.merge_kernel_size
-        if isinstance(image_features, (list, tuple)):
-            processed_features = list()
-            for image_feature, image_grid in zip(image_features, image_grid_thw):
-                image_feature = self.pre_norm(image_feature)
-                t, h, w = image_grid
-
-                image_feature = rearrange(
-                    image_feature,
-                    "(t h p1 w p2) d -> (t h w) (p1 p2 d)",
-                    t=t,
-                    h=h // m1,
-                    p1=m1,
-                    w=w // m2,
-                    p2=m2,
-                )
-                hidden_states, _ = self.linear_1(image_feature)
-                hidden_states = self.act(hidden_states)
-                hidden_states, _ = self.linear_2(hidden_states)
-                processed_features.append(hidden_states)
-
-            return processed_features
-
-        dims = image_features.shape[:-1]
-        dim = image_features.shape[-1]
-        image_features = image_features.view(np.prod(dims), dim)
-        hidden_states = self.pre_norm(image_features).view(-1, self.hidden_size)
-        hidden_states = self.linear_1(hidden_states)
-        hidden_states = self.act(hidden_states)
-        hidden_states = self.linear_2(hidden_states)
-
-        return hidden_states.view(*dims, -1)
-
-
-def _keye_field_config(
-    hf_inputs: Mapping[str, torch.Tensor],
-):
-    image_grid_thw = hf_inputs.get("image_grid_thw", torch.empty((0, 3)))
-    image_grid_sizes = image_grid_thw.prod(-1)
-
-    video_grid_thw = hf_inputs.get("video_grid_thw", torch.empty((0, 3)))
-    video_grid_sizes = video_grid_thw.prod(-1)
-
-    return dict(
-        pixel_values=MultiModalFieldConfig.flat_from_sizes("image", image_grid_sizes),
-        image_embeds=MultiModalFieldConfig.flat_from_sizes("image", image_grid_sizes),
-        image_grid_thw=MultiModalFieldConfig.batched("image", keep_on_cpu=True),
-        pixel_values_videos=MultiModalFieldConfig.flat_from_sizes(
-            "video", video_grid_sizes
-        ),
-        video_embeds=MultiModalFieldConfig.flat_from_sizes("video", video_grid_sizes),
-        video_grid_thw=MultiModalFieldConfig.batched("video", keep_on_cpu=True),
-    )
-
-
-class KeyeMultiModalDataParser(MultiModalDataParser):
-    # The patch grid is what sizes the placeholder range.
-    embedding_fields = {
-        "image": {"image_embeds": "values", "image_grid_thw": "metadata"},
-        "video": {"video_embeds": "values", "video_grid_thw": "metadata"},
-    }
-
-    def _parse_image_data(
-        self,
-        data: dict[str, torch.Tensor] | ModalityData[ImageItem],
-    ) -> ModalityDataItems[Any, Any] | None:
-        if isinstance(data, dict):
-            required, optional = self.embedding_field_sets("image")
-            return DictEmbeddingItems(
-                data,
-                modality="image",
-                required_fields=required,
-                optional_fields=optional,
-                fields_factory=_keye_field_config,
-            )
-
-        return super()._parse_image_data(data)
-
-    def _parse_video_data(
-        self,
-        data: dict[str, torch.Tensor] | ModalityData[VideoItem],
-    ) -> ModalityDataItems[Any, Any] | None:
-        if isinstance(data, dict):
-            required, optional = self.embedding_field_sets("video")
-            return DictEmbeddingItems(
-                data,
-                modality="video",
-                required_fields=required,
-                optional_fields=optional,
-                fields_factory=_keye_field_config,
-            )
-
-        return super()._parse_video_data(data)
-
-
-class KeyeProcessingInfo(BaseProcessingInfo):
-    def get_max_image_size(self) -> int:
-        return 9999999  # _MAX_IMAGE_SIZE
-
-    def get_max_frame_per_video(self) -> int:
-        return 16  # _MAX_FRAMES_PER_VIDEO
-
-    def get_image_processor(self, **kwargs: object):
-        return self.get_hf_processor(**kwargs).image_processor
-
-    def get_data_parser(self):
-        return KeyeMultiModalDataParser(
-            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 {"image": None, "video": None}
-
-    def get_mm_max_tokens_per_item(
-        self,
-        seq_len: int,
-        mm_counts: Mapping[str, int],
-    ) -> Mapping[str, int]:
-        return {
-            "image": self.get_max_image_tokens(),
-            "video": self.get_max_video_tokens(seq_len),
-        }
-
-    def _get_vision_info(
-        self,
-        *,
-        image_width: int,
-        image_height: int,
-        num_frames: int = 1,
-        do_resize: bool = True,
-        image_processor: BaseImageProcessor,
-        mm_kwargs: Mapping[str, object],
-    ) -> tuple[ImageSize, int]:
-        hf_config = self.get_hf_config()
-        vision_config = hf_config.vision_config
-        patch_size = vision_config.patch_size
-        merge_size = vision_config.spatial_merge_size
-        temporal_patch_size = 1
-
-        mm_kwargs = self.ctx.get_merged_mm_kwargs(mm_kwargs)
-        size = image_processor.size
-        if override_size := mm_kwargs.get("size"):
-            size = size | override_size
-        if (override_min_pixels := mm_kwargs.get("min_pixels")) is not None:
-            size = size | {"min_pixels": override_min_pixels}
-        if (override_max_pixels := mm_kwargs.get("max_pixels")) is not None:
-            size = size | {"max_pixels": override_max_pixels}
-
-        if do_resize:
-            resized_height, resized_width = smart_resize(
-                height=image_height,
-                width=image_width,
-                factor=patch_size * merge_size,
-                min_pixels=size["min_pixels"],
-                max_pixels=size["max_pixels"],
-            )
-            preprocessed_size = ImageSize(width=resized_width, height=resized_height)
-        else:
-            preprocessed_size = ImageSize(width=image_width, height=image_height)
-
-        padded_num_frames = num_frames + (-num_frames % temporal_patch_size)
-
-        grid_t = max(padded_num_frames // temporal_patch_size, 1)
-        grid_h = preprocessed_size.height // patch_size
-        grid_w = preprocessed_size.width // patch_size
-
-        num_patches = grid_t * grid_h * grid_w
-        num_vision_tokens = num_patches // (merge_size**2)
-
-        return preprocessed_size, num_vision_tokens
-
-    def get_num_image_tokens(
-        self,
-        *,
-        image_width: int,
-        image_height: int,
-        image_processor: BaseImageProcessor,
-        mm_kwargs: Mapping[str, object],
-    ) -> int:
-        _, num_image_tokens = self._get_vision_info(
-            image_width=image_width,
-            image_height=image_height,
-            image_processor=image_processor,
-            mm_kwargs=mm_kwargs,
-        )
-        return num_image_tokens
-
-    def get_num_video_tokens(
-        self,
-        *,
-        image_width: int,
-        image_height: int,
-        num_frames: int,
-        image_processor: BaseImageProcessor,
-        mm_kwargs: Mapping[str, object],
-    ) -> int:
-        _, num_video_tokens = self._get_vision_info(
-            image_width=image_width,
-            image_height=image_height,
-            num_frames=num_frames,
-            image_processor=image_processor,
-            mm_kwargs=mm_kwargs,
-        )
-        return num_video_tokens
-
-    def get_image_size_with_most_features(self) -> ImageSize:
-        image_processor = self.get_image_processor()
-
-        max_image_size, _ = self._get_vision_info(
-            image_width=self.get_max_image_size(),
-            image_height=self.get_max_image_size(),
-            image_processor=image_processor,
-            mm_kwargs={},
-        )
-        return max_image_size
-
-    def get_max_image_tokens(self) -> int:
-        image_processor = self.get_image_processor()
-        target_width, target_height = self.get_image_size_with_most_features()
-
-        return self.get_num_image_tokens(
-            image_width=target_width,
-            image_height=target_height,
-            image_processor=image_processor,
-            mm_kwargs={},
-        )
-
-    def _get_max_video_frames(self, max_tokens: int) -> int:
-        image_processor = self.get_image_processor()
-        target_width, target_height = self.get_image_size_with_most_features()
-
-        num_frames = 0
-
-        while True:
-            next_num_frames = num_frames + 1
-            next_max_tokens = self.get_num_video_tokens(
-                image_width=target_width,
-                image_height=target_height,
-                num_frames=next_num_frames,
-                image_processor=image_processor,
-                mm_kwargs={},
-            )
-
-            if next_max_tokens > max_tokens:
-                break
-
-            num_frames = next_num_frames
-
-        return num_frames
-
-    def get_num_frames_with_most_features(self, seq_len: int) -> int:
-        mm_config = self.ctx.get_mm_config()
-        max_images = mm_config.get_limit_per_prompt("image")
-        max_videos = mm_config.get_limit_per_prompt("video")
-
-        max_image_tokens = self.get_max_image_tokens() * max_images
-        max_total_frames = self._get_max_video_frames(seq_len - max_image_tokens)
-        max_frames_per_video = min(
-            max_total_frames // max(max_videos, 1),
-            self.get_max_frame_per_video(),
-        )
-
-        return max(max_frames_per_video, 1)
-
-    def get_max_video_tokens(self, seq_len: int) -> int:
-        image_processor = self.get_image_processor()
-        target_width, target_height = self.get_image_size_with_most_features()
-
-        return self.get_num_video_tokens(
-            image_width=target_width,
-            image_height=target_height,
-            num_frames=self.get_num_frames_with_most_features(seq_len),
-            image_processor=image_processor,
-            mm_kwargs={},
-        )
-
-
-_I = TypeVar("_I", bound=KeyeProcessingInfo)
-
-
-class KeyeBaseDummyInputsBuilder(BaseDummyInputsBuilder[_I]):
-    def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
-        num_images = mm_counts.get("image", 0)
-        num_videos = mm_counts.get("video", 0)
-
-        hf_processor = self.info.get_hf_processor()
-        image_token: str = hf_processor.image_token
-        video_token: str = hf_processor.video_token
-
-        return image_token * num_images + video_token * num_videos
-
-    def get_dummy_mm_data(
-        self,
-        seq_len: int,
-        mm_counts: Mapping[str, int],
-        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_data = {
-            "image": self._get_dummy_images(
-                width=target_width,
-                height=target_height,
-                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=mm_counts.get("video", 0),
-                overrides=mm_options.get("video"),
-            ),
-        }
-
-        return mm_data
-
-
-class KeyeDummyInputsBuilder(KeyeBaseDummyInputsBuilder[KeyeProcessingInfo]):
-    pass
-
-
-class KeyeMultiModalProcessor(BaseMultiModalProcessor[KeyeProcessingInfo]):
-    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,
-        mm_items: MultiModalDataItems,
-        hf_processor_mm_kwargs: Mapping[str, Any],
-        out_mm_kwargs: MultiModalKwargsItems,
-    ) -> Sequence[PromptUpdate]:
-        hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)
-        image_processor = self.info.get_image_processor(**hf_processor_mm_kwargs)
-        tokenizer = self.info.get_tokenizer()
-        vocab = tokenizer.get_vocab()
-
-        placeholder = {
-            "image": vocab[hf_processor.image_token],
-            "video": vocab[hf_processor.video_token],
-        }
-
-        merge_length = image_processor.merge_size**2
-
-        def get_replacement_keye(item_idx: int, modality: str):
-            out_item = out_mm_kwargs[modality][item_idx]
-            grid_thw = out_item[f"{modality}_grid_thw"].data
-            assert isinstance(grid_thw, torch.Tensor)
-
-            num_tokens = int(grid_thw.prod()) // merge_length
-            return [placeholder[modality]] * num_tokens
-
-        return [
-            PromptReplacement(
-                modality=modality,
-                target=[placeholder[modality]],
-                replacement=partial(get_replacement_keye, modality=modality),
-            )
-            for modality in ("image", "video")
-        ]
-
-    def _get_mm_fields_config(
-        self,
-        hf_inputs: BatchFeature,
-        hf_processor_mm_kwargs: Mapping[str, object],
-    ) -> Mapping[str, MultiModalFieldConfig]:
-        return _keye_field_config(hf_inputs)
-
-
-class BaseKeyeModule(nn.Module, SupportsMultiModal):
-    packed_modules_mapping = {
-        "qkv_proj": [
-            "q_proj",
-            "k_proj",
-            "v_proj",
-        ],
-        "gate_up_proj": [
-            "gate_proj",
-            "up_proj",
-        ],
-    }
-
-    hf_to_vllm_mapper = WeightsMapper(
-        orig_to_new_prefix={
-            "lm_head.": "language_model.lm_head.",
-            "model.": "language_model.model.",
-        }
-    )
-
-    supports_tower_connector_lora = True
-
-    @classmethod
-    def get_placeholder_str(cls, modality: str, i: int) -> str | None:
-        if modality.startswith("image"):
-            return "<|vision_start|><|image_pad|><|vision_end|>"
-        if modality.startswith("video"):
-            return "<|vision_start|><|video_pad|><|vision_end|>"
-
-        raise ValueError("Only image or video modality is supported")
-
-    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
-        super().__init__()
-        config: PreTrainedConfig = vllm_config.model_config.hf_config
-        quant_config = vllm_config.quant_config
-
-        self.config = config
-
-        with self._mark_tower_model(vllm_config, {"image", "video"}):
-            self.visual = KeyeSiglipVisionModel(
-                config.vision_config,
-                quant_config=quant_config,
-                prefix=maybe_prefix(prefix, "visual"),
-            )
-            self.mlp_AR = self._build_projector(
-                config,
-                config.vision_config,
-                quant_config=quant_config,
-                prefix=maybe_prefix(prefix, "mlp_AR"),
-            )
-
-        with self._mark_language_model(vllm_config):
-            self.language_model = init_vllm_registered_model(
-                vllm_config=vllm_config,
-                prefix=maybe_prefix(prefix, "language_model"),
-                architectures=["Qwen3ForCausalLM"],
-            )
-
-        self.make_empty_intermediate_tensors = (
-            self.language_model.make_empty_intermediate_tensors
-        )
-
-    @abstractmethod
-    def _build_projector(
-        self,
-        text_config: PreTrainedConfig,
-        vision_config: PreTrainedConfig,
-        quant_config: QuantizationConfig | None = None,
-        prefix: str = "",
-    ) -> nn.Module:
-        raise NotImplementedError("Need projector")
-
-    def _process_image_input(self, image_input: Any) -> tuple[torch.Tensor, ...]:
-        siglip_position_ids = list()
-        image_grid_hws = list()
-        sample_indices = list()
-        cu_seqlens = [0]
-
-        image_grid_thw = image_input["image_grid_thw"]
-        assert image_grid_thw.ndim == 2
-
-        image_grid_thw_list = image_grid_thw.tolist()
-        for idx, thw in enumerate(image_grid_thw_list):
-            thw_tuple = tuple(thw)
-            numel = np.prod(thw_tuple)
-            image_grid_hws.append(thw_tuple)
-            image_position_ids = torch.arange(numel) % np.prod(thw_tuple[1:])
-            siglip_position_ids.append(image_position_ids)
-            sample_indices.append(torch.full((numel,), idx, dtype=torch.int64))
-            cu_seqlens.append(cu_seqlens[-1] + numel)
-
-        if image_input["type"] == "image_embeds":
-            raise ValueError(
-                "Image embeddings are not supported for this processing path."
-            )
-        else:
-            pixel_values = image_input["pixel_values"].type(self.visual.dtype)
-            # These are all built on the host; concat straight into pinned
-            # buffers so the H2D copies stay non-blocking.
-            siglip_position_ids = torch.concat(
-                siglip_position_ids,
-                dim=0,
-                out=torch.empty(
-                    sum(t.numel() for t in siglip_position_ids),
-                    dtype=torch.int64,
-                    pin_memory=PIN_MEMORY,
-                ),
-            ).to(pixel_values.device, non_blocking=True)
-            cu_seqlens = async_tensor_h2d(
-                cu_seqlens, dtype=torch.int32, device=pixel_values.device
-            )
-            sample_indices = torch.concat(
-                sample_indices,
-                dim=0,
-                out=torch.empty(
-                    sum(t.numel() for t in sample_indices),
-                    dtype=torch.int64,
-                    pin_memory=PIN_MEMORY,
-                ),
-            ).to(pixel_values.device, non_blocking=True)
-
-            image_embeds = self.visual(
-                pixel_values=pixel_values,
-                image_grid_thw=image_grid_hws,
-                position_ids=siglip_position_ids,
-                vision_return_embed_list=False,
-                interpolate_pos_encoding=True,
-                sample_indices=sample_indices,
-                cu_seqlens=cu_seqlens,
-                use_rope=True,
-                window_size=-1,
-            )
-            image_embeds = tuple(self.mlp_AR(image_embeds, image_grid_thw))
-            return image_embeds
-
-    def _process_video_embeds(
-        self,
-        video_type: Literal["video_embeds", "pixel_values_videos"],
-        video_grid_thw: torch.Tensor,
-        pixel_values_videos: torch.Tensor | None = None,
-    ) -> torch.Tensor | list[torch.Tensor]:
-        siglip_position_ids = list()
-        video_grid_hws = list()
-        sample_indices = list()
-        cu_seqlens = [0]
-
-        assert video_grid_thw.ndim == 2
-        for idx, sub_thw in enumerate(video_grid_thw):
-            thw_tuple = tuple(sub_thw.detach().cpu().numpy().tolist())
-            numel = np.prod(thw_tuple)
-
-            video_grid_hws.append(thw_tuple)
-            video_position_ids = torch.arange(numel) % np.prod(thw_tuple[1:])
-            siglip_position_ids.append(video_position_ids)
-            sample_indices.append(torch.full((numel,), idx, dtype=torch.int64))
-            cu_seqlens.append(cu_seqlens[-1] + numel)
-
-        if video_type == "video_embeds":
-            raise ValueError(
-                "Video embeddings are not supported for this processing path."
-            )
-        else:
-            assert pixel_values_videos is not None
-            pixel_values_videos = pixel_values_videos.type(self.visual.dtype)
-            # These are host-built; concat straight into pinned buffers so
-            # the H2D copies stay non-blocking.
-            siglip_position_ids = torch.concat(
-                siglip_position_ids,
-                dim=0,
-                out=torch.empty(
-                    sum(t.numel() for t in siglip_position_ids),
-                    dtype=torch.int64,
-                    pin_memory=PIN_MEMORY,
-                ),
-            ).to(pixel_values_videos.device, non_blocking=True)
-            cu_seqlens = async_tensor_h2d(
-                cu_seqlens, dtype=torch.int32, device=pixel_values_videos.device
-            )
-            sample_indices = torch.concat(
-                sample_indices,
-                dim=0,
-                out=torch.empty(
-                    sum(t.numel() for t in sample_indices),
-                    dtype=torch.int64,
-                    pin_memory=PIN_MEMORY,
-                ),
-            ).to(pixel_values_videos.device, non_blocking=True)
-
-            video_embeds = self.visual(
-                pixel_values=pixel_values_videos,
-                image_grid_thw=video_grid_hws,
-                position_ids=siglip_position_ids,
-                vision_return_embed_list=True,
-                interpolate_pos_encoding=True,
-                sample_indices=sample_indices,
-                cu_seqlens=cu_seqlens,
-                use_rope=True,
-                window_size=-1,
-            )
-            video_embeds = self.mlp_AR(video_embeds, video_grid_thw)
-            return video_embeds
-
-    def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict:
-        modalities = {}
-
-        for input_key in kwargs:
-            if (
-                input_key in ("pixel_values", "image_embeds")
-                and "images" not in modalities
-            ):
-                modalities["images"] = self._parse_and_validate_image_input(**kwargs)
-            if (
-                input_key in ("pixel_values_videos", "video_embeds")
-                and "videos" not in modalities
-            ):
-                modalities["videos"] = self._parse_and_validate_video_input(**kwargs)
-
-        return modalities
-
-    def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings | None:
-        modalities = self._parse_and_validate_multimodal_inputs(**kwargs)
-        if not modalities:
-            return None
-
-        multimodal_embeddings: tuple[torch.Tensor, ...] = ()
-
-        for modality in modalities:
-            if modality == "images":
-                image_input = modalities["images"]
-                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)
-                multimodal_embeddings += tuple(video_embeddings)
-        return multimodal_embeddings
-
-    def forward(
-        self,
-        input_ids: torch.Tensor | None,
-        positions: torch.Tensor,
-        intermediate_tensors: IntermediateTensors | None = None,
-        inputs_embeds: torch.Tensor | None = None,
-        **kwargs: object,
-    ) -> torch.Tensor | IntermediateTensors:
-        """Run forward pass for Keye-VL.
-
-        Args:
-            input_ids: Flattened (concatenated) input_ids corresponding to a
-                batch.
-            positions: Flattened (concatenated) position ids corresponding to a
-                batch.
-                **NOTE**: If mrope is enabled (default setting for Qwen2-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.
-            **kwargs: Multimodal inputs for this batch, forwarded to the
-                multimodal embedding path.
-
-        """
-        if intermediate_tensors is not None:
-            inputs_embeds = None
-
-        hidden_states = self.language_model.model(
-            input_ids=input_ids,
-            positions=positions,
-            intermediate_tensors=intermediate_tensors,
-            inputs_embeds=inputs_embeds,
-        )
-
-        return hidden_states
-
-    def compute_logits(
-        self,
-        hidden_states: torch.Tensor,
-    ) -> torch.Tensor | None:
-        return self.language_model.compute_logits(hidden_states)
-
-    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)
-
-    def get_mm_mapping(self) -> MultiModelKeys:
-        """Get the module prefix in multimodal models."""
-        return MultiModelKeys.from_string_field(
-            language_model="language_model",
-            connector="mlp_AR.",
-            tower_model="visual.",
-        )
-
-    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_mm_embeds * merge_size**2, num_mm_embeds
-
-
-@MULTIMODAL_REGISTRY.register_processor(
-    KeyeMultiModalProcessor,
-    info=KeyeProcessingInfo,
-    dummy_inputs=KeyeDummyInputsBuilder,
-)
-class KeyeForConditionalGeneration(
-    BaseKeyeModule, SupportsMultiModal, SupportsLoRA, SupportsPP, SupportsMRoPE
-):
-    def _build_projector(
-        self,
-        text_config: PreTrainedConfig,
-        vision_config: PreTrainedConfig,
-        quant_config: QuantizationConfig | None = None,
-        prefix: str = "",
-    ) -> nn.Module:
-        return Projector(text_config, vision_config, quant_config, prefix)
-
-    def _parse_and_validate_image_input(
-        self, **kwargs: object
-    ) -> KeyeImageInputs | None:
-        pixel_values = kwargs.pop("pixel_values", None)
-        image_embeds = kwargs.pop("image_embeds", None)
-        image_grid_thw = kwargs.pop("image_grid_thw", None)
-
-        if pixel_values is None and image_embeds is None:
-            return None
-
-        if pixel_values is not None:
-            return KeyeImagePixelInputs(
-                type="pixel_values",
-                pixel_values=pixel_values,
-                image_grid_thw=image_grid_thw,
-            )
-
-        return KeyeImageEmbeddingInputs(
-            type="image_embeds",
-            image_embeds=image_embeds,
-            image_grid_thw=image_grid_thw,
-        )
-
-    def _parse_and_validate_video_input(
-        self, **kwargs: object
-    ) -> KeyeVideoInputs | 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)
-
-        if pixel_values_videos is None and video_embeds is None:
-            return None
-
-        if pixel_values_videos is not None:
-            return KeyeVideoPixelInputs(
-                type="pixel_values_videos",
-                pixel_values_videos=pixel_values_videos,
-                video_grid_thw=video_grid_thw,
-            )
-
-        return KeyeVideoEmbeddingInputs(
-            type="video_embeds",
-            video_embeds=video_embeds,
-            video_grid_thw=video_grid_thw,
-        )
-
-    def _process_video_input(
-        self, video_input: KeyeVideoInputs
-    ) -> tuple[torch.Tensor, ...]:
-        video_type = video_input["type"]
-        video_grid_thw = video_input["video_grid_thw"]
-        pixel_values_videos = video_input.get("pixel_values_videos", None)
-
-        return tuple(
-            self._process_video_embeds(video_type, video_grid_thw, pixel_values_videos)
-        )
-
-    @staticmethod
-    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.
-
-        This preserves Keye's current M-RoPE behavior, where a video is emitted
-        as consecutive frame-level multimodal blocks rather than a single block
-        spanning the whole video.
-        """
-        if isinstance(grid_thw, list):
-            if len(grid_thw) == 0:
-                return []
-            if isinstance(grid_thw[0], int):
-                grid_thw = torch.tensor([grid_thw], dtype=torch.long)
-            else:
-                grid_thw = torch.tensor(grid_thw, dtype=torch.long)
-        elif grid_thw.ndim == 1:
-            grid_thw = grid_thw.unsqueeze(0)
-
-        if grid_thw.numel() == 0:
-            return []
-
-        t, hw = grid_thw[:, 0], grid_thw[:, 1:]
-        ones = torch.ones_like(hw[:, :1])
-        out = torch.cat([ones, hw], dim=1).repeat_interleave(t, dim=0)
-        return out.tolist()
-
-    def iter_mm_grid_thw(
-        self, mm_features: list[MultiModalFeatureSpec]
-    ) -> Iterator[tuple[int, int, int, int]]:
-        spatial_merge_size = self.config.vision_config.spatial_merge_size
-
-        for mm_feature in sorted(mm_features, key=lambda f: f.mm_position.offset):
-            if mm_feature.data is None:
-                raise ValueError("M-RoPE calculation requires multimodal feature data")
-
-            if mm_feature.modality == "image":
-                grid_thw = mm_feature.data["image_grid_thw"].data
-                if isinstance(grid_thw, torch.Tensor):
-                    if grid_thw.ndim == 2:
-                        assert grid_thw.shape[0] == 1
-                        t, h, w = grid_thw[0].tolist()
-                    else:
-                        t, h, w = grid_thw.tolist()
-                else:
-                    if isinstance(grid_thw[0], list):
-                        assert len(grid_thw) == 1
-                        t, h, w = grid_thw[0]
-                    else:
-                        t, h, w = grid_thw
-
-                yield (
-                    mm_feature.mm_position.offset,
-                    t,
-                    h // spatial_merge_size,
-                    w // spatial_merge_size,
-                )
-            elif mm_feature.modality == "video":
-                current_offset = mm_feature.mm_position.offset
-                for t, h, w in self._split_video_grid_thw(
-                    mm_feature.data["video_grid_thw"].data
-                ):
-                    llm_grid_h = h // spatial_merge_size
-                    llm_grid_w = w // spatial_merge_size
-                    yield (current_offset, t, llm_grid_h, llm_grid_w)
-                    current_offset += t * llm_grid_h * llm_grid_w
-            else:
-                raise ValueError(f"Unsupported modality: {mm_feature.modality}")
-
-    def get_mrope_input_positions(
-        self,
-        input_tokens: list[int],
-        mm_features: list[MultiModalFeatureSpec],
-    ) -> tuple[torch.Tensor, int]:
-        llm_pos_ids_list: list = []
-        st = 0
-
-        for (
-            offset,
-            llm_grid_t,
-            llm_grid_h,
-            llm_grid_w,
-        ) in self.iter_mm_grid_thw(mm_features):
-            text_len = offset - st
-
-            st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
-            llm_pos_ids_list.append(
-                torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx
-            )
-
-            t_index = (
-                (
-                    torch.arange(llm_grid_t)
-                    .view(-1, 1)
-                    .expand(-1, llm_grid_h * llm_grid_w)
-                )
-                .long()
-                .flatten()
-            )
-
-            h_index = (
-                torch.arange(llm_grid_h)
-                .view(1, -1, 1)
-                .expand(llm_grid_t, -1, llm_grid_w)
-                .flatten()
-            )
-            w_index = (
-                torch.arange(llm_grid_w)
-                .view(1, 1, -1)
-                .expand(llm_grid_t, llm_grid_h, -1)
-                .flatten()
-            )
-            llm_pos_ids_list.append(
-                torch.stack([t_index, h_index, w_index]) + text_len + st_idx
-            )
-            st = offset + llm_grid_t * llm_grid_h * llm_grid_w
-
-        if st < len(input_tokens):
-            st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
-            text_len = len(input_tokens) - st
-            llm_pos_ids_list.append(
-                torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx
-            )
-
-        llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)
-        mrope_position_delta = (llm_positions.max() + 1 - len(input_tokens)).item()
-
-        return llm_positions, mrope_position_delta
diff --git a/vllm/model_executor/models/keye_vl1_5.py b/vllm/model_executor/models/keye_vl1_5.py
deleted file mode 100644
index fc08035fa7dc..000000000000
--- a/vllm/model_executor/models/keye_vl1_5.py
+++ /dev/null
@@ -1,707 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import itertools
-from collections.abc import Iterator, Mapping, Sequence
-from functools import partial
-from typing import Annotated, Any, Literal, TypeAlias
-
-import numpy as np
-import torch
-import torch.nn as nn
-from einops import rearrange
-from transformers import PreTrainedConfig
-from transformers.activations import GELUActivation
-from transformers.feature_extraction_utils import BatchFeature
-
-from vllm.config import VllmConfig
-from vllm.inputs import ModalityData
-from vllm.logger import init_logger
-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.inputs import (
-    ImageItem,
-    MultiModalFeatureSpec,
-    MultiModalFieldConfig,
-    MultiModalKwargsItems,
-    VideoItem,
-)
-from vllm.multimodal.parse import (
-    DictEmbeddingItems,
-    ModalityDataItems,
-    MultiModalDataItems,
-    MultiModalDataParser,
-)
-from vllm.multimodal.processing import (
-    PromptReplacement,
-    PromptUpdate,
-    PromptUpdateDetails,
-)
-from vllm.utils.tensor_schema import TensorSchema, TensorShape
-
-from .interfaces import SupportsLoRA, SupportsMRoPE, SupportsMultiModal, SupportsPP
-from .keye import (
-    BaseKeyeModule,
-    BaseMultiModalProcessor,
-    KeyeBaseDummyInputsBuilder,
-    KeyeProcessingInfo,
-)
-
-logger = init_logger(__name__)
-
-
-def split_thw(grid_thw: torch.Tensor) -> torch.Tensor:
-    """Split grid_thw in t dimension.
-
-    Args:
-        grid_thw: [N, 3] tensor of [t, h, w]
-
-    Returns:
-        [Σt, 3] tensor where each row is [1, h, w]
-
-    Example:
-    >>> grid_thw = torch.tensor([[2, 3, 4], [1, 5, 6]])
-    >>> split_thw(grid_thw)
-    tensor([[1, 3, 4],
-           [1, 3, 4],
-           [1, 5, 6]])
-
-    """
-    t = grid_thw[:, 0]
-    h_w = grid_thw[:, 1:]
-    ones = torch.ones_like(h_w[:, :1])
-    return torch.cat([ones, h_w], dim=1).repeat_interleave(t, dim=0)
-
-
-def get_num_patches(
-    grid_thw: torch.Tensor, num_frames: list[int] | torch.Tensor
-) -> torch.Tensor:
-    """Return num_patches per video.
-
-    Args:
-        grid_thw: Tensor with shape [N, 3] containing temporal, height, width
-            dimensions
-        num_frames: List or tensor indicating the number of frames per video
-
-    Returns:
-        Tensor of the number of patches for each video
-
-    Examples:
-        >>> # Suppose there are 2 videos with a total of 3 grids
-        >>> grid_thw = torch.tensor(
-        ...     [
-        ...         [2, 2, 2],  # grid 0: 2*2*2=8 patches
-        ...         [2, 2, 2],  # grid 1: 2*2*2=8 patches
-        ...         [1, 1, 1],
-        ...     ]
-        ... )  # grid 2: 1*1*1=1 patches
-        >>> num_frames = [2, 1]  # The first video contains 2 grids,
-                                   the second contains 1 grid.
-        >>> 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()
-
-    num_grids_per_frame = grid_thw.prod(dim=1)
-    start_idx_per_video = [0, *itertools.accumulate(num_frames)]
-    num_patches = [
-        num_grids_per_frame[start_idx_per_video[i] : start_idx_per_video[i + 1]].sum()
-        for i in range(len(num_frames))
-    ]
-    return (
-        torch.stack(num_patches)
-        if num_patches
-        else torch.zeros(0, dtype=grid_thw.dtype, device=grid_thw.device)
-    )
-
-
-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)
-    """
-
-    type: Literal["pixel_values"]
-
-    pixel_values: Annotated[
-        torch.Tensor, TensorShape("bnp", 3, "ps", "ps", dynamic_dims={"bnp"})
-    ]
-
-    image_grid_thw: Annotated[torch.Tensor, TensorShape("ni", 3)]
-
-
-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)
-    """
-
-    type: Literal["image_embeds"]
-    image_embeds: Annotated[torch.Tensor, TensorShape("nf", "hs")]
-    image_grid_thw: Annotated[torch.Tensor, TensorShape("ni", 3)]
-
-
-KeyeVL1_5ImageInputs: TypeAlias = (
-    KeyeVL1_5ImagePixelInputs | KeyeVL1_5ImageEmbeddingInputs
-)
-
-
-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)
-    """
-
-    type: Literal["pixel_values_videos"]
-    pixel_values_videos: Annotated[
-        torch.Tensor, TensorShape("bnp", 3, "ps", "ps", dynamic_dims={"bnp"})
-    ]
-    video_grid_thw: Annotated[torch.Tensor, TensorShape("nv", 3)]
-
-    num_frames: torch.Tensor
-
-
-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)
-    """
-
-    type: Literal["video_embeds"]
-    video_embeds: Annotated[torch.Tensor, TensorShape("nf", "hs")]
-    video_grid_thw: Annotated[torch.Tensor, TensorShape("nv", 3)]
-    num_frames: torch.Tensor
-
-
-KeyeVL1_5VideoInputs: TypeAlias = (
-    KeyeVL1_5VideoPixelInputs | KeyeVL1_5VideoEmbeddingInputs
-)
-
-
-class KeyeVL1_5Projector(nn.Module):
-    def __init__(
-        self,
-        text_config: PreTrainedConfig,
-        vision_config: PreTrainedConfig,
-        quant_config: QuantizationConfig | None = None,
-        prefix: str = "",
-    ):
-        super().__init__()
-        self.text_config = text_config
-        self.vision_config = vision_config
-        self.merge_kernel_size = (2, 2)
-
-        self.hidden_size = (
-            self.vision_config.hidden_size
-            * self.merge_kernel_size[0]
-            * self.merge_kernel_size[1]
-        )
-
-        self.pre_norm = torch.nn.LayerNorm(self.hidden_size, eps=1e-05)
-        self.act = GELUActivation()
-
-        self.linear_1 = ColumnParallelLinear(
-            self.hidden_size,
-            self.hidden_size,
-            bias=True,
-            quant_config=quant_config,
-            prefix=f"{prefix}.linear_1",
-        )
-        self.linear_2 = RowParallelLinear(
-            self.hidden_size,
-            self.text_config.hidden_size,
-            bias=True,
-            quant_config=quant_config,
-            prefix=f"{prefix}.linear_2",
-        )
-
-    def forward(
-        self,
-        image_features: torch.Tensor | tuple[torch.Tensor] | list[torch.Tensor],
-        image_grid_thw: list[tuple[int, int, int]],
-    ) -> torch.Tensor | list[torch.Tensor]:
-        m1, m2 = self.merge_kernel_size
-        if isinstance(image_features, (list, tuple)):
-            processed_features = list()
-            for image_feature, image_grid in zip(image_features, image_grid_thw):
-                t, h, w = image_grid
-                image_feature = rearrange(
-                    image_feature,
-                    "(t h p1 w p2) d -> (t h w) (p1 p2 d)",
-                    t=t,
-                    h=h // m1,
-                    p1=m1,
-                    w=w // m2,
-                    p2=m2,
-                )
-                image_feature = self.pre_norm(image_feature)
-                hidden_states, _ = self.linear_1(image_feature)
-                hidden_states = self.act(hidden_states)
-                hidden_states, _ = self.linear_2(hidden_states)
-                processed_features.append(hidden_states)
-
-            return processed_features
-
-        dims = image_features.shape[:-1]
-        dim = image_features.shape[-1]
-        image_features = image_features.view(np.prod(dims), dim)
-        hidden_states = self.pre_norm(image_features.view(-1, self.hidden_size))
-        hidden_states = self.linear_1(hidden_states)
-        hidden_states = self.act(hidden_states)
-        hidden_states = self.linear_2(hidden_states)
-
-        return hidden_states.view(*dims, -1)
-
-
-def _keye_field_config(
-    hf_inputs: Mapping[str, torch.Tensor],
-):
-    image_grid_thw = hf_inputs.get(
-        "image_grid_thw", torch.empty((0, 3), dtype=torch.int64)
-    )
-    image_grid_sizes = image_grid_thw.prod(-1)
-
-    video_grid_thw = hf_inputs.get(
-        "video_grid_thw", torch.empty((0, 3), dtype=torch.int64)
-    )
-    video_grid_thw = split_thw(video_grid_thw)
-    num_frames = hf_inputs.get("num_frames", video_grid_thw[:, 0]).clone().tolist()
-
-    video_num_patches = get_num_patches(video_grid_thw, num_frames)
-
-    video_num_grids = []
-    if len(num_frames) > 0:
-        i = 0
-        j = 1
-        cur_frames = num_frames[i]
-        for t, _, _ in video_grid_thw.tolist():
-            cur_frames -= t
-            if cur_frames == 0:
-                video_num_grids.append(j)
-                i += 1
-                if i < len(num_frames):
-                    cur_frames = num_frames[i]
-                j = 1
-            else:
-                j += 1
-    video_num_grids = torch.tensor(video_num_grids)
-    return dict(
-        pixel_values=MultiModalFieldConfig.flat_from_sizes("image", image_grid_sizes),
-        image_embeds=MultiModalFieldConfig.flat_from_sizes("image", image_grid_sizes),
-        image_grid_thw=MultiModalFieldConfig.batched("image", keep_on_cpu=True),
-        pixel_values_videos=MultiModalFieldConfig.flat_from_sizes(
-            "video", video_num_patches
-        ),
-        video_embeds=MultiModalFieldConfig.flat_from_sizes("video", video_num_patches),
-        video_grid_thw=MultiModalFieldConfig.flat_from_sizes(
-            "video", video_num_grids, keep_on_cpu=True
-        ),
-        num_frames=MultiModalFieldConfig.batched("video", keep_on_cpu=True),
-    )
-
-
-class KeyeVL1_5MultiModalDataParser(MultiModalDataParser):
-    # The patch grid is what sizes the placeholder range.
-    embedding_fields = {
-        "image": {"image_embeds": "values", "image_grid_thw": "metadata"},
-        "video": {"video_embeds": "values", "video_grid_thw": "metadata"},
-    }
-
-    def _parse_image_data(
-        self,
-        data: dict[str, torch.Tensor] | ModalityData[ImageItem],
-    ) -> ModalityDataItems[Any, Any] | None:
-        if isinstance(data, dict):
-            required, optional = self.embedding_field_sets("image")
-            return DictEmbeddingItems(
-                data,
-                modality="image",
-                required_fields=required,
-                optional_fields=optional,
-                fields_factory=_keye_field_config,
-            )
-
-        return super()._parse_image_data(data)
-
-    def _parse_video_data(
-        self,
-        data: dict[str, torch.Tensor] | ModalityData[VideoItem],
-    ) -> ModalityDataItems[Any, Any] | None:
-        if isinstance(data, dict):
-            required, optional = self.embedding_field_sets("video")
-            return DictEmbeddingItems(
-                data,
-                modality="video",
-                required_fields=required,
-                optional_fields=optional,
-                fields_factory=_keye_field_config,
-            )
-
-        return super()._parse_video_data(data)
-
-
-class KeyeVL1_5ProcessingInfo(KeyeProcessingInfo):
-    def get_data_parser(self):
-        return KeyeVL1_5MultiModalDataParser(
-            expected_hidden_size=self._get_expected_hidden_size(),
-            allow_missing_mm_embeddings=self.allow_missing_mm_embeddings,
-        )
-
-    def get_max_frame_per_video(self) -> int:
-        return 2048
-
-    def get_supported_mm_limits(
-        self,
-    ) -> Mapping[str, int | None]:
-        return {"image": None, "video": 1}
-
-
-class KeyeVL1_5MultiModalProcessor(BaseMultiModalProcessor[KeyeVL1_5ProcessingInfo]):
-    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,
-        mm_items: MultiModalDataItems,
-        hf_processor_mm_kwargs: Mapping[str, Any],
-        out_mm_kwargs: MultiModalKwargsItems,
-    ) -> Sequence[PromptUpdate]:
-        hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)
-        image_processor = self.info.get_image_processor(**hf_processor_mm_kwargs)
-        tokenizer = self.info.get_tokenizer()
-        vocab = tokenizer.get_vocab()
-        image_token_id = vocab[hf_processor.image_token]
-        video_token_id = vocab[hf_processor.video_token]
-        placeholder = {"image": image_token_id, "video": video_token_id}
-        merge_length = image_processor.merge_size**2
-
-        out_mm_kwargs_data = out_mm_kwargs.get_data()
-        frame_types: list[torch.Tensor] = hf_processor_mm_kwargs.get(
-            "frame_types", None
-        )
-        timestamps: list[torch.Tensor] = hf_processor_mm_kwargs.get("timestamps", None)
-        num_videos = mm_items.get_count("video", strict=False)
-
-        if frame_types is None:
-            frame_types = [None] * num_videos
-        assert len(frame_types) == num_videos, (
-            f"Number of frame_types={len(frame_types)} "
-            f"doesn't equal to number of videos={num_videos}"
-        )
-        if timestamps is None:
-            timestamps = [None] * num_videos
-        assert len(timestamps) == num_videos, (
-            f"Number of timestamps={len(timestamps)} "
-            f"doesn't equal to number of videos={num_videos}"
-        )
-
-        video_grid_thw = out_mm_kwargs_data.get(
-            "video_grid_thw", torch.empty((0, 3), dtype=torch.int64)
-        )
-        num_frames = out_mm_kwargs_data.get(
-            "num_frames", torch.tensor([], dtype=torch.int64)
-        )
-
-        assert len(num_frames) == num_videos, (
-            f"Size of num_frames={len(num_frames)} "
-            f"doesn't equal to number of videos={num_videos}"
-        )
-
-        video_grid_hws = split_thw(video_grid_thw)
-        assert int(num_frames.sum().tolist()) == video_grid_hws.shape[0], (
-            f"The first dimension of `video_grid_hws`={video_grid_hws.shape[0]}"
-            f"doesn't equal to num of frames."
-        )
-
-        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
-
-            """
-            if modality == "image":
-                out_item = out_mm_kwargs[modality][item_idx]
-                grid_thw = out_item[f"{modality}_grid_thw"].data
-                assert isinstance(grid_thw, torch.Tensor)
-
-                num_tokens = int(grid_thw.prod()) // merge_length
-                return [image_token_id] * num_tokens
-            elif modality == "video":
-                placeholders = []
-                video_timestamps = timestamps[item_idx]
-                video_frame_types = frame_types[item_idx]
-                grid_thw = video_grid_hws[
-                    cu_seqlens[item_idx] : cu_seqlens[item_idx + 1]
-                ]
-
-                nframes = grid_thw.shape[0]
-
-                if video_timestamps is None:
-                    video_timestamps = [""] * nframes
-                else:
-                    video_timestamps = [format(ts, ".1f") for ts in video_timestamps]
-
-                if video_frame_types is None:
-                    video_frame_types = [0] * nframes
-                for i, sub_thw in enumerate(grid_thw):
-                    s = f"{hf_processor.frame_token}{video_timestamps[i]}"
-                    if video_frame_types[i] == 1:
-                        s += hf_processor.fast_start
-                    placeholders.extend(tokenizer.encode(s))
-                    num_frame_tokens = int(sub_thw.prod()) // merge_length
-                    placeholders.extend([video_token_id] * num_frame_tokens)
-                    if video_frame_types[i] == 1:
-                        placeholders.append(vocab[hf_processor.fast_end])
-
-                return PromptUpdateDetails.select_token_id(
-                    placeholders, embed_token_id=video_token_id
-                )
-            else:
-                raise ValueError(f"Unsupported modality {modality}")
-
-        return [
-            PromptReplacement(
-                modality=modality,
-                target=[placeholder[modality]],
-                replacement=partial(get_replacement_keye, modality=modality),
-            )
-            for modality in ("image", "video")
-        ]
-
-    def _get_mm_fields_config(
-        self,
-        hf_inputs: BatchFeature,
-        hf_processor_mm_kwargs: Mapping[str, object],
-    ) -> Mapping[str, MultiModalFieldConfig]:
-        return _keye_field_config(hf_inputs)
-
-
-class KeyeVL1_5DummyInputsBuilder(
-    KeyeBaseDummyInputsBuilder[KeyeVL1_5ProcessingInfo]
-): ...
-
-
-@MULTIMODAL_REGISTRY.register_processor(
-    KeyeVL1_5MultiModalProcessor,
-    info=KeyeVL1_5ProcessingInfo,
-    dummy_inputs=KeyeVL1_5DummyInputsBuilder,
-)
-class KeyeVL1_5ForConditionalGeneration(
-    BaseKeyeModule, SupportsMultiModal, SupportsLoRA, SupportsPP, SupportsMRoPE
-):
-    def _build_projector(
-        self,
-        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
-        self.merge_size = config.vision_config.spatial_merge_size
-        super().__init__(vllm_config=vllm_config, prefix=prefix)
-
-    def _parse_and_validate_image_input(
-        self, **kwargs: object
-    ) -> KeyeVL1_5ImageInputs | None:
-        pixel_values = kwargs.pop("pixel_values", None)
-        image_embeds = kwargs.pop("image_embeds", None)
-        image_grid_thw = kwargs.pop("image_grid_thw", None)
-
-        if pixel_values is None and image_embeds is None:
-            return None
-
-        if pixel_values is not None:
-            return KeyeVL1_5ImagePixelInputs(
-                type="pixel_values",
-                pixel_values=pixel_values,
-                image_grid_thw=image_grid_thw,
-            )
-
-        return KeyeVL1_5ImageEmbeddingInputs(
-            type="image_embeds",
-            image_embeds=image_embeds,
-            image_grid_thw=image_grid_thw,
-        )
-
-    def _parse_and_validate_video_input(
-        self, **kwargs: object
-    ) -> KeyeVL1_5VideoInputs | 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)
-        num_frames = kwargs.pop("num_frames", None)
-
-        if pixel_values_videos is None and video_embeds is None:
-            return None
-
-        if pixel_values_videos is not None:
-            return KeyeVL1_5VideoPixelInputs(
-                type="pixel_values_videos",
-                pixel_values_videos=pixel_values_videos,
-                video_grid_thw=video_grid_thw,
-                num_frames=num_frames,
-            )
-
-        return KeyeVL1_5VideoEmbeddingInputs(
-            type="video_embeds",
-            video_embeds=video_embeds,
-            video_grid_thw=video_grid_thw,
-            num_frames=num_frames,
-        )
-
-    def _process_video_input(
-        self, video_input: KeyeVL1_5VideoInputs
-    ) -> tuple[torch.Tensor, ...]:
-        video_type = video_input["type"]
-        video_grid_thw = split_thw(video_input["video_grid_thw"])
-        pixel_values_videos = video_input.get("pixel_values_videos", None)
-
-        video_embeds = self._process_video_embeds(
-            video_type, video_grid_thw, pixel_values_videos
-        )
-        video_embeds = torch.concat(video_embeds, dim=0)
-
-        num_frames = video_input["num_frames"].clone().tolist()
-
-        num_patches = get_num_patches(video_grid_thw, num_frames).tolist()
-
-        patch_cu_seqlens = torch.cumsum(
-            torch.tensor([0] + num_patches).detach().clone(), dim=-1
-        )
-        patch_cu_seqlens = torch.div(
-            patch_cu_seqlens, self.merge_size**2, rounding_mode="floor"
-        )
-
-        new_video_embeds = []
-        for idx in range(patch_cu_seqlens.shape[0] - 1):
-            start = patch_cu_seqlens[idx]
-            end = patch_cu_seqlens[idx + 1]
-            new_video_embeds.append(video_embeds[start:end])
-        return tuple(new_video_embeds)
-
-    def iter_mm_grid_thw(
-        self, mm_features: list[MultiModalFeatureSpec]
-    ) -> Iterator[tuple[int, int, int, int]]:
-        spatial_merge_size = self.config.vision_config.spatial_merge_size
-
-        for mm_feature in sorted(mm_features, key=lambda f: f.mm_position.offset):
-            if mm_feature.data is None:
-                raise ValueError("M-RoPE calculation requires multimodal feature data")
-
-            embed_ranges = mm_feature.mm_position.extract_embeds_range()
-            if mm_feature.modality == "image":
-                assert len(embed_ranges) == 1
-                grid_thw = mm_feature.data["image_grid_thw"].data
-                if isinstance(grid_thw, torch.Tensor):
-                    if grid_thw.ndim == 2:
-                        assert grid_thw.shape[0] == 1
-                        t, h, w = grid_thw[0].tolist()
-                    else:
-                        t, h, w = grid_thw.tolist()
-                else:
-                    if isinstance(grid_thw[0], list):
-                        assert len(grid_thw) == 1
-                        t, h, w = grid_thw[0]
-                    else:
-                        t, h, w = grid_thw
-
-                yield (
-                    embed_ranges[0][0],
-                    t,
-                    h // spatial_merge_size,
-                    w // spatial_merge_size,
-                )
-            elif mm_feature.modality == "video":
-                split_video_grids = split_thw(mm_feature.data["video_grid_thw"].data)
-                assert len(embed_ranges) == split_video_grids.shape[0]
-                for (start_idx, end_idx), (t, h, w) in zip(
-                    embed_ranges, split_video_grids.tolist()
-                ):
-                    llm_grid_h = h // spatial_merge_size
-                    llm_grid_w = w // spatial_merge_size
-                    num_mm_tokens = t * llm_grid_h * llm_grid_w
-                    assert end_idx - start_idx + 1 == num_mm_tokens
-                    yield (start_idx, t, llm_grid_h, llm_grid_w)
-            else:
-                raise ValueError(f"Unsupported modality: {mm_feature.modality}")
-
-    def get_mrope_input_positions(
-        self,
-        input_tokens: list[int],
-        mm_features: list[MultiModalFeatureSpec],
-    ) -> tuple[torch.Tensor, int]:
-        llm_pos_ids_list: list = []
-        st = 0
-
-        for (
-            offset,
-            llm_grid_t,
-            llm_grid_h,
-            llm_grid_w,
-        ) in self.iter_mm_grid_thw(mm_features):
-            text_len = offset - st
-
-            st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
-            llm_pos_ids_list.append(
-                torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx
-            )
-
-            t_index = (
-                (
-                    torch.arange(llm_grid_t)
-                    .view(-1, 1)
-                    .expand(-1, llm_grid_h * llm_grid_w)
-                )
-                .long()
-                .flatten()
-            )
-
-            h_index = (
-                torch.arange(llm_grid_h)
-                .view(1, -1, 1)
-                .expand(llm_grid_t, -1, llm_grid_w)
-                .flatten()
-            )
-            w_index = (
-                torch.arange(llm_grid_w)
-                .view(1, 1, -1)
-                .expand(llm_grid_t, llm_grid_h, -1)
-                .flatten()
-            )
-            llm_pos_ids_list.append(
-                torch.stack([t_index, h_index, w_index]) + text_len + st_idx
-            )
-            st = offset + llm_grid_t * llm_grid_h * llm_grid_w
-
-        if st < len(input_tokens):
-            st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
-            text_len = len(input_tokens) - st
-            llm_pos_ids_list.append(
-                torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx
-            )
-
-        llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)
-        mrope_position_delta = (llm_positions.max() + 1 - len(input_tokens)).item()
-
-        return llm_positions, mrope_position_delta
diff --git a/vllm/model_executor/models/kimi_k25_vit.py b/vllm/model_executor/models/kimi_k25_vit.py
index bb1be1ba32f8..902e5eb346cf 100644
--- a/vllm/model_executor/models/kimi_k25_vit.py
+++ b/vllm/model_executor/models/kimi_k25_vit.py
@@ -40,7 +40,6 @@
     is_vit_use_data_parallel,
     run_dp_sharded_mrope_vision_model,
 )
-from vllm.platforms import current_platform
 from vllm.transformers_utils.configs.kimi_k3 import KimiK3VisionConfig
 from vllm.triton_utils import HAS_TRITON
 from vllm.utils.torch_utils import async_tensor_h2d
@@ -63,25 +62,6 @@ def _apply_rope_input_validation(x, freqs_cis):
     assert freqs_cis.dtype == torch.complex64, freqs_cis.dtype
 
 
-def get_rope_shape_decorate(func):
-    _get_rope_shape_first_call_flag = set()
-
-    def wrapper(org, interpolation_mode, shape):
-        key = (org.requires_grad, torch.is_grad_enabled(), interpolation_mode)
-        if key not in _get_rope_shape_first_call_flag:
-            _get_rope_shape_first_call_flag.add(key)
-            _ = func(org, interpolation_mode, shape=(64, 64))
-        return func(org, interpolation_mode, shape)
-
-    return wrapper
-
-
-@get_rope_shape_decorate
-@torch.compile(
-    dynamic=True,
-    backend=current_platform.simple_compile_backend,
-    disable=current_platform.simple_compile_backend == "tpu",
-)
 def get_rope_shape(org, interpolation_mode, shape):
     return (
         F.interpolate(
diff --git a/vllm/model_executor/models/laguna_dflash.py b/vllm/model_executor/models/laguna_dflash.py
index 02af55c00033..14d6078ad6c1 100644
--- a/vllm/model_executor/models/laguna_dflash.py
+++ b/vllm/model_executor/models/laguna_dflash.py
@@ -29,6 +29,7 @@
 )
 from vllm.model_executor.models.interfaces import EagleModelMixin, SupportsEagle3
 from vllm.multimodal.inputs import NestedTensors
+from vllm.v1.worker.gpu.spec_decode.dflash.config import DFlashSpeculatorConfig
 
 from .laguna import LagunaDecoderLayer
 from .qwen3_dflash import DFlashQwen3Model
@@ -236,6 +237,9 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
 
 
 class DFlashLagunaForCausalLM(nn.Module, SupportsEagle3):
+    # The speculator this is a draft model of (see SpeculativeDraftModel).
+    speculator_config_class = DFlashSpeculatorConfig
+
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         nn.Module.__init__(self)
         speculative_config = vllm_config.speculative_config
diff --git a/vllm/model_executor/models/lilicorr.py b/vllm/model_executor/models/lilicorr.py
index 9901a68615f6..8615fc4a12cb 100644
--- a/vllm/model_executor/models/lilicorr.py
+++ b/vllm/model_executor/models/lilicorr.py
@@ -23,6 +23,7 @@
 from vllm.model_executor.layers.quantization.modelopt import ModelOptLinearMethod
 from vllm.model_executor.layers.quantization.utils.quant_utils import kNvfp4Static
 from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
+from vllm.v1.worker.gpu.spec_decode.lilicorr.config import LiLiCorrSpeculatorConfig
 
 from .nemotron_h import NemotronHMLP
 from .qwen3_dflash import DFlashQwen3ForCausalLM, DFlashQwen3Model
@@ -463,6 +464,9 @@ def __init__(
 
 
 class LiLiCorrForCausalLM(DFlashQwen3ForCausalLM):
+    # The speculator this is a draft model of (see SpeculativeDraftModel).
+    speculator_config_class = LiLiCorrSpeculatorConfig
+
     model_cls = LiLiCorr
 
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
diff --git a/vllm/model_executor/models/mimo_v2.py b/vllm/model_executor/models/mimo_v2.py
index d35e4faf22f9..28f89fb99e8e 100644
--- a/vllm/model_executor/models/mimo_v2.py
+++ b/vllm/model_executor/models/mimo_v2.py
@@ -317,15 +317,24 @@ def __init__(
 
         sliding_window = sliding_window_size if sliding_window_size > -1 else None
 
-        # Use DiffKV backend when V has a different head dim than K.
-        # Auto-pick FA-DiffKV when FA3/4 is usable on this device, else fall
-        # back to TRITON_ATTN_DIFFKV.  Users can force a choice via
-        # `--attention-backend <FLASH_ATTN_DIFFKV|TRITON_ATTN_DIFFKV>`.
+        # V has a different head dim than K. Honor an explicit backend that
+        # declares it can serve that geometry; anything else keeps the auto-pick
+        # below, with a warning.
         if self.v_head_dim != self.head_dim:
             requested = get_current_vllm_config().attention_config.backend
-            if requested is not None and requested.name.endswith("_DIFFKV"):
+            if requested is not None and requested.get_class().supports_head_size_v(
+                self.head_dim, self.v_head_dim
+            ):
                 backend_enum = requested
             else:
+                if requested is not None:
+                    logger.warning_once(
+                        "%s cannot serve head_size_v=%d != head_size=%d; "
+                        "using the DiffKV auto-pick instead.",
+                        requested.name,
+                        self.v_head_dim,
+                        self.head_dim,
+                    )
                 fa_backend = AttentionBackendEnum.FLASH_ATTN_DIFFKV.get_class()
                 assert hasattr(fa_backend, "is_supported_on_current_device")
                 if fa_backend.is_supported_on_current_device(
@@ -337,8 +346,11 @@ def __init__(
                 else:
                     backend_enum = AttentionBackendEnum.TRITON_ATTN_DIFFKV
             attn_backend = backend_enum.get_class()
-            assert hasattr(attn_backend, "set_head_size_v")
-            attn_backend.set_head_size_v(self.v_head_dim)
+            if backend_enum.name.endswith("_DIFFKV"):
+                # DiffKV takes V's head size as a backend class attribute;
+                # unified attention reads it from the layer's KV spec.
+                assert hasattr(attn_backend, "set_head_size_v")
+                attn_backend.set_head_size_v(self.v_head_dim)
             logger.info_once("Using %s for attention.", attn_backend.get_name())
         else:
             attn_backend = None
diff --git a/vllm/model_executor/models/mimo_v2_omni.py b/vllm/model_executor/models/mimo_v2_omni.py
index 390af45739f8..01bb1781814c 100644
--- a/vllm/model_executor/models/mimo_v2_omni.py
+++ b/vllm/model_executor/models/mimo_v2_omni.py
@@ -162,7 +162,12 @@ def __init__(
             if use_data_parallel
             else parallel_state.get_tensor_model_parallel_world_size()
         )
-        self.tp_rank = parallel_state.get_tensor_model_parallel_rank()
+        # Data parallel runs the whole encoder on every rank, so the rank has to
+        # collapse with the size: otherwise the per-head slices below start past
+        # the end of the head range on ranks > 0.
+        self.tp_rank = (
+            0 if use_data_parallel else parallel_state.get_tensor_model_parallel_rank()
+        )
 
         self.num_heads = num_heads
         self.num_kv_heads = num_kv_heads
@@ -1229,6 +1234,8 @@ def get_dummy_mm_data(
 class MiMoV2OmniForCausalLM(
     nn.Module, SupportsMultiModal, SupportsPP, SupportsQuant, SupportsEagle3
 ):
+    supports_encoder_tp_data = True
+
     # To ensure correct weight loading and mapping.
     hf_to_vllm_mapper = WeightsMapper(
         orig_to_new_prefix={
diff --git a/vllm/model_executor/models/ming_audio.py b/vllm/model_executor/models/ming_audio.py
new file mode 100644
index 000000000000..ab29c0a65c2f
--- /dev/null
+++ b/vllm/model_executor/models/ming_audio.py
@@ -0,0 +1,151 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+# Adapted from
+# https://huggingface.co/inclusionAI/Ming-3.0-asr-tiny/blob/ac95e7b1678c08ea6f4a2581a051e432ed63e9b7/modeling_ming_v3_asr_encoder.py
+# https://huggingface.co/inclusionAI/Ming-3.0-asr-tiny/blob/ac95e7b1678c08ea6f4a2581a051e432ed63e9b7/modeling_whisper_encoder.py
+# Copyright 2026 Ant Group. All rights reserved.
+#
+# 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.
+"""Reusable Ming audio tower and checkpoint-compatible projector."""
+
+from collections.abc import Iterable
+
+import torch
+from torch import nn
+from transformers import WhisperConfig
+
+from vllm.config import VllmConfig
+from vllm.transformers_utils.configs.ming_v3_asr import MingAudioConfig
+
+from .utils import AutoWeightsLoader, WeightsMapper
+from .whisper import WhisperEncoder, _create_fake_bias_for_k_proj
+
+
+def _get_audio_projector_params(config: MingAudioConfig) -> list[tuple[int, int, int]]:
+    kernels = config.downsampler_kernel_size_list
+    strides = config.downsampler_stride_list
+    paddings = config.downsampler_padding_list
+    if not kernels or not (len(kernels) == len(strides) == len(paddings)):
+        raise ValueError("Audio convolution lists must have equal positive length")
+    params = list(zip(kernels, strides, paddings))
+    if any(
+        kernel <= 0 or stride <= 0 or padding < 0 for kernel, stride, padding in params
+    ):
+        raise ValueError(
+            "Audio kernels and strides must be positive and padding nonnegative"
+        )
+    return params
+
+
+def get_audio_token_count(
+    config: MingAudioConfig, num_frames: int, chunk_frames: int
+) -> int:
+    """Count projected frames with independent encoder positions for each chunk."""
+    if num_frames <= 0:
+        raise ValueError("Audio must contain at least one mel frame")
+    params = _get_audio_projector_params(config)
+
+    def output_length(length: int) -> int:
+        length = (length + 1) // 2
+        for kernel, stride, padding in params:
+            length = (length + 2 * padding - kernel) // stride + 1
+        if length <= 0:
+            raise ValueError("Audio chunk is too short for its projector")
+        return length
+
+    num_chunks, tail = divmod(num_frames, chunk_frames)
+    return (num_chunks * output_length(chunk_frames) if num_chunks else 0) + (
+        output_length(tail) if tail else 0
+    )
+
+
+class MingAudioEncoder(nn.Module):
+    hf_to_vllm_mapper = WeightsMapper(
+        orig_to_new_substr={
+            "blocks.": "layers.",
+            ".attn.query.": ".self_attn.q_proj.",
+            ".attn.key.": ".self_attn.k_proj.",
+            ".attn.value.": ".self_attn.v_proj.",
+            ".attn.out.": ".self_attn.out_proj.",
+            ".attn_ln.": ".self_attn_layer_norm.",
+            ".mlp_ln.": ".final_layer_norm.",
+            ".mlp.0.": ".mlp.fc1.",
+            ".mlp.2.": ".mlp.fc2.",
+            "ln_post.": "layer_norm.",
+        },
+        orig_to_new_suffix={"positional_embedding": "embed_positions.weight"},
+        orig_to_new_stacked={
+            ".self_attn.q_proj": (".self_attn.qkv_proj", "q"),
+            ".self_attn.k_proj": (".self_attn.qkv_proj", "k"),
+            ".self_attn.v_proj": (".self_attn.qkv_proj", "v"),
+        },
+    )
+
+    def __init__(
+        self, *, vllm_config: VllmConfig, config: MingAudioConfig, prefix: str
+    ) -> None:
+        super().__init__()
+        encoder_config = WhisperConfig(
+            num_mel_bins=config.n_mels,
+            d_model=config.n_audio_state,
+            encoder_attention_heads=config.n_audio_head,
+            encoder_layers=config.n_audio_layer,
+            encoder_ffn_dim=4 * config.n_audio_state,
+            max_source_positions=config.n_audio_ctx,
+            activation_function="gelu",
+            pos_embed="learned",
+        )
+        self.encoder = WhisperEncoder(
+            vllm_config=vllm_config.with_hf_config(encoder_config),
+            prefix=prefix,
+            enable_pp=False,
+            init_in_fp32=True,
+        )
+
+    def forward(self, features: torch.Tensor) -> torch.Tensor:
+        return self.encoder([features.transpose(0, 1)])[0]
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        weights = _create_fake_bias_for_k_proj(
+            weights, ".attn.key.weight", out_features=self.encoder.conv1.out_channels
+        )
+        loaded = AutoWeightsLoader(self.encoder).load_weights(
+            weights, mapper=self.hf_to_vllm_mapper
+        )
+        return {f"encoder.{name}" for name in loaded}
+
+
+class Transpose(nn.Module):
+    def forward(self, value: torch.Tensor) -> torch.Tensor:
+        return value.transpose(-1, -2)
+
+
+class MingAudioProjector(nn.Module):
+    def __init__(self, config: MingAudioConfig, hidden_size: int) -> None:
+        super().__init__()
+        layers: list[nn.Module] = []
+        width = config.n_audio_state
+        for kernel, stride, padding in _get_audio_projector_params(config):
+            layers.extend(
+                [
+                    nn.Conv1d(width, hidden_size, kernel, stride, padding),
+                    nn.GELU(),
+                ]
+            )
+            width = hidden_size
+        layers.extend([Transpose(), nn.Linear(hidden_size, hidden_size)])
+        self.proj = nn.Sequential(*layers)
+
+    def forward(self, embeddings: torch.Tensor) -> torch.Tensor:
+        return self.proj(embeddings.transpose(-1, -2))
diff --git a/vllm/model_executor/models/ming_v3_asr.py b/vllm/model_executor/models/ming_v3_asr.py
new file mode 100644
index 000000000000..384e3c2bb06b
--- /dev/null
+++ b/vllm/model_executor/models/ming_v3_asr.py
@@ -0,0 +1,497 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+# Adapted from
+# https://huggingface.co/inclusionAI/Ming-3.0-asr-tiny/blob/ac95e7b1678c08ea6f4a2581a051e432ed63e9b7/modeling_ming_v3_asr.py
+# https://huggingface.co/inclusionAI/Ming-3.0-asr-tiny/blob/ac95e7b1678c08ea6f4a2581a051e432ed63e9b7/processing_ming_v3_asr.py
+# https://github.com/inclusionAI/Ming-3.0-asr/blob/6141e8c552815db724d8df0f72c9d865c547a1ea/examples/hf_runtime.py
+# Copyright 2025 Antgroup and The HuggingFace Inc. team. All rights reserved.
+#
+# The HF model/processor sources retain Apache-2.0 headers.
+# ASR helpers from hf_runtime.py are covered by the MIT notice below.
+#
+# 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.
+#
+# MIT License
+#
+# Copyright (c) 2026 inclusionAI
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+"""Inference-only Ming V3 ASR with the native Bailing hybrid decoder."""
+
+from collections.abc import Iterable, Mapping, Sequence
+from functools import cached_property
+from numbers import Integral
+from typing import cast
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from transformers import BatchFeature
+
+from vllm.config import ModelConfig, SpeechToTextConfig, VllmConfig
+from vllm.config.multimodal import MultiModalDummyOptions
+from vllm.config.speech_to_text import SpeechToTextParams
+from vllm.inputs import MultiModalDataDict, TokensPrompt
+from vllm.multimodal import MULTIMODAL_REGISTRY
+from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargsItems
+from vllm.multimodal.parse import MultiModalDataItems, MultiModalDataParser
+from vllm.multimodal.processing import (
+    BaseDummyInputsBuilder,
+    BaseMultiModalProcessor,
+    BaseProcessingInfo,
+    PromptReplacement,
+)
+from vllm.sequence import IntermediateTensors
+from vllm.tokenizers import cached_tokenizer_from_config
+from vllm.transformers_utils.configs.ming_v3_asr import (
+    MingAudioConfig,
+    MingV3ASRConfig,
+)
+from vllm.transformers_utils.processors.ming_audio import MingASRFeatureExtractor
+
+from .bailing_moe_v3 import BailingMoeV3ForCausalLM
+from .interfaces import (
+    HasInnerState,
+    IsHybrid,
+    MultiModalEmbeddings,
+    StreamingTranscriptionPostProcessor,
+    SupportsMultiModal,
+    SupportsPP,
+    SupportsTranscription,
+)
+from .ming_audio import MingAudioEncoder, MingAudioProjector, get_audio_token_count
+from .module_mapping import MultiModelKeys
+from .utils import AutoWeightsLoader, WeightsMapper, maybe_prefix
+
+ASR_INSTRUCTION = (
+    "Please recognize the language of this speech and transcribe it, "
+    "placing a `\t` between the identified language and the transcribe."
+)
+_LANGUAGES = {"zh": "Chinese", "en": "English"}
+_OUTPUT_LANGUAGES = ("Chinese", "English", "Canton", "Japanese", "Korean")
+_AUDIO_FORMS = (
+    ("<audio>", "<audioPatch>", "</audio>"),
+    ("<|audio_start|>", "<|audio_pad|>", "<|audio_end|>"),
+)
+
+
+def _audio_tokens(tokenizer) -> tuple[str, str, str]:
+    vocab = tokenizer.get_vocab()
+    return next(form for form in _AUDIO_FORMS if all(t in vocab for t in form))
+
+
+def _get_audio_chunk_frames(
+    config: MingAudioConfig, options: Mapping[str, object], samples: int
+) -> int:
+    chunk_frames = options.get("audio_chunk_frames")
+    context_chunks = options.get("context_chunks", False)
+    if context_chunks not in (False, True):
+        raise ValueError("context_chunks must be a boolean or 0/1")
+    if context_chunks:
+        if chunk_frames is not None:
+            raise ValueError("context_chunks and audio_chunk_frames are exclusive")
+        if samples > 60 * 16000:
+            raise ValueError("Context ASR chunk policy supports at most 60 seconds")
+        chunk_frames = 6000 if samples > 30 * 16000 else 3000
+    if chunk_frames is None:
+        chunk_frames = config.chunk_frames
+    if (
+        isinstance(chunk_frames, bool)
+        or not isinstance(chunk_frames, Integral)
+        or not 0 < int(chunk_frames) <= 2 * config.n_audio_ctx
+    ):
+        raise ValueError(
+            "audio_chunk_frames must be a positive integer no larger than "
+            f"{2 * config.n_audio_ctx}"
+        )
+    return int(chunk_frames)
+
+
+def _transcript(text: str) -> str:
+    text = text.strip()
+    if "\t" in text:
+        return text.split("\t", 1)[1].strip()
+    for language in _OUTPUT_LANGUAGES:
+        if text.startswith(language):
+            return text[len(language) :].strip()
+    return text
+
+
+class MingASRStreamingPostProcessor(StreamingTranscriptionPostProcessor):
+    def __init__(self) -> None:
+        self.raw_text = ""
+        self.emitted_text = ""
+
+    def process_delta(self, text_delta: str, finished: bool) -> str:
+        self.raw_text += text_delta
+        raw = self.raw_text.strip()
+        if not finished and any(
+            language.startswith(raw) for language in (*_OUTPUT_LANGUAGES, "Cantonese")
+        ):
+            return ""
+        text = _transcript(self.raw_text)
+        delta = text[len(self.emitted_text) :]
+        self.emitted_text = text
+        return delta
+
+
+class MingASRProcessingInfo(BaseProcessingInfo):
+    def get_hf_config(self) -> MingV3ASRConfig:
+        return self.ctx.get_hf_config(MingV3ASRConfig)
+
+    @cached_property
+    def _feature_extractor(self) -> MingASRFeatureExtractor:
+        return MingASRFeatureExtractor(self.get_hf_config().audio_config.n_mels)
+
+    def get_feature_extractor(self, **kwargs: object) -> MingASRFeatureExtractor:
+        return self._feature_extractor
+
+    def get_supported_mm_processor_kwargs(self) -> dict[str, set[str]]:
+        return {"audio_kwargs": {"audio_chunk_frames", "context_chunks"}}
+
+    def get_audio_options(self, kwargs: Mapping[str, object]) -> Mapping[str, object]:
+        resolved_kwargs = self._merge_and_resolve_mm_processor_kwargs(kwargs)
+        audio_options = resolved_kwargs.get("audio_kwargs", {})
+        if not isinstance(audio_options, Mapping):
+            raise ValueError("audio_kwargs must be a mapping")
+
+        return audio_options
+
+    def get_data_parser(self) -> MultiModalDataParser:
+        return MultiModalDataParser(target_sr=16000, target_channels=1)
+
+    def get_supported_mm_limits(self) -> Mapping[str, int | None]:
+        return {"audio": 1}
+
+    def get_mm_max_tokens_per_item(
+        self, seq_len: int, mm_counts: Mapping[str, int]
+    ) -> Mapping[str, int]:
+        tokenizer = self.get_tokenizer()
+        start, _, end = _audio_tokens(tokenizer)
+        overhead = len(tokenizer.encode(ASR_INSTRUCTION + start + end))
+        return {"audio": max(1, seq_len - overhead)}
+
+
+class MingASRDummyInputsBuilder(BaseDummyInputsBuilder[MingASRProcessingInfo]):
+    def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
+        return "".join(_audio_tokens(self.info.get_tokenizer())) * mm_counts.get(
+            "audio", 0
+        )
+
+    def get_dummy_mm_data(
+        self,
+        seq_len: int,
+        mm_counts: Mapping[str, int],
+        mm_options: MultiModalDummyOptions,
+    ) -> MultiModalDataDict:
+        audio_options = self.info.get_audio_options({})
+        config = self.info.get_hf_config().audio_config
+        chunk_frames = _get_audio_chunk_frames(config, audio_options, 60 * 16000)
+        budget = self.info.get_mm_max_tokens_per_item(seq_len, mm_counts)["audio"]
+        low, high = 1, budget * 4 + chunk_frames
+        if audio_options.get("context_chunks"):
+            high = min(high, 6000)
+        while low < high:
+            middle = (low + high + 1) // 2
+            if get_audio_token_count(config, middle, chunk_frames) <= budget:
+                low = middle
+            else:
+                high = middle - 1
+        return {
+            "audio": self._get_dummy_audios(
+                length=max(201, low * 160),
+                num_audios=mm_counts.get("audio", 0),
+                overrides=mm_options.get("audio"),
+            )
+        }
+
+
+class MingASRMultiModalProcessor(BaseMultiModalProcessor[MingASRProcessingInfo]):
+    def _call_hf_processor(
+        self, hf_data: Mapping[str, object], hf_kwargs: Mapping[str, object]
+    ) -> BatchFeature:
+        config = self.info.get_hf_config().audio_config
+        feature_extractor = self.info.get_feature_extractor()
+        audio_options = self.info.get_audio_options(hf_kwargs)
+        features, lengths, chunk_frames = [], [], []
+        audios = hf_data.get("audio", [])
+        assert isinstance(audios, list)
+        for audio in audios:
+            frames = _get_audio_chunk_frames(config, audio_options, len(audio))
+            waveform = torch.as_tensor(audio)
+            if waveform.ndim != 1 or waveform.numel() <= 200:
+                raise ValueError(
+                    "Ming audio requires more than 200 mono samples at 16 kHz"
+                )
+            if waveform.is_complex() or not torch.isfinite(waveform).all():
+                raise ValueError("Audio must contain finite real samples")
+            inputs = feature_extractor(
+                waveform.detach().to(device="cpu", dtype=torch.float32).numpy(),
+                sampling_rate=feature_extractor.sampling_rate,
+                padding=False,
+                truncation=False,
+                return_tensors="pt",
+            )
+            feature = inputs.input_features[0].T.to(torch.bfloat16)
+            features.append(feature)
+            lengths.append(len(feature))
+            chunk_frames.append(frames)
+        return BatchFeature(
+            {
+                "audio_feats": features,
+                "audio_feats_lengths": torch.tensor(lengths),
+                "audio_chunk_frames": torch.tensor(chunk_frames),
+            }
+        )
+
+    def _get_mm_fields_config(
+        self,
+        hf_inputs: BatchFeature,
+        hf_processor_mm_kwargs: Mapping[str, object],
+    ) -> Mapping[str, MultiModalFieldConfig]:
+        return {
+            name: MultiModalFieldConfig.batched(
+                "audio", keep_on_cpu=name != "audio_feats"
+            )
+            for name in ("audio_feats", "audio_feats_lengths", "audio_chunk_frames")
+        }
+
+    def _get_prompt_updates(
+        self,
+        mm_items: MultiModalDataItems,
+        hf_processor_mm_kwargs: Mapping[str, object],
+        out_mm_kwargs: MultiModalKwargsItems,
+    ) -> Sequence[PromptReplacement]:
+        config = self.info.get_hf_config()
+        data = out_mm_kwargs.get_data()
+        lengths = torch.as_tensor(data.get("audio_feats_lengths", [])).flatten()
+        frames = torch.as_tensor(data.get("audio_chunk_frames", [])).flatten()
+        counts = [
+            get_audio_token_count(config.audio_config, int(length), int(chunk))
+            for length, chunk in zip(lengths, frames, strict=True)
+        ]
+        return [
+            PromptReplacement(
+                modality="audio",
+                target=[config.audio_token_id],
+                replacement=lambda index: [config.audio_token_id] * counts[index],
+            )
+        ]
+
+
+@MULTIMODAL_REGISTRY.register_processor(
+    MingASRMultiModalProcessor,
+    info=MingASRProcessingInfo,
+    dummy_inputs=MingASRDummyInputsBuilder,
+)
+class MingV3ASRForConditionalGeneration(
+    nn.Module,
+    SupportsMultiModal,
+    SupportsTranscription,
+    HasInnerState,
+    IsHybrid,
+    SupportsPP,
+):
+    supported_languages = _LANGUAGES
+    supports_transcription_only = True
+    hf_to_vllm_mapper = WeightsMapper(
+        orig_to_new_prefix={
+            "model.audio.": "audio_encoder.",
+            "linear_proj_audio.": "audio_projector.proj.",
+            "model.visual.": None,
+            "linear_proj.": None,
+            "model.": "language_model.model.",
+            "lm_head.": "language_model.lm_head.",
+        }
+    )
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
+        super().__init__()
+        self.config = vllm_config.model_config.hf_config
+        self.dtype = vllm_config.model_config.dtype
+        with self._mark_tower_model(vllm_config, "audio"):
+            self.audio_encoder = MingAudioEncoder(
+                vllm_config=vllm_config,
+                config=self.config.audio_config,
+                prefix=maybe_prefix(prefix, "audio_encoder"),
+            )
+            self.audio_projector = MingAudioProjector(
+                self.config.audio_config, self.config.text_config.hidden_size
+            )
+        with self._mark_language_model(vllm_config):
+            self.language_model = BailingMoeV3ForCausalLM(
+                vllm_config=vllm_config.with_hf_config(self.config.text_config),
+                prefix=maybe_prefix(prefix, "language_model"),
+            )
+        self.make_empty_intermediate_tensors = (
+            self.language_model.make_empty_intermediate_tensors
+        )
+
+    @classmethod
+    def get_placeholder_str(cls, modality: str, i: int) -> str:
+        if not modality.startswith("audio"):
+            raise ValueError("Ming V3 ASR only supports audio")
+        return "<audio><audioPatch></audio>"
+
+    def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
+        features = cast(
+            torch.Tensor | list[torch.Tensor] | None, kwargs.get("audio_feats")
+        )
+        if features is None:
+            return ()
+        lengths = torch.as_tensor(kwargs["audio_feats_lengths"]).flatten().tolist()
+        chunk_frames = torch.as_tensor(kwargs["audio_chunk_frames"]).flatten().tolist()
+        outputs = []
+        for feature, length, frames in zip(
+            features, lengths, chunk_frames, strict=True
+        ):
+            chunks = []
+            feature = feature[:length].to(dtype=self.dtype)
+            for chunk in feature.split(frames):
+                encoded = self.audio_encoder(chunk)
+                projected = self.audio_projector(encoded.unsqueeze(0))[0]
+                chunks.append(projected)
+            outputs.append(F.normalize(torch.cat(chunks), dim=-1))
+        return outputs
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        **kwargs: object,
+    ) -> torch.Tensor | IntermediateTensors:
+        return self.language_model(
+            input_ids, positions, intermediate_tensors, inputs_embeds
+        )
+
+    def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.language_model.compute_logits(hidden_states)
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        return AutoWeightsLoader(self).load_weights(
+            weights, mapper=self.hf_to_vllm_mapper
+        )
+
+    def get_mm_mapping(self) -> MultiModelKeys:
+        return MultiModelKeys.from_string_field(
+            language_model="language_model",
+            tower_model=["audio_encoder.", "audio_projector."],
+        )
+
+    @classmethod
+    def get_mamba_state_shape_from_config(cls, vllm_config: VllmConfig):
+        return BailingMoeV3ForCausalLM.get_mamba_state_shape_from_config(
+            vllm_config.with_hf_config(vllm_config.model_config.hf_config.text_config)
+        )
+
+    @classmethod
+    def get_mamba_state_dtype_from_config(cls, vllm_config: VllmConfig):
+        return BailingMoeV3ForCausalLM.get_mamba_state_dtype_from_config(
+            vllm_config.with_hf_config(vllm_config.model_config.hf_config.text_config)
+        )
+
+    @classmethod
+    def get_mamba_state_copy_func(cls):
+        return BailingMoeV3ForCausalLM.get_mamba_state_copy_func()
+
+    @classmethod
+    def validate_language(cls, language: str | None) -> str | None:
+        if language is not None and language not in cls.supported_languages:
+            raise ValueError("Ming V3 ASR language hint must be zh or en")
+        return language
+
+    @classmethod
+    def get_speech_to_text_config(
+        cls, model_config: ModelConfig, task_type: str
+    ) -> SpeechToTextConfig:
+        return SpeechToTextConfig(
+            sample_rate=16000,
+            audio_resample_method="torchaudio",
+            max_audio_clip_s=None,
+            min_energy_split_window_size=None,
+        )
+
+    @classmethod
+    def get_generation_prompt(cls, stt_params: SpeechToTextParams) -> TokensPrompt:
+        if stt_params.task_type != "transcribe" or stt_params.to_language is not None:
+            raise ValueError("Ming V3 ASR does not support translation")
+        language = cls.validate_language(stt_params.language)
+        prompt = stt_params.request_prompt
+        if not prompt.startswith(ASR_INSTRUCTION):
+            prompt = ASR_INSTRUCTION + prompt
+        reserved = ("\n", "\r", "<AUDIO>") + tuple(
+            token for form in _AUDIO_FORMS for token in form
+        )
+        for text in (prompt, stt_params.hotwords or ""):
+            if any(token in text for token in reserved):
+                raise ValueError(
+                    "ASR context and hotwords must be contiguous text "
+                    "without audio placeholders"
+                )
+        if stt_params.hotwords:
+            prompt += (
+                "This audio may contains the following words or phrases:"
+                + stt_params.hotwords
+            )
+        tokenizer = cached_tokenizer_from_config(stt_params.model_config)
+        start, audio_token, end = _audio_tokens(tokenizer)
+        ids = tokenizer.encode(
+            prompt + start + audio_token + end, add_special_tokens=False
+        )
+        if language is not None:
+            ids += tokenizer.encode(
+                _LANGUAGES[language] + "\t", add_special_tokens=False
+            )
+        options = stt_params.extra_args or {}
+        processor_options = {
+            key: options[key]
+            for key in ("audio_chunk_frames", "context_chunks")
+            if key in options
+        }
+        _get_audio_chunk_frames(
+            stt_params.model_config.hf_config.audio_config,
+            processor_options,
+            len(stt_params.audio),
+        )
+        return TokensPrompt(
+            prompt_token_ids=ids,
+            multi_modal_data={"audio": stt_params.audio},
+            mm_processor_kwargs=processor_options,
+        )
+
+    @classmethod
+    def post_process_output(cls, text: str) -> str:
+        return _transcript(text)
+
+    @classmethod
+    def get_streaming_post_processor_cls(cls):
+        return MingASRStreamingPostProcessor
diff --git a/vllm/model_executor/models/ming_v3_asr_logits_processor.py b/vllm/model_executor/models/ming_v3_asr_logits_processor.py
new file mode 100644
index 000000000000..5d81515ccf34
--- /dev/null
+++ b/vllm/model_executor/models/ming_v3_asr_logits_processor.py
@@ -0,0 +1,76 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Opt-in full-history n-gram blocking for Ming ASR on Model Runner V2."""
+
+from numbers import Integral
+
+from vllm.v1.worker.gpu.sample.logits_processor.interface import LogitsProcessor
+
+
+class NoRepeatNGramLogitsProcessor(LogitsProcessor):
+    @classmethod
+    def validate_params(cls, sampling_params):
+        size = (sampling_params.extra_args or {}).get("no_repeat_ngram_size", 0)
+        if (
+            isinstance(size, bool)
+            or not isinstance(size, Integral)
+            or not 0 <= int(size) < 2**31
+        ):
+            raise ValueError("no_repeat_ngram_size must be an integer in [0, 2**31)")
+
+    def __init__(self, vllm_config, req_states):
+        import torch
+
+        from vllm.v1.worker.gpu.buffer_utils import UvaBackedTensor
+
+        if vllm_config.speculative_config is not None:
+            raise ValueError(
+                "NoRepeatNGramLogitsProcessor does not support speculation"
+            )
+        self.req_states = req_states
+        self.sizes = UvaBackedTensor(req_states.max_num_reqs, dtype=torch.int32)
+
+    def add_request(self, req_idx, sampling_params):
+        self.validate_params(sampling_params)
+        self.sizes.np[req_idx] = int(
+            (sampling_params.extra_args or {}).get("no_repeat_ngram_size", 0)
+        )
+        return bool(self.sizes.np[req_idx] > 0)
+
+    def apply_staged_writes(self):
+        self.sizes.copy_to_uva()
+
+    def apply(self, logits, ctx):
+        import torch
+
+        sizes = set(self.sizes.np[ctx.idx_mapping_np].tolist())
+        if sizes == {0}:
+            return logits
+        slots = ctx.expanded_idx_mapping.long()
+        history = self.req_states.all_token_ids.gpu.index_select(0, slots)
+        lengths = ctx.pos + 1
+        for size in sizes:
+            if size == 0 or size > history.shape[1]:
+                continue
+            windows = history.unfold(1, size, 1)
+            valid = (
+                torch.arange(windows.shape[1], device=logits.device)[None, :]
+                <= lengths[:, None] - size
+            )
+            valid &= self.sizes.gpu[slots, None] == size
+            if size > 1:
+                positions = (
+                    lengths[:, None]
+                    - size
+                    + 1
+                    + torch.arange(size - 1, device=logits.device)[None, :]
+                )
+                prefix = history.gather(1, positions.clamp(min=0).long())
+                valid &= (windows[:, :, :-1] == prefix[:, None, :]).all(dim=2)
+            token_ids = windows[:, :, -1].long()
+            # Input embeddings may have a larger vocabulary than output logits.
+            valid &= (token_ids >= 0) & (token_ids < logits.shape[1])
+            penalties = torch.where(valid, -float("inf"), 0.0)
+            logits.scatter_add_(1, token_ids.clamp(0, logits.shape[1] - 1), penalties)
+
+        return logits
diff --git a/vllm/model_executor/models/minicpm_ocr.py b/vllm/model_executor/models/minicpm_ocr.py
index 75d72896aa43..e737f64218c9 100644
--- a/vllm/model_executor/models/minicpm_ocr.py
+++ b/vllm/model_executor/models/minicpm_ocr.py
@@ -147,6 +147,7 @@ def _compute_visual_tokens(
                 hf_image_size,
                 scale_res,
                 patch_size,
+                downsample_mode=image_processor.downsample_mode,
                 allow_upscale=True,
             )
             source_tokens = best_size[0] * best_size[1] // tokens_per_patch
@@ -156,6 +157,7 @@ def _compute_visual_tokens(
             hf_image_size,
             scale_res,
             patch_size,
+            downsample_mode=image_processor.downsample_mode,
         )
         source_tokens = best_resize[0] * best_resize[1] // tokens_per_patch
 
@@ -164,6 +166,7 @@ def _compute_visual_tokens(
             grids,
             scale_res,
             patch_size,
+            downsample_mode=image_processor.downsample_mode,
             allow_upscale=True,
         )
         patch_h = refine_size[0] // grids[0]
diff --git a/vllm/model_executor/models/minicpmv4_6.py b/vllm/model_executor/models/minicpmv4_6.py
index 333ab8eedd66..5b43a2c43c51 100644
--- a/vllm/model_executor/models/minicpmv4_6.py
+++ b/vllm/model_executor/models/minicpmv4_6.py
@@ -7,6 +7,8 @@
 
 import numpy as np
 import torch
+import transformers
+from packaging.version import Version
 from PIL import Image as PILImage
 from torch import nn
 from transformers import MiniCPMV4_6Config
@@ -122,6 +124,12 @@ def _expand_flags(value: object, num_items: int) -> list[bool | None]:
     return [any(bool(v) for v in raw)] * num_items
 
 
+# Transformers >= 5.21 sizes images by the `downsample_mode` merge factor
+RESIZE_TAKES_DOWNSAMPLE_MODE = Version(transformers.__version__) >= Version(
+    "5.21.0.dev0"
+)
+
+
 def _vision_downsample_mode(flag: bool | None) -> str | None:
     if flag is None:
         return None
@@ -739,6 +747,9 @@ def _compute_visual_tokens(
         scale_res = image_processor.scale_resolution
         downsample_mode = self._get_downsample_mode(downsample_mode)
         token_divisor = 4 if downsample_mode == "4x" else 16
+        resize_kwargs = (
+            {"downsample_mode": downsample_mode} if RESIZE_TAKES_DOWNSAMPLE_MODE else {}
+        )
 
         # vLLM ImageSize is (width, height); transformers expects (height, width)
         hf_image_size = (image_size.height, image_size.width)
@@ -762,6 +773,7 @@ def _compute_visual_tokens(
                 scale_res,
                 patch_size,
                 allow_upscale=True,
+                **resize_kwargs,
             )
             source_tokens = (
                 best_size[0] * best_size[1] // (patch_size * patch_size * token_divisor)
@@ -772,6 +784,7 @@ def _compute_visual_tokens(
             hf_image_size,
             scale_res,
             patch_size,
+            **resize_kwargs,
         )
         source_tokens = (
             best_resize[0] * best_resize[1] // (patch_size * patch_size * token_divisor)
@@ -782,6 +795,7 @@ def _compute_visual_tokens(
             scale_res,
             patch_size,
             allow_upscale=True,
+            **resize_kwargs,
         )
         patch_w = refine_size[0] // grids[0]
         patch_h = refine_size[1] // grids[1]
diff --git a/vllm/model_executor/models/mllama4.py b/vllm/model_executor/models/mllama4.py
index fbf393b554b6..968bed610ff0 100644
--- a/vllm/model_executor/models/mllama4.py
+++ b/vllm/model_executor/models/mllama4.py
@@ -45,7 +45,6 @@
 from vllm.model_executor.layers.linear import (
     ColumnParallelLinear,
     QKVParallelLinear,
-    ReplicatedLinear,
     RowParallelLinear,
 )
 from vllm.model_executor.layers.quantization import QuantizationConfig
@@ -266,38 +265,24 @@ def __init__(
             prefix=f"{prefix}.attn",
         )
 
-        if use_data_parallel:
-            self.qkv_proj = ReplicatedLinear(
-                self.embed_dim,
-                self.q_size + 2 * self.kv_size,
-                bias=True,
-                quant_config=quant_config,
-                prefix=f"{prefix}.qkv_proj",
-            )
-            self.o_proj = ReplicatedLinear(
-                self.num_heads * self.head_dim,
-                self.embed_dim,
-                bias=True,
-                quant_config=quant_config,
-                prefix=f"{prefix}.o_proj",
-            )
-        else:
-            self.qkv_proj = QKVParallelLinear(
-                self.embed_dim,
-                self.head_dim,
-                self.num_heads,
-                bias=True,
-                quant_config=quant_config,
-                prefix=f"{prefix}.qkv_proj",
-            )
-            self.o_proj = RowParallelLinear(
-                self.num_heads * self.head_dim,
-                self.embed_dim,
-                bias=True,
-                input_is_parallel=True,
-                quant_config=quant_config,
-                prefix=f"{prefix}.o_proj",
-            )
+        self.qkv_proj = QKVParallelLinear(
+            self.embed_dim,
+            self.head_dim,
+            self.num_heads,
+            bias=True,
+            quant_config=quant_config,
+            prefix=f"{prefix}.qkv_proj",
+            disable_tp=use_data_parallel,
+        )
+        self.o_proj = RowParallelLinear(
+            self.num_heads * self.head_dim,
+            self.embed_dim,
+            bias=True,
+            input_is_parallel=True,
+            quant_config=quant_config,
+            prefix=f"{prefix}.o_proj",
+            disable_tp=use_data_parallel,
+        )
 
         rope_parameters = {
             "rope_type": "mllama4",
@@ -1038,30 +1023,6 @@ def compute_logits(
     ) -> torch.Tensor | None:
         return self.language_model.compute_logits(hidden_states)
 
-    def _consolidate_qkv_weights(
-        self, weights: Iterable[tuple[str, torch.Tensor]]
-    ) -> Iterable[tuple[str, torch.Tensor]]:
-        qkv_idx_mappings = {
-            ".self_attn.q_proj": 0,
-            ".self_attn.k_proj": 1,
-            ".self_attn.v_proj": 2,
-        }
-        qkv_weights = {}
-        for name, loaded_weight in weights:
-            for weight_name, idx in qkv_idx_mappings.items():
-                if weight_name not in name:
-                    continue
-                new_name = name.replace(weight_name, ".self_attn.qkv_proj")
-                if new_name not in qkv_weights:
-                    qkv_weights[new_name] = [None] * 3
-                qkv_weights[new_name][idx] = loaded_weight
-                break
-            else:
-                yield name, loaded_weight
-        for key, weight in qkv_weights.items():
-            qkv_weight = torch.cat(weight, dim=0)
-            yield key, qkv_weight
-
     def _rename_weight_for_modelopt_checkpoint(self, name: str) -> str:
         """Rename weights from ModelOpt llama4 fp8 checkpoints to vLLM
         format."""
@@ -1115,13 +1076,10 @@ def _load_other_weights(
         """Load non-language-model weights with stacking support."""
         updated_params = set()
 
-        if self.use_data_parallel:
-            other_weights = self._consolidate_qkv_weights(other_weights)
-
         for name, loaded_weight in other_weights:
             # Try stacked parameter mapping first
             for param_name, weight_name, shard_id in stacked_params_mapping:
-                if weight_name not in name or self.use_data_parallel:
+                if weight_name not in name:
                     continue
                 name = name.replace(weight_name, param_name)
                 param = params_dict[name]
diff --git a/vllm/model_executor/models/muse_glimmer.py b/vllm/model_executor/models/muse_glimmer.py
index ef5ec5b608f2..d6be6c222ec6 100644
--- a/vllm/model_executor/models/muse_glimmer.py
+++ b/vllm/model_executor/models/muse_glimmer.py
@@ -54,6 +54,7 @@
 from vllm.logger import init_logger
 from vllm.model_executor.layers.activation import SiluAndMul
 from vllm.model_executor.layers.attention import Attention, MMEncoderAttention
+from vllm.model_executor.layers.layernorm import GemmaRMSNorm, RMSNorm
 from vllm.model_executor.layers.linear import (
     ColumnParallelLinear,
     MergedColumnParallelLinear,
@@ -343,41 +344,6 @@ def video_replacement(item_idx: int) -> PromptUpdateDetails:
         ]
 
 
-class MuseGlimmerRMSNorm(nn.Module):
-    """RMSNorm mirroring HF MuseGlimmer exactly (fp32 compute, cast at the end).
-
-    ``normed = _norm(x.float()) * (w.float() + weight_offset)`` cast back to the
-    input dtype. When ``with_scale`` is False the layer is weightless (used for
-    QK-norm and the token-embedding norm).
-    """
-
-    def __init__(
-        self,
-        dim: int | None = None,
-        eps: float = 1e-6,
-        with_scale: bool = True,
-        weight_offset: int = 0,
-    ) -> None:
-        super().__init__()
-        self.eps = eps
-        self.with_scale = with_scale
-        self.weight_offset = weight_offset
-        if with_scale:
-            assert dim is not None
-            self.weight = nn.Parameter(torch.zeros(dim))
-        else:
-            self.register_parameter("weight", None)
-
-    def _norm(self, x: torch.Tensor) -> torch.Tensor:
-        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
-
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        out = self._norm(hidden_states.float())
-        if self.with_scale:
-            out = out * (self.weight.float() + self.weight_offset)
-        return out.type_as(hidden_states)
-
-
 class MuseGlimmerVisionAttention(nn.Module):
     def __init__(
         self,
@@ -919,7 +885,7 @@ def __init__(
         self.use_rope = bool(config.layer_rope_theta[self.layer_idx])
 
         # Weightless, computed in fp32, applied per head over head_dim.
-        self.qk_norm = MuseGlimmerRMSNorm(eps=config.rms_norm_eps, with_scale=False)
+        self.qk_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps, has_weight=False)
         self.scale_query_by = config.qk_scale_factor
 
         self.qkv_proj = QKVParallelLinear(
@@ -1028,17 +994,15 @@ def __init__(
         )
         # Sandwich norms with baked +1 offset. Pre-norms use rms_norm_eps; the
         # post-norms use the (typically smaller) post_norm_eps.
-        self.input_layernorm = MuseGlimmerRMSNorm(
-            config.hidden_size, eps=config.rms_norm_eps, weight_offset=1
+        self.input_layernorm = GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = GemmaRMSNorm(
+            config.hidden_size, eps=config.post_norm_eps
         )
-        self.post_attention_layernorm = MuseGlimmerRMSNorm(
-            config.hidden_size, eps=config.post_norm_eps, weight_offset=1
+        self.pre_feedforward_layernorm = GemmaRMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
         )
-        self.pre_feedforward_layernorm = MuseGlimmerRMSNorm(
-            config.hidden_size, eps=config.rms_norm_eps, weight_offset=1
-        )
-        self.post_feedforward_layernorm = MuseGlimmerRMSNorm(
-            config.hidden_size, eps=config.post_norm_eps, weight_offset=1
+        self.post_feedforward_layernorm = GemmaRMSNorm(
+            config.hidden_size, eps=config.post_norm_eps
         )
 
     def forward(
@@ -1078,7 +1042,9 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         )
         # MuseGlimmer normalizes token embeddings with a weightless RMSNorm instead of
         # Gemma's sqrt(hidden_size) multiplier.
-        self.embed_norm = MuseGlimmerRMSNorm(eps=config.rms_norm_eps, with_scale=False)
+        self.embed_norm = RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps, has_weight=False
+        )
 
         self.start_layer, self.end_layer, self.layers = make_layers(
             config.num_hidden_layers,
@@ -1088,7 +1054,9 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
             prefix=f"{prefix}.layers",
         )
         # Final norm: weight-as-scale, no offset.
-        self.norm = MuseGlimmerRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.norm = RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps, dtype=torch.float32
+        )
         self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
             ["hidden_states", "residual"], config.hidden_size
         )
@@ -1250,7 +1218,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
 
         self.vision_encoder: MuseGlimmerVisionEncoder | None
         self.vision_adapter: MuseGlimmerVisionAdapter | None
-        self.perception_emb_norm: MuseGlimmerRMSNorm | None
+        self.perception_emb_norm: RMSNorm | None
         if self.has_vision:
             with self._mark_tower_model(vllm_config, {"image", "video"}):
                 self.vision_encoder = MuseGlimmerVisionEncoder(
@@ -1263,8 +1231,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
                     text_config.hidden_size,
                     bias=False,
                 )
-                self.perception_emb_norm = MuseGlimmerRMSNorm(
-                    eps=text_config.rms_norm_eps, with_scale=False
+                self.perception_emb_norm = RMSNorm(
+                    text_config.hidden_size,
+                    eps=text_config.rms_norm_eps,
+                    has_weight=False,
                 )
         else:
             self.vision_encoder = None
diff --git a/vllm/model_executor/models/nano_nemotron_vl.py b/vllm/model_executor/models/nano_nemotron_vl.py
index b77206295d5c..77038b363f68 100644
--- a/vllm/model_executor/models/nano_nemotron_vl.py
+++ b/vllm/model_executor/models/nano_nemotron_vl.py
@@ -28,6 +28,7 @@
 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.fusion.mm_input_norm import FusedMMInputNorm
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 from vllm.model_executor.models.interfaces import (
@@ -934,6 +935,7 @@ class NemotronH_Nano_VL_V2(
     SupportsMultiModalPruning,
     SupportsLoRA,
 ):
+    supports_mm_device_do_normalize = True
     requires_sequential_video_encoding = True
     """Temporarily needed for dynamic res video w/ conv3d, doesn't support bs>1 yet"""
 
@@ -992,12 +994,23 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         assert isinstance(llm_dtype, torch.dtype)
         self.llm_dtype = llm_dtype
         with self._mark_tower_model(vllm_config, {"image", "video", "audio"}):
+            self.input_norm: FusedMMInputNorm | None = (
+                FusedMMInputNorm(
+                    image_mean=config.norm_mean,
+                    image_std=config.norm_std,
+                    rescale_factor=1.0 / 255.0,
+                )
+                if multimodal_config.mm_device_do_normalize
+                else None
+            )
             self.vision_model = self.get_vit_model_from_radio_config(config).to(
                 llm_dtype
             )
 
             # Construct the vision projection.
-            vit_hidden_size = config.vit_hidden_size
+            vit_hidden_size = getattr(config, "vision_hidden_size", None)
+            if vit_hidden_size is None:
+                vit_hidden_size = config.vit_hidden_size
             vision_projection_hidden_size = config.projector_hidden_size
             llm_hidden_size = config.text_config.hidden_size
 
@@ -1099,10 +1112,23 @@ def pixel_shuffle_dynamic_res(
 
         return x
 
+    def _normalize_pixel_values(self, pixel_values: torch.Tensor) -> torch.Tensor:
+        if self.input_norm is None:
+            return pixel_values
+        shape = pixel_values.shape
+        if pixel_values.ndim == 4:
+            flattened = pixel_values.reshape(shape[0], -1)
+        else:
+            assert pixel_values.ndim == 3
+            flattened = pixel_values.reshape(-1, shape[-1])
+        normalized = self.input_norm(flattened, self.llm_dtype)
+        return normalized.reshape(shape)
+
     def extract_feature_dynamic(
         self, pixel_values: torch.Tensor, imgs_sizes: list[tuple[int, int]]
     ):
         """Dynamic resolution extract_feature for images."""
+        pixel_values = self._normalize_pixel_values(pixel_values)
         _, vit_embeds = self.vision_model(pixel_values, imgs_sizes=imgs_sizes)
         vit_embeds = vit_embeds.to(dtype=torch.bfloat16)
         if self.vision_final_layernorm is not None:
@@ -1134,6 +1160,7 @@ def extract_feature(
         vit_embeds_list = []
         for i in range(0, N, micro_batch_size):
             chunk = pixel_values[i : i + micro_batch_size]
+            chunk = self._normalize_pixel_values(chunk)
             if num_frames is not None and T > 1:
                 _, vit_embeds = self.vision_model(chunk, num_frames=chunk.shape[0])
             else:
@@ -1620,7 +1647,11 @@ def is_adapter_weights(weight: tuple[str, torch.Tensor]):
             return weight[0].startswith("mlp1")
 
         def is_vision_weights(name: str) -> bool:
-            return name.startswith("vision_model.radio_model.")
+            # The whole RADIO tower lives under ``vision_model.``: legacy
+            # remote-code checkpoints nest it as ``vision_model.radio_model.*``,
+            # native Transformers checkpoints as ``vision_model.embeddings.*`` /
+            # ``vision_model.encoder.*``. RadioModel.load_weights handles both.
+            return name.startswith("vision_model.")
 
         def is_sound_weights(name: str) -> bool:
             return name.startswith("sound")
@@ -1653,7 +1684,9 @@ def llm_weights_gen():
                 elif is_vision_weights(name):
                     if not load_multimodal_weights:
                         continue
-                    # Convert: vision_model.radio_model.* → radio_model.*
+                    # Strip the ``vision_model.`` prefix; the remainder is the
+                    # RADIO key (legacy ``radio_model.*`` or native
+                    # ``embeddings.*``/``encoder.*``).
                     hf_key = name[len("vision_model.") :]
                     vision_weights.append((hf_key, w.detach().clone()))
                 elif is_sound_weights(name):
diff --git a/vllm/model_executor/models/nemotron3_5_transcribe.py b/vllm/model_executor/models/nemotron3_5_transcribe.py
new file mode 100644
index 000000000000..6ec2147418ec
--- /dev/null
+++ b/vllm/model_executor/models/nemotron3_5_transcribe.py
@@ -0,0 +1,670 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+"""Inference-only NVIDIA Nemotron-3.5-Transcribe.
+
+A NeMo speech-language model. The audio tower (log-mel, Parallel Expert
+Encoder, projection; see ``nemotron3_5_transcribe_audio.py``) turns each
+80 ms of audio into one embedding, and those embeddings replace the
+``<|audio|>`` placeholders in the prompt of a NemotronH decoder (hybrid
+Mamba, MoE and attention layers).
+
+The checkpoint uses NeMo names: the audio tower under ``perception.*``, the
+decoder under ``llm.*``, the multi-token-prediction head under ``llm.mtp.*``
+and an optional CTC timestamp head, which is not loaded, under
+``ctc_timestamp.*``.
+"""
+
+from collections.abc import Iterable, Iterator, Mapping, Sequence
+from typing import Annotated, Any, Literal, cast
+
+import numpy as np
+import torch
+from torch import nn
+from transformers import BatchFeature
+
+from vllm.config import ModelConfig, SpeechToTextConfig, VllmConfig
+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
+from vllm.multimodal import MULTIMODAL_REGISTRY
+from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargsItems
+from vllm.multimodal.parse import MultiModalDataItems, MultiModalDataParser
+from vllm.multimodal.processing import (
+    BaseDummyInputsBuilder,
+    BaseMultiModalProcessor,
+    BaseProcessingInfo,
+    PromptReplacement,
+    PromptUpdate,
+    PromptUpdateDetails,
+)
+from vllm.sequence import IntermediateTensors
+from vllm.tokenizers import cached_tokenizer_from_config
+from vllm.transformers_utils.configs.nemotron3_5_transcribe import (
+    Nemotron3_5TranscribeConfig,
+)
+from vllm.utils.tensor_schema import TensorSchema, TensorShape
+
+from .interfaces import (
+    IsHybrid,
+    MultiModalEmbeddings,
+    SupportsEagle3,
+    SupportsMambaPrefixCaching,
+    SupportsMultiModal,
+    SupportsMultiModalEmbeddings,
+    SupportsPP,
+    SupportsTranscription,
+)
+from .module_mapping import MultiModelKeys
+from .nemotron3_5_transcribe_audio import Nemotron3_5TranscribeAudioTower
+from .nemotron_h import NemotronHForCausalLM
+from .nemotron_h_mtp import NemotronHMTP
+from .utils import (
+    WeightsMapper,
+    _merge_multimodal_embeddings,
+    init_vllm_registered_model,
+    maybe_prefix,
+)
+
+logger = init_logger(__name__)
+
+AUDIO_PLACEHOLDER = "<|audio|>"
+# Default longest audio used to profile memory, unless --limit-mm-per-prompt
+# sets an audio length.
+_DEFAULT_MAX_AUDIO_DURATION_S = 3600.0
+# Prompt tokens kept free beside the longest dummy audio.
+_DUMMY_TEXT_TOKEN_RESERVE = 256
+
+# The transcription instruction of the model card's evaluation prompts. The
+# model was trained with the audio after the instruction.
+_TRANSCRIPTION_PROMPT = (
+    "Audio transcription. Language: as spoken. "
+    "Wording: verbatim, including fillers and repetitions. "
+    "Formatting: punctuation and capitalization. Output: transcript only."
+)
+# Training audio was at most one hour, so transcription sends pieces of at most an
+# hour; vLLM splits longer audio at quiet points.
+_MAX_TRANSCRIPTION_CLIP_S = 60 * 60
+# Context kept free for the transcript of each second of audio: a fast speaker,
+# with speaker tags.
+_TRANSCRIPT_TOKENS_PER_S = 4.0
+
+# The languages of the model card's FLEURS evaluation.
+ISO639_1_SUPPORTED_LANGS = {
+    "ar": "Arabic",
+    "bg": "Bulgarian",
+    "cs": "Czech",
+    "da": "Danish",
+    "de": "German",
+    "el": "Greek",
+    "en": "English",
+    "es": "Spanish",
+    "et": "Estonian",
+    "fi": "Finnish",
+    "fr": "French",
+    "he": "Hebrew",
+    "hi": "Hindi",
+    "hr": "Croatian",
+    "hu": "Hungarian",
+    "id": "Indonesian",
+    "it": "Italian",
+    "ja": "Japanese",
+    "ko": "Korean",
+    "lt": "Lithuanian",
+    "lv": "Latvian",
+    "mt": "Maltese",
+    "nl": "Dutch",
+    "no": "Norwegian",
+    "pl": "Polish",
+    "pt": "Portuguese",
+    "ro": "Romanian",
+    "ru": "Russian",
+    "sk": "Slovak",
+    "sl": "Slovenian",
+    "sv": "Swedish",
+    "th": "Thai",
+    "tl": "Filipino",
+    "tr": "Turkish",
+    "uk": "Ukrainian",
+    "vi": "Vietnamese",
+    "zh": "Chinese",
+}
+
+
+def _num_audio_embeddings(audio_config: Mapping[str, Mapping], num_samples: int) -> int:
+    """Embeddings the audio tower produces for ``num_samples`` samples."""
+    preprocessor = audio_config["preprocessor"]
+    factor = audio_config["asr_encoder"]["subsampling_factor"]
+    n_fft = preprocessor["n_fft"]
+    num_frames = (num_samples + n_fft // 2 * 2 - n_fft) // preprocessor["hop_length"]
+    return (num_frames + factor - 1) // factor
+
+
+def _transcription_clip_s(
+    audio_config: Mapping[str, Mapping], max_model_len: int
+) -> int:
+    """Longest piece ``/v1/audio/transcriptions`` sends: at most an hour, and short
+    enough that its embeddings and transcript fit in the context."""
+    sample_rate = audio_config["preprocessor"]["sample_rate"]
+    hour = 3600
+    tokens_per_s = (
+        _num_audio_embeddings(audio_config, hour * sample_rate) / hour
+        + _TRANSCRIPT_TOKENS_PER_S
+    )
+    fits_s = int((max_model_len - _DUMMY_TEXT_TOKEN_RESERVE) / tokens_per_s)
+    return max(1, min(_MAX_TRANSCRIPTION_CLIP_S, fits_s))
+
+
+def _max_audio_samples(
+    audio_config: Mapping[str, Mapping], seq_len: int, override_samples: int | None
+) -> int:
+    """Longest audio whose embeddings fit in ``seq_len`` beside a short prompt."""
+    preprocessor = audio_config["preprocessor"]
+    min_samples = preprocessor["hop_length"]
+    limit = int(
+        override_samples or _DEFAULT_MAX_AUDIO_DURATION_S * preprocessor["sample_rate"]
+    )
+
+    def num_tokens(num_samples: int) -> int:
+        return _num_audio_embeddings(audio_config, max(num_samples, min_samples))
+
+    max_tokens = max(1, seq_len - _DUMMY_TEXT_TOKEN_RESERVE)
+    if num_tokens(limit) <= max_tokens:
+        return limit
+    low, high = min_samples, limit
+    while low < high:
+        mid = (low + high + 1) // 2
+        if num_tokens(mid) <= max_tokens:
+            low = mid
+        else:
+            high = mid - 1
+    return low
+
+
+class Nemotron3_5TranscribeAudioInputs(TensorSchema):
+    """Dimensions:
+    - b: Number of audio clips
+    - t: Number of samples (16 kHz)
+    """
+
+    type: Literal["audio_signal"] = "audio_signal"
+    audio_signal: Annotated[
+        torch.Tensor | list[torch.Tensor],
+        TensorShape("b", "t", dynamic_dims={"t"}),
+    ]
+    audio_signal_length: Annotated[torch.Tensor, TensorShape("b")]
+
+
+class Nemotron3_5TranscribeProcessingInfo(BaseProcessingInfo):
+    def get_hf_config(self) -> Nemotron3_5TranscribeConfig:
+        return self.ctx.get_hf_config(Nemotron3_5TranscribeConfig)
+
+    def get_supported_mm_limits(self) -> Mapping[str, int | None]:
+        return {"audio": None}
+
+    def get_data_parser(self) -> MultiModalDataParser:
+        return MultiModalDataParser(
+            target_sr=self.get_sampling_rate(), target_channels=1
+        )
+
+    def get_sampling_rate(self) -> int:
+        return self.get_hf_config().audio_config["preprocessor"]["sample_rate"]
+
+    def get_min_num_samples(self) -> int:
+        """Shorter clips are zero-padded so that they yield one embedding."""
+        return self.get_hf_config().audio_config["preprocessor"]["hop_length"]
+
+    def get_num_audio_tokens(self, num_samples: int) -> int:
+        num_samples = max(num_samples, self.get_min_num_samples())
+        return _num_audio_embeddings(self.get_hf_config().audio_config, num_samples)
+
+    def get_audio_token_id(self) -> int:
+        return self.get_tokenizer().convert_tokens_to_ids(AUDIO_PLACEHOLDER)
+
+    def get_max_audio_samples(
+        self, seq_len: int, override_samples: int | None = None
+    ) -> int:
+        """Longest audio whose embeddings fit in ``seq_len`` beside a short prompt."""
+        return _max_audio_samples(
+            self.get_hf_config().audio_config, seq_len, override_samples
+        )
+
+    def get_audio_length_override(self) -> int | None:
+        """The audio length set with ``--limit-mm-per-prompt``, in samples."""
+        options = self.ctx.get_mm_config().limit_per_prompt.get("audio")
+        return getattr(options, "length", None)
+
+    def get_mm_max_tokens_per_item(
+        self, seq_len: int, mm_counts: Mapping[str, int]
+    ) -> Mapping[str, int]:
+        # The same longest clip as the dummy inputs, so the encoder cache that
+        # vLLM sizes from this number admits every clip the limits allow.
+        samples = self.get_max_audio_samples(seq_len, self.get_audio_length_override())
+        return {"audio": self.get_num_audio_tokens(samples)}
+
+
+class Nemotron3_5TranscribeDummyInputsBuilder(
+    BaseDummyInputsBuilder[Nemotron3_5TranscribeProcessingInfo]
+):
+    def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
+        return AUDIO_PLACEHOLDER * mm_counts.get("audio", 0)
+
+    def get_dummy_mm_data(
+        self,
+        seq_len: int,
+        mm_counts: Mapping[str, int],
+        mm_options: MultiModalDummyOptions,
+    ) -> MultiModalDataDict:
+        audio_options = mm_options.get("audio") if mm_options else None
+        override = getattr(audio_options, "length", None)
+        if override is None:
+            override = self.info.get_audio_length_override()
+        return {
+            "audio": self._get_dummy_audios(
+                length=self.info.get_max_audio_samples(seq_len, override),
+                num_audios=mm_counts.get("audio", 0),
+            )
+        }
+
+
+class Nemotron3_5TranscribeMultiModalProcessor(
+    BaseMultiModalProcessor[Nemotron3_5TranscribeProcessingInfo]
+):
+    def _call_hf_processor(
+        self,
+        hf_data: Mapping[str, object],
+        hf_kwargs: Mapping[str, object],
+    ) -> BatchFeature:
+        # There is no Hugging Face processor: the model computes features
+        # from the waveform, so pass each clip through as float32 samples.
+        min_samples = self.info.get_min_num_samples()
+        signals, lengths = [], []
+        for audio in cast(Iterable[Any], hf_data.get("audio", [])):
+            signal = torch.as_tensor(np.asarray(audio, dtype=np.float32)).flatten()
+            if signal.shape[0] < min_samples:
+                signal = torch.nn.functional.pad(
+                    signal, (0, min_samples - signal.shape[0])
+                )
+            signals.append(signal)
+            lengths.append(signal.shape[0])
+        return BatchFeature(
+            {
+                "audio_signal": signals,
+                "audio_signal_length": torch.tensor(lengths, dtype=torch.long),
+            }
+        )
+
+    def _get_mm_fields_config(
+        self,
+        hf_inputs: BatchFeature,
+        hf_processor_mm_kwargs: Mapping[str, object],
+    ) -> Mapping[str, MultiModalFieldConfig]:
+        return dict(
+            audio_signal=MultiModalFieldConfig.batched("audio"),
+            audio_signal_length=MultiModalFieldConfig.batched(
+                "audio", keep_on_cpu=True
+            ),
+        )
+
+    def _get_prompt_updates(
+        self,
+        mm_items: MultiModalDataItems,
+        hf_processor_mm_kwargs: Mapping[str, object],
+        out_mm_kwargs: MultiModalKwargsItems,
+    ) -> Sequence[PromptUpdate]:
+        audio_token_id = self.info.get_audio_token_id()
+        lengths = out_mm_kwargs.get_data().get("audio_signal_length")
+        num_tokens = (
+            []
+            if lengths is None
+            else [
+                self.info.get_num_audio_tokens(int(n))
+                for n in torch.as_tensor(lengths).tolist()
+            ]
+        )
+
+        def get_replacement(item_idx: int) -> PromptUpdateDetails:
+            return PromptUpdateDetails.select_token_id(
+                [audio_token_id] * num_tokens[item_idx], embed_token_id=audio_token_id
+            )
+
+        return [
+            PromptReplacement(
+                modality="audio",
+                target=[audio_token_id],
+                replacement=get_replacement,
+            )
+        ]
+
+
+def _unpack_experts(
+    name: str, tensor: torch.Tensor
+) -> Iterator[tuple[str, torch.Tensor]]:
+    """NeMo packs the routed experts of a MoE layer into one rank-3 tensor per
+    projection; vLLM loads one transposed matrix per expert."""
+    for suffix, projection in (
+        (".experts.down_projs", "down_proj"),
+        (".experts.gate_and_up_projs", "up_proj"),
+    ):
+        if name.endswith(suffix):
+            prefix = name[: -len(suffix)]
+            for expert in range(tensor.shape[0]):
+                yield (
+                    f"{prefix}.experts.{expert}.{projection}.weight",
+                    tensor[expert].t(),
+                )
+            return
+    yield name, tensor
+
+
+def _pad_vocab_rows(tensor: torch.Tensor, vocab_size: int) -> torch.Tensor:
+    if tensor.shape[0] >= vocab_size:
+        return tensor
+    return torch.cat(
+        [tensor, tensor.new_zeros(vocab_size - tensor.shape[0], tensor.shape[1])], dim=0
+    )
+
+
+def _decoder_weights(
+    weights: Iterable[tuple[str, torch.Tensor]], vocab_size: int
+) -> Iterator[tuple[str, torch.Tensor]]:
+    """Shape the mapped decoder tensors for ``NemotronHForCausalLM``.
+
+    NeMo packs the routed experts of each MoE layer into two tensors; the FP8 and
+    NVFP4 checkpoints have one tensor per expert, with its quantization scales.
+    The embedding and LM head are zero-padded to the decoder's padded vocabulary.
+    """
+    for name, tensor in weights:
+        if name in (
+            "model.embed_tokens.weight",
+            "model.embeddings.weight",
+            "lm_head.weight",
+        ):
+            tensor = _pad_vocab_rows(tensor, vocab_size)
+        yield from _unpack_experts(name, tensor)
+
+
+def _mtp_weights(
+    weights: Iterable[tuple[str, torch.Tensor]],
+) -> Iterator[tuple[str, torch.Tensor]]:
+    """Select the MTP head's NeMo tensors and rename them for ``NemotronHMTP``.
+
+    The head is ``llm.mtp.*``; the drafter shares the target's embedding and LM
+    head.
+    """
+    for name, tensor in weights:
+        if name.startswith("llm.mtp.") and not name.endswith("._extra_state"):
+            yield from _unpack_experts(name[len("llm.") :], tensor)
+
+
+@MULTIMODAL_REGISTRY.register_processor(
+    Nemotron3_5TranscribeMultiModalProcessor,
+    info=Nemotron3_5TranscribeProcessingInfo,
+    dummy_inputs=Nemotron3_5TranscribeDummyInputsBuilder,
+)
+class Nemotron3_5TranscribeForConditionalGeneration(
+    nn.Module,
+    SupportsMultiModal,
+    SupportsPP,
+    IsHybrid,
+    SupportsMambaPrefixCaching,
+    SupportsEagle3,
+    SupportsTranscription,
+):
+    packed_modules_mapping = NemotronHForCausalLM.packed_modules_mapping
+    # The BF16 checkpoint names the decoder in NeMo's layout (``llm.*``; Mamba
+    # ``A_log``/``D``/``dt_bias`` in a float32 holder; final norm ``norm``). The FP8
+    # and NVFP4 checkpoints, and their quantization configs, use Hugging Face names.
+    hf_to_vllm_mapper = WeightsMapper(
+        orig_to_new_substr={".mixer._fp32_params.": ".mixer."},
+        orig_to_new_prefix={
+            "llm.model.norm.": "language_model.model.norm_f.",
+            "llm.model.": "language_model.model.",
+            "llm.lm_head.": "language_model.lm_head.",
+            "backbone.": "language_model.model.",
+            "lm_head": "language_model.lm_head",
+            "perception.": "audio_tower.",
+        },
+    )
+    supported_languages = ISO639_1_SUPPORTED_LANGS
+
+    @classmethod
+    def get_placeholder_str(cls, modality: str, i: int) -> str | None:
+        if modality.startswith("audio"):
+            return AUDIO_PLACEHOLDER
+        raise ValueError("Only audio modality is supported")
+
+    @classmethod
+    def get_speech_to_text_config(
+        cls, model_config: ModelConfig, task_type: str
+    ) -> SpeechToTextConfig:
+        audio_config = model_config.hf_config.audio_config
+        return SpeechToTextConfig(
+            sample_rate=audio_config["preprocessor"]["sample_rate"],
+            max_audio_clip_s=_transcription_clip_s(
+                audio_config, model_config.max_model_len
+            ),
+        )
+
+    @classmethod
+    def get_generation_prompt(cls, stt_params: SpeechToTextParams) -> PromptType:
+        """The model card's transcription prompt, whatever the request's language.
+
+        The prompt says "Language: as spoken" and the model detects the
+        language. No published prompt names it, so the request's language is
+        only validated.
+        """
+        if stt_params.task_type != "transcribe":
+            raise ValueError(
+                "Nemotron-3.5-Transcribe does not support task "
+                f"{stt_params.task_type!r} yet."
+            )
+        if stt_params.request_prompt or stt_params.hotwords:
+            logger.warning_once(
+                "Nemotron-3.5-Transcribe ignores the prompt and hotwords of "
+                "transcription requests."
+            )
+        tokenizer = cached_tokenizer_from_config(stt_params.model_config)
+        prompt = tokenizer.apply_chat_template(
+            [
+                {
+                    "role": "user",
+                    "content": f"{_TRANSCRIPTION_PROMPT} {AUDIO_PLACEHOLDER}",
+                }
+            ],
+            tokenize=False,
+            add_generation_prompt=True,
+            enable_thinking=False,
+        )
+        return TokensPrompt(
+            prompt_token_ids=tokenizer.encode(prompt, add_special_tokens=False),
+            multi_modal_data={
+                "audio": (stt_params.audio, stt_params.stt_config.sample_rate)
+            },
+        )
+
+    @classmethod
+    def get_num_audio_tokens(
+        cls,
+        audio_duration_s: float,
+        stt_config: SpeechToTextConfig,
+        model_config: ModelConfig,
+    ) -> int:
+        audio_config = model_config.hf_config.audio_config
+        num_samples = max(
+            int(audio_duration_s * stt_config.sample_rate),
+            audio_config["preprocessor"]["hop_length"],
+        )
+        return _num_audio_embeddings(audio_config, num_samples)
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+        config: Nemotron3_5TranscribeConfig = vllm_config.model_config.hf_config
+        self.config = config
+
+        with self._mark_language_model(vllm_config):
+            self.language_model = init_vllm_registered_model(
+                vllm_config=vllm_config,
+                hf_config=config.text_config,
+                prefix=maybe_prefix(prefix, "language_model"),
+            )
+        with self._mark_tower_model(vllm_config, "audio"):
+            self.audio_tower = Nemotron3_5TranscribeAudioTower(
+                config.audio_config,
+                quant_config=vllm_config.quant_config,
+                prefix=maybe_prefix(prefix, "audio_tower"),
+            )
+
+        self.make_empty_intermediate_tensors = (
+            self.language_model.make_empty_intermediate_tensors
+        )
+
+    def _parse_audio_input(
+        self, **kwargs: object
+    ) -> Nemotron3_5TranscribeAudioInputs | None:
+        audio_signal = kwargs.pop("audio_signal", None)
+        if audio_signal is None:
+            return None
+        return Nemotron3_5TranscribeAudioInputs(
+            audio_signal=audio_signal,
+            audio_signal_length=kwargs.pop("audio_signal_length"),
+        )
+
+    def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
+        audio_input = self._parse_audio_input(**kwargs)
+        if audio_input is None:
+            return []
+        signals = audio_input.audio_signal
+        if isinstance(signals, torch.Tensor):
+            signals = list(signals)
+        return tuple(
+            self.audio_tower.encode(signals, audio_input.audio_signal_length.flatten())
+        )
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None,
+        positions: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        **kwargs: object,
+    ) -> torch.Tensor | IntermediateTensors:
+        if intermediate_tensors is not None:
+            inputs_embeds = None
+        return self.language_model(
+            input_ids=input_ids,
+            positions=positions,
+            intermediate_tensors=intermediate_tensors,
+            inputs_embeds=inputs_embeds,
+            **kwargs,
+        )
+
+    def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:
+        logits = self.language_model.compute_logits(hidden_states)
+        if logits is not None:
+            # Rows past the trained vocabulary (the audio placeholder and the
+            # embedding headroom) are never generated.
+            logits[..., self.config.output_vocab_size :] = -torch.inf
+        return logits
+
+    def get_mm_mapping(self) -> MultiModelKeys:
+        return MultiModelKeys.from_string_field(
+            language_model="language_model",
+            connector="audio_tower.proj",
+            tower_model=["audio_tower.preprocessor", "audio_tower.encoder"],
+        )
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        load_audio = isinstance(self.audio_tower, Nemotron3_5TranscribeAudioTower)
+        audio_weights: list[tuple[str, torch.Tensor]] = []
+
+        def decoder_source() -> Iterator[tuple[str, torch.Tensor]]:
+            for name, tensor in self.hf_to_vllm_mapper.apply(weights):
+                if name.endswith("._extra_state"):
+                    continue
+                if name.startswith("audio_tower."):
+                    if load_audio:
+                        # Copied so the tensor outlives a reusable read buffer.
+                        audio_weights.append(
+                            (name[len("audio_tower.") :], tensor.detach().clone())
+                        )
+                elif name.startswith("llm.mtp."):
+                    continue  # loaded by the multi-token-prediction drafter
+                elif name.startswith("ctc_timestamp."):
+                    continue  # the CTC timestamp head and its tokenizer
+                elif name.startswith("language_model."):
+                    yield name[len("language_model.") :], tensor
+                else:
+                    raise ValueError(f"Unexpected checkpoint weight {name!r}.")
+
+        decoder_weights = _decoder_weights(
+            decoder_source(), self.config.text_config.vocab_size
+        )
+        loaded = {
+            f"language_model.{name}"
+            for name in self.language_model.load_weights(decoder_weights)
+        }
+        if load_audio:
+            loaded |= {
+                f"audio_tower.{name}"
+                for name in self.audio_tower.load_weights(audio_weights)
+            }
+        return loaded
+
+    @classmethod
+    def get_mamba_state_shape_from_config(cls, vllm_config: VllmConfig):
+        text_config = vllm_config.model_config.hf_config.text_config
+        return NemotronHForCausalLM.get_mamba_state_shape_from_config(
+            vllm_config.with_hf_config(text_config)
+        )
+
+    @classmethod
+    def get_mamba_state_dtype_from_config(cls, vllm_config: VllmConfig):
+        text_config = vllm_config.model_config.hf_config.text_config
+        return NemotronHForCausalLM.get_mamba_state_dtype_from_config(
+            vllm_config.with_hf_config(text_config)
+        )
+
+    @classmethod
+    def get_mamba_state_copy_func(cls):
+        return NemotronHForCausalLM.get_mamba_state_copy_func()
+
+
+class Nemotron3_5TranscribeMTP(NemotronHMTP, SupportsMultiModalEmbeddings):
+    """Multi-token-prediction draft head of Nemotron-3.5-Transcribe.
+
+    One attention-plus-MoE block, trained to be reused for every draft token.
+    It was trained on the decoder's input stream, so it receives the audio
+    embeddings at the ``<|audio|>`` positions, not the placeholder embedding.
+    """
+
+    def embed_input_ids(
+        self,
+        input_ids: torch.Tensor,
+        multimodal_embeddings: MultiModalEmbeddings | None = None,
+        *,
+        is_multimodal: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        inputs_embeds = self.model.get_input_embeddings(input_ids)
+        if multimodal_embeddings is None or len(multimodal_embeddings) == 0:
+            return inputs_embeds
+        if is_multimodal is None:
+            raise ValueError("is_multimodal is required with multimodal_embeddings.")
+        return _merge_multimodal_embeddings(
+            inputs_embeds, multimodal_embeddings, is_multimodal
+        )
+
+    def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:
+        logits = super().compute_logits(hidden_states)
+        if logits is not None:
+            logits[..., self.config.speechlm_output_vocab_size :] = -torch.inf
+        return logits
+
+    def is_unused_checkpoint_weight(self, name: str) -> bool:
+        # Shard selection passes the raw llm.mtp.* names; load_weights passes
+        # the mtp.* names that _mtp_weights produces.
+        return super().is_unused_checkpoint_weight(name.removeprefix("llm."))
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        return super().load_weights(_mtp_weights(weights))
diff --git a/vllm/model_executor/models/nemotron3_5_transcribe_audio.py b/vllm/model_executor/models/nemotron3_5_transcribe_audio.py
new file mode 100644
index 000000000000..65b54f6d52a0
--- /dev/null
+++ b/vllm/model_executor/models/nemotron3_5_transcribe_audio.py
@@ -0,0 +1,1111 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+"""Audio tower of Nemotron-3.5-Transcribe.
+
+A native port of the inference path of NeMo's ``AudioPerceptionModule`` with a
+``ParallelExpertEncoder``:
+
+    waveform -> log-mel -> Parallel Expert Encoder -> Linear -> decoder width
+
+The Parallel Expert Encoder runs an ASR ``TransformerEncoder`` and a streaming
+Sortformer speaker encoder over the same mel features. Both run in windows of
+``window_frames`` encoder frames with left and right context, and the speaker
+encoder carries a speaker cache and a FIFO between windows. Thresholded speaker
+activity is added to the normalized ASR states through a fixed sinusoidal
+kernel. Module names follow NeMo so the checkpoint loads by prefix.
+"""
+
+import math
+from collections.abc import Iterable, Sequence
+from dataclasses import dataclass
+from typing import Any, cast
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from vllm.model_executor.layers.attention.mm_encoder_attention import (
+    MMEncoderAttention,
+)
+from vllm.model_executor.layers.linear import ReplicatedLinear
+from vllm.model_executor.layers.quantization import QuantizationConfig
+from vllm.utils.torch_utils import async_tensor_h2d
+
+from .utils import maybe_prefix
+from .vision import get_multimodal_config
+
+# Checkpoint tensors that the inference path does not use. The speaker
+# encoder's own preprocessor is unused because both encoders read the mel
+# features of the perception preprocessor; ``hidden_to_spks`` is a frozen
+# legacy head.
+_UNUSED_CHECKPOINT_PREFIXES = (
+    "encoder.diarization_model.preprocessor.",
+    "encoder.diarization_model.sortformer_modules.hidden_to_spks.",
+)
+
+
+class _TowerLinear(ReplicatedLinear):
+    """A replicated linear layer over the last dimension of any input shape.
+
+    vLLM's quantized kernels take ``(tokens, features)``; the tower's activations
+    are ``(batch, frames, features)``.
+    """
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # Built with return_bias=False, so the layer returns only the output.
+        out = cast(torch.Tensor, super().forward(x.reshape(-1, x.shape[-1])))
+        return out.view(*x.shape[:-1], out.shape[-1])
+
+
+def _linear(
+    in_features: int,
+    out_features: int,
+    *,
+    bias: bool,
+    quant_config: QuantizationConfig | None,
+    prefix: str,
+) -> _TowerLinear:
+    """A linear layer the checkpoint's quantization config applies to by name."""
+    return _TowerLinear(
+        in_features,
+        out_features,
+        bias=bias,
+        quant_config=quant_config,
+        prefix=prefix,
+        return_bias=False,
+        disable_tp=True,
+    )
+
+
+class MelSpectrogram(nn.Module):
+    """Log-mel features as NeMo ``FilterbankFeatures`` computes them at inference.
+
+    Dither is off and no normalization is applied; frames past each row's
+    length are set to ``pad_value``. The filterbank and window come from the
+    checkpoint and are held in the model dtype, as in the reference, and are
+    upcast to float32 for the computation.
+    """
+
+    def __init__(self, cfg: dict[str, Any]) -> None:
+        super().__init__()
+        self.n_fft = cfg["n_fft"]
+        self.hop_length = cfg["hop_length"]
+        self.win_length = cfg["win_length"]
+        self.preemph = cfg["preemph"]
+        self.log_zero_guard_value = cfg["log_zero_guard_value"]
+        self.mag_power = cfg["mag_power"]
+        self.pad_value = cfg["pad_value"]
+        self.register_buffer("fb", torch.zeros(1, cfg["n_mels"], self.n_fft // 2 + 1))
+        self.register_buffer("window", torch.zeros(self.win_length))
+
+    def get_seq_len(self, num_samples: torch.Tensor) -> torch.Tensor:
+        pad_amount = self.n_fft // 2 * 2
+        seq_len = torch.floor_divide(
+            num_samples + pad_amount - self.n_fft, self.hop_length
+        ).to(torch.long)
+        return torch.where(num_samples == 0, torch.zeros_like(seq_len), seq_len)
+
+    @torch.no_grad()
+    def forward(
+        self, audio: torch.Tensor, num_samples: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        seq_len = self.get_seq_len(num_samples)
+        x = audio.float()
+        if self.preemph:
+            time_mask = torch.arange(x.shape[1], device=x.device).unsqueeze(
+                0
+            ) < num_samples.unsqueeze(1)
+            x = torch.cat((x[:, :1], x[:, 1:] - self.preemph * x[:, :-1]), dim=1)
+            x = x.masked_fill(~time_mask, 0.0)
+        x = torch.stft(
+            x,
+            n_fft=self.n_fft,
+            hop_length=self.hop_length,
+            win_length=self.win_length,
+            center=True,
+            window=self.window.float(),
+            return_complex=True,
+            pad_mode="constant",
+        )
+        x = torch.sqrt(torch.view_as_real(x).pow(2).sum(-1))
+        if self.mag_power != 1.0:
+            x = x.pow(self.mag_power)
+        x = torch.matmul(self.fb.float(), x)
+        x = torch.log(x + self.log_zero_guard_value)
+        pad_mask = torch.arange(x.shape[-1], device=x.device).unsqueeze(
+            0
+        ) >= seq_len.unsqueeze(1)
+        x = x.masked_fill(pad_mask.unsqueeze(1), self.pad_value)
+        return x, seq_len
+
+
+class AudioPreprocessor(nn.Module):
+    """Holds the featurizer under NeMo's ``preprocessor.featurizer`` name."""
+
+    def __init__(self, cfg: dict[str, Any]) -> None:
+        super().__init__()
+        self.featurizer = MelSpectrogram(cfg)
+
+    def forward(
+        self, audio: torch.Tensor, num_samples: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        return self.featurizer(audio, num_samples)
+
+
+def normalize_per_feature(x: torch.Tensor, seq_len: torch.Tensor) -> torch.Tensor:
+    """NeMo ``normalize_batch(x, seq_len, "per_feature")`` on ``(B, C, T)`` features."""
+    batch_size, _, max_time = x.shape
+    time_steps = (
+        torch.arange(max_time, device=x.device)
+        .unsqueeze(0)
+        .expand(batch_size, max_time)
+    )
+    valid_mask = time_steps < seq_len.unsqueeze(1)
+    denominator = valid_mask.sum(dim=1)
+    reference = x[:, :, 0] if max_time else x.new_zeros((batch_size, x.shape[1]))
+    reference = reference.masked_fill((denominator == 0).unsqueeze(1), 0.0)
+    centered = torch.where(valid_mask.unsqueeze(1), x - reference.unsqueeze(2), 0.0)
+    mean = reference + centered.sum(dim=2) / denominator.clamp_min(1).unsqueeze(1)
+    variance_sum = torch.sum(
+        torch.where(valid_mask.unsqueeze(1), x - mean.unsqueeze(2), 0.0) ** 2, dim=2
+    )
+    std = torch.sqrt(variance_sum / (denominator.unsqueeze(1) - 1.0))
+    std = std.masked_fill(std.isnan(), 0.0) + 1e-5
+    normalized = (x - mean.unsqueeze(2)) / std.unsqueeze(2)
+    return normalized.masked_fill(~valid_mask.unsqueeze(1), 0.0)
+
+
+class RotaryEmbedding:
+    """NeMo ``RotaryPositionalEncoding``: NeoX rotation of the first ``rot_dim`` dims.
+
+    Angles are computed in float32 on the CPU and stored in the activation
+    dtype before use, as the reference does, so rotations match bit for bit.
+    """
+
+    def __init__(self, head_dim: int, rotary_fraction: float, base: float) -> None:
+        self.rot_dim = int(head_dim * rotary_fraction)
+        if self.rot_dim < 2 or self.rot_dim % 2:
+            raise ValueError(
+                f"Rotary dim must be a positive even number, got {self.rot_dim}."
+            )
+        self.base = base
+        self._cache: dict[
+            tuple[torch.dtype, torch.device], tuple[torch.Tensor, torch.Tensor]
+        ] = {}
+
+    def cos_sin(
+        self, length: int, dtype: torch.dtype, device: torch.device
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        key = (dtype, device)
+        cached = self._cache.get(key)
+        if cached is None or cached[0].shape[0] < length:
+            inv_freq = 1.0 / (
+                self.base
+                ** (
+                    torch.arange(0, self.rot_dim, 2, dtype=torch.float32, device="cpu")
+                    / self.rot_dim
+                )
+            )
+            positions = torch.arange(
+                0, max(length, 1024), dtype=torch.float32, device="cpu"
+            )
+            freqs = torch.outer(positions, inv_freq)
+            emb = torch.cat((freqs, freqs), dim=-1)
+            cached = (
+                emb.cos().to(device=device, dtype=dtype),
+                emb.sin().to(device=device, dtype=dtype),
+            )
+            self._cache[key] = cached
+        return cached[0][:length], cached[1][:length]
+
+    def apply(
+        self, t: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
+    ) -> torch.Tensor:
+        def rotate(x: torch.Tensor) -> torch.Tensor:
+            half = x.shape[-1] // 2
+            return x * cos + torch.cat((-x[..., half:], x[..., :half]), dim=-1) * sin
+
+        if self.rot_dim == t.shape[-1]:
+            return rotate(t)
+        return torch.cat(
+            (rotate(t[..., : self.rot_dim]), t[..., self.rot_dim :]), dim=-1
+        )
+
+
+class FeatureStacking(nn.Module):
+    """Stack ``factor`` consecutive frames and project them to the model width."""
+
+    def __init__(
+        self,
+        factor: int,
+        feat_in: int,
+        feat_out: int,
+        quant_config: QuantizationConfig | None,
+        prefix: str,
+    ) -> None:
+        super().__init__()
+        self.factor = factor
+        self.proj = _linear(
+            factor * feat_in,
+            feat_out,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.proj",
+        )
+
+    def forward(
+        self, x: torch.Tensor, lengths: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        x = x.transpose(1, 2)
+        batch_size, num_frames, channels = x.shape
+        pad = (self.factor - num_frames % self.factor) % self.factor
+        if pad:
+            x = F.pad(x, (0, 0, 0, pad))
+        x = x.reshape(
+            batch_size, (num_frames + pad) // self.factor, channels * self.factor
+        )
+        return self.proj(x), self.output_lengths(lengths)
+
+    def output_lengths(self, lengths: torch.Tensor) -> torch.Tensor:
+        return (lengths + self.factor - 1) // self.factor
+
+
+@dataclass
+class PackedFrames:
+    """The valid frames of a padded ``(B, T)`` batch, packed row after row.
+
+    ``rows`` indexes the flattened batch and ``positions`` holds each frame's
+    position in its row. ``cu_seqlens``, ``max_seqlen`` and ``sequence_lengths``
+    are in the format of the attention backend.
+    """
+
+    rows: torch.Tensor
+    positions: torch.Tensor
+    cu_seqlens: torch.Tensor
+    max_seqlen: torch.Tensor
+    sequence_lengths: torch.Tensor | None
+
+
+class MultiHeadAttention(nn.Module):
+    def __init__(
+        self,
+        cfg: dict[str, Any],
+        rope: RotaryEmbedding,
+        quant_config: QuantizationConfig | None,
+        prefix: str,
+    ) -> None:
+        super().__init__()
+        self.d_model = cfg["d_model"]
+        self.n_heads = cfg["n_heads"]
+        self.head_dim = self.d_model // self.n_heads
+        self.rope = rope
+        self.w_qkv = _linear(
+            self.d_model,
+            3 * self.d_model,
+            bias=cfg["qkv_bias"],
+            quant_config=quant_config,
+            prefix=f"{prefix}.w_qkv",
+        )
+        self.out_proj = _linear(
+            self.d_model,
+            self.d_model,
+            bias=True,
+            quant_config=quant_config,
+            prefix=f"{prefix}.out_proj",
+        )
+        self.qk_norm = cfg["qk_norm"]
+        if self.qk_norm:
+            self.q_norm = nn.LayerNorm(self.head_dim)
+            self.k_norm = nn.LayerNorm(self.head_dim)
+        self.attn = MMEncoderAttention(
+            num_heads=self.n_heads,
+            head_size=self.head_dim,
+            prefix=f"{prefix}.attn",
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        packed: PackedFrames,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+    ) -> torch.Tensor:
+        """Attend within each row of the packed ``(N, d_model)`` frames."""
+        num_frames = x.shape[0]
+        qkv = self.w_qkv(x).view(num_frames, 3, self.n_heads, self.head_dim)
+        q, k, v = qkv.unbind(1)
+        if self.qk_norm:
+            q = self.q_norm(q).to(v.dtype)
+            k = self.k_norm(k).to(v.dtype)
+        q = self.rope.apply(q, cos.to(q.dtype), sin.to(q.dtype))
+        k = self.rope.apply(k, cos.to(k.dtype), sin.to(k.dtype))
+        out = self.attn(
+            q.unsqueeze(0),
+            k.unsqueeze(0),
+            v.unsqueeze(0),
+            cu_seqlens=packed.cu_seqlens,
+            max_seqlen=packed.max_seqlen,
+            sequence_lengths=packed.sequence_lengths,
+        )
+        return self.out_proj(out.reshape(num_frames, self.d_model))
+
+
+class TransformerBlock(nn.Module):
+    def __init__(
+        self,
+        cfg: dict[str, Any],
+        rope: RotaryEmbedding,
+        quant_config: QuantizationConfig | None,
+        prefix: str,
+    ) -> None:
+        super().__init__()
+        d_model = cfg["d_model"]
+        self.norm1 = nn.LayerNorm(d_model)
+        self.attn = MultiHeadAttention(cfg, rope, quant_config, f"{prefix}.attn")
+        self.norm2 = nn.LayerNorm(d_model)
+        ff_hidden = int(cfg["ff_expansion"] * d_model)
+        # Indices 0 and 3 match the checkpoint's ``ffn.net.{0,3}``.
+        self.ffn = nn.Module()
+        self.ffn.net = nn.Sequential(
+            _linear(
+                d_model,
+                ff_hidden,
+                bias=True,
+                quant_config=quant_config,
+                prefix=f"{prefix}.ffn.net.0",
+            ),
+            nn.GELU(),
+            nn.Identity(),
+            _linear(
+                ff_hidden,
+                d_model,
+                bias=True,
+                quant_config=quant_config,
+                prefix=f"{prefix}.ffn.net.3",
+            ),
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        packed: PackedFrames,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+    ) -> torch.Tensor:
+        x = x + self.attn(self.norm1(x), packed, cos, sin)
+        return x + self.ffn.net(self.norm2(x))
+
+
+class TransformerEncoder(nn.Module):
+    """NeMo ``TransformerEncoder`` with RoPE and full attention, at inference.
+
+    ``forward`` takes ``(B, C, T)`` mel features, or ``(B, T, d_model)``
+    pre-encoded embeddings with ``bypass_pre_encode=True``, and returns
+    ``(B, d_model, T')`` states with their lengths. Only each row's valid frames
+    run through the layers, packed for variable-length attention; the padded
+    positions of the output are zero. ``lengths_cpu``, the lengths on the host,
+    lets the packing skip a device-to-host copy.
+    """
+
+    def __init__(
+        self,
+        cfg: dict[str, Any],
+        quant_config: QuantizationConfig | None,
+        prefix: str,
+    ) -> None:
+        super().__init__()
+        self.d_model = cfg["d_model"]
+        self.subsampling_factor = cfg["subsampling_factor"]
+        self.pre_encode = FeatureStacking(
+            self.subsampling_factor,
+            cfg["feat_in"],
+            self.d_model,
+            quant_config,
+            f"{prefix}.pre_encode",
+        )
+        self.xscale = math.sqrt(self.d_model) if cfg["xscaling"] else None
+        self.rope = RotaryEmbedding(
+            self.d_model // cfg["n_heads"], cfg["rotary_fraction"], cfg["rope_base"]
+        )
+        self.embed_norm = (
+            nn.LayerNorm(self.d_model) if cfg["pre_block_norm"] else nn.Identity()
+        )
+        self.layers = nn.ModuleList(
+            [
+                TransformerBlock(cfg, self.rope, quant_config, f"{prefix}.layers.{i}")
+                for i in range(cfg["n_layers"])
+            ]
+        )
+        self.final_norm = nn.LayerNorm(self.d_model)
+        self.attn_backend = self.layers[0].attn.attn.attn_backend
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        lengths: torch.Tensor,
+        bypass_pre_encode: bool = False,
+        lengths_cpu: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if not bypass_pre_encode:
+            x, lengths = self.pre_encode(x, lengths)
+            if lengths_cpu is not None:
+                lengths_cpu = self.pre_encode.output_lengths(lengths_cpu)
+        lengths = lengths.to(torch.int64)
+        if lengths_cpu is None:
+            lengths_cpu = lengths.cpu()
+        batch_size, num_frames, width = x.shape
+        out = x.new_zeros(batch_size * num_frames, width)
+        packed = self.pack(lengths_cpu, num_frames, x.device)
+        if packed.rows.numel():
+            h = x.reshape(-1, width).index_select(0, packed.rows)
+            if self.xscale:
+                h = h * self.xscale
+            h = self.embed_norm(h)
+            cos, sin = self.rope.cos_sin(num_frames, x.dtype, x.device)
+            cos = cos.index_select(0, packed.positions).unsqueeze(1)
+            sin = sin.index_select(0, packed.positions).unsqueeze(1)
+            for layer in self.layers:
+                h = layer(h, packed, cos, sin)
+            out.index_copy_(0, packed.rows, self.final_norm(h))
+        return out.view(batch_size, num_frames, width).transpose(1, 2), lengths
+
+    def pack(
+        self, lengths_cpu: torch.Tensor, num_frames: int, device: torch.device
+    ) -> PackedFrames:
+        """Index each row's valid frames; rows without any are left out."""
+        lengths_cpu = lengths_cpu.to(torch.int64).clamp(min=0, max=num_frames)
+        valid = torch.arange(num_frames) < lengths_cpu.unsqueeze(1)
+        rows = valid.flatten().nonzero().squeeze(1)
+        positions = torch.arange(num_frames).expand_as(valid)[valid]
+        seq_lens = lengths_cpu[lengths_cpu > 0].numpy()
+        cu_seqlens = np.concatenate(([0], np.cumsum(seq_lens))).astype(np.int32)
+        backend = self.attn_backend
+        return PackedFrames(
+            rows=async_tensor_h2d(rows, device=device),
+            positions=async_tensor_h2d(positions, device=device),
+            cu_seqlens=MMEncoderAttention.maybe_recompute_cu_seqlens(
+                backend, cu_seqlens, self.d_model, 1, device
+            ),
+            max_seqlen=torch.tensor(
+                MMEncoderAttention.compute_max_seqlen(backend, cu_seqlens),
+                dtype=torch.int32,
+            ),
+            sequence_lengths=MMEncoderAttention.maybe_compute_seq_lens(
+                backend, cu_seqlens, device
+            ),
+        )
+
+
+def downsample_preds(preds: torch.Tensor, factor: int) -> torch.Tensor:
+    """Average ``(B, T, S)`` probabilities over windows of ``factor`` frames."""
+    if factor == 1 or preds.shape[1] == 0:
+        return preds
+    pooled = F.avg_pool1d(
+        preds.transpose(1, 2),
+        kernel_size=factor,
+        stride=factor,
+        ceil_mode=True,
+        count_include_pad=False,
+    )
+    return pooled.transpose(1, 2)
+
+
+def align_frames(preds: torch.Tensor, target_len: int) -> torch.Tensor:
+    """Pad with the last frame, or truncate, ``(B, T, S)`` to ``target_len`` frames."""
+    current_len = preds.shape[1]
+    if current_len < target_len:
+        if current_len == 0:
+            raise ValueError(
+                "Speaker activity cannot be empty when the encoder output is not."
+            )
+        preds = torch.cat(
+            [preds, preds[:, -1:].repeat(1, target_len - current_len, 1)], dim=1
+        )
+    return preds[:, :target_len]
+
+
+@dataclass
+class SpeakerCacheState:
+    """Synchronous streaming state of the Sortformer speaker encoder."""
+
+    spkcache: torch.Tensor
+    fifo: torch.Tensor
+    spkcache_preds: torch.Tensor | None = None
+    spkcache_compressed: bool = False
+
+
+class SortformerModules(nn.Module):
+    """Speaker head and speaker cache of NeMo ``SortformerModules`` at inference."""
+
+    def __init__(
+        self,
+        cfg: dict[str, Any],
+        online: dict[str, Any],
+        upsample_factor: int,
+        quant_config: QuantizationConfig | None,
+        prefix: str,
+    ) -> None:
+        super().__init__()
+        self.n_spk = cfg["num_speakers"]
+        self.upsample_factor = upsample_factor
+        hidden = cfg["tf_d_model"]
+        self.encoder_proj = _linear(
+            cfg["fc_d_model"],
+            hidden,
+            bias=True,
+            quant_config=quant_config,
+            prefix=f"{prefix}.encoder_proj",
+        )
+        self.subpixel_upsample = nn.Conv1d(
+            hidden, hidden * upsample_factor, kernel_size=3, padding=1
+        )
+        self.first_hidden_to_hidden = _linear(
+            hidden,
+            hidden,
+            bias=True,
+            quant_config=quant_config,
+            prefix=f"{prefix}.first_hidden_to_hidden",
+        )
+        self.single_hidden_to_spks = _linear(
+            hidden,
+            self.n_spk,
+            bias=True,
+            quant_config=quant_config,
+            prefix=f"{prefix}.single_hidden_to_spks",
+        )
+        self.learnable_sil_emb = nn.Parameter(torch.zeros(cfg["fc_d_model"]))
+
+        self.fifo_len = online["fifo_len"]
+        self.spkcache_len = online["spkcache_len"]
+        self.spkcache_update_period = online["spkcache_update_period"]
+        self.sil_frames_per_spk = cfg["spkcache_sil_frames_per_spk"]
+        self.pred_score_threshold = cfg["pred_score_threshold"]
+        self.max_index = cfg["max_index"]
+        self.scores_boost_latest = cfg["scores_boost_latest"]
+        self.strong_boost_rate = cfg["strong_boost_rate"]
+        self.weak_boost_rate = cfg["weak_boost_rate"]
+        self.min_pos_scores_rate = cfg["min_pos_scores_rate"]
+        if self.spkcache_len < (1 + self.sil_frames_per_spk) * self.n_spk:
+            raise ValueError("spkcache_len is too small for the number of speakers.")
+
+    def speaker_probs(
+        self, hidden: torch.Tensor, lengths: torch.Tensor
+    ) -> torch.Tensor:
+        """Upsampled speaker probabilities, ``(B, T * upsample_factor, n_spk)``."""
+        batch_size, num_frames, width = hidden.shape
+        mask = torch.arange(num_frames, device=hidden.device).unsqueeze(
+            0
+        ) < lengths.unsqueeze(1)
+        # The convolution reads neighbouring frames: zero the padding so a clip's
+        # last frame does not depend on the clips batched with it.
+        hidden = hidden.masked_fill(~mask.unsqueeze(-1), 0.0)
+        upsampled = self.subpixel_upsample(hidden.transpose(1, 2)).transpose(1, 2)
+        upsampled = upsampled.reshape(
+            batch_size, num_frames, self.upsample_factor, width
+        ).reshape(batch_size, num_frames * self.upsample_factor, width)
+        logits = self.single_hidden_to_spks(
+            F.relu(self.first_hidden_to_hidden(F.relu(upsampled)))
+        )
+        output_mask = mask.repeat_interleave(self.upsample_factor, dim=1)
+        return torch.sigmoid(logits) * output_mask.unsqueeze(-1)
+
+    def update(
+        self,
+        state: SpeakerCacheState,
+        chunk: torch.Tensor,
+        preds: torch.Tensor,
+        lc: int,
+        rc: int,
+    ) -> SpeakerCacheState:
+        """NeMo ``SortformerModules.streaming_update`` (synchronous)."""
+        spkcache_len, fifo_len = state.spkcache.shape[1], state.fifo.shape[1]
+        chunk_len = chunk.shape[1] - lc - rc
+        fifo_preds = preds[:, spkcache_len : spkcache_len + fifo_len]
+        chunk = chunk[:, lc : chunk_len + lc]
+        chunk_preds = preds[
+            :, spkcache_len + fifo_len + lc : spkcache_len + fifo_len + chunk_len + lc
+        ]
+        fifo = torch.cat([state.fifo, chunk], dim=1)
+        fifo_preds = torch.cat([fifo_preds, chunk_preds], dim=1)
+        spkcache, spkcache_preds, compressed = (
+            state.spkcache,
+            state.spkcache_preds,
+            state.spkcache_compressed,
+        )
+
+        if fifo_len + chunk_len > self.fifo_len:
+            pop_out_len = max(
+                self.spkcache_update_period, chunk_len - self.fifo_len + fifo_len
+            )
+            pop_out_len = min(pop_out_len, fifo_len + chunk_len)
+            pop_out_embs, pop_out_preds = (
+                fifo[:, :pop_out_len],
+                fifo_preds[:, :pop_out_len],
+            )
+            fifo = fifo[:, pop_out_len:]
+            spkcache = torch.cat([spkcache, pop_out_embs], dim=1)
+            if compressed:
+                spkcache_preds = torch.cat([spkcache_preds, pop_out_preds], dim=1)
+            else:
+                spkcache_preds = torch.cat(
+                    [preds[:, :spkcache_len], pop_out_preds], dim=1
+                )
+            if spkcache.shape[1] > self.spkcache_len:
+                spkcache, spkcache_preds = self._compress_spkcache(
+                    spkcache, spkcache_preds
+                )
+                compressed = True
+        return SpeakerCacheState(spkcache, fifo, spkcache_preds, compressed)
+
+    def _compress_spkcache(
+        self, emb_seq: torch.Tensor, preds: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Keep the ``spkcache_len`` most informative frames, as NeMo does."""
+        batch_size, _, n_spk = preds.shape
+        per_spk = self.spkcache_len // n_spk - self.sil_frames_per_spk
+        strong_boost = math.floor(per_spk * self.strong_boost_rate)
+        weak_boost = math.floor(per_spk * self.weak_boost_rate)
+        min_pos_scores = math.floor(per_spk * self.min_pos_scores_rate)
+
+        threshold = self.pred_score_threshold
+        log_probs = torch.log(torch.clamp(preds, min=threshold))
+        log_1_probs = torch.log(torch.clamp(1.0 - preds, min=threshold))
+        log_1_probs_sum = log_1_probs.sum(dim=2).unsqueeze(-1).expand(-1, -1, n_spk)
+        scores = log_probs - log_1_probs + log_1_probs_sum - math.log(0.5)
+
+        is_speech = preds > 0.5
+        scores = torch.where(is_speech, scores, float("-inf"))
+        is_pos = scores > 0
+        replace = (
+            (~is_pos) * is_speech * (is_pos.sum(dim=1).unsqueeze(1) >= min_pos_scores)
+        )
+        scores = torch.where(replace, float("-inf"), scores)
+
+        if self.scores_boost_latest > 0:
+            scores[:, self.spkcache_len :, :] += self.scores_boost_latest
+        scores = self._boost_topk_scores(scores, strong_boost, scale_factor=2)
+        scores = self._boost_topk_scores(scores, weak_boost, scale_factor=1)
+        if self.sil_frames_per_spk > 0:
+            pad = torch.full(
+                (batch_size, self.sil_frames_per_spk, n_spk),
+                float("inf"),
+                dtype=scores.dtype,
+                device=scores.device,
+            )
+            scores = torch.cat([scores, pad], dim=1)
+
+        num_frames = scores.shape[1]
+        flat = scores.permute(0, 2, 1).reshape(batch_size, -1)
+        topk_values, topk_indices = torch.topk(
+            flat, self.spkcache_len, dim=1, sorted=False
+        )
+        topk_indices = torch.where(
+            topk_values != float("-inf"), topk_indices, self.max_index
+        )
+        topk_indices, _ = torch.sort(topk_indices, dim=1)
+        is_disabled = topk_indices == self.max_index
+        topk_indices = torch.remainder(topk_indices, num_frames)
+        is_disabled += topk_indices >= num_frames - self.sil_frames_per_spk
+        topk_indices[is_disabled] = 0
+
+        sil_emb = (
+            self.learnable_sil_emb.to(dtype=emb_seq.dtype)
+            .unsqueeze(0)
+            .expand(batch_size, -1)
+        )
+        emb_gathered = torch.gather(
+            emb_seq, 1, topk_indices.unsqueeze(-1).expand(-1, -1, emb_seq.shape[2])
+        )
+        emb_gathered = torch.where(
+            is_disabled.unsqueeze(-1),
+            sil_emb.unsqueeze(1).expand(-1, self.spkcache_len, -1),
+            emb_gathered,
+        )
+        preds_gathered = torch.gather(
+            preds, 1, topk_indices.unsqueeze(-1).expand(-1, -1, n_spk)
+        )
+        preds_gathered = torch.where(is_disabled.unsqueeze(-1), 0.0, preds_gathered)
+        return emb_gathered, preds_gathered
+
+    @staticmethod
+    def _boost_topk_scores(
+        scores: torch.Tensor, n_boost_per_spk: int, scale_factor: float
+    ) -> torch.Tensor:
+        batch_size, _, n_spk = scores.shape
+        _, topk_indices = torch.topk(
+            scores, n_boost_per_spk, dim=1, largest=True, sorted=False
+        )
+        batch_indices = torch.arange(batch_size, device=scores.device).view(-1, 1, 1)
+        speaker_indices = torch.arange(n_spk, device=scores.device).view(1, 1, -1)
+        scores[batch_indices, topk_indices, speaker_indices] -= scale_factor * math.log(
+            0.5
+        )
+        return scores
+
+
+class SortformerSpeakerEncoder(nn.Module):
+    """Streaming inference of the high-resolution Sortformer in NeMo's PEE.
+
+    Mirrors ``SortformerEncLabelModel.forward_streaming_step`` for synchronous
+    streaming without speaker permutation. Each step returns speaker
+    probabilities on the encoder frame grid for the window's core frames.
+    """
+
+    def __init__(
+        self,
+        encoder_cfg: dict[str, Any],
+        sortformer_cfg: dict[str, Any],
+        online: dict[str, Any],
+        quant_config: QuantizationConfig | None,
+        prefix: str,
+    ) -> None:
+        super().__init__()
+        self.encoder = TransformerEncoder(
+            encoder_cfg, quant_config, f"{prefix}.encoder"
+        )
+        self.subsampling_factor = encoder_cfg["subsampling_factor"]
+        self.sortformer_modules = SortformerModules(
+            sortformer_cfg,
+            online,
+            self.subsampling_factor,
+            quant_config,
+            f"{prefix}.sortformer_modules",
+        )
+        self.fc_d_model = sortformer_cfg["fc_d_model"]
+        self.n_spk = sortformer_cfg["num_speakers"]
+
+    def init_state(
+        self, batch_size: int, dtype: torch.dtype, device: torch.device
+    ) -> SpeakerCacheState:
+        empty = torch.zeros(
+            (batch_size, 0, self.fc_d_model), dtype=dtype, device=device
+        )
+        return SpeakerCacheState(spkcache=empty, fifo=empty.clone())
+
+    def step(
+        self,
+        features: torch.Tensor,
+        lengths: torch.Tensor,
+        state: SpeakerCacheState,
+        left_offset: int,
+        right_offset: int,
+        lengths_cpu: torch.Tensor | None = None,
+    ) -> tuple[SpeakerCacheState, torch.Tensor]:
+        """Process one window of ``(B, C, T)`` mel features.
+
+        Returns the new state and the core frames' speaker probabilities on the
+        encoder grid.
+        """
+        modules = self.sortformer_modules
+        factor = self.subsampling_factor
+        chunk_embs, chunk_lengths = self.encoder.pre_encode(features, lengths)
+        embs = torch.cat([state.spkcache, state.fifo, chunk_embs], dim=1)
+        context_len = state.spkcache.shape[1] + state.fifo.shape[1]
+        emb_lengths = context_len + chunk_lengths
+        emb_lengths_cpu = (
+            None
+            if lengths_cpu is None
+            else context_len + self.encoder.pre_encode.output_lengths(lengths_cpu)
+        )
+        states, state_lengths = self.encoder(
+            embs, emb_lengths, bypass_pre_encode=True, lengths_cpu=emb_lengths_cpu
+        )
+        hidden = modules.encoder_proj(states.transpose(1, 2))
+        high_res_preds = modules.speaker_probs(hidden, state_lengths)
+
+        lc = round(left_offset / factor)
+        rc = math.ceil(right_offset / factor)
+        preds = downsample_preds(high_res_preds, modules.upsample_factor)
+        valid = torch.arange(preds.shape[1], device=preds.device).view(
+            1, -1, 1
+        ) < state_lengths.view(-1, 1, 1)
+        preds = torch.where(valid, preds, 0.0)
+        state = modules.update(state, chunk_embs, preds, lc, rc)
+
+        chunk_len = chunk_embs.shape[1] - lc - rc
+        start = (context_len + lc) * modules.upsample_factor
+        chunk_preds = high_res_preds[
+            :, start : start + chunk_len * modules.upsample_factor
+        ]
+        return state, downsample_preds(chunk_preds, factor)
+
+
+class ParallelExpertEncoder(nn.Module):
+    """NeMo ``ParallelExpertEncoder`` online inference with a transformer ASR encoder.
+
+    ``forward`` takes unnormalized ``(B, C, T)`` mel features and returns the
+    fused ``(B, d_model, T')`` states.
+    """
+
+    def __init__(
+        self,
+        audio_config: dict[str, Any],
+        quant_config: QuantizationConfig | None,
+        prefix: str,
+    ) -> None:
+        super().__init__()
+        online = audio_config["online"]
+        fusion = audio_config["fusion"]
+        self.asr_encoder = TransformerEncoder(
+            audio_config["asr_encoder"],
+            quant_config,
+            f"{prefix}.asr_encoder",
+        )
+        self.diarization_model = SortformerSpeakerEncoder(
+            audio_config["speaker_encoder"],
+            audio_config["sortformer"],
+            online,
+            quant_config,
+            f"{prefix}.diarization_model",
+        )
+        self.subsampling_factor = self.asr_encoder.subsampling_factor
+        self.window_feat_len = online["window_frames"] * self.subsampling_factor
+        self.left_ctx_feat_len = online["left_context_frames"] * self.subsampling_factor
+        self.right_ctx_feat_len = (
+            online["right_context_frames"] * self.subsampling_factor
+        )
+        self.speaker_feature_mode = fusion["speaker_feature_mode"]
+        self.speaker_activity_threshold = fusion["speaker_activity_threshold"]
+        self.spk_kernel_scale = fusion["spk_kernel_scale"]
+        n_spk = self.diarization_model.n_spk
+        self.asr_norm = nn.LayerNorm(self.asr_encoder.d_model)
+        self.diar_norm = nn.LayerNorm(n_spk)
+        self.register_buffer(
+            "diar_kernel",
+            self._sinusoid_kernel(n_spk, self.asr_encoder.d_model),
+            persistent=False,
+        )
+
+    @staticmethod
+    def _sinusoid_kernel(num_speakers: int, width: int) -> torch.Tensor:
+        position = torch.arange(num_speakers, dtype=torch.float32).unsqueeze(1)
+        div_term = torch.exp(
+            torch.arange(0, width, 2, dtype=torch.float32)
+            * -(math.log(10000.0) / width)
+        )
+        kernel = torch.zeros(num_speakers, width, dtype=torch.float32)
+        kernel[:, 0::2] = torch.sin(position * div_term)
+        kernel[:, 1::2] = torch.cos(position * div_term)
+        return kernel
+
+    def _output_boundary(self, feat_boundary: int) -> int:
+        return (feat_boundary + self.subsampling_factor - 1) // self.subsampling_factor
+
+    @property
+    def dtype(self) -> torch.dtype:
+        """The model dtype, which quantized layers' weights may not have."""
+        return self.asr_norm.weight.dtype
+
+    def forward(
+        self,
+        mel: torch.Tensor,
+        mel_lengths: torch.Tensor,
+        max_feat_len: int,
+        mel_lengths_cpu: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        total_feat_len = min(mel.shape[-1], max_feat_len)
+        num_windows = max(1, math.ceil(total_feat_len / self.window_feat_len))
+        asr_signal = normalize_per_feature(mel, mel_lengths).to(self.dtype)
+        diar_signal = mel.to(self.dtype)
+        state = self.diarization_model.init_state(mel.shape[0], self.dtype, mel.device)
+
+        asr_windows: list[torch.Tensor] = []
+        speaker_windows: list[torch.Tensor] = []
+        for index in range(num_windows):
+            start = index * self.window_feat_len
+            end = min(start + self.window_feat_len, total_feat_len)
+            context_start = max(start - self.left_ctx_feat_len, 0)
+            context_end = min(end + self.right_ctx_feat_len, total_feat_len)
+            left_offset, right_offset = start - context_start, context_end - end
+            window_lengths = (mel_lengths - context_start).clamp(
+                min=0, max=context_end - context_start
+            )
+            window_lengths_cpu = (
+                None
+                if mel_lengths_cpu is None
+                else (mel_lengths_cpu - context_start).clamp(
+                    min=0, max=context_end - context_start
+                )
+            )
+
+            encoded, _ = self.asr_encoder(
+                asr_signal[:, :, context_start:context_end],
+                window_lengths,
+                lengths_cpu=window_lengths_cpu,
+            )
+            left_drop = left_offset // self.subsampling_factor
+            core_len = self._output_boundary(end) - self._output_boundary(start)
+            core_len = max(0, min(core_len, encoded.shape[-1] - left_drop))
+            asr_windows.append(encoded[:, :, left_drop : left_drop + core_len])
+
+            state, window_preds = self.diarization_model.step(
+                diar_signal[:, :, context_start:context_end],
+                window_lengths,
+                state,
+                left_offset,
+                right_offset,
+                lengths_cpu=window_lengths_cpu,
+            )
+            speaker_windows.append(align_frames(window_preds, core_len))
+
+        asr_encoded = torch.cat(asr_windows, dim=2)
+        speaker_probs = torch.cat(speaker_windows, dim=1)
+        return self._fuse(asr_encoded, speaker_probs)
+
+    def _fuse(
+        self, asr_encoded: torch.Tensor, speaker_probs: torch.Tensor
+    ) -> torch.Tensor:
+        states = asr_encoded.transpose(1, 2)
+        speaker_probs = align_frames(speaker_probs, states.shape[1]).to(
+            device=states.device, dtype=states.dtype
+        )
+        if self.speaker_feature_mode == "thresholded":
+            speaker_features = (speaker_probs > self.speaker_activity_threshold).to(
+                states.dtype
+            )
+        else:
+            speaker_features = speaker_probs
+        normalized_states = self.asr_norm(states)
+        normalized_speakers = self.diar_norm(speaker_features)
+        infusion = torch.matmul(
+            normalized_speakers, self.diar_kernel.to(normalized_speakers.dtype)
+        )
+        return (normalized_states + self.spk_kernel_scale * infusion).transpose(1, 2)
+
+
+class Nemotron3_5TranscribeAudioTower(nn.Module):
+    """Waveform to decoder-width embeddings: mel, Parallel Expert Encoder, Linear."""
+
+    def __init__(
+        self,
+        audio_config: dict[str, Any],
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        mm_config = get_multimodal_config()
+        if mm_config is not None and mm_config.mm_encoder_attn_dtype == "fp8":
+            raise ValueError(
+                "Nemotron-3.5-Transcribe does not support "
+                "mm_encoder_attn_dtype='fp8': the audio tower's attention runs in "
+                "the model dtype. The FP8 checkpoints quantize the tower's linear "
+                "layers without this setting."
+            )
+        self.preprocessor = AudioPreprocessor(audio_config["preprocessor"])
+        self.encoder = ParallelExpertEncoder(
+            audio_config, quant_config, maybe_prefix(prefix, "encoder")
+        )
+        projection = audio_config["projection"]
+        self.proj = _linear(
+            projection["input_dim"],
+            projection["output_dim"],
+            bias=True,
+            quant_config=quant_config,
+            prefix=maybe_prefix(prefix, "proj"),
+        )
+        self.subsampling_factor = self.encoder.subsampling_factor
+
+    def encode(
+        self, signals: Sequence[torch.Tensor], num_samples_cpu: torch.Tensor
+    ) -> list[torch.Tensor]:
+        """Encode each clip into its embeddings, returned in input order.
+
+        Every row of a batch runs every window of the batch's longest clip, so a
+        clip is batched only with clips that span as many windows.
+        """
+        window = self.encoder.window_feat_len
+        num_mel_frames = self.preprocessor.featurizer.get_seq_len(num_samples_cpu)
+        groups: dict[int, list[int]] = {}
+        for index, num_windows in enumerate(
+            ((num_mel_frames + window - 1) // window).tolist()
+        ):
+            groups.setdefault(num_windows, []).append(index)
+
+        device = self.encoder.asr_norm.weight.device
+        embeddings: dict[int, torch.Tensor] = {}
+        for indices in groups.values():
+            lengths_cpu = num_samples_cpu[indices]
+            audio = torch.zeros(
+                len(indices), int(lengths_cpu.max()), dtype=torch.float32, device=device
+            )
+            for row, index in enumerate(indices):
+                length = int(lengths_cpu[row])
+                audio[row, :length] = (
+                    signals[index]
+                    .flatten()[:length]
+                    .to(device=device, dtype=torch.float32)
+                )
+            lengths = (
+                lengths_cpu
+                if device.type == "cpu"
+                else async_tensor_h2d(lengths_cpu, device=device, dtype=torch.long)
+            )
+            for index, emb in zip(indices, self(audio, lengths, lengths_cpu)):
+                embeddings[index] = emb
+        return [embeddings[index] for index in range(len(signals))]
+
+    def forward(
+        self,
+        audio: torch.Tensor,
+        num_samples: torch.Tensor,
+        num_samples_cpu: torch.Tensor,
+    ) -> list[torch.Tensor]:
+        """Encode padded ``(B, N)`` float32 audio into one embedding tensor per row.
+
+        ``num_samples_cpu`` holds the lengths on the host, so the window bounds
+        and the output lengths need no device-to-host copy.
+        """
+        mel, mel_lengths = self.preprocessor(audio, num_samples)
+        # NeMo's AudioPreprocessor computes in float32 but returns features in
+        # the module dtype, so the encoders normalize model-dtype features.
+        mel = mel.to(self.encoder.dtype)
+        mel_lengths_cpu = self.preprocessor.featurizer.get_seq_len(
+            num_samples_cpu
+        ).clamp(max=mel.shape[-1])
+        encoded = self.encoder(
+            mel, mel_lengths, int(mel_lengths_cpu.max()), mel_lengths_cpu
+        )
+        encoded_lengths = (
+            mel_lengths_cpu + self.subsampling_factor - 1
+        ) // self.subsampling_factor
+        embeddings = self.proj(encoded.transpose(1, 2))
+        return [
+            emb[:length] for emb, length in zip(embeddings, encoded_lengths.tolist())
+        ]
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        """Load NeMo ``perception.*`` tensors, named without the ``perception.``."""
+        targets = dict(self.named_parameters())
+        targets.update(dict(self.named_buffers()))
+        targets.pop("encoder.diar_kernel", None)
+        loaded: set[str] = set()
+        for name, tensor in weights:
+            if name.endswith("._extra_state") or name.startswith(
+                _UNUSED_CHECKPOINT_PREFIXES
+            ):
+                continue
+            target = targets.get(name)
+            if target is None:
+                raise ValueError(f"Unexpected audio tower weight {name!r}.")
+            weight_loader = getattr(target, "weight_loader", None)
+            if weight_loader is not None:
+                # vLLM's linear layers: quantized parameters and scalar scales.
+                weight_loader(target, tensor)
+            elif target.shape != tensor.shape:
+                raise ValueError(
+                    f"Shape mismatch for {name!r}: {tuple(tensor.shape)} vs "
+                    f"{tuple(target.shape)}."
+                )
+            else:
+                with torch.no_grad():
+                    target.copy_(tensor.to(device=target.device, dtype=target.dtype))
+            loaded.add(name)
+        missing = sorted(set(targets) - loaded)
+        if missing:
+            raise ValueError(
+                f"Audio tower weights missing from the checkpoint ({len(missing)}): "
+                f"{missing[:8]}"
+            )
+        return loaded
diff --git a/vllm/model_executor/models/nemotron_parse3_5.py b/vllm/model_executor/models/nemotron_parse3_5.py
index 1ff0f1d0ac02..bb87d65d2c37 100644
--- a/vllm/model_executor/models/nemotron_parse3_5.py
+++ b/vllm/model_executor/models/nemotron_parse3_5.py
@@ -76,22 +76,22 @@
 
 def _content_normalized_positions(
     patches: torch.Tensor,
-    patch_generator: nn.Module,
+    embeddings: nn.Module,
     image_size: tuple[int, int],
     content_size: tuple[int, int],
 ) -> torch.Tensor:
     """Apply Parse 3.5's content-normalized positional field to one page."""
     height, width = image_size
     content_height, content_width = content_size
-    patch_size = int(patch_generator.patch_size)
+    patch_size = int(embeddings.patch_size)
     rows, cols = height // patch_size, width // patch_size
     if height % patch_size or width % patch_size:
         raise ValueError(f"image size {image_size} is not patch aligned")
     if not (0 < content_height <= height and 0 < content_width <= width):
         raise ValueError(f"content size {content_size} must be within {image_size}")
 
-    positional_grid = patch_generator.pos_embed.reshape(
-        1, patch_generator.num_rows, patch_generator.num_cols, -1
+    positional_grid = embeddings.position_embedding.reshape(
+        1, embeddings.num_rows, embeddings.num_cols, -1
     ).permute(0, 3, 1, 2)
     content_rows = (content_height + patch_size - 1) // patch_size
     content_cols = (content_width + patch_size - 1) // patch_size
@@ -590,12 +590,12 @@ def forward(
             raise ValueError("image_sizes and content_sizes must have equal length")
         if pixel_values.ndim == 2:
             pixel_values = pixel_values.unsqueeze(0)
-        patch_generator = self.model_encoder.model.patch_generator
-        backbone = self.model_encoder.model
-        patch_size = int(patch_generator.patch_size)
-        num_skip = int(patch_generator.num_skip)
-        num_cls = int(patch_generator.num_cls_tokens)
-        patches = patch_generator.embedder(pixel_values)
+        embeddings = self.model_encoder.embeddings
+        backbone = self.model_encoder
+        patch_size = int(embeddings.patch_size)
+        num_skip = int(embeddings.num_skip)
+        num_cls = int(embeddings.num_cls_tokens)
+        patches = embeddings.patch_projection(pixel_values)
         cursor = 0
         positioned_pages = []
         for image_size, content_size in zip(image_sizes, content_sizes):
@@ -607,20 +607,19 @@ def forward(
             cursor += patch_count
             positioned_pages.append(
                 _content_normalized_positions(
-                    page, patch_generator, image_size, content_size
+                    page, embeddings, image_size, content_size
                 )
             )
         if not positioned_pages or cursor != patches.shape[1]:
             raise ValueError("packed patch sequence does not match image_sizes")
-        hidden = patch_generator.cls_token_dynamic(
+        hidden = embeddings._cls_token_dynamic(
             torch.cat(positioned_pages, dim=1), imgs_sizes=image_sizes
         )
-        hidden = patch_generator.patch_normalizer(hidden)
-        mask_meta = None
-        if len(image_sizes) > 1:
-            mask_meta = backbone.inter_image_mask_metadata(
-                image_sizes, device=hidden.device
-            )
+        hidden = embeddings.patch_normalizer(hidden)
+        # One image gives cu_seqlens [0, n], i.e. unmasked attention over one sequence.
+        mask_meta = backbone._inter_image_mask_metadata(
+            image_sizes, device=hidden.device
+        )
         hidden = backbone.encoder(inputs_embeds=hidden, mask_meta=mask_meta)
 
         cursor = 0
@@ -876,11 +875,15 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         synthetic_layerscales = {
             name: param
             for name, param in self.encoder.model_encoder.named_parameters()
-            if name.endswith((".ls1", ".ls2")) and name not in radio_loaded
+            if name.endswith((".layer_scale1.lambda1", ".layer_scale2.lambda1"))
+            and name not in radio_loaded
         }
         if synthetic_layerscales:
+            layerscale_markers = (".ls1", ".ls2", ".layer_scale1", ".layer_scale2")
             source_layerscales = {
-                name for name in radio_source_names if ".ls1" in name or ".ls2" in name
+                name
+                for name in radio_source_names
+                if any(marker in name for marker in layerscale_markers)
             }
             if radio_model_name not in identity_layerscale_models:
                 raise RuntimeError(
diff --git a/vllm/model_executor/models/olmoe.py b/vllm/model_executor/models/olmoe.py
deleted file mode 100644
index 175f25259ea2..000000000000
--- a/vllm/model_executor/models/olmoe.py
+++ /dev/null
@@ -1,405 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-# 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 OLMoE model compatible with HuggingFace weights."""
-
-from collections.abc import Iterable
-from functools import partial
-from itertools import islice
-
-import torch
-from torch import nn
-
-from vllm.compilation.decorators import support_torch_compile
-from vllm.config import VllmConfig
-from vllm.distributed import (
-    get_pp_group,
-    get_tensor_model_parallel_rank,
-    get_tensor_model_parallel_world_size,
-    tensor_model_parallel_all_gather,
-)
-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,
-    GateLinear,
-)
-from vllm.model_executor.layers.layernorm import RMSNorm
-from vllm.model_executor.layers.linear import (
-    QKVParallelLinear,
-    RowParallelLinear,
-)
-from vllm.model_executor.layers.logits_processor import LogitsProcessor
-from vllm.model_executor.layers.quantization import QuantizationConfig
-from vllm.model_executor.layers.rotary_embedding import get_rope
-from vllm.model_executor.layers.vocab_parallel_embedding import (
-    ParallelLMHead,
-    VocabParallelEmbedding,
-)
-from vllm.sequence import IntermediateTensors
-
-from .interfaces import SupportsLoRA, SupportsPP
-from .utils import (
-    AutoWeightsLoader,
-    WeightsMapper,
-    make_empty_intermediate_tensors_factory,
-    make_layers,
-    maybe_prefix,
-)
-
-logger = init_logger(__name__)
-
-
-class OlmoeMoE(nn.Module):
-    """A tensor-parallel MoE implementation for Olmoe that shards each expert
-    across all ranks.
-
-    Each expert's weights are sharded across all ranks and a fused MoE
-    kernel is used for the forward pass, and finally we reduce the outputs
-    across ranks.
-    """
-
-    def __init__(
-        self,
-        num_experts: int,
-        top_k: int,
-        hidden_size: int,
-        intermediate_size: int,
-        params_dtype: torch.dtype | None = None,
-        quant_config: QuantizationConfig | None = None,
-        tp_size: int | None = None,
-        prefix: str = "",
-    ):
-        super().__init__()
-        self.hidden_size = hidden_size
-
-        # Gate always runs at half / full precision for now.
-        self.gate = GateLinear(
-            hidden_size,
-            num_experts,
-            prefix=f"{prefix}.gate",
-        )
-
-        self.experts = FusedMoEFactory(
-            num_experts=num_experts,
-            top_k=top_k,
-            hidden_size=hidden_size,
-            intermediate_size=intermediate_size,
-            renormalize=False,
-            quant_config=quant_config,
-            tp_size=tp_size,
-            prefix=f"{prefix}.experts",
-        )
-
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        # NOTE: hidden_states can have either 1D or 2D shape.
-        orig_shape = hidden_states.shape
-        hidden_dim = hidden_states.shape[-1]
-        hidden_states = hidden_states.view(-1, hidden_dim)
-        # router_logits: (num_tokens, n_experts)
-        router_logits, _ = self.gate(hidden_states)
-        final_hidden_states = self.experts(
-            hidden_states=hidden_states, router_logits=router_logits
-        )
-        return final_hidden_states.view(orig_shape)
-
-
-class OlmoeAttention(nn.Module):
-    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
-        super().__init__()
-
-        config = vllm_config.model_config.hf_config
-        cache_config = vllm_config.cache_config
-        quant_config = vllm_config.quant_config
-
-        self.hidden_size = config.hidden_size
-        max_position_embeddings = getattr(config, "max_position_embeddings", 4096)
-
-        num_heads = config.num_attention_heads
-        num_kv_heads = config.num_key_value_heads
-
-        tp_size = get_tensor_model_parallel_world_size()
-        self.total_num_heads = num_heads
-        assert self.total_num_heads % tp_size == 0
-        self.num_heads = self.total_num_heads // tp_size
-        self.total_num_kv_heads = num_kv_heads
-        if self.total_num_kv_heads >= tp_size:
-            # Number of KV heads is greater than TP size, so we partition
-            # the KV heads across multiple tensor parallel GPUs.
-            assert self.total_num_kv_heads % tp_size == 0
-        else:
-            # Number of KV heads is less than TP size, so we replicate
-            # the KV heads across multiple tensor parallel GPUs.
-            assert tp_size % self.total_num_kv_heads == 0
-        self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
-        self.head_dim = self.hidden_size // self.total_num_heads
-        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.max_position_embeddings = max_position_embeddings
-
-        self.qkv_proj = QKVParallelLinear(
-            self.hidden_size,
-            self.head_dim,
-            self.total_num_heads,
-            self.total_num_kv_heads,
-            bias=False,
-            quant_config=quant_config,
-            prefix=f"{prefix}.qkv_proj",
-        )
-        self.tp_size = tp_size
-        self.tp_rank = get_tensor_model_parallel_rank()
-        self.q_norm = RMSNorm(self.total_num_heads * self.head_dim, eps=1e-5)
-        self.k_norm = RMSNorm(self.total_num_kv_heads * self.head_dim, eps=1e-5)
-        self.o_proj = RowParallelLinear(
-            self.total_num_heads * self.head_dim,
-            self.hidden_size,
-            bias=False,
-            quant_config=quant_config,
-            prefix=f"{prefix}.o_proj",
-        )
-
-        self.rotary_emb = get_rope(
-            self.head_dim,
-            max_position=max_position_embeddings,
-            rope_parameters=config.rope_parameters,
-            is_neox_style=True,
-        )
-        self.attn = Attention(
-            self.num_heads,
-            self.head_dim,
-            self.scaling,
-            num_kv_heads=self.num_kv_heads,
-            cache_config=cache_config,
-            quant_config=quant_config,
-            prefix=f"{prefix}.attn",
-        )
-
-    def _apply_qk_norm(
-        self, q: torch.Tensor, k: torch.Tensor
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        if self.tp_size > 1:
-            q = tensor_model_parallel_all_gather(q.contiguous())
-            k = tensor_model_parallel_all_gather(k.contiguous())
-        q = self.q_norm(q)
-        k = self.k_norm(k)
-        if self.tp_size > 1:
-            splitter = partial(split_tensor_along_last_dim, num_partitions=self.tp_size)
-            q = splitter(q)[self.tp_rank]
-            k = splitter(k)[self.tp_rank]
-        return q, k
-
-    def forward(
-        self,
-        positions: torch.Tensor,
-        hidden_states: torch.Tensor,
-    ) -> torch.Tensor:
-        qkv, _ = self.qkv_proj(hidden_states)
-        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
-        q, k = self._apply_qk_norm(q, k)
-        q, k = self.rotary_emb(positions, q, k)
-        attn_output = self.attn(q, k, v)
-        output, _ = self.o_proj(attn_output)
-        return output
-
-
-class OlmoeDecoderLayer(nn.Module):
-    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
-        super().__init__()
-        config = vllm_config.model_config.hf_config
-        quant_config = vllm_config.quant_config
-
-        self.hidden_size = config.hidden_size
-
-        self.self_attn = OlmoeAttention(
-            vllm_config=vllm_config,
-            prefix=f"{prefix}.self_attn",
-        )
-
-        self.mlp = OlmoeMoE(
-            num_experts=config.num_experts,
-            top_k=config.num_experts_per_tok,
-            hidden_size=config.hidden_size,
-            intermediate_size=config.intermediate_size,
-            quant_config=quant_config,
-            prefix=f"{prefix}.mlp",
-        )
-        self.input_layernorm = RMSNorm(config.hidden_size, eps=1e-5)
-        self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=1e-5)
-
-    def forward(
-        self,
-        positions: torch.Tensor,
-        hidden_states: torch.Tensor,
-        residual: torch.Tensor | None,
-    ) -> torch.Tensor:
-        # Self Attention
-        if residual is None:
-            residual = hidden_states
-            hidden_states = self.input_layernorm(hidden_states)
-        else:
-            hidden_states, residual = self.input_layernorm(hidden_states, residual)
-
-        hidden_states = self.self_attn(
-            positions=positions,
-            hidden_states=hidden_states,
-        )
-
-        # Fully Connected
-        hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
-        hidden_states = self.mlp(hidden_states)
-        return hidden_states, residual
-
-
-@support_torch_compile
-class OlmoeModel(nn.Module):
-    def __init__(
-        self,
-        *,
-        vllm_config: VllmConfig,
-        prefix: str = "",
-        layer_type: type[nn.Module] = OlmoeDecoderLayer,
-    ):
-        super().__init__()
-
-        config = vllm_config.model_config.hf_config
-        quant_config = vllm_config.quant_config
-
-        self.vocab_size = config.vocab_size
-        self.config = config
-        self.embed_tokens = VocabParallelEmbedding(
-            config.vocab_size,
-            config.hidden_size,
-            quant_config=quant_config,
-        )
-        self.start_layer, self.end_layer, self.layers = make_layers(
-            config.num_hidden_layers,
-            lambda prefix: layer_type(vllm_config=vllm_config, prefix=prefix),
-            prefix=f"{prefix}.layers",
-        )
-        self.norm = RMSNorm(config.hidden_size, eps=1e-5)
-
-        self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
-            ["hidden_states", "residual"], config.hidden_size
-        )
-
-    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
-        return self.embed_tokens(input_ids)
-
-    def forward(
-        self,
-        input_ids: torch.Tensor | None,
-        positions: torch.Tensor,
-        intermediate_tensors: IntermediateTensors | 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"]
-            residual = 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:
-            return IntermediateTensors(
-                {"hidden_states": hidden_states, "residual": residual}
-            )
-
-        if residual is not None:
-            hidden_states, _ = self.norm(hidden_states, residual)
-        else:
-            hidden_states = self.norm(hidden_states)
-        return hidden_states
-
-
-class OlmoeForCausalLM(nn.Module, SupportsPP, SupportsLoRA):
-    hf_to_vllm_mapper = WeightsMapper(
-        orig_to_new_stacked={
-            # weight_name: (param_name, shard_id)
-            ".q_proj": (".qkv_proj", "q"),
-            ".k_proj": (".qkv_proj", "k"),
-            ".v_proj": (".qkv_proj", "v"),
-        }
-    )
-    packed_modules_mapping = {
-        "qkv_proj": [
-            "q_proj",
-            "k_proj",
-            "v_proj",
-        ]
-    }
-
-    def __init__(
-        self,
-        *,
-        vllm_config: VllmConfig,
-        prefix: str = "",
-        layer_type: type[nn.Module] = OlmoeDecoderLayer,
-    ):
-        super().__init__()
-        config = vllm_config.model_config.hf_config
-        quant_config = vllm_config.quant_config
-        self.config = config
-        self.quant_config = quant_config
-        self.model = OlmoeModel(
-            vllm_config=vllm_config,
-            prefix=maybe_prefix(prefix, "model"),
-            layer_type=layer_type,
-        )
-        self.lm_head = ParallelLMHead(
-            config.vocab_size,
-            config.hidden_size,
-            quant_config=quant_config,
-            prefix=maybe_prefix(prefix, "lm_head"),
-        )
-        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 | None,
-        positions: torch.Tensor,
-        intermediate_tensors: IntermediateTensors | None = None,
-        inputs_embeds: torch.Tensor | None = None,
-    ) -> torch.Tensor | IntermediateTensors:
-        hidden_states = self.model(
-            input_ids, positions, intermediate_tensors, inputs_embeds
-        )
-        return hidden_states
-
-    def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        logits = self.logits_processor(self.lm_head, hidden_states)
-        return logits
-
-    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)
diff --git a/vllm/model_executor/models/qwen2_moe.py b/vllm/model_executor/models/qwen2_moe.py
index 1c095d6f091c..46320b2c8eae 100644
--- a/vllm/model_executor/models/qwen2_moe.py
+++ b/vllm/model_executor/models/qwen2_moe.py
@@ -44,6 +44,7 @@
     FusedMoEFactory,
     GateLinear,
 )
+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 (
     MergedColumnParallelLinear,
@@ -114,7 +115,7 @@ def __init__(
 
     def forward(self, x):
         gate_up, _ = self.gate_up_proj(x)
-        out = self.act_fn(gate_up)
+        out = maybe_fused_act_quant(self.act_fn, gate_up, self.down_proj)
         out, _ = self.down_proj(out)
 
         if self.expert_gate is not None:
diff --git a/vllm/model_executor/models/qwen3.py b/vllm/model_executor/models/qwen3.py
index 31d1a094ebef..bc6e1a9cd7a8 100644
--- a/vllm/model_executor/models/qwen3.py
+++ b/vllm/model_executor/models/qwen3.py
@@ -343,6 +343,12 @@ 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 load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         loader = AutoWeightsLoader(self)
         return loader.load_weights(weights)
diff --git a/vllm/model_executor/models/qwen3_5.py b/vllm/model_executor/models/qwen3_5.py
index f105948f9283..6e370f180da5 100644
--- a/vllm/model_executor/models/qwen3_5.py
+++ b/vllm/model_executor/models/qwen3_5.py
@@ -42,9 +42,6 @@
     get_pp_group,
 )
 from vllm.logger import init_logger
-from vllm.model_executor.layers.fused_moe.utils import (
-    is_model_fused_shared_expert_compatible,
-)
 from vllm.model_executor.layers.fusion.mm_input_norm import build_mm_input_norm
 from vllm.model_executor.layers.layernorm import GemmaRMSNorm as Qwen3_5RMSNorm
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
@@ -99,6 +96,7 @@
     WeightsMapper,
     _merge_multimodal_embeddings,
     extract_layer_index,
+    get_fused_shared_expert_layers,
     make_empty_intermediate_tensors_factory,
     make_layers,
     maybe_fuse_shared_experts,
@@ -263,10 +261,10 @@ def get_layer(prefix: str):
         self.start_layer, self.end_layer, self.layers = make_layers(
             config.num_hidden_layers, get_layer, prefix=f"{prefix}.layers"
         )
-        self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible(
-            self.layers,
+        self.fused_shared_expert_layers = get_fused_shared_expert_layers(
+            self,
             Qwen3NextSparseMoeBlock,
-            "mlp",
+            enable_eplb=vllm_config.parallel_config.enable_eplb,
         )
         self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
             ["hidden_states", "residual"], config.hidden_size
@@ -286,7 +284,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         if "moe" in self.config.model_type:
             weights = maybe_fuse_shared_experts(
                 weights,
-                enabled=self.is_fused_shared_expert_enabled,
+                fused_shared_expert_layers=self.fused_shared_expert_layers,
                 n_routed_experts=self.config.num_experts,
                 n_shared_experts=1,
                 ckpt_prefix="mlp.shared_expert",
diff --git a/vllm/model_executor/models/qwen3_5_mtp.py b/vllm/model_executor/models/qwen3_5_mtp.py
index 2acac6b5b7c5..b6b2c3851c91 100644
--- a/vllm/model_executor/models/qwen3_5_mtp.py
+++ b/vllm/model_executor/models/qwen3_5_mtp.py
@@ -12,9 +12,6 @@
 from vllm.config import VllmConfig
 from vllm.distributed import get_pp_group, tensor_model_parallel_all_gather
 from vllm.logger import init_logger
-from vllm.model_executor.layers.fused_moe.utils import (
-    is_model_fused_shared_expert_compatible,
-)
 from vllm.model_executor.layers.linear import ColumnParallelLinear
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
 from vllm.model_executor.layers.vocab_parallel_embedding import (
@@ -31,6 +28,9 @@
     Qwen3NextSparseMoeBlock,
     QwenNextMixtureOfExperts,
 )
+from vllm.model_executor.models.qwen3_next_mtp import (
+    should_torch_compile_mtp_drafter,
+)
 from vllm.model_executor.models.utils import sequence_parallel_chunk
 from vllm.sequence import IntermediateTensors
 
@@ -44,6 +44,7 @@
     AutoWeightsLoader,
     PPMissingLayer,
     _merge_multimodal_embeddings,
+    get_fused_shared_expert_layers,
     make_empty_intermediate_tensors_factory,
     maybe_fuse_shared_experts,
     maybe_prefix,
@@ -61,7 +62,8 @@
         "intermediate_tensors": 0,
         "inputs_embeds": 0,
         "hidden_states": 0,
-    }
+    },
+    enable_if=should_torch_compile_mtp_drafter,
 )
 class Qwen3_5MultiTokenPredictor(nn.Module):
     hf_to_vllm_mapper = Qwen3_5Model.hf_to_vllm_mapper
@@ -124,10 +126,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
             for idx in range(self.num_mtp_layers)
         )
         vllm_config.quant_config = original_quant
-        self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible(
-            self.layers,
+        self.fused_shared_expert_layers = get_fused_shared_expert_layers(
+            self,
             Qwen3NextSparseMoeBlock,
-            "mlp",
+            enable_eplb=vllm_config.parallel_config.enable_eplb,
         )
         self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
             ["hidden_states", "residual"], config.hidden_size
@@ -193,7 +195,7 @@ def forward(
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         weights = maybe_fuse_shared_experts(
             weights,
-            enabled=self.is_fused_shared_expert_enabled,
+            fused_shared_expert_layers=self.fused_shared_expert_layers,
             n_routed_experts=getattr(self.config, "num_experts", 0),
             n_shared_experts=1,
             ckpt_prefix="mlp.shared_expert",
@@ -211,7 +213,8 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         "intermediate_tensors": 0,
         "inputs_embeds": 0,
         "hidden_states": 0,
-    }
+    },
+    enable_if=should_torch_compile_mtp_drafter,
 )
 class Qwen3_5MTP(LocalArgmaxMixin, nn.Module, SupportsMultiModal, SupportsPP):
     packed_modules_mapping = {
@@ -285,10 +288,15 @@ def forward(  # type: ignore[override]
         hidden_states: torch.Tensor,
         intermediate_tensors: IntermediateTensors | None = None,
         inputs_embeds: torch.Tensor | None = None,
-        **kwargs: object,
+        spec_step_idx: int = 0,
     ):
         hidden_states = self.model(
-            input_ids, positions, hidden_states, intermediate_tensors, inputs_embeds
+            input_ids,
+            positions,
+            hidden_states,
+            intermediate_tensors,
+            inputs_embeds,
+            spec_step_idx,
         )
         return hidden_states
 
@@ -299,6 +307,14 @@ def compute_logits(
     ) -> torch.Tensor | None:
         return self.logits_processor(self.lm_head, hidden_states)
 
+    def get_top_tokens(
+        self,
+        hidden_states: torch.Tensor,
+        spec_step_idx: int = 0,
+    ) -> torch.Tensor:
+        # One lm_head is shared by all MTP modules, as in compute_logits.
+        return super().get_top_tokens(hidden_states)
+
     def is_unused_checkpoint_weight(self, name: str) -> bool:
         return not name.startswith("mtp.") and not any(
             key in name for key in ["embed_tokens", "lm_head"]
diff --git a/vllm/model_executor/models/qwen3_dflash.py b/vllm/model_executor/models/qwen3_dflash.py
index b7959cfc856b..f8cb63bf62bb 100644
--- a/vllm/model_executor/models/qwen3_dflash.py
+++ b/vllm/model_executor/models/qwen3_dflash.py
@@ -27,17 +27,22 @@
 )
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
 from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
+from vllm.model_executor.layers.quantization.modelopt import ModelOptMxFp8Config
 from vllm.model_executor.layers.rotary_embedding import get_rope
 from vllm.model_executor.layers.vocab_parallel_embedding import (
     ParallelLMHead,
     VocabParallelEmbedding,
 )
 from vllm.model_executor.model_loader.weight_utils import make_attention_sink
+from vllm.model_executor.utils import (
+    weights_already_processed,
+)
 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
+from vllm.v1.worker.gpu.spec_decode.dflash.config import DFlashSpeculatorConfig
 from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import (
     get_eagle3_aux_layers_from_config,
 )
@@ -59,6 +64,25 @@
 _SLIDING_ATTENTION = "sliding_attention"
 
 
+def _is_mxfp8_linear(layer: nn.Module) -> bool:
+    return isinstance(
+        getattr(layer, "quant_config", None), ModelOptMxFp8Config
+    ) and not isinstance(layer.quant_method, UnquantizedLinearMethod)
+
+
+_QKV_PARAM = re.compile(
+    r"^layers\.(\d+)\.self_attn\.qkv_proj\.(weight|weight_scale|bias)$"
+)
+
+
+def _load_context_kv_rows(
+    param: torch.Tensor, loaded_weight: torch.Tensor, row_offset: int
+) -> None:
+    """Weight loader of the MXFP8 context K/V projection: one layer's K or V
+    rows, at their offset in the stacked projection."""
+    param.data.narrow(0, row_offset, loaded_weight.shape[0]).copy_(loaded_weight)
+
+
 def _dflash_layer_causal(config: Qwen3Config, layer_idx: int) -> bool:
     """Resolve explicit causality before falling back to legacy layer defaults."""
     is_causal = getattr(config, "is_causal", None)
@@ -81,6 +105,13 @@ def dflash_has_any_non_causal(config: Qwen3Config) -> bool:
     )
 
 
+def _reads_raw_weights(projs: list[nn.Module]) -> bool:
+    """Whether post-load `.weight` is plain (not quantized, not CPU-repacked)."""
+    return not current_platform.is_cpu() and all(
+        isinstance(p.quant_method, UnquantizedLinearMethod) for p in projs
+    )
+
+
 def _get_dflash_fc_input_size(vllm_config: VllmConfig) -> int:
     spec_config = vllm_config.speculative_config
     assert spec_config is not None
@@ -182,6 +213,7 @@ def __init__(
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
         attn_type: str = AttentionType.DECODER,
+        query_pre_quantized: bool = False,
     ) -> None:
         super().__init__()
         self.layer_name = prefix
@@ -242,6 +274,7 @@ def __init__(
             prefix=f"{prefix}.attn",
             attn_type=attn_type,
             sinks=self.attention_sink_bias,
+            query_pre_quantized=query_pre_quantized,
         )
         self.causal = causal
         if current_platform.is_rocm():
@@ -408,6 +441,9 @@ class DFlashQwen3Model(nn.Module):
             ".up_proj": (".gate_up_proj", 1),
         },
     )
+    # Maps module paths in the draft checkpoint's quantization exclusions to
+    # this model's module paths.
+    quant_config_mapper: WeightsMapper | None = None
 
     def __init__(
         self,
@@ -422,6 +458,8 @@ def __init__(
         self.config = speculative_config.draft_model_config.hf_config
         self.vocab_size = self.config.vocab_size
         self.quant_config = get_draft_quant_config(vllm_config)
+        if self.quant_config is not None and self.quant_config_mapper is not None:
+            self.quant_config.apply_vllm_mapper(self.quant_config_mapper)
         map_draft_quant_config(
             self.quant_config,
             self.hf_to_vllm_mapper,
@@ -493,6 +531,10 @@ def __init__(
             self.config.hidden_size,
             eps=self.config.rms_norm_eps,
         )
+        layers_attn = [getattr(layer, "self_attn", None) for layer in self.layers]
+        self._context_kv_proj: ReplicatedLinear | None = None
+        if all(_is_mxfp8_linear(getattr(a, "qkv_proj", None)) for a in layers_attn):
+            self._context_kv_proj = self._make_mxfp8_context_kv_proj(layers_attn)
 
     def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
         embeds = self.embed_tokens(input_ids)
@@ -511,10 +553,7 @@ def _build_context_kv_buffers(
         self._context_qkv_projs = [a.qkv_proj for a in layers_attn]
         self._context_q_sizes = [a.q_size for a in layers_attn]
 
-        if all(
-            isinstance(proj.quant_method, UnquantizedLinearMethod)
-            for proj in self._context_qkv_projs
-        ):
+        if _reads_raw_weights(self._context_qkv_projs):
             # KV projection weights: [num_layers * 2 * kv_size, hidden_size]
             kv_weights = [a.qkv_proj.weight[a.q_size :] for a in layers_attn]
             self._fused_kv_weight: torch.Tensor | None = torch.cat(kv_weights, dim=0)
@@ -525,7 +564,8 @@ def _build_context_kv_buffers(
                 self._fused_kv_bias = None
         else:
             # Quantized linear weights may use packed storage that cannot be
-            # consumed by F.linear. Run each projection through its quant method.
+            # consumed by F.linear. Run each projection through its quant method
+            # (or the MXFP8 context projection, when there is one).
             self._fused_kv_weight = None
             self._fused_kv_bias = None
 
@@ -535,6 +575,51 @@ def _build_context_kv_buffers(
             [a.k_norm.weight.data for a in layers_attn], dim=0
         ).contiguous()
 
+    def _make_mxfp8_context_kv_proj(
+        self, layers_attn: list[nn.Module]
+    ) -> ReplicatedLinear:
+        """One MXFP8 GEMM over the K/V rows of every layer's QKV projection.
+
+        MXFP8 block scales are per output row, so it holds the K/V rows of the
+        checkpoint's weights and scales as-is. They are loaded from the
+        checkpoint's K/V tensors (`load_weights`), so the projection goes
+        through post-load processing and weight reloads like any other layer.
+        """
+        first = layers_attn[0].qkv_proj
+        proj = ReplicatedLinear(
+            first.input_size,
+            sum(2 * a.kv_size for a in layers_attn),
+            bias=first.bias is not None,
+            params_dtype=first.params_dtype,
+            quant_config=first.quant_config,
+            prefix=first.prefix,
+            return_bias=False,
+            disable_tp=True,  # holds this rank's K/V rows
+        )
+        for param in proj.parameters():
+            param.weight_loader = _load_context_kv_rows
+        offsets = [0]
+        for a in layers_attn:
+            offsets.append(offsets[-1] + 2 * a.kv_size)
+        self._context_kv_offsets = offsets
+        return proj
+
+    def _load_context_kv(
+        self, layer_idx: int, shard_id: str, param_name: str, weight: torch.Tensor
+    ) -> None:
+        """Load this rank's rows of a layer's K or V checkpoint tensor into the
+        MXFP8 context projection, sharded as the layer's QKV projection does."""
+        assert self._context_kv_proj is not None
+        attn = self.layers[layer_idx].self_attn
+        qkv = attn.qkv_proj
+        rank = qkv.tp_rank // qkv.num_kv_head_replicas
+        rows = weight.narrow(0, rank * attn.kv_size, attn.kv_size)
+        offset = self._context_kv_offsets[layer_idx]
+        if shard_id == "v":
+            offset += attn.kv_size
+        param = getattr(self._context_kv_proj, param_name)
+        param.weight_loader(param, rows, offset)
+
     def _build_fused_kv_buffers(self) -> None:
         """Build fused weight buffers for precompute_and_store_context_kv.
 
@@ -578,15 +663,7 @@ def _build_fused_kv_buffers(self) -> None:
         # References to inner Attention layers for direct cache writes
         self._attn_layers = [layer.self_attn.attn for layer in self.layers]
 
-    def _project_context_kv(
-        self,
-        context_states: torch.Tensor,
-        num_ctx: int,
-        num_layers: int,
-        num_kv_heads: int,
-        head_dim: int,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        # --- Fused KV projection (one GEMM for all layers) ---
+    def _normalize_context_states(self, context_states: torch.Tensor) -> torch.Tensor:
         normed_context_states = torch.empty_like(context_states)
         ops.rms_norm(
             normed_context_states,
@@ -594,11 +671,28 @@ def _project_context_kv(
             self._hidden_norm_weight,
             self._rms_norm_eps,
         )
+        return normed_context_states
+
+    def _project_context_kv_token_major(
+        self,
+        context_states: torch.Tensor,
+        num_ctx: int,
+        num_layers: int,
+        num_kv_heads: int,
+        head_dim: int,
+    ) -> torch.Tensor:
+        """Context K/V of every layer, ``[num_ctx, L, 2, nkv, hd]``."""
+        # --- Fused KV projection (one GEMM for all layers) ---
+        normed_context_states = self._normalize_context_states(context_states)
         if self._fused_kv_weight is not None:
             all_kv_flat = F.linear(
                 normed_context_states, self._fused_kv_weight, self._fused_kv_bias
             )
             all_kv = all_kv_flat.view(num_ctx, num_layers, 2, num_kv_heads, head_dim)
+        elif self._context_kv_proj is not None:
+            all_kv = self._context_kv_proj(normed_context_states).view(
+                num_ctx, num_layers, 2, num_kv_heads, head_dim
+            )
         else:
             layer_kv = []
             for proj, q_size in zip(
@@ -609,7 +703,19 @@ def _project_context_kv(
                     qkv[:, q_size:].view(num_ctx, 2, num_kv_heads, head_dim)
                 )
             all_kv = torch.stack(layer_kv, dim=1)
+        return all_kv
 
+    def _project_context_kv(
+        self,
+        context_states: torch.Tensor,
+        num_ctx: int,
+        num_layers: int,
+        num_kv_heads: int,
+        head_dim: int,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        all_kv = self._project_context_kv_token_major(
+            context_states, num_ctx, num_layers, num_kv_heads, head_dim
+        )
         # Single contiguous copy that separates K/V and transposes to
         # layer-major layout. Result: [2, L, num_ctx, nkv, hd] contiguous.
         # Indexing dim-0 gives contiguous [L, num_ctx, nkv, hd] for K and V.
@@ -653,7 +759,8 @@ def precompute_and_store_context_kv(
                 "DFlash buffer initialization was skipped. If dummy weights are not "
                 "in use, this may indicate an error in weight loading."
             )
-            self._build_fused_kv_buffers()
+            with weights_already_processed():
+                self._build_fused_kv_buffers()
 
         num_ctx = context_states.shape[0]
         L = self._num_attn_layers
@@ -729,11 +836,35 @@ def forward(
         return hidden_states
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        proj = self._context_kv_proj
+        if proj is not None:
+            weights = self._tee_context_kv(weights)
         loader = AutoWeightsLoader(self)
-        return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
+        loaded = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
+        if proj is not None:
+            loaded |= {
+                f"_context_kv_proj.{name}" for name, _ in proj.named_parameters()
+            }
+        return loaded
+
+    def _tee_context_kv(
+        self, weights: Iterable[tuple[str, torch.Tensor]]
+    ) -> Iterable[tuple[str, torch.Tensor]]:
+        for name, weight in weights:
+            mapped = self.hf_to_vllm_mapper._map_name_with_shard(name)
+            if mapped is not None and mapped[1] in ("k", "v"):
+                match = _QKV_PARAM.match(mapped[0])
+                if match is not None:
+                    self._load_context_kv(
+                        int(match[1]), str(mapped[1]), match[2], weight
+                    )
+            yield name, weight
 
 
 class DFlashQwen3ForCausalLM(Qwen3ForCausalLM):
+    # The speculator this is a draft model of (see SpeculativeDraftModel).
+    speculator_config_class = DFlashSpeculatorConfig
+
     model_cls = DFlashQwen3Model
 
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
@@ -890,6 +1021,8 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
         mapper = WeightsMapper(orig_to_new_substr=orig_to_new_substr)
         loader = AutoWeightsLoader(self)
         loader.load_weights(model_weights.items(), mapper=mapper)
+
+    def process_weights_after_loading(self) -> None:
         self.model._build_fused_kv_buffers()
 
     def _read_mask_embedding(self) -> torch.Tensor | None:
diff --git a/vllm/model_executor/models/qwen3_dflash2.py b/vllm/model_executor/models/qwen3_dflash2.py
index 861a7c5412ea..80f8de794a8f 100644
--- a/vllm/model_executor/models/qwen3_dflash2.py
+++ b/vllm/model_executor/models/qwen3_dflash2.py
@@ -13,6 +13,7 @@
 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 vllm.v1.worker.gpu.spec_decode.dflash2.config import DFlash2SpeculatorConfig
 
 from .qwen3_dflash import (
     DFlashQwen3DecoderLayer,
@@ -296,6 +297,10 @@ def _score_edges(
     anchor_token_ids: torch.Tensor,
     top_k: int,
 ) -> torch.Tensor:
+    # A request still in chunked prefill anchors on its next prompt token, which
+    # can lie past the draft vocabulary (an audio placeholder, say). Its draft is
+    # discarded, so only keep the codebook lookup in range.
+    anchor_token_ids = anchor_token_ids.clamp(0, predecessor_table.shape[0] - 1)
     successors = successor_table[candidate_ids]
     predecessor_ids = torch.cat(
         (
@@ -393,14 +398,19 @@ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
 
 
 class DFlash2Qwen3ForCausalLM(DFlashQwen3ForCausalLM):
+    # The speculator this is a draft model of (see SpeculativeDraftModel).
+    speculator_config_class = DFlash2SpeculatorConfig
+
     model_cls = DFlash2Qwen3Model
 
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
         super().__init__(vllm_config=vllm_config, prefix=prefix)
         draft_config = self.config.dflash_config
         softcap = float(draft_config.get("final_logit_softcapping") or 0.0)
+        # Candidates index the selector codebooks, which have the draft's
+        # vocabulary; the target's vocabulary can be larger.
         self.candidate_logits_processor = LogitsProcessor(
-            vllm_config.model_config.get_vocab_size(),
+            min(vllm_config.model_config.get_vocab_size(), self.config.vocab_size),
             scale=float(draft_config.get("output_multiplier", 1.0)),
             soft_cap=softcap if softcap > 0 else None,
         )
diff --git a/vllm/model_executor/models/qwen3_dspark.py b/vllm/model_executor/models/qwen3_dspark.py
index 156f55c373ad..6f902170546e 100644
--- a/vllm/model_executor/models/qwen3_dspark.py
+++ b/vllm/model_executor/models/qwen3_dspark.py
@@ -32,7 +32,11 @@
 from vllm.model_executor.layers.vocab_parallel_embedding import (
     ParallelLMHead,
 )
+from vllm.model_executor.model_loader.mtp_validation import (
+    is_mtp_completeness_check_enabled,
+)
 from vllm.platforms import current_platform
+from vllm.v1.worker.gpu.spec_decode.dspark.config import DSparkSpeculatorConfig
 
 from .qwen3_dflash import DFlashQwen3ForCausalLM, DFlashQwen3Model
 from .utils import (
@@ -231,6 +235,10 @@ def __init__(
 
 
 class Qwen3DSparkForCausalLM(DFlashQwen3ForCausalLM):
+    # The speculator this is a draft model of (see SpeculativeDraftModel).
+    speculator_config_class = DSparkSpeculatorConfig
+    model_cls = Qwen3DSparkModel
+
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
         nn.Module.__init__(self)
         speculative_config = vllm_config.speculative_config
@@ -240,7 +248,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
         if getattr(self.config, "draft_vocab_size", None) is None:
             self.config.draft_vocab_size = getattr(self.config, "vocab_size", None)
         target_layer_num = vllm_config.model_config.get_total_num_hidden_layers()
-        self.model = Qwen3DSparkModel(
+        self.model = self.model_cls(
             vllm_config=vllm_config,
             prefix=maybe_prefix(prefix, "model"),
             start_layer_id=target_layer_num,
@@ -275,6 +283,10 @@ def compute_draft_logits(self, hidden_states: torch.Tensor) -> torch.Tensor:
         # Markov bias in draft space, then remaps via map_draft_to_target.
         return self.logits_processor(self.lm_head, hidden_states)
 
+    def compute_draft_logits_local(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        """This TP rank's vocab shard of compute_draft_logits, not gathered."""
+        return self.logits_processor(self.lm_head, hidden_states, skip_gather=True)
+
     def map_draft_to_target(self, draft_ids: torch.Tensor) -> torch.Tensor:
         # Map draft-vocab ids to target ids (identity for full-vocab drafts).
         if self.draft_id_to_target_id is None:
@@ -340,17 +352,18 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
         # 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
-        if uses_expanded_input_vocab and not includes_embed_tokens:
+        complete = is_mtp_completeness_check_enabled()
+        if complete and uses_expanded_input_vocab and not includes_embed_tokens:
             raise ValueError(
                 "Qwen3 DSpark checkpoints whose input vocab_size is larger than "
                 "the target vocabulary must include embed_tokens weights."
             )
-        if uses_reduced_vocab and not includes_lm_head:
+        if complete and uses_reduced_vocab and not includes_lm_head:
             raise ValueError(
                 "Reduced-vocabulary Qwen3 DSpark checkpoints must include "
                 "lm_head weights; the full target lm_head cannot be shared."
             )
-        if uses_reduced_vocab and not includes_draft_id_mapping:
+        if complete and uses_reduced_vocab and not includes_draft_id_mapping:
             raise ValueError(
                 "Reduced-vocabulary Qwen3 DSpark checkpoints must include a "
                 "d2t mapping so sampled draft ids can be converted to target ids."
@@ -363,10 +376,10 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
             orig_to_new_substr["lm_head"] = None
         if not includes_draft_id_mapping:
             orig_to_new_substr["draft_id_to_target_id"] = None
-        if self.model.confidence_head is None or not includes_confidence_head:
+        if complete and not includes_confidence_head:
             self.model.confidence_head = None
+        if self.model.confidence_head is None:
             orig_to_new_substr["confidence_head"] = None
         mapper = WeightsMapper(orig_to_new_substr=orig_to_new_substr)
         loader = AutoWeightsLoader(self)
         loader.load_weights(model_weights.items(), mapper=mapper)
-        self.model._build_fused_kv_buffers()
diff --git a/vllm/model_executor/models/qwen3_next.py b/vllm/model_executor/models/qwen3_next.py
index ee5c05d0a4ae..de566b7b06a4 100644
--- a/vllm/model_executor/models/qwen3_next.py
+++ b/vllm/model_executor/models/qwen3_next.py
@@ -18,16 +18,20 @@
     tensor_model_parallel_all_gather,
     tensor_model_parallel_reduce_scatter,
 )
+from vllm.kernels.triton.fused_sigmoid_gate_mul import fused_sigmoid_gate_mul
 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,
     resolve_layer_fused_shared_expert,
 )
 from vllm.model_executor.layers.fused_qk_norm_rope import fused_qk_rmsnorm_rope_gate
+from vllm.model_executor.layers.fusion.quant_activation import (
+    QuantizedActivation,
+    get_input_quant_key,
+)
 from vllm.model_executor.layers.layernorm import (
     GemmaRMSNorm as Qwen3NextRMSNorm,
 )
@@ -50,6 +54,9 @@
 from vllm.model_executor.layers.quantization.utils.config_utils import (
     get_quark_ocp_mx_group_size,
 )
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    kFp8StaticTensorSym,
+)
 from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding, get_rope
 from vllm.model_executor.layers.vocab_parallel_embedding import (
     ParallelLMHead,
@@ -75,6 +82,7 @@
     PPMissingLayer,
     WeightsMapper,
     extract_layer_index,
+    get_fused_shared_expert_layers,
     make_empty_intermediate_tensors_factory,
     make_layers,
     maybe_fuse_shared_experts,
@@ -393,6 +401,10 @@ def __init__(
             and (text_only or supports_mrope)
         )
 
+        # Off by default; Qwen3.5 enables the fused sigmoid-gate-mul Triton
+        # kernel post-construction (see Qwen3_5DecoderLayer).
+        self.use_fused_output_gate = False
+
     def _project_qkv_gate(
         self,
         qkv: torch.Tensor,
@@ -458,13 +470,26 @@ def _project_qkv_gate(
     def forward(
         self,
         positions: torch.Tensor,
-        hidden_states: torch.Tensor,
+        hidden_states: torch.Tensor | QuantizedActivation,
     ) -> torch.Tensor:
         qkv, _ = self.qkv_proj(hidden_states)
         q, k, v, gate = self._project_qkv_gate(qkv, positions)
         attn_output = self.attn(q, k, v)
         if gate is not None:
-            attn_output = attn_output * torch.sigmoid(gate)
+            if not self.use_fused_output_gate:
+                attn_output = attn_output * torch.sigmoid(gate)
+            elif get_input_quant_key(self.o_proj) == kFp8StaticTensorSym:
+                # Sigmoid gate, fused with o_proj's static FP8 input quant.
+                scale = self.o_proj.input_scale
+                attn_output = QuantizedActivation(
+                    data=fused_sigmoid_gate_mul(attn_output, gate, scale),
+                    scale=scale,
+                    orig_dtype=attn_output.dtype,
+                    orig_shape=attn_output.shape,
+                    quant_key=kFp8StaticTensorSym,
+                )
+            else:
+                attn_output = fused_sigmoid_gate_mul(attn_output, gate)
         output, _ = self.o_proj(attn_output)
         return output
 
@@ -672,10 +697,10 @@ def get_layer(prefix: str):
         self.start_layer, self.end_layer, self.layers = make_layers(
             config.num_hidden_layers, get_layer, prefix=f"{prefix}.layers"
         )
-        self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible(
-            self.layers,
+        self.fused_shared_expert_layers = get_fused_shared_expert_layers(
+            self,
             Qwen3NextSparseMoeBlock,
-            "mlp",
+            enable_eplb=parallel_config.enable_eplb,
         )
         self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
             ["hidden_states", "residual"], config.hidden_size
@@ -756,7 +781,7 @@ def forward(
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         weights = maybe_fuse_shared_experts(
             weights,
-            enabled=self.is_fused_shared_expert_enabled,
+            fused_shared_expert_layers=self.fused_shared_expert_layers,
             n_routed_experts=getattr(self.config, "num_experts", 0),
             n_shared_experts=1,
             ckpt_prefix="mlp.shared_expert",
diff --git a/vllm/model_executor/models/qwen3_next_mtp.py b/vllm/model_executor/models/qwen3_next_mtp.py
index 4ae8f21840a1..9180dc876fce 100644
--- a/vllm/model_executor/models/qwen3_next_mtp.py
+++ b/vllm/model_executor/models/qwen3_next_mtp.py
@@ -12,9 +12,6 @@
 from vllm.config import VllmConfig
 from vllm.distributed import get_pp_group, tensor_model_parallel_all_gather
 from vllm.logger import init_logger
-from vllm.model_executor.layers.fused_moe.utils import (
-    is_model_fused_shared_expert_compatible,
-)
 from vllm.model_executor.layers.linear import ColumnParallelLinear
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
 from vllm.model_executor.layers.vocab_parallel_embedding import (
@@ -33,6 +30,7 @@
 
 from .utils import (
     AutoWeightsLoader,
+    get_fused_shared_expert_layers,
     make_empty_intermediate_tensors_factory,
     maybe_fuse_shared_experts,
     maybe_prefix,
@@ -43,7 +41,17 @@
 KVCache = tuple[torch.Tensor, torch.Tensor]
 
 
-@support_torch_compile
+def should_torch_compile_mtp_drafter(vllm_config: VllmConfig) -> bool:
+    """`enable_if` for MTP drafters that pick their module by `spec_step_idx`.
+
+    The drafter is compiled once with guards dropped, so a compiled
+    multi-module drafter would reuse the first-traced module for every step.
+    """
+    spec_config = vllm_config.speculative_config
+    return spec_config is None or not spec_config.use_multi_module_mtp()
+
+
+@support_torch_compile(enable_if=should_torch_compile_mtp_drafter)
 class Qwen3NextMultiTokenPredictor(nn.Module):
     hf_to_vllm_mapper = Qwen3NextModel.hf_to_vllm_mapper
 
@@ -95,10 +103,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
             for idx in range(self.num_mtp_layers)
         )
 
-        self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible(
-            self.layers,
+        self.fused_shared_expert_layers = get_fused_shared_expert_layers(
+            self,
             Qwen3NextSparseMoeBlock,
-            "mlp",
+            enable_eplb=vllm_config.parallel_config.enable_eplb,
         )
         self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
             ["hidden_states", "residual"], config.hidden_size
@@ -164,7 +172,7 @@ def forward(
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         weights = maybe_fuse_shared_experts(
             weights,
-            enabled=self.is_fused_shared_expert_enabled,
+            fused_shared_expert_layers=self.fused_shared_expert_layers,
             n_routed_experts=self.config.num_experts,
             n_shared_experts=1,
             ckpt_prefix="mlp.shared_expert",
@@ -173,7 +181,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
 
 
-@support_torch_compile
+@support_torch_compile(enable_if=should_torch_compile_mtp_drafter)
 class Qwen3NextMTP(nn.Module, QwenNextMixtureOfExperts):
     packed_modules_mapping = {
         "qkv_proj": [
@@ -213,10 +221,15 @@ def forward(
         hidden_states: torch.Tensor,
         intermediate_tensors: IntermediateTensors | None = None,
         inputs_embeds: torch.Tensor | None = None,
-        **kwargs: object,
+        spec_step_idx: int = 0,
     ):
         hidden_states = self.model(
-            input_ids, positions, hidden_states, intermediate_tensors, inputs_embeds
+            input_ids,
+            positions,
+            hidden_states,
+            intermediate_tensors,
+            inputs_embeds,
+            spec_step_idx,
         )
         return hidden_states
 
diff --git a/vllm/model_executor/models/qwen3_vl.py b/vllm/model_executor/models/qwen3_vl.py
index 527e2b0f3e2e..5779e44773eb 100644
--- a/vllm/model_executor/models/qwen3_vl.py
+++ b/vllm/model_executor/models/qwen3_vl.py
@@ -120,6 +120,7 @@
 from vllm.utils.collection_utils import is_list_of
 from vllm.utils.math_utils import cdiv, round_up
 from vllm.utils.torch_utils import PIN_MEMORY
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
 from vllm.v1.worker.encoder_cudagraph_defs import EncoderCudaGraphReplayBuffers
 
 from ...utils.gpu_sync_debug import gpu_sync_allowed
@@ -144,6 +145,8 @@
     Qwen2_5_VLVideoEmbeddingInputs,
     Qwen2_5_VLVideoInputs,
     Qwen2_5_VLVideoPixelInputs,
+    _pad_cumulative_seqlens_buffer,
+    _pad_flashinfer_cu_seqlens_buffer,
 )
 from .qwen2_vl import (
     Qwen2VLMultiModalDataParser,
@@ -2146,6 +2149,19 @@ def get_encoder_cudagraph_config(self):
                 "max_seqlen",
                 "sequence_lengths",
             ],
+            # FlashInfer packs the q/k/o and v offsets into one cu_seqlens
+            # tensor, split at its midpoint. Replays usually pad to a smaller
+            # batch bucket than capture, so each half has to be copied into
+            # its own half of the capture buffer. Other backends repeat the
+            # last offset so padded sequences have zero length, as in
+            # Qwen2.5-VL.
+            padding_logics={
+                "cu_seqlens": (
+                    _pad_flashinfer_cu_seqlens_buffer
+                    if self.visual.attn_backend == AttentionBackendEnum.FLASHINFER
+                    else _pad_cumulative_seqlens_buffer
+                )
+            },
             out_hidden_size=self.visual.out_hidden_size,
             max_frames_per_video=max_frames,
         )
diff --git a/vllm/model_executor/models/qwen3_xpress.py b/vllm/model_executor/models/qwen3_xpress.py
index b410dc17ee1f..7ad671199da7 100644
--- a/vllm/model_executor/models/qwen3_xpress.py
+++ b/vllm/model_executor/models/qwen3_xpress.py
@@ -12,6 +12,7 @@
 from vllm.model_executor.layers.logits_processor import LogitsProcessor, _topk
 from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
 from vllm.v1.worker.gpu.spec_decode.xpress import kernels
+from vllm.v1.worker.gpu.spec_decode.xpress.config import XPressSpeculatorConfig
 
 from .qwen3_dflash import DFlashQwen3ForCausalLM, DFlashQwen3Model
 from .utils import AutoWeightsLoader, maybe_prefix, process_eagle_weight
@@ -271,6 +272,9 @@ def __init__(
 
 
 class Qwen3XPressForCausalLM(DFlashQwen3ForCausalLM):
+    # The speculator this is a draft model of (see SpeculativeDraftModel).
+    speculator_config_class = XPressSpeculatorConfig
+
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
         nn.Module.__init__(self)
         assert vllm_config.speculative_config is not None
diff --git a/vllm/model_executor/models/radio.py b/vllm/model_executor/models/radio.py
index d01e0f846418..1d5722756abe 100644
--- a/vllm/model_executor/models/radio.py
+++ b/vllm/model_executor/models/radio.py
@@ -1,16 +1,32 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-# Copyright (c) 2023-2024, NVIDIA CORPORATION.  All rights reserved.
-#
-# NVIDIA CORPORATION and its licensors retain all intellectual property
-# and proprietary rights in and to this software, related documentation
-# and any modifications thereto.  Any use, reproduction, disclosure or
-# distribution of this software and related documentation without an express
-# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+# adapted from https://github.com/NVlabs/RADIO and
+# https://github.com/huggingface/transformers/blob/main/src/transformers/models/radio/modeling_radio.py
+# --------------------------------------------------------
+# RADIO
+# Copyright (c) 2023-2026, NVIDIA CORPORATION.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# --------------------------------------------------------
+"""RADIO vision encoder, matching the ``RadioModel`` in the Transformers library.
+
+The module tree mirrors the Transformers ``RadioModel`` (``embeddings`` and an
+``encoder.layer`` stack of ``norm1``/``attention``/``layer_scale1``/``norm2``/
+``mlp``/``layer_scale2`` blocks). The vLLM-specific deviations are the fused
+``attention.qkv`` (a :class:`~vllm.model_executor.layers.linear.QKVParallelLinear`
+for TP/quantization, vs the split ``query``/``key``/``value`` in HF) and running
+input normalization in the processor instead of an ``input_conditioner`` module;
+the attention backend and MLP reuse vLLM's :class:`MMEncoderAttention` and
+:class:`~vllm.model_executor.models.intern_vit.InternMLP`. ``load_weights`` maps
+both native Transformers and legacy remote-code checkpoint keys onto this tree
+via :data:`hf_to_vllm_mapper`.
+"""
 
 import math
+from collections.abc import Iterable
 from dataclasses import dataclass
-from itertools import accumulate
+from functools import partial
+from itertools import accumulate, repeat
 from typing import TypeAlias
 
 import torch
@@ -19,185 +35,134 @@
 from einops import rearrange
 from transformers import PreTrainedConfig
 
-from vllm.model_executor.layers.quantization import QuantizationConfig
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
-from vllm.model_executor.models.intern_vit import (
-    InternParallelAttention,
-    InternVisionEncoder,
-    InternVisionEncoderLayer,
+from vllm.compilation.decorators import (
+    should_torch_compile_mm_encoder,
+    support_torch_compile,
 )
+from vllm.distributed import (
+    divide,
+    get_tensor_model_parallel_rank,
+    get_tensor_model_parallel_world_size,
+    split_tensor_along_last_dim,
+    tensor_model_parallel_all_gather,
+)
+from vllm.logger import init_logger
+from vllm.model_executor.layers.attention import MMEncoderAttention
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.layers.linear import QKVParallelLinear, RowParallelLinear
+from vllm.model_executor.layers.quantization import QuantizationConfig
+from vllm.model_executor.models.intern_vit import InternMLP
+from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper
+from vllm.model_executor.models.vision import is_vit_use_data_parallel
+from vllm.utils.torch_utils import async_tensor_h2d
 
-input_dim_t: TypeAlias = int | tuple[int, int]
-
-
-def calc_seq_len(size: tuple[int, int], patch_size: int) -> int:
-    h, w = size
-    return (h // patch_size) * (w // patch_size)
+logger = init_logger(__name__)
 
+input_dim_t: TypeAlias = int | tuple[int, int]
+norm_t: TypeAlias = tuple[float, float, float] | torch.Tensor
 
-def calc_seq_lens(sizes: list[tuple[int, int]], patch_size: int) -> list[int]:
-    return [calc_seq_len(size, patch_size) for size in sizes]
 
+def _ntuple(n):
+    def parse(x):
+        if isinstance(x, Iterable) and not isinstance(x, str):
+            return tuple(x)
+        return tuple(repeat(x, n))
 
-class ClsToken(nn.Module):
-    def __init__(
-        self,
-        ndim: int,
-        num_tokens: int = 1,
-        enabled: bool = True,
-        register_multiple: int | None = None,
-        num_registers: int | None = None,
-    ):
-        super().__init__()
+    return parse
 
-        self.ndim = ndim
-        self.enabled = enabled
-        self.num_registers = 0
-        self.num_tokens = num_tokens
-        if enabled:
-            if num_registers:
-                self.num_registers = num_registers
-            elif register_multiple:
-                self.num_registers = register_multiple - (
-                    num_tokens % register_multiple
-                )
 
-            scale = ndim**-0.5
-            self.token = nn.Parameter(
-                torch.randn(num_tokens + self.num_registers, ndim) * scale
-            )
+to_1tuple = _ntuple(1)
+to_2tuple = _ntuple(2)
+to_3tuple = _ntuple(3)
+to_4tuple = _ntuple(4)
+to_ntuple = _ntuple
 
-        else:
-            self.token = None
 
-        self.num_patches = self.num_tokens + self.num_registers
+def calc_seq_len(size: tuple[int, int], patch_size: int) -> int:
+    h, w = size
+    return (h // patch_size) * (w // patch_size)
 
-    def forward(self, x: torch.Tensor):
-        if self.token is None:
-            return x
 
-        token = self.token.unsqueeze(0).expand(x.shape[0], -1, -1)
-        x = torch.cat(
-            [
-                token,
-                x,
-            ],
-            dim=1,
-        )
+def calc_seq_lens(sizes: list[tuple[int, int]], patch_size: int) -> list[int]:
+    return [calc_seq_len(size, patch_size) for size in sizes]
 
-        return x
 
+class RadioPatchEmbeddings(nn.Module):
+    """Cropped Position Embedding (CPE) patch generator, matching Transformers'
+    ``RadioPatchEmbeddings``, with additional support for dynamic multi-resolution
+    packing and video temporal compression.
+    """
 
-class ViTPatchGenerator(nn.Module):
-    def __init__(
-        self,
-        #  config: PreTrainedConfig,
-        patch_size: int,
-        embed_dim: int,
-        input_dims: input_dim_t,
-        abs_pos: bool = True,
-        normalize_patches: bool = False,
-        cls_token: bool = False,
-        max_input_dims: input_dim_t | None = None,
-        pos_dropout: float = 0.0,
-        return_pos_enc: bool = False,
-        num_cls_tokens: int = 1,
-        register_multiple: int | None = None,
-        num_registers: int | None = None,
-        patch_bias: bool = False,
-        temporal_patch_size: int = 1,
-        separate_video_embedder: bool = True,
-        device=None,
-        dtype=None,
-    ):
+    def __init__(self, config: PreTrainedConfig):
         super().__init__()
-        if isinstance(input_dims, int):
-            input_dims = (input_dims, input_dims)
-
-        if max_input_dims is None:
-            max_input_dims = input_dims
-        if isinstance(max_input_dims, int):
-            max_input_dims = (max_input_dims, max_input_dims)
+        patch_size = config.patch_size
+        embed_dim = config.hidden_size
+        num_channels = config.num_channels
 
-        max_input_dims = (
-            int(math.ceil(max_input_dims[0] / patch_size) * patch_size),
-            int(math.ceil(max_input_dims[1] / patch_size) * patch_size),
+        max_input_dims = to_2tuple(config.max_img_size)
+        max_input_dims = tuple(
+            int(math.ceil(d / patch_size) * patch_size) for d in max_input_dims
         )
-
-        self.cpe_mode = max_input_dims != input_dims
-        self.pos_dropout = pos_dropout
-        self.return_pos_enc = return_pos_enc
-
-        factory = dict(device=device, dtype=dtype)
+        input_dims = to_2tuple(config.image_size)
 
         self.patch_size = patch_size
-        self.abs_pos = abs_pos
         self.embed_dim = embed_dim
-        self.temporal_patch_size = temporal_patch_size
+        self.num_cls_tokens = config.num_cls_tokens
+        self.num_registers = config.num_registers
+        self.temporal_patch_size = config.video_temporal_patch_size
 
+        self.cpe_mode = max_input_dims != input_dims
         self.num_rows = max_input_dims[0] // patch_size
         self.num_cols = max_input_dims[1] // patch_size
-        self.input_dims = (
-            input_dims[0] // patch_size,
-            input_dims[1] // patch_size,
-        )
+        self.input_dims = tuple(d // patch_size for d in input_dims)
         self.num_patches = self.num_rows * self.num_cols
-        self.max_input_dims = max_input_dims
 
         self.im_to_patches = Im2Patches(patch_size)
-        self.embedder = ViTPatchLinear(
-            patch_size, embed_dim, bias=patch_bias, **factory
+        self.patch_projection = ViTPatchLinear(
+            patch_size, embed_dim, num_channels=num_channels, bias=False
         )
 
-        if temporal_patch_size > 1:
-            if not separate_video_embedder:
+        if self.temporal_patch_size > 1:
+            if not config.separate_video_embedder:
                 raise NotImplementedError(
                     "Only separate_video_embedder=True is supported for"
-                    " temporal compression (temporal_patch_size > 1)"
+                    " temporal compression (video_temporal_patch_size > 1)"
                 )
-            self.video_embedder = ViTPatchLinear(
+            self.video_patch_projection = ViTPatchLinear(
                 patch_size,
                 embed_dim,
-                bias=patch_bias,
-                temporal_patch_size=temporal_patch_size,
-                **factory,
+                num_channels=num_channels,
+                bias=False,
+                temporal_patch_size=self.temporal_patch_size,
             )
 
-        if abs_pos:
-            scale = embed_dim**-0.5
-            self.pos_embed = nn.Parameter(
-                torch.randn(1, self.num_patches, embed_dim, **factory) * scale
-            )
-
-        self.cls_token = ClsToken(
-            embed_dim,
-            num_tokens=num_cls_tokens,
-            enabled=cls_token,
-            register_multiple=register_multiple,
-            num_registers=num_registers,
+        scale = embed_dim**-0.5
+        self.position_embedding = nn.Parameter(
+            torch.randn(1, self.num_patches, embed_dim) * scale
         )
-
-        self.patch_normalizer = (
-            nn.LayerNorm(embed_dim) if normalize_patches else nn.Identity()
+        self.cls_register_token = nn.Parameter(
+            torch.randn(self.num_cls_tokens + self.num_registers, embed_dim) * scale
         )
+        self.patch_normalizer = nn.Identity()
+
+    @property
+    def num_skip(self) -> int:
+        return self.num_cls_tokens + self.num_registers
 
     def forward(
         self, x: torch.Tensor, imgs_sizes: list[tuple[int, int]] | None = None
     ) -> torch.Tensor:
         if imgs_sizes is not None:
-            patches = self.embedder(x)
-            patches, pos_enc = self.apply_pos_enc_dynamic(
-                patches, imgs_sizes=imgs_sizes
-            )
-            patches = self.cls_token_dynamic(patches, imgs_sizes=imgs_sizes)
+            # Dynamic multi-resolution: ``x`` is already patchified and packed
+            # to [B, total_patches, C*P*P].
+            patches = self.patch_projection(x)
+            patches = self._apply_pos_enc_dynamic(patches, imgs_sizes=imgs_sizes)
+            patches = self._cls_token_dynamic(patches, imgs_sizes=imgs_sizes)
         else:
-            patches = self.embed_patches(x)
-            patches, pos_enc = self.apply_pos_enc(patches, input_size=x.shape[2:])
-            patches = self.cls_token(patches)
-        patches = self.patch_normalizer(patches)
-        if self.return_pos_enc:
-            return patches, pos_enc
-        return patches
+            patches = self.patch_projection(self.im_to_patches(x))
+            patches = self._apply_pos_enc(patches, input_size=x.shape[2:])
+            patches = self._prepend_cls_token(patches)
+        return self.patch_normalizer(patches)
 
     def forward_video(self, x: torch.Tensor) -> torch.Tensor:
         """Process video frames with temporal compression.
@@ -238,148 +203,75 @@ def forward_video(self, x: torch.Tensor) -> torch.Tensor:
             feat=feat_dim,
         )
 
-        patches = self.video_embedder(patches)
+        patches = self.video_patch_projection(patches)
+        patches = self._apply_pos_enc(patches, input_size=input_size)
+        patches = self._prepend_cls_token(patches)
+        return self.patch_normalizer(patches)
 
-        patches, pos_enc = self.apply_pos_enc(patches, input_size=input_size)
+    def _prepend_cls_token(self, x: torch.Tensor) -> torch.Tensor:
+        token = self.cls_register_token.unsqueeze(0).expand(x.shape[0], -1, -1)
+        return torch.cat([token, x], dim=1)
 
-        patches = self.cls_token(patches)
-
-        patches = self.patch_normalizer(patches)
-        if self.return_pos_enc:
-            return patches, pos_enc
-        return patches
-
-    def apply_pos_enc_dynamic(
-        self, patches: torch.Tensor, imgs_sizes: list[tuple[int, int]]
-    ) -> tuple[torch.Tensor, torch.Tensor | None]:
-        if not self.abs_pos:
-            return patches, None
-
-        current_length = 0
-        pos_enc_list = []
-
-        for size in imgs_sizes:
-            seq_length = calc_seq_len(size, self.patch_size)
-
-            img_patches = patches[:, current_length : current_length + seq_length, :]
-            pos_enc = self.get_pos_enc(patches.shape[0], input_size=size)
-            img_patches_with_pos = img_patches + pos_enc
-
-            patches = torch.cat(
-                [
-                    patches[:, :current_length, :],
-                    img_patches_with_pos,
-                    patches[:, current_length + seq_length :, :],
-                ],
-                dim=1,
-            )
-            pos_enc_list.append(pos_enc)
-            current_length += seq_length
-
-        full_pos_enc = torch.cat(pos_enc_list, dim=1) if pos_enc_list else None
-        return patches, full_pos_enc
-
-    def cls_token_dynamic(
+    def _cls_token_dynamic(
         self, patches: torch.Tensor, imgs_sizes: list[tuple[int, int]]
     ) -> torch.Tensor:
-        if not self.cls_token.enabled:
-            return patches
-
         out = []
         current_length = 0
-
         for seq_len in calc_seq_lens(imgs_sizes, self.patch_size):
-            class_token = self.cls_token.token.unsqueeze(0).expand(
+            class_token = self.cls_register_token.unsqueeze(0).expand(
                 patches.shape[0], -1, -1
             )
             out.append(class_token)
             out.append(patches[:, current_length : current_length + seq_len, :])
             current_length += seq_len
-
         return torch.cat(out, dim=1)
 
-    @property
-    def apply_cls_token(self):
-        return self.cls_token.enabled
-
-    @property
-    def num_cls_tokens(self):
-        return self.cls_token.num_tokens
-
-    @property
-    def num_cls_patches(self):
-        return self.cls_token.num_patches
-
-    @property
-    def num_registers(self):
-        return self.cls_token.num_registers
-
-    @property
-    def num_skip(self):
-        return self.num_cls_tokens + self.num_registers
-
-    def embed_patches(self, x: torch.Tensor) -> torch.Tensor:
-        patches = self.im_to_patches(x)
-        patches = self.embedder(patches)
-        return patches
-
-    def apply_pos_enc(
+    def _apply_pos_enc(
         self,
         patches: torch.Tensor,
-        patch_idxs: torch.Tensor | None = None,
         input_size: tuple[int, int] | None = None,
     ) -> torch.Tensor:
-        if not self.abs_pos:
-            return patches
+        pos_enc = self._get_pos_enc(patches.shape[0], input_size=input_size)
+        return patches + pos_enc
 
-        pos_enc = self.get_pos_enc(patches.shape[0], patch_idxs, input_size)
-
-        if self.training and self.pos_dropout > 0:
-            keeps = (
-                torch.rand(
-                    patches.shape[0], 1, 1, dtype=pos_enc.dtype, device=pos_enc.device
-                )
-                > self.pos_dropout
+    def _apply_pos_enc_dynamic(
+        self, patches: torch.Tensor, imgs_sizes: list[tuple[int, int]]
+    ) -> torch.Tensor:
+        current_length = 0
+        for size in imgs_sizes:
+            seq_length = calc_seq_len(size, self.patch_size)
+            img_patches = patches[:, current_length : current_length + seq_length, :]
+            pos_enc = self._get_pos_enc(patches.shape[0], input_size=size)
+            img_patches_with_pos = img_patches + pos_enc
+            patches = torch.cat(
+                [
+                    patches[:, :current_length, :],
+                    img_patches_with_pos,
+                    patches[:, current_length + seq_length :, :],
+                ],
+                dim=1,
             )
-            pos_enc_drop = torch.where(keeps, pos_enc, 0)
-        else:
-            pos_enc_drop = pos_enc
-
-        return patches + pos_enc_drop, pos_enc
+            current_length += seq_length
+        return patches
 
-    def get_pos_enc(
+    def _get_pos_enc(
         self,
         batch_size: int,
-        patch_idxs: torch.Tensor | None = None,
         input_size: tuple[int, int] | None = None,
     ) -> torch.Tensor:
         if input_size is None:
             input_dims = self.input_dims
         else:
-            input_dims = (
-                input_size[0] // self.patch_size,
-                input_size[1] // self.patch_size,
-            )
-
-        pos_embed = self._get_pos_embeddings(batch_size, input_dims)
-
-        if patch_idxs is None:
-            return pos_embed
-
-        exp_patch_idxs = patch_idxs.unsqueeze(-1).expand(-1, -1, pos_embed.shape[-1])
-
-        pos_embed = torch.gather(
-            pos_embed.expand(patch_idxs.shape[0], -1, -1), dim=1, index=exp_patch_idxs
-        )
-        return pos_embed
+            input_dims = tuple(d // self.patch_size for d in input_size)
+        return self._get_pos_embeddings(batch_size, input_dims)
 
     def _get_pos_embeddings(self, batch_size: int, input_dims: tuple[int, int]):
         if (self.num_rows, self.num_cols) == input_dims:
-            return self.pos_embed
+            return self.position_embedding
 
-        pos_embed = self.pos_embed.reshape(1, self.num_rows, self.num_cols, -1).permute(
-            0, 3, 1, 2
-        )
+        pos_embed = self.position_embedding.reshape(
+            1, self.num_rows, self.num_cols, -1
+        ).permute(0, 3, 1, 2)
 
         def window_select(pos_embed):
             if input_dims[0] < pos_embed.shape[-2]:
@@ -440,12 +332,16 @@ def __init__(
         self,
         patch_size: int,
         embed_dim: int,
+        num_channels: int = 3,
         bias: bool = False,
         temporal_patch_size: int = 1,
         **factory,
     ):
         super().__init__(
-            3 * temporal_patch_size * (patch_size**2), embed_dim, bias=bias, **factory
+            num_channels * temporal_patch_size * (patch_size**2),
+            embed_dim,
+            bias=bias,
+            **factory,
         )
         self.patch_size = patch_size
         self.temporal_patch_size = temporal_patch_size
@@ -457,67 +353,288 @@ class MaskMetadata:
     max_seqlen: torch.Tensor
 
 
-class RadioParallelAttention(InternParallelAttention):
+class RadioAttention(nn.Module):
+    """RADIO self-attention: a fused ``qkv`` projection feeding the vLLM
+    attention backend, with optional QK-normalization and dynamic-resolution
+    packed attention (via ``mask_meta``). The fused counterpart of the
+    Transformers ``RadioAttention`` (which keeps ``query``/``key``/``value`` and
+    ``output.dense`` split); ``load_weights`` bridges the two."""
+
+    def __init__(
+        self,
+        config: PreTrainedConfig,
+        quant_config: QuantizationConfig | None = None,
+        *,
+        num_dummy_heads: int = 0,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.embed_dim = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.head_dim = self.embed_dim // self.num_heads
+        if self.head_dim * self.num_heads != self.embed_dim:
+            raise ValueError(
+                f"embed_dim must be divisible by num_heads (got embed_dim="
+                f"{self.embed_dim} and num_heads={self.num_heads})."
+            )
+
+        use_data_parallel = is_vit_use_data_parallel()
+        # Disable attention TP when the head count is not divisible by tp_size.
+        tp_size = 1 if use_data_parallel else get_tensor_model_parallel_world_size()
+        use_data_parallel = (
+            use_data_parallel or (self.num_heads + num_dummy_heads) % tp_size != 0
+        )
+        self.tp_size = 1 if use_data_parallel else tp_size
+        self.tp_rank = 0 if use_data_parallel else get_tensor_model_parallel_rank()
+
+        # Dummy heads pad the head count so TP divides evenly on common GPU counts.
+        self.dummy_dim = (num_dummy_heads + self.num_heads) * self.head_dim
+        self.num_heads_per_partition = divide(
+            num_dummy_heads + self.num_heads, self.tp_size
+        )
+        self.scale = self.head_dim**-0.5
+
+        self.qkv = QKVParallelLinear(
+            self.embed_dim,
+            self.head_dim,
+            num_dummy_heads + self.num_heads,
+            bias=config.qkv_bias,
+            quant_config=quant_config,
+            prefix=f"{prefix}.qkv",
+            disable_tp=use_data_parallel,
+        )
+        self.qk_normalization = config.qk_normalization
+        if self.qk_normalization:
+            self.q_norm = RMSNorm(
+                self.dummy_dim,
+                eps=config.layer_norm_eps,
+                var_hidden_size=self.embed_dim,
+            )
+            self.k_norm = RMSNorm(
+                self.dummy_dim,
+                eps=config.layer_norm_eps,
+                var_hidden_size=self.embed_dim,
+            )
+        self.proj = RowParallelLinear(
+            self.dummy_dim,
+            self.embed_dim,
+            quant_config=quant_config,
+            prefix=f"{prefix}.proj",
+            disable_tp=use_data_parallel,
+        )
+        self.attn = MMEncoderAttention(
+            self.num_heads_per_partition,
+            self.head_dim,
+            self.scale,
+            prefix=f"{prefix}.attn",
+        )
+
+    def _apply_qk_norm(
+        self, q: torch.Tensor, k: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if self.tp_size > 1:
+            q = tensor_model_parallel_all_gather(q.contiguous())
+            k = tensor_model_parallel_all_gather(k.contiguous())
+        q = self.q_norm(q)
+        k = self.k_norm(k)
+        if self.tp_size > 1:
+            splitter = partial(split_tensor_along_last_dim, num_partitions=self.tp_size)
+            q = splitter(q)[self.tp_rank]
+            k = splitter(k)[self.tp_rank]
+        return q, k
+
     def forward(
-        self, x: torch.Tensor, mask_meta: MaskMetadata | None = None
+        self,
+        hidden_states: torch.Tensor,
+        cu_seqlens: torch.Tensor | None = None,
+        max_seqlen: torch.Tensor | None = None,
     ) -> torch.Tensor:
-        qkv, _ = self.qkv(x)
+        qkv, _ = self.qkv(hidden_states)
         q, k, v = qkv.chunk(3, dim=-1)
 
         if self.qk_normalization:
             q, k = self._apply_qk_norm(q, k)
 
-        cu_seqlens, max_seqlen = None, None
-        if mask_meta is not None:
-            cu_seqlens = mask_meta.cu_seqlens
-            max_seqlen = mask_meta.max_seqlen
         out = self.attn(q, k, v, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen)
         out, _ = self.proj(out)
         return out
 
 
-class RadioVisionEncoderLayer(InternVisionEncoderLayer):
-    def __init__(self, *args, **kwargs) -> None:
-        super().__init__(*args, attn_cls=RadioParallelAttention, **kwargs)
+class RadioLayerScale(nn.Module):
+    """LayerScale matching the Transformers ``RadioLayerScale``: a ``lambda1``
+    parameter scaling the residual branch."""
+
+    def __init__(self, config: PreTrainedConfig) -> None:
+        super().__init__()
+        self.lambda1 = nn.Parameter(
+            config.layerscale_value * torch.ones(config.hidden_size)
+        )
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return hidden_states * self.lambda1
+
+
+@support_torch_compile(
+    # Batch is always 1 here; the packed tokens in dim 1 and cu_seqlens are what vary.
+    dynamic_arg_dims={"hidden_states": 1, "cu_seqlens": 0},
+    enable_if=should_torch_compile_mm_encoder,
+    is_encoder=True,
+)
+class RadioLayer(nn.Module):
+    """RADIO transformer block, mirroring the Transformers ``RadioLayer``
+    (``norm1`` / ``attention`` / ``layer_scale1`` / ``norm2`` / ``mlp`` /
+    ``layer_scale2``). ``drop_path`` is an inference-time ``Identity``; the MLP
+    reuses :class:`InternMLP` (``fc1``/``fc2``, matching ``RadioMLP``)."""
+
+    def __init__(
+        self,
+        config: PreTrainedConfig,
+        quant_config: QuantizationConfig | None = None,
+        *,
+        num_dummy_heads: int = 0,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.norm1 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.attention = RadioAttention(
+            config,
+            quant_config=quant_config,
+            num_dummy_heads=num_dummy_heads,
+            prefix=f"{prefix}.attention",
+        )
+        self.layer_scale1 = RadioLayerScale(config)
+        self.drop_path = nn.Identity()
+        self.norm2 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.mlp = InternMLP(config, quant_config=quant_config, prefix=f"{prefix}.mlp")
+        self.layer_scale2 = RadioLayerScale(config)
 
     def forward(
         self,
         hidden_states: torch.Tensor,
-        mask_meta: MaskMetadata | None = None,
-    ):
-        hidden_states = (
-            hidden_states
-            + self.attn(self.norm1(hidden_states), mask_meta=mask_meta) * self.ls1
+        cu_seqlens: torch.Tensor | None = None,
+        max_seqlen: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        attn_out = self.attention(
+            self.norm1(hidden_states), cu_seqlens=cu_seqlens, max_seqlen=max_seqlen
         )
-
-        hidden_states = hidden_states + self.mlp(self.norm2(hidden_states)) * self.ls2
-
+        hidden_states = self.drop_path(self.layer_scale1(attn_out)) + hidden_states
+        mlp_out = self.mlp(self.norm2(hidden_states))
+        hidden_states = self.drop_path(self.layer_scale2(mlp_out)) + hidden_states
         return hidden_states
 
 
-class RadioVisionEncoder(InternVisionEncoder):
-    def __init__(self, *args, **kwargs) -> None:
-        super().__init__(*args, layer_cls=RadioVisionEncoderLayer, **kwargs)
+class RadioVisionEncoder(nn.Module):
+    """RADIO encoder stack, mirroring the Transformers ``RadioEncoder`` module
+    tree (a ``layer`` ModuleList of :class:`RadioLayer`)."""
+
+    def __init__(
+        self,
+        config: PreTrainedConfig,
+        quant_config: QuantizationConfig | None = None,
+        *,
+        num_hidden_layers_override: int | None = None,
+        num_dummy_heads: int = 0,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        if num_hidden_layers_override is None:
+            num_hidden_layers = config.num_hidden_layers
+        else:
+            num_hidden_layers = num_hidden_layers_override
+        self.layer = nn.ModuleList(
+            [
+                RadioLayer(
+                    config,
+                    quant_config=quant_config,
+                    num_dummy_heads=num_dummy_heads,
+                    prefix=f"{prefix}.layer.{layer_idx}",
+                )
+                for layer_idx in range(num_hidden_layers)
+            ]
+        )
 
     def forward(
         self,
         inputs_embeds: torch.Tensor,
         mask_meta: MaskMetadata | None = None,
     ):
+        # Metadata implies a packed batch, and the compiled graph bakes that dim 0 == 1.
+        assert mask_meta is None or inputs_embeds.shape[0] == 1
         hidden_states = inputs_embeds
-        for encoder_layer in self.layers:
-            hidden_states = encoder_layer(hidden_states, mask_meta=mask_meta)
+        cu_seqlens = mask_meta.cu_seqlens if mask_meta is not None else None
+        max_seqlen = mask_meta.max_seqlen if mask_meta is not None else None
+        for encoder_layer in self.layer:
+            hidden_states = encoder_layer(
+                hidden_states, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen
+            )
         return hidden_states
 
 
-class RadioInternVisionModel(nn.Module):
+class RadioModel(nn.Module):
     packed_modules_mapping = {
         "qkv": ["qkv"],
     }
 
+    # Map both RADIO checkpoint layouts onto this module tree, which mirrors the
+    # Transformers ``RadioModel`` (``embeddings.*`` and ``encoder.layer.N.*``):
+    #   * native Transformers checkpoints load onto it directly; only the split
+    #     attention projections fuse into ``attention.qkv`` via
+    #     ``orig_to_new_stacked`` and the output projection is renamed;
+    #   * legacy remote-code checkpoints nest weights under
+    #     ``radio_model.model.{patch_generator,blocks}.*`` with a fused
+    #     ``attn.qkv`` (renamed to ``attention.qkv`` and loaded directly).
+    # ``video_patch_projection`` must precede ``embedder`` (applied in order);
+    # keys mapping to ``None`` are intentionally not loaded (input normalization
+    # runs in the processor and ``summary_idxs`` comes from the config).
+    hf_to_vllm_mapper = WeightsMapper(
+        orig_to_new_substr={
+            # legacy remote-code layout onto the native module tree
+            "radio_model.model.patch_generator.video_embedder": (
+                "embeddings.video_patch_projection"
+            ),
+            "radio_model.model.patch_generator.embedder": (
+                "embeddings.patch_projection"
+            ),
+            "radio_model.model.patch_generator.pos_embed": (
+                "embeddings.position_embedding"
+            ),
+            "radio_model.model.patch_generator.cls_token.token": (
+                "embeddings.cls_register_token"
+            ),
+            "radio_model.model.blocks": "encoder.layer",
+            # legacy names onto the native tree: ``attn`` -> ``attention`` (its
+            # fused ``qkv``/``proj`` load directly) and bare ``ls1``/``ls2`` ->
+            # the ``layer_scale{1,2}`` LayerScale modules.
+            ".attn.": ".attention.",
+            ".ls1": ".layer_scale1.lambda1",
+            ".ls2": ".layer_scale2.lambda1",
+            # native attention output projection: ``attention.output.dense`` is
+            # current at transformers 5.x (Dinov2 naming); ``attention.o_proj``
+            # covers the post-refactor spelling on transformers main.
+            "attention.output.dense": "attention.proj",
+            "attention.o_proj": "attention.proj",
+            # dropped in both layouts: normalization runs in the processor and
+            # ``summary_idxs`` comes from the config.
+            "input_conditioner": None,
+            "summary_idxs": None,
+        },
+        orig_to_new_stacked={
+            # Native split q/k/v projections fuse into the ``qkv``
+            # QKVParallelLinear. ``attention.attention.{query,key,value}`` is
+            # current at transformers 5.x; ``attention.{q,k,v}_proj`` covers the
+            # post-refactor spelling on transformers main.
+            "attention.attention.query": ("attention.qkv", "q"),
+            "attention.attention.key": ("attention.qkv", "k"),
+            "attention.attention.value": ("attention.qkv", "v"),
+            "attention.q_proj": ("attention.qkv", "q"),
+            "attention.k_proj": ("attention.qkv", "k"),
+            "attention.v_proj": ("attention.qkv", "v"),
+        },
+    )
+
     def __init__(
         self,
-        config: PreTrainedConfig = None,
+        config: PreTrainedConfig,
         quant_config: QuantizationConfig | None = None,
         *,
         num_hidden_layers_override: int | None = None,
@@ -527,27 +644,9 @@ def __init__(
         super().__init__()
 
         self.config = config
-        self.img_size, self.grid_size, self.num_patches = self._init_img_size(
-            config.patch_size, config.image_size
-        )
-        assert self.img_size is not None
-        max_img_size = int(
-            round(config.cpe_max_size / config.patch_size) * config.patch_size
-        )
         self.temporal_patch_size = config.video_temporal_patch_size
-        unique_teachers = set(t["name"] for t in config.teachers)
-        self.patch_generator = ViTPatchGenerator(
-            config.patch_size,
-            config.hidden_size,
-            input_dims=self.img_size,
-            max_input_dims=max_img_size,
-            cls_token=True,
-            num_cls_tokens=len(unique_teachers) if config.cls_token_per_teacher else 1,
-            register_multiple=config.register_multiple,
-            temporal_patch_size=self.temporal_patch_size,
-            separate_video_embedder=config.separate_video_embedder,
-        )
 
+        self.embeddings = RadioPatchEmbeddings(config)
         self.encoder = RadioVisionEncoder(
             config=config,
             quant_config=quant_config,
@@ -556,200 +655,95 @@ def __init__(
             prefix=f"{prefix}.encoder",
         )
 
-    def _init_img_size(
-        self,
-        patch_size: input_dim_t,
-        img_size: input_dim_t | None,
-    ) -> tuple[tuple[int, int] | None, tuple[int, int] | None, int | None]:
-        if img_size is None:
-            return None, None, None
-        if isinstance(img_size, int):
-            img_size = (img_size, img_size)
-        if isinstance(patch_size, int):
-            patch_size = (patch_size, patch_size)
-        grid_size = (
-            img_size[0] // patch_size[0],
-            img_size[1] // patch_size[1],
+        # summary_idxs selects which class tokens to gather: None (no teachers)
+        # Only a compiled layer needs this; TORCH_SDPA pays one call per item for it.
+        self.pack_uniform_batches = bool(self.encoder.layer) and not (
+            self.encoder.layer[0].do_not_compile
         )
-        num_patches = grid_size[0] * grid_size[1]
-        return img_size, grid_size, num_patches
 
-    def get_input_embeddings(self):
-        return self.embeddings
-
-    def inter_image_mask_metadata(
-        self, imgs_sizes: list[tuple[int, int]], device: torch.device
-    ) -> MaskMetadata:
-        """Build mask metadata from image pixel sizes. Adds num_skip to each
-        sequence length (cls/register tokens) to match patch generator output."""
-        patch_size = self.patch_generator.patch_size
-        num_skip = self.patch_generator.num_skip
-
-        seq_lens = calc_seq_lens(imgs_sizes, patch_size)
-        adjusted = [s + num_skip for s in seq_lens]
-        return self._inter_image_mask_metadata_from_seq_lens(adjusted, device=device)
-
-    def _inter_image_mask_metadata_from_seq_lens(
-        self, seq_lens: list[int], device: torch.device
-    ) -> MaskMetadata:
-        """Build mask metadata from actual sequence lengths (already including
-        cls/register tokens, i.e. patch_count + num_skip per item).
-        Use inter_image_mask_metadata() when you only have imgs_sizes."""
-        assert len(seq_lens) > 0
-        cu_seqlens = torch.tensor(
-            list(accumulate(seq_lens, initial=0)), dtype=torch.int32, device=device
-        )
-        # Keep max_seqlen on CPU to avoid .item() sync
-        # See: https://github.com/vllm-project/vllm/blob/20b6b01/vllm/v1/attention/ops/vit_attn_wrappers.py#L48
-        max_seqlen = torch.tensor(max(seq_lens), dtype=torch.int32)
-        return MaskMetadata(cu_seqlens=cu_seqlens, max_seqlen=max_seqlen)
+        # keeps all class tokens; a list (possibly empty) gathers those indices.
+        if config.summary_idxs is not None:
+            self.register_buffer(
+                "summary_idxs", torch.tensor(config.summary_idxs, dtype=torch.long)
+            )
+        else:
+            self.summary_idxs = None
 
     def forward(
         self,
-        x: torch.Tensor,
+        pixel_values: torch.Tensor | None = None,
+        *,
         imgs_sizes: list[tuple[int, int]] | None = None,
         num_frames: int | None = None,
-    ) -> torch.FloatTensor:
+    ) -> tuple[torch.FloatTensor, torch.FloatTensor]:
         T = self.temporal_patch_size
 
-        # Build packed-sequence metadata for MMEncoderAttention when needed.
-        mask_meta = None
-        packed_batch_size = None  # Original batch size before packing
+        packed_batch_size = None  # Original batch size before packing.
 
+        # A compiled layer bakes the traced cu_seqlens form and ignores the others.
         if num_frames is not None and T > 1:
             # Conv3d video: all tubelets have the same sequence length.
-            # Pack [num_tubelets, seq_per_tubelet, hidden] → [1, total, hidden]
-            hidden_states = self.patch_generator.forward_video(x)
-            packed_batch_size, seq_per_tubelet, hidden_dim = hidden_states.shape
-            hidden_states = hidden_states.reshape(1, -1, hidden_dim)
-            mask_meta = self._inter_image_mask_metadata_from_seq_lens(
-                [seq_per_tubelet] * packed_batch_size, device=hidden_states.device
+            hidden_states = self.embeddings.forward_video(pixel_values)
+            hidden_states, mask_meta, packed_batch_size, seq_len = (
+                self._pack_uniform_batch(hidden_states)
             )
         else:
-            # Images for any model, or video for non-conv3d model
-            hidden_states = self.patch_generator(x, imgs_sizes=imgs_sizes)
-            if imgs_sizes is not None and len(imgs_sizes) > 1:
-                # Dynamic resolution w/ > 1 image, create attn mask
-                mask_meta = self.inter_image_mask_metadata(
+            hidden_states = self.embeddings(pixel_values, imgs_sizes=imgs_sizes)
+            if imgs_sizes is not None:
+                # Already packed; one image gives cu_seqlens [0, n], i.e. unmasked.
+                mask_meta = self._inter_image_mask_metadata(
                     imgs_sizes, device=hidden_states.device
                 )
+            elif self.pack_uniform_batches:
+                hidden_states, mask_meta, packed_batch_size, seq_len = (
+                    self._pack_uniform_batch(hidden_states)
+                )
+            else:
+                mask_meta = None
 
-        encoder_outputs = self.encoder(inputs_embeds=hidden_states, mask_meta=mask_meta)
+        encoder_outputs = self.encoder(hidden_states, mask_meta=mask_meta)
 
-        # Unpack back to original batch shape if we packed for video
+        # Unpack back to the original batch shape.
         if packed_batch_size is not None:
-            encoder_outputs = encoder_outputs.reshape(
-                packed_batch_size, seq_per_tubelet, -1
-            )
-
-        return encoder_outputs
+            encoder_outputs = encoder_outputs.reshape(packed_batch_size, seq_len, -1)
 
+        return self._extract_final(encoder_outputs, imgs_sizes=imgs_sizes)
 
-class RadioModel(nn.Module):
-    packed_modules_mapping = {
-        "qkv": ["qkv"],
-    }
-
-    def __init__(
-        self,
-        config: PreTrainedConfig,
-        quant_config: QuantizationConfig | None = None,
-        *,
-        num_hidden_layers_override: int | None = None,
-        num_dummy_heads: int = 0,
-        prefix: str = "",
-    ) -> None:
-        super().__init__()
-
-        self.config = config
-        self.model = RadioInternVisionModel(
-            config=config,
-            quant_config=quant_config,
-            num_hidden_layers_override=num_hidden_layers_override,
-            num_dummy_heads=num_dummy_heads,
-            prefix=prefix,
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        """Load RADIO checkpoint weights, remapping legacy remote-code and
+        native Transformers keys onto this module tree via
+        :data:`hf_to_vllm_mapper`."""
+        if isinstance(weights, dict):
+            weights = weights.items()
+        # Keep only tensors belonging to the RADIO tower (legacy ``radio_model.*``
+        # or native top-level keys); ignore any auxiliary keys a fuller checkpoint
+        # may carry (an unmapped tower key still errors).
+        prefixes = (
+            "radio_model.",
+            "embeddings.",
+            "encoder.",
+            "input_conditioner.",
+            "summary_idxs",
         )
-
-        summary_idxs = None
-        if config.teachers:
-            summary_idxs = torch.tensor(
-                [i for i, t in enumerate(config.teachers) if t.get("use_summary", True)]
+        weights = ((name, w) for name, w in weights if name.startswith(prefixes))
+        loader = AutoWeightsLoader(self)
+        loaded = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
+        if not loaded:
+            logger.warning(
+                "RadioModel.load_weights matched no checkpoint tensors; the "
+                "vision tower will keep its random initialization. Expected "
+                "legacy 'radio_model.*' or native 'embeddings.*'/'encoder.*' "
+                "keys."
             )
-            if summary_idxs.numel() > 0:
-                self.register_buffer("summary_idxs", summary_idxs)
-        self.summary_idxs = summary_idxs
-
-    def forward(
-        self,
-        pixel_values: torch.Tensor | None = None,
-        pixel_embeds: torch.Tensor | None = None,
-        *,
-        imgs_sizes: list[tuple[int, int]] | None = None,
-        num_frames: int | None = None,
-    ) -> tuple[torch.FloatTensor, torch.FloatTensor]:
-        y = self.model(
-            pixel_values,
-            imgs_sizes=imgs_sizes,
-            num_frames=num_frames,
-        )
-        return self._extract_final(y, imgs_sizes=imgs_sizes)
-
-    def load_weights(self, weights) -> set[str]:
-        loaded_params: set[str] = set()
-        params_dict = dict(self.named_parameters())
-
-        if isinstance(weights, dict):
-            weights_list = list(weights.items())
-        else:
-            weights_list = list(weights)
-
-        for name, weight in weights_list:
-            if not name.startswith("radio_model."):
-                # Skip non-radio weights
-                continue
-
-            sub = name[len("radio_model.") :]  # drop "radio_model." prefix
-
-            # Skip buffers not used in vLLM
-            if sub in {"summary_idxs"}:
-                continue
-            if sub.startswith("input_conditioner."):
-                # we normalize in the input processor,
-                # based on norm and std values from the config
-                continue
-
-            vllm_key = None
-            if sub.startswith("model.patch_generator."):
-                vllm_key = f"model.patch_generator.{sub.split('.', 2)[-1]}"
-            elif sub.startswith("input_conditioner."):
-                vllm_key = f"input_conditioner.{sub.split('.', 1)[-1]}"
-            elif sub.startswith("model.blocks."):
-                # Encoder blocks: HF 'model.blocks.{i}.' ->
-                # vLLM 'model.encoder.layers.{i}.'
-                parts = sub.split(".")
-                if len(parts) >= 4:
-                    layer_idx = parts[2]
-                    suffix = ".".join(parts[3:])
-                    # Skip layer-scale entries that vLLM doesn't use
-                    if suffix in {"ls1", "ls2"} or suffix.startswith(("ls1.", "ls2.")):
-                        continue
-                    vllm_key = f"model.encoder.layers.{layer_idx}.{suffix}"
-
-            if vllm_key and vllm_key in params_dict:
-                param = params_dict[vllm_key]
-                weight_loader = getattr(param, "weight_loader", default_weight_loader)
-                weight_loader(param, weight)
-                loaded_params.add(vllm_key)
-
-        return loaded_params
+        return loaded
 
     def _extract_final(
         self, y: torch.Tensor, imgs_sizes: list[tuple[int, int]] | None = None
     ) -> tuple[torch.FloatTensor, torch.FloatTensor]:
-        # Remove CLS + REGISTERS tokens
-        num_skip = self.model.patch_generator.num_skip
-        patch_size = self.model.patch_generator.patch_size
-        num_cls_tokens = self.model.patch_generator.num_cls_tokens
+        # Remove CLS + REGISTER tokens.
+        num_skip = self.embeddings.num_skip
+        patch_size = self.embeddings.patch_size
+        num_cls_tokens = self.embeddings.num_cls_tokens
         if imgs_sizes is None:
             all_summary = y[:, :num_cls_tokens]
             all_feat = y[:, num_skip:]
@@ -773,3 +767,55 @@ def _extract_final(
         else:
             bb_summary = all_summary
         return bb_summary.flatten(1), all_feat
+
+    def get_input_embeddings(self):
+        return self.embeddings
+
+    def _pack_uniform_batch(
+        self, hidden_states: torch.Tensor
+    ) -> tuple[torch.Tensor, MaskMetadata, int, int]:
+        """Pack [B, n, hidden] into [1, B*n, hidden] plus its varlen metadata.
+
+        FA/Triton synthesise this cu_seqlens anyway; TORCH_SDPA splits it per item.
+        """
+        batch_size, seq_len, hidden_dim = hidden_states.shape
+        hidden_states = hidden_states.reshape(1, -1, hidden_dim)
+        # A Python list here would be a blocking host copy on every encoder forward.
+        cu_seqlens = torch.arange(
+            0,
+            (batch_size + 1) * seq_len,
+            seq_len,
+            dtype=torch.int32,
+            device=hidden_states.device,
+        )
+        mask_meta = MaskMetadata(
+            cu_seqlens=cu_seqlens,
+            max_seqlen=torch.tensor(seq_len, dtype=torch.int32),
+        )
+        return hidden_states, mask_meta, batch_size, seq_len
+
+    def _inter_image_mask_metadata(
+        self, imgs_sizes: list[tuple[int, int]], device: torch.device
+    ) -> MaskMetadata:
+        """Build mask metadata from image pixel sizes. Adds ``num_skip`` to each
+        sequence length (cls/register tokens) to match patch generator output."""
+        patch_size = self.embeddings.patch_size
+        num_skip = self.embeddings.num_skip
+        seq_lens = calc_seq_lens(imgs_sizes, patch_size)
+        adjusted = [s + num_skip for s in seq_lens]
+        return self._mask_metadata_from_seq_lens(adjusted, device=device)
+
+    def _mask_metadata_from_seq_lens(
+        self, seq_lens: list[int], device: torch.device
+    ) -> MaskMetadata:
+        """Build mask metadata from sequence lengths (already including
+        cls/register tokens, i.e. ``patch_count + num_skip`` per item)."""
+        assert len(seq_lens) > 0
+        # torch.tensor(list, device=cuda) here is a blocking host copy per forward.
+        cu_seqlens = async_tensor_h2d(
+            list(accumulate(seq_lens, initial=0)), device=device, dtype=torch.int32
+        )
+        # Keep max_seqlen on CPU to avoid .item() sync.
+        # See: https://github.com/vllm-project/vllm/blob/20b6b01/vllm/v1/attention/ops/vit_attn_wrappers.py#L48
+        max_seqlen = torch.tensor(max(seq_lens), dtype=torch.int32)
+        return MaskMetadata(cu_seqlens=cu_seqlens, max_seqlen=max_seqlen)
diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py
index d7f3c3bbf4a2..da9d0ce9737e 100644
--- a/vllm/model_executor/models/registry.py
+++ b/vllm/model_executor/models/registry.py
@@ -187,7 +187,6 @@
     "NemotronHPuzzleForCausalLM": ("nemotron_h", "NemotronHForCausalLM"),
     "MuseGlimmerForCausalLM": ("muse_glimmer", "MuseGlimmerForCausalLM"),
     "OlmoHybridForCausalLM": ("olmo_hybrid", "OlmoHybridForCausalLM"),
-    "OlmoeForCausalLM": ("olmoe", "OlmoeForCausalLM"),
     "OPTForCausalLM": ("opt", "OPTForCausalLM"),
     "OrionForCausalLM": ("orion", "OrionForCausalLM"),
     "PanguEmbeddedForCausalLM": ("openpangu", "PanguEmbeddedForCausalLM"),
@@ -201,8 +200,8 @@
     "Qwen2MoeForCausalLM": ("qwen2_moe", "Qwen2MoeForCausalLM"),
     "Qwen3ForCausalLM": ("qwen3", "Qwen3ForCausalLM"),
     "Qwen3MoeForCausalLM": ("qwen3_moe", "Qwen3MoeForCausalLM"),
-    "Qwen3_5ForCausalLM": ("qwen3_5", "Qwen3_5ForCausalLM"),
-    "Qwen3_5MoeForCausalLM": ("qwen3_5", "Qwen3_5MoeForCausalLM"),
+    "Qwen3_5ForCausalLM": ("vllm.models.qwen3_5", "Qwen3_5ForCausalLM"),
+    "Qwen3_5MoeForCausalLM": ("vllm.models.qwen3_5", "Qwen3_5MoeForCausalLM"),
     "SarvamMoEForCausalLM": ("sarvam", "SarvamMoEForCausalLM"),
     "SarvamMLAForCausalLM": ("sarvam", "SarvamMLAForCausalLM"),
     "Step1ForCausalLM": ("step1", "Step1ForCausalLM"),
@@ -478,11 +477,6 @@
     ),
     "IsaacForConditionalGeneration": ("isaac", "IsaacForConditionalGeneration"),
     "KananaVForConditionalGeneration": ("kanana_v", "KananaVForConditionalGeneration"),
-    "KeyeForConditionalGeneration": ("keye", "KeyeForConditionalGeneration"),
-    "KeyeVL1_5ForConditionalGeneration": (
-        "keye_vl1_5",
-        "KeyeVL1_5ForConditionalGeneration",
-    ),
     "KimiVLForConditionalGeneration": ("kimi_vl", "KimiVLForConditionalGeneration"),
     "KimiK25ForConditionalGeneration": ("kimi_k25", "KimiK25ForConditionalGeneration"),
     "KimiK3ForConditionalGeneration": (
@@ -520,6 +514,10 @@
     ),
     "MiDashengLMModel": ("midashenglm", "MiDashengLMModel"),
     "MiMoV2OmniForCausalLM": ("mimo_v2_omni", "MiMoV2OmniForCausalLM"),
+    "MingV3ASRForConditionalGeneration": (
+        "ming_v3_asr",
+        "MingV3ASRForConditionalGeneration",
+    ),
     "MiniMaxM3SparseForConditionalGeneration": (
         "vllm.models.minimax_m3",
         "MiniMaxM3SparseForConditionalGeneration",
@@ -551,6 +549,10 @@
     "NemotronH_Nano_Omni_Reasoning_V3": ("nano_nemotron_vl", "NemotronH_Nano_VL_V2"),
     "NemotronH_Super_Omni_Reasoning_V3": ("nano_nemotron_vl", "NemotronH_Nano_VL_V2"),
     "NemotronH_Omni_Reasoning_V3": ("nano_nemotron_vl", "NemotronH_Nano_VL_V2"),
+    "Nemotron3_5TranscribeForConditionalGeneration": (
+        "nemotron3_5_transcribe",
+        "Nemotron3_5TranscribeForConditionalGeneration",
+    ),
     "NVLM_D": ("nvlm_d", "NVLM_D_Model"),
     "MuseGlimmerForConditionalGeneration": ("muse_glimmer", "MuseGlimmerForCausalLM"),
     "MuseSparkForConditionalGeneration": (
@@ -614,9 +616,12 @@
         "qwen3_vl_moe",
         "Qwen3VLMoeForConditionalGeneration",
     ),
-    "Qwen3_5ForConditionalGeneration": ("qwen3_5", "Qwen3_5ForConditionalGeneration"),
+    "Qwen3_5ForConditionalGeneration": (
+        "vllm.models.qwen3_5",
+        "Qwen3_5ForConditionalGeneration",
+    ),
     "Qwen3_5MoeForConditionalGeneration": (
-        "qwen3_5",
+        "vllm.models.qwen3_5",
         "Qwen3_5MoeForConditionalGeneration",
     ),
     "Qwen4ExpForConditionalGeneration": (
@@ -637,6 +642,7 @@
     "VoxtralForConditionalGeneration": ("voxtral", "VoxtralForConditionalGeneration"),
     "VoxtralRealtimeGeneration": ("voxtral_realtime", "VoxtralRealtimeGeneration"),
     # [Encoder-decoder]
+    "CanaryForConditionalGeneration": ("canary", "CanaryForConditionalGeneration"),
     "CohereAsrForConditionalGeneration": (
         "cohere_asr",
         "CohereAsrForConditionalGeneration",
@@ -678,6 +684,10 @@
         "vllm.models.deepseek_v41",
         "DSparkDeepseekV4ForCausalLM",
     ),
+    "DSparkMiniMaxDraftModel": (
+        "vllm.models.minimax_m3.common.dspark",
+        "MiniMaxM3DSparkForCausalLM",
+    ),
     "Qwen3DSparkModel": ("qwen3_dspark", "Qwen3DSparkForCausalLM"),
     "Qwen3OmniDSparkModel": ("qwen3_dspark", "Qwen3DSparkForCausalLM"),
     # Serving under method="dflash" means EAGLEConfig rewrites the draft
@@ -725,6 +735,10 @@
     "ExaoneMoeMTP": ("exaone_moe_mtp", "ExaoneMoeMTP"),
     "Exaone4_5_MTP": ("exaone4_5_mtp", "Exaone4_5_MTP"),
     "NemotronHMTPModel": ("nemotron_h_mtp", "NemotronHMTP"),
+    "Nemotron3_5TranscribeMTPModel": (
+        "nemotron3_5_transcribe",
+        "Nemotron3_5TranscribeMTP",
+    ),
     "LongCatFlashMTPModel": ("longcat_flash_mtp", "LongCatFlashMTP"),
     "Glm4MoeMTPModel": ("glm4_moe_mtp", "Glm4MoeMTP"),
     "Glm4MoeLiteMTPModel": ("glm4_moe_lite_mtp", "Glm4MoeLiteMTP"),
@@ -736,8 +750,8 @@
     "Step5MTP": ("vllm.models.step5", "Step5MTP"),
     "Qwen4ExpMTP": ("vllm.models.qwen4_exp", "Qwen4ExpMTP"),
     "Step3p5MTP": ("step3p5_mtp", "Step3p5MTP"),
-    "Qwen3_5MTP": ("qwen3_5_mtp", "Qwen3_5MTP"),
-    "Qwen3_5MoeMTP": ("qwen3_5_mtp", "Qwen3_5MoeMTP"),
+    "Qwen3_5MTP": ("vllm.models.qwen3_5", "Qwen3_5MTP"),
+    "Qwen3_5MoeMTP": ("vllm.models.qwen3_5", "Qwen3_5MoeMTP"),
     "InternS2MobiusMTP": ("interns2_mobius", "InternS2MobiusMTP"),
     "HYV3MTPModel": ("hy_v3_mtp", "HYV3MTP"),
     "HYV4MTPModel": ("vllm.models.hy_v4", "HYV4MTP"),
@@ -761,6 +775,7 @@
     "OlmoForCausalLM": ("transformers", "TransformersForCausalLM"),
     "Olmo2ForCausalLM": ("transformers", "TransformersForCausalLM"),
     "Olmo3ForCausalLM": ("transformers", "TransformersForCausalLM"),
+    "OlmoeForCausalLM": ("transformers", "TransformersMoEForCausalLM"),
     "PhiForCausalLM": ("transformers", "TransformersForCausalLM"),
     "SeedOssForCausalLM": ("transformers", "TransformersForCausalLM"),
     "SmolLM3ForCausalLM": ("transformers", "TransformersForCausalLM"),
@@ -785,6 +800,9 @@
     # Text generation models
     "TransformersForCausalLM": ("transformers", "TransformersForCausalLM"),
     "TransformersMoEForCausalLM": ("transformers", "TransformersMoEForCausalLM"),
+    # Multi-token prediction draft models
+    "TransformersForMTP": ("transformers", "TransformersForMTP"),
+    "TransformersMoEForMTP": ("transformers", "TransformersMoEForMTP"),
     # Multimodal models
     "TransformersMultiModalForCausalLM": (
         "transformers",
@@ -889,6 +907,8 @@
     "StableLMEpochForCausalLM": "0.28.0",
     # Superseded by the generic "Terratorch" architecture
     "PrithviGeoSpatialMAE": "0.28.0",
+    "KeyeForConditionalGeneration": "0.33.0",
+    "KeyeVL1_5ForConditionalGeneration": "0.33.0",
 }
 
 _OOT_SUPPORTED_MODELS = {
@@ -926,6 +946,10 @@ class _ModelInfo:
     supports_kv_cache_grouping: bool
     supported_video_pruning_methods: tuple[str, ...]
     supports_mm_device_do_normalize: bool
+    # Fully-qualified name of `SpeculativeDraftModel.speculator_config_class`.
+    speculator_config_class: str | None
+    # Whether a draft model may reuse the target's input embedding under PP.
+    may_use_target_embed: bool
 
     @staticmethod
     def from_model_cls(model: type[nn.Module]) -> "_ModelInfo":
@@ -964,6 +988,12 @@ def from_model_cls(model: type[nn.Module]) -> "_ModelInfo":
             supports_mm_device_do_normalize=getattr(
                 model, "supports_mm_device_do_normalize", False
             ),
+            speculator_config_class=(
+                f"{config_cls.__module__}.{config_cls.__qualname__}"
+                if (config_cls := getattr(model, "speculator_config_class", None))
+                else None
+            ),
+            may_use_target_embed=hasattr(model, "has_own_embed_tokens"),
         )
 
 
diff --git a/vllm/model_executor/models/transformers/__init__.py b/vllm/model_executor/models/transformers/__init__.py
index e19de6d4e891..2e50ed4e1f6a 100644
--- a/vllm/model_executor/models/transformers/__init__.py
+++ b/vllm/model_executor/models/transformers/__init__.py
@@ -25,6 +25,7 @@
 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
+from vllm.model_executor.models.transformers.mtp import MTPMixin
 from vllm.model_executor.models.transformers.multimodal import (
     MultiModalDummyInputsBuilder,
     MultiModalMixin,
@@ -132,6 +133,13 @@ class TransformersForCausalLM(CausalMixin): ...
 class TransformersMoEForCausalLM(MoEMixin, CausalMixin): ...
 
 
+# Multi-token prediction draft models
+class TransformersForMTP(MTPMixin): ...
+
+
+class TransformersMoEForMTP(MoEMixin, MTPMixin): ...
+
+
 # Multimodal models
 @MULTIMODAL_REGISTRY.register_processor(
     MultiModalProcessor,
diff --git a/vllm/model_executor/models/transformers/base.py b/vllm/model_executor/models/transformers/base.py
index 43898125ec7a..9985a1b5115e 100644
--- a/vllm/model_executor/models/transformers/base.py
+++ b/vllm/model_executor/models/transformers/base.py
@@ -32,6 +32,7 @@
 from transformers import AutoModel, PreTrainedConfig
 from transformers.conversion_mapping import (
     WeightRenaming,
+    WeightTransform,
     get_checkpoint_conversion_mapping,
     get_model_conversion_mapping,
 )
@@ -176,8 +177,7 @@ def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""):
             self._mark_model_components(vllm_config),
             init_on_device_without_buffers("meta"),
         ):
-            from_config_kwargs = self._from_config_kwargs
-            self.model: PreTrainedModel = AutoModel.from_config(**from_config_kwargs)
+            self.model: PreTrainedModel = self._create_model()
 
         # Create weight name to module qualname mapper
         self._create_hf_to_vllm_mapper()
@@ -228,6 +228,10 @@ def _from_config_kwargs(self) -> dict[str, Any]:
             trust_remote_code=self.model_config.trust_remote_code,
         )
 
+    def _create_model(self) -> "PreTrainedModel":
+        """Create the Transformers model that `self.model` wraps."""
+        return AutoModel.from_config(**self._from_config_kwargs)
+
     def _find_encoder_classes(
         self, model: "PreTrainedModel"
     ) -> dict[str, type["PreTrainedModel"]]:
@@ -316,14 +320,7 @@ def _create_hf_to_vllm_mapper(self):
         orig_to_new_renaming: list[WeightRenaming] = []
         orig_to_new_regex: dict[re.Pattern, str | None] = {}
 
-        mappings = get_model_conversion_mapping(self.model)
-        # Ensure layers that are part of the model, but not AutoModel, get their mapping
-        instantiated = {type(m).__name__ for m in self.model.modules()}
-        for arch in self.config.architectures or []:
-            if arch not in instantiated:
-                mappings.extend(get_checkpoint_conversion_mapping(arch) or [])
-
-        for mapping in mappings:
+        for mapping in self._weight_conversions():
             # Handle weights which have been renamed in Transformers
             if isinstance(mapping, WeightRenaming):
                 orig_to_new_renaming.append(mapping)
@@ -365,6 +362,20 @@ def _create_hf_to_vllm_mapper(self):
         # Apply mapping to quantization config if needed
         self._maybe_apply_model_mapping()
 
+    def _weight_conversions(self) -> list[WeightTransform]:
+        """The Transformers weight conversions that apply to `self.model`."""
+        mappings = get_model_conversion_mapping(self.model)
+        # Ensure layers that are part of the model, but not AutoModel, get their mapping
+        instantiated = {type(m).__name__ for m in self.model.modules()}
+        for arch in self.config.architectures or []:
+            if arch not in instantiated:
+                mappings.extend(get_checkpoint_conversion_mapping(arch) or [])
+        return mappings
+
+    def _tp_plan(self) -> dict:
+        """The tensor parallel plan of `self.model`, relative to `self.model`."""
+        return self.model.tp_plan or {}
+
     def _get_tie_word_embeddings(self):
         """Check if the model has tied word embeddings."""
         # Models created with Transformers v4 and v5 will store this in different places
@@ -423,11 +434,7 @@ def pipeline_parallel(self):
             attrsetter(name)(module, PPMissingLayer())
 
         # Module list
-        start_layer, end_layer = get_pp_indices(
-            self.text_config.num_hidden_layers,
-            self.pp_group.rank_in_group,
-            self.pp_group.world_size,
-        )
+        start_layer, end_layer = self._pp_indices()
         layers_name = names[module_list_idx]
         # attrgetter in case the module is nested (e.g. "text_model.layers")
         layers = attrgetter(layers_name)(module)
@@ -481,7 +488,7 @@ def recursive_replace(self):
         - RMSNorm (detected from their dataflow) with vLLM's `RMSNorm`or `GemmaRMSNorm`
         - `nn.LayerNorm` with vLLM's `LayerNorm`
         """
-        tp_plan = self.model.tp_plan or {}
+        tp_plan = self._tp_plan()
 
         if not tp_plan and self.tp_group.world_size > 1:
             tip = get_feature_request_tip(
@@ -591,6 +598,14 @@ def _recursive_replace(module: nn.Module, prefix: str):
 
         self.hf_to_vllm_mapper |= WeightsMapper(orig_to_new_stacked=orig_to_new_stacked)
 
+    def _pp_indices(self) -> tuple[int, int]:
+        """The range of decoder layers on this pipeline parallel rank."""
+        return get_pp_indices(
+            self.text_config.num_hidden_layers,
+            self.pp_group.rank_in_group,
+            self.pp_group.world_size,
+        )
+
     def _create_attention_instances(self, prefix: str = ""):
         """Create `Attention` instances to inform KV cache allocation."""
         text_config = self.text_config
@@ -610,9 +625,7 @@ def _create_attention_instances(self, prefix: str = ""):
         logits_soft_cap = getattr(text_config, "attn_logit_softcapping", None)
         layer_types = getattr(text_config, "layer_types", None)
 
-        pp_rank = self.pp_group.rank_in_group
-        pp_size = self.pp_group.world_size
-        start, end = get_pp_indices(text_config.num_hidden_layers, pp_rank, pp_size)
+        start, end = self._pp_indices()
 
         for i in range(start, end):
             if i not in self.attention_fusers:
diff --git a/vllm/model_executor/models/transformers/layers.py b/vllm/model_executor/models/transformers/layers.py
index 4ff3651dfc37..eefac0a6346f 100644
--- a/vllm/model_executor/models/transformers/layers.py
+++ b/vllm/model_executor/models/transformers/layers.py
@@ -1,42 +1,19 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Layer provider resolution for the Transformers modeling backend.
+"""Layer resolution for the Transformers modeling backend.
 
-When ``VLLM_USE_HW_AGNOSTIC`` is set, layer symbols are imported from
-``vllm.model_executor.hw_agnostic.layers.<module>``, falling back to
-``vllm.model_executor.layers.<module>`` for anything not yet ported. The
-resolved source of every symbol is logged so it is clear which layers run
-hw-agnostic and which fell back to vLLM.
+Layers may come from the in-tree and the hw-agnostic path,
+resolved via `vllm.model_executor.hw_agnostic.resolve`.
 """
 
-import importlib
-
 import vllm.envs as envs
 from vllm.logger import init_logger
+from vllm.model_executor import hw_agnostic
 
 logger = init_logger(__name__)
 
-_HW_PKG = "vllm.model_executor.hw_agnostic.layers"
-_VLLM_PKG = "vllm.model_executor.layers"
-
-
-def _resolve(module: str, name: str):
-    """Return `name` from the hw-agnostic `module` when enabled and available,
-    else from vLLM. Logs which source was used."""
-    if envs.VLLM_USE_HW_AGNOSTIC:
-        try:
-            obj = getattr(importlib.import_module(f"{_HW_PKG}.{module}"), name)
-            logger.info("Using hw-agnostic layer: %s", name)
-            return obj
-        except (ImportError, AttributeError):
-            logger.warning(
-                "hw-agnostic layer %s is not available; falling back to default", name
-            )
-    return getattr(importlib.import_module(f"{_VLLM_PKG}.{module}"), name)
-
-
-RMSNorm = _resolve("layernorm", "RMSNorm")
-GemmaRMSNorm = _resolve("layernorm", "GemmaRMSNorm")
+RMSNorm = hw_agnostic.resolve("layernorm", "RMSNorm")
+GemmaRMSNorm = hw_agnostic.resolve("layernorm", "GemmaRMSNorm")
 
 
 def get_act_and_mul_fn(act_fn_name: str):
diff --git a/vllm/model_executor/models/transformers/moe.py b/vllm/model_executor/models/transformers/moe.py
index a426a2c9d993..7c01c5634f6a 100644
--- a/vllm/model_executor/models/transformers/moe.py
+++ b/vllm/model_executor/models/transformers/moe.py
@@ -16,6 +16,7 @@
 # limitations under the License.
 """Transformers modeling backend mixin for Mixture of Experts (MoE) models."""
 
+from collections.abc import Iterable
 from dataclasses import dataclass
 from functools import partial
 from typing import TYPE_CHECKING, Any
@@ -165,6 +166,9 @@ def _hf_topk_router(
 
 
 class TransformersRoutedExperts(RoutedExperts):
+    _is_mixtral_format: bool | None = None
+    """Does the checkpoint use Mixtral style naming. `None` until weights are loaded."""
+
     def get_expert_mapping(
         self,
         ckpt_gate_proj_name: str | None = None,
@@ -180,11 +184,27 @@ def get_expert_mapping(
         )
         mixtral_map = super().get_expert_mapping("w1", "w2", "w3", include_fused)
         if not include_fused:
+            # LoRA expects one triplet per expert, so only return the loaded naming
+            if self._is_mixtral_format is not None:
+                return mixtral_map if self._is_mixtral_format else common_map
             return common_map + mixtral_map
         common_fused, common_unfused = common_map[:3], common_map[3:]
         mixtral_fused, mixtral_unfused = mixtral_map[:3], mixtral_map[3:]
         return common_fused + mixtral_fused + common_unfused + mixtral_unfused
 
+    def load_weights(
+        self, weights: Iterable[tuple[str, torch.Tensor]]
+    ) -> Iterable[str]:
+        return super().load_weights(self._detect_mixtral_format(weights))
+
+    def _detect_mixtral_format(
+        self, weights: Iterable[tuple[str, torch.Tensor]]
+    ) -> Iterable[tuple[str, torch.Tensor]]:
+        for name, weight in weights:
+            uses_w123 = any(f".{w}." in f".{name}" for w in ("w1", "w2", "w3"))
+            self._is_mixtral_format = bool(self._is_mixtral_format) or uses_w123
+            yield name, weight
+
 
 class MoEMixin(MixtureOfExperts, Base):
     mlp_layers: list[nn.Module]
diff --git a/vllm/model_executor/models/transformers/mtp.py b/vllm/model_executor/models/transformers/mtp.py
new file mode 100644
index 000000000000..05a5d5b458fb
--- /dev/null
+++ b/vllm/model_executor/models/transformers/mtp.py
@@ -0,0 +1,139 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+# Copyright 2026 The vLLM team.
+#
+# 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.
+"""Transformers modeling backend mixin for multi-token prediction draft models."""
+
+from collections.abc import Iterable
+from typing import TYPE_CHECKING
+
+import torch
+from transformers import AutoModelForCausalLM
+from transformers.conversion_mapping import WeightTransform
+
+from vllm.model_executor.models.transformers.causal import CausalMixin
+from vllm.sequence import IntermediateTensors
+
+if TYPE_CHECKING:
+    from transformers import PreTrainedModel
+
+
+class MTPMixin(CausalMixin):
+    """Serve the MTP layers of a model which opts into the Transformers MTP standard
+    (i.e. its config sets `num_mtp_layers`) as a draft model.
+
+    `self.model` is the Transformers `MtpModel`, whose embedding and LM head the
+    proposer replaces with the target's.
+    """
+
+    def _patch_config(self):
+        # `vllm_config.model_config` is the target's, which must not be patched
+        assert self.vllm_config.speculative_config is not None
+        self.model_config = self.vllm_config.speculative_config.draft_model_config
+        self.config = self.model_config.hf_config
+        self.text_config = self.config.get_text_config()
+        super()._patch_config()
+
+    def _create_model(self) -> "PreTrainedModel":
+        self.check_version("5.20.0", "MTP")
+        from transformers.modeling_layers import MtpModel
+
+        main_model = AutoModelForCausalLM.from_config(**self._from_config_kwargs)
+        num_mtp_layers = getattr(self.text_config, "num_mtp_layers", None)
+        mtp_model = MtpModel(main_model, num_mtp_layers)
+        self.mtp_weights = MtpModel.get_mtp_weights_regex(main_model)
+        self.mtp_weight_conversions = mtp_model.get_weight_conversions(main_model)
+
+        # The proposer shares the target's LM head with the draft
+        del mtp_model.shared_head
+        # With PP, the last rank has no target embedding to share with the draft
+        if self.pp_group.world_size > 1:
+            mtp_model._keys_to_ignore_on_load_unexpected = [
+                r"layers\..*embed_tokens\.weight" if "embed_tokens" in k else k
+                for k in mtp_model._keys_to_ignore_on_load_unexpected
+            ]
+        self.text_config = mtp_model.config
+        return mtp_model
+
+    def _weight_conversions(self) -> list[WeightTransform]:
+        return self.mtp_weight_conversions
+
+    def _tp_plan(self) -> dict:
+        tp_plan = self.text_config.base_model_tp_plan or {}
+        return {
+            k.replace("layers.*.", "layers.*.mtp_block.", 1): v
+            for k, v in tp_plan.items()
+            if k.startswith("layers.*.")
+        }
+
+    # The draft runs whole on the last PP rank
+    def pipeline_parallel(self):
+        pass
+
+    def _pp_indices(self) -> tuple[int, int]:
+        return 0, self.text_config.num_hidden_layers
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        embed_prefix = None
+        if self.pp_group.world_size > 1:
+            embed = self.model.get_input_embeddings()
+            embed_prefix = next(n for n, m in self.named_modules() if m is embed) + "."
+
+        def is_draft_weight(name: str) -> bool:
+            if self.mtp_weights.search(name):
+                return True
+            if embed_prefix is None:
+                return False
+            mapped = self.hf_to_vllm_mapper.map_name(name)
+            return mapped is not None and mapped.startswith(embed_prefix)
+
+        loaded = super().load_weights(w for w in weights if is_draft_weight(w[0]))
+        if embed_prefix is not None:
+            self.has_own_embed_tokens = any(n.startswith(embed_prefix) for n in loaded)
+        return loaded
+
+    def forward(  # type: ignore[override]
+        self,
+        input_ids: torch.Tensor | None,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        spec_step_idx: int = 0,
+        **kwargs,
+    ) -> torch.Tensor:
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_input_ids(input_ids)
+        # Add batch dimension before entering Transformers modules
+        inputs_embeds = inputs_embeds[None, ...]
+        hidden_states = hidden_states[None, ...]
+
+        hidden_states = self.model.forward_layer(
+            spec_step_idx % len(self.model.layers),
+            hidden_states,
+            positions[None, ...],
+            mtp_cache=None,
+            inputs_embeds=inputs_embeds,
+            **kwargs,
+        )
+
+        # Remove batch dimension after exiting Transformers modules
+        return hidden_states[0, ...]
+
+    def compute_logits(
+        self, hidden_states: torch.Tensor, spec_step_idx: int = 0
+    ) -> torch.Tensor | None:
+        # The MTP layers share the LM head
+        return super().compute_logits(hidden_states)
diff --git a/vllm/model_executor/models/utils.py b/vllm/model_executor/models/utils.py
index a86496b48a9b..05315b1c2c68 100644
--- a/vllm/model_executor/models/utils.py
+++ b/vllm/model_executor/models/utils.py
@@ -12,13 +12,16 @@
 import torch.nn as nn
 from torch.nn.modules.module import register_module_module_registration_hook
 
-from vllm.config import VllmConfig
+from vllm.config import VllmConfig, get_current_vllm_config_or_none
 from vllm.distributed import (
     get_tensor_model_parallel_rank,
     get_tensor_model_parallel_world_size,
 )
 from vllm.logger import init_logger
 from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
+from vllm.model_executor.model_loader.mtp_validation import (
+    is_mtp_completeness_check_enabled,
+)
 from vllm.model_executor.model_loader.reload import (
     support_quantized_model_reload_from_hp_weights,
 )
@@ -44,12 +47,21 @@
 ShardIds: TypeAlias = ShardId | list[ShardId]
 """One shard, or a list of shards the same weight is loaded into in turn."""
 
+WeightsTransformFn: TypeAlias = Callable[
+    [Iterable[tuple[str, torch.Tensor]]], Iterable[tuple[str, torch.Tensor]]
+]
+"""A lazy transform over a raw checkpoint ``(name, tensor)`` stream, passing
+unrecognized tensors through unchanged."""
+
 
 @dataclass
 class WeightsMapper:
     """Maps the name of each weight if they match the following patterns.
 
-    If a key maps to a value of `None`, the corresponding weight is ignored."""
+    If a key maps to a value of `None`, the corresponding weight is ignored.
+
+    ``weights_transforms`` rewrite the raw weight stream in list order before
+    any name mapping."""
 
     orig_to_new_renaming: list["WeightRenaming"] = field(default_factory=list)
     orig_to_new_regex: Mapping[re.Pattern, str | None] = field(default_factory=dict)
@@ -59,6 +71,7 @@ class WeightsMapper:
     )
     orig_to_new_prefix: Mapping[str, str | None] = field(default_factory=dict)
     orig_to_new_suffix: Mapping[str, str | None] = field(default_factory=dict)
+    weights_transforms: list[WeightsTransformFn] = field(default_factory=list)
 
     def __or__(self, other: "WeightsMapper") -> "WeightsMapper":
         """Combine two `WeightsMapper`s by merging their mappings."""
@@ -75,6 +88,7 @@ def __or__(self, other: "WeightsMapper") -> "WeightsMapper":
             },
             orig_to_new_prefix={**self.orig_to_new_prefix, **other.orig_to_new_prefix},
             orig_to_new_suffix={**self.orig_to_new_suffix, **other.orig_to_new_suffix},
+            weights_transforms=[*self.weights_transforms, *other.weights_transforms],
         )
 
     def map_name(self, key: str) -> str | None:
@@ -147,6 +161,8 @@ def _map_name_with_shard(self, key: str) -> tuple[str, ShardIds | None] | None:
     def apply(
         self, weights: Iterable[tuple[str, torch.Tensor]]
     ) -> Iterable[tuple[str, torch.Tensor]]:
+        for transform in self.weights_transforms:
+            weights = transform(weights)
         for name, data in weights:
             result = self._map_name_with_shard(name)
             if result is None:
@@ -196,6 +212,7 @@ def get_rename_mapper(self) -> "WeightsMapper":
             orig_to_new_stacked={},
             orig_to_new_prefix=remove_none(self.orig_to_new_prefix),
             orig_to_new_suffix=remove_none(self.orig_to_new_suffix),
+            weights_transforms=[],
         )
 
 
@@ -490,7 +507,10 @@ def _filter_skipped(
 
     def _check_skipped_aliases(self, autoloaded_weights: set[str]) -> None:
         """Guard against skipping an alias whose canonical name never loads."""
-        if not self._loaded_params_are_complete:
+        if (
+            not self._loaded_params_are_complete
+            or not is_mtp_completeness_check_enabled()
+        ):
             return
         for alias, canonical in self._skipped_aliases.items():
             if canonical not in autoloaded_weights:
@@ -502,13 +522,88 @@ def _check_skipped_aliases(self, autoloaded_weights: set[str]) -> None:
                 )
 
 
+def get_fused_shared_expert_layers(
+    module: nn.Module,
+    moe_cls: type[nn.Module],
+    enable_eplb: bool = False,
+) -> frozenset[int]:
+    """Resolve fused-shared-expert layer indices from the live module tree.
+
+    The index is taken from each MoE module's qualified name, so it follows the
+    same index space as the caller's weight names: 0-based for modules stored in
+    an ``nn.ModuleList``, and absolute for ``nn.ModuleDict`` MTP layers keyed by
+    layer index. No loader needs to hard-code an offset.
+
+    Also rejects mixed fusion when EPLB is enabled or the ``nixl_ep`` all2all
+    backend is selected, and logs the resolved policy.
+    """
+    fused_layers: set[int] = set()
+    moe_layers: set[int] = set()
+    for name, mod in module.named_modules():
+        if not isinstance(mod, moe_cls):
+            continue
+        idx = extract_layer_index(name)
+        moe_layers.add(idx)
+        if getattr(mod, "is_fused_shared_expert_enabled", False):
+            fused_layers.add(idx)
+
+    # EPLB bookkeeping is per rank/model: `EplbState` builds each model's
+    # `expert_weights`/`num_moe_layers` from the local `MixtureOfExperts` (one
+    # pipeline stage here), while main and draft models have separate
+    # `EplbModelState`s keyed by model config. The uniformity requirement is
+    # therefore stage/model-local by design: mixed within a stage is rejected,
+    # mixed across PP stages (or main vs draft) is allowed.
+    if 0 < len(fused_layers) < len(moe_layers):
+        if enable_eplb:
+            raise NotImplementedError(
+                "Mixed per-layer shared-expert fusion within a pipeline "
+                "stage/model is not supported with EPLB"
+            )
+        vllm_config = get_current_vllm_config_or_none()
+        parallel_config = (
+            vllm_config.parallel_config if vllm_config is not None else None
+        )
+        if parallel_config is not None and parallel_config.enable_expert_parallel:
+            raise NotImplementedError(
+                "Mixed per-layer shared-expert fusion with expert parallelism is "
+                "not tested. Please open an issue if you need this combination."
+            )
+        # These backends size one shared buffer from the first MoE layer's
+        # geometry (all2all.py NixlEPAll2AllManager._init_buffer/get_handle).
+        all2all_backend = (
+            parallel_config.all2all_backend if parallel_config is not None else None
+        )
+        if all2all_backend == "nixl_ep":
+            raise NotImplementedError(
+                "Mixed per-layer shared-expert fusion is not supported with "
+                f"all2all_backend={all2all_backend!r}"
+            )
+
+    if fused_layers:
+        logger.info(
+            "AITER fused shared experts: enabled on MoE layer(s) %s; "
+            "disabled on MoE layer(s) %s.",
+            sorted(fused_layers),
+            sorted(moe_layers - fused_layers) or "none",
+        )
+    elif moe_layers:
+        from vllm._aiter_ops import rocm_aiter_ops
+
+        if rocm_aiter_ops.is_fusion_moe_shared_experts_enabled():
+            logger.info(
+                "AITER fused shared experts: disabled on all %d MoE layer(s).",
+                len(moe_layers),
+            )
+    return frozenset(fused_layers)
+
+
 def maybe_fuse_shared_experts(
     weights: Iterable[tuple[str, torch.Tensor]],
     *,
     n_routed_experts: int,
     n_shared_experts: int,
+    fused_shared_expert_layers: frozenset[int],
     ckpt_prefix: str = "mlp.shared_experts",
-    enabled: bool | None = None,
 ) -> Iterable[tuple[str, torch.Tensor]]:
     """Route AITER fused-shared-expert checkpoint weights into fused slots.
 
@@ -524,21 +619,16 @@ def maybe_fuse_shared_experts(
         n_routed_experts: Number of routed experts; offsets the fused slots.
         n_shared_experts: Number of shared experts packed into the tensor.
         ckpt_prefix: Checkpoint module name of the shared experts.
-        enabled: Whether AITER fused-shared-experts is active. Defaults to
-            `rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()`; pass an
-            explicit value only when the model gates on something more (e.g.
-            quant-spec compatibility) and it must match its construction-time
-            decision.
+        fused_shared_expert_layers: Layer indices whose shared expert was fused
+            at construction time, as returned by
+            `get_fused_shared_expert_layers`. Only weights belonging to these
+            layers are routed to fused slots.
 
     Yields:
         `(name, tensor)` pairs with shared experts routed to fused slots.
 
     """
-    if enabled is None:
-        from vllm._aiter_ops import rocm_aiter_ops
-
-        enabled = rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
-    if not enabled:
+    if not fused_shared_expert_layers:
         yield from weights
         return
 
@@ -549,6 +639,9 @@ def maybe_fuse_shared_experts(
         if prefix not in name:
             yield name, loaded_weight
             continue
+        if extract_layer_index(name) not in fused_shared_expert_layers:
+            yield name, loaded_weight
+            continue
         # gate/up split on the output dim; down_proj on the input dim.
         split_dim = 1 if ("down_proj.weight" in name and loaded_weight.ndim > 1) else 0
         total = loaded_weight.shape[split_dim]
@@ -766,6 +859,10 @@ def __init__(self, stage_name: str, module: nn.Module | None = None) -> None:
     def __getattr__(self, name: str):
         return getattr(self.__dict__["module"], name)
 
+    def tie_weights(self, embed_tokens: nn.Module) -> "StageMissingLayer":
+        """Omitted layers do not need to share embedding weights."""
+        return self
+
     def __call__(self, *args, **kwargs):
         raise RuntimeError(f"{self} should not be called")
 
@@ -860,24 +957,17 @@ def forward(self, *args, **kwargs):
         return args[0] if args else next(iter(kwargs.values()))
 
 
-def spec_decode_needs_target_embed(
-    vllm_config: VllmConfig, *, include_mtp: bool = False
-) -> bool:
-    """Whether the last PP rank needs the target input embedding.
-
-    MTP targets opt in when their drafter supports sharing this embedding.
-    """
+def spec_decode_needs_target_embed(vllm_config: VllmConfig) -> bool:
+    """Whether the last PP rank needs the target input embedding, for a
+    hidden-state drafter that may reuse it (see `maybe_share_target_embed`)."""
     from vllm.distributed.parallel_state import get_pp_group
 
     speculative_config = vllm_config.speculative_config
-    if speculative_config is None:
-        return False
-    if speculative_config.method not in (
-        "eagle",
-        "eagle3",
-        "dflash",
-        "dspark",
-    ) and not (include_mtp and speculative_config.method == "mtp"):
+    if (
+        speculative_config is None
+        or not speculative_config.uses_target_hidden_states
+        or not speculative_config.draft_model_config.may_use_target_embed
+    ):
         return False
     pp = get_pp_group()
     return pp.world_size > 1 and pp.is_last_rank
@@ -952,7 +1042,93 @@ def is_pp_missing_parameter(name: str, model: torch.nn.Module) -> bool:
     )
 
 
-def make_empty_intermediate_tensors_factory(keys: list[str], hidden_size: int):
+def compute_skip_topk_layers(config: "PreTrainedConfig") -> set[int]:
+    """Backbone layers that reuse the top-k of the closest preceding indexer
+    layer instead of building their own indexer."""
+    if not hasattr(config, "index_topk"):
+        return set()
+
+    num_hidden_layers = config.num_hidden_layers
+    indexer_types = getattr(config, "indexer_types", None)
+    if indexer_types is not None:
+        if len(indexer_types) != num_hidden_layers:
+            raise ValueError(
+                "indexer_types must contain one entry per hidden layer: "
+                f"expected {num_hidden_layers}, got {len(indexer_types)}."
+            )
+        invalid_types = sorted(set(indexer_types) - {"full", "shared"})
+        if invalid_types:
+            raise ValueError(
+                f"indexer_types only supports 'full' and 'shared', got {invalid_types}."
+            )
+        return {
+            layer_idx
+            for layer_idx, indexer_type in enumerate(indexer_types)
+            if indexer_type == "shared"
+        }
+
+    freq = getattr(config, "index_topk_freq", 1)
+    if not isinstance(freq, int) or freq <= 0:
+        raise ValueError(f"index_topk_freq must be a positive integer, got {freq!r}.")
+    pattern = getattr(config, "index_topk_pattern", None)
+    offset = getattr(config, "index_skip_topk_offset", 2)
+    skip_layers: set[int] = set()
+    for layer_idx in range(num_hidden_layers):
+        if pattern is None:
+            if max(layer_idx - offset + 1, 0) % freq != 0:
+                skip_layers.add(layer_idx)
+        elif 0 <= layer_idx < len(pattern) and pattern[layer_idx] == "S":
+            skip_layers.add(layer_idx)
+    return skip_layers
+
+
+def pp_splits_layer_sharing(
+    num_hidden_layers: int, users: Iterable[int], consumers: set[int]
+) -> bool:
+    """Whether a later PP stage's first ``users`` layer is one of ``consumers``,
+    which read state an earlier layer wrote in the same forward pass."""
+    from vllm.distributed.parallel_state import get_pp_group
+    from vllm.distributed.utils import get_pp_indices
+
+    users = sorted(users)
+    pp_size = get_pp_group().world_size
+    for rank in range(1, pp_size):
+        start, _ = get_pp_indices(num_hidden_layers, rank, pp_size)
+        if next((layer for layer in users if layer >= start), None) in consumers:
+            return True
+    return False
+
+
+def pp_topk_buffers(
+    config: "PreTrainedConfig",
+    topk_indices_buffer: torch.Tensor,
+    topk_layers: Iterable[int],
+) -> dict[str, torch.Tensor]:
+    """The shared top-k, if a later PP stage starts on a layer that reuses it."""
+    topk_layers = set(topk_layers)
+    skip_layers = compute_skip_topk_layers(config) & topk_layers
+    if pp_splits_layer_sharing(config.num_hidden_layers, topk_layers, skip_layers):
+        return {"topk_indices": topk_indices_buffer}
+    return {}
+
+
+def copy_pp_buffers(
+    intermediate_tensors: IntermediateTensors, buffers: Mapping[str, torch.Tensor]
+) -> None:
+    """Copy cross-layer state received from the previous PP stage into place."""
+    for key, buffer in buffers.items():
+        tensor = intermediate_tensors[key]
+        buffer[: tensor.shape[0]].copy_(tensor)
+
+
+def make_empty_intermediate_tensors_factory(
+    keys: list[str],
+    hidden_size: int,
+    buffers: Mapping[str, torch.Tensor] | None = None,
+):
+    """``buffers`` are token-major cross-layer state sent along with ``keys``."""
+    buffers = buffers or {}
+
     def make_empty_intermediate_tensors(
         batch_size: int,
         dtype: torch.dtype,
@@ -963,6 +1139,12 @@ def make_empty_intermediate_tensors(
                 key: torch.zeros((batch_size, hidden_size), dtype=dtype, device=device)
                 for key in keys
             }
+            | {
+                key: torch.zeros(
+                    (batch_size, *buffer.shape[1:]), dtype=buffer.dtype, device=device
+                )
+                for key, buffer in buffers.items()
+            }
         )
 
     return make_empty_intermediate_tensors
diff --git a/vllm/model_executor/models/whisper.py b/vllm/model_executor/models/whisper.py
index 9e5d74cba29b..fd024da88b45 100644
--- a/vllm/model_executor/models/whisper.py
+++ b/vllm/model_executor/models/whisper.py
@@ -823,6 +823,13 @@ class WhisperForConditionalGeneration(
             ".encoder_attn.v_proj": (".encoder_attn.kv_proj", 1),
         },
         orig_to_new_prefix={"proj_out.": None},
+        # Add fake zeros bias to k_proj for qkv packed
+        weights_transforms=[
+            lambda weights: _create_fake_bias_for_k_proj(weights, ".k_proj.weight"),
+            lambda weights: _create_fake_bias_for_k_proj(
+                weights, ".k_proj.weight_packed"
+            ),
+        ],
     )
 
     # Whisper only supports audio-conditioned generation.
@@ -1038,15 +1045,6 @@ def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor:
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         loader = AutoWeightsLoader(self)
-
-        # add fake zeros bias for k_proj to state_dict. Covers the
-        # ``.k_proj.weight`` and the compressed-tensors packed
-        # ``.k_proj.weight_packed`` checkpoints.
-        k_proj_weight_suffixes = (".k_proj.weight", ".k_proj.weight_packed")
-        for suffix in k_proj_weight_suffixes:
-            weights = _create_fake_bias_for_k_proj(
-                weights, suffix, out_features=self.config.d_model
-            )
         return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
 
 
diff --git a/vllm/model_executor/offloader/prefetch.py b/vllm/model_executor/offloader/prefetch.py
index b87535ac376e..98cc0f06f49f 100644
--- a/vllm/model_executor/offloader/prefetch.py
+++ b/vllm/model_executor/offloader/prefetch.py
@@ -467,7 +467,7 @@ def sync_cpu_storage(self):
             param_offloader.sync_cpu_storage()
 
         # Remove offloaders whose parameter was deleted during
-        # process_weights_after_loading (e.g. k_scale / v_scale).
+        # process_weights_after_loading (e.g. online requantization sources).
         deleted = [
             name
             for name, offloader in self._param_offloaders.items()
@@ -639,9 +639,8 @@ def __init__(self, module: nn.Module, param_name: str):
         self._cpu_storage: torch.Tensor | None = None
         self._gpu_buffer: torch.Tensor | None = None  # Store reference to GPU buffer
         # Set to True if the underlying nn.Parameter was deleted by
-        # process_weights_after_loading (e.g. transient KV-cache scale params
-        # such as k_scale/v_scale created by BaseKVCacheMethod.create_weights
-        # and deleted after copying into permanent _k_scale buffers).
+        # process_weights_after_loading (e.g. online requantization source
+        # params released after conversion).
         self._param_deleted: bool = False
 
         # Offload to CPU immediately to free GPU memory during model loading
@@ -750,12 +749,10 @@ def sync_cpu_storage(self) -> None:
         2. device_loading_context creates NEW CPU tensors when moving back
         3. Our old _cpu_storage would have pre-processed or stale data
 
-        If the parameter no longer exists on the module (e.g. transient
-        KV-cache scale parameters such as k_scale/v_scale that are created
-        by BaseKVCacheMethod.create_weights() and then deleted by
-        process_weights_after_loading() after copying their values into
-        permanent _k_scale buffers), the offloader marks itself as deleted
-        and skips the sync.  The caller (_ModuleOffloader.sync_cpu_storage)
+        If the parameter no longer exists on the module (e.g. online
+        requantization source parameters released by
+        process_weights_after_loading()), the offloader marks itself as
+        deleted and skips the sync.  The caller (_ModuleOffloader.sync_cpu_storage)
         is responsible for removing these stale entries.
         """
         try:
diff --git a/vllm/model_executor/utils.py b/vllm/model_executor/utils.py
index e73e4a7ec569..cc572b360e49 100644
--- a/vllm/model_executor/utils.py
+++ b/vllm/model_executor/utils.py
@@ -68,6 +68,65 @@ def set_weight_attrs(
         setattr(weight, key, value)
 
 
+def _extent(t: torch.Tensor) -> int:
+    """Number of elements spanned by ``t`` in its storage."""
+    if t.numel() == 0:
+        return 0
+    return 1 + sum((n - 1) * s for n, s in zip(t.shape, t.stride()))
+
+
+def storage_view(
+    storage: torch.UntypedStorage,
+    offset: int,
+    shape: tuple[int, ...],
+    stride: tuple[int, ...],
+    dtype: torch.dtype,
+) -> torch.Tensor:
+    """A ``dtype`` tensor on ``storage`` at element ``offset``. ``set_`` does
+    not bounds-check: callers check that the view fits."""
+    view = torch.empty(0, dtype=dtype, device=storage.device)
+    return view.set_(storage, offset, shape, stride)
+
+
+def _in_own_storage(
+    old: torch.Tensor, new: torch.Tensor, whole_storage: bool = False
+) -> torch.Tensor:
+    """``new`` copied into the bytes ``old`` spans if it fits, so the loaded
+    storage is kept (with its ``data_ptr``); else ``new`` itself. An equal-size
+    result fits too: a repack to the same byte count (e.g. FP8 → Marlin int32)
+    then keeps the address captured CUDA graphs read.
+
+    With ``whole_storage`` (direct reload), ``new`` may use the rest of
+    ``old``'s storage, so a weight that post-load processing grew keeps its
+    address on a reload too, as do a padded view with gaps (ROCm FP8 padding)
+    and an emptied weight (the checkpoint bytes stay for the next reload)."""
+    storage = old.untyped_storage()
+    offset, size = old.storage_offset() * old.element_size(), new.element_size()
+    span = _extent(new)
+    if whole_storage:
+        fits = (
+            span >= new.numel()  # not overlapping
+            and offset + span * size <= storage.nbytes()
+        )
+    else:
+        fits = (
+            new.numel() > 0
+            and span == new.numel()  # neither overlapping nor with gaps
+            # its gaps may hold another tensor's data
+            and _extent(old) == old.numel()
+            and new.numel() * size <= old.numel() * old.element_size()
+        )
+    if (
+        not fits
+        or old.device != new.device
+        or new.untyped_storage().data_ptr() == storage.data_ptr()
+        or offset % size
+    ):
+        return new
+    view = storage_view(storage, offset // size, new.shape, new.stride(), new.dtype)
+    return view.copy_(new)
+
+
 def replace_parameter(
     layer: torch.nn.Module,
     param_name: str,
@@ -79,7 +138,8 @@ def replace_parameter(
     The existing parameter object is kept (and so its class and loader attributes
     for weight reloading): ``new_tensor`` is copied into it when shape, stride,
     dtype and device match, which preserves ``data_ptr`` for captured CUDA graphs;
-    otherwise the parameter's ``.data`` is rebound to ``new_tensor``. A new
+    otherwise a smaller ``new_tensor`` is copied into the parameter's own storage
+    (so the loaded storage never shrinks), else ``.data`` is rebound to it. A new
     parameter is registered only if none exists under ``param_name``, or if the
     existing one is an online requantization source (released after conversion).
 
@@ -87,6 +147,10 @@ def replace_parameter(
     ``is_shuffled``) are carried over to the parameter, except ``weight_loader``,
     which is always taken from the existing parameter.
 
+    Under direct reload, a ``new_tensor`` larger than the parameter at cold start
+    may be written into the parameter's existing storage on a reload, so do not
+    read a local alias of the old data (``w = layer.weight.data``) after this call.
+
     This function should not be called on weights which are tied/shared
 
     Args:
@@ -113,6 +177,9 @@ def replace_parameter(
     # a loader riding along under that name has to be dropped as well.
     new_tensor_attrs.pop("weight_loader", None)
     new_tensor_attrs.pop("_weight_loader", None)
+    # A checkpoint record belongs to the parameter it was recorded on (direct
+    # reload); one on new_tensor (another checkpoint parameter) must not overwrite it.
+    new_tensor_attrs.pop("_ckpt", None)
 
     if isinstance(new_tensor, torch.nn.Parameter):
         new_tensor = new_tensor.data
@@ -131,8 +198,16 @@ def replace_parameter(
             and old_param.device == new_tensor.device
         ):
             old_param.data.copy_(new_tensor)
+        elif getattr(getattr(layer, "quant_method", None), "uses_meta_device", False):
+            old_param.data = new_tensor  # online quantization frees the source
         else:
-            old_param.data = new_tensor
+            from vllm.model_executor.model_loader.reload.direct import (
+                direct_reload_enabled,
+            )
+
+            old_param.data = _in_own_storage(
+                old_param, new_tensor, whole_storage=direct_reload_enabled()
+            )
         for attr_name, attr in new_tensor_attrs.items():
             setattr(old_param, attr_name, attr)
         return
@@ -149,6 +224,74 @@ def replace_parameter(
     setattr(layer, param_name, new_param)
 
 
+def move_parameter(layer: torch.nn.Module, src_name: str, dst_name: str) -> None:
+    """Register the ``src_name`` parameter object under ``dst_name`` as well.
+
+    Both names then refer to one object and so to one storage: the loader writes
+    through the checkpoint name, kernels read the runtime name, and a later
+    `replace_parameter` on either name updates both. ``dst_name`` becomes the
+    name `kernel_parameters` reports for the object."""
+    param = getattr(layer, src_name)
+    existing = layer._parameters.get(dst_name)
+    if existing is not None and existing is not param:
+        raise ValueError(f"{dst_name!r} is already registered as another parameter")
+    layer.register_parameter(dst_name, param)
+    param.runtime_name = dst_name
+
+
+def retire_parameter(layer: torch.nn.Module, name: str) -> None:
+    """Mark a loaded parameter that no kernel reads after post-processing.
+
+    It stays registered with its loaded data (so it can be loaded again);
+    `kernel_parameters` and `get_live_param` skip it."""
+    getattr(layer, name).retired = True
+
+
+def set_kernel_parameter(
+    layer: torch.nn.Module, name: str, tensor: torch.Tensor
+) -> None:
+    """`replace_parameter`, for a tensor kernels read under ``name``: the parameter
+    is live even if it was retired before, and ``name`` becomes its runtime name
+    if it is registered under several names (see `move_parameter`)."""
+    replace_parameter(layer, name, tensor)
+    param = getattr(layer, name)
+    param.__dict__.pop("retired", None)
+    if hasattr(param, "runtime_name"):
+        param.runtime_name = name
+
+
+def _is_live(param: torch.Tensor | None) -> bool:
+    return (
+        param is not None and not getattr(param, "retired", False) and param.numel() > 0
+    )
+
+
+def get_live_param(layer: torch.nn.Module, name: str) -> torch.Tensor | None:
+    """``getattr(layer, name, None)``, treating retired or empty params as absent."""
+    tensor = getattr(layer, name, None)
+    if isinstance(tensor, torch.nn.Parameter) and not _is_live(tensor):
+        return None
+    return tensor
+
+
+def kernel_parameters(layer: torch.nn.Module) -> dict[str, torch.nn.Parameter]:
+    """The layer's own parameters as kernels see them: live parameters only, each
+    object once, under its runtime name (see `move_parameter`)."""
+    names: dict[int, str] = {}
+    params: dict[str, torch.nn.Parameter] = {}
+    for name, param in layer._parameters.items():
+        if not _is_live(param):
+            continue
+        first = names.setdefault(id(param), name)
+        if first == name:
+            params[name] = param
+        elif getattr(param, "runtime_name", None) == name:
+            del params[first]
+            params[name] = param
+            names[id(param)] = name
+    return params
+
+
 def get_packed_modules_mapping(model: torch.nn.Module) -> dict[str, list[str]]:
     parent_map = getattr(model, "packed_modules_mapping", None)
     parent_map = copy.deepcopy(parent_map) if parent_map is not None else {}
diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py
index a51f3bb2373c..f7171219bfd7 100644
--- a/vllm/model_executor/warmup/kernel_warmup.py
+++ b/vllm/model_executor/warmup/kernel_warmup.py
@@ -59,6 +59,46 @@
 _LL_BF16_WARMUP_M_RANGE = range(1, 17)
 
 
+def _warmup_flashinfer_attention(worker: "Worker") -> None:
+    from vllm.distributed.parallel_state import get_pp_group
+
+    # Only warmup if the model has FlashInfer attention groups
+    # and is not a pooling model
+    def _is_flashinfer_backend(backend):
+        try:
+            return backend.get_name() == "FLASHINFER"
+        except NotImplementedError:
+            return False
+
+    check = any if worker.use_v2_model_runner else all
+    local_flashinfer_warmup = (
+        not worker.model_runner.is_pooling_model
+        and bool(worker.model_runner.attn_groups)
+        and check(
+            _is_flashinfer_backend(group.backend)
+            for groups in worker.model_runner.attn_groups
+            for group in groups
+        )
+    )
+    # The mixed warmup uses worker.execute_model, which participates in PP
+    # send/recv operations. Every PP rank must make the same warmup decision.
+    pp_group = get_pp_group()
+    eligible_pp_ranks = int(local_flashinfer_warmup)
+    if pp_group.world_size > 1:
+        eligibility = torch.tensor([eligible_pp_ranks], dtype=torch.int32, device="cpu")
+        pp_group.cpu_all_reduce_(eligibility)
+        eligible_pp_ranks = int(eligibility.item())
+
+    if worker.use_v2_model_runner:
+        should_warmup = eligible_pp_ranks > 0
+    else:
+        should_warmup = eligible_pp_ranks == pp_group.world_size
+
+    if should_warmup:
+        logger.info_once("Warming up FlashInfer attention.")
+        mixed_batch_attention_warmup(worker)
+
+
 def _ll_bf16_router_shapes_from_model(
     model: torch.nn.Module,
 ) -> tuple[tuple[int, int], ...]:
@@ -104,6 +144,27 @@ def _warmup_ll_bf16_router_gemm(model: torch.nn.Module) -> None:
     )
 
 
+def _warmup_skinny_gemm(model: torch.nn.Module) -> None:
+    from vllm.model_executor.kernels.linear.cute_dsl.skinny_gemm import (
+        SkinnyGemmConfig,
+        shape_dynamic_skinny_gemm,
+    )
+    from vllm.model_executor.kernels.linear.cute_dsl.skinny_gemm_plans import (
+        get_skinny_gemm_plan,
+    )
+
+    configs: set[SkinnyGemmConfig] = set()
+    for module in model.modules():
+        if not getattr(module, "use_skinny_gemm", False):
+            continue
+        plan = get_skinny_gemm_plan(*module.weight.shape)
+        if plan is not None:
+            configs.update(plan.values())
+    if configs:
+        # Compiled by cutedsl_warmup() below.
+        shape_dynamic_skinny_gemm.request_warmup_configs(torch.bfloat16, configs)
+
+
 def _warmup_bf16x3_router_gemm(
     model: torch.nn.Module,
     max_num_tokens: int,
@@ -162,9 +223,6 @@ def _autotune_kimi_k3_kda_qkvg(model: torch.nn.Module) -> None:
 
 
 def kernel_warmup(worker: "Worker", *, process_local_only: bool = False):
-    from vllm.model_executor.warmup.minimax_m3_msa_warmup import (
-        minimax_m3_msa_warmup,
-    )
     from vllm.model_executor.warmup.step5_dsa_warmup import step5_dsa_warmup
 
     if not worker.use_v2_model_runner:
@@ -222,6 +280,7 @@ def kernel_warmup(worker: "Worker", *, process_local_only: bool = False):
         # TODO(roberto): Remove after registered CuTeDSL warmups are migrated
         # to the shared JIT warmup infrastructure.
         # https://github.com/vllm-project/vllm/pull/47451
+        _warmup_skinny_gemm(worker.get_model())
         cutedsl_warmup()
 
     if process_local_only:
@@ -249,7 +308,13 @@ def kernel_warmup(worker: "Worker", *, process_local_only: bool = False):
 
     b12x_warmup(worker, cudagraph_capture_sizes)
 
-    minimax_m3_msa_warmup(worker)
+    if current_platform.is_cuda():
+        # The MSA kernels, and the CuTe DSL their modules import, are CUDA-only.
+        from vllm.model_executor.warmup.minimax_m3_msa_warmup import (
+            minimax_m3_msa_warmup,
+        )
+
+        minimax_m3_msa_warmup(worker)
 
     # Allocate the exact decode-sized workspace, autotune cache misses, and
     # resolve every CUDA Graph bucket before capture begins.
@@ -271,29 +336,7 @@ def kernel_warmup(worker: "Worker", *, process_local_only: bool = False):
     elif has_flashinfer() and current_platform.has_device_capability(90):
         flashinfer_autotune(worker.model_runner)
 
-    # FlashInfer attention warmup
-    # Only warmup if the model has FlashInfer attention groups
-    # and is not a pooling model
-    def _is_flashinfer_backend(backend):
-        try:
-            return backend.get_name() == "FLASHINFER"
-        except NotImplementedError:
-            return False
-
-    if (
-        not worker.model_runner.is_pooling_model
-        and worker.model_runner.attn_groups
-        # NOTE: This should be `any` instead of `all` but other hybrid attention
-        # backends don't support this dummy run. Once we remove
-        # `build_for_cudagraph_capture`, we can change it to `any`.
-        and all(
-            _is_flashinfer_backend(group.backend)
-            for groups in worker.model_runner.attn_groups
-            for group in groups
-        )
-    ):
-        logger.info_once("Warming up FlashInfer attention.")
-        mixed_batch_attention_warmup(worker)
+    _warmup_flashinfer_attention(worker)
 
 
 def _flashinfer_autotune_skip_ops(runner: "GPUModelRunner") -> set[str] | None:
@@ -304,6 +347,11 @@ def _flashinfer_autotune_skip_ops(runner: "GPUModelRunner") -> set[str] | None:
         FlashInferCuteDslNvFp4LinearKernel,
     )
 
+    skip_ops: set[str] = set()
+    if runner.cache_config.enable_extensible_kv_cache:
+        # The trtllm-gen MLA decode sweep reads arbitrary block ids, which
+        # an extensible KV cache has not committed yet.
+        skip_ops.add("trtllm_batch_decode_mla")
     for module in runner.get_model().modules():
         for holder_name in ("quant_method", "scheme"):
             kernel = getattr(getattr(module, holder_name, None), "kernel", None)
@@ -311,8 +359,8 @@ def _flashinfer_autotune_skip_ops(runner: "GPUModelRunner") -> set[str] | None:
             # fallback is already the heuristic; all mm_fp4 backends share
             # the "fp4_gemm" op name, so skip only when cute-dsl is selected.
             if isinstance(kernel, FlashInferCuteDslNvFp4LinearKernel):
-                return {"fp4_gemm"}
-    return None
+                skip_ops.add("fp4_gemm")
+    return skip_ops or None
 
 
 def _flashinfer_deferred_moe_token_counts(
@@ -347,6 +395,7 @@ def _run_flashinfer_autotune_dummy_runs(
     runner: "GPUModelRunner", *, skip_attn: bool = False
 ) -> None:
     import vllm.utils.flashinfer as fi_utils
+    from vllm.models.deepseek_v41.attention import autotune_indexer_wq_b
 
     dummy_run_kwargs = {"skip_attn": True} if skip_attn else {}
     for num_tokens in _flashinfer_autotune_token_counts(runner):
@@ -383,6 +432,7 @@ def _run_flashinfer_autotune_dummy_runs(
                 randomize_inputs=True,
                 **dummy_run_kwargs,
             )
+            autotune_indexer_wq_b(runner.get_model(), max_tuning_tokens)
 
 
 def _autotune_cache_fingerprint(path: Path) -> tuple[str, int] | None:
diff --git a/vllm/model_executor/warmup/minimax_m3_msa_warmup.py b/vllm/model_executor/warmup/minimax_m3_msa_warmup.py
index 2c8996a93341..699cad10af08 100644
--- a/vllm/model_executor/warmup/minimax_m3_msa_warmup.py
+++ b/vllm/model_executor/warmup/minimax_m3_msa_warmup.py
@@ -1,21 +1,83 @@
 # 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
 
 from vllm.logger import init_logger
 from vllm.model_executor.warmup.attention_warmup import (
     mixed_batch_attention_warmup,
 )
+from vllm.models.minimax_m3.nvidia.indexer_msa import MiniMaxM3IndexerMSAImpl
+from vllm.models.minimax_m3.nvidia.indexer_nvfp4 import MiniMaxM3IndexerNvfp4Impl
 from vllm.models.minimax_m3.nvidia.model import MiniMaxM3SparseAttention
+from vllm.models.minimax_m3.nvidia.msa_cutlass_sparse_decode import (
+    NVFP4_BLOCK_SCALE_SHIFT,
+    uses_q8kv4_decode,
+)
+from vllm.models.minimax_m3.nvidia.sparse_attention_msa import (
+    MiniMaxM3SparseMSABaseImpl,
+)
 from vllm.platforms import current_platform
 from vllm.tracing import instrument
 
 if TYPE_CHECKING:
+    from vllm.config import VllmConfig
     from vllm.v1.worker.gpu_worker import Worker
 
 logger = init_logger(__name__)
 
+# KV-cache dtypes fmha_sm100.warmup plans attention kernels for.
+_MSA_KV_CACHE_DTYPES = ("auto", "bfloat16", "fp8", "fp8_e4m3", "nvfp4")
+
+
+def _msa_prebuild_kwargs(
+    vllm_config: "VllmConfig", module: MiniMaxM3SparseAttention
+) -> dict[str, Any] | None:
+    """``fmha_sm100.warmup`` arguments for the kernels this rank can request."""
+    impl = module.impl
+    indexer = module.indexer.impl
+    msa = (
+        impl
+        if isinstance(impl, MiniMaxM3SparseMSABaseImpl)
+        and module.kv_cache_dtype in _MSA_KV_CACHE_DTYPES
+        else None
+    )
+    attention = msa is not None
+    if isinstance(indexer, MiniMaxM3IndexerNvfp4Impl):
+        index_cache_dtype = "nvfp4"
+    elif isinstance(indexer, MiniMaxM3IndexerMSAImpl):
+        index_cache_dtype = module.indexer_kv_dtype
+    else:
+        index_cache_dtype = None
+    if not attention and index_cache_dtype is None:
+        return None
+
+    spec_config = vllm_config.speculative_config
+    num_spec = spec_config.num_speculative_tokens if spec_config is not None else 0
+    q8kv4_decode = attention and uses_q8kv4_decode(
+        kv_cache_dtype=module.kv_cache_dtype,
+        num_q_heads=module.num_heads,
+        num_kv_heads=module.num_kv_heads,
+    )
+    return dict(
+        kv_cache_dtype=module.kv_cache_dtype if attention else "auto",
+        index_cache_dtype=index_cache_dtype,
+        num_q_heads=module.num_heads,
+        num_kv_heads=module.num_kv_heads,
+        topk=impl.topk_blocks,
+        decode_query_lens=range(1, num_spec + 2),
+        num_index_heads=indexer.num_index_heads,
+        page_size=impl.block_size,
+        decode_backend="q8kv4" if q8kv4_decode else "kv_mode3",
+        block_scale_shift=NVFP4_BLOCK_SCALE_SHIFT,
+        sparse_decode=msa is not None and msa.use_cutlass_decode,
+        sparse_prefill=attention,
+        prefill_backend=msa.prefill_backend if msa is not None else "cute_dsl",
+        device=torch.device("cuda", torch.accelerator.current_device_index()),
+    )
+
 
 @instrument(span_name="MiniMax M3 MSA warmup")
 def minimax_m3_msa_warmup(worker: "Worker") -> None:
@@ -36,5 +98,14 @@ def minimax_m3_msa_warmup(worker: "Worker") -> None:
 
     logger.info("Warming up MiniMax M3 MSA kernels.")
 
-    # Cover sparse prefill through the normal model path.
+    # Build every native kernel the configuration can request at once instead
+    # of compiling each one serially on first use.
+    kwargs = _msa_prebuild_kwargs(worker.vllm_config, sparse_module)
+    if kwargs is not None:
+        from vllm.third_party.fmha_sm100 import warmup
+
+        logger.info("MiniMax M3 MSA kernels: %s", warmup(**kwargs))
+
+    # Cover sparse prefill (and its CuTe-DSL kernels) through the normal model
+    # path.
     mixed_batch_attention_warmup(worker)
diff --git a/vllm/model_executor/warmup/qwen4_exp_qsa_warmup.py b/vllm/model_executor/warmup/qwen4_exp_qsa_warmup.py
index 2065b3833822..e5126a766635 100644
--- a/vllm/model_executor/warmup/qwen4_exp_qsa_warmup.py
+++ b/vllm/model_executor/warmup/qwen4_exp_qsa_warmup.py
@@ -19,7 +19,7 @@
 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")
+    qsa_module = sys.modules.get("vllm.models.qwen4_exp.nvidia.qsa_indexer")
     attn_module = sys.modules.get("vllm.models.qwen4_exp.nvidia.qsa")
     if qsa_module is None or attn_module is None:
         return
diff --git a/vllm/model_executor/warmup/qwen_triton_warmup.py b/vllm/model_executor/warmup/qwen_triton_warmup.py
index dc512eee02d9..00a76899d7f1 100644
--- a/vllm/model_executor/warmup/qwen_triton_warmup.py
+++ b/vllm/model_executor/warmup/qwen_triton_warmup.py
@@ -49,6 +49,8 @@ class _QwenGDNWarmupConfig:
     state_stride_token: int
     state_dtype: torch.dtype
     fused_spec_decode: bool = False
+    # The FlashInfer prefill takes exp(g) straight from fused_post_conv_prep.
+    post_conv_g_exp: bool = False
 
     @property
     def conv_dim(self) -> int:
@@ -141,6 +143,7 @@ def _qwen_gdn_warmup_config(
             fused_spec_decode=bool(
                 getattr(layer, "enable_fused_gdn_spec_decode", False)
             ),
+            post_conv_g_exp=getattr(layer, "gdn_prefill_backend", None) == "flashinfer",
         )
 
     if found_layer:
@@ -235,7 +238,7 @@ def _warm_fused_post_conv_kernel(
             config.k,
             config.v,
             apply_l2norm=True,
-            output_g_exp=False,
+            output_g_exp=config.post_conv_g_exp,
         )
 
 
diff --git a/vllm/model_executor/warmup/replayssm_warmup.py b/vllm/model_executor/warmup/replayssm_warmup.py
index 28c45804ca7e..094bca938b9a 100644
--- a/vllm/model_executor/warmup/replayssm_warmup.py
+++ b/vllm/model_executor/warmup/replayssm_warmup.py
@@ -36,16 +36,19 @@ def _replayssm_autotune_kwargs(
             "is unavailable."
         )
         return None
+    from vllm.v1.worker.extensible_kv_cache import num_committable_kv_blocks
+
     v2_runner: Any = runner
     query_len = (
         v2_runner.decode_query_len
         if config.use_v2_model_runner
         else runner.uniform_decode_query_len
     )
+    # Block 0 is the null block; the dummy slots must be committed memory.
     max_num_reqs = min(
         runner.scheduler_config.max_num_seqs,
         runner.max_num_tokens // query_len,
-        runner.kv_cache_config.num_blocks - 1,
+        num_committable_kv_blocks(v2_runner) - 1,
     )
     if max_num_reqs <= 0:
         logger.warning_once(
@@ -168,6 +171,8 @@ def _temporary_replayssm_autotune_state(
 
 
 def replayssm_autotune_warmup(runner: "GPUModelRunner") -> None:
+    from vllm.v1.worker.extensible_kv_cache import ensure_kv_cache_blocks
+
     max_token_prefill_kwargs = {
         "num_tokens": runner.scheduler_config.max_num_batched_tokens,
         "skip_eplb": True,
@@ -178,5 +183,7 @@ def replayssm_autotune_warmup(runner: "GPUModelRunner") -> None:
     if autotune is None:
         return
     max_num_reqs, decode_kwargs = autotune
+    v2_runner: Any = runner
+    ensure_kv_cache_blocks(v2_runner, 1 + max_num_reqs)
     with _temporary_replayssm_autotune_state(runner, max_num_reqs):
         runner._dummy_run(**decode_kwargs)
diff --git a/vllm/model_executor/warmup/step5_dsa_warmup.py b/vllm/model_executor/warmup/step5_dsa_warmup.py
index 658caeaf7260..5837c2544e28 100644
--- a/vllm/model_executor/warmup/step5_dsa_warmup.py
+++ b/vllm/model_executor/warmup/step5_dsa_warmup.py
@@ -11,6 +11,7 @@
 from typing import TYPE_CHECKING, cast
 
 from vllm.logger import init_logger
+from vllm.platforms import current_platform
 from vllm.v1.worker.gpu.warmup import run_mixed_prefill_decode_warmup
 
 if TYPE_CHECKING:
@@ -32,13 +33,32 @@ def _step5_dsa_impl(worker: "Worker"):
 
 
 def step5_dsa_warmup(worker: "Worker") -> None:
+    if worker.vllm_config.model_config.hf_text_config.model_type != "step5":
+        return
+    if current_platform.is_rocm():
+        from vllm.models.step5.amd.warmup import step5_csa_warmup
+
+        step5_csa_warmup(worker)
+        return
     impl = _step5_dsa_impl(worker)
     if impl is None or impl.adapter is None:
         return
+    if not run_step5_scored_prefill_warmup(worker, impl.sparse_topk):
+        return
+    impl.adapter.warmup_prefill_scoring(
+        impl.max_num_regions,
+        impl.sparse_topk,
+        impl.indexer_heads if impl.indexer_heads * 2 == 8 else 0,
+        worker.device,
+    )
+
+
+def run_step5_scored_prefill_warmup(worker: "Worker", topk: int) -> bool:
+    """Run one mixed step whose prefill row scores more than ``topk`` regions."""
     runner = worker.model_runner
     # One region past "select every region" (top-k x 8 tokens), plus the
     # decode token of the mixed step.
-    num_tokens = impl.sparse_topk * 8 + 17
+    num_tokens = topk * 8 + 17
     if num_tokens > worker.scheduler_config.max_num_batched_tokens or (
         num_tokens > runner.max_model_len
     ):
@@ -46,7 +66,7 @@ def step5_dsa_warmup(worker: "Worker") -> None:
             "Skipping Step5 DSA scored-prefill warmup: needs %d tokens per step.",
             num_tokens,
         )
-        return
+        return False
     logger.info("Warming up Step5 DSA scored prefill with %d tokens.", num_tokens)
     if runner.max_num_reqs >= 2:
         run_mixed_prefill_decode_warmup(
@@ -58,9 +78,4 @@ def step5_dsa_warmup(worker: "Worker") -> None:
         )
     else:
         runner._dummy_run(num_tokens=num_tokens, skip_eplb=True, is_profile=True)
-    impl.adapter.warmup_prefill_scoring(
-        impl.max_num_regions,
-        impl.sparse_topk,
-        impl.indexer_heads if impl.indexer_heads * 2 == 8 else 0,
-        worker.device,
-    )
+    return True
diff --git a/vllm/models/beam/nvidia/dflash2.py b/vllm/models/beam/nvidia/dflash2.py
index cf9874ab72be..86f31c2bd5c9 100644
--- a/vllm/models/beam/nvidia/dflash2.py
+++ b/vllm/models/beam/nvidia/dflash2.py
@@ -44,6 +44,7 @@
     WeightsMapper,
     maybe_prefix,
 )
+from vllm.v1.worker.gpu.spec_decode.dflash2.config import DFlash2SpeculatorConfig
 
 from .dflash2_ops import candidate_edges
 
@@ -540,6 +541,9 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
 
 
 class DFlash2BeamForCausalLM(nn.Module):
+    # The speculator this is a draft model of (see SpeculativeDraftModel).
+    speculator_config_class = DFlash2SpeculatorConfig
+
     # The checkpoint ships neither, so the speculator always aliases the target's.
     has_own_embed_tokens = False
     has_own_lm_head = False
@@ -618,6 +622,8 @@ def compute_candidates(
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         loaded = {f"model.{name}" for name in self.model.load_weights(weights)}
-        self.model.build_context_kv_weights()
         # Aliased to the target's after loading.
         return loaded | {"model.embed_tokens.weight", "lm_head.weight"}
+
+    def process_weights_after_loading(self) -> None:
+        self.model.build_context_kv_weights()
diff --git a/vllm/models/common/ops/moe_tail_fusion/__init__.py b/vllm/models/common/ops/moe_tail_fusion/__init__.py
new file mode 100644
index 000000000000..3a391caf17df
--- /dev/null
+++ b/vllm/models/common/ops/moe_tail_fusion/__init__.py
@@ -0,0 +1,5 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Building blocks shared by the CuTe DSL kernels that fuse a MoE layer's tail
+(top-k finalize, TP all-reduce, norms) across models."""
diff --git a/vllm/models/common/ops/moe_tail_fusion/primitives.py b/vllm/models/common/ops/moe_tail_fusion/primitives.py
new file mode 100644
index 000000000000..53a2384bf265
--- /dev/null
+++ b/vllm/models/common/ops/moe_tail_fusion/primitives.py
@@ -0,0 +1,671 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Adapted from FlashInfer's flashinfer/comm/mnnvl_cutedsl/cute_dsl_primitives.py
+# and runtime.py (flashinfer-ai/flashinfer@139af6f6).
+
+"""CuTe DSL primitives shared by the fused MoE-tail kernels: PTX loads and
+stores, NVLS multicast stores, Lamport mailbox helpers, BF16 packing, block
+reductions and launch helpers. This module defines no kernel."""
+
+import cuda.bindings.driver as cuda
+import cutlass
+import cutlass.cute as cute
+import torch
+from cutlass import BFloat16, Float32, Int32, Int64, Uint16, Uint32
+from cutlass._mlir import ir
+from cutlass._mlir.dialects import llvm, vector
+from cutlass.cute.runtime import from_dlpack, make_fake_compact_tensor
+from cutlass.cutlass_dsl import T, dsl_user_op
+
+WARP_SIZE = 32
+VEC_BF16 = 8
+QUAD_BF16 = 4
+PACKED_BYTES = 16
+NEGATIVE_ZERO_BF16_BITS = 0x8000
+NEGATIVE_ZERO_BF16_PAIR = 0x80008000
+NEGATIVE_ZERO_F32_BITS = 0x80000000
+
+# Lamport mailboxes rotate through three generations; the stage-state words
+# hold the next generation to publish into and the one being consumed.
+LAMPORT_GENERATIONS = 3
+NEXT_STAGE = 0
+ACTIVE_STAGE = 1
+
+
+def _asm(result, operands, text, constraints, *, side_effects, loc, ip):
+    return llvm.inline_asm(
+        result,
+        operands,
+        text,
+        constraints,
+        has_side_effects=side_effects,
+        is_align_stack=False,
+        asm_dialect=llvm.AsmDialect.AD_ATT,
+        loc=loc,
+        ip=ip,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Global memory
+# ---------------------------------------------------------------------------
+
+
+@dsl_user_op
+def load_global_u32x4(
+    pointer: cute.Pointer,
+    *,
+    volatile: cutlass.Constexpr[bool] = False,
+    loc=None,
+    ip=None,
+):
+    """Load one 128-bit fragment as four u32 registers.
+
+    The volatile form is the Lamport polling load. Marking the asm
+    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"
+    loaded = _asm(
+        llvm.StructType.get_literal([T.i32()] * 4),
+        [address.ir_value(loc=loc, ip=ip)],
+        f"{opcode} {{$0, $1, $2, $3}}, [$4];",
+        "=r,=r,=r,=r,l",
+        side_effects=volatile,
+        loc=loc,
+        ip=ip,
+    )
+    packed = vector.from_elements(
+        ir.VectorType.get([4], T.i32(), loc=loc),
+        [
+            llvm.extractvalue(T.i32(), loaded, [index], loc=loc, ip=ip)
+            for index in range(4)
+        ],
+        loc=loc,
+        ip=ip,
+    )
+    return cute.TensorSSA(packed, 4, Uint32)
+
+
+@dsl_user_op
+def load_global_u32x2(pointer: cute.Pointer, *, loc=None, ip=None):
+    address = pointer.toint(loc=loc, ip=ip)
+    loaded = _asm(
+        llvm.StructType.get_literal([T.i32()] * 2),
+        [address.ir_value(loc=loc, ip=ip)],
+        "ld.global.v2.u32 {$0, $1}, [$2];",
+        "=r,=r,l",
+        side_effects=False,
+        loc=loc,
+        ip=ip,
+    )
+    packed = vector.from_elements(
+        ir.VectorType.get([2], T.i32(), loc=loc),
+        [
+            llvm.extractvalue(T.i32(), loaded, [index], loc=loc, ip=ip)
+            for index in range(2)
+        ],
+        loc=loc,
+        ip=ip,
+    )
+    return cute.TensorSSA(packed, 2, Uint32)
+
+
+@dsl_user_op
+def load_global_bf16_as_f32(address: Int64, *, loc=None, ip=None) -> Float32:
+    return Float32(
+        _asm(
+            T.f32(),
+            [address.ir_value(loc=loc, ip=ip)],
+            "{\n\t.reg .b16 bits;\n\tld.global.b16 bits, [$1];\n\t"
+            "cvt.f32.bf16 $0, bits;\n\t}",
+            "=f,l",
+            side_effects=False,
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
+@dsl_user_op
+def load_volatile_u32(pointer: cute.Pointer, *, loc=None, ip=None) -> Uint32:
+    address = pointer.toint(loc=loc, ip=ip)
+    return Uint32(
+        _asm(
+            T.i32(),
+            [address.ir_value(loc=loc, ip=ip)],
+            "ld.volatile.global.u32 $0, [$1];",
+            "=r,l",
+            side_effects=True,
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
+@dsl_user_op
+def store_global_u32x4(
+    address: Int64,
+    packed,
+    *,
+    volatile: cutlass.Constexpr[bool] = False,
+    loc=None,
+    ip=None,
+) -> None:
+    """Store four packed words to an ordinary or NVLS multicast global VA."""
+    words = [packed[index].ir_value(loc=loc, ip=ip) for index in range(4)]
+    opcode = "st.volatile.global.v4.u32" if volatile else "st.global.v4.u32"
+    _asm(
+        None,
+        [address.ir_value(loc=loc, ip=ip), *words],
+        f"{opcode} [$0], {{$1, $2, $3, $4}};",
+        "l,r,r,r,r",
+        side_effects=True,
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def store_global_u32x2(address: Int64, packed, *, loc=None, ip=None) -> None:
+    words = [packed[index].ir_value(loc=loc, ip=ip) for index in range(2)]
+    _asm(
+        None,
+        [address.ir_value(loc=loc, ip=ip), *words],
+        "st.global.v2.u32 [$0], {$1, $2};",
+        "l,r,r",
+        side_effects=True,
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def store_global_u32(address: Int64, word: Uint32, *, loc=None, ip=None) -> None:
+    _asm(
+        None,
+        [address.ir_value(loc=loc, ip=ip), word.ir_value(loc=loc, ip=ip)],
+        "st.global.u32 [$0], $1;",
+        "l,r",
+        side_effects=True,
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def red_async_release_gpu_add_u32(
+    pointer: cute.Pointer, value: Uint32, *, loc=None, ip=None
+) -> None:
+    """The SM100 arrival primitive FlashInfer's LamportFlags use."""
+    address = pointer.toint(loc=loc, ip=ip)
+    _asm(
+        None,
+        [address.ir_value(loc=loc, ip=ip), value.ir_value(loc=loc, ip=ip)],
+        "red.async.release.global.gpu.add.u32 [$0], $1;",
+        "l,r",
+        side_effects=True,
+        loc=loc,
+        ip=ip,
+    )
+
+
+# ---------------------------------------------------------------------------
+# NVLS multicast stores
+# ---------------------------------------------------------------------------
+
+
+@dsl_user_op
+def stmc_bf16x2(address: Int64, packed: Uint32, *, loc=None, ip=None) -> None:
+    _asm(
+        None,
+        [address.ir_value(loc=loc, ip=ip), packed.ir_value(loc=loc, ip=ip)],
+        "multimem.st.relaxed.sys.global.bf16x2 [$0], $1;",
+        "l,r",
+        side_effects=True,
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def stmc_bf16x4(address: Int64, packed, *, loc=None, ip=None) -> None:
+    words = [packed[index].ir_value(loc=loc, ip=ip) for index in range(2)]
+    _asm(
+        None,
+        [address.ir_value(loc=loc, ip=ip), *words],
+        "multimem.st.relaxed.sys.global.v2.bf16x2 [$0], {$1, $2};",
+        "l,r,r",
+        side_effects=True,
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def stmc_bf16x8(address: Int64, packed, *, loc=None, ip=None) -> None:
+    """Publish eight BF16 values (four packed words) through NVLS multicast."""
+    words = [packed[index].ir_value(loc=loc, ip=ip) for index in range(4)]
+    _asm(
+        None,
+        [address.ir_value(loc=loc, ip=ip), *words],
+        "multimem.st.relaxed.sys.global.v4.bf16x2 [$0], {$1, $2, $3, $4};",
+        "l,r,r,r,r",
+        side_effects=True,
+        loc=loc,
+        ip=ip,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Lamport sentinels
+#
+# A BF16 mailbox marks an empty slot with BF16 negative zero in every lane; an
+# FP32 mailbox marks it with FP32 negative zero in every word. Senders replace
+# real negative zeros with positive zeros so they never read as empty.
+# ---------------------------------------------------------------------------
+
+
+@dsl_user_op
+def store_lamport_sentinel_u32x4(address: Int64, *, loc=None, ip=None) -> None:
+    """Reset one BF16 Lamport fragment to BF16 negative zeros."""
+    sentinel = Uint32(NEGATIVE_ZERO_BF16_PAIR).ir_value(loc=loc, ip=ip)
+    _asm(
+        None,
+        [address.ir_value(loc=loc, ip=ip), sentinel, sentinel, sentinel, sentinel],
+        "st.global.v4.u32 [$0], {$1, $2, $3, $4};",
+        "l,r,r,r,r",
+        side_effects=True,
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def store_lamport_sentinel_128(pointer: cute.Pointer, *, loc=None, ip=None) -> None:
+    """Reset one FP32 Lamport fragment to four FP32 negative-zero bit patterns."""
+    address = pointer.toint(loc=loc, ip=ip)
+    sentinel = Uint32(NEGATIVE_ZERO_F32_BITS).ir_value(loc=loc, ip=ip)
+    _asm(
+        None,
+        [address.ir_value(loc=loc, ip=ip), sentinel, sentinel, sentinel, sentinel],
+        "st.global.v4.u32 [$0], {$1, $2, $3, $4};",
+        "l,r,r,r,r",
+        side_effects=True,
+        loc=loc,
+        ip=ip,
+    )
+
+
+@cute.jit
+def sanitize_negative_zero_u32(word: Uint32) -> Uint32:
+    low = Uint16(word & Uint32(0xFFFF))
+    high = Uint16(word >> Uint32(16))
+    if low == Uint16(NEGATIVE_ZERO_BF16_BITS):
+        word = word & Uint32(0xFFFF0000)
+    if high == Uint16(NEGATIVE_ZERO_BF16_BITS):
+        word = word & Uint32(0x0000FFFF)
+    return word
+
+
+@cute.jit
+def sanitize_negative_zero_u32x2(packed):
+    sanitized = cute.make_rmem_tensor(cute.make_layout((2,)), Uint32)
+    for index in cutlass.range_constexpr(2):
+        sanitized[index] = sanitize_negative_zero_u32(packed[index])
+    return sanitized.load()
+
+
+@cute.jit
+def sanitize_negative_zero_u32x4(packed):
+    """Turn real BF16 -0 into +0 so it cannot equal the empty sentinel."""
+    sanitized = cute.make_rmem_tensor(cute.make_layout((4,)), Uint32)
+    for index in cutlass.range_constexpr(4):
+        sanitized[index] = sanitize_negative_zero_u32(packed[index])
+    return sanitized.load()
+
+
+@cute.jit
+def fragment_has_negative_zero(packed):
+    """Whether a BF16 fragment still holds the BF16 sentinel in any lane."""
+    dirty = False
+    for index in cutlass.range_constexpr(4):
+        word = packed[index]
+        dirty = (
+            dirty
+            | (Uint16(word & Uint32(0xFFFF)) == Uint16(NEGATIVE_ZERO_BF16_BITS))
+            | (Uint16(word >> Uint32(16)) == Uint16(NEGATIVE_ZERO_BF16_BITS))
+        )
+    return dirty
+
+
+@cute.jit
+def fragment_is_dirty(packed):
+    """Whether an FP32 fragment still holds the FP32 sentinel in any word."""
+    dirty = packed[0] == Uint32(NEGATIVE_ZERO_F32_BITS)
+    for index in cutlass.range_constexpr(1, 4):
+        dirty = dirty | (packed[index] == Uint32(NEGATIVE_ZERO_F32_BITS))
+    return dirty
+
+
+# ---------------------------------------------------------------------------
+# Shared and distributed shared memory
+# ---------------------------------------------------------------------------
+
+
+@dsl_user_op
+def map_shared_to_peer(
+    smem_pointer: cute.Pointer, peer_rank: Int32, *, loc=None, ip=None
+) -> Int32:
+    """Map a local shared-memory slot to the same slot in a peer CTA."""
+    address = smem_pointer.toint(loc=loc, ip=ip).ir_value(loc=loc, ip=ip)
+    return Int32(
+        _asm(
+            T.i32(),
+            [address, peer_rank.ir_value(loc=loc, ip=ip)],
+            "mapa.shared::cluster.u32 $0, $1, $2;",
+            "=r,r,r",
+            side_effects=False,
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
+@dsl_user_op
+def store_shared_cluster_f32(
+    remote_address: Int32, value: Float32, *, loc=None, ip=None
+) -> None:
+    _asm(
+        None,
+        [remote_address.ir_value(loc=loc, ip=ip), value.ir_value(loc=loc, ip=ip)],
+        "st.shared::cluster.f32 [$0], $1;",
+        "r,f",
+        side_effects=True,
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def load_shared_f32x2(pointer: cute.Pointer, *, loc=None, ip=None):
+    """Load two aligned FP32 partials from local shared memory."""
+    address = Int32(pointer.toint(loc=loc, ip=ip))
+    loaded = _asm(
+        llvm.StructType.get_literal([T.f32()] * 2),
+        [address.ir_value(loc=loc, ip=ip)],
+        "ld.shared.v2.f32 {$0, $1}, [$2];",
+        "=f,=f,r",
+        side_effects=True,
+        loc=loc,
+        ip=ip,
+    )
+    return tuple(
+        Float32(llvm.extractvalue(T.f32(), loaded, [index], loc=loc, ip=ip))
+        for index in range(2)
+    )
+
+
+@dsl_user_op
+def load_shared_f32x4(pointer: cute.Pointer, *, loc=None, ip=None):
+    """Load four aligned FP32 partials from local shared memory."""
+    address = Int32(pointer.toint(loc=loc, ip=ip))
+    loaded = _asm(
+        llvm.StructType.get_literal([T.f32()] * 4),
+        [address.ir_value(loc=loc, ip=ip)],
+        "ld.shared.v4.f32 {$0, $1, $2, $3}, [$4];",
+        "=f,=f,=f,=f,r",
+        side_effects=True,
+        loc=loc,
+        ip=ip,
+    )
+    return tuple(
+        Float32(llvm.extractvalue(T.f32(), loaded, [index], loc=loc, ip=ip))
+        for index in range(4)
+    )
+
+
+# ---------------------------------------------------------------------------
+# BF16 packing and arithmetic
+# ---------------------------------------------------------------------------
+
+
+@dsl_user_op
+def packed_u32x4_to_bf16x8(packed, *, loc=None, ip=None):
+    values = llvm.bitcast(
+        ir.VectorType.get([VEC_BF16], BFloat16.mlir_type, loc=loc),
+        packed.ir_value(loc=loc, ip=ip),
+        loc=loc,
+        ip=ip,
+    )
+    return cute.TensorSSA(values, VEC_BF16, BFloat16)
+
+
+@dsl_user_op
+def packed_u32x2_to_bf16x4(packed, *, loc=None, ip=None):
+    values = llvm.bitcast(
+        ir.VectorType.get([QUAD_BF16], BFloat16.mlir_type, loc=loc),
+        packed.ir_value(loc=loc, ip=ip),
+        loc=loc,
+        ip=ip,
+    )
+    return cute.TensorSSA(values, QUAD_BF16, BFloat16)
+
+
+@dsl_user_op
+def bf16x8_to_packed_u32x4(values, *, loc=None, ip=None):
+    packed = llvm.bitcast(
+        ir.VectorType.get([4], T.i32(), loc=loc),
+        values.ir_value(loc=loc, ip=ip),
+        loc=loc,
+        ip=ip,
+    )
+    return cute.TensorSSA(packed, 4, Uint32)
+
+
+@dsl_user_op
+def bf16x4_to_packed_u32x2(values, *, loc=None, ip=None):
+    packed = llvm.bitcast(
+        ir.VectorType.get([2], T.i32(), loc=loc),
+        values.ir_value(loc=loc, ip=ip),
+        loc=loc,
+        ip=ip,
+    )
+    return cute.TensorSSA(packed, 2, Uint32)
+
+
+@dsl_user_op
+def bf16x2_to_u32(values, *, loc=None, ip=None) -> Uint32:
+    return Uint32(
+        llvm.bitcast(T.i32(), values.ir_value(loc=loc, ip=ip), loc=loc, ip=ip)
+    )
+
+
+@dsl_user_op
+def f32_to_bf16_bits(value: Float32, *, loc=None, ip=None) -> Uint32:
+    return Uint32(
+        _asm(
+            T.i32(),
+            [value.ir_value(loc=loc, ip=ip)],
+            "{\n\t.reg .b16 bits;\n\tcvt.rn.bf16.f32 bits, $1;\n\t"
+            "cvt.u32.u16 $0, bits;\n\t}",
+            "=r,f",
+            side_effects=False,
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
+@dsl_user_op
+def fma_f32_bf16(
+    a: BFloat16, b: BFloat16, acc: Float32, *, loc=None, ip=None
+) -> Float32:
+    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)
+    return Float32(
+        _asm(
+            T.f32(),
+            [a_bits, b_bits, acc.ir_value(loc=loc, ip=ip)],
+            "fma.rn.f32.bf16 $0, $1, $2, $3;",
+            "=f,h,h,f",
+            side_effects=False,
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
+# ---------------------------------------------------------------------------
+# Warp and block reductions
+# ---------------------------------------------------------------------------
+
+
+@dsl_user_op
+def shuffle_sync_idx_u32(
+    value: Uint32, source_lane: Int32, *, loc=None, ip=None
+) -> Uint32:
+    return Uint32(
+        _asm(
+            T.i32(),
+            [value.ir_value(loc=loc, ip=ip), source_lane.ir_value(loc=loc, ip=ip)],
+            "shfl.sync.idx.b32 $0, $1, $2, 0x1f, 0xffffffff;",
+            "=r,r,r",
+            # Preserve full-warp execution across later divergent consumers.
+            side_effects=True,
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
+@cute.jit
+def group_leader_block_sum(
+    value: Float32,
+    warp_sums: cute.Tensor,
+    warps: cutlass.Constexpr[int],
+    leader_stride: cutlass.Constexpr[int],
+) -> Float32:
+    """Block sum of the values held by every ``leader_stride``-th lane."""
+    lane = cute.arch.lane_idx()
+    warp = cute.arch.warp_idx()
+    for offset in cutlass.range_constexpr(1, WARP_SIZE):
+        if cutlass.const_expr(offset >= leader_stride and (offset & (offset - 1)) == 0):
+            value = value + cute.arch.shuffle_sync_bfly(
+                value,
+                offset=offset,
+                mask=-1,
+                mask_and_clamp=31,
+            )
+    if lane == 0:
+        cute.arch.store((warp_sums + warp).llvm_ptr, value)
+    cute.arch.barrier()
+
+    result = Float32(0.0)
+    if warp == 0:
+        if lane < Int32(warps):
+            result = cute.arch.load(
+                (warp_sums + lane).llvm_ptr,
+                Float32,
+            )
+        result = cute.arch.warp_reduction_sum(result)
+        if lane == 0:
+            cute.arch.store(warp_sums.llvm_ptr, result)
+    cute.arch.barrier()
+    return cute.arch.load(warp_sums.llvm_ptr, Float32)
+
+
+@cute.jit
+def warp_sum_specialized(
+    value: Float32,
+    warp_idx: Int32,
+    lane: Int32,
+    warps: cutlass.Constexpr[int],
+    last_warp_lanes: cutlass.Constexpr[int],
+    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.
+            if cutlass.const_expr(offset in (16, 8, 4, 2, 1)):
+                other = cute.arch.shuffle_sync_bfly(
+                    value,
+                    offset=offset,
+                    mask=last_warp_mask,
+                    mask_and_clamp=31,
+                )
+                if (lane ^ Int32(offset)) < Int32(last_warp_lanes):
+                    value = value + other
+    else:
+        for offset in cutlass.range_constexpr(16, 0, -1):
+            if cutlass.const_expr(offset in (16, 8, 4, 2, 1)):
+                value = value + cute.arch.shuffle_sync_bfly(
+                    value,
+                    offset=offset,
+                    mask=-1,
+                    mask_and_clamp=31,
+                )
+    return value
+
+
+# ---------------------------------------------------------------------------
+# Launch helpers
+# ---------------------------------------------------------------------------
+
+
+class CUDAGraphCompatibleWrapper:
+    """DLPack view that does not synchronize with the producer stream."""
+
+    __slots__ = ("tensor",)
+
+    def __init__(self, tensor: torch.Tensor) -> None:
+        self.tensor = tensor
+
+    def __dlpack__(self, stream=None):
+        # stream=-1 skips producer sync; CuTe launches on the current stream,
+        # which is the capturing one under CUDA graphs.
+        return self.tensor.__dlpack__(stream=-1)
+
+    def __dlpack_device__(self):
+        return self.tensor.__dlpack_device__()
+
+
+def to_cute(tensor: torch.Tensor, assumed_align: int = 16) -> cute.Tensor:
+    return from_dlpack(
+        CUDAGraphCompatibleWrapper(tensor.detach()), assumed_align=assumed_align
+    )
+
+
+def to_cute_dynamic(
+    tensor: torch.Tensor, alignment: int, *, divisibility: int
+) -> cute.Tensor:
+    """A 1-D view whose length is a runtime shape divisible by ``divisibility``."""
+    return to_cute(tensor, alignment).mark_compact_shape_dynamic(
+        mode=0, divisibility=divisibility
+    )
+
+
+def to_cute_dynamic_m(
+    tensor: torch.Tensor, *, mode: int, assumed_align: int = 16
+) -> cute.Tensor:
+    """Expose exactly one compact tensor mode as a runtime shape.
+
+    Model dimensions remain part of the compiled tensor type. Only the token
+    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()
+    )
+
+
+def make_fake_dynamic_compact_tensor(dtype, *, alignment: int, divisibility: int):
+    return make_fake_compact_tensor(
+        dtype, (cute.sym_int32(divisibility=divisibility),), assumed_align=alignment
+    )
+
+
+def current_cu_stream() -> cuda.CUstream:
+    return cuda.CUstream(torch.cuda.current_stream().cuda_stream)
diff --git a/vllm/models/common/ops/sequence_parallel.py b/vllm/models/common/ops/sequence_parallel.py
index 0f712c617e5a..b224fcd4637f 100644
--- a/vllm/models/common/ops/sequence_parallel.py
+++ b/vllm/models/common/ops/sequence_parallel.py
@@ -23,13 +23,17 @@ def _custom_collective(name: str, x: torch.Tensor) -> torch.Tensor | None:
 
 
 @lru_cache
-def is_sequence_parallel_active(is_enabled: bool, num_tokens: int) -> bool:
+def is_sequence_parallel_active(
+    is_enabled: bool, num_tokens: int, min_tokens: int = 0
+) -> bool:
     """Whether a batch of ``num_tokens`` shards the stream along the sequence.
 
     Below one token per rank the shards are mostly padding, so the batch
-    stays replicated.
+    stays replicated. A model may raise this lower bound with ``min_tokens``.
     """
-    return is_enabled and num_tokens >= get_tensor_model_parallel_world_size()
+    return is_enabled and num_tokens >= max(
+        get_tensor_model_parallel_world_size(), min_tokens
+    )
 
 
 def sp_all_gather(x: torch.Tensor) -> torch.Tensor:
diff --git a/vllm/models/deepseek_v32/amd/model.py b/vllm/models/deepseek_v32/amd/model.py
index 225e9b0ffe10..ce7713535494 100644
--- a/vllm/models/deepseek_v32/amd/model.py
+++ b/vllm/models/deepseek_v32/amd/model.py
@@ -7,14 +7,16 @@
 
 import torch
 
+from vllm._aiter_ops import rocm_aiter_ops
 from vllm.config import VllmConfig
-from vllm.distributed import get_pp_group, get_tensor_model_parallel_rank
+from vllm.distributed import (
+    get_pp_group,
+    get_tensor_model_parallel_rank,
+    tensor_model_parallel_all_reduce,
+)
 from vllm.model_executor.layers.fused_moe import (
     fused_moe_make_expert_params_mapping,
 )
-from vllm.model_executor.layers.fused_moe.utils import (
-    is_model_fused_shared_expert_compatible,
-)
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.vocab_parallel_embedding import (
     VocabParallelEmbedding,
@@ -32,10 +34,14 @@
 )
 from vllm.model_executor.models.utils import (
     PPMissingLayer,
+    copy_pp_buffers,
+    extract_layer_index,
+    get_fused_shared_expert_layers,
     get_pp_missing_layer_names,
     is_pp_missing_parameter,
     make_empty_intermediate_tensors_factory,
     make_layers,
+    pp_topk_buffers,
 )
 from vllm.models.common.ops.fused_allreduce_rms_norm import fused_allreduce_rms_norm
 from vllm.sequence import IntermediateTensors
@@ -87,7 +93,7 @@ def __init__(
             pc = moe_cfg.moe_parallel_config
             # Validated on MoRI only. DeepEP, pplx and the other combines
             # keep the previous always-skip behavior.
-            if pc.use_mori_kernels:
+            if pc.use_mori_epv1_kernels or pc.use_mori_epv2_kernels:
                 # Combine is already all-reduced; the next fused norm sums TP.
                 self.zero_mlp_output = get_tensor_model_parallel_rank() > 0
             else:
@@ -145,7 +151,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
 
         self.vocab_size = config.vocab_size
         self.is_v32 = True
-        topk_indices_buffer = torch.empty(
+        self.topk_indices_buffer = torch.empty(
             vllm_config.scheduler_config.max_num_batched_tokens,
             config.index_topk,
             dtype=torch.int32,
@@ -167,23 +173,26 @@ 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,
             ),
             prefix=f"{prefix}.layers",
         )
 
-        self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible(
-            self.layers,
+        self.fused_shared_expert_layers = get_fused_shared_expert_layers(
+            self,
             DeepseekV2MoE,
-            "mlp",
+            enable_eplb=vllm_config.parallel_config.enable_eplb,
         )
 
         if get_pp_group().is_last_rank:
             self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
         else:
             self.norm = PPMissingLayer()
+        self.pp_buffers = pp_topk_buffers(
+            config, self.topk_indices_buffer, range(config.num_hidden_layers)
+        )
         self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
-            ["hidden_states", "residual"], config.hidden_size
+            ["hidden_states"], config.hidden_size, self.pp_buffers
         )
 
         self.aux_hidden_state_layers = tuple[int, ...]()
@@ -207,11 +216,12 @@ def forward(
             else:
                 assert input_ids is not None
                 hidden_states = self.embed_input_ids(input_ids)
-            residual = None
         else:
             assert intermediate_tensors is not None
             hidden_states = intermediate_tensors["hidden_states"]
-            residual = intermediate_tensors["residual"]
+            copy_pp_buffers(intermediate_tensors, self.pp_buffers)
+        # Later PP stages receive the TP-reduced residual stream, like an embedding.
+        residual = None
 
         aux_hidden_states = []
         for idx, layer in enumerate(
@@ -219,12 +229,19 @@ def forward(
             start=self.start_layer,
         ):
             if idx in self.aux_hidden_state_layers:
-                aux_hidden_states.append(hidden_states + residual)
+                aux_hidden_states.append(
+                    hidden_states if residual is None else hidden_states + residual
+                )
             hidden_states, residual = layer(positions, hidden_states, residual)
 
         if not get_pp_group().is_last_rank:
+            # PP send/recv needs TP-replicated tensors; reduce the partial sum.
+            hidden_states = residual + tensor_model_parallel_all_reduce(hidden_states)
             return IntermediateTensors(
-                {"hidden_states": hidden_states, "residual": residual}
+                {
+                    "hidden_states": hidden_states,
+                    **{k: v[: positions.shape[0]] for k, v in self.pp_buffers.items()},
+                }
             )
 
         hidden_states, _ = fused_allreduce_rms_norm(hidden_states, residual, self.norm)
@@ -246,12 +263,9 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
             ckpt_gate_proj_name="gate_proj",
             ckpt_down_proj_name="down_proj",
             ckpt_up_proj_name="up_proj",
+            # NOTE: global despite per-layer FSE; non-fused layers skip the extras.
             num_experts=self.config.n_routed_experts
-            + (
-                self.config.n_shared_experts
-                if self.is_fused_shared_expert_enabled
-                else 0
-            ),
+            + (self.config.n_shared_experts if self.fused_shared_expert_layers else 0),
             num_redundant_experts=self.num_redundant_experts,
         )
 
@@ -269,8 +283,11 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
                 continue
             if get_spec_layer_idx_from_weight_name(self.config, name) is not None:
                 continue
+            # Only parse the layer index when FSE is enabled.
             is_fusion_moe_shared_experts_layer = (
-                self.is_fused_shared_expert_enabled and ("mlp.shared_experts" in name)
+                "mlp.shared_experts" in name
+                and bool(rocm_aiter_ops.is_fusion_moe_shared_experts_enabled())
+                and extract_layer_index(name) in self.fused_shared_expert_layers
             )
             if _try_load_fp8_indexer_wk(
                 name,
diff --git a/vllm/models/deepseek_v32/amd/mtp.py b/vllm/models/deepseek_v32/amd/mtp.py
index d5066b39de8d..2821bb7afa99 100644
--- a/vllm/models/deepseek_v32/amd/mtp.py
+++ b/vllm/models/deepseek_v32/amd/mtp.py
@@ -8,19 +8,20 @@
 import torch
 import torch.nn as nn
 
+from vllm._aiter_ops import rocm_aiter_ops
 from vllm.config import VllmConfig
 from vllm.distributed import tensor_model_parallel_all_reduce
 from vllm.model_executor.layers.fused_moe import (
     fused_moe_make_expert_params_mapping,
 )
-from vllm.model_executor.layers.fused_moe.utils import (
-    is_model_fused_shared_expert_compatible,
-)
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
 from vllm.model_executor.layers.vocab_parallel_embedding import (
     VocabParallelEmbedding,
 )
+from vllm.model_executor.model_loader.mtp_validation import (
+    is_mtp_completeness_check_enabled,
+)
 from vllm.model_executor.model_loader.weight_utils import (
     default_weight_loader,
     maybe_remap_kv_scale_name,
@@ -34,6 +35,8 @@
 )
 from vllm.model_executor.models.interfaces import SupportsPP
 from vllm.model_executor.models.utils import (
+    extract_layer_index,
+    get_fused_shared_expert_layers,
     get_pp_missing_layer_names,
     make_empty_intermediate_tensors_factory,
     maybe_prefix,
@@ -183,10 +186,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         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(),
+        self.fused_shared_expert_layers = get_fused_shared_expert_layers(
+            self,
             DeepseekV2MoE,
-            "mtp_block.mlp",
+            enable_eplb=vllm_config.parallel_config.enable_eplb,
         )
 
     def set_moe_parameters(self):
@@ -266,17 +269,15 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
             ("wk_weights_proj", "wk", 0),
             ("wk_weights_proj", "weights_proj", 1),
         ]
+        fused_shared_expert_layers = self.fused_shared_expert_layers
         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",
+            # NOTE: global despite per-layer FSE; non-fused layers skip the extras.
             num_experts=self.config.n_routed_experts
-            + (
-                self.config.n_shared_experts
-                if self.is_fused_shared_expert_enabled
-                else 0
-            ),
+            + (self.config.n_shared_experts if fused_shared_expert_layers else 0),
         )
 
         pp_missing_layer_names = get_pp_missing_layer_names(self)
@@ -303,8 +304,11 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
                     weight_loader(param, loaded_weight)
                     loaded_params.add(name)
                 continue
+            # Only parse the layer index when FSE is enabled.
             is_fusion_moe_shared_experts_layer = (
-                self.is_fused_shared_expert_enabled and ("mlp.shared_experts" in name)
+                "mlp.shared_experts" in name
+                and bool(rocm_aiter_ops.is_fusion_moe_shared_experts_enabled())
+                and extract_layer_index(name) in fused_shared_expert_layers
             )
             name = self._rewrite_spec_layer_name(spec_layer, name)
 
@@ -422,7 +426,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
             self.model.mtp_start_layer_idx,
             self.model.mtp_start_layer_idx + self.model.num_mtp_layers,
         ):
-            if layer_idx not in loaded_layers:
+            if layer_idx not in loaded_layers and is_mtp_completeness_check_enabled():
                 raise ValueError(
                     f"MTP speculative decoding layer {layer_idx} weights "
                     f"missing from checkpoint."
diff --git a/vllm/models/deepseek_v32/attention.py b/vllm/models/deepseek_v32/attention.py
index ed079204d283..de4121b63ca3 100644
--- a/vllm/models/deepseek_v32/attention.py
+++ b/vllm/models/deepseek_v32/attention.py
@@ -31,7 +31,10 @@
     DeepseekV32IndexerCache,
     yarn_get_mscale,
 )
-from vllm.model_executor.models.utils import extract_layer_index
+from vllm.model_executor.models.utils import (
+    compute_skip_topk_layers,
+    extract_layer_index,
+)
 from vllm.models.deepseek_v32.common.kernels import fused_norm_rope, fused_q
 from vllm.platforms import current_platform
 from vllm.v1.attention.ops.pcp import (
@@ -164,17 +167,7 @@ def __init__(
             scaling = scaling * mscale * mscale
 
         layer_id = extract_layer_index(prefix)
-        index_topk_freq = getattr(config, "index_topk_freq", 1)
-        index_topk_pattern = getattr(config, "index_topk_pattern", None)
-        index_skip_topk_offset = getattr(config, "index_skip_topk_offset", 2)
-        if index_topk_pattern is None:
-            skip_topk = (
-                max(layer_id - index_skip_topk_offset + 1, 0) % index_topk_freq != 0
-            )
-        elif 0 <= layer_id < len(index_topk_pattern):
-            skip_topk = index_topk_pattern[layer_id] == "S"
-        else:
-            skip_topk = False
+        skip_topk = layer_id in compute_skip_topk_layers(config)
 
         num_hidden_layers = getattr(config, "num_hidden_layers", None)
         is_mtp_layer = num_hidden_layers is not None and layer_id >= num_hidden_layers
diff --git a/vllm/models/deepseek_v32/nvidia/glm52_low_latency_gemm.py b/vllm/models/deepseek_v32/nvidia/glm52_low_latency_gemm.py
index dc247729f525..4e884885460d 100644
--- a/vllm/models/deepseek_v32/nvidia/glm52_low_latency_gemm.py
+++ b/vllm/models/deepseek_v32/nvidia/glm52_low_latency_gemm.py
@@ -1,10 +1,11 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""GLM-5.2 decode GEMM selection for unquantized BF16 on SM103/SM107.
+"""GLM-5.2 decode GEMM selection for unquantized BF16 on SM100/SM103/SM107.
 
 Dispatch is by local ``(N, K)`` weight shape and token count ``M``.
-:data:`GLM52_PROJECTIONS` was measured on B300 (SM103) and
-:data:`GLM52_PROJECTIONS_SM107` on VRNVL72 (SM107). The SM107 table also covers
+:data:`GLM52_PROJECTIONS` was measured on B300 (SM103),
+:data:`GLM52_PROJECTIONS_SM107` on VRNVL72 (SM107) and
+:data:`GLM52_PROJECTIONS_SM100` on B200 (SM100). The SM107 table also covers
 the TP4 projections (q_b, o_proj, shared expert) and the replicated indexer
 projections, which run as BF16 GEMMs at decode batch sizes.
 """
@@ -208,12 +209,134 @@ def build_plan(self) -> dict[int, ResolvedCall]:
 }
 
 
+# SM100 (B200) table: M=1-4, 6 and 8 were measured; each M is listed only where it
+# beat cuBLAS by >= 5 % (CUPTI, CUDA graph, cold and hot L2), others use cuBLAS.
+# Times are table backend / cuBLAS, cold L2. o_proj at TP4 and the shared expert
+# down projections stay on cuBLAS.
+# q_a + kv_a, replicated: 8.42 / 12.00 us at M=1.
+_SM100_QKV_A_PROJECTION = GLM52ProjectionSpec(
+    n=2624,
+    k=6144,
+    cute_configs=(
+        (1, SkinnyGemmConfig(1, 64, 4, k_unroll=4)),
+        (2, SkinnyGemmConfig(2, 64, 2, k_unroll=4)),
+        (3, SkinnyGemmConfig(3, 128, 2, k_unroll=2, vector_width=4)),
+        (4, SkinnyGemmConfig(4, 64, 2, k_unroll=2)),
+    ),
+    dsv3_tokens=frozenset({6, 8}),
+)
+
+# q_b_proj at TP8: 4.16 / 6.05 us at M=1.
+_SM100_Q_B_TP8_PROJECTION = GLM52ProjectionSpec(
+    n=2048,
+    k=2048,
+    cute_configs=(
+        (1, SkinnyGemmConfig(1, 128, 4, static_k=2048)),
+        (2, SkinnyGemmConfig(2, 128, 2, k_unroll=4, vector_width=4)),
+        (3, SkinnyGemmConfig(3, 128, 2, k_unroll=4, vector_width=4)),
+    ),
+    dsv3_tokens=frozenset({4, 6, 8}),
+)
+
+# q_b_proj at TP4 and the indexer wq_b (replicated): 5.76 / 6.88 us at M=1.
+_SM100_Q_B_TP4_PROJECTION = GLM52ProjectionSpec(
+    n=4096,
+    k=2048,
+    cute_configs=(
+        (1, SkinnyGemmConfig(1, 128, 2, k_unroll=4, vector_width=4)),
+        (2, SkinnyGemmConfig(2, 64, 2, k_unroll=4)),
+        (3, SkinnyGemmConfig(3, 64, 2, k_unroll=2)),
+        (4, SkinnyGemmConfig(4, 64, 2, k_unroll=4, vector_width=4)),
+    ),
+)
+
+# o_proj at TP8: 6.94 / 7.97 us at M=1.
+_SM100_O_PROJ_TP8_PROJECTION = GLM52ProjectionSpec(
+    n=6144,
+    k=2048,
+    cute_configs=(
+        (1, SkinnyGemmConfig(1, 64, 4, k_unroll=2)),
+        (2, SkinnyGemmConfig(2, 64, 4, k_unroll=2)),
+    ),
+)
+
+# indexer wk + weights_proj, replicated: 2.72 / 6.30 us at M=1.
+_SM100_WK_WEIGHTS_PROJECTION = GLM52ProjectionSpec(
+    n=160,
+    k=6144,
+    cute_configs=(
+        (1, SkinnyGemmConfig(1, 256, 1, k_unroll=2, static_k=6144)),
+        (2, SkinnyGemmConfig(2, 256, 1, k_unroll=2, static_k=6144)),
+        (3, SkinnyGemmConfig(3, 256, 1, k_unroll=2)),
+        (4, SkinnyGemmConfig(4, 256, 1, k_unroll=2)),
+        (6, SkinnyGemmConfig(6, 256, 1, k_unroll=2)),
+        (8, SkinnyGemmConfig(8, 256, 1, k_unroll=2)),
+    ),
+)
+
+# shared expert gate_up at TP8: 4.19 / 7.71 us at M=1.
+_SM100_SHARED_GATE_UP_TP8_PROJECTION = GLM52ProjectionSpec(
+    n=512,
+    k=6144,
+    cute_configs=(
+        (1, SkinnyGemmConfig(1, 256, 1, k_unroll=2)),
+        (2, SkinnyGemmConfig(2, 256, 1, k_unroll=2)),
+        (3, SkinnyGemmConfig(3, 256, 1, k_unroll=2)),
+        (4, SkinnyGemmConfig(4, 256, 1, k_unroll=2)),
+        (6, SkinnyGemmConfig(6, 128, 1, k_unroll=2)),
+        (8, SkinnyGemmConfig(8, 128, 1, k_unroll=4, vector_width=4)),
+    ),
+)
+
+# shared expert gate_up at TP4: 5.28 / 8.67 us at M=1.
+_SM100_SHARED_GATE_UP_TP4_PROJECTION = GLM52ProjectionSpec(
+    n=1024,
+    k=6144,
+    cute_configs=(
+        (1, SkinnyGemmConfig(1, 256, 2, k_unroll=2)),
+        (2, SkinnyGemmConfig(2, 256, 2, k_unroll=2)),
+        (3, SkinnyGemmConfig(3, 128, 1, k_unroll=2)),
+        (4, SkinnyGemmConfig(4, 128, 1, k_unroll=2)),
+        (6, SkinnyGemmConfig(6, 128, 2, k_unroll=4, vector_width=4)),
+        (8, SkinnyGemmConfig(8, 128, 2, k_unroll=2)),
+    ),
+)
+
+# MTP eh_proj (draft M = number of requests): 26.34 / 37.31 us at M=1.
+_SM100_EH_PROJECTION = GLM52ProjectionSpec(
+    n=6144,
+    k=12288,
+    cute_configs=(
+        (1, SkinnyGemmConfig(1, 256, 1, k_unroll=2)),
+        (2, SkinnyGemmConfig(2, 128, 1, k_unroll=4)),
+        (3, SkinnyGemmConfig(3, 128, 2, k_unroll=4, vector_width=4)),
+        (4, SkinnyGemmConfig(4, 128, 2, k_unroll=4, vector_width=4)),
+    ),
+)
+
+GLM52_PROJECTIONS_SM100 = {
+    (spec.n, spec.k): spec
+    for spec in (
+        _SM100_QKV_A_PROJECTION,
+        _SM100_Q_B_TP8_PROJECTION,
+        _SM100_Q_B_TP4_PROJECTION,
+        _SM100_O_PROJ_TP8_PROJECTION,
+        _SM100_WK_WEIGHTS_PROJECTION,
+        _SM100_SHARED_GATE_UP_TP8_PROJECTION,
+        _SM100_SHARED_GATE_UP_TP4_PROJECTION,
+        _SM100_EH_PROJECTION,
+    )
+}
+
+
 def _low_latency_table() -> dict[tuple[int, int], GLM52ProjectionSpec] | None:
     """Measured dispatch table for the current device, or None if unsupported."""
     if current_platform.is_device_capability((10, 7)):
         return GLM52_PROJECTIONS_SM107
     if current_platform.is_device_capability((10, 3)):
         return GLM52_PROJECTIONS
+    if current_platform.is_device_capability((10, 0)):
+        return GLM52_PROJECTIONS_SM100
     return None
 
 
diff --git a/vllm/models/deepseek_v32/nvidia/model.py b/vllm/models/deepseek_v32/nvidia/model.py
index 7720dfc5d182..ead6af480a9c 100644
--- a/vllm/models/deepseek_v32/nvidia/model.py
+++ b/vllm/models/deepseek_v32/nvidia/model.py
@@ -8,7 +8,7 @@
 
 import vllm.envs as envs
 from vllm.config import VllmConfig
-from vllm.distributed import get_pp_group
+from vllm.distributed import get_pp_group, tensor_model_parallel_all_reduce
 from vllm.forward_context import get_forward_context, is_forward_context_available
 from vllm.model_executor.layers.fused_embed_norm import (
     fused_embed_norm,
@@ -32,10 +32,12 @@
 )
 from vllm.model_executor.models.utils import (
     PPMissingLayer,
+    copy_pp_buffers,
     get_pp_missing_layer_names,
     is_pp_missing_parameter,
     make_empty_intermediate_tensors_factory,
     make_layers,
+    pp_topk_buffers,
 )
 from vllm.models.common.ops.fused_allreduce_rms_norm import fused_allreduce_rms_norm
 from vllm.models.common.ops.sequence_parallel import (
@@ -228,8 +230,11 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
             self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
         else:
             self.norm = PPMissingLayer()
+        self.pp_buffers = pp_topk_buffers(
+            config, self.topk_indices_buffer, range(config.num_hidden_layers)
+        )
         self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
-            ["hidden_states", "residual"], config.hidden_size
+            ["hidden_states"], config.hidden_size, self.pp_buffers
         )
 
         self.aux_hidden_state_layers = tuple[int, ...]()
@@ -263,11 +268,12 @@ def forward(
             else:
                 assert input_ids is not None
                 hidden_states = self.embed_input_ids(input_ids)
-            residual = None
         else:
             assert intermediate_tensors is not None
             hidden_states = intermediate_tensors["hidden_states"]
-            residual = intermediate_tensors["residual"]
+            copy_pp_buffers(intermediate_tensors, self.pp_buffers)
+        # Later PP stages receive the TP-reduced residual stream, like an embedding.
+        residual = None
 
         full_num_tokens = positions.shape[0]
         if self.use_sequence_parallel:
@@ -279,7 +285,6 @@ def forward(
             hidden_states = sp_shard(hidden_states)
             if attn_in is not None:
                 attn_in = sp_shard(attn_in)
-            assert residual is None, "Currently, SP is not supported with PP"
 
         aux_hidden_states = []
         for idx, layer in enumerate(
@@ -297,8 +302,13 @@ def forward(
             assert not self.use_sequence_parallel, (
                 "Currently, SP is not supported with PP"
             )
+            # PP send/recv needs TP-replicated tensors; reduce the partial sum.
+            hidden_states = residual + tensor_model_parallel_all_reduce(hidden_states)
             return IntermediateTensors(
-                {"hidden_states": hidden_states, "residual": residual}
+                {
+                    "hidden_states": hidden_states,
+                    **{k: v[:full_num_tokens] for k, v in self.pp_buffers.items()},
+                }
             )
 
         if self.use_sequence_parallel:
diff --git a/vllm/models/deepseek_v4/amd/dspark.py b/vllm/models/deepseek_v4/amd/dspark.py
index a316b6927e22..2f337c90b077 100644
--- a/vllm/models/deepseek_v4/amd/dspark.py
+++ b/vllm/models/deepseek_v4/amd/dspark.py
@@ -38,6 +38,9 @@
     ParallelLMHead,
     VocabParallelEmbedding,
 )
+from vllm.model_executor.model_loader.mtp_validation import (
+    is_mtp_completeness_check_enabled,
+)
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 from vllm.model_executor.models.qwen3_dspark import (
     DSparkConfidenceHead,
@@ -477,7 +480,11 @@ 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:
+        if (
+            self.model.confidence_head is not None
+            and not loaded_confidence_head
+            and is_mtp_completeness_check_enabled()
+        ):
             self.model.confidence_head = None
         logger.info_once("DSpark draft model loaded: %d params", len(loaded_params))
         return loaded_params
diff --git a/vllm/models/deepseek_v4/amd/model.py b/vllm/models/deepseek_v4/amd/model.py
index b0d798033f11..01a826f3b9e1 100644
--- a/vllm/models/deepseek_v4/amd/model.py
+++ b/vllm/models/deepseek_v4/amd/model.py
@@ -13,7 +13,7 @@
 
 import vllm.envs as envs
 from vllm._aiter_ops import rocm_aiter_ops
-from vllm.config import ParallelConfig, VllmConfig
+from vllm.config import VllmConfig
 from vllm.distributed import (
     get_pp_group,
     get_tensor_model_parallel_rank,
@@ -35,7 +35,6 @@
     fused_topk_bias,
 )
 from vllm.model_executor.layers.fused_moe.utils import (
-    is_model_fused_shared_expert_compatible,
     resolve_layer_fused_shared_expert,
 )
 from vllm.model_executor.layers.layernorm import RMSNorm
@@ -54,9 +53,6 @@
 )
 from vllm.model_executor.layers.quantization import QuantizationConfig
 from vllm.model_executor.layers.quantization.mxfp4 import Mxfp4MoEMethod
-from vllm.model_executor.layers.quantization.utils.config_utils import (
-    is_shared_expert_quant_fse_compatible,
-)
 from vllm.model_executor.layers.vocab_parallel_embedding import (
     ParallelLMHead,
     VocabParallelEmbedding,
@@ -72,6 +68,7 @@
     PPMissingLayer,
     WeightsMapper,
     extract_layer_index,
+    get_fused_shared_expert_layers,
     is_pp_missing_parameter,
     make_layers,
     maybe_prefix,
@@ -525,31 +522,6 @@ def forward_modular(
         return routed + shared_out
 
 
-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 parallel_config.enable_expert_parallel
-    )
-
-
 def _validate_aiter_mega_moe_config(vllm_config: VllmConfig) -> None:
     parallel_config = vllm_config.parallel_config
     if not rocm_aiter_ops.is_fused_moe_enabled():
@@ -572,10 +544,11 @@ def _validate_aiter_mega_moe_config(vllm_config: VllmConfig) -> None:
     if parallel_config.all2all_backend not in (
         "mori_high_throughput",
         "mori_low_latency",
-    ):
+        "mori",
+    ) or (parallel_config.all2all_backend == "mori" and parallel_config.mori_ep_v2):
         raise ValueError(
-            "AITER MegaMoE runs on MoRI shmem; add "
-            "--all2all-backend mori_high_throughput."
+            "AITER MegaMoE runs on MoRI EPv1 shmem; add --all2all-backend mori "
+            "(without VLLM_MORI_EP_V2=1)."
         )
     if parallel_config.enable_eplb:
         raise NotImplementedError("AITER MegaMoE does not support EPLB.")
@@ -658,28 +631,31 @@ def __init__(
 
         self.n_shared_experts = config.n_shared_experts
 
-        # TODO: Historically, only `VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=1`
-        # is checked to enable FSE for DeepSeek-v4, despite AITER not being used.
-        # 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, vllm_config.parallel_config)
+        self.is_fused_shared_expert_enabled = False
+        if (
+            getattr(config, "n_shared_experts", None)
             and not self.fuse_heterogeneous_shared_expert
-        )
-        fse_compatible = False
-        if fse_requested:
-            fse_compatible, fse_reason = is_shared_expert_quant_fse_compatible(
-                quant_config,
-                f"{prefix}.experts",
-                f"{prefix}.shared_experts",
-            )
-            if not fse_compatible:
-                logger.warning(
-                    "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS is enabled but "
-                    "cannot be enabled: %s.",
-                    fse_reason,
+        ):
+            # Fused shared experts are not supported under data parallelism:
+            # the fused path cannot load the FP8 shared expert into the MXFP4
+            # routed slot. Fail fast with an actionable message.
+            if (
+                not vllm_config.parallel_config.enable_expert_parallel
+                and vllm_config.parallel_config.data_parallel_size > 1
+                and rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
+            ):
+                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."
                 )
-        self.is_fused_shared_expert_enabled = fse_requested and fse_compatible
+            self.is_fused_shared_expert_enabled = resolve_layer_fused_shared_expert(
+                quant_config, prefix, disable_under_ep=True
+            )
 
         self.use_mega_moe = vllm_config.kernel_config.moe_backend == "aiter_mega_moe"
         fuse_mega_moe_shared_expert = (
@@ -1160,10 +1136,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
             ),
             prefix=f"{prefix}.layers",
         )
-        self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible(
-            self.layers,
+        self.fused_shared_expert_layers = get_fused_shared_expert_layers(
+            self,
             DeepseekV4MoE,
-            "ffn",
+            enable_eplb=vllm_config.parallel_config.enable_eplb,
         )
 
         if get_pp_group().is_last_rank:
@@ -1344,16 +1320,12 @@ def _resolve_param_name(name: str) -> str:
         # at init) as the single source of truth, so the redirect below cannot
         # diverge from how the module was built if per-layer quantization ever
         # mixes fused and non-fused layers.
-        fuse_by_layer = {
-            extract_layer_index(mod_name): mod.is_fused_shared_expert_enabled
-            for mod_name, mod in self.named_modules()
-            if isinstance(mod, DeepseekV4MoE)
-        }
+        fuse_by_layer = self.fused_shared_expert_layers
         n_routed = self.config.n_routed_experts
         # The redirect below maps the single shared-expert tensor group to one
         # appended slot; multiple shared experts would need per-expert slicing
         # (see deepseek_v2.py). DeepSeek-V4 has n_shared_experts == 1.
-        if any(fuse_by_layer.values()) and self.config.n_shared_experts != 1:
+        if fuse_by_layer and self.config.n_shared_experts != 1:
             raise NotImplementedError(
                 "deepseek-v4 fused shared-expert loading supports only "
                 f"n_shared_experts == 1, got {self.config.n_shared_experts}"
@@ -1364,13 +1336,7 @@ def _resolve_param_name(name: str) -> str:
             # into appended routed-expert slot ``.ffn.experts.{n_routed}``
             # so the MXFP4-quantized shared expert loads through the routed
             # expert loader (grouped GEMM). Single shared expert only.
-            if ".ffn.shared_experts.w" in name and fuse_by_layer.get(
-                extract_layer_index(name), False
-            ):
-                name = name.replace(
-                    ".ffn.shared_experts.w",
-                    f".ffn.experts.{n_routed}.w",
-                )
+            name = _redirect_shared_expert_name(name, fuse_by_layer, n_routed)
 
             for param_name, weight_name, shard_id in stacked_params_mapping:
                 # Skip non-stacked layers and experts (experts handled below).
@@ -1448,7 +1414,7 @@ def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
         # redirected shared-expert weights route through the expert loader.
         n_shared = getattr(self.config, "n_shared_experts", 0) or 0
         num_experts = self.config.n_routed_experts + (
-            n_shared if self.is_fused_shared_expert_enabled else 0
+            n_shared if self.fused_shared_expert_layers else 0
         )
         return fused_moe_make_expert_params_mapping(
             self,
@@ -1460,6 +1426,24 @@ def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
         )
 
 
+def _redirect_shared_expert_name(
+    name: str, fuse_by_layer: frozenset[int], n_routed_experts: int
+) -> str:
+    """Route a DSV4 shared-expert checkpoint name to its final parameter name.
+
+    Fused layers go to the appended routed-expert slot. When any layer is fused
+    the weight mapper leaves the checkpoint ``.w2`` names untouched, so
+    non-fused layers still need the native ``down_proj`` parameter name.
+    """
+    if ".ffn.shared_experts.w" not in name:
+        return name
+    if extract_layer_index(name) in fuse_by_layer:
+        return name.replace(
+            ".ffn.shared_experts.w", f".ffn.experts.{n_routed_experts}.w"
+        )
+    return name.replace(".ffn.shared_experts.w2", ".ffn.shared_experts.down_proj")
+
+
 def _make_deepseek_v4_weights_mapper(
     expert_dtype: str, fuse_shared_experts: bool = False
 ) -> WeightsMapper:
@@ -1514,7 +1498,6 @@ 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".
@@ -1533,10 +1516,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         self.model = self.model_cls(
             vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
         )
-        if expert_dtype != "fp4" or self.model.is_fused_shared_expert_enabled:
+        if expert_dtype != "fp4" or self.model.fused_shared_expert_layers:
             self.hf_to_vllm_mapper = _make_deepseek_v4_weights_mapper(
                 expert_dtype,
-                fuse_shared_experts=self.model.is_fused_shared_expert_enabled,
+                fuse_shared_experts=bool(self.model.fused_shared_expert_layers),
             )
         if get_pp_group().is_last_rank:
             self.lm_head = ParallelLMHead(
diff --git a/vllm/models/deepseek_v4/amd/mono/runner.py b/vllm/models/deepseek_v4/amd/mono/runner.py
index 47bea49be826..0468f0acb84c 100644
--- a/vllm/models/deepseek_v4/amd/mono/runner.py
+++ b/vllm/models/deepseek_v4/amd/mono/runner.py
@@ -21,6 +21,8 @@
 
 import torch
 
+from vllm.v1.attention.ops.flydsl_kernels import _run_compiled
+
 from .common.plan import BLOCKS
 from .dims import (
     EXPERTS,
@@ -359,7 +361,8 @@ def forward(
                 comp_block_table.stride(0),
                 topk_indices,
             )
-        k1(
+        _run_compiled(
+            k1,
             residual.data_ptr(),
             x.data_ptr(),
             post_mix.data_ptr(),
@@ -400,10 +403,11 @@ def forward(
             scratch.data_ptr(),
             self.epoch.data_ptr(),
             self.tl[0].data_ptr() if self.tl is not None else 0,
-            stream=st,
+            st,
         )
         select(self.normed[:M], self.qr[:M])
-        k2(
+        _run_compiled(
+            k2,
             self.q.data_ptr(),
             swa_cache.data_ptr(),
             swa_cache.stride(0),
@@ -458,6 +462,6 @@ def forward(
             self.rank,
             self.epoch.data_ptr(),
             self.tl[1].data_ptr() if self.tl is not None else 0,
-            stream=st,
+            st,
         )
         return outs
diff --git a/vllm/models/deepseek_v4/amd/mtp.py b/vllm/models/deepseek_v4/amd/mtp.py
index d49769c843d9..1144bbde9604 100644
--- a/vllm/models/deepseek_v4/amd/mtp.py
+++ b/vllm/models/deepseek_v4/amd/mtp.py
@@ -302,6 +302,18 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         self.model = DeepSeekV4MultiTokenPredictor(
             vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
         )
+        if any(
+            getattr(
+                layer.mtp_block.ffn,
+                "is_fused_shared_expert_enabled",
+                False,
+            )
+            for layer in self.model.layers.values()
+        ):
+            raise NotImplementedError(
+                "DeepSeek-V4 MTP fused shared experts are not supported yet; "
+                "please open an issue."
+            )
 
     def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
         return self.model.embed_input_ids(input_ids)
diff --git a/vllm/models/deepseek_v4/common/vl_model.py b/vllm/models/deepseek_v4/common/vl_model.py
index 7868c59f3393..eaa888a7754c 100644
--- a/vllm/models/deepseek_v4/common/vl_model.py
+++ b/vllm/models/deepseek_v4/common/vl_model.py
@@ -79,24 +79,6 @@ def _make_deepseek_v4_vl_weights_mapper(
     )
 
 
-def stream_language_model_first(
-    weights: Iterable[tuple[str, torch.Tensor]],
-) -> Iterable[tuple[str, torch.Tensor]]:
-    """Yield ``language_model.*`` weights as they arrive, then the rest.
-
-    Keeps the language model's weights in one contiguous group without
-    materializing the stream. Deferred weights are cloned so they do not keep
-    the loader's per-shard storage alive until the end of the load.
-    """
-    rest = []
-    for name, tensor in weights:
-        if name.startswith("language_model."):
-            yield name, tensor
-        else:
-            rest.append((name, tensor.clone()))
-    yield from rest
-
-
 @MULTIMODAL_REGISTRY.register_processor(
     DeepseekV4VLMultiModalProcessor,
     info=DeepseekV4VLProcessingInfo,
@@ -203,7 +185,6 @@ def __init__(self, *, vllm_config, prefix: str = "") -> None:
         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)
@@ -345,26 +326,8 @@ def get_mtp_target_hidden_states(self) -> torch.Tensor | None:
         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 = stream_language_model_first(mapped)
         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
+        return loader.load_weights(self.hf_to_vllm_mapper.apply(weights))
 
     def get_mm_mapping(self) -> MultiModelKeys:
         """Get the module prefixes in the multimodal model."""
diff --git a/vllm/models/deepseek_v4/cpu/cpu_mla.py b/vllm/models/deepseek_v4/cpu/cpu_mla.py
index 52d4f4c3b017..430886590d7c 100644
--- a/vllm/models/deepseek_v4/cpu/cpu_mla.py
+++ b/vllm/models/deepseek_v4/cpu/cpu_mla.py
@@ -87,9 +87,10 @@ def _split_indexer_prefill_chunks(
         workspace_size: int,
         max_logits_bytes: int,
         request_offset: int = 0,
+        fused_mqa_topk_eligible: bool = False,
     ) -> list[tuple[slice, slice]]:
-        # workspace_size/max_logits_bytes are unused: no flat-gather
-        # workspace or dense M*N buffer is allocated.
+        # workspace_size/max_logits_bytes/fused_mqa_topk_eligible are unused: no
+        # flat-gather workspace or dense M*N buffer is allocated.
         num_reqs = compressed_seq_lens_cpu.shape[0]
         total_query_len = int(prefill_query_lens_cpu.sum().item())
         return [
diff --git a/vllm/models/deepseek_v4/cpu/cpu_sparse.py b/vllm/models/deepseek_v4/cpu/cpu_sparse.py
index 021bef5b9c64..88210ec3e9a1 100644
--- a/vllm/models/deepseek_v4/cpu/cpu_sparse.py
+++ b/vllm/models/deepseek_v4/cpu/cpu_sparse.py
@@ -15,7 +15,7 @@
 
 from vllm import _custom_ops as ops
 from vllm.forward_context import get_forward_context
-from vllm.model_executor.utils import replace_parameter
+from vllm.model_executor.utils import replace_parameter, retire_parameter
 from vllm.models.deepseek_v4.attention import (
     DeepseekV4Attention,
     DeepseekV4Indexer,
@@ -283,7 +283,7 @@ def __init__(self, *args, **kwargs) -> None:
         self._wo_a_packed: torch.Tensor | None = None
         self._wrap_wo_a_process_weights_after_loading()
 
-    def process_weights_after_loading(self, act_dtype: torch.dtype) -> None:
+    def process_weights_after_loading(self) -> None:
         """Cache fp32-contiguous copies of tensors CPU kernels want that way
         but that arrive bf16 (rotary cos/sin table, compressor RMSNorm
         weights) -- cast once here instead of per forward call.
@@ -317,18 +317,17 @@ def _capture_and_pack(layer: torch.nn.Module) -> None:
             dequant = dequant.view(self.n_local_groups, self.o_lora_rank, k_dim)
             self._wo_a_packed = torch.ops._C.convert_weight_packed(dequant.contiguous())
             orig_pwal(layer)
-            # orig_pwal VNNI-repacks layer.weight/scale for ordinary-linear
-            # use that _o_proj never needs -- free them, same pattern as
-            # dispatch_cpu_unquantized_gemm's remove_weight=True.
-            replace_parameter(layer, "weight", layer.weight.new_empty(0))
+            # _o_proj reads only the packed copy: it becomes the loaded weight's
+            # storage, and the (small) scale keeps its data.
+            replace_parameter(layer, "weight", self._wo_a_packed)
+            self._wo_a_packed = layer.weight.data
             scale_attr = (
                 "weight_scale_inv"
                 if hasattr(layer, "weight_scale_inv")
                 else "weight_scale"
             )
             if hasattr(layer, scale_attr):
-                scale = getattr(layer, scale_attr)
-                replace_parameter(layer, scale_attr, scale.new_empty(0))
+                retire_parameter(layer, scale_attr)
 
         quant_method.process_weights_after_loading = (  # type: ignore[method-assign]
             _capture_and_pack
diff --git a/vllm/models/deepseek_v4/cpu/model.py b/vllm/models/deepseek_v4/cpu/model.py
index 63ec337fa48b..2e6f76c09e07 100644
--- a/vllm/models/deepseek_v4/cpu/model.py
+++ b/vllm/models/deepseek_v4/cpu/model.py
@@ -21,9 +21,6 @@
     GateLinear,
     fused_moe_make_expert_params_mapping,
 )
-from vllm.model_executor.layers.fused_moe.utils import (
-    is_model_fused_shared_expert_compatible,
-)
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import (
     MergedColumnParallelLinear,
@@ -493,11 +490,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
             ),
             prefix=f"{prefix}.layers",
         )
-        self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible(
-            self.layers,
-            DeepseekV4MoE,
-            "ffn",
-        )
+        self.is_fused_shared_expert_enabled = False
 
         if get_pp_group().is_last_rank:
             self.norm = RMSNorm(config.hidden_size, self.rms_norm_eps)
diff --git a/vllm/models/deepseek_v4/nvidia/dspark.py b/vllm/models/deepseek_v4/nvidia/dspark.py
index dda1026a53ef..252c639f8dde 100644
--- a/vllm/models/deepseek_v4/nvidia/dspark.py
+++ b/vllm/models/deepseek_v4/nvidia/dspark.py
@@ -37,6 +37,9 @@
     ParallelLMHead,
     VocabParallelEmbedding,
 )
+from vllm.model_executor.model_loader.mtp_validation import (
+    is_mtp_completeness_check_enabled,
+)
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 from vllm.model_executor.models.qwen3_dspark import (
     DSparkConfidenceHead,
@@ -423,6 +426,12 @@ def compute_draft_logits(self, hidden_states: torch.Tensor) -> torch.Tensor:
         # Full-vocab draft: base logits, no d2t scatter.
         return self.compute_logits(hidden_states)
 
+    def compute_draft_logits_local(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        """This TP rank's vocab shard of compute_draft_logits, not gathered."""
+        return self.logits_processor(
+            self.lm_head, self.model.norm(hidden_states), skip_gather=True
+        )
+
     def map_draft_to_target(self, draft_ids: torch.Tensor) -> torch.Tensor:
         return draft_ids  # full-vocab: draft ids are target ids
 
@@ -564,7 +573,11 @@ 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:
+        if (
+            self.model.confidence_head is not None
+            and not loaded_confidence_head
+            and is_mtp_completeness_check_enabled()
+        ):
             self.model.confidence_head = None
         logger.info_once("DSpark draft model loaded: %d params", len(loaded_params))
         return loaded_params
diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py
index 04f7ae21f95d..cc2d2d6c994d 100644
--- a/vllm/models/deepseek_v4/nvidia/model.py
+++ b/vllm/models/deepseek_v4/nvidia/model.py
@@ -1323,7 +1323,6 @@ class DeepseekV4ForCausalLM(
     SupportsLoRA,
     DeepseekV4MixtureOfExperts,
 ):
-    finalizes_weights_during_load = True
     model_cls = DeepseekV4Model
 
     # Default mapper assumes the original FP4-expert checkpoint layout.
@@ -1411,9 +1410,7 @@ def get_mtp_target_hidden_states(self) -> torch.Tensor | None:
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         loader = AutoWeightsLoader(self)
-        loaded_params = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
-        self.process_weights_after_loading()
-        return loaded_params
+        return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
 
     def process_weights_after_loading(self) -> None:
         self.model.finalize_mhc_broadcast_weights()
diff --git a/vllm/models/deepseek_v4/xpu/dspark.py b/vllm/models/deepseek_v4/xpu/dspark.py
index 41ec1b4a1674..fc95e5c5ee01 100644
--- a/vllm/models/deepseek_v4/xpu/dspark.py
+++ b/vllm/models/deepseek_v4/xpu/dspark.py
@@ -388,11 +388,10 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
                 weight_loader(param, loaded_weight)
                 loaded_params.add(name)
 
-        self._finalize_moe()
         logger.info_once("DSpark XPU draft model loaded: %d params", len(loaded_params))
         return loaded_params
 
-    def _finalize_moe(self) -> None:
+    def process_weights_after_loading(self) -> None:
         for layer in self.model.layers:
             layer.ffn.finalize_mega_moe_weights()
 
diff --git a/vllm/models/deepseek_v4/xpu/model.py b/vllm/models/deepseek_v4/xpu/model.py
index a2fb37a1c71b..5b819352ad16 100644
--- a/vllm/models/deepseek_v4/xpu/model.py
+++ b/vllm/models/deepseek_v4/xpu/model.py
@@ -541,10 +541,6 @@ def _run_mega_moe(
             symm_buffer.topk_weights[:num_tokens],
         )
 
-        # This method must have been already called during the weight loading phase.
-        # We call it again here to cover the dummy weight loading case.
-        self.finalize_weights()
-
         assert self._transformed_l1_weights is not None
         assert self._transformed_l2_weights is not None
         deep_gemm.fp8_fp4_mega_moe(
@@ -1367,9 +1363,10 @@ def get_mtp_target_hidden_states(self) -> torch.Tensor | None:
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         loader = AutoWeightsLoader(self)
-        loaded_params = 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.finalize_mega_moe_weights()
-        return loaded_params
 
     def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
         return self.model.get_expert_mapping()
diff --git a/vllm/models/deepseek_v4/xpu/mtp.py b/vllm/models/deepseek_v4/xpu/mtp.py
index dfbb79722ab0..b0fa52e53c19 100644
--- a/vllm/models/deepseek_v4/xpu/mtp.py
+++ b/vllm/models/deepseek_v4/xpu/mtp.py
@@ -466,11 +466,10 @@ def _find_mtp_layer_idx(name: str) -> int:
                     f"Use a checkpoint that includes MTP layer weights, "
                     f"or disable speculative decoding."
                 )
-        self.finalize_mega_moe_weights()
         logger.info_once("MTP draft model loaded: %d params", len(loaded_params))
         return loaded_params
 
-    def finalize_mega_moe_weights(self) -> None:
+    def process_weights_after_loading(self) -> None:
         for layer in self.model.layers.values():
             layer.mtp_block.ffn.finalize_mega_moe_weights()
 
diff --git a/vllm/models/deepseek_v41/amd/dspark.py b/vllm/models/deepseek_v41/amd/dspark.py
index 62b663ff86b3..9dc264565957 100644
--- a/vllm/models/deepseek_v41/amd/dspark.py
+++ b/vllm/models/deepseek_v41/amd/dspark.py
@@ -32,6 +32,9 @@
     ParallelLMHead,
     VocabParallelEmbedding,
 )
+from vllm.model_executor.model_loader.mtp_validation import (
+    is_mtp_completeness_check_enabled,
+)
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 from vllm.model_executor.models.qwen3_dspark import (
     DSparkConfidenceHead,
@@ -510,17 +513,15 @@ 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:
+        if (
+            self.model.confidence_head is not None
+            and not loaded_confidence_head
+            and is_mtp_completeness_check_enabled()
+        ):
             self.model.confidence_head = None
-        self.process_weights_after_loading()
         logger.info_once("DSpark draft model loaded: %d params", len(loaded_params))
         return loaded_params
 
-    def process_weights_after_loading(self) -> None:
-        # ROCm linears and fused-MoE runners finalize their own quantized
-        # parameters. NVIDIA-only MegaMoE/WO-A requantization is not needed.
-        return
-
     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/amd/model.py b/vllm/models/deepseek_v41/amd/model.py
index e7c513047140..9c7538007d6d 100644
--- a/vllm/models/deepseek_v41/amd/model.py
+++ b/vllm/models/deepseek_v41/amd/model.py
@@ -11,7 +11,6 @@
 import vllm.envs as envs
 from vllm.config import VllmConfig
 from vllm.distributed import (
-    get_engram_dp_size,
     get_pp_group,
     get_tensor_model_parallel_world_size,
     tensor_model_parallel_all_reduce,
@@ -576,7 +575,9 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
             prefix=f"{prefix}.layers",
         )
 
-        # Config flag is the switch: ROCm forces swa_cache_layer.bounded_replay off.
+        # Config flag is the switch. On ROCm, attention forces
+        # swa_cache_layer.bounded_replay off only when the CDNA version is
+        # not 4. CDNA4 keeps the window in bounded replay.
         self.decoder_replay_layers: DecoderReplayLayers | None = None
         self.decoder_replay_start = self.end_layer
         cut = max(config.kv_source_layer_ids)
@@ -604,6 +605,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
                 self.layers[cut].write_kv,
                 metadata_prefixes,
                 self.layers[cut].attn.swa_cache_layer.prefix,
+                lambda: any(i < cut for i in self.aux_hidden_state_layers),
             )
             logger.info_once(
                 "Decoder SWA bounded replay: in eager prefill steps, layer %d "
@@ -705,9 +707,8 @@ def forward(
 
         # Engram n-gram hashes for the whole (flattened) batch, computed once
         # on the full token stream — before any sequence-parallel sharding —
-        # and consumed by the engram layers (1 and 14) below. Skipped on
-        # profile runs (KV cache unbound -> no slot space to key the rolling
-        # compressed-id cache into).
+        # and consumed by the engram layers (1 and 14) below. Profile runs
+        # (KV cache unbound) look up dead rows instead.
         engram_hashes: torch.Tensor | None = None
         engram_mask: torch.Tensor | None = None
         if (
@@ -748,9 +749,10 @@ def forward(
                     swa_metadata.slot_mapping,
                     swa_metadata.block_table,
                 )
-            elif not self.engram_dp_shared_memory and get_engram_dp_size() > 1:
+            else:
                 # DP-sharded lookups are collective, so a replica skipping the
-                # hash still has to reach them.
+                # hash still has to reach them; the memory profile also has to
+                # see the Engram activations.
                 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.
@@ -1343,6 +1345,18 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         self.model = self.model_cls(
             vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
         )
+        if any(
+            getattr(
+                getattr(layer, "ffn", None),
+                "is_fused_shared_expert_enabled",
+                False,
+            )
+            for layer in self.model.layers
+        ):
+            raise NotImplementedError(
+                "DeepSeek-V4.1 AMD fused shared experts are not supported yet; "
+                "please open an issue."
+            )
         if get_pp_group().is_last_rank:
             self.lm_head = ParallelLMHead(
                 config.vocab_size,
@@ -1435,9 +1449,7 @@ def get_mtp_target_hidden_states(self) -> torch.Tensor | None:
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         loader = AutoWeightsLoader(self)
-        loaded_params = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
-        self.process_weights_after_loading()
-        return loaded_params
+        return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
 
     def process_weights_after_loading(self) -> None:
         self.model.finalize_mega_moe_weights()
diff --git a/vllm/models/deepseek_v41/amd/mono/README.md b/vllm/models/deepseek_v41/amd/mono/README.md
index 2f9be2e15ef4..3c1ebde3cce3 100644
--- a/vllm/models/deepseek_v41/amd/mono/README.md
+++ b/vllm/models/deepseek_v41/amd/mono/README.md
@@ -25,7 +25,7 @@ new, reading vLLM's KV records and weights as loaded.
   the block's non-causal SWA rows).
   The paths interleave freely (same tensors at the layer boundary).
 - Steps: decode only (FULL graph or eager), M <= 48 rows (8 requests x 6
-  DSpark tokens), TP2 / TP4, `fp8_ds_mla` KV records (ROCm), text tokens (the
+  DSpark tokens), TP2 / TP4 / TP8 (gfx950), `fp8_ds_mla` KV records (ROCm), text tokens (the
   routing bias is `e_score_correction_bias`).
 - Contract: a drop-in for `DeepseekV4DecoderLayer.forward`:
   in  `x` (the previous FFN output, reduced, bf16 [M, 5120]), `residual`
diff --git a/vllm/models/deepseek_v41/amd/mono/attention/back.py b/vllm/models/deepseek_v41/amd/mono/attention/back.py
index 4fa23619bc6f..8985efa67e9c 100644
--- a/vllm/models/deepseek_v41/amd/mono/attention/back.py
+++ b/vllm/models/deepseek_v41/amd/mono/attention/back.py
@@ -370,7 +370,7 @@ def stage_split(c, unit):
                 fx.ptr_store(mw, st + ((hgi * nht + ht) * nkt + kt) * HEAD_TILE + r16)
             svs.append(sv)
         gpu.barrier()
-        lst = st + d.heads * nkt
+        lst = st + d.heads_pad * nkt
         for ht in range_constexpr(nht):
             hti = hgi * nht + ht
             m = fx.ptr_load(st + (hti * nkt) * HEAD_TILE + r16)
@@ -489,13 +489,22 @@ def stage_split(c, unit):
                 w2 = v4.to(fx.BFloat16).bitcast(fx.Int32)
                 words += [w2[0], w2[1]]
             off = (((t * nchunk + ch) * d.heads + hh) * HEAD_DIM + dbase + 16 * g) // 2
-            for hlf in range_constexpr(2):
-                bo.buffer_store(
-                    fx.Vector.from_elements(words[4 * hlf : 4 * hlf + 4], fx.Int32),
-                    rsrc(c["part"]),
-                    off + 4 * hlf,
-                    cache_modifier=CM_DEV,
-                )
+
+            def store_part(off=off, words=words):
+                for hlf in range_constexpr(2):
+                    bo.buffer_store(
+                        fx.Vector.from_elements(words[4 * hlf : 4 * hlf + 4], fx.Int32),
+                        rsrc(c["part"]),
+                        off + 4 * hlf,
+                        cache_modifier=CM_DEV,
+                    )
+
+            if const_expr(d.heads < d.heads_pad):
+                # TP8: the padded columns of the head tile are not stored
+                if hh < d.heads:
+                    store_part()
+            else:
+                store_part()
         publish(c, c["srdy"], t * nchunk + ch)
         stamp(c, 7)
 
@@ -720,10 +729,10 @@ def back_lds_members(s: int, d: Dims) -> dict:
 
     @fx.struct
     class SplitLds:
-        q: fx.Array[fx.Int32, d.heads * Q_ROW, 16]
+        q: fx.Array[fx.Int32, d.heads_pad * Q_ROW, 16]
         raw: fx.Array[fx.Int32, ck * REC // 4, 16]
-        p: fx.Array[fx.Int32, d.heads * p_row(d, s), 16]
-        st: fx.Array[fx.Float32, 2 * d.heads * (WAVES // head_groups(d, s)), 16]
+        p: fx.Array[fx.Int32, d.heads_pad * p_row(d, s), 16]
+        st: fx.Array[fx.Float32, 2 * d.heads_pad * (WAVES // head_groups(d, s)), 16]
 
     @fx.struct
     class GemvLds:
diff --git a/vllm/models/deepseek_v41/amd/mono/attention/plan.py b/vllm/models/deepseek_v41/amd/mono/attention/plan.py
index 4eb137e8be42..a5954ff4d600 100644
--- a/vllm/models/deepseek_v41/amd/mono/attention/plan.py
+++ b/vllm/models/deepseek_v41/amd/mono/attention/plan.py
@@ -57,7 +57,7 @@ class Dims:
     tp: int
 
     def __post_init__(self):
-        assert self.tp in (2, 4), self.tp
+        assert self.tp in (2, 4, 8), self.tp
 
     @property
     def heads(self) -> int:
@@ -65,7 +65,16 @@ def heads(self) -> int:
 
     @property
     def head_tiles(self) -> int:
-        return self.heads // HEAD_TILE
+        """16-head MFMA tiles a rank: TP8's 8 heads take one, half of it
+        padding (``heads_pad``)."""
+        return cdiv(self.heads, HEAD_TILE)
+
+    @property
+    def heads_pad(self) -> int:
+        """The heads the split stage's LDS and MFMAs hold: whole head tiles.
+        The padded heads' columns are computed (each head is its own score
+        maximum and its own PV column) but never stored."""
+        return self.head_tiles * HEAD_TILE
 
     @property
     def groups(self) -> int:
diff --git a/vllm/models/deepseek_v41/amd/mono/layer.py b/vllm/models/deepseek_v41/amd/mono/layer.py
index f4f976b9a33e..82f85a7fa262 100644
--- a/vllm/models/deepseek_v41/amd/mono/layer.py
+++ b/vllm/models/deepseek_v41/amd/mono/layer.py
@@ -631,6 +631,7 @@ def mark(point, core_clock=False):
         s, rs, lds.moe.peek(), amb, bases, mlayout, scratch, margs, rank, own
     )
     cb["x_cm"], cb["x_ready"] = CM_DEV, True
+    cb["push_bases"] = a.get("push_bases") if tp == 4 else None
     cb["tag"] = ep & 255
     for i, region in enumerate(COUNTERS):
         cb[region] = sreg(a["epoch"], 4 * (EPOCH_COUNTERS + i * moe.TAGS), region)
@@ -859,6 +860,7 @@ def push(t, col, v0, v1):
             out=out,
             scratch=scratch,
             epoch=epoch,
+            push_bases=lambda: [b + par for b in peer_bases(peers, tp)],
         )
         run_ffn(
             key, lds, ffn, amb, bases, own, rank, ep, tl=tl if key.timeline else None
diff --git a/vllm/models/deepseek_v41/amd/mono/runner.py b/vllm/models/deepseek_v41/amd/mono/runner.py
index 6cf144db52f5..46743f63cf34 100644
--- a/vllm/models/deepseek_v41/amd/mono/runner.py
+++ b/vllm/models/deepseek_v41/amd/mono/runner.py
@@ -22,6 +22,8 @@
 
 import torch
 
+from vllm.v1.attention.ops.flydsl_kernels import _run_compiled
+
 from .attention.front import FRONT_POINTS
 from .attention.plan import HEAD_DIM, HIDDEN, KEYS, Q_RANK, WINDOW, Dims
 from .common.arch import GFX942
@@ -386,7 +388,8 @@ def forward(
         else:
             comp, comp_stride, comp_block = self._dummy, 0, 1
             bt, bt_stride, topk = self._dummy, 0, self._dummy
-        k1(
+        _run_compiled(
+            k1,
             residual.data_ptr(),
             x.data_ptr(),
             post_mix.data_ptr(),
@@ -427,7 +430,7 @@ def forward(
             self.scratch(M).data_ptr(),
             self.epoch.data_ptr(),
             self.tl1.data_ptr() if self.timeline else 0,
-            stream=st,
+            st,
         )
         if select is not None:
             select(self.normed[:M], self.qr[:M])
@@ -450,7 +453,8 @@ def forward(
                     comp_block_table=bt,
                     comp_block=comp_block,
                 )
-        k2(
+        _run_compiled(
+            k2,
             self.q.data_ptr(),
             swa_cache.data_ptr(),
             swa_cache.stride(0),
@@ -505,7 +509,7 @@ def forward(
             self.rank,
             self.epoch.data_ptr(),
             self.tl2.data_ptr() if self.timeline else 0,
-            stream=st,
+            st,
         )
         return outs
 
@@ -533,7 +537,8 @@ def ffn(
         if outs is None:
             outs = self._outs(M, residual)
         out, res_out, post_out, comb_out, pre_out = outs
-        self.ffn_kernel(M, w.gate_w.shape[0], topk)(
+        _run_compiled(
+            self.ffn_kernel(M, w.gate_w.shape[0], topk),
             part.data_ptr(),
             residual.data_ptr(),
             post_mix.data_ptr(),
@@ -564,6 +569,6 @@ def ffn(
             self.rank,
             self.epoch.data_ptr(),
             self.tl3.data_ptr() if self.timeline else 0,
-            stream=torch.cuda.current_stream(),
+            torch.cuda.current_stream(),
         )
         return outs
diff --git a/vllm/models/deepseek_v41/amd/mono/stages/gemv.py b/vllm/models/deepseek_v41/amd/mono/stages/gemv.py
index 0627071ea073..80751f81e18b 100644
--- a/vllm/models/deepseek_v41/amd/mono/stages/gemv.py
+++ b/vllm/models/deepseek_v41/amd/mono/stages/gemv.py
@@ -96,8 +96,10 @@ def _fp8_step_load(c, r_w, r_ws, k, rg, st, row_major=False):
     shuffle lays a row group out in 32-column blocks of 16 rows, a chunk two of
     them (lanes 0-31, 32-63); ``k`` % 64 = 32 (a last quarter step): the lanes
     of a block past the row read its last block. ``row_major`` (vLLM's dense
-    layout, k % 64 == 0): the same registers from a row-major [N, K] weight --
-    lane l row 16 rg + l % 16, K bytes 128 st + 16 (l // 16) .. and + 64."""
+    layout, k % 32 == 0): the same registers from a row-major [N, K] weight --
+    lane l row 16 rg + l % 16, K bytes 128 st + 16 (l // 16) .. and + 64; with
+    k % 64 == 32 (TP8's shared expert, K = 288) the last step's upper lanes are
+    clamped into the row, as in the shuffled layout."""
     lane = c["lane"]
     j = lane // 16
     kg = k // 32
@@ -105,9 +107,15 @@ def _fp8_step_load(c, r_w, r_ws, k, rg, st, row_major=False):
         assert row_major, "gfx942 reads the runner's row-major FP8 copies only"
         return _fp8_step_load_942(r_w, r_ws, k, rg, st, lane)
     if const_expr(row_major):
-        assert k % 64 == 0
+        assert k % 32 == 0
         row = rg * 16 + lane % 16
-        lo_at = row * (k // 4) + st * 32 + j * 4
+        if const_expr(k % 64 == 0):
+            lo_at = row * (k // 4) + st * 32 + j * 4
+        else:
+            # K % 64 == 32 (TP8's shared expert, K = 288): the last step's
+            # upper lanes clamped into the row (``_fp8_step_mfma`` zeroes
+            # them on both sides)
+            lo_at = row * (k // 4) + fx.min(st * 32 + j * 4, k // 4 - 4)
         hi_at = row * (k // 4) + fx.min(2 * st + 1, k // 64 - 1) * 16 + j * 4
         lo = fx.Vector(
             bo.buffer_load(r_w, lo_at, vec_width=4, dtype=T.i32, cache_modifier=CM_NT)
diff --git a/vllm/models/deepseek_v41/amd/mono/stages/moe.py b/vllm/models/deepseek_v41/amd/mono/stages/moe.py
index 3c7e10bc8b19..2914bb31dd94 100644
--- a/vllm/models/deepseek_v41/amd/mono/stages/moe.py
+++ b/vllm/models/deepseek_v41/amd/mono/stages/moe.py
@@ -427,13 +427,15 @@ def stage_route(c, t):
             for p in range_constexpr(ROUTER_PARTS):
                 tot = tot + parts[p].bitcast(fx.Float32)
             logits.append(tot)
+        neg = fx.Float32(float("-inf"))
         vals, orig, idxs = [], [], []
         for i in range_constexpr(ept):
             e = lane + 64 * i
             sc = _score(logits[i])
             # aiter's NaN guards: a NaN score weighs 0 and never wins
             orig.append(fx.Float32(fx.isnan(sc).select(fx.Float32(0.0), sc)))
-            vals.append(fx.max(sc + ld_f32(a["bias"], e), fx.Float32(float("-inf"))))
+            val = fx.max(sc + ld_f32(a["bias"], e), neg)
+            vals.append(fx.isnan(val).select(neg, val))
             idxs.append(fx.Int32(e))
         for p, q in c["rs"].sort_net:
             sw = vals[p] < vals[q]
@@ -453,7 +455,7 @@ def stage_route(c, t):
                     at.select(orig[i], co),
                 )
             best = butterfly(cv, (32, 16, 8, 4, 2, 1), fx.max)
-            ball = ballot(cv == best)
+            ball = ballot((cv == best) & (cursor < ept))
             win = fx.Int32(
                 (ball == fx.Int64(0)).select(fx.Int32(0), fx.Int32(fmath.cttz(ball)))
             )
@@ -1067,9 +1069,16 @@ def load_x8_rows(c, words_mb, codes_mb, xl, xsl, token_of, live):
     if tid < TILE:
         c["poll"]([(c["x8rdy"], token_of(fx.min(tid, live - 1)), 1)])
     gpu.barrier()
+    got, x8all = [], c["push_bases"] and not GFX942 and isinstance(live, int)
     for mb, width, dst in ((words_mb, X8_WORDS, xl), (codes_mb, HIDDEN // 32, xsl)):
         quads = width // 4
-        for i in range_constexpr(-(-TILE * quads // THREADS)):
+        for i in range_constexpr(-(-live * quads // THREADS) if x8all else 0):
+            q = fx.Uint32(fx.min(tid + THREADS * i, live * quads - 1))
+            row = fx.Int32(q // quads)
+            r, off = rsrc(c[mb].value), token_of(row) * width + fx.Int32(q % quads) * 4
+            x = fx.Vector(bo.buffer_load(r, off, dtype=T.i32, cache_modifier=CM_DEV))
+            got.append((tid + THREADS * i, row, x, dst, lds_row(width), quads))
+        for i in range_constexpr(0 if x8all else -(-TILE * quads // THREADS)):
             q = tid + THREADS * i
             row = q // quads
             if (q < TILE * quads) & (row < live):
@@ -1083,6 +1092,9 @@ def load_x8_rows(c, words_mb, codes_mb, xl, xsl, token_of, live):
                     )
                 )
                 fx.ptr_store(v, dst + (row * lds_row(width) + q % quads * 4))
+    for q, row, x, dst, stride, quads in got:
+        if q < live * quads:
+            fx.ptr_store(x, dst + (row * stride + fx.Int32(fx.Uint32(q) % quads) * 4))
     gpu.barrier()
 
 
@@ -1609,6 +1621,7 @@ def stage_down(c, tasks):
                                 avs[j * slots + i],
                                 us[i] < nu,
                             )
+    c = dict(c, peer_addr=c["push_bases"]().__getitem__) if c["push_bases"] else c
     gpu.barrier()
     for t0, nn in token_tiles(s):
         for j in range_constexpr(n):
@@ -2014,6 +2027,8 @@ def moe_context(s, rs, lds, mb, bases, layout, scratch, args, rank, sym):
         # x (``ffn_normed``) before this launch: plain loads, no readiness wait
         "x_cm": 0,
         "x_ready": False,
+        # TP4 K2: the push's peer bases, loaded late; also batches the X8 copies
+        "push_bases": None,
     }
     for name, off in stage_lds(s, rs)["offsets"].items():
         c[name] = lds.stage.ptr + off
diff --git a/vllm/models/deepseek_v41/amd/mono_decode.py b/vllm/models/deepseek_v41/amd/mono_decode.py
index bc3fdce616b2..25e4fb88444e 100644
--- a/vllm/models/deepseek_v41/amd/mono_decode.py
+++ b/vllm/models/deepseek_v41/amd/mono_decode.py
@@ -354,8 +354,8 @@ def _create(
         why = None
         if get_cdna_version() not in CDNA_VERSIONS:
             why = f"needs CDNA{'/'.join(map(str, CDNA_VERSIONS))}"
-        elif get_tensor_model_parallel_world_size() not in (2, 4):
-            why = "needs tensor parallel size 2 or 4"
+        elif get_tensor_model_parallel_world_size() not in (2, 4, 8):
+            why = "needs tensor parallel size 2, 4 or 8"
         else:
             why = _unsupported_moe(layer, vllm_config)
         if why is None:
diff --git a/vllm/models/deepseek_v41/amd/ops/__init__.py b/vllm/models/deepseek_v41/amd/ops/__init__.py
index 7ed592c0ba3a..3be417427779 100644
--- a/vllm/models/deepseek_v41/amd/ops/__init__.py
+++ b/vllm/models/deepseek_v41/amd/ops/__init__.py
@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""gfx942 (MI300X, MI325X) decode helpers for DeepSeek-V4.1.
+"""gfx942 (MI300X, MI325X) decode helpers for DeepSeek-V4.1. swa_decode_meta
+also runs on gfx950 (MI355X).
 
 ``topk`` holds the sparse indexer's top-512 hooks, built from
 ``topk512_gfx942.cu``, and ``cand_logits`` the indexer logits of the DSpark
@@ -8,7 +9,8 @@
 decode steps. ``topk_pack`` and ``swa_decode_meta`` build the decode
 attention metadata, ``fused_markov`` is DSpark's Markov sampling in one
 launch, and ``draft_lm_head`` keeps the DSpark draft's 5-row logits GEMM off
-wvSplitK. All of them run only when ``enabled()`` is true.
+wvSplitK. All of them run only when ``enabled()`` is true, except
+``swa_decode_meta``, which runs when ``swa_decode_meta_enabled()`` is true.
 """
 
 import functools
@@ -28,15 +30,29 @@ def record_monokernel(vllm_config) -> None:
     global _monokernel
     _monokernel = bool(vllm_config.kernel_config.monokernel)
     enabled.cache_clear()
+    swa_decode_meta_enabled.cache_clear()
 
 
 @functools.cache
 def enabled() -> bool:
     """True with --monokernel on a gfx942 GPU. Every helper keeps the result
     of vLLM's op it replaces, so each one returns False for a call it cannot
-    take and vLLM then runs its own op. gfx950 always keeps vLLM's ops here."""
+    take and vLLM then runs its own op. gfx950 keeps vLLM's ops here, except
+    for swa_decode_meta (swa_decode_meta_enabled)."""
     if not _monokernel or not current_platform.is_rocm():
         return False
     from vllm.platforms.rocm import on_gfx942
 
     return on_gfx942()
+
+
+@functools.cache
+def swa_decode_meta_enabled() -> bool:
+    """True with --monokernel on a CDNA3 (gfx942) or CDNA4 (gfx950) GPU.
+    swa_decode_meta is plain Triton and writes the buffers that the attention
+    of both reads."""
+    if not _monokernel or not current_platform.is_rocm():
+        return False
+    from vllm.platforms.rocm import get_cdna_version
+
+    return get_cdna_version() in (3, 4)
diff --git a/vllm/models/deepseek_v41/amd/ops/fused_markov.py b/vllm/models/deepseek_v41/amd/ops/fused_markov.py
index 72408eeb96b3..aaf8e389019b 100644
--- a/vllm/models/deepseek_v41/amd/ops/fused_markov.py
+++ b/vllm/models/deepseek_v41/amd/ops/fused_markov.py
@@ -225,8 +225,8 @@ def _why_not(spec, num_reqs: int) -> str | None:
         return "Markov weights not bf16"
     if w1.shape[1] != w2.shape[1] or not (w1.is_contiguous() and w2.is_contiguous()):
         return "Markov weight layout"
-    if spec.draft_logits.dtype != torch.bfloat16:
-        return f"a {spec.draft_logits.dtype} draft logits cache"
+    if spec.draft_logits.values.dtype != torch.bfloat16:
+        return f"a {spec.draft_logits.values.dtype} draft logits cache"
     if head.markov_w2.quant_method.__class__.__name__ != "UnquantizedEmbeddingMethod":
         return "a quantized Markov head"
     return None
@@ -260,7 +260,7 @@ def sample_sequential(spec, num_reqs: int, head_hidden: torch.Tensor) -> bool:
     idx_map = spec.sample_idx_mapping[:num_sample].view(num_reqs, n_spec)
     sample_pos = spec.sample_pos[:num_sample].view(num_reqs, n_spec)
     assert idx_map.stride(1) == 1 and sample_pos.stride(1) == 1
-    cache = spec.draft_logits
+    cache = spec.draft_logits.values
     tokens = spec.draft_tokens
     use_fp64 = bool(spec.use_fp64_gumbel)
     bufs = _scratch(base_logits.device, vocab, use_fp64)
diff --git a/vllm/models/deepseek_v41/amd/ops/swa_decode_meta.py b/vllm/models/deepseek_v41/amd/ops/swa_decode_meta.py
index 7fd6e277de14..0cd87063d2ff 100644
--- a/vllm/models/deepseek_v41/amd/ops/swa_decode_meta.py
+++ b/vllm/models/deepseek_v41/amd/ops/swa_decode_meta.py
@@ -22,16 +22,17 @@
 ComputeDSparkNoncausalSWAIndicesKernel in place of the window kernel. That
 is 3 more builds of 8 launches after every target step.
 
-On gfx942 with --monokernel (``ops.enabled()``), decode
-batches of up to MAX_ROWS rows, causal or the draft block, use it. Other
-batches take the rest of DeepseekV41ROCMAiterSparseSWAMetadataBuilder.build.
+On gfx942 and gfx950 with --monokernel
+(``ops.swa_decode_meta_enabled()``), decode batches of up to MAX_ROWS rows,
+causal or the draft block, use it. Other batches take the rest of
+DeepseekV41ROCMAiterSparseSWAMetadataBuilder.build.
 tests/kernels/models/amd/test_dsv41_gfx942_decode_meta.py compares it with the vLLM ops
 that it replaces.
 """
 
 import torch
 
-from vllm.models.deepseek_v41.amd.ops import enabled
+from vllm.models.deepseek_v41.amd.ops import swa_decode_meta_enabled
 from vllm.triton_utils import tl, triton
 
 # One program holds all rows of the batch in registers, as a ROWS x WIDTH
@@ -182,7 +183,7 @@ def build(builder, common_prefix_len, cam, fast_build, replay_start, metadata_cl
     """Return the ROCm SWA metadata of a decode batch, causal or the DSpark
     draft block, built with one launch, or None when the batch needs the rest
     of DeepseekV41ROCMAiterSparseSWAMetadataBuilder.build."""
-    if not enabled():
+    if not swa_decode_meta_enabled():
         return None
     from vllm.v1.attention.backends.mla import sparse_swa
     from vllm.v1.attention.backends.utils import split_decodes_and_prefills
diff --git a/vllm/models/deepseek_v41/amd/rocm.py b/vllm/models/deepseek_v41/amd/rocm.py
index de59e04da05d..3d694b2b9106 100644
--- a/vllm/models/deepseek_v41/amd/rocm.py
+++ b/vllm/models/deepseek_v41/amd/rocm.py
@@ -380,7 +380,7 @@ def build(
         fast_build: bool = False,
         replay_start: torch.Tensor | None = None,
     ) -> DeepseekV41ROCMAiterSparseSWAMetadata:
-        # On gfx942 with --monokernel, a decode batch gets this
+        # On gfx942 and gfx950 with --monokernel, a decode batch gets this
         # metadata from one launch instead of 8 launches per SWA group
         # (swa_decode_meta.py). Other batches return None there and take
         # the build below.
diff --git a/vllm/models/deepseek_v41/amd/vl_model.py b/vllm/models/deepseek_v41/amd/vl_model.py
index 60e4e55cde26..d11c93a8648b 100644
--- a/vllm/models/deepseek_v41/amd/vl_model.py
+++ b/vllm/models/deepseek_v41/amd/vl_model.py
@@ -41,7 +41,6 @@
     DeepseekV4ViT,
     run_dp_sharded_vision_tower,
 )
-from vllm.models.deepseek_v4.common.vl_model import stream_language_model_first
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.utils.tensor_schema import TensorSchema, TensorShape
 
@@ -336,19 +335,5 @@ def get_mtp_target_hidden_states(self) -> torch.Tensor | None:
         return self.language_model.get_mtp_target_hidden_states()
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        # The child finalizes in load_weights, so its weights must reach it as
-        # one contiguous group.
-        mapped = stream_language_model_first(self.hf_to_vllm_mapper.apply(weights))
         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()
+        return loader.load_weights(self.hf_to_vllm_mapper.apply(weights))
diff --git a/vllm/models/deepseek_v41/attention.py b/vllm/models/deepseek_v41/attention.py
index 778920aefc03..b7d815ebbc6c 100644
--- a/vllm/models/deepseek_v41/attention.py
+++ b/vllm/models/deepseek_v41/attention.py
@@ -60,6 +60,12 @@
 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.model_executor.warmup.jit_warmup_triton_helper import (
+    DispatchSpec,
+    TritonWarmupTensor,
+    triton_kernel_dispatcher_with_warmup,
+)
+from vllm.models.deepseek_v41.common.ops.fused_layout import Q_CHUNK
 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
@@ -90,6 +96,78 @@
 # shard per TP rank and all-gathers the result.
 _SHARD_FUSED_WQA_WKV_MIN_TOKENS = 2048
 
+# From this many tokens, eager steps build the padded interleaved Q with
+# _pad_interleaved_q instead of in the fused KV-insert op. Set from a
+# 64-16384-token sweep on B300: by GPU time alone the copy wins from about
+# 256 tokens, but it costs one extra eager launch, and with that launch gap
+# fully exposed the crossover moves to about 2048. This value keeps the
+# large-step win without betting on launch overlap.
+_Q_PAD_COPY_MIN_TOKENS = 2048
+
+
+@triton.jit
+def _pad_interleaved_q_kernel(
+    q_ptr,
+    out_ptr,
+    num_rows,
+    LIVE: tl.constexpr,
+    PAD: tl.constexpr,
+    BLOCK_ROWS: tl.constexpr,
+):
+    # One row per (token, 16-dim chunk): LIVE copied elements, then PAD zeros.
+    pid = tl.program_id(0)
+    rows = pid.to(tl.int64) * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS).to(tl.int64)
+    row_mask = rows[:, None] < num_rows
+    cols = tl.arange(0, LIVE)
+    x = tl.load(q_ptr + rows[:, None] * LIVE + cols[None, :], mask=row_mask)
+    out_rows = out_ptr + rows[:, None] * (LIVE + PAD)
+    tl.store(out_rows + cols[None, :], x, mask=row_mask)
+    pad_cols = tl.arange(0, PAD)
+    tl.store(
+        out_rows + LIVE + pad_cols[None, :],
+        tl.zeros([BLOCK_ROWS, PAD], dtype=x.dtype),
+        mask=row_mask,
+    )
+
+
+def _pad_interleaved_q_warmup_inputs(
+    *,
+    live_heads: int,
+    padded_heads: int,
+    head_dim: int,
+    dtype: torch.dtype,
+) -> dict[str, object]:
+    return dict(
+        q=TritonWarmupTensor(dtype, shape=(1, live_heads, head_dim)),
+        out=TritonWarmupTensor(dtype, shape=(1, padded_heads, head_dim)),
+    )
+
+
+@triton_kernel_dispatcher_with_warmup(
+    kernel=_pad_interleaved_q_kernel,
+    warmup_inputs=_pad_interleaved_q_warmup_inputs,
+)
+def _pad_interleaved_q(q: torch.Tensor, out: torch.Tensor) -> DispatchSpec:
+    """Copy chunk-interleaved ``q`` [N, H, D] into ``out`` [N, HP, D].
+
+    Writes the fused KV-insert op's interleaved Q layout: in each 16-dim chunk,
+    the H live heads copied bit for bit, then +0.0 in the HP - H padding slots.
+    Both tensors are dense and share a dtype; ``16 * H`` and ``16 * (HP - H)``
+    must be powers of two.
+    """
+    num_tokens, live_heads, head_dim = q.shape
+    padded_heads = out.shape[1]
+    assert out.shape == (num_tokens, padded_heads, head_dim) and out.dtype == q.dtype
+    num_rows = num_tokens * (head_dim // Q_CHUNK)
+    block_rows = 4
+    return (triton.cdiv(num_rows, block_rows),), dict(
+        num_rows=num_rows,
+        LIVE=live_heads * Q_CHUNK,
+        PAD=(padded_heads - live_heads) * Q_CHUNK,
+        BLOCK_ROWS=block_rows,
+        num_warps=4,
+    )
+
 
 def _replace_layer_index(prefix: str, layer_id: int) -> str:
     """Swap the layer index inside a ``...layers.<idx>...`` prefix."""
@@ -550,15 +628,21 @@ def __init__(
         ):
             logger.warning_once("SWA bounded replay with PCP requires prefix caching.")
             swa_bounded_replay = False
+        # CDNA4 (MI350/MI355, gfx950) is where prefix-cache SWA replay is
+        # validated. The AITER sparse-MLA path and the default Triton gather
+        # path both clamp to replay_start. CDNA3 and other ROCm archs have
+        # not been validated, so the sliding-window cache stays in the prefix
+        # cache there. The CDNA version is imported only on ROCm: loading the
+        # ROCm platform resolves the GPU arch and fails when no GPU is visible.
         if swa_bounded_replay and current_platform.is_rocm():
-            logger.warning_once(
-                "Prefix-cache SWA 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
+            from vllm.platforms.rocm import get_cdna_version
+
+            if get_cdna_version() != 4:
+                logger.warning_once(
+                    "Prefix-cache SWA replay is only enabled on ROCm CDNA4; 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,
@@ -614,6 +698,15 @@ def __init__(
         else:
             self.compressed_cache_prefix = None
 
+        # _pad_interleaved_q needs power-of-two chunk widths: 32 live heads of
+        # 64 (e.g. TP2) qualify, 16 or 8 (768 or 896 pad columns) do not.
+        self._pads_q_by_copy = (
+            self.accepts_unnormed_unroped_query
+            and self.n_local_heads == 32
+            and self.padded_heads == 64
+            and self.max_num_batched_tokens >= _Q_PAD_COPY_MIN_TOKENS
+        )
+
         if vllm_config.kernel_config.enable_jit_warmup:
             from vllm.v1.attention.backends.mla.sparse_swa import (
                 _COMPUTE_PREFILL_METADATA_KERNEL,
@@ -626,6 +719,14 @@ def __init__(
                 block_size=self.swa_cache_layer.block_size,
             )
 
+            if self._pads_q_by_copy:
+                _pad_interleaved_q.register_warmup(
+                    live_heads=self.n_local_heads,
+                    padded_heads=self.padded_heads,
+                    head_dim=self.head_dim,
+                    dtype=vllm_config.model_config.dtype,
+                )
+
             if self.compress_ratio > 1:
                 from vllm.v1.attention.backends.mla.compressor_utils import (
                     _COMPRESSED_SLOT_MAPPING_KERNEL,
@@ -1072,6 +1173,8 @@ def _fused_qnorm_rope_kv_insert(
         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.
+        On large eager steps of a ``_pads_q_by_copy`` layer, a KV-only launch
+        leaves that pad to ``_pad_interleaved_q``.
         """
         if not isinstance(attn_metadata, dict):
             # Profile run: kernel doesn't fire; produce a padded tensor so
@@ -1137,6 +1240,12 @@ def _fused_qnorm_rope_kv_insert(
                 and self.n_local_heads == self.padded_heads
                 else self.padded_heads
             )
+            # Captured graphs keep the fused pad.
+            copy_q = (
+                self._pads_q_by_copy
+                and q.shape[0] >= _Q_PAD_COPY_MIN_TOKENS
+                and not torch.cuda.is_current_stream_capturing()
+            )
             q_padded = torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
                 q,
                 kv,
@@ -1144,7 +1253,7 @@ def _fused_qnorm_rope_kv_insert(
                 slot_mapping,
                 positions,
                 cos_sin_cache,
-                pad_to,
+                0 if copy_q else pad_to,
                 self.eps,
                 swa_metadata.block_size,
                 False,  # apply_q_norm: qr is normed before wq_b
@@ -1152,6 +1261,10 @@ def _fused_qnorm_rope_kv_insert(
                 not self.accepts_unnormed_unroped_query,  # apply_q_rope
                 self.accepts_unnormed_unroped_query,  # is_q_interleaved
             )
+            if copy_q:
+                out = q.new_empty((q.shape[0], self.padded_heads, q.shape[2]))
+                _pad_interleaved_q(q, out)
+                return out
             return q if pad_to == 0 else q_padded
 
         assert not self.accepts_unnormed_unroped_query, (
@@ -1604,3 +1717,18 @@ def _wq_b_proj(
         # ReplicatedLinear returns (output, bias); bias is None.
         q, _ = self.wq_b(qr)
         return q
+
+
+def autotune_indexer_wq_b(model: nn.Module, num_tokens: int) -> None:
+    """Short-context dummy runs skip wq_b, so tune it directly."""
+    for module in model.modules():
+        if isinstance(module, DeepseekV4Indexer):
+            # Run every wq_b: FlashInfer keys on op, runner, shapes and dtype,
+            # and skips profiling on a cache hit, so duplicates are cheap.
+            qr = torch.randn(
+                num_tokens,
+                module.q_lora_rank,
+                dtype=module.vllm_config.model_config.dtype,
+                device=next(module.weights_proj.parameters()).device,
+            )
+            module.wq_b(qr)
diff --git a/vllm/models/deepseek_v41/common/engram.py b/vllm/models/deepseek_v41/common/engram.py
index ecbcc46c0703..b88293037ea2 100644
--- a/vllm/models/deepseek_v41/common/engram.py
+++ b/vllm/models/deepseek_v41/common/engram.py
@@ -37,6 +37,7 @@
 import ctypes
 import mmap
 import os
+import platform
 import sys
 import tempfile
 import weakref
@@ -75,6 +76,7 @@
 )
 from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
+from vllm.utils.cpu_resource_utils import get_available_memory
 from vllm.utils.platform_utils import is_uva_available
 from vllm.utils.torch_utils import get_accelerator_view_from_cpu_tensor
 
@@ -778,46 +780,132 @@ def gather_engram_hashes(
     return dp_group.all_gather(hash_ids, dim=0)
 
 
+def describe_error(exc: Exception) -> str:
+    return f"{type(exc).__name__}: {exc}"
+
+
+def gather_errors(group: GroupCoordinator, error: str | None) -> str:
+    """Each rank's error, joined; empty if every rank succeeded."""
+    errors: list[str | None] = [None] * group.world_size
+    torch.distributed.all_gather_object(errors, error, group=group.cpu_group)
+    return "; ".join(
+        f"EDP rank {rank}: {error}" for rank, error in enumerate(errors) if error
+    )
+
+
+def _shm_shortfall(num_bytes: int) -> str | None:
+    """Why /dev/shm cannot hold a table of this size, or None if it can."""
+    if not os.path.isdir(SHM_PATH):
+        return f"{SHM_PATH} is not mounted"
+    try:
+        check_shm_free_space(num_bytes, allocation_name="Engram tables")
+    except Exception as exc:
+        return describe_error(exc)
+    # tmpfs pages beyond available host memory are faulted in by the OOM killer.
+    available = get_available_memory()
+    if num_bytes > available:
+        return (
+            f"{num_bytes / 2**30:.1f} GiB exceeds the "
+            f"{available / 2**30:.1f} GiB of available host memory"
+        )
+    return None
+
+
+def _allow_any_numa_node(address: int, num_bytes: int) -> None:
+    """Let a mapping's pages spill to any NUMA node despite --numa-bind.
+
+    mbind(MPOL_LOCAL) overrides the task's membind for this range: pages still
+    go to the faulting CPU's node first, then to any other node.
+    """
+    sys_mbind = {"aarch64": 235, "x86_64": 237}.get(platform.machine())
+    libc = ctypes.CDLL(None, use_errno=True)
+    libc.syscall.restype = ctypes.c_long
+    if sys_mbind is None or libc.syscall(
+        ctypes.c_long(sys_mbind),
+        ctypes.c_void_p(address),
+        ctypes.c_ulong(num_bytes),
+        ctypes.c_int(4),  # MPOL_LOCAL
+        None,
+        ctypes.c_ulong(0),
+        ctypes.c_uint(0),
+    ):
+        logger.warning_once(
+            "Engram mbind(MPOL_LOCAL) failed; /dev/shm tables follow the "
+            "process memory policy: %s",
+            os.strerror(ctypes.get_errno()),
+        )
+
+
 class DPSharedEngramStorage:
-    """Registered host weights shared by a node-local DP group with one writer."""
+    """Host weights shared by a node-local DP group with one writer: host-NUMA
+    cuMem where FABRIC handles work, else a registered /dev/shm mapping."""
 
     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
+        num_bytes = weight_bytes + weight_bytes // block_size
+
+        def split_table(table: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+            weight = table[:weight_bytes].view(torch.float8_e4m3fn).view(-1, dim)
+            return weight, table[weight_bytes:].view(-1, dim // block_size)
 
-    def _allocate(self, num_bytes: int) -> torch.Tensor:
-        """Map and register one physical allocation across a node-local DP group."""
+        views = None
+        if current_platform.is_cuda():
+            from vllm.models.deepseek_v41.nvidia.ops import engram_host_table
+
+            views = engram_host_table.share_host_table(num_bytes, group)
+        # The GPU maps cuMem host memory with 2M pages, /dev/shm with 64K pages.
+        self.huge_pages = views is not None
+        host, gpu = views or self._allocate_shm(num_bytes)
+        self.weight, self.weight_scale_inv = split_table(host)
+        self._views = split_table(gpu)
+
+    def _allocate_shm(self, num_bytes: int) -> tuple[torch.Tensor, torch.Tensor]:
+        """Map and register one /dev/shm file across a node-local DP group.
+
+        Returns CPU and GPU views of the file.
+        """
         group = self.group
         with ExitStack() as stack:
             path = error = None
-            if group.rank_in_group == 0:
+            if group.rank_in_group == 0 and not (error := _shm_shortfall(num_bytes)):
                 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}"
+                    error = describe_error(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
+            mapping = None
             try:
                 with open(path, "r+b") as file:
                     mapping = mmap.mmap(file.fileno(), num_bytes, flags=mmap.MAP_SHARED)
+            except Exception as exc:
+                error = describe_error(exc)
+            if failures := gather_errors(group, error):
+                del mapping
+                raise RuntimeError(
+                    "Engram shared-memory initialization failed: " + failures
+                )
+            # Unlink while the file is still sparse, so a crash during the slow
+            # registration below cannot orphan the table.
+            stack.close()
+            assert mapping is not None
+
+            owner = tensor = None
+            try:
                 owner = np.frombuffer(mapping, dtype=np.uint8)
                 pointer = owner.ctypes.data
+                _allow_any_numa_node(pointer, num_bytes)
                 tensor = torch.from_numpy(owner)
                 result = torch.cuda.cudart().cudaHostRegister(pointer, num_bytes, 0)
                 if result.value != 0:
@@ -830,24 +918,16 @@ def _allocate(self, num_bytes: int) -> torch.Tensor:
                         "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:
+                error = describe_error(exc)
+
+            if failures := gather_errors(group, error):
                 # 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
+            return tensor, get_accelerator_view_from_cpu_tensor(tensor)
 
     @staticmethod
     def _unregister(mapping: mmap.mmap, pointer: int) -> None:
@@ -874,16 +954,11 @@ def get_views(
             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
 
 
 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."""
+    """Whether co-located DP replicas exist and can hold the full tables."""
     if get_engram_dp_size() == 1:
         logger.warning_once(
             "Engram DP replicas are not co-located on one node; "
@@ -895,16 +970,9 @@ def can_share_engram_tables(layout: EngramLayout, block_size: int = 32) -> bool:
     )
     group = get_engram_dp_group()
     assert group is not None
-    # Only the leader allocates, so every replica follows its check.
-    error = None
-    if group.rank_in_group == 0:
-        if not os.path.isdir(SHM_PATH):
-            error = f"{SHM_PATH} is not mounted"
-        else:
-            try:
-                check_shm_free_space(num_bytes, allocation_name="Engram tables")
-            except Exception as exc:
-                error = f"{type(exc).__name__}: {exc}"
+    # Only the leader allocates, so every replica follows its check. /dev/shm
+    # must fit even where cuMem will, since it is the fallback.
+    error = _shm_shortfall(num_bytes) if group.rank_in_group == 0 else None
     error = group.broadcast_object(error)
     if error is not None:
         logger.warning_once(
@@ -1069,7 +1137,10 @@ def __init__(
                 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",
+                "per rank"
+                if self._shared_memory is None
+                else "shared across DP replicas via "
+                + ("cuMem" if self._shared_memory.huge_pages else "/dev/shm"),
             )
 
     def _get_shard_info(self) -> tuple[int, int]:
@@ -1162,6 +1233,13 @@ def collapse_huge_pages(self) -> None:
                 os.strerror(ctypes.get_errno()),
             )
 
+    @property
+    def _uses_small_host_pages(self) -> bool:
+        """Whether the GPU maps this host table with pages too small for its TLB."""
+        if self._shared_memory is not None:
+            return not self._shared_memory.huge_pages
+        return self.cpu_offload and self._packed is None
+
     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)
@@ -1191,8 +1269,7 @@ def lookup(
         ids_stride_t, ids_stride_h = indices.stride()
         sort_min_rows, _ = _engram_lookup_thresholds(indices.device)
         sort = (
-            self.cpu_offload
-            and self._packed is None
+            self._uses_small_host_pages
             and sort_min_rows is not None
             and rows >= sort_min_rows
         )
diff --git a/vllm/models/deepseek_v41/common/ops/cache_utils.py b/vllm/models/deepseek_v41/common/ops/cache_utils.py
index 1ab4d8def091..23aa3d3fe595 100644
--- a/vllm/models/deepseek_v41/common/ops/cache_utils.py
+++ b/vllm/models/deepseek_v41/common/ops/cache_utils.py
@@ -602,8 +602,12 @@ def _dequantize_and_gather_k_nvfp4_kernel(
     quant_block: tl.constexpr,  # 16
     cache_block_size: tl.constexpr,
     block_stride: tl.constexpr,
+    TILE: tl.constexpr,  # tokens per iteration
 ):
-    """Gather and dequantize V4.1 NVFP4 rows into a bf16 workspace."""
+    """Gather and dequantize V4.1 NVFP4 rows into a bf16 workspace, TILE tokens
+    per iteration. A nibble placed as fp16 bits (sign at 15, exponent and
+    mantissa at 11:9) reads as e2m1 * 2^-14; times scale * 2^14 the fp32 product
+    is exact, so the bf16 result matches a direct decode."""
     batch_idx = tl.program_id(0)
     worker_id = tl.program_id(1)
     num_workers = tl.num_programs(1)
@@ -616,35 +620,52 @@ def _dequantize_and_gather_k_nvfp4_kernel(
     start_pos = seq_len - gather_len
 
     packed_bytes: tl.constexpr = head_dim // 2
+    group_bytes: tl.constexpr = quant_block // 2
+    # Byte b of a record packs elements 2b (low nibble) and 2b + 1, which share
+    # scale b // group_bytes: index the bytes as [scale, byte within group].
+    byte = (
+        tl.arange(0, scale_dim)[:, None] * group_bytes
+        + tl.arange(0, group_bytes)[None, :]
+    )
+    t = tl.arange(0, TILE)
     d = tl.arange(0, head_dim)
-    for i in range(worker_id, gather_len, num_workers):
+    block_table_row_ptr = block_table_ptr + batch_idx * max_blocks_per_seq
+    out_row_ptr = out_ptr + batch_idx.to(tl.int64) * out_stride0
+    for tile_start in range(worker_id * TILE, gather_len, num_workers * TILE):
+        i = tile_start + t
+        valid = i < gather_len
         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
+            block_table_row_ptr + pos // cache_block_size, mask=valid, other=0
         )
         page = k_cache_ptr + physical_block_idx.to(tl.int64) * block_stride
-        token_ptr = page + pos_in_block * (packed_bytes + scale_dim)
-
-        packed = tl.load(token_ptr + 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(token_ptr + packed_bytes + 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))
+        token_ptr = page + (pos % cache_block_size) * (packed_bytes + scale_dim)
+
+        packed = tl.load(
+            token_ptr[:, None, None] + byte[None, :, :],
+            mask=valid[:, None, None],
+            other=0,
+        ).to(tl.int32)
+        sf = tl.load(
+            token_ptr[:, None] + packed_bytes + tl.arange(0, scale_dim)[None, :],
+            mask=valid[:, None],
+            other=0,
+        )
+        scale = sf.to(tl.float8e4nv, bitcast=True).to(tl.float32) * 16384.0
+        scale = scale[:, :, None]
+
+        lo = packed & 0xF
+        hi = packed >> 4
+        lo = ((lo & 7) << 9) | ((lo & 8) << 12)
+        hi = ((hi & 7) << 9) | ((hi & 8) << 12)
+        lo = lo.to(tl.int16).to(tl.float16, bitcast=True).to(tl.float32) * scale
+        hi = hi.to(tl.int16).to(tl.float16, bitcast=True).to(tl.float32) * scale
+        dequant = tl.reshape(tl.join(lo, hi), (TILE, head_dim))
+
+        out_ptrs = (
+            out_row_ptr + (offset + i).to(tl.int64)[:, None] * out_stride1 + d[None, :]
+        )
+        tl.store(out_ptrs, dequant.to(tl.bfloat16), mask=valid[:, None])
 
 
 # The per-token loop is a chain of dependent loads, so the grid has to hold
@@ -653,6 +674,7 @@ def _dequantize_and_gather_k_nvfp4_kernel(
 _GATHER_STEPS_PER_WORKER = 8
 _GATHER_MIN_WORKERS = 128
 _GATHER_MAX_WORKERS = 8192
+_NVFP4_GATHER_TILE = 8
 
 
 def gather_num_workers(max_gather_len: int | None) -> int:
@@ -682,9 +704,11 @@ def dequantize_and_gather_k_cache_triton(
     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.
+        # One tile per worker up to the cap: the free width of ``out`` bounds
+        # every request's gather, and the kernel loops over any remainder.
         num_workers = min(
-            max(128, 2048 // max(1, num_reqs)), max(128, out.shape[1] - offset)
+            cdiv(max(1, out.shape[1] - offset), _NVFP4_GATHER_TILE),
+            _GATHER_MAX_WORKERS,
         )
         _dequantize_and_gather_k_nvfp4_kernel[(num_reqs, num_workers)](
             out,
@@ -701,6 +725,8 @@ def dequantize_and_gather_k_cache_triton(
             quant_block=V41_NVFP4_QUANT_BLOCK,
             cache_block_size=block_size,
             block_stride=k_cache.stride(0),
+            TILE=_NVFP4_GATHER_TILE,
+            num_warps=2,
         )
         return
 
diff --git a/vllm/models/deepseek_v41/common/ops/query_quant.py b/vllm/models/deepseek_v41/common/ops/query_quant.py
index 97714e29f151..729f55114e93 100644
--- a/vllm/models/deepseek_v41/common/ops/query_quant.py
+++ b/vllm/models/deepseek_v41/common/ops/query_quant.py
@@ -140,12 +140,19 @@ def can_fuse_query_quant(linears: list[torch.nn.Module]) -> bool:
         FlashInferCutedslMxfp8LinearKernel,
         FlashInferCutlassMxfp8LinearKernel,
     )
+    from vllm.model_executor.kernels.linear.mxfp8.pytorch import (
+        TorchMxFp8LinearKernel,
+    )
 
     # QuantKey does not encode scale layout; restrict this producer to the
     # consumers that accept F8_128x4 swizzled scales.
     return all(
         get_input_quant_key(linear) == kMxfp8Dynamic
         and type(getattr(getattr(linear, "quant_method", None), "kernel", None))
-        in (FlashInferCutedslMxfp8LinearKernel, FlashInferCutlassMxfp8LinearKernel)
+        in (
+            FlashInferCutedslMxfp8LinearKernel,
+            FlashInferCutlassMxfp8LinearKernel,
+            TorchMxFp8LinearKernel,
+        )
         for linear in linears
     )
diff --git a/vllm/models/deepseek_v41/decoder_replay_layers.py b/vllm/models/deepseek_v41/decoder_replay_layers.py
index b109d8634a1a..a08cd8063597 100644
--- a/vllm/models/deepseek_v41/decoder_replay_layers.py
+++ b/vllm/models/deepseek_v41/decoder_replay_layers.py
@@ -46,6 +46,8 @@ class DecoderReplayLayers:
     every row. ``metadata_prefixes`` are the attention metadata keys the layers
     read, so the replay batch builds only those. ``first_swa_prefix`` keys the
     first replay layer's SWA metadata, built with the batch's window starts.
+    ``has_early_aux`` tells whether layers before the replay layers capture aux
+    hidden states, which the forward holds across the replay layers.
     """
 
     def __init__(
@@ -55,12 +57,14 @@ def __init__(
         write_batch_kv: Callable[..., None],
         metadata_prefixes: set[str],
         first_swa_prefix: str,
+        has_early_aux: Callable[[], bool] = lambda: False,
     ) -> None:
         self.window = window
         self.run_layers = run_layers
         self.write_batch_kv = write_batch_kv
         self.metadata_prefixes = metadata_prefixes
         self.first_swa_prefix = first_swa_prefix
+        self.has_early_aux = has_early_aux
         # Set by the model state every step; None runs the layers on the batch.
         self.replay_batch: ReplayBatch | None = None
         # Set by the model state with replay graphs: PIECEWISE model graphs of
diff --git a/vllm/models/deepseek_v41/nvidia/decoder_replay_cudagraph.py b/vllm/models/deepseek_v41/nvidia/decoder_replay_cudagraph.py
index c985f973bb49..e8a8e2c33539 100644
--- a/vllm/models/deepseek_v41/nvidia/decoder_replay_cudagraph.py
+++ b/vllm/models/deepseek_v41/nvidia/decoder_replay_cudagraph.py
@@ -68,9 +68,21 @@ def __init__(self, vllm_config: VllmConfig, device, layers: DecoderReplayLayers)
         ]
 
     def capture_replay_graphs(
-        self, prepare: Callable[[BatchExecutionDescriptor], None]
+        self,
+        prepare: Callable[[BatchExecutionDescriptor], None],
+        stream: torch.cuda.Stream | None = None,
     ) -> None:
-        """Capture each size on the replay batch ``prepare`` sets on the layers."""
+        """Capture each size on the replay batch ``prepare`` sets on the layers.
+
+        On the model graphs' ``stream`` they share its free blocks of the graph
+        pool: a replay graph runs in a model graph's eager break, after it has
+        gathered its inputs into ``input_buffers`` and before the model graph
+        reads its outputs from the break's own buffers, so no model graph
+        tensor in the pool is live while it runs. Aux hidden states captured
+        before the replay layers would be, so with those it keeps its own.
+        """
+        if self.layers.has_early_aux():
+            stream = None
 
         def create_forward_fn(desc: BatchExecutionDescriptor, warmup: bool):
             prepare(desc)
@@ -86,7 +98,7 @@ def forward_fn(cg_mode: CUDAGraphMode) -> None:
 
             return forward_fn
 
-        self.capture(create_forward_fn, "Capturing decoder replay CUDA graphs")
+        self.capture(create_forward_fn, "Capturing decoder replay CUDA graphs", stream)
         self.layers.replay_batch = None
 
     def run(self, rows: torch.Tensor, inputs: tuple) -> tuple[torch.Tensor, ...]:
diff --git a/vllm/models/deepseek_v41/nvidia/dspark.py b/vllm/models/deepseek_v41/nvidia/dspark.py
index c0fbfa553c30..021b92f1c853 100644
--- a/vllm/models/deepseek_v41/nvidia/dspark.py
+++ b/vllm/models/deepseek_v41/nvidia/dspark.py
@@ -38,6 +38,9 @@
     ParallelLMHead,
     VocabParallelEmbedding,
 )
+from vllm.model_executor.model_loader.mtp_validation import (
+    is_mtp_completeness_check_enabled,
+)
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 from vllm.model_executor.models.qwen3_dspark import (
     DSparkConfidenceHead,
@@ -358,6 +361,12 @@ def compute_draft_logits(self, hidden_states: torch.Tensor) -> torch.Tensor:
         # Full-vocab draft: base logits, no d2t scatter.
         return self.compute_logits(hidden_states)
 
+    def compute_draft_logits_local(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        """This TP rank's vocab shard of compute_draft_logits, not gathered."""
+        return self.logits_processor(
+            self.lm_head, self.model.norm(hidden_states), skip_gather=True
+        )
+
     def map_draft_to_target(self, draft_ids: torch.Tensor) -> torch.Tensor:
         return draft_ids  # full-vocab: draft ids are target ids
 
@@ -503,7 +512,11 @@ 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:
+        if (
+            self.model.confidence_head is not None
+            and not loaded_confidence_head
+            and is_mtp_completeness_check_enabled()
+        ):
             self.model.confidence_head = None
         logger.info_once("DSpark draft model loaded: %d params", len(loaded_params))
         return loaded_params
diff --git a/vllm/models/deepseek_v41/nvidia/flash_mla_mega_attn.py b/vllm/models/deepseek_v41/nvidia/flash_mla_mega_attn.py
index 9d4f03abb37b..9b7f7c6e755d 100644
--- a/vllm/models/deepseek_v41/nvidia/flash_mla_mega_attn.py
+++ b/vllm/models/deepseek_v41/nvidia/flash_mla_mega_attn.py
@@ -5,7 +5,7 @@
 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
+leaving Q 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.
 
@@ -282,7 +282,7 @@ def forward_mqa(
             return
 
         assert isinstance(attn_metadata, dict)
-        # q was zero-padded to the kernel's head count by the fused KV insert.
+        # _fused_qnorm_rope_kv_insert zero-padded q to the kernel's head count.
         assert q.shape[1] == self.padded_heads
         flashmla_metadata = cast(
             DeepseekV4FlashMLAMetadata | None,
@@ -293,14 +293,12 @@ def forward_mqa(
         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:],
+                positions[num_decode_tokens:],
                 flashmla_metadata,
                 swa_metadata,
                 output,
@@ -309,7 +307,7 @@ def forward_mqa(
         if swa_metadata.num_decodes > 0:
             self._forward_decode_mega(
                 q[:num_decode_tokens],
-                positions_int32[:num_decode_tokens],
+                positions[:num_decode_tokens],
                 flashmla_metadata,
                 swa_metadata,
                 _token_slice(output, 0, num_decode_tokens),
@@ -324,9 +322,18 @@ def _reserve_prefill_workspace(self, q: torch.Tensor) -> None:
         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)
+        # One request covers the KV chunk (retained at the tail at runtime) and
+        # the index scratch (taken from the head) side by side.
         current_workspace_manager().get_simultaneous(
             ((self.PREFILL_CHUNK_SIZE, m, q.shape[-1]), torch.bfloat16),
+            *self._prefill_index_workspace_specs(top_k),
+        )
+
+    def _prefill_index_workspace_specs(
+        self, top_k: int
+    ) -> tuple[tuple[tuple[int, ...], torch.dtype], ...]:
+        combined_topk = round_up(top_k + self.window_size, 128)
+        return (
             ((self.max_num_batched_tokens, combined_topk), torch.int32),
             ((self.max_num_batched_tokens,), torch.int32),
         )
@@ -355,7 +362,7 @@ def _decode_compressed_kv_and_topk(
     def _forward_decode_mega(
         self,
         q: torch.Tensor,
-        positions_int32: torch.Tensor,
+        positions: torch.Tensor,
         flashmla_metadata: DeepseekV4FlashMLAMetadata | None,
         swa_metadata: "DeepseekSparseSWAMetadata",
         out: QuantizedActivation,
@@ -384,7 +391,7 @@ def _forward_decode_mega(
             extra_len,
             False,  # enable_q_norm
             0.0,  # rms_norm_eps
-            positions_int32,
+            positions,
             *_ROPE_ARGS,
             self.rotary_emb.cos_sin_cache,
             self.n_wv_group,
@@ -396,7 +403,7 @@ def _forward_decode_mega(
     def _forward_prefill_mega(
         self,
         q: torch.Tensor,
-        positions_int32: torch.Tensor,
+        positions: torch.Tensor,
         flashmla_metadata: DeepseekV4FlashMLAMetadata | None,
         swa_metadata: "DeepseekSparseSWAMetadata",
         out: QuantizedActivation,
@@ -424,27 +431,39 @@ def _forward_prefill_mega(
         )
         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:
+        index_specs = self._prefill_index_workspace_specs(top_k)
+        for chunk in chunk_plan:
+            chunk_start, chunk_end, chunk_n, chunk_m = chunk
             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:
+            kv_spec = ((chunk_size, chunk_m, q.shape[-1]), torch.bfloat16)
+            if swa_only:
+                kv_ws, idx_ws, lens_ws = workspace_manager.get_simultaneous(
+                    kv_spec, *index_specs
+                )
+            else:
                 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,
+                idx_ws, lens_ws = workspace_manager.get_simultaneous(*index_specs)
+                # The chunk's compressed columns stay in the workspace's tail
+                # span, which the scratch of the layers in between leaves
+                # alone. A kv source has just inserted this step's rows, so it
+                # always gathers; the consumers after it in the same forward
+                # reuse its columns while the span still holds this source and
+                # chunk, and rewrite only their SWA columns.
+                kv_ws, intact = workspace_manager.get_retained(
+                    (self.compressed_cache_prefix, chunk), *kv_spec
                 )
+                if self.is_kv_source or not intact:
+                    dequantize_and_gather_k_cache(
+                        kv_ws[:, :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,
@@ -485,7 +504,7 @@ def _forward_prefill_mega(
                 combined_lens,
                 False,  # enable_q_norm
                 0.0,  # rms_norm_eps
-                positions_int32[qs:qe],
+                positions[qs:qe],
                 *_ROPE_ARGS,
                 self.rotary_emb.cos_sin_cache,
                 self.n_wv_group,
diff --git a/vllm/models/deepseek_v41/nvidia/model.py b/vllm/models/deepseek_v41/nvidia/model.py
index b71dadfe9755..4cb39ac4d92d 100644
--- a/vllm/models/deepseek_v41/nvidia/model.py
+++ b/vllm/models/deepseek_v41/nvidia/model.py
@@ -13,7 +13,6 @@
 from vllm.config import VllmConfig
 from vllm.config.kernel import MEGA_MOE_BACKENDS, NATIVE_MEGA_MOE_BACKENDS
 from vllm.distributed import (
-    get_engram_dp_size,
     get_pcp_group,
     get_pp_group,
     get_tensor_model_parallel_world_size,
@@ -111,6 +110,7 @@
 from ..common.mm_preprocess import IMAGE_SENTINEL_BASE_ID, image_sentinel_mask
 from .ops.mhc import (
     MHC_OVERLAP_MAX_TOKENS,
+    hc_expand,
     init_mhc_all_reduce,
     mhc_pre_delayed_overlap,
     mhc_shifted_post_pre,
@@ -265,8 +265,7 @@ def _use_sequence_parallel(vllm_config: VllmConfig) -> bool:
         vllm_config.kernel_config.moe_backend in MEGA_MOE_BACKENDS
     )
     return (
-        parallel_config.pipeline_parallel_size == 1
-        and parallel_config.enable_expert_parallel
+        parallel_config.enable_expert_parallel
         and parallel_config.tensor_parallel_size > 1
         and (moe_needs_token_sharded_input or parallel_config.data_parallel_size > 1)
     )
@@ -529,7 +528,7 @@ def forward(
                 # First layer: the stream is the embedding broadcast to hc
                 # 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()
+                residual = hc_expand(x, self.hc_mult)
                 post_mix, res_mix, x, attn_pre = mhc_pre(
                     residual,
                     self.hc_attn_fn_broadcast,
@@ -583,6 +582,7 @@ def forward(
                 engram_hashes[:, self.engram.layer_hash_index],
                 engram_mask,
             )
+            del previous_post
             post_mix, res_mix, x, attn_pre = mhc_pre(
                 residual,
                 self.hc_attn_fn,
@@ -707,14 +707,13 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         self.quant_config = quant_config
         self.parallel_config = vllm_config.parallel_config
         if self.parallel_config.prefill_context_parallel_size > 1 and (
-            self.parallel_config.pipeline_parallel_size != 1
-            or self.parallel_config.data_parallel_size != 1
-            or self.parallel_config.decode_context_parallel_size != 1
+            self.parallel_config.decode_context_parallel_size != 1
             or self.parallel_config.use_ubatching
             or (
                 vllm_config.speculative_config is not None
                 and (
-                    not vllm_config.speculative_config.use_dspark()
+                    self.parallel_config.pipeline_parallel_size > 1
+                    or not vllm_config.speculative_config.use_dspark()
                     or vllm_config.speculative_config.enable_adaptive_verification
                 )
             )
@@ -726,9 +725,9 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
             or vllm_config.kernel_config.moe_backend in MEGA_MOE_BACKENDS
         ):
             raise NotImplementedError(
-                "DeepSeek-V4.1 PCP currently requires PP=DP=DCP=1, with no "
+                "DeepSeek-V4.1 PCP currently requires DCP=1, with no "
                 "microbatching, KV transfer, or MegaMoE. "
-                "Speculative decoding requires fixed-length DSpark. "
+                "Speculative decoding requires fixed-length DSpark with PP=1. "
                 "Prefix caching requires SWA bounded replay."
             )
         self.use_native_mega_moe = (
@@ -889,6 +888,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
                 self.layers[cut].write_kv,
                 metadata_prefixes,
                 self.layers[cut].attn.swa_cache_layer.prefix,
+                lambda: any(i < cut for i in self.aux_hidden_state_layers),
             )
             logger.info_once(
                 "Decoder SWA bounded replay: in eager prefill steps, layer %d "
@@ -999,8 +999,8 @@ def forward(
 
         # Engram n-gram hashes for the whole (flattened) batch, computed once
         # on the full token stream — before any sequence-parallel sharding —
-        # and consumed by the engram layers (1 and 14) below. Skipped on
-        # profile runs (KV cache unbound).
+        # and consumed by the engram layers (1 and 14) below. Profile runs
+        # (KV cache unbound) look up dead rows instead.
         engram_hashes: torch.Tensor | None = None
         engram_mask: torch.Tensor | None = None
         if (
@@ -1041,9 +1041,10 @@ def forward(
                     swa_metadata.slot_mapping,
                     swa_metadata.block_table,
                 )
-            elif not self.engram_dp_shared_memory and get_engram_dp_size() > 1:
+            else:
                 # DP-sharded lookups are collective, so a replica skipping the
-                # hash still has to reach them.
+                # hash still has to reach them; the memory profile also has to
+                # see the Engram activations.
                 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.
@@ -1089,6 +1090,8 @@ def forward(
         if not get_pp_group().is_first_rank:
             assert intermediate_tensors is not None
             pre_mix = intermediate_tensors["pre_mix"]
+            if self.use_sequence_parallel:
+                pre_mix = sp_shard(pre_mix)
         aux_hidden_by_layer: dict[int, torch.Tensor] = {}
         hidden_states, residual, post_mix, res_mix, pre_mix = self._run_layers(
             range(self.start_layer, self.decoder_replay_start),
@@ -1133,6 +1136,10 @@ def forward(
         ] + late_aux
 
         if not get_pp_group().is_last_rank:
+            # PP transport expects tensors replicated across TP ranks.
+            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]
             return IntermediateTensors(
                 {"hidden_states": hidden_states, "pre_mix": pre_mix}
             )
@@ -1761,8 +1768,10 @@ def get_mtp_target_hidden_states(self) -> torch.Tensor | None:
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         loader = AutoWeightsLoader(self)
-        loaded_params = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
-        self.process_weights_after_loading()
+        return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
+
+    def process_weights_after_loading(self) -> None:
+        self.model.finalize_mhc_broadcast_weights()
         config = self.model.vllm_config
         if config.engram_config and config.engram_config.use_thp:
             # Loading weights refills the file cache; release it before MADV_COLLAPSE.
@@ -1777,10 +1786,6 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
             ):
                 if layer.engram is not None:
                     layer.engram.embed_tokens.collapse_huge_pages()
-        return loaded_params
-
-    def process_weights_after_loading(self) -> None:
-        self.model.finalize_mhc_broadcast_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
index ed09ae3549ec..a434716dff2f 100644
--- a/vllm/models/deepseek_v41/nvidia/model_state.py
+++ b/vllm/models/deepseek_v41/nvidia/model_state.py
@@ -234,9 +234,9 @@ def __init__(
             self._keeps_rows_np = np.zeros(self.max_num_reqs, dtype=np.bool_)
             self._dropped_before = UvaBufferPool(max_local_reqs + 1, torch.int32)
             # The replay batch's inputs: its rows of the batch, query offsets,
-            # positions, slot mappings (sized with the KV cache groups on first
-            # use) and the position from which each request holds replay-layer
-            # window KV.
+            # positions, MoE padding mask, slot mappings (sized with the KV cache
+            # groups on first use) and the position from which each request holds
+            # replay-layer window KV.
             self._kept_rows = torch.zeros(
                 self.max_num_tokens, dtype=torch.int64, device=device
             )
@@ -246,6 +246,9 @@ def __init__(
             self._kept_positions = torch.zeros(
                 self.max_num_tokens, dtype=torch.int64, device=device
             )
+            self._kept_is_padding = torch.ones(
+                self.max_num_tokens, dtype=torch.bool, device=device
+            )
             self._kept_slot_mappings: torch.Tensor | None = None
             self._kept_kv_start = torch.zeros(
                 max_local_reqs, dtype=torch.int32, device=device
@@ -612,18 +615,27 @@ def _prepare_replay_batch(
             kv_cache_config,
             model_specific_attn_metadata=ReplayAttnMetadata(replay_start),
         )[first_swa]
+        if self.pcp_size > 1:
+            is_padding = kept_batch.is_padding
+        else:
+            # Dummy batches and graph-padding rows are padding, as in the outer
+            # batch. Filled in place: replay graphs read the captured buffer.
+            num_real = 0 if is_dummy else num_tokens
+            is_padding = self._kept_is_padding[:num_padded]
+            is_padding[:num_real].fill_(False)
+            is_padding[num_real:].fill_(True)
         layers.replay_batch = ReplayBatch(
             self._kept_rows[:num_tokens],
             create_forward_context(
                 attn_metadata,
                 self.vllm_config,
                 dp_metadata=dp_metadata,
+                is_padding=is_padding,
                 cudagraph_runtime_mode=cg_mode,
                 batch_descriptor=BatchDescriptor(num_tokens=num_padded),
                 slot_mapping=build_slot_mappings_by_layer(
                     kept_slot_mappings, kv_cache_config
                 ),
-                is_padding=kept_batch.is_padding if self.pcp_size > 1 else None,
             ),
             run_graph,
             num_padded if self.pcp_size > 1 else None,
@@ -635,6 +647,7 @@ def capture_inner_cudagraphs(
         block_tables: BlockTables,
         attn_groups: list[list[AttentionGroup]],
         kv_cache_config: KVCacheConfig,
+        stream: torch.cuda.Stream | None = None,
     ) -> None:
         def prepare(desc: BatchExecutionDescriptor) -> None:
             # Capture reuses the runtime path on a dummy batch, so the graph reads
@@ -654,7 +667,7 @@ def prepare(desc: BatchExecutionDescriptor) -> None:
         if self.replay_cudagraphs is not None:
             # profile_only captures use pre-KV-cache groups: drop the cache.
             self._replay_attn_groups = None
-            self.replay_cudagraphs.capture_replay_graphs(prepare)
+            self.replay_cudagraphs.capture_replay_graphs(prepare, stream)
 
     def _kept_input_batch(
         self,
diff --git a/vllm/models/deepseek_v41/nvidia/ops/cute_dsl/all_reduce_mhc.py b/vllm/models/deepseek_v41/nvidia/ops/cute_dsl/all_reduce_mhc.py
index 191a928fe2e9..3bbd13cf2f49 100644
--- a/vllm/models/deepseek_v41/nvidia/ops/cute_dsl/all_reduce_mhc.py
+++ b/vllm/models/deepseek_v41/nvidia/ops/cute_dsl/all_reduce_mhc.py
@@ -34,10 +34,11 @@
 from cutlass.cute.runtime import make_fake_compact_tensor
 
 from vllm.distributed import get_tp_group
-from vllm.models.deepseek_v41.common.ops.query_quant import mxfp8_scale_bytes
-
-from .primitives import (
+from vllm.models.common.ops.moe_tail_fusion.primitives import (
+    ACTIVE_STAGE,
+    LAMPORT_GENERATIONS,
     NEGATIVE_ZERO_BF16_BITS,
+    NEXT_STAGE,
     QUAD_BF16,
     VEC_BF16,
     WARP_SIZE,
@@ -46,6 +47,7 @@
     current_cu_stream,
     f32_to_bf16_bits,
     fragment_has_negative_zero,
+    group_leader_block_sum,
     load_global_bf16_as_f32,
     load_global_u32x2,
     load_global_u32x4,
@@ -67,10 +69,8 @@
     to_cute,
     to_cute_dynamic,
 )
+from vllm.models.deepseek_v41.common.ops.query_quant import mxfp8_scale_bytes
 
-LAMPORT_GENERATIONS = 3
-NEXT_STAGE = 0
-ACTIVE_STAGE = 1
 # One UE8M0 scale per 32 values: four adjacent threads' BF16 fragments.
 MXFP8_BLOCK = 32
 FRAGMENTS_PER_MXFP8_BLOCK = MXFP8_BLOCK // VEC_BF16
@@ -79,41 +79,6 @@
 SCALE_TILE_COLS = 4
 
 
-@cute.jit
-def _group_leader_block_sum(
-    value: Float32,
-    warp_sums: cute.Tensor,
-    warps: cutlass.Constexpr[int],
-    leader_stride: cutlass.Constexpr[int],
-) -> Float32:
-    lane = cute.arch.lane_idx()
-    warp = cute.arch.warp_idx()
-    for offset in cutlass.range_constexpr(1, WARP_SIZE):
-        if cutlass.const_expr(offset >= leader_stride and (offset & (offset - 1)) == 0):
-            value = value + cute.arch.shuffle_sync_bfly(
-                value,
-                offset=offset,
-                mask=-1,
-                mask_and_clamp=31,
-            )
-    if lane == 0:
-        cute.arch.store((warp_sums + warp).llvm_ptr, value)
-    cute.arch.barrier()
-
-    result = Float32(0.0)
-    if warp == 0:
-        if lane < Int32(warps):
-            result = cute.arch.load(
-                (warp_sums + lane).llvm_ptr,
-                Float32,
-            )
-        result = cute.arch.warp_reduction_sum(result)
-        if lane == 0:
-            cute.arch.store(warp_sums.llvm_ptr, result)
-    cute.arch.barrier()
-    return cute.arch.load(warp_sums.llvm_ptr, Float32)
-
-
 class _SharedOnlyPublishDeviceKernel:
     def __init__(
         self,
@@ -227,7 +192,7 @@ def kernel(
         if cutlass.const_expr(self.release_before_store):
             if block == 0 and tidx == 0:
                 store_global_u32(
-                    stage_state.iterator + ACTIVE_STAGE,
+                    (stage_state.iterator + ACTIVE_STAGE).toint(),
                     stage,
                 )
             cute.arch.barrier()
@@ -270,7 +235,7 @@ def kernel(
         if cutlass.const_expr(not self.release_before_store):
             if block == 0 and tidx == 0:
                 store_global_u32(
-                    stage_state.iterator + ACTIVE_STAGE,
+                    (stage_state.iterator + ACTIVE_STAGE).toint(),
                     stage,
                 )
             cute.arch.barrier()
@@ -523,7 +488,7 @@ def kernel(
 
         if block == 0 and tidx == 0:
             store_global_u32(
-                stage_state.iterator + ACTIVE_STAGE,
+                (stage_state.iterator + ACTIVE_STAGE).toint(),
                 stage,
             )
         cute.arch.barrier()
@@ -826,7 +791,7 @@ def kernel(
 
         if token == 0 and cluster_rank == 0 and tidx == 0:
             store_global_u32(
-                stage_state.iterator + NEXT_STAGE,
+                (stage_state.iterator + NEXT_STAGE).toint(),
                 (active_stage + Uint32(1)) % Uint32(LAMPORT_GENERATIONS),
             )
         if cutlass.const_expr(self.enable_pdl):
@@ -861,7 +826,7 @@ def kernel(
             smem = cutlass.utils.SmemAllocator()
             warp_sums = smem.allocate_array(Float32, self.warps)
             cluster_sums = smem.allocate_array(Float32, self.cluster_size)
-            cta_sum = _group_leader_block_sum(
+            cta_sum = group_leader_block_sum(
                 thread_sum,
                 warp_sums,
                 self.warps,
diff --git a/vllm/models/deepseek_v41/nvidia/ops/cute_dsl/primitives.py b/vllm/models/deepseek_v41/nvidia/ops/cute_dsl/primitives.py
deleted file mode 100644
index 16d5f4721a33..000000000000
--- a/vllm/models/deepseek_v41/nvidia/ops/cute_dsl/primitives.py
+++ /dev/null
@@ -1,396 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-# Adapted from FlashInfer's flashinfer/comm/mnnvl_cutedsl/cute_dsl_primitives.py
-# and runtime.py (flashinfer-ai/flashinfer@139af6f6).
-
-"""PTX building blocks and launch helpers for the DSV4.1 LL all-reduce."""
-
-import cuda.bindings.driver as cuda
-import cutlass
-import cutlass.cute as cute
-import torch
-from cutlass import Float32, Int32, Int64, Uint16, Uint32
-from cutlass._mlir import ir
-from cutlass._mlir.dialects import llvm, vector
-from cutlass.cute.runtime import from_dlpack, make_fake_compact_tensor
-from cutlass.cutlass_dsl import T, dsl_user_op
-
-WARP_SIZE = 32
-VEC_BF16 = 8
-QUAD_BF16 = 4
-NEGATIVE_ZERO_BF16_BITS = 0x8000
-NEGATIVE_ZERO_BF16_PAIR = 0x80008000
-
-
-def _asm(result, operands, text, constraints, *, side_effects, loc, ip):
-    return llvm.inline_asm(
-        result,
-        operands,
-        text,
-        constraints,
-        has_side_effects=side_effects,
-        is_align_stack=False,
-        asm_dialect=llvm.AsmDialect.AD_ATT,
-        loc=loc,
-        ip=ip,
-    )
-
-
-@dsl_user_op
-def load_global_u32x4(
-    pointer: cute.Pointer,
-    *,
-    volatile: cutlass.Constexpr[bool] = False,
-    loc=None,
-    ip=None,
-):
-    address = pointer.toint(loc=loc, ip=ip)
-    opcode = "ld.volatile.global.v4.u32" if volatile else "ld.global.v4.u32"
-    loaded = _asm(
-        llvm.StructType.get_literal([T.i32()] * 4),
-        [address.ir_value(loc=loc, ip=ip)],
-        f"{opcode} {{$0, $1, $2, $3}}, [$4];",
-        "=r,=r,=r,=r,l",
-        side_effects=volatile,
-        loc=loc,
-        ip=ip,
-    )
-    packed = vector.from_elements(
-        ir.VectorType.get([4], T.i32(), loc=loc),
-        [
-            llvm.extractvalue(T.i32(), loaded, [index], loc=loc, ip=ip)
-            for index in range(4)
-        ],
-        loc=loc,
-        ip=ip,
-    )
-    return cute.TensorSSA(packed, 4, Uint32)
-
-
-@dsl_user_op
-def load_global_u32x2(pointer: cute.Pointer, *, loc=None, ip=None):
-    address = pointer.toint(loc=loc, ip=ip)
-    loaded = _asm(
-        llvm.StructType.get_literal([T.i32()] * 2),
-        [address.ir_value(loc=loc, ip=ip)],
-        "ld.global.v2.u32 {$0, $1}, [$2];",
-        "=r,=r,l",
-        side_effects=False,
-        loc=loc,
-        ip=ip,
-    )
-    packed = vector.from_elements(
-        ir.VectorType.get([2], T.i32(), loc=loc),
-        [
-            llvm.extractvalue(T.i32(), loaded, [index], loc=loc, ip=ip)
-            for index in range(2)
-        ],
-        loc=loc,
-        ip=ip,
-    )
-    return cute.TensorSSA(packed, 2, Uint32)
-
-
-@dsl_user_op
-def load_global_bf16_as_f32(address: Int64, *, loc=None, ip=None) -> Float32:
-    return Float32(
-        _asm(
-            T.f32(),
-            [address.ir_value(loc=loc, ip=ip)],
-            "{\n\t.reg .b16 bits;\n\tld.global.b16 bits, [$1];\n\t"
-            "cvt.f32.bf16 $0, bits;\n\t}",
-            "=f,l",
-            side_effects=False,
-            loc=loc,
-            ip=ip,
-        )
-    )
-
-
-@dsl_user_op
-def store_global_u32x4(address: Int64, packed, *, loc=None, ip=None) -> None:
-    words = [packed[index].ir_value(loc=loc, ip=ip) for index in range(4)]
-    _asm(
-        None,
-        [address.ir_value(loc=loc, ip=ip), *words],
-        "st.global.v4.u32 [$0], {$1, $2, $3, $4};",
-        "l,r,r,r,r",
-        side_effects=True,
-        loc=loc,
-        ip=ip,
-    )
-
-
-@dsl_user_op
-def store_lamport_sentinel_u32x4(address: Int64, *, loc=None, ip=None) -> None:
-    sentinel = Uint32(NEGATIVE_ZERO_BF16_PAIR).ir_value(loc=loc, ip=ip)
-    _asm(
-        None,
-        [address.ir_value(loc=loc, ip=ip), sentinel, sentinel, sentinel, sentinel],
-        "st.global.v4.u32 [$0], {$1, $2, $3, $4};",
-        "l,r,r,r,r",
-        side_effects=True,
-        loc=loc,
-        ip=ip,
-    )
-
-
-@dsl_user_op
-def load_volatile_u32(pointer: cute.Pointer, *, loc=None, ip=None) -> Uint32:
-    address = pointer.toint(loc=loc, ip=ip)
-    return Uint32(
-        _asm(
-            T.i32(),
-            [address.ir_value(loc=loc, ip=ip)],
-            "ld.volatile.global.u32 $0, [$1];",
-            "=r,l",
-            side_effects=True,
-            loc=loc,
-            ip=ip,
-        )
-    )
-
-
-@dsl_user_op
-def store_global_u32(
-    pointer: cute.Pointer, value: Uint32, *, loc=None, ip=None
-) -> None:
-    address = pointer.toint(loc=loc, ip=ip)
-    _asm(
-        None,
-        [address.ir_value(loc=loc, ip=ip), value.ir_value(loc=loc, ip=ip)],
-        "st.global.u32 [$0], $1;",
-        "l,r",
-        side_effects=True,
-        loc=loc,
-        ip=ip,
-    )
-
-
-@dsl_user_op
-def packed_u32x4_to_bf16x8(packed, *, loc=None, ip=None):
-    values = llvm.bitcast(
-        ir.VectorType.get([VEC_BF16], cutlass.BFloat16.mlir_type, loc=loc),
-        packed.ir_value(loc=loc, ip=ip),
-        loc=loc,
-        ip=ip,
-    )
-    return cute.TensorSSA(values, VEC_BF16, cutlass.BFloat16)
-
-
-@dsl_user_op
-def packed_u32x2_to_bf16x4(packed, *, loc=None, ip=None):
-    values = llvm.bitcast(
-        ir.VectorType.get([QUAD_BF16], cutlass.BFloat16.mlir_type, loc=loc),
-        packed.ir_value(loc=loc, ip=ip),
-        loc=loc,
-        ip=ip,
-    )
-    return cute.TensorSSA(values, QUAD_BF16, cutlass.BFloat16)
-
-
-@dsl_user_op
-def bf16x4_to_packed_u32x2(values, *, loc=None, ip=None):
-    packed = llvm.bitcast(
-        ir.VectorType.get([2], T.i32(), loc=loc),
-        values.ir_value(loc=loc, ip=ip),
-        loc=loc,
-        ip=ip,
-    )
-    return cute.TensorSSA(packed, 2, Uint32)
-
-
-@dsl_user_op
-def bf16x8_to_packed_u32x4(values, *, loc=None, ip=None):
-    packed = llvm.bitcast(
-        ir.VectorType.get([4], T.i32(), loc=loc),
-        values.ir_value(loc=loc, ip=ip),
-        loc=loc,
-        ip=ip,
-    )
-    return cute.TensorSSA(packed, 4, Uint32)
-
-
-@cute.jit
-def sanitize_negative_zero_u32(word: Uint32) -> Uint32:
-    low = Uint16(word & Uint32(0xFFFF))
-    high = Uint16(word >> Uint32(16))
-    if low == Uint16(NEGATIVE_ZERO_BF16_BITS):
-        word = word & Uint32(0xFFFF0000)
-    if high == Uint16(NEGATIVE_ZERO_BF16_BITS):
-        word = word & Uint32(0x0000FFFF)
-    return word
-
-
-@cute.jit
-def sanitize_negative_zero_u32x4(packed):
-    sanitized = cute.make_rmem_tensor(cute.make_layout((4,)), Uint32)
-    for index in cutlass.range_constexpr(4):
-        sanitized[index] = sanitize_negative_zero_u32(packed[index])
-    return sanitized.load()
-
-
-@cute.jit
-def sanitize_negative_zero_u32x2(packed):
-    sanitized = cute.make_rmem_tensor(cute.make_layout((2,)), Uint32)
-    for index in cutlass.range_constexpr(2):
-        sanitized[index] = sanitize_negative_zero_u32(packed[index])
-    return sanitized.load()
-
-
-@cute.jit
-def fragment_has_negative_zero(packed):
-    dirty = False
-    for index in cutlass.range_constexpr(4):
-        word = packed[index]
-        dirty = (
-            dirty
-            | (Uint16(word & Uint32(0xFFFF)) == Uint16(NEGATIVE_ZERO_BF16_BITS))
-            | (Uint16(word >> Uint32(16)) == Uint16(NEGATIVE_ZERO_BF16_BITS))
-        )
-    return dirty
-
-
-@dsl_user_op
-def map_shared_to_peer(
-    smem_pointer: cute.Pointer, peer_rank: Int32, *, loc=None, ip=None
-) -> Int32:
-    address = smem_pointer.toint(loc=loc, ip=ip).ir_value(loc=loc, ip=ip)
-    return Int32(
-        _asm(
-            T.i32(),
-            [address, peer_rank.ir_value(loc=loc, ip=ip)],
-            "mapa.shared::cluster.u32 $0, $1, $2;",
-            "=r,r,r",
-            side_effects=False,
-            loc=loc,
-            ip=ip,
-        )
-    )
-
-
-@dsl_user_op
-def store_shared_cluster_f32(
-    remote_address: Int32, value: Float32, *, loc=None, ip=None
-) -> None:
-    _asm(
-        None,
-        [remote_address.ir_value(loc=loc, ip=ip), value.ir_value(loc=loc, ip=ip)],
-        "st.shared::cluster.f32 [$0], $1;",
-        "r,f",
-        side_effects=True,
-        loc=loc,
-        ip=ip,
-    )
-
-
-@dsl_user_op
-def f32_to_bf16_bits(value: Float32, *, loc=None, ip=None) -> Uint32:
-    return Uint32(
-        _asm(
-            T.i32(),
-            [value.ir_value(loc=loc, ip=ip)],
-            "{\n\t.reg .b16 bits;\n\tcvt.rn.bf16.f32 bits, $1;\n\t"
-            "cvt.u32.u16 $0, bits;\n\t}",
-            "=r,f",
-            side_effects=False,
-            loc=loc,
-            ip=ip,
-        )
-    )
-
-
-@dsl_user_op
-def shuffle_sync_idx_u32(
-    value: Uint32, source_lane: Int32, *, loc=None, ip=None
-) -> Uint32:
-    return Uint32(
-        _asm(
-            T.i32(),
-            [value.ir_value(loc=loc, ip=ip), source_lane.ir_value(loc=loc, ip=ip)],
-            "shfl.sync.idx.b32 $0, $1, $2, 0x1f, 0xffffffff;",
-            "=r,r,r",
-            # Preserve full-warp execution across later divergent consumers.
-            side_effects=True,
-            loc=loc,
-            ip=ip,
-        )
-    )
-
-
-@dsl_user_op
-def stmc_bf16x2(address: Int64, packed: Uint32, *, loc=None, ip=None) -> None:
-    _asm(
-        None,
-        [address.ir_value(loc=loc, ip=ip), packed.ir_value(loc=loc, ip=ip)],
-        "multimem.st.relaxed.sys.global.bf16x2 [$0], $1;",
-        "l,r",
-        side_effects=True,
-        loc=loc,
-        ip=ip,
-    )
-
-
-@dsl_user_op
-def stmc_bf16x4(address: Int64, values, *, loc=None, ip=None) -> None:
-    words = [values[index].ir_value(loc=loc, ip=ip) for index in range(2)]
-    _asm(
-        None,
-        [address.ir_value(loc=loc, ip=ip), *words],
-        "multimem.st.relaxed.sys.global.v2.bf16x2 [$0], {$1, $2};",
-        "l,r,r",
-        side_effects=True,
-        loc=loc,
-        ip=ip,
-    )
-
-
-@dsl_user_op
-def stmc_bf16x8(address: Int64, values, *, loc=None, ip=None) -> None:
-    words = [values[index].ir_value(loc=loc, ip=ip) for index in range(4)]
-    _asm(
-        None,
-        [address.ir_value(loc=loc, ip=ip), *words],
-        "multimem.st.relaxed.sys.global.v4.bf16x2 [$0], {$1, $2, $3, $4};",
-        "l,r,r,r,r",
-        side_effects=True,
-        loc=loc,
-        ip=ip,
-    )
-
-
-class _GraphSafeDLPack:
-    __slots__ = ("tensor",)
-
-    def __init__(self, tensor: torch.Tensor) -> None:
-        self.tensor = tensor
-
-    def __dlpack__(self, stream=None):
-        # stream=-1 skips producer sync; CuTe launches on the current stream,
-        # which is the capturing one under CUDA graphs.
-        return self.tensor.__dlpack__(stream=-1)
-
-    def __dlpack_device__(self):
-        return self.tensor.__dlpack_device__()
-
-
-def to_cute(tensor: torch.Tensor, alignment: int) -> cute.Tensor:
-    return from_dlpack(_GraphSafeDLPack(tensor.detach()), assumed_align=alignment)
-
-
-def to_cute_dynamic(
-    tensor: torch.Tensor, alignment: int, *, divisibility: int
-) -> cute.Tensor:
-    return to_cute(tensor, alignment).mark_compact_shape_dynamic(
-        mode=0, divisibility=divisibility
-    )
-
-
-def make_fake_dynamic_compact_tensor(dtype, *, alignment: int, divisibility: int):
-    return make_fake_compact_tensor(
-        dtype, (cute.sym_int32(divisibility=divisibility),), assumed_align=alignment
-    )
-
-
-def current_cu_stream() -> cuda.CUstream:
-    return cuda.CUstream(torch.cuda.current_stream().cuda_stream)
diff --git a/vllm/models/deepseek_v41/nvidia/ops/engram_host_table.py b/vllm/models/deepseek_v41/nvidia/ops/engram_host_table.py
new file mode 100644
index 000000000000..ca480de6056b
--- /dev/null
+++ b/vllm/models/deepseek_v41/nvidia/ops/engram_host_table.py
@@ -0,0 +1,241 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Engram host tables in host-NUMA cuMem, shared across DP ranks by FABRIC handle.
+
+GPUs map cuMem host memory with 2M pages instead of the 64K pages of a registered
+/dev/shm mapping, so random row lookups miss the TLB far less.
+"""
+
+import ctypes
+import weakref
+from functools import cache
+from types import SimpleNamespace
+from typing import NamedTuple
+
+import torch
+from cuda.bindings import driver as cu
+
+from vllm.distributed.parallel_state import GroupCoordinator
+from vllm.logger import init_logger
+from vllm.models.deepseek_v41.common.engram import describe_error, gather_errors
+from vllm.utils.cpu_resource_utils import get_memory_node_info
+from vllm.utils.math_utils import round_up
+
+logger = init_logger(__name__)
+
+FABRIC_HANDLE = cu.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_FABRIC
+# Tables are created in chunks, so the creator can stop before draining its node.
+CHUNK_BYTES = 2 << 30
+# Fraction of the NUMA node left to other allocations.
+NUMA_RESERVE_FRACTION = 0.05
+
+
+class ExportedHostTable(NamedTuple):
+    """What a peer needs to map a table."""
+
+    numa_node: int
+    padded_bytes: int
+    fabric_handles: list[bytes]
+
+
+def _chunk_ranges(padded_bytes: int) -> list[tuple[int, int]]:
+    return [
+        (offset, min(CHUNK_BYTES, padded_bytes - offset))
+        for offset in range(0, padded_bytes, CHUNK_BYTES)
+    ]
+
+
+def _call_driver(fn, *args):
+    """Call a CUDA driver API and return its result, raising on failure."""
+    error, *result = fn(*args)
+    if error != cu.CUresult.CUDA_SUCCESS:
+        raise RuntimeError(f"{fn.__name__} failed: {error.name}")
+    return result[0] if result else None
+
+
+def _gpu_numa_node() -> int:
+    return _call_driver(
+        cu.cuDeviceGetAttribute,
+        cu.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_HOST_NUMA_ID,
+        torch.accelerator.current_device_index(),
+    )
+
+
+def _host_allocation_prop(numa_node: int) -> cu.CUmemAllocationProp:
+    prop = cu.CUmemAllocationProp()
+    prop.type = cu.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_PINNED
+    prop.location.type = cu.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST_NUMA
+    prop.location.id = numa_node
+    prop.requestedHandleTypes = FABRIC_HANDLE
+    return prop
+
+
+def _allocation_granularity(prop: cu.CUmemAllocationProp) -> int:
+    return _call_driver(cu.cuMemGetAllocationGranularity, prop, 0)
+
+
+def _export_fabric_handle(handle) -> bytes:
+    exported = _call_driver(cu.cuMemExportToShareableHandle, handle, FABRIC_HANDLE, 0)
+    return bytes(exported.data)
+
+
+def _release(handles: list) -> None:
+    for handle in handles:
+        cu.cuMemRelease(handle)
+
+
+def _read_write_access(location_type, location_id: int) -> cu.CUmemAccessDesc:
+    access = cu.CUmemAccessDesc()
+    access.location.type = location_type
+    access.location.id = location_id
+    access.flags = cu.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READWRITE
+    return access
+
+
+def _unmap(address: int, size: int, ranges: list[tuple[int, int]]) -> None:
+    for offset, length in ranges:
+        cu.cuMemUnmap(address + offset, length)
+    cu.cuMemAddressFree(address, size)
+
+
+@cache
+def fabric_host_memory_supported() -> bool:
+    """Whether host-NUMA cuMem allocations can export FABRIC handles here."""
+    try:
+        prop = _host_allocation_prop(_gpu_numa_node())
+        handle = _call_driver(cu.cuMemCreate, _allocation_granularity(prop), prop, 0)
+        try:
+            _export_fabric_handle(handle)
+        finally:
+            cu.cuMemRelease(handle)
+    except Exception as exc:
+        logger.info_once(
+            "Host-NUMA cuMem FABRIC handles are unavailable (%s); Engram DP "
+            "replicas share tables through /dev/shm.",
+            exc,
+        )
+        return False
+    return True
+
+
+def _create_table(num_bytes: int) -> tuple[list, ExportedHostTable]:
+    """Allocate and export a table on this GPU's host NUMA node.
+
+    Stops once the node would drop below its reserve. The kernel compacts to
+    find 2M pages and fails the chunk, rather than OOM-killing, if it cannot.
+    """
+    numa_node = _gpu_numa_node()
+    prop = _host_allocation_prop(numa_node)
+    padded_bytes = round_up(num_bytes, _allocation_granularity(prop))
+    reserve = get_memory_node_info(numa_node).total_memory * NUMA_RESERVE_FRACTION
+    handles: list = []
+    try:
+        for offset, length in _chunk_ranges(padded_bytes):
+            handles.append(_call_driver(cu.cuMemCreate, length, prop, 0))
+            if get_memory_node_info(numa_node).available_memory < reserve:
+                raise MemoryError(
+                    f"NUMA node {numa_node} fell below its "
+                    f"{reserve / 2**30:.1f} GiB reserve after "
+                    f"{(offset + length) / 2**30:.1f} of "
+                    f"{padded_bytes / 2**30:.1f} GiB"
+                )
+        exported = [_export_fabric_handle(handle) for handle in handles]
+    except BaseException:
+        _release(handles)
+        raise
+    return handles, ExportedHostTable(numa_node, padded_bytes, exported)
+
+
+def _map_table(
+    table: ExportedHostTable, num_bytes: int, handles: list
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Map a table, importing its handles unless given, and release the handles.
+
+    Returns CPU and GPU uint8 views; either keeps the mapping alive.
+    """
+    device = torch.accelerator.current_device_index()
+    address = None
+    mapped: list[tuple[int, int]] = []
+    try:
+        if not handles:
+            for exported in table.fabric_handles:
+                handles.append(
+                    _call_driver(
+                        cu.cuMemImportFromShareableHandle, exported, FABRIC_HANDLE
+                    )
+                )
+        address = int(_call_driver(cu.cuMemAddressReserve, table.padded_bytes, 0, 0, 0))
+        for (offset, length), handle in zip(_chunk_ranges(table.padded_bytes), handles):
+            _call_driver(cu.cuMemMap, address + offset, length, 0, handle, 0)
+            mapped.append((offset, length))
+        location = cu.CUmemLocationType
+        access = [
+            _read_write_access(location.CU_MEM_LOCATION_TYPE_DEVICE, device),
+            _read_write_access(
+                location.CU_MEM_LOCATION_TYPE_HOST_NUMA, table.numa_node
+            ),
+        ]
+        _call_driver(
+            cu.cuMemSetAccess, address, table.padded_bytes, access, len(access)
+        )
+    except BaseException:
+        if address is not None:
+            _unmap(address, table.padded_bytes, mapped)
+        raise
+    finally:
+        # The mappings keep the allocations alive.
+        _release(handles)
+
+    host_buffer = (ctypes.c_uint8 * num_bytes).from_address(address)
+    finalizer = weakref.finalize(
+        host_buffer, _unmap, address, table.padded_bytes, mapped
+    )
+    finalizer.atexit = False  # type: ignore[misc]
+    # Torch retains this object, so the GPU view keeps the buffer alive too.
+    # Name the device: an index-less "cuda" makes torch copy the table.
+    gpu_view = torch.as_tensor(
+        SimpleNamespace(
+            __cuda_array_interface__={
+                "shape": (num_bytes,),
+                "typestr": "|u1",
+                "data": (address, False),
+                "version": 3,
+            },
+            owner=host_buffer,
+        ),
+        device=torch.device("cuda", device),
+    )
+    return torch.frombuffer(host_buffer, dtype=torch.uint8), gpu_view
+
+
+def share_host_table(
+    num_bytes: int, group: GroupCoordinator
+) -> tuple[torch.Tensor, torch.Tensor] | None:
+    """Map one table on every rank of a node-local group, or None on all of them.
+
+    Rank 0 creates the table and broadcasts its handles; the others import them.
+    Returns CPU and GPU uint8 views of the table.
+    """
+    handles: list = []
+    table = error = None
+    if group.rank_in_group == 0 and fabric_host_memory_supported():
+        try:
+            handles, table = _create_table(num_bytes)
+        except Exception as exc:
+            error = describe_error(exc)
+    table, error = group.broadcast_object((table, error))
+    if table is not None:
+        views = None
+        try:
+            views = _map_table(table, num_bytes, handles)
+        except Exception as exc:
+            error = describe_error(exc)
+        if not (error := gather_errors(group, error)):
+            return views
+    if error and group.rank_in_group == 0:
+        logger.warning(
+            "Engram host table cannot use host-NUMA cuMem (%s); falling back "
+            "to /dev/shm, whose 64K GPU pages make lookups slower.",
+            error,
+        )
+    return None
diff --git a/vllm/models/deepseek_v41/nvidia/ops/mhc.py b/vllm/models/deepseek_v41/nvidia/ops/mhc.py
index edfdbd1ee40d..a2207d1c4d70 100644
--- a/vllm/models/deepseek_v41/nvidia/ops/mhc.py
+++ b/vllm/models/deepseek_v41/nvidia/ops/mhc.py
@@ -19,6 +19,7 @@
     UnfinalizedMoEOutput,
 )
 from vllm.platforms import current_platform
+from vllm.triton_utils import tl, triton
 from vllm.utils.deep_gemm import is_deep_gemm_supported
 
 from .mega_mhc import (
@@ -320,3 +321,35 @@ def mhc_shifted_post_pre(
         norm_eps=norm_eps,
         capture_aux=capture_aux,
     )
+
+
+@triton.jit
+def _hc_expand_kernel(
+    x,
+    out,
+    numel,
+    H: tl.constexpr,
+    N: tl.constexpr,
+    STRIDE_M: tl.constexpr,
+    STRIDE_H: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    index = tl.program_id(0).to(tl.int64) * BLOCK + tl.arange(0, BLOCK)
+    mask = index < numel
+    row = index // H
+    col = index % H
+    value = tl.load(x + row * STRIDE_M + col * STRIDE_H, mask=mask, other=0)
+    for stream in tl.static_range(N):
+        tl.store(out + (row * N + stream) * H + col, value, mask=mask)
+
+
+def hc_expand(x: torch.Tensor, hc_mult: int) -> torch.Tensor:
+    """``x.unsqueeze(1).expand(-1, hc_mult, -1).contiguous()``, reading
+    ``x`` once instead of once per copy."""
+    m, h = x.shape
+    out = x.new_empty(m, hc_mult, h)
+    if out.numel():
+        _hc_expand_kernel[(triton.cdiv(m * h, 1024),)](
+            x, out, m * h, h, hc_mult, *x.stride(), BLOCK=1024, num_warps=4
+        )
+    return out
diff --git a/vllm/models/deepseek_v41/nvidia/ops/o_proj.py b/vllm/models/deepseek_v41/nvidia/ops/o_proj.py
index cd1a12fabf0d..9d7df287d14b 100644
--- a/vllm/models/deepseek_v41/nvidia/ops/o_proj.py
+++ b/vllm/models/deepseek_v41/nvidia/ops/o_proj.py
@@ -10,6 +10,7 @@
 from vllm.model_executor.kernels.linear.mxfp8.flashinfer import (
     FlashInferCutedslMxfp8LinearKernel,
 )
+from vllm.model_executor.kernels.linear.mxfp8.pytorch import TorchMxFp8LinearKernel
 from vllm.model_executor.layers.fusion.quant_activation import (
     QuantizedActivation,
     get_input_quant_key,
@@ -38,7 +39,8 @@ def _can_fuse_wo_a(layer: nn.Module) -> bool:
         and layer._einsum_recipe == (1, 1, 32)
         and get_input_quant_key(layer.wo_b) == kMxfp8Dynamic
         and isinstance(
-            getattr(method, "kernel", None), FlashInferCutedslMxfp8LinearKernel
+            getattr(method, "kernel", None),
+            (FlashInferCutedslMxfp8LinearKernel, TorchMxFp8LinearKernel),
         )
         # A group's heads form one portable (<= 8 CTA) cluster.
         and layer.n_local_heads // layer.n_local_groups <= 8
diff --git a/vllm/models/deepseek_v41/nvidia/specialized_decode/AGENTS.md b/vllm/models/deepseek_v41/nvidia/specialized_decode/AGENTS.md
index 0986fd485687..724b3aee63c1 100644
--- a/vllm/models/deepseek_v41/nvidia/specialized_decode/AGENTS.md
+++ b/vllm/models/deepseek_v41/nvidia/specialized_decode/AGENTS.md
@@ -36,22 +36,28 @@ collectives and would hang. Which forward steps use the executor depends
 only on the batch, which every TP rank receives identically, so that needs no
 collective.
 
-Four features run only on flashinfer-python 0.7.2 (`FLASHINFER_VERSION`, the
-version `requirements/cuda.txt` pins). Three use FlashInfer internals:
+Two features use FlashInfer internals and run on flashinfer-python 0.7.2 or
+newer (`MIN_FLASHINFER_VERSION`):
 
 - the publishing `wo_b` GEMM: the CuTe DSL low-latency block-scaled GEMM
   kernel, and `mm_mxfp8`'s private runner and tuning config
   (`flashinfer.gemm.gemm_base`) to look up its tuned tactic;
 - the fused kernel's MXFP8 layer input: the CuTe DSL MXFP8 quantize helpers
   (`flashinfer.quantization.quantization_cute_dsl_utils`,
-  `flashinfer.cute_dsl.fp4_common`);
-- the pinned MXFP8 GEMM tactics: the low-latency GEMM's `autotune_tactics`.
+  `flashinfer.cute_dsl.fp4_common`).
 
-The fourth, the sparse-MLA decode's persistent KV counters, passes the
-`multi_ctas_kv_counter_buffer` argument that 0.7.1 added.
+On an older version each logs a warning once and the executor goes without
+it. `test_numerics.py` compares both with FlashInfer's own GEMM and quantize
+bitwise; run it when `requirements/cuda.txt` moves to a new FlashInfer.
 
-On any other version each logs a warning once and the executor goes without
-it; the MXFP8 GEMMs then tune over every tactic.
+Two more depend on FlashInfer's interface, not its version:
+
+- the pinned MXFP8 GEMM tactics (GB300) wrap the low-latency GEMM's
+  `autotune_tactics`; without it, or when it does not list a pinned tactic,
+  the GEMMs tune over every tactic;
+- the sparse-MLA decode's persistent KV counters need the
+  `multi_ctas_kv_counter_buffer` argument (0.7.1 added it); without it each
+  decode allocates its own.
 
 ## Structure
 
diff --git a/vllm/models/deepseek_v41/nvidia/specialized_decode/__init__.py b/vllm/models/deepseek_v41/nvidia/specialized_decode/__init__.py
index 1d93d5e64f7d..a08a284cdac5 100644
--- a/vllm/models/deepseek_v41/nvidia/specialized_decode/__init__.py
+++ b/vllm/models/deepseek_v41/nvidia/specialized_decode/__init__.py
@@ -16,6 +16,7 @@
 
 import torch
 import torch.distributed as dist
+from packaging.version import Version
 
 import vllm.envs as envs
 from vllm.config import CompilationMode, CUDAGraphMode, VllmConfig
@@ -37,10 +38,11 @@
 # Rows of the low-batch GEMMs: one DSpark verification step of up to 8 rows
 # (6 at one request).
 LOW_BATCH_MAX_ROWS = 8
-# The FlashInfer (requirements/cuda.txt) whose internals the publishing wo_b
-# GEMM, the MXFP8 layer input and the pinned MXFP8 tactics were checked
-# against. Its sparse-MLA decode takes caller-owned KV counters (since 0.7.1).
-FLASHINFER_VERSION = "0.7.2"
+# The oldest FlashInfer whose internals the publishing wo_b GEMM and the MXFP8
+# layer input were checked against. Newer releases run them too: test_numerics.py
+# compares both with FlashInfer's own GEMM and quantize bitwise on the release
+# requirements/cuda.txt pins.
+MIN_FLASHINFER_VERSION = "0.7.2"
 
 
 @dataclass(frozen=True)
@@ -54,19 +56,19 @@ class DecodeConfig:
 
 @cache
 def uses_checked_flashinfer(feature: str) -> bool:
-    """Whether the installed FlashInfer is the version whose internals
-    ``feature`` uses; logs once per feature when it is not."""
+    """Whether the installed FlashInfer is at least the version whose internals
+    ``feature`` was checked against; logs once per feature when it is not."""
     try:
         version = importlib.metadata.version("flashinfer-python")
     except importlib.metadata.PackageNotFoundError:
         version = None
-    if version == FLASHINFER_VERSION:
+    if version is not None and Version(version) >= Version(MIN_FLASHINFER_VERSION):
         return True
     logger.warning(
-        "DeepSeek-V4.1 specialized decode: %s needs flashinfer-python %s "
+        "DeepSeek-V4.1 specialized decode: %s needs flashinfer-python >= %s "
         "(found %s); using the stock path instead.",
         feature,
-        FLASHINFER_VERSION,
+        MIN_FLASHINFER_VERSION,
         version,
     )
     return False
diff --git a/vllm/models/deepseek_v41/nvidia/specialized_decode/attention/mxfp8.py b/vllm/models/deepseek_v41/nvidia/specialized_decode/attention/mxfp8.py
index 4014147ba0be..18a79fb02758 100644
--- a/vllm/models/deepseek_v41/nvidia/specialized_decode/attention/mxfp8.py
+++ b/vllm/models/deepseek_v41/nvidia/specialized_decode/attention/mxfp8.py
@@ -33,7 +33,7 @@
 from vllm.platforms import current_platform
 from vllm.utils import flashinfer as vllm_flashinfer
 
-from .. import LOW_BATCH_MAX_ROWS, uses_checked_flashinfer
+from .. import LOW_BATCH_MAX_ROWS
 
 BACKEND = "cutedsl_low_latency"
 
@@ -134,12 +134,10 @@ def run_zeros(linear: nn.Module, rows: int) -> torch.Tensor:
 @contextmanager
 def pinned_tactics() -> Iterator[None]:
     """Offer FlashInfer's tuner only the pinned tactic for pinned problems. The
-    tuner's tactic list is a FlashInfer internal: off the checked version,
-    every tactic is tuned."""
-    if not (
-        current_platform.is_device_capability(103)
-        and uses_checked_flashinfer("pinning the MXFP8 GEMM tactics")
-    ):
+    tuner's tactic list is a FlashInfer internal: when it has no
+    ``autotune_tactics`` or does not list a pinned tactic, every tactic is
+    tuned."""
+    if not current_platform.is_device_capability(103):
         yield
         return
     module = importlib.import_module(
diff --git a/vllm/models/deepseek_v41/nvidia/specialized_decode/attention/sparse_mla.py b/vllm/models/deepseek_v41/nvidia/specialized_decode/attention/sparse_mla.py
index 383093c27ec6..a0e68f271032 100644
--- a/vllm/models/deepseek_v41/nvidia/specialized_decode/attention/sparse_mla.py
+++ b/vllm/models/deepseek_v41/nvidia/specialized_decode/attention/sparse_mla.py
@@ -18,11 +18,14 @@
   backend.
 """
 
+import inspect
+from functools import cache
 from typing import TYPE_CHECKING, cast
 
 import torch
 
 from vllm.forward_context import get_forward_context
+from vllm.logger import init_logger
 from vllm.models.deepseek_v41.sparse_mla import DeepseekV4FlashMLAMetadata
 from vllm.utils.flashinfer import flashinfer_trtllm_batch_decode_sparse_mla_dsv4
 
@@ -30,7 +33,8 @@
     DeepseekV4FlashInferMLAAttention,
     _get_flashinfer_dsv4_workspace,
 )
-from .. import uses_checked_flashinfer
+
+logger = init_logger(__name__)
 
 if TYPE_CHECKING:
     from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata
@@ -41,17 +45,40 @@
 _kv_counters_by_device: dict[torch.device, torch.Tensor] = {}
 
 
+@cache
+def _takes_kv_counters() -> bool:
+    """Whether the installed FlashInfer's sparse-MLA decode takes caller-owned
+    KV counters (``multi_ctas_kv_counter_buffer``, since 0.7.1); logs once when
+    it does not."""
+    try:
+        import flashinfer.utils
+        from flashinfer.decode import trtllm_batch_decode_sparse_mla_dsv4
+    except ImportError:
+        return False
+    if "multi_ctas_kv_counter_buffer" in inspect.signature(
+        trtllm_batch_decode_sparse_mla_dsv4
+    ).parameters and hasattr(
+        flashinfer.utils, "get_trtllm_gen_multi_ctas_kv_counter_bytes"
+    ):
+        return True
+    logger.warning(
+        "DeepSeek-V4.1 specialized decode: the installed FlashInfer's sparse-MLA "
+        "decode takes no KV counter buffer; each decode allocates its own."
+    )
+    return False
+
+
 def kv_counter_buffer(max_rows: int) -> torch.Tensor | None:
     """Zeroed multi-CTA KV counters on the current device for decode launches
-    of up to ``max_rows`` rows (requests x heads), or None off the checked
-    FlashInfer version.
+    of up to ``max_rows`` rows (requests x heads), or None when FlashInfer's
+    decode takes no counter buffer.
 
     Allocated when the executor is built, outside any CUDA graph capture. Every
     layer shares the buffer, as it shares the workspace: their decode launches
     run one after another on one stream. Decode launches that overlap on one
     device would race on it.
     """
-    if not uses_checked_flashinfer("the sparse-MLA decode's persistent KV counters"):
+    if not _takes_kv_counters():
         return None
     from flashinfer.utils import (
         get_device_sm_count,
diff --git a/vllm/models/deepseek_v41/nvidia/specialized_decode/boundary/publishing_gemm.py b/vllm/models/deepseek_v41/nvidia/specialized_decode/boundary/publishing_gemm.py
index 87e7e45fa099..3677a3432b3c 100644
--- a/vllm/models/deepseek_v41/nvidia/specialized_decode/boundary/publishing_gemm.py
+++ b/vllm/models/deepseek_v41/nvidia/specialized_decode/boundary/publishing_gemm.py
@@ -101,10 +101,7 @@
     QuantizedActivation,
     as_quantized_activation,
 )
-from vllm.models.deepseek_v41.common.ops.query_quant import mxfp8_scale_bytes
-
-from ...ops.cute_dsl.all_reduce_mhc import ACTIVE_STAGE, NEXT_STAGE
-from ...ops.cute_dsl.primitives import (
+from vllm.models.common.ops.moe_tail_fusion.primitives import (
     VEC_BF16,
     _asm,
     current_cu_stream,
@@ -113,6 +110,9 @@
     shuffle_sync_idx_u32,
     stmc_bf16x8,
 )
+from vllm.models.deepseek_v41.common.ops.query_quant import mxfp8_scale_bytes
+
+from ...ops.cute_dsl.all_reduce_mhc import ACTIVE_STAGE, NEXT_STAGE
 from .. import LOW_BATCH_MAX_ROWS
 from ..attention.mxfp8 import Tactic, mxfp8_eligible
 from .all_reduce import BoundaryAllReduce
diff --git a/vllm/models/deepseek_v41/nvidia/specialized_decode/moe/routed_expert_front_cute.py b/vllm/models/deepseek_v41/nvidia/specialized_decode/moe/routed_expert_front_cute.py
index 99a87c22bcc7..871b3032ddef 100644
--- a/vllm/models/deepseek_v41/nvidia/specialized_decode/moe/routed_expert_front_cute.py
+++ b/vllm/models/deepseek_v41/nvidia/specialized_decode/moe/routed_expert_front_cute.py
@@ -74,7 +74,8 @@
 from cutlass._mlir.dialects import llvm
 from cutlass.cutlass_dsl import T, dsl_user_op
 
-from ...ops.cute_dsl.primitives import load_global_u32x4
+from vllm.models.common.ops.moe_tail_fusion.primitives import load_global_u32x4
+
 from .routed_expert_route_cute import (
     MAX_EXPERTS,
     WORDS,
diff --git a/vllm/models/deepseek_v41/nvidia/specialized_decode/moe/routed_expert_gemm.py b/vllm/models/deepseek_v41/nvidia/specialized_decode/moe/routed_expert_gemm.py
index f7120d8948dc..590fe2675a2a 100644
--- a/vllm/models/deepseek_v41/nvidia/specialized_decode/moe/routed_expert_gemm.py
+++ b/vllm/models/deepseek_v41/nvidia/specialized_decode/moe/routed_expert_gemm.py
@@ -55,8 +55,8 @@
 
 from vllm.logger import init_logger
 from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput
+from vllm.models.common.ops.moe_tail_fusion.primitives import current_cu_stream
 
-from ...ops.cute_dsl.primitives import current_cu_stream
 from .routed_expert_front_cute import RoutedExpertFront
 from .routed_expert_gemm_cute import (
     RoutedExpertGemm1,
diff --git a/vllm/models/deepseek_v41/nvidia/specialized_decode/runtime/model_state.py b/vllm/models/deepseek_v41/nvidia/specialized_decode/runtime/model_state.py
index 2fc7ea5c8ff1..251acde01edc 100644
--- a/vllm/models/deepseek_v41/nvidia/specialized_decode/runtime/model_state.py
+++ b/vllm/models/deepseek_v41/nvidia/specialized_decode/runtime/model_state.py
@@ -154,6 +154,7 @@ def capture_inner_cudagraphs(
         block_tables: BlockTables,
         attn_groups: list[list[AttentionGroup]],
         kv_cache_config: KVCacheConfig,
+        stream: torch.cuda.Stream | None = None,
     ) -> None:
         # Record the metadata build of every decode step the FULL decode graphs
         # serve, on dummy batches of that step's shape.
@@ -173,5 +174,5 @@ def capture_inner_cudagraphs(
                 finally:
                     self._recording = None
         super().capture_inner_cudagraphs(
-            input_buffers, block_tables, attn_groups, kv_cache_config
+            input_buffers, block_tables, attn_groups, kv_cache_config, stream
         )
diff --git a/vllm/models/deepseek_v41/nvidia/vl_model.py b/vllm/models/deepseek_v41/nvidia/vl_model.py
index 19fddfd20bd2..a605369595ed 100644
--- a/vllm/models/deepseek_v41/nvidia/vl_model.py
+++ b/vllm/models/deepseek_v41/nvidia/vl_model.py
@@ -41,7 +41,6 @@
     DeepseekV4ViT,
     run_dp_sharded_vision_tower,
 )
-from vllm.models.deepseek_v4.common.vl_model import stream_language_model_first
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.utils.tensor_schema import TensorSchema, TensorShape
 
@@ -315,19 +314,5 @@ def get_mtp_target_hidden_states(self) -> torch.Tensor | None:
         return self.language_model.get_mtp_target_hidden_states()
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        # The child finalizes in load_weights, so its weights must reach it as
-        # one contiguous group.
-        mapped = stream_language_model_first(self.hf_to_vllm_mapper.apply(weights))
         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()
+        return loader.load_weights(self.hf_to_vllm_mapper.apply(weights))
diff --git a/vllm/models/dots3_note/nvidia/model.py b/vllm/models/dots3_note/nvidia/model.py
index 15128132befe..58a57f5e521b 100644
--- a/vllm/models/dots3_note/nvidia/model.py
+++ b/vllm/models/dots3_note/nvidia/model.py
@@ -44,6 +44,7 @@
     PPMissingLayer,
     make_empty_intermediate_tensors_factory,
     make_layers,
+    pp_topk_buffers,
     sequence_parallel_chunk,
 )
 from vllm.models.deepseek_v32.nvidia.model import (
@@ -582,7 +583,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
         self.use_sequence_parallel = False
         self.is_v32 = True
         self._weight_block_size = getattr(quant_config, "weight_block_size", None)
-        topk_indices_buffer = torch.empty(
+        self.topk_indices_buffer = torch.empty(
             vllm_config.scheduler_config.max_num_batched_tokens,
             config.index_topk,
             dtype=torch.int32,
@@ -603,7 +604,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
             lambda prefix: Dots3NoteDecoderLayer(
                 vllm_config=vllm_config,
                 prefix=prefix,
-                topk_indices_buffer=topk_indices_buffer,
+                topk_indices_buffer=self.topk_indices_buffer,
             ),
             prefix=f"{prefix}.layers",
         )
@@ -611,8 +612,11 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
             self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
         else:
             self.norm = PPMissingLayer()
+        self.pp_buffers = pp_topk_buffers(
+            config, self.topk_indices_buffer, range(config.num_hidden_layers)
+        )
         self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
-            ["hidden_states", "residual"], config.hidden_size
+            ["hidden_states"], config.hidden_size, self.pp_buffers
         )
         self.aux_hidden_state_layers = tuple[int, ...]()
         self.num_redundant_experts = (
diff --git a/vllm/models/dots3_note/nvidia/multimodal.py b/vllm/models/dots3_note/nvidia/multimodal.py
index 421d185bb2dd..38cc19854cf8 100644
--- a/vllm/models/dots3_note/nvidia/multimodal.py
+++ b/vllm/models/dots3_note/nvidia/multimodal.py
@@ -318,10 +318,6 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
             mapper=self.hf_to_vllm_mapper,
         )
 
-    def process_weights_after_loading(self) -> None:
-        if self.visual is not None:
-            self.visual.process_weights_after_loading()
-
     def get_mm_mapping(self) -> MultiModelKeys:
         return MultiModelKeys.from_string_field(
             language_model="language_model",
diff --git a/vllm/models/dots3_note/nvidia/vision.py b/vllm/models/dots3_note/nvidia/vision.py
index 774d88830125..e49b1ab605f6 100644
--- a/vllm/models/dots3_note/nvidia/vision.py
+++ b/vllm/models/dots3_note/nvidia/vision.py
@@ -547,17 +547,18 @@ def device(self) -> torch.device:
         return self.patch_embed.proj.weight.device
 
     def process_weights_after_loading(self) -> None:
-        for block in self.blocks:
-            if isinstance(block.mlp, MoESwiGLUFFNFP8):
-                block.mlp.process_weights_after_loading()
         if self.config.enable_torch_compile:
             self.compile_block_modules()
 
     def compile_block_modules(self, **compile_kwargs: Any) -> None:
+        # In-place compile keeps checkpoint parameter names (no ``_orig_mod``);
+        # post-load processing reruns on every reload, so compile once.
+        if getattr(self, "_blocks_compiled", False):
+            return
         for block in self.blocks:
-            block.attn = torch.compile(block.attn, **compile_kwargs)
-            block.norm_1 = torch.compile(block.norm_1, **compile_kwargs)
-            block.norm_2 = torch.compile(block.norm_2, **compile_kwargs)
+            for module in (block.attn, block.norm_1, block.norm_2):
+                module.compile(**compile_kwargs)
+        self._blocks_compiled = True
 
     def get_pos_ids_by_grid(self, grid_thw_cpu: torch.Tensor):
         # Mirrors ``cybertron`` ``AIMv2NativeModel.rot_pos_emb``: when ``pre_pixel_shuffle``
diff --git a/vllm/models/glm5next/common/kda.py b/vllm/models/glm5next/common/kda.py
index f7103e75fd45..702d799433ce 100644
--- a/vllm/models/glm5next/common/kda.py
+++ b/vllm/models/glm5next/common/kda.py
@@ -424,11 +424,12 @@ def _flashkda_prefill(
         )
         if out is None:
             out = workspace_out[:, : q.shape[1]]
-        # FlashKDA hardcodes dense q/k/v/g strides; beta may be row-strided.
+        # FlashKDA accepts row-strided q/k/v (contiguous H/D rows, aligned
+        # token stride), so the packed q|k|v views are passed without a copy.
         torch.ops._flashkda_C.fwd(
-            q.contiguous(),
-            k.contiguous(),
-            v.contiguous(),
+            q,
+            k,
+            v,
             g.contiguous(),
             beta,
             self.head_dim**-0.5,
diff --git a/vllm/models/glm5next/common/model.py b/vllm/models/glm5next/common/model.py
index 6b0cbe4ad61b..e944ad38f151 100644
--- a/vllm/models/glm5next/common/model.py
+++ b/vllm/models/glm5next/common/model.py
@@ -8,6 +8,7 @@
 from torch import nn
 from transformers import Glm5NextTextConfig
 
+from vllm._aiter_ops import rocm_aiter_ops
 from vllm.config import ParallelConfig, VllmConfig
 from vllm.distributed import (
     get_ep_group,
@@ -24,7 +25,6 @@
     fused_moe_make_expert_params_mapping,
 )
 from vllm.model_executor.layers.fused_moe.utils import (
-    is_model_fused_shared_expert_compatible,
     resolve_layer_fused_shared_expert,
 )
 from vllm.model_executor.layers.fusion.mm_input_norm import build_mm_input_norm
@@ -76,6 +76,8 @@
 from vllm.model_executor.models.utils import (
     AutoWeightsLoader,
     PPMissingLayer,
+    extract_layer_index,
+    get_fused_shared_expert_layers,
     init_vllm_registered_model,
     is_pp_missing_parameter,
     make_layers,
@@ -749,8 +751,10 @@ def get_layer(prefix: str):
         # The active slice is fixed after construction; cache it so forward
         # doesn't rebuild the slice (a fresh list) every step.
         self._active_layers = self.layers[self.start_layer : self.end_layer]
-        self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible(
-            self.layers, Glm5NextMoE, "mlp"
+        self.fused_shared_expert_layers = get_fused_shared_expert_layers(
+            self,
+            Glm5NextMoE,
+            enable_eplb=vllm_config.parallel_config.enable_eplb,
         )
 
         if get_pp_group().is_last_rank:
@@ -894,13 +898,15 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
                 0,
             )
             num_fused_shared = _num_fused_shared_experts(
-                self.config.n_shared_experts, self.is_fused_shared_expert_enabled
+                self.config.n_shared_experts,
+                bool(self.fused_shared_expert_layers),
             )
             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",
+                # NOTE: global despite per-layer FSE; non-fused layers skip the extras.
                 num_experts=self.config.n_routed_experts + num_fused_shared,
                 num_redundant_experts=num_redundant_experts,
             )
@@ -930,7 +936,12 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
                 # Models trained using ColossalAI may include these tensors in
                 # the checkpoint. Skip them.
                 continue
-            if self.is_fused_shared_expert_enabled:
+            # Only parse the layer index when FSE is enabled.
+            if (
+                "mlp.shared_experts." in name
+                and bool(rocm_aiter_ops.is_fusion_moe_shared_experts_enabled())
+                and extract_layer_index(name) in self.fused_shared_expert_layers
+            ):
                 name = _fused_shared_expert_name(name, self.config.n_routed_experts)
 
             # Handle FP8 indexer WK: dequantize to BF16 for fusion with
diff --git a/vllm/models/glm5next/common/mtp.py b/vllm/models/glm5next/common/mtp.py
index 6e6950ca700f..cc738fdf4d00 100644
--- a/vllm/models/glm5next/common/mtp.py
+++ b/vllm/models/glm5next/common/mtp.py
@@ -6,25 +6,30 @@
 import torch
 import torch.nn as nn
 
+from vllm._aiter_ops import rocm_aiter_ops
 from vllm.config import VllmConfig
 from vllm.model_executor.layers.fused_moe import (
     fused_moe_make_expert_params_mapping,
 )
-from vllm.model_executor.layers.fused_moe.utils import (
-    is_model_fused_shared_expert_compatible,
-)
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
 from vllm.model_executor.layers.vocab_parallel_embedding import (
     VocabParallelEmbedding,
 )
+from vllm.model_executor.model_loader.mtp_validation import (
+    is_mtp_completeness_check_enabled,
+)
 from vllm.model_executor.model_loader.weight_utils import (
     default_weight_loader,
     maybe_remap_kv_scale_name,
 )
 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.model_executor.models.utils import (
+    extract_layer_index,
+    get_fused_shared_expert_layers,
+    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
@@ -212,8 +217,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
             vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
         )
         self.set_moe_parameters()
-        self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible(
-            self.model.layers.values(), Glm5NextMoE, "mtp_block.mlp"
+        self.fused_shared_expert_layers = get_fused_shared_expert_layers(
+            self,
+            Glm5NextMoE,
+            enable_eplb=vllm_config.parallel_config.enable_eplb,
         )
 
     def set_moe_parameters(self):
@@ -296,14 +303,17 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
             ("wk_weights_proj", "wk", 0),
             ("wk_weights_proj", "weights_proj", 1),
         ]
+        fused_shared_expert_layers = self.fused_shared_expert_layers
         num_fused_shared = _num_fused_shared_experts(
-            self.config.n_shared_experts, self.is_fused_shared_expert_enabled
+            self.config.n_shared_experts,
+            bool(fused_shared_expert_layers),
         )
         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",
+            # NOTE: global despite per-layer FSE; non-fused layers skip the extras.
             num_experts=self.config.n_routed_experts + num_fused_shared,
             num_redundant_experts=self.num_redundant_experts,
         )
@@ -329,7 +339,12 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
             if spec_layer is None:
                 continue
             name = self._rewrite_spec_layer_name(spec_layer, name)
-            if self.is_fused_shared_expert_enabled:
+            # Only parse the layer index when FSE is enabled.
+            if (
+                "mlp.shared_experts." in name
+                and bool(rocm_aiter_ops.is_fusion_moe_shared_experts_enabled())
+                and extract_layer_index(name) in fused_shared_expert_layers
+            ):
                 name = _fused_shared_expert_name(name, self.config.n_routed_experts)
 
             if _try_load_fp8_indexer_wk(
@@ -426,7 +441,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
             self.model.mtp_start_layer_idx,
             self.model.mtp_start_layer_idx + self.model.num_mtp_layers,
         ):
-            if layer_idx not in loaded_layers:
+            if layer_idx not in loaded_layers and is_mtp_completeness_check_enabled():
                 raise ValueError(
                     f"MTP speculative decoding layer {layer_idx} weights "
                     f"missing from checkpoint."
diff --git a/vllm/models/glm5next/common/multimodal.py b/vllm/models/glm5next/common/multimodal.py
index 23c90cce7326..88286929ffa1 100644
--- a/vllm/models/glm5next/common/multimodal.py
+++ b/vllm/models/glm5next/common/multimodal.py
@@ -713,13 +713,15 @@ def _get_video_second_idx_glm46v(
             getattr(video_processor, "temporal_patch_size", 1) or 1
         )
         video_fps = float(metadata["fps"])
+        do_sample_frames = metadata.get("do_sample_frames", True)
 
-        if not metadata.get("do_sample_frames", True):
-            # The loader already picked the frames and the processor keeps them
-            # as they are, so the frame count is what was handed in and
-            # ``frames_indices`` only supplies the timestamp values.
+        if not do_sample_frames:
+            # Loader or HF write-back already picked the frames; use those
+            # indices for timestamps. After HF write-back, do_sample_frames is
+            # cleared so the count is len(frames_indices); a pre-sampled loader
+            # hands in frames whose count matches frames_indices.
             frame_indices = [int(idx) for idx in metadata["frames_indices"]]
-            num_frames = int(total_frames)
+            num_frames = len(frame_indices)
         else:
             frame_indices = [
                 int(idx)
@@ -820,10 +822,7 @@ def _get_vision_info(
 class Glm5NextMultiModalProcessor(Glm4vMultiModalProcessor):
     """``Glm5NextProcessor`` extracts image/video features
     only and passes the prompt text through unchanged, so prompt expansion
-    (image token repeat, video frame/timestamp structure) is owned by vLLM's
-    prompt-update machinery — the inherited ``_get_prompt_updates`` builds
-    the replacement content and the placeholder scan validates against
-    exactly that."""
+    is owned by the inherited GLM-4V prompt-update machinery."""
 
     def _hf_processor_applies_updates(
         self,
diff --git a/vllm/models/glm5next/common/sparse_indexer.py b/vllm/models/glm5next/common/sparse_indexer.py
index ad3845457063..5fcd473c00e8 100644
--- a/vllm/models/glm5next/common/sparse_indexer.py
+++ b/vllm/models/glm5next/common/sparse_indexer.py
@@ -16,8 +16,10 @@ def _build_decode_scatter_indices(
     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).
+    batch. ``n`` is the graph-padded token count and may be larger than
+    ``decode_lens.sum()`` under adaptive verification. All graph-padding tokens
+    are mapped to a dummy request that is appended as the last one
+    (``idx=num_requests``).
 
     Shared by every ``_scatter_decode_tokens_by_request`` call in a step:
     building it per call would repeat the same repeat_interleave/cumsum chain
@@ -25,20 +27,34 @@ def _build_decode_scatter_indices(
     """
     device = decode_lens.device
     dl = decode_lens.to(torch.int64)
+    # FULL CUDA graphs pad to the captured token count. Adaptive verification
+    # changes the live per-request lengths only on the GPU, so the padding
+    # amount must remain device-resident too.
+    graph_padding = n - dl.sum()
+    dl_with_dummy = torch.cat((dl, graph_padding.reshape(1)))
     req_id = torch.repeat_interleave(
-        torch.arange(num_requests, device=device, dtype=torch.int64),
-        dl,
+        torch.arange(num_requests + 1, device=device, dtype=torch.int64),
+        dl_with_dummy,
         output_size=n,
     )
     req_starts = torch.cumsum(
-        torch.cat([torch.zeros(1, device=device, dtype=torch.int64), dl[:-1]]),
+        torch.cat(
+            [
+                torch.zeros(1, device=device, dtype=torch.int64),
+                dl_with_dummy[:-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)
+    # dl_with_dummy.sum()) so each token's intra-request index subtracts its
+    # own request's start.
+    starts = torch.repeat_interleave(req_starts, dl_with_dummy, output_size=n)
     intra = torch.arange(n, device=device, dtype=torch.int64) - starts
+    # Padding may produce arbitrarily large intra-request indices, while the
+    # the dummy request has a fixed width. Since padding values are ignored,
+    # map every padding token to column zero.
+    intra = torch.where(req_id == num_requests, 0, intra)
     return req_id, intra
 
 
@@ -59,14 +75,17 @@ def _scatter_decode_tokens_by_request(
     reshape.
     """
     req_id, intra = scatter_indices
+    # The final row is the dummy request. It receives graph-padding tokens and
+    # is removed from the result. This avoids out-of-bounds advanced indexing
+    # without letting padding overwrite a live request's first token.
     out = torch.full(
-        (num_requests, lmax, *tokens.shape[1:]),
+        (num_requests + 1, lmax, *tokens.shape[1:]),
         pad_value,
         dtype=tokens.dtype,
         device=tokens.device,
     )
     out[req_id, intra] = tokens
-    return out
+    return out[:num_requests]
 
 
 def _decode_topk_seq_lens(
diff --git a/vllm/models/glm5next/nvidia/sparse_indexer.py b/vllm/models/glm5next/nvidia/sparse_indexer.py
index 65ba7be1fdec..0b8d4fa788f2 100644
--- a/vllm/models/glm5next/nvidia/sparse_indexer.py
+++ b/vllm/models/glm5next/nvidia/sparse_indexer.py
@@ -11,6 +11,13 @@
 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.kernels.attention.dsa.fused_mqa_topk import (
+    fused_mqa_topk_chunk,
+    fused_mqa_topk_key_scales,
+    fused_mqa_topk_prefill,
+    get_fused_mqa_topk_workspaces,
+    use_fused_mqa_topk,
+)
 from vllm.model_executor.layers.indexer_topk import get_indexer_topk
 from vllm.model_executor.layers.sparse_attn_indexer import (
     PrefillRowShard,
@@ -44,6 +51,7 @@
 )
 from vllm.v1.attention.backends.mla.indexer import (
     DeepseekV32IndexerMetadata,
+    dense_logits_rows,
 )
 from vllm.v1.attention.ops.common import pack_seq_triton, unpack_seq_triton
 from vllm.v1.worker.workspace import current_workspace_manager
@@ -144,6 +152,13 @@ def sparse_attn_indexer_kpool(
     fp8_dtype = current_platform.fp8_dtype()
     k_cache_prefix = _resolve_layer_name(k_cache_prefix)
 
+    fused_mqa_topk = use_fused_mqa_topk(
+        query_shape=q_quant.shape[1:],
+        topk=topk_tokens // max(index_kpool, 1),
+        fp4_cache=use_fp4_cache,
+        full_context=dcp_world_size == 1,
+    )
+
     # assert isinstance(attn_metadata, dict)
     if not isinstance(attn_metadata, dict):
         # Reserve workspace for indexer during profiling run
@@ -155,6 +170,10 @@ def sparse_attn_indexer_kpool(
             scales_spec,
             ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8),
         )
+        # Allocate the largest fused_mqa_topk workspaces so they are accounted
+        # for (freed again, they overlap the dense logits below).
+        if fused_mqa_topk:
+            get_fused_mqa_topk_workspaces(q_quant.device, q_quant.shape[0])
 
         # Reserve profiler-visible memory for the worst-case decode logits,
         # whose shape is [B * next_n, max_pool_len]. This profiling branch
@@ -355,51 +374,75 @@ def sparse_attn_indexer_kpool(
                 ).squeeze(-1)
 
             num_rows = q_slice.shape[0]
-            if chunk.local_total_seq_lens == 0:
-                logits = q_slice.new_empty((num_rows, 0), dtype=torch.float32)
-            else:
-                q_scale_slice, chunk_weights = mqa_logits_mx_inputs(
-                    q_slice_cast, q_scale_slice, weights[row_start:row_end]
-                )
-                # Compressed: row i starts at K position cu_seqlen_ks[i].
-                row_ends = cu_seqlen_ke - cu_seqlen_ks
-                logits = as_fp32_logits(
-                    fp8_fp4_mqa_logits(
-                        (q_slice_cast, q_scale_slice),
-                        (k_quant_cast, k_scale_cast),
-                        chunk_weights,
-                        cu_seqlen_ks,
-                        cu_seqlen_ke,
-                        chunk.max_seqlen_k,
-                    ),
-                    row_ends,
-                )
-
             # 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
+                    (num_rows, select_k), dtype=torch.int32, device=q_slice.device
                 )
                 topk_dst = pool_topk
             else:
                 topk_dst = topk_indices_buffer[row_start:row_end, :topk_tokens]
 
-            if logits.shape[1] == 0:
+            # Rows per dense logits call under the logits budget.
+            dense_rows = dense_logits_rows(chunk.max_seqlen_k)
+            if chunk.local_total_seq_lens == 0:
+                logits = q_slice.new_empty((num_rows, 0), dtype=torch.float32)
                 topk_dst.fill_(-1)
-            else:
-                torch.ops._C.top_k_per_row_prefill(
-                    logits,
-                    torch.zeros_like(row_ends),
-                    row_ends,
+            elif fused_mqa_topk and fused_mqa_topk_chunk(
+                chunk.max_seqlen_k, k_quant.shape[0], select_k, num_rows, dense_rows
+            ):
+                fused_mqa_topk_prefill(
+                    q_slice,
+                    k_quant,
+                    fused_mqa_topk_key_scales(k_scale),
+                    weights[row_start:row_end],
+                    cu_seqlen_ks,
+                    cu_seqlen_ke,
                     topk_dst,
-                    num_rows,
-                    logits.stride(0),
-                    logits.stride(1),
-                    select_k,
                 )
+                logits = q_slice.new_empty((num_rows, 0), dtype=torch.float32)
+            else:
+                # Chunks planned for fused_mqa_topk are not bounded by the logits
+                # budget: split their rows (budget-planned chunks, the only ones
+                # under DCP, are one window, so the merge below sees all the
+                # rows).
+                # One pass also for 0 rows: the merge is collective.
+                for sub in range(0, max(1, num_rows), dense_rows):
+                    rows = slice(sub, min(sub + dense_rows, num_rows))
+                    q_scale_rows, chunk_weights = mqa_logits_mx_inputs(
+                        q_slice_cast[rows],
+                        q_scale_slice[rows] if q_scale_slice is not None else None,
+                        weights[row_start + rows.start : row_start + rows.stop],
+                    )
+                    # Compressed: row i starts at K position cu_seqlen_ks[i].
+                    row_ends = cu_seqlen_ke[rows] - cu_seqlen_ks[rows]
+                    logits = as_fp32_logits(
+                        fp8_fp4_mqa_logits(
+                            (q_slice_cast[rows], q_scale_rows),
+                            (k_quant_cast, k_scale_cast),
+                            chunk_weights,
+                            cu_seqlen_ks[rows],
+                            cu_seqlen_ke[rows],
+                            chunk.max_seqlen_k,
+                        ),
+                        row_ends,
+                    )
+                    if logits.shape[1] == 0:
+                        topk_dst[rows].fill_(-1)
+                    else:
+                        torch.ops._C.top_k_per_row_prefill(
+                            logits,
+                            torch.zeros_like(row_ends),
+                            row_ends,
+                            topk_dst[rows],
+                            logits.shape[0],
+                            logits.stride(0),
+                            logits.stride(1),
+                            select_k,
+                        )
             # Local pool ids -> request-relative global pool ids.
             _merge_dcp_topk_global(
                 logits,
diff --git a/vllm/models/hy_v4/amd/model.py b/vllm/models/hy_v4/amd/model.py
index ceb061136806..0aee05f07e18 100644
--- a/vllm/models/hy_v4/amd/model.py
+++ b/vllm/models/hy_v4/amd/model.py
@@ -53,19 +53,22 @@
     AutoWeightsLoader,
     PPMissingLayer,
     WeightsMapper,
+    compute_skip_topk_layers,
+    copy_pp_buffers,
     get_pp_missing_layer_names,
     is_pp_missing_parameter,
     make_empty_intermediate_tensors_factory,
     make_layers,
     maybe_prefix,
+    pp_topk_buffers,
 )
 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,
+    sparse_layer_ids,
 )
 from ..nvidia.moe import HYV4FeedForward, HYV4MoEFused
 from .hc import (
@@ -287,6 +290,9 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
             self.norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
         else:
             self.norm = PPMissingLayer()
+        self.pp_buffers = pp_topk_buffers(
+            config, self.topk_indices_buffer, sparse_layer_ids(config)
+        )
         if self.enable_ihc:
             # iHC head layer: merge the residual channels into one stream.
             if get_pp_group().is_last_rank:
@@ -301,13 +307,15 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
                 self.hc_head = PPMissingLayer()
             self.make_empty_intermediate_tensors = (
                 make_empty_intermediate_tensors_factory(
-                    ["hidden_states"], config.hc_mult * config.hidden_size
+                    ["hidden_states"],
+                    config.hc_mult * config.hidden_size,
+                    self.pp_buffers,
                 )
             )
         else:
             self.make_empty_intermediate_tensors = (
                 make_empty_intermediate_tensors_factory(
-                    ["hidden_states", "residual"], config.hidden_size
+                    ["hidden_states", "residual"], config.hidden_size, self.pp_buffers
                 )
             )
 
@@ -398,17 +406,21 @@ def forward(
             # 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"]
+            copy_pp_buffers(intermediate_tensors, self.pp_buffers)
 
         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:
+            buffers = {k: v[: positions.shape[0]] for k, v in self.pp_buffers.items()}
             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.flatten(1), **buffers}
+                )
             return IntermediateTensors(
-                {"hidden_states": hidden_states, "residual": residual}
+                {"hidden_states": hidden_states, "residual": residual, **buffers}
             )
         if self.enable_ihc:
             hidden_states = self.hc_head(hidden_states)
@@ -728,8 +740,5 @@ 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 | drop_mapper)
 
-    def process_weights_after_loading(self) -> None:
-        self.model.process_weights_after_loading()
-
     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 d10d3afe9025..a9efd38c92fd 100644
--- a/vllm/models/hy_v4/nvidia/attention.py
+++ b/vllm/models/hy_v4/nvidia/attention.py
@@ -42,6 +42,7 @@
 from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer
 from vllm.model_executor.model_loader.weight_utils import make_attention_sink
 from vllm.model_executor.models.deepseek_v2 import DeepseekV32IndexerCache
+from vllm.model_executor.models.utils import compute_skip_topk_layers
 from vllm.platforms import current_platform
 from vllm.v1.attention.backend import AttentionBackend, AttentionType
 from vllm.v1.attention.selector import get_attn_backend
@@ -58,59 +59,12 @@
 _WEIGHT_LAYER_INDEX_RE = re.compile(r"(?:^|\.)layers\.(\d+)(?:\.|$)")
 
 
-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
-    by the closest preceding "full" indexer layer, so it does not build its own
-    indexer and its checkpoint indexer weights must be skipped.
-
-    Args:
-        config: The model config.
-
-    Returns:
-        The set of layer indices that share another layer's top-k indices.
-
-    Raises:
-        ValueError: If ``indexer_types`` has the wrong length or an unknown
-            entry, or if ``index_topk_freq`` is not a positive integer.
-
-    """
+def sparse_layer_ids(config: PreTrainedConfig) -> set[int]:
+    """Layers that run sparse attention, MTP layers included."""
     if not hasattr(config, "index_topk"):
         return set()
-
-    num_hidden_layers = config.num_hidden_layers
-    indexer_types = getattr(config, "indexer_types", None)
-    if indexer_types is not None:
-        if len(indexer_types) != num_hidden_layers:
-            raise ValueError(
-                "indexer_types must contain one entry per hidden layer: "
-                f"expected {num_hidden_layers}, got {len(indexer_types)}."
-            )
-        invalid_types = sorted(set(indexer_types) - {"full", "shared"})
-        if invalid_types:
-            raise ValueError(
-                f"indexer_types only supports 'full' and 'shared', got {invalid_types}."
-            )
-        return {
-            layer_idx
-            for layer_idx, indexer_type in enumerate(indexer_types)
-            if indexer_type == "shared"
-        }
-
-    freq = getattr(config, "index_topk_freq", 1)
-    if not isinstance(freq, int) or freq <= 0:
-        raise ValueError(f"index_topk_freq must be a positive integer, got {freq!r}.")
-    pattern = getattr(config, "index_topk_pattern", None)
-    offset = getattr(config, "index_skip_topk_offset", 2)
-    skip_layers: set[int] = set()
-    for layer_idx in range(num_hidden_layers):
-        if pattern is None:
-            if max(layer_idx - offset + 1, 0) % freq != 0:
-                skip_layers.add(layer_idx)
-        elif 0 <= layer_idx < len(pattern) and pattern[layer_idx] == "S":
-            skip_layers.add(layer_idx)
-    return skip_layers
+    layer_types = getattr(config, "layer_types", None) or ()
+    return {i for i, t in enumerate(layer_types) if t in _SPARSE_LAYER_TYPES}
 
 
 def is_skip_topk_indexer_weight(weight_name: str, skip_topk_layers: set[int]) -> bool:
@@ -319,13 +273,7 @@ def __init__(
         self.num_local_heads = num_heads // tp_size
 
         self.layer_id = int(prefix.split(".")[-2])
-        layer_types = getattr(config, "layer_types", None)
-        requested_sparse = (
-            hasattr(config, "index_topk")
-            and layer_types is not None
-            and self.layer_id < len(layer_types)
-            and layer_types[self.layer_id] in _SPARSE_LAYER_TYPES
-        )
+        requested_sparse = self.layer_id in sparse_layer_ids(config)
         # Only actual sparse layers may share another layer's top-k indices.
         self.skip_topk = requested_sparse and self.layer_id in compute_skip_topk_layers(
             config
diff --git a/vllm/models/hy_v4/nvidia/model.py b/vllm/models/hy_v4/nvidia/model.py
index 6e91ec242221..946b64cfe3c0 100644
--- a/vllm/models/hy_v4/nvidia/model.py
+++ b/vllm/models/hy_v4/nvidia/model.py
@@ -52,19 +52,22 @@
     AutoWeightsLoader,
     PPMissingLayer,
     WeightsMapper,
+    compute_skip_topk_layers,
+    copy_pp_buffers,
     get_pp_missing_layer_names,
     is_pp_missing_parameter,
     make_empty_intermediate_tensors_factory,
     make_layers,
     maybe_prefix,
+    pp_topk_buffers,
 )
 from vllm.platforms import current_platform
 from vllm.sequence import IntermediateTensors
 
 from .attention import (
     HYV4MLAAttention,
-    compute_skip_topk_layers,
     is_skip_topk_indexer_weight,
+    sparse_layer_ids,
 )
 from .hc import HYV4HCHeadLayer, HYV4HCLayer
 from .moe import HYV4FeedForward, HYV4MoEFused
@@ -260,6 +263,9 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
             self.norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
         else:
             self.norm = PPMissingLayer()
+        self.pp_buffers = pp_topk_buffers(
+            config, self.topk_indices_buffer, sparse_layer_ids(config)
+        )
         if self.enable_ihc:
             # iHC head layer: merge the residual channels into one stream.
             if get_pp_group().is_last_rank:
@@ -274,13 +280,15 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
                 self.hc_head = PPMissingLayer()
             self.make_empty_intermediate_tensors = (
                 make_empty_intermediate_tensors_factory(
-                    ["hidden_states"], config.hc_mult * config.hidden_size
+                    ["hidden_states"],
+                    config.hc_mult * config.hidden_size,
+                    self.pp_buffers,
                 )
             )
         else:
             self.make_empty_intermediate_tensors = (
                 make_empty_intermediate_tensors_factory(
-                    ["hidden_states", "residual"], config.hidden_size
+                    ["hidden_states", "residual"], config.hidden_size, self.pp_buffers
                 )
             )
 
@@ -371,17 +379,21 @@ def forward(
             # 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"]
+            copy_pp_buffers(intermediate_tensors, self.pp_buffers)
 
         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:
+            buffers = {k: v[: positions.shape[0]] for k, v in self.pp_buffers.items()}
             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.flatten(1), **buffers}
+                )
             return IntermediateTensors(
-                {"hidden_states": hidden_states, "residual": residual}
+                {"hidden_states": hidden_states, "residual": residual, **buffers}
             )
         if self.enable_ihc:
             hidden_states = self.hc_head(hidden_states)
diff --git a/vllm/models/kimi_k3/amd/kda.py b/vllm/models/kimi_k3/amd/kda.py
index 48ac78338571..896bd29823db 100644
--- a/vllm/models/kimi_k3/amd/kda.py
+++ b/vllm/models/kimi_k3/amd/kda.py
@@ -73,6 +73,35 @@
 logger = init_logger(__name__)
 
 
+class _PaddedKimiGDNMergedColumnParallelLinear(_KimiGDNMergedColumnParallelLinear):
+    """Merged projection with ``padding_size`` zero rows appended per rank."""
+
+    def __init__(self, *args, padding_size: int = 0, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        self.padding_size = padding_size
+        if padding_size:
+            assert isinstance(self.weight, torch.Tensor)
+            # Exclude padding rows from the load-completion count.
+            set_weight_attrs(
+                self.weight,
+                {"weight_loader_numel": self.weight[:-padding_size].numel()},
+            )
+            self.weight.data[-padding_size:].zero_()
+
+    def weight_loader(self, param, loaded_weight, loaded_shard_id=None) -> None:
+        super().weight_loader(param, loaded_weight, loaded_shard_id)
+        self._zero_padding(param)
+
+    def weight_loader_v2(self, param, loaded_weight, loaded_shard_id=None) -> None:
+        super().weight_loader_v2(param, loaded_weight, loaded_shard_id)
+        self._zero_padding(param)
+
+    def _zero_padding(self, param: torch.Tensor) -> None:
+        # Online processing materializes the weight after __init__ zeroing.
+        if self.padding_size and param is self.weight:
+            param.data[-self.padding_size :].zero_()
+
+
 class KimiK3DeltaAttention(GatedDeltaNetAttention):
     def get_attn_backend(self) -> type[AttentionBackend]:
         return KimiK3ROCmKDABackend
@@ -98,10 +127,18 @@ def get_state_shape(self) -> tuple[tuple[int, ...], tuple[int, ...]]:
     def get_kv_cache_spec(self, vllm_config: VllmConfig) -> MambaSpec:
         spec = super().get_kv_cache_spec(vllm_config)
         assert isinstance(spec, MambaSpec)
-        # Only the fused path exports the checkpoint snapshots
+        # Only the fused path exports the checkpoint snapshots. AITER's
+        # chunk_kda cannot, so the gluon backend hands the steps that carry a
+        # checkpoint to the fused kernel when the build has it.
         can_checkpoint = (
             self.use_prefill_checkpoint
-            and self.kda_prefill_backend == "fused"
+            and (
+                self.kda_prefill_backend == "fused"
+                or (
+                    self.kda_prefill_backend == "gluon"
+                    and is_fused_kda_chunk_supported()
+                )
+            )
             and self.use_safe_gate
         )
         if can_checkpoint and vllm_config.cache_config.mamba_cache_mode == "align":
@@ -159,17 +196,16 @@ def __init__(
         if self.in_proj_padding:
             in_proj_output_sizes.append(self.in_proj_padding * self.tp_size)
 
-        self.in_proj_qkvgfab = _KimiGDNMergedColumnParallelLinear(
+        self.in_proj_qkvgfab = _PaddedKimiGDNMergedColumnParallelLinear(
             self.hidden_size,
             in_proj_output_sizes,
             replicated_shard_id=4,
             tp_size=self.tp_size,
+            padding_size=self.in_proj_padding,
             bias=False,
             quant_config=self.quant_config,
             prefix=f"{prefix}.in_proj_qkvgfab",
         )
-        if self.in_proj_padding:
-            self.in_proj_qkvgfab.weight.data[-self.in_proj_padding :].zero_()
 
         self.f_b_proj = ColumnParallelLinear(
             self.head_dim,
@@ -309,7 +345,7 @@ def __init__(
         compilation_config.static_forward_context[prefix] = self
 
     @torch.no_grad()
-    def process_weights_after_loading(self, act_dtype: torch.dtype) -> None:
+    def process_weights_after_loading(self) -> None:
         if is_weights_pre_processed():
             return
         if self.decode_conv1d_weight is not None:
@@ -732,7 +768,12 @@ def _prefill_conv(
                     cu_seqlens=prefill_query_start_loc,
                     chunk_indices=m.chunk_indices,
                     chunk_offsets=m.chunk_offsets,
-                    backend=self.kda_prefill_backend,
+                    # Checkpoint metadata exists only when get_kv_cache_spec
+                    # reserved the block, which it does only with the fused
+                    # kernel available.
+                    backend=(
+                        "fused" if checkpoint is not None else self.kda_prefill_backend
+                    ),
                     out=core_attn_out[:, nd_tok:num_actual_tokens] if direct else None,
                     checkpoint_state=checkpoint_state,
                     checkpoint_offsets=checkpoint_offsets,
diff --git a/vllm/models/kimi_k3/amd/linear.py b/vllm/models/kimi_k3/amd/linear.py
index 0671a69d293e..9d4d47edb3f4 100644
--- a/vllm/models/kimi_k3/amd/linear.py
+++ b/vllm/models/kimi_k3/amd/linear.py
@@ -1171,11 +1171,6 @@ def get_mamba_state_copy_func(
     ) -> tuple[MambaStateCopyFunc, MambaStateCopyFunc]:
         return MambaStateCopyFuncCalculator.kda_state_copy_func()
 
-    def process_weights_after_loading(self) -> None:
-        for module in self.modules():
-            if isinstance(module, KimiK3MultiHeadLatentAttentionWrapper):
-                module.process_weights_after_loading()
-
     def compute_logits(
         self,
         hidden_states: torch.Tensor,
diff --git a/vllm/models/kimi_k3/amd/mla.py b/vllm/models/kimi_k3/amd/mla.py
index 73f7e506b76d..0c30e0d634b8 100644
--- a/vllm/models/kimi_k3/amd/mla.py
+++ b/vllm/models/kimi_k3/amd/mla.py
@@ -7,6 +7,7 @@
 import torch
 
 from vllm._aiter_ops import rocm_aiter_ops
+from vllm.config import get_current_vllm_config
 from vllm.model_executor.layers.attention.attention import get_attention_context
 from vllm.model_executor.layers.attention.mla_attention import split_kv_b_proj
 from vllm.model_executor.layers.fusion.quant_activation import (
@@ -46,6 +47,18 @@ def __init__(self, *args, **kwargs) -> None:
             self._fused_q_heads = AiterMLAHelper.get_actual_mla_num_heads(
                 self.num_heads
             )
+        # Positions fed to the fused kernel's identity RoPE; see _fused_decode.
+        # Allocated once here so every captured CUDA graph shares one zeroed
+        # buffer instead of capturing a memset per layer.
+        zero_positions = None
+        if self._fused_qk_prep:
+            vllm_config = get_current_vllm_config()
+            max_tokens = max(
+                vllm_config.scheduler_config.max_num_batched_tokens,
+                vllm_config.compilation_config.max_cudagraph_capture_size or 0,
+            )
+            zero_positions = torch.zeros(max_tokens, dtype=torch.int64)
+        self.register_buffer("_zero_positions", zero_positions, persistent=False)
         # The fold reads q_b_proj.weight as plain bf16/fp16 rows, so a quantized
         # q_b_proj keeps the unfused path.
         self._fold_qb_uk = (
@@ -174,7 +187,6 @@ def _fused_decode(
         q: torch.Tensor,
         kv_c_normed: torch.Tensor,
         k_pe: torch.Tensor,
-        positions: torch.Tensor,
         output_shape: tuple[int, int],
         attn_metadata,
         layer,
@@ -228,9 +240,9 @@ def _fused_decode(
 
         if self._identity_rope is None:
             # cos = 1, sin = 0 makes the kernel's RoPE the identity, which is
-            # what a NoPE model needs. One row suffices because the `_opt`
-            # kernel clamps `pos` into [0, cos_cache.size(0)); a full-length
-            # cache would cost ~268 MB at K3's 1M max_position_embeddings.
+            # what a NoPE model needs. A full-length cache would cost ~268 MB at
+            # K3's 1M max_position_embeddings, so it holds one row and the
+            # kernel is given all-zero positions below.
             half = self.qk_rope_head_dim // 2
             opts = {"dtype": kv_c_normed.dtype, "device": kv_c_normed.device}
             self._identity_rope = (
@@ -253,6 +265,15 @@ def _fused_decode(
             device=ql_nope.device,
         )
         q_out = q_out_padded[:, :num_heads]
+        # The kernel must not see the real positions: not every AITER variant
+        # clamps them into the one-row cos/sin cache. The head-grouped kernel
+        # (AITER v0.1.26) reads cos = sin = 0 past the cache at one head per
+        # block, zeroing q_pe and the cached k_pe, and faults at more heads per
+        # block. RoPE is the identity, so position 0 is exact for every token.
+        # The kernel takes its token count from slot_mapping.
+        slot_mapping = slot_mapping.flatten()
+        num_slots = slot_mapping.shape[0]
+        assert num_slots <= self._zero_positions.shape[0]
         fused_qk_rope_concat_and_cache_mla(
             ql_nope,
             q_pe,
@@ -260,10 +281,10 @@ def _fused_decode(
             k_pe.squeeze(1),
             kv_cache.view(kv_cache.shape[0], -1, head_size),
             q_out,
-            slot_mapping.flatten(),
+            slot_mapping,
             layer._k_scale,
             layer._q_scale,
-            positions,
+            self._zero_positions[:num_slots],
             cos_cache,
             sin_cache,
             is_neox=True,
@@ -383,7 +404,6 @@ def forward(
                 q,
                 kv_c_normed,
                 k_pe,
-                positions,
                 output_shape,
                 attn_metadata,
                 layer,
diff --git a/vllm/models/kimi_k3/amd/mtp.py b/vllm/models/kimi_k3/amd/mtp.py
index 8f719adb1055..20c8c54c6020 100644
--- a/vllm/models/kimi_k3/amd/mtp.py
+++ b/vllm/models/kimi_k3/amd/mtp.py
@@ -20,6 +20,9 @@
     ParallelLMHead,
     VocabParallelEmbedding,
 )
+from vllm.model_executor.model_loader.mtp_validation import (
+    is_mtp_completeness_check_enabled,
+)
 from vllm.model_executor.model_loader.weight_utils import (
     default_weight_loader,
     maybe_remap_kv_scale_name,
@@ -361,7 +364,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
             self.model.mtp_start_layer_idx,
             self.model.mtp_start_layer_idx + self.model.num_mtp_layers,
         ):
-            if layer_idx not in loaded_layers:
+            if layer_idx not in loaded_layers and is_mtp_completeness_check_enabled():
                 raise ValueError(
                     f"MTP speculative decoding layer {layer_idx} weights "
                     f"missing from checkpoint. The checkpoint may not include "
diff --git a/vllm/models/kimi_k3/amd/specialized_decode/README.md b/vllm/models/kimi_k3/amd/specialized_decode/README.md
index 8417964eb20e..79f00504a7d8 100644
--- a/vllm/models/kimi_k3/amd/specialized_decode/README.md
+++ b/vllm/models/kimi_k3/amd/specialized_decode/README.md
@@ -1,24 +1,36 @@
 # Kimi-K3 specialized decode (MI355X)
 
 An opt-in decode path for Kimi-K3 at TP8 on MI355X (gfx950), following RFC #60904.
-A pure spec-verify step of at most 8 rows runs each MoE decoder layer as persistent
+A pure spec-verify step of at most 16 rows runs each MoE decoder layer as persistent
 FlyDSL MonoKernels (256 CTAs, all resident), in "layer mode", inside the stock
 layer loop. The TP all-reduces run in-kernel over peer memory.
 
 - **Whole layers** (KDA): K1 (AttnRes, in_proj, f_b, spec conv, KDA recurrence,
   gated RMSNorm), then K2.
-- **FFN layers** (MLA): vLLM's AttnRes and MLA up to o_proj's input, then K2.
+- **FFN layers** (MLA, DCP1 or DCP8): M1a (AttnRes, qkv_a and g_proj, q / kv
+  RMSNorm, the fp8 KV write of this rank's tokens, q_b, MXFP4 W_UK, the fp8 query
+  gathered to every rank's peer memory), AITER's MLA decode on this rank's KV
+  shard, then K2, which runs M1b first (the LSE combine of every rank's partials,
+  MXFP4 W_UV, the sigmoid gate). At DCP1 every rank holds the whole KV: M1a writes
+  its 12 heads' query directly and M1b reads AITER's output as is, with no gather
+  and no combine.
+- **FFN layers M1 cannot serve** (a non-fp8 KV cache, `VLLM_ROCM_USE_AITER_FP4BMM`
+  unset, or MLA weights M1 does not read): vLLM's AttnRes and MLA attention ops
+  write o_proj's input, then K2.
 - **K2**: o_proj + all-reduce + MLP AttnRes, then the latent MoE (router,
   top-16, MXFP4 experts, shared experts, latent all-reduce + RMSNorm, up_proj,
   final all-reduce).
 
-The kernels read vLLM's loaded weights and KDA states in place. Nothing is copied
-or re-quantized.
+The kernels read vLLM's loaded weights, KDA states and MLA KV cache in place.
+Nothing is copied or re-quantized. The kernels tile up to 64 rows (`MAX_TOKENS`),
+but at 32 and 64 rows they lose to vLLM's decode, so a served step is capped at 16
+(`MAX_ROWS`).
 
 ## Enabling
 
 ```bash
-VLLM_USE_V2_MODEL_RUNNER=1 vllm serve <Kimi-K3> -tp 8 \
+VLLM_USE_V2_MODEL_RUNNER=1 VLLM_ROCM_USE_AITER_FP4BMM=1 vllm serve <Kimi-K3> -tp 8 \
+  [--decode-context-parallel-size 8 --dcp-comm-backend a2a] --kv-cache-dtype fp8 \
   --speculative-config '{"method": "dspark", "num_speculative_tokens": 7, ...}' \
   --additional-config '{"decode_backend": "latency_specialized"}'
 ```
@@ -28,8 +40,14 @@ There are no other options. `decode_backend=auto` (the default) runs vLLM's deco
 ## Supported deployments
 
 TP8 on MI355X with all 256 CUs, decode context parallel 1 or 8, speculative
-decoding with at most 7 speculative tokens, Model Runner V2, BF16 conv / FP32
-recurrent KDA states.
+decoding with at most 15 speculative tokens, Model Runner V2, BF16 conv / FP32
+recurrent KDA states. M1 assumes no KV cache interleave (it writes the slots vLLM
+maps), but at DCP8 vLLM's AITER MLA backend serves DSpark at interleave 1 only.
+
+M1 also needs an fp8 KV cache and MLA's W_UK / W_UV as AITER MXFP4 bmms
+(`VLLM_ROCM_USE_AITER_FP4BMM=1`). Without them the MLA layers run vLLM's
+attention ops, then K2, and startup logs why. The KDA layers run K1 and K2 either
+way.
 
 KV transfer support is limited to built-in MoRIIO READ and WRITE consumers.
 Keep the producer on `decode_backend=auto` and use matching cache layouts.
@@ -42,11 +60,14 @@ If you request the specialized decode for anything else, startup fails with one
   routed-expert capture, sleep mode, CPU weight offload, weight updates,
   no speculative decoding,
   adaptive verification, dynamic speculative decoding, ReplaySSM / RecoverSSM,
-  the V1 model runner, `STOCK_TORCH_COMPILE`, another model's shapes, another GPU.
+  the V1 model runner, `STOCK_TORCH_COMPILE`, another model's shapes, another
+  GPU.
 - **Loaded weights, when the model state builds the executor:** each kernel
-  layer's weight shapes and dtypes, and the MoE backend (`AITER_MXFP4_BF16`, K3
-  SiTU).
-- **KV caches, at every bind:** the KDA state dtypes and layout.
+  layer's weight shapes and dtypes and the MoE backend (`AITER_MXFP4_BF16`, K3
+  SiTU). Each MLA layer's shapes, gate, DCP setup, fp8 cache and bmm weights
+  decide M1 or vLLM's ops (not a refusal).
+- **KV caches, at every bind:** the KDA state dtypes and layout, and with M1 the
+  MLA cache's fp8 dense rows.
 
 Every TP rank agrees on each verdict before anything collective runs.
 
@@ -54,16 +75,22 @@ Every TP rank agrees on each verdict before anything collective runs.
 
 `plan()` runs once per step. It picks the kernels for a pure spec-verify step of
 whole uniform requests whose width is one of the built widths: every multiple of
-`num_speculative_tokens + 1` up to 8 rows. Every other step runs vLLM's layers:
-prefill, mixed, wider and PIECEWISE steps.
+`num_speculative_tokens + 1` up to 16 rows. With M1 the MLA layers' metadata must
+agree (pure decode, the same width, AITER's non-Gluon decode). Every other step
+runs vLLM's layers: prefill, mixed, wider and PIECEWISE steps.
 
 Under FULL graphs the plan is baked in at capture. vLLM's capture sizes are not
 changed. A replay with fewer requests than its graph pads whole requests. Those
 requests get the reserved null state slot and one accepted token, and the kernels
-serve them like real ones. For example, at `num_speculative_tokens=1` three
-requests replay the 4-request graph. To capture every width exactly, pass them in
+serve them like real ones. To capture every width exactly, pass them in
 `--cudagraph-capture-sizes`.
 
+Every launch's mailbox tag is the step epoch and a slot: K1 and M1a use the layer
+index, K2 (and the M1b it runs) 128 + the layer index. The peer buffer holds K2's
+regions, then M1a's gathered query, then M1b's partials (both M1 at DCP8 only). A
+region is reused by the next layer only after K2's all-reduce, which every rank
+reaches past its M1b.
+
 ## Fail-stop
 
 Every in-kernel wait gives up after 10 s. The rank records the wait in its health
@@ -89,22 +116,27 @@ Everything lives in this folder. The model files only call in: `linear.py`
 (the config check in `__init__`, `plan_specialized_decode` in `forward`,
 `step.layer(i)` in the layer loop, `after_specialized_decode_step` in
 `compute_logits`), `model.py` (`get_model_state_cls`) and `mla.py`
-(`skip_o_proj`, o_proj's input for the MLA layers' K2).
+(`skip_o_proj`). Nothing outside the model changes: M1a writes the query
+AITER's cprr decode pads to (e.g. 96 -> 128 heads), and `runtime/mla.py` hands
+it to `forward_mqa` in place of the backend's padding copy.
 
 ```text
 hooks.py         SpecializedDecodeMixin: the model-side hooks
 model_state.py   KimiK3ModelState: builds the executor after load, binds KV caches
 runtime/
-    __init__.py  DecodeConfig, specialized_decode_config(), unsupported()
+    __init__.py  DecodeConfig, specialized_decode_config(), unsupported(),
+                 m1_reasons()
     decode.py    SpecializedDecode: plan() -> Step, Step.layer(i), prepare_kv(),
                  after_step()
     guards.py    weight and KV cache refusals, the fail-stop watch
     layers.py    the layer kernels: vLLM tensors -> launches
+    mla.py       the MLA layers' decode as vLLM ops: the step's context,
+                 AITER's decode (the M1a / M1b reference)
     runner.py    per-rank launch state: scratch, step epoch, health, peers
     peer.py      peer memory and TP agreement (adapted from ROCm/ATOM, MIT)
 mono/            kernel source
     layer.py     K2
-    attention/   kda.py (K1), back.py
+    attention/   kda.py (K1), mla.py (M1a), mla_back.py (M1b), back.py
     stages/      moe.py, gemv.py
     common/      device ops, mailbox sync, MX helpers, plan / build keys
                  (adapted from ROCm/ATOM atom/mono, MIT)
diff --git a/vllm/models/kimi_k3/amd/specialized_decode/__init__.py b/vllm/models/kimi_k3/amd/specialized_decode/__init__.py
index 4a88231b5408..84ff4d501814 100644
--- a/vllm/models/kimi_k3/amd/specialized_decode/__init__.py
+++ b/vllm/models/kimi_k3/amd/specialized_decode/__init__.py
@@ -1,7 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Kimi-K3 specialized decode for ROCm (MI355X, TP8), layer mode: FlyDSL
-MonoKernels per MoE decoder layer on pure spec-verify steps, inside the stock layer
+MonoKernels per MoE decoder layer on pure spec-verify steps (K1 / K2 on the KDA
+layers, M1a / M1b / K2 on the MLA layers, DCP1 or DCP8), inside the stock layer
 loop. Enabled by ``additional_config={"decode_backend": "latency_specialized"}``.
 
 ``hooks.py`` holds the model-side hooks (``SpecializedDecodeMixin``) and
diff --git a/vllm/models/kimi_k3/amd/specialized_decode/mono/attention/back.py b/vllm/models/kimi_k3/amd/specialized_decode/mono/attention/back.py
index 9efe5d7f478e..bf4bb255c9dc 100644
--- a/vllm/models/kimi_k3/amd/specialized_decode/mono/attention/back.py
+++ b/vllm/models/kimi_k3/amd/specialized_decode/mono/attention/back.py
@@ -22,7 +22,7 @@
 
 import flydsl.expr as fx
 from aiter.ops.flydsl.kernels import buffer_ops as bo
-from flydsl.expr import gpu, range_constexpr
+from flydsl.expr import const_expr, gpu, range_constexpr
 from flydsl.expr.typing import T
 
 from vllm.models.kimi_k3.amd.specialized_decode.mono.attention.kda import (
@@ -31,11 +31,19 @@
     LDS_PAD,
     NA,
     PROJ,
+    QC,
     ROWS,
+    bf4_words,
+    bf16r,
 )
 from vllm.models.kimi_k3.amd.specialized_decode.mono.common.ops import (
     CM_NT,
+    bf2_f32,
+    bf16_pair,
     bf16_round,
+    bf_hi,
+    bf_lo,
+    butterfly,
     mfma_bf16,
     row_sum,
     rsrc,
@@ -48,18 +56,81 @@
     pair_layout,
 )
 from vllm.models.kimi_k3.amd.specialized_decode.mono.common.ranks import sum_partials
+from vllm.models.kimi_k3.amd.specialized_decode.mono.common.sync import (
+    Addr,
+    poll_sent,
+    put_sent,
+    rearm_sent,
+    sent_word,
+)
 
 TP = 8
+OWN_COLS = HIDDEN // TP  # two-shot: the columns a rank reduces (7 slices)
+
+
+def two_shot(s):
+    """Above 8 rows the AR scatters each partial to its columns' owner rank,
+    which reduces and broadcasts: 2/8 of the one-shot's bytes a link."""
+    return s > 8
+
+
+def peer_at(peers, p):
+    """``peers[p]`` (an Addr a rank) for a traced rank ``p``."""
+    v = peers[0].value
+    for q in range(1, TP):
+        v = (p == q).select(peers[q].value, v)
+    return Addr(peers[0].region, peers[0].space, v)
+
+
 OROWS = HIDDEN  # o_proj rows (full hidden, a rank's K slice)
 OTASKS = OROWS // ROWS  # 448
 OK = PROJ  # K = 1536 a rank
 OKCH = OK // 128  # 12 chunks
 GEMV0 = NA  # the first gemv CTA
 GEMV_N = BLOCKS - NA  # 200
+TL_POINTS = 6  # start, gemv, pushed, reduced, slice/gate, mix
 
 _STREAM = fx.Stream(None)
 
 
+def sent_on(c, which="ar"):
+    """(OPT_SENTS) This build's all-reduce ``which`` ("ar": o_proj's,
+    "latr", "final") goes through its sentinel region (``layer.k2_sent``; a
+    plain call: a method call on ``c`` inside a traced branch would make FlyDSL
+    carry ``c`` as branch state)."""
+    return c.get("sent_" + which, False)
+
+
+def sent_ix(src, col, t, s, width):
+    """(OPT_SENT) Word index of (src rank, pair at ``col``, token t), blocked
+    `[src][16-col group][token][8 pairs]`; shared by push and reduce."""
+    return ((src * (width // 16) + col // 16) * s + t) * 8 + (col % 16) // 2
+
+
+def sent_item(idx, s):
+    """(OPT_SENT) A reduce item's (token, column pair), so consecutive items
+    read consecutive ``sent_ix`` words."""
+    g = idx // (8 * s)
+    r = idx % (8 * s)
+    return r // 8, g * 8 + r % 8
+
+
+def sent_sum(c, own, col, t, width, rearm=True):
+    """The ``TP`` ranks' sentinel words of (col pair, t) in ``own`` summed in
+    fp32 in rank order (as ``sum_partials``), re-armed unless ``rearm`` is
+    off (the caller re-arms them later)."""
+    s = c["S"]
+    ws = poll_sent(
+        c["mb"], own, [sent_ix(src, col, t, s, width) for src in range(TP)], rearm=rearm
+    )
+    acc0, acc1 = bf2_f32(ws[0])
+    for src in range_constexpr(1, TP):
+        lo, hi = bf2_f32(ws[src])
+        acc0 = acc0 + lo
+        acc1 = acc1 + hi
+    return acc0, acc1
+
+
 def ar_half_pairs(s):
     """One parity's AR region: ``[src rank][token][column pair]``."""
     return TP * s * HIDDEN // 2
@@ -105,15 +176,22 @@ def o_loads(c, rg):
 
 
 @traced
-def load_core(c):
-    """The core rows (S x 1536 bf16) -> LDS, rows padded."""
-    s, tid = c["S"], c["tid"]
-    words = s * OK // 8
+def load_core(c, t0, n, cm=0):
+    """Core rows t0 .. t0 + n (of S x 1536 bf16) -> LDS, rows padded; ``cm``
+    CM_DEV when this launch's own CTAs wrote them (K2's fused M1b)."""
+    tid = c["tid"]
+    words = n * OK // 8
     row = OK + LDS_PAD
     for i in range_constexpr((words + THREADS - 1) // THREADS):
         e = fx.min(tid + THREADS * i, words - 1)
         v = fx.Vector(
-            bo.buffer_load(rsrc(c["args"]["core"]), e * 4, vec_width=4, dtype=T.i32)
+            bo.buffer_load(
+                rsrc(c["args"]["core"]),
+                (t0 * (OK // 8) + e) * 4,
+                vec_width=4,
+                dtype=T.i32,
+                cache_modifier=cm,
+            )
         )
         t = e // (OK // 8)
         k = e % (OK // 8) * 8
@@ -123,12 +201,13 @@ def load_core(c):
 
 
 @traced
-def stage_o(c, rg, ops):
-    """Rows 16 rg ..: GEMV -> bf16 partials -> every rank's AR region, this
-    rank's slot (thread: token, row pair; a put a peer)."""
+def stage_o(c, rg, ops, t0, n):
+    """Rows 16 rg .. for the core tile t0 .. t0 + n: GEMV -> bf16 partials ->
+    every rank's AR region, this rank's slot (thread: token, row pair; a put a
+    peer)."""
     s, tid, lane, wave, red = c["S"], c["tid"], c["lane"], c["wave"], c["red"]
     g = lane // ROWS
-    t = fx.min(lane % ROWS, s - 1)
+    t = fx.min(lane % ROWS, n - 1)
     row = OK + LDS_PAD
     acc = fx.Vector.filled(4, 0.0, fx.Float32)
     for i in range_constexpr(2):
@@ -150,14 +229,28 @@ def stage_o(c, rg, ops):
             )
     fx.ptr_store(acc, red + (wave * 64 + lane) * 4)
     gpu.barrier()
-    if tid < s * (ROWS // 2):
+    if tid < n * (ROWS // 2):
         tt = tid // (ROWS // 2)
         rp = tid % (ROWS // 2)
         v0 = bf16_round(row_sum(red, 2 * rp, tt))
         v1 = bf16_round(row_sum(red, 2 * rp + 1, tt))
-        at = (c["rank"] * s + tt) * HIDDEN + rg * ROWS + 2 * rp
-        for p in range_constexpr(TP):
-            c["put_bf"](c["ar_peer"][p], at, [v0, v1])
+        at = (c["rank"] * s + t0 + tt) * HIDDEN + rg * ROWS + 2 * rp
+        if const_expr(sent_on(c)):
+            # (OPT_SENT) A plain word to the sentinel region
+            si = sent_ix(c["rank"], rg * ROWS + 2 * rp, t0 + tt, s, HIDDEN)
+            word = sent_word(bf16_pair(v0, v1))
+            if const_expr(two_shot(s)):
+                owner = rg * ROWS // OWN_COLS
+                put_sent(peer_at(c["ars_peer"], owner), si, word)
+            else:
+                for p in range_constexpr(TP):
+                    put_sent(c["ars_peer"][p], si, word)
+        elif const_expr(two_shot(s)):
+            owner = rg * ROWS // OWN_COLS
+            c["put_bf"](peer_at(c["ar_peer"], owner), at, [v0, v1])
+        else:
+            for p in range_constexpr(TP):
+                c["put_bf"](c["ar_peer"][p], at, [v0, v1])
     gpu.barrier()
 
 
@@ -166,13 +259,200 @@ def stage_reduce(c, task):
     """Columns 128 task ..: the 8 ranks' partials, rank order, -> bf16 delta in
     LDS (thread: token, column pair)."""
     s, tid = c["S"], c["tid"]
-    if tid < s * (COLS // 2):
-        t = tid // (COLS // 2)
-        cp = tid % (COLS // 2)
-        col = task * COLS + 2 * cp
-        lo, hi = sum_partials(
-            c["poll"], c["ar_own"], lambda src: ((src * s + t) * HIDDEN + col) // 2, TP
+    n = s * (COLS // 2)
+    for i in range_constexpr((n + THREADS - 1) // THREADS):
+        idx = tid + THREADS * i
+        if idx < n:
+            if const_expr(sent_on(c)):
+                t, cp = sent_item(idx, s)
+            else:
+                t = idx // (COLS // 2)
+                cp = idx % (COLS // 2)
+            col = task * COLS + 2 * cp
+            if const_expr(two_shot(s)):
+                lo, hi = reduce_two_shot(c, task, t, col)
+            elif const_expr(sent_on(c)):
+                # re-armed after the mix (``rearm_ar``)
+                lo, hi = sent_sum(c, c["ars_own"], col, t, HIDDEN, rearm=False)
+            else:
+                lo, hi = sum_partials(
+                    c["poll"],
+                    c["ar_own"],
+                    lambda src, t=t, col=col: ((src * s + t) * HIDDEN + col) // 2,
+                    TP,
+                )
+            fx.ptr_store(bf16_round(lo), c["delta_lds"] + (t * COLS + 2 * cp))
+            fx.ptr_store(bf16_round(hi), c["delta_lds"] + (t * COLS + 2 * cp + 1))
+    gpu.barrier()
+
+
+@traced
+def reduce_two_shot(c, task, t, col):
+    """The owner of ``col`` sums the 8 scattered partials (rank order, as the
+    one-shot) and puts the bf16 sum to the other ranks' slot [owner]; they
+    poll it there (a slot the scatter leaves free off the owner)."""
+    s = c["S"]
+    owner = task * COLS // OWN_COLS
+    at = (owner * s + t) * HIDDEN + col
+    lo, hi = fx.Float32(0.0), fx.Float32(0.0)
+    if const_expr(sent_on(c)):
+        # (OPT_SENT) The scatter and the broadcast through the
+        # sentinel region, the same slots as the tagged pairs'
+        si = sent_ix(owner, col, t, s, HIDDEN)
+        if owner == c["rank"]:
+            a0, a1 = sent_sum(c, c["ars_own"], col, t, HIDDEN, rearm=False)
+            word = sent_word(bf16_pair(a0, a1))
+            for p in range_constexpr(TP):
+                if p != c["rank"]:
+                    put_sent(c["ars_peer"][p], si, word)
+            lo, hi = a0, a1
+        else:
+            lo, hi = bf2_f32(poll_sent(c["mb"], c["ars_own"], [si], rearm=False)[0])
+        return lo, hi
+    if owner == c["rank"]:
+        a0, a1 = sum_partials(
+            c["poll"],
+            c["ar_own"],
+            lambda src: ((src * s + t) * HIDDEN + col) // 2,
+            TP,
         )
-        fx.ptr_store(bf16_round(lo), c["delta_lds"] + (t * COLS + 2 * cp))
-        fx.ptr_store(bf16_round(hi), c["delta_lds"] + (t * COLS + 2 * cp + 1))
+        for p in range_constexpr(TP):
+            if p != c["rank"]:
+                c["put_bf"](c["ar_peer"][p], at, [a0, a1])
+        lo, hi = a0, a1
+    else:
+        lo, hi = bf2_f32(c["poll"]([(c["ar_own"], at // 2, 1)])[0][0])
+    return lo, hi
+
+
+# ---------------------------------------------------------------- pre-AR preload
+# (``layer.OPT_PRELOAD``)
+# The slice's block / prefix / norm / qk loads (independent of the
+# all-reduce) issued before ``stage_reduce``'s poll; ``stage_slice_pre`` is
+# ``kda.stage_slice`` on the preloaded words, same arithmetic and order.
+
+
+def _raw4(ptr, i):
+    """bf16 elements i .. i + 4 as two raw words (``kda.ld_bf4``'s load)."""
+    return fx.Vector(bo.buffer_load(rsrc(ptr), i // 2, vec_width=2, dtype=T.i32))
+
+
+def _f4(w):
+    return [bf_lo(w[0]), bf_hi(w[0]), bf_lo(w[1]), bf_hi(w[1])]
+
+
+def slice_loads(c, task):
+    """``stage_slice``'s loads of slice ``task`` that need no hand-off: each
+    round's block and prefix words, then the round-invariant ares_nw / ares_qk
+    words (a live thread's columns are the same every round: THREADS % QC == 0)."""
+    s, ns, nb, tid = c["S"], c["NS"], c["NB"], c["tid"]
+    a = c["args"]
+    c0 = task * COLS
+    n = s * ns * QC
+    blk, pre = [], []
+    for i in range(-(-n // THREADS)):
+        e = fx.min(tid + THREADS * i, n - 1)
+        row = e // QC
+        t = row // ns
+        src = row % ns
+        col = c0 + (e % QC) * 4
+        blk.append(
+            _raw4(
+                a["blocks"], t * a["blk_sm"] + fx.min(src, nb - 1) * a["blk_sr"] + col
+            )
+        )
+        if const_expr(not c.get("reset", False)):
+            pre.append(_raw4(a["prefix"], t * HIDDEN + col))
+    col0 = c0 + (fx.min(tid, n - 1) % QC) * 4
+    return {
+        "blk": blk,
+        "pre": pre,
+        "w": _raw4(a["ares_nw"], col0),
+        "q": _raw4(a["ares_qk"], col0),
+    }
+
+
+@traced
+def stage_slice_pre(c, task, ld):
+    """``kda.stage_slice`` for K2 (delta in LDS, no block write) on the
+    preloaded words ``ld`` (``slice_loads``)."""
+    s, ns, nb, tid, lane = c["S"], c["NS"], c["NB"], c["tid"], c["lane"]
+    a = c["args"]
+    assert c["delta"] and c.get("delta_lds") is not None
+    assert c.get("write_idx", -1) < 0
+    c0 = task * COLS
+    n = s * ns * QC
+    w = _f4(ld["w"])
+    q = _f4(ld["q"])
+    for i in range_constexpr((n + THREADS - 1) // THREADS):
+        e = fx.min(tid + THREADS * i, n - 1)
+        live = (tid + THREADS * i) < n
+        row = e // QC
+        t = row // ns
+        src = row % ns
+        col = c0 + (e % QC) * 4
+        is_pre = src == nb
+        blk = _f4(ld["blk"][i])
+        dl = [
+            fx.ptr_load(c["delta_lds"] + (t * COLS + (e % QC) * 4 + j))
+            for j in range(4)
+        ]
+        if const_expr(c.get("reset", False)):
+            pre = dl
+        else:
+            pre = _f4(ld["pre"][i])
+            pre = [bf16r(pre[j] + dl[j]) for j in range(4)]
+        if live & is_pre:
+            bo.buffer_store(bf4_words(pre), rsrc(a["prefix"]), (t * HIDDEN + col) // 2)
+        v = [
+            live.select(is_pre.select(pre[j], blk[j]), fx.Float32(0.0))
+            for j in range(4)
+        ]
+        if const_expr(c["vals_all"]):
+            if live:
+                for j in range_constexpr(4):
+                    fx.ptr_store(v[j], c["vals"] + (row * COLS + (e % QC) * 4 + j))
+        else:
+            if live & is_pre:
+                for j in range_constexpr(4):
+                    fx.ptr_store(v[j], c["vals"] + (t * COLS + (e % QC) * 4 + j))
+        sq = v[0] * v[0] + v[1] * v[1] + v[2] * v[2] + v[3] * v[3]
+        dt = (
+            v[0] * (w[0] * q[0])
+            + v[1] * (w[1] * q[1])
+            + v[2] * (w[2] * q[2])
+            + v[3] * (w[3] * q[3])
+        )
+        sq = butterfly(sq, (16, 8, 4, 2, 1))
+        dt = butterfly(dt, (16, 8, 4, 2, 1))
+        if live & (lane % 32 == 0):
+            at = ((task * s + t) * ns + src) * 2
+            c["put"](c["part"], at, sq)
+            c["put"](c["part"], at + 1, dt)
     gpu.barrier()
+
+
+@traced
+def rearm_ar(c, task):
+    """(OPT_SENT) Re-arm the words ``stage_reduce`` read for ``task``,
+    after the mix (off the critical path)."""
+    s, tid = c["S"], c["tid"]
+    n = s * (COLS // 2)
+    for i in range_constexpr((n + THREADS - 1) // THREADS):
+        idx = tid + THREADS * i
+        if idx < n:
+            t, cp = sent_item(idx, s)
+            col = task * COLS + 2 * cp
+            if const_expr(two_shot(s)):
+                owner = task * COLS // OWN_COLS
+                if owner == c["rank"]:
+                    rearm_sent(
+                        c["ars_own"],
+                        [sent_ix(src, col, t, s, HIDDEN) for src in range(TP)],
+                    )
+                else:
+                    rearm_sent(c["ars_own"], [sent_ix(owner, col, t, s, HIDDEN)])
+            else:
+                rearm_sent(
+                    c["ars_own"], [sent_ix(src, col, t, s, HIDDEN) for src in range(TP)]
+                )
diff --git a/vllm/models/kimi_k3/amd/specialized_decode/mono/attention/kda.py b/vllm/models/kimi_k3/amd/specialized_decode/mono/attention/kda.py
index 43a624f7c723..d02c15a98b4b 100644
--- a/vllm/models/kimi_k3/amd/specialized_decode/mono/attention/kda.py
+++ b/vllm/models/kimi_k3/amd/specialized_decode/mono/attention/kda.py
@@ -41,6 +41,9 @@
 
 from vllm.models.kimi_k3.amd.specialized_decode.mono.common.debug import (
     region_ids,
+    stamp,
+    stamp_begin,
+    stamp_flush,
 )
 from vllm.models.kimi_k3.amd.specialized_decode.mono.common.ops import (
     CM_DEV,
@@ -73,6 +76,8 @@
     publish,
     sreg,
 )
+from vllm.models.kimi_k3.amd.specialized_decode.mono.stages.gemv import token_tiles
+from vllm.v1.attention.ops.flydsl_kernels import _run_compiled
 
 HIDDEN = 7168
 NA = 56  # attn CTAs
@@ -98,11 +103,20 @@
 CONV_W = 4
 LOG2E = 1.0 / math.log(2.0)
 LDS_PAD = 8  # bf16 a padded x row
+MAX_S = 64  # rows a launch
+XT = 8  # x rows a gemv tile past 8 rows (16 padded 7168-wide rows exceed LDS)
+# every source's slice in LDS up to this; past it the prefix only
+VALS_BYTES = 96 * 1024
+MSQ0 = 1024  # red: the mix's sums of squares (the weights below)
+# timeline points: start, slice, gate, mix, load_x, gemv0, gemv1, k.stage,
+# k.conv, k.gates, k.recur, k.norm
+TL_POINTS = 15
 
 assert NA + GEMV_CTAS <= BLOCKS and NA == 8 * GSL
 
 _STREAM = fx.Stream(None)
 REGIONS = ("part", "wgt", "msq", "xrdy", "proj", "qk", "onrm")
+TAG_BATCH = 14  # tagged-row specs a poll batch (56 registers; ``poll_rows_to_lds``)
 # every mono source file: in every build's JIT cache key
 _SOURCES = source_digest(".")
 
@@ -118,6 +132,7 @@ class KdaPreBuild:
     out_eps: float = 1e-5  # input_layernorm
     onorm_eps: float = 1e-5
     lower_bound: float = -5.0
+    timeline: bool = False
     write_idx: int = -1  # a block-write layer's block: the updated prefix stored there
     kda0: int = 0  # the first KDA task's CTA
 
@@ -153,6 +168,7 @@ class KdaPreBuild:
         "epoch",
         "health",
         "alarm",
+        "tl",
     )
 )
 
@@ -191,6 +207,16 @@ def step_tag(epoch, slot):
     return (e % (1 << 22)) * TAG_SLOTS + slot + 1
 
 
+def vals_all(s, ns, extra=0):
+    """Whether every (token, source)'s slice stays in LDS ``vals`` (else the
+    prefix only; ``stage_mix`` re-reads the blocks), ``extra`` bytes beside it."""
+    return s * ns * COLS * 4 + extra <= VALS_BYTES
+
+
+def vals_words(s, ns, extra=0):
+    return s * (ns if vals_all(s, ns, extra) else 1) * COLS
+
+
 def ld_bf(ptr, i, cm=0):
     return fx.Float32(
         fx.BFloat16(
@@ -291,9 +317,14 @@ def stage_slice(c, task):
         q = ld_bf4(a["ares_qk"], col)
         # a thread past the end holds a clamped copy of the last element: it
         # must not store (it raced the owner with zeros)
-        if live:
-            for j in range_constexpr(4):
-                fx.ptr_store(v[j], c["vals"] + (row * COLS + (e % QC) * 4 + j))
+        if const_expr(c["vals_all"]):
+            if live:
+                for j in range_constexpr(4):
+                    fx.ptr_store(v[j], c["vals"] + (row * COLS + (e % QC) * 4 + j))
+        else:
+            if live & is_pre:
+                for j in range_constexpr(4):
+                    fx.ptr_store(v[j], c["vals"] + (t * COLS + (e % QC) * 4 + j))
         sq = v[0] * v[0] + v[1] * v[1] + v[2] * v[2] + v[3] * v[3]
         dt = (
             v[0] * (w[0] * q[0])
@@ -350,59 +381,126 @@ def stage_gate(c, t):
     gpu.barrier()
 
 
+def mix_src(c, t, src, cc, c0):
+    """Token t's source ``src``, slice columns cc .. cc + 4: from ``vals``, or
+    (prefix only in LDS) a block re-read from ``blocks`` -- the same bf16 values
+    ``stage_slice`` read."""
+    if const_expr(c["vals_all"]):
+        return [
+            fx.ptr_load(c["vals"] + ((t * c["NS"] + src) * COLS + cc + j))
+            for j in range(4)
+        ]
+    if const_expr(src == c["NB"]):
+        return [fx.ptr_load(c["vals"] + (t * COLS + cc + j)) for j in range(4)]
+    a = c["args"]
+    col = c0 + cc
+    return ld_bf4(a["blocks"], t * a["blk_sm"] + src * a["blk_sr"] + col)
+
+
 @traced
 def stage_mix(c, task):
-    """Columns 128 task ..: the softmax mix of ``vals`` (4 columns a thread), its
-    sum of squares (MSQ), every slice's -> rstd, then bf16 x = mixed rstd w ->
-    the plain x rows, and the slice's XRDY flag."""
+    """Columns 128 task ..: the softmax mix of the sources (4 columns a thread,
+    16 tokens a round), its sum of squares (MSQ), every slice's -> rstd, then
+    bf16 x = mixed rstd w -> the plain x rows, and the slice's XRDY flag."""
     s, ns, tid, lane, red = c["S"], c["NS"], c["tid"], c["lane"], c["red"]
     a = c["args"]
     c0 = task * COLS
     nw = s * (ns + 1)
-    if tid < nw:
-        v = c["poll"]([(c["wgt"], tid, 1)])[0][0]
-        fx.ptr_store(f32_of(v), red + tid)
+    assert nw <= MSQ0 and 8 * s <= THREADS
+    # read here, not in a traced branch (a method call on ``c`` there makes
+    # FlyDSL carry ``c`` as branch state)
+    tag_x = c.get("tag_x", False)
+    for i in range_constexpr((nw + THREADS - 1) // THREADS):
+        w = tid + THREADS * i
+        if w < nw:
+            v = c["poll"]([(c["wgt"], w, 1)])[0][0]
+            fx.ptr_store(f32_of(v), red + w)
     gpu.barrier()
-    mine = tid < s * QC
-    t = fx.min(tid // QC, s - 1)
+    tpr = THREADS // QC  # tokens a round
     cc = (tid % QC) * 4
-    mixed = []
-    for j in range_constexpr(4):
-        acc = fx.Float32(0.0)
+    rounds = []
+    for rr in range_constexpr((s + tpr - 1) // tpr):
+        tr = rr * tpr + tid // QC
+        mine = tr < s
+        t = fx.min(tr, s - 1)
+        mixed = [fx.Float32(0.0) for _ in range(4)]
         for src in range_constexpr(ns):
-            acc = acc + fx.ptr_load(red + (t * (ns + 1) + src)) * fx.ptr_load(
-                c["vals"] + ((t * ns + src) * COLS + cc + j)
-            )
-        mixed.append(acc / fx.ptr_load(red + (t * (ns + 1) + ns)))
-    sq = (
-        mixed[0] * mixed[0]
-        + mixed[1] * mixed[1]
-        + mixed[2] * mixed[2]
-        + mixed[3] * mixed[3]
-    )
-    sq = butterfly(sq, (16, 8, 4, 2, 1))
-    if mine & (lane % 32 == 0):
-        c["put"](c["msq"], task * s + t, sq)
+            wt = fx.ptr_load(red + (t * (ns + 1) + src))
+            v = mix_src(c, t, src, cc, c0)
+            mixed = [mixed[j] + wt * v[j] for j in range(4)]
+        den = fx.ptr_load(red + (t * (ns + 1) + ns))
+        mixed = [mixed[j] / den for j in range(4)]
+        sq = (
+            mixed[0] * mixed[0]
+            + mixed[1] * mixed[1]
+            + mixed[2] * mixed[2]
+            + mixed[3] * mixed[3]
+        )
+        sq = butterfly(sq, (16, 8, 4, 2, 1))
+        if mine & (lane % 32 == 0):
+            c["put"](c["msq"], task * s + t, sq)
+        rounds.append((mine, t, mixed))
     gpu.barrier()
     if tid < 8 * s:
         g = tid // s
         tt = tid % s
         vals = c["poll"]([(c["msq"], (g * GSL + k) * s + tt, 1) for k in range(GSL)])
-        fx.ptr_store(sum_polled(vals), red + (128 + g * s + tt))
+        fx.ptr_store(sum_polled(vals), red + (MSQ0 + g * s + tt))
     gpu.barrier()
-    tot = fx.ptr_load(red + (128 + t))
-    for g in range_constexpr(1, 8):
-        tot = tot + fx.ptr_load(red + (128 + g * s + t))
-    r = hw_rsq(tot * fx.Float32(1.0 / HIDDEN) + fx.Float32(c["out_eps"]))
     wn = ld_bf4(a["in_nw"], c0 + cc)
-    xw = bf4_words([mixed[j] * r * wn[j] for j in range(4)])
-    if mine:
-        at = (t * HIDDEN + c0 + cc) // 2
-        bo.buffer_store(xw, rsrc(c["xbuf"]), at, cache_modifier=CM_DEV)
+    for mine, t, mixed in rounds:
+        tot = fx.ptr_load(red + (MSQ0 + t))
+        for g in range_constexpr(1, 8):
+            tot = tot + fx.ptr_load(red + (MSQ0 + g * s + t))
+        r = hw_rsq(tot * fx.Float32(1.0 / HIDDEN) + fx.Float32(c["out_eps"]))
+        xw = bf4_words([mixed[j] * r * wn[j] for j in range(4)])
+        if mine:
+            at = (t * HIDDEN + c0 + cc) // 2
+            if const_expr(tag_x):
+                # K2 (``layer.OPT_TAG_X``): the rows also as tagged pairs,
+                # which the router and latent GEMVs poll (no XRDY wait)
+                c["put_words"](c["xt"], at, [xw[0], xw[1]])
+            bo.buffer_store(xw, rsrc(c["xbuf"]), at, cache_modifier=CM_DEV)
     publish(c["put"], c["xrdy"], task, fx.Int32(1), tid == 0)
     gpu.barrier()
 
 
+@traced
+def poll_rows_to_lds(c, region, k, n, xl, xrow, ldk, k0, t0):
+    """``gemv.rows_to_lds`` from tagged pairs: columns k0 .. k0 + k of rows
+    t0 .. t0 + n of the [S, ldk] bf16 rows a producer put as (word, tag) pairs
+    into ``region``, polled (two pairs a spec, every thread's specs one batch)
+    -> LDS rows 0 .. n - 1."""
+    tid = c["tid"]
+    quads = n * k // 4  # 4 bf16 = 2 pairs a spec
+    per = (quads + THREADS - 1) // THREADS
+    # a batch at most TAG_BATCH specs (4 registers each), stored before the next
+    for b0 in range_constexpr(0, per, TAG_BATCH):
+        ii = list(range(b0, min(per, b0 + TAG_BATCH)))
+        es = [fx.min(tid + THREADS * i, quads - 1) for i in ii]
+        tk = [(e // (k // 4), e % (k // 4) * 4) for e in es]
+        got = c["poll"](
+            [(region, ((t0 + t) * ldk + k0 + kk) // 2, 2) for t, kk in tk],
+            batch=len(ii),
+        )
+        for j in range_constexpr(len(ii)):
+            t, kk = tk[j]
+            if tid + THREADS * ii[j] < quads:
+                fx.ptr_store(
+                    fx.Vector.from_elements(got[j], fx.Int32), xl + (t * xrow + kk) // 2
+                )
+    gpu.barrier()
+
+
+@traced
+def stage_gates(c, bid):
+    """The gates of tokens bid, bid + 56, .. (< S) on attn CTA ``bid``."""
+    s = c["S"]
+    for rr in range_constexpr((s + NA - 1) // NA):
+        if bid + NA * rr < s:
+            stage_gate(c, bid + NA * rr)
+
+
 # ---------------------------------------------------------------- gemv CTAs
 def gemv_loads(c, rg, chunks):
     """This wave's weight operands of rows 16 rg .. for its K chunks ``chunks``
@@ -434,19 +532,28 @@ def gemv_loads(c, rg, chunks):
 
 
 @traced
-def load_x(c):
-    """Every slice's x ready -> the bf16 x rows into LDS (rows padded)."""
-    s, tid = c["S"], c["tid"]
-    if tid < NA:
-        c["poll"]([(c["xrdy"], tid, 1)])
+def wait_x(c):
+    """Every slice's x ready."""
+    if c["tid"] < NA:
+        c["poll"]([(c["xrdy"], c["tid"], 1)])
     gpu.barrier()
-    words = s * HIDDEN // 8  # 16 B chunks
+
+
+@traced
+def load_x(c, t0, n):
+    """The bf16 x rows t0 .. t0 + n into LDS rows 0 .. (rows padded)."""
+    tid = c["tid"]
+    words = n * HIDDEN // 8  # 16 B chunks
     row = HIDDEN + LDS_PAD
     for i in range_constexpr((words + THREADS - 1) // THREADS):
         e = fx.min(tid + THREADS * i, words - 1)
         v = fx.Vector(
             bo.buffer_load(
-                rsrc(c["xbuf"]), e * 4, vec_width=4, dtype=T.i32, cache_modifier=CM_DEV
+                rsrc(c["xbuf"]),
+                (t0 * (HIDDEN // 8) + e) * 4,
+                vec_width=4,
+                dtype=T.i32,
+                cache_modifier=CM_DEV,
             )
         )
         t = e // (HIDDEN // 8)
@@ -457,11 +564,12 @@ def load_x(c):
 
 
 @traced
-def stage_gemv(c, rg, ops):
-    """Rows 16 rg .. of in_proj against the LDS x -> bf16 PROJ."""
-    s, tid, lane, wave, red = c["S"], c["tid"], c["lane"], c["wave"], c["red"]
+def stage_gemv(c, rg, ops, t0, n):
+    """Rows 16 rg .. of in_proj against the LDS x (tokens t0 .. t0 + n) -> bf16
+    PROJ."""
+    tid, lane, wave, red = c["tid"], c["lane"], c["wave"], c["red"]
     g = lane // ROWS
-    t = fx.min(lane % ROWS, s - 1)
+    t = fx.min(lane % ROWS, n - 1)
     row = HIDDEN + LDS_PAD
     acc = fx.Vector.filled(4, 0.0, fx.Float32)
     for i in range_constexpr(PER_WAVE):
@@ -478,11 +586,11 @@ def stage_gemv(c, rg, ops):
             )
     fx.ptr_store(acc, red + (wave * 64 + lane) * 4)
     gpu.barrier()
-    if (tid < ROWS * s) & (rg < GEMV_TASKS):
+    if (tid < ROWS * n) & (rg < GEMV_TASKS):
         r = tid % ROWS
         tt = tid // ROWS
         v = bf16r(row_sum(red, r, tt))
-        c["put"](c["proj"], tt * NPROJ + rg * ROWS + r, v)
+        c["put"](c["proj"], (t0 + tt) * NPROJ + rg * ROWS + r, v)
     gpu.barrier()
 
 
@@ -546,11 +654,13 @@ def kda_conv(c, k, req, h, ch):
     conv state; q / k out through QK, v to LDS."""
     L, sl, tid = c["L"], c["SL"], c["tid"]
     a = c["args"]
+    stamp(c["on"], c["tls"], tid, 12)
     if tid < 3 * VCH:
         xs = c["poll"](
             [(c["proj"], (req * L + j) * NPROJ + k["chan"], 1) for j in range(L)]
         )
         xv = [f32_of(xs[j][0]) for j in range(L)]
+        stamp(c["on"], c["tls"], tid, 13)
         rs = rsrc(k["cs"])
         if k["slot_c"] > 0:
             for j in range_constexpr(sl):
@@ -558,6 +668,7 @@ def kda_conv(c, k, req, h, ch):
                 bo.buffer_store(
                     fx.Float32(nv).to(fx.BFloat16), rs, k["co"] + j * a["cs_tok"]
                 )
+        stamp(c["on"], c["tls"], tid, 14)
         seq = k["win"] + xv
         for j in range_constexpr(L):
             acc = fx.Float32(0.0)
@@ -707,17 +818,22 @@ def kda_norm(c, req, h, ch):
 
 
 @traced
-def kda_tasks(c):
-    s, L, bid = c["S"], c["L"], c["bid"]
+def kda_tasks(c, on, tls):
+    s, L, tid, bid = c["S"], c["L"], c["tid"], c["bid"]
     for task in range(first_task(bid, c["kda0"]), (s // L) * NH * NVCH, BLOCKS):
         req = task // (NH * NVCH)
         h = (task // NVCH) % NH
         ch = task % NVCH
         k = kda_stage(c, req, h, ch)
+        stamp(on, tls, tid, 7)
         kda_conv(c, k, req, h, ch)
+        stamp(on, tls, tid, 8)
         kda_gates(c, k, req, h)
+        stamp(on, tls, tid, 9)
         kda_recur(c, k, req, h, ch)
+        stamp(on, tls, tid, 10)
         kda_norm(c, req, h, ch)
+        stamp(on, tls, tid, 11)
 
 
 # ---------------------------------------------------------------- kernel
@@ -731,16 +847,17 @@ def build(key: KdaPreBuild):
     ns = key.nblocks + 1
     lay = scratch_layout(key)
     xrow = HIDDEN + LDS_PAD
-    assert s * xrow * 2 <= 136 * 1024, "x rows exceed LDS"
-    assert L <= 16 and key.nblocks >= 1
+    assert 1 <= s <= MAX_S and L <= 16 and key.nblocks >= 1
+    # the mix re-reads blocks 0 .. nblocks - 1: none may be this launch's write
+    assert vals_all(s, ns) or key.write_idx < 0 or key.write_idx >= key.nblocks
 
     @fx.struct
     class AttnLds:
-        vals: fx.Array[fx.Float32, s * ns * COLS, 16]
+        vals: fx.Array[fx.Float32, vals_words(s, ns), 16]
 
     @fx.struct
     class GemvLds:
-        xl: fx.Array[fx.Int32, s * xrow // 2, 16]
+        xl: fx.Array[fx.Int32, min(s, XT) * xrow // 2, 16]
 
     @fx.struct
     class KdaLds:
@@ -760,9 +877,13 @@ class RoleLds:
     @fx.struct
     class Smem:
         red: fx.Array[fx.Float32, WAVES * 64 * 4, 16]
+        tls: fx.Array[fx.Int64, TL_POINTS, 16]
 
     keyed = key_tuple(key, _SOURCES)
-    name = f"k3_mono_kda_pre_s{s}_l{L}_nb{key.nblocks}_d{int(key.delta)}"
+    name = (
+        f"k3_mono_kda_pre_s{s}_l{L}_nb{key.nblocks}_d{int(key.delta)}"
+        f"_t{int(key.timeline)}"
+    )
 
     @flyc.kernel(name=name, known_block_size=[THREADS, 1, 1])
     def kda_pre(
@@ -795,6 +916,7 @@ def kda_pre(
         epoch: Int64,
         health: Int64,
         alarm: Int64,
+        tl: Int64,
     ):
         _ = keyed
         tid = fx.thread_idx.x
@@ -821,6 +943,7 @@ def kda_pre(
             "lb": key.lower_bound,
             "kda0": key.kda0,
             "write_idx": key.write_idx,
+            "vals_all": vals_all(s, ns),
             "put": mb.put,
             "poll": mb.poll,
             "red": lds.red.ptr,
@@ -862,23 +985,46 @@ def kda_pre(
         }
         for region in REGIONS:
             c[region] = sreg(scratch, lay[region][0], region)
+        on, tls = key.timeline, lds.tls.ptr
+        c["on"], c["tls"] = on, tls
+        stamp_begin(on, tls, tid, TL_POINTS)
         if bid < NA:
             stage_slice(c, bid)
-            if bid < s:
-                stage_gate(c, bid)
+            stamp(on, tls, tid, 1)
+            stage_gates(c, bid)
+            stamp(on, tls, tid, 2)
             stage_mix(c, bid)
+            stamp(on, tls, tid, 3)
         if bid >= NA:
             gi = bid - NA
             rg0 = 2 * gi
             rg1 = 2 * gi + 1
             ops0 = gemv_loads(c, rg0, range(PER_WAVE))
-            load_x(c)
-            # the second group's first half in flight under the first's MFMAs
-            ops1 = gemv_loads(c, rg1, range(4))
-            stage_gemv(c, rg0, ops0)
-            ops1 = ops1 + gemv_loads(c, rg1, range(4, PER_WAVE))
-            stage_gemv(c, rg1, ops1)
-        kda_tasks(c)
+            wait_x(c)
+            if const_expr(s <= XT):
+                load_x(c, 0, s)
+                stamp(on, tls, tid, 4)
+                # the second group's first half in flight under the first's MFMAs
+                ops1 = gemv_loads(c, rg1, range(4))
+                stage_gemv(c, rg0, ops0, 0, s)
+                stamp(on, tls, tid, 5)
+                ops1 = ops1 + gemv_loads(c, rg1, range(4, PER_WAVE))
+                stage_gemv(c, rg1, ops1, 0, s)
+                stamp(on, tls, tid, 6)
+            else:
+                # x a tile of XT rows at a time, a group's weights held over them
+                stamp(on, tls, tid, 4)
+                for t0, n in token_tiles(s, XT):
+                    load_x(c, t0, n)
+                    stage_gemv(c, rg0, ops0, t0, n)
+                stamp(on, tls, tid, 5)
+                ops1 = gemv_loads(c, rg1, range(PER_WAVE))
+                for t0, n in token_tiles(s, XT):
+                    load_x(c, t0, n)
+                    stage_gemv(c, rg1, ops1, t0, n)
+                stamp(on, tls, tid, 6)
+        kda_tasks(c, on, tls)
+        stamp_flush(on, tls, tl, tid, bid, TL_POINTS)
 
     @flyc.jit
     def launch(
@@ -911,6 +1057,7 @@ def launch(
         epoch: Int64,
         health: Int64,
         alarm: Int64,
+        tl: Int64,
         stream: fx.Stream = _STREAM,
     ):
         _ = keyed
@@ -944,6 +1091,7 @@ def launch(
             epoch,
             health,
             alarm,
+            tl,
         ).launch(grid=(BLOCKS,), block=(THREADS,), stream=stream)
 
     ABI.check(kda_pre, launch)
@@ -980,6 +1128,7 @@ def kda_pre(
     epoch,
     health,
     alarm,
+    tl=None,
 ):
     """Host launcher. ``conv_state`` [slots, dim, state_len] or a transposed view
     of [slots, state_len, dim] (strides read); ``rstate`` [slots, 12, 128, 128]
@@ -987,7 +1136,8 @@ def kda_pre(
     layer (tags are a layer's); ``epoch`` a device int32 bumped once a step and
     ``layer`` this launch's slot below ``TAG_SLOTS``."""
     f = build(key)
-    f(
+    _run_compiled(
+        f,
         prefix.data_ptr(),
         ptr(delta),
         blocks.data_ptr(),
@@ -1017,5 +1167,6 @@ def kda_pre(
         epoch.data_ptr(),
         health.data_ptr(),
         alarm,
-        stream=torch.cuda.current_stream(),
+        ptr(tl),
+        torch.cuda.current_stream(),
     )
diff --git a/vllm/models/kimi_k3/amd/specialized_decode/mono/attention/mla.py b/vllm/models/kimi_k3/amd/specialized_decode/mono/attention/mla.py
new file mode 100644
index 000000000000..0bfd4454ebcf
--- /dev/null
+++ b/vllm/models/kimi_k3/amd/specialized_decode/mono/attention/mla.py
@@ -0,0 +1,869 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""M1a ``mla_pre``: one Kimi-K3 MLA layer from the AttnRes seam to AITER's
+decode input, for a spec-verify step (S rows) at TP8 + DCP8 (the DCP group the
+TP group: rank r's 12 heads are heads 12 r .. of the gathered query) or TP8 +
+DCP1 (``ranks`` 1: the query is this rank's 12 heads, no gather).
+
+One launch a layer, ``BLOCKS`` CTAs x ``THREADS`` threads:
+
+    attn CTAs 0 .. 55   AttnRes + input_layernorm as K1 (``kda.stage_*``) ->
+                        bf16 x rows, an XRDY flag a slice
+    gemv CTAs 56 ..     fused_qkv_a (132 x 16 rows) and g_proj (96 x 16), task
+                        gi and gi + 200: q_c | kv_c | k_pe -> QA (bf16 pairs);
+                        the gate bf16(sigmoid(bf16 g)) -> ``gate_out``
+    norm (token t on CTA t)
+                        q_a / kv_a RMSNorm: q_c -> plain bf16 rows, a QCR flag
+                        a token; kv_c (normed) | k_pe -> fp8 cache entry when
+                        the token's slot is this rank's (>= 0)
+    q_b (16 rows, CTAs 64 ..)
+                        q_c rows in LDS: q_nope -> QN (bf16 pairs); q_pe -> fp8
+                        -> every rank's QG (DCP1: ``q_out``)
+    W_UK (a wave a (token tile, head, 16 latent rows), CTAs 208 ..)
+                        q_nope -> MXFP4 (aiter's even scale, a lane a 32-block),
+                        one 16x16x128 FP4 MFMA against W_K -> bf16 -> fp8 ->
+                        every rank's QG (DCP1: ``q_out``)
+                        (``qbf16``, DCP1: bf16 -> ``q_out``, M1c's)
+    unpack (every CTA, DCP8)
+                        this rank's QG -> the dense gathered ``q_out``
+                        [S, ranks 12, 576] fp8
+
+Rounding follows the vLLM path where it is elementwise (bf16 hand-offs, the
+fp8 quants, the MXFP4 quant of the W_UK input); split-K sums differ only in
+their order.
+"""
+
+from dataclasses import dataclass
+
+import flydsl.compiler as flyc
+import flydsl.expr as fx
+import torch
+from aiter.ops.flydsl.kernels import buffer_ops as bo
+from flydsl.expr import const_expr, gpu, range_constexpr
+from flydsl.expr.typing import Int32, Int64, T
+
+from vllm.models.kimi_k3.amd.specialized_decode.mono.attention import kda
+from vllm.models.kimi_k3.amd.specialized_decode.mono.attention.kda import (
+    HIDDEN,
+    LDS_PAD,
+    MAX_S,
+    NA,
+    XT,
+    bf16r,
+    ld_bf4,
+    sigmoid,
+    stage_gates,
+    stage_mix,
+    stage_slice,
+    step_tag,
+    vals_all,
+    vals_words,
+)
+from vllm.models.kimi_k3.amd.specialized_decode.mono.common.debug import (
+    region_ids,
+    stamp,
+    stamp_begin,
+    stamp_flush,
+)
+from vllm.models.kimi_k3.amd.specialized_decode.mono.common.mx import (
+    FP4,
+    clamp_fp8,
+    mfma_scaled,
+    pow2,
+)
+from vllm.models.kimi_k3.amd.specialized_decode.mono.common.ops import (
+    CM_DEV,
+    bf16_pair,
+    bf_hi,
+    bf_lo,
+    butterfly,
+    fp8_pack4,
+    hw_rsq,
+    ld_i32,
+    load_ptr64,
+    row_sum,
+    rsrc,
+    traced,
+    uniform,
+)
+from vllm.models.kimi_k3.amd.specialized_decode.mono.common.plan import (
+    BLOCKS,
+    THREADS,
+    WAVES,
+    KernelAbi,
+    first_task,
+    key_tuple,
+    pair_layout,
+    source_digest,
+)
+from vllm.models.kimi_k3.amd.specialized_decode.mono.common.ranks import peer_bases
+from vllm.models.kimi_k3.amd.specialized_decode.mono.common.sync import (
+    Mailbox,
+    preg,
+    publish,
+    sreg,
+)
+from vllm.models.kimi_k3.amd.specialized_decode.mono.stages import gemv
+
+NH = 12  # local heads at TP8
+QL = 1536  # q_lora_rank
+KVL = 512  # kv_lora_rank
+NOPE = 128
+ROPE = 64
+QHD = NOPE + ROPE  # 192
+QKV_A = QL + KVL + ROPE  # 2112: fused_qkv_a's rows
+GROWS = NH * 128  # 1536: g_proj's rows (heads x v_head_dim)
+QB = NH * QHD  # 2304: q_b's rows
+ENTRY = KVL + ROPE  # 576: a query row / cache entry (fp8 bytes)
+QW = ENTRY // 4  # 144 words
+ROWS = 16
+A_TASKS = QKV_A // ROWS  # 132
+GEMV_TASKS = A_TASKS + GROWS // ROWS  # 228
+GEMV_CTAS = BLOCKS - NA  # 200
+QB_SUB = QHD // ROWS  # 12 tasks a head: 8 nope, 4 pe
+QB_TASKS = NH * QB_SUB  # 144
+QB0 = 64  # q_b's first CTA (the norm tasks on CTAs 0 .. S - 1)
+UK_GROUPS = KVL // ROWS  # 32 latent row groups a head
+UK0 = QB0 + QB_TASKS  # W_UK's first CTA (208)
+QCROW = QL + LDS_PAD  # an LDS q_c row
+# timeline points: start, slice, gate, mix, gemv0, gemv1, norm, q_b, W_UK,
+# unpack
+TL_POINTS = 10
+
+assert GEMV_TASKS <= 2 * GEMV_CTAS and QB0 >= MAX_S and UK0 <= BLOCKS
+assert NH * QHD == QB and QL % (4 * 384) == 0 and KVL == 4 * 128
+
+_STREAM = fx.Stream(None)
+REGIONS = ("part", "wgt", "msq", "xrdy", "qa", "qcr", "qn", "qg")
+_SOURCES = source_digest(".")
+
+
+@dataclass(frozen=True)
+class MlaPreBuild:
+    tokens: int  # S
+    nblocks: int  # AttnRes block sources (the prefix is one more)
+    delta: bool
+    eps: float = 1e-5  # AttnRes norm
+    out_eps: float = 1e-5  # input_layernorm
+    q_eps: float = 1e-6  # q_a_layernorm
+    kv_eps: float = 1e-6  # kv_a_layernorm
+    ranks: int = 8  # the DCP (= TP) group; 1: no DCP, no peer memory
+    peer_off: int = 0  # QG's byte offset in a rank's peer buffer
+    timeline: bool = False
+    # q_out's heads a token (DCP8; 0: ranks x 12). More: AITER cprr's padded
+    # query, heads ranks x 12.. tiling heads 0.. as forward_mqa's pad would
+    qheads: int = 0
+    # DCP1: ``q_out`` bf16 [S, 12, 576] (AITER's mla_gluon's query) for fp8
+    qbf16: bool = False
+
+
+ABI = KernelAbi(
+    (
+        "prefix",
+        "delta",
+        "blocks",
+        "blk_sm",
+        "blk_sr",
+        "ares_nw",
+        "ares_qk",
+        "in_nw",
+        "w_qkv",
+        "w_g",
+        "q_nw",
+        "kv_nw",
+        "w_qb",
+        "w_k",
+        "wk_sh",
+        "wk_sn",
+        "w_ks",
+        "ks_sh",
+        "ks_sn",
+        "ks_sk",
+        "kv_cache",
+        "kv_sb",
+        "kv_ss",
+        "kv_bs",
+        "slots",
+        "gate_out",
+        "q_out",
+        "scratch",
+        "peers",
+        "rank",
+        "layer",
+        "epoch",
+        "health",
+        "alarm",
+        "tl",
+    )
+)
+
+
+def scratch_layout(key: MlaPreBuild) -> dict:
+    s, ns = key.tokens, key.nblocks + 1
+    pairs = {
+        "part": NA * s * ns * 2,
+        "wgt": s * (ns + 1),
+        "msq": NA * s,
+        "xrdy": NA,
+        "qa": s * QKV_A // 2,
+        "qcr": s,
+        "qn": s * NH * NOPE // 2,
+    }
+    lay = pair_layout(pairs.items())
+    end = max(o + n for o, n in lay.values())
+    # the plain bf16 rows (not pairs), 16 B aligned: x, then the normed q_c
+    lay["x"] = ((end + 15) // 16 * 16, s * HIDDEN * 2)
+    lay["qc"] = (lay["x"][0] + s * HIDDEN * 2, s * QL * 2)
+    return lay
+
+
+def scratch_bytes(key: MlaPreBuild) -> int:
+    return max(o + n for o, n in scratch_layout(key).values())
+
+
+def q_heads(key: MlaPreBuild) -> int:
+    return key.qheads or key.ranks * NH
+
+
+def peer_pairs(s: int, ranks: int) -> int:
+    """QG: every source rank's [S, 12, 144] fp8 words, a pair each."""
+    return ranks * s * NH * QW
+
+
+def peer_bytes(s: int, ranks: int) -> int:
+    return peer_pairs(s, ranks) * 8
+
+
+def qg_pair(c, src, t, h, w):
+    return ((src * c["S"] + t) * NH + h) * QW + w
+
+
+def put_qg(c, t, h, w, word):
+    """fp8 word ``w`` of (token t, local head h) -> every rank's QG; at DCP1
+    straight to ``q_out`` (nothing in this launch reads it)."""
+    if const_expr(c["R"] == 1):
+        bo.buffer_store(word, rsrc(c["args"]["q_out"]), (t * NH + h) * QW + w)
+    if const_expr(c["R"] > 1):
+        i = qg_pair(c, c["rank"], t, h, w)
+        for p in range_constexpr(c["R"]):
+            c["put_words"](c["qg"][p], i, [word])
+
+
+def put_q(c, t, h, w, v):
+    """Query values ``4 w ..`` (four f32, bf16-exact) of (token t, local head h):
+    fp8 -> QG / ``q_out``, or with ``qbf16`` bf16 pairs -> ``q_out``."""
+    if const_expr(c["qbf16"]):
+        pair = fx.Vector.from_elements(
+            [bf16_pair(v[0], v[1]), bf16_pair(v[2], v[3])], fx.Float32
+        )
+        bo.buffer_store(pair, rsrc(c["args"]["q_out"]), ((t * NH + h) * QW + w) * 2)
+    if const_expr(not c["qbf16"]):
+        put_qg(c, t, h, w, fp8_word(v))
+
+
+def fp8_word(v):
+    """Four f32 -> four E4M3 (vLLM's saturating quant, scale 1)."""
+    return fp8_pack4(*[clamp_fp8(x) for x in v])
+
+
+def bf_words(ptr, i, n):
+    """``n`` dwords of a bf16 row at word ``i`` (device scope)."""
+    return fx.Vector(
+        bo.buffer_load(rsrc(ptr), i, vec_width=n, dtype=T.i32, cache_modifier=CM_DEV)
+    )
+
+
+# ---------------------------------------------------------------- gemv CTAs
+def proj_loads(c, task):
+    """Task ``task``'s 16 weight rows (fused_qkv_a, then g_proj) for this wave;
+    a task past the last loads the last."""
+    a = c["args"]
+    task = fx.min(task, GEMV_TASKS - 1)
+    is_g = uniform((task >= A_TASKS).select(fx.Int32(1), fx.Int32(0)))
+    w = a["w_qkv"] + fx.Int64(is_g) * (a["w_g"] - a["w_qkv"])
+    row0 = (task - is_g * A_TASKS) * ROWS
+    return gemv.loads(c["lane"], c["wave"], w, row0, HIDDEN)
+
+
+@traced
+def stage_proj(c, task, ops, t0, n):
+    """Task ``task`` against the LDS x (tokens t0 .. t0 + n): q_c | kv_c | k_pe
+    rows -> QA, g rows -> the gate."""
+    tid, lane, wave, red = c["tid"], c["lane"], c["wave"], c["red"]
+    acc = gemv.mfmas(lane, wave, c["xl"], HIDDEN + LDS_PAD, HIDDEN, n, ops)
+    fx.ptr_store(acc, red + (wave * 64 + lane) * 4)
+    gpu.barrier()
+    r2 = tid % 8
+    tt = tid // 8
+    live = (tid < 8 * n) & (task < GEMV_TASKS)
+    v0 = bf16r(row_sum(red, 2 * r2, fx.min(tt, n - 1)))
+    v1 = bf16r(row_sum(red, 2 * r2 + 1, fx.min(tt, n - 1)))
+    t = t0 + tt
+    if live & (task < A_TASKS):
+        c["put_bf"](c["qa"], t * QKV_A + task * ROWS + 2 * r2, [v0, v1])
+    if live & (task >= A_TASKS):
+        col = (task - A_TASKS) * ROWS + 2 * r2
+        g = bf16_pair(sigmoid(v0), sigmoid(v1))
+        bo.buffer_store(g, rsrc(c["args"]["gate_out"]), (t * GROWS + col) // 2)
+    gpu.barrier()
+
+
+@traced
+def gemv_tasks(c, gi):
+    """Gemv CTA ``gi``: task gi (weights issued before x lands), then gi + 200."""
+    s, on, tls, tid = c["S"], c["on"], c["tls"], c["tid"]
+    ops0 = proj_loads(c, gi)
+    kda.wait_x(c)
+    for t0, n in gemv.token_tiles(s, XT):
+        kda.load_x(c, t0, n)
+        stage_proj(c, gi, ops0, t0, n)
+    stamp(on, tls, tid, 4)
+    if gi + GEMV_CTAS < GEMV_TASKS:
+        ops1 = proj_loads(c, gi + GEMV_CTAS)
+        for t0, n in gemv.token_tiles(s, XT):
+            if const_expr(s > XT):
+                kda.load_x(c, t0, n)
+            stage_proj(c, gi + GEMV_CTAS, ops1, t0, n)
+    stamp(on, tls, tid, 5)
+
+
+# ---------------------------------------------------------------- norm
+@traced
+def block_sum2(c, v):
+    """Waves 0 .. 5's sum and waves 6, 7's sum of ``v`` (q / kv threads)."""
+    tid, lane, wave, red = c["tid"], c["lane"], c["wave"], c["red"]
+    v = butterfly(v, (32, 16, 8, 4, 2, 1))
+    gpu.barrier()
+    if lane == 0:
+        fx.ptr_store(v, red + wave)
+    gpu.barrier()
+    q = fx.ptr_load(red + 0)
+    for w in range_constexpr(1, 6):
+        q = q + fx.ptr_load(red + w)
+    kv = fx.ptr_load(red + 6) + fx.ptr_load(red + 7)
+    _ = tid
+    return q, kv
+
+
+@traced
+def stage_norm(c, t):
+    """Token t: q_c (threads 0 .. 383, 4 columns each) and kv_c (384 ..) RMSNorm;
+    q_c -> plain bf16 row + QCR; kv_c | k_pe (threads 0 .. 15) -> fp8 entry."""
+    tid, a = c["tid"], c["args"]
+    is_q = tid < QL // 4
+    col = is_q.select(tid * 4, QL + (tid - QL // 4) * 4)
+    pe = QL + KVL + fx.min(tid, ROPE // 4 - 1) * 4
+    got = c["poll"](
+        [(c["qa"], (t * QKV_A + col) // 2, 2), (c["qa"], (t * QKV_A + pe) // 2, 2)]
+    )
+    w = got[0]
+    v = [bf_lo(w[0]), bf_hi(w[0]), bf_lo(w[1]), bf_hi(w[1])]
+    sq = v[0] * v[0] + v[1] * v[1] + v[2] * v[2] + v[3] * v[3]
+    qs, kvs = block_sum2(c, sq)
+    rq = hw_rsq(qs * fx.Float32(1.0 / QL) + fx.Float32(c["q_eps"]))
+    rkv = hw_rsq(kvs * fx.Float32(1.0 / KVL) + fx.Float32(c["kv_eps"]))
+    nw = ld_bf4(a["q_nw"], fx.min(tid, QL // 4 - 1) * 4)
+    kw = ld_bf4(a["kv_nw"], fx.max(tid - QL // 4, 0) * 4)
+    qn = [bf16r(v[j] * rq * nw[j]) for j in range(4)]
+    kn = [bf16r(v[j] * rkv * kw[j]) for j in range(4)]
+    if is_q:
+        bo.buffer_store(
+            kda.bf4_words(qn),
+            rsrc(c["qc"]),
+            (t * QL + tid * 4) // 2,
+            cache_modifier=CM_DEV,
+        )
+    slot = uniform(ld_i32(a["slots"], t * 2))
+    if slot >= 0:
+        blk = slot // a["kv_bs"]
+        off = slot % a["kv_bs"]
+        base = (
+            a["kv_cache"]
+            + fx.Int64(blk) * fx.Int64(a["kv_sb"])
+            + fx.Int64(off) * fx.Int64(a["kv_ss"])
+        )
+        if tid >= QL // 4:
+            bo.buffer_store(fp8_word(kn), rsrc(base), tid - QL // 4)
+        if tid < ROPE // 4:
+            p = got[1]
+            pv = [bf_lo(p[0]), bf_hi(p[0]), bf_lo(p[1]), bf_hi(p[1])]
+            bo.buffer_store(fp8_word(pv), rsrc(base), KVL // 4 + tid)
+    publish(c["put"], c["qcr"], t, fx.Int32(1), tid == 0)
+    gpu.barrier()
+
+
+# ---------------------------------------------------------------- q_b
+@traced
+def stage_qb(c, task):
+    """q_b rows 16 task ..: head h = task / 12, its nope rows (sub < 8) -> QN or
+    pe rows -> fp8 -> QG."""
+    s, tid, lane, wave, red = c["S"], c["tid"], c["lane"], c["wave"], c["red"]
+    a = c["args"]
+    h = task // QB_SUB
+    sub = task % QB_SUB
+    ops = gemv.loads(lane, wave, a["w_qb"], task * ROWS, QL)
+    if tid < s:
+        c["poll"]([(c["qcr"], tid, 1)])
+    gpu.barrier()
+    for t0, n in gemv.token_tiles(s):
+        gemv.dev_rows_to_lds(tid, c["qc"], QL, n, c["xl"], QCROW, t0=t0)
+        acc = gemv.mfmas(lane, wave, c["xl"], QCROW, QL, n, ops)
+        fx.ptr_store(acc, red + (wave * 64 + lane) * 4)
+        gpu.barrier()
+        r2 = tid % 8
+        tt = fx.min(tid // 8, n - 1)
+        if (tid < 8 * n) & (sub < NOPE // ROWS):
+            v0 = row_sum(red, 2 * r2, tt)
+            v1 = row_sum(red, 2 * r2 + 1, tt)
+            at = ((t0 + tt) * NH + h) * NOPE + sub * ROWS + 2 * r2
+            c["put_bf"](c["qn"], at, [v0, v1])
+        r4 = tid % 4
+        t4 = fx.min(tid // 4, n - 1)
+        if (tid < 4 * n) & (sub >= NOPE // ROWS):
+            v = [bf16r(row_sum(red, 4 * r4 + j, t4)) for j in range(4)]
+            w = KVL // 4 + (sub - NOPE // ROWS) * 4 + r4
+            put_q(c, t0 + t4, h, w, v)
+        gpu.barrier()
+
+
+# ---------------------------------------------------------------- W_UK
+def even_code(am):
+    """Aiter's ``_mxfp4_scale_from_amax`` code from |bf16| bits ``am``: the
+    exponent of amax rounded at mantissa bit 21, minus 2, clamped at 0."""
+    e = (((am << 16) + fx.Int32(0x200000)) >> 23) & 0xFF
+    return fx.max(e - 2, fx.Int32(0))
+
+
+def pk_fp4_bf16(old, word, scale, j):
+    from vllm.models.kimi_k3.amd.specialized_decode.mono.stages.moe import (
+        pk_fp4_bf16 as pk,
+    )
+
+    return pk(old, word, scale, j)
+
+
+@traced
+def stage_uk(c, task):
+    """Wave task (tile, head, 16 latent rows): W_K's rows against the 16 tokens'
+    MXFP4 q_nope -> bf16 -> fp8 -> QG (a lane 4 rows of one token)."""
+    s, lane = c["S"], c["lane"]
+    a = c["args"]
+    ng = task % UK_GROUPS
+    h = (task // UK_GROUPS) % NH
+    t0 = task // (UK_GROUPS * NH) * 16
+    n = ng * ROWS + lane % 16
+    kb = lane // 16
+    wo = h * a["wk_sh"] + n * a["wk_sn"] + kb * 16
+    wv = fx.Vector(bo.buffer_load(rsrc(a["w_k"]), wo // 4, vec_width=4, dtype=T.i32))
+    so = h * a["ks_sh"] + n * a["ks_sn"] + kb * a["ks_sk"]
+    sw = fx.Int32(bo.buffer_load(rsrc(a["w_ks"]), so // 4, vec_width=1, dtype=T.i32))
+    sa = (sw >> ((so % 4) * 8)) & 0xFF
+    t = t0 + lane % 16
+    tc = fx.min(t, s - 1)
+    p0 = (((tc * NH + h) * NOPE) + kb * 32) // 2
+    got = c["poll"]([(c["qn"], p0 + 2 * j, 2) for j in range(8)])
+    ws = [got[j // 2][j % 2] for j in range(16)]
+    am = fx.Int32(0)
+    for w in ws:
+        am = fx.max(am, w & 0x7FFF)
+        am = fx.max(am, (w >> 16) & 0x7FFF)
+    # code 0 (amax < 2^-125) quantizes by 2^-126: the tiny lane stays ~0
+    code = fx.max(even_code(am), fx.Int32(1))
+    scale = pow2(code)
+    xs = []
+    for q in range_constexpr(4):
+        pk = fx.Int32(0)
+        for j in range_constexpr(4):
+            pk = pk_fp4_bf16(pk, ws[4 * q + j], scale, j)
+        xs.append(pk)
+    xv = fx.Vector.from_elements(xs, fx.Int32)
+    acc = fx.Vector.filled(4, 0.0, fx.Float32)
+    acc = mfma_scaled(wv, xv, acc, sa, code, FP4, FP4, 0, 0)
+    v = [bf16r(acc[j]) for j in range(4)]
+    if t < s:
+        put_q(c, t, h, ng * 4 + kb, v)
+
+
+@traced
+def uk_tasks(c):
+    s, bid, wave = c["S"], c["bid"], c["wave"]
+    ntask = ((s + 15) // 16) * NH * UK_GROUPS
+    first = first_task(bid, UK0) * WAVES + wave
+    for task in range(first, ntask, BLOCKS * WAVES):
+        stage_uk(c, task)
+
+
+# ---------------------------------------------------------------- unpack
+@traced
+def stage_unpack(c):
+    """This rank's QG -> ``q_out`` [S, QH, 576]: a thread 2 words a round (QG
+    pair 2 u, u = global thread + round x grid). Gathered head g also lands
+    at R x 12 + g while that is a pad head (g < QH - R x 12)."""
+    s, R, tid, bid, QH = c["S"], c["R"], c["tid"], c["bid"], c["QH"]
+    pad = QH - R * NH
+    nu = R * s * NH * QW // 2
+    grid = BLOCKS * THREADS
+    rounds = (nu + grid - 1) // grid
+    g = bid * THREADS + tid
+    us = [fx.min(g + grid * i, nu - 1) for i in range(rounds)]
+    got = c["poll"]([(c["qg_own"], 2 * u, 2) for u in us])
+    for i in range_constexpr(rounds):
+        u = us[i]
+        w2 = u % (QW // 2)
+        h = (u // (QW // 2)) % NH
+        t = (u // (NH * QW // 2)) % s
+        src = u // (s * NH * QW // 2)
+        gh = src * NH + h
+        at = (t * QH + gh) * QW + 2 * w2
+        if g + grid * i < nu:
+            word = fx.Vector.from_elements(got[i], fx.Int32)
+            bo.buffer_store(word, rsrc(c["args"]["q_out"]), at)
+            if const_expr(pad > 0):  # noqa: SIM102 (compile-time, then traced)
+                if gh < pad:
+                    bo.buffer_store(word, rsrc(c["args"]["q_out"]), at + R * NH * QW)
+
+
+# ---------------------------------------------------------------- kernel
+_BUILDS: dict = {}
+
+
+def build(key: MlaPreBuild):
+    if key in _BUILDS:
+        return _BUILDS[key]
+    s, R = key.tokens, key.ranks
+    ns = key.nblocks + 1
+    lay = scratch_layout(key)
+    xrow = HIDDEN + LDS_PAD
+    assert 1 <= s <= MAX_S and key.nblocks >= 1 and 1 <= R <= 8
+    qh = q_heads(key)
+    assert R * NH <= qh <= 2 * R * NH and (R > 1 or qh == NH)
+    assert not key.qbf16 or R == 1
+
+    @fx.struct
+    class AttnLds:
+        vals: fx.Array[fx.Float32, vals_words(s, ns), 16]
+
+    @fx.struct
+    class GemvLds:
+        xl: fx.Array[fx.Int32, max(min(s, XT) * xrow, 16 * QCROW) // 2, 16]
+
+    @fx.union
+    class RoleLds:
+        attn: AttnLds
+        gemv: GemvLds
+
+    @fx.struct
+    class Smem:
+        red: fx.Array[fx.Float32, WAVES * 64 * 4, 16]
+        tls: fx.Array[fx.Int64, TL_POINTS, 16]
+
+    keyed = key_tuple(key, _SOURCES)
+    name = (
+        f"k3_mono_mla_pre_s{s}_nb{key.nblocks}_d{int(key.delta)}_r{R}"
+        f"_t{int(key.timeline)}_qh{qh}{'_qbf' if key.qbf16 else ''}"
+    )
+
+    @flyc.kernel(name=name, known_block_size=[THREADS, 1, 1])
+    def mla_pre(
+        prefix: Int64,
+        delta: Int64,
+        blocks: Int64,
+        blk_sm: Int32,
+        blk_sr: Int32,
+        ares_nw: Int64,
+        ares_qk: Int64,
+        in_nw: Int64,
+        w_qkv: Int64,
+        w_g: Int64,
+        q_nw: Int64,
+        kv_nw: Int64,
+        w_qb: Int64,
+        w_k: Int64,
+        wk_sh: Int32,
+        wk_sn: Int32,
+        w_ks: Int64,
+        ks_sh: Int32,
+        ks_sn: Int32,
+        ks_sk: Int32,
+        kv_cache: Int64,
+        kv_sb: Int64,
+        kv_ss: Int32,
+        kv_bs: Int32,
+        slots: Int64,
+        gate_out: Int64,
+        q_out: Int64,
+        scratch: Int64,
+        peers: Int64,
+        rank: Int32,
+        layer: Int32,
+        epoch: Int64,
+        health: Int64,
+        alarm: Int64,
+        tl: Int64,
+    ):
+        _ = keyed
+        tid = fx.thread_idx.x
+        bid = fx.block_idx.x
+        alloc = fx.SharedAllocator()
+        lds = alloc.allocate(Smem).peek()
+        role = alloc.allocate(RoleLds)
+        al, gl = role.attn.peek(), role.gemv.peek()
+        mb = Mailbox(step_tag(epoch, layer) - 1, health, alarm, region_ids(REGIONS))
+        # DCP1 reads no peer memory (``peers`` may be 0)
+        bases = peer_bases(peers, R) if R > 1 else []
+        own = load_ptr64(peers, rank) if R > 1 else peers
+        c = {
+            "S": s,
+            "R": R,
+            "QH": qh,
+            "qbf16": key.qbf16,
+            "NS": ns,
+            "NB": key.nblocks,
+            "tid": tid,
+            "bid": bid,
+            "lane": tid % 64,
+            "wave": tid // 64,
+            "rank": rank,
+            "delta": key.delta,
+            "eps": key.eps,
+            "out_eps": key.out_eps,
+            "q_eps": key.q_eps,
+            "kv_eps": key.kv_eps,
+            "vals_all": vals_all(s, ns),
+            "put": mb.put,
+            "put_bf": mb.put_bf,
+            "put_words": mb.put_words,
+            "poll": mb.poll,
+            "red": lds.red.ptr,
+            "vals": al.vals.ptr,
+            "xl": gl.xl.ptr,
+            "xbuf": scratch + fx.Int64(lay["x"][0]),
+            "qc": scratch + fx.Int64(lay["qc"][0]),
+            "qg": [preg(b, key.peer_off, "qg") for b in bases],
+            "qg_own": preg(own, key.peer_off, "qg"),
+            "args": {
+                "prefix": prefix,
+                "delta": delta,
+                "blocks": blocks,
+                "blk_sm": blk_sm,
+                "blk_sr": blk_sr,
+                "ares_nw": ares_nw,
+                "ares_qk": ares_qk,
+                "in_nw": in_nw,
+                "w_qkv": w_qkv,
+                "w_g": w_g,
+                "q_nw": q_nw,
+                "kv_nw": kv_nw,
+                "w_qb": w_qb,
+                "w_k": w_k,
+                "wk_sh": wk_sh,
+                "wk_sn": wk_sn,
+                "w_ks": w_ks,
+                "ks_sh": ks_sh,
+                "ks_sn": ks_sn,
+                "ks_sk": ks_sk,
+                "kv_cache": kv_cache,
+                "kv_sb": kv_sb,
+                "kv_ss": kv_ss,
+                "kv_bs": kv_bs,
+                "slots": slots,
+                "gate_out": gate_out,
+                "q_out": q_out,
+            },
+        }
+        for region in REGIONS:
+            if const_expr(region != "qg"):
+                c[region] = sreg(scratch, lay[region][0], region)
+        on, tls = key.timeline, lds.tls.ptr
+        c["on"], c["tls"] = on, tls
+        stamp_begin(on, tls, tid, TL_POINTS)
+        if bid < NA:
+            stage_slice(c, bid)
+            stamp(on, tls, tid, 1)
+            stage_gates(c, bid)
+            stamp(on, tls, tid, 2)
+            stage_mix(c, bid)
+            stamp(on, tls, tid, 3)
+        if bid >= NA:
+            gemv_tasks(c, bid - NA)
+        for t in range(first_task(bid, 0), s, BLOCKS):
+            stage_norm(c, t)
+        stamp(on, tls, tid, 6)
+        for task in range(first_task(bid, QB0), QB_TASKS, BLOCKS):
+            stage_qb(c, task)
+        stamp(on, tls, tid, 7)
+        uk_tasks(c)
+        stamp(on, tls, tid, 8)
+        if const_expr(R > 1):
+            stage_unpack(c)
+        stamp(on, tls, tid, 9)
+        stamp_flush(on, tls, tl, tid, bid, TL_POINTS)
+
+    @flyc.jit
+    def launch(
+        prefix: Int64,
+        delta: Int64,
+        blocks: Int64,
+        blk_sm: Int32,
+        blk_sr: Int32,
+        ares_nw: Int64,
+        ares_qk: Int64,
+        in_nw: Int64,
+        w_qkv: Int64,
+        w_g: Int64,
+        q_nw: Int64,
+        kv_nw: Int64,
+        w_qb: Int64,
+        w_k: Int64,
+        wk_sh: Int32,
+        wk_sn: Int32,
+        w_ks: Int64,
+        ks_sh: Int32,
+        ks_sn: Int32,
+        ks_sk: Int32,
+        kv_cache: Int64,
+        kv_sb: Int64,
+        kv_ss: Int32,
+        kv_bs: Int32,
+        slots: Int64,
+        gate_out: Int64,
+        q_out: Int64,
+        scratch: Int64,
+        peers: Int64,
+        rank: Int32,
+        layer: Int32,
+        epoch: Int64,
+        health: Int64,
+        alarm: Int64,
+        tl: Int64,
+        stream: fx.Stream = _STREAM,
+    ):
+        _ = keyed
+        mla_pre(
+            prefix,
+            delta,
+            blocks,
+            blk_sm,
+            blk_sr,
+            ares_nw,
+            ares_qk,
+            in_nw,
+            w_qkv,
+            w_g,
+            q_nw,
+            kv_nw,
+            w_qb,
+            w_k,
+            wk_sh,
+            wk_sn,
+            w_ks,
+            ks_sh,
+            ks_sn,
+            ks_sk,
+            kv_cache,
+            kv_sb,
+            kv_ss,
+            kv_bs,
+            slots,
+            gate_out,
+            q_out,
+            scratch,
+            peers,
+            rank,
+            layer,
+            epoch,
+            health,
+            alarm,
+            tl,
+        ).launch(grid=(BLOCKS,), block=(THREADS,), stream=stream)
+
+    ABI.check(mla_pre, launch)
+    _BUILDS[key] = launch
+    return launch
+
+
+def ptr(t):
+    return 0 if t is None else t.data_ptr()
+
+
+def mla_pre(
+    key: MlaPreBuild,
+    *,
+    prefix,
+    delta,
+    blocks,
+    ares_nw,
+    ares_qk,
+    in_nw,
+    w_qkv,
+    w_g,
+    q_nw,
+    kv_nw,
+    w_qb,
+    w_k,
+    w_k_scale,
+    kv_cache,
+    slots,
+    gate_out,
+    q_out,
+    scratch,
+    peers,
+    rank,
+    layer,
+    epoch,
+    health,
+    alarm,
+    tl=None,
+):
+    """Host launcher. ``w_k`` (12, 512, 64) u8 MXFP4 and ``w_k_scale`` (12, 512,
+    4) e8m0, any strides but W_K's last; ``kv_cache`` [blocks, block_size, 576]
+    one-byte, entries dense; ``slots`` [S] int64 (-1: another rank's token);
+    ``gate_out`` [S, 1536] bf16; ``q_out`` [S, q_heads(key), 576] fp8 (dense; bf16 with
+    ``qbf16``);
+    ``peers`` every rank's peer buffer base (i64 a rank; unread at ranks 1,
+    may be None); ``layer`` a slot below
+    ``kda.TAG_SLOTS``."""
+    assert w_k.stride(-1) == 1 and kv_cache.stride(-1) == 1
+    assert slots.dtype == torch.int64 and slots.size(0) == key.tokens
+    assert q_out.is_contiguous() and gate_out.is_contiguous()
+    assert q_out.shape[1:] == (q_heads(key), ENTRY), q_out.shape
+    assert q_out.element_size() == (2 if key.qbf16 else 1), q_out.dtype
+    f = build(key)
+    f(
+        prefix.data_ptr(),
+        ptr(delta),
+        blocks.data_ptr(),
+        blocks.stride(0),
+        blocks.stride(1),
+        ares_nw.data_ptr(),
+        ares_qk.data_ptr(),
+        in_nw.data_ptr(),
+        w_qkv.data_ptr(),
+        w_g.data_ptr(),
+        q_nw.data_ptr(),
+        kv_nw.data_ptr(),
+        w_qb.data_ptr(),
+        w_k.data_ptr(),
+        w_k.stride(0),
+        w_k.stride(1),
+        w_k_scale.data_ptr(),
+        w_k_scale.stride(0),
+        w_k_scale.stride(1),
+        w_k_scale.stride(2),
+        kv_cache.data_ptr(),
+        kv_cache.stride(0),
+        kv_cache.stride(1),
+        kv_cache.size(1),
+        slots.data_ptr(),
+        gate_out.data_ptr(),
+        q_out.data_ptr(),
+        scratch.data_ptr(),
+        ptr(peers),
+        rank,
+        layer,
+        epoch.data_ptr(),
+        health.data_ptr(),
+        alarm,
+        ptr(tl),
+        stream=torch.cuda.current_stream(),
+    )
diff --git a/vllm/models/kimi_k3/amd/specialized_decode/mono/attention/mla_attn.py b/vllm/models/kimi_k3/amd/specialized_decode/mono/attention/mla_attn.py
new file mode 100644
index 000000000000..c2a6078bb0f4
--- /dev/null
+++ b/vllm/models/kimi_k3/amd/specialized_decode/mono/attention/mla_attn.py
@@ -0,0 +1,1424 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""M1c ``mla_attn``: Kimi-K3's MLA decode attention at TP8 + DCP1, the bf16
+query (M1a ``qbf16``) against the fp8 latent cache, split over the KV: one pass
+a split, every split's partial (bf16 o, f32 LSE) for a combine (M1b's).
+
+One launch, ``BLOCKS`` CTAs x ``THREADS`` threads. A CTA a (head group, split):
+16 query rows (``16 / qlen`` heads x qlen positions) against a chunk of the
+request's KV. The split's CTAs (every head group) share an XCD, so a chunk
+streams from HBM once into that XCD's L2.
+
+    block of 128 tokens (wave w its tile of 16, the waves in step), its K
+    rows DMA'd into one of two LDS slots a block ahead (global_load_lds, the
+    chunks spread over the block before: ``DMA``; this wave's chunks waited
+    for one at a time, so the top barrier only frees the slot)
+        K tile LDS -> registers (a lane 32 B a line of a token: the MFMA's A, the K
+        order permuted alike in Q's B), S^T = K Q^T: the fp8 K as is against
+        Q's exact fp8 hi + lo (``q_split``), 9 16x16x128 MFMAs a row tile
+        the block's row max over the waves (LDS) -> P = exp2(S - m) 2^PLOG,
+        fp8 hi + lo -> LDS; the running O, l rescaled (online softmax)
+        O^T += V^T P^T: wave w its 64 latents, the fp8 V^T straight from the
+        slot (``ds_read_b64_tr_b8``) against P's hi and lo, 4 16x16x128 MFMAs
+        a row tile and part
+    end: l summed over the waves -> o = O / l (bf16), LSE = m + log l
+
+A split with no token of the request (or none before its causal bound) leaves
+LSE -inf and o 0. Rows are causal: position i of the request's qlen sees its
+KV up to ``seq_len - qlen + i``.
+"""
+
+from dataclasses import dataclass, replace
+
+import flydsl.compiler as flyc
+import flydsl.expr as fx
+import torch
+from aiter.ops.flydsl.kernels import buffer_ops as bo
+from flydsl._mlir import ir
+from flydsl._mlir.dialects import llvm
+from flydsl.expr import const_expr, gpu, range_constexpr, rocdl
+from flydsl.expr import math as fmath
+from flydsl.expr.typing import Int32, Int64, T
+
+from vllm.models.kimi_k3.amd.specialized_decode.mono.attention import mla
+from vllm.models.kimi_k3.amd.specialized_decode.mono.attention.kda import LOG2E
+from vllm.models.kimi_k3.amd.specialized_decode.mono.attention.mla import (
+    ENTRY,
+    KVL,
+)
+from vllm.models.kimi_k3.amd.specialized_decode.mono.attention.mla_attn_geom import (
+    NH,
+    RTILES,
+    TILE,
+    XCDS,
+    geometry,
+    splits,
+)
+from vllm.models.kimi_k3.amd.specialized_decode.mono.common.ops import (
+    _hw_f32,
+    as_ir_value,
+    bf16_pair,
+    butterfly,
+    hw_exp2,
+    hw_rcp,
+    ld_i32,
+    memrealtime,
+    permlane_swap,
+    rsrc,
+    traced,
+    uniform,
+)
+from vllm.models.kimi_k3.amd.specialized_decode.mono.common.plan import (
+    BLOCKS,
+    THREADS,
+    WAVES,
+    key_tuple,
+    source_digest,
+)
+
+assert NH == mla.NH
+BLK = TILE * WAVES  # tokens a block
+CHUNK = 64  # bytes a token a K chunk (a DMA: 64 lanes x 16 B a tile)
+KLOADS = ENTRY // CHUNK  # 9 DMAs a tile: 8 of 8 tokens' whole 128 B lines (a
+#   token's 8 lanes, a quarter the addresses of 4 B a lane), the tail's 16 x 64
+QSTEPS = KLOADS // 2  # 4 fp8 MFMA steps of QK a Q part (a lane 32 B: K
+#   line s, chunks 2 s, 2 s + 1), then the tail (chunk 8 twice against its
+#   hi | lo)
+TILEW = TILE * CHUNK // 4  # 256 dwords: a tile's chunk p, one DMA
+SLOT = KLOADS * WAVES * TILEW  # a block's K (dwords): [chunk p][tile][TILEW]
+LINE = 2 * CHUNK  # 128 B: chunks 2 s, 2 s + 1 hold line s (bytes 128 s ..) of
+#   the tile's tokens 0..7, 8..15 (a DMA 8 whole lines, ``lines_at``); chunk 8
+#   (the tail, 64 B a token) its 16 tokens (``kpos``)
+VCOLS = KVL // WAVES  # 64 latents a wave
+LN2 = 0.6931471805599453
+NEG = -1.0e30  # the running max before any token (finite: no inf - inf)
+
+assert ENTRY % CHUNK == 0 and KVL % CHUNK == 0 and KVL // 4 % 16 == 0
+assert VCOLS == 64 and THREADS == WAVES * 64 and KVL // CHUNK == WAVES
+
+# Block b + 1's K chunks go out spread over block b (one at a time): issued
+# together they hold the waves at the issue for the whole block's DMA (~1 us,
+# the CU's load path), and issued late their landing shows at the next wait.
+# Hook -> chunks: "q<s>" after QK step s (``QSTEPS``: the tail), "p" after
+# the P barrier, "v<i>" before latent tile i's V^T reads, "m<k>" before PV's
+# MFMAs of (row tile k // 2, P part k % 2); a schedule by row tiles (a
+# spread over QK's first steps, P and PV within ~2%; bunched 5-20% slower).
+# Among PV's MFMAs is the ~2% faster, but in the fused K2 builds at 2 row
+# tiles it leaves a Q part (QK's first B) spilled: a scratch reload a block,
+# whose vmcnt(0) waits out every DMA in flight.
+# At 1 row tile (the block's K the longer wait) two each after QK's first
+# steps: 2-3% faster than one.
+DMA = {
+    1: {"q0": 2, "q1": 2, "q2": 1, "q3": 1, "p": 1, "m0": 2},
+    2: {"q0": 1, "q1": 1, "q2": 1, "q3": 1, "p": 2, "v0": 1, "v1": 1, "v2": 1},
+}
+assert all(sum(d.values()) == KLOADS for d in DMA.values())
+
+# P to the PV MFMA as two fp8 terms (hi + lo, of p 2^PLOG: p <= 1, hi's
+# rounding error in lo's 3 more bits, ~bf16 overall), against the slot's raw
+# fp8 V: fp8 MFMAs, no V conversion
+PLOG = 8  # P 2^PLOG up (p <= 1: <= 256 of E4M3's 448), lo normal to p 2^-10
+
+# Warp specialization (the row tiles in ``WSPEC``): a SIMD's two waves on the
+# two halves of a block, software-pipelined. Waves 0..QKW-1 QK and softmax of
+# block b (wave w tiles 2 w, 2 w + 1, both row tiles), waves QKW.. PV of block
+# b - 1 (wave QKW + v latents 128 v .., V^T in registers, read a block ahead):
+# one wave's softmax VALU under the other's MFMAs. A step two barriers: B1 the
+# max exchange (the PV waves amid their MFMAs, P(b - 1) read), B2 P(b) and
+# alpha(b) written, V^T(b) read and block b + 1 landed. P in ``xp`` (read by
+# B1(b + 1), written after it). Every wave DMAs its tile (as ``k_dma``) of
+# block b + 2 into block b's slot, chunks ..WS_K1 after B1(b) (its V^T read
+# before it, its QK done), the rest early in step b + 1 (the DMAs over the
+# step, not a burst a step), then the indices of block b + 3
+WSPEC = (2,)
+QKW = WAVES // 2
+WS_B1 = 4  # PV's latent tiles (of 8) before B1
+WS_K1 = 3  # a tile's chunks DMA'd after B1 two blocks ahead (the rest one)
+
+
+_STREAM = fx.Stream(None)
+_SOURCES = source_digest(".")
+
+
+@dataclass(frozen=True)
+class MlaAttnBuild:
+    tokens: int  # S = requests x qlen
+    qlen: int  # 4 or 8
+    stamps: bool = False  # debug: phase timestamps (memrealtime) over split 0
+    #   of part_o: 2 + 4 b + (wait done, QK + max barrier, P barrier, PV end)
+    dbg: int = 0  # debug: 1 no K DMA in the loop (stale slots), 2 no part_o stores
+    rt: int = 0  # row tiles a CTA (0: ``RTILES``)
+    wide: bool = False  # a KV pool over 4 GiB (``wide_kv``): 64-bit DMA
+    #   addresses (else a 32-bit offset off the pool's SGPR base)
+    ws: int = -1  # warp specialization (-1: by ``WSPEC``)
+
+
+def wspec(rt: int, ws: int = -1) -> bool:
+    return rt in WSPEC if ws < 0 else bool(ws)
+
+
+@traced
+def stamp(c, slot):
+    # a trace-time if around a device one (not one condition)
+    if const_expr(c["stamps"]):  # noqa: SIM102
+        if (c["wave"] == 0) & (c["lane"] == 0):
+            bo.buffer_store(
+                fx.Int32(memrealtime() & 0x7FFFFFFF),
+                rsrc(c["args"]["part_o"]),
+                c["bid"] * 64 + fx.min(slot, 63),
+            )
+
+
+def hw_log2(x):
+    return _hw_f32("llvm.amdgcn.log.f32", x)
+
+
+def mfma_f8(a, b, c):
+    """16x16x128 E4M3 x E4M3 MFMA, unscaled (scale operands 0: LLVM's
+    v_mfma_f32_16x16x128_f8f6f4): ``a`` / ``b`` 8 dwords (QK and PV)."""
+    zero = as_ir_value(fx.Int32(0))
+    return fx.Vector(
+        rocdl.mfma_scale_f32_16x16x128_f8f6f4(
+            T.vec(4, T.f32),
+            [as_ir_value(a), as_ir_value(b), as_ir_value(c), 0, 0, 0, zero, 0, zero],
+        )
+    )
+
+
+def fp8_unpack4(w):
+    """One dword of four E4M3 -> four f32 (exact)."""
+    out = []
+    for hi in (False, True):
+        v = fx.Vector(rocdl.cvt_pk_f32_fp8(T.vec(2, T.f32), as_ir_value(w), hi))
+        out += [v[0], v[1]]
+    return out
+
+
+def fp8_pack4(a, b, c, d):
+    """Four f32 (in range) -> one dword of four E4M3, RNE (``ops``' but the
+    first pair over a poison dword: no zeroing move)."""
+    lo = fx.Int32(
+        rocdl.cvt_pk_fp8_f32(T.i32, a, b, fx.Int32(llvm.mlir_poison(T.i32)), False)
+    )
+    return fx.Int32(rocdl.cvt_pk_fp8_f32(T.i32, c, d, lo, True))
+
+
+def fp8_split4(x):
+    """Four f32 (in range) -> (hi, lo) dwords of four E4M3: x ~ hi + lo (the
+    differences as pairs, as the unpack gives them: no shuffling moves)."""
+    hi = fp8_pack4(*x)
+    lo = []
+    for k in range(2):
+        hf = fx.Vector(rocdl.cvt_pk_f32_fp8(T.vec(2, T.f32), as_ir_value(hi), k == 1))
+        d = fx.Vector.from_elements([x[2 * k], x[2 * k + 1]], fx.Float32) - hf
+        lo += [d[0], d[1]]
+    return hi, fp8_pack4(*lo)
+
+
+def xp_at(r, h, k, lane):
+    """The P dword of (row tile r, hi / lo h, tiles 4 k ..) of ``lane``."""
+    return (((r * 2 + h) * 2 + k) * 64 + lane) * 4
+
+
+@traced
+def q_split(c, rows, at):
+    """The CTA's bf16 Q rows (slot 1: ``[row][ENTRY]``) -> E4M3 hi, lo of q 2^k
+    (a power of two a row: its |max| in [128, 256)): hi + lo is q exactly,
+    bf16's 8 bits E4M3's 4 + 4 (lo q - hi, a multiple of hi's ulp / 16; but
+    below 2^-9 of 2^k: 2^-17 of the row's |max|). At dword ``at``:
+    ``[hi | lo][row][ENTRY]`` bytes, the tail chunk's also at ``qt``; a row's
+    2^-k at ``xm`` + row. A row 16 threads, 4 latents a unit (``tid % 16`` +
+    16 i)."""
+    tid = c["tid"]
+    rr = tid // 16
+    if rr < rows:
+        src = c["ks"] + SLOT + rr * (ENTRY // 2)
+        us = [tid % 16 + 16 * i for i in range(ENTRY // 4 // 16)]
+        qs = [
+            fx.Vector(
+                fx.ptr_load(src + 2 * u, result_type=fx.Vector.make_type(2, fx.Int32))
+            ).bitcast(fx.BFloat16)
+            for u in us
+        ]
+        x = [fx.Float32(q[e]) for q in qs for e in range(4)]
+        am = fx.Float32(1.0e-30)
+        for k in range_constexpr(len(x)):
+            am = fx.max(am, fx.max(x[k], fx.Float32(0.0) - x[k]))
+        am = butterfly(am, (8, 4, 2, 1), fx.max)
+        e = (am.bitcast(fx.Int32) >> 23) & 0xFF
+        e = fx.max(fx.min(e, fx.Int32(253)), fx.Int32(8))
+        qi = ((261 - e) << 23).bitcast(fx.Float32)
+        if tid % 16 == 0:
+            fx.ptr_store(((e - 7) << 23).bitcast(fx.Float32), c["xm"] + rr)
+        for i in range_constexpr(len(us)):
+            u = us[i]
+            y = [x[4 * i + j] * qi for j in range(4)]
+            hi = fp8_pack4(*y)
+            lo = fp8_pack4(
+                *[
+                    y[j] - fx.Float32(rocdl.cvt_f32_fp8(T.f32, as_ir_value(hi), j))
+                    for j in range(4)
+                ]
+            )
+            fx.ptr_store(hi, c["ks"] + at + rr * (ENTRY // 4) + u)
+            fx.ptr_store(lo, c["ks"] + at + (rows + rr) * (ENTRY // 4) + u)
+            if const_expr(i == len(us) - 1):
+                # the tail chunk's hi | lo kept apart (``qt``: read a block)
+                tb = c["qt"] + rr * (2 * CHUNK // 4) + u - 2 * QSTEPS * CHUNK // 4
+                fx.ptr_store(hi, tb)
+                fx.ptr_store(lo, tb + CHUNK // 4)
+
+
+def q_parts(c, r, at, rows):
+    """Row tile r's B operands from ``q_split``'s: QK step s's hi, lo
+    (s < ``QSTEPS``: line s), 8 dwords each (a lane 32 B at 32 g of the line,
+    K's order); the tail's stay in LDS (``qt``)."""
+    row = r * TILE + c["lane"] % TILE
+    g = c["lane"] // TILE
+
+    def ld(part, m, k):
+        return fx.Vector(
+            fx.ptr_load(
+                c["ks"]
+                + at
+                + (part * rows + row) * (ENTRY // 4)
+                + (
+                    32 * m + 8 * g + 4 * k
+                    if const_expr(c["lines"])
+                    else 16 * (2 * m + k) + 4 * g
+                ),
+                result_type=fx.Vector.make_type(4, fx.Int32),
+            )
+        )
+
+    def cat(a, b):
+        return fx.Vector.from_elements(
+            [a[i] for i in range(4)] + [b[i] for i in range(4)], fx.Int32
+        )  # noqa: E501
+
+    parts = []
+    for m in range_constexpr(QSTEPS):
+        for h in range_constexpr(2):
+            parts.append(cat(ld(h, m, 0), ld(h, m, 1)))
+    return parts
+
+
+def kpos(t, cc):
+    """The 16 B position (of 64) in a tile's tail chunk of token t's 16 B cc:
+    a lane group's K read (tokens 0..15, one cc) over the 64 banks."""
+    return 16 * cc + 4 * ((t % 4 + cc) % 4) + t // 4
+
+
+def sig(t, tp):
+    """The 16 B column swizzle (XOR) of token t (of 16) of a tile of parity
+    tp in a line chunk: a K read (tokens 0..15, one column) and a V^T read
+    (tokens 4 g + j, tiles 2 n, 2 n + 1) both over the 64 banks."""
+    return ((t >> 1) & 3) | ((((t >> 3) ^ tp) & 1) << 2)
+
+
+# tokens t, t + 8 of a lane (``dma_offs``): columns 64 B apart
+assert all(sig(t + 8, tp) == sig(t, tp) ^ 4 for t in range(8) for tp in (0, 1))
+
+
+def lines_at(t, tp, col):
+    """Token t's 16 B ``col`` of line s: its byte in the slot (from line s's
+    chunk 2 s, a tile 1 KB: 8 tokens x 128 B), the DMA lane ``8 (t % 8) +
+    (col ^ sig)`` of chunk 2 s + t // 8."""
+    return (t // 8) * WAVES * TILEW * 4 + (t % 8) * LINE + 16 * (col ^ sig(t, tp))
+
+
+def dma_lane(lane):
+    """(token, 16 B cc) tail DMA lane ``lane`` carries (it lands at position
+    lane)."""
+    cc = lane // 16
+    return 4 * (lane % 4) + ((lane // 4) % 4 + 4 - cc) % 4, cc
+
+
+def vtr_lane(lane, wave, i, lines=True):
+    """This lane's byte in a block's (line wave // 2) of the 8 B it gives a
+    V^T transpose read of latent tile i (``tr_read``; tiles 2 n .. at + 2 n
+    TILEW): lane k of a 16 its half k % 2 of the 16 B 4 (wave % 2) + i of
+    token 4 g + k // 2 % 4 of tile k // 8 -> lane m: latent 64 wave + 16 i +
+    m, tokens 4 g .. of tiles 2 n, 2 n + 1."""
+    g, k = lane // 16, lane % 16
+    t = 4 * g + (k // 2) % 4
+    if not lines:
+        # the 64 B chunk layout: latent tile i at chunk wave
+        return (k // 8) * TILEW * 4 + 16 * kpos(t, i) + 8 * (k % 2)
+    col = 4 * (wave % 2) + i
+    return (k // 8) * TILEW * 4 + lines_at(t, k // 8, col) + 8 * (k % 2)
+
+
+def tr_read(addr):
+    """ds_read_b64_tr_b8 at LDS byte ``addr``: a lane m of 16 byte m % 8 of
+    the 8 B of lanes m // 8 + 2 e (byte e), as two dwords."""
+    p3 = llvm.inttoptr(ir.Type.parse("!llvm.ptr<3>"), as_ir_value(fx.Int32(addr)))
+    return fx.Vector(rocdl.ds_read_tr8_b64(T.vec(2, T.i32), p3).result)
+
+
+def idx_load(c, b):
+    """The page indices of this lane's DMA tokens of block b's tile of this
+    wave (the tail's, lines' tokens lane // 8, + 8), clamped into the chunk
+    (a block past the last loads its rows)."""
+    t0 = c["c0"] + (b * WAVES + c["wave"]) * TILE
+    if const_expr(not c["lines"]):
+        return (ld_i32(c["args"]["pages"], c["p0"] + fx.min(t0 + c["dt"], c["last"])),)
+    return tuple(
+        ld_i32(c["args"]["pages"], c["p0"] + fx.min(t0 + t, c["last"]))
+        for t in (c["dt"], c["lane"] // 8, c["lane"] // 8 + 8)
+    )
+
+
+def dma_offs(c, idx):
+    """``k_dma``'s sources of ``idx_load``'s indices: the tail chunk's (16 B
+    dcc), line chunks' tokens lane // 8 and + 8 (their 16 B column
+    ``sig``-swizzled); bytes off the pool (``wide``: addresses)."""
+    cols = (
+        (c["dcc"] * 16, c["dcol"], c["dcol"] ^ 64)
+        if const_expr(c["lines"])
+        else (c["dcc"] * 16,)
+    )
+    if const_expr(c["wide"]):
+        return tuple(
+            c["args"]["kv"] + fx.Int64(i) * ENTRY + fx.Int64(col)
+            for i, col in zip(idx, cols)
+        )
+    return tuple(fx.Int32(i) * ENTRY + col for i, col in zip(idx, cols))
+
+
+def k_dma(c, b, offs, ps=None, tile=None):
+    """Block b's K tile of this wave (or ``tile``; ``dma_offs``) -> slot
+    b % 2: a DMA a chunk (64 lanes x 16 B: line s's 8 tokens, the tail's 16).
+    Inline asm: LLVM's own wait for an LDS DMA would hold this block's slot
+    reads for the next block's DMA too; ``vmwait`` / ``wait_dma`` are the only
+    waits."""
+    ps = range(KLOADS) if ps is None else ps
+    tile = c["wave"] if tile is None else tile
+    dst = c["ks_b"] + ((b % 2) * SLOT + tile * TILEW) * 4
+
+    def at(p):
+        # (offset, immediate): chunk 2 s + h line s of tokens 8 h ..
+        if const_expr(p == KLOADS - 1 or not c["lines"]):
+            return offs[0], CHUNK * p
+        return offs[1 + p % 2], LINE * (p // 2)
+
+    if const_expr(c["wide"]):
+        for p in ps:
+            o, imm = at(p)
+            lds_dma(o + imm, dst + p * WAVES * TILEW * 4)
+    else:
+        # a 32-bit offset off the pool's SGPR base, the line in the
+        # immediate (half the address VGPRs a DMA reads; the immediate
+        # offsets the LDS side too, M0 takes it back)
+        for p in ps:
+            o, imm = at(p)
+            llvm.inline_asm(
+                None,
+                [
+                    as_ir_value(o),
+                    as_ir_value(fx.Int64(c["args"]["kv"])),
+                    as_ir_value(uniform(dst + p * WAVES * TILEW * 4 - imm)),
+                ],
+                "s_mov_b32 m0, $2\n\ts_nop 0\n\t"
+                f"global_load_lds_dwordx4 $0, $1 offset:{imm}",
+                "v,s,s,~{m0},~{memory}",
+                has_side_effects=True,
+            )
+
+
+def lds_dma(src, dst):
+    """A global_load_lds_dwordx4: lane i's 16 B at ``src`` -> LDS byte
+    ``dst`` (uniform) + 16 i. (Every LDS DMA here an ``s_nop`` between its M0
+    write and the load: right after the write the load could read a stale M0,
+    the warp-specialized waves' K / index DMAs then landing in another slot:
+    nondeterministic outputs.)"""
+    llvm.inline_asm(
+        None,
+        [as_ir_value(fx.Int64(src)), as_ir_value(uniform(dst))],
+        "s_mov_b32 m0, $1\n\ts_nop 0\n\tglobal_load_lds_dwordx4 $0, off",
+        "v,s,~{m0},~{memory}",
+        has_side_effects=True,
+    )
+
+
+@traced
+def idx_dma(c, b, t0=None, n=1):
+    """Block b's page indices of this wave's tile (or tiles t0 .. t0 + n) ->
+    its LDS index slot b % 2 (lanes 0 .. 16 n, token lane's): a DMA, so it
+    lands in order with the K DMAs (a load would make LLVM wait on every DMA
+    before it to read it)."""
+    t0 = c["wave"] if t0 is None else t0
+    if c["lane"] < n * TILE:
+        tok = c["c0"] + (b * WAVES + t0) * TILE + c["lane"]
+        src = c["args"]["pages"] + fx.Int64(c["p0"] + fx.min(tok, c["last"])) * 4
+        dst = c["xi_b"] + ((b % 2) * WAVES + t0) * TILE * 4
+        llvm.inline_asm(
+            None,
+            [as_ir_value(fx.Int64(src)), as_ir_value(uniform(dst))],
+            "s_mov_b32 m0, $1\n\ts_nop 0\n\tglobal_load_lds_dword $0, off",
+            "v,s,~{m0},~{memory}",
+            has_side_effects=True,
+        )
+
+
+def idx_lds(c, b):
+    """This lane's DMA tokens' page indices of block b (``idx_dma``'s; as
+    ``idx_load``)."""
+    xb = c["xi"] + ((b % 2) * WAVES + c["wave"]) * TILE
+    if const_expr(not c["lines"]):
+        return (fx.ptr_load(xb + c["dt"]),)
+    return (
+        fx.ptr_load(xb + c["dt"]),
+        fx.ptr_load(xb + c["lane"] // 8),
+        fx.ptr_load(xb + c["lane"] // 8 + 8),
+    )
+
+
+def vmwait(n):
+    """s_waitcnt vmcnt(n): at most n of this wave's loads / DMAs in flight
+    (they land in order)."""
+    llvm.inline_asm(
+        None, [], f"s_waitcnt vmcnt({n})", "~{memory}", has_side_effects=True
+    )
+
+
+def wait_dma():
+    """This wave's DMAs (and loads) landed; the CTA's after the barrier."""
+    llvm.inline_asm(None, [], "s_waitcnt vmcnt(0)", "", has_side_effects=True)
+    gpu.barrier()
+
+
+def xm_row(c, r, row):
+    """The waves' values of row ``row`` of row tile r in xm
+    (``[r][row][wave]``: two b128 reads)."""
+    v = [
+        fx.Vector(
+            fx.ptr_load(
+                c["xm"] + (r * TILE + row) * WAVES + 4 * k,
+                result_type=fx.Vector.make_type(4, fx.Float32),
+            )
+        )
+        for k in range(2)
+    ]
+    return [v[w // 4][w % 4] for w in range(WAVES)]
+
+
+def v_reads(c, b, dma):
+    """V^T of the slot's block b by transpose reads, the A of a 16x16x128 fp8
+    MFMA a latent tile i: a lane's 32 B tile w's tokens 4 g .. at 4 w (P's
+    order); ``dma`` the hook ``v<i>``'s K chunks before tile i's reads."""
+    vw = (c["wave"] // 2) * 2 if const_expr(c["lines"]) else c["wave"]
+    vt = c["ks_b"] + ((b % 2) * SLOT + vw * WAVES * TILEW) * 4
+    va = []
+    for i in range_constexpr(4):
+        dma(f"v{i}")
+        tr = [tr_read(vt + c["vtr"][i] + n * 2 * TILEW * 4) for n in range(4)]
+        va.append(
+            fx.Vector.from_elements(
+                [tr[n][e] for n in range(4) for e in range(2)], fx.Int32
+            )
+        )
+    return va
+
+
+@traced
+def stage_block(c, b, st, dma, kwait, masked):
+    """One block (its K in slot b % 2) into the running (m, l, O) of each of
+    the CTA's row tiles (K and V^T read once for all); ``dma(hook)``
+    issues the hook's of block b + 1's K chunks (``DMA``), ``kwait(p)`` waits
+    for this wave's chunk p of block b; ``masked``: a block with a token past
+    the chunk or a row's causal bound (else every score kept)."""
+    lane, wave = c["lane"], c["wave"]
+    row = lane % TILE
+    rt = c["rt"]
+    m_run, l_run, o = st
+    kt = c["ks"] + (b % 2) * SLOT + wave * TILEW
+
+    def kld(at):
+        return fx.Vector(
+            fx.ptr_load(kt + at, result_type=fx.Vector.make_type(4, fx.Int32))
+        )
+
+    def kstep(s):
+        """Step s's K words: line s's 16 B columns 2 g, 2 g + 1 (chunks 2 s,
+        2 s + 1; the tail: chunk 8's g), once landed."""
+        if const_expr(not c["lines"]):
+            ats = [
+                p * WAVES * TILEW + c["qpos"][2]
+                for p in ([2 * s, 2 * s + 1] if s < QSTEPS else [2 * s])
+            ]
+        else:
+            ats = (
+                [2 * s * WAVES * TILEW + c["qpos"][k] for k in range(2)]
+                if s < QSTEPS
+                else [2 * s * WAVES * TILEW + c["qpos"][2]]
+            )
+        kwait(2 * s + 1 if s < QSTEPS else 2 * s)
+        return [w for at in ats for w in (lambda v: [v[e] for e in range(4)])(kld(at))]
+
+    # S^T: token (4 g + j) x row (lane % 16), a tile r of rows each; fp8 K
+    # straight into the MFMA against Q's hi and lo (``q_parts``), a step's K
+    # read during the step before (``sched_barrier``: no hoisting past it)
+    acc = [fx.Vector.filled(4, 0.0, fx.Float32) for _ in range(rt)]
+    nxt = kstep(0)
+    for s in range_constexpr(QSTEPS + 1):
+        words = nxt
+        if const_expr(s < QSTEPS):
+            nxt = kstep(s + 1)
+            ka = fx.Vector.from_elements(words, fx.Int32)
+            for r in range_constexpr(rt):
+                for h in range_constexpr(2):
+                    acc[r] = mfma_f8(ka, c["q"][r][2 * s + h], acc[r])
+        else:
+            # the tail: chunk 8 twice against its hi | lo (from ``qt``: in
+            # VGPRs they'd spill, the reload a stall a block)
+            ka = fx.Vector.from_elements(words + words, fx.Int32)
+            for r in range_constexpr(rt):
+                qb = c["qt"] + (r * TILE + row) * (2 * CHUNK // 4) + 4 * (lane // 16)
+                qtl = [
+                    fx.Vector(
+                        fx.ptr_load(
+                            qb + h * CHUNK // 4,
+                            result_type=fx.Vector.make_type(4, fx.Int32),
+                        )
+                    )
+                    for h in range(2)
+                ]
+                qtv = fx.Vector.from_elements(
+                    [qtl[h][e] for h in range(2) for e in range(4)], fx.Int32
+                )
+                acc[r] = mfma_f8(ka, qtv, acc[r])
+        dma(f"q{s}")
+        rocdl.sched_barrier(0)
+    # scores (log2 units: acc qscale, qscale > 0, so the max on acc's),
+    # causal + chunk mask
+    sc = [[acc[r][j] for j in range(4)] for r in range(rt)]
+    if const_expr(masked):
+        t0 = c["c0"] + (b * WAVES + wave) * TILE + 4 * (lane // TILE)
+        ok = [(t0 + j <= c["last"]) & (t0 + j <= c["bound"]) for j in range(4)]
+        ninf = fx.Float32(float("-inf"))
+        sc = [[ok[j].select(sc[r][j], ninf) for j in range(4)] for r in range(rt)]
+    mx = [
+        fx.max(fx.max(sc[r][0], sc[r][1]), fx.max(sc[r][2], sc[r][3]))
+        for r in range(rt)
+    ]
+    if const_expr(rt == 2):
+        # the two row tiles' halves in one swap: lanes 32 r .. row tile r's
+        x, y = permlane_swap(32, mx[0].bitcast(fx.Int32), mx[1].bitcast(fx.Int32))
+        m2 = butterfly(
+            fx.max(x.bitcast(fx.Float32), y.bitcast(fx.Float32)), (16,), fx.max
+        )
+        if lane % 32 < TILE:
+            fx.ptr_store(m2, c["xm"] + ((lane // 32) * TILE + row) * WAVES + wave)
+    else:
+        for r in range_constexpr(rt):
+            m1 = butterfly(mx[r], (32, 16), fx.max)
+            if lane < TILE:
+                fx.ptr_store(m1, c["xm"] + (r * TILE + row) * WAVES + wave)
+    gpu.barrier()
+    stamp(c, 2 + 4 * b + 1)
+    dma("p")
+    mb = []
+    for r in range_constexpr(rt):
+        xw = xm_row(c, r, row)
+        mb.append(
+            fx.max(
+                fx.max(fx.max(xw[0], xw[1]), fx.max(xw[2], xw[3])),
+                fx.max(fx.max(xw[4], xw[5]), fx.max(xw[6], xw[7])),
+            )
+            * c["qscale"][r]
+        )
+    m_new, alpha = [], []
+    for r in range_constexpr(rt):
+        m_new.append(fx.max(m_run[r], mb[r]))
+        alpha.append(hw_exp2(m_run[r] - m_new[r]))
+        # p 2^PLOG (l, O alike: the normalization undoes it)
+        nmp = fx.Float32(PLOG) - m_new[r]
+        p = [hw_exp2(fmath.fma(sc[r][j], c["qscale"][r], nmp)) for j in range(4)]
+        # l a lane's (its 4 tokens'), summed over the lanes at the end
+        # (pairs, as the split's: no shuffling moves)
+        pp = [
+            fx.Vector.from_elements(p[2 * k : 2 * k + 2], fx.Float32) for k in range(2)
+        ]
+        ps = pp[0] + pp[1]
+        l_run[r] = l_run[r] * alpha[r] + (ps[0] + ps[1])
+        # [r][hi, lo][tile // 4][lane][tile % 4]: a lane's 8 tiles two b128
+        pw = fp8_split4(p)
+        for h in range_constexpr(2):
+            fx.ptr_store(pw[h], c["xp"] + xp_at(r, h, wave // 4, lane) + wave % 4)
+    gpu.barrier()
+    stamp(c, 2 + 4 * b + 2)
+    # O^T += V^T P^T, the block an MFMA a (latent tile i, P part) (k: tile
+    # w's tokens 4 g .. at 4 w): latents 64 wave + 16 i + 4 g + j, chunk p =
+    # wave of the slot; o[4 r + i] row tile r's (P's hi and lo into it alike)
+    o = [
+        fx.Vector.from_elements(
+            [o[4 * r + i][j] * alpha[r] for j in range(4)], fx.Float32
+        )
+        for r in range(rt)
+        for i in range(4)
+    ]
+    va = v_reads(c, b, dma)
+    for r in range_constexpr(rt):
+        for h in range_constexpr(2):
+            pb = [
+                fx.Vector(
+                    fx.ptr_load(
+                        c["xp"] + xp_at(r, h, k, lane),
+                        result_type=fx.Vector.make_type(4, fx.Int32),
+                    )
+                )
+                for k in range(2)
+            ]
+            pv = fx.Vector.from_elements(
+                [pb[w // 4][w % 4] for w in range(WAVES)], fx.Int32
+            )
+            dma(f"m{2 * r + h}")
+            for i in range_constexpr(4):
+                o[4 * r + i] = mfma_f8(va[i], pv, o[4 * r + i])
+    # the PV MFMAs before the next block's index wait (in order, it would
+    # hold them for this block's DMAs)
+    rocdl.sched_barrier(0)
+    stamp(c, 2 + 4 * b + 3)
+    return m_new, l_run, o
+
+
+def block(c, b, st, masked):
+    """Block b (``stage_block``) behind the top barrier, its DMA hooks."""
+    rt = c["rt"]
+    # the slot block b + 1's DMA takes was the last block's (its PV
+    # over): the waves in step; block b's own K waits a chunk
+    gpu.barrier()
+    stamp(c, 2 + 4 * b)
+    idx_dma(c, b + 2)
+    # block b + 1's indices landed (newer: block b's 9 chunks, b + 2's
+    # indices)
+    vmwait(KLOADS + 1)
+    i1 = dma_offs(c, idx_lds(c, b + 1))
+    # in flight (oldest first): block b's K chunks, block b + 2's
+    # indices, then block b + 1's chunks as sent
+    sent = [0]
+
+    def kwait(p):
+        # newer than chunk p: chunks p + 1 .., the indices, the sent
+        vmwait(KLOADS - p + sent[0])
+
+    def dma(hook):
+        n = DMA[min(rt, 2)].get(hook, 0)
+        if const_expr(c["dbg"] != 1 and n > 0):
+            # pinned where the schedule puts it
+            rocdl.sched_barrier(0)
+            k_dma(c, b + 1, i1, range(sent[0], sent[0] + n))
+            rocdl.sched_barrier(0)
+        sent[0] += n
+
+    return stage_block(c, b, st, dma, kwait, masked)
+
+
+@traced
+def run(c):
+    """The blocks (block 0's DMA and block 1's indices issued) and the
+    partial's stores."""
+    lane, wave, g = c["lane"], c["wave"], c["lane"] // TILE
+    row = lane % TILE
+    rt = c["rt"]
+    m_run = [fx.Float32(NEG) for _ in range(rt)]
+    l_run = [fx.Float32(0.0) for _ in range(rt)]
+    o = [fx.Vector.filled(4, 0.0, fx.Float32) for _ in range(4 * rt)]
+    # K a block ahead (slot b % 2), the indices two (``xi``, DMA'd too)
+    stamp(c, 1)
+    # the blocks every score of is kept (their last token at most the
+    # chunk's and position 0's bound) unmasked, then the rest
+    b = fx.Int32(0)
+    while b < c["nfull"]:
+        m_run, l_run, o = block(c, b, (m_run, l_run, o), False)
+        b = b + 1
+    while b < c["nblk"]:
+        m_run, l_run, o = block(c, b, (m_run, l_run, o), True)
+        b = b + 1
+    # l over the lanes, then the waves (each its tiles' sum, at the common
+    # running max)
+    l_run = [butterfly(l_run[r], (32, 16)) for r in range(rt)]
+    if lane < TILE:
+        for r in range_constexpr(rt):
+            fx.ptr_store(l_run[r], c["xm"] + (r * TILE + row) * WAVES + wave)
+    gpu.barrier()
+    a = c["args"]
+    for r in range_constexpr(rt):
+        xw = xm_row(c, r, row)
+        lt = xw[0]
+        for w in range_constexpr(1, WAVES):
+            lt = lt + xw[w]
+        empty = lt == fx.Float32(0.0)
+        inv = empty.select(fx.Float32(0.0), hw_rcp(lt) * c["vscale"])
+        # lane: latents 64 wave + 16 i + 4 g + j of its row; tiles 2 n,
+        # 2 n + 1 traded over the 16-lane rows (``permlane_swap``): an even
+        # g its 8 latents 4 g .. of tile 2 n, an odd g 4 (g - 1) .. of 2 n + 1
+        th = (c["split"] * c["S"] + c["t"]) * NH + c["h"][r]
+        pk = [
+            [
+                bf16_pair(o[4 * r + i][j] * inv, o[4 * r + i][j + 1] * inv).bitcast(
+                    fx.Int32
+                )
+                for j in (0, 2)
+            ]
+            for i in range(4)
+        ]
+        for n in range_constexpr(2):
+            sw = [permlane_swap(16, pk[2 * n][e], pk[2 * n + 1][e]) for e in range(2)]
+            lat = wave * VCOLS + 16 * (2 * n + g % 2) + 4 * (g - g % 2)
+            v = fx.Vector.from_elements(
+                [sw[0][0], sw[1][0], sw[0][1], sw[1][1]], fx.Int32
+            )
+            # stamps: split 0's rows hold them; the rest still store (PV live)
+            live = (c["split"] > 0) if const_expr(c["stamps"]) else True
+            if const_expr(c["dbg"] != 2):  # noqa: SIM102
+                if live:
+                    bo.buffer_store(
+                        v,
+                        rsrc(a["part_o"]),
+                        (th * KVL + lat) // 2,
+                        cache_modifier=c["cm"],
+                    )
+        if (wave == 0) & (lane < TILE):
+            lse = m_run[r] + (hw_log2(lt) - fx.Float32(PLOG))
+            lse = empty.select(fx.Float32(float("-inf")), lse * fx.Float32(LN2))
+            bo.buffer_store(lse, rsrc(a["part_lse"]), th, cache_modifier=c["cm"])
+    stamp(c, 63)
+
+
+def chunk_dma(c, src, imm, dst):
+    """One chunk DMA (``k_dma``'s): ``src`` + ``imm`` (an offset off the pool,
+    or ``wide`` an address) -> LDS byte ``dst`` (uniform; lane i its 16 B at
+    + 16 i)."""
+    if const_expr(c["wide"]):
+        lds_dma(src + imm, dst)
+    else:
+        llvm.inline_asm(
+            None,
+            [
+                as_ir_value(src),
+                as_ir_value(fx.Int64(c["args"]["kv"])),
+                as_ir_value(uniform(dst - imm)),
+            ],
+            "s_mov_b32 m0, $2\n\ts_nop 0\n\t"
+            f"global_load_lds_dwordx4 $0, $1 offset:{imm}",
+            "v,s,s,~{m0},~{memory}",
+            has_side_effects=True,
+        )
+
+
+def ws_offs(c, b, t):
+    """This lane's DMA source of tile t of block b (``dma_offs``, its index
+    from the index slot)."""
+    return dma_offs(
+        c, (fx.ptr_load(c["xi"] + ((b % 2) * WAVES + t) * TILE + c["dt"]),)
+    )[0]
+
+
+def ws_dmas(c, bb, ps, hooks=("all",)):
+    """This wave's DMA hook of block bb's tile's chunks ``ps`` (``k_dma``'s),
+    spread over ``hooks`` (the first ones one more)."""
+    if const_expr(c["dbg"] == 1):
+        return lambda hook: None
+    offs = (ws_offs(c, bb, c["wave"]),)
+    ps = list(ps)
+    nh = len(hooks)
+    cut = [0] + [len(ps) - (len(ps) * (nh - i)) // nh for i in range(1, nh + 1)]
+    at = {h: ps[cut[i] : cut[i + 1]] for i, h in enumerate(hooks)}
+
+    def dma(hook):
+        if const_expr(len(at.get(hook, [])) > 0):
+            rocdl.sched_barrier(0)
+            k_dma(c, bb, offs, at[hook])
+            rocdl.sched_barrier(0)
+
+    return dma
+
+
+def ws_vtile(c, b, i):
+    """V^T latent tile i (of 8: 128 v + 16 i ..) of the slot's block b
+    (``v_reads``'): 8 dwords."""
+    vt = c["vb"] + ((b % 2) * SLOT + (i // 4) * WAVES * TILEW) * 4 + c["vtr"][i % 4]
+    tr = [tr_read(vt + n * 2 * TILEW * 4) for n in range(4)]
+    return fx.Vector.from_elements(
+        [tr[n][e] for n in range(4) for e in range(2)], fx.Int32
+    )
+
+
+@traced
+def ws_qk_block(c, b, st, masked):
+    """QK wave step b: S of block b's tiles 2 w, 2 w + 1 (``stage_block``'s
+    QK), the max over the QK waves (B1), block b + 2's tile DMA'd, P(b) -> xp,
+    alpha -> xm (B2: block b + 1's landed)."""
+    lane, wave = c["lane"], c["wave"]
+    row = lane % TILE
+    g = lane // TILE
+    rt = c["rt"]
+    m_run, l_run = st
+    # block b + 1's chunks WS_K1.. over the QK steps, then block b + 3's
+    # indices (its slot's read by then)
+    early = ws_dmas(c, b + 1, range(WS_K1, KLOADS), [f"k{s}" for s in range(QSTEPS)])
+    slot = c["ks"] + (b % 2) * SLOT
+    kt = slot + 2 * wave * TILEW + c["qpos"][2]
+
+    def kstep(s):
+        ps = [2 * s, 2 * s + 1] if s < QSTEPS else [2 * s]
+        return [
+            [
+                v[e]
+                for p in ps
+                for v in [
+                    fx.Vector(
+                        fx.ptr_load(
+                            kt + p * WAVES * TILEW + j * TILEW,
+                            result_type=fx.Vector.make_type(4, fx.Int32),
+                        )
+                    )
+                ]
+                for e in range(4)
+            ]
+            for j in range(2)
+        ]
+
+    acc = [[fx.Vector.filled(4, 0.0, fx.Float32) for _ in range(rt)] for _ in range(2)]
+    nxt = kstep(0)
+    for s in range_constexpr(QSTEPS + 1):
+        words = nxt
+        if const_expr(s < QSTEPS):
+            nxt = kstep(s + 1)
+            early(f"k{s}")
+            if const_expr(s == QSTEPS - 1):
+                idx_dma(c, b + 3)
+            for j in range_constexpr(2):
+                ka = fx.Vector.from_elements(words[j], fx.Int32)
+                for r in range_constexpr(rt):
+                    for h in range_constexpr(2):
+                        acc[j][r] = mfma_f8(ka, c["q"][r][2 * s + h], acc[j][r])
+        else:
+            for r in range_constexpr(rt):
+                qb = c["qt"] + (r * TILE + row) * (2 * CHUNK // 4) + 4 * g
+                qtl = [
+                    fx.Vector(
+                        fx.ptr_load(
+                            qb + h * CHUNK // 4,
+                            result_type=fx.Vector.make_type(4, fx.Int32),
+                        )
+                    )
+                    for h in range_constexpr(2)
+                ]
+                qtv = fx.Vector.from_elements(
+                    [qtl[h][e] for h in range(2) for e in range(4)], fx.Int32
+                )
+                for j in range_constexpr(2):
+                    ka = fx.Vector.from_elements(words[j] + words[j], fx.Int32)
+                    acc[j][r] = mfma_f8(ka, qtv, acc[j][r])
+        rocdl.sched_barrier(0)
+    sc = [[[acc[j][r][e] for e in range(4)] for r in range(rt)] for j in range(2)]
+    if const_expr(masked):
+        ninf = fx.Float32(float("-inf"))
+        for j in range_constexpr(2):
+            t0 = c["c0"] + (b * WAVES + 2 * wave + j) * TILE + 4 * g
+            ok = [(t0 + e <= c["last"]) & (t0 + e <= c["bound"]) for e in range(4)]
+            sc[j] = [
+                [ok[e].select(sc[j][r][e], ninf) for e in range(4)] for r in range(rt)
+            ]
+    mx = []
+    for r in range_constexpr(rt):
+        m4 = [fx.max(sc[0][r][e], sc[1][r][e]) for e in range(4)]
+        mx.append(fx.max(fx.max(m4[0], m4[1]), fx.max(m4[2], m4[3])))
+    x, y = permlane_swap(32, mx[0].bitcast(fx.Int32), mx[1].bitcast(fx.Int32))
+    m2 = butterfly(fx.max(x.bitcast(fx.Float32), y.bitcast(fx.Float32)), (16,), fx.max)
+    if lane % 32 < TILE:
+        fx.ptr_store(m2, c["xm"] + ((lane // 32) * TILE + row) * WAVES + wave)
+    gpu.barrier()
+    ws_dmas(c, b + 2, range(WS_K1))("all")
+    m_new, alpha = [], []
+    for r in range_constexpr(rt):
+        xw = fx.Vector(
+            fx.ptr_load(
+                c["xm"] + (r * TILE + row) * WAVES,
+                result_type=fx.Vector.make_type(4, fx.Float32),
+            )
+        )
+        mb = fx.max(fx.max(xw[0], xw[1]), fx.max(xw[2], xw[3])) * c["qscale"][r]
+        m_new.append(fx.max(m_run[r], mb))
+        alpha.append(hw_exp2(m_run[r] - m_new[r]))
+        nmp = fx.Float32(PLOG) - m_new[r]
+        lsum = None
+        for j in range_constexpr(2):
+            p = [hw_exp2(fmath.fma(sc[j][r][e], c["qscale"][r], nmp)) for e in range(4)]
+            pp = [
+                fx.Vector.from_elements(p[2 * k : 2 * k + 2], fx.Float32)
+                for k in range(2)
+            ]
+            lsum = pp[0] + pp[1] if lsum is None else lsum + pp[0] + pp[1]
+            pw = fp8_split4(p)
+            t = 2 * wave + j
+            for h in range_constexpr(2):
+                fx.ptr_store(pw[h], c["xp"] + xp_at(r, h, t // 4, lane) + t % 4)
+        assert lsum is not None
+        l_run[r] = l_run[r] * alpha[r] + (lsum[0] + lsum[1])
+    if (wave == 0) & (lane < TILE):
+        for r in range_constexpr(rt):
+            fx.ptr_store(alpha[r], c["xm"] + (r * TILE + row) * WAVES + QKW)
+    vmwait(0 if const_expr(c["dbg"] == 1) else WS_K1)
+    gpu.barrier()
+    return m_new, l_run
+
+
+@traced
+def ws_qk(c):
+    """The QK waves' steps (``ws_qk_block``), the last block's PV's barriers,
+    then the l over the waves and the LSE."""
+    lane, wave = c["lane"], c["wave"]
+    row = lane % TILE
+    rt = c["rt"]
+    m_run = [fx.Float32(NEG) for _ in range(rt)]
+    l_run = [fx.Float32(0.0) for _ in range(rt)]
+    # block 1's tile's chunks ..WS_K1 (the rest step 0's)
+    ws_dmas(c, 1, range(WS_K1))("all")
+    b = fx.Int32(0)
+    while b < c["nfull"]:
+        m_run, l_run = ws_qk_block(c, b, (m_run, l_run), False)
+        b = b + 1
+    while b < c["nblk"]:
+        m_run, l_run = ws_qk_block(c, b, (m_run, l_run), True)
+        b = b + 1
+    vmwait(0)
+    gpu.barrier()
+    gpu.barrier()
+    l_run = [butterfly(l_run[r], (32, 16)) for r in range(rt)]
+    if lane < TILE:
+        for r in range_constexpr(rt):
+            fx.ptr_store(l_run[r], c["xm"] + (r * TILE + row) * WAVES + wave)
+    gpu.barrier()
+    if (wave == 0) & (lane < TILE):
+        for r in range_constexpr(rt):
+            lt = ws_l(c, r, row)
+            empty = lt == fx.Float32(0.0)
+            th = (c["split"] * c["S"] + c["t"]) * NH + c["h"][r]
+            lse = m_run[r] + (hw_log2(lt) - fx.Float32(PLOG))
+            lse = empty.select(fx.Float32(float("-inf")), lse * fx.Float32(LN2))
+            bo.buffer_store(
+                lse, rsrc(c["args"]["part_lse"]), th, cache_modifier=c["cm"]
+            )
+
+
+def ws_l(c, r, row):
+    """Row ``row`` of row tile r's l: the QK waves' (xm) summed."""
+    xw = fx.Vector(
+        fx.ptr_load(
+            c["xm"] + (r * TILE + row) * WAVES,
+            result_type=fx.Vector.make_type(4, fx.Float32),
+        )
+    )
+    return (xw[0] + xw[1]) + (xw[2] + xw[3])
+
+
+@traced
+def ws_pv_block(c, b, st):
+    """PV wave step b: O += V^T P^T of block b - 1 (its V^T ``va`` read in
+    step b - 1, P from xp, alpha from xm), B1 after ``WS_B1``
+    latent tiles, V^T(b) read before it (block b's slot then free); after B1
+    block b + 2's chunks DMA'd (block b + 1's landed by B2)."""
+    lane = c["lane"]
+    row = lane % TILE
+    rt = c["rt"]
+    o, va = st
+    early = ws_dmas(c, b + 1, range(WS_K1, KLOADS), [f"m{i}" for i in range(WS_B1)])
+    xp = c["xp"]
+    al = [fx.ptr_load(c["xm"] + (r * TILE + row) * WAVES + QKW) for r in range(rt)]
+    pv = []
+    for r in range_constexpr(rt):
+        for h in range_constexpr(2):
+            pb = [
+                fx.Vector(
+                    fx.ptr_load(
+                        xp + xp_at(r, h, k, lane),
+                        result_type=fx.Vector.make_type(4, fx.Int32),
+                    )
+                )
+                for k in range(2)
+            ]
+            pv.append(
+                fx.Vector.from_elements(
+                    [pb[w // 4][w % 4] for w in range(WAVES)], fx.Int32
+                )
+            )
+    o = [
+        fx.Vector.from_elements([o[8 * r + i][j] * al[r] for j in range(4)], fx.Float32)
+        for r in range(rt)
+        for i in range(8)
+    ]
+    nva = []
+    dma = None
+    for i in range_constexpr(8):
+        if const_expr(i == WS_B1):
+            nva += [ws_vtile(c, b, k) for k in range(WS_B1, 8)]
+            idx_dma(c, b + 3)
+            gpu.barrier()
+            dma = ws_dmas(c, b + 2, range(WS_K1), [f"m{k}" for k in range(WS_B1, 8)])
+        for r in range_constexpr(rt):
+            for h in range_constexpr(2):
+                o[8 * r + i] = mfma_f8(va[i], pv[2 * r + h], o[8 * r + i])
+        if const_expr(i < WS_B1):
+            nva.append(ws_vtile(c, b, i))
+            early(f"m{i}")
+        else:
+            assert dma is not None
+            dma(f"m{i}")
+    if const_expr(WS_B1 >= 8):
+        idx_dma(c, b + 3)
+        gpu.barrier()
+        ws_dmas(c, b + 2, range(WS_K1))("all")
+    vmwait(0 if const_expr(c["dbg"] == 1) else WS_K1)
+    gpu.barrier()
+    return o, nva
+
+
+@traced
+def ws_pv(c):
+    """The PV waves' steps: step 0 V^T(0) and block 1's DMA, then
+    ``ws_pv_block`` a block (one behind the QK waves), then o = O / l."""
+    lane, wave, g = c["lane"], c["wave"], c["lane"] // TILE
+    row = lane % TILE
+    rt = c["rt"]
+    o = [fx.Vector.filled(4, 0.0, fx.Float32) for _ in range(8 * rt)]
+    # step 0: block 1 (landed by B2), block 3's indices (after block 1's
+    # read theirs: the slot's), V^T(0), block 2 after B1
+    ws_dmas(c, 1, range(KLOADS))("all")
+    idx_dma(c, 3)
+    va = [ws_vtile(c, 0, i) for i in range(8)]
+    gpu.barrier()
+    ws_dmas(c, 2, range(WS_K1))("all")
+    vmwait(0 if const_expr(c["dbg"] == 1) else WS_K1)
+    gpu.barrier()
+    b = fx.Int32(1)
+    while b < c["nblk"] + 1:
+        o, va = ws_pv_block(c, b, (o, va))
+        b = b + 1
+    # the DMAs past the last block (clamped) land before the CTA's LDS goes
+    vmwait(0)
+    gpu.barrier()
+    a = c["args"]
+    v = wave - QKW
+    for r in range_constexpr(rt):
+        lt = ws_l(c, r, row)
+        empty = lt == fx.Float32(0.0)
+        inv = empty.select(fx.Float32(0.0), hw_rcp(lt) * c["vscale"])
+        th = (c["split"] * c["S"] + c["t"]) * NH + c["h"][r]
+        for hf in range_constexpr(2):
+            pk = [
+                [
+                    bf16_pair(
+                        o[8 * r + 4 * hf + i][j] * inv,
+                        o[8 * r + 4 * hf + i][j + 1] * inv,
+                    ).bitcast(fx.Int32)
+                    for j in (0, 2)
+                ]
+                for i in range_constexpr(4)
+            ]
+            for n in range_constexpr(2):
+                sw = [
+                    permlane_swap(16, pk[2 * n][e], pk[2 * n + 1][e]) for e in range(2)
+                ]
+                lat = (2 * v + hf) * VCOLS + 16 * (2 * n + g % 2) + 4 * (g - g % 2)
+                vv = fx.Vector.from_elements(
+                    [sw[0][0], sw[1][0], sw[0][1], sw[1][1]], fx.Int32
+                )
+                if const_expr(c["dbg"] != 2):
+                    bo.buffer_store(
+                        vv,
+                        rsrc(a["part_o"]),
+                        (th * KVL + lat) // 2,
+                        cache_modifier=c["cm"],
+                    )
+
+
+@traced
+def ws_run(c):
+    """Warp-specialized blocks (``WSPEC``): the QK waves', the PV waves'."""
+    wv = uniform(c["wave"])
+    if wv < QKW:
+        ws_qk(c)
+    if wv >= QKW:
+        ws_pv(c)
+
+
+def check(s: int, q: int, rt: int = 0):
+    assert q in (4, 8) and s % q == 0
+    # slot 1 holds the bf16 Q rows and their fp8 hi | lo (``q_split``)
+    assert (rt or RTILES[q]) * TILE * ENTRY <= SLOT
+    assert NH % geometry(q, rt)[0] == 0
+    # warp-specialized: 2 row tiles
+    assert not wspec(rt or RTILES[q]) or (rt or RTILES[q]) == 2
+    assert splits(q, rt) % (s // q) == 0, (splits(q, rt), s // q)
+
+
+def smem(qlen: int, rt: int = 0, ws: int = -1):
+    """The LDS of a CTA (``cta``'s ks / xm / xp / xi; xp P's hi | lo, then
+    the tail Q's tail hi | lo, ``qt``)."""
+    rt = rt or RTILES[qlen]
+    xpn = rt * WAVES * 64 * 2
+
+    @fx.struct
+    class Smem:
+        ks: fx.Array[fx.Int32, 2 * SLOT, 16]
+        xm: fx.Array[fx.Float32, rt * WAVES * TILE, 16]
+        xp: fx.Array[fx.Int32, xpn + rt * TILE * 2 * CHUNK // 4, 16]
+        xi: fx.Array[fx.Int32, 2 * WAVES * TILE, 16]
+
+    return Smem
+
+
+@traced
+def cta(c):
+    """CTA ``c["bid"]``'s (head group, split): its partial stored (cache
+    modifier ``c["cm"]``), then ``c["done"](flag)`` when given (``ctas``).
+    ``c["args"]``: qb, kv, pages, indptr, part_o, part_lse, scale, vscale."""
+    s, q, a = c["S"], c["Q"], c["args"]
+    rt = c.get("rt") or RTILES[q]
+    ws = wspec(rt, c.get("ws", -1))
+    ht, hg, spx = geometry(q, rt)
+    nreq = s // q
+    per_req = splits(q, rt) // nreq
+    bid, tid = c["bid"], c["tid"]
+    lane = tid % 64
+    wave = tid // 64
+    g = lane // TILE
+    row = lane % TILE
+    x = bid % XCDS
+    u = bid // XCDS
+    if u < spx * hg:
+        gs = x * spx + u // hg  # global split
+        req = gs % nreq
+        split = gs // nreq
+        hd = [(u % hg) * ht + r * (TILE // q) + row // q for r in range(rt)]
+        pos = row % q
+        t = req * q + pos
+        p0 = uniform(ld_i32(a["indptr"], req))
+        L = uniform(ld_i32(a["indptr"], req + 1)) - p0
+        # an empty request's p0 may be pages' end: its (unused) block-0 loads
+        # read pages[0] instead
+        p0 = (L > 0).select(p0, fx.Int32(0))
+        ct = ((L + per_req - 1) // per_req + TILE - 1) // TILE * TILE
+        c0 = split * ct
+        c1 = fx.min(L, c0 + ct)
+        ntile = (fx.max(c1 - c0, 0) + TILE - 1) // TILE
+        dt, dcc = dma_lane(lane)
+        # a new dict (the branch would carry a mutated ``c`` out as values)
+        d = dict(
+            c,
+            rt=rt,
+            # the line layout (a DMA 8 tokens' whole lines) at 1 row tile;
+            # at 2 its extra VGPRs spill in the fused K2 builds (a reload a
+            # block, whose vmcnt(0) waits out every DMA): 64 B chunks
+            lines=rt == 1,
+            lane=lane,
+            wave=wave,
+            dt=dt,
+            dcc=dcc,
+            dcol=16 * ((lane % 8) ^ sig(lane // 8, wave % 2)) if rt == 1 else 0,
+            ks_b=fx.Int32(fx.ptrtoint(c["ks"])),
+            xi_b=fx.Int32(fx.ptrtoint(c["xi"])),
+            qt=c["xp"] + rt * WAVES * 64 * 2,
+            p0=p0,
+            c0=c0,
+            last=fx.max(c1 - 1, 0),
+        )
+        # block 0's DMA before the Q loads (one in-order vmcnt: the indices'
+        # wait would hold for them too)
+        k_dma(d, 0, dma_offs(d, idx_load(d, 0)))
+        idx_dma(d, 1)
+        if const_expr(ws):
+            idx_dma(d, 2)
+        # the CTA's Q rows (16 r + row: [r][row][ENTRY] bf16) -> slot 1
+        # once, a DMA a 1 KB (a wave every WAVES-th; the last repeated):
+        # the waves' own loads of it would outweigh a K block 8 x
+        rowb = 2 * ENTRY
+        nq = rt * TILE * rowb // 1024
+        for i in range_constexpr((nq + WAVES - 1) // WAVES):
+            k = fx.min(wave + WAVES * i, nq - 1) * 64 + lane
+            lr = k // (rowb // 16)
+            qh = (u % hg) * ht + (lr // TILE) * (TILE // q) + (lr % TILE) // q
+            qt = req * q + (lr % TILE) % q
+            lds_dma(
+                a["qb"] + fx.Int64((qt * NH + qh) * rowb + k % (rowb // 16) * 16),
+                d["ks_b"] + (SLOT * 4 + (k - lane) * 16),
+            )
+        wait_dma()
+        # Q -> fp8 hi / lo (slot 1 after the bf16 rows), then Q^T's B operands
+        rows = rt * TILE
+        at = SLOT + rows * ENTRY // 2
+        q_split(d, rows, at)
+        gpu.barrier()
+        qv = [q_parts(d, r, at, rows) for r in range(rt)]
+        scale = fx.Int32(a["scale"]).bitcast(fx.Float32) * fx.Float32(LOG2E)
+        qsc = [scale * fx.ptr_load(c["xm"] + r * TILE + row) for r in range(rt)]
+        # every wave's reads before block 1's DMA into slot 1
+        gpu.barrier()
+        d.update(
+            qpos=(
+                [lines_at(row, wave % 2, 2 * g + k) // 4 for k in range(2)]
+                if rt == 1
+                else [None, None]
+            )
+            + [4 * kpos(row, g)],
+            vtr=[vtr_lane(lane, wave, i, rt == 1) for i in range(4)],
+            vb=d["ks_b"] + 2 * (wave - QKW) * WAVES * TILEW * 4,
+            q=qv,
+            qscale=qsc,
+            t=t,
+            h=hd,
+            split=req * per_req + split,
+            bound=L - q + pos,
+            nblk=(ntile + WAVES - 1) // WAVES,
+            nfull=fx.max(fx.min(c1, L - q + 1) - c0, 0) // BLK,
+            vscale=fx.Int32(a["vscale"]).bitcast(fx.Float32),
+        )
+        stamp(d, 0)
+        if const_expr(ws):
+            ws_run(d)
+        else:
+            run(d)
+        if const_expr("done" in d):
+            d["done"]((req * per_req + split) * hg + u % hg)
+
+
+_BUILDS: dict = {}
+
+
+def build(key: MlaAttnBuild):
+    if key in _BUILDS:
+        return _BUILDS[key]
+    s, q = key.tokens, key.qlen
+    check(s, q, key.rt)
+    Smem = smem(q, key.rt, key.ws)
+
+    keyed = key_tuple(key, _SOURCES)
+    name = f"k3_mono_mla_attn_s{s}_q{q}_r{key.rt or RTILES[q]}"
+    name += "_w" if key.wide else ""
+    name += "_ws" if wspec(key.rt or RTILES[q], key.ws) else ""
+
+    @flyc.kernel(name=name, known_block_size=[THREADS, 1, 1])
+    def mla_attn(
+        qb: Int64,
+        kv: Int64,
+        pages: Int64,
+        indptr: Int64,
+        part_o: Int64,
+        part_lse: Int64,
+        scale: Int32,
+        vscale: Int32,
+    ):
+        _ = keyed
+        tid = fx.thread_idx.x
+        bid = fx.block_idx.x
+        lds = fx.SharedAllocator().allocate(Smem).peek()
+        cta(
+            {
+                "S": s,
+                "Q": q,
+                "stamps": key.stamps,
+                "dbg": key.dbg,
+                "rt": key.rt,
+                "wide": key.wide,
+                "ws": key.ws,
+                "bid": bid,
+                "tid": tid,
+                "cm": 0,
+                "ks": lds.ks.ptr,
+                "xm": lds.xm.ptr,
+                "xp": lds.xp.ptr,
+                "xi": lds.xi.ptr,
+                "args": {
+                    "qb": qb,
+                    "kv": kv,
+                    "pages": pages,
+                    "indptr": indptr,
+                    "part_o": part_o,
+                    "part_lse": part_lse,
+                    "scale": scale,
+                    "vscale": vscale,
+                },
+            }
+        )
+
+    @flyc.jit
+    def launch(
+        qb: Int64,
+        kv: Int64,
+        pages: Int64,
+        indptr: Int64,
+        part_o: Int64,
+        part_lse: Int64,
+        scale: Int32,
+        vscale: Int32,
+        stream: fx.Stream = _STREAM,
+    ):
+        _ = keyed
+        mla_attn(qb, kv, pages, indptr, part_o, part_lse, scale, vscale).launch(
+            grid=(BLOCKS,), block=(THREADS,), stream=stream
+        )
+
+    _BUILDS[key] = launch
+    return launch
+
+
+def partial_shapes(key: MlaAttnBuild):
+    n = splits(key.qlen, key.rt)
+    return (n, key.tokens, NH, KVL), (n, key.tokens, NH)
+
+
+def f32_bits(x: float) -> int:
+    return int(torch.tensor(float(x), dtype=torch.float32).view(torch.int32))
+
+
+def mla_attn(
+    key: MlaAttnBuild, *, q, kv, pages, indptr, part_o, part_lse, scale, kv_scale=1.0
+):
+    """Host launcher. ``q`` [S, 12, 576] bf16; ``kv`` [tokens, 576] fp8 rows
+    (dense); ``pages`` int32 flat token indices, ``indptr`` [requests + 1]
+    int32 (a request's KV ``pages[indptr[r]:indptr[r + 1]]``); ``part_o`` /
+    ``part_lse`` of ``partial_shapes(key)``, bf16 / f32. ``scale`` the softmax
+    scale, ``kv_scale`` the cache's dequant (folded into QK and o)."""
+    so, sl = partial_shapes(key)
+    assert part_o.shape == so and part_o.dtype == torch.bfloat16
+    assert part_lse.shape == sl and part_lse.dtype == torch.float32
+    assert q.shape == (key.tokens, NH, ENTRY), q.shape
+    a = host_args(
+        q=q, kv=kv, pages=pages, indptr=indptr, scale=scale, kv_scale=kv_scale
+    )
+    build(replace(key, wide=key.wide or wide_kv(kv)))(
+        *a[:4],
+        part_o.data_ptr(),
+        part_lse.data_ptr(),
+        *a[4:],
+        stream=torch.cuda.current_stream(),
+    )
+
+
+def wide_kv(kv) -> bool:
+    """The build's ``wide`` for KV pool ``kv``: a DMA's 32-bit offset off the
+    pool's base (a token row's end at most) won't reach past 4 GiB."""
+    return kv.numel() > 2**32
+
+
+def host_args(*, q, kv, pages, indptr, scale, kv_scale=1.0):
+    """(qb, kv, pages, indptr, scale, vscale) of the kernel (``mla_attn``'s
+    inputs but the partials; K2's fused M1c's)."""
+    assert (
+        q.dtype == torch.bfloat16 and q.is_contiguous() and q.shape[1:] == (NH, ENTRY)
+    )
+    assert kv.element_size() == 1 and kv.stride(-1) == 1 and kv.size(-1) == ENTRY
+    assert kv.stride(0) == ENTRY
+    assert pages.dtype == torch.int32 and indptr.dtype == torch.int32
+    return [
+        q.data_ptr(),
+        kv.data_ptr(),
+        pages.data_ptr(),
+        indptr.data_ptr(),
+        f32_bits(scale * kv_scale),
+        f32_bits(kv_scale),
+    ]
diff --git a/vllm/models/kimi_k3/amd/specialized_decode/mono/attention/mla_attn_geom.py b/vllm/models/kimi_k3/amd/specialized_decode/mono/attention/mla_attn_geom.py
new file mode 100644
index 000000000000..99445569e11b
--- /dev/null
+++ b/vllm/models/kimi_k3/amd/specialized_decode/mono/attention/mla_attn_geom.py
@@ -0,0 +1,33 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""M1c ``mla_attn``'s launch geometry, without the kernel's FlyDSL imports: the
+runtime's config checks (``runtime.m1c_serves``) size M1c's KV splits on any
+platform."""
+
+from vllm.models.kimi_k3.amd.specialized_decode.mono.common.plan import BLOCKS
+
+NH = 12  # local heads at TP8 (``mla.NH``)
+XCDS = 8
+TILE = 16  # tokens a wave's tile
+RTILES = {4: 1, 8: 2}  # row tiles a CTA by qlen (3 head groups either way)
+
+
+def splits(qlen: int, rt: int = 0) -> int:
+    """KV splits of the launch (every request's together): the head groups'
+    CTAs an XCD hold a whole number of splits."""
+    return XCDS * geometry(qlen, rt)[2]
+
+
+def geometry(qlen: int, rt: int = 0):
+    """(heads a CTA, head groups, splits an XCD); a CTA ``rt`` row tiles of 16
+    (``RTILES`` by default)."""
+    ht = TILE * (rt or RTILES[qlen]) // qlen
+    hg = NH // ht
+    return ht, hg, BLOCKS // XCDS // hg
+
+
+def ctas(qlen: int, rt: int = 0) -> int:
+    """The working CTAs (0 .. ctas): a done flag each when fused (``cta``'s
+    ``done``), at (split, head group) ``split * hg + group``."""
+    _, hg, spx = geometry(qlen, rt)
+    return XCDS * spx * hg
diff --git a/vllm/models/kimi_k3/amd/specialized_decode/mono/attention/mla_back.py b/vllm/models/kimi_k3/amd/specialized_decode/mono/attention/mla_back.py
new file mode 100644
index 000000000000..fbcd857de92b
--- /dev/null
+++ b/vllm/models/kimi_k3/amd/specialized_decode/mono/attention/mla_back.py
@@ -0,0 +1,760 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""M1b ``mla_post``: one Kimi-K3 MLA layer from AITER's decode output to
+o_proj's input, at TP8 + DCP8 (the DCP group the TP group) or TP8 + DCP1.
+
+AITER leaves every rank the partial attention of its KV shard for all the
+gathered heads (``o`` [S, ranks x 12, 512] bf16 and ``lse`` [S, ranks x 12]
+f32, strided views). As ``dcp_a2a_lse_reduce`` -> ``_v_up_proj`` -> the gate:
+
+    push (a wave a (token, head), over every CTA)
+                        o / lse of rank p's heads -> rank p's CMB: the o row
+                        untagged, then the LSE pair as its flag (an LSE of
+                        -inf for a row with no local KV or past the batch, as
+                        ``_dcp_a2a_pack_send_kernel`` masks it)
+    combine (a CTA a (head, 4 tokens))
+                        the ranks' partials in rank order, LSE-weighted in
+                        fp32 -> bf16 latents in LDS
+    W_UV (a wave 16 of the head's 128 rows)
+                        the latents -> MXFP4 (aiter's even scale), four FP4
+                        MFMAs against W_V -> bf16, times the gate -> ``core``
+
+At DCP1 (``ranks`` 1) ``o`` [S, 12, 512] is already the attention's output, as
+vLLM hands it to ``_v_up_proj``: no push, no LSE; the combine stage loads it.
+With ``splits`` (M1c ``mla_attn``'s partials: ``o`` / ``lse`` split 0 of dense
+[splits, S, 12, 512] bf16 / [splits, S, 12] f32) the combine stage merges a
+token's request's splits, LSE-weighted.
+
+K2 (``layer``, ``K2Build.m1b_qlen``) runs these stages in its own launch.
+"""
+
+from dataclasses import dataclass
+
+import flydsl.compiler as flyc
+import flydsl.expr as fx
+import torch
+from aiter.ops.flydsl.kernels import buffer_ops as bo
+from flydsl.expr import const_expr, gpu, range_constexpr, rocdl
+from flydsl.expr.typing import Int32, Int64, T
+
+from vllm.models.kimi_k3.amd.specialized_decode.mono.attention.kda import (
+    LOG2E,
+    MAX_S,
+    bf16r,
+    step_tag,
+)
+from vllm.models.kimi_k3.amd.specialized_decode.mono.attention.mla import (
+    GROWS,
+    KVL,
+    NH,
+    even_code,
+    pk_fp4_bf16,
+)
+from vllm.models.kimi_k3.amd.specialized_decode.mono.common.debug import (
+    region_ids,
+    stamp,
+    stamp_begin,
+    stamp_flush,
+)
+from vllm.models.kimi_k3.amd.specialized_decode.mono.common.mx import (
+    FP4,
+    mfma_scaled,
+    pow2,
+)
+from vllm.models.kimi_k3.amd.specialized_decode.mono.common.ops import (
+    CM_SYS,
+    bf16_pair,
+    bf_hi,
+    bf_lo,
+    butterfly,
+    hw_exp2,
+    hw_rcp,
+    ld_i32,
+    load_ptr64,
+    rsrc,
+    traced,
+    uniform,
+)
+from vllm.models.kimi_k3.amd.specialized_decode.mono.common.plan import (
+    BLOCKS,
+    THREADS,
+    WAVES,
+    KernelAbi,
+    first_task,
+    key_tuple,
+    source_digest,
+)
+from vllm.models.kimi_k3.amd.specialized_decode.mono.common.ranks import peer_bases
+from vllm.models.kimi_k3.amd.specialized_decode.mono.common.sync import Mailbox, preg
+
+VD = 128  # v_head_dim
+CW = KVL // 2 + 2  # CMB pairs a (source, token, head): 256 o words, the LSE, pad
+OU = KVL // 4  # the LSE's pair in a CMB row
+TT = 4  # tokens a combine task
+KS = KVL // 128  # W_UV's 128-k MFMA steps
+LROW = KVL // 2 + 4  # an LDS latent row (words, padded)
+WMAX = 128  # M1c splits a request at most (an LDS weight row)
+CB = 8  # combine_splits' o loads in flight a thread (1 with M1c in the launch)
+# timeline points: start, push, combine, W_UV
+TL_POINTS = 4
+
+assert VD == WAVES * 16 and THREADS == TT * (KVL // 4) and KVL // 2 == 64 * 4
+
+_STREAM = fx.Stream(None)
+REGIONS = ("cmb",)
+_SOURCES = source_digest(".")
+
+
+@dataclass(frozen=True)
+class MlaPostBuild:
+    tokens: int  # S
+    qlen: int  # tokens a request (1 + speculative)
+    ranks: int = 8  # 1: DCP1, no peer memory
+    peer_off: int = 0  # CMB's byte offset in a rank's peer buffer
+    timeline: bool = False
+    splits: int = 0  # DCP1: M1c's KV splits (0: ``o`` the attention's output)
+    ctok: int = 0  # tokens a combine task (0: ``ctok``'s)
+
+
+ABI = KernelAbi(
+    (
+        "o",
+        "o_sb",
+        "o_sh",
+        "lse",
+        "lse_sb",
+        "lse_sh",
+        "seq_lens",
+        "qsl",
+        "w_v",
+        "wv_sh",
+        "wv_sn",
+        "w_vs",
+        "vs_sh",
+        "vs_sn",
+        "vs_sk",
+        "gate",
+        "core",
+        "peers",
+        "rank",
+        "layer",
+        "epoch",
+        "health",
+        "alarm",
+        "tl",
+    )
+)
+
+
+def peer_pairs(s: int, ranks: int) -> int:
+    return ranks * s * NH * CW
+
+
+def peer_bytes(s: int, ranks: int) -> int:
+    return peer_pairs(s, ranks) * 8
+
+
+def cmb_pair(c, src, t, h):
+    return ((src * c["S"] + t) * NH + h) * CW
+
+
+# ---------------------------------------------------------------- push
+@traced
+def stage_push(c):
+    """Rank p's heads of this rank's partials -> rank p's CMB: a wave a (token,
+    head), the waves over every CTA first. The o row goes untagged (a lane 16 B
+    a destination, half the bytes of tagged pairs: the push is xGMI-bound), then
+    once the wave's stores drained, the LSE pair, its tag the row's flag."""
+    s, q, bid, a = c["S"], c["Q"], c["bid"], c["args"]
+    lane, wave = c["lane"], c["wave"]
+    nr = s * NH
+    nw = BLOCKS * WAVES
+    for i in range_constexpr((nr + nw - 1) // nw):
+        g = (wave + WAVES * i) * BLOCKS + bid
+        if g < nr:
+            h = g % NH
+            t = g // NH
+            # the request's local KV empty, or a row past the batch: LSE -inf
+            rows = ld_i32(a["qsl"], s // q)
+            empty = (ld_i32(a["seq_lens"], t // q) == 0) | (t >= rows)
+            at = cmb_pair(c, c["rank"], t, h)
+            lses = []
+            for p in range_constexpr(c["R"]):
+                gh = p * NH + h
+                w = bo.buffer_load(
+                    rsrc(a["o"]),
+                    (t * a["o_sb"] + gh * a["o_sh"]) // 2 + 4 * lane,
+                    vec_width=4,
+                    dtype=T.i32,
+                )
+                bo.buffer_store(
+                    w, rsrc(c["cmb"][p].value), at * 2 + 4 * lane, cache_modifier=CM_SYS
+                )
+                lse = fx.Int32(
+                    bo.buffer_load(
+                        rsrc(a["lse"]),
+                        t * a["lse_sb"] + gh * a["lse_sh"],
+                        vec_width=1,
+                        dtype=T.i32,
+                    )
+                )
+                lses.append(empty.select(fx.Int32(-8388608), lse))  # -inf
+            # the wave's row stores acknowledged before its flags (a wave's
+            # own stores: no barrier)
+            rocdl.s_waitcnt(vmcnt=0)
+            for p in range_constexpr(c["R"]):
+                if lane == p:
+                    c["put_words"](c["cmb"][p], at + 2 * OU, [lses[p], fx.Int32(0)])
+
+
+# ---------------------------------------------------------------- combine
+def sane(lse):
+    """NaN / +inf -> -inf (``_dcp_a2a_unpack_combine_kernel``)."""
+    bad = (lse != lse) | (lse == fx.Float32(float("inf")))
+    return bad.select(fx.Float32(float("-inf")), lse)
+
+
+@traced
+def stage_combine(c, h, t0):
+    """Tokens t0 .. t0 + 4 of head h: thread (token tid / 128, latents
+    4 (tid % 128) ..) combines the ranks' partials into the LDS row."""
+    s, R, tid = c["S"], c["R"], c["tid"]
+    tt = tid // (KVL // 4)
+    cc = tid % (KVL // 4)
+    t = fx.min(t0 + tt, s - 1)
+    if const_expr(R == 1 and c["splits"] > 0):
+        combine_splits(c, t0, h)
+    if const_expr(R == 1 and c["splits"] == 0):
+        a = c["args"]
+        o = fx.Vector(
+            bo.buffer_load(
+                rsrc(a["o"]),
+                (t * a["o_sb"] + h * a["o_sh"]) // 2 + 2 * cc,
+                vec_width=2,
+                dtype=T.i32,
+            )
+        )
+        live = t0 + tt < s
+        base = c["lat"] + tt * LROW + 2 * cc
+        fx.ptr_store(live.select(o[0], fx.Int32(0)), base)
+        fx.ptr_store(live.select(o[1], fx.Int32(0)), base + 1)
+    if const_expr(R > 1):
+        combine_ranks(c, h, t0, tt, cc, t)
+    gpu.barrier()
+
+
+@traced
+def combine_splits(c, t0, h):
+    """``stage_combine`` at DCP1 over M1c's partials: token t's request's splits
+    (``splits / requests`` of them, request-major), LSE-weighted. Wave tt < tc
+    (the task's tokens) turns token t0 + tt's LSEs into weights (LDS ``cw``);
+    then a token's WAVES / tc waves (parts) sum the splits part, part + parts,
+    .. a thread 8 latents of a token (16 B loads, CB in flight), the parts'
+    sums halved through ``cw`` (the upper half's handed over) to part 0."""
+    s, q, n, a, tid, tc = c["S"], c["Q"], c["splits"], c["args"], c["tid"], c["TC"]
+    per = n // (s // q)
+    sp0 = t0 // q * per
+    wv, lane = tid // 64, tid % 64
+    if const_expr("ardy" in c):
+        # M1c in this launch (its VGPRs at the limit): the wave index scalar
+        wv = uniform(wv)
+        # its CTAs' flags (split, head group of M1c's ``ht`` heads), a thread each
+        ht = c["m1c_ht"]
+        if tid < per:
+            c["poll"]([(c["ardy"], (sp0 + tid) * (NH // ht) + h // ht, 1)])
+        gpu.barrier()
+    if wv < tc:
+        t = fx.min(t0 + wv, s - 1)
+        lse = []
+        for k in range_constexpr((per + 63) // 64):
+            i = lane + 64 * k
+            v = fx.Int32(
+                bo.buffer_load(
+                    rsrc(a["lse"]),
+                    ((sp0 + fx.min(i, per - 1)) * s + t) * NH + h,
+                    vec_width=1,
+                    dtype=T.i32,
+                )
+            ).bitcast(fx.Float32)
+            lse.append((i < per).select(sane(v), fx.Float32(float("-inf"))))
+        m = lse[0]
+        for k in range_constexpr(1, len(lse)):
+            m = fx.max(m, lse[k])
+        m = butterfly(m, (32, 16, 8, 4, 2, 1), fx.max)
+        m = (m == fx.Float32(float("-inf"))).select(fx.Float32(0.0), m)
+        e = [hw_exp2((x - m) * fx.Float32(LOG2E)) for x in lse]
+        tot = e[0]
+        for k in range_constexpr(1, len(e)):
+            tot = tot + e[k]
+        tot = butterfly(tot, (32, 16, 8, 4, 2, 1))
+        # a padded request (no KV): every split empty, its rows 0
+        inv = (tot == fx.Float32(0.0)).select(fx.Float32(0.0), hw_rcp(tot))
+        for k in range_constexpr(len(e)):
+            fx.ptr_store(e[k] * inv, c["cw"] + wv * WMAX + lane + 64 * k)
+    gpu.barrier()
+    parts = WAVES // tc
+    part, tt = wv // tc, wv % tc
+    t = fx.min(t0 + tt, s - 1)
+    nj = per // parts  # splits a part
+    src = (((sp0 + part) * s + t) * NH + h) * (KVL // 2) + 4 * lane
+    wrow = c["cw"] + tt * WMAX + part
+    stride = parts * s * NH * KVL * 2  # bytes between a part's splits
+    acc = [fx.Float32(0.0)] * 8
+
+    def batch(acc, j0, nb):
+        os_ = [
+            fx.Vector(
+                bo.buffer_load(
+                    rsrc(a["o"]),
+                    src,
+                    vec_width=4,
+                    dtype=T.i32,
+                    soffset_bytes=(j0 + j) * stride,
+                )
+            )
+            for j in range(nb)
+        ]
+        # an empty split: weight 0, its o 0
+        wts = [fx.ptr_load(wrow + (j0 + j) * parts) for j in range(nb)]
+        acc = list(acc)
+        for j in range_constexpr(nb):
+            o, wt = os_[j], wts[j]
+            for k in range_constexpr(4):
+                acc[2 * k] = acc[2 * k] + bf_lo(o[k]) * wt
+                acc[2 * k + 1] = acc[2 * k + 1] + bf_hi(o[k]) * wt
+        return acc
+
+    if const_expr("ardy" in c):
+        # with M1c: unrolled, a load a batch (a loop, or wider batches, spill
+        # M1c's Q fragment in its loop)
+        for j0 in range_constexpr(nj):
+            acc = batch(acc, j0, 1)
+    else:
+        # CB loads in flight a batch (a loop over the batches: the scheduler
+        # hoists no further), the rest unrolled
+        jb = fx.Int32(0)
+        while jb < nj // CB * CB:
+            acc = batch(acc, jb, CB)
+            jb = jb + CB
+        if const_expr(nj % CB):
+            acc = batch(acc, nj // CB * CB, nj % CB)
+    # parts p >= hf hand their sums to p - hf, halving
+    for hf in [parts >> k for k in range(1, parts.bit_length())]:
+        red = c["cw"] + TT * WMAX + (tt * 64 + lane) * 8
+        if (part >= hf) & (part < 2 * hf):
+            at = red + (part - hf) * (tc * 64 * 8)
+            for k in range_constexpr(2):
+                fx.ptr_store(
+                    fx.Vector.from_elements(acc[4 * k : 4 * k + 4], fx.Float32),
+                    at + 4 * k,
+                )
+        gpu.barrier()
+        if const_expr(hf > 1):
+            at = red + fx.min(part, hf - 1) * (tc * 64 * 8)
+            acc = [
+                (part < hf).select(acc[i] + fx.ptr_load(at + i), acc[i])
+                for i in range(8)
+            ]
+            gpu.barrier()  # read before the next half's stores
+        else:
+            acc = [acc[i] + fx.ptr_load(red + i) for i in range(8)]
+    if part == 0:
+        live = t0 + tt < s
+        for k in range_constexpr(4):
+            fx.ptr_store(
+                live.select(
+                    bf16_pair(acc[2 * k], acc[2 * k + 1]).bitcast(fx.Int32),
+                    fx.Int32(0),
+                ),
+                c["lat"] + tt * LROW + 4 * lane + k,
+            )
+
+
+@traced
+def combine_ranks(c, h, t0, tt, cc, t):
+    """``stage_combine`` at DCP8: the ranks' CMB partials, LSE-weighted."""
+    s, R = c["S"], c["R"]
+    specs = [(c["cmb_own"], cmb_pair(c, src, t, h) + 2 * OU, 1) for src in range(R)]
+    got = c["poll"](specs)
+    lse = [sane(got[r][0].bitcast(fx.Float32)) for r in range(R)]
+    # each source's row, untagged: its LSE flag above announced it
+    rows = [
+        fx.Vector(
+            bo.buffer_load(
+                rsrc(c["cmb_own"].value),
+                cmb_pair(c, src, t, h) * 2 + 2 * cc,
+                vec_width=2,
+                dtype=T.i32,
+                cache_modifier=CM_SYS,
+            )
+        )
+        for src in range(R)
+    ]
+    m = lse[0]
+    for r in range_constexpr(1, R):
+        m = fx.max(m, lse[r])
+    m = (m == fx.Float32(float("-inf"))).select(fx.Float32(0.0), m)
+    e = [hw_exp2((lse[r] - m) * fx.Float32(LOG2E)) for r in range(R)]
+    tot = e[0]
+    for r in range_constexpr(1, R):
+        tot = tot + e[r]
+    inv = hw_rcp(tot)
+    acc = [fx.Float32(0.0)] * 4
+    for r in range_constexpr(R):
+        wt = (e[r] == fx.Float32(0.0)).select(fx.Float32(0.0), e[r] * inv)
+        o = rows[r]
+        pv = [bf_lo(o[0]), bf_hi(o[0]), bf_lo(o[1]), bf_hi(o[1])]
+        for j in range_constexpr(4):
+            pj = (wt == fx.Float32(0.0)).select(fx.Float32(0.0), pv[j])
+            acc[j] = acc[j] + pj * wt
+    live = t0 + tt < s
+    w0 = bf16_pair(acc[0], acc[1]).bitcast(fx.Int32)
+    w1 = bf16_pair(acc[2], acc[3]).bitcast(fx.Int32)
+    w0 = live.select(w0, fx.Int32(0))
+    w1 = live.select(w1, fx.Int32(0))
+    base = c["lat"] + tt * LROW + 2 * cc
+    fx.ptr_store(w0, base)
+    fx.ptr_store(w1, base + 1)
+
+
+# ---------------------------------------------------------------- W_UV
+def wv_loads(c, h):
+    """This lane's W_V operands for the wave's 16 rows: (16 B, scale) a step."""
+    a, lane, wave = c["args"], c["lane"], c["wave"]
+    v = wave * 16 + lane % 16
+    kb = lane // 16
+    ops = []
+    for st in range_constexpr(KS):
+        wo = h * a["wv_sh"] + v * a["wv_sn"] + st * 64 + kb * 16
+        wv = fx.Vector(
+            bo.buffer_load(rsrc(a["w_v"]), wo // 4, vec_width=4, dtype=T.i32)
+        )
+        so = h * a["vs_sh"] + v * a["vs_sn"] + (st * 4 + kb) * a["vs_sk"]
+        sw = fx.Int32(
+            bo.buffer_load(rsrc(a["w_vs"]), so // 4, vec_width=1, dtype=T.i32)
+        )
+        ops.append((wv, (sw >> ((so % 4) * 8)) & 0xFF))
+    return ops
+
+
+@traced
+def stage_uv(c, h, t0, ops):
+    """The LDS latents of tokens t0 .. t0 + 4 against W_V[h]'s 16 rows of this
+    wave -> bf16, times the gate -> ``core``."""
+    s, lane, wave, a, tc = c["S"], c["lane"], c["wave"], c["args"], c["TC"]
+    tl = lane % 16
+    kb = lane // 16
+    row = c["lat"] + fx.min(tl, tc - 1) * LROW
+    acc = fx.Vector.filled(4, 0.0, fx.Float32)
+    for st in range_constexpr(KS):
+        ws = [
+            (tl < tc).select(fx.ptr_load(row + st * 64 + kb * 16 + i), fx.Int32(0))
+            for i in range(16)
+        ]
+        am = fx.Int32(0)
+        for w in ws:
+            am = fx.max(am, w & 0x7FFF)
+            am = fx.max(am, (w >> 16) & 0x7FFF)
+        code = fx.max(even_code(am), fx.Int32(1))
+        scale = pow2(code)
+        xs = []
+        for q in range_constexpr(4):
+            pk = fx.Int32(0)
+            for j in range_constexpr(4):
+                pk = pk_fp4_bf16(pk, ws[4 * q + j], scale, j)
+            xs.append(pk)
+        xv = fx.Vector.from_elements(xs, fx.Int32)
+        wv, sa = ops[st]
+        acc = mfma_scaled(wv, xv, acc, sa, code, FP4, FP4, 0, 0)
+    t = t0 + tl
+    v0 = wave * 16 + 4 * kb
+    col = h * VD + v0
+    tg = fx.min(t, s - 1)
+    gw = fx.Vector(
+        bo.buffer_load(rsrc(a["gate"]), (tg * GROWS + col) // 2, vec_width=2,
+                       dtype=T.i32)
+    )  # fmt: skip
+    g = [bf_lo(gw[0]), bf_hi(gw[0]), bf_lo(gw[1]), bf_hi(gw[1])]
+    out = [bf16r(acc[j]) * g[j] for j in range(4)]
+    if (tl < tc) & (t < s):
+        # fused into K2 (``c["core_cm"]``): read by other CTAs, device scope
+        bo.buffer_store(
+            fx.Vector.from_elements(
+                [bf16_pair(out[0], out[1]), bf16_pair(out[2], out[3])], fx.Float32
+            ),
+            rsrc(a["core"]),
+            (t * GROWS + col) // 2,
+            **({"cache_modifier": c["core_cm"]} if "core_cm" in c else {}),
+        )
+    gpu.barrier()
+
+
+def ctok(s: int, q: int, splits: int, ctas: int = BLOCKS) -> int:
+    """Tokens a combine task: TT, or with M1c's splits the fewest of 1, 2, 4
+    whose parts (WAVES / it a token) divide a request's splits and whose tasks
+    fit ``ctas`` (a CTA's reads, a part's splits a thread, over more CTAs)."""
+    if not splits:
+        return TT
+    per = splits // (s // q)
+    ok = [tc for tc in (1, 2, 4) if per % (WAVES // tc) == 0]
+    return next((tc for tc in ok if tasks(s, tc) <= ctas), ok[-1])
+
+
+def tasks(s: int, tc: int = TT) -> int:
+    """Combine tasks: a (head, ``tc`` tokens), head-minor (task h + 12 tile)."""
+    return NH * ((s + tc - 1) // tc)
+
+
+@traced
+def combine_tasks(c, ctas=BLOCKS, done=None):
+    """The combine + W_UV tasks over CTAs 0 .. ``ctas``; ``done(task)`` after
+    each task's ``core`` rows are stored (K2's fused M1b: their flag)."""
+    s, bid, tc = c["S"], c["bid"], c["TC"]
+    for task in range(first_task(bid, 0), tasks(s, tc), ctas):
+        h = task % NH
+        t0 = task // NH * tc
+        ops = wv_loads(c, h)
+        stage_combine(c, h, t0)
+        stage_uv(c, h, t0, ops)
+        if const_expr(done is not None):
+            done(task)
+
+
+# ---------------------------------------------------------------- kernel
+_BUILDS: dict = {}
+
+
+def build(key: MlaPostBuild):
+    if key in _BUILDS:
+        return _BUILDS[key]
+    s, R = key.tokens, key.ranks
+    assert 1 <= s <= MAX_S and s % key.qlen == 0 and 1 <= R <= 8
+
+    @fx.struct
+    class Smem:
+        lat: fx.Array[fx.Int32, TT * LROW, 16]
+        cw: fx.Array[fx.Float32, TT * (WMAX + KVL) if key.splits else 1, 16]
+        tls: fx.Array[fx.Int64, TL_POINTS, 16]
+
+    keyed = key_tuple(key, _SOURCES)
+    name = f"k3_mono_mla_post_s{s}_q{key.qlen}_r{R}_t{int(key.timeline)}"
+    tc = key.ctok or ctok(s, key.qlen, key.splits)
+    assert tc in (1, 2, 4) and (
+        not key.splits or key.splits // (s // key.qlen) % (WAVES // tc) == 0
+    )
+    name += f"_sp{key.splits}_c{tc}" if key.splits else ""
+    assert not key.splits or (R == 1 and key.splits % (s // key.qlen) == 0)
+    assert key.splits // (s // key.qlen) % 2 == 0
+    assert key.splits // (s // key.qlen) <= WMAX
+
+    @flyc.kernel(name=name, known_block_size=[THREADS, 1, 1])
+    def mla_post(
+        o: Int64,
+        o_sb: Int32,
+        o_sh: Int32,
+        lse: Int64,
+        lse_sb: Int32,
+        lse_sh: Int32,
+        seq_lens: Int64,
+        qsl: Int64,
+        w_v: Int64,
+        wv_sh: Int32,
+        wv_sn: Int32,
+        w_vs: Int64,
+        vs_sh: Int32,
+        vs_sn: Int32,
+        vs_sk: Int32,
+        gate: Int64,
+        core: Int64,
+        peers: Int64,
+        rank: Int32,
+        layer: Int32,
+        epoch: Int64,
+        health: Int64,
+        alarm: Int64,
+        tl: Int64,
+    ):
+        _ = keyed
+        tid = fx.thread_idx.x
+        bid = fx.block_idx.x
+        lds = fx.SharedAllocator().allocate(Smem).peek()
+        mb = Mailbox(step_tag(epoch, layer) - 1, health, alarm, region_ids(REGIONS))
+        # DCP1 reads no peer memory (``peers`` may be 0)
+        bases = peer_bases(peers, R) if R > 1 else []
+        own = load_ptr64(peers, rank) if R > 1 else peers
+        c = {
+            "S": s,
+            "Q": key.qlen,
+            "R": R,
+            "splits": key.splits,
+            "TC": tc,
+            "tid": tid,
+            "bid": bid,
+            "lane": tid % 64,
+            "wave": tid // 64,
+            "rank": uniform(rank),
+            "put_words": mb.put_words,
+            "poll": mb.poll,
+            "lat": lds.lat.ptr,
+            "cw": lds.cw.ptr,
+            "cmb": [preg(b, key.peer_off, "cmb") for b in bases],
+            "cmb_own": preg(own, key.peer_off, "cmb"),
+            "args": {
+                "o": o,
+                "o_sb": o_sb,
+                "o_sh": o_sh,
+                "lse": lse,
+                "lse_sb": lse_sb,
+                "lse_sh": lse_sh,
+                "seq_lens": seq_lens,
+                "qsl": qsl,
+                "w_v": w_v,
+                "wv_sh": wv_sh,
+                "wv_sn": wv_sn,
+                "w_vs": w_vs,
+                "vs_sh": vs_sh,
+                "vs_sn": vs_sn,
+                "vs_sk": vs_sk,
+                "gate": gate,
+                "core": core,
+            },
+        }
+        on, tls = key.timeline, lds.tls.ptr
+        stamp_begin(on, tls, tid, TL_POINTS)
+        if const_expr(R > 1):
+            stage_push(c)
+        stamp(on, tls, tid, 1)
+        combine_tasks(c)
+        stamp(on, tls, tid, 3)
+        stamp_flush(on, tls, tl, tid, bid, TL_POINTS)
+
+    @flyc.jit
+    def launch(
+        o: Int64,
+        o_sb: Int32,
+        o_sh: Int32,
+        lse: Int64,
+        lse_sb: Int32,
+        lse_sh: Int32,
+        seq_lens: Int64,
+        qsl: Int64,
+        w_v: Int64,
+        wv_sh: Int32,
+        wv_sn: Int32,
+        w_vs: Int64,
+        vs_sh: Int32,
+        vs_sn: Int32,
+        vs_sk: Int32,
+        gate: Int64,
+        core: Int64,
+        peers: Int64,
+        rank: Int32,
+        layer: Int32,
+        epoch: Int64,
+        health: Int64,
+        alarm: Int64,
+        tl: Int64,
+        stream: fx.Stream = _STREAM,
+    ):
+        _ = keyed
+        mla_post(
+            o,
+            o_sb,
+            o_sh,
+            lse,
+            lse_sb,
+            lse_sh,
+            seq_lens,
+            qsl,
+            w_v,
+            wv_sh,
+            wv_sn,
+            w_vs,
+            vs_sh,
+            vs_sn,
+            vs_sk,
+            gate,
+            core,
+            peers,
+            rank,
+            layer,
+            epoch,
+            health,
+            alarm,
+            tl,
+        ).launch(grid=(BLOCKS,), block=(THREADS,), stream=stream)
+
+    ABI.check(mla_post, launch)
+    _BUILDS[key] = launch
+    return launch
+
+
+def host_args(ranks, o, lse, seq_lens, qsl, w_v, w_v_scale, gate) -> list:
+    """ABI's ``o`` .. ``gate`` (``mla_post``'s, or K2's fused M1b's)."""
+    assert o.stride(-1) == 1 and w_v.stride(-1) == 1
+    assert o.dtype == torch.bfloat16
+    assert lse is not None or ranks == 1  # DCP1: M1c's partials carry one
+    assert lse is None or lse.dtype == torch.float32
+    assert seq_lens.dtype == torch.int32 and qsl.dtype == torch.int32
+    assert gate.is_contiguous()
+    return [
+        o.data_ptr(),
+        o.stride(0),
+        o.stride(1),
+        0 if lse is None else lse.data_ptr(),
+        0 if lse is None else lse.stride(0),
+        0 if lse is None else lse.stride(1),
+        seq_lens.data_ptr(),
+        qsl.data_ptr(),
+        w_v.data_ptr(),
+        w_v.stride(0),
+        w_v.stride(1),
+        w_v_scale.data_ptr(),
+        w_v_scale.stride(0),
+        w_v_scale.stride(1),
+        w_v_scale.stride(2),
+        gate.data_ptr(),
+    ]
+
+
+def mla_post(
+    key: MlaPostBuild,
+    *,
+    o,
+    lse,
+    seq_lens,
+    qsl,
+    w_v,
+    w_v_scale,
+    gate,
+    core,
+    peers,
+    rank,
+    layer,
+    epoch,
+    health,
+    alarm,
+    tl=None,
+    scratch=None,
+):
+    """Host launcher. ``o`` [S, >= ranks x 12, 512] bf16 and ``lse`` [S, >=
+    ranks x 12] f32, rows and heads strided (ranks 1: no ``lse``, ``peers``
+    may be None); ``seq_lens`` [requests] and ``qsl``
+    [requests + 1] int32 (the local KV lengths, the query starts); ``w_v`` (12,
+    128, 256) u8 MXFP4 and ``w_v_scale`` (12, 128, 16) e8m0, any strides but
+    W_V's last; ``gate`` / ``core`` [S, 1536] bf16, dense."""
+    del scratch
+    assert core.is_contiguous()
+    f = build(key)
+    f(
+        *host_args(key.ranks, o, lse, seq_lens, qsl, w_v, w_v_scale, gate),
+        core.data_ptr(),
+        0 if peers is None else peers.data_ptr(),
+        rank,
+        layer,
+        epoch.data_ptr(),
+        health.data_ptr(),
+        alarm,
+        0 if tl is None else tl.data_ptr(),
+        stream=torch.cuda.current_stream(),
+    )
diff --git a/vllm/models/kimi_k3/amd/specialized_decode/mono/common/debug.py b/vllm/models/kimi_k3/amd/specialized_decode/mono/common/debug.py
index 18cdcff38bdc..a967c3566afc 100644
--- a/vllm/models/kimi_k3/amd/specialized_decode/mono/common/debug.py
+++ b/vllm/models/kimi_k3/amd/specialized_decode/mono/common/debug.py
@@ -2,17 +2,61 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 # Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at a526f0d, MIT License,
 # Copyright (C) 2026, Advanced Micro Devices, Inc.:
+# atom/mono/device/stamps.py
 # atom/mono/runtime/debug.py
-"""Wait records: which mailbox region a CTA's given-up wait was on."""
+"""A timeline build's per-CTA phase stamps (debug only; ``ATOM_MONO_TIMELINE``,
+saved by ``atom.mono.runtime.timeline``): the s_memrealtime (100 MHz) at which
+the CTA passes point k, kept in LDS until the kernel ends -- a global store per
+stamp would sit on every later ``s_waitcnt vmcnt(0)``.
+
+``on`` is the build's compile-time switch: off, every call traces to nothing.
+"""
 
 import zlib
 
+import flydsl.expr as fx
 import torch
+from flydsl.expr import const_expr
 
 from vllm.models.kimi_k3.amd.specialized_decode.mono.abi import RECORD_WORDS
+from vllm.models.kimi_k3.amd.specialized_decode.mono.common.ops import (
+    memrealtime,
+    traced,
+)
 from vllm.models.kimi_k3.amd.specialized_decode.mono.common.plan import BLOCKS
 
 
+@traced
+def stamp_begin(on, tls, tid, points):
+    """Zero the LDS record, then point 0."""
+    if const_expr(on):  # noqa: SIM102
+        if tid < points:  # the stamping thread's wave: ordered
+            fx.ptr_store(fx.Int64(0), tls + tid)
+    stamp(on, tls, tid, 0)
+
+
+@traced
+def stamp(on, tls, tid, k):
+    # compile-time gate outside, traced condition inside: they cannot be one `and`
+    if const_expr(on):  # noqa: SIM102
+        if tid == 0:
+            fx.ptr_store(memrealtime(), tls + k)
+
+
+@traced
+def stamp_flush(on, tls, tl, tid, bid, points):
+    """The CTA's record -> ``tl`` ``[CTA][points]`` int64."""
+    if const_expr(on):  # noqa: SIM102
+        if tid < points:
+            fx.generic_store(
+                fx.inttoptr(
+                    fx.PointerType.get(fx.Int64.ir_type, fx.AddressSpace.Global, 8),
+                    tl + fx.Int64((bid * points + tid) * 8),
+                ),
+                fx.ptr_load(tls + tid),
+            )
+
+
 def region_id(name: str) -> int:
     """A mailbox region's id in a wait record: its name's CRC-32 (31 bits), the
     same in every kernel and table that names it."""
diff --git a/vllm/models/kimi_k3/amd/specialized_decode/mono/common/sync.py b/vllm/models/kimi_k3/amd/specialized_decode/mono/common/sync.py
index 4dfa087e3cd2..cf653b99cfff 100644
--- a/vllm/models/kimi_k3/amd/specialized_decode/mono/common/sync.py
+++ b/vllm/models/kimi_k3/amd/specialized_decode/mono/common/sync.py
@@ -253,3 +253,94 @@ def publish(put, a, i, v, who):
     gpu.barrier()
     if who:
         put(a, i, v)
+
+
+# ---------------------------------------------------------------- sentinel all-reduce
+# (``layer.OPT_SENTS``)
+# Push: plain bf16 pair words into the receiver's region, which holds the
+# sentinel (all ones, a -NaN pair; a pushed word equal to it is remapped).
+# Read: poll off the sentinel, sum, re-arm. Regions alternate by layer parity
+# (``layer.sent_layout``), so a word is re-pushed only after its reader ended.
+SENT = -1
+SENT_NAN = -4128832  # 0xFFC0FFC0
+
+
+def sent_word(w):
+    """A pushed bf16-pair word (an f32-typed ``bf16_pair``), never the
+    sentinel."""
+    w = w.bitcast(fx.Int32) if isinstance(w, fx.Float32) else fx.Int32(w)
+    return (w == fx.Int32(SENT)).select(fx.Int32(SENT_NAN), w)
+
+
+def put_sent(a, i, word):
+    """Word ``i`` of the peer sentinel region ``a`` := ``word`` (``sent_word``),
+    system scope."""
+    bo.buffer_store(
+        fx.Int32(word), rsrc(_noted(a, "put")), i, cache_modifier=_SCOPE_CM[a.space]
+    )
+
+
+def poll_sent(mb, a, idxs, lives=None, rearm=True):
+    """Poll words ``idxs`` of this rank's sentinel region ``a`` until none is
+    the sentinel (with fail-stop), then re-arm them. ``lives``: dead indices
+    are not waited for or re-armed. A read is destructive: one thread a word.
+    ``rearm=False``: the caller re-arms later (``rearm_sent``)."""
+    base = _noted(a, "poll")
+    cm = _SCOPE_CM[a.space]
+    n = len(idxs)
+    if const_expr(n == 1):
+        idxs = [idxs[0], idxs[0]]  # _spin's record reads v[1]
+        if const_expr(lives is not None):
+            lives = [lives[0], lives[0]]
+
+    def load_all():
+        return fx.Vector.from_elements(
+            [
+                fx.Int32(
+                    bo.buffer_load(
+                        rsrc(base),
+                        fx.Int32(i),
+                        vec_width=1,
+                        dtype=T.i32,
+                        cache_modifier=cm,
+                    )
+                )
+                for i in idxs
+            ],
+            fx.Int32,
+        )
+
+    def pending(v):
+        bad = None
+        for e in range_constexpr(len(idxs)):
+            b = v[e] == fx.Int32(SENT)
+            if const_expr(lives is not None):
+                b = b & lives[e]
+            bad = b if bad is None else bad | b
+        return bad
+
+    rid = mb.region_ids.get(a.region, -1)
+    v = _spin(load_all, pending, mb.health, mb.alarm, rid, fx.Int32(idxs[0]), SENT)
+    if const_expr(rearm):
+        rearm_sent(a, idxs[:n], None if lives is None else lives[:n], noted=False)
+    return [v[e] for e in range(n)]
+
+
+def rearm_sent(a, idxs, lives=None, noted=True):
+    """Words ``idxs`` of this rank's sentinel region ``a`` := the sentinel (a
+    dead index skipped), after this thread read them (``poll_sent``)."""
+    base = _noted(a, "put") if noted else a.value
+    cm = _SCOPE_CM[a.space]
+    for e in range_constexpr(len(idxs)):
+        if const_expr(lives is None):
+            bo.buffer_store(
+                fx.Int32(SENT), rsrc(base), fx.Int32(idxs[e]), cache_modifier=cm
+            )
+        else:
+            _rearm(base, idxs[e], lives[e], cm)
+
+
+@traced
+def _rearm(base, i, live, cm):
+    if live:
+        bo.buffer_store(fx.Int32(SENT), rsrc(base), fx.Int32(i), cache_modifier=cm)
diff --git a/vllm/models/kimi_k3/amd/specialized_decode/mono/layer.py b/vllm/models/kimi_k3/amd/specialized_decode/mono/layer.py
index c89c2b309bcd..b9535c6f6714 100644
--- a/vllm/models/kimi_k3/amd/specialized_decode/mono/layer.py
+++ b/vllm/models/kimi_k3/amd/specialized_decode/mono/layer.py
@@ -16,38 +16,53 @@
     CTAs 160 .. 243  router parts (28) and latent parts (56): weights first
     every CTA        ug; down (0 .. 223) or the shared down (224 .. 255);
                      lnorm (56 .. 83), up (0 .. 55), reduce (200 .. 255)
+
+An MLA layer's K2 runs its M1b (``mla_back``) first, in the same launch: every
+CTA pushes, CTAs 0 .. 55 combine and write ``core`` (a flag a head and 4
+tokens), and the o_proj CTAs, their weights in flight, poll the flags of a core
+tile before they load it. That removes the M1b -> K2 launch boundary.
 """
 
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
 
 import flydsl.compiler as flyc
 import flydsl.expr as fx
 import torch
-from flydsl.expr import const_expr, range_constexpr
+from flydsl.expr import const_expr, gpu, range_constexpr
 from flydsl.expr.typing import Int32, Int64
 
 from vllm.models.kimi_k3.amd.specialized_decode.mono.attention import back as k2a
+from vllm.models.kimi_k3.amd.specialized_decode.mono.attention import mla_attn as m1c
+from vllm.models.kimi_k3.amd.specialized_decode.mono.attention import mla_back as m1b
 from vllm.models.kimi_k3.amd.specialized_decode.mono.attention.kda import (
     _SOURCES,
     COLS,
     HIDDEN,
     NA,
-    stage_gate,
+    stage_gates,
     stage_mix,
     stage_slice,
     step_tag,
+    vals_all,
+    vals_words,
 )
 from vllm.models.kimi_k3.amd.specialized_decode.mono.common.debug import (
     region_ids,
+    stamp,
+    stamp_begin,
+    stamp_flush,
 )
 from vllm.models.kimi_k3.amd.specialized_decode.mono.common.ops import (
+    CM_DEV,
     ld_i32,
+    lds_bytes,
     load_ptr64,
     traced,
     uniform,
 )
 from vllm.models.kimi_k3.amd.specialized_decode.mono.common.plan import (
     BLOCKS,
+    LDS_BYTES,
     THREADS,
     WAVES,
     KernelAbi,
@@ -58,12 +73,14 @@
 from vllm.models.kimi_k3.amd.specialized_decode.mono.common.sync import (
     Mailbox,
     preg,
+    publish,
     shift,
     sreg,
     stall,
 )
 from vllm.models.kimi_k3.amd.specialized_decode.mono.stages import gemv as gemv
 from vllm.models.kimi_k3.amd.specialized_decode.mono.stages import moe as k2b
+from vllm.v1.attention.ops.flydsl_kernels import _run_compiled
 
 O0, O1 = NA, BLOCKS  # o_proj CTAs: every CTA past AttnRes's
 O_N = O1 - O0  # 200
@@ -72,10 +89,49 @@
 E0 = 160  # router / latent CTAs (their weights in flight under their o_proj)
 H0 = 8  # the shared SiTU's CTAs (after their AttnRes work)
 SD0 = k2b.DOWN_TASKS  # 224: the shared down's CTAs
+# timeline points: 0 start, 1 o_proj done / AR reduced, 2 early done / slice,
+# 3 ug, 4 down/sdown, 5 lnorm, 6 up, 7 reduce, 8 gate, 9 mix (x published),
+# 10 route, 11 route table, 12 latent quant, 13 the shared SiTU's h done
+TL_POINTS = 18  # + 14 the fused M1b's push, 15 INTER loaded, 16 down done,
+# 17 the fused M1b's combine
 
 assert E0 + k2b.R_TASKS + k2b.L_TASKS <= BLOCKS and k2b.G_TASKS <= E0 - O0
 
-REGIONS = ("part", "wgt", "msq", "xrdy", "ar") + k2b.REGIONS
+REGIONS: tuple[str, ...] = (
+    "part",
+    "wgt",
+    "msq",
+    "xrdy",
+    "ar",
+    "crdy",
+    "cmb",
+    "ardy",
+) + k2b.REGIONS
+# the all-reduce sentinel regions, RLOGP and the tagged rows (``K2Build.opt``),
+# past every region above
+REGIONS = REGIONS + ("ars", "latrs", "finals", "xt", "rlogp")
+
+# K2 options (``K2Build.opt``), all bit-exact vs opt=0; K2_OPT is the default.
+#   1 / 16 / 32  o_proj / latent / final all-reduce via sentinel regions
+#                (plain bf16 pair words, reader re-arms; ``sync.put_sent``)
+#   2    MLP AttnRes slice loads issued before the all-reduce poll
+#   8    x rows also as tagged pairs to the router / latent GEMVs (no XRDY)
+#   64   shared-down tasks that up waits on run first
+#   256  router K parts push fp32 partials to every rank's RLOGP
+#   512  batched polls (route, down, up)
+OPT_SENT = 1
+OPT_PRELOAD = 2
+OPT_TAG_X = 8
+OPT_SENT_LATR = 16
+OPT_SENT_FINAL = 32
+OPT_SENTS = OPT_SENT | OPT_SENT_LATR | OPT_SENT_FINAL
+OPT_SDOWN_FIRST = 64
+OPT_RLOG_PARTS = 256
+OPT_BATCH = 512
+K2_OPT = (
+    OPT_SENTS | OPT_PRELOAD | OPT_TAG_X | OPT_SDOWN_FIRST | OPT_RLOG_PARTS | OPT_BATCH
+)
+SENT_ROWS = 16  # the served widths (runtime.MAX_ROWS)
 
 _STREAM = fx.Stream(None)
 
@@ -95,11 +151,147 @@ class K2Build:
     fault_epoch: int = 0
     fault_layer: int = 0
     fault_ticks: int = 0
+    timeline: bool = False
+    # an MLA layer's M1b in this launch (``mla_back``) when m1b_qlen > 0: its
+    # tokens a request, its ranks and its CMB's peer byte offset. CTAs 0 .. 55
+    # combine (and push, at ranks 8) before their AttnRes work; the o_proj
+    # CTAs poll a flag a (head, 4 tokens) before they load the core tile.
+    m1b_qlen: int = 0
+    m1b_ranks: int = 1
+    m1b_off: int = 0
+    m1b_splits: int = 0  # DCP1: M1c's KV splits (``MlaPostBuild.splits``)
+    # DCP1: M1c (``mla_attn.cta``) in this launch too, first, on its CTAs: the
+    # partials stored device-scope, a flag a CTA (``ardy``) the combine polls
+    m1c: bool = False
+    m1c_wide: bool = False  # M1c's ``wide`` (its KV pool over 4 GiB, ``k2_launch``)
+    # the OPT_* bits
+    opt: int = K2_OPT
+
+
+def k2_sent(key: K2Build) -> bool:
+    """(OPT_SENTS) Whether build ``key``'s all-reduces go through the sentinel
+    regions. The one predicate the push and the reduce of every one of them
+    use: a push the reduce does not take would leave words neither polled nor
+    re-armed, and the next launch's reader would hang."""
+    return bool(key.opt & OPT_SENTS) and key.tokens <= SENT_ROWS
+
+
+def sent_layout() -> dict:
+    """(OPT_SENTS) One layer parity's sentinel regions, name -> (byte
+    offset, bytes): a word a (rank, bf16 pair, row), SENT_ROWS rows."""
+    n = {
+        "ars": TP_ * SENT_ROWS * HIDDEN // 2 * 4,
+        "latrs": TP_ * SENT_ROWS * k2b.LAT // 2 * 4,
+        "finals": TP_ * SENT_ROWS * HIDDEN // 2 * 4,
+    }
+    out, off = {}, 0
+    for name, b in n.items():
+        out[name] = (off, b)
+        off += -(-b // 256) * 256
+    return out
+
+
+def sent_half_bytes() -> int:
+    return max(o + n for o, n in sent_layout().values())
+
+
+def sent_offset() -> int:
+    """(OPT_SENTS) Byte offset (from the kernels' peer base) of the sentinel
+    regions: past K2's regions and M1's at their largest (DCP8), so no other
+    region moves. Two layer parities."""
+    from vllm.models.kimi_k3.amd.specialized_decode.mono.attention import mla
+
+    end = (
+        peer_bytes(k2b.MAX_S)
+        + mla.peer_bytes(k2b.MAX_S, TP_)
+        + m1b.peer_bytes(k2b.MAX_S, TP_)
+    )
+    return -(-end // 4096) * 4096
+
+
+def sent_bytes() -> int:
+    return 2 * sent_half_bytes()
+
+
+def rlogp_half_bytes() -> int:
+    """(OPT_RLOG_PARTS) One epoch parity's RLOGP, tagged pairs
+    `[src rank][part][token][router row]`, SENT_ROWS rows."""
+    return TP_ * k2b.R_PARTS * SENT_ROWS * (k2b.E // TP_) * 8
+
+
+def rlogp_offset() -> int:
+    return -(-(sent_offset() + sent_bytes()) // 4096) * 4096
+
+
+def extra_peer_end() -> int:
+    """The end (from the kernels' peer base) of every region the OPT_* bits
+    add to the peer buffer."""
+    return rlogp_offset() + 2 * rlogp_half_bytes()
+
+
+def k2_rlogp(key: K2Build) -> bool:
+    if not (key.opt & OPT_RLOG_PARTS) or key.tokens > SENT_ROWS:
+        return False
+    vextra = key.tokens * COLS * 4
+    return vals_words(key.tokens, key.nblocks + 1, vextra) >= k2b.R_PARTS * k2b.E
+
+
+def tag_layout() -> dict:
+    """(OPT_TAG_X) The tagged rows' scratch region (name -> (byte offset,
+    bytes)), past every launch's layout at its largest (K1's, M1a's,
+    K2's: as ``MonoRunner`` sizes the scratch), so no other launch's plain
+    word sits where these keep tags."""
+    from vllm.models.kimi_k3.amd.specialized_decode.mono.attention import kda, mla
+
+    S = k2b.MAX_S
+    end = max(
+        kda.scratch_bytes(
+            kda.KdaPreBuild(tokens=S, qlen=S, nblocks=8, delta=True, state_len=3 + S)
+        ),
+        scratch_bytes(K2Build(tokens=S, nblocks=8)),
+        mla.scratch_bytes(mla.MlaPreBuild(tokens=S, nblocks=8, delta=True)),
+    )
+    # + 1 MiB: a fused K2's M1b flags (``crdy``) can outgrow the unfused key's
+    off = -(-(end + (1 << 20)) // 4096) * 4096
+    return {"xt": (off, S * HIDDEN // 2 * 8)}
+
+
+def tag_end() -> int:
+    return max(o + n for o, n in tag_layout().values())
+
+
+# set by the runner once it has allocated (and filled) the regions: a launch
+# whose build uses them without them asserts, as out-of-bounds peer stores would
+# corrupt another rank's memory
+SENT_READY = False
+TAG_READY = False
 
 
 ABI = KernelAbi(
     (
         "core",
+        "o",
+        "o_sb",
+        "o_sh",
+        "lse",
+        "lse_sb",
+        "lse_sh",
+        "seq_lens",
+        "qsl",
+        "w_v",
+        "wv_sh",
+        "wv_sn",
+        "w_vs",
+        "vs_sh",
+        "vs_sn",
+        "vs_sk",
+        "gate",
+        "m1c_q",
+        "m1c_kv",
+        "m1c_pages",
+        "m1c_indptr",
+        "m1c_scale",
+        "m1c_vscale",
         "w_o",
         "prefix",
         "blocks",
@@ -127,6 +319,7 @@ class K2Build:
         "layer",
         "health",
         "alarm",
+        "tl",
     )
 )
 
@@ -152,23 +345,58 @@ def scratch_layout(key: K2Build) -> dict:
         lay[name] = ((end + 15) // 16 * 16 + o, n)
     end = max(o + n for o, n in lay.values())
     lay["xrow"] = ((end + 15) // 16 * 16, s * HIDDEN * 2)
+    end = max(o + n for o, n in lay.values())
+    # the fused M1b's flags, a combine task each
+    lay["crdy"] = ((end + 15) // 16 * 16, m1b.tasks(s, m1b_ctok(key)) * 8)
+    # the fused M1c's flags, a CTA each
+    lay["ardy"] = ((lay["crdy"][0] + lay["crdy"][1] + 15) // 16 * 16, BLOCKS * 8)
     return lay
 
 
+def m1b_ctok(key: K2Build) -> int:
+    """The fused M1b's tokens a combine task (over its NA CTAs)."""
+    if not key.m1b_qlen:
+        return m1b.TT
+    return m1b.ctok(key.tokens, key.m1b_qlen, key.m1b_splits, NA)
+
+
 def scratch_bytes(key: K2Build) -> int:
     return max(o + n for o, n in scratch_layout(key).values())
 
 
+def xl_words(s):
+    """The int32 words of the bf16 row tile LDS ``xl`` the GEMV stages share:
+    the core tile (o_proj), the early stages' x tiles, every h row (the shared
+    down) and the LN tile (up)."""
+    tiled = [k2a.OK, HIDDEN // k2b.R_PARTS, HIDDEN // k2b.L_PARTS, HIDDEN, k2b.LAT]
+    rows = [(k, min(s, gemv.tile_rows(k))) for k in tiled] + [(k2b.SI, s)]
+    return max(n * (k + gemv.LDS_PAD) // 2 for k, n in rows)
+
+
 @traced
 def stage_oproj(c, gi):
     """o_proj tasks gi, gi + 200, gi + 400 (< 448): every task's weights in
-    flight at once (one memory latency, not three), then the MFMAs."""
+    flight at once (one memory latency, not three), then the MFMAs a core tile
+    of 16 rows at a time."""
     ops = [k2a.o_loads(c, fx.min(gi + O_N * j, OTASKS - 1)) for j in range(O_ROUNDS)]
-    k2a.load_core(c)
-    for j in range_constexpr(O_ROUNDS):
-        t = gi + O_N * j
-        if t < OTASKS:
-            k2a.stage_o(c, t, ops[j])
+    for t0, n in gemv.token_tiles(c["S"]):
+        cm = 0
+        if const_expr(c["m1b"]):
+            # the fused M1b's tasks of these rows (a head, ``TC`` tokens), a
+            # thread a flag
+            tc = c["TC"]
+            f0 = t0 // tc * m1b.NH
+            nf = ((t0 + n + tc - 1) // tc) * m1b.NH - f0
+            tid = c["tid"]
+            if tid < nf:
+                c["poll"]([(c["crdy"], f0 + tid, 1)])
+            gpu.barrier()
+            cm = CM_DEV
+        k2a.load_core(c, t0, n, cm)
+        for j in range_constexpr(O_ROUNDS):
+            t = gi + O_N * j
+            if t < OTASKS:
+                k2a.stage_o(c, t, ops[j], t0, n)
 
 
 _BUILDS: dict = {}
@@ -178,53 +406,132 @@ def build(key: K2Build):
     if key in _BUILDS:
         return _BUILDS[key]
     s, ns = key.tokens, key.nblocks + 1
-    assert 1 <= s <= 8
+    assert 1 <= s <= k2b.MAX_S
+    fused = key.m1b_qlen > 0
+    assert not fused or (s % key.m1b_qlen == 0 and key.m1b_ranks in (1, TP_))
+    assert not key.m1c or (fused and key.m1b_ranks == 1 and key.m1b_splits)
+    assert key.m1c or not key.m1c_wide
+    if key.m1c:
+        m1c.check(s, key.m1b_qlen)
+        assert key.m1b_splits == m1c.splits(key.m1b_qlen)
+    # the AttnRes stages' delta tile beside their slices
+    vextra = s * COLS * 4
     lay = scratch_layout(key)
     play = peer_layout(s)
     half = peer_bytes(s) // 2
 
+    # the route table: a slot's kof row (``few_rows``), else its pick tiles
+    few = k2b.few_rows(s)
+    nkof = k2b.max_u(s) * s if few else 1
+    nslot = 1 if few else k2b.max_u(s) + 1
+    npick = 1 if few else s * k2b.TOPK
+    ntile = 1 if few else k2b.max_tiles(s)
+
     @fx.struct
     class RouteLds:
         route: fx.Array[fx.Int32, s * 2 * k2b.TOPK, 16]
-        flag: fx.Array[fx.Int32, k2b.E, 16]
-        expert: fx.Array[fx.Int32, k2b.MAX_U, 16]
-        kof: fx.Array[fx.Int32, k2b.MAX_U * 8, 16]
+        eslot: fx.Array[fx.Int32, k2b.E, 16]
+        expert: fx.Array[fx.Int32, k2b.max_u(s), 16]
+        kof: fx.Array[fx.Int32, nkof, 16]
+        first: fx.Array[fx.Int32, nslot, 16]
+        tfirst: fx.Array[fx.Int32, nslot, 16]
+        order: fx.Array[fx.Int32, npick, 16]
+        tslot: fx.Array[fx.Int32, ntile, 16]
         scan: fx.Array[fx.Int32, WAVES, 16]
 
+    # red, and as words the route table's token bitmaps (``load_route``)
+    @fx.struct
+    class RedF:
+        v: fx.Array[fx.Float32, WAVES * 64 * 4, 16]
+
+    @fx.struct
+    class RedI:
+        v: fx.Array[fx.Int32, WAVES * 64 * 4, 16]
+
+    @fx.union
+    class RedLds:
+        f: RedF
+        i: RedI
+
     @fx.struct
     class EarlyLds:
-        xl: fx.Array[fx.Int32, s * (HIDDEN + gemv.LDS_PAD) // 2, 16]
+        xl: fx.Array[fx.Int32, xl_words(s), 16]
+        av: fx.Array[fx.Float32, THREADS, 16]
 
     @fx.struct
     class AttnLds:
-        vals: fx.Array[fx.Float32, s * ns * COLS, 16]
+        vals: fx.Array[fx.Float32, vals_words(s, ns, vextra), 16]
         dl: fx.Array[fx.Float32, s * COLS, 16]
 
     @fx.struct
     class UgLds:
         latq: fx.Array[fx.Int32, s * k2b.LAT // 8, 16]
-        latsc: fx.Array[fx.Int32, s * k2b.LAT // 32, 16]
-        av: fx.Array[fx.Float32, THREADS, 16]
+        latsc: fx.Array[fx.Int32, k2b.latsc_words(s), 16]
+
+    # every pick's INTER in LDS, or (polled at use) none
+    nq = s * k2b.TOPK if k2b.inter_lds(s) else 1
 
     @fx.struct
     class DownLds:
-        interq: fx.Array[fx.Int32, s * k2b.TOPK * 48, 16]
-        intsc: fx.Array[fx.Int32, s * k2b.TOPK * k2b.UG_GROUPS, 16]
+        interq: fx.Array[fx.Int32, nq * 48, 16]
+        intsc: fx.Array[fx.Int32, nq * k2b.UG_GROUPS, 16]
         contrib: fx.Array[fx.Float32, s * k2b.TOPK * k2b.ROWS, 16]
 
+    @fx.struct
+    class M1bLds:
+        lat: fx.Array[fx.Int32, m1b.TT * m1b.LROW if fused else 1, 16]
+        cw: fx.Array[
+            fx.Float32, m1b.TT * (m1b.WMAX + m1b.KVL) if key.m1b_splits else 1, 16
+        ]
+
+    @fx.struct
+    class NoLds:
+        v: fx.Array[fx.Int32, 1, 16]
+
+    M1cLds = m1c.smem(key.m1b_qlen) if key.m1c else NoLds
+
     @fx.union
     class StageLds:
         early: EarlyLds
         attn: AttnLds
         ug: UgLds
         down: DownLds
+        m1b: M1bLds
+
+    # the stages' LDS: M1c runs first in a CTA and is done (a barrier) before
+    # any of it is touched, so M1c's LDS shares it all
+    @fx.struct
+    class K2Lds:
+        red: RedLds
+        route: RouteLds
+        stage: StageLds
+
+    @fx.union
+    class MainLds:
+        k2: K2Lds
+        m1c: M1cLds  # type: ignore[valid-type]
 
     @fx.struct
     class Smem:
-        red: fx.Array[fx.Float32, WAVES * 64 * 4, 16]
+        tls: fx.Array[fx.Int64, TL_POINTS, 16]
+
+    # LLVM's own over-the-limit error doesn't stop a build (COMPILE_ONLY)
+    assert lds_bytes(Smem) + lds_bytes(MainLds) <= LDS_BYTES, lds_bytes(
+        Smem
+    ) + lds_bytes(MainLds)
 
     keyed = key_tuple(key, _SOURCES)
-    name = f"k3_mono_k2_s{s}_nb{key.nblocks}_x{int(key.reset)}_r{int(key.reduce)}"
+    name = (
+        f"k3_mono_k2_s{s}_nb{key.nblocks}_x{int(key.reset)}_r{int(key.reduce)}"
+        f"_t{int(key.timeline)}"
+    )
+    if fused:
+        name += f"_m{key.m1b_qlen}_{key.m1b_ranks}"
+        name += f"_sp{key.m1b_splits}" if key.m1b_splits else ""
+        name += "_c" if key.m1c else ""
+        name += "w" if key.m1c_wide else ""
+    if key.opt:
+        name += f"_o{key.opt}"
     if key.fault_rank >= 0:
         name += (
             f"_f{key.fault_rank}_{key.fault_epoch}_{key.fault_layer}_{key.fault_ticks}"
@@ -233,6 +540,28 @@ class Smem:
     @flyc.kernel(name=name, known_block_size=[THREADS, 1, 1])
     def k2(
         core: Int64,
+        o: Int64,
+        o_sb: Int32,
+        o_sh: Int32,
+        lse: Int64,
+        lse_sb: Int32,
+        lse_sh: Int32,
+        seq_lens: Int64,
+        qsl: Int64,
+        w_v: Int64,
+        wv_sh: Int32,
+        wv_sn: Int32,
+        w_vs: Int64,
+        vs_sh: Int32,
+        vs_sn: Int32,
+        vs_sk: Int32,
+        gate: Int64,
+        m1c_q: Int64,
+        m1c_kv: Int64,
+        m1c_pages: Int64,
+        m1c_indptr: Int64,
+        m1c_scale: Int32,
+        m1c_vscale: Int32,
         w_o: Int64,
         prefix: Int64,
         blocks: Int64,
@@ -260,14 +589,17 @@ def k2(
         layer: Int32,
         health: Int64,
         alarm: Int64,
+        tl: Int64,
     ):
         _ = keyed
         tid = fx.thread_idx.x
         bid = fx.block_idx.x
         alloc = fx.SharedAllocator()
         lds = alloc.allocate(Smem).peek()
-        rl = alloc.allocate(RouteLds).peek()
-        st = alloc.allocate(StageLds)
+        main = alloc.allocate(MainLds)
+        rd = main.k2.red
+        rl = main.k2.route.peek()
+        st = main.k2.stage
         el, al, ul, dl = st.early.peek(), st.attn.peek(), st.ug.peek(), st.down.peek()
         if const_expr(key.fault_rank >= 0):
             stall(
@@ -284,7 +616,7 @@ def k2(
         bases = peer_bases(peers, TP_)
         own = load_ptr64(peers, rank)
         peer = {}
-        for region in ("rlog", "lat", "latr", "final", "ar"):
+        for region in ("rlog", "lat", "latq4", "latr", "final", "ar"):
             off = play[region][0]
             peer[region] = [
                 shift(preg(b, off, region), fx.Int64(par) * half) for b in bases
@@ -301,6 +633,7 @@ def k2(
             "rank": rank,
             "delta": True,
             "reset": key.reset,
+            "vals_all": vals_all(s, ns, vextra),
             "eps": key.eps,
             "out_eps": key.out_eps,
             "ln_eps": key.ln_eps,
@@ -308,7 +641,8 @@ def k2(
             "put_bf": mb.put_bf,
             "put_words": mb.put_words,
             "poll": mb.poll,
-            "red": lds.red.ptr,
+            "red": rd.f.peek().v.ptr,
+            "redi": rd.i.peek().v.ptr,
             "peer": peer,
             "ar_peer": peer["ar"],
             "ar_own": peer["ar_own"],
@@ -317,21 +651,28 @@ def k2(
             "delta_lds": al.dl.ptr,
             "latq": ul.latq.ptr,
             "latsc": ul.latsc.ptr,
-            "av": ul.av.ptr,
+            "av": el.av.ptr,
             "interq": dl.interq.ptr,
             "intsc": dl.intsc.ptr,
             "contrib": dl.contrib.ptr,
             "scan": rl.scan.ptr,
             "rt": {
                 "route": rl.route.ptr,
-                "flag": rl.flag.ptr,
+                "eslot": rl.eslot.ptr,
                 "expert": rl.expert.ptr,
                 "kof": rl.kof.ptr,
+                "first": rl.first.ptr,
+                "tfirst": rl.tfirst.ptr,
+                "order": rl.order.ptr,
+                "tslot": rl.tslot.ptr,
             },
             "xbuf": scratch + fx.Int64(lay["xrow"][0]),
             "xflags": NA,
+            "tl_on": key.timeline,
+            "tls": lds.tls.ptr,
             "ln": scratch + fx.Int64(lay["ln"][0]),
             "shp": scratch + fx.Int64(lay["shp"][0]),
+            "m1b": fused,
             "args": {
                 "core": core,
                 "w_o": w_o,
@@ -357,6 +698,31 @@ def k2(
                 "out": out,
             },
         }
+        # the OPT_* bits' switches and regions
+        c["tag_x"] = bool(key.opt & OPT_TAG_X)
+        c["batch"] = bool(key.opt & OPT_BATCH)
+        rlogp = k2_rlogp(key)
+        c["rlog_parts"] = rlogp
+        if const_expr(rlogp):
+            roff = rlogp_offset()
+            rhalf = fx.Int64(par) * rlogp_half_bytes()
+            c["rlogp"] = [shift(preg(b, roff, "rlogp"), rhalf) for b in bases]
+            c["rlogp_own"] = shift(preg(own, roff, "rlogp"), rhalf)
+        if const_expr(key.opt & OPT_TAG_X):
+            c["xt"] = sreg(scratch, tag_layout()["xt"][0], "xt")
+        sent = k2_sent(key)
+        if const_expr(sent):
+            # this layer parity's sentinel regions
+            lpar = fx.Int64(layer % 2) * sent_half_bytes()
+            sl = sent_layout()
+            c["sent_ar"] = bool(key.opt & OPT_SENT)
+            c["sent_latr"] = bool(key.opt & OPT_SENT_LATR)
+            c["sent_final"] = bool(key.opt & OPT_SENT_FINAL)
+            c["mb"] = mb
+            for region in ("ars", "latrs", "finals"):
+                off = sent_offset() + sl[region][0]
+                c[region + "_peer"] = [shift(preg(b, off, region), lpar) for b in bases]
+                c[region + "_own"] = shift(preg(own, off, region), lpar)
         for region in (
             "part",
             "wgt",
@@ -369,11 +735,91 @@ def k2(
             "route",
             "inter",
             "isc",
-            "lmsq",
             "lnrdy",
             "shrdy",
         ):
             c[region] = sreg(scratch, lay[region][0], region)
+        on, tls = key.timeline, lds.tls.ptr
+        stamp_begin(on, tls, tid, TL_POINTS)
+        # ---- the fused M1b: core's rows, a flag a combine task
+        if const_expr(fused):
+            R = key.m1b_ranks
+            c["crdy"] = sreg(scratch, lay["crdy"][0], "crdy")
+            c.update(
+                Q=key.m1b_qlen,
+                R=R,
+                splits=key.m1b_splits,
+                TC=m1b_ctok(key),
+                lat=st.m1b.peek().lat.ptr,
+                cw=st.m1b.peek().cw.ptr,
+                core_cm=CM_DEV,
+                cmb=[preg(b, key.m1b_off, "cmb") for b in bases] if R > 1 else [],
+                cmb_own=preg(own, key.m1b_off, "cmb"),
+            )
+            c["args"].update(
+                o=o,
+                o_sb=o_sb,
+                o_sh=o_sh,
+                lse=lse,
+                lse_sb=lse_sb,
+                lse_sh=lse_sh,
+                seq_lens=seq_lens,
+                qsl=qsl,
+                w_v=w_v,
+                wv_sh=wv_sh,
+                wv_sn=wv_sn,
+                w_vs=w_vs,
+                vs_sh=vs_sh,
+                vs_sn=vs_sn,
+                vs_sk=vs_sk,
+                gate=gate,
+            )
+            if const_expr(key.m1c):
+                c["ardy"] = sreg(scratch, lay["ardy"][0], "ardy")
+                c["m1c_ht"] = m1c.geometry(key.m1b_qlen)[0]
+                ml = main.m1c.peek()
+                m1c.cta(
+                    {
+                        "S": s,
+                        "Q": key.m1b_qlen,
+                        "stamps": False,
+                        "dbg": 0,
+                        "wide": key.m1c_wide,
+                        # no warp specialization in here: with it the fused
+                        # TP8 builds fault (memory access) where M1c alone
+                        # passes (not root-caused)
+                        "ws": 0,
+                        "bid": bid,
+                        "tid": tid,
+                        "cm": CM_DEV,
+                        "ks": ml.ks.ptr,
+                        "xm": ml.xm.ptr,
+                        "xp": ml.xp.ptr,
+                        "xi": ml.xi.ptr,
+                        "args": {
+                            "qb": m1c_q,
+                            "kv": m1c_kv,
+                            "pages": m1c_pages,
+                            "indptr": m1c_indptr,
+                            "part_o": o,
+                            "part_lse": lse,
+                            "scale": m1c_scale,
+                            "vscale": m1c_vscale,
+                        },
+                        "done": lambda f: publish(c["put"], c["ardy"], f, 1, tid == 0),
+                    }
+                )
+                # M1c's LDS is the stages' (``MainLds``): its last DMAs (a
+                # block past its end, prefetched) landed, every wave done
+                m1c.wait_dma()
+            if const_expr(R > 1):
+                m1b.stage_push(c)
+            stamp(on, tls, tid, 14)
+            if bid < NA:
+                m1b.combine_tasks(
+                    c, NA, lambda task: publish(c["put"], c["crdy"], task, 1, tid == 0)
+                )
+            stamp(on, tls, tid, 17)
         # ---- o_proj on every CTA past AttnRes's
         eb = bid - E0
         lb = eb - k2b.R_TASKS
@@ -383,33 +829,65 @@ def k2(
         # retire in order), the early weights after: they land before x does
         if is_r:
             stage_oproj(c, bid - O0)
+            stamp(on, tls, tid, 1)
             rops = k2b.early_loads(c, 0, eb // k2b.R_PARTS, eb % k2b.R_PARTS)
             k2b.early_compute(c, 0, eb // k2b.R_PARTS, eb % k2b.R_PARTS, rops)
         if is_l:
             stage_oproj(c, bid - O0)
+            stamp(on, tls, tid, 1)
             lops = k2b.early_loads(c, 1, lb // k2b.L_PARTS, lb % k2b.L_PARTS)
             k2b.early_compute(c, 1, lb // k2b.L_PARTS, lb % k2b.L_PARTS, lops)
         if (bid >= O0) & ~is_r & ~is_l:
             stage_oproj(c, bid - O0)
+            stamp(on, tls, tid, 1)
             gi = bid - O0
             if gi < k2b.G_TASKS:
                 k2b.early_compute(c, 2, gi, 0, k2b.early_loads(c, 2, gi, 0))
         # ---- AttnRes: the reduced o_proj -> x; then route / SiTU
         if bid < NA:
-            k2a.stage_reduce(c, bid)
-            stage_slice(c, bid)
-            if bid < s:
-                stage_gate(c, bid)
+            if const_expr(key.opt & OPT_PRELOAD):
+                # the slice's loads before the AR poll
+                sld = k2a.slice_loads(c, bid)
+                k2a.stage_reduce(c, bid)
+                stamp(on, tls, tid, 1)
+                k2a.stage_slice_pre(c, bid, sld)
+            else:
+                k2a.stage_reduce(c, bid)
+                stamp(on, tls, tid, 1)
+                stage_slice(c, bid)
+            stamp(on, tls, tid, 2)
+            stage_gates(c, bid)
+            stamp(on, tls, tid, 8)
             stage_mix(c, bid)
-            if bid < s:
-                k2b.stage_route(c, bid)
-            if (bid >= H0) & (bid < H0 + s):
-                k2b.stage_h(c, bid - H0)
+            if const_expr(sent and (key.opt & OPT_SENT)):
+                # the sentinel words stage_reduce read, off its critical path
+                k2a.rearm_ar(c, bid)
+            stamp(on, tls, tid, 9)
+            # the shared SiTU of token t on CTA (t + H0) % 56, first: SGU lands
+            # before the logits
+            hb = (bid + NA - H0) % NA
+            for rr in range_constexpr((s + NA - 1) // NA):
+                if hb + NA * rr < s:
+                    k2b.stage_h(c, hb + NA * rr)
+            stamp(on, tls, tid, 13)
+            for rr in range_constexpr((s + NA - 1) // NA):
+                if bid + NA * rr < s:
+                    if const_expr(rlogp):
+                        k2b.route_parts(c, bid + NA * rr)
+                    k2b.stage_route(c, bid + NA * rr)
+            stamp(on, tls, tid, 10)
+        if bid >= NA:
+            stamp(on, tls, tid, 2)
         # ---- the routed experts
         # the latent (ready ~10 us before the route) quantized first
-        k2b.load_latq(c)
-        nu = k2b.load_route(c)
-        total = nu * k2b.UG_GROUPS
+        if const_expr(k2b.few_rows(s)):
+            k2b.load_latq4(c)
+        else:
+            k2b.load_latq(c)
+        stamp(on, tls, tid, 12)
+        nt = k2b.load_route(c)
+        stamp(on, tls, tid, 11)
+        total = nt * k2b.UG_GROUPS
         last = total - 1
         # reversed placement: the round's leftover tasks go to the high CTAs,
         # whose early work ends first (the AttnRes / route CTAs start ug last)
@@ -421,23 +899,58 @@ def k2(
             k2b.stage_ug(c, task, ops)
             task = nxt
             ops = ops_n
+        stamp(on, tls, tid, 3)
         if bid < k2b.DOWN_TASKS:
-            dops = k2b.down_loads(c, bid, nu, fx.min(c["wave"], nu - 1))
-            k2b.load_inter(c)
-            k2b.stage_down(c, bid, nu, dops)
-        k2b.run_sdown_rest(c, bid, w_sd)
-        k2b.run_sdown(c, bid, w_sd)
+            dops = k2b.down_loads(c, bid, nt, fx.min(c["wave"], nt - 1))
+            if const_expr(k2b.inter_lds(s)):
+                k2b.load_inter(c)
+            stamp(on, tls, tid, 15)
+            k2b.stage_down(c, bid, nt, dops)
+            stamp(on, tls, tid, 16)
+        if const_expr(key.opt & OPT_SDOWN_FIRST):
+            k2b.run_sdown_first(c, bid, w_sd)
+            k2b.run_sdown_rest(c, bid, w_sd)
+        else:
+            k2b.run_sdown_rest(c, bid, w_sd)
+            k2b.run_sdown(c, bid, w_sd)
+        stamp(on, tls, tid, 4)
         if (bid >= k2b.LN0) & (bid < k2b.LN0 + k2b.LN_TASKS):
-            k2b.stage_lnorm(c, bid - k2b.LN0)
-        if bid < k2b.UP_TASKS:
-            k2b.stage_up(c, bid)
+            k2b.run_lnorm(c, bid)
+        stamp(on, tls, tid, 5)
+        if bid < k2b.UP_TASKS * k2b.up_splits(s):
+            k2b.run_up(c, bid)
+        stamp(on, tls, tid, 6)
         if const_expr(key.reduce):  # noqa: SIM102 (traced: `and` does not combine traced values)
             if (bid >= k2b.RED0) & (bid < k2b.RED0 + k2b.RED_TASKS):
                 k2b.stage_reduce(c, bid - k2b.RED0)
+        stamp(on, tls, tid, 7)
+        stamp_flush(on, tls, tl, tid, bid, TL_POINTS)
 
     @flyc.jit
     def launch(
         core: Int64,
+        o: Int64,
+        o_sb: Int32,
+        o_sh: Int32,
+        lse: Int64,
+        lse_sb: Int32,
+        lse_sh: Int32,
+        seq_lens: Int64,
+        qsl: Int64,
+        w_v: Int64,
+        wv_sh: Int32,
+        wv_sn: Int32,
+        w_vs: Int64,
+        vs_sh: Int32,
+        vs_sn: Int32,
+        vs_sk: Int32,
+        gate: Int64,
+        m1c_q: Int64,
+        m1c_kv: Int64,
+        m1c_pages: Int64,
+        m1c_indptr: Int64,
+        m1c_scale: Int32,
+        m1c_vscale: Int32,
         w_o: Int64,
         prefix: Int64,
         blocks: Int64,
@@ -465,11 +978,34 @@ def launch(
         layer: Int32,
         health: Int64,
         alarm: Int64,
+        tl: Int64,
         stream: fx.Stream = _STREAM,
     ):
         _ = keyed
         k2(
             core,
+            o,
+            o_sb,
+            o_sh,
+            lse,
+            lse_sb,
+            lse_sh,
+            seq_lens,
+            qsl,
+            w_v,
+            wv_sh,
+            wv_sn,
+            w_vs,
+            vs_sh,
+            vs_sn,
+            vs_sk,
+            gate,
+            m1c_q,
+            m1c_kv,
+            m1c_pages,
+            m1c_indptr,
+            m1c_scale,
+            m1c_vscale,
             w_o,
             prefix,
             blocks,
@@ -497,6 +1033,7 @@ def launch(
             layer,
             health,
             alarm,
+            tl,
         ).launch(grid=(BLOCKS,), block=(THREADS,), stream=stream)
 
     ABI.check(k2, launch)
@@ -507,6 +1044,10 @@ def launch(
 TP_ = k2b.TP
 
 
+def ptr(t):
+    return 0 if t is None else t.data_ptr()
+
+
 def k2_launch(
     key: K2Build,
     *,
@@ -536,12 +1077,33 @@ def k2_launch(
     layer,
     health,
     alarm,
+    tl=None,
+    m1b_in=None,
+    m1c_in=None,
 ):
     """Host launcher: ``core`` [S, 1536] (o_proj's input), ``prefix`` updated in
-    place, ``out`` the MoE output [S, 7168]; MoE weights as ``moe``."""
+    place, ``out`` the MoE output [S, 7168]; MoE weights as ``moe``. With
+    ``key.m1b_qlen``, ``m1b_in`` is ``mla_back.mla_post``'s (o, lse, seq_lens,
+    qsl, w_v, w_v_scale, gate) and this launch writes ``core`` first. With
+    ``key.m1c``, ``m1c_in`` is ``mla_attn.mla_attn``'s (q, kv, pages, indptr,
+    scale, kv_scale), its partials M1b's o / lse (split 0's views)."""
+    assert (m1b_in is not None) == (key.m1b_qlen > 0)
+    assert (m1c_in is not None) == key.m1c
+    # the regions the OPT_* bits use exist
+    assert SENT_READY or not (k2_sent(key) or k2_rlogp(key)), (
+        "OPT_SENTS / OPT_RLOG_PARTS without the runner's peer regions"
+    )
+    assert TAG_READY or not (key.opt & OPT_TAG_X), (
+        "OPT_TAG_X without the runner's region"
+    )
+    if m1c_in is not None and m1c.wide_kv(m1c_in["kv"]):
+        key = replace(key, m1c_wide=True)
     f = build(key)
-    f(
+    _run_compiled(
+        f,
         core.data_ptr(),
+        *(m1b.host_args(key.m1b_ranks, **m1b_in) if m1b_in is not None else [0] * 16),
+        *(m1c.host_args(**m1c_in) if m1c_in is not None else [0] * 6),
         w_o.data_ptr(),
         prefix.data_ptr(),
         blocks.data_ptr(),
@@ -569,5 +1131,6 @@ def k2_launch(
         layer,
         health.data_ptr(),
         alarm,
-        stream=torch.cuda.current_stream(),
+        ptr(tl),
+        torch.cuda.current_stream(),
     )
diff --git a/vllm/models/kimi_k3/amd/specialized_decode/mono/stages/gemv.py b/vllm/models/kimi_k3/amd/specialized_decode/mono/stages/gemv.py
index 26acddb5ae39..6cde23731a63 100644
--- a/vllm/models/kimi_k3/amd/specialized_decode/mono/stages/gemv.py
+++ b/vllm/models/kimi_k3/amd/specialized_decode/mono/stages/gemv.py
@@ -26,6 +26,22 @@
 ROWS = 16
 KCH = 128
 LDS_PAD = 8  # bf16 a padded token row
+TILE = 16  # token rows an MFMA's B operand holds
+XTILE_BYTES = 116 * 1024  # a tile of x rows in LDS at most
+
+
+def tile_rows(k):
+    """Token rows an LDS tile of K = ``k`` (padded) bf16 rows holds: 16, or 8
+    for the 7168-wide rows (16 of them are 229 KB)."""
+    n = TILE
+    while n > 1 and n * (k + LDS_PAD) * 2 > XTILE_BYTES:
+        n //= 2
+    return n
+
+
+def token_tiles(s, n=TILE):
+    """(first token, rows) of each tile of ``n`` rows over S = ``s``."""
+    return [(t0, min(n, s - t0)) for t0 in range(0, s, n)]
 
 
 def chunks(k):
@@ -98,9 +114,9 @@ def const_live(per, n):
 
 
 @traced
-def rows_to_lds(tid, src, k, s, xl, xrow, cm=0, ldk=None, k0=0):
-    """Columns k0 .. k0 + k of token rows 0 .. s - 1 of the bf16 [s, ldk] at
-    ``src`` -> LDS (``xrow`` a row), 16 B a lane."""
+def rows_to_lds(tid, src, k, s, xl, xrow, cm=0, ldk=None, k0=0, t0=0):
+    """Columns k0 .. k0 + k of token rows t0 .. t0 + s - 1 of the bf16 [S, ldk]
+    at ``src`` -> LDS rows 0 .. s - 1 (``xrow`` a row), 16 B a lane."""
     ldk = k if ldk is None else ldk
     words = s * k // 8
     for i in range_constexpr((words + THREADS - 1) // THREADS):
@@ -110,7 +126,7 @@ def rows_to_lds(tid, src, k, s, xl, xrow, cm=0, ldk=None, k0=0):
         v = fx.Vector(
             bo.buffer_load(
                 rsrc(src),
-                (t * ldk + k0 + kk) // 2,
+                ((t0 + t) * ldk + k0 + kk) // 2,
                 vec_width=4,
                 dtype=T.i32,
                 cache_modifier=cm,
@@ -121,6 +137,6 @@ def rows_to_lds(tid, src, k, s, xl, xrow, cm=0, ldk=None, k0=0):
     gpu.barrier()
 
 
-def dev_rows_to_lds(tid, src, k, s, xl, xrow):
+def dev_rows_to_lds(tid, src, k, s, xl, xrow, t0=0):
     """``rows_to_lds`` at device scope: rows another CTA of this launch wrote."""
-    rows_to_lds(tid, src, k, s, xl, xrow, CM_DEV)
+    rows_to_lds(tid, src, k, s, xl, xrow, CM_DEV, t0=t0)
diff --git a/vllm/models/kimi_k3/amd/specialized_decode/mono/stages/moe.py b/vllm/models/kimi_k3/amd/specialized_decode/mono/stages/moe.py
index 01a0f861e6f8..bacfcbf96e25 100644
--- a/vllm/models/kimi_k3/amd/specialized_decode/mono/stages/moe.py
+++ b/vllm/models/kimi_k3/amd/specialized_decode/mono/stages/moe.py
@@ -3,7 +3,7 @@
 """K2b ``moe``: one Kimi-K3 latent-MoE layer, from the MoE input row x to the
 layer's output, every TP reduction inside the kernel.
 
-One launch a layer and rank, ``BLOCKS`` x ``THREADS``, S <= 8 rows (a
+One launch a layer and rank, ``BLOCKS`` x ``THREADS``, S <= 64 rows (a
 spec-verify step). The original path (vLLM ROCm, see MOE_SPEC.md) replicates the
 router and the latent down projection on every rank; here each rank computes
 its eighth and the rows are all-gathered over the peer buffer.
@@ -14,16 +14,18 @@
       latent  (28 x 16 rows of this rank's 448): bf16 -> every rank's LAT
       sgu     (96 x 16 rows of the shared gate | up): bf16 -> SGU
     route   (S, CTAs 131 ..): sigmoid (v_rcp), + bias, top-16, the pre-bias
-            sigmoids renormalized (IEEE 1 / sum) -> ROUTE (every rank alike)
+            sigmoids renormalized (IEEE 1 / sum) -> ROUTE (every rank alike);
+            every CTA then lays out U and its pick tiles in LDS (``load_route``)
     sdown   (448 x 16 rows, K 768, CTAs 131 ..): SiTU of SGU (vLLM's kernel) ->
             bf16 -> GEMV -> bf16 partial; this rank's up columns -> SHP (plain,
             SHRDY flags), the rest -> every rank's FINAL
-    ug      (|U| x 12, every CTA): the union U of the step's picks; an expert's
-            32 gate + 32 up rows against the MXFP4 latent (aiter's inline quant),
-            SiTUv2 (aiter's fp32 epilogue), MXFP4 -> INTER / ISC a pick
-    down    (224 x 16 latent rows, every CTA): w2 of every expert of U against
-            its picks' INTER, bf16(w acc) folded in each token's top-k order with
-            a bf16 rounding an add -> every rank's LATR
+    ug      (pick tiles x 12, every CTA): the union U of the step's picks, a
+            slot's picks in tiles of 16; an expert's 32 gate + 32 up rows against
+            its tile's MXFP4 latent rows (aiter's inline quant), SiTUv2 (aiter's
+            fp32 epilogue), MXFP4 -> INTER / ISC a pick
+    down    (224 x 16 latent rows, every CTA): w2 of every pick tile's expert
+            against its picks' INTER, bf16(w acc) folded in each token's top-k
+            order with a bf16 rounding an add -> every rank's LATR
     lnorm   (28 x 128 columns): the ranks' LATR in rank order, bf16, RMSNorm
             (the fused AR + RMSNorm form) -> LN (plain, LNRDY flags)
     up      (56 x 16 rows of this rank's 896): LN GEMV + SHP -> bf16 -> every
@@ -40,6 +42,17 @@
 from flydsl.expr import const_expr, gpu, range_constexpr, rocdl
 from flydsl.expr.typing import T, as_ir_value
 
+from vllm.models.kimi_k3.amd.specialized_decode.mono.attention.back import (
+    peer_at,
+    sent_item,
+    sent_ix,
+    sent_on,
+    sent_sum,
+    two_shot,
+)
+from vllm.models.kimi_k3.amd.specialized_decode.mono.attention.kda import (
+    poll_rows_to_lds,
+)
 from vllm.models.kimi_k3.amd.specialized_decode.mono.common.mx import (
     FP4,
     UNIT_SCALE,
@@ -50,6 +63,8 @@
 from vllm.models.kimi_k3.amd.specialized_decode.mono.common.ops import (
     CM_DEV,
     ballot,
+    bf2_f32,
+    bf16_pair,
     bf16_round,
     bf_hi,
     bf_lo,
@@ -60,10 +75,12 @@
     hw_rsq,
     lanes_below,
     popcount,
+    readlane,
     row_sum,
     rsrc,
     traced,
     uniform,
+    unlikely,
     wave_sum,
 )
 from vllm.models.kimi_k3.amd.specialized_decode.mono.common.plan import (
@@ -73,7 +90,13 @@
     pair_layout,
 )
 from vllm.models.kimi_k3.amd.specialized_decode.mono.common.ranks import sum_partials
-from vllm.models.kimi_k3.amd.specialized_decode.mono.common.sync import publish
+from vllm.models.kimi_k3.amd.specialized_decode.mono.common.sync import (
+    poll_sent,
+    publish,
+    put_sent,
+    rearm_sent,
+    sent_word,
+)
 from vllm.models.kimi_k3.amd.specialized_decode.mono.stages import gemv as gemv
 
 REGIONS = (
@@ -84,11 +107,11 @@
     "route",
     "inter",
     "isc",
-    "lmsq",
     "lnrdy",
     "shrdy",
     "rlog",
     "lat",
+    "latq4",
     "latr",
     "final",
 )
@@ -99,6 +122,8 @@
 LAT = 3584
 E = 896
 TOPK = 16
+TOPK_BITS = TOPK.bit_length() - 1
+assert 1 << TOPK_BITS == TOPK
 RI = 384  # routed intermediate a rank
 SI = 768  # shared intermediate a rank
 ROWS = 16
@@ -129,7 +154,43 @@
 W2_STEPS = RI // 128  # 3
 DOWN_TASKS = LAT // ROWS  # 224
 DOWN_BATCH = 5  # a wave's experts whose w2 is in flight together
-MAX_U = 8 * TOPK
+MASK_WORDS = 2  # an expert's token bitmap in the route table
+MAX_S = 32 * MASK_WORDS  # 64 rows
+PTILE = 16  # picks a pick tile: the ug and down MFMAs' N
+INTER_LDS_S = 16  # rows whose every pick's INTER down holds in LDS
+
+
+def max_u(s):
+    """The union U's slots at most: every pick distinct, at most every expert."""
+    return min(E, s * TOPK)
+
+
+def max_tiles(s):
+    """U's pick tiles at most: a slot's partial tile each, then the whole
+    tiles, never more than a tile a pick."""
+    return min(s * TOPK, max_u(s) + s * TOPK // PTILE)
+
+
+def inter_lds(s):
+    """Down reads the picks' INTER from LDS (``load_inter``), else polls each
+    pick's at its use."""
+    return s <= INTER_LDS_S
+
+
+def few_rows(s):
+    """At most a pick tile of rows: a slot's picks fit one tile, so a slot is
+    a tile and the MFMAs' columns are the tokens. The route table is then
+    expert flags and each slot's top-k index a token (``kof``), and the
+    latent's scale codes a word each: none of the bitmap, the pick order or
+    the packed codes' atomics."""
+    return s <= PTILE
+
+
+def latsc_words(s):
+    return s * LAT // 32 if few_rows(s) else s * LAT // 128
+
+
+TL_POINTS = 8  # start, early, sdown, ug, down, lnorm, up, reduce
 
 _STREAM = fx.Stream(None)
 
@@ -139,6 +200,7 @@ def peer_pairs(s):
     return {
         "rlog": TP * s * (E // TP),
         "lat": s * LAT // 2,
+        "latq4": s * LAT // 8 + s * LAT // 32,  # ``few_rows``: fp4 words, codes
         "latr": TP * s * LAT // 2,
         "final": TP * s * HIDDEN // 2,
     }
@@ -153,8 +215,7 @@ def scratch_layout(s: int) -> dict:
         "route": s * 2 * TOPK,
         "inter": s * TOPK * 48,
         "isc": s * TOPK * UG_GROUPS,
-        "lmsq": LN_TASKS * s,
-        "lnrdy": LN_TASKS,
+        "lnrdy": s,
         "shrdy": UP_TASKS,
     }
     lay = pair_layout(pairs.items())
@@ -169,6 +230,11 @@ def f32_of(w):
     return w.bitcast(fx.Float32)
 
 
+def tag_x_on(c):
+    """``layer.OPT_TAG_X``, a plain call (see ``back.sent_on``)."""
+    return c.get("tag_x", False)
+
+
 def tanh_libm(x):
     """The ``tanhf`` of x to a few ulp: (1 - e) / (1 + e), e = exp(-2|x|), x's sign."""
     e = hw_exp2(abs(x) * fx.Float32(-2.0 * LOG2E))
@@ -236,26 +302,38 @@ def pk_fp4_bf16(old, word, scale, j):
 
 # ---------------------------------------------------------------- early
 @traced
-def early_out(c, kind, i, q, red):
-    """Task (i, q) of ``kind`` (0 router, 1 latent, 2 sgu): its 16 rows. A K-split
-    kind's parts q > 0 publish fp32 partials; part 0 adds them in part order and
-    publishes the rows (router: fp32 to every rank's RLOG, latent: bf16 to every
-    rank's LAT); sgu bf16 to SGU."""
+def early_out(c, kind, i, q, red, t0, n):
+    """Task (i, q) of ``kind`` (0 router, 1 latent, 2 sgu): its 16 rows for
+    tokens t0 .. t0 + n. A K-split kind's parts q > 0 publish fp32 partials;
+    part 0 adds them in part order and publishes the rows (router: fp32 to every
+    rank's RLOG, latent: bf16 to every rank's LAT); sgu bf16 to SGU."""
     s, tid = c["S"], c["tid"]
-    if const_expr(kind == 2):
-        if tid < (ROWS // 2) * s:
+    if const_expr((kind == 1) and few_rows(s)):
+        latent_mx_out(c, i, q, red, t0, n)
+    elif const_expr(kind == 2):
+        if tid < (ROWS // 2) * n:
             rp = tid % (ROWS // 2)
-            t = tid // (ROWS // 2)
-            v0 = bf16_round(row_sum(red, 2 * rp, t))
-            v1 = bf16_round(row_sum(red, 2 * rp + 1, t))
-            c["put_bf"](c["sgu"], t * 2 * SI + i * ROWS + 2 * rp, [v0, v1])
+            tl = tid // (ROWS // 2)
+            v0 = bf16_round(row_sum(red, 2 * rp, tl))
+            v1 = bf16_round(row_sum(red, 2 * rp + 1, tl))
+            c["put_bf"](c["sgu"], (t0 + tl) * 2 * SI + i * ROWS + 2 * rp, [v0, v1])
+    elif const_expr((kind == 0) and rlogp_on(c)):
+        # (OPT_RLOG_PARTS) Every part's fp32 partial straight to
+        # every rank's RLOGP (the route sums them), no part-0 hop
+        if tid < ROWS * n:
+            r = tid % ROWS
+            t = t0 + tid // ROWS
+            v = row_sum(red, r, tid // ROWS)
+            for p in range_constexpr(TP):
+                c["put"](c["rlogp"][p], rlogp_pair(c["rank"], q, s, t, i * ROWS + r), v)
     else:
         parts = R_PARTS if kind == 0 else L_PARTS
         region = c["rpart"] if kind == 0 else c["lpart"]
-        if tid < ROWS * s:
+        if tid < ROWS * n:
             r = tid % ROWS
-            t = tid // ROWS
-            v = row_sum(red, r, t)
+            tl = tid // ROWS
+            t = t0 + tl
+            v = row_sum(red, r, tl)
             at = lambda qq: ((i * parts + qq) * s + t) * ROWS + r  # noqa: E731
             if q > 0:
                 c["put"](region, at(q), v)
@@ -266,9 +344,9 @@ def early_out(c, kind, i, q, red):
                 fx.ptr_store(v, c["av"] + tid)
         gpu.barrier()
         if const_expr(kind == 0):
-            if (q == 0) & (tid < ROWS * s):
+            if (q == 0) & (tid < ROWS * n):
                 r = tid % ROWS
-                t = tid // ROWS
+                t = t0 + tid // ROWS
                 v = fx.ptr_load(c["av"] + tid)
                 for p in range_constexpr(TP):
                     c["put"](
@@ -277,17 +355,73 @@ def early_out(c, kind, i, q, red):
                         v,
                     )
         else:
-            if (q == 0) & (tid < (ROWS // 2) * s):
+            if (q == 0) & (tid < (ROWS // 2) * n):
                 rp = tid % (ROWS // 2)
-                t = tid // (ROWS // 2)
-                v0 = bf16_round(fx.ptr_load(c["av"] + (t * ROWS + 2 * rp)))
-                v1 = bf16_round(fx.ptr_load(c["av"] + (t * ROWS + 2 * rp + 1)))
+                tl = tid // (ROWS // 2)
+                t = t0 + tl
+                v0 = bf16_round(fx.ptr_load(c["av"] + (tl * ROWS + 2 * rp)))
+                v1 = bf16_round(fx.ptr_load(c["av"] + (tl * ROWS + 2 * rp + 1)))
                 col = c["rank"] * (LAT // TP) + i * ROWS + 2 * rp
                 for p in range_constexpr(TP):
                     c["put_bf"](c["peer"]["lat"][p], t * LAT + col, [v0, v1])
     gpu.barrier()
 
 
+@traced
+def latent_mx_out(c, i, q, red, t0, n):
+    """``early_out`` of the latent at ``few_rows``: an even group's part 0 ends
+    the MXFP4 block of its rows and the odd group's (32 rows): the 4 partials in
+    part order, bf16, then aiter's inline quant as ``load_latq`` -> every rank's
+    LATQ4 (fp4 words, then a code a word). The other tasks publish partials."""
+    s, tid = c["S"], c["tid"]
+    region = c["lpart"]
+    at = lambda ii, qq, t, r: ((ii * L_PARTS + qq) * s + t) * ROWS + r  # noqa: E731
+    fin = (i % 2 == 0) & (q == 0)
+    if (~fin) & (tid < ROWS * n):
+        c["put"](
+            region,
+            at(i, q, t0 + tid // ROWS, tid % ROWS),
+            row_sum(red, tid % ROWS, tid // ROWS),
+        )
+    if fin & (tid < 2 * ROWS * n):
+        tl = tid // (2 * ROWS)
+        r = tid % ROWS
+        odd = tid % (2 * ROWS) >= ROWS
+        t = t0 + tl
+        a0 = odd.select(at(i + 1, 0, t, r), at(i, 1, t, r))
+        a1 = odd.select(at(i + 1, 1, t, r), at(i, 1, t, r))
+        got = c["poll"]([(region, a0, 1), (region, a1, 1)])
+        p0, p1 = f32_of(got[0][0]), f32_of(got[1][0])
+        v = odd.select(p0 + p1, row_sum(red, r, tl) + p0)
+        fx.ptr_store(bf16_round(v), c["av"] + tid)
+    gpu.barrier()
+    if fin & (tid < 4 * n):
+        tl = tid // 4
+        qd = tid % 4
+        t = t0 + tl
+        vs = [fx.ptr_load(c["av"] + (tl * 2 * ROWS + k)) for k in range(2 * ROWS)]
+        am = fx.Int32(0)
+        for v in vs:
+            am = fx.max(am, (v.bitcast(fx.Int32) >> 16) & 0x7FFF)
+        code = e8m0_roundup((am << 16).bitcast(fx.Float32))
+        scale = (code << 23).bitcast(fx.Float32)
+        pk = fx.Int32(0)
+        for j in range_constexpr(4):
+            w = bf16_pair(
+                fx.ptr_load(c["av"] + (tl * 2 * ROWS + 8 * qd + 2 * j)),
+                fx.ptr_load(c["av"] + (tl * 2 * ROWS + 8 * qd + 2 * j + 1)),
+            )
+            pk = pk_fp4_bf16(pk, w.bitcast(fx.Int32), scale, j)
+        kb = c["rank"] * (LAT // TP // 32) + i // 2
+        for p in range_constexpr(TP):
+            c["put_words"](c["peer"]["latq4"][p], t * (LAT // 8) + kb * 4 + qd, [pk])
+            if qd == 0:
+                c["put_words"](
+                    c["peer"]["latq4"][p], s * LAT // 8 + t * (LAT // 32) + kb, [code]
+                )
+    gpu.barrier()
+
+
 def early_shape(kind):
     """(K a part, parts) of an early kind: router a quarter, latent a half, sgu all."""
     return {
@@ -318,21 +452,30 @@ def early_compute(c, kind, i, q, ops):
     s, tid, lane, wave, red = c["S"], c["tid"], c["lane"], c["wave"], c["red"]
     klen, _ = early_shape(kind)
     xrow = klen + gemv.LDS_PAD
-    if const_expr(c.get("xflags") is not None):
+    # (OPT_TAG_X) The router's and the latent's x rows polled as
+    # tagged pairs (the shared gate/up keeps XRDY and the plain rows)
+    tagged = tag_x_on(c) and kind < 2
+    if const_expr(tagged):
+        src, cm = None, 0
+    elif const_expr(c.get("xflags") is not None):
         if tid < c["xflags"]:
             c["poll"]([(c["xrdy"], tid, 1)])
         gpu.barrier()
-        gemv.rows_to_lds(
-            tid, c["xbuf"], klen, s, c["xl"], xrow, CM_DEV, ldk=HIDDEN, k0=q * klen
-        )
+        src, cm = c["xbuf"], CM_DEV
     else:
-        gemv.rows_to_lds(
-            tid, c["args"]["x"], klen, s, c["xl"], xrow, ldk=HIDDEN, k0=q * klen
-        )
-    acc = gemv.mfmas(lane, wave, c["xl"], xrow, klen, s, ops)
-    fx.ptr_store(acc, red + (wave * 64 + lane) * 4)
-    gpu.barrier()
-    early_out(c, kind, i, q, red)
+        src, cm = c["args"]["x"], 0
+    # a tile of x rows at a time, the weights held over them
+    for t0, n in gemv.token_tiles(s, gemv.tile_rows(klen)):
+        if const_expr(tagged):
+            poll_rows_to_lds(c, c["xt"], klen, n, c["xl"], xrow, HIDDEN, q * klen, t0)
+        else:
+            gemv.rows_to_lds(
+                tid, src, klen, n, c["xl"], xrow, cm, ldk=HIDDEN, k0=q * klen, t0=t0
+            )
+        acc = gemv.mfmas(lane, wave, c["xl"], xrow, klen, n, ops)
+        fx.ptr_store(acc, red + (wave * 64 + lane) * 4)
+        gpu.barrier()
+        early_out(c, kind, i, q, red, t0, n)
 
 
 @traced
@@ -369,31 +512,104 @@ def _okey(v):
 
 
 @traced
-def stage_route(c, t):
-    """Token t (wave 0): 896 logits, 14 a lane (expert lane + 64 i) -> sigmoid
-    (v_rcp) + bias -> the top-16 by a radix select (the 16th largest key by a
-    32-step bit search of ballot counts; equal keys: the lower experts) -> the
-    16 ranked by score, descending (aiter's slot order) -> the pre-bias
-    sigmoids renormalized in that order (IEEE 1 / sum)."""
-    lane, wave, red = c["lane"], c["wave"], c["red"]
-    a = c["args"]
-    if wave == 0:
-        s = c["S"]
+def _route_out(c, t, red, sel, rank, me, ms):
+    """The picks' sigmoids in rank order (red[192 ..]), summed serially (aiter's
+    renorm, IEEE 1 / sum) -> ROUTE (expert, weight) at the pick's rank."""
+    if sel:
+        fx.ptr_store(ms, red + (192 + rank))
+    total = fx.Float32(0.0)
+    for j in range_constexpr(TOPK):
+        total = total + fx.ptr_load(red + (192 + j))
+    inv = fx.Float32(1.0) / total
+    if sel:
+        c["put"](
+            c["route"], t * 2 * TOPK + rank, fx.min(fx.max(me, 0), fx.Int32(E - 1))
+        )
+        c["put"](c["route"], t * 2 * TOPK + TOPK + rank, ms * inv)
+
+
+@traced
+def route_parts(c, t):
+    """(OPT_RLOG_PARTS) Token t's router partials (every rank's 4 K
+    parts, RLOGP) -> LDS `vals[part][expert]`, wave q < R_PARTS polling part
+    q (one batch: the 14 experts of its lane); ``stage_route`` sums each
+    expert's parts in part order, as the router's part 0 did."""
+    s, lane, wave = c["S"], c["lane"], c["wave"]
+    if wave < R_PARTS:
         per = E // 64
         specs = []
         for i in range_constexpr(per):
             e = lane + 64 * i
             src = e // (E // TP)
             specs.append(
-                (c["peer"]["rlog_own"], (src * s + t) * (E // TP) + e % (E // TP), 1)
+                (c["rlogp_own"], rlogp_pair(src, wave, s, t, e % (E // TP)), 1)
             )
-        got = c["poll"](specs)
+        got = c["poll"](specs, batch=per)
+        for i in range_constexpr(per):
+            fx.ptr_store(f32_of(got[i][0]), c["vals"] + (wave * E + lane + 64 * i))
+    gpu.barrier()
+
+
+def rlogp_pair(src, q, s, t, r):
+    """(OPT_RLOG_PARTS) RLOGP's pair of (source rank, K part q,
+    token t, the rank's router row r)."""
+    return ((src * R_PARTS + q) * s + t) * (E // TP) + r
+
+
+def rlogp_on(c):
+    return c.get("rlog_parts", False)
+
+
+def batch_on(c):
+    return c.get("batch", False)
+
+
+@traced
+def stage_route(c, t):
+    """Token t (wave 0): 896 logits, 14 a lane (expert lane + 64 i) -> sigmoid
+    (v_rcp) + bias -> the top-16 by (key desc, lower expert first) -> the
+    pre-bias sigmoids renormalized in that order (IEEE 1 / sum). The top-16 is
+    found from a lower bound, the 16th largest lane max (16 lanes hold a key >=
+    it: so does the 16th largest key), by one ballot a bit; the keys >= it (the
+    candidates, ~20-30) are compacted in expert order and each ranked against
+    the others (readlane), its rank the pick's slot. More than 64 candidates
+    (a warmup's NaN row): the full radix select over all 14 keys a lane."""
+    lane, wave, red = c["lane"], c["wave"], c["red"]
+    a = c["args"]
+    if wave == 0:
+        s = c["S"]
+        per = E // 64
+        # the bias loads in flight across the poll
+        bias = [bo_ld_f32(a["bias"], lane + 64 * i) for i in range(per)]
+        if const_expr(rlogp_on(c)):
+            # (OPT_RLOG_PARTS) The parts' sums from LDS (``route_parts``)
+            lgs = []
+            for i in range_constexpr(per):
+                e = lane + 64 * i
+                v = fx.ptr_load(c["vals"] + e)
+                for q in range_constexpr(1, R_PARTS):
+                    v = v + fx.ptr_load(c["vals"] + (q * E + e))
+                lgs.append(v)
+        else:
+            specs = []
+            for i in range_constexpr(per):
+                e = lane + 64 * i
+                src = e // (E // TP)
+                specs.append(
+                    (
+                        c["peer"]["rlog_own"],
+                        (src * s + t) * (E // TP) + e % (E // TP),
+                        1,
+                    )
+                )
+            # (OPT_BATCH) One batch (POLL_MAX would split the 14)
+            got = c["poll"](specs, batch=per) if batch_on(c) else c["poll"](specs)
+            lgs = [f32_of(got[i][0]) for i in range(per)]
         keys, sigs = [], []
         for i in range_constexpr(per):
-            e = lane + 64 * i
-            lg = f32_of(got[i][0])
+            lg = lgs[i]
             sg = hw_rcp(fx.Float32(1.0) + hw_exp2(lg * fx.Float32(-LOG2E)))
-            sc = sg + bo_ld_f32(a["bias"], e)
+            sc = sg + bias[i]
             # a NaN logit (a warmup's garbage row) ranks last
             keys.append(_okey((sc == sc).select(sc, fx.Float32(-3.0e38))))
             sigs.append((sg == sg).select(sg, fx.Float32(0.0)))
@@ -404,54 +620,80 @@ def count_ge(th):
                 n = n + popcount(ballot(keys[i] >= th))
             return n
 
-        th = fx.Uint32(0)
+        lm = keys[0]
+        for i in range_constexpr(1, per):
+            lm = (keys[i] > lm).select(keys[i], lm)
+        t0 = fx.Uint32(0)
         for bit in range_constexpr(31, -1, -1):
-            cand = th | fx.Uint32(1 << bit)
-            th = (count_ge(cand) >= TOPK).select(cand, th)
-        # selected: key > th, then the lowest experts with key == th
-        n_gt = fx.Int32(0)
-        for i in range_constexpr(per):
-            n_gt = n_gt + popcount(ballot(keys[i] > th))
-        need = fx.Int32(TOPK) - n_gt
-        eq_before = fx.Int32(0)
-        sel_before = fx.Int32(0)
-        for i in range_constexpr(per):
-            eqb = ballot(keys[i] == th)
-            eq_rank = eq_before + lanes_below(eqb, lane)
-            sel = (keys[i] > th) | ((keys[i] == th) & (eq_rank < need))
-            selb = ballot(sel)
-            pos = sel_before + lanes_below(selb, lane)
-            if sel:
-                fx.ptr_store(keys[i].bitcast(fx.Float32), red + pos * 3)
-                fx.ptr_store(
-                    fx.Int32(lane + 64 * i).bitcast(fx.Float32), red + (pos * 3 + 1)
-                )
-                fx.ptr_store(sigs[i], red + (pos * 3 + 2))
-            eq_before = eq_before + popcount(eqb)
-            sel_before = sel_before + popcount(selb)
-        # rank the 16 by key descending, the lower expert first on a tie
-        q = fx.min(lane, TOPK - 1)
-        mk = fx.Uint32(fx.ptr_load(red + q * 3).bitcast(fx.Int32))
-        me = fx.ptr_load(red + (q * 3 + 1)).bitcast(fx.Int32)
-        ms = fx.ptr_load(red + (q * 3 + 2))
-        rank = fx.Int32(0)
-        for j in range_constexpr(TOPK):
-            ok = fx.Uint32(fx.ptr_load(red + j * 3).bitcast(fx.Int32))
-            oe = fx.ptr_load(red + (j * 3 + 1)).bitcast(fx.Int32)
-            before = (ok > mk) | ((ok == mk) & (oe < me))
-            rank = rank + before.select(fx.Int32(1), fx.Int32(0))
-        # the sigmoids in rank order, summed serially (aiter's renorm)
-        if lane < TOPK:
-            fx.ptr_store(ms, red + (64 + rank))
-        total = fx.Float32(0.0)
-        for j in range_constexpr(TOPK):
-            total = total + fx.ptr_load(red + (64 + j))
-        inv = fx.Float32(1.0) / total
-        if lane < TOPK:
-            c["put"](
-                c["route"], t * 2 * TOPK + rank, fx.min(fx.max(me, 0), fx.Int32(E - 1))
-            )
-            c["put"](c["route"], t * 2 * TOPK + TOPK + rank, ms * inv)
+            cand = t0 | fx.Uint32(1 << bit)
+            t0 = (popcount(ballot(lm >= cand)) >= TOPK).select(cand, t0)
+        n_c = count_ge(t0)
+        if unlikely(n_c > 64):
+            # the 16th largest key by a 32-step bit search; selected: key > th,
+            # then the lowest experts with key == th
+            th = fx.Uint32(0)
+            for bit in range_constexpr(31, -1, -1):
+                cand = th | fx.Uint32(1 << bit)
+                th = (count_ge(cand) >= TOPK).select(cand, th)
+            n_gt = fx.Int32(0)
+            for i in range_constexpr(per):
+                n_gt = n_gt + popcount(ballot(keys[i] > th))
+            need = fx.Int32(TOPK) - n_gt
+            eq_before = fx.Int32(0)
+            sel_before = fx.Int32(0)
+            for i in range_constexpr(per):
+                eqb = ballot(keys[i] == th)
+                eq_rank = eq_before + lanes_below(eqb, lane)
+                sel = (keys[i] > th) | ((keys[i] == th) & (eq_rank < need))
+                selb = ballot(sel)
+                pos = sel_before + lanes_below(selb, lane)
+                if sel:
+                    fx.ptr_store(keys[i].bitcast(fx.Float32), red + pos * 3)
+                    fx.ptr_store(
+                        fx.Int32(lane + 64 * i).bitcast(fx.Float32),
+                        red + (pos * 3 + 1),
+                    )
+                    fx.ptr_store(sigs[i], red + (pos * 3 + 2))
+                eq_before = eq_before + popcount(eqb)
+                sel_before = sel_before + popcount(selb)
+            # rank the 16 by key descending, the lower expert first on a tie
+            q = fx.min(lane, TOPK - 1)
+            mk = fx.Uint32(fx.ptr_load(red + q * 3).bitcast(fx.Int32))
+            me = fx.ptr_load(red + (q * 3 + 1)).bitcast(fx.Int32)
+            ms = fx.ptr_load(red + (q * 3 + 2))
+            rank = fx.Int32(0)
+            for j in range_constexpr(TOPK):
+                ok = fx.Uint32(fx.ptr_load(red + j * 3).bitcast(fx.Int32))
+                oe = fx.ptr_load(red + (j * 3 + 1)).bitcast(fx.Int32)
+                before = (ok > mk) | ((ok == mk) & (oe < me))
+                rank = rank + before.select(fx.Int32(1), fx.Int32(0))
+            _route_out(c, t, red, lane < TOPK, rank, me, ms)
+        else:
+            # the candidates -> red (keys, experts, sigmoids), in expert order
+            cbefore = fx.Int32(0)
+            for i in range_constexpr(per):
+                cb = ballot(keys[i] >= t0)
+                pos = cbefore + lanes_below(cb, lane)
+                if keys[i] >= t0:
+                    fx.ptr_store(keys[i].bitcast(fx.Float32), red + pos)
+                    fx.ptr_store(
+                        fx.Int32(lane + 64 * i).bitcast(fx.Float32), red + (64 + pos)
+                    )
+                    fx.ptr_store(sigs[i], red + (128 + pos))
+                cbefore = cbefore + popcount(cb)
+            # lane j: candidate j, its rank = the candidates before it
+            q = fx.min(lane, n_c - 1)
+            mki = fx.ptr_load(red + q).bitcast(fx.Int32)
+            mk = fx.Uint32(mki)
+            me = fx.ptr_load(red + (64 + q)).bitcast(fx.Int32)
+            ms = fx.ptr_load(red + (128 + q))
+            rank = fx.Int32(0)
+            for j in range_constexpr(64):
+                ok = fx.Uint32(readlane(mki, j))
+                oe = readlane(me, j)
+                before = ((ok > mk) | ((ok == mk) & (oe < me))) & (j < n_c)
+                rank = rank + before.select(fx.Int32(1), fx.Int32(0))
+            _route_out(c, t, red, (lane < n_c) & (rank < TOPK), rank, me, ms)
     gpu.barrier()
 
 
@@ -459,70 +701,217 @@ def bo_ld_f32(ptr, i):
     return fx.Float32(bo.buffer_load(rsrc(ptr), i, vec_width=1, dtype=T.f32))
 
 
+def route_id(rt, j):
+    """Pick j's (t TOPK + k) expert, clamped to [0, E)."""
+    return fx.max(
+        fx.min(fx.ptr_load(rt["route"] + (j // TOPK * 2 * TOPK + j % TOPK)), E - 1), 0
+    )
+
+
 @traced
-def load_route(c):
-    """ROUTE -> LDS (ids, weights); the union U: ``slot_of[e]``,
-    ``expert_of[slot]``, ``kof[slot][t]`` (the token's top-k index, -1 if not
-    picked); returns |U|."""
+def poll_route(c):
+    """ROUTE -> LDS (ids, weights)."""
+    s, tid, rt = c["S"], c["tid"], c["rt"]
+    n = s * 2 * TOPK
+    per = (n + THREADS - 1) // THREADS
+    got = c["poll"](
+        [(c["route"], fx.min(tid + THREADS * i, n - 1), 1) for i in range(per)]
+    )
+    for i in range_constexpr(per):
+        if tid + THREADS * i < n:
+            fx.ptr_store(got[i][0], rt["route"] + (tid + THREADS * i))
+
+
+@traced
+def load_route_few(c):
+    """``few_rows``: ROUTE -> LDS, then U, the experts with picks (ascending):
+    each expert's slot (``eslot``, a flag first), each slot's expert, and
+    ``kof[slot][t]``, token t's top-k index of the slot's expert (-1: not
+    picked). Returns |U|, the pick tiles' count."""
     s, tid, lane, wave = c["S"], c["tid"], c["lane"], c["wave"]
     rt = c["rt"]
-    n = s * 2 * TOPK
-    if tid < n:
-        v = c["poll"]([(c["route"], tid, 1)])[0][0]
-        fx.ptr_store(v, rt["route"] + tid)
+    npk = s * TOPK
+    assert npk <= THREADS and 2 * THREADS >= E
+    poll_route(c)
     for i in range_constexpr((E + THREADS - 1) // THREADS):
         e = tid + THREADS * i
         if e < E:
-            fx.ptr_store(fx.Int32(0), rt["flag"] + e)
-    for i in range_constexpr((MAX_U * 8 + THREADS - 1) // THREADS):
+            fx.ptr_store(fx.Int32(0), rt["eslot"] + e)
+    nk = max_u(s) * s
+    for i in range_constexpr((nk + THREADS - 1) // THREADS):
         j = tid + THREADS * i
-        if j < MAX_U * 8:
+        if j < nk:
             fx.ptr_store(fx.Int32(-1), rt["kof"] + j)
     gpu.barrier()
-    if tid < s * TOPK:
-        t = tid // TOPK
-        k = tid % TOPK
-        e = fx.max(fx.min(fx.ptr_load(rt["route"] + (t * 2 * TOPK + k)), E - 1), 0)
-        fx.ptr_store(fx.Int32(1), rt["flag"] + e)
+    if tid < npk:
+        fx.ptr_store(fx.Int32(1), rt["eslot"] + route_id(rt, tid))
     gpu.barrier()
     e0 = fx.min(2 * tid, E - 1)
-    f0 = (2 * tid < E).select(fx.ptr_load(rt["flag"] + e0), fx.Int32(0))
+    f0 = (2 * tid < E).select(fx.ptr_load(rt["eslot"] + e0), fx.Int32(0))
     f1 = (2 * tid + 1 < E).select(
-        fx.ptr_load(rt["flag"] + fx.min(2 * tid + 1, E - 1)), fx.Int32(0)
+        fx.ptr_load(rt["eslot"] + fx.min(2 * tid + 1, E - 1)), fx.Int32(0)
     )
     off, tot = block_excl_scan(f0 + f1, lane, wave, c["scan"])
     gpu.barrier()
     if f0 == 1:
-        fx.ptr_store(off, rt["flag"] + 2 * tid)
+        fx.ptr_store(off, rt["eslot"] + 2 * tid)
         fx.ptr_store(fx.Int32(2 * tid), rt["expert"] + off)
     if f1 == 1:
-        fx.ptr_store(off + f0, rt["flag"] + 2 * tid + 1)
+        fx.ptr_store(off + f0, rt["eslot"] + 2 * tid + 1)
         fx.ptr_store(fx.Int32(2 * tid + 1), rt["expert"] + off + f0)
     gpu.barrier()
-    if tid < s * TOPK:
-        t = tid // TOPK
-        k = tid % TOPK
-        e = fx.max(fx.min(fx.ptr_load(rt["route"] + (t * 2 * TOPK + k)), E - 1), 0)
-        slot = fx.ptr_load(rt["flag"] + e)
-        fx.ptr_store(fx.Int32(k), rt["kof"] + (slot * 8 + t))
+    if tid < npk:
+        slot = fx.ptr_load(rt["eslot"] + route_id(rt, tid))
+        fx.ptr_store(tid % TOPK, rt["kof"] + (slot * s + tid // TOPK))
     gpu.barrier()
     return uniform(tot)
 
 
+@traced
+def load_route(c):
+    """ROUTE -> LDS (ids, weights), then U, the experts with picks (ascending),
+    and their pick tiles: each expert's slot, each slot's expert, first pick
+    and first tile, the picks (t TOPK + k) in slot order (a slot's tokens
+    ascending) and each tile's slot. A pick ORs its token's bit into its
+    expert's bitmap (``redi``): a slot's picks are the bitmap's popcount, a
+    pick's rank in its slot the bits below its token. The bitmaps are the same
+    whatever order the ORs land in, so every CTA lays out the same table.
+    Returns the pick tiles' count. At ``few_rows``, ``load_route_few``."""
+    s, tid, lane, wave = c["S"], c["tid"], c["lane"], c["wave"]
+    if const_expr(few_rows(s)):
+        return load_route_few(c)
+    rt, mask = c["rt"], c["redi"]
+    assert s <= MAX_S and 2 * THREADS >= E and MASK_WORDS * E <= WAVES * 256
+    poll_route(c)
+    for i in range_constexpr((MASK_WORDS * E + THREADS - 1) // THREADS):
+        j = tid + THREADS * i
+        if j < MASK_WORDS * E:
+            fx.ptr_store(fx.Int32(0), mask + j)
+    gpu.barrier()
+    npk = s * TOPK
+    for i in range_constexpr((npk + THREADS - 1) // THREADS):
+        j = tid + THREADS * i
+        if j < npk:
+            t = j // TOPK
+            fx.llvm.atomic_or(
+                mask + (route_id(rt, j) * MASK_WORDS + t // 32),
+                fx.Int32(1) << (t % 32),
+                syncscope="workgroup",
+            )
+    gpu.barrier()
+    # experts 2 tid, 2 tid + 1: one scan of (slots, picks, tiles) packed in
+    # 10, 11 and 10 bits (totals <= 896, 1024, 960)
+    assert E < 1 << 10 and npk < 1 << 11 and max_tiles(s) < 1 << 10
+    cnt = []
+    for j in range_constexpr(2):
+        e = fx.min(2 * tid + j, E - 1)
+        n_e = fx.Int32(0)
+        for w in range_constexpr(MASK_WORDS):
+            n_e = n_e + popcount(fx.ptr_load(mask + (e * MASK_WORDS + w)))
+        cnt.append((2 * tid + j < E).select(n_e, fx.Int32(0)))
+    packed = [
+        (n_e > 0).select(fx.Int32(1), fx.Int32(0))
+        + (n_e << 10)
+        + (((n_e + PTILE - 1) // PTILE) << 21)
+        for n_e in cnt
+    ]
+    off, tot = block_excl_scan(packed[0] + packed[1], lane, wave, c["scan"])
+    for j in range_constexpr(2):
+        e = 2 * tid + j
+        at = off if j == 0 else off + packed[0]
+        slot = at & 1023
+        tile0 = (at >> 21) & 1023
+        if cnt[j] > 0:
+            fx.ptr_store(slot, rt["eslot"] + e)
+            fx.ptr_store(fx.Int32(e), rt["expert"] + slot)
+            fx.ptr_store((at >> 10) & 2047, rt["first"] + slot)
+            fx.ptr_store(tile0, rt["tfirst"] + slot)
+            for q in range_constexpr((s + PTILE - 1) // PTILE):
+                if q * PTILE < cnt[j]:
+                    fx.ptr_store(slot, rt["tslot"] + (tile0 + q))
+    if tid == 0:
+        fx.ptr_store(fx.Int32(npk), rt["first"] + (tot & 1023))
+        fx.ptr_store((tot >> 21) & 1023, rt["tfirst"] + (tot & 1023))
+    gpu.barrier()
+    for i in range_constexpr((npk + THREADS - 1) // THREADS):
+        j = tid + THREADS * i
+        if j < npk:
+            t = j // TOPK
+            e = route_id(rt, j)
+            slot = fx.ptr_load(rt["eslot"] + e)
+            # the tokens below t that picked e: the picks before j in its slot
+            rank = fx.Int32(0)
+            for w in range_constexpr(MASK_WORDS):
+                word = fx.ptr_load(mask + (e * MASK_WORDS + w))
+                lo = t - 32 * w
+                keep = (lo >= 32).select(
+                    fx.Int32(-1),
+                    (lo > 0).select((fx.Int32(1) << fx.max(lo, 0)) - 1, fx.Int32(0)),
+                )
+                rank = rank + popcount(word & keep)
+            fx.ptr_store(j, rt["order"] + (fx.ptr_load(rt["first"] + slot) + rank))
+    gpu.barrier()
+    return uniform((tot >> 21) & 1023)
+
+
+def pick_tile(c, tile):
+    """Pick tile ``tile``: (its slot, its first pick's position in the pick
+    order, its picks)."""
+    rt = c["rt"]
+    slot = uniform(fx.ptr_load(rt["tslot"] + tile))
+    first = fx.ptr_load(rt["first"] + slot)
+    p0 = first + (tile - fx.ptr_load(rt["tfirst"] + slot)) * PTILE
+    n = fx.min(fx.ptr_load(rt["first"] + (slot + 1)) - p0, PTILE)
+    return slot, uniform(p0), uniform(n)
+
+
+@traced
+def load_latq4(c):
+    """``load_latq`` at ``few_rows``: the producer quantized
+    (``latent_mx_out``); a thread a block's 4 words and its code, one poll batch."""
+    s, tid = c["S"], c["tid"]
+    nblk = s * (LAT // 32)
+    its = (nblk + THREADS - 1) // THREADS
+    bs = [fx.min(tid + THREADS * i, nblk - 1) for i in range(its)]
+    lq = c["peer"]["latq4_own"]
+    ws = [b // (LAT // 32) * (LAT // 8) + b % (LAT // 32) * 4 for b in bs]
+    specs = [
+        sp
+        for b, w0 in zip(bs, ws)
+        for sp in ((lq, w0, 2), (lq, w0 + 2, 2), (lq, s * LAT // 8 + b, 1))
+    ]
+    got = c["poll"](specs, batch=3 * its)
+    for i in range_constexpr(its):
+        if tid + THREADS * i < nblk:
+            v = fx.Vector.from_elements(got[3 * i] + got[3 * i + 1], fx.Int32)
+            fx.ptr_store(v, c["latq"] + ws[i])
+            fx.ptr_store(got[3 * i + 2][0], c["latsc"] + bs[i])
+    gpu.barrier()
+
+
 @traced
 def load_latq(c):
     """LAT (every rank's 448 columns) -> MXFP4 as aiter's inline quant: a
     thread a (token, 32-block): amax of |bf16| bits, the ceil code of amax / 6,
-    v_cvt_scalef32_pk_fp4_bf16 -> latq words, latsc codes."""
+    v_cvt_scalef32_pk_fp4_bf16 -> latq words, latsc codes (a word each at
+    ``few_rows``, else 4 a word)."""
     s, tid = c["S"], c["tid"]
     nblk = s * (LAT // 32)
+    if const_expr(not few_rows(s)):
+        # latsc: a word 4 codes (K blocks 4 st ..), ORed in
+        for i in range_constexpr((nblk // 4 + THREADS - 1) // THREADS):
+            j = tid + THREADS * i
+            if j < nblk // 4:
+                fx.ptr_store(fx.Int32(0), c["latsc"] + j)
+        gpu.barrier()
     for i in range_constexpr((nblk + THREADS - 1) // THREADS):
         b = fx.min(tid + THREADS * i, nblk - 1)
         t = b // (LAT // 32)
         kb = b % (LAT // 32)
         base = (t * LAT + kb * 32) // 2
-        got = c["poll"]([(c["peer"]["lat_own"], base + j, 1) for j in range(16)])
-        ws = [got[j][0] for j in range(16)]
+        # the block's 16 pairs as 8 x 16 B loads: one poll batch
+        got = c["poll"]([(c["peer"]["lat_own"], base + 2 * j, 2) for j in range(8)])
+        ws = [got[j // 2][j % 2] for j in range(16)]
         am = fx.Int32(0)
         for w in ws:
             am = fx.max(am, w & 0x7FFF)
@@ -536,7 +925,14 @@ def load_latq(c):
                 for j in range_constexpr(4):
                     pk = pk_fp4_bf16(pk, ws[4 * q + j], scale, j)
                 fx.ptr_store(pk, c["latq"] + (t * (LAT // 8) + kb * 4 + q))
-            fx.ptr_store(code, c["latsc"] + (t * (LAT // 32) + kb))
+            if const_expr(few_rows(s)):
+                fx.ptr_store(code, c["latsc"] + (t * (LAT // 32) + kb))
+            else:
+                fx.llvm.atomic_or(
+                    c["latsc"] + (t * (LAT // 128) + kb // 4),
+                    code << ((kb % 4) * 8),
+                    syncscope="workgroup",
+                )
     gpu.barrier()
 
 
@@ -559,7 +955,10 @@ def ug_loads(c, task):
     and their scale words: issued a task ahead (``stage_ug`` takes them)."""
     lane, wave = c["lane"], c["wave"]
     a, rt = c["args"], c["rt"]
-    slot = task // UG_GROUPS
+    if const_expr(few_rows(c["S"])):
+        slot = task // UG_GROUPS
+    else:
+        slot = uniform(fx.ptr_load(rt["tslot"] + task // UG_GROUPS))
     g = task % UG_GROUPS
     e = uniform(fx.ptr_load(rt["expert"] + slot))
     rgl = wave % 4  # 0, 1 gate; 2, 3 up
@@ -583,13 +982,18 @@ def ug_loads(c, task):
 
 @traced
 def stage_ug(c, task, ops):
-    """Slot task / 12, 32-column group task % 12: gate rows 32 g .. and up rows
-    384 + 32 g .. of the slot's expert (``ops``: ``ug_loads``) against the MXFP4
-    latent; SiTUv2; MXFP4 a 32-column group -> INTER / ISC of each picking
-    token."""
+    """Pick tile task / 12, 32-column group task % 12: gate rows 32 g .. and up
+    rows 384 + 32 g .. of the tile's expert (``ops``: ``ug_loads``) against its
+    picks' MXFP4 latent rows (a lane's column a pick); SiTUv2; MXFP4 a 32-column
+    group -> INTER / ISC of each pick. At ``few_rows`` the slot's tile is
+    every token (a lane's column a token), its picks those with a ``kof``."""
     s, tid, lane, wave, red = c["S"], c["tid"], c["lane"], c["wave"], c["red"]
     rt = c["rt"]
-    slot = task // UG_GROUPS
+    few = few_rows(s)
+    if const_expr(few):
+        slot = task // UG_GROUPS
+    else:
+        _, p0, n = pick_tile(c, task // UG_GROUPS)
     g = task % UG_GROUPS
     rgl = wave % 4
     half = wave // 4
@@ -597,7 +1001,10 @@ def stage_ug(c, task, ops):
     steps = [half * (W13_STEPS // 2) + i for i in range(W13_STEPS // 2)]
     av = ops[: len(steps)]
     sw = ops[len(steps) :]
-    t = fx.min(lane % 16, s - 1)
+    if const_expr(few):
+        t = fx.min(lane % 16, s - 1)
+    else:
+        t = fx.ptr_load(rt["order"] + (p0 + fx.min(lane % 16, n - 1))) // TOPK
     kb = lane // 16
     acc = fx.Vector.filled(4, 0.0, fx.Float32)
     for i in range_constexpr(len(steps)):
@@ -608,17 +1015,26 @@ def stage_ug(c, task, ops):
                 result_type=fx.Vector.make_type(4, fx.Int32),
             )
         )
-        sb = fx.ptr_load(c["latsc"] + (t * (LAT // 32) + st * 4 + kb))
+        if const_expr(few):
+            sb = fx.ptr_load(c["latsc"] + (t * (LAT // 32) + st * 4 + kb))
+        else:
+            sb = (fx.ptr_load(c["latsc"] + (t * (LAT // 128) + st)) >> (kb * 8)) & 0xFF
         # the code's byte (st % 2) 2 + rg % 2 of the word; rg % 2 is the wave's
         sa = (sw[i // 2] >> ((st % 2) * 16 + (rg % 2) * 8)) & 0xFF
         acc = mfma_scaled(av[i], bw, acc, sa, sb, FP4, FP4, 0, 0)
     fx.ptr_store(acc, red + (wave * 64 + lane) * 4)
     gpu.barrier()
-    # thread (token, column): g / u summed over the K halves, SiTUv2
+    # thread (pick, column): g / u summed over the K halves, SiTUv2
     tt = tid // 32
     col = tid % 32
-    live = tt < s
-    tt = fx.min(tt, s - 1)
+    if const_expr(few):
+        live = tt < s
+        tt = fx.min(tt, s - 1)
+        k = fx.ptr_load(rt["kof"] + (slot * s + tt))
+        live = live & (k >= 0)
+    else:
+        live = tt < n
+        tt = fx.min(tt, n - 1)
     first = col < 16  # gate / up row group 0 (waves 0, 2) or 1 (1, 3); +4: K's 2nd half
     r = col % 16
     gv = first.select(
@@ -630,18 +1046,21 @@ def stage_ug(c, task, ops):
     v = situ_aiter(gv, uv)
     amax = butterfly(abs(v), (16, 8, 4, 2, 1), fx.max)
     code = e8m0_roundup(amax)
+    # the values through red, its MFMA sums read
     gpu.barrier()
-    fx.ptr_store(v, c["av"] + tid)
+    fx.ptr_store(v, red + tid)
     gpu.barrier()
-    k = fx.ptr_load(rt["kof"] + (slot * 8 + tt))
-    if live & (col < 4) & (k >= 0):
+    if live & (col < 4):
         scale = (code << 23).bitcast(fx.Float32)
         pk = fx.Int32(0)
         for j in range_constexpr(4):
-            x0 = fx.ptr_load(c["av"] + (tt * 32 + col * 8 + 2 * j))
-            x1 = fx.ptr_load(c["av"] + (tt * 32 + col * 8 + 2 * j + 1))
+            x0 = fx.ptr_load(red + (tt * 32 + col * 8 + 2 * j))
+            x1 = fx.ptr_load(red + (tt * 32 + col * 8 + 2 * j + 1))
             pk = pk_fp4_f32(pk, x0, x1, scale, j)
-        pick = tt * TOPK + k
+        if const_expr(few):
+            pick = tt * TOPK + k
+        else:
+            pick = fx.ptr_load(rt["order"] + (p0 + tt))
         c["put_words"](c["inter"], pick * 48 + g * 4 + col, [pk])
         if col == 0:
             c["put"](c["isc"], pick * UG_GROUPS + g, code)
@@ -655,17 +1074,23 @@ def load_inter(c):
     s, tid = c["S"], c["tid"]
     nw = s * TOPK * 48
     per = (nw + THREADS - 1) // THREADS
-    got = c["poll"](
-        [(c["inter"], fx.min(tid + THREADS * i, nw - 1), 1) for i in range(per)]
-    )
+    nc = s * TOPK * UG_GROUPS
+    perc = (nc + THREADS - 1) // THREADS
+    ispecs = [(c["inter"], fx.min(tid + THREADS * i, nw - 1), 1) for i in range(per)]
+    cspecs = [(c["isc"], fx.min(tid + THREADS * i, nc - 1), 1) for i in range(perc)]
+    if const_expr(batch_on(c) and per + perc <= 16):
+        # (OPT_BATCH) INTER and ISC in one poll batch (one round
+        # trip once the last ug task landed, not two)
+        both = c["poll"](ispecs + cspecs, batch=per + perc)
+        got, gc = both[:per], both[per:]
+    else:
+        got = c["poll"](ispecs)
+        gc = None
     for i in range_constexpr(per):
         if tid + THREADS * i < nw:
             fx.ptr_store(got[i][0], c["interq"] + (tid + THREADS * i))
-    nc = s * TOPK * UG_GROUPS
-    perc = (nc + THREADS - 1) // THREADS
-    gc = c["poll"](
-        [(c["isc"], fx.min(tid + THREADS * i, nc - 1), 1) for i in range(perc)]
-    )
+    if const_expr(gc is None):
+        gc = c["poll"](cspecs)
     for i in range_constexpr(perc):
         if tid + THREADS * i < nc:
             fx.ptr_store(gc[i][0], c["intsc"] + (tid + THREADS * i))
@@ -697,14 +1122,20 @@ def poll_inter(c, pick, kb):
     )
 
 
-def down_loads(c, rgd, nu, b0):
-    """The w2 tiles and scale codes of a wave's batch from slot b0 (slots b0,
-    b0 + 8, ...; past |U| the last slot's): flat, ``stage_down`` unpacks."""
+def down_loads(c, rgd, nt, b0):
+    """The w2 tiles and scale codes of a wave's batch from pick tile b0 (tiles
+    b0, b0 + 8, ...; past the last, the last's): flat, ``stage_down`` unpacks."""
     lane = c["lane"]
     r_w = rsrc(c["args"]["w2"])
     rt = c["rt"]
-    slots = [fx.min(b0 + WAVES * j, nu - 1) for j in range(DOWN_BATCH)]
-    es = [uniform(fx.ptr_load(rt["expert"] + sl)) for sl in slots]
+    tiles = [fx.min(b0 + WAVES * j, nt - 1) for j in range(DOWN_BATCH)]
+    if const_expr(few_rows(c["S"])):
+        es = [uniform(fx.ptr_load(rt["expert"] + tl)) for tl in tiles]
+    else:
+        es = [
+            uniform(fx.ptr_load(rt["expert"] + fx.ptr_load(rt["tslot"] + tl)))
+            for tl in tiles
+        ]
     av = [
         fp4_tile_load(r_w, es[j] * (LAT * RI // 8), rgd, RI, st, lane * 4)
         for j in range(DOWN_BATCH)
@@ -719,39 +1150,54 @@ def down_loads(c, rgd, nu, b0):
 
 
 @traced
-def stage_down(c, rgd, nu, ops0):
-    """Latent rows 16 rgd ..: every slot's w2 against its picks (wave w the slots
-    w, w + 8, ...; a batch of DOWN_BATCH in flight, the next issued under this
-    one's MFMAs; ``ops0`` the first, issued before INTER was waited for) ->
-    bf16(w acc) a (token, k) in LDS -> folded in top-k order with a bf16
-    rounding an add -> every rank's LATR."""
+def stage_down(c, rgd, nt, ops0):
+    """Latent rows 16 rgd ..: every pick tile's w2 against its picks (wave w
+    the tiles w, w + 8, ...; a batch of DOWN_BATCH in flight, the next issued
+    under this one's MFMAs; ``ops0`` the first, issued before INTER was waited
+    for) -> bf16(w acc) a (token, k) in LDS -> folded in top-k order with a
+    bf16 rounding an add -> every rank's LATR. INTER from LDS (``inter_lds``),
+    else a tile's polled at its use."""
     s, tid, lane, wave = c["S"], c["tid"], c["lane"], c["wave"]
     rt = c["rt"]
-    t = fx.min(lane % 16, s - 1)
     kb = lane // 16
     nbs = DOWN_BATCH * W2_STEPS
     b0 = wave
     ops = ops0
-    while b0 < nu:
+    while b0 < nt:
         nb = b0 + WAVES * DOWN_BATCH
-        ops_n = down_loads(c, rgd, nu, fx.min(nb, nu - 1))
-        slots = [fx.min(b0 + WAVES * j, nu - 1) for j in range(DOWN_BATCH)]
+        ops_n = down_loads(c, rgd, nt, fx.min(nb, nt - 1))
         for j in range_constexpr(DOWN_BATCH):
-            live = (b0 + WAVES * j) < nu
-            k = fx.ptr_load(rt["kof"] + (slots[j] * 8 + t))
-            picked = (k >= 0) & (lane % 16 < s)
-            pick = t * TOPK + fx.max(k, 0)
+            live = (b0 + WAVES * j) < nt
+            if const_expr(few_rows(s)):
+                # the slot's tile: lane % 16 a token, picked if it has a kof
+                tq = fx.min(lane % 16, s - 1)
+                kq = fx.ptr_load(rt["kof"] + (fx.min(b0 + WAVES * j, nt - 1) * s + tq))
+                picked = (kq >= 0) & (lane % 16 < s)
+                t, k = tq, fx.max(kq, 0)
+                pick = t * TOPK + k
+            else:
+                _, p0, cnt = pick_tile(c, fx.min(b0 + WAVES * j, nt - 1))
+                picked = lane % 16 < cnt
+                pick = fx.ptr_load(rt["order"] + (p0 + fx.min(lane % 16, cnt - 1)))
+                # a pick is >= 0: shifts, not signed division
+                t, k = pick >> TOPK_BITS, pick & (TOPK - 1)
+            if const_expr(not inter_lds(s)):
+                got = poll_inter(c, pick, kb)
             acc = fx.Vector.filled(4, 0.0, fx.Float32)
-            # this slot's picks' INTER, polled here (not all of U first): a slot
-            # whose ug tasks are done is multiplied while later ones still run
             for st in range_constexpr(W2_STEPS):
-                bw = fx.Vector(
-                    fx.ptr_load(
-                        c["interq"] + (pick * 48 + st * 16 + kb * 4),
-                        result_type=fx.Vector.make_type(4, fx.Int32),
+                if const_expr(inter_lds(s)):
+                    bw = fx.Vector(
+                        fx.ptr_load(
+                            c["interq"] + (pick * 48 + st * 16 + kb * 4),
+                            result_type=fx.Vector.make_type(4, fx.Int32),
+                        )
                     )
-                )
-                scode = fx.ptr_load(c["intsc"] + (pick * UG_GROUPS + st * 4 + kb))
+                    scode = fx.ptr_load(c["intsc"] + (pick * UG_GROUPS + st * 4 + kb))
+                else:
+                    bw = fx.Vector.from_elements(
+                        [got[st * 4 + d][0] for d in range(4)], fx.Int32
+                    )
+                    scode = got[W2_STEPS * 4 + st][0]
                 bw = fx.Vector.from_elements(
                     [picked.select(bw[d], fx.Int32(0)) for d in range(4)], fx.Int32
                 )
@@ -767,7 +1213,7 @@ def stage_down(c, rgd, nu, ops0):
                     0,
                     0,
                 )
-            # lane: rows 4 (lane / 16) + i of token lane % 16
+            # lane: rows 4 (lane / 16) + i of pick lane % 16
             if live & picked:
                 wt = fx.ptr_load(rt["route"] + (t * 2 * TOPK + TOPK + k)).bitcast(
                     fx.Float32
@@ -775,30 +1221,36 @@ def stage_down(c, rgd, nu, ops0):
                 for i in range_constexpr(4):
                     row = 4 * (lane // 16) + i
                     fx.ptr_store(
-                        bf16_round(acc[i] * wt),
-                        c["contrib"] + ((t * TOPK + k) * ROWS + row),
+                        bf16_round(acc[i] * wt), c["contrib"] + (pick * ROWS + row)
                     )
         b0 = nb
         ops = ops_n
     gpu.barrier()
-    if tid < (ROWS // 2) * s:
-        rp = tid % (ROWS // 2)
-        tt = tid // (ROWS // 2)
-        y0 = fx.Float32(0.0)
-        y1 = fx.Float32(0.0)
-        for k in range_constexpr(TOPK):
-            y0 = bf16_round(
-                y0 + fx.ptr_load(c["contrib"] + ((tt * TOPK + k) * ROWS + 2 * rp))
-            )
-            y1 = bf16_round(
-                y1 + fx.ptr_load(c["contrib"] + ((tt * TOPK + k) * ROWS + 2 * rp + 1))
-            )
-        for p in range_constexpr(TP):
-            c["put_bf"](
-                c["peer"]["latr"][p],
-                (c["rank"] * s + tt) * LAT + rgd * ROWS + 2 * rp,
-                [y0, y1],
-            )
+    n = (ROWS // 2) * s
+    for i in range_constexpr((n + THREADS - 1) // THREADS):
+        idx = tid + THREADS * i
+        if idx < n:
+            rp = idx % (ROWS // 2)
+            tt = idx // (ROWS // 2)
+            y0 = fx.Float32(0.0)
+            y1 = fx.Float32(0.0)
+            for k in range_constexpr(TOPK):
+                at = (tt * TOPK + k) * ROWS + 2 * rp
+                y0 = bf16_round(y0 + fx.ptr_load(c["contrib"] + at))
+                y1 = bf16_round(y1 + fx.ptr_load(c["contrib"] + (at + 1)))
+            if const_expr(sent_on(c, "latr")):
+                # (OPT_SENT_LATR) Plain words to the sentinel region
+                si = sent_ix(c["rank"], rgd * ROWS + 2 * rp, tt, s, LAT)
+                word = sent_word(bf16_pair(y0, y1))
+                for p in range_constexpr(TP):
+                    put_sent(c["latrs_peer"][p], si, word)
+            else:
+                for p in range_constexpr(TP):
+                    c["put_bf"](
+                        c["peer"]["latr"][p],
+                        (c["rank"] * s + tt) * LAT + rgd * ROWS + 2 * rp,
+                        [y0, y1],
+                    )
     gpu.barrier()
 
 
@@ -843,18 +1295,39 @@ def load_h(c):
 def stage_sdown(c, j, ops):
     """Shared down rows 16 j .. (K 768) against h in LDS (``load_h``): bf16
     partial -> SHP (this rank's up columns) or every rank's FINAL."""
+    s = c["S"]
+    for i, (t0, n) in enumerate(gemv.token_tiles(s)):
+        if const_expr(i > 0):
+            gpu.barrier()  # the last tile's red read
+        sdown_tile(c, j, ops, t0, n)
+    lo = c["rank"] * UP_TASKS
+    mine = (j >= lo) & (j < lo + UP_TASKS)
+    publish(
+        c["put"],
+        c["shrdy"],
+        fx.max(fx.min(j - lo, UP_TASKS - 1), 0),
+        fx.Int32(1),
+        (c["tid"] == 0) & mine,
+    )
+    gpu.barrier()
+
+
+@traced
+def sdown_tile(c, j, ops, t0, n):
+    """``stage_sdown`` for the h rows t0 .. t0 + n."""
     s, tid, lane, wave, red = c["S"], c["tid"], c["lane"], c["wave"], c["red"]
     hrow = SI + gemv.LDS_PAD
-    acc = gemv.mfmas(lane, wave, c["xl"], hrow, SI, s, ops)
+    acc = gemv.mfmas(lane, wave, c["xl"] + t0 * hrow // 2, hrow, SI, n, ops)
     fx.ptr_store(acc, red + (wave * 64 + lane) * 4)
     gpu.barrier()
     lo = c["rank"] * UP_TASKS
     mine = (j >= lo) & (j < lo + UP_TASKS)
-    if tid < (ROWS // 2) * s:
+    if tid < (ROWS // 2) * n:
         rp = tid % (ROWS // 2)
-        tt = tid // (ROWS // 2)
-        v0 = bf16_round(row_sum(red, 2 * rp, tt))
-        v1 = bf16_round(row_sum(red, 2 * rp + 1, tt))
+        tl = tid // (ROWS // 2)
+        tt = t0 + tl
+        v0 = bf16_round(row_sum(red, 2 * rp, tl))
+        v1 = bf16_round(row_sum(red, 2 * rp + 1, tl))
         if mine:
             w = (
                 fx.Vector.from_elements([v0, v1], fx.Float32)
@@ -868,20 +1341,22 @@ def stage_sdown(c, j, ops):
                 cache_modifier=CM_DEV,
             )
         else:
-            for p in range_constexpr(TP):
-                c["put_bf"](
-                    c["peer"]["final"][p],
-                    (c["rank"] * s + tt) * HIDDEN + j * ROWS + 2 * rp,
-                    [v0, v1],
-                )
-    publish(
-        c["put"],
-        c["shrdy"],
-        fx.max(fx.min(j - lo, UP_TASKS - 1), 0),
-        fx.Int32(1),
-        (tid == 0) & mine,
-    )
-    gpu.barrier()
+            at = (c["rank"] * s + tt) * HIDDEN + j * ROWS + 2 * rp
+            if const_expr(sent_on(c, "final")):
+                # (OPT_SENT_FINAL) Plain words to the sentinel region
+                si = sent_ix(c["rank"], j * ROWS + 2 * rp, tt, s, HIDDEN)
+                word = sent_word(bf16_pair(v0, v1))
+                if const_expr(two_shot(s)):
+                    put_sent(peer_at(c["finals_peer"], j // UP_TASKS), si, word)
+                else:
+                    for p in range_constexpr(TP):
+                        put_sent(c["finals_peer"][p], si, word)
+            elif const_expr(two_shot(s)):
+                # to the columns' owner only: its up reduces them
+                c["put_bf"](peer_at(c["peer"]["final"], j // UP_TASKS), at, [v0, v1])
+            else:
+                for p in range_constexpr(TP):
+                    c["put_bf"](c["peer"]["final"][p], at, [v0, v1])
 
 
 @traced
@@ -897,6 +1372,28 @@ def run_sdown(c, bid, w_sd):
             stage_sdown(c, j, gemv.loads(c["lane"], c["wave"], w_sd, j * ROWS, SI))
 
 
+@traced
+def run_sdown_first(c, bid, w_sd):
+    """(OPT_SDOWN_FIRST) ``run_sdown``'s tasks (the up stage polls
+    their SHRDY) on CTAs 224 .., each CTA's two tasks' weights issued before H
+    is polled, ahead of its ``run_sdown_rest`` tasks (which only the final
+    reduce waits on). Placement only: the same tasks, the same arithmetic."""
+    lo = c["rank"] * UP_TASKS
+    nsd = BLOCKS - DOWN_TASKS  # 32
+    if bid >= DOWN_TASKS:
+        k0 = bid - DOWN_TASKS
+        k1 = fx.min(k0 + nsd, UP_TASKS - 1)
+        j0 = (k0 + lo) % SD_TASKS
+        j1 = (k1 + lo) % SD_TASKS
+        ops0 = gemv.loads(c["lane"], c["wave"], w_sd, j0 * ROWS, SI)
+        ops1 = gemv.loads(c["lane"], c["wave"], w_sd, j1 * ROWS, SI)
+        load_h(c)
+        stage_sdown(c, j0, ops0)
+        if k0 + nsd < UP_TASKS:
+            stage_sdown(c, j1, ops1)
+    assert UP_TASKS <= 2 * (BLOCKS - DOWN_TASKS)
+
+
 @traced
 def run_sdown_rest(c, bid, w_sd):
     """The other 392 shared-down tasks (only the final reduce waits on them), on
@@ -904,81 +1401,167 @@ def run_sdown_rest(c, bid, w_sd):
     up critical path."""
     r0 = LN0 + LN_TASKS  # 84: past the lnorm / up CTAs
     if bid >= r0:
-        load_h(c)
         lo = c["rank"] * UP_TASKS
-        for k in range(UP_TASKS + bid - r0, SD_TASKS, BLOCKS - r0):
-            j = (k + lo) % SD_TASKS
-            stage_sdown(c, j, gemv.loads(c["lane"], c["wave"], w_sd, j * ROWS, SI))
+        last = SD_TASKS - 1
+
+        def loads(k):
+            j = (fx.min(k, last) + lo) % SD_TASKS
+            return gemv.loads(c["lane"], c["wave"], w_sd, j * ROWS, SI)
+
+        # a task's weights issued a task ahead (the first before H is polled)
+        k = UP_TASKS + bid - r0
+        ops = loads(k)
+        load_h(c)
+        while k < SD_TASKS:
+            kn = k + (BLOCKS - r0)
+            ops_n = loads(kn)
+            stage_sdown(c, (k + lo) % SD_TASKS, ops)
+            k = kn
+            ops = ops_n
 
 
 # ---------------------------------------------------------------- lnorm / up / reduce
 @traced
-def stage_lnorm(c, j):
-    """Columns 128 j ..: the ranks' LATR (rank order, fp32 -> bf16), every
-    task's sum of squares -> rstd -> bf16(acc rstd w) -> LN, LNRDY."""
-    s, tid, lane, red = c["S"], c["tid"], c["lane"], c["red"]
+def stage_lnorm(c, t):
+    """Token t's latent row: the ranks' LATR (rank order, fp32 -> bf16; wave w
+    the 128-column chunks w, w + 8, ..), a chunk's sum of squares, summed in
+    chunk order -> rstd -> bf16(acc rstd w) -> LN row t, LNRDY[t]. The whole
+    row on one CTA: no partial sums of squares to exchange."""
+    s, tid, lane, wave, red = c["S"], c["tid"], c["lane"], c["wave"], c["red"]
     a = c["args"]
-    mine = tid < s * 64
-    t = fx.min(tid // 64, s - 1)
-    cp = tid % 64
-    col = j * 128 + 2 * cp
-    lo, hi = sum_partials(
-        c["poll"],
-        c["peer"]["latr_own"],
-        lambda src: ((src * s + t) * LAT + col) // 2,
-        TP,
-    )
-    lo, hi = bf16_round(lo), bf16_round(hi)
-    sq = wave_sum(lo * lo + hi * hi)
-    if mine & (lane == 0):
-        c["put"](c["lmsq"], j * s + t, sq)
-    gpu.barrier()
-    if mine & (lane < LN_TASKS):
-        v = c["poll"]([(c["lmsq"], lane * s + t, 1)])[0][0]
-        fx.ptr_store(f32_of(v), red + (t * 32 + lane))
+    nch = (LN_TASKS + WAVES - 1) // WAVES  # a wave's chunks
+    js = [fx.min(wave + WAVES * i, LN_TASKS - 1) for i in range(nch)]
+    lives = [wave + WAVES * i < LN_TASKS for i in range(nch)]
+    cols = [j * 128 + 2 * lane for j in js]
+    rw = rsrc(a["ln_w"])
+    ws = [
+        fx.Float32(fx.BFloat16(bo.buffer_load(rw, col + d, vec_width=1, dtype=T.bf16)))
+        for col in cols
+        for d in range(2)
+    ]
+    # every chunk's TP partials in one poll batch
+    sidx = [sent_ix(src, col, t, s, LAT) for col in cols for src in range(TP)]
+    # a dead chunk is chunk 27's clamped copy: another wave's
+    slive = [lives[i] for i in range(nch) for _ in range(TP)]
+    if const_expr(sent_on(c, "latr")):
+        # (OPT_SENT_LATR) The sentinel words, as [pair] lists;
+        # re-armed after the LN rows are out
+        got = [
+            [w] for w in poll_sent(c["mb"], c["latrs_own"], sidx, slive, rearm=False)
+        ]
+    else:
+        got = c["poll"](
+            [
+                (c["peer"]["latr_own"], ((src * s + t) * LAT + col) // 2, 1)
+                for col in cols
+                for src in range(TP)
+            ],
+            batch=nch * TP,
+        )
+    vals = []
+    for i in range_constexpr(nch):
+        acc0, acc1 = bf2_f32(got[i * TP][0])
+        for src in range_constexpr(1, TP):
+            lo, hi = bf2_f32(got[i * TP + src][0])
+            acc0 = acc0 + lo
+            acc1 = acc1 + hi
+        lo, hi = bf16_round(acc0), bf16_round(acc1)
+        sq = wave_sum(lo * lo + hi * hi)
+        if lives[i] & (lane == 0):
+            fx.ptr_store(sq, red + js[i])
+        vals.append((lo, hi))
     gpu.barrier()
-    tot = fx.ptr_load(red + t * 32)
+    tot = fx.ptr_load(red + 0)
     for q in range_constexpr(1, LN_TASKS):
-        tot = tot + fx.ptr_load(red + (t * 32 + q))
+        tot = tot + fx.ptr_load(red + q)
     r = hw_rsq(tot * fx.Float32(1.0 / LAT) + fx.Float32(c["ln_eps"]))
-    w0 = fx.Float32(
-        fx.BFloat16(bo.buffer_load(rsrc(a["ln_w"]), col, vec_width=1, dtype=T.bf16))
-    )
-    w1 = fx.Float32(
-        fx.BFloat16(bo.buffer_load(rsrc(a["ln_w"]), col + 1, vec_width=1, dtype=T.bf16))
-    )
-    y = (
-        fx.Vector.from_elements([lo * r * w0, hi * r * w1], fx.Float32)
-        .to(fx.BFloat16)
-        .bitcast(fx.Int32)[0]
-    )
-    if mine:
-        bo.buffer_store(y, rsrc(c["ln"]), (t * LAT + col) // 2, cache_modifier=CM_DEV)
-    publish(c["put"], c["lnrdy"], j, fx.Int32(1), tid == 0)
+    for i in range_constexpr(nch):
+        lo, hi = vals[i]
+        y = (
+            fx.Vector.from_elements(
+                [lo * r * ws[2 * i], hi * r * ws[2 * i + 1]], fx.Float32
+            )
+            .to(fx.BFloat16)
+            .bitcast(fx.Int32)[0]
+        )
+        if lives[i]:
+            bo.buffer_store(
+                y, rsrc(c["ln"]), (t * LAT + cols[i]) // 2, cache_modifier=CM_DEV
+            )
+    publish(c["put"], c["lnrdy"], t, fx.Int32(1), tid == 0)
+    if const_expr(sent_on(c, "latr")):
+        rearm_sent(c["latrs_own"], sidx, slive)
     gpu.barrier()
 
 
 @traced
-def stage_up(c, j):
+def run_lnorm(c, bid):
+    """The S latent rows on lnorm CTA bid (LN0 .. LN0 + LN_TASKS), a row a
+    round."""
+    s = c["S"]
+    for rr in range_constexpr((s + LN_TASKS - 1) // LN_TASKS):
+        t = bid - LN0 + LN_TASKS * rr
+        if t < s:
+            stage_lnorm(c, t)
+
+
+def up_splits(s):
+    """Above 16 rows the up tasks split by 16-row token tile (56 CTAs a tile):
+    the LN GEMV is MFMA-bound on 56 CTAs."""
+    n = gemv.tile_rows(LAT)
+    return s // n if s > n and s % n == 0 else 1
+
+
+@traced
+def run_up(c, bid):
+    s = c["S"]
+    if const_expr(up_splits(s) == 1):
+        stage_up(c, bid)
+    else:
+        stage_up(c, bid % UP_TASKS, bid // UP_TASKS)
+
+
+@traced
+def stage_up(c, j, tile=None):
     """This rank's up rows 16 j .. of its shard (hidden rows 896 rank ..): LN GEMV
     + SHP -> bf16 -> FINAL."""
-    s, tid, lane, wave, red = c["S"], c["tid"], c["lane"], c["wave"], c["red"]
+    s, tid, lane, wave = c["S"], c["tid"], c["lane"], c["wave"]
     a = c["args"]
-    row0 = c["rank"] * UP_N + j * ROWS  # the hidden column; w_up is the shard
     ops = gemv.loads(lane, wave, a["w_up"], j * ROWS, LAT)
-    if tid < LN_TASKS:
-        c["poll"]([(c["lnrdy"], tid, 1)])
-    if tid == 0:
-        c["poll"]([(c["shrdy"], j, 1)])
+    if const_expr(batch_on(c)):
+        # (OPT_BATCH) thread 0's two flags in one poll batch
+        if tid == 0:
+            c["poll"]([(c["lnrdy"], 0, 1), (c["shrdy"], j, 1)])
+        if (tid > 0) & (tid < s):
+            c["poll"]([(c["lnrdy"], tid, 1)])
+    else:
+        if tid < s:
+            c["poll"]([(c["lnrdy"], tid, 1)])
+        if tid == 0:
+            c["poll"]([(c["shrdy"], j, 1)])
     gpu.barrier()
+    if const_expr(tile is None):
+        for t0, n in gemv.token_tiles(s, gemv.tile_rows(LAT)):
+            up_tile(c, j, ops, t0, n)
+    else:
+        n = gemv.tile_rows(LAT)
+        up_tile(c, j, ops, tile * n, n)
+
+
+@traced
+def up_tile(c, j, ops, t0, n):
+    """``stage_up`` for the LN rows t0 .. t0 + n."""
+    s, tid, lane, wave, red = c["S"], c["tid"], c["lane"], c["wave"], c["red"]
+    row0 = c["rank"] * UP_N + j * ROWS  # the hidden column; w_up is the shard
     lrow = LAT + gemv.LDS_PAD
-    gemv.dev_rows_to_lds(tid, c["ln"], LAT, s, c["xl"], lrow)
-    acc = gemv.mfmas(lane, wave, c["xl"], lrow, LAT, s, ops)
+    gemv.dev_rows_to_lds(tid, c["ln"], LAT, n, c["xl"], lrow, t0=t0)
+    acc = gemv.mfmas(lane, wave, c["xl"], lrow, LAT, n, ops)
     fx.ptr_store(acc, red + (wave * 64 + lane) * 4)
     gpu.barrier()
-    if tid < (ROWS // 2) * s:
+    if tid < (ROWS // 2) * n:
         rp = tid % (ROWS // 2)
-        tt = tid // (ROWS // 2)
+        tl = tid // (ROWS // 2)
+        tt = t0 + tl
         w = fx.Int32(
             bo.buffer_load(
                 rsrc(c["shp"]),
@@ -988,14 +1571,58 @@ def stage_up(c, j):
                 cache_modifier=CM_DEV,
             )
         )
-        v0 = bf16_round(bf_lo(w) + row_sum(red, 2 * rp, tt))
-        v1 = bf16_round(bf_hi(w) + row_sum(red, 2 * rp + 1, tt))
-        for p in range_constexpr(TP):
-            c["put_bf"](
-                c["peer"]["final"][p],
-                (c["rank"] * s + tt) * HIDDEN + row0 + 2 * rp,
-                [v0, v1],
+        v0 = bf16_round(bf_lo(w) + row_sum(red, 2 * rp, tl))
+        v1 = bf16_round(bf_hi(w) + row_sum(red, 2 * rp + 1, tl))
+        at = (c["rank"] * s + tt) * HIDDEN + row0 + 2 * rp
+        if const_expr(sent_on(c, "final")):
+            # (OPT_SENT_FINAL) The sentinel region; this rank's own word
+            # through the region too (rank order, as the tagged pairs)
+            col = row0 + 2 * rp
+            si = sent_ix(c["rank"], col, tt, s, HIDDEN)
+            word = sent_word(bf16_pair(v0, v1))
+            if const_expr(two_shot(s)):
+                put_sent(c["finals_own"], si, word)
+                lo, hi = sent_sum(c, c["finals_own"], col, tt, HIDDEN, rearm=False)
+                y = fx.Vector.from_elements([lo, hi], fx.Float32).to(fx.BFloat16)
+                bo.buffer_store(
+                    y.bitcast(fx.Int32)[0],
+                    rsrc(c["args"]["out"]),
+                    (tt * HIDDEN + col) // 2,
+                )
+                word = sent_word(bf16_pair(lo, hi))
+                for p in range_constexpr(TP):
+                    if p != c["rank"]:
+                        put_sent(c["finals_peer"][p], si, word)
+                # after this thread's pushes
+                rearm_sent(
+                    c["finals_own"],
+                    [sent_ix(src, col, tt, s, HIDDEN) for src in range(TP)],
+                )
+            else:
+                for p in range_constexpr(TP):
+                    put_sent(c["finals_peer"][p], si, word)
+        elif const_expr(two_shot(s)):
+            # this rank owns the columns: its slot, the ranks' sum in rank order
+            # (as ``stage_reduce``) -> out, and to the other ranks' slot [rank]
+            c["put_bf"](c["peer"]["final_own"], at, [v0, v1])
+            lo, hi = sum_partials(
+                c["poll"],
+                c["peer"]["final_own"],
+                lambda src: ((src * s + tt) * HIDDEN + row0 + 2 * rp) // 2,
+                TP,
+            )
+            y = fx.Vector.from_elements([lo, hi], fx.Float32).to(fx.BFloat16)
+            bo.buffer_store(
+                y.bitcast(fx.Int32)[0],
+                rsrc(c["args"]["out"]),
+                (tt * HIDDEN + row0 + 2 * rp) // 2,
             )
+            for p in range_constexpr(TP):
+                if p != c["rank"]:
+                    c["put_bf"](c["peer"]["final"][p], at, [lo, hi])
+        else:
+            for p in range_constexpr(TP):
+                c["put_bf"](c["peer"]["final"][p], at, [v0, v1])
     gpu.barrier()
 
 
@@ -1003,22 +1630,67 @@ def stage_up(c, j):
 def stage_reduce(c, j):
     """Columns 128 j ..: FINAL in rank order -> bf16 out."""
     s, tid = c["S"], c["tid"]
-    if tid < s * 64:
-        t = tid // 64
-        cp = tid % 64
-        col = j * 128 + 2 * cp
-        lo, hi = sum_partials(
-            c["poll"],
-            c["peer"]["final_own"],
-            lambda src: ((src * s + t) * HIDDEN + col) // 2,
-            TP,
-        )
-        w = (
-            fx.Vector.from_elements([lo, hi], fx.Float32)
-            .to(fx.BFloat16)
-            .bitcast(fx.Int32)[0]
-        )
-        bo.buffer_store(w, rsrc(c["args"]["out"]), (t * HIDDEN + col) // 2)
+    for i in range_constexpr((s * 64 + THREADS - 1) // THREADS):
+        idx = tid + THREADS * i
+        # two-shot: the owner's up wrote its columns
+        live = idx < s * 64
+        if const_expr(two_shot(s)):
+            live = live & (j * 128 // UP_N != c["rank"])
+        if live:
+            if const_expr(sent_on(c, "final")):
+                t, cp = sent_item(idx, s)
+            else:
+                t = idx // 64
+                cp = idx % 64
+            col = j * 128 + 2 * cp
+            if const_expr(sent_on(c, "final") and two_shot(s)):
+                # (OPT_SENT_FINAL) The owner's sum in its sentinel
+                # slot; every round's words re-armed after the last round
+                si = sent_ix(j * 128 // UP_N, col, t, s, HIDDEN)
+                lo, hi = bf2_f32(
+                    poll_sent(c["mb"], c["finals_own"], [si], rearm=False)[0]
+                )
+            elif const_expr(sent_on(c, "final")):
+                lo, hi = sent_sum(c, c["finals_own"], col, t, HIDDEN, rearm=False)
+            elif const_expr(two_shot(s)):
+                # the owner's up reduced the columns: its slot [owner]
+                at = ((j * 128 // UP_N) * s + t) * HIDDEN + col
+                lo, hi = bf2_f32(
+                    c["poll"]([(c["peer"]["final_own"], at // 2, 1)])[0][0]
+                )
+            else:
+                lo, hi = sum_partials(
+                    c["poll"],
+                    c["peer"]["final_own"],
+                    lambda src, t=t, col=col: ((src * s + t) * HIDDEN + col) // 2,
+                    TP,
+                )
+            w = (
+                fx.Vector.from_elements([lo, hi], fx.Float32)
+                .to(fx.BFloat16)
+                .bitcast(fx.Int32)[0]
+            )
+            bo.buffer_store(w, rsrc(c["args"]["out"]), (t * HIDDEN + col) // 2)
+    if const_expr(sent_on(c, "final")):
+        # (OPT_SENT_FINAL) The words every round read, re-armed
+        # after the last round's poll (not between the rounds' polls)
+        for i in range_constexpr((s * 64 + THREADS - 1) // THREADS):
+            idx = tid + THREADS * i
+            live = idx < s * 64
+            if const_expr(two_shot(s)):
+                live = live & (j * 128 // UP_N != c["rank"])
+            if live:
+                t, cp = sent_item(idx, s)
+                col = j * 128 + 2 * cp
+                if const_expr(two_shot(s)):
+                    rearm_sent(
+                        c["finals_own"], [sent_ix(j * 128 // UP_N, col, t, s, HIDDEN)]
+                    )
+                else:
+                    rearm_sent(
+                        c["finals_own"],
+                        [sent_ix(src, col, t, s, HIDDEN) for src in range(TP)],
+                    )
     gpu.barrier()
 
 
diff --git a/vllm/models/kimi_k3/amd/specialized_decode/runtime/__init__.py b/vllm/models/kimi_k3/amd/specialized_decode/runtime/__init__.py
index f1a3509db513..d1dc61f6ea9f 100644
--- a/vllm/models/kimi_k3/amd/specialized_decode/runtime/__init__.py
+++ b/vllm/models/kimi_k3/amd/specialized_decode/runtime/__init__.py
@@ -19,8 +19,17 @@
 
 AUTO, SPECIALIZED = "auto", "latency_specialized"
 TP = 8
-# a launch's x rows fit one CTA's LDS
-MAX_TOKENS = 8
+DCP = TP  # the MLA layers' KV shards: one a TP rank
+DCP_SIZES = (1, DCP)  # M1's builds: the whole KV a rank, or DCP over the TP group
+# the kernels' rows: the MoE route table's token bitmaps (scratch and peer memory
+# are sized for it)
+MAX_TOKENS = 64
+# a served step's rows: at 32 and 64 rows the kernels lose to vLLM's decode
+MAX_ROWS = 16
+# the served max_num_seqs: a full step of each is one of the widths (MAX_ROWS)
+MAX_SEQS = (1, 2, 4, 8)
+# M1c's qlens: DCP1 runs only them (AITER's asm decode at DCP1 is not served)
+M1C_QLENS = (4, 8)
 WHOLE, FFN = "whole", "ffn"
 # Kimi-K3's text shapes the kernels are written for (hidden size, routed experts,
 # KDA heads, KDA head dim, KDA short-conv width)
@@ -31,14 +40,18 @@
 class DecodeConfig:
     # rows a request verifies: num_speculative_tokens + 1
     q: int
+    # the MLA layers run M1 (else vLLM's attention ops, then K2; ``m1_reasons``)
+    m1: bool = True
+    # the DCP KV cache interleave (AITER's decode path, M1a's query pad)
+    interleave: int = 1
 
     @property
     def widths(self) -> tuple[int, ...]:
-        """The rows a launch serves: whole requests, at most MAX_TOKENS rows. Every
-        uniform verify step up to MAX_TOKENS rows is one of them, so an eager step
+        """The rows a launch serves: whole requests, at most MAX_ROWS rows. Every
+        uniform verify step up to MAX_ROWS rows is one of them, so an eager step
         never pads; a FULL replay pads whole requests the way vLLM's stock decode
         does."""
-        return tuple(range(self.q, MAX_TOKENS + 1, self.q))
+        return tuple(range(self.q, MAX_ROWS + 1, self.q))
 
 
 def decode_backend(vllm_config) -> str:
@@ -70,14 +83,44 @@ def specialized_decode_config(vllm_config) -> DecodeConfig | None:
             "Kimi-K3 specialized decode: PIECEWISE-only CUDA graphs capture every "
             "decode step as stock; use FULL decode graphs"
         )
+    m1 = m1_reasons(vllm_config)
+    if m1:
+        logger.info(
+            "Kimi-K3 specialized decode: the MLA layers run vLLM's attention ops, "
+            "then K2 (M1 needs %s)",
+            "; ".join(m1),
+        )
     spec = vllm_config.speculative_config
-    return DecodeConfig(q=spec.num_speculative_tokens + 1)
+    return DecodeConfig(
+        q=spec.num_speculative_tokens + 1,
+        m1=not m1,
+        interleave=vllm_config.parallel_config.cp_kv_cache_interleave_size,
+    )
+
+
+def m1_reasons(vllm_config) -> list[str]:
+    """Why the MLA layers cannot run M1 (empty: they can). At DCP8 not a refusal:
+    those layers run vLLM's attention ops instead, then K2 (``step_reasons``)."""
+    import vllm.envs as envs
+
+    cc = vllm_config.cache_config
+    checks = (
+        # M1a writes the MLA layers' fp8 KV and hands AITER an fp8 query
+        (not cc.cache_dtype.startswith("fp8"), "an fp8 KV cache"),
+        # M1 runs W_UK / W_UV as AITER's MXFP4 bmms
+        (
+            not envs.VLLM_ROCM_USE_AITER_FP4BMM,
+            "AITER's MXFP4 bmms (VLLM_ROCM_USE_AITER_FP4BMM=1)",
+        ),
+    )
+    return [why for bad, why in checks if bad]
 
 
 def unsupported(vllm_config) -> list[str]:
     """Every reason this deployment cannot run the specialized decode (empty: it
     can). Config only: the loaded weights are checked when the executor is built.
-    Default-deny: every feature below is refused, the kernels serve none of them."""
+    Default-deny: every feature below is refused, the kernels serve none of them.
+    What only M1 needs is not refused at DCP8 (``m1_reasons``)."""
     from vllm.config import CompilationMode
     from vllm.model_executor.layers.mamba.mamba_utils import (
         MambaStateDtypeCalculator,
@@ -97,8 +140,8 @@ def unsupported(vllm_config) -> list[str]:
         (pc.enable_expert_parallel, "expert parallelism"),
         (pc.enable_eplb, "EPLB"),
         (
-            pc.decode_context_parallel_size not in (1, pc.tensor_parallel_size),
-            "decode context parallelism other than 1 or TP",
+            pc.decode_context_parallel_size not in DCP_SIZES,
+            f"decode context parallelism other than 1 or {DCP} (over the TP group)",
         ),
         (pc.prefill_context_parallel_size != 1, "prefill context parallelism"),
         # the kernels read one batch's attention metadata per step
@@ -128,8 +171,8 @@ def unsupported(vllm_config) -> list[str]:
         # the kernels serve spec-verify steps of whole uniform requests only
         (spec is None, "decoding without speculative decoding"),
         (
-            spec is not None and spec.num_speculative_tokens + 1 > MAX_TOKENS,
-            f"more than {MAX_TOKENS - 1} speculative tokens",
+            spec is not None and spec.num_speculative_tokens + 1 > MAX_ROWS,
+            f"more than {MAX_ROWS - 1} speculative tokens",
         ),
         (
             spec is not None and spec.enable_adaptive_verification,
@@ -169,7 +212,46 @@ def unsupported(vllm_config) -> list[str]:
         )
     if getattr(mc.hf_text_config, "attn_res_block_size", None) is None:
         reasons.append("models without AttnRes")
-    return reasons + platform_reasons()
+    platform = platform_reasons()
+    if spec is not None and spec.num_speculative_tokens + 1 <= MAX_ROWS:
+        reasons += step_reasons(vllm_config, check_m1c=not platform)
+    return reasons + platform
+
+
+def m1c_serves(q: int, seqs: int) -> bool:
+    """M1c runs every step of 1 .. ``seqs`` requests at qlen ``q``: its KV splits
+    share out evenly over the requests (``layers.m1c_at``)."""
+    if q not in M1C_QLENS:
+        return False
+    from vllm.models.kimi_k3.amd.specialized_decode.mono.attention import (
+        mla_attn_geom,
+    )
+
+    return all(mla_attn_geom.splits(q) % n == 0 for n in range(1, seqs + 1))
+
+
+def step_reasons(vllm_config, check_m1c: bool = True) -> list[str]:
+    """The (max_num_seqs, speculative tokens) pairs served: a full step fits one
+    launch, and at DCP1 every step runs M1c (``check_m1c``: M1c's geometry needs
+    the kernels' modules, so only on a supported platform)."""
+    k = vllm_config.speculative_config.num_speculative_tokens
+    seqs = vllm_config.scheduler_config.max_num_seqs
+    if seqs not in MAX_SEQS or seqs * (k + 1) > MAX_ROWS:
+        return [
+            f"max_num_seqs {seqs} at {k} speculative tokens (max_num_seqs 1, 2, 4 "
+            f"or 8 with max_num_seqs x (k + 1) <= {MAX_ROWS})"
+        ]
+    if vllm_config.parallel_config.decode_context_parallel_size != 1:
+        return []
+    m1 = m1_reasons(vllm_config)
+    if m1:
+        return ["DCP1 without M1 (" + "; ".join(m1) + ")"]
+    if check_m1c and not m1c_serves(k + 1, seqs):
+        return [
+            f"DCP1 at {k} speculative tokens, max_num_seqs {seqs} (DCP1 runs M1c "
+            "only: 3 or 7 speculative tokens, max_num_seqs 1 or 2; use DCP8)"
+        ]
+    return []
 
 
 def text_shapes(text) -> tuple:
diff --git a/vllm/models/kimi_k3/amd/specialized_decode/runtime/decode.py b/vllm/models/kimi_k3/amd/specialized_decode/runtime/decode.py
index fa62155fd2e5..2844bddde73f 100644
--- a/vllm/models/kimi_k3/amd/specialized_decode/runtime/decode.py
+++ b/vllm/models/kimi_k3/amd/specialized_decode/runtime/decode.py
@@ -17,11 +17,13 @@
 import torch
 
 from vllm.config import CUDAGraphMode
-from vllm.distributed import get_tp_group
+from vllm.distributed import get_dcp_group, get_tp_group
 from vllm.forward_context import get_forward_context, is_forward_context_available
 from vllm.logger import init_logger
 from vllm.models.kimi_k3.amd.specialized_decode.mono.abi import FAILED
 from vllm.models.kimi_k3.amd.specialized_decode.runtime import (
+    DCP,
+    FFN,
     WHOLE,
     DecodeConfig,
     layer_kinds,
@@ -29,19 +31,22 @@
 from vllm.models.kimi_k3.amd.specialized_decode.runtime.guards import (
     FailStop,
     cache_reasons,
+    mla_weight_reasons,
     weight_reasons,
 )
 from vllm.models.kimi_k3.amd.specialized_decode.runtime.layers import (
     FfnLayer,
+    MlaOpsLayer,
     WholeLayer,
     compile_build,
+    m1c_at,
 )
 from vllm.models.kimi_k3.amd.specialized_decode.runtime.peer import tp_agree
 from vllm.models.kimi_k3.amd.specialized_decode.runtime.runner import MonoRunner
 
 logger = init_logger(__name__)
 
-Kernel = WholeLayer | FfnLayer
+Kernel = WholeLayer | FfnLayer | MlaOpsLayer
 
 
 class Step(NamedTuple):
@@ -57,10 +62,12 @@ def layer(self, idx: int) -> Callable | None:
         return None if kernel is None else partial(kernel.launch, self)
 
 
-def step_reason(attn_metadata, kda_prefix: str, width: int) -> str | None:
-    """None if a step's KDA metadata (identical on every TP rank) is a pure
-    spec-verify step of whole uniform requests filling ``width`` rows, else why
-    not."""
+def step_reason(
+    attn_metadata, kda_prefix: str, width: int, mla_prefix: str = ""
+) -> str | None:
+    """None if a step's KDA and MLA metadata (identical on every TP rank) are a
+    pure spec-verify step of whole uniform requests filling ``width`` rows, else
+    why not."""
     if not isinstance(attn_metadata, dict):
         return "no metadata"
     m: Any = attn_metadata.get(kda_prefix)
@@ -71,6 +78,19 @@ def step_reason(attn_metadata, kda_prefix: str, width: int) -> str | None:
     q = m.spec_state_indices_tensor.size(-1)
     if m.uniform_spec_sequence_length != q or m.num_spec_decodes * q != width:
         return "not uniform"
+    if mla_prefix:
+        a: Any = attn_metadata.get(mla_prefix)
+        if a is None or a.decode is None:
+            return "no MLA decode metadata"
+        if a.num_prefills or a.num_decode_tokens != width:
+            return "MLA: not pure decode"
+        if a.decode.max_qo_len != q:
+            return "MLA: not uniform"
+        # Gluon serves bf16 caches only; M1a hands AITER an fp8 query
+        if getattr(a.decode, "use_gluon_decode", False) or getattr(
+            a.decode, "use_gluon_verify", False
+        ):
+            return "MLA: Gluon decode"
     return None
 
 
@@ -90,22 +110,47 @@ class SpecializedDecode:
     plans every step."""
 
     def __init__(self, backbone, config: DecodeConfig):
-        tp = get_tp_group()
-        agree(weight_reasons(backbone), tp.cpu_group, "the loaded weights")
+        tp, dcp = get_tp_group(), get_dcp_group()
+        reasons = weight_reasons(backbone)
+        if dcp.world_size > 1 and (
+            dcp.world_size != DCP or dcp.rank_in_group != tp.rank_in_group
+        ):
+            reasons.append(
+                f"DCP{dcp.world_size} rank {dcp.rank_in_group} on TP rank "
+                f"{tp.rank_in_group} (the DCP group must be the TP group)"
+            )
+        agree(reasons, tp.cpu_group, "the loaded weights")
         self.backbone, self.config = backbone, config
         kinds = {
             i: kind
             for i, kind in layer_kinds(backbone.config).items()
             if backbone.start_layer <= i < backbone.end_layer
         }
-        runner = self.runner = MonoRunner(next(backbone.parameters()).device)
+        ffn = [i for i, kind in kinds.items() if kind == FFN]
+        m1 = config.m1 and bool(ffn)
+        if m1:
+            why = mla_weight_reasons(backbone)
+            m1 = tp_agree(not why, tp.cpu_group)
+            if not m1:
+                logger.info(
+                    "Kimi-K3 specialized decode: the MLA layers run vLLM's "
+                    "attention ops, then K2 (M1 does not support: %s)",
+                    "; ".join(why) if why else "another TP rank's MLA layers",
+                )
+        runner = self.runner = MonoRunner(next(backbone.parameters()).device, m1=m1)
         runner.num_spec = config.q - 1
+        runner.interleave = config.interleave
         runner.first_layer = min(kinds)
         whole = [i for i, kind in kinds.items() if kind == WHOLE]
         runner.kda_prefix = backbone.layers[whole[0]].self_attn.prefix
+        if m1:
+            # the M1 layers' metadata: their MLAAttention's
+            wrapper = backbone.layers[ffn[0]].self_attn.mla_attn
+            runner.mla_prefix = wrapper.mla_attn.layer_name
+        mla_kernel = FfnLayer if m1 else MlaOpsLayer
         self.kernels: dict[int, dict[int, Kernel]] = {
             w: {
-                i: (WholeLayer if kind == WHOLE else FfnLayer)(
+                i: (WholeLayer if kind == WHOLE else mla_kernel)(
                     backbone.layers[i], runner, w
                 )
                 for i, kind in kinds.items()
@@ -116,6 +161,10 @@ def __init__(self, backbone, config: DecodeConfig):
             k for k in self.kernels[config.widths[0]].values()
             if isinstance(k, WholeLayer)
         ]  # fmt: skip
+        self.mla: list[FfnLayer] = [
+            k for k in self.kernels[config.widths[0]].values()
+            if isinstance(k, FfnLayer)
+        ]  # fmt: skip
         builds = {
             b for by in self.kernels.values() for k in by.values() for b in k.builds()
         }
@@ -126,10 +175,13 @@ def __init__(self, backbone, config: DecodeConfig):
         )
         # (data_ptr, numel) of the first KDA layer's state prepare_kv bound
         self._kv: tuple[int, int] | None = None
+        m1c = [w for w in config.widths if m1c_at(runner, w)]
+        mla = f"M1 (M1c in K2 at widths {m1c})" if m1c else "M1"
         logger.info(
-            "Kimi-K3 specialized decode: %d layers, widths %s, %d builds; other "
-            "steps run vLLM's decode",
+            "Kimi-K3 specialized decode: %d layers (MLA: %s), widths %s, %d "
+            "builds; other steps run vLLM's decode",
             len(kinds),
+            mla if m1 else "vLLM's ops",
             list(config.widths),
             len(builds),
         )
@@ -143,12 +195,13 @@ def _kv_key(self) -> tuple[int, int] | None:
         return state[1].data_ptr(), state[1].numel()
 
     def prepare_kv(self) -> None:
-        """Checks and binds the KDA states vLLM just bound (the CUDA-graph memory
-        profiling cache, then the real one)."""
+        """Checks and binds the KDA states and MLA caches vLLM just bound (the
+        CUDA-graph memory profiling cache, then the real one)."""
         self._kv = None
         if self._kv_key() is None:
             return
         caches = {k.layer.self_attn.prefix: k.caches() for k in self.kda}
+        caches.update({k.wrapper.mla_attn.layer_name: k.caches() for k in self.mla})
         agree(cache_reasons(caches), get_tp_group().cpu_group, "the KV caches")
         self._kv = self._kv_key()
 
@@ -159,7 +212,7 @@ def plan(self, positions: torch.Tensor, intermediate_tensors) -> Step | None:
         - pipeline-parallel stage inputs, and steps before a KV cache is bound
           (profiling, dummy runs);
         - PIECEWISE steps (captured without attention metadata);
-        - widths without kernels (more than MAX_TOKENS rows);
+        - widths without kernels (more than MAX_ROWS rows);
         - steps that are not pure spec-verify steps of whole uniform requests
           (``step_reason``).
 
@@ -181,7 +234,8 @@ def plan(self, positions: torch.Tensor, intermediate_tensors) -> Step | None:
         if kernels is None:
             return None
         md = context.attn_metadata
-        if step_reason(md, self.runner.kda_prefix, width) is not None:
+        runner = self.runner
+        if step_reason(md, runner.kda_prefix, width, runner.mla_prefix) is not None:
             return None
         assert isinstance(md, dict)
         if self._kv_key() != self._kv:
diff --git a/vllm/models/kimi_k3/amd/specialized_decode/runtime/guards.py b/vllm/models/kimi_k3/amd/specialized_decode/runtime/guards.py
index 9d0808d721d2..be14fa91e70d 100644
--- a/vllm/models/kimi_k3/amd/specialized_decode/runtime/guards.py
+++ b/vllm/models/kimi_k3/amd/specialized_decode/runtime/guards.py
@@ -9,18 +9,19 @@
 
 import torch
 
-from vllm.models.kimi_k3.amd.specialized_decode.runtime import WHOLE, layer_kinds
+from vllm.models.kimi_k3.amd.specialized_decode.runtime import (
+    DCP,
+    DCP_SIZES,
+    FFN,
+    layer_kinds,
+)
 
+MLA_HEADS, MLA_Q, MLA_KV, MLA_NOPE, MLA_ROPE, MLA_V = 12, 1536, 512, 128, 64, 128
 
-def weight_reasons(model) -> list[str]:
-    """``model``: the loaded ``KimiLinearModel``; every kernel layer's weights as
-    the kernels read them, and the MoE backend that lays them out."""
-    from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import Mxfp4MoeBackend
-    from vllm.models.kimi_k3.amd.specialized_decode.mono.attention import back, kda
-    from vllm.models.kimi_k3.amd.specialized_decode.mono.stages import moe
 
-    reasons: list[str] = []
-    bf, f32 = torch.bfloat16, torch.float32
+def _need(reasons: list[str]):
+    """``need(name, t, shape, dtype)``: appends to ``reasons`` unless ``t`` is
+    a contiguous ``shape`` ``dtype`` tensor."""
 
     def need(name, t, shape, dtype):
         if t is None:
@@ -32,6 +33,20 @@ def need(name, t, shape, dtype):
                 f" kernel reads contiguous {shape} {dtype}"
             )
 
+    return need
+
+
+def weight_reasons(model) -> list[str]:
+    """``model``: the loaded ``KimiLinearModel``; every kernel layer's weights as
+    K1 and K2 read them, and the MoE backend that lays them out. The MLA layers'
+    attention is ``mla_weight_reasons``'."""
+    from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import Mxfp4MoeBackend
+    from vllm.models.kimi_k3.amd.specialized_decode.mono.attention import back, kda
+    from vllm.models.kimi_k3.amd.specialized_decode.mono.stages import moe
+
+    reasons: list[str] = []
+    bf, f32 = torch.bfloat16, torch.float32
+    need = _need(reasons)
     kinds = layer_kinds(model.config)
     if not kinds:
         return ["no layer has an AttnRes latent MoE"]
@@ -74,7 +89,12 @@ def need(name, t, shape, dtype):
         need(f"{p} post_attention_layernorm", layer.post_attention_layernorm.weight,
              (H,), bf)  # fmt: skip
         need(f"{p} o_proj", layer.self_attn.o_proj.weight, (H, back.OK), bf)
-        if kind != WHOLE:
+        need(f"{p} input_layernorm", layer.input_layernorm.weight, (H,), bf)
+        need(f"{p} self_attention_res_norm", layer.self_attention_res_norm.weight,
+             (H,), bf)  # fmt: skip
+        need(f"{p} self_attention_res_proj", layer.self_attention_res_proj.weight,
+             (1, H), bf)  # fmt: skip
+        if kind == FFN:
             continue
         a = layer.self_attn
         if a.local_num_heads != kda.NH or a.head_dim != kda.HD or a.conv_size != 4:
@@ -86,20 +106,99 @@ def need(name, t, shape, dtype):
         need(f"{p} conv1d", a.conv1d.weight.view(-1, 4), (kda.QKV, 4), f32)
         need(f"{p} A_log", a.A_log.view(-1), (kda.NH,), f32)
         need(f"{p} dt_bias", a.dt_bias.view(-1), (kda.PROJ,), f32)
-        need(f"{p} input_layernorm", layer.input_layernorm.weight, (H,), bf)
-        need(f"{p} self_attention_res_norm", layer.self_attention_res_norm.weight,
-             (H,), bf)  # fmt: skip
-        need(f"{p} self_attention_res_proj", layer.self_attention_res_proj.weight,
-             (1, H), bf)  # fmt: skip
+    return reasons
+
+
+def mla_weight_reasons(model) -> list[str]:
+    """Why the loaded MLA layers cannot run M1 (empty: they can). Not a refusal:
+    those layers run vLLM's attention ops instead, then K2."""
+    reasons: list[str] = []
+    for i, kind in layer_kinds(model.config).items():
+        if kind == FFN and model.start_layer <= i < model.end_layer:
+            w = model.layers[i].self_attn.mla_attn
+            reasons += _mla_reasons(f"layer {i}", w, _need(reasons))
+    return reasons
+
+
+def _mla_reasons(p: str, w, need) -> list[str]:
+    """An MLA layer as M1a / M1b (and ``runtime/mla.py``) run it: NoPE, gated,
+    DCP1 or DCP8, an fp8 cache, W_UK / W_UV as MXFP4 bmms."""
+    from vllm.models.kimi_k3.amd.specialized_decode.mono.attention.kda import (
+        HIDDEN as H,
+    )
+    from vllm.models.kimi_k3.amd.specialized_decode.runtime.mla import (
+        qkv_gate_weights,
+    )
+
+    reasons = []
+    a = w.mla_attn
+    dims = (w.num_heads, w.q_lora_rank, a.kv_lora_rank, w.qk_nope_head_dim,
+            w.qk_rope_head_dim, w.v_head_dim)  # fmt: skip
+    if dims != (MLA_HEADS, MLA_Q, MLA_KV, MLA_NOPE, MLA_ROPE, MLA_V):
+        return [f"{p}: MLA heads/q_lora/kv_lora/nope/rope/v {dims}"]
+    if w.rotary_emb is not None or w.indexer is not None or w.dcp_q_replicate:
+        reasons.append(f"{p}: needs NoPE, dense, unreplicated-query MLA")
+    w_qkv, w_g = qkv_gate_weights(w)
+    if w_g is None:
+        reasons.append(f"{p}: needs the output gate (g_proj)")
+    dcp = a.impl.dcp_world_size
+    if dcp not in DCP_SIZES or a.use_pcp or a.hisparse_cache is not None:
+        reasons.append(f"{p}: needs DCP1 or DCP{DCP}, no PCP / HiSparse")
+    elif dcp > 1 and (a.dcp_manager is None or a.dcp_manager.query_gather is None):
+        reasons.append(f"{p}: no DCP query gather")
+    if not a.kv_cache_dtype.startswith("fp8"):
+        reasons.append(f"{p}: KV cache {a.kv_cache_dtype}, needs fp8")
+    if not a.is_aiter_triton_fp4_bmm_enabled:
+        reasons.append(
+            f"{p}: W_UK / W_UV must be AITER MXFP4 bmms (VLLM_ROCM_USE_AITER_FP4BMM=1)"
+        )
+    if float(a._q_scale_float) != 1.0 or float(a._k_scale_float) != 1.0:
+        reasons.append(f"{p}: q / k scales must be 1.0")
+    bf = torch.bfloat16
+    need(f"{p} fused_qkv_a_proj", w_qkv,
+         (MLA_Q + MLA_KV + MLA_ROPE, H), bf)  # fmt: skip
+    need(f"{p} q_a_layernorm", w.q_a_layernorm.weight, (MLA_Q,), bf)
+    need(f"{p} kv_a_layernorm", w.kv_a_layernorm.weight, (MLA_KV,), bf)
+    need(f"{p} q_b_proj", w.q_b_proj.weight,
+         (MLA_HEADS * (MLA_NOPE + MLA_ROPE), MLA_Q), bf)  # fmt: skip
+    if w_g is not None:
+        need(f"{p} g_proj", w_g, (MLA_HEADS * MLA_V, H), bf)
+    if a.is_aiter_triton_fp4_bmm_enabled:
+        # the bmm weights are transposed views (rows strided, K dense); the
+        # e8m0 scales are aiter's column-major blocks, any strides
+        for name, t, shape, dense in (
+            ("W_K", a.W_K, (MLA_HEADS, MLA_KV, MLA_NOPE // 2), True),
+            ("W_K_scale", a.W_K_scale, (MLA_HEADS, MLA_KV, MLA_NOPE // 32), False),
+            ("W_V", a.W_V, (MLA_HEADS, MLA_V, MLA_KV // 2), True),
+            ("W_V_scale", a.W_V_scale, (MLA_HEADS, MLA_V, MLA_KV // 32), False),
+        ):
+            if t is None or tuple(t.shape) != shape:
+                reasons.append(f"{p} {name}: needs {shape}")
+            elif dense and t.stride(-1) != 1:
+                reasons.append(f"{p} {name}: needs the last dim dense")
+            elif t.element_size() != 1:
+                reasons.append(f"{p} {name}: {t.dtype}, needs bytes (torch.uint8)")
     return reasons
 
 
 def cache_reasons(caches: Mapping[str, Any]) -> list[str]:
-    """Every KDA layer's (conv, recurrent) state as K1 reads them."""
+    """Every KDA layer's (conv, recurrent) state as K1 reads them; every MLA
+    layer's fp8 cache, rows of 576 bytes (the blocks may be strided)."""
     from vllm.models.kimi_k3.amd.specialized_decode.mono.attention import kda
 
     reasons = []
-    for name, (conv_state, rstate) in caches.items():
+    for name, c in caches.items():
+        if isinstance(c, torch.Tensor):
+            row = MLA_KV + MLA_ROPE
+            if c.ndim != 3 or c.size(-1) != row or c.element_size() != 1:
+                reasons.append(
+                    f"{name}: MLA cache {tuple(c.shape)} {c.dtype}, needs fp8 "
+                    f"[blocks, block_size, {row}]"
+                )
+            elif c.stride(-1) != 1 or c.stride(1) != row:
+                reasons.append(f"{name}: MLA cache rows must be dense")
+            continue
+        conv_state, rstate = c
         if conv_state.dtype != torch.bfloat16:
             reasons.append(f"{name}: conv state {conv_state.dtype}, needs bf16")
         if rstate.dtype != torch.float32:
diff --git a/vllm/models/kimi_k3/amd/specialized_decode/runtime/layers.py b/vllm/models/kimi_k3/amd/specialized_decode/runtime/layers.py
index 843e66dc50cd..f00e95533bd3 100644
--- a/vllm/models/kimi_k3/amd/specialized_decode/runtime/layers.py
+++ b/vllm/models/kimi_k3/amd/specialized_decode/runtime/layers.py
@@ -1,13 +1,17 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""The layer kernels: each takes and returns what ``KimiDecoderLayer.forward``
-does with AttnRes, ``(prefix_sum, block_residual, prefix_delta)``.
+"""The layer kernels: each takes and returns what
+``KimiDecoderLayer.forward`` does with AttnRes, ``(prefix_sum, block_residual,
+prefix_delta)``.
 
 - ``WholeLayer`` (KDA layers): K1 (AttnRes, in_proj, f_b, spec conv, the KDA
   recurrence, the gated norm) writes o_proj's input; K2 runs o_proj, its
   all-reduce, the MLP AttnRes and the latent MoE.
-- ``FfnLayer`` (MLA layers): vLLM's AttnRes and MLA up to o_proj's input, then
-  K2.
+- ``FfnLayer`` (MLA layers): M1a (AttnRes, the query prologue, the KV write
+  and, at DCP8, the query gather), AITER's decode, M1b (at DCP8 the LSE
+  combine, W_UV and the gate), then K2.
+- ``MlaOpsLayer`` (MLA layers M1 cannot serve, ``m1_reasons``): vLLM's AttnRes
+  and MLA attention ops, then K2.
 
 Weights are views of the loaded parameters, taken at every launch.
 """
@@ -23,10 +27,22 @@
     from vllm.models.kimi_k3.amd.specialized_decode.runtime.decode import Step
 
 
-def _k2_key(layer, runner: MonoRunner, tokens: int, reset: bool):
+def _k2_key(layer, runner: MonoRunner, tokens: int, reset: bool, m1b: bool = False):
+    """``m1b``: an MLA layer's K2 that runs its M1b first (``runner.fuse_m1b``)."""
     from vllm.models.kimi_k3.amd.specialized_decode.mono import layer as k2
 
     m = layer.mlp
+    fused = (
+        dict(
+            m1b_qlen=runner.num_spec + 1,
+            m1b_ranks=runner.dcp,
+            m1b_off=runner.mla_post_off,
+            m1b_splits=m1c_splits(runner, tokens),
+            m1c=m1c_at(runner, tokens),
+        )
+        if m1b
+        else {}
+    )
     return k2.K2Build(
         tokens=tokens,
         nblocks=layer.prev_valid_blocks + (1 if reset else 0),
@@ -34,13 +50,46 @@ def _k2_key(layer, runner: MonoRunner, tokens: int, reset: bool):
         eps=layer.mlp_res_norm.variance_epsilon,
         out_eps=layer.post_attention_layernorm.variance_epsilon,
         ln_eps=m.routed_expert_norm.variance_epsilon,
+        **fused,
         **runner.fault,
     )
 
 
-def _k2(layer, runner: MonoRunner, core, prefix, blocks, reset=False):
+def m1c_at(runner: MonoRunner, width: int) -> bool:
+    """``runner.m1c`` at this width: M1c's KV splits must share out evenly over
+    the requests (qlen 4 has 80: a 12-row width, 3 requests, runs the asm
+    decode instead)."""
+    from vllm.models.kimi_k3.amd.specialized_decode.mono.attention import (
+        mla_attn_geom,
+    )
+
+    q = runner.num_spec + 1
+    return runner.m1c and mla_attn_geom.splits(q) % (width // q) == 0
+
+
+def m1c_splits(runner: MonoRunner, width: int) -> int:
+    """M1b's ``splits``: M1c's, when K2 runs it at this width (else 0)."""
+    from vllm.models.kimi_k3.amd.specialized_decode.mono.attention import (
+        mla_attn_geom,
+    )
+
+    return mla_attn_geom.splits(runner.num_spec + 1) if m1c_at(runner, width) else 0
+
+
+def _k2(
+    layer,
+    runner: MonoRunner,
+    core,
+    prefix,
+    blocks,
+    reset=False,
+    m1b_in=None,
+    m1c_in=None,
+):
     """K2 from ``core`` (o_proj's input): ``prefix`` updated in place (written,
-    when ``reset``), the MoE output returned."""
+    when ``reset``), the MoE output returned. ``m1b_in``: ``mla_post``'s inputs,
+    when this launch runs M1b and writes ``core`` itself; ``m1c_in``: M1c's
+    (``mla_attn.host_args``'s), when it runs M1c before."""
     from vllm.models.kimi_k3.amd.specialized_decode.mono import layer as k2
     from vllm.models.kimi_k3.amd.specialized_decode.mono.attention.kda import HIDDEN
     from vllm.models.kimi_k3.amd.specialized_decode.mono.stages import moe
@@ -52,7 +101,7 @@ def _k2(layer, runner: MonoRunner, core, prefix, blocks, reset=False):
         w_up = w_up.narrow(0, runner.rank * moe.UP_N, moe.UP_N)
     out = torch.empty_like(prefix)
     k2.k2_launch(
-        _k2_key(layer, runner, core.size(0), reset),
+        _k2_key(layer, runner, core.size(0), reset, m1b_in is not None),
         core=core,
         w_o=layer.self_attn.o_proj.weight,
         prefix=prefix,
@@ -73,6 +122,8 @@ def _k2(layer, runner: MonoRunner, core, prefix, blocks, reset=False):
         w_up=w_up,
         out=out,
         layer=128 + layer.layer_idx,
+        m1b_in=m1b_in,
+        m1c_in=m1c_in,
         **runner.launch_args(),
         **runner.peer_args(),
     )
@@ -80,14 +131,24 @@ def _k2(layer, runner: MonoRunner, core, prefix, blocks, reset=False):
 
 
 def compile_build(key) -> None:
-    """Compiles one launch's kernel (a ``K2Build`` or ``KdaPreBuild``) without
-    running it: FlyDSL's compile-only mode returns before the launch."""
+    """Compiles one launch's kernel (a ``K2Build``, ``KdaPreBuild``,
+    ``MlaPreBuild`` or ``MlaPostBuild``) without running it: FlyDSL's
+    compile-only mode returns before the launch."""
     import os
 
     from vllm.models.kimi_k3.amd.specialized_decode.mono import layer as k2
-    from vllm.models.kimi_k3.amd.specialized_decode.mono.attention import kda
+    from vllm.models.kimi_k3.amd.specialized_decode.mono.attention import (
+        kda,
+        mla,
+        mla_back,
+    )
 
-    mod = kda if isinstance(key, kda.KdaPreBuild) else k2
+    mod = {
+        kda.KdaPreBuild: kda,
+        k2.K2Build: k2,
+        mla.MlaPreBuild: mla,
+        mla_back.MlaPostBuild: mla_back,
+    }[type(key)]
     launch = mod.build(key)
     # FlyDSL's compile-only switch has no FLYDSL_ prefix
     old = os.environ.get("COMPILE_ONLY")
@@ -198,6 +259,154 @@ def launch(self, step: "Step", positions, prefix_sum, block_residual, prefix_del
 
 
 class FfnLayer(_Layer):
+    def __init__(self, layer, runner: MonoRunner, width: int):
+        super().__init__(layer, runner, width)
+        self.wrapper = layer.self_attn.mla_attn
+
+    def _m1a_key(self, delta: bool, qheads: int | None = None):
+        from vllm.models.kimi_k3.amd.specialized_decode.mono.attention import mla
+
+        layer, w = self.layer, self.wrapper
+        if qheads is None:
+            qheads = self._cprr_heads()
+        return mla.MlaPreBuild(
+            tokens=self.width,
+            nblocks=layer.prev_valid_blocks,
+            delta=delta,
+            eps=layer.self_attention_res_norm.variance_epsilon,
+            out_eps=layer.input_layernorm.variance_epsilon,
+            q_eps=w.q_a_layernorm.variance_epsilon,
+            kv_eps=w.kv_a_layernorm.variance_epsilon,
+            ranks=self.runner.dcp,
+            peer_off=self.runner.mla_peer_off,
+            qheads=qheads,
+            qbf16=m1c_at(self.runner, self.width),
+        )
+
+    def _cprr_heads(self) -> int:
+        """The head count AITER's cprr decode pads the gathered query to, 0 if
+        it doesn't pad. The warm-up's guess; ``launch`` reads the step's."""
+        from vllm.models.kimi_k3.amd.specialized_decode.mono.attention.mla import NH
+        from vllm.v1.attention.backends.mla import rocm_aiter_mla as am
+
+        n = self.runner.dcp * NH
+        asm = am._asm_dcp_verify_configured(
+            self.runner.dcp, self.runner.interleave, True
+        )
+        if not asm or self.runner.num_spec + 1 < am._MIN_CPRR_QLEN:
+            return 0
+        heads = am._asm_dcp_verify_heads(n)
+        return heads if heads > n else 0
+
+    def _m1b_key(self, qlen: int):
+        from vllm.models.kimi_k3.amd.specialized_decode.mono.attention import mla_back
+
+        return mla_back.MlaPostBuild(
+            tokens=self.width,
+            qlen=qlen,
+            ranks=self.runner.dcp,
+            peer_off=self.runner.mla_post_off,
+            splits=m1c_splits(self.runner, self.width),
+        )
+
+    def builds(self) -> list:
+        # a layer after another always has a pending delta
+        if self.runner.fuse_m1b:
+            k2 = _k2_key(self.layer, self.runner, self.width, False, True)
+            return [self._m1a_key(True), k2]
+        m1b = self._m1b_key(self.runner.num_spec + 1)
+        return [self._m1a_key(True), m1b, *super().builds()]
+
+    def caches(self) -> torch.Tensor:
+        """The fp8 MLA cache [blocks, block_size, 576]."""
+        return self.wrapper.mla_attn.kv_cache
+
+    def launch(self, step: "Step", positions, prefix_sum, block_residual, prefix_delta):
+        from vllm.models.kimi_k3.amd.specialized_decode.mono.attention import mla as m1
+        from vllm.models.kimi_k3.amd.specialized_decode.mono.attention import mla_back
+        from vllm.models.kimi_k3.amd.specialized_decode.mono.attention.kda import HIDDEN
+        from vllm.models.kimi_k3.amd.specialized_decode.runtime import mla
+        from vllm.platforms import current_platform
+
+        layer, w, runner = self.layer, self.wrapper, self.runner
+        a = w.mla_attn
+        if self.first:
+            runner.step_begin()
+        md, _, kv_cache, slots = mla.context(w)
+        w_qkv, w_g = mla.qkv_gate_weights(w)
+        s = prefix_sum.size(0)
+        m1c = m1c_at(runner, self.width)
+        # AITER's cprr pad (e.g. 96 -> 128 heads) M1a writes in place of
+        # forward_mqa's copy (mla.attend hands it over)
+        qh = md.decode.asm_decode_num_heads if runner.dcp > 1 else 0
+        qh = qh if qh > runner.dcp * m1.NH else 0
+        q8 = prefix_sum.new_empty(
+            s, qh or runner.dcp * m1.NH, m1.ENTRY,
+            dtype=torch.bfloat16 if m1c else current_platform.fp8_dtype(),
+        )  # fmt: skip
+        gate = prefix_sum.new_empty(s, m1.GROWS)
+        m1.mla_pre(
+            self._m1a_key(prefix_delta is not None, qh),
+            prefix=prefix_sum,
+            delta=prefix_delta,
+            blocks=block_residual,
+            ares_nw=layer.self_attention_res_norm.weight,
+            ares_qk=layer.self_attention_res_proj.weight.view(HIDDEN),
+            in_nw=layer.input_layernorm.weight,
+            w_qkv=w_qkv,
+            w_g=w_g,
+            q_nw=w.q_a_layernorm.weight,
+            kv_nw=w.kv_a_layernorm.weight,
+            w_qb=w.q_b_proj.weight,
+            w_k=a.W_K,
+            w_k_scale=a.W_K_scale,
+            kv_cache=kv_cache,
+            slots=slots[:s],
+            gate_out=gate,
+            q_out=q8,
+            layer=layer.layer_idx,
+            **runner.launch_args(),
+            **runner.peer_args(),
+        )
+        # DCP1: K2 runs the attention (M1c), its KV splits' partials split 0's views
+        m1c_in = None
+        if m1c:
+            m1c_in, po, pl = mla.m1c_inputs(w, q8, kv_cache, md)
+            o, lse = po[0], pl[0]
+        else:
+            o, lse = mla.attend(w, q8, kv_cache, md)
+        core = torch.empty_like(gate)
+        m1b_in = dict(
+            o=o,
+            lse=lse,
+            seq_lens=md.decode.seq_lens,
+            qsl=md.query_start_loc,
+            w_v=a.W_V,
+            w_v_scale=a.W_V_scale,
+            gate=gate,
+        )
+        if runner.fuse_m1b:
+            out = _k2(
+                layer, runner, core, prefix_sum, block_residual, m1b_in=m1b_in,
+                m1c_in=m1c_in,
+            )  # fmt: skip
+            return prefix_sum, block_residual, out
+        mla_back.mla_post(
+            self._m1b_key(runner.num_spec + 1),
+            core=core,
+            layer=layer.layer_idx,
+            **m1b_in,
+            **runner.launch_args(),
+            **runner.peer_args(),
+        )
+        out = _k2(layer, runner, core, prefix_sum, block_residual)
+        return prefix_sum, block_residual, out
+
+
+class MlaOpsLayer(_Layer):
+    """An MLA layer whose attention runs as vLLM's ops (o_proj's input), then
+    K2: for deployments M1 cannot serve."""
+
     def launch(self, step: "Step", positions, prefix_sum, block_residual, prefix_delta):
         from vllm.models.kimi_k3.amd.linear import _apply_attn_res
 
diff --git a/vllm/models/kimi_k3/amd/specialized_decode/runtime/mla.py b/vllm/models/kimi_k3/amd/specialized_decode/runtime/mla.py
new file mode 100644
index 000000000000..5d36466c0848
--- /dev/null
+++ b/vllm/models/kimi_k3/amd/specialized_decode/runtime/mla.py
@@ -0,0 +1,131 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""An MLA layer's decode, split where the mono kernels split it (vLLM's ops,
+the reference M1a / M1b reproduce):
+
+- ``pre``: qkv_a, the q / kv norms, q_b, W_UK (fp4), the query's fp8 quant,
+  this rank's KV-cache rows, the DCP query all-gather (DCP8), the output gate.
+- ``attend``: AITER's MLA decode on the (gathered) query (``forward_mqa``),
+  on M1a's padded query as is. At DCP1 with qlen 4 / 8 K2 runs the decode
+  itself (M1c, ``m1c_inputs``).
+- ``post``: the DCP LSE combine (DCP8), W_UV (fp4), the gate: o_proj's input.
+
+Pure decode only (every row a decode row), as the mono step is.
+"""
+
+import torch
+
+from vllm._aiter_ops import rocm_aiter_ops
+from vllm.model_executor.layers.attention.attention import get_attention_context
+
+
+def context(wrapper):
+    """``(metadata, MLAAttention, kv_cache, slot_mapping)`` of the step."""
+    return get_attention_context(wrapper.mla_attn.layer_name)
+
+
+def qkv_gate_weights(wrapper):
+    """fused_qkv_a's [2112, 7168] and g_proj's [1536, 7168] weights: row
+    slices of one [q_a | kv_a | gate] weight where vLLM merges the gate in
+    (``KimiK3MergedQKVGateLinear``); the gate None without one."""
+    w = wrapper.fused_qkv_a_proj.weight
+    if wrapper.g_proj is not None:
+        return w, wrapper.g_proj.weight
+    if not getattr(wrapper, "_qkv_gate_fused", False):
+        return w, None
+    n = wrapper.q_lora_rank + wrapper.mla_attn.kv_lora_rank + wrapper.qk_rope_head_dim
+    return w[:n], w[n:]
+
+
+def pre(wrapper, hidden, kv_cache, slots, bf16=False):
+    """``hidden``: input_layernorm's output [S, 7168]. Returns the gathered fp8
+    query [S, heads * dcp, 576] (``bf16``, DCP1: bf16, unquantized) and the
+    gate sigmoid(g_proj) [S, 1536]; writes the rows whose slot is this
+    rank's."""
+    attn = wrapper.mla_attn
+    nope, rope = wrapper.qk_nope_head_dim, wrapper.qk_rope_head_dim
+    qkv = wrapper.fused_qkv_a_proj(hidden)[0]
+    if wrapper.g_proj is None:  # [q_a | kv_a | gate] in one GEMM
+        qkv, g = qkv.split([wrapper.q_lora_rank + attn.kv_lora_rank + rope, -1], -1)
+    else:
+        g = wrapper.g_proj(hidden)[0]
+    q_c, kv_lora = qkv.split([wrapper.q_lora_rank, attn.kv_lora_rank + rope], -1)
+    kv_c, k_pe = kv_lora.split([attn.kv_lora_rank, rope], -1)
+    q_c, kv_c = wrapper._normalize_q_kv(q_c, kv_c)
+    q = wrapper.q_b_proj(q_c)[0].view(-1, wrapper.num_heads, nope + rope)
+    q_nope, q_pe = q.split([nope, rope], -1)
+    ql_nope = rocm_aiter_ops.batched_gemm_a16wfp4(
+        q_nope.transpose(0, 1),
+        attn.W_K,
+        attn.W_K_scale,
+        transpose_bm=True,
+        prequant=True,
+        y_scale=None if bf16 else attn._q_scale,
+    )
+    if bf16:
+        q8 = torch.cat([ql_nope, q_pe], -1)
+    else:
+        q8 = attn._decode_concat_quant_fp8_op(ql_nope, q_pe, attn._q_scale)
+    attn.impl.do_kv_cache_update(
+        kv_c, k_pe.unsqueeze(1), kv_cache, slots, attn.kv_cache_dtype, attn._k_scale
+    )
+    if attn.impl.dcp_world_size > 1:
+        q8 = attn.dcp_manager.query_gather(q8)
+    gate = g.sigmoid()
+    return q8, gate
+
+
+def attend(wrapper, q8, kv_cache, md):
+    """The partial (this rank's KV) attention: ``o`` [S, heads * dcp, 512] and
+    its LSE [S, heads * dcp] (None at DCP1), both possibly strided views.
+
+    ``q8`` may carry more heads than the layer: AITER's cprr pad (e.g. 96 ->
+    128, M1a writes it). forward_mqa gets the real heads, a view whose row
+    stride spans the pad, and AITER's ``get_mla_padded_q`` widens it back to
+    ``q8`` instead of copying."""
+    from vllm.platforms import current_platform
+
+    attn = wrapper.mla_attn
+    kv = kv_cache.view(current_platform.fp8_dtype())
+    n = attn.impl._decode_num_heads
+    return attn.impl.forward_mqa(q8[:, :n], kv, md, attn)
+
+
+def m1c_inputs(wrapper, q, kv_cache, md):
+    """K2's ``m1c_in`` for M1a's bf16 query ``q`` [S, 12, 576] and its new
+    partials, o [splits, S, 12, 512] bf16 and LSE [splits, S, 12] f32 (M1b
+    merges the splits)."""
+    from vllm.models.kimi_k3.amd.specialized_decode.mono.attention import mla_attn
+    from vllm.platforms import current_platform
+
+    attn = wrapper.mla_attn
+    decode = md.decode
+    kv = kv_cache.view(current_platform.fp8_dtype())
+    key = mla_attn.MlaAttnBuild(tokens=q.size(0), qlen=int(decode.max_qo_len))
+    so, sl = mla_attn.partial_shapes(key)
+    m1c_in = dict(
+        q=q,
+        kv=kv.view(-1, kv.size(-1)),  # the flat token rows
+        pages=decode.paged_kv_indices,
+        indptr=decode.paged_kv_indptr,
+        scale=attn.impl.scale,
+        kv_scale=float(attn._k_scale_float),
+    )
+    po = q.new_empty(so)
+    return m1c_in, po, torch.empty(sl, dtype=torch.float32, device=q.device)
+
+
+def post(wrapper, o, lse, md, gate):
+    """o_proj's input [S, heads * 128] bf16."""
+    attn = wrapper.mla_attn
+    latent = o
+    if attn.impl.dcp_world_size > 1:
+        latent = attn.dcp_manager.combine(
+            o,
+            lse,
+            seq_lens=md.decode.seq_lens,
+            query_start_loc=md.query_start_loc[: md.num_decodes + 1],
+        )
+    core = gate.new_empty(gate.shape)
+    attn._v_up_proj(latent, out=core)
+    return core * gate
diff --git a/vllm/models/kimi_k3/amd/specialized_decode/runtime/runner.py b/vllm/models/kimi_k3/amd/specialized_decode/runtime/runner.py
index 85119c7bdc90..84ef02dd6874 100644
--- a/vllm/models/kimi_k3/amd/specialized_decode/runtime/runner.py
+++ b/vllm/models/kimi_k3/amd/specialized_decode/runtime/runner.py
@@ -14,19 +14,25 @@
 
 import torch
 
-from vllm.distributed import get_tp_group
+from vllm.distributed import get_dcp_group, get_tp_group
 from vllm.models.kimi_k3.amd.specialized_decode.mono.abi import PEER_HEAD
 from vllm.models.kimi_k3.amd.specialized_decode.mono.common.plan import HEALTH_WORDS
-from vllm.models.kimi_k3.amd.specialized_decode.runtime import MAX_TOKENS
+from vllm.models.kimi_k3.amd.specialized_decode.runtime import M1C_QLENS, MAX_TOKENS
 from vllm.models.kimi_k3.amd.specialized_decode.runtime.peer import PeerBuffer
 
 
 class MonoRunner:
     """Collective: every TP rank builds it at the same point."""
 
-    def __init__(self, device: torch.device):
+    def __init__(self, device: torch.device, dcp: int | None = None, m1: bool = True):
+        """``dcp``: M1's ranks, 1 or the TP group (default: the DCP group's).
+        ``m1``: the MLA layers run M1 (else vLLM's attention ops)."""
         from vllm.models.kimi_k3.amd.specialized_decode.mono import layer
-        from vllm.models.kimi_k3.amd.specialized_decode.mono.attention import kda
+        from vllm.models.kimi_k3.amd.specialized_decode.mono.attention import (
+            kda,
+            mla,
+            mla_back,
+        )
 
         self.scratch = torch.zeros(
             max(
@@ -40,26 +46,63 @@ def __init__(self, device: torch.device):
                     )
                 ),
                 layer.scratch_bytes(layer.K2Build(tokens=MAX_TOKENS, nblocks=8)),
+                mla.scratch_bytes(
+                    mla.MlaPreBuild(tokens=MAX_TOKENS, nblocks=8, delta=True)
+                ),
+                # K2's tagged x rows (layer.OPT_TAG_X), past every layout
+                layer.tag_end(),
             ),
             dtype=torch.uint8,
             device=device,
         )
+        layer.TAG_READY = True
         self.epoch = torch.zeros(1, dtype=torch.int32, device=device)
         self.health = torch.zeros(HEALTH_WORDS, dtype=torch.int32, device=device)
 
         tp = get_tp_group()
         self.rank = tp.rank_in_group
-        nbytes = PEER_HEAD + layer.peer_bytes(MAX_TOKENS)
+        self.world = tp.world_size
+        self.dcp = get_dcp_group().world_size if dcp is None else dcp
+        # K2's regions, M1a's gathered query (QG), M1b's partials (CMB); M1
+        # takes no peer memory at DCP1
+        self.m1 = m1
+        m1 = m1 and self.dcp > 1
+        self.mla_peer_off = layer.peer_bytes(MAX_TOKENS)
+        self.mla_post_off = self.mla_peer_off + (
+            mla.peer_bytes(MAX_TOKENS, self.dcp) if m1 else 0
+        )
+        nbytes = (
+            PEER_HEAD
+            + self.mla_post_off
+            + (mla_back.peer_bytes(MAX_TOKENS, self.dcp) if m1 else 0)
+        )
+        # K2's sentinel regions (layer.OPT_SENTS, all ones, both layer
+        # parities) and RLOGP (layer.OPT_RLOG_PARTS), past every other region
+        sent_lo = PEER_HEAD + layer.sent_offset()
+        nbytes = max(nbytes, PEER_HEAD + layer.extra_peer_end())
         self.peers = PeerBuffer(nbytes, tp.cpu_group, self.rank, tp.world_size, device)
         self.peers.bytes.zero_()
+        self.peers.bytes[sent_lo : sent_lo + layer.sent_bytes()].fill_(0xFF)
+        torch.accelerator.synchronize()
+        torch.distributed.barrier(group=tp.cpu_group)
+        layer.SENT_READY = True
         self.peer_table = self.peers.addresses + PEER_HEAD
         self.alarm = int(self.peers.addresses[0].item()) + 4 * self.rank
         self.flags = self.peers.bytes[: 4 * tp.world_size].view(torch.int32)
         self.fault: dict = {}  # K2Build's fault_* fields (tests only)
+        self.fuse_m1b = True  # an MLA layer's M1b runs in its K2's launch
         self.first_layer = -1  # the step's first kernel layer: bumps the epoch
         self.kda_prefix = ""  # a KDA layer's name: the step's metadata
+        self.mla_prefix = ""  # an M1 layer's name: its decode metadata
+        self.interleave = 1  # the DCP KV cache interleave
         self.num_spec = 0  # speculative tokens a request
 
+    @property
+    def m1c(self) -> bool:
+        """DCP1 at qlen 4 / 8 (M1b in K2): M1a writes a bf16 query, K2's launch
+        runs the decode attention (M1c) first, its M1b merging the KV splits."""
+        return self.dcp == 1 and self.num_spec + 1 in M1C_QLENS and self.fuse_m1b
+
     def step_begin(self) -> None:
         """Once a forward, before its first mono launch (a graph captures it)."""
         self.epoch.add_(1)
diff --git a/vllm/models/kimi_k3/nvidia/dspark_mla.py b/vllm/models/kimi_k3/nvidia/dspark_mla.py
index d591c69a21e1..42bb3af7e8b6 100644
--- a/vllm/models/kimi_k3/nvidia/dspark_mla.py
+++ b/vllm/models/kimi_k3/nvidia/dspark_mla.py
@@ -29,6 +29,7 @@
 from vllm.models.kimi_k3.nvidia.mla import MultiHeadLatentAttention
 from vllm.models.kimi_k3.nvidia.model import KimiMLP
 from vllm.v1.kv_cache_interface import is_quantized_kv_cache
+from vllm.v1.worker.gpu.spec_decode.dspark.config import DSparkSpeculatorConfig
 from vllm.v1.worker.workspace import current_workspace_manager
 
 _GROUPED_KV_CACHE_DTYPES = frozenset(
@@ -431,6 +432,9 @@ def forward(
 
 
 class K3DSparkForCausalLM(nn.Module):
+    # The speculator this is a draft model of (see SpeculativeDraftModel).
+    speculator_config_class = DSparkSpeculatorConfig
+
     has_own_embed_tokens = False
     has_own_lm_head = False
     draft_id_to_target_id = None
diff --git a/vllm/models/kimi_k3/nvidia/kda.py b/vllm/models/kimi_k3/nvidia/kda.py
index a921d1bd2879..d3329771181a 100644
--- a/vllm/models/kimi_k3/nvidia/kda.py
+++ b/vllm/models/kimi_k3/nvidia/kda.py
@@ -73,7 +73,10 @@
 )
 from vllm.v1.attention.backend import AttentionBackend
 from vllm.v1.kv_cache_interface import MambaSpec
-from vllm.v1.worker.workspace import current_workspace_manager
+from vllm.v1.worker.workspace import (
+    current_workspace_manager,
+    is_workspace_manager_initialized,
+)
 
 logger = init_logger(__name__)
 
@@ -857,6 +860,18 @@ def __init__(
                 self.model_config.dtype,
             )
 
+        # Profiling skips attention, so reserve its runtime layout before the
+        # KV cache is sized. Layers share the manager's grow-only allocation.
+        if is_workspace_manager_initialized():
+            if self._flashkda_buffer_specs is not None:
+                current_workspace_manager().get_simultaneous(
+                    *self._flashkda_buffer_specs
+                )
+            elif self._flashinfer_kda_output_spec is not None:
+                current_workspace_manager().get_simultaneous(
+                    self._flashinfer_kda_output_spec
+                )
+
         self.o_norm = FusedRMSNormGated(self.head_dim, activation="sigmoid")
         decode_norm_weight = None
         if decode_conv1d_weight is not None:
diff --git a/vllm/models/kimi_k3/nvidia/mla.py b/vllm/models/kimi_k3/nvidia/mla.py
index 2e940b02993c..ab7d1036151b 100644
--- a/vllm/models/kimi_k3/nvidia/mla.py
+++ b/vllm/models/kimi_k3/nvidia/mla.py
@@ -321,6 +321,7 @@ def __init__(
             self.kv_cache_dtype = "auto"
 
         dtype = torch.get_default_dtype()
+        self.act_dtype = dtype
         self.attn_backend = get_attn_backend(
             self.head_size,
             dtype,
@@ -436,7 +437,7 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
             non_causal_multi_token_decode=self.non_causal_multi_token_decode,
         )
 
-    def process_weights_after_loading(self, act_dtype: torch.dtype) -> None:
+    def process_weights_after_loading(self) -> None:
         """Absorb ``kv_b_proj`` into decode-time ``W_UK_T`` / ``W_UV`` bmm weights.
 
         ``kv_b_proj`` produces ``[k_nope; v]`` per head from the ``kv_lora_rank``
@@ -449,7 +450,7 @@ def process_weights_after_loading(self, act_dtype: torch.dtype) -> None:
         if not (is_weights_pre_processed() and getattr(self, "W_UV", None) is not None):
             W_UK, W_UV = split_kv_b_proj(
                 self.kv_b_proj,
-                act_dtype,
+                self.act_dtype,
                 self.kv_lora_rank,
                 self.num_local_heads,
                 self.qk_nope_head_dim,
diff --git a/vllm/models/kimi_k3/nvidia/mtp.py b/vllm/models/kimi_k3/nvidia/mtp.py
index 06a0d3377843..067fd89359b2 100644
--- a/vllm/models/kimi_k3/nvidia/mtp.py
+++ b/vllm/models/kimi_k3/nvidia/mtp.py
@@ -23,6 +23,9 @@
     ParallelLMHead,
     VocabParallelEmbedding,
 )
+from vllm.model_executor.model_loader.mtp_validation import (
+    is_mtp_completeness_check_enabled,
+)
 from vllm.model_executor.model_loader.weight_utils import (
     default_weight_loader,
     maybe_remap_kv_scale_name,
@@ -413,7 +416,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
             self.model.mtp_start_layer_idx,
             self.model.mtp_start_layer_idx + self.model.num_mtp_layers,
         ):
-            if layer_idx not in loaded_layers:
+            if layer_idx not in loaded_layers and is_mtp_completeness_check_enabled():
                 raise ValueError(
                     f"MTP speculative decoding layer {layer_idx} weights "
                     f"missing from checkpoint. The checkpoint may not include "
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 67d5cfdf081d..41b10f7d1dc9 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
@@ -15,15 +15,16 @@
 import torch.distributed as dist
 import torch.distributed._symmetric_memory as symm_mem
 from cutlass import BFloat16, Float32, Int32, Int64, Uint32
+from cutlass._mlir import ir
+from cutlass._mlir.dialects import llvm, vector
+from cutlass.cutlass_dsl import T, dsl_user_op
 
 from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput
-
-from .primitives import (
-    NUM_LAMPORT_BUFFERS,
+from vllm.models.common.ops.moe_tail_fusion.primitives import (
+    LAMPORT_GENERATIONS,
     PACKED_BYTES,
     VEC_BF16,
     bf16x8_to_packed_u32x4,
-    finalize_top16_bf16,
     fragment_is_dirty,
     load_global_u32x4,
     load_shared_f32x2,
@@ -32,7 +33,7 @@
     map_shared_to_peer,
     packed_u32x4_to_bf16x8,
     red_async_release_gpu_add_u32,
-    sanitize_negative_zero,
+    sanitize_negative_zero_u32x4,
     stmc_bf16x8,
     store_global_u32x4,
     store_lamport_sentinel_128,
@@ -45,6 +46,103 @@
 _SEVEN_CTA_MAX_M = 4
 
 
+def _make_top16_bf16_finalize_asm() -> str:
+    lines = [
+        "{",
+        ".reg .pred _valid<16>;",
+        ".reg .s32 _index<16>;",
+        ".reg .b32 _weight_words<8>;",
+        ".reg .b16 _weights<16>;",
+        ".reg .b32 _data<64>;",
+        ".reg .b16 _lo, _hi;",
+        ".reg .f32 _acc<8>;",
+        ".reg .u64 _address<16>;",
+        "ld.global.v4.u32 {_index0, _index1, _index2, _index3}, [$5];",
+        "ld.global.v4.u32 {_index4, _index5, _index6, _index7}, [$5+16];",
+        "ld.global.v4.u32 {_index8, _index9, _index10, _index11}, [$5+32];",
+        "ld.global.v4.u32 {_index12, _index13, _index14, _index15}, [$5+48];",
+        "ld.global.v4.u32 "
+        "{_weight_words0, _weight_words1, _weight_words2, "
+        "_weight_words3}, [$6];",
+        "ld.global.v4.u32 "
+        "{_weight_words4, _weight_words5, _weight_words6, "
+        "_weight_words7}, [$6+16];",
+    ]
+    for pair in range(8):
+        lines.append(
+            f"mov.b32 {{_weights{2 * pair}, _weights{2 * pair + 1}}}, "
+            f"_weight_words{pair};"
+        )
+    for element in range(8):
+        lines.append(f"mov.f32 _acc{element}, 0f00000000;")
+    for route in range(16):
+        data = 4 * route
+        lines.extend(
+            [
+                f"setp.ge.s32 _valid{route}, _index{route}, 0;",
+                f"@_valid{route} mad.wide.s32 _address{route}, "
+                f"_index{route}, 7168, $4;",
+                f"@_valid{route} ld.global.v4.u32 "
+                f"{{_data{data}, _data{data + 1}, _data{data + 2}, "
+                f"_data{data + 3}}}, [_address{route}];",
+            ]
+        )
+    for route in range(16):
+        for pair in range(4):
+            data = 4 * route + pair
+            element = 2 * pair
+            lines.extend(
+                [
+                    f"@_valid{route} mov.b32 {{_lo, _hi}}, _data{data};",
+                    f"@_valid{route} fma.rn.f32.bf16 _acc{element}, _lo, "
+                    f"_weights{route}, _acc{element};",
+                    f"@_valid{route} fma.rn.f32.bf16 _acc{element + 1}, _hi, "
+                    f"_weights{route}, _acc{element + 1};",
+                ]
+            )
+    for pair in range(4):
+        lines.append(f"cvt.rn.bf16x2.f32 ${pair}, _acc{2 * pair + 1}, _acc{2 * pair};")
+    lines.append("}")
+    return "\n".join(lines)
+
+
+_TOP16_BF16_FINALIZE_ASM = _make_top16_bf16_finalize_asm()
+
+
+@dsl_user_op
+def finalize_top16_bf16(
+    gemm2_vector: cute.Pointer,
+    route_indices: cute.Pointer,
+    route_weights: cute.Pointer,
+    *,
+    loc=None,
+    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)
+    ]
+    out = llvm.inline_asm(
+        llvm.StructType.get_literal([T.i32()] * 4),
+        addresses,
+        _TOP16_BF16_FINALIZE_ASM,
+        "=r,=r,=r,=r,l,l,l",
+        has_side_effects=False,
+        is_align_stack=False,
+        asm_dialect=llvm.AsmDialect.AD_ATT,
+        loc=loc,
+        ip=ip,
+    )
+    packed = vector.from_elements(
+        ir.VectorType.get([4], T.i32(), loc=loc),
+        [llvm.extractvalue(T.i32(), out, [i], loc=loc, ip=ip) for i in range(4)],
+        loc=loc,
+        ip=ip,
+    )
+    return cute.TensorSSA(packed, 4, Uint32)
+
+
 def _mapping(
     tp_size: int,
     latent_dim: int,
@@ -367,7 +465,7 @@ def _token_device(
                         cute.AddressSpace.gmem,
                         assumed_align=16,
                     )
-                    local_packed = sanitize_negative_zero(
+                    local_packed = sanitize_negative_zero_u32x4(
                         finalize_top16_bf16(
                             gemm2_vector,
                             route_indices,
@@ -402,7 +500,7 @@ def _token_device(
                                     local_values[element]
                                     + values[element].to(Float32) * weight
                                 )
-                    local_packed = sanitize_negative_zero(
+                    local_packed = sanitize_negative_zero_u32x4(
                         bf16x8_to_packed_u32x4(local_values.load().to(BFloat16))
                     )
             else:
@@ -412,7 +510,7 @@ def _token_device(
                     cute.AddressSpace.gmem,
                     assumed_align=16,
                 )
-                local_packed = sanitize_negative_zero(
+                local_packed = sanitize_negative_zero_u32x4(
                     load_global_u32x4(local_ptr, volatile=False)
                 )
             multicast_offset = (
@@ -592,7 +690,7 @@ def _token_device(
                     arrived = load_volatile_u32(access_counter)
                     while arrived < Uint32(m):
                         arrived = load_volatile_u32(access_counter)
-                next_index = (current_index + Uint32(1)) % Uint32(NUM_LAMPORT_BUFFERS)
+                next_index = (current_index + Uint32(1)) % Uint32(LAMPORT_GENERATIONS)
                 actual_bytes = Uint32(m) * Uint32(self.tp_size * self.latent_dim * 2)
                 cute.arch.store((latent_flags.iterator + 0).llvm_ptr, next_index)
                 cute.arch.store((latent_flags.iterator + 1).llvm_ptr, current_index)
@@ -653,7 +751,7 @@ def _token_device(
                             cute.AddressSpace.gmem,
                             assumed_align=16,
                         )
-                        local_packed = sanitize_negative_zero(
+                        local_packed = sanitize_negative_zero_u32x4(
                             load_global_u32x4(source_ptr, volatile=False)
                         )
                         destination_element = current_elements + (
@@ -771,7 +869,7 @@ def _token_device(
                 target = Uint32(m) * Uint32(self.shared_roles)
                 while arrived < target:
                     arrived = load_volatile_u32(access_counter)
-                next_index = (current_index + Uint32(1)) % Uint32(NUM_LAMPORT_BUFFERS)
+                next_index = (current_index + Uint32(1)) % Uint32(LAMPORT_GENERATIONS)
                 actual_bytes = Uint32(m) * Uint32(self.tp_size * self.shard_dim * 2)
                 cute.arch.store((shared_flags.iterator + 0).llvm_ptr, next_index)
                 cute.arch.store((shared_flags.iterator + 1).llvm_ptr, current_index)
@@ -1071,7 +1169,7 @@ def __init__(
         )
 
         bytes_per_routed_buffer = max_m * tp_size * latent_dim * 2
-        routed_bytes = NUM_LAMPORT_BUFFERS * bytes_per_routed_buffer
+        routed_bytes = LAMPORT_GENERATIONS * bytes_per_routed_buffer
         self._routed_workspace = symm_mem.empty(
             routed_bytes // 4,
             dtype=torch.float32,
@@ -1080,7 +1178,7 @@ def __init__(
         self._routed_symm_mem = symm_mem.rendezvous(self._routed_workspace, group)
         self._routed_workspace.fill_(-0.0)
         actual_bytes_per_buffer = (
-            self._routed_symm_mem.buffer_size // NUM_LAMPORT_BUFFERS // 16 * 16
+            self._routed_symm_mem.buffer_size // LAMPORT_GENERATIONS // 16 * 16
         )
         if actual_bytes_per_buffer < bytes_per_routed_buffer:
             raise RuntimeError("routed symmetric workspace is too small")
@@ -1105,7 +1203,7 @@ def __init__(
         self._shared_shard = self._shared_output[:, shard_start:shard_end]
 
         self._shared_workspace = symm_mem.empty(
-            (NUM_LAMPORT_BUFFERS, max_m, tp_size, self.shard_dim),
+            (LAMPORT_GENERATIONS, max_m, tp_size, self.shard_dim),
             dtype=torch.bfloat16,
             device=device,
         )
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 ea9d86d6fa85..fd91fad1f398 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
@@ -52,10 +52,11 @@
 from cutlass.cute.runtime import from_dlpack
 from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
 
+from vllm.models.common.ops.moe_tail_fusion.primitives import CUDAGraphCompatibleWrapper
+
 from .fused_add_multicast_skinny_gemm import (
     FusedAddMulticastSkinnyGemmKernel,
 )
-from .primitives import CUDAGraphCompatibleWrapper
 
 
 def _as_cute(tensor: torch.Tensor, *, dynamic_m: bool = False):
diff --git a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/fused_add_multicast_skinny_gemm.py b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/fused_add_multicast_skinny_gemm.py
index 21b27b278f60..635939e7fcee 100644
--- a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/fused_add_multicast_skinny_gemm.py
+++ b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/fused_add_multicast_skinny_gemm.py
@@ -14,15 +14,15 @@
 from cutlass import BFloat16, Float32, Int64, const_expr
 from cutlass.cute.runtime import from_dlpack
 
-from .primitives import (
+from vllm.models.common.ops.moe_tail_fusion.primitives import (
     CUDAGraphCompatibleWrapper,
     bf16x2_to_u32,
     bf16x4_to_packed_u32x2,
     bf16x8_to_packed_u32x4,
     fma_f32_bf16,
-    sanitize_negative_zero,
     sanitize_negative_zero_u32,
     sanitize_negative_zero_u32x2,
+    sanitize_negative_zero_u32x4,
     store_global_u32,
     store_global_u32x2,
     store_global_u32x4,
@@ -283,7 +283,7 @@ def kernel(
                         packed,
                     )
                 else:
-                    packed = sanitize_negative_zero(
+                    packed = sanitize_negative_zero_u32x4(
                         bf16x8_to_packed_u32x4(fused.load())
                     )
                     store_global_u32x4(
diff --git a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/lamport_copy.py b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/lamport_copy.py
index b2a9100b711c..0eccdd9b8740 100644
--- a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/lamport_copy.py
+++ b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/lamport_copy.py
@@ -13,7 +13,7 @@
 import torch
 from cutlass.cute.runtime import make_fake_compact_tensor, make_fake_stream
 
-from .primitives import (
+from vllm.models.common.ops.moe_tail_fusion.primitives import (
     VEC_BF16,
     fragment_is_dirty,
     load_global_u32x4,
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
deleted file mode 100644
index 3840caf24a27..000000000000
--- a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/primitives.py
+++ /dev/null
@@ -1,576 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-"""Shared CuTe DSL primitives; this module does not define a CUDA kernel."""
-
-from __future__ import annotations
-
-import cutlass
-import cutlass.cute as cute
-import torch
-from cutlass import BFloat16, Float32, Int32, Int64, Uint16, Uint32
-from cutlass._mlir import ir
-from cutlass._mlir.dialects import llvm, vector
-from cutlass.cute.runtime import from_dlpack
-from cutlass.cutlass_dsl import T, dsl_user_op
-
-VEC_BF16 = 8
-PACKED_BYTES = 16
-NUM_LAMPORT_BUFFERS = 3
-NEG_ZERO_F32_BITS = 0x80000000
-NEG_ZERO_BF16_BITS = 0x8000
-
-
-class CUDAGraphCompatibleWrapper:
-    """DLPack view that does not synchronize with the producer stream."""
-
-    def __init__(self, tensor: torch.Tensor):
-        self.tensor = tensor
-
-    def __dlpack__(self, stream=None):
-        return self.tensor.__dlpack__(stream=-1)
-
-    def __dlpack_device__(self):
-        return self.tensor.__dlpack_device__()
-
-
-def to_cute(tensor: torch.Tensor, assumed_align: int = 16) -> cute.Tensor:
-    return from_dlpack(
-        CUDAGraphCompatibleWrapper(tensor.detach()), assumed_align=assumed_align
-    )
-
-
-def to_cute_dynamic_m(
-    tensor: torch.Tensor,
-    *,
-    mode: int,
-    assumed_align: int = 16,
-) -> cute.Tensor:
-    """Expose exactly one compact tensor mode as a runtime shape.
-
-    Model dimensions remain part of the compiled tensor type. Only the token
-    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(),
-    )
-
-
-@dsl_user_op
-def fma_f32_bf16(
-    a: BFloat16,
-    b: BFloat16,
-    acc: Float32,
-    *,
-    loc=None,
-    ip=None,
-) -> Float32:
-    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)],
-        "fma.rn.f32.bf16 $0, $1, $2, $3;",
-        "=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)
-
-
-@dsl_user_op
-def load_global_u32x4(
-    pointer: cute.Pointer,
-    *,
-    volatile: cutlass.Constexpr[bool] = False,
-    loc=None,
-    ip=None,
-):
-    """Load one 128-bit fragment as four u32 registers.
-
-    The volatile form is the Lamport polling load.  Marking the asm
-    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(
-        llvm.StructType.get_literal([T.i32()] * 4),
-        [address.ir_value(loc=loc, ip=ip)],
-        f"{opcode} {{$0, $1, $2, $3}}, [$4];",
-        "=r,=r,=r,=r,l",
-        has_side_effects=volatile,
-        is_align_stack=False,
-        asm_dialect=llvm.AsmDialect.AD_ATT,
-        loc=loc,
-        ip=ip,
-    )
-    packed = vector.from_elements(
-        ir.VectorType.get([4], T.i32(), loc=loc),
-        [llvm.extractvalue(T.i32(), out, [i], loc=loc, ip=ip) for i in range(4)],
-        loc=loc,
-        ip=ip,
-    )
-    return cute.TensorSSA(packed, 4, Uint32)
-
-
-def _make_top16_bf16_finalize_asm() -> str:
-    lines = [
-        "{",
-        ".reg .pred _valid<16>;",
-        ".reg .s32 _index<16>;",
-        ".reg .b32 _weight_words<8>;",
-        ".reg .b16 _weights<16>;",
-        ".reg .b32 _data<64>;",
-        ".reg .b16 _lo, _hi;",
-        ".reg .f32 _acc<8>;",
-        ".reg .u64 _address<16>;",
-        "ld.global.v4.u32 {_index0, _index1, _index2, _index3}, [$5];",
-        "ld.global.v4.u32 {_index4, _index5, _index6, _index7}, [$5+16];",
-        "ld.global.v4.u32 {_index8, _index9, _index10, _index11}, [$5+32];",
-        "ld.global.v4.u32 {_index12, _index13, _index14, _index15}, [$5+48];",
-        "ld.global.v4.u32 "
-        "{_weight_words0, _weight_words1, _weight_words2, "
-        "_weight_words3}, [$6];",
-        "ld.global.v4.u32 "
-        "{_weight_words4, _weight_words5, _weight_words6, "
-        "_weight_words7}, [$6+16];",
-    ]
-    for pair in range(8):
-        lines.append(
-            f"mov.b32 {{_weights{2 * pair}, _weights{2 * pair + 1}}}, "
-            f"_weight_words{pair};"
-        )
-    for element in range(8):
-        lines.append(f"mov.f32 _acc{element}, 0f00000000;")
-    for route in range(16):
-        data = 4 * route
-        lines.extend(
-            [
-                f"setp.ge.s32 _valid{route}, _index{route}, 0;",
-                f"@_valid{route} mad.wide.s32 _address{route}, "
-                f"_index{route}, 7168, $4;",
-                f"@_valid{route} ld.global.v4.u32 "
-                f"{{_data{data}, _data{data + 1}, _data{data + 2}, "
-                f"_data{data + 3}}}, [_address{route}];",
-            ]
-        )
-    for route in range(16):
-        for pair in range(4):
-            data = 4 * route + pair
-            element = 2 * pair
-            lines.extend(
-                [
-                    f"@_valid{route} mov.b32 {{_lo, _hi}}, _data{data};",
-                    f"@_valid{route} fma.rn.f32.bf16 _acc{element}, _lo, "
-                    f"_weights{route}, _acc{element};",
-                    f"@_valid{route} fma.rn.f32.bf16 _acc{element + 1}, _hi, "
-                    f"_weights{route}, _acc{element + 1};",
-                ]
-            )
-    for pair in range(4):
-        lines.append(f"cvt.rn.bf16x2.f32 ${pair}, _acc{2 * pair + 1}, _acc{2 * pair};")
-    lines.append("}")
-    return "\n".join(lines)
-
-
-_TOP16_BF16_FINALIZE_ASM = _make_top16_bf16_finalize_asm()
-
-
-@dsl_user_op
-def finalize_top16_bf16(
-    gemm2_vector: cute.Pointer,
-    route_indices: cute.Pointer,
-    route_weights: cute.Pointer,
-    *,
-    loc=None,
-    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)
-    ]
-    out = llvm.inline_asm(
-        llvm.StructType.get_literal([T.i32()] * 4),
-        addresses,
-        _TOP16_BF16_FINALIZE_ASM,
-        "=r,=r,=r,=r,l,l,l",
-        has_side_effects=False,
-        is_align_stack=False,
-        asm_dialect=llvm.AsmDialect.AD_ATT,
-        loc=loc,
-        ip=ip,
-    )
-    packed = vector.from_elements(
-        ir.VectorType.get([4], T.i32(), loc=loc),
-        [llvm.extractvalue(T.i32(), out, [i], loc=loc, ip=ip) for i in range(4)],
-        loc=loc,
-        ip=ip,
-    )
-    return cute.TensorSSA(packed, 4, Uint32)
-
-
-@dsl_user_op
-def store_global_u32x4(
-    address: Int64,
-    packed,
-    *,
-    volatile: cutlass.Constexpr[bool] = False,
-    loc=None,
-    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(
-        None,
-        [address.ir_value(loc=loc, ip=ip), *words],
-        f"{opcode} [$0], {{$1, $2, $3, $4}};",
-        "l,r,r,r,r",
-        has_side_effects=True,
-        is_align_stack=False,
-        asm_dialect=llvm.AsmDialect.AD_ATT,
-        loc=loc,
-        ip=ip,
-    )
-
-
-@dsl_user_op
-def store_global_u32x2(
-    address: Int64,
-    packed,
-    *,
-    loc=None,
-    ip=None,
-) -> None:
-    words = [packed[i].ir_value(loc=loc, ip=ip) for i in range(2)]
-    llvm.inline_asm(
-        None,
-        [address.ir_value(loc=loc, ip=ip), *words],
-        "st.global.v2.u32 [$0], {$1, $2};",
-        "l,r,r",
-        has_side_effects=True,
-        is_align_stack=False,
-        asm_dialect=llvm.AsmDialect.AD_ATT,
-        loc=loc,
-        ip=ip,
-    )
-
-
-@dsl_user_op
-def store_global_u32(
-    address: Int64,
-    word: Uint32,
-    *,
-    loc=None,
-    ip=None,
-) -> None:
-    llvm.inline_asm(
-        None,
-        [
-            address.ir_value(loc=loc, ip=ip),
-            word.ir_value(loc=loc, ip=ip),
-        ],
-        "st.global.u32 [$0], $1;",
-        "l,r",
-        has_side_effects=True,
-        is_align_stack=False,
-        asm_dialect=llvm.AsmDialect.AD_ATT,
-        loc=loc,
-        ip=ip,
-    )
-
-
-@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,
-        [address.ir_value(loc=loc, ip=ip), *words],
-        "multimem.st.relaxed.sys.global.v4.bf16x2 [$0], {$1, $2, $3, $4};",
-        "l,r,r,r,r",
-        has_side_effects=True,
-        is_align_stack=False,
-        asm_dialect=llvm.AsmDialect.AD_ATT,
-        loc=loc,
-        ip=ip,
-    )
-
-
-@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(
-        None,
-        [address.ir_value(loc=loc, ip=ip), value, value, value, value],
-        "st.global.v4.u32 [$0], {$1, $2, $3, $4};",
-        "l,r,r,r,r",
-        has_side_effects=True,
-        is_align_stack=False,
-        asm_dialect=llvm.AsmDialect.AD_ATT,
-        loc=loc,
-        ip=ip,
-    )
-
-
-@dsl_user_op
-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,
-        [
-            address.ir_value(loc=loc, ip=ip),
-            value.ir_value(loc=loc, ip=ip),
-        ],
-        "red.async.release.global.gpu.add.u32 [$0], $1;",
-        "l,r",
-        has_side_effects=True,
-        is_align_stack=False,
-        asm_dialect=llvm.AsmDialect.AD_ATT,
-        loc=loc,
-        ip=ip,
-    )
-
-
-@dsl_user_op
-def load_volatile_u32(pointer: cute.Pointer, *, loc=None, ip=None) -> Uint32:
-    address = pointer.toint(loc=loc, ip=ip)
-    return Uint32(
-        llvm.inline_asm(
-            T.i32(),
-            [address.ir_value(loc=loc, ip=ip)],
-            "ld.volatile.global.u32 $0, [$1];",
-            "=r,l",
-            has_side_effects=True,
-            is_align_stack=False,
-            asm_dialect=llvm.AsmDialect.AD_ATT,
-            loc=loc,
-            ip=ip,
-        )
-    )
-
-
-@dsl_user_op
-def map_shared_to_peer(
-    smem_ptr: cute.Pointer,
-    peer_rank: Int32,
-    *,
-    loc=None,
-    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(
-            T.i32(),
-            [smem_address, peer_rank.ir_value(loc=loc, ip=ip)],
-            "mapa.shared::cluster.u32 $0, $1, $2;",
-            "=r,r,r",
-            has_side_effects=False,
-            is_align_stack=False,
-            asm_dialect=llvm.AsmDialect.AD_ATT,
-            loc=loc,
-            ip=ip,
-        )
-    )
-
-
-@dsl_user_op
-def store_shared_cluster_f32(
-    remote_address: Int32,
-    value: Float32,
-    *,
-    loc=None,
-    ip=None,
-) -> None:
-    llvm.inline_asm(
-        None,
-        [
-            remote_address.ir_value(loc=loc, ip=ip),
-            value.ir_value(loc=loc, ip=ip),
-        ],
-        "st.shared::cluster.f32 [$0], $1;",
-        "r,f",
-        has_side_effects=True,
-        is_align_stack=False,
-        asm_dialect=llvm.AsmDialect.AD_ATT,
-        loc=loc,
-        ip=ip,
-    )
-
-
-@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),
-        [address.ir_value(loc=loc, ip=ip)],
-        "ld.shared.v2.f32 {$0, $1}, [$2];",
-        "=f,=f,r",
-        has_side_effects=True,
-        is_align_stack=False,
-        asm_dialect=llvm.AsmDialect.AD_ATT,
-        loc=loc,
-        ip=ip,
-    )
-    return tuple(
-        Float32(llvm.extractvalue(T.f32(), out, [i], loc=loc, ip=ip)) for i in range(2)
-    )
-
-
-@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),
-        [address.ir_value(loc=loc, ip=ip)],
-        "ld.shared.v4.f32 {$0, $1, $2, $3}, [$4];",
-        "=f,=f,=f,=f,r",
-        has_side_effects=True,
-        is_align_stack=False,
-        asm_dialect=llvm.AsmDialect.AD_ATT,
-        loc=loc,
-        ip=ip,
-    )
-    return tuple(
-        Float32(llvm.extractvalue(T.f32(), out, [i], loc=loc, ip=ip)) for i in range(4)
-    )
-
-
-@dsl_user_op
-def packed_u32x4_to_bf16x8(packed, *, loc=None, ip=None):
-    target_type = ir.VectorType.get([VEC_BF16], BFloat16.mlir_type, loc=loc)
-    values = llvm.bitcast(
-        target_type,
-        packed.ir_value(loc=loc, ip=ip),
-        loc=loc,
-        ip=ip,
-    )
-    return cute.TensorSSA(values, VEC_BF16, BFloat16)
-
-
-@dsl_user_op
-def bf16x8_to_packed_u32x4(values, *, loc=None, ip=None):
-    target_type = ir.VectorType.get([4], T.i32(), loc=loc)
-    packed = llvm.bitcast(
-        target_type,
-        values.ir_value(loc=loc, ip=ip),
-        loc=loc,
-        ip=ip,
-    )
-    return cute.TensorSSA(packed, 4, Uint32)
-
-
-@dsl_user_op
-def bf16x4_to_packed_u32x2(values, *, loc=None, ip=None):
-    target_type = ir.VectorType.get([2], T.i32(), loc=loc)
-    packed = llvm.bitcast(
-        target_type,
-        values.ir_value(loc=loc, ip=ip),
-        loc=loc,
-        ip=ip,
-    )
-    return cute.TensorSSA(packed, 2, Uint32)
-
-
-@dsl_user_op
-def bf16x2_to_u32(values, *, loc=None, ip=None):
-    packed = llvm.bitcast(
-        T.i32(),
-        values.ir_value(loc=loc, ip=ip),
-        loc=loc,
-        ip=ip,
-    )
-    return Uint32(packed)
-
-
-@cute.jit
-def sanitize_negative_zero_u32(word):
-    low = Uint16(word & Uint32(0xFFFF))
-    high = Uint16(word >> Uint32(16))
-    if low == Uint16(NEG_ZERO_BF16_BITS):
-        word = word & Uint32(0xFFFF0000)
-    if high == Uint16(NEG_ZERO_BF16_BITS):
-        word = word & Uint32(0x0000FFFF)
-    return word
-
-
-@cute.jit
-def sanitize_negative_zero_u32x2(packed):
-    result = cute.make_rmem_tensor(cute.make_layout((2,)), Uint32)
-    for i in cutlass.range_constexpr(2):
-        result[i] = sanitize_negative_zero_u32(packed[i])
-    return result.load()
-
-
-@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]
-        low = Uint16(word & Uint32(0xFFFF))
-        high = Uint16(word >> Uint32(16))
-        if low == Uint16(NEG_ZERO_BF16_BITS):
-            word = word & Uint32(0xFFFF0000)
-        if high == Uint16(NEG_ZERO_BF16_BITS):
-            word = word & Uint32(0x0000FFFF)
-        result[i] = word
-    return result.load()
-
-
-@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))
-    return dirty
-
-
-@cute.jit
-def warp_sum_specialized(
-    value: Float32,
-    warp_idx: Int32,
-    lane: Int32,
-    warps: cutlass.Constexpr[int],
-    last_warp_lanes: cutlass.Constexpr[int],
-    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.
-            if cutlass.const_expr(offset in (16, 8, 4, 2, 1)):
-                other = cute.arch.shuffle_sync_bfly(
-                    value,
-                    offset=offset,
-                    mask=last_warp_mask,
-                    mask_and_clamp=31,
-                )
-                if (lane ^ Int32(offset)) < Int32(last_warp_lanes):
-                    value = value + other
-    else:
-        for offset in cutlass.range_constexpr(16, 0, -1):
-            if cutlass.const_expr(offset in (16, 8, 4, 2, 1)):
-                value = value + cute.arch.shuffle_sync_bfly(
-                    value,
-                    offset=offset,
-                    mask=-1,
-                    mask_and_clamp=31,
-                )
-    return value
diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/host.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/host.py
index 2833c3a37d13..c6efad7cc3d3 100644
--- a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/host.py
+++ b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/host.py
@@ -11,7 +11,7 @@
 import torch
 from cutlass.cute.runtime import make_fake_tensor
 
-from vllm.models.kimi_k3.nvidia.ops.cute_dsl.latent_moe_tail.primitives import to_cute
+from vllm.models.common.ops.moe_tail_fusion.primitives import to_cute
 
 _NAMES = ("BFloat16", "Float32", "Uint8", "Uint32", "Int32", "Int64")
 _DTYPES = {getattr(torch, name.lower()): getattr(cutlass, name) for name in _NAMES}
diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/small_batch/front_tp16.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/small_batch/front_tp16.py
index 7e80d430488e..292734d746da 100644
--- a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/small_batch/front_tp16.py
+++ b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/small_batch/front_tp16.py
@@ -12,7 +12,8 @@
 from cutlass import BFloat16, Float32, Int32, Uint8, Uint32, Uint64
 from cutlass.cute.runtime import make_fake_stream, make_fake_tensor
 
-from .....ops.cute_dsl.latent_moe_tail.primitives import to_cute
+from vllm.models.common.ops.moe_tail_fusion.primitives import to_cute
+
 from ...common.primitives import (
     alloc,
     asm,
diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/runtime/prefetch.py b/vllm/models/kimi_k3/nvidia/specialized_decode/runtime/prefetch.py
index 8ca3e1e5277c..d739e39bb58c 100644
--- a/vllm/models/kimi_k3/nvidia/specialized_decode/runtime/prefetch.py
+++ b/vllm/models/kimi_k3/nvidia/specialized_decode/runtime/prefetch.py
@@ -36,7 +36,8 @@
 import cuda.bindings.driver as cuda
 import torch
 
-from ...ops.cute_dsl.latent_moe_tail.primitives import to_cute
+from vllm.models.common.ops.moe_tail_fusion.primitives import to_cute
+
 from ..cutedsl.common.host import _checked, resolve_kernel
 from ..cutedsl.common.policy import batch_table
 from ..cutedsl.common.prefetch import compile_kernel
diff --git a/vllm/models/minimax_m3/amd/indexer_msa.py b/vllm/models/minimax_m3/amd/indexer_msa.py
index 8d9b9032d522..3617bb03bc4c 100644
--- a/vllm/models/minimax_m3/amd/indexer_msa.py
+++ b/vllm/models/minimax_m3/amd/indexer_msa.py
@@ -147,6 +147,9 @@ class MiniMaxM3IndexerMSAMetadata(MiniMaxM3IndexerMetadata):
     prefill_row_req_id: torch.Tensor | None = None
     # [num_prefill_tokens] int32, causal token count (position + 1) per row.
     prefill_kv_lens: torch.Tensor | None = None
+    # Decode CP's 4-row requests (block table, seq_lens): built in the forward by
+    # the step's first NVFP4 indexer layer, so a full cudagraph records it.
+    cp_requests: tuple[torch.Tensor, torch.Tensor] | None = None
 
 
 class MiniMaxM3IndexerMSAMetadataBuilder(MiniMaxM3IndexerMetadataBuilder):
diff --git a/vllm/models/minimax_m3/amd/indexer_nvfp4.py b/vllm/models/minimax_m3/amd/indexer_nvfp4.py
new file mode 100644
index 000000000000..bd46a1186b4c
--- /dev/null
+++ b/vllm/models/minimax_m3/amd/indexer_nvfp4.py
@@ -0,0 +1,257 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""MiniMax M3 lightning indexer over NVFP4 index-K pages on ROCm (gfx950).
+
+Keys score as ``e4m3_rn_satfinite(e2m1 * scale)`` with AITER's init / local
+sentinels; the top-k writes logical block ids for the NVFP4 attend, no page
+table. With ``indexer_decode_cp`` at TP4, decode steps 4/8 rows
+per request (DSpark 3/7) and at most 256 4-row requests score
+context-parallel (``ops/index_cp``).
+"""
+
+import torch
+
+from vllm.config import get_current_vllm_config
+from vllm.forward_context import get_forward_context
+from vllm.logger import init_logger
+from vllm.models.minimax_m3.amd.indexer_msa import (
+    MSA_INDEX_HEAD_DIM,
+    MSA_SCORE_TYPE,
+    MSA_SPARSE_BLOCK_SIZE,
+    MSA_TOPK_BLOCKS,
+    MiniMaxM3IndexerMSABackend,
+    MiniMaxM3IndexerMSAMetadata,
+    _score_width,
+)
+from vllm.models.minimax_m3.common.indexer import (
+    MiniMaxM3IndexerDecodeMetadata,
+    MiniMaxM3IndexerImpl,
+)
+from vllm.platforms import current_platform
+from vllm.v1.worker.ubatching import dbo_enabled
+
+logger = init_logger(__name__)
+
+# Index heads per rank (TP4 / TP2 / TP1).
+NUM_INDEX_HEADS = (1, 2, 4)
+# Decode CP: rows per request (DSpark 3); DSpark 7 splits each request in two.
+CP_QUERY_LEN = 4
+
+
+def _cp_requests(
+    d: MiniMaxM3IndexerDecodeMetadata,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """The CP scorer's 4-row requests. DSpark 7's 8-row request scores as two on
+    its block table row, the first 4 tokens shorter, so every row keeps its
+    causal length."""
+    if d.decode_query_len == CP_QUERY_LEN:
+        return d.block_table, d.seq_lens
+    first = (d.seq_lens - CP_QUERY_LEN).clamp_(min=0)  # padding stays 0
+    seq_lens = torch.stack((first, d.seq_lens), dim=1).flatten()
+    return d.block_table.repeat_interleave(2, dim=0), seq_lens
+
+
+def rocm_nvfp4_indexer_unsupported_reason(
+    *,
+    topk_blocks: int,
+    sparse_block_size: int,
+    num_index_heads: int,
+    index_head_dim: int,
+    score_type: str,
+) -> str | None:
+    """Why the ROCm NVFP4 indexer cannot serve this config, None when it can."""
+    if not current_platform.is_rocm():
+        return "needs ROCm"
+    from vllm.platforms.rocm import on_gfx950
+
+    if not on_gfx950():
+        return "needs gfx950"
+    if (topk_blocks, sparse_block_size, index_head_dim, score_type) != (
+        MSA_TOPK_BLOCKS,
+        MSA_SPARSE_BLOCK_SIZE,
+        MSA_INDEX_HEAD_DIM,
+        MSA_SCORE_TYPE,
+    ):
+        return (
+            f"needs topk_blocks={MSA_TOPK_BLOCKS}, block_size="
+            f"{MSA_SPARSE_BLOCK_SIZE}, index_head_dim={MSA_INDEX_HEAD_DIM} and "
+            f"score_type={MSA_SCORE_TYPE!r}"
+        )
+    if num_index_heads not in NUM_INDEX_HEADS:
+        return f"needs num_index_heads in {NUM_INDEX_HEADS}, got {num_index_heads}"
+    return None
+
+
+def init_rocm_nvfp4_indexer_decode_cp(max_model_len: int) -> bool:
+    """Set up decode CP if configured (a TP collective); returns whether it is on."""
+    from vllm.distributed import get_tensor_model_parallel_world_size
+
+    from .ops.index_cp import host as index_cp
+
+    if not get_current_vllm_config().attention_config.indexer_decode_cp:
+        return False
+    if get_tensor_model_parallel_world_size() != index_cp.WORLD:
+        logger.warning_once(
+            "indexer_decode_cp needs TP %d on ROCm; scoring decode locally.",
+            index_cp.WORLD,
+        )
+        return False
+    index_cp.init(
+        torch.device("cuda", torch.accelerator.current_device_index()),
+        _score_width(max_model_len, MSA_SPARSE_BLOCK_SIZE),
+    )
+    if index_cp.get() is None:
+        return False
+    logger.info_once("MiniMax M3 NVFP4 indexer: decode context parallelism on (TP4)")
+    return True
+
+
+class MiniMaxM3IndexerRocmNvfp4Impl(MiniMaxM3IndexerImpl):
+    """FlyDSL NVFP4 index scores + top-k block ids for prefill and decode."""
+
+    # Its builder materializes the prefill rows' block counts.
+    indexer_backend_cls = MiniMaxM3IndexerMSABackend
+
+    def __init__(self, **kwargs) -> None:
+        super().__init__(**kwargs)
+        if self.indexer_kv_dtype != "nvfp4":
+            raise ValueError(
+                f"{type(self).__name__} requires indexer_kv_dtype='nvfp4', got "
+                f"{self.indexer_kv_dtype!r}"
+            )
+        reason = rocm_nvfp4_indexer_unsupported_reason(
+            topk_blocks=self.topk_blocks,
+            sparse_block_size=self.block_size,
+            num_index_heads=self.num_index_heads,
+            index_head_dim=self.index_head_dim,
+            score_type=self.score_type,
+        )
+        if reason is not None:
+            raise ValueError(f"MiniMax M3 NVFP4 indexer on ROCm {reason}")
+
+    def _new_score(self, rows: int, max_seq_len: int) -> torch.Tensor:
+        """Uninitialized: the scorer writes every block the top-k reads."""
+        return torch.empty(
+            (
+                self.num_index_heads,
+                rows,
+                _score_width(max_seq_len, self.block_size),
+            ),
+            dtype=torch.float32,
+            device=self.index_cache.kv_cache.device,
+        )
+
+    def _decode_cp(self, num_decodes: int, query_len: int):
+        from .ops.index_cp import host as index_cp
+
+        cp = index_cp.get()
+        if (
+            cp is None
+            # The push has no back-pressure; under DBO no all-reduce orders it.
+            or dbo_enabled()
+            or self.num_index_heads != 1
+            or query_len not in (CP_QUERY_LEN, 2 * CP_QUERY_LEN)
+            or num_decodes * query_len > cp.max_rows
+        ):
+            return None
+        return cp
+
+    def forward(
+        self,
+        index_query: torch.Tensor,
+        **_aiter_table_args,
+    ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
+        from .ops.index_topk_decode.host import (
+            minimax_m3_index_topk_decode_indices,
+            minimax_m3_index_topk_prefill_indices,
+        )
+        from .ops.nvfp4_index_decode.host import minimax_m3_nvfp4_index_score_decode
+        from .ops.nvfp4_index_prefill.host import (
+            minimax_m3_nvfp4_index_score_prefill,
+        )
+
+        attn_metadata = get_forward_context().attn_metadata
+        if not isinstance(attn_metadata, dict):
+            return None, None  # profiling run; caches unbound
+        md = attn_metadata[self.index_cache.prefix]
+        assert isinstance(md, MiniMaxM3IndexerMSAMetadata)
+        num_tokens = md.num_actual_tokens
+        nd = md.num_decode_tokens
+        iq = index_query[:num_tokens].view(
+            -1, self.num_index_heads, self.index_head_dim
+        )
+        kv = self.index_cache.kv_cache
+        buf = self.topk_indices_buffer
+        if buf is None:
+            buf = torch.empty(
+                (self.num_index_heads, num_tokens, self.topk_blocks),
+                dtype=torch.int32,
+                device=iq.device,
+            )
+
+        decode_topk: torch.Tensor | None = None
+        prefill_topk: torch.Tensor | None = None
+        if md.num_decodes > 0:
+            d = md.decode
+            assert d is not None
+            cp = self._decode_cp(md.num_decodes, d.decode_query_len)
+            if cp is not None:
+                if md.cp_requests is None:
+                    md.cp_requests = _cp_requests(d)
+                block_table, seq_lens = md.cp_requests
+                # A view of the receive buffer, valid until the next call.
+                score = cp.score_push(
+                    iq[:nd, 0],
+                    kv,
+                    block_table,
+                    seq_lens,
+                    self.init_blocks,
+                    self.local_blocks,
+                    _score_width(d.max_seq_len, self.block_size),
+                )
+            else:
+                score = self._new_score(nd, d.max_seq_len)
+                minimax_m3_nvfp4_index_score_decode(
+                    iq[:nd],
+                    kv,
+                    score,
+                    d.block_table,
+                    d.seq_lens,
+                    init_blocks=self.init_blocks,
+                    local_blocks=self.local_blocks,
+                    query_len=d.decode_query_len,
+                    max_seq_len=d.max_seq_len,
+                )
+            decode_topk = buf[:, :nd, :]
+            minimax_m3_index_topk_decode_indices(
+                score,
+                decode_topk,
+                d.seq_lens,
+                max_seq_len=d.max_seq_len,
+                query_len=d.decode_query_len,
+            )
+
+        if md.num_prefills > 0:
+            p = md.prefill
+            assert p is not None and md.prefill_num_valid_pages is not None
+            score = self._new_score(num_tokens - nd, p.max_seq_len)
+            minimax_m3_nvfp4_index_score_prefill(
+                iq[nd:],
+                kv,
+                score,
+                p.block_table,
+                p.cu_seqlens_q,
+                p.seq_lens,
+                init_blocks=self.init_blocks,
+                local_blocks=self.local_blocks,
+                max_query_len=p.max_query_len,
+                max_seq_len=p.max_seq_len,
+            )
+            prefill_topk = buf[:, nd:num_tokens, :]
+            minimax_m3_index_topk_prefill_indices(
+                score,
+                prefill_topk,
+                max_seq_len=p.max_seq_len,
+                num_valid_pages=md.prefill_num_valid_pages,
+            )
+        return decode_topk, prefill_topk
diff --git a/vllm/models/minimax_m3/amd/model.py b/vllm/models/minimax_m3/amd/model.py
index 3222280b65ab..541e90d19197 100644
--- a/vllm/models/minimax_m3/amd/model.py
+++ b/vllm/models/minimax_m3/amd/model.py
@@ -36,6 +36,9 @@
 from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
 from vllm.forward_context import get_forward_context
 from vllm.logger import init_logger
+from vllm.model_executor.kernels.linear.mxfp8.rocm_native import (
+    RocmDotScaledMxfp8LinearKernel,
+)
 from vllm.model_executor.layers.attention import Attention
 from vllm.model_executor.layers.attention.attention import set_default_quant_scales
 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
@@ -45,10 +48,15 @@
 from vllm.model_executor.layers.fused_moe import (
     FusedMoEFactory,
     GateLinear,
+    RoutedExperts,
     fused_moe_make_expert_params_mapping,
 )
 from vllm.model_executor.layers.fused_moe.utils import (
-    is_model_fused_shared_expert_compatible,
+    resolve_layer_fused_shared_expert,
+)
+from vllm.model_executor.layers.fusion.quant_activation import (
+    QuantizedActivation,
+    get_input_quant_key,
 )
 from vllm.model_executor.layers.linear import (
     MergedColumnParallelLinear,
@@ -58,9 +66,8 @@
 )
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
 from vllm.model_executor.layers.quantization import QuantizationConfig
-from vllm.model_executor.layers.quantization.utils.config_utils import (
-    is_shared_expert_quant_fse_compatible,
-)
+from vllm.model_executor.layers.quantization.modelopt import ModelOptMxFp8Config
+from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp8Dynamic
 from vllm.model_executor.layers.rotary_embedding import get_rope
 from vllm.model_executor.layers.vocab_parallel_embedding import (
     ParallelLMHead,
@@ -81,11 +88,14 @@
     AutoWeightsLoader,
     PPMissingLayer,
     WeightsMapper,
+    extract_layer_index,
+    get_fused_shared_expert_layers,
     init_vllm_registered_model,
     is_pp_missing_parameter,
     make_empty_intermediate_tensors_factory,
     make_layers,
     maybe_prefix,
+    pp_splits_layer_sharing,
     spec_decode_needs_target_embed,
 )
 from vllm.model_executor.models.vision import run_dp_sharded_mrope_vision_model
@@ -95,8 +105,14 @@
     msa_indexer_cp_enabled,
     select_msa_indexer_impl_cls,
 )
+from vllm.models.minimax_m3.amd.indexer_nvfp4 import (
+    MiniMaxM3IndexerRocmNvfp4Impl,
+    init_rocm_nvfp4_indexer_decode_cp,
+)
+from vllm.models.minimax_m3.amd.mono.envs import envs as mono_envs
 from vllm.models.minimax_m3.amd.ops import (
     gemma_fused_add_rmsnorm,
+    gemma_fused_add_rmsnorm_mxfp8,
     gemma_rmsnorm,
     swiglu_oai_split,
 )
@@ -132,8 +148,12 @@
 )
 from vllm.models.minimax_m3.common.vision_tower import MiniMaxVLVisionModel
 from vllm.multimodal import MULTIMODAL_REGISTRY
+from vllm.platforms.rocm import on_gfx942
 from vllm.sequence import IntermediateTensors
 from vllm.utils.torch_utils import kv_cache_dtype_str_to_dtype
+from vllm.v1.attention.ops.triton_reshape_and_cache_flash import (
+    triton_reshape_and_cache_flash,
+)
 from vllm.v1.kv_cache_interface import (
     FullAttentionSpec,
     KVCacheSpec,
@@ -146,6 +166,9 @@
 
 logger = init_logger(__name__)
 
+# The fused QK-norm/RoPE/KV-insert HIP op segfaults on gfx942.
+_USE_GFX942_NATIVE_QKNORM_ROPE = on_gfx942()
+
 
 def _sparse_attention_layer_ids(config: PreTrainedConfig) -> set[int]:
     """Layer ids whose attention runs the extra sparse "index" branch."""
@@ -365,6 +388,30 @@ def __init__(
                 f"the number of experts {config.num_local_experts}."
             )
 
+        # MiniMax-M3.1 like checkpoints (`quant_method: mxfp8`, `moe_quant_algo:
+        # NVFP4`) have no native NVFP4 MoE backend on ROCm: the routed experts must be
+        # re-quantized online.
+        experts_prefix = f"{prefix}.experts"
+        if (
+            isinstance(quant_config, ModelOptMxFp8Config)
+            and quant_config.moe_quant_algo == "NVFP4"
+            and not quant_config.is_layer_excluded(experts_prefix)
+            and (
+                quant_config.online_quantization_config is None
+                or quant_config.online_quantization_config.resolve_quant_method_cls(
+                    RoutedExperts, experts_prefix
+                )
+                is None
+            )
+        ):
+            raise NotImplementedError(
+                "This MiniMax model with NVFP4 MoE experts "
+                "(`moe_quant_algo: NVFP4`, `moe_quant_format: nvfp4-pack-quantized`) "
+                "is not supported on ROCm. Please use `--quantization-config.moe "
+                "mxfp8` (or `--quantization-config.moe mxfp4`) to re-quantize the "
+                "MoE to MXFP8 (or MXFP4) when loading to serve on ROCm."
+            )
+
         self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0)
         self.n_shared_experts = getattr(config, "n_shared_experts", None)
 
@@ -396,32 +443,11 @@ def __init__(
         # MoE call as the last expert slot, so we don't build a separate module.
         # On gfx950 with aiter MoE the append is fused inside aiter's grouped
         # top-k kernel; otherwise it goes through the vLLM top-k bias router.
-        # It is disabled under expert parallelism (the shared slot is appended to
-        # the routed top-k, which the EP expert-mapping path does not handle).
-        # TODO: Historically, only `VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=1`
-        # is checked to enable FSE for MiniMax-M3, despite AITER not being used.
-        # This should be cleaned up and use `resolve_layer_fused_shared_expert`.
-        fse_requested = envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS
-
         self.is_fused_shared_expert_enabled = False
-        if (
-            fse_requested
-            and bool(getattr(config, "n_shared_experts", None))
-            and not get_current_vllm_config().parallel_config.enable_expert_parallel
-        ):
-            fse_compatible, fse_reason = is_shared_expert_quant_fse_compatible(
-                quant_config,
-                f"{prefix}.experts",
-                f"{prefix}.shared_experts",
+        if bool(getattr(config, "n_shared_experts", None)):
+            self.is_fused_shared_expert_enabled = resolve_layer_fused_shared_expert(
+                quant_config, prefix, disable_under_ep=True
             )
-            if not fse_compatible:
-                logger.warning(
-                    "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS is enabled but "
-                    "cannot be enabled: %s.",
-                    fse_reason,
-                )
-            else:
-                self.is_fused_shared_expert_enabled = True
 
         # When additionally on gfx950 with an active aiter MoE backend, the shared
         # expert is appended inside aiter's an active aiter MoE backend, the shared
@@ -569,9 +595,22 @@ def __init__(
     def forward(
         self,
         positions: torch.Tensor,
-        hidden_states: torch.Tensor,
+        hidden_states: torch.Tensor | QuantizedActivation,
     ) -> torch.Tensor:
         qkv, _ = self.qkv_proj(hidden_states)
+        if _USE_GFX942_NATIVE_QKNORM_ROPE:
+            num_tokens = qkv.shape[0]
+            q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+            q = self.q_norm(q.view(num_tokens, self.num_heads, self.head_dim))
+            k = self.k_norm(k.view(num_tokens, self.num_kv_heads, self.head_dim))
+            q, k = self.rotary_emb.forward_native(positions, q, k)
+            assert k is not None
+            q = q.reshape(num_tokens, self.q_size)
+            k = k.reshape(num_tokens, self.kv_size)
+            attn_output = self.attn(q, k, v)
+            output, _ = self.o_proj(attn_output)
+            return output
+
         # Fused per-head Gemma QK-norm + partial NeoX RoPE on q/k, in place (dense
         # mode: no index branch, no KV-cache insert). Matches nvidia/model.py and
         # replaces the unfused split -> q_norm/k_norm -> rotary_emb chain; verified
@@ -748,6 +787,9 @@ def __init__(
         # Shared top-k buffer: the indexer writes the selected blocks into it and
         # the attend impl reads them back (no Python value crosses the break).
         self.topk_indices_buffer = topk_indices_buffer
+        # One fused kernel inserts both caches, so NVFP4 K/V needs an NVFP4
+        # index cache.
+        self.use_nvfp4_kv = get_kv_quant_mode(self.kv_cache_dtype).is_nvfp4
         # Indexer and main attention are separate impls. This indexer is
         # ROCm-only, so the model selects it rather than the neutral
         # MiniMaxM3Indexer, which knows nothing about it and would pick its own.
@@ -778,8 +820,10 @@ def __init__(
             topk_blocks=sparse_cfg["sparse_topk_blocks"],
             sparse_block_size=sparse_cfg["sparse_block_size"],
         )
-        self.use_aiter_sparse_pa = minimax_m3_use_aiter_sparse_pa(
-            self.num_kv_heads, emits_sparse_block_table=self.indexer_emits_table
+        self.use_aiter_sparse_pa = not self.use_nvfp4_kv and (
+            minimax_m3_use_aiter_sparse_pa(
+                self.num_kv_heads, emits_sparse_block_table=self.indexer_emits_table
+            )
         )
         self.kv_cache_k = torch.tensor([])
         self.kv_cache_v = torch.tensor([])
@@ -806,7 +850,11 @@ def __init__(
         self.sparse_ctx_buffer = sparse_ctx_buffer
         # Same condition the forward dispatches on, so the indexer built here
         # and the branch that calls it cannot disagree about which one this is.
-        if msa_indexer_impl_cls is None:
+        if self.use_nvfp4_kv:
+            self.indexer = MiniMaxM3MSAIndexer(
+                impl_cls=MiniMaxM3IndexerRocmNvfp4Impl, **indexer_kwargs
+            )
+        elif msa_indexer_impl_cls is None:
             # Self-contained nn.Module: owns its side cache, selects its impl.
             self.indexer = MiniMaxM3Indexer(**indexer_kwargs)
         else:
@@ -977,10 +1025,122 @@ def _insert_aiter_sparse_pa_kv(
             return
         minimax_m3_insert_index_cache(index_k, index_cache, index_slot_mapping)
 
+    def _native_gfx942_qknorm_rope_cache(
+        self,
+        qkv: torch.Tensor,
+        positions: torch.Tensor,
+        fwd_slot_mapping: dict[str, torch.Tensor],
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        """Run QK norm, RoPE, and cache insertion without the fused HIP op."""
+        from vllm.models.minimax_m3.amd.ops.sparse_pa import (
+            minimax_m3_insert_index_cache,
+        )
+
+        num_tokens = qkv.shape[0]
+        q, k, v, index_q_raw, index_k_raw = qkv.split(
+            [
+                self.q_size,
+                self.kv_size,
+                self.kv_size,
+                self.index_q_size,
+                self.idx_head_dim,
+            ],
+            dim=-1,
+        )
+
+        q = self.q_norm(q.view(num_tokens, self.num_heads, self.head_dim))
+        k = self.k_norm(k.view(num_tokens, self.num_kv_heads, self.head_dim))
+        q, k = self.rotary_emb.forward_native(positions, q, k)
+        assert k is not None
+        q = q.reshape(num_tokens, self.q_size)
+
+        main_slot_mapping = fwd_slot_mapping[self.layer_name]
+        if self.kv_cache.numel() != 0:
+            # [block, head, token, 2 * dim] -> [block, token, head, dim].
+            key_cache, value_cache = self.kv_cache.transpose(1, 2).split(
+                self.head_dim, dim=-1
+            )
+            triton_reshape_and_cache_flash(
+                k,
+                v.view(num_tokens, self.num_kv_heads, self.head_dim),
+                key_cache,
+                value_cache,
+                main_slot_mapping,
+                self.kv_cache_dtype,
+                self._k_scale,
+                self._v_scale,
+            )
+
+        if self.skip_index_topk:
+            return q, None
+
+        index_q = self.index_q_norm(
+            index_q_raw.view(num_tokens, self.num_idx_heads, self.idx_head_dim)
+        )
+        index_k = self.index_k_norm(index_k_raw.view(num_tokens, self.idx_head_dim))
+        index_q, index_k = self.index_rotary_emb.forward_native(
+            positions, index_q, index_k
+        )
+        assert index_k is not None
+        index_q = index_q.reshape(num_tokens, self.index_q_size).to(
+            self.indexer.index_cache.dtype
+        )
+
+        index_cache = self.indexer.index_cache.kv_cache
+        if index_cache.numel() != 0:
+            index_slot_mapping = fwd_slot_mapping[self.indexer.index_cache.prefix]
+            minimax_m3_insert_index_cache(
+                index_k.reshape(num_tokens, self.idx_head_dim),
+                index_cache,
+                index_slot_mapping,
+            )
+        return q, index_q
+
+    def _norm_rope_nvfp4_insert(
+        self,
+        qkv: torch.Tensor,
+        positions: torch.Tensor,
+        slot_mapping: torch.Tensor,
+        index_slot_mapping: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Fused norm + RoPE + NVFP4 K/V and index-K insert; returns q, index-q."""
+        from vllm.models.minimax_m3.amd.ops.norm_rope_nvfp4_insert.host import (
+            minimax_m3_norm_rope_nvfp4_insert,
+        )
+
+        num_tokens = qkv.shape[0]
+        q = qkv.new_empty((num_tokens, self.q_size))
+        index_q = qkv.new_empty(
+            (num_tokens, self.index_q_size),
+            dtype=self.indexer.index_cache.query_dtype,
+        )
+        minimax_m3_norm_rope_nvfp4_insert(
+            qkv,
+            self.q_norm.weight,
+            self.k_norm.weight,
+            self.index_q_norm.weight,
+            self.index_k_norm.weight,
+            self.rotary_emb.cos_sin_cache,
+            positions,
+            slot_mapping,
+            index_slot_mapping,
+            self.kv_cache,
+            self.indexer.index_cache.kv_cache,
+            self._k_scale,
+            self._v_scale,
+            q,
+            index_q,
+            num_heads=self.num_heads,
+            num_kv_heads=self.num_kv_heads,
+            num_index_heads=self.num_idx_heads,
+            eps=self.q_norm.variance_epsilon,
+        )
+        return q, index_q
+
     def forward(
         self,
         positions: torch.Tensor,
-        hidden_states: torch.Tensor,
+        hidden_states: torch.Tensor | QuantizedActivation,
     ) -> torch.Tensor:
         # Single fused projection emitting [q | k | v | index_q | index_k].
         qkv, _ = self.qkv_proj(hidden_states)
@@ -1006,7 +1166,29 @@ def forward(
             # Memory-profiling run: caches not yet bound, slot_mapping is empty.
             return qkv.new_zeros((num_tokens, self.hidden_size))
 
+        if _USE_GFX942_NATIVE_QKNORM_ROPE:
+            q, index_q = self._native_gfx942_qknorm_rope_cache(
+                qkv, positions, fwd_slot_mapping
+            )
+            output = torch.empty_like(q)
+            attn_output = self._run_attention(q, index_q, output)
+            output, _ = self.o_proj(attn_output)
+            return output
+
         main_slot_mapping = fwd_slot_mapping[self.layer_name]
+        if self.use_nvfp4_kv:
+            q, index_q = self._norm_rope_nvfp4_insert(
+                qkv,
+                positions,
+                main_slot_mapping,
+                fwd_slot_mapping[self.indexer.index_cache.prefix],
+            )
+            output = torch.empty_like(q)
+            attn_output = self._run_attention(
+                q, None if self.skip_index_topk else index_q, output
+            )
+            output, _ = self.o_proj(attn_output)
+            return output
         q = qkv.new_empty((num_tokens, self.q_size))
         key_cache = value_cache = page16_slot_mapping = None
         k_scale = v_scale = None
@@ -1283,6 +1465,24 @@ def _run_attention(
         return self.impl.forward(self, query, self.kv_cache, output)
 
 
+def _can_fuse_input_norm_mxfp8(qkv_proj: nn.Module) -> bool:
+    """Whether ``input_layernorm`` can hand ``qkv_proj`` an MXFP8 activation.
+
+    The quant key does not encode the scale layout (FlashInfer wants swizzled
+    scales under the same key), so this also requires the ROCm ``dot_scaled``
+    kernel, which reads the row-major scales the fused norm writes.
+    """
+    if envs.VLLM_BATCH_INVARIANT or get_input_quant_key(qkv_proj) != kMxfp8Dynamic:
+        return False
+    # Quark keeps the kernel on the layer's scheme, ModelOpt on the method.
+    kernel = getattr(getattr(qkv_proj, "scheme", None), "ocp_mx_linear", None)
+    if kernel is None:
+        kernel = getattr(getattr(qkv_proj, "quant_method", None), "kernel", None)
+    return isinstance(
+        kernel, RocmDotScaledMxfp8LinearKernel
+    ) and kernel.supports_quantized_input(qkv_proj)
+
+
 class MiniMaxM3DecoderLayer(nn.Module):
     def __init__(
         self,
@@ -1356,6 +1556,9 @@ def __init__(
         self.post_attention_layernorm = MiniMAXGemmaRMSNorm(
             config.hidden_size, eps=config.rms_norm_eps
         )
+        self.fuse_input_norm_mxfp8 = _can_fuse_input_norm_mxfp8(self.self_attn.qkv_proj)
+        if self.fuse_input_norm_mxfp8:
+            logger.info_once("Fusing input_layernorm with qkv_proj's MXFP8 input quant")
 
     def forward(
         self,
@@ -1364,14 +1567,25 @@ def forward(
         residual: torch.Tensor | None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         # Self Attention
+        attn_input: torch.Tensor | QuantizedActivation
         if residual is None:
             residual = hidden_states
-            hidden_states = self.input_layernorm(hidden_states)
+            attn_input = self.input_layernorm(hidden_states)
+        elif self.fuse_input_norm_mxfp8:
+            x_q, x_scale, residual = gemma_fused_add_rmsnorm_mxfp8(
+                hidden_states,
+                residual,
+                self.input_layernorm.weight,
+                self.input_layernorm.variance_epsilon,
+            )
+            attn_input = QuantizedActivation(
+                x_q, x_scale, hidden_states.dtype, hidden_states.shape, kMxfp8Dynamic
+            )
         else:
-            hidden_states, residual = self.input_layernorm(hidden_states, residual)
+            attn_input, residual = self.input_layernorm(hidden_states, residual)
         hidden_states = self.self_attn(
             positions=positions,
-            hidden_states=hidden_states,
+            hidden_states=attn_input,
         )
 
         hidden_states, residual = fused_allreduce_gemma_rms_norm(
@@ -1384,6 +1598,7 @@ def forward(
 
 class MiniMaxM3Model(nn.Module, EagleModelMixin):
     fall_back_to_pt_during_load = False
+    supports_aux_hidden_states_over_pp = True
 
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()
@@ -1397,9 +1612,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
 
         self.vocab_size = config.vocab_size
 
-        if get_pp_group().is_first_rank or spec_decode_needs_target_embed(
-            vllm_config, include_mtp=vllm_config.use_v2_model_runner
-        ):
+        if get_pp_group().is_first_rank or spec_decode_needs_target_embed(vllm_config):
             self.embed_tokens = VocabParallelEmbedding(
                 config.vocab_size,
                 config.hidden_size,
@@ -1432,6 +1645,9 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
             indexer_kv_dtype = vllm_config.attention_config.resolve_indexer_kv_dtype(
                 "bf16"
             )
+            use_nvfp4_kv = get_kv_quant_mode(
+                vllm_config.cache_config.cache_dtype
+            ).is_nvfp4
             # The only MSA indexer probe. It has to live here rather than in
             # the layer because the buffers it gates are one allocation shared
             # by all of them, so the answer is needed before any layer exists,
@@ -1441,13 +1657,18 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
             # this indexer and there is no such choice until it exists; the
             # probe owes nothing to CP in return, since what it checks is the
             # platform and the kernels' shape contract, neither of which a
-            # head count enters.
-            impl_cls = select_msa_indexer_impl_cls(
-                topk_blocks=sparse_cfg["sparse_topk_blocks"],
-                sparse_block_size=sparse_cfg["sparse_block_size"],
-                index_head_dim=sparse_cfg["sparse_index_dim"],
-                indexer_kv_dtype=indexer_kv_dtype,
-                score_type=sparse_cfg.get("sparse_score_type", "max"),
+            # head count enters. NVFP4 K/V runs its own indexer, which emits
+            # no page table.
+            impl_cls = (
+                None
+                if use_nvfp4_kv
+                else select_msa_indexer_impl_cls(
+                    topk_blocks=sparse_cfg["sparse_topk_blocks"],
+                    sparse_block_size=sparse_cfg["sparse_block_size"],
+                    index_head_dim=sparse_cfg["sparse_index_dim"],
+                    indexer_kv_dtype=indexer_kv_dtype,
+                    score_type=sparse_cfg.get("sparse_score_type", "max"),
+                )
             )
             if impl_cls is not None and minimax_m3_use_aiter_sparse_pa(
                 num_kv_heads, emits_sparse_block_table=True
@@ -1479,8 +1700,24 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
                 vllm_config,
                 msa_indexer_selected=self.msa_indexer_impl_cls is not None,
             )
+            if use_nvfp4_kv:
+                # A TP collective (IPC buffers), outside graph capture.
+                init_rocm_nvfp4_indexer_decode_cp(
+                    vllm_config.model_config.max_model_len
+                )
         else:
             self.topk_indices_buffer = None
+        sparse_layers = _sparse_attention_layer_ids(config)
+        skip_layers = {i for i in sparse_layers if _should_skip_index_topk(config, i)}
+        if pp_splits_layer_sharing(
+            config.num_hidden_layers, sparse_layers, skip_layers
+        ):
+            # The top-k is sharded over TP index heads, so PP cannot send it.
+            raise NotImplementedError(
+                "use_index_cache with a PP stage that starts on a top-k sharing "
+                "layer is not supported. Set VLLM_PP_LAYER_PARTITION so every "
+                "stage starts on a top-k computing layer."
+            )
 
         self.start_layer, self.end_layer, self.layers = make_layers(
             config.num_hidden_layers,
@@ -1497,10 +1734,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
             ),
             prefix=f"{prefix}.layers",
         )
-        self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible(
-            self.layers,
+        self.fused_shared_expert_layers = get_fused_shared_expert_layers(
+            self,
             MiniMaxM3MoE,
-            "block_sparse_moe",
+            enable_eplb=vllm_config.parallel_config.enable_eplb,
         )
 
         if get_pp_group().is_last_rank:
@@ -1510,6 +1747,13 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
             ["hidden_states", "residual"], config.hidden_size
         )
+        # VLLM_ROCM_MINIMAX_M3_MONO=1: supported decode steps run one FlyDSL
+        # launch per layer (mono kernels)
+        self._mono = None
+        if mono_envs.MONO_ENABLE:
+            from vllm.models.minimax_m3.amd.mono.dispatch import MonoDecode
+
+            self._mono = MonoDecode(self, vllm_config)
 
     def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
         return self.embed_tokens(input_ids)
@@ -1529,6 +1773,10 @@ def forward(
         intermediate_tensors: IntermediateTensors | None,
         inputs_embeds: torch.Tensor | None = None,
     ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]:
+        if self._mono is not None and self._mono.supports(
+            input_ids, positions, intermediate_tensors, inputs_embeds
+        ):
+            return self._mono.forward(input_ids, positions, inputs_embeds)
         if get_pp_group().is_first_rank:
             if inputs_embeds is not None:
                 hidden_states = inputs_embeds
@@ -1541,10 +1789,15 @@ def forward(
             residual = intermediate_tensors["residual"]
 
         decode = self.specialized_decode
-        # EAGLE3 is not yet compatible with pipeline parallel
-        aux_hidden_states = self._maybe_add_hidden_state([], 0, hidden_states, residual)
+        remote_aux = self.collect_remote_aux_hidden_states(intermediate_tensors)
+        aux_hidden_states: list[torch.Tensor] = []
+        if get_pp_group().is_first_rank:
+            self._maybe_add_hidden_state(
+                aux_hidden_states, self.start_layer, hidden_states, residual
+            )
+        layers = self.layers[self.start_layer : self.end_layer]
         if decode is None:
-            for idx, layer in enumerate(self.layers[self.start_layer : self.end_layer]):
+            for idx, layer in enumerate(layers, start=self.start_layer):
                 hidden_states, residual = layer(positions, hidden_states, residual)
                 self._maybe_add_hidden_state(
                     aux_hidden_states, idx + 1, hidden_states, residual
@@ -1553,8 +1806,7 @@ def forward(
             step = decode.plan(positions, intermediate_tensors)
             failure: Exception | None = None
             try:
-                layers = self.layers[self.start_layer : self.end_layer]
-                for idx, layer in enumerate(layers):
+                for idx, layer in enumerate(layers, start=self.start_layer):
                     kernel = step.layer(layer.layer_id) if step else None
                     hidden_states, residual = (kernel or layer)(
                         positions, hidden_states, residual
@@ -1568,11 +1820,16 @@ def forward(
 
         if not get_pp_group().is_last_rank:
             return IntermediateTensors(
-                {"hidden_states": hidden_states, "residual": residual}
+                {
+                    "hidden_states": hidden_states,
+                    "residual": residual,
+                    **self.pack_local_aux_hidden_states(aux_hidden_states),
+                }
             )
 
         hidden_states, _ = self.norm(hidden_states, residual)
 
+        aux_hidden_states = remote_aux + aux_hidden_states
         if len(aux_hidden_states) > 0:
             return hidden_states, aux_hidden_states
         return hidden_states
@@ -1581,8 +1838,9 @@ def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
         # Checkpoint experts use w1=gate, w2=down, w3=up. When fusing the shared
         # expert, include the appended slot (id == num_local_experts).
         n_shared = getattr(self.config, "n_shared_experts", 0) or 0
+        # NOTE: global despite per-layer FSE; non-fused layers skip the extras.
         num_experts = self.config.num_local_experts + (
-            n_shared if self.is_fused_shared_expert_enabled else 0
+            n_shared if self.fused_shared_expert_layers else 0
         )
         return fused_moe_make_expert_params_mapping(
             self,
@@ -1631,7 +1889,12 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
             # down->w2) so it loads via the routed expert loader. Runs before the
             # stacked/dense mappings so shared_experts.gate_proj/up_proj are not
             # captured by the dense gate_up_proj mapping.
-            if self.is_fused_shared_expert_enabled and ".shared_experts." in name:
+            # Only parse the layer index when FSE is enabled.
+            if (
+                ".shared_experts." in name
+                and bool(rocm_aiter_ops.is_fusion_moe_shared_experts_enabled())
+                and extract_layer_index(name) in self.fused_shared_expert_layers
+            ):
                 sid = self.config.num_local_experts
                 name = name.replace(".shared_experts.gate_proj.", f".experts.{sid}.w1.")
                 name = name.replace(".shared_experts.up_proj.", f".experts.{sid}.w3.")
diff --git a/vllm/models/minimax_m3/amd/mono/__init__.py b/vllm/models/minimax_m3/amd/mono/__init__.py
new file mode 100644
index 000000000000..7a79e4b7c0da
--- /dev/null
+++ b/vllm/models/minimax_m3/amd/mono/__init__.py
@@ -0,0 +1,4 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""MiniMax-M3 fused per-layer decode of up to ``config.MAX_TOKENS`` tokens; entered
+through ``dispatch.MonoDecode`` (``VLLM_ROCM_MINIMAX_M3_MONO=1``)."""
diff --git a/vllm/models/minimax_m3/amd/mono/check.py b/vllm/models/minimax_m3/amd/mono/check.py
new file mode 100644
index 000000000000..bef5dbe7d736
--- /dev/null
+++ b/vllm/models/minimax_m3/amd/mono/check.py
@@ -0,0 +1,68 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Debug tooling for the mono path.
+
+``forward_checked`` (``VLLM_ROCM_MINIMAX_M3_MONO_CHECK=1``, meant for
+``--enforce-eager``) runs every sparse layer twice on the same (ar, res): first
+the mono layer kernel, then the original decoder layer. Both insert the same cache
+slots (the kernel writes the same bytes up to the re-quantized attention weights),
+so running both is side-effect free. Rank 0 logs, per layer and token, the
+difference of the residual stream and of the reduced MoE output.
+"""
+
+import torch
+
+from vllm.logger import init_logger
+from vllm.models.minimax_m3.amd.mono.config import HIDDEN
+
+logger = init_logger(__name__)
+
+
+def _diff(a: torch.Tensor, b: torch.Tensor) -> str:
+    a, b = a.float().flatten(), b.float().flatten()
+    d = (a - b).abs()
+    rel = (d.norm() / b.norm().clamp(min=1e-30)).item()
+    return f"max {d.max().item():.3e} rel {rel:.3e}"
+
+
+def forward_checked(
+    runner,
+    input_ids: torch.Tensor | None,
+    positions: torch.Tensor,
+    inputs_embeds: torch.Tensor | None,
+) -> torch.Tensor | tuple[torch.Tensor, list[torch.Tensor]]:
+    """``runner.forward`` with the original decoder layer run beside every mono
+    layer; the mono chain's output is returned."""
+    aux_hidden_states: list[torch.Tensor] | None = (
+        [] if runner.model.aux_hidden_state_layers else None
+    )
+    res = runner.start(input_ids, positions, inputs_embeds, aux_hidden_states)
+    n = res.shape[0]
+    tables = runner.step_tables(n)
+    ars = runner.ars
+    for i, lw in enumerate(runner.sparse):
+        layer = runner.model.layers[lw.layer_id]
+        h_in = ars[i % 2][:n].clone()
+        res_in = res.clone()
+        res = runner.run_sparse_layer(i, lw, tables, positions, res)
+        ref_h, ref_res = layer(positions, h_in, res_in)
+        if aux_hidden_states is not None:
+            runner.model._maybe_add_hidden_state(
+                aux_hidden_states,
+                lw.layer_id + 1,
+                ars[(i + 1) % 2][:n],
+                res,
+            )
+        if runner.rank == 0:
+            for k in range(n):
+                logger.info(
+                    "mono check layer %d token %d/%d pos %d seq_len %d: residual %s | "
+                    "reduced out %s",
+                    lw.layer_id, k, n, int(positions[k]), int(tables.seq_lens[k]),
+                    _diff(res[k], ref_res.view(n, HIDDEN)[k]),
+                    _diff(ars[(i + 1) % 2][k], ref_h.view(n, HIDDEN)[k]),
+                )  # fmt: skip
+    hidden = runner.run_final_norm(res)
+    if aux_hidden_states is not None:
+        return hidden, aux_hidden_states
+    return hidden
diff --git a/vllm/models/minimax_m3/amd/mono/config.py b/vllm/models/minimax_m3/amd/mono/config.py
new file mode 100644
index 000000000000..9ed7777f3cef
--- /dev/null
+++ b/vllm/models/minimax_m3/amd/mono/config.py
@@ -0,0 +1,108 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+# SPDX-License-Identifier: MIT
+# Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved.
+"""Shapes the mono kernels are compiled for: one MiniMax-M3 TP4 rank.
+
+Every value here is checked against the loaded model in ``weights.py``; a
+model that does not match is never routed to the mono path.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+# One CTA per MI355X CU; the kernels rely on every CTA being co-resident.
+BLOCKS = 256
+THREADS = 512
+WAVES = THREADS // 64
+
+HIDDEN = 6144
+HEAD_DIM = 128
+ROTARY_DIM = 64
+LOCAL_Q_HEADS = 16
+
+
+def qkv_rows(idx_heads: int) -> int:
+    """Rows of a rank's fused projection, q | k | v | index_q | index_k: 16 q
+    heads, one k and v head, ``idx_heads`` index q heads (the rank's one, or every
+    one under indexer context parallelism) and the index k head."""
+    return (LOCAL_Q_HEADS + 3 + idx_heads) * HEAD_DIM
+
+
+O_K = LOCAL_Q_HEADS * HEAD_DIM
+
+N_ROUTED = 128
+TOP_K = 4
+SHARED_EXPERT = N_ROUTED  # the fused shared expert's id
+MOE_SLOTS = TOP_K + 1
+INTER = 768  # expert intermediate per rank
+
+SPARSE_BLOCK = 128
+TOPK_BLOCKS = 16
+MAX_SPARSE_KEYS = SPARSE_BLOCK * TOPK_BLOCKS
+# the longest context served: the score region spans it
+MAX_CONTEXT = 1 << 20
+MAX_INDEX_BLOCKS = MAX_CONTEXT // SPARSE_BLOCK
+PAGE16 = 16
+
+MAX_TOKENS = 16  # tokens one mono step serves (the MFMA B operand holds 16)
+TP = 4
+# indexer context parallelism: a rank computes every index q head
+MAX_QKV_ROWS = qkv_rows(TP)
+# indexer context parallelism serves a step's requests past this many index
+# blocks: its selection costs a fixed ~7 us (S = 1: ~43.5 us/layer at any length)
+# that the one-head path's O(n^2) ranking passes between 258 and 297 blocks
+# (S = 1 and S = 16 alike; logs/cp_threshold*.log)
+INDEX_CP_FROM_BLOCKS = 288
+
+
+@dataclass(frozen=True)
+class IndexHeads:
+    """The index q heads in a rank's fused projection, a build parameter:
+    ``count`` of them (1, or all TP in head order under indexer context
+    parallelism) and ``own``, the one this rank's selection scores."""
+
+    count: int = 1
+    own: int = 0
+
+    def __post_init__(self):
+        assert self.count in (1, TP) and 0 <= self.own < self.count
+
+    @property
+    def rows(self) -> int:
+        return qkv_rows(self.count)
+
+    @property
+    def iq_off(self) -> int:
+        return (LOCAL_Q_HEADS + 2 + self.own) * HEAD_DIM
+
+    @property
+    def ik_off(self) -> int:
+        return (LOCAL_Q_HEADS + 2 + self.count) * HEAD_DIM
+
+
+# without indexer context parallelism: the rank's own index q head only
+ONE_INDEX_HEAD = IndexHeads()
+
+LAYER_SLOTS = 128  # mailbox epochs: step * LAYER_SLOTS + layer + 1
+
+# host K/V cache geometry, set by the runner before the kernels are built: the
+# page-16 ids one sparse block spans (8: K and V as separate planes; 16: a block
+# interleaves K|V and V is reached from the same id at +8 pages, as vLLM lays
+# it out) and
+# whether fp8 K/V carry the host's fixed unit scale instead of per-token scales
+KV_PAGE_STRIDE = SPARSE_BLOCK // PAGE16
+STATIC_KV_SCALE = False
+
+
+def configure(*, kv_page_stride: int, static_kv_scale: bool) -> None:
+    global KV_PAGE_STRIDE, STATIC_KV_SCALE
+    assert kv_page_stride in (SPARSE_BLOCK // PAGE16, 2 * SPARSE_BLOCK // PAGE16)
+    KV_PAGE_STRIDE = kv_page_stride
+    STATIC_KV_SCALE = static_kv_scale
+
+
+class MonoUnsupported(Exception):
+    """The loaded model or runtime configuration is outside what mono serves."""
diff --git a/vllm/models/minimax_m3/amd/mono/dispatch.py b/vllm/models/minimax_m3/amd/mono/dispatch.py
new file mode 100644
index 000000000000..386f9d023bc9
--- /dev/null
+++ b/vllm/models/minimax_m3/amd/mono/dispatch.py
@@ -0,0 +1,216 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Which MiniMax-M3 forwards take the mono path (``VLLM_ROCM_MINIMAX_M3_MONO=1``).
+
+The deployment is checked once (``_config_refusal``); a step is routed to the
+kernels only when it is a pure decode of 1..MAX_TOKENS tokens whose tables the
+kernels can read (``_step_refusal``). Every per-step condition is decided from
+TP-replicated state: the kernels all-reduce in-kernel, so ranks that disagreed
+would wait on each other forever.
+
+The runner is built on the first eligible step, which under cudagraphs is the
+warmup run before the first capture (its allocations and IPC handshake cannot be
+captured), and rebuilt when vLLM re-binds the KV cache after its memory-profiling
+capture round.
+"""
+
+from typing import Any, cast
+
+import torch
+
+from vllm.config import VllmConfig
+from vllm.forward_context import ForwardContext, get_forward_context
+from vllm.logger import init_logger
+from vllm.models.minimax_m3.amd.mono.check import forward_checked
+from vllm.models.minimax_m3.amd.mono.config import (
+    MAX_CONTEXT,
+    MAX_TOKENS,
+    SPARSE_BLOCK,
+    TP,
+    MonoUnsupported,
+)
+from vllm.models.minimax_m3.amd.mono.envs import envs
+from vllm.models.minimax_m3.amd.mono.runner import MonoDecodeRunner, dense_prefix
+
+logger = init_logger(__name__)
+
+
+def _dense_prefix_ok(layers) -> bool:
+    try:
+        dense_prefix(layers)
+    except MonoUnsupported:
+        return False
+    return True
+
+
+def _separate_shared_ok(layers) -> bool:
+    """Every MoE layer keeps its shared expert as a separate MLP (the kernels then
+    run it from that MLP's MXFP8 weights, ``weights.mxfp8_a_tiles``)."""
+    moes = [lyr.block_sparse_moe for lyr in layers if lyr.is_moe_layer]
+    return bool(moes) and all(
+        not m.is_fused_shared_expert_enabled
+        and getattr(m, "shared_experts", None) is not None
+        for m in moes
+    )
+
+
+def _config_refusal(vllm_config: VllmConfig, model) -> str | None:
+    """Why this deployment cannot use the kernels at all, or None."""
+    pc = vllm_config.parallel_config
+    cc = vllm_config.cache_config
+    cg = vllm_config.compilation_config.cudagraph_mode
+    max_len = vllm_config.model_config.max_model_len
+    index_dtype = vllm_config.attention_config.resolve_indexer_kv_dtype("bf16")
+    checks = (
+        (pc.tensor_parallel_size == TP, f"TP {pc.tensor_parallel_size} != {TP}"),
+        (pc.data_parallel_size == 1, "DP > 1"),
+        (pc.pipeline_parallel_size == 1, "PP > 1"),
+        (not pc.enable_expert_parallel, "EP"),
+        (str(cc.cache_dtype).startswith("fp8"), f"kv cache {cc.cache_dtype}"),
+        (index_dtype == "fp8", f"index cache {index_dtype}"),
+        (cc.block_size == SPARSE_BLOCK, f"block size {cc.block_size}"),
+        (max_len <= MAX_CONTEXT, f"max_model_len {max_len} > {MAX_CONTEXT}"),
+        (cg is None or not cg.requires_piecewise_compilation(), "piecewise cudagraph"),
+        (
+            getattr(model, "is_fused_shared_expert_enabled", False)
+            or _separate_shared_ok(model.layers),
+            "shared expert neither fused nor a separate MLP",
+        ),
+        (_dense_prefix_ok(model.layers), "layer mix"),
+    )
+    for ok, why in checks:
+        if not ok:
+            return why
+    return None
+
+
+def _step_refusal(
+    fwd: ForwardContext, layer_name: str, n: int, positions
+) -> str | None:
+    """Why this step cannot run on the kernels, or None when it can."""
+    if not 1 <= n <= MAX_TOKENS:
+        return f"{n} tokens"
+    if positions.dtype != torch.int64:  # read as an int64's low word
+        return f"positions {positions.dtype}"
+    if fwd.ubatch_slices is not None:
+        return "microbatched"
+    if not isinstance(fwd.attn_metadata, dict) or not isinstance(
+        fwd.slot_mapping, dict
+    ):
+        return "no per-layer metadata"
+    md = cast(Any, fwd.attn_metadata.get(layer_name))
+    if md is None or md.decode is None:
+        return "no decode metadata"
+    if md.num_prefills:
+        return "prefill in the step"
+    q = md.decode.decode_query_len
+    if q < 1 or n % q:
+        return f"{n} tokens for {q} per request"
+    bt, seq_lens = md.decode.block_table, md.decode.seq_lens
+    if bt.dtype != torch.int32 or seq_lens.dtype != torch.int32:
+        return "tables not int32"
+    if bt.shape[0] < n // q or seq_lens.shape[0] < n // q:
+        return "tables shorter than the batch"
+    if bt.stride(0) != bt.shape[1] or bt.stride(1) != 1:  # rows bt_width apart
+        return "block table not row-contiguous"
+    slots = fwd.slot_mapping.get(layer_name)
+    if slots is None or slots.dtype != torch.int64 or slots.numel() < n:
+        return "slot mapping"
+    return None
+
+
+class MonoDecode:
+    """Routes supported decode steps of one ``MiniMaxM3Model`` to the kernels."""
+
+    def __init__(self, model, vllm_config: VllmConfig) -> None:
+        self._model = model
+        self._config = vllm_config.model_config.hf_text_config
+        self._runner: MonoDecodeRunner | None = None
+        self._enabled = False
+        self._slots_verified = False
+        if not envs.MONO_ENABLE:
+            return
+        why = _config_refusal(vllm_config, model)
+        if why is not None:
+            logger.info("MiniMax-M3 mono decode off: %s", why)
+            return
+        attn0 = model.layers[dense_prefix(model.layers)].self_attn
+        self._layer_name = attn0.layer_name
+        self._index_prefix = attn0.indexer.index_cache.prefix
+        self._enabled = True
+
+    def _disable(self, why: str) -> None:
+        logger.warning("MiniMax-M3 mono decode off: %s", why)
+        self._enabled = False
+
+    def _ensure_runner(self) -> bool:
+        """Build (or rebuild after a KV cache re-bind) the runner outside capture."""
+        if self._runner is not None and self._runner.stale():
+            if torch.cuda.is_current_stream_capturing():
+                return False
+            self._runner.close()
+            self._runner = None
+            self._slots_verified = False
+            logger.info("MiniMax-M3 mono decode: KV cache re-bound, rebuilding")
+        if self._runner is not None:
+            return True
+        if torch.cuda.is_current_stream_capturing():
+            return False
+        try:
+            self._runner = MonoDecodeRunner(self._model, self._config)
+        except MonoUnsupported as why:
+            self._disable(str(why))
+            return False
+        except Exception:
+            # TP-uniform: every rank fails the same way, so all fall back
+            logger.exception("MiniMax-M3 mono decode: runner build failed")
+            self._enabled = False
+            return False
+        if self._runner.shared_fp8:
+            copy = sum(lw.shared_bytes() for lw in self._runner.sparse)
+            how = f"separate MXFP8 ({copy / 2**30:.2f} GiB tiled copy per rank)"
+        else:
+            how = "fused into the MXFP4 expert tables"
+        logger.info("MiniMax-M3 mono decode on: shared expert %s", how)
+        return True
+
+    def _slots_shared(self, fwd: ForwardContext, n: int) -> bool:
+        """The kernels insert K/V and the index key through one page-128 slot
+        mapping, so the index cache must share the main cache's block ids.
+        Checked once (a device sync)."""
+        if self._slots_verified:
+            return True
+        slot_mapping = cast(dict[str, Any], fwd.slot_mapping)
+        slots = slot_mapping[self._layer_name]
+        index_slots = slot_mapping.get(self._index_prefix)
+        if index_slots is None or not torch.equal(index_slots[:n], slots[:n]):
+            self._disable("index cache does not share the KV cache's block ids")
+            return False
+        self._slots_verified = True
+        return True
+
+    def supports(
+        self,
+        input_ids: torch.Tensor | None,
+        positions: torch.Tensor,
+        intermediate_tensors,
+        inputs_embeds: torch.Tensor | None,
+    ) -> bool:
+        if not self._enabled or intermediate_tensors is not None:
+            return False
+        fwd = get_forward_context()
+        n = positions.numel()
+        if _step_refusal(fwd, self._layer_name, n, positions) is not None:
+            return False
+        return self._ensure_runner() and self._slots_shared(fwd, n)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None,
+        positions: torch.Tensor,
+        inputs_embeds: torch.Tensor | None,
+    ) -> torch.Tensor:
+        assert self._runner is not None
+        if envs.MONO_CHECK:
+            return forward_checked(self._runner, input_ids, positions, inputs_embeds)
+        return self._runner.forward(input_ids, positions, inputs_embeds)
diff --git a/vllm/models/minimax_m3/amd/mono/envs.py b/vllm/models/minimax_m3/amd/mono/envs.py
new file mode 100644
index 000000000000..0ef17463c786
--- /dev/null
+++ b/vllm/models/minimax_m3/amd/mono/envs.py
@@ -0,0 +1,20 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Environment switches of the mono decode path (read once at import)."""
+
+import os
+
+
+def _flag(name: str) -> bool:
+    return os.getenv(name, "0") in ("1", "true", "True")
+
+
+class _Envs:
+    # route supported decode steps to the mono kernels
+    MONO_ENABLE = _flag("VLLM_ROCM_MINIMAX_M3_MONO")
+    # run every sparse layer twice (kernels, then the original modules) and log the
+    # per-token differences; the layer-parity self-check, meant for --enforce-eager
+    MONO_CHECK = _flag("VLLM_ROCM_MINIMAX_M3_MONO_CHECK")
+
+
+envs = _Envs()
diff --git a/vllm/models/minimax_m3/amd/mono/kernels/__init__.py b/vllm/models/minimax_m3/amd/mono/kernels/__init__.py
new file mode 100644
index 000000000000..c8834edb71b7
--- /dev/null
+++ b/vllm/models/minimax_m3/amd/mono/kernels/__init__.py
@@ -0,0 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+# SPDX-License-Identifier: MIT
+# Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved.
+"""FlyDSL kernels of the mono decode path: K1 ``pre_attn`` and K4 ``post_attn``."""
diff --git a/vllm/models/minimax_m3/amd/mono/kernels/common.py b/vllm/models/minimax_m3/amd/mono/kernels/common.py
new file mode 100644
index 000000000000..b42f1e375e3a
--- /dev/null
+++ b/vllm/models/minimax_m3/amd/mono/kernels/common.py
@@ -0,0 +1,418 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+# SPDX-License-Identifier: MIT
+# Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved.
+"""Device-side building blocks shared by the mono kernels.
+
+Ported from the FlyDSL GLM-5 layer kernel (ROCm/FlyDSL branch
+``glm5-mla-moe-layer-monokernel`` @ 1e28fd8a,
+``kernels/mla_moe_layer/shared_reuse_moe_kernel.py``): wave reductions, tagged
+mailboxes and the MFMA GEMV loop. Only published flydsl APIs and aiter's vendored
+``buffer_ops`` / ``dpp_utils`` are used, so nothing here needs the FlyDSL repo.
+
+Every function but ``kernel_symbol`` is traced inside a ``@flyc.kernel`` body.
+
+Raw boundaries the kernels keep, for want of an equivalent on the pinned FlyDSL:
+``buffer_ops`` loads / stores (the kernels take raw Int64 addresses, no tensors
+for ``make_buffer_tensor``); raw MFMA intrinsics (hand-laid operands: permlane
+transposes, fp8 packing, MXFP4 scales); the LLVM intrinsics here (``expect``,
+``s_memrealtime``, the approximate rsq / rcp / exp2) and the struct unpacking of
+``permlane*_swap``; ``inttoptr`` for the timeline's global stores.
+"""
+
+from __future__ import annotations
+
+import flydsl.expr as fx
+from aiter.ops.flydsl.kernels import buffer_ops as bo
+from aiter.ops.flydsl.kernels.dpp_utils import update_dpp_i32
+from aiter.ops.flydsl.kernels.kernels_common import kernel_signature
+from flydsl._mlir import ir
+from flydsl._mlir.dialects import llvm
+from flydsl.compiler.ast_rewriter import ASTRewriter
+from flydsl.expr import const_expr, gpu, range_constexpr, rocdl
+from flydsl.expr import math as fmath
+from flydsl.expr.typing import T, as_ir_value
+
+from vllm.models.minimax_m3.amd.mono.config import LAYER_SLOTS, WAVES
+
+# gfx95x cache policy bits (LLVM CPol): SC0 = 1, SC1 = 16. SC1 alone is device
+# scope (past the per-XCD caches); SC0 | SC1 is system scope (peers over XGMI).
+CM_DEV = 16
+CM_SYS = 17
+POLL_MAX = 12  # mailbox specs polled per batch (bounds live registers)
+
+
+def kernel_symbol(stem, **params):
+    """Kernel symbol that profilers demangle to ``vllm::<stem>_<params>``. The
+    amdhsa assembler rejects ``::`` in a symbol, so it is Itanium-mangled, as
+    aiter's ``_ZN5aiter...E`` kernels are."""
+    name = f"{stem}_{kernel_signature(**params)}"
+    return f"_ZN4vllm{len(name)}{name}E"
+
+
+def rsrc(addr, nbytes=None):
+    return bo.create_buffer_resource_from_addr(addr, num_records_bytes=nbytes)
+
+
+def permlane_swap(off, x, y):
+    """v_permlane{32,16}_swap on registers (x, y): off 32 trades x's upper-half
+    lanes with y's lower-half lanes; off 16 trades x's odd 16-lane rows with y's
+    even rows (each within its half). -> (x', y')."""
+    swap = rocdl.permlane32_swap if off == 32 else rocdl.permlane16_swap
+    pr = swap(
+        llvm.StructType.get_literal([T.i32, T.i32]),
+        as_ir_value(fx.Int32(x)),
+        as_ir_value(fx.Int32(y)),
+        False,
+        False,
+    )
+    return tuple(fx.Int32(llvm.extractvalue(T.i32, pr, [j])) for j in range(2))
+
+
+def lane_gather(v, src_lane):
+    """Lane ``src_lane``'s value of v, a source lane a lane (ds_bpermute)."""
+    return fx.Int32(
+        rocdl.ds_bpermute(T.i32, fx.Int32(src_lane) * 4, fx.Int32(v).ir_value())
+    )
+
+
+def readlane(v, src_lane):
+    """Lane ``src_lane``'s value of v (wave-uniform)."""
+    bits = v.bitcast(fx.Int32) if isinstance(v, fx.Float32) else fx.Int32(v)
+    r = fx.Int32(rocdl.readlane(T.i32, bits, fx.Int32(src_lane)))
+    return r.bitcast(fx.Float32) if isinstance(v, fx.Float32) else r
+
+
+def memrealtime():
+    """s_memrealtime: the 100 MHz device clock (no FlyDSL wrapper)."""
+    return fx.Int64(llvm.call_intrinsic(T.i64, "llvm.amdgcn.s.memrealtime", [], [], []))
+
+
+def unlikely(cond):
+    """``cond`` marked cold (llvm.expect): LLVM lays the branch out of the hot
+    path. A rare path's code inline among the hot stages costs them even when it
+    never runs (the long-context selection: S = 16, 3k context +1.4 us)."""
+    return fx.Boolean(
+        llvm.call_intrinsic(
+            ir.IntegerType.get_signless(1),
+            "llvm.expect.i1",
+            [as_ir_value(cond), as_ir_value(fx.Boolean(False))],
+            [],
+            [],
+        )
+    )
+
+
+def uniform(v):
+    return fx.Int32(rocdl.readfirstlane(T.i32, fx.Int32(v).ir_value()))
+
+
+def _hw_f32(name, x):
+    return fx.Float32(
+        llvm.call_intrinsic(T.f32, name, [fx.Float32(x).ir_value()], [], [])
+    )
+
+
+def hw_rsq(x):
+    """v_rsq_f32 (~1 ulp, no libm range fixups)."""
+    return _hw_f32("llvm.amdgcn.rsq.f32", x)
+
+
+def hw_rcp(x):
+    return _hw_f32("llvm.amdgcn.rcp.f32", x)
+
+
+def hw_exp2(x):
+    return _hw_f32("llvm.amdgcn.exp2.f32", x)
+
+
+def div_rn(x, d, r):
+    """``x / d`` rounded to nearest, bit-identical to the IEEE division, given
+    ``r = 1.0 / d`` from one true division per divisor (Markstein: a quotient
+    within 1 ulp plus its exact fma residual times RN(1/d) rounds correctly)."""
+    q = x * r
+    return fx.Float32(fmath.fma(fx.Float32(fmath.fma(-q, d, x)), r, q))
+
+
+def xshfl(v, off):
+    """Value of lane ``lane ^ off``: offsets 32 / 16 via v_permlane*_swap, in-row
+    offsets via DPP (VALU latency, not LDS)."""
+    if off >= 16:
+        return v.shuffle_xor(off, 64)
+    is_f = isinstance(v, fx.Float32)
+    x = v.bitcast(fx.Int32) if is_f else fx.Int32(v)
+    if off == 8:
+        y = fx.Int32(update_dpp_i32(x, x, 0x118, 0xF, 0xC, False))
+        y = fx.Int32(update_dpp_i32(y, x, 0x108, 0xF, 0x3, False))
+    elif off == 4:
+        y = fx.Int32(update_dpp_i32(x, x, 0x114, 0xF, 0xA, False))
+        y = fx.Int32(update_dpp_i32(y, x, 0x104, 0xF, 0x5, False))
+    elif off == 2:
+        y = fx.Int32(update_dpp_i32(x, x, 0x4E, 0xF, 0xF, False))
+    else:
+        y = fx.Int32(update_dpp_i32(x, x, 0xB1, 0xF, 0xF, False))
+    return y.bitcast(fx.Float32) if is_f else y
+
+
+def xred(v, off, op):
+    """op(v, lane ^ off) for a symmetric op; 32 / 16 take both halves of one swap."""
+    if off < 16:
+        return op(v, xshfl(v, off))
+    is_f = isinstance(v, fx.Float32)
+    x = v.bitcast(fx.Int32) if is_f else fx.Int32(v)
+    a, b = permlane_swap(off, x, x)
+    if is_f:
+        return op(a.bitcast(fx.Float32), b.bitcast(fx.Float32))
+    return op(type(v)(a), type(v)(b))
+
+
+def _add(a, b):
+    return a + b
+
+
+def butterfly(v, offsets, op=_add):
+    """Butterfly reduction over the given xor offsets, in the given order.
+
+    The order is the numerical contract: the aiter per-head norm reduces 32
+    lanes with offsets 16, 8, 4, 2, 1, and matching it keeps rstd bit-exact."""
+    for off in offsets:
+        v = xred(v, off, op)
+    return v
+
+
+def wave_sum(v):
+    return butterfly(v, (32, 16, 8, 4, 2, 1))
+
+
+def traced(fn):
+    """Apply flydsl's AST rewriting to a helper, so its ``while`` / ``if`` over
+    traced values lower to scf ops. Kernels rewrite only their own body."""
+    return ASTRewriter.transform(fn)
+
+
+def wave_max(v):
+    return butterfly(v, (32, 16, 8, 4, 2, 1), fx.max)
+
+
+@traced
+def block_max(v, lane, wave, red):
+    """Max of v over the CTA's WAVES waves; ``red`` needs WAVES words."""
+    wv = wave_max(v)
+    gpu.barrier()
+    if lane == 0:
+        fx.ptr_store(wv, red + wave)
+    gpu.barrier()
+    t = fx.ptr_load(red + 0)
+    for w in range_constexpr(1, WAVES):
+        t = fx.max(t, fx.ptr_load(red + w))
+    gpu.barrier()
+    return t
+
+
+def mfma_bf16(a, b, c):
+    return fx.Vector(rocdl.mfma_f32_16x16x32_bf16(T.vec(4, T.f32), [a, b, c]))
+
+
+def fp8x8_bf16_pk(words):
+    """Eight fp8 of two dwords -> bf16 MFMA operand (exact) by
+    v_cvt_scalef32_pk_bf16_fp8 (scale 1), two values an instruction."""
+    out = []
+    for w in words:
+        for hi in (False, True):
+            p = rocdl.cvt_scalef32_pk_bf16_fp8(
+                T.bf16x2, as_ir_value(fx.Int32(w)), as_ir_value(fx.Float32(1.0)), hi
+            )
+            out.append(fx.Vector(p).bitcast(fx.Int32)[0])
+    return fx.Vector.from_elements(out, fx.Int32).bitcast(fx.BFloat16)
+
+
+def bf16_round(a):
+    return fx.Float32(fx.Float32(a).to(fx.BFloat16))
+
+
+def bf16_pair(a, b):
+    """Two f32 -> one f32-typed word holding (bf16(a), bf16(b))."""
+    return (
+        fx.Vector.from_elements([a, b], fx.Float32)
+        .to(fx.BFloat16)
+        .bitcast(fx.Float32)[0]
+    )
+
+
+def bf2_f32(w):
+    """Packed bf16 pair word -> (f32 low, f32 high)."""
+    return (w << 16).bitcast(fx.Float32), (w & fx.Int32(-65536)).bitcast(fx.Float32)
+
+
+def fp8_pack4(a, b, c, d):
+    """Four f32 (already in range) -> one dword of four E4M3 bytes, RNE."""
+    lo = fx.Int32(rocdl.cvt_pk_fp8_f32(T.i32, a, b, fx.Int32(0), False))
+    return fx.Int32(rocdl.cvt_pk_fp8_f32(T.i32, c, d, lo, True))
+
+
+@traced
+def _spin(load_all, pending):
+    """Re-issue ``load_all`` until ``pending`` clears; a side-effecting op in the
+    loop keeps the loads from being hoisted."""
+    v = load_all()
+    while pending(v):
+        rocdl.s_sleep(0)
+        v = load_all()
+    return v
+
+
+class Mailbox:
+    """Tagged-pair hand-off between CTAs (device scope) and GPUs (system scope).
+
+    Every 32-bit value is stored next to this launch's epoch tag; a consumer
+    polls the payload until all tags match, so a hand-off costs one round trip
+    (no store drain, no separate flag). The epoch is ``step * LAYER_SLOTS +
+    layer + 1`` with ``step`` a device counter, so the sequence is graph-safe.
+    """
+
+    def __init__(self, scratch, step_ptr, layer):
+        self.scratch = scratch
+        self.tag = (
+            uniform(bo.buffer_load(rsrc(step_ptr), 0, vec_width=1, dtype=T.i32))
+            * LAYER_SLOTS
+            + layer
+            + 1
+        )
+
+    def addr(self, offset):
+        return self.scratch + fx.Int64(offset)
+
+    def put(self, base_addr, i, v, cm=CM_DEV):
+        bits = v.bitcast(fx.Int32) if isinstance(v, fx.Float32) else fx.Int32(v)
+        bo.buffer_store(
+            fx.Vector.from_elements([bits, self.tag], fx.Int32),
+            rsrc(base_addr),
+            i * 2,
+            cache_modifier=cm,
+        )
+
+    def put_bf(self, base_addr, i, vs, cm=CM_DEV):
+        """Store 2 or 4 elements as aligned packed bf16 pairs."""
+        words = []
+        for j in range_constexpr(len(vs) // 2):
+            words += [bf16_pair(vs[2 * j], vs[2 * j + 1]).bitcast(fx.Int32), self.tag]
+        bo.buffer_store(
+            fx.Vector.from_elements(words, fx.Int32),
+            rsrc(base_addr),
+            i,
+            cache_modifier=cm,
+        )
+
+    def put_words(self, base_addr, i, words, cm=CM_DEV):
+        """Raw 32-bit words (1 or 2) at pairs i.. as (word, tag). Integer payloads
+        (fp8 words) must come this way: through an f32 value a NaN bit pattern is
+        not kept."""
+        vals = []
+        for w in words:
+            vals += [fx.Int32(w), self.tag]
+        bo.buffer_store(
+            fx.Vector.from_elements(vals, fx.Int32),
+            rsrc(base_addr),
+            i * 2,
+            cache_modifier=cm,
+        )
+
+    def poll(self, specs, scope="agent", batch=POLL_MAX):
+        """Poll ``(base_addr, pair index, npairs)`` specs in one batch.
+
+        All pairs are re-loaded together while any tag is stale, so a batch costs
+        one round trip after its last producer lands. ``scope``: "agent" for pairs
+        another CTA of this GPU wrote, "system" for pairs a peer GPU wrote. Returns
+        the value words."""
+        if const_expr(len(specs) == 0):
+            return []
+        if const_expr(len(specs) > batch):
+            return self.poll(specs[:batch], scope, batch) + self.poll(
+                specs[batch:], scope, batch
+            )
+        assert scope in ("agent", "system"), scope
+        cm = CM_DEV if const_expr(scope == "agent") else CM_SYS
+        tag = self.tag
+
+        def load_all():
+            words = []
+            for b, i, n in specs:
+                w = fx.Vector(
+                    bo.buffer_load(
+                        rsrc(b),
+                        fx.Int32(i) * 2,
+                        vec_width=2 * n,
+                        dtype=T.i32,
+                        cache_modifier=cm,
+                    )
+                )
+                words += [w[e] for e in range(2 * n)]
+            return fx.Vector.from_elements(words, fx.Int32)
+
+        nw = sum(2 * n for _, _, n in specs)
+
+        def pending(v):
+            bad = v[1] != tag
+            for e in range_constexpr(3, nw, 2):
+                bad = bad | (v[e] != tag)
+            return bad
+
+        v = _spin(load_all, pending)
+        outs, e = [], 0
+        for _, _, n in specs:
+            outs.append([v[e + 2 * q] for q in range(n)])
+            e += 2 * n
+        return outs
+
+
+def ar_pack_sumsq(vals):
+    """One logical thread of the custom all-reduce epilogue: its 8 elements'
+    squares accumulated in order (hipcc contracts ``acc += v * v`` into fma)."""
+    acc = fx.Float32(0.0)
+    for v in vals:
+        acc = fmath.fma(v, v, acc)
+    return acc
+
+
+@traced
+def ar_block_sum(p0, p1, lane, wave, red):
+    """Sum of the 768 per-pack partials of a 6144-wide bf16 row in the order the
+    1-stage fused all-reduce + RMSNorm reduces them: 24 warps of 32, a 16..1
+    butterfly per warp, then the 24 warp sums butterflied again in one warp.
+
+    Logical warp ``3 * wave + lane // 32`` holds ``p0``; logical warp
+    ``3 * wave + 2`` holds ``p1`` in lanes < 32 (``WAVES`` = 8). ``red`` needs
+    24 words. Matching this order keeps rstd bit-exact with the original path.
+    """
+    return ar_block_sums([(p0, p1)], lane, wave, red)[0]
+
+
+@traced
+def ar_block_sums(parts, lane, wave, red):
+    """``ar_block_sum`` of several rows ``parts = [(p0, p1), ...]`` behind one set
+    of barriers; ``red`` needs 24 words per row."""
+    bs = [
+        (butterfly(p0, (16, 8, 4, 2, 1)), butterfly(p1, (16, 8, 4, 2, 1)))
+        for p0, p1 in parts
+    ]
+    gpu.barrier()
+    for k in range_constexpr(len(bs)):
+        b0, b1 = bs[k]
+        if lane % 32 == 0:
+            fx.ptr_store(b0, red + (24 * k + wave * 3 + lane // 32))
+        if lane == 0:
+            fx.ptr_store(b1, red + (24 * k + wave * 3 + 2))
+    gpu.barrier()
+    li = lane % 32
+    tots = [
+        butterfly(
+            (li < 24).select(
+                fx.ptr_load(red + (24 * k + fx.min(li, 23))), fx.Float32(0.0)
+            ),
+            (16, 8, 4, 2, 1),
+        )
+        for k in range(len(parts))
+    ]
+    gpu.barrier()
+    return tots
diff --git a/vllm/models/minimax_m3/amd/mono/kernels/index_score.py b/vllm/models/minimax_m3/amd/mono/kernels/index_score.py
new file mode 100644
index 000000000000..ac95896137e4
--- /dev/null
+++ b/vllm/models/minimax_m3/amd/mono/kernels/index_score.py
@@ -0,0 +1,474 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+# SPDX-License-Identifier: MIT
+# Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved.
+"""Indexer block scores for K4's sparse selection (``select_blocks`` polls them).
+
+The Triton decode scorer's math, one 128-key block of index-K per task: bf16(fp8
+K) x bf16 index q on the MFMA, max over the causal keys, times fp32(sm_scale *
+log2 e). A context of <= TOPK_BLOCKS blocks keeps every block and is not scored;
+nor are the blocks the selector pins (the first ``init_blocks``, the last
+``local_blocks``): their keys are forced. A scored block waits for the index keys
+K1 inserts this step only when it may hold one (``wait_new_keys``).
+
+A task is one block of one request and runs on one wave: the request's q tokens
+(a speculative verify's rows) are the MFMA's B columns, so the block's index-K is
+read once for all of them, and no task needs the CTA (no barrier).
+
+Emitted at the end of K1 (``emit_index_scores``), which makes index q: K4's
+selector then finds the scores in place instead of scoring after its launch.
+``build_index_score_kernel`` runs the same tasks alone, for K4's harnesses.
+"""
+
+import struct
+from typing import NamedTuple
+
+import flydsl.compiler as flyc
+import flydsl.expr as fx
+from aiter.ops.flydsl.kernels import buffer_ops as bo
+from aiter.ops.flydsl.kernels.kernels_common import LOG2E
+from flydsl.expr import const_expr, gpu, range_constexpr
+from flydsl.expr.typing import Int32, Int64, T
+
+from vllm.models.minimax_m3.amd.mono.config import (
+    BLOCKS,
+    HEAD_DIM,
+    INDEX_CP_FROM_BLOCKS,
+    MAX_INDEX_BLOCKS,
+    MAX_TOKENS,
+    ONE_INDEX_HEAD,
+    SPARSE_BLOCK,
+    THREADS,
+    TOPK_BLOCKS,
+    TP,
+    WAVES,
+    IndexHeads,
+)
+from vllm.models.minimax_m3.amd.mono.kernels.common import (
+    Mailbox,
+    butterfly,
+    fp8x8_bf16_pk,
+    kernel_symbol,
+    lane_gather,
+    mfma_bf16,
+    permlane_swap,
+    readlane,
+    rsrc,
+    traced,
+    uniform,
+    unlikely,
+)
+
+INDEX_BLOCK_BYTES = SPARSE_BLOCK * HEAD_DIM  # one fp8 index-K block
+KEY_TILES = SPARSE_BLOCK // 16  # MFMA row tiles of a block
+_CURRENT_STREAM = fx.Stream(None)
+
+
+def _f32(x: float) -> float:
+    return struct.unpack("f", struct.pack("f", x))[0]
+
+
+def index_scale_log2e(sm_scale: float) -> float:
+    """fp32(sm_scale * log2 e), the scorer's scale."""
+    return _f32(_f32(sm_scale) * _f32(LOG2E))
+
+
+def scored_count(n_blk, init_blocks: int, local_blocks: int):
+    """How many blocks (from block init_blocks on) a context of n_blk blocks scores."""
+    return (n_blk > TOPK_BLOCKS).select(n_blk - (init_blocks + local_blocks), 0)
+
+
+class StepRows(NamedTuple):
+    """A decode step's rows as the indexer sees them (``step_rows``)."""
+
+    seq: list  # each row's context, >= 1 (a graph's pad rows carry 0)
+    n_blks: list  # each row's index blocks
+    long: list  # each row's request is long
+
+    def any_long(self):
+        v = self.long[0]
+        for lg in self.long[1:]:
+            v = v | lg
+        return v
+
+
+def step_rows(seq_lens, tokens: int, q_len, heads: IndexHeads) -> StepRows:
+    """The one place a request is decided long: past INDEX_CP_FROM_BLOCKS index
+    blocks with every index head (its scores take the context-parallel layout,
+    ``emit_index_scores``), past THREADS with one (more blocks than
+    ``rank_blocks`` ranks); either way K4 selects it past the split stage.
+
+    Decided by the request's last row (a request is ``q_len`` consecutive rows):
+    a speculative verify's rows can sit either side of the threshold, and a
+    request's rows must all read the layout its scores were written in."""
+    seq = [
+        fx.max(uniform(bo.buffer_load(rsrc(seq_lens), k, vec_width=1, dtype=T.i32)), 1)
+        for k in range(tokens)
+    ]
+    n_blks = [(s + SPARSE_BLOCK - 1) // SPARSE_BLOCK for s in seq]
+    req_blks = list(n_blks)
+    for k in range(tokens - 2, -1, -1):
+        req_blks[k] = (fx.Int32(k) % q_len == q_len - 1).select(
+            n_blks[k], req_blks[k + 1]
+        )
+    past = INDEX_CP_FROM_BLOCKS if heads.count > 1 else THREADS
+    return StepRows(seq, n_blks, [nb > past for nb in req_blks])
+
+
+def _pick(vals, k):
+    """vals[k] for a traced index k."""
+    v = vals[0]
+    for j in range_constexpr(1, len(vals)):
+        v = (k == j).select(vals[j], v)
+    return v
+
+
+@traced
+def emit_index_scores(
+    slot, lane, wave, red, mb_put, iscore, q_frag, tokens, q_len, init_blocks,
+    local_blocks, scale, index_cache, block_table, seq_lens, bt_width,
+    wait_new_keys=None, prepare=None, heads=ONE_INDEX_HEAD, q_frag_head=None,
+    prepare_heads=None,
+):  # fmt: skip
+    """This CTA's (``slot``) share of the step's score tasks: up to BLOCKS tasks,
+    task u runs on CTA u, else on wave (u // BLOCKS) % WAVES of CTA u % BLOCKS
+    (``red``: WAVES * 16 floats of LDS). The tokens form requests of ``q_len``
+    consecutive rows sharing a block table; the last one has the longest context
+    and lists the request's scored blocks. ``q_frag(k, s)``: this lane's B operand
+    for token k, k-chunk s (bf16x8 of index q). Pair k * MAX_INDEX_BLOCKS + b of
+    mailbox ``iscore`` <- the score of token k's block b.
+
+    ``wait_new_keys()`` (K1 fused into the layer kernel; wave-local) waits for the
+    index keys this launch inserts: a block within MAX_TOKENS keys of a token's
+    end may hold an earlier token of its request (speculative decode rows).
+    Without it, every key this launch inserts must sit in a pinned block
+    (one-token requests). ``prepare(first, last)`` (with CTA barriers) runs on a
+    CTA with tasks before any index q is read: K1 makes rows first..last's
+    there.
+
+    Indexer context parallelism (``heads.count`` > 1): a long request
+    (``step_rows``) scores only this rank's blocks, b = heads.own
+    (mod TP) -- every index q head of them, a task's key block read once for all:
+    pair k * MAX_INDEX_BLOCKS + h * (MAX_INDEX_BLOCKS / TP) + b / TP <- the score
+    of token k's block b for head h. ``q_frag_head(k, s, h)``: head h's B operand
+    (every lane its own h); ``prepare_heads()`` (with CTA barriers, on every CTA)
+    makes those requests' rows' heads first. The other requests score as above."""
+    assert local_blocks >= 1, "the block of the current token must be pinned"
+    l16 = lane % 16
+    g4 = lane // 16
+    r_ic = rsrc(index_cache)
+    rows = step_rows(seq_lens, tokens, q_len, heads)
+    seq = rows.seq
+    # every request's scored blocks in one list, a task per wave across the grid;
+    # context-parallel requests' (this rank's blocks) in a second
+    cp = heads.count > 1
+    first_own = init_blocks + (heads.own - init_blocks) % TP
+    pre = [fx.Int32(0)]
+    pre_cp = [fx.Int32(0)]
+    for k in range_constexpr(tokens):
+        last = fx.Int32(k) % q_len == q_len - 1
+        n_scored = scored_count(rows.n_blks[k], init_blocks, local_blocks)
+        if const_expr(cp):
+            long = rows.long[k]
+            n_own = fx.max(init_blocks + n_scored - first_own + TP - 1, 0) // TP
+            pre.append(pre[-1] + (last & ~long).select(n_scored, fx.Int32(0)))
+            pre_cp.append(pre_cp[-1] + (last & long).select(n_own, fx.Int32(0)))
+        else:
+            pre.append(pre[-1] + last.select(n_scored, fx.Int32(0)))
+    total = pre[tokens]
+    total_cp = pre_cp[-1]  # 0 without context parallelism
+
+    r_bt = rsrc(block_table)
+
+    def task_block(u, pre=pre, first=init_blocks, stride=1):
+        """Task u of list ``pre`` -> (the request's last token, block b); per lane
+        if u is."""
+        tk = fx.Int32(0)
+        for j in range_constexpr(1, tokens):
+            tk = tk + (u >= pre[j]).select(fx.Int32(1), fx.Int32(0))
+        return tk, first + stride * (u - _pick(pre[:tokens], tk))
+
+    def page_of(tk, b):
+        return fx.Int32(
+            bo.buffer_load(r_bt, tk * bt_width + b, vec_width=1, dtype=T.i32)
+        )
+
+    def new_keys_in(tk, b):
+        """Block b of the request ending at token tk may hold keys this launch
+        inserts (a speculative verify's earlier rows)."""
+        return (b + 1) * SPARSE_BLOCK > _pick(seq, tk) - MAX_TOKENS
+
+    def cols_of(tk):
+        """B column l16's token (the request's token l16, the last one repeated)
+        and its context."""
+        col = tk - (q_len - 1) + fx.min(l16, q_len - 1)
+        return col, _pick(seq, col)
+
+    def load_raw(page, tiles):
+        """Key tiles ``tiles`` (16 keys each) of a block, in flight: lane g4 =
+        lane / 16 reads its key's bytes 32 g4 .. 32 g4 + 32 with two 16 B loads
+        (half the load instructions of eight-byte operand reads, scattered over
+        the rows: 1M context, S = 16, scores 125 -> 100 us); ``split_tiles``
+        makes them MFMA operands."""
+        words = []
+        for i in tiles:
+            row = (page * INDEX_BLOCK_BYTES + (16 * i + l16) * HEAD_DIM + 32 * g4) // 4
+            for h in range_constexpr(2):
+                w = fx.Vector(
+                    bo.buffer_load(r_ic, row + 4 * h, vec_width=4, dtype=T.i32)
+                )
+                words += [w[e] for e in range(4)]
+        return fx.Vector.from_elements(words, fx.Int32)
+
+    def split_tiles(raw, tiles):
+        """``load_raw``'s words -> per tile the MFMA A operands: tile i holds keys
+        16 i + lane % 16, k = 32 s + 8 (lane / 16). A 4 x 4 transpose of the 8 B
+        parts across the key's four lanes (permlane swaps) gives each lane part
+        g4 of every chunk s."""
+        out = []
+        for n in range_constexpr(len(tiles)):
+            d = [raw[8 * n + e] for e in range(8)]
+            # part j = dwords 2 j, 2 j + 1; lanes g4 and g4 ^ 2, then g4 ^ 1, trade
+            for off, pairs in ((32, ((0, 2), (1, 3))), (16, ((0, 1), (2, 3)))):
+                for a, b in pairs:
+                    for w in range_constexpr(2):
+                        d[2 * a + w], d[2 * b + w] = permlane_swap(
+                            off, d[2 * a + w], d[2 * b + w]
+                        )
+            out.append(
+                [
+                    fx.Vector.from_elements([d[2 * s], d[2 * s + 1]], fx.Int32)
+                    for s in range(HEAD_DIM // 32)
+                ]
+            )
+        return out
+
+    def tiles_max(b, qb, col_seq, tiles, kw):
+        """Max over key tiles ``tiles`` of block b (``kw``: ``split_tiles``) of this
+        lane's column's causal scores (its B operand ``qb``, per k-chunk; the
+        column's context ``col_seq``), the same in every lane of the column."""
+        v = fx.Float32(float("-inf"))
+        for n, i in enumerate(tiles):
+            c = fx.Vector.filled(4, 0.0, fx.Float32)
+            for s in range_constexpr(HEAD_DIM // 32):
+                c = mfma_bf16(fp8x8_bf16_pk(kw[n][s]), qb[s], c)
+            # lane holds keys 16 i + 4 (lane / 16) + e of column lane % 16
+            key0 = b * SPARSE_BLOCK + 16 * i + 4 * g4
+            for e in range_constexpr(4):
+                v = fx.max(v, (key0 + e < col_seq).select(c[e], float("-inf")))
+        return butterfly(v, (32, 16), fx.max)
+
+    def put(b, col, v):
+        if (lane < 16) & (l16 < q_len):
+            mb_put(iscore, col * MAX_INDEX_BLOCKS + b, v * scale)
+
+    def emit_cp_tasks():
+        """The context-parallel requests' tasks, a task a wave: B column f of pass
+        p (f = 16 p + lane % 16) is head f / q_len's row f % q_len of the request,
+        the key block's operands kept for every pass."""
+        if const_expr(prepare_heads is not None):
+            prepare_heads()
+        tiles = range(KEY_TILES)
+        seq_lane = _pick(seq, fx.min(lane, tokens - 1))  # lane k: token k's context
+        n_cols = TP * q_len
+
+        def score_cp(tk, b, page):
+            """Block b of the request ending at token tk, at ``page``: every head."""
+            if const_expr(wait_new_keys is not None):  # noqa: SIM102
+                if (b + 1) * SPARSE_BLOCK > readlane(seq_lane, tk) - MAX_TOKENS:
+                    wait_new_keys()
+            kw = split_tiles(load_raw(page, tiles), tiles)
+            # every pass unrolled, the unused ones skipped: a runtime pass loop
+            # slowed the whole kernel's short contexts (S = 1, 3k: +0.6 us)
+            for p in range_constexpr(TP * MAX_TOKENS // 16):
+                if p * 16 < n_cols:
+                    f = p * 16 + l16
+                    head = fx.min(f // q_len, TP - 1)
+                    col = tk - (q_len - 1) + f % q_len
+                    qb = [q_frag_head(col, s, head) for s in range(HEAD_DIM // 32)]
+                    v = tiles_max(b, qb, lane_gather(seq_lane, col), tiles, kw)
+                    if (lane < 16) & (f < n_cols):
+                        mb_put(
+                            iscore,
+                            col * MAX_INDEX_BLOCKS
+                            + head * (MAX_INDEX_BLOCKS // TP)
+                            + b // TP,
+                            v * scale,
+                        )
+
+        # as the wave mode below: task u on wave (u // BLOCKS) % WAVES of CTA
+        # u % BLOCKS, lane j decoding task c0 + step j
+        step = BLOCKS * WAVES
+        for c0 in range(slot + BLOCKS * wave, total_cp, step * 64):
+            c0 = fx.Int32(c0)
+            tk_j, b_j = task_block(
+                fx.min(c0 + step * lane, total_cp - 1), pre_cp, first_own, TP
+            )
+            pages = page_of(tk_j, b_j)
+            for i in range(fx.min((total_cp - c0 + step - 1) // step, 64)):
+                i = fx.Int32(i)
+                score_cp(readlane(tk_j, i), readlane(b_j, i), readlane(pages, i))
+
+    # Up to BLOCKS tasks: one a CTA, a key tile a wave, the columns' maxima met in
+    # LDS; the keys go out before prepare() (only the task's request's rows), so
+    # their latency hides under the index q's: the scores are on the step's
+    # critical path (keys after: S = 1, 3k context +1 us). More: a task a wave,
+    # every row's index q first (a wave's task count differs: no barrier inside),
+    # the pages of a wave's next 64 tasks fetched at once, a lane each.
+    if total <= BLOCKS:
+        for u in range(slot, total, BLOCKS):
+            tk, b = task_block(fx.Int32(u))
+            page = uniform(page_of(tk, b))
+            if const_expr(wait_new_keys is not None):  # noqa: SIM102
+                if new_keys_in(tk, b):
+                    wait_new_keys()
+            raw = load_raw(page, [wave])
+            if const_expr(prepare is not None):
+                prepare(tk - (q_len - 1), tk)
+            col, col_seq = cols_of(tk)
+            qb = [q_frag(col, s) for s in range(HEAD_DIM // 32)]
+            v = tiles_max(b, qb, col_seq, [wave], split_tiles(raw, [wave]))
+            if lane < 16:
+                fx.ptr_store(v, red + (wave * 16 + l16))
+            gpu.barrier()
+            if wave == 0:
+                for w in range_constexpr(1, WAVES):
+                    v = fx.max(v, fx.ptr_load(red + (w * 16 + l16)))
+                put(b, col, v)
+            gpu.barrier()
+    if total > BLOCKS:
+        if const_expr(prepare is not None):
+            prepare(fx.Int32(0), fx.Int32(tokens - 1))
+        tiles = range(KEY_TILES)
+        seq_lane = _pick(seq, fx.min(lane, tokens - 1))  # lane k: token k's context
+
+        def score_task(tk, b, page):
+            """Block b of the request ending at token tk, at ``page``."""
+            if const_expr(wait_new_keys is not None):  # noqa: SIM102
+                if (b + 1) * SPARSE_BLOCK > readlane(seq_lane, tk) - MAX_TOKENS:
+                    wait_new_keys()
+            raw = load_raw(page, tiles)
+            col = tk - (q_len - 1) + fx.min(l16, q_len - 1)
+            col_seq = lane_gather(seq_lane, col)
+            qb = [q_frag(col, s) for s in range(HEAD_DIM // 32)]
+            put(b, col, tiles_max(b, qb, col_seq, tiles, split_tiles(raw, tiles)))
+
+        # Every wave the same task count: task u on wave
+        # (u // BLOCKS) % WAVES of CTA u % BLOCKS; lane j decodes task c0 + step j
+        # once -- token, block, page -- read back per task (a task's select chains
+        # sat between one task's keys and the next: S = 16, +8 us).
+        step = BLOCKS * WAVES
+        for c0 in range(slot + BLOCKS * wave, total, step * 64):
+            c0 = fx.Int32(c0)
+            tk_j, b_j = task_block(fx.min(c0 + step * lane, total - 1))
+            pages = page_of(tk_j, b_j)
+            for i in range(fx.min((total - c0 + step - 1) // step, 64)):
+                i = fx.Int32(i)
+                score_task(readlane(tk_j, i), readlane(b_j, i), readlane(pages, i))
+    if const_expr(cp):  # noqa: SIM102
+        # cold: a step without context-parallel requests runs none of it
+        if unlikely(total_cp > 0):
+            emit_cp_tasks()
+
+
+def q_frag_from(iq, heads=ONE_INDEX_HEAD):
+    """(`q_frag`, `q_frag_head`) reading index q from memory:
+    `[token][head][HEAD_DIM]` bf16, `heads.count` heads a token."""
+
+    def q_frag_head(tk, s, head):
+        lane = fx.thread_idx.x % 64
+        return fx.Vector(
+            bo.buffer_load(
+                rsrc(iq),
+                ((tk * heads.count + head) * HEAD_DIM + 32 * s + 8 * (lane // 16)) // 2,
+                vec_width=4,
+                dtype=T.i32,
+            )
+        ).bitcast(fx.BFloat16)
+
+    def q_frag(tk, s):
+        return q_frag_head(tk, s, heads.own)
+
+    return q_frag, q_frag_head
+
+
+def build_index_score_kernel(
+    sm_scale: float,
+    init_blocks: int,
+    local_blocks: int,
+    tokens: int = 1,
+    heads: IndexHeads = ONE_INDEX_HEAD,
+):
+    """``@flyc.jit`` launcher of the score tasks alone (index q from memory,
+    ``heads.count`` heads a token), for K4's harnesses: scores into mailbox
+    ``iscore`` with K4's (step, layer) tag."""
+    assert 1 <= tokens <= MAX_TOKENS
+    scale = index_scale_log2e(sm_scale)
+    # the JIT cache key holds scalar closure values only: the heads go in as ints
+    ih_count, ih_own = heads.count, heads.own
+
+    @fx.struct
+    class Smem:
+        red: fx.Array[fx.Float32, WAVES * 16, 16]
+
+    kernel_name = kernel_symbol(
+        "minimax_m3_index_score",
+        s=tokens,
+        ib=init_blocks,
+        lb=local_blocks,
+        ih=heads.count,
+        io=heads.own,
+    )
+
+    @flyc.kernel(name=kernel_name, known_block_size=[THREADS, 1, 1])
+    def index_score_kernel(
+        iq: Int64,
+        index_cache: Int64,
+        block_table: Int64,
+        seq_lens: Int64,
+        bt_width: Int32,
+        q_len: Int32,
+        iscore: Int64,
+        step: Int64,
+        layer: Int32,
+    ):
+        tid = fx.thread_idx.x
+        lds = fx.SharedAllocator().allocate(Smem).peek()
+        red = lds.red.ptr
+        mbox = Mailbox(iscore, step, layer)
+        mb_put = mbox.put
+        k_heads = IndexHeads(ih_count, ih_own)
+        q_frag, q_frag_head = q_frag_from(iq, k_heads)
+        emit_index_scores(
+            fx.block_idx.x, tid % 64, uniform(tid // 64), red, mb_put, iscore,
+            q_frag, tokens, q_len, init_blocks, local_blocks, scale, index_cache,
+            block_table, seq_lens, bt_width, heads=k_heads, q_frag_head=q_frag_head,
+        )  # fmt: skip
+
+    @flyc.jit
+    def launch_index_score(
+        iq: Int64,
+        index_cache: Int64,
+        block_table: Int64,
+        seq_lens: Int64,
+        bt_width: Int32,
+        q_len: Int32,
+        iscore: Int64,
+        step: Int64,
+        layer: Int32,
+        stream: fx.Stream = _CURRENT_STREAM,
+    ):
+        index_score_kernel(
+            iq,
+            index_cache,
+            block_table,
+            seq_lens,
+            bt_width,
+            q_len,
+            iscore,
+            step,
+            layer,
+        ).launch(grid=(BLOCKS,), block=(THREADS,), stream=stream)
+
+    return launch_index_score
diff --git a/vllm/models/minimax_m3/amd/mono/kernels/post_attn.py b/vllm/models/minimax_m3/amd/mono/kernels/post_attn.py
new file mode 100644
index 000000000000..7824d15b8030
--- /dev/null
+++ b/vllm/models/minimax_m3/amd/mono/kernels/post_attn.py
@@ -0,0 +1,2963 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+# SPDX-License-Identifier: MIT
+# Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved.
+"""K4 ``mono_post``: one MiniMax-M3 sparse MoE layer, attention through FFN reduce.
+
+One launch per layer and rank, ``BLOCKS`` CTAs x ``THREADS`` threads, for a step of
+S = 1..MAX_TOKENS tokens (each token's stage tasks side by side); for one token::
+
+    index : K1, fused in, scored every index-K block (the Triton decode scorer's
+            block score; with indexer context parallelism a long context's
+            blocks of this rank, for every index head); every split CTA then
+            runs the Triton selector's top-k and page-16 table emission itself
+            (ranking, no sort; for a long request the split tasks' shares'
+            best TOPK_BLOCKS are ranked, gathered from the peers under CP)
+    split : one 256-key context partition x 16 heads per task (the gluon decode's
+            partitioning) over that page-16 table, FP8 QK / PV with its scale
+            placement
+    o     : combine the partitions (the gluon reduce), per-token FP8 (standalone
+            quant formula) -> o_proj GEMV
+            -> push bf16 partial rows to every peer -> rank-ordered sum
+            -> acc = bf16(sum) + h   (h_mid = bf16(acc), mailbox ``a`` = acc)
+    router: gemma RMSNorm(acc) in the fused all-reduce's reduction order -> bf16
+            gate GEMV -> logits; publishes its slice of the normalized input
+    moe   : sigmoid + bias top-4 (+ the fused shared expert) -> MXFP4 a16w4
+            up/gate on aiter's shuffled layout -> swiglu -> mid (bf16);
+            the CTA's down task reuses that routing, its weights prefetched
+            under the up/gate GEMV: route-weighted down GEMV -> push bf16
+            partial rows -> rank-ordered sum -> ar_out = bf16(sum)
+
+From ``WIDE_FROM`` tokens the MoE groups its work by expert (``wide_moe``): the
+router tasks publish xn as MXFP8, each up / gate task pair runs one distinct
+expert's 32 rows for every token routed to it (a8w4 MFMA, a B column per token)
+and publishes MXFP8 mid rows, and a down task sums over the step's (expert,
+token) mids. From ``AG_RS_FROM`` tokens both all-reduces are a reduce-scatter to
+the row group's owner rank, which sums in rank order and all-gathers the result.
+
+The next layer's K1 takes (ar_out, h_mid) exactly as the original path's fused
+all-reduce + RMSNorm takes (partials, residual).
+
+Every stage is a task list; task t runs on CTA (base + t) % BLOCKS and every CTA
+walks the stages in order. Dependencies only point to earlier stages and all
+CTAs are co-resident, so every spin wait makes progress.
+"""
+
+import struct
+
+import flydsl.compiler as flyc
+import flydsl.expr as fx
+from aiter.ops.flydsl.kernels import buffer_ops as bo
+from aiter.ops.flydsl.kernels.act import swiglu_mul_batch
+from aiter.ops.flydsl.kernels.kernels_common import LOG2E
+from flydsl.expr import const_expr, gpu, range_constexpr, rocdl
+from flydsl.expr import math as fmath
+from flydsl.expr.typing import Int32, Int64, T, as_ir_value
+
+from vllm.models.minimax_m3.amd.mono import config as mono_config
+from vllm.models.minimax_m3.amd.mono.config import (
+    BLOCKS,
+    HEAD_DIM,
+    HIDDEN,
+    INDEX_CP_FROM_BLOCKS,
+    INTER,
+    MAX_INDEX_BLOCKS,
+    MAX_TOKENS,
+    MOE_SLOTS,
+    N_ROUTED,
+    ONE_INDEX_HEAD,
+    PAGE16,
+    SHARED_EXPERT,
+    SPARSE_BLOCK,
+    THREADS,
+    TOP_K,
+    TOPK_BLOCKS,
+    TP,
+    WAVES,
+    IndexHeads,
+)
+from vllm.models.minimax_m3.amd.mono.kernels.common import (
+    CM_DEV,
+    CM_SYS,
+    Mailbox,
+    bf2_f32,
+    bf16_pair,
+    bf16_round,
+    butterfly,
+    fp8_pack4,
+    hw_exp2,
+    hw_rcp,
+    hw_rsq,
+    kernel_symbol,
+    memrealtime,
+    rsrc,
+    uniform,
+    unlikely,
+    wave_max,
+    wave_sum,
+    xshfl,
+)
+from vllm.models.minimax_m3.amd.mono.kernels.common import (
+    block_max as block_max_of,
+)
+from vllm.models.minimax_m3.amd.mono.kernels.index_score import (
+    index_scale_log2e,
+    step_rows,
+)
+from vllm.models.minimax_m3.amd.mono.kernels.pre_attn import (
+    K1_ARGS,
+    emit_k1,
+)
+from vllm.models.minimax_m3.amd.mono.kernels.pre_attn import (
+    N_HEAD_TASKS as K1_HEAD_TASKS,
+)
+from vllm.models.minimax_m3.amd.mono.kernels.pre_attn import (
+    SCRATCH_HDONE as K1_SCRATCH_HDONE,
+)
+from vllm.models.minimax_m3.amd.mono.kernels.pre_attn import (
+    SCRATCH_RDONE as K1_SCRATCH_RDONE,
+)
+from vllm.models.minimax_m3.amd.mono.layout import (
+    AG_RS_FROM,
+    DN_ROWS,
+    MID_ROW,
+    MIDSC_ROW,
+    N_DN,
+    N_O,
+    N_ROUTER,
+    N_SPLIT,
+    O_K,
+    O_ROW,
+    O_ROWS,
+    PAGES_PER_BLOCK,
+    SCORE_COPIES,
+    SCRATCH,
+    SH_ROWS,
+    SH_TASKS,
+    SPARSE_TABLE_WORDS,
+    SPLIT_KEYS,
+    UG_PER_SLOT,
+    UG_ROWS,
+    XN8_ROW,
+    XSC_ROW,
+    H,
+    pool_words,
+    shared_after_routing,
+    stage_bases,
+    sym_layout,
+    ug_ctas_per_token,
+    ug_tasks_of,
+    wide_moe,
+)
+
+FP8_MAX = 448.0
+NEG = -3.4e38  # the gluon decode's masked-score value
+PAGE_BYTES = PAGE16 * HEAD_DIM
+XN_SLICE = HIDDEN // N_ROUTER  # normalized-input elements each router task publishes
+W13_ROWS = 2 * INTER  # gate rows then up rows (GGUU)
+W13_BYTES = W13_ROWS * HIDDEN // 2  # one expert, fp4
+W2_BYTES = HIDDEN * INTER // 2
+W13_SCALE_COLS = HIDDEN // 32
+W2_SCALE_COLS = INTER // 32
+# the separate MXFP8 shared expert (weights.mxfp8_a_tiles): 16-row x 128-k tiles,
+# its four tables reached through the ``sh_args`` pointer table (SH_ARGS order)
+F8_TILE = 16 * 128
+SH8_W13_BYTES = W13_ROWS * HIDDEN
+SH8_W2_BYTES = HIDDEN * INTER
+SH8_S13_BYTES = W13_ROWS * W13_SCALE_COLS
+SH8_S2_BYTES = HIDDEN * W2_SCALE_COLS
+SH_ARGS = ("sh_w13", "sh_s13", "sh_w2", "sh_s2")
+SH_W13, SH_S13, SH_W2, SH_S2 = range(len(SH_ARGS))
+
+# The launch goes on the caller's current stream (graph capture included).
+_CURRENT_STREAM = fx.Stream(None)
+
+
+def _f32(x: float) -> float:
+    return struct.unpack("f", struct.pack("f", x))[0]
+
+
+TL_POINTS = 32  # timeline stamps per CTA
+K1_STAMPS = {1: 27, 2: 28, 3: 29, 4: 30, 6: 31}  # fuse_k1: K1's points -> K4's
+RANK_SPLIT = 8  # threads ranking one indexer block
+RANK_LANES = (4, 2, 1)  # butterfly over those RANK_SPLIT lanes
+KEY_NONE = -(2**31)  # the sort key below every block's
+# blocks a thread keys in its split task's eighth of a context (<= MAX_INDEX_BLOCKS)
+CAND_BATCH = 2
+assert N_SPLIT * THREADS * CAND_BATCH >= MAX_INDEX_BLOCKS
+# Indexer context parallelism: a long context's split task p scans index head
+# p // 2's scores of half p % 2 of this rank's blocks -- as many as a TP split task
+# scans -- for the head's rank; a request is long past INDEX_CP_FROM_BLOCKS blocks
+# (``step_rows``; every long request context-parallel: no TP long selection in that
+# build), its rows ending less than a block apart are >= INDEX_CP_FROM_BLOCKS blocks
+# each, and a split task then scans >= TOPK_BLOCKS blocks
+assert N_SPLIT == 2 * TP
+assert N_SPLIT * TOPK_BLOCKS <= INDEX_CP_FROM_BLOCKS <= THREADS
+assert MAX_TOKENS <= SPARSE_BLOCK
+
+
+def scale_index(row, col, cols):
+    """Byte of e8m0 scale (row, col) in aiter's ``shuffle_scale`` layout:
+    ``view(rows/32, 2, 16, cols/8, 2, 4).permute(0, 3, 5, 2, 4, 1)``."""
+    r32, a, b = row // 32, (row // 16) % 2, row % 16
+    c8, d, e = col // 8, (col // 4) % 2, col % 4
+    return ((((r32 * (cols // 8) + c8) * 4 + e) * 16 + b) * 2 + d) * 2 + a
+
+
+def build_post_attn_kernel(
+    npes: int,
+    sm_scale: float,
+    eps: float,
+    route_scale: float,
+    shared_weight: float,
+    swiglu_limit: float,
+    init_blocks: int,
+    local_blocks: int,
+    tokens: int = 1,
+    timeline: bool = False,
+    fuse_k1: bool = False,
+    heads: IndexHeads = ONE_INDEX_HEAD,
+    shared_fp8: bool = False,
+):
+    """``@flyc.jit`` launcher of K4 for one rank of an ``npes``-way TP group and a
+    decode step of ``tokens`` (<= MAX_TOKENS) rows, ``q_len`` consecutive rows a
+    request (a speculative verify's), with contexts up to MAX_CONTEXT.
+    ``fuse_k1``: the layer's K1 runs first in the same launch (``K1_ARGS``,
+    ``positions``, ``slot_mapping``, ``res``; 0 otherwise), for a fused projection
+    of index q ``heads`` (``IndexHeads``). ``shared_fp8``: the expert tables hold
+    the routed experts only and the shared expert is MXFP8, its tables behind
+    the ``sh_args`` pointer table (``SH_ARGS``; 0 otherwise), its GEMVs fp8 x fp8
+    MFMAs."""
+    assert 1 <= tokens <= MAX_TOKENS
+    N_E = N_ROUTED if shared_fp8 else N_ROUTED + 1  # experts in w13 .. s2
+    E_TOP = N_E - 1  # the bound a routed expert id is clamped to
+    S13_BYTES = N_E * W13_ROWS * W13_SCALE_COLS  # e8m0, every expert
+    S2_BYTES = N_E * HIDDEN * W2_SCALE_COLS
+    W = npes
+    G = BLOCKS
+    BASE = stage_bases(tokens)
+    SY = sym_layout(W)
+    # fuse_k1: K1's outputs are read device-coherent (K1 stores them so)
+    CM_K1 = CM_DEV if fuse_k1 else 0
+    WIDE = wide_moe(tokens)
+    POOL = pool_words(tokens)
+    U_MAX = TOP_K * tokens + 1  # distinct routed experts at most, then the shared
+
+    @fx.struct
+    class Smem:
+        # activations, bf16 pairs / fp8 quads: every token's xn, o_proj input or
+        # mid; wide: also the split's PV partials (the stages take turns)
+        x: fx.Array[fx.Int32, POOL, 16]
+        red: fx.Array[fx.Float32, WAVES * 64 * 4, 16]
+        p8: fx.Array[fx.Int32, H * SPLIT_KEYS // 4, 16]
+        q8: fx.Array[fx.Int32, O_K // 4, 16]
+        opart: fx.Array[fx.Float32, 1 if WIDE else WAVES * O_K, 16]  # PV partials
+        # wide MoE: the step's distinct experts (the shared one last), token masks
+        uexp: fx.Array[fx.Int32, U_MAX + 2, 16]
+        umask: fx.Array[fx.Int32, U_MAX + 2, 16]
+        # and per (expert, B column) the down stage's mid row and route weight
+        urow: fx.Array[fx.Int32, U_MAX * 16 if WIDE else 1, 16]
+        uwt: fx.Array[fx.Float32, U_MAX * 16 if WIDE else 1, 16]
+        hst: fx.Array[
+            fx.Float32, 3 * WAVES * H + WAVES, 16
+        ]  # head max / sum, max, v_scale max
+        misc: fx.Array[fx.Float32, O_ROWS * MAX_TOKENS, 16]  # a task's output rows
+        route: fx.Array[fx.Int32, 8 * MAX_TOKENS, 16]  # token k's slots at 8 k
+        rwt: fx.Array[fx.Float32, 8 * MAX_TOKENS, 16]  # and their routing weights
+        # indexer sort keys, a block a thread (a long request: the waves'
+        # candidate (key, block) pairs) and their ranks
+        keys: fx.Array[fx.Int32, THREADS, 16]
+        ranks: fx.Array[fx.Int32, THREADS, 16]
+        blk: fx.Array[fx.Int32, 32, 16]  # this split's 16 pages, n_ctx, tail flag
+        tls: fx.Array[fx.Int64, TL_POINTS if timeline else 1, 16]  # timeline stamps
+        # fuse_k1: index q (every head's: in the pool)
+        qs: fx.Array[fx.Float32, HEAD_DIM // 2 * tokens if heads.count == 1 else 1, 16]
+
+    kernel_name = kernel_symbol(
+        "minimax_m3_mono_layer" if fuse_k1 else "minimax_m3_post_attn",
+        tp=npes,
+        s=tokens,
+        ib=init_blocks,
+        lb=local_blocks,
+        ih=heads.count,
+        io=heads.own,
+        tl=timeline,
+        kps=mono_config.KV_PAGE_STRIDE,
+        sks=mono_config.STATIC_KV_SCALE,
+        sf8=shared_fp8,
+    )
+
+    # the JIT cache key holds the kernel's scalar closure values, not objects: the
+    # index heads go in as ints (a shared cache otherwise serves one rank's build
+    # to every rank)
+    ih_count, ih_own = heads.count, heads.own
+    # fuse_k1 with every index head: K1's index q of every head (qs) after its x8
+    # rows in the pool (K1 is done with the pool before K4's stages use it)
+    QS_CP = tokens * HIDDEN // 4
+    assert not (fuse_k1 and heads.count > 1) or (
+        QS_CP + HEAD_DIM // 2 * tokens * heads.count <= POOL
+    )
+
+    @flyc.kernel(name=kernel_name, known_block_size=[THREADS, 1, 1])
+    def post_attn_kernel(
+        h_in: Int64,
+        q: Int64,
+        block_table: Int64,
+        seq_lens: Int64,
+        k_cache: Int64,
+        v_cache: Int64,
+        k_scale: Int64,
+        v_scale: Int64,
+        w_o: Int64,
+        s_o: Int64,
+        g_post: Int64,
+        w_gate: Int64,
+        bias: Int64,
+        w13: Int64,
+        s13: Int64,
+        w2: Int64,
+        s2: Int64,
+        sh_args: Int64,
+        h_mid: Int64,
+        ar_out: Int64,
+        scratch: Int64,
+        sym: Int64,
+        peers: Int64,
+        step: Int64,
+        rank: Int32,
+        layer: Int32,
+        bt_width: Int32,
+        q_len: Int32,
+        tl: Int64,
+        k1_args: Int64,
+        positions: Int64,
+        slot_mapping: Int64,
+        res: Int64,
+    ):
+        tid = fx.thread_idx.x
+        bid = fx.block_idx.x
+        lane = tid % 64
+        wave = tid // 64
+        g4 = lane // 16
+        l16 = lane % 16
+        lds = fx.SharedAllocator().allocate(Smem).peek()
+        xs = lds.x.ptr
+        red = lds.red.ptr
+        p8 = lds.p8.ptr
+        q8 = lds.q8.ptr
+        # the pool read and written as f32 (wide: the split's PV partials live there)
+        fpool = fx.recast_iter(fx.Float32, xs)
+        opart = fpool if WIDE else lds.opart.ptr
+        uexp = lds.uexp.ptr
+        umask = lds.umask.ptr
+        urow = lds.urow.ptr
+        uwt = lds.uwt.ptr
+        hst = lds.hst.ptr
+        misc = lds.misc.ptr
+        route = lds.route.ptr
+        rwt = lds.rwt.ptr
+        keys = lds.keys.ptr
+        ranks = lds.ranks.ptr
+        blk = lds.blk.ptr
+        tls = lds.tls.ptr
+        v4f = fx.Vector.make_type(4, fx.Float32)
+        v2i = fx.Vector.make_type(2, fx.Int32)
+
+        mbox = Mailbox(scratch, step, layer)
+        # Plain names, not method calls: flydsl's AST rewriter carries every local
+        # whose method is called inside a dynamic if/for as loop state.
+        mb_put = mbox.put
+        mb_put_bf = mbox.put_bf
+        mb_put_words = mbox.put_words
+        mb_poll = mbox.poll
+
+        def mb(name):
+            return scratch + fx.Int64(SCRATCH[name])
+
+        def acquire():
+            fx.memory_fence(
+                syncscope=rocdl.SyncScope.Workgroup, ordering=fx.AtomicOrdering.Acquire
+            )
+
+        def scores_copy(c, k=0):
+            """Copy c of token k's router outputs: expert e's (routing key, sigmoid)
+            at pairs 2 e, 2 e + 1."""
+            return mb("scores") + (fx.Int64(c) * MAX_TOKENS + fx.Int64(k)) * (
+                N_ROUTED * 16
+            )
+
+        def route_key(v, e):
+            """Signed-orderable i32 of score ``v`` with its low 7 bits replaced by
+            127 - expert ``e``: a pick is one integer max, a tie goes to the lower
+            expert."""
+            b = v.bitcast(fx.Int32)
+            b = (b < 0).select(b ^ fx.Int32(0x7FFFFFFF), b)
+            return (b & fx.Int32(-128)) | (127 - e)
+
+        # each wave pushes to one peer
+        pv = fx.Vector(
+            bo.buffer_load(
+                rsrc(peers), fx.min(wave, W - 1) * 2, vec_width=2, dtype=T.i32
+            )
+        )
+        peer_dst = (fx.Int64(uniform(pv[1])) << 32) | fx.Int64(
+            fx.Uint32(uniform(pv[0]))
+        )
+        AG_RS = tokens >= AG_RS_FROM
+
+        def peer_addr(w):
+            pw = fx.Vector(bo.buffer_load(rsrc(peers), w * 2, vec_width=2, dtype=T.i32))
+            return (fx.Int64(uniform(pw[1])) << 32) | fx.Int64(
+                fx.Uint32(uniform(pw[0]))
+            )
+
+        def start(name):
+            return (bid + (G - BASE[name])) % G
+
+        def stamp(k):
+            """``timeline``: s_memrealtime (100 MHz) of this CTA passing point k, kept
+            in LDS until the kernel ends -- a global store per stamp would put its
+            write latency on every later ``s_waitcnt vmcnt(0)``."""
+            # The compile-time and traced conditions cannot share one ``and``.
+            if const_expr(timeline):  # noqa: SIM102
+                if tid == 0:
+                    fx.ptr_store(memrealtime(), tls + k)
+
+        def flush_stamps():
+            if const_expr(timeline):  # noqa: SIM102
+                if tid < TL_POINTS:
+                    fx.generic_store(
+                        fx.inttoptr(
+                            fx.PointerType.get(
+                                fx.Int64.ir_type, fx.AddressSpace.Global, 8
+                            ),
+                            tl + fx.Int64((bid * TL_POINTS + tid) * 8),
+                        ),
+                        fx.ptr_load(tls + tid),
+                    )
+
+        def ld_bf16x4(r, k, cm=0):
+            w = fx.Vector(
+                bo.buffer_load(r, k // 2, vec_width=2, dtype=T.i32, cache_modifier=cm)
+            )
+            v = w.bitcast(fx.BFloat16).to(fx.Float32)
+            return [v[j] for j in range(4)]
+
+        def mfma_fp8(a64, b64, c):
+            return fx.Vector(
+                rocdl.mfma_f32_16x16x32_fp8_fp8(T.vec(4, T.f32), [a64, b64, c, 0, 0, 0])
+            )
+
+        def block_max(v):
+            return block_max_of(v, lane, wave, red)
+
+        def push_rows(region, row0, n_rows, only=None):
+            """This task's bf16 partial rows (misc: token k's n_rows from k n_rows) to
+            the peers' ``region``, a pair a lane, wave w -> peer w (itself too);
+            ``only``: to that peer alone."""
+            half = n_rows // 2
+            to = (wave < W) if only is None else (wave == only)
+            for c in range_constexpr((half * tokens + 63) // 64):
+                i = lane + 64 * c
+                if to & (i < half * tokens):
+                    mb_put_bf(
+                        peer_dst + fx.Int64(SY[region]),
+                        (rank * MAX_TOKENS + i // half) * HIDDEN
+                        + row0
+                        + (i % half) * 2,
+                        [fx.ptr_load(misc + i * 2), fx.ptr_load(misc + (i * 2 + 1))],
+                        CM_SYS,
+                    )
+
+        if const_expr(timeline):  # noqa: SIM102
+            if tid < TL_POINTS:  # same wave as the stamping thread: ordered
+                fx.ptr_store(fx.Int64(0), tls + tid)
+        stamp(0)
+
+        if const_expr(fuse_k1):
+            # this layer's K1 first: its GEMV / head / norm tasks and the scores
+            def k1_ptr(i):
+                v = fx.Vector(
+                    bo.buffer_load(rsrc(k1_args), 2 * i, vec_width=2, dtype=T.i32)
+                )
+                return (fx.Int64(uniform(v[1])) << 32) | fx.Int64(
+                    fx.Uint32(uniform(v[0]))
+                )
+
+            def k1_stamp(k):
+                """K1's timeline points 1-4 and 6 on K4's free stamps 27-31."""
+                if const_expr(k in K1_STAMPS):
+                    stamp(K1_STAMPS[k])
+
+            (
+                ar_in, g_in, w_qkv, s_qkv, g_q, g_k, g_iq, g_ik, cos_sin, index_cache,
+                iq_out, scratch1,
+            ) = [k1_ptr(i) for i in range(len(K1_ARGS))]  # fmt: skip
+            qs_k1 = lds.qs.ptr
+            if const_expr(ih_count > 1):
+                qs_k1 = fpool + QS_CP
+            emit_k1(
+                tid, bid, lane, wave, xs, red, qs_k1, k1_stamp, step, layer,
+                eps, index_scale_log2e(sm_scale), tokens, q_len, init_blocks,
+                local_blocks, IndexHeads(ih_count, ih_own),
+                (
+                    ar_in, res, h_in, g_in, w_qkv, s_qkv, g_q, g_k, g_iq, g_ik,
+                    cos_sin, positions, slot_mapping, k_cache, v_cache, k_scale,
+                    v_scale, index_cache, q, iq_out, block_table, seq_lens,
+                    mb("iscore"), scratch1,
+                ),
+                bt_width,
+                True,
+            )  # fmt: skip
+            gpu.barrier()
+            k1_mbox = Mailbox(scratch1, step, layer)
+            k1_poll = k1_mbox.poll
+            hdone_mb = k1_mbox.addr(K1_SCRATCH_HDONE)
+            rdone_mb = k1_mbox.addr(K1_SCRATCH_RDONE)
+
+        # ============================================================ 0. index
+        # Indexer: K1 scored every block the selection does not pin (index_score);
+        # here only the selection runs.
+        # a graph's pad rows carry seq_len 0 and an all-zero block-table row: as one
+        # key of page 0 they stay finite and select a real page
+        seq_lens_k, n_blks, long_rows = step_rows(
+            seq_lens, tokens, q_len, IndexHeads(ih_count, ih_own)
+        )
+        # Up to TOPK_BLOCKS blocks the top-k keeps every block and the scores would
+        # only order them. The original selector's own contract calls either order
+        # equally correct (index_topk._pack_score_key), so a short context is not
+        # scored and takes the blocks in id order -- the tail block is last either
+        # way -- instead of putting the scoring chain in front of the attention.
+        scored = [nb > TOPK_BLOCKS for nb in n_blks]
+
+        def bt_row(k):
+            """Token k's block-table row (rows are ``bt_width`` apart)."""
+            return block_table + fx.Int64(k) * fx.Int64(bt_width) * 4
+
+        def pick(vals, k):
+            """vals[k] for a traced token index k."""
+            v = vals[0]
+            for i in range_constexpr(1, tokens):
+                v = (k == i).select(vals[i], v)
+            return v
+
+        def index_key(b, word, n_blk):
+            """Sort key of block b's raw score word: forced blocks pinned, then a
+            signed int ordered like the float (-0 == 0, NaN below -inf)."""
+            sc = word.bitcast(fx.Float32)
+            sc = (b < init_blocks).select(fx.Float32(1e30), sc)
+            sc = (b >= n_blk - local_blocks).select(fx.Float32(1e29), sc)
+            bits = sc.bitcast(fx.Int32)
+            bits = (bits == fx.Int32(-(2**31))).select(fx.Int32(0), bits)
+            key = (bits < 0).select(bits ^ 0x7FFFFFFF, bits)
+            return ((bits & 0x7FFFFFFF) > 0x7F800000).select(fx.Int32(-(2**31)), key)
+
+        def rank_blocks(k, iscore, n_blk, r_bt, scored_k):
+            """Up to THREADS blocks: thread b keys block b, and a block's rank
+            counts the blocks that beat it. -> this thread's (key, block, rank,
+            live, the tail block's key, page). ``scored_k``: the scores are read
+            (else the blocks rank by id)."""
+            b = tid
+            live = b < n_blk
+            bb = fx.min(b, n_blk - 1)
+            page = fx.Int32(bo.buffer_load(r_bt, bb, vec_width=1, dtype=T.i32))
+            # K1 (or the harnesses' scorer) wrote every score before this launch, so
+            # a plain load reads them: it goes out before the context length is in.
+            # A pinned block's slot holds no score (index_key forces its key); past
+            # n_blk the keys are never read
+            if const_expr(fuse_k1):
+                # the scorers are CTAs of this launch: poll (a pinned block's slot
+                # polls a scored block's)
+                word = fx.Int32(0)
+                if scored_k:
+                    word = mb_poll(
+                        [
+                            (
+                                iscore,
+                                fx.min(
+                                    fx.max(tid, init_blocks),
+                                    n_blk - 1 - local_blocks,
+                                ),
+                                1,
+                            )
+                        ]
+                    )[0][0]
+            else:
+                word = fx.Int32(
+                    bo.buffer_load(rsrc(iscore), 2 * tid, vec_width=1, dtype=T.i32)
+                )
+            fx.ptr_store(scored_k.select(index_key(tid, word, n_blk), -tid), keys + tid)
+            stamp(18)
+            gpu.barrier()
+            stamp(20)
+            # rank = the keys that beat it. RANK_SPLIT threads share a block, each
+            # comparing every RANK_SPLIT-th key: one lane walking all 61 keys of a
+            # 7.8k context took 1.9 us (a serial compare -> SGPR mask chain)
+            part = tid % RANK_SPLIT
+            for b0 in range(0, n_blk, THREADS // RANK_SPLIT):
+                rb = fx.Int32(b0) + tid // RANK_SPLIT
+                rbc = fx.min(rb, n_blk - 1)
+                kb = fx.ptr_load(keys + rbc)
+                cnt = fx.Int32(0)
+                for j0 in range(0, n_blk, RANK_SPLIT):
+                    j = fx.Int32(j0) + part
+                    kj = fx.ptr_load(keys + fx.min(j, n_blk - 1))
+                    beats = (kj > kb) | ((kj == kb) & (j > rbc))
+                    cnt = cnt + (beats & (j < n_blk)).select(fx.Int32(1), fx.Int32(0))
+                cnt = butterfly(cnt, RANK_LANES)
+                if (part == 0) & (rb < n_blk):
+                    fx.ptr_store(cnt, ranks + rb)
+            gpu.barrier()
+            key = fx.ptr_load(keys + bb)
+            rank = fx.ptr_load(ranks + bb)
+            stamp(19)
+            kt = fx.ptr_load(keys + (n_blk - 1))
+            return key, bb, rank, live, kt, page
+
+        def ballot(pred):
+            return fx.Int64(rocdl.ballot(T.i64, pred))
+
+        def beats(ah, al, bh, bl):
+            """(key, block) pair a orders before b."""
+            return (ah > bh) | ((ah == bh) & (al > bl))
+
+        def choose_pair(key, block, peer_key, peer_block, take_high):
+            partner_high = beats(peer_key, peer_block, key, block)
+            high_h = partner_high.select(peer_key, key)
+            high_l = partner_high.select(peer_block, block)
+            low_h = partner_high.select(key, peer_key)
+            low_l = partner_high.select(block, peer_block)
+            return take_high.select(high_h, low_h), take_high.select(high_l, low_l)
+
+        def merge16_best(ah, al, bh, bl):
+            """Return the top half of two descending 16-item lists."""
+            b_hi = beats(bh, bl, ah, al)
+            key, block = b_hi.select(bh, ah), b_hi.select(bl, al)
+            for off in (8, 4, 2, 1):
+                ph, pl = xshfl(key, off), xshfl(block, off)
+                p_hi = beats(ph, pl, key, block)
+                hi_h, hi_l = p_hi.select(ph, key), p_hi.select(pl, block)
+                lo_h, lo_l = p_hi.select(key, ph), p_hi.select(block, pl)
+                take_hi = (l16 & off) == 0
+                key = take_hi.select(hi_h, lo_h)
+                block = take_hi.select(hi_l, lo_l)
+            return key, block
+
+        def sort16_desc(key, block):
+            for size, strides in (
+                (2, (1,)),
+                (4, (2, 1)),
+                (8, (4, 2, 1)),
+                (16, (8, 4, 2, 1)),
+            ):
+                for stride in strides:
+                    ph, pl = xshfl(key, stride), xshfl(block, stride)
+                    asc = (l16 & size) != 0
+                    low = (l16 & stride) == 0
+                    key, block = choose_pair(key, block, ph, pl, low ^ asc)
+            return key, block
+
+        def wave_best(h0, l0, h1, l1):
+            """Return the exact top 16 of 128 pairs in lanes 0..15."""
+            h0, l0 = sort16_desc(h0, l0)
+            h1, l1 = sort16_desc(h1, l1)
+            h0, l0 = merge16_best(h0, l0, xshfl(h0, 31), xshfl(l0, 31))
+            h1, l1 = merge16_best(h1, l1, xshfl(h1, 31), xshfl(l1, 31))
+            h0, l0 = merge16_best(h0, l0, xshfl(h0, 47), xshfl(l0, 47))
+            h1, l1 = merge16_best(h1, l1, xshfl(h1, 47), xshfl(l1, 47))
+            h0, l0 = merge16_best(
+                h0,
+                l0,
+                h1.shuffle_xor(15, 64),
+                l1.shuffle_xor(15, 64),
+            )
+            return h0, l0
+
+        def rank_cands():
+            """Merge eight sorted candidate rows to the exact top 16."""
+            key, block = fx.Int32(KEY_NONE), fx.Int32(-1)
+            if wave == 0:
+                group = lane // TOPK_BLOCKS
+                index = l16
+                ai = (2 * group) * TOPK_BLOCKS + index
+                bi = (2 * group + 1) * TOPK_BLOCKS + (TOPK_BLOCKS - 1 - index)
+                key, block = merge16_best(
+                    fx.ptr_load(keys + 2 * ai),
+                    fx.ptr_load(keys + (2 * ai + 1)),
+                    fx.ptr_load(keys + 2 * bi),
+                    fx.ptr_load(keys + (2 * bi + 1)),
+                )
+                key, block = merge16_best(key, block, xshfl(key, 31), xshfl(block, 31))
+                key, block = merge16_best(key, block, xshfl(key, 47), xshfl(block, 47))
+            return key, block, tid, (tid < TOPK_BLOCKS) & (block >= 0)
+
+        def place(key, bb, rank, live, kt, n_blk, seq_len, active):
+            """The ranked blocks' slots, as the Triton selector emits them: full
+            blocks by rank, the tail block (the one holding the current token)
+            last. -> (selected, slot, sparse context length); ``active``: this
+            CTA places (a uniform flag)."""
+            n_sel = fx.min(n_blk, TOPK_BLOCKS)
+            tail = n_blk - 1
+            tail_first = (kt > key) | ((kt == key) & (tail > bb))
+            sel = live & (rank < n_sel) & active
+            is_tail = sel & (bb == tail)
+            if is_tail:
+                fx.ptr_store(fx.Int32(1), blk + 17)
+            gpu.barrier()
+            tail_sel = fx.ptr_load(blk + 17) != 0
+            n_full = n_sel - tail_sel.select(fx.Int32(1), fx.Int32(0))
+            slot = is_tail.select(
+                n_full, rank - tail_first.select(fx.Int32(1), fx.Int32(0))
+            )
+            n_ctx = tail_sel.select(
+                n_full * SPARSE_BLOCK + seq_len - tail * SPARSE_BLOCK,
+                fx.min(n_sel * SPARSE_BLOCK, seq_len),
+            )
+            return sel, slot, n_ctx
+
+        def table_of(k):
+            return rsrc(mb("sparse_table") + fx.Int64(k) * (SPARSE_TABLE_WORDS * 4))
+
+        def select_blocks(k, t):
+            """blk[0:16] := the pages of selection slots 2t, 2t+1 and blk[16] := the
+            sparse context length, exactly as the Triton selector emits them: top
+            TOPK_BLOCKS by (score, block id) descending with the init blocks pinned
+            to 1e30 and the local ones to 1e29, full blocks in that order and the
+            tail block (the one holding the current token) last. A block's slot
+            follows from its rank, the count of blocks that beat it. k: token.
+            For a long request ``stage_select_long`` left blk: a one-block
+            stand-in is ranked, placing nothing."""
+            iscore = mb("iscore") + fx.Int64(k) * (MAX_INDEX_BLOCKS * 8)
+            short = ~pick(long_rows, k)
+            n_blk = short.select(pick(n_blks, k), fx.Int32(1))
+            seq_len = pick(seq_lens_k, k)
+            r_bt = rsrc(bt_row(k))
+            table = table_of(k)
+            if (tid < 2 * PAGES_PER_BLOCK + 2) & short:
+                fx.ptr_store(fx.Int32(0), blk + tid)
+            key, bb, rank, live, kt, page = rank_blocks(
+                k, iscore, n_blk, r_bt, pick(scored, k) & short
+            )
+            sel, slot, n_ctx = place(key, bb, rank, live, kt, n_blk, seq_len, short)
+            if sel & (slot // 2 == t):
+                for j in range_constexpr(PAGES_PER_BLOCK):
+                    fx.ptr_store(
+                        page * mono_config.KV_PAGE_STRIDE + j,
+                        blk + ((slot % 2) * PAGES_PER_BLOCK + j),
+                    )
+            if sel & (t == 0):
+                for j in range_constexpr(PAGES_PER_BLOCK):
+                    bo.buffer_store(
+                        page * mono_config.KV_PAGE_STRIDE + j,
+                        table,
+                        slot * PAGES_PER_BLOCK + j,
+                    )
+            if (tid == 0) & short:
+                fx.ptr_store(n_ctx, blk + 16)
+                if t == 0:
+                    bo.buffer_store(n_ctx, table, TOPK_BLOCKS * PAGES_PER_BLOCK)
+            gpu.barrier()
+
+        def select_long(k, part):
+            """Split task (k, part) of a long request's token, before the
+            split stage. Its share of the blocks (a thread keys CAND_BATCH) -> each
+            wave's TOPK_BLOCKS best -> the CTA's, to a candidate mailbox; then the
+            N_SPLIT shares' candidates (the context's TOPK_BLOCKS best are among
+            them) are ranked and placed -> blk as ``select_blocks`` leaves it, and
+            (part 0) the sparse table. The share: an eighth of the blocks, to this
+            rank's sel_long; with every index head (indexer context parallelism)
+            head part // 2's scores of half part % 2 of this rank's blocks, to the
+            head's rank's sel_cp -- the shares its own head's candidates come in."""
+            n_blk = pick(n_blks, k)
+            iscore = mb("iscore") + fx.Int64(k) * (MAX_INDEX_BLOCKS * 8)
+            first = k * (N_SPLIT * TOPK_BLOCKS)  # token k's candidates (pair pairs)
+            table = table_of(k)
+            if tid == 0:  # the tail flag ``place`` reads: LDS, not zeroed on this path
+                fx.ptr_store(fx.Int32(0), blk + 17)
+            if const_expr(ih_count > 1):
+                own0 = init_blocks + (ih_own - init_blocks) % TP  # this rank's first
+                n_own = (n_blk - own0 + TP - 1) // TP
+                # the last scored one: the local blocks' slots hold no score
+                last = own0 + TP * ((n_blk - local_blocks - own0 + TP - 1) // TP - 1)
+                span = (n_own + 1) // 2
+                lo = (part % 2) * span
+                hi = fx.min(lo + span, n_own)
+                js = [lo + tid + THREADS * i for i in range(CAND_BATCH)]
+                blocks = [own0 + TP * j for j in js]
+                live = [j < hi for j in js]
+                src = iscore + fx.Int64(part // 2) * (MAX_INDEX_BLOCKS // TP * 8)
+                slots = [fx.min(b, last) // TP for b in blocks]
+                out, out_slot = peer_addr(part // 2) + fx.Int64(SY["sel_cp"]), 2 * rank
+                out_slot = out_slot + part % 2
+                got_from, got_scope = sym + fx.Int64(SY["sel_cp"]), "system"
+            else:
+                span = (n_blk + N_SPLIT - 1) // N_SPLIT
+                lo = part * span
+                hi = fx.min(lo + span, n_blk)
+                blocks = [lo + tid + THREADS * i for i in range(CAND_BATCH)]
+                live = [b < hi for b in blocks]
+                src = iscore
+                slots = [
+                    fx.min(fx.max(b, init_blocks), n_blk - 1 - local_blocks)
+                    for b in blocks
+                ]
+                out, out_slot = mb("sel_long"), part
+                got_from, got_scope = out, "agent"
+            if const_expr(fuse_k1):
+                words = [w[0] for w in mb_poll([(src, sl, 1) for sl in slots])]
+            else:
+                words = [
+                    fx.Int32(
+                        bo.buffer_load(rsrc(src), 2 * sl, vec_width=1, dtype=T.i32)
+                    )
+                    for sl in slots
+                ]
+            hs = [
+                lv.select(index_key(b, w, n_blk), fx.Int32(KEY_NONE))
+                for b, w, lv in zip(blocks, words, live)
+            ]
+            ls = [lv.select(b, fx.Int32(-1)) for b, lv in zip(blocks, live)]
+            stamp(18)
+            rh, rl = wave_best(hs[0], ls[0], hs[1], ls[1])
+            if lane < TOPK_BLOCKS:
+                fx.ptr_store(rh, keys + 2 * (wave * TOPK_BLOCKS + lane))
+                fx.ptr_store(rl, keys + (2 * (wave * TOPK_BLOCKS + lane) + 1))
+            gpu.barrier()
+            ch, cl, cnt, held = rank_cands()
+            # a share holds >= TOPK_BLOCKS blocks: its 16 best are blocks of
+            # distinct ranks
+            if held & (cnt < TOPK_BLOCKS):
+                mb_put_words(
+                    out,
+                    2 * (first + out_slot * TOPK_BLOCKS + cnt),
+                    [ch, cl],
+                    CM_SYS if ih_count > 1 else CM_DEV,
+                )
+            gpu.barrier()
+            stamp(20)
+            if tid < N_SPLIT * TOPK_BLOCKS:
+                got = mb_poll([(got_from, 2 * (first + tid), 2)], got_scope)[0]
+                fx.ptr_store(got[0], keys + 2 * tid)
+                fx.ptr_store(got[1], keys + (2 * tid + 1))
+            gpu.barrier()
+            ch, cl, cnt, held = rank_cands()
+            page = fx.Int32(
+                bo.buffer_load(
+                    rsrc(bt_row(k), MAX_INDEX_BLOCKS * 4),
+                    cl,
+                    vec_width=1,
+                    dtype=T.i32,
+                    mask=held,
+                )
+            )
+            stamp(19)
+            # the tail block is a local one: its key is forced
+            kt = index_key(n_blk - 1, fx.Int32(0), n_blk)
+            sel, slot, n_ctx = place(
+                ch, cl, cnt, held, kt, n_blk, pick(seq_lens_k, k), pick(long_rows, k)
+            )
+            if sel & (slot // 2 == part):
+                for j in range_constexpr(PAGES_PER_BLOCK):
+                    fx.ptr_store(
+                        page * mono_config.KV_PAGE_STRIDE + j,
+                        blk + ((slot % 2) * PAGES_PER_BLOCK + j),
+                    )
+            if sel & (part == 0):
+                for j in range_constexpr(PAGES_PER_BLOCK):
+                    bo.buffer_store(
+                        page * mono_config.KV_PAGE_STRIDE + j,
+                        table,
+                        slot * PAGES_PER_BLOCK + j,
+                    )
+            if tid == 0:
+                fx.ptr_store(n_ctx, blk + 16)
+                if part == 0:
+                    bo.buffer_store(n_ctx, table, TOPK_BLOCKS * PAGES_PER_BLOCK)
+            gpu.barrier()
+
+        def stage_select_long():
+            """The split tasks of a long request's token select before the
+            split stage (``select_long``), each leaving its blk for its split task
+            (a CTA runs at most one). Out of the split stage and cold: inline
+            there, this code slowed the split for every context (64k: +1 us)."""
+            ts = start("split")
+            k = ts // N_SPLIT
+            if unlikely((ts < N_SPLIT * tokens) & pick(long_rows, k)):
+                select_long(k, ts % N_SPLIT)
+
+        # Each stage is its own function: flydsl carries a local reassigned inside a
+        # dynamic if / for as region state, so names must not be shared across stages.
+        # ============================================================ 1. split
+        def stage_split():
+            r_k, r_v = rsrc(k_cache), rsrc(v_cache)
+            r_ks, r_vs = rsrc(k_scale), rsrc(v_scale)
+            PPW = SPLIT_KEYS // WAVES // PAGE16  # pages per wave
+            for ts in range(start("split"), N_SPLIT * tokens, G):
+                ts = fx.Int32(ts)  # token tok's partition t: ts = tok N_SPLIT + t
+                tok = ts // N_SPLIT
+                t = ts % N_SPLIT
+                select_blocks(tok, t)
+                stamp(17)
+                if const_expr(fuse_k1):
+                    # this token's q heads and every token's new K / V: a request's
+                    # tokens of one step (speculative decode) attend to each other's
+                    if wave == 0:
+                        kv = lane - H
+                        k1_poll(
+                            [
+                                (
+                                    hdone_mb,
+                                    (lane < H).select(
+                                        tok * K1_HEAD_TASKS + lane,
+                                        fx.min(kv // 2, tokens - 1) * K1_HEAD_TASKS
+                                        + H
+                                        + kv % 2,
+                                    ),
+                                    1,
+                                )
+                            ]
+                        )
+                    gpu.barrier()
+                    acquire()
+                # Every independent load goes out first: this wave's pages and q.
+                # Pages past the context read junk masked below.
+                n_ctx = uniform(fx.ptr_load(blk + 16))
+                pages = [
+                    uniform(fx.ptr_load(blk + (wave * PPW + j))) for j in range(PPW)
+                ]
+                qv = ld_bf16x4(rsrc(q), tok * O_K + tid * 4, CM_K1)
+                # K (QK B operand, key = lane % 16 of page j), V (PV B operand: dims
+                # 16 jd + lane % 16, keys 8 (lane / 16) .. + 8), per-token scales
+                kw = []
+                for j in range_constexpr(PPW):
+                    for s in range_constexpr(HEAD_DIM // 32):
+                        kw.append(
+                            fx.Vector(
+                                bo.buffer_load(
+                                    r_k,
+                                    (
+                                        pages[j] * PAGE_BYTES
+                                        + (2 * s + g4 // 2) * 256
+                                        + l16 * 16
+                                        + (g4 % 2) * 8
+                                    )
+                                    // 4,
+                                    vec_width=2,
+                                    dtype=T.i32,
+                                    cache_modifier=CM_K1,
+                                )
+                            ).bitcast(fx.Int64)[0]
+                        )
+                vpg = (g4 // 2 == 0).select(pages[0], pages[1])
+                vw = [
+                    fx.Vector(
+                        bo.buffer_load(
+                            r_v,
+                            (vpg * PAGE_BYTES + (16 * jd + l16) * PAGE16 + (g4 % 2) * 8)
+                            // 4,
+                            vec_width=2,
+                            dtype=T.i32,
+                            cache_modifier=CM_K1,
+                        )
+                    ).bitcast(fx.Int64)[0]
+                    for jd in range(HEAD_DIM // 16)
+                ]
+                kss = [
+                    fx.Float32(
+                        bo.buffer_load(
+                            r_ks,
+                            pages[j] * PAGE16 + l16,
+                            vec_width=1,
+                            dtype=T.f32,
+                            cache_modifier=CM_K1,
+                        )
+                    )
+                    for j in range(PPW)
+                ]
+                vss = [
+                    fx.Float32(
+                        bo.buffer_load(
+                            r_vs,
+                            pages[j] * PAGE16 + l16,
+                            vec_width=1,
+                            dtype=T.f32,
+                            cache_modifier=CM_K1,
+                        )
+                    )
+                    for j in range(PPW)
+                ]
+                fx.ptr_store(fp8_pack4(qv[0], qv[1], qv[2], qv[3]), q8 + tid)
+                gpu.barrier()
+                stamp(12)
+                # QK per page: A=q (heads), B=K (keys); a lane holds heads
+                # 4*g4+e and key l16.
+                key0 = t * SPLIT_KEYS + wave * (PPW * PAGE16)
+                valid = [key0 + j * PAGE16 + l16 < n_ctx for j in range(PPW)]
+                sc = []
+                for j in range_constexpr(PPW):
+                    c = fx.Vector.filled(4, 0.0, fx.Float32)
+                    for s in range_constexpr(HEAD_DIM // 32):
+                        qw = fx.Vector(
+                            fx.ptr_load(
+                                q8 + (l16 * 32 + 8 * s + 2 * g4), result_type=v2i
+                            )
+                        )
+                        c = mfma_fp8(
+                            qw.bitcast(fx.Int64)[0], kw[j * (HEAD_DIM // 32) + s], c
+                        )
+                    qk = sm_scale * kss[j]
+                    sc.append(
+                        [valid[j].select(qk * c[e], fx.Float32(NEG)) for e in range(4)]
+                    )
+                # partition max per head and max v_scale: in-wave butterfly, then LDS
+                vloc = fx.max(
+                    valid[0].select(vss[0], fx.Float32(0.0)),
+                    valid[1].select(vss[1], fx.Float32(0.0)),
+                )
+                vmax_w = wave_max(vloc)
+                for e in range_constexpr(4):
+                    mh = butterfly(fx.max(sc[0][e], sc[1][e]), (8, 4, 2, 1), fx.max)
+                    if l16 == 0:
+                        fx.ptr_store(mh, hst + (wave * H + g4 * 4 + e))
+                if lane == 0:
+                    fx.ptr_store(vmax_w, hst + (3 * WAVES * H + wave))
+                gpu.barrier()
+                stamp(13)
+                # every wave partial is read before the first p8 store below: LDS
+                # stores in between serialize each read behind its own wait
+                vmax = fx.ptr_load(hst + 3 * WAVES * H)
+                for w in range_constexpr(1, WAVES):
+                    vmax = fx.max(vmax, fx.ptr_load(hst + (3 * WAVES * H + w)))
+                mxs = []
+                for e in range_constexpr(4):
+                    mx = fx.ptr_load(hst + (g4 * 4 + e))
+                    for w in range_constexpr(1, WAVES):
+                        mx = fx.max(mx, fx.ptr_load(hst + (w * H + g4 * 4 + e)))
+                    mxs.append(mx)
+                fscale = FP8_MAX * hw_rcp(vmax + 1e-8)
+                pscale = vmax * (1.0 / FP8_MAX)
+                for e in range_constexpr(4):
+                    head = g4 * 4 + e
+                    mx = mxs[e]
+                    ls = fx.Float32(0.0)
+                    for j in range_constexpr(PPW):
+                        pj = hw_exp2((sc[j][e] - mx) * LOG2E)
+                        ls = ls + pj
+                        b8 = (
+                            fp8_pack4(
+                                (valid[j].select(vss[j], fx.Float32(0.0)) * fscale)
+                                * pj,
+                                0.0,
+                                0.0,
+                                0.0,
+                            )
+                            & 0xFF
+                        )
+                        w8 = b8 | (xshfl(b8, 1) << 8)
+                        w8 = w8 | (xshfl(w8, 2) << 16)
+                        if l16 % 4 == 0:
+                            kloc = wave * (PPW * PAGE16) + j * PAGE16 + l16
+                            fx.ptr_store(
+                                w8, p8 + (head * (SPLIT_KEYS // 4) + kloc // 4)
+                            )
+                    ls = butterfly(ls, (8, 4, 2, 1))
+                    if l16 == 0:
+                        fx.ptr_store(ls, hst + (WAVES * H + wave * H + head))
+                        fx.ptr_store(mx, hst + (2 * WAVES * H + head))
+                gpu.barrier()
+                stamp(14)
+                # PV over this wave's 32 keys: A = P (heads), B = V (dims)
+                pw = fx.Vector(
+                    fx.ptr_load(
+                        p8
+                        + (
+                            l16 * (SPLIT_KEYS // 4)
+                            + (wave * (PPW * PAGE16) + 8 * g4) // 4
+                        ),
+                        result_type=v2i,
+                    )
+                ).bitcast(fx.Int64)[0]
+                # keys past the context carry P = 0, but a NaN V byte would still poison
+                nv = n_ctx - (key0 + 8 * g4)
+                vmask = (nv >= 8).select(
+                    fx.Int64(-1),
+                    (nv <= 0).select(
+                        fx.Int64(0), (fx.Int64(1) << (fx.Int64(nv) * 8)) - 1
+                    ),
+                )
+                for jd in range_constexpr(HEAD_DIM // 16):
+                    cv = mfma_fp8(
+                        pw, vw[jd] & vmask, fx.Vector.filled(4, 0.0, fx.Float32)
+                    )
+                    for e in range_constexpr(4):
+                        fx.ptr_store(
+                            cv[e],
+                            opart
+                            + (wave * O_K + (g4 * 4 + e) * HEAD_DIM + 16 * jd + l16),
+                        )
+                gpu.barrier()
+                stamp(15)
+                # Sum wave partials: acc=prob_scale*PV, then /exp_sum, bf16.
+                e0 = tid * 4
+                head = e0 // HEAD_DIM
+                lsum = fx.ptr_load(hst + (WAVES * H + head))
+                for w in range_constexpr(1, WAVES):
+                    lsum = lsum + fx.ptr_load(hst + (WAVES * H + w * H + head))
+                inv_l = 1.0 / (lsum > 0.0).select(lsum, fx.Float32(1.0))
+                ov = []
+                for k in range_constexpr(4):
+                    acc = fx.ptr_load(opart + (e0 + k))
+                    for w in range_constexpr(1, WAVES):
+                        acc = acc + fx.ptr_load(opart + (w * O_K + e0 + k))
+                    ov.append((pscale * acc) * inv_l)
+                mb_put_bf(mb("sp_o"), ts * O_K + e0, ov)
+                if tid < H:
+                    mb_put(
+                        mb("sp_m"),
+                        ts * H + tid,
+                        fx.ptr_load(hst + (2 * WAVES * H + tid)),
+                    )
+                    lt = fx.ptr_load(hst + (WAVES * H + tid))
+                    for w in range_constexpr(1, WAVES):
+                        lt = lt + fx.ptr_load(hst + (WAVES * H + w * H + tid))
+                    mb_put(mb("sp_l"), ts * H + tid, lt)
+                gpu.barrier()
+
+        stage_select_long()
+        stage_split()
+        stamp(1)
+
+        def merge_token(k):
+            """Token k's partitions merged (the gluon decode's reduce; thread = 4 dims
+            of one head) and quantized per token (standalone quant: scale = amax *
+            (1 / FP8_MAX), x * rcp(scale)) -> (this thread's fp8 word, scale)."""
+            hm = tid // (HEAD_DIM // 4)
+            e0 = tid * 4
+            sp0 = k * N_SPLIT  # token k's partitions
+            # wave 0 waits on every partition's l, then the CTA polls once: 512
+            # threads x 24 loads spinning through the whole split load the memory
+            # system the split itself reads (-0.4 us)
+            if wave == 0:
+                mb_poll(
+                    [
+                        (mb("sp_l"), sp0 * H + lane + 64 * i, 1)
+                        for i in range(N_SPLIT * H // 64)
+                    ]
+                )
+            gpu.barrier()
+            # one batch: every partition's (m, l) and this thread's 4 outputs
+            got = mb_poll(
+                [(mb("sp_m"), (sp0 + i) * H + hm, 1) for i in range(N_SPLIT)]
+                + [(mb("sp_l"), (sp0 + i) * H + hm, 1) for i in range(N_SPLIT)]
+                + [
+                    (mb("sp_o"), ((sp0 + i) * O_K + e0) // 2, 2) for i in range(N_SPLIT)
+                ],
+                batch=3 * N_SPLIT,
+            )
+            mls = got[: 2 * N_SPLIT]
+            ovs = got[2 * N_SPLIT :]
+            stamp(3)
+            m_i = [mls[i][0].bitcast(fx.Float32) for i in range(N_SPLIT)]
+            l_i = [mls[N_SPLIT + i][0].bitcast(fx.Float32) for i in range(N_SPLIT)]
+            gm = m_i[0]
+            for i in range_constexpr(1, N_SPLIT):
+                gm = fx.max(gm, m_i[i])
+            sl = [l_i[i] * hw_exp2((m_i[i] - gm) * LOG2E) for i in range(N_SPLIT)]
+            gs = sl[0]
+            for i in range_constexpr(1, N_SPLIT):
+                gs = gs + sl[i]
+            gs = (gs > 0.0).select(gs, fx.Float32(1.0))
+            av = [fx.Float32(0.0) for _ in range(4)]
+            for i in range_constexpr(N_SPLIT):
+                wi = sl[i] / gs
+                o01 = bf2_f32(ovs[i][0])
+                o23 = bf2_f32(ovs[i][1])
+                av = [
+                    av[0] + wi * o01[0],
+                    av[1] + wi * o01[1],
+                    av[2] + wi * o23[0],
+                    av[3] + wi * o23[1],
+                ]
+            a0, a1, a2, a3 = (bf16_round(x) for x in av)
+            amax = block_max(
+                fx.max(
+                    fx.max(fmath.absf(a0), fmath.absf(a1)),
+                    fx.max(fmath.absf(a2), fmath.absf(a3)),
+                )
+            )
+            x_scale = amax * (1.0 / FP8_MAX)
+            # an all-zero row (a pad row) quantizes to zeros, not 0 * inf
+            inv = (amax == 0.0).select(fx.Float32(0.0), hw_rcp(x_scale))
+            return fp8_pack4(a0 * inv, a1 * inv, a2 * inv, a3 * inv), x_scale
+
+        # ======================================= 2b. attention merge (S > 1)
+        # One task per token, in parallel on its own CTA: merged in the o tasks, the
+        # tokens' merges ran one after another in every one of them (+14 us at S = 4)
+        def stage_merge():
+            if const_expr(tokens > 1):
+                for k in range(start("merge"), tokens, G):
+                    k = fx.Int32(k)
+                    word, x_scale = merge_token(k)
+                    # plain words, then the scale: its pair is the flag
+                    bo.buffer_store(
+                        word,
+                        rsrc(mb("attnq")),
+                        k * (O_K // 4) + tid,
+                        cache_modifier=CM_DEV,
+                    )
+                    fx.memory_fence(
+                        syncscope=rocdl.SyncScope.Workgroup,
+                        ordering=fx.AtomicOrdering.Release,
+                    )
+                    gpu.barrier()
+                    if tid == 0:
+                        mb_put(mb("attnq_s"), k, x_scale)
+
+        stage_merge()
+
+        # ================================================ 3. o_proj + attn reduce
+        def stage_o():
+            r_wo = rsrc(w_o)
+            O_CPW = O_K // 64 // 4  # 64-k chunks per wave: 4 waves per 16-row group
+            for t in range(start("o"), N_O, G):
+                t = fx.Int32(t)
+                rg = t * 2 + wave // 4
+                wts = []
+                for cc in range_constexpr(O_CPW):
+                    kc = (wave % 4) * O_CPW + cc
+                    wts.append(
+                        fx.Vector(
+                            bo.buffer_load(
+                                r_wo,
+                                ((rg * (O_K // 32) + kc * 2) * 512 + lane * 16) // 4,
+                                vec_width=4,
+                                dtype=T.i32,
+                            )
+                        )
+                    )
+                # this task's row scales, before the wait
+                ws = fx.Float32(
+                    bo.buffer_load(
+                        rsrc(s_o), t * O_ROWS + tid % O_ROWS, vec_width=1, dtype=T.f32
+                    )
+                )
+
+                if const_expr(tokens == 1):
+                    # the merge folded in: the partitions -> this CTA's fp8 input
+                    word, x_scale = merge_token(0)
+                    x_scales = [x_scale]
+                    fx.ptr_store(word, xs + tid)
+                else:
+                    # the merge tasks' fp8 inputs: every token's scale (published
+                    # after its plain words), then the words, 16 B a load
+                    got = mb_poll([(mb("attnq_s"), k, 1) for k in range(tokens)])
+                    stamp(3)
+                    x_scales = [got[k][0].bitcast(fx.Float32) for k in range(tokens)]
+                    acquire()
+                    n_v = tokens * O_K // 16
+                    for i in range_constexpr((n_v + THREADS - 1) // THREADS):
+                        v = tid + THREADS * i
+                        if v < n_v:
+                            fx.ptr_store(
+                                fx.Vector(
+                                    bo.buffer_load(
+                                        rsrc(mb("attnq")),
+                                        v * 4,
+                                        vec_width=4,
+                                        dtype=T.i32,
+                                        cache_modifier=CM_DEV,
+                                    )
+                                ),
+                                xs + (v * 4 // (O_K // 4) * O_ROW + v * 4 % (O_K // 4)),
+                            )
+                gpu.barrier()
+                # B column l % 16 is token l % 16 (token 0 past the last): C column k
+                # is token k's rows, the weight stream is the same for any token count
+                col_tok = fx.min(l16, tokens - 1)
+                c = fx.Vector.filled(4, 0.0, fx.Float32)
+                # chunks kc, kc + 1 in one 16x16x128 FP8 MFMA (unit scales): its lane
+                # holds k 16 (l / 16) .. + 16 of each 64-k half -- the chunks' packing
+                for p in range_constexpr(O_CPW // 2):
+                    xb = [
+                        fx.Vector(
+                            fx.ptr_load(
+                                xs
+                                + (
+                                    col_tok * O_ROW
+                                    + ((wave % 4) * O_CPW + 2 * p + h) * 16
+                                    + g4 * 4
+                                ),
+                                result_type=v4f,
+                            )
+                        ).bitcast(fx.Int32)
+                        for h in range(2)
+                    ]
+                    c = fx.Vector(
+                        rocdl.mfma_scale_f32_16x16x128_f8f6f4(
+                            T.vec(4, T.f32),
+                            [
+                                fx.Vector.from_elements(
+                                    [
+                                        wts[2 * p + h][e]
+                                        for h in range(2)
+                                        for e in range(4)
+                                    ],
+                                    fx.Int32,
+                                ),
+                                fx.Vector.from_elements(
+                                    [xb[h][e] for h in range(2) for e in range(4)],
+                                    fx.Int32,
+                                ),
+                                c,
+                                0,
+                                0,
+                                0,
+                                fx.Int32(127),
+                                0,
+                                fx.Int32(127),
+                            ],
+                        )
+                    )
+                fx.ptr_store(c, red + (wave * 64 + lane) * 4)
+                gpu.barrier()
+                if tid < O_ROWS * tokens:
+                    tk = tid // O_ROWS
+                    grp = (tid % O_ROWS) // 16
+                    rr = tid % 16
+                    tot = fx.Float32(0.0)
+                    for w in range_constexpr(4):
+                        tot = tot + fx.ptr_load(
+                            red
+                            + (((grp * 4 + w) * 64 + 16 * (rr // 4) + tk) * 4 + rr % 4)
+                        )
+                    fx.ptr_store(bf16_round(tot * pick(x_scales, tk) * ws), misc + tid)
+                gpu.barrier()
+                # push bf16 partial rows, wave w -> peer w (to itself too); with AG_RS
+                # to the row group's owner rank only
+                push_rows("attn", t * O_ROWS, O_ROWS, t % W if AG_RS else None)
+                gpu.barrier()
+                if const_expr(AG_RS):  # noqa: SIM102
+                    # the owner sums every peer's partials in rank order and gathers
+                    # the sums to every rank's a_ag (h_mid: the router tasks)
+                    if (t % W == rank) & (tid < O_ROWS // 2 * tokens):
+                        ptk = tid // (O_ROWS // 2)
+                        row = t * O_ROWS + (tid % (O_ROWS // 2)) * 2
+                        own = sym + fx.Int64(SY["attn"])
+                        prt = mb_poll(
+                            [
+                                (own, ((src * MAX_TOKENS + ptk) * HIDDEN + row) // 2, 1)
+                                for src in range(W)
+                            ],
+                            "system",
+                        )
+                        t0 = fx.Float32(0.0)
+                        t1 = fx.Float32(0.0)
+                        for src in range_constexpr(W):
+                            f0, f1 = bf2_f32(prt[src][0])
+                            t0 = t0 + f0
+                            t1 = t1 + f1
+                        for w in range_constexpr(W):
+                            mb_put_bf(
+                                peer_addr(w) + fx.Int64(SY["a_ag"]),
+                                ptk * HIDDEN + row,
+                                [t0, t1],
+                                CM_SYS,
+                            )
+                if const_expr(not AG_RS):  # noqa: SIM102
+                    # this task's rows of the all-reduce, finished here (GLM's
+                    # peer_reduce): sum peer partials in rank order, bf16, then
+                    # add the residual into h_mid and the 'a' mailbox. A router
+                    # reads 24 KB device-scope, not 4 peers' 48 KB system-scope.
+                    if tid < O_ROWS // 2 * tokens:
+                        ptk = tid // (O_ROWS // 2)
+                        row = t * O_ROWS + (tid % (O_ROWS // 2)) * 2
+                        if const_expr(tokens != 1):
+                            if const_expr(fuse_k1):
+                                k1_poll([(rdone_mb, ptk, 1)])
+                                acquire()
+                            hr = fx.Int32(
+                                bo.buffer_load(
+                                    rsrc(h_in),
+                                    (ptk * HIDDEN + row) // 2,
+                                    vec_width=1,
+                                    dtype=T.i32,
+                                    cache_modifier=CM_K1,
+                                )
+                            )
+                        own = sym + fx.Int64(SY["attn"])
+                        prt = mb_poll(
+                            [
+                                (own, ((src * MAX_TOKENS + ptk) * HIDDEN + row) // 2, 1)
+                                for src in range(W)
+                            ],
+                            "system",
+                        )
+                        t0 = fx.Float32(0.0)
+                        t1 = fx.Float32(0.0)
+                        for src in range_constexpr(W):
+                            f0, f1 = bf2_f32(prt[src][0])
+                            t0 = t0 + f0
+                            t1 = t1 + f1
+                        # the rank sum only: the router adds the residual in f32, so the
+                        # norm reads bf16(sum) plus the unrounded residual.
+                        mb_put_bf(mb("a"), ptk * HIDDEN + row, [t0, t1])
+                        if const_expr(tokens != 1):
+                            r0, r1 = bf2_f32(hr)
+                            a0 = bf16_round(t0) + r0
+                            a1 = bf16_round(t1) + r1
+                            bo.buffer_store(
+                                fx.Vector.from_elements([a0, a1], fx.Float32).to(
+                                    fx.BFloat16
+                                ),
+                                rsrc(h_mid),
+                                ptk * HIDDEN + row,
+                            )
+
+        stage_o()
+        stamp(4)
+
+        # ======================================================= MoE helpers
+        def e8m0_load(r_s, byte_idx):
+            """An e8m0 scale byte, still in flight (the MFMA takes it as is)."""
+            return fx.Int32(bo.buffer_load(r_s, byte_idx, vec_width=1, dtype=T.i8))
+
+        def group_scale_byte(expert, rows, rg, kc, cols):
+            """scale_index(expert rows + 16 rg + l16, 4 kc + g4, cols): this lane's
+            byte for row group rg, 128-k chunk kc, folded to wave-uniform terms plus
+            the lane (the general form's per-lane arithmetic cost the S = 16 up /
+            gate loop ~8 us)."""
+            base = (expert * (rows // 32) + rg // 2) * (cols // 8) + kc // 2
+            return (base * 64 + lane) * 4 + (kc % 2) * 2 + rg % 2
+
+        def route_top4(k, k0, sc0, k1, sc1, sig, wave_mask=None):
+            """route[8k : 8k+4] / rwt[..] := token k's sigmoid top-k gating from the
+            router tasks' routing keys ``k0`` / ``k1`` (experts ``lane`` / ``lane +
+            64``, sigmoid + bias, see route_key) and unbiased sigmoids (wave k % WAVES
+            computes it, looking the sigmoids up in LDS at ``sig``): a pick is one
+            wave max; the weights are the picks' sigmoids renormalized and scaled by
+            route_scale. Slot 4 is the fused shared expert. ``wave_mask``: this
+            wave's expert -> token mask, the picks' bit k set (the picks are
+            distinct, and the wave's tokens run in program order). The caller
+            syncs."""
+            if wave == k % WAVES:
+                # the unbiased weights, looked up by expert once the picks are known
+                fx.ptr_store(sc0, sig + (k * N_ROUTED + lane))
+                fx.ptr_store(sc1, sig + (k * N_ROUTED + 64 + lane))
+                pid = fx.Int32(0)
+                for p in range_constexpr(TOP_K):
+                    m = butterfly(fx.max(k0, k1), (32, 16, 8, 4, 2, 1), fx.max)
+                    hit0 = k0 == m
+                    k0 = hit0.select(fx.Int32(-(2**31)), k0)
+                    k1 = ((k1 == m) & (k0 != m)).select(fx.Int32(-(2**31)), k1)
+                    pid = (lane == p).select(127 - (m & 127), pid)
+                if lane < TOP_K:
+                    w = fx.ptr_load(sig + (k * N_ROUTED + pid))
+                    tot = butterfly(w, (1, 2))
+                    fx.ptr_store(pid, route + (8 * k + lane))
+                    fx.ptr_store(
+                        w * (route_scale / fx.max(tot, 1e-20)), rwt + (8 * k + lane)
+                    )
+                    if const_expr(wave_mask is not None):
+                        was = fx.Int32(fx.ptr_load(wave_mask + pid))
+                        fx.ptr_store(was | (fx.Int32(1) << k), wave_mask + pid)
+                if lane == 0:
+                    fx.ptr_store(fx.Int32(SHARED_EXPERT), route + (8 * k + TOP_K))
+                    fx.ptr_store(fx.Float32(shared_weight), rwt + (8 * k + TOP_K))
+
+        # a8w4 operands in xs (Int32 words, never through f32 arithmetic): the up /
+        # gate stage holds every token's MXFP8 xn then its E8M0 scales; the down
+        # stage every token's MXFP8 mid then its scales
+        XN8_WORDS = HIDDEN // 4
+        XSC_WORDS = HIDDEN // 32
+        XSC_LDS = tokens * XN8_ROW
+        MID16_WORDS = MOE_SLOTS * INTER // 2  # bf16 pairs, as the mids land
+        MID8_WORDS = MOE_SLOTS * INTER // 4
+        MID8_LDS = tokens * MID16_WORDS
+        MIDSC_WORDS = MOE_SLOTS * INTER // 32
+        MIDSC_LDS = MID8_LDS + tokens * MID8_WORDS
+
+        def mx8_scale(vals, offsets):
+            """E8M0 byte (>= 1) of the 32-block the lanes xor-``offsets`` apart share,
+            aiter's MXFP8 RoundUp rule ceil_pow2(amax / 448), and its reciprocal."""
+            amax = fx.max(vals[0], -vals[0])
+            for v in vals[1:]:
+                amax = fx.max(amax, fx.max(v, -v))
+            amax = butterfly(amax, offsets, fx.max)
+            bits = (amax * fx.Float32(1.0 / FP8_MAX)).bitcast(fx.Int32)
+            e = (bits >> 23) & 0xFF
+            e = e + ((bits & 0x7FFFFF) != 0).select(fx.Int32(1), fx.Int32(0))
+            e = fx.max(fx.min(e, fx.Int32(254)), fx.Int32(1))
+            return e, hw_rcp((e << 23).bitcast(fx.Float32))
+
+        def mx8_q(v, inv):
+            return fx.max(fx.min(v * inv, fx.Float32(FP8_MAX)), fx.Float32(-FP8_MAX))
+
+        def lds_b8(chunk_word):
+            """This lane's B operand of the 128-k chunk at LDS word ``chunk_word``:
+            fp8 k 16 g .. + 16 then 64 + 16 g .. + 16 (g = lane // 16), the
+            f8f6f4 MFMA's B layout (harness/mfma_f4f8_probe.py)."""
+            h = [
+                fx.Vector(
+                    fx.ptr_load(xs + (chunk_word + 16 * j + 4 * g4), result_type=v4f)
+                ).bitcast(fx.Int32)
+                for j in range(2)
+            ]
+            return fx.Vector.from_elements(
+                [h[j][e] for j in range(2) for e in range(4)], fx.Int32
+            )
+
+        def bf16x2(word):
+            return fx.Vector.from_elements([word], fx.Int32).bitcast(fx.BFloat16)
+
+        def lds_i32(word):
+            return fx.ptr_load(xs + word).bitcast(fx.Int32)
+
+        def mfma_f4f8(w4, sa, b8, sb, c, sa_byte=None):
+            """16x16x128 scaled MFMA: A = this lane's 32 fp4 (row l % 16, k block
+            l / 16) with E8M0 ``sa``, B = its 32 fp8 with E8M0 ``sb``. ``sa_byte``:
+            ``sa`` is a scale dword, the MFMA takes that byte of it."""
+            if const_expr(sa_byte is None):
+                sa, sa_byte = sa & 0xFF, 0
+            return fx.Vector(
+                rocdl.mfma_scale_f32_16x16x128_f8f6f4(
+                    T.vec(4, T.f32), [w4, b8, c, 4, 0, sa_byte, sa, 0, sb]
+                )
+            )
+
+        def mfma_f8f8(w8, sa, b8, sb, c, sa_byte=None):
+            """``mfma_f4f8`` with an fp8 A: this lane's 32 fp8 of row l % 16 in the B
+            operand's k order (``f8_tile_load``); the scales are placed alike."""
+            if const_expr(sa_byte is None):
+                sa, sa_byte = sa & 0xFF, 0
+            return fx.Vector(
+                rocdl.mfma_scale_f32_16x16x128_f8f6f4(
+                    T.vec(4, T.f32), [w8, b8, c, 0, 0, sa_byte, sa, 0, sb]
+                )
+            )
+
+        def f8_tile_load(r_w, tile, slot):
+            """Lane ``slot``'s 32 fp8 of 16 x 128 tile ``tile`` (weights.mxfp8_a_tiles)
+            as the f8f6f4 MFMA's 8-dword A operand."""
+            word = (tile * F8_TILE + slot * 32) // 4
+            h = [
+                fx.Vector(bo.buffer_load(r_w, word + 4 * j, vec_width=4, dtype=T.i32))
+                for j in range(2)
+            ]
+            return fx.Vector.from_elements(
+                [h[j][e] for j in range(2) for e in range(4)], fx.Int32
+            )
+
+        def sh_ptr(i):
+            """The MXFP8 shared expert's table ``SH_ARGS[i]``, read at use from the
+            ``sh_args`` pointer table (as K1's are): as four more kernel arguments
+            they widened the kernarg pointer tuple, which ROCm LLVM's partial-
+            register rewrite spilled incompletely (a fault at S = 2)."""
+            v = fx.Vector(
+                bo.buffer_load(rsrc(sh_args), 2 * i, vec_width=2, dtype=T.i32)
+            )
+            return (fx.Int64(uniform(v[1])) << 32) | fx.Int64(fx.Uint32(uniform(v[0])))
+
+        def w_tile(r_w, rg, kc, slot, k, fp8):
+            """Lane ``slot``'s A operand of row group rg, 128-k chunk kc of a table
+            of k columns: fp4 (two 16 x 64 tiles of 512 B, shuffle_weight(16, 16))
+            or MXFP8 (``f8_tile_load``)."""
+            if const_expr(fp8):
+                return f8_tile_load(r_w, rg * (k // 128) + kc, slot)
+            return fx.Vector(
+                bo.buffer_load(
+                    r_w,
+                    ((rg * (k // 64) + kc * 2) * 512 + slot * 16) // 4,
+                    vec_width=4,
+                    dtype=T.i32,
+                )
+            )
+
+        mfma_sh = mfma_f8f8 if shared_fp8 else mfma_f4f8  # the shared expert's
+
+        # ============================================================ 4. router
+        def stage_router():
+            r_g = rsrc(g_post)
+            r_gate = rsrc(w_gate)
+            # this thread's 12 consecutive elements of the row (512 x 12 = HIDDEN):
+            # the poll is 3 specs per rank, 12 in one batch, every thread busy
+            RE = HIDDEN // THREADS
+            rw0 = tid * (RE // 2)
+
+            def ld_bf16x12(r, word, cm=0):
+                """12 bf16 from ``word`` (a 16 B then an 8 B load)."""
+                a8 = fx.Vector(
+                    bo.buffer_load(r, word, vec_width=4, dtype=T.i32, cache_modifier=cm)
+                ).bitcast(fx.BFloat16)
+                a4 = fx.Vector(
+                    bo.buffer_load(
+                        r, word + 4, vec_width=2, dtype=T.i32, cache_modifier=cm
+                    )
+                ).bitcast(fx.BFloat16)
+                return [a8[e] for e in range(8)] + [a4[e] for e in range(4)]
+
+            # one task per (token, 8 experts): the tokens' norms and dot products run
+            # on different CTAs, not one after another in each (S = 4: 74.7 -> 64.1 us)
+            for tt in range(start("router"), N_ROUTER * tokens, G):
+                tt = fx.Int32(tt)
+                k = tt // N_ROUTER
+                t = tt % N_ROUTER
+                e_r = t * WAVES + wave
+                bias_r = fx.Float32(
+                    bo.buffer_load(rsrc(bias), e_r, vec_width=1, dtype=T.f32)
+                )
+                gw = [  # bf16 pairs
+                    fx.Vector(
+                        bo.buffer_load(
+                            r_gate,
+                            (e_r * HIDDEN + (i * 64 + lane) * 8) // 2,
+                            vec_width=4,
+                            dtype=T.i32,
+                        )
+                    )
+                    for i in range(HIDDEN // 512)
+                ]
+                gp = ld_bf16x12(r_g, rw0)
+                if const_expr(fuse_k1):
+                    # the residual (K1's norm / GEMV task)
+                    if wave == 0:
+                        k1_poll([(rdone_mb, k, 1)])
+                    gpu.barrier()
+                    acquire()
+                hres = ld_bf16x12(rsrc(h_in), k * (HIDDEN // 2) + rw0, CM_K1)
+                # the all-reduced post-attention row (the o tasks finished it; with
+                # AG_RS the owner ranks gathered it into this rank's a_ag)
+                a_src, a_scope = (
+                    (sym + fx.Int64(SY["a_ag"]), "system")
+                    if AG_RS
+                    else (mb("a"), "agent")
+                )
+                arow = mb_poll(
+                    [
+                        (a_src, (k * HIDDEN + tid * RE) // 2 + 2 * h, 2)
+                        for h in range(RE // 4)
+                    ],
+                    a_scope,
+                )
+                stamp(22)
+                v = [
+                    f + fx.Float32(hres[4 * h + 2 * qq + i])
+                    for h in range(RE // 4)
+                    for qq in range(2)
+                    for i, f in enumerate(bf2_f32(arow[h][qq]))
+                ]
+                if const_expr(AG_RS or tokens == 1):  # noqa: SIM102
+                    # h_mid = bf16(bf16(sum) + residual): task t writes slice t
+                    if tid // (XN_SLICE // RE) == t:
+                        hw = (k * HIDDEN + tid * RE) // 2
+                        hb = fx.Vector.from_elements(v, fx.Float32).to(fx.BFloat16)
+                        hi = hb.bitcast(fx.Int32)
+                        bo.buffer_store(
+                            fx.Vector.from_elements(
+                                [hi[e] for e in range(4)], fx.Int32
+                            ),
+                            rsrc(h_mid),
+                            hw,
+                        )
+                        bo.buffer_store(
+                            fx.Vector.from_elements([hi[4], hi[5]], fx.Int32),
+                            rsrc(h_mid),
+                            hw + 4,
+                        )
+                # the row's sum of squares: a wave sum each, one barrier, every lane
+                # adds the 8
+                ssw = fx.Float32(0.0)
+                for e in range_constexpr(RE):
+                    ssw = fmath.fma(v[e], v[e], ssw)
+                ssw = wave_sum(ssw)
+                if lane == 0:
+                    fx.ptr_store(ssw, red + wave)
+                gpu.barrier()
+                ss = fx.ptr_load(red)
+                for w in range_constexpr(1, WAVES):
+                    ss = ss + fx.ptr_load(red + w)
+                stamp(24)
+                rstd = hw_rsq(ss / float(HIDDEN) + eps)
+                for j in range_constexpr(RE // 2):
+                    fx.ptr_store(
+                        bf16_pair(
+                            v[2 * j] * rstd * (fx.Float32(gp[2 * j]) + 1.0),
+                            v[2 * j + 1] * rstd * (fx.Float32(gp[2 * j + 1]) + 1.0),
+                        ),
+                        xs + (rw0 + j),
+                    )
+                gpu.barrier()
+                stamp(25)
+                if tid < XN_SLICE // 4:  # 4 elements a lane, 8 lanes a 32-block
+                    w0 = t * (XN_SLICE // 2) + tid * 2
+                    f = [
+                        v
+                        for j in range(2)
+                        for v in bf2_f32(fx.ptr_load(xs + (w0 + j)).bitcast(fx.Int32))
+                    ]
+                    e, inv = mx8_scale(f, (1, 2, 4))
+                    word = fp8_pack4(*[mx8_q(v, inv) for v in f])
+                    xw = k * XN8_WORDS + t * (XN_SLICE // 4) + tid
+                    xsw = k * XSC_WORDS + t * (XN_SLICE // 32) + tid // 8
+                    if const_expr(WIDE):
+                        # plain words, flagged once the task is done (below)
+                        bo.buffer_store(
+                            word, rsrc(mb("xn8p")), xw, cache_modifier=CM_DEV
+                        )
+                        if tid % 8 == 0:
+                            bo.buffer_store(
+                                e, rsrc(mb("xscp")), xsw, cache_modifier=CM_DEV
+                            )
+                    else:
+                        mb_put_words(mb("xn8"), xw, [word])
+                        if tid % 8 == 0:
+                            mb_put_words(mb("xsc"), xsw, [e])
+                stamp(26)
+                # bf16 pair dot products, one accumulator per pair position: four
+                # 12-long chains instead of one 96-long FMA chain behind 192 widenings
+                accs = [fx.Float32(0.0) for _ in range(4)]
+                for i in range_constexpr(HIDDEN // 512):
+                    xv = fx.Vector(
+                        fx.ptr_load(xs + (i * 64 + lane) * 4, result_type=v4f)
+                    ).bitcast(fx.Int32)
+                    for j in range_constexpr(4):
+                        accs[j] = fx.Float32(
+                            rocdl.fdot2_f32_bf16(
+                                T.f32,
+                                as_ir_value(bf16x2(gw[i][j])),
+                                as_ir_value(bf16x2(xv[j])),
+                                as_ir_value(accs[j]),
+                                clamp=False,
+                            ).result
+                        )
+                logit = wave_sum((accs[0] + accs[1]) + (accs[2] + accs[3]))
+                stamp(23)
+                # the routing key and sigmoid go out, not the logit: the consumers'
+                # routing starts at the max rounds
+                sc = hw_rcp(1.0 + hw_exp2(-bf16_round(logit) * LOG2E))
+                rk = route_key(sc + bias_r, e_r)
+                if lane < SCORE_COPIES:
+                    mb_put_words(
+                        scores_copy(lane, k), e_r * 2, [rk, sc.bitcast(fx.Int32)]
+                    )
+                if const_expr(WIDE):
+                    # this task's xn slice is out: its flag (after the logits, not
+                    # in front of them)
+                    fx.memory_fence(
+                        syncscope=rocdl.SyncScope.Workgroup,
+                        ordering=fx.AtomicOrdering.Release,
+                    )
+                    gpu.barrier()
+                    if tid == 0:
+                        mb_put(mb("xndone"), tt, fx.Int32(1))
+
+        stage_router()
+        stamp(5)
+
+        # ================================ 5. experts: up / gate -> down -> FFN reduce
+        UG_KC = HIDDEN // 128  # 128-k chunks
+        UG_CPW = UG_KC // WAVES  # 128-k chunks per wave: each wave a k eighth
+        DN_KC = INTER // 128  # 128-k chunks of one slot
+        DN_UNITS = TOP_K * DN_KC  # routed (slot, 128-k chunk) units per row group
+        DN_UPW = (DN_UNITS + 3) // 4  # 4 waves per row group
+        SH_UPW = (DN_KC + 3) // 4  # the shared expert's chunks, once for all tokens
+        # routed (token, slot, row group) tasks t < ROUTED_TASKS; the shared expert's
+        # row groups (one GEMV with a column per token) are tasks ROUTED_TASKS + j,
+        # run as task j of stage "shared" while the routing is still on its way, so
+        # the routed tasks alone set the rounds (S = 4: 768 = 3 x 256)
+        ROUTED_TASKS = TOP_K * UG_PER_SLOT * tokens
+        TPT = TOP_K * UG_PER_SLOT  # routed tasks per token
+        CPT = ug_ctas_per_token(tokens)
+        UG_ITERS = max(ug_tasks_of(c, tokens) for c in range(G))  # per CTA, at most
+
+        def ug_task_id(it):
+            """(this CTA's it-th routed task, clamped; whether it exists)."""
+            if const_expr(CPT is not None):
+                u = bid % CPT + it * CPT
+                return (bid // CPT) * TPT + fx.min(u, TPT - 1), u < TPT
+            t = start("ug") + it * G
+            return fx.min(t, ROUTED_TASKS - 1), t < ROUTED_TASKS
+
+        def ug_task(t):
+            """(shared expert?, token, slot, row group) of up / gate task t."""
+            shared = t >= ROUTED_TASKS
+            r = fx.min(t, ROUTED_TASKS - 1)
+            tk = r // (TOP_K * UG_PER_SLOT)
+            slot = shared.select(
+                fx.Int32(TOP_K), (r % (TOP_K * UG_PER_SLOT)) // UG_PER_SLOT
+            )
+            return shared, tk, slot, t % UG_PER_SLOT
+
+        def ug_load(t):
+            """Up / gate weights and scales of task t."""
+            shared, tk, slot, cg = ug_task(t)
+            # the bound keeps a corrupt route from addressing past the expert table
+            # (a shared task never gets here: ug_task_id clamps to the routed tasks
+            # and sh_load runs the shared expert)
+            e_u = shared.select(
+                fx.Int32(min(SHARED_EXPERT, E_TOP)),
+                fx.min(uniform(fx.ptr_load(route + (8 * tk + slot))), E_TOP),
+            )
+            r_w = rsrc(w13 + fx.Int64(e_u) * W13_BYTES)
+            r_s = rsrc(s13)
+            wts: list[list] = []  # [gate, up][chunk]
+            scs: list[list] = []
+            for gu in range_constexpr(2):
+                rg = gu * UG_PER_SLOT + cg
+                wts.append([])
+                scs.append([])
+                for i in range_constexpr(UG_CPW):
+                    kc = wave * UG_CPW + i
+                    wts[gu].append(
+                        fx.Vector(
+                            bo.buffer_load(
+                                r_w,
+                                ((rg * (HIDDEN // 64) + kc * 2) * 512 + lane * 16) // 4,
+                                vec_width=4,
+                                dtype=T.i32,
+                            )
+                        )
+                    )
+                    scs[gu].append(
+                        e8m0_load(
+                            r_s,
+                            group_scale_byte(e_u, W13_ROWS, rg, kc, W13_SCALE_COLS),
+                        )
+                    )
+            return wts, scs
+
+        def ug_compute(t, wts, scs):
+            """Up / gate GEMV of task t -> swiglu -> its 16 mid values (of every token
+            for the shared expert: B column l % 16 is token l % 16)."""
+            shared, tk, slot, cg = ug_task(t)
+            col = shared.select(fx.min(l16, tokens - 1), tk)
+            cg_ = fx.Vector.filled(4, 0.0, fx.Float32)
+            cu_ = fx.Vector.filled(4, 0.0, fx.Float32)
+            # every chunk's B scale up front: read at use, each MFMA waited out an LDS
+            # latency (S = 4 up / gate 3.5 us over the unscaled floor)
+            sbs = [
+                lds_i32(XSC_LDS + col * XSC_ROW + (wave * UG_CPW + i) * 4 + g4)
+                for i in range(UG_CPW)
+            ]
+            for i in range_constexpr(UG_CPW):
+                kc = wave * UG_CPW + i
+                xb = lds_b8(col * XN8_ROW + kc * 32)
+                cg_ = mfma_f4f8(wts[0][i], scs[0][i], xb, sbs[i], cg_)
+                cu_ = mfma_f4f8(wts[1][i], scs[1][i], xb, sbs[i], cu_)
+            stamp(9)
+            # gate partials in red, up partials in opart (free once the split is done)
+            fx.ptr_store(cg_, red + (wave * 64 + lane) * 4)
+            fx.ptr_store(cu_, opart + (wave * 64 + lane) * 4)
+            gpu.barrier()
+            # row tid % 16; a shared task's thread tid reads token tid / 16's column
+            if tid < shared.select(fx.Int32(UG_ROWS * tokens), fx.Int32(UG_ROWS)):
+                r = tid % UG_ROWS
+                col = tid // UG_ROWS
+                kk = shared.select(col, tk)
+                gsum = fx.Float32(0.0)
+                usum = fx.Float32(0.0)
+                for w in range_constexpr(WAVES):
+                    gsum = gsum + fx.ptr_load(
+                        red + ((w * 64 + 16 * (r // 4) + col) * 4 + r % 4)
+                    )
+                    usum = usum + fx.ptr_load(
+                        opart + ((w * 64 + 16 * (r // 4) + col) * 4 + r % 4)
+                    )
+                y = swiglu_mul_batch([gsum], [usum], fx.Float32(-swiglu_limit))[0]
+                mb_put(
+                    mb("mid"),
+                    (kk * MOE_SLOTS + slot) * INTER + cg * UG_ROWS + r,
+                    bf16_round(y),
+                )
+            gpu.barrier()
+
+        def sh_load(j, has=None):
+            """The shared expert's half task j: gate rows 8 j .. + 8 on lanes l16 < 8,
+            the up rows alike on the rest -- one MFMA tile; the waves split k."""
+            row = (l16 < 8).select(SH_ROWS * j + l16, INTER + SH_ROWS * j + l16 - 8)
+            rg = row // 16
+            if const_expr(shared_fp8):
+                w_base, s_base = sh_ptr(SH_W13), sh_ptr(SH_S13)
+                w_bytes, s_bytes, e_rows = SH8_W13_BYTES, SH8_S13_BYTES, 0
+            else:
+                w_base, s_base = w13 + fx.Int64(SHARED_EXPERT) * W13_BYTES, s13
+                w_bytes, s_bytes = W13_BYTES, S13_BYTES
+                e_rows = SHARED_EXPERT * W13_ROWS
+            if const_expr(has is None):
+                r_w = rsrc(w_base)
+                r_s = rsrc(s_base)
+            else:  # a CTA without the task reads through zero-sized buffers
+                r_w = rsrc(w_base, has.select(fx.Int32(w_bytes), fx.Int32(0)))
+                r_s = rsrc(s_base, has.select(fx.Int32(s_bytes), fx.Int32(0)))
+            wts = []
+            scs = []
+            # Row offset; each chunk adds 256 * (kc // 2) + 2 * (kc % 2).
+            sc_row = scale_index(e_rows + row, g4, W13_SCALE_COLS)
+            for i in range_constexpr(UG_CPW):
+                kc = wave * UG_CPW + i
+                wts.append(w_tile(r_w, rg, kc, g4 * 16 + row % 16, HIDDEN, shared_fp8))
+                scs.append(e8m0_load(r_s, sc_row + kc // 2 * 256 + kc % 2 * 2))
+            return wts, scs
+
+        def sh_compute(j, wts, scs):
+            """Half task j's GEMV (B column l % 16 = token l % 16) -> swiglu -> the
+            shared slot's mid rows 8 j .. + 8 of every token."""
+            col = fx.min(l16, tokens - 1)
+            c = fx.Vector.filled(4, 0.0, fx.Float32)
+            sbs = [
+                lds_i32(XSC_LDS + col * XSC_ROW + (wave * UG_CPW + i) * 4 + g4)
+                for i in range(UG_CPW)
+            ]
+            for i in range_constexpr(UG_CPW):
+                kc = wave * UG_CPW + i
+                c = mfma_sh(wts[i], scs[i], lds_b8(col * XN8_ROW + kc * 32), sbs[i], c)
+            fx.ptr_store(c, red + (wave * 64 + lane) * 4)
+            gpu.barrier()
+            # mid row r = tid % 8 of token tid / 8: gate in tile row r, up in r + 8
+            if tid < SH_ROWS * tokens:
+                r = tid % SH_ROWS
+                col = tid // SH_ROWS
+                gsum = fx.Float32(0.0)
+                usum = fx.Float32(0.0)
+                for w in range_constexpr(WAVES):
+                    gsum = gsum + fx.ptr_load(
+                        red + ((w * 64 + 16 * (r // 4) + col) * 4 + r % 4)
+                    )
+                    usum = usum + fx.ptr_load(
+                        red + ((w * 64 + 16 * (r // 4 + 2) + col) * 4 + r % 4)
+                    )
+                y = swiglu_mul_batch([gsum], [usum], fx.Float32(-swiglu_limit))[0]
+                mb_put(
+                    mb("mid"),
+                    (col * MOE_SLOTS + TOP_K) * INTER + j * SH_ROWS + r,
+                    bf16_round(y),
+                )
+            gpu.barrier()
+
+        def down_load(k, td, has_dn):
+            """Token k's route-weighted down units of row group td; a CTA without a
+            down task reads through zero-sized buffers."""
+            dn_bytes = has_dn.select(fx.Int32(W2_BYTES), fx.Int32(0))
+            dn_scale_bytes = has_dn.select(fx.Int32(S2_BYTES), fx.Int32(0))
+            rg_d = td * 2 + wave // 4
+            units = []
+            for i in range_constexpr(DN_UPW):
+                j = (wave % 4) + 4 * i
+                jj = fx.min(j, DN_UNITS - 1)
+                sl = jj // DN_KC
+                kc = jj % DN_KC
+                e_d = fx.min(uniform(fx.ptr_load(route + (8 * k + sl))), E_TOP)
+                wv = fx.Vector(
+                    bo.buffer_load(
+                        rsrc(w2 + fx.Int64(e_d) * W2_BYTES, dn_bytes),
+                        ((rg_d * (INTER // 64) + kc * 2) * 512 + lane * 16) // 4,
+                        vec_width=4,
+                        dtype=T.i32,
+                    )
+                )
+                sc = e8m0_load(
+                    rsrc(s2, dn_scale_bytes),
+                    group_scale_byte(e_d, HIDDEN, rg_d, kc, W2_SCALE_COLS),
+                )
+                wgt = (j < DN_UNITS).select(
+                    fx.ptr_load(rwt + (8 * k + sl)), fx.Float32(0.0)
+                )
+                units.append((wv, sc, sl, kc, wgt))
+            return units
+
+        def down_load_shared(td, has_dn):
+            """The shared expert's down units of row group td (every token's)."""
+            rg_d = td * 2 + wave // 4
+            if const_expr(shared_fp8):
+                w_base, s_base, w_bytes, s_bytes = (
+                    sh_ptr(SH_W2), sh_ptr(SH_S2), SH8_W2_BYTES, SH8_S2_BYTES,
+                )  # fmt: skip
+                e = 0
+            else:
+                w_base, s_base, w_bytes, s_bytes = (
+                    w2 + fx.Int64(SHARED_EXPERT) * W2_BYTES, s2, W2_BYTES, S2_BYTES,
+                )  # fmt: skip
+                e = SHARED_EXPERT
+            r_w = rsrc(w_base, has_dn.select(fx.Int32(w_bytes), fx.Int32(0)))
+            r_s = rsrc(s_base, has_dn.select(fx.Int32(s_bytes), fx.Int32(0)))
+            units = []
+            for i in range_constexpr(SH_UPW):
+                j = (wave % 4) + 4 * i
+                kc = fx.min(j, DN_KC - 1)
+                wv = w_tile(r_w, rg_d, kc, lane, INTER, shared_fp8)
+                sc = e8m0_load(
+                    r_s, group_scale_byte(e, HIDDEN, rg_d, kc, W2_SCALE_COLS)
+                )
+                wgt = (j < DN_KC).select(fx.Float32(shared_weight), fx.Float32(0.0))
+                units.append((wv, sc, kc, wgt))
+            return units
+
+        def ffn_finish(t):
+            """Misc (row group t's bf16 partial rows, token k's DN_ROWS from k DN_ROWS)
+            -> every peer -> rank-ordered sum -> ar_out."""
+            # every token's rows: reduce threads take token tid // 16, pair tid % 16
+            row_tok = tid // (DN_ROWS // 2)
+            row_pair = tid % (DN_ROWS // 2)
+            if const_expr(W > 1):  # push bf16 partial rows: wave w -> peer w
+                # (with AG_RS to the row group's owner rank only)
+                push_rows("ffn", t * DN_ROWS, DN_ROWS, t % W if AG_RS else None)
+                gpu.barrier()
+            if tid < DN_ROWS // 2 * tokens:
+                row = t * DN_ROWS + row_pair * 2
+                if const_expr(W == 1):
+                    t0 = fx.ptr_load(misc + tid * 2)
+                    t1 = fx.ptr_load(misc + (tid * 2 + 1))
+                else:
+                    own = sym + fx.Int64(SY["ffn"])
+
+                    def rank_sum():
+                        parts = mb_poll(
+                            [
+                                (
+                                    own,
+                                    ((src * MAX_TOKENS + row_tok) * HIDDEN + row) // 2,
+                                    1,
+                                )
+                                for src in range(W)
+                            ],
+                            "system",
+                        )
+                        s0 = fx.Float32(0.0)
+                        s1 = fx.Float32(0.0)
+                        for src in range_constexpr(W):
+                            f0, f1 = bf2_f32(parts[src][0])
+                            s0 = s0 + f0
+                            s1 = s1 + f1
+                        return s0, s1
+
+                    if const_expr(AG_RS):
+                        t0 = fx.Float32(0.0)
+                        t1 = fx.Float32(0.0)
+                        if t % W == rank:  # the owner: sum, gather to every rank
+                            t0, t1 = rank_sum()
+                            for w in range_constexpr(W):
+                                mb_put_bf(
+                                    peer_addr(w) + fx.Int64(SY["ffn_ag"]),
+                                    row_tok * HIDDEN + row,
+                                    [t0, t1],
+                                    CM_SYS,
+                                )
+                        else:  # the owner's sums (bf16, as ar_out stores them)
+                            got = mb_poll(
+                                [
+                                    (
+                                        sym + fx.Int64(SY["ffn_ag"]),
+                                        (row_tok * HIDDEN + row) // 2,
+                                        1,
+                                    )
+                                ],
+                                "system",
+                            )
+                            t0, t1 = bf2_f32(got[0][0])
+                    else:
+                        t0, t1 = rank_sum()
+                bo.buffer_store(
+                    fx.Vector.from_elements([t0, t1], fx.Float32).to(fx.BFloat16),
+                    rsrc(ar_out),
+                    row_tok * HIDDEN + row,
+                )
+
+        # ======================================== 5w. experts, wide (S > 4)
+        # LDS pool words: MXFP8 xn, its scales, the up partials, router sigmoids; the
+        # down stage reuses the pool for MXFP8 mid rows (token k's slot s at row
+        # k MOE_SLOTS + s, a zero row last) and their scales
+        W_UPP = tokens * (XN8_ROW + XSC_ROW)
+        W_SIG = W_UPP + WAVES * 64 * 4
+        MID_ROWS = MOE_SLOTS * tokens  # the zero row's index
+        MID_ROW_WORDS = INTER // 4
+        MIDSC_ROW_WORDS = INTER // 32
+        W_MIDSC = (MID_ROWS + 1) * MID_ROW
+        # up / gate tasks go in pairs, a CTA runs both: an MXFP8 block is 32 rows
+        UG_PAIRS = UG_PER_SLOT // 2
+        UGW_ITERS = 2 * ((U_MAX * UG_PAIRS + G - 1) // G)  # up / gate tasks a CTA runs
+        # distinct experts the expert-table loops run (shared_fp8: the shared
+        # expert's MXFP8 tasks / units follow them, compiled apart)
+        U_TBL = U_MAX - 1 if shared_fp8 else U_MAX
+        # an MXFP8 shared task runs as SH_PARTS k parts, a part's weights in flight
+        # as many as a table task's: the whole task's and the next's were 192 VGPRs
+        # (256 with spills, and the table's tasks ran slower too)
+        SH_PARTS = 2
+        assert UG_CPW % SH_PARTS == 0
+
+        def sh_ks(part):
+            n = UG_CPW // SH_PARTS
+            return tuple(range(part * n, (part + 1) * n))
+
+        DN_WITERS = (U_TBL * DN_KC + 3) // 4  # (expert, chunk) units a wave runs
+        DN_DEPTH = 12  # units whose weights are in flight
+        ROUTE_ITERS = (tokens + WAVES - 1) // WAVES  # tokens a wave routes
+
+        def popc64(x):
+            return fx.Int32(fx.ctpop(x))
+
+        def slot_of(k, e):
+            """Token k's slot routed to expert e (TOP_K: the shared expert)."""
+            s = fx.Int32(TOP_K)
+            for j in range_constexpr(TOP_K):
+                s = (fx.Int32(fx.ptr_load(route + (8 * k + j))) == e).select(
+                    fx.Int32(j), s
+                )
+            return s
+
+        def build_experts():
+            """Uexp / umask[0 .. n_u) := the tokens' distinct routed experts in
+            expert order and their token masks, [n_u] := the shared expert (every
+            token), uexp[U_MAX + 1] := n_u. The waves' masks must be in (synced)."""
+            m = fx.Int32(0)
+            if tid < N_ROUTED:
+                for w in range_constexpr(WAVES):
+                    m = m | fx.Int32(fx.ptr_load(xs + (W_UPP + w * N_ROUTED + tid)))
+            bal = ballot(m != 0)
+            below = popc64(bal & ((fx.Int64(1) << fx.Int64(lane)) - 1))
+            if tid == 0:
+                fx.ptr_store(popc64(bal), umask + (U_MAX + 1))
+            gpu.barrier()
+            if tid < N_ROUTED:
+                off = (wave == 0).select(
+                    fx.Int32(0), fx.Int32(fx.ptr_load(umask + (U_MAX + 1)))
+                )
+                rank_ = off + below
+                if m != 0:
+                    fx.ptr_store(fx.Int32(tid), uexp + rank_)
+                    fx.ptr_store(m, umask + rank_)
+                if tid == N_ROUTED - 1:
+                    nu_ = off + popc64(bal)
+                    fx.ptr_store(nu_, uexp + (U_MAX + 1))
+                    fx.ptr_store(fx.Int32(SHARED_EXPERT), uexp + nu_)
+                    fx.ptr_store(fx.Int32((1 << tokens) - 1), umask + nu_)
+            gpu.barrier()
+
+        UG_KS = tuple(range(UG_CPW))
+
+        def wug_load(t, has, scales=True, fp8=False, ks=UG_KS):
+            """Up / gate weights (this wave's 128-k chunks ``ks``; and ``scales``:
+            the scale dwords, which hold the pair's both halves) of task t =
+            (distinct expert t / UG_PER_SLOT, row group t % UG_PER_SLOT); without
+            the task, reads go through zero-sized buffers. ``fp8``: the MXFP8
+            shared expert's task (weights.mxfp8_a_tiles; the scales -- every
+            chunk's -- in the routed tables' order)."""
+            u = t // UG_PER_SLOT
+            cg = t % UG_PER_SLOT
+            if const_expr(fp8):
+                e_u = 0
+                r_w = rsrc(
+                    sh_ptr(SH_W13), has.select(fx.Int32(SH8_W13_BYTES), fx.Int32(0))
+                )
+                r_s = rsrc(
+                    sh_ptr(SH_S13), has.select(fx.Int32(SH8_S13_BYTES), fx.Int32(0))
+                )
+            else:
+                e_u = fx.min(uniform(fx.ptr_load(uexp + u)), E_TOP)
+                r_w = rsrc(
+                    w13 + fx.Int64(e_u) * W13_BYTES,
+                    has.select(fx.Int32(W13_BYTES), fx.Int32(0)),
+                )
+                r_s = rsrc(s13, has.select(fx.Int32(S13_BYTES), fx.Int32(0)))
+            wts: list[list] = []
+            scs: list[list] = []
+            for gu in range_constexpr(2):
+                rg = gu * UG_PER_SLOT + cg
+                wts.append([])
+                scs.append([])
+                for i in range_constexpr(UG_CPW):
+                    kc = wave * UG_CPW + i
+                    if const_expr(i in ks):
+                        wts[gu].append(w_tile(r_w, rg, kc, lane, HIDDEN, fp8))
+                    if const_expr(scales and i % 2 == 0):
+                        # the dword holding this lane's scales of chunks kc, kc + 1
+                        # of both 16-row halves of the 32-row block (bytes 2 d + a)
+                        scs[gu].append(
+                            fx.Int32(
+                                bo.buffer_load(
+                                    r_s,
+                                    group_scale_byte(
+                                        e_u, W13_ROWS, rg, kc, W13_SCALE_COLS
+                                    )
+                                    // 4,
+                                    vec_width=1,
+                                    dtype=T.i32,
+                                )
+                            )
+                        )
+            return wts, scs
+
+        def wug_acc(wts, scs, half, cg_, cu_, ks=UG_KS, fp8=False):
+            """A task's up / gate MFMAs over its 128-k chunks ``ks`` of this wave
+            (wts[gu][j]: chunk ks[j]) for every token (B column l % 16 = token
+            l % 16), into (cg_, cu_). ``half``: the task's 16-row half of its
+            32-row scale block (compile time); ``fp8``: an MXFP8 task
+            (``wug_load(fp8=True)``)."""
+            mfma_ = mfma_f8f8 if fp8 else mfma_f4f8
+            col = fx.min(l16, tokens - 1)
+            sbs = [
+                lds_i32(XSC_LDS + col * XSC_ROW + (wave * UG_CPW + i) * 4 + g4)
+                for i in ks
+            ]
+            for j in range_constexpr(len(ks)):
+                i = ks[j]
+                kc = wave * UG_CPW + i
+                xb = lds_b8(col * XN8_ROW + kc * 32)
+                sel = (i % 2) * 2 + half
+                cg_ = mfma_(wts[0][j], scs[0][i // 2], xb, sbs[j], cg_, sel)
+                cu_ = mfma_(wts[1][j], scs[1][i // 2], xb, sbs[j], cu_, sel)
+            return cg_, cu_
+
+        def wug_y(wts, scs, half):
+            """A task's up / gate GEMV (``wug_acc``) -> swiglu: thread (row r,
+            token k)'s mid value."""
+            zero = fx.Vector.filled(4, 0.0, fx.Float32)
+            return wug_swiglu(*wug_acc(wts, scs, half, zero, zero))
+
+        def wug_swiglu(cg_, cu_):
+            """The waves' up / gate partials -> swiglu (``wug_y``)."""
+            fx.ptr_store(cg_, red + (wave * 64 + lane) * 4)
+            fx.ptr_store(cu_, fpool + (W_UPP + (wave * 64 + lane) * 4))
+            gpu.barrier()
+            r = tid % UG_ROWS
+            k = fx.min(tid // UG_ROWS, tokens - 1)
+            gsum = fx.Float32(0.0)
+            usum = fx.Float32(0.0)
+            for w in range_constexpr(WAVES):
+                gsum = gsum + fx.ptr_load(
+                    red + ((w * 64 + 16 * (r // 4) + k) * 4 + r % 4)
+                )
+                usum = usum + fx.ptr_load(
+                    fpool + (W_UPP + (w * 64 + 16 * (r // 4) + k) * 4 + r % 4)
+                )
+            gpu.barrier()  # red and the up partials are free for the next task
+            return swiglu_mul_batch([gsum], [usum], fx.Float32(-swiglu_limit))[0]
+
+        def fp8_lanes4(q):
+            """Lane r (r % 4 == 0): the E4M3 dword of lanes r .. r + 3's values."""
+            lo = fx.Int32(
+                rocdl.cvt_pk_fp8_f32(T.i32, q, xshfl(q, 1), fx.Int32(0), False)
+            )
+            return lo | (xshfl(lo, 2) << 16)
+
+        def wug_publish(pr, y_lo, y_hi):
+            """Task pair pr = (distinct expert pr / UG_PAIRS, 32-row block
+            pr % UG_PAIRS): its mid rows (thread (r, k) holds rows r and 16 + r of
+            the block) as the down stage's MXFP8 -- the rounding its bf16 staging
+            did, one 32-block a 16-lane group -- then the pair's flag."""
+            u = pr // UG_PAIRS
+            blk_ = pr % UG_PAIRS
+            r = tid % UG_ROWS
+            k = fx.min(tid // UG_ROWS, tokens - 1)
+            lo = bf16_round(y_lo)
+            hi = bf16_round(y_hi)
+            sc, inv = mx8_scale([lo, hi], (1, 2, 4, 8))
+            words = [fp8_lanes4(mx8_q(v, inv)) for v in (lo, hi)]
+            e_u = fx.Int32(fx.ptr_load(uexp + u))
+            m_u = fx.Int32(fx.ptr_load(umask + u))
+            row = k * MOE_SLOTS + slot_of(k, e_u)
+            if (tid < UG_ROWS * tokens) & (((m_u >> k) & 1) != 0):
+                if r % 4 == 0:
+                    for h in range_constexpr(2):
+                        bo.buffer_store(
+                            words[h],
+                            rsrc(mb("mid8p")),
+                            row * MID_ROW_WORDS + blk_ * 8 + h * 4 + r // 4,
+                            cache_modifier=CM_DEV,
+                        )
+                if r == 0:
+                    bo.buffer_store(
+                        sc,
+                        rsrc(mb("midsp")),
+                        row * MIDSC_ROW_WORDS + blk_,
+                        cache_modifier=CM_DEV,
+                    )
+            fx.memory_fence(
+                syncscope=rocdl.SyncScope.Workgroup, ordering=fx.AtomicOrdering.Release
+            )
+            gpu.barrier()
+            if tid == 0:
+                mb_put(mb("ugdone"), pr, fx.Int32(1))
+
+        def down_tables():
+            """Urow / uwt := per (distinct expert, B column) the down stage's mid row
+            (the zero row where the column's token is not routed there) and route
+            weight. Routing only: built while the first up / gate weights land."""
+            nu = uniform(fx.ptr_load(uexp + (U_MAX + 1)))
+            for i in range_constexpr((U_MAX * 16 + THREADS - 1) // THREADS):
+                ix = tid + THREADS * i
+                if ix < U_MAX * 16:
+                    xu = ix // 16
+                    xk = ix % 16
+                    xe = fx.Int32(fx.ptr_load(uexp + fx.min(xu, nu)))
+                    xm = fx.Int32(fx.ptr_load(umask + fx.min(xu, nu)))
+                    xt = fx.min(xk, tokens - 1)
+                    xs_ = slot_of(xt, xe)
+                    xh = (((xm >> xk) & 1) != 0) & (xu <= nu)
+                    fx.ptr_store(
+                        xh.select(xt * MOE_SLOTS + xs_, fx.Int32(MID_ROWS)), urow + ix
+                    )
+                    fx.ptr_store(
+                        xh.select(fx.ptr_load(rwt + (8 * xt + xs_)), fx.Float32(0.0)),
+                        uwt + ix,
+                    )
+
+        def stage_moe_wide():
+            td = start("down")  # this CTA's down task (the split tasks come first)
+            has_dn = td < N_DN
+            td = fx.min(td, N_DN - 1)
+            # every token's routing: wave w routes tokens w, w + WAVES, .., their
+            # scores polled in one batch
+            mine = bid % SCORE_COPIES
+            # wave w's expert -> mask of the tokens it routed there, in the up
+            # partials' words (free until the first up / gate task)
+            wave_mask = xs + (W_UPP + wave * N_ROUTED)
+            fx.ptr_store(fx.Int32(0), wave_mask + lane)
+            fx.ptr_store(fx.Int32(0), wave_mask + (lane + 64))
+            wave_toks = [
+                fx.min(wave + WAVES * j, tokens - 1) for j in range(ROUTE_ITERS)
+            ]
+            sg = mb_poll(
+                [
+                    (scores_copy(mine, k), (lane + 64 * h) * 2, 2)
+                    for k in wave_toks
+                    for h in range(2)
+                ]
+            )
+            for j in range_constexpr(ROUTE_ITERS):
+                if wave + WAVES * j < tokens:
+                    route_top4(
+                        wave_toks[j],
+                        sg[2 * j][0],
+                        sg[2 * j][1].bitcast(fx.Float32),
+                        sg[2 * j + 1][0],
+                        sg[2 * j + 1][1].bitcast(fx.Float32),
+                        fpool + W_SIG,
+                        wave_mask,
+                    )
+            stamp(2)
+            gpu.barrier()
+            build_experts()
+            stamp(6)
+            if (bid == BASE["down"]) & (tid < MOE_SLOTS * tokens):
+                sk = tid // MOE_SLOTS
+                sj = tid % MOE_SLOTS
+                mb_put(
+                    mb("sel"),
+                    sk * 2 * MOE_SLOTS + sj,
+                    fx.ptr_load(route + (8 * sk + sj)),
+                )
+                mb_put(
+                    mb("sel"),
+                    sk * 2 * MOE_SLOTS + MOE_SLOTS + sj,
+                    fx.ptr_load(rwt + (8 * sk + sj)),
+                )
+            nu_ug = uniform(fx.ptr_load(uexp + (U_MAX + 1)))
+            n_pairs = (nu_ug if shared_fp8 else nu_ug + 1) * UG_PAIRS
+            # shared_fp8: shared pair j runs after the table's pairs on CTA
+            # (sh0 + j) % G. It costs two of theirs: from a full round of pairs on,
+            # the table's skip slots hole0 .. hole0 + UG_PAIRS - 1 (slot s: CTA s %
+            # G's (s / G)-th pair), the shared CTAs' last full round, so that those
+            # carry one pair fewer (otherwise the up / gate stage ran a pair longer
+            # than the fused build's at every S; S = 8: 25.8 -> 31.6 us slowest CTA)
+            sh0 = n_pairs
+            if const_expr(shared_fp8):
+                hole0 = n_pairs + (UG_PAIRS - G)
+                holes = hole0 >= 0
+                sh0 = n_pairs + holes.select(fx.Int32(UG_PAIRS), fx.Int32(0))
+            sh_j = (bid + (G - sh0 % G)) % G
+            has_shp = sh_j < UG_PAIRS
+            sh_pr = n_pairs + fx.min(sh_j, UG_PAIRS - 1)
+            UG_ALL = UGW_ITERS + (2 * SH_PARTS if shared_fp8 else 0)
+
+            def task(it):
+                """Iteration it's task: slot bid + (it / 2) G's pair, half it % 2
+                (past UGW_ITERS: the shared pair's, half (it - UGW_ITERS) / SH_PARTS
+                a k part an iteration)."""
+                if const_expr(it >= UGW_ITERS):
+                    return sh_pr * 2 + (it - UGW_ITERS) // SH_PARTS, has_shp
+                pr = fx.Int32(bid + it // 2 * G)
+                if const_expr(shared_fp8):
+                    # a hole runs nothing (pair n_pairs); past them, UG_PAIRS lower
+                    pr = (holes & (pr >= hole0)).select(
+                        (pr < hole0 + UG_PAIRS).select(n_pairs, pr - UG_PAIRS), pr
+                    )
+                return fx.min(pr, n_pairs - 1) * 2 + it % 2, pr < n_pairs
+
+            def task_load(it, t, has, scales=True):
+                if const_expr(it >= UGW_ITERS):
+                    # the shared pair's scales come with its first part
+                    q = it - UGW_ITERS
+                    return wug_load(
+                        t, has, scales=q == 0, fp8=True, ks=sh_ks(q % SH_PARTS)
+                    )
+                return wug_load(t, has, scales)
+
+            t0, h0 = task(0)
+            cur = task_load(0, t0, h0)  # task 0: a lo half, with the pair's scales
+            down_tables()
+            stamp(7)
+            # the first task's weights are in flight: now the xn
+            # every token's MXFP8 xn and scales -> the pool, once every router task
+            # flagged its slice (plain words, 16 B a load)
+            n_rt = N_ROUTER * tokens
+            if wave == 0:
+                mb_poll(
+                    [
+                        (mb("xndone"), fx.min(lane + 64 * c, n_rt - 1), 1)
+                        for c in range((n_rt + 63) // 64)
+                    ]
+                )
+            gpu.barrier()
+            acquire()
+            for base_w, n_row, lds_w, lds_row in (
+                (mb("xn8p"), XN8_WORDS, 0, XN8_ROW),
+                (mb("xscp"), XSC_WORDS, XSC_LDS, XSC_ROW),
+            ):
+                n_v = tokens * n_row // 4
+                for i in range_constexpr((n_v + THREADS - 1) // THREADS):
+                    v = tid + THREADS * i
+                    if v < n_v:
+                        fx.ptr_store(
+                            fx.Vector(
+                                bo.buffer_load(
+                                    rsrc(base_w),
+                                    v * 4,
+                                    vec_width=4,
+                                    dtype=T.i32,
+                                    cache_modifier=CM_DEV,
+                                )
+                            ),
+                            xs + (lds_w + v * 4 // n_row * lds_row + v * 4 % n_row),
+                        )
+            gpu.barrier()
+            y_lo = fx.Float32(0.0)  # the pair's first half, until the second's
+            if const_expr(shared_fp8):
+                # a shared half's up / gate accumulators, part to part
+                sh_cg = fx.Vector.filled(4, 0.0, fx.Float32)
+                sh_cu = fx.Vector.filled(4, 0.0, fx.Float32)
+            for it in range_constexpr(UG_ALL):
+                t, has_t = task(it)
+                if const_expr(it + 1 < UG_ALL):
+                    tn, hn = task(it + 1)
+                    # a hi half reads no scales: its lo half's dwords hold them
+                    nxt = task_load(it + 1, tn, hn, scales=(it + 1) % 2 == 0)
+                if const_expr(it < UGW_ITERS):
+                    if const_expr(it % 2 == 0):
+                        pair_scs = cur[1]
+                    if has_t:
+                        # a pair: row groups cg, cg + 1
+                        y = wug_y(cur[0], pair_scs, it % 2)
+                        if const_expr(it % 2 == 0):
+                            y_lo = y
+                        else:
+                            wug_publish(t // 2, y_lo, y)
+                else:
+                    half, part = divmod(it - UGW_ITERS, SH_PARTS)
+                    if const_expr(half == 0 and part == 0):
+                        pair_scs = cur[1]
+                    if has_t:
+                        if const_expr(part == 0):
+                            sh_cg = fx.Vector.filled(4, 0.0, fx.Float32)
+                            sh_cu = fx.Vector.filled(4, 0.0, fx.Float32)
+                        sh_cg, sh_cu = wug_acc(
+                            cur[0], pair_scs, half, sh_cg, sh_cu, sh_ks(part), fp8=True
+                        )
+                        if const_expr(part == SH_PARTS - 1):
+                            y = wug_swiglu(sh_cg, sh_cu)
+                            if const_expr(half == 0):
+                                y_lo = y
+                            else:
+                                wug_publish(t // 2, y_lo, y)
+                if const_expr(it + 1 < UG_ALL):
+                    cur = nxt
+            stamp(8)
+            if has_dn:
+                wide_down(td)
+                stamp(11)
+
+        def wide_down(td):
+            """Row group td's down GEMV over every distinct expert (B column l % 16 =
+            token l % 16's mid for it, the zero row where the token is not routed
+            there), route-weighted per column -> ffn_finish."""
+            nu = uniform(fx.ptr_load(uexp + (U_MAX + 1)))
+            rg_d = td * 2 + wave // 4
+            n_units = (nu if shared_fp8 else nu + 1) * DN_KC
+
+            def unit_load(i):
+                """(weights, scale, unit) of this wave's i-th (expert, 128-k chunk)
+                unit, zero-sized reads past the last."""
+                j = wave % 4 + 4 * i
+                ok = j < n_units
+                # clamped into the table: a unit past the last reads a real or zero
+                # row, never past urow (a garbage row is garbage B, and 0 * NaN)
+                du = fx.min(j // DN_KC, U_MAX - 1)
+                dkc = j % DN_KC
+                de = fx.min(uniform(fx.ptr_load(uexp + du)), E_TOP)
+                dwv = fx.Vector(
+                    bo.buffer_load(
+                        rsrc(
+                            w2 + fx.Int64(de) * W2_BYTES,
+                            ok.select(fx.Int32(W2_BYTES), fx.Int32(0)),
+                        ),
+                        ((rg_d * (INTER // 64) + dkc * 2) * 512 + lane * 16) // 4,
+                        vec_width=4,
+                        dtype=T.i32,
+                    )
+                )
+                dsc = e8m0_load(
+                    rsrc(s2, ok.select(fx.Int32(S2_BYTES), fx.Int32(0))),
+                    group_scale_byte(de, HIDDEN, rg_d, dkc, W2_SCALE_COLS),
+                )
+                return dwv, dsc, du, dkc, ok
+
+            def with_rows(u):
+                """A loaded unit plus its B row and route weight (the tables are in),
+                read ahead so the unit's B loads wait on one LDS trip."""
+                dwv, dsc, du, dkc, ok = u
+                drow = fx.Int32(fx.ptr_load(urow + (du * 16 + l16)))
+                dw = ok.select(fx.ptr_load(uwt + (du * 16 + l16)), fx.Float32(0.0))
+                return dwv, dsc, drow, dkc, dw, ok
+
+            # software pipeline: the first units' weights go out before the mids are
+            # in; unit i + DN_DEPTH's before unit i's MFMA
+            acc = [fx.Float32(0.0) for _ in range(4)]
+            if const_expr(shared_fp8):
+                # the MXFP8 shared expert's units (wave % 4 + 4 i): no routing needed,
+                # their weights go out first
+                r_shared_w = rsrc(sh_ptr(SH_W2))
+                r_shared_s = rsrc(sh_ptr(SH_S2))
+                sh_dn = []
+                for i in range_constexpr(SH_UPW):
+                    skc = fx.min(wave % 4 + 4 * i, DN_KC - 1)
+                    sh_dn.append(
+                        (
+                            f8_tile_load(r_shared_w, rg_d * (INTER // 128) + skc, lane),
+                            e8m0_load(
+                                r_shared_s,
+                                group_scale_byte(0, HIDDEN, rg_d, skc, W2_SCALE_COLS),
+                            ),
+                            skc,
+                            wave % 4 + 4 * i < DN_KC,
+                        )
+                    )
+            inflight = [unit_load(i) for i in range(min(DN_DEPTH, DN_WITERS))]
+            # every up / gate task pair's flag, then the MXFP8 mid rows and their
+            # scales, the pool's layout (plain words, 16 B a load)
+            n_ug = (nu + 1) * UG_PAIRS
+            mb_poll(
+                [
+                    (mb("ugdone"), fx.min(tid + THREADS * i, n_ug - 1), 1)
+                    for i in range((U_MAX * UG_PAIRS + THREADS - 1) // THREADS)
+                ]
+            )
+            gpu.barrier()
+            acquire()
+            for base_w, n_row, lds_w, lds_row in (
+                (mb("mid8p"), MID_ROW_WORDS, 0, MID_ROW),
+                (mb("midsp"), MIDSC_ROW_WORDS, W_MIDSC, MIDSC_ROW),
+            ):
+                n_v = MID_ROWS * n_row // 4
+                for i in range_constexpr((n_v + THREADS - 1) // THREADS):
+                    v = tid + THREADS * i
+                    if v < n_v:
+                        fx.ptr_store(
+                            fx.Vector(
+                                bo.buffer_load(
+                                    rsrc(base_w),
+                                    v * 4,
+                                    vec_width=4,
+                                    dtype=T.i32,
+                                    cache_modifier=CM_DEV,
+                                )
+                            ),
+                            xs + (lds_w + v * 4 // n_row * lds_row + v * 4 % n_row),
+                        )
+            if tid < MID_ROW_WORDS:
+                fx.ptr_store(fx.Int32(0), xs + (MID_ROWS * MID_ROW + tid))
+            if tid < MIDSC_ROW_WORDS:
+                fx.ptr_store(fx.Int32(0), xs + (W_MIDSC + MID_ROWS * MIDSC_ROW + tid))
+            gpu.barrier()
+            stamp(10)
+            for i in range_constexpr(DN_WITERS):
+                if const_expr(i == 0):  # the tables are in: the early units' rows
+                    inflight = [with_rows(u) for u in inflight]
+                dwv, dsc, drow, dkc, dw, ok = inflight.pop(0)
+                if const_expr(i + DN_DEPTH < DN_WITERS):
+                    inflight.append(with_rows(unit_load(i + DN_DEPTH)))
+                dcu = mfma_f4f8(
+                    dwv,
+                    dsc,
+                    lds_b8(drow * MID_ROW + dkc * 32),
+                    lds_i32(W_MIDSC + drow * MIDSC_ROW + dkc * 4 + g4),
+                    fx.Vector.filled(4, 0.0, fx.Float32),
+                )
+                # a unit past the last adds nothing, whatever its MFMA gave
+                acc = [
+                    acc[e] + ok.select(dcu[e] * dw, fx.Float32(0.0)) for e in range(4)
+                ]
+            if const_expr(shared_fp8):
+                # the shared expert's B column l16: its mid row (slot TOP_K) and weight
+                srow = fx.Int32(fx.ptr_load(urow + (nu * 16 + l16)))
+                sw = fx.ptr_load(uwt + (nu * 16 + l16))
+                for swv, ssc, skc, sok in sh_dn:
+                    dcu = mfma_f8f8(
+                        swv,
+                        ssc,
+                        lds_b8(srow * MID_ROW + skc * 32),
+                        lds_i32(W_MIDSC + srow * MIDSC_ROW + skc * 4 + g4),
+                        fx.Vector.filled(4, 0.0, fx.Float32),
+                    )
+                    acc = [
+                        acc[e] + sok.select(dcu[e] * sw, fx.Float32(0.0))
+                        for e in range(4)
+                    ]
+            a0, a1, a2, a3 = acc
+            fx.ptr_store(
+                fx.Vector.from_elements([a0, a1, a2, a3], fx.Float32),
+                red + (wave * 64 + lane) * 4,
+            )
+            gpu.barrier()
+            if tid < DN_ROWS * tokens:
+                k = tid // DN_ROWS
+                grp = (tid % DN_ROWS) // 16
+                rr = tid % 16
+                tot = fx.Float32(0.0)
+                for w in range_constexpr(4):
+                    tot = tot + fx.ptr_load(
+                        red + (((grp * 4 + w) * 64 + 16 * (rr // 4) + k) * 4 + rr % 4)
+                    )
+                fx.ptr_store(bf16_round(tot), misc + tid)
+            gpu.barrier()
+            ffn_finish(td)
+
+        def stage_moe():
+            td = bid - BASE["down"]  # this CTA's down task
+            has_dn = (td >= 0) & (td < N_DN)
+            td = fx.min(fx.max(td, 0), N_DN - 1)
+            has_sh = start("shared") < SH_TASKS
+            n_xw = XN8_WORDS // THREADS
+
+            def xn_specs(k):
+                return [
+                    (mb("xn8"), k * XN8_WORDS + tid + i * THREADS, 1)
+                    for i in range(n_xw)
+                ] + [(mb("xsc"), k * XSC_WORDS + fx.min(tid, XSC_WORDS - 1), 1)]
+
+            def stage_xn(k, got):
+                for i in range_constexpr(n_xw):
+                    fx.ptr_store(got[i][0], xs + (k * XN8_ROW + tid + i * THREADS))
+                if tid < XSC_WORDS:
+                    fx.ptr_store(got[n_xw][0], xs + (XSC_LDS + k * XSC_ROW + tid))
+
+            def fetch_xn(tok=None):
+                """Token tok's (default: every token's) MXFP8 xn and scales -> xs. Wave
+                0 first polls the last word of each router task's slice: all 512
+                threads of 256 CTAs spinning on the whole of it (S = 4) held the
+                shared-expert CTAs' xn back 3 us."""
+                gate_tok = fx.min(lane // N_ROUTER, tokens - 1) if tok is None else tok
+                if wave == 0:
+                    mb_poll(
+                        [
+                            (
+                                mb("xn8"),
+                                gate_tok * XN8_WORDS
+                                + (lane % N_ROUTER + 1) * (XN_SLICE // 4)
+                                - 1,
+                                1,
+                            )
+                        ]
+                    )
+                gpu.barrier()
+                if tok is None:
+                    got = mb_poll([sp for k in range(tokens) for sp in xn_specs(k)])
+                    for k in range_constexpr(tokens):
+                        stage_xn(k, got[k * (n_xw + 1) : (k + 1) * (n_xw + 1)])
+                else:
+                    stage_xn(tok, mb_poll(xn_specs(tok)))
+
+            sh_late = shared_after_routing(tokens)
+            if const_expr(sh_late):
+                sh_j = fx.min(start("shared"), SH_TASKS - 1)
+                sh_w = sh_load(sh_j, has_sh)
+                if has_sh:
+                    fetch_xn()
+            elif has_sh:  # the shared expert needs xn, not the routing
+                sh_j0 = start("shared")
+                sh_w0 = sh_load(sh_j0)
+                fetch_xn()
+                gpu.barrier()
+                sh_compute(sh_j0, sh_w0[0], sh_w0[1])
+            if ug_task_id(0)[1] | has_dn:
+                # xn lands before the logits: stage it while waiting (a shared-expert
+                # CTA has it already); a CTA's routed tasks are one token's
+                if start("shared") >= SH_TASKS:
+                    fetch_xn(bid // CPT if CPT is not None else None)
+                # wave k polls token k's router logits and routes it right away
+                mine = bid % SCORE_COPIES
+                for k in range_constexpr(tokens):
+                    if wave == k:
+                        got = mb_poll(
+                            [
+                                (scores_copy(mine, k), (lane + 64 * h) * 2, 2)
+                                for h in range(2)
+                            ]
+                        )
+                        stamp(21)
+                        route_top4(
+                            k,
+                            got[0][0],
+                            got[0][1].bitcast(fx.Float32),
+                            got[1][0],
+                            got[1][1].bitcast(fx.Float32),
+                            opart,
+                        )
+                stamp(2)
+                gpu.barrier()
+                stamp(6)
+                if (bid == BASE["down"]) & (tid < MOE_SLOTS * tokens):
+                    sk = tid // MOE_SLOTS
+                    sj = tid % MOE_SLOTS
+                    mb_put(
+                        mb("sel"),
+                        sk * 2 * MOE_SLOTS + sj,
+                        fx.ptr_load(route + (8 * sk + sj)),
+                    )
+                    mb_put(
+                        mb("sel"),
+                        sk * 2 * MOE_SLOTS + MOE_SLOTS + sj,
+                        fx.ptr_load(rwt + (8 * sk + sj)),
+                    )
+                cur = ug_load(ug_task_id(0)[0])
+                stamp(7)
+                if const_expr(sh_late):  # noqa: SIM102
+                    if has_sh:  # while the first routed task's weights are in flight
+                        sh_compute(sh_j, sh_w[0], sh_w[1])
+                # task i + 1's weights go out before task i's GEMV (software pipeline);
+                # past the last task the load index is clamped and the result dropped
+                for it in range_constexpr(UG_ITERS):
+                    t, has_t = ug_task_id(it)
+                    if const_expr(it + 1 < UG_ITERS):
+                        nxt = ug_load(ug_task_id(it + 1)[0])
+                    if has_t:
+                        ug_compute(t, cur[0], cur[1])
+                    if const_expr(it == 0):
+                        # token 0's down weights stream under the mid exchange; issued
+                        # with the first task's up / gate weights they share the CU's
+                        # bandwidth and delay them (K4 on 4 GPUs 34.3 -> 32.8 us)
+                        dn_units = {0: down_load(0, td, has_dn)}
+                        sh_units = down_load_shared(td, has_dn)
+                    if const_expr(tokens > 1 and it == UG_ITERS - 1):
+                        # two tokens in flight when the down GEMVs start
+                        dn_units[1] = down_load(1, td, has_dn)
+                    if const_expr(it + 1 < UG_ITERS):
+                        cur = nxt
+                stamp(8)
+                if has_dn:
+                    moe_down(td, has_dn, dn_units, sh_units)
+                    stamp(11)
+
+        def moe_down(t, has_dn, dn_units, sh_units):
+            """Every token's route-weighted down GEMV of row group ``t`` -> push bf16
+            partial rows -> rank-ordered sum -> ar_out."""
+
+            def stage_mids(k):
+                """Token k's mid rows (every slot) -> xs, as bf16 pairs."""
+                mids = mb_poll(
+                    [
+                        (
+                            mb("mid"),
+                            k * MOE_SLOTS * INTER
+                            + fx.min(tid * 2 + i * 2 * THREADS, MOE_SLOTS * INTER - 2),
+                            2,
+                        )
+                        for i in range(4)
+                    ]
+                )
+                for i in range_constexpr(4):
+                    e = tid * 2 + i * 2 * THREADS
+                    if e < MOE_SLOTS * INTER:
+                        fx.ptr_store(
+                            bf16_pair(
+                                mids[i][0].bitcast(fx.Float32),
+                                mids[i][1].bitcast(fx.Float32),
+                            ),
+                            xs + (k * (MOE_SLOTS * INTER // 2) + e // 2),
+                        )
+
+            # every token's mid in one go: they land together (S = 4: -1.1 us)
+            for k in range_constexpr(tokens):
+                stage_mids(k)
+            stamp(10)
+            gpu.barrier()
+            # then MXFP8 in a pass of its own: quantizing inside stage_mids (after
+            # each poll) cost S = 1 / 4 1.75 / 6 us, though the arithmetic is tiny
+            for k in range_constexpr(tokens):
+                if tid < MID8_WORDS // 2:  # 8 elements a lane, 4 lanes a 32-block
+                    w = fx.Vector(
+                        fx.ptr_load(xs + (k * MID16_WORDS + tid * 4), result_type=v4f)
+                    ).bitcast(fx.Int32)
+                    f = [v for j in range(4) for v in bf2_f32(w[j])]
+                    sc, inv = mx8_scale(f, (1, 2))
+                    qv = [mx8_q(v, inv) for v in f]
+                    fx.ptr_store(
+                        fp8_pack4(*qv[:4]), xs + (MID8_LDS + k * MID8_WORDS + tid * 2)
+                    )
+                    fx.ptr_store(
+                        fp8_pack4(*qv[4:]),
+                        xs + (MID8_LDS + k * MID8_WORDS + tid * 2 + 1),
+                    )
+                    if tid % 4 == 0:
+                        fx.ptr_store(sc, xs + (MIDSC_LDS + k * MIDSC_WORDS + tid // 4))
+            gpu.barrier()
+            # the shared expert once, B column l % 16 = token l % 16's mid -> opart
+            # (free once the split is done); each token's reduce adds its column
+            col_tok = fx.min(l16, tokens - 1)
+            sh = [fx.Float32(0.0) for _ in range(4)]
+            for wv, sc, kc, wgt in sh_units:
+                cu = mfma_sh(
+                    wv,
+                    sc,
+                    lds_b8(
+                        MID8_LDS
+                        + col_tok * MID8_WORDS
+                        + (TOP_K * INTER + kc * 128) // 4
+                    ),
+                    lds_i32(
+                        MIDSC_LDS
+                        + col_tok * MIDSC_WORDS
+                        + (TOP_K * INTER + kc * 128) // 32
+                        + g4
+                    ),
+                    fx.Vector.filled(4, 0.0, fx.Float32),
+                )
+                sh = [sh[j4] + cu[j4] * wgt for j4 in range(4)]
+            fx.ptr_store(
+                fx.Vector.from_elements(sh, fx.Float32), opart + (wave * 64 + lane) * 4
+            )
+            # every token's GEMV first, partials in opart past the shared ones: one
+            # barrier and one reduce for all tokens, not two barriers per token
+            for k in range_constexpr(tokens):
+                # token k + 2's weights go out before token k's GEMV
+                if const_expr(k + 2 < tokens):
+                    dn_units[k + 2] = down_load(k + 2, t, has_dn)
+                acc = [fx.Float32(0.0) for _ in range(4)]
+                for wv, sc, sl, kc, wgt in dn_units.pop(k):
+                    cu = mfma_f4f8(
+                        wv,
+                        sc,
+                        lds_b8(
+                            MID8_LDS + k * MID8_WORDS + (sl * INTER + kc * 128) // 4
+                        ),
+                        lds_i32(
+                            MIDSC_LDS
+                            + k * MIDSC_WORDS
+                            + (sl * INTER + kc * 128) // 32
+                            + g4
+                        ),
+                        fx.Vector.filled(4, 0.0, fx.Float32),
+                    )
+                    acc = [acc[j4] + cu[j4] * wgt for j4 in range(4)]
+                fx.ptr_store(
+                    fx.Vector.from_elements(acc, fx.Float32),
+                    opart + ((k + 1) * WAVES * 64 + wave * 64 + lane) * 4,
+                )
+            gpu.barrier()
+            if tid < DN_ROWS * tokens:
+                k = tid // DN_ROWS
+                grp = (tid % DN_ROWS) // 16
+                rr = tid % 16
+                tot = fx.Float32(0.0)
+                for w in range_constexpr(4):
+                    tot = tot + fx.ptr_load(
+                        opart
+                        + (
+                            ((k + 1) * WAVES * 64 + (grp * 4 + w) * 64 + 16 * (rr // 4))
+                            * 4
+                            + rr % 4
+                        )
+                    )
+                for w in range_constexpr(4):
+                    tot = tot + fx.ptr_load(
+                        opart + (((grp * 4 + w) * 64 + 16 * (rr // 4) + k) * 4 + rr % 4)
+                    )
+                fx.ptr_store(bf16_round(tot), misc + tid)
+            gpu.barrier()
+            ffn_finish(t)
+
+        (stage_moe_wide if WIDE else stage_moe)()
+        flush_stamps()
+
+    @flyc.jit
+    def launch_post_attn(
+        h_in: Int64,
+        q: Int64,
+        block_table: Int64,
+        seq_lens: Int64,
+        k_cache: Int64,
+        v_cache: Int64,
+        k_scale: Int64,
+        v_scale: Int64,
+        w_o: Int64,
+        s_o: Int64,
+        g_post: Int64,
+        w_gate: Int64,
+        bias: Int64,
+        w13: Int64,
+        s13: Int64,
+        w2: Int64,
+        s2: Int64,
+        sh_args: Int64,
+        h_mid: Int64,
+        ar_out: Int64,
+        scratch: Int64,
+        sym: Int64,
+        peers: Int64,
+        step: Int64,
+        rank: Int32,
+        layer: Int32,
+        bt_width: Int32,
+        q_len: Int32,
+        tl: Int64,
+        k1_args: Int64,
+        positions: Int64,
+        slot_mapping: Int64,
+        res: Int64,
+        stream: fx.Stream = _CURRENT_STREAM,
+    ):
+        post_attn_kernel(
+            h_in,
+            q,
+            block_table,
+            seq_lens,
+            k_cache,
+            v_cache,
+            k_scale,
+            v_scale,
+            w_o,
+            s_o,
+            g_post,
+            w_gate,
+            bias,
+            w13,
+            s13,
+            w2,
+            s2,
+            sh_args,
+            h_mid,
+            ar_out,
+            scratch,
+            sym,
+            peers,
+            step,
+            rank,
+            layer,
+            bt_width,
+            q_len,
+            tl,
+            k1_args,
+            positions,
+            slot_mapping,
+            res,
+        ).launch(grid=(G,), block=(THREADS,), stream=stream)
+
+    return launch_post_attn
diff --git a/vllm/models/minimax_m3/amd/mono/kernels/pre_attn.py b/vllm/models/minimax_m3/amd/mono/kernels/pre_attn.py
new file mode 100644
index 000000000000..b8b9b2133e39
--- /dev/null
+++ b/vllm/models/minimax_m3/amd/mono/kernels/pre_attn.py
@@ -0,0 +1,902 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+# SPDX-License-Identifier: MIT
+# Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved.
+"""K1 ``mono_pre``: one MiniMax-M3 sparse layer, input norm through cache insert.
+
+One launch per layer and rank, ``BLOCKS`` CTAs x ``THREADS`` threads, for S =
+``tokens`` decode rows, ``q_len`` consecutive rows a request (each row below is
+per token)::
+
+    acc = ar + res (f32)  -> residual out bf16(acc)
+      -> gemma RMSNorm -> per-token FP8 (amax / FP8_MAX, true divide)
+      -> qkv | index_q | index_k GEMV on the preshuffled ptpc-FP8 weight
+      -> per-head gemma norm + partial NeoX RoPE
+      -> q, index_q out (bf16); K / V -> page-16 SHUFFLE cache with per-token
+         scales; index_k -> index cache (E4M3, unit scale)
+
+It reads exactly what the original path leaves in memory (the ``ar`` / ``res``
+pair of the custom all-reduce, the loaded weights, the bound caches) and writes
+what the original ``fused_qknorm_idxrqknorm`` writes, so either path can run any
+layer. The numerical contracts (reduction orders, where values round to bf16)
+follow the aiter kernels the original path uses; see each stage.
+"""
+
+import flydsl.compiler as flyc
+import flydsl.expr as fx
+from aiter.ops.flydsl.kernels import buffer_ops as bo
+from flydsl.expr import const_expr, gpu, range_constexpr, rocdl
+from flydsl.expr import math as fmath
+from flydsl.expr.typing import Int32, Int64, T
+
+from vllm.models.minimax_m3.amd.mono import config as mono_config
+from vllm.models.minimax_m3.amd.mono.config import (
+    BLOCKS,
+    HEAD_DIM,
+    HIDDEN,
+    LOCAL_Q_HEADS,
+    MAX_QKV_ROWS,
+    MAX_TOKENS,
+    ONE_INDEX_HEAD,
+    PAGE16,
+    ROTARY_DIM,
+    THREADS,
+    TP,
+    WAVES,
+    IndexHeads,
+    qkv_rows,
+)
+from vllm.models.minimax_m3.amd.mono.kernels.common import (
+    CM_DEV,
+    Mailbox,
+    ar_block_sums,
+    ar_pack_sumsq,
+    bf16_pair,
+    bf16_round,
+    butterfly,
+    div_rn,
+    fp8_pack4,
+    hw_rsq,
+    kernel_symbol,
+    memrealtime,
+    rsrc,
+    traced,
+    uniform,
+    wave_max,
+)
+from vllm.models.minimax_m3.amd.mono.kernels.index_score import (
+    emit_index_scores,
+    index_scale_log2e,
+    step_rows,
+)
+
+FP8_MAX = 448.0
+ROWS_PER_TASK = 16
+K_CHUNKS = HIDDEN // 64  # 64-k chunks: one dwordx4 of weight per lane
+CHUNKS_PER_WAVE = K_CHUNKS // WAVES
+Q_OFF, K_OFF, V_OFF = 0, LOCAL_Q_HEADS * HEAD_DIM, (LOCAL_Q_HEADS + 1) * HEAD_DIM
+N_HEAD_TASKS = LOCAL_Q_HEADS + 4  # q heads, k, v, the rank's index_q, index_k
+IK_TASK = LOCAL_Q_HEADS + 3
+HEAD_ROW_GROUPS = HEAD_DIM // ROWS_PER_TASK  # GEMV tasks a head
+ELEMS = 4  # per lane in a head task: 32 lanes x 4 = 128, the aiter kernel's split
+PAGE_BYTES = PAGE16 * HEAD_DIM  # one page-16 of one kv head (fp8)
+
+
+def page16_of(slot):
+    """Page-16 id of a page-128 slot under the host's K/V block page stride."""
+    kps = mono_config.KV_PAGE_STRIDE
+    if kps == mono_config.SPARSE_BLOCK // PAGE16:
+        return slot // PAGE16
+    blk, t = slot // mono_config.SPARSE_BLOCK, slot % mono_config.SPARSE_BLOCK
+    return blk * kps + t // PAGE16
+
+
+# GEMV tasks, head tasks (every token of a head on one CTA) and the tokens' norm
+# tasks each own a CTA, in this order; a head CTA's half waves hold the tokens.
+# The rank's own rows (q | k | v | its index q | index k) are the first ONE_GEMV
+# tasks under indexer context parallelism too, and its other index q heads' the
+# OTHER_GEMV after the norm tasks, run only in a step that needs them: the CTA
+# layout of a step without long requests is the one-head build's.
+ONE_GEMV = qkv_rows(1) // ROWS_PER_TASK
+OTHER_GEMV = (TP - 1) * HEAD_ROW_GROUPS
+HT0 = ONE_GEMV
+NT0 = HT0 + N_HEAD_TASKS
+XT0 = NT0 + MAX_TOKENS
+assert XT0 + OTHER_GEMV <= BLOCKS
+assert MAX_TOKENS * 32 <= THREADS
+assert HIDDEN // 4 % THREADS == 0
+
+# The launch goes on the caller's current stream (graph capture included).
+_CURRENT_STREAM = fx.Stream(None)
+
+SCRATCH_QKV = 0  # (value, tag) pairs, one per qkv row
+SCRATCH_X8 = MAX_TOKENS * MAX_QKV_ROWS * 8  # S > 2: the norm tasks' x8 rows, words
+SCRATCH_X8S = SCRATCH_X8 + MAX_TOKENS * HIDDEN // 4 * 8  # and their scales
+# fused layer kernel only: a head task's / the residual's done flags
+SCRATCH_HDONE = SCRATCH_X8S + MAX_TOKENS * 8
+SCRATCH_RDONE = SCRATCH_HDONE + MAX_TOKENS * N_HEAD_TASKS * 8
+SCRATCH_BYTES = SCRATCH_RDONE + MAX_TOKENS * 8
+
+
+def _ld_bf16x4(r, k):
+    w = fx.Vector(bo.buffer_load(r, k // 2, vec_width=2, dtype=T.i32))
+    v = w.bitcast(fx.BFloat16).to(fx.Float32)
+    return [v[j] for j in range(4)]
+
+
+def ld_raw8(r, lt):
+    """Logical thread lt's 8 bf16 of a row, as loaded (for ``bf16x8_f32``)."""
+    return fx.Vector(bo.buffer_load(r, lt * 4, vec_width=4, dtype=T.i32))
+
+
+def bf16x8_f32(raw):
+    v = raw.bitcast(fx.BFloat16)
+    return [fx.Float32(v[j]) for j in range(8)]
+
+
+def head_rope_inputs(tk, gamma, d0, cos_sin, positions):
+    """Token tk's head gamma and cos / sin at dims d0 .. d0 + ELEMS (as f32)."""
+    # positions are int64: this token's low word
+    pos = fx.Int32(bo.buffer_load(rsrc(positions), 2 * tk, vec_width=1, dtype=T.i32))
+    half = ROTARY_DIM // 2
+    ib = (d0 < half).select(d0, d0 - half)
+    r_cs = rsrc(cos_sin)
+    return (
+        _ld_bf16x4(rsrc(gamma), d0),
+        _ld_bf16x4(r_cs, pos * ROTARY_DIM + fx.min(ib, half - ELEMS)),
+        _ld_bf16x4(r_cs, pos * ROTARY_DIM + half + fx.min(ib, half - ELEMS)),
+    )
+
+
+def head_norm_rope(e, lane, gw4, cs, sn, eps):
+    """One head's gemma RMSNorm + partial NeoX RoPE on a half wave, lane l of it
+    holding dims 4 (l % 32) .. + 4 (the aiter kernel's split and reduction order)."""
+    d0 = (lane % 32) * ELEMS
+    half = ROTARY_DIM // 2
+    first = d0 < half
+    ss = fx.Float32(0.0)
+    for j in range_constexpr(ELEMS):
+        ss = ss + e[j] * e[j]
+    ss = butterfly(ss, (16, 8, 4, 2, 1))
+    rr = hw_rsq(ss / float(HEAD_DIM) + eps)
+    n = [e[j] * rr * (gw4[j] + 1.0) for j in range(ELEMS)]
+    # partial NeoX RoPE over the first ROTARY_DIM dims: the partner is in this half
+    in_rope = d0 < ROTARY_DIM
+    partner_lane = first.select(lane + half // ELEMS, lane - half // ELEMS)
+    out = []
+    for j in range_constexpr(ELEMS):
+        p = fx.Int32(
+            rocdl.ds_bpermute(
+                T.i32,
+                (fx.min(fx.max(partner_lane, 0), 63) * 4).ir_value(),
+                n[j].bitcast(fx.Int32).ir_value(),
+            )
+        ).bitcast(fx.Float32)
+        # hipcc contracts ``x*c - p*s`` into fma(x, c, -(p*s)); match it
+        ps = p * sn[j]
+        rot = fmath.fma(n[j], cs[j], first.select(-ps, ps))
+        out.append(in_rope.select(rot, n[j]))
+    return out
+
+
+@traced
+def emit_pre_attn(
+    tid, bid, lane, wave, x8, red, mb_put, mb_put_words, mb_poll, qkv_mb, x8_mb,
+    x8s_mb, stamp, eps, tokens, heads, ptrs, signal=None, long_step=None,
+):  # fmt: skip
+    """K1's body: the GEMV tasks (ROWS_PER_TASK of the ``heads`` projection's rows
+    each: the rank's own on CTAs 0 .. ONE_GEMV, its other index q heads' on CTAs
+    XT0 .., in a step where ``long_step()``), the head tasks (CTAs HT0 .., a token
+    per half wave) and, with S > 2, one norm task per token (CTAs NT0 ..) feeding
+    the GEMV tasks."""
+    rows = heads.rows
+    (
+        ar,
+        res,
+        h_out, g_in, w_qkv, s_qkv, g_q, g_k, g_iq, g_ik, cos_sin, positions,
+        slot_mapping, k_cache, v_cache, k_scale, v_scale, index_cache, q_out, iq_out,
+    ) = ptrs  # fmt: skip
+
+    # stores K4 reads in the same launch (signal): device-coherent like the
+    # mailboxes, so the done flags need no L2 writeback
+    cm_out = CM_DEV if signal is not None else 0
+
+    def publish_done(mb_addr, idx, who=None):
+        """The CTA's stores done, then flag idx of ``mb_addr`` (from thread 0, or
+        from every thread where ``who``)."""
+        if const_expr(signal is not None):
+            fx.memory_fence(
+                syncscope=rocdl.SyncScope.Workgroup, ordering=fx.AtomicOrdering.Release
+            )
+            gpu.barrier()
+            if (tid == 0) if who is None else who:
+                mb_put(mb_addr, idx, fx.Int32(1))
+
+    # acc = bf16(sum of partials) + residual, f32; the norm reads acc unrounded.
+    # Elements are laid out as the fused all-reduce's logical threads (8 each): this
+    # lane holds lt0 and, in lanes < 32, lt1, so rstd reduces in its order.
+    lts = [wave * 96 + lane, wave * 96 + 64 + fx.min(lane, 31)]
+
+    def load_norm_inputs(toks):
+        """(gamma, [token][lt] -> (all-reduce, residual)) as loaded."""
+        raw_g = [ld_raw8(rsrc(g_in), lt) for lt in lts]
+        raw_ar = [
+            [
+                (
+                    ld_raw8(rsrc(ar), tk * (HIDDEN // 8) + lt),
+                    ld_raw8(rsrc(res), tk * (HIDDEN // 8) + lt),
+                )
+                for lt in lts
+            ]
+            for tk in toks
+        ]
+        return raw_g, raw_ar
+
+    def norm_quant(toks, raw_g, raw_ar, store_residual):
+        """Tokens toks: the residual out (where ``store_residual(j)``), gemma RMSNorm
+        and per-token FP8 into x8 row j (the j-th of toks) -> the x scales."""
+        gs = [bf16x8_f32(raw) for raw in raw_g]
+        accs = [  # [j][lt][8]
+            [
+                [a + r for a, r in zip(bf16x8_f32(ra), bf16x8_f32(rr))]
+                for ra, rr in raw_ar[j]
+            ]
+            for j in range(len(toks))
+        ]
+        for j in range_constexpr(len(toks)):
+            if store_residual(j):
+                r_h = rsrc(h_out)
+                bo.buffer_store(
+                    fx.Vector.from_elements(accs[j][0], fx.Float32).to(fx.BFloat16),
+                    r_h,
+                    toks[j] * HIDDEN + lts[0] * 8,
+                    cache_modifier=cm_out,
+                )
+                if lane < 32:
+                    bo.buffer_store(
+                        fx.Vector.from_elements(accs[j][1], fx.Float32).to(fx.BFloat16),
+                        r_h,
+                        toks[j] * HIDDEN + lts[1] * 8,
+                        cache_modifier=cm_out,
+                    )
+        # every token's sum of squares behind one set of barriers
+        tots = ar_block_sums(
+            [
+                (
+                    ar_pack_sumsq(accs[j][0]),
+                    (lane < 32).select(ar_pack_sumsq(accs[j][1]), fx.Float32(0.0)),
+                )
+                for j in range(len(toks))
+            ],
+            lane,
+            wave,
+            red,
+        )
+        stamp(1)
+        normed = []
+        for j in range_constexpr(len(toks)):
+            rstd = hw_rsq(tots[j] / float(HIDDEN) + eps)
+            normed.append(
+                [
+                    [accs[j][i][e] * rstd * (gs[i][e] + 1.0) for e in range(8)]
+                    for i in range(2)
+                ]
+            )
+        for j in range_constexpr(len(toks)):
+            amax = fx.Float32(0.0)
+            for e in range_constexpr(8):
+                amax = fx.max(amax, fmath.absf(normed[j][0][e]))
+                amax = fx.max(
+                    amax,
+                    (lane < 32).select(fmath.absf(normed[j][1][e]), fx.Float32(0.0)),
+                )
+            w_max = wave_max(amax)
+            if lane == 0:
+                fx.ptr_store(w_max, red + (j * WAVES + wave))
+        gpu.barrier()
+        x_scales = []
+        for j in range_constexpr(len(toks)):
+            amax = fx.ptr_load(red + j * WAVES)
+            for w in range_constexpr(1, WAVES):
+                amax = fx.max(amax, fx.ptr_load(red + (j * WAVES + w)))
+            x_scale = (amax == 0.0).select(fx.Float32(1.0), amax / FP8_MAX)
+            x_scales.append(x_scale)
+            x_rcp = 1.0 / x_scale
+            for i in range_constexpr(2):
+                q = [div_rn(normed[j][i][e], x_scale, x_rcp) for e in range(8)]
+                w01 = [
+                    fp8_pack4(q[0], q[1], q[2], q[3]),
+                    fp8_pack4(q[4], q[5], q[6], q[7]),
+                ]
+                xo = x8 + (j * (HIDDEN // 4) + lts[i] * 2)
+                if (i == 0) | (lane < 32):
+                    fx.ptr_store(w01[0].bitcast(fx.Float32), xo)
+                    fx.ptr_store(w01[1].bitcast(fx.Float32), xo + 1)
+        gpu.barrier()
+        stamp(2)
+        return x_scales
+
+    # ============================================ norm tasks (S > 2)
+    # (S <= 2 the GEMV tasks normalize every token themselves: S = 2 fused layer
+    # 43.8 -> 43.6 us; S = 4 57.5 -> 59.4 us the other way round)
+    if const_expr(tokens > 2):  # noqa: SIM102
+        if (bid >= NT0) & (bid < NT0 + tokens):
+            tk = bid - NT0
+            raw_g, raw_ar = load_norm_inputs([tk])
+            x_scale = norm_quant([tk], raw_g, raw_ar, lambda j: True)[0]
+            if const_expr(signal is not None):
+                publish_done(signal[1], tk)  # the residual
+            # the x8 row as plain words, then its scale: the scale's pair is the flag
+            for i in range_constexpr(2):
+                if (i == 0) | (lane < 32):
+                    bo.buffer_store(
+                        fx.Vector(
+                            fx.ptr_load(
+                                x8 + lts[i] * 2,
+                                result_type=fx.Vector.make_type(2, fx.Int32),
+                            )
+                        ),
+                        rsrc(x8_mb),
+                        tk * (HIDDEN // 4) + lts[i] * 2,
+                        cache_modifier=CM_DEV,
+                    )
+            fx.memory_fence(
+                syncscope=rocdl.SyncScope.Workgroup, ordering=fx.AtomicOrdering.Release
+            )
+            gpu.barrier()
+            if tid == 0:
+                mb_put(x8s_mb, tk, x_scale)
+
+    # ============================================ GEMV tasks (one per CTA)
+    # CTA bid < ONE_GEMV: row group bid of q | k | v, then the own index q head's
+    # and index k's (the identity in the one-head build)
+    own_q = (LOCAL_Q_HEADS + 2) * HEAD_ROW_GROUPS
+    rg = (bid < own_q).select(
+        bid,
+        (bid < own_q + HEAD_ROW_GROUPS).select(
+            heads.iq_off // ROWS_PER_TASK + bid - own_q,
+            heads.ik_off // ROWS_PER_TASK + bid - own_q - HEAD_ROW_GROUPS,
+        ),
+    )
+    gemv = bid < ONE_GEMV
+    if const_expr(heads.count > 1):
+        xt = bid - XT0
+        if (xt >= 0) & (xt < OTHER_GEMV):
+            other = xt // HEAD_ROW_GROUPS  # the other heads in head order, own skipped
+            head = other + (other >= heads.own).select(1, 0)
+            rg = (LOCAL_Q_HEADS + 2 + head) * HEAD_ROW_GROUPS + xt % HEAD_ROW_GROUPS
+            gemv = long_step()
+    if gemv:
+        if const_expr(tokens <= 2):
+            raw_g, raw_ar = load_norm_inputs(list(range(tokens)))
+        # weight stream after the norm's inputs: the norm waits on its loads in issue
+        # order (vmcnt), so it would otherwise wait for the 96 KB of weights too; the
+        # scheduling barrier keeps the compiler from interleaving them again (nothing
+        # before it uses a loaded value, so it does not wait either)
+        rocdl.sched_barrier(0)
+        r_w = rsrc(w_qkv)
+        # 64-k chunk kc of row group rg is 1 KB at ((rg*K/32 + 2kc)*512); lane l
+        # holds row l%16, k = 64kc + 16(l/16) + [0, 16).
+        wts = []
+        for c in range_constexpr(CHUNKS_PER_WAVE):
+            kc = wave * CHUNKS_PER_WAVE + c
+            wts.append(
+                fx.Vector(
+                    bo.buffer_load(
+                        r_w,
+                        ((rg * (HIDDEN // 32) + kc * 2) * 512 + lane * 16) // 4,
+                        vec_width=4,
+                        dtype=T.i32,
+                    )
+                )
+            )
+        ws = fx.Float32(  # row tid % 16's scale, sampled below by tid < 16 tokens
+            bo.buffer_load(
+                rsrc(s_qkv),
+                rg * ROWS_PER_TASK + tid % ROWS_PER_TASK,
+                vec_width=1,
+                dtype=T.f32,
+            )
+        )
+        if const_expr(tokens <= 2):
+            x_scales = norm_quant(
+                list(range(tokens)), raw_g, raw_ar, lambda j: bid == j
+            )
+            if const_expr(signal is not None):  # noqa: SIM102
+                if bid < tokens:
+                    publish_done(signal[1], bid)  # the residual
+        else:
+            # the norm tasks' scales (each published after its row's plain x8 words),
+            # then every row, 16 B a load
+            got = mb_poll([(x8s_mb, tk, 1) for tk in range(tokens)])
+            x_scales = [got[tk][0].bitcast(fx.Float32) for tk in range(tokens)]
+            fx.memory_fence(
+                syncscope=rocdl.SyncScope.Workgroup, ordering=fx.AtomicOrdering.Acquire
+            )
+            n_v = tokens * HIDDEN // 16  # 16 B loads
+            for i in range_constexpr((n_v + THREADS - 1) // THREADS):
+                v = tid + THREADS * i
+                if v < n_v:
+                    fx.ptr_store(
+                        fx.Vector(
+                            bo.buffer_load(
+                                rsrc(x8_mb),
+                                v * 4,
+                                vec_width=4,
+                                dtype=T.i32,
+                                cache_modifier=CM_DEV,
+                            )
+                        ),
+                        x8 + v * 4,
+                    )
+            gpu.barrier()
+            stamp(2)
+        # 12 chunks per wave: two 16x16x32 FP8 MFMAs each (k halves of 8 bytes).
+        # B column l % 16 is token l % 16 (token 0 past the last): C column t is
+        # token t's rows, the weight stream is the same for any token count.
+        col_tok = fx.min(lane % 16, tokens - 1)
+        c = fx.Vector.filled(4, 0.0, fx.Float32)
+        for cc in range_constexpr(CHUNKS_PER_WAVE):
+            kc = wave * CHUNKS_PER_WAVE + cc
+            xb = fx.Vector(
+                fx.ptr_load(
+                    x8 + (col_tok * (HIDDEN // 4) + kc * 16 + (lane // 16) * 4),
+                    result_type=fx.Vector.make_type(4, fx.Float32),
+                )
+            ).bitcast(fx.Int64)
+            wa = wts[cc].bitcast(fx.Int64)
+            for h in range_constexpr(2):
+                c = fx.Vector(
+                    rocdl.mfma_f32_16x16x32_fp8_fp8(
+                        T.vec(4, T.f32), [wa[h], xb[h], c, 0, 0, 0]
+                    )
+                )
+        fx.ptr_store(c, red + (wave * 64 + lane) * 4)
+        gpu.barrier()
+        # sample column t of C: lanes t, 16 + t, 32 + t, 48 + t hold rows 4 (l/16) + e
+        if tid < ROWS_PER_TASK * tokens:
+            r = tid % ROWS_PER_TASK
+            tk = tid // ROWS_PER_TASK
+            tot_r = fx.Float32(0.0)
+            for w in range_constexpr(WAVES):
+                tot_r = tot_r + fx.ptr_load(
+                    red + ((w * 64 + 16 * (r // 4) + tk) * 4 + r % 4)
+                )
+            x_scale = x_scales[0]
+            for k in range_constexpr(1, tokens):
+                x_scale = (tk == k).select(x_scales[k], x_scale)
+            mb_put(
+                qkv_mb,
+                tk * rows + rg * ROWS_PER_TASK + r,
+                bf16_round(tot_r * x_scale * ws),
+            )
+
+        stamp(3)
+
+    # ======================================= head tasks: norm / rope / cache
+    # one CTA per head, token k on the half wave of threads 32 k .. 32 k + 31
+    if (bid >= HT0) & (bid < HT0 + N_HEAD_TASKS):
+        ht = bid - HT0
+        tk = tid // 32
+        l32 = tid % 32
+        is_q = ht < LOCAL_Q_HEADS
+        is_k = ht == LOCAL_Q_HEADS
+        is_v = ht == LOCAL_Q_HEADS + 1
+        is_iq = ht == LOCAL_Q_HEADS + 2
+        base = tk * rows + is_q.select(
+            ht * HEAD_DIM,
+            is_k.select(
+                K_OFF, is_v.select(V_OFF, is_iq.select(heads.iq_off, heads.ik_off))
+            ),
+        )
+        if tk < tokens:
+            d0 = l32 * ELEMS
+            # slot_mapping is int64: this token's low word
+            slot = fx.Int32(
+                bo.buffer_load(rsrc(slot_mapping), 2 * tk, vec_width=1, dtype=T.i32)
+            )
+            # gamma and this position's cos / sin go out before the wait
+            gw = is_q.select(g_q, is_k.select(g_k, is_iq.select(g_iq, g_ik)))
+            gw4, cs, sn = head_rope_inputs(tk, gw, d0, cos_sin, positions)
+            got = mb_poll([(qkv_mb, base + d0, 2), (qkv_mb, base + d0 + 2, 2)])
+            stamp(4)
+            e = [
+                got[0][0].bitcast(fx.Float32),
+                got[0][1].bitcast(fx.Float32),
+                got[1][0].bitcast(fx.Float32),
+                got[1][1].bitcast(fx.Float32),
+            ]
+            if is_v:
+                # raw V: per-token amax over the head, scale = amax / FP8_MAX
+                if const_expr(mono_config.STATIC_KV_SCALE):
+                    vs = fx.Float32(1.0)
+                else:
+                    am = fx.Float32(0.0)
+                    for j in range_constexpr(ELEMS):
+                        am = fx.max(am, fmath.absf(e[j]))
+                    am = butterfly(am, (16, 8, 4, 2, 1), fx.max)
+                    vs = (am > 0.0).select(am / FP8_MAX, fx.Float32(1.0))
+                if slot >= 0:
+                    b = page16_of(slot)
+                    t = slot % PAGE16
+                    if l32 == 0:
+                        bo.buffer_store(
+                            vs, rsrc(v_scale), b * PAGE16 + t, cache_modifier=cm_out
+                        )
+                    r_v = rsrc(v_cache)
+                    for j in range_constexpr(ELEMS):
+                        q8 = fp8_pack4(e[j] / vs, 0.0, 0.0, 0.0) & 0xFF
+                        bo.buffer_store(
+                            fx.Int8(q8),
+                            r_v,
+                            b * PAGE_BYTES + (d0 + j) * PAGE16 + t,
+                            cache_modifier=cm_out,
+                        )
+            else:
+                out = head_norm_rope(e, lane, gw4, cs, sn, eps)
+                if is_q:
+                    bo.buffer_store(
+                        fx.Vector.from_elements(out, fx.Float32).to(fx.BFloat16),
+                        rsrc(q_out),
+                        (tk * LOCAL_Q_HEADS + ht) * HEAD_DIM + d0,
+                        cache_modifier=cm_out,
+                    )
+                if is_iq:
+                    bo.buffer_store(
+                        fx.Vector.from_elements(out, fx.Float32).to(fx.BFloat16),
+                        rsrc(iq_out),
+                        tk * HEAD_DIM + d0,
+                    )
+                if is_k:
+                    if const_expr(mono_config.STATIC_KV_SCALE):
+                        ksc = fx.Float32(1.0)
+                    else:
+                        am = fx.Float32(0.0)
+                        for j in range_constexpr(ELEMS):
+                            am = fx.max(am, fmath.absf(out[j]))
+                        am = butterfly(am, (16, 8, 4, 2, 1), fx.max)
+                        ksc = (am > 0.0).select(am / FP8_MAX, fx.Float32(1.0))
+                    if slot >= 0:
+                        b = page16_of(slot)
+                        t = slot % PAGE16
+                        if l32 == 0:
+                            bo.buffer_store(
+                                ksc,
+                                rsrc(k_scale),
+                                b * PAGE16 + t,
+                                cache_modifier=cm_out,
+                            )
+                        qk = [bf16_round(out[j]) / ksc for j in range(ELEMS)]
+                        bo.buffer_store(
+                            fp8_pack4(qk[0], qk[1], qk[2], qk[3]),
+                            rsrc(k_cache),
+                            (
+                                b * PAGE_BYTES
+                                + (d0 // 16) * (PAGE16 * 16)
+                                + t * 16
+                                + d0 % 16
+                            )
+                            // 4,
+                            cache_modifier=cm_out,
+                        )
+                if (ht == IK_TASK) & (slot >= 0):
+                    qi = [bf16_round(out[j]) for j in range(ELEMS)]
+                    bo.buffer_store(
+                        fp8_pack4(qi[0], qi[1], qi[2], qi[3]),
+                        rsrc(index_cache),
+                        (slot * HEAD_DIM + d0) // 4,
+                        cache_modifier=cm_out,
+                    )
+
+        if const_expr(signal is not None):
+            # every token's flag once the whole CTA's stores are out
+            publish_done(signal[0], tid * N_HEAD_TASKS + ht, tid < tokens)
+
+
+# the fused layer kernel's K1 pointers that do not change per step, in its device
+# array (the rest are K4's own, or kernel arguments)
+K1_ARGS = (
+    "ar", "g_in", "w_qkv", "s_qkv", "g_q", "g_k", "g_iq", "g_ik", "cos_sin",
+    "index_cache", "iq_out", "scratch",
+)  # fmt: skip
+
+
+@traced
+def emit_k1(
+    tid, bid, lane, wave, x8, red, qs, stamp, step, layer, eps, index_scale, tokens,
+    q_len, init_blocks, local_blocks, heads, ptrs, bt_width, signal=False,
+):  # fmt: skip
+    """K1 whole: the GEMV / head / norm tasks, then every CTA's indexer scores.
+    ``x8``: tokens * HIDDEN / 4 words of LDS, ``red``: WAVES * 256 floats, ``qs``:
+    tokens * heads.count * HEAD_DIM / 2 floats. ``q_len``: tokens per request
+    (speculative verify rows). ``heads``: the index q heads of the fused projection
+    (``IndexHeads``). ``signal`` (the fused layer kernel): the head tasks and
+    residual writers publish done flags (SCRATCH_HDONE / SCRATCH_RDONE)."""
+    (
+        ar, res, h_out, g_in, w_qkv, s_qkv, g_q, g_k, g_iq, g_ik, cos_sin, positions,
+        slot_mapping, k_cache, v_cache, k_scale, v_scale, index_cache, q_out, iq_out,
+        block_table, seq_lens, iscore, scratch,
+    ) = ptrs  # fmt: skip
+    mbox = Mailbox(scratch, step, layer)
+    qkv_mb = mbox.addr(SCRATCH_QKV)
+    x8_mb = mbox.addr(SCRATCH_X8)
+    x8s_mb = mbox.addr(SCRATCH_X8S)
+    # Plain names, not method calls: flydsl's AST rewriter carries every local whose
+    # method is called inside a dynamic if/for as loop state.
+    mb_put = mbox.put
+    mb_put_words = mbox.put_words
+    mb_poll = mbox.poll
+
+    def index_q_row(tk, head):
+        """Token tk's index q of the projection's index head ``head`` into qs (row
+        tk * heads.count + head, HEAD_DIM / 2 words), made from the GEMV's raw rows
+        with the index-q head's math, by this wave's lanes < 32."""
+        d0 = lane * ELEMS
+        gw4, cs, sn = head_rope_inputs(tk, g_iq, d0, cos_sin, positions)
+        base = tk * heads.rows + IndexHeads(heads.count, 0).iq_off + head * HEAD_DIM
+        got = mb_poll([(qkv_mb, base + d0, 2), (qkv_mb, base + d0 + 2, 2)])
+        e = [got[h][j].bitcast(fx.Float32) for h in range(2) for j in range(2)]
+        out = head_norm_rope(e, lane, gw4, cs, sn, eps)
+        row = qs + (tk * heads.count + head) * (HEAD_DIM // 2)
+        fx.ptr_store(bf16_pair(out[0], out[1]), row + lane * 2)
+        fx.ptr_store(bf16_pair(out[2], out[3]), row + (lane * 2 + 1))
+
+    def make_index_q(first, last):
+        """Tokens first..last's index q of this rank's head into qs: wave w makes
+        tokens w, w + WAVES, ..; run by a CTA with score tasks only."""
+        for k0 in range_constexpr(0, tokens, WAVES):
+            tk = wave + k0
+            if (tk < tokens) & (tk >= first) & (tk <= last) & (lane < 32):
+                index_q_row(tk, heads.own)
+        gpu.barrier()
+
+    def make_index_q_heads():
+        """Every token's index q of every head into qs (indexer context
+        parallelism): (token, head) job j on wave j % WAVES."""
+        for j0 in range_constexpr(0, tokens * heads.count, WAVES):
+            j = wave + j0
+            if (j < tokens * heads.count) & (lane < 32):
+                index_q_row(j // heads.count, j % heads.count)
+        gpu.barrier()
+
+    def q_frag_head(tk, s, head):
+        """This lane's MFMA B operand of token tk's index q of head ``head``,
+        k-chunk s."""
+        row = (tk * heads.count + head) * (HEAD_DIM // 2)
+        return fx.Vector(
+            fx.ptr_load(
+                qs + (row + (32 * s + 8 * (lane // 16)) // 2),
+                result_type=fx.Vector.make_type(4, fx.Float32),
+            )
+        ).bitcast(fx.BFloat16)
+
+    def q_frag(tk, s):
+        """This lane's MFMA B operand of token tk's index q (this rank's head),
+        k-chunk s."""
+        return q_frag_head(tk, s, heads.own)
+
+    def long_step():
+        """A request of the step is long (``step_rows``): it scores every index q
+        head (indexer context parallelism)."""
+        return step_rows(seq_lens, tokens, q_len, heads).any_long()
+
+    emit_pre_attn(
+        tid, bid, lane, wave, x8, red, mb_put, mb_put_words, mb_poll, qkv_mb, x8_mb,
+        x8s_mb, stamp, eps, tokens, heads,
+        (
+            ar, res, h_out, g_in, w_qkv, s_qkv, g_q, g_k, g_iq, g_ik,
+            cos_sin, positions, slot_mapping, k_cache, v_cache, k_scale,
+            v_scale, index_cache, q_out, iq_out,
+        ),
+        (mbox.addr(SCRATCH_HDONE), mbox.addr(SCRATCH_RDONE)) if signal else None,
+        long_step,
+    )  # fmt: skip
+    stamp(5)
+    # the scorers: the idle CTAs first (the norm tasks' come first but finish
+    # early), then the GEMV CTAs, the head CTAs last
+    score_base = ONE_GEMV + N_HEAD_TASKS
+    hdone_mb = mbox.addr(SCRATCH_HDONE)
+
+    def wait_new_keys():
+        """Every token's index_k head task done, its key visible to this wave."""
+        mb_poll([(hdone_mb, fx.min(lane, tokens - 1) * N_HEAD_TASKS + IK_TASK, 1)])
+        fx.memory_fence(
+            syncscope=rocdl.SyncScope.Agent, ordering=fx.AtomicOrdering.Acquire
+        )
+
+    emit_index_scores(
+        (bid + (BLOCKS - score_base)) % BLOCKS, lane, uniform(wave), red,
+        mb_put, iscore, q_frag, tokens, q_len, init_blocks, local_blocks,
+        index_scale, index_cache, block_table, seq_lens, bt_width,
+        wait_new_keys if signal else None,
+        make_index_q, heads, q_frag_head, make_index_q_heads,
+    )  # fmt: skip
+    stamp(6)
+
+
+TL_POINTS = 8  # timeline stamps per CTA
+
+
+def build_pre_attn_kernel(
+    eps: float,
+    sm_scale: float,
+    init_blocks: int,
+    local_blocks: int,
+    tokens: int = 1,
+    timeline: bool = False,
+    heads: IndexHeads = ONE_INDEX_HEAD,
+):
+    """``@flyc.jit`` launcher of K1 for this model's RMSNorm epsilon, attention
+    scale and sparse pinned blocks, a decode batch of ``tokens`` (<= MAX_TOKENS)
+    rows and the fused projection's index q ``heads``."""
+    assert 1 <= tokens <= MAX_TOKENS
+    index_scale = index_scale_log2e(sm_scale)
+
+    @fx.struct
+    class Smem:
+        # FP8 activations, natural k order, one row per token
+        x8: fx.Array[fx.Int32, tokens * HIDDEN // 4, 16]
+        red: fx.Array[fx.Float32, WAVES * 64 * 4, 16]
+        tls: fx.Array[fx.Int64, TL_POINTS if timeline else 1, 16]  # timeline stamps
+        # every token's index q for the scorers, bf16 pairs
+        qs: fx.Array[fx.Float32, HEAD_DIM // 2 * tokens * heads.count, 16]
+
+    kernel_name = kernel_symbol(
+        "minimax_m3_pre_attn",
+        s=tokens,
+        ib=init_blocks,
+        lb=local_blocks,
+        ih=heads.count,
+        io=heads.own,
+        tl=timeline,
+        kps=mono_config.KV_PAGE_STRIDE,
+        sks=mono_config.STATIC_KV_SCALE,
+    )
+
+    # the JIT cache key holds the kernel's scalar closure values, not objects: the
+    # index heads go in as ints (a shared cache otherwise serves one rank's build
+    # to every rank)
+    ih_count, ih_own = heads.count, heads.own
+
+    @flyc.kernel(name=kernel_name, known_block_size=[THREADS, 1, 1])
+    def pre_attn_kernel(
+        ar: Int64,
+        res: Int64,
+        h_out: Int64,
+        g_in: Int64,
+        w_qkv: Int64,
+        s_qkv: Int64,
+        g_q: Int64,
+        g_k: Int64,
+        g_iq: Int64,
+        g_ik: Int64,
+        cos_sin: Int64,
+        positions: Int64,
+        slot_mapping: Int64,
+        k_cache: Int64,
+        v_cache: Int64,
+        k_scale: Int64,
+        v_scale: Int64,
+        index_cache: Int64,
+        q_out: Int64,
+        iq_out: Int64,
+        block_table: Int64,
+        seq_lens: Int64,
+        bt_width: Int32,
+        q_len: Int32,
+        iscore: Int64,
+        scratch: Int64,
+        step: Int64,
+        layer: Int32,
+        tl: Int64,
+    ):
+        tid = fx.thread_idx.x
+        bid = fx.block_idx.x
+        lane = tid % 64
+        wave = tid // 64
+        lds = fx.SharedAllocator().allocate(Smem).peek()
+        x8 = lds.x8.ptr
+        red = lds.red.ptr
+        tls = lds.tls.ptr
+        qs = lds.qs.ptr
+
+        def stamp(k):
+            """``timeline``: s_memrealtime (100 MHz) of this CTA passing point k, kept
+            in LDS until the kernel ends."""
+            # The compile-time and traced conditions cannot share one ``and``.
+            if const_expr(timeline):  # noqa: SIM102
+                if tid == 0:
+                    fx.ptr_store(memrealtime(), tls + k)
+
+        if const_expr(timeline):  # noqa: SIM102
+            if tid < TL_POINTS:  # same wave as the stamping thread: ordered
+                fx.ptr_store(fx.Int64(0), tls + tid)
+        stamp(0)
+        emit_k1(
+            tid, bid, lane, wave, x8, red, qs, stamp, step, layer, eps, index_scale,
+            tokens, q_len, init_blocks, local_blocks, IndexHeads(ih_count, ih_own),
+            (
+                ar, res, h_out, g_in, w_qkv, s_qkv, g_q, g_k, g_iq, g_ik,
+                cos_sin, positions, slot_mapping, k_cache, v_cache, k_scale,
+                v_scale, index_cache, q_out,
+                iq_out, block_table, seq_lens, iscore, scratch,
+            ),
+            bt_width,
+        )  # fmt: skip
+        if const_expr(timeline):  # noqa: SIM102
+            if tid < TL_POINTS:
+                fx.generic_store(
+                    fx.inttoptr(
+                        fx.PointerType.get(fx.Int64.ir_type, fx.AddressSpace.Global, 8),
+                        tl + fx.Int64((bid * TL_POINTS + tid) * 8),
+                    ),
+                    fx.ptr_load(tls + tid),
+                )
+
+    @flyc.jit
+    def launch_pre_attn(
+        ar: Int64,
+        res: Int64,
+        h_out: Int64,
+        g_in: Int64,
+        w_qkv: Int64,
+        s_qkv: Int64,
+        g_q: Int64,
+        g_k: Int64,
+        g_iq: Int64,
+        g_ik: Int64,
+        cos_sin: Int64,
+        positions: Int64,
+        slot_mapping: Int64,
+        k_cache: Int64,
+        v_cache: Int64,
+        k_scale: Int64,
+        v_scale: Int64,
+        index_cache: Int64,
+        q_out: Int64,
+        iq_out: Int64,
+        block_table: Int64,
+        seq_lens: Int64,
+        bt_width: Int32,
+        q_len: Int32,
+        iscore: Int64,
+        scratch: Int64,
+        step: Int64,
+        layer: Int32,
+        tl: Int64,
+        stream: fx.Stream = _CURRENT_STREAM,
+    ):
+        pre_attn_kernel(
+            ar,
+            res,
+            h_out,
+            g_in,
+            w_qkv,
+            s_qkv,
+            g_q,
+            g_k,
+            g_iq,
+            g_ik,
+            cos_sin,
+            positions,
+            slot_mapping,
+            k_cache,
+            v_cache,
+            k_scale,
+            v_scale,
+            index_cache,
+            q_out,
+            iq_out,
+            block_table,
+            seq_lens,
+            bt_width,
+            q_len,
+            iscore,
+            scratch,
+            step,
+            layer,
+            tl,
+        ).launch(grid=(BLOCKS,), block=(THREADS,), stream=stream)
+
+    return launch_pre_attn
diff --git a/vllm/models/minimax_m3/amd/mono/layout.py b/vllm/models/minimax_m3/amd/mono/layout.py
new file mode 100644
index 000000000000..bd186e2d1f47
--- /dev/null
+++ b/vllm/models/minimax_m3/amd/mono/layout.py
@@ -0,0 +1,234 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+# SPDX-License-Identifier: MIT
+# Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved.
+"""Where the mono layer kernel (K4) puts its tasks and data, per decode token count.
+
+Plain Python, evaluated at build time: which CTA runs each stage's tasks, the
+scratch and peer-symmetric regions, and the LDS row pitches. Kept free of FlyDSL
+so it is testable without a GPU.
+"""
+
+from vllm.models.minimax_m3.amd.mono.config import (
+    BLOCKS,
+    HEAD_DIM,
+    HIDDEN,
+    INTER,
+    LOCAL_Q_HEADS,
+    MAX_INDEX_BLOCKS,
+    MAX_SPARSE_KEYS,
+    MAX_TOKENS,
+    MOE_SLOTS,
+    N_ROUTED,
+    PAGE16,
+    SPARSE_BLOCK,
+    TOP_K,
+    TOPK_BLOCKS,
+    WAVES,
+)
+
+H = LOCAL_Q_HEADS
+SPLIT_KEYS = 256  # one gluon context partition per split task
+N_SPLIT = MAX_SPARSE_KEYS // SPLIT_KEYS
+O_K = H * HEAD_DIM
+O_ROWS = 32  # hidden rows per o_proj / attention-reduce task
+# from this many tokens the attention all-reduce is a reduce-scatter to the row
+# group's owner rank (t % npes) plus an all-gather of the sums: half the bytes over
+# each link for one more hop (smaller batches are latency-bound: one hop wins)
+AG_RS_FROM = 12
+N_O = HIDDEN // O_ROWS
+PAGES_PER_BLOCK = SPARSE_BLOCK // PAGE16
+N_ROUTER = N_ROUTED // WAVES  # one expert per wave
+UG_ROWS = 16
+UG_PER_SLOT = INTER // UG_ROWS
+SH_ROWS = 8  # mid rows per shared-expert up / gate task (a half task)
+SH_TASKS = INTER // SH_ROWS
+DN_ROWS = 32
+N_DN = HIDDEN // DN_ROWS
+# The router logits are written SCORE_COPIES times and CTA b polls copy
+# b % SCORE_COPIES: 240 expert CTAs spinning on one 1 KB region serialize on its
+# lines. K4 on 4 GPUs 37.3 -> 35.9 us; copies of xn / mid measured neutral and of
+# the split outputs +0.3 us (the split CTAs pay the extra stores).
+SCORE_COPIES = 8
+
+
+def _layout(items, align=256):
+    """Byte offsets of (value, tag) pair regions, 8 bytes per pair."""
+    off, out = 0, {}
+    for name, pairs in items:
+        out[name] = off
+        off += (pairs * 8 + align - 1) // align * align
+    out["_bytes"] = off
+    return out
+
+
+# plain words, not pairs: the selected pages then the context length, as the
+# original selector emits them (read by the tests); one row per token
+SPARSE_TABLE_WORDS = TOPK_BLOCKS * PAGES_PER_BLOCK + 2
+# a context past THREADS blocks, one-head build: each split task's TOPK_BLOCKS
+# best (key, block) (indexer context parallelism: the peers' sel_cp)
+SEL_LONG_PAIRS = N_SPLIT * TOPK_BLOCKS * 2
+
+# every region holds MAX_TOKENS rows (token k at k * the one-token size)
+SCRATCH = _layout(
+    [
+        # indexer block scores, token k's block b at pair k * MAX_INDEX_BLOCKS + b
+        ("iscore", MAX_INDEX_BLOCKS * MAX_TOKENS),
+        ("sparse_table", SPARSE_TABLE_WORDS // 2 * MAX_TOKENS),
+        ("sp_o", N_SPLIT * H * HEAD_DIM // 2 * MAX_TOKENS),  # packed bf16 pairs
+        ("sp_m", N_SPLIT * H * MAX_TOKENS),
+        ("sp_l", N_SPLIT * H * MAX_TOKENS),
+        ("attnq", O_K // 4 * MAX_TOKENS // 2),  # merged attention, plain FP8 words
+        ("attnq_s", MAX_TOKENS),  # and its per-token scale
+        # the finished attention all-reduce + residual, bf16 pairs (o tasks -> router)
+        ("a", HIDDEN // 2 * MAX_TOKENS),
+        ("xn8", HIDDEN // 4 * MAX_TOKENS),  # MXFP8 xn, 4 fp8 per word
+        ("xsc", HIDDEN // 32 * MAX_TOKENS),  # and its per-1x32 E8M0 scales
+        # copy c, token k at pairs [(c * MAX_TOKENS + k) * N_ROUTED, ..)
+        # per expert (routing key, sigmoid) pairs, see route_key
+        ("scores", 2 * N_ROUTED * MAX_TOKENS * SCORE_COPIES),
+        ("mid", MOE_SLOTS * INTER * MAX_TOKENS),
+        # routing record read by the tests: token k's expert ids, then weights
+        ("sel", 2 * MOE_SLOTS * MAX_TOKENS),
+        # token k's split task p's candidates: rank r at pairs 2 ((k N_SPLIT + p)
+        # TOPK_BLOCKS + r) + (0 key, 1 block)
+        ("sel_long", SEL_LONG_PAIRS * MAX_TOKENS),
+        # wide: the MXFP8 xn and scales as plain words (every token's row), and a
+        # done flag per router task
+        ("xn8p", HIDDEN // 4 * MAX_TOKENS // 2),
+        ("xscp", HIDDEN // 32 * MAX_TOKENS // 2),
+        ("xndone", N_ROUTER * MAX_TOKENS),
+        # and the mid rows as MXFP8 words, their E8M0 scales a word each (the down
+        # stage's LDS layout), a done flag per up / gate task pair
+        ("mid8p", MOE_SLOTS * INTER * MAX_TOKENS // 8),
+        ("midsp", MOE_SLOTS * INTER // 32 * MAX_TOKENS // 2),
+        ("ugdone", (TOP_K * MAX_TOKENS + 1) * UG_PER_SLOT // 2),
+    ]
+)
+SCRATCH_BYTES = SCRATCH["_bytes"]
+
+
+def sym_layout(npes: int):
+    """Per-rank symmetric buffer: attn / ffn partial regions,
+    `[src rank][token][HIDDEN]`, and the peers' index-selection candidates
+    (``sel_cp``)."""
+    rows = npes * MAX_TOKENS * HIDDEN // 2
+    # a_ag / ffn_ag: the finished attention / FFN all-reduce rows, each written by
+    # its owner rank (the reduce-scatter + all-gather form, S >= AG_RS_FROM)
+    ag = MAX_TOKENS * HIDDEN // 2
+    return _layout(
+        [
+            ("attn", rows),
+            ("ffn", rows),
+            ("a_ag", ag),
+            ("ffn_ag", ag),
+            # indexer context parallelism: token k's candidates for this rank's
+            # index head, source rank s's half h's r-th best at pairs
+            # 2 ((k N_SPLIT + 2 s + h) TOPK_BLOCKS + r) + (0 key, 1 block)
+            ("sel_cp", SEL_LONG_PAIRS * MAX_TOKENS),
+        ]
+    )
+
+
+def stage_bases(tokens: int):
+    """Task t of a stage runs on CTA (base + t) % BLOCKS; the merge (S > 1), o and
+    down tasks start past the tokens' split CTAs."""
+    n_split = N_SPLIT * tokens
+    n_merge = tokens if tokens > 1 else 0
+    router = (n_split + n_merge + N_O) % BLOCKS
+    return {
+        "split": 0,
+        "merge": n_split,
+        "o": n_split + n_merge,
+        "router": router,
+        "shared": _shared_base(tokens, router),
+        "ug": 0,
+        "down": n_split,
+    }
+
+
+def ug_ctas_per_token(tokens: int):
+    """CTAs per token of the routed up / gate tasks, token k on CTAs k C .. (k + 1)
+    C - 1: a CTA stages one token's xn, not all (fused layer, S = 4: 57.9 -> 57.6
+    us). None: round-robin -- S < 4 (S = 2: +0.4 us) or S not dividing the grid."""
+    return BLOCKS // tokens if tokens >= 4 and BLOCKS % tokens == 0 else None
+
+
+def ug_tasks_of(cta: int, tokens: int) -> int:
+    """Routed up / gate tasks CTA ``cta`` runs."""
+    per_token = TOP_K * UG_PER_SLOT
+    cpt = ug_ctas_per_token(tokens)
+    if cpt is None:
+        return (per_token * tokens - cta + BLOCKS - 1) // BLOCKS
+    return (per_token - cta % cpt + cpt - 1) // cpt
+
+
+def _shared_base(tokens: int, router: int) -> int:
+    """First of the SH_TASKS consecutive CTAs that run the shared expert's up /
+    gate before the routing lands. Ranked by the most routed up / gate tasks a
+    window CTA carries (they queue behind it); then, only if that is the most any
+    CTA carries, by its router CTAs (they start it late and would delay the last
+    routed tasks: S = 4, 4 us); then by its down CTAs (S = 2: 43.1 vs 43.7 us);
+    then by its router CTAs anyway."""
+    n_router = N_ROUTER * tokens
+    n_split = N_SPLIT * tokens
+    most = max(ug_tasks_of(c, tokens) for c in range(BLOCKS))
+
+    def cost(base):
+        ctas = [(base + j) % BLOCKS for j in range(SH_TASKS)]
+        busiest = max(ug_tasks_of(c, tokens) for c in ctas)
+        in_router = sum((c - router) % BLOCKS < n_router for c in ctas)
+        in_down = sum(n_split <= c < n_split + N_DN for c in ctas)
+        return busiest, in_router if busiest == most else 0, in_down, in_router
+
+    return min(range(BLOCKS), key=cost)
+
+
+def shared_after_routing(tokens: int) -> bool:
+    """Every CTA carries the same number of routed up / gate tasks (S = 4): a
+    shared-expert CTA routes first and runs its shared half task while its first
+    routed task's weights are in flight, instead of holding its routing back
+    (S = 4 54.51 -> 54.09 us; S = 2, where the shared window sits on CTAs with one
+    task fewer, +0.1)."""
+    loads = [ug_tasks_of(c, tokens) for c in range(BLOCKS)]
+    return min(loads) == max(loads)
+
+
+WIDE_FROM = 5  # token counts from here run the expert-grouped MoE
+
+
+def wide_moe(tokens: int) -> bool:
+    """The MoE stage groups its work by expert (every distinct expert's weights
+    read once, a B column per token) instead of by (token, slot): past four tokens
+    the per-token tasks re-read the experts the tokens share, and their operands
+    no longer fit in LDS."""
+    return tokens >= WIDE_FROM
+
+
+# LDS rows of a token's MXFP8 xn and of its E8M0 scales, padded by 16 B: at the
+# natural 1536 / 192-word rows the 16 tokens a MFMA B operand gathers sit in one
+# bank group (16-way conflicts; ug spent ~3 us a task in them at S = 16)
+XN8_ROW = HIDDEN // 4 + 4
+XSC_ROW = HIDDEN // 32 + 4
+# and the down stage's MXFP8 mid rows / their scales (16 lanes, 16 rows): 196 and
+# 28 words put the rows' 16 B reads, and 4 B scale reads, in distinct banks
+MID_ROW = INTER // 4 + 4
+MIDSC_ROW = INTER // 32 + 4
+# and the o stage's FP8 attention rows (512 words: the same 16-way conflict)
+O_ROW = O_K // 4 + 4
+
+
+def pool_words(tokens: int) -> int:
+    """LDS words of the activation pool ``x``. Narrow: every token's bf16 xn (the
+    largest stage operand). Wide: the largest of the split's PV partials, the up /
+    gate stage's MXFP8 xn + scales + up partials + router sigmoids, and the down
+    stage's MXFP8 mids + scales + a zero row."""
+    if not wide_moe(tokens):
+        return HIDDEN // 2 * tokens
+    rows = MOE_SLOTS * tokens + 1
+    return max(
+        WAVES * O_K,
+        tokens * (XN8_ROW + XSC_ROW + N_ROUTED) + WAVES * 64 * 4,
+        rows * (MID_ROW + MIDSC_ROW),
+    )
diff --git a/vllm/models/minimax_m3/amd/mono/peer_buffer.py b/vllm/models/minimax_m3/amd/mono/peer_buffer.py
new file mode 100644
index 000000000000..60766486a665
--- /dev/null
+++ b/vllm/models/minimax_m3/amd/mono/peer_buffer.py
@@ -0,0 +1,52 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+# SPDX-License-Identifier: MIT
+# Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved.
+"""Symmetric peer memory for the mono kernels' in-kernel all-reduces."""
+
+import torch
+import torch.distributed as dist
+from aiter.ops.flydsl.quick_allreduce_int4_ipc import UncachedIpcHeap
+
+
+class PeerBuffer:
+    """One uncached buffer per rank, every rank holding every peer's address.
+
+    K4 pushes a layer's attention partials into every peer's ``attn`` region and
+    its FFN partials into the ``ffn`` region, then polls its own. Regions are
+    reused by every layer and step (the mailbox epoch tells a stale pair from a
+    fresh one) and no rank can overwrite a peer's unread pairs: to push layer L's
+    FFN partials a rank must have passed layer L's attention reduce, which needs
+    every peer's layer-L attention push, and a peer only pushes that after it
+    finished reading layer L-1's FFN region. So ranks are at most half a layer
+    apart and the two regions double-buffer each other.
+    """
+
+    def __init__(self, nbytes: int, group, rank: int, npes: int, device: torch.device):
+        self.local = UncachedIpcHeap.alloc_uncached(nbytes)
+        self._opened: list[int] = []
+        addresses = [self.local]
+        if npes > 1:
+            handles = UncachedIpcHeap.gather_object_list_via_broadcast(
+                group, UncachedIpcHeap.get_mem_handle_bytes(self.local)
+            )
+            addresses = []
+            for peer, handle in enumerate(handles):
+                if peer == rank:
+                    addresses.append(self.local)
+                else:
+                    base = UncachedIpcHeap.open_mem_handle(handle)
+                    self._opened.append(base)
+                    addresses.append(base)
+            dist.barrier(group=group)
+        self.addresses = torch.tensor(addresses, dtype=torch.int64, device=device)
+
+    def close(self) -> None:
+        """Close the peer mappings and free the local buffer; idempotent."""
+        opened, self._opened = self._opened, []
+        for base in opened:
+            UncachedIpcHeap.close_mem_handle(base)
+        if self.local:
+            UncachedIpcHeap.free_device_mem(self.local)
+            self.local = 0
diff --git a/vllm/models/minimax_m3/amd/mono/runner.py b/vllm/models/minimax_m3/amd/mono/runner.py
new file mode 100644
index 000000000000..89b9b84afd8b
--- /dev/null
+++ b/vllm/models/minimax_m3/amd/mono/runner.py
@@ -0,0 +1,334 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""One decode step of 1..MAX_TOKENS tokens on the mono kernels.
+
+Every sparse MoE layer is one launch of the fused layer kernel; it takes
+``(ar, res)`` -- the pair vLLM's input norm would receive (reduced MoE output,
+residual) -- and leaves ``(ar, h_mid)`` for the next layer, alternating between
+two buffers by layer parity. The embedding, the leading dense-FFN layers and the
+final norm run the original modules. The K/V cache is read in vLLM's interleaved
+page-16 layout under the fixed unit fp8 scale; the step's tables come from the
+forward context.
+"""
+
+from typing import Any, NamedTuple, cast
+
+import torch
+
+from vllm.distributed import (
+    get_tensor_model_parallel_rank,
+    get_tensor_model_parallel_world_size,
+    get_tp_group,
+)
+from vllm.forward_context import get_forward_context
+from vllm.model_executor.layers.fused_moe.experts import rocm_aiter_moe
+from vllm.models.minimax_m3.amd.mono import config as mono_config
+from vllm.models.minimax_m3.amd.mono.config import (
+    BLOCKS,
+    HEAD_DIM,
+    HIDDEN,
+    LOCAL_Q_HEADS,
+    MAX_TOKENS,
+    N_ROUTED,
+    PAGE16,
+    TOP_K,
+    TOPK_BLOCKS,
+    TP,
+    IndexHeads,
+    MonoUnsupported,
+)
+from vllm.models.minimax_m3.amd.mono.kernels.post_attn import build_post_attn_kernel
+from vllm.models.minimax_m3.amd.mono.kernels.pre_attn import K1_ARGS
+from vllm.models.minimax_m3.amd.mono.kernels.pre_attn import (
+    SCRATCH_BYTES as K1_SCRATCH_BYTES,
+)
+from vllm.models.minimax_m3.amd.mono.layout import SCRATCH_BYTES as K4_SCRATCH_BYTES
+from vllm.models.minimax_m3.amd.mono.layout import sym_layout
+from vllm.models.minimax_m3.amd.mono.peer_buffer import PeerBuffer
+from vllm.models.minimax_m3.amd.mono.weights import SparseMoeLayer, build_sparse_layer
+from vllm.platforms import current_platform
+
+
+def _ptr(t: torch.Tensor) -> int:
+    return t.data_ptr()
+
+
+class StepTables(NamedTuple):
+    """The per-step inputs every layer launch of a decode step shares."""
+
+    block_table: torch.Tensor  # int32 [rows, width], one row per token
+    seq_lens: torch.Tensor  # int32 [rows]
+    slot_mapping: torch.Tensor  # int64, page-128 slots (K/V and index cache)
+    q_len: int  # query tokens per request (speculative verify rows)
+
+
+def token_rows(
+    block_table: torch.Tensor, seq_lens: torch.Tensor, q: int, n: int
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """One (block table row, seq_len) per token of the step: a request of q query
+    tokens becomes q rows, token j seeing ``seq_len - (q - 1 - j)`` keys."""
+    if q == 1:
+        return block_table, seq_lens
+    reqs = n // q
+    back = torch.arange(q - 1, -1, -1, dtype=seq_lens.dtype, device=seq_lens.device)
+    return (
+        block_table[:reqs].repeat_interleave(q, dim=0),
+        (seq_lens[:reqs, None] - back).reshape(-1),
+    )
+
+
+def dense_prefix(layers) -> int:
+    """Leading layers without an MoE FFN; every later layer must be sparse MoE."""
+    n = 0
+    while n < len(layers) and not layers[n].is_moe_layer:
+        n += 1
+    if n == len(layers) or not all(
+        lyr.is_moe_layer and hasattr(lyr.self_attn, "indexer") for lyr in layers[n:]
+    ):
+        raise MonoUnsupported("layer mix: dense layers must precede the MoE layers")
+    return n
+
+
+def _fused_shared_weight() -> float:
+    """The fused shared expert's slot weight, from aiter's routing metadata."""
+    meta = rocm_aiter_moe.aiter_topK_meta_data
+    if meta is None:
+        raise MonoUnsupported("fused shared expert routing metadata")
+    total_w, total_ids = meta
+    if int(total_ids[0, TOP_K].item()) != N_ROUTED:
+        raise MonoUnsupported("fused shared expert id")
+    return float(total_w[0, TOP_K].item())
+
+
+class MonoDecodeRunner:
+    """Owns the per-layer weight views, scratch, peer buffers and compiled kernels."""
+
+    def __init__(self, model, config) -> None:
+        npes = get_tensor_model_parallel_world_size()
+        if npes != TP:
+            raise MonoUnsupported(f"TP {npes}")
+        self.model = model
+        self.rank = get_tensor_model_parallel_rank()
+        layers = list(model.layers)
+        # Enable indexer context parallelism: each rank computes all TP index
+        # heads, allowing local block scoring without cross-rank exchange.
+        heads = IndexHeads(count=TP, own=self.rank)
+        tp_group = get_tp_group().device_group
+        self.n_dense = n_dense = dense_prefix(layers)
+        self.sparse = [
+            build_sparse_layer(layers[i], i, heads, rank=self.rank, tp_group=tp_group)
+            for i in range(n_dense, len(layers))
+        ]
+        attn0 = layers[n_dense].self_attn
+        moe0 = layers[n_dense].block_sparse_moe
+        self.shared_fp8 = self.sparse[0].shared_fp8
+        if any(lw.shared_fp8 != self.shared_fp8 for lw in self.sparse):
+            raise MonoUnsupported("shared expert fused in some layers only")
+        # a separate MXFP8 shared expert adds unscaled: routed_scaling_factor goes
+        # on the routed weights alone (route_scale), as the model's runner does
+        shared_weight = 1.0 if self.shared_fp8 else _fused_shared_weight()
+        self.sm_scale = attn0.scaling
+        if abs(self.sm_scale - HEAD_DIM**-0.5) > 1e-12:
+            raise MonoUnsupported(f"softmax scale {self.sm_scale}")
+        sel = attn0.indexer.impl
+        if sel.topk_blocks != TOPK_BLOCKS:
+            raise MonoUnsupported(f"indexer top-{sel.topk_blocks} blocks")
+        self.layer_name = attn0.layer_name
+        self.index_prefix = attn0.indexer.index_cache.prefix
+        # the caches this build viewed: vLLM re-binds them after its cudagraph
+        # memory-profiling capture, which must rebuild the runner
+        self._cache_ptrs = self.cache_ptrs(model, n_dense)
+        # the K/V page views (V offset to its side of the block) and how many page
+        # ids a block spans; the kernels are compiled for that geometry
+        self.kv_views = [lyr.self_attn.get_aiter_sparse_pa_kv_cache() for lyr in layers]
+        stride = attn0._aiter_sparse_pa_block_page_stride
+        mono_config.configure(kv_page_stride=stride, static_kv_scale=True)
+        self._build(
+            npes, heads, float(config.rms_norm_eps), float(moe0.routed_scaling_factor),
+            shared_weight,
+            float(getattr(moe0.experts, "routed_experts", moe0.experts).swiglu_limit),
+            sel.init_blocks, sel.local_blocks,
+        )  # fmt: skip
+        dev = self.h.device
+        n16 = self.kv_views[0][0].shape[0]
+        # the host's unit fp8 scale for every page-16 slot (K1 writes 1.0 back)
+        self.kv_scale = torch.ones(n16, 1, PAGE16, dtype=torch.float32, device=dev)
+        self.res0 = torch.zeros(MAX_TOKENS, HIDDEN, dtype=torch.bfloat16, device=dev)
+
+    @staticmethod
+    def cache_ptrs(model, n_dense: int) -> tuple[int, int]:
+        """Addresses of the first MoE layer's KV cache and index cache."""
+        attn = model.layers[n_dense].self_attn
+        return attn.kv_cache.data_ptr(), attn.indexer.index_cache.kv_cache.data_ptr()
+
+    def stale(self) -> bool:
+        """Whether vLLM re-bound the caches this runner viewed (profiling capture)."""
+        return self.cache_ptrs(self.model, self.n_dense) != self._cache_ptrs
+
+    def step_tables(self, n: int) -> StepTables:
+        """The current step's tables for ``n`` tokens, from the forward context."""
+        fwd = get_forward_context()
+        decode = cast(Any, fwd.attn_metadata[self.layer_name]).decode
+        q = decode.decode_query_len
+        block_table, seq_lens = token_rows(decode.block_table, decode.seq_lens, q, n)
+        return StepTables(block_table, seq_lens, fwd.slot_mapping[self.layer_name], q)
+
+    def start(
+        self,
+        input_ids: torch.Tensor | None,
+        positions: torch.Tensor,
+        inputs_embeds: torch.Tensor | None,
+        aux_hidden_states: list[torch.Tensor] | None = None,
+    ) -> torch.Tensor:
+        """Embedding and the dense layers on the original modules; leaves the first
+        MoE layer's (ar, res) in ``ars[0][:n]`` and the returned residual (zeros
+        when layer 0 is already an MoE layer: (embeddings, 0))."""
+        n = positions.numel()
+        hidden = (
+            inputs_embeds
+            if inputs_embeds is not None
+            else self.model.embed_input_ids(input_ids)
+        )
+        residual = None
+        if aux_hidden_states is not None:
+            self.model._maybe_add_hidden_state(aux_hidden_states, 0, hidden, residual)
+        for i in range(self.n_dense):
+            hidden, residual = self.model.layers[i](positions, hidden, residual)
+            if aux_hidden_states is not None:
+                self.model._maybe_add_hidden_state(
+                    aux_hidden_states, i + 1, hidden, residual
+                )
+        self.ars[0][:n].copy_(hidden.view(n, HIDDEN))
+        return self.res0[:n] if residual is None else residual.view(n, HIDDEN)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None,
+        positions: torch.Tensor,
+        inputs_embeds: torch.Tensor | None = None,
+    ) -> torch.Tensor | tuple[torch.Tensor, list[torch.Tensor]]:
+        """The final hidden states of one decode step (what the model returns)."""
+        aux_hidden_states: list[torch.Tensor] | None = (
+            [] if self.model.aux_hidden_state_layers else None
+        )
+        res = self.start(input_ids, positions, inputs_embeds, aux_hidden_states)
+        tables = self.step_tables(res.shape[0])
+        for i, lw in enumerate(self.sparse):
+            res = self.run_sparse_layer(i, lw, tables, positions, res)
+            if aux_hidden_states is not None:
+                self.model._maybe_add_hidden_state(
+                    aux_hidden_states,
+                    lw.layer_id + 1,
+                    self.ars[(i + 1) % 2][: res.shape[0]],
+                    res,
+                )
+        hidden = self.run_final_norm(res)
+        if aux_hidden_states is not None:
+            return hidden, aux_hidden_states
+        return hidden
+
+    def run_sparse_layer(
+        self,
+        i: int,
+        lw: SparseMoeLayer,
+        tables: StepTables,
+        positions: torch.Tensor,
+        res: torch.Tensor,
+    ) -> torch.Tensor:
+        """Sparse layer `i`: (`ars[i % 2][:n]`, res) -> (`ars[(i + 1) % 2][:n]`, the
+        returned residual)."""
+        n = res.shape[0]
+        block_table, seq_lens, slots, q_len = tables
+        k16, v16 = self.kv_views[lw.layer_id]
+        self.k4[n](
+            _ptr(self.h), _ptr(self.q), _ptr(block_table), _ptr(seq_lens),
+            _ptr(k16), _ptr(v16), _ptr(self.kv_scale), _ptr(self.kv_scale),
+            _ptr(lw.w_o), _ptr(lw.s_o), _ptr(lw.g_post),
+            _ptr(lw.gate), _ptr(lw.bias), _ptr(lw.w13), _ptr(lw.s13),
+            _ptr(lw.w2), _ptr(lw.s2),
+            _ptr(lw.sh_args) if lw.shared_fp8 else 0,
+            _ptr(self.h_mids[(i + 1) % 2]), _ptr(self.ars[(i + 1) % 2]),
+            _ptr(self.scratch4), self.peers.local, _ptr(self.peers.addresses),
+            _ptr(self.step), self.rank, lw.layer_id,
+            block_table.shape[1], q_len,
+            0,  # timeline stamps (kernel debug build only)
+            _ptr(self.k1_args[i]), _ptr(positions), _ptr(slots), _ptr(res),
+            stream=torch.cuda.current_stream(),  # flydsl rejects torch.Stream
+        )  # fmt: skip
+        return self.h_mids[(i + 1) % 2][:n]
+
+    def _build(
+        self, npes, heads, eps, route_scale, shared_weight, swiglu_limit,
+        init_blocks, local_blocks,
+    ) -> None:  # fmt: skip
+        """Compile the layer kernels and allocate scratch, peer buffers and the
+        per-step activation buffers; ``self.sparse`` and ``self.sm_scale`` are set."""
+        device_index = torch.accelerator.current_device_index()
+        dev = torch.device("cuda", device_index)
+        cus = current_platform.get_device_properties(device_index).multi_processor_count
+        if cus != BLOCKS:
+            raise MonoUnsupported(
+                f"{cus} CUs: the kernels need all {BLOCKS} CTAs co-resident"
+            )
+        # one layer kernel per decode batch size (every graph captures its own)
+        self.k4 = {
+            s: build_post_attn_kernel(
+                npes,
+                self.sm_scale,
+                eps,
+                route_scale,
+                shared_weight,
+                swiglu_limit,
+                init_blocks,
+                local_blocks,
+                s,
+                fuse_k1=True,
+                heads=heads,
+                shared_fp8=self.shared_fp8,
+            )
+            for s in range(1, MAX_TOKENS + 1)
+        }
+        self.scratch1 = torch.zeros(K1_SCRATCH_BYTES, dtype=torch.uint8, device=dev)
+        self.scratch4 = torch.zeros(K4_SCRATCH_BYTES, dtype=torch.uint8, device=dev)
+        self.peers = PeerBuffer(
+            sym_layout(npes)["_bytes"], get_tp_group().cpu_group, self.rank, npes, dev
+        )
+        self.step = torch.zeros(1, dtype=torch.int32, device=dev)
+        bf16 = torch.bfloat16
+        # row k = token k of the step; sparse layer i reads ars[i % 2] (the
+        # previous layer's output) and writes ars[(i + 1) % 2], likewise h_mids
+        self.ars = [
+            torch.empty(MAX_TOKENS, HIDDEN, dtype=bf16, device=dev) for _ in range(2)
+        ]
+        self.h_mids = [
+            torch.empty(MAX_TOKENS, HIDDEN, dtype=bf16, device=dev) for _ in range(2)
+        ]
+        self.h = torch.empty(MAX_TOKENS, HIDDEN, dtype=bf16, device=dev)
+        self.q = torch.empty(
+            MAX_TOKENS, LOCAL_Q_HEADS * HEAD_DIM, dtype=bf16, device=dev
+        )
+        self.iq = torch.empty(MAX_TOKENS, 1, HEAD_DIM, dtype=bf16, device=dev)
+        # each layer's K1 pointers that do not change per step (K1_ARGS order)
+        self.k1_args = []
+        for i, lw in enumerate(self.sparse):
+            ptrs = {
+                "ar": _ptr(self.ars[i % 2]), "g_in": _ptr(lw.g_in),
+                "w_qkv": _ptr(lw.w_qkv),
+                "s_qkv": _ptr(lw.s_qkv), "g_q": _ptr(lw.g_q), "g_k": _ptr(lw.g_k),
+                "g_iq": _ptr(lw.g_iq), "g_ik": _ptr(lw.g_ik),
+                "cos_sin": _ptr(lw.cos_sin),
+                "index_cache": _ptr(lw.index_cache), "iq_out": _ptr(self.iq),
+                "scratch": _ptr(self.scratch1),
+            }  # fmt: skip
+            self.k1_args.append(
+                torch.tensor([ptrs[a] for a in K1_ARGS], dtype=torch.int64, device=dev)
+            )
+
+    def close(self) -> None:
+        self.peers.close()
+
+    def run_final_norm(self, res: torch.Tensor) -> torch.Tensor:
+        ar = self.ars[len(self.sparse) % 2]
+        hidden, _ = self.model.norm(ar[: res.shape[0]], res)
+        self.step.add_(1)
+        return hidden
diff --git a/vllm/models/minimax_m3/amd/mono/weights.py b/vllm/models/minimax_m3/amd/mono/weights.py
new file mode 100644
index 000000000000..ad7446e601cf
--- /dev/null
+++ b/vllm/models/minimax_m3/amd/mono/weights.py
@@ -0,0 +1,377 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Zero-copy views (and two converted copies) of one vLLM ``MiniMaxM3DecoderLayer``
+for the mono kernels.
+
+The kernels read tensors through raw pointers in fixed layouts: a differently
+laid-out weight would not fault, it would produce plausible garbage. So every
+layout assumption is asserted once and a model that does not match is refused
+(``MonoUnsupported``). The experts already have the kernels' layout (MXFP4, aiter
+``shuffle_weight(16, 16)`` + ``e8m0_shuffle`` scales, fused shared expert in slot
+128). The attention linears are MXFP8 in vLLM and the router gate is fp32, so
+those are converted once when the runner is built: MXFP8 -> bf16 -> per-channel
+fp8, preshuffled (16, 16); gate -> bf16.
+
+Without shared-expert fusion the routed table holds the 128 routed experts and
+the shared expert stays the separate MXFP8 MLP; its weights are copied once into
+the f8f6f4 MFMA's fp8 A-operand order (``mxfp8_a_tiles``, the same bytes) with
+``e8m0_shuffle``-ordered scales, so the kernels run it in MXFP8.
+"""
+
+from dataclasses import dataclass
+
+import torch
+import torch.distributed as dist
+from aiter import dtypes
+from aiter.ops.quant import pertoken_quant
+from aiter.ops.shuffle import shuffle_weight
+
+from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
+    dequant_mxfp8_to_bf16,
+)
+from vllm.models.minimax_m3.amd.mono.config import (
+    HEAD_DIM,
+    HIDDEN,
+    INTER,
+    LOCAL_Q_HEADS,
+    N_ROUTED,
+    O_K,
+    ONE_INDEX_HEAD,
+    ROTARY_DIM,
+    TOP_K,
+    TOPK_BLOCKS,
+    TP,
+    IndexHeads,
+    MonoUnsupported,
+)
+from vllm.models.minimax_m3.amd.mono.kernels.post_attn import SH_ARGS
+
+
+def _need(ok: bool, what: str) -> None:
+    if not ok:
+        raise MonoUnsupported(what)
+
+
+def _bf16_vec(t: torch.Tensor, n: int, name: str) -> torch.Tensor:
+    _need(
+        t.dtype == torch.bfloat16 and t.numel() == n and t.is_contiguous(),
+        f"{name}: {t.dtype} {tuple(t.shape)}",
+    )
+    return t.data
+
+
+@dataclass(frozen=True)
+class SparseMoeLayer:
+    """Everything the layer kernel reads for one sparse-attention MoE layer."""
+
+    layer_id: int
+    attn: object  # MiniMaxM3SparseAttention: layer name, KV cache views, scales
+    index_cache: torch.Tensor  # [blocks, block, head_dim] fp8, shared block ids
+    g_in: torch.Tensor
+    w_qkv: torch.Tensor
+    s_qkv: torch.Tensor
+    g_q: torch.Tensor
+    g_k: torch.Tensor
+    g_iq: torch.Tensor
+    g_ik: torch.Tensor
+    cos_sin: torch.Tensor
+    w_o: torch.Tensor
+    s_o: torch.Tensor
+    g_post: torch.Tensor
+    gate: torch.Tensor
+    bias: torch.Tensor
+    w13: torch.Tensor
+    s13: torch.Tensor
+    w2: torch.Tensor
+    s2: torch.Tensor
+    # the separate MXFP8 shared expert (``mxfp8_a_tiles``) and the int64 table of
+    # their addresses the kernels take (``SH_ARGS`` order); None when it is fused
+    sh_w13: torch.Tensor | None = None
+    sh_s13: torch.Tensor | None = None
+    sh_w2: torch.Tensor | None = None
+    sh_s2: torch.Tensor | None = None
+    sh_args: torch.Tensor | None = None
+
+    @property
+    def shared_fp8(self) -> bool:
+        return self.sh_args is not None
+
+    def shared_bytes(self) -> int:
+        """Bytes of the kernels' copy of the separate shared expert."""
+        if not self.shared_fp8:
+            return 0
+        nbytes = 0
+        for t in (self.sh_w13, self.sh_s13, self.sh_w2, self.sh_s2):
+            assert t is not None
+            nbytes += t.numel()
+        return nbytes
+
+
+def _ptpc_fp8_from_mxfp8(
+    linear, rows: int, cols: int, name: str
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """MXFP8 linear (fp8 [rows, cols], e8m0 [rows, cols/32]) -> per-channel fp8
+    preshuffled (16, 16) + fp32 scale [rows], as the kernels' PTPC GEMV reads."""
+    w = getattr(linear, "weight", None)
+    s = getattr(linear, "weight_scale", None)
+    _need(
+        w is not None and w.dtype == dtypes.fp8 and tuple(w.shape) == (rows, cols),
+        f"{name}: weight {None if w is None else (w.dtype, tuple(w.shape))}",
+    )
+    _need(
+        s is not None and tuple(s.shape) == (rows, cols // 32),
+        f"{name}: scale {None if s is None else tuple(s.shape)}",
+    )
+    assert w is not None and s is not None
+    bf = dequant_mxfp8_to_bf16(w.data, s.data.view(torch.uint8))
+    q, scale = pertoken_quant(bf, quant_dtype=dtypes.fp8)
+    return (
+        shuffle_weight(q, layout=(16, 16)).contiguous(),
+        scale.view(rows).to(torch.float32).contiguous(),
+    )
+
+
+def mxfp8_a_tiles(
+    linear, rows: int, cols: int, name: str
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """MXFP8 linear (fp8 [rows, cols], e8m0 [rows, cols/32]) -> its bytes in the
+    f8f6f4 MFMA's fp8 A order and its scales in ``e8m0_shuffle`` order.
+
+    A 16-row x 128-k tile is 2 KiB; lane g * 16 + r of a tile holds row r's k
+    16 g .. + 16 then 64 + 16 g .. + 16 (the B operand's order), 32 bytes."""
+    w = getattr(linear, "weight", None)
+    s = getattr(linear, "weight_scale", None)
+    _need(
+        w is not None and w.dtype == dtypes.fp8 and tuple(w.shape) == (rows, cols),
+        f"{name}: weight {None if w is None else (w.dtype, tuple(w.shape))}",
+    )
+    _need(
+        s is not None and tuple(s.shape) == (rows, cols // 32),
+        f"{name}: scale {None if s is None else tuple(s.shape)}",
+    )
+    assert w is not None and s is not None
+    # the row-major layout the native MXFP8 linear kernel keeps; a kernel that
+    # preshuffles or pads would make the tiles silently wrong
+    _need(
+        not getattr(w, "is_shuffled", False) and not getattr(s, "is_shuffled", False),
+        f"{name}: weight preshuffled",
+    )
+    _need(rows % 32 == 0 and cols % 256 == 0, f"{name}: shape {rows} x {cols}")
+    tiles = (
+        w.data.contiguous()
+        .view(torch.uint8)
+        .view(rows // 16, 16, cols // 128, 2, 4, 16)
+        .permute(0, 2, 4, 1, 3, 5)
+        .contiguous()
+        .view(-1)
+    )
+    scales = (
+        s.data.contiguous()
+        .view(torch.uint8)
+        .view(rows // 32, 2, 16, cols // 256, 2, 4)
+        .permute(0, 3, 5, 2, 4, 1)
+        .contiguous()
+        .view(-1)
+    )
+    return tiles, scales
+
+
+def _ptpc_fp8_from_bf16(
+    bf: torch.Tensor, rows: int
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """bf16 [rows, cols] -> per-channel fp8 preshuffled (16, 16) + fp32 scale [rows]."""
+    q, scale = pertoken_quant(bf, quant_dtype=dtypes.fp8)
+    return (
+        shuffle_weight(q, layout=(16, 16)).contiguous(),
+        scale.view(rows).to(torch.float32).contiguous(),
+    )
+
+
+def _build_cp_qkv(
+    linear, rank: int, tp_group, layer_id: int
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """The wider fused projection of indexer context parallelism: every rank's
+    index q head, all-gathered once at build time (``IndexHeads(TP, rank)``).
+
+    One rank's rows (2560): q(2048) | k(128) | v(128) | index_q(128) | index_k(128)
+    CP rows (2944):         q(2048) | k(128) | v(128) | iq_rank0..3 | index_k(128)
+
+    The gather's list order is the TP group's rank order, which ``IndexHeads.own``
+    indexes; the rank's own slot is checked against its local rows."""
+    one_rows = (LOCAL_Q_HEADS + 3 + 1) * HEAD_DIM  # 2560
+    w = getattr(linear, "weight", None)
+    s = getattr(linear, "weight_scale", None)
+    _need(
+        w is not None
+        and w.dtype == dtypes.fp8
+        and tuple(w.shape) == (one_rows, HIDDEN),
+        f"layer {layer_id} qkv_proj CP: weight "
+        f"{None if w is None else (w.dtype, tuple(w.shape))}",
+    )
+    _need(
+        s is not None and tuple(s.shape) == (one_rows, HIDDEN // 32),
+        f"layer {layer_id} qkv_proj CP: scale {None if s is None else tuple(s.shape)}",
+    )
+    assert w is not None and s is not None
+    _need(
+        dist.get_rank(tp_group) == rank and dist.get_world_size(tp_group) == TP,
+        f"layer {layer_id} qkv_proj CP: TP group rank {dist.get_rank(tp_group)}"
+        f" of {dist.get_world_size(tp_group)}, not {rank} of {TP}",
+    )
+    bf = dequant_mxfp8_to_bf16(w.data, s.data.view(torch.uint8))
+
+    iq_start = (LOCAL_Q_HEADS + 2) * HEAD_DIM  # 2304
+    iq_end = iq_start + HEAD_DIM  # 2432
+    my_iq = bf[iq_start:iq_end].contiguous()
+
+    all_iq = [torch.empty_like(my_iq) for _ in range(TP)]
+    dist.all_gather(all_iq, my_iq, group=tp_group)
+    _need(
+        torch.equal(all_iq[rank], my_iq),
+        f"layer {layer_id} qkv_proj CP: gathered index q heads not in rank order",
+    )
+
+    cp_bf = torch.cat([bf[:iq_start], *all_iq, bf[iq_end:]], dim=0)
+    cp_rows = (LOCAL_Q_HEADS + 3 + TP) * HEAD_DIM  # 2944
+    assert cp_bf.shape[0] == cp_rows, f"CP QKV rows {cp_bf.shape[0]} != {cp_rows}"
+    return _ptpc_fp8_from_bf16(cp_bf, cp_rows)
+
+
+def _cos_sin(attn, layer_id: int) -> torch.Tensor:
+    t = attn.rotary_emb.cos_sin_cache
+    _need(
+        t.dim() == 2 and t.shape[1] == ROTARY_DIM,
+        f"layer {layer_id}: cos/sin cache {tuple(t.shape)}",
+    )
+    if t.dtype != torch.bfloat16:
+        t = t.to(torch.bfloat16)
+    return t.contiguous()
+
+
+def build_sparse_layer(
+    layer,
+    layer_id: int,
+    heads: IndexHeads = ONE_INDEX_HEAD,
+    rank: int = 0,
+    tp_group=None,
+) -> SparseMoeLayer:
+    attn = layer.self_attn
+    _need(
+        hasattr(attn, "indexer") and layer.is_moe_layer,
+        f"layer {layer_id}: not sparse MoE",
+    )
+    # the model has one index head per rank; under indexer context parallelism
+    # the fused projection is built wider (every rank's), see _build_cp_qkv
+    _need(
+        (attn.num_heads, attn.num_kv_heads, attn.num_idx_heads, attn.head_dim)
+        == (LOCAL_Q_HEADS, 1, 1, HEAD_DIM),
+        f"layer {layer_id}: heads "
+        f"{attn.num_heads}/{attn.num_kv_heads}/{attn.num_idx_heads}",
+    )
+    _need(attn.rotary_emb.rotary_dim == ROTARY_DIM, f"layer {layer_id}: rotary_dim")
+    sel = attn.indexer.impl  # the selection config lives on the indexer impl
+    _need(
+        (sel.topk_blocks, sel.init_blocks, sel.local_blocks, attn.skip_index_topk)
+        == (TOPK_BLOCKS, 0, 1, False),
+        f"layer {layer_id}: sparse selection config",
+    )
+    _need(
+        attn.use_aiter_sparse_pa and str(attn.kv_cache_dtype).startswith("fp8"),
+        f"layer {layer_id}: needs the AITER sparse PA path with an fp8 KV cache",
+    )
+    index_cache = attn.indexer.index_cache.kv_cache
+    _need(
+        index_cache.numel() > 0
+        and index_cache.dtype == dtypes.fp8
+        and index_cache.dim() == 3
+        and index_cache.shape[2] == HEAD_DIM
+        and index_cache.is_contiguous(),
+        f"layer {layer_id}: index cache {index_cache.dtype} {tuple(index_cache.shape)}",
+    )
+
+    if heads.count > 1:
+        _need(tp_group is not None, f"layer {layer_id}: CP requires tp_group")
+        w_qkv, s_qkv = _build_cp_qkv(attn.qkv_proj, rank, tp_group, layer_id)
+    else:
+        w_qkv, s_qkv = _ptpc_fp8_from_mxfp8(
+            attn.qkv_proj, heads.rows, HIDDEN, f"layer {layer_id} qkv_proj"
+        )
+
+    w_o, s_o = _ptpc_fp8_from_mxfp8(
+        attn.o_proj, HIDDEN, O_K, f"layer {layer_id} o_proj"
+    )
+
+    moe = layer.block_sparse_moe
+    # FusedMoEFactory returns a MoERunner wrapping the RoutedExperts module
+    experts = getattr(moe.experts, "routed_experts", moe.experts)
+    fused = moe.is_fused_shared_expert_enabled
+    E = N_ROUTED + 1 if fused else N_ROUTED
+    w13, w2 = experts.w13_weight, experts.w2_weight
+    _need(
+        w13.shape[0] == E,
+        f"layer {layer_id}: {'fused' if fused else 'separate'} shared expert "
+        f"(experts {w13.shape[0]})",
+    )
+    shared = {}
+    if not fused:
+        mlp = moe.shared_experts
+        _need(mlp is not None, f"layer {layer_id}: no shared expert")
+        shared["sh_w13"], shared["sh_s13"] = mxfp8_a_tiles(
+            mlp.gate_up_proj, 2 * INTER, HIDDEN, f"layer {layer_id} shared gate_up"
+        )
+        shared["sh_w2"], shared["sh_s2"] = mxfp8_a_tiles(
+            mlp.down_proj, HIDDEN, INTER, f"layer {layer_id} shared down"
+        )
+        shared["sh_args"] = torch.tensor(
+            [shared[a].data_ptr() for a in SH_ARGS],
+            dtype=torch.int64,
+            device=shared["sh_w13"].device,
+        )
+    _need(experts.top_k == TOP_K, f"layer {layer_id}: top_k {experts.top_k}")
+    _need(
+        getattr(w13, "is_shuffled", False) and getattr(w2, "is_shuffled", False),
+        f"layer {layer_id}: experts not in aiter shuffled layout",
+    )
+    _need(
+        w13.numel() * w13.element_size() == E * 2 * INTER * HIDDEN // 2
+        and w2.numel() * w2.element_size() == E * HIDDEN * INTER // 2,
+        f"layer {layer_id}: expert weight bytes",
+    )
+    s13, s2 = experts.w13_weight_scale, experts.w2_weight_scale
+    _need(
+        s13.numel() == E * 2 * INTER * HIDDEN // 32
+        and s2.numel() == E * HIDDEN * INTER // 32,
+        f"layer {layer_id}: expert scale bytes",
+    )
+    gate = moe.gate.weight.data
+    _need(tuple(gate.shape) == (N_ROUTED, HIDDEN), f"layer {layer_id}: gate")
+    bias = moe.e_score_correction_bias
+    _need(
+        bias is not None and bias.dtype == torch.float32 and bias.numel() == N_ROUTED,
+        f"layer {layer_id}: bias",
+    )
+    return SparseMoeLayer(
+        layer_id=layer_id,
+        attn=attn,
+        index_cache=index_cache,
+        g_in=_bf16_vec(layer.input_layernorm.weight, HIDDEN, "input_layernorm"),
+        w_qkv=w_qkv,
+        s_qkv=s_qkv,
+        g_q=_bf16_vec(attn.q_norm.weight, HEAD_DIM, "q_norm"),
+        g_k=_bf16_vec(attn.k_norm.weight, HEAD_DIM, "k_norm"),
+        g_iq=_bf16_vec(attn.index_q_norm.weight, HEAD_DIM, "index_q_norm"),
+        g_ik=_bf16_vec(attn.index_k_norm.weight, HEAD_DIM, "index_k_norm"),
+        cos_sin=_cos_sin(attn, layer_id),
+        w_o=w_o,
+        s_o=s_o,
+        g_post=_bf16_vec(
+            layer.post_attention_layernorm.weight, HIDDEN, "post_attention_layernorm"
+        ),
+        gate=gate.to(torch.bfloat16).contiguous(),
+        bias=bias.data,
+        w13=w13.data,
+        s13=s13.data,
+        w2=w2.data,
+        s2=s2.data,
+        **shared,
+    )
diff --git a/vllm/models/minimax_m3/amd/mtp.py b/vllm/models/minimax_m3/amd/mtp.py
index 4c63b3ff01f7..f3a8fc0b6a59 100644
--- a/vllm/models/minimax_m3/amd/mtp.py
+++ b/vllm/models/minimax_m3/amd/mtp.py
@@ -265,10 +265,17 @@ def forward(
 
 
 class MiniMaxM3MTP(nn.Module):
+    # Shares the target's input embedding unless loading finds its own (only
+    # tracked under V2 PP, where the target's may not be available).
+    has_own_embed_tokens = False
+
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()
 
-        if vllm_config.use_v2_model_runner and get_pp_group().world_size > 1:
+        self._tracks_own_embed = (
+            vllm_config.use_v2_model_runner and get_pp_group().world_size > 1
+        )
+        if self._tracks_own_embed:
             # Preserve embeddings from loaders that bypass load_weights.
             self.has_own_embed_tokens = vllm_config.load_config.load_format != "dummy"
 
@@ -441,7 +448,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
                     f"Failed to load MTP layer {layer_idx} weights from checkpoint."
                 )
 
-        if hasattr(self, "has_own_embed_tokens"):
+        if self._tracks_own_embed:
             if "model.embed_tokens.weight" in loaded_params:
                 self.has_own_embed_tokens = True
             elif is_mtp_completeness_check_enabled():
diff --git a/vllm/models/minimax_m3/amd/ops/__init__.py b/vllm/models/minimax_m3/amd/ops/__init__.py
index 22d96f9de979..a9261bdfd236 100644
--- a/vllm/models/minimax_m3/amd/ops/__init__.py
+++ b/vllm/models/minimax_m3/amd/ops/__init__.py
@@ -9,6 +9,7 @@
 
 from vllm.models.minimax_m3.amd.ops.gemma_rmsnorm import (
     gemma_fused_add_rmsnorm,
+    gemma_fused_add_rmsnorm_mxfp8,
     gemma_rmsnorm,
 )
 from vllm.models.minimax_m3.amd.ops.swiglu_oai import (
@@ -19,6 +20,7 @@
 __all__ = [
     "gemma_rmsnorm",
     "gemma_fused_add_rmsnorm",
+    "gemma_fused_add_rmsnorm_mxfp8",
     "swiglu_oai_split",
     "swiglu_oai_quantize_mxfp8",
 ]
diff --git a/vllm/models/minimax_m3/amd/ops/flydsl_utils.py b/vllm/models/minimax_m3/amd/ops/flydsl_utils.py
new file mode 100644
index 000000000000..890ab3714e74
--- /dev/null
+++ b/vllm/models/minimax_m3/amd/ops/flydsl_utils.py
@@ -0,0 +1,196 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Helpers shared by the MiniMax-M3 FlyDSL kernels (gfx950)."""
+
+import functools
+
+import flydsl.expr as fx
+import torch
+from flydsl._mlir import ir
+from flydsl._mlir.dialects import arith, llvm
+from flydsl.expr import rocdl
+from flydsl.expr.typing import T
+from flydsl.expr.typing import Vector as Vec
+
+# Cache-policy bits of buffer loads and stores: sc1 is coherent across the XCDs'
+# L2s, sc0 | sc1 across the system, nt streams.
+SC0, NT, SC1 = 0x1, 0x2, 0x10
+LOG2E = 1.4426950409  # softmax runs in exp2
+# The first 16 dwords of kernel arguments arrive in SGPRs, so a kernel's first
+# loads do not wait for s_load.
+PRELOAD_HINTS = {
+    "llvm_options": {
+        "amdgpu-kernarg-preload": True,
+        "amdgpu-kernarg-preload-count": 16,
+    }
+}
+
+
+def run_compiled(exe, *args, waves_per_eu: int | None = None):
+    """Compiles and runs ``exe`` on the first call, then dispatches the cache."""
+    import flydsl.compiler as flyc
+    from flydsl.compiler.kernel_function import CompilationContext
+
+    cf = getattr(exe, "_cf", None)
+    if cf is None:
+        hints = {"waves_per_eu": waves_per_eu} if waves_per_eu else {}
+        with CompilationContext.compile_hints(hints):
+            exe._cf = flyc.compile(exe, *args)
+    else:
+        cf(*args)
+
+
+@functools.cache
+def cu_count(device: torch.device) -> int:
+    return torch.cuda.get_device_properties(device).multi_processor_count
+
+
+@functools.cache
+def unit_scale(device: torch.device) -> torch.Tensor:
+    return torch.ones(1, dtype=torch.float32, device=device)
+
+
+def global_i32_ptr(addr_i64):
+    ptr_ty = fx.PointerType.get(
+        T.i32, address_space=fx.AddressSpace.Global, alignment=4
+    )
+    return fx.inttoptr(ptr_ty, fx.Int64(addr_i64))
+
+
+def bits(x):
+    return fx.Float32(x).bitcast(fx.Int32)
+
+
+def as_f32(x):
+    return fx.Int32(x).bitcast(fx.Float32)
+
+
+def permlane_swap(op, d_a, d_b):
+    """v_permlane{16,32}_swap; fx has no wrapper."""
+    pair_ty = ir.Type.parse("!llvm.struct<(i32, i32)>")
+    res = op(pair_ty, fx.as_ir_value(d_a), fx.as_ir_value(d_b), False, False)
+    return (
+        fx.Int32(llvm.extractvalue(T.i32, res, [0])),
+        fx.Int32(llvm.extractvalue(T.i32, res, [1])),
+    )
+
+
+def maxf_nn(a, b):
+    """v_max_f32 with nnan; without it LLVM canonicalizes both inputs first."""
+    fm = ir.Attribute.parse("#arith.fastmath<nnan>")
+    return fx.Float32(
+        arith.MaxNumFOp(fx.as_ir_value(a), fx.as_ir_value(b), fastmath=fm).result
+    )
+
+
+def xlane4(x, op):
+    """``op`` over lanes L, L^16, L^32, L^48, broadcast to all four."""
+    a, b = permlane_swap(rocdl.permlane32_swap, bits(x), bits(x))
+    m = op(as_f32(a), as_f32(b))
+    a, b = permlane_swap(rocdl.permlane16_swap, bits(m), bits(m))
+    return op(as_f32(a), as_f32(b))
+
+
+def perm(hi, lo, sel):
+    """v_perm_b32: byte i of the result is byte sel_i of (hi:lo)."""
+    return fx.Int32(rocdl.perm_b32(fx.Int32(hi), fx.Int32(lo), fx.Int32(sel)))
+
+
+def exp2(x):
+    """v_exp_f32 without denormal range reduction."""
+    return fx.Float32(fx.rocdl.exp2(fx.Float32.ir_type, fx.Float32(x).ir_value()))
+
+
+def pk_bf16(a, b):
+    """Two fp32 -> one dword of two bf16 (RNE), ``a`` low."""
+    return fx.Int32(
+        fx.rocdl.cvt_pk_bf16_f32(
+            fx.as_ir_value(fx.Float32(a)), fx.as_ir_value(fx.Float32(b))
+        )
+    )
+
+
+def pk_f16(a, b):
+    """Two fp32 -> one dword of two fp16 (RNE), ``a`` low."""
+    return (
+        Vec.from_elements([fx.Float32(a), fx.Float32(b)], fx.Float32)
+        .to(fx.Float16)
+        .bitcast(fx.Int32)[0]
+    )
+
+
+def cvt_fp4x2_f16(word, sel):
+    """The two e2m1 codes of byte ``sel`` of ``word`` as two fp16 (exact)."""
+    return Vec(
+        fx.rocdl.cvt_scalef32_pk_f16_fp4(
+            Vec.make_type(2, fx.Float16),
+            fx.as_ir_value(fx.Int32(word)),
+            fx.as_ir_value(fx.Float32(1.0)),
+            sel,
+        )
+    )
+
+
+def cvt_fp8x2_f16(word, hi):
+    """E4M3 bytes 0, 1 (``hi``: 2, 3) of ``word`` as two fp16 (exact)."""
+    return Vec(
+        fx.rocdl.cvt_scalef32_pk_f16_fp8(
+            Vec.make_type(2, fx.Float16),
+            fx.as_ir_value(fx.Int32(word)),
+            fx.as_ir_value(fx.Float32(1.0)),
+            hi,
+        )
+    )
+
+
+def mfma_16x16x32(a4, b4, acc, fp16: bool):
+    """v_mfma_f32_16x16x32_{bf16,f16} on 4-dword operands."""
+    dtype = fx.Float16 if fp16 else fx.BFloat16
+    op = fx.rocdl.mfma_f32_16x16x32_f16 if fp16 else fx.rocdl.mfma_f32_16x16x32_bf16
+    a = Vec(a4).bitcast(dtype)
+    b = Vec(b4).bitcast(dtype)
+    return Vec(op(Vec.make_type(4, fx.Float32), [a, b, acc, 0, 0, 0]))
+
+
+def mfma_fp8_16x16x128(a8, b8, acc4):
+    """v_mfma_f32_16x16x128_f8f6f4 on e4m3 operands; zero scales select the
+    unscaled instruction."""
+    return Vec(
+        fx.rocdl.mfma_scale_f32_16x16x128_f8f6f4(
+            Vec.make_type(4, fx.Float32), [a8, b8, acc4, 0, 0, 0, 0, 0, 0]
+        )
+    )
+
+
+def pack8(h0, h1):
+    """Two 4 x i32 halves -> the 8 x i32 operand of the fp8 16x16x128 MFMA."""
+    a, b = Vec(h0), Vec(h1)
+    return Vec.from_elements(
+        [a[k] for k in range(4)] + [b[k] for k in range(4)], fx.Int32
+    )
+
+
+# Byte selectors of the even and odd dims in key order.
+PERM_EVEN, PERM_ODD = 0x06040200, 0x07050301
+
+
+def key_order4(d):
+    """4 dwords of e4m3 (16 dims) -> even dims, then odd dims per 8 (key order)."""
+    return Vec.from_elements(
+        [
+            perm(d[1], d[0], PERM_EVEN),
+            perm(d[1], d[0], PERM_ODD),
+            perm(d[3], d[2], PERM_EVEN),
+            perm(d[3], d[2], PERM_ODD),
+        ],
+        fx.Int32,
+    )
+
+
+def dequant8(c, t_lo, t_hi):
+    """8 e2m1 codes (low nibble the even dim) and the group's e4m3 table ->
+    the e4m3 keys of the even dims, then of the odd dims."""
+    ml = perm(t_hi, t_lo, c & fx.Int32(0x07070707))
+    mh = perm(t_hi, t_lo, (c >> 4) & fx.Int32(0x07070707))
+    sign = fx.Int32(0x80808080 - (1 << 32))
+    return ml ^ ((c << 4) & sign), mh ^ (c & sign)
diff --git a/vllm/models/minimax_m3/amd/ops/gemma_rmsnorm.py b/vllm/models/minimax_m3/amd/ops/gemma_rmsnorm.py
index cb74877f6824..208f3432d0fa 100644
--- a/vllm/models/minimax_m3/amd/ops/gemma_rmsnorm.py
+++ b/vllm/models/minimax_m3/amd/ops/gemma_rmsnorm.py
@@ -8,9 +8,13 @@
 pow, mean, rsqrt, two muls, cast) — each a separate kernel launch materializing
 fp32 intermediates. These kernels collapse that into a single pass per row.
 
-Two entry points:
+Entry points:
   * ``gemma_rmsnorm(x, w, eps)``                 -> normalized tensor
   * ``gemma_fused_add_rmsnorm(x, res, w, eps)``  -> (normalized, x + res)
+  * ``gemma_fused_add_rmsnorm_mxfp8(x, res, w, eps)``
+        -> (MXFP8 values, E8M0 scales, x + res): the fused add + norm with the
+           per-32-block MXFP8 activation quant as an epilogue, so a following
+           MXFP8 linear need not quantize its input again.
 
 Both normalize over the last dim and broadcast ``weight`` (shape [N]) over it,
 so they serve both the full-hidden norms (input/post-attn/final) and the
@@ -21,6 +25,11 @@
 
 import torch
 
+from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
+    MXFP8_BLOCK_SIZE,
+    MXFP8_SCALE_DTYPE,
+    MXFP8_VALUE_DTYPE,
+)
 from vllm.triton_utils import tl, triton
 
 
@@ -59,6 +68,7 @@ def _gemma_fused_add_rmsnorm_kernel(
     w_ptr,
     out_ptr,
     res_out_ptr,
+    scale_ptr,
     n_cols,
     stride_xrow,
     stride_xcol,
@@ -66,6 +76,9 @@ def _gemma_fused_add_rmsnorm_kernel(
     stride_rcol,
     eps,
     BLOCK_N: tl.constexpr,
+    QUANT_MXFP8: tl.constexpr,
+    FP8_MAX: tl.constexpr,
+    TINY: tl.constexpr,
 ):
     row = tl.program_id(0)
     cols = tl.arange(0, BLOCK_N)
@@ -87,11 +100,35 @@ def _gemma_fused_add_rmsnorm_kernel(
     rstd = 1.0 / tl.sqrt(var + eps)
     w = tl.load(w_ptr + cols, mask=mask, other=0.0).to(tl.float32)
     out = s * rstd * (1.0 + w)
-    tl.store(
-        out_ptr + row * n_cols + cols,
-        out.to(out_ptr.dtype.element_ty),
-        mask=mask,
-    )
+    if QUANT_MXFP8:
+        # Same math as mxfp8_e4m3_quantize on the input-dtype normed row, so
+        # the result is bit-identical to norm -> bf16 -> quant.
+        y = tl.reshape(
+            out.to(x_ptr.dtype.element_ty).to(tl.float32), (BLOCK_N // 32, 32)
+        )
+        amax = tl.maximum(tl.max(tl.abs(y), axis=1), TINY)
+        sb = tl.ceil(tl.log2(amax / FP8_MAX)) + 127.0
+        sb = tl.minimum(tl.maximum(sb, 0.0), 254.0)
+        # exp2(sb - 127) is subnormal at sb == 0 and flushes to zero on CDNA, so
+        # scale by the (exact, normal) reciprocal instead of dividing.
+        q = y * tl.exp2(127.0 - sb)[:, None]
+        tl.store(
+            out_ptr + row * n_cols + cols,
+            tl.reshape(q, (BLOCK_N,)).to(out_ptr.dtype.element_ty),
+            mask=mask,
+        )
+        blocks = tl.arange(0, BLOCK_N // 32)
+        tl.store(
+            scale_ptr + row * (n_cols // 32) + blocks,
+            sb.to(scale_ptr.dtype.element_ty),
+            mask=blocks < n_cols // 32,
+        )
+    else:
+        tl.store(
+            out_ptr + row * n_cols + cols,
+            out.to(out_ptr.dtype.element_ty),
+            mask=mask,
+        )
 
 
 def _num_warps(block_n: int) -> int:
@@ -129,12 +166,57 @@ def gemma_fused_add_rmsnorm(
     weight: torch.Tensor,
     eps: float,
 ) -> tuple[torch.Tensor, torch.Tensor]:
+    out, _, res_out = _gemma_fused_add_rmsnorm(x, residual, weight, eps, False)
+    return out.reshape(x.shape), res_out
+
+
+def gemma_fused_add_rmsnorm_mxfp8(
+    x: torch.Tensor,
+    residual: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """``gemma_fused_add_rmsnorm`` with the normed output quantized to MXFP8.
+
+    Bit-identical to ``mxfp8_e4m3_quantize(gemma_fused_add_rmsnorm(...)[0])``.
+
+    Returns:
+        ``(x_q, x_scale, residual_out)``: ``x_q`` float8_e4m3fn ``[M, N]`` and
+        ``x_scale`` uint8 E8M0 ``[M, N // 32]`` (row-major, not swizzled), with
+        the leading dims of ``x`` flattened into ``M``; ``residual_out`` has
+        the shape of ``x``.
+
+    """
+    n = x.shape[-1]
+    assert n % MXFP8_BLOCK_SIZE == 0, (
+        f"MXFP8 needs the hidden size to be a multiple of {MXFP8_BLOCK_SIZE}, got {n}"
+    )
+    x_q, x_scale, res_out = _gemma_fused_add_rmsnorm(x, residual, weight, eps, True)
+    assert x_scale is not None
+    return x_q, x_scale, res_out
+
+
+def _gemma_fused_add_rmsnorm(
+    x: torch.Tensor,
+    residual: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float,
+    quant_mxfp8: bool,
+) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor]:
     orig_shape = x.shape
     n = orig_shape[-1]
     x2 = x.reshape(-1, n)
     r2 = residual.reshape(-1, n)
     m = x2.shape[0]
-    out = torch.empty((m, n), dtype=x.dtype, device=x.device)
+    out_dtype = MXFP8_VALUE_DTYPE if quant_mxfp8 else x.dtype
+    out = torch.empty((m, n), dtype=out_dtype, device=x.device)
+    scale = (
+        torch.empty(
+            (m, n // MXFP8_BLOCK_SIZE), dtype=MXFP8_SCALE_DTYPE, device=x.device
+        )
+        if quant_mxfp8
+        else None
+    )
     res_out = torch.empty((m, n), dtype=x.dtype, device=x.device)
     block_n = triton.next_power_of_2(n)
     _gemma_fused_add_rmsnorm_kernel[(m,)](
@@ -143,6 +225,7 @@ def gemma_fused_add_rmsnorm(
         weight,
         out,
         res_out,
+        scale,
         n,
         x2.stride(0),
         x2.stride(1),
@@ -150,6 +233,9 @@ def gemma_fused_add_rmsnorm(
         r2.stride(1),
         eps,
         BLOCK_N=block_n,
+        QUANT_MXFP8=quant_mxfp8,
+        FP8_MAX=float(torch.finfo(MXFP8_VALUE_DTYPE).max),
+        TINY=float(torch.finfo(torch.float32).tiny),
         num_warps=_num_warps(block_n),
     )
-    return out.reshape(orig_shape), res_out.reshape(orig_shape)
+    return out, scale, res_out.reshape(orig_shape)
diff --git a/vllm/models/minimax_m3/amd/ops/index_cp/__init__.py b/vllm/models/minimax_m3/amd/ops/index_cp/__init__.py
new file mode 100644
index 000000000000..208f01a7cb5e
--- /dev/null
+++ b/vllm/models/minimax_m3/amd/ops/index_cp/__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/minimax_m3/amd/ops/index_cp/host.py b/vllm/models/minimax_m3/amd/ops/index_cp/host.py
new file mode 100644
index 000000000000..1094bdb5fbe5
--- /dev/null
+++ b/vllm/models/minimax_m3/amd/ops/index_cp/host.py
@@ -0,0 +1,120 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Buffers of the context-parallel decode indexer (TP4, gfx950).
+
+Rank ``me`` holds index head ``me``; ``kernel_cp_push`` scores a quarter of every
+request's blocks for all four heads and stores head h's scores into rank h's
+receive buffer. Buffers and flags are uncached memory mapped by the peers; the
+call counts stay on the device, so the handshake replays in CUDA graphs.
+"""
+
+import torch
+import torch.distributed as dist
+
+from vllm.distributed.device_communicators.custom_all_reduce import CustomAllreduce
+from vllm.models.minimax_m3.amd.ops.nvfp4_index_decode import (
+    kernel_cp_push as _kernel,
+)
+
+WORLD = _kernel.WORLD
+
+
+class _DevicePointer:
+    """A raw device allocation as a tensor (``__cuda_array_interface__``)."""
+
+    def __init__(self, ptr: int, nbytes: int):
+        self.__cuda_array_interface__ = {
+            "shape": (nbytes,),
+            "typestr": "|u1",
+            "data": (ptr, False),
+            "version": 3,
+        }
+
+
+class IndexCP:
+    def __init__(
+        self, rank: int, group, device: torch.device, max_rows: int, max_width: int
+    ):
+        self.rank, self.device = rank, device
+        self.max_rows, self.max_width = max_rows, max_width
+        score_bytes = max_rows * max_width * 4
+        iq_bytes = WORLD * max_rows * _kernel.D
+        # Each rank's buffer, mapped by every peer (uncached on ROCm).
+        self.score_ptrs, self.iq_ptrs, self.flag_ptrs = (
+            CustomAllreduce.create_shared_buffer(n, group=group)
+            for n in (score_bytes, iq_bytes, _kernel.FLAG_BYTES)
+        )
+        self.score_rcv = torch.as_tensor(
+            _DevicePointer(self.score_ptrs[rank], score_bytes), device=device
+        ).view(torch.float32)
+        flags = torch.as_tensor(
+            _DevicePointer(self.flag_ptrs[rank], _kernel.FLAG_BYTES), device=device
+        )
+        flags.zero_()
+        # Call counts (the flag values), equal on every rank.
+        self.counts = torch.zeros(_kernel.CNT_WORDS, dtype=torch.int32, device=device)
+        torch.accelerator.synchronize(device)
+        dist.barrier(group=group)
+
+    def score_push(
+        self,
+        iq: torch.Tensor,
+        index_cache: torch.Tensor,
+        block_table: torch.Tensor,
+        seq_lens: torch.Tensor,
+        init_blocks: int,
+        local_blocks: int,
+        width: int,
+        grid: int = 0,
+    ) -> torch.Tensor:
+        """Score this rank's blocks of every head from its index-q rows ``iq``.
+
+        Returns this rank's head over every block, [1, rows, width]: a view of
+        the receive buffer, valid until the next call.
+        """
+        from vllm.models.minimax_m3.amd.ops.nvfp4_index_decode.host import (
+            minimax_m3_nvfp4_index_score_push_cp,
+        )
+
+        rows = iq.shape[0]
+        assert rows <= self.max_rows and width <= self.max_width
+        minimax_m3_nvfp4_index_score_push_cp(
+            iq,
+            index_cache,
+            self.score_ptrs,
+            width,
+            rows * width * 4,
+            block_table,
+            seq_lens,
+            init_blocks,
+            local_blocks,
+            self.rank,
+            self.counts,
+            self.flag_ptrs,
+            self.iq_ptrs,
+            self.max_rows,
+            grid=grid,
+        )
+        return self.score_rcv[: rows * width].view(1, rows, width)
+
+
+# None means the single-head decode path.
+_CP: IndexCP | None = None
+
+
+def init(device: torch.device, max_width: int) -> None:
+    """Create this rank's ``IndexCP``; a collective, called once outside capture."""
+    global _CP
+    if _CP is not None:
+        return
+    from vllm.distributed import get_tp_group
+    from vllm.platforms.rocm import on_gfx950
+
+    tp = get_tp_group()
+    if tp.world_size != WORLD or not on_gfx950():
+        return
+    _CP = IndexCP(tp.rank_in_group, tp.cpu_group, device, _kernel.MAX_ROWS, max_width)
+
+
+def get() -> IndexCP | None:
+    return _CP
diff --git a/vllm/models/minimax_m3/amd/ops/index_topk_decode/__init__.py b/vllm/models/minimax_m3/amd/ops/index_topk_decode/__init__.py
new file mode 100644
index 000000000000..208f01a7cb5e
--- /dev/null
+++ b/vllm/models/minimax_m3/amd/ops/index_topk_decode/__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/minimax_m3/amd/ops/index_topk_decode/host.py b/vllm/models/minimax_m3/amd/ops/index_topk_decode/host.py
new file mode 100644
index 000000000000..f3669b733c96
--- /dev/null
+++ b/vllm/models/minimax_m3/amd/ops/index_topk_decode/host.py
@@ -0,0 +1,113 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Launch wrappers of the FlyDSL index top-k: the block ids AITER's
+``pa_sparse_block_topk`` selects, in its order."""
+
+import functools
+
+import torch
+
+from vllm.models.minimax_m3.amd.ops.flydsl_utils import (
+    run_compiled,
+)
+
+from . import prefill_kernel as _kernel
+
+MAX_SLOTS = 128  # AITER's TOPK_MAX_SLOTS: 8192 blocks
+
+
+def _pow2_ceil(n: int) -> int:
+    return 1 << (n - 1).bit_length() if n >= 1 else 1
+
+
+@functools.lru_cache(maxsize=len(_kernel.PREFILL_STRIPS))
+def _get_prefill_kernel(strips: int):
+    return _kernel.compile_topk_prefill(strips)
+
+
+@functools.lru_cache(maxsize=1)
+def _get_decode_kernel():
+    return _kernel.compile_topk_prefill(26, decode=True)
+
+
+def supports_index_topk(topk_blocks: int, block_size: int) -> bool:
+    from vllm.platforms.rocm import on_gfx950
+
+    return on_gfx950() and topk_blocks == _kernel.TOPK and block_size == _kernel.BLK
+
+
+def _check(score: torch.Tensor, topk_idx: torch.Tensor, max_seq_len: int) -> int:
+    """Validate the layouts; returns the score columns the kernel walks."""
+    heads, rows, width = score.shape
+    assert score.dtype == torch.float32 and score.stride(2) == 1
+    assert topk_idx.dtype == torch.int32 and topk_idx.stride(2) == 1
+    assert topk_idx.shape == (heads, rows, _kernel.TOPK)
+    slots = _pow2_ceil((max_seq_len + _kernel.BLK * 64 - 1) // (_kernel.BLK * 64))
+    assert slots <= MAX_SLOTS and width >= slots * 64
+    return slots * 64
+
+
+def minimax_m3_index_topk_prefill_indices(
+    score: torch.Tensor,  # [H, rows, S] fp32
+    topk_idx: torch.Tensor,  # [H, rows, 16] int32
+    *,
+    max_seq_len: int,
+    num_valid_pages: torch.Tensor,  # [>= rows] int32: row n's visible blocks
+) -> None:
+    """Top-16 blocks of ragged prefill rows."""
+    heads, rows, _ = score.shape
+    cols = _check(score, topk_idx, max_seq_len)
+    assert num_valid_pages.dtype == torch.int32
+    assert num_valid_pages.is_contiguous() and num_valid_pages.numel() >= rows
+    if rows == 0:
+        return
+    tiles = _kernel.PREFILL_STRIPS
+    strips = next((s for s in tiles if max_seq_len <= s * _kernel.BLK * 64), tiles[-1])
+    run_compiled(
+        _get_prefill_kernel(strips),
+        score.data_ptr(),
+        score.stride(0),
+        score.stride(1),
+        cols,
+        num_valid_pages.data_ptr(),
+        1,
+        topk_idx.data_ptr(),
+        topk_idx.stride(0),
+        topk_idx.stride(1),
+        rows,
+        heads,
+        torch.cuda.current_stream(),
+    )
+
+
+def minimax_m3_index_topk_decode_indices(
+    score: torch.Tensor,  # [H, num_reqs * query_len, S] fp32
+    topk_idx: torch.Tensor,  # [H, num_reqs * query_len, 16] int32
+    seq_lens: torch.Tensor,  # [>= num_reqs] int32
+    *,
+    max_seq_len: int,
+    query_len: int,
+) -> None:
+    """Top-16 blocks of decode rows, ``query_len`` rows per request."""
+    heads, rows, _ = score.shape
+    cols = _check(score, topk_idx, max_seq_len)
+    assert query_len > 0 and rows % query_len == 0
+    assert seq_lens.dtype == torch.int32 and seq_lens.is_contiguous()
+    assert seq_lens.numel() >= rows // query_len
+    if rows == 0:
+        return
+    run_compiled(
+        _get_decode_kernel(),
+        score.data_ptr(),
+        score.stride(0),
+        score.stride(1),
+        cols,
+        seq_lens.data_ptr(),
+        query_len,
+        topk_idx.data_ptr(),
+        topk_idx.stride(0),
+        topk_idx.stride(1),
+        rows,
+        heads,
+        torch.cuda.current_stream(),
+    )
diff --git a/vllm/models/minimax_m3/amd/ops/index_topk_decode/prefill_kernel.py b/vllm/models/minimax_m3/amd/ops/index_topk_decode/prefill_kernel.py
new file mode 100644
index 000000000000..b18db79e055a
--- /dev/null
+++ b/vllm/models/minimax_m3/amd/ops/index_topk_decode/prefill_kernel.py
@@ -0,0 +1,387 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Index top-k for MiniMax-M3 (FlyDSL, gfx950).
+
+One 64-lane workgroup per (row, head) returns the 16 highest-scoring block ids
+in AITER's order, walking the row in windows of STRIPS x 64 blocks.
+"""
+
+import flydsl.compiler as flyc
+import flydsl.expr as fx
+from aiter.ops.flydsl.kernels import buffer_ops
+from flydsl.expr import const_expr, range_constexpr
+from flydsl.expr.typing import T
+from flydsl.expr.typing import Vector as Vec
+
+from vllm.models.minimax_m3.amd.ops.flydsl_utils import PRELOAD_HINTS
+
+BLK = 128  # tokens per block
+TOPK = 16
+PREFILL_STRIPS = (4, 8, 16, 20, 26)
+CAP = 64  # candidates one rank pass holds (one per lane)
+IDX_BITS = 13  # block index bits (slots * 64 <= 8192)
+INT_MIN = -(1 << 31)
+KEY_NEG_INF = (0xFF800000 ^ 0x7FFFFFFF) - (1 << 32)  # signed key of -inf (NaN's)
+KEY_POS_INF = 0x7F800000
+
+
+def _s32(v):
+    v &= 0xFFFFFFFF
+    return v - (1 << 32) if v >= 1 << 31 else v
+
+
+def _key(bits):
+    """Signed key in AITER's order; NaNs land above +inf or below -inf."""
+    return bits ^ ((bits >> 31) & 0x7FFFFFFF)
+
+
+def _fold_nan(k):
+    return ((k > KEY_POS_INF) | (k < KEY_NEG_INF)).select(fx.Int32(KEY_NEG_INF), k)
+
+
+def _ballot(c):
+    return fx.Int64(fx.rocdl.ballot(T.i64, c))
+
+
+def _popc(m):
+    """Set bits of a wave mask; two 32-bit counts keep the compare scalar."""
+    lo = fx.math.ctpop(m.to(fx.Int32)).to(fx.Int32)
+    hi = fx.math.ctpop((m >> 32).to(fx.Int32)).to(fx.Int32)
+    return fx.Int32(lo) + fx.Int32(hi)
+
+
+def _below(mask, base):
+    """The set bits of ``mask`` below this lane, plus ``base`` (v_mbcnt)."""
+    lo = mask.to(fx.Int32)
+    hi = (mask >> 32).to(fx.Int32)
+    r = fx.rocdl.mbcnt_lo(T.i32, lo.ir_value(), fx.Int32(base).ir_value())
+    return fx.Int32(fx.rocdl.mbcnt_hi(T.i32, hi.ir_value(), r))
+
+
+def _imax(a, b):
+    return (a > b).select(a, b)
+
+
+def _one(c):
+    return c.select(fx.Int32(1), fx.Int32(0))
+
+
+def _comp(k, i):
+    """(key, index) as one i64 whose signed order is AITER's."""
+    return (k.to(fx.Int64) << 32) | (i.to(fx.Int64) & 0xFFFFFFFF)
+
+
+def compile_topk_prefill(strips: int, *, decode: bool = False):
+    assert strips == 26 if decode else strips in PREFILL_STRIPS
+    STRIPS = strips
+    WIN = STRIPS * 64
+    RANK_BLOCK = 8
+    # Candidates, rank-loop padding, then one discard slot per lane.
+    TRASH = CAP + RANK_BLOCK
+    CAND = TRASH + 64
+    name = f"m3_index_topk_{'decode' if decode else 'prefill'}_indices_w{STRIPS}"
+
+    @fx.struct
+    class Shared:
+        cand: fx.Array[fx.Int32, 2 * CAND, 16]
+        sel: fx.Array[fx.Int32, 4 * TOPK, 16]
+
+    @flyc.kernel(name=name, known_block_size=[64, 1, 1])
+    def kernel(
+        arg_score: fx.Int64,
+        i64_s_head: fx.Int64,
+        i64_s_row: fx.Int64,
+        i32_width: fx.Int32,
+        arg_nvp: fx.Int64,
+        i32_query_len: fx.Int32,
+        arg_topk: fx.Int64,
+        i64_t_head: fx.Int64,
+        i64_t_row: fx.Int64,
+    ):
+        smem = fx.SharedAllocator().allocate(Shared).peek()
+        n, h = fx.block_idx.x, fx.block_idx.y
+        lane = fx.thread_idx.x
+        n64 = fx.Int32(n).to(fx.Int64)
+        h64 = fx.Int32(h).to(fx.Int64)
+
+        def rsrc(base, nbytes):
+            """Buffer at an i64 base; per-row rebasing keeps offsets small."""
+            return buffer_ops.create_buffer_resource_from_addr(
+                base, num_records_bytes=fx.Int64(nbytes)
+            )
+
+        # Reads past the row's width return 0.
+        score_row = arg_score + (h64 * i64_s_head + n64 * i64_s_row) * 4
+        srs = rsrc(score_row, i32_width.to(fx.Int64) * 4)
+
+        def load_window(srs, wbase):
+            return Vec.from_elements(
+                [
+                    fx.Int32(
+                        buffer_ops.buffer_load(
+                            srs,
+                            wbase + (j * 64) + lane,
+                            vec_width=1,
+                            dtype=fx.Int32,
+                        )
+                    )
+                    for j in range_constexpr(STRIPS)
+                ],
+                fx.Int32,
+            )
+
+        # The first window is loaded before the row's length is known.
+        bits = load_window(srs, 0)
+
+        def scalar_at(base, idx64, nbytes):
+            return fx.Int32(
+                buffer_ops.buffer_load(
+                    rsrc(base + idx64 * 4, nbytes),
+                    fx.Int32(0),
+                    vec_width=1,
+                    is_scalar=True,
+                )
+            )
+
+        if const_expr(decode):
+            req = (fx.Int32(n) // i32_query_len).to(fx.Int64)
+            seq = scalar_at(arg_nvp, req, fx.Int64(4))
+            kv0 = seq - i32_query_len + fx.Int32(n) % i32_query_len + 1
+            kv = (kv0 < 0).select(fx.Int32(0), kv0)
+            nblk = (kv + BLK - 1) // BLK
+        else:
+            nblk = scalar_at(arg_nvp, n64, fx.Int64(4))
+        nlive = (nblk < i32_width).select(nblk, i32_width)
+        # Later windows read only the row's live blocks.
+        srs_live = rsrc(score_row, nlive.to(fx.Int64) * 4)
+
+        def view(ptr, n_dw, width):
+            return fx.logical_divide(
+                fx.make_view(fx.recast_iter(fx.Int32, ptr), fx.make_layout(n_dw, 1)),
+                fx.make_layout(width, 1),
+            )
+
+        pairs = view(smem.cand.ptr, 2 * CAND, 2)
+        quads = view(smem.cand.ptr, 2 * CAND, 4)
+        sels = view(smem.sel.ptr, 4 * TOPK, 4)
+        atom64 = fx.make_copy_atom(fx.UniversalCopy64b(), fx.Int32)
+        atom128 = fx.make_copy_atom(fx.UniversalCopy128b(), fx.Int32)
+
+        def lds_put(atom, v, w, vals):
+            r = fx.make_rmem_tensor(fx.make_layout(w, 1), fx.Int32)
+            r.store(Vec.from_elements(vals, fx.Int32))
+            fx.copy(atom, r, v)
+
+        def lds_get(atom, v, w):
+            r = fx.make_rmem_tensor(fx.make_layout(w, 1), fx.Int32)
+            fx.copy(atom, v, r)
+            x = Vec(r.load())
+            return [fx.Int32(x[e]) for e in range_constexpr(w)]
+
+        def put_pair(pos, i, k):
+            lds_put(atom64, fx.slice(pairs, (None, pos)), 2, [i, k])
+
+        def wave_sum(c):
+            """Sum over the wave of a per-lane count in [0, 63]."""
+            t = fx.Int32(0)
+            for b in range_constexpr(6):
+                t = t + (_popc(_ballot((c & (1 << b)) != 0)) << b)
+            return t
+
+        def compact(keys, wbase, base, pred):
+            """Append the keys ``pred`` keeps at ``base`` if they fit below CAP."""
+            oks = [
+                _one(pred(keys[j], wbase + (j * 64) + lane))
+                for j in range_constexpr(STRIPS)
+            ]
+            cnt = oks[0]
+            for j in range_constexpr(1, STRIPS):
+                cnt = cnt + oks[j]
+            pos = fx.Int32(base)
+            end = fx.Int32(base)
+            for b in range_constexpr(6):  # cnt <= 32
+                bm = _ballot((cnt & (1 << b)) != 0)
+                pos = pos + (_below(bm, 0) << b)
+                end = end + (_popc(bm) << b)
+            # On overflow nothing is written; the caller rebuilds the candidates.
+            if end <= CAP:
+                for j in range_constexpr(STRIPS):
+                    put_pair(
+                        (oks[j] != 0).select(pos, TRASH + lane),
+                        wbase + (j * 64) + lane,
+                        keys[j],
+                    )
+                    pos = pos + oks[j]
+            return end
+
+        def exact_cut(keys, wbase):
+            """The window's 16th key v and the index cut ie among the keys tied
+            with it: (k > v) | (k == v & i >= ie) keeps exactly its top 16."""
+            pe = fx.Int32(0)
+            b = fx.Int32(31)
+            while b >= 0:
+                t = (pe | (fx.Int32(1) << b)) ^ INT_MIN
+                c = fx.Int32(0)
+                for j in range_constexpr(STRIPS):
+                    c = c + _one(keys[j] >= t)
+                pe = (wave_sum(c) >= TOPK).select(t ^ INT_MIN, pe)
+                b = b - 1
+            v = pe ^ INT_MIN
+            c = fx.Int32(0)
+            for j in range_constexpr(STRIPS):
+                c = c + _one(keys[j] > v)
+            need = TOPK - wave_sum(c)
+            ie = fx.Int32(0)
+            b = fx.Int32(IDX_BITS - 1)
+            while b >= 0:
+                t = ie | (fx.Int32(1) << b)
+                c = fx.Int32(0)
+                for j in range_constexpr(STRIPS):
+                    c = c + _one((keys[j] == v) & ((wbase + (j * 64) + lane) >= t))
+                ie = (wave_sum(c) >= need).select(t, ie)
+                b = b - 1
+            return v, ie
+
+        def lane_max_bound(m):
+            p = fx.Int32(0)
+            for d in range_constexpr(12):
+                t = p | _s32(1 << (31 - d))
+                count = _popc(_ballot(m >= (t ^ INT_MIN)))
+                p = (count >= TOPK).select(t, p)
+            return p ^ INT_MIN
+
+        def window(bits, wbase, nrun, tprev):
+            """One window of keys merged into the running top 16 (sel)."""
+            keys = []
+            for j in range_constexpr(STRIPS):
+                live = (wbase + (j * 64) + lane) < nlive
+                keys.append(live.select(_key(bits[j]), fx.Int32(INT_MIN)))
+            m = keys[0]
+            for j in range_constexpr(1, STRIPS):
+                m = _imax(m, keys[j])
+            # Bound: the lane maxima in the first window, the running 16th later.
+            tw = fx.Int32(INT_MIN)
+            if nrun < TOPK:
+                tw = lane_max_bound(m)
+            ta = _imax(tw, tprev)
+            total = compact(keys, wbase, nrun, lambda k, i: k >= ta)
+            # Exact only if 16 keys above -inf bound it and no NaN is above +inf.
+            nan_hi = _ballot(m > KEY_POS_INF) != fx.Int64(0)
+            if (total > CAP) | (ta <= KEY_NEG_INF) | nan_hi:
+                wl = nlive - wbase
+                if wl <= CAP - nrun:
+                    # Every live block of the window is a candidate.
+                    put_pair(
+                        (lane < wl).select(nrun + lane, TRASH + lane),
+                        wbase + lane,
+                        _fold_nan(keys[0]),
+                    )
+                    total = nrun + wl
+                else:
+                    # NaNs to -inf and re-bound; exact search if still too many.
+                    fixed = []
+                    for j in range_constexpr(STRIPS):
+                        live = (wbase + (j * 64) + lane) < nlive
+                        fixed.append(live.select(_fold_nan(keys[j]), fx.Int32(INT_MIN)))
+                    mf = fixed[0]
+                    for j in range_constexpr(1, STRIPS):
+                        mf = _imax(mf, fixed[j])
+                    tf = _imax(lane_max_bound(mf), tprev)
+                    total = compact(fixed, wbase, nrun, lambda k, i: k >= tf)
+                    if total > CAP:
+                        v, ie = exact_cut(fixed, wbase)
+                        total = compact(
+                            fixed,
+                            wbase,
+                            nrun,
+                            lambda k, i: (k > v) | ((k == v) & (i >= ie)),
+                        )
+            if lane < RANK_BLOCK:
+                put_pair(total + lane, fx.Int32(0), fx.Int32(INT_MIN))
+            fx.gpu.barrier()
+            mi, mk = lds_get(atom64, fx.slice(pairs, (None, lane)), 2)
+            mw = _comp(mk, mi)
+            rank = fx.Int32(0)
+            for g in range_constexpr(CAP // RANK_BLOCK):
+                if total > RANK_BLOCK * g:
+                    es = [
+                        lds_get(
+                            atom128, fx.slice(quads, (None, RANK_BLOCK // 2 * g + u)), 4
+                        )
+                        for u in range_constexpr(RANK_BLOCK // 2)
+                    ]
+                    # All of a block's reads in flight before the compares.
+                    fx.rocdl.sched_barrier(0)
+                    for e in es:
+                        rank = rank + _one(_comp(e[1], e[0]) > mw)
+                        rank = rank + _one(_comp(e[3], e[2]) > mw)
+            if (lane < total) & (rank < TOPK):
+                lds_put(atom128, fx.slice(sels, (None, rank)), 4, [mi, mk, 0, mi])
+            return (total < TOPK).select(total, fx.Int32(TOPK))
+
+        wbase = fx.Int32(0)
+        nrun = fx.Int32(0)
+        tprev = fx.Int32(INT_MIN + 1)  # dead keys are INT_MIN
+        while (wbase == 0) | (wbase < nlive):
+            if wbase > 0:
+                bits = load_window(srs_live, wbase)
+                # The running top 16 back into cand; its 16th key bounds the window.
+                fx.gpu.barrier()
+                r = lds_get(atom128, fx.slice(sels, (None, lane & (TOPK - 1))), 4)
+                if lane < nrun:
+                    put_pair(lane, r[0], r[1])
+                k15 = fx.Int32(fx.rocdl.readlane(T.i32, r[1], fx.Int32(TOPK - 1)))
+                tprev = (nrun >= TOPK).select(k15, fx.Int32(INT_MIN + 1))
+                fx.gpu.barrier()
+            nrun = window(
+                [fx.Int32(bits[j]) for j in range_constexpr(STRIPS)],
+                wbase,
+                nrun,
+                tprev,
+            )
+            wbase = wbase + WIN
+
+        fx.gpu.barrier()
+        r = lds_get(atom128, fx.slice(sels, (None, lane & (TOPK - 1))), 4)
+        blk = (lane < nrun).select(r[0], fx.Int32(-1))
+        buffer_ops.buffer_store(
+            blk,
+            rsrc(arg_topk + (h64 * i64_t_head + n64 * i64_t_row) * 4, TOPK * 4),
+            lane,
+            mask=lane < TOPK,
+        )
+
+    @flyc.jit
+    def launch(
+        arg_score: fx.Int64,
+        i64_s_head: fx.Int64,
+        i64_s_row: fx.Int64,
+        i32_width: fx.Int32,
+        arg_nvp: fx.Int64,
+        i32_query_len: fx.Int32,
+        arg_topk: fx.Int64,
+        i64_t_head: fx.Int64,
+        i64_t_row: fx.Int64,
+        i32_rows: fx.Int32,
+        i32_heads: fx.Int32,
+        stream: fx.Stream,
+    ):
+        kernel(
+            arg_score,
+            i64_s_head,
+            i64_s_row,
+            i32_width,
+            arg_nvp,
+            i32_query_len,
+            arg_topk,
+            i64_t_head,
+            i64_t_row,
+        ).launch(
+            grid=(fx.Int64(i32_rows), fx.Int64(i32_heads), 1),
+            block=(64, 1, 1),
+            stream=stream,
+        )
+
+    launch.kernel_name = name
+    launch.compile_hints = PRELOAD_HINTS
+    return launch
diff --git a/vllm/models/minimax_m3/amd/ops/norm_rope_nvfp4_insert/__init__.py b/vllm/models/minimax_m3/amd/ops/norm_rope_nvfp4_insert/__init__.py
new file mode 100644
index 000000000000..208f01a7cb5e
--- /dev/null
+++ b/vllm/models/minimax_m3/amd/ops/norm_rope_nvfp4_insert/__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/minimax_m3/amd/ops/norm_rope_nvfp4_insert/host.py b/vllm/models/minimax_m3/amd/ops/norm_rope_nvfp4_insert/host.py
new file mode 100644
index 000000000000..d6590ee458a4
--- /dev/null
+++ b/vllm/models/minimax_m3/amd/ops/norm_rope_nvfp4_insert/host.py
@@ -0,0 +1,91 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Launch wrapper of the fused norm / RoPE / NVFP4 insert."""
+
+import functools
+import struct
+
+import torch
+
+from vllm.models.minimax_m3.amd.ops.flydsl_utils import (
+    run_compiled,
+)
+
+from . import kernel as _kernel
+
+
+@functools.cache
+def _get_kernel(nq: int, nkv: int, niq: int, rotary_dim: int, tp: int):
+    return _kernel.compile_norm_rope_insert(nq, nkv, niq, rotary_dim, tp)
+
+
+def _token_pairs_per_wave(n: int) -> int:
+    return 1 if n <= 512 else 4
+
+
+def minimax_m3_norm_rope_nvfp4_insert(
+    qkv: torch.Tensor,  # [n, (nq + 2 nkv + niq + 1) * 128] bf16
+    q_norm_weight: torch.Tensor,
+    k_norm_weight: torch.Tensor,
+    index_q_norm_weight: torch.Tensor,
+    index_k_norm_weight: torch.Tensor,
+    cos_sin_cache: torch.Tensor,  # [max_pos, rotary_dim] bf16
+    positions: torch.Tensor,  # [n] int64
+    slot_mapping: torch.Tensor,  # [n] int64, -1 skips the token
+    index_slot_mapping: torch.Tensor,  # [n] int64, -1 skips the token
+    kv_cache: torch.Tensor,  # [blocks, 2 nkv, 128, 72] uint8
+    index_cache: torch.Tensor,  # [blocks, 128, 72] uint8
+    k_scale: torch.Tensor,  # fp32 global scales
+    v_scale: torch.Tensor,
+    q_out: torch.Tensor,  # [n, nq * 128] bf16
+    index_q_out: torch.Tensor,  # [n, niq * 128] e4m3, row stride free
+    num_heads: int,
+    num_kv_heads: int,
+    num_index_heads: int,
+    eps: float,
+) -> None:
+    """Norm + RoPE of q / k / index-q / index-k and the NVFP4 K/V and index-K
+    inserts in one launch; k and index-k are not written back into ``qkv``."""
+    n = qkv.shape[0]
+    if n == 0:
+        return
+    nq, nkv, niq = num_heads, num_kv_heads, num_index_heads
+    rotary_dim = cos_sin_cache.shape[1]
+    assert eps >= 2.0**-126  # the kernel's rsqrt skips ocml's denormal scaling
+    assert qkv.dtype == torch.bfloat16 and qkv.is_contiguous()
+    assert qkv.shape[1] == (nq + 2 * nkv + niq + 1) * _kernel.D
+    assert cos_sin_cache.dtype == torch.bfloat16 and cos_sin_cache.is_contiguous()
+    for w in (q_norm_weight, k_norm_weight, index_q_norm_weight, index_k_norm_weight):
+        assert w.dtype == torch.bfloat16 and w.is_contiguous()
+    assert q_out.dtype == torch.bfloat16 and q_out.is_contiguous()
+    assert index_q_out.dtype == torch.float8_e4m3fn and index_q_out.stride(1) == 1
+    assert index_q_out.stride(0) % 4 == 0
+    assert kv_cache.shape[1:] == (2 * nkv, _kernel.BLK, 72)
+    assert kv_cache.stride(1) == _kernel.BLK * 72 and kv_cache.stride(3) == 1
+    assert index_cache.shape[1:] == (_kernel.BLK, 72) and index_cache.stride(2) == 1
+    assert index_cache.stride(1) == 72
+    run_compiled(
+        _get_kernel(nq, nkv, niq, rotary_dim, _token_pairs_per_wave(n)),
+        qkv.data_ptr(),
+        q_out.data_ptr(),
+        index_q_out.data_ptr(),
+        q_norm_weight.data_ptr(),
+        k_norm_weight.data_ptr(),
+        index_q_norm_weight.data_ptr(),
+        index_k_norm_weight.data_ptr(),
+        cos_sin_cache.data_ptr(),
+        positions.data_ptr(),
+        slot_mapping.data_ptr(),
+        index_slot_mapping.data_ptr(),
+        kv_cache.data_ptr(),
+        index_cache.data_ptr(),
+        k_scale.data_ptr(),
+        v_scale.data_ptr(),
+        n,
+        index_q_out.stride(0),
+        kv_cache.stride(0),
+        index_cache.stride(0),
+        cos_sin_cache.shape[0],
+        struct.unpack("<i", struct.pack("<f", eps))[0],
+        torch.cuda.current_stream(),
+    )
diff --git a/vllm/models/minimax_m3/amd/ops/norm_rope_nvfp4_insert/kernel.py b/vllm/models/minimax_m3/amd/ops/norm_rope_nvfp4_insert/kernel.py
new file mode 100644
index 000000000000..e728cc983d3b
--- /dev/null
+++ b/vllm/models/minimax_m3/amd/ops/norm_rope_nvfp4_insert/kernel.py
@@ -0,0 +1,581 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Gemma norm + RoPE of q / k / index-q / index-k and the NVFP4 K/V + index-K
+page inserts in one launch (FlyDSL, gfx950), bit for bit the ROCm
+``fused_minimax_m3_qknorm_rope_kv_insert`` followed by the NVFP4 page insert.
+
+Waves of 2 * tp tokens: [q heads | k heads | index-k | index-q heads | v heads x tp];
+a norm wave holds 4 dims of a token per lane (the HIP kernel's lane map).
+"""
+
+import flydsl.compiler as flyc
+import flydsl.expr as fx
+from aiter.ops.flydsl.kernels import buffer_ops
+from flydsl._mlir.dialects import arith as _arith
+from flydsl._mlir.dialects import llvm as _llvm
+from flydsl.expr import range_constexpr
+from flydsl.expr.typing import T
+from flydsl.expr.typing import Vector as Vec
+
+from vllm.models.minimax_m3.amd.ops.flydsl_utils import (
+    PRELOAD_HINTS,
+    as_f32,
+    bits,
+    permlane_swap,
+    pk_bf16,
+)
+
+NW = 4
+D = 128
+BLK = 128  # tokens per NVFP4 page
+RCP6 = 0x3E2AAAAB  # RN(1 / 6)
+F448 = 0x43E00000
+F448N = 0xC3E00000 - (1 << 32)
+
+
+def _bf16x4(v):
+    """4 fp32 of a dword pair of bf16 (low half first)."""
+    w0, w1 = fx.Int32(v[0]), fx.Int32(v[1])
+    return [
+        (w0 << 16).bitcast(fx.Float32),
+        (w0 & -65536).bitcast(fx.Float32),
+        (w1 << 16).bitcast(fx.Float32),
+        (w1 & -65536).bitcast(fx.Float32),
+    ]
+
+
+def _norm_rope_out4(xws_t, css_t, w, eps, hi16, first, in_rope, half4):
+    """Norm + RoPE of one token's 4 dims in this lane -> 4 fp32."""
+    cos, sin = css_t
+    e = _bf16x4(xws_t)
+    ss = e[0] * e[0] + e[1] * e[1]
+    ss = ss + e[2] * e[2]
+    ss = ss + e[3] * e[3]
+    for m in (16, 8, 4, 2, 1):
+        ss = _xsum(ss, m)
+    r = _rsq(fx.math.fma(ss, fx.Float32(0.0078125), eps))
+    nrm = [(e[i] * r) * (fx.Float32(1.0) + w[i]) for i in range_constexpr(4)]
+    out = []
+    for i in range_constexpr(4):
+        pr = _xor_lane(bits(nrm[i]), half4, hi16).bitcast(fx.Float32)
+        ps = pr * sin[i]
+        roped = fx.math.fma(nrm[i], cos[i], first.select(-ps, ps))
+        out.append(in_rope.select(roped, nrm[i]))
+    return out
+
+
+def _maxnum(a, b):
+    return fx.Float32(_arith.MaxNumFOp(fx.as_ir_value(a), fx.as_ir_value(b)).result)
+
+
+def _minnum(a, b):
+    return fx.Float32(_arith.MinNumFOp(fx.as_ir_value(a), fx.as_ir_value(b)).result)
+
+
+def _dpp(src, ctrl, old=None, bank_mask=0xF):
+    """v_mov_b32_dpp; an undef ``old`` lets the move fold into its consumer."""
+    o = _llvm.mlir_undef(T.i32) if old is None else fx.as_ir_value(old)
+    full = old is None and bank_mask == 0xF
+    return fx.Int32(
+        fx.rocdl.update_dpp(T.i32, o, fx.as_ir_value(src), ctrl, 0xF, bank_mask, full)
+    )
+
+
+def _swap16(v):
+    """permlane16_swap(v, v): (own, lane ^ 16) in rows 0 / 2, swapped in 1 / 3."""
+    return permlane_swap(fx.rocdl.permlane16_swap, v, v)
+
+
+def _xor_lane(v, m, hi16):
+    """Lane l ^ m of the 32-lane half; ``hi16`` is this lane's bit 4."""
+    if m == 1:
+        return _dpp(v, 0xB1)  # quad_perm [1, 0, 3, 2]
+    if m == 2:
+        return _dpp(v, 0x4E)  # quad_perm [2, 3, 0, 1]
+    if m == 4:  # row_shr:4 into banks 1, 3, then row_shl:4 into banks 0, 2
+        return _dpp(v, 0x104, old=_dpp(v, 0x114, bank_mask=0xA), bank_mask=0x5)
+    if m == 8:
+        return _dpp(v, 0x128)  # row_ror:8
+    assert m == 16
+    a, b = _swap16(v)
+    return hi16.select(a, b)
+
+
+def _xsum(x, m):
+    """Adds lane l ^ m's value to ``x``."""
+    if m == 16:
+        a, b = _swap16(bits(x))
+        return a.bitcast(fx.Float32) + b.bitcast(fx.Float32)
+    return x + _xor_lane(bits(x), m, None).bitcast(fx.Float32)
+
+
+def _xmax(x, m):
+    return _maxnum(x, _xor_lane(bits(x), m, None).bitcast(fx.Float32))
+
+
+def _rsq(x):
+    """v_rsq_f32; ocml's denormal scaling never applies as x >= eps >= 2^-126."""
+    return fx.Float32(fx.rocdl.rsq(T.f32, fx.as_ir_value(x)))
+
+
+def _e4m3_byte(x):
+    """RNE e4m3 byte of ``x``."""
+    pk = fx.rocdl.cvt_scalef32_pk_fp8_f32(
+        T.vec(2, T.i16),
+        fx.as_ir_value(Vec.filled(2, 0, fx.Int16)),
+        fx.as_ir_value(x),
+        fx.as_ir_value(x),
+        fx.as_ir_value(fx.Float32(1.0)),
+        False,
+    )
+    return Vec(pk).bitcast(fx.Int32)[0] & 0xFF
+
+
+def _e4m3_value(byte):
+    pk = fx.rocdl.cvt_scalef32_pk_f32_fp8(
+        T.vec(2, T.f32),
+        fx.as_ir_value(fx.Int32(byte)),
+        fx.as_ir_value(fx.Float32(1.0)),
+        False,
+    )
+    return fx.Float32(Vec(pk)[0])
+
+
+def _iq_fp8(a, b, old, hi):
+    """Two fp32 -> e4m3 bytes in the low (``hi``: high) half of ``old``."""
+    return fx.Int32(
+        fx.rocdl.cvt_pk_fp8_f32(
+            T.i32,
+            fx.as_ir_value(fx.Float32(a)),
+            fx.as_ir_value(fx.Float32(b)),
+            fx.as_ir_value(fx.Int32(old)),
+            hi,
+        )
+    )
+
+
+def _e2m1x2(a, b, old, sel):
+    """E2M1 codes of ``a``, ``b`` into byte ``sel`` of ``old``, ``a`` low."""
+    return fx.Int32(
+        fx.rocdl.cvt_scalef32_pk_fp4_f32(
+            T.i32,
+            fx.as_ir_value(fx.Int32(old)),
+            fx.as_ir_value(a),
+            fx.as_ir_value(b),
+            fx.as_ir_value(fx.Float32(1.0)),
+            sel,
+        )
+    )
+
+
+def _quant4(xs, gs, inv_gs):
+    """E2M1 codes of this lane's 4 values (16 bits) and the e4m3 scale byte of
+    their 16-value group (lanes l ^ 0..3)."""
+    a = [fx.math.absf(x) for x in xs]
+    m = _maxnum(_maxnum(a[0], a[1]), _maxnum(a[2], a[3]))
+    m = _xmax(_xmax(m, 1), 2)
+    sf = _minnum((m * as_f32(RCP6)) * inv_gs, as_f32(F448))
+    byte = _e4m3_byte(sf)
+    sv = _e4m3_value(byte)
+    quant = (sv != 0.0).select(fx.Float32(1.0) / (gs * sv), fx.Float32(0.0))
+    y = [x * quant for x in xs]
+    return _e2m1x2(y[2], y[3], _e2m1x2(y[0], y[1], fx.Int32(0), 0), 1), byte
+
+
+def _v_row(state, row, r, n_run, p0, vs, inv_vs):
+    """Adds row ``r`` of a V run at octet positions p0.. to ``state``: new
+    nibbles and scale bytes, with their masks."""
+    words, written = state[0:4], state[4]
+    sc, sc_mask = state[5:7], state[7:9]
+    nib, sf = _quant4(_bf16x4(row), vs, inv_vs)
+    in_run = n_run > r
+    p = p0 + r
+    sh = p * 4
+    words = [
+        wd | in_run.select(((nib >> (4 * i)) & 15) << sh, fx.Int32(0))
+        for i, wd in enumerate(words)
+    ]
+    written = written | in_run.select(fx.Int32(15) << sh, fx.Int32(0))
+    bsh = (p % 4) * 8
+    for k in range_constexpr(2):
+        here = in_run & ((p // 4) == k)
+        sc[k] = sc[k] | here.select(sf << bsh, fx.Int32(0))
+        sc_mask[k] = sc_mask[k] | here.select(fx.Int32(255) << bsh, fx.Int32(0))
+    return [*words, written, *sc, *sc_mask]
+
+
+def _gptr(addr, ty, align):
+    return fx.inttoptr(
+        fx.PointerType.get(ty, address_space=fx.AddressSpace.Global, alignment=align),
+        fx.Int64(addr),
+    )
+
+
+def _k_scale_off(tok, grp):
+    """Byte of (token, group) in a head's K scale region."""
+    r = tok % 32
+    n = (r // 8) * 4 + r % 4
+    return ((((tok // 32) * 4 + grp // 2) * 16 + n) * 2 + (r % 8) // 4) * 2 + grp % 2
+
+
+def compile_norm_rope_insert(nq: int, nkv: int, niq: int, rotary_dim: int, tp: int):
+    """Build the kernel; ``tp`` token pairs per norm / RoPE wave, one per V wave."""
+    assert rotary_dim in (8, 16, 32, 64, 128)
+    NR = nq + nkv + 1 + niq  # norm + RoPE waves
+    S = NR + nkv * tp  # waves per 2 * tp tokens
+    ROW = (nq + 2 * nkv + niq + 1) * D
+    HALF = rotary_dim // 2
+    SIDE = nkv * BLK * 72  # bytes of one side of a main page
+    TPW = 2 * tp
+    name = f"m3_norm_rope_nvfp4_insert_q{nq}_kv{nkv}_iq{niq}_r{rotary_dim}_t{tp}"
+
+    @flyc.kernel(name=name, known_block_size=[64 * NW, 1, 1])
+    def kernel(
+        i32_n: fx.Int32,
+        i32_cs_rows: fx.Int32,
+        arg_qkv: fx.Int64,  # [n, ROW] bf16
+        arg_pos: fx.Int64,  # positions [n] i64
+        arg_slot: fx.Int64,  # main slots [n] i64, -1 skips
+        arg_islot: fx.Int64,  # index slots [n] i64, -1 skips
+        arg_cs: fx.Int64,  # cos_sin_cache [i32_cs_rows, rotary_dim] bf16
+        arg_qw: fx.Int64,  # norm weights [128] bf16: q, k, index-q, index-k
+        arg_kw: fx.Int64,
+        arg_iqw: fx.Int64,
+        arg_ikw: fx.Int64,
+        arg_q: fx.Int64,  # q_out [n, nq * 128] bf16
+        arg_iq: fx.Int64,  # index_q_out [n, niq * 128] e4m3, row stride i32_iq_st
+        arg_kv: fx.Int64,  # [blocks, 2 * nkv, 128, 72] u8
+        arg_idx: fx.Int64,  # [blocks, 128, 72] u8
+        arg_ks: fx.Int64,  # fp32 global scales
+        arg_vs: fx.Int64,
+        i64_kv_st: fx.Int64,
+        i64_idx_st: fx.Int64,
+        i32_iq_st: fx.Int32,
+        i32_eps: fx.Int32,  # fp32 bits
+    ):
+        tx, bx = fx.thread_idx.x, fx.block_idx.x
+        lane = tx % 64
+        wave = fx.Int32(fx.rocdl.readfirstlane(T.i32, tx // 64))
+        wid = bx * NW + wave
+        s = wid % S
+        hl = lane % 32
+        hi16 = (hl & 16) != 0
+        dim0 = hl * 4
+        toks = [(wid // S) * TPW + t * 2 + lane // 32 for t in range_constexpr(tp)]
+        lives = [tk < i32_n for tk in toks]
+
+        qkv = buffer_ops.create_buffer_resource_from_addr(
+            arg_qkv, num_records_bytes=fx.Int64(i32_n) * (ROW * 2)
+        )
+        sl_r = buffer_ops.create_buffer_resource_from_addr(
+            arg_slot, num_records_bytes=fx.Int64(i32_n) * 8
+        )
+
+        if s < NR:
+            is_q = s < nq
+            is_k = (s >= nq) & (s < nq + nkv)
+            is_ik = s == nq + nkv
+            is_kq = is_k | is_ik
+            # This wave's head column in the qkv row.
+            col = (s <= nq + nkv).select(
+                is_ik.select(fx.Int32(nq + 2 * nkv + niq), s),
+                s - (nq + nkv + 1) + (nq + 2 * nkv),
+            )
+            w_addr = is_q.select(
+                arg_qw, is_k.select(arg_kw, is_ik.select(arg_ikw, arg_iqw))
+            )
+            wr = buffer_ops.create_buffer_resource_from_addr(
+                w_addr, num_records_bytes=D * 2
+            )
+            pos_r = buffer_ops.create_buffer_resource_from_addr(
+                arg_pos, num_records_bytes=fx.Int64(i32_n) * 8
+            )
+            isl_r = buffer_ops.create_buffer_resource_from_addr(
+                arg_islot, num_records_bytes=fx.Int64(i32_n) * 8
+            )
+            cs_r = buffer_ops.create_buffer_resource_from_addr(
+                arg_cs, num_records_bytes=fx.Int64(i32_cs_rows) * (rotary_dim * 2)
+            )
+            xws, poss, slots = [], [], []
+            for t in range_constexpr(tp):
+                xws.append(
+                    Vec(
+                        buffer_ops.buffer_load(
+                            qkv,
+                            toks[t] * (ROW // 2) + col * (D // 2) + hl * 2,
+                            vec_width=2,
+                            dtype=fx.Int32,
+                            mask=lives[t],
+                        )
+                    )
+                )
+                poss.append(
+                    fx.Int32(
+                        buffer_ops.buffer_load(
+                            pos_r,
+                            toks[t] * 2,
+                            vec_width=1,
+                            dtype=fx.Int32,
+                            mask=lives[t],
+                        )
+                    )
+                )
+                kslot = fx.Int32(
+                    buffer_ops.buffer_load(
+                        sl_r,
+                        toks[t] * 2,
+                        vec_width=1,
+                        dtype=fx.Int32,
+                        mask=lives[t] & is_k,
+                    )
+                )
+                islot = fx.Int32(
+                    buffer_ops.buffer_load(
+                        isl_r,
+                        toks[t] * 2,
+                        vec_width=1,
+                        dtype=fx.Int32,
+                        mask=lives[t] & is_ik,
+                    )
+                )
+                slots.append(is_k.select(kslot, islot))
+            w = _bf16x4(
+                Vec(buffer_ops.buffer_load(wr, hl * 2, vec_width=2, dtype=fx.Int32))
+            )
+            first = dim0 < HALF
+            in_rope = dim0 < rotary_dim
+            ib = first.select(dim0, dim0 - HALF)
+            css = []
+            for t in range_constexpr(tp):
+                cs_off = (poss[t] * rotary_dim + ib) // 2
+                css.append(
+                    (
+                        _bf16x4(
+                            Vec(
+                                buffer_ops.buffer_load(
+                                    cs_r,
+                                    cs_off,
+                                    vec_width=2,
+                                    dtype=fx.Int32,
+                                    mask=lives[t] & in_rope,
+                                )
+                            )
+                        ),
+                        _bf16x4(
+                            Vec(
+                                buffer_ops.buffer_load(
+                                    cs_r,
+                                    cs_off + HALF // 2,
+                                    vec_width=2,
+                                    dtype=fx.Int32,
+                                    mask=lives[t] & in_rope,
+                                )
+                            )
+                        ),
+                    )
+                )
+            eps = i32_eps.bitcast(fx.Float32)
+            qvs = []
+            for t in range_constexpr(tp):
+                out = _norm_rope_out4(
+                    xws[t], css[t], w, eps, hi16, first, in_rope, HALF // 4
+                )
+                lo, hi = pk_bf16(out[0], out[1]), pk_bf16(out[2], out[3])
+                qvs.append(Vec.from_elements([lo, hi], fx.Int32))
+            if is_q:
+                q_r = buffer_ops.create_buffer_resource_from_addr(
+                    arg_q, num_records_bytes=fx.Int64(i32_n) * (nq * D * 2)
+                )
+                for t in range_constexpr(tp):
+                    buffer_ops.buffer_store(
+                        qvs[t],
+                        q_r,
+                        toks[t] * (nq * D // 2) + s * (D // 2) + hl * 2,
+                        mask=lives[t],
+                    )
+            if s > nq + nkv:
+                iq_r = buffer_ops.create_buffer_resource_from_addr(
+                    arg_iq, num_records_bytes=fx.Int64(i32_n) * fx.Int64(i32_iq_st)
+                )
+                iq_col = (s - (nq + nkv + 1)) * D
+                for t in range_constexpr(tp):
+                    f = [
+                        _minnum(_maxnum(x, as_f32(F448N)), as_f32(F448))
+                        for x in _bf16x4(qvs[t])
+                    ]
+                    word = _iq_fp8(
+                        f[2], f[3], _iq_fp8(f[0], f[1], fx.Int32(0), False), True
+                    )
+                    buffer_ops.buffer_store(
+                        word,
+                        iq_r,
+                        toks[t] * i32_iq_st + iq_col + dim0,
+                        mask=lives[t],
+                        offset_is_bytes=True,
+                    )
+            if is_kq:
+                gs = is_k.select(
+                    fx.Float32(_gptr(arg_ks, T.f32, 4)[0]), fx.Float32(1.0)
+                )
+                inv_gs = fx.Float32(1.0) / gs
+                quants = [
+                    _quant4(_bf16x4(qvs[t]), gs, inv_gs) for t in range_constexpr(tp)
+                ]
+                h = s - nq
+                grp = hl // 4
+                for t in range_constexpr(tp):
+                    slot = slots[t]
+                    tp_ = slot & (BLK - 1)
+                    page = is_k.select(
+                        arg_kv + fx.Int64(slot >> 7) * i64_kv_st,
+                        arg_idx + fx.Int64(slot >> 7) * i64_idx_st,
+                    )
+                    d_off = is_k.select(h * (BLK * 64), fx.Int32(0)) + tp_ * 64 + hl * 2
+                    s_off = is_k.select(
+                        nkv * BLK * 64 + h * (BLK * 8) + _k_scale_off(tp_, grp),
+                        BLK * 64 + tp_ * 8 + grp,
+                    )
+                    d_ptr = _gptr(page + fx.Int64(d_off), T.i16, 2)
+                    s_ptr = _gptr(page + fx.Int64(s_off), T.i8, 1)
+                    packed = quants[t][0].to(fx.Int16)
+                    sf8 = quants[t][1].to(fx.Int8)
+                    if lives[t] & (slot >= 0):
+                        d_ptr[0] = packed
+                        if hl % 4 == 0:
+                            s_ptr[0] = sf8
+        else:
+            vh = (s - NR) // tp
+            tok = (wid // S) * TPW + ((s - NR) % tp) * 2 + lane // 32
+            live = tok < i32_n
+            vcol = (nq + nkv + vh) * (D // 2)  # V head column (dwords)
+            vs = fx.Float32(_gptr(arg_vs, T.f32, 4)[0])
+            near_t = tok - 7 + hl
+            has_near = (hl < 16) & (near_t >= 0) & (near_t < i32_n)
+            near = has_near.select(
+                fx.Int32(
+                    buffer_ops.buffer_load(
+                        sl_r, near_t * 2, vec_width=1, dtype=fx.Int32, mask=has_near
+                    )
+                ),
+                fx.Int32(-1),
+            )
+            slot = fx.Int32(
+                buffer_ops.buffer_load(
+                    sl_r, tok * 2, vec_width=1, dtype=fx.Int32, mask=live
+                )
+            )
+            has_prev = live & (tok > 0)
+            prev = has_prev.select(
+                fx.Int32(
+                    buffer_ops.buffer_load(
+                        sl_r, tok * 2 - 2, vec_width=1, dtype=fx.Int32, mask=has_prev
+                    )
+                ),
+                fx.Int32(-1),
+            )
+            # The first token of a run in one 8-token octet writes the run and
+            # keeps the octet's other nibbles and scales. Rows it can cover load
+            # before its slot is known.
+            rows = [
+                Vec(
+                    buffer_ops.buffer_load(
+                        qkv,
+                        (tok + r) * (ROW // 2) + vcol + hl * 2,
+                        vec_width=2,
+                        dtype=fx.Int32,
+                        mask=tok + r < i32_n,
+                    )
+                )
+                for r in range_constexpr(8)
+            ]
+            base = slot - slot % 8
+            in_oct = (near >= base) & (near < base + 8)
+            votes = fx.Int64(fx.rocdl.ballot(T.i64, in_oct))
+            mine = (votes >> fx.Int64((lane // 32) * 32)).to(fx.Int32) & 0xFFFF
+            n_run = fx.Int32(fx.math.ctpop(mine))
+            run_first = (prev < 0) | ((prev >> 3) != (slot >> 3))
+            tp_ = slot & (BLK - 1)
+            v_side = arg_kv + fx.Int64(slot >> 7) * i64_kv_st + SIDE
+            v_data = v_side + fx.Int64(vh * (BLK * 64) + (tp_ // 8) * (D * 4))
+            v_sf = v_side + fx.Int64(nkv * BLK * 64 + vh * (BLK * 8) + tp_ - tp_ % 8)
+            dp = _gptr(v_data + fx.Int64(dim0 * 4), T.i32, 4)
+            sp = _gptr(v_sf + fx.Int64((hl // 4) * BLK), T.i32, 4)
+            go = live & (slot >= 0) & run_first
+            inv_vs = fx.Float32(1.0) / vs
+            # Runs of <= 4 tokens (decode) take the short tile.
+            for run_rows, take in (
+                (rows[:4], go & (n_run <= 4)),
+                (rows, go & (n_run > 4)),
+            ):
+                if take:
+                    old = [fx.Int32(dp[i]) for i in range_constexpr(4)]
+                    old_sc = [fx.Int32(sp[i]) for i in range_constexpr(2)]
+                    st = [fx.Int32(0)] * 9
+                    for r in range_constexpr(len(run_rows)):
+                        st = _v_row(st, run_rows[r], r, n_run, tp_ % 8, vs, inv_vs)
+                    keep = ~st[4]
+                    for i in range_constexpr(4):
+                        dp[i] = st[i] | (old[i] & keep)
+                    sc0 = st[5] | (old_sc[0] & ~st[7])
+                    sc1 = st[6] | (old_sc[1] & ~st[8])
+                    if hl % 4 == 0:
+                        sp[0] = sc0
+                        sp[1] = sc1
+
+    @flyc.jit
+    def launch(
+        arg_qkv: fx.Int64,
+        arg_q: fx.Int64,
+        arg_iq: fx.Int64,
+        arg_qw: fx.Int64,
+        arg_kw: fx.Int64,
+        arg_iqw: fx.Int64,
+        arg_ikw: fx.Int64,
+        arg_cs: fx.Int64,
+        arg_pos: fx.Int64,
+        arg_slot: fx.Int64,
+        arg_islot: fx.Int64,
+        arg_kv: fx.Int64,
+        arg_idx: fx.Int64,
+        arg_ks: fx.Int64,
+        arg_vs: fx.Int64,
+        i32_n: fx.Int32,
+        i32_iq_st: fx.Int32,
+        i64_kv_st: fx.Int64,
+        i64_idx_st: fx.Int64,
+        i32_cs_rows: fx.Int32,
+        i32_eps: fx.Int32,
+        stream: fx.Stream,
+    ):
+        waves = (i32_n + (TPW - 1)) // TPW * S
+        kernel(
+            i32_n,
+            i32_cs_rows,
+            arg_qkv,
+            arg_pos,
+            arg_slot,
+            arg_islot,
+            arg_cs,
+            arg_qw,
+            arg_kw,
+            arg_iqw,
+            arg_ikw,
+            arg_q,
+            arg_iq,
+            arg_kv,
+            arg_idx,
+            arg_ks,
+            arg_vs,
+            i64_kv_st,
+            i64_idx_st,
+            i32_iq_st,
+            i32_eps,
+        ).launch(
+            grid=(fx.Int64((waves + (NW - 1)) // NW), 1, 1),
+            block=(64 * NW, 1, 1),
+            stream=stream,
+        )
+
+    launch.kernel_name = name
+    launch.compile_hints = PRELOAD_HINTS
+    return launch
diff --git a/vllm/models/minimax_m3/amd/ops/nvfp4_index_decode/__init__.py b/vllm/models/minimax_m3/amd/ops/nvfp4_index_decode/__init__.py
new file mode 100644
index 000000000000..208f01a7cb5e
--- /dev/null
+++ b/vllm/models/minimax_m3/amd/ops/nvfp4_index_decode/__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/minimax_m3/amd/ops/nvfp4_index_decode/host.py b/vllm/models/minimax_m3/amd/ops/nvfp4_index_decode/host.py
new file mode 100644
index 000000000000..37e260514bb0
--- /dev/null
+++ b/vllm/models/minimax_m3/amd/ops/nvfp4_index_decode/host.py
@@ -0,0 +1,222 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Launch wrappers of the FlyDSL NVFP4 decode index scorers."""
+
+import functools
+
+import torch
+
+from vllm.models.minimax_m3.amd.ops.flydsl_utils import run_compiled
+from vllm.models.minimax_m3.amd.ops.nvfp4_index_prefill.host import nvfp4_e4m3_lut
+from vllm.utils.math_utils import cdiv
+
+from . import kernel as _kernel
+
+# 158 VGPRs = 3 waves per SIMD; more waves add a second round of workgroups.
+TARGET_WAVES = 3072
+# Scores are buffered in LDS and stored by whole rows once the batch's score
+# footprint reaches this; smaller footprints store directly.
+BUFFER_MIN_BYTES = 8 << 20
+UNITS_PER_SLOT = 4
+
+
+def _end_bytes(t: torch.Tensor) -> int:
+    """Bytes from the data pointer to the end of the last element."""
+    return (
+        sum((n - 1) * s for n, s in zip(t.shape, t.stride())) + 1
+    ) * t.element_size()
+
+
+def _unit(max_blocks: int, slots: int, umax: int, footprint: int) -> int:
+    """Consecutive blocks per buffered unit (a power of two, 4..umax) giving
+    every wave slot >= UNITS_PER_SLOT units, or 0 for direct stores."""
+    if footprint < BUFFER_MIN_BYTES:
+        return 0
+    u = umax
+    while u > 4 and cdiv(max_blocks, u) < UNITS_PER_SLOT * slots:
+        u //= 2
+    return u if cdiv(max_blocks, u) >= UNITS_PER_SLOT * slots else 0
+
+
+@functools.cache
+def _get_kernel(rows: int, heads: int, buffered: bool):
+    return _kernel.compile_score_decode(rows, heads, buffered)
+
+
+def _supported(
+    q_idx: torch.Tensor,
+    index_cache: torch.Tensor,
+    score: torch.Tensor,
+    block_table: torch.Tensor,
+    seq_lens: torch.Tensor,
+    query_len: int,
+) -> bool:
+    from vllm.platforms.rocm import on_gfx950
+
+    num_heads = q_idx.shape[1]
+    return (
+        on_gfx950()
+        and num_heads in (1, 2, 4)
+        and 1 <= num_heads * query_len <= _kernel.MAX_ROWS
+        and q_idx.dtype == torch.float8_e4m3fn
+        and q_idx.shape[2] == _kernel.D
+        and q_idx.stride(2) == 1
+        and q_idx.stride(0) % 16 == 0
+        and q_idx.stride(1) % 16 == 0
+        and q_idx.data_ptr() % 16 == 0
+        and index_cache.dtype == torch.uint8
+        and index_cache.stride(2) == 1
+        and index_cache.stride(1) == _kernel.D // 2 + _kernel.D // 16
+        and index_cache.stride(0) % 16 == 0
+        and index_cache.data_ptr() % 16 == 0
+        and score.dtype == torch.float32
+        and score.stride(2) == 1
+        and block_table.dtype == torch.int32
+        and block_table.stride(1) == 1
+        and seq_lens.dtype == torch.int32
+        and seq_lens.is_contiguous()
+    )
+
+
+def minimax_m3_nvfp4_index_score_decode(
+    q_idx: torch.Tensor,  # [num_reqs * query_len, H, 128] e4m3
+    index_cache: torch.Tensor,  # [num_blocks, 128, 72] uint8
+    score: torch.Tensor,  # [H, num_reqs * query_len, >= max_blocks] fp32
+    block_table: torch.Tensor,  # [num_reqs, >= max_blocks] int32
+    seq_lens: torch.Tensor,  # [num_reqs] int32
+    init_blocks: int,
+    local_blocks: int,
+    query_len: int,
+    max_seq_len: int,
+) -> None:
+    """Block scores of every decode row for blocks below cdiv(seq_len, 128)."""
+    total_q, num_heads, _ = q_idx.shape
+    num_reqs = total_q // query_len
+    if num_reqs == 0:
+        return
+    assert _supported(q_idx, index_cache, score, block_table, seq_lens, query_len)
+    rows = num_heads * query_len
+    max_blocks = max(1, cdiv(max_seq_len, _kernel.BLK))
+    # A chunk of NW x (pages per wave) requests; with at most one page per
+    # wave, chunks of one request and each wave on its own positions.
+    pages_per_wave = cdiv(num_reqs * max_blocks, TARGET_WAVES)
+    s_bytes = _end_bytes(score)
+    assert s_bytes < 2**31 * 4, "score offsets must fit in int32 elements"
+    footprint = num_reqs * rows * max_blocks * 4
+    chunk = 1 if pages_per_wave <= 1 else _kernel.NW * pages_per_wave
+    chunk_cap = _kernel.chunk_max(_kernel.padded_rows(rows))
+    chunk = max(1, min(chunk, chunk_cap, num_reqs))
+    nch = cdiv(num_reqs, chunk)
+    chunk = cdiv(num_reqs, nch)
+    per_wave = chunk < _kernel.NW
+    wanted = max(1, TARGET_WAVES // (_kernel.NW * nch))
+    # Past 16 rows the buffered variant's flush costs more than it saves.
+    unit = (
+        _unit(
+            max_blocks,
+            wanted * (_kernel.NW if per_wave else 1),
+            _kernel.UMAX,
+            footprint,
+        )
+        if _kernel.padded_rows(rows) <= 16
+        else 0
+    )
+    units = cdiv(max_blocks, max(1, unit))
+    positions = cdiv(units, _kernel.NW) if per_wave else units
+    splits = max(1, min(positions, wanted))
+    run_compiled(
+        _get_kernel(rows, num_heads, unit > 0),
+        q_idx.data_ptr(),
+        index_cache.data_ptr(),
+        nvfp4_e4m3_lut(q_idx.device).data_ptr(),
+        score.data_ptr(),
+        block_table.data_ptr(),
+        seq_lens.data_ptr(),
+        int(index_cache.stride(0)),
+        int(q_idx.stride(0)),
+        int(q_idx.stride(1)),
+        _end_bytes(q_idx),
+        int(score.stride(0)),
+        int(score.stride(1)),
+        s_bytes,
+        int(block_table.stride(0)),
+        _end_bytes(block_table),
+        int(num_reqs),
+        int(chunk),
+        int(nch),
+        int(splits),
+        int(init_blocks),
+        int(local_blocks),
+        int(per_wave),
+        int(unit),
+        torch.cuda.current_stream(),
+        # Keep the buffered variant at the plain one's occupancy.
+        waves_per_eu=3 if unit > 0 and _kernel.padded_rows(rows) <= 16 else None,
+    )
+
+
+@functools.cache
+def _get_kernel_cp_push():
+    from . import kernel_cp_push
+
+    return kernel_cp_push.compile_score_push()
+
+
+def minimax_m3_nvfp4_index_score_push_cp(
+    iq: torch.Tensor,  # this rank's index head [num_reqs * 4, 128] e4m3
+    index_cache: torch.Tensor,  # [num_blocks, 128, 72] uint8
+    dests: list[int],  # head h's scores [rows, row_stride] fp32 on rank h
+    row_stride: int,
+    dest_bytes: int,
+    block_table: torch.Tensor,  # [num_reqs, >= max_blocks] int32
+    seq_lens: torch.Tensor,  # [num_reqs] int32
+    init_blocks: int,
+    local_blocks: int,
+    rank: int,
+    counts: torch.Tensor,  # index_cp's device counters
+    flags: list[int],  # rank p's flag words
+    iq_bufs: list[int],  # rank p's index-q receive buffer [4, iq_rows, 128] e4m3
+    iq_rows: int,
+    grid: int = 0,
+) -> None:
+    """This rank's share of all four heads' decode scores, pushed to the heads'
+    ranks; returns once every rank's scores have arrived."""
+    from . import kernel_cp_push
+
+    rows = iq.shape[0]
+    assert rows % 4 == 0 and rows <= iq_rows
+    assert iq.dtype == torch.float8_e4m3fn and iq.shape[1] == kernel_cp_push.D
+    assert iq.stride(1) == 1 and iq.stride(0) == kernel_cp_push.D
+    num_reqs = rows // 4
+    assert num_reqs <= kernel_cp_push.MAX_REQS
+    if not grid:
+        small = num_reqs <= kernel_cp_push.GRID_SMALL_REQS
+        grid = kernel_cp_push.GRID_SMALL if small else kernel_cp_push.GRID
+    assert 0 < grid <= kernel_cp_push.SC_MAX_WG
+    if num_reqs == 0:
+        return
+    assert dest_bytes < 2**31 * 4, "score offsets must fit in int32 elements"
+    run_compiled(
+        _get_kernel_cp_push(),
+        iq.data_ptr(),
+        index_cache.data_ptr(),
+        nvfp4_e4m3_lut(iq.device).data_ptr(),
+        *dests,
+        block_table.data_ptr(),
+        seq_lens.data_ptr(),
+        counts.data_ptr(),
+        *flags,
+        *iq_bufs,
+        int(iq_rows),
+        int(index_cache.stride(0)),
+        int(row_stride),
+        int(dest_bytes),
+        int(block_table.stride(0)),
+        _end_bytes(block_table),
+        int(num_reqs),
+        int(init_blocks),
+        int(local_blocks),
+        int(rank),
+        int(grid),
+        torch.cuda.current_stream(),
+    )
diff --git a/vllm/models/minimax_m3/amd/ops/nvfp4_index_decode/kernel.py b/vllm/models/minimax_m3/amd/ops/nvfp4_index_decode/kernel.py
new file mode 100644
index 000000000000..a80f10463a51
--- /dev/null
+++ b/vllm/models/minimax_m3/amd/ops/nvfp4_index_decode/kernel.py
@@ -0,0 +1,564 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""NVFP4 decode index scorer for MiniMax-M3 (FlyDSL, gfx950).
+
+Per (head, query row, block): the max of q . k over the row's visible tokens,
+with AITER's init / local sentinels, on e4m3_rn_satfinite(e2m1 x scale) keys. A
+workgroup takes a chunk of up to 64 requests (lane r holds request r) and every
+splits-th block position. The distinct pages at a position (a shared prefix page
+is one page for the whole chunk) go round robin over the 4 waves; a wave
+dequantizes its page into the A fragments of v_mfma_f32_16x16x128_f8f6f4, issuing
+the next page's loads into the freed registers, and scores it against every
+16-column group of Q rows holding the page.
+"""
+
+import flydsl.compiler as flyc
+import flydsl.expr as fx
+from aiter.ops.flydsl.kernels import buffer_ops
+from flydsl.expr import const_expr, range_constexpr
+from flydsl.expr.typing import T
+from flydsl.expr.typing import Vector as Vec
+
+from vllm.models.minimax_m3.amd.ops.flydsl_utils import (
+    NT,
+    as_f32,
+    bits,
+    dequant8,
+    global_i32_ptr,
+    key_order4,
+    maxf_nn,
+    mfma_fp8_16x16x128,
+    pack8,
+    xlane4,
+)
+from vllm.models.minimax_m3.amd.ops.nvfp4_index_prefill.kernel import with_sentinels
+
+NW = 4
+MAX_REQS = 64  # requests per chunk, one per lane
+BLK = 128  # tokens per block
+D = 128  # index head dim
+PAGE_CODES = BLK * D // 2  # page: [128 x 64 B] e2m1 codes, [128 x 8 B] e4m3 scales
+PAGE_BYTES = PAGE_CODES + BLK * 8
+MT = BLK // 16  # token tiles per page
+UMAX = 8  # max consecutive blocks per LDS-buffered unit
+MAX_ROWS = 32  # H * query_len
+
+
+def padded_rows(rows: int) -> int:
+    """Columns a request takes in the 16-column groups: rows rounded up to a
+    power of two."""
+    return 1 << (rows - 1).bit_length()
+
+
+def chunk_max(rows: int) -> int:
+    """Requests per chunk: one per lane, and 256 Q rows (32 KB) in LDS."""
+    return min(MAX_REQS, 256 // rows)
+
+
+def block_sentinels(mx, visible, b, i32_init, i32_local):
+    """AITER's sentinels on block b of a row seeing ``visible`` tokens."""
+    nb = ((visible < 0).select(fx.Int32(0), visible) + (BLK - 1)) // BLK
+    return with_sentinels(mx, b, nb, i32_init, i32_local)
+
+
+def compile_score_decode(rows: int, heads: int, buffered: bool = False):
+    """Build the scorer.
+
+    Args:
+        rows: H * query_len (1..32).
+        heads: index heads (1, 2 or 4).
+        buffered: gather a unit's scores in LDS and store whole rows (large
+            score footprints); otherwise store them directly.
+
+    """
+    assert 1 <= rows <= MAX_ROWS and heads in (1, 2, 4) and rows % heads == 0
+    H = heads
+    QL = rows // H
+    R = padded_rows(rows)  # a request's columns; column j = query row j // H
+    RPC = max(1, 16 // R)  # requests per 16-column group
+    CPR = max(1, R // 16)  # 16-column groups per request
+    GMASK = (1 << RPC) - 1
+    QROWS = chunk_max(R) * R
+    QROW_BYTES = D
+    name = f"m3_nvfp4_index_score_decode_r{rows}h{H}" + ("_buf" if buffered else "")
+    SB_STRIDE = UMAX + 1  # padded against LDS bank conflicts
+
+    @fx.struct
+    class Shared:
+        q: fx.Array[fx.Uint8, QROWS * QROW_BYTES, 16]
+        lut: fx.Array[fx.Int32, 512, 16]
+        seq: fx.Array[fx.Int32, MAX_REQS]
+        # The unit's scores as f32 bits, plus 64 dummy slots.
+        sbuf: fx.Array[fx.Int32, (QROWS * SB_STRIDE + 64) if buffered else 16, 16]
+
+    @flyc.kernel(name=name, known_block_size=[64 * NW, 1, 1])
+    def kernel(
+        arg_q: fx.Int64,  # [num_reqs * QL, H, 128] e4m3
+        arg_kv: fx.Int64,  # [num_blocks, 128, 72] uint8
+        arg_lut: fx.Int64,  # nvfp4_e4m3_lut [256] int64
+        arg_score: fx.Int64,  # [H, num_reqs * QL, S] fp32
+        arg_bt: fx.Int64,  # [num_reqs, >= max_blocks] int32
+        arg_seq: fx.Int64,  # [num_reqs] int32
+        i64_page_stride: fx.Int64,
+        i32_q_st: fx.Int32,  # q strides (bytes)
+        i32_q_sh: fx.Int32,
+        i64_q_bytes: fx.Int64,
+        i32_s_sh: fx.Int32,  # score strides (elements)
+        i32_s_st: fx.Int32,
+        i64_s_bytes: fx.Int64,
+        i32_bt_stride: fx.Int32,
+        i64_bt_bytes: fx.Int64,
+        i32_num_reqs: fx.Int32,
+        i32_chunk: fx.Int32,  # requests per chunk
+        i32_nch: fx.Int32,
+        i32_splits: fx.Int32,  # workgroups per chunk
+        i32_init: fx.Int32,
+        i32_local: fx.Int32,
+        i32_per_wave: fx.Int32,  # 1: each wave walks its own positions
+        i32_unit: fx.Int32,  # consecutive positions per unit, 0: none
+    ):
+        smem = fx.SharedAllocator().allocate(Shared).peek()
+        tx, pid = fx.thread_idx.x, fx.block_idx.x
+        lane = tx % 64
+        wave = fx.Int32(fx.rocdl.readfirstlane(T.i32, tx // 64))
+        col, kl = lane % 16, lane // 16
+        ch = pid % i32_nch
+        split = pid // i32_nch
+        req0 = ch * i32_chunk
+        left = i32_num_reqs - req0
+        nr = (left < i32_chunk).select(left, i32_chunk)
+
+        # Lane r holds request req0 + r.
+        r_ok = lane < nr
+        seq_r = fx.Int32(global_i32_ptr(arg_seq)[r_ok.select(req0 + lane, fx.Int32(0))])
+        seq_r = r_ok.select(seq_r, fx.Int32(0))
+        nblk_r = (seq_r + (BLK - 1)) // BLK
+        first_row = seq_r - (QL - 1)
+        # Every row sees blocks before rbf_r whole: init / local ones there only
+        # need their sentinel.
+        rbf_r = ((first_row < 0).select(fx.Int32(0), first_row) + (BLK - 1)) // BLK
+        ll = nblk_r - i32_local
+        loc_last_r = (ll < 0).select(fx.Int32(0), ll)
+        fm_r = seq_r - QL  # tokens past this are masked for some row
+
+        btr = buffer_ops.create_buffer_resource_from_addr(
+            arg_bt, num_records_bytes=i64_bt_bytes
+        )
+        sr = buffer_ops.create_buffer_resource_from_addr(
+            arg_score, num_records_bytes=i64_s_bytes
+        )
+        qr = buffer_ops.create_buffer_resource_from_addr(
+            arg_q, num_records_bytes=i64_q_bytes
+        )
+
+        q16 = fx.logical_divide(
+            fx.make_view(
+                fx.recast_iter(fx.Int32, smem.q.ptr),
+                fx.make_layout(QROWS * QROW_BYTES // 4, 1),
+            ),
+            fx.make_layout(4, 1),
+        )
+        atom16 = fx.make_copy_atom(fx.UniversalCopy128b(), fx.Int32)
+        seq_l = smem.seq.ptr
+        scb = smem.sbuf.ptr
+        per_wave = i32_per_wave != fx.Int32(0)
+        if const_expr(buffered):
+            direct = per_wave | (i32_unit == fx.Int32(0))
+            to_lds = (i32_per_wave == fx.Int32(0)) & (i32_unit != fx.Int32(0))
+            usz = (i32_unit == fx.Int32(0)).select(fx.Int32(1), i32_unit)
+        else:
+            direct = i32_unit >= fx.Int32(0)
+            to_lds = i32_unit < fx.Int32(0)
+            usz = fx.Int32(1)
+
+        def prologue():
+            """Q rows, the e4m3 table and the request lengths -> LDS.
+
+            Row i = request i // R, column j = i % R (query row j // H, head
+            j % H); 16 B column c8 sits at c8 ^ (i & 7), in the key's dim order.
+            """
+            for k in range_constexpr(QROWS * 8 // (64 * NW)):
+                u = tx + k * 64 * NW
+                i = u // 8
+                c8 = u % 8
+                rr = i // R
+                j = i % R
+                # Missing requests' and padding rows stay unwritten; their
+                # scores are never stored.
+                if (rr < nr) & (j // H < QL):
+                    tq = (req0 + rr) * QL + j // H
+                    off = (tq * i32_q_st + (j % H) * i32_q_sh + c8 * 16) // 4
+                    d = Vec(
+                        buffer_ops.buffer_load(
+                            qr, off, vec_width=4, dtype=fx.Int32, mask=rr < nr
+                        )
+                    )
+                    d = [fx.Int32(d[e]) for e in range_constexpr(4)]
+                    r = fx.make_rmem_tensor(fx.make_layout(4, 1), fx.Int32)
+                    r.store(key_order4(d))
+                    fx.copy(atom16, r, fx.slice(q16, (None, i * 8 + (c8 ^ (i & 7)))))
+            lut_g = global_i32_ptr(arg_lut)
+            lut_l = smem.lut.ptr
+            for e in range_constexpr(2):
+                lut_l[tx * 2 + e] = fx.Int32(lut_g[tx * 2 + e])
+            seq_l[lane] = seq_r
+
+        prologue()
+        fx.gpu.barrier()
+
+        lut8 = fx.logical_divide(
+            fx.make_view(
+                fx.recast_iter(fx.Int32, smem.lut.ptr), fx.make_layout(512, 1)
+            ),
+            fx.make_layout(2, 1),
+        )
+        lut_atom = fx.make_copy_atom(fx.UniversalCopy64b(), fx.Int32)
+
+        def lut_load(sbyte):
+            r = fx.make_rmem_tensor(fx.make_layout(2, 1), fx.Int32)
+            fx.copy(lut_atom, fx.slice(lut8, (None, sbyte)), r)
+            v = Vec(r.load())
+            return fx.Int32(v[0]), fx.Int32(v[1])
+
+        def q_frag(cg):
+            """B operand of column group cg (dims 32 kl .. +31 of Q row cg * 16 +
+            col)."""
+            i = cg * 16 + col
+            halves = []
+            for h in range_constexpr(2):
+                r = fx.make_rmem_tensor(fx.make_layout(4, 1), fx.Int32)
+                fx.copy(
+                    atom16, fx.slice(q16, (None, i * 8 + ((2 * kl + h) ^ (i & 7)))), r
+                )
+                halves.append(r.load())
+            return pack8(halves[0], halves[1])
+
+        code_off = (col * 64 + kl * 16) // 4  # dwords
+        scale_off = col * 4 + kl  # u16: groups 2 kl, 2 kl + 1 of token col
+
+        def page_rsrc(page):
+            return buffer_ops.create_buffer_resource_from_addr(
+                arg_kv + fx.Int64(page) * i64_page_stride, num_records_bytes=PAGE_BYTES
+            )
+
+        def load_tile(kr, ok, t):
+            """Tile t of a page: this lane's 16 B of codes (token 16 t + col) and
+            its two group scales; no memory access when not ``ok``."""
+            c = Vec(
+                buffer_ops.buffer_load(
+                    kr,
+                    code_off,
+                    vec_width=4,
+                    dtype=fx.Int32,
+                    mask=ok,
+                    soffset_bytes=t * 1024,
+                    cache_modifier=NT,
+                )
+            )
+            sc = fx.Int32(
+                fx.Int16(
+                    buffer_ops.buffer_load(
+                        kr,
+                        scale_off,
+                        vec_width=1,
+                        dtype=fx.Int16,
+                        mask=ok,
+                        soffset_bytes=PAGE_CODES + t * 128,
+                        cache_modifier=NT,
+                    )
+                ).to(fx.Int32)
+            )
+            return c, sc
+
+        def dequant_tile(c, sc):
+            """The A fragment of a tile (e4m3 keys in the key's dim order)."""
+            t0 = lut_load(sc & 0xFF)
+            t1 = lut_load((sc >> 8) & 0xFF)
+            f = []
+            for e in range_constexpr(4):
+                lo, hi = dequant8(fx.Int32(Vec(c)[e]), *(t0 if e < 2 else t1))
+                f += [lo, hi]
+            return Vec.from_elements(f, fx.Int32)
+
+        zero4 = fx.Vector.filled(4, 0.0, fx.Float32)
+        neg_inf = fx.Float32(float("-inf"))
+
+        def frag_max(frags, qfs, masked, lims):
+            """Per column group: the max of its column's scores over the page.
+
+            Tile t's results are reduced under tile t + 1's MFMAs.
+            """
+            ng = len(qfs)
+            runs = [neg_inf] * ng
+            accs = [[None] * ng for _ in range(MT)]
+
+            def reduce(t, runs):
+                out = []
+                for g in range_constexpr(ng):
+                    run = runs[g]
+                    v = accs[t][g]
+                    for e in range_constexpr(4):
+                        x = v[e]
+                        if const_expr(masked):
+                            x = (lims[g] >= t * 16 + e).select(x, neg_inf)
+                        run = maxf_nn(run, x)
+                    out.append(run)
+                return out
+
+            for t in range_constexpr(MT):
+                for g in range_constexpr(ng):
+                    accs[t][g] = mfma_fp8_16x16x128(frags[t], qfs[g], zero4)
+                if const_expr(t >= 1):
+                    runs = reduce(t - 1, runs)
+            runs = reduce(MT - 1, runs)
+            return [xlane4(r, maxf_nn) for r in runs]
+
+        def row_value(mx, seq, q_off, b):
+            return block_sentinels(mx, seq - QL + q_off + 1, b, i32_init, i32_local)
+
+        def put(i_row, jpos, v, cond, addr):
+            """Score v of Q row i_row at unit position jpos, to memory or to the
+            LDS buffer; lanes with nothing to buffer write a dummy slot."""
+            if const_expr(not buffered):
+                buffer_ops.buffer_store(v, sr, addr, mask=cond)
+            else:
+                buffer_ops.buffer_store(v, sr, addr, mask=cond & direct)
+                idx = (cond & to_lds).select(
+                    i_row * SB_STRIDE + jpos, QROWS * SB_STRIDE + lane
+                )
+                (scb + idx).store(bits(v))  # `scb[idx] = v` would rebind scb
+
+        def flush(ub):
+            """Store the unit at block ub: 4 consecutive blocks per thread and
+            pass."""
+            NCH = UMAX // 4
+            k = fx.Int32(0)
+            while k < fx.Int32(QROWS * NCH // (64 * NW)):
+                item = tx + k * (64 * NW)
+                i = item // NCH
+                c = item % NCH
+                rr = i // R
+                jc = i % R
+                qoff = jc // H
+                nb_c = (fx.Int32(seq_l[rr]) + (BLK - 1)) // BLK
+                b0c = ub + c * 4
+                ok = (rr < nr) & (qoff < QL) & (c * 4 < i32_unit)
+                addr = (jc % H) * i32_s_sh + ((req0 + rr) * QL + qoff) * i32_s_st + b0c
+                for e in range_constexpr(4):
+                    v = as_f32(fx.Int32(scb[i * SB_STRIDE + c * 4 + e]))
+                    buffer_ops.buffer_store(
+                        v,
+                        sr,
+                        addr + e,
+                        mask=ok & (b0c + e < nb_c) & (c * 4 + e < i32_unit),
+                    )
+                k = k + 1
+
+        def score_cgs(frags, hm, masked_any, cgs, b, jpos):
+            """Scores of column groups cgs, masked where a request is at its last
+            blocks, with sentinels."""
+            info = []
+            for cg in cgs:
+                i = cg * 16 + col
+                crr = i // R
+                st_j = i % R
+                info.append((cg, crr, st_j // H, st_j % H))
+            qfs = [q_frag(cg) for cg, _, _, _ in info]
+            seqs = [fx.Int32(seq_l[crr]) for _, crr, _, _ in info]
+            mxs = [neg_inf for _ in info]
+            if masked_any:
+                lims = [
+                    seq_c - QL + qoff - b * BLK - kl * 4
+                    for seq_c, (_, _, qoff, _) in zip(seqs, info)
+                ]
+                mxs = frag_max(frags, qfs, True, lims)
+            else:
+                mxs = frag_max(frags, qfs, False, None)
+            for mx, seq_c, (cg, crr, qoff, head) in zip(mxs, seqs, info):
+                hit = ((hm >> fx.Int64(crr)) & fx.Int64(1)) != fx.Int64(0)
+                ov = row_value(mx, seq_c, qoff, b)
+                otq = (req0 + crr) * QL + qoff
+                put(
+                    cg * 16 + col,
+                    jpos,
+                    ov,
+                    hit & (kl == 0) & (qoff < QL),
+                    head * i32_s_sh + otq * i32_s_st + b,
+                )
+
+        def score_groups(frags, hm, mmask, b, jpos):
+            """Every column group holding the page (requests hm)."""
+            left_m = hm
+            while left_m != fx.Int64(0):
+                lead_req = fx.Int32(fx.math.cttz(left_m).to(fx.Int32))
+                if const_expr(CPR == 1):
+                    cg = lead_req // RPC
+                    gbits = fx.Int64(GMASK) << fx.Int64(cg * RPC)
+                    left_m = left_m & ~gbits
+                    score_cgs(frags, hm, (mmask & gbits) != fx.Int64(0), [cg], b, jpos)
+                else:
+                    gbits = fx.Int64(1) << fx.Int64(lead_req)
+                    left_m = left_m & ~gbits
+                    masked_any = (mmask & gbits) != fx.Int64(0)
+                    score_cgs(
+                        frags,
+                        hm,
+                        masked_any,
+                        [lead_req * CPR + g for g in range(CPR)],
+                        b,
+                        jpos,
+                    )
+
+        # The waves share positions in units of i32_unit blocks; page k of
+        # position `it` goes to wave (k + it) % NW. In per-wave mode (fewer
+        # requests than waves) wave w walks its own units instead.
+        u_step = per_wave.select(i32_splits * NW, i32_splits)
+        ub = per_wave.select(split * NW + wave, split) * usz
+        it = fx.Int32(0)
+        while fx.Int64(fx.rocdl.ballot(T.i64, ub < nblk_r)) != fx.Int64(0):
+            b = ub
+            jpos = fx.Int32(0)
+            while (jpos < usz) & (
+                fx.Int64(fx.rocdl.ballot(T.i64, b < nblk_r)) != fx.Int64(0)
+            ):
+                in_r = b < nblk_r
+                skip = (b < rbf_r) & ((b < i32_init) | (b >= loc_last_r))
+                need_r = in_r & ~skip
+                bid = fx.Int32(
+                    buffer_ops.buffer_load(
+                        btr,
+                        need_r.select((req0 + lane) * i32_bt_stride + b, fx.Int32(0)),
+                        vec_width=1,
+                        dtype=fx.Int32,
+                    )
+                )
+                bid = need_r.select(bid, fx.Int32(-1))
+                # Sentinel-only blocks, stored by wave 0.
+                if (fx.Int64(fx.rocdl.ballot(T.i64, in_r & skip)) != fx.Int64(0)) & (
+                    (wave == 0) | per_wave
+                ):
+                    # A runtime loop: unrolled, all store addresses stay live.
+                    jr = fx.Int32(0)
+                    while jr < fx.Int32(QL * H):
+                        sv = row_value(neg_inf, seq_r, jr // H, b)
+                        stq = (req0 + lane) * QL + jr // H
+                        put(
+                            lane * R + jr,
+                            jpos,
+                            sv,
+                            in_r & skip,
+                            (jr % H) * i32_s_sh + stq * i32_s_st + b,
+                        )
+                        jr = jr + 1
+                pend = fx.Int64(fx.rocdl.ballot(T.i64, need_r))
+                mmask = fx.Int64(
+                    fx.rocdl.ballot(T.i64, need_r & (b * BLK + (BLK - 1) > fm_r))
+                )
+                # This wave's pages, as the mask of their first lanes.
+                mine = fx.Int64(0)
+                k = it
+                while pend != fx.Int64(0):
+                    lead = fx.Int32(fx.math.cttz(pend).to(fx.Int32))
+                    page = fx.Int32(fx.rocdl.readlane(T.i32, bid, lead))
+                    pend = pend & ~fx.Int64(fx.rocdl.ballot(T.i64, bid == page))
+                    mine = mine | (((k % NW) == wave) | per_wave).select(
+                        fx.Int64(1) << fx.Int64(lead), fx.Int64(0)
+                    )
+                    k = k + 1
+                ok0 = mine != fx.Int64(0)
+                l0 = fx.Int32(fx.math.cttz(mine).to(fx.Int32))
+                kr0 = page_rsrc(
+                    fx.Int32(fx.rocdl.readlane(T.i32, bid, ok0.select(l0, fx.Int32(0))))
+                )
+                tiles = [load_tile(kr0, ok0, t) for t in range_constexpr(MT)]
+                while mine != fx.Int64(0):
+                    lead = fx.Int32(fx.math.cttz(mine).to(fx.Int32))
+                    page = fx.Int32(fx.rocdl.readlane(T.i32, bid, lead))
+                    hm = fx.Int64(fx.rocdl.ballot(T.i64, bid == page))
+                    mine = mine & ~(fx.Int64(1) << fx.Int64(lead))
+                    okn = mine != fx.Int64(0)
+                    ln = fx.Int32(fx.math.cttz(mine).to(fx.Int32))
+                    krn = page_rsrc(
+                        fx.Int32(
+                            fx.rocdl.readlane(T.i32, bid, okn.select(ln, fx.Int32(0)))
+                        )
+                    )
+                    # Tile t of the next page loads into the registers tile t of
+                    # this page frees; the barriers keep LLVM from hoisting the
+                    # loads.
+                    frags, nxt = [], []
+                    for t in range_constexpr(MT):
+                        frags.append(dequant_tile(*tiles[t]))
+                        nxt.append(load_tile(krn, okn, t))
+                        fx.rocdl.sched_barrier(0)
+                    tiles = nxt
+                    score_groups(frags, hm, mmask, b, jpos)
+                b = b + 1
+                jpos = jpos + 1
+                it = it + 1
+            if const_expr(buffered):  # noqa: SIM102 (the outer if is static)
+                if to_lds:
+                    fx.gpu.barrier()
+                    flush(ub)
+                    fx.gpu.barrier()
+            ub = ub + u_step * usz
+
+    @flyc.jit
+    def launch(
+        arg_q: fx.Int64,
+        arg_kv: fx.Int64,
+        arg_lut: fx.Int64,
+        arg_score: fx.Int64,
+        arg_bt: fx.Int64,
+        arg_seq: fx.Int64,
+        i64_page_stride: fx.Int64,
+        i32_q_st: fx.Int32,
+        i32_q_sh: fx.Int32,
+        i64_q_bytes: fx.Int64,
+        i32_s_sh: fx.Int32,
+        i32_s_st: fx.Int32,
+        i64_s_bytes: fx.Int64,
+        i32_bt_stride: fx.Int32,
+        i64_bt_bytes: fx.Int64,
+        i32_num_reqs: fx.Int32,
+        i32_chunk: fx.Int32,
+        i32_nch: fx.Int32,
+        i32_splits: fx.Int32,
+        i32_init: fx.Int32,
+        i32_local: fx.Int32,
+        i32_per_wave: fx.Int32,
+        i32_unit: fx.Int32,
+        stream: fx.Stream,
+    ):
+        kernel(
+            arg_q,
+            arg_kv,
+            arg_lut,
+            arg_score,
+            arg_bt,
+            arg_seq,
+            i64_page_stride,
+            i32_q_st,
+            i32_q_sh,
+            i64_q_bytes,
+            i32_s_sh,
+            i32_s_st,
+            i64_s_bytes,
+            i32_bt_stride,
+            i64_bt_bytes,
+            i32_num_reqs,
+            i32_chunk,
+            i32_nch,
+            i32_splits,
+            i32_init,
+            i32_local,
+            i32_per_wave,
+            i32_unit,
+        ).launch(
+            grid=(fx.Int64(i32_nch * i32_splits), 1, 1),
+            block=(64 * NW, 1, 1),
+            stream=stream,
+        )
+
+    launch.kernel_name = name
+    return launch
diff --git a/vllm/models/minimax_m3/amd/ops/nvfp4_index_decode/kernel_cp_push.py b/vllm/models/minimax_m3/amd/ops/nvfp4_index_decode/kernel_cp_push.py
new file mode 100644
index 000000000000..df1e15e8b5db
--- /dev/null
+++ b/vllm/models/minimax_m3/amd/ops/nvfp4_index_decode/kernel_cp_push.py
@@ -0,0 +1,613 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Context-parallel decode index scorer for MiniMax-M3 (FlyDSL, gfx950, TP4).
+
+Rank ``me`` scores blocks b with (b // CP_CH) % 4 == me for all four index heads
+(bit for bit as ``kernel.py``) and stores head h's scores into rank h's buffer
+(``index_cp.IndexCP``); it returns once every rank's scores have arrived.
+"""
+
+import flydsl.compiler as flyc
+import flydsl.expr as fx
+from aiter.ops.flydsl.kernels import buffer_ops
+from flydsl._mlir.dialects import llvm as _llvm
+from flydsl.expr import const_expr, range_constexpr
+from flydsl.expr.typing import T
+from flydsl.expr.typing import Vector as Vec
+
+from vllm.models.minimax_m3.amd.ops.flydsl_utils import (
+    NT,
+    SC0,
+    SC1,
+    bits,
+    dequant8,
+    global_i32_ptr,
+    key_order4,
+    maxf_nn,
+    mfma_fp8_16x16x128,
+    pack8,
+    perm,
+    xlane4,
+)
+from vllm.models.minimax_m3.amd.ops.nvfp4_index_decode.kernel import block_sentinels
+
+WORLD = 4  # TP ranks = index heads, one per rank
+BLK = 128  # tokens per block
+D = 128
+CP_CH = 16  # blocks per run of one rank
+PAGE_CODES = BLK * D // 2  # e2m1 codes, then [128 x 8 B] e4m3 group scales
+PAGE_BYTES = PAGE_CODES + BLK * 8
+MT = BLK // 16  # token tiles per page
+SYS = SC0 | SC1
+MAX_ROWS = 1024  # decode rows (max_num_seqs x query_len) the receive buffers hold
+# Flag words per rank: [SC_FLAGS + p] source p's scores arrived, [SC_DONE + w]
+# workgroup w done, [SC_IQ + p * SC_MAX_REQS + r] source p's index-q of request r.
+SC_FLAGS = 64
+SC_DONE = 1024
+SC_MAX_WG = 1024
+SC_IQ = 2048
+SC_MAX_REQS = 256
+FLAG_BYTES = 12288
+# Device counters (int32); [CNT_SC] holds the call count.
+CNT_SC = 128
+CNT_WORDS = 512
+
+
+def _flag_read(addr):
+    """An atomic add of 0 at system scope (never hoisted out of a spin loop)."""
+    return fx.Int32(
+        _llvm.AtomicRMWOp(
+            _llvm.AtomicBinOp.add,
+            global_i32_ptr(addr).llvm_ptr,
+            fx.Int32(0).ir_value(),
+            _llvm.AtomicOrdering.monotonic,
+            syncscope="",
+            alignment=4,
+        ).result
+    )
+
+
+def _sel4(p, a0, a1, a2, a3):
+    return (p == 0).select(a0, (p == 1).select(a1, (p == 2).select(a2, a3)))
+
+
+GRID_SMALL, GRID = 512, 768
+GRID_SMALL_REQS = 32
+NW = 4  # waves per workgroup: quarter w of every run
+QD = CP_CH // NW  # blocks per quarter
+WIN = 64  # pages per block-table window (one VGPR)
+MAX_REQS = 64 * NW  # one request per thread in the prefix sum
+QL = 4  # DSpark rows per request
+HEADS = 4
+
+
+def _uni(x):
+    return fx.Int32(fx.rocdl.readfirstlane(T.i32, fx.Int32(x)))
+
+
+def _min(a, b):
+    return (a < b).select(a, b)
+
+
+def _max(a, b):
+    return (a > b).select(a, b)
+
+
+def _wave_scan(x, lane):
+    """Inclusive prefix sum over the wave's 64 lanes."""
+    for d in range_constexpr(6):
+        y = fx.Int32(fx.rocdl.ds_bpermute(T.i32, (lane - (1 << d)) * 4, x))
+        x = x + (lane >= (1 << d)).select(y, fx.Int32(0))
+    return x
+
+
+def compile_score_push():
+    name = "m3_nvfp4_index_score_push_cp"
+
+    @fx.struct
+    class Shared:
+        lut: fx.Array[fx.Int32, 512, 16]
+        stage: fx.Array[fx.Int32, 2 * 16 * NW * QD, 16]  # [run % 2][col][quarter]
+        incl: fx.Array[fx.Int32, MAX_REQS]  # inclusive prefix of the run counts
+        seq: fx.Array[fx.Int32, MAX_REQS]
+        wsum: fx.Array[fx.Int32, NW]
+
+    @flyc.kernel(name=name, known_block_size=[64 * NW, 1, 1])
+    def kernel(
+        arg_iq: fx.Int64,  # this rank's head: [num_reqs * QL, 128] e4m3
+        arg_kv: fx.Int64,  # NVFP4 index pages [num_blocks, 128, 72] u8
+        arg_lut: fx.Int64,  # nvfp4_e4m3_lut [256] i64
+        arg_d0: fx.Int64,  # head h's scores [rows, S] fp32 (rank h's receive buffer)
+        arg_d1: fx.Int64,
+        arg_d2: fx.Int64,
+        arg_d3: fx.Int64,
+        arg_bt: fx.Int64,  # block table [num_reqs, >= max_blocks] i32
+        arg_seq: fx.Int64,  # seq_lens [num_reqs] i32
+        arg_cnt: fx.Int64,  # device counters (CNT_*)
+        arg_f0: fx.Int64,  # rank p's flag words
+        arg_f1: fx.Int64,
+        arg_f2: fx.Int64,
+        arg_f3: fx.Int64,
+        arg_q0: fx.Int64,  # rank p's index-q receive buffer [4, rmax, 128] e4m3
+        arg_q1: fx.Int64,
+        arg_q2: fx.Int64,
+        arg_q3: fx.Int64,
+        i32_rmax: fx.Int32,
+        i64_page_stride: fx.Int64,
+        i32_s: fx.Int32,  # score row stride (elements)
+        i64_d_bytes: fx.Int64,
+        i32_bt_stride: fx.Int32,
+        i64_bt_bytes: fx.Int64,
+        i32_num_reqs: fx.Int32,
+        i32_init: fx.Int32,  # leading blocks forced in (sentinel 1e30)
+        i32_local: fx.Int32,  # trailing blocks forced in (sentinel 1e29)
+        i32_me: fx.Int32,
+        i32_nwg: fx.Int32,  # the grid
+    ):
+        smem = fx.SharedAllocator().allocate(Shared).peek()
+        tx, pid = fx.thread_idx.x, fx.block_idx.x
+        lane = tx % 64
+        wave = _uni(tx // 64)
+        col, kl = lane % 16, lane // 16
+        st_qoff = col // HEADS  # this lane's column: row st_qoff of head st_head
+        st_head = col % HEADS
+        incl_l, seq_l = smem.incl.ptr, smem.seq.ptr
+
+        cr = buffer_ops.create_buffer_resource_from_addr(
+            arg_cnt, num_records_bytes=CNT_WORDS * 4
+        )
+        calls = (
+            _uni(
+                buffer_ops.buffer_load(
+                    cr, CNT_SC, vec_width=1, dtype=fx.Int32, cache_modifier=SC1
+                )
+            )
+            + 1
+        )
+        btr = buffer_ops.create_buffer_resource_from_addr(
+            arg_bt, num_records_bytes=i64_bt_bytes
+        )
+        qbytes = fx.Int64(i32_rmax) * (HEADS * D)
+        qs = (arg_q0, arg_q1, arg_q2, arg_q3)
+        fs = (arg_f0, arg_f1, arg_f2, arg_f3)
+        mine = _sel4(i32_me, arg_f0, arg_f1, arg_f2, arg_f3)
+        qr = buffer_ops.create_buffer_resource_from_addr(
+            _sel4(i32_me, arg_q0, arg_q1, arg_q2, arg_q3), num_records_bytes=qbytes
+        )
+        drs = [
+            buffer_ops.create_buffer_resource_from_addr(
+                d, num_records_bytes=i64_d_bytes
+            )
+            for d in (arg_d0, arg_d1, arg_d2, arg_d3)
+        ]
+
+        # LUT to LDS; per request this rank's run count, prefix-summed.
+        lut_g = global_i32_ptr(arg_lut)
+        lut_l = smem.lut.ptr
+        for i in range_constexpr(2):
+            lut_l[tx * 2 + i] = fx.Int32(lut_g[tx * 2 + i])
+        r_ok = tx < i32_num_reqs
+        seq_t = fx.Int32(global_i32_ptr(arg_seq)[r_ok.select(tx, fx.Int32(0))])
+        seq_t = r_ok.select(seq_t, fx.Int32(0))
+        # This rank's index-q rows of requests pid, pid + nwg, ... to slot me of
+        # every rank; their flags follow below.
+        ir = buffer_ops.create_buffer_resource_from_addr(
+            arg_iq, num_records_bytes=fx.Int64(i32_num_reqs) * (QL * D)
+        )
+        rs = pid
+        while (wave == 0) & (rs < i32_num_reqs):
+            v = buffer_ops.buffer_load(
+                ir,
+                (rs * QL * D + lane * 16) // 4,
+                vec_width=4,
+                dtype=fx.Int32,
+                mask=lane < QL * D // 16,
+            )
+            for p in range_constexpr(WORLD):
+                buffer_ops.buffer_store(
+                    v,
+                    buffer_ops.create_buffer_resource_from_addr(
+                        qs[p], num_records_bytes=qbytes
+                    ),
+                    ((i32_me * i32_rmax + rs * QL) * D + lane * 16) // 4,
+                    mask=lane < QL * D // 16,
+                    cache_modifier=SYS,
+                )
+            rs = rs + i32_nwg
+        nb_t = (seq_t + (BLK - 1)) // BLK
+        runs_t = (nb_t // (WORLD * CP_CH)) + (
+            nb_t % (WORLD * CP_CH) > i32_me * CP_CH
+        ).select(fx.Int32(1), fx.Int32(0))
+        x = _wave_scan(runs_t, lane)
+        if lane == 63:
+            smem.wsum.ptr[wave] = x
+        seq_l[tx] = seq_t
+        fx.gpu.barrier()
+        for w in range_constexpr(NW - 1):
+            x = x + (wave > w).select(fx.Int32(smem.wsum.ptr[w]), fx.Int32(0))
+        incl_l[tx] = x
+        fx.gpu.barrier()
+
+        # Runs pid, pid + grid, ...: the busy workgroups stay one per CU.
+        total = _uni(incl_l[MAX_REQS - 1])
+        nruns = (pid < total).select((total - pid + i32_nwg - 1) // i32_nwg, 0)
+        nslots = nruns * QD  # pages of this wave
+
+        lut8 = fx.logical_divide(
+            fx.make_view(
+                fx.recast_iter(fx.Int32, smem.lut.ptr), fx.make_layout(512, 1)
+            ),
+            fx.make_layout(2, 1),
+        )
+        lut_atom = fx.make_copy_atom(fx.UniversalCopy64b(), fx.Int32)
+        stage16 = fx.logical_divide(
+            fx.make_view(
+                fx.recast_iter(fx.Int32, smem.stage.ptr),
+                fx.make_layout(2 * 16 * NW * QD, 1),
+            ),
+            fx.make_layout(4, 1),
+        )
+        atom16 = fx.make_copy_atom(fx.UniversalCopy128b(), fx.Int32)
+
+        def lut_load(sbyte):
+            r = fx.make_rmem_tensor(fx.make_layout(2, 1), fx.Int32)
+            fx.copy(lut_atom, fx.slice(lut8, (None, sbyte)), r)
+            v = Vec(r.load())
+            return fx.Int32(v[0]), fx.Int32(v[1])
+
+        code_off = (col * 64 + kl * 16) // 4  # dwords
+        # Scales 2 kl, 2 kl + 1 of token col as a dword: a u16 load's sign
+        # extension made the loop head wait for the whole next page.
+        scale_off = col * 2 + kl // 2  # dwords
+        # v_perm selectors: scale byte 2 (kl % 2) (+ 1) of the dword, zeros above
+        scale_sel = (kl % 2) * 2 + 0x0C0C0C00
+
+        def page_rsrc(page):
+            return buffer_ops.create_buffer_resource_from_addr(
+                arg_kv + fx.Int64(page) * i64_page_stride, num_records_bytes=PAGE_BYTES
+            )
+
+        def load_tile(kr, ok, t):
+            """Tile t: this lane's codes and scales; no memory access unless ok."""
+            c = Vec(
+                buffer_ops.buffer_load(
+                    kr,
+                    code_off,
+                    vec_width=4,
+                    dtype=fx.Int32,
+                    mask=ok,
+                    soffset_bytes=t * 1024,
+                    cache_modifier=NT,
+                )
+            )
+            sc = fx.Int32(
+                buffer_ops.buffer_load(
+                    kr,
+                    scale_off,
+                    vec_width=1,
+                    dtype=fx.Int32,
+                    mask=ok,
+                    soffset_bytes=PAGE_CODES + t * 128,
+                    cache_modifier=NT,
+                )
+            )
+            return c, sc
+
+        def dequant_tile(c, sc):
+            """The A fragment of a tile (e4m3 keys, even dims then odd dims per 8)"""
+            t0 = lut_load(perm(sc, sc, scale_sel))
+            t1 = lut_load(perm(sc, sc, scale_sel + 1))
+            f = []
+            for e in range_constexpr(4):
+                lo, hi = dequant8(fx.Int32(Vec(c)[e]), *(t0 if e < 2 else t1))
+                f += [lo, hi]
+            return Vec.from_elements(f, fx.Int32)
+
+        def q_load(r, ok):
+            """This lane's 32 B of Q: row r * QL + st_qoff, head st_head."""
+            off = ((st_head * i32_rmax + r * QL + st_qoff) * D + kl * 32) // 4
+            a = Vec(
+                buffer_ops.buffer_load(qr, off, vec_width=4, dtype=fx.Int32, mask=ok)
+            )
+            b = Vec(
+                buffer_ops.buffer_load(
+                    qr, off + 4, vec_width=4, dtype=fx.Int32, mask=ok
+                )
+            )
+            return a, b
+
+        def q_frag(a, b):
+            """B operand: the 32 B in the key's dim order."""
+            halves = []
+            for v in (a, b):
+                d = [fx.Int32(Vec(v)[e]) for e in range_constexpr(4)]
+                halves.append(key_order4(d))
+            return pack8(halves[0], halves[1])
+
+        zero4 = fx.Vector.filled(4, 0.0, fx.Float32)
+        neg_inf = fx.Float32(float("-inf"))
+
+        def reduce(v, t, run, masked, lim):
+            for e in range_constexpr(4):
+                x = v[e]
+                if const_expr(masked):
+                    x = (lim >= t * 16 + e).select(x, neg_inf)
+                run = maxf_nn(run, x)
+            return run
+
+        def page_pass(tiles, qf, krn, okn, masked, lim):
+            """Max of this lane's column over the page; tile t of the next page
+            loads into tile t's registers (sched barriers stop LLVM hoisting)."""
+            accs = []
+            nxt = []
+            for t in range_constexpr(MT):
+                f = dequant_tile(*tiles[t])
+                nxt.append(load_tile(krn, okn, t))
+                accs.append(mfma_fp8_16x16x128(f, qf, zero4))
+                fx.rocdl.sched_barrier(0)
+            mx = neg_inf
+            if masked:
+                run_m = neg_inf
+                for t in range_constexpr(MT):
+                    run_m = reduce(accs[t], t, run_m, True, lim)
+                mx = xlane4(run_m, maxf_nn)
+            else:
+                run_u = neg_inf
+                for t in range_constexpr(MT):
+                    run_u = reduce(accs[t], t, run_u, False, None)
+                mx = xlane4(run_u, maxf_nn)
+            return mx, nxt
+
+        def row_value(mx, seq, q_off, b):
+            return block_sentinels(mx, seq - QL + q_off + 1, b, i32_init, i32_local)
+
+        fx.rocdl.s_waitcnt(vmcnt=0)
+        rs = pid
+        while (wave == 0) & (rs < i32_num_reqs):
+            for p in range_constexpr(WORLD):
+                buffer_ops.buffer_store(
+                    calls,
+                    buffer_ops.create_buffer_resource_from_addr(
+                        fs[p], num_records_bytes=FLAG_BYTES
+                    ),
+                    SC_IQ + i32_me * SC_MAX_REQS + rs,
+                    mask=lane == 0,
+                    cache_modifier=SYS,
+                )
+            rs = rs + i32_nwg
+
+        w0 = fx.Int32(0)
+        while w0 < nslots:
+            # One page per lane: page s is block 4 w + s % 4 of run pid + (s // 4)
+            # * grid, of request rj (binary search in the prefix).
+            s = w0 + lane
+            k = pid + (s // QD) * i32_nwg
+            rj = fx.Int32(0)
+            for bs in range_constexpr(8):
+                step = MAX_REQS >> (bs + 1)
+                rj = (fx.Int32(incl_l[rj + step - 1]) <= k).select(rj + step, rj)
+            valid = s < nslots
+            rj = valid.select(rj, fx.Int32(0))
+            excl = (rj > 0).select(fx.Int32(incl_l[_max(rj - 1, 0)]), fx.Int32(0))
+            seqj = fx.Int32(seq_l[rj])
+            nblkj = (seqj + (BLK - 1)) // BLK
+            bj = ((k - excl) * WORLD + i32_me) * CP_CH + wave * QD + s % QD
+            frow = seqj - (QL - 1)
+            rbf = ((frow < 0).select(fx.Int32(0), frow) + (BLK - 1)) // BLK
+            ll = nblkj - i32_local
+            loc_last = (ll < 0).select(fx.Int32(0), ll)
+            # Init / local blocks every row sees whole only need their sentinel.
+            skip = (bj < rbf) & ((bj < i32_init) | (bj >= loc_last))
+            need = valid & (bj < nblkj) & ~skip
+            pgj = fx.Int32(
+                buffer_ops.buffer_load(
+                    btr,
+                    need.select(rj * i32_bt_stride + bj, fx.Int32(0)),
+                    vec_width=1,
+                    dtype=fx.Int32,
+                )
+            )
+            pgj = need.select(pgj, fx.Int32(-1))
+            n = _min(nslots - w0, WIN)
+
+            p_0 = fx.Int32(fx.rocdl.readlane(T.i32, pgj, 0))
+            r_0 = fx.Int32(fx.rocdl.readlane(T.i32, rj, 0))
+            ok0 = p_0 >= 0
+            kr0 = page_rsrc(ok0.select(p_0, fx.Int32(0)))
+            fx.rocdl.sched_barrier(0)
+            tiles = []
+            for t in range_constexpr(MT):
+                tiles.append(load_tile(kr0, ok0, t))
+                fx.rocdl.sched_barrier(0)
+            # Wait for the index-q rows; lane i polls source i % 4 of its request.
+            word = mine + fx.Int64((SC_IQ + (lane % WORLD) * SC_MAX_REQS + rj) * 4)
+            want = lane < n
+            f = calls
+            if want:
+                f = _flag_read(word)
+            while fx.Int64(fx.rocdl.ballot(T.i64, want & (f < calls))) != 0:
+                if want & (f < calls):
+                    f = _flag_read(word)
+            qa, qb = q_load(r_0, n > 0)
+            # Nothing pending where the loop head merges with the back edge.
+            fx.rocdl.s_waitcnt(vmcnt=0)
+            qf = q_frag(qa, qb)
+            qna, qnb = qa, qb
+            sw = n < 0  # false: Q of the next page already loaded
+            vq = [neg_inf] * QD
+            j = fx.Int32(0)
+            while j < n:
+                r = fx.Int32(fx.rocdl.readlane(T.i32, rj, j))
+                b = fx.Int32(fx.rocdl.readlane(T.i32, bj, j))
+                seq = fx.Int32(fx.rocdl.readlane(T.i32, seqj, j))
+                if sw:
+                    qf = q_frag(qna, qnb)
+                jn = j + 1
+                okj = jn < n
+                jc = okj.select(jn, fx.Int32(0))
+                pgn = fx.Int32(fx.rocdl.readlane(T.i32, pgj, jc))
+                rn = fx.Int32(fx.rocdl.readlane(T.i32, rj, jc))
+                okn = okj & (pgn >= 0)
+                krn = page_rsrc(okn.select(pgn, fx.Int32(0)))
+                sw = okj & (rn != r)
+                qna, qnb = q_load(rn, sw)
+                fx.rocdl.sched_barrier(0)
+                mx, tiles = page_pass(
+                    tiles,
+                    qf,
+                    krn,
+                    okn,
+                    b * BLK + (BLK - 1) > seq - QL,
+                    seq - QL + st_qoff - b * BLK - kl * 4,
+                )
+                ov = row_value(mx, seq, st_qoff, b)
+                e = (w0 + j) % QD
+                vq = [(e == i).select(ov, v) for i, v in enumerate(vq)]
+                if e == QD - 1:
+                    # The run's quarters meet in LDS; wave m % NW stores it.
+                    m = (w0 + j) // QD
+                    slot = m % 2
+                    val = fx.make_rmem_tensor(fx.make_layout(4, 1), fx.Int32)
+                    val.store(Vec.from_elements([bits(v) for v in vq], fx.Int32))
+                    if kl == 0:
+                        fx.copy(
+                            atom16,
+                            val,
+                            fx.slice(stage16, (None, (slot * 16 + col) * NW + wave)),
+                        )
+                    fx.gpu.barrier()
+                    if wave == m % NW:
+                        got = fx.make_rmem_tensor(fx.make_layout(4, 1), fx.Int32)
+                        fx.copy(
+                            atom16,
+                            fx.slice(stage16, (None, (slot * 16 + col) * NW + kl)),
+                            got,
+                        )
+                        off = (r * QL + st_qoff) * i32_s + (b - b % CP_CH) + kl * QD
+                        for h in range_constexpr(HEADS):
+                            buffer_ops.buffer_store(
+                                got.load(),
+                                drs[h],
+                                off,
+                                mask=st_head == h,
+                                cache_modifier=SYS,
+                            )
+                j = jn
+            w0 = w0 + WIN
+
+        # sc0 sc1 stores are system-visible once vmcnt drains; a release fence
+        # per workgroup would write back L2. The last workgroup (dispatched last,
+        # so its spin cannot starve the others) collects the done words, signals
+        # every rank and waits for all four.
+        fx.rocdl.s_waitcnt(vmcnt=0)
+        fx.gpu.barrier()
+        own = buffer_ops.create_buffer_resource_from_addr(
+            mine, num_records_bytes=FLAG_BYTES
+        )
+        buffer_ops.buffer_store(
+            calls, own, SC_DONE + pid, mask=tx == 0, cache_modifier=SYS
+        )
+        if pid == i32_nwg - 1:
+            # Thread t polls done words t, t + 256, ...
+            def done_min():
+                lo = calls
+                for i in range_constexpr(SC_MAX_WG // (64 * NW)):
+                    w = tx + i * 64 * NW
+                    f = _flag_read(
+                        mine + fx.Int64((SC_DONE + _min(w, SC_MAX_WG - 1)) * 4)
+                    )
+                    lo = _min(lo, (w < i32_nwg).select(f, calls))
+                return lo
+
+            lo = done_min()
+            while fx.Int64(fx.rocdl.ballot(T.i64, lo < calls)) != 0:
+                lo = done_min()
+            fx.gpu.barrier()
+            _llvm.FenceOp(_llvm.AtomicOrdering.release, syncscope="")
+            for p in range_constexpr(WORLD):
+                fr = buffer_ops.create_buffer_resource_from_addr(
+                    fs[p], num_records_bytes=FLAG_BYTES
+                )
+                buffer_ops.buffer_store(
+                    calls, fr, SC_FLAGS + i32_me, mask=tx == 0, cache_modifier=SYS
+                )
+            buffer_ops.buffer_store(calls, cr, CNT_SC, mask=tx == 0, cache_modifier=SYS)
+            if tx < WORLD:
+                f = _flag_read(mine + fx.Int64((SC_FLAGS + tx) * 4))
+                while f < calls:
+                    f = _flag_read(mine + fx.Int64((SC_FLAGS + tx) * 4))
+            fx.gpu.barrier()
+            _llvm.FenceOp(_llvm.AtomicOrdering.acquire, syncscope="")
+
+    @flyc.jit
+    def launch(
+        arg_iq: fx.Int64,
+        arg_kv: fx.Int64,
+        arg_lut: fx.Int64,
+        arg_d0: fx.Int64,
+        arg_d1: fx.Int64,
+        arg_d2: fx.Int64,
+        arg_d3: fx.Int64,
+        arg_bt: fx.Int64,
+        arg_seq: fx.Int64,
+        arg_cnt: fx.Int64,
+        arg_f0: fx.Int64,
+        arg_f1: fx.Int64,
+        arg_f2: fx.Int64,
+        arg_f3: fx.Int64,
+        arg_q0: fx.Int64,
+        arg_q1: fx.Int64,
+        arg_q2: fx.Int64,
+        arg_q3: fx.Int64,
+        i32_rmax: fx.Int32,
+        i64_page_stride: fx.Int64,
+        i32_s: fx.Int32,
+        i64_d_bytes: fx.Int64,
+        i32_bt_stride: fx.Int32,
+        i64_bt_bytes: fx.Int64,
+        i32_num_reqs: fx.Int32,
+        i32_init: fx.Int32,
+        i32_local: fx.Int32,
+        i32_me: fx.Int32,
+        i32_nwg: fx.Int32,
+        stream: fx.Stream,
+    ):
+        kernel(
+            arg_iq,
+            arg_kv,
+            arg_lut,
+            arg_d0,
+            arg_d1,
+            arg_d2,
+            arg_d3,
+            arg_bt,
+            arg_seq,
+            arg_cnt,
+            arg_f0,
+            arg_f1,
+            arg_f2,
+            arg_f3,
+            arg_q0,
+            arg_q1,
+            arg_q2,
+            arg_q3,
+            i32_rmax,
+            i64_page_stride,
+            i32_s,
+            i64_d_bytes,
+            i32_bt_stride,
+            i64_bt_bytes,
+            i32_num_reqs,
+            i32_init,
+            i32_local,
+            i32_me,
+            i32_nwg,
+        ).launch(grid=(fx.Int64(i32_nwg), 1, 1), block=(64 * NW, 1, 1), stream=stream)
+
+    launch.kernel_name = name
+    return launch
+
+
+assert CNT_SC < CNT_WORDS and SC_DONE + SC_MAX_WG <= FLAG_BYTES // 4
+# Rows that fit the receive buffers also fit the request limits.
+assert MAX_ROWS == MAX_REQS * QL and MAX_REQS <= SC_MAX_REQS
diff --git a/vllm/models/minimax_m3/amd/ops/nvfp4_index_prefill/__init__.py b/vllm/models/minimax_m3/amd/ops/nvfp4_index_prefill/__init__.py
new file mode 100644
index 000000000000..208f01a7cb5e
--- /dev/null
+++ b/vllm/models/minimax_m3/amd/ops/nvfp4_index_prefill/__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/minimax_m3/amd/ops/nvfp4_index_prefill/host.py b/vllm/models/minimax_m3/amd/ops/nvfp4_index_prefill/host.py
new file mode 100644
index 000000000000..2530d9f87eb5
--- /dev/null
+++ b/vllm/models/minimax_m3/amd/ops/nvfp4_index_prefill/host.py
@@ -0,0 +1,152 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Launch wrapper of the FlyDSL NVFP4 prefill index scorers."""
+
+import functools
+import math
+
+import torch
+
+from vllm.models.minimax_m3.amd.ops.flydsl_utils import cu_count, run_compiled
+from vllm.utils.math_utils import cdiv
+
+from . import kernel as _kernel
+
+# Workgroups to aim for: more, shorter segments balance the tail of the causal
+# windows; a segment keeps >= GROUP blocks.
+TARGET_WORKGROUPS = 2048
+# Segments are added until the last round of workgroups is at least this full.
+MIN_ROUND_FILL = 0.9
+WORKGROUPS_PER_CU = 2
+MAX_SEGMENTS = 1024
+# Small work (<= 128 query rows, or below SMALL_WORK) takes ``kernel_shared``
+# with 128-row tiles, aiming at SMALL_TARGET workgroups below MID_WORK.
+SMALL_WORK = 4096
+MID_WORK = 16384
+SMALL_TARGET = 512
+# Tests: pin the kernel (True: kernel_shared) or its 32-row tiles per wave.
+PIN_SHARED: bool | None = None
+PIN_MQ: int | None = None
+
+_NVFP4_E4M3_LUT: dict[torch.device, torch.Tensor] = {}
+
+
+def nvfp4_e4m3_lut(device: torch.device) -> torch.Tensor:
+    """[256] int64: byte j of entry s is e4m3_rn_satfinite(m_j * e4m3(s)) for the
+    E2M1 magnitudes m = 0, 0.5, 1, 1.5, 2, 3, 4, 6."""
+    lut = _NVFP4_E4M3_LUT.get(device)
+    if lut is None:
+        # Device ops only, so a first call under graph capture is fine.
+        scales = torch.arange(256, dtype=torch.uint8, device=device)
+        scales = scales.view(torch.float8_e4m3fn).float()
+        code = torch.arange(8, device=device)
+        exp, man = code // 2, (code % 2).float()
+        mags = torch.where(
+            exp == 0, man * 0.5, torch.exp2((exp - 1).float()) * (1.0 + man * 0.5)
+        )
+        # e2m1 * e4m3 is exact in fp32; saturate, then round once to E4M3.
+        values = (scales[:, None] * mags[None, :]).clamp(-448.0, 448.0)
+        lut = values.to(torch.float8_e4m3fn).view(torch.int64).reshape(256)
+        _NVFP4_E4M3_LUT[device] = lut
+    return lut
+
+
+@functools.cache
+def _get_kernel(heads: int):
+    return _kernel.compile_score_prefill(heads)
+
+
+def _kernel_shared():
+    # Imported late: the decode host imports this module.
+    from . import kernel_shared
+
+    return kernel_shared
+
+
+@functools.cache
+def _get_kernel_shared(heads: int, mq: int):
+    return _kernel_shared().compile_score_prefill_shared(heads, mq)
+
+
+def _supported(q_idx, index_cache, score, *int32_tensors) -> bool:
+    from vllm.platforms.rocm import on_gfx950
+
+    heads = q_idx.shape[1]
+    return (
+        on_gfx950()
+        and q_idx.dtype == torch.float8_e4m3fn
+        and heads in (1, 2, 4)
+        and q_idx.shape[2] == _kernel.D
+        and q_idx.stride(2) == 1
+        and q_idx.stride(1) == _kernel.D
+        and q_idx.stride(0) == heads * _kernel.D
+        and score.shape[0] == heads
+        and score.stride(2) == 1
+        and index_cache.dtype == torch.uint8
+        and index_cache.stride(2) == 1
+        and index_cache.stride(1) == _kernel.D // 2 + _kernel.D // 16
+        and index_cache.stride(0) % 16 == 0
+        and index_cache.data_ptr() % 16 == 0
+        and all(t.dtype == torch.int32 and t.is_contiguous() for t in int32_tensors)
+    )
+
+
+def minimax_m3_nvfp4_index_score_prefill(
+    q_idx: torch.Tensor,  # [total_q, H, 128] e4m3
+    index_cache: torch.Tensor,  # [num_blocks, 128, 72] uint8
+    score: torch.Tensor,  # [H, total_q, >= cdiv(max_seq_len, 128)] fp32
+    block_table: torch.Tensor,  # [batch, >= max_blocks] int32
+    cu_seqlens_q: torch.Tensor,  # [batch + 1] int32
+    seq_lens: torch.Tensor,  # [batch] int32
+    init_blocks: int,
+    local_blocks: int,
+    max_query_len: int,
+    max_seq_len: int,
+) -> None:
+    """Block scores of every prefill row for 1, 2 or 4 index heads."""
+    batch = cu_seqlens_q.shape[0] - 1
+    total_q, heads, _ = q_idx.shape
+    if total_q == 0 or block_table.shape[0] == 0:
+        return
+    assert _supported(q_idx, index_cache, score, block_table, cu_seqlens_q, seq_lens)
+    max_blocks = cdiv(max_seq_len, _kernel.BLK)
+    assert max_blocks <= score.shape[2]
+    work = max_blocks * cdiv(max_query_len, 256) * batch * heads
+    small = max_query_len * heads <= 128 or work < SMALL_WORK
+    shared = small if PIN_SHARED is None else PIN_SHARED
+    mq = PIN_MQ or (1 if small else 2)
+    target = SMALL_TARGET if shared and work < MID_WORK else TARGET_WORKGROUPS
+    tile_q = _kernel_shared().tile_q(mq) if shared else _kernel.TILE_Q
+    qt = cdiv(max_query_len * heads, tile_q)
+    min_seg = 1 if shared else _kernel.GROUP
+    max_nseg = max(1, min(MAX_SEGMENTS, max_blocks // min_seg))
+    nseg = max(1, min(max_nseg, target // max(1, qt * batch)))
+    slots = WORKGROUPS_PER_CU * cu_count(q_idx.device)
+    while nseg < max_nseg:
+        rounds = qt * batch * nseg / slots
+        if rounds / math.ceil(rounds) >= MIN_ROUND_FILL:
+            break
+        nseg += 1
+    lut = () if shared else (nvfp4_e4m3_lut(q_idx.device).data_ptr(),)
+    run_compiled(
+        _get_kernel_shared(heads, mq) if shared else _get_kernel(heads),
+        q_idx.data_ptr(),
+        index_cache.data_ptr(),
+        *lut,
+        score.data_ptr(),
+        block_table.data_ptr(),
+        cu_seqlens_q.data_ptr(),
+        seq_lens.data_ptr(),
+        int(index_cache.stride(0)),
+        int(block_table.stride(0)),
+        int(score.stride(1)),
+        int(nseg),
+        int(qt),
+        int(total_q),
+        int(batch),
+        int(init_blocks),
+        int(local_blocks),
+        int(score.stride(0)),
+        int(cdiv(qt * nseg * batch, 8) * 8),
+        torch.cuda.current_stream(),
+    )
diff --git a/vllm/models/minimax_m3/amd/ops/nvfp4_index_prefill/kernel.py b/vllm/models/minimax_m3/amd/ops/nvfp4_index_prefill/kernel.py
new file mode 100644
index 000000000000..248fc990f954
--- /dev/null
+++ b/vllm/models/minimax_m3/amd/ops/nvfp4_index_prefill/kernel.py
@@ -0,0 +1,614 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""NVFP4 prefill index scorer for MiniMax-M3 (FlyDSL, gfx950).
+
+Per (query row, block): the max of q . k over the row's causal tokens, with
+AITER's init / local sentinels, on e4m3_rn_satfinite(e2m1 x scale) keys. A
+workgroup takes 512 query rows (token * H + head) x a segment of the request's
+blocks; Q^T stays in VGPRs and each 16-token tile is one
+v_mfma_f32_16x16x128_f8f6f4 per 16-row N-tile. Raw NVFP4 pages arrive in LDS by
+DMA one block ahead, and block b + 1 is dequantized into the other key slot
+under block b's MFMAs.
+"""
+
+import flydsl.compiler as flyc
+import flydsl.expr as fx
+from aiter.ops.flydsl.kernels import buffer_ops
+from flydsl._mlir import ir
+from flydsl._mlir.dialects import llvm
+from flydsl.expr import const_expr, range_constexpr
+from flydsl.expr import rocdl as _rocdl
+from flydsl.expr.typing import T
+from flydsl.expr.typing import Vector as Vec
+
+from vllm.models.minimax_m3.amd.ops.flydsl_utils import (
+    as_f32,
+    bits,
+    dequant8,
+    global_i32_ptr,
+    key_order4,
+    maxf_nn,
+    mfma_fp8_16x16x128,
+    pack8,
+    permlane_swap,
+)
+
+NW = 4
+ROWS_PER_WAVE = 128
+TILE_Q = NW * ROWS_PER_WAVE
+BLK = 128  # tokens per block
+D = 128  # index head dim
+ROW_BYTES = D  # one e4m3 key row in LDS
+BLOCK_BYTES = BLK * ROW_BYTES
+RS_T = ROW_BYTES // 16  # 16 B tiles per LDS row
+SLOT_T = BLK * RS_T  # 16 B tiles per key slot
+N_SLOTS = 2
+LDS_BYTES = N_SLOTS * BLOCK_BYTES
+PAGE_CODES = BLK * D // 2  # page: [128 x 64 B] e2m1 codes, [128 x 8 B] e4m3 scales
+PAGE_BYTES = PAGE_CODES + BLK * 8
+GROUP = 4  # blocks per score store (16 B per row)
+# SPAN groups of a row are stored back to back (the first SPAN - 1 wait in LDS):
+# half-written 128 B score lines get evicted from L2 and cost extra write bytes.
+SPAN = 4
+STAGE_T = (SPAN - 1) * NW * 2 * 64  # 16 B LDS tiles: (group, wave, row half, lane)
+NI = ROWS_PER_WAVE // 16  # N-tiles per wave
+MT = BLK // 16  # M-tiles per block
+SENTINEL_INIT = 1e30  # AITER's forced-selection scores
+SENTINEL_LOCAL = 1e29
+STAGE_W = 2048 + 256  # raw bytes per wave and block: 32 tokens' codes + scales
+RAW_SLOTS = 2
+# M-tile -> code dwords dequantized under its MFMAs.
+DQ_AT = {1: (0,), 2: (1,), 3: (2,), 4: (3,), 5: (4,), 6: (5,), 7: (6, 7)}
+
+
+def _ir(v):
+    return v.ir_value() if hasattr(v, "ir_value") else v
+
+
+def _dma(rsrc, lds_addr, voff, size):
+    """buffer_load ... lds: ``size`` bytes per lane into LDS at ``lds_addr``."""
+    lds = llvm.IntToPtrOp(ir.Type.parse("!llvm.ptr<3>"), _ir(fx.Int32(lds_addr)))
+    fx.rocdl.raw_ptr_buffer_load_lds(
+        rsrc,
+        lds.result,
+        fx.Int32(size),
+        fx.Int32(voff),
+        fx.Int32(0),
+        fx.Int32(0),
+        aux=0,
+    )
+
+
+def _xlane_max4_pair(x, y):
+    """Max over lanes L, L^16, L^32, L^48 of two values at once, broadcast."""
+    a, b = permlane_swap(_rocdl.permlane32_swap, bits(x), bits(y))
+    m = bits(maxf_nn(as_f32(a), as_f32(b)))
+    a, b = permlane_swap(_rocdl.permlane16_swap, m, m)
+    big = bits(maxf_nn(as_f32(a), as_f32(b)))
+    a, b = permlane_swap(_rocdl.permlane32_swap, big, big)
+    return as_f32(a), as_f32(b)
+
+
+def tile_rows(pid, i32_qt, i32_nseg, i32_nb, arg_cu, arg_seq, heads, tile_q):
+    """Workgroup pid's request b, query tile and context segment.
+
+    Returns (live, b, seg, seq_start, seq_len, prefix, q_rows, row0); padded
+    workgroups read request 0 and are not live.
+    """
+    n_wg = i32_qt * i32_nseg * i32_nb
+    q_tile = pid % i32_qt
+    seg = (pid // i32_qt) % i32_nseg
+    b = pid // (i32_qt * i32_nseg)
+    bc = (pid < n_wg).select(b, fx.Int32(0))
+    cu = global_i32_ptr(arg_cu)
+    seq_start = fx.Int32(cu[bc])
+    q_len = fx.Int32(cu[bc + 1]) - seq_start
+    seq_len = fx.Int32(global_i32_ptr(arg_seq)[bc])
+    prefix = seq_len - q_len
+    q_rows = q_len * heads
+    row0 = q_tile * tile_q
+    live = (pid < n_wg) & (row0 < q_rows)
+    return live, b, seg, seq_start, seq_len, prefix, q_rows, row0
+
+
+def segment_blocks(prefix, row0, seq_len, seg, i32_nseg, heads, tile_q, align):
+    """The tile's nblk causal blocks and its segment [blk0, blk1) of them."""
+    hi_pos = prefix + (row0 + tile_q) // heads
+    hi = (hi_pos < seq_len).select(hi_pos, seq_len)
+    nblk = (hi + (BLK - 1)) // BLK
+    seg_len = (nblk + i32_nseg - 1) // i32_nseg
+    if align > 1:
+        seg_len = (seg_len + (align - 1)) & ~(align - 1)
+    blk0 = seg * seg_len
+    blk1e = blk0 + seg_len
+    blk1 = (blk1e < nblk).select(blk1e, nblk)
+    return nblk, blk0, blk1
+
+
+def score_resources(arg_q, arg_score, i32_total_q, i32_hs, i32_S, heads):
+    """Buffer resources of the Q rows and the scores."""
+    qr = buffer_ops.create_buffer_resource_from_addr(
+        arg_q, num_records_bytes=fx.Int64(i32_total_q) * (heads * D)
+    )
+    sr = buffer_ops.create_buffer_resource_from_addr(
+        arg_score,
+        num_records_bytes=(
+            fx.Int64(i32_hs) * (heads - 1) + fx.Int64(i32_total_q) * fx.Int64(i32_S)
+        )
+        * 4,
+    )
+    return qr, sr
+
+
+def page_lookup(arg_bt, bt_row, end):
+    """Page of a block, clamped to the last block before ``end``.
+
+    A constant-space s_load: before an s_buffer_load LLVM would wait for the
+    in-flight LDS DMAs.
+    """
+    bt = llvm.IntToPtrOp(ir.Type.parse("!llvm.ptr<4>"), _ir(fx.Int64(arg_bt))).result
+
+    def page_of(blk):
+        bi = (blk < end).select(blk, end - 1)
+        p = buffer_ops.get_element_ptr(
+            bt, byte_offset=_ir((bt_row + bi) * 4), elem_type=T.i8
+        )
+        return fx.Int32(llvm.LoadOp(T.i32, p, alignment=4, invariant=True).result)
+
+    return page_of
+
+
+def with_sentinels(v, blk, nb, i32_init, i32_local):
+    """AITER's init / local sentinels on block blk of a row of nb causal blocks."""
+    is_local = (blk >= nb - i32_local) & (blk < nb)
+    is_init = (blk < i32_init) & (blk < nb)
+    return is_local.select(
+        fx.Float32(SENTINEL_LOCAL), is_init.select(fx.Float32(SENTINEL_INIT), v)
+    )
+
+
+def compile_score_prefill(heads: int = 1):
+    """Build the scorer for 1, 2 or 4 index heads (sharing one index-K head)."""
+    assert heads in (1, 2, 4), heads
+    H = heads
+    name = f"m3_nvfp4_index_score_prefill_tq{TILE_Q}_h{H}"
+
+    @fx.struct
+    class Shared:
+        a: fx.Array[fx.Uint8, LDS_BYTES, 16]  # e4m3 key slots
+        lut: fx.Array[fx.Int32, 512, 16]
+        raw: fx.Array[fx.Uint8, RAW_SLOTS * NW * STAGE_W, 16]  # DMA'd pages
+        stage: fx.Array[fx.Int32, STAGE_T * 4, 16]  # staged score groups
+
+    @flyc.kernel(name=name, known_block_size=[64 * NW, 1, 1])
+    def kernel(
+        arg_q: fx.Int64,  # [total_q, H, 128] e4m3
+        arg_kv: fx.Int64,  # [num_blocks, 128, 72] uint8
+        arg_lut: fx.Int64,  # nvfp4_e4m3_lut [256] int64
+        arg_score: fx.Int64,  # [H, total_q, S] fp32
+        arg_bt: fx.Int64,  # [batch, >= max_blocks] int32
+        arg_cu: fx.Int64,  # cu_seqlens_q [batch + 1] int32
+        arg_seq: fx.Int64,  # seq_lens [batch] int32
+        i64_page_stride: fx.Int64,
+        i32_bt_stride: fx.Int32,
+        i32_S: fx.Int32,
+        i32_nseg: fx.Int32,
+        i32_qt: fx.Int32,  # query tiles per request
+        i32_total_q: fx.Int32,
+        i32_nb: fx.Int32,
+        i32_init: fx.Int32,
+        i32_local: fx.Int32,
+        i32_hs: fx.Int32,  # score head stride
+    ):
+        smem = fx.SharedAllocator().allocate(Shared).peek()
+        tx, pid = fx.thread_idx.x, fx.block_idx.x
+        lane = tx % 64
+        wave = fx.Int32(fx.rocdl.readfirstlane(T.i32, tx // 64))
+        l16, q16 = lane % 16, lane // 16
+        live, b, seg, seq_start, seq_len, prefix, q_rows, row0 = tile_rows(
+            pid, i32_qt, i32_nseg, i32_nb, arg_cu, arg_seq, H, TILE_Q
+        )
+        if live:
+            nblk, blk0, blk1 = segment_blocks(
+                prefix, row0, seq_len, seg, i32_nseg, H, TILE_Q, GROUP
+            )
+            if blk0 < blk1:
+                qr, sr = score_resources(
+                    arg_q, arg_score, i32_total_q, i32_hs, i32_S, H
+                )
+                page_of = page_lookup(arg_bt, b * i32_bt_stride, nblk)
+
+                # Lane l of wave w owns token 32 w + l / 2, dims 64 (l % 2) ..
+                # +63. The DMA writes LDS lane-linearly, so each lane reads back
+                # its own code and scale bytes.
+                raw_base_s = fx.Int32(
+                    fx.rocdl.readfirstlane(T.i32, fx.Int32(fx.ptrtoint(smem.raw.ptr)))
+                )
+                code_voff = [wave * 2048 + lane * 32 + h * 16 for h in range(2)]
+                scale_voff = PAGE_CODES + wave * 256 + lane * 4
+
+                def dma_raw(page, st):
+                    kr = buffer_ops.create_buffer_resource_from_addr(
+                        arg_kv + fx.Int64(page) * i64_page_stride,
+                        num_records_bytes=PAGE_BYTES,
+                    )
+                    base = raw_base_s + (st * NW + wave) * STAGE_W
+                    for h in range_constexpr(2):
+                        _dma(kr, base + h * 1024, code_voff[h], 16)
+                    _dma(kr, base + 2048, scale_voff, 4)
+
+                raw_t = fx.logical_divide(
+                    fx.make_view(
+                        fx.recast_iter(fx.Int32, smem.raw.ptr),
+                        fx.make_layout(RAW_SLOTS * NW * STAGE_W // 4, 1),
+                    ),
+                    fx.make_layout(4, 1),
+                )
+                raw_d = smem.raw.ptr
+
+                def read_raw(st):
+                    """This lane's 8 code dwords and its scale dword."""
+                    b4 = (st * NW + wave) * (STAGE_W // 4)
+                    out = []
+                    for h in range_constexpr(2):
+                        r = fx.make_rmem_tensor(fx.make_layout(4, 1), fx.Int32)
+                        fx.copy(
+                            fx.make_copy_atom(fx.UniversalCopy128b(), fx.Int32),
+                            fx.slice(raw_t, (None, b4 // 4 + h * 64 + lane)),
+                            r,
+                        )
+                        v = Vec(r.load())
+                        out += [fx.Int32(v[i]) for i in range_constexpr(4)]
+                    sc = fx.Int32(
+                        fx.make_view(
+                            fx.recast_iter(fx.Int32, raw_d),
+                            fx.make_layout(RAW_SLOTS * NW * STAGE_W // 4, 1),
+                        )[b4 + 512 + lane]
+                    )
+                    return out + [sc]
+
+                pg = [page_of(blk0 + i) for i in range_constexpr(2)]
+                for i in range_constexpr(2):
+                    dma_raw(pg[i], i)
+                pg_next = page_of(blk0 + 2)
+
+                lut_g = global_i32_ptr(arg_lut)
+                lut_l = smem.lut.ptr
+                for i in range_constexpr(2):
+                    lut_l[tx * 2 + i] = fx.Int32(lut_g[tx * 2 + i])
+
+                # Rows past q_len * H read the next request's Q and are masked at
+                # the store.
+                rows = [row0 + wave * ROWS_PER_WAVE + ni * 16 + l16 for ni in range(NI)]
+                qpos = [prefix + r // H for r in rows]
+                # Q^T in the key's dim order (even dims, then odd dims per 8).
+                qf = []
+                for ni in range_constexpr(NI):
+                    qrow = (seq_start * H + rows[ni]) * ROW_BYTES
+                    halves = []
+                    for h in range_constexpr(2):
+                        d = Vec(
+                            buffer_ops.buffer_load(
+                                qr,
+                                (qrow + h * 64 + q16 * 16) // 4,
+                                vec_width=4,
+                                dtype=fx.Int32,
+                            )
+                        )
+                        d = [fx.Int32(d[i]) for i in range_constexpr(4)]
+                        halves.append(key_order4(d))
+                    qf.append(pack8(halves[0], halves[1]))
+
+                lds16 = fx.logical_divide(
+                    fx.make_view(
+                        fx.recast_iter(fx.Int32, smem.a.ptr),
+                        fx.make_layout(LDS_BYTES // 4, 1),
+                    ),
+                    fx.make_layout(4, 1),
+                )
+                lds_atom = fx.make_copy_atom(fx.UniversalCopy128b(), fx.Int32)
+                lut8 = fx.logical_divide(
+                    fx.make_view(
+                        fx.recast_iter(fx.Int32, smem.lut.ptr), fx.make_layout(512, 1)
+                    ),
+                    fx.make_layout(2, 1),
+                )
+                lut_atom = fx.make_copy_atom(fx.UniversalCopy64b(), fx.Int32)
+
+                def lds_load16(tile):
+                    r = fx.make_rmem_tensor(fx.make_layout(4, 1), fx.Int32)
+                    fx.copy(lds_atom, fx.slice(lds16, (None, tile)), r)
+                    return r.load()
+
+                def lds_store16(tile, vals):
+                    r = fx.make_rmem_tensor(fx.make_layout(4, 1), fx.Int32)
+                    r.store(Vec.from_elements(vals, fx.Int32))
+                    fx.copy(lds_atom, r, fx.slice(lds16, (None, tile)))
+
+                def lut_load(sbyte):
+                    r = fx.make_rmem_tensor(fx.make_layout(2, 1), fx.Int32)
+                    fx.copy(lut_atom, fx.slice(lut8, (None, sbyte)), r)
+                    v = Vec(r.load())
+                    return fx.Int32(v[0]), fx.Int32(v[1])
+
+                # 16 B column c of key row t sits at c ^ ((t >> 1) & 7).
+                t_row = wave * 32 + (lane >> 1)
+                hh = lane & 1
+                wr_base = t_row * RS_T
+                wr_sw = (t_row >> 1) & 7
+
+                def dq_tables(raw):
+                    sc = raw[8]
+                    return [
+                        lut_load((sc >> (8 * g)) & 0xFF) for g in range_constexpr(4)
+                    ]
+
+                def dq_step(raw, tabs, dslot, k, pend):
+                    """Code dword k -> 8 keys; every second one stores 16 B."""
+                    lo, hi = dequant8(raw[k], *tabs[k // 2])
+                    if const_expr(k % 2 == 0):
+                        return [lo, hi]
+                    col = (hh * 4 + k // 2) ^ wr_sw
+                    lds_store16(dslot * SLOT_T + wr_base + col, pend + [lo, hi])
+                    return []
+
+                zero4 = fx.Vector.filled(4, 0.0, fx.Float32)
+                neg_inf = fx.Float32(float("-inf"))
+                rd_row = l16 * RS_T
+                rd_col = [q16 ^ (l16 >> 1), (4 + q16) ^ (l16 >> 1)]
+
+                def compute_block(slot, blk, masked, st, dslot):
+                    """This wave's scores of ``blk``; dequantizes the next block's
+                    raw page ``st`` into key slot ``dslot`` on the way."""
+                    if const_expr(masked):
+                        vis = [
+                            qpos[ni] - (blk * BLK + q16 * 4)
+                            for ni in range_constexpr(NI)
+                        ]
+                    run = [neg_inf for _ in range(NI)]
+                    raw = read_raw(st)
+                    tabs = dq_tables(raw)
+
+                    def read_a(mt):
+                        base = slot * SLOT_T + mt * 16 * RS_T + rd_row
+                        return pack8(
+                            lds_load16(base + rd_col[0]), lds_load16(base + rd_col[1])
+                        )
+
+                    accs = [[None] * NI, [None] * NI]
+
+                    def issue(mt, a):
+                        acc = accs[mt % 2]
+                        for ni in range_constexpr(NI):
+                            acc[ni] = mfma_fp8_16x16x128(a, qf[ni], zero4)
+
+                    def reduce(mt):
+                        acc = accs[mt % 2]
+                        for ni in range_constexpr(NI):
+                            v = acc[ni]
+                            for i in range_constexpr(4):
+                                x = v[i]
+                                if const_expr(masked):
+                                    x = (vis[ni] >= mt * 16 + i).select(x, neg_inf)
+                                run[ni] = maxf_nn(run[ni], x)
+
+                    pend: list = []
+                    a_next = read_a(0)
+                    for mt in range_constexpr(MT):
+                        a = a_next
+                        if const_expr(mt + 1 < MT):
+                            a_next = read_a(mt + 1)
+                        issue(mt, a)
+                        if const_expr(mt >= 1):
+                            reduce(mt - 1)
+                        ks = DQ_AT.get(mt, ())
+                        for k in ks:
+                            pend = dq_step(raw, tabs, dslot, k, pend)
+                        n_valu = (16 if not masked else 48) + 10 * len(ks)
+                        n_dsw = sum(k % 2 for k in ks)
+                        for i in range_constexpr(NI):
+                            fx.rocdl.sched_group_barrier(0x008, 1, 0)
+                            fx.rocdl.sched_group_barrier(
+                                0x002, (n_valu + NI - 1) // NI, 0
+                            )
+                            if const_expr(i < 2):
+                                fx.rocdl.sched_group_barrier(0x100, 1, 0)
+                            if const_expr(i == 4 and n_dsw):
+                                fx.rocdl.sched_group_barrier(0x200, n_dsw, 0)
+                        fx.rocdl.sched_barrier(0)
+                    reduce(MT - 1)
+                    out = []
+                    for p in range_constexpr(NI // 2):
+                        x, y = _xlane_max4_pair(run[2 * p], run[2 * p + 1])
+                        out += [x, y]
+                    return out
+
+                stage16 = fx.logical_divide(
+                    fx.make_view(
+                        fx.recast_iter(fx.Int32, smem.stage.ptr),
+                        fx.make_layout(STAGE_T * 4, 1),
+                    ),
+                    fx.make_layout(4, 1),
+                )
+
+                def stage_tile(k, h):
+                    return ((k * NW + wave) * 2 + h) * 64 + lane
+
+                def store_group(g0, sc):
+                    """Scores sc[j][ni] of blocks g0 .. g0 + 3, sentinels applied.
+
+                    Lane q16 = k stores rows 2k and 2k + 1. A group waits in LDS
+                    unless it ends its SPAN piece or the segment; that one stores
+                    the piece's staged groups and itself back to back.
+                    """
+                    full = g0 + GROUP <= blk1
+                    pos = (g0 >> 2) & (SPAN - 1)
+                    piece0 = g0 & ~(SPAN * GROUP - 1)
+                    first = (piece0 < blk0).select(
+                        (blk0 >> 2) & (SPAN - 1), fx.Int32(0)
+                    )
+                    flush = (pos == SPAN - 1) | (g0 + GROUP >= blk1)
+
+                    def store_now(vals, off, valid):
+                        """16 B per row, single dwords for a partial tail."""
+                        if full:
+                            buffer_ops.buffer_store(
+                                fx.Vector.from_elements(vals, fx.Float32),
+                                sr,
+                                off,
+                                mask=valid,
+                            )
+                        else:
+                            for j in range_constexpr(GROUP):
+                                buffer_ops.buffer_store(
+                                    vals[j], sr, off + j, mask=valid & (g0 + j < blk1)
+                                )
+
+                    def flush_piece(h, vals, off, valid):
+                        for k in range_constexpr(SPAN - 1):
+                            t = fx.make_rmem_tensor(fx.make_layout(4, 1), fx.Int32)
+                            fx.copy(
+                                lds_atom, fx.slice(stage16, (None, stage_tile(k, h))), t
+                            )
+                            buffer_ops.buffer_store(
+                                Vec(t.load()),
+                                sr,
+                                off - (pos - k) * GROUP,
+                                mask=valid & (first <= k) & (pos > k),
+                            )
+                        store_now(vals, off, valid)
+
+                    def stage_group(h, vals):
+                        t = fx.make_rmem_tensor(fx.make_layout(4, 1), fx.Int32)
+                        t.store(Vec.from_elements([bits(x) for x in vals], fx.Int32))
+                        fx.copy(
+                            lds_atom, t, fx.slice(stage16, (None, stage_tile(pos, h)))
+                        )
+
+                    for h in range_constexpr(2):
+                        row = rows[h]
+                        vals = [sc[j][h] for j in range_constexpr(GROUP)]
+                        for k in range_constexpr(1, NW):
+                            is_k = q16 == k
+                            row = is_k.select(rows[2 * k + h], row)
+                            vals = [
+                                is_k.select(sc[j][2 * k + h], vals[j])
+                                for j in range_constexpr(GROUP)
+                            ]
+                        valid = row < q_rows
+                        tok = row // H
+                        nb_row = (prefix + tok) // BLK + 1
+                        for j in range_constexpr(GROUP):
+                            blk = g0 + j
+                            vals[j] = with_sentinels(
+                                vals[j], blk, nb_row, i32_init, i32_local
+                            )
+                        off = (row % H) * i32_hs + (seq_start + tok) * i32_S + g0
+                        if flush:
+                            flush_piece(h, vals, off, valid)
+                        else:
+                            stage_group(h, vals)
+
+                n_groups = (blk1 - blk0 + (GROUP - 1)) // GROUP
+                # Groups every row of the tile sees whole need no causal mask.
+                v = prefix + row0 // H - (BLK - 1)
+                last_full = (v < 0).select(fx.Int32(-1), v // BLK)
+                d = last_full - (GROUP - 1) - blk0
+                nf = (d < 0).select(fx.Int32(0), d // GROUP + 1)
+                n_full = (nf < n_groups).select(nf, n_groups)
+
+                # Prologue: block blk0 dequantized into slot 0.
+                fx.rocdl.s_waitcnt(vmcnt=0, lgkmcnt=0)
+                fx.gpu.barrier()
+                raw0 = read_raw(0)
+                tabs0 = dq_tables(raw0)
+                pend: list = []
+                for k in range_constexpr(8):
+                    pend = dq_step(raw0, tabs0, 0, k, pend)
+                fx.gpu.barrier()
+
+                def group(g0, pg_next, masked):
+                    """Blocks g0 .. g0 + 3: per block b, DMA block b + 2's page,
+                    score b while dequantizing b + 1, then wait and barrier."""
+                    sc = []
+                    for j in range_constexpr(GROUP):
+                        blk = g0 + j
+                        dma_raw(pg_next, (j + 2) % RAW_SLOTS)
+                        pg_next = page_of(blk + 3)
+                        sc.append(
+                            compute_block(
+                                j % N_SLOTS,
+                                blk,
+                                masked,
+                                (j + 1) % RAW_SLOTS,
+                                (j + 1) % N_SLOTS,
+                            )
+                        )
+                        fx.rocdl.s_waitcnt(vmcnt=0, lgkmcnt=0)
+                        fx.gpu.barrier()
+                    store_group(g0, sc)
+                    return pg_next
+
+                def one(st):
+                    return fx.Int32(st[0] if isinstance(st, (list, tuple)) else st)
+
+                for g, st in range(  # type: ignore[call-overload]
+                    fx.Index(0), fx.Index(n_full), fx.Index(1), init=[_ir(pg_next)]
+                ):
+                    pgn = group(blk0 + fx.Int32(g) * GROUP, one(st), False)
+                    res = yield [_ir(pgn)]
+                for g, st in range(  # type: ignore[call-overload]
+                    fx.Index(n_full),
+                    fx.Index(n_groups),
+                    fx.Index(1),
+                    init=[_ir(one(res))],
+                ):
+                    pgn = group(blk0 + fx.Int32(g) * GROUP, one(st), True)
+                    res = yield [_ir(pgn)]
+
+    @flyc.jit
+    def launch(
+        arg_q: fx.Int64,
+        arg_kv: fx.Int64,
+        arg_lut: fx.Int64,
+        arg_score: fx.Int64,
+        arg_bt: fx.Int64,
+        arg_cu: fx.Int64,
+        arg_seq: fx.Int64,
+        i64_page_stride: fx.Int64,
+        i32_bt_stride: fx.Int32,
+        i32_S: fx.Int32,
+        i32_nseg: fx.Int32,
+        i32_qt: fx.Int32,
+        i32_total_q: fx.Int32,
+        i32_nb: fx.Int32,
+        i32_init: fx.Int32,
+        i32_local: fx.Int32,
+        i32_hs: fx.Int32,
+        i32_grid: fx.Int32,
+        stream: fx.Stream,
+    ):
+        kernel(
+            arg_q,
+            arg_kv,
+            arg_lut,
+            arg_score,
+            arg_bt,
+            arg_cu,
+            arg_seq,
+            i64_page_stride,
+            i32_bt_stride,
+            i32_S,
+            i32_nseg,
+            i32_qt,
+            i32_total_q,
+            i32_nb,
+            i32_init,
+            i32_local,
+            i32_hs,
+            # Without the hint LLVM picks 264 VGPRs, one wave per SIMD.
+            value_attrs={"rocdl.waves_per_eu": 2},
+        ).launch(grid=(fx.Int64(i32_grid), 1, 1), block=(64 * NW, 1, 1), stream=stream)
+
+    launch.kernel_name = name
+    return launch
diff --git a/vllm/models/minimax_m3/amd/ops/nvfp4_index_prefill/kernel_shared.py b/vllm/models/minimax_m3/amd/ops/nvfp4_index_prefill/kernel_shared.py
new file mode 100644
index 000000000000..da68e37f49ca
--- /dev/null
+++ b/vllm/models/minimax_m3/amd/ops/nvfp4_index_prefill/kernel_shared.py
@@ -0,0 +1,382 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""NVFP4 prefill index scorer for small work (FlyDSL, gfx950).
+
+The scores of ``kernel.py`` with one dequantized block per workgroup: a
+workgroup is 4 waves x 32 MQ query rows x a segment of the causal window. Per
+block, wave w dequantizes token tile w into e4m3 A fragments in LDS (double
+buffered, one barrier per block) while the MFMAs of the previous block run;
+every wave multiplies all four token tiles against its query rows with
+v_mfma_f32_32x32x64_f8f6f4. Blocks before the wave's first query position skip
+the causal mask in a loop of their own.
+"""
+
+import flydsl.compiler as flyc
+import flydsl.expr as fx
+from aiter.ops.flydsl.kernels import buffer_ops
+from flydsl._mlir import ir as _mlir_ir
+from flydsl._mlir.dialects import arith as _arith
+from flydsl.expr import const_expr, range_constexpr
+from flydsl.expr import rocdl as _rocdl
+from flydsl.expr.typing import T
+from flydsl.expr.typing import Vector as Vec
+
+from vllm.models.minimax_m3.amd.ops.flydsl_utils import (
+    as_f32,
+    bits,
+    cvt_fp4x2_f16,
+    cvt_fp8x2_f16,
+    maxf_nn,
+    permlane_swap,
+)
+from vllm.models.minimax_m3.amd.ops.nvfp4_index_prefill.kernel import (
+    BLK,
+    PAGE_BYTES,
+    PAGE_CODES,
+    D,
+    page_lookup,
+    score_resources,
+    segment_blocks,
+    tile_rows,
+    with_sentinels,
+)
+
+NW = 4
+MQ_CHOICES = (1, 2)  # 32-row query tiles per wave
+TILES = BLK // 32  # 32-token tiles per block
+KS = 2  # MFMA k-steps per key row
+FL = 32  # fragment bytes per lane
+FRAG_BYTES = TILES * KS * 64 * FL  # one block's A fragments
+RING = 4  # fragment reads in flight ahead of the MFMA
+
+
+def _sat448_fp8x2(x):
+    """Clamp 2 x f16 to +-448; the hardware fp8 convert does not saturate."""
+    fm = _mlir_ir.Attribute.parse("#arith.fastmath<nnan>")
+    hi = fx.as_ir_value(fx.Vector.filled(2, 448.0, fx.Float16))
+    lo = fx.as_ir_value(fx.Vector.filled(2, -448.0, fx.Float16))
+    v = _arith.MaxNumFOp(fx.as_ir_value(x), lo, fastmath=fm).result
+    return _arith.MinNumFOp(v, hi, fastmath=fm).result
+
+
+def _fp8_pair(w, sv):
+    """Code dword w x splat scale sv -> 2 dwords of e4m3."""
+    out = []
+    for h in range_constexpr(2):
+        p = fx.as_ir_value(fx.Vector.filled(2, 0, fx.Int16))
+        for k in range_constexpr(2):
+            x = _sat448_fp8x2(cvt_fp4x2_f16(w, 2 * h + k) * sv)
+            p = fx.rocdl.cvt_scalef32_pk_fp8_f16(
+                Vec.make_type(2, fx.Int16),
+                p,
+                x,
+                fx.as_ir_value(fx.Float32(1.0)),
+                k == 1,
+            )
+        out.append(fx.Int32(Vec(p).bitcast(fx.Int32)[0]))
+    return out
+
+
+def _mfma_fp8_32x32x64(a8, b8, acc16):
+    """v_mfma_f32_32x32x64_f8f6f4 on e4m3; lane (r, h) gets
+    `D[8 (i // 4) + 4 h + i % 4][r]`."""
+    return Vec(
+        fx.rocdl.mfma_scale_f32_32x32x64_f8f6f4(
+            Vec.make_type(16, fx.Float32), [a8, b8, acc16, 0, 0, 0, 0, 0, 0]
+        )
+    )
+
+
+def _xlane_max32(x):
+    """max(x[L], x[L ^ 32])."""
+    a, b = permlane_swap(_rocdl.permlane32_swap, bits(x), bits(x))
+    return maxf_nn(as_f32(a), as_f32(b))
+
+
+def tile_q(mq: int) -> int:
+    """Query rows per workgroup."""
+    return NW * 32 * mq
+
+
+def compile_score_prefill_shared(heads: int = 1, mq: int = 2):
+    """Build the scorer for 1, 2 or 4 index heads and ``mq`` 32-row query tiles
+    per wave."""
+    assert heads in (1, 2, 4) and mq in MQ_CHOICES, (heads, mq)
+    H, MQ = heads, mq
+    ROWS_PER_WAVE = 32 * MQ
+    TILE_Q = tile_q(MQ)
+    name = f"m3_nvfp4_index_score_prefill_fp8_shared_tq{TILE_Q}_h{H}"
+
+    @fx.struct
+    class Shared:
+        frag: fx.Array[fx.Uint8, 2 * FRAG_BYTES, 16]  # A fragments of 2 blocks
+
+    @flyc.kernel(name=name, known_block_size=[64 * NW, 1, 1])
+    def kernel(
+        arg_q: fx.Int64,  # [total_q, H, 128] e4m3
+        arg_kv: fx.Int64,  # [num_blocks, 128, 72] uint8
+        arg_score: fx.Int64,  # [H, total_q, S] fp32
+        arg_bt: fx.Int64,  # [batch, >= max_blocks] int32
+        arg_cu: fx.Int64,  # cu_seqlens_q [batch + 1] int32
+        arg_seq: fx.Int64,  # seq_lens [batch] int32
+        i64_page_stride: fx.Int64,
+        i32_bt_stride: fx.Int32,
+        i32_S: fx.Int32,
+        i32_nseg: fx.Int32,
+        i32_qt: fx.Int32,  # query tiles per request
+        i32_total_q: fx.Int32,
+        i32_nb: fx.Int32,
+        i32_init: fx.Int32,
+        i32_local: fx.Int32,
+        i32_hs: fx.Int32,  # score head stride
+    ):
+        smem = fx.SharedAllocator().allocate(Shared).peek()
+        tx, pid = fx.thread_idx.x, fx.block_idx.x
+        lane = tx % 64
+        wave = fx.Int32(fx.rocdl.readfirstlane(T.i32, tx // 64))
+        r32, h = lane % 32, lane // 32
+        live, b, seg, seq_start, seq_len, prefix, q_rows, row0 = tile_rows(
+            pid, i32_qt, i32_nseg, i32_nb, arg_cu, arg_seq, H, TILE_Q
+        )
+        if live:
+            nblk, blk0, blk1 = segment_blocks(
+                prefix, row0, seq_len, seg, i32_nseg, H, TILE_Q, 1
+            )
+            if blk0 < blk1:
+                qr, sr = score_resources(
+                    arg_q, arg_score, i32_total_q, i32_hs, i32_S, H
+                )
+                page_of = page_lookup(arg_bt, b * i32_bt_stride, blk1)
+
+                tok = wave * 32 + r32
+
+                def load_raw(blk):
+                    """Token tok's codes and scales of dims 64 h .. +63; zeros past
+                    the segment."""
+                    ok = blk < blk1
+                    kr = buffer_ops.create_buffer_resource_from_addr(
+                        arg_kv + fx.Int64(page_of(blk)) * i64_page_stride,
+                        num_records_bytes=PAGE_BYTES,
+                    )
+                    w = []
+                    for i in range_constexpr(2):
+                        v = Vec(
+                            buffer_ops.buffer_load(
+                                kr,
+                                (tok * 64 + h * 32 + i * 16) // 4,
+                                vec_width=4,
+                                dtype=fx.Int32,
+                                mask=ok,
+                            )
+                        )
+                        w += [fx.Int32(v[e]) for e in range_constexpr(4)]
+                    sc = fx.Int32(
+                        buffer_ops.buffer_load(
+                            kr,
+                            (PAGE_CODES + tok * 8 + h * 4) // 4,
+                            vec_width=1,
+                            dtype=fx.Int32,
+                            mask=ok,
+                        )
+                    )
+                    return w + [sc]
+
+                frag128 = fx.logical_divide(
+                    fx.make_view(
+                        fx.recast_iter(fx.Int32, smem.frag.ptr),
+                        fx.make_layout(2 * FRAG_BYTES // 4, 1),
+                    ),
+                    fx.make_layout(4, 1),
+                )
+                atom16 = fx.make_copy_atom(fx.UniversalCopy128b(), fx.Int32)
+
+                def frag_slot(buf, n, s, half=0):
+                    return (
+                        buf * (FRAG_BYTES // 16)
+                        + ((n * KS + s) * 64 + lane) * (FL // 16)
+                        + half
+                    )
+
+                def store16(words, slot):
+                    r = fx.make_rmem_tensor(fx.make_layout(4, 1), fx.Int32)
+                    r.store(Vec.from_elements(words, fx.Int32))
+                    fx.copy(atom16, r, fx.slice(frag128, (None, slot)))
+
+                def load16(slot):
+                    r = fx.make_rmem_tensor(fx.make_layout(4, 1), fx.Int32)
+                    fx.copy(atom16, fx.slice(frag128, (None, slot)), r)
+                    return [fx.Int32(x) for x in Vec(r.load())]
+
+                def dequant_step(raw, buf, s):
+                    """K-step s (dims 64 h + 32 s .. +31) of this wave's tile."""
+                    scl = cvt_fp8x2_f16(raw[8], s == 1)
+                    sv = [
+                        Vec.from_elements([scl[g], scl[g]], fx.Float16) for g in (0, 1)
+                    ]
+                    k = []
+                    for e in range_constexpr(4):
+                        k += _fp8_pair(raw[4 * s + e], sv[e // 2])
+                    store16(k[:4], frag_slot(buf, wave, s, 0))
+                    store16(k[4:], frag_slot(buf, wave, s, 1))
+
+                def read_frag(buf, n, s):
+                    return Vec.from_elements(
+                        load16(frag_slot(buf, n, s, 0))
+                        + load16(frag_slot(buf, n, s, 1)),
+                        fx.Int32,
+                    )
+
+                # Rows past q_len * H are zero and never stored.
+                rows = [row0 + wave * ROWS_PER_WAVE + 32 * m + r32 for m in range(MQ)]
+                qpos = [prefix + rw // H for rw in rows]
+                qf = []
+                for m in range_constexpr(MQ):
+                    qrow = (seq_start * H + rows[m]) * D
+                    d = []
+                    for i in range_constexpr(4):
+                        v = Vec(
+                            buffer_ops.buffer_load(
+                                qr,
+                                (qrow + h * 64 + i * 16) // 4,
+                                vec_width=4,
+                                dtype=fx.Int32,
+                                mask=rows[m] < q_rows,
+                            )
+                        )
+                        d += [fx.Int32(v[e]) for e in range_constexpr(4)]
+                    qf.append(
+                        [
+                            Vec.from_elements(d[8 * s : 8 * s + 8], fx.Int32)
+                            for s in range(KS)
+                        ]
+                    )
+
+                zero16 = fx.Vector.filled(16, 0.0, fx.Float32)
+                neg_inf = fx.Float32(float("-inf"))
+                st_tok = [rw // H for rw in rows]
+                st_out = [
+                    (rw % H) * i32_hs + (seq_start + t) * i32_S
+                    for rw, t in zip(rows, st_tok)
+                ]
+                st_nb = [qp // BLK + 1 for qp in qpos]
+                qpos_min = prefix + (row0 + wave * ROWS_PER_WAVE) // H
+
+                def max16(acc):
+                    v = acc[0]
+                    for i in range_constexpr(1, 16):
+                        v = maxf_nn(v, acc[i])
+                    return v
+
+                def max16_masked(acc, lim):
+                    v = neg_inf
+                    for i in range_constexpr(16):
+                        x = (lim >= 8 * (i // 4) + i % 4).select(acc[i], neg_inf)
+                        v = maxf_nn(v, x)
+                    return v
+
+                def block(blk, buf, raw_next, masked):
+                    """Scores of block blk for this wave's rows; dequantizes
+                    raw_next into the other buffer meanwhile."""
+                    runs = [neg_inf] * MQ
+                    jobs = [(n, s) for n in range(TILES) for s in range(KS)]
+                    ring = [None] * len(jobs)
+                    for j in range_constexpr(RING - 1):
+                        ring[j] = read_frag(buf, *jobs[j])
+                    for n in range_constexpr(TILES):
+                        accs = [zero16] * MQ
+                        for s in range_constexpr(KS):
+                            j = n * KS + s
+                            if const_expr(j + RING - 1 < len(jobs)):
+                                ring[j + RING - 1] = read_frag(buf, *jobs[j + RING - 1])
+                            if const_expr(j % TILES == 0):
+                                dequant_step(raw_next, 1 - buf, j // TILES)
+                            for m in range_constexpr(MQ):
+                                accs[m] = _mfma_fp8_32x32x64(ring[j], qf[m][s], accs[m])
+                        base = blk * BLK + 32 * n + 4 * h
+                        for m in range_constexpr(MQ):
+                            if const_expr(masked):
+                                mx = max16_masked(accs[m], qpos[m] - base)
+                            else:
+                                mx = max16(accs[m])
+                            runs[m] = maxf_nn(runs[m], mx)
+                    for m in range_constexpr(MQ):
+                        v = _xlane_max32(runs[m])
+                        v = with_sentinels(v, blk, st_nb[m], i32_init, i32_local)
+                        buffer_ops.buffer_store(
+                            v,
+                            sr,
+                            st_out[m] + blk,
+                            mask=(rows[m] < q_rows) & (h == 0),
+                        )
+
+                raw0 = load_raw(blk0)
+                for s in range_constexpr(KS):
+                    dequant_step(raw0, 0, s)
+                raw = load_raw(blk0 + 1)
+                fx.gpu.barrier()
+
+                # Blocks below dsplit need no causal mask. Two loops rather than a
+                # branch per block keep the wave within 256 registers; waves split
+                # at different blocks but pass the same barriers.
+                dfirst = (qpos_min + 1) // BLK
+                dsplit = (dfirst < blk0).select(blk0, dfirst)
+                dsplit = (dsplit > blk1).select(blk1, dsplit)
+                blk = blk0
+                buf = fx.Int32(0)
+                while blk < dsplit:
+                    nxt = load_raw(blk + 2)
+                    block(blk, buf, raw, False)
+                    raw = nxt
+                    fx.gpu.barrier()
+                    blk = blk + 1
+                    buf = 1 - buf
+                while blk < blk1:
+                    nxt = load_raw(blk + 2)
+                    block(blk, buf, raw, True)
+                    raw = nxt
+                    fx.gpu.barrier()
+                    blk = blk + 1
+                    buf = 1 - buf
+
+    @flyc.jit
+    def launch(
+        arg_q: fx.Int64,
+        arg_kv: fx.Int64,
+        arg_score: fx.Int64,
+        arg_bt: fx.Int64,
+        arg_cu: fx.Int64,
+        arg_seq: fx.Int64,
+        i64_page_stride: fx.Int64,
+        i32_bt_stride: fx.Int32,
+        i32_S: fx.Int32,
+        i32_nseg: fx.Int32,
+        i32_qt: fx.Int32,
+        i32_total_q: fx.Int32,
+        i32_nb: fx.Int32,
+        i32_init: fx.Int32,
+        i32_local: fx.Int32,
+        i32_hs: fx.Int32,
+        i32_grid: fx.Int32,
+        stream: fx.Stream,
+    ):
+        kernel(
+            arg_q,
+            arg_kv,
+            arg_score,
+            arg_bt,
+            arg_cu,
+            arg_seq,
+            i64_page_stride,
+            i32_bt_stride,
+            i32_S,
+            i32_nseg,
+            i32_qt,
+            i32_total_q,
+            i32_nb,
+            i32_init,
+            i32_local,
+            i32_hs,
+        ).launch(grid=(fx.Int64(i32_grid), 1, 1), block=(64 * NW, 1, 1), stream=stream)
+
+    launch.kernel_name = name
+    return launch
diff --git a/vllm/models/minimax_m3/amd/ops/nvfp4_prefill/__init__.py b/vllm/models/minimax_m3/amd/ops/nvfp4_prefill/__init__.py
new file mode 100644
index 000000000000..208f01a7cb5e
--- /dev/null
+++ b/vllm/models/minimax_m3/amd/ops/nvfp4_prefill/__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/minimax_m3/amd/ops/nvfp4_prefill/host.py b/vllm/models/minimax_m3/amd/ops/nvfp4_prefill/host.py
new file mode 100644
index 000000000000..27052b662673
--- /dev/null
+++ b/vllm/models/minimax_m3/amd/ops/nvfp4_prefill/host.py
@@ -0,0 +1,102 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Launch wrapper of the FlyDSL NVFP4 sparse prefill attention."""
+
+import functools
+
+import numpy as np
+import torch
+
+from vllm.models.minimax_m3.amd.ops.flydsl_utils import (
+    LOG2E,
+    cu_count,
+    run_compiled,
+    unit_scale,
+)
+
+from . import kernel as _kernel
+from .kernel import BLK, HG, TOPK, D
+
+
+@functools.cache
+def _get_kernel(rounds: int):
+    return _kernel.compile_nvfp4_prefill_fp16(rounds)
+
+
+def supports_nvfp4_sparse_prefill(
+    num_heads: int, num_kv_heads: int, head_dim: int, topk_blocks: int, block_size: int
+) -> bool:
+    from vllm.platforms.rocm import on_gfx950
+
+    return (
+        on_gfx950()
+        and head_dim == D
+        and num_heads == num_kv_heads * HG
+        and topk_blocks == TOPK
+        and block_size == BLK
+    )
+
+
+def minimax_m3_nvfp4_sparse_prefill(
+    q: torch.Tensor,  # [T, H, 128] bf16
+    kv_cache: torch.Tensor,  # [pages, 2 * Hkv, 128, 72] uint8
+    topk: torch.Tensor,  # [Hkv, T, 16] int32 logical blocks, -1 = none
+    block_table: torch.Tensor,  # [reqs, max_blocks] int32
+    query_req_id: torch.Tensor,  # [T] int32
+    query_abs_pos: torch.Tensor,  # [T] int32
+    num_kv_heads: int,
+    sm_scale: float,
+    out: torch.Tensor,  # [T, H, 128] bf16
+    k_scale: torch.Tensor | None = None,
+    v_scale: torch.Tensor | None = None,
+) -> None:
+    """Token t attends to the causal tokens of its ``topk`` blocks."""
+    n_tok, n_heads, head_dim = q.shape
+    if n_tok == 0:
+        return
+    assert head_dim == D and n_heads == num_kv_heads * HG
+    assert q.dtype == torch.bfloat16 and out.dtype == torch.bfloat16
+    assert out.shape == q.shape
+    assert q.is_contiguous() and out.is_contiguous()
+    assert kv_cache.dtype == torch.uint8 and tuple(kv_cache.shape[1:]) == (
+        2 * num_kv_heads,
+        BLK,
+        72,
+    )
+    assert kv_cache.stride(1) == BLK * 72 and kv_cache.stride(3) == 1
+    assert topk.dtype == torch.int32 and tuple(topk.shape) == (
+        num_kv_heads,
+        n_tok,
+        TOPK,
+    )
+    assert topk.stride(2) == 1 and topk.stride(1) == TOPK
+    for t in (block_table, query_req_id, query_abs_pos):
+        assert t.dtype == torch.int32 and t.stride(-1) == 1
+    assert kv_cache.stride(0) < 2**31
+    unit = unit_scale(q.device)
+    ks = unit if k_scale is None else k_scale
+    vs = unit if v_scale is None else v_scale
+    c = float(np.float32(sm_scale) * np.float32(LOG2E))
+    # Up to one workgroup per CU the one-round combine is ~1 us faster.
+    n_wg = (n_tok + 3) // 4 * num_kv_heads
+    rounds = 2 if n_wg > cu_count(q.device) else 1
+    run_compiled(
+        _get_kernel(rounds),
+        q.data_ptr(),
+        kv_cache.data_ptr(),
+        topk.data_ptr(),
+        block_table.data_ptr(),
+        query_req_id.data_ptr(),
+        query_abs_pos.data_ptr(),
+        out.data_ptr(),
+        ks.data_ptr(),
+        vs.data_ptr(),
+        n_tok,
+        n_heads,
+        num_kv_heads,
+        topk.stride(0),
+        block_table.stride(0),
+        kv_cache.stride(0),
+        c,
+        torch.cuda.current_stream(q.device),
+    )
diff --git a/vllm/models/minimax_m3/amd/ops/nvfp4_prefill/kernel.py b/vllm/models/minimax_m3/amd/ops/nvfp4_prefill/kernel.py
new file mode 100644
index 000000000000..bedf654f12c4
--- /dev/null
+++ b/vllm/models/minimax_m3/amd/ops/nvfp4_prefill/kernel.py
@@ -0,0 +1,507 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""NVFP4 sparse prefill attention for MiniMax-M3 (FlyDSL, gfx950).
+
+A workgroup takes 4 query tokens x the 16 query heads of one KV head and walks
+the union of the pages they selected. K and V are the exact e2m1 x e4m3
+products in fp16; Q and P are fp16; scores, softmax and accumulation are fp32;
+the output is bf16.
+"""
+
+import flydsl.compiler as flyc
+import flydsl.expr as fx
+from aiter.ops.flydsl.kernels import buffer_ops
+from flydsl.expr import const_expr, range_constexpr
+from flydsl.expr.typing import T
+from flydsl.expr.typing import Vector as Vec
+
+from vllm.models.minimax_m3.amd.ops.flydsl_utils import (
+    cvt_fp4x2_f16,
+    cvt_fp8x2_f16,
+    exp2,
+    global_i32_ptr,
+    maxf_nn,
+    mfma_16x16x32,
+    pk_bf16,
+    pk_f16,
+    xlane4,
+)
+
+BLK = 128  # tokens per page
+D = 128
+HG = 16  # query heads per KV head (the MFMA N dimension)
+NW = 4
+KPW = BLK // NW  # keys per wave
+TOK_PER_WG = 4
+TOPK = 16
+MAXU = TOK_PER_WG * TOPK  # union capacity: one lane per (token, slot)
+DATA_B = BLK * D // 2  # e2m1 bytes per head per side
+SCALE_B = BLK * D // 16  # e4m3 scale bytes per head per side
+NDT = D // 16  # d-tiles of O^T
+NKS = D // 32  # k-steps of S^T
+RAW = 2 * 4 + 1 + 8 + 2  # NVFP4 dwords a lane loads per page
+Q_ROW = D * 2 // 16 + 1  # 16 B LDS tiles per Q^T row, padded by one
+
+
+def _bf16_to_f16(v):
+    return Vec(v).bitcast(fx.BFloat16).to(fx.Float32).to(fx.Float16).bitcast(fx.Int32)
+
+
+def kv_lane_loader(arg_cache, page_bytes, kvh, hkv, w, n, g):
+    """The loader of this lane's NVFP4 K/V dwords of a page.
+
+    S^T tile st row r is key 32 w + 8 (r // 4) + 4 st + r % 4, so lane
+    (n, g)'s P^T / V^T key slots are keys 32 w + 8 g + j.
+    """
+    side_b = hkv * (DATA_B + SCALE_B)
+    k_data = kvh * DATA_B
+    k_scl = hkv * DATA_B + kvh * SCALE_B
+    v_data = side_b + kvh * DATA_B
+    v_scl = side_b + hkv * DATA_B + kvh * SCALE_B
+    k_rows = [w * KPW + 8 * (n // 4) + 4 * st + n % 4 for st in range_constexpr(2)]
+    k_load = [(k_data + r * (D // 2) + g * 16) // 4 for r in k_rows]
+    k_scale_load = (k_scl + ((w * 4 + g) * 16 + n) * 4) // 4
+    v_load = (v_data + ((w * KPW + 8 * g) // 8 * 16 + n) * 32) // 4
+    v_scale_load = (v_scl + (n // 2) * BLK + w * KPW + 8 * g) // 4
+
+    def load(page):
+        """This lane's RAW NVFP4 dwords of ``page``."""
+        pr = buffer_ops.create_buffer_resource_from_addr(
+            arg_cache + fx.Int64(page) * fx.Int64(page_bytes),
+            num_records_bytes=page_bytes,
+        )
+        raw = []
+        for st in range_constexpr(2):
+            kd = fx.Vector(
+                buffer_ops.buffer_load(pr, k_load[st], vec_width=4, dtype=fx.Int32)
+            )
+            raw += [kd[i] for i in range_constexpr(4)]
+        raw.append(
+            fx.Int32(
+                buffer_ops.buffer_load(pr, k_scale_load, vec_width=1, dtype=fx.Int32)
+            )
+        )
+        for h in range_constexpr(2):
+            vd = fx.Vector(
+                buffer_ops.buffer_load(pr, v_load + 4 * h, vec_width=4, dtype=fx.Int32)
+            )
+            raw += [vd[i] for i in range_constexpr(4)]
+        vs = fx.Vector(
+            buffer_ops.buffer_load(pr, v_scale_load, vec_width=2, dtype=fx.Int32)
+        )
+        raw += [vs[i] for i in range_constexpr(2)]
+        return raw
+
+    return load
+
+
+def compile_nvfp4_prefill_fp16(rounds: int = 2):
+    """Build the prefill kernel.
+
+    Args:
+        rounds: rounds of the cross-wave combine. 2 halves its LDS so two
+            workgroups share a CU; 1 keeps Q^T in registers and combines all
+            tokens at once (grids of at most one workgroup per CU).
+
+    """
+    assert rounds in (1, 2)
+    tpr = TOK_PER_WG // rounds  # tokens per combine round
+    wpe = rounds  # waves per SIMD
+    name = f"m3_nvfp4_sparse_prefill_fp16matched_tq{TOK_PER_WG}_ksplit"
+    if rounds == 1:
+        name += "_r1"
+
+    @fx.struct
+    class Shared:
+        # Q^T during the page loop, then the partial O per (wave, token, d-tile,
+        # lane) in 16 B tiles; (max, sum) per (wave, token, lane).
+        o: fx.Array[fx.Int32, NW * tpr * NDT * 64 * 4, 16]
+        ml: fx.Array[fx.Float32, NW * TOK_PER_WG * 2 * 64]
+
+    @flyc.kernel(name=name, known_block_size=[64 * NW, 1, 1])
+    def kernel(
+        arg_q: fx.Int64,  # [T, H, 128] bf16
+        arg_cache: fx.Int64,  # [pages, 2 * Hkv, 128, 72] uint8
+        arg_topk: fx.Int64,  # [Hkv, T, 16] int32 logical blocks, -1 = none
+        arg_bt: fx.Int64,  # [reqs, stride] int32
+        arg_req: fx.Int64,  # [T] int32
+        arg_pos: fx.Int64,  # [T] int32
+        arg_out: fx.Int64,  # [T, H, 128] bf16
+        arg_ks: fx.Int64,  # fp32 global K scale
+        arg_vs: fx.Int64,  # fp32 global V scale
+        i32_T: fx.Int32,
+        i32_H: fx.Int32,
+        i32_hkv: fx.Int32,
+        i32_topk_hs: fx.Int32,  # topk stride over KV heads
+        i32_bt_s: fx.Int32,
+        i32_page_byte: fx.Int32,
+        f32_c: fx.Float32,  # sm_scale * log2(e)
+    ):
+        smem = fx.SharedAllocator().allocate(Shared).peek()
+        tx = fx.thread_idx.x
+        bidx, kvh = fx.block_idx.x, fx.block_idx.y
+        lane = tx % 64
+        w = fx.Int32(fx.rocdl.readfirstlane(T.i32, tx // 64))
+        n, g = lane % 16, lane // 16
+        topk = global_i32_ptr(arg_topk)
+        bt = global_i32_ptr(arg_bt)
+        req_p = global_i32_ptr(arg_req)
+        pos_p = global_i32_ptr(arg_pos)
+        f32p = fx.PointerType.get(
+            T.f32, address_space=fx.AddressSpace.Global, alignment=4
+        )
+        c = f32_c * fx.Float32(fx.inttoptr(f32p, fx.Int64(arg_ks))[0])
+        v_scale = fx.Float32(fx.inttoptr(f32p, fx.Int64(arg_vs))[0])
+        neg_inf = fx.Float32(float("-inf"))
+        topk_base = kvh * i32_topk_hs
+
+        lds = fx.logical_divide(
+            fx.make_view(
+                fx.recast_iter(fx.Int32, smem.o.ptr),
+                fx.make_layout(NW * tpr * NDT * 64 * 4, 1),
+            ),
+            fx.make_layout(4, 1),
+        )
+        lds_atom = fx.make_copy_atom(fx.UniversalCopy128b(), fx.Int32)
+        qr = buffer_ops.create_buffer_resource_from_addr(
+            arg_q, num_records_bytes=fx.Int64(i32_T) * fx.Int64(i32_H) * (D * 2)
+        )
+        if const_expr(rounds == 2):
+            # Q^T in LDS; thread (token, head, quarter) stages 64 B.
+            tq, hq, qq = tx // 64, (tx // 4) % HG, tx % 4
+            tok_q = bidx * TOK_PER_WG + tq
+            tok_q = (tok_q < i32_T).select(tok_q, fx.Int32(0))
+            q_src = (tok_q * i32_H + kvh * HG + hq) * (D // 2) + qq * 16
+            for i in range_constexpr(4):
+                r = fx.make_rmem_tensor(fx.make_layout(4, 1), fx.Int32)
+                q4 = buffer_ops.buffer_load(
+                    qr, q_src + 4 * i, vec_width=4, dtype=fx.Int32
+                )
+                r.store(_bf16_to_f16(q4))
+                q_tile = (tq * HG + hq) * Q_ROW + qq * 4 + i
+                fx.copy(lds_atom, r, fx.slice(lds, (None, q_tile)))
+
+        def q_frags(t):
+            """Token t's Q^T, one B operand per k-step."""
+            if const_expr(rounds == 1):
+                return q_regs[t]
+            out = []
+            for ks in range_constexpr(NKS):
+                r = fx.make_rmem_tensor(fx.make_layout(4, 1), fx.Int32)
+                q_tile = (t * HG + n) * Q_ROW + g * 4 + ks
+                fx.copy(lds_atom, fx.slice(lds, (None, q_tile)), r)
+                out.append(Vec(r.load()))
+            return out
+
+        # Per token: the page of its own block and its causal limit.
+        diag, lim, q_regs = [], [], []
+        for t in range_constexpr(TOK_PER_WG):
+            tok = bidx * TOK_PER_WG + t
+            ok = tok < i32_T
+            tc = ok.select(tok, fx.Int32(0))
+            pos = fx.Int32(fx.rocdl.readfirstlane(T.i32, fx.Int32(pos_p[tc])))
+            req = fx.Int32(fx.rocdl.readfirstlane(T.i32, fx.Int32(req_p[tc])))
+            diag.append(
+                fx.Int32(
+                    fx.rocdl.readfirstlane(
+                        T.i32, fx.Int32(bt[req * i32_bt_s + pos // BLK])
+                    )
+                )
+            )
+            lim.append(pos % BLK - w * KPW)
+            if const_expr(rounds == 1):
+                q_row = (tc * i32_H + kvh * HG + n) * (D // 2) + g * 16
+                q_regs.append(
+                    [
+                        _bf16_to_f16(
+                            buffer_ops.buffer_load(
+                                qr, q_row + ks * 4, vec_width=4, dtype=fx.Int32
+                            )
+                        )
+                        for ks in range_constexpr(NKS)
+                    ]
+                )
+
+        load_page = kv_lane_loader(arg_cache, i32_page_byte, kvh, i32_hkv, w, n, g)
+
+        # Lane (token lane // 16, slot lane % 16): the page it selected, -1 none.
+        lane_tok = bidx * TOK_PER_WG + lane // TOPK
+        lane_tok_ok = lane_tok < i32_T
+        lane_tok_clamped = lane_tok_ok.select(lane_tok, fx.Int32(0))
+        blk = fx.Int32(topk[topk_base + lane_tok_clamped * TOPK + lane % TOPK])
+        lane_blk_ok = lane_tok_ok & (blk >= 0)
+        lane_page = fx.Int32(
+            bt[
+                fx.Int32(req_p[lane_tok_clamped]) * i32_bt_s
+                + lane_blk_ok.select(blk, fx.Int32(0))
+            ]
+        )
+        lane_page = lane_blk_ok.select(lane_page, fx.Int32(-1))
+        # The union's first page is the lowest valid lane's: load it before the
+        # dedup runs.
+        valid = fx.Int64(fx.rocdl.ballot(T.i64, lane_page >= 0))
+        first = (valid != fx.Int64(0)).select(
+            fx.Int32(fx.math.cttz(valid).to(fx.Int32)), fx.Int32(0)
+        )
+
+        def page_at(mask):
+            """Page of the lowest lane left in ``mask``, the first page when empty."""
+            idx = (mask != fx.Int64(0)).select(
+                fx.Int32(fx.math.cttz(mask).to(fx.Int32)), first
+            )
+            pg = fx.Int32(fx.rocdl.readlane(T.i32, lane_page, idx))
+            return (pg >= 0).select(pg, fx.Int32(0))
+
+        raw0 = load_page(page_at(valid))
+        fx.rocdl.sched_barrier(0)
+        # Union: the first lane of each distinct page.
+        n_dup = fx.Int32(0)
+        for j in range_constexpr(MAXU - 1):
+            pj = fx.Int32(fx.rocdl.readlane(T.i32, lane_page, j))
+            n_dup = n_dup + ((lane_page == pj) & (lane > j)).select(
+                fx.Int32(1), fx.Int32(0)
+            )
+        umask = fx.Int64(fx.rocdl.ballot(T.i64, (lane_page >= 0) & (n_dup == 0)))
+        n_u = fx.Int32(fx.math.ctpop(umask).to(fx.Int32))
+
+        def dequant(raw):
+            """K frags kf[st][ks] and V^T frags vf[dt] of a page in fp16."""
+            kf = []
+            for st in range_constexpr(2):
+                ksc = cvt_fp8x2_f16(raw[8], st == 1)
+                row = []
+                for ks in range_constexpr(NKS):
+                    sv = Vec.from_elements([ksc[ks // 2], ksc[ks // 2]], fx.Float16)
+                    row.append(
+                        Vec.from_elements(
+                            [
+                                (cvt_fp4x2_f16(raw[st * 4 + ks], i) * sv).bitcast(
+                                    fx.Int32
+                                )[0]
+                                for i in range_constexpr(4)
+                            ],
+                            fx.Int32,
+                        )
+                    )
+                kf.append(row)
+            vs = []
+            for h in range_constexpr(2):
+                vs += [
+                    cvt_fp8x2_f16(raw[17 + h], False),
+                    cvt_fp8x2_f16(raw[17 + h], True),
+                ]
+            # A V byte holds two consecutive keys: already the operand's fp16 pair.
+            vf = [
+                Vec.from_elements(
+                    [
+                        (cvt_fp4x2_f16(raw[9 + dt], i) * vs[i]).bitcast(fx.Int32)[0]
+                        for i in range_constexpr(4)
+                    ],
+                    fx.Int32,
+                )
+                for dt in range_constexpr(NDT)
+            ]
+            return kf, vf
+
+        zero4 = Vec.filled(4, 0.0, fx.Float32)
+        key_off = [
+            8 * g + 4 * st + e for st in range_constexpr(2) for e in range_constexpr(4)
+        ]
+        if const_expr(rounds == 2):
+            fx.gpu.barrier()  # Q^T staged
+        # Iteration u dequantizes page u and loads page u + 1.
+        init = [umask]
+        init += [neg_inf for _ in range_constexpr(TOK_PER_WG)] + [
+            fx.Float32(0.0) for _ in range_constexpr(TOK_PER_WG)
+        ]
+        init += [zero4 for _ in range_constexpr(TOK_PER_WG * NDT)]
+        init += raw0
+        n_st = 1 + 2 * TOK_PER_WG + TOK_PER_WG * NDT
+        results = init
+        for u, st_ in range(  # type: ignore[call-overload]
+            fx.Index(0),
+            fx.Index(n_u),
+            fx.Index(1),
+            init=[fx.as_ir_value(x) for x in init],
+        ):
+            rem = fx.Int64(st_[0])
+            m = [fx.Float32(st_[1 + t]) for t in range_constexpr(TOK_PER_WG)]
+            l_run = [
+                fx.Float32(st_[1 + TOK_PER_WG + t]) for t in range_constexpr(TOK_PER_WG)
+            ]
+            o = [
+                [
+                    Vec(st_[1 + 2 * TOK_PER_WG + t * NDT + dt])
+                    for dt in range_constexpr(NDT)
+                ]
+                for t in range_constexpr(TOK_PER_WG)
+            ]
+            raw = [fx.Int32(st_[n_st + i]) for i in range_constexpr(RAW)]
+            page = page_at(rem)
+            rem = rem & (rem - fx.Int64(1))
+            raw_next = load_page(page_at(rem))
+            kf, vf = dequant(raw)
+            hits = fx.Int64(fx.rocdl.ballot(T.i64, lane_page == page))
+            for t in range_constexpr(TOK_PER_WG):
+                m_t, l_t, o_t = m[t], l_run[t], o[t]
+                if (hits.shrui(fx.Int64(TOPK * t)) & fx.Int64(0xFFFF)) != fx.Int64(0):
+                    qf = q_frags(t)
+                    a = []
+                    for st in range_constexpr(2):
+                        acc = zero4
+                        for ks in range_constexpr(NKS):
+                            acc = mfma_16x16x32(kf[st][ks], qf[ks], acc, True)
+                        a += [acc[e] for e in range_constexpr(4)]
+                    if page == diag[t]:
+                        a = [
+                            (key_off[i] <= lim[t]).select(a[i], neg_inf)
+                            for i in range_constexpr(8)
+                        ]
+                    mx = a[0]
+                    for i in range_constexpr(1, 8):
+                        mx = maxf_nn(mx, a[i])
+                    # The running max is kept in scaled units (c > 0).
+                    m_new = maxf_nn(m_t, xlane4(mx, maxf_nn) * c)
+                    m_use = (m_new == neg_inf).select(fx.Float32(0.0), m_new)
+                    alpha = exp2(m_t - m_use)
+                    p = [exp2(fx.math.fma(x, c, -m_use)) for x in a]
+                    ps = p[0]
+                    for i in range_constexpr(1, 8):
+                        ps = ps + p[i]
+                    l_t = l_t * alpha + ps
+                    pf = Vec.from_elements(
+                        [pk_f16(p[2 * j], p[2 * j + 1]) for j in range_constexpr(4)],
+                        fx.Int32,
+                    )
+                    # alpha is exactly 1 unless some row's max grew.
+                    if fx.Int64(fx.rocdl.ballot(T.i64, m_new != m_t)) != fx.Int64(0):
+                        av = Vec.filled(4, alpha, fx.Float32)
+                        o_t = [x * av for x in o_t]
+                    o_t = [
+                        mfma_16x16x32(vf[dt], pf, o_t[dt], True)
+                        for dt in range_constexpr(NDT)
+                    ]
+                    m_t = m_new
+                m[t], l_run[t], o[t] = m_t, l_t, o_t
+            state = (
+                [rem] + m + l_run + [x for token_out in o for x in token_out] + raw_next
+            )
+            results = yield [fx.as_ir_value(x) for x in state]
+
+        # Combine the waves' partials, tpr tokens per round; wave w finalizes
+        # token w in round w // tpr.
+        ml = smem.ml.ptr
+        if const_expr(rounds == 2):
+            fx.gpu.barrier()
+        for t in range_constexpr(TOK_PER_WG):
+            ml[((w * TOK_PER_WG + t) * 2) * 64 + lane] = fx.Float32(results[1 + t])
+            ml[((w * TOK_PER_WG + t) * 2 + 1) * 64 + lane] = fx.Float32(
+                results[1 + TOK_PER_WG + t]
+            )
+        for rnd in range_constexpr(rounds):
+            for tt in range_constexpr(tpr):
+                for dt in range_constexpr(NDT):
+                    r = fx.make_rmem_tensor(fx.make_layout(4, 1), fx.Int32)
+                    ri = 1 + 2 * TOK_PER_WG + (tpr * rnd + tt) * NDT + dt
+                    r.store(Vec(results[ri]).bitcast(fx.Int32))
+                    fx.copy(
+                        lds_atom,
+                        r,
+                        fx.slice(lds, (None, ((w * tpr + tt) * NDT + dt) * 64 + lane)),
+                    )
+            fx.gpu.barrier()
+            if w // tpr == rnd:
+                t = w
+                ms = [
+                    fx.Float32(ml[((ww * TOK_PER_WG + t) * 2) * 64 + lane])
+                    for ww in range_constexpr(NW)
+                ]
+                big = ms[0]
+                for ww in range_constexpr(1, NW):
+                    big = maxf_nn(big, ms[ww])
+                big = (big == neg_inf).select(fx.Float32(0.0), big)
+                wt = [
+                    fx.Float32(fx.math.exp2(ms[ww] - big)) for ww in range_constexpr(NW)
+                ]
+                l_tot = fx.Float32(0.0)
+                for ww in range_constexpr(NW):
+                    l_w = ml[((ww * TOK_PER_WG + t) * 2 + 1) * 64 + lane]
+                    l_tot = l_tot + fx.Float32(l_w) * wt[ww]
+                l_tot = xlane4(l_tot, lambda a, b: a + b)
+                inv = (l_tot > fx.Float32(0.0)).select(v_scale / l_tot, fx.Float32(0.0))
+                acc = [zero4 for _ in range_constexpr(NDT)]
+                for ww in range_constexpr(NW):
+                    wv = Vec.filled(4, wt[ww], fx.Float32)
+                    for dt in range_constexpr(NDT):
+                        r = fx.make_rmem_tensor(fx.make_layout(4, 1), fx.Int32)
+                        o_tile = ((ww * tpr + t - tpr * rnd) * NDT + dt) * 64 + lane
+                        fx.copy(lds_atom, fx.slice(lds, (None, o_tile)), r)
+                        acc[dt] = acc[dt] + Vec(r.load()).bitcast(fx.Float32) * wv
+                # Lane (n, g) holds dims [32 g, 32 g + 32) of head n in order.
+                tok = bidx * TOK_PER_WG + t
+                words = [
+                    [
+                        pk_bf16(acc[2 * j][e] * inv, acc[2 * j + 1][e] * inv)
+                        for j in range_constexpr(4)
+                    ]
+                    for e in range_constexpr(4)
+                ]
+                outr = buffer_ops.create_buffer_resource_from_addr(
+                    arg_out,
+                    num_records_bytes=fx.Int64(i32_T) * fx.Int64(i32_H) * (D * 2),
+                )
+                o_row = (tok * i32_H + kvh * HG + n) * (D // 2) + g * 16
+                for e in range_constexpr(4):
+                    val = Vec.from_elements(words[e], fx.Int32)
+                    buffer_ops.buffer_store(val, outr, o_row + e * 4, mask=tok < i32_T)
+            fx.gpu.barrier()
+
+    @flyc.jit
+    def launch(
+        arg_q: fx.Int64,
+        arg_cache: fx.Int64,
+        arg_topk: fx.Int64,
+        arg_bt: fx.Int64,
+        arg_req: fx.Int64,
+        arg_pos: fx.Int64,
+        arg_out: fx.Int64,
+        arg_ks: fx.Int64,
+        arg_vs: fx.Int64,
+        i32_T: fx.Int32,
+        i32_H: fx.Int32,
+        i32_hkv: fx.Int32,
+        i32_topk_hs: fx.Int32,
+        i32_bt_s: fx.Int32,
+        i32_page_b: fx.Int32,
+        f32_c: fx.Float32,
+        stream: fx.Stream,
+    ):
+        kernel(
+            arg_q,
+            arg_cache,
+            arg_topk,
+            arg_bt,
+            arg_req,
+            arg_pos,
+            arg_out,
+            arg_ks,
+            arg_vs,
+            i32_T,
+            i32_H,
+            i32_hkv,
+            i32_topk_hs,
+            i32_bt_s,
+            i32_page_b,
+            f32_c,
+            value_attrs={
+                "rocdl.waves_per_eu": wpe,
+                "rocdl.flat_work_group_size": "256,256",
+            },
+        ).launch(
+            grid=((i32_T + TOK_PER_WG - 1) // TOK_PER_WG, i32_hkv, 1),
+            block=(64 * NW, 1, 1),
+            stream=stream,
+        )
+
+    launch.kernel_name = name
+    return launch
diff --git a/vllm/models/minimax_m3/amd/ops/nvfp4_sparse_decode/__init__.py b/vllm/models/minimax_m3/amd/ops/nvfp4_sparse_decode/__init__.py
new file mode 100644
index 000000000000..208f01a7cb5e
--- /dev/null
+++ b/vllm/models/minimax_m3/amd/ops/nvfp4_sparse_decode/__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/minimax_m3/amd/ops/nvfp4_sparse_decode/host.py b/vllm/models/minimax_m3/amd/ops/nvfp4_sparse_decode/host.py
new file mode 100644
index 000000000000..d266c4e9d00c
--- /dev/null
+++ b/vllm/models/minimax_m3/amd/ops/nvfp4_sparse_decode/host.py
@@ -0,0 +1,182 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Launch wrapper of the FlyDSL NVFP4 sparse decode."""
+
+import ctypes
+import functools
+
+import numpy as np
+import torch
+
+from vllm.logger import init_logger
+from vllm.models.minimax_m3.amd.ops.flydsl_utils import (
+    LOG2E,
+    cu_count,
+    run_compiled,
+    unit_scale,
+)
+from vllm.models.minimax_m3.amd.ops.nvfp4_index_prefill.host import nvfp4_e4m3_lut
+
+from . import kernel as _kernel
+
+logger = init_logger(__name__)
+
+_HIP_DEVICE_ATTRIBUTE_NUMBER_OF_XCCS = 10018
+_SUPPORTED_XCD_COUNT = 8
+
+MAX_SPLITS = 8
+# Launches with more (row, KV head) workgroup rows than this do not split.
+MAX_SPLIT_ROWS = 1024
+
+
+@functools.cache
+def _get_kernel(splits: int, topk: int, xcds: int):
+    return _kernel.compile_nvfp4_sparse_decode(splits, topk, xcds)
+
+
+def _query_xcd_count(device_index: int) -> int:
+    """The device's XCD count, or zero when HIP cannot report it."""
+    try:
+        get_attribute = ctypes.CDLL("libamdhip64.so").hipDeviceGetAttribute
+        get_attribute.argtypes = [
+            ctypes.POINTER(ctypes.c_int),
+            ctypes.c_int,
+            ctypes.c_int,
+        ]
+        get_attribute.restype = ctypes.c_int
+        value = ctypes.c_int()
+        status = get_attribute(
+            ctypes.byref(value), _HIP_DEVICE_ATTRIBUTE_NUMBER_OF_XCCS, device_index
+        )
+        return value.value if status == 0 and value.value > 0 else 0
+    except (AttributeError, OSError, TypeError, ctypes.ArgumentError):
+        return 0
+
+
+@functools.cache
+def _xcd_count(device_index: int) -> int:
+    """8 on gfx950 (the topology the XCD grid was tuned on), else 0."""
+    from vllm.platforms.rocm import on_gfx950
+
+    queried = _query_xcd_count(device_index)
+    count = queried if on_gfx950() and queried == _SUPPORTED_XCD_COUNT else 0
+    logger.info(
+        "MiniMax-M3 NVFP4 sparse decode: XCD-aware grid %s on device %d "
+        "(reported XCDs=%d)",
+        "on" if count else "off",
+        device_index,
+        queried,
+    )
+    return count
+
+
+@functools.cache
+def _split_buffers(device: torch.device):
+    """Split partials and tickets; allocated in the warm-up, before capture."""
+    part = torch.empty(
+        MAX_SPLIT_ROWS * MAX_SPLITS * 64 * _kernel.PART_F,
+        dtype=torch.float32,
+        device=device,
+    )
+    return part, torch.zeros(MAX_SPLIT_ROWS, dtype=torch.int32, device=device)
+
+
+def decode_splits(num_rows: int, num_blocks: int, cus: int) -> int:
+    """Splits of a row's blocks.
+
+    The kernel is VALU-bound, so the cost model is rounds of workgroups per CU
+    times blocks per split plus about one block of fixed cost; the fewest splits
+    among the minima, each split at least two blocks.
+    """
+    best_cost, best = None, 1
+    for s in range(1, MAX_SPLITS + 1):
+        per = -(-num_blocks // s)
+        if s > 1 and per < 2:
+            break
+        cost = -(-(num_rows * s) // cus) * (per + 1)
+        if best_cost is None or cost < best_cost:
+            best_cost, best = cost, s
+    return best
+
+
+def minimax_m3_nvfp4_sparse_decode_fly(
+    q: torch.Tensor,  # [rows, H, 128] bf16
+    kv_cache: torch.Tensor,  # [pages, 2 * Hkv, 128, 72] uint8
+    topk_idx: torch.Tensor,  # [Hkv, rows, topk] int32 logical blocks, -1 = none
+    block_table: torch.Tensor,  # [reqs, max_blocks] int32
+    seq_lens: torch.Tensor,  # [reqs] int32
+    sm_scale: float,
+    output: torch.Tensor,  # [rows, H, 128] bf16
+    k_scale: torch.Tensor | None = None,
+    v_scale: torch.Tensor | None = None,
+    decode_query_len: int = 1,
+    splits: int | None = None,
+    xcd_count: int | None = None,
+) -> None:
+    """Sparse decode attention over the NVFP4 K/V pages.
+
+    Row i is query row i % decode_query_len of request i // decode_query_len
+    and attends to the tokens of its ``topk_idx`` blocks up to its position.
+    """
+    rows, num_heads, head_dim = q.shape
+    if rows == 0:
+        return
+    num_kv_heads = kv_cache.shape[1] // 2
+    topk = topk_idx.shape[-1]
+    assert rows == seq_lens.shape[0] * decode_query_len
+    assert head_dim == _kernel.D and num_heads == num_kv_heads * _kernel.HG
+    assert q.dtype == torch.bfloat16 and output.dtype == torch.bfloat16
+    assert q.is_contiguous() and output.is_contiguous()
+    assert kv_cache.dtype == torch.uint8 and kv_cache.stride(1) == _kernel.BLK * 72
+    assert kv_cache.stride(3) == 1 and kv_cache.stride(0) < 2**31
+    assert topk_idx.dtype == torch.int32 and topk_idx.shape[0] == num_kv_heads
+    assert topk_idx.shape[1] >= rows
+    assert topk_idx.stride(2) == 1 and topk_idx.stride(1) == topk
+    for t in (block_table, seq_lens):
+        assert t.dtype == torch.int32 and t.stride(-1) == 1
+    if xcd_count is None:
+        xcd_count = _xcd_count(q.device.index)
+    xcds = xcd_count if xcd_count and xcd_count > 1 else 0
+    if xcds:
+        group = xcds * decode_query_len
+        grid_x = -(-rows // group) * group
+    else:
+        grid_x = rows
+    if splits is None:
+        wg_rows = grid_x * num_kv_heads
+        splits = (
+            1
+            if wg_rows > MAX_SPLIT_ROWS
+            else decode_splits(rows * num_kv_heads, topk, cu_count(q.device))
+        )
+    assert splits == 1 or grid_x * num_kv_heads <= MAX_SPLIT_ROWS
+    part, cnt = _split_buffers(q.device)
+    unit = unit_scale(q.device)
+    ks = unit if k_scale is None else k_scale
+    vs = unit if v_scale is None else v_scale
+    c = float(np.float32(sm_scale) * np.float32(LOG2E))
+    run_compiled(
+        _get_kernel(splits, topk, xcds),
+        q.data_ptr(),
+        kv_cache.data_ptr(),
+        topk_idx.data_ptr(),
+        block_table.data_ptr(),
+        seq_lens.data_ptr(),
+        output.data_ptr(),
+        ks.data_ptr(),
+        vs.data_ptr(),
+        nvfp4_e4m3_lut(q.device).data_ptr(),
+        part.data_ptr(),
+        cnt.data_ptr(),
+        part.numel() * 4,
+        rows,
+        decode_query_len,
+        num_heads,
+        num_kv_heads,
+        topk_idx.stride(0),
+        block_table.stride(0),
+        kv_cache.stride(0),
+        c,
+        grid_x,
+        torch.cuda.current_stream(q.device),
+    )
diff --git a/vllm/models/minimax_m3/amd/ops/nvfp4_sparse_decode/kernel.py b/vllm/models/minimax_m3/amd/ops/nvfp4_sparse_decode/kernel.py
new file mode 100644
index 000000000000..96188426f25c
--- /dev/null
+++ b/vllm/models/minimax_m3/amd/ops/nvfp4_sparse_decode/kernel.py
@@ -0,0 +1,576 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""NVFP4 sparse decode attention for MiniMax-M3 (FlyDSL, gfx950).
+
+One workgroup = one query row x the 16 query heads of one KV head x one split of
+the row's top-k blocks. K and V are dequantized in registers to
+e4m3_rn_satfinite(e2m1 x block scale), the values AITER's FP8 sparse PA reads,
+and widened to bf16; scores and softmax are fp32, P is bf16. Wave w owns keys
+[32 w, 32 w + 32) of every block:
+
+  S^T = K Q^T      v_mfma_f32_16x16x32_bf16, A = K, B = Q^T
+  O^T += V^T P^T   A = V^T, B = P^T straight from the S^T accumulators
+
+The waves combine through LDS. With ``splits`` > 1 every split stores its
+unnormalized partial and the last split of a row to arrive combines them in
+split order, so the result does not depend on arrival order.
+"""
+
+import flydsl.compiler as flyc
+import flydsl.expr as fx
+from aiter.ops.flydsl.kernels import buffer_ops
+from flydsl._mlir import ir as _ir
+from flydsl._mlir.dialects import arith as _arith
+from flydsl._mlir.dialects import llvm as _llvm
+from flydsl.expr import const_expr, range_constexpr
+from flydsl.expr.typing import T
+from flydsl.expr.typing import Vector as Vec
+
+from vllm.models.minimax_m3.amd.ops.flydsl_utils import (
+    PRELOAD_HINTS,
+    SC1,
+    as_f32,
+    bits,
+    cvt_fp4x2_f16,
+    cvt_fp8x2_f16,
+    exp2,
+    global_i32_ptr,
+    maxf_nn,
+    mfma_16x16x32,
+    perm,
+    pk_bf16,
+    xlane4,
+)
+from vllm.models.minimax_m3.amd.ops.nvfp4_prefill.kernel import (
+    BLK,
+    HG,
+    KPW,
+    NDT,
+    NKS,
+    NW,
+    RAW,
+    D,
+    kv_lane_loader,
+)
+
+PART_F = 36  # fp32 per (split, row, lane) partial: O (NDT x 4), max, sum, pad
+SIGN4 = 0x80808080 - (1 << 32)
+MAG4 = 0x07070707
+
+
+def _nnan():
+    return _ir.Attribute.parse("#arith.fastmath<nnan>")
+
+
+def _bf16x2_e4m3(word, hi):
+    """E4M3 bytes 0, 1 (``hi``: 2, 3) of ``word`` -> one dword of two bf16."""
+    return Vec(
+        fx.rocdl.cvt_scalef32_pk_bf16_fp8(
+            Vec.make_type(2, fx.BFloat16),
+            fx.as_ir_value(fx.Int32(word)),
+            fx.as_ir_value(fx.Float32(1.0)),
+            hi,
+        )
+    ).bitcast(fx.Int32)[0]
+
+
+def _k_bf16x8(codes, tab_lo, tab_hi):
+    """Eight K codes under one block scale -> four dwords of bf16 pairs.
+
+    ``tab_lo`` / ``tab_hi`` hold the e4m3 keys of magnitudes 0-3 / 4-7 for the
+    block scale (``nvfp4_e4m3_lut``).
+    """
+    even = perm(tab_hi, tab_lo, codes & MAG4) ^ ((codes << 4) & SIGN4)
+    odd = perm(tab_hi, tab_lo, (codes >> 4) & MAG4) ^ (codes & SIGN4)
+    lo = perm(odd, even, 0x05010400)
+    hi = perm(odd, even, 0x07030602)
+    return [
+        _bf16x2_e4m3(lo, False),
+        _bf16x2_e4m3(lo, True),
+        _bf16x2_e4m3(hi, False),
+        _bf16x2_e4m3(hi, True),
+    ]
+
+
+def _v_bf16x8(codes, scale_pairs):
+    """One V dword (eight keys of one dim) -> four dwords of bf16 pairs.
+
+    The fp16 products are exact; they are clamped to +-448 and rounded once to
+    e4m3.
+    """
+    big = Vec.filled(2, 448.0, fx.Float16)
+    small = Vec.filled(2, -448.0, fx.Float16)
+    f8 = [None, None]
+    for j in range_constexpr(4):
+        v = cvt_fp4x2_f16(codes, j) * scale_pairs[j]
+        v = _arith.MinNumFOp(
+            fx.as_ir_value(v), fx.as_ir_value(big), fastmath=_nnan()
+        ).result
+        v = _arith.MaxNumFOp(v, fx.as_ir_value(small), fastmath=_nnan()).result
+        old = Vec.filled(2, 0, fx.Int16) if j % 2 == 0 else f8[j // 2]
+        f8[j // 2] = Vec(
+            fx.rocdl.cvt_scalef32_pk_fp8_f16(
+                Vec.make_type(2, fx.Int16),
+                fx.as_ir_value(old),
+                v,
+                fx.as_ir_value(fx.Float32(1.0)),
+                j % 2 == 1,
+            )
+        )
+    out = []
+    for h in range_constexpr(2):
+        word = f8[h].bitcast(fx.Int32)[0]
+        out += [_bf16x2_e4m3(word, False), _bf16x2_e4m3(word, True)]
+    return out
+
+
+def _global_atomic_add(addr_i64, value):
+    """Relaxed agent-scope atomic add on a global i32; returns the old value."""
+    return fx.Int32(
+        _llvm.AtomicRMWOp(
+            _llvm.AtomicBinOp.add,
+            global_i32_ptr(addr_i64).llvm_ptr,
+            fx.Int32(value).ir_value(),
+            _llvm.AtomicOrdering.monotonic,
+            syncscope="agent",
+            alignment=4,
+        ).result
+    )
+
+
+def compile_nvfp4_sparse_decode(splits: int = 1, topk: int = 16, xcds: int = 0):
+    """Build the decode kernel.
+
+    Args:
+        splits: workgroups per (row, KV head), each over a contiguous share of
+            the row's blocks.
+        topk: top-k width (one lane per slot).
+        xcds: if > 0, workgroup w runs on XCD w % xcds and a request's rows and
+            a row's splits are placed on one XCD.
+
+    """
+    assert 1 <= topk <= 64 and splits >= 1
+    SPLITS, TOPK, XCDS = splits, topk, xcds
+    # Split partials need sc1 unless all splits of a row run on one XCD.
+    PART_SCOPE = 0 if XCDS else SC1
+    name = f"m3_nvfp4_sparse_decode_k{TOPK}_s{SPLITS}" + (f"_x{XCDS}" if XCDS else "")
+
+    @fx.struct
+    class Shared:
+        o: fx.Array[fx.Int32, NW * NDT * 64 * 4, 16]
+        ml: fx.Array[fx.Float32, NW * 2 * 64]
+        lut: fx.Array[fx.Int32, 512]
+        ticket: fx.Array[fx.Int32, 4]
+
+    @flyc.kernel(name=name, known_block_size=[64 * NW, 1, 1])
+    def kernel(
+        i32_rows: fx.Int32,
+        i32_ql: fx.Int32,  # rows per request
+        i32_topk_hs: fx.Int32,  # topk stride over KV heads
+        i32_bt_s: fx.Int32,
+        i32_H: fx.Int32,
+        i32_hkv: fx.Int32,
+        arg_seq: fx.Int64,  # [reqs] int32
+        arg_topk: fx.Int64,  # [Hkv, >= rows, TOPK] int32 blocks, -1 = none
+        arg_bt: fx.Int64,  # [reqs, stride] int32
+        arg_lut: fx.Int64,  # [256] int64 K magnitude tables
+        arg_q: fx.Int64,  # [rows, H, 128] bf16
+        arg_cache: fx.Int64,  # [pages, 2 * Hkv, 128, 72] uint8
+        arg_ks: fx.Int64,  # fp32 global K scale
+        arg_vs: fx.Int64,  # fp32 global V scale
+        arg_out: fx.Int64,  # [rows, H, 128] bf16
+        arg_part: fx.Int64,  # [grid rows * Hkv, S, 64, PART_F] fp32
+        arg_cnt: fx.Int64,  # [grid rows * Hkv] int32 tickets, zero
+        i64_part_bytes: fx.Int64,
+        i32_page_byte: fx.Int32,
+        f32_c: fx.Float32,  # sm_scale * log2(e)
+    ):
+        smem = fx.SharedAllocator().allocate(Shared).peek()
+        tx = fx.thread_idx.x
+        bx, kvh = fx.block_idx.x, fx.block_idx.y
+        lane = tx % 64
+        w = fx.Int32(fx.rocdl.readfirstlane(T.i32, tx // 64))
+        n, g = lane % 16, lane // 16
+        topk = global_i32_ptr(arg_topk)
+        bt = global_i32_ptr(arg_bt)
+        seq_p = global_i32_ptr(arg_seq)
+        f32p = fx.PointerType.get(
+            T.f32, address_space=fx.AddressSpace.Global, alignment=4
+        )
+        c = f32_c * fx.Float32(fx.inttoptr(f32p, fx.Int64(arg_ks))[0])
+        v_scale = fx.Float32(fx.inttoptr(f32p, fx.Int64(arg_vs))[0])
+        neg_inf = fx.Float32(float("-inf"))
+        lut_g = global_i32_ptr(arg_lut)
+        smem.lut.ptr[tx] = fx.Int32(lut_g[tx])
+        smem.lut.ptr[tx + 256] = fx.Int32(lut_g[tx + 256])
+
+        if const_expr(XCDS > 0):
+            # grid.x is padded to XCDS x ql so a request's rows share an XCD.
+            xcd = bx % XCDS
+            slot = bx // XCDS
+            row = ((slot // i32_ql) * XCDS + xcd) * i32_ql + slot % i32_ql
+        else:
+            row = bx
+        row_ok = row < i32_rows
+        rowc = row_ok.select(row, fx.Int32(0))
+        req = rowc // i32_ql
+        pos = fx.Int32(seq_p[req]) - i32_ql + rowc % i32_ql
+        # Padding rows (seq_len 0 or past i32_rows) see nothing.
+        pos = (row_ok & (pos >= 0)).select(pos, fx.Int32(-1))
+        pos = fx.Int32(fx.rocdl.readfirstlane(T.i32, pos))
+        req = fx.Int32(fx.rocdl.readfirstlane(T.i32, req))
+        self_blk = (pos >= 0).select(pos // BLK, fx.Int32(-1))
+
+        # Lane = top-k slot.
+        slot_ok = lane < TOPK
+        blk = fx.Int32(
+            topk[kvh * i32_topk_hs + rowc * TOPK + slot_ok.select(lane, fx.Int32(0))]
+        )
+        ok = slot_ok & (blk >= 0) & (blk <= self_blk)
+        page = fx.Int32(bt[req * i32_bt_s + ok.select(blk, fx.Int32(0))])
+        page = ok.select(page, fx.Int32(-1))
+        blk = ok.select(blk, fx.Int32(-1))
+        vmask = fx.Int64(fx.rocdl.ballot(T.i64, ok))
+        n_v = fx.Int32(fx.math.ctpop(vmask).to(fx.Int32))
+        if const_expr(SPLITS > 1):
+            below = (fx.Int64(1) << fx.Int64(lane)) - fx.Int64(1)
+            rank = fx.Int32(fx.math.ctpop(vmask & below).to(fx.Int32))
+            chunk = (n_v + (SPLITS - 1)) // SPLITS
+            lo = fx.block_idx.z * chunk
+            vmask = fx.Int64(
+                fx.rocdl.ballot(T.i64, ok & (rank >= lo) & (rank < lo + chunk))
+            )
+            n_v = fx.Int32(fx.math.ctpop(vmask).to(fx.Int32))
+
+        lim = pos % BLK - w * KPW
+        qr = buffer_ops.create_buffer_resource_from_addr(
+            arg_q, num_records_bytes=fx.Int64(i32_rows) * fx.Int64(i32_H) * (D * 2)
+        )
+        q_row = (rowc * i32_H + kvh * HG + n) * (D // 2) + g * 16
+        qf = [
+            fx.Vector(
+                buffer_ops.buffer_load(qr, q_row + ks * 4, vec_width=4, dtype=fx.Int32)
+            )
+            for ks in range_constexpr(NKS)
+        ]
+
+        load_block = kv_lane_loader(arg_cache, i32_page_byte, kvh, i32_hkv, w, n, g)
+
+        def dequant_k(raw):
+            kf = []
+            for st in range_constexpr(2):
+                tabs = []
+                for grp in range_constexpr(2):
+                    sb = raw[8].shrui(fx.Int32(8 * (2 * st + grp))) & fx.Int32(0xFF)
+                    tabs.append(
+                        (
+                            fx.Int32(smem.lut.ptr[sb * 2]),
+                            fx.Int32(smem.lut.ptr[sb * 2 + 1]),
+                        )
+                    )
+                kf.append(
+                    [
+                        Vec.from_elements(
+                            _k_bf16x8(raw[st * 4 + ks], *tabs[ks // 2]), fx.Int32
+                        )
+                        for ks in range_constexpr(NKS)
+                    ]
+                )
+            return kf
+
+        def dequant_v(raw):
+            sp = [
+                cvt_fp8x2_f16(raw[17 + j // 2], j % 2 == 1) for j in range_constexpr(4)
+            ]
+            return [
+                Vec.from_elements(_v_bf16x8(raw[9 + dt], sp), fx.Int32)
+                for dt in range_constexpr(NDT)
+            ]
+
+        fx.gpu.barrier()  # LUT in LDS
+
+        zero4 = Vec.filled(4, 0.0, fx.Float32)
+        key_off = [
+            8 * g + 4 * st + e for st in range_constexpr(2) for e in range_constexpr(4)
+        ]
+        first = (n_v > 0).select(
+            fx.Int32(fx.math.cttz(vmask).to(fx.Int32)), fx.Int32(0)
+        )
+
+        def slot_at(mask):
+            """The lowest slot left in ``mask``, the first one when empty."""
+            return (mask != fx.Int64(0)).select(
+                fx.Int32(fx.math.cttz(mask).to(fx.Int32)), first
+            )
+
+        def page_of(slot):
+            pg = fx.Int32(fx.rocdl.readlane(T.i32, page, slot))
+            return (pg >= 0).select(pg, fx.Int32(0))
+
+        def attend(kf, vf, cur_blk, live, m_t, l_t, o_t):
+            """Fold one block into (max, sum, O); a dead block is fully masked."""
+            a = []
+            for st in range_constexpr(2):
+                acc = zero4
+                for ks in range_constexpr(NKS):
+                    acc = mfma_16x16x32(kf[st][ks], qf[ks], acc, False)
+                a += [acc[e] for e in range_constexpr(4)]
+            diag = cur_blk == self_blk
+            dead = live == fx.Int32(0)
+            a = [
+                (dead | (diag & (key_off[i] > lim))).select(neg_inf, a[i])
+                for i in range_constexpr(8)
+            ]
+            mx = a[0]
+            for i in range_constexpr(1, 8):
+                mx = maxf_nn(mx, a[i])
+            # The running max is kept in scaled units (c > 0).
+            m_new = maxf_nn(m_t, xlane4(mx, maxf_nn) * c)
+            m_use = (m_new == neg_inf).select(fx.Float32(0.0), m_new)
+            alpha = exp2(m_t - m_use)
+            p = [exp2(fx.math.fma(x, c, -m_use)) for x in a]
+            ps = p[0]
+            for i in range_constexpr(1, 8):
+                ps = ps + p[i]
+            pf = Vec.from_elements(
+                [pk_bf16(p[2 * j], p[2 * j + 1]) for j in range_constexpr(4)],
+                fx.Int32,
+            )
+            av = Vec.filled(4, alpha, fx.Float32)
+            o_t = [
+                mfma_16x16x32(vf[dt], pf, o_t[dt] * av, False)
+                for dt in range_constexpr(NDT)
+            ]
+            return m_new, l_t * alpha + ps, o_t
+
+        # Two blocks per iteration, the next pair's loads in flight; an odd last
+        # block is masked.
+        init = [vmask, neg_inf, fx.Float32(0.0)]
+        init += [zero4 for _ in range_constexpr(NDT)]
+        init += load_block(page_of(slot_at(vmask)))
+        n_st = 3 + NDT
+        results = init
+        for u, st_ in range(  # type: ignore[call-overload]
+            fx.Index(0),
+            fx.Index((n_v + 1) // 2),
+            fx.Index(1),
+            init=[fx.as_ir_value(x) for x in init],
+        ):
+            rem = fx.Int64(st_[0])
+            m_t = fx.Float32(st_[1])
+            l_t = fx.Float32(st_[2])
+            o_t = [Vec(st_[3 + dt]) for dt in range_constexpr(NDT)]
+            raw_a = [fx.Int32(st_[n_st + i]) for i in range_constexpr(RAW)]
+            blk_a = fx.Int32(fx.rocdl.readlane(T.i32, blk, slot_at(rem)))
+            rem_b = rem & (rem - fx.Int64(1))
+            raw_b = load_block(page_of(slot_at(rem_b)))
+            fx.rocdl.sched_barrier(0)
+            kf_a = dequant_k(raw_a)
+            vf_a = dequant_v(raw_a)
+            rem_n = rem_b & (rem_b - fx.Int64(1))
+            fx.rocdl.sched_barrier(0)
+            raw_n = load_block(page_of(slot_at(rem_n)))
+            fx.rocdl.sched_barrier(0)
+            res = attend(kf_a, vf_a, blk_a, fx.Int32(1), m_t, l_t, o_t)
+            m_t, l_t, o_t = res[0], res[1], res[2]
+            blk_b = fx.Int32(fx.rocdl.readlane(T.i32, blk, slot_at(rem_b)))
+            live_b = (rem_b != fx.Int64(0)).select(fx.Int32(1), fx.Int32(0))
+            kf_b = dequant_k(raw_b)
+            vf_b = dequant_v(raw_b)
+            res = attend(kf_b, vf_b, blk_b, live_b, m_t, l_t, o_t)
+            m_t, l_t, o_t = res[0], res[1], res[2]
+            state = [rem_n, m_t, l_t] + o_t + raw_n
+            results = yield [fx.as_ir_value(x) for x in state]
+
+        def store_out(acc, inv):
+            outr = buffer_ops.create_buffer_resource_from_addr(
+                arg_out,
+                num_records_bytes=fx.Int64(i32_rows) * fx.Int64(i32_H) * (D * 2),
+            )
+            o_row = (row * i32_H + kvh * HG + n) * (D // 2) + g * 16
+            for e in range_constexpr(4):
+                val = Vec.from_elements(
+                    [
+                        pk_bf16(acc[2 * j][e] * inv, acc[2 * j + 1][e] * inv)
+                        for j in range_constexpr(4)
+                    ],
+                    fx.Int32,
+                )
+                buffer_ops.buffer_store(val, outr, o_row + e * 4, mask=row < i32_rows)
+
+        # Combine the waves' partials in wave 0.
+        o_lds = fx.logical_divide(
+            fx.make_view(
+                fx.recast_iter(fx.Int32, smem.o.ptr),
+                fx.make_layout(NW * NDT * 64 * 4, 1),
+            ),
+            fx.make_layout(4, 1),
+        )
+        lds_atom = fx.make_copy_atom(fx.UniversalCopy128b(), fx.Int32)
+        ml = smem.ml.ptr
+        ml[(w * 2) * 64 + lane] = fx.Float32(results[1])
+        ml[(w * 2 + 1) * 64 + lane] = fx.Float32(results[2])
+        for dt in range_constexpr(NDT):
+            r = fx.make_rmem_tensor(fx.make_layout(4, 1), fx.Int32)
+            r.store(Vec(results[3 + dt]).bitcast(fx.Int32))
+            fx.copy(lds_atom, r, fx.slice(o_lds, (None, (w * NDT + dt) * 64 + lane)))
+        fx.gpu.barrier()
+        if w == 0:
+            ms = [fx.Float32(ml[(ww * 2) * 64 + lane]) for ww in range_constexpr(NW)]
+            big = ms[0]
+            for ww in range_constexpr(1, NW):
+                big = maxf_nn(big, ms[ww])
+            big_raw = big
+            big = (big == neg_inf).select(fx.Float32(0.0), big)
+            wt = [fx.Float32(fx.math.exp2(ms[ww] - big)) for ww in range_constexpr(NW)]
+            l_tot = fx.Float32(0.0)
+            for ww in range_constexpr(NW):
+                l_tot = l_tot + fx.Float32(ml[(ww * 2 + 1) * 64 + lane]) * wt[ww]
+            l_tot = xlane4(l_tot, lambda x, y: x + y)
+            acc = [zero4 for _ in range_constexpr(NDT)]
+            for ww in range_constexpr(NW):
+                wv = Vec.filled(4, wt[ww], fx.Float32)
+                for dt in range_constexpr(NDT):
+                    r = fx.make_rmem_tensor(fx.make_layout(4, 1), fx.Int32)
+                    fx.copy(
+                        lds_atom,
+                        fx.slice(o_lds, (None, (ww * NDT + dt) * 64 + lane)),
+                        r,
+                    )
+                    acc[dt] = acc[dt] + Vec(r.load()).bitcast(fx.Float32) * wv
+            if const_expr(SPLITS > 1):
+                tid = row * i32_hkv + kvh
+                pr = buffer_ops.create_buffer_resource_from_addr(
+                    arg_part, num_records_bytes=i64_part_bytes
+                )
+                row_f = (tid * SPLITS + fx.block_idx.z) * 64 + lane
+                for dt in range_constexpr(NDT):
+                    buffer_ops.buffer_store(
+                        acc[dt].bitcast(fx.Int32),
+                        pr,
+                        row_f * PART_F + dt * 4,
+                        cache_modifier=PART_SCOPE,
+                    )
+                buffer_ops.buffer_store(
+                    Vec.from_elements([bits(big_raw), bits(l_tot)], fx.Int32),
+                    pr,
+                    row_f * PART_F + NDT * 4,
+                    cache_modifier=PART_SCOPE,
+                )
+                # The partial is stored before the ticket is taken.
+                fx.rocdl.s_waitcnt(vmcnt=0)
+                if lane == 0:
+                    smem.ticket.ptr[0] = _global_atomic_add(
+                        arg_cnt + fx.Int64(tid) * 4, 1
+                    )
+                fx.rocdl.s_waitcnt(lgkmcnt=0)
+                if (
+                    fx.Int32(
+                        fx.rocdl.readfirstlane(T.i32, fx.Int32(smem.ticket.ptr[0]))
+                    )
+                    == SPLITS - 1
+                ):
+                    rows_f = [
+                        (tid * SPLITS + sp) * 64 + lane
+                        for sp in range_constexpr(SPLITS)
+                    ]
+                    mls = [
+                        Vec(
+                            buffer_ops.buffer_load(
+                                pr,
+                                rf * PART_F + NDT * 4,
+                                vec_width=2,
+                                dtype=fx.Int32,
+                                cache_modifier=PART_SCOPE,
+                            )
+                        )
+                        for rf in rows_f
+                    ]
+                    bigs = [as_f32(x[0]) for x in mls]
+                    top = bigs[0]
+                    for sp in range_constexpr(1, SPLITS):
+                        top = maxf_nn(top, bigs[sp])
+                    top = (top == neg_inf).select(fx.Float32(0.0), top)
+                    ws = [fx.Float32(fx.math.exp2(b - top)) for b in bigs]
+                    l_all = fx.Float32(0.0)
+                    for sp in range_constexpr(SPLITS):
+                        l_all = l_all + as_f32(mls[sp][1]) * ws[sp]
+                    inv_all = (l_all > fx.Float32(0.0)).select(
+                        v_scale / l_all, fx.Float32(0.0)
+                    )
+                    acc_all = [zero4 for _ in range_constexpr(NDT)]
+                    for sp in range_constexpr(SPLITS):
+                        wv = Vec.filled(4, ws[sp], fx.Float32)
+                        for dt in range_constexpr(NDT):
+                            part = Vec(
+                                buffer_ops.buffer_load(
+                                    pr,
+                                    rows_f[sp] * PART_F + dt * 4,
+                                    vec_width=4,
+                                    dtype=fx.Int32,
+                                    cache_modifier=PART_SCOPE,
+                                )
+                            ).bitcast(fx.Float32)
+                            acc_all[dt] = acc_all[dt] + part * wv
+                    store_out(acc_all, inv_all)
+                    if lane == 0:
+                        global_i32_ptr(arg_cnt + fx.Int64(tid) * 4)[0] = fx.Int32(0)
+            else:
+                inv = (l_tot > fx.Float32(0.0)).select(v_scale / l_tot, fx.Float32(0.0))
+                store_out(acc, inv)
+
+    @flyc.jit
+    def launch(
+        arg_q: fx.Int64,
+        arg_cache: fx.Int64,
+        arg_topk: fx.Int64,
+        arg_bt: fx.Int64,
+        arg_seq: fx.Int64,
+        arg_out: fx.Int64,
+        arg_ks: fx.Int64,
+        arg_vs: fx.Int64,
+        arg_lut: fx.Int64,
+        arg_part: fx.Int64,
+        arg_cnt: fx.Int64,
+        i64_part_bytes: fx.Int64,
+        i32_rows: fx.Int32,
+        i32_ql: fx.Int32,
+        i32_H: fx.Int32,
+        i32_hkv: fx.Int32,
+        i32_topk_hs: fx.Int32,
+        i32_bt_s: fx.Int32,
+        i32_page_b: fx.Int32,
+        f32_c: fx.Float32,
+        i32_grid_x: fx.Int32,
+        stream: fx.Stream,
+    ):
+        kernel(
+            i32_rows,
+            i32_ql,
+            i32_topk_hs,
+            i32_bt_s,
+            i32_H,
+            i32_hkv,
+            arg_seq,
+            arg_topk,
+            arg_bt,
+            arg_lut,
+            arg_q,
+            arg_cache,
+            arg_ks,
+            arg_vs,
+            arg_out,
+            arg_part,
+            arg_cnt,
+            i64_part_bytes,
+            i32_page_b,
+            f32_c,
+        ).launch(
+            grid=(i32_grid_x, i32_hkv, SPLITS),
+            block=(64 * NW, 1, 1),
+            stream=stream,
+        )
+
+    launch.kernel_name = name
+    launch.compile_hints = PRELOAD_HINTS
+    return launch
diff --git a/vllm/models/minimax_m3/amd/sparse_attention_msa.py b/vllm/models/minimax_m3/amd/sparse_attention_msa.py
index a2c49d13abf0..44ad045db0d9 100644
--- a/vllm/models/minimax_m3/amd/sparse_attention_msa.py
+++ b/vllm/models/minimax_m3/amd/sparse_attention_msa.py
@@ -126,7 +126,6 @@ def forward(
         kv_cache: torch.Tensor,
         output: torch.Tensor,
         *,
-        query_fp8: torch.Tensor | None = None,
         decode_sparse_table: tuple[torch.Tensor, torch.Tensor] | None = None,
     ) -> torch.Tensor:
         from vllm.models.minimax_m3.amd.ops.sparse_pa import (
diff --git a/vllm/models/minimax_m3/amd/sparse_attention_nvfp4.py b/vllm/models/minimax_m3/amd/sparse_attention_nvfp4.py
new file mode 100644
index 000000000000..72f816956438
--- /dev/null
+++ b/vllm/models/minimax_m3/amd/sparse_attention_nvfp4.py
@@ -0,0 +1,170 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""MiniMax M3 block-sparse attention over an NVFP4 K/V cache on ROCm (gfx950).
+
+The cache spec is NVIDIA's, but the bytes inside a page are laid out for the
+FlyDSL kernels, so pages are not interchangeable with NVIDIA's.
+"""
+
+from dataclasses import replace
+from typing import ClassVar
+
+import torch
+
+from vllm.config.cache import CacheDType
+from vllm.forward_context import get_forward_context
+from vllm.models.minimax_m3.common.sparse_attention import (
+    MiniMaxM3SparseBackend,
+    MiniMaxM3SparseImpl,
+    MiniMaxM3SparseMetadata,
+    MiniMaxM3SparseMetadataBuilder,
+    minimax_m3_query_token_positions,
+)
+from vllm.utils.torch_utils import get_dtype_size, nvfp4_kv_cache_full_dim
+from vllm.v1.attention.backend import AttentionLayer, CommonAttentionMetadata
+from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheLayout
+
+
+class MiniMaxM3SparseRocmNvfp4MetadataBuilder(MiniMaxM3SparseMetadataBuilder):
+    """Adds each prefill token's request and absolute position."""
+
+    def build(
+        self,
+        common_prefix_len: int,
+        common_attn_metadata: CommonAttentionMetadata,
+        fast_build: bool = False,
+    ) -> MiniMaxM3SparseMetadata:
+        metadata = super().build(common_prefix_len, common_attn_metadata, fast_build)
+        p = metadata.prefill
+        if p is not None and p.query_req_id is None:
+            p.query_req_id, p.query_abs_pos = minimax_m3_query_token_positions(
+                p.cu_seqlens_q, p.context_lens, metadata.num_prefill_tokens
+            )
+        return metadata
+
+
+class MiniMaxM3SparseRocmNvfp4Backend(MiniMaxM3SparseBackend):
+    """MiniMax M3 sparse attend over NVFP4 K/V pages (ROCm, gfx950)."""
+
+    supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = ["nvfp4"]
+
+    @staticmethod
+    def get_builder_cls() -> type["MiniMaxM3SparseRocmNvfp4MetadataBuilder"]:
+        return MiniMaxM3SparseRocmNvfp4MetadataBuilder
+
+    @classmethod
+    def customize_spec(cls, spec: AttentionSpec) -> AttentionSpec:
+        """Store each head's K and V as two slots of packed data + block scales."""
+        if spec.state_content_bytes is not None or not spec.kv_quant_mode.is_nvfp4:
+            return spec
+        return replace(
+            spec,
+            num_head_slots=2 * spec.num_kv_heads,
+            state_content_bytes=nvfp4_kv_cache_full_dim(spec.head_size)
+            * get_dtype_size(spec.dtype),
+        )
+
+    @classmethod
+    def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...] | None:
+        # The kernels read the page stride from stride(0) and address pages in
+        # 64 bits, so block-outer packing works too (opt in with BLHNC).
+        return (KVCacheLayout.LBHNC, KVCacheLayout.BLHNC)
+
+
+def rocm_nvfp4_unsupported_reason(
+    num_heads: int, num_kv_heads: int, head_dim: int, topk_blocks: int, block_size: int
+) -> str | None:
+    """Why the ROCm NVFP4 attend cannot serve this layer, None when it can."""
+    from vllm.models.minimax_m3.amd.ops.nvfp4_prefill.host import (
+        supports_nvfp4_sparse_prefill,
+    )
+
+    if not supports_nvfp4_sparse_prefill(
+        num_heads, num_kv_heads, head_dim, topk_blocks, block_size
+    ):
+        return (
+            "needs gfx950, head_dim 128, 16 query heads per KV head and 16 "
+            f"selected 128-token blocks (got num_heads={num_heads}, "
+            f"num_kv_heads={num_kv_heads}, head_dim={head_dim}, "
+            f"topk_blocks={topk_blocks}, block_size={block_size})"
+        )
+    return None
+
+
+class MiniMaxM3SparseRocmNvfp4Impl(MiniMaxM3SparseImpl):
+    """FlyDSL NVFP4 sparse decode (bf16 compute) and prefill (fp16 compute)."""
+
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        reason = rocm_nvfp4_unsupported_reason(
+            self.num_heads,
+            self.num_kv_heads,
+            self.head_size,
+            self.topk_blocks,
+            self.block_size,
+        )
+        if reason is not None:
+            raise ValueError(f"MiniMax M3 NVFP4 KV cache on ROCm {reason}")
+
+    def forward(
+        self,
+        layer: AttentionLayer,
+        query: torch.Tensor,
+        kv_cache: torch.Tensor,
+        output: torch.Tensor,
+    ) -> torch.Tensor:
+        from vllm.models.minimax_m3.amd.ops.nvfp4_prefill.host import (
+            minimax_m3_nvfp4_sparse_prefill,
+        )
+        from vllm.models.minimax_m3.amd.ops.nvfp4_sparse_decode.host import (
+            minimax_m3_nvfp4_sparse_decode_fly,
+        )
+
+        attn_metadata = get_forward_context().attn_metadata
+        if not isinstance(attn_metadata, dict):
+            return output  # profiling run; caches unbound
+        md = attn_metadata[layer.layer_name]  # type: ignore[attr-defined]
+        assert isinstance(md, MiniMaxM3SparseMetadata)
+
+        nd = md.num_decode_tokens
+        num_tokens = md.num_actual_tokens
+        topk = layer.topk_indices_buffer  # type: ignore[attr-defined]
+        assert topk is not None
+        q = query[:num_tokens].view(-1, self.num_heads, self.head_size)
+        out = output[:num_tokens].view(-1, self.num_heads, self.head_size)
+        k_scale = layer._k_scale  # type: ignore[attr-defined]
+        v_scale = layer._v_scale  # type: ignore[attr-defined]
+
+        if md.num_decodes > 0:
+            d = md.decode
+            assert d is not None
+            minimax_m3_nvfp4_sparse_decode_fly(
+                q[:nd].contiguous(),
+                kv_cache,
+                topk[:, :nd, :],
+                d.block_table,
+                d.seq_lens,
+                self.scale,
+                out[:nd],
+                k_scale=k_scale,
+                v_scale=v_scale,
+                decode_query_len=d.decode_query_len,
+            )
+        if md.num_prefills > 0:
+            p = md.prefill
+            assert p is not None
+            assert p.query_req_id is not None and p.query_abs_pos is not None
+            minimax_m3_nvfp4_sparse_prefill(
+                q[nd:],
+                kv_cache,
+                topk[:, nd:num_tokens, :],
+                p.block_table,
+                p.query_req_id,
+                p.query_abs_pos,
+                self.num_kv_heads,
+                self.scale,
+                out[nd:],
+                k_scale=k_scale,
+                v_scale=v_scale,
+            )
+        return output
diff --git a/vllm/models/minimax_m3/common/dspark.py b/vllm/models/minimax_m3/common/dspark.py
new file mode 100644
index 000000000000..58fecf2bc353
--- /dev/null
+++ b/vllm/models/minimax_m3/common/dspark.py
@@ -0,0 +1,413 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""MiniMax-M3.1 DSpark draft model.
+
+The Qwen3 DSpark block drafter (context-KV precompute, non-causal query block,
+Markov and confidence heads) built from MiniMax-M3 layers: Gemma RMSNorm
+(``x * (1 + w)``) everywhere, partial RoPE, and a dense SwiGLU-OAI MLP. The
+checkpoint ships no token embedding or LM head; both are shared with the target.
+"""
+
+from collections.abc import Iterable
+
+import torch
+from torch import nn
+
+from vllm import _custom_ops as ops
+from vllm.config import CacheConfig, VllmConfig
+from vllm.model_executor.layers.activation import SiluAndMulWithClamp
+from vllm.model_executor.layers.attention.attention import get_attention_context
+from vllm.model_executor.layers.fusion.quant_activation import get_input_quant_key
+from vllm.model_executor.layers.linear import (
+    MergedColumnParallelLinear,
+    RowParallelLinear,
+)
+from vllm.model_executor.layers.quantization import QuantizationConfig
+from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp8Dynamic
+from vllm.model_executor.models.qwen3_dflash import (
+    DFlashQwen3Attention,
+    DFlashQwen3DecoderLayer,
+    _resolve_layer_attention,
+)
+from vllm.model_executor.models.qwen3_dspark import (
+    Qwen3DSparkForCausalLM,
+    Qwen3DSparkModel,
+)
+from vllm.model_executor.models.utils import WeightsMapper
+from vllm.models.minimax_m3.common.ops.dspark_context_kv import (
+    DSparkContextKVWriter,
+    rotary_dims_supported,
+)
+from vllm.platforms import current_platform
+from vllm.transformers_utils.configs.minimax_m3 import MiniMaxM3TextConfig
+from vllm.utils.torch_utils import (
+    LayerNameType,
+    _encode_layer_name,
+    _resolve_layer_name,
+    direct_register_custom_op,
+)
+
+_CHECKPOINT_PREFIX = "language_model.model.dspark."
+
+
+def _shared_slot_mapping(
+    slot_mapping: torch.Tensor | list[torch.Tensor | None] | None,
+) -> torch.Tensor | None:
+    """The context slot mapping when every draft layer shares it, else None
+    (dummy runs, and drafts whose layers span several KV cache groups)."""
+    if slot_mapping is None or isinstance(slot_mapping, torch.Tensor):
+        return slot_mapping
+    first = slot_mapping[0]
+    if first is None or any(
+        s is None or s.data_ptr() != first.data_ptr() for s in slot_mapping
+    ):
+        return None
+    return first
+
+
+def _gemma_rms_norm(hidden_size: int, eps: float) -> nn.Module:
+    """The target's Gemma RMSNorm (fused FlashInfer kernels on CUDA).
+
+    vLLM's GemmaRMSNorm hands the IR op an fp32 ``1 + w`` weight, which the
+    vllm_c kernel rejects for bf16 inputs, so every call ran the unfused native
+    path (about ten kernels).
+    """
+    if current_platform.is_rocm():
+        from vllm.models.minimax_m3.amd.model import MiniMAXGemmaRMSNorm
+    else:
+        from vllm.models.minimax_m3.nvidia.model import MiniMAXGemmaRMSNorm
+    return MiniMAXGemmaRMSNorm(hidden_size, eps=eps)
+
+
+def dspark_qknorm_rope_kv_insert(
+    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,
+    q_fp8_out: torch.Tensor | None,
+    layer_name: LayerNameType,
+) -> None:
+    """The fused QK-norm/RoPE on ``qkv``, also storing the draft tokens' k/v.
+
+    A custom op so the slot mapping is read when it runs: read in the compiled
+    forward, it is fixed at trace time, when the draft layers have none.
+    """
+    _, layer, kv_cache, slot_mapping = get_attention_context(
+        _resolve_layer_name(layer_name)
+    )
+    insert = slot_mapping is not None
+    ops.fused_minimax_m3_qknorm_rope_kv_insert(
+        qkv,
+        q_norm_weight,
+        k_norm_weight,
+        cos_sin_cache,
+        positions,
+        num_heads,
+        num_kv_heads,
+        rotary_dim,
+        eps,
+        slot_mapping=slot_mapping,
+        kv_cache=kv_cache if insert else None,
+        block_size=kv_cache.size(2) if insert else 0,
+        kv_cache_dtype="fp8" if insert and kv_cache.dtype == torch.uint8 else "auto",
+        q_fp8_out=q_fp8_out,
+        q_fp8_scale=layer._q_scale_float,
+    )
+
+
+def dspark_qknorm_rope_kv_insert_fake(
+    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,
+    q_fp8_out: torch.Tensor | None,
+    layer_name: LayerNameType,
+) -> None:
+    return
+
+
+direct_register_custom_op(
+    op_name="minimax_m3_dspark_qknorm_rope_kv_insert",
+    op_func=dspark_qknorm_rope_kv_insert,
+    mutates_args=["qkv", "q_fp8_out"],
+    fake_impl=dspark_qknorm_rope_kv_insert_fake,
+)
+
+
+class MiniMaxM3DSparkAttention(DFlashQwen3Attention):
+    def __init__(
+        self,
+        config: MiniMaxM3TextConfig,
+        sliding_window: int | None,
+        causal: bool,
+        cache_config: CacheConfig | None,
+        quant_config: QuantizationConfig | None,
+        prefix: str,
+        query_pre_quantized: bool = False,
+    ) -> None:
+        super().__init__(
+            hidden_size=config.hidden_size,
+            num_heads=config.num_attention_heads,
+            num_kv_heads=config.num_key_value_heads,
+            rope_parameters={
+                "rope_theta": config.rope_theta,
+                "partial_rotary_factor": config.partial_rotary_factor,
+            },
+            max_position=config.max_position_embeddings,
+            head_dim=config.head_dim,
+            rms_norm_eps=config.rms_norm_eps,
+            sliding_window=sliding_window,
+            causal=causal,
+            cache_config=cache_config,
+            quant_config=quant_config,
+            prefix=prefix,
+            query_pre_quantized=query_pre_quantized,
+        )
+        self.q_norm = _gemma_rms_norm(self.head_dim, eps=config.rms_norm_eps)
+        self.k_norm = _gemma_rms_norm(self.head_dim, eps=config.rms_norm_eps)
+        # With an FP8 KV cache the backend wants an FP8 query, and the fused
+        # kernel below can write one. Producing it there instead of letting the
+        # attention layer quantize afterwards saves a pass over q.
+        self._emit_fp8_query = query_pre_quantized
+        # FLASHINFER_TS writes an FP8 query's output as FP16 (FA4 only bf16);
+        # an MXFP8 o_proj quantizes FP16 as readily and writes the model dtype.
+        self._attn_output_dtype = (
+            torch.float16
+            if query_pre_quantized
+            and self.attn.attn_backend.get_name() == "FLASHINFER_TS"
+            and get_input_quant_key(self.o_proj) == kMxfp8Dynamic
+            else None
+        )
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+    ) -> torch.Tensor:
+        qkv, _ = self.qkv_proj(hidden_states)
+        # The target's fused per-head Gemma QK-norm + partial NeoX RoPE on q/k,
+        # in place: one launch instead of the q-norm, k-norm and RoPE kernels.
+        # q_fp8_out additionally stores the normed/roped q as E4M3, which is
+        # what an FP8-KV backend consumes.
+        q_fp8 = (
+            torch.empty(
+                (qkv.shape[0], self.q_size),
+                dtype=current_platform.fp8_dtype(),
+                device=qkv.device,
+            )
+            if self._emit_fp8_query
+            else None
+        )
+        # Also stores the draft tokens' k/v, so the attention layer gets q alone.
+        torch.ops.vllm.minimax_m3_dspark_qknorm_rope_kv_insert(
+            qkv,
+            self.q_norm.weight,
+            self.k_norm.weight,
+            self.rotary_emb.cos_sin_cache,
+            positions,
+            self.num_heads,
+            self.num_kv_heads,
+            self.rotary_emb.rotary_dim,
+            self.q_norm.variance_epsilon,
+            q_fp8,
+            _encode_layer_name(self.attn.layer_name),
+        )
+        q = qkv[:, : self.q_size] if q_fp8 is None else q_fp8
+        # output_dtype must be stated: it otherwise follows the query, which
+        # may be FP8.
+        attn_output = self.attn(
+            q, None, None, output_dtype=self._attn_output_dtype or hidden_states.dtype
+        )
+        output, _ = self.o_proj(attn_output)
+        return output
+
+
+class MiniMaxM3DSparkMLP(nn.Module):
+    """Dense SwiGLU-OAI MLP: gate * sigmoid(alpha * gate) * (up + beta)."""
+
+    def __init__(
+        self,
+        config: MiniMaxM3TextConfig,
+        quant_config: QuantizationConfig | None,
+        prefix: str,
+    ) -> None:
+        super().__init__()
+        if config.hidden_act != "swigluoai":
+            raise ValueError(
+                f"Unsupported MiniMax DSpark activation {config.hidden_act!r}; "
+                "only swigluoai is supported."
+            )
+        intermediate_size = config.dspark_ffn_hidden_size
+        self.gate_up_proj = MergedColumnParallelLinear(
+            config.hidden_size,
+            [intermediate_size] * 2,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.gate_up_proj",
+        )
+        self.down_proj = RowParallelLinear(
+            intermediate_size,
+            config.hidden_size,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.down_proj",
+        )
+        self.act_fn = SiluAndMulWithClamp(
+            swiglu_limit=config.swiglu_limit,
+            alpha=config.swiglu_alpha,
+            beta=config.swiglu_beta,
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        gate_up, _ = self.gate_up_proj(x)
+        x, _ = self.down_proj(self.act_fn(gate_up))
+        return x
+
+
+class MiniMaxM3DSparkDecoderLayer(DFlashQwen3DecoderLayer):
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        *,
+        config: MiniMaxM3TextConfig,
+        layer_idx: int,
+        cache_config: CacheConfig | None = None,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        nn.Module.__init__(self)
+        if not getattr(config, "dspark_use_dense_ffn", False):
+            raise NotImplementedError("MiniMax DSpark supports dense FFN drafts only.")
+        self.hidden_size = config.hidden_size
+        sliding_window, causal = _resolve_layer_attention(config, layer_idx)
+        # An FP8 draft cache means the backend reads an FP8 query, which the
+        # fused QK-norm/RoPE kernel emits directly.
+        kv_cache_dtype = getattr(cache_config, "cache_dtype", "auto")
+        self.self_attn = MiniMaxM3DSparkAttention(
+            config,
+            sliding_window=sliding_window,
+            causal=causal,
+            cache_config=cache_config,
+            quant_config=quant_config,
+            prefix=f"{prefix}.self_attn",
+            query_pre_quantized=str(kv_cache_dtype).startswith("fp8"),
+        )
+        self.mlp = MiniMaxM3DSparkMLP(config, quant_config, prefix=f"{prefix}.mlp")
+        hidden_size, eps = config.hidden_size, config.rms_norm_eps
+        self.input_layernorm = _gemma_rms_norm(hidden_size, eps=eps)
+        self.post_attention_layernorm = _gemma_rms_norm(hidden_size, eps=eps)
+
+
+class MiniMaxM3DSparkModel(Qwen3DSparkModel):
+    decoder_layer_cls = MiniMaxM3DSparkDecoderLayer
+    quant_config_mapper = WeightsMapper(
+        orig_to_new_prefix={_CHECKPOINT_PREFIX: "model."}
+    )
+
+    def __init__(
+        self,
+        *,
+        vllm_config: VllmConfig,
+        start_layer_id: int = 0,
+        prefix: str = "",
+    ) -> None:
+        super().__init__(
+            vllm_config=vllm_config, start_layer_id=start_layer_id, prefix=prefix
+        )
+        hidden_size, eps = self.config.hidden_size, self.config.rms_norm_eps
+        self.hidden_norm = _gemma_rms_norm(hidden_size, eps=eps)
+        self.norm = _gemma_rms_norm(hidden_size, eps=eps)
+
+    def _normalize_context_states(self, context_states: torch.Tensor) -> torch.Tensor:
+        return self.hidden_norm(context_states)
+
+    def _normalize_context_k(self, all_k: torch.Tensor) -> torch.Tensor:
+        return torch.stack(
+            [layer.self_attn.k_norm(k) for layer, k in zip(self.layers, all_k)]
+        )
+
+    def _context_kv_writer(self) -> DSparkContextKVWriter | None:
+        """The fused context K/V writer, or None where the base path runs."""
+        if not hasattr(self, "_ctx_kv_writer"):
+            if any(a.kv_cache.numel() == 0 for a in self._attn_layers):
+                return None  # KV caches not bound yet: decide once they are
+            rotary_dim = self._rope_cos_sin_cache.shape[-1]
+            supported = (
+                self._rope_is_neox
+                and self.layers[0].self_attn.v_scale is None
+                and rotary_dims_supported(rotary_dim, self._head_dim)
+                and self._rope_cos_sin_cache.dtype == torch.bfloat16
+                and DSparkContextKVWriter.supported(self._attn_layers)
+            )
+            self._ctx_kv_writer = (
+                DSparkContextKVWriter(self._attn_layers) if supported else None
+            )
+        return self._ctx_kv_writer
+
+    def precompute_and_store_context_kv(
+        self,
+        context_states: torch.Tensor,
+        context_positions: torch.Tensor,
+        context_slot_mapping: torch.Tensor | list[torch.Tensor | None] | None = None,
+    ) -> None:
+        """One kernel norms, ropes and stores the context K/V of every layer,
+        straight from the projection output."""
+        slot_mapping = _shared_slot_mapping(context_slot_mapping)
+        writer = (
+            self._context_kv_writer()
+            if slot_mapping is not None
+            and hasattr(self, "_num_attn_layers")
+            and context_states.dtype == torch.bfloat16
+            else None
+        )
+        if writer is None:
+            super().precompute_and_store_context_kv(
+                context_states, context_positions, context_slot_mapping
+            )
+            return
+        all_kv = self._project_context_kv_token_major(
+            context_states,
+            context_states.shape[0],
+            self._num_attn_layers,
+            self._num_kv_heads,
+            self._head_dim,
+        )
+        writer(
+            all_kv,
+            self._k_norm_weights,
+            self._rms_norm_eps,
+            context_positions,
+            self._rope_cos_sin_cache,
+            self._rope_cos_sin_cache.shape[-1],
+            slot_mapping,
+        )
+
+
+class MiniMaxM3DSparkForCausalLM(Qwen3DSparkForCausalLM):
+    model_cls = MiniMaxM3DSparkModel
+
+    @staticmethod
+    def _rename_checkpoint_weights(
+        weights: Iterable[tuple[str, torch.Tensor]],
+    ) -> Iterable[tuple[str, torch.Tensor]]:
+        for name, tensor in weights:
+            name = name.removeprefix(_CHECKPOINT_PREFIX)
+            name = name.replace(".decoder_layer.", ".")
+            if name.startswith("final_norm."):
+                name = "norm." + name.removeprefix("final_norm.")
+            # MiniMax stores the MXFP8 E8M0 block scales as weight_scale_inv.
+            name = name.replace(".weight_scale_inv", ".weight_scale")
+            yield name, tensor
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
+        return super().load_weights(self._rename_checkpoint_weights(weights))
diff --git a/vllm/models/minimax_m3/common/indexer.py b/vllm/models/minimax_m3/common/indexer.py
index 0d7ee1d2f9ba..c86d1fcd7b33 100644
--- a/vllm/models/minimax_m3/common/indexer.py
+++ b/vllm/models/minimax_m3/common/indexer.py
@@ -115,20 +115,30 @@ def __init__(
         backend_cls: type[AttentionBackend] = MiniMaxM3IndexerBackend,
     ) -> None:
         super().__init__()
+        # Per-token storage width in elements of ``dtype``.
+        self.head_size = head_dim
         if indexer_kv_dtype in ("fp8", "fp8_e4m3"):
-            cache_dtype = torch.float8_e4m3fn
+            cache_dtype = query_dtype = torch.float8_e4m3fn
         elif indexer_kv_dtype == "bf16":
-            cache_dtype = torch.bfloat16
+            cache_dtype = query_dtype = torch.bfloat16
+        elif indexer_kv_dtype == "nvfp4":
+            # Packed NVFP4 pages scored by the MSA Q8KV4 indexer with e4m3 Q.
+            cache_dtype, query_dtype = torch.uint8, torch.float8_e4m3fn
+            self.head_size = head_dim // 2 + head_dim // 16
         else:
             raise NotImplementedError(
                 f"indexer_kv_dtype={indexer_kv_dtype!r} is not supported by the "
-                "MiniMax M3 indexer cache (only 'bf16' or 'fp8'/'fp8_e4m3')."
+                "MiniMax M3 indexer cache (only 'bf16', 'fp8'/'fp8_e4m3' or "
+                "'nvfp4')."
             )
         self.kv_cache = torch.tensor([])
         self.head_dim = head_dim
         self.indexer_kv_dtype = indexer_kv_dtype
-        # Side-cache storage dtype: bf16, or e4m3 for the fp8 score path.
+        # Side-cache storage dtype: bf16, e4m3 for the fp8 score path, or uint8
+        # for packed NVFP4.
         self.dtype = cache_dtype
+        # Index-query dtype the fused qk-norm/rope kernel emits.
+        self.query_dtype = query_dtype
         self.prefix = prefix
         self.cache_config = cache_config
         # Impl-chosen backend -> each impl gets its own builder (get_attn_backend).
@@ -147,7 +157,7 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
         return MLAAttentionSpec(
             block_size=vllm_config.cache_config.block_size,
             num_kv_heads=1,
-            head_size=self.head_dim,
+            head_size=self.head_size,
             dtype=self.dtype,
             cache_role=SparseCacheRole.INDEXER,
         )
@@ -392,6 +402,11 @@ def __init__(
             backend_cls=type(self).indexer_backend_cls,
         )
 
+    def reserve_workspace(self) -> None:
+        """Reserve the workspace-manager scratch ``forward`` needs; called from
+        the memory-profiling run, which skips ``forward``."""
+        return
+
     def forward(
         self,
         index_query: torch.Tensor,
@@ -501,8 +516,9 @@ def select_indexer_impl_cls(
 
     On Blackwell (SM100) with ``topk_blocks == 16`` (the only width fmha_sm100's
     ``sparse_topk_select`` kernel supports), the fmha_sm100 score + top-k path is
-    used for both bf16 and fp8 index caches. Everything else falls back to the
-    Triton indexer, with fp8 restricted to CUDA platforms that advertise fp8
+    used for both bf16 and fp8 index caches. An nvfp4 index cache uses the MSA
+    Q8KV4/Q8KV8 paged indexers (SM100/SM103/SM107 only). Everything else falls back to
+    the Triton indexer, with fp8 restricted to CUDA platforms that advertise fp8
     support.
     """
     if get_current_vllm_config().attention_config.minimax_m3_indexer_cp:
@@ -530,10 +546,9 @@ def select_indexer_impl_cls(
             tp_size,
         )
         return MiniMaxM3IndexerTritonCPImpl
-    if indexer_kv_dtype in ("mxfp4", "nvfp4"):
+    if indexer_kv_dtype == "mxfp4":
         raise NotImplementedError(
-            f"indexer_kv_dtype={indexer_kv_dtype!r} needs the (not-yet-added) "
-            "CuteDSL indexer impl."
+            "indexer_kv_dtype='mxfp4' is not supported by the MiniMax M3 indexer."
         )
     is_cuda = current_platform.is_cuda()
     is_fp8 = indexer_kv_dtype in ("fp8", "fp8_e4m3")
@@ -542,6 +557,28 @@ def select_indexer_impl_cls(
             f"indexer_kv_dtype={indexer_kv_dtype!r} requires CUDA fp8 support."
         )
     is_sm100 = is_cuda and current_platform.is_device_capability_family(100)
+    if indexer_kv_dtype == "nvfp4":
+        if not (
+            is_cuda
+            and (
+                current_platform.is_device_capability(100)
+                or current_platform.is_device_capability(103)
+                or current_platform.is_device_capability(107)
+            )
+        ):
+            raise NotImplementedError(
+                "indexer_kv_dtype='nvfp4' requires an SM100, SM103 or SM107 GPU."
+            )
+        from vllm.models.minimax_m3.nvidia.indexer_nvfp4 import (
+            MiniMaxM3IndexerNvfp4Impl,
+        )
+
+        logger.info_once(
+            "MiniMax M3 indexer: selected NVFP4 (MSA Q8KV4 decode, Q8KV8 prefill) "
+            "[topk_blocks=%d]",
+            topk_blocks,
+        )
+        return MiniMaxM3IndexerNvfp4Impl
     use_msa = (
         is_sm100
         and topk_blocks == 16
@@ -628,6 +665,9 @@ def index_cache(self) -> MiniMaxM3IndexerCache:
     def num_index_heads(self) -> int:
         return self.impl.num_index_heads
 
+    def reserve_workspace(self) -> None:
+        self.impl.reserve_workspace()
+
     def forward(
         self,
         index_query: torch.Tensor,
diff --git a/vllm/models/minimax_m3/common/mm_preprocess.py b/vllm/models/minimax_m3/common/mm_preprocess.py
index 958e4f189b78..602aee361f46 100644
--- a/vllm/models/minimax_m3/common/mm_preprocess.py
+++ b/vllm/models/minimax_m3/common/mm_preprocess.py
@@ -40,24 +40,28 @@
 )
 from vllm.transformers_utils.configs.minimax_m3 import MiniMaxM3Config
 from vllm.transformers_utils.processors.minimax_m3 import (
-    MIN_SHORT_SIDE_PIXEL,
     MiniMaxM3VLImageProcessor,
     MiniMaxM3VLVideoProcessor,
     MiniMaxVLProcessor,
-    smart_resize,
+    compute_aligned_size,
+    resolve_max_long_side_pixel,
 )
 
+# Worst-case long-side caps in pixels, matching the processor's detail=high
+# tiers. Profiling budgets are sized against these; request-level
+# ``detail``/``max_long_side_pixel`` kwargs cannot exceed them.
+_IMAGE_MAX_LONG_SIDE_PIXEL = 3584
+_VIDEO_MAX_LONG_SIDE_PIXEL = 2188
+
 # Upper bound on the number of frames used to build the dummy video during
 # memory profiling. Sized to the worst-case video the processor accepts:
-# ``max_total_pixels // max_pixels_per_frame`` = 301,056,000 // 602,112 = 500
-# frames, each at the video processor's per-frame ``max_pixels`` (768 * 28 * 28
-# = 602,112). This reaches the true worst-case ~192,000 vision tokens, but only
-# because the dummy video is sized via ``get_video_size_with_most_features()``
-# (the video ``max_pixels`` bound), not the smaller image bound. Without a cap,
+# ``max_total_pixels // max_pixels_per_frame`` = 301,056,000 // (2184 * 2184)
+# = 63 frames, each a square at the video long-side cap (2188 px floored to
+# the patch*merge factor, 28). Without a cap,
 # ``_get_max_video_frames(seq_len)`` with M3's large ``max_model_len`` yields
 # ~1400 frames, producing a multi-GB dummy tensor that overflows the
 # multimodal encoder cache.
-_MAX_FRAMES_PER_VIDEO = 500
+_MAX_FRAMES_PER_VIDEO = 63
 
 
 class MiniMaxM3VLProcessingInfo(BaseProcessingInfo):
@@ -101,34 +105,46 @@ def _get_vision_info(
         image_height: int,
         num_frames: int,
         image_processor,
+        mm_kwargs: Mapping[str, object],
+        modality: str,
     ) -> tuple[ImageSize, int]:
         """Compute resized image size and number of vision tokens.
 
-        Mirrors the processor's Qwen-style ``smart_resize`` (area bound by
-        ``max_pixels``) so token counts match the actual processor output.
+        Mirrors the processor's ``compute_aligned_size`` so token counts
+        match the actual processor output.
         """
         patch_size: int = image_processor.patch_size
         merge_size: int = image_processor.merge_size
         temporal_patch_size: int = image_processor.temporal_patch_size
         factor = patch_size * merge_size
-        max_pixels: int = image_processor.max_pixels
-        # Long-side resize spec (opt-in). ``image_processor`` is the *video*
-        # processor when counting video tokens, so read the bounds off it.
-        max_long_side_pixel = getattr(image_processor, "max_long_side_pixel", None)
-        min_short_side_pixel = getattr(
-            image_processor, "min_short_side_pixel", MIN_SHORT_SIDE_PIXEL
+
+        # ``image_processor`` is the *video* processor when counting video
+        # tokens, so modality-scoped request kwargs (flat keys overlaid with
+        # ``images_kwargs``/``videos_kwargs`` via ``modality=``) and processor
+        # defaults are read per modality. Per-key fallback mirrors the
+        # processor, whose own defaults feed ``resolve_max_long_side_pixel``
+        # the same way.
+        merged = self.ctx.get_merged_mm_kwargs(mm_kwargs)
+        # Kwargs scoped to this modality (``images_kwargs`` / ``videos_kwargs``)
+        # override the flat ones.
+        scoped = merged.get("images_kwargs" if modality == "image" else "videos_kwargs")
+        if isinstance(scoped, Mapping):
+            merged = {**merged, **scoped}
+        detail = merged.get("detail", getattr(image_processor, "detail", None))
+        max_long_side_pixel = merged.get(
+            "max_long_side_pixel", getattr(image_processor, "max_long_side_pixel", None)
+        )
+        resolved_long_side = resolve_max_long_side_pixel(
+            detail, max_long_side_pixel, modality=modality
         )
 
-        new_h, new_w = smart_resize(
+        # No volumetric cap here: token counting must not raise; the video
+        # processor enforces its volume cap on the real inputs.
+        new_h, new_w = compute_aligned_size(
             image_height,
             image_width,
-            factor=factor,
-            max_pixels=max_pixels,
-            max_long_side_pixel=max_long_side_pixel,
-            min_short_side_pixel=min_short_side_pixel,
-            # Token counting must not raise; the volumetric/area cap is enforced
-            # in the processor's _preprocess on the real inputs.
-            max_total_pixels=None,
+            factor,
+            max_long_side_pixel=resolved_long_side,
         )
         grid_h = new_h // patch_size
         grid_w = new_w // patch_size
@@ -153,6 +169,8 @@ def get_num_image_tokens(
             image_height=image_height,
             num_frames=1,
             image_processor=image_processor,
+            mm_kwargs=mm_kwargs,
+            modality="image",
         )
         return n
 
@@ -170,48 +188,27 @@ def get_num_video_tokens(
             image_height=image_height,
             num_frames=num_frames,
             image_processor=image_processor,
+            mm_kwargs=mm_kwargs,
+            modality="video",
         )
         return n
 
     def get_image_size_with_most_features(self) -> ImageSize:
-        # Largest square (a multiple of patch_size*merge_size) whose area is
-        # within the image processor's bound — this yields the most vision
-        # tokens for one image. With the long-side spec the square side is
-        # capped by ``max_long_side_pixel`` (and the fixed ``max_total_pixels``);
-        # otherwise it is bound by the ``max_pixels`` area.
+        # Worst case is a detail=high image: a square whose side is the max
+        # long side (3584 px, a multiple of patch_size*merge_size).
         image_processor = self.get_image_processor()
         factor = image_processor.patch_size * image_processor.merge_size
-        max_long_side_pixel = getattr(image_processor, "max_long_side_pixel", None)
-        if max_long_side_pixel is not None:
-            side_px = min(
-                max_long_side_pixel,
-                math.isqrt(image_processor.max_total_pixels),
-            )
-        else:
-            side_px = math.isqrt(image_processor.max_pixels)
-        side = max(factor, (side_px // factor) * factor)
+        side = (_IMAGE_MAX_LONG_SIDE_PIXEL // factor) * factor
         return ImageSize(width=side, height=side)
 
     def get_video_size_with_most_features(self) -> ImageSize:
-        # Per-frame size that yields the most vision tokens, bound by the
-        # *video* processor's ``max_pixels`` (which differs from the image
-        # bound). Token count depends only on area, so maximize the area
-        # achievable with both sides a multiple of patch_size*merge_size rather
-        # than picking the largest square — a square (e.g. 756x756 for M3's
-        # 602,112 bound) leaves area on the table, undercounting frames.
+        # Worst single frame is a detail=high square capped by the video
+        # long-side max (2188 px), floored to a patch_size*merge_size
+        # multiple (2184).
         video_processor = self.get_video_processor()
         factor = video_processor.patch_size * video_processor.merge_size
-        per_frame_pixels = video_processor.max_pixels
-        max_long_side_pixel = getattr(video_processor, "max_long_side_pixel", None)
-        if max_long_side_pixel is not None:
-            # Long-side spec: a frame's worst case is a square capped by
-            # ``max_long_side_pixel`` (per-frame area, not the volumetric cap).
-            per_frame_pixels = min(per_frame_pixels, max_long_side_pixel**2)
-        units = per_frame_pixels // (factor * factor)  # h_u * w_u
-        h_u = math.isqrt(units)
-        while units % h_u:
-            h_u -= 1
-        return ImageSize(width=(units // h_u) * factor, height=h_u * factor)
+        side = (_VIDEO_MAX_LONG_SIDE_PIXEL // factor) * factor
+        return ImageSize(width=side, height=side)
 
     def get_max_image_tokens(self) -> int:
         image_processor = self.get_image_processor()
@@ -220,7 +217,9 @@ def get_max_image_tokens(self) -> int:
             image_width=size.width,
             image_height=size.height,
             image_processor=image_processor,
-            mm_kwargs={},
+            # Budgets must cover the worst case (detail=high); the default
+            # detail=auto cap would undercount.
+            mm_kwargs={"max_long_side_pixel": _IMAGE_MAX_LONG_SIDE_PIXEL},
         )
 
     def _get_max_video_frames(self, max_tokens: int) -> int:
@@ -233,7 +232,8 @@ def _get_max_video_frames(self, max_tokens: int) -> int:
                 image_height=size.height,
                 num_frames=num_frames + 1,
                 image_processor=video_processor,
-                mm_kwargs={},
+                # Worst case (detail=high); see get_max_image_tokens.
+                mm_kwargs={"max_long_side_pixel": _VIDEO_MAX_LONG_SIDE_PIXEL},
             )
             if next_n > max_tokens:
                 break
@@ -265,7 +265,8 @@ def get_max_video_tokens(
             image_height=size.height,
             num_frames=self.get_num_frames_with_most_features(seq_len, mm_counts),
             image_processor=video_processor,
-            mm_kwargs={},
+            # Worst case (detail=high); see get_max_image_tokens.
+            mm_kwargs={"max_long_side_pixel": _VIDEO_MAX_LONG_SIDE_PIXEL},
         )
 
 
@@ -487,6 +488,11 @@ def load_bytes(
         backend: VideoDecoderBackend = "opencv",
         **kwargs,
     ) -> tuple[npt.NDArray, dict[str, Any]]:
+        # fps=-1 (unspecified) falls back to the default 1 fps.
+        if fps == -1:
+            fps = 1
+        if not 0.2 <= fps <= 5:
+            raise ValueError(f"video fps must be in [0.2, 5], got {fps}")
         return super().load_bytes(
             data,
             num_frames=num_frames,
@@ -526,6 +532,15 @@ def compute_frames_index_to_sample(
                 break
             indices.append(target_frame)
             prev_kept_ts = target_frame / video_fps
+
+        # Always keep the last frame when it is far enough (>= 1/fps, up to
+        # eps) from the previously kept frame, matching the SFT frame
+        # sampler's constant-mode behavior.
+        last_frame_idx = total_frames - 1
+        last_ts = last_frame_idx / video_fps
+        if indices and indices[-1] != last_frame_idx and last_ts - prev_kept_ts > eps:
+            indices.append(last_frame_idx)
+
         if not indices:
             indices = [0]
         return indices
diff --git a/vllm/models/minimax_m3/common/ops/dspark_context_kv.py b/vllm/models/minimax_m3/common/ops/dspark_context_kv.py
new file mode 100644
index 000000000000..514f052051d6
--- /dev/null
+++ b/vllm/models/minimax_m3/common/ops/dspark_context_kv.py
@@ -0,0 +1,170 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Fused DSpark context-KV write for all draft layers in one CUDA kernel.
+
+Replaces the permute copy, per-layer K-norm launches + stack, positions.repeat,
+rotary_embedding and per-layer reshape_and_cache_flash calls in
+``DFlashQwen3Model.precompute_and_store_context_kv``. K/V are read straight
+from the fused projection output ``all_kv`` ``[num_ctx, L, 2, nkv, 128]`` and
+written into each layer's paged cache (FLASH_ATTN, FLASHINFER_TS)
+``[num_blocks, nkv, block_size, 2 * 128]`` (K at ``[..., :128]``, V at
+``[..., 128:]``), matching ``reshape_and_cache_flash``: fp8 e4m3 of
+``x / scale`` for kv_cache_dtype fp8, raw bf16 otherwise; slots < 0 are padding
+and are not written.
+
+K gets the Gemma RMSNorm of its layer (``x * rsqrt(mean(x^2) + eps) * (1 + w)``
+in fp32, materialized to bf16 like the unfused flashinfer path) followed by
+partial NeoX RoPE with the token's position (see
+``fused_minimax_m3_dspark_context_kv``). Each layer's cache is a separate
+allocation, addressed as layer 0's cache plus a per-layer element offset, and
+each fp8 scale as layer 0's k scale plus a float offset; the offset tables are
+built once and cached (keyed by the data_ptrs) so the call stays CUDA-graph
+safe.
+"""
+
+import torch
+
+from vllm import _custom_ops as ops
+from vllm.platforms import current_platform
+
+_SUPPORTED_KV_DTYPES = ("auto", "bfloat16", "fp8", "fp8_e4m3")
+_HEAD_DIM = 128
+
+
+def rotary_dims_supported(rotary_dim: int, head_dim: int) -> bool:
+    """The kernel rotates the leading 16, 32 or 64 dims of a 128-dim head."""
+    return head_dim == _HEAD_DIM and rotary_dim in (16, 32, 64)
+
+
+class DSparkContextKVWriter:
+    """Writes precomputed context K/V into every draft layer's KV cache.
+
+    Built from the draft's ``Attention`` layers; the cache- and scale-offset
+    tables are cached, keyed by their data_ptrs, so calls captured in a CUDA
+    graph keep addressing the same memory.
+    """
+
+    def __init__(self, attn_layers: list) -> None:
+        layers = list(attn_layers)
+        error = self._validate(layers)
+        if error is not None:
+            raise ValueError(f"unsupported DSpark context-KV layout: {error}")
+        self._layers = layers
+        cache = layers[0].kv_cache
+        self._num_layers = len(layers)
+        self._num_kv_heads = cache.shape[1]
+        self._kv_cache_dtype = "fp8" if cache.dtype == torch.uint8 else "auto"
+        self._table_key: tuple | None = None
+        self._cache_offsets: torch.Tensor | None = None
+        self._scale_offsets: torch.Tensor | None = None
+
+    @staticmethod
+    def _validate(layers: list) -> str | None:
+        if not layers:
+            return "no layers"
+        ref = None
+        for a in layers:
+            kv_dtype = getattr(a, "kv_cache_dtype", None)
+            if kv_dtype not in _SUPPORTED_KV_DTYPES:
+                return f"kv_cache_dtype {kv_dtype!r}"
+            cache = getattr(a, "kv_cache", None)
+            if not isinstance(cache, torch.Tensor) or cache.numel() == 0:
+                return "KV cache not bound"
+            if cache.dim() != 4 or cache.shape[3] != 2 * _HEAD_DIM:
+                return f"unexpected cache shape {tuple(cache.shape)}"
+            block_size = cache.shape[2]
+            if block_size & (block_size - 1):
+                return f"block size {block_size} is not a power of two"
+            if not cache.is_contiguous():
+                return "cache is not contiguous"
+            if kv_dtype in ("fp8", "fp8_e4m3"):
+                if cache.dtype != torch.uint8:
+                    return f"fp8 cache has dtype {cache.dtype}"
+            elif cache.dtype != torch.bfloat16:
+                return f"cache has dtype {cache.dtype}"
+            for name in ("_k_scale", "_v_scale"):
+                scale = getattr(a, name, None)
+                if (
+                    not isinstance(scale, torch.Tensor)
+                    or scale.dtype != torch.float32
+                    or not scale.is_cuda
+                ):
+                    return f"{name} is not a CUDA fp32 tensor"
+            key = (cache.shape, cache.stride(), cache.dtype, kv_dtype)
+            if ref is None:
+                ref = key
+            elif key != ref:
+                return "layers do not share cache shape/stride/dtype"
+            # Caches are addressed from layer 0's as 16-element aligned offsets.
+            delta = cache.data_ptr() - layers[0].kv_cache.data_ptr()
+            if delta % (16 * cache.element_size()):
+                return "layer caches are not 16-element aligned to each other"
+        return None
+
+    @staticmethod
+    def supported(attn_layers: list) -> bool:
+        # The fp8 store is OCP E4M3; ROCm fnuz caches keep the base path.
+        if not current_platform.is_cuda() or not hasattr(
+            torch.ops._C, "fused_minimax_m3_dspark_context_kv"
+        ):
+            return False
+        return DSparkContextKVWriter._validate(list(attn_layers)) is None
+
+    def _ensure_tables(self) -> None:
+        layers = self._layers
+        key = tuple(
+            (a.kv_cache.data_ptr(), a._k_scale.data_ptr(), a._v_scale.data_ptr())
+            for a in layers
+        )
+        if key == self._table_key:
+            return
+        cache0, scale0 = layers[0].kv_cache, layers[0]._k_scale
+
+        def table(offsets: list[int]) -> torch.Tensor:
+            return torch.tensor(offsets, dtype=torch.int64, device=cache0.device)
+
+        self._cache_offsets = table(
+            [
+                (a.kv_cache.data_ptr() - cache0.data_ptr()) // cache0.element_size()
+                for a in layers
+            ]
+        )
+        # [2, L]: k scales, then v scales, as float offsets from layer 0's k.
+        self._scale_offsets = table(
+            [
+                (getattr(a, name).data_ptr() - scale0.data_ptr()) // 4
+                for name in ("_k_scale", "_v_scale")
+                for a in layers
+            ]
+        )
+        self._table_key = key
+
+    def __call__(
+        self,
+        all_kv: torch.Tensor,
+        k_norm_weight: torch.Tensor,
+        eps: float,
+        positions: torch.Tensor,
+        cos_sin_cache: torch.Tensor,
+        rotary_dim: int,
+        slot_mapping: torch.Tensor,
+    ) -> None:
+        num_ctx, L, _, nkv, _ = all_kv.shape
+        assert self._num_layers == L and self._num_kv_heads == nkv
+        if num_ctx == 0:
+            return
+        self._ensure_tables()
+        ops.fused_minimax_m3_dspark_context_kv(
+            all_kv,
+            k_norm_weight,
+            cos_sin_cache,
+            positions,
+            slot_mapping,
+            self._layers[0].kv_cache,
+            self._cache_offsets,
+            self._layers[0]._k_scale,
+            self._scale_offsets,
+            rotary_dim,
+            eps,
+            self._kv_cache_dtype,
+        )
diff --git a/vllm/models/minimax_m3/common/ops/index_topk.py b/vllm/models/minimax_m3/common/ops/index_topk.py
index 2d612c12f254..bc24a9a3e285 100644
--- a/vllm/models/minimax_m3/common/ops/index_topk.py
+++ b/vllm/models/minimax_m3/common/ops/index_topk.py
@@ -703,15 +703,19 @@ def minimax_m3_index_score(
     )
     BLOCK_SIZE_Q = 64
     n_q_tiles = triton.cdiv(max_query_len, BLOCK_SIZE_Q)
-    SCORE_TARGET_GRID = 48
+    # Splits write disjoint score columns, preserving the unsplit result.
+    if current_platform.is_cuda() and current_platform.is_device_capability((12, 0)):
+        score_target_grid = 48
+    elif current_platform.is_cuda() and current_platform.is_device_capability_family(
+        90
+    ):
+        score_target_grid = 1056
+    else:
+        score_target_grid = 1
     split_k = max(
         1,
-        min(max_block, SCORE_TARGET_GRID // max(1, n_q_tiles * batch * num_idx_heads)),
+        min(max_block, score_target_grid // max(1, n_q_tiles * batch * num_idx_heads)),
     )
-    if not (
-        current_platform.is_cuda() and current_platform.is_device_capability((12, 0))
-    ):
-        split_k = 1
     grid_score = (n_q_tiles, batch * num_idx_heads, split_k)
     _index_block_score_kernel[grid_score](
         idx_q,
diff --git a/vllm/models/minimax_m3/common/sparse_attention.py b/vllm/models/minimax_m3/common/sparse_attention.py
index 73ab067eb94f..4cc28f0ffd1e 100644
--- a/vllm/models/minimax_m3/common/sparse_attention.py
+++ b/vllm/models/minimax_m3/common/sparse_attention.py
@@ -264,7 +264,7 @@ def __init__(
     ) -> None:
         super().__init__(kv_cache_spec, layer_names, vllm_config, device)
         self._init_reorder_batch_threshold(1, supports_spec_as_decode=True)
-        # Stable context-length buffer for decode cudagraph replays.
+        # Stable context-length buffer for the prefill metadata.
         self.context_len_buffer = torch.empty(
             vllm_config.scheduler_config.max_num_batched_tokens,
             dtype=torch.int32,
@@ -306,14 +306,13 @@ def build(
         assert num_decodes + num_prefills == num_reqs
         assert num_decode_tokens + num_prefill_tokens == num_tokens
 
-        # Decode-first batch: context lengths into the stable cudagraph buffer.
-        context_lens = self.context_len_buffer[:num_reqs]
-        context_lens.copy_(
-            common_attn_metadata.compute_num_computed_tokens(), non_blocking=True
-        )
-
         prefill_metadata: MiniMaxM3SparsePrefillMetadata | None = None
         if num_prefills > 0:
+            # Only prefill reads the context lengths; decode-only steps skip it.
+            context_lens = self.context_len_buffer[:num_reqs]
+            context_lens.copy_(
+                common_attn_metadata.compute_num_computed_tokens(), non_blocking=True
+            )
             seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound
             assert seq_lens_cpu is not None
             prefill_seq_lens_cpu = seq_lens_cpu[num_decodes:]
@@ -439,11 +438,12 @@ def forward(
         query: torch.Tensor,
         kv_cache: torch.Tensor,
         output: torch.Tensor,
-        *,
-        query_fp8: torch.Tensor | None = None,
     ) -> torch.Tensor:
         """Attend the queries to the indexer-selected blocks. Per kernel.
 
+        ``query`` is E4M3 for impls with ``use_cutlass_decode``, else the model
+        dtype.
+
         The indexer has already written the top-k block ids into
         ``layer.topk_indices_buffer`` (decode at ``[:, :nd]``, prefill at
         ``[:, nd:num_tokens]``); the attend reads them from there.
@@ -463,8 +463,6 @@ def forward(
         query: torch.Tensor,
         kv_cache: torch.Tensor,
         output: torch.Tensor,
-        *,
-        query_fp8: torch.Tensor | None = None,
     ) -> torch.Tensor:
         attn_metadata = get_forward_context().attn_metadata
         if not isinstance(attn_metadata, dict):
@@ -563,6 +561,17 @@ def select_main_backend_and_impl_cls(
             "MiniMax M3 sparse attention supports only the 'nvfp4' NVFP4 KV "
             f"cache dtype, got {kv_cache_dtype}"
         )
+    if use_nvfp4 and current_platform.is_rocm():
+        from vllm.models.minimax_m3.amd.sparse_attention_nvfp4 import (
+            MiniMaxM3SparseRocmNvfp4Backend,
+            MiniMaxM3SparseRocmNvfp4Impl,
+        )
+
+        logger.info_once(
+            "MiniMax M3 sparse attention selected ROCm NVFP4 (topk_blocks=%s)",
+            topk_blocks,
+        )
+        return MiniMaxM3SparseRocmNvfp4Backend, MiniMaxM3SparseRocmNvfp4Impl
     if use_nvfp4 and (use_aiter_sparse_pa or not use_msa):
         raise ValueError(
             "MiniMax M3 sparse attention supports an NVFP4 KV cache only with "
@@ -590,10 +599,11 @@ def select_main_backend_and_impl_cls(
             MiniMaxM3SparseMSABackend,
             MiniMaxM3SparseMSAImpl,
             MiniMaxM3SparseMSANvfp4Backend,
+            MiniMaxM3SparseMSANvfp4Impl,
         )
 
         if use_nvfp4:
-            return MiniMaxM3SparseMSANvfp4Backend, MiniMaxM3SparseMSAImpl
+            return MiniMaxM3SparseMSANvfp4Backend, MiniMaxM3SparseMSANvfp4Impl
         return MiniMaxM3SparseMSABackend, MiniMaxM3SparseMSAImpl
     return MiniMaxM3SparseBackend, MiniMaxM3SparseTritonImpl
 
diff --git a/vllm/models/minimax_m3/common/vision_tower.py b/vllm/models/minimax_m3/common/vision_tower.py
index c7aba484f2a3..b730b07093bb 100644
--- a/vllm/models/minimax_m3/common/vision_tower.py
+++ b/vllm/models/minimax_m3/common/vision_tower.py
@@ -59,11 +59,12 @@ def __init__(self, config: PreTrainedConfig) -> None:
             stride=(temporal_patch_size, patch_size, patch_size),
             bias=False,
         )
+        # The checkpoint stores this weight in fp32 (the vision tower's only
+        # fp32 tensor); keep it that way so the conv runs at full precision.
+        self.patch_embedding.to(torch.float32)
 
     def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
         # pixel_values: (N, C * temporal_patch_size * patch_size²)
-        if self.patch_embedding.weight.dtype != pixel_values.dtype:
-            self.patch_embedding = self.patch_embedding.to(pixel_values.dtype)
         x = pixel_values.reshape(
             pixel_values.shape[0],
             self.num_channels,
@@ -71,7 +72,8 @@ def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
             self.patch_size,
             self.patch_size,
         )
-        return self.patch_embedding(x).reshape(x.shape[0], -1)
+        out = self.patch_embedding(x.to(self.patch_embedding.weight.dtype))
+        return out.reshape(x.shape[0], -1).to(pixel_values.dtype)
 
 
 class MiniMaxVLAttention(nn.Module):
@@ -657,7 +659,9 @@ def __init__(
             prefix=maybe_prefix(prefix, "patch_merge_mlp"),
         )
 
-        self.dtype = self.vision_model.embeddings.patch_embedding.weight.dtype
+        # Tower compute dtype; the patch embedding weight stays fp32 (see
+        # MiniMaxVLPatchEmbed), so take it from the projector instead.
+        self.dtype = self.multi_modal_projector.linear_1.weight.dtype
         self.out_hidden_size = text_hidden_size
 
     def forward(
diff --git a/vllm/models/minimax_m3/nvidia/indexer_nvfp4.py b/vllm/models/minimax_m3/nvidia/indexer_nvfp4.py
new file mode 100644
index 000000000000..501d81f772a7
--- /dev/null
+++ b/vllm/models/minimax_m3/nvidia/indexer_nvfp4.py
@@ -0,0 +1,481 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""NVFP4 index-cache indexer impl for MiniMax M3 (SM100/SM103/SM107).
+
+The index-K side cache holds packed NVFP4 pages ``[num_blocks, 128, 72]`` (E2M1
+codes, then E4M3 block scales, no global scale) written by the fused
+qk-norm/rope insert, and index-Q is E4M3. Both sides use the MSA paged indexers
+of ``fmha_sm100`` with fused top-k, which write 15 score-ranked history pages
+plus the query's local page last (``init_blocks=0, local_blocks=1``) as
+logical block ids:
+
+* Decode scores the NVFP4 pages with the Q8KV4 decode indexer, which reads the
+  1..8 queries of each request in place (they take the last slots of its
+  8-query tile) and writes their top-k straight into the shared buffer.
+  Uniform decode batches are captured in full CUDA graphs, so each padded
+  batch size and query length gets its own CUDA-graph wrapper whose buffers
+  never move; the builder re-plans it outside capture every step.
+* Prefill has no NVFP4 kernel. Each layer dequantizes the history pages of the
+  prefill requests into compact FP8 workspace pages (what the decode kernel
+  scores) and runs the Q8KV8 prefill indexer on them, in request chunks whose
+  pages fit the gather budget.
+
+All index heads of the rank share the index keys and select their own pages.
+Prefill and decode score every head of the rank in one launch. Prefill plans
+are rebuilt every step and only run eagerly (mixed batches are never captured
+whole).
+"""
+
+from dataclasses import dataclass, field
+from typing import Any, ClassVar
+
+import torch
+
+from vllm.config import VllmConfig, get_current_vllm_config
+from vllm.forward_context import get_forward_context
+from vllm.models.minimax_m3.common.indexer import (
+    MiniMaxM3IndexerBackend,
+    MiniMaxM3IndexerImpl,
+    MiniMaxM3IndexerMetadata,
+    MiniMaxM3IndexerMetadataBuilder,
+)
+from vllm.models.minimax_m3.nvidia.indexer_nvfp4_decode_cp import (
+    Nvfp4DecodeCPLayer,
+    Nvfp4DecodeCPPlan,
+    get_nvfp4_indexer_decode_cp,
+)
+from vllm.models.minimax_m3.nvidia.ops.gather_nvfp4_index_k import (
+    minimax_m3_gather_nvfp4_index_k_fp8,
+)
+from vllm.utils.math_utils import cdiv
+from vllm.v1.attention.backend import (
+    AttentionBackend,
+    CommonAttentionMetadata,
+    max_decode_query_len,
+)
+from vllm.v1.attention.backends.utils import split_decodes_and_prefills
+from vllm.v1.kv_cache_interface import AttentionSpec
+from vllm.v1.worker.workspace import current_workspace_manager
+
+PAGE_SIZE = 128
+HEAD_DIM = 128
+TOPK_BLOCKS = 16
+# Queries per request the Q8KV4 decode indexer scores.
+DECODE_QUERY_LEN = 8
+# Longest block table the MSA indexers accept.
+MAX_PAGES = 8192
+# Index heads the Q8KV4 decode indexer scores in one launch.
+DECODE_NUM_HEADS = (1, 2, 4)
+
+
+def _msa_indexers():
+    """Import the MSA indexer wrappers.
+
+    Through ``fmha_sm100.sparse``, which binds the package's csrc JIT loader.
+    """
+    from vllm.third_party.fmha_sm100.sparse import (
+        BatchDecodeIndexerQ8KV4Wrapper,
+        BatchPrefillIndexerQ8KV8Wrapper,
+    )
+
+    return BatchDecodeIndexerQ8KV4Wrapper, BatchPrefillIndexerQ8KV8Wrapper
+
+
+def max_gather_pages(vllm_config: VllmConfig) -> int:
+    """FP8 pages one prefill chunk may gather; always fits one full request."""
+    tokens = max(
+        vllm_config.model_config.max_model_len,
+        4 * vllm_config.scheduler_config.max_num_batched_tokens,
+    )
+    return cdiv(tokens, PAGE_SIZE)
+
+
+class MiniMaxM3IndexerNvfp4Backend(MiniMaxM3IndexerBackend):
+    """Indexer side-cache backend selecting the NVFP4 builder."""
+
+    @staticmethod
+    def get_builder_cls() -> type["MiniMaxM3IndexerNvfp4MetadataBuilder"]:
+        return MiniMaxM3IndexerNvfp4MetadataBuilder
+
+
+@dataclass
+class MiniMaxM3IndexerNvfp4DecodeMetadata:
+    wrapper: Any  # planned CUDA-graph wrapper, None when decode CP is on
+    num_reqs: int  # includes zero-length full-CUDA-graph padding requests
+    query_len: int  # uniform, <= DECODE_QUERY_LEN
+    # Set instead of the wrapper when decode CP is on.
+    cp_plan: Nvfp4DecodeCPPlan | None = None
+
+
+@dataclass
+class MiniMaxM3IndexerNvfp4PrefillChunk:
+    wrapper: Any  # planned BatchPrefillIndexerQ8KV8Wrapper
+    # Token range within the prefill region (offsets from num_decode_tokens).
+    token_start: int
+    token_end: int
+    block_table: torch.Tensor  # [n, max_blocks] physical pages
+    num_pages: torch.Tensor  # [n] int32 history pages to gather
+    page_offsets: torch.Tensor  # [n] int32 first compact page of each request
+    total_pages: int
+    max_pages: int
+
+
+@dataclass
+class MiniMaxM3IndexerNvfp4Metadata(MiniMaxM3IndexerMetadata):
+    """The base ``prefill``/``decode`` fields are unused on this path."""
+
+    decode_q8kv4: MiniMaxM3IndexerNvfp4DecodeMetadata | None = None
+    prefill_chunks: list[MiniMaxM3IndexerNvfp4PrefillChunk] = field(
+        default_factory=list
+    )
+
+
+class MiniMaxM3IndexerNvfp4MetadataBuilder(MiniMaxM3IndexerMetadataBuilder):
+    """Plans the MSA decode and prefill indexers once per step."""
+
+    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)
+        if self.max_decode_query_len > DECODE_QUERY_LEN:
+            raise NotImplementedError(
+                "The MiniMax M3 NVFP4 indexer scores at most "
+                f"{DECODE_QUERY_LEN} decode queries per request, got "
+                f"{self.max_decode_query_len} (use num_speculative_tokens <= "
+                f"{DECODE_QUERY_LEN - 1})."
+            )
+        if self.num_index_heads not in DECODE_NUM_HEADS:
+            raise NotImplementedError(
+                "The MiniMax M3 NVFP4 indexer serves "
+                f"{DECODE_NUM_HEADS} index heads per rank, got "
+                f"{self.num_index_heads}."
+            )
+        self.max_gather_pages = max_gather_pages(vllm_config)
+        # Decode wrapper per padded decode batch size and query length.
+        self.decode_wrappers: dict[tuple[int, int], Any] = {}
+        # Decode CP (opt-in); static, so all ranks and graphs agree.
+        self.decode_cp = get_nvfp4_indexer_decode_cp(
+            num_index_heads=self.num_index_heads,
+            max_decode_query_len=self.max_decode_query_len,
+        )
+        # Pages a request can span; decode scores at most this many.
+        self.max_decode_pages = min(
+            cdiv(vllm_config.model_config.max_model_len, PAGE_SIZE), MAX_PAGES
+        )
+        # One page-score scratch for every decode wrapper: they never run at
+        # once, and a buffer per batch size would hold 256 KiB per request at a
+        # 1M context for every batch size the engine sees. Decode batches,
+        # padded to a CUDA graph size, never exceed max_decode_reqs.
+        self.decode_score_buffer: torch.Tensor | None = None
+        if self.decode_cp is None:
+            max_decode_reqs = max(
+                vllm_config.scheduler_config.max_num_seqs,
+                vllm_config.compilation_config.max_cudagraph_capture_size or 0,
+            )
+            decode_wrapper_cls, _ = _msa_indexers()
+            self.decode_score_buffer = torch.empty(
+                decode_wrapper_cls.score_buffer_size(
+                    max_decode_reqs,
+                    self.max_decode_pages,
+                    num_heads=self.num_index_heads,
+                ),
+                dtype=torch.float32,
+                device=device,
+            )
+
+    def _decode_wrapper(
+        self, block_table: torch.Tensor, seq_lens: torch.Tensor, query_len: int
+    ) -> Any:
+        num_reqs = seq_lens.shape[0]
+        wrapper = self.decode_wrappers.get((num_reqs, query_len))
+        if wrapper is None:
+            # The runner's block table and seq_lens are persistent buffers, so
+            # the wrapper reads them in place and plan() copies nothing. A
+            # strided (width-capped) block table gets its own buffer.
+            if not block_table.is_contiguous():
+                block_table = torch.empty(
+                    block_table.shape, dtype=torch.int32, device=self.device
+                )
+            decode_wrapper_cls, _ = _msa_indexers()
+            wrapper = self.decode_wrappers[num_reqs, query_len] = decode_wrapper_cls(
+                num_heads=self.num_index_heads,
+                use_cuda_graph=True,
+                block_table_buffer=block_table,
+                seq_lens_buffer=seq_lens,
+                score_buffer=self.decode_score_buffer,
+            )
+        return wrapper
+
+    def build(
+        self,
+        common_prefix_len: int,
+        common_attn_metadata: CommonAttentionMetadata,
+        fast_build: bool = False,
+    ) -> MiniMaxM3IndexerNvfp4Metadata:
+        num_reqs = common_attn_metadata.num_reqs
+        num_tokens = common_attn_metadata.num_actual_tokens
+        seq_lens = common_attn_metadata.seq_lens
+        block_table = common_attn_metadata.block_table_tensor
+        qsl_cpu = common_attn_metadata.query_start_loc_cpu
+
+        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=True,
+            )
+        )
+        assert num_decodes + num_prefills == num_reqs
+        assert num_decode_tokens + num_prefill_tokens == num_tokens
+
+        decode: MiniMaxM3IndexerNvfp4DecodeMetadata | None = None
+        if num_decode_tokens > 0:
+            # num_decodes counts the zero-length requests that pad a full CUDA
+            # graph batch, so it is the batch size the graph was captured with.
+            query_lens_cpu = qsl_cpu[1 : num_decodes + 1] - qsl_cpu[:num_decodes]
+            query_len = int(query_lens_cpu[0])
+            assert torch.all((query_lens_cpu == query_len) | (query_lens_cpu == 0))
+            num_pages = min(block_table.shape[1], self.max_decode_pages)
+            if self.decode_cp is not None:
+                decode = MiniMaxM3IndexerNvfp4DecodeMetadata(
+                    wrapper=None,
+                    num_reqs=num_decodes,
+                    query_len=query_len,
+                    cp_plan=self.decode_cp.plan(
+                        block_table[:num_decodes, :num_pages],
+                        seq_lens[:num_decodes],
+                    ),
+                )
+            else:
+                decode_block_table = block_table[:num_decodes, :num_pages]
+                decode_seq_lens = seq_lens[:num_decodes]
+                wrapper = self._decode_wrapper(
+                    decode_block_table, decode_seq_lens, query_len
+                )
+                wrapper.plan(decode_block_table, decode_seq_lens, query_len=query_len)
+                decode = MiniMaxM3IndexerNvfp4DecodeMetadata(
+                    wrapper=wrapper,
+                    num_reqs=num_decodes,
+                    query_len=query_len,
+                )
+
+        prefill_chunks: list[MiniMaxM3IndexerNvfp4PrefillChunk] = []
+        if num_prefills > 0:
+            prefill_chunks = self._build_prefill_chunks(
+                common_attn_metadata, num_decodes
+            )
+
+        return MiniMaxM3IndexerNvfp4Metadata(
+            seq_lens=seq_lens,
+            max_seq_len=common_attn_metadata.max_seq_len,
+            slot_mapping=common_attn_metadata.slot_mapping,
+            num_actual_tokens=num_tokens,
+            num_decodes=num_decodes,
+            num_decode_tokens=num_decode_tokens,
+            num_prefills=num_prefills,
+            num_prefill_tokens=num_prefill_tokens,
+            decode_q8kv4=decode,
+            prefill_chunks=prefill_chunks,
+        )
+
+    def _build_prefill_chunks(
+        self, common_attn_metadata: CommonAttentionMetadata, num_decodes: int
+    ) -> list[MiniMaxM3IndexerNvfp4PrefillChunk]:
+        _, prefill_wrapper_cls = _msa_indexers()
+        num_reqs = common_attn_metadata.num_reqs
+        seq_lens = common_attn_metadata.seq_lens
+        block_table = common_attn_metadata.block_table_tensor
+        qsl_cpu = common_attn_metadata.query_start_loc_cpu[num_decodes : num_reqs + 1]
+        # Host-side lengths keep planning free of device syncs. An upper bound
+        # (async spec decode) only over-sizes the plan and the gather: the
+        # kernels mask by the device seq_lens and never score the extra pages.
+        seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound
+        assert seq_lens_cpu is not None
+        seq_lens_cpu = seq_lens_cpu[num_decodes:num_reqs].to(torch.int32)
+        query_lens_cpu = (qsl_cpu[1:] - qsl_cpu[:-1]).to(torch.int32)
+        # Only pages before the last query's local page are ever scored.
+        history_pages = (seq_lens_cpu - 1) // PAGE_SIZE
+        device = seq_lens.device
+
+        chunks = []
+        start, num_prefills = 0, seq_lens_cpu.shape[0]
+        while start < num_prefills:
+            end, total_pages = start, 0
+            while end < num_prefills and (
+                end == start
+                or total_pages + int(history_pages[end]) <= self.max_gather_pages
+            ):
+                total_pages += int(history_pages[end])
+                end += 1
+            assert total_pages <= self.max_gather_pages
+            num_chunk_reqs = end - start
+            max_seq_len = int(seq_lens_cpu[start:end].max())
+            width = cdiv(max_seq_len, PAGE_SIZE)
+            # [cu_seqlens_q | num_pages | page_offsets | compact block table]
+            # in one pinned buffer, so a single H2D copy is truly asynchronous
+            # (a pageable one may synchronize with the stream).
+            staging = torch.empty(
+                3 * num_chunk_reqs + 1 + num_chunk_reqs * width,
+                dtype=torch.int32,
+                pin_memory=True,
+            )
+            cu_seqlens_q, pages, offsets, table = staging.split(
+                [
+                    num_chunk_reqs + 1,
+                    num_chunk_reqs,
+                    num_chunk_reqs,
+                    num_chunk_reqs * width,
+                ]
+            )
+            torch.sub(qsl_cpu[start : end + 1], qsl_cpu[start], out=cu_seqlens_q)
+            pages.copy_(history_pages[start:end])
+            torch.cumsum(pages, 0, dtype=torch.int32, out=offsets)
+            offsets.sub_(pages)
+            cols = torch.arange(width, dtype=torch.int32)
+            table.view(num_chunk_reqs, width).copy_(
+                torch.where(cols < pages[:, None], offsets[:, None] + cols, 0)
+            )
+            staging_gpu = staging.to(device, non_blocking=True)
+            cu_seqlens_q_gpu, pages_gpu, offsets_gpu, table_gpu = staging_gpu.split(
+                [
+                    num_chunk_reqs + 1,
+                    num_chunk_reqs,
+                    num_chunk_reqs,
+                    num_chunk_reqs * width,
+                ]
+            )
+            lo, hi = num_decodes + start, num_decodes + end
+            token_start = int(qsl_cpu[start] - qsl_cpu[0])
+            token_end = int(qsl_cpu[end] - qsl_cpu[0])
+            wrapper = prefill_wrapper_cls(num_heads=self.num_index_heads)
+            wrapper.plan(
+                cu_seqlens_q_gpu,
+                seq_lens[lo:hi],
+                table_gpu.view(num_chunk_reqs, width),
+                total_q=token_end - token_start,
+                max_seqlen_q=int(query_lens_cpu[start:end].max()),
+                max_seqlen_k=max_seq_len,
+            )
+            chunks.append(
+                MiniMaxM3IndexerNvfp4PrefillChunk(
+                    wrapper=wrapper,
+                    token_start=token_start,
+                    token_end=token_end,
+                    block_table=block_table[lo:hi],
+                    num_pages=pages_gpu,
+                    page_offsets=offsets_gpu,
+                    total_pages=total_pages,
+                    max_pages=int(pages.max()),
+                )
+            )
+            start = end
+        return chunks
+
+
+class MiniMaxM3IndexerNvfp4Impl(MiniMaxM3IndexerImpl):
+    """Decode: Q8KV4 on NVFP4 pages. Prefill: FP8 gather + Q8KV8."""
+
+    indexer_backend_cls: ClassVar[type[AttentionBackend]] = MiniMaxM3IndexerNvfp4Backend
+
+    def __init__(self, **kwargs) -> None:
+        super().__init__(**kwargs)
+        config = (
+            self.topk_blocks,
+            self.init_blocks,
+            self.local_blocks,
+            self.block_size,
+            self.index_head_dim,
+            self.score_type,
+        )
+        if config != (TOPK_BLOCKS, 0, 1, PAGE_SIZE, HEAD_DIM, "max"):
+            raise NotImplementedError(
+                "The MiniMax M3 NVFP4 indexer supports topk_blocks=16, "
+                "init_blocks=0, local_blocks=1, 128-token blocks, index head "
+                "dim 128 and max scoring; got (topk_blocks, init_blocks, "
+                "local_blocks, block_size, index_head_dim, score_type)="
+                f"{config}."
+            )
+        vllm_config = get_current_vllm_config()
+        self.max_gather_pages = max_gather_pages(vllm_config)
+        # Per-layer decode-CP buffers (collective: same layer order on all ranks).
+        self.decode_cp = get_nvfp4_indexer_decode_cp(
+            num_index_heads=self.num_index_heads,
+            max_decode_query_len=max_decode_query_len(vllm_config),
+        )
+        self.decode_cp_layer: Nvfp4DecodeCPLayer | None = (
+            None if self.decode_cp is None else self.decode_cp.new_layer()
+        )
+
+    def reserve_workspace(self) -> None:
+        current_workspace_manager().get_simultaneous(
+            ((self.max_gather_pages, PAGE_SIZE, HEAD_DIM), torch.float8_e4m3fn)
+        )
+
+    def forward(
+        self,
+        index_query: torch.Tensor,
+    ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
+        attn_metadata = get_forward_context().attn_metadata
+        if not isinstance(attn_metadata, dict):
+            return None, None  # profiling run; caches unbound
+        md = attn_metadata[self.index_cache.prefix]
+        assert isinstance(md, MiniMaxM3IndexerNvfp4Metadata)
+
+        nd = md.num_decode_tokens
+        # Not sliced to the actual tokens: a padded full-CUDA-graph decode
+        # batch also scores its padding rows, and every captured tensor must
+        # keep its address and shape.
+        index_q = index_query.view(-1, self.num_index_heads, self.index_head_dim)
+        kv = self.index_cache.kv_cache
+        # The kernels write logical block ids straight into the shared
+        # [T, H, 16] buffer the attend reads.
+        buf = self.topk_indices_buffer
+        assert buf is not None
+
+        if md.decode_q8kv4 is not None:
+            d = md.decode_q8kv4
+            rows = d.num_reqs * d.query_len
+            if d.cp_plan is not None:
+                assert self.decode_cp is not None and self.decode_cp_layer is not None
+                self.decode_cp.forward(
+                    self.decode_cp_layer,
+                    d.cp_plan,
+                    index_q[:rows],
+                    kv,
+                    buf[:rows],
+                    d.query_len,
+                )
+            else:
+                d.wrapper.run(index_q[:rows], kv, out=buf[:rows])
+        for chunk in md.prefill_chunks:
+            self._prefill(
+                index_q[nd + chunk.token_start : nd + chunk.token_end],
+                kv,
+                chunk,
+                buf[nd + chunk.token_start : nd + chunk.token_end],
+            )
+        return None, None
+
+    def _prefill(
+        self,
+        index_q: torch.Tensor,
+        kv: torch.Tensor,
+        chunk: MiniMaxM3IndexerNvfp4PrefillChunk,
+        out: torch.Tensor,
+    ) -> None:
+        (k_pages,) = current_workspace_manager().get_simultaneous(
+            ((max(chunk.total_pages, 1), PAGE_SIZE, HEAD_DIM), torch.float8_e4m3fn)
+        )
+        minimax_m3_gather_nvfp4_index_k_fp8(
+            kv,
+            chunk.block_table,
+            chunk.num_pages,
+            chunk.page_offsets,
+            chunk.max_pages,
+            k_pages,
+        )
+        chunk.wrapper.run(index_q, k_pages, out=out)
diff --git a/vllm/models/minimax_m3/nvidia/indexer_nvfp4_decode_cp.py b/vllm/models/minimax_m3/nvidia/indexer_nvfp4_decode_cp.py
new file mode 100644
index 000000000000..012515855628
--- /dev/null
+++ b/vllm/models/minimax_m3/nvidia/indexer_nvfp4_decode_cp.py
@@ -0,0 +1,513 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Context-parallel decode for the MiniMax M3 NVFP4 indexer (opt-in, TP4).
+
+At TP4 each rank owns one of the 4 index heads, but the index-K cache is
+replicated, so every rank reads the whole cache per sparse layer. With
+``--attention-config '{"indexer_decode_cp": true}'``, rank ``r``
+instead:
+
+1. exchanges index-q heads so it holds all 4 (``_q_exchange_kernel``);
+2. scores all 4 heads over pages ``j * CP + r`` with the stock Q8KV4 kernel
+   on a compact block table, reading 1/CP of the index-K;
+3. takes the exact local top 16 per (query, head) and sends it to the head's
+   owner (``_local_topk_push_kernel``);
+4. merges the owner's ``16 * CP`` candidates into the indexer output: 15
+   pages by (score desc, page asc), then the own page (``_merge_kernel``).
+
+The merge is exact: the global top 16 lies in the union of per-rank top 16s.
+The stock decode ranks 16-bit quantized scores, so the two may differ at ties.
+
+The chunking is written against one contract, the MSA Q8KV4 decode scorer's:
+eight query slots per request, ``(seq_len - 1) / 128`` pages scored, and only
+the pages below a query's own page stored. The three places that depend on it
+are ``_compact_block_table_kernel`` (the length it is handed),
+``_local_topk_push_kernel`` (the cut that length drops) and ``_set_width``
+(room for the page that length invents).
+
+Both exchanges are NVLink stores into symmetric buffers whose payload is the
+flag (``0xFF`` E4M3 NaN for queries, ``-1`` for keys); consumers restore it.
+Buffers are per layer, shapes are fixed per batch size, and nothing syncs with
+the host, so the path is CUDA-graph safe. It needs symmetric memory across the
+tensor-parallel ranks and falls back to the stock decode without it.
+
+The three Triton kernels are PDL-launched and wait for their producer, so
+each starts its prologue in the tail of the kernel before it. Only the push
+releases its dependent early, right after its own wait, which is what lets
+the merge overlap it; the exchange must not, because the score kernel after
+it reads the whole tile.
+"""
+
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+
+from vllm.config import get_current_vllm_config
+from vllm.distributed import (
+    get_tensor_model_parallel_rank,
+    get_tensor_model_parallel_world_size,
+    get_tp_group,
+)
+from vllm.logger import init_logger
+from vllm.triton_utils import tl, triton
+from vllm.utils.math_utils import cdiv
+from vllm.v1.attention.ops.cp_common import direct_cp_enabled
+
+logger = init_logger(__name__)
+
+PAGE_SIZE = 128
+HEAD_DIM = 128
+TOPK_BLOCKS = 16
+DECODE_QUERY_LEN = 8  # query slots per request in a Q8KV4 decode tile
+MAX_PAGES = 8192
+NUM_INDEX_HEADS = 4
+SUPPORTED_TP_SIZES = (4,)
+# Rows with at most this many ranked pages take a single top-k.
+TOPK_CHUNK = 256
+# Groups of the two-level top-k over longer rows.
+MAX_TOPK_GROUPS = 128
+
+_TOPK = tl.constexpr(TOPK_BLOCKS)
+_RANKED = tl.constexpr(TOPK_BLOCKS - 1)
+_PAGE = tl.constexpr(PAGE_SIZE)
+_SLOTS = tl.constexpr(DECODE_QUERY_LEN)
+_WORDS = tl.constexpr(HEAD_DIM // 4)  # int32 words per index-q head
+
+
+@triton.jit
+def _peer_ptr(ptrs, rank, dtype: tl.constexpr):
+    return tl.load(ptrs + rank).to(tl.pointer_type(dtype))
+
+
+@triton.jit
+def _compact_block_table_kernel(
+    bt_ptr,  # [B, P] int32
+    seq_lens_ptr,  # [B] int32
+    cbt_ptr,  # [B, W] int32 out: compact page j -> global page j * CP + RANK
+    cseq_ptr,  # [B] int32 out: lengths the score kernel sees
+    seq_out_ptr,  # [B] int32 out: real lengths (stable buffer)
+    P,
+    W,
+    stride_bt,
+    stride_cbt,
+    CP: tl.constexpr,
+    RANK: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    """Compact page j = global page ``j * CP + RANK``, and the length the score
+    kernel is handed for it.
+
+    That length is a lie: the scorer stores only the pages below a query's own
+    page, so ``n_local * 128 + 8`` puts all eight query slots one page past
+    the chunk and stores every owned page, uncausally. The real cut comes back
+    in ``_local_topk_push_kernel``. ``W`` keeps the invented page addressable,
+    since a scorer that clamps the own page to its last score column would
+    drop the chunk's last page."""
+    b = tl.program_id(0)
+    seq = tl.load(seq_lens_ptr + b)
+    tl.store(seq_out_ptr + b, seq)
+    hist = tl.minimum(tl.maximum((seq - 1) // _PAGE, 0), P)
+    n_local = tl.minimum((hist + (CP - 1 - RANK)) // CP, W - 1)
+    for start in range(0, W, BLOCK):
+        j = start + tl.arange(0, BLOCK)
+        src = tl.minimum(j * CP + RANK, P - 1)
+        v = tl.load(bt_ptr + b.to(tl.int64) * stride_bt + src, mask=j < W)
+        tl.store(cbt_ptr + b.to(tl.int64) * stride_cbt + j, v, mask=j < W)
+    tl.store(cseq_ptr + b, n_local * _PAGE + _SLOTS)
+
+
+@triton.jit
+def _has_ff_byte(v):
+    u = v.to(tl.uint32, bitcast=True)
+    ff = (
+        ((u & 0xFF) == 0xFF)
+        | (((u >> 8) & 0xFF) == 0xFF)
+        | (((u >> 16) & 0xFF) == 0xFF)
+        | ((u >> 24) == 0xFF)
+    )
+    return tl.max(ff.to(tl.int32), axis=0) > 0
+
+
+@triton.jit
+def _q_exchange_kernel(
+    q_ptr,  # [rows, 32] int32 view of this rank's e4m3 index-q head
+    tile_ptrs,  # [CP] int64: q tile [B, 8, H, 128] bytes of every rank
+    H: tl.constexpr,
+    CP: tl.constexpr,
+    RANK: tl.constexpr,
+    HEAD: tl.constexpr,
+    QLEN: tl.constexpr,
+):
+    """Store this rank's index-q head into every rank's tile, then wait for
+    the other heads. The PDL wait is for index-q's producer."""
+    tl.extra.cuda.gdc_wait()
+    t = tl.program_id(0)
+    b = t // QLEN
+    slot = (_SLOTS - QLEN) + t % QLEN
+    w = tl.arange(0, _WORDS)
+    row = tl.load(q_ptr + t * _WORDS + w)
+    dst_off = ((b * _SLOTS + slot) * H + HEAD) * _WORDS + w
+    for dst in tl.static_range(CP):
+        tl.store(_peer_ptr(tile_ptrs, dst, tl.int32) + dst_off, row)
+    local = _peer_ptr(tile_ptrs, RANK, tl.int32)
+    hw = tl.arange(0, H * _WORDS)
+    off = (b * _SLOTS + slot) * H * _WORDS + hw
+    v = tl.load(local + off, volatile=True)
+    pending = _has_ff_byte(v)
+    while pending:
+        v = tl.load(local + off, volatile=True)
+        pending = _has_ff_byte(v)
+
+
+@triton.jit
+def _page_keys(row_ptr, j, n, CP: tl.constexpr, RANK: tl.constexpr):
+    """Key ``1 << 48 | ordered(score) << 16 | (0xFFFF - page)`` for ``j < n``,
+    else 0. Unique; orders by score desc, then page asc."""
+    valid = j < n
+    page = j * CP + RANK
+    score = tl.load(row_ptr + j, mask=valid, other=0.0)
+    bits = score.to(tl.uint32, bitcast=True)
+    ordered = bits ^ tl.where((bits >> 31) != 0, 0xFFFFFFFF, 0x80000000)
+    key = (1 << 48) | (ordered.to(tl.int64) << 16) | (0xFFFF - page).to(tl.int64)
+    return tl.where(valid, key, 0)
+
+
+@triton.jit
+def _local_topk_push_kernel(
+    score_ptr,  # [B, 8 * H, W] fp32 compact scores (row = slot * H + head)
+    seq_lens_ptr,  # [B] int32
+    recv_ptrs,  # [CP] int64: candidates [CP, rows_max, 16] int64 of every rank
+    rows_max,
+    stride_sb,
+    stride_sr,
+    H: tl.constexpr,
+    CP: tl.constexpr,
+    RANK: tl.constexpr,
+    QLEN: tl.constexpr,
+    WIDTH: tl.constexpr,
+    CHUNK: tl.constexpr,
+    NUM_GROUPS: tl.constexpr,
+    GROUP: tl.constexpr,
+):
+    """Per (query, head): exact top 16 of the owned pages, sent to the head's
+    owner. Long rows use a two-level select (top 16 groups, then top 16).
+
+    PDL-launched: the wait is for the scores, and releasing the merge only
+    after it keeps the merge's own wait a guarantee that the score kernel has
+    read the q tile."""
+    tl.extra.cuda.gdc_wait()
+    tl.extra.cuda.gdc_launch_dependents()
+    t = tl.program_id(0)
+    h = tl.program_id(1)
+    b = t // QLEN
+    i = t % QLEN
+    slot = (_SLOTS - QLEN) + i
+    own = (tl.load(seq_lens_ptr + b) - QLEN + i) // _PAGE
+    if own + 1 > _TOPK:
+        # The causal cut the compact length dropped: the rank's pages below
+        # the query's own page, bounded by the width like the scored pages (a
+        # length past the page table ranks what it holds, as the stock TopK).
+        n_rank = tl.minimum((own + (CP - 1 - RANK)) // CP, WIDTH - 1)
+        row_ptr = score_ptr + b.to(tl.int64) * stride_sb + (slot * H + h) * stride_sr
+        if n_rank <= CHUNK:
+            j = tl.arange(0, CHUNK)
+            best = tl.topk(_page_keys(row_ptr, j, n_rank, CP, RANK), _TOPK)
+        else:
+            g = tl.arange(0, NUM_GROUPS)
+            s = tl.arange(0, GROUP)
+            keys = _page_keys(
+                row_ptr, g[:, None] * GROUP + s[None, :], n_rank, CP, RANK
+            )
+            top_g = tl.topk(tl.max(keys, axis=1), _TOPK)
+            page = (0xFFFF - (top_g & 0xFFFF)).to(tl.int32)
+            grp = tl.where(top_g > 0, ((page - RANK) // CP) // GROUP, NUM_GROUPS)
+            cand = _page_keys(
+                row_ptr, grp[:, None] * GROUP + s[None, :], n_rank, CP, RANK
+            )
+            best = tl.topk(tl.reshape(cand, [_TOPK * GROUP]), _TOPK)
+        k = tl.arange(0, _TOPK)
+        dst = _peer_ptr(recv_ptrs, h, tl.int64)
+        tl.store(dst + (RANK * rows_max + t) * _TOPK + k, best)
+
+
+@triton.jit
+def _merge_kernel(
+    seq_lens_ptr,  # [B] int32
+    recv_ptrs,  # [CP] int64
+    tile_ptrs,  # [CP] int64
+    out_ptr,  # [rows, 1, 16] int32
+    rows_max,
+    out_stride,
+    H: tl.constexpr,
+    CP: tl.constexpr,
+    RANK: tl.constexpr,
+    QLEN: tl.constexpr,
+):
+    """Merge all ranks' candidates for this rank's head and restore sentinels.
+
+    PDL-launched after every push program has started, so polling here never
+    blocks a push."""
+    t = tl.program_id(0)
+    b = t // QLEN
+    i = t % QLEN
+    own = (tl.load(seq_lens_ptr + b) - QLEN + i) // _PAGE
+    k = tl.arange(0, _TOPK)
+    out = out_ptr + t.to(tl.int64) * out_stride + k
+    if own + 1 <= _TOPK:
+        tl.store(out, tl.where(k < own + 1, k, -1))
+        # No push for this row: wait for the score before resetting the q tile.
+        tl.extra.cuda.gdc_wait()
+    else:
+        recv = _peer_ptr(recv_ptrs, RANK, tl.int64)
+        p = tl.arange(0, CP)
+        off = (p[:, None] * rows_max + t) * _TOPK + k[None, :]
+        keys = tl.load(recv + off, volatile=True)
+        pending = tl.min(tl.min(keys, axis=1), axis=0) == -1
+        while pending:
+            keys = tl.load(recv + off, volatile=True)
+            pending = tl.min(tl.min(keys, axis=1), axis=0) == -1
+        best = tl.topk(tl.reshape(keys, [CP * _TOPK]), _TOPK)
+        page = (0xFFFF - (best & 0xFFFF)).to(tl.int32)
+        tl.store(out, tl.where(k == _RANKED, own, page))
+        tl.store(recv + off, tl.full([CP, _TOPK], -1, tl.int64))
+    # The score kernel has read this q tile row; re-arm it.
+    slot = (_SLOTS - QLEN) + i
+    tile = _peer_ptr(tile_ptrs, RANK, tl.int32)
+    hw = tl.arange(0, H * _WORDS)
+    tl.store(
+        tile + (b * _SLOTS + slot) * H * _WORDS + hw,
+        tl.full(hw.shape, -1, tl.int32),
+    )
+
+
+@dataclass
+class Nvfp4DecodeCPPlan:
+    """Per padded batch size: the 4-head score wrapper and its buffers."""
+
+    # The wrapper owns the compact [B, 8 * H, width] fp32 scores.
+    wrapper: Any  # BatchDecodeIndexerQ8KV4Wrapper(num_heads=4), CUDA-graph mode
+    block_table: torch.Tensor  # [B, width] compact pages
+    score_seq_lens: torch.Tensor  # [B] lengths the score kernel sees
+    seq_lens: torch.Tensor  # [B] real lengths
+    num_reqs: int
+
+
+@dataclass
+class Nvfp4DecodeCPLayer:
+    """One sparse layer's symmetric buffers."""
+
+    # Keep both alive: freeing one unmaps it while peers still write to it.
+    tile: torch.Tensor  # local q tile, int32 words
+    recv: torch.Tensor  # local candidates [CP, rows_max, 16] int64
+    tile_ptrs: torch.Tensor  # [CP] int64
+    recv_ptrs: torch.Tensor  # [CP] int64
+
+
+class Nvfp4IndexerDecodeCP:
+    """Per-process state shared by the builder and every sparse layer."""
+
+    def __init__(self, *, max_reqs: int, device: torch.device) -> None:
+        from vllm.third_party.fmha_sm100.sparse import BatchDecodeIndexerQ8KV4Wrapper
+
+        self._wrapper_cls = BatchDecodeIndexerQ8KV4Wrapper
+        self.cp = get_tensor_model_parallel_world_size()
+        self.rank = get_tensor_model_parallel_rank()
+        self.head = self.rank  # one index head per rank at TP4
+        # Compact page width, fixed by the first plan() (see _set_width).
+        self.width = 0
+        self.num_groups = 0
+        self.group = 0
+        self.max_reqs = max_reqs
+        self.rows_max = max_reqs * DECODE_QUERY_LEN
+        self.device = device
+        self._plans: dict[int, Nvfp4DecodeCPPlan] = {}
+        self._tile_words = max_reqs * DECODE_QUERY_LEN * NUM_INDEX_HEADS * HEAD_DIM // 4
+        self._recv_words = self.cp * self.rows_max * TOPK_BLOCKS
+
+    def new_layer(self) -> Nvfp4DecodeCPLayer:
+        """Collective: every TP rank builds its sparse layers in the same order."""
+        import torch.distributed._symmetric_memory as symm_mem
+
+        group_name = get_tp_group().device_group.group_name
+        tile = symm_mem.empty(self._tile_words, dtype=torch.int32, device=self.device)
+        tile.fill_(-1)
+        recv = symm_mem.empty(self._recv_words, dtype=torch.int64, device=self.device)
+        recv.fill_(-1)
+        tile_ptrs = symm_mem.rendezvous(tile, group_name).buffer_ptrs
+        recv_ptrs = symm_mem.rendezvous(recv, group_name).buffer_ptrs
+        return Nvfp4DecodeCPLayer(
+            tile=tile,
+            recv=recv,
+            tile_ptrs=torch.tensor(tile_ptrs, dtype=torch.int64, device=self.device),
+            recv_ptrs=torch.tensor(recv_ptrs, dtype=torch.int64, device=self.device),
+        )
+
+    def _set_width(self, max_pages: int) -> None:
+        """Fix the compact page width and the two-level top-k shape.
+
+        One column past the pages a rank can own, so the page the compact
+        length invents stays inside the score buffer and the chunk's last page
+        is stored (see ``_compact_block_table_kernel``). Both levels of the
+        top-k select over power-of-two tiles that cover the width.
+        """
+        width = cdiv(min(max_pages, MAX_PAGES), self.cp) + 1
+        if self.width == 0:
+            padded = 1 << (width - 1).bit_length()
+            self.width = width
+            self.group = max(1, padded // MAX_TOPK_GROUPS)
+            self.num_groups = padded // self.group
+        assert width <= self.width, "the block table grew after the first plan"
+
+    def plan(
+        self, block_table: torch.Tensor, seq_lens: torch.Tensor
+    ) -> Nvfp4DecodeCPPlan:
+        """Outside CUDA graph capture, once per step."""
+        self._set_width(block_table.shape[1])
+        num_reqs = seq_lens.shape[0]
+        p = self._plans.get(num_reqs)
+        if p is None:
+            assert num_reqs <= self.max_reqs
+            i32 = dict(dtype=torch.int32, device=self.device)
+            bt = torch.empty(num_reqs, self.width, **i32)
+            score_seq_lens = torch.empty(num_reqs, **i32)
+            p = self._plans[num_reqs] = Nvfp4DecodeCPPlan(
+                wrapper=self._wrapper_cls(
+                    num_heads=NUM_INDEX_HEADS,
+                    use_cuda_graph=True,
+                    block_table_buffer=bt,
+                    seq_lens_buffer=score_seq_lens,
+                ),
+                block_table=bt,
+                score_seq_lens=score_seq_lens,
+                seq_lens=torch.empty(num_reqs, **i32),
+                num_reqs=num_reqs,
+            )
+        _compact_block_table_kernel[(num_reqs,)](
+            block_table,
+            seq_lens,
+            p.block_table,
+            p.score_seq_lens,
+            p.seq_lens,
+            min(block_table.shape[1], MAX_PAGES),
+            self.width,
+            block_table.stride(0),
+            p.block_table.stride(0),
+            CP=self.cp,
+            RANK=self.rank,
+            BLOCK=1024,
+        )
+        p.wrapper.plan(p.block_table, p.score_seq_lens)
+        return p
+
+    def forward(
+        self,
+        layer: Nvfp4DecodeCPLayer,
+        plan: Nvfp4DecodeCPPlan,
+        index_q: torch.Tensor,  # [rows, 1, 128] e4m3, this rank's head
+        k_cache: torch.Tensor,
+        out: torch.Tensor,  # [rows, 1, 16] int32
+        query_len: int,
+    ) -> None:
+        rows = index_q.shape[0]
+        assert rows == plan.num_reqs * query_len and out.shape[0] == rows
+        _q_exchange_kernel[(rows,)](
+            index_q.view(torch.int32).view(rows, HEAD_DIM // 4),
+            layer.tile_ptrs,
+            H=NUM_INDEX_HEADS,
+            CP=self.cp,
+            RANK=self.rank,
+            HEAD=self.head,
+            QLEN=query_len,
+            num_warps=1,
+            launch_pdl=True,
+        )
+        q_tile = layer.tile[
+            : plan.num_reqs * DECODE_QUERY_LEN * NUM_INDEX_HEADS * HEAD_DIM // 4
+        ]
+        scores = plan.wrapper._run_scores(
+            q_tile.view(torch.float8_e4m3fn).view(-1, NUM_INDEX_HEADS, HEAD_DIM),
+            k_cache,
+        )
+        _local_topk_push_kernel[(rows, NUM_INDEX_HEADS)](
+            scores,
+            plan.seq_lens,
+            layer.recv_ptrs,
+            self.rows_max,
+            scores.stride(0),
+            scores.stride(1),
+            H=NUM_INDEX_HEADS,
+            CP=self.cp,
+            RANK=self.rank,
+            QLEN=query_len,
+            WIDTH=self.width,
+            CHUNK=TOPK_CHUNK,
+            NUM_GROUPS=self.num_groups,
+            GROUP=self.group,
+            num_warps=4,
+            launch_pdl=True,
+        )
+        _merge_kernel[(rows,)](
+            plan.seq_lens,
+            layer.recv_ptrs,
+            layer.tile_ptrs,
+            out,
+            self.rows_max,
+            out.stride(0),
+            H=NUM_INDEX_HEADS,
+            CP=self.cp,
+            RANK=self.rank,
+            QLEN=query_len,
+            num_warps=1,
+            launch_pdl=True,
+        )
+
+
+_STATE: Nvfp4IndexerDecodeCP | None = None
+_RESOLVED = False
+
+
+def get_nvfp4_indexer_decode_cp(
+    *, num_index_heads: int, max_decode_query_len: int
+) -> Nvfp4IndexerDecodeCP | None:
+    """The process's decode-CP state, or None (stock decode). Built once."""
+    global _STATE, _RESOLVED
+    if _RESOLVED:
+        return _STATE
+    _RESOLVED = True
+    vllm_config = get_current_vllm_config()
+    if not vllm_config.attention_config.indexer_decode_cp:
+        return None
+    tp = get_tensor_model_parallel_world_size()
+    text_config = vllm_config.model_config.hf_text_config
+    total_heads = text_config.sparse_attention_config["sparse_num_index_heads"]
+    why = None
+    if tp not in SUPPORTED_TP_SIZES:
+        why = f"tensor parallel size {tp} (needs one of {SUPPORTED_TP_SIZES})"
+    elif total_heads != NUM_INDEX_HEADS or num_index_heads != 1:
+        why = f"{total_heads} index heads, {num_index_heads} per rank"
+    elif vllm_config.parallel_config.decode_context_parallel_size > 1:
+        why = "decode context parallelism is enabled"
+    elif max_decode_query_len > DECODE_QUERY_LEN:
+        why = f"decode query length {max_decode_query_len} > {DECODE_QUERY_LEN}"
+    elif not direct_cp_enabled(get_tp_group(), torch.int32, use_direct=None):
+        why = "the tensor-parallel ranks have no symmetric memory between them"
+    if why is not None:
+        logger.warning(
+            "indexer_decode_cp ignored: %s; using the stock NVFP4 indexer decode.",
+            why,
+        )
+        return None
+    sched = vllm_config.scheduler_config
+    _STATE = Nvfp4IndexerDecodeCP(
+        max_reqs=sched.max_num_seqs,
+        device=torch.device("cuda", torch.accelerator.current_device_index()),
+    )
+    logger.info_once(
+        "MiniMax M3 NVFP4 indexer: decode context parallelism over %d TP ranks "
+        "(each scores 1/%d of the index-K pages for all %d index heads).",
+        tp,
+        tp,
+        NUM_INDEX_HEADS,
+    )
+    return _STATE
diff --git a/vllm/models/minimax_m3/nvidia/model.py b/vllm/models/minimax_m3/nvidia/model.py
index d67e9ea19d90..83eae0bef382 100644
--- a/vllm/models/minimax_m3/nvidia/model.py
+++ b/vllm/models/minimax_m3/nvidia/model.py
@@ -13,16 +13,24 @@
 """
 
 from collections.abc import Iterable
+from typing import TYPE_CHECKING
 
 import torch
 from torch import nn
 from transformers import PreTrainedConfig
 
+import vllm.envs as envs
 from vllm import _custom_ops as ops
 from vllm.compilation.breakable_cudagraph import eager_break_during_capture
 from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config
-from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
-from vllm.forward_context import get_forward_context
+from vllm.distributed import (
+    get_ep_group,
+    get_pp_group,
+    get_tensor_model_parallel_rank,
+    get_tensor_model_parallel_world_size,
+)
+from vllm.forward_context import get_forward_context, is_forward_context_available
+from vllm.logger import init_logger
 from vllm.model_executor.layers.activation import SiluAndMulWithClamp
 from vllm.model_executor.layers.attention import Attention
 from vllm.model_executor.layers.attention.attention import set_default_quant_scales
@@ -35,6 +43,28 @@
     GateLinear,
     fused_moe_make_expert_params_mapping,
 )
+from vllm.model_executor.layers.fused_moe.deep_gemm_mega_moe import (
+    MEGA_MOE_NVFP4,
+    DeepGemmMegaMoEExperts,
+    MegaMoEActivation,
+)
+from vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe import (
+    TrtLlmNvFp4ExpertsModular,
+    TrtLlmNvFp4ExpertsMonolithic,
+)
+from vllm.model_executor.layers.fused_moe.moe_output import (
+    MoEOutput,
+    UnfinalizedMoEOutput,
+)
+from vllm.model_executor.layers.fused_moe.router.fused_topk_bias_router import (
+    fused_topk_bias,
+)
+from vllm.model_executor.layers.fused_moe.runner.moe_runner import _unpack
+from vllm.model_executor.layers.fusion.quant_activation import (
+    QuantizedActivation,
+    as_quantized_activation,
+    get_input_quant_key,
+)
 from vllm.model_executor.layers.linear import (
     MergedColumnParallelLinear,
     MinimaxM3QKVParallelLinearWithIndexer,
@@ -43,6 +73,13 @@
 )
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
 from vllm.model_executor.layers.quantization import QuantizationConfig
+from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
+    mxfp8_e4m3_quantize,
+)
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    kMxfp8Dynamic,
+    kNvfp4DynamicToken,
+)
 from vllm.model_executor.layers.rotary_embedding import get_rope
 from vllm.model_executor.layers.vocab_parallel_embedding import (
     ParallelLMHead,
@@ -63,6 +100,7 @@
     AutoWeightsLoader,
     PPMissingLayer,
     WeightsMapper,
+    extract_layer_index,
     init_vllm_registered_model,
     is_pp_missing_parameter,
     make_empty_intermediate_tensors_factory,
@@ -71,6 +109,12 @@
     spec_decode_needs_target_embed,
 )
 from vllm.model_executor.models.vision import run_dp_sharded_mrope_vision_model
+from vllm.models.common.ops.sequence_parallel import (
+    sp_all_gather,
+    sp_padding_mask,
+    sp_reduce_scatter,
+    sp_shard,
+)
 from vllm.models.minimax_m3.common.encoder_cudagraph import (
     MiniMaxM3EncoderCudaGraphMixin,
 )
@@ -87,7 +131,9 @@
 )
 from vllm.models.minimax_m3.common.vision_tower import MiniMaxVLVisionModel
 from vllm.multimodal import MULTIMODAL_REGISTRY
+from vllm.platforms import current_platform
 from vllm.sequence import IntermediateTensors
+from vllm.utils.flashinfer import is_flashinfer_autotuning
 from vllm.utils.torch_utils import kv_cache_dtype_str_to_dtype
 from vllm.v1.kv_cache_interface import (
     FullAttentionSpec,
@@ -95,6 +141,52 @@
     get_kv_quant_mode,
 )
 
+if TYPE_CHECKING:
+    from vllm.model_executor.kernels.linear.cute_dsl.gemm_rs_ar import GemmRsAr
+
+logger = init_logger(__name__)
+
+# Largest batch whose MoE top-k finalize runs inside the next layer's fused
+# all-reduce + input norm + MXFP8 quantization. Larger batches finalize in the
+# MoE and use the regular fused all-reduce + norm.
+FUSED_MOE_FINALIZE_MAX_TOKENS = 32
+
+_QKV_FUSED_EPILOGUE_MIN_TOKENS = 512
+
+# Smallest batch whose o_proj/down_proj runs the fused GEMM-RS kernel. Below
+# it, the GEMM followed by sp_reduce_scatter is as fast (TP4 GB300, see
+# benchmarks/kernels/benchmark_gemm_rs_ar.py --model minimax_m3).
+_GEMM_RS_MIN_TOKENS = 1024
+
+
+def _supports_qkv_fused_epilogue(attn: "MiniMaxM3SparseAttention") -> bool:
+    """Whether minimax_m3_qkv_index_fused_epilogue can replace this layer's
+    qkv_proj + norm/RoPE/cache-insert chain."""
+    from vllm.model_executor.kernels.linear.mxfp8.flashinfer import (
+        FlashInferCutedslMxfp8LinearKernel,
+        FlashInferCutlassMxfp8LinearKernel,
+    )
+
+    if not current_platform.is_device_capability_family(100):
+        return False
+    # Both keep the weight K-major and its scales F8_128x4-swizzled.
+    kernel = getattr(attn.qkv_proj.quant_method, "kernel", None)
+    if not isinstance(
+        kernel, (FlashInferCutedslMxfp8LinearKernel, FlashInferCutlassMxfp8LinearKernel)
+    ):
+        return False
+    nkv = attn.num_kv_heads
+    return (
+        attn.use_nvfp4_kv
+        and attn.indexer_kv_dtype == "nvfp4"
+        and nkv in (1, 2, 4)
+        and attn.num_heads == 16 * nkv
+        and attn.num_idx_heads == nkv
+        and attn.head_dim == attn.idx_head_dim == 128
+        and attn.rotary_emb.rotary_dim == 64
+        and attn.hidden_size % 128 == 0
+    )
+
 
 def _sparse_attention_layer_ids(config: PreTrainedConfig) -> set[int]:
     """Layer ids whose attention runs the extra sparse "index" branch."""
@@ -115,8 +207,164 @@ def _is_moe_layer(config: PreTrainedConfig, layer_id: int) -> bool:
     return moe_layer_freq[layer_id] != 0
 
 
+def _use_mega_moe(vllm_config: VllmConfig) -> bool:
+    """Whether MoE layers run the DeepGEMM NVFP4 MegaMoE kernel.
+
+    Opt in with ``--kernel-config '{"moe_backend": "deep_gemm_mega_moe"}'``. Only
+    checkpoints with NVFP4 routed experts (MiniMax-M3.1) are supported.
+    """
+    if vllm_config.kernel_config.moe_backend != "deep_gemm_mega_moe":
+        return False
+    moe_quant_algo = getattr(vllm_config.quant_config, "moe_quant_algo", None)
+    if moe_quant_algo != "NVFP4":
+        raise NotImplementedError(
+            "MiniMax M3 MegaMoE requires a checkpoint with NVFP4 routed experts "
+            f"(moe_quant_algo: NVFP4), got moe_quant_algo={moe_quant_algo!r}. "
+            "Pick a different moe backend for this checkpoint."
+        )
+    if not vllm_config.parallel_config.enable_expert_parallel:
+        raise NotImplementedError(
+            "MiniMax M3 MegaMoE requires expert parallel. Enable it with "
+            "--enable-expert-parallel."
+        )
+    return True
+
+
+def _use_sequence_parallel(vllm_config: VllmConfig) -> bool:
+    """Whether the residual stream stays sharded by token across TP ranks.
+
+    MegaMoE routes each rank's token shard, so the decoder keeps that shard
+    between layers: attention all-gathers its input and reduce-scatters its
+    output, replacing the two all-reduces per layer.
+    """
+    if vllm_config.kernel_config.moe_backend != "deep_gemm_mega_moe":
+        return False
+    parallel_config = vllm_config.parallel_config
+    return (
+        parallel_config.pipeline_parallel_size == 1
+        and parallel_config.enable_expert_parallel
+        and parallel_config.tensor_parallel_size > 1
+    )
+
+
+def _maybe_init_gemm_rs(vllm_config: VllmConfig, use_sequence_parallel: bool) -> bool:
+    if not (use_sequence_parallel and envs.VLLM_ENABLE_GEMM_RS):
+        return False
+    max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens
+    if max_num_batched_tokens < _GEMM_RS_MIN_TOKENS:
+        logger.info_once(
+            "MiniMax-M3 GEMM-RS is disabled because max_num_batched_tokens=%d is "
+            "below its %d-token dispatch threshold.",
+            max_num_batched_tokens,
+            _GEMM_RS_MIN_TOKENS,
+        )
+        return False
+
+    from vllm.model_executor.kernels.linear.cute_dsl.gemm_rs_ar import (
+        maybe_init_gemm_rs_ar,
+    )
+
+    hidden_size = vllm_config.model_config.hf_text_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 MiniMax-M3 GEMM-RS, set VLLM_ENABLE_GEMM_RS=0.")
+    return True
+
+
+def _bind_gemm_rs(linear: RowParallelLinear) -> "GemmRsAr | None":
+    from vllm.model_executor.kernels.linear.cute_dsl.gemm_rs_ar import (
+        get_gemm_rs_ar,
+    )
+
+    gemm_rs = get_gemm_rs_ar()
+    if gemm_rs.can_run(linear):
+        return gemm_rs
+    gemm_rs.warn_incompatible_projection()
+    return None
+
+
+def _row_parallel_reduce_scatter(
+    x: torch.Tensor | QuantizedActivation,
+    linear: RowParallelLinear,
+    gemm_rs: "GemmRsAr | None",
+) -> torch.Tensor:
+    """``linear`` followed by a reduce-scatter to this rank's token shard."""
+    num_tokens = x.orig_shape[0] if isinstance(x, QuantizedActivation) else x.shape[0]
+    if gemm_rs is not None and num_tokens >= _GEMM_RS_MIN_TOKENS:
+        return gemm_rs.apply(x, linear)
+    x, _ = linear(x)
+    return sp_reduce_scatter(x)
+
+
+# MiniMax-M3.1 QAT quantizes the FC2 input with a fixed global scale of 16;
+# checkpoints carry a placeholder ``w2.input_global_scale`` of 1.0.
+MEGA_MOE_FC2_INPUT_SCALE = 1.0 / 16.0
+
+
+def _mega_moe_expert_params_mapping(
+    num_experts: int,
+) -> list[tuple[str, str, int, str]]:
+    """Map ``experts.{e}.w{1,2,3}.*`` to the MegaMoE ``w13_*``/``w2_*`` params."""
+    return [
+        (
+            "experts.w2_" if shard_id == "w2" else "experts.w13_",
+            f"experts.{expert_id}.{shard_id}.",
+            expert_id,
+            shard_id,
+        )
+        for expert_id in range(num_experts)
+        for shard_id in ("w1", "w2", "w3")
+    ]
+
+
+def _mega_moe_expert_weight(
+    name: str, weight: torch.Tensor
+) -> tuple[str, torch.Tensor] | None:
+    """Convert a compressed-tensors NVFP4 routed-expert tensor to the ModelOpt
+    convention ``DeepGemmMegaMoEExperts`` loads, or None to drop it.
+
+    The global scales become dequantization scales. The FC1 input is quantized
+    per token, so its input scale is dropped, and the FC2 input scale is the
+    fixed QAT scale whatever the checkpoint says.
+    """
+    prefix, _, suffix = name.rpartition(".")
+    if suffix == "weight_packed":
+        return f"{prefix}.weight", weight
+    if suffix == "weight_global_scale":
+        return f"{prefix}.weight_scale_2", 1.0 / weight.float()
+    if suffix == "input_global_scale":
+        if not prefix.endswith(".w2"):
+            return None
+        scale = torch.full_like(weight, MEGA_MOE_FC2_INPUT_SCALE, dtype=torch.float32)
+        return f"{prefix}.input_scale", scale
+    return name, weight
+
+
+def _consumes_mxfp8_128x4(linear: nn.Module) -> bool:
+    """Whether ``linear`` takes a pre-quantized MXFP8 activation whose UE8M0
+    scales are in FlashInfer's 128x4 swizzled layout (what the MXFP8 norm
+    kernels write)."""
+    if get_input_quant_key(linear) != kMxfp8Dynamic:
+        return False
+    from vllm.model_executor.kernels.linear.mxfp8.flashinfer import (
+        FlashInferCutedslMxfp8LinearKernel,
+        FlashInferCutlassMxfp8LinearKernel,
+    )
+
+    kernel = getattr(getattr(linear, "quant_method", None), "kernel", None)
+    return type(kernel) in (
+        FlashInferCutedslMxfp8LinearKernel,
+        FlashInferCutlassMxfp8LinearKernel,
+    )
+
+
 class MiniMAXGemmaRMSNorm(nn.Module):
-    """Gemma-style RMS normalization backed by FlashInfer kernels.
+    """Gemma-style RMS normalization.
+
+    2-D inputs run the Triton kernel in ``ops/norm_mxfp8.py``; with
+    ``mxfp8=True`` the normed activation is emitted as MXFP8 (a
+    ``QuantizedActivation``) for a consumer linear instead of BF16. Other
+    shapes (the dspark drafter's per-head norms) keep the FlashInfer kernels.
 
     When ``residual`` is given, the fused add + norm runs in place and the
     updated ``(x, residual)`` pair is returned.
@@ -135,7 +383,33 @@ def forward(
         self,
         x: torch.Tensor,
         residual: torch.Tensor | None = None,
+        mxfp8: bool = False,
     ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        if x.dim() == 2:
+            from vllm.models.minimax_m3.nvidia.ops.norm_mxfp8 import (
+                gemma_rmsnorm_mxfp8,
+            )
+
+            out, scale = gemma_rmsnorm_mxfp8(
+                x,
+                self.weight,
+                self.variance_epsilon,
+                residual=residual,
+                mxfp8=mxfp8,
+            )
+            if mxfp8:
+                out = QuantizedActivation(
+                    data=out,
+                    scale=scale,
+                    orig_dtype=x.dtype,
+                    orig_shape=x.shape,
+                    quant_key=kMxfp8Dynamic,
+                )
+            if residual is None:
+                return out
+            return out, residual
+
+        assert not mxfp8
         from flashinfer.norm import gemma_fused_add_rmsnorm, gemma_rmsnorm
 
         if residual is None:
@@ -155,14 +429,19 @@ def __init__(
         intermediate_size: int,
         quant_config: QuantizationConfig | None = None,
         reduce_results: bool = True,
+        disable_tp: bool = False,
+        use_sequence_parallel: bool = False,
+        run_gemm_rs: bool = False,
         prefix: str = "",
     ) -> None:
         super().__init__()
+        self.use_sequence_parallel = use_sequence_parallel
         self.gate_up_proj = MergedColumnParallelLinear(
             config.hidden_size,
             [intermediate_size] * 2,
             bias=False,
             quant_config=quant_config,
+            disable_tp=disable_tp,
             prefix=f"{prefix}.gate_up_proj",
         )
         self.down_proj = RowParallelLinear(
@@ -171,8 +450,14 @@ def __init__(
             bias=False,
             quant_config=quant_config,
             reduce_results=reduce_results,
+            disable_tp=disable_tp,
             prefix=f"{prefix}.down_proj",
         )
+        self.gemm_rs = (
+            _bind_gemm_rs(self.down_proj)
+            if use_sequence_parallel and run_gemm_rs
+            else None
+        )
         if config.hidden_act != "swigluoai":
             raise ValueError(
                 f"Unsupported activation: {config.hidden_act}. "
@@ -184,10 +469,39 @@ def __init__(
             alpha=config.swiglu_alpha,
             beta=config.swiglu_beta,
         )
+        from vllm.models.minimax_m3.nvidia.ops import dense_gemm1_mxfp8
+
+        # At or below the shared-experts stream threshold the persistent kernel
+        # would contend with the routed experts on the aux stream.
+        self.fused_gemm1_min_tokens: int | None = None
+        if dense_gemm1_mxfp8.can_fuse_gemm1(
+            self.gate_up_proj, self.down_proj, self.act_fn
+        ):
+            self.fused_gemm1_min_tokens = max(
+                dense_gemm1_mxfp8.MIN_TOKENS,
+                envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD,
+            )
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        gate_up, _ = self.gate_up_proj(x)
-        x = self.act_fn(gate_up)
+        if self.use_sequence_parallel:
+            x = sp_all_gather(x)
+        num_tokens = x.shape[0]
+        # Run unfused under FlashInfer autotuning so gate_up_proj is tuned too
+        if (
+            self.fused_gemm1_min_tokens is not None
+            and num_tokens > self.fused_gemm1_min_tokens
+            and not is_flashinfer_autotuning()
+        ):
+            from vllm.models.minimax_m3.nvidia.ops.dense_gemm1_mxfp8 import (
+                fused_gemm1,
+            )
+
+            x = fused_gemm1(x, self.gate_up_proj, self.act_fn)
+        else:
+            gate_up, _ = self.gate_up_proj(x)
+            x = self.act_fn(gate_up)
+        if self.use_sequence_parallel:
+            return _row_parallel_reduce_scatter(x, self.down_proj, self.gemm_rs)
         x, _ = self.down_proj(x)
         return x
 
@@ -203,15 +517,22 @@ def __init__(
         quant_config: QuantizationConfig | None = None,
         reduce_results: bool = True,
         prefix: str = "",
+        use_mega_moe: bool = False,
+        input_is_sequence_sharded: bool = False,
     ) -> None:
         super().__init__()
         self.tp_size = get_tensor_model_parallel_world_size()
+        self.input_is_sequence_sharded = input_is_sequence_sharded
         if self.tp_size > config.num_local_experts:
             raise ValueError(
                 f"Tensor parallel size {self.tp_size} is greater than "
                 f"the number of experts {config.num_local_experts}."
             )
 
+        self.use_mega_moe = use_mega_moe
+        self.reduce_results = reduce_results
+        self.scoring_func = config.scoring_func
+        self.top_k = config.num_experts_per_tok
         self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0)
         self.n_shared_experts = getattr(config, "n_shared_experts", None)
 
@@ -239,6 +560,10 @@ def __init__(
         )
 
         self.shared_experts: MiniMaxM3MLP | None = None
+        if self.use_mega_moe:
+            self._init_mega_moe_experts(config, quant_config, prefix)
+            return
+
         if self.n_shared_experts:
             self.shared_experts = MiniMaxM3MLP(
                 config=config,
@@ -265,20 +590,127 @@ def __init__(
             swiglu_beta=config.swiglu_beta,
             routed_scaling_factor=self.routed_scaling_factor,
             router_logits_dtype=self.gate.out_dtype,
+            # The NVFP4 routed experts are QAT'd with a per-token FC1 input
+            # scale; their ``input_global_scale`` is a 1.0 placeholder that
+            # the quantization config cannot tell from a calibrated scale.
+            qat_activation_key=kNvfp4DynamicToken,
             shared_experts=self.shared_experts,
             quant_config=quant_config,
             reduce_results=reduce_results,
             prefix=f"{prefix}.experts",
         )
 
+    def _init_mega_moe_experts(
+        self,
+        config: PreTrainedConfig,
+        quant_config: QuantizationConfig | None,
+        prefix: str,
+    ) -> None:
+        if config.hidden_act != "swigluoai":
+            raise ValueError(
+                f"MiniMax M3 MegaMoE requires swigluoai, got {config.hidden_act}."
+            )
+        # Replicated hidden states are sharded here, so each rank routes only
+        # its own tokens through the EP-wide kernel.
+        self.shard_replicated_input = (
+            self.tp_size > 1 and not self.input_is_sequence_sharded
+        )
+
+        def replicated_shared_experts() -> MiniMaxM3MLP:
+            # Replicated, so the kernel can fuse it.
+            return MiniMaxM3MLP(
+                config=config,
+                intermediate_size=config.intermediate_size * self.n_shared_experts,
+                quant_config=quant_config,
+                reduce_results=False,
+                disable_tp=True,
+                prefix=f"{prefix}.shared_experts",
+            )
+
+        shared_experts = None
+        if self.n_shared_experts:
+            # Only describes the weights: a fused layer allocates its own.
+            with torch.device("meta"):
+                shared_experts = replicated_shared_experts()
+        ep_group = get_ep_group()
+        num_experts = config.num_local_experts
+        if num_experts % ep_group.world_size != 0:
+            raise ValueError(
+                f"MiniMax M3 MegaMoE needs num_experts={num_experts} divisible "
+                f"by the EP size {ep_group.world_size}."
+            )
+        num_local_experts = num_experts // ep_group.world_size
+        self.experts = DeepGemmMegaMoEExperts(
+            get_current_vllm_config(),
+            num_experts=num_experts,
+            num_local_experts=num_local_experts,
+            experts_start_idx=ep_group.rank_in_group * num_local_experts,
+            top_k=self.top_k,
+            hidden_size=config.hidden_size,
+            intermediate_size=config.intermediate_size,
+            activation=MegaMoEActivation(
+                "swiglu",
+                clamp=config.swiglu_limit,
+                alpha=config.swiglu_alpha,
+                beta=config.swiglu_beta,
+            ),
+            weight_format=MEGA_MOE_NVFP4,
+            shared_experts=shared_experts,
+            prefix=f"{prefix}.experts",
+        )
+        if shared_experts is not None and not self.experts.has_fused_shared_experts:
+            self.shared_experts = replicated_shared_experts()
+
     @staticmethod
     def ebias_weight_loader(param: nn.Parameter, loaded_weight: torch.Tensor) -> None:
         assert param.size() == loaded_weight.size()
         param.data.copy_(loaded_weight.to(torch.float32))
 
+    @property
+    def ffn_all_reduce_deferred(self) -> bool:
+        if self.use_mega_moe:
+            return not self.reduce_results
+        return self.experts.moe_config.skip_final_all_reduce
+
+    def defer_finalize(self, max_num_tokens: int) -> bool:
+        """Leave the routed top-k reduction of calls up to ``max_num_tokens``
+        to the consumer; return whether the experts accepted.
+
+        Only the TRTLLM-Gen NVFP4 experts can stop after GEMM2.
+        """
+        if self.use_mega_moe:
+            return False
+        experts_cls = getattr(
+            self.experts.routed_experts.quant_method, "experts_cls", None
+        )
+        if experts_cls not in (TrtLlmNvFp4ExpertsMonolithic, TrtLlmNvFp4ExpertsModular):
+            return False
+        moe_config = self.experts.moe_config
+        moe_config.defer_moe_finalize(max_num_tokens)
+        return moe_config.use_deferred_moe_finalize
+
+    def defers_finalize(self, num_tokens: int) -> bool:
+        return (
+            not self.use_mega_moe
+            and self.experts.moe_config.should_defer_moe_finalize(num_tokens)
+        )
+
+    def forward_unfinalized(self, hidden_states: torch.Tensor) -> MoEOutput:
+        """``forward`` with the routed top-k reduction and all-reduce left open."""
+        router_logits, _ = self.gate(hidden_states)
+        # The runner's custom op returns tensors only, so run its body directly.
+        shared_output, routed = _unpack(
+            self.experts._forward_impl(hidden_states, router_logits, hidden_states)
+        )
+        assert shared_output is not None
+        assert isinstance(routed, UnfinalizedMoEOutput)
+        return MoEOutput(routed=routed, shared_output=shared_output)
+
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         num_tokens, hidden_dim = hidden_states.shape
         hidden_states = hidden_states.view(-1, hidden_dim)
+        if self.use_mega_moe:
+            return self._forward_mega_moe(hidden_states)
 
         # router_logits: (num_tokens, n_experts); GateLinear casts to fp32.
         router_logits, _ = self.gate(hidden_states)
@@ -288,6 +720,54 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
 
         return final_hidden_states.view(num_tokens, hidden_dim)
 
+    def _forward_mega_moe(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        if not self.shard_replicated_input:
+            return self._forward_mega_moe_local(hidden_states)
+
+        num_tokens = hidden_states.shape[0]
+        # Routing and the experts read the padding mask from the forward
+        # context; hand them this rank's shard of it.
+        forward_context = None
+        if envs.VLLM_MOE_SKIP_PADDING and is_forward_context_available():
+            forward_context = get_forward_context()
+            is_padding = forward_context.is_padding
+            forward_context.is_padding = sp_padding_mask(
+                is_padding[:num_tokens] if is_padding is not None else None,
+                hidden_states,
+            )
+        try:
+            local_output = self._forward_mega_moe_local(sp_shard(hidden_states))
+        finally:
+            if forward_context is not None:
+                forward_context.is_padding = is_padding
+        if self.reduce_results:
+            return sp_all_gather(local_output)[:num_tokens]
+
+        # Leave the other ranks' rows zero, so the deferred TP all-reduce (fused
+        # into the next RMSNorm) gathers them.
+        shard_size, hidden_dim = local_output.shape
+        tp_rank = get_tensor_model_parallel_rank()
+        output = local_output.new_zeros(shard_size * self.tp_size, hidden_dim)
+        output[tp_rank * shard_size : (tp_rank + 1) * shard_size] = local_output
+        return output[:num_tokens]
+
+    def _forward_mega_moe_local(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        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.e_score_correction_bias.data
+            if self.e_score_correction_bias is not None
+            else None,
+            topk=self.top_k,
+            renormalize=True,
+            routed_scaling_factor=self.routed_scaling_factor,
+        )
+        return self.experts(
+            hidden_states, topk_weights, topk_ids, shared_experts=self.shared_experts
+        )
+
 
 class MiniMaxM3Attention(nn.Module):
     """Dense attention with per-head QK norm and partial RoPE."""
@@ -299,6 +779,8 @@ def __init__(
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
         cache_config: CacheConfig | None = None,
+        use_sequence_parallel: bool = False,
+        run_gemm_rs: bool = False,
     ) -> None:
         super().__init__()
         self.hidden_size = config.hidden_size
@@ -327,9 +809,9 @@ def __init__(
             quant_config=quant_config,
             prefix=f"{prefix}.qkv_proj",
         )
-        # reduce_results=False: the attention all-reduce is fused with the
-        # following post_attention_layernorm (GemmaRMSNorm) in the decoder layer
-        # via fused_allreduce_gemma_rms_norm.
+        # reduce_results=False: under sequence parallelism o_proj reduce-scatters
+        # in forward; otherwise the all-reduce is fused with the following
+        # post_attention_layernorm via fused_allreduce_gemma_rms_norm.
         self.o_proj = RowParallelLinear(
             self.total_num_heads * self.head_dim,
             self.hidden_size,
@@ -338,6 +820,12 @@ def __init__(
             quant_config=quant_config,
             prefix=f"{prefix}.o_proj",
         )
+        self.use_sequence_parallel = use_sequence_parallel
+        self.gemm_rs = (
+            _bind_gemm_rs(self.o_proj)
+            if use_sequence_parallel and run_gemm_rs
+            else None
+        )
 
         # Per-head QK norm (qk_norm_type == "per_head", use_gemma_norm == True).
         self.q_norm = MiniMAXGemmaRMSNorm(self.head_dim, eps=config.rms_norm_eps)
@@ -385,6 +873,8 @@ def forward(
         )
         q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
         attn_output = self.attn(q, k, v)
+        if self.use_sequence_parallel:
+            return _row_parallel_reduce_scatter(attn_output, self.o_proj, self.gemm_rs)
         output, _ = self.o_proj(attn_output)
         return output
 
@@ -413,6 +903,8 @@ def __init__(
         prefix: str = "",
         cache_config: CacheConfig | None = None,
         topk_indices_buffer: torch.Tensor | None = None,
+        use_sequence_parallel: bool = False,
+        run_gemm_rs: bool = False,
     ) -> None:
         super().__init__()
         self.hidden_size = config.hidden_size
@@ -453,9 +945,9 @@ def __init__(
             quant_config=quant_config,
             prefix=f"{prefix}.qkv_proj",
         )
-        # reduce_results=False: the attention all-reduce is fused with the
-        # following post_attention_layernorm (GemmaRMSNorm) in the decoder layer
-        # via fused_allreduce_gemma_rms_norm.
+        # reduce_results=False: under sequence parallelism o_proj reduce-scatters
+        # in forward; otherwise the all-reduce is fused with the following
+        # post_attention_layernorm via fused_allreduce_gemma_rms_norm.
         self.o_proj = RowParallelLinear(
             self.total_num_heads * self.head_dim,
             self.hidden_size,
@@ -464,6 +956,12 @@ def __init__(
             quant_config=quant_config,
             prefix=f"{prefix}.o_proj",
         )
+        self.use_sequence_parallel = use_sequence_parallel
+        self.gemm_rs = (
+            _bind_gemm_rs(self.o_proj)
+            if use_sequence_parallel and run_gemm_rs
+            else None
+        )
 
         # Per-head QK norm (qk_norm_type == "per_head", use_gemma_norm == True).
         self.q_norm = MiniMAXGemmaRMSNorm(self.head_dim, eps=config.rms_norm_eps)
@@ -556,6 +1054,26 @@ def __init__(
         compilation_config.static_forward_context[self.layer_name] = self
         self.kv_cache = torch.tensor([])  # replaced by bind_kv_cache
 
+        self.use_qkv_fused_epilogue = _supports_qkv_fused_epilogue(self)
+        if self.use_qkv_fused_epilogue:
+            from vllm.models.minimax_m3.nvidia.ops.qkv_index_fused_epilogue import (
+                register_qkv_index_fused_epilogue_warmup,
+            )
+
+            logger.info_once("MiniMax-M3: QKV+indexer GEMM with fused epilogue enabled")
+            block_size = vllm_config.cache_config.block_size
+            # The query is E4M3-only when CUTLASS decode serves the layer.
+            q_is_fp8 = getattr(self.impl, "use_cutlass_decode", False)
+            if vllm_config.kernel_config.enable_jit_warmup and block_size is not None:
+                register_qkv_index_fused_epilogue_warmup(
+                    has_q=not q_is_fp8,
+                    has_qfp8=q_is_fp8,
+                    nkv=self.num_kv_heads,
+                    block_size=block_size,
+                    weight_bytes=self.qkv_proj.output_size_per_partition
+                    * self.hidden_size,
+                )
+
     def get_attn_backend(self) -> type[MiniMaxM3SparseBackend]:
         return self.attn_backend
 
@@ -570,37 +1088,74 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None:
             kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype),
         )
 
-    def _allocate_query_fp8(self, qkv: torch.Tensor) -> torch.Tensor | None:
-        if not getattr(self.impl, "use_cutlass_decode", False):
-            return None
+    def _qkv_proj_fused_epilogue(
+        self,
+        hidden_states: torch.Tensor | QuantizedActivation,
+        positions: torch.Tensor,
+        q: torch.Tensor,
+        index_q: torch.Tensor,
+        slot_mapping: torch.Tensor,
+        index_slot_mapping: torch.Tensor,
+    ) -> None:
+        """qkv_proj + QK-norm/RoPE + KV/index-cache insert as one MXFP8 kernel,
+        writing q (bf16 or E4M3, as allocated) and index_q."""
+        from vllm.models.minimax_m3.nvidia.ops.qkv_index_fused_epilogue import (
+            minimax_m3_qkv_index_fused_epilogue,
+        )
+
+        weight = self.qkv_proj.weight
+        if weight.shape[1] != self.hidden_size:
+            weight = weight.t()  # the CuTe-DSL kernel keeps a [K, N] view
+        qa = as_quantized_activation(hidden_states, kMxfp8Dynamic)
+        if qa is not None:
+            xq, xsf = qa.data, qa.scale
+        else:
+            assert isinstance(hidden_states, torch.Tensor)
+            xq, xsf = mxfp8_e4m3_quantize(hidden_states, is_sf_swizzled_layout=True)
+        q_is_fp8 = q.dtype == torch.float8_e4m3fn
+        minimax_m3_qkv_index_fused_epilogue(
+            xq,
+            xsf,
+            weight,
+            self.qkv_proj.weight_scale,
+            positions,
+            self.rotary_emb.cos_sin_cache,
+            self.q_norm.weight,
+            self.k_norm.weight,
+            self.index_q_norm.weight,
+            self.index_k_norm.weight,
+            self.q_norm.variance_epsilon,
+            slot_mapping=slot_mapping,
+            index_slot_mapping=index_slot_mapping,
+            kv_cache=self.kv_cache,
+            index_cache=self.indexer.index_cache.kv_cache,
+            q_out=None if q_is_fp8 else q,
+            index_q_out=index_q,
+            q_fp8_out=q if q_is_fp8 else None,
+            q_fp8_scale=self._q_scale_float,
+            k_scale=self._k_scale_float,
+            v_scale=self._v_scale_float,
+        )
+
+    def _allocate_query(
+        self, num_tokens: int, dtype: torch.dtype, device: torch.device
+    ) -> torch.Tensor:
+        # CUTLASS decode reads an E4M3 query, so the insert emits only that.
+        fp8 = getattr(self.impl, "use_cutlass_decode", False)
         return torch.empty(
-            (qkv.shape[0], self.q_size),
-            dtype=torch.float8_e4m3fn,
-            device=qkv.device,
+            (num_tokens, self.q_size),
+            dtype=torch.float8_e4m3fn if fp8 else dtype,
+            device=device,
         )
 
     def forward(
         self,
         positions: torch.Tensor,
-        hidden_states: torch.Tensor,
+        hidden_states: torch.Tensor | QuantizedActivation,
     ) -> torch.Tensor:
-        # Single fused projection emitting [q | k | v | index_q | index_k].
-        qkv, _ = self.qkv_proj(hidden_states)
-
-        # Horizontally-fused per-head Gemma QK-norm + partial NeoX RoPE on the
-        # main (q/k) and index (index_q/index_k) branches, all read straight out
-        # of the single fused ``qkv`` tensor (the "5 results").  Once the paged
-        # caches are bound the kernel also inserts k/v and the index key into
-        # them; the initial memory-profiling run (caches unbound, no slot_mapping)
-        # short-circuits to zeros below. k/v and index_k are rewritten in place
-        # inside qkv (and scatter-inserted into the caches); q and index_q are
-        # de-interleaved
-        # straight into the dedicated contiguous ``q``/``index_q`` buffers below.
-
         cos_sin_cache = self.rotary_emb.cos_sin_cache
         rotary_dim = self.rotary_emb.rotary_dim
         eps = self.q_norm.variance_epsilon
-        num_tokens = qkv.shape[0]
 
         fwd_slot_mapping = get_forward_context().slot_mapping
         if (
@@ -608,49 +1163,86 @@ def forward(
             or self.layer_name not in fwd_slot_mapping
         ):
             # Memory-profiling run: caches not yet bound, slot_mapping is empty.
-            return qkv.new_zeros((num_tokens, self.hidden_size))
+            qkv, _ = self.qkv_proj(hidden_states)
+            self.indexer.reserve_workspace()
+            output = qkv.new_zeros((qkv.shape[0], self.hidden_size))
+            return sp_shard(output) if self.use_sequence_parallel else output
 
         main_slot_mapping = fwd_slot_mapping[self.layer_name]
         index_slot_mapping = fwd_slot_mapping[self.indexer.index_cache.prefix]
-        q = qkv.new_empty((num_tokens, self.q_size))
-        # With query_fp8, q is emitted only in fp8; the attend dequantizes it
-        # into ``q`` for the tokens that run on Triton decode.
-        query_fp8 = self._allocate_query_fp8(qkv)
-        # index_q matches the index-K cache dtype (e4m3 for the fp8 score path);
-        # the fused kernel emits fp8 directly when this buffer is e4m3.
-        index_q = qkv.new_empty(
-            (num_tokens, self.index_q_size),
-            dtype=self.indexer.index_cache.dtype,
+        # FlashInfer autotunes the GEMMs its dummy runs call: take the unfused
+        # path there so qkv_proj is tuned for the M < 512 fallback.
+        if isinstance(hidden_states, QuantizedActivation):
+            num_tokens = hidden_states.orig_shape[0]
+            dtype, device = hidden_states.orig_dtype, hidden_states.data.device
+        else:
+            num_tokens = hidden_states.shape[0]
+            dtype, device = hidden_states.dtype, hidden_states.device
+        fused = (
+            self.use_qkv_fused_epilogue
+            and num_tokens >= _QKV_FUSED_EPILOGUE_MIN_TOKENS
+            and not is_flashinfer_autotuning()
         )
-        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,
-            main_slot_mapping,
-            index_slot_mapping,
-            self.kv_cache,
-            self.indexer.index_cache.kv_cache,
-            self.kv_cache.size(2),  # paged-cache block size
-            q if query_fp8 is None else None,
-            index_q,
-            self.kv_cache_dtype,
-            q_fp8_out=query_fp8,
-            q_fp8_scale=self._q_scale_float,
-            kv_k_scale=self._k_scale if self.use_nvfp4_kv else None,
-            kv_v_scale=self._v_scale if self.use_nvfp4_kv else None,
+        if not fused:
+            # Single fused projection emitting [q | k | v | index_q | index_k].
+            qkv, _ = self.qkv_proj(hidden_states)
+        q = self._allocate_query(num_tokens, dtype, device)
+        q_is_fp8 = q.dtype == torch.float8_e4m3fn
+        # index_q is e4m3 for the fp8 and nvfp4 index caches, else bf16; the
+        # fused kernel emits fp8 directly when this buffer is e4m3.
+        index_q = torch.empty(
+            (num_tokens, self.index_q_size),
+            dtype=self.indexer.index_cache.query_dtype,
+            device=device,
         )
+        if fused:
+            self._qkv_proj_fused_epilogue(
+                hidden_states,
+                positions,
+                q,
+                index_q,
+                main_slot_mapping,
+                index_slot_mapping,
+            )
+        else:
+            # Horizontally-fused per-head Gemma QK-norm + partial NeoX RoPE on
+            # the main (q/k) and index (index_q/index_k) branches, all read
+            # straight out of qkv. k/v and index_k are scatter-inserted into the
+            # caches; q and index_q are de-interleaved into the dedicated
+            # contiguous buffers above. With an E4M3 q nothing reads qkv
+            # afterwards, so k and index_k are not rewritten in place.
+            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,
+                main_slot_mapping,
+                index_slot_mapping,
+                self.kv_cache,
+                self.indexer.index_cache.kv_cache,
+                self.kv_cache.size(2),  # paged-cache block size
+                None if q_is_fp8 else q,
+                index_q,
+                self.kv_cache_dtype,
+                q_fp8_out=q if q_is_fp8 else None,
+                q_fp8_scale=self._q_scale_float,
+                kv_k_scale=self._k_scale if self.use_nvfp4_kv else None,
+                kv_v_scale=self._v_scale if self.use_nvfp4_kv else None,
+                write_qkv=not q_is_fp8,
+            )
 
-        output = torch.empty_like(q)
-        attn_output = self._run_attention(q, query_fp8, index_q, output)
+        output = torch.empty((num_tokens, self.q_size), dtype=dtype, device=device)
+        attn_output = self._run_attention(q, index_q, output)
+        if self.use_sequence_parallel:
+            return _row_parallel_reduce_scatter(attn_output, self.o_proj, self.gemm_rs)
         output, _ = self.o_proj(attn_output)
         return output
 
@@ -658,7 +1250,6 @@ def forward(
     def _run_attention(
         self,
         query: torch.Tensor,
-        query_fp8: torch.Tensor | None,
         index_query: torch.Tensor,
         output: torch.Tensor,
     ) -> torch.Tensor:
@@ -666,13 +1257,7 @@ def _run_attention(
         # metadata and can't be captured into a cudagraph. The indexer writes its
         # top-k into the shared ``topk_indices_buffer``; the attend reads it back.
         self.indexer(index_query)
-        return self.impl.forward(
-            self,
-            query,
-            self.kv_cache,
-            output,
-            query_fp8=query_fp8,
-        )
+        return self.impl.forward(self, query, self.kv_cache, output)
 
 
 class MiniMaxM3DecoderLayer(nn.Module):
@@ -685,8 +1270,12 @@ def __init__(
         force_moe: bool = False,
         is_mtp_block: bool = False,
         topk_indices_buffer: torch.Tensor | None = None,
+        run_gemm_rs: bool = False,
     ) -> None:
         super().__init__()
+        self.use_sequence_parallel = (
+            _use_sequence_parallel(vllm_config) and not is_mtp_block
+        )
         if is_mtp_block:
             assert vllm_config.speculative_config is not None
             config = vllm_config.speculative_config.draft_model_config.hf_config
@@ -717,6 +1306,8 @@ def __init__(
                 prefix=f"{prefix}.self_attn",
                 cache_config=cache_config,
                 topk_indices_buffer=topk_indices_buffer,
+                use_sequence_parallel=self.use_sequence_parallel,
+                run_gemm_rs=run_gemm_rs,
             )
         else:
             self.self_attn = MiniMaxM3Attention(
@@ -725,6 +1316,8 @@ def __init__(
                 quant_config=quant_config,
                 prefix=f"{prefix}.self_attn",
                 cache_config=cache_config,
+                use_sequence_parallel=self.use_sequence_parallel,
+                run_gemm_rs=run_gemm_rs,
             )
 
         # Dense layers store the FFN under `mlp`; MoE layers under
@@ -737,12 +1330,19 @@ def __init__(
         )
         self.is_moe_layer = force_moe or _is_moe_layer(config, layer_id)
         if self.is_moe_layer:
+            use_mega_moe = _use_mega_moe(vllm_config)
+            if use_mega_moe and is_mtp_block:
+                raise NotImplementedError(
+                    "MiniMax M3 MegaMoE does not support MTP draft layers."
+                )
             self.block_sparse_moe = MiniMaxM3MoE(
                 config=config,
                 layer_id=layer_id,
                 quant_config=quant_config,
                 reduce_results=reduce_results,
                 prefix=f"{prefix}.block_sparse_moe",
+                use_mega_moe=use_mega_moe,
+                input_is_sequence_sharded=self.use_sequence_parallel,
             )
         else:
             self.mlp = MiniMaxM3MLP(
@@ -751,6 +1351,8 @@ def __init__(
                 quant_config=quant_config,
                 prefix=f"{prefix}.mlp",
                 reduce_results=reduce_results,
+                use_sequence_parallel=self.use_sequence_parallel,
+                run_gemm_rs=run_gemm_rs,
             )
 
         # config.use_gemma_norm is True for M3 -> Gemma-style RMSNorm.
@@ -764,27 +1366,73 @@ def __init__(
     def forward(
         self,
         positions: torch.Tensor,
-        hidden_states: torch.Tensor,
+        hidden_states: torch.Tensor | MoEOutput,
         residual: torch.Tensor | None,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
+    ) -> tuple[torch.Tensor | MoEOutput, torch.Tensor]:
+        hidden_states, residual = self.forward_input_norm(hidden_states, residual)
+        return self.forward_after_input_norm(positions, hidden_states, residual)
+
+    def forward_input_norm(
+        self, hidden_states: torch.Tensor | MoEOutput, residual: torch.Tensor | None
+    ) -> tuple[torch.Tensor | QuantizedActivation, torch.Tensor]:
+        if isinstance(hidden_states, MoEOutput):
+            # The preceding MoE left a small batch unfinalized: finalize,
+            # all-reduce, norm and MXFP8-quantize the qkv_proj input at once.
+            assert residual is not None
+            from vllm.models.minimax_m3.nvidia.ops.finalize_allreduce_norm import (
+                finalize_allreduce_gemma_rms_norm_mxfp8,
+            )
+
+            return finalize_allreduce_gemma_rms_norm_mxfp8(
+                hidden_states,
+                residual,
+                self.input_layernorm.weight,
+                self.input_layernorm.variance_epsilon,
+            )
         if self.fuse_input_allreduce and residual is not None:
             hidden_states, residual = fused_allreduce_gemma_rms_norm(
                 hidden_states, residual, self.input_layernorm
             )
         else:
+            # MXFP8 output when qkv_proj consumes it (checked per call: LoRA
+            # wrappers installed after construction change the key).
+            mxfp8 = _consumes_mxfp8_128x4(self.self_attn.qkv_proj)
             if residual is None:
                 residual = hidden_states
-                hidden_states = self.input_layernorm(hidden_states)
+                hidden_states = self.input_layernorm(hidden_states, mxfp8=mxfp8)
             else:
-                hidden_states, residual = self.input_layernorm(hidden_states, residual)
+                hidden_states, residual = self.input_layernorm(
+                    hidden_states, residual, mxfp8=mxfp8
+                )
+        return hidden_states, residual
+
+    def forward_after_input_norm(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor | QuantizedActivation,
+        residual: torch.Tensor,
+    ) -> tuple[torch.Tensor | MoEOutput, torch.Tensor]:
+        if self.use_sequence_parallel:
+            hidden_states = sp_all_gather(hidden_states)[: positions.shape[0]]
         hidden_states = self.self_attn(
             positions=positions,
             hidden_states=hidden_states,
         )
 
-        hidden_states, residual = fused_allreduce_gemma_rms_norm(
-            hidden_states, residual, self.post_attention_layernorm
-        )
+        if self.use_sequence_parallel:
+            # Attention already reduce-scattered its output to this rank's shard.
+            hidden_states, residual = self.post_attention_layernorm(
+                hidden_states, residual
+            )
+        else:
+            hidden_states, residual = fused_allreduce_gemma_rms_norm(
+                hidden_states, residual, self.post_attention_layernorm
+            )
+        if self.is_moe_layer and self.block_sparse_moe.defers_finalize(
+            hidden_states.shape[0]
+        ):
+            # The next layer's input norm finalizes it with the all-reduce.
+            return self.block_sparse_moe.forward_unfinalized(hidden_states), residual
         ffn = self.block_sparse_moe if self.is_moe_layer else self.mlp
         hidden_states = ffn(hidden_states)
         return hidden_states, residual
@@ -793,8 +1441,10 @@ def forward(
     def ffn_all_reduce_deferred(self) -> bool:
         """This layer's FFN output is left un-reduced; the caller fuses the
         all-reduce into the next RMSNorm."""
+        if self.use_sequence_parallel:
+            return False
         if self.is_moe_layer:
-            return self.block_sparse_moe.experts.moe_config.skip_final_all_reduce
+            return self.block_sparse_moe.ffn_all_reduce_deferred
         return not self.mlp.down_proj.reduce_results
 
 
@@ -808,12 +1458,13 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         config = vllm_config.model_config.hf_text_config
         quant_config = vllm_config.quant_config
         self.config = config
+        self.use_mega_moe = _use_mega_moe(vllm_config)
+        self.use_sequence_parallel = _use_sequence_parallel(vllm_config)
+        run_gemm_rs = _maybe_init_gemm_rs(vllm_config, self.use_sequence_parallel)
 
         self.vocab_size = config.vocab_size
 
-        if get_pp_group().is_first_rank or spec_decode_needs_target_embed(
-            vllm_config, include_mtp=vllm_config.use_v2_model_runner
-        ):
+        if get_pp_group().is_first_rank or spec_decode_needs_target_embed(vllm_config):
             self.embed_tokens = VocabParallelEmbedding(
                 config.vocab_size,
                 config.hidden_size,
@@ -851,6 +1502,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
                 vllm_config=vllm_config,
                 prefix=prefix,
                 topk_indices_buffer=self.topk_indices_buffer,
+                run_gemm_rs=run_gemm_rs,
             ),
             prefix=f"{prefix}.layers",
         )
@@ -871,6 +1523,54 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
             layer.fuse_input_allreduce = idx > 0 and prev_defers
             prev_defers = layer.ffn_all_reduce_deferred
         self.fuse_final_norm_allreduce = prev_defers
+        if current_platform.is_cuda():
+            self._init_fused_moe_finalize(vllm_config)
+
+    def _init_fused_moe_finalize(self, vllm_config: VllmConfig) -> None:
+        """Let small batches leave a MoE's top-k finalize to the next layer's
+        input norm, which already fuses the TP all-reduce, and have it quantize
+        the normed rows to the MXFP8 its qkv_proj reads.
+
+        Each MoE whose successor can take the output gets the hint here, and
+        ``defers_finalize`` decides per call. The final norm feeds the LM head
+        in BF16, so the last layer keeps its own finalize. Builds a collective
+        kernel: every TP rank decides alike.
+        """
+        from vllm.models.minimax_m3.nvidia.ops.finalize_allreduce_norm import (
+            init_fused_moe_finalize,
+            supports_fused_moe_finalize,
+        )
+
+        if not supports_fused_moe_finalize(vllm_config):
+            return
+        layers = self.layers[self.start_layer : self.end_layer]
+        deferring = [
+            layer.block_sparse_moe
+            for layer, successor in zip(layers, layers[1:])
+            if layer.is_moe_layer
+            and successor.fuse_input_allreduce
+            and get_input_quant_key(successor.self_attn.qkv_proj) == kMxfp8Dynamic
+            and layer.block_sparse_moe.defer_finalize(FUSED_MOE_FINALIZE_MAX_TOKENS)
+        ]
+        if not deferring:
+            return
+        # The modular experts return the router's FP32 routing weights; the
+        # monolithic kernel routes itself in BF16.
+        experts_cls = getattr(
+            deferring[0].experts.routed_experts.quant_method, "experts_cls", None
+        )
+        init_fused_moe_finalize(
+            hidden_size=self.config.hidden_size,
+            max_num_tokens=FUSED_MOE_FINALIZE_MAX_TOKENS,
+            top_k=self.config.num_experts_per_tok,
+            fp32_weights=experts_cls is TrtLlmNvFp4ExpertsModular,
+        )
+        logger.info_once(
+            "MiniMax M3: MoE top-k finalize fused into the next all-reduce + "
+            "input norm + MXFP8 quantization for up to %d tokens (%d layers).",
+            FUSED_MOE_FINALIZE_MAX_TOKENS,
+            len(deferring),
+        )
 
     def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
         return self.embed_tokens(input_ids)
@@ -893,24 +1593,47 @@ def forward(
             hidden_states = intermediate_tensors["hidden_states"]
             residual = intermediate_tensors["residual"]
 
+        num_tokens = hidden_states.shape[0]
+        forward_context = None
+        if self.use_sequence_parallel:
+            # MegaMoE reads the padding mask from the forward context; give it
+            # this rank's shard while the residual stream is sharded.
+            if envs.VLLM_MOE_SKIP_PADDING and is_forward_context_available():
+                forward_context = get_forward_context()
+                is_padding = forward_context.is_padding
+                forward_context.is_padding = sp_padding_mask(
+                    is_padding[:num_tokens] if is_padding is not None else None,
+                    hidden_states,
+                )
+            hidden_states = sp_shard(hidden_states)
+
         remote_aux = self.collect_remote_aux_hidden_states(intermediate_tensors)
+
+        # Aux hidden state k is the residual stream entering layer k. Under TP
+        # it is only complete after layer k's input norm (or the final norm)
+        # finishes the previous FFN's deferred all-reduce, so it is read there.
+        # Cloned, as the norms may update `residual` in place.
         aux_hidden_states: list[torch.Tensor] = []
-        if get_pp_group().is_first_rank:
-            self._maybe_add_hidden_state(
-                aux_hidden_states, self.start_layer, hidden_states, residual
-            )
-            # Preserve the entry tap before in-place residual updates.
-            if aux_hidden_states:
-                aux_hidden_states[0] = aux_hidden_states[0].clone()
         for idx, layer in enumerate(
             self.layers[self.start_layer : self.end_layer], start=self.start_layer
         ):
-            hidden_states, residual = layer(positions, hidden_states, residual)
-            self._maybe_add_hidden_state(
-                aux_hidden_states, idx + 1, hidden_states, residual
+            hidden_states, residual = layer.forward_input_norm(hidden_states, residual)
+            # A later PP stage's start boundary is captured by the stage before.
+            if idx in self.aux_hidden_state_layers and (
+                idx > self.start_layer or get_pp_group().is_first_rank
+            ):
+                aux_hidden_states.append(residual.clone())
+            hidden_states, residual = layer.forward_after_input_norm(
+                positions, hidden_states, residual
             )
+        if forward_context is not None:
+            forward_context.is_padding = is_padding
 
         if not get_pp_group().is_last_rank:
+            # Under PP every FFN reduces its own output: no all-reduce is pending.
+            self._maybe_add_hidden_state(
+                aux_hidden_states, self.end_layer, hidden_states, residual
+            )
             return IntermediateTensors(
                 {
                     "hidden_states": hidden_states,
@@ -920,11 +1643,22 @@ def forward(
             )
 
         if self.fuse_final_norm_allreduce:
-            hidden_states, _ = fused_allreduce_gemma_rms_norm(
+            hidden_states, residual = fused_allreduce_gemma_rms_norm(
                 hidden_states, residual, self.norm
             )
         else:
-            hidden_states, _ = self.norm(hidden_states, residual)
+            hidden_states, residual = self.norm(hidden_states, residual)
+        if self.end_layer in self.aux_hidden_state_layers:
+            aux_hidden_states.append(residual.clone())
+
+        if self.use_sequence_parallel and aux_hidden_states:
+            # One all-gather for the output and every aux hidden state.
+            hidden_size = hidden_states.shape[-1]
+            packed = torch.cat([hidden_states, *aux_hidden_states], dim=-1)
+            packed = sp_all_gather(packed)[:num_tokens]
+            hidden_states, *aux_hidden_states = packed.split(hidden_size, dim=-1)
+        elif self.use_sequence_parallel:
+            hidden_states = sp_all_gather(hidden_states)[:num_tokens]
 
         aux_hidden_states = remote_aux + aux_hidden_states
         if len(aux_hidden_states) > 0:
@@ -932,6 +1666,8 @@ def forward(
         return hidden_states
 
     def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
+        if self.use_mega_moe:
+            return _mega_moe_expert_params_mapping(self.config.num_local_experts)
         # Checkpoint experts use w1=gate, w2=down, w3=up.
         return fused_moe_make_expert_params_mapping(
             self,
@@ -975,6 +1711,23 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
             if "weight_scale_inv" in name:
                 name = name.replace("weight_scale_inv", "weight_scale")
 
+            if self.use_mega_moe and ".block_sparse_moe.experts." in name:
+                converted = _mega_moe_expert_weight(name, loaded_weight)
+                if converted is None:
+                    continue
+                name, loaded_weight = converted
+            if self.use_mega_moe and ".shared_experts." in name:
+                if is_pp_missing_parameter(name, self):
+                    continue
+                moe = self.layers[extract_layer_index(name)].block_sparse_moe
+                if moe.experts.has_fused_shared_experts:
+                    name = name.replace(".gate_proj.", ".w1.")
+                    name = name.replace(".up_proj.", ".w3.")
+                    loaded_params.add(
+                        moe.experts.load_shared_expert_weight(name, loaded_weight)
+                    )
+                    continue
+
             for param_name, weight_name, shard_id in stacked_params_mapping:
                 if weight_name not in name:
                     continue
diff --git a/vllm/models/minimax_m3/nvidia/msa_cutlass_sparse_decode.py b/vllm/models/minimax_m3/nvidia/msa_cutlass_sparse_decode.py
index 970b31c70233..f12232c0ee99 100644
--- a/vllm/models/minimax_m3/nvidia/msa_cutlass_sparse_decode.py
+++ b/vllm/models/minimax_m3/nvidia/msa_cutlass_sparse_decode.py
@@ -22,11 +22,22 @@
 # fmha_sm100 plans one row per query head. Keep every cached plan within the
 # fixed planner allocation used by the MSA decode kernel.
 _MAX_QUERY_HEAD_ROWS = 65536
+# The Q8KV4 kernel (every NVFP4 cache) plans without that allocation. Without
+# split KV it launches sequences x query tokens in gridDim.z, which also keeps
+# the sequence count under its 16-bit work-item field.
+_MAX_Q8KV4_QUERY_TOKENS = 65535
 _MAX_DECODE_QUERY_LEN = 32
 _DEFAULT_MIN_CUTLASS_BATCH_SIZE = 16
 _B300_EAGLE3_QUERY_HEAD_ROWS = 1024
 _B300_EAGLE3_MAX_QUERY_LEN = 8
 _B300_EAGLE3_HEAD_GEOMETRIES = {(64, 4), (32, 2), (16, 1)}
+# Q heads per KV head the Q8KV4 NVFP4 decode kernel serves.
+_Q8KV4_GQA_RATIOS = (8, 16)
+# Q8KV4 block-scale staging of the NVFP4 KV cache. MiniMax-M3.1 writes it with
+# unit global scales, so block scales are e4m3(amax / 6) and every
+# code * block_scale already fits E4M3: read it unstaged. Shift 3 is for block
+# scales that use the full E4M3 range next to a calibrated global scale.
+NVFP4_BLOCK_SCALE_SHIFT = 0
 
 
 def is_nvfp4_kv_cache(kv_cache_dtype: str) -> bool:
@@ -61,13 +72,37 @@ class MSACutlassDecodeMetadata:
     page_table: torch.Tensor
 
 
+def uses_q8kv4_decode(
+    *,
+    kv_cache_dtype: str,
+    num_q_heads: int,
+    num_kv_heads: int,
+) -> bool:
+    """Return whether sparse decode runs on the MSA Q8KV4 kernel: always for an
+    NVFP4 cache, which is never read with a BF16 query."""
+    if not is_nvfp4_kv_cache(kv_cache_dtype):
+        return False
+    if (
+        num_q_heads % num_kv_heads != 0
+        or num_q_heads // num_kv_heads not in _Q8KV4_GQA_RATIOS
+    ):
+        raise ValueError(
+            "MiniMax M3 NVFP4 KV cache decodes only on the MSA Q8KV4 kernel, "
+            f"which serves {_Q8KV4_GQA_RATIOS} Q heads per KV head; got "
+            f"{num_q_heads} Q heads over {num_kv_heads} KV heads"
+        )
+    return True
+
+
 @triton.jit
 def _update_runtime_metadata_kernel(
     seq_lens_ptr,
     kv_segment_lens_ptr,
     qo_offset_ptr,
+    request_seq_lens_ptr,
     num_rows: tl.constexpr,
     decode_query_len: tl.constexpr,
+    HAS_REQUEST_SEQ_LENS: tl.constexpr,
     BLOCK_SIZE: tl.constexpr,
 ):
     offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
@@ -81,6 +116,10 @@ def _update_runtime_metadata_kernel(
         seq_len - decode_query_len + local_query,
         mask=mask,
     )
+    if HAS_REQUEST_SEQ_LENS:
+        tl.store(
+            request_seq_lens_ptr + request, seq_len, mask=mask & (local_query == 0)
+        )
 
 
 @dataclass
@@ -101,11 +140,14 @@ def _build_plan(
         num_kv_heads: int,
         page_size: int,
         topk_blocks: int,
+        q8kv4: bool,
     ) -> Any:
         from vllm.third_party.fmha_sm100.api import fmha_sm100_plan
 
         qo_lens_cpu = torch.full((batch,), decode_query_len, dtype=torch.int32)
-        kv_lens_cpu = initial_seq_lens_cpu
+        # Zero-length full-CUDA-graph padding requests would keep the Q8KV4
+        # kernel from planning the batch; lengths are refreshed every step.
+        kv_lens_cpu = initial_seq_lens_cpu.clamp(min=1)
         plan = fmha_sm100_plan(
             qo_lens_cpu,
             kv_lens_cpu,
@@ -120,12 +162,20 @@ def _build_plan(
             use_fp8_kvcache=True,
             split_prefill_decode=False,
             device=device,
-            # The Q8KV4 route keeps plan-time lengths and its own page offsets,
-            # which this cached plan does not refresh.
-            decode_backend="kv_mode3",
+            decode_backend="q8kv4" if q8kv4 else "kv_mode3",
+            block_scale_shift=NVFP4_BLOCK_SCALE_SHIFT,
         )
 
         plan_info = plan[3]
+        # The Q8KV4 plan addresses pages per request, through the same flat
+        # block table as the per-token rows below.
+        q8kv4_plan = plan_info.get("q8kv4")
+        if q8kv4_plan is not None:
+            q8kv4_plan["kv_indptr"].copy_(
+                torch.arange(batch + 1, dtype=torch.int32, device=device).mul_(
+                    page_table_stride
+                )
+            )
         row_starts = (
             torch.arange(batch, dtype=torch.int32, device=device)
             .mul_(page_table_stride)
@@ -155,6 +205,7 @@ def prepare(
         num_kv_heads: int,
         page_size: int,
         topk_blocks: int,
+        q8kv4: bool = False,
     ) -> MSACutlassDecodeMetadata:
         batch = int(seq_lens.shape[0])
         if (
@@ -196,6 +247,7 @@ def prepare(
             num_kv_heads,
             page_size,
             topk_blocks,
+            q8kv4,
         )
         plan = self.plans.get(key)
         if plan is None:
@@ -209,17 +261,21 @@ def prepare(
                 num_kv_heads=num_kv_heads,
                 page_size=page_size,
                 topk_blocks=topk_blocks,
+                q8kv4=q8kv4,
             )
             self.plans[key] = plan
 
         plan_info = plan[3]
+        q8kv4_plan = plan_info.get("q8kv4")
         num_rows = batch * decode_query_len
         _update_runtime_metadata_kernel[(triton.cdiv(num_rows, 128),)](
             seq_lens,
             plan_info["kv_segment_lens"],
             plan_info["qo_offset"],
+            q8kv4_plan["seq_lens"] if q8kv4_plan is not None else seq_lens,
             num_rows=num_rows,
             decode_query_len=decode_query_len,
+            HAS_REQUEST_SEQ_LENS=q8kv4_plan is not None,
             BLOCK_SIZE=128,
         )
         return MSACutlassDecodeMetadata(
@@ -288,6 +344,7 @@ def should_prepare_decode_metadata(
 ) -> bool:
     """Return whether a graph shape can use the CUTLASS decode path."""
     total_q = batch_size * decode_query_len
+    nvfp4 = is_nvfp4_kv_cache(kv_cache_dtype)
     if min_batch_size is None:
         min_batch_size = _min_cutlass_batch_size(
             decode_query_len, num_q_heads, num_kv_heads
@@ -302,8 +359,12 @@ def should_prepare_decode_metadata(
             topk_blocks=topk_blocks,
         )
         and 1 <= decode_query_len <= _MAX_DECODE_QUERY_LEN
-        and (batch_size >= min_batch_size or is_nvfp4_kv_cache(kv_cache_dtype))
-        and total_q * num_q_heads <= _MAX_QUERY_HEAD_ROWS
+        and (batch_size >= min_batch_size or nvfp4)
+        and (
+            total_q <= _MAX_Q8KV4_QUERY_TOKENS
+            if nvfp4
+            else total_q * num_q_heads <= _MAX_QUERY_HEAD_ROWS
+        )
     )
 
 
@@ -318,9 +379,14 @@ def prepare_decode_metadata(
     num_kv_heads: int,
     page_size: int,
     topk_blocks: int,
+    q8kv4: bool = False,
     plan_cache: MSACutlassDecodePlanCache | None = None,
 ) -> MSACutlassDecodeMetadata:
-    """Prepare graph-stable runtime metadata for one sparse decode step."""
+    """Prepare graph-stable runtime metadata for one sparse decode step.
+
+    ``q8kv4`` plans the Q8KV4 kernel, which serves every NVFP4 cache (see
+    ``uses_q8kv4_decode``).
+    """
     cache = plan_cache or MSACutlassDecodePlanCache()
     return cache.prepare(
         block_table,
@@ -331,6 +397,7 @@ def prepare_decode_metadata(
         num_kv_heads=num_kv_heads,
         page_size=page_size,
         topk_blocks=topk_blocks,
+        q8kv4=q8kv4,
     )
 
 
diff --git a/vllm/models/minimax_m3/nvidia/mtp.py b/vllm/models/minimax_m3/nvidia/mtp.py
index 9c10a52b7cae..039c08bc6504 100644
--- a/vllm/models/minimax_m3/nvidia/mtp.py
+++ b/vllm/models/minimax_m3/nvidia/mtp.py
@@ -144,10 +144,17 @@ def forward(
 
 
 class MiniMaxM3MTP(nn.Module):
+    # Shares the target's input embedding unless loading finds its own (only
+    # tracked under V2 PP, where the target's may not be available).
+    has_own_embed_tokens = False
+
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()
 
-        if vllm_config.use_v2_model_runner and get_pp_group().world_size > 1:
+        self._tracks_own_embed = (
+            vllm_config.use_v2_model_runner and get_pp_group().world_size > 1
+        )
+        if self._tracks_own_embed:
             # Preserve embeddings from loaders that bypass load_weights.
             self.has_own_embed_tokens = vllm_config.load_config.load_format != "dummy"
 
@@ -320,7 +327,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
                     f"Failed to load MTP layer {layer_idx} weights from checkpoint."
                 )
 
-        if hasattr(self, "has_own_embed_tokens"):
+        if self._tracks_own_embed:
             if "model.embed_tokens.weight" in loaded_params:
                 self.has_own_embed_tokens = True
             elif is_mtp_completeness_check_enabled():
diff --git a/vllm/models/minimax_m3/nvidia/ops/dense_gemm1_mxfp8.py b/vllm/models/minimax_m3/nvidia/ops/dense_gemm1_mxfp8.py
new file mode 100644
index 000000000000..f18819269172
--- /dev/null
+++ b/vllm/models/minimax_m3/nvidia/ops/dense_gemm1_mxfp8.py
@@ -0,0 +1,664 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""MXFP8 dense GEMM1 (gate_up) with fused clamped-SwiGLU + MXFP8 quant epilogue.
+
+One persistent SM100 tcgen05 MXFP8 GEMM tuned for prefill shapes (BM=128,
+BN=128 per gate/up half, BK=128, cta_group=2; 6-warp CTA: warps 0-3
+epilogue, warp 4 MMA, warp 5 TMA) computes ``gate = X @ Wg^T`` and
+``up = X @ Wu^T`` and applies clamped SwiGLU and MXFP8 quantization. No
+BF16 intermediate.
+
+The tmem accumulator is double-buffered with a 32-col overlap (even tiles
+[0,256), odd [224,480), SF fixed [500,512)): the epilogue reads the shared
+32-col group first, so the next tile's MMA waits on a 32-col read instead
+of 256. The single epilogue warpgroup processes a tile in 32-col chunks
+(load -> act+quant -> store), overlap chunk first.
+
+Tile order: flat (n-fastest) when the B operand fits in half of L2, else
+grouped m-fastest (GROUP_M=32) so B panels stay L2-resident (dense TP1:
+B=151MB > 126MB L2).
+"""
+
+from functools import cache
+
+import cutlass
+import torch
+from cuda.bindings.driver import CUstream
+from cutlass import (
+    Float8E4M3FN,
+    Float8E8M0FNU,
+    Float32,
+    Int32,
+    Int64,
+    Uint8,
+    Uint16,
+    Uint32,
+    Uint64,
+    cute,
+)
+from cutlass.cute.nvgpu import cpasync, tcgen05
+from cutlass.cute.runtime import make_fake_stream, make_fake_tensor
+
+from vllm.cute_utils import _tcgen05, mbarrier, simple_tma_copy, to_cta0_smem
+from vllm.cute_utils.cvt import fp32x4_to_fp8x4
+from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation
+from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
+    mxfp8_e4m3_quantize,
+)
+from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp8Dynamic
+from vllm.platforms import current_platform
+
+
+class Sm100DenseGemm1Mxfp8:
+    def __init__(self, GROUP_M: int, use_pdl: bool = True) -> None:
+        BM, BN, BK = 128, 128, 128
+        self.cta_tile = (BM, BN, BK)
+        self.GROUP_M = GROUP_M
+        self.use_pdl = use_pdl
+
+        smem_bytes = cutlass.memory.SmemAllocator.capacity_in_bytes()
+        # A + Bg/2 + Bu/2 tiles + SFA + SFBg + SFBu (2-CTA: B halves, SF full)
+        self.stage_size = (BM + BN) * BK + (BM + 256) * (BK // 32)
+        self.num_stages = smem_bytes // self.stage_size
+
+    @cute.jit
+    def prepare_AB(self, A: cute.Tensor, BM: int, BK: int) -> cpasync.TmaInfo:
+        tma_op = cpasync.CopyBulkTensorTileG2SOp(cta_group=tcgen05.CtaGroup.TWO)
+        swizzle_128B = cute.make_swizzle(3, 4, 3)
+        s_layout = cute.make_layout((BM, BK, self.num_stages), stride=(BK, 1, BM * BK))
+        s_layout = cute.make_composed_layout(swizzle_128B, 0, s_layout)
+        return cpasync.make_tiled_tma_atom(tma_op, A, s_layout, (BM, BK))
+
+    @cute.jit
+    def prepare_SF(
+        self,
+        SF: cute.Tensor,
+        R: Int32,
+        Ksf: Int32,
+        BR: int,
+        BKsf: int,
+        tma_op: cpasync.CopyG2STileBaseOp,
+    ) -> cpasync.TmaInfo:
+        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)
+        s_layout = cute.make_layout((*tiler, self.num_stages))
+        return cpasync.make_tiled_tma_atom(tma_op, SF_i64, s_layout, tiler)
+
+    @cute.jit
+    def __call__(
+        self,
+        A: cute.Tensor,
+        B: cute.Tensor,
+        SFA: cute.Tensor,
+        SFB: cute.Tensor,
+        Q: cute.Tensor,
+        SFQ: cute.Tensor,
+        limit: Float32,
+        alpha: Float32,
+        beta: Float32,
+        grid_size: Int32,
+        stream: CUstream,
+    ) -> None:
+        M, K = A.shape
+        N2, _ = B.shape  # 2 * N (gate|up)
+        BM, BN, BK = self.cta_tile
+
+        A_tma = self.prepare_AB(A, BM, BK)
+        B_tma = self.prepare_AB(B, BN, BK)
+        SFA_tma = self.prepare_SF(
+            SFA,
+            M,
+            K // 32,
+            BM,
+            BK // 32,
+            cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.TWO),
+        )
+        SFB_tma = self.prepare_SF(
+            SFB,
+            N2,
+            K // 32,
+            BN,
+            BK // 32,
+            cpasync.CopyBulkTensorTileG2SMulticastOp(tcgen05.CtaGroup.TWO),
+        )
+
+        self.kernel(A_tma, B_tma, SFA_tma, SFB_tma, Q, SFQ, limit, alpha, beta).launch(
+            grid=(grid_size, 1, 1),
+            block=(6 * 32, 1, 1),
+            cluster=(2, 1, 1),
+            stream=stream,
+            use_pdl=self.use_pdl,
+        )
+
+    @cute.kernel
+    def kernel(
+        self,
+        A_tma: cpasync.TmaInfo,
+        B_tma: cpasync.TmaInfo,
+        SFA_tma: cpasync.TmaInfo,
+        SFB_tma: cpasync.TmaInfo,
+        Q_tensor: cute.Tensor,
+        SFQ_tensor: cute.Tensor,
+        limit: Float32,
+        alpha: Float32,
+        beta: Float32,
+    ) -> None:
+        tid, _, _ = cute.arch.thread_idx()
+        raw_bid, _, _ = cute.arch.block_idx()
+        num_bids, _, _ = cute.arch.grid_dim()
+        warp_id = cute.arch.make_warp_uniform(tid // 32)
+        cta_rank = raw_bid % 2
+
+        BM, BN, BK = self.cta_tile
+        GROUP_M = self.GROUP_M
+        num_stages = self.num_stages
+
+        smem = cutlass.memory.SmemAllocator()
+        sA = smem.allocate_tensor(
+            Float8E4M3FN,
+            A_tma.smem_layout.outer,
+            byte_alignment=128,
+            swizzle=A_tma.smem_layout.inner,
+        )
+        sB = smem.allocate_tensor(
+            Float8E4M3FN,
+            B_tma.smem_layout.outer,
+            byte_alignment=128,
+            swizzle=B_tma.smem_layout.inner,
+        )
+        sSFA = smem.allocate_tensor(Int64, SFA_tma.smem_layout, byte_alignment=128)
+        # SFB slots: rank r's multicast copy fills slot r of BOTH CTAs;
+        # slice shape/stride per slot matches the atom's (64,1):(1,64) tile
+        sSFB = smem.allocate_tensor(
+            Int64,
+            cute.make_layout((64, 1, 2, num_stages), stride=(1, 64, 64, 128)),
+            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, 2)
+        tmem_partial_mbar = smem.allocate_array(Int64, 1)
+        tmem_empty_mbar = smem.allocate_array(Int64, 2)
+        taddr = smem.allocate(Int32, 4)
+
+        # allocate named barrier
+        BAR_EPI = 1
+        BAR_TMEM_ALLOC = 2
+
+        M, K = A_tma.tma_tensor.shape
+        N2, _ = B_tma.tma_tensor.shape
+        N = N2 // 2
+        grid_m = cute.ceil_div(M, BM * 2) * 2
+        grid_n = N // BN
+        total_tiles = grid_m * grid_n
+        # output SF swizzle: NUM_K_TILES = ceil(N/32/4) tiles of 512B per 128 rows
+        NUM_KT = cute.ceil_div(N // 32, 4)
+
+        def compute_bid(bid, GROUP_M, grid_m, grid_n):
+            if cutlass.const_expr(GROUP_M == 0):
+                bid_m = (bid // (grid_n * 2)) * 2 + bid % 2
+                bid_n = (bid // 2) % grid_n
+            else:
+                group_size = GROUP_M * grid_n
+                group_id = bid // group_size
+                group_m = min(grid_m - group_id * GROUP_M, GROUP_M)
+                local_bid = bid - group_id * group_size
+                bid_m = group_id * GROUP_M + local_bid % group_m
+                bid_n = local_bid // group_m
+            return bid_m, bid_n
+
+        if warp_id == 0:
+            with cute.arch.elect_one():
+                for i in cutlass.range_constexpr(num_stages):
+                    cute.arch.mbarrier_init(tma_full_mbar + i, 2)
+                    cute.arch.mbarrier_init(tma_empty_mbar + i, 1)
+                for i in cutlass.range_constexpr(2):
+                    cute.arch.mbarrier_init(tmem_full_mbar + i, 1)
+                    cute.arch.mbarrier_init(tmem_empty_mbar + i, 256)
+                cute.arch.mbarrier_init(tmem_partial_mbar, 256)
+                cute.arch.mbarrier_init_fence()
+        elif warp_id == 1:
+            cpasync.prefetch_descriptor(A_tma.atom)
+            cpasync.prefetch_descriptor(B_tma.atom)
+            cpasync.prefetch_descriptor(SFA_tma.atom)
+            cpasync.prefetch_descriptor(SFB_tma.atom)
+
+        cute.arch.cluster_arrive_relaxed()
+        cute.arch.cluster_wait()
+
+        if cutlass.const_expr(self.use_pdl):
+            cute.arch.griddepcontrol_wait()
+
+        if warp_id == 5:
+            # TMA warp
+            tma_stage = 0
+            parity = 1
+            tma_full_mbar_ = to_cta0_smem(tma_full_mbar)
+
+            gA_tiles = cute.zipped_divide(A_tma.tma_tensor, (BM, BK))
+            gB_tiles = cute.zipped_divide(B_tma.tma_tensor, (BN, BK))
+            gSFA_tiles = cute.zipped_divide(
+                SFA_tma.tma_tensor, (128 * (BK // 32) // 8, BM // 128)
+            )
+            gSFB_tiles = cute.zipped_divide(
+                SFB_tma.tma_tensor, (128 * (BK // 32) // 8, BN // 128)
+            )
+
+            for work_id in range(raw_bid, total_tiles, num_bids):
+                bid_m, bid_n = compute_bid(work_id, GROUP_M, grid_m, grid_n)
+
+                for iter_k in range(cute.ceil_div(K, BK)):
+                    mbar = tma_full_mbar_ + tma_stage
+                    cute.arch.mbarrier_wait(tma_empty_mbar + tma_stage, parity)
+                    with cute.arch.elect_one():
+                        mbarrier.arrive_expect_tx(mbar, self.stage_size, "cluster")
+                    simple_tma_copy(
+                        A_tma.atom,
+                        gA_tiles[None, (bid_m, iter_k)],
+                        sA[None, None, tma_stage],
+                        mbar,
+                    )
+                    simple_tma_copy(
+                        B_tma.atom,
+                        gB_tiles[None, (cta_rank * grid_n + bid_n, iter_k)],
+                        sB[None, None, tma_stage],
+                        mbar,
+                    )
+                    simple_tma_copy(
+                        SFA_tma.atom,
+                        gSFA_tiles[None, (iter_k, bid_m)],
+                        sSFA[None, None, tma_stage],
+                        mbar,
+                    )
+                    # TMA multicast. rank0 loads gate, rank1 loads up
+                    simple_tma_copy(
+                        SFB_tma.atom,
+                        gSFB_tiles[None, (iter_k, bid_n + cta_rank * grid_n)],
+                        sSFB[None, None, cta_rank, tma_stage],
+                        mbar,
+                        mcast_mask=Uint16(3),
+                    )
+
+                    tma_stage = (tma_stage + 1) % num_stages
+                    if tma_stage == 0:
+                        parity ^= 1
+
+            if cutlass.const_expr(self.use_pdl):
+                cute.arch.griddepcontrol_launch_dependents()
+
+        elif warp_id == 4:
+            # MMA warp (2-CTA collective alloc; only cta_rank 0 issues)
+            cute.arch.barrier(barrier_id=BAR_TMEM_ALLOC, number_of_threads=2 * 32)
+
+            tma_stage = 0
+            tma_full_parity = 0
+
+            sdesc = _tcgen05.make_sdesc_128B_swizzle(0)
+            # SF smem descriptor: no swizzle, SBO = 8 * 16 bytes.
+            sdesc_sf = Uint64((8 << 32) | (1 << 46))
+            idesc = _tcgen05.make_mxfp8_idesc(BM * 2, BN * 2)
+            multicast_mask = Uint16(3)
+
+            # each [128,128] MXFP8 tile needs 4 tmem columns for SF
+            sfa_tmem = 500
+            sfbg_tmem = sfa_tmem + 4
+            sfbu_tmem = sfbg_tmem + 4
+
+            if cta_rank == 0:
+                tmem_stage = 0
+                tmem_parity = 1
+                for _ in range(raw_bid, total_tiles, num_bids):
+                    acc_base = 224 * tmem_stage
+                    cute.arch.mbarrier_wait(tmem_empty_mbar + tmem_stage, tmem_parity)
+                    # wait for the previous tile's overlap-region read; the
+                    # single partial barrier phases once per tile
+                    cute.arch.mbarrier_wait(tmem_partial_mbar, tmem_stage ^ 1)
+                    _tcgen05.fence_after_thread_sync()
+
+                    for iter_k in cutlass.range(cute.ceil_div(K, BK), unroll=1):
+                        a_desc = sdesc | (
+                            sA[None, None, tma_stage].iterator.toint() >> 4
+                        )
+                        b_desc = sdesc | (
+                            sB[None, None, tma_stage].iterator.toint() >> 4
+                        )
+                        sfa_desc = sdesc_sf | (
+                            sSFA[None, None, tma_stage].iterator.toint() >> 4
+                        )
+                        sfbg_desc = sdesc_sf | (
+                            sSFB[None, None, 0, tma_stage].iterator.toint() >> 4
+                        )
+                        sfbu_desc = sdesc_sf | (
+                            sSFB[None, None, 1, tma_stage].iterator.toint() >> 4
+                        )
+
+                        cute.arch.mbarrier_wait(
+                            tma_full_mbar + tma_stage, tma_full_parity
+                        )
+                        _tcgen05.fence_after_thread_sync()
+
+                        _tcgen05.cp(sfa_tmem, sfa_desc, "32x128b", "warpx4", 2)
+                        _tcgen05.cp(sfbg_tmem, sfbg_desc, "32x128b", "warpx4", 2)
+                        _tcgen05.cp(sfbu_tmem, sfbu_desc, "32x128b", "warpx4", 2)
+
+                        for k in cutlass.range_constexpr(BK // 32):
+                            enable_input_d = iter_k > 0 or k > 0
+                            # select this k-block's SF byte inside the staged
+                            # 4-col tmem atoms (idesc SF-id fields)
+                            idesc_k = idesc + ((k << 4) | (k << 29))
+                            _tcgen05.mma_mxfp8(
+                                acc_base,
+                                a_desc,
+                                b_desc,
+                                idesc_k,
+                                sfa_tmem,
+                                sfbg_tmem,
+                                enable_input_d,
+                                2,
+                            )
+                            a_desc += 32 >> 4
+                            b_desc += 32 >> 4
+                        _tcgen05.commit(tma_empty_mbar + tma_stage, multicast_mask, 2)
+
+                        tma_stage = (tma_stage + 1) % num_stages
+                        if tma_stage == 0:
+                            tma_full_parity ^= 1
+
+                    _tcgen05.commit(tmem_full_mbar + tmem_stage, multicast_mask, 2)
+                    tmem_stage ^= 1
+                    if tmem_stage == 0:
+                        tmem_parity ^= 1
+
+        else:
+            # epilogue warps (one warpgroup; tid < 128 indexes the tmem lane)
+            if warp_id == 0:
+                _tcgen05.alloc(taddr, cta_group=2)
+                cute.arch.barrier(barrier_id=BAR_TMEM_ALLOC, number_of_threads=2 * 32)
+
+            exp_scale = -1.4426950216293334961 * alpha
+
+            # ((32,4), (M, N/BN))
+            Q_ = cute.zipped_divide(Q_tensor, (1, cute.make_layout((32, 4))))
+            Q_ = Q_[(0, None), None]
+
+            f8x32_atom = cute.make_copy_atom(
+                cute.nvgpu.CopyUniversalOp(),
+                Float8E4M3FN,
+                num_bits_per_copy=256,
+                l1c_evict_priority=cute.nvgpu.CacheEvictionPriority.NO_ALLOCATE,
+            )
+
+            partial_ = to_cta0_smem(tmem_partial_mbar)
+            empty_ = to_cta0_smem(tmem_empty_mbar)
+
+            tmem_stage = 0
+            tmem_parity = 0
+            for work_id in range(raw_bid, total_tiles, num_bids):
+                bid_m, bid_n = compute_bid(work_id, GROUP_M, grid_m, grid_n)
+                m = bid_m * BM + tid
+                acc_base = 224 * tmem_stage
+
+                if warp_id == 0:
+                    cute.arch.mbarrier_wait(tmem_full_mbar + tmem_stage, tmem_parity)
+                cute.arch.barrier(barrier_id=BAR_EPI, number_of_threads=4 * 32)
+                _tcgen05.fence_after_thread_sync()
+
+                rGate = cute.make_rmem_tensor(32, Float32)
+                rUp = cute.make_rmem_tensor(32, Float32)
+
+                # cc=0 is the overlap chunk: cols [224,256) (up[96:128] on
+                # even tiles, gate[0:32] on odd) plus the matching 32 cols of
+                # the other operand. cc=1..3: the remaining chunk-pairs in
+                # order. Each chunk: load -> act+quant -> store.
+                for cc in cutlass.range_constexpr(4):
+                    if cutlass.const_expr(cc == 0):
+                        c = 3 - 3 * tmem_stage  # even tiles: 3; odd: 0
+                        gate_col = 96 + 128 * tmem_stage
+                        up_col = 224 + 128 * tmem_stage
+                    else:
+                        c = cc - 1 + tmem_stage  # even: 0,1,2; odd: 1,2,3
+                        gate_col = acc_base + c * 32
+                        up_col = acc_base + 128 + c * 32
+                    rGate.store(_tcgen05.ld(tid // 32 * 32, gate_col, "32x32b", 32))
+                    rUp.store(_tcgen05.ld(tid // 32 * 32, up_col, "32x32b", 32))
+                    _tcgen05.wait_ld()
+                    if cutlass.const_expr(cc == 0):
+                        _tcgen05.fence_before_thread_sync()
+                        mbarrier.arrive(partial_, "cluster")
+                    if cutlass.const_expr(cc == 3):
+                        _tcgen05.fence_before_thread_sync()
+                        mbarrier.arrive(empty_ + tmem_stage, "cluster")
+
+                    q_rmem = cute.make_rmem_tensor(32, Float8E4M3FN)
+                    q_i32 = cute.recast_tensor(q_rmem, Uint32)
+
+                    # activation
+                    for j in cutlass.range_constexpr(32):
+                        rGate[j] = cute.arch.fmin(rGate[j], limit)
+                        rUp[j] = cute.arch.fmin(cute.arch.fmax(rUp[j], -limit), limit)
+
+                    af = cute.make_rmem_tensor(32, Float32)
+                    for j in cutlass.range(32, vectorize=True):
+                        e = cute.math.exp2(rGate[j] * exp_scale, fastmath=True)
+                        gate = rGate[j] * cute.arch.rcp_approx(1.0 + e)
+                        af[j] = gate * (rUp[j] + beta)
+
+                    # quantization
+                    amax0 = cute.math.abs(af[0])
+                    amax1 = cute.math.abs(af[1])
+                    for j in cutlass.range_constexpr(15):
+                        amax0 = cute.arch.fmax(amax0, cute.math.abs(af[j * 2 + 2]))
+                        amax1 = cute.arch.fmax(amax1, cute.math.abs(af[j * 2 + 3]))
+                    amax = cute.arch.fmax(amax0, amax1)
+
+                    # Scale byte via exact integer math:
+                    # ceil(log2(amax/448)) = E - 8 + (amax > 1.75 * 2^E),
+                    # E = exponent(amax); amax=0 (all-zero block) clamps to 0.
+                    bits = amax.bitcast(Int32)
+                    exp_bits = bits & 0x7F800000
+                    thr = (exp_bits | 0x600000).bitcast(Float32)  # 1.75 * 2^E
+                    sb = (exp_bits >> 23) - 8 + Int32(amax > thr)
+                    sb = cutlass.min(cutlass.max(sb, Int32(0)), Int32(254))
+
+                    inv = ((Int32(254) - sb) << 23).bitcast(Float32)
+                    for j in cutlass.range(32, vectorize=True):
+                        af[j] *= inv
+
+                    for j in cutlass.range_constexpr(8):
+                        q_i32[j] = fp32x4_to_fp8x4(*[af[4 * j + t] for t in range(4)])
+
+                    kb = bid_n * (BN // 32) + c
+                    sf_off = (
+                        (bid_m * NUM_KT + kb // 4) * 512
+                        + (m % 32) * 16
+                        + ((m // 32) % 4) * 4
+                        + kb % 4
+                    )
+                    if m < M:
+                        cute.copy(f8x32_atom, q_rmem, Q_[(None, c), (m, bid_n)])
+                        SFQ_tensor[sf_off] = Uint8(sb)
+
+                tmem_stage ^= 1
+                if tmem_stage == 0:
+                    tmem_parity ^= 1
+
+            cute.arch.cluster_arrive_relaxed()
+            cute.arch.cluster_wait()
+            if warp_id == 0:
+                _tcgen05.dealloc(cta_group=2)
+
+    @cache
+    @staticmethod
+    def compile(GROUP_M: int, use_pdl: bool = True):
+        M = cute.sym_int()
+        K = cute.sym_int()
+        N2 = cute.sym_int()
+
+        A = make_fake_tensor(
+            Float8E4M3FN,
+            (M, K),
+            (cute.sym_int64(divisibility=16), 1),
+            assumed_align=16,
+        )
+        B = make_fake_tensor(
+            Float8E4M3FN,
+            (N2, K),
+            (cute.sym_int64(divisibility=16), 1),
+            assumed_align=16,
+        )
+        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
+        )
+        Q = make_fake_tensor(
+            Float8E4M3FN,
+            (M, cute.sym_int(divisibility=256)),
+            (cute.sym_int64(divisibility=256), 1),
+            assumed_align=32,
+        )
+        SFQ = make_fake_tensor(
+            Uint8, (cute.sym_int(divisibility=512),), (1,), assumed_align=16
+        )
+
+        stream = make_fake_stream(use_tvm_ffi_env_stream=True)
+        kernel = Sm100DenseGemm1Mxfp8(GROUP_M, use_pdl)
+        return cute.compile(
+            kernel,
+            A,
+            B,
+            SFA,
+            SFB,
+            Q,
+            SFQ,
+            Float32(7.0),
+            Float32(1.702),
+            Float32(1.0),
+            148,
+            stream,
+            options="--enable-tvm-ffi",
+        )
+
+
+# Below this the unfused GEMM -> act -> quant chain is as fast (TP4 shared
+# expert, I=768: break-even at ~1024 tokens).
+MIN_TOKENS = 1024
+
+
+def _linear_kernel(linear):
+    method = getattr(linear, "scheme", getattr(linear, "quant_method", None))
+    return getattr(method, "kernel", None)
+
+
+def can_fuse_gemm1(gate_up_proj, down_proj, act_fn) -> bool:
+    """Whether fused_gemm1 can replace gate_up_proj + act_fn for down_proj.
+
+    Called at init, before weight loading: kernels are selected at
+    create_weights time, so the decision does not depend on the weights.
+    """
+    from vllm.model_executor.kernels.linear.mxfp8.flashinfer import (
+        FlashInferCutedslMxfp8LinearKernel,
+        FlashInferCutlassMxfp8LinearKernel,
+    )
+    from vllm.model_executor.layers.activation import SiluAndMulWithClamp
+
+    if not (
+        current_platform.is_cuda() and current_platform.is_device_capability_family(100)
+    ):
+        return False
+    if not isinstance(act_fn, SiluAndMulWithClamp):
+        return False
+    # The kernel reads gate_up's fp8 weight and F8_128x4-swizzled weight
+    # scale directly, so it depends on these kernels' weight layout.
+    if not isinstance(
+        _linear_kernel(gate_up_proj),
+        (FlashInferCutedslMxfp8LinearKernel, FlashInferCutlassMxfp8LinearKernel),
+    ):
+        return False
+    down_kernel = _linear_kernel(down_proj)
+    if down_kernel is None or down_kernel.input_quant_key() != kMxfp8Dynamic:
+        return False
+    # I % 256 == 0 (output row stride), K % 128 == 0 (BK).
+    n = gate_up_proj.output_size_per_partition // 2
+    k = gate_up_proj.input_size_per_partition
+    return n % 256 == 0 and k % 128 == 0
+
+
+def fused_gemm1(x: torch.Tensor, gate_up_proj, act_fn) -> QuantizedActivation:
+    """act_fn(gate_up_proj(x)), MXFP8-quantized for an MXFP8 down_proj.
+
+    Requires can_fuse_gemm1(gate_up_proj, down_proj, act_fn).
+    """
+    from vllm.model_executor.kernels.linear.mxfp8.flashinfer import (
+        FlashInferCutedslMxfp8LinearKernel,
+    )
+
+    assert x.ndim == 2 and x.dtype == torch.bfloat16
+    weight = gate_up_proj.weight
+    # FlashInferCutedsl stores the weight as a column-major [K, 2I] view.
+    if isinstance(_linear_kernel(gate_up_proj), FlashInferCutedslMxfp8LinearKernel):
+        weight = weight.t()
+
+    M = x.shape[0]
+    N = weight.shape[0] // 2
+    q = torch.empty((M, N), dtype=torch.float8_e4m3fn, device=x.device)
+    sf = torch.empty(
+        ((M + 127) // 128) * 128 * (N // 32), dtype=torch.uint8, device=x.device
+    )
+
+    xq, xsf = mxfp8_e4m3_quantize(x, is_sf_swizzled_layout=True)
+    run(
+        xq,
+        weight,
+        xsf.view(-1).view(torch.float8_e8m0fnu),
+        gate_up_proj.weight_scale.view(torch.float8_e8m0fnu),
+        q,
+        sf,
+        act_fn.swiglu_limit,
+        act_fn.alpha,
+        act_fn.beta,
+    )
+    return QuantizedActivation(
+        data=q,
+        scale=sf,
+        orig_dtype=x.dtype,
+        orig_shape=x.shape[:-1] + (N,),
+        quant_key=kMxfp8Dynamic,
+    )
+
+
+def run(
+    xq: torch.Tensor,
+    weight: torch.Tensor,
+    x_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    q: torch.Tensor,
+    sf: torch.Tensor,
+    limit: float,
+    alpha: float,
+    beta: float,
+    grid: int | None = None,
+) -> None:
+    """Functional interface to the fused kernel (hides compile + grid sizing).
+
+    xq: [M, K] fp8e4m3; weight: [2N, K] fp8e4m3 row-major; x_scale/weight_scale:
+    F8_128x4-swizzled E8M0 (flat, any uint8-viewable dtype); q: [M, N] fp8e4m3;
+    sf: flat uint8 scale buffer (round_up(M,128) * N/32 bytes; rows >= M are
+    not written).
+    """
+    M, K = xq.shape
+    _, N = q.shape
+    if grid is None:
+        grid_m2 = ((M + 127) // 128 + 1) // 2 * 2
+        num_sms = torch.cuda.get_device_properties(xq.device).multi_processor_count
+        grid = min(num_sms, grid_m2 * (N // 128))
+
+    # Grouped (m-fastest) rasterization only pays when B (n2*k fp8 bytes)
+    # exceeds half of L2; below that the flat order is faster (no div/mod
+    # overhead in compute_bid)
+    GROUP_M = 32 if N * 2 * K > (63 << 20) else 0
+    kernel = Sm100DenseGemm1Mxfp8.compile(GROUP_M)
+    kernel(xq, weight, x_scale, weight_scale, q, sf, limit, alpha, beta, grid)
diff --git a/vllm/models/minimax_m3/nvidia/ops/finalize_allreduce_norm.py b/vllm/models/minimax_m3/nvidia/ops/finalize_allreduce_norm.py
new file mode 100644
index 000000000000..9c75e71e3249
--- /dev/null
+++ b/vllm/models/minimax_m3/nvidia/ops/finalize_allreduce_norm.py
@@ -0,0 +1,781 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# The consumer is adapted from FlashInfer's low-latency MNNVL CuTe DSL
+# all-reduce (flashinfer/comm/mnnvl_cutedsl/kernel_ll,
+# _LamportResidualRMSNormDeviceKernel), as DSV4.1's AllReduceMHC is.
+
+"""MoE top-k finalize + TP all-reduce + Gemma RMSNorm + MXFP8 quantization.
+
+FlashInfer's Lamport all-reduce: the finalize publish kernel sums each token's
+routed expert rows and its shared-expert row, then multicasts that TP-partial
+row into a three-generation mailbox; a PDL-chained consumer spins until every
+rank's fragment lands, reduces them, adds the residual and applies the Gemma
+RMSNorm. Instead of a BF16 row, the consumer writes the normalized row as MXFP8
+(E4M3 values and UE8M0 scales of 32-element blocks in the 128x4 swizzled
+layout), so the MXFP8 linear that follows skips its own input quantization.
+"""
+
+import math
+
+import cuda.bindings.driver as cuda
+import cutlass
+import cutlass.cute as cute
+import torch
+import torch.distributed as dist
+import torch.distributed._symmetric_memory as symm_mem
+from cutlass import BFloat16, Float32, Int32, Int64, Uint8, Uint32
+from cutlass._mlir.dialects import llvm
+from cutlass.cute.runtime import make_fake_compact_tensor
+from cutlass.cutlass_dsl import T, dsl_user_op
+
+from vllm.config import VllmConfig
+from vllm.distributed import get_tensor_model_parallel_world_size, get_tp_group
+from vllm.model_executor.layers.fused_moe.moe_output import (
+    MoEOutput,
+    UnfinalizedMoEOutput,
+)
+from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation
+from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp8Dynamic
+from vllm.models.common.ops.moe_tail_fusion.primitives import (
+    ACTIVE_STAGE,
+    LAMPORT_GENERATIONS,
+    NEXT_STAGE,
+    VEC_BF16,
+    WARP_SIZE,
+    bf16x8_to_packed_u32x4,
+    current_cu_stream,
+    fragment_has_negative_zero,
+    group_leader_block_sum,
+    load_global_u32x4,
+    load_volatile_u32,
+    make_fake_dynamic_compact_tensor,
+    map_shared_to_peer,
+    packed_u32x4_to_bf16x8,
+    store_global_u32,
+    store_global_u32x4,
+    store_lamport_sentinel_u32x4,
+    store_shared_cluster_f32,
+    to_cute,
+    to_cute_dynamic,
+)
+from vllm.models.deepseek_v41.nvidia.ops.cute_dsl.all_reduce_mhc import (
+    _QuadFinalizePublishDeviceKernel,
+)
+from vllm.platforms import current_platform
+
+MXFP8_BLOCK = 32
+# Fragments (of VEC_BF16 values) per MXFP8 block, held by adjacent lanes.
+_FRAGMENTS_PER_BLOCK = MXFP8_BLOCK // VEC_BF16
+# The swizzled scale layout tiles 128 rows by 4 block columns.
+_SCALE_ROW_TILE = 128
+_SCALE_COL_TILE = 4
+
+
+@dsl_user_op
+def _ue8m0_ceil(value: Float32, *, loc=None, ip=None) -> Uint32:
+    """FlashInfer's float_to_ue8m0_fast: the biased exponent rounded up, 0 for
+    non-positive input, as its MXFP8 quantization kernels compute it."""
+    return Uint32(
+        llvm.inline_asm(
+            T.i32(),
+            [Float32(value).ir_value(loc=loc, ip=ip)],
+            "{\n\t"
+            ".reg .pred p_zero, p_has_mant, p_exp_zero, p_tiny_sub, p_ovf;\n\t"
+            ".reg .u32 bits, exp_biased, mantissa, bump, result;\n\t"
+            "setp.le.f32 p_zero, $1, 0f00000000;\n\t"
+            "mov.b32 bits, $1;\n\t"
+            "shr.b32 exp_biased, bits, 23;\n\t"
+            "and.b32 exp_biased, exp_biased, 255;\n\t"
+            "and.b32 mantissa, bits, 0x7FFFFF;\n\t"
+            "setp.ne.u32 p_has_mant, mantissa, 0;\n\t"
+            "selp.u32 bump, 1, 0, p_has_mant;\n\t"
+            "setp.eq.u32 p_exp_zero, exp_biased, 0;\n\t"
+            "setp.le.u32 p_tiny_sub, mantissa, 0x400000;\n\t"
+            "and.pred p_tiny_sub, p_exp_zero, p_tiny_sub;\n\t"
+            "@p_tiny_sub mov.u32 bump, 0;\n\t"
+            "add.u32 result, exp_biased, bump;\n\t"
+            "setp.gt.u32 p_ovf, result, 254;\n\t"
+            "selp.u32 result, 254, result, p_ovf;\n\t"
+            "selp.u32 $0, 0, result, p_zero;\n\t"
+            "}",
+            "=r,f",
+            has_side_effects=False,
+            is_align_stack=False,
+            asm_dialect=llvm.AsmDialect.AD_ATT,
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
+@dsl_user_op
+def _ue8m0_inverse(exponent: Uint32, *, loc=None, ip=None) -> Float32:
+    """2^(127 - exponent), or 0 for exponent 0."""
+    return Float32(
+        llvm.inline_asm(
+            T.f32(),
+            [Uint32(exponent).ir_value(loc=loc, ip=ip)],
+            "{\n\t"
+            ".reg .s32 new_exp;\n\t"
+            ".reg .pred p_zero;\n\t"
+            "setp.eq.u32 p_zero, $1, 0;\n\t"
+            "sub.s32 new_exp, 254, $1;\n\t"
+            "max.s32 new_exp, new_exp, 0;\n\t"
+            "shl.b32 $0, new_exp, 23;\n\t"
+            "@p_zero mov.b32 $0, 0;\n\t"
+            "}",
+            "=f,r",
+            has_side_effects=False,
+            is_align_stack=False,
+            asm_dialect=llvm.AsmDialect.AD_ATT,
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
+@dsl_user_op
+def _f32x4_to_e4m3x4(
+    a: Float32, b: Float32, c: Float32, d: Float32, *, loc=None, ip=None
+) -> Uint32:
+    """Four values to E4M3, ``a`` in the low byte, clamped to +-448 first so
+    that inf saturates like FlashInfer's MXFP8 kernels and torch."""
+    return Uint32(
+        llvm.inline_asm(
+            T.i32(),
+            [Float32(v).ir_value(loc=loc, ip=ip) for v in (a, b, c, d)],
+            "{\n\t"
+            ".reg .f32 f0, f1, f2, f3;\n\t"
+            ".reg .b16 lo, hi;\n\t"
+            "min.f32 f0, $1, 0f43E00000;\n\t"
+            "max.f32 f0, f0, 0fC3E00000;\n\t"
+            "min.f32 f1, $2, 0f43E00000;\n\t"
+            "max.f32 f1, f1, 0fC3E00000;\n\t"
+            "min.f32 f2, $3, 0f43E00000;\n\t"
+            "max.f32 f2, f2, 0fC3E00000;\n\t"
+            "min.f32 f3, $4, 0f43E00000;\n\t"
+            "max.f32 f3, f3, 0fC3E00000;\n\t"
+            "cvt.rn.satfinite.e4m3x2.f32 lo, f1, f0;\n\t"
+            "cvt.rn.satfinite.e4m3x2.f32 hi, f3, f2;\n\t"
+            "mov.b32 $0, {lo, hi};\n\t"
+            "}",
+            "=r,f,f,f,f",
+            has_side_effects=False,
+            is_align_stack=False,
+            asm_dialect=llvm.AsmDialect.AD_ATT,
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
+@dsl_user_op
+def _store_global_u32x2(
+    address: Int64, low: Uint32, high: Uint32, *, loc=None, ip=None
+) -> None:
+    llvm.inline_asm(
+        None,
+        [
+            address.ir_value(loc=loc, ip=ip),
+            low.ir_value(loc=loc, ip=ip),
+            high.ir_value(loc=loc, ip=ip),
+        ],
+        "st.global.v2.u32 [$0], {$1, $2};",
+        "l,r,r",
+        has_side_effects=True,
+        is_align_stack=False,
+        asm_dialect=llvm.AsmDialect.AD_ATT,
+        loc=loc,
+        ip=ip,
+    )
+
+
+@cute.jit
+def _swizzled_scale_offset(row: Int32, col: Int32, padded_cols: Int32) -> Int32:
+    """Byte offset of block ``col`` of ``row`` in the 128x4 swizzled layout."""
+    return (
+        col % Int32(4)
+        + (col // Int32(4)) * Int32(512)
+        + (row % Int32(32)) * Int32(16)
+        + ((row % Int32(128)) // Int32(32)) * Int32(4)
+        + (row // Int32(128)) * Int32(128) * padded_cols
+    )
+
+
+class _LamportGemmaRMSNormMxfp8Kernel:
+    """FlashInfer's ``_LamportResidualRMSNormDeviceKernel`` (residual add and
+    copy, one rank lane) with a Gemma RMSNorm whose BF16 result is quantized to
+    MXFP8 instead of stored.
+
+    A 32-element scale block is four consecutive fragments, which adjacent lanes
+    hold, so two shuffles reduce its amax. The scales of the padding rows below
+    the next multiple of 128 are zeroed, as FlashInfer's quantization does.
+    """
+
+    def __init__(
+        self,
+        *,
+        hidden: int,
+        tp: int,
+        capacity_m: int,
+        cluster_size: int,
+        threads: int,
+        enable_pdl: bool,
+    ) -> None:
+        if hidden <= 0 or hidden % (MXFP8_BLOCK * _SCALE_COL_TILE):
+            raise ValueError("hidden must be a positive multiple of 128")
+        if threads <= 0 or threads % WARP_SIZE:
+            raise ValueError("threads must be a positive multiple of the warp size")
+        self.hidden = hidden
+        self.tp = tp
+        self.capacity_m = capacity_m
+        self.cluster_size = cluster_size
+        self.threads = threads
+        self.enable_pdl = enable_pdl
+        self.fragments = hidden // VEC_BF16
+        self.fragment_stride = cluster_size * threads
+        self.trips = math.ceil(self.fragments / self.fragment_stride)
+        self.warps = threads // WARP_SIZE
+        self.scale_cols = hidden // MXFP8_BLOCK
+
+    def smem_size_in_bytes(self) -> int:
+        return (self.warps + self.cluster_size) * 4
+
+    @cute.jit
+    def __call__(
+        self,
+        contribution_mailbox: cute.Tensor,
+        residual_source: cute.Tensor,
+        gamma: cute.Tensor,
+        residual_output: cute.Tensor,
+        quant_output: cute.Tensor,
+        scale_output: cute.Tensor,
+        stage_state: cute.Tensor,
+        eps: Float32,
+        m: Int32,
+        stream: cuda.CUstream,
+    ) -> None:
+        self.kernel(
+            contribution_mailbox,
+            residual_source,
+            gamma,
+            residual_output,
+            quant_output,
+            scale_output,
+            stage_state,
+            eps,
+            m,
+        ).launch(
+            grid=(m, self.cluster_size, 1),
+            block=(self.threads, 1, 1),
+            cluster=(1, self.cluster_size, 1),
+            smem=self.smem_size_in_bytes(),
+            stream=stream,
+            use_pdl=self.enable_pdl,
+        )
+
+    @cute.kernel
+    def kernel(
+        self,
+        contribution_mailbox: cute.Tensor,
+        residual_source: cute.Tensor,
+        gamma: cute.Tensor,
+        residual_output: cute.Tensor,
+        quant_output: cute.Tensor,
+        scale_output: cute.Tensor,
+        stage_state: cute.Tensor,
+        eps: Float32,
+        m: Int32,
+    ) -> None:
+        tidx, _, _ = cute.arch.thread_idx()
+        token, _, _ = cute.arch.block_idx()
+        cluster_rank = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster())
+        base_fragment = cluster_rank * self.threads + tidx
+
+        prenorm_fragments = cute.make_rmem_tensor(
+            cute.make_layout((self.trips, VEC_BF16), stride=(VEC_BF16, 1)),
+            BFloat16,
+        )
+        prenorm_fragments.fill(BFloat16(0.0))
+        gamma_fragments = cute.make_rmem_tensor(
+            cute.make_layout((self.trips, VEC_BF16), stride=(VEC_BF16, 1)),
+            BFloat16,
+        )
+        gamma_fragments.fill(BFloat16(0.0))
+        residual_fragments = cute.make_rmem_tensor(
+            cute.make_layout((self.trips, VEC_BF16), stride=(VEC_BF16, 1)),
+            BFloat16,
+        )
+        residual_fragments.fill(BFloat16(0.0))
+
+        for trip in cutlass.range_constexpr(self.trips):
+            fragment = base_fragment + trip * self.fragment_stride
+            if fragment < self.fragments:
+                gamma_element = Int64(fragment) * VEC_BF16
+                gamma_pointer = cute.make_ptr(
+                    BFloat16,
+                    (gamma.iterator + gamma_element).llvm_ptr,
+                    cute.AddressSpace.gmem,
+                    assumed_align=16,
+                )
+                gamma_fragments[trip, None].store(
+                    packed_u32x4_to_bf16x8(load_global_u32x4(gamma_pointer))
+                )
+                residual_element = Int64(token) * self.hidden + gamma_element
+                residual_pointer = cute.make_ptr(
+                    BFloat16,
+                    (residual_source.iterator + residual_element).llvm_ptr,
+                    cute.AddressSpace.gmem,
+                    assumed_align=16,
+                )
+                residual_fragments[trip, None].store(
+                    packed_u32x4_to_bf16x8(load_global_u32x4(residual_pointer))
+                )
+
+        if cutlass.const_expr(self.enable_pdl):
+            cute.arch.griddepcontrol_wait()
+
+        active_stage = load_volatile_u32(stage_state.iterator + ACTIVE_STAGE)
+        for trip in cutlass.range_constexpr(self.trips):
+            fragment = base_fragment + trip * self.fragment_stride
+            lane_packed = cute.make_rmem_tensor(cute.make_layout((self.tp, 4)), Uint32)
+            lane_packed.fill(Uint32(0))
+            dirty = fragment < self.fragments
+            while dirty:
+                dirty = False
+                for source_rank in cutlass.range_constexpr(self.tp):
+                    if fragment < self.fragments:
+                        source_element = (
+                            (Int64(active_stage) * self.tp + Int64(source_rank))
+                            * self.capacity_m
+                            + Int64(token)
+                        ) * self.hidden + Int64(fragment) * VEC_BF16
+                        source_pointer = cute.make_ptr(
+                            BFloat16,
+                            (contribution_mailbox.iterator + source_element).llvm_ptr,
+                            cute.AddressSpace.gmem,
+                            assumed_align=16,
+                        )
+                        packed = load_global_u32x4(source_pointer, volatile=True)
+                        dirty = dirty | fragment_has_negative_zero(packed)
+                        for word in cutlass.range_constexpr(4):
+                            lane_packed[source_rank, word] = packed[word]
+
+            if fragment < self.fragments:
+                prenorm = cute.make_rmem_tensor(cute.make_layout((VEC_BF16,)), Float32)
+                prenorm.fill(Float32(0.0))
+                for source_rank in cutlass.range_constexpr(self.tp):
+                    packed = cute.make_rmem_tensor(cute.make_layout((4,)), Uint32)
+                    for word in cutlass.range_constexpr(4):
+                        packed[word] = lane_packed[source_rank, word]
+                    prenorm.store(
+                        prenorm.load()
+                        + packed_u32x4_to_bf16x8(packed.load()).to(Float32)
+                    )
+                prenorm_bf16 = (
+                    prenorm.load() + residual_fragments[trip, None].load().to(Float32)
+                ).to(BFloat16)
+                prenorm_fragments[trip, None].store(prenorm_bf16)
+                output_element = Int64(token) * self.hidden + Int64(fragment) * VEC_BF16
+                store_global_u32x4(
+                    Int64((residual_output.iterator + output_element).toint()),
+                    bf16x8_to_packed_u32x4(prenorm_bf16),
+                )
+
+        if token == 0 and cluster_rank == 0 and tidx == 0:
+            store_global_u32(
+                (stage_state.iterator + NEXT_STAGE).toint(),
+                (active_stage + Uint32(1)) % Uint32(LAMPORT_GENERATIONS),
+            )
+        if cutlass.const_expr(self.enable_pdl):
+            cute.arch.griddepcontrol_launch_dependents()
+
+        for trip in cutlass.range_constexpr(self.trips):
+            fragment = base_fragment + trip * self.fragment_stride
+            for source_rank in cutlass.range_constexpr(self.tp):
+                if fragment < self.fragments:
+                    source_element = (
+                        (Int64(active_stage) * self.tp + Int64(source_rank))
+                        * self.capacity_m
+                        + Int64(token)
+                    ) * self.hidden + Int64(fragment) * VEC_BF16
+                    store_lamport_sentinel_u32x4(
+                        Int64((contribution_mailbox.iterator + source_element).toint())
+                    )
+
+        thread_sum = Float32(0.0)
+        for trip in cutlass.range_constexpr(self.trips):
+            fragment = base_fragment + trip * self.fragment_stride
+            if fragment < self.fragments:
+                values = prenorm_fragments[trip, None].load().to(Float32)
+                thread_sum = thread_sum + (values * values).reduce(
+                    cute.ReductionOp.ADD,
+                    init_val=Float32(0.0),
+                    reduction_profile=0,
+                )
+
+        smem = cutlass.utils.SmemAllocator()
+        warp_sums = smem.allocate_array(Float32, self.warps)
+        cluster_sums = smem.allocate_array(Float32, self.cluster_size)
+        cta_sum = group_leader_block_sum(thread_sum, warp_sums, self.warps, 1)
+        if tidx < self.cluster_size:
+            local_slot = cluster_sums + cluster_rank
+            remote_slot = map_shared_to_peer(local_slot, Int32(tidx))
+            store_shared_cluster_f32(remote_slot, cta_sum)
+        cute.arch.cluster_arrive()
+        cute.arch.cluster_wait()
+
+        full_sum = Float32(0.0)
+        for peer in cutlass.range_constexpr(self.cluster_size):
+            full_sum = full_sum + cute.arch.load(
+                (cluster_sums + peer).llvm_ptr,
+                Float32,
+            )
+        inv_rms = cute.math.rsqrt(
+            full_sum / Float32(self.hidden) + eps,
+            fastmath=True,
+        )
+        padded_cols = Int32(self.scale_cols)
+        for trip in cutlass.range_constexpr(self.trips):
+            fragment = base_fragment + trip * self.fragment_stride
+            # Gemma RMSNorm scales by (1 + weight); the unfused norm returns BF16.
+            normed = (
+                (
+                    prenorm_fragments[trip, None].load().to(Float32)
+                    * inv_rms
+                    * (gamma_fragments[trip, None].load().to(Float32) + Float32(1.0))
+                )
+                .to(BFloat16)
+                .to(Float32)
+            )
+            values = cute.make_rmem_tensor(cute.make_layout((VEC_BF16,)), Float32)
+            values.store(normed)
+            amax = Float32(0.0)
+            for element in cutlass.range_constexpr(VEC_BF16):
+                amax = cute.arch.fmax(
+                    amax, cute.arch.fmax(values[element], -values[element])
+                )
+            # Out-of-range fragments are zero; every lane joins the shuffles.
+            for offset in cutlass.range_constexpr(1, _FRAGMENTS_PER_BLOCK):
+                if cutlass.const_expr((offset & (offset - 1)) == 0):
+                    amax = cute.arch.fmax(
+                        amax,
+                        cute.arch.shuffle_sync_bfly(
+                            amax, offset=offset, mask=-1, mask_and_clamp=31
+                        ),
+                    )
+            exponent = _ue8m0_ceil(amax * Float32(1.0 / 448.0))
+            inverse = _ue8m0_inverse(exponent)
+            if fragment < self.fragments:
+                low = _f32x4_to_e4m3x4(
+                    values[0] * inverse,
+                    values[1] * inverse,
+                    values[2] * inverse,
+                    values[3] * inverse,
+                )
+                high = _f32x4_to_e4m3x4(
+                    values[4] * inverse,
+                    values[5] * inverse,
+                    values[6] * inverse,
+                    values[7] * inverse,
+                )
+                quant_element = Int64(token) * self.hidden + Int64(fragment) * VEC_BF16
+                _store_global_u32x2(
+                    Int64((quant_output.iterator + quant_element).toint()), low, high
+                )
+                if fragment % _FRAGMENTS_PER_BLOCK == 0:
+                    col = fragment // _FRAGMENTS_PER_BLOCK
+                    scale_output[_swizzled_scale_offset(token, col, padded_cols)] = (
+                        exponent.to(Uint8)
+                    )
+
+        # Zero the scales of the last row tile's padding rows. A 4-byte word of
+        # a 512-byte block holds the 4 blocks of one row, so every thread of
+        # the grid clears whole words.
+        tile_words = Int32(self.scale_cols * _SCALE_ROW_TILE // 4)
+        last_tile = (m - Int32(1)) // Int32(_SCALE_ROW_TILE)
+        word = (token * Int32(self.cluster_size) + cluster_rank) * Int32(
+            self.threads
+        ) + tidx
+        while word < tile_words:
+            block_word = word % Int32(128)
+            row = (
+                last_tile * Int32(_SCALE_ROW_TILE)
+                + block_word // Int32(4)
+                + block_word % Int32(4) * Int32(32)
+            )
+            if row >= m:
+                byte = (Int64(last_tile) * Int64(tile_words) + Int64(word)) * 4
+                store_global_u32(
+                    cute.make_ptr(
+                        Uint32,
+                        (scale_output.iterator + byte).llvm_ptr,
+                        cute.AddressSpace.gmem,
+                        assumed_align=4,
+                    ).toint(),
+                    Uint32(0),
+                )
+            word = word + m * Int32(self.cluster_size * self.threads)
+
+
+def mxfp8_scale_numel(num_tokens: int, hidden_size: int) -> int:
+    """Bytes of the 128x4 swizzled UE8M0 scales of ``num_tokens`` rows."""
+    rows = math.ceil(num_tokens / _SCALE_ROW_TILE) * _SCALE_ROW_TILE
+    cols = math.ceil(hidden_size // MXFP8_BLOCK / _SCALE_COL_TILE) * _SCALE_COL_TILE
+    return rows * cols
+
+
+class FinalizeAllReduceGemmaRMSNormMxfp8:
+    """MoE top-k finalize, TP all-reduce, residual add, Gemma RMSNorm and MXFP8
+    quantization in two PDL-chained launches.
+
+    ``__call__`` takes an unfinalized MoE output plus its TP-partial shared-expert
+    output and returns the MXFP8 normalized rows with the new residual. Each
+    instance owns a Lamport mailbox for ``max_num_tokens`` rows that all layers
+    share; construction is collective over the TP group.
+    """
+
+    def __init__(
+        self,
+        *,
+        hidden_size: int,
+        max_num_tokens: int,
+        top_k: int,
+        fp32_weights: bool,
+        device: torch.device,
+    ) -> None:
+        group = get_tp_group().device_group
+        self.hidden_size = hidden = hidden_size
+        self.capacity = capacity = max_num_tokens
+        self.top_k = top_k
+        self.fp32_weights = fp32_weights
+        tp = dist.get_world_size(group)
+        rank = dist.get_rank(group)
+        device = torch.device(device)
+
+        # A cluster of 128-thread CTAs covers a token's fragments in one trip
+        # (six for hidden 6144); the finalize stages all top_k routes at once.
+        cluster_size = math.ceil(hidden // VEC_BF16 / 128)
+        with torch.accelerator.device_index(device.index):
+            self._publish = cute.compile(
+                _QuadFinalizePublishDeviceKernel(
+                    hidden=hidden,
+                    top_k=top_k,
+                    tp=tp,
+                    rank=rank,
+                    capacity_m=capacity,
+                    threads=128,
+                    # The routing weights already carry routed_scaling_factor.
+                    routed_scaling_factor=1.0,
+                    include_shared_expert=True,
+                    load_shared_expert_before_pdl=False,
+                    enable_pdl=True,
+                    prefetch_group=top_k,
+                    fp32_weights=fp32_weights,
+                ),
+                make_fake_dynamic_compact_tensor(
+                    BFloat16, alignment=16, divisibility=hidden
+                ),
+                make_fake_dynamic_compact_tensor(
+                    Float32 if fp32_weights else BFloat16,
+                    alignment=4 if fp32_weights else 2,
+                    divisibility=top_k,
+                ),
+                make_fake_dynamic_compact_tensor(
+                    Int32, alignment=4, divisibility=top_k
+                ),
+                make_fake_dynamic_compact_tensor(
+                    BFloat16, alignment=16, divisibility=hidden
+                ),
+                make_fake_compact_tensor(Int32, (2,), assumed_align=4),
+                Int64(0),
+                Int32(capacity),
+                current_cu_stream(),
+            )
+            self._collective = cute.compile(
+                _LamportGemmaRMSNormMxfp8Kernel(
+                    hidden=hidden,
+                    tp=tp,
+                    capacity_m=capacity,
+                    cluster_size=cluster_size,
+                    threads=128,
+                    enable_pdl=True,
+                ),
+                make_fake_compact_tensor(
+                    BFloat16,
+                    (LAMPORT_GENERATIONS * tp * capacity * hidden,),
+                    assumed_align=16,
+                ),
+                make_fake_dynamic_compact_tensor(
+                    BFloat16, alignment=16, divisibility=hidden
+                ),
+                make_fake_compact_tensor(BFloat16, (hidden,), assumed_align=16),
+                make_fake_dynamic_compact_tensor(
+                    BFloat16, alignment=16, divisibility=hidden
+                ),
+                make_fake_dynamic_compact_tensor(
+                    Uint8, alignment=16, divisibility=hidden
+                ),
+                make_fake_dynamic_compact_tensor(
+                    Uint8, alignment=16, divisibility=_SCALE_ROW_TILE
+                ),
+                make_fake_compact_tensor(Int32, (2,), assumed_align=4),
+                Float32(0.0),
+                Int32(capacity),
+                current_cu_stream(),
+            )
+
+            self._mailbox = symm_mem.empty(
+                (LAMPORT_GENERATIONS, tp, capacity, hidden),
+                dtype=torch.bfloat16,
+                device=device,
+            )
+            self._mailbox_handle = symm_mem.rendezvous(self._mailbox, group)
+            multicast = int(self._mailbox_handle.multicast_ptr or 0)
+            if not multicast or multicast % 16:
+                raise RuntimeError("NVLink multicast mapping is unavailable")
+            self._multicast = multicast
+            # Every slot starts as the Lamport sentinel, BF16 negative zero.
+            self._mailbox.view(torch.int16).fill_(-32768)
+            self._stage_state = torch.zeros(2, dtype=torch.int32, device=device)
+            torch.accelerator.synchronize()
+        dist.barrier(group=group)
+
+    def __call__(
+        self,
+        moe_output: MoEOutput,
+        residual: torch.Tensor,
+        norm_weight: torch.Tensor,
+        norm_eps: float,
+    ) -> tuple[QuantizedActivation, torch.Tensor]:
+        """Return (MXFP8 normalized rows, new residual).
+
+        Args:
+            moe_output: Unfinalized TP-partial routed output and the TP-partial
+                shared-expert output, with routed_scaling_factor in the weights.
+            residual: [M, hidden] BF16 residual stream.
+            norm_weight: [hidden] BF16 Gemma RMSNorm weight (scale is 1 + w).
+            norm_eps: RMSNorm epsilon.
+
+        """
+        routed = moe_output.routed
+        shared_output = moe_output.shared_output
+        assert isinstance(routed, UnfinalizedMoEOutput)
+        assert shared_output is not None
+        assert moe_output.routed_scaling_factor == 1.0
+        hidden, top_k = self.hidden_size, self.top_k
+        m = residual.shape[0]
+        if not 1 <= m <= self.capacity:
+            raise ValueError(f"M={m} is outside [1, {self.capacity}]")
+        if residual.shape != (m, hidden) or shared_output.shape != (m, hidden):
+            raise ValueError("residual and shared_output must be [M, hidden]")
+        weights_dtype = torch.float32 if self.fp32_weights else torch.bfloat16
+        if (
+            routed.expert_weights.dtype != weights_dtype
+            or routed.expert_weights.numel() != m * top_k
+        ):
+            raise ValueError(f"expert_weights must be {weights_dtype} [M, top_k]")
+        self._publish(
+            to_cute_dynamic(routed.gemm2_permuted.flatten(), 16, divisibility=hidden),
+            to_cute_dynamic(
+                routed.expert_weights.flatten(),
+                4 if self.fp32_weights else 2,
+                divisibility=top_k,
+            ),
+            to_cute_dynamic(
+                routed.expanded_idx_to_permuted_idx.flatten(), 4, divisibility=top_k
+            ),
+            to_cute_dynamic(shared_output.flatten(), 16, divisibility=hidden),
+            to_cute(self._stage_state, 4),
+            Int64(self._multicast),
+            Int32(m),
+            current_cu_stream(),
+        )
+        residual_output = torch.empty_like(residual)
+        quantized = torch.empty(
+            (m, hidden), dtype=torch.float8_e4m3fn, device=residual.device
+        )
+        scales = torch.empty(
+            mxfp8_scale_numel(m, hidden), dtype=torch.uint8, device=residual.device
+        )
+        self._collective(
+            to_cute(self._mailbox.flatten(), 16),
+            to_cute_dynamic(residual.flatten(), 16, divisibility=hidden),
+            to_cute(norm_weight, 16),
+            to_cute_dynamic(residual_output.flatten(), 16, divisibility=hidden),
+            to_cute_dynamic(
+                quantized.view(torch.uint8).flatten(), 16, divisibility=hidden
+            ),
+            to_cute_dynamic(scales, 16, divisibility=_SCALE_ROW_TILE),
+            to_cute(self._stage_state, 4),
+            Float32(norm_eps),
+            Int32(m),
+            current_cu_stream(),
+        )
+        return (
+            QuantizedActivation(
+                quantized, scales, residual.dtype, quantized.shape, kMxfp8Dynamic
+            ),
+            residual_output,
+        )
+
+
+_fused_moe_finalize: FinalizeAllReduceGemmaRMSNormMxfp8 | None = None
+
+
+def supports_fused_moe_finalize(vllm_config: VllmConfig) -> bool:
+    """Whether to build the fused MoE finalize kernel: tensor parallel without
+    expert parallel, a BF16 model on SM100-family GPUs, and NVLink multicast on
+    every TP rank. Every TP rank calls it; the multicast check is collective.
+    """
+    from vllm.distributed.device_communicators.custom_all_reduce import (
+        _all_ranks_true,
+        _has_local_multicast_support,
+    )
+
+    if (
+        get_tensor_model_parallel_world_size() <= 1
+        or vllm_config.parallel_config.enable_expert_parallel
+        or not current_platform.is_cuda()
+        or not current_platform.is_device_capability_family(100)
+        or vllm_config.model_config.dtype != torch.bfloat16
+    ):
+        return False
+    device = torch.device(
+        current_platform.device_type, torch.accelerator.current_device_index()
+    )
+    return _all_ranks_true(
+        get_tp_group().cpu_group, _has_local_multicast_support(device)
+    )
+
+
+def init_fused_moe_finalize(
+    *, hidden_size: int, max_num_tokens: int, top_k: int, fp32_weights: bool
+) -> None:
+    """Build the kernel every deferred MoE output is finalized with.
+
+    Collective over the TP group, so every rank calls it, and only when
+    ``supports_fused_moe_finalize`` holds.
+    """
+    global _fused_moe_finalize
+    _fused_moe_finalize = FinalizeAllReduceGemmaRMSNormMxfp8(
+        hidden_size=hidden_size,
+        max_num_tokens=max_num_tokens,
+        top_k=top_k,
+        fp32_weights=fp32_weights,
+        device=torch.device(
+            current_platform.device_type, torch.accelerator.current_device_index()
+        ),
+    )
+
+
+def finalize_allreduce_gemma_rms_norm_mxfp8(
+    moe_output: MoEOutput,
+    residual: torch.Tensor,
+    norm_weight: torch.Tensor,
+    norm_eps: float,
+) -> tuple[QuantizedActivation, torch.Tensor]:
+    """Finalize a deferred MoE output into the next layer's MXFP8 input and
+    new residual."""
+    assert _fused_moe_finalize is not None, "init_fused_moe_finalize was not called"
+    return _fused_moe_finalize(moe_output, residual, norm_weight, norm_eps)
diff --git a/vllm/models/minimax_m3/nvidia/ops/gather_nvfp4_index_k.py b/vllm/models/minimax_m3/nvidia/ops/gather_nvfp4_index_k.py
new file mode 100644
index 000000000000..0c9f477180ff
--- /dev/null
+++ b/vllm/models/minimax_m3/nvidia/ops/gather_nvfp4_index_k.py
@@ -0,0 +1,174 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Gather NVFP4 index-K pages into compact FP8 pages for the Q8KV8 prefill indexer.
+
+A packed NVFP4 page ``[128, 72]`` holds the E2M1 codes of all 128 tokens (64 B
+per token, low nibble first) followed by their E4M3 16-element block scales
+(``token * 8 + group``), with no global scale. Each element dequantizes to
+``e4m3_rn_satfinite(e2m1 * scale)``, the value the Q8KV4 decode indexer scores,
+so prefill and decode rank pages from identical keys.
+
+Where Triton can emit it, each four elements are one QMUL4
+(``mul.rn.satfinite.e4m3x4.e2m1x4.e4m3x4``), the instruction the decode kernel
+dequantizes with; otherwise the product is formed exactly in fp32 and rounded
+once, which gives the same bytes.
+"""
+
+import functools
+
+import torch
+
+from vllm.platforms import current_platform
+from vllm.triton_utils import tl, triton
+
+PAGE_SIZE = 128
+HEAD_DIM = 128
+E4M3_MAX = 448.0
+# QMUL4 needs PTX ISA 9.4 (ptxas 13.4) and an sm_100a/sm_103a target.
+QMUL4_MIN_PTX_VERSION = 94
+
+
+@functools.cache
+def _qmul4_supported(device_index: int) -> bool:
+    """Whether Triton emits PTX that can use QMUL4 on this device.
+
+    Triton takes the PTX ISA version from its Blackwell ptxas: the bundled one
+    unless ``TRITON_PTXAS_BLACKWELL_PATH`` points at a CUDA 13.4+ ptxas.
+    """
+    capability = current_platform.get_device_capability(device_index)
+    if capability is None or capability.to_int() not in (100, 103):
+        return False
+    try:
+        from triton.backends.nvidia.compiler import get_ptxas, ptx_get_version
+
+        ptxas_version = get_ptxas(capability.to_int()).version
+        return ptx_get_version(ptxas_version) >= QMUL4_MIN_PTX_VERSION
+    except Exception:
+        # Triton-internal helpers; without them keep the portable path.
+        return False
+
+
+@triton.jit
+def _e2m1_to_f32(code):
+    # |code| 0..7 -> 0, 0.5, 1, 1.5, 2, 3, 4, 6; bit 3 is the sign.
+    exponent = (code >> 1) & 3
+    mantissa = (code & 1).to(tl.float32)
+    normal = (2.0 + mantissa) * (1 << exponent).to(tl.float32) * 0.25
+    value = tl.where(exponent == 0, 0.5 * mantissa, normal)
+    # Multiply (not negate) so that code 8 stays -0, as in the decode kernel.
+    return value * tl.where((code & 8) != 0, -1.0, 1.0)
+
+
+@triton.jit
+def _gather_nvfp4_index_k_fp8_kernel(
+    kv_cache,
+    block_table,
+    num_pages,
+    page_offsets,
+    out,
+    kv_page_stride,
+    block_table_stride,
+    PAGE_SIZE: tl.constexpr,
+    HEAD_DIM: tl.constexpr,
+    E4M3_MAX: tl.constexpr,
+    USE_QMUL4: tl.constexpr,
+):
+    req = tl.program_id(0)
+    col = tl.program_id(1)
+    if col >= tl.load(num_pages + req):
+        return
+    page = tl.load(block_table + req * block_table_stride + col).to(tl.int64)
+    dst_page = (tl.load(page_offsets + req) + col).to(tl.int64)
+
+    src = kv_cache + page * kv_page_stride
+    dst = out + dst_page * (PAGE_SIZE * HEAD_DIM)
+    tokens = tl.arange(0, PAGE_SIZE)[:, None]
+    scales = src + PAGE_SIZE * (HEAD_DIM // 2) + tokens * (HEAD_DIM // 16)
+    if USE_QMUL4:
+        # A 16-bit word holds four elements of one group; multiply it by the
+        # group's scale byte broadcast to E4M3x4.
+        words = tl.arange(0, HEAD_DIM // 4)[None, :]
+        packed4 = tl.load(
+            src.to(tl.pointer_type(tl.uint16)) + tokens * (HEAD_DIM // 4) + words
+        )
+        scale4 = tl.load(scales + words // 4).to(tl.uint32) * 0x01010101
+        values4 = tl.inline_asm_elementwise(
+            "{ .reg .b16 v; cvt.u16.u32 v, $1; "
+            "mul.rn.satfinite.e4m3x4.e2m1x4.e4m3x4 $0, v, $2; }",
+            "=r,r,r",
+            [packed4.to(tl.uint32), scale4],
+            dtype=tl.uint32,
+            is_pure=True,
+            pack=1,
+        )
+        tl.store(
+            dst.to(tl.pointer_type(tl.uint32)) + tokens * (HEAD_DIM // 4) + words,
+            values4,
+        )
+        return
+
+    packed_cols = tl.arange(0, HEAD_DIM // 2)[None, :]
+    packed = tl.load(src + tokens * (HEAD_DIM // 2) + packed_cols)
+    # Packed byte b holds elements 2b and 2b+1, both in group b // 8.
+    scale = tl.load(scales + packed_cols // 8)
+    scale = scale.to(tl.float8e4nv, bitcast=True).to(tl.float32)
+    # e2m1 * e4m3 is exact in fp32; clamp, then round once to e4m3.
+    low = _e2m1_to_f32(packed & 0xF) * scale
+    high = _e2m1_to_f32(packed >> 4) * scale
+    values = tl.interleave(low, high)
+    values = tl.minimum(tl.maximum(values, -E4M3_MAX), E4M3_MAX)
+    tl.store(
+        dst + tokens * HEAD_DIM + tl.arange(0, HEAD_DIM)[None, :],
+        values.to(tl.float8e4nv),
+    )
+
+
+def minimax_m3_gather_nvfp4_index_k_fp8(
+    kv_cache: torch.Tensor,
+    block_table: torch.Tensor,
+    num_pages: torch.Tensor,
+    page_offsets: torch.Tensor,
+    max_num_pages: int,
+    out: torch.Tensor,
+) -> torch.Tensor:
+    """Dequantize each request's leading NVFP4 pages into contiguous FP8 pages.
+
+    Args:
+        kv_cache: ``[num_blocks, 128, 72]`` uint8 packed NVFP4 pages.
+        block_table: ``[num_reqs, max_blocks]`` int32 physical page ids.
+        num_pages: ``[num_reqs]`` int32 pages to gather per request (its first
+            ``num_pages[r]`` logical pages).
+        page_offsets: ``[num_reqs]`` int32 first output page of each request.
+        max_num_pages: Host upper bound of ``num_pages``.
+        out: ``[>= sum(num_pages), 128, 128]`` contiguous float8_e4m3fn.
+
+    Returns:
+        ``out``.
+
+    """
+    assert kv_cache.dtype == torch.uint8 and kv_cache.shape[1:] == (
+        PAGE_SIZE,
+        HEAD_DIM // 2 + HEAD_DIM // 16,
+    )
+    assert kv_cache.stride(1) == kv_cache.shape[2] and kv_cache.stride(2) == 1
+    assert out.dtype == torch.float8_e4m3fn and out.is_contiguous()
+    assert out.shape[1:] == (PAGE_SIZE, HEAD_DIM)
+    assert block_table.stride(1) == 1
+    num_reqs = block_table.shape[0]
+    if num_reqs == 0 or max_num_pages == 0:
+        return out
+    _gather_nvfp4_index_k_fp8_kernel[(num_reqs, max_num_pages)](
+        kv_cache,
+        block_table,
+        num_pages,
+        page_offsets,
+        out,
+        kv_cache.stride(0),
+        block_table.stride(0),
+        PAGE_SIZE=PAGE_SIZE,
+        HEAD_DIM=HEAD_DIM,
+        E4M3_MAX=E4M3_MAX,
+        USE_QMUL4=_qmul4_supported(kv_cache.device.index or 0),
+        num_warps=4,
+    )
+    return out
diff --git a/vllm/models/minimax_m3/nvidia/ops/norm_mxfp8.py b/vllm/models/minimax_m3/nvidia/ops/norm_mxfp8.py
new file mode 100644
index 000000000000..406f523ca8e0
--- /dev/null
+++ b/vllm/models/minimax_m3/nvidia/ops/norm_mxfp8.py
@@ -0,0 +1,164 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Gemma RMSNorm (optionally fused with the residual add) in one Triton pass.
+
+With an ``sf`` scale buffer the output is E4M3 data plus UE8M0 scales in the
+128x4 swizzled layout, replacing the BF16 norm followed by a separate MXFP8
+quantize (which re-reads the full normed activation). Without it, the output
+is the plain BF16 norm. The residual is bitwise the BF16 norm's; the
+quantized data and scales match the reference quantize within one E4M3 step
+(the rstd reduction order differs).
+"""
+
+import torch
+
+from vllm.platforms import current_platform
+from vllm.triton_utils import tl, triton
+from vllm.utils.math_utils import round_up
+
+MXFP8_BLOCK = 32
+SF_ROW_TILE = 128
+
+
+@triton.jit(do_not_specialize=["M"])
+def _gemma_rmsnorm_mxfp8_kernel(
+    x_ptr,
+    res_ptr,
+    w_ptr,
+    out_ptr,
+    sf_ptr,
+    M,
+    x_stride0,
+    res_stride0,
+    eps,
+    H: tl.constexpr,
+    launch_pdl: tl.constexpr,
+):
+    SF_ROW_TILE: tl.constexpr = 128
+    G: tl.constexpr = H // 32
+    GP: tl.constexpr = triton.next_power_of_2(H) // 32
+    PADDED_COLS: tl.constexpr = (G + 3) // 4 * 4
+    row = tl.program_id(0).to(tl.int64)
+    g = tl.arange(0, GP)
+    if launch_pdl:
+        # All CTAs wait for the predecessor before any global access (the
+        # padding-row stores included), then release dependents before the
+        # loads: the consumer (the weight-loading GEMM) overlaps.
+        tl.extra.cuda.gdc_wait()
+        tl.extra.cuda.gdc_launch_dependents()
+    if sf_ptr is not None:
+        # Swizzled 128x4 scale layout [m_tile, k_tile, 32, 4, 4].
+        sf_off = (
+            (g % 4)
+            + (g // 4) * 512
+            + (row % 32) * 16
+            + ((row % SF_ROW_TILE) // 32) * 4
+            + (row // SF_ROW_TILE) * (SF_ROW_TILE * PADDED_COLS)
+        )
+        if row >= M:
+            # Zero the padding rows of the swizzled scale tensor.
+            tl.store(sf_ptr + sf_off, tl.zeros((GP,), tl.uint8), mask=g < G)
+            return
+
+    cols = g[:, None] * 32 + tl.arange(0, 32)[None, :]
+    mask = cols < H
+    x = tl.load(x_ptr + row * x_stride0 + cols, mask=mask, other=0.0)
+    w = tl.load(w_ptr + cols, mask=mask, other=0.0)
+    if res_ptr is not None:
+        r = tl.load(res_ptr + row * res_stride0 + cols, mask=mask, other=0.0)
+    h = x.to(tl.float32)
+    if res_ptr is not None:
+        h += r.to(tl.float32)
+        tl.store(res_ptr + row * res_stride0 + cols, h, mask=mask)
+    var = tl.sum(h * h) / H
+    rstd = tl.rsqrt(var + eps)
+    y = (h * rstd * (1.0 + w.to(tl.float32))).to(tl.bfloat16)
+
+    if sf_ptr is not None:
+        amax = tl.max(tl.abs(y.to(tl.float32)), axis=1)
+        scale = amax * (1.0 / 448.0)
+        # UE8M0 scale: ceil to a power of two (exponent bump when the
+        # mantissa is nonzero, except tiny subnormals), clamped to [0, 254].
+        bits = scale.to(tl.uint32, bitcast=True)
+        exp = (bits >> 23) & 255
+        mant = bits & 0x7FFFFF
+        bump = (mant != 0) & ~((exp == 0) & (mant <= 0x400000))
+        e = tl.minimum(exp + bump.to(tl.uint32), 254)
+        e = tl.where(scale <= 0.0, 0, e)
+        # Inverse scale 2^(127 - e) (0 for e == 0), always a power of two.
+        inv_bits = tl.where(e == 0, 0, (254 - e) << 23)
+        inv = inv_bits.to(tl.float32, bitcast=True)
+        # Bitwise the reference conversion (saturating; NaN -> +448);
+        # checked exhaustively over all BF16 values and power-of-two scales.
+        q = tl.clamp(y.to(tl.float32) * inv[:, None], -448.0, 448.0)
+        tl.store(
+            out_ptr + row * H + cols,
+            q.to(tl.float8e4nv).to(tl.uint8, bitcast=True),
+            mask=mask,
+        )
+        tl.store(sf_ptr + sf_off, e.to(tl.uint8), mask=g < G)
+    else:
+        tl.store(out_ptr + row * H + cols, y, mask=mask)
+
+
+def gemma_rmsnorm_mxfp8(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float,
+    residual: torch.Tensor | None = None,
+    mxfp8: bool = True,
+) -> tuple[torch.Tensor, torch.Tensor | None]:
+    """GemmaRMSNorm of ``x`` (plus ``residual`` in place when given).
+
+    Args:
+        x: (num_tokens, hidden) BF16 activation, contiguous along the row
+            (any row stride). Read only.
+        weight: (hidden,) norm weight; the norm scales by (1 + weight).
+        eps: RMS norm epsilon.
+        residual: same shape/dtype/stride(1) as x; updated in place to
+            x + residual. None for a plain norm.
+        mxfp8: return MXFP8 (E4M3 data + UE8M0 scales) or the plain BF16
+            normed activation.
+
+    Returns:
+        ``(data, scale)``: with ``mxfp8=True``, E4M3 data (num_tokens,
+        hidden) and UE8M0 scales in the 128x4 swizzled layout with zeroed
+        padding rows; with ``mxfp8=False``, the BF16 normed activation and
+        None.
+
+    """
+    assert x.dim() == 2 and x.stride(1) == 1, x.shape
+    assert weight.is_contiguous()
+    assert residual is None or (residual.shape == x.shape and residual.stride(1) == 1)
+    M, H = x.shape
+    sf = None
+    if mxfp8:
+        padded_m = round_up(M, SF_ROW_TILE)
+        out = torch.empty((M, H), dtype=torch.uint8, device=x.device)
+        sf = torch.empty(
+            padded_m * round_up(H // MXFP8_BLOCK, 4), dtype=torch.uint8, device=x.device
+        )
+        if M == 0:
+            return out.view(torch.float8_e4m3fn), sf
+    else:
+        padded_m = M
+        out = torch.empty((M, H), dtype=torch.bfloat16, device=x.device)
+        if M == 0:
+            return out, None
+    _gemma_rmsnorm_mxfp8_kernel[(padded_m,)](
+        x,
+        residual,
+        weight,
+        out,
+        sf,
+        M,
+        x.stride(0),
+        residual.stride(0) if residual is not None else x.stride(0),
+        eps,
+        H=H,
+        launch_pdl=current_platform.is_arch_support_pdl(),
+        num_warps=4,
+    )
+    if mxfp8:
+        return out.view(torch.float8_e4m3fn), sf
+    return out, None
diff --git a/vllm/models/minimax_m3/nvidia/ops/qkv_index_fused_epilogue.py b/vllm/models/minimax_m3/nvidia/ops/qkv_index_fused_epilogue.py
new file mode 100644
index 000000000000..e531f4f9e792
--- /dev/null
+++ b/vllm/models/minimax_m3/nvidia/ops/qkv_index_fused_epilogue.py
@@ -0,0 +1,1101 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""MXFP8 QKV+indexer GEMM with a fused QK-norm/RoPE + KV-insert epilogue.
+
+Replaces mm_mxfp8 + fused_minimax_m3_qknorm_rope_kv_insert for the MiniMax-M3
+qkv_proj with one persistent SM100 kernel. Per 128-col head tile, the
+epilogue applies Gemma RMSNorm + partial NeoX RoPE to the f32 accumulator
+(no bf16 round trip) and writes q (bf16, optionally fp8) and index_q (fp8),
+or NVFP4-quantizes k / v / index_k into the paged caches. The qkv projection
+is never materialized.
+"""
+
+from dataclasses import dataclass
+from functools import cache
+
+import cutlass
+import torch
+from cuda.bindings.driver import CUstream
+from cutlass import (
+    BFloat16,
+    Float8E4M3FN,
+    Float8E8M0FNU,
+    Float32,
+    Int32,
+    Int64,
+    Uint8,
+    Uint16,
+    Uint32,
+    Uint64,
+    cute,
+)
+from cutlass.cute.nvgpu import cpasync, tcgen05
+from cutlass.cute.runtime import make_fake_tensor
+
+from vllm.cute_utils import _tcgen05, mbarrier, simple_tma_copy, to_cta0_smem
+from vllm.cute_utils.cvt import (
+    bf16x2_to_fp32x2,
+    fp32x8_to_fp4x8,
+)
+from vllm.model_executor.warmup.jit_warmup_cutedsl_helper import VllmCuTeDSLJitKernel
+
+HEAD_DIM = 128  # one 128-col N-tile per head
+# epilogue outputs are consumed by other kernels: stores bypass L1
+_NO_L1_ALLOC = dict(level1_eviction_priority="evict_no_allocate")
+
+
+@cute.jit
+def _rms_norm(acc, sW, eps):
+    """Gemma RMS-norm the 128-col head row in acc in place; sW is the head
+    type's bf16 weight as (8, 16) tiles."""
+    acc1 = cute.make_tensor(acc.iterator, cute.make_layout(HEAD_DIM))
+    # 4 independent fma chains; starting from 0 makes every step
+    # fma(x, x, ss) (x0 * x0 + x1 * x1 would let ptxas fuse either product)
+    ss = cute.full(4, 0.0, Float32)
+    for j in cutlass.range_constexpr(HEAD_DIM // 4):
+        x = cute.local_tile(acc1, (4,), (j,)).load()
+        ss = ss + x * x
+    t = (ss[0] + ss[1] + ss[2] + ss[3]) * Float32(1.0 / HEAD_DIM) + eps
+    # t >= eps, so rsqrt.approx needs no subnormal handling
+    rms = cute.math.rsqrt(t, approx=True)
+
+    # load 8 norm weights at a time -> 16-byte loads
+    for j in cutlass.range_constexpr(HEAD_DIM // 8):
+        w_f32 = bf16x2_to_fp32x2(sW[None, j].load())
+        x = cute.local_tile(acc1, (8,), (j,))
+        n = x.load() * rms
+        x.store(n + n * w_f32)
+
+
+def _rope64(acc, CS, pos):
+    """NeoX-rope cols [0, 64) of acc in place, rotating chunk 0 (x) with
+    chunk 1 (y); CS is the bf16 cos | sin cache as [32, (pos, 2)]."""
+    cos = bf16x2_to_fp32x2(CS[None, (pos, 0)].load())
+    sin = bf16x2_to_fp32x2(CS[None, (pos, 1)].load())
+    x = acc[None, 0].load()
+    y = acc[None, 1].load()
+    acc[None, 0].store(x * cos - y * sin)
+    acc[None, 1].store(y * cos + x * sin)
+
+
+def _nvfp4_cache_views(C, bs, nslots):
+    """Views of the [num_pages, page_bytes] NVFP4 cache C. A page holds nslots
+    slots of bs tokens' E2M1 data (64B each), then their E4M3 scales (8B
+    each). Returns the data as u32 words [8, half, token, slot, page], the
+    linear scales as [8, token, slot, page], and the V-swizzled scales as
+    [(2, 4), token, slot, page]: byte h of token t at
+    (t // 4 * 4 + h // 2) * 8 + h % 2 * 4 + t % 4."""
+    slot_bytes = bs * (HEAD_DIM // 2 + HEAD_DIM // 16)
+    sf_off = bs * HEAD_DIM // 2
+    nb, page = C.shape[0], C.stride[0]
+    data = cute.make_tensor(
+        C.iterator,
+        cute.make_layout((32, 2, bs, nslots, nb), stride=(1, 32, 64, slot_bytes, page)),
+    )
+    sf = cute.make_tensor(
+        C.iterator + sf_off,
+        cute.make_layout((8, bs, nslots, nb), stride=(1, 8, slot_bytes, page)),
+    )
+    vsf = cute.make_tensor(
+        C.iterator + sf_off,
+        cute.make_layout(
+            ((2, 4), (4, bs // 4), nslots, nb),
+            stride=((4, 8), (1, 32), slot_bytes, page),
+        ),
+    )
+    return cute.recast_tensor(data, Uint32), sf, vsf
+
+
+def _nvfp4_quant32(vals, words, sf, sf_scale):
+    """NVFP4-quantize 32 f32 (two 16-value groups) into 4 data words and the
+    2 scale bytes `sf`, following
+    fused_minimax_m3_qknorm_rope_kv_insert's quantizer. sf_scale is
+    1 / (global dequant scale), or None for unit scale."""
+    # E4M3(amax / 6) floored at 2^-9 (satfinite caps it at 448); the unfused
+    # op's 1e-12 amax floor lies below that and is omitted. Both groups'
+    # scales share one e4m3x2 cvt.
+    amax_6 = cute.make_rmem_tensor(2, Float32)
+    for g in range(2):
+        gmax = Float32(0.0)
+        for i in range(16 * g, 16 * g + 16):
+            gmax = cute.arch.fmax(gmax, cute.math.abs(vals[i]))
+        a = gmax * Float32(1.0 / 6.0)
+        if sf_scale is not None:
+            a = a * sf_scale
+        amax_6[g] = cute.arch.fmax(a, Float32(2.0**-9))
+
+    sf8 = amax_6.load().to(Float8E4M3FN)
+    sf.store(sf8.bitcast(Uint8))
+    sf_f32 = sf8.to(Float32)
+
+    q = cute.make_rmem_tensor(32, Float32)
+    for g in range(2):
+        # rcp.approx.ftz is exact on every positive E4M3 value
+        r = cute.arch.rcp_approx(sf_f32[g])
+        if sf_scale is not None:
+            r = r * sf_scale
+        # rz: r >= 1 / sf, so a rounded product could step past an exact E2M1
+        # midpoint and round away from even; truncation lands back on it (raw
+        # v accumulators do hit exact midpoints, rarely)
+        for i in range(16 * g, 16 * g + 16, 2):
+            q[i], q[i + 1] = cute.arch.mul_packed_f32x2(
+                (vals[i], vals[i + 1]), (r, r), rnd="rz"
+            )
+
+    # clear the sign of zero-magnitude nibbles (+7 carries into bit 3 exactly
+    # when the 3-bit magnitude is nonzero)
+    m = Uint32(0x77777777)
+    for j in range(4):
+        w = fp32x8_to_fp4x8(q, 8 * j)
+        words[j] = w & (((w & m) + m) | m)
+
+
+def _nvfp4_quant64(vals, sf_scale=None):
+    """NVFP4-quantize 64 f32 (a 64-col half, one 32-col chunk at a time);
+    returns the 8 data words and the 4 scale bytes."""
+    words = cute.make_rmem_tensor(8, Uint32)
+    sf = cute.make_rmem_tensor(4, Uint8)
+    for c in range(2):
+        _nvfp4_quant32(
+            cute.local_tile(vals, (32,), (c,)),
+            cute.local_tile(words, (4,), (c,)),
+            cute.local_tile(sf, (2,), (c,)),
+            sf_scale,
+        )
+    return words.load(), sf.load()
+
+
+class QkvIndexFusedEpilogueKernel:
+    """Persistent MXFP8 GEMM with a BM x BN = 128 x (128 * cta_group) output
+    tile per CTA (one head per epilogue warpgroup). 2-SM (cta_group=2): a CTA
+    pair runs one m256n256 MMA per k-step over two adjacent M-tiles, each CTA
+    loading half of the BN rows of B. 1-SM (cta_group=1): one m128n128 MMA per
+    CTA (finer tiles for small M). Warps: 4 * cta_group epilogue, then MMA and
+    TMA; cta_group=2 adds 2 padding warps so setmaxnreg can move registers to
+    the epilogue (at cta_group=1, 6 warps already fit 255 regs each).
+    """
+
+    def __init__(
+        self,
+        nkv: int,
+        block_size: int,
+        cta_group: int = 2,
+        group_m: int = 1,
+        use_pdl: bool = True,
+    ) -> None:
+        assert nkv in (1, 2, 4)
+        assert block_size % 4 == 0  # V scales swizzle over 4-token quads
+        assert cta_group in (1, 2)
+        self.cta_group = cta_group
+        self.group_m = group_m
+        self.num_warps = 12 if cta_group == 2 else 6
+        self.tma_group = (
+            tcgen05.CtaGroup.TWO if cta_group == 2 else tcgen05.CtaGroup.ONE
+        )
+        self.use_pdl = use_pdl
+        self.nkv = nkv
+        self.block_size = block_size
+
+        # The N dim is a list of 128-col heads, [q x 16nkv | k x nkv |
+        # v x nkv | index_q x nkv | index_k x 1]; first head of each role
+        # (q starts at 0). An odd head count leaves a phantom head in the
+        # last tile at cta_group=2 (TMA zero-fill, never stored).
+        self.k_start = 16 * nkv
+        self.v_start = self.k_start + nkv
+        self.iq_start = self.v_start + nkv
+        self.ik_start = self.iq_start + nkv
+
+        BM, BN, BK = 128, cta_group * HEAD_DIM, 128
+        self.cta_tile = (BM, BN, BK)
+        # per CTA and stage: A + this CTA's share of B (cta_group=2 splits B
+        # across the pair) + SFA + full SFB (each CTA multicasts its half)
+        self.stage_size = (BM + BN // cta_group) * BK + (BM + BN) * (BK // 32)
+        self.num_stages = (
+            cutlass.memory.SmemAllocator.capacity_in_bytes() // self.stage_size
+        )
+
+    @cute.jit
+    def prepare_AB(self, A: cute.Tensor, BM: int, BK: int) -> cpasync.TmaInfo:
+        tma_op = cpasync.CopyBulkTensorTileG2SOp(cta_group=self.tma_group)
+        swizzle_128B = cute.make_swizzle(3, 4, 3)
+        s_layout = cute.make_layout((BM, BK, self.num_stages), stride=(BK, 1, BM * BK))
+        s_layout = cute.make_composed_layout(swizzle_128B, 0, s_layout)
+        return cpasync.make_tiled_tma_atom(tma_op, A, s_layout, (BM, BK))
+
+    @cute.jit
+    def prepare_SF(
+        self,
+        SF: cute.Tensor,
+        R: Int32,
+        Ksf: Int32,
+        tma_op: cpasync.CopyG2STileBaseOp,
+    ) -> cpasync.TmaInfo:
+        # SF grid [R, Ksf] padded to R % 128 == 0 and permuted as
+        # [R/128, Ksf/4, 32, 4, 4]; the [32, 4, 4] atom is tcgen05's opaque
+        # 512B block, so treat it as bytes. One box = one atom (128 rows x
+        # BK / 32 = 4 K-groups) = 64 Int64 x 1 row-block.
+        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)
+        s_layout = cute.make_layout((64, 1, self.num_stages))
+        return cpasync.make_tiled_tma_atom(tma_op, SF_i64, s_layout, (64, 1))
+
+    @cute.jit
+    def __call__(
+        self,
+        A: cute.Tensor,
+        B: cute.Tensor,
+        SFA: cute.Tensor,
+        SFB: cute.Tensor,
+        Q: cute.Tensor | None,
+        WQ: cute.Tensor,
+        WK: cute.Tensor,
+        WIQ: cute.Tensor,
+        WIK: cute.Tensor,
+        COS_SIN: cute.Tensor,
+        POS: cute.Tensor,
+        KVC: cute.Tensor,
+        IDXC: cute.Tensor,
+        SLOT: cute.Tensor,
+        IDX_SLOT: cute.Tensor,
+        IQFP8: cute.Tensor,
+        QFP8: cute.Tensor | None,
+        eps: Float32,
+        q_fp8_scale: Float32,
+        k_scale: Float32,
+        v_scale: Float32,
+        grid_size: Int32,
+        stream: CUstream,
+    ) -> None:
+        M, K = A.shape
+        N, _ = B.shape
+        BM, BN, BK = self.cta_tile
+
+        A_tma = self.prepare_AB(A, BM, BK)
+        B_tma = self.prepare_AB(B, BN // self.cta_group, BK)
+        SFA_tma = self.prepare_SF(
+            SFA, M, K // 32, cpasync.CopyBulkTensorTileG2SOp(self.tma_group)
+        )
+        sfb_op = (
+            cpasync.CopyBulkTensorTileG2SMulticastOp(self.tma_group)
+            if self.cta_group == 2
+            else cpasync.CopyBulkTensorTileG2SOp(self.tma_group)
+        )
+        SFB_tma = self.prepare_SF(SFB, N, K // 32, sfb_op)
+
+        self.kernel(
+            A_tma,
+            B_tma,
+            SFA_tma,
+            SFB_tma,
+            Q,
+            WQ,
+            WK,
+            WIQ,
+            WIK,
+            COS_SIN,
+            POS,
+            KVC,
+            IDXC,
+            SLOT,
+            IDX_SLOT,
+            IQFP8,
+            QFP8,
+            eps,
+            q_fp8_scale,
+            k_scale,
+            v_scale,
+        ).launch(
+            grid=(grid_size, 1, 1),
+            block=(self.num_warps * 32, 1, 1),
+            cluster=(self.cta_group, 1, 1),
+            min_blocks_per_mp=1,
+            stream=stream,
+            use_pdl=self.use_pdl,
+        )
+
+    @cute.kernel
+    def kernel(
+        self,
+        A_tma: cpasync.TmaInfo,
+        B_tma: cpasync.TmaInfo,
+        SFA_tma: cpasync.TmaInfo,
+        SFB_tma: cpasync.TmaInfo,
+        Q_tensor: cute.Tensor | None,
+        WQ_tensor: cute.Tensor,
+        WK_tensor: cute.Tensor,
+        WIQ_tensor: cute.Tensor,
+        WIK_tensor: cute.Tensor,
+        COS_SIN_tensor: cute.Tensor,
+        POS_tensor: cute.Tensor,
+        KVC_tensor: cute.Tensor,
+        IDXC_tensor: cute.Tensor,
+        SLOT_tensor: cute.Tensor,
+        IDX_SLOT_tensor: cute.Tensor,
+        IQFP8_tensor: cute.Tensor,
+        QFP8_tensor: cute.Tensor | None,
+        eps: Float32,
+        q_fp8_scale: Float32,
+        k_scale: Float32,
+        v_scale: Float32,
+    ) -> None:
+        tid, _, _ = cute.arch.thread_idx()
+        raw_bid, _, _ = cute.arch.block_idx()
+        num_bids, _, _ = cute.arch.grid_dim()
+        warp_id = cute.arch.make_warp_uniform(tid // 32)
+        cg = self.cta_group
+        cta_rank = raw_bid % cg
+
+        BM, BN, BK = self.cta_tile
+        num_stages = self.num_stages
+
+        # tmem (512 cols): double-buffered accumulators, then SFA/SFB0/SFB1.
+        # At cta_group=2, 2 x 256 + SF cols don't fit, so they overlap by 32
+        # cols (even tiles at [0, 256), odd at [224, 480)).
+        acc_odd = BN - 32 if cg == 2 else BN
+        sfa_tmem = 496
+        sfb0_tmem = sfa_tmem + 4
+        sfb1_tmem = sfb0_tmem + 4
+        # q / k / index_q / index_k norm weight index in sW
+        W_Q, W_K, W_IQ, W_IK = range(4)
+        bs = self.block_size
+
+        BAR_EPI = 1
+        BAR_TMEM_ALLOC = 2
+        cta_pair = Uint16(3) if cg == 2 else None  # multicast / commit mask
+
+        smem = cutlass.memory.SmemAllocator()
+        sA = smem.allocate_tensor(
+            Float8E4M3FN,
+            A_tma.smem_layout.outer,
+            byte_alignment=128,
+            swizzle=A_tma.smem_layout.inner,
+        )
+        sB = smem.allocate_tensor(
+            Float8E4M3FN,
+            B_tma.smem_layout.outer,
+            byte_alignment=128,
+            swizzle=B_tma.smem_layout.inner,
+        )
+        sSFA = smem.allocate_tensor(Int64, SFA_tma.smem_layout, byte_alignment=128)
+        # SFB slots: CTA r's multicast copy (head tile 2 * bid_n + r) fills slot
+        # r of both CTAs; each slot matches the atom's (64, 1):(1, 64) tile
+        sSFB = smem.allocate_tensor(
+            Int64,
+            cute.make_layout((64, 1, cg, num_stages), stride=(1, 64, 64, 64 * cg)),
+            byte_alignment=128,
+        )
+        # norm weights as (8, 16) tiles per head type
+        sW = smem.allocate_tensor(
+            BFloat16, cute.make_layout((8, 16, 4)), byte_alignment=16
+        )
+
+        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, 2)
+        tmem_partial_mbar = smem.allocate_array(Int64, 1)
+        tmem_empty_mbar = smem.allocate_array(Int64, 2)
+        taddr = smem.allocate(Int32, 4)
+
+        M, K = A_tma.tma_tensor.shape
+        N, _ = B_tma.tma_tensor.shape
+
+        def num_tiles(M, N, cg):
+            # per-CTA M-tiles (a multiple of cg) x BN-wide N-tiles. Each role
+            # recomputes these after setmaxnreg: a value kept live across it
+            # is spilled.
+            grid_n = cute.ceil_div(N, 128 * cg)
+            grid_mp = cute.ceil_div(M, 128 * cg)
+            return grid_mp, grid_n, grid_mp * cg * grid_n
+
+        def tile_coords(bid, grid_mp, grid_n, cg, G):
+            # CTA group p = bid // cg covers M-tiles cg * m + rank (one per
+            # CTA) and N-tile n. Rasterized in groups of G M-tile groups, m
+            # fastest within each N-tile of a group (G = 1: n fastest).
+            p = bid // cg
+            if cutlass.const_expr(G == 1):
+                return (p // grid_n) * cg + bid % cg, p % grid_n
+            group = p // (G * grid_n)
+            first = group * G
+            size = cutlass.min(grid_mp - first, G)
+            local = p - group * (G * grid_n)
+            return (first + local % size) * cg + bid % cg, local // size
+
+        if warp_id == 0:
+            # counts: tma_full cg CTAs; tmem_empty cg CTAs x 128 * cg epilogue
+            # threads; tmem_partial 2 CTAs x 128 (the overlap warpgroup)
+            with cute.arch.elect_one():
+                for i in cutlass.range_constexpr(num_stages):
+                    cute.arch.mbarrier_init(tma_full_mbar + i, cg)
+                    cute.arch.mbarrier_init(tma_empty_mbar + i, 1)
+                for i in cutlass.range_constexpr(2):
+                    cute.arch.mbarrier_init(tmem_full_mbar + i, 1)
+                    cute.arch.mbarrier_init(tmem_empty_mbar + i, 128 * cg * cg)
+                cute.arch.mbarrier_init(tmem_partial_mbar, 256)
+                cute.arch.mbarrier_init_fence()
+        elif warp_id == 1:
+            cpasync.prefetch_descriptor(A_tma.atom)
+            cpasync.prefetch_descriptor(B_tma.atom)
+            cpasync.prefetch_descriptor(SFA_tma.atom)
+            cpasync.prefetch_descriptor(SFB_tma.atom)
+        elif warp_id < 4:
+            # norm weights (static, so before the PDL wait): 4 x 256B, one 16B
+            # cp.async per thread of warps 2-3
+            t = tid - 64
+            cp_atom = cute.make_copy_atom(
+                cpasync.CopyG2SOp(), BFloat16, num_bits_per_copy=128
+            )
+            for i, Wt in enumerate((WQ_tensor, WK_tensor, WIQ_tensor, WIK_tensor)):
+                if t // 16 == i:
+                    cute.copy(
+                        cp_atom,
+                        cute.zipped_divide(Wt, (8,))[None, t % 16],
+                        sW[None, t % 16, i],
+                    )
+            cute.arch.cp_async_commit_group()
+            cute.arch.cp_async_wait_group(0)
+        cute.arch.sync_threads()  # sW visible (cluster_arrive is relaxed)
+        cute.arch.cluster_arrive_relaxed()
+        cute.arch.cluster_wait()
+
+        if cutlass.const_expr(self.use_pdl):
+            cute.arch.griddepcontrol_wait()
+
+        # warps: 4 * cg epilogue (one warpgroup per head), then MMA, TMA (and 2
+        # padding warps at cg=2, completing the producer warpgroup)
+        MMA_WARP = 4 * cg
+        TMA_WARP = MMA_WARP + 1
+        # setmaxnreg (cg=2): the producer warpgroup drops to 40 regs, the
+        # epilogue rises to 232 (2*128*232 + 128*40 <= 64K)
+        REG_DEC = 40
+        REG_INC = 232
+
+        if warp_id > TMA_WARP:
+            # padding warps
+            if cutlass.const_expr(cg == 2):
+                cute.arch.setmaxregister_decrease(REG_DEC)
+
+        elif warp_id == TMA_WARP:
+            # TMA warp; both CTAs load, copies arrive on CTA0's mbarrier.
+            # SFB is multicast: the pair's MMA needs all 256 rows in each CTA.
+            if cutlass.const_expr(cg == 2):
+                cute.arch.setmaxregister_decrease(REG_DEC)
+            grid_mp, grid_n, total_tiles = num_tiles(M, N, cg)
+            tma_stage = 0
+            tma_empty_parity = 1
+            tma_full_mbar_ = to_cta0_smem(tma_full_mbar)
+
+            gA_tiles = cute.zipped_divide(A_tma.tma_tensor, (BM, BK))
+            gB_tiles = cute.zipped_divide(B_tma.tma_tensor, (BN // cg, BK))
+            # SF boxes: one [32, 4, 4] atom = (64, 1) Int64 (see prepare_SF)
+            gSFA_tiles = cute.zipped_divide(SFA_tma.tma_tensor, (64, 1))
+            gSFB_tiles = cute.zipped_divide(SFB_tma.tma_tensor, (64, 1))
+
+            for bid in range(raw_bid, total_tiles, num_bids):
+                bid_m, bid_n = tile_coords(bid, grid_mp, grid_n, cg, self.group_m)
+
+                for iter_k in cutlass.range(cute.ceil_div(K, BK), unroll=1):
+                    mbar = tma_full_mbar_ + tma_stage
+                    cute.arch.mbarrier_wait(
+                        tma_empty_mbar + tma_stage, tma_empty_parity
+                    )
+                    with cute.arch.elect_one():
+                        mbarrier.arrive_expect_tx(mbar, self.stage_size, "cluster")
+                    simple_tma_copy(
+                        A_tma.atom,
+                        gA_tiles[None, (bid_m, iter_k)],
+                        sA[None, None, tma_stage],
+                        mbar,
+                    )
+                    simple_tma_copy(
+                        B_tma.atom,
+                        gB_tiles[None, (cg * bid_n + cta_rank, iter_k)],
+                        sB[None, None, tma_stage],
+                        mbar,
+                    )
+                    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, cg * bid_n + cta_rank)],
+                        sSFB[None, None, cta_rank, tma_stage],
+                        mbar,
+                        mcast_mask=cta_pair,
+                    )
+
+                    tma_stage = (tma_stage + 1) % num_stages
+                    if tma_stage == 0:
+                        tma_empty_parity ^= 1
+
+            if cutlass.const_expr(self.use_pdl):
+                cute.arch.griddepcontrol_launch_dependents()
+
+        elif warp_id == MMA_WARP:
+            # MMA warp (2-SM collective; only cta_rank 0 issues). tmem is
+            # allocated by epilogue warp 0.
+            if cutlass.const_expr(cg == 2):
+                cute.arch.setmaxregister_decrease(REG_DEC)
+            cute.arch.barrier(barrier_id=BAR_TMEM_ALLOC, number_of_threads=2 * 32)
+
+            grid_mp, grid_n, total_tiles = num_tiles(M, N, cg)
+            tma_stage = 0
+            tma_full_parity = 0
+
+            sdesc = _tcgen05.make_sdesc_128B_swizzle(0)
+            # SF smem descriptor: no swizzle, SBO = 8 * 16 bytes
+            sdesc_sf = Uint64((8 << 32) | (1 << 46))
+            idesc = _tcgen05.make_mxfp8_idesc(cg * BM, BN)
+
+            if cta_rank == 0:
+                tmem_stage = 0
+                tmem_parity = 1
+                for _ in range(raw_bid, total_tiles, num_bids):
+                    acc_base = acc_odd * tmem_stage
+                    cute.arch.mbarrier_wait(tmem_empty_mbar + tmem_stage, tmem_parity)
+                    # the previous tile's overlap-region read (cols [224, 256));
+                    # the single partial barrier phases once per tile
+                    if cutlass.const_expr(cg == 2):
+                        cute.arch.mbarrier_wait(tmem_partial_mbar, tmem_stage ^ 1)
+                    _tcgen05.fence_after_thread_sync()
+
+                    for iter_k in cutlass.range(cute.ceil_div(K, BK), unroll=1):
+                        a_desc = sdesc | (
+                            sA[None, None, tma_stage].iterator.toint() >> 4
+                        )
+                        b_desc = sdesc | (
+                            sB[None, None, tma_stage].iterator.toint() >> 4
+                        )
+                        sfa_desc = sdesc_sf | (
+                            sSFA[None, None, tma_stage].iterator.toint() >> 4
+                        )
+                        sfb0_desc = sdesc_sf | (
+                            sSFB[None, None, 0, tma_stage].iterator.toint() >> 4
+                        )
+
+                        cute.arch.mbarrier_wait(
+                            tma_full_mbar + tma_stage, tma_full_parity
+                        )
+                        _tcgen05.fence_after_thread_sync()
+
+                        _tcgen05.cp(sfa_tmem, sfa_desc, "32x128b", "warpx4", cg)
+                        _tcgen05.cp(sfb0_tmem, sfb0_desc, "32x128b", "warpx4", cg)
+                        if cutlass.const_expr(cg == 2):
+                            sfb1_desc = sdesc_sf | (
+                                sSFB[None, None, 1, tma_stage].iterator.toint() >> 4
+                            )
+                            _tcgen05.cp(sfb1_tmem, sfb1_desc, "32x128b", "warpx4", 2)
+
+                        for k in cutlass.range_constexpr(BK // 32):
+                            enable_input_d = iter_k > 0 or k > 0
+                            # select this k-block's SF byte inside the staged
+                            # 4-col tmem atoms (idesc SF-id fields)
+                            idesc_k = idesc + ((k << 4) | (k << 29))
+                            _tcgen05.mma_mxfp8(
+                                acc_base,
+                                a_desc,
+                                b_desc,
+                                idesc_k,
+                                sfa_tmem,
+                                sfb0_tmem,
+                                enable_input_d,
+                                cg,
+                            )
+                            a_desc += 32 >> 4
+                            b_desc += 32 >> 4
+                        _tcgen05.commit(tma_empty_mbar + tma_stage, cta_pair, cg)
+
+                        tma_stage = (tma_stage + 1) % num_stages
+                        if tma_stage == 0:
+                            tma_full_parity ^= 1
+
+                    _tcgen05.commit(tmem_full_mbar + tmem_stage, cta_pair, cg)
+                    tmem_stage ^= 1
+                    if tmem_stage == 0:
+                        tmem_parity ^= 1
+
+        else:
+            # epilogue warps
+            if cutlass.const_expr(cg == 2):
+                cute.arch.setmaxregister_increase(REG_INC)
+            if warp_id == 0:
+                _tcgen05.alloc(taddr, cta_group=cg)
+                cute.arch.barrier(barrier_id=BAR_TMEM_ALLOC, number_of_threads=2 * 32)
+
+            # warp w reads tmem rows (w % 4) * 32 of head cg * bid_n + w // 4;
+            # thread t owns row t % 128. Padding rows (m >= M) compute garbage
+            # from clamped metadata and skip only the stores.
+            warp_e = warp_id % 4
+            wg_id = warp_id // 4
+            tid_ = tid % 128
+
+            # Q / QFP8 are each optional (None -> compiled out)
+            has_q = cutlass.const_expr(Q_tensor is not None)
+            has_qfp8 = cutlass.const_expr(QFP8_tensor is not None)
+            # [32, (pos, cos | sin)]
+            CS_CHUNKS = cute.zipped_divide(COS_SIN_tensor, (1, 32))[(0, None), None]
+            # KV head h's K / V go to slots 2h / 2h + 1; index_k has one slot
+            KV_DATA, KV_SF, KV_VSF = _nvfp4_cache_views(KVC_tensor, bs, 2 * self.nkv)
+            IDX_DATA, IDX_SF, _ = _nvfp4_cache_views(IDXC_tensor, bs, 1)
+
+            partial_ = to_cta0_smem(tmem_partial_mbar)
+            empty_ = to_cta0_smem(tmem_empty_mbar)
+
+            grid_mp, grid_n, total_tiles = num_tiles(M, N, cg)
+            tmem_stage = 0
+            tmem_parity = 0
+            for bid in range(raw_bid, total_tiles, num_bids):
+                bid_m, bid_n = tile_coords(bid, grid_mp, grid_n, cg, self.group_m)
+                acc_base = acc_odd * tmem_stage
+
+                if warp_id == 0:
+                    cute.arch.mbarrier_wait(tmem_full_mbar + tmem_stage, tmem_parity)
+                cute.arch.barrier(barrier_id=BAR_EPI, number_of_threads=cg * 128)
+                _tcgen05.fence_after_thread_sync()
+
+                head = cg * bid_n + wg_id  # this warpgroup's head
+                m = bid_m * BM + tid_
+
+                # Drain this warp's 32 rows of its head into registers
+                row = warp_e * 32
+                base = acc_base + wg_id * HEAD_DIM
+                acc = cute.make_rmem_tensor((32, 4), Float32)  # 4 x 32 cols
+                owns_overlap = False
+                if cutlass.const_expr(cg == 2):
+                    owns_overlap = wg_id == 1 - tmem_stage
+                if owns_overlap:
+                    # This warpgroup holds cols [224, 256), which the next
+                    # tile's MMA reuses: load and release them first. They are
+                    # chunk 0 of the head on odd tiles, chunk 3 on even.
+                    if tmem_stage == 1:
+                        acc[None, 0].store(_tcgen05.ld(row, base, "32x32b", 32))
+                        _tcgen05.wait_ld()
+                        mbarrier.arrive(partial_, "cluster")
+                        for c in cutlass.range_constexpr(1, 4):
+                            col = base + 32 * c
+                            acc[None, c].store(_tcgen05.ld(row, col, "32x32b", 32))
+                    else:
+                        acc[None, 3].store(_tcgen05.ld(row, base + 96, "32x32b", 32))
+                        _tcgen05.wait_ld()
+                        mbarrier.arrive(partial_, "cluster")
+                        for c in cutlass.range_constexpr(3):
+                            col = base + 32 * c
+                            acc[None, c].store(_tcgen05.ld(row, col, "32x32b", 32))
+                else:
+                    acc.store(_tcgen05.ld(row, base, "32x32b", HEAD_DIM))
+                _tcgen05.wait_ld()
+                _tcgen05.fence_before_thread_sync()
+                mbarrier.arrive(empty_ + tmem_stage, "cluster")
+
+                # rope dim is 64: store the unrotated half (cols 64-127) first
+                # to free its registers before rope. Only the stores are
+                # predicated on the row / slot.
+                halves = cute.make_tensor(acc.iterator, cute.make_layout((64, 2)))
+                m_safe = cutlass.min(m, M - 1)
+
+                if head < self.k_start:
+                    # q: bf16 and / or scaled e4m3
+                    _rms_norm(acc, sW[None, None, W_Q], eps)
+                    if cutlass.const_expr(has_qfp8):
+                        qfp8_inv = cute.math.rcp(q_fp8_scale)
+
+                    for h in (1, 0):  # NoPE, then RoPE
+                        if cutlass.const_expr(h == 0):
+                            _rope64(acc, CS_CHUNKS, POS_tensor[m_safe].to(Int32))
+
+                        x = halves[None, h].load()
+                        col = HEAD_DIM * head + 64 * h
+                        if cutlass.const_expr(has_q):
+                            assert Q_tensor is not None
+                            q = x.to(BFloat16)
+                            if m < M:
+                                dst = Q_tensor[m, None].iterator + col
+                                cute.arch.store(dst, q, **_NO_L1_ALLOC)
+                        # fp8 is stored as u32 words: fp8 vectors don't lower
+                        # in cute.arch.store (NVIDIA/cutlass#3712)
+                        if cutlass.const_expr(has_qfp8):
+                            assert QFP8_tensor is not None
+                            q = (x * qfp8_inv).to(Float8E4M3FN).bitcast(Uint32)
+                            if m < M:
+                                dst = QFP8_tensor[m, None].iterator + col
+                                dst = cute.recast_ptr(dst, dtype=Uint32)
+                                cute.arch.store(dst, q, **_NO_L1_ALLOC)
+
+                elif head < self.v_start:
+                    # k: NVFP4 into its KV head's K slot, linear scale bytes
+                    _rms_norm(acc, sW[None, None, W_K], eps)
+                    kv_slot = 2 * (head - self.k_start)
+                    slot = SLOT_tensor[m_safe]
+                    page, token = slot // bs, slot % bs
+                    k_sf = cute.math.rcp(k_scale)
+                    sf = cute.make_rmem_tensor(8, Uint8)
+
+                    for h in (1, 0):  # NoPE, then RoPE
+                        if cutlass.const_expr(h == 0):
+                            _rope64(acc, CS_CHUNKS, POS_tensor[m_safe].to(Int32))
+
+                        words, sf_h = _nvfp4_quant64(halves[None, h], k_sf)
+                        cute.local_tile(sf, (4,), (h,)).store(sf_h)
+                        if m < M and slot >= 0:
+                            dst = KV_DATA[None, h, token, kv_slot, page].iterator
+                            cute.arch.store(dst, words, **_NO_L1_ALLOC)
+
+                    if m < M and slot >= 0:
+                        dst = KV_SF[None, token, kv_slot, page].iterator
+                        cute.arch.store(dst, sf.load(), **_NO_L1_ALLOC)
+
+                elif head < self.iq_start:
+                    # v: raw accumulator, NVFP4 into its KV head's V slot,
+                    # swizzled scale bytes
+                    kv_slot = 2 * (head - self.v_start) + 1
+                    slot = SLOT_tensor[m_safe]
+                    page, token = slot // bs, slot % bs
+                    v_sf = cute.math.rcp(v_scale)
+
+                    for h in (0, 1):
+                        words, sf = _nvfp4_quant64(halves[None, h], v_sf)
+                        if m < M and slot >= 0:
+                            dst = KV_DATA[None, h, token, kv_slot, page].iterator
+                            cute.arch.store(dst, words, **_NO_L1_ALLOC)
+                            vsf = KV_VSF[None, token, kv_slot, page]
+                            cute.local_tile(vsf, (4,), (h,)).store(sf)
+
+                elif head < self.ik_start:
+                    # index_q: e4m3
+                    _rms_norm(acc, sW[None, None, W_IQ], eps)
+                    iq_col = HEAD_DIM * (head - self.iq_start)
+
+                    for h in (1, 0):  # NoPE, then RoPE
+                        if cutlass.const_expr(h == 0):
+                            _rope64(acc, CS_CHUNKS, POS_tensor[m_safe].to(Int32))
+
+                        # u32 words: see the q branch
+                        q = halves[None, h].load().to(Float8E4M3FN).bitcast(Uint32)
+                        if m < M:
+                            dst = IQFP8_tensor[m, None].iterator + iq_col + 64 * h
+                            dst = cute.recast_ptr(dst, dtype=Uint32)
+                            cute.arch.store(dst, q, **_NO_L1_ALLOC)
+
+                elif head == self.ik_start:
+                    # index_k: NVFP4 into the index cache (one slot per page,
+                    # linear scale bytes)
+                    _rms_norm(acc, sW[None, None, W_IK], eps)
+                    slot = IDX_SLOT_tensor[m_safe]
+                    page, token = slot // bs, slot % bs
+                    sf = cute.make_rmem_tensor(8, Uint8)
+
+                    for h in (1, 0):  # NoPE, then RoPE
+                        if cutlass.const_expr(h == 0):
+                            _rope64(acc, CS_CHUNKS, POS_tensor[m_safe].to(Int32))
+
+                        words, sf_h = _nvfp4_quant64(halves[None, h])
+                        cute.local_tile(sf, (4,), (h,)).store(sf_h)
+                        if m < M and slot >= 0:
+                            dst = IDX_DATA[None, h, token, 0, page].iterator
+                            cute.arch.store(dst, words, **_NO_L1_ALLOC)
+
+                    if m < M and slot >= 0:
+                        dst = IDX_SF[None, token, 0, page].iterator
+                        cute.arch.store(dst, sf.load(), **_NO_L1_ALLOC)
+
+                # else: phantom head (odd head count), nothing to store
+
+                tmem_stage ^= 1
+                if tmem_stage == 0:
+                    tmem_parity ^= 1
+
+            cute.arch.cluster_arrive_relaxed()
+            cute.arch.cluster_wait()
+            if warp_id == 0:
+                _tcgen05.dealloc(cta_group=cg)
+
+
+class _QkvIndexFusedEpilogueJit(
+    VllmCuTeDSLJitKernel["_QkvIndexFusedEpilogueJit.CompileKey"]
+):
+    """Compile keys, warmup (fake) inputs and the compiled executors of
+    QkvIndexFusedEpilogueKernel."""
+
+    @dataclass(frozen=True)
+    class CompileKey:
+        has_q: bool
+        has_qfp8: bool
+        nkv: int
+        block_size: int
+        group_m: int
+        cta_group: int
+        use_pdl: bool = True
+
+    @staticmethod
+    def kernel(key: CompileKey) -> QkvIndexFusedEpilogueKernel:
+        return QkvIndexFusedEpilogueKernel(
+            key.nkv, key.block_size, key.cta_group, key.group_m, key.use_pdl
+        )
+
+    def get_warmup_keys(self, **fields) -> list[CompileKey]:
+        # a layer's keys differ only in the tile config, picked per batch size
+        return [self.CompileKey(cta_group=cg, **fields) for cg in (1, 2)]
+
+    def warmup_inputs(self, key: CompileKey) -> tuple:
+        M = cute.sym_int()
+        N = cute.sym_int()
+        K = cute.sym_int()
+        P = cute.sym_int()
+        slot_bytes = key.block_size * (HEAD_DIM // 2 + HEAD_DIM // 16)
+
+        A = make_fake_tensor(
+            Float8E4M3FN, (M, K), (cute.sym_int64(divisibility=16), 1), assumed_align=16
+        )
+        B = make_fake_tensor(
+            Float8E4M3FN, (N, K), (cute.sym_int64(divisibility=16), 1), assumed_align=16
+        )
+        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
+        )
+        q_cols = 16 * key.nkv * HEAD_DIM
+        Q = (
+            make_fake_tensor(BFloat16, (M, q_cols), (q_cols, 1), assumed_align=32)
+            if key.has_q
+            else None
+        )
+        WQ, WK, WIQ, WIK = (
+            make_fake_tensor(BFloat16, (HEAD_DIM,), (1,), assumed_align=16)
+            for _ in range(4)
+        )
+        COS_SIN = make_fake_tensor(BFloat16, (P, 64), (64, 1), assumed_align=16)
+        POS = make_fake_tensor(Int64, (M,), (1,), assumed_align=16)
+        # caches as [nb, page_bytes] with a dynamic page stride: no dynamic dim
+        # may exceed int32, and cross-layer layouts space pages wider than a page
+        KVC = make_fake_tensor(
+            Uint8,
+            (cute.sym_int(), 2 * key.nkv * slot_bytes),
+            (cute.sym_int64(divisibility=32), 1),
+            assumed_align=32,
+        )
+        IDXC = make_fake_tensor(
+            Uint8,
+            (cute.sym_int(), slot_bytes),
+            (cute.sym_int64(divisibility=32), 1),
+            assumed_align=32,
+        )
+        SLOT = make_fake_tensor(Int64, (M,), (1,), assumed_align=16)
+        IDX_SLOT = make_fake_tensor(Int64, (M,), (1,), assumed_align=16)
+        IQFP8 = make_fake_tensor(
+            Float8E4M3FN,
+            (M, key.nkv * HEAD_DIM),
+            (key.nkv * HEAD_DIM, 1),
+            assumed_align=32,
+        )
+        QFP8 = (
+            make_fake_tensor(Float8E4M3FN, (M, q_cols), (q_cols, 1), assumed_align=32)
+            if key.has_qfp8
+            else None
+        )
+
+        return (
+            A,
+            B,
+            SFA,
+            SFB,
+            Q,
+            WQ,
+            WK,
+            WIQ,
+            WIK,
+            COS_SIN,
+            POS,
+            KVC,
+            IDXC,
+            SLOT,
+            IDX_SLOT,
+            IQFP8,
+            QFP8,
+            # runtime scalars (eps, q/k/v scales, grid): placeholder values
+            Float32(1e-6),
+            Float32(1.0),
+            Float32(1.0),
+            Float32(1.0),
+            148,
+        )
+
+
+_JIT = _QkvIndexFusedEpilogueJit()
+
+
+@cache
+def _device_info(device: int) -> tuple[int, int]:
+    """(SM count, L2 bytes) of a CUDA device."""
+    props = torch.cuda.get_device_properties(device)
+    return props.multi_processor_count, props.L2_cache_size
+
+
+def _group_m(weight_bytes: int, l2_bytes: int) -> int:
+    # Group M-tiles 16 at a time when the weight is a large fraction of L2, so
+    # a wave reads a few N-panels instead of all of B (GB300: helps TP1's 60MB
+    # weight by 1.5-5%; slows TP2 / TP4, whose weights stay L2-resident anyway)
+    return 16 if weight_bytes >= l2_bytes // 4 else 1
+
+
+def register_qkv_index_fused_epilogue_warmup(
+    *, has_q: bool, has_qfp8: bool, nkv: int, block_size: int, weight_bytes: int
+) -> None:
+    """Compile this layer's kernels (both tile configs) during JIT warmup."""
+    _, l2_bytes = _device_info(torch.accelerator.current_device_index())
+    _JIT.register_warmup(
+        has_q=has_q,
+        has_qfp8=has_qfp8,
+        nkv=nkv,
+        block_size=block_size,
+        group_m=_group_m(weight_bytes, l2_bytes),
+    )
+
+
+def minimax_m3_qkv_index_fused_epilogue(
+    xq: torch.Tensor,
+    xsf: torch.Tensor,
+    wq: torch.Tensor,
+    wsf: torch.Tensor,
+    positions: torch.Tensor,
+    cos_sin_cache: torch.Tensor,
+    q_norm_weight: torch.Tensor,
+    k_norm_weight: torch.Tensor,
+    index_q_norm_weight: torch.Tensor,
+    index_k_norm_weight: torch.Tensor,
+    eps: float,
+    *,
+    slot_mapping: torch.Tensor,
+    index_slot_mapping: torch.Tensor,
+    kv_cache: torch.Tensor,
+    index_cache: torch.Tensor,
+    index_q_out: torch.Tensor,
+    q_out: torch.Tensor | None = None,
+    q_fp8_out: torch.Tensor | None = None,
+    q_fp8_scale: float = 1.0,
+    k_scale: float = 1.0,
+    v_scale: float = 1.0,
+    cta_group: int | None = None,
+    use_pdl: bool = True,
+    grid: int | None = None,
+) -> None:
+    """QKV+indexer projection with fused epilogue: writes q (bf16 and / or
+    fp8; at least one) and index_q, and NVFP4-inserts k / v / index_k into the
+    paged caches.
+    nkv in {1, 2, 4} KV heads per rank is inferred from N = (19 * nkv + 1) * 128.
+
+    Args:
+        xq: MXFP8 activation [M, K].
+        xsf: xq's F8_128x4-swizzled E8M0 scales.
+        wq: K-major MXFP8 weight [N, K]: [q | k | v | index_q | index_k].
+        wsf: wq's F8_128x4-swizzled E8M0 scales.
+        positions: [M] int64 token positions.
+        cos_sin_cache: [max_pos, 64] bf16 RoPE cache (cos | sin).
+        q_norm_weight: [128] bf16 Gemma RMSNorm weight; likewise the next three.
+        k_norm_weight: see q_norm_weight.
+        index_q_norm_weight: see q_norm_weight.
+        index_k_norm_weight: see q_norm_weight.
+        eps: RMSNorm epsilon.
+        slot_mapping: [M] int64 KV-cache slots (-1 = skip).
+        index_slot_mapping: [M] int64 index-cache slots (-1 = skip).
+        kv_cache: NVFP4 [nb, 2 * nkv, block_size, 72] uint8; packed pages at
+            any page stride. block_size % 4 == 0.
+        index_cache: NVFP4 [nb, block_size, 72] uint8; packed pages at any
+            page stride.
+        index_q_out: [M, nkv * 128] e4m3 normed + roped index_q.
+        q_out: optional [M, 16 * nkv * 128] bf16 normed + roped q.
+        q_fp8_out: optional [M, 16 * nkv * 128] e4m3 normed + roped q.
+        q_fp8_scale: dequantization scale of q_fp8_out.
+        k_scale: dequantization (global) scale of the NVFP4 k entries.
+        v_scale: dequantization (global) scale of the NVFP4 v entries.
+        cta_group: 2 (2-SM: CTA pairs, 128 x 256 tiles), 1 (1-SM: 128 x 128
+            tiles) or None (pick by wave count).
+        use_pdl: launch with programmatic dependent launch.
+        grid: persistent grid size override (testing).
+
+    """
+    M, K = xq.shape
+    N = wq.shape[0]
+    nkv = (N // HEAD_DIM - 1) // 19
+    assert nkv in (1, 2, 4) and (19 * nkv + 1) * HEAD_DIM == N, N
+    assert K % 128 == 0 and wq.shape[1] == K and wq.stride() == (K, 1)
+    assert index_q_out.shape == (M, nkv * HEAD_DIM)
+    assert index_q_out.dtype == torch.float8_e4m3fn and index_q_out.is_contiguous()
+    assert q_out is not None or q_fp8_out is not None, "no q output"
+    for t in (positions, slot_mapping, index_slot_mapping):
+        assert t.shape == (M,) and t.dtype == torch.int64
+    assert cos_sin_cache.dtype == torch.bfloat16 and cos_sin_cache.is_contiguous()
+    assert cos_sin_cache.dim() == 2 and cos_sin_cache.shape[1] == 64
+    for w in (q_norm_weight, k_norm_weight, index_q_norm_weight, index_k_norm_weight):
+        assert w.shape == (HEAD_DIM,) and w.dtype == torch.bfloat16
+    for q, dtype in ((q_out, torch.bfloat16), (q_fp8_out, torch.float8_e4m3fn)):
+        if q is not None:
+            assert q.shape == (M, 16 * nkv * HEAD_DIM)
+            assert q.dtype == dtype and q.is_contiguous()
+    assert kv_cache.dtype == torch.uint8 and index_cache.dtype == torch.uint8
+    block_size, token_bytes = kv_cache.shape[2:]
+    assert token_bytes == HEAD_DIM // 2 + HEAD_DIM // 16
+    assert kv_cache.shape[1] == 2 * nkv
+    assert kv_cache.stride()[1:] == (block_size * token_bytes, token_bytes, 1)
+    assert index_cache.shape[1:] == (block_size, token_bytes)
+    assert index_cache.stride()[1:] == (token_bytes, 1)
+    for c in (kv_cache, index_cache):  # 32B NVFP4 data stores
+        assert c.data_ptr() % 32 == 0 and c.stride(0) % 32 == 0
+    num_sms, l2_bytes = _device_info(xq.device.index)
+
+    def num_tiles(cg):
+        tile = 128 * cg  # M rows per CTA group, N cols per tile
+        return -(-M // tile) * cg * -(-N // tile)
+
+    if cta_group is None:
+        # a 128 x 128 tile costs ~2/3 of a 128 x 256 one (GB300): take 1-SM
+        # when that wins after wave quantization (odd M-tile counts, few tiles)
+        waves1, waves2 = (-(-num_tiles(cg) // num_sms) for cg in (1, 2))
+        cta_group = 1 if 2 * waves1 < 3 * waves2 else 2
+    if grid is None:
+        grid = min(num_sms, num_tiles(cta_group))
+    # caches as [nb, page_bytes], keeping the page stride (a flat 1-D view
+    # exceeds the DSL's int32 shapes at production sizes)
+    slot_bytes = block_size * token_bytes
+    kvc = kv_cache.as_strided(
+        (kv_cache.size(0), 2 * nkv * slot_bytes), (kv_cache.stride(0), 1)
+    )
+    idxc = index_cache.as_strided(
+        (index_cache.size(0), slot_bytes), (index_cache.stride(0), 1)
+    )
+    key = _JIT.CompileKey(
+        has_q=q_out is not None,
+        has_qfp8=q_fp8_out is not None,
+        nkv=nkv,
+        block_size=block_size,
+        group_m=_group_m(wq.numel(), l2_bytes),
+        cta_group=cta_group,
+        use_pdl=use_pdl,
+    )
+    _JIT._get_or_compile(key)(
+        xq,
+        wq,
+        xsf.view(-1).view(torch.float8_e8m0fnu),
+        wsf.view(-1).view(torch.float8_e8m0fnu),
+        q_out,
+        q_norm_weight,
+        k_norm_weight,
+        index_q_norm_weight,
+        index_k_norm_weight,
+        cos_sin_cache,
+        positions,
+        kvc,
+        idxc,
+        slot_mapping,
+        index_slot_mapping,
+        index_q_out,
+        q_fp8_out,
+        eps,
+        q_fp8_scale,
+        k_scale,
+        v_scale,
+        grid,
+    )
diff --git a/vllm/models/minimax_m3/nvidia/sparse_attention_msa.py b/vllm/models/minimax_m3/nvidia/sparse_attention_msa.py
index a5598c719bfa..7443b449f776 100644
--- a/vllm/models/minimax_m3/nvidia/sparse_attention_msa.py
+++ b/vllm/models/minimax_m3/nvidia/sparse_attention_msa.py
@@ -1,16 +1,21 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""MSA (SM100/Blackwell) block-sparse attention for MiniMax M3.
+"""MSA (SM100 family) block-sparse attention for MiniMax M3.
 
-Prefill attends with ``fmha_sm100`` (``build_k2q_csr`` + ``sparse_atten_func``).
-Decode uses Triton split-K by default, with an opt-in CUTLASS ``fmha_sm100``
-path for regular decode and speculative verification. NVFP4 KV caches read the
-vLLM packed pages directly in both prefill and CUTLASS decode.
+``MiniMaxM3SparseMSAImpl`` serves BF16 and FP8 KV caches: prefill attends with
+``fmha_sm100`` (``build_k2q_csr`` + ``sparse_atten_func``), and decode uses the
+CUTLASS ``fmha_sm100`` path by default for regular decode and speculative
+verification, with Triton split-K as the fallback.
+
+``MiniMaxM3SparseMSANvfp4Impl`` serves NVFP4 KV caches, reading vLLM's packed
+pages directly with only the E4M3 query: decode and prefill both run the CUTLASS
+Q8KV4 kernels.
 """
 
+import functools
 import os
-from dataclasses import dataclass, replace
-from typing import ClassVar
+from dataclasses import dataclass, fields, replace
+from typing import TYPE_CHECKING, ClassVar
 
 import torch
 
@@ -29,8 +34,10 @@
     MiniMaxM3SparseImpl,
     MiniMaxM3SparseMetadata,
     MiniMaxM3SparseMetadataBuilder,
+    MiniMaxM3SparsePrefillMetadata,
 )
 from vllm.models.minimax_m3.nvidia.msa_cutlass_sparse_decode import (
+    NVFP4_BLOCK_SCALE_SHIFT,
     MSACutlassDecodeMetadata,
     MSACutlassDecodePlanCache,
     is_nvfp4_kv_cache,
@@ -39,7 +46,9 @@
     prepare_decode_metadata,
     should_prepare_decode_metadata,
     supports_cutlass_sparse_decode,
+    uses_q8kv4_decode,
 )
+from vllm.platforms import current_platform
 from vllm.utils.torch_utils import get_dtype_size, nvfp4_kv_cache_full_dim
 from vllm.v1.attention.backend import (
     AttentionLayer,
@@ -47,19 +56,40 @@
 )
 from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheLayout
 
+if TYPE_CHECKING:
+    from vllm.third_party.fmha_sm100.prefill_q8kv4 import PrefillWorkspace
+
 logger = init_logger(__name__)
 
 
+@functools.cache
+def _q8kv4_prefill_blocker(
+    num_q_heads: int, num_kv_heads: int, page_size: int, topk_blocks: int
+) -> str | None:
+    """Why the Q8KV4 kernel cannot prefill this layout on this device, or None."""
+    from vllm.third_party.fmha_sm100.q8kv4_prefill_adapter import _plan_blocker
+
+    return _plan_blocker(
+        num_qo_heads=num_q_heads,
+        num_kv_heads=num_kv_heads,
+        page_size=page_size,
+        kv_block_num=topk_blocks,
+        causal=True,
+        output_maxscore=False,
+        device=None,
+    )
+
+
 def _get_msa_cutlass_min_batch_size() -> int | None:
     value = os.getenv("VLLM_MINIMAX_M3_MSA_CUTLASS_MIN_BATCH")
     return int(value) if value is not None else None
 
 
 def _dequantize_query(
-    query_fp8: torch.Tensor, q_scale: float, out: torch.Tensor
+    query_fp8: torch.Tensor, q_scale: float, dtype: torch.dtype
 ) -> torch.Tensor:
-    """Write ``query_fp8 * q_scale`` into ``out`` (E4M3 is exact in BF16)."""
-    out.copy_(query_fp8)
+    """``query_fp8 * q_scale`` in ``dtype`` (E4M3 is exact in BF16)."""
+    out = query_fp8.to(dtype)
     if q_scale != 1.0:
         out.mul_(q_scale)
     return out
@@ -118,6 +148,12 @@ class MiniMaxM3SparseMSADecodeMetadata(MiniMaxM3SparseDecodeMetadata):
     msa_cutlass: MSACutlassDecodeMetadata | None = None
 
 
+@dataclass
+class MiniMaxM3SparseMSAPrefillMetadata(MiniMaxM3SparsePrefillMetadata):
+    # The builder's Q8KV4 prefill split workspace; None off Rubin.
+    workspace: "PrefillWorkspace | None" = None
+
+
 class MiniMaxM3SparseMSAMetadataBuilder(MiniMaxM3SparseMetadataBuilder):
     """Prepare MSA plans only for decode shapes supported by ``fmha_sm100``."""
 
@@ -138,8 +174,29 @@ def __init__(
         # format to distinguish E4M3 (supported) from E5M2 before planning.
         self.kv_cache_dtype = vllm_config.cache_config.cache_dtype
         self.decode_backend = vllm_config.attention_config.minimax_m3_msa_decode_backend
+        self.q8kv4_decode = uses_q8kv4_decode(
+            kv_cache_dtype=self.kv_cache_dtype,
+            num_q_heads=self.num_q_heads,
+            num_kv_heads=self.num_kv_heads,
+        )
         self.msa_cutlass_plan_cache = MSACutlassDecodePlanCache()
         self.msa_cutlass_min_batch_size = _get_msa_cutlass_min_batch_size()
+        # Rubin only: Q8KV4 prefill split partials striped over the memory
+        # locality domains, reserved for the largest batch.
+        self.prefill_workspace: "PrefillWorkspace | None" = None  # noqa: UP037
+        if is_nvfp4_kv_cache(
+            self.kv_cache_dtype
+        ) and current_platform.is_device_capability(107):
+            from vllm.third_party.fmha_sm100.prefill_q8kv4 import PrefillWorkspace
+
+            self.prefill_workspace = PrefillWorkspace(
+                device, block_scale_shift=NVFP4_BLOCK_SCALE_SHIFT
+            )
+            self.prefill_workspace.split_buffers(
+                self.topk_blocks,
+                vllm_config.scheduler_config.max_num_batched_tokens,
+                self.num_q_heads,
+            )
 
     def build(
         self,
@@ -152,6 +209,12 @@ def build(
             common_attn_metadata,
             fast_build,
         )
+        prefill = metadata.prefill
+        if prefill is not None and self.prefill_workspace is not None:
+            metadata.prefill = MiniMaxM3SparseMSAPrefillMetadata(
+                **{f.name: getattr(prefill, f.name) for f in fields(prefill)},
+                workspace=self.prefill_workspace,
+            )
         decode = metadata.decode
         if decode is None:
             return metadata
@@ -179,6 +242,7 @@ def build(
                 num_kv_heads=self.num_kv_heads,
                 page_size=SPARSE_BLOCK_SIZE,
                 topk_blocks=self.topk_blocks,
+                q8kv4=self.q8kv4_decode,
                 plan_cache=self.msa_cutlass_plan_cache,
             )
         metadata.decode = MiniMaxM3SparseMSADecodeMetadata(
@@ -190,8 +254,58 @@ def build(
         return metadata
 
 
-class MiniMaxM3SparseMSAImpl(MiniMaxM3SparseImpl):
-    """MSA block-sparse attention with guarded CUTLASS sparse decode."""
+class MiniMaxM3SparseMSABaseImpl(MiniMaxM3SparseImpl):
+    """What the MSA impls share: CUTLASS decode plans and the prefill CSR."""
+
+    use_cutlass_decode: bool
+    # The prefill kernels ``fmha_sm100.warmup`` prebuilds for this impl.
+    prefill_backend: str
+
+    def should_use_msa_decode(self, layer_name: str) -> bool:
+        if not self.use_cutlass_decode:
+            return False
+        attn_metadata = get_forward_context().attn_metadata
+        if not isinstance(attn_metadata, dict):
+            return False
+        main_md = attn_metadata[layer_name]
+        if not isinstance(main_md, MiniMaxM3SparseMetadata):
+            return False
+        decode = main_md.decode
+        return (
+            isinstance(decode, MiniMaxM3SparseMSADecodeMetadata)
+            and decode.msa_cutlass is not None
+        )
+
+    def _build_prefill_csr(
+        self,
+        topk: torch.Tensor,
+        nd: int,
+        num_tokens: int,
+        p: MiniMaxM3SparsePrefillMetadata,
+    ):
+        from vllm.third_party.fmha_sm100.sparse import build_k2q_csr
+
+        # [H, prefill, MK] transposed view; build_k2q_csr consumes the strided
+        # view directly (topK stays innermost-contiguous).
+        return build_k2q_csr(
+            topk[nd:num_tokens].transpose(0, 1),
+            p.cu_seqlens_q,
+            p.cu_seqlens_k,
+            SPARSE_BLOCK_SIZE,
+            total_k=0,
+            max_seqlen_k=p.max_seq_len,
+            max_seqlen_q=p.max_query_len,
+            total_rows=p.total_kv_blocks,
+            qhead_per_kv=self.num_heads // self.num_kv_heads,
+            return_schedule=True,
+        )
+
+
+class MiniMaxM3SparseMSAImpl(MiniMaxM3SparseMSABaseImpl):
+    """MSA block-sparse attention over a BF16 or FP8 KV cache, with guarded
+    CUTLASS sparse decode."""
+
+    prefill_backend = "cute_dsl"
 
     def __init__(
         self,
@@ -203,7 +317,7 @@ def __init__(
         *,
         topk_blocks: int,
         sparse_block_size: int,
-        msa_decode_backend: MiniMaxM3MSADecodeBackend = "triton",
+        msa_decode_backend: MiniMaxM3MSADecodeBackend = "cutlass",
     ) -> None:
         super().__init__(
             num_heads,
@@ -214,7 +328,6 @@ def __init__(
             topk_blocks=topk_blocks,
             sparse_block_size=sparse_block_size,
         )
-        self.use_nvfp4_kv = is_nvfp4_kv_cache(kv_cache_dtype)
         self.use_cutlass_decode = supports_cutlass_sparse_decode(
             decode_backend=msa_decode_backend,
             num_q_heads=self.num_heads,
@@ -223,31 +336,18 @@ def __init__(
             page_size=self.block_size,
             topk_blocks=self.topk_blocks,
         )
-        if self.use_nvfp4_kv and not self.use_cutlass_decode:
-            raise ValueError(
-                "MiniMax M3 NVFP4 KV cache requires CUTLASS MSA sparse decode "
-                f"(num_heads={self.num_heads}, num_kv_heads={self.num_kv_heads}, "
-                f"block_size={self.block_size}, topk_blocks={self.topk_blocks})"
-            )
         logger.info_once(
             "MiniMax M3 MSA sparse decode selected %s",
             "CUTLASS" if self.use_cutlass_decode else "Triton",
         )
 
-    def should_use_msa_decode(self, layer_name: str) -> bool:
+    def _bf16_query(
+        self, q: torch.Tensor, q_scale: float, dtype: torch.dtype
+    ) -> torch.Tensor:
+        # With CUTLASS decode the query is E4M3; the other kernels read BF16.
         if not self.use_cutlass_decode:
-            return False
-        attn_metadata = get_forward_context().attn_metadata
-        if not isinstance(attn_metadata, dict):
-            return False
-        main_md = attn_metadata[layer_name]
-        if not isinstance(main_md, MiniMaxM3SparseMetadata):
-            return False
-        decode = main_md.decode
-        return (
-            isinstance(decode, MiniMaxM3SparseMSADecodeMetadata)
-            and decode.msa_cutlass is not None
-        )
+            return q
+        return _dequantize_query(q, q_scale, dtype)
 
     def forward(
         self,
@@ -255,8 +355,6 @@ def forward(
         query: torch.Tensor,
         kv_cache: torch.Tensor,
         output: torch.Tensor,
-        *,
-        query_fp8: torch.Tensor | None = None,
     ) -> torch.Tensor:
         attn_metadata = get_forward_context().attn_metadata
         if not isinstance(attn_metadata, dict):
@@ -273,16 +371,8 @@ def forward(
         hd = self.head_size
         q = query[:num_tokens].view(-1, self.num_heads, hd)
         out = output[:num_tokens].view(-1, self.num_heads, hd)
-        # Given query_fp8, it is the only q the fused insert wrote: CUTLASS
-        # decode and prefill read it directly, and Triton decode reads it
-        # dequantized into ``query``.
-        q_fp8 = (
-            None
-            if query_fp8 is None
-            else query_fp8[:num_tokens].view(-1, self.num_heads, hd)
-        )
         q_scale_float = getattr(layer, "_q_scale_float", 1.0)
-        if self.use_fp8_kv and not self.use_nvfp4_kv:
+        if self.use_fp8_kv:
             kv_cache = kv_cache.view(self.kv_cache_fp8_dtype)
         k_scale = getattr(layer, "_k_scale", None) if self.use_fp8_kv else None
         v_scale = getattr(layer, "_v_scale", None) if self.use_fp8_kv else None
@@ -290,16 +380,11 @@ def forward(
         # Decode [:nd]: CUTLASS for planned shapes, otherwise Triton.
         if main_md.num_decodes > 0:
             d = main_md.decode
-            assert d is not None
-            msa_metadata = (
-                d.msa_cutlass
-                if isinstance(d, MiniMaxM3SparseMSADecodeMetadata)
-                else None
-            )
+            assert isinstance(d, MiniMaxM3SparseMSADecodeMetadata)
+            msa_metadata = d.msa_cutlass
             if self.use_cutlass_decode and msa_metadata is not None:
-                assert q_fp8 is not None
                 msa_cutlass_sparse_decode(
-                    q_fp8[:nd],
+                    q[:nd],
                     kv_cache,
                     topk[:nd],
                     out[:nd],
@@ -308,20 +393,10 @@ def forward(
                     q_scale_float=q_scale_float,
                     k_scale_float=getattr(layer, "_k_scale_float", 1.0),
                     v_scale_float=getattr(layer, "_v_scale_float", 1.0),
-                    k_scale=k_scale if self.use_nvfp4_kv else None,
-                    v_scale=v_scale if self.use_nvfp4_kv else None,
-                )
-            elif self.use_nvfp4_kv:
-                raise RuntimeError(
-                    "MiniMax M3 NVFP4 KV cache has no Triton decode fallback; "
-                    f"CUTLASS decode was not planned for {main_md.num_decodes} "
-                    f"requests with query length {d.decode_query_len}"
                 )
             else:
                 minimax_m3_sparse_attn_decode(
-                    q[:nd]
-                    if q_fp8 is None
-                    else _dequantize_query(q_fp8[:nd], q_scale_float, q[:nd]),
+                    self._bf16_query(q[:nd], q_scale_float, output.dtype),
                     kv_cache,
                     topk[:nd].transpose(0, 1),
                     d.block_table,
@@ -336,61 +411,18 @@ def forward(
 
         # Prefill [nd:]: MSA sparse FMHA over the selected blocks.
         if main_md.num_prefills > 0:
-            from vllm.third_party.fmha_sm100.sparse import (
-                build_k2q_csr,
-                sparse_atten_func,
-                sparse_atten_nvfp4_kv_func,
-            )
+            from vllm.third_party.fmha_sm100.sparse import sparse_atten_func
 
             p = main_md.prefill
             assert p is not None
-            # [H, prefill, MK] transposed view; build_k2q_csr consumes the
-            # strided view directly (topK stays innermost-contiguous).
-            prefill_topk = topk[nd:num_tokens].transpose(0, 1)
+            k2q_row_ptr, k2q_q_indices, schedule = self._build_prefill_csr(
+                topk, nd, num_tokens, p
+            )
             # An E4M3 query attends as is, its scale folded into softmax_scale.
-            if q_fp8 is None:
-                qp, softmax_scale = q[nd:], self.scale
-            else:
-                qp, softmax_scale = q_fp8[nd:], self.scale * q_scale_float
-            k2q_row_ptr, k2q_q_indices, schedule = build_k2q_csr(
-                prefill_topk,
-                p.cu_seqlens_q,
-                p.cu_seqlens_k,
-                SPARSE_BLOCK_SIZE,
-                total_k=0,
-                max_seqlen_k=p.max_seq_len,
-                max_seqlen_q=p.max_query_len,
-                total_rows=p.total_kv_blocks,
-                qhead_per_kv=qp.shape[1] // self.num_kv_heads,
-                return_schedule=True,
+            qp = q[nd:]
+            softmax_scale = (
+                self.scale * q_scale_float if self.use_cutlass_decode else self.scale
             )
-            if self.use_nvfp4_kv:
-                k_data, k_sf, v_data, v_sf = nvfp4_kv_cache_views(kv_cache)
-                sparse_atten_nvfp4_kv_func(
-                    qp,
-                    k_data,
-                    v_data,
-                    k_sf,
-                    v_sf,
-                    k_scale,
-                    v_scale,
-                    k2q_row_ptr,
-                    k2q_q_indices,
-                    self.topk_blocks,
-                    cu_seqlens_q=p.cu_seqlens_q,
-                    cu_seqlens_k=p.cu_seqlens_k,
-                    max_seqlen_q=p.max_query_len,
-                    max_seqlen_k=p.max_seq_len,
-                    blk_kv=SPARSE_BLOCK_SIZE,
-                    causal=True,
-                    softmax_scale=softmax_scale,
-                    page_table=p.block_table,
-                    seqused_k=p.seq_lens,
-                    schedule=schedule,
-                    out=out[nd:],
-                    kv_layout="vllm",
-                )
-                return output
             k_cache, v_cache = kv_cache.split(self.head_size, dim=-1)
             sparse_atten_func(
                 qp,
@@ -412,3 +444,158 @@ def forward(
                 out=out[nd:],
             )
         return output
+
+
+class MiniMaxM3SparseMSANvfp4Impl(MiniMaxM3SparseMSABaseImpl):
+    """MSA block-sparse attention over an NVFP4 KV cache, SM100 family only.
+
+    Decode and prefill run the CUTLASS Q8KV4 kernels, which read the E4M3
+    query; there is no BF16-query kernel and no Triton fallback.
+    """
+
+    use_cutlass_decode = True
+    prefill_backend = "q8kv4"
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        scale: float,
+        num_kv_heads: int | None = None,
+        kv_cache_dtype: str = "nvfp4",
+        *,
+        topk_blocks: int,
+        sparse_block_size: int,
+        msa_decode_backend: MiniMaxM3MSADecodeBackend = "cutlass",
+    ) -> None:
+        super().__init__(
+            num_heads,
+            head_size,
+            scale,
+            num_kv_heads,
+            kv_cache_dtype,
+            topk_blocks=topk_blocks,
+            sparse_block_size=sparse_block_size,
+        )
+        if not (
+            current_platform.is_cuda()
+            and current_platform.is_device_capability_family(100)
+        ):
+            raise ValueError("MiniMax M3 NVFP4 KV cache requires an SM100-family GPU")
+        # NVFP4 decodes only on CUTLASS, whatever msa_decode_backend says.
+        if not supports_cutlass_sparse_decode(
+            decode_backend="cutlass",
+            num_q_heads=self.num_heads,
+            num_kv_heads=self.num_kv_heads,
+            kv_cache_dtype=self.kv_cache_dtype,
+            page_size=self.block_size,
+            topk_blocks=self.topk_blocks,
+        ):
+            raise ValueError(
+                "MiniMax M3 NVFP4 KV cache requires CUTLASS MSA sparse decode "
+                f"(num_heads={self.num_heads}, num_kv_heads={self.num_kv_heads}, "
+                f"block_size={self.block_size}, topk_blocks={self.topk_blocks})"
+            )
+        blocker = _q8kv4_prefill_blocker(
+            self.num_heads, self.num_kv_heads, self.block_size, self.topk_blocks
+        )
+        if blocker is not None:
+            raise ValueError(
+                "MiniMax M3 NVFP4 KV cache requires the CUTLASS Q8KV4 prefill, "
+                f"which cannot serve this layout or device: {blocker}"
+            )
+        logger.info_once("MiniMax M3 NVFP4 MSA: CUTLASS Q8KV4 decode and prefill")
+        # NVFP4_BLOCK_SCALE_SHIFT assumes unit global scales; checked on the
+        # first forward, once the checkpoint's KV scales are loaded.
+        self._check_global_scales = True
+
+    def forward(
+        self,
+        layer: AttentionLayer,
+        query: torch.Tensor,
+        kv_cache: torch.Tensor,
+        output: torch.Tensor,
+    ) -> torch.Tensor:
+        assert query.dtype == torch.float8_e4m3fn
+        attn_metadata = get_forward_context().attn_metadata
+        if not isinstance(attn_metadata, dict):
+            return output  # profiling run; caches unbound
+        main_md = attn_metadata[layer.layer_name]  # type: ignore[attr-defined]
+        assert isinstance(main_md, MiniMaxM3SparseMetadata)
+
+        nd = main_md.num_decode_tokens
+        num_tokens = main_md.num_actual_tokens
+        topk = layer.topk_indices_buffer  # type: ignore[attr-defined]
+        assert topk is not None
+        hd = self.head_size
+        q_fp8 = query[:num_tokens].view(-1, self.num_heads, hd)
+        out = output[:num_tokens].view(-1, self.num_heads, hd)
+        q_scale_float = getattr(layer, "_q_scale_float", 1.0)
+        k_scale = getattr(layer, "_k_scale", None)
+        v_scale = getattr(layer, "_v_scale", None)
+        if self._check_global_scales:
+            k_global = getattr(layer, "_k_scale_float", 1.0)
+            v_global = getattr(layer, "_v_scale_float", 1.0)
+            if k_global != 1.0 or v_global != 1.0:
+                raise ValueError(
+                    "MiniMax M3 reads its NVFP4 KV cache on the Q8KV4 kernels with "
+                    f"block_scale_shift={NVFP4_BLOCK_SCALE_SHIFT}, which needs unit "
+                    f"global scales; got k_scale={k_global}, v_scale={v_global}"
+                )
+            self._check_global_scales = False
+
+        if main_md.num_decodes > 0:
+            d = main_md.decode
+            assert isinstance(d, MiniMaxM3SparseMSADecodeMetadata)
+            msa_metadata = d.msa_cutlass
+            if msa_metadata is None:
+                raise RuntimeError(
+                    "MiniMax M3 NVFP4 KV cache has no Triton decode fallback; "
+                    f"CUTLASS decode was not planned for {main_md.num_decodes} "
+                    f"requests with query length {d.decode_query_len}"
+                )
+            msa_cutlass_sparse_decode(
+                q_fp8[:nd],
+                kv_cache,
+                topk[:nd],
+                out[:nd],
+                msa_metadata,
+                scale=self.scale,
+                q_scale_float=q_scale_float,
+                k_scale_float=getattr(layer, "_k_scale_float", 1.0),
+                v_scale_float=getattr(layer, "_v_scale_float", 1.0),
+                k_scale=k_scale,
+                v_scale=v_scale,
+            )
+
+        if main_md.num_prefills > 0:
+            from vllm.third_party.fmha_sm100.prefill_q8kv4 import run_prefill
+
+            p = main_md.prefill
+            assert p is not None
+            k2q_row_ptr, _, schedule = self._build_prefill_csr(topk, nd, num_tokens, p)
+            k_data, k_sf, v_data, v_sf = nvfp4_kv_cache_views(kv_cache)
+            run_prefill(
+                q_fp8[nd:],
+                (k_data, v_data),
+                (k_sf, v_sf),
+                kv_indices=p.block_table,
+                kv_indptr=None,
+                cu_seqlens_q=p.cu_seqlens_q,
+                cu_seqlens_k=p.cu_seqlens_k,
+                k2q_row_ptr=k2q_row_ptr,
+                schedule=schedule,
+                sm_scale=self.scale * q_scale_float,
+                topk=self.topk_blocks,
+                seqused_k=p.seq_lens,
+                k_global_scale=k_scale,
+                v_global_scale=v_scale,
+                block_scale_shift=NVFP4_BLOCK_SCALE_SHIFT,
+                out=out[nd:],
+                workspace=(
+                    p.workspace
+                    if isinstance(p, MiniMaxM3SparseMSAPrefillMetadata)
+                    else None
+                ),
+            )
+        return output
diff --git a/vllm/models/muse_spark/amd/model.py b/vllm/models/muse_spark/amd/model.py
index 0da05872b59b..79eedffba256 100644
--- a/vllm/models/muse_spark/amd/model.py
+++ b/vllm/models/muse_spark/amd/model.py
@@ -69,6 +69,7 @@
     MuseSparkTextConfig,
 )
 
+from .mono import MuseSparkMono, maybe_attach_mono
 from .muse_moe_runner import (
     MuseMoERunner,
     MuseSparkQuarkOCPMXFusedMoE,
@@ -665,7 +666,10 @@ def forward(
             attn_output = self.attn(q, None, None)
         else:
             attn_output = self.attn(q, k, v)
+        return self.back(attn_output, og)
 
+    def back(self, attn_output: torch.Tensor, og: torch.Tensor | None) -> torch.Tensor:
+        """O-norm, O-gate and the output projection."""
         if self.fuse_o_norm_gate:
             assert og is not None
             attn_output = attn_output.reshape(-1, self.num_heads, self.head_dim)
@@ -844,6 +848,12 @@ def forward(
         assert stream.dtype == self.residual_dtype
 
         attn_out = self.self_attn(positions=positions, hidden_states=attn_input)
+        return self.after_attention(stream, attn_out)
+
+    def after_attention(
+        self, stream: torch.Tensor, attn_out: torch.Tensor
+    ) -> MuseSpark1xResidualStream:
+        """Attention boundary and FFN; the FFN boundary is returned unblended."""
         ffn_in, stream = norm_residual_then_norm(
             MuseSpark1xResidualStream(
                 stream=stream,
@@ -916,6 +926,8 @@ def __init__(
 
         self.fuse_across_layers = get_pp_group().world_size == 1
         self.make_empty_intermediate_tensors = self._make_empty_intermediate_tensors
+        # Set after weight loading under --monokernel.
+        self.mono: MuseSparkMono | None = None
 
     def _make_empty_intermediate_tensors(
         self, batch_size: int, dtype: torch.dtype, device: torch.device
@@ -957,8 +969,12 @@ def forward(  # type: ignore[override]
         aux_hidden_states: list[torch.Tensor] = []
         if get_pp_group().is_first_rank:
             self._maybe_add_hidden_state(aux_hidden_states, self.start_layer, x, None)
+        mono_out = None
+        if self.mono is not None:
+            mono_out = self.mono.try_forward(self, positions, x, aux_hidden_states)
         for idx, layer in enumerate(
-            self.layers[self.start_layer : self.end_layer], start=self.start_layer
+            self.layers[self.start_layer : self.end_layer] if mono_out is None else [],
+            start=self.start_layer,
         ):
             x = layer(positions, x)
             if not self.fuse_across_layers and isinstance(x, MuseSpark1xResidualStream):
@@ -970,6 +986,8 @@ def forward(  # type: ignore[override]
                     else x
                 )
                 aux_hidden_states.append(state)
+        if mono_out is not None:
+            x = mono_out
 
         if not get_pp_group().is_last_rank:
             assert isinstance(x, torch.Tensor)  # fuse_across_layers is off here
@@ -1034,6 +1052,7 @@ def __init__(
             layer_type = self.decoder_layer_cls
 
         super().__init__(vllm_config=vllm_config, prefix=prefix, layer_type=layer_type)
+        self.vllm_config = vllm_config
 
         # qkv_proj fuses Q+K+V+OG, so LoRA must treat it as one module.
         self.packed_modules_mapping = {  # type: ignore[misc]
@@ -1126,6 +1145,8 @@ def compute_logits(
         self,
         hidden_states: torch.Tensor,
     ) -> torch.Tensor | None:
+        if self.model.mono is not None:
+            self.model.mono.poll_health()
         hidden_states = hidden_states.to(self.lm_head.weight.dtype)
         logits = self.logits_processor(self.lm_head, hidden_states)
         if logits is None:
@@ -1133,6 +1154,8 @@ def compute_logits(
         return self._apply_output_transform(logits)
 
     def compute_logits_local(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        if self.model.mono is not None:
+            self.model.mono.poll_health()
         hidden_states = hidden_states.to(self.lm_head.weight.dtype)
         logits = self.logits_processor(self.lm_head, hidden_states, skip_gather=True)
         return self._apply_output_transform(logits)
@@ -1179,15 +1202,19 @@ def interleaved_weights() -> Iterable[tuple[str, torch.Tensor]]:
                 yield name, weight
 
         loader = AutoWeightsLoader(self)
-        loaded = loader.load_weights(
-            interleaved_weights(), mapper=self.hf_to_vllm_mapper
-        )
+        return loader.load_weights(interleaved_weights(), mapper=self.hf_to_vllm_mapper)
+
+    def process_weights_after_loading(self) -> None:
         n_gates = materialize_residual_path_constants(self)
         logger.info(
             "Materialized residual-path constants for %d gated residual connections",
             n_gates,
         )
-        return loaded
+        # The experts are in their kernel layout only after the per-layer
+        # quant finalize, which runs between load_weights and this hook. The
+        # monokernel also needs the cached coefficients and the folded router
+        # weights from above.
+        maybe_attach_mono(self)
 
 
 __all__ = [
diff --git a/vllm/models/muse_spark/amd/mono/README.md b/vllm/models/muse_spark/amd/mono/README.md
new file mode 100644
index 000000000000..4977a71370b5
--- /dev/null
+++ b/vllm/models/muse_spark/amd/mono/README.md
@@ -0,0 +1,82 @@
+# Muse Spark decode monokernel (MI355X, FlyDSL)
+
+`--monokernel` for `MuseSpark1xForCausalLM` on MI355X: decoder layers 1-60 of
+a speculative verify step run as persistent FlyDSL launches, from one layer's
+attention output to the next layer's q, o-gate and KV insert. vLLM's own
+attention op runs between the launches, except where a launch runs the
+attention itself (below); layer 0, layer 61 and every step the kernels do not
+cover run vLLM's stock path.
+
+## Scope
+
+- Deployment (checked at startup; a deployment that cannot run the kernels
+  raises): gfx950 with 256 CUs, tensor parallel size 8 with
+  `--enable-expert-parallel` (linear placement, 32 experts a rank, no EPLB),
+  no PP / DP / DCP / PCP, Quark OCP-MXFP4 W4A4 experts on `AITER_MXFP4_MXFP4`
+  (shuffled weights and scales), the fused QK-norm / RoPE / KV insert,
+  `ROCM_AITER_FA` with a bf16 KV cache in the LBHNC layout (block size 16),
+  DSpark with 7 draft tokens, no LoRA, sleep mode or weight transfer, no KV
+  connector or offloading.
+- Steps (checked every forward; anything else runs stock): uniform decode
+  verify batches of 1-16 requests x 8 tokens, FULL CUDA graph or eager.
+  8 / 16 rows run K3 (S = 8 / 16); 24-128 rows run HA + HM + HB with the rows
+  padded up to S = 32 / 64 / 96 / 128 (`nrows` = the real rows).
+
+## Launches (layer i, 1 <= i <= 60)
+
+    K3  o-norm x gate, o_proj, all-reduce, attention boundary (post-attn norm,
+        gated residual, pre-FFN norm), router + pre-expert projection, top-8,
+        pre-expert norm, the local experts (gate-up, SwiGLU, down, post-expert
+        norm, weighted sum), EP all-reduce, post-expert projection, FFN
+        boundary (post-FFN norm, gated residual, layer i+1's input norm),
+        layer i+1's QKV + O-gate GEMV, QK-norm, RoPE, KV insert
+    HA  K3 up to the pre-expert norm; the residual stream token-sharded
+    HM  the local experts; the partials pushed to their owner ranks
+    HB  EP reduce, post-expert projection, FFN boundary, QKV, KV insert; the
+        stream all-gathered only where a full copy is read (DSpark aux layers,
+        layer 61)
+
+Fused attention: at 8, 16 and 24-32 rows every layer from layer 2 runs its
+verify attention at the head of K3_i / HA_i (`kernels/attn_frag.py`) instead
+of vLLM's op: split over all 256 CTAs by
+(request, kv head, key range), then merged before o_proj. It reads layer i's
+q from the previous launch, layer i's bf16 KV cache through its KV cache
+group's `block_table` (int32) and `seq_lens` (int32), the softmax scale of the
+layer's attention impl, and for sliding-window layers vLLM's window (2047, 0)
+(only blocks inside the window are read). The builds stage block ids for up
+to 1M keys a request; with a larger `max_model_len` every layer runs vLLM's
+op. 40-128 rows keep vLLM's op.
+
+## Contract with vLLM
+
+- Weights: read in place from vLLM's parameters (`kernels/` builds take
+  vLLM's layouts): QKV + O-gate rows `[q | k | v | o_gate]` per rank,
+  `o_proj`, the fp32 router with the output multiplier folded in (rows
+  `[32 r, 32 r + 32)`), `pooler` / `unpooler` row slices, the experts after
+  `shuffle_weight(layout=(16, 16))` and the scales after `e8m0_shuffle`. The
+  norm multipliers are small bf16 copies (`weight + gain_center`).
+- KV: layer i+1's `kv_cache` `[block][kv head][16 tokens][K 64 | V 64]` bf16
+  and its group's `slot_mapping` (`slot < 0` = graph padding, skipped); the
+  block stride is a launch argument.
+- Fused attention: one descriptor table a build and step shape (10 words a
+  layer: q, KV cache and block stride, block table and row stride, seq_lens,
+  output, scratch, request count, scale). Tables are keyed by their contents
+  and written once (a pinned-buffer copy when first seen during graph capture).
+- q leaves without the QK scale: vLLM's attention applies
+  `0.125 * qk_scale_factor` as its softmax scale.
+- Peers: one `cudaMalloc` buffer a rank and build, opened by every rank through
+  IPC handles exchanged on the TP CPU group (plain device memory, not the
+  uncached custom all-reduce buffers). Buffers are allocated at load time and
+  every launch argument is a fixed pointer, so the launches are captured in
+  vLLM's FULL graphs.
+- Health: a wait on a peer rank times out instead of hanging and leaves a
+  record in the build's health words; `compute_logits` copies them back
+  asynchronously every 64 steps and raises on a timeout.
+
+## Numerics
+
+The kernels keep the MoE activations in bf16 (MXFP4 weights x bf16
+activations). The stock W4A4 path quantizes the expert inputs and the SwiGLU
+output to MXFP4, so `--monokernel` outputs differ from the stock path (and
+prefill, which runs stock, uses W4A4). Other differences are reduction order
+(router, GEMVs, all-reduces in rank order) and the bf16 norm multipliers.
diff --git a/vllm/models/muse_spark/amd/mono/__init__.py b/vllm/models/muse_spark/amd/mono/__init__.py
new file mode 100644
index 000000000000..50bfc03b333f
--- /dev/null
+++ b/vllm/models/muse_spark/amd/mono/__init__.py
@@ -0,0 +1,146 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""FlyDSL decode monokernel for MuseSpark1x on MI355X (``--monokernel``)."""
+
+import importlib
+import os
+from pathlib import Path
+
+import torch
+import torch.distributed as dist
+from torch import nn
+
+from vllm.distributed import get_tp_group
+from vllm.logger import init_logger
+from vllm.platforms import current_platform
+
+from .decoder import (
+    FIRST_LAYER,
+    LAST_LAYER,
+    NUM_CUS,
+    MonoDims,
+    MonoKernels,
+    MuseSparkMono,
+    _routed_experts,
+    built_buckets,
+    config_refusals,
+    deployment_refusals,
+    fused_attention_layers,
+    layer_refusals,
+)
+
+logger = init_logger(__name__)
+
+
+def _ensure_writable_flydsl_cache() -> None:
+    """AITER points FlyDSL at its bundled, read-only cache; new kernels need a
+    writable one."""
+    cur = os.environ.get("FLYDSL_RUNTIME_CACHE_DIR")
+    if cur and os.access(cur, os.W_OK):
+        return
+    path = Path.home() / ".flydsl" / "cache"
+    path.mkdir(parents=True, exist_ok=True)
+    os.environ["FLYDSL_RUNTIME_CACHE_DIR"] = str(path)
+
+
+def _load_kernels() -> tuple[MonoKernels | None, str | None]:
+    _ensure_writable_flydsl_cache()
+    try:
+        modules = {
+            name: importlib.import_module(f"{__name__}.kernels.{name}")
+            for name in ("k3", "k5", "k5m", "attn_frag")
+        }
+    except ImportError as e:
+        return None, f"needs FlyDSL and AITER's FlyDSL helpers ({e})"
+    return MonoKernels(**modules), None
+
+
+def static_refusals(lm: nn.Module, tp_rank: int, tp_size: int) -> list[str]:
+    """Every reason this deployment cannot run the monokernel; empty = supported."""
+    if not current_platform.is_rocm():
+        return ["not ROCm"]
+    from vllm.platforms.rocm import on_gfx950
+
+    reasons = [] if on_gfx950() else ["not gfx950"]
+    cus = torch.cuda.get_device_properties(
+        torch.accelerator.current_device_index()
+    ).multi_processor_count
+    if cus != NUM_CUS:
+        reasons.append(f"{cus} CUs (kernel: {NUM_CUS})")
+    text_config = lm.vllm_config.model_config.hf_text_config
+    reasons += config_refusals(text_config)
+    if reasons:
+        return reasons
+    experts = _routed_experts(lm.model.layers[FIRST_LAYER])
+    moe_parallel = getattr(
+        getattr(experts, "moe_config", None), "moe_parallel_config", None
+    )
+    reasons = deployment_refusals(lm.vllm_config, tp_rank, tp_size, moe_parallel)
+    return reasons or layer_refusals(
+        lm.model, MonoDims.from_config(text_config, tp_size)
+    )
+
+
+def maybe_attach_mono(lm: nn.Module) -> None:
+    """Attach the monokernel to ``lm.model`` after the weights are final.
+
+    Collective over the TP group: every rank calls it. A deployment that any
+    rank cannot run raises on every rank, with every reason of this rank.
+    """
+    if not lm.vllm_config.kernel_config.monokernel:
+        return
+    model = lm.model
+    if model.mono is not None:
+        model.mono.refresh_weights(model)
+        return
+    group = get_tp_group()
+    rank, world_size = group.rank_in_group, group.world_size
+    text_config = lm.vllm_config.model_config.hf_text_config
+    reasons = static_refusals(lm, rank, world_size)
+    kernels = None
+    if not reasons:
+        kernels, missing = _load_kernels()
+        if missing is not None:
+            reasons.append(missing)
+    verdict = torch.tensor([0 if reasons else 1], dtype=torch.int32)
+    dist.all_reduce(verdict, op=dist.ReduceOp.MIN, group=group.cpu_group)
+    if int(verdict.item()) == 0:
+        raise ValueError(
+            "--monokernel: Muse Spark cannot run the monokernel: "
+            + ("; ".join(reasons) or "a peer rank refused the deployment")
+        )
+    assert kernels is not None
+    buckets = built_buckets(lm.vllm_config.scheduler_config.max_num_seqs)
+    max_model_len = lm.vllm_config.model_config.max_model_len
+    fused = fused_attention_layers(
+        model.layers, max_model_len, kernels.attn_frag.VLLM_MAX_KEYS
+    )
+    if not fused:
+        logger.info(
+            "Muse Spark monokernel: verify attention unfused (max_model_len %d, "
+            "fused builds cover %d)",
+            max_model_len,
+            kernels.attn_frag.VLLM_MAX_KEYS,
+        )
+    model.mono = MuseSparkMono(
+        model,
+        kernels,
+        rank=rank,
+        group=group.cpu_group,
+        device=torch.device("cuda", torch.accelerator.current_device_index()),
+        dims=MonoDims.from_config(text_config, world_size),
+        buckets=buckets,
+        fused_attention=fused,
+    )
+    logger.info(
+        "Muse Spark monokernel attached: layers %d-%d of verify steps with "
+        "up to %d rows run as FlyDSL launches (builds S=%s; MoE activations "
+        "in bf16)",
+        FIRST_LAYER,
+        LAST_LAYER,
+        buckets[-1],
+        ",".join(map(str, buckets)),
+    )
+
+
+__all__ = ["MuseSparkMono", "maybe_attach_mono"]
diff --git a/vllm/models/muse_spark/amd/mono/decoder.py b/vllm/models/muse_spark/amd/mono/decoder.py
new file mode 100644
index 000000000000..bb99f53ba49a
--- /dev/null
+++ b/vllm/models/muse_spark/amd/mono/decoder.py
@@ -0,0 +1,1018 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""MuseSpark1x decode monokernel: layers 1-60 of a verify step as FlyDSL launches.
+
+See README.md. Every rank must take the same path (the kernel waits on its
+peers), so the decision uses only rank-invariant data: the static deployment
+checked at attach, the batch shape and the attention metadata counts.
+"""
+
+import dataclasses
+from collections.abc import Sequence
+from dataclasses import dataclass
+from types import ModuleType
+from typing import Any
+
+import torch
+from torch import nn
+
+from vllm.config import CUDAGraphMode
+from vllm.forward_context import get_forward_context
+from vllm.logger import init_logger
+from vllm.model_executor.layers.attention.attention import get_attention_context
+from vllm.models.muse_spark.common.layers import (
+    MuseSpark1xNormPreservingResidualConnection,
+    MuseSpark1xRMSNormWithInputScale,
+    norm_then_gated_residual,
+)
+
+from ..muse_moe_runner import MuseSparkQuarkOCPMXFusedMoE
+from .peer import PeerBuffer
+
+logger = init_logger(__name__)
+
+FIRST_LAYER = 1
+LAST_LAYER = 60
+ROWS_PER_REQUEST = 8
+K3_ROWS = (8, 16)
+K5_ROWS = (32, 64, 96, 128)
+WORLD_SIZE = 8
+NUM_CUS = 256
+KV_BLOCK_SIZE = 16
+HEALTH_POLL_INTERVAL = 64
+# Attention types whose verify attention runs inside the launch, per build.
+# Other builds and layer 1 run vLLM's attention op.
+FUSED_ATTENTION = {8: ("swa", "full"), 16: ("swa", "full"), 32: ("swa", "full")}
+SLIDING_WINDOW = (2047, 0)
+
+# The model shape the K3 build is specialized for.
+KERNEL_CONFIG: dict[str, float] = {
+    "hidden_size": 8192,
+    "moe_hidden_size": 4096,
+    "moe_intermediate_size": 4096,
+    "num_local_experts": 256,
+    "num_experts_per_tok": 8,
+    "head_dim": 64,
+    "num_attention_heads": 128,
+    "num_key_value_heads": 16,
+    "num_hidden_layers": 62,
+    "routed_scaling_factor": 1.0,
+    "rms_norm_eps": 1e-5,
+    "post_norm_eps": 1e-8,
+    "residual_gate_temperature": 0.3,
+}
+
+
+@dataclass(frozen=True)
+class MonoDims:
+    """Per-rank tensor sizes of a TP = EP rank."""
+
+    world_size: int
+    hidden_size: int
+    pooled_size: int
+    intermediate_size: int
+    num_experts: int
+    num_heads: int
+    num_kv_heads: int
+    head_dim: int
+
+    @classmethod
+    def from_config(cls, config: Any, world_size: int) -> "MonoDims":
+        return cls(
+            world_size=world_size,
+            hidden_size=config.hidden_size,
+            pooled_size=config.moe_hidden_size,
+            intermediate_size=config.moe_intermediate_size,
+            num_experts=config.num_local_experts,
+            num_heads=config.num_attention_heads,
+            num_kv_heads=config.num_key_value_heads,
+            head_dim=config.head_dim,
+        )
+
+    @property
+    def local_experts(self) -> int:
+        return self.num_experts // self.world_size
+
+    @property
+    def local_kv_heads(self) -> int:
+        return self.num_kv_heads // self.world_size
+
+    @property
+    def q_rows(self) -> int:
+        return self.num_heads // self.world_size * self.head_dim
+
+    @property
+    def kv_rows(self) -> int:
+        return self.local_kv_heads * self.head_dim
+
+    @property
+    def qkvg_rows(self) -> int:
+        return 2 * self.q_rows + 2 * self.kv_rows
+
+    def expert_shapes(self) -> dict[str, tuple[int, ...]]:
+        """MXFP4 expert tensors of the local experts, as uint8."""
+        e, h, i = self.local_experts, self.pooled_size, self.intermediate_size
+        return {
+            "w13_weight": (e, 2 * i, h // 2),
+            "w13_weight_scale": (e, 2 * i, h // 32),
+            "w2_weight": (e, h, i // 2),
+            "w2_weight_scale": (e, h, i // 32),
+        }
+
+
+def bucket_of(rows: int) -> int | None:
+    """The kernel build that serves ``rows`` (padded up); None = stock."""
+    if rows <= 0 or rows % ROWS_PER_REQUEST:
+        return None
+    if rows in K3_ROWS:
+        return rows
+    return next((s for s in K5_ROWS if rows <= s), None)
+
+
+@dataclass
+class LayerWeights:
+    """Views of vLLM's parameters (and small norm copies) for layer i's launches."""
+
+    w_o: torch.Tensor
+    alpha_a: torch.Tensor
+    beta_a: torch.Tensor
+    w_pffn: torch.Tensor
+    w_r: torch.Tensor
+    bias: torch.Tensor
+    w_pre: torch.Tensor
+    w_pren: torch.Tensor
+    w13: torch.Tensor
+    w13s: torch.Tensor
+    w2: torch.Tensor
+    w2s: torch.Tensor
+    w_post: torch.Tensor
+    w_pp: torch.Tensor
+    w_pff: torch.Tensor
+    alpha: torch.Tensor
+    beta: torch.Tensor
+    w_in: torch.Tensor
+    w_qkvg: torch.Tensor
+    next_rope: bool
+
+    def __post_init__(self) -> None:
+        def ptrs(names: str) -> tuple[int, ...]:
+            return tuple(getattr(self, n).data_ptr() for n in names.split())
+
+        # HA (and K3): attention boundary, router and pre-expert projection;
+        # HM: the local experts; HB: post-expert projection, FFN boundary and
+        # the next QKV. Without EPLB replicas, the replica arguments are the
+        # own experts.
+        self.ha = ptrs("w_o alpha_a beta_a w_pffn w_r bias w_pre w_pren")
+        self.hm = ptrs("w13 w13s w2 w2s w_post")
+        self.hb = ptrs("w_pp w_pff alpha beta w_in w_qkvg")
+        self.replicas = ptrs("w13 w13s w2 w2s w13")
+
+
+def effective_norm_weight(norm: nn.Module) -> torch.Tensor:
+    """The bf16 multiplier the kernel applies: ``weight (+ gain_center)``."""
+    weight = norm.weight.detach().float()
+    if norm.zero_centered_gamma:
+        weight = weight + norm.gain_center
+    return weight.to(torch.bfloat16).contiguous()
+
+
+def _routed_experts(layer: nn.Module) -> nn.Module | None:
+    experts = getattr(layer.feed_forward, "experts", None)
+    return getattr(experts, "routed_experts", None)
+
+
+def _rows(t: torch.Tensor, rank: int, n: int) -> torch.Tensor:
+    return t.detach()[rank * n : (rank + 1) * n]
+
+
+def collect_layer_weights(
+    layers: Sequence[nn.Module], i: int, rank: int, dims: MonoDims
+) -> LayerWeights:
+    """Layer i's kernel weights from decoder layers ``i`` and ``i + 1``; no copy."""
+    layer, nxt = layers[i], layers[i + 1]
+    moe = layer.feed_forward
+    experts = _routed_experts(layer)
+    assert experts is not None
+    w = dims.world_size
+    return LayerWeights(
+        w_o=layer.self_attn.o_proj.weight.detach(),
+        alpha_a=layer.add_attn.alpha,
+        beta_a=layer.add_attn.beta,
+        w_pffn=effective_norm_weight(layer.pre_ffn_norm),
+        w_r=_rows(moe.router.weight, rank, dims.num_experts // w),
+        bias=moe.expert_bias,
+        w_pre=_rows(moe.pooler.weight, rank, dims.pooled_size // w),
+        w_pren=effective_norm_weight(moe.pre_expert_norm),
+        w13=experts.w13_weight.detach().view(torch.uint8),
+        w13s=experts.w13_weight_scale.detach().view(torch.uint8),
+        w2=experts.w2_weight.detach().view(torch.uint8),
+        w2s=experts.w2_weight_scale.detach().view(torch.uint8),
+        w_post=effective_norm_weight(moe.post_expert_norm),
+        w_pp=_rows(moe.unpooler.weight, rank, dims.hidden_size // w),
+        w_pff=effective_norm_weight(layer.post_ffn_norm),
+        alpha=layer.add_ffn.alpha,
+        beta=layer.add_ffn.beta,
+        w_in=effective_norm_weight(nxt.self_attn.pre_qkv_proj_norm),
+        w_qkvg=nxt.self_attn.qkv_proj.weight.detach(),
+        next_rope=nxt.self_attn.rotary_emb is not None,
+    )
+
+
+def config_refusals(config: Any) -> list[str]:
+    reasons = []
+    for name, want in KERNEL_CONFIG.items():
+        got = getattr(config, name, None)
+        if not isinstance(got, (int, float)) or abs(got - want) > 1e-12 * max(
+            1.0, abs(want)
+        ):
+            reasons.append(f"config.{name}={got} (kernel: {want})")
+    return reasons
+
+
+def deployment_refusals(
+    vllm_config: Any, tp_rank: int, tp_size: int, moe_parallel: Any
+) -> list[str]:
+    """Parallel layout and server features the monokernel does not support."""
+    reasons = []
+    parallel = vllm_config.parallel_config
+    if tp_size != WORLD_SIZE:
+        reasons.append(f"tensor_parallel_size={tp_size} (kernel: {WORLD_SIZE})")
+    for name in (
+        "pipeline_parallel_size",
+        "data_parallel_size",
+        "decode_context_parallel_size",
+        "prefill_context_parallel_size",
+    ):
+        if getattr(parallel, name, 1) != 1:
+            reasons.append(f"{name}={getattr(parallel, name)}")
+    if not parallel.enable_expert_parallel:
+        reasons.append("expert parallelism off (needs --enable-expert-parallel)")
+    if parallel.enable_eplb:
+        reasons.append("EPLB on")
+    if parallel.expert_placement_strategy != "linear":
+        reasons.append(f"expert placement {parallel.expert_placement_strategy}")
+    if moe_parallel is None or not moe_parallel.use_ep:
+        reasons.append("MoE layers not expert parallel")
+    elif moe_parallel.ep_size != tp_size or moe_parallel.ep_rank != tp_rank:
+        reasons.append(
+            f"EP rank {moe_parallel.ep_rank}/{moe_parallel.ep_size} "
+            f"!= TP rank {tp_rank}/{tp_size}"
+        )
+    spec = vllm_config.speculative_config
+    if spec is None or spec.num_speculative_tokens != ROWS_PER_REQUEST - 1:
+        reasons.append(
+            f"needs speculative decoding with {ROWS_PER_REQUEST - 1} draft tokens"
+        )
+    if vllm_config.lora_config is not None:
+        reasons.append("LoRA on")
+    if vllm_config.model_config.enable_sleep_mode:
+        reasons.append("sleep mode on")
+    if vllm_config.weight_transfer_config is not None:
+        reasons.append("weight transfer on")
+    if vllm_config.kv_transfer_config is not None:
+        reasons.append("KV connector on")
+    cache = vllm_config.cache_config
+    if getattr(cache, "kv_offloading_size", None):
+        reasons.append("KV offloading on")
+    if cache.cache_dtype not in ("auto", "bfloat16"):
+        reasons.append(f"kv_cache_dtype={cache.cache_dtype} (kernel: bf16)")
+    if vllm_config.model_config.dtype != torch.bfloat16:
+        reasons.append(f"model dtype {vllm_config.model_config.dtype}")
+    return reasons
+
+
+def _shape(t: torch.Tensor | None) -> tuple[int, ...] | None:
+    return None if t is None else tuple(t.shape)
+
+
+def layer_refusals(model: nn.Module, dims: MonoDims) -> list[str]:
+    """Module and tensor guards for layers FIRST_LAYER .. LAST_LAYER + 1."""
+    reasons: list[str] = []
+    if getattr(model, "residual_dtype", None) != torch.float32:
+        reasons.append("residual dtype not fp32")
+    if not getattr(model, "fuse_across_layers", False):
+        reasons.append("layer fusion off")
+    layers = model.layers
+    if len(layers) != LAST_LAYER + 2:
+        return reasons + [f"{len(layers)} layers"]
+    bad = reasons.append
+    h, p = dims.hidden_size, dims.pooled_size
+
+    def check_bf16(i: int, name: str, w: torch.Tensor, shape: tuple) -> None:
+        if _shape(w) != shape or w.dtype != torch.bfloat16 or not w.is_contiguous():
+            bad(f"layer {i}: {name} {_shape(w)} {w.dtype}")
+
+    for i in range(FIRST_LAYER, LAST_LAYER + 2):
+        layer, attn = layers[i], layers[i].self_attn
+        if not (
+            attn.use_qk_norm
+            and attn.use_o_norm
+            and attn.use_o_gate
+            and attn.fuse_o_norm_gate
+            and attn.fused_qk_norm_rope_kv_insert
+        ) or getattr(attn, "use_attn_temperature_tuning", False):
+            bad(f"layer {i}: attention options (needs the fused KV insert)")
+        backend = getattr(getattr(attn.attn, "attn_backend", None), "__name__", None)
+        if backend != "AiterFlashAttentionBackend":
+            bad(f"layer {i}: attention backend {backend}")
+        check_bf16(i, "qkv_proj", attn.qkv_proj.weight, (dims.qkvg_rows, h))
+        rope = attn.rotary_emb
+        if rope is not None and (
+            rope.cos_sin_cache.dtype != torch.float32
+            or rope.cos_sin_cache.shape[-1] != dims.head_dim
+            or rope.rotary_dim != dims.head_dim
+            or rope.is_neox_style
+        ):
+            bad(f"layer {i}: RoPE cache {_shape(rope.cos_sin_cache)}")
+        if not layer.self_attn.pre_qkv_proj_norm.has_weight:
+            bad(f"layer {i}: weightless pre_qkv_proj_norm")
+        if i > LAST_LAYER:
+            continue
+        check_bf16(i, "o_proj", attn.o_proj.weight, (h, dims.q_rows))
+        for name in ("add_attn", "add_ffn"):
+            conn = getattr(layer, name)
+            if not isinstance(
+                conn, MuseSpark1xNormPreservingResidualConnection
+            ) or not getattr(conn, "_coefficients_ready", False):
+                bad(f"layer {i}: {name} coefficients not materialized")
+        if layer.post_attn_norm.has_weight:
+            bad(f"layer {i}: post_attn_norm has a weight")
+        if not isinstance(layer.post_ffn_norm, MuseSpark1xRMSNormWithInputScale):
+            bad(f"layer {i}: post_ffn_norm type")
+        if not getattr(layer, "is_moe_layer", False):
+            bad(f"layer {i}: dense FFN")
+            continue
+        moe = layer.feed_forward
+        r = moe.router.weight
+        if _shape(r) != (dims.num_experts, h) or r.dtype != torch.float32:
+            bad(f"layer {i}: router {_shape(r)} {r.dtype}")
+        if moe.router_output_multiplier is not None and not getattr(
+            moe, "_router_multiplier_folded", False
+        ):
+            bad(f"layer {i}: router multiplier not folded")
+        check_bf16(i, "pooler", moe.pooler.weight, (p, h))
+        check_bf16(i, "unpooler", moe.unpooler.weight, (h, p))
+        experts = _routed_experts(layer)
+        if experts is None or not isinstance(
+            getattr(experts, "quant_method", None), MuseSparkQuarkOCPMXFusedMoE
+        ):
+            bad(f"layer {i}: experts not MXFP4 W4A4")
+            continue
+        for name, want in dims.expert_shapes().items():
+            t = getattr(experts, name)
+            if _shape(t) != want or t.element_size() != 1 or not t.is_contiguous():
+                bad(f"layer {i}: {name} {_shape(t)} {t.dtype}")
+        if not (
+            getattr(experts.w13_weight, "is_shuffled", False)
+            and getattr(experts.w2_weight, "is_shuffled", False)
+        ):
+            bad(f"layer {i}: expert weights not shuffled")
+        for e_n in (2 * dims.intermediate_size, dims.pooled_size):
+            if dims.local_experts * e_n % 256:
+                bad(f"layer {i}: padded scale shuffle ({dims.local_experts} x {e_n})")
+    return reasons
+
+
+def kv_cache_refusal(kv_cache: torch.Tensor, kv_heads: int) -> str | None:
+    """K3 writes bf16 LBHNC pages: ``[block][kv head][16 tokens][K | V]``."""
+    if kv_cache.numel() == 0:
+        return "kv:unbound"
+    if kv_cache.dtype != torch.bfloat16:
+        return f"kv:{kv_cache.dtype}"
+    want_shape = (kv_heads, KV_BLOCK_SIZE, 2 * 64)
+    if tuple(kv_cache.shape[1:]) != want_shape:
+        return f"kv:shape{tuple(kv_cache.shape)}"
+    if tuple(kv_cache.stride()[1:]) != (KV_BLOCK_SIZE * 128, 128, 1):
+        return f"kv:stride{tuple(kv_cache.stride())}"
+    return None
+
+
+def attention_type(layer: nn.Module) -> str | None:
+    """``"swa"`` / ``"full"`` when the fused attention can run layer's
+    attention (vLLM's ``sliding_window`` (2047, 0) or none); else None."""
+    window = getattr(
+        getattr(layer.self_attn.attn, "impl", None), "sliding_window", None
+    )
+    if window == SLIDING_WINDOW:
+        return "swa"
+    if window in (None, (-1, -1)):
+        return "full"
+    return None
+
+
+def fused_attention_layers(
+    layers: Sequence[nn.Module], max_model_len: int, max_keys: int
+) -> dict[int, str]:
+    """Layers whose attention the launches can fuse: ``{i: type}``.
+
+    Empty when a request may hold more keys than a build stages block ids for.
+    """
+    if max_model_len > max_keys:
+        return {}
+    fused = {}
+    for i in range(FIRST_LAYER + 1, LAST_LAYER + 1):
+        kind = attention_type(layers[i])
+        if kind is not None:
+            fused[i] = kind
+    return fused
+
+
+def batch_refusal(
+    num_tokens: int, forward_context: Any, layer_names: Sequence[str]
+) -> str | None:
+    """Why this step runs stock, from rank-invariant data only; None = mono.
+
+    Reads the attention metadata of the kernels' layers (``layer_names``) only:
+    the dict also holds the draft model's layers.
+    """
+    if bucket_of(num_tokens) is None:
+        return f"rows:{num_tokens}"
+    mode = forward_context.cudagraph_runtime_mode
+    if mode not in (CUDAGraphMode.NONE, CUDAGraphMode.FULL):
+        return f"cudagraph:{mode.name}"
+    metadata = forward_context.attn_metadata
+    if not isinstance(metadata, dict) or not metadata:
+        return "attn-metadata"
+    seen: set[int] = set()
+    for name in layer_names:
+        md = metadata.get(name)
+        if md is None:
+            return "attn-metadata"
+        if id(md) in seen:
+            continue
+        seen.add(id(md))
+        decode = getattr(md, "decode_metadata", None)
+        if (
+            getattr(md, "num_prefills", None) != 0
+            or getattr(md, "num_extends", None) != 0
+            or decode is None
+        ):
+            return "not-decode"
+        if decode.max_query_len != ROWS_PER_REQUEST:
+            return f"query-len:{decode.max_query_len}"
+        if md.num_decodes * ROWS_PER_REQUEST != num_tokens:
+            return f"width:{num_tokens}/{md.num_decodes}"
+    return None
+
+
+@dataclass(frozen=True)
+class MonoKernels:
+    """The FlyDSL kernel modules (``kernels.k3``, ``k5``, ``k5m``, ``attn_frag``)."""
+
+    k3: ModuleType
+    k5: ModuleType
+    k5m: ModuleType
+    attn_frag: ModuleType
+
+
+class _FusedAttention:
+    """The fused attention of one build: its output, scratch and descriptor
+    tables (``kernels/attn_frag.py``; row i = layer i).
+
+    A table holds pointers and counts only (q, KV cache, block table,
+    seq_lens, output, request count, scale). Tables are keyed by their
+    contents, so one is written once and never changes after a CUDA graph
+    captured a launch that reads it.
+    """
+
+    def __init__(
+        self,
+        attn_frag: ModuleType,
+        layers: dict[int, str],
+        scales: dict[int, float],
+        rows: int,
+        dims: MonoDims,
+        device: torch.device,
+    ) -> None:
+        self.af = attn_frag
+        self.layers = layers
+        self.scales = scales
+        self.out = torch.zeros(rows, dims.q_rows, dtype=torch.bfloat16, device=device)
+        self.part = torch.zeros(
+            attn_frag.part_bytes() // 4, dtype=torch.int32, device=device
+        )
+        self.words = attn_frag.TABLE_ROWS * attn_frag.DESC_WORDS
+        self.device = device
+        self.tables: dict[tuple, torch.Tensor] = {}
+        # Pinned sources of the table copies captured in CUDA graphs: a pinned
+        # allocation is not allowed while a stream captures.
+        self._pinned = [
+            torch.zeros(self.words, dtype=torch.int64, pin_memory=True)
+            for _ in range(8)
+        ]
+        self._keep: list[torch.Tensor] = []
+
+    def table(
+        self, layers: Sequence[nn.Module], q_of: Any, num_reqs: int
+    ) -> torch.Tensor:
+        """The descriptor table of this step; ``q_of(i)`` = layer i's q."""
+        rows = {}
+        for i in self.layers:
+            md, _, kv_cache, _ = get_attention_context(
+                layers[i].self_attn.attn.layer_name
+            )
+            rows[i] = self.af.vllm_desc(
+                q_of(i),
+                kv_cache,
+                md.block_table,
+                md.seq_lens,
+                self.out,
+                self.part,
+                num_reqs,
+                self.scales[i],
+            )
+        key = tuple((i, *w) for i, w in rows.items())
+        table = self.tables.get(key)
+        if table is not None:
+            return table
+        host = torch.zeros(self.words, dtype=torch.int64)
+        n = self.af.DESC_WORDS
+        for i, w in rows.items():
+            host[i * n : (i + 1) * n] = torch.tensor(w, dtype=torch.int64)
+        table = torch.zeros(self.words, dtype=torch.int64, device=self.device)
+        if torch.cuda.is_current_stream_capturing():
+            if not self._pinned:
+                raise RuntimeError(
+                    "Muse Spark monokernel: more fused-attention tables than "
+                    "pinned buffers"
+                )
+            pinned = self._pinned.pop()
+            pinned.copy_(host)
+            self._keep.append(pinned)
+            table.copy_(pinned, non_blocking=True)
+        else:
+            table.copy_(host)
+        self.tables[key] = table
+        return table
+
+    def row(self, table: torch.Tensor, i: int) -> int:
+        return table.data_ptr() + i * self.af.DESC_WORDS * table.element_size()
+
+
+class _Bucket:
+    """Launchers and persistent buffers of one kernel build (row count)."""
+
+    q: list[torch.Tensor]
+
+    def __init__(
+        self,
+        kernel: ModuleType,
+        peer_bytes: int,
+        dims: MonoDims,
+        device: torch.device,
+        group: Any,
+    ) -> None:
+        self.dims = dims
+        self.blocks = kernel.BLOCKS
+        self.peers = PeerBuffer(peer_bytes + 256, group)
+        self.epoch = torch.zeros(4, dtype=torch.int32, device=device)
+        self.health = torch.zeros(kernel.HEALTH_WORDS, dtype=torch.int32, device=device)
+        self.timeline = torch.zeros(
+            kernel.BLOCKS * kernel.TL_POINTS, dtype=torch.int64, device=device
+        )
+        self.no_rope = torch.zeros(64, dtype=torch.float32, device=device)
+        self.fa: _FusedAttention | None = None
+
+    def _common(self) -> tuple[int, ...]:
+        return (
+            self.epoch.data_ptr(),
+            self.health.data_ptr(),
+            self.timeline.data_ptr(),
+            *self.peers.pointers,
+        )
+
+    def run(self, mono, layers, positions, o, g, stream, aux_layers, aux_hidden_states):
+        """Layers FIRST_LAYER..LAST_LAYER; returns layer 61's (stream, q, g)."""
+        q, fa, table = o, self.fa, None
+        if fa is not None:
+            nreq = stream.shape[0] // ROWS_PER_REQUEST
+            table = fa.table(layers, lambda i: self.q[(i - 1) % 2], nreq)
+        for i in range(FIRST_LAYER, LAST_LAYER + 1):
+            kind = fa.layers.get(i) if fa is not None else None
+            if fa is not None and kind is not None:
+                attn_o = fa.row(table, i)
+            else:
+                if i > FIRST_LAYER:
+                    o = layers[i].self_attn.attn(q, None, None)
+                attn_o = o.data_ptr()
+            nxt = layers[i + 1].self_attn
+            _, _, kv_cache, slots = get_attention_context(nxt.attn.layer_name)
+            rope = nxt.rotary_emb
+            cos_sin = rope.cos_sin_cache if rope is not None else self.no_rope
+            ptrs = (slots.data_ptr(), cos_sin.data_ptr(), positions.data_ptr())
+            # The kernels take the block stride (elements) in the V slot.
+            kv = (kv_cache.data_ptr(), kv_cache.stride(0), *ptrs)
+            stream, q, g = self.layer(mono, i, kind, attn_o, g, stream, kv, aux_layers)
+            if i + 1 in aux_layers:
+                aux_hidden_states.append(stream.clone())
+        return stream, q, g
+
+    def layer(self, mono, i, kind, attn_o, g, stream, kv, aux_layers):
+        """Layer i's launches; returns layer i + 1's (stream, q, g)."""
+        raise NotImplementedError
+
+
+class _K3Bucket(_Bucket):
+    """8 / 16 rows: one K3 launch a layer."""
+
+    def __init__(
+        self,
+        k3: ModuleType,
+        rows: int,
+        dims: MonoDims,
+        device: torch.device,
+        group: Any,
+        fa: _FusedAttention | None = None,
+        fused_kinds: Sequence[tuple[bool, str]] = (),
+    ) -> None:
+        w = dims.world_size
+        super().__init__(k3, k3.peer_layout(rows, w)[1], dims, device, group)
+        self.launchers = {
+            next_rope: k3.build(k3.build_vllm(S=rows, W=w, next_rope=next_rope))
+            for next_rope in (True, False)
+        }
+        self.fa = fa
+        self.fused_launchers = {
+            (next_rope, kind): k3.build(
+                k3.build_vllm(S=rows, W=w, next_rope=next_rope, attn=kind)
+            )
+            for next_rope, kind in fused_kinds
+        }
+        self.scratch = torch.zeros(
+            k3.scratch_layout(rows, w)[1] // 4 + 64, dtype=torch.int32, device=device
+        )
+        bf16 = torch.bfloat16
+        self.outs = [
+            {
+                "stream": torch.zeros(
+                    rows, dims.hidden_size, dtype=torch.float32, device=device
+                ),
+                "xn": torch.zeros(rows, dims.hidden_size, dtype=bf16, device=device),
+                "part": torch.zeros(rows, dims.pooled_size, dtype=bf16, device=device),
+                "q": torch.zeros(rows, dims.q_rows, dtype=bf16, device=device),
+                "g": torch.zeros(rows, dims.q_rows, dtype=bf16, device=device),
+            }
+            for _ in range(2)
+        ]
+        self.q = [out["q"] for out in self.outs]
+
+    def layer(self, mono, i, kind, attn_o, g, stream, kv, aux_layers):
+        a, out = mono.weights[i], self.outs[i % 2]
+        launch = (
+            self.launchers[a.next_rope]
+            if kind is None
+            else self.fused_launchers[(a.next_rope, kind)]
+        )
+        launch(
+            attn_o,
+            g.data_ptr(),
+            stream.data_ptr(),
+            *a.ha,
+            *a.hm,
+            *a.hb,
+            out["stream"].data_ptr(),
+            out["xn"].data_ptr(),
+            out["part"].data_ptr(),
+            out["q"].data_ptr(),
+            out["g"].data_ptr(),
+            *kv,
+            self.scratch.data_ptr(),
+            *self._common(),
+            i,
+            mono.rank,
+            g.stride(0),
+            *a.replicas,
+            stream=torch.cuda.current_stream(),
+        )
+        return out["stream"], out["q"], out["g"]
+
+
+class _K5Bucket(_Bucket):
+    """24-128 rows: HA (attention boundary, router, pre-expert projection),
+    HM (the local experts) and HB (FFN boundary, next QKV) a layer; between
+    the launches the residual stream lives token-sharded unless the build
+    replicates it (``key.bnd``)."""
+
+    def __init__(
+        self,
+        k5: ModuleType,
+        k5m: ModuleType,
+        rows: int,
+        dims: MonoDims,
+        device: torch.device,
+        group: Any,
+        fa: _FusedAttention | None = None,
+        fused_kinds: Sequence[str] = (),
+    ) -> None:
+        w = dims.world_size
+        key = k5.build_vllm(S=rows, W=w, next_rope=True)
+        super().__init__(k5, k5.peer_layout(rows, w, key.bnd)[1], dims, device, group)
+        self.bnd = key.bnd
+        self.launch_a, launch_b = k5.build(key)
+        self.launch_b = {
+            True: launch_b,
+            False: k5.build(dataclasses.replace(key, next_rope=False))[1],
+        }
+        self.launch_m = k5m.build(k5m.build_vllm(S=rows, W=w))
+        self.fa = fa
+        self.fused_launch_a = {
+            kind: k5.build(k5.build_vllm(S=rows, W=w, next_rope=True, attn=kind))[0]
+            for kind in fused_kinds
+        }
+        i32, f32, bf16 = torch.int32, torch.float32, torch.bfloat16
+        self.scratch = torch.zeros(
+            k5.scratch_layout(rows, w, key.bnd)[1] // 4 + 64, dtype=i32, device=device
+        )
+        self.scratch_m = torch.zeros(
+            k5m.scratch_layout(rows)[1] // 4 + 64, dtype=i32, device=device
+        )
+        self.timeline_m = torch.zeros(k5.BLOCKS * 8, dtype=torch.int64, device=device)
+        topk = int(KERNEL_CONFIG["num_experts_per_tok"])
+        self.h = torch.zeros(rows, dims.pooled_size, dtype=bf16, device=device)
+        self.ids = torch.zeros(rows, topk, dtype=i32, device=device)
+        self.scores = torch.zeros(rows, topk, dtype=f32, device=device)
+        self.part = torch.zeros(rows, dims.pooled_size, dtype=bf16, device=device)
+        sharded_rows = rows // w if key.bnd < 2 else rows
+        self.ts = [
+            torch.zeros(sharded_rows, dims.hidden_size, dtype=f32, device=device)
+            for _ in range(2)
+        ]
+        self.sfull = torch.zeros(rows, dims.hidden_size, dtype=f32, device=device)
+        self.next_in: tuple[torch.Tensor, int] = (self.sfull, 1)
+        self.q = [
+            torch.zeros(rows, dims.q_rows, dtype=bf16, device=device) for _ in range(2)
+        ]
+        self.g = [
+            torch.zeros(rows, dims.q_rows, dtype=bf16, device=device) for _ in range(2)
+        ]
+
+    def layer(self, mono, i, kind, attn_o, g, stream, kv, aux_layers):
+        nrows, a, cur = stream.shape[0], mono.weights[i], torch.cuda.current_stream()
+        s_in, s_full = (stream, 1) if i == FIRST_LAYER else self.next_in
+        launch_a = self.launch_a if kind is None else self.fused_launch_a[kind]
+        launch_a(
+            attn_o,
+            g.data_ptr(),
+            s_in.data_ptr(),
+            *a.ha,
+            self.h.data_ptr(),
+            self.ids.data_ptr(),
+            self.scores.data_ptr(),
+            self.scratch.data_ptr(),
+            *self._common(),
+            i,
+            mono.rank,
+            g.stride(0),
+            s_full,
+            nrows,
+            stream=cur,
+        )
+        self.launch_m(
+            self.h.data_ptr(),
+            self.ids.data_ptr(),
+            self.scores.data_ptr(),
+            *a.hm,
+            self.part.data_ptr(),
+            self.scratch_m.data_ptr(),
+            self.epoch.data_ptr(),
+            self.health.data_ptr(),
+            self.timeline_m.data_ptr(),
+            *self.peers.pointers,
+            i,
+            mono.rank,
+            nrows,
+            *a.replicas,
+            stream=cur,
+        )
+        gather = int(i == LAST_LAYER or i + 1 in aux_layers or self.bnd == 1)
+        ts, q_out, g_out = self.ts[i % 2], self.q[i % 2], self.g[i % 2]
+        self.launch_b[a.next_rope](
+            self.part.data_ptr(),
+            *a.hb,
+            ts.data_ptr(),
+            self.sfull.data_ptr(),
+            q_out.data_ptr(),
+            g_out.data_ptr(),
+            *kv,
+            self.scratch.data_ptr(),
+            *self._common(),
+            i,
+            mono.rank,
+            self.dims.hidden_size,
+            0,
+            gather,
+            nrows,
+            stream=cur,
+        )
+        self.next_in = (self.sfull, 1) if self.bnd == 1 else (ts, int(self.bnd >= 2))
+        return self.sfull[:nrows], q_out[:nrows], g_out[:nrows]
+
+
+def built_buckets(max_num_seqs: int) -> list[int]:
+    """The builds a server with ``max_num_seqs`` requests can use."""
+    max_rows = min(max_num_seqs, K5_ROWS[-1] // ROWS_PER_REQUEST) * ROWS_PER_REQUEST
+    return sorted(
+        {
+            b
+            for rows in range(ROWS_PER_REQUEST, max_rows + 1, ROWS_PER_REQUEST)
+            if (b := bucket_of(rows)) is not None
+        }
+    )
+
+
+class MuseSparkMono:
+    """Runs decoder layers 1-60 of eligible verify steps through the kernels."""
+
+    def __init__(
+        self,
+        model: nn.Module,
+        kernels: MonoKernels,
+        rank: int,
+        group: Any,
+        device: torch.device,
+        dims: MonoDims,
+        buckets: Sequence[int],
+        fused_attention: dict[int, str],
+    ) -> None:
+        self.rank = rank
+        self.dims = dims
+        self.refresh_weights(model)
+        self.buckets: dict[int, _Bucket] = {}
+        for b in buckets:
+            kinds = FUSED_ATTENTION.get(b, ())
+            layers = {i: t for i, t in fused_attention.items() if t in kinds}
+            fa = None
+            if layers:
+                scales = {
+                    i: float(model.layers[i].self_attn.attn.impl.scale) for i in layers
+                }
+                fa = _FusedAttention(kernels.attn_frag, layers, scales, b, dims, device)
+            if b in K3_ROWS:
+                k3_kinds = sorted(
+                    {(self.weights[i].next_rope, t) for i, t in layers.items()}
+                )
+                self.buckets[b] = _K3Bucket(
+                    kernels.k3, b, dims, device, group, fa, k3_kinds
+                )
+            else:
+                self.buckets[b] = _K5Bucket(
+                    kernels.k5,
+                    kernels.k5m,
+                    b,
+                    dims,
+                    device,
+                    group,
+                    fa,
+                    sorted(set(layers.values())),
+                )
+        self._validated: set[int] = set()
+        self._fallbacks: set[str] = set()
+        self._engaged: set[int] = set()
+        self._steps = 0
+        self._health_host = {
+            b: torch.zeros_like(buf.health, device="cpu").pin_memory()
+            for b, buf in self.buckets.items()
+        }
+        self._health_event: torch.cuda.Event | None = None
+
+    def refresh_weights(self, model: nn.Module) -> None:
+        """Re-take the views, e.g. after a weight reload replaced parameters."""
+        self.weights = {
+            i: collect_layer_weights(model.layers, i, self.rank, self.dims)
+            for i in range(FIRST_LAYER, LAST_LAYER + 1)
+        }
+
+    def _step_refusal(
+        self, model: nn.Module, positions: torch.Tensor, x: Any
+    ) -> str | None:
+        if not isinstance(x, torch.Tensor):
+            return "input"
+        rows = x.shape[0]
+        names = [
+            model.layers[i].self_attn.attn.layer_name
+            for i in range(FIRST_LAYER, LAST_LAYER + 2)
+        ]
+        reason = batch_refusal(rows, get_forward_context(), names)
+        if reason is not None:
+            return reason
+        if bucket_of(rows) not in self.buckets:
+            return f"rows:{rows}"
+        if positions.dtype != torch.int64 or positions.shape != (rows,):
+            return "positions"
+        for name in names[1:]:
+            _, _, kv_cache, slots = get_attention_context(name)
+            reason = kv_cache_refusal(kv_cache, self.dims.local_kv_heads)
+            if reason is not None:
+                return reason
+            if (
+                slots is None
+                or slots.dtype != torch.int64
+                or slots.shape[0] < rows
+                or slots.stride(-1) != 1
+            ):
+                return "slot-mapping"
+        return None
+
+    def try_forward(
+        self,
+        model: nn.Module,
+        positions: torch.Tensor,
+        x: Any,
+        aux_hidden_states: list[torch.Tensor],
+    ) -> Any:
+        """Returns the residual stream after the last layer, or None for stock."""
+        reason = self._step_refusal(model, positions, x)
+        if reason is not None:
+            if reason not in self._fallbacks:
+                self._fallbacks.add(reason)
+                logger.debug("Muse Spark monokernel: stock step (%s)", reason)
+            return None
+        n, rows = x.shape[0], bucket_of(x.shape[0])
+        assert rows is not None
+        if n not in self._engaged:
+            self._engaged.add(n)
+            kernel = "K3" if rows in K3_ROWS else "HA+HM+HB"
+            logger.info(
+                "Muse Spark monokernel engaged: %d rows, %s S=%d", n, kernel, rows
+            )
+        layers, aux_layers = model.layers, model.aux_hidden_state_layers
+        bucket = self.buckets[rows]
+        bucket.epoch.add_(1)
+
+        y = layers[0](positions, x)
+        if 1 in aux_layers:
+            aux_hidden_states.append(norm_then_gated_residual(y))
+        from ..model import norm_residual_then_norm  # model.py imports this package
+
+        attn = layers[FIRST_LAYER].self_attn
+        attn_input, stream = norm_residual_then_norm(
+            y, attn.pre_qkv_proj_norm, out_dtype=layers[FIRST_LAYER].default_dtype
+        )
+        qkv, _ = attn.qkv_proj(attn_input)
+        q, _, _, g = attn._qk_norm_rope_kv_insert(qkv, positions)
+        o = attn.attn(q, None, None)
+        self._validate(o, g, stream)
+        stream, q, g = bucket.run(
+            self, layers, positions, o, g, stream, aux_layers, aux_hidden_states
+        )
+
+        last = layers[LAST_LAYER + 1]
+        o = last.self_attn.attn(q, None, None)
+        y = last.after_attention(stream, last.self_attn.back(o, g))
+        if LAST_LAYER + 2 in aux_layers:
+            aux_hidden_states.append(norm_then_gated_residual(y))
+        return y
+
+    def _validate(
+        self, attn_o: torch.Tensor, attn_g: torch.Tensor, stream: torch.Tensor
+    ) -> None:
+        """Check the stock layer-1 tensors the kernels read, once a row count."""
+        rows, d = stream.shape[0], self.dims
+        if rows in self._validated:
+            return
+        for name, t, dtype, cols in (
+            ("attn_o", attn_o, torch.bfloat16, d.q_rows),
+            ("o-gate", attn_g, torch.bfloat16, d.q_rows),
+            ("stream", stream, torch.float32, d.hidden_size),
+        ):
+            # The o-gate is a column slice of the QKV + O-gate output.
+            if (
+                t.dtype != dtype
+                or tuple(t.shape) != (rows, cols)
+                or not (t.stride(1) == 1 if t is attn_g else t.is_contiguous())
+            ):
+                raise RuntimeError(
+                    f"Muse Spark monokernel: {name} is {t.dtype} "
+                    f"{tuple(t.shape)} {t.stride()}, expected {dtype} {(rows, cols)}"
+                )
+        self._validated.add(rows)
+
+    def poll_health(self) -> None:
+        """Raise if a launch timed out waiting for a peer rank.
+
+        Reads the health words with an asynchronous copy every
+        HEALTH_POLL_INTERVAL calls and checks the previous copy, so the step
+        never waits for the device.
+        """
+        if torch.cuda.is_current_stream_capturing():
+            return
+        self._steps += 1
+        if self._health_event is not None and self._health_event.query():
+            self._health_event = None
+            for b, host in self._health_host.items():
+                if int(host[0]) != 0:
+                    records = host[8:].view(self.buckets[b].blocks, 8)
+                    stuck = records[records[:, 0] != 0][:8].tolist()
+                    raise RuntimeError(
+                        f"Muse Spark monokernel: rank {self.rank}, S={b}: a launch "
+                        f"timed out waiting for a peer rank {stuck}"
+                    )
+        if self._health_event is None and self._steps % HEALTH_POLL_INTERVAL == 0:
+            for b, buf in self.buckets.items():
+                self._health_host[b].copy_(buf.health, non_blocking=True)
+            self._health_event = torch.cuda.Event()
+            self._health_event.record()
+
+    def close(self) -> None:
+        for buf in self.buckets.values():
+            buf.peers.close()
diff --git a/vllm/models/muse_spark/amd/mono/kernels/__init__.py b/vllm/models/muse_spark/amd/mono/kernels/__init__.py
new file mode 100644
index 000000000000..2d87d0bb534f
--- /dev/null
+++ b/vllm/models/muse_spark/amd/mono/kernels/__init__.py
@@ -0,0 +1,14 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""FlyDSL sources of the Muse Spark decode monokernel (see ../README.md).
+
+``k3.build_vllm(S, W, next_rope)`` (8 / 16 rows), ``k5.build_vllm`` (HA / HB)
+and ``k5m.build_vllm`` (HM) for 32-128 rows are the builds vLLM launches:
+TP = EP = 8, the router output multiplier folded into the router weight, q
+without the QK scale, bf16 KV pages in vLLM's LBHNC layout, and vLLM's
+shuffled MXFP4 expert weights and scales read in place. With ``attn="swa"`` /
+``"full"`` the K3 / HA builds also run the layer's verify attention at their
+head (``attn_frag.make``: vLLM's bf16 KV cache, block table and seq_lens).
+Importing these modules needs FlyDSL and AITER's FlyDSL
+helpers.
+"""
diff --git a/vllm/models/muse_spark/amd/mono/kernels/attn_frag.py b/vllm/models/muse_spark/amd/mono/kernels/attn_frag.py
new file mode 100644
index 000000000000..71059fb8f71a
--- /dev/null
+++ b/vllm/models/muse_spark/amd/mono/kernels/attn_frag.py
@@ -0,0 +1,678 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# mypy: ignore-errors
+# ruff: noqa: B002, E501, E741, SIM102, SIM108
+
+# Vendored from AMD's Muse Spark FlyDSL monokernel kernels (mk/attn_frag.py):
+# SPDX-License-Identifier: MIT
+# Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+"""The verify attention of one Muse Spark layer as a fragment of the persistent kernels: K3
+and HA run it at their head (attn = "swa" / "full"), so the layer's attention and its decoder
+work are one launch.
+
+Shapes (a TP8 rank): 16 q heads, 2 kv heads (QPK = 8), head dim 64, 8 query tokens a request,
+NR requests (K3 S8: 1, S16: 2; HA S32: 4). K / V bf16 in the layer's vLLM LBHNC cache
+`[block][2][16][K 64 | V 64]` (block 16); the page table is the layer group's block_table
+(int32 `[NR][stride]`); seq_used = vLLM's seq_lens (int32 [NR], keys incl. the 8 new ones).
+q [8 NR, 16, 64] bf16 (row stride 1024, without the QK scale); out attn_o [8 NR, 1024] bf16.
+SWA: vLLM's window (2047, 0), keys qpos - key < 2048, and only keys in the window are read
+(vLLM frees the blocks before it: their block-table entries are stale).
+
+Split (every CTA): group = (request, kv head), CPG CTAs a group; CTA j of a group takes tiles
+[lo_t + j tps, +tps) of 64 keys (tps = ceil(n_t / CPG) at run time) for all 64 rows (8 tokens x 8
+heads) of the kv head. 8 waves = 4 row quarters (16 rows) x 2 tile lanes; a step = 2 tiles.
+K / V tiles: global -> VGPR (PD steps ahead; one 64-bit block base a wave, a wave loads the 16
+keys of one block) -> LDS. S^T = K Q^T (16x16x32 bf16 MFMA, a lane holds one q row), online
+softmax (exp2 domain, the -inf -> 0 rule), O^T = V^T P^T (P^T lane-local by a key permutation
+inside the 32-key k block; V^T from LDS with ds_read_b64_tr_b16). The 2 tile lanes merge in LDS;
+the CTA writes (m, l, unnormalized o) for 64 rows (fp32, device scope) + one tagged flag.
+
+Merge: item = (group, row) -> one output row of 64 d. NR = 1: CTA 2i merges item i with 512
+threads (32 segment lanes x 16 d-quads); NR >= 2: NR / 2 items a CTA. Poll the group's flags,
+load the partials, merge (max trick) per thread, over lanes and over waves, out = o / l -> bf16
+-> attn_o and the same rows as tagged pairs (P_AO), which the consumer (K3 stage O, HA ng)
+polls directly.
+
+Descriptor (int64 words, device memory; vllm_desc): 0 q, 1 the KV cache base, 2 its block stride
+in elements (a value), 3 block table, 4 seq_lens, 5 attn_o (out), 6 part (scratch of
+part_bytes()), 7 block-table row stride (elements), 8 real requests (requests >= it are graph
+padding: no split, zero output rows, their flags still set), 9 the softmax scale x log2(e) as
+fp32 bits. A fragment's descriptor is row `layer` of a table of TABLE_ROWS rows.
+"""
+
+from types import SimpleNamespace
+
+import flydsl.expr as fx
+from aiter.ops.flydsl.kernels import buffer_ops as bo
+from aiter.ops.flydsl.kernels.kernels_common import LOG2E
+from flydsl.expr import const_expr, gpu, range_constexpr, rocdl
+from flydsl.expr.typing import T
+
+from .dev import (
+    BLOCKS,
+    CM_DEV,
+    CM_NT,
+    OOB_BOUND,
+    THREADS,
+    bf16_pair,
+    hw_exp2,
+    ld_i32,
+    ld_i32_if,
+    ld_vec,
+    ld_vec_if,
+    mb_poll,
+    mb_put,
+    mb_put2,
+    mfma_bf16,
+    permlane_swap,
+    rsrc,
+    traced,
+    uniform,
+    xred,
+)
+
+HKV, QPK, HD, BS, TILE, NQ = 2, 8, 64, 16, 64, 8
+NH = HKV * QPK
+WIN_SWA = 2048  # vLLM's window (2047, 0) -> keys qpos - key < 2048
+KP = 72  # LDS row stride of a K / V tile (bf16)
+KPW = KP // 2
+NEG = float("-inf")
+PGMAX = 512  # page ids a CTA staged in one round (128 tiles)
+# keys a request the page staging covers (max_position_embeddings)
+VLLM_MAX_KEYS = 1 << 20
+REGION_IDS = {"aseg": 40}
+DESC_WORDS = 10
+TABLE_ROWS = 64
+
+# parts of the scratch buffer (bytes); the segment count is always BLOCKS (G x CPG <= 256)
+P_O = 0  # partial o, fp32 [seg][row][64]
+P_ML = BLOCKS * 64 * HD * 4  # (m, l) [seg][row]
+P_SF = P_ML + BLOCKS * 64 * 2 * 4  # one tagged flag a segment
+P_IF = BLOCKS * 64 * (HD + 4) * 8  # one tagged flag an output row
+P_AO = P_IF + 16 * 128 * 8  # attn_o as tagged pairs [row][512 words]
+P_END = P_AO + 128 * 512 * 8
+assert P_SF + BLOCKS * 8 <= P_IF
+
+
+def part_bytes():
+    return P_END
+
+
+def default_cpg(NR, swa):
+    g = NR * HKV
+    if swa:
+        return min(BLOCKS // g, 17)  # 2 tiles a CTA (34 tiles at most in the window)
+    return BLOCKS // g
+
+
+def make(NR, swa, maxk=VLLM_MAX_KEYS):
+    """The fragment for NR requests; swa = SWA (window 2048) or full attention. It stages enough
+    page ids a CTA for maxk keys a request (PGN)."""
+    # K / V load cache policy: cached for SWA, nontemporal for full
+    KCM = 0 if swa else CM_NT
+
+    G = NR * HKV
+    CPG = default_cpg(NR, swa)
+    # steps a CTA at the window / at 32k context
+    NST = -(-(-(-(34 if swa else 512) // CPG)) // 2)
+    PD = max(1, min(4, NST))  # K / V steps in flight
+    assert G * CPG <= BLOCKS
+    NI = NR * NQ * NH  # output rows
+    if NI >= BLOCKS:
+        IPC = NI // BLOCKS
+        STR = 1
+    else:
+        IPC = 1
+        STR = BLOCKS // NI
+    TPI = THREADS // IPC
+    WPI = TPI // 64
+    SL = TPI // 16
+    SPT = -(-CPG // SL)
+    assert IPC * TPI == THREADS and WPI >= 1 and NI <= 16 * 128 and PGMAX == THREADS
+    WIN = WIN_SWA if swa else 0
+    PGN = PGMAX  # page ids staged a CTA (LDS); full attention: maxk keys a request (PGN // THREADS rounds)
+    if not swa:
+        tiles = -(-(-(-maxk // TILE)) // CPG)  # tiles a CTA at maxk keys
+        PGN = max(PGMAX, -(-tiles * (TILE // BS) // THREADS) * THREADS)
+
+    @fx.struct
+    class Lds:
+        kv: fx.Array[fx.Int32, 2 * 2 * 64 * KPW, 16]  # [tile lane][K | V][key][KP bf16]
+        # tile lane 1 -> 0: [quarter][lane][m, l, acc 16]
+        mrg: fx.Array[fx.Float32, 4 * 64 * 18, 16]
+        mg2: fx.Array[fx.Float32, 8 * 16 * 6, 16]  # merge: [wave][d quad][M, L, A 4]
+        # the CTA's page ids, staged once (off the K / V vmcnt chain)
+        pg: fx.Array[fx.Int32, PGN, 16]
+
+    def fmax(a, b):
+        return (a > b).select(a, b)
+
+    def fadd(a, b):
+        return a + b
+
+    def p64(dv, i):
+        lo = uniform(dv[2 * i])
+        hi = uniform(dv[2 * i + 1])
+        return (fx.Int64(hi) << fx.Int64(32)) | (fx.Int64(lo) & fx.Int64(0xFFFFFFFF))
+
+    def desc(dptr):
+        dv = []
+        for i in range(DESC_WORDS // 2):
+            v4 = ld_vec(dptr, 4 * i, 4)
+            dv += [v4[e] for e in range(4)]
+        return SimpleNamespace(
+            q=p64(dv, 0),
+            k=p64(dv, 1),
+            v=p64(dv, 2),
+            pt=p64(dv, 3),
+            used=p64(dv, 4),
+            out=p64(dv, 5),
+            part=p64(dv, 6),
+            pts=uniform(dv[14]),
+            nreq=uniform(dv[16]),
+            qks=uniform(dv[18]).bitcast(fx.Float32),
+        )
+
+    def geom(d, req):
+        """Run-time key-tile range of request req (uniform)."""
+        seq_len = uniform(ld_i32(d.used, req))
+        ctx = seq_len - NQ
+        if const_expr(swa):
+            lo = fx.max(ctx - WIN + 1, 0)
+        else:
+            lo = fx.Int32(0)
+        lo_t = lo // TILE
+        hi_t = (seq_len + TILE - 1) // TILE
+        n_t = hi_t - lo_t
+        tps = fx.max((n_t + CPG - 1) // CPG, 1)
+        nseg = (n_t + tps - 1) // tps
+        return seq_len, ctx, lo_t, hi_t, tps, nseg
+
+    def tr16(kvb, byte):
+        ptr = fx.to_llvm_ptr(fx.inttoptr(kvb.type, fx.Int64(byte)))
+        return fx.Vector(
+            rocdl.ds_read_tr16_b64(T.vec(4, T.bf16), ptr).result, (4,), fx.BFloat16
+        )
+
+    def kv_load(d, t0, t1, seq_len, kvh, step, pgl, ltl, lkey, lch, wlo=None):
+        tile = t0 + 2 * step + ltl
+        kg = tile * TILE + lkey
+        ok = (tile < t1) & (kg < seq_len)
+        pg = fx.ptr_load(
+            pgl + fx.min((2 * step + ltl) * (TILE // BS) + lkey // BS, PGN - 1)
+        )
+        if const_expr(swa):
+            # keys before the window: no access (vLLM frees their blocks)
+            ok = ok & (kg >= wlo)
+        base = d.k + (
+            fx.Int64(uniform(pg)) * d.v + fx.Int64(kvh * (BS * 2 * HD))
+        ) * fx.Int64(2)
+        # dwords: row (kg % 16) of 128 bf16, channels [16 lch, 16 lch + 16)
+        off = (kg % BS) * HD + 8 * lch
+        kv = [
+            ld_vec_if(ok, base, off + o, 4, OOB_BOUND, KCM)
+            for o in (0, 4, HD // 2, HD // 2 + 4)
+        ]
+        return fx.Vector.from_elements(
+            [kv[j][i] for j in range(4) for i in range(4)], fx.Int32
+        )
+
+    def step(
+        d,
+        s,
+        kvc,
+        m,
+        l,
+        acc,
+        t0,
+        t1,
+        seq_len,
+        kvh,
+        pgl,
+        ltl,
+        lkey,
+        lch,
+        wtl,
+        r16,
+        g,
+        qpos,
+        qb,
+        kvl,
+        kvbase,
+        wlo=None,
+    ):
+        """One step (2 tiles): K / V of step s (kvc) -> LDS, the loads of step s + PD issued, then this wave's tile."""
+        gpu.barrier()
+        for kind in range_constexpr(2):
+            words = [kvc[8 * kind + i] for i in range(8)]
+            b0 = ((ltl * 2 + kind) * 64 + lkey) * KPW + 8 * lch
+            fx.ptr_store(fx.Vector.from_elements(words[:4], fx.Int32), kvl + b0)
+            fx.ptr_store(fx.Vector.from_elements(words[4:], fx.Int32), kvl + b0 + 4)
+        gpu.barrier()
+        rocdl.sched_barrier(0)
+        kvn = kv_load(d, t0, t1, seq_len, kvh, s + PD, pgl, ltl, lkey, lch, wlo)
+        rocdl.sched_barrier(0)
+        tile = t0 + 2 * s + wtl
+        kb = tile * TILE
+        tin = tile < t1
+        # S^T = K Q^T: lane holds q row r16, keys 16 n + 4 g + i
+        xs = []
+        mx = fx.Float32(NEG)
+        for n in range_constexpr(4):
+            sv = fx.Vector.filled(4, 0.0, fx.Float32)
+            for c in range_constexpr(2):
+                a = fx.Vector(
+                    fx.ptr_load(
+                        kvl
+                        + ((wtl * 2) * 64 + 16 * n + r16) * KPW
+                        + (32 * c + 8 * g) // 2,
+                        result_type=fx.Vector.make_type(4, fx.Int32),
+                    )
+                ).bitcast(fx.BFloat16)
+                sv = mfma_bf16(a, qb[c], sv)
+            for i in range_constexpr(4):
+                key = kb + 16 * n + 4 * g + i
+                ok = tin & (key <= qpos)
+                if const_expr(swa):
+                    ok = ok & ((qpos - key) < WIN)
+                x = ok.select(sv[i] * d.qks, fx.Float32(NEG))
+                xs.append(x)
+                mx = fmax(mx, x)
+        mx = xred(mx, 16, fmax)
+        mx = xred(mx, 32, fmax)
+        mj = fmax(m, mx)
+        mj = (mj > fx.Float32(NEG)).select(mj, fx.Float32(0.0))
+        ps = [hw_exp2(x - mj) for x in xs]
+        alpha = hw_exp2(m - mj)
+        psum = fx.Float32(0.0)
+        for x in ps:
+            psum = psum + x
+        l = l * alpha + psum
+        acc = [
+            fx.Vector.from_elements([acc[j][i] * alpha for i in range(4)], fx.Float32)
+            for j in range(4)
+        ]
+        m = mj
+        # O^T += V^T P^T
+        for c in range_constexpr(2):
+            pb = fx.Vector.from_elements(
+                [
+                    bf16_pair(ps[(2 * c) * 4 + 0], ps[(2 * c) * 4 + 1]),
+                    bf16_pair(ps[(2 * c) * 4 + 2], ps[(2 * c) * 4 + 3]),
+                    bf16_pair(ps[(2 * c + 1) * 4 + 0], ps[(2 * c + 1) * 4 + 1]),
+                    bf16_pair(ps[(2 * c + 1) * 4 + 2], ps[(2 * c + 1) * 4 + 3]),
+                ],
+                fx.Int32,
+            ).bitcast(fx.BFloat16)
+            for j in range_constexpr(4):
+                row1 = 32 * c + 4 * g + r16 // 4
+                col = 16 * j + 4 * (r16 % 4)
+                vb = kvbase + (((wtl * 2 + 1) * 64 + row1) * KP + col) * 2
+                lo_ = tr16(kvl, vb)
+                hi_ = tr16(kvl, vb + 16 * KP * 2)
+                a = lo_.shuffle(hi_, [0, 1, 2, 3, 4, 5, 6, 7])
+                acc[j] = mfma_bf16(a, pb, acc[j])
+        return m, l, acc, kvn
+
+    @traced
+    def split(d, tag, tid, bid, lane, wave, L):
+        kvl = L.kv.ptr
+        kvbase = fx.Int32(fx.ptrtoint(kvl))
+        grp = bid // CPG
+        jseg = bid % CPG
+        req = grp // HKV
+        kvh = grp % HKV
+        seq_len, ctx, lo_t, hi_t, tps, nseg = geom(d, req)
+        t0 = lo_t + jseg * tps
+        t1 = fx.min(t0 + tps, hi_t)
+        wlo = None
+        if const_expr(swa):
+            wlo = fx.max(ctx - WIN + 1, 0)  # the first key in the window
+        if (bid < G * CPG) & (req < d.nreq):
+            if t0 < t1:
+                nst = (t1 - t0 + 1) // 2
+                # load lanes: tile lane, key, 16 B chunk
+                ltl = tid // 256
+                lkey = (tid % 256) // 4
+                lch = tid % 4
+
+                # this wave: row quarter qr (16 q rows), tile lane wtl
+                qr = wave % 4
+                wtl = wave // 4
+                r16 = lane % 16
+                g = lane // 16
+                m_row = 16 * qr + r16
+                tok = m_row // QPK
+                head = kvh * QPK + m_row % QPK
+                qpos = ctx + tok
+                qb = [
+                    ld_vec(
+                        d.q,
+                        ((NQ * req + tok) * (NH * HD) + head * HD + 32 * c + 8 * g)
+                        // 2,
+                        4,
+                    ).bitcast(fx.BFloat16)
+                    for c in range(2)
+                ]
+                pgl = L.pg.ptr
+                # full attention past 128 tiles a CTA
+                for r in range_constexpr(1, PGN // THREADS):
+                    tr = tid + r * THREADS
+                    pk = t0 * (TILE // BS) + tr
+                    pok = (tr < (t1 - t0) * (TILE // BS)) & (pk * BS < seq_len)
+                    fx.ptr_store(
+                        ld_i32_if(pok, d.pt, req * d.pts + pk, OOB_BOUND), pgl + tr
+                    )
+                pk = t0 * (TILE // BS) + tid  # page tid of the CTA's range
+                pok = (tid < (t1 - t0) * (TILE // BS)) & (pk * BS < seq_len)
+                if const_expr(swa):
+                    # blocks before the window: stale ids, not read
+                    pok = pok & ((pk + 1) * BS > wlo)
+                pgv = ld_i32_if(pok, d.pt, req * d.pts + pk, OOB_BOUND)
+                fx.ptr_store(pgv, pgl + fx.min(tid, PGMAX - 1))
+                gpu.barrier()
+                kvq = []
+                # in step order: the scheduler must not move slot 0 behind the others
+                for i in range_constexpr(PD):
+                    kvq.append(
+                        kv_load(d, t0, t1, seq_len, kvh, i, pgl, ltl, lkey, lch, wlo)
+                    )
+                    rocdl.sched_barrier(0)
+                m = fx.Float32(NEG)
+                l = fx.Float32(0.0)
+                acc = [fx.Vector.filled(4, 0.0, fx.Float32) for _ in range(4)]
+                s = fx.Int32(0)
+                while s < nst:
+                    # PD steps a trip with fixed register slots (a rotated list makes LLVM copy the slots at the back
+                    # edge behind a vmcnt wait); steps past nst are fully masked (no loads, no change of m / l / acc)
+                    for j in range_constexpr(PD):
+                        m, l, acc, kvn = step(
+                            d,
+                            s + j,
+                            kvq[j],
+                            m,
+                            l,
+                            acc,
+                            t0,
+                            t1,
+                            seq_len,
+                            kvh,
+                            pgl,
+                            ltl,
+                            lkey,
+                            lch,
+                            wtl,
+                            r16,
+                            g,
+                            qpos,
+                            qb,
+                            kvl,
+                            kvbase,
+                            wlo,
+                        )
+                        kvq = kvq[:j] + [kvn] + kvq[j + 1 :]
+                    s = s + PD
+                lt = xred(l, 16, fadd)
+                lt = xred(lt, 32, fadd)
+                mrg = L.mrg.ptr
+                mb = (qr * 64 + lane) * 18
+                if wtl == 1:
+                    fx.ptr_store(m, mrg + mb)
+                    fx.ptr_store(lt, mrg + mb + 1)
+                    for j in range_constexpr(4):
+                        for i in range_constexpr(4):
+                            fx.ptr_store(acc[j][i], mrg + mb + 2 + 4 * j + i)
+                gpu.barrier()
+                if wtl == 0:
+                    m1 = fx.ptr_load(mrg + mb)
+                    l1 = fx.ptr_load(mrg + mb + 1)
+                    mm = fmax(m, m1)
+                    a0 = hw_exp2(m - mm)
+                    a1 = hw_exp2(m1 - mm)
+                    lf = lt * a0 + l1 * a1
+                    seg = grp * CPG + jseg
+                    po = d.part + fx.Int64(P_O)
+                    for j in range_constexpr(4):
+                        ov = [
+                            acc[j][i] * a0 + fx.ptr_load(mrg + mb + 2 + 4 * j + i) * a1
+                            for i in range(4)
+                        ]
+                        bo.buffer_store(
+                            fx.Vector.from_elements(
+                                [v_.bitcast(fx.Int32) for v_ in ov], fx.Int32
+                            ),
+                            rsrc(po),
+                            ((seg * 64 + m_row) * HD + 16 * j + 4 * g),
+                            cache_modifier=CM_DEV,
+                        )
+                    if g == 0:
+                        bo.buffer_store(
+                            fx.Vector.from_elements(
+                                [mm.bitcast(fx.Int32), lf.bitcast(fx.Int32)], fx.Int32
+                            ),
+                            rsrc(d.part + fx.Int64(P_ML)),
+                            (seg * 64 + m_row) * 2,
+                            cache_modifier=CM_DEV,
+                        )
+                rocdl.s_waitcnt(vmcnt=0, expcnt=0, lgkmcnt=0)
+                gpu.barrier()
+                if tid == 0:
+                    mb_put(tag, d.part + fx.Int64(P_SF), grp * CPG + jseg, fx.Int32(1))
+
+    def comb(Ma, La, Aa, Mb, Lb, Ab):
+        M = fmax(Ma, Mb)
+        wa = (Ma > fx.Float32(NEG)).select(hw_exp2(Ma - M), fx.Float32(0.0))
+        wb = (Mb > fx.Float32(NEG)).select(hw_exp2(Mb - M), fx.Float32(0.0))
+        return M, La * wa + Lb * wb, [Aa[i] * wa + Ab[i] * wb for i in range(4)]
+
+    def swap(off, v):
+        a, b = permlane_swap(off, v.bitcast(fx.Int32), v.bitcast(fx.Int32))
+        return a.bitcast(fx.Float32), b.bitcast(fx.Float32)
+
+    @traced
+    def merge(d, tag, health, tid, bid, lane, wave, L):
+        if const_expr(STR > 1):
+            has = (bid % STR) == 0
+            item = bid // STR
+        else:
+            has = bid >= 0
+            item = bid * IPC + tid // TPI
+        if has:
+            itl = tid // TPI
+            tt = tid % TPI
+            sl = tt // 16
+            dq = tt % 16
+            grp = item // 64
+            row = item % 64
+            req = grp // HKV
+            kvh = grp % HKV
+            real = req < d.nreq
+            seq_len, ctx, lo_t, hi_t, tps, nseg = geom(d, fx.min(req, d.nreq - 1))
+            mls, os_ = [], []
+            sf = d.part + fx.Int64(P_SF)
+            specs, oks = [], []
+            for k in range_constexpr(SPT):
+                sg = sl + SL * k
+                specs.append((sf, grp * CPG + sg, 1))
+                oks.append((sg < nseg) & real)
+            mb_poll(tag, health, specs, REGION_IDS["aseg"], oks=oks)
+            for k in range_constexpr(SPT):
+                seg = grp * CPG + sl + SL * k
+                mls.append(
+                    ld_vec_if(
+                        oks[k],
+                        d.part + fx.Int64(P_ML),
+                        (seg * 64 + row) * 2,
+                        2,
+                        P_SF - P_ML,
+                        CM_DEV,
+                    )
+                )
+                os_.append(
+                    ld_vec_if(
+                        oks[k],
+                        d.part + fx.Int64(P_O),
+                        (seg * 64 + row) * HD + 4 * dq,
+                        4,
+                        P_ML,
+                        CM_DEV,
+                    )
+                )
+            M = fx.Float32(NEG)
+            for k in range_constexpr(SPT):
+                M = oks[k].select(fmax(M, mls[k][0].bitcast(fx.Float32)), M)
+            Ls = fx.Float32(0.0)
+            A = [fx.Float32(0.0) for _ in range(4)]
+            for k in range_constexpr(SPT):
+                w = oks[k].select(
+                    hw_exp2(mls[k][0].bitcast(fx.Float32) - M), fx.Float32(0.0)
+                )
+                Ls = Ls + w * mls[k][1].bitcast(fx.Float32)
+                A = [A[i] + w * os_[k][i].bitcast(fx.Float32) for i in range(4)]
+            for off in (16, 32):
+                ma, mb_ = swap(off, M)
+                la, lb = swap(off, Ls)
+                aa, ab = [], []
+                for i in range(4):
+                    x_, y_ = swap(off, A[i])
+                    aa.append(x_)
+                    ab.append(y_)
+                M, Ls, A = comb(ma, la, aa, mb_, lb, ab)
+            g2 = L.mg2.ptr
+            wi = tt // 64
+            if lane < 16:
+                b = ((itl * WPI + wi) * 16 + dq) * 6
+                fx.ptr_store(M, g2 + b)
+                fx.ptr_store(Ls, g2 + b + 1)
+                for i in range_constexpr(4):
+                    fx.ptr_store(A[i], g2 + b + 2 + i)
+            gpu.barrier()
+            if tt < 16:
+                b0 = ((itl * WPI) * 16 + dq) * 6
+                M = fx.ptr_load(g2 + b0)
+                Ls = fx.ptr_load(g2 + b0 + 1)
+                A = [fx.ptr_load(g2 + b0 + 2 + i) for i in range(4)]
+                for w_ in range_constexpr(1, WPI):
+                    b = ((itl * WPI + w_) * 16 + dq) * 6
+                    M, Ls, A = comb(
+                        M,
+                        Ls,
+                        A,
+                        fx.ptr_load(g2 + b),
+                        fx.ptr_load(g2 + b + 1),
+                        [fx.ptr_load(g2 + b + 2 + i) for i in range(4)],
+                    )
+                f = fx.Float32(1.0) / (Ls > fx.Float32(0.0)).select(Ls, fx.Float32(1.0))
+                tok = row // QPK
+                head = kvh * QPK + row % QPK
+                w2 = fx.Vector.from_elements(
+                    [bf16_pair(A[0] * f, A[1] * f), bf16_pair(A[2] * f, A[3] * f)],
+                    fx.Int32,
+                )
+                bo.buffer_store(
+                    w2,
+                    rsrc(d.out),
+                    ((NQ * req + tok) * (NH * HD) + head * HD + 4 * dq) // 2,
+                    cache_modifier=CM_DEV,
+                )
+                mb_put2(
+                    tag,
+                    d.part + fx.Int64(P_AO),
+                    ((NQ * req + tok) * (NH * HD) + head * HD + 4 * dq) // 2,
+                    w2[0],
+                    w2[1],
+                )
+                rocdl.s_waitcnt(vmcnt=0, expcnt=0, lgkmcnt=0)
+                if tt == 0:
+                    mb_put(tag, d.part + fx.Int64(P_IF), item, fx.Int32(1))
+
+    def run(dptr, tag, health, tid, bid, lane, wave, L):
+        """The whole fragment; returns (the attn_o pointer, the tagged attn_o rows)."""
+        d = desc(dptr)
+        split(d, tag, tid, bid, lane, wave, L)
+        merge(d, tag, health, tid, bid, lane, wave, L)
+        return d.out, d.part + fx.Int64(P_AO)
+
+    return SimpleNamespace(
+        Lds=Lds,
+        run=run,
+        key=f"attnfrag_nr{NR}_{'swa' if swa else 'full'}_c{CPG}_p{PD}_k{KCM}_{_SRC_HASH}",
+    )
+
+
+with open(__file__, "rb") as _f:
+    import hashlib as _hl
+
+    _SRC_HASH = _hl.sha1(_f.read()).hexdigest()[:12]
+
+
+def vllm_scale_word(scale):
+    """Descriptor word 9: scale x log2(e) as fp32 bits (an int64 word)."""
+    import struct
+
+    return struct.unpack("<i", struct.pack("<f", float(scale) * LOG2E))[0]
+
+
+def vllm_desc(q, kv_cache, block_table, seq_lens, out, part, nreq, scale):
+    """The 10 descriptor words of one layer: q [8 nreq.., 16 x 64] bf16 rows,
+    kv_cache the layer's `[block][2][16][128]` bf16 view (block stride = kv_cache.stride(0)), block_table int32
+    `[>= nreq][stride]` of the layer's group, seq_lens int32 [>= nreq], out attn_o [rows, 1024] bf16, part the scratch
+    (int32, part_bytes()), nreq the requests of the step (graph padding included: seq_len 0 = no keys, zero rows)."""
+    import torch
+
+    assert (
+        q.dtype == torch.bfloat16
+        and q.stride(-1) == 1
+        and q.shape[-1] * (q.shape[-2] if q.dim() == 3 else 1) % (NH * HD) == 0
+    )
+    assert kv_cache.dtype == torch.bfloat16 and tuple(kv_cache.shape[1:]) == (
+        HKV,
+        BS,
+        2 * HD,
+    ), kv_cache.shape
+    assert tuple(kv_cache.stride()[1:]) == (BS * 2 * HD, 2 * HD, 1), kv_cache.stride()
+    assert (
+        block_table.dtype == torch.int32
+        and block_table.stride(1) == 1
+        and block_table.shape[0] >= nreq
+    )
+    assert (
+        seq_lens.dtype == torch.int32
+        and seq_lens.is_contiguous()
+        and seq_lens.shape[0] >= nreq
+    )
+    assert (
+        out.dtype == torch.bfloat16
+        and out.is_contiguous()
+        and part.numel() * part.element_size() >= part_bytes()
+    )
+    return [
+        q.data_ptr(),
+        kv_cache.data_ptr(),
+        kv_cache.stride(0),
+        block_table.data_ptr(),
+        seq_lens.data_ptr(),
+        out.data_ptr(),
+        part.data_ptr(),
+        block_table.stride(0),
+        nreq,
+        vllm_scale_word(scale),
+    ]
diff --git a/vllm/models/muse_spark/amd/mono/kernels/dev.py b/vllm/models/muse_spark/amd/mono/kernels/dev.py
new file mode 100644
index 000000000000..1339a247f88b
--- /dev/null
+++ b/vllm/models/muse_spark/amd/mono/kernels/dev.py
@@ -0,0 +1,429 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# mypy: ignore-errors
+# ruff: noqa: E501
+
+# Vendored from AMD's Muse Spark FlyDSL monokernel kernels (mk/dev.py):
+# SPDX-License-Identifier: MIT
+# Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+"""Device-side primitives of the Muse Spark monokernel.
+
+Adapted from ATOM's ``atom/mono/device`` (ops.py, sync.py; MIT, AMD). Differences
+from ATOM that matter:
+
+* Mailbox tags are epoch based (``tag = epoch * 64 + slot + 1`` with the epoch a
+  device word bumped once per forward inside the CUDA graph) instead of ATOM's
+  zero-at-step-start discipline: nothing is ever zeroed.
+* Every spin is bounded by a wall-clock timeout (s_memrealtime, 100 MHz) that
+  writes a diagnostic record and sets a sticky FAILED word, so a bug turns into
+  a report instead of a hang.
+"""
+
+from __future__ import annotations
+
+import flydsl.expr as fx
+from aiter.ops.flydsl.kernels import buffer_ops as bo
+from aiter.ops.flydsl.kernels.dpp_utils import update_dpp_i32
+from flydsl._mlir.dialects import llvm
+from flydsl.compiler.ast_rewriter import ASTRewriter
+from flydsl.expr import range_constexpr, rocdl
+from flydsl.expr.typing import T, as_ir_value
+
+BLOCKS = 256
+THREADS = 512
+WAVES = THREADS // 64
+
+# gfx950 cache-policy bits (LLVM CPol): SC0 = 1, SC1 = 16, NT = 2.
+CM_DEV = 16  # device scope: past the per-XCD L2
+CM_SYS = 17  # system scope
+CM_NT = 2  # nontemporal: weights streamed once
+
+
+def traced(fn):
+    """FlyDSL AST rewriting for a helper (``if`` / ``while`` on traced values)."""
+    return ASTRewriter.transform(fn)
+
+
+def rsrc(addr, nbytes=None):
+    return bo.create_buffer_resource_from_addr(addr, num_records_bytes=nbytes)
+
+
+def memrealtime():
+    return fx.Int64(llvm.call_intrinsic(T.i64, "llvm.amdgcn.s.memrealtime", [], [], []))
+
+
+def uniform(v):
+    return fx.Int32(rocdl.readfirstlane(T.i32, fx.Int32(v).ir_value()))
+
+
+def readlane(v, src_lane):
+    bits = v.bitcast(fx.Int32) if isinstance(v, fx.Float32) else fx.Int32(v)
+    r = fx.Int32(rocdl.readlane(T.i32, bits, fx.Int32(src_lane)))
+    return r.bitcast(fx.Float32) if isinstance(v, fx.Float32) else r
+
+
+def fresh(v):
+    """An opaque v_mov: stops LLVM hoisting / CSE of lane math into long-lived
+    registers (ATOM ops.fresh)."""
+    return fx.Int32(
+        llvm.inline_asm(
+            T.i32, [fx.Int32(v).ir_value()], "v_mov_b32 $0, $1", "=v,v",
+            has_side_effects=True
+        )
+    )  # fmt: skip
+
+
+def _hw_f32(name, x):
+    return fx.Float32(
+        llvm.call_intrinsic(T.f32, name, [fx.Float32(x).ir_value()], [], [])
+    )
+
+
+def hw_rsq(x):
+    return _hw_f32("llvm.amdgcn.rsq.f32", x)
+
+
+def hw_exp2(x):
+    return _hw_f32("llvm.amdgcn.exp2.f32", x)
+
+
+def hw_sqrt(x):
+    return _hw_f32("llvm.amdgcn.sqrt.f32", x)
+
+
+def permlane_swap(off, x, y):
+    swap = rocdl.permlane32_swap if off == 32 else rocdl.permlane16_swap
+    pr = swap(
+        llvm.StructType.get_literal([T.i32, T.i32]),
+        as_ir_value(fx.Int32(x)),
+        as_ir_value(fx.Int32(y)),
+        False,
+        False,
+    )
+    return tuple(fx.Int32(llvm.extractvalue(T.i32, pr, [j])) for j in range(2))
+
+
+def xshfl(v, off):
+    assert off in (1, 2, 4, 8, 16, 32), off
+    if off >= 16:
+        return v.shuffle_xor(off, 64)
+    is_f = isinstance(v, fx.Float32)
+    x = v.bitcast(fx.Int32) if is_f else fx.Int32(v)
+    if off == 8:
+        y = fx.Int32(update_dpp_i32(x, x, 0x118, 0xF, 0xC, False))
+        y = fx.Int32(update_dpp_i32(y, x, 0x108, 0xF, 0x3, False))
+    elif off == 4:
+        y = fx.Int32(update_dpp_i32(x, x, 0x114, 0xF, 0xA, False))
+        y = fx.Int32(update_dpp_i32(y, x, 0x104, 0xF, 0x5, False))
+    elif off == 2:
+        y = fx.Int32(update_dpp_i32(x, x, 0x4E, 0xF, 0xF, False))
+    else:
+        y = fx.Int32(update_dpp_i32(x, x, 0xB1, 0xF, 0xF, False))
+    return y.bitcast(fx.Float32) if is_f else y
+
+
+def xred(v, off, op):
+    if off < 16:
+        return op(v, xshfl(v, off))
+    is_f = isinstance(v, fx.Float32)
+    x = v.bitcast(fx.Int32) if is_f else fx.Int32(v)
+    a, b = permlane_swap(off, x, x)
+    if is_f:
+        return op(a.bitcast(fx.Float32), b.bitcast(fx.Float32))
+    return op(type(v)(a), type(v)(b))
+
+
+def _add(a, b):
+    return a + b
+
+
+def butterfly(v, offsets, op=_add):
+    for off in offsets:
+        v = xred(v, off, op)
+    return v
+
+
+def wave_sum(v):
+    return butterfly(v, (32, 16, 8, 4, 2, 1))
+
+
+def ballot(pred):
+    return fx.Int64(rocdl.ballot(T.i64, as_ir_value(pred)))
+
+
+def popcount(x):
+    return fx.Int32(fx.ctpop(x))
+
+
+def lanes_below(bal, lane):
+    return popcount(bal & ((fx.Int64(1) << fx.Int64(lane)) - 1))
+
+
+def _xpartner(v, off, lane):
+    if off < 16:
+        return xshfl(v, off)
+    a, b = permlane_swap(off, v, v)
+    return ((lane & off) != 0).select(a, b)
+
+
+def wave_incl_scan(v, lane):
+    blk = v
+    for off in (1, 2, 4, 8, 16, 32):
+        o = _xpartner(blk, off, lane)
+        v = v + ((lane & off) != 0).select(o, 0)
+        blk = blk + o
+    return v
+
+
+def bpermute(v, src_lane):
+    """Lane ``src_lane``'s i32 v (ds_bpermute: the LDS crossbar, no LDS memory)."""
+    return fx.Int32(
+        rocdl.ds_bpermute(T.i32, fx.Int32(src_lane) * 4, fx.Int32(v).ir_value())
+    )
+
+
+def mfma_bf16(a, b, c):
+    """16x16x32 bf16 MFMA: a, b bf16x8 vectors, c f32x4."""
+    return fx.Vector(rocdl.mfma_f32_16x16x32_bf16(T.vec(4, T.f32), [a, b, c]))
+
+
+def bf16_round(a):
+    return fx.Float32(fx.Float32(a).to(fx.BFloat16))
+
+
+def bf16_pair(a, b):
+    """Two f32 -> one i32 word holding (bf16(a) low, bf16(b) high), RNE."""
+    return (
+        fx.Vector.from_elements([a, b], fx.Float32).to(fx.BFloat16).bitcast(fx.Int32)[0]
+    )
+
+
+def bf_lo(w):
+    return (w << 16).bitcast(fx.Float32)
+
+
+def bf_hi(w):
+    return (w & fx.Int32(-65536)).bitcast(fx.Float32)
+
+
+def fp4x8_bf16(word, scale_f32):
+    """One dword of 8 E2M1 (low nibble first) times ``scale_f32`` -> bf16x8
+    (exact: v_cvt_scalef32_pk_bf16_fp4, the aiter Muse a16w4 decode)."""
+    parts = []
+    for sel in range_constexpr(4):
+        pair = rocdl.cvt_scalef32_pk_bf16_fp4(
+            T.vec(2, T.bf16), as_ir_value(fx.Int32(word)), as_ir_value(scale_f32), sel
+        )
+        parts.append(fx.Vector(pair).bitcast(fx.Int32)[0])
+    return fx.Vector.from_elements(parts, fx.Int32).bitcast(fx.BFloat16)
+
+
+def e8m0_f32(code):
+    """E8M0 biased exponent (an i32 holding 1..254) -> the f32 power of two."""
+    return (fx.Int32(code) << 23).bitcast(fx.Float32)
+
+
+def ld_i32(ptr, i, cm=0):
+    return fx.Int32(
+        bo.buffer_load(rsrc(ptr), i, vec_width=1, dtype=T.i32, cache_modifier=cm)
+    )
+
+
+def ld_f32(ptr, i, cm=0):
+    return fx.Float32(
+        bo.buffer_load(rsrc(ptr), i, vec_width=1, dtype=T.f32, cache_modifier=cm)
+    )
+
+
+def ld_vec(ptr, i, words, cm=0, nbytes=None):
+    return fx.Vector(
+        bo.buffer_load(
+            rsrc(ptr, nbytes), i, vec_width=words, dtype=T.i32, cache_modifier=cm
+        )
+    )
+
+
+OOB_DW = 0x3FFFFFFF  # a dword index past any buffer's num_records: the load returns 0, no memory access
+OOB_BOUND = 0x80000000  # num_records of a mailbox resource: every real offset is below it, OOB_DW*4 above
+
+
+def ld_vec_if(ok, ptr, i, words, nbytes, cm=0):
+    """``words`` dwords at dword ``i``, or zeros without a memory access when not ``ok``
+    (out-of-range buffer offset; ``nbytes`` must bound the buffer)."""
+    return fx.Vector(
+        bo.buffer_load(
+            rsrc(ptr, nbytes),
+            ok.select(i, fx.Int32(OOB_DW)),
+            vec_width=words,
+            dtype=T.i32,
+            cache_modifier=cm,
+        )
+    )
+
+
+def ld_i32_if(ok, ptr, i, nbytes, cm=0):
+    return fx.Int32(
+        bo.buffer_load(
+            rsrc(ptr, nbytes),
+            ok.select(i, fx.Int32(OOB_DW)),
+            vec_width=1,
+            dtype=T.i32,
+            cache_modifier=cm,
+        )
+    )
+
+
+# ------------------------------------------------------------------ mailbox
+DIAG_WORDS = 8
+TIMEOUT_TICKS = 2 * 100_000_000  # 2 s at 100 MHz
+
+
+@traced
+def _spin_bounded(load_all, pending, health, region_id, index, tag):
+    """Re-load until ``pending`` clears or TIMEOUT_TICKS pass; on a timeout the
+    CTA's first record goes to ``health`` (word 0 = sticky FAILED, words
+    8 + 8 bid .. = record) and later waits of this launch exit at once."""
+    t0 = memrealtime()
+    v = load_all()
+    late = fx.Int32(0)
+    failed = ld_i32(health, 0, CM_DEV)
+    while pending(v) & (late == 0) & (failed == 0):
+        rocdl.s_sleep(1)
+        v = load_all()
+        late = (memrealtime() - t0 > TIMEOUT_TICKS).select(fx.Int32(1), fx.Int32(0))
+    if late != 0:
+        bid = fx.block_idx.x
+        at = fx.Int32(memrealtime() & 0x7FFFFFFF)
+        record = [1, region_id, index, tag, v[1], bid, at, 0]
+        for h in range_constexpr(2):
+            bo.buffer_store(
+                fx.Vector.from_elements(
+                    [fx.Int32(w) for w in record[4 * h : 4 * h + 4]], fx.Int32
+                ),
+                rsrc(health),
+                DIAG_WORDS + bid * DIAG_WORDS + 4 * h,
+                cache_modifier=CM_DEV,
+            )
+        bo.buffer_store(fx.Int32(1), rsrc(health), 0, cache_modifier=CM_DEV)
+    return v
+
+
+# Mailbox helpers are plain functions, not methods: FlyDSL's region
+# analyzer carries every local whose method is called inside an ``if`` / ``while``
+# as scf state, which an object of Python callables cannot be.
+#
+# ``tag`` is a traced Int32 (epoch-based, unique per launch and stage), so a pair
+# from an earlier launch never matches; no region is ever zeroed.
+
+
+def mb_put(tag, base, i, v, cm=CM_DEV):
+    """Pair i := (v, tag), device scope (CM_SYS: system scope, a peer GPU's buffer)."""
+    bits = v.bitcast(fx.Int32) if isinstance(v, fx.Float32) else fx.Int32(v)
+    bo.buffer_store(
+        fx.Vector.from_elements([bits, tag], fx.Int32),
+        rsrc(base),
+        i * 2,
+        cache_modifier=cm,
+    )
+
+
+def mb_put_sys(tag, base, i, v):
+    mb_put(tag, base, i, v, CM_SYS)
+
+
+def mb_put2(tag, base, i, w0, w1):
+    """Pairs i, i+1 := (w0, tag), (w1, tag): one 16 B store (i even)."""
+    bo.buffer_store(
+        fx.Vector.from_elements([fx.Int32(w0), tag, fx.Int32(w1), tag], fx.Int32),
+        rsrc(base),
+        i * 2,
+        cache_modifier=CM_DEV,
+    )
+
+
+def mb_poll(tag, health, specs, region_id=0, batch=16, oks=None, cm=CM_DEV):
+    """[(base, pair index, npairs in {1, 2})] -> value words, all re-loaded
+    together until every tag matches (one round trip after the last lands).
+    ``oks`` (optional, one traced bool a spec): a false spec is not loaded (an
+    out-of-range buffer offset, no memory access), returns zeros, and is not
+    waited for."""
+    if len(specs) == 0:
+        return []
+    if len(specs) > batch:
+        o1 = None if oks is None else oks[:batch]
+        o2 = None if oks is None else oks[batch:]
+        return mb_poll(tag, health, specs[:batch], region_id, batch, o1, cm) + mb_poll(
+            tag, health, specs[batch:], region_id, batch, o2, cm
+        )
+
+    def load_all():
+        words = []
+        for si, (b, i, n) in enumerate(specs):
+            idx = fx.Int32(i) * 2
+            if oks is not None:
+                idx = oks[si].select(idx, fx.Int32(OOB_DW))
+            w = fx.Vector(
+                bo.buffer_load(
+                    rsrc(b, OOB_BOUND),
+                    idx,
+                    vec_width=2 * n,
+                    dtype=T.i32,
+                    cache_modifier=cm,
+                )
+            )
+            words += [w[e] for e in range(2 * n)]
+        return fx.Vector.from_elements(words, fx.Int32)
+
+    def pending(v):
+        bad = None
+        e = 0
+        for si, (_, _, n) in enumerate(specs):
+            for q in range(n):
+                b1 = v[e + 2 * q + 1] != tag
+                if oks is not None:
+                    b1 = b1 & oks[si]
+                bad = b1 if bad is None else (bad | b1)
+            e += 2 * n
+        return bad
+
+    v = _spin_bounded(load_all, pending, health, region_id, specs[0][1], tag)
+    outs, e = [], 0
+    for _, _, n in specs:
+        outs.append([v[e + 2 * q] for q in range(n)])
+        e += 2 * n
+    return outs
+
+
+def st_if(ok, v, base, i, cm=0):
+    """buffer_store of v (i32 or i32 vector) at dword i, dropped (out-of-range offset, no
+    branch) when not ``ok``: keeps the waitcnt pass from draining at a join."""
+    bo.buffer_store(
+        v,
+        rsrc(base, OOB_BOUND),
+        ok.select(fx.Int32(i), fx.Int32(OOB_DW)),
+        cache_modifier=cm,
+    )
+
+
+def mb_put_if(ok, tag, base, i, v, cm=CM_DEV):
+    """Pair i := (v, tag) when ``ok`` (branch-free, see st_if)."""
+    bits = v.bitcast(fx.Int32) if isinstance(v, fx.Float32) else fx.Int32(v)
+    st_if(ok, fx.Vector.from_elements([bits, tag], fx.Int32), base, fx.Int32(i) * 2, cm)
diff --git a/vllm/models/muse_spark/amd/mono/kernels/k3.py b/vllm/models/muse_spark/amd/mono/kernels/k3.py
new file mode 100644
index 000000000000..55ebb4f56329
--- /dev/null
+++ b/vllm/models/muse_spark/amd/mono/kernels/k3.py
@@ -0,0 +1,2522 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# mypy: ignore-errors
+# ruff: noqa: B008, E501
+
+# Vendored from AMD's Muse Spark FlyDSL monokernel kernels (mk/k3.py):
+# SPDX-License-Identifier: MIT
+# Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+"""K3: one Muse Spark decoder layer for S = 8 / 16 rows as one persistent launch.
+
+One rank of a TP = EP = 8 group (rank r owns experts [32 r, 32 r + 32)), 256 CTAs x 512
+threads, everything between two attention cores:
+
+    [attention]  optional: this layer's verify attention (attn_frag.py) at the head
+    O      o-norm x sigmoid(gate) -> o_proj rows 32 b .. over this rank's K slice; the bf16
+           partial goes to the owner rank of the column block (b // 32): reduce-scatter
+    AB     the owner sums the 8 partials in rank order and all-gathers the reduced o and its
+           norm partials; every CTA: post-attention norm, gated residual on the fp32 stream,
+           pre-FFN norm -> ffn_in (one token-reduction round for both norms)
+    A      router logits (fp32) + pre_expert_proj of this rank's shard, K in 8 parts over
+           8-CTA groups; logit partials and h_raw go to every rank's peer buffer
+    route  every CTA routes every token (identical top-8), the local experts' picks
+    C, D   gate_up / down of the local experts: MXFP4 weights in vLLM's shuffled layouts,
+           decoded in registers, bf16 MFMA
+    E      per-pick post-expert norm and weighted sum -> the EP partial, reduce-scattered to
+           the block owner, summed in rank order, all-gathered
+    F      post_expert_proj (this rank's 1024 rows, K in 4 parts) -> every rank
+    G      post-FFN norm, gated residual, the next layer's input norm -> stream (fp32) and
+           the next attention input (bf16)
+    QKV    the next layer's qkv + gate GEMV (K in eighths); the owners finish whole heads:
+           qk-norm, RoPE (next_rope), q out, bf16 K / V into vLLM's LBHNC KV cache
+
+Peer regions are double buffered by tag parity; tag = epoch * 64 + layer + 1 with the
+epoch bumped in lockstep by every rank once per forward.
+"""
+
+from dataclasses import dataclass
+
+import flydsl.compiler as flyc
+import flydsl.expr as fx
+from aiter.ops.flydsl.kernels import buffer_ops as bo
+from aiter.ops.flydsl.kernels.kernels_common import LOG2E
+from flydsl.expr import const_expr, gpu, range_constexpr, rocdl
+from flydsl.expr.typing import Int32, Int64
+
+from .dev import (
+    BLOCKS,
+    CM_DEV,
+    CM_NT,
+    CM_SYS,
+    THREADS,
+    WAVES,
+    ballot,
+    bf16_pair,
+    bf16_round,
+    bf_hi,
+    bf_lo,
+    bpermute,
+    butterfly,
+    e8m0_f32,
+    fp4x8_bf16,
+    fresh,
+    hw_exp2,
+    hw_rsq,
+    hw_sqrt,
+    lanes_below,
+    ld_f32,
+    ld_i32,
+    ld_i32_if,
+    ld_vec,
+    ld_vec_if,
+    mb_poll,
+    mb_put,
+    mb_put_sys,
+    mfma_bf16,
+    permlane_swap,
+    popcount,
+    readlane,
+    rsrc,
+    st_if,
+    traced,
+    uniform,
+    wave_incl_scan,
+    wave_sum,
+)
+
+HID = 8192
+DM = 4096
+INTER = 4096
+E = 256
+TOPK = 8
+W = 8  # ranks in the TP = EP group
+EL = E // W  # local experts
+ROUTE_EPS = 1e-15
+PRE_EPS = 1e-5
+POST_EPS = 1e-8
+XPAD = 8
+HROW = DM + XPAD
+AROW = INTER + XPAD
+KCH = 128
+FCH = 256
+W2_BYTES = EL * DM * (INTER // 2)
+W2S_BYTES = EL * DM * (INTER // 32)
+KQ = 8  # stage-A K split (CTAs a group)
+NG = BLOCKS // KQ  # stage-A groups
+XQ = HID // KQ  # stage-A K a CTA
+XQROW = XQ + XPAD
+TL_POINTS = 16
+REGION_IDS = {
+    "logit": 1,
+    "hf": 2,
+    "a": 3,
+    "z": 4,
+    "ss": 5,
+    "px": 6,
+    "moep": 7,
+    "mf": 8,
+    "py": 9,
+    "y": 10,
+}
+HEALTH_WORDS = 8 + 8 * BLOCKS
+HQ, HKV, HD = 128 // W, 16 // W, 64  # q / kv heads a rank
+QROWS, KROWS = HQ * HD, HKV * HD
+QKV_ROWS = QROWS + 2 * KROWS  # qkvg rows a rank: q 1024 | k 128 | v 128 | gate 1024
+QKVG_ROWS = QKV_ROWS + QROWS
+OK_K = HID // W  # o_proj K a rank
+NQTILES = QKVG_ROWS // 16  # qkvg 16-row tiles a rank
+# qkvg tiles a CTA group (32 groups x 8 K eighths; the last group padded)
+NQT = -(-NQTILES // 32)
+
+
+def bf_tile_load_mma(wbase, row16, ldk, k, lane, cm=CM_NT):
+    """16 rows (row16 ..) x 128 K (k ..) of a row-major bf16 [*, ldk] weight, loaded straight in the
+    16x16x32 bf16 MFMA A-operand layout: element qq (of 4) is the operand of MFMA step qq (lane l = row
+    l % 16, K 32 qq + 8 (l / 16) .. + 8), 16 rows x 64 B an instruction. The same operand values as
+    bf_tile_load + bf_mma_operand (bit-identical results) without the 8 ds_bpermute a step."""
+    rowm = row16 + lane % 16
+    kk = k + 8 * (lane // 16)
+    return [ld_vec(wbase, (rowm * ldk + kk + 32 * qq) // 2, 4, cm) for qq in range(4)]
+
+
+def bf_mma_operand(ops, qq, lane):
+    """Coalesced bf16 weight registers (4 instructions: rows 8 blk + lane / 8, 16 B at
+    128 h + 16 (lane % 8)) -> the MFMA A operand of step qq (K 32 qq + 8 g of row r)."""
+    r = lane % 16
+    g = lane // 16
+    h = qq // 2
+    src = 8 * (r % 8) + 4 * (qq % 2) + g
+    hi = r >= 8
+    ws = []
+    for jj in range_constexpr(4):
+        pa = bpermute(ops[h][jj], src)
+        pb = bpermute(ops[2 + h][jj], src)
+        ws.append(hi.select(pb, pa))
+    return fx.Vector.from_elements(ws, fx.Int32)
+
+
+def bf_tile_load(wbase, row16, ldk, k, lane):
+    """16 rows (row16 ..) x 128 K (k ..) of a row-major bf16 [*, ldk] weight, coalesced."""
+    out = []
+    for blk in range_constexpr(2):
+        rowc = row16 + 8 * blk + lane // 8
+        for h in range_constexpr(2):
+            out.append(
+                ld_vec(wbase, (rowc * ldk + k + 64 * h + 8 * (lane % 8)) // 2, 4, CM_NT)
+            )
+    return out
+
+
+def okey(v):
+    b = v.bitcast(fx.Int32)
+    neg = b < fx.Int32(0)
+    return fx.Uint32(neg.select(~b, b ^ fx.Int32(-2147483648)))
+
+
+@traced
+def route_token(tkns, bias_v, lgss, tid, lane, rt):
+    """Tokens ``tkns`` (this wave's, all valid), interleaved for ILP: sigmoid +
+    bias, radix top-8, rank, renorm -> rid / rw."""
+    T_ = len(tkns)
+    keys, sigs = [], []
+    for x_ in range_constexpr(T_):
+        kx, sx = [], []
+        for i in range_constexpr(4):
+            lg = lgss[x_][i].bitcast(fx.Float32)
+            sg = fx.Float32(1.0) / (fx.Float32(1.0) + hw_exp2(lg * fx.Float32(-LOG2E)))
+            sc = sg + bias_v[i]
+            live = sc == sc  # NaN ranks last
+            kx.append(okey(live.select(sc, fx.Float32(-3.0e38))))
+            sx.append((sg == sg).select(sg, fx.Float32(0.0)))
+        keys.append(kx)
+        sigs.append(sx)
+    th = [fx.Uint32(0) for _ in range(T_)]
+    for bit in range_constexpr(31, -1, -1):
+        for x_ in range_constexpr(T_):
+            cand = th[x_] | fx.Uint32(1 << bit)
+            n = fx.Int32(0)
+            for i in range_constexpr(4):
+                n = n + popcount(ballot(keys[x_][i] >= cand))
+            th[x_] = (n >= TOPK).select(cand, th[x_])
+    for x_ in range_constexpr(T_):
+        tkn = tkns[x_]
+        n_gt = fx.Int32(0)
+        for i in range_constexpr(4):
+            n_gt = n_gt + popcount(ballot(keys[x_][i] > th[x_]))
+        need = fx.Int32(TOPK) - n_gt
+        eq_before = fx.Int32(0)
+        sel_before = fx.Int32(0)
+        sel_base = rt["sel"] + tkn * (TOPK * 3)
+        for i in range_constexpr(4):
+            eqb = ballot(keys[x_][i] == th[x_])
+            eq_rank = eq_before + lanes_below(eqb, lane)
+            sel = (keys[x_][i] > th[x_]) | ((keys[x_][i] == th[x_]) & (eq_rank < need))
+            selb = ballot(sel)
+            pos = sel_before + lanes_below(selb, lane)
+            if sel:
+                fx.ptr_store(keys[x_][i].bitcast(fx.Int32), sel_base + pos * 3)
+                fx.ptr_store(fx.Int32(lane + 64 * i), sel_base + (pos * 3 + 1))
+                fx.ptr_store(sigs[x_][i].bitcast(fx.Int32), sel_base + (pos * 3 + 2))
+            eq_before = eq_before + popcount(eqb)
+            sel_before = sel_before + popcount(selb)
+    q = fx.min(lane, TOPK - 1)
+    mks, mes, mss, ranks = [], [], [], []
+    for x_ in range_constexpr(T_):
+        sel_base = rt["sel"] + tkns[x_] * (TOPK * 3)
+        mk = fx.Uint32(fx.ptr_load(sel_base + q * 3))
+        me = fx.ptr_load(sel_base + (q * 3 + 1))
+        ms = fx.ptr_load(sel_base + (q * 3 + 2)).bitcast(fx.Float32)
+        rank_ = fx.Int32(0)
+        for j in range_constexpr(TOPK):
+            ok = fx.Uint32(fx.ptr_load(sel_base + j * 3))
+            oe = fx.ptr_load(sel_base + (j * 3 + 1))
+            before = (ok > mk) | ((ok == mk) & (oe < me))
+            rank_ = rank_ + before.select(fx.Int32(1), fx.Int32(0))
+        mks.append(mk)
+        mes.append(me)
+        mss.append(ms)
+        ranks.append(rank_)
+    # serial sum of the 8 sigmoids in rank order (+ eps), IEEE division
+    for x_ in range_constexpr(T_):
+        if lane < TOPK:
+            fx.ptr_store(mss[x_], rt["rw"] + (tkns[x_] * TOPK + ranks[x_]))
+    for x_ in range_constexpr(T_):
+        tkn = tkns[x_]
+        total = fx.Float32(0.0)
+        for j in range_constexpr(TOPK):
+            total = total + fx.ptr_load(rt["rw"] + (tkn * TOPK + j))
+        total = total + fx.Float32(ROUTE_EPS)
+        if lane < TOPK:
+            fx.ptr_store(
+                fx.min(fx.max(mes[x_], 0), fx.Int32(E - 1)),
+                rt["rid"] + (tkn * TOPK + ranks[x_]),
+            )
+            fx.ptr_store(mss[x_] / total, rt["rw"] + (tkn * TOPK + ranks[x_]))
+
+
+def sel_peer(pl, d):
+    p = pl[0]
+    for r_ in range_constexpr(1, W):
+        p = (d == r_).select(pl[r_], p)
+    return p
+
+
+def fp4_mma(ops, base, bl, brow, lane, acc, c, jpar, rows, row0=0):
+    """Chunk c (256 K) of a 16-row MXFP4 group: 8 MFMA steps, B column n = LDS row row0 + n (at
+    most rows - 1: columns past the expert's picks read some row, their outputs are never stored).
+    ``jpar``: the task's 16-row group parity (j & 1), for the shuffled scales."""
+    g = lane // 16
+    t = fx.min(row0 + lane % 16, rows - 1)
+    sw = ops[base + 2]
+    # a = the dword of the lane's lower-half partner (K block % 4 = 2 (g % 2)), b = of the upper-half one
+    # (2 (g % 2) + 1); the byte = 2 (g // 2) (K block + 4) + the row group parity (rows + 16)
+    sa_, sb_ = permlane_swap(32, sw, sw)
+    shs = (g // 2) * 16 + jpar * 8
+    s0 = e8m0_f32((sa_ >> shs) & 0xFF)
+    s1 = e8m0_f32((sb_ >> shs) & 0xFF)
+    # loaded (shuffled): vector q lane 16 p + r' = row r', piece 4 q + p; MFMA lane 16 g + r wants row r
+    # pieces 2 g, 2 g + 1: X0 = [V0 p0, V0 p2, V1 p0, V1 p2], X1 = [V0 p1, V0 p3, V1 p1, V1 p3] (16-lane
+    # rows) = permlane32_swap then permlane16_swap, no bpermute
+    x0, x1 = [], []
+    for jj in range_constexpr(4):
+        a_, b_ = permlane_swap(32, ops[base][jj], ops[base + 1][jj])
+        p_, q_ = permlane_swap(16, a_, b_)
+        x0.append(p_)
+        x1.append(q_)
+    xs = [x0, x1]
+    for d in range_constexpr(8):
+        word = xs[d // 4][d % 4]
+        a = fp4x8_bf16(word, s0 if d < 4 else s1)
+        kk = FCH * c + 64 * g + 8 * d
+        b = fx.Vector(
+            fx.ptr_load(
+                bl + (t * brow + kk) // 2,
+                result_type=fx.Vector.make_type(4, fx.Int32),
+            )
+        ).bitcast(fx.BFloat16)
+        acc = mfma_bf16(a, b, acc)
+    return acc
+
+
+def o_gate(xv, gv):
+    """o-norm (over the head's 8 lanes) x sigmoid(gate) of 8 bf16 (4 pair words xv, gv) -> 4 words."""
+    ss = fx.Float32(0.0)
+    for d in range_constexpr(4):
+        ss = ss + bf_lo(xv[d]) * bf_lo(xv[d]) + bf_hi(xv[d]) * bf_hi(xv[d])
+    ss = butterfly(ss, (1, 2, 4))
+    rstd = fx.Float32(1.0) / hw_sqrt(ss * fx.Float32(1.0 / HD) + fx.Float32(PRE_EPS))
+    sig = lambda x: fx.Float32(1.0) / (
+        fx.Float32(1.0) + hw_exp2(x * fx.Float32(-LOG2E))
+    )
+    ow = []
+    for d in range_constexpr(4):
+        s0 = sig(bf_lo(gv[d]))
+        s1 = sig(bf_hi(gv[d]))
+        ow.append(bf16_pair(bf_lo(xv[d]) * rstd * s0, bf_hi(xv[d]) * rstd * s1))
+    return ow
+
+
+def mx_chunk(ok, w, ws, rbase, k, c, lane, g, r, nbytes, sbytes):
+    """Rows rbase .. rbase + 16 x K chunk c (256) of a shuffled MXFP4 [*, k] weight w and its e8m0
+    scales ws: 2 dwordx4 + 1 scale dword. vLLM shuffle_weight((16, 16)): lane L reads the 16 B at
+    16 L of the 1 KB block (16 rows x K bytes 64 (2 c + q) ..): one instruction = 1 KB contiguous;
+    fp4_mma's permlane swaps map the pieces to the MFMA lanes. e8m0_shuffle: lane 16 g + r reads
+    the dword (rbase // 32, c, K block % 4 = 2 (g % 2) + g // 2, row r): rows r, r + 16 x K blocks
+    of 4 apart; 64 lanes = 256 B contiguous (kernel_s_dword)."""
+    wb = (rbase * (k // 2) + (2 * c) * 1024 + lane * 16) // 4
+    ops = [
+        ld_vec_if(ok, w, wb, 4, nbytes, CM_NT),
+        ld_vec_if(ok, w, wb + 256, 4, nbytes, CM_NT),
+    ]
+    sdw = (rbase >> 5) * (8 * (k // 32)) + 64 * c + 16 * (2 * (g % 2) + g // 2) + r
+    return ops + [ld_i32_if(ok, ws, sdw, sbytes)]
+
+
+def dn_loads(w2, w2s, w_post, etab, task, lane, wave, ok, nv=None):
+    """Task (u, j)'s w2 operands for this wave (expert etab[u], u clamped below nv)."""
+    r = lane % 16
+    g = lane // 16
+    task = ok.select(task, fx.Int32(0))
+    u = task >> 8
+    j = task & 255
+    e = uniform(fx.ptr_load(etab + (u if nv is None else fx.min(u, nv - 1))))
+    rbase = e * DM + j * 16
+    ops = []
+    for i in range_constexpr(2):
+        c = 2 * wave + i
+        ops += mx_chunk(ok, w2, w2s, rbase, INTER, c, lane, g, r, W2_BYTES, W2S_BYTES)
+    # the epilogue's w_post pair (rows 16j + 2 (lane % 8), +1): loaded with the weights,
+    # not in the epilogue, where it would wait behind the whole prefetch queue
+    ops.append(ld_i32_if(ok, w_post, 8 * j + lane % 8, DM * 2))
+    return ops
+
+
+@traced
+def task_ring(load, compute, q, task, it, end, unr, *args):
+    """Tasks task .. end - 1: compute(q[0], task, it, *args) with the next len(q) tasks' loads
+    (load(task)) in flight. ``unr`` (len(q) = 2): three tasks an iteration in fixed slots (q[0],
+    q[1], qc): two tasks' loads stay in flight during each compute, no rotation copy waits for its
+    data (the tail: the second loop)."""
+    if const_expr(unr):
+        while task + 2 < end:
+            qc = load(task + 2)
+            compute(q[0], task, it, *args)
+            qa = load(task + 3)
+            compute(q[1], task + 1, it + 1, *args)
+            qb = load(task + 4)
+            compute(qc, task + 2, it + 2, *args)
+            q = [qa, qb]
+            task = task + 3
+            it = it + 3
+    while task < end:
+        nxt = load(task + len(q))
+        compute(q[0], task, it, *args)
+        q = q[1:] + [nxt]
+        task = task + 1
+        it = it + 1
+    return q, task, it
+
+
+REGION_IDS.update({"orr": 24, "mrr": 25, "ssa4p": 26})
+REGION_IDS.update(
+    {"op": 13, "xf": 16, "xnf": 17, "qx": 18, "qkv": 19, "ssa4": 20, "ssg4": 21}
+)
+
+
+@dataclass(frozen=True)
+class Build:
+    S: int = 8  # rows (8 or 16)
+    next_rope: bool = True  # the next layer has RoPE (SWA) or not (full, NoPE)
+    attn: str = ""  # this layer's verify attention fused at the head ("swa" / "full"); "" = run before the launch
+
+
+def _lay(items):
+    out, off = {}, 0
+    for name, n in items:
+        off = (off + 255) // 256 * 256
+        out[name] = (off, n)
+        off += n
+    return out, (off + 4095) // 4096 * 4096
+
+
+def scratch_layout(S, world=W):
+    """This rank's scratch regions. ss, ss1, ss2g, ssa1 and ssa2 are unused; they keep the
+    offsets of the regions after them."""
+    assert world == W, world
+    P = S * TOPK
+    return _lay(
+        [
+            ("px", NG * KQ * 64 * S * 8),  # [group][K part][row < RT <= 64][t]
+            ("a", P * (INTER // 2) * 8),
+            ("z", P * (DM // 2) * 8),
+            ("ss", P * 256 * 8),
+            ("ss2", P * 256 * 8),
+            ("m", S * DM * 2),  # all-reduced MoE output (bf16, plain)
+            ("mf", BLOCKS * 8),  # m tile flags
+            ("py", 512 * 16 * S * 8),  # post-proj K-part partials
+            ("ss1", (S * BLOCKS + S) * 8),
+            ("ss2g", (S * BLOCKS + S) * 8),
+            ("route", P * 8),
+            ("xin", S * HID * 2),  # ffn_in (bf16 plain), CTA b's 32 columns each
+            ("xf", BLOCKS * 8),  # its flags
+            ("smid", S * HID * 4),  # stream after the attention residual (fp32 plain)
+            ("ssa1", (S * BLOCKS + S) * 8),
+            ("ssa2", (S * BLOCKS + S) * 8),
+            ("xn", S * HID * 2),  # next layer's attention input (bf16 plain)
+            ("xnf", BLOCKS * 8),
+            # AB: 4 partials [t][CTA][4], then (r1, r2)[t]
+            ("ssa4", (4 * S * BLOCKS + 2 * S) * 8),
+            ("ssg4", (4 * S * BLOCKS + 2 * S) * 8),  # G
+            # qkvg K-eighth partials [group][src][tile row][t]
+            ("qx", 32 * 8 * NQT * 16 * S * 8),
+            # q / k / v rows (tagged bf16 pairs)
+            ("qkv", S * (QKV_ROWS // 2) * 8),
+        ]
+    )
+
+
+def peer_layout(S, world=W):
+    """Every rank's peer buffer (IPC-shared), [parity] first."""
+    assert world == W, world
+    return _lay(
+        [
+            ("logit", 2 * KQ * S * E * 8),  # [par][kq][t][e] tagged
+            ("hraw", 2 * S * DM * 2),  # [par][t][col] bf16 plain
+            ("hf", 2 * W * BLOCKS * 8),  # [par][src][CTA] tagged flags
+            ("moep", 2 * W * S * (DM // 2) * 8),  # [par][src][t][pair] tagged
+            ("y", 2 * S * (HID // 2) * 8),  # [par][t][pair] tagged
+            # o_proj partials [par][src][t][pair] tagged
+            ("op", 2 * W * S * (HID // 2) * 8),
+            # the owner's reduced o_proj [par][t][pair] tagged (all-gather)
+            ("orr", 2 * S * (HID // 2) * 8),
+            # the owner's reduced EP sum [par][t][pair] tagged (all-gather)
+            ("mrr", 2 * S * (DM // 2) * 8),
+            # the AB owners' norm partials [par][t][CTA][4] tagged (all-gather)
+            ("ssa4p", 2 * 4 * S * BLOCKS * 8),
+        ]
+    )
+
+
+_BUILDS = {}
+
+
+def build(key: Build):
+    if key in _BUILDS:
+        return _BUILDS[key]
+    # FlyDSL keys its disk cache on sources + *scalar* closure values; the Build
+    # dataclass is not a scalar, so every knob read as key.x would share one binary.
+    # Referencing this string from the launcher puts the whole key into the cache key.
+    _KEYSTR = repr(key)
+    S = key.S
+    assert S in (8, 16), S
+    XQROW = XQ + XPAD
+    CA = XQ // KCH // WAVES  # stage-A pre-proj chunks a wave
+    CO = OK_K // KCH // WAVES  # o_proj chunks a wave
+    QPF = 4  # qkvg tiles issued before stage G (of NQT)
+    P = S * TOPK
+    RROUNDS = (S + WAVES - 1) // WAVES
+    lay, _ = scratch_layout(S)
+    play, _ = peer_layout(S)
+    PF_GU = 2 if S == 8 else 1  # gate_up tasks in flight
+    PF_DN = 2  # down tasks in flight
+    # gate_up: the two 16-row halves of a task in separate MFMA passes
+    GUSPLIT = S == 16
+    W13_BYTES = EL * 2 * INTER * (DM // 2)
+    W13S_BYTES = EL * 2 * INTER * (DM // 32)
+    # stage A: a KQ-CTA group (K parts) covers RT pre-proj rows and EG router experts
+    RS_PRE = DM // W
+    RT = RS_PRE // NG  # 512 / 32 = 16
+    NT_A = RT // 16  # MFMA tiles a CTA
+    OWN = RT // KQ  # rows a CTA owns after the K-part exchange
+    EG = (E // W) // NG  # router experts a group
+    TPE = THREADS // EG  # threads an expert
+    KPT = XQ // TPE  # router K a thread
+    assert (NT_A, OWN, EG, KPT) == (1, 2, 1, 2), (NT_A, OWN, EG, KPT)
+    # stage F: post_expert_proj rows RS_POST a rank, 16-row tiles
+    RS_POST = HID // W
+    NTL = RS_POST // 16
+    TPC = 1  # tiles a CTA
+    KSP = BLOCKS // NTL  # K split of a tile
+    assert KSP == 4, KSP
+    KR = DM // KSP  # K a CTA reads
+    CPW = KR // KCH // WAVES  # chunks a wave a tile
+    MROW = KR + XPAD
+    tag_par = lambda tag: tag & 1
+
+    @fx.struct
+    class RouteLds:
+        rid: fx.Array[fx.Int32, S * TOPK, 16]
+        rw: fx.Array[fx.Float32, S * TOPK, 16]
+        uexp: fx.Array[fx.Int32, EL, 16]
+        kof: fx.Array[fx.Int32, EL * S, 16]
+        nu: fx.Array[fx.Int32, 4, 16]
+        sel: fx.Array[fx.Int32, max(S * TOPK * 3, EL), 16]
+        plist: fx.Array[fx.Int32, P, 16]
+        pidx: fx.Array[fx.Int32, P, 16]
+        np_: fx.Array[fx.Int32, 4, 16]
+        rstd: fx.Array[fx.Float32, P, 16]
+        uslot: fx.Array[fx.Int32, EL, 16]
+        unused: fx.Array[fx.Int32, 3 * E, 16]  # keeps the tuned LDS offsets
+        ssacc: fx.Array[fx.Float32, 2 * S, 16]
+        upk: fx.Array[fx.Int32, 2 * S, 16]
+        upn: fx.Array[fx.Int32, 4, 16]
+        tok: fx.Array[fx.Float32, WAVES * S + 8, 16]
+        rs: fx.Array[fx.Float32, 2 * S + 32, 16]
+
+    @fx.struct
+    class StageA:
+        x: fx.Array[fx.Int32, S * XQROW // 2, 16]
+        red: fx.Array[fx.Float32, WAVES * NT_A * 64 * 4, 16]
+        own: fx.Array[fx.Float32, OWN * S, 16]
+
+    @fx.struct
+    class StageC:
+        h: fx.Array[fx.Int32, S * HROW // 2, 16]
+        red: fx.Array[fx.Float32, 2 * (1 if GUSPLIT else 2) * WAVES * 64 * 4, 16]
+
+    @fx.struct
+    class StageD:
+        a: fx.Array[fx.Int32, S * AROW // 2, 16]
+        red: fx.Array[fx.Float32, 2 * WAVES * 64 * 4, 16]
+
+    @fx.struct
+    class StageE:
+        z: fx.Array[fx.Float32, P * 16, 16]
+
+    @fx.struct
+    class StageF:
+        mm: fx.Array[fx.Int32, S * MROW // 2, 16]
+        red: fx.Array[fx.Float32, WAVES * TPC * 64 * 4, 16]
+
+    YROW = OK_K + XPAD
+    XQ8 = HID // 8  # qkvg K eighth
+    assert CA >= 1 and CO >= 1 and NQT * 32 * 16 >= QKVG_ROWS
+    X8ROW = XQ8 + XPAD
+
+    @fx.struct
+    class StageO:
+        y: fx.Array[fx.Int32, S * YROW // 2, 16]
+        red: fx.Array[fx.Float32, WAVES * 2 * 64 * 4, 16]
+
+    @fx.struct
+    class StageQ:
+        xq: fx.Array[fx.Int32, S * X8ROW // 2, 16]
+        red: fx.Array[fx.Float32, WAVES * NQT * 64 * 4, 16]
+        own: fx.Array[fx.Float32, NQT * 2 * S, 16]
+
+    FA = None
+    if key.attn:  # the fused attention fragment (attn_frag.py)
+        from . import attn_frag
+
+        assert key.attn in ("swa", "full"), key
+        FA = attn_frag.make(S // 8, key.attn == "swa")
+        FaLds = FA.Lds
+        # the fragment's source hash: the disk cache does not follow FA.run
+        _KEYSTR = _KEYSTR + FA.key
+
+    @fx.union
+    class StageLds:
+        o_: StageO
+        q_: StageQ
+        a_: StageA
+        c_: StageC
+        d_: StageD
+        e_: StageE
+        f_: StageF
+        if FA is not None:
+            fa_: FaLds
+
+    @fx.struct
+    class Reserved:  # 128 B at the LDS base, unused: keeps the tuned LDS offsets
+        pad: fx.Array[fx.Int64, TL_POINTS, 16]
+
+    UNR = S == 8  # gate_up / down task loops as a ring (unroll 3, fixed slots)
+    name = f"muse_k3_s{S}_r{int(key.next_rope)}" + (
+        f"_fa{key.attn}" if key.attn else ""
+    )
+    CPRX = BLOCKS // W  # CTA blocks a rank owns (owner of block b = b // CPRX)
+
+    # ------------------------------------------------------------ stage A (sharded)
+    @traced
+    def stage_a(
+        x,
+        w_r,
+        bias,
+        w_pre,
+        scratch,
+        peer,
+        pl,
+        tag,
+        health,
+        tid,
+        bid,
+        lane,
+        wave,
+        rank,
+        xl,
+        red,
+        own,
+        tok,
+    ):
+        """CTA b = (group rb = b / KQ, K part kq = b % KQ): router experts EG rb .. and
+        pre-proj rows RT rb .. of this rank's shard over K part kq."""
+        tid = fresh(tid)  # no cross-stage CSE of tid-derived offsets (S=16 spills)
+        lane = fresh(lane)
+        g = lane // 16
+        kq = bid % KQ
+        rb = bid // KQ
+        k0 = kq * XQ
+        par = tag_par(tag)
+        nx = S * XQ // 8
+        # weights first (data independent), then the ffn_in quarter's 64 flags, then x
+        ei = tid // TPE
+        kr = (tid % TPE) * KPT
+        wr = ld_vec(w_r, (EG * rb + ei) * HID + k0 + kr, 2)  # KPT = 2 router weights
+        bias_v = [ld_f32(bias, lane + 64 * i) for i in range(4)]
+        wp = []
+        for nt in range_constexpr(NT_A):
+            for ci in range_constexpr(CA):
+                c = wave + WAVES * ci
+                wp.append(
+                    bf_tile_load(w_pre, rb * RT + nt * 16, HID, k0 + c * KCH, lane)
+                )
+        if tid < BLOCKS // KQ:  # the CTAs owning this K part's columns (32 a CTA)
+            mb_poll(
+                tag,
+                health,
+                [(scratch + fx.Int64(lay["xf"][0]), kq * (BLOCKS // KQ) + tid, 1)],
+                REGION_IDS["xf"],
+            )
+        gpu.barrier()
+        xv = []
+        for i in range_constexpr((nx + THREADS - 1) // THREADS):
+            e = fx.min(tid + THREADS * i, nx - 1)
+            t = e // (XQ // 8)
+            kk = (e % (XQ // 8)) * 8
+            xv.append(
+                ld_vec(
+                    scratch + fx.Int64(lay["xin"][0]),
+                    (t * HID + k0 + kk) // 2,
+                    4,
+                    CM_DEV,
+                )
+            )
+        for i in range_constexpr((nx + THREADS - 1) // THREADS):
+            e = tid + THREADS * i
+            if e < nx:
+                t = e // (XQ // 8)
+                kk = (e % (XQ // 8)) * 8
+                fx.ptr_store(xv[i], xl + (t * XQROW + kk) // 2)
+        gpu.barrier()
+        wf = [wr[0].bitcast(fx.Float32), wr[1].bitcast(fx.Float32)]
+        for t in range_constexpr(S):
+            pr = fx.Float32(0.0)
+            xw1 = fx.ptr_load(xl + (t * XQROW + kr) // 2)
+            pr = pr + bf_lo(xw1) * wf[0]
+            pr = pr + bf_hi(xw1) * wf[1]
+            pr = wave_sum(pr)
+            if lane == 0:
+                fx.ptr_store(pr, tok + (wave * S + t))
+        gpu.barrier()
+        if tid < EG * S:
+            e2 = tid // S
+            t = tid % S
+            v = fx.Float32(0.0)
+            for w in range_constexpr(WAVES // EG):
+                v = v + fx.ptr_load(tok + ((e2 * (WAVES // EG) + w) * S + t))
+            eg = rank * (E // W) + EG * rb + e2
+            idx = ((par * KQ + kq) * S + t) * E + eg
+            for pb in pl:
+                mb_put_sys(tag, pb + fx.Int64(play["logit"][0]), idx, v)
+        accs = [fx.Vector.filled(4, 0.0, fx.Float32) for _ in range(NT_A)]
+        t = fx.min(lane % 16, S - 1)
+        for nt in range_constexpr(NT_A):
+            for ci in range_constexpr(CA):
+                c = wave + WAVES * ci
+                for qq in range_constexpr(4):
+                    kk = c * KCH + 8 * g + 32 * qq
+                    b = fx.Vector(
+                        fx.ptr_load(
+                            xl + (t * XQROW + kk) // 2,
+                            result_type=fx.Vector.make_type(4, fx.Int32),
+                        )
+                    ).bitcast(fx.BFloat16)
+                    accs[nt] = mfma_bf16(
+                        bf_mma_operand(wp[nt * CA + ci], qq, lane).bitcast(fx.BFloat16),
+                        b,
+                        accs[nt],
+                    )
+        for nt in range_constexpr(NT_A):
+            fx.ptr_store(accs[nt], red + ((wave * NT_A + nt) * 64 + lane) * 4)
+        gpu.barrier()
+        # thread (row pair rp, token tt) over the group's RT rows: this quarter's partial
+        if tid < (RT // 2) * S:
+            rp = tid // S
+            tt = tid % S
+            v = []
+            for h2 in range_constexpr(2):
+                rr = 2 * rp + h2
+                s_ = fx.Float32(0.0)
+                for w in range_constexpr(WAVES):
+                    s_ = s_ + fx.ptr_load(
+                        red
+                        + (
+                            ((w * NT_A + rr // 16) * 64 + 16 * ((rr % 16) // 4) + tt)
+                            * 4
+                            + rr % 4
+                        )
+                    )
+                v.append(s_)
+            owner = (2 * rp) // OWN
+            if owner == kq:
+                fx.ptr_store(v[0], own + ((2 * rp - kq * OWN) * S + tt))
+                fx.ptr_store(v[1], own + ((2 * rp + 1 - kq * OWN) * S + tt))
+            else:
+                for h2 in range_constexpr(2):
+                    mb_put(
+                        tag,
+                        scratch + fx.Int64(lay["px"][0]),
+                        ((rb * KQ + kq) * 64 + 2 * rp + h2) * S + tt,
+                        v[h2],
+                    )
+        gpu.barrier()
+        # own rows: the KQ parts in order -> bf16 h_raw -> every rank's HRAW
+        if tid < (OWN // 2) * S:
+            rp = tid // S
+            tt = tid % S
+            specs, oks = [], []
+            for src in range_constexpr(KQ):
+                for h2 in range_constexpr(2):
+                    specs.append(
+                        (
+                            scratch + fx.Int64(lay["px"][0]),
+                            ((rb * KQ + src) * 64 + kq * OWN + 2 * rp + h2) * S + tt,
+                            1,
+                        )
+                    )
+                    oks.append(fx.Int32(src) != kq)
+            got = mb_poll(tag, health, specs, REGION_IDS["px"], oks=oks)
+            tot = []
+            for h2 in range_constexpr(2):
+                s_ = fx.Float32(0.0)
+                for src in range_constexpr(KQ):
+                    mine = fx.ptr_load(own + ((2 * rp + h2) * S + tt))
+                    s_ = s_ + (fx.Int32(src) == kq).select(
+                        mine, got[2 * src + h2][0].bitcast(fx.Float32)
+                    )
+                tot.append(s_)
+            col = (rank * RS_PRE) + rb * RT + kq * OWN + 2 * rp
+            hw = bf16_pair(tot[0], tot[1])
+            idx = (par * S * DM + tt * DM + col) // 2
+            for pb in pl:
+                bo.buffer_store(
+                    hw, rsrc(pb + fx.Int64(play["hraw"][0])), idx, cache_modifier=CM_SYS
+                )
+        rocdl.s_waitcnt(vmcnt=0, expcnt=0, lgkmcnt=0)
+        gpu.barrier()
+        if tid == 0:
+            fidx = (par * W + rank) * BLOCKS + bid
+            for pb in pl:
+                mb_put_sys(tag, pb + fx.Int64(play["hf"][0]), fidx, fx.Int32(1))
+        return bias_v
+
+    @traced
+    def route_poll(peer, tag, health, lane, wave):
+        """This wave's tokens' 256 logits from the local peer buffer: KQ K-part partials
+        each, summed in part order (identical on every CTA of every rank)."""
+        par = tag_par(tag)
+        specs = [
+            (
+                peer + fx.Int64(play["logit"][0]),
+                ((par * KQ + kq) * S + fx.min(wave + WAVES * rr, S - 1)) * E
+                + lane
+                + 64 * i,
+                1,
+            )
+            for rr in range(RROUNDS)
+            for i in range(4)
+            for kq in range(KQ)
+        ]
+        got = mb_poll(
+            tag,
+            health,
+            specs,
+            REGION_IDS["logit"],
+            batch=min(len(specs), 32),
+            cm=CM_SYS,
+        )  # W = 8 (KQ = 8): 2 polls of 32
+        out = []
+        for rr in range_constexpr(RROUNDS):
+            lgs = []
+            for i in range_constexpr(4):
+                base = (rr * 4 + i) * KQ
+                v = got[base][0].bitcast(fx.Float32)
+                for kq in range_constexpr(1, KQ):
+                    v = v + got[base + kq][0].bitcast(fx.Float32)
+                lgs.append(v.bitcast(fx.Int32))
+            out.append(lgs)
+        return out
+
+    @traced
+    def h_wait_issue(peer, w_pren, tag, health, tid, rank):
+        """Every producing rank's HRAW flags, then this thread's h_raw (cols kk .. kk+8
+        of every token) from the local peer buffer and its w_pre values."""
+        par = tag_par(tag)
+        NSRC = W
+        nf = NSRC * BLOCKS
+        specs, oks = [], []
+        for i in range_constexpr((nf + THREADS - 1) // THREADS):
+            f = fx.min(tid + THREADS * i, nf - 1)
+            src = f // BLOCKS
+            specs.append(
+                (
+                    peer + fx.Int64(play["hf"][0]),
+                    (par * W + src) * BLOCKS + f % BLOCKS,
+                    1,
+                )
+            )
+            oks.append((tid + THREADS * i) < nf)
+        mb_poll(tag, health, specs, REGION_IDS["hf"], oks=oks, cm=CM_SYS)
+        gpu.barrier()
+        kk = (tid % (DM // 8)) * 8
+        wv = ld_vec(w_pren, kk // 2, 4)
+        hv = []
+        for t in range_constexpr(S):
+            hv.append(
+                ld_vec(
+                    peer + fx.Int64(play["hraw"][0]),
+                    (par * S * DM + t * DM + kk) // 2,
+                    4,
+                    CM_SYS,
+                )
+            )
+        return hv, wv
+
+    @traced
+    def stage_route(scratch, bias_v, lgs_r, tag, health, rank, tid, lane, wave, rt):
+        """Every token's routing (wave w: tokens w, w + 8, ...), then the local union."""
+        if const_expr(S <= WAVES):
+            if wave < S:
+                route_token([wave], bias_v, [lgs_r[0]], tid, lane, rt)
+        else:
+            # S a multiple of 8: wave w routes tokens w, w + 8, .. interleaved
+            route_token(
+                [wave + WAVES * rr for rr in range(RROUNDS)],
+                bias_v,
+                lgs_r,
+                tid,
+                lane,
+                rt,
+            )
+        gpu.barrier()
+        # the union U of local experts (sorted), kof[u][t]; picks list
+        e0 = rank * EL
+        # parallel: picks flag their experts, a scan numbers U, picks scatter kof
+        for ii in range_constexpr((EL * S + THREADS - 1) // THREADS):
+            ix = tid + THREADS * ii
+            if ix < EL * S:
+                fx.ptr_store(fx.Int32(-1), rt["kof"] + ix)
+        if tid < EL:
+            fx.ptr_store(fx.Int32(0), rt["sel"] + tid)
+        gpu.barrier()
+        for ii in range_constexpr((P + THREADS - 1) // THREADS):
+            pq = tid + THREADS * ii
+            if pq < P:
+                eid = slot_of(fx.ptr_load(rt["rid"] + pq), e0)
+                if eid >= 0:
+                    fx.ptr_store(fx.Int32(1), rt["sel"] + eid)
+        gpu.barrier()
+        if wave == 0:
+            ucarry = fx.Int32(0)
+            for hh in range_constexpr((EL + 63) // 64):
+                el_ = lane + 64 * hh
+                f = fx.ptr_load(rt["sel"] + fx.min(el_, EL - 1))
+                f = (el_ < EL).select(f, fx.Int32(0))
+                incl = wave_incl_scan(f, lane) + ucarry
+                slot = incl - f
+                if f != 0:
+                    fx.ptr_store(el_, rt["uexp"] + slot)
+                if el_ < EL:
+                    fx.ptr_store(slot, rt["uslot"] + el_)
+                ucarry = readlane(incl, 63)
+            if lane == 63:
+                fx.ptr_store(ucarry, rt["nu"] + 0)
+        if wave == 1:
+            # local picks p = t*8 + k in (t, k) order (lane p, then p + 64, ...)
+            carry = fx.Int32(0)
+            for hh in range_constexpr((P + 63) // 64):
+                pl_ = lane + 64 * hh
+                pp = fx.min(pl_, P - 1)
+                eid = slot_of(fx.ptr_load(rt["rid"] + pp), e0)
+                loc = ((eid >= 0) & (pl_ < P)).select(fx.Int32(1), fx.Int32(0))
+                incl = wave_incl_scan(loc, lane) + carry
+                if loc != 0:
+                    fx.ptr_store(pl_, rt["plist"] + (incl - loc))
+                if pl_ < P:
+                    fx.ptr_store(incl - loc, rt["pidx"] + pp)
+                carry = readlane(incl, 63)
+            if lane == 63:
+                fx.ptr_store(carry, rt["np_"] + 0)
+        gpu.barrier()
+        for ii in range_constexpr((P + THREADS - 1) // THREADS):
+            pq = tid + THREADS * ii
+            if pq < P:
+                eid = slot_of(fx.ptr_load(rt["rid"] + pq), e0)
+                if eid >= 0:
+                    slot = fx.ptr_load(rt["uslot"] + eid)
+                    fx.ptr_store(
+                        fx.Int32(pq % TOPK), rt["kof"] + (slot * S + pq // TOPK)
+                    )
+        gpu.barrier()
+
+    def slot_of(e, e0):
+        """Local slot of global expert e on this rank (-1: not computed here)."""
+        el_ = e - e0
+        return ((el_ >= 0) & (el_ < EL)).select(el_, fx.Int32(-1))
+
+    def gu_loads(w13, w13s, rt, task, lane, wave, ok):
+        """Task (u, j)'s w13 operands for this wave: gate / up row groups x the
+        wave's 2 chunks of 256 K: 2 dwordx4 + 1 scale dword each. ``ok`` false:
+        zeros, no memory access (a lookahead past the CTA's last task)."""
+        r = lane % 16
+        g = lane // 16
+        task = ok.select(task, fx.Int32(0))
+        u = task >> 8
+        j = task & 255
+        e = uniform(fx.ptr_load(rt["uexp"] + u))
+        ops = []
+        if const_expr(GUSPLIT):
+            # waves 0-3 the gate rows, 4-7 the up rows; wave w the chunks 4 (w % 4) + i
+            rgs = [wave // 4]
+            cs = [4 * (wave % 4) + i for i in range(4)]
+        else:
+            rgs = list(range(2))
+            cs = [2 * wave + i for i in range(2)]
+        for rg in rgs:
+            rbase = e * (2 * INTER) + rg * INTER + j * 16
+            for c in cs:
+                ops += mx_chunk(
+                    ok, w13, w13s, rbase, DM, c, lane, g, r, W13_BYTES, W13S_BYTES
+                )
+        return ops
+
+    @traced
+    def load_h_norm(hv, wv, tid, lane, wave, hl, tok):
+        """Rstd per token from the full h_raw, h = bf16(h_raw rstd w) -> LDS."""
+        kk = (tid % (DM // 8)) * 8
+        for t in range_constexpr(S):
+            ss = fx.Float32(0.0)
+            for d in range_constexpr(4):
+                lo = bf_lo(hv[t][d])
+                hi = bf_hi(hv[t][d])
+                ss = ss + lo * lo + hi * hi
+            ss = wave_sum(ss)
+            if lane == 0:
+                fx.ptr_store(ss, tok + (wave * S + t))
+        gpu.barrier()
+        for t in range_constexpr(S):
+            tot = fx.Float32(0.0)
+            for w in range_constexpr(WAVES):
+                tot = tot + fx.ptr_load(tok + (w * S + t))
+            rs = hw_rsq(tot * fx.Float32(1.0 / DM) + fx.Float32(PRE_EPS))
+            outw = []
+            for d in range_constexpr(4):
+                a0 = bf_lo(hv[t][d]) * rs * bf_lo(wv[d])
+                a1 = bf_hi(hv[t][d]) * rs * bf_hi(wv[d])
+                outw.append(bf16_pair(a0, a1))
+            fx.ptr_store(
+                fx.Vector.from_elements(outw, fx.Int32), hl + (t * HROW + kk) // 2
+            )
+        gpu.barrier()
+
+    @traced
+    def gu_compute(ops, task, it, scratch, tag, rt, tid, lane, wave, hl, red):
+        if const_expr(GUSPLIT):
+            acc = fx.Vector.filled(4, 0.0, fx.Float32)
+            for i in range_constexpr(4):
+                acc = fp4_mma(
+                    ops, 3 * i, hl, HROW, lane, acc, 4 * (wave % 4) + i, task & 1, S
+                )
+            par = (it & 1) * (WAVES * 64 * 4)
+            fx.ptr_store(acc, red + (par + (wave * 64 + lane) * 4))
+        else:
+            acc_g = fx.Vector.filled(4, 0.0, fx.Float32)
+            acc_u = fx.Vector.filled(4, 0.0, fx.Float32)
+            for i in range_constexpr(2):
+                acc_g = fp4_mma(
+                    ops, 3 * i, hl, HROW, lane, acc_g, 2 * wave + i, task & 1, S
+                )
+                acc_u = fp4_mma(
+                    ops, 6 + 3 * i, hl, HROW, lane, acc_u, 2 * wave + i, task & 1, S
+                )
+            par = (it & 1) * (2 * WAVES * 64 * 4)
+            fx.ptr_store(acc_g, red + (par + (wave * 64 + lane) * 4))
+            fx.ptr_store(acc_u, red + (par + WAVES * 64 * 4 + (wave * 64 + lane) * 4))
+        gpu.barrier()
+        if tid < 8 * S:
+            rp = tid % 8
+            tt = tid // 8
+            u = task >> 8
+            j = task & 255
+            k = fx.ptr_load(rt["kof"] + (u * S + tt))
+            av = []
+            for h2 in range_constexpr(2):
+                rr = 2 * rp + h2
+                gs = fx.Float32(0.0)
+                us = fx.Float32(0.0)
+                if const_expr(GUSPLIT):
+                    for w in range_constexpr(4):
+                        off = par + (w * 64 + 16 * (rr // 4) + tt) * 4 + rr % 4
+                        gs = gs + fx.ptr_load(red + off)
+                        us = us + fx.ptr_load(red + (off + 4 * 64 * 4))
+                else:
+                    for w in range_constexpr(WAVES):
+                        off = par + (w * 64 + 16 * (rr // 4) + tt) * 4 + rr % 4
+                        gs = gs + fx.ptr_load(red + off)
+                        us = us + fx.ptr_load(red + (off + WAVES * 64 * 4))
+                gb = bf16_round(gs)
+                ub = bf16_round(us)
+                sg = bf16_round(
+                    gb / (fx.Float32(1.0) + hw_exp2(gb * fx.Float32(-LOG2E)))
+                )
+                av.append(sg * ub)
+            if k >= 0:
+                pick = tt * TOPK + k
+                mb_put(
+                    tag,
+                    scratch + fx.Int64(lay["a"][0]),
+                    pick * (INTER // 2) + j * 8 + rp,
+                    bf16_pair(av[0], av[1]),
+                )
+
+    @traced
+    def dn_build_picks(rt, u0, u1, tid):
+        """upk[h][n]: the picks (t*8 + k, token order) of expert slot u0 + h; upn[h]."""
+        if tid < 2:
+            u = fx.min(u0 + tid, EL - 1)
+            vu = (tid == 0) | (u1 != u0)
+            n = fx.Int32(0)
+            for tt in range_constexpr(S):
+                k = fx.ptr_load(rt["kof"] + (u * S + tt))
+                ok = vu & (k >= 0)
+                if ok:
+                    fx.ptr_store(fx.Int32(tt * TOPK) + k, rt["upk"] + (tid * S + n))
+                n = n + ok.select(fx.Int32(1), fx.Int32(0))
+            fx.ptr_store(n, rt["upn"] + tid)
+        gpu.barrier()
+
+    @traced
+    def load_ac(scratch, tag, health, rt, second, buf, tid):
+        """Compact A: LDS row m = u0's pick m (m < n0), then u1's picks (when both fit
+        in S rows); ``second``: u1's picks from row 0 (the overflow reload)."""
+        n0 = uniform(fx.ptr_load(rt["upn"] + 0))
+        n1 = uniform(fx.ptr_load(rt["upn"] + 1))
+        fit = (n0 + n1) <= S
+        specs, oks, dsts = [], [], []
+        for m in range_constexpr(S):
+            if const_expr(second):
+                pk = fx.ptr_load(rt["upk"] + (S + fx.Int32(m)))
+                ok = fx.Int32(m) < n1
+            else:
+                in0 = fx.Int32(m) < n0
+                pk = in0.select(
+                    fx.ptr_load(rt["upk"] + m),
+                    fx.ptr_load(rt["upk"] + (S + fx.max(fx.Int32(m) - n0, 0))),
+                )
+                ok = in0 | (fit & (fx.Int32(m) < n0 + n1))
+            pk = ok.select(pk, fx.Int32(0))
+            for qq in range_constexpr(INTER // 4 // THREADS):
+                sidx = tid + THREADS * qq
+                specs.append(
+                    (scratch + fx.Int64(lay["a"][0]), pk * (INTER // 2) + 2 * sidx, 2)
+                )
+                oks.append(ok)
+                dsts.append(buf + (m * AROW + 4 * sidx) // 2)
+        got = mb_poll(tag, health, specs, REGION_IDS["a"], batch=16, oks=oks)
+        for n in range_constexpr(len(specs)):
+            fx.ptr_store(
+                fx.Vector.from_elements([got[n][0], got[n][1]], fx.Int32), dsts[n]
+            )
+
+    @traced
+    def dn_switch(scratch, tag, health, rt, task, u0, u1, buf, tid):
+        """Overflow (u0's + u1's picks > S rows): reload u1's rows at its first task."""
+        n0 = uniform(fx.ptr_load(rt["upn"] + 0))
+        n1 = uniform(fx.ptr_load(rt["upn"] + 1))
+        if (u1 != u0) & ((n0 + n1) > S):
+            gpu.barrier()
+            load_ac(scratch, tag, health, rt, True, buf, tid)
+            gpu.barrier()
+
+    @traced
+    def dn_compute_c(ops, task, it, ubase, scratch, tag, rt, tid, lane, wave, al, red):
+        """Compact-A down task: B column n = pick n of the task's expert."""
+        u = task >> 8
+        j = task & 255
+        h = (u - ubase) & 1
+        n0 = uniform(fx.ptr_load(rt["upn"] + 0))
+        n1 = uniform(fx.ptr_load(rt["upn"] + 1))
+        fit = (n0 + n1) <= S
+        rowb = (h == 0).select(fx.Int32(0), fit.select(n0, fx.Int32(0)))
+        nv = (h == 0).select(n0, n1)
+        acc = fx.Vector.filled(4, 0.0, fx.Float32)
+        for i in range_constexpr(2):
+            acc = fp4_mma(
+                ops, 3 * i, al, AROW, lane, acc, 2 * wave + i, task & 1, S, rowb
+            )
+        par = (it & 1) * (WAVES * 64 * 4)
+        fx.ptr_store(acc, red + (par + (wave * 64 + lane) * 4))
+        gpu.barrier()
+        if tid < 8 * S:
+            rp = tid % 8
+            n = tid // 8
+            pk = fx.ptr_load(rt["upk"] + (h * S + fx.min(n, S - 1)))
+            tt = pk >> 3
+            k = pk & 7
+            wpw = ops[6]
+            zv = []
+            for h2 in range_constexpr(2):
+                rr = 2 * rp + h2
+                ys = fx.Float32(0.0)
+                for w in range_constexpr(WAVES):
+                    ys = ys + fx.ptr_load(
+                        red + (par + (w * 64 + 16 * (rr // 4) + n) * 4 + rr % 4)
+                    )
+                wp = bf_lo(wpw) if h2 == 0 else bf_hi(wpw)
+                zv.append(bf16_round(wp * bf16_round(ys)))
+            ss = zv[0] * zv[0] + zv[1] * zv[1]
+            ss = butterfly(ss, (1, 2, 4))
+            if n < nv:
+                pick = tt * TOPK + k
+                mb_put(
+                    tag,
+                    scratch + fx.Int64(lay["z"][0]),
+                    pick * (DM // 2) + j * 8 + rp,
+                    bf16_pair(zv[0], zv[1]),
+                )
+                if rp == 0:
+                    sa_ = rt["ssacc"] + (h * S + tt)
+                    fx.ptr_store(fx.ptr_load(sa_) + ss, sa_)
+
+    @traced
+    def stage_final(scratch, out, tag, health, rt, tid, bid, lane, zl):
+        npk = uniform(fx.ptr_load(rt["np_"] + 0))
+        part = tid % 8
+        for hh in range_constexpr((P * 8 + THREADS - 1) // THREADS):
+            pi = tid // 8 + (THREADS // 8) * hh
+            if pi < npk:
+                final_pick(scratch, tag, health, rt, pi, part, bid, zl)
+
+    @traced
+    def final_pick(scratch, tag, health, rt, pi, part, bid, zl):
+        pick = fx.ptr_load(rt["plist"] + pi)
+        zspec = (scratch + fx.Int64(lay["z"][0]), pick * (DM // 2) + bid * 8 + part, 1)
+        # the CTA-aggregated partials of the pick's expert: CTAs cf .. cl
+        nu_ = uniform(fx.ptr_load(rt["nu"] + 0))
+        el_ = slot_of(fx.ptr_load(rt["rid"] + pick), rt["rank"] * EL)
+        u = fx.ptr_load(rt["uslot"] + fx.min(fx.max(el_, 0), EL - 1))
+        cf = (u * 256) // nu_
+        nc = (u * 256 + 255) // nu_ - cf + 1
+        specs = [
+            (
+                scratch + fx.Int64(lay["ss2"][0]),
+                pick * 256 + fx.min(part + 8 * q, 255),
+                1,
+            )
+            for q in range(32)
+        ]
+        oks = [(part + 8 * q) < nc for q in range(32)]
+        got = mb_poll(
+            tag,
+            health,
+            specs + [zspec],
+            REGION_IDS["ss"],
+            batch=33,
+            oks=oks + [fx.Boolean(True)],
+        )
+        s_ = fx.Float32(0.0)
+        for q in range_constexpr(32):
+            s_ = s_ + got[q][0].bitcast(fx.Float32)
+        zw = got[32][0]
+        s_ = butterfly(s_, (1, 2, 4))
+        fx.ptr_store(bf_lo(zw), zl + (pi * 16 + 2 * part))
+        fx.ptr_store(bf_hi(zw), zl + (pi * 16 + 2 * part + 1))
+        if part == 0:
+            fx.ptr_store(
+                hw_rsq(s_ * fx.Float32(1.0 / DM) + fx.Float32(POST_EPS)),
+                rt["rstd"] + pi,
+            )
+
+    # ------------------------------------------------------------ stage E': partial push + EP all-reduce
+    @traced
+    def final_push(scratch, part_out, pl, tag, rt, tid, bid, lane, zl, rank):
+        """CTA b's [S x 16] EP partial (slot order, the reference rounding) -> part_out
+        (local) and slot [parity][rank] of the block owner's MOEP."""
+        gpu.barrier()
+        par = tag_par(tag)
+        if tid < 8 * S:
+            rp = tid % 8
+            tt = tid // 8
+            ov = []
+            for h2 in range_constexpr(2):
+                rr = 2 * rp + h2
+                acc = fx.Float32(0.0)
+                for k in range_constexpr(TOPK):
+                    p = tt * TOPK + k
+                    ix = fx.ptr_load(rt["pidx"] + p)
+                    eid = fx.ptr_load(rt["rid"] + p)
+                    isl = slot_of(eid, rank * EL) >= 0
+                    ixc = fx.min(fx.max(ix, 0), P - 1)
+                    zz = fx.ptr_load(zl + (ixc * 16 + rr))
+                    rs = fx.ptr_load(rt["rstd"] + ixc)
+                    yn = bf16_round(zz * rs)
+                    acc = acc + isl.select(
+                        yn * fx.ptr_load(rt["rw"] + p), fx.Float32(0.0)
+                    )
+                ov.append(acc)
+            w = bf16_pair(ov[0], ov[1])
+            bo.buffer_store(w, rsrc(part_out), (tt * DM + bid * 16 + 2 * rp) // 2)
+            pd = sel_peer(pl, bid // CPRX)
+            mb_put_sys(
+                tag,
+                pd + fx.Int64(play["moep"][0]),
+                ((par * W + rank) * S + tt) * (DM // 2) + bid * 8 + rp,
+                w,
+            )
+
+    @traced
+    def ar_reduce(scratch, own_peer, tag, health, tid, bid, pl=None, rank=None):
+        """The W ranks' slots of this CTA's columns, summed in rank order (fp32) ->
+        bf16 m (local, plain) -> drain -> this CTA's MF flag."""
+        par = tag_par(tag)
+        # the block owner: W slots in rank order -> every rank
+        if (tid < 8 * S) & ((bid // CPRX) == rank):
+            rp = tid % 8
+            tt = tid // 8
+            specs = [
+                (
+                    own_peer + fx.Int64(play["moep"][0]),
+                    ((par * W + src) * S + tt) * (DM // 2) + bid * 8 + rp,
+                    1,
+                )
+                for src in range(W)
+            ]
+            got = mb_poll(tag, health, specs, REGION_IDS["moep"], cm=CM_SYS)
+            lo = bf_lo(got[0][0])
+            hi = bf_hi(got[0][0])
+            for src in range_constexpr(1, W):
+                lo = lo + bf_lo(got[src][0])
+                hi = hi + bf_hi(got[src][0])
+            wr = bf16_pair(lo, hi)
+            for pb in pl:
+                mb_put_sys(
+                    tag,
+                    pb + fx.Int64(play["mrr"][0]),
+                    (par * S + tt) * (DM // 2) + bid * 8 + rp,
+                    wr,
+                )
+        if tid < 8 * S:
+            rp = tid % 8
+            tt = tid // 8
+            got = mb_poll(
+                tag,
+                health,
+                [
+                    (
+                        own_peer + fx.Int64(play["mrr"][0]),
+                        (par * S + tt) * (DM // 2) + bid * 8 + rp,
+                        1,
+                    )
+                ],
+                REGION_IDS["mrr"],
+                cm=CM_SYS,
+            )
+            lo = bf_lo(got[0][0])
+            hi = bf_hi(got[0][0])
+            bo.buffer_store(
+                bf16_pair(lo, hi),
+                rsrc(scratch + fx.Int64(lay["m"][0])),
+                (tt * DM + bid * 16 + 2 * rp) // 2,
+                cache_modifier=CM_DEV,
+            )
+        rocdl.s_waitcnt(vmcnt=0, expcnt=0, lgkmcnt=0)
+        gpu.barrier()
+        if tid == 0:
+            mb_put(tag, scratch + fx.Int64(lay["mf"][0]), bid, fx.Int32(1))
+
+    # ------------------------------------------------------------ stage F: post_expert_proj
+    def pp_tile(bid, i):
+        return bid // KSP
+
+    def pp_loads(w_pp, bid, lane, wave):
+        k0 = (bid % KSP) * KR
+        ops = []
+        for i in range_constexpr(TPC):
+            for ci in range_constexpr(CPW):
+                c = wave + WAVES * ci
+                ops.append(
+                    bf_tile_load(w_pp, pp_tile(bid, i) * 16, DM, k0 + c * KCH, lane)
+                )
+        return ops
+
+    @traced
+    def post_proj(
+        w_pp_ops, scratch, peer, pl, tag, health, tid, bid, lane, wave, rank, ml, red
+    ):
+        # lane math through an opaque move: with MROW == XQROW LLVM otherwise reuses stage A's
+        # LDS addresses here, holding them across the whole kernel (spills)
+        lane = fresh(lane)
+        par = tag_par(tag)
+        g = lane // 16
+        kh = bid % KSP
+        k0 = kh * KR
+        nfl = KR // 16
+        specs, oks = [], []
+        for i in range_constexpr((nfl + THREADS - 1) // THREADS):
+            f = fx.min(tid + THREADS * i, nfl - 1)
+            specs.append((scratch + fx.Int64(lay["mf"][0]), k0 // 16 + f, 1))
+            oks.append((tid + THREADS * i) < nfl)
+        mb_poll(tag, health, specs, REGION_IDS["mf"], oks=oks)
+        gpu.barrier()
+        nm = S * KR // 8
+        for i in range_constexpr((nm + THREADS - 1) // THREADS):
+            e = fx.min(tid + THREADS * i, nm - 1)
+            t = e // (KR // 8)
+            kk = (e % (KR // 8)) * 8
+            v = ld_vec(
+                scratch + fx.Int64(lay["m"][0]), (t * DM + k0 + kk) // 2, 4, CM_DEV
+            )
+            if tid + THREADS * i < nm:
+                fx.ptr_store(v, ml + (t * MROW + kk) // 2)
+        gpu.barrier()
+        t = fx.min(lane % 16, S - 1)
+        accs = [fx.Vector.filled(4, 0.0, fx.Float32) for _ in range(TPC)]
+        for i in range_constexpr(TPC):
+            for ci in range_constexpr(CPW):
+                c = wave + WAVES * ci
+                for qq in range_constexpr(4):
+                    kk = c * KCH + 8 * g + 32 * qq
+                    b = fx.Vector(
+                        fx.ptr_load(
+                            ml + (t * MROW + kk) // 2,
+                            result_type=fx.Vector.make_type(4, fx.Int32),
+                        )
+                    ).bitcast(fx.BFloat16)
+                    accs[i] = mfma_bf16(
+                        bf_mma_operand(w_pp_ops[i * CPW + ci], qq, lane).bitcast(
+                            fx.BFloat16
+                        ),
+                        b,
+                        accs[i],
+                    )
+        for i in range_constexpr(TPC):
+            fx.ptr_store(accs[i], red + ((wave * TPC + i) * 64 + lane) * 4)
+        gpu.barrier()
+        if tid < TPC * 8 * S:
+            i = tid // (8 * S)
+            rp = (tid // S) % 8
+            tt = tid % S
+            v = []
+            for h2 in range_constexpr(2):
+                rr = 2 * rp + h2
+                s_ = fx.Float32(0.0)
+                for w in range_constexpr(WAVES):
+                    s_ = s_ + fx.ptr_load(
+                        red + (((w * TPC + i) * 64 + 16 * (rr // 4) + tt) * 4 + rr % 4)
+                    )
+                v.append(s_)
+            tile = bid // KSP
+            # K parts 1 .. KSP-1 put their partial (slot kh - 1); part 0 sums all parts in order
+            if kh != 0:
+                for h2 in range_constexpr(2):
+                    mb_put(
+                        tag,
+                        scratch + fx.Int64(lay["py"][0]),
+                        ((kh - 1) * NTL * 16 + tile * 16 + 2 * rp + h2) * S + tt,
+                        v[h2],
+                    )
+            else:
+                got = mb_poll(
+                    tag,
+                    health,
+                    [
+                        (
+                            scratch + fx.Int64(lay["py"][0]),
+                            ((q - 1) * NTL * 16 + tile * 16 + 2 * rp + h2) * S + tt,
+                            1,
+                        )
+                        for q in range(1, KSP)
+                        for h2 in range(2)
+                    ],
+                    REGION_IDS["py"],
+                )
+                y0 = v[0]
+                y1 = v[1]
+                for q in range_constexpr(KSP - 1):
+                    y0 = y0 + got[2 * q][0].bitcast(fx.Float32)
+                    y1 = y1 + got[2 * q + 1][0].bitcast(fx.Float32)
+                col = (rank * RS_POST) + tile * 16 + 2 * rp
+                idx = (par * S + tt) * (HID // 2) + col // 2
+                for pb in pl:
+                    mb_put_sys(tag, pb + fx.Int64(play["y"][0]), idx, bf16_pair(y0, y1))
+
+    # ------------------------------------------------------------ stage G: post-FFN norm, residual, next norm
+    @traced
+    def tok_rstd2(
+        scratch,
+        region,
+        tag,
+        health,
+        tid,
+        lane,
+        rs_out,
+        eps1,
+        eps2,
+        bid,
+        red,
+        pbase=None,
+        pcm=CM_DEV,
+    ):
+        """One round for two norms of a token. Every CTA put 4 partials of token t at
+        ((t * BLOCKS + CTA) * 4 + j): P1 = sum o^2, P2 = sum (a s)^2, P3 = sum (a s)(b o),
+        P4 = sum (b o)^2 for s' = a s + b (o r1). Reducer CTA 255 - t sums the 256 CTAs in a
+        fixed order: r1 = rsqrt(P1 / HID + eps1), r2 = rsqrt((P2 + 2 r1 P3 + r1^2 P4) / HID + eps2)
+        (r2 from the unrounded o r1). rs_out[t] = r1, rs_out[S + t] = r2 on every CTA."""
+        rbase = scratch + fx.Int64(lay[region][0])
+        # stage AB: the partials come from the own peer buffer (the block owners)
+        prb = rbase if pbase is None else pbase
+        tred = bid >= BLOCKS - S
+        if tred:
+            t = BLOCKS - 1 - bid
+            va = fx.Float32(0.0)
+            vb = fx.Float32(0.0)
+            if tid < 2 * BLOCKS:
+                got = mb_poll(
+                    tag,
+                    health,
+                    [(prb, (t * BLOCKS + tid // 2) * 4 + 2 * (tid % 2), 2)],
+                    REGION_IDS[region],
+                    cm=pcm,
+                )
+                va = got[0][0].bitcast(fx.Float32)
+                vb = got[0][1].bitcast(fx.Float32)
+            va = butterfly(va, (32, 16, 8, 4, 2))  # lanes of one parity: 32 CTAs
+            vb = butterfly(vb, (32, 16, 8, 4, 2))
+            if lane < 2:
+                fx.ptr_store(va, red + (tid // 64) * 4 + 2 * lane)
+                fx.ptr_store(vb, red + (tid // 64) * 4 + 2 * lane + 1)
+            gpu.barrier()
+            if tid == 0:
+                p = [fx.Float32(0.0) for _ in range(4)]
+                for w in range_constexpr(WAVES):
+                    for j in range_constexpr(4):
+                        p[j] = p[j] + fx.ptr_load(red + w * 4 + j)
+                # red slot order: lane 0 -> (P1, P2), lane 1 -> (P3, P4)
+                r1 = hw_rsq(p[0] * fx.Float32(1.0 / HID) + fx.Float32(eps1))
+                q = p[1] + fx.Float32(2.0) * r1 * p[2] + r1 * r1 * p[3]
+                r2 = hw_rsq(q * fx.Float32(1.0 / HID) + fx.Float32(eps2))
+                mb_put(tag, rbase, 4 * S * BLOCKS + 2 * t, r1)
+                mb_put(tag, rbase, 4 * S * BLOCKS + 2 * t + 1, r2)
+        if tid < S:
+            got2 = mb_poll(
+                tag, health, [(rbase, 4 * S * BLOCKS + 2 * tid, 2)], REGION_IDS[region]
+            )
+            fx.ptr_store(got2[0][0].bitcast(fx.Float32), rs_out + tid)
+            fx.ptr_store(got2[0][1].bitcast(fx.Float32), rs_out + (S + tid))
+        gpu.barrier()
+
+    def put4(tag, scratch, region, tt, bid, q1, q2, q3, q4):
+        base = scratch + fx.Int64(lay[region][0])
+        i0 = (tt * BLOCKS + bid) * 4
+        bo.buffer_store(
+            fx.Vector.from_elements(
+                [q1.bitcast(fx.Int32), tag, q2.bitcast(fx.Int32), tag], fx.Int32
+            ),
+            rsrc(base),
+            i0 * 2,
+            cache_modifier=CM_DEV,
+        )
+        bo.buffer_store(
+            fx.Vector.from_elements(
+                [q3.bitcast(fx.Int32), tag, q4.bitcast(fx.Int32), tag], fx.Int32
+            ),
+            rsrc(base),
+            (i0 + 2) * 2,
+            cache_modifier=CM_DEV,
+        )
+
+    @traced
+    def stage_g(
+        peer,
+        scratch,
+        stream_in,
+        w_pff,
+        alpha,
+        beta,
+        w_in,
+        stream_out,
+        xn_out,
+        tag,
+        health,
+        tid,
+        bid,
+        lane,
+        rsl,
+    ):
+        tid = fresh(tid)  # no cross-stage CSE of tid-derived offsets (S=16 spills)
+        par = tag_par(tag)
+        pp = tid % 16
+        tt = fx.min(tid // 16, S - 1)
+        live = tid < 16 * S
+        c = bid * 32 + 2 * pp
+        # data-independent loads first
+        st = ld_vec(scratch + fx.Int64(lay["smid"][0]), tt * HID + c, 2, CM_DEV)
+        wpf = ld_i32(w_pff, c // 2)
+        al = ld_vec(alpha, c, 2)
+        be = ld_vec(beta, c, 2)
+        wi = ld_i32(w_in, c // 2)
+        z0 = fx.Float32(0.0)
+        z1 = fx.Float32(0.0)
+        if live:
+            got = mb_poll(
+                tag,
+                health,
+                [
+                    (
+                        peer + fx.Int64(play["y"][0]),
+                        (par * S + tt) * (HID // 2) + bid * 16 + pp,
+                        1,
+                    )
+                ],
+                REGION_IDS["y"],
+                cm=CM_SYS,
+            )
+            yw = got[0][0]
+            z0 = bf16_round(bf_lo(wpf) * bf_lo(yw))
+            z1 = bf16_round(bf_hi(wpf) * bf_hi(yw))
+        a0 = al[0].bitcast(fx.Float32) * st[0].bitcast(fx.Float32)
+        a1 = al[1].bitcast(fx.Float32) * st[1].bitcast(fx.Float32)
+        b0 = be[0].bitcast(fx.Float32) * z0
+        b1 = be[1].bitcast(fx.Float32) * z1
+        q1 = butterfly(z0 * z0 + z1 * z1, (1, 2, 4, 8))
+        q2 = butterfly(a0 * a0 + a1 * a1, (1, 2, 4, 8))
+        q3 = butterfly(a0 * b0 + a1 * b1, (1, 2, 4, 8))
+        q4 = butterfly(b0 * b0 + b1 * b1, (1, 2, 4, 8))
+        if live & (pp == 0):
+            put4(tag, scratch, "ssg4", tt, bid, q1, q2, q3, q4)
+        tok_rstd2(
+            scratch,
+            "ssg4",
+            tag,
+            health,
+            tid,
+            lane,
+            rsl,
+            POST_EPS,
+            PRE_EPS,
+            bid,
+            rsl + 2 * S,
+        )
+        r1 = fx.ptr_load(rsl + tt)
+        m0 = bf16_round(z0 * r1)
+        m1 = bf16_round(z1 * r1)
+        s0 = a0 + be[0].bitcast(fx.Float32) * m0
+        s1 = a1 + be[1].bitcast(fx.Float32) * m1
+        if live:
+            bo.buffer_store(
+                fx.Vector.from_elements(
+                    [s0.bitcast(fx.Int32), s1.bitcast(fx.Int32)], fx.Int32
+                ),
+                rsrc(stream_out),
+                tt * HID + c,
+            )
+        r2 = fx.ptr_load(rsl + (S + tt))
+        if live:
+            xw = bf16_pair((s0 * r2) * bf_lo(wi), (s1 * r2) * bf_hi(wi))
+            bo.buffer_store(xw, rsrc(xn_out), (tt * HID + c) // 2)
+            bo.buffer_store(
+                xw,
+                rsrc(scratch + fx.Int64(lay["xn"][0])),
+                (tt * HID + c) // 2,
+                cache_modifier=CM_DEV,
+            )
+        rocdl.s_waitcnt(vmcnt=0, expcnt=0, lgkmcnt=0)
+        gpu.barrier()
+        if tid == 0:
+            mb_put(tag, scratch + fx.Int64(lay["xnf"][0]), bid, fx.Int32(1))
+
+    # ------------------------------------------------------------ stage O: o-norm x gate, o_proj, push
+    def o_loads(w_o, bid, lane, wave):
+        """o_proj rows 32b .. 32b+32 (2 tiles) x this wave's CO chunks of K = OK_K (2048 / 1024)."""
+        ops = []
+        for i in range_constexpr(2):
+            for ci in range_constexpr(CO):
+                c = wave + WAVES * ci
+                ops.append(
+                    bf_tile_load_mma(w_o, bid * 32 + 16 * i, OK_K, c * KCH, lane)
+                )
+        return ops
+
+    @traced
+    def stage_o(
+        o_ops,
+        attn_o,
+        attn_g,
+        g_stride,
+        pl,
+        tag,
+        tid,
+        bid,
+        lane,
+        wave,
+        rank,
+        yl,
+        red,
+        ao=None,
+    ):
+        """Y = bf16(o rms(o) sigmoid(g)) per 64-wide head (the gated o-norm: fp32,
+        one rounding) into LDS for every CTA, then o_proj rows 32b.. (this rank's K slice),
+        the bf16 partial pushed into slot [par][rank] of the block owner's OP."""
+        tid = fresh(tid)  # no cross-stage CSE of tid-derived offsets (S=16 spills)
+        lane = fresh(lane)
+        par = tag_par(tag)
+        g = lane // 16
+        npc = S * OK_K // 8
+        for i in range_constexpr((npc + THREADS - 1) // THREADS):
+            e = tid + THREADS * i
+            t = fx.min(e // (OK_K // 8), S - 1)
+            kk = (e % (OK_K // 8)) * 8
+            if const_expr(ao is None):
+                xv = ld_vec(attn_o, (t * OK_K + kk) // 2, 4)
+            else:  # the fused attention: the merged rows as tagged pairs
+                got = mb_poll(
+                    tag,
+                    ao[1],
+                    [
+                        (ao[0], (t * OK_K + kk) // 2, 2),
+                        (ao[0], (t * OK_K + kk) // 2 + 2, 2),
+                    ],
+                    43,
+                )
+                xv = got[0] + got[1]
+            gv = ld_vec(attn_g, (t * g_stride + kk) // 2, 4)
+            ow = o_gate(xv, gv)
+            if e < npc:
+                fx.ptr_store(
+                    fx.Vector.from_elements(ow, fx.Int32), yl + (t * YROW + kk) // 2
+                )
+        gpu.barrier()
+        t = fx.min(lane % 16, S - 1)
+        accs = [fx.Vector.filled(4, 0.0, fx.Float32) for _ in range(2)]
+        for i in range_constexpr(2):
+            for ci in range_constexpr(CO):
+                c = wave + WAVES * ci
+                for qq in range_constexpr(4):
+                    kk = c * KCH + 8 * g + 32 * qq
+                    b = fx.Vector(
+                        fx.ptr_load(
+                            yl + (t * YROW + kk) // 2,
+                            result_type=fx.Vector.make_type(4, fx.Int32),
+                        )
+                    ).bitcast(fx.BFloat16)
+                    aop = o_ops[i * CO + ci][qq]
+                    accs[i] = mfma_bf16(aop.bitcast(fx.BFloat16), b, accs[i])
+        for i in range_constexpr(2):
+            fx.ptr_store(accs[i], red + ((wave * 2 + i) * 64 + lane) * 4)
+        gpu.barrier()
+        if tid < 16 * S:
+            pp = tid % 16  # row pair of the CTA's 32 rows
+            tt = tid // 16
+            v = []
+            for h2 in range_constexpr(2):
+                rr = 2 * pp + h2
+                s_ = fx.Float32(0.0)
+                for w in range_constexpr(WAVES):
+                    s_ = s_ + fx.ptr_load(
+                        red
+                        + (
+                            ((w * 2 + rr // 16) * 64 + 16 * ((rr % 16) // 4) + tt) * 4
+                            + rr % 4
+                        )
+                    )
+                v.append(s_)
+            w2 = bf16_pair(v[0], v[1])
+            pd = sel_peer(pl, bid // CPRX)
+            mb_put_sys(
+                tag,
+                pd + fx.Int64(play["op"][0]),
+                ((par * W + rank) * S + tt) * (HID // 2) + bid * 16 + pp,
+                w2,
+            )
+
+    # ------------------------------------------------------------ stage AB: attention boundary
+    @traced
+    def stage_ab(
+        peer,
+        scratch,
+        stream_pre,
+        alpha_a,
+        beta_a,
+        w_pffn,
+        tag,
+        health,
+        tid,
+        bid,
+        lane,
+        rsl,
+        pl=None,
+        rank=None,
+    ):
+        """CTA b's hidden columns 32b ..: the W o_proj partials in rank order (fp32 -> bf16,
+        the all-reduce output), post_attn_norm (weightless, 1e-8), gated residual on the
+        fp32 stream, pre_feedforward_layernorm -> ffn_in (local, + flag) and the mid stream."""
+        tid = fresh(tid)  # no cross-stage CSE of tid-derived offsets (S=16 spills)
+        par = tag_par(tag)
+        pp = tid % 16
+        tt = fx.min(tid // 16, S - 1)
+        live = tid < 16 * S
+        c = bid * 32 + 2 * pp
+        st = ld_vec(stream_pre, tt * HID + c, 2)
+        al = ld_vec(alpha_a, c, 2)
+        be = ld_vec(beta_a, c, 2)
+        wf = ld_i32(w_pffn, c // 2)
+        o0 = fx.Float32(0.0)
+        o1 = fx.Float32(0.0)
+        # the block owner: W partials in rank order -> bf16 -> every rank
+        if live & ((bid // CPRX) == rank):
+            specs = [
+                (
+                    peer + fx.Int64(play["op"][0]),
+                    ((par * W + src) * S + tt) * (HID // 2) + bid * 16 + pp,
+                    1,
+                )
+                for src in range(W)
+            ]
+            got = mb_poll(tag, health, specs, REGION_IDS["op"], cm=CM_SYS)
+            a0 = bf_lo(got[0][0])
+            a1 = bf_hi(got[0][0])
+            for src in range_constexpr(1, W):
+                a0 = a0 + bf_lo(got[src][0])
+                a1 = a1 + bf_hi(got[src][0])
+            wr = bf16_pair(a0, a1)
+            oo0 = bf_lo(wr)
+            oo1 = bf_hi(wr)
+            pa0 = al[0].bitcast(fx.Float32) * st[0].bitcast(fx.Float32)
+            pa1 = al[1].bitcast(fx.Float32) * st[1].bitcast(fx.Float32)
+            pb0 = be[0].bitcast(fx.Float32) * oo0
+            pb1 = be[1].bitcast(fx.Float32) * oo1
+            p1 = butterfly(oo0 * oo0 + oo1 * oo1, (1, 2, 4, 8))
+            p2 = butterfly(pa0 * pa0 + pa1 * pa1, (1, 2, 4, 8))
+            p3 = butterfly(pa0 * pb0 + pa1 * pb1, (1, 2, 4, 8))
+            p4 = butterfly(pb0 * pb0 + pb1 * pb1, (1, 2, 4, 8))
+            if pp == 0:
+                i0 = ((par * S + tt) * BLOCKS + bid) * 4
+                for pb in pl:
+                    base = pb + fx.Int64(play["ssa4p"][0])
+                    bo.buffer_store(
+                        fx.Vector.from_elements(
+                            [p1.bitcast(fx.Int32), tag, p2.bitcast(fx.Int32), tag],
+                            fx.Int32,
+                        ),
+                        rsrc(base),
+                        i0 * 2,
+                        cache_modifier=CM_SYS,
+                    )
+                    bo.buffer_store(
+                        fx.Vector.from_elements(
+                            [p3.bitcast(fx.Int32), tag, p4.bitcast(fx.Int32), tag],
+                            fx.Int32,
+                        ),
+                        rsrc(base),
+                        (i0 + 2) * 2,
+                        cache_modifier=CM_SYS,
+                    )
+            for pb in pl:
+                mb_put_sys(
+                    tag,
+                    pb + fx.Int64(play["orr"][0]),
+                    (par * S + tt) * (HID // 2) + bid * 16 + pp,
+                    wr,
+                )
+        if live:
+            got = mb_poll(
+                tag,
+                health,
+                [
+                    (
+                        peer + fx.Int64(play["orr"][0]),
+                        (par * S + tt) * (HID // 2) + bid * 16 + pp,
+                        1,
+                    )
+                ],
+                REGION_IDS["orr"],
+                cm=CM_SYS,
+            )
+            o0 = bf_lo(got[0][0])
+            o1 = bf_hi(got[0][0])
+        a0 = al[0].bitcast(fx.Float32) * st[0].bitcast(fx.Float32)
+        a1 = al[1].bitcast(fx.Float32) * st[1].bitcast(fx.Float32)
+        tok_rstd2(
+            scratch,
+            "ssa4",
+            tag,
+            health,
+            tid,
+            lane,
+            rsl,
+            POST_EPS,
+            PRE_EPS,
+            bid,
+            rsl + 2 * S,
+            pbase=peer
+            + fx.Int64(play["ssa4p"][0])
+            + fx.Int64(par * S * BLOCKS * 4 * 8),
+            pcm=CM_SYS,
+        )
+        r1 = fx.ptr_load(rsl + tt)
+        n0 = bf16_round(o0 * r1)
+        n1 = bf16_round(o1 * r1)
+        s0 = a0 + be[0].bitcast(fx.Float32) * n0
+        s1 = a1 + be[1].bitcast(fx.Float32) * n1
+        if live:
+            bo.buffer_store(
+                fx.Vector.from_elements(
+                    [s0.bitcast(fx.Int32), s1.bitcast(fx.Int32)], fx.Int32
+                ),
+                rsrc(scratch + fx.Int64(lay["smid"][0])),
+                tt * HID + c,
+                cache_modifier=CM_DEV,
+            )
+        r2 = fx.ptr_load(rsl + (S + tt))
+        if live:
+            bo.buffer_store(
+                bf16_pair((s0 * r2) * bf_lo(wf), (s1 * r2) * bf_hi(wf)),
+                rsrc(scratch + fx.Int64(lay["xin"][0])),
+                (tt * HID + c) // 2,
+                cache_modifier=CM_DEV,
+            )
+        rocdl.s_waitcnt(vmcnt=0, expcnt=0, lgkmcnt=0)
+        gpu.barrier()
+        if tid == 0:
+            mb_put(tag, scratch + fx.Int64(lay["xf"][0]), bid, fx.Int32(1))
+
+    # ------------------------------------------------------------ stage QKV: next layer's qkv+gate
+    # Head-aligned tile map: group gq < NH = HQ + 2 HKV holds the 4 tiles of head gq (q heads, then k,
+    # then v) + gate tile NQKVT + gq; groups NH .. hold the other gate tiles (5 a group, the tail padded).
+    # After the K-split partial sums, CTA k8 of the group owns the tokens t = k8 (mod 8) for all its rows,
+    # so it holds whole heads and finishes them (qk-norm, RoPE, q / KV stores). The qkv rows also go to
+    # the scratch region. The x eighth is the MFMA B operand straight from registers; the weights load in
+    # the MFMA A-operand layout (bf_tile_load_mma).
+    NH = HQ + 2 * HKV
+    NQKVT = QKV_ROWS // 16
+    assert NQT == 5 and NH <= 32 and S % 8 == 0, (NQT, NH, S)
+
+    def q8_tile(gq, n):
+        """Tile index (traced) of chunk-tile n of group gq (the head-aligned map)."""
+        head = gq < NH
+        if const_expr(n < 4):
+            return head.select(4 * gq + n, NQKVT + NH + NQT * (gq - NH) + n)
+        return head.select(NQKVT + gq, NQKVT + NH + NQT * (gq - NH) + n)
+
+    def q8_loads(w_qkvg, bid, lane, wave, n0, n1):
+        k8 = bid % 8
+        gq = bid // 8
+        out = []
+        for n in range_constexpr(n0, n1):
+            tl = q8_tile(gq, n)
+            row16 = fx.min(tl, NQTILES - 1) * 16
+            out.append(
+                bf_tile_load_mma(w_qkvg, row16, HID, k8 * XQ8 + wave * KCH, lane)
+            )
+        return out
+
+    @traced
+    def vllm_finish(w2, hh, tt, p, q_out, k_cache, v_cache, slots, cos_sin, pos):
+        """The owner's head finish: head hh (uniform), token tt (uniform in the wave), lane p = pair
+        (2p, 2p + 1) of the head; w2 = the bf16 pair of the qkv GEMV.
+        q / k as vLLM's fused_muse_spark_qknorm_rope_kv_insert: weightless qk-norm, then GPT-J RoPE in
+        fp32 with the fp32 cos / sin, one RNE bf16 rounding (bf16_pair); q without qk_scale. V passes
+        through (the bf16 GEMV). K / V pairs go as one bf16 dword into the LBHNC cache, element
+        ((slot >> 4) * bstride + (kv * 16 + (slot & 15)) * 128 + 64 isv + 2p), bstride = v_cache (the
+        block stride in elements, a value, not a pointer); slot < 0 (padded rows) skipped. The block
+        base is 64-bit and wave-uniform (one buffer resource a wave)."""
+        x0 = bf_lo(w2)
+        x1 = bf_hi(w2)
+        ss = butterfly(x0 * x0 + x1 * x1, (1, 2, 4, 8, 16))
+        rs = hw_rsq(ss * fx.Float32(1.0 / HD) + fx.Float32(PRE_EPS))
+        isv = hh >= HQ + HKV
+        isq = hh < HQ
+        n0 = x0 * rs
+        n1 = x1 * rs
+        if const_expr(key.next_rope):
+            ps = ld_i32(pos, tt * 2)  # int64 positions: low word
+            c = ld_f32(cos_sin, ps * HD + p)
+            sn = ld_f32(cos_sin, ps * HD + HD // 2 + p)
+            r0 = n0 * c - n1 * sn
+            r1 = n0 * sn + n1 * c
+            n0 = r0
+            n1 = r1
+        o0 = isv.select(x0, n0)
+        o1 = isv.select(x1, n1)
+        row0 = (hh * HD) // 2 + p
+        if isq:
+            bo.buffer_store(bf16_pair(o0, o1), rsrc(q_out), (tt * QROWS) // 2 + row0)
+        else:
+            kv = hh - HQ - isv.select(fx.Int32(HKV), fx.Int32(0))
+            # int64 slot_mapping: low word (-1 = padding)
+            slot = uniform(ld_i32(slots, tt * 2))
+            if slot >= 0:
+                eoff = fx.Int64(slot >> 4) * v_cache + fx.Int64(
+                    (kv * 16 + (slot & 15)) * 128
+                    + isv.select(fx.Int32(HD), fx.Int32(0))
+                )
+                bo.buffer_store(
+                    bf16_pair(o0, o1), rsrc(k_cache + eoff * fx.Int64(2)), p
+                )
+
+    @traced
+    def stage_qkv8(
+        q_ops,
+        scratch,
+        q_out,
+        g_out,
+        k_cache,
+        v_cache,
+        slots,
+        cos_sin,
+        pos,
+        tag,
+        health,
+        tid,
+        bid,
+        lane,
+        wave,
+        xl,
+        red,
+        own,
+    ):
+        tid = fresh(tid)  # no cross-stage CSE of tid-derived offsets (S=16 spills)
+        lane = fresh(lane)
+        g = lane // 16
+        k8 = bid % 8
+        gq = bid // 8
+        if tid < 32:
+            mb_poll(
+                tag,
+                health,
+                [(scratch + fx.Int64(lay["xnf"][0]), k8 * 32 + tid, 1)],
+                REGION_IDS["xnf"],
+            )
+        gpu.barrier()
+        t = fx.min(lane % 16, S - 1)
+        xb = []
+        for qq in range_constexpr(4):
+            kk = wave * KCH + 8 * g + 32 * qq
+            xb.append(
+                ld_vec(
+                    scratch + fx.Int64(lay["xn"][0]),
+                    (t * HID + k8 * XQ8 + kk) // 2,
+                    4,
+                    CM_DEV,
+                )
+            )
+        rocdl.sched_barrier(0)  # keep the x loads issued here, before the MFMA loop
+        for n in range_constexpr(NQT):
+            acc = fx.Vector.filled(4, 0.0, fx.Float32)
+            for qq in range_constexpr(4):
+                b = xb[qq].bitcast(fx.BFloat16)
+                aop = q_ops[n][qq]
+                acc = mfma_bf16(aop.bitcast(fx.BFloat16), b, acc)
+            fx.ptr_store(acc, red + ((wave * NQT + n) * 64 + lane) * 4)
+        gpu.barrier()
+        SO = S // 8  # tokens an owner (t = k8 + 8 j)
+        for it_ in range_constexpr((NQT * 8 * S + THREADS - 1) // THREADS):
+            e = tid + THREADS * it_
+            if e < NQT * 8 * S:
+                n = e // (8 * S)
+                rp = (e // S) % 8
+                tt = e % S
+                v = []
+                for h2 in range_constexpr(2):
+                    rr = 2 * rp + h2
+                    s_ = fx.Float32(0.0)
+                    for w in range_constexpr(WAVES):
+                        s_ = s_ + fx.ptr_load(
+                            red
+                            + (((w * NQT + n) * 64 + 16 * (rr // 4) + tt) * 4 + rr % 4)
+                        )
+                    v.append(s_)
+                if (tt % 8) == k8:
+                    for h2 in range_constexpr(2):
+                        fx.ptr_store(
+                            v[h2], own + ((n * 16 + 2 * rp + h2) * SO + tt // 8)
+                        )
+                else:
+                    for h2 in range_constexpr(2):
+                        mb_put(
+                            tag,
+                            scratch + fx.Int64(lay["qx"][0]),
+                            (((gq * 8 + k8) * NQT + n) * 16 + 2 * rp + h2) * S + tt,
+                            v[h2],
+                        )
+        gpu.barrier()
+        # owner wave j: token tt = k8 + 8 j; lane p < 40 = (tile n = p / 8, row pair p % 8); a head's 32 pairs
+        # are lanes 0 .. 31 (pair p of the head = lane p)
+        if (tid < 64 * SO) & (lane < NQT * 8):
+            j = tid // 64
+            p = lane
+            n = p // 8
+            rp = p % 8
+            tt = k8 + 8 * j
+            specs, oks = [], []
+            for src in range_constexpr(8):
+                for h2 in range_constexpr(2):
+                    specs.append(
+                        (
+                            scratch + fx.Int64(lay["qx"][0]),
+                            (((gq * 8 + src) * NQT + n) * 16 + 2 * rp + h2) * S + tt,
+                            1,
+                        )
+                    )
+                    oks.append(fx.Int32(src) != k8)
+            got = mb_poll(tag, health, specs, REGION_IDS["qx"], oks=oks)
+            tot = []
+            for h2 in range_constexpr(2):
+                s_ = fx.Float32(0.0)
+                for src in range_constexpr(8):
+                    mine = fx.ptr_load(own + ((n * 16 + 2 * rp + h2) * SO + j))
+                    s_ = s_ + (fx.Int32(src) == k8).select(
+                        mine, got[2 * src + h2][0].bitcast(fx.Float32)
+                    )
+                tot.append(s_)
+            head = gq < NH
+            tl = head.select(4 * gq + n, NQKVT + NH + NQT * (gq - NH) + n)
+            tl = (head & (n == 4)).select(NQKVT + gq, tl)
+            row = tl * 16 + 2 * rp
+            w2 = bf16_pair(tot[0], tot[1])
+            if row < QKV_ROWS:
+                mb_put(
+                    tag,
+                    scratch + fx.Int64(lay["qkv"][0]),
+                    tt * (QKV_ROWS // 2) + row // 2,
+                    w2,
+                )
+            else:
+                st_if(row < QKVG_ROWS, w2, g_out, (tt * QROWS + row - QKV_ROWS) // 2)
+            if head & (p < 32):
+                vllm_finish(w2, gq, tt, p, q_out, k_cache, v_cache, slots, cos_sin, pos)
+
+    # ------------------------------------------------------------ kernel
+    @traced
+    def front2(
+        own_peer,
+        bias_v,
+        w13,
+        w13s,
+        w_pren,
+        scratch,
+        tag,
+        health,
+        rank,
+        tid,
+        lane,
+        wave,
+        rt,
+        bid,
+    ):
+        lgs_r = route_poll(own_peer, tag, health, lane, wave)
+        stage_route(scratch, bias_v, lgs_r, tag, health, rank, tid, lane, wave, rt)
+        nu = uniform(fx.ptr_load(rt["nu"] + 0))
+        t0 = bid * nu
+        t1 = t0 + nu
+        last = t1 - 1
+        q = [gu_loads(w13, w13s, rt, t0, lane, wave, t0 < t1)]
+        hv, wv = h_wait_issue(own_peer, w_pren, tag, health, tid, rank)
+        q = q + [
+            gu_loads(w13, w13s, rt, t0 + d, lane, wave, (t0 + d) < t1)
+            for d in range(1, PF_GU)
+        ]
+        return q, hv, wv, t0, t1, last
+
+    # The kernel and its launcher. tl (a timeline buffer) and w13r .. eptab (expert replicas) are
+    # part of the launch signature the decoder passes; this build does not read them.
+
+    @flyc.kernel(name=name, known_block_size=[THREADS, 1, 1])
+    def k3(
+        attn_o: Int64,
+        attn_g: Int64,
+        stream_pre: Int64,
+        w_o: Int64,
+        alpha_a: Int64,
+        beta_a: Int64,
+        w_pffn: Int64,
+        w_r: Int64,
+        bias: Int64,
+        w_pre: Int64,
+        w_pren: Int64,
+        w13: Int64,
+        w13s: Int64,
+        w2: Int64,
+        w2s: Int64,
+        w_post: Int64,
+        w_pp: Int64,
+        w_pff: Int64,
+        alpha: Int64,
+        beta: Int64,
+        w_in: Int64,
+        w_qkvg: Int64,
+        stream_out: Int64,
+        xn_out: Int64,
+        part_out: Int64,
+        q_out: Int64,
+        g_out: Int64,
+        k_cache: Int64,
+        v_cache: Int64,
+        slots: Int64,
+        cos_sin: Int64,
+        pos: Int64,
+        scratch: Int64,
+        epoch: Int64,
+        health: Int64,
+        tl: Int64,
+        peer0: Int64,
+        peer1: Int64,
+        peer2: Int64,
+        peer3: Int64,
+        peer4: Int64,
+        peer5: Int64,
+        peer6: Int64,
+        peer7: Int64,
+        layer: Int32,
+        rank: Int32,
+        g_stride: Int32,
+        w13r: Int64,
+        w13sr: Int64,
+        w2r: Int64,
+        w2sr: Int64,
+        eptab: Int64,
+    ):
+        tid = fx.thread_idx.x
+        bid = fx.block_idx.x
+        lane = tid % 64
+        wave = tid // 64
+        alloc = fx.SharedAllocator()
+        alloc.allocate(Reserved)
+        rl = alloc.allocate(RouteLds).peek()
+        sl = alloc.allocate(StageLds)
+        sa, sc_, sd, se, sf = (
+            sl.a_.peek(),
+            sl.c_.peek(),
+            sl.d_.peek(),
+            sl.e_.peek(),
+            sl.f_.peek(),
+        )
+        so, sq = sl.o_.peek(), sl.q_.peek()
+        rt = {
+            "rid": rl.rid.ptr,
+            "rw": rl.rw.ptr,
+            "uexp": rl.uexp.ptr,
+            "kof": rl.kof.ptr,
+            "nu": rl.nu.ptr,
+            "sel": rl.sel.ptr,
+            "plist": rl.plist.ptr,
+            "pidx": rl.pidx.ptr,
+            "np_": rl.np_.ptr,
+            "rstd": rl.rstd.ptr,
+            "rank": rank,
+            "uslot": rl.uslot.ptr,
+            "ssacc": rl.ssacc.ptr,
+            "upk": rl.upk.ptr,
+            "upn": rl.upn.ptr,
+        }
+        tok = rl.tok.ptr
+        pl = [peer0, peer1, peer2, peer3, peer4, peer5, peer6, peer7]
+        own_peer = sel_peer(pl, rank)
+        ep = uniform(ld_i32(epoch, 0, CM_DEV))
+        tag = (ep & 0x1FFFFFF) * 64 + layer + 1
+        ao = None
+        if const_expr(FA is not None):
+            # attn_o is the fragment's descriptor; stage O polls the merged rows (tagged pairs)
+            attn_o, ao_t = FA.run(
+                attn_o, tag, health, tid, bid, lane, wave, sl.fa_.peek()
+            )
+            ao = (ao_t, health)
+        o_ops = o_loads(w_o, bid, lane, wave)
+        stage_o(
+            o_ops,
+            attn_o,
+            attn_g,
+            g_stride,
+            pl,
+            tag,
+            tid,
+            bid,
+            lane,
+            wave,
+            rank,
+            so.y.ptr,
+            so.red.ptr,
+            ao,
+        )
+        stage_ab(
+            own_peer,
+            scratch,
+            stream_pre,
+            alpha_a,
+            beta_a,
+            w_pffn,
+            tag,
+            health,
+            tid,
+            bid,
+            lane,
+            rl.rs.ptr,
+            pl,
+            rank,
+        )
+        bias_v = stage_a(
+            0,
+            w_r,
+            bias,
+            w_pre,
+            scratch,
+            own_peer,
+            pl,
+            tag,
+            health,
+            tid,
+            bid,
+            lane,
+            wave,
+            rank,
+            sa.x.ptr,
+            sa.red.ptr,
+            sa.own.ptr,
+            tok,
+        )
+        q, hv, wv, t0, t1, last = front2(
+            own_peer,
+            bias_v,
+            w13,
+            w13s,
+            w_pren,
+            scratch,
+            tag,
+            health,
+            rank,
+            tid,
+            lane,
+            wave,
+            rt,
+            bid,
+        )
+        load_h_norm(hv, wv, tid, lane, wave, sc_.h.ptr, tok)
+        gu_ld = lambda t: gu_loads(w13, w13s, rt, t, lane, wave, t < t1)
+        gu_c = lambda ops, t, i: gu_compute(
+            ops, t, i, scratch, tag, rt, tid, lane, wave, sc_.h.ptr, sc_.red.ptr
+        )
+        task_ring(gu_ld, gu_c, q, t0, fx.Int32(0), t1, UNR)
+        dn_ld = lambda t: dn_loads(w2, w2s, w_post, rt["uexp"], t, lane, wave, t < t1)
+        dq = [dn_loads(w2, w2s, w_post, rt["uexp"], t0, lane, wave, t0 < t1)]
+        gpu.barrier()
+        u0 = t0 >> 8
+        u1 = last >> 8
+        dn_build_picks(rt, u0, u1, tid)
+        load_ac(scratch, tag, health, rt, False, sd.a.ptr, tid)
+        dq = dq + [dn_ld(t0 + d) for d in range(1, PF_DN)]
+        if tid < 2 * S:
+            fx.ptr_store(fx.Float32(0.0), rt["ssacc"] + tid)
+        gpu.barrier()
+        tsplit = fx.min(t1, (u0 + 1) << 8)
+        dn_c = lambda ops, t, i: dn_compute_c(
+            ops, t, i, u0, scratch, tag, rt, tid, lane, wave, sd.a.ptr, sd.red.ptr
+        )
+        dq, task, it = task_ring(dn_ld, dn_c, dq, t0, fx.Int32(0), tsplit, UNR)
+        dn_switch(scratch, tag, health, rt, tsplit, u0, u1, sd.a.ptr, tid)
+        task_ring(dn_ld, dn_c, dq, task, it, t1, UNR)
+        gpu.barrier()
+        if tid < 2 * S:
+            h2 = tid // S
+            tt = tid % S
+            u = u0 + h2
+            k = fx.ptr_load(rt["kof"] + (fx.min(u, EL - 1) * S + tt))
+            if ((h2 == 0) | (u1 != u0)) & (k >= 0):
+                cf = (u * 256) // (t1 - t0)
+                mb_put(
+                    tag,
+                    scratch + fx.Int64(lay["ss2"][0]),
+                    (tt * TOPK + k) * 256 + (bid - cf),
+                    fx.ptr_load(rt["ssacc"] + tid),
+                )
+        # post_expert_proj weights in flight before the all-reduce wait (pp_early: before
+        # the final stage's polls too)
+        ppo = pp_loads(w_pp, bid, lane, wave)
+        gpu.barrier()
+        stage_final(scratch, part_out, tag, health, rt, tid, bid, lane, se.z.ptr)
+        final_push(scratch, part_out, pl, tag, rt, tid, bid, lane, se.z.ptr, rank)
+        ar_reduce(scratch, own_peer, tag, health, tid, bid, pl, rank)
+        post_proj(
+            ppo,
+            scratch,
+            own_peer,
+            pl,
+            tag,
+            health,
+            tid,
+            bid,
+            lane,
+            wave,
+            rank,
+            sf.mm.ptr,
+            sf.red.ptr,
+        )
+        qo = q8_loads(w_qkvg, bid, lane, wave, 0, QPF)
+        stage_g(
+            own_peer,
+            scratch,
+            0,
+            w_pff,
+            alpha,
+            beta,
+            w_in,
+            stream_out,
+            xn_out,
+            tag,
+            health,
+            tid,
+            bid,
+            lane,
+            rl.rs.ptr,
+        )
+        qo = qo + q8_loads(w_qkvg, bid, lane, wave, QPF, NQT)
+        stage_qkv8(
+            qo,
+            scratch,
+            q_out,
+            g_out,
+            k_cache,
+            v_cache,
+            slots,
+            cos_sin,
+            pos,
+            tag,
+            health,
+            tid,
+            bid,
+            lane,
+            wave,
+            sq.xq.ptr,
+            sq.red.ptr,
+            sq.own.ptr,
+        )
+        if (bid == 0) & (tid < P):
+            bo.buffer_store(
+                fx.ptr_load(rt["rid"] + tid),
+                rsrc(scratch + fx.Int64(lay["route"][0])),
+                tid,
+            )
+            bo.buffer_store(
+                fx.ptr_load(rt["rw"] + tid),
+                rsrc(scratch + fx.Int64(lay["route"][0])),
+                P + tid,
+            )
+
+    @flyc.jit
+    def launch(
+        attn_o: Int64,
+        attn_g: Int64,
+        stream_pre: Int64,
+        w_o: Int64,
+        alpha_a: Int64,
+        beta_a: Int64,
+        w_pffn: Int64,
+        w_r: Int64,
+        bias: Int64,
+        w_pre: Int64,
+        w_pren: Int64,
+        w13: Int64,
+        w13s: Int64,
+        w2: Int64,
+        w2s: Int64,
+        w_post: Int64,
+        w_pp: Int64,
+        w_pff: Int64,
+        alpha: Int64,
+        beta: Int64,
+        w_in: Int64,
+        w_qkvg: Int64,
+        stream_out: Int64,
+        xn_out: Int64,
+        part_out: Int64,
+        q_out: Int64,
+        g_out: Int64,
+        k_cache: Int64,
+        v_cache: Int64,
+        slots: Int64,
+        cos_sin: Int64,
+        pos: Int64,
+        scratch: Int64,
+        epoch: Int64,
+        health: Int64,
+        tl: Int64,
+        peer0: Int64,
+        peer1: Int64,
+        peer2: Int64,
+        peer3: Int64,
+        peer4: Int64,
+        peer5: Int64,
+        peer6: Int64,
+        peer7: Int64,
+        layer: Int32,
+        rank: Int32,
+        g_stride: Int32,
+        w13r: Int64,
+        w13sr: Int64,
+        w2r: Int64,
+        w2sr: Int64,
+        eptab: Int64,
+        stream: fx.Stream = fx.Stream(None),
+    ):
+        _ = _KEYSTR
+        k3(
+            attn_o,
+            attn_g,
+            stream_pre,
+            w_o,
+            alpha_a,
+            beta_a,
+            w_pffn,
+            w_r,
+            bias,
+            w_pre,
+            w_pren,
+            w13,
+            w13s,
+            w2,
+            w2s,
+            w_post,
+            w_pp,
+            w_pff,
+            alpha,
+            beta,
+            w_in,
+            w_qkvg,
+            stream_out,
+            xn_out,
+            part_out,
+            q_out,
+            g_out,
+            k_cache,
+            v_cache,
+            slots,
+            cos_sin,
+            pos,
+            scratch,
+            epoch,
+            health,
+            tl,
+            peer0,
+            peer1,
+            peer2,
+            peer3,
+            peer4,
+            peer5,
+            peer6,
+            peer7,
+            layer,
+            rank,
+            g_stride,
+            w13r,
+            w13sr,
+            w2r,
+            w2sr,
+            eptab,
+        ).launch(grid=(BLOCKS,), block=(THREADS,), stream=stream)
+
+    _BUILDS[key] = launch
+    return launch
+
+
+def build_vllm(S=8, W=8, next_rope=True, attn=""):
+    """The K3 build vLLM launches for S rows (8 / 16) of a TP = EP = 8 group; next_rope: the next
+    layer has RoPE; attn = "swa" / "full": this layer's verify attention fused at the head."""
+    assert W == 8 and S in (8, 16), (S, W)
+    return Build(S=S, next_rope=bool(next_rope), attn=attn)
diff --git a/vllm/models/muse_spark/amd/mono/kernels/k5.py b/vllm/models/muse_spark/amd/mono/kernels/k5.py
new file mode 100644
index 000000000000..81f9588e5957
--- /dev/null
+++ b/vllm/models/muse_spark/amd/mono/kernels/k5.py
@@ -0,0 +1,2389 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# mypy: ignore-errors
+# ruff: noqa: B008, E501, SIM102
+
+# Vendored from AMD's Muse Spark FlyDSL monokernel kernels (mk/k5.py):
+# SPDX-License-Identifier: MIT
+# Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+"""HA and HB: one Muse Spark decoder layer for S = 32 / 64 / 96 / 128 rows as two persistent
+launches around HM (k5m.py, the local experts); one rank of a TP = EP = 8 group.
+
+    [attention] -> HA -> HM -> HB
+
+HA  optional: this layer's verify attention (attn_frag.py) at the head; o-norm x gate
+    (distributed) -> o_proj (rows 32 b .., K = this rank's 1024) -> the bf16 partials go to
+    the owner rank -> attention boundary (post_attn_norm, gated residual, pre-FFN norm) ->
+    ffn_in -> router (fp32, 32 experts a rank) and pre_expert_proj (512 rows a rank), K in 8
+    parts over 8-CTA groups -> logits and h_raw all-gathered -> top-8 (CTA t routes token t)
+    -> h = pre_expert_norm(h_raw) (CTA t: token t) -> h, top-k ids / weights for HM.
+HB  EP all-reduce of HM's partial (reduce-scatter to the column-block owner + all-gather) ->
+    post_expert_proj (this rank's 1024 rows, K in 4 parts) -> post-FFN norm, gated residual,
+    next input norm -> xn -> qkv + gate (K eighths) -> head finishers (qk-norm, RoPE,
+    q out, bf16 K / V into vLLM's LBHNC KV cache).
+
+The boundaries (Build.bnd). S > 32, token owners: rank r owns tokens t = r (mod 8); the o_proj
+partials and the post_expert_proj outputs go to the owner rank of each token, one CTA a token
+does the whole row (all-reduce sum, both norms, gated residual) and the ffn_in / xn rows are
+all-gathered. The residual stream lives token-sharded (rank r: its S / 8 rows, fp32, all 8192
+columns); HA reads it as rows so the first layer can read a full replicated stream (s_full);
+HB writes it the same way and (``gather``) all-gathers it. S = 32, column blocks: the o_proj
+partials go to the owner of the column block, which reduces and all-gathers them; both
+boundaries run on all 256 CTAs (token groups x column groups, one-hop token norms) and the
+stream is replicated.
+
+Every row tile is 16 tokens (MFMA N); dense stages keep their weights in VGPRs and loop over
+the S / 16 row tiles with the activations read straight from L2 / peer memory in MFMA
+layout. Tags: tag = epoch * 64 + layer + 1 (HA and HB of a layer share it, regions are
+disjoint); peer regions double buffered by tag parity.
+"""
+
+from dataclasses import dataclass
+
+import flydsl.compiler as flyc
+import flydsl.expr as fx
+from aiter.ops.flydsl.kernels import buffer_ops as bo
+from flydsl.expr import const_expr, gpu, range_constexpr, rocdl
+from flydsl.expr.typing import Int32, Int64
+
+from .dev import (
+    BLOCKS,
+    CM_DEV,
+    CM_SYS,
+    THREADS,
+    WAVES,
+    _add,
+    bf16_pair,
+    bf16_round,
+    bf_hi,
+    bf_lo,
+    butterfly,
+    fresh,
+    hw_rsq,
+    ld_f32,
+    ld_i32,
+    ld_vec,
+    mb_poll,
+    mb_put,
+    mb_put_if,
+    mb_put_sys,
+    mfma_bf16,
+    readlane,
+    rsrc,
+    st_if,
+    traced,
+    uniform,
+    wave_sum,
+    xred,
+)
+from .k3 import _lay, bf_mma_operand, bf_tile_load, o_gate, route_token, sel_peer
+
+W = 8  # ranks in the TP = EP group
+HID = 8192
+QH = HID // W  # hidden columns a rank owns (= stage-A K part = post-proj rows)
+DM = 4096
+E = 256
+EL = E // W  # local experts = router experts a rank = stage-A groups
+KQ = BLOCKS // EL  # stage-A K parts
+OWN = 16 // KQ  # pre-proj rows a CTA owns after the K-part exchange
+NRP = OWN // 2  # their row pairs
+NTILES = QH // 16  # post-proj 16-row tiles a rank
+KSP = BLOCKS // NTILES  # post-proj K split
+TOPK = 8
+PRE_EPS = 1e-5
+POST_EPS = 1e-8
+KCH = 128
+HQ, HKV, HD = 128 // W, 16 // W, 64  # q / kv heads a rank
+QROWS, KROWS = HQ * HD, HKV * HD
+QKV_ROWS = QROWS + 2 * KROWS
+QKVG_ROWS = QKV_ROWS + QROWS
+OK_K = HID // W  # o_proj K a rank
+NQTILES = QKVG_ROWS // 16  # qkvg 16-row tiles a rank
+# qkvg tiles a CTA group (32 groups x 8 K eighths; 144 of 160 real)
+NQT = -(-NQTILES // 32)
+TL_POINTS = 32
+# loads of peer-written data (after its system-scope flag): device scope = past the per-XCD L2,
+# served by the memory-side cache, which every write to this memory passes
+XCM = CM_DEV
+YCM = CM_DEV  # activations re-read by many CTAs
+ACM = CM_DEV  # stage A's x tiles
+HEALTH_WORDS = 8 + 8 * BLOCKS
+REGION_IDS = {
+    "yf": 1,
+    "opf": 2,
+    "xf": 6,
+    "logit": 7,
+    "px": 8,
+    "hf": 9,
+    "mf": 10,
+    "mmf": 11,
+    "py": 12,
+    "ytf": 13,
+    "xnf": 16,
+    "qx": 17,
+    "qkv": 18,
+    "sgf": 19,
+    # the column-block boundaries (bnd)
+    "opb": 20,
+    "orb": 21,
+    "yb": 22,
+    "xinbf": 23,
+    "xnbf": 24,
+    "ssa4": 25,
+    "ssg4": 26,
+    "gpa": 27,
+    "gpg": 28,
+}
+
+
+@dataclass(frozen=True)
+class Build:
+    S: int = 32  # rows (32 / 64 / 96 / 128)
+    next_rope: bool = True  # the next layer has RoPE (SWA) or not (full, NoPE)
+    # this layer's verify attention fused at the HA head ("swa" / "full")
+    attn: str = ""
+
+    @property
+    def bnd(self):
+        """The boundaries: 3 at S = 32 (o reduce-scattered to the column-block owner and
+        all-gathered, both boundaries on all 256 CTAs with one-hop token norms over column
+        groups, the stream replicated); 0 (S > 32): the token-owner boundaries (one CTA a
+        token, the stream token-sharded and all-gathered by HB)."""
+        return 3 if self.S == 32 else 0
+
+
+def scratch_layout(S, world=W, bnd=0):
+    assert world == W, world
+    TS = S // W
+    extra = (
+        []
+        if not bnd
+        else [
+            # the mid stream, replicated [t][col] fp32 (each CTA its 32 columns)
+            ("smidb", S * HID * 4),
+            ("xinb", S * HID * 2),
+            ("xinbf", BLOCKS * 8),  # ffn_in [t][col] bf16 + a flag a CTA
+            ("xnb", S * HID * 2),
+            ("xnbf", BLOCKS * 8),  # the next layer's input [t][col] bf16 + a flag a CTA
+            ("ssa4", (4 * S * BLOCKS + 2 * S) * 8),
+            ("ssg4", (4 * S * BLOCKS + 2 * S) * 8),  # partials, rstd
+            ("gpa", BLOCKS * 16 * 8),
+            # [CTA][4 tokens][4] partials of the AB / G
+            ("gpg", BLOCKS * 16 * 8),
+        ]
+    )
+    return _lay(
+        [
+            ("y", S * OK_K * 2),
+            ("yf", BLOCKS * 8),
+            ("smid", TS * HID * 4),
+            ("px", EL * KQ * 16 * S * 8),
+            ("m", S * DM * 2),
+            ("mmf", BLOCKS * 8),
+            # [K part - 1][tile][row][t]
+            ("py", max(128, (KSP - 1) * NTILES) * 16 * S * 8),
+            ("qx", 32 * 8 * NQT * 16 * S * 8),
+            ("qkv", S * (QKV_ROWS // 2) * 8),
+        ]
+        + extra
+    )
+
+
+def peer_layout(S, world=W, bnd=0):
+    assert world == W, world
+    TS = S // W
+    extra = (
+        []
+        if not bnd
+        else [
+            # o_proj partials to the block owner [par][src][t][own blk][pair]
+            ("opb", 2 * W * S * (BLOCKS // W) * 16 * 8),
+            # the owner's reduced o, all-gathered [par][t][pair] tagged
+            ("orb", 2 * S * (HID // 2) * 8),
+            # stage F's y, all-gathered [par][t][pair] tagged
+            ("yb", 2 * S * (HID // 2) * 8),
+        ]
+    )
+    # the bnd regions go last: HM (k5m) addresses moep / mf with the layout of the plain build
+    return _lay(
+        [
+            ("op", 2 * W * TS * (HID // 2) * 4),
+            ("opf", 2 * W * BLOCKS * 8),
+            ("xin", 2 * S * (HID // 2) * 4),
+            ("xf", 2 * S * 8),
+            ("logit", 2 * KQ * S * E * 8),
+            ("hraw", 2 * S * (DM // 2) * 4),
+            ("hf", 2 * W * BLOCKS * 8),
+            ("moep", 2 * W * S * (DM // 2) * 4),
+            ("mf", 2 * W * BLOCKS * 8),
+            ("moem", 2 * S * (DM // 2) * 4),
+            ("mmf", 2 * BLOCKS * 8),
+            ("yt", 2 * TS * (HID // 2) * 4),
+            ("ytf", 2 * W * NTILES * 8),
+            ("xn", 2 * S * (HID // 2) * 4),
+            ("xnf", 2 * S * 8),
+            ("sg", 2 * S * HID * 4),
+            ("sgf", 2 * S * 8),
+        ]
+        + extra
+    )
+
+
+_BUILDS = {}
+
+
+def build(key: Build):
+    if key in _BUILDS:
+        return _BUILDS[key]
+    _KEYSTR = repr(key)  # FlyDSL disk-cache key (see k3.build)
+    S = key.S
+    assert S in (32, 64, 96, 128), S
+    NT = S // 16
+    TS = S // W  # tokens a rank owns in the token-owner boundaries
+    CO = OK_K // KCH // WAVES  # o_proj chunks a wave
+    CA = QH // KCH // WAVES  # stage-A pre-proj chunks a wave
+    CP = (DM // KSP) // KCH // WAVES  # post-proj chunks a wave
+    KPT = QH // THREADS  # router K a thread
+    NFQ = W * BLOCKS // THREADS  # producer flags a thread (opf / hf polls)
+    QPF = 3  # qkvg tiles issued before stage G (of NQT)
+    CPR = BLOCKS // W  # CTAs (16-column blocks) a rank owns in the EP reduce-scatter
+    assert TS * W == S and NFQ * THREADS == W * BLOCKS and W * NTILES * 8 >= 512
+    BND = key.bnd != 0  # the column-block boundaries (S = 32)
+    QR1 = S == 32  # QKV5 reduces both 16-token row tiles in one LDS round
+    lay, _ = scratch_layout(S, W, key.bnd)
+    play, _ = peer_layout(S, W, key.bnd)
+    tag_par = lambda tag: tag & 1
+
+    @fx.struct
+    class RouteLds:
+        rid: fx.Array[fx.Int32, TOPK, 16]
+        rw: fx.Array[fx.Float32, TOPK, 16]
+        sel: fx.Array[fx.Int32, TOPK * 3 + 8, 16]
+
+    @fx.struct
+    class HaO:
+        red: fx.Array[fx.Float32, NT * WAVES * 2 * 64 * 4, 16]
+
+    XAROW = QH + 8  # LDS row of the x slice tile (bf16)
+
+    @fx.struct
+    class HaA:
+        x: fx.Array[fx.Int32, 16 * XAROW // 2, 16]
+        red: fx.Array[fx.Float32, WAVES * 64 * 4, 16]
+        tok: fx.Array[fx.Float32, WAVES * 16 + 8 + S, 16]
+        own: fx.Array[fx.Float32, NT * OWN * 16, 16]
+
+    FA = None
+    if key.attn:  # the fused attention fragment (attn_frag.py)
+        from . import attn_frag
+
+        assert key.attn in ("swa", "full"), key
+        FA = attn_frag.make(S // 8, key.attn == "swa")
+        FaLds = FA.Lds
+        # the fragment's source hash: the disk cache does not follow FA.run
+        _KEYSTR = _KEYSTR + FA.key
+
+    @fx.union
+    class HaLds:
+        o_: HaO
+        a_: HaA
+        if FA is not None:
+            fa_: FaLds
+
+    @fx.struct
+    class HbF:
+        red: fx.Array[fx.Float32, NT * WAVES * 64 * 4, 16]
+
+    @fx.struct
+    class HbQ:
+        red: fx.Array[fx.Float32, (NT if QR1 else 1) * WAVES * NQT * 64 * 4, 16]
+        own: fx.Array[fx.Float32, NT * NQT * 2 * 16, 16]
+
+    @fx.union
+    class HbLds:
+        f_: HbF
+        q_: HbQ
+
+    @fx.struct
+    class Misc:
+        pad: fx.Array[fx.Int64, TL_POINTS, 16]  # unused: keeps the tuned LDS offsets
+        # bnd: tok_rstd2 uses rs[2S .. 2S + 32)
+        rs: fx.Array[fx.Float32, 2 * S + (32 if BND else 8), 16]
+        # bnd: the polled group partials + wave sums
+        pr: fx.Array[fx.Float32, 512 + 64 if BND else 8, 16]
+        red4: fx.Array[fx.Float32, 32, 16]
+
+    name_a = f"muse_k5a_s{S}" + (f"_fa{key.attn}" if key.attn else "")
+    name_b = f"muse_k5b_s{S}_r{int(key.next_rope)}"
+
+    # ------------------------------------------------------------ shared helpers
+    def red_at(red, base_tile, rr, tt):
+        """LDS offset of row rr (0..15), token tt (0..15) of MFMA tile slot base_tile."""
+        return red + ((base_tile * 64 + 16 * (rr // 4) + tt) * 4 + rr % 4)
+
+    @traced
+    def block_sum(v, lane, wave, red4, slot):
+        """Sum of v over the CTA's 512 threads (fixed order: lanes, then waves)."""
+        v = wave_sum(v)
+        if lane == 0:
+            fx.ptr_store(v, red4 + (slot * 8 + wave))
+        gpu.barrier()
+        tot = fx.Float32(0.0)
+        for w in range_constexpr(WAVES):
+            tot = tot + fx.ptr_load(red4 + (slot * 8 + w))
+        return tot
+
+    # ============================================================ HA
+    def o_loads(w_o, bid, lane, wave):
+        ops = []
+        for i in range_constexpr(2):
+            for ci in range_constexpr(CO):
+                c = wave + WAVES * ci
+                ops.append(bf_tile_load(w_o, bid * 32 + 16 * i, OK_K, c * KCH, lane))
+        return ops
+
+    @traced
+    def stage_ng(attn_o, attn_g, g_stride, scratch, tag, tid, bid, nrows, ao=None):
+        """Y = bf16(o rms(o) sigmoid(g)) per 64-wide head, distributed: 8 lanes a (token, head)."""
+        tid = fresh(tid)
+        gt = bid * THREADS + tid
+        item = gt // 8
+        if item < S * HQ:
+            t = item // HQ
+            kk = (item % HQ) * HD + (gt % 8) * 8
+            # padded rows (t >= nrows) repeat the last real row
+            tl_ = fx.min(t, nrows - 1)
+            if const_expr(ao is None):
+                xv = ld_vec(attn_o, (tl_ * OK_K + kk) // 2, 4)
+            else:  # the fused attention: the merged rows as tagged pairs
+                got = mb_poll(
+                    tag,
+                    ao[1],
+                    [
+                        (ao[0], (tl_ * OK_K + kk) // 2, 2),
+                        (ao[0], (tl_ * OK_K + kk) // 2 + 2, 2),
+                    ],
+                    43,
+                )
+                xv = got[0] + got[1]
+            gv = ld_vec(attn_g, (tl_ * g_stride + kk) // 2, 4)
+            ow = o_gate(xv, gv)
+            bo.buffer_store(
+                fx.Vector.from_elements(ow, fx.Int32),
+                rsrc(scratch + fx.Int64(lay["y"][0])),
+                (t * OK_K + kk) // 2,
+                cache_modifier=CM_DEV,
+            )
+        rocdl.s_waitcnt(vmcnt=0, expcnt=0, lgkmcnt=0)
+        gpu.barrier()
+        if (tid == 0) & (bid < S * HQ * 8 // THREADS):
+            mb_put(tag, scratch + fx.Int64(lay["yf"][0]), bid, fx.Int32(1))
+
+    @traced
+    def stage_o(o_ops, scratch, pl, tag, health, tid, bid, lane, wave, rank, red):
+        """o_proj rows 32b .. (this rank's K slice) for all S tokens; the bf16 partial goes to
+        the owner rank of the column block (bnd: b // CPR) or of each token (t % W)."""
+        tid = fresh(tid)
+        lane = fresh(lane)
+        par = tag_par(tag)
+        g = lane // 16
+        t = lane % 16
+        if tid < S * HQ * 8 // THREADS:
+            mb_poll(
+                tag,
+                health,
+                [(scratch + fx.Int64(lay["yf"][0]), tid, 1)],
+                REGION_IDS["yf"],
+            )
+        gpu.barrier()
+        ybase = scratch + fx.Int64(lay["y"][0])
+        # MFMA A operands once (reused by every row tile)
+        oa = [
+            [
+                bf_mma_operand(o_ops[i * CO + ci], qq, lane).bitcast(fx.BFloat16)
+                for qq in range(4)
+            ]
+            for i in range(2)
+            for ci in range(CO)
+        ]
+
+        def ob_loads(rt):
+            return [
+                ld_vec(
+                    ybase,
+                    ((16 * rt + t) * OK_K + (wave + WAVES * ci) * KCH + 8 * g + 32 * qq)
+                    // 2,
+                    4,
+                    YCM,
+                ).bitcast(fx.BFloat16)
+                for ci in range(CO)
+                for qq in range(4)
+            ]
+
+        # row groups of 2 tiles (32 tokens): MFMA -> barrier -> (next group's loads) -> reduce + push,
+        # so the XGMI pushes of a group overlap the next group's loads and MFMAs
+        NG = NT // 2
+        nxt = ob_loads(0) + ob_loads(1)
+        for gi in range_constexpr(NG):
+            cur = nxt
+            for r2 in range_constexpr(2):
+                rt = 2 * gi + r2
+                accs = [fx.Vector.filled(4, 0.0, fx.Float32) for _ in range(2)]
+                for ci in range_constexpr(CO):
+                    for i in range_constexpr(2):
+                        for qq in range_constexpr(4):
+                            accs[i] = mfma_bf16(
+                                oa[i * CO + ci][qq],
+                                cur[r2 * (4 * CO) + ci * 4 + qq],
+                                accs[i],
+                            )
+                for i in range_constexpr(2):
+                    fx.ptr_store(
+                        accs[i], red + (((rt * WAVES + wave) * 2 + i) * 64 + lane) * 4
+                    )
+            gpu.barrier()
+            if const_expr(gi + 1 < NG):
+                nxt = ob_loads(2 * gi + 2) + ob_loads(2 * gi + 3)
+            if const_expr(BND):
+                # the CTA's whole [32 columns x tokens] tile to the owner of block b (b // 32), tagged pairs
+                e = tid
+                pp = e % 16
+                tt32 = e // 16
+                rt_ = 2 * gi + tt32 // 16
+                t16 = tt32 % 16
+                v = []
+                for h2 in range_constexpr(2):
+                    rr = 2 * pp + h2
+                    s_ = fx.Float32(0.0)
+                    for w in range_constexpr(WAVES):
+                        s_ = s_ + fx.ptr_load(
+                            red_at(red, (rt_ * WAVES + w) * 2 + rr // 16, rr % 16, t16)
+                        )
+                    v.append(s_)
+                pdo = sel_peer(pl, uniform(bid // CPR))
+                mb_put_sys(
+                    tag,
+                    pdo + fx.Int64(play["opb"][0]),
+                    ((par * W + rank) * S + 16 * rt_ + t16) * (CPR * 16)
+                    + (bid % CPR) * 16
+                    + pp,
+                    bf16_pair(v[0], v[1]),
+                )
+            else:
+                e = tid
+                # owner t % 8: wave w pushes owner w's 2 tokens of both tiles of the group
+                rt_ = 2 * gi + (e // 32) % 2
+                pp = e % 16
+                dw = uniform(wave)
+                tt = dw + W * ((e // 16) % 2)
+                v = []
+                for h2 in range_constexpr(2):
+                    rr = 2 * pp + h2
+                    s_ = fx.Float32(0.0)
+                    for w in range_constexpr(WAVES):
+                        s_ = s_ + fx.ptr_load(
+                            red_at(red, (rt_ * WAVES + w) * 2 + rr // 16, rr % 16, tt)
+                        )
+                    v.append(s_)
+                t_ = 16 * rt_ + tt
+                # one destination a wave (SGPR peer select): no branch around the store (a branch makes
+                # the waitcnt pass drain everything), and the links carry every group in parallel
+                idx = (((par * W + rank) * BLOCKS + bid) * TS + t_ // W) * 16 + pp
+                w2 = bf16_pair(v[0], v[1])
+                bo.buffer_store(
+                    w2,
+                    rsrc(sel_peer(pl, dw) + fx.Int64(play["op"][0])),
+                    idx,
+                    cache_modifier=CM_SYS,
+                )
+        if const_expr(not BND):  # bnd: tagged pairs, no flags
+            rocdl.s_waitcnt(vmcnt=0, expcnt=0, lgkmcnt=0)
+            gpu.barrier()
+            if tid == 0:
+                for pb in pl:
+                    mb_put_sys(
+                        tag,
+                        pb + fx.Int64(play["opf"][0]),
+                        (par * W + rank) * BLOCKS + bid,
+                        fx.Int32(1),
+                    )
+
+    @traced
+    def stage_ab(
+        own_peer,
+        pl,
+        scratch,
+        stream_in,
+        s_full,
+        nrows,
+        alpha_a,
+        beta_a,
+        w_pffn,
+        tag,
+        health,
+        tid,
+        bid,
+        lane,
+        wave,
+        rank,
+        red4,
+    ):
+        """CTA j < S / W: token t = W j + rank, all 8192 columns (16 a thread): the W o_proj
+        partials (rank order) -> o, post_attn_norm (weightless, 1e-8), gated residual (mid stream,
+        local), pre_feedforward_layernorm -> the ffn_in row pushed to every rank + flag."""
+        tid = fresh(tid)
+        par = tag_par(tag)
+        j = bid
+        t = W * j + rank
+        # replicated [rows, 8192] or token-sharded
+        srow = (s_full != 0).select(fx.min(t, nrows - 1), j)
+        st, al, be, wfs = [], [], [], []
+        for h in range_constexpr(2):
+            cq = 8 * (tid + THREADS * h)
+            wfs.append(ld_vec(w_pffn, cq // 2, 4))
+            for h2 in range_constexpr(2):
+                v4 = ld_vec(stream_in, srow * HID + cq + 4 * h2, 4)
+                a4 = ld_vec(alpha_a, cq + 4 * h2, 4)
+                b4 = ld_vec(beta_a, cq + 4 * h2, 4)
+                for d in range_constexpr(4):
+                    st.append(v4[d].bitcast(fx.Float32))
+                    al.append(a4[d].bitcast(fx.Float32))
+                    be.append(b4[d].bitcast(fx.Float32))
+        specs = [
+            (own_peer + fx.Int64(play["opf"][0]), par * W * BLOCKS + NFQ * tid + q, 1)
+            for q in range(NFQ)
+        ]
+        mb_poll(tag, health, specs, REGION_IDS["opf"], cm=CM_SYS)
+        gpu.barrier()
+        a = [fx.Float32(0.0) for _ in range(16)]
+        for src in range_constexpr(W):
+            for h in range_constexpr(2):
+                # 8-column chunk: columns 8q .. 8q + 8 (block q / 4)
+                q = tid + THREADS * h
+                wv = ld_vec(
+                    own_peer + fx.Int64(play["op"][0]),
+                    (((par * W + src) * BLOCKS + q // 4) * TS + j) * 16 + 4 * (q % 4),
+                    4,
+                    XCM,
+                )
+                for d in range_constexpr(4):
+                    k = 8 * h + 2 * d
+                    if const_expr(src == 0):
+                        a[k] = bf_lo(wv[d])
+                        a[k + 1] = bf_hi(wv[d])
+                    else:
+                        a[k] = a[k] + bf_lo(wv[d])
+                        a[k + 1] = a[k + 1] + bf_hi(wv[d])
+        o = [bf16_round(x) for x in a]
+        ss = fx.Float32(0.0)
+        for k in range_constexpr(16):
+            ss = ss + o[k] * o[k]
+        r1 = hw_rsq(
+            block_sum(ss, lane, wave, red4, 0) * fx.Float32(1.0 / HID)
+            + fx.Float32(POST_EPS)
+        )
+        sv = [al[k] * st[k] + be[k] * bf16_round(o[k] * r1) for k in range(16)]
+        for h in range_constexpr(2):
+            cq = 8 * (tid + THREADS * h)
+            for h2 in range_constexpr(2):
+                bo.buffer_store(
+                    fx.Vector.from_elements(
+                        [sv[8 * h + 4 * h2 + d].bitcast(fx.Int32) for d in range(4)],
+                        fx.Int32,
+                    ),
+                    rsrc(scratch + fx.Int64(lay["smid"][0])),
+                    j * HID + cq + 4 * h2,
+                    cache_modifier=CM_DEV,
+                )
+        ss2 = fx.Float32(0.0)
+        for k in range_constexpr(16):
+            ss2 = ss2 + sv[k] * sv[k]
+        r2 = hw_rsq(
+            block_sum(ss2, lane, wave, red4, 1) * fx.Float32(1.0 / HID)
+            + fx.Float32(PRE_EPS)
+        )
+        for h in range_constexpr(2):
+            cq = 8 * (tid + THREADS * h)
+            wf = wfs[h]
+            xw = [
+                bf16_pair(
+                    (sv[8 * h + 2 * d] * r2) * bf_lo(wf[d]),
+                    (sv[8 * h + 2 * d + 1] * r2) * bf_hi(wf[d]),
+                )
+                for d in range(4)
+            ]
+            xv = fx.Vector.from_elements(xw, fx.Int32)
+            for pb in pl:
+                bo.buffer_store(
+                    xv,
+                    rsrc(pb + fx.Int64(play["xin"][0])),
+                    (par * S + t) * (HID // 2) + cq // 2,
+                    cache_modifier=CM_SYS,
+                )
+        rocdl.s_waitcnt(vmcnt=0, expcnt=0, lgkmcnt=0)
+        gpu.barrier()
+        if tid == 0:
+            for pb in pl:
+                mb_put_sys(tag, pb + fx.Int64(play["xf"][0]), par * S + t, fx.Int32(1))
+
+    def a_loads(w_r, w_pre, bid, tid, lane, wave):
+        """CTA b = (group rb = b / KQ, K part kq = b % KQ): router expert rb of this rank (fp32,
+        KPT K a thread) and pre_expert_proj rows 16 rb .. (CA chunk tiles a wave)."""
+        kq = bid % KQ
+        rb = bid // KQ
+        wr = ld_vec(w_r, rb * HID + kq * QH + KPT * tid, KPT)
+        wp = []
+        for ci in range_constexpr(CA):
+            c = wave + WAVES * ci
+            wp.append(bf_tile_load(w_pre, rb * 16, HID, kq * QH + c * KCH, lane))
+        return wr, wp
+
+    @traced
+    def stage_a(
+        wr,
+        wp,
+        own_peer,
+        pl,
+        scratch,
+        tag,
+        health,
+        tid,
+        bid,
+        lane,
+        wave,
+        rank,
+        red,
+        tok,
+        own,
+        xl,
+    ):
+        tid = fresh(tid)
+        lane = fresh(lane)
+        par = tag_par(tag)
+        g = lane // 16
+        t = lane % 16
+        kq = bid % KQ
+        rb = bid // KQ
+        # ffn_in is local; the 32 CTAs of this K part wrote its columns
+        if const_expr(BND):
+            if tid < 32:
+                fa = (tid // 4) * 32 + 4 * kq + tid % 4
+                mb_poll(
+                    tag,
+                    health,
+                    [(scratch + fx.Int64(lay["xinbf"][0]), fa, 1)],
+                    REGION_IDS["xinbf"],
+                )
+        else:
+            if tid < S:
+                mb_poll(
+                    tag,
+                    health,
+                    [(own_peer + fx.Int64(play["xf"][0]), par * S + tid, 1)],
+                    REGION_IDS["xf"],
+                    cm=CM_SYS,
+                )
+        gpu.barrier()
+        if const_expr(BND):
+            xb = scratch + fx.Int64(lay["xinb"][0])
+            xpar = 0
+        else:
+            xb = own_peer + fx.Int64(play["xin"][0])
+            xpar = par
+        wf = [wr[d].bitcast(fx.Float32) for d in range(KPT)]
+        wa = [
+            [bf_mma_operand(wp[ci], qq, lane).bitcast(fx.BFloat16) for qq in range(4)]
+            for ci in range(CA)
+        ]
+        QV = QH // 8  # 16-B vectors a token row of the x tile
+        NXQ = 16 * QV // THREADS  # of them a thread
+
+        def ax_loads(rt):
+            return [
+                ld_vec(
+                    xb,
+                    (xpar * S + 16 * rt + (tid + THREADS * q) // QV) * (HID // 2)
+                    + (kq * QH + 8 * ((tid + THREADS * q) % QV)) // 2,
+                    4,
+                    ACM,
+                )
+                for q in range(NXQ)
+            ]
+
+        xnext = ax_loads(0)
+        for rt in range_constexpr(NT):
+            # the tile's x slice (16 tokens x QH) -> LDS, NXQ x 16 B a thread
+            xs = xnext
+            for q in range_constexpr(NXQ):
+                e = tid + THREADS * q
+                fx.ptr_store(xs[q], xl + ((e // QV) * XAROW + 8 * (e % QV)) // 2)
+            gpu.barrier()
+            if const_expr(rt + 1 < NT):
+                xnext = ax_loads(rt + 1)
+            # router: thread's 4 K of every token of the tile; all LDS reads first, then the
+            # 16 wave reductions interleaved (ILP)
+            xws = [fx.ptr_load(xl + (t16 * XAROW + 2 * tid) // 2) for t16 in range(16)]
+            bvs = []
+            for ci in range_constexpr(CA):
+                c = wave + WAVES * ci
+                for qq in range_constexpr(4):
+                    kk = c * KCH + 8 * g + 32 * qq
+                    bvs.append(
+                        fx.Vector(
+                            fx.ptr_load(
+                                xl + (t * XAROW + kk) // 2,
+                                result_type=fx.Vector.make_type(4, fx.Int32),
+                            )
+                        ).bitcast(fx.BFloat16)
+                    )
+            prs = [bf_lo(xw) * wf[0] + bf_hi(xw) * wf[1] for xw in xws]
+            for off in (32, 16, 8, 4, 2, 1):
+                prs = [xred(v, off, _add) for v in prs]
+            if lane < 16:
+                v_ = prs[0]
+                for t16 in range_constexpr(1, 16):
+                    v_ = (lane == t16).select(prs[t16], v_)
+                fx.ptr_store(v_, tok + (wave * 16 + lane))
+            acc = fx.Vector.filled(4, 0.0, fx.Float32)
+            for ci in range_constexpr(CA):
+                for qq in range_constexpr(4):
+                    acc = mfma_bf16(wa[ci][qq], bvs[ci * 4 + qq], acc)
+            fx.ptr_store(acc, red + (wave * 64 + lane) * 4)
+            gpu.barrier()
+            # branch-free stores (dropped out-of-range offsets): no waitcnt drain at a join
+            t16 = tid % 16
+            v = fx.Float32(0.0)
+            for w in range_constexpr(WAVES):
+                v = v + fx.ptr_load(tok + (w * 16 + t16))
+            if tid < 16:
+                # logits of the CTA's expert, pushed after the loop
+                fx.ptr_store(v, tok + (WAVES * 16 + 8 + rt * 16 + t16))
+            rp = (tid // 16) % 8
+            tt = tid % 16
+            for h2 in range_constexpr(2):
+                rr = 2 * rp + h2
+                s_ = fx.Float32(0.0)
+                for w in range_constexpr(WAVES):
+                    s_ = s_ + fx.ptr_load(red_at(red, w, rr, tt))
+                mine = (rr // OWN) == kq
+                if (tid < 128) & mine:
+                    fx.ptr_store(s_, own + ((rt * OWN + rr - OWN * kq) * 16 + tt))
+                mb_put_if(
+                    (tid < 128) & (~mine),
+                    tag,
+                    scratch + fx.Int64(lay["px"][0]),
+                    ((rb * KQ + kq) * 16 + rr) * S + 16 * rt + tt,
+                    s_,
+                )
+        gpu.barrier()
+        if tid < S:
+            lv = fx.ptr_load(tok + (WAVES * 16 + 8 + tid))
+            for pb in pl:
+                mb_put_sys(
+                    tag,
+                    pb + fx.Int64(play["logit"][0]),
+                    ((par * KQ + kq) * E + rank * EL + rb) * S + tid,
+                    lv,
+                )
+        for it_ in range_constexpr((NRP * 16 * NT + THREADS - 1) // THREADS):
+            e = tid + THREADS * it_
+            if e < NRP * 16 * NT:
+                rt_ = e // (NRP * 16)
+                rp = e % NRP
+                tt = (e // NRP) % 16
+                specs, oks = [], []
+                for src in range_constexpr(KQ):
+                    for h2 in range_constexpr(2):
+                        specs.append(
+                            (
+                                scratch + fx.Int64(lay["px"][0]),
+                                ((rb * KQ + src) * 16 + OWN * kq + 2 * rp + h2) * S
+                                + 16 * rt_
+                                + tt,
+                                1,
+                            )
+                        )
+                        oks.append(fx.Int32(src) != kq)
+                got = mb_poll(tag, health, specs, REGION_IDS["px"], oks=oks)
+                tot = []
+                for h2 in range_constexpr(2):
+                    mine = fx.ptr_load(own + ((rt_ * OWN + 2 * rp + h2) * 16 + tt))
+                    s_ = fx.Float32(0.0)
+                    for src in range_constexpr(KQ):
+                        s_ = s_ + (fx.Int32(src) == kq).select(
+                            mine, got[2 * src + h2][0].bitcast(fx.Float32)
+                        )
+                    tot.append(s_)
+                hw = bf16_pair(tot[0], tot[1])
+                hidx = (
+                    (((par * W + rank) * EL + rb) * KQ + kq) * S + 16 * rt_ + tt
+                ) * NRP + rp
+                for pb in pl:
+                    bo.buffer_store(
+                        hw,
+                        rsrc(pb + fx.Int64(play["hraw"][0])),
+                        hidx,
+                        cache_modifier=CM_SYS,
+                    )
+        rocdl.s_waitcnt(vmcnt=0, expcnt=0, lgkmcnt=0)
+        gpu.barrier()
+        if tid == 0:
+            for pb in pl:
+                mb_put_sys(
+                    tag,
+                    pb + fx.Int64(play["hf"][0]),
+                    (par * W + rank) * BLOCKS + bid,
+                    fx.Int32(1),
+                )
+
+    @traced
+    def stage_route_h(
+        own_peer,
+        bias,
+        w_pren,
+        h_out,
+        ids_out,
+        sc_out,
+        tag,
+        health,
+        tid,
+        bid,
+        lane,
+        wave,
+        rt,
+        tok,
+    ):
+        """CTA t < S: token t's top-8 (sigmoid + bias, radix top-8, renorm) -> ids / weights,
+        then h = pre_expert_norm(h_raw) of token t."""
+        tid = fresh(tid)
+        par = tag_par(tag)
+        t = bid
+        if wave == 0:
+            bias_v = [ld_f32(bias, lane + 64 * i) for i in range(4)]
+            specs = [
+                (
+                    own_peer + fx.Int64(play["logit"][0]),
+                    ((par * KQ + kq) * E + lane + 64 * i) * S + t,
+                    1,
+                )
+                for i in range(4)
+                for kq in range(KQ)
+            ]
+            got = mb_poll(tag, health, specs, REGION_IDS["logit"], cm=CM_SYS)
+            lgs = []
+            for i in range_constexpr(4):
+                v = got[i * KQ][0].bitcast(fx.Float32)
+                for kq in range_constexpr(1, KQ):
+                    v = v + got[i * KQ + kq][0].bitcast(fx.Float32)
+                lgs.append(v.bitcast(fx.Int32))
+            route_token([0], bias_v, [lgs], tid, lane, rt)
+            if lane < TOPK:
+                bo.buffer_store(
+                    fx.ptr_load(rt["rid"] + lane), rsrc(ids_out), t * TOPK + lane
+                )
+                bo.buffer_store(
+                    fx.ptr_load(rt["rw"] + lane), rsrc(sc_out), t * TOPK + lane
+                )
+        # h_raw of token t from every rank (4 x 256 producer CTAs)
+        specs = [
+            (
+                own_peer + fx.Int64(play["hf"][0]),
+                (par * W + (NFQ * tid + q) // BLOCKS) * BLOCKS
+                + (NFQ * tid + q) % BLOCKS,
+                1,
+            )
+            for q in range(NFQ)
+        ]
+        mb_poll(tag, health, specs, REGION_IDS["hf"], cm=CM_SYS)
+        gpu.barrier()
+        kk = tid * 8
+        hsrc = kk // (DM // W)
+        hloc = kk % (DM // W)
+        hb0 = (
+            (((par * W + hsrc) * EL + hloc // 16) * KQ + (hloc % 16) // OWN) * S + t
+        ) * NRP
+        # 8 columns = 4 owner CTAs (K parts) of 2 rows each: one pair a part
+        hv = fx.Vector.from_elements(
+            [
+                ld_i32(own_peer + fx.Int64(play["hraw"][0]), hb0 + o * S, XCM)
+                for o in range(4)
+            ],
+            fx.Int32,
+        )
+        wv = ld_vec(w_pren, kk // 2, 4)
+        ss = fx.Float32(0.0)
+        for d in range_constexpr(4):
+            ss = ss + bf_lo(hv[d]) * bf_lo(hv[d]) + bf_hi(hv[d]) * bf_hi(hv[d])
+        ss = wave_sum(ss)
+        if lane == 0:
+            fx.ptr_store(ss, tok + wave)
+        gpu.barrier()
+        tot = fx.Float32(0.0)
+        for w in range_constexpr(WAVES):
+            tot = tot + fx.ptr_load(tok + w)
+        rs = hw_rsq(tot * fx.Float32(1.0 / DM) + fx.Float32(PRE_EPS))
+        outw = []
+        for d in range_constexpr(4):
+            outw.append(
+                bf16_pair(
+                    bf_lo(hv[d]) * rs * bf_lo(wv[d]), bf_hi(hv[d]) * rs * bf_hi(wv[d])
+                )
+            )
+        bo.buffer_store(
+            fx.Vector.from_elements(outw, fx.Int32), rsrc(h_out), (t * DM + kk) // 2
+        )
+
+    # ------------------------------------------------------------ bnd: column-group boundaries
+    # CTA b = token group tg = b / 32 (tokens 4 tg .. 4 tg + 4) x column group cb = b % 32 (columns 256 cb .. + 256);
+    # thread = (token tq = tid / 128, column pair pp = tid % 128). The token norms reduce in the CTA (2 waves a
+    # token), then over the 32 CTAs of the token group in ONE hop: every CTA of the group reads the group's 32 x 4 x 4
+    # partials (fixed order cb = 0 .. 31) instead of a reducer CTA and a second hop.
+    @traced
+    def group_rstd2(
+        scratch,
+        region,
+        tag,
+        health,
+        tid,
+        lane,
+        wave,
+        bid,
+        q1,
+        q2,
+        q3,
+        q4,
+        eps1,
+        eps2,
+        pr,
+    ):
+        """-> (r1, r2) of this thread's token from the 4 norm partials of every CTA of the token group."""
+        tg = bid // 32
+        base = scratch + fx.Int64(lay[region][0])
+        qs = [wave_sum(q) for q in (q1, q2, q3, q4)]
+        if lane == 0:
+            for k in range_constexpr(4):
+                fx.ptr_store(qs[k], pr + (512 + wave * 4 + k))
+        gpu.barrier()
+        # token tq = tid / 4, partial k = tid % 4: the token's 2 waves in order
+        if tid < 16:
+            v = fx.ptr_load(pr + (512 + (2 * (tid // 4)) * 4 + tid % 4)) + fx.ptr_load(
+                pr + (512 + (2 * (tid // 4) + 1) * 4 + tid % 4)
+            )
+            mb_put(tag, base, bid * 16 + tid, v)
+        src = tid // 16  # (column group, token, partial) of the group, one a thread
+        got = mb_poll(
+            tag,
+            health,
+            [(base, (tg * 32 + src) * 16 + tid % 16, 1)],
+            REGION_IDS[region],
+        )
+        fx.ptr_store(got[0][0].bitcast(fx.Float32), pr + tid)
+        gpu.barrier()
+        if tid < 4:  # token tid: the 32 groups in order, then the two norms
+            p = []
+            for k in range_constexpr(4):
+                a_ = fx.ptr_load(pr + (tid * 4 + k))
+                for c_ in range_constexpr(1, 32):
+                    a_ = a_ + fx.ptr_load(pr + (c_ * 16 + tid * 4 + k))
+                p.append(a_)
+            r1 = hw_rsq(p[0] * fx.Float32(1.0 / HID) + fx.Float32(eps1))
+            q = p[1] + fx.Float32(2.0) * r1 * p[2] + r1 * r1 * p[3]
+            r2 = hw_rsq(q * fx.Float32(1.0 / HID) + fx.Float32(eps2))
+            fx.ptr_store(r1, pr + (512 + 32 + tid))
+            fx.ptr_store(r2, pr + (512 + 36 + tid))
+        gpu.barrier()
+        tq = tid // 128
+        return fx.ptr_load(pr + (512 + 32 + tq)), fx.ptr_load(pr + (512 + 36 + tq))
+
+    @traced
+    def stage_ab_c(
+        own_peer,
+        pl,
+        scratch,
+        stream_in,
+        nrows,
+        alpha_a,
+        beta_a,
+        w_pffn,
+        tag,
+        health,
+        tid,
+        bid,
+        lane,
+        wave,
+        rank,
+        pr,
+    ):
+        """The attention boundary with the column-group mapping: the block owners (b / 32 == rank, block
+        b = 32 columns, all tokens) reduce the o_proj partials and all-gather the result; CTA (tg, cb) then
+        does the boundary of tokens 4 tg .. x columns 256 cb .."""
+        tid = fresh(tid)
+        par = tag_par(tag)
+        if (
+            bid // CPR
+        ) == rank:  # the block owner: 8 partials in rank order -> bf16 -> every rank
+            pp_ = tid % 16
+            tt_ = tid // 16
+            specs = [
+                (
+                    own_peer + fx.Int64(play["opb"][0]),
+                    ((par * W + src) * S + tt_) * (CPR * 16) + (bid % CPR) * 16 + pp_,
+                    1,
+                )
+                for src in range(W)
+            ]
+            got = mb_poll(tag, health, specs, REGION_IDS["opb"], cm=CM_SYS)
+            a0 = bf_lo(got[0][0])
+            a1 = bf_hi(got[0][0])
+            for src in range_constexpr(1, W):
+                a0 = a0 + bf_lo(got[src][0])
+                a1 = a1 + bf_hi(got[src][0])
+            for pb in pl:
+                mb_put_sys(
+                    tag,
+                    pb + fx.Int64(play["orb"][0]),
+                    (par * S + tt_) * (HID // 2) + bid * 16 + pp_,
+                    bf16_pair(a0, a1),
+                )
+        tg = bid // 32
+        cb = bid % 32
+        tt = 4 * tg + tid // 128
+        pp = tid % 128
+        c = cb * 256 + 2 * pp
+        st = ld_vec(stream_in, fx.min(tt, nrows - 1) * HID + c, 2)
+        al = ld_vec(alpha_a, c, 2)
+        be = ld_vec(beta_a, c, 2)
+        wf = ld_i32(w_pffn, c // 2)
+        got = mb_poll(
+            tag,
+            health,
+            [
+                (
+                    own_peer + fx.Int64(play["orb"][0]),
+                    (par * S + tt) * (HID // 2) + c // 2,
+                    1,
+                )
+            ],
+            REGION_IDS["orb"],
+            cm=CM_SYS,
+        )
+        o0 = bf_lo(got[0][0])
+        o1 = bf_hi(got[0][0])
+        a0 = al[0].bitcast(fx.Float32) * st[0].bitcast(fx.Float32)
+        a1 = al[1].bitcast(fx.Float32) * st[1].bitcast(fx.Float32)
+        b0 = be[0].bitcast(fx.Float32) * o0
+        b1 = be[1].bitcast(fx.Float32) * o1
+        r1, r2 = group_rstd2(
+            scratch,
+            "gpa",
+            tag,
+            health,
+            tid,
+            lane,
+            wave,
+            bid,
+            o0 * o0 + o1 * o1,
+            a0 * a0 + a1 * a1,
+            a0 * b0 + a1 * b1,
+            b0 * b0 + b1 * b1,
+            POST_EPS,
+            PRE_EPS,
+            pr,
+        )
+        n0 = bf16_round(o0 * r1)
+        n1 = bf16_round(o1 * r1)
+        s0 = a0 + be[0].bitcast(fx.Float32) * n0
+        s1 = a1 + be[1].bitcast(fx.Float32) * n1
+        bo.buffer_store(
+            fx.Vector.from_elements(
+                [s0.bitcast(fx.Int32), s1.bitcast(fx.Int32)], fx.Int32
+            ),
+            rsrc(scratch + fx.Int64(lay["smidb"][0])),
+            tt * HID + c,
+            cache_modifier=CM_DEV,
+        )
+        bo.buffer_store(
+            bf16_pair((s0 * r2) * bf_lo(wf), (s1 * r2) * bf_hi(wf)),
+            rsrc(scratch + fx.Int64(lay["xinb"][0])),
+            (tt * HID + c) // 2,
+            cache_modifier=CM_DEV,
+        )
+        rocdl.s_waitcnt(vmcnt=0, expcnt=0, lgkmcnt=0)
+        gpu.barrier()
+        if tid == 0:
+            mb_put(tag, scratch + fx.Int64(lay["xinbf"][0]), bid, fx.Int32(1))
+
+    @traced
+    def stage_g_c(
+        own_peer,
+        scratch,
+        w_pff,
+        alpha,
+        beta,
+        w_in,
+        stream_out,
+        so_stride,
+        so_row,
+        sfull_out,
+        gather,
+        tag,
+        health,
+        tid,
+        bid,
+        lane,
+        wave,
+        pr,
+    ):
+        """The FFN boundary with the column-group mapping (CTA (tg, cb): tokens 4 tg .. x columns
+        256 cb ..)."""
+        tid = fresh(tid)
+        par = tag_par(tag)
+        tg = bid // 32
+        cb = bid % 32
+        tt = 4 * tg + tid // 128
+        pp = tid % 128
+        c = cb * 256 + 2 * pp
+        st = ld_vec(scratch + fx.Int64(lay["smidb"][0]), tt * HID + c, 2, CM_DEV)
+        wpf = ld_i32(w_pff, c // 2)
+        al = ld_vec(alpha, c, 2)
+        be = ld_vec(beta, c, 2)
+        wi = ld_i32(w_in, c // 2)
+        got = mb_poll(
+            tag,
+            health,
+            [
+                (
+                    own_peer + fx.Int64(play["yb"][0]),
+                    (par * S + tt) * (HID // 2) + c // 2,
+                    1,
+                )
+            ],
+            REGION_IDS["yb"],
+            cm=CM_SYS,
+        )
+        yw = got[0][0]
+        z0 = bf16_round(bf_lo(wpf) * bf_lo(yw))
+        z1 = bf16_round(bf_hi(wpf) * bf_hi(yw))
+        a0 = al[0].bitcast(fx.Float32) * st[0].bitcast(fx.Float32)
+        a1 = al[1].bitcast(fx.Float32) * st[1].bitcast(fx.Float32)
+        b0 = be[0].bitcast(fx.Float32) * z0
+        b1 = be[1].bitcast(fx.Float32) * z1
+        r1, r2 = group_rstd2(
+            scratch,
+            "gpg",
+            tag,
+            health,
+            tid,
+            lane,
+            wave,
+            bid,
+            z0 * z0 + z1 * z1,
+            a0 * a0 + a1 * a1,
+            a0 * b0 + a1 * b1,
+            b0 * b0 + b1 * b1,
+            POST_EPS,
+            PRE_EPS,
+            pr,
+        )
+        m0 = bf16_round(z0 * r1)
+        m1 = bf16_round(z1 * r1)
+        s0 = a0 + be[0].bitcast(fx.Float32) * m0
+        s1 = a1 + be[1].bitcast(fx.Float32) * m1
+        sw = fx.Vector.from_elements(
+            [s0.bitcast(fx.Int32), s1.bitcast(fx.Int32)], fx.Int32
+        )
+        bo.buffer_store(sw, rsrc(stream_out), tt * so_stride + so_row + c)
+        if gather != 0:
+            bo.buffer_store(sw, rsrc(sfull_out), tt * HID + c)
+        xw = bf16_pair((s0 * r2) * bf_lo(wi), (s1 * r2) * bf_hi(wi))
+        bo.buffer_store(
+            xw,
+            rsrc(scratch + fx.Int64(lay["xnb"][0])),
+            (tt * HID + c) // 2,
+            cache_modifier=CM_DEV,
+        )
+        rocdl.s_waitcnt(vmcnt=0, expcnt=0, lgkmcnt=0)
+        gpu.barrier()
+        if tid == 0:
+            mb_put(tag, scratch + fx.Int64(lay["xnbf"][0]), bid, fx.Int32(1))
+
+    # ============================================================ HB
+    def pp_loads(w_pp, bid, lane, wave):
+        """post_expert_proj tile b / KSP (16 of this rank's QH rows), K part b % KSP."""
+        kh = bid % KSP
+        ops = []
+        for ci in range_constexpr(CP):
+            c = wave + WAVES * ci
+            ops.append(
+                bf_tile_load(
+                    w_pp, (bid // KSP) * 16, DM, kh * (DM // KSP) + c * KCH, lane
+                )
+            )
+        return ops
+
+    @traced
+    def stage_ar(own_peer, pl, tag, health, tid, bid, rank):
+        """CTA b of the column-block owner (b / CPR == rank): the W ranks' partials of columns
+        16b .., rank order (fp32) -> bf16 m -> every rank's MOEM (all-gather, plain) + flag."""
+        tid = fresh(tid)
+        par = tag_par(tag)
+        if tid < W:
+            mb_poll(
+                tag,
+                health,
+                [
+                    (
+                        own_peer + fx.Int64(play["mf"][0]),
+                        (par * W + tid) * BLOCKS + bid,
+                        1,
+                    )
+                ],
+                REGION_IDS["mf"],
+                cm=CM_SYS,
+            )
+        gpu.barrier()
+        for it_ in range_constexpr((8 * S + THREADS - 1) // THREADS):
+            e0 = tid + THREADS * it_
+            ok = e0 < 8 * S
+            e = fx.min(e0, 8 * S - 1)
+            t = e // 8
+            p = e % 8
+            a0 = fx.Float32(0.0)
+            a1 = fx.Float32(0.0)
+            for src in range_constexpr(W):
+                w_ = ld_i32(
+                    own_peer + fx.Int64(play["moep"][0]),
+                    (((par * W + src) * BLOCKS + bid) * S + t) * 8 + p,
+                    XCM,
+                )
+                if const_expr(src == 0):
+                    a0 = bf_lo(w_)
+                    a1 = bf_hi(w_)
+                else:
+                    a0 = a0 + bf_lo(w_)
+                    a1 = a1 + bf_hi(w_)
+            mw = bf16_pair(a0, a1)
+            for pb in pl:
+                st_if(
+                    ok,
+                    mw,
+                    pb + fx.Int64(play["moem"][0]),
+                    ((par * BLOCKS + bid) * S + t) * 8 + p,
+                    CM_SYS,
+                )
+        rocdl.s_waitcnt(vmcnt=0, expcnt=0, lgkmcnt=0)
+        gpu.barrier()
+        if tid == 0:
+            for pb in pl:
+                mb_put_sys(
+                    tag, pb + fx.Int64(play["mmf"][0]), par * BLOCKS + bid, fx.Int32(1)
+                )
+
+    @traced
+    def stage_post(ppo, pl, scratch, tag, health, tid, bid, lane, wave, rank, red):
+        """post_expert_proj tile b / KSP over K part b % KSP for all S tokens; the K parts are
+        summed by the kh = 0 CTA -> y (this rank's QH hidden columns) + flag."""
+        tid = fresh(tid)
+        lane = fresh(lane)
+        par = tag_par(tag)
+        g = lane // 16
+        t = lane % 16
+        kh = bid % KSP
+        tile = bid // KSP
+        own_peer = sel_peer(pl, rank)
+        if tid < BLOCKS // KSP:
+            mb_poll(
+                tag,
+                health,
+                [
+                    (
+                        own_peer + fx.Int64(play["mmf"][0]),
+                        par * BLOCKS + kh * (BLOCKS // KSP) + tid,
+                        1,
+                    )
+                ],
+                REGION_IDS["mmf"],
+                cm=CM_SYS,
+            )
+        gpu.barrier()
+        mb_ = own_peer + fx.Int64(play["moem"][0])
+        pa = [
+            [bf_mma_operand(ppo[ci], qq, lane).bitcast(fx.BFloat16) for qq in range(4)]
+            for ci in range(CP)
+        ]
+
+        def pb_loads(rt):
+            # MOEM block layout [par][b][t][16 columns]: columns kk .. kk+8 = block kk / 16, half (kk % 16) / 8
+            return [
+                ld_vec(
+                    mb_,
+                    (
+                        (
+                            par * BLOCKS
+                            + (
+                                kh * (DM // KSP)
+                                + (wave + WAVES * ci) * KCH
+                                + 8 * g
+                                + 32 * qq
+                            )
+                            // 16
+                        )
+                        * S
+                        + 16 * rt
+                        + t
+                    )
+                    * 8
+                    + 4 * (g % 2),
+                    4,
+                    XCM,
+                ).bitcast(fx.BFloat16)
+                for ci in range(CP)
+                for qq in range(4)
+            ]
+
+        nxt = pb_loads(0)
+        for rt in range_constexpr(NT):
+            cur = nxt
+            if const_expr(rt + 1 < NT):
+                nxt = pb_loads(rt + 1)
+            acc = fx.Vector.filled(4, 0.0, fx.Float32)
+            for ci in range_constexpr(CP):
+                for qq in range_constexpr(4):
+                    acc = mfma_bf16(pa[ci][qq], cur[ci * 4 + qq], acc)
+            fx.ptr_store(acc, red + ((rt * WAVES + wave) * 64 + lane) * 4)
+        gpu.barrier()
+        if const_expr(BND):
+            # bnd: y of the tile (16 columns of this rank x every token) to every rank, tagged pairs at the global
+            # column rank * QH + 16 tile + 2 rp; thread = (row tile rt_, token tt, row pair rp)
+            e0 = tid
+            if e0 < NT * 128:
+                rt_ = e0 // 128
+                tt = (e0 // 8) % 16
+                rp = e0 % 8
+                v = []
+                for h2 in range_constexpr(2):
+                    rr = 2 * rp + h2
+                    s_ = fx.Float32(0.0)
+                    for w in range_constexpr(WAVES):
+                        s_ = s_ + fx.ptr_load(red_at(red, rt_ * WAVES + w, rr, tt))
+                    v.append(s_)
+                if kh != 0:
+                    for h2 in range_constexpr(2):
+                        mb_put(
+                            tag,
+                            scratch + fx.Int64(lay["py"][0]),
+                            (((kh - 1) * NTILES + tile) * 16 + 2 * rp + h2) * S
+                            + 16 * rt_
+                            + tt,
+                            v[h2],
+                        )
+                else:
+                    got = mb_poll(
+                        tag,
+                        health,
+                        [
+                            (
+                                scratch + fx.Int64(lay["py"][0]),
+                                (((q - 1) * NTILES + tile) * 16 + 2 * rp + h2) * S
+                                + 16 * rt_
+                                + tt,
+                                1,
+                            )
+                            for q in range(1, KSP)
+                            for h2 in range(2)
+                        ],
+                        REGION_IDS["py"],
+                    )
+                    y0 = v[0]
+                    y1 = v[1]
+                    for q in range_constexpr(KSP - 1):
+                        y0 = y0 + got[2 * q][0].bitcast(fx.Float32)
+                        y1 = y1 + got[2 * q + 1][0].bitcast(fx.Float32)
+                    w2 = bf16_pair(y0, y1)
+                    for pb in pl:
+                        mb_put_sys(
+                            tag,
+                            pb + fx.Int64(play["yb"][0]),
+                            (par * S + 16 * rt_ + tt) * (HID // 2)
+                            + (rank * QH + tile * 16 + 2 * rp) // 2,
+                            w2,
+                        )
+        else:
+            for it_ in range_constexpr((NT + 3) // 4):
+                e = tid
+                # 4 tiles a round; wave w: owner w, tiles 4 it + k (k < 4), the owner's 2 tokens of each
+                dw = uniform(wave)
+                k_ = (e // 16) % 4
+                rt_ = 4 * it_ + k_
+                rp = e % 8
+                tt = dw + W * ((e // 8) % 2)
+                if rt_ < NT:
+                    v = []
+                    for h2 in range_constexpr(2):
+                        rr = 2 * rp + h2
+                        s_ = fx.Float32(0.0)
+                        for w in range_constexpr(WAVES):
+                            s_ = s_ + fx.ptr_load(red_at(red, rt_ * WAVES + w, rr, tt))
+                        v.append(s_)
+                    # K parts 1 .. KSP-1 put their partial (slot kh - 1); part 0 sums the parts in order
+                    if kh != 0:
+                        for h2 in range_constexpr(2):
+                            mb_put(
+                                tag,
+                                scratch + fx.Int64(lay["py"][0]),
+                                (((kh - 1) * NTILES + tile) * 16 + 2 * rp + h2) * S
+                                + 16 * rt_
+                                + tt,
+                                v[h2],
+                            )
+                    else:
+                        got = mb_poll(
+                            tag,
+                            health,
+                            [
+                                (
+                                    scratch + fx.Int64(lay["py"][0]),
+                                    (((q - 1) * NTILES + tile) * 16 + 2 * rp + h2) * S
+                                    + 16 * rt_
+                                    + tt,
+                                    1,
+                                )
+                                for q in range(1, KSP)
+                                for h2 in range(2)
+                            ],
+                            REGION_IDS["py"],
+                        )
+                        y0 = v[0]
+                        y1 = v[1]
+                        for q in range_constexpr(KSP - 1):
+                            y0 = y0 + got[2 * q][0].bitcast(fx.Float32)
+                            y1 = y1 + got[2 * q + 1][0].bitcast(fx.Float32)
+                        t_ = 16 * rt_ + tt
+                        idx = (
+                            ((par * W + rank) * NTILES + tile) * TS + t_ // W
+                        ) * 8 + rp
+                        w2 = bf16_pair(y0, y1)
+                        bo.buffer_store(
+                            w2,
+                            rsrc(sel_peer(pl, dw) + fx.Int64(play["yt"][0])),
+                            idx,
+                            cache_modifier=CM_SYS,
+                        )
+            rocdl.s_waitcnt(vmcnt=0, expcnt=0, lgkmcnt=0)
+            gpu.barrier()
+            if (tid == 0) & (kh == 0):
+                for pb in pl:
+                    mb_put_sys(
+                        tag,
+                        pb + fx.Int64(play["ytf"][0]),
+                        (par * W + rank) * NTILES + tile,
+                        fx.Int32(1),
+                    )
+
+    @traced
+    def stage_g(
+        own_peer,
+        pl,
+        scratch,
+        w_pff,
+        alpha,
+        beta,
+        w_in,
+        stream_out,
+        so_stride,
+        so_row,
+        gather,
+        tag,
+        health,
+        tid,
+        bid,
+        lane,
+        wave,
+        rank,
+        red4,
+    ):
+        """CTA j < S / W: token t = W j + rank, all 8192 columns: post_feedforward_layernorm
+        (input-scaled, eps 1e-8), gated residual on the mid stream, next input_layernorm ->
+        stream row (token-sharded; ``gather``: to every rank too), xn row to every rank + flag."""
+        tid = fresh(tid)
+        par = tag_par(tag)
+        j = bid
+        t = W * j + rank
+        st, al, be, wps, wis = [], [], [], [], []
+        for h in range_constexpr(2):
+            cq = 8 * (tid + THREADS * h)
+            wps.append(ld_vec(w_pff, cq // 2, 4))
+            wis.append(ld_vec(w_in, cq // 2, 4))
+            for h2 in range_constexpr(2):
+                v4 = ld_vec(
+                    scratch + fx.Int64(lay["smid"][0]), j * HID + cq + 4 * h2, 4, CM_DEV
+                )
+                a4 = ld_vec(alpha, cq + 4 * h2, 4)
+                b4 = ld_vec(beta, cq + 4 * h2, 4)
+                for d in range_constexpr(4):
+                    st.append(v4[d].bitcast(fx.Float32))
+                    al.append(a4[d].bitcast(fx.Float32))
+                    be.append(b4[d].bitcast(fx.Float32))
+        mb_poll(
+            tag,
+            health,
+            [(own_peer + fx.Int64(play["ytf"][0]), par * W * NTILES + tid, 1)],
+            REGION_IDS["ytf"],
+            cm=CM_SYS,
+        )
+        gpu.barrier()
+        z = []
+        for h in range_constexpr(2):
+            q = tid + THREADS * h
+            ysrc = (8 * q) // QH
+            yloc = (8 * q) % QH
+            yv = ld_vec(
+                own_peer + fx.Int64(play["yt"][0]),
+                (((par * W + ysrc) * NTILES + yloc // 16) * TS + j) * 8 + 4 * (q % 2),
+                4,
+                XCM,
+            )
+            wv = wps[h]
+            for d in range_constexpr(4):
+                z.append(bf16_round(bf_lo(wv[d]) * bf_lo(yv[d])))
+                z.append(bf16_round(bf_hi(wv[d]) * bf_hi(yv[d])))
+        ss = fx.Float32(0.0)
+        for k in range_constexpr(16):
+            ss = ss + z[k] * z[k]
+        r1 = hw_rsq(
+            block_sum(ss, lane, wave, red4, 0) * fx.Float32(1.0 / HID)
+            + fx.Float32(POST_EPS)
+        )
+        sv = [al[k] * st[k] + be[k] * bf16_round(z[k] * r1) for k in range(16)]
+        for h in range_constexpr(2):
+            cq = 8 * (tid + THREADS * h)
+            for h2 in range_constexpr(2):
+                swv = fx.Vector.from_elements(
+                    [sv[8 * h + 4 * h2 + d].bitcast(fx.Int32) for d in range(4)],
+                    fx.Int32,
+                )
+                bo.buffer_store(
+                    swv, rsrc(stream_out), j * so_stride + so_row + cq + 4 * h2
+                )
+                if gather != 0:
+                    for pb in pl:
+                        bo.buffer_store(
+                            swv,
+                            rsrc(pb + fx.Int64(play["sg"][0])),
+                            (par * S + t) * HID + cq + 4 * h2,
+                            cache_modifier=CM_SYS,
+                        )
+        ss2 = fx.Float32(0.0)
+        for k in range_constexpr(16):
+            ss2 = ss2 + sv[k] * sv[k]
+        r2 = hw_rsq(
+            block_sum(ss2, lane, wave, red4, 1) * fx.Float32(1.0 / HID)
+            + fx.Float32(PRE_EPS)
+        )
+        for h in range_constexpr(2):
+            cq = 8 * (tid + THREADS * h)
+            wi = wis[h]
+            xw = [
+                bf16_pair(
+                    (sv[8 * h + 2 * d] * r2) * bf_lo(wi[d]),
+                    (sv[8 * h + 2 * d + 1] * r2) * bf_hi(wi[d]),
+                )
+                for d in range(4)
+            ]
+            xv = fx.Vector.from_elements(xw, fx.Int32)
+            for pb in pl:
+                bo.buffer_store(
+                    xv,
+                    rsrc(pb + fx.Int64(play["xn"][0])),
+                    (par * S + t) * (HID // 2) + cq // 2,
+                    cache_modifier=CM_SYS,
+                )
+        rocdl.s_waitcnt(vmcnt=0, expcnt=0, lgkmcnt=0)
+        gpu.barrier()
+        if tid == 0:
+            for pb in pl:
+                mb_put_sys(tag, pb + fx.Int64(play["xnf"][0]), par * S + t, fx.Int32(1))
+                if gather != 0:
+                    mb_put_sys(
+                        tag, pb + fx.Int64(play["sgf"][0]), par * S + t, fx.Int32(1)
+                    )
+
+    def q_loads(w_qkvg, bid, lane, wave, n0, n1):
+        k8 = bid % 8
+        gq = bid // 8
+        return [
+            bf_tile_load(
+                w_qkvg,
+                fx.min(NQT * gq + n, NQTILES - 1) * 16,
+                HID,
+                k8 * (HID // 8) + wave * KCH,
+                lane,
+            )
+            for n in range(n0, n1)
+        ]
+
+    @traced
+    def stage_qkv(
+        q_ops, own_peer, scratch, g_out, tag, health, tid, bid, lane, wave, red, own
+    ):
+        tid = fresh(tid)
+        lane = fresh(lane)
+        par = tag_par(tag)
+        g = lane // 16
+        t = lane % 16
+        k8 = bid % 8
+        gq = bid // 8
+        # xn is local; the 32 CTAs of this K eighth wrote its columns
+        if const_expr(BND):
+            if tid < 32:
+                fq = (tid // 4) * 32 + 4 * k8 + tid % 4
+                mb_poll(
+                    tag,
+                    health,
+                    [(scratch + fx.Int64(lay["xnbf"][0]), fq, 1)],
+                    REGION_IDS["xnbf"],
+                )
+        else:
+            if tid < S:
+                mb_poll(
+                    tag,
+                    health,
+                    [(own_peer + fx.Int64(play["xnf"][0]), par * S + tid, 1)],
+                    REGION_IDS["xnf"],
+                    cm=CM_SYS,
+                )
+        gpu.barrier()
+        if const_expr(BND):
+            xb = scratch + fx.Int64(lay["xnb"][0])
+            xpar = 0
+        else:
+            xb = own_peer + fx.Int64(play["xn"][0])
+            xpar = par
+        qa = [
+            [bf_mma_operand(q_ops[n], qq, lane).bitcast(fx.BFloat16) for qq in range(4)]
+            for n in range(NQT)
+        ]
+
+        def qb_loads(rt):
+            return [
+                ld_vec(
+                    xb,
+                    (xpar * S + 16 * rt + t) * (HID // 2)
+                    + (k8 * (HID // 8) + wave * KCH + 8 * g + 32 * qq) // 2,
+                    4,
+                    XCM,
+                ).bitcast(fx.BFloat16)
+                for qq in range(4)
+            ]
+
+        if const_expr(QR1):
+            # the MFMAs of both row tiles first (red holds both), then one reduction round
+            bnext = qb_loads(0)
+            for rt in range_constexpr(NT):
+                bs = bnext
+                if const_expr(rt + 1 < NT):
+                    bnext = qb_loads(rt + 1)
+                for n in range_constexpr(NQT):
+                    acc = fx.Vector.filled(4, 0.0, fx.Float32)
+                    for qq in range_constexpr(4):
+                        acc = mfma_bf16(qa[n][qq], bs[qq], acc)
+                    fx.ptr_store(
+                        acc, red + (((rt * WAVES + wave) * NQT + n) * 64 + lane) * 4
+                    )
+            gpu.barrier()
+            for it_ in range_constexpr((NT * NQT * 128 + THREADS - 1) // THREADS):
+                e0 = tid + THREADS * it_
+                okq = e0 < NT * NQT * 128
+                e = fx.min(e0, NT * NQT * 128 - 1)
+                rt = e // (NQT * 128)
+                n = (e // 128) % NQT
+                rp = (e // 16) % 8
+                tt = e % 16
+                v = []
+                for h2 in range_constexpr(2):
+                    rr = 2 * rp + h2
+                    s_ = fx.Float32(0.0)
+                    for w in range_constexpr(WAVES):
+                        s_ = s_ + fx.ptr_load(
+                            red_at(red, (rt * WAVES + w) * NQT + n, rr, tt)
+                        )
+                    v.append(s_)
+                if okq & (rp == k8):
+                    fx.ptr_store(v[0], own + (((rt * NQT + n) * 2) * 16 + tt))
+                    fx.ptr_store(v[1], own + (((rt * NQT + n) * 2 + 1) * 16 + tt))
+                for h2 in range_constexpr(2):
+                    mb_put_if(
+                        okq & (rp != k8),
+                        tag,
+                        scratch + fx.Int64(lay["qx"][0]),
+                        (((gq * 8 + k8) * NQT + n) * 16 + 2 * rp + h2) * S
+                        + 16 * rt
+                        + tt,
+                        v[h2],
+                    )
+        else:
+            bnext = qb_loads(0)
+            for rt in range_constexpr(NT):
+                bs = bnext
+                if const_expr(rt + 1 < NT):
+                    bnext = qb_loads(rt + 1)
+                for n in range_constexpr(NQT):
+                    acc = fx.Vector.filled(4, 0.0, fx.Float32)
+                    for qq in range_constexpr(4):
+                        acc = mfma_bf16(qa[n][qq], bs[qq], acc)
+                    fx.ptr_store(acc, red + ((wave * NQT + n) * 64 + lane) * 4)
+                gpu.barrier()
+                for it_ in range_constexpr((NQT * 128 + THREADS - 1) // THREADS):
+                    e0 = tid + THREADS * it_
+                    okq = e0 < NQT * 128
+                    e = fx.min(e0, NQT * 128 - 1)
+                    n = e // 128
+                    rp = (e // 16) % 8
+                    tt = e % 16
+                    v = []
+                    for h2 in range_constexpr(2):
+                        rr = 2 * rp + h2
+                        s_ = fx.Float32(0.0)
+                        for w in range_constexpr(WAVES):
+                            s_ = s_ + fx.ptr_load(red_at(red, w * NQT + n, rr, tt))
+                        v.append(s_)
+                    if okq & (rp == k8):
+                        fx.ptr_store(v[0], own + (((rt * NQT + n) * 2) * 16 + tt))
+                        fx.ptr_store(v[1], own + (((rt * NQT + n) * 2 + 1) * 16 + tt))
+                    for h2 in range_constexpr(2):
+                        mb_put_if(
+                            okq & (rp != k8),
+                            tag,
+                            scratch + fx.Int64(lay["qx"][0]),
+                            (((gq * 8 + k8) * NQT + n) * 16 + 2 * rp + h2) * S
+                            + 16 * rt
+                            + tt,
+                            v[h2],
+                        )
+                gpu.barrier()  # red is rewritten by the next row tile
+        gpu.barrier()
+        for it_ in range_constexpr((NQT * 16 * NT + THREADS - 1) // THREADS):
+            e = tid + THREADS * it_
+            if e < NQT * 16 * NT:
+                rt_ = e // (NQT * 16)
+                n = (e // 16) % NQT
+                tt = e % 16
+                specs, oks = [], []
+                for src in range_constexpr(8):
+                    for h2 in range_constexpr(2):
+                        specs.append(
+                            (
+                                scratch + fx.Int64(lay["qx"][0]),
+                                (((gq * 8 + src) * NQT + n) * 16 + 2 * k8 + h2) * S
+                                + 16 * rt_
+                                + tt,
+                                1,
+                            )
+                        )
+                        oks.append(fx.Int32(src) != k8)
+                got = mb_poll(tag, health, specs, REGION_IDS["qx"], oks=oks)
+                tot = []
+                for h2 in range_constexpr(2):
+                    mine = fx.ptr_load(own + (((rt_ * NQT + n) * 2 + h2) * 16 + tt))
+                    s_ = fx.Float32(0.0)
+                    for src in range_constexpr(8):
+                        s_ = s_ + (fx.Int32(src) == k8).select(
+                            mine, got[2 * src + h2][0].bitcast(fx.Float32)
+                        )
+                    tot.append(s_)
+                row = (NQT * gq + n) * 16 + 2 * k8
+                w2 = bf16_pair(tot[0], tot[1])
+                if row < QKV_ROWS:
+                    mb_put(
+                        tag,
+                        scratch + fx.Int64(lay["qkv"][0]),
+                        (16 * rt_ + tt) * (QKV_ROWS // 2) + row // 2,
+                        w2,
+                    )
+                else:
+                    st_if(
+                        row < QKVG_ROWS,
+                        w2,
+                        g_out,
+                        ((16 * rt_ + tt) * QROWS + row - QKV_ROWS) // 2,
+                    )
+
+    @traced
+    def finish_vllm(
+        scratch,
+        q_out,
+        k_cache,
+        v_cache,
+        slots,
+        cos_sin,
+        pos,
+        tag,
+        health,
+        tid,
+        hh,
+        p,
+        tt,
+        nrows,
+    ):
+        """The finish of head hh (uniform), token tt (lanes 0-31 / 32-63 of a wave: two tokens): q / k as
+        vLLM's fused_muse_spark_qknorm_rope_kv_insert (weightless qk-norm, GPT-J RoPE with fp32 cos / sin,
+        all fp32, one RNE bf16 rounding; no qk_scale); K / V as one bf16 dword a lane into the LBHNC cache,
+        element (slot >> 4) * bstride + (kv * 16 + (slot & 15)) * 128 + 64 isv + 2p, bstride = v_cache (a
+        value); each half wave stores with its own wave-uniform 64-bit block base (readlane of its slot);
+        slot < 0 and tt >= nrows skipped. Rows tt >= nrows are not real rows (their q is written, their KV
+        is not)."""
+        row0 = (hh * HD) // 2 + p
+        got = mb_poll(
+            tag,
+            health,
+            [(scratch + fx.Int64(lay["qkv"][0]), tt * (QKV_ROWS // 2) + row0, 1)],
+            REGION_IDS["qkv"],
+        )
+        x0 = bf_lo(got[0][0])
+        x1 = bf_hi(got[0][0])
+        ss = butterfly(x0 * x0 + x1 * x1, (1, 2, 4, 8, 16))
+        rs = hw_rsq(ss * fx.Float32(1.0 / HD) + fx.Float32(PRE_EPS))
+        isv = hh >= HQ + HKV
+        isq = hh < HQ
+        n0 = x0 * rs
+        n1 = x1 * rs
+        if const_expr(key.next_rope):
+            ps = ld_i32(pos, fx.min(tt, nrows - 1) * 2)
+            c = ld_f32(cos_sin, ps * HD + p)
+            sn = ld_f32(cos_sin, ps * HD + HD // 2 + p)
+            r0 = n0 * c - n1 * sn
+            r1 = n0 * sn + n1 * c
+            n0 = r0
+            n1 = r1
+        o0 = isv.select(x0, n0)
+        o1 = isv.select(x1, n1)
+        if isq:
+            bo.buffer_store(bf16_pair(o0, o1), rsrc(q_out), (tt * QROWS) // 2 + row0)
+        else:
+            kv = hh - HQ - isv.select(fx.Int32(HKV), fx.Int32(0))
+            slot = ld_i32(slots, fx.min(tt, nrows - 1) * 2)
+            for half in range_constexpr(2):
+                sl = readlane(slot, 32 * half)
+                tl_ = readlane(tt, 32 * half)
+                mine = ((tid % 64) >= 32) if half == 1 else ((tid % 64) < 32)
+                if (sl >= 0) & (tl_ < nrows):
+                    base = k_cache + (
+                        fx.Int64(sl >> 4) * v_cache
+                        + fx.Int64(
+                            (kv * 16 + (sl & 15)) * 128
+                            + isv.select(fx.Int32(HD), fx.Int32(0))
+                        )
+                    ) * fx.Int64(2)
+                    if mine:
+                        bo.buffer_store(bf16_pair(o0, o1), rsrc(base), p)
+
+    @traced
+    def head_finish(
+        scratch,
+        q_out,
+        k_cache,
+        v_cache,
+        slots,
+        cos_sin,
+        pos,
+        tag,
+        health,
+        tid,
+        bid,
+        nrows,
+    ):
+        """CTA hh < HQ + 2 HKV finishes head hh (q heads, then k, then v) of the next layer for all
+        S tokens."""
+        tid = fresh(tid)
+        hh = bid
+        p = tid % 32
+        for tg in range_constexpr(NT):
+            if bid < HQ + 2 * HKV:
+                finish_vllm(
+                    scratch,
+                    q_out,
+                    k_cache,
+                    v_cache,
+                    slots,
+                    cos_sin,
+                    pos,
+                    tag,
+                    health,
+                    tid,
+                    hh,
+                    p,
+                    16 * tg + tid // 32,
+                    nrows,
+                )
+
+    @traced
+    def stage_gather(own_peer, sfull_out, tag, health, tid, bid):
+        """CTA t < S copies the all-gathered stream row t into sfull_out."""
+        tid = fresh(tid)
+        par = tag_par(tag)
+        t = bid
+        if tid == 0:
+            mb_poll(
+                tag,
+                health,
+                [(own_peer + fx.Int64(play["sgf"][0]), par * S + t, 1)],
+                REGION_IDS["sgf"],
+                cm=CM_SYS,
+            )
+        gpu.barrier()
+        for h in range_constexpr(4):
+            v = ld_vec(
+                own_peer + fx.Int64(play["sg"][0]),
+                (par * S + t) * HID + tid * 16 + 4 * h,
+                4,
+                XCM,
+            )
+            bo.buffer_store(v, rsrc(sfull_out), t * HID + tid * 16 + 4 * h)
+
+    # The kernels and their launchers. tl (a timeline buffer) is part of the launch signature the
+    # decoder passes; these builds do not read it.
+
+    @flyc.kernel(name=name_a, known_block_size=[THREADS, 1, 1])
+    def ha(
+        attn_o: Int64,
+        attn_g: Int64,
+        stream_in: Int64,
+        w_o: Int64,
+        alpha_a: Int64,
+        beta_a: Int64,
+        w_pffn: Int64,
+        w_r: Int64,
+        bias: Int64,
+        w_pre: Int64,
+        w_pren: Int64,
+        h_out: Int64,
+        ids_out: Int64,
+        sc_out: Int64,
+        scratch: Int64,
+        epoch: Int64,
+        health: Int64,
+        tl: Int64,
+        peer0: Int64,
+        peer1: Int64,
+        peer2: Int64,
+        peer3: Int64,
+        peer4: Int64,
+        peer5: Int64,
+        peer6: Int64,
+        peer7: Int64,
+        layer: Int32,
+        rank: Int32,
+        g_stride: Int32,
+        s_full: Int32,
+        nrows: Int32,
+    ):
+        tid = fx.thread_idx.x
+        bid = fx.block_idx.x
+        lane = tid % 64
+        wave = tid // 64
+        alloc = fx.SharedAllocator()
+        misc = alloc.allocate(Misc).peek()
+        rl = alloc.allocate(RouteLds).peek()
+        sl = alloc.allocate(HaLds)
+        so, sa = sl.o_.peek(), sl.a_.peek()
+        rt = {"rid": rl.rid.ptr, "rw": rl.rw.ptr, "sel": rl.sel.ptr}
+        pl = [peer0, peer1, peer2, peer3, peer4, peer5, peer6, peer7]
+        own_peer = sel_peer(pl, rank)
+        ep = uniform(ld_i32(epoch, 0, CM_DEV))
+        tag = (ep & 0x1FFFFFF) * 64 + layer + 1
+        ao = None
+        if const_expr(FA is None):
+            o_ops = o_loads(w_o, bid, lane, wave)
+        else:
+            # attn_o is the fragment's descriptor; ng polls the merged rows (tagged pairs)
+            attn_o, ao_t = FA.run(
+                attn_o, tag, health, tid, bid, lane, wave, sl.fa_.peek()
+            )
+            ao = (ao_t, health)
+            o_ops = o_loads(w_o, bid, lane, wave)
+        stage_ng(attn_o, attn_g, g_stride, scratch, tag, tid, bid, nrows, ao)
+        stage_o(o_ops, scratch, pl, tag, health, tid, bid, lane, wave, rank, so.red.ptr)
+        wr, wp = a_loads(w_r, w_pre, bid, tid, lane, wave)
+        if const_expr(BND):  # all 256 CTAs, column groups
+            stage_ab_c(
+                own_peer,
+                pl,
+                scratch,
+                stream_in,
+                nrows,
+                alpha_a,
+                beta_a,
+                w_pffn,
+                tag,
+                health,
+                tid,
+                bid,
+                lane,
+                wave,
+                rank,
+                misc.pr.ptr,
+            )
+        else:
+            if bid < TS:
+                stage_ab(
+                    own_peer,
+                    pl,
+                    scratch,
+                    stream_in,
+                    s_full,
+                    nrows,
+                    alpha_a,
+                    beta_a,
+                    w_pffn,
+                    tag,
+                    health,
+                    tid,
+                    bid,
+                    lane,
+                    wave,
+                    rank,
+                    misc.red4.ptr,
+                )
+        stage_a(
+            wr,
+            wp,
+            own_peer,
+            pl,
+            scratch,
+            tag,
+            health,
+            tid,
+            bid,
+            lane,
+            wave,
+            rank,
+            sa.red.ptr,
+            sa.tok.ptr,
+            sa.own.ptr,
+            sa.x.ptr,
+        )
+        if bid < S:
+            stage_route_h(
+                own_peer,
+                bias,
+                w_pren,
+                h_out,
+                ids_out,
+                sc_out,
+                tag,
+                health,
+                tid,
+                bid,
+                lane,
+                wave,
+                rt,
+                sa.tok.ptr,
+            )
+
+    @flyc.kernel(name=name_b, known_block_size=[THREADS, 1, 1])
+    def hb(
+        part: Int64,
+        w_pp: Int64,
+        w_pff: Int64,
+        alpha: Int64,
+        beta: Int64,
+        w_in: Int64,
+        w_qkvg: Int64,
+        stream_out: Int64,
+        sfull_out: Int64,
+        q_out: Int64,
+        g_out: Int64,
+        k_cache: Int64,
+        v_cache: Int64,
+        slots: Int64,
+        cos_sin: Int64,
+        pos: Int64,
+        scratch: Int64,
+        epoch: Int64,
+        health: Int64,
+        tl: Int64,
+        peer0: Int64,
+        peer1: Int64,
+        peer2: Int64,
+        peer3: Int64,
+        peer4: Int64,
+        peer5: Int64,
+        peer6: Int64,
+        peer7: Int64,
+        layer: Int32,
+        rank: Int32,
+        so_stride: Int32,
+        so_off: Int32,
+        gather: Int32,
+        nrows: Int32,
+    ):
+        tid = fx.thread_idx.x
+        bid = fx.block_idx.x
+        lane = tid % 64
+        wave = tid // 64
+        alloc = fx.SharedAllocator()
+        misc = alloc.allocate(Misc).peek()
+        sl = alloc.allocate(HbLds)
+        sf, sq = sl.f_.peek(), sl.q_.peek()
+        pl = [peer0, peer1, peer2, peer3, peer4, peer5, peer6, peer7]
+        own_peer = sel_peer(pl, rank)
+        ep = uniform(ld_i32(epoch, 0, CM_DEV))
+        tag = (ep & 0x1FFFFFF) * 64 + layer + 1
+        ppo = pp_loads(w_pp, bid, lane, wave)
+        if bid // CPR == rank:
+            stage_ar(own_peer, pl, tag, health, tid, bid, rank)
+        stage_post(
+            ppo, pl, scratch, tag, health, tid, bid, lane, wave, rank, sf.red.ptr
+        )
+        qo = q_loads(w_qkvg, bid, lane, wave, 0, QPF)
+        if const_expr(BND):  # all 256 CTAs, column groups
+            stage_g_c(
+                own_peer,
+                scratch,
+                w_pff,
+                alpha,
+                beta,
+                w_in,
+                stream_out,
+                so_stride,
+                so_off,
+                sfull_out,
+                gather,
+                tag,
+                health,
+                tid,
+                bid,
+                lane,
+                wave,
+                misc.pr.ptr,
+            )
+        else:
+            if bid < TS:
+                stage_g(
+                    own_peer,
+                    pl,
+                    scratch,
+                    w_pff,
+                    alpha,
+                    beta,
+                    w_in,
+                    stream_out,
+                    so_stride,
+                    so_off,
+                    gather,
+                    tag,
+                    health,
+                    tid,
+                    bid,
+                    lane,
+                    wave,
+                    rank,
+                    misc.red4.ptr,
+                )
+        qo = qo + q_loads(w_qkvg, bid, lane, wave, QPF, NQT)
+        stage_qkv(
+            qo,
+            own_peer,
+            scratch,
+            g_out,
+            tag,
+            health,
+            tid,
+            bid,
+            lane,
+            wave,
+            sq.red.ptr,
+            sq.own.ptr,
+        )
+        head_finish(
+            scratch,
+            q_out,
+            k_cache,
+            v_cache,
+            slots,
+            cos_sin,
+            pos,
+            tag,
+            health,
+            tid,
+            bid,
+            nrows,
+        )
+        if const_expr(not BND):  # the column-block G wrote sfull_out itself
+            if (gather != 0) & (bid < nrows):
+                stage_gather(own_peer, sfull_out, tag, health, tid, bid)
+
+    @flyc.jit
+    def launch_a(
+        attn_o: Int64,
+        attn_g: Int64,
+        stream_in: Int64,
+        w_o: Int64,
+        alpha_a: Int64,
+        beta_a: Int64,
+        w_pffn: Int64,
+        w_r: Int64,
+        bias: Int64,
+        w_pre: Int64,
+        w_pren: Int64,
+        h_out: Int64,
+        ids_out: Int64,
+        sc_out: Int64,
+        scratch: Int64,
+        epoch: Int64,
+        health: Int64,
+        tl: Int64,
+        peer0: Int64,
+        peer1: Int64,
+        peer2: Int64,
+        peer3: Int64,
+        peer4: Int64,
+        peer5: Int64,
+        peer6: Int64,
+        peer7: Int64,
+        layer: Int32,
+        rank: Int32,
+        g_stride: Int32,
+        s_full: Int32,
+        nrows: Int32,
+        stream: fx.Stream = fx.Stream(None),
+    ):
+        _ = _KEYSTR
+        ha(
+            attn_o,
+            attn_g,
+            stream_in,
+            w_o,
+            alpha_a,
+            beta_a,
+            w_pffn,
+            w_r,
+            bias,
+            w_pre,
+            w_pren,
+            h_out,
+            ids_out,
+            sc_out,
+            scratch,
+            epoch,
+            health,
+            tl,
+            peer0,
+            peer1,
+            peer2,
+            peer3,
+            peer4,
+            peer5,
+            peer6,
+            peer7,
+            layer,
+            rank,
+            g_stride,
+            s_full,
+            nrows,
+        ).launch(grid=(BLOCKS,), block=(THREADS,), stream=stream)
+
+    @flyc.jit
+    def launch_b(
+        part: Int64,
+        w_pp: Int64,
+        w_pff: Int64,
+        alpha: Int64,
+        beta: Int64,
+        w_in: Int64,
+        w_qkvg: Int64,
+        stream_out: Int64,
+        sfull_out: Int64,
+        q_out: Int64,
+        g_out: Int64,
+        k_cache: Int64,
+        v_cache: Int64,
+        slots: Int64,
+        cos_sin: Int64,
+        pos: Int64,
+        scratch: Int64,
+        epoch: Int64,
+        health: Int64,
+        tl: Int64,
+        peer0: Int64,
+        peer1: Int64,
+        peer2: Int64,
+        peer3: Int64,
+        peer4: Int64,
+        peer5: Int64,
+        peer6: Int64,
+        peer7: Int64,
+        layer: Int32,
+        rank: Int32,
+        so_stride: Int32,
+        so_off: Int32,
+        gather: Int32,
+        nrows: Int32,
+        stream: fx.Stream = fx.Stream(None),
+    ):
+        _ = _KEYSTR
+        hb(
+            part,
+            w_pp,
+            w_pff,
+            alpha,
+            beta,
+            w_in,
+            w_qkvg,
+            stream_out,
+            sfull_out,
+            q_out,
+            g_out,
+            k_cache,
+            v_cache,
+            slots,
+            cos_sin,
+            pos,
+            scratch,
+            epoch,
+            health,
+            tl,
+            peer0,
+            peer1,
+            peer2,
+            peer3,
+            peer4,
+            peer5,
+            peer6,
+            peer7,
+            layer,
+            rank,
+            so_stride,
+            so_off,
+            gather,
+            nrows,
+        ).launch(grid=(BLOCKS,), block=(THREADS,), stream=stream)
+
+    _BUILDS[key] = (launch_a, launch_b)
+    return _BUILDS[key]
+
+
+def build_vllm(S=32, W=8, next_rope=True, attn=""):
+    """The HA / HB build vLLM launches for S rows (32 / 64 / 96 / 128) of a TP = EP = 8 group;
+    next_rope: the next layer has RoPE (HB); attn = "swa" / "full": this layer's verify attention
+    fused at the HA head."""
+    assert W == 8 and S in (32, 64, 96, 128), (S, W)
+    return Build(S=S, next_rope=bool(next_rope), attn=attn)
diff --git a/vllm/models/muse_spark/amd/mono/kernels/k5m.py b/vllm/models/muse_spark/amd/mono/kernels/k5m.py
new file mode 100644
index 000000000000..63d81483ac7b
--- /dev/null
+++ b/vllm/models/muse_spark/amd/mono/kernels/k5m.py
@@ -0,0 +1,698 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# mypy: ignore-errors
+# ruff: noqa: B008, E501, SIM102
+
+# Vendored from AMD's Muse Spark FlyDSL monokernel kernels (mk/k5m.py):
+# SPDX-License-Identifier: MIT
+# Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+"""HM: this rank's routed experts for S = 32 / 64 / 96 / 128 rows as one persistent launch
+between HA and HB (k5.py); one rank of a TP = EP = 8 group.
+
+Inputs from HA: h [S, 4096] bf16 (pre_expert_norm output), top-8 ids [S, 8] and weights [S, 8].
+Every CTA builds the same tables: the picks of each local expert in token order, split into
+virtual experts of at most 16 picks (one MFMA N tile). Tasks: (virtual expert v, row tile j),
+256 a virtual expert (gate_up: 16 gate + 16 up rows; down: 16 output rows); CTA b owns tasks
+[b V, (b+1) V) (V virtual experts), so it spans at most two virtual experts and keeps their
+picks' h (gate_up) / a (down) rows in LDS. Weights MXFP4 (w13 / w2 + E8M0 scales) in vLLM's
+shuffled layouts, streamed with a prefetch queue, decoded in registers
+(v_cvt_scalef32_pk_bf16_fp4), bf16 MFMA.
+
+  gate_up: a = bf16(silu(bf16 g) bf16 u) -> scratch A (tagged pairs, pick rows)
+  down:    z = bf16(w_post bf16(y)) -> scratch Z (tagged), per-pick sums of squares a CTA
+  rstd:    CTA c: the post-expert-norm rstd of picks c, c + 256, .. (tagged)
+  final:   CTA b, columns 16b ..: part[t] = sum over t's local picks (k order) of
+           bf16(z rstd) w -> pushed to the owner rank of the column block (b / 32) = HB's
+           reduce-scatter input (+ flags), and part_out (a local copy).
+Same tag as HA / HB of the layer (epoch * 64 + layer + 1).
+"""
+
+from dataclasses import dataclass
+
+import flydsl.compiler as flyc
+import flydsl.expr as fx
+from aiter.ops.flydsl.kernels.kernels_common import LOG2E
+from flydsl.expr import gpu, range_constexpr, rocdl
+from flydsl.expr.typing import Int32, Int64
+
+from . import k5
+from .dev import (
+    BLOCKS,
+    CM_DEV,
+    CM_SYS,
+    THREADS,
+    WAVES,
+    ballot,
+    bf16_pair,
+    bf16_round,
+    bf_hi,
+    bf_lo,
+    butterfly,
+    fresh,
+    hw_exp2,
+    hw_rsq,
+    lanes_below,
+    ld_f32,
+    ld_i32,
+    ld_vec_if,
+    mb_poll,
+    mb_put,
+    mb_put_sys,
+    popcount,
+    readlane,
+    st_if,
+    traced,
+    uniform,
+    wave_incl_scan,
+)
+from .k3 import _lay, dn_loads, fp4_mma, mx_chunk, sel_peer, task_ring
+
+DM = 4096
+INTER = 4096
+E = 256
+W = 8  # ranks in the TP = EP group
+EL = E // W  # local experts
+TOPK = 8
+POST_EPS = 1e-8
+HROW = DM + 8
+AROW = INTER + 8
+SSN = 257  # ss partial slots a pick (CTAs covering one virtual expert, at most 256 / V + 1)
+REGION = {"a": 31, "z": 32, "ss": 33, "rs": 34}
+
+
+@dataclass(frozen=True)
+class BuildM:
+    S: int = 32  # rows (32 / 64 / 96 / 128)
+
+
+def build_vllm(S, W=8):
+    """The HM build vLLM launches for S rows of a TP = EP = 8 group."""
+    assert S in (32, 64, 96, 128) and W == 8, (S, W)
+    return BuildM(S=S)
+
+
+def scratch_layout(S):
+    P = TOPK * S
+    return _lay(
+        [
+            ("a", P * (INTER // 2) * 8),
+            ("z", P * (DM // 2) * 8),
+            ("ss", P * SSN * 8),
+            ("rs", P * 8),
+        ]
+    )
+
+
+_BUILDS = {}
+
+
+def build(key: BuildM):
+    if key in _BUILDS:
+        return _BUILDS[key]
+    _KEYSTR = repr(key)
+    S = key.S
+    assert S in (32, 64, 96, 128), S
+    CPR = BLOCKS // W  # 16-column blocks a rank owns in the EP reduce-scatter
+    P = TOPK * S
+    assert EL % WAVES == 0
+    VMAX = EL + P // 16
+    lay, _ = scratch_layout(S)
+    play, _ = k5.peer_layout(S, W)  # HM addresses moep / mf of the token-owner layout
+    PF_GU, PF_DN = 1, 2  # gate_up / down tasks in flight
+    UNR = S != 96  # down task loop as a ring (unroll 3, fixed slots)
+    W13_BYTES = EL * 2 * INTER * (DM // 2)
+    W13S_BYTES = EL * 2 * INTER * (DM // 32)
+    tag_par = lambda tag: tag & 1
+
+    @fx.struct
+    class RouteM:
+        nu: fx.Array[fx.Int32, EL, 16]
+        vstart: fx.Array[fx.Int32, EL, 16]
+        unused: fx.Array[fx.Int32, 3 * 256, 16]  # keeps the tuned LDS offsets
+        ve_u: fx.Array[fx.Int32, VMAX, 16]
+        ve_n: fx.Array[fx.Int32, VMAX, 16]
+        ve_off: fx.Array[fx.Int32, VMAX, 16]
+        nve: fx.Array[fx.Int32, 4, 16]
+        pinfo: fx.Array[fx.Int32, P, 16]
+        pk: fx.Array[fx.Int32, 2 * 16, 16]
+        ssacc: fx.Array[fx.Float32, 2 * 16, 16]
+        rsl: fx.Array[fx.Float32, P, 16]
+        pad: fx.Array[fx.Int64, 8, 16]  # unused: keeps the tuned LDS offsets
+
+    @fx.struct
+    class MSetup:
+        lists: fx.Array[fx.Int32, EL * S, 16]
+
+    @fx.struct
+    class MGu:
+        h: fx.Array[fx.Int32, 16 * HROW // 2, 16]
+        red: fx.Array[fx.Float32, 2 * WAVES * 64 * 4, 16]
+
+    @fx.struct
+    class MDn:
+        a: fx.Array[fx.Int32, 16 * AROW // 2, 16]
+        red: fx.Array[fx.Float32, 2 * WAVES * 64 * 4, 16]
+
+    @fx.union
+    class MLds:
+        s_: MSetup
+        g_: MGu
+        d_: MDn
+
+    name = f"muse_k5m_s{S}"
+
+    # ------------------------------------------------------------ tables
+    @traced
+    def setup(ids, rank, tid, lane, wave, rt, lists, nrows):
+        tid = fresh(tid)
+        NC = P // 64
+        for q in range_constexpr((P + THREADS - 1) // THREADS):
+            p = tid + THREADS * q
+            if p < P:
+                fx.ptr_store(fx.Int32(-1), rt["pinfo"] + p)
+        # every wave: the weight slot of every pick, 64 picks a register
+        ev = []
+        for c in range_constexpr(NC):
+            real = ((c * 64 + lane) // TOPK) < nrows  # padded rows route nowhere
+            e = ld_i32(ids, c * 64 + lane) - rank * EL
+            ev.append(((e >= 0) & (e < EL) & real).select(e, fx.Int32(-1)))
+        # wave w: slots w, w + 8, ..: pick lists in pick (= token) order
+        for uu in range_constexpr(EL // WAVES):
+            u = wave + WAVES * uu
+            cnt = fx.Int32(0)
+            for c in range_constexpr(NC):
+                hit = ev[c] == u
+                b = ballot(hit)
+                pos = cnt + lanes_below(b, lane)
+                if hit:
+                    fx.ptr_store(fx.Int32(c * 64) + lane, lists + (u * S + pos))
+                cnt = cnt + popcount(b)
+            if lane == 0:
+                fx.ptr_store(cnt, rt["nu"] + u)
+        gpu.barrier()
+        # wave 0, lane u (+ 64 hh): groups of 16 -> virtual experts (slot order)
+        if wave == 0:
+            vcarry = fx.Int32(0)
+            for hh in range_constexpr((EL + 63) // 64):
+                u = lane + 64 * hh
+                n = (u < EL).select(
+                    fx.ptr_load(rt["nu"] + fx.min(u, EL - 1)), fx.Int32(0)
+                )
+                ng = (n + 15) // 16
+                incl = wave_incl_scan(ng, lane) + vcarry
+                st = incl - ng
+                if u < EL:
+                    fx.ptr_store(st, rt["vstart"] + u)
+                for gg in range_constexpr(S // 16):
+                    if gg < ng:
+                        v = st + gg
+                        fx.ptr_store(u, rt["ve_u"] + v)
+                        fx.ptr_store(fx.Int32(16 * gg), rt["ve_off"] + v)
+                        fx.ptr_store(fx.min(n - 16 * gg, 16), rt["ve_n"] + v)
+                vcarry = readlane(incl, 63)
+            if lane == 0:
+                fx.ptr_store(vcarry, rt["nve"] + 0)
+        gpu.barrier()
+        # pinfo[pick] = virtual expert * 32 + slot
+        for q in range_constexpr((EL * S + THREADS - 1) // THREADS):
+            x = tid + THREADS * q
+            u = fx.min(x // S, EL - 1)
+            i = x % S
+            if (x < EL * S) & (i < fx.ptr_load(rt["nu"] + u)):
+                p = fx.ptr_load(lists + (u * S + i))
+                v = fx.ptr_load(rt["vstart"] + u) + i // 16
+                fx.ptr_store(v * 32 + i % 16, rt["pinfo"] + p)
+        gpu.barrier()
+
+    @traced
+    def own_picks(rt, lists, v0, v1, tid):
+        """pk[h][n]: the picks of the CTA's virtual experts v0 (h = 0), v1 (h = 1)."""
+        if tid < 32:
+            h = tid // 16
+            n = tid % 16
+            v = (h == 0).select(v0, v1)
+            ok = (v < uniform(fx.ptr_load(rt["nve"] + 0))) & (
+                n < fx.ptr_load(rt["ve_n"] + fx.min(v, VMAX - 1))
+            )
+            pk = fx.Int32(0)
+            if ok:
+                vv = fx.min(v, VMAX - 1)
+                pk = fx.ptr_load(
+                    lists
+                    + (
+                        fx.ptr_load(rt["ve_u"] + vv) * S
+                        + fx.ptr_load(rt["ve_off"] + vv)
+                        + n
+                    )
+                )
+            fx.ptr_store(ok.select(pk, fx.Int32(-1)), rt["pk"] + tid)
+            fx.ptr_store(fx.Float32(0.0), rt["ssacc"] + tid)
+        gpu.barrier()
+
+    # ------------------------------------------------------------ gate_up
+    def gu_loads(w13, w13s, rt, task, lane, wave, ok):
+        """Task (v, j): waves 0-3 the 16 gate rows 16j .., waves 4-7 the up rows; wave w the
+        K chunks 4 (w % 4) + i (256 K each): 2 x dwordx4 + 1 scale dword a chunk."""
+        r = lane % 16
+        g = lane // 16
+        task = ok.select(task, fx.Int32(0))
+        v = task >> 8
+        j = task & 255
+        e = uniform(fx.ptr_load(rt["ve_u"] + fx.min(v, VMAX - 1)))
+        rg = wave // 4
+        ops = []
+        for i in range_constexpr(4):
+            c = 4 * (wave % 4) + i
+            rbase = e * (2 * INTER) + rg * INTER + j * 16
+            ops += mx_chunk(
+                ok, w13, w13s, rbase, DM, c, lane, g, r, W13_BYTES, W13S_BYTES
+            )
+        return ops
+
+    @traced
+    def load_rows_h(h, rt, hsel, hl, tid):
+        """LDS row n := h row of the token of pick pk[hsel][n] (zeros for an empty slot)."""
+        tid = fresh(tid)
+        NQ = 16 * (DM // 8) // THREADS
+        for half in range_constexpr(2):  # 8 x 16 B in flight a thread (VGPR budget)
+            vs = []
+            for q in range_constexpr(NQ // 2):
+                x = tid + THREADS * (half * (NQ // 2) + q)
+                n = x // (DM // 8)
+                kk = (x % (DM // 8)) * 8
+                pk = fx.ptr_load(rt["pk"] + (hsel * 16 + n))
+                vs.append(
+                    ld_vec_if(pk >= 0, h, ((pk >> 3) * DM + kk) // 2, 4, S * DM * 2)
+                )
+            for q in range_constexpr(NQ // 2):
+                x = tid + THREADS * (half * (NQ // 2) + q)
+                n = x // (DM // 8)
+                kk = (x % (DM // 8)) * 8
+                fx.ptr_store(vs[q], hl + (n * HROW + kk) // 2)
+        gpu.barrier()
+
+    @traced
+    def gu_compute(ops, task, it, hsel, scratch, tag, rt, tid, lane, wave, hl, red):
+        tid = fresh(tid)
+        acc = fx.Vector.filled(4, 0.0, fx.Float32)
+        for i in range_constexpr(4):
+            acc = fp4_mma(
+                ops, 3 * i, hl, HROW, lane, acc, 4 * (wave % 4) + i, task & 1, 16
+            )
+        par = (it & 1) * (WAVES * 64 * 4)
+        fx.ptr_store(acc, red + (par + (wave * 64 + lane) * 4))
+        gpu.barrier()
+        if tid < 128:
+            rp = tid % 8
+            n = tid // 8
+            j = task & 255
+            pk = fx.ptr_load(rt["pk"] + (hsel * 16 + n))
+            av = []
+            for h2 in range_constexpr(2):
+                rr = 2 * rp + h2
+                gs = fx.Float32(0.0)
+                us = fx.Float32(0.0)
+                for w in range_constexpr(4):
+                    off = par + (w * 64 + 16 * (rr // 4) + n) * 4 + rr % 4
+                    gs = gs + fx.ptr_load(red + off)
+                    us = us + fx.ptr_load(red + (off + 4 * 64 * 4))
+                gb = bf16_round(gs)
+                ub = bf16_round(us)
+                sg = bf16_round(
+                    gb / (fx.Float32(1.0) + hw_exp2(gb * fx.Float32(-LOG2E)))
+                )
+                av.append(sg * ub)
+            if pk >= 0:
+                mb_put(
+                    tag,
+                    scratch + fx.Int64(lay["a"][0]),
+                    pk * (INTER // 2) + j * 8 + rp,
+                    bf16_pair(av[0], av[1]),
+                )
+
+    # ------------------------------------------------------------ down
+    @traced
+    def load_rows_a(scratch, tag, health, rt, hsel, al, tid):
+        """LDS row n := A row of pick pk[hsel][n] (polled; zeros for an empty slot)."""
+        tid = fresh(tid)
+        NQ = 16 * (INTER // 4) // THREADS
+        for qb in range_constexpr(NQ // 8):  # 8 specs (16 pairs) a round: VGPR budget
+            specs, oks, dsts = [], [], []
+            for q in range_constexpr(8):
+                x = tid + THREADS * (qb * 8 + q)
+                n = x // (INTER // 4)
+                sidx = x % (INTER // 4)
+                pk = fx.ptr_load(rt["pk"] + (hsel * 16 + n))
+                specs.append(
+                    (
+                        scratch + fx.Int64(lay["a"][0]),
+                        fx.max(pk, 0) * (INTER // 2) + 2 * sidx,
+                        2,
+                    )
+                )
+                oks.append(pk >= 0)
+                dsts.append(al + (n * AROW + 4 * sidx) // 2)
+            got = mb_poll(tag, health, specs, REGION["a"], batch=8, oks=oks)
+            for q in range_constexpr(8):
+                fx.ptr_store(
+                    fx.Vector.from_elements([got[q][0], got[q][1]], fx.Int32), dsts[q]
+                )
+        gpu.barrier()
+
+    @traced
+    def dn_compute(ops, task, it, hsel, scratch, tag, rt, tid, lane, wave, al, red):
+        tid = fresh(tid)
+        acc = fx.Vector.filled(4, 0.0, fx.Float32)
+        for i in range_constexpr(2):
+            acc = fp4_mma(ops, 3 * i, al, AROW, lane, acc, 2 * wave + i, task & 1, 16)
+        par = (it & 1) * (WAVES * 64 * 4)
+        fx.ptr_store(acc, red + (par + (wave * 64 + lane) * 4))
+        gpu.barrier()
+        if tid < 128:
+            rp = tid % 8
+            n = tid // 8
+            j = task & 255
+            pk = fx.ptr_load(rt["pk"] + (hsel * 16 + n))
+            wpw = ops[6]
+            zv = []
+            for h2 in range_constexpr(2):
+                rr = 2 * rp + h2
+                ys = fx.Float32(0.0)
+                for w in range_constexpr(WAVES):
+                    ys = ys + fx.ptr_load(
+                        red + (par + (w * 64 + 16 * (rr // 4) + n) * 4 + rr % 4)
+                    )
+                wp = bf_lo(wpw) if h2 == 0 else bf_hi(wpw)
+                zv.append(bf16_round(wp * bf16_round(ys)))
+            ss = butterfly(zv[0] * zv[0] + zv[1] * zv[1], (1, 2, 4))
+            if pk >= 0:
+                mb_put(
+                    tag,
+                    scratch + fx.Int64(lay["z"][0]),
+                    pk * (DM // 2) + j * 8 + rp,
+                    bf16_pair(zv[0], zv[1]),
+                )
+                if rp == 0:
+                    sa = rt["ssacc"] + (hsel * 16 + n)
+                    fx.ptr_store(fx.ptr_load(sa) + ss, sa)
+
+    # ------------------------------------------------------------ rstd + final
+    @traced
+    def stage_rstd(scratch, tag, health, rt, tid, bid):
+        """CTA c: picks c, c + 256, ..: rstd = rsqrt(sum of the per-CTA partials / 4096 + eps),
+        8 threads a pick (slots part + 8 q)."""
+        tid = fresh(tid)
+        V = uniform(fx.ptr_load(rt["nve"] + 0))
+        part = tid % 8
+        for hh in range_constexpr((P // BLOCKS) * 8 // THREADS + 1):
+            pi = bid + BLOCKS * (tid // 8 + (THREADS // 8) * hh)
+            if pi < P:
+                pinf = fx.ptr_load(rt["pinfo"] + pi)
+                if pinf >= 0:
+                    v = pinf >> 5
+                    cf = (v * 256) // V
+                    nc = (v * 256 + 255) // V - cf + 1
+                    specs = [
+                        (
+                            scratch + fx.Int64(lay["ss"][0]),
+                            pi * SSN + fx.min(part + 8 * q, SSN - 1),
+                            1,
+                        )
+                        for q in range(SSN // 8 + 1)
+                    ]
+                    oks = [(part + 8 * q) < nc for q in range(SSN // 8 + 1)]
+                    got = mb_poll(tag, health, specs, REGION["ss"], batch=17, oks=oks)
+                    s_ = fx.Float32(0.0)
+                    for q in range_constexpr(SSN // 8 + 1):
+                        s_ = s_ + got[q][0].bitcast(fx.Float32)
+                    s_ = butterfly(s_, (1, 2, 4))
+                    if part == 0:
+                        mb_put(
+                            tag,
+                            scratch + fx.Int64(lay["rs"][0]),
+                            pi,
+                            hw_rsq(s_ * fx.Float32(1.0 / DM) + fx.Float32(POST_EPS)),
+                        )
+
+    @traced
+    def stage_final(scratch, sc, part_out, pl, tag, health, rt, tid, bid, rank):
+        tid = fresh(tid)
+        par = tag_par(tag)
+        # rstd of every local pick (polled), then CTA b's 16 columns of every token
+        for q in range_constexpr((P + THREADS - 1) // THREADS):
+            p = tid + THREADS * q
+            if p < P:
+                if fx.ptr_load(rt["pinfo"] + p) >= 0:
+                    got = mb_poll(
+                        tag,
+                        health,
+                        [(scratch + fx.Int64(lay["rs"][0]), p, 1)],
+                        REGION["rs"],
+                    )
+                    fx.ptr_store(got[0][0].bitcast(fx.Float32), rt["rsl"] + p)
+        gpu.barrier()
+        pd = sel_peer(pl, uniform(bid // CPR))
+        for it_ in range_constexpr((8 * S + THREADS - 1) // THREADS):
+            e0 = tid + THREADS * it_
+            ok = e0 < 8 * S
+            e = fx.min(e0, 8 * S - 1)
+            t = e // 8
+            rp = e % 8
+            specs, oks, wts, rss = [], [], [], []
+            for k in range_constexpr(TOPK):
+                p = t * TOPK + k
+                loc = fx.ptr_load(rt["pinfo"] + p) >= 0
+                specs.append(
+                    (scratch + fx.Int64(lay["z"][0]), p * (DM // 2) + bid * 8 + rp, 1)
+                )
+                oks.append(loc & ok)
+                wts.append(loc.select(ld_f32(sc, p), fx.Float32(0.0)))
+                rss.append(fx.ptr_load(rt["rsl"] + p))
+            got = mb_poll(tag, health, specs, REGION["z"], oks=oks)
+            a0 = fx.Float32(0.0)
+            a1 = fx.Float32(0.0)
+            for k in range_constexpr(TOPK):
+                zw = got[k][0]
+                a0 = a0 + oks[k].select(
+                    bf16_round(bf_lo(zw) * rss[k]) * wts[k], fx.Float32(0.0)
+                )
+                a1 = a1 + oks[k].select(
+                    bf16_round(bf_hi(zw) * rss[k]) * wts[k], fx.Float32(0.0)
+                )
+            w2 = bf16_pair(a0, a1)
+            st_if(ok, w2, part_out, (t * DM + bid * 16 + 2 * rp) // 2)
+            st_if(
+                ok,
+                w2,
+                pd + fx.Int64(play["moep"][0]),
+                (((par * W + rank) * BLOCKS + bid) * S + t) * 8 + rp,
+                CM_SYS,
+            )
+        rocdl.s_waitcnt(vmcnt=0, expcnt=0, lgkmcnt=0)
+        gpu.barrier()
+        if tid == 0:
+            mb_put_sys(
+                tag,
+                pd + fx.Int64(play["mf"][0]),
+                (par * W + rank) * BLOCKS + bid,
+                fx.Int32(1),
+            )
+
+    # The kernel and its launcher. tl (a timeline buffer) and w13r .. eptab (expert replicas) are
+    # part of the launch signature the decoder passes; this build does not read them.
+
+    @flyc.kernel(name=name, known_block_size=[THREADS, 1, 1])
+    def hm(
+        h: Int64,
+        ids: Int64,
+        sc: Int64,
+        w13: Int64,
+        w13s: Int64,
+        w2: Int64,
+        w2s: Int64,
+        w_post: Int64,
+        part_out: Int64,
+        scratch: Int64,
+        epoch: Int64,
+        health: Int64,
+        tl: Int64,
+        peer0: Int64,
+        peer1: Int64,
+        peer2: Int64,
+        peer3: Int64,
+        peer4: Int64,
+        peer5: Int64,
+        peer6: Int64,
+        peer7: Int64,
+        layer: Int32,
+        rank: Int32,
+        nrows: Int32,
+        w13r: Int64,
+        w13sr: Int64,
+        w2r: Int64,
+        w2sr: Int64,
+        eptab: Int64,
+    ):
+        tid = fx.thread_idx.x
+        bid = fx.block_idx.x
+        lane = tid % 64
+        wave = tid // 64
+        alloc = fx.SharedAllocator()
+        rl = alloc.allocate(RouteM).peek()
+        ml = alloc.allocate(MLds)
+        ms, mg, md = ml.s_.peek(), ml.g_.peek(), ml.d_.peek()
+        rt = {
+            "nu": rl.nu.ptr,
+            "vstart": rl.vstart.ptr,
+            "ve_u": rl.ve_u.ptr,
+            "ve_n": rl.ve_n.ptr,
+            "ve_off": rl.ve_off.ptr,
+            "nve": rl.nve.ptr,
+            "pinfo": rl.pinfo.ptr,
+            "pk": rl.pk.ptr,
+            "ssacc": rl.ssacc.ptr,
+            "rsl": rl.rsl.ptr,
+        }
+        pl = [peer0, peer1, peer2, peer3, peer4, peer5, peer6, peer7]
+        ep = uniform(ld_i32(epoch, 0, CM_DEV))
+        tag = (ep & 0x1FFFFFF) * 64 + layer + 1
+        setup(ids, rank, tid, lane, wave, rt, ms.lists.ptr, nrows)
+        V = uniform(fx.ptr_load(rt["nve"] + 0))
+        t0 = bid * V
+        t1 = t0 + V
+        v0 = t0 >> 8
+        v1 = (fx.max(t1, 1) - 1) >> 8
+        gu_ld = lambda t: gu_loads(w13, w13s, rt, t, lane, wave, t < t1)
+        q = [gu_ld(t0 + d) for d in range(PF_GU)]
+        own_picks(rt, ms.lists.ptr, v0, v1, tid)
+        # gate_up over virtual expert v0, then (if the range crosses) v1
+        load_rows_h(h, rt, 0, mg.h.ptr, tid)
+        tsplit = fx.min(t1, (v0 + 1) << 8)
+        gu_c = lambda ops, t, i, v: gu_compute(
+            ops, t, i, v, scratch, tag, rt, tid, lane, wave, mg.h.ptr, mg.red.ptr
+        )
+        q, task, it = task_ring(gu_ld, gu_c, q, t0, fx.Int32(0), tsplit, False, 0)
+        if task < t1:
+            gpu.barrier()
+            load_rows_h(h, rt, 1, mg.h.ptr, tid)
+        task_ring(gu_ld, gu_c, q, task, it, t1, False, 1)
+        # down: same tasks
+        dn_ld = lambda t: dn_loads(
+            w2, w2s, w_post, rt["ve_u"], t, lane, wave, t < t1, VMAX
+        )
+        dq = [dn_ld(t0 + d) for d in range(PF_DN)]
+        gpu.barrier()
+        load_rows_a(scratch, tag, health, rt, 0, md.a.ptr, tid)
+        dn_c = lambda ops, t, i, v: dn_compute(
+            ops, t, i, v, scratch, tag, rt, tid, lane, wave, md.a.ptr, md.red.ptr
+        )
+        dq, task, it = task_ring(dn_ld, dn_c, dq, t0, fx.Int32(0), tsplit, UNR, 0)
+        if task < t1:
+            gpu.barrier()
+            load_rows_a(scratch, tag, health, rt, 1, md.a.ptr, tid)
+        task_ring(dn_ld, dn_c, dq, task, it, t1, UNR, 1)
+        gpu.barrier()
+        # per-CTA sums of squares of the picks of v0 / v1 -> slot (b - first CTA of the virtual expert)
+        # (the down's epilogue: ssacc is the down's CTA-local sum)
+        if (tid < 32) & (V > 0):
+            hs = tid // 16
+            pk = fx.ptr_load(rt["pk"] + tid)
+            v = (hs == 0).select(v0, v1)
+            okh = (pk >= 0) & ((hs == 0) | (v1 != v0))
+            cf = (v * 256) // fx.max(V, 1)
+            if okh:
+                mb_put(
+                    tag,
+                    scratch + fx.Int64(lay["ss"][0]),
+                    pk * SSN + (bid - cf),
+                    fx.ptr_load(rt["ssacc"] + tid),
+                )
+        stage_rstd(scratch, tag, health, rt, tid, bid)
+        stage_final(scratch, sc, part_out, pl, tag, health, rt, tid, bid, rank)
+
+    @flyc.jit
+    def launch(
+        h: Int64,
+        ids: Int64,
+        sc: Int64,
+        w13: Int64,
+        w13s: Int64,
+        w2: Int64,
+        w2s: Int64,
+        w_post: Int64,
+        part_out: Int64,
+        scratch: Int64,
+        epoch: Int64,
+        health: Int64,
+        tl: Int64,
+        peer0: Int64,
+        peer1: Int64,
+        peer2: Int64,
+        peer3: Int64,
+        peer4: Int64,
+        peer5: Int64,
+        peer6: Int64,
+        peer7: Int64,
+        layer: Int32,
+        rank: Int32,
+        nrows: Int32,
+        w13r: Int64,
+        w13sr: Int64,
+        w2r: Int64,
+        w2sr: Int64,
+        eptab: Int64,
+        stream: fx.Stream = fx.Stream(None),
+    ):
+        _ = _KEYSTR
+        hm(
+            h,
+            ids,
+            sc,
+            w13,
+            w13s,
+            w2,
+            w2s,
+            w_post,
+            part_out,
+            scratch,
+            epoch,
+            health,
+            tl,
+            peer0,
+            peer1,
+            peer2,
+            peer3,
+            peer4,
+            peer5,
+            peer6,
+            peer7,
+            layer,
+            rank,
+            nrows,
+            w13r,
+            w13sr,
+            w2r,
+            w2sr,
+            eptab,
+        ).launch(grid=(BLOCKS,), block=(THREADS,), stream=stream)
+
+    _BUILDS[key] = launch
+    return launch
diff --git a/vllm/models/muse_spark/amd/mono/peer.py b/vllm/models/muse_spark/amd/mono/peer.py
new file mode 100644
index 000000000000..ab85de285945
--- /dev/null
+++ b/vllm/models/muse_spark/amd/mono/peer.py
@@ -0,0 +1,62 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Device buffers that every rank of a group maps through IPC handles."""
+
+import ctypes
+
+import torch.distributed as dist
+from torch.distributed import ProcessGroup
+
+from vllm.distributed.device_communicators.cuda_wrapper import (
+    CudaRTLibrary,
+    cudaIpcMemHandle_t,
+)
+
+
+class PeerBuffer:
+    """One zeroed ``cudaMalloc`` buffer a rank, mapped by every rank of a group.
+
+    The monokernel polls these mailboxes with system-scope loads, so they come
+    from a plain device allocation: not torch's caching allocator (pool
+    offsets) and not the uncached custom all-reduce buffers. Freed by
+    ``close()`` or when the object is collected (like ``CustomAllreduce``).
+    """
+
+    def __init__(self, nbytes: int, group: ProcessGroup) -> None:
+        self.nbytes = nbytes
+        self.pointers: list[int] = []
+        self._opened: list[ctypes.c_void_p] = []
+        self.local: ctypes.c_void_p | None = None
+        self.lib = CudaRTLibrary()
+        self.rank = dist.get_rank(group=group)
+        self.local = self.lib.cudaMalloc(nbytes)
+        self.lib.cudaMemset(self.local, 0, nbytes)
+        self.lib.cudaDeviceSynchronize()
+        handle = self.lib.cudaIpcGetMemHandle(self.local)
+
+        handles: list[bytes | None] = [None] * dist.get_world_size(group=group)
+        dist.all_gather_object(handles, bytes(handle), group=group)
+
+        for rank, raw in enumerate(handles):
+            if rank == self.rank:
+                pointer = self.local
+            else:
+                assert raw is not None
+                pointer = self.lib.cudaIpcOpenMemHandle(
+                    cudaIpcMemHandle_t.from_buffer_copy(raw)
+                )
+                self._opened.append(pointer)
+            assert pointer.value is not None
+            self.pointers.append(pointer.value)
+
+    def close(self) -> None:
+        """Unmap the peers' buffers and free this rank's; idempotent."""
+        while self._opened:
+            self.lib.cudaIpcCloseMemHandle(self._opened.pop())
+        if self.local is not None:
+            self.lib.cudaFree(self.local)
+            self.local = None
+        self.pointers = []
+
+    def __del__(self) -> None:
+        self.close()
diff --git a/vllm/models/muse_spark/common/dspark.py b/vllm/models/muse_spark/common/dspark.py
index 884c807a2164..82243be6523c 100644
--- a/vllm/models/muse_spark/common/dspark.py
+++ b/vllm/models/muse_spark/common/dspark.py
@@ -372,12 +372,13 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
             },
             orig_to_new_stacked=self.hf_to_vllm_mapper.orig_to_new_stacked,
         )
-        loaded = AutoWeightsLoader(self).load_weights(
+        return AutoWeightsLoader(self).load_weights(
             model_weights.items(), mapper=mapper
         )
-        self.model._build_fused_kv_buffers()
+
+    def process_weights_after_loading(self) -> None:
+        super().process_weights_after_loading()
         materialize_residual_path_constants(self.model)
-        return loaded
 
 
 EntryClass = MuseSparkDSparkForCausalLM
diff --git a/vllm/models/muse_spark/common/moe.py b/vllm/models/muse_spark/common/moe.py
index b9862feee1ca..69e62b01db17 100644
--- a/vllm/models/muse_spark/common/moe.py
+++ b/vllm/models/muse_spark/common/moe.py
@@ -65,9 +65,14 @@ def _compute_routing(
         indices_type: torch.dtype | None,
         *,
         input_ids: torch.Tensor | None = None,
+        is_padding: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         if not router_logits.is_cuda:
-            return self._compute_routing_native(router_logits, indices_type)
+            weights, ids = self._compute_routing_native(router_logits, indices_type)
+            if is_padding is not None:
+                weights = weights.masked_fill(is_padding[:, None], 0)
+                ids = ids.masked_fill(is_padding[:, None], -1)
+            return weights, ids
 
         topk_weights, topk_ids = fused_musespark1x_routing(
             router_logits,
@@ -77,6 +82,7 @@ def _compute_routing(
             router_output_multiplier=self.router_output_multiplier,
             router_score_normalization=self.router_score_normalization,
             router_score_division_eps=self.router_score_division_eps,
+            is_padding=is_padding,
         )
         if indices_type is not None and indices_type != torch.int32:
             topk_ids = topk_ids.to(indices_type)
diff --git a/vllm/models/muse_spark/common/ops/fused_routing.py b/vllm/models/muse_spark/common/ops/fused_routing.py
index d2a37e718c98..fd881df93040 100644
--- a/vllm/models/muse_spark/common/ops/fused_routing.py
+++ b/vllm/models/muse_spark/common/ops/fused_routing.py
@@ -1,9 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Triton kernel for MuseSpark1x's top-k router, one program per token."""
+"""Fused CUDA and Triton kernels for MuseSpark1x's top-k router."""
 
 import torch
 
+from vllm import _custom_ops as ops
 from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
 
@@ -114,17 +115,18 @@ def _routing_kernel(
     if HAS_EPS:
         score_sum = score_sum + division_eps
     inv_sum_scaled = route_scale / score_sum
+    weights = scores * inv_sum_scaled
     if is_padding is not None:
         padding = tl.load(is_padding + token_id * is_padding_stride)
         ids = tl.where(padding, -1, ids)
-        inv_sum_scaled = tl.where(padding, 0.0, inv_sum_scaled)
+        weights = tl.where(padding, 0.0, weights)
 
     tl.store(
         topk_ids_ptr + token_id * stride_im + slots * stride_ik, ids, mask=slot_mask
     )
     tl.store(
         topk_weights_ptr + token_id * stride_wm + slots * stride_wk,
-        scores * inv_sum_scaled,
+        weights,
         mask=slot_mask,
     )
 
@@ -186,6 +188,24 @@ def fused_musespark1x_routing(
     weights_view = topk_weights[:num_tokens, :topk]
     ids_view = topk_ids[:num_tokens, :topk]
 
+    if (
+        current_platform.is_cuda()
+        and num_experts == 256
+        and topk == 8
+        and normalization_mode == 0
+    ):
+        ops.muse_spark_routing(
+            weights_view,
+            ids_view,
+            gating_output,
+            bias,
+            router_output_multiplier if router_output_multiplier is not None else 1.0,
+            route_scale,
+            router_score_division_eps if router_score_division_eps is not None else 0.0,
+            is_padding,
+        )
+        return weights_view, ids_view
+
     enable_pdl = current_platform.is_cuda() and current_platform.is_arch_support_pdl()
     _routing_kernel[(num_tokens,)](
         gating_output,
diff --git a/vllm/models/muse_spark/nvidia/hybrid_moe.py b/vllm/models/muse_spark/nvidia/hybrid_moe.py
new file mode 100644
index 000000000000..87c23cbd8912
--- /dev/null
+++ b/vllm/models/muse_spark/nvidia/hybrid_moe.py
@@ -0,0 +1,159 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Token-threshold TRT-LLM execution over Muse MegaMoE's shared NVFP4 weights."""
+
+import functools
+
+import torch
+
+from vllm.distributed import get_tensor_model_parallel_world_size
+from vllm.model_executor.layers.fused_moe.moe_output import (
+    UnfinalizedMoEOutput,
+    convert_flashinfer_moe_output,
+)
+from vllm.models.muse_spark.nvidia.mega_moe import MuseSparkMegaMoEExperts
+from vllm.models.muse_spark.nvidia.ops.mnnvl_fusions import (
+    maybe_create_moe_finalize_all_reduce,
+)
+from vllm.models.muse_spark.nvidia.ops.prepare_trtllm_nvfp4 import (
+    norm_quant_trtllm_nvfp4,
+)
+
+
+@functools.cache
+def _get_decode_workspace(hidden_size: int, top_k: int, max_tokens: int):
+    workspace = maybe_create_moe_finalize_all_reduce(
+        hidden_size, top_k, max_tokens, enable_pdl=True
+    )
+    if workspace is None:
+        raise RuntimeError("Muse hybrid MoE requires the fused EP finalize workspace")
+    return workspace
+
+
+class MuseSparkHybridMoEExperts(MuseSparkMegaMoEExperts):
+    """TRT-LLM for small replicated batches; MegaMoE for sequence shards."""
+
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        self._decode_workspace = None
+        if self.is_nvfp4:
+            if get_tensor_model_parallel_world_size() not in (4, 8):
+                raise ValueError(
+                    "Muse hybrid MoE currently supports single-node TEP4/8"
+                )
+            self._decode_workspace = _get_decode_workspace(
+                self.hidden_size, self.top_k, self.trtllm_token_threshold
+            )
+            norm = self.post_expert_norm
+            self._decode_workspace.compile(
+                row_stride=self.hidden_size,
+                weight_dtype=norm.weight.dtype,
+                eps=norm.variance_epsilon,
+                gain_center=norm.fused_gain_center,
+                weight_before_norm=norm.applies_weight_before_norm,
+            )
+
+    def finalize_weights(self) -> None:
+        if self._transformed_l1 is not None:
+            return
+        super().finalize_weights()
+        if self.is_nvfp4:
+            # Preserve MegaMoE's per-expert scales. TRT follows its existing
+            # static-activation recipe, with one maximum per local expert bank.
+            a13 = self.w13_input_scale.max().float().reshape(1)
+            a2 = self.w2_input_scale.max().float().reshape(1)
+            self._input_global_scale = a13.reciprocal()
+            self._g1_gate = (a13 * self.w13_weight_scale_2[:, 0]).contiguous()
+            self._g1_up = (a13 * self.w13_weight_scale_2[:, 1] / a2).contiguous()
+            self._g2_scale = (a2 * self.w2_weight_scale_2).contiguous()
+
+    def use_trtllm(self, num_tokens: int, sp_active: bool) -> bool:
+        return (
+            self.is_nvfp4
+            and not sp_active
+            and 0 < num_tokens <= self.trtllm_token_threshold
+        )
+
+    @property
+    def nvfp4_input_global_scale(self) -> torch.Tensor:
+        return self._input_global_scale
+
+    def prepare_decode_inputs(self, pooled, norm, buffer, num_tokens: int):
+        return norm_quant_trtllm_nvfp4(
+            pooled,
+            num_tokens,
+            norm.weight,
+            norm.variance_epsilon,
+            norm.fused_gain_center,
+            self._input_global_scale,
+            out=(
+                buffer.x[:num_tokens].view(torch.uint8),
+                buffer.x_sf[:num_tokens].view(torch.uint8),
+            ),
+        )
+
+    def forward_trtllm(
+        self,
+        quantized_input: tuple[torch.Tensor, torch.Tensor],
+        buffer,
+        topk_ids: torch.Tensor,
+    ) -> torch.Tensor:
+        import flashinfer
+        from flashinfer.comm.mnnvl_cutedsl_moe_finalize import (
+            moe_finalize_rmsnorm_allreduce,
+        )
+
+        packed, scales = quantized_input
+        num_tokens = packed.shape[0]
+        raw = flashinfer.fused_moe.trtllm_fp4_block_scale_routed_moe(
+            topk_ids=(
+                topk_ids,
+                buffer.topk_weights[:num_tokens],
+            ),
+            routing_bias=None,
+            hidden_states=packed,
+            hidden_states_scale=scales.view(torch.float8_e4m3fn),
+            gemm1_weights=self.w13_weight,
+            gemm1_weights_scale=self.w13_weight_scale,
+            gemm1_bias=None,
+            gemm1_alpha=None,
+            gemm1_beta=None,
+            gemm1_clamp_limit=None,
+            gemm2_weights=self.w2_weight,
+            gemm2_weights_scale=self.w2_weight_scale,
+            gemm2_bias=None,
+            output1_scale_scalar=self._g1_up,
+            output1_scale_gate_scalar=self._g1_gate,
+            output2_scale_scalar=self._g2_scale,
+            num_experts=self.num_experts,
+            top_k=self.top_k,
+            n_group=0,
+            topk_group=0,
+            intermediate_size=self.intermediate_size,
+            local_expert_offset=self.experts_start_idx,
+            local_num_experts=self.num_local_experts,
+            routed_scaling_factor=None,
+            routing_method_type=7,
+            do_finalize=False,
+            tune_max_num_tokens=self.trtllm_token_threshold,
+        )
+        result = convert_flashinfer_moe_output(
+            raw, do_finalize=False, num_tokens=num_tokens, top_k=self.top_k
+        )
+        assert isinstance(result, UnfinalizedMoEOutput)
+        norm = self.post_expert_norm
+        y = torch.empty(
+            (num_tokens, self.hidden_size), device=packed.device, dtype=torch.bfloat16
+        )
+        moe_finalize_rmsnorm_allreduce(
+            result.gemm2_permuted,
+            result.expert_weights,
+            result.expanded_idx_to_permuted_idx,
+            norm.weight,
+            y,
+            self._decode_workspace,
+            eps=norm.variance_epsilon,
+            gain_center=norm.fused_gain_center,
+            weight_before_norm=norm.applies_weight_before_norm,
+        )
+        return y
diff --git a/vllm/models/muse_spark/nvidia/mega_moe.py b/vllm/models/muse_spark/nvidia/mega_moe.py
index d5d34c96f7f8..d59ee52e6081 100644
--- a/vllm/models/muse_spark/nvidia/mega_moe.py
+++ b/vllm/models/muse_spark/nvidia/mega_moe.py
@@ -22,6 +22,7 @@
 )
 from vllm.forward_context import get_forward_context, is_forward_context_available
 from vllm.logger import init_logger
+from vllm.model_executor.layers.quantization.modelopt import ModelOptNvFp4Config
 from vllm.model_executor.utils import set_weight_attrs
 from vllm.models.common.ops.sequence_parallel import is_sequence_parallel_active
 from vllm.models.muse_spark.nvidia.ops.prepare_megamoe_bf16 import (
@@ -37,6 +38,9 @@
 
 logger = init_logger(__name__)
 
+_TRTLLM_TOKEN_THRESHOLD_TEP4 = 128
+_TRTLLM_TOKEN_THRESHOLD_TEP8 = 256
+
 
 @functools.cache
 def _get_deep_gemm():
@@ -47,15 +51,33 @@ def _get_deep_gemm():
         or not hasattr(module, "bf16_mega_moe")
         or "post_expert_norm_gamma"
         not in inspect.signature(module.bf16_mega_moe).parameters
+        or "weight_layout" not in inspect.signature(module.fp4_fp4_mega_moe).parameters
     ):
         raise RuntimeError(
             "Muse Spark MegaMoE requires the pinned DeepGEMM build "
-            "with BF16 post-expert normalization. "
+            "with BF16 post-expert normalization and TRT-LLM NVFP4 weight layout. "
             "Install it with tools/install_deepgemm.sh or rebuild the vLLM wheel."
         )
     return module
 
 
+def get_trtllm_token_threshold(vllm_config: VllmConfig) -> int:
+    parallel = vllm_config.parallel_config
+    if (
+        vllm_config.kernel_config.moe_backend == "deep_gemm_mega_moe"
+        and isinstance(vllm_config.quant_config, ModelOptNvFp4Config)
+        and parallel.tensor_parallel_size in (4, 8)
+        and parallel.data_parallel_size == 1
+        and parallel.pipeline_parallel_size == 1
+    ):
+        return (
+            _TRTLLM_TOKEN_THRESHOLD_TEP4
+            if parallel.tensor_parallel_size == 4
+            else _TRTLLM_TOKEN_THRESHOLD_TEP8
+        )
+    return 0
+
+
 def _shard_padding_mask(is_padding: torch.Tensor, num_tokens: int) -> torch.Tensor:
     """This rank's chunk of a full-length padding mask; the round-up tail pads."""
     tp_size = get_tensor_model_parallel_world_size()
@@ -98,7 +120,11 @@ def __init__(
         self.is_sequence_parallel = is_sequence_parallel
         self.post_expert_norm = post_expert_norm
         self.is_nvfp4 = is_nvfp4
-
+        self.trtllm_token_threshold = get_trtllm_token_threshold(vllm_config)
+        self.sp_min_tokens = self.trtllm_token_threshold + 1
+        self.weight_layout = (
+            "trtllm" if is_nvfp4 and self.trtllm_token_threshold else "native"
+        )
         if hidden_size % 256 != 0 or intermediate_size % 256 != 0:
             raise ValueError(
                 "MuseSpark1x MegaMoE requires hidden and intermediate "
@@ -138,6 +164,16 @@ def param(*shape: int, dtype: torch.dtype) -> nn.Parameter:
             self.w2_input_scale = param(num_local_experts, dtype=torch.float32)
             self.w13_input_scale = param(num_local_experts, 2, dtype=torch.float32)
 
+            for scale, value in (
+                (self.w13_weight_scale, 1.0),
+                (self.w2_weight_scale, 1.0),
+                (self.w13_weight_scale_2, 1e-3),
+                (self.w2_weight_scale_2, 1e-3),
+                (self.w13_input_scale, 1e-2),
+                (self.w2_input_scale, 1e-2),
+            ):
+                set_weight_attrs(scale, {"dummy_weight_value": value})
+
         else:
             self.w13_weight = param(
                 num_local_experts,
@@ -199,7 +235,9 @@ def finalize_weights(self) -> None:
             return
         deep_gemm = _get_deep_gemm()
 
-        if self.is_nvfp4:
+        if self.weight_layout == "trtllm":
+            self._prepare_trtllm_weights()
+        elif self.is_nvfp4:
 
             def sf_to_kernel_layout(sf: torch.Tensor) -> torch.Tensor:
                 # Four E4M3 block scales per int32, then MN-major for the TMA loads.
@@ -233,6 +271,40 @@ def sf_to_kernel_layout(sf: torch.Tensor) -> torch.Tensor:
             gamma = gamma + norm.gain_center
         self._norm_gamma = gamma.contiguous()
 
+    def _prepare_trtllm_weights(self) -> None:
+        from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import (
+            prepare_static_weights_for_trtllm_fp4_moe,
+            reorder_w1w3_to_w3w1,
+        )
+
+        w13, sf13 = reorder_w1w3_to_w3w1(
+            self.w13_weight.data, self.w13_weight_scale.data
+        )
+        w13, sf13, w2, sf2 = prepare_static_weights_for_trtllm_fp4_moe(
+            w13,
+            self.w2_weight.data,
+            sf13,
+            self.w2_weight_scale.data,
+            self.hidden_size,
+            self.intermediate_size,
+            self.num_local_experts,
+            True,
+        )
+        self.w13_weight.data = w13
+        self.w13_weight_scale.data = sf13
+        self.w2_weight.data = w2
+        self.w2_weight_scale.data = sf2
+        e, h, i = self.num_local_experts, self.hidden_size, self.intermediate_size
+        # Views describe logical shapes; DeepGEMM reads the physical TRT layout.
+        self._transformed_l1 = (
+            w13.view(torch.int8),
+            sf13.view(torch.int32).view(e, 2 * i, h // 64),
+        )
+        self._transformed_l2 = (
+            w2.view(torch.int8),
+            sf2.view(torch.int32).view(e, h, i // 64),
+        )
+
     def _drop_retransformed_weights(self) -> None:
         """Release the checkpoint tensors the kernel no longer reads."""
         assert isinstance(self._transformed_l2, tuple)
@@ -300,7 +372,9 @@ def get_padding_mask(self, num_tokens: int) -> torch.Tensor | None:
         is_padding = get_forward_context().is_padding
         if is_padding is None:
             return None
-        if is_sequence_parallel_active(self.is_sequence_parallel, is_padding.shape[0]):
+        if is_sequence_parallel_active(
+            self.is_sequence_parallel, is_padding.shape[0], self.sp_min_tokens
+        ):
             return _shard_padding_mask(is_padding, num_tokens)
         return is_padding[:num_tokens]
 
@@ -372,6 +446,7 @@ def forward_prepared(self, hidden_states: torch.Tensor, buffer) -> torch.Tensor:
             a2_scales=self.w2_input_scale.data,
             post_expert_norm_gamma=self._norm_gamma,
             post_expert_norm_eps=self.post_expert_norm.variance_epsilon,
+            weight_layout=self.weight_layout,
         )
         return y
 
diff --git a/vllm/models/muse_spark/nvidia/model.py b/vllm/models/muse_spark/nvidia/model.py
index 779f86043225..94d824f3c092 100644
--- a/vllm/models/muse_spark/nvidia/model.py
+++ b/vllm/models/muse_spark/nvidia/model.py
@@ -61,6 +61,7 @@ def supports_o_norm_sigmoid_gate(*_args: object) -> bool:
 from vllm.logger import init_logger
 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.router.base_router import get_moe_padding_mask
 from vllm.model_executor.layers.linear import (
     QKVParallelLinear,
     ReplicatedLinear,
@@ -106,7 +107,10 @@ def supports_o_norm_sigmoid_gate(*_args: object) -> bool:
     supports_norm_residual_norm,
 )
 from vllm.models.muse_spark.common.ops.fused_routing import fused_musespark1x_routing
-from vllm.models.muse_spark.nvidia.mega_moe import MuseSparkMegaMoEExperts
+from vllm.models.muse_spark.nvidia.mega_moe import (
+    MuseSparkMegaMoEExperts,
+    get_trtllm_token_threshold,
+)
 from vllm.models.muse_spark.nvidia.muse_moe_runner import (
     MuseMoERunner,
     MuseSpark1xRouter,
@@ -411,10 +415,7 @@ def __init__(
         self.n_local_physical_experts = self.n_physical_experts // self.ep_size
         self.n_shared_experts = config.num_shared_experts
 
-        self.fp32_router_input = (
-            getattr(config, "muse_spark_fp32_router_input", True)
-            and torch.backends.cuda.matmul.allow_tf32
-        )
+        self.fp32_router_input = torch.backends.cuda.matmul.allow_tf32
         if self.fp32_router_input and not HAS_ROUTER_OUTPUT:
             raise RuntimeError(
                 "FP32 router input requires FlashInfer dual-norm router output"
@@ -494,6 +495,19 @@ def __init__(
                 "--enable-expert-parallel, or drop the moe backend."
             )
         self.is_sequence_parallel = is_sequence_parallel_enabled(vllm_config)
+        threshold = get_trtllm_token_threshold(vllm_config)
+        self.sp_min_tokens = threshold + 1
+        self.use_hybrid_moe = threshold > 0
+        if self.use_hybrid_moe and (
+            not self.use_mega_moe
+            or self.tp_size not in (4, 8)
+            or parallel_config.data_parallel_size != 1
+            or parallel_config.pipeline_parallel_size != 1
+        ):
+            raise ValueError(
+                "Muse Spark TRT-LLM token threshold requires MegaMoE with "
+                "TEP4/8, DP1 and PP1"
+            )
         # Only the megakernel takes the shard; FusedMoE gets replicated tokens.
         self.use_sequence_parallel_moe = (
             False
@@ -505,23 +519,21 @@ def __init__(
         # loaded, so larger batches keep the replicated projections.
         replicated_tokens = (
             self.tp_size > 1
-            and not self.use_mega_moe
-            and not self.is_sequence_parallel
+            and (
+                (self.use_hybrid_moe and self.is_nvfp4_mega_moe)
+                or (not self.use_mega_moe and not self.is_sequence_parallel)
+            )
             and not self.use_sequence_parallel_moe
             and parallel_config.data_parallel_size == 1
             and not envs.VLLM_BATCH_INVARIANT
         )
-        tep_default = (
+        use_tep_optimizations = (
             replicated_tokens
             and current_platform.is_cuda()
             and current_platform.has_device_capability(100)
         )
-        self.shard_pooler = replicated_tokens and getattr(
-            config, "muse_spark_tep_shard_pooler", tep_default
-        )
-        self.shard_unpooler = replicated_tokens and getattr(
-            config, "muse_spark_tep_shard_unpooler", tep_default
-        )
+        self.shard_pooler = use_tep_optimizations
+        self.shard_unpooler = use_tep_optimizations
         if self.shard_pooler or self.shard_unpooler:
             logger.info_once(
                 "Sharding MuseSpark projection outputs over TP: pooler=%s, "
@@ -541,12 +553,8 @@ def __init__(
         # NVFP4 TRT-LLM layers can take inputs prepared outside the experts:
         # the pre-expert norm fused with the FP4 quant, and/or top-k routing
         # finished on the router's aux stream.
-        self.fused_input_quant = not self.use_mega_moe and getattr(
-            config, "muse_spark_tep_fused_moe_input_quant", tep_default
-        )
-        self.aux_stream_routing = not self.use_mega_moe and getattr(
-            config, "muse_spark_tep_aux_stream_routing", tep_default
-        )
+        self.fused_input_quant = not self.use_mega_moe and use_tep_optimizations
+        self.aux_stream_routing = not self.use_mega_moe and use_tep_optimizations
         self.all_gather_norm_quant: Any | None = None
         if self.use_mega_moe and self.tp_size > 1 and not self.is_sequence_parallel:
             raise NotImplementedError(
@@ -576,9 +584,20 @@ def __init__(
         if self.use_mega_moe:
             max_tokens = vllm_config.scheduler_config.max_num_batched_tokens
             if self.is_sequence_parallel:
-                # Unsharded small batches carry up to tp_size - 1 tokens.
-                max_tokens = max(cdiv(max_tokens, self.tp_size), self.tp_size - 1)
-            self.experts: nn.Module = MuseSparkMegaMoEExperts(
+                # Reserve room for small batches that stay replicated.
+                max_tokens = max(
+                    cdiv(max_tokens, self.tp_size),
+                    self.tp_size - 1,
+                    self.sp_min_tokens - 1,
+                )
+            experts_cls = MuseSparkMegaMoEExperts
+            if self.use_hybrid_moe:
+                from vllm.models.muse_spark.nvidia.hybrid_moe import (
+                    MuseSparkHybridMoEExperts,
+                )
+
+                experts_cls = MuseSparkHybridMoEExperts
+            self.experts: nn.Module = experts_cls(
                 vllm_config,
                 num_experts=self.n_physical_experts,
                 num_local_experts=self.n_local_physical_experts,
@@ -625,9 +644,7 @@ def __init__(
             runner_cls=self.post_norm_runner_cls,
             runner_args={
                 "post_expert_norm": self.post_expert_norm,
-                "fuse_all_reduce": getattr(
-                    config, "muse_spark_tep_fused_moe_all_reduce", tep_default
-                ),
+                "fuse_all_reduce": use_tep_optimizations,
             },
             moe_config_cls=MuseSparkMoEConfig,
         )
@@ -645,15 +662,20 @@ def fold_router_output_multiplier(self) -> None:
         self.router_output_multiplier = None
         self.moe_router.router_output_multiplier = None
 
-    def shards_pooler(self, num_tokens: int) -> bool:
-        return self.shard_pooler and num_tokens <= _SHARDED_PROJECTION_MAX_TOKENS
+    def _can_shard_projection(self, num_tokens: int, sp_active: bool) -> bool:
+        return num_tokens <= _SHARDED_PROJECTION_MAX_TOKENS and (
+            not self.use_mega_moe
+            or (self.use_hybrid_moe and self.experts.use_trtllm(num_tokens, sp_active))
+        )
+
+    def shards_pooler(self, num_tokens: int, sp_active: bool = True) -> bool:
+        return self.shard_pooler and self._can_shard_projection(num_tokens, sp_active)
 
-    def shards_unpooler(self, num_tokens: int) -> bool:
+    def shards_unpooler(self, num_tokens: int, sp_active: bool = True) -> bool:
         return (
             self.shard_unpooler
-            and self.min_sharded_unpooler_tokens
-            <= num_tokens
-            <= _SHARDED_PROJECTION_MAX_TOKENS
+            and self.min_sharded_unpooler_tokens <= num_tokens
+            and self._can_shard_projection(num_tokens, sp_active)
         )
 
     def _project_local_columns(
@@ -743,9 +765,12 @@ def forward(
         hidden_states: torch.Tensor,
         router_input: torch.Tensor | None = None,
         defer_column_gather: bool = False,
+        sp_active: bool = True,
     ) -> torch.Tensor:
         """With ``defer_column_gather``, returns only this rank's unpooler columns."""
-        return self._routed_forward(hidden_states, router_input, defer_column_gather)
+        return self._routed_forward(
+            hidden_states, router_input, defer_column_gather, sp_active
+        )
 
     def _can_fuse_megamoe_inputs(self, hidden_states: torch.Tensor) -> bool:
         norm = self.pre_expert_norm
@@ -762,16 +787,36 @@ def _can_fuse_megamoe_inputs(self, hidden_states: torch.Tensor) -> bool:
         )
 
     def _stage_megamoe_inputs(
-        self, hidden_states: torch.Tensor, router_input: torch.Tensor
+        self, hidden_states: torch.Tensor, router_input: torch.Tensor, sp_active: bool
     ):
         num_tokens = hidden_states.shape[0]
+        use_trtllm = self.use_hybrid_moe and self.experts.use_trtllm(
+            num_tokens, sp_active
+        )
         buffer, padding = self.experts.input_buffer(num_tokens)
+        topk_ids = buffer.topk_idx
+        if use_trtllm:
+            # Reuse the storage for contiguous int32 IDs, without a dtype cast.
+            topk_ids = topk_ids.view(torch.int32).view(-1, topk_ids.shape[-1])
+        topk_ids = topk_ids[:num_tokens]
         norm = self.pre_expert_norm
         pooler = self.pooler
         assert norm is not None and pooler is not None
 
         def pool_and_quantize():
-            pooled, _ = pooler(hidden_states)
+            if use_trtllm and self.shards_pooler(num_tokens, sp_active):
+                pooled = self._project_local_columns(pooler, hidden_states)
+                if fused := self._all_gather_norm_quant(
+                    pooled, norm, self.experts.nvfp4_input_global_scale
+                ):
+                    return None, fused
+                pooled = sp_all_gather(pooled)
+            else:
+                pooled, _ = pooler(hidden_states)
+            if use_trtllm:
+                return None, self.experts.prepare_decode_inputs(
+                    pooled, norm, buffer, num_tokens
+                )
             prepare_megamoe_nvfp4_with_norm(
                 pooled,
                 norm.weight,
@@ -780,7 +825,7 @@ def pool_and_quantize():
                 buffer.x[:num_tokens],
                 buffer.x_sf[:num_tokens].view(torch.uint8),
             )
-            return pooled
+            return pooled, None
 
         def route():
             logits, _ = self.router(router_input)
@@ -793,16 +838,20 @@ def route():
                 router_output_multiplier=router.router_output_multiplier,
                 router_score_normalization=router.router_score_normalization,
                 router_score_division_eps=router.router_score_division_eps,
-                topk_ids=buffer.topk_idx[:num_tokens],
+                topk_ids=topk_ids,
                 topk_weights=buffer.topk_weights[:num_tokens],
                 is_padding=padding,
             )
 
         # The fork event also orders both writers after the previous layer's
         # MegaMoE, which may have consumed this same symmetric buffer.
-        pooled, _ = maybe_execute_in_parallel(
+        (pooled, quantized_input), _ = maybe_execute_in_parallel(
             pool_and_quantize, route, *self._router_events, self._router_aux_stream
         )
+        if use_trtllm:
+            assert quantized_input is not None
+            return self.experts.forward_trtllm(quantized_input, buffer, topk_ids)
+        assert pooled is not None
         return self.experts.forward_prepared(pooled, buffer)
 
     def _prepares_nvfp4_inputs(self, num_tokens: int) -> bool:
@@ -892,8 +941,13 @@ def route() -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
             logits, _ = self.router(router_input)
             if not self.aux_stream_routing:
                 return logits
+            # This path calls TRTLLM's routed experts directly, which skip -1.
+            # An SP all-gather can add rows beyond the runner's padding mask.
+            padding = get_moe_padding_mask(num_tokens)
+            if padding is not None and padding.shape[0] != num_tokens:
+                padding = None
             weights, ids = self.moe_router._compute_routing(
-                hidden_states, logits, torch.int32
+                hidden_states, logits, torch.int32, is_padding=padding
             )
             return ids, weights
 
@@ -913,6 +967,7 @@ def _routed_forward(
         hidden_states: torch.Tensor,
         router_input: torch.Tensor | None = None,
         defer_column_gather: bool = False,
+        sp_active: bool = True,
     ) -> torch.Tensor:
         num_tokens = hidden_states.shape[0]
         if self.use_sequence_parallel_moe:
@@ -923,7 +978,9 @@ def _routed_forward(
             router_input = hidden_states
 
         if self._can_fuse_megamoe_inputs(hidden_states):
-            experts_out = self._stage_megamoe_inputs(hidden_states, router_input)
+            experts_out = self._stage_megamoe_inputs(
+                hidden_states, router_input, sp_active
+            )
         elif self._prepares_nvfp4_inputs(num_tokens):
             experts_out = self._prepared_nvfp4_forward(hidden_states, router_input)
         elif self.use_mega_moe:
@@ -945,7 +1002,7 @@ def _routed_forward(
             return self._project_local_columns(self.unpooler, experts_out)
         if self.unpooler is not None:
             experts_out = self._project(
-                self.unpooler, experts_out, self.shards_unpooler(num_tokens)
+                self.unpooler, experts_out, self.shards_unpooler(num_tokens, sp_active)
             )
 
         if self.use_sequence_parallel_moe:
@@ -1092,6 +1149,7 @@ def __init__(
             prefix=f"{prefix}.o_proj",
         )
         self.is_sequence_parallel_output = False
+        self.sp_min_tokens = 0
         self.gemm_rs = None
 
         self.rotary_emb = (
@@ -1283,7 +1341,7 @@ def forward(
 
     def _project_output(self, attn_output: torch.Tensor) -> torch.Tensor:
         sp_active = is_sequence_parallel_active(
-            self.is_sequence_parallel_output, attn_output.shape[0]
+            self.is_sequence_parallel_output, attn_output.shape[0], self.sp_min_tokens
         )
         if (
             sp_active
@@ -1411,8 +1469,10 @@ def __init__(
         self.all_gather_dual_norm: Any | None = None
 
         self.is_sequence_parallel = is_sequence_parallel_enabled(vllm_config)
+        self.sp_min_tokens = get_trtllm_token_threshold(vllm_config) + 1
         self.is_sequence_parallel_input = self.is_sequence_parallel
         self.self_attn.is_sequence_parallel_output = self.is_sequence_parallel
+        self.self_attn.sp_min_tokens = self.sp_min_tokens
         if self.is_sequence_parallel:
             self.self_attn.o_proj.reduce_results = False
 
@@ -1453,7 +1513,7 @@ def forward(
         assert stream.dtype == self.residual_dtype
 
         sp_active = is_sequence_parallel_active(
-            self.is_sequence_parallel, positions.shape[0]
+            self.is_sequence_parallel, positions.shape[0], self.sp_min_tokens
         )
         if sp_active and self.is_sequence_parallel_input:
             # Positions carries the unpadded length of the gathered sequence.
@@ -1480,11 +1540,14 @@ def forward(
         if sp_active and not self.feed_forward.use_mega_moe:
             # Non-megakernel experts need every token.
             branch = sp_shard(self.feed_forward(sp_all_gather(ffn_in)))
-        elif self._defers_column_gather(ffn_in.shape[0]):
+        elif self._defers_column_gather(ffn_in.shape[0], sp_active):
             return MuseSpark1xColumnShardedStream(
                 stream=stream,
                 branch=self.feed_forward(
-                    ffn_in, router_input=router_input, defer_column_gather=True
+                    ffn_in,
+                    router_input=router_input,
+                    defer_column_gather=True,
+                    sp_active=sp_active,
                 ),
                 post_norm=self.post_ffn_norm,
                 connection=self.add_ffn,
@@ -1492,8 +1555,10 @@ def forward(
             )
         else:
             branch = (
-                self.feed_forward(ffn_in, router_input=router_input)
-                if router_input is not None
+                self.feed_forward(
+                    ffn_in, router_input=router_input, sp_active=sp_active
+                )
+                if isinstance(self.feed_forward, MuseSpark1xMoE)
                 else self.feed_forward(ffn_in)
             )
         return MuseSpark1xResidualStream(
@@ -1503,11 +1568,11 @@ def forward(
             connection=self.add_ffn,
         )
 
-    def _defers_column_gather(self, num_tokens: int) -> bool:
+    def _defers_column_gather(self, num_tokens: int, sp_active: bool) -> bool:
         # Only sharded MoE layers get a workspace (see _configure_column_gather).
         return self.all_gather_dual_norm is not None and cast(
             MuseSpark1xMoE, self.feed_forward
-        ).shards_unpooler(num_tokens)
+        ).shards_unpooler(num_tokens, sp_active)
 
 
 # Runs uncompiled under breakable cudagraphs; drop LlamaModel's torch.compile.
@@ -1545,6 +1610,7 @@ def __init__(
 
         self.fuse_across_layers = get_pp_group().world_size == 1
         self.is_sequence_parallel = is_sequence_parallel_enabled(vllm_config)
+        self.sp_min_tokens = get_trtllm_token_threshold(vllm_config) + 1
         self.make_empty_intermediate_tensors = self._make_empty_intermediate_tensors
         self.is_sequence_parallel_input = self.is_sequence_parallel
         self._configure_replicated_prefix()
@@ -1596,16 +1662,12 @@ def _configure_column_gather(self) -> None:
             if isinstance(layer.feed_forward, MuseSpark1xMoE)
             and layer.feed_forward.shard_unpooler
         ]
-        if (
-            not sharded
-            or not self.fuse_across_layers
-            or not getattr(self.config, "muse_spark_tep_fused_boundary", True)
-        ):
+        if not sharded or not self.fuse_across_layers:
             return
         all_gather = maybe_create_all_gather_dual_norm(
             self.config.hidden_size,
             _SHARDED_PROJECTION_MAX_TOKENS,
-            enable_pdl=getattr(self.config, "muse_spark_tep_fused_boundary_pdl", True),
+            enable_pdl=True,
         )
         if all_gather is None:
             return
@@ -1630,7 +1692,10 @@ def _configure_pooler_gather(self) -> None:
             for layer in islice(self.layers, self.start_layer, self.end_layer)
             if isinstance(layer.feed_forward, MuseSpark1xMoE)
             and layer.feed_forward.shard_pooler
-            and layer.feed_forward.fused_input_quant
+            and (
+                layer.feed_forward.fused_input_quant
+                or layer.feed_forward.use_hybrid_moe
+            )
             and layer.feed_forward.norm_fuses_into_nvfp4_quant
         ]
         if not moes:
@@ -1696,7 +1761,7 @@ def forward(  # type: ignore[override]
             self.residual_dtype
         )
         sp_active = is_sequence_parallel_active(
-            self.is_sequence_parallel, positions.shape[0]
+            self.is_sequence_parallel, positions.shape[0], self.sp_min_tokens
         )
         if sp_active and self.is_sequence_parallel_input:
             x = sp_shard(x)
@@ -1930,13 +1995,12 @@ def interleaved_weights() -> Iterable[tuple[str, torch.Tensor]]:
                 yield name, weight
 
         loader = AutoWeightsLoader(self)
-        loaded = loader.load_weights(
-            interleaved_weights(), mapper=self.hf_to_vllm_mapper
-        )
+        return loader.load_weights(interleaved_weights(), mapper=self.hf_to_vllm_mapper)
+
+    def process_weights_after_loading(self) -> None:
         self.model.finalize_mega_moe_weights()
         n_gates = materialize_residual_path_constants(self)
         logger.info(
             "Materialized residual-path constants for %d gated residual connections",
             n_gates,
         )
-        return loaded
diff --git a/vllm/models/muse_spark/nvidia/ops/prepare_trtllm_nvfp4.py b/vllm/models/muse_spark/nvidia/ops/prepare_trtllm_nvfp4.py
index 7918f4c9e7b2..e68a8a41c406 100644
--- a/vllm/models/muse_spark/nvidia/ops/prepare_trtllm_nvfp4.py
+++ b/vllm/models/muse_spark/nvidia/ops/prepare_trtllm_nvfp4.py
@@ -117,6 +117,8 @@ def norm_quant_trtllm_nvfp4(
     norm_eps: float,
     gain_center: float,
     global_scale: torch.Tensor,
+    *,
+    out: tuple[torch.Tensor, torch.Tensor] | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """RMSNorm + NVFP4 quant of ``[M, H]`` pooled states, in any rank-block layout.
 
@@ -128,10 +130,20 @@ def norm_quant_trtllm_nvfp4(
     shard = pooled.shape[1]
     assert pooled.dtype == torch.bfloat16 and pooled.is_contiguous()
     assert pooled.shape[0] * shard == num_tokens * hidden and hidden % _BLOCK_K == 0
-    packed = torch.empty(
-        num_tokens, hidden // 2, dtype=torch.uint8, device=pooled.device
-    )
-    sf = torch.empty(num_tokens, hidden // 16, dtype=torch.uint8, device=pooled.device)
+    if out is None:
+        packed = torch.empty(
+            num_tokens, hidden // 2, dtype=torch.uint8, device=pooled.device
+        )
+        sf = torch.empty(
+            num_tokens, hidden // 16, dtype=torch.uint8, device=pooled.device
+        )
+    else:
+        packed, sf = out
+        assert packed.shape == (num_tokens, hidden // 2)
+        assert sf.shape == (num_tokens, hidden // 16)
+        assert packed.dtype == sf.dtype == torch.uint8
+        assert packed.is_contiguous() and sf.is_contiguous()
+        assert packed.device == sf.device == pooled.device
     if num_tokens == 0:
         return packed, sf
     enable_pdl = _enable_pdl()
diff --git a/vllm/models/qwen3_5/__init__.py b/vllm/models/qwen3_5/__init__.py
new file mode 100644
index 000000000000..b6def8434743
--- /dev/null
+++ b/vllm/models/qwen3_5/__init__.py
@@ -0,0 +1,48 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Qwen3.5 model — hardware-isolated entry point.
+
+Only the NVIDIA variant migrates to the new eager + breakable CUDA-graph
+modelling code with explicit fused kernels. ROCm and all other platforms keep
+the previous torch.compile implementation in ``vllm/model_executor/models/``
+so their behavior is unchanged.
+"""
+
+from typing import TYPE_CHECKING
+
+from vllm.platforms import current_platform
+
+if TYPE_CHECKING or current_platform.is_cuda():
+    from .nvidia.model import (
+        Qwen3_5ForCausalLM,
+        Qwen3_5ForCausalLMBase,
+        Qwen3_5ForConditionalGeneration,
+        Qwen3_5Model,
+        Qwen3_5MoeForCausalLM,
+        Qwen3_5MoeForConditionalGeneration,
+    )
+    from .nvidia.mtp import Qwen3_5MoeMTP, Qwen3_5MTP
+else:
+    from vllm.model_executor.models.qwen3_5 import (  # type: ignore[assignment]
+        Qwen3_5ForCausalLM,
+        Qwen3_5ForCausalLMBase,
+        Qwen3_5ForConditionalGeneration,
+        Qwen3_5Model,
+        Qwen3_5MoeForCausalLM,
+        Qwen3_5MoeForConditionalGeneration,
+    )
+    from vllm.model_executor.models.qwen3_5_mtp import (  # type: ignore[assignment]
+        Qwen3_5MoeMTP,
+        Qwen3_5MTP,
+    )
+
+__all__ = [
+    "Qwen3_5ForCausalLM",
+    "Qwen3_5ForCausalLMBase",
+    "Qwen3_5ForConditionalGeneration",
+    "Qwen3_5Model",
+    "Qwen3_5MoeForCausalLM",
+    "Qwen3_5MoeForConditionalGeneration",
+    "Qwen3_5MoeMTP",
+    "Qwen3_5MTP",
+]
diff --git a/vllm/models/qwen3_5/nvidia/__init__.py b/vllm/models/qwen3_5/nvidia/__init__.py
new file mode 100644
index 000000000000..208f01a7cb5e
--- /dev/null
+++ b/vllm/models/qwen3_5/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/qwen3_5/nvidia/model.py b/vllm/models/qwen3_5/nvidia/model.py
new file mode 100644
index 000000000000..2a07e89d339b
--- /dev/null
+++ b/vllm/models/qwen3_5/nvidia/model.py
@@ -0,0 +1,1048 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+# Copyright 2025 The vLLM team.
+# Copyright 2025 The Qwen Team.
+# Copyright 2025 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 Qwen3.5 Series compatible with HuggingFace weights."""
+
+from collections.abc import Iterable
+from itertools import islice
+from typing import ClassVar, cast
+
+import torch
+from torch import nn
+from transformers import (
+    Qwen3_5Config,
+    Qwen3_5MoeConfig,
+    Qwen3_5MoeTextConfig,
+    Qwen3_5TextConfig,
+)
+
+from vllm.config import MultiModalConfig, VllmConfig
+from vllm.distributed import (
+    get_pp_group,
+    get_tensor_model_parallel_world_size,
+    tensor_model_parallel_all_gather,
+    tensor_model_parallel_all_reduce,
+)
+from vllm.logger import init_logger
+from vllm.model_executor.layers.fused_allreduce_gemma_rms_norm import (
+    flashinfer_trtllm_fused_allreduce_norm,
+    fused_allreduce_gemma_rms_norm,
+)
+from vllm.model_executor.layers.fusion.mm_input_norm import build_mm_input_norm
+from vllm.model_executor.layers.fusion.quant_activation import (
+    QuantizedActivation,
+    get_input_quant_key,
+)
+from vllm.model_executor.layers.linear import LinearBase
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
+from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import (
+    QwenGatedDeltaNetAttention,
+)
+from vllm.model_executor.layers.mamba.mamba_utils import (
+    MambaStateCopyFunc,
+    MambaStateCopyFuncCalculator,
+    MambaStateDtypeCalculator,
+    MambaStateShapeCalculator,
+)
+from vllm.model_executor.layers.quantization.utils.nvfp4_utils import (
+    nvfp4_weight_padding_bytes,
+)
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    QuantKey,
+    kFp8StaticTensorSym,
+    kNvfp4Dynamic,
+)
+from vllm.model_executor.layers.vocab_parallel_embedding import (
+    ParallelLMHead,
+    VocabParallelEmbedding,
+)
+from vllm.model_executor.models.interfaces import (
+    EagleModelMixin,
+    HasInnerState,
+    IsHybrid,
+    MixtureOfExperts,
+    MultiModalEmbeddings,
+    SupportsEagle3,
+    SupportsLoRA,
+    SupportsMRoPE,
+    SupportsPP,
+    _require_is_multimodal,
+)
+from vllm.model_executor.models.qwen2_moe import Qwen2MoeMLP as Qwen3NextMLP
+from vllm.model_executor.models.qwen3_5 import (
+    Qwen3_5MoeProcessingInfo,
+    Qwen3_5ProcessingInfo,
+)
+from vllm.model_executor.models.qwen3_next import (
+    Qwen3NextAttention,
+    Qwen3NextDecoderLayer,
+    Qwen3NextModel,
+    Qwen3NextSparseMoeBlock,
+    QwenNextMixtureOfExperts,
+)
+from vllm.model_executor.models.qwen3_vl import (
+    Qwen3_VisionTransformer,
+    Qwen3VLDummyInputsBuilder,
+    Qwen3VLForConditionalGeneration,
+    Qwen3VLMultiModalProcessor,
+    mrope_positions_factor,
+)
+from vllm.model_executor.models.utils import (
+    AutoWeightsLoader,
+    PPMissingLayer,
+    WeightsMapper,
+    _merge_multimodal_embeddings,
+    extract_layer_index,
+    make_empty_intermediate_tensors_factory,
+    make_layers,
+    maybe_prefix,
+    sequence_parallel_chunk,
+)
+from vllm.model_executor.utils import set_weight_attrs
+from vllm.multimodal import MULTIMODAL_REGISTRY
+from vllm.multimodal.inputs import MultiModalFeatureSpec
+from vllm.sequence import IntermediateTensors
+from vllm.tokenizers.registry import cached_tokenizer_from_config
+
+from .ops.fused_gdn_gated_norm import fused_gdn_gated_rms_norm
+from .ops.fused_norm_quant import FP8_DTYPE, gemma_add_rms_norm_quant
+
+logger = init_logger(__name__)
+
+
+def _norm_quant_key(linear: nn.Module) -> QuantKey | None:
+    """The input quant key of ``linear`` that the fused norms can produce.
+
+    Checked per call: keys are exposed after weight loading, and LoRA hides
+    them (requires_unquantized_input).
+    """
+    key = get_input_quant_key(linear)
+    if key == kFp8StaticTensorSym:
+        return key
+    # The norm writes block scales for exactly K = hidden_size, in whole
+    # 4-column scale tiles (weights only require K % 32 == 0).
+    if (
+        key == kNvfp4Dynamic
+        and nvfp4_weight_padding_bytes(linear) == 0
+        and linear.input_size_per_partition % 64 == 0
+    ):
+        return key
+    return None
+
+
+def _fused_gdn_gated_rms_norm_quant(
+    *,
+    x: torch.Tensor,
+    out: torch.Tensor,
+    out_proj: LinearBase,
+    **kwargs,
+) -> torch.Tensor | QuantizedActivation:
+    """GDN gated norm, quantized for out_proj when it takes static FP8."""
+    if get_input_quant_key(out_proj) != kFp8StaticTensorSym:
+        fused_gdn_gated_rms_norm(x=x, out=out, **kwargs)
+        return out
+    scale = out_proj.input_scale
+    q = torch.empty(x.shape, dtype=FP8_DTYPE, device=x.device)
+    fused_gdn_gated_rms_norm(x=x, out=q, scale=scale, **kwargs)
+    return QuantizedActivation(
+        data=q,
+        scale=scale,
+        orig_dtype=x.dtype,
+        orig_shape=x.shape,
+        quant_key=kFp8StaticTensorSym,
+    )
+
+
+class Qwen3_5RMSNorm(nn.Module):
+    """Gemma-style RMS normalization backed by flashinfer kernels.
+
+    When ``residual`` is given, the fused add + norm runs in place and the
+    updated ``(x, residual)`` pair is returned.
+    """
+
+    def __init__(self, hidden_size: int, eps: float = 1e-6) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.zeros(hidden_size))
+        set_weight_attrs(self.weight, {"dummy_weight_value": 0.0})
+        self.variance_epsilon = eps
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        residual: torch.Tensor | None = None,
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        from flashinfer.norm import gemma_fused_add_rmsnorm, gemma_rmsnorm
+
+        if residual is None:
+            return gemma_rmsnorm(x, self.weight, self.variance_epsilon)
+
+        # gemma_fused_add_rmsnorm mutates x and residual in place.
+        gemma_fused_add_rmsnorm(x, residual, self.weight, self.variance_epsilon)
+        return x, residual
+
+    def forward_quant(
+        self,
+        x: torch.Tensor,
+        residual: torch.Tensor | None,
+        linear: LinearBase,
+        key: QuantKey,
+        write_y: bool = False,
+    ) -> tuple[QuantizedActivation, torch.Tensor | None, torch.Tensor]:
+        """``forward`` fused with ``linear``'s input quantization ``key``.
+
+        x is not modified. Returns the quantized normed activation, the
+        normed activation in x's dtype when ``write_y`` (else None), and the
+        residual (x itself when ``residual`` is None).
+        """
+        nvfp4 = key == kNvfp4Dynamic
+        scale = linear.input_global_scale_inv if nvfp4 else linear.input_scale
+        q, block_scale, y = gemma_add_rms_norm_quant(
+            x, residual, self.weight, self.variance_epsilon, scale, nvfp4, write_y
+        )
+        qa = QuantizedActivation(
+            data=q,
+            scale=block_scale,
+            orig_dtype=x.dtype,
+            orig_shape=x.shape,
+            quant_key=key,
+        )
+        return qa, y, x if residual is None else residual
+
+
+class Qwen3_5DecoderLayer(Qwen3NextDecoderLayer):
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        layer_type: str,
+        prefix: str = "",
+    ) -> None:
+        super(Qwen3NextDecoderLayer, self).__init__()
+
+        config = vllm_config.model_config.hf_text_config
+        model_config = vllm_config.model_config
+        cache_config = vllm_config.cache_config
+        parallel_config = vllm_config.parallel_config
+        quant_config = vllm_config.quant_config
+
+        self.layer_type = layer_type
+        self.layer_idx = extract_layer_index(prefix)
+        is_moe_layer = config.model_type == "qwen3_5_moe_text"
+        self.use_attn_reduce_scatter_for_moe = (
+            parallel_config.use_sequence_parallel_moe
+            and parallel_config.pipeline_parallel_size == 1
+            and is_moe_layer
+        )
+        # Defer row-parallel all-reduces to the following norm
+        # (deepseek_v32 pattern); exclude PP and SP.
+        self.use_fused_allreduce_norm = (
+            get_tensor_model_parallel_world_size() > 1
+            and parallel_config.pipeline_parallel_size == 1
+            and not self.use_attn_reduce_scatter_for_moe
+            and flashinfer_trtllm_fused_allreduce_norm is not None
+        )
+        # Whether this layer's MLP/MoE leaves an un-reduced partial, i.e.
+        # whether the *next* norm (or the final norm) must do the all-reduce.
+        self.mlp_output_is_partial = self.use_fused_allreduce_norm
+        reduce_results = not (
+            self.use_attn_reduce_scatter_for_moe or self.use_fused_allreduce_norm
+        )
+
+        if self.layer_type == "linear_attention":
+            self.linear_attn = QwenGatedDeltaNetAttention(
+                config=config,
+                vllm_config=vllm_config,
+                prefix=f"{prefix}.linear_attn",
+                gqa_interleaved_layout=False,
+                reduce_results=reduce_results,
+                # This path runs outside torch.compile, as the overlap requires.
+                overlap_ba_proj=True,
+                # Recover the inductor gated-norm fusion (z read strided
+                # from packed qkvz), with out_proj's input quant fused in.
+                fused_gdn_gated_rms_norm=_fused_gdn_gated_rms_norm_quant,
+            )
+        elif self.layer_type == "full_attention":
+            self.self_attn = Qwen3NextAttention(
+                config,
+                model_config=model_config,
+                cache_config=cache_config,
+                quant_config=quant_config,
+                prefix=f"{prefix}.self_attn",
+                reduce_results=reduce_results,
+                mrope_positions_factor=mrope_positions_factor(vllm_config),
+            )
+            self.self_attn.use_fused_output_gate = True
+        else:
+            raise ValueError(f"Invalid layer_type {self.layer_type}")
+
+        if config.model_type == "qwen3_5_moe_text":
+            self.mlp = Qwen3NextSparseMoeBlock(
+                vllm_config=vllm_config,
+                prefix=f"{prefix}.mlp",
+                reduce_results=not self.use_fused_allreduce_norm,
+            )
+            if self.use_fused_allreduce_norm:
+                # reduce_results=False is only honored on the plain-TP late-AR
+                # path (skip_final_all_reduce); EP/all2all, SP, zero experts,
+                # and replicated shared experts all produce or add a complete
+                # output, which a downstream fused norm would all-reduce again.
+                # NOTE: moe_config is shared across layers; the condition is
+                # layer-uniform, so this reset is idempotent.
+                moe_config = self.mlp.experts.moe_config
+                moe_partial_ok = (
+                    moe_config.skip_final_all_reduce
+                    and moe_config.tp_size > 1
+                    and moe_config.ep_size == 1
+                    and moe_config.dp_size == 1
+                    and not self.mlp.replicate_shared_expert
+                )
+                if moe_config.skip_final_all_reduce and not moe_partial_ok:
+                    moe_config.skip_final_all_reduce = False
+                self.mlp_output_is_partial = moe_partial_ok
+        elif config.model_type == "qwen3_5_text":
+            self.mlp = Qwen3NextMLP(
+                hidden_size=config.hidden_size,
+                intermediate_size=config.intermediate_size,
+                hidden_act=config.hidden_act,
+                quant_config=quant_config,
+                prefix=f"{prefix}.mlp",
+                reduce_results=not self.use_fused_allreduce_norm,
+            )
+        else:
+            raise ValueError(f"Invalid model_type {config.model_type}")
+
+        self.input_layernorm = Qwen3_5RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+        self.post_attention_layernorm = Qwen3_5RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+
+        self.layer_scale = getattr(config, "layer_scale", False)
+        if self.layer_scale:
+            self.attn_layer_scale = torch.nn.Parameter(
+                torch.zeros(
+                    1,
+                    1,
+                    config.hidden_size,
+                ),
+            )
+            self.ffn_layer_scale = torch.nn.Parameter(
+                torch.zeros(
+                    1,
+                    1,
+                    config.hidden_size,
+                ),
+            )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        residual: torch.Tensor | None,
+        positions: torch.Tensor,
+        **kwargs: object,
+    ):
+        if self.use_attn_reduce_scatter_for_moe:
+            return super().forward(hidden_states, residual, positions, **kwargs)
+
+        # Each norm absorbs the preceding row-parallel linear's deferred
+        # all-reduce (use_fused_allreduce_norm), or quantizes for its
+        # consumer. layer_scale commutes with the deferred all-reduce.
+        is_gdn = self.layer_type == "linear_attention"
+        in_proj = self.linear_attn.in_proj_qkvz if is_gdn else self.self_attn.qkv_proj
+        in_proj_input = None
+        if residual is not None and self.mlp_output_is_partial:
+            hidden_states, residual = fused_allreduce_gemma_rms_norm(
+                hidden_states, residual, self.input_layernorm
+            )
+        elif (key := _norm_quant_key(in_proj)) is not None:
+            # GDN's in_proj_ba reads the unquantized normed activation.
+            in_proj_input, hidden_states, residual = self.input_layernorm.forward_quant(
+                hidden_states, residual, in_proj, key, write_y=is_gdn
+            )
+        elif residual is None:
+            residual = hidden_states
+            hidden_states = self.input_layernorm(hidden_states)
+        else:
+            hidden_states, residual = self.input_layernorm(hidden_states, residual)
+
+        if not is_gdn:
+            hidden_states = self.self_attn(
+                hidden_states=hidden_states if in_proj_input is None else in_proj_input,
+                positions=positions,
+            )
+        elif in_proj_input is not None:
+            hidden_states = self.linear_attn.forward_cuda(
+                hidden_states, qkvz_input=in_proj_input
+            )
+        else:
+            hidden_states = self.linear_attn(hidden_states=hidden_states)
+
+        if self.layer_scale:
+            if len(hidden_states.shape) == 2:
+                hidden_states = hidden_states * (
+                    self.attn_layer_scale.to(hidden_states.dtype)[0] + 1
+                )
+            else:
+                hidden_states = hidden_states * (
+                    self.attn_layer_scale.to(hidden_states.dtype) + 1
+                )
+
+        # The MoE router needs the unquantized activation.
+        gate_up = self.mlp.gate_up_proj if isinstance(self.mlp, Qwen3NextMLP) else None
+        if self.use_fused_allreduce_norm:
+            hidden_states, residual = fused_allreduce_gemma_rms_norm(
+                hidden_states, residual, self.post_attention_layernorm
+            )
+        elif gate_up is not None and (key := _norm_quant_key(gate_up)) is not None:
+            hidden_states, _, residual = self.post_attention_layernorm.forward_quant(
+                hidden_states, residual, gate_up, key
+            )
+        else:
+            hidden_states, residual = self.post_attention_layernorm(
+                hidden_states, residual
+            )
+        hidden_states = self.mlp(hidden_states)
+
+        if self.layer_scale:
+            if len(hidden_states.shape) == 2:
+                hidden_states = hidden_states * (
+                    self.ffn_layer_scale.to(hidden_states.dtype)[0] + 1
+                )
+            else:
+                assert len(hidden_states.shape) == len(self.ffn_layer_scale.shape), (
+                    f"shape must be the same {len(hidden_states.shape)}, "
+                    f"{len(self.ffn_layer_scale.shape)}"
+                )
+                hidden_states = hidden_states * (
+                    self.ffn_layer_scale.to(hidden_states.dtype) + 1
+                )
+
+        return hidden_states, residual
+
+
+class Qwen3_5Model(nn.Module, EagleModelMixin):
+    # Qwen3.5 ships the GDN in_proj checkpoints separately (qwen3-next
+    # pre-fuses them); fuse them on top of the qwen3-next QKV/gate_up mapping.
+    hf_to_vllm_mapper = Qwen3NextModel.hf_to_vllm_mapper | WeightsMapper(
+        orig_to_new_stacked={
+            ".in_proj_qkv": (".in_proj_qkvz", (0, 1, 2)),
+            ".in_proj_z": (".in_proj_qkvz", 3),
+            ".in_proj_b": (".in_proj_ba", 0),
+            ".in_proj_a": (".in_proj_ba", 1),
+        }
+    )
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+
+        config: Qwen3_5TextConfig | Qwen3_5MoeTextConfig = (
+            vllm_config.model_config.hf_text_config
+        )
+        parallel_config = vllm_config.parallel_config
+
+        eplb_config = parallel_config.eplb_config
+        self.num_redundant_experts = eplb_config.num_redundant_experts
+
+        self.config = config
+        self.quant_config = vllm_config.quant_config
+
+        self.vocab_size = config.vocab_size
+
+        self.embed_tokens = VocabParallelEmbedding(
+            self.vocab_size,
+            config.hidden_size,
+        )
+
+        def get_layer(prefix: str):
+            return Qwen3_5DecoderLayer(
+                vllm_config,
+                layer_type=config.layer_types[extract_layer_index(prefix)],
+                prefix=prefix,
+            )
+
+        self.start_layer, self.end_layer, self.layers = make_layers(
+            config.num_hidden_layers, get_layer, prefix=f"{prefix}.layers"
+        )
+        # Mirror of the decoder layers' flag (uniform across layers): when
+        # True, every layer returns its MLP/MoE output as an un-reduced TP
+        # partial, so the final norm and eagle aux captures must all-reduce.
+        self.mlp_output_is_partial = getattr(
+            self.layers[self.start_layer], "mlp_output_is_partial", False
+        )
+        self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
+            ["hidden_states", "residual"], config.hidden_size
+        )
+
+        if get_pp_group().is_last_rank:
+            self.norm = Qwen3_5RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        else:
+            self.norm = PPMissingLayer()
+
+        self.aux_hidden_state_layers: tuple[int, ...] = ()
+
+    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.embed_tokens(input_ids)
+
+    @property
+    def use_sequence_parallel(self) -> bool:
+        return self.layers[self.start_layer].use_attn_reduce_scatter_for_moe
+
+    def _maybe_add_hidden_state(
+        self,
+        aux_hidden_states: list[torch.Tensor],
+        layer_idx: int,
+        hidden_states: torch.Tensor,
+        residual: torch.Tensor | None,
+    ) -> list[torch.Tensor]:
+        if residual is None and layer_idx in self.aux_hidden_state_layers:
+            # The layer-0 capture aliases the embeddings, which become the
+            # first residual and are mutated in place by the next layer's
+            # fused-add norm; capture by value.
+            hidden_states = hidden_states.clone()
+        if (
+            residual is not None
+            and self.mlp_output_is_partial
+            and layer_idx in self.aux_hidden_state_layers
+        ):
+            # hidden_states is an un-reduced TP partial; the aux capture needs
+            # the complete hidden state.
+            hidden_states = tensor_model_parallel_all_reduce(hidden_states)
+        return super()._maybe_add_hidden_state(
+            aux_hidden_states, layer_idx, hidden_states, residual
+        )
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None,
+        positions: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+    ) -> 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
+            else:
+                hidden_states = self.embed_input_ids(input_ids)
+            residual = None
+        else:
+            assert intermediate_tensors is not None
+            hidden_states = intermediate_tensors["hidden_states"]
+            residual = intermediate_tensors["residual"]
+
+        full_num_tokens = positions.shape[-1]
+        if self.use_sequence_parallel:
+            hidden_states = sequence_parallel_chunk(hidden_states)
+            assert residual is None
+
+        aux_hidden_states = self._maybe_add_hidden_state([], 0, hidden_states, residual)
+        for layer_idx, layer in enumerate(
+            islice(self.layers, self.start_layer, self.end_layer),
+            start=self.start_layer,
+        ):
+            hidden_states, residual = layer(
+                positions=positions,
+                hidden_states=hidden_states,
+                residual=residual,
+            )
+            self._maybe_add_hidden_state(
+                aux_hidden_states, layer_idx + 1, hidden_states, residual
+            )
+
+        if not get_pp_group().is_last_rank:
+            return IntermediateTensors(
+                {"hidden_states": hidden_states, "residual": residual}
+            )
+        if self.mlp_output_is_partial:
+            # The last layer's MLP/MoE output is an un-reduced TP partial;
+            # fuse its all-reduce into the final norm.
+            hidden_states, _ = fused_allreduce_gemma_rms_norm(
+                hidden_states, residual, self.norm
+            )
+        else:
+            hidden_states, _ = self.norm(hidden_states, residual)
+        if self.use_sequence_parallel:
+            if aux_hidden_states:
+                hidden_size = hidden_states.shape[-1]
+                hidden_states = torch.cat([hidden_states, *aux_hidden_states], dim=-1)
+                hidden_states = tensor_model_parallel_all_gather(hidden_states, 0)
+                hidden_states = hidden_states[:full_num_tokens]
+                hidden_states, *aux_hidden_states = hidden_states.split(
+                    hidden_size, dim=-1
+                )
+            else:
+                hidden_states = tensor_model_parallel_all_gather(hidden_states, 0)
+                hidden_states = hidden_states[:full_num_tokens]
+        if aux_hidden_states:
+            return hidden_states, aux_hidden_states
+        return hidden_states
+
+    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 Qwen3_5ForCausalLMBase(
+    nn.Module,
+    HasInnerState,
+    IsHybrid,
+    SupportsEagle3,
+    SupportsLoRA,
+    SupportsMRoPE,
+    SupportsPP,
+):
+    packed_modules_mapping = {
+        "qkv_proj": [
+            "q_proj",
+            "k_proj",
+            "v_proj",
+        ],
+        "gate_up_proj": ["gate_proj", "up_proj"],
+        # GDN fused projections.
+        "in_proj_qkvz": ["in_proj_qkv", "in_proj_z"],
+        "in_proj_ba": ["in_proj_b", "in_proj_a"],
+    }
+    # Maps PEFT embed/lm_head LoRA targets onto vLLM embedding wrappers.
+    embedding_modules = {
+        "embed_tokens": "input_embeddings",
+        "lm_head": "output_embeddings",
+    }
+
+    # Some community text-only checkpoints keep the extraneous
+    # `model.language_model.` prefix inherited from the VL training stack.
+    # Strip it so both prefixed and clean checkpoints load correctly.
+    hf_to_vllm_mapper = WeightsMapper(
+        orig_to_new_prefix={"model.language_model.": "model.", "mtp.": None},
+    )
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        config = vllm_config.model_config.hf_text_config
+        self.vllm_config = vllm_config
+        self.model_config = vllm_config.model_config
+
+        scheduler_config = vllm_config.scheduler_config
+        self.quant_config = vllm_config.quant_config
+
+        super().__init__()
+        self.config = config
+        self.scheduler_config = scheduler_config
+        self.model = Qwen3_5Model(
+            vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
+        )
+
+        if get_pp_group().is_last_rank:
+            if config.tie_word_embeddings:
+                self.lm_head = self.model.embed_tokens
+            else:
+                self.lm_head = ParallelLMHead(
+                    config.vocab_size,
+                    config.hidden_size,
+                    quant_config=self.quant_config,
+                    prefix=maybe_prefix(prefix, "lm_head"),
+                )
+        else:
+            self.lm_head = PPMissingLayer()
+
+        self.logits_processor = LogitsProcessor(config.vocab_size)
+        self.make_empty_intermediate_tensors = (  # type: ignore[method-assign]
+            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 set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None:
+        self.model.aux_hidden_state_layers = layers
+
+    def get_eagle3_aux_hidden_state_layers(self) -> tuple[int, ...]:
+        num_layers = len(self.model.layers)
+        return (2, num_layers // 2, num_layers - 3)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        **kwargs: object,
+    ):
+        hidden_states = self.model(
+            input_ids, positions, intermediate_tensors, inputs_embeds
+        )
+
+        return hidden_states
+
+    @classmethod
+    def get_mamba_state_dtype_from_config(
+        cls,
+        vllm_config: "VllmConfig",
+    ) -> tuple[torch.dtype, torch.dtype]:
+        return MambaStateDtypeCalculator.gated_delta_net_state_dtype(
+            vllm_config.model_config.dtype,
+            vllm_config.cache_config.mamba_cache_dtype,
+            vllm_config.cache_config.mamba_ssm_cache_dtype,
+        )
+
+    @classmethod
+    def get_mamba_state_shape_from_config(
+        cls, vllm_config: "VllmConfig"
+    ) -> tuple[tuple[int, int], tuple[int, int, int]]:
+        parallel_config = vllm_config.parallel_config
+        hf_config = vllm_config.model_config.hf_text_config
+        tp_size = parallel_config.tensor_parallel_size
+        num_spec = (
+            vllm_config.speculative_config.num_speculative_tokens
+            if vllm_config.speculative_config
+            else 0
+        )
+        return MambaStateShapeCalculator.gated_delta_net_state_shape(
+            tp_size,
+            hf_config.linear_num_key_heads,
+            hf_config.linear_num_value_heads,
+            hf_config.linear_key_head_dim,
+            hf_config.linear_value_head_dim,
+            hf_config.linear_conv_kernel_dim,
+            num_spec,
+        )
+
+    @classmethod
+    def get_mamba_state_copy_func(
+        cls,
+    ) -> tuple[MambaStateCopyFunc, MambaStateCopyFunc]:
+        return MambaStateCopyFuncCalculator.gated_delta_net_state_copy_func()
+
+    def compute_logits(
+        self,
+        hidden_states: torch.Tensor,
+    ) -> torch.Tensor | None:
+        return self.logits_processor(self.lm_head, hidden_states)
+
+    def compute_logits_local(
+        self,
+        hidden_states: torch.Tensor,
+    ) -> torch.Tensor:
+        return self.logits_processor(self.lm_head, hidden_states, skip_gather=True)
+
+    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)
+
+    def get_mrope_input_positions(
+        self,
+        input_tokens: list[int],
+        mm_features: list[MultiModalFeatureSpec],
+    ) -> tuple[torch.Tensor, int]:
+        positions = torch.arange(len(input_tokens), dtype=torch.long)
+        return positions.unsqueeze(0).expand(3, -1), 0
+
+
+class Qwen3_5ForCausalLM(Qwen3_5ForCausalLMBase):
+    pass
+
+
+class Qwen3_5MoeForCausalLM(Qwen3_5ForCausalLMBase, QwenNextMixtureOfExperts):
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__(vllm_config=vllm_config, prefix=prefix)
+
+        # set MoE hyperparameters
+        self.set_moe_parameters()
+
+
+########################################################
+# Qwen3_5-Dense
+########################################################
+
+
+@MULTIMODAL_REGISTRY.register_processor(
+    Qwen3VLMultiModalProcessor,
+    info=Qwen3_5ProcessingInfo,
+    dummy_inputs=Qwen3VLDummyInputsBuilder,
+)
+class Qwen3_5ForConditionalGeneration(Qwen3VLForConditionalGeneration, IsHybrid):
+    supports_multimodal_pruning = True
+
+    hf_to_vllm_mapper = (
+        Qwen3VLForConditionalGeneration.hf_to_vllm_mapper
+        | WeightsMapper(orig_to_new_prefix={"mtp.": None})
+    )
+
+    packed_modules_mapping = Qwen3VLForConditionalGeneration.packed_modules_mapping | {
+        "in_proj_qkvz": ["in_proj_qkv", "in_proj_z"],
+        "in_proj_ba": ["in_proj_b", "in_proj_a"],
+    }
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "model"):
+        # protocols have not __init__ method, so we need to use nn.Module.__init__
+        nn.Module.__init__(self)
+        config: Qwen3_5Config = vllm_config.model_config.hf_config
+        quant_config = vllm_config.quant_config
+        multimodal_config = cast(
+            MultiModalConfig, vllm_config.model_config.multimodal_config
+        )
+
+        self.config = config
+        self.model_config = vllm_config.model_config
+        self.multimodal_config = multimodal_config
+        self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data"
+        self.is_multimodal_pruning_enabled = (
+            multimodal_config.is_multimodal_pruning_enabled()
+        )
+        self.video_pruning_rate = self.multimodal_config.video_pruning_rate
+        self._tokenizer = cached_tokenizer_from_config(vllm_config.model_config)
+
+        # attributes needed by EVS-related functions inherited from Qwen3-VL
+        self.use_deepstack = hasattr(config.vision_config, "deepstack_visual_indexes")
+        self.deepstack_num_level = (
+            len(config.vision_config.deepstack_visual_indexes)
+            if self.use_deepstack
+            else 0
+        )
+        self.visual_dim = config.vision_config.out_hidden_size
+        self.multiscale_dim = self.visual_dim * self.deepstack_num_level
+
+        with self._mark_tower_model(vllm_config, {"image", "video"}):
+            self.visual = Qwen3_VisionTransformer(
+                config.vision_config,
+                norm_eps=getattr(config, "rms_norm_eps", 1e-6),
+                quant_config=quant_config,
+                input_norm=build_mm_input_norm(self.model_config),
+                prefix=maybe_prefix(prefix, "visual"),
+            )
+
+        with self._mark_language_model(vllm_config):
+            self.language_model = Qwen3_5ForCausalLM(
+                vllm_config=vllm_config, prefix=maybe_prefix(prefix, "language_model")
+            )
+
+        self.make_empty_intermediate_tensors = (  # type: ignore[method-assign]
+            self.language_model.make_empty_intermediate_tensors
+        )
+
+    def embed_input_ids(
+        self,
+        input_ids: torch.Tensor,
+        multimodal_embeddings: MultiModalEmbeddings | None = None,
+        *,
+        is_multimodal: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        inputs_embeds = self._embed_text_input_ids(
+            input_ids,
+            self.language_model.embed_input_ids,
+            is_multimodal=is_multimodal,
+        )
+
+        if multimodal_embeddings is None or len(multimodal_embeddings) == 0:
+            return inputs_embeds
+
+        is_multimodal = _require_is_multimodal(is_multimodal)
+
+        inputs_embeds = _merge_multimodal_embeddings(
+            inputs_embeds=inputs_embeds,
+            multimodal_embeddings=multimodal_embeddings,
+            is_multimodal=is_multimodal,
+        )
+
+        return inputs_embeds
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        **kwargs: object,
+    ) -> torch.Tensor | IntermediateTensors:
+        """Run forward pass for Qwen3.5.
+
+        With mrope enabled (default for Qwen3VL-class models), ``positions``
+        has shape `(3, seq_len)`, otherwise `(seq_len,)`.
+        """
+        if intermediate_tensors is not None:
+            inputs_embeds = None
+
+        hidden_states = self.language_model.model(
+            input_ids=input_ids,
+            positions=positions,
+            intermediate_tensors=intermediate_tensors,
+            inputs_embeds=inputs_embeds,
+        )
+
+        return hidden_states
+
+    def compute_logits_local(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.language_model.compute_logits_local(hidden_states)
+
+    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)
+
+    @classmethod
+    def get_mamba_state_dtype_from_config(
+        cls,
+        vllm_config: "VllmConfig",
+    ) -> tuple[torch.dtype, torch.dtype]:
+        return MambaStateDtypeCalculator.gated_delta_net_state_dtype(
+            vllm_config.model_config.dtype,
+            vllm_config.cache_config.mamba_cache_dtype,
+            vllm_config.cache_config.mamba_ssm_cache_dtype,
+        )
+
+    @classmethod
+    def get_mamba_state_shape_from_config(
+        cls, vllm_config: "VllmConfig"
+    ) -> tuple[tuple[int, int], tuple[int, int, int]]:
+        parallel_config = vllm_config.parallel_config
+        hf_config = vllm_config.model_config.hf_text_config
+        tp_size = parallel_config.tensor_parallel_size
+        num_spec = (
+            vllm_config.speculative_config.num_speculative_tokens
+            if vllm_config.speculative_config
+            else 0
+        )
+        return MambaStateShapeCalculator.gated_delta_net_state_shape(
+            tp_size,
+            hf_config.linear_num_key_heads,
+            hf_config.linear_num_value_heads,
+            hf_config.linear_key_head_dim,
+            hf_config.linear_value_head_dim,
+            hf_config.linear_conv_kernel_dim,
+            num_spec,
+        )
+
+    @classmethod
+    def get_mamba_state_copy_func(cls) -> tuple[MambaStateCopyFunc, MambaStateCopyFunc]:
+        return MambaStateCopyFuncCalculator.gated_delta_net_state_copy_func()
+
+
+########################################################
+# Qwen3_5-MoE
+########################################################
+
+
+class Qwen3_5_MoeMixtureOfExperts(MixtureOfExperts):
+    language_model: Qwen3_5MoeForCausalLM
+    num_local_physical_experts: int
+
+    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.language_model.model.layers:
+            if isinstance(layer.mlp, Qwen3NextSparseMoeBlock):
+                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 set_moe_parameters(self):
+        self.moe_layers = []
+        example_moe = None
+        for layer in self.language_model.model.layers:
+            if isinstance(layer, Qwen3_5DecoderLayer) and isinstance(
+                layer.mlp, Qwen3NextSparseMoeBlock
+            ):
+                example_moe = layer.mlp
+                self.moe_layers.append(layer.mlp.experts)
+
+        if example_moe is None:
+            raise RuntimeError(
+                "No Qwen3_5 layer found in the language_model.model.layers."
+            )
+
+        # Set MoE hyperparameters
+        self.num_moe_layers = len(self.moe_layers)
+        self.num_expert_groups = 1
+        self.num_shared_experts = 0
+        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
+        self.num_routed_experts = example_moe.n_routed_experts
+        self.num_redundant_experts = example_moe.n_redundant_experts
+
+
+@MULTIMODAL_REGISTRY.register_processor(
+    Qwen3VLMultiModalProcessor,
+    info=Qwen3_5MoeProcessingInfo,
+    dummy_inputs=Qwen3VLDummyInputsBuilder,
+)
+class Qwen3_5MoeForConditionalGeneration(
+    Qwen3_5ForConditionalGeneration, Qwen3_5_MoeMixtureOfExperts
+):
+    # For MoE LoRA weights loading
+    is_3d_moe_weight: ClassVar[bool] = True
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "model"):
+        # protocols have not __init__ method, so we need to use nn.Module.__init__
+        nn.Module.__init__(self)
+        config: Qwen3_5MoeConfig = vllm_config.model_config.hf_config
+        quant_config = vllm_config.quant_config
+        multimodal_config = cast(
+            MultiModalConfig, vllm_config.model_config.multimodal_config
+        )
+
+        self.config = config
+        self.model_config = vllm_config.model_config
+        self.multimodal_config = multimodal_config
+        self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data"
+        self.is_multimodal_pruning_enabled = (
+            multimodal_config.is_multimodal_pruning_enabled()
+        )
+        self.video_pruning_rate = self.multimodal_config.video_pruning_rate
+        self._tokenizer = cached_tokenizer_from_config(vllm_config.model_config)
+
+        # attributes needed by EVS-related functions inherited from Qwen3-VL
+        self.use_deepstack = hasattr(config.vision_config, "deepstack_visual_indexes")
+        self.deepstack_num_level = (
+            len(config.vision_config.deepstack_visual_indexes)
+            if self.use_deepstack
+            else 0
+        )
+        self.visual_dim = config.vision_config.out_hidden_size
+        self.multiscale_dim = self.visual_dim * self.deepstack_num_level
+
+        with self._mark_tower_model(vllm_config, {"image", "video"}):
+            self.visual = Qwen3_VisionTransformer(
+                config.vision_config,
+                norm_eps=getattr(config, "rms_norm_eps", 1e-6),
+                quant_config=quant_config,
+                input_norm=build_mm_input_norm(self.model_config),
+                prefix=maybe_prefix(prefix, "visual"),
+            )
+
+        with self._mark_language_model(vllm_config):
+            self.language_model = Qwen3_5MoeForCausalLM(
+                vllm_config=vllm_config, prefix=maybe_prefix(prefix, "language_model")
+            )
+
+        self.make_empty_intermediate_tensors = (  # type: ignore[method-assign]
+            self.language_model.make_empty_intermediate_tensors
+        )
+
+        # set MoE hyperparameters
+        self.set_moe_parameters()
diff --git a/vllm/models/qwen3_5/nvidia/mtp.py b/vllm/models/qwen3_5/nvidia/mtp.py
new file mode 100644
index 000000000000..dfdee2d90bc1
--- /dev/null
+++ b/vllm/models/qwen3_5/nvidia/mtp.py
@@ -0,0 +1,308 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Inference-only Qwen3_5 MTP model."""
+
+from collections.abc import Iterable
+
+import torch
+from torch import nn
+from transformers import Qwen3_5MoeTextConfig, Qwen3_5TextConfig
+
+from vllm.config import VllmConfig
+from vllm.distributed import tensor_model_parallel_all_gather
+from vllm.distributed.parallel_state import get_pp_group
+from vllm.kernels.triton.fused_pre_fc_norms import fused_pre_fc_norms
+from vllm.logger import init_logger
+from vllm.model_executor.layers.fused_allreduce_gemma_rms_norm import (
+    fused_allreduce_gemma_rms_norm,
+)
+from vllm.model_executor.layers.linear import ColumnParallelLinear
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
+from vllm.model_executor.layers.vocab_parallel_embedding import (
+    ParallelLMHead,
+    VocabParallelEmbedding,
+)
+from vllm.model_executor.models.interfaces import (
+    LocalArgmaxMixin,
+    MultiModalEmbeddings,
+    SupportsMultiModal,
+    SupportsPP,
+    _require_is_multimodal,
+)
+from vllm.model_executor.models.qwen3_next import QwenNextMixtureOfExperts
+from vllm.model_executor.models.utils import (
+    AutoWeightsLoader,
+    PPMissingLayer,
+    _merge_multimodal_embeddings,
+    make_empty_intermediate_tensors_factory,
+    maybe_prefix,
+    sequence_parallel_chunk,
+)
+from vllm.sequence import IntermediateTensors
+
+from .model import Qwen3_5DecoderLayer, Qwen3_5Model, Qwen3_5RMSNorm
+
+logger = init_logger(__name__)
+
+
+class Qwen3_5MultiTokenPredictor(nn.Module):
+    hf_to_vllm_mapper = Qwen3_5Model.hf_to_vllm_mapper
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+
+        model_config = vllm_config.model_config
+        quant_config = vllm_config.quant_config
+
+        config: Qwen3_5TextConfig | Qwen3_5MoeTextConfig = model_config.hf_text_config
+
+        self.config = config
+
+        self.vocab_size = config.vocab_size
+
+        self.mtp_start_layer_idx = config.num_hidden_layers
+        self.num_mtp_layers = getattr(config, "mtp_num_hidden_layers", 1)
+
+        self.embed_tokens = VocabParallelEmbedding(
+            self.vocab_size,
+            config.hidden_size,
+        )
+
+        # Workaround: mtp.fc is stored as BF16 in NVFP4 checkpoints but is
+        # missing from hf_quant_config.json exclude_modules. Force unquantized.
+        # Ref: https://github.com/vllm-project/vllm/pull/38650
+        # Ref: https://github.com/NVIDIA/Model-Optimizer/pull/1124
+        fc_quant = (
+            None
+            if (quant_config and quant_config.get_name() == "modelopt_fp4")
+            else quant_config
+        )
+        self.fc = ColumnParallelLinear(
+            self.config.hidden_size * 2,
+            self.config.hidden_size,
+            gather_output=True,
+            bias=False,
+            return_bias=False,
+            quant_config=fc_quant,
+            prefix=f"{prefix}.fc",
+        )
+
+        # GPTQ: quantized checkpoints may exclude MTP from quantization via
+        # quantization_config.dynamic with "-:pattern" entries. When detected,
+        # disable quantization for MTP layers so they use unquantized params.
+        original_quant = vllm_config.quant_config
+        if quant_config and quant_config.get_name() not in ("modelopt_fp4",):
+            hf_qc = getattr(model_config.hf_config, "quantization_config", None)
+            if isinstance(hf_qc, dict):
+                dynamic = hf_qc.get("dynamic", {})
+                if any(k.startswith("-:") and "mtp" in k for k in dynamic):
+                    vllm_config.quant_config = None
+        self.layers = torch.nn.ModuleList(
+            Qwen3_5DecoderLayer(
+                vllm_config,
+                layer_type="full_attention",
+                prefix=f"{prefix}.layers.{idx}",
+            )
+            for idx in range(self.num_mtp_layers)
+        )
+        vllm_config.quant_config = original_quant
+        self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
+            ["hidden_states", "residual"], config.hidden_size
+        )
+        self.norm = Qwen3_5RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.pre_fc_norm_hidden = Qwen3_5RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+        self.pre_fc_norm_embedding = Qwen3_5RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+
+    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.embed_tokens(input_ids)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        spec_step_idx: int = 0,
+    ) -> torch.Tensor:
+        # 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]
+            hidden_states = fused_pre_fc_norms(
+                inputs_embeds,
+                hidden_states,
+                self.pre_fc_norm_embedding.weight,
+                self.pre_fc_norm_hidden.weight,
+                self.pre_fc_norm_embedding.variance_epsilon,
+            )
+            hidden_states = self.fc(hidden_states)
+            residual = None
+        else:
+            hidden_states = intermediate_tensors["hidden_states"]
+            residual = intermediate_tensors["residual"]
+
+        current_step_idx = spec_step_idx % self.num_mtp_layers
+        mtp_layer = self.layers[current_step_idx]
+        if mtp_layer.use_attn_reduce_scatter_for_moe:
+            assert hidden_states.shape[0] == positions.shape[-1]
+            hidden_states = sequence_parallel_chunk(hidden_states)
+            assert residual is None
+        hidden_states, residual = mtp_layer(
+            positions=positions,
+            hidden_states=hidden_states,
+            residual=residual,
+        )
+
+        if not get_pp_group().is_last_rank:
+            return IntermediateTensors(
+                {"hidden_states": hidden_states, "residual": residual}
+            )
+
+        if mtp_layer.mlp_output_is_partial:
+            # The MTP layer's MLP/MoE output is an un-reduced TP partial;
+            # fuse its all-reduce into the final norm.
+            hidden_states, _ = fused_allreduce_gemma_rms_norm(
+                hidden_states, residual, self.norm
+            )
+        else:
+            hidden_states, _ = self.norm(hidden_states, residual)
+        if mtp_layer.use_attn_reduce_scatter_for_moe:
+            hidden_states = tensor_model_parallel_all_gather(hidden_states, 0)
+            hidden_states = hidden_states[: positions.shape[-1]]
+        return hidden_states
+
+    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 Qwen3_5MTP(LocalArgmaxMixin, nn.Module, SupportsMultiModal, SupportsPP):
+    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 = ""):
+        config = vllm_config.model_config.hf_text_config
+        self.vllm_config = vllm_config
+        self.quant_config = vllm_config.quant_config
+
+        super().__init__()
+        self.config = config
+        self.model = Qwen3_5MultiTokenPredictor(
+            vllm_config=vllm_config, prefix=maybe_prefix(prefix, "mtp")
+        )
+
+        if get_pp_group().is_last_rank:
+            self.lm_head = ParallelLMHead(
+                config.vocab_size,
+                config.hidden_size,
+                quant_config=self.quant_config,
+                prefix=maybe_prefix(prefix, "lm_head"),
+            )
+            if config.tie_word_embeddings:
+                self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)
+        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,
+        multimodal_embeddings: MultiModalEmbeddings | None = None,
+        *,
+        is_multimodal: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        inputs_embeds = self._embed_text_input_ids(
+            input_ids,
+            self.model.embed_input_ids,
+            is_multimodal=is_multimodal,
+        )
+
+        if multimodal_embeddings is None or len(multimodal_embeddings) == 0:
+            return inputs_embeds
+
+        is_multimodal = _require_is_multimodal(is_multimodal)
+
+        inputs_embeds = _merge_multimodal_embeddings(
+            inputs_embeds=inputs_embeds,
+            multimodal_embeddings=multimodal_embeddings,
+            is_multimodal=is_multimodal,
+        )
+
+        return inputs_embeds
+
+    def forward(  # type: ignore[override]
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        spec_step_idx: int = 0,
+    ):
+        hidden_states = self.model(
+            input_ids,
+            positions,
+            hidden_states,
+            intermediate_tensors,
+            inputs_embeds,
+            spec_step_idx,
+        )
+        return hidden_states
+
+    def compute_logits(
+        self,
+        hidden_states: torch.Tensor,
+        spec_step_idx: int = 0,
+    ) -> torch.Tensor | None:
+        return self.logits_processor(self.lm_head, hidden_states)
+
+    def get_top_tokens(
+        self,
+        hidden_states: torch.Tensor,
+        spec_step_idx: int = 0,
+    ) -> torch.Tensor:
+        # One lm_head is shared by all MTP modules, as in compute_logits.
+        return super().get_top_tokens(hidden_states)
+
+    def is_unused_checkpoint_weight(self, name: str) -> bool:
+        return not name.startswith("mtp.") and not any(
+            key in name for key in ["embed_tokens", "lm_head"]
+        )
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        def remap_weight_names(weights):
+            for name, weight in weights:
+                if self.is_unused_checkpoint_weight(name):
+                    continue
+                if name.startswith("mtp."):
+                    name = name.replace("mtp.", "model.")
+                elif "embed_tokens" in name:
+                    name = name.replace("language_model.", "")
+                yield name, weight
+
+        loader = AutoWeightsLoader(self)
+        return loader.load_weights(remap_weight_names(weights))
+
+
+class Qwen3_5MoeMTP(Qwen3_5MTP, QwenNextMixtureOfExperts):
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__(vllm_config=vllm_config, prefix=prefix)
+        self.set_moe_parameters()
diff --git a/vllm/models/qwen3_5/nvidia/ops/__init__.py b/vllm/models/qwen3_5/nvidia/ops/__init__.py
new file mode 100644
index 000000000000..208f01a7cb5e
--- /dev/null
+++ b/vllm/models/qwen3_5/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/qwen3_5/nvidia/ops/fused_gdn_gated_norm.py b/vllm/models/qwen3_5/nvidia/ops/fused_gdn_gated_norm.py
new file mode 100644
index 000000000000..b090d6d5c201
--- /dev/null
+++ b/vllm/models/qwen3_5/nvidia/ops/fused_gdn_gated_norm.py
@@ -0,0 +1,138 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import torch
+
+from vllm.third_party.flash_linear_attention.ops.layernorm_guard import (
+    calc_rows_per_block,
+)
+from vllm.triton_utils import tl, triton
+
+
+@triton.jit(do_not_specialize=["M"])
+def _fused_gdn_gated_rms_norm_kernel(
+    x_ptr,
+    z_ptr,
+    w_ptr,
+    out_ptr,
+    scale_ptr,
+    stride_x_tok,
+    stride_z_tok,
+    stride_out_tok,
+    M,  # number of tokens
+    eps,
+    N: tl.constexpr,  # head_v_dim
+    BLOCK_N: tl.constexpr,
+    ROWS_PER_BLOCK: tl.constexpr,
+    NORM_BEFORE_GATE: tl.constexpr,
+    ACTIVATION: tl.constexpr,
+    QUANT_FP8: tl.constexpr,
+):
+    """FLA-equivalent per-head gated RMS norm with token-strided z."""
+    row_start = tl.program_id(0) * ROWS_PER_BLOCK
+    head = tl.program_id(1)
+
+    rows = row_start + tl.arange(0, ROWS_PER_BLOCK)
+    cols = tl.arange(0, BLOCK_N)
+    row_mask = rows[:, None] < M
+    col_mask = cols[None, :] < N
+    mask = row_mask & col_mask
+
+    x_ptrs = x_ptr + rows[:, None] * stride_x_tok + head * N + cols[None, :]
+    z_ptrs = z_ptr + rows[:, None] * stride_z_tok + head * N + cols[None, :]
+
+    x = tl.load(x_ptrs, mask=mask, other=0.0).to(tl.float32)
+
+    if not NORM_BEFORE_GATE:
+        z = tl.load(z_ptrs, mask=mask, other=0.0).to(tl.float32)
+        if ACTIVATION == "swish" or ACTIVATION == "silu":
+            x *= z * tl.sigmoid(z)
+        elif ACTIVATION == "sigmoid":
+            x *= tl.sigmoid(z)
+
+    xbar = tl.where(mask, x, 0.0)
+    var = tl.sum(xbar * xbar, axis=1) / N
+    rstd = tl.rsqrt(var + eps)
+
+    w = tl.load(w_ptr + cols, mask=cols < N, other=0.0).to(tl.float32)
+    y = x * rstd[:, None] * w[None, :]
+
+    if NORM_BEFORE_GATE:
+        z = tl.load(z_ptrs, mask=mask, other=0.0).to(tl.float32)
+        if ACTIVATION == "swish" or ACTIVATION == "silu":
+            y *= z * tl.sigmoid(z)
+        elif ACTIVATION == "sigmoid":
+            y *= tl.sigmoid(z)
+
+    out_ptrs = out_ptr + rows[:, None] * stride_out_tok + head * N + cols[None, :]
+    if QUANT_FP8:
+        # Static FP8 from the bf16-rounded output, like scaled_fp8_quant.
+        y = y.to(x_ptr.dtype.element_ty).to(tl.float32)
+        y = y * tl.math.div_rn(1.0, tl.load(scale_ptr))
+        y = tl.clamp(y, -448.0, 448.0)
+    tl.store(out_ptrs, y.to(out_ptr.dtype.element_ty), mask=mask)
+
+
+def fused_gdn_gated_rms_norm(
+    x: torch.Tensor,
+    z: torch.Tensor,
+    weight: torch.Tensor,
+    out: torch.Tensor,
+    eps: float,
+    num_v_heads: int,
+    head_v_dim: int,
+    norm_before_gate: bool,
+    activation: str,
+    scale: torch.Tensor | None = None,
+) -> None:
+    """Gated RMS norm ``out = rms_norm(x) * act(z)`` for Qwen3.5 GDN.
+
+    Args:
+        x: (num_tokens, num_v_heads * head_v_dim), contiguous. May alias out.
+        z: (num_tokens, num_v_heads * head_v_dim) gate values, typically a
+            strided view into the packed in_proj_qkvz output. Only the last
+            dim must be contiguous.
+        weight: (head_v_dim,) shared per-head norm weight.
+        out: (num_tokens, num_v_heads * head_v_dim), contiguous.
+        eps: RMS norm epsilon.
+        num_v_heads: number of V heads per TP rank.
+        head_v_dim: head_v_dim (per-row norm size).
+        norm_before_gate: if True, out = norm(x) * act(z), else
+            out = norm(x * act(z)).
+        activation: "swish"/"silu" or "sigmoid".
+        scale: optional static per-tensor FP8 scale; when given, out is
+            float8_e4m3fn and must not alias x.
+
+    """
+    assert x.shape == z.shape == out.shape and x.dim() == 2
+    assert x.stride(-1) == 1 and z.stride(-1) == 1 and out.stride(-1) == 1
+    assert weight.shape == (head_v_dim,) and weight.stride(0) == 1
+    assert x.shape[1] == num_v_heads * head_v_dim
+    quant_fp8 = scale is not None
+    if scale is not None:
+        assert out.dtype == torch.float8_e4m3fn and scale.numel() == 1
+
+    num_tokens = x.shape[0]
+    if num_tokens == 0:
+        return
+    block_n = triton.next_power_of_2(head_v_dim)
+    rows_per_block = calc_rows_per_block(num_tokens, x.device)
+    grid = (triton.cdiv(num_tokens, rows_per_block), num_v_heads)
+    _fused_gdn_gated_rms_norm_kernel[grid](
+        x,
+        z,
+        weight,
+        out,
+        scale if quant_fp8 else weight,
+        x.stride(0),
+        z.stride(0),
+        out.stride(0),
+        num_tokens,
+        eps,
+        N=head_v_dim,
+        BLOCK_N=block_n,
+        ROWS_PER_BLOCK=rows_per_block,
+        NORM_BEFORE_GATE=norm_before_gate,
+        ACTIVATION=activation,
+        QUANT_FP8=quant_fp8,
+        num_warps=min(max(block_n // 256, 1), 8),
+    )
diff --git a/vllm/models/qwen3_5/nvidia/ops/fused_norm_quant.py b/vllm/models/qwen3_5/nvidia/ops/fused_norm_quant.py
new file mode 100644
index 000000000000..65abf238b832
--- /dev/null
+++ b/vllm/models/qwen3_5/nvidia/ops/fused_norm_quant.py
@@ -0,0 +1,216 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Gemma (1 + w) add-RMSNorm fused with the consumer linear's input quant.
+
+Replaces ``gemma_fused_add_rmsnorm`` followed by a standalone static FP8
+(``scaled_fp8_quant``) or NVFP4 (``scaled_fp4_quant``) kernel. The quantized
+values are computed from the bf16-rounded normed activation with the same
+scale arithmetic as those kernels, so codes match the unfused chain.
+"""
+
+import torch
+
+from vllm.platforms import current_platform
+from vllm.triton_utils import tl, triton
+from vllm.utils.math_utils import round_up
+
+FP8_DTYPE = torch.float8_e4m3fn
+FP8_MAX = tl.constexpr(448.0)
+
+
+@triton.jit
+def _rcp_approx_ftz(x):
+    # Matches reciprocal_approximate_ftz in csrc/quantization/fp4/nvfp4_utils.cuh.
+    return tl.inline_asm_elementwise(
+        "rcp.approx.ftz.f32 $0, $1;",
+        "=r,r",
+        [x],
+        dtype=tl.float32,
+        is_pure=True,
+        pack=1,
+    )
+
+
+@triton.jit
+def _cvt_e2m1x2(lo, hi):
+    """Pack fp32 (lo, hi) pairs into E2M1 bytes (lo in the low nibble)."""
+    return tl.inline_asm_elementwise(
+        """
+        {
+        .reg .b8 b0, b1, b2, b3;
+        cvt.rn.satfinite.e2m1x2.f32 b0, $5, $1;
+        cvt.rn.satfinite.e2m1x2.f32 b1, $6, $2;
+        cvt.rn.satfinite.e2m1x2.f32 b2, $7, $3;
+        cvt.rn.satfinite.e2m1x2.f32 b3, $8, $4;
+        mov.b32 $0, {b0, b1, b2, b3};
+        }
+        """,
+        "=r,r,r,r,r,r,r,r,r",
+        [lo, hi],
+        dtype=tl.uint8,
+        is_pure=True,
+        pack=4,
+    )
+
+
+@triton.jit
+def _gemma_add_rms_norm_quant_kernel(
+    x_ptr,
+    res_ptr,
+    w_ptr,
+    scale_ptr,
+    q_ptr,
+    sf_ptr,
+    y_ptr,
+    eps,
+    H: tl.constexpr,
+    HAS_RESIDUAL: tl.constexpr,
+    QUANT_NVFP4: tl.constexpr,
+    WRITE_Y: tl.constexpr,
+    launch_pdl: tl.constexpr,
+    TRIGGER_EARLY: tl.constexpr,
+):
+    # Index the row as [G, 8, 2]: G groups of 16 (one NVFP4 block each) of 8
+    # (lo, hi) pairs, so block reductions and fp4 packing need no reshapes.
+    G: tl.constexpr = triton.next_power_of_2(H) // 16
+    row = tl.program_id(0).to(tl.int64)
+    g = tl.arange(0, G)
+    pair = tl.arange(0, 8)
+    cols = (
+        g[:, None, None] * 16 + pair[None, :, None] * 2 + tl.arange(0, 2)[None, None, :]
+    )
+    mask = cols < H
+
+    if launch_pdl:
+        tl.extra.cuda.gdc_wait()
+        if TRIGGER_EARLY:
+            # Launch dependents before the loads: the consumer (e.g. the
+            # weight-loading GEMM) overlaps with this kernel.
+            tl.extra.cuda.gdc_launch_dependents()
+    x = tl.load(x_ptr + row * H + cols, mask=mask, other=0.0)
+    if HAS_RESIDUAL:
+        r = tl.load(res_ptr + row * H + cols, mask=mask, other=0.0)
+    w = tl.load(w_ptr + cols, mask=mask, other=0.0)
+    scale = tl.load(scale_ptr)
+
+    h = x.to(tl.float32)
+    if HAS_RESIDUAL:
+        h += r.to(tl.float32)
+        # The norm uses the fp32 sum, like flashinfer's fused_add_rmsnorm.
+        tl.store(res_ptr + row * H + cols, h, mask=mask)
+    var = tl.sum(h * h) / H
+    rstd = tl.rsqrt(var + eps)
+    # Round to the activation dtype, like the unfused norm output.
+    y = (h * rstd * (1.0 + w.to(tl.float32))).to(x_ptr.dtype.element_ty)
+    if WRITE_Y:
+        tl.store(y_ptr + row * H + cols, y, mask=mask)
+    y = y.to(tl.float32)
+
+    if QUANT_NVFP4:
+        # scaled_fp4_quant numerics; `scale` is the global scale (448 * 6 / amax).
+        vmax = tl.max(tl.max(tl.abs(y), axis=2), axis=1)
+        sf = scale * (vmax * _rcp_approx_ftz(tl.full((G,), 6.0, tl.float32)))
+        sf_fp8 = tl.minimum(sf, FP8_MAX).to(tl.float8e4nv)
+        sf32 = sf_fp8.to(tl.float32)
+        out_scale = tl.where(
+            sf32 != 0.0, _rcp_approx_ftz(sf32 * _rcp_approx_ftz(scale)), 0.0
+        )
+        lo, hi = tl.split(y * out_scale[:, None, None])
+        pcols = g[:, None] * 8 + pair[None, :]
+        tl.store(
+            q_ptr + row * (H // 2) + pcols, _cvt_e2m1x2(lo, hi), mask=pcols < H // 2
+        )
+
+        # Swizzled 128x4 scale layout [m_tile, k_tile, 32, 4, 4].
+        num_k_tiles: tl.constexpr = (H // 16 + 3) // 4
+        sf_off = (
+            ((row >> 7) * num_k_tiles + (g >> 2)) * 512
+            + (row & 31) * 16
+            + ((row >> 5) & 3) * 4
+            + (g & 3)
+        )
+        tl.store(sf_ptr + sf_off, sf_fp8, mask=g < H // 16)
+    else:
+        inv_scale = tl.math.div_rn(1.0, scale)
+        q = tl.clamp(y * inv_scale, -FP8_MAX, FP8_MAX)
+        tl.store(q_ptr + row * H + cols, q, mask=mask)
+
+    if launch_pdl and not TRIGGER_EARLY:
+        tl.extra.cuda.gdc_launch_dependents()
+
+
+def gemma_add_rms_norm_quant(
+    x: torch.Tensor,
+    residual: torch.Tensor | None,
+    weight: torch.Tensor,
+    eps: float,
+    scale: torch.Tensor,
+    nvfp4: bool,
+    write_y: bool = False,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
+    """Gemma RMSNorm of ``x (+ residual)`` quantized for the next linear.
+
+    Args:
+        x: (num_tokens, hidden) bf16/fp16 activation, contiguous. Read only.
+        residual: same shape as x, or None. Updated in place to x + residual.
+        weight: (hidden,) Gemma norm weight; the norm scales by (1 + weight).
+        eps: RMS norm epsilon.
+        scale: FP8: the static per-tensor input scale. NVFP4: the global
+            scale (``input_global_scale_inv``).
+        nvfp4: quantize to NVFP4 with swizzled 128x4 block scales instead of
+            static FP8.
+        write_y: also return the bf16 normed activation.
+
+    Returns:
+        ``(q, block_scale, y)``: quantized data (fp8, or packed fp4 uint8),
+        NVFP4 block scales (fp8, ``(round_up(M, 128), round_up(H // 16, 4))``;
+        rows >= M are left unwritten) or ``scale`` for FP8, and the normed
+        activation (or None).
+
+    """
+    assert x.dim() == 2 and x.is_contiguous()
+    assert x.dtype in (torch.bfloat16, torch.float16)
+    assert residual is None or (residual.shape == x.shape and residual.is_contiguous())
+    assert weight.shape == (x.shape[1],) and weight.is_contiguous()
+    assert scale.numel() == 1 and scale.dtype == torch.float32
+    M, H = x.shape
+    if nvfp4:
+        # Whole 4-column scale tiles: padded scale columns are not written.
+        assert H % 64 == 0
+        q = torch.empty((M, H // 2), dtype=torch.uint8, device=x.device)
+        block_scale = torch.empty(
+            (round_up(M, 128), round_up(H // 16, 4) // 4),
+            dtype=torch.int32,
+            device=x.device,
+        ).view(FP8_DTYPE)
+    else:
+        q = torch.empty((M, H), dtype=FP8_DTYPE, device=x.device)
+        block_scale = scale
+    y = torch.empty_like(x) if write_y else None
+    if M == 0:
+        return q, block_scale, y
+
+    # Swept on GB300 at H=5120, M=1..8192: within 3% of the best of 2/4/8/16
+    # warps at every swept M except NVFP4 M=129 (4%).
+    max_m_16_warps = 128 if nvfp4 else 384
+    _gemma_add_rms_norm_quant_kernel[(M,)](
+        x,
+        residual if residual is not None else x,
+        weight,
+        scale,
+        q,
+        block_scale,
+        y if y is not None else x,
+        eps,
+        H=H,
+        HAS_RESIDUAL=residual is not None,
+        QUANT_NVFP4=nvfp4,
+        WRITE_Y=write_y,
+        num_warps=16 if max_m_16_warps >= M else 4,
+        launch_pdl=current_platform.is_arch_support_pdl(),
+        # Overlaps the consumer's weight load with this kernel. Swept on
+        # GB300: ~2% win at M <= 128; the prelaunched consumer grid steals
+        # SMs at M >= 192 (fp8 +4-7%).
+        TRIGGER_EARLY=M <= 128,
+    )
+    return q, block_scale, y
diff --git a/vllm/models/qwen4_exp/amd/model.py b/vllm/models/qwen4_exp/amd/model.py
index 5715d038dd1a..675b9a9a6cff 100644
--- a/vllm/models/qwen4_exp/amd/model.py
+++ b/vllm/models/qwen4_exp/amd/model.py
@@ -12,9 +12,6 @@
 from vllm.compilation.decorators import support_torch_compile
 from vllm.config import VllmConfig
 from vllm.distributed import get_pp_group
-from vllm.model_executor.layers.fused_moe.utils import (
-    is_model_fused_shared_expert_compatible,
-)
 from vllm.model_executor.layers.fusion.mm_input_norm import build_mm_input_norm
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
 from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import (
@@ -64,6 +61,7 @@
     WeightsMapper,
     _merge_multimodal_embeddings,
     extract_layer_index,
+    get_fused_shared_expert_layers,
     make_empty_intermediate_tensors_factory,
     make_layers,
     maybe_fuse_shared_experts,
@@ -446,10 +444,10 @@ def get_layer(prefix: str) -> Qwen4ExpDecoderLayer:
         self.start_layer, self.end_layer, self.layers = make_layers(
             config.num_hidden_layers, get_layer, prefix=f"{prefix}.layers"
         )
-        self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible(
-            self.layers,
+        self.fused_shared_expert_layers = get_fused_shared_expert_layers(
+            self,
             Qwen4ExpSparseMoeBlock,
-            "mlp",
+            enable_eplb=vllm_config.parallel_config.enable_eplb,
         )
         intermediate_size = config.hidden_size * config.hc_count
         self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
@@ -620,7 +618,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         )
         weights = maybe_fuse_shared_experts(
             weights,
-            enabled=self.is_fused_shared_expert_enabled,
+            fused_shared_expert_layers=self.fused_shared_expert_layers,
             n_routed_experts=getattr(self.config, "num_experts", 0) or 0,
             n_shared_experts=1,
             ckpt_prefix="mlp.shared_expert",
diff --git a/vllm/models/qwen4_exp/amd/mtp.py b/vllm/models/qwen4_exp/amd/mtp.py
index 0805d102b0c9..7a73ac8f33ed 100644
--- a/vllm/models/qwen4_exp/amd/mtp.py
+++ b/vllm/models/qwen4_exp/amd/mtp.py
@@ -23,9 +23,6 @@
 from vllm.compilation.decorators import support_torch_compile
 from vllm.config import VllmConfig, replace, set_current_vllm_config
 from vllm.distributed import get_pp_group
-from vllm.model_executor.layers.fused_moe.utils import (
-    is_model_fused_shared_expert_compatible,
-)
 from vllm.model_executor.layers.layernorm import GemmaRMSNorm
 from vllm.model_executor.layers.linear import ColumnParallelLinear
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
@@ -41,6 +38,7 @@
     PPMissingLayer,
     WeightsMapper,
     get_draft_quant_config,
+    get_fused_shared_expert_layers,
     make_empty_intermediate_tensors_factory,
     maybe_fuse_shared_experts,
     maybe_prefix,
@@ -229,10 +227,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
                 )
                 for idx in range(self.num_mtp_layers)
             )
-        self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible(
-            self.layers,
+        self.fused_shared_expert_layers = get_fused_shared_expert_layers(
+            self,
             Qwen4ExpSparseMoeBlock,
-            "mlp",
+            enable_eplb=vllm_config.parallel_config.enable_eplb,
         )
 
         self.pre_fc_norm_embedding = GemmaRMSNorm(
@@ -373,7 +371,7 @@ def forward(
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         weights = maybe_fuse_shared_experts(
             weights,
-            enabled=self.is_fused_shared_expert_enabled,
+            fused_shared_expert_layers=self.fused_shared_expert_layers,
             n_routed_experts=getattr(self.config, "num_experts", 0) or 0,
             n_shared_experts=1,
             ckpt_prefix="mlp.shared_expert",
diff --git a/vllm/models/qwen4_exp/amd/ops/hc.py b/vllm/models/qwen4_exp/amd/ops/hc.py
index 1faf89b020a3..0f8017f53674 100644
--- a/vllm/models/qwen4_exp/amd/ops/hc.py
+++ b/vllm/models/qwen4_exp/amd/ops/hc.py
@@ -27,7 +27,7 @@ def _grouped_gemma_rmsnorm_kernel(
 
     pid = tl.program_id(0)
     group_id = pid % NUM_GROUPS
-    row = pid // NUM_GROUPS
+    row = (pid // NUM_GROUPS).to(tl.int64)
 
     offs_g = tl.arange(0, BLOCK_SIZE)
     offsets = group_id * GROUP_DIM + offs_g
@@ -90,7 +90,7 @@ def _hc_silu_kernel(
 ) -> None:
     BLOCK_SIZE: tl.constexpr = triton.next_power_of_2(DIM)
 
-    row = tl.program_id(0)
+    row = tl.program_id(0).to(tl.int64)
     offs = tl.arange(0, BLOCK_SIZE)
     mask = offs < DIM
 
@@ -137,7 +137,7 @@ def _hc_gate_mix_kernel(
 ) -> None:
     HC_DIM: tl.constexpr = DIM // HC
 
-    row = tl.program_id(0)
+    row = tl.program_id(0).to(tl.int64)
     tile_id = tl.program_id(1)
     offs_inner = tile_id * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
     mask = offs_inner < HC_DIM
@@ -202,7 +202,7 @@ def _hc_combine_kernel(
 ) -> None:
     HC_PAD: tl.constexpr = triton.next_power_of_2(HC)
 
-    row = tl.program_id(0)
+    row = tl.program_id(0).to(tl.int64)
     tile_id = tl.program_id(1)
 
     offs_inner = tile_id * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
@@ -287,7 +287,7 @@ def _hc_combine_norm_kernel(
     NUM_TILES: tl.constexpr = triton.cdiv(HC_DIM, BLOCK_SIZE)
     NUM_TILES_PAD: tl.constexpr = triton.next_power_of_2(NUM_TILES)
 
-    row = tl.program_id(0)
+    row = tl.program_id(0).to(tl.int64)
     stream = tl.program_id(1)
     offs_hc = tl.arange(0, HC_PAD)
     mask_hc = offs_hc < HC
diff --git a/vllm/models/qwen4_exp/amd/ops/qsa.py b/vllm/models/qwen4_exp/amd/ops/qsa.py
index 6bdb1fe592df..1ad1f8323ac8 100644
--- a/vllm/models/qwen4_exp/amd/ops/qsa.py
+++ b/vllm/models/qwen4_exp/amd/ops/qsa.py
@@ -220,6 +220,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)
@@ -292,6 +293,14 @@ 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 (an fp8
+            # QK dot measured slower and less accurate on the NVIDIA port).
+            keys = keys.to(query.dtype)
+            # Dequant V to fp16 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)
         scores = tl.dot(query, keys)
         # Scaling scores avoids re-quantizing a scaled query to BF16.
         scores *= softmax_scale_log2
@@ -825,7 +834,7 @@ def qsa_sparse_paged_attention(
     token_to_req: torch.Tensor,
     out: torch.Tensor | None = None,
 ) -> torch.Tensor:
-    """Run sparse GQA directly over paged BF16 K/V caches."""
+    """Run sparse GQA directly over paged bf16 or fp8 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:
@@ -842,7 +851,12 @@ 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
+    is_fp8 = k_cache.dtype in (torch.float8_e4m3fn, torch.float8_e4m3fnuz)
+    assert q.dtype == torch.bfloat16
+    assert k_cache.dtype == v_cache.dtype
+    assert is_fp8 or k_cache.dtype == torch.bfloat16, (
+        "QSA sparse attention requires BF16 or FP8-e4m3 K/V"
+    )
     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
@@ -938,6 +952,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=partial_stages,
     )
diff --git a/vllm/models/qwen4_exp/amd/ops/qsa_prepare.py b/vllm/models/qwen4_exp/amd/ops/qsa_prepare.py
index bc2fee3b3a47..80fe084031ad 100644
--- a/vllm/models/qwen4_exp/amd/ops/qsa_prepare.py
+++ b/vllm/models/qwen4_exp/amd/ops/qsa_prepare.py
@@ -124,7 +124,7 @@ def _to_dst_dtype(x, dst, scale):
     """Round to BF16 like the unfused path, then scale for an FP8 destination."""
     out_ty = dst.dtype.element_ty
     x = x.to(tl.bfloat16)
-    if out_ty == tl.float8e4nv:
+    if out_ty.is_fp8():
         x = x.to(tl.float32) / scale
     return x.to(out_ty)
 
diff --git a/vllm/models/qwen4_exp/amd/qsa.py b/vllm/models/qwen4_exp/amd/qsa.py
index 6a246dd31ec8..d0203089d4ae 100644
--- a/vllm/models/qwen4_exp/amd/qsa.py
+++ b/vllm/models/qwen4_exp/amd/qsa.py
@@ -10,6 +10,7 @@
 from torch import nn
 from transformers import Qwen4ExpTextConfig
 
+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 import get_tensor_model_parallel_world_size
@@ -35,17 +36,10 @@
 )
 from vllm.v1.attention.backend import (
     AttentionBackend,
-    AttentionCGSupport,
+    AttentionImplBase,
     AttentionType,
     MultipleOf,
 )
-from vllm.v1.attention.backends.fa_utils import is_flash_attn_varlen_func_available
-from vllm.v1.attention.backends.flash_attn import (
-    FlashAttentionBackend,
-    FlashAttentionImpl,
-    FlashAttentionMetadata,
-    FlashAttentionMetadataBuilder,
-)
 from vllm.v1.attention.ops import rocm_octave
 from vllm.v1.kv_cache_interface import (
     AttentionSpec,
@@ -58,9 +52,10 @@
     is_workspace_manager_initialized,
 )
 
-from ..common.qsa_cache import QSAForwardMetadata
+from ..common.qsa import Qwen4ExpQSAMetadata, Qwen4ExpQSAMetadataBuilder
+from ..common.qsa_indexer import QSAIndexerMetadata, qsa_kv_cache_dtype
 from . import model
-from .indexer_qsa import QSAIndexer
+from .qsa_indexer import QSAIndexer
 
 
 def _octave_mid_o_capacity(vllm_config: VllmConfig) -> int:
@@ -80,10 +75,8 @@ def _octave_mid_o_capacity(vllm_config: VllmConfig) -> int:
 _OCTAVE_SINK_ENTRIES = 2048
 
 
-class Qwen4ExpQSAMetadataBuilder(FlashAttentionMetadataBuilder):
-    """Flash metadata supporting uniform decode and target-verify graphs."""
-
-    _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
+class Qwen4ExpROCmQSAMetadataBuilder(Qwen4ExpQSAMetadataBuilder):
+    """Reserve the Octave workspace for ROCm QSA."""
 
     def __init__(
         self,
@@ -112,10 +105,16 @@ def __init__(
     "octave_k3v3",
     "octave_k3v3_compact",
 )
-_QSA_CACHE_DTYPES: tuple[CacheDType, ...] = ("auto", "bfloat16", *_OCTAVE_CACHE_DTYPES)
+_QSA_CACHE_DTYPES: tuple[CacheDType, ...] = (
+    "auto",
+    "bfloat16",
+    "fp8",
+    "fp8_e4m3",
+    *_OCTAVE_CACHE_DTYPES,
+)
 
 
-class Qwen4ExpQSAFlashAttentionBackend(FlashAttentionBackend):
+class Qwen4ExpQSABackend(AttentionBackend):
     """FullAttentionSpec backend used by the merged QSA owner.
 
     With an Octave ``kv_cache_dtype`` the main K/V cache holds one packed
@@ -124,7 +123,11 @@ class Qwen4ExpQSAFlashAttentionBackend(FlashAttentionBackend):
     """
 
     supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16]
+    # fp8/fp8_e4m3: uint8 storage read by the QSA Triton kernel.
+    # octave_*: packed slots handled by the Octave kernels.
     supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = list(_QSA_CACHE_DTYPES)
+    # The owner writes K/V before calling the sparse attention implementation.
+    forward_includes_kv_cache_update: bool = False
 
     @classmethod
     def customize_spec(cls, spec: AttentionSpec) -> AttentionSpec:
@@ -144,12 +147,12 @@ def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleO
         return [MultipleOf(16)]
 
     @staticmethod
-    def get_impl_cls() -> type[Qwen4ExpQSAFlashAttentionImpl]:
-        return Qwen4ExpQSAFlashAttentionImpl
+    def get_impl_cls() -> type[Qwen4ExpQSAImpl]:
+        return Qwen4ExpQSAImpl
 
     @staticmethod
-    def get_builder_cls() -> type[Qwen4ExpQSAMetadataBuilder]:
-        return Qwen4ExpQSAMetadataBuilder
+    def get_builder_cls() -> type[Qwen4ExpROCmQSAMetadataBuilder]:
+        return Qwen4ExpROCmQSAMetadataBuilder
 
     @classmethod
     def is_sparse(cls) -> bool:
@@ -160,25 +163,31 @@ def supports_kv_connector(cls) -> bool:
         return False
 
 
-class Qwen4ExpQSAFlashAttentionImpl(FlashAttentionImpl):
+class Qwen4ExpQSAImpl(AttentionImplBase[Qwen4ExpQSAMetadata]):
     """Run paged sparse GQA with the QSA Triton kernel."""
 
-    supports_dcp: bool = False
-    supports_pcp: bool = False
-
-    def __init__(self, *args, **kwargs) -> None:
-        super().__init__(*args, **kwargs)
-        if not is_flash_attn_varlen_func_available():
-            raise NotImplementedError("Qwen4Exp QSA requires FlashAttention")
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        scale: float,
+        num_kv_heads: int,
+        kv_cache_dtype: str,
+    ) -> None:
+        """Initialize QSA state and preserve the optional Octave cache path."""
         if self.dcp_world_size != 1:
             raise NotImplementedError(
                 "Qwen4Exp QSA does not support decode context parallelism"
             )
-        if self.kv_cache_dtype not in _QSA_CACHE_DTYPES:
+        if kv_cache_dtype not in _QSA_CACHE_DTYPES:
             raise NotImplementedError(
-                "Qwen4Exp QSA requires a BF16 or Octave main KV cache"
+                "Qwen4Exp QSA requires a BF16, FP8-e4m3, or Octave main KV cache"
             )
-        self.supports_quant_query_input = False
+        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
         self.octave_fmt: rocm_octave.OctaveFormat | None = None
         if self.kv_cache_dtype in _OCTAVE_CACHE_DTYPES:
             self.octave_fmt = rocm_octave.OctaveFormat.from_cache_dtype(
@@ -245,8 +254,21 @@ def do_kv_cache_update(
         kv_cache: torch.Tensor,
         slot_mapping: torch.Tensor,
     ) -> None:
+        """Store BF16 K/V directly or use the existing Octave cache writer."""
         if self.octave_fmt is None:
-            super().do_kv_cache_update(layer, key, value, kv_cache, slot_mapping)
+            key_cache, value_cache = kv_cache.transpose(1, 2).split(
+                self.head_size, dim=-1
+            )
+            ops.reshape_and_cache_flash(
+                key,
+                value,
+                key_cache,
+                value_cache,
+                slot_mapping,
+                self.kv_cache_dtype,
+                layer._k_scale,
+                layer._v_scale,
+            )
             return
         n = slot_mapping.shape[0]
         if n == 0:
@@ -275,27 +297,17 @@ def do_kv_cache_update(
             self._slocks,
         )
 
-    def forward_qsa(
+    def forward(
         self,
         layer: torch.nn.Module,
         query: torch.Tensor,
-        key: torch.Tensor | None,
-        value: torch.Tensor | None,
         kv_cache: torch.Tensor,
-        attn_metadata: FlashAttentionMetadata,
+        attn_metadata: Qwen4ExpQSAMetadata,
         output: torch.Tensor,
+        *,
         token_to_req: torch.Tensor,
-        output_scale: torch.Tensor | None = None,
-        output_block_scale: torch.Tensor | None = None,
     ) -> torch.Tensor:
-        del key, value
-        if output_scale is not None or output_block_scale is not None:
-            raise NotImplementedError("QSA does not support fused output quantization")
-        if self.alibi_slopes is not None or self.sinks is not None:
-            raise NotImplementedError("QSA does not support ALiBi or attention sinks")
-        if self.sliding_window != (-1, -1):
-            raise NotImplementedError("QSA does not support sliding-window attention")
-
+        """Run sparse attention over the owner's indexer-selected K/V rows."""
         num_tokens = attn_metadata.num_actual_tokens
         output.zero_()
         if num_tokens == 0:
@@ -332,8 +344,24 @@ def forward_qsa(
         key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1)
         key_cache = canonicalize_singleton_dim_strides(key_cache)
         value_cache = canonicalize_singleton_dim_strides(value_cache)
-        if key_cache.dtype != torch.bfloat16 or query.dtype != torch.bfloat16:
-            raise NotImplementedError("Qwen4Exp QSA requires BF16 Q/K/V")
+        if self.kv_cache_dtype in ("fp8", "fp8_e4m3"):
+            assert key_cache.dtype == torch.uint8, (
+                "fp8 KV cache must be uint8 storage; a config mismatch would "
+                "silently double the view's shape"
+            )
+            fp8_dtype = current_platform.fp8_dtype()
+            key_cache = key_cache.view(fp8_dtype)
+            value_cache = value_cache.view(fp8_dtype)
+        if query.dtype != torch.bfloat16:
+            raise NotImplementedError("Qwen4Exp QSA requires BF16 queries")
+        if key_cache.dtype != torch.bfloat16 and (
+            layer._k_scale_float != 1.0 or layer._v_scale_float != 1.0
+        ):
+            raise NotImplementedError(
+                "Qwen4Exp QSA with an FP8 KV cache requires unit scales on "
+                "the AMD path (the NVIDIA port's host-side dequant fold is "
+                "not yet ported here). Use bf16 KV caches or unit scales."
+            )
 
         from .ops.qsa import qsa_sparse_paged_attention
 
@@ -350,7 +378,7 @@ def forward_qsa(
 
 
 class Qwen4ExpQSAAttention(Qwen3NextAttention, AttentionLayerBase):
-    """Merged Qwen full-attention owner with a QSA index side branch."""
+    """Merged Qwen full-attention owner with a QSA indexer branch."""
 
     supports_dcp = False
 
@@ -371,9 +399,10 @@ 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 _QSA_CACHE_DTYPES:
+        cache_dtype = qsa_kv_cache_dtype(cache_config, prefix)
+        if cache_dtype not in _QSA_CACHE_DTYPES:
             raise NotImplementedError(
-                "Qwen4Exp QSA requires a BF16 or Octave main KV cache"
+                "Qwen4Exp QSA requires a BF16, FP8-e4m3, or Octave main KV cache"
             )
         if getattr(quant_config, "kv_cache_scheme", None) is not None:
             raise NotImplementedError("Qwen4Exp QSA does not support KV quantization")
@@ -452,31 +481,28 @@ def __init__(
 
         self.layer_name = f"{prefix}.attn"
         self.attn_type = AttentionType.DECODER
-        self.kv_cache_dtype = cache_config.cache_dtype
+        self.kv_cache_dtype = cache_dtype
         self.kv_cache_torch_dtype = kv_cache_dtype_str_to_dtype(
             self.kv_cache_dtype, model_config
         )
         if (
-            self.kv_cache_dtype not in _OCTAVE_CACHE_DTYPES
+            self.kv_cache_dtype not in _QSA_CACHE_DTYPES
             and self.kv_cache_torch_dtype != torch.bfloat16
         ):
-            raise NotImplementedError("Qwen4Exp QSA requires BF16 cache storage")
+            raise NotImplementedError(
+                "Qwen4Exp QSA requires a BF16, FP8-e4m3, or Octave main KV cache"
+            )
         self.kv_sharing_target_layer_name = None
         self.kv_cache = torch.tensor([])
         set_default_quant_scales(self, register_buffer=True)
 
-        self.attn_backend = Qwen4ExpQSAFlashAttentionBackend
-        self.impl = Qwen4ExpQSAFlashAttentionImpl(
+        self.attn_backend = Qwen4ExpQSABackend
+        self.impl: Qwen4ExpQSAImpl = Qwen4ExpQSAImpl(  # type: ignore[assignment]
             self.num_heads,
             self.head_dim,
             self.scaling,
             self.num_kv_heads,
-            None,
-            None,
             self.kv_cache_dtype,
-            None,
-            AttentionType.DECODER,
-            None,
         )
         self.indexer = QSAIndexer(
             vllm_config=vllm_config,
@@ -486,8 +512,8 @@ def __init__(
             quant_config=quant_config,
             prefix=f"{prefix}.indexer",
         )
-        # The fused prepare writes raw BF16 K/V; an Octave cache is packed by
-        # do_kv_cache_update instead.
+        # The fused prepare writes bf16 or platform-fp8 K/V; an Octave cache
+        # is packed by do_kv_cache_update instead.
         self.use_fused_qsa_prepare = (
             self.attn_output_gate
             and self.indexer.use_fused_pre_indexer
@@ -544,17 +570,15 @@ def _run_qsa(
             if gate is not None:
                 gate.zero_()
             return
-        main_metadata = cast(FlashAttentionMetadata, metadata[self.layer_name])
+        main_metadata = cast(Qwen4ExpQSAMetadata, metadata[self.layer_name])
         if self.kv_cache.numel() == 0:
             raise RuntimeError("QSA main K/V cache is not bound")
 
         num_tokens = main_metadata.num_actual_tokens
-        side_metadata = cast(
-            QSAForwardMetadata,
+        indexer_metadata = cast(
+            QSAIndexerMetadata,
             metadata[self.indexer.raw_key_cache.prefix],
         )
-        if side_metadata.num_actual_tokens != num_tokens:
-            raise RuntimeError("QSA main and side metadata token counts disagree")
         selected, main_outputs = self.indexer(
             hidden_states,
             positions,
@@ -569,7 +593,7 @@ def _run_qsa(
             self.indexer.output_width,
         ):
             raise RuntimeError("QSA indexer returned an invalid selection shape")
-        impl = cast(Qwen4ExpQSAFlashAttentionImpl, self.impl)
+        impl = self.impl
         if main_outputs is None:
             assert key is not None and value is not None
             if impl.octave_fmt is not None:
@@ -585,19 +609,17 @@ def _run_qsa(
         else:
             query, _ = main_outputs
             # The gate is applied to every row outside this op; padded rows
-            # must stay finite so they keep the zero output forward_qsa gives.
+            # must stay finite so they keep the zero output forward gives.
             assert gate is not None
             gate[num_tokens:].zero_()
         assert query is not None
-        impl.forward_qsa(
+        impl.forward(
             self,
             query,
-            key,
-            value,
             self.kv_cache,
             main_metadata,
             output,
-            token_to_req=side_metadata.token_to_req,
+            token_to_req=indexer_metadata.token_to_req,
         )
 
     def forward(
@@ -733,8 +755,8 @@ def qwen4_exp_qsa_prepare_with_output(
 __all__ = [
     "QSAIndexer",
     "Qwen4ExpQSAAttention",
-    "Qwen4ExpQSAFlashAttentionBackend",
-    "Qwen4ExpQSAFlashAttentionImpl",
+    "Qwen4ExpQSABackend",
+    "Qwen4ExpQSAImpl",
     "qwen4_exp_qsa_prepare_with_output",
     "qwen4_exp_qsa_with_output",
 ]
diff --git a/vllm/models/qwen4_exp/amd/indexer_qsa.py b/vllm/models/qwen4_exp/amd/qsa_indexer.py
similarity index 88%
rename from vllm/models/qwen4_exp/amd/indexer_qsa.py
rename to vllm/models/qwen4_exp/amd/qsa_indexer.py
index 09cb39acf919..19c7780d6c3a 100644
--- a/vllm/models/qwen4_exp/amd/indexer_qsa.py
+++ b/vllm/models/qwen4_exp/amd/qsa_indexer.py
@@ -16,11 +16,12 @@
 from vllm.model_executor.layers.linear import ReplicatedLinear
 from vllm.model_executor.layers.quantization import QuantizationConfig
 from vllm.model_executor.layers.rotary_embedding.mrope import triton_mrope
+from vllm.platforms import current_platform
 
-from ..common.qsa_cache import (
+from ..common.qsa_indexer import (
     QSACompressedKeyCache,
-    QSAForwardMetadata,
-    QSAKeyStateCache,
+    QSAIndexerMetadata,
+    QSARawKeyCache,
     canonical_qsa_rope_positions,
 )
 from .ops.qsa_prepare import qsa_prepare, supports_fused_pre_indexer
@@ -110,7 +111,7 @@ def __init__(
         )
         self.prefix = prefix
         # MTP step 0 selects the target-aligned rows; later steps reuse them
-        # while continuing to update the QSA side cache.
+        # while continuing to update the QSA indexer cache.
         self.skip_topk = False
 
         self.index_qk_proj = ReplicatedLinear(
@@ -131,7 +132,20 @@ def __init__(
 
         cache_config = vllm_config.cache_config
         cache_prefix = f"{prefix}." if prefix else ""
-        self.raw_key_cache = QSAKeyStateCache(
+        self.indexer_kv_dtype = vllm_config.attention_config.resolve_indexer_kv_dtype(
+            default="bf16"
+        )
+        if self.indexer_kv_dtype == "fp8":
+            indexer_dtype = current_platform.fp8_dtype()
+        elif self.indexer_kv_dtype == "bf16":
+            indexer_dtype = torch.bfloat16
+        else:
+            raise NotImplementedError(
+                f"indexer_kv_dtype={self.indexer_kv_dtype!r} is not supported "
+                "by the Qwen4Exp QSA indexer (only 'bf16' or 'fp8')."
+            )
+        self.indexer_dtype = indexer_dtype
+        self.raw_key_cache = QSARawKeyCache(
             head_size=self.index_head_dim,
             dtype=torch.bfloat16,
             cache_rope_positions=vllm_config.model_config.uses_mrope,
@@ -142,7 +156,7 @@ def __init__(
         )
         self.compressed_key_cache = QSACompressedKeyCache(
             head_size=self.index_head_dim,
-            dtype=torch.bfloat16,
+            dtype=indexer_dtype,
             compress_ratio=self.compress_ratio,
             prefix=f"{cache_prefix}compressed_key_cache",
             cache_config=cache_config,
@@ -180,6 +194,7 @@ def project_qk(
             q.reshape(-1, self.index_head_dim),
         ).reshape_as(q)
         q = apply_qsa_rope(self.rotary_emb, positions, q)
+        q = q.to(self.indexer_dtype)
         return q, token_k.reshape(-1, 1, self.index_head_dim)
 
     def normalize_compressed_keys(
@@ -200,30 +215,24 @@ def normalize_compressed_keys(
 
     def _metadata(
         self,
-    ) -> tuple[QSAForwardMetadata, QSAForwardMetadata] | None:
+    ) -> tuple[QSAIndexerMetadata, QSAIndexerMetadata] | None:
         metadata = get_forward_context().attn_metadata
         if isinstance(metadata, list):
             metadata = metadata[0]
         if not isinstance(metadata, dict):
             return None
-        raw = cast(QSAForwardMetadata, metadata[self.raw_key_cache.prefix])
+        raw = cast(QSAIndexerMetadata, metadata[self.raw_key_cache.prefix])
         compressed = cast(
-            QSAForwardMetadata, metadata[self.compressed_key_cache.prefix]
+            QSAIndexerMetadata, metadata[self.compressed_key_cache.prefix]
         )
-        if raw.num_actual_tokens != compressed.num_actual_tokens:
-            raise RuntimeError("QSA side-cache metadata token counts disagree")
-        if not raw.logical_positions.is_cuda and (
-            not torch.equal(raw.logical_positions, compressed.logical_positions)
-        ):
-            raise RuntimeError("QSA side-cache metadata positions disagree")
         return raw, compressed
 
     def _update_and_compress(
         self,
         token_k: torch.Tensor,
         positions: torch.Tensor,
-        raw_metadata: QSAForwardMetadata,
-        compressed_metadata: QSAForwardMetadata,
+        raw_metadata: QSAIndexerMetadata,
+        compressed_metadata: QSAIndexerMetadata,
     ) -> None:
         num_tokens = raw_metadata.num_actual_tokens
         raw_key_cache = self.raw_key_cache.key_cache
@@ -273,8 +282,8 @@ def _fused_prepare(
         projected_q: torch.Tensor,
         raw_keys: torch.Tensor,
         positions: torch.Tensor,
-        raw_metadata: QSAForwardMetadata,
-        compressed_metadata: QSAForwardMetadata,
+        raw_metadata: QSAIndexerMetadata,
+        compressed_metadata: QSAIndexerMetadata,
         attn: Qwen4ExpQSAAttention,
         qkv: torch.Tensor,
         slot_mapping: torch.Tensor,
@@ -291,7 +300,19 @@ def _fused_prepare(
         """
         num_tokens = raw_metadata.num_actual_tokens
         raw_key_cache = self.raw_key_cache
-        q = projected_q.new_empty(num_tokens, self.index_n_heads, self.index_head_dim)
+        q = projected_q.new_empty(
+            num_tokens,
+            self.index_n_heads,
+            self.index_head_dim,
+            dtype=self.indexer_dtype,
+        )
+        main_kv_cache = attn.kv_cache.transpose(1, 2)
+        if attn.kv_cache_dtype in ("fp8", "fp8_e4m3"):
+            # Reinterpret the uint8 cache in the platform's native fp8 flavor,
+            # matching the C++ cache-update op and the attention backend's
+            # read-side views.
+            fp8_dtype = current_platform.fp8_dtype()
+            main_kv_cache = main_kv_cache.view(fp8_dtype)
         main_outputs = qsa_prepare(
             projected_q,
             raw_keys,
@@ -320,7 +341,7 @@ def _fused_prepare(
             main_q_norm_weight=attn.q_norm.weight,
             main_k_norm_weight=attn.k_norm.weight,
             main_eps=attn.q_norm.variance_epsilon,
-            main_kv_cache=attn.kv_cache.transpose(1, 2),
+            main_kv_cache=main_kv_cache,
             main_slot_mapping=slot_mapping[:num_tokens],
             main_k_scale=attn._k_scale_float,
             main_v_scale=attn._v_scale_float,
@@ -331,7 +352,7 @@ def _fused_prepare(
     def _select(
         self,
         q: torch.Tensor,
-        metadata: QSAForwardMetadata,
+        metadata: QSAIndexerMetadata,
         out: torch.Tensor | None,
     ) -> torch.Tensor:
         from .ops.qsa import qsa_select_paged_tokens
diff --git a/vllm/models/qwen4_exp/common/ngram_embedding.py b/vllm/models/qwen4_exp/common/ngram_embedding.py
index d5bd9f689270..8dbefee09d9b 100644
--- a/vllm/models/qwen4_exp/common/ngram_embedding.py
+++ b/vllm/models/qwen4_exp/common/ngram_embedding.py
@@ -159,6 +159,10 @@ def _select_embeddings(
             return embeddings
         return embeddings.narrow(0, slot_offset, local_num_tokens)
 
+    def bind_file_shards(self) -> None:
+        """Bind the checkpoint files of a disk-offloaded table after loading."""
+        return None
+
     @abstractmethod
     def start_prefetch(self, ngram_ids: torch.Tensor) -> None:
         """Start an asynchronous lookup when supported."""
@@ -282,6 +286,11 @@ def lookup_from_pinned(
             BLOCK_D=triton.next_power_of_2(row_bytes),
         )
 
+    def process_weights_after_loading(self, layer: nn.Module) -> None:
+        """Bind disk-offloaded tables once every checkpoint shard is loaded."""
+        if isinstance(layer, Qwen4ExpPLEEmbedding):
+            layer.bind_file_shards()
+
     @abstractmethod
     def dequantize(
         self,
@@ -371,6 +380,7 @@ def process_weights_after_loading(self, layer: nn.Module) -> None:
         sentinel = torch.finfo(torch.float32).min
         if torch.any(layer.weight_scale == sentinel):
             raise ValueError("FP8 PLE checkpoint is missing its global scale")
+        super().process_weights_after_loading(layer)
 
     def dequantize(
         self,
@@ -454,6 +464,7 @@ def process_weights_after_loading(self, layer: nn.Module) -> None:
         sentinel = torch.finfo(torch.float32).min
         if torch.any(layer.weight_scale_2 == sentinel):
             raise ValueError("NVFP4 PLE checkpoint is missing its global scale")
+        super().process_weights_after_loading(layer)
 
     def lookup_dtype(self, layer: nn.Module) -> torch.dtype:
         return layer.params_dtype
@@ -474,20 +485,41 @@ def embedding(self, layer: nn.Module, input_: torch.Tensor) -> torch.Tensor:
             dtype=layer.params_dtype,
             device=input_.device,
         )
+        # Device rows are already local, so every ID in range is owned.
+        self.decode_rows(
+            layer.weight,
+            layer.weight_scale,
+            layer.weight_scale_2,
+            ids,
+            output,
+            0,
+            layer.weight.shape[0],
+        )
+        return output
+
+    @staticmethod
+    def decode_rows(
+        weight: torch.Tensor,
+        weight_scale: torch.Tensor,
+        weight_scale_2: torch.Tensor,
+        ids: torch.Tensor,
+        output: torch.Tensor,
+        vocab_start: int,
+        vocab_end: int,
+    ) -> None:
+        """Decode the rows of `ids` in `[vocab_start, vocab_end)` into `output`."""
         if ids.numel():
-            # Device rows are already local, so every ID in range is owned.
             _lookup_nvfp4_ple_embedding_kernel[(ids.numel(),)](
-                layer.weight,
-                layer.weight_scale.view(torch.uint8),
-                layer.weight_scale_2,
+                weight,
+                weight_scale.view(torch.uint8),
+                weight_scale_2,
                 ids,
                 output,
-                layer.embedding_dim,
-                0,
-                layer.weight.shape[0],
-                BLOCK_D=triton.next_power_of_2(layer.embedding_dim),
+                output.shape[-1],
+                vocab_start,
+                vocab_end,
+                BLOCK_D=triton.next_power_of_2(output.shape[-1]),
             )
-        return output
 
     def lookup_from_pinned(
         self,
@@ -495,16 +527,14 @@ def lookup_from_pinned(
         ids: torch.Tensor,
         output: torch.Tensor,
     ) -> None:
-        _lookup_nvfp4_ple_embedding_kernel[(ids.numel(),)](
+        self.decode_rows(
             layer._uva_weight,
-            self._uva_weight_scale.view(torch.uint8),
+            self._uva_weight_scale,
             layer.weight_scale_2,
             ids,
             output,
-            layer.embedding_dim,
             layer.shard_indices.org_vocab_start_index,
             layer.shard_indices.org_vocab_end_index,
-            BLOCK_D=triton.next_power_of_2(layer.embedding_dim),
         )
 
     def dequantize(
diff --git a/vllm/models/qwen4_exp/common/ops/ple.py b/vllm/models/qwen4_exp/common/ops/ple.py
index 994b799ecdcd..915a42982f35 100644
--- a/vllm/models/qwen4_exp/common/ops/ple.py
+++ b/vllm/models/qwen4_exp/common/ops/ple.py
@@ -199,7 +199,7 @@ def _ple_gate_kernel(
     BLOCK_H: tl.constexpr,
     launch_pdl: tl.constexpr,
 ):
-    t = tl.program_id(0)
+    t = tl.program_id(0).to(tl.int64)
     s = tl.program_id(1)
     lanes = tl.arange(0, BLOCK_H)
     mask = lanes < H
diff --git a/vllm/models/qwen4_exp/common/qsa.py b/vllm/models/qwen4_exp/common/qsa.py
new file mode 100644
index 000000000000..ace69efcdb31
--- /dev/null
+++ b/vllm/models/qwen4_exp/common/qsa.py
@@ -0,0 +1,67 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Main-attention metadata shared by the QSA platform implementations."""
+
+from dataclasses import dataclass, replace
+from typing import ClassVar
+
+import torch
+
+from vllm.config import VllmConfig
+from vllm.v1.attention.backend import (
+    AttentionCGSupport,
+    AttentionMetadata,
+    AttentionMetadataBuilder,
+    CommonAttentionMetadata,
+)
+from vllm.v1.kv_cache_interface import AttentionSpec
+
+
+@dataclass
+class Qwen4ExpQSAMetadata(AttentionMetadata):
+    num_actual_tokens: int
+    max_query_len: int
+    block_table: torch.Tensor
+    slot_mapping: torch.Tensor
+
+
+class Qwen4ExpQSAMetadataBuilder(AttentionMetadataBuilder[Qwen4ExpQSAMetadata]):
+    _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
+    supports_update_block_table = True
+    supports_draft_decode_metadata_update = 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)
+
+    def build(
+        self,
+        common_prefix_len: int,
+        common_attn_metadata: CommonAttentionMetadata,
+        fast_build: bool = False,
+    ) -> Qwen4ExpQSAMetadata:
+        """Keep cache addresses shared with the current runner batch."""
+        return Qwen4ExpQSAMetadata(
+            num_actual_tokens=common_attn_metadata.num_actual_tokens,
+            max_query_len=common_attn_metadata.max_query_len,
+            block_table=common_attn_metadata.block_table_tensor,
+            slot_mapping=common_attn_metadata.slot_mapping,
+        )
+
+    def update_block_table(
+        self,
+        metadata: Qwen4ExpQSAMetadata,
+        blk_table: torch.Tensor,
+        slot_mapping: torch.Tensor,
+    ) -> Qwen4ExpQSAMetadata:
+        """Reuse batch bounds while replacing cache-group-specific addresses."""
+        return replace(metadata, block_table=blk_table, slot_mapping=slot_mapping)
+
+    def update_draft_decode_metadata(self, metadata: Qwen4ExpQSAMetadata) -> None:
+        """Decode batch bounds stay fixed and addresses reference runner buffers."""
+        pass
diff --git a/vllm/models/qwen4_exp/common/qsa_cache.py b/vllm/models/qwen4_exp/common/qsa_indexer.py
similarity index 94%
rename from vllm/models/qwen4_exp/common/qsa_cache.py
rename to vllm/models/qwen4_exp/common/qsa_indexer.py
index 5eba092eec25..165757738063 100644
--- a/vllm/models/qwen4_exp/common/qsa_cache.py
+++ b/vllm/models/qwen4_exp/common/qsa_indexer.py
@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Paged side-cache ownership and metadata for Qwen4Exp QSA.
+"""Paged indexer cache ownership and metadata for Qwen4Exp QSA.
 
 Each QSA layer keeps a fixed circular buffer of raw index keys (the
 compressor state) and one compressed key. MRoPE models pack exact three-axis
@@ -399,7 +399,7 @@ def build_qsa_metadata_triton(
     k_work_metadata_buffer: torch.Tensor | None = None,
     request_capacity: int | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-    """Build QSA side-cache and optional pre-indexer work metadata."""
+    """Build QSA indexer cache and optional pre-indexer work metadata."""
     num_tokens = common_attn_metadata.num_actual_tokens
     num_mapped_tokens = int(common_attn_metadata.query_start_loc_cpu[-1])
     token_to_req = token_to_req_buffer[:num_tokens]
@@ -573,17 +573,22 @@ def _build_qsa_metadata_torch(
 
 
 @dataclass
-class QSAForwardMetadata(AttentionMetadata):
-    """Common per-forward metadata for one QSA side cache."""
+class QSAIndexerMetadata(AttentionMetadata):
+    """Per-forward metadata for one raw or compressed QSA indexer cache."""
 
     block_table: torch.Tensor
     slot_mapping: torch.Tensor
     seq_lens: torch.Tensor
     query_start_loc: torch.Tensor
     token_to_req: torch.Tensor
-    logical_positions: torch.Tensor
-    visible_blocks: torch.Tensor
-    k_work_metadata: torch.Tensor
+    # Positions in the original uncompressed sequence, padded with -1.
+    logical_positions: torch.Tensor  # [num_tokens]
+    # Complete compression groups visible to each query.
+    visible_blocks: torch.Tensor  # [num_tokens]
+    # K compression work items: (request index, local task index), -1 for padding.
+    # Task 0 per request also saves the raw-K suffix, even without a complete group.
+    # Empty for raw or uncompressed caches.
+    k_work_metadata: torch.Tensor  # [num_k_work, 2]
     num_actual_tokens: int
     num_decodes: int
     num_decode_tokens: int
@@ -592,11 +597,12 @@ class QSAForwardMetadata(AttentionMetadata):
     max_query_len: int
     decode_query_len: int
     max_seq_len: int
+    # Key rows per physical cache block.
     storage_block_size: int
     compress_ratio: int
 
 
-class QSAMetadataBuilder(AttentionMetadataBuilder[QSAForwardMetadata]):
+class QSAIndexerMetadataBuilder(AttentionMetadataBuilder[QSAIndexerMetadata]):
     """Build QSA metadata from vLLM's cache-group-specific common metadata."""
 
     _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
@@ -653,7 +659,7 @@ def build(
         common_prefix_len: int,
         common_attn_metadata: CommonAttentionMetadata,
         fast_build: bool = False,
-    ) -> QSAForwardMetadata:
+    ) -> QSAIndexerMetadata:
         del common_prefix_len, fast_build
         num_tokens = common_attn_metadata.num_actual_tokens
         decode_threshold = self.reorder_batch_threshold
@@ -702,7 +708,7 @@ def build(
                 request_capacity=request_capacity,
             )
         )
-        return QSAForwardMetadata(
+        return QSAIndexerMetadata(
             block_table=common_attn_metadata.block_table_tensor,
             slot_mapping=slot_mapping,
             seq_lens=common_attn_metadata.seq_lens,
@@ -724,8 +730,8 @@ def build(
         )
 
 
-class QSAStateBackend(AttentionBackend):
-    """Key-only dummy backend for out-of-band QSA side-cache operations."""
+class QSAIndexerBackend(AttentionBackend):
+    """Key-only dummy backend for out-of-band QSA indexer cache operations."""
 
     supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16]
     # fp8 entries allow the optional e4m3 compressed indexer cache.
@@ -738,17 +744,17 @@ class QSAStateBackend(AttentionBackend):
 
     @staticmethod
     def get_name() -> str:
-        return "QWEN4_EXP_EXP_QSA_STATE"
+        return "QWEN4_EXP_QSA_INDEXER"
 
     @staticmethod
     def get_impl_cls():
         raise NotImplementedError(
-            "QSA state caches run out-of-band and have no attention impl"
+            "QSA indexer caches run out-of-band and have no attention impl"
         )
 
     @staticmethod
-    def get_builder_cls() -> type[QSAMetadataBuilder]:
-        return QSAMetadataBuilder
+    def get_builder_cls() -> type[QSAIndexerMetadataBuilder]:
+        return QSAIndexerMetadataBuilder
 
     @classmethod
     def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...]:
@@ -756,7 +762,7 @@ def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...]:
         return (KVCacheLayout.BLNHC, KVCacheLayout.BLHNC)
 
 
-class _QSAStateCache(nn.Module, AttentionLayerBase):
+class _QSAIndexerCache(nn.Module, AttentionLayerBase):
     supports_dcp = False
 
     def __init__(
@@ -795,20 +801,20 @@ def forward(self) -> None: ...
     def bind_kv_cache(self, kv_cache: torch.Tensor) -> None:
         """Adapt the unified [B, H, N, C] view to QSA's [B, N, H, C]."""
         if kv_cache.ndim != 4 or kv_cache.shape[1] != 1:
-            raise ValueError("QSA state cache must be [blocks, 1, states, width]")
+            raise ValueError("QSA indexer cache must be [blocks, 1, states, width]")
         if kv_cache.dtype != self.dtype or kv_cache.shape[3] != self.head_size:
             raise ValueError(
-                f"QSA state cache does not match its spec "
+                f"QSA indexer cache does not match its spec "
                 f"(dtype {kv_cache.dtype} != {self.dtype} or width "
                 f"{kv_cache.shape[3]} != {self.head_size})"
             )
         super().bind_kv_cache(kv_cache.transpose(1, 2))
 
     def get_attn_backend(self) -> type[AttentionBackend]:
-        return QSAStateBackend
+        return QSAIndexerBackend
 
 
-class QSAKeyStateCache(_QSAStateCache):
+class QSARawKeyCache(_QSAIndexerCache):
     """Raw BF16 key, optionally followed by exact int64 MRoPE positions."""
 
     _BF16_PER_INT64 = 4
@@ -860,7 +866,7 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
         )
 
 
-class QSACompressedKeyCache(_QSAStateCache):
+class QSACompressedKeyCache(_QSAIndexerCache):
     """Normed, group-first-RoPE key at one row per complete group."""
 
     def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
@@ -876,11 +882,21 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
 
 __all__ = [
     "QSACompressedKeyCache",
-    "QSAForwardMetadata",
-    "QSAKeyStateCache",
-    "QSAMetadataBuilder",
-    "QSAStateBackend",
+    "QSAIndexerMetadata",
+    "QSARawKeyCache",
+    "QSAIndexerMetadataBuilder",
+    "QSAIndexerBackend",
     "canonical_qsa_rope_positions",
     "circular_qsa_slot_mapping",
     "compressed_qsa_slot_mapping",
+    "qsa_kv_cache_dtype",
 ]
+
+
+def qsa_kv_cache_dtype(cache_config: CacheConfig, prefix: str) -> CacheDType | str:
+    """The layer's KV cache dtype, honoring ``--kv-cache-dtype-skip-layers``."""
+    from vllm.model_executor.models.utils import extract_layer_index
+
+    if str(extract_layer_index(prefix)) in cache_config.kv_cache_dtype_skip_layers:
+        return "auto"
+    return cache_config.cache_dtype
diff --git a/vllm/models/qwen4_exp/nvidia/low_latency_gemm.py b/vllm/models/qwen4_exp/nvidia/low_latency_gemm.py
index ea56750702fc..156faeb99ebf 100644
--- a/vllm/models/qwen4_exp/nvidia/low_latency_gemm.py
+++ b/vllm/models/qwen4_exp/nvidia/low_latency_gemm.py
@@ -63,6 +63,38 @@
         1: SkinnyGemmConfig(1, 64, 4, k_unroll=2),
         2: SkinnyGemmConfig(2, 32, 4, k_unroll=2),
     },
+    # GDN and QSA output projections, TP=1 (DGX Station GB300).
+    (2560, 6144): {
+        1: SkinnyGemmConfig(1, 128, 2, k_unroll=2, vector_width=4, static_k=6144),
+        2: SkinnyGemmConfig(2, 128, 2, static_k=6144),
+    },
+    # GDN fused B/A projection, TP=1.
+    (96, 2560): {
+        1: SkinnyGemmConfig(1, 128, 1, k_unroll=4, vector_width=4, static_k=2560),
+        2: SkinnyGemmConfig(2, 128, 1, k_unroll=4, vector_width=4, static_k=2560),
+        4: SkinnyGemmConfig(4, 160, 1, k_unroll=2, static_k=2560),
+    },
+    # QSA fused QKV/gate + replicated indexer Q/K, TP=1.
+    (13952, 2560): {
+        1: SkinnyGemmConfig(1, 128, 2, k_unroll=6, vector_width=2, static_k=2560),
+    },
+    # Shared-expert fused gate/up projection, TP=1.
+    (1280, 2560): {
+        1: SkinnyGemmConfig(1, 128, 2, k_unroll=2, vector_width=4, static_k=2560),
+        2: SkinnyGemmConfig(2, 64, 2, k_unroll=2, static_k=2560),
+        4: SkinnyGemmConfig(4, 64, 2, k_unroll=4, vector_width=4, static_k=2560),
+    },
+    # Shared-expert down projection, TP=1.
+    (2560, 640): {
+        1: SkinnyGemmConfig(1, 160, 4, k_unroll=2, vector_width=4),
+        2: SkinnyGemmConfig(2, 160, 4, vector_width=4),
+    },
+    # HC final down projection, TP=1.
+    (320, 10240): {
+        1: SkinnyGemmConfig(1, 256, 1, k_unroll=2, static_k=10240),
+        2: SkinnyGemmConfig(2, 256, 1, k_unroll=2, static_k=10240),
+        4: SkinnyGemmConfig(4, 256, 1, k_unroll=2),
+    },
 }
 
 # H200 plans selected by exhaustive CUDA graph replay measurements over
@@ -118,7 +150,7 @@
 }
 
 
-# B200 plans; the winning configs differ from the B300 table.
+# B200 plans for TP=4 and TP=1; the winning configs differ from the B300 table.
 QWEN4_EXP_SM100_GEMM_PLANS: dict[tuple[int, int], dict[int, SkinnyGemmConfig]] = {
     # GDN fused B/A projection, TP=4.
     (24, 2560): {
@@ -165,6 +197,28 @@
     (62080, 2560): {
         1: SkinnyGemmConfig(1, 128, 2, k_unroll=4, vector_width=4),
     },
+    # GDN fused B/A projection, TP=1.
+    (96, 2560): {
+        1: SkinnyGemmConfig(1, 128, 4, k_unroll=2, vector_width=2, static_k=2560),
+    },
+    # Router, TP=1.
+    (512, 2560): {
+        1: SkinnyGemmConfig(1, 256, 1, k_unroll=4, vector_width=2, static_k=2560),
+    },
+    # Shared-expert fused gate/up projection, TP=1.
+    (1280, 2560): {
+        1: SkinnyGemmConfig(1, 32, 1, k_unroll=2, static_k=2560),
+    },
+    # Shared-expert down projection, TP=1.
+    (2560, 640): {
+        1: SkinnyGemmConfig(1, 160, 4, k_unroll=2, vector_width=4),
+        2: SkinnyGemmConfig(2, 160, 4, vector_width=4),
+    },
+    # GDN and QSA output projections, TP=1.
+    (2560, 6144): {
+        1: SkinnyGemmConfig(1, 256, 2, k_unroll=2, static_k=6144),
+        2: SkinnyGemmConfig(2, 128, 2, static_k=6144),
+    },
 }
 
 
diff --git a/vllm/models/qwen4_exp/nvidia/model.py b/vllm/models/qwen4_exp/nvidia/model.py
index ed9aaa927733..08ad38361e14 100644
--- a/vllm/models/qwen4_exp/nvidia/model.py
+++ b/vllm/models/qwen4_exp/nvidia/model.py
@@ -13,9 +13,6 @@
 from vllm.config import VllmConfig
 from vllm.distributed import get_pp_group
 from vllm.forward_context import get_forward_context, is_forward_context_available
-from vllm.model_executor.layers.fused_moe.utils import (
-    is_model_fused_shared_expert_compatible,
-)
 from vllm.model_executor.layers.fusion.mm_input_norm import build_mm_input_norm
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
 from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import (
@@ -65,6 +62,7 @@
     WeightsMapper,
     _merge_multimodal_embeddings,
     extract_layer_index,
+    get_fused_shared_expert_layers,
     make_empty_intermediate_tensors_factory,
     make_layers,
     maybe_fuse_shared_experts,
@@ -446,10 +444,10 @@ def get_layer(prefix: str) -> Qwen4ExpDecoderLayer:
         self.start_layer, self.end_layer, self.layers = make_layers(
             config.num_hidden_layers, get_layer, prefix=f"{prefix}.layers"
         )
-        self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible(
-            self.layers,
+        self.fused_shared_expert_layers = get_fused_shared_expert_layers(
+            self,
             Qwen4ExpSparseMoeBlock,
-            "mlp",
+            enable_eplb=vllm_config.parallel_config.enable_eplb,
         )
         intermediate_size = config.hidden_size * config.hc_count
         self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
@@ -656,7 +654,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         )
         weights = maybe_fuse_shared_experts(
             weights,
-            enabled=self.is_fused_shared_expert_enabled,
+            fused_shared_expert_layers=self.fused_shared_expert_layers,
             n_routed_experts=getattr(self.config, "num_experts", 0) or 0,
             n_shared_experts=1,
             ckpt_prefix="mlp.shared_expert",
diff --git a/vllm/models/qwen4_exp/nvidia/model_state.py b/vllm/models/qwen4_exp/nvidia/model_state.py
index 4d43cddd08f0..520d258b2191 100644
--- a/vllm/models/qwen4_exp/nvidia/model_state.py
+++ b/vllm/models/qwen4_exp/nvidia/model_state.py
@@ -13,6 +13,12 @@
 from vllm.v1.worker.gpu.model_states.mamba_hybrid import MambaHybridModelState
 from vllm.v1.worker.gpu.states import RequestState
 
+from .ngram_embedding import (
+    Qwen4ExpNGramEmbedding,
+    Qwen4ExpPLEDiskOffloadEmbedding,
+    stage_checkpoint_rows,
+)
+
 
 class Qwen4ExpModelState(MambaHybridModelState):
     """Add rollback-safe PLE n-gram context to the model inputs."""
@@ -55,6 +61,12 @@ def __init__(
             dtype=torch.int32,
             device=self.device,
         )
+        self.disk_offload_modules = [
+            (module, module.ngram_embedding)
+            for module in model.modules()
+            if isinstance(module, Qwen4ExpNGramEmbedding)
+            and isinstance(module.ngram_embedding, Qwen4ExpPLEDiskOffloadEmbedding)
+        ]
 
     def _prepare_ngram_context(
         self,
@@ -98,9 +110,17 @@ def prepare_inputs(
         query_start_loc[: num_reqs_padded + 1].copy_(input_batch.query_start_loc)
         # Represent unused capacity as trailing zero-length requests.
         query_start_loc[num_reqs_padded + 1 :].copy_(input_batch.query_start_loc[-1])
+        ngram_context = self._prepare_ngram_context(input_batch, req_states)
+        if self.disk_offload_modules:
+            stage_checkpoint_rows(
+                self.disk_offload_modules,
+                input_batch.input_ids,
+                query_start_loc,
+                ngram_context,
+            )
         model_inputs.update(
             query_start_loc=query_start_loc,
-            ngram_context=self._prepare_ngram_context(input_batch, req_states),
+            ngram_context=ngram_context,
         )
         return model_inputs
 
diff --git a/vllm/models/qwen4_exp/nvidia/mtp.py b/vllm/models/qwen4_exp/nvidia/mtp.py
index fb15d076467a..30401ca56dcb 100644
--- a/vllm/models/qwen4_exp/nvidia/mtp.py
+++ b/vllm/models/qwen4_exp/nvidia/mtp.py
@@ -24,9 +24,6 @@
 from vllm.config import VllmConfig, replace, set_current_vllm_config
 from vllm.distributed import get_pp_group
 from vllm.forward_context import get_forward_context, is_forward_context_available
-from vllm.model_executor.layers.fused_moe.utils import (
-    is_model_fused_shared_expert_compatible,
-)
 from vllm.model_executor.layers.layernorm import GemmaRMSNorm
 from vllm.model_executor.layers.linear import ColumnParallelLinear
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
@@ -42,6 +39,7 @@
     PPMissingLayer,
     WeightsMapper,
     get_draft_quant_config,
+    get_fused_shared_expert_layers,
     make_empty_intermediate_tensors_factory,
     maybe_fuse_shared_experts,
     maybe_prefix,
@@ -228,10 +226,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
                 )
                 for idx in range(self.num_mtp_layers)
             )
-        self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible(
-            self.layers,
+        self.fused_shared_expert_layers = get_fused_shared_expert_layers(
+            self,
             Qwen4ExpSparseMoeBlock,
-            "mlp",
+            enable_eplb=vllm_config.parallel_config.enable_eplb,
         )
 
         self.pre_fc_norm_embedding = GemmaRMSNorm(
@@ -384,7 +382,7 @@ def forward(
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         weights = maybe_fuse_shared_experts(
             weights,
-            enabled=self.is_fused_shared_expert_enabled,
+            fused_shared_expert_layers=self.fused_shared_expert_layers,
             n_routed_experts=getattr(self.config, "num_experts", 0) or 0,
             n_shared_experts=1,
             ckpt_prefix="mlp.shared_expert",
diff --git a/vllm/models/qwen4_exp/nvidia/ngram_embedding.py b/vllm/models/qwen4_exp/nvidia/ngram_embedding.py
index 7b0d9ed8c2f7..9dfdf2f4d180 100644
--- a/vllm/models/qwen4_exp/nvidia/ngram_embedding.py
+++ b/vllm/models/qwen4_exp/nvidia/ngram_embedding.py
@@ -1,7 +1,9 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Qwen4Exp n-gram embeddings with device and pinned-host storage."""
+"""Qwen4Exp n-gram embeddings with device, pinned-host and disk storage."""
 
+import mmap
+import os
 from collections.abc import Iterable
 
 import torch
@@ -17,6 +19,7 @@
 from vllm.model_executor.layers.quantization.utils.fp8_utils import is_fp8
 from vllm.model_executor.models.utils import AutoWeightsLoader
 from vllm.models.common.ops.sequence_parallel import sp_reduce_scatter, sp_shard
+from vllm.utils.gpu_sync_debug import gpu_sync_allowed
 from vllm.utils.math_utils import round_up
 
 from ..common.ngram_embedding import (
@@ -36,6 +39,7 @@
     "Qwen4ExpPLEDeviceEmbedding",
     "Qwen4ExpPLEEmbedding",
     "Qwen4ExpPLEEmbeddingMethod",
+    "Qwen4ExpPLEDiskOffloadEmbedding",
     "Qwen4ExpPLEFp8EmbeddingMethod",
     "Qwen4ExpPLENvFp4EmbeddingMethod",
     "Qwen4ExpPLEPinnedHostEmbedding",
@@ -44,6 +48,219 @@
 ]
 
 
+def _maps_entry(ptr: int) -> tuple[str, int]:
+    """Return the path and file offset of `ptr` from `/proc/self/maps`."""
+    with open("/proc/self/maps") as maps:
+        for line in maps:
+            fields = line.split(maxsplit=5)
+            low, high = (int(bound, 16) for bound in fields[0].split("-"))
+            if low <= ptr < high:
+                path = fields[5].rstrip("\n") if len(fields) > 5 else ""
+                return path, ptr - low + int(fields[2], 16)
+    return "", 0
+
+
+class Qwen4ExpPLEDiskOffloadEmbedding(Qwen4ExpPLEEmbedding):
+    """PLE rows read on the host from the checkpoint shard files.
+
+    No table is allocated. Each step, `stage_checkpoint_rows` reads the rows
+    of that step's n-gram ids from the safetensors files and copies them to
+    fixed device buffers, which the forward reads, so CUDA graphs replay it.
+    Every ETP rank stages full rows, so no ETP reduction is needed.
+    """
+
+    def __init__(
+        self,
+        num_embeddings: int,
+        embedding_dim: int,
+        *,
+        params_dtype: torch.dtype,
+        padding_size: int,
+        prefix: str,
+        embedding_method: Qwen4ExpPLEEmbeddingMethod,
+        num_ngram_heads: int = 1,
+        max_total_tokens: int = 0,
+        data_parallel_rank: int = 0,
+        reduce_results: bool = True,
+    ) -> None:
+        super().__init__(
+            num_embeddings,
+            embedding_dim,
+            params_dtype=params_dtype,
+            padding_size=padding_size,
+            prefix=prefix,
+            embedding_method=embedding_method,
+            data_parallel_rank=data_parallel_rank,
+            reduce_results=reduce_results,
+        )
+        self._dummy = get_current_vllm_config().load_config.load_format == "dummy"
+        # NVFP4 rows are packed E2M1 bytes plus a separate plane of block scales.
+        self._nvfp4 = isinstance(embedding_method, Qwen4ExpPLENvFp4EmbeddingMethod)
+        device, pin = self.weight.device, self.weight.is_cuda
+        rows = (max_total_tokens, num_ngram_heads)
+        self._host_ids = torch.empty(
+            rows, dtype=torch.long, device="cpu", pin_memory=pin
+        )
+        # Per plane: device staging, pinned host rows (written only after the
+        # step's stream sync, so after the previous H2D) and distinct-row bytes.
+        self._planes: dict[str, tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = {}
+        for name in ("weight", "weight_scale") if self._nvfp4 else ("weight",):
+            plane = getattr(self, name)
+            shape, dtype = (*rows, plane.shape[1]), plane.dtype
+            host = torch.empty(shape, dtype=dtype, device="cpu", pin_memory=pin)
+            self._planes[name] = (
+                torch.zeros(shape, dtype=dtype, device=device),
+                host,
+                torch.empty_like(host).view(torch.uint8).flatten(0, 1),
+            )
+        if self._nvfp4:
+            self._output = torch.zeros(
+                (*rows, embedding_dim), dtype=params_dtype, device=device
+            )
+            self._row_ids = torch.arange(rows[0] * rows[1], device=device)
+        self._shards: dict[str, dict[int, torch.Tensor]] = {n: {} for n in self._planes}
+        self._shard_size = 0
+        self._bound = False
+        self._fds: dict[str, int] = {}
+        self._sources: dict[str, tuple[torch.Tensor, torch.Tensor]] = {}
+
+    def allocate_embedding_weight(
+        self,
+        num_embeddings: int,
+        embedding_dim: int,
+        dtype: torch.dtype,
+    ) -> torch.Tensor:
+        """Allocate no rows; they stay in the checkpoint files."""
+        return torch.empty(0, embedding_dim, dtype=dtype)
+
+    def accept_checkpoint_shard(
+        self,
+        shard_index: int,
+        loaded_weight: torch.Tensor,
+        shard_size: int,
+        name: str,
+    ) -> None:
+        """Keep a checkpoint shard as a zero-copy byte view of its storage."""
+        if self._bound:
+            raise RuntimeError("PLE disk offload does not support weight reload")
+        dtype = getattr(self, name).dtype
+        if loaded_weight.dtype != dtype:
+            raise ValueError(
+                f"PLE shard {name} dtype {loaded_weight.dtype} must match the "
+                f"embedding dtype {dtype}; disk offload serves rows as stored"
+            )
+        if shard_index in self._shards[name]:
+            raise ValueError(f"Duplicate PLE embedding shard {shard_index} {name}")
+        self._shards[name][shard_index] = loaded_weight.view(torch.uint8)
+        self._shard_size = shard_size
+
+    def bind_file_shards(self) -> None:
+        """Record each shard's file and offset and release the loader's views."""
+        if self._bound or self._dummy:
+            self._bound = True
+            return
+        for name, shards in self._shards.items():
+            rows = sum(shard.shape[0] for shard in shards.values())
+            if rows != self.org_vocab_size:
+                raise ValueError(
+                    f"PLE {name} shards cover {rows} of {self.org_vocab_size} rows"
+                )
+            fds, bases = torch.zeros(2, len(shards)).long()
+            for index, shard in shards.items():
+                # The loader hands over mmap views, so find their files here.
+                path, offset = _maps_entry(shard.data_ptr())
+                if not path or not os.path.isfile(path):
+                    raise ValueError(
+                        f"PLE embedding shard {index} {name} is not a "
+                        f"file-backed view ({path or 'anonymous'})"
+                    )
+                if path not in self._fds:
+                    self._fds[path] = os.open(path, os.O_RDONLY)
+                fds[index], bases[index] = self._fds[path], offset
+            self._sources[name] = (fds, bases)
+        self._shards.clear()
+        self._bound = True
+
+    def copy_ids_to_host(self, ids: torch.Tensor) -> None:
+        self._host_ids[: ids.shape[0]].copy_(ids, non_blocking=True)
+
+    def stage_rows(self, num_tokens: int) -> None:
+        """Read the rows of `_host_ids[:num_tokens]` and copy them to staging."""
+        ids = self._host_ids[:num_tokens].flatten()
+        if not self._fds:
+            for _, host, _ in self._planes.values():
+                host[:num_tokens].view(torch.uint8).zero_()
+        else:
+            ids, inverse = ids.unique(return_inverse=True)
+            if ids.numel() and not 0 <= ids[0] <= ids[-1] < self.org_vocab_size:
+                raise IndexError(f"PLE id out of range for {self.org_vocab_size} rows")
+            shard = ids // self._shard_size
+            local = ids - shard * self._shard_size
+            for name, (_, host, rows) in self._planes.items():
+                out = host[:num_tokens].flatten(0, 1).view(torch.uint8)
+                self._read_rows(name, shard, local, rows[: ids.numel()])
+                torch.index_select(rows[: ids.numel()], 0, inverse, out=out)
+        for staging, host, _ in self._planes.values():
+            staging[:num_tokens].copy_(host[:num_tokens], non_blocking=True)
+
+    def _read_rows(
+        self,
+        name: str,
+        shard: torch.Tensor,
+        local: torch.Tensor,
+        rows: torch.Tensor,
+    ) -> None:
+        """Read one plane's rows for sorted distinct ids into `rows`."""
+        fds, bases = self._sources[name]
+        page, row_bytes = mmap.PAGESIZE, rows.shape[1]
+        start = bases[shard] + local * row_bytes
+        first, last = start // page, (start + row_bytes - 1) // page
+        run = torch.ones_like(first, dtype=torch.bool)
+        run[1:] = (first[1:] > last[:-1] + 1) | (shard[1:] != shard[:-1])
+        # Queue every page run before the serial reads.
+        for fd, lo, hi in zip(
+            fds[shard[run]].tolist(),
+            first[run].tolist(),
+            last[run.roll(-1)].tolist(),
+        ):
+            os.posix_fadvise(
+                fd, lo * page, (hi - lo + 1) * page, os.POSIX_FADV_WILLNEED
+            )
+        for row, fd, offset in zip(rows.numpy(), fds[shard].tolist(), start.tolist()):
+            if os.preadv(fd, [row], offset) != row_bytes:
+                raise ValueError(
+                    f"PLE shard file {os.readlink(f'/proc/self/fd/{fd}')} is "
+                    f"truncated at offset {offset}"
+                )
+
+    def staged_rows(self, num_tokens: int) -> torch.Tensor:
+        """Return this step's staged rows, decoded if they are NVFP4."""
+        weight, *scale = (
+            staging[:num_tokens] for staging, _, _ in self._planes.values()
+        )
+        if not self._nvfp4:
+            return weight.flatten(-2)
+        num_rows = weight.shape[0] * weight.shape[1]
+        output = self._output[:num_tokens]
+        # Staged rows are compact, so the i-th staged row has id i.
+        Qwen4ExpPLENvFp4EmbeddingMethod.decode_rows(
+            weight,
+            scale[0],
+            self.weight_scale_2,
+            self._row_ids[:num_rows],
+            output,
+            0,
+            num_rows,
+        )
+        return output.flatten(-2)
+
+    def start_prefetch(self, ngram_ids: torch.Tensor) -> None:
+        return None
+
+    def finalize_prefetch(self, total_tokens: int) -> torch.Tensor:
+        raise RuntimeError("PLE disk offload does not support prefetch")
+
+
 class Qwen4ExpNGramEmbedding(nn.Module):
     _MASK64 = (1 << 64) - 1
     _SPLITMIX_GAMMA = 0x9E3779B97F4A7C15
@@ -215,14 +432,18 @@ def __init__(
             params_dtype = torch.get_default_dtype()
         engram_config = get_current_vllm_config().engram_config
         embedding_cls = (
-            Qwen4ExpPLEPinnedHostEmbedding
+            Qwen4ExpPLEDiskOffloadEmbedding
+            if engram_config is not None and engram_config.disk_offload
+            else Qwen4ExpPLEPinnedHostEmbedding
             if engram_config is not None and engram_config.cpu_offload
             else Qwen4ExpPLEDeviceEmbedding
         )
         # A TP reduce-scatter covers all vocabulary owners only when ETP=TP.
+        # Disk-offload ranks hold full rows, so they only shard tokens.
         self.use_reduce_scatter = (
             self.use_sequence_parallel
             and get_etp_group().world_size == get_tp_group().world_size
+            and embedding_cls is not Qwen4ExpPLEDiskOffloadEmbedding
         )
         self.ngram_embedding = embedding_cls(
             padded_vocab_size,
@@ -388,6 +609,8 @@ def forward(
         if embedding.supports_prefetch:
             # Prefetch contains full tokens even when hidden states are sharded.
             embeddings = embedding.finalize_prefetch(total_tokens=input_ids.shape[0])
+        elif isinstance(embedding, Qwen4ExpPLEDiskOffloadEmbedding):
+            embeddings = embedding.staged_rows(input_ids.shape[0])
         else:
             ngram_ids = self.compute_ngram_ids(
                 input_ids, query_start_loc, ngram_context
@@ -480,11 +703,16 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
                         f"NVFP4 PLE shard {shard_index} {suffix} requires "
                         f"{parameter.dtype}, got {loaded_weight.dtype}"
                     )
-                parameter.weight_loader(
-                    parameter,
-                    loaded_weight,
-                    checkpoint_start=checkpoint_start,
-                )
+                if isinstance(embedding, Qwen4ExpPLEDiskOffloadEmbedding):
+                    embedding.accept_checkpoint_shard(
+                        shard_index, loaded_weight, shard_size, suffix
+                    )
+                else:
+                    parameter.weight_loader(
+                        parameter,
+                        loaded_weight,
+                        checkpoint_start=checkpoint_start,
+                    )
                 loaded.add(f"ngram_embedding.{suffix}")
                 continue
             if nvfp4_method is not None and name in (
@@ -500,3 +728,25 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         if regular_weights:
             loaded.update(AutoWeightsLoader(self).load_weights(regular_weights))
         return loaded
+
+
+def stage_checkpoint_rows(
+    modules: list[tuple[Qwen4ExpNGramEmbedding, Qwen4ExpPLEDiskOffloadEmbedding]],
+    input_ids: torch.Tensor,
+    query_start_loc: torch.Tensor,
+    ngram_context: torch.Tensor,
+) -> None:
+    """Stage disk-offload PLE rows for one step behind one stream sync.
+
+    The ids hash this step's tokens, which the previous step samples on the
+    GPU, so the host read has to wait for them.
+    """
+    num_tokens = input_ids.shape[0]
+    for module, embedding in modules:
+        ids = module.compute_ngram_ids(input_ids, query_start_loc, ngram_context)
+        embedding.copy_ids_to_host(ids)
+    if input_ids.device.type != "cpu":
+        with gpu_sync_allowed():
+            torch.accelerator.current_stream().synchronize()
+    for _, embedding in modules:
+        embedding.stage_rows(num_tokens)
diff --git a/vllm/models/qwen4_exp/nvidia/ops/hc.py b/vllm/models/qwen4_exp/nvidia/ops/hc.py
index 3dee9db436b7..dfe95d3722b3 100644
--- a/vllm/models/qwen4_exp/nvidia/ops/hc.py
+++ b/vllm/models/qwen4_exp/nvidia/ops/hc.py
@@ -26,7 +26,7 @@ def _grouped_gemma_rmsnorm_kernel(
 
     pid = tl.program_id(0)
     group_id = pid % NUM_GROUPS
-    row = pid // NUM_GROUPS
+    row = (pid // NUM_GROUPS).to(tl.int64)
 
     offs_g = tl.arange(0, BLOCK_SIZE)
     offsets = group_id * GROUP_DIM + offs_g
@@ -89,7 +89,7 @@ def _hc_silu_kernel(
 ) -> None:
     BLOCK_SIZE: tl.constexpr = triton.next_power_of_2(DIM)
 
-    row = tl.program_id(0)
+    row = tl.program_id(0).to(tl.int64)
     offs = tl.arange(0, BLOCK_SIZE)
     mask = offs < DIM
 
@@ -136,7 +136,7 @@ def _hc_gate_mix_kernel(
 ) -> None:
     HC_DIM: tl.constexpr = DIM // HC
 
-    row = tl.program_id(0)
+    row = tl.program_id(0).to(tl.int64)
     tile_id = tl.program_id(1)
     offs_inner = tile_id * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
     mask = offs_inner < HC_DIM
@@ -201,7 +201,7 @@ def _hc_combine_kernel(
 ) -> None:
     HC_PAD: tl.constexpr = triton.next_power_of_2(HC)
 
-    row = tl.program_id(0)
+    row = tl.program_id(0).to(tl.int64)
     tile_id = tl.program_id(1)
 
     offs_inner = tile_id * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
@@ -293,7 +293,7 @@ def _hc_combine_norm_kernel(
     NUM_TILES_PAD: tl.constexpr = triton.next_power_of_2(NUM_TILES)
 
     pid = tl.program_id(0)
-    row = pid // HC
+    row = (pid // HC).to(tl.int64)
     stream = pid % HC
     offs_hc = tl.arange(0, HC_PAD)
     mask_hc = offs_hc < HC
diff --git a/vllm/models/qwen4_exp/nvidia/qsa.py b/vllm/models/qwen4_exp/nvidia/qsa.py
index 36e4864570ee..6ca0a05cdf69 100644
--- a/vllm/models/qwen4_exp/nvidia/qsa.py
+++ b/vllm/models/qwen4_exp/nvidia/qsa.py
@@ -12,9 +12,10 @@
 from transformers import Qwen4ExpTextConfig
 
 import vllm.envs as envs
+from vllm import _custom_ops as ops
 from vllm.compilation.breakable_cudagraph import eager_break_during_capture
 from vllm.config import VllmConfig, get_current_vllm_config_or_none
-from vllm.config.cache import CacheConfig, CacheDType
+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.model_executor.layers.attention.attention import (
@@ -34,37 +35,29 @@
 from vllm.model_executor.models.utils import (
     AutoWeightsLoader,
     WeightsMapper,
-    extract_layer_index,
 )
 from vllm.platforms import current_platform
-from vllm.platforms.interface import DeviceCapability
 from vllm.utils.math_utils import cdiv
 from vllm.utils.torch_utils import (
     kv_cache_dtype_str_to_dtype,
 )
 from vllm.v1.attention.backend import (
     AttentionBackend,
-    AttentionCGSupport,
+    AttentionImplBase,
     AttentionType,
     MultipleOf,
 )
-from vllm.v1.attention.backends.fa_utils import is_flash_attn_varlen_func_available
-from vllm.v1.attention.backends.flash_attn import (
-    FlashAttentionBackend,
-    FlashAttentionImpl,
-    FlashAttentionMetadata,
-    FlashAttentionMetadataBuilder,
-)
 from vllm.v1.kv_cache_interface import (
     FullAttentionSpec,
     KVCacheSpec,
     get_kv_quant_mode,
 )
 
-from ..common.qsa_cache import QSAForwardMetadata
+from ..common.qsa import Qwen4ExpQSAMetadata, Qwen4ExpQSAMetadataBuilder
+from ..common.qsa_indexer import QSAIndexerMetadata, qsa_kv_cache_dtype
 from . import model
-from .indexer_qsa import QSAIndexer
 from .ops.qsa_indexer import _TOPK_WORKSPACE_BYTES
+from .qsa_indexer import QSAIndexer
 
 
 class Qwen4ExpQSAQKVIndexerLinear(MergedColumnParallelLinear):
@@ -133,27 +126,13 @@ def remap_shards():
         return super().load_weights(remap_shards())
 
 
-def qsa_kv_cache_dtype(cache_config: CacheConfig, prefix: str) -> CacheDType | str:
-    """The layer's KV cache dtype, honoring ``--kv-cache-dtype-skip-layers``."""
-    if str(extract_layer_index(prefix)) in cache_config.kv_cache_dtype_skip_layers:
-        return "auto"
-    return cache_config.cache_dtype
-
-
-class Qwen4ExpQSAMetadataBuilder(FlashAttentionMetadataBuilder):
-    """Flash metadata supporting uniform decode and target-verify graphs."""
-
-    _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
-
-
-class Qwen4ExpQSAFlashAttentionBackend(FlashAttentionBackend):
+class Qwen4ExpQSABackend(AttentionBackend):
     """FullAttentionSpec backend used by the merged QSA owner."""
 
     supported_dtypes: ClassVar[list[torch.dtype]] = [torch.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).
+    # Triton kernel.
     supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
         "auto",
         "bfloat16",
@@ -161,27 +140,8 @@ class Qwen4ExpQSAFlashAttentionBackend(FlashAttentionBackend):
         "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
+    # The owner writes K/V before calling the sparse attention implementation.
+    forward_includes_kv_cache_update: bool = False
 
     @staticmethod
     def get_name() -> str:
@@ -193,8 +153,8 @@ def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleO
         return [MultipleOf(16)]
 
     @staticmethod
-    def get_impl_cls() -> type[Qwen4ExpQSAFlashAttentionImpl]:
-        return Qwen4ExpQSAFlashAttentionImpl
+    def get_impl_cls() -> type[Qwen4ExpQSAImpl]:
+        return Qwen4ExpQSAImpl
 
     @staticmethod
     def get_builder_cls() -> type[Qwen4ExpQSAMetadataBuilder]:
@@ -209,84 +169,65 @@ def supports_kv_connector(cls) -> bool:
         return False
 
 
-class Qwen4ExpQSAFlashAttentionImpl(FlashAttentionImpl):
+class Qwen4ExpQSAImpl(AttentionImplBase[Qwen4ExpQSAMetadata]):
     """Run paged sparse GQA with the QSA Triton kernel."""
 
-    supports_dcp: bool = False
-    supports_pcp: bool = False
-
     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", "fp8", "fp8_e4m3"):
+        if 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
+        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
 
-    def forward_qsa(
+    def do_kv_cache_update(
         self,
         layer: torch.nn.Module,
-        query: torch.Tensor,
         key: torch.Tensor,
         value: torch.Tensor,
         kv_cache: torch.Tensor,
-        attn_metadata: FlashAttentionMetadata,
+        slot_mapping: torch.Tensor,
+    ) -> None:
+        """Scatter K/V using the same layout and scales as the fused writer."""
+        key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1)
+        ops.reshape_and_cache_flash(
+            key,
+            value,
+            key_cache,
+            value_cache,
+            slot_mapping,
+            self.kv_cache_dtype,
+            layer._k_scale,
+            layer._v_scale,
+        )
+
+    def forward(
+        self,
+        layer: torch.nn.Module,
+        query: torch.Tensor,
+        kv_cache: torch.Tensor,
+        attn_metadata: Qwen4ExpQSAMetadata,
         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:
-        del key, value
-        if output_scale is not None or output_block_scale is not None:
-            raise NotImplementedError("QSA does not support fused output quantization")
-        if self.alibi_slopes is not None or self.sinks is not None:
-            raise NotImplementedError("QSA does not support ALiBi or attention sinks")
-        if self.sliding_window != (-1, -1):
-            raise NotImplementedError("QSA does not support sliding-window attention")
-
+        """Run sparse attention over the owner's indexer-selected K/V rows."""
         num_tokens = attn_metadata.num_actual_tokens
         output.zero_()
         if num_tokens == 0:
@@ -338,7 +279,7 @@ def forward_qsa(
 
 
 class Qwen4ExpQSAAttention(Qwen3NextAttention, AttentionLayerBase):
-    """Merged Qwen full-attention owner with a QSA index side branch."""
+    """Merged Qwen full-attention owner with a QSA indexer branch."""
 
     supports_dcp = False
 
@@ -475,18 +416,13 @@ def __init__(
         self.kv_cache = torch.tensor([])
         set_default_quant_scales(self, register_buffer=True)
 
-        self.attn_backend = Qwen4ExpQSAFlashAttentionBackend
-        self.impl = Qwen4ExpQSAFlashAttentionImpl(
+        self.attn_backend = Qwen4ExpQSABackend
+        self.impl: Qwen4ExpQSAImpl = Qwen4ExpQSAImpl(  # type: ignore[assignment]
             self.num_heads,
             self.head_dim,
             self.scaling,
             self.num_kv_heads,
-            None,
-            None,
             self.kv_cache_dtype,
-            None,
-            AttentionType.DECODER,
-            None,
         )
         self.indexer = QSAIndexer(
             vllm_config=vllm_config,
@@ -614,17 +550,15 @@ def _run_qsa(
             )
             output.zero_()
             return
-        main_metadata = cast(FlashAttentionMetadata, metadata[self.layer_name])
+        main_metadata = cast(Qwen4ExpQSAMetadata, metadata[self.layer_name])
         if self.kv_cache.numel() == 0:
             raise RuntimeError("QSA main K/V cache is not bound")
 
         num_tokens = main_metadata.num_actual_tokens
-        side_metadata = cast(
-            QSAForwardMetadata,
+        indexer_metadata = cast(
+            QSAIndexerMetadata,
             metadata[self.indexer.raw_key_cache.prefix],
         )
-        if side_metadata.num_actual_tokens != num_tokens:
-            raise RuntimeError("QSA main and side metadata token counts disagree")
         selected, main_outputs = self.indexer(
             projected_qk,
             positions,
@@ -635,10 +569,9 @@ def _run_qsa(
         )
         if selected.shape != (num_tokens, self.indexer.packed_output_width):
             raise RuntimeError("QSA indexer returned an invalid selection shape")
-        impl = cast(Qwen4ExpQSAFlashAttentionImpl, self.impl)
         if main_outputs is None:
             assert key is not None and value is not None
-            impl.do_kv_cache_update(
+            self.impl.do_kv_cache_update(
                 self,
                 key,
                 value,
@@ -648,15 +581,13 @@ def _run_qsa(
         else:
             query, output_gate = main_outputs
         assert query is not None and output_gate is not None
-        impl.forward_qsa(
+        self.impl.forward(
             self,
             query,
-            key,
-            value,
             self.kv_cache,
             main_metadata,
             output,
-            token_to_req=side_metadata.token_to_req,
+            token_to_req=indexer_metadata.token_to_req,
             use_prefill_config=main_metadata.max_query_len > self._max_decode_query_len,
             output_gate=output_gate,
         )
@@ -702,6 +633,6 @@ def forward(
 __all__ = [
     "QSAIndexer",
     "Qwen4ExpQSAAttention",
-    "Qwen4ExpQSAFlashAttentionBackend",
-    "Qwen4ExpQSAFlashAttentionImpl",
+    "Qwen4ExpQSABackend",
+    "Qwen4ExpQSAImpl",
 ]
diff --git a/vllm/models/qwen4_exp/nvidia/indexer_qsa.py b/vllm/models/qwen4_exp/nvidia/qsa_indexer.py
similarity index 92%
rename from vllm/models/qwen4_exp/nvidia/indexer_qsa.py
rename to vllm/models/qwen4_exp/nvidia/qsa_indexer.py
index c8ffa0dbfbe3..2c7cebbb119a 100644
--- a/vllm/models/qwen4_exp/nvidia/indexer_qsa.py
+++ b/vllm/models/qwen4_exp/nvidia/qsa_indexer.py
@@ -15,10 +15,10 @@
 from vllm.model_executor.layers.quantization import QuantizationConfig
 from vllm.model_executor.layers.rotary_embedding.mrope import triton_mrope
 
-from ..common.qsa_cache import (
+from ..common.qsa_indexer import (
     QSACompressedKeyCache,
-    QSAForwardMetadata,
-    QSAKeyStateCache,
+    QSAIndexerMetadata,
+    QSARawKeyCache,
     canonical_qsa_rope_positions,
 )
 from .ops.qsa_prepare import qsa_prepare
@@ -88,7 +88,7 @@ def _supports_fused_pre_indexer(
 
 
 class QSAIndexer(nn.Module):
-    """QSA projection weights, side caches, and paged, weight-free selection.
+    """QSA projection weights, indexer caches, and paged, weight-free selection.
 
     ``prefix`` must be the checkpoint's indexer prefix, normally
     ``model.layers.N.self_attn.indexer``.  Consequently the trainable names are
@@ -126,7 +126,7 @@ def __init__(
         )
         self.prefix = prefix
         # MTP step 0 selects the target-aligned rows; later steps reuse them
-        # while continuing to update the QSA side cache.
+        # while continuing to update the QSA indexer cache.
         self.skip_topk = False
 
         self.index_qk_proj = ReplicatedLinear(
@@ -162,7 +162,7 @@ def __init__(
                 "by the Qwen4Exp QSA indexer (only 'bf16' or 'fp8')."
             )
         self.indexer_dtype = indexer_dtype
-        self.raw_key_cache = QSAKeyStateCache(
+        self.raw_key_cache = QSARawKeyCache(
             head_size=self.index_head_dim,
             dtype=torch.bfloat16,
             cache_rope_positions=vllm_config.model_config.uses_mrope,
@@ -197,38 +197,16 @@ def packed_output_width(self) -> int:
 
     def _metadata(
         self,
-    ) -> tuple[QSAForwardMetadata, QSAForwardMetadata] | None:
+    ) -> tuple[QSAIndexerMetadata, QSAIndexerMetadata] | None:
         metadata = get_forward_context().attn_metadata
         if isinstance(metadata, list):
             metadata = metadata[0]
         if not isinstance(metadata, dict):
             return None
-        raw = cast(QSAForwardMetadata, metadata[self.raw_key_cache.prefix])
+        raw = cast(QSAIndexerMetadata, metadata[self.raw_key_cache.prefix])
         compressed = cast(
-            QSAForwardMetadata, metadata[self.compressed_key_cache.prefix]
+            QSAIndexerMetadata, metadata[self.compressed_key_cache.prefix]
         )
-        if raw.num_actual_tokens != compressed.num_actual_tokens:
-            raise RuntimeError("QSA side-cache metadata token counts disagree")
-        raw_split = (
-            raw.num_decodes,
-            raw.num_decode_tokens,
-            raw.num_prefills,
-            raw.num_prefill_tokens,
-            raw.decode_query_len,
-        )
-        compressed_split = (
-            compressed.num_decodes,
-            compressed.num_decode_tokens,
-            compressed.num_prefills,
-            compressed.num_prefill_tokens,
-            compressed.decode_query_len,
-        )
-        if raw_split != compressed_split:
-            raise RuntimeError("QSA side-cache metadata batch splits disagree")
-        if not raw.logical_positions.is_cuda and (
-            not torch.equal(raw.logical_positions, compressed.logical_positions)
-        ):
-            raise RuntimeError("QSA side-cache metadata positions disagree")
         return raw, compressed
 
     def forward(
@@ -241,7 +219,7 @@ def forward(
         qkv: torch.Tensor | None = None,
         slot_mapping: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor] | None]:
-        """Update side caches and select token indices from pre-projected Q/K.
+        """Update indexer caches and select token indices from pre-projected Q/K.
 
         Returns the packed buffer of shape [num_tokens, output_width + 1]:
         the leading ``output_width`` columns are ``-1``-padded
diff --git a/vllm/models/step5/__init__.py b/vllm/models/step5/__init__.py
index e6b43b2f8183..e07b30a68dcd 100644
--- a/vllm/models/step5/__init__.py
+++ b/vllm/models/step5/__init__.py
@@ -1,9 +1,34 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Step5 entry point for the NVIDIA implementation."""
+"""Step5 entry point: the NVIDIA or AMD implementation, picked on lookup."""
+
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from .nvidia.model import Step5ForCausalLM
+    from .nvidia.mtp import Step5MTP
+    from .nvidia.vl import Step5ForConditionalGeneration
+
+
+def __getattr__(name: str) -> Any:
+    if name not in __all__:
+        raise AttributeError(name)
+    from vllm.platforms import current_platform
+
+    if current_platform.is_rocm():
+        from .amd.model import Step5ForCausalLM
+        from .amd.mtp import Step5MTP
+        from .amd.vl import Step5ForConditionalGeneration
+    else:
+        from .nvidia.model import Step5ForCausalLM
+        from .nvidia.mtp import Step5MTP
+        from .nvidia.vl import Step5ForConditionalGeneration
+
+    return {
+        "Step5ForCausalLM": Step5ForCausalLM,
+        "Step5ForConditionalGeneration": Step5ForConditionalGeneration,
+        "Step5MTP": Step5MTP,
+    }[name]
 
-from .nvidia.model import Step5ForCausalLM
-from .nvidia.mtp import Step5MTP
-from .nvidia.vl import Step5ForConditionalGeneration
 
 __all__ = ["Step5ForCausalLM", "Step5ForConditionalGeneration", "Step5MTP"]
diff --git a/vllm/models/step5/amd/__init__.py b/vllm/models/step5/amd/__init__.py
new file mode 100644
index 000000000000..208f01a7cb5e
--- /dev/null
+++ b/vllm/models/step5/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/step5/amd/checks.py b/vllm/models/step5/amd/checks.py
new file mode 100644
index 000000000000..94e441bcb32f
--- /dev/null
+++ b/vllm/models/step5/amd/checks.py
@@ -0,0 +1,70 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Startup checks for the ROCm Step5 sparse attention layers."""
+
+from vllm.config import KVTransferConfig, VllmConfig
+
+
+def verify_kv_transfer_compatibility(vllm_config: VllmConfig) -> None:
+    config = vllm_config.kv_transfer_config
+    if config is None or not config.is_kv_transfer_instance:
+        return
+    configs = [config]
+    if config.kv_connector == "MultiConnector":
+        children = config.kv_connector_extra_config["connectors"]
+        configs = [KVTransferConfig(**child) for child in children]
+    if not all(
+        c.kv_connector == "MooncakeStoreConnector"
+        or (c.kv_connector == "NixlConnector" and c.kv_buffer_device == "cuda")
+        for c in configs
+    ):
+        raise ValueError(
+            "Step5 sparse summaries are only transferred by "
+            "MooncakeStoreConnector, or by NixlConnector with device buffers, "
+            "directly or under a MultiConnector."
+        )
+
+
+def _groups_per_rank(total_groups: int, tp_size: int) -> int:
+    if total_groups % tp_size == 0:
+        return total_groups // tp_size
+    if tp_size % total_groups == 0:
+        return 1
+    raise ValueError(
+        f"Step5 tensor parallel size {tp_size} is incompatible with "
+        f"group count {total_groups}: one must divide the other."
+    )
+
+
+def validate_parallel_geometry(
+    *,
+    total_num_heads: int,
+    total_num_kv_heads: int,
+    indexer_num_heads: int,
+    index_tp_size: int,
+    tp_size: int,
+) -> None:
+    """The indexer scores regions per provider group and attention reads the
+    matching KV group, so both partitions must agree globally and per rank."""
+    if index_tp_size != total_num_kv_heads:
+        raise ValueError(
+            "Step5 sparse provider groups must match attention KV groups: "
+            f"{index_tp_size} != {total_num_kv_heads}"
+        )
+    if indexer_num_heads % index_tp_size:
+        raise ValueError(
+            f"{indexer_num_heads} sparse indexer heads do not divide into "
+            f"{index_tp_size} provider groups."
+        )
+    if total_num_heads % tp_size:
+        raise ValueError(
+            f"Step5 attention heads ({total_num_heads}) must be divisible by "
+            f"the tensor parallel size ({tp_size})."
+        )
+    local_kv_groups = _groups_per_rank(total_num_kv_heads, tp_size)
+    q_heads_per_kv, rem = divmod(total_num_heads // tp_size, local_kv_groups)
+    if rem or q_heads_per_kv not in (4, 8, 16):
+        raise ValueError(
+            "Step5 sparse attention needs 4, 8 or 16 local query heads per KV "
+            f"group (num_heads={total_num_heads}, tp_size={tp_size})."
+        )
diff --git a/vllm/models/step5/amd/csa_attention.py b/vllm/models/step5/amd/csa_attention.py
new file mode 100644
index 000000000000..74c1e149ca9c
--- /dev/null
+++ b/vllm/models/step5/amd/csa_attention.py
@@ -0,0 +1,210 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Step-5 CSA as one graph-splitting custom op: region summaries, scoring,
+top-k and sparse attention over the paged K/V."""
+
+from typing import TYPE_CHECKING, Any, cast
+
+import torch
+
+from vllm.compilation.breakable_cudagraph import eager_break_during_capture
+from vllm.forward_context import get_forward_context
+from vllm.models.step5.amd.csa_cache import Step5CSAMetadata
+from vllm.models.step5.amd.ops.csa_paged import sparse_attention_paged
+from vllm.models.step5.common.ops.sparse_decode_attention import (
+    sparse_decode_attention,
+    sparse_decode_splits,
+)
+from vllm.models.step5.common.ops.sparse_indexer import (
+    pack_summary_pages,
+    select_all_regions,
+    select_regions,
+)
+from vllm.models.step5.common.ops.sparse_indexer_scores import (
+    decode_summary_scores,
+    paged_mqa_logits_sm90,
+)
+from vllm.platforms import current_platform
+from vllm.utils.torch_utils import (
+    LayerNameType,
+    _resolve_layer_name,
+    direct_register_custom_op,
+)
+from vllm.v1.attention.backends.rocm_attn import RocmAttentionMetadata
+
+if TYPE_CHECKING:
+    from vllm.models.step5.amd.model import Step5Attention
+
+
+@eager_break_during_capture
+def step5_csa_attention(
+    hidden_states: torch.Tensor,
+    positions: torch.Tensor,
+    query: torch.Tensor,
+    output: torch.Tensor,
+    layer_name: LayerNameType,
+    gate: torch.Tensor | None = None,
+    kv_cache_dummy_dep: torch.Tensor | None = None,
+) -> None:
+    del kv_cache_dummy_dep
+    layer: Step5Attention = get_forward_context().no_compile_layers[
+        _resolve_layer_name(layer_name)
+    ]
+    indexer = layer.indexer
+    assert indexer is not None
+    metadata = get_forward_context().attn_metadata
+    if isinstance(metadata, list):
+        metadata = metadata[0]
+    if not isinstance(metadata, dict):
+        # Profiling run: caches are not bound yet; claim the scoring workspace.
+        torch.empty(indexer.workspace_bytes, dtype=torch.uint8, device=output.device)
+        output.zero_()
+        return
+    md = metadata[layer.sparse_summary_cache.layer_name]
+    main_md = metadata[layer.attn.layer_name]
+    assert isinstance(md, Step5CSAMetadata)
+    assert isinstance(main_md, RocmAttentionMetadata)
+    n = md.num_actual_tokens
+    output[n:].zero_()
+    if n == 0:
+        return
+
+    cache = cast(Any, layer.attn.impl).summary_cache
+    ring = layer.csa_tail_ring
+    slots = md.slot_mapping[:n]
+    uniform = md.prefill_groups is None
+    # Without drafts every decode token closes at most its own region, so the
+    # summary update fuses into the indexer kernel.
+    fuse = uniform and indexer.num_speculative_tokens == 0
+    q8, k, z, w = indexer(
+        hidden_states[:n], positions[:n], (cache, ring, slots) if fuse else None
+    )
+    groups = q8.shape[1]
+    rows = n * groups
+
+    kv_cache = layer.attn.kv_cache
+    k_scale = v_scale = None
+    if kv_cache.dtype == torch.uint8:
+        kv_cache = kv_cache.view(current_platform.fp8_dtype())
+        k_scale, v_scale = layer.attn._k_scale, layer.attn._v_scale
+
+    if uniform:
+        # Decode or MTP verify: each request owns max_query_len consecutive rows.
+        query_len = md.max_query_len
+        reqs = n // query_len
+        table = md.block_table[:reqs]
+        if indexer.num_speculative_tokens:
+            cache.update_verify(
+                k, z, slots, md.query_start_loc[: reqs + 1], md.seq_lens[:reqs], ring
+            )
+        capacity = min(indexer.max_regions, table.shape[1] * cache.mean.shape[1])
+        scores, decode = decode_summary_scores(
+            q8,
+            w,
+            cache.mean,
+            table,
+            slots,
+            md.seq_lens[:reqs],
+            capacity,
+            num_sms=indexer.num_sms,
+            row_groups=groups,
+            query_len=query_len,
+        )
+        regions = select_regions(
+            scores,
+            decode.current,
+            indexer.topk,
+            indexer.topk_backend,
+            topk_lengths=decode.topk_lengths,
+            sort=False,
+        )
+        sparse_decode_attention(
+            query[:n].reshape(rows, -1, layer.head_dim),
+            kv_cache,
+            regions,
+            main_md.block_table,
+            decode.positions,
+            decode.active,
+            output[:n].view(rows, -1, layer.head_dim),
+            scale=layer.scaling,
+            query_len=query_len,
+            splits=sparse_decode_splits(rows),
+            k_scale=k_scale,
+            v_scale=v_scale,
+            gate=None if gate is None else gate[:n],
+        )
+        return
+
+    assert md.prefill_groups is not None
+    cache.update_prefill(k, z, md.prefill_groups, md.positions, ring)
+    if md.prefill_indexer is None:
+        # Every earlier region fits in top-k: nothing to score.
+        own = torch.div(md.positions, indexer.region_size, rounding_mode="floor")
+        regions = select_all_regions(
+            own.repeat_interleave(groups) if groups > 1 else own, indexer.topk
+        ).view(n, groups, -1)
+    else:
+        regions = torch.empty(
+            (n, groups, indexer.topk + 1), dtype=torch.int64, device=query.device
+        )
+        packed, _ = pack_summary_pages(
+            cache.mean, md.block_table[: md.num_reqs], md.num_regions
+        )
+        start = 0
+        for chunk in md.prefill_indexer:
+            end = start + chunk.current.numel()
+            for g in range(groups):
+                scores = paged_mqa_logits_sm90(
+                    q=q8[start:end, g].unsqueeze(1),
+                    kv_fused=packed,
+                    weights=w[start:end, g],
+                    seq_lens=chunk.lengths,
+                    block_tables=chunk.block_table,
+                    max_seq_len=md.num_regions,
+                )
+                regions[start:end, g] = select_regions(
+                    scores,
+                    chunk.current,
+                    indexer.topk,
+                    indexer.topk_backend,
+                    topk_lengths=chunk.topk_lengths,
+                    sort=False,
+                )
+            start = end
+    key_cache, value_cache = kv_cache.transpose(1, 2).split(layer.head_dim, dim=-1)
+    sparse_attention_paged(
+        query[:n],
+        key_cache,
+        value_cache,
+        regions,
+        main_md.block_table,
+        md.token_req,
+        md.positions,
+        output[:n],
+        region_size=indexer.region_size,
+        scale=layer.scaling,
+        k_scale=k_scale,
+        v_scale=v_scale,
+    )
+    if gate is not None:
+        output[:n] *= gate[:n].unsqueeze(-1).sigmoid().to(output.dtype)
+
+
+def step5_csa_attention_fake(
+    hidden_states: torch.Tensor,
+    positions: torch.Tensor,
+    query: torch.Tensor,
+    output: torch.Tensor,
+    layer_name: LayerNameType,
+    gate: torch.Tensor | None = None,
+    kv_cache_dummy_dep: torch.Tensor | None = None,
+) -> None:
+    return
+
+
+direct_register_custom_op(
+    op_name="step5_csa_attention",
+    op_func=step5_csa_attention,
+    mutates_args=["output"],
+    fake_impl=step5_csa_attention_fake,
+)
diff --git a/vllm/models/step5/amd/csa_cache.py b/vllm/models/step5/amd/csa_cache.py
new file mode 100644
index 000000000000..8661f89625f4
--- /dev/null
+++ b/vllm/models/step5/amd/csa_cache.py
@@ -0,0 +1,154 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Step-5 CSA region summaries as their own KV cache group, and the per-step
+metadata the CSA op reads from it."""
+
+from dataclasses import dataclass
+from typing import ClassVar
+
+import torch
+
+from vllm.config import VllmConfig
+from vllm.config.cache import CacheDType
+from vllm.models.step5.common.config import get_step5_sparse_config
+from vllm.models.step5.common.ops.prefill_indexer_metadata import (
+    PrefillIndexerMetadata,
+    prepare_prefill_indexer,
+)
+from vllm.models.step5.common.ops.summary_update import build_prefill_groups
+from vllm.models.step5.common.sparse_summary_cache import (
+    Step5SparseSummaryCacheLayer,
+)
+from vllm.utils.math_utils import cdiv
+from vllm.v1.attention.backend import (
+    AttentionBackend,
+    AttentionCGSupport,
+    AttentionMetadata,
+    AttentionMetadataBuilder,
+    CommonAttentionMetadata,
+)
+from vllm.v1.kv_cache_interface import AttentionSpec
+
+
+@dataclass
+class Step5CSAMetadata(AttentionMetadata):
+    # Block table and token slots of the summary group.
+    block_table: torch.Tensor
+    slot_mapping: torch.Tensor
+    query_start_loc: torch.Tensor
+    seq_lens: torch.Tensor
+    num_actual_tokens: int
+    num_reqs: int
+    max_query_len: int
+    token_req: torch.Tensor
+    positions: torch.Tensor
+    num_regions: int = 0
+    # None for uniform decode and MTP verify batches.
+    prefill_groups: torch.Tensor | None = None
+    # None when every earlier region fits in top-k.
+    prefill_indexer: list[PrefillIndexerMetadata] | None = None
+
+
+class Step5CSAMetadataBuilder(AttentionMetadataBuilder[Step5CSAMetadata]):
+    _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.NEVER
+
+    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)
+        sparse = get_step5_sparse_config(vllm_config.model_config.hf_text_config)
+        assert sparse is not None
+        self.region_size = sparse.region_block_size
+        self.topk = sparse.topk
+        spec = vllm_config.speculative_config
+        self.max_uniform_query_len = 1 + (
+            spec.num_speculative_tokens if spec is not None else 0
+        )
+        self.num_sms = torch.cuda.get_device_properties(device).multi_processor_count
+        self.token_req_buffer = torch.empty(
+            vllm_config.scheduler_config.max_num_batched_tokens,
+            dtype=torch.int32,
+            device=device,
+        )
+
+    def build(
+        self,
+        common_prefix_len: int,
+        common_attn_metadata: CommonAttentionMetadata,
+        fast_build: bool = False,
+    ) -> Step5CSAMetadata:
+        del common_prefix_len, fast_build
+        m = common_attn_metadata
+        n = m.num_actual_tokens
+        token_req = m.token_to_req_indices(self.token_req_buffer)
+        if m.positions is None:
+            positions = (
+                m.compute_num_computed_tokens().index_select(0, token_req)
+                + torch.arange(n, device=token_req.device)
+                - m.query_start_loc.index_select(0, token_req)
+            )
+        else:
+            positions = m.positions[:n]
+        md = Step5CSAMetadata(
+            block_table=m.block_table_tensor,
+            slot_mapping=m.slot_mapping[:n],
+            query_start_loc=m.query_start_loc,
+            seq_lens=m.seq_lens,
+            num_actual_tokens=n,
+            num_reqs=m.num_reqs,
+            max_query_len=m.max_query_len,
+            token_req=token_req,
+            positions=positions.to(torch.int64),
+        )
+        uniform = (
+            n == m.num_reqs * m.max_query_len
+            and m.max_query_len <= self.max_uniform_query_len
+        )
+        if n and not uniform:
+            md.prefill_groups = build_prefill_groups(
+                md.slot_mapping,
+                md.positions,
+                token_req,
+                m.query_start_loc[: m.num_reqs + 1],
+            )
+            # Upper bound: scores past each row's history are masked.
+            md.num_regions = cdiv(m.max_seq_len, self.region_size)
+            if md.num_regions > self.topk:
+                md.prefill_indexer = prepare_prefill_indexer(
+                    token_req, md.positions, md.num_regions, self.topk, self.num_sms
+                )
+        return md
+
+
+class Step5CSABackend(AttentionBackend):
+    supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16]
+    supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = ["auto"]
+
+    @staticmethod
+    def get_name() -> str:
+        return "STEP5_CSA_SUMMARY"
+
+    @staticmethod
+    def get_impl_cls():
+        raise NotImplementedError("CSA summaries have no attention impl")
+
+    @staticmethod
+    def get_builder_cls() -> type[Step5CSAMetadataBuilder]:
+        return Step5CSAMetadataBuilder
+
+    @classmethod
+    def is_sparse(cls) -> bool:
+        return True
+
+
+class Step5CSASummaryCacheLayer(Step5SparseSummaryCacheLayer):
+    """FP8 region summaries in their own KV cache group: under the
+    layer-outer ROCm KV layout they use a 384-token block, so the CSA op
+    reads them through that group's metadata."""
+
+    def get_attn_backend(self) -> type[AttentionBackend]:
+        return Step5CSABackend
diff --git a/vllm/models/step5/amd/model.py b/vllm/models/step5/amd/model.py
new file mode 100644
index 000000000000..f76273755f47
--- /dev/null
+++ b/vllm/models/step5/amd/model.py
@@ -0,0 +1,780 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Step5 on ROCm: compressed sparse attention (CSA) and an FP32 residual."""
+
+from collections.abc import Iterable
+
+import torch
+from torch import nn
+
+import vllm.envs as envs
+from vllm._aiter_ops import rocm_aiter_ops
+from vllm.compilation.decorators import support_torch_compile
+from vllm.config import VllmConfig, get_current_vllm_config
+from vllm.distributed import (
+    get_ep_group,
+    get_pp_group,
+    get_tensor_model_parallel_rank,
+    get_tensor_model_parallel_world_size,
+)
+from vllm.model_executor.layers.activation import SiluAndMul, SwigluStepAndMul
+from vllm.model_executor.layers.attention import Attention
+from vllm.model_executor.layers.attention.attention import unified_kv_cache_update
+from vllm.model_executor.layers.fused_moe import FusedMoEFactory, GateLinear
+from vllm.model_executor.layers.layernorm import GemmaRMSNorm
+from vllm.model_executor.layers.linear import (
+    ColumnParallelLinear,
+    MergedColumnParallelLinear,
+    QKVParallelLinear,
+    ReplicatedLinear,
+    RowParallelLinear,
+)
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
+from vllm.model_executor.layers.rotary_embedding import get_rope
+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.interfaces import MixtureOfExperts, SupportsPP
+from vllm.model_executor.models.utils import (
+    AutoWeightsLoader,
+    PPMissingLayer,
+    WeightsMapper,
+    extract_layer_index,
+    is_pp_missing_parameter,
+    make_layers,
+    maybe_prefix,
+    skip_spec_layers,
+)
+from vllm.platforms import current_platform
+from vllm.platforms.rocm import on_gfx1250
+from vllm.sequence import IntermediateTensors
+from vllm.transformers_utils.configs.step5 import Step5Config
+from vllm.utils.math_utils import cdiv, round_up
+from vllm.utils.torch_utils import _encode_layer_name, kv_cache_dtype_str_to_dtype
+from vllm.v1.attention.backends.rocm_aiter_unified_attn import (
+    RocmAiterUnifiedAttentionBackend,
+)
+
+from ..common.config import get_step5_sparse_config
+from ..common.csa_tail_ring import Step5CSATailRing
+from ..common.ops.indexer_norm_rope import indexer_norm_rope
+from .checks import validate_parallel_geometry, verify_kv_transfer_compatibility
+from .csa_attention import step5_csa_attention
+from .csa_cache import Step5CSASummaryCacheLayer
+
+
+class Step5SparseIndexer(nn.Module):
+    """MQA indexer with ``heads_per_group`` query heads per KV group.
+
+    Q and W are sharded like the attention KV heads, so each rank computes only
+    its own groups; the shared single-head K and Z stay replicated.
+    """
+
+    def __init__(self, config: Step5Config, layer_idx: int, prefix: str) -> None:
+        super().__init__()
+        sparse = get_step5_sparse_config(config)
+        assert sparse is not None
+        hidden = config.hidden_size
+        self.proxy_dim = sparse.proxy_dim
+        self.region_size = sparse.region_block_size
+        self.topk = sparse.topk
+        self.num_groups = sparse.index_tp_size
+        self.heads_per_group = sparse.sparse_indexer_num_heads // self.num_groups
+        self.eps = config.rms_norm_eps
+        self.zero_centered = config.zero_centered
+
+        tp_size = get_tensor_model_parallel_world_size()
+        group_shards = min(tp_size, self.num_groups)
+        group_rank = get_tensor_model_parallel_rank() * group_shards // tp_size
+        self.num_local_groups = self.num_groups // group_shards
+
+        def linear(name: str, out: int):
+            return ReplicatedLinear(
+                hidden, out, bias=False, quant_config=None, prefix=f"{prefix}.{name}"
+            )
+
+        def group_linear(name: str, out: int):
+            return ColumnParallelLinear(
+                hidden,
+                out,
+                bias=False,
+                quant_config=None,
+                prefix=f"{prefix}.{name}",
+                tp_rank=group_rank,
+                tp_size=group_shards,
+            )
+
+        heads = self.num_groups * self.heads_per_group
+        self.sparse_indexer_q = group_linear("sparse_indexer_q", heads * self.proxy_dim)
+        self.sparse_indexer_k = linear("sparse_indexer_k", self.proxy_dim)
+        self.sparse_indexer_z = linear("sparse_indexer_z", self.proxy_dim)
+        self.sparse_indexer_w = group_linear("sparse_indexer_w", heads)
+        self.sparse_indexer_q_norm = nn.Module()
+        self.sparse_indexer_q_norm.weight = nn.Parameter(
+            torch.zeros(self.proxy_dim), requires_grad=False
+        )
+        self.sparse_indexer_k_norm = nn.LayerNorm(self.proxy_dim, eps=self.eps)
+
+        self.rotary_emb = get_rope(
+            head_size=self.proxy_dim,
+            max_position=config.max_position_embeddings,
+            rope_parameters=_rope_parameters(config, layer_idx)
+            | {
+                "partial_rotary_factor": sparse.sparse_indexer_rope_dim / self.proxy_dim
+            },
+        )
+
+        self.native_rope = not self.rotary_emb.enabled()
+        vllm_config = get_current_vllm_config()
+        self.topk_backend = vllm_config.kernel_config.sparse_indexer_topk_backend
+        spec = vllm_config.speculative_config
+        self.num_speculative_tokens = spec.num_speculative_tokens if spec else 0
+        self.max_regions = cdiv(
+            vllm_config.model_config.max_model_len, self.region_size
+        )
+        self.num_sms = current_platform.num_compute_units()
+        # Peak scoring workspace: decode scores or bounded prefill logits, plus
+        # the per-step prefill page table.
+        sched = vllm_config.scheduler_config
+        decode_rows = (
+            sched.max_num_seqs
+            * (1 + self.num_speculative_tokens)
+            * self.num_local_groups
+        )
+        decode_scores = decode_rows * round_up(self.max_regions, 256) * 4
+        prefill_table = sched.max_num_batched_tokens * cdiv(self.max_regions, 128) * 16
+        self.workspace_bytes = (
+            max(envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB << 20, decode_scores)
+            + prefill_table
+        )
+
+    def forward(
+        self, hidden_states: torch.Tensor, positions: torch.Tensor, decode=None
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Returns ``(q [T, G, 8, P] fp8, k [T, 1, P], z [T, 1, P],
+        w [T, G, 8])`` for this rank's ``G`` local groups, heads padded to 8.
+
+        With ``decode=(summary cache, ring, slots)`` it also stores each
+        token's open-region row and summarizes the regions it closes.
+        """
+        q, _ = self.sparse_indexer_q(hidden_states)
+        k, _ = self.sparse_indexer_k(hidden_states)
+        z, _ = self.sparse_indexer_z(hidden_states)
+        w, _ = self.sparse_indexer_w(hidden_states)
+        fused = {}
+        if decode is not None:
+            cache, ring, slots = decode
+            fused = dict(
+                decode_cache=cache, decode_ring=ring, decode_z=z, decode_slots=slots
+            )
+        q, k, w = indexer_norm_rope(
+            q,
+            k,
+            w,
+            self.sparse_indexer_q_norm.weight,
+            self.sparse_indexer_k_norm.weight,
+            self.sparse_indexer_k_norm.bias,
+            self.rotary_emb.cos_sin_cache,
+            positions,
+            self.eps,
+            self.zero_centered,
+            native_rope=self.native_rope,
+            num_groups=self.num_local_groups,
+            **fused,
+        )
+        return q, k, z.view(-1, 1, self.proxy_dim), w
+
+
+def _layer_types(config: Step5Config) -> list[str]:
+    # `layer_types` covers only the dense stack; MTP layers index past it.
+    return config.layer_types_with_mtp or config.layer_types or []
+
+
+def _num_attention_heads(config: Step5Config, layer_type: str) -> int:
+    if layer_type == "sliding_attention":
+        if num_heads := getattr(config, "num_sliding_attention_heads", None):
+            return num_heads
+        other = config.attention_other_setting or {}
+        if other.get("attention_type") == layer_type:
+            return other["num_attention_heads"]
+    return config.num_attention_heads
+
+
+def _rope_parameters(config: Step5Config, layer_idx: int) -> dict:
+    rope_parameters = dict(_rope_scaling(config, layer_idx) or {})
+    rope_parameters.setdefault("rope_type", "default")
+    theta = config.rope_theta
+    rope_parameters["rope_theta"] = (
+        theta[layer_idx] if isinstance(theta, list) else theta
+    )
+    factors = config.partial_rotary_factors
+    rope_parameters["partial_rotary_factor"] = factors[layer_idx] if factors else 1.0
+    return rope_parameters
+
+
+def _is_moe_layer(config: Step5Config, layer_idx: int) -> bool:
+    if (moe_layers := getattr(config, "moe_layers_enum", None)) is not None:
+        return layer_idx in {int(i) for i in moe_layers.strip().split(",")}
+    return 0 < layer_idx < config.num_hidden_layers
+
+
+def _swiglu_limit(limits: list | None, layer_idx: int) -> float | None:
+    limit = limits[layer_idx] if limits and layer_idx < len(limits) else None
+    return float(limit) if limit else None
+
+
+class Step5MLP(nn.Module):
+    def __init__(
+        self,
+        config: Step5Config,
+        intermediate_size: int,
+        quant_config,
+        reduce_results: bool,
+        prefix: str,
+    ) -> None:
+        super().__init__()
+        self.gate_up_proj = MergedColumnParallelLinear(
+            config.hidden_size,
+            [intermediate_size] * 2,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.gate_up_proj",
+        )
+        self.down_proj = RowParallelLinear(
+            intermediate_size,
+            config.hidden_size,
+            bias=False,
+            quant_config=quant_config,
+            reduce_results=reduce_results,
+            prefix=f"{prefix}.down_proj",
+        )
+        limit = _swiglu_limit(config.swiglu_limits_shared, extract_layer_index(prefix))
+        self.act_fn = SwigluStepAndMul(limit=limit) if limit else SiluAndMul()
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        gate_up, _ = self.gate_up_proj(hidden_states)
+        output, _ = self.down_proj(self.act_fn(gate_up))
+        return output
+
+
+class Step5MoE(nn.Module):
+    """Routed experts with the shared expert fused into one FusedMoE call, so
+    the combined output is all-reduced once in the model dtype."""
+
+    def __init__(self, vllm_config: VllmConfig, prefix: str) -> None:
+        super().__init__()
+        config = vllm_config.model_config.hf_config
+        parallel_config = vllm_config.parallel_config
+        layer_idx = extract_layer_index(prefix)
+        self.n_routed_experts = config.moe_num_experts
+        self.n_logical_experts = self.n_routed_experts
+        self.n_redundant_experts = parallel_config.eplb_config.num_redundant_experts
+        if self.n_redundant_experts:
+            raise NotImplementedError(
+                "Step5 on ROCm does not support redundant experts"
+            )
+        self.n_physical_experts = self.n_logical_experts + self.n_redundant_experts
+        self.n_local_physical_experts = (
+            self.n_physical_experts // get_ep_group().device_group.size()
+        )
+        if not (config.use_moe_router_bias and config.need_fp32_gate):
+            raise ValueError("Step5 MoE requires an FP32 gate with a router bias.")
+        limit = _swiglu_limit(config.swiglu_limits, layer_idx)
+        if limit not in (None, 7.0):
+            raise ValueError(f"Step5 fused MoE supports swiglu_limit 7.0, got {limit}")
+
+        self.gate = GateLinear(
+            config.hidden_size,
+            config.moe_num_experts,
+            out_dtype=torch.float32,
+            params_dtype=torch.float32,
+            prefix=f"{prefix}.gate",
+        )
+        self.router_bias = nn.Parameter(
+            torch.zeros(config.moe_num_experts, dtype=torch.float32),
+            requires_grad=False,
+        )
+        self.share_expert = Step5MLP(
+            config,
+            config.share_expert_dim,
+            vllm_config.quant_config,
+            reduce_results=False,
+            prefix=f"{prefix}.share_expert",
+        )
+        self.experts = FusedMoEFactory(
+            shared_experts=self.share_expert,
+            gate=self.gate,
+            num_experts=config.moe_num_experts,
+            top_k=config.moe_top_k,
+            hidden_size=config.hidden_size,
+            intermediate_size=config.moe_intermediate_size,
+            renormalize=config.norm_expert_weight,
+            quant_config=vllm_config.quant_config,
+            activation="swiglustep" if limit else "silu",
+            prefix=f"{prefix}.experts",
+            scoring_func=config.moe_router_activation,
+            e_score_correction_bias=self.router_bias,
+            routed_scaling_factor=config.moe_router_scaling_factor,
+            enable_eplb=parallel_config.enable_eplb,
+            num_redundant_experts=self.n_redundant_experts,
+            router_logits_dtype=torch.float32,
+        )
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.experts(hidden_states=hidden_states, router_logits=hidden_states)
+
+
+def _is_sparse_layer(config: Step5Config, layer_idx: int) -> bool:
+    sparse = get_step5_sparse_config(config)
+    layer_types = _layer_types(config)
+    return (
+        sparse is not None
+        and layer_idx < len(layer_types)
+        and layer_types[layer_idx] in sparse.apply_to_layer_types
+    )
+
+
+def _rope_scaling(config: Step5Config, layer_idx: int) -> dict | None:
+    yarn_only = getattr(config, "yarn_only_types", None)
+    layer_types = _layer_types(config)
+    if yarn_only and layer_types and layer_types[layer_idx] not in yarn_only:
+        return None
+    return getattr(config, "rope_scaling", None)
+
+
+class Step5Attention(nn.Module):
+    def __init__(self, vllm_config: VllmConfig, prefix: str) -> None:
+        super().__init__()
+        config = vllm_config.model_config.hf_config
+        cache_config = vllm_config.cache_config
+        quant_config = vllm_config.quant_config
+        self.layer_idx = extract_layer_index(prefix)
+        layer_type = _layer_types(config)[self.layer_idx]
+        tp_size = get_tensor_model_parallel_world_size()
+        self.total_num_heads = _num_attention_heads(config, layer_type)
+        self.total_num_kv_heads = config.num_attention_groups
+        self.num_heads = self.total_num_heads // tp_size
+        self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
+        self.head_dim = getattr(config, "head_dim", None) or (
+            config.hidden_size // self.total_num_heads
+        )
+        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
+        use_rope_layers = getattr(config, "use_rope_layers", None)
+        self.use_rope = use_rope_layers[self.layer_idx] if use_rope_layers else True
+        self.use_head_wise_attn_gate = getattr(config, "use_head_wise_attn_gate", False)
+
+        self.qkv_proj = QKVParallelLinear(
+            config.hidden_size,
+            self.head_dim,
+            self.total_num_heads,
+            self.total_num_kv_heads,
+            bias=getattr(config, "attention_bias", False),
+            quant_config=quant_config,
+            prefix=f"{prefix}.qkv_proj",
+        )
+        self.o_proj = RowParallelLinear(
+            self.total_num_heads * self.head_dim,
+            config.hidden_size,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.o_proj",
+        )
+        if self.use_head_wise_attn_gate:
+            self.g_proj = ColumnParallelLinear(
+                config.hidden_size,
+                self.total_num_heads,
+                bias=False,
+                quant_config=quant_config,
+                prefix=f"{prefix}.g_proj",
+            )
+        self.q_norm = GemmaRMSNorm(self.head_dim, config.rms_norm_eps)
+        self.k_norm = GemmaRMSNorm(self.head_dim, config.rms_norm_eps)
+        self.rotary_emb = get_rope(
+            head_size=self.head_dim,
+            max_position=config.max_position_embeddings,
+            rope_parameters=_rope_parameters(config, self.layer_idx),
+        )
+
+        sparse = _is_sparse_layer(config, self.layer_idx)
+        self.attn = Attention(
+            self.num_heads,
+            self.head_dim,
+            self.scaling,
+            num_kv_heads=self.num_kv_heads,
+            cache_config=cache_config,
+            quant_config=quant_config,
+            prefix=f"{prefix}.attn",
+            per_layer_sliding_window=getattr(config, "sliding_window", None)
+            if layer_type == "sliding_attention"
+            else None,
+            # The CSA kernels read this backend's paged K/V layout.
+            attn_backend=RocmAiterUnifiedAttentionBackend if sparse else None,
+        )
+
+        self.indexer = None
+        if not sparse:
+            return
+        self._validate_sparse(vllm_config)
+        self.indexer = Step5SparseIndexer(config, self.layer_idx, prefix=prefix)
+        static_context = vllm_config.compilation_config.static_forward_context
+        self.sparse_summary_cache = Step5CSASummaryCacheLayer(
+            prefix=f"{prefix}.attn.sparse_summary_cache",
+            target_impl=self.attn.impl,
+            static_forward_context=static_context,
+        )
+        self.csa_tail_ring = Step5CSATailRing(
+            prefix=f"{prefix}.attn.csa_tail_ring",
+            main_layer_name=self.attn.layer_name,
+            num_speculative_tokens=self.indexer.num_speculative_tokens,
+            static_forward_context=static_context,
+        )
+        self.csa_layer_name = f"{prefix}.csa"
+        if self.csa_layer_name in static_context:
+            raise ValueError(f"Duplicate layer name: {self.csa_layer_name}")
+        static_context[self.csa_layer_name] = self
+
+    def _validate_sparse(self, vllm_config: VllmConfig) -> None:
+        cache_dtype = vllm_config.cache_config.cache_dtype
+        if (
+            cache_dtype not in ("fp8", "fp8_e4m3")
+            and kv_cache_dtype_str_to_dtype(cache_dtype, vllm_config.model_config)
+            != torch.bfloat16
+        ):
+            raise NotImplementedError(
+                "Step-5 sparse attention reads a BF16 or FP8 E4M3 KV cache; "
+                f"got {cache_dtype}"
+            )
+        if not self.use_rope:
+            raise NotImplementedError("Step-5 sparse layers require RoPE")
+        if self.head_dim != 192:
+            raise ValueError(
+                f"Step-5 sparse attention requires head_dim 192, got {self.head_dim}"
+            )
+        if vllm_config.parallel_config.use_ubatching:
+            raise ValueError(
+                "Step-5 sparse attention does not support dual-batch overlap or "
+                "microbatching yet. Disable --enable-dbo and set --ubatch-size 1."
+            )
+        verify_kv_transfer_compatibility(vllm_config)
+        if rocm_aiter_ops.is_shuffle_kv_cache_enabled() and on_gfx1250():
+            raise NotImplementedError(
+                "Step-5 sparse attention reads the unshuffled KV cache layout."
+            )
+        sparse = get_step5_sparse_config(vllm_config.model_config.hf_config)
+        assert sparse is not None
+        validate_parallel_geometry(
+            total_num_heads=self.total_num_heads,
+            total_num_kv_heads=self.total_num_kv_heads,
+            indexer_num_heads=sparse.sparse_indexer_num_heads,
+            index_tp_size=sparse.index_tp_size,
+            tp_size=get_tensor_model_parallel_world_size(),
+        )
+
+    def forward(
+        self, positions: torch.Tensor, hidden_states: torch.Tensor
+    ) -> torch.Tensor:
+        qkv, _ = self.qkv_proj(hidden_states)
+        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+        q = self.q_norm(q.view(-1, self.num_heads, self.head_dim)).view(q.shape)
+        k = self.k_norm(k.view(-1, self.num_kv_heads, self.head_dim)).view(k.shape)
+        if self.use_rope:
+            q, k = self.rotary_emb(positions, q, k)
+        gate = None
+        if self.use_head_wise_attn_gate:
+            gate, _ = self.g_proj(hidden_states)
+        if self.indexer is None:
+            out = self.attn(q, k, v)
+            if gate is not None:
+                out = (
+                    out.view(-1, self.num_heads, self.head_dim)
+                    * gate.unsqueeze(-1).sigmoid()
+                ).view(out.shape)
+            output, _ = self.o_proj(out)
+            return output
+
+        query = q.view(-1, self.num_heads, self.head_dim)
+        k = k.view(-1, self.num_kv_heads, self.head_dim)
+        v = v.view(-1, self.num_kv_heads, self.head_dim)
+        out = torch.empty_like(query)
+        attn_name = _encode_layer_name(self.attn.layer_name)
+        csa_name = _encode_layer_name(self.csa_layer_name)
+        if current_platform.opaque_attention_op():
+            kv_dep = torch.ops.vllm.unified_kv_cache_update(k, v, attn_name)
+            torch.ops.vllm.step5_csa_attention(
+                hidden_states, positions, query, out, csa_name, gate, kv_dep
+            )
+        else:
+            kv_dep = unified_kv_cache_update(k, v, attn_name)
+            step5_csa_attention(
+                hidden_states, positions, query, out, csa_name, gate, kv_dep
+            )
+        output, _ = self.o_proj(out.view(-1, self.q_size))
+        return output
+
+
+def use_fp32_norm_weights(module: nn.Module, config: Step5Config) -> None:
+    """The checkpoint keeps every RMSNorm weight in ``norm_dtype``."""
+    norm_dtype = (getattr(config, "norm_dtype", None) or "").lower()
+    if norm_dtype not in ("fp32", "float32", "float"):
+        return
+    for norm in module.modules():
+        if isinstance(norm, GemmaRMSNorm):
+            norm.weight.data = norm.weight.data.float()
+
+
+class Step5DecoderLayer(nn.Module):
+    def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None:
+        super().__init__()
+        config = vllm_config.model_config.hf_config
+        layer_idx = extract_layer_index(prefix)
+        if vllm_config.cache_config is not None:
+            vllm_config.cache_config.sliding_window = None
+        self.self_attn = Step5Attention(vllm_config, prefix=f"{prefix}.self_attn")
+        self.use_moe = _is_moe_layer(config, layer_idx)
+        if self.use_moe:
+            self.moe = Step5MoE(vllm_config, prefix=f"{prefix}.moe")
+        else:
+            self.mlp = Step5MLP(
+                config,
+                config.intermediate_size,
+                vllm_config.quant_config,
+                reduce_results=True,
+                prefix=f"{prefix}.mlp",
+            )
+        self.input_layernorm = GemmaRMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.post_attention_layernorm = GemmaRMSNorm(
+            config.hidden_size, config.rms_norm_eps
+        )
+        self.act_dtype = vllm_config.model_config.dtype
+
+    def forward(
+        self, positions: torch.Tensor, hidden_states: torch.Tensor
+    ) -> torch.Tensor:
+        residual = hidden_states
+        hidden = self.input_layernorm(residual).to(self.act_dtype)
+        residual = residual + self.self_attn(positions, hidden).float()
+        hidden = self.post_attention_layernorm(residual).to(self.act_dtype)
+        ffn = self.moe(hidden) if self.use_moe else self.mlp(hidden)
+        return residual + ffn.float()
+
+
+_STACKED_PARAMS: tuple[tuple[str, str, str | int], ...] = (
+    ("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),
+)
+# The checkpoint packs each MoE projection as one [experts, out, in] tensor.
+_EXPERT_PARAMS = tuple(
+    (f"moe.experts.routed_experts.{param}{suffix}", f"moe.{proj}.{suffix}", shard)
+    for proj, param, shard in (
+        ("gate_proj", "w13_", "w1"),
+        ("up_proj", "w13_", "w3"),
+        ("down_proj", "w2_", "w2"),
+    )
+    for suffix in ("weight_scale_inv", "weight")
+)
+
+
+def load_step5_weights(
+    model: nn.Module, weights: Iterable[tuple[str, torch.Tensor]]
+) -> set[str]:
+    params = dict(model.named_parameters())
+    loaded: set[str] = set()
+    for name, weight in weights:
+        for param_name, ckpt_name, shard in _EXPERT_PARAMS:
+            if ckpt_name in name:
+                name = name.replace(ckpt_name, param_name)
+                if is_pp_missing_parameter(name, model):
+                    break
+                param = params[name]
+                for expert_id in range(weight.shape[0]):
+                    param.weight_loader(
+                        param,
+                        weight[expert_id],
+                        name,
+                        shard_id=shard,
+                        expert_id=expert_id,
+                    )
+                loaded.add(name)
+                break
+        else:
+            for param_name, ckpt_name, shard_id in _STACKED_PARAMS:
+                if ckpt_name not in name or ".moe.experts." in name:
+                    continue
+                name = name.replace(ckpt_name, param_name)
+                if is_pp_missing_parameter(name, model):
+                    break
+                if name not in params:
+                    raise ValueError(f"Unsupported Step5 checkpoint tensor: {name}")
+                params[name].weight_loader(params[name], weight, shard_id)
+                loaded.add(name)
+                break
+            else:
+                if is_pp_missing_parameter(name, model):
+                    continue
+                if name not in params:
+                    raise ValueError(f"Unsupported Step5 checkpoint tensor: {name}")
+                param = params[name]
+                getattr(param, "weight_loader", default_weight_loader)(param, weight)
+                loaded.add(name)
+    return loaded
+
+
+@support_torch_compile
+class Step5Model(nn.Module):
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
+        super().__init__()
+        config = vllm_config.model_config.hf_config
+        self.config = config
+        self.moe_num_experts = config.moe_num_experts
+        self.act_dtype = vllm_config.model_config.dtype
+        if config.tie_word_embeddings:
+            raise NotImplementedError("Step5 does not tie word embeddings")
+        if get_pp_group().is_first_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: Step5DecoderLayer(vllm_config, prefix=prefix),
+            prefix=f"{prefix}.layers",
+        )
+        if get_pp_group().is_last_rank:
+            self.norm = GemmaRMSNorm(config.hidden_size, config.rms_norm_eps)
+        else:
+            self.norm = PPMissingLayer()
+        use_fp32_norm_weights(self, config)
+
+    def make_empty_intermediate_tensors(
+        self, batch_size: int, dtype: torch.dtype, device: torch.device
+    ) -> IntermediateTensors:
+        # The residual stream stays FP32 across pipeline stages.
+        return IntermediateTensors(
+            {
+                "hidden_states": torch.zeros(
+                    (batch_size, self.config.hidden_size),
+                    dtype=torch.float32,
+                    device=device,
+                )
+            }
+        )
+
+    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.embed_tokens(input_ids).float()
+
+    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.float()
+            else:
+                hidden_states = self.embed_input_ids(input_ids)
+        else:
+            assert intermediate_tensors is not None
+            hidden_states = intermediate_tensors["hidden_states"]
+        for layer in self.layers[self.start_layer : self.end_layer]:
+            hidden_states = layer(positions, hidden_states)
+        if not get_pp_group().is_last_rank:
+            return IntermediateTensors({"hidden_states": hidden_states})
+        return hidden_states.to(self.act_dtype)
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        return load_step5_weights(self, skip_spec_layers(weights, self.config))
+
+
+class Step5ForCausalLM(nn.Module, SupportsPP, MixtureOfExperts):
+    packed_modules_mapping = {
+        "qkv_proj": ["q_proj", "k_proj", "v_proj"],
+        "gate_up_proj": ["gate_proj", "up_proj"],
+    }
+    hf_to_vllm_mapper = WeightsMapper(
+        orig_to_new_substr={
+            ".share_expert.": ".moe.share_expert.",
+            ".self_attn.sparse_indexer_": ".self_attn.indexer.sparse_indexer_",
+            ".self_attn.ssmax_s": None,
+        },
+        orig_to_new_prefix={"vision_model.": None, "vit_large_projector.": None},
+    )
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
+        super().__init__()
+        config = vllm_config.model_config.hf_config
+        self.model = Step5Model(
+            vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
+        )
+        if get_pp_group().is_last_rank:
+            self.lm_head = ParallelLMHead(
+                config.vocab_size,
+                config.hidden_size,
+                quant_config=vllm_config.quant_config,
+                prefix=maybe_prefix(prefix, "lm_head"),
+            )
+            self.logits_processor = LogitsProcessor(config.vocab_size)
+        else:
+            self.lm_head = PPMissingLayer()
+        self.make_empty_intermediate_tensors = (
+            self.model.make_empty_intermediate_tensors
+        )
+
+        moe = [
+            layer.moe
+            for layer in self.model.layers
+            if isinstance(layer, Step5DecoderLayer) and layer.use_moe
+        ]
+        self.moe_layers = [m.experts for m in moe]
+        self.num_moe_layers = len(moe)
+        self.num_expert_groups = 1
+        self.num_shared_experts = 0
+        # A pipeline stage may hold only dense layers.
+        example = moe[0] if moe else None
+        self.num_logical_experts = example.n_logical_experts if example else 0
+        self.num_physical_experts = example.n_physical_experts if example else 0
+        self.num_local_physical_experts = (
+            example.n_local_physical_experts if example else 0
+        )
+        self.num_routed_experts = example.n_routed_experts if example else 0
+        self.num_redundant_experts = example.n_redundant_experts if example else 0
+
+    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:
+        return self.logits_processor(self.lm_head, self.model.norm(hidden_states))
+
+    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.model.embed_input_ids(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_redundant_experts = num_physical_experts - self.num_logical_experts
+        for layer in self.model.layers:
+            if isinstance(layer, Step5DecoderLayer) and layer.use_moe:
+                layer.moe.n_physical_experts = num_physical_experts
+                layer.moe.n_redundant_experts = self.num_redundant_experts
+                layer.moe.experts.update_expert_map()
+
+    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)
diff --git a/vllm/models/step5/amd/mtp.py b/vllm/models/step5/amd/mtp.py
new file mode 100644
index 000000000000..7d824281a798
--- /dev/null
+++ b/vllm/models/step5/amd/mtp.py
@@ -0,0 +1,240 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Step5 MTP on ROCm."""
+
+from collections.abc import Iterable
+
+import torch
+from torch import nn
+
+from vllm.config import VllmConfig
+from vllm.model_executor.layers.layernorm import GemmaRMSNorm
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
+from vllm.model_executor.layers.vocab_parallel_embedding import (
+    ParallelLMHead,
+    VocabParallelEmbedding,
+)
+from vllm.model_executor.model_loader.mtp_validation import (
+    is_mtp_completeness_check_enabled,
+)
+from vllm.model_executor.models.interfaces import (
+    MultiModalEmbeddings,
+    SupportsMultiModal,
+    _require_is_multimodal,
+)
+from vllm.model_executor.models.utils import (
+    _merge_multimodal_embeddings,
+    get_draft_quant_config,
+    get_spec_layer_idx_from_weight_name,
+    maybe_prefix,
+)
+from vllm.sequence import IntermediateTensors
+
+from .model import (
+    Step5DecoderLayer,
+    Step5ForCausalLM,
+    load_step5_weights,
+    use_fp32_norm_weights,
+)
+
+# Per-MTP-layer weights; everything else in the layer is its decoder block.
+_MTP_LAYER_WEIGHTS = ("embed_tokens", "enorm", "hnorm", "eh_proj", "shared_head")
+
+
+class Step5MTPSharedHead(nn.Module):
+    def __init__(self, config, quant_config, prefix: str) -> None:
+        super().__init__()
+        self.norm = GemmaRMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.head = ParallelLMHead(
+            config.vocab_size,
+            config.hidden_size,
+            quant_config=quant_config,
+            prefix=f"{prefix}.head",
+        )
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.norm(hidden_states)
+
+
+class Step5MultiTokenPredictorLayer(nn.Module):
+    def __init__(self, vllm_config: VllmConfig, prefix: str) -> None:
+        super().__init__()
+        config = vllm_config.model_config.hf_config
+        self.enorm = GemmaRMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.hnorm = GemmaRMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.eh_proj = nn.Linear(config.hidden_size * 2, config.hidden_size, bias=False)
+        self.shared_head = Step5MTPSharedHead(
+            config, vllm_config.quant_config, prefix=f"{prefix}.shared_head"
+        )
+        self.mtp_block = Step5DecoderLayer(vllm_config, prefix=f"{prefix}.mtp_block")
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        previous_hidden_states: torch.Tensor,
+        inputs_embeds: torch.Tensor,
+    ) -> torch.Tensor:
+        hidden_states = self.eh_proj(
+            torch.cat(
+                [self.enorm(inputs_embeds), self.hnorm(previous_hidden_states)],
+                dim=-1,
+            )
+        )
+        return self.mtp_block(positions, hidden_states.float()).to(hidden_states.dtype)
+
+
+class Step5MultiTokenPredictor(nn.Module):
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
+        super().__init__()
+        config = vllm_config.model_config.hf_config
+        self.embed_tokens = VocabParallelEmbedding(
+            config.vocab_size, config.hidden_size
+        )
+        self.mtp_start_layer_idx = config.num_hidden_layers
+        self.num_mtp_layers = config.num_nextn_predict_layers
+        self.layers = nn.ModuleDict(
+            {
+                str(idx): Step5MultiTokenPredictorLayer(
+                    vllm_config, f"{prefix}.layers.{idx}"
+                )
+                for idx in range(
+                    self.mtp_start_layer_idx,
+                    self.mtp_start_layer_idx + self.num_mtp_layers,
+                )
+            }
+        )
+        self.logits_processor = LogitsProcessor(config.vocab_size)
+        use_fp32_norm_weights(self, config)
+
+    def _layer(self, spec_step_idx: int) -> Step5MultiTokenPredictorLayer:
+        idx = self.mtp_start_layer_idx + spec_step_idx % self.num_mtp_layers
+        return self.layers[str(idx)]
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        previous_hidden_states: torch.Tensor,
+        inputs_embeds: torch.Tensor | None = None,
+        spec_step_idx: int = 0,
+    ) -> torch.Tensor:
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+        return self._layer(spec_step_idx)(
+            positions, previous_hidden_states, inputs_embeds
+        )
+
+    def compute_logits(
+        self, hidden_states: torch.Tensor, spec_step_idx: int = 0
+    ) -> torch.Tensor:
+        head = self._layer(spec_step_idx).shared_head
+        return self.logits_processor(head.head, head(hidden_states))
+
+    def get_top_tokens(
+        self, hidden_states: torch.Tensor, spec_step_idx: int = 0
+    ) -> torch.Tensor:
+        head = self._layer(spec_step_idx).shared_head
+        return self.logits_processor.get_top_tokens(head.head, head(hidden_states))
+
+
+class Step5MTP(nn.Module, SupportsMultiModal):
+    # Each MTP layer's shared_head carries its own LM head weights.
+    has_own_lm_head = True
+    supports_multimodal_embeddings = True
+    packed_modules_mapping = Step5ForCausalLM.packed_modules_mapping
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
+        super().__init__()
+        assert vllm_config.speculative_config is not None
+        draft_quant_config = get_draft_quant_config(vllm_config)
+        vllm_config = vllm_config.with_hf_config(
+            vllm_config.speculative_config.draft_model_config.hf_config
+        )
+        vllm_config.quant_config = draft_quant_config
+        self.config = vllm_config.model_config.hf_config
+        self.model = Step5MultiTokenPredictor(
+            vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
+        )
+
+    def embed_input_ids(
+        self,
+        input_ids: torch.Tensor,
+        multimodal_embeddings: MultiModalEmbeddings | None = None,
+        *,
+        is_multimodal: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        inputs_embeds = self._embed_text_input_ids(
+            input_ids, self.model.embed_tokens, is_multimodal=is_multimodal
+        )
+        if not multimodal_embeddings:
+            return inputs_embeds
+        return _merge_multimodal_embeddings(
+            inputs_embeds=inputs_embeds,
+            multimodal_embeddings=multimodal_embeddings,
+            is_multimodal=_require_is_multimodal(is_multimodal),
+        )
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        spec_step_idx: int = 0,
+    ) -> torch.Tensor:
+        return self.model(
+            input_ids, positions, hidden_states, inputs_embeds, spec_step_idx
+        )
+
+    def compute_logits(
+        self, hidden_states: torch.Tensor, spec_step_idx: int = 0
+    ) -> torch.Tensor | None:
+        return self.model.compute_logits(hidden_states, spec_step_idx)
+
+    def get_top_tokens(
+        self, hidden_states: torch.Tensor, spec_step_idx: int = 0
+    ) -> torch.Tensor:
+        return self.model.get_top_tokens(hidden_states, spec_step_idx)
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        def mtp_weights():
+            for name, weight in weights:
+                if name.startswith("model.embed_tokens."):
+                    yield "model." + name.split(".", 1)[1], weight
+                    continue
+                layer = get_spec_layer_idx_from_weight_name(self.config, name)
+                if layer is None:
+                    continue
+                name = name.replace(".transformer.", ".").replace(
+                    "shared_head.output", "shared_head.head"
+                )
+                if not any(w in name for w in _MTP_LAYER_WEIGHTS):
+                    name = name.replace(
+                        f"model.layers.{layer}.", f"model.layers.{layer}.mtp_block."
+                    )
+                yield name, weight
+
+        loaded = {
+            f"model.{name}"
+            for name in load_step5_weights(
+                self.model,
+                (
+                    (name.removeprefix("model."), weight)
+                    for name, weight in mtp_weights()
+                ),
+            )
+        }
+        if is_mtp_completeness_check_enabled():
+            missing = {
+                name
+                for name, _ in self.named_parameters()
+                if not name.endswith(
+                    (".k_scale", ".v_scale", ".q_scale", ".prob_scale")
+                )
+            } - loaded
+            if missing:
+                raise RuntimeError(
+                    f"Step5 MTP checkpoint is missing parameters: {sorted(missing)}"
+                )
+        return loaded
diff --git a/vllm/models/step5/amd/ops/__init__.py b/vllm/models/step5/amd/ops/__init__.py
new file mode 100644
index 000000000000..208f01a7cb5e
--- /dev/null
+++ b/vllm/models/step5/amd/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/step5/amd/ops/csa_paged.py b/vllm/models/step5/amd/ops/csa_paged.py
new file mode 100644
index 000000000000..723eafefa56b
--- /dev/null
+++ b/vllm/models/step5/amd/ops/csa_paged.py
@@ -0,0 +1,202 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Paged Step-5 CSA kernels; each cache has its own block table and page size."""
+
+import torch
+
+from vllm.triton_utils import tl, triton
+
+
+@triton.jit
+def _attention_paged_kernel(
+    q_ptr,
+    k_cache_ptr,
+    v_cache_ptr,
+    regions_ptr,
+    block_table_ptr,
+    req_ptr,
+    pos_ptr,
+    out_ptr,
+    k_scale_ptr,
+    v_scale_ptr,
+    scale,
+    stride_q_token,
+    stride_q_head,
+    stride_k_page,
+    stride_k_row,
+    stride_k_head,
+    stride_v_page,
+    stride_v_row,
+    stride_v_head,
+    stride_r_token,
+    stride_r_group,
+    stride_bt,
+    stride_o_token,
+    stride_o_head,
+    HEADS_PER_GROUP: tl.constexpr,
+    HEAD_DIM: tl.constexpr,
+    PAGE_TOKENS: tl.constexpr,
+    REGION_SIZE: tl.constexpr,
+    SLOTS: tl.constexpr,
+    BLOCK_H: tl.constexpr,
+    BLOCK_D: tl.constexpr,
+    BLOCK_S: tl.constexpr,
+    FP8: tl.constexpr,
+):
+    token = tl.program_id(0)
+    group = tl.program_id(1)
+    req = tl.load(req_ptr + token).to(tl.int64)
+    position = tl.load(pos_ptr + token)
+    if FP8:
+        scale = scale * tl.load(k_scale_ptr)
+
+    heads = tl.arange(0, BLOCK_H)
+    dims = tl.arange(0, BLOCK_D)
+    head_ok = heads < HEADS_PER_GROUP
+    dim_ok = dims < HEAD_DIM
+    q = tl.load(
+        q_ptr
+        + token.to(tl.int64) * stride_q_token
+        + (group * HEADS_PER_GROUP + heads)[:, None] * stride_q_head
+        + dims[None, :],
+        mask=head_ok[:, None] & dim_ok[None, :],
+        other=0.0,
+    )
+    row_max = tl.full((BLOCK_H,), float("-inf"), dtype=tl.float32)
+    row_sum = tl.zeros((BLOCK_H,), dtype=tl.float32)
+    acc = tl.zeros((BLOCK_H, BLOCK_D), dtype=tl.float32)
+    in_region = tl.arange(0, REGION_SIZE)
+
+    for base in tl.range(0, SLOTS, BLOCK_S):
+        slots = base + tl.arange(0, BLOCK_S)
+        ids = tl.load(
+            regions_ptr
+            + token.to(tl.int64) * stride_r_token
+            + group * stride_r_group
+            + slots,
+            mask=slots < SLOTS,
+            other=-1,
+        )
+        cols = tl.reshape(
+            ids[:, None] * REGION_SIZE + in_region[None, :], (BLOCK_S * REGION_SIZE,)
+        )
+        live = tl.reshape(
+            (ids[:, None] >= 0) & (in_region[None, :] >= 0), (BLOCK_S * REGION_SIZE,)
+        )
+        live = live & (cols <= position)
+        safe = tl.where(live, cols, 0)
+        page = tl.load(
+            block_table_ptr + req * stride_bt + safe // PAGE_TOKENS,
+            mask=live,
+            other=0,
+        ).to(tl.int64)
+        row = (safe % PAGE_TOKENS).to(tl.int64)
+        kv_mask = live[:, None] & dim_ok[None, :]
+        k = tl.load(
+            k_cache_ptr
+            + page[:, None] * stride_k_page
+            + row[:, None] * stride_k_row
+            + group * stride_k_head
+            + dims[None, :],
+            mask=kv_mask,
+            other=0.0,
+        )
+        v = tl.load(
+            v_cache_ptr
+            + page[:, None] * stride_v_page
+            + row[:, None] * stride_v_row
+            + group * stride_v_head
+            + dims[None, :],
+            mask=kv_mask,
+            other=0.0,
+        )
+        if FP8:
+            k = k.to(q.dtype)
+            v = v.to(q.dtype)
+        s = tl.dot(q, tl.trans(k), out_dtype=tl.float32) * scale
+        s = tl.where(live[None, :], s, float("-inf"))
+        tile_max = tl.maximum(row_max, tl.max(s, axis=1))
+        empty = tile_max == float("-inf")
+        shift = tl.where(empty, 0.0, tile_max)
+        rescale = tl.where(empty, 1.0, tl.exp(row_max - shift))
+        p = tl.exp(s - shift[:, None])
+        row_sum = row_sum * rescale + tl.sum(p, axis=1)
+        acc = acc * rescale[:, None] + tl.dot(p.to(v.dtype), v, out_dtype=tl.float32)
+        row_max = tile_max
+
+    if FP8:
+        acc = acc * tl.load(v_scale_ptr)
+    out = tl.where(row_sum[:, None] > 0, acc / row_sum[:, None], 0.0)
+    tl.store(
+        out_ptr
+        + token.to(tl.int64) * stride_o_token
+        + (group * HEADS_PER_GROUP + heads)[:, None] * stride_o_head
+        + dims[None, :],
+        out.to(out_ptr.dtype.element_ty),
+        mask=head_ok[:, None] & dim_ok[None, :],
+    )
+
+
+def sparse_attention_paged(
+    query: torch.Tensor,
+    key_cache: torch.Tensor,
+    value_cache: torch.Tensor,
+    regions: torch.Tensor,
+    block_table: torch.Tensor,
+    token_req: torch.Tensor,
+    positions: torch.Tensor,
+    output: torch.Tensor,
+    *,
+    region_size: int,
+    scale: float,
+    k_scale: torch.Tensor | None = None,
+    v_scale: torch.Tensor | None = None,
+) -> None:
+    """Attention over each query's selected regions of the paged K/V, into
+    ``output``. ``key_cache``/``value_cache`` are ``[pages, page_tokens, kv_heads,
+    head_dim]`` views; ``token_req`` and ``positions`` locate every query row."""
+    tokens, heads, head_dim = query.shape
+    if tokens == 0:
+        return
+    groups = key_cache.shape[2]
+    heads_per_group = heads // groups
+    fp8 = key_cache.dtype in (torch.float8_e4m3fn, torch.float8_e4m3fnuz)
+    assert value_cache.dtype == key_cache.dtype
+    assert not fp8 or (k_scale is not None and v_scale is not None)
+    _attention_paged_kernel[(tokens, groups)](
+        query,
+        key_cache,
+        value_cache,
+        regions,
+        block_table,
+        token_req,
+        positions,
+        output,
+        k_scale,
+        v_scale,
+        scale,
+        query.stride(0),
+        query.stride(1),
+        key_cache.stride(0),
+        key_cache.stride(1),
+        key_cache.stride(2),
+        value_cache.stride(0),
+        value_cache.stride(1),
+        value_cache.stride(2),
+        regions.stride(0),
+        regions.stride(1),
+        block_table.stride(0),
+        output.stride(0),
+        output.stride(1),
+        HEADS_PER_GROUP=heads_per_group,
+        HEAD_DIM=head_dim,
+        PAGE_TOKENS=key_cache.shape[1],
+        REGION_SIZE=region_size,
+        SLOTS=regions.shape[-1],
+        BLOCK_H=max(16, triton.next_power_of_2(heads_per_group)),
+        BLOCK_D=triton.next_power_of_2(head_dim),
+        # 8 regions x 8 tokens x 256 padded dims fits gfx942's 64 KiB LDS.
+        BLOCK_S=8,
+        num_stages=1,
+        FP8=fp8,
+    )
diff --git a/vllm/models/step5/amd/vl.py b/vllm/models/step5/amd/vl.py
new file mode 100644
index 000000000000..a00d08ad5023
--- /dev/null
+++ b/vllm/models/step5/amd/vl.py
@@ -0,0 +1,23 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Inference-only Step5 vision-language model on ROCm."""
+
+from vllm.model_executor.models.module_mapping import MultiModelKeys
+from vllm.model_executor.models.step3p7 import Step3p7ForConditionalGeneration
+
+from .model import Step5ForCausalLM
+
+
+class Step5ForConditionalGeneration(Step3p7ForConditionalGeneration):
+    """Step5 VLM using the Step3p7 multimodal frontend and Step5 text model.
+
+    Step5-specific weight names are mapped by the text model's own mapper."""
+
+    packed_modules_mapping = Step5ForCausalLM.packed_modules_mapping
+
+    def get_mm_mapping(self) -> MultiModelKeys:
+        return MultiModelKeys.from_string_field(
+            language_model="language_model",
+            connector="vit_large_projector",
+            tower_model="vision_model",
+        )
diff --git a/vllm/models/step5/amd/warmup.py b/vllm/models/step5/amd/warmup.py
new file mode 100644
index 000000000000..4b0f83f13f95
--- /dev/null
+++ b/vllm/models/step5/amd/warmup.py
@@ -0,0 +1,21 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from typing import TYPE_CHECKING
+
+from vllm.model_executor.warmup.step5_dsa_warmup import (
+    run_step5_scored_prefill_warmup,
+)
+
+from .model import Step5Attention
+
+if TYPE_CHECKING:
+    from vllm.v1.worker.gpu_worker import Worker
+
+
+def step5_csa_warmup(worker: "Worker") -> None:
+    """JIT the CSA scored-prefill kernels before the first long prompt."""
+    context = worker.vllm_config.compilation_config.static_forward_context
+    for layer in context.values():
+        if isinstance(layer, Step5Attention) and layer.indexer is not None:
+            run_step5_scored_prefill_warmup(worker, layer.indexer.topk)
+            return
diff --git a/vllm/models/step5/common/__init__.py b/vllm/models/step5/common/__init__.py
new file mode 100644
index 000000000000..208f01a7cb5e
--- /dev/null
+++ b/vllm/models/step5/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/step5/nvidia/config.py b/vllm/models/step5/common/config.py
similarity index 100%
rename from vllm/models/step5/nvidia/config.py
rename to vllm/models/step5/common/config.py
diff --git a/vllm/models/step5/nvidia/csa_tail_ring.py b/vllm/models/step5/common/csa_tail_ring.py
similarity index 100%
rename from vllm/models/step5/nvidia/csa_tail_ring.py
rename to vllm/models/step5/common/csa_tail_ring.py
diff --git a/vllm/models/step5/common/ops/__init__.py b/vllm/models/step5/common/ops/__init__.py
new file mode 100644
index 000000000000..208f01a7cb5e
--- /dev/null
+++ b/vllm/models/step5/common/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/step5/nvidia/ops/indexer_norm_rope.py b/vllm/models/step5/common/ops/indexer_norm_rope.py
similarity index 100%
rename from vllm/models/step5/nvidia/ops/indexer_norm_rope.py
rename to vllm/models/step5/common/ops/indexer_norm_rope.py
diff --git a/vllm/models/step5/nvidia/ops/prefill_indexer_metadata.py b/vllm/models/step5/common/ops/prefill_indexer_metadata.py
similarity index 100%
rename from vllm/models/step5/nvidia/ops/prefill_indexer_metadata.py
rename to vllm/models/step5/common/ops/prefill_indexer_metadata.py
diff --git a/vllm/models/step5/nvidia/ops/sparse_decode_attention.py b/vllm/models/step5/common/ops/sparse_decode_attention.py
similarity index 96%
rename from vllm/models/step5/nvidia/ops/sparse_decode_attention.py
rename to vllm/models/step5/common/ops/sparse_decode_attention.py
index c0c94e401234..8ef3e0590433 100644
--- a/vllm/models/step5/nvidia/ops/sparse_decode_attention.py
+++ b/vllm/models/step5/common/ops/sparse_decode_attention.py
@@ -11,11 +11,13 @@
 import torch
 
 from vllm.platforms import current_platform
-from vllm.triton_utils import tl, triton
+from vllm.triton_utils import tl, tldevice, triton
 
 # PDL: kernels launch during the producer's tail and gdc_wait before reading
 # its outputs. No early trigger, so their writes stay visible downstream.
 _PDL = current_platform.is_arch_support_pdl()
+# The AMD Triton backend rejects the launch_pdl keyword, even when False.
+_LAUNCH = {"launch_pdl": True} if _PDL else {}
 
 
 @triton.jit(do_not_specialize=["SCALE"])
@@ -234,7 +236,7 @@ def _gate(Gate, token, head, mask, GATED: tl.constexpr, GATE_ROW, GATE_HEAD):
         gate = tl.load(
             Gate + token.to(tl.int64) * GATE_ROW + head * GATE_HEAD, mask, other=0.0
         ).to(tl.float32)
-        sigmoid = tl.div_rn(1.0, 1.0 + tl.extra.cuda.libdevice.exp(-gate))
+        sigmoid = tl.div_rn(1.0, 1.0 + tldevice.exp(-gate))
         return sigmoid.to(Gate.dtype.element_ty).to(tl.float32)
     return tl.full(head.shape, 1.0, tl.float32)
 
@@ -283,10 +285,11 @@ def sparse_decode_attention(
     assert dim == 192 and cache.shape[-1] == 384 and heads <= 16
     assert query.stride(2) == output.stride(2) == cache.stride(3) == 1
     assert regions.stride(1) == 1 and rows == regions.shape[0]
-    fp8 = cache.dtype in (torch.uint8, torch.float8_e4m3fn)
+    fp8 = cache.dtype in (torch.uint8, torch.float8_e4m3fn, torch.float8_e4m3fnuz)
     if fp8:
         assert k_scale is not None and v_scale is not None
-        cache = cache.view(torch.float8_e4m3fn)
+        if cache.dtype == torch.uint8:
+            cache = cache.view(current_platform.fp8_dtype())
     if not rows:
         return
     partial = torch.empty(
@@ -332,7 +335,7 @@ def sparse_decode_attention(
         fp8,
         _PDL,
         num_warps=4,
-        launch_pdl=_PDL,
+        **_LAUNCH,
     )
     if splits > 1:
         _sparse_decode_combine[(rows, heads)](
@@ -353,5 +356,5 @@ def sparse_decode_attention(
             gate.stride(1) if gate is not None else 0,
             _PDL,
             num_warps=4,
-            launch_pdl=_PDL,
+            **_LAUNCH,
         )
diff --git a/vllm/models/step5/nvidia/ops/sparse_indexer.py b/vllm/models/step5/common/ops/sparse_indexer.py
similarity index 100%
rename from vllm/models/step5/nvidia/ops/sparse_indexer.py
rename to vllm/models/step5/common/ops/sparse_indexer.py
diff --git a/vllm/models/step5/nvidia/ops/sparse_indexer_scores.py b/vllm/models/step5/common/ops/sparse_indexer_scores.py
similarity index 100%
rename from vllm/models/step5/nvidia/ops/sparse_indexer_scores.py
rename to vllm/models/step5/common/ops/sparse_indexer_scores.py
diff --git a/vllm/models/step5/nvidia/ops/summary_update.py b/vllm/models/step5/common/ops/summary_update.py
similarity index 99%
rename from vllm/models/step5/nvidia/ops/summary_update.py
rename to vllm/models/step5/common/ops/summary_update.py
index 12a9dafd4066..85aded03fa5a 100644
--- a/vllm/models/step5/nvidia/ops/summary_update.py
+++ b/vllm/models/step5/common/ops/summary_update.py
@@ -12,9 +12,9 @@
 
 import torch
 
-from vllm.triton_utils import tl, triton
+from vllm.triton_utils import tl, tldevice, triton
 
-libdevice = tl.extra.cuda.libdevice
+libdevice = tldevice
 
 
 @triton.jit
diff --git a/vllm/models/step5/nvidia/sparse_summary_cache.py b/vllm/models/step5/common/sparse_summary_cache.py
similarity index 100%
rename from vllm/models/step5/nvidia/sparse_summary_cache.py
rename to vllm/models/step5/common/sparse_summary_cache.py
diff --git a/vllm/models/step5/nvidia/flashinfer_sparse.py b/vllm/models/step5/nvidia/flashinfer_sparse.py
index f0c963650785..bf3862ccd61b 100644
--- a/vllm/models/step5/nvidia/flashinfer_sparse.py
+++ b/vllm/models/step5/nvidia/flashinfer_sparse.py
@@ -7,17 +7,17 @@
 
 from vllm.config import get_current_vllm_config_or_none
 
-from .ops.prefill_indexer_metadata import prepare_prefill_indexer
-from .ops.sparse_decode_attention import (
+from ..common.ops.prefill_indexer_metadata import prepare_prefill_indexer
+from ..common.ops.sparse_decode_attention import (
     sparse_decode_attention,
     sparse_decode_splits,
 )
-from .ops.sparse_indexer import (
+from ..common.ops.sparse_indexer import (
     pack_decode_summary_pages,
     pack_summary_pages,
     select_regions,
 )
-from .ops.sparse_indexer_scores import decode_summary_scores
+from ..common.ops.sparse_indexer_scores import decode_summary_scores
 from .ops.sparse_prefill import PrefillChunk, gather_prefill_kv
 from .ops.sparse_prefill_pages import CompactPagedPrefill
 
@@ -44,7 +44,7 @@ def __init__(self, device: torch.device):
             config.kernel_config.sparse_indexer_topk_backend if config else "auto"
         )
         if torch.cuda.get_device_capability(device) == (9, 0):
-            from .ops.sparse_indexer_scores import paged_mqa_logits_sm90
+            from ..common.ops.sparse_indexer_scores import paged_mqa_logits_sm90
 
             self._logits = paged_mqa_logits_sm90
         else:
@@ -171,7 +171,7 @@ def warmup_prefill_scoring(self, max_regions, topk, paired_heads, device):
         """
         from flashinfer.attn_scores import compute_paged_mqa_logits_schedule
 
-        from .ops.prefill_indexer_metadata import prefill_indexer_chunk_size
+        from ..common.ops.prefill_indexer_metadata import prefill_indexer_chunk_size
 
         rows = torch.arange(2, device=device)
         owners = torch.zeros(2, dtype=torch.long, device=device)
diff --git a/vllm/models/step5/nvidia/model.py b/vllm/models/step5/nvidia/model.py
index ea6b1b7273bf..3414103347db 100644
--- a/vllm/models/step5/nvidia/model.py
+++ b/vllm/models/step5/nvidia/model.py
@@ -86,7 +86,7 @@
 from vllm.sequence import IntermediateTensors
 from vllm.v1.attention.backend import AttentionBackend, AttentionType
 
-from .config import (
+from ..common.config import (
     Step5SparseConfig,
     checkpoint_has_step5_sparse_config,
     get_step5_sparse_config,
@@ -1208,8 +1208,8 @@ def __init__(
             **sparse_impl_args,
         )
         if self.use_dsa_backend:
-            from .csa_tail_ring import Step5CSATailRing
-            from .sparse_summary_cache import Step5SparseSummaryCacheLayer
+            from ..common.csa_tail_ring import Step5CSATailRing
+            from ..common.sparse_summary_cache import Step5SparseSummaryCacheLayer
 
             sparse_config = typing.cast(Step5SparseConfig, sparse_config)
             proxy_dim = int(sparse_config.proxy_dim)
@@ -1644,7 +1644,7 @@ def _project_sparse_indexer(
         torch.Tensor,
         bool,
     ]:
-        from .ops.indexer_norm_rope import indexer_norm_rope
+        from ..common.ops.indexer_norm_rope import indexer_norm_rope
 
         owner = self.owner
         index_q, index_k, index_z, extra_dims = owner.qkv_indexer_proj.split_indexer(
diff --git a/vllm/models/step5/nvidia/ops/qkv_indexer_norm_rope.py b/vllm/models/step5/nvidia/ops/qkv_indexer_norm_rope.py
index 05dbdf27e1d3..d07093efe4ed 100644
--- a/vllm/models/step5/nvidia/ops/qkv_indexer_norm_rope.py
+++ b/vllm/models/step5/nvidia/ops/qkv_indexer_norm_rope.py
@@ -6,8 +6,8 @@
 
 from vllm.triton_utils import tl, triton
 
-from .indexer_norm_rope import _indexer_norm_rope
-from .summary_update import summary_layout
+from ...common.ops.indexer_norm_rope import _indexer_norm_rope
+from ...common.ops.summary_update import summary_layout
 
 
 @triton.jit(do_not_specialize=["SLOT_ROWS"])
diff --git a/vllm/models/step5/nvidia/ops/sparse_prefill_pages.py b/vllm/models/step5/nvidia/ops/sparse_prefill_pages.py
index a3921a0aba17..d1d40db05153 100644
--- a/vllm/models/step5/nvidia/ops/sparse_prefill_pages.py
+++ b/vllm/models/step5/nvidia/ops/sparse_prefill_pages.py
@@ -9,7 +9,7 @@
 
 from vllm.triton_utils import tl, triton
 
-from .sparse_indexer import PagedPrefillMetadata
+from ...common.ops.sparse_indexer import PagedPrefillMetadata
 
 
 @triton.jit(do_not_specialize=["REGION_COUNT", "REGION_ROW"])
diff --git a/vllm/models/step5/nvidia/ops/swa_decode.py b/vllm/models/step5/nvidia/ops/swa_decode.py
index 4ea52520ea33..2e95b5da956b 100644
--- a/vllm/models/step5/nvidia/ops/swa_decode.py
+++ b/vllm/models/step5/nvidia/ops/swa_decode.py
@@ -12,7 +12,7 @@
 
 from vllm.triton_utils import tl, triton
 
-from .sparse_decode_attention import _PDL, _gate
+from ...common.ops.sparse_decode_attention import _PDL, _gate
 
 
 @triton.jit(do_not_specialize=["SCALE"])
diff --git a/vllm/models/step5/nvidia/sparse_attention.py b/vllm/models/step5/nvidia/sparse_attention.py
index 07d0d2da2d90..1f231f8097a3 100644
--- a/vllm/models/step5/nvidia/sparse_attention.py
+++ b/vllm/models/step5/nvidia/sparse_attention.py
@@ -22,6 +22,7 @@
 from vllm.v1.kv_cache_interface import AttentionSpec, is_quantized_kv_cache
 from vllm.v1.kv_cache_layout import KVCacheLayout
 
+from ..common.ops.summary_update import build_prefill_groups
 from .ops.attention_gate import headwise_sigmoid_gate
 from .ops.sparse_prefill import make_prefill_chunks
 from .ops.sparse_prefill_pages import (
@@ -29,7 +30,6 @@
     plan_staged_prefill_pages,
     stage_prefill_pages,
 )
-from .ops.summary_update import build_prefill_groups
 
 # Transient dense BF16 copy of a prefill batch's live pages for one KV group.
 _PREFILL_STAGING_BYTES = 512 << 20
diff --git a/vllm/multimodal/cache/base.py b/vllm/multimodal/cache/base.py
index f613325f8550..287cf7d5823d 100644
--- a/vllm/multimodal/cache/base.py
+++ b/vllm/multimodal/cache/base.py
@@ -345,6 +345,14 @@ def is_cached(self, mm_hashes: list[str]) -> list[bool]:
         """
         return [self.is_cached_item(mm_hash) for mm_hash in mm_hashes]
 
+    def get_cached_item(self, mm_hash: str) -> MultiModalProcessorCacheInItem:
+        """Return a locally stored payload without updating eviction order.
+
+        Metadata-only and shared-memory sender caches do not retain a local
+        payload, so they return ``None``.
+        """
+        return None
+
     def invalidate(self, mm_hash: str) -> None:
         """Drop ``mm_hash`` from this P0 shadow cache to recover from P0/P1 drift.
 
@@ -416,6 +424,12 @@ def __init__(self, model_config: ModelConfig) -> None:
     def is_cached_item(self, mm_hash: str) -> bool:
         return mm_hash in self._cache
 
+    @override
+    def get_cached_item(self, mm_hash: str) -> MultiModalProcessorCacheInItem:
+        if (cached_item := self._cache.peek(mm_hash)) is not None:
+            return cached_item.item, cached_item.prompt_updates
+        return None
+
     @override
     def get_and_update_item(
         self,
diff --git a/vllm/multimodal/processing/processor.py b/vllm/multimodal/processing/processor.py
index 5b46267dff54..ea3cf450c1c3 100644
--- a/vllm/multimodal/processing/processor.py
+++ b/vllm/multimodal/processing/processor.py
@@ -1416,6 +1416,14 @@ def _merge_mm_kwargs(
         mm_missing_kwargs: MultiModalKwargsItems,
         mm_missing_prompt_updates: MultiModalPromptUpdates,
     ) -> tuple[MultiModalKwargsOptionalItems, MultiModalPromptUpdates]:
+        # A request can exceed the cache capacity. Retain local payloads before
+        # inserting misses can evict hits that occur later in this request.
+        cached_items = {
+            item_hash: cache.get_cached_item(item_hash)
+            for hashes in mm_hashes.values()
+            for item_hash in hashes
+        }
+
         # Need to touch all mm hashes before update to avoid hash in updated
         # list evict during update
         for hashes in mm_hashes.values():
@@ -1433,6 +1441,7 @@ def _merge_mm_kwargs(
             missing_prompt_updates = mm_missing_prompt_updates.get(modality, [])
 
             for item_idx, item_hash in enumerate(hashes):
+                item = cached_items.get(item_hash)
                 if not mm_is_cached[modality][item_idx]:
                     missing_next_idx = mm_missing_next_idx[modality]
                     missing_kwargs_item = missing_kwargs[missing_next_idx]
@@ -1441,8 +1450,6 @@ def _merge_mm_kwargs(
                     mm_missing_next_idx[modality] += 1
 
                     item = missing_kwargs_item, missing_updates_item
-                else:
-                    item = None
 
                 kwargs, updates = cache.get_and_update_item(item, item_hash)
 
diff --git a/vllm/multimodal/video_decoders/pynvvideocodec.py b/vllm/multimodal/video_decoders/pynvvideocodec.py
index 1d5b5d72e57b..83f4107f701a 100644
--- a/vllm/multimodal/video_decoders/pynvvideocodec.py
+++ b/vllm/multimodal/video_decoders/pynvvideocodec.py
@@ -5,7 +5,8 @@
 import tempfile
 import threading
 from contextlib import contextmanager, suppress
-from typing import ClassVar, NamedTuple
+from functools import lru_cache
+from typing import Any, ClassVar, NamedTuple
 
 import numpy as np
 import numpy.typing as npt
@@ -112,6 +113,24 @@ def _decoder_surface_bytes(width: int, height: int) -> int:
     return width * height * 3 * PYNVVIDEOCODEC_DECODER_CACHE_SIZE
 
 
+# Distinct (device, codec, chroma format, bit depth) combinations; real deployments
+# see a handful, so this bound only guards against unbounded growth.
+PYNVVIDEOCODEC_DECODER_CAPS_CACHE_SIZE = 64
+
+
+@lru_cache(maxsize=PYNVVIDEOCODEC_DECODER_CAPS_CACHE_SIZE)
+def _get_decoder_caps(
+    nvc, device_index: int, codec, chroma_format, bit_depth: int
+) -> dict[str, Any]:
+    """``nvc.GetDecoderCaps``, cached per device and stream format.
+
+    Each call takes ~180 ms and holds the GIL, which stalls every other request in
+    the API process; the result only depends on these arguments. Callers must not
+    mutate the returned dict.
+    """
+    return nvc.GetDecoderCaps(device_index, codec, chroma_format, bit_depth)
+
+
 def _pynvvideocodec_exception_types(nvc) -> tuple[type[Exception], ...]:
     return tuple(
         exception_type
@@ -343,7 +362,8 @@ def _validate_decoder_caps(
         demuxer = nvc.CreateDemuxer(file_path)
         width = demuxer.Width()
         height = demuxer.Height()
-        caps = nvc.GetDecoderCaps(
+        caps = _get_decoder_caps(
+            nvc,
             cls._DEVICE_INDEX,
             demuxer.GetNvCodecId(),
             demuxer.ChromaFormat(),
diff --git a/vllm/parser/deepseek_v32.py b/vllm/parser/deepseek_v32.py
index de6634cf7368..ae5f33a09961 100644
--- a/vllm/parser/deepseek_v32.py
+++ b/vllm/parser/deepseek_v32.py
@@ -28,6 +28,7 @@
     DSML_PARAM_CLOSE,
     DSML_PARAM_START,
     _dsml_arg_converter,
+    _dsml_arg_value_extender,
     _unwrap_wrapper_args,
 )
 from vllm.parser.engine.events import EventType
@@ -115,6 +116,7 @@ def deepseek_v32_config() -> ParserEngineConfig:
         },
         arg_converter=_dsml_arg_converter,
         arg_structural_chars=frozenset(">"),
+        arg_value_extender=_dsml_arg_value_extender,
         strip_content_whitespace_with_tools=False,
         tool_args_json=False,
     )
diff --git a/vllm/parser/deepseek_v4.py b/vllm/parser/deepseek_v4.py
index 49ef63736531..f18b4ad46f9b 100644
--- a/vllm/parser/deepseek_v4.py
+++ b/vllm/parser/deepseek_v4.py
@@ -28,6 +28,7 @@
 from vllm.parser.engine.events import EventType
 from vllm.parser.engine.parser_engine import ParserEngine
 from vllm.parser.engine.parser_engine_config import (
+    ArgValueExtender,
     ParserEngineConfig,
     ParserState,
     Transition,
@@ -64,24 +65,46 @@ def _param_patterns(
 
     The ``string`` attribute is optional: the model sometimes omits it, and
     dropping such a parameter hands the client a tool call with no arguments.
+    The trailing-partial regex captures the open value up to the end of the
+    text; ``_match_open_param`` holds back an unfinished next parameter tag
+    so a streamed string value never includes it.
     """
     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),
+        re.compile(rf"{head}(.*)\Z", re.DOTALL),
     )
 
 
 _PARAM_RE, _PARTIAL_PARAM_RE = _param_patterns(DSML_PARAM_START, DSML_PARAM_CLOSE)
 
 
+def _match_open_param(
+    raw_args: str, pos: int, partial_param_re: re.Pattern, param_start: str
+) -> tuple[str, str | None, str, bool] | None:
+    """Return (name, is_str, value, held_back) for the trailing open param.
+
+    An unfinished next parameter tag (``param_start`` with no ``>`` after it)
+    is cut from the value; ``held_back`` reports whether that happened.
+    """
+    pm = partial_param_re.search(raw_args, pos)
+    if pm is None:
+        return None
+    value = pm.group(3)
+    cut = value.find(param_start, value.rfind(">") + 1)
+    if cut >= 0:
+        return pm.group(1), pm.group(2), value[:cut], True
+    return pm.group(1), pm.group(2), value, False
+
+
 def _dsml_arg_converter(
     raw_args: str,
     partial: bool,
     *,
     param_re: re.Pattern = _PARAM_RE,
     partial_param_re: re.Pattern = _PARTIAL_PARAM_RE,
+    param_start: str = DSML_PARAM_START,
 ) -> str:
     params: dict[str, object] = {}
 
@@ -100,9 +123,11 @@ def _dsml_arg_converter(
         last_end = m.end()
 
     if partial:
-        pm = partial_param_re.search(raw_args, last_end)
-        if pm:
-            name, is_str, value = pm.group(1), pm.group(2), pm.group(3)
+        open_param = _match_open_param(
+            raw_args, last_end, partial_param_re, param_start
+        )
+        if open_param:
+            name, is_str, value, _ = open_param
             if is_str == "true":
                 params[name] = value
             else:
@@ -112,6 +137,44 @@ def _dsml_arg_converter(
     return json.dumps(params, ensure_ascii=False)
 
 
+def _dsml_arg_value_extender(
+    raw_args: str,
+    *,
+    param_re: re.Pattern = _PARAM_RE,
+    partial_param_re: re.Pattern = _PARTIAL_PARAM_RE,
+    param_start: str = DSML_PARAM_START,
+    param_close: str = DSML_PARAM_CLOSE,
+) -> ArgValueExtender | None:
+    # ``arg_structural_chars`` skips every delta without ``>``, so a long
+    # string value would wait for its closing tag. Extend the open trailing
+    # ``string="true"`` value instead. The lexer can flush a held partial
+    # tag as text, so keep a tail to catch a tag split across deltas.
+    names = set()
+    last_end = 0
+    for m in param_re.finditer(raw_args):
+        names.add(m.group(1))
+        last_end = m.end()
+    open_param = _match_open_param(raw_args, last_end, partial_param_re, param_start)
+    if open_param is None:
+        return None
+    name, is_str, _, held_back = open_param
+    if is_str != "true" or name in names or held_back:
+        return None
+
+    keep = max(len(param_start), len(param_close)) - 1
+    tail = raw_args[-keep:]
+
+    def extend(delta: str) -> str | None:
+        nonlocal tail
+        text = tail + delta
+        if param_start in text or param_close in text:
+            return None
+        tail = text[-keep:]
+        return json.dumps(delta, ensure_ascii=False)[1:-1]
+
+    return extend
+
+
 def _unwrap_wrapper_args(
     args_json: str,
     tools: list[Tool] | None,
@@ -238,6 +301,7 @@ def deepseek_v4_config(thinking: bool = False) -> ParserEngineConfig:
         },
         arg_converter=_dsml_arg_converter,
         arg_structural_chars=frozenset(">"),
+        arg_value_extender=_dsml_arg_value_extender,
         strip_content_whitespace_with_tools=False,
         tool_args_json=False,
     )
diff --git a/vllm/parser/deepseek_v41.py b/vllm/parser/deepseek_v41.py
index 84216005102a..76c7ad19bfd3 100644
--- a/vllm/parser/deepseek_v41.py
+++ b/vllm/parser/deepseek_v41.py
@@ -8,6 +8,7 @@
 from vllm.parser.deepseek_v4 import (
     DeepSeekV4Parser,
     _dsml_arg_converter,
+    _dsml_arg_value_extender,
     _param_patterns,
     deepseek_v4_config,
 )
@@ -46,6 +47,14 @@ def deepseek_v41_config(thinking: bool = False) -> ParserEngineConfig:
             _dsml_arg_converter,
             param_re=_PARAM_RE,
             partial_param_re=_PARTIAL_PARAM_RE,
+            param_start=DSML_PARAM_START,
+        ),
+        arg_value_extender=functools.partial(
+            _dsml_arg_value_extender,
+            param_re=_PARAM_RE,
+            partial_param_re=_PARTIAL_PARAM_RE,
+            param_start=DSML_PARAM_START,
+            param_close=DSML_PARAM_CLOSE,
         ),
     )
 
diff --git a/vllm/parser/engine/parser_engine.py b/vllm/parser/engine/parser_engine.py
index 22ca626f5afd..5d45902c1fa2 100644
--- a/vllm/parser/engine/parser_engine.py
+++ b/vllm/parser/engine/parser_engine.py
@@ -1132,6 +1132,11 @@ def _compute_arg_delta(self, idx: int, raw_delta: str) -> str | None:
         if not safe_json:
             return None
         if prev and not safe_json.startswith(prev):
+            if self._arg_value_extender is not None:
+                # A repeated key replaced a streamed value. Deltas that only
+                # extend the new value skip the conversion until it closes.
+                slot.extend_open_value = self._arg_value_extender(slot.args)
+                slot.open_value_streams = False
             return None
 
         slot.streamed_json = safe_json
diff --git a/vllm/parser/glm47_moe.py b/vllm/parser/glm47_moe.py
index 44377e62b707..6f2a6321ca35 100644
--- a/vllm/parser/glm47_moe.py
+++ b/vllm/parser/glm47_moe.py
@@ -64,18 +64,47 @@
 )
 
 
+class _ClosedArgs:
+    """Closed key/value pairs of a raw argument text, reused by longer texts.
+
+    Raw arguments only grow while a call streams, so a parse of a text that
+    extends ``text`` resumes after its last closed value.
+    """
+
+    __slots__ = ("text", "params", "outside")
+
+    def __init__(self) -> None:
+        self.text = ""  # raw text up to the end of the last closed value
+        self.params: dict[str, object] = {}
+        self.outside = ""  # text of ``self.text`` outside the closed pairs
+
+
 def _glm47_parse_args(
-    raw_args: str, partial: bool
+    raw_args: str, partial: bool, closed: _ClosedArgs | None = None
 ) -> tuple[dict[str, object], str | None]:
     """Return the arguments and the key of the open trailing value, if any."""
     params: dict[str, object] = {}
-
-    for match in _ARG_RE.finditer(raw_args):
+    outside: list[str] = []
+    pos = 0
+    if closed is not None and raw_args.startswith(closed.text):
+        params.update(closed.params)
+        outside.append(closed.outside)
+        pos = len(closed.text)
+    start = pos
+
+    for match in _ARG_RE.finditer(raw_args, pos):
+        outside.append(raw_args[pos : match.start()])
         params[match.group("key").strip()] = match.group("value")
+        pos = match.end()
+
+    if closed is not None and pos > start:
+        closed.text = raw_args[:pos]
+        closed.params = dict(params)
+        closed.outside = "".join(outside)
 
     open_key = None
     if partial:
-        remaining = _ARG_RE.sub("", raw_args)
+        remaining = "".join(outside) + raw_args[pos:]
         match = _PARTIAL_ARG_RE.search(remaining)
         if match:
             key = match.group("key").strip()
@@ -86,14 +115,20 @@ def _glm47_parse_args(
     return params, open_key
 
 
-def _glm47_arg_converter(raw_args: str, partial: bool) -> str:
-    return json.dumps(_glm47_parse_args(raw_args, partial)[0], ensure_ascii=False)
+def _glm47_arg_converter(
+    raw_args: str, partial: bool, closed: _ClosedArgs | None = None
+) -> str:
+    return json.dumps(
+        _glm47_parse_args(raw_args, partial, closed)[0], ensure_ascii=False
+    )
 
 
-def _glm47_arg_value_extender(raw_args: str) -> ArgValueExtender | None:
+def _glm47_arg_value_extender(
+    raw_args: str, closed: _ClosedArgs | None = None
+) -> ArgValueExtender | None:
     # The open value runs to the end of raw_args and is the last JSON value,
     # so a delta appends its JSON-escaped text unless it closes the value.
-    params, open_key = _glm47_parse_args(raw_args, partial=True)
+    params, open_key = _glm47_parse_args(raw_args, partial=True, closed=closed)
     if open_key is None or next(reversed(params)) != open_key:
         return None
 
@@ -225,6 +260,7 @@ def glm47_moe_config(thinking: bool = True) -> ParserEngineConfig:
         },
         turn_boundary_tokens=GLM_TURN_BOUNDARIES,
         arg_converter=_glm47_arg_converter,
+        arg_structural_chars=frozenset(">"),
         arg_value_extender=_glm47_arg_value_extender,
         stream_arg_deltas=True,
         tool_args_json=False,
@@ -264,6 +300,13 @@ def __init__(
             glm47_moe_config(thinking=self.thinking_enabled),
         )
         super().__init__(tokenizer, tools, **kwargs)
+        # Each streaming conversion resumes after the last closed pair.
+        if self._arg_converter is _glm47_arg_converter:
+            closed = _ClosedArgs()
+            self._arg_converter = functools.partial(_glm47_arg_converter, closed=closed)
+            self._arg_value_extender = functools.partial(
+                _glm47_arg_value_extender, closed=closed
+            )
 
     def _emit_name_delta(self, idx: int, deltas, name: str | None) -> None:
         if name is not None:
diff --git a/vllm/parser/mistral.py b/vllm/parser/mistral.py
index f6d3435bd7c5..74d76bfa15b9 100644
--- a/vllm/parser/mistral.py
+++ b/vllm/parser/mistral.py
@@ -13,20 +13,12 @@
 
 import ijson
 import regex as re
-from mistral_common.protocol.instruct.tool_calls import (
-    NamedToolChoice as MistralNamedToolChoice,
-)
-from mistral_common.protocol.instruct.tool_calls import (
-    Tool as MistralTool,
-)
-from mistral_common.protocol.instruct.tool_calls import (
-    ToolChoice as MistralToolChoice,
-)
-from mistral_common.protocol.instruct.tool_calls import (
-    ToolChoiceEnum as MistralToolChoiceEnum,
+from mistral_common.exceptions import InvalidRequestException
+from mistral_common.protocol.instruct.request import (
+    ChatCompletionRequest as MistralChatCompletionRequest,
 )
 from mistral_common.tokens.tokenizers.base import SpecialTokens
-from pydantic import Field
+from pydantic import Field, ValidationError
 
 from vllm.entrypoints.generate.base.protocol import (
     DeltaFunctionCall,
@@ -40,6 +32,7 @@
     ChatCompletionNamedToolChoiceParam,
 )
 from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
+from vllm.exceptions import VLLMValidationError
 from vllm.logger import init_logger
 from vllm.parser.engine.events import EventType
 from vllm.parser.engine.parser_engine import ParserEngine
@@ -63,9 +56,6 @@
 logger = init_logger(__name__)
 
 _ALPHANUMERIC = ascii_letters + digits
-_DEFAULT_JSON_SCHEMA: dict[str, Any] = {
-    "anyOf": [{"type": "object"}, {"type": "array"}]
-}
 
 # Special token strings as emitted/decoded by MistralTokenizer.
 _TOOL_CALLS = SpecialTokens.tool_calls.value
@@ -110,6 +100,41 @@ def is_valid_id(id: str) -> bool:
         return id.isalnum() and len(id) == 9
 
 
+def _mistral_response_format(request: ChatCompletionRequest) -> dict[str, Any]:
+    """Return the OpenAI-format response format to pass to mistral-common.
+
+    `structured_outputs` (`json`, `json_object`) takes precedence over
+    `response_format`, as in `structured_outputs_from_response_format`.
+    """
+    structured_outputs = request.structured_outputs
+    if structured_outputs is not None:
+        if structured_outputs.json_object is not None:
+            return {"type": "json_object"}
+        if structured_outputs.json is None:
+            raise ValueError(
+                "Unsupported request.structured_outputs for MistralParser. "
+                "Only `json` and `json_object` are supported."
+            )
+        schema = structured_outputs.json
+        if isinstance(schema, str):
+            check_json_nesting(schema)
+            schema = json.loads(schema)
+        # `structured_outputs.json` asks vLLM to enforce the schema.
+        return {
+            "type": "json_schema",
+            "json_schema": {"name": "response", "schema": schema, "strict": True},
+        }
+
+    if request.response_format is None:
+        return {"type": "text"}
+    if request.response_format.type not in ("text", "json_object", "json_schema"):
+        raise ValueError(
+            "MistralParser only accepts `text`, `json_object` or "
+            f"`json_schema`, got {request.response_format=}"
+        )
+    return request.response_format.model_dump(by_alias=True, exclude_none=True)
+
+
 def _is_pre_v11_tokeniser(model_tokenizer: TokenizerLike) -> bool:
     if is_mistral_tokenizer(model_tokenizer):
         return model_tokenizer.version < 11
@@ -371,11 +396,9 @@ def adjust_request(
             or any_so_non_supported_active
             or response_format_non_supported_active
         ):
-            request = super().adjust_request(request)
-            if request.tools and request.tool_choice != "none":
-                # Keep special tokens so the [TOOL_CALLS] marker
-                # survives for tool detection.
-                request.skip_special_tokens = False
+            # MistralTokenizer keeps the parser's special tokens when skipping.
+            if not is_mistral_tokenizer(self.model_tokenizer):
+                request = super().adjust_request(request)
             # Inject a guided JSON schema for pre-v11 required/named tool choice
             # so the model emits a well-formed bare JSON array instead of
             # rambling. tool_choice forces a tool call, so a json_object /
@@ -407,82 +430,32 @@ def adjust_request(
                     request.response_format = None
             return request
 
-        json_schema: dict[str, Any] | None = None
-        if request.structured_outputs is not None:
-            if request.structured_outputs.json_object is not None:
-                json_schema = _DEFAULT_JSON_SCHEMA
-            elif request.structured_outputs.json is not None:
-                if isinstance(request.structured_outputs.json, str):
-                    check_json_nesting(request.structured_outputs.json)
-                    json_schema = json.loads(request.structured_outputs.json)
-                else:
-                    json_schema = request.structured_outputs.json
-            else:
-                raise ValueError(
-                    "Unsupported request.structured_outputs for MistralParser. "
-                    "Only `json` and `json_object` are supported."
-                )
-        elif (
-            request.response_format is not None
-            and request.response_format.type != "text"
-        ):
-            if request.response_format.type == "json_object":
-                json_schema = _DEFAULT_JSON_SCHEMA
-            elif request.response_format.type == "json_schema":
-                if request.response_format.json_schema is not None:
-                    json_schema = request.response_format.json_schema.json_schema
-                else:
-                    json_schema = _DEFAULT_JSON_SCHEMA
-            else:
-                raise ValueError(
-                    "MistralParser only accepts `text`, `json_object` or "
-                    f"`json_schema`, got {request.response_format=}"
-                )
-            request.response_format = None
+        response_format = _mistral_response_format(request)
+        # The grammar replaces response_format as the engine constraint.
+        request.response_format = None
 
         grammar_factory = self.model_tokenizer.grammar_factory
 
         # Rendering grammar is cached in mistral-common given tools, template and mode.
         template = grammar_factory.select_jinja_template()
 
-        mistral_tools = (
-            [MistralTool.from_openai(tool.model_dump()) for tool in request.tools]
-            if request.tools is not None
-            else None
-        )
-
-        tool_choice: MistralToolChoice
-        match request.tool_choice:
-            case "none" | "auto" | "required":
-                tool_choice = MistralToolChoiceEnum(request.tool_choice)
-            case None:
-                tool_choice = MistralToolChoiceEnum.auto
-            # _ == Named tool choice
-            case _:
-                tool_choice = MistralNamedToolChoice.model_validate(
-                    {
-                        "type": "function",
-                        "function": {"name": request.tool_choice.function.name},
-                    }
-                )
-
-        match tool_choice, json_schema is not None:
-            case MistralToolChoiceEnum.none, True:
-                lark_grammar = grammar_factory.get_lark_for_json_schema(
-                    template=template, json_schema=json_schema
-                )
-            case _, _:
-                lark_grammar = grammar_factory.get_lark_from_jinja(
-                    template=template,
-                    mode=tool_choice,
-                    tools=mistral_tools,
-                    json_schema=json_schema,
-                    parallel_tool_calls=request.parallel_tool_calls,
-                    json_only=False,
-                )
+        try:
+            mistral_request = MistralChatCompletionRequest.from_openai(
+                **request.model_dump(
+                    include={"messages", "tools", "tool_choice", "parallel_tool_calls"},
+                    exclude_none=True,
+                ),
+                response_format=response_format,
+            )
+            lark_grammar = grammar_factory.get_grammar_for_chat_request(
+                request=mistral_request, template=template
+            )
+        except (ValidationError, InvalidRequestException) as e:
+            raise VLLMValidationError(str(e)) from e
 
-        request.structured_outputs = StructuredOutputsParams(grammar=lark_grammar)
-        request._grammar_from_parser = True
+        if lark_grammar is not None:
+            request.structured_outputs = StructuredOutputsParams(grammar=lark_grammar)
+            request._grammar_from_parser = True
         return request
 
     def _build_guided_schema_pre_v11(
diff --git a/vllm/parser/qwen3.py b/vllm/parser/qwen3.py
index 577e01edf6ed..f99d54c0c8d8 100644
--- a/vllm/parser/qwen3.py
+++ b/vllm/parser/qwen3.py
@@ -56,6 +56,7 @@
     re.DOTALL,
 )
 _PARTIAL_PARAM_RE = re.compile(r"<\s*parameter\s*=\s*([^>]+)>(.*)$", re.DOTALL)
+_PARAM_OPENER_RE = re.compile(r"<\s*parameter\s*=")
 
 
 def _trim_wrapping_newlines(value: str) -> str:
@@ -67,22 +68,60 @@ def _trim_wrapping_newlines(value: str) -> str:
     return value
 
 
-def _qwen3_arg_converter(raw_args: str, partial: bool) -> str:
-    params: dict[str, object] = {}
+class _ClosedParams:
+    """Closed parameters of a raw argument text, reused by longer texts.
 
-    for match in _PARAM_RE.finditer(raw_args):
-        name = match.group(1)
-        value = match.group(2)
-        params[name] = _trim_wrapping_newlines(value)
+    Raw arguments only grow while a call streams, so a parse of a text that
+    extends ``text`` resumes after its last closed parameter.
+    """
 
+    __slots__ = ("text", "end", "params", "outside")
+
+    def __init__(self) -> None:
+        self.text = ""  # raw text the closed parameters depend on
+        self.end = 0  # end of the last closed parameter in ``text``
+        self.params: dict[str, str] = {}
+        self.outside = ""  # text before ``end`` outside the closed parameters
+
+
+def _qwen3_parse_params(
+    raw_args: str, partial: bool, closed: _ClosedParams | None = None
+) -> tuple[dict[str, str], re.Match | None]:
+    """Return the closed parameters, trimmed, and the open trailing one."""
+    params: dict[str, str] = {}
+    outside: list[str] = []
+    pos = 0
+    if closed is not None and raw_args.startswith(closed.text):
+        params.update(closed.params)
+        outside.append(closed.outside)
+        pos = closed.end
+    start = pos
+
+    for match in _PARAM_RE.finditer(raw_args, pos):
+        outside.append(raw_args[pos : match.start()])
+        params[match.group(1)] = _trim_wrapping_newlines(match.group(2))
+        pos = match.end()
+
+    if closed is not None and pos > start:
+        # A value closed by the next opener holds only while that opener follows.
+        opener = _PARAM_OPENER_RE.match(raw_args, pos)
+        closed.text = raw_args[: opener.end() if opener else pos]
+        closed.end = pos
+        closed.params = dict(params)
+        closed.outside = "".join(outside)
+
+    open_param = None
     if partial:
-        remaining = _PARAM_RE.sub("", raw_args)
-        m = _PARTIAL_PARAM_RE.search(remaining)
-        if m:
-            name = m.group(1)
-            value = m.group(2)
-            if name:
-                params[name] = _trim_wrapping_newlines(value)
+        open_param = _PARTIAL_PARAM_RE.search("".join(outside) + raw_args[pos:])
+    return params, open_param
+
+
+def _qwen3_arg_converter(
+    raw_args: str, partial: bool, closed: _ClosedParams | None = None
+) -> str:
+    params, open_param = _qwen3_parse_params(raw_args, partial, closed)
+    if open_param is not None and open_param.group(1):
+        params[open_param.group(1)] = _trim_wrapping_newlines(open_param.group(2))
 
     return json.dumps(params, ensure_ascii=False)
 
@@ -101,7 +140,7 @@ def _param_value_end_window(text: str) -> str:
 
 
 def _param_value_extender(
-    raw_args: str, trim_newlines: bool
+    raw_args: str, trim_newlines: bool, closed: _ClosedParams | None = None
 ) -> ArgValueExtender | None:
     """Extend the open trailing ``<parameter=...>`` value of *raw_args*.
 
@@ -109,12 +148,12 @@ def _param_value_extender(
     With *trim_newlines* it follows ``_trim_wrapping_newlines``: it drops a
     leading newline of the value and holds back a trailing one.
     """
-    match = _PARTIAL_PARAM_RE.search(_PARAM_RE.sub("", raw_args))
+    closed_params, match = _qwen3_parse_params(raw_args, True, closed)
     if match is None:
         return None
     name, value = match.group(1), match.group(2)
     # A repeated key keeps its first position, so its value is not last.
-    if any(m.group(1) == name for m in _PARAM_RE.finditer(raw_args)):
+    if name in closed_params:
         return None
 
     window = _param_value_end_window(value)
@@ -141,8 +180,10 @@ def extend(delta: str) -> str | None:
     return extend
 
 
-def _qwen3_arg_value_extender(raw_args: str) -> ArgValueExtender | None:
-    return _param_value_extender(raw_args, trim_newlines=True)
+def _qwen3_arg_value_extender(
+    raw_args: str, closed: _ClosedParams | None = None
+) -> ArgValueExtender | None:
+    return _param_value_extender(raw_args, trim_newlines=True, closed=closed)
 
 
 @functools.cache
@@ -333,6 +374,13 @@ def __init__(
             tools,
             **kwargs,
         )
+        # Each streaming conversion resumes after the last closed parameter.
+        if self._arg_converter is _qwen3_arg_converter:
+            closed = _ClosedParams()
+            self._arg_converter = functools.partial(_qwen3_arg_converter, closed=closed)
+            self._arg_value_extender = functools.partial(
+                _qwen3_arg_value_extender, closed=closed
+            )
 
     def extract_reasoning(
         self,
diff --git a/vllm/platforms/cpu.py b/vllm/platforms/cpu.py
index 6b0d2958f85c..9d13e6eb8fb9 100644
--- a/vllm/platforms/cpu.py
+++ b/vllm/platforms/cpu.py
@@ -649,6 +649,13 @@ def opaque_attention_op(cls) -> bool:
     def support_hybrid_kv_cache(cls) -> bool:
         return True
 
+    @classmethod
+    def uses_standalone_cpu_model_runner(cls) -> bool:
+        # Only non-x86 CPUs have a standalone no-Triton MRV2 runner.
+        from vllm.triton_utils import HAS_TRITON
+
+        return not HAS_TRITON and cls.get_cpu_architecture() != CpuArchEnum.X86
+
     @classmethod
     def num_compute_units(cls, device_id: int = 0) -> int:
         return torch.get_num_threads()
diff --git a/vllm/platforms/cuda.py b/vllm/platforms/cuda.py
index c76587787132..87330e6d5b11 100644
--- a/vllm/platforms/cuda.py
+++ b/vllm/platforms/cuda.py
@@ -235,6 +235,10 @@ class CudaPlatformBase(Platform):
         "RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES",
     ]
 
+    @classmethod
+    def get_hisparse_backend_cls(cls) -> str | None:
+        return "vllm.v1.hisparse.backend.CUDAHiSparseBackend"
+
     @classmethod
     def import_kernels(cls) -> None:
         """Import CUDA kernel extensions (_C_stable_libtorch, optional _qutlass_C)."""
diff --git a/vllm/platforms/interface.py b/vllm/platforms/interface.py
index de899f7a1bf2..533f23d13eb0 100644
--- a/vllm/platforms/interface.py
+++ b/vllm/platforms/interface.py
@@ -196,6 +196,11 @@ def supported_dtypes(self) -> list[torch.dtype]:
     def is_cuda(self) -> bool:
         return self._enum == PlatformEnum.CUDA
 
+    @classmethod
+    def get_hisparse_backend_cls(cls) -> str | None:
+        """Return the device implementation for the common HiSparse runtime."""
+        return None
+
     def is_rocm(self) -> bool:
         return self._enum == PlatformEnum.ROCM
 
@@ -1355,6 +1360,12 @@ def support_hybrid_kv_cache(cls) -> bool:
         """Returns if the hybrid kv cache is supported by the current platform."""
         return False
 
+    @classmethod
+    def uses_standalone_cpu_model_runner(cls) -> bool:
+        """Returns if Model Runner V2 runs as the standalone (no-Triton) CPU
+        runner instead of the Triton-based one."""
+        return False
+
     @classmethod
     def support_static_graph_mode(cls) -> bool:
         """Returns if the graph mode is supported by the current platform."""
diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py
index 83ad6e3bcf19..000ac9cc18a2 100644
--- a/vllm/platforms/rocm.py
+++ b/vllm/platforms/rocm.py
@@ -1207,6 +1207,18 @@ def check_and_update_config(cls, vllm_config: "VllmConfig") -> None:
 
         _maybe_default_dsv4_mxfp4_indexer(vllm_config)
 
+        if parallel_config.all2all_backend == "mori":
+            from vllm.config.parallel import resolve_mori_all2all
+
+            resolve_mori_all2all(
+                parallel_config,
+                model_config,
+                vllm_config.kernel_config.moe_backend,
+                lsa_spans_hosts=on_gfx1250(),
+                sdma_picks_epv1=not on_gfx1250(),
+                validated_arch=on_gfx950() or on_gfx1250(),
+            )
+
     @classmethod
     def verify_model_arch(cls, model_arch: str) -> None:
         if model_arch in _ROCM_UNSUPPORTED_MODELS:
diff --git a/vllm/reasoning/minimax_m3_reasoning_parser.py b/vllm/reasoning/minimax_m3_reasoning_parser.py
index d49f9311f460..b6a8e48e49ef 100644
--- a/vllm/reasoning/minimax_m3_reasoning_parser.py
+++ b/vllm/reasoning/minimax_m3_reasoning_parser.py
@@ -39,6 +39,9 @@ def __init__(self, tokenizer, *args, **kwargs):
         super().__init__(tokenizer, *args, **kwargs)
         self._start_token_ids = self._encode_marker(self.start_token)
         self._end_token_ids = self._encode_marker(self.end_token)
+        # The chat template opens every message, including the generation
+        # prompt, with this token.
+        self._message_begin_token_id = self.vocab.get("]~b]")
         chat_kwargs = kwargs.get("chat_template_kwargs", {}) or {}
         self._initial_in_reasoning = chat_kwargs.get("thinking_mode") == "enabled"
         self._reasoning_ended_streaming = False
@@ -117,6 +120,16 @@ def _ends_with_token_sequence_prefix(
                 return True
         return False
 
+    def _last_message_ids(self, input_ids: Sequence[int]) -> Sequence[int]:
+        # The system message's thinking instructions and earlier turns contain
+        # reasoning markers; only the last message (the generation prompt)
+        # tells whether generation starts inside reasoning. Output-only ids
+        # contain no message-begin token and are returned whole.
+        for i in range(len(input_ids) - 1, -1, -1):
+            if input_ids[i] == self._message_begin_token_id:
+                return input_ids[i + 1 :]
+        return input_ids
+
     @staticmethod
     def _strip_partial_marker_suffix(text: str, marker: str) -> str:
         max_len = min(len(text), len(marker) - 1)
@@ -204,6 +217,7 @@ def is_reasoning_end_streaming(
             return False
 
         delta_ids = tuple(delta_ids)
+        input_ids = self._last_message_ids(input_ids)
         if self._contains_token_sequence(delta_ids, self._end_token_ids):
             return True
         if self._contains_token_sequence(input_ids, self._end_token_ids):
@@ -311,6 +325,7 @@ def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int:
         return count
 
     def is_reasoning_end(self, input_ids: Sequence[int]) -> bool:
+        input_ids = self._last_message_ids(input_ids)
         start_index = self._rfind_token_sequence(input_ids, self._start_token_ids)
         end_index = self._rfind_token_sequence(input_ids, self._end_token_ids)
         if end_index < 0:
diff --git a/vllm/renderers/hf.py b/vllm/renderers/hf.py
index fa8a1db71689..b1e096cd2230 100644
--- a/vllm/renderers/hf.py
+++ b/vllm/renderers/hf.py
@@ -1479,6 +1479,7 @@ def _apply_prompt_embeds_to_engine_input(
         for tensor, (start, length) in zip(
             prompt_embeds_tensors, positions, strict=True
         ):
+            token_ids[start : start + length] = [0] * length
             pe_kwargs_items.append(
                 MultiModalKwargsItem(
                     {
@@ -1500,6 +1501,7 @@ def _apply_prompt_embeds_to_engine_input(
                 PlaceholderRange(offset=start, length=length, is_embed=None)
             )
 
+        engine_input["prompt_token_ids"] = token_ids
         cast(
             MultiModalKwargsItems[MultiModalKwargsItem | None],
             engine_input["mm_kwargs"],
diff --git a/vllm/renderers/online_derenderer.py b/vllm/renderers/online_derenderer.py
index df0d21e42bc2..8ca82910e7e1 100644
--- a/vllm/renderers/online_derenderer.py
+++ b/vllm/renderers/online_derenderer.py
@@ -40,6 +40,7 @@
     GenerateLogProbs,
     GenerateTokensResponse,
     GenerateTokensStreamResponse,
+    ParseContext,
 )
 from vllm.entrypoints.serve.engine.protocol import UsageInfo
 from vllm.entrypoints.serve.utils.request_logger import RequestLogger
@@ -47,6 +48,7 @@
     resolve_finish_reason,
 )
 from vllm.logger import init_logger
+from vllm.logprobs import Logprob, PromptLogprobs
 from vllm.outputs import CompletionOutput
 from vllm.parser import Parser, ParserManager
 from vllm.renderers import BaseRenderer
@@ -149,7 +151,7 @@ async def derender_chat(
         generate_response: GenerateTokensResponse,
         chat_request: ChatCompletionRequest | None = None,
         prompt_token_ids: list[int] | None = None,
-    ) -> list[ChatCompletionResponseChoice]:
+    ) -> tuple[list[ChatCompletionResponseChoice], PromptLogprobs | None]:
         return await self._derender_chat_async(
             generate_response, chat_request, prompt_token_ids
         )
@@ -260,7 +262,7 @@ def _derender_chat(
         generate_response: GenerateTokensResponse,
         chat_request: ChatCompletionRequest | None = None,
         prompt_token_ids: list[int] | None = None,
-    ) -> list[ChatCompletionResponseChoice]:
+    ) -> tuple[list[ChatCompletionResponseChoice], PromptLogprobs | None]:
         tokenizer = self.renderer.get_tokenizer()
         choices: list[ChatCompletionResponseChoice] = []
 
@@ -355,7 +357,10 @@ def _derender_chat(
                 )
             )
 
-        return choices
+        prompt_logprobs = _resolve_prompt_logprobs(
+            generate_response.prompt_logprobs, tokenizer
+        )
+        return choices, prompt_logprobs
 
     async def derender_text_choices(
         self,
@@ -409,6 +414,48 @@ def _derender_text_choices(
             results.append((message, finish_reason))
         return results
 
+    def make_stream_processor(
+        self,
+        chat_request: ChatCompletionRequest | None,
+        history_tool_call_cnt: int,
+    ) -> ChatStreamProcessor:
+        """Build the request-scoped processor for one streamed choice.
+
+        Args:
+            chat_request: Request from `ParseContext.to_chat_request()`. Unused
+                when the server has no parser.
+            history_tool_call_cnt: Tool calls already in the conversation, for
+                deterministic tool call IDs.
+
+        Returns:
+            A processor holding its own parser, so call this once per choice.
+
+        Raises:
+            ValueError: If the parser rejects the request.
+            VLLMValidationError: If the parser rejects the request.
+
+        """
+        if self.parser is None or chat_request is None:
+            return ChatStreamProcessor(
+                ParseContext().to_chat_request(),
+                None,
+                enable_auto_tools=self.enable_auto_tools,
+                parser_cls=None,
+            )
+        parser = self.parser(
+            self.renderer.get_tokenizer(),
+            chat_request.tools,
+            chat_template_kwargs=chat_request.chat_template_kwargs or {},
+            model_config=self.model_config,
+        )
+        parser.set_history_tool_call_count(history_tool_call_cnt)
+        return ChatStreamProcessor(
+            chat_request,
+            parser,
+            enable_auto_tools=self.enable_auto_tools,
+            parser_cls=self.parser,
+        )
+
     def _resolve_chat_template_kwargs(
         self, chat_request: ChatCompletionRequest
     ) -> dict[str, Any]:
@@ -956,6 +1003,7 @@ def _derender_completion(
             seed_state = _seed_stream_state(
                 tokenizer, seed_ids, skip_special_tokens=skip_special
             )
+            prompt_logprobs = _resolve_prompt_logprobs(gen.prompt_logprobs, tokenizer)
             for choice in gen.choices:
                 if not choice.token_ids:
                     raise ValueError(
@@ -994,6 +1042,7 @@ def _derender_completion(
                             gen.prompt_token_ids if return_token_ids else None
                         ),
                         token_ids=choice.token_ids if return_token_ids else None,
+                        prompt_logprobs=prompt_logprobs,
                     )
                 )
                 total_completion_tokens += len(choice.token_ids)
@@ -1284,6 +1333,56 @@ def _completion_top_logprobs_limit(
     )
 
 
+def _resolve_prompt_logprobs(
+    prompt_logprobs: PromptLogprobs | None, tokenizer: TokenizerLike
+) -> PromptLogprobs | None:
+    """Fill missing ``decoded_token`` values in prompt logprobs.
+
+    A ``--tokens-only`` generate server has no tokenizer, so it returns
+    ``decoded_token=None``. Mirrors LogprobsProcessor._update_prompt_logprobs
+    in v1/engine/logprobs.py: decode each token id, then fix U+FFFD using the
+    preceding prompt tokens (the first entry of each earlier position).
+    """
+    if prompt_logprobs is None:
+        return None
+
+    missing = list(
+        dict.fromkeys(
+            tid
+            for entry in prompt_logprobs
+            if entry is not None
+            for tid, lp in entry.items()
+            if lp.decoded_token is None
+        )
+    )
+    if not missing:
+        return prompt_logprobs
+    # Ids with no vocab entry decode to "", as for output logprobs.
+    decoded = {
+        tid: text
+        for tid, (text, _) in zip(missing, decode_token_ids(missing, tokenizer))
+    }
+
+    resolved: PromptLogprobs = []
+    context_token_ids: list[int] = []
+    for entry in prompt_logprobs:
+        if not entry:
+            resolved.append(entry)
+            continue
+        new_entry: dict[int, Logprob] = {}
+        for tid, lp in entry.items():
+            if lp.decoded_token is None:
+                text = decoded[tid]
+                if text.endswith("\ufffd"):
+                    text = _correct_decoded_token(tid, context_token_ids, tokenizer)
+                lp = Logprob(lp.logprob, lp.rank, text)
+            new_entry[tid] = lp
+        resolved.append(new_entry)
+        # The prompt token is always inserted first by the engine.
+        context_token_ids.append(next(iter(entry)))
+    return resolved
+
+
 def _resolve_logprobs(
     logprobs: GenerateLogProbs,
     tokenizer: TokenizerLike,
diff --git a/vllm/renderers/registry.py b/vllm/renderers/registry.py
index cd687f650140..14429fb5100b 100644
--- a/vllm/renderers/registry.py
+++ b/vllm/renderers/registry.py
@@ -26,6 +26,7 @@
     "deepseek_v4": ("deepseek_v4", "DeepseekV4Renderer"),
     "deepseek_v41": ("deepseek_v4", "DeepseekV41Renderer"),
     "hf": ("hf", "HfRenderer"),
+    "glm5_next": ("hf", "HfRenderer"),
     "kimi_audio": ("hf", "HfRenderer"),
     "kimi_k3": ("kimi_k3", "KimiK3Renderer"),
     "mistral": ("mistral", "MistralRenderer"),
diff --git a/vllm/sampling_params.py b/vllm/sampling_params.py
index 823840312864..73dd2d06d4b9 100644
--- a/vllm/sampling_params.py
+++ b/vllm/sampling_params.py
@@ -1125,10 +1125,6 @@ def _validate_trace_replay(
             raise ValueError(
                 "trace_decode_token_ids is not supported with prompt_logprob_token_ids."
             )
-        if speculative_config is not None:
-            raise ValueError(
-                "trace_decode_token_ids is not supported with speculative decoding."
-            )
         if self.structured_outputs is not None:
             raise ValueError(
                 "trace_decode_token_ids is not supported with structured outputs."
diff --git a/vllm/third_party/flash_linear_attention/ops/chunk_delta_h.py b/vllm/third_party/flash_linear_attention/ops/chunk_delta_h.py
index cc27437fac58..cde69bef7403 100644
--- a/vllm/third_party/flash_linear_attention/ops/chunk_delta_h.py
+++ b/vllm/third_party/flash_linear_attention/ops/chunk_delta_h.py
@@ -94,20 +94,21 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64(
         b_h4 = tl.zeros([BV, 64], dtype=tl.float32)
 
     # calculate offset
-    h += ((boh * H + i_h) * V * K).to(tl.int64)
-    v += ((bos * H + i_h) * V).to(tl.int64)
-    k += ((bos * Hg + i_h // (H // Hg)) * K).to(tl.int64)
-    w += ((bos * H + i_h) * K).to(tl.int64)
+    # Cast before multiplying: the products overflow int32 for long batches.
+    h += (boh * H + i_h).to(tl.int64) * V * K
+    v += (bos * H + i_h).to(tl.int64) * V
+    k += (bos * Hg + i_h // (H // Hg)).to(tl.int64) * K
+    w += (bos * H + i_h).to(tl.int64) * K
     if SAVE_NEW_VALUE:
-        v_new += ((bos * H + i_h) * V).to(tl.int64)
+        v_new += (bos * H + i_h).to(tl.int64) * V
     stride_v = H * V
     stride_h = H * V * K
     stride_k = Hg * K
     stride_w = H * K
     if USE_INITIAL_STATE:
-        h0 = h0 + i_nh * V * K
+        h0 = h0 + i_nh.to(tl.int64) * V * K
     if STORE_FINAL_STATE:
-        ht = ht + i_nh * V * K
+        ht = ht + i_nh.to(tl.int64) * V * K
 
     # load initial state
     o_v = i_v * BV + tl.arange(0, BV)
@@ -291,7 +292,7 @@ def chunk_gated_delta_rule_fwd_h(
     chunk_indices: torch.Tensor | None = None,
     chunk_offsets: torch.Tensor | None = None,
     use_exp2: bool = False,
-) -> tuple[torch.Tensor, torch.Tensor]:
+) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]:
     # This kernel is slightly different from fla to support Q/K with different head numbers.
     # In fla, Q/K always have the same head number, so Hg is always equal to H.
     B, T, Hg, K, V = *k.shape, u.shape[-1]
diff --git a/vllm/third_party/flash_linear_attention/ops/fused_gdn_prefill_post_conv.py b/vllm/third_party/flash_linear_attention/ops/fused_gdn_prefill_post_conv.py
index 4807c78e7b1a..2975ea15722f 100644
--- a/vllm/third_party/flash_linear_attention/ops/fused_gdn_prefill_post_conv.py
+++ b/vllm/third_party/flash_linear_attention/ops/fused_gdn_prefill_post_conv.py
@@ -63,8 +63,20 @@ def _fused_post_conv_kernel(
 
     HK: tl.constexpr = H * K
 
-    offs_t = i_tb * BLOCK_T + tl.arange(0, BLOCK_T)  # [BLOCK_T]
-    mask_t = offs_t < L
+    # Move every pointer to the block's first token in int64, so that
+    # token * stride does not overflow int32 for large batches.
+    t_start = i_tb.to(tl.int64) * BLOCK_T
+    mixed_qkv_ptr += t_start * stride_x_tok
+    a_ptr += t_start * stride_a_tok
+    b_ptr += t_start * stride_b_tok
+    q_ptr += t_start * stride_q_tok
+    k_ptr += t_start * stride_k_tok
+    v_ptr += t_start * stride_v_tok
+    g_ptr += t_start * HV
+    beta_ptr += t_start * HV
+
+    offs_t = tl.arange(0, BLOCK_T)  # [BLOCK_T], relative to t_start
+    mask_t = t_start + offs_t < L
 
     if i_head < H:
         # ============ Q/K head processing ============
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 997df5a173b1..b51bcdabcb98 100644
--- a/vllm/third_party/flash_linear_attention/ops/layernorm_guard.py
+++ b/vllm/third_party/flash_linear_attention/ops/layernorm_guard.py
@@ -95,7 +95,7 @@ def layer_norm_fwd_kernel(
     ACTIVATION: tl.constexpr,
 ):
     # Map the program id to the starting row of X and Y it should compute.
-    row_start = tl.program_id(0) * ROWS_PER_BLOCK
+    row_start = tl.program_id(0).to(tl.int64) * ROWS_PER_BLOCK
     group = tl.program_id(1)
 
     # Create 2D tile: [ROWS_PER_BLOCK, BLOCK_N]
diff --git a/vllm/third_party/flash_linear_attention/ops/solve_tril.py b/vllm/third_party/flash_linear_attention/ops/solve_tril.py
index 93ef770be627..b0dbaf775977 100644
--- a/vllm/third_party/flash_linear_attention/ops/solve_tril.py
+++ b/vllm/third_party/flash_linear_attention/ops/solve_tril.py
@@ -65,8 +65,8 @@ def solve_tril_16x16_kernel(
     m_A = o_i[:, None] > o_i[None, :]
     m_I = o_i[:, None] == o_i[None, :]
 
-    A = A + (bos * H + i_h) * BT
-    Ai = Ai + (bos * H + i_h) * 16
+    A = A + (bos * H + i_h).to(tl.int64) * BT
+    Ai = Ai + (bos * H + i_h).to(tl.int64) * 16
 
     offset = (i_t * 16) % BT
     o_t = (i_t * 16 + o_i).to(tl.int64)
@@ -83,7 +83,7 @@ def solve_tril_16x16_kernel(
 
     for i in range(2, min(16, T - i_t * 16)):
         # [16]
-        b_a = -tl.load(A + (i_t * 16 + i) * H * BT + o_i + offset)
+        b_a = -tl.load(A + (i_t * 16 + i).to(tl.int64) * H * BT + o_i + offset)
         b_a = b_a + tl.sum(b_a[:, None] * b_A, 0)
         b_A = tl.where((o_i == i)[:, None], b_a, b_A)
     b_A += m_I
@@ -138,8 +138,8 @@ def merge_16x16_to_32x32_inverse_kernel(
     o_i = tl.arange(0, 16)
     m_A = o_i[:, None] > o_i[None, :]
     m_I = o_i[:, None] == o_i[None, :]
-    A += (bos * H + i_h) * BT
-    Ai += (bos * H + i_h) * BT
+    A += (bos * H + i_h).to(tl.int64) * BT
+    Ai += (bos * H + i_h).to(tl.int64) * BT
 
     o_t = (i_t * BT + o_i).to(tl.int64)
     if not USE_TMA:
@@ -160,11 +160,11 @@ def merge_16x16_to_32x32_inverse_kernel(
     b_Ai_22 = -tl.where(m_A, b_Ai_22, 0)
 
     for i in range(2, min(16, T - i_t * BT)):
-        b_a_11 = -tl.load(A + (i_t * BT + i) * H * BT + o_i)
+        b_a_11 = -tl.load(A + (i_t * BT + i).to(tl.int64) * H * BT + o_i)
         b_a_11 += tl.sum(b_a_11[:, None] * b_Ai_11, 0)
         b_Ai_11 = tl.where((o_i == i)[:, None], b_a_11, b_Ai_11)
     for i in range(16 + 2, min(32, T - i_t * BT)):
-        b_a_22 = -tl.load(A + (i_t * BT + i) * H * BT + o_i + 16)
+        b_a_22 = -tl.load(A + (i_t * BT + i).to(tl.int64) * H * BT + o_i + 16)
         b_a_22 += tl.sum(b_a_22[:, None] * b_Ai_22, 0)
         b_Ai_22 = tl.where((o_i == i - 16)[:, None], b_a_22, b_Ai_22)
 
@@ -254,8 +254,8 @@ def merge_16x16_to_64x64_inverse_kernel(
     o_i = tl.arange(0, 16)
     m_A = o_i[:, None] > o_i[None, :]
     m_I = o_i[:, None] == o_i[None, :]
-    A += (bos * H + i_h) * BT
-    Ai += (bos * H + i_h) * BT
+    A += (bos * H + i_h).to(tl.int64) * BT
+    Ai += (bos * H + i_h).to(tl.int64) * BT
 
     o_t = (i_t * BT + o_i).to(tl.int64)
     if not USE_TMA:
@@ -288,19 +288,19 @@ def merge_16x16_to_64x64_inverse_kernel(
     b_Ai_44 = -tl.where(m_A, b_Ai_44, 0)
 
     for i in range(2, min(16, T - i_t * BT)):
-        b_a_11 = -tl.load(A + (i_t * BT + i) * H * BT + o_i)
+        b_a_11 = -tl.load(A + (i_t * BT + i).to(tl.int64) * H * BT + o_i)
         b_a_11 += tl.sum(b_a_11[:, None] * b_Ai_11, 0)
         b_Ai_11 = tl.where((o_i == i)[:, None], b_a_11, b_Ai_11)
     for i in range(16 + 2, min(32, T - i_t * BT)):
-        b_a_22 = -tl.load(A + (i_t * BT + i) * H * BT + o_i + 16)
+        b_a_22 = -tl.load(A + (i_t * BT + i).to(tl.int64) * H * BT + o_i + 16)
         b_a_22 += tl.sum(b_a_22[:, None] * b_Ai_22, 0)
         b_Ai_22 = tl.where((o_i == i - 16)[:, None], b_a_22, b_Ai_22)
     for i in range(32 + 2, min(48, T - i_t * BT)):
-        b_a_33 = -tl.load(A + (i_t * BT + i) * H * BT + o_i + 32)
+        b_a_33 = -tl.load(A + (i_t * BT + i).to(tl.int64) * H * BT + o_i + 32)
         b_a_33 += tl.sum(b_a_33[:, None] * b_Ai_33, 0)
         b_Ai_33 = tl.where((o_i == i - 32)[:, None], b_a_33, b_Ai_33)
     for i in range(48 + 2, min(64, T - i_t * BT)):
-        b_a_44 = -tl.load(A + (i_t * BT + i) * H * BT + o_i + 48)
+        b_a_44 = -tl.load(A + (i_t * BT + i).to(tl.int64) * H * BT + o_i + 48)
         b_a_44 += tl.sum(b_a_44[:, None] * b_Ai_44, 0)
         b_Ai_44 = tl.where((o_i == i - 48)[:, None], b_a_44, b_Ai_44)
     b_Ai_11 += m_I
diff --git a/vllm/tokenizers/deepseek_v41.py b/vllm/tokenizers/deepseek_v41.py
index 966dc7f9b9cf..c0efec6dc07c 100644
--- a/vllm/tokenizers/deepseek_v41.py
+++ b/vllm/tokenizers/deepseek_v41.py
@@ -29,7 +29,7 @@ def _normalize_messages(
     messages: list[ChatCompletionMessageParam],
 ) -> list[dict[str, Any]]:
     # vLLM additions to the reference: Responses API part types, image_pil,
-    # and ValueError (400) instead of the reference's asserts.
+    # tool_reference, and ValueError (400) instead of the reference's asserts.
     result = [dict(message) for message in copy.deepcopy(messages)]
     image_idx = 0
     for message in result:
@@ -54,12 +54,18 @@ def _normalize_messages(
                     parts.append(IMAGE_PLACEHOLDER)
                     image_indices.append(image_idx)
                     image_idx += 1
+                elif part_type == "tool_reference":
+                    # The system prompt already lists every tool, deferred
+                    # ones included, so the name is all the model needs.
+                    _reject_image_placeholder(block.get("name"))
+                    parts.append(block.get("name", ""))
                 else:
                     raise ValueError(
                         "DeepSeek V4.1 supports text and image content only; "
                         f"got {part_type!r}"
                     )
             message["content"] = "\n\n".join(parts)
+
     return result
 
 
diff --git a/vllm/tokenizers/registry.py b/vllm/tokenizers/registry.py
index b4351a03be0a..6f9d63310cce 100644
--- a/vllm/tokenizers/registry.py
+++ b/vllm/tokenizers/registry.py
@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import contextlib
+from copy import copy
 from dataclasses import dataclass, field
 from functools import lru_cache
 from pathlib import Path
@@ -52,6 +53,7 @@
     "deepseek_v4": ("deepseek_v4", "DeepseekV4Tokenizer"),
     "deepseek_v41": ("deepseek_v41", "DeepseekV41Tokenizer"),
     "hf": ("hf", "CachedHfTokenizer"),
+    "glm5_next": ("hf", "CachedHfTokenizer"),
     "kimi_audio": ("kimi_audio", "KimiAudioTokenizer"),
     "kimi_k3": ("hf", "CachedHfTokenizer"),
     "mistral": ("mistral", "MistralTokenizer"),
@@ -270,6 +272,10 @@ def get_tokenizer(
         tokenizer_cls_ = TokenizersBackend
 
     if config is not None and tokenizer_cls_ is CachedHfTokenizer:
+        if not envs.VLLM_USE_MODELSCOPE and str(tokenizer_name) != requested_name:
+            # AutoTokenizer selects some tokenizer classes by Hub ID
+            config = copy(config)
+            config._name_or_path = requested_name
         # AutoTokenizer otherwise reloads config.json internally. Reuse the
         # config that get_config just loaded successfully so a concurrent Hub
         # cache refresh cannot invalidate the file between the two reads.
diff --git a/vllm/tool_parsers/functiongemma_tool_parser.py b/vllm/tool_parsers/functiongemma_tool_parser.py
index 6fcae6e55220..4d809bb42029 100644
--- a/vllm/tool_parsers/functiongemma_tool_parser.py
+++ b/vllm/tool_parsers/functiongemma_tool_parser.py
@@ -25,6 +25,109 @@
 
 logger = init_logger(__name__)
 
+_ESCAPE = "<escape>"
+_BARE_LITERALS = {"true": True, "false": False, "null": None, "None": None}
+
+
+class _IncompleteValue(Exception):
+    """The argument text ends before the value does."""
+
+
+class _MalformedValue(Exception):
+    """A delimiter appears where a value or key should start."""
+
+
+def _skip_whitespace(text: str, pos: int) -> int:
+    while pos < len(text) and text[pos] in " \n":
+        pos += 1
+    return pos
+
+
+def _decode_escaped(raw: str):
+    try:
+        return json.loads(raw)
+    except json.JSONDecodeError:
+        return raw
+
+
+def _decode_bare(raw: str):
+    raw = raw.strip()
+    if raw in _BARE_LITERALS:
+        return _BARE_LITERALS[raw]
+    try:
+        return json.loads(raw)
+    except json.JSONDecodeError:
+        return raw
+
+
+def _parse_value(text: str, pos: int, partial: bool) -> tuple[object, int]:
+    """Parse one value at ``pos`` and return it with the position after it.
+
+    Raises:
+        _IncompleteValue: ``text`` ends inside the value. A bare scalar at the
+            end of a ``partial`` string counts as incomplete, since more digits
+            or letters may still follow.
+        _MalformedValue: a delimiter stands where the value should start.
+
+    """
+    pos = _skip_whitespace(text, pos)
+    if text.startswith(_ESCAPE, pos):
+        start = pos + len(_ESCAPE)
+        end = text.find(_ESCAPE, start)
+        if end < 0:
+            raise _IncompleteValue
+        return _decode_escaped(text[start:end]), end + len(_ESCAPE)
+    if text.startswith("{", pos):
+        return _parse_object(text, pos + 1, partial)
+    if text.startswith("[", pos):
+        return _parse_array(text, pos + 1, partial)
+    end = pos
+    while end < len(text) and text[end] not in ",}]":
+        end += 1
+    if end == pos:
+        raise _MalformedValue
+    if end == len(text) and partial:
+        raise _IncompleteValue
+    return _decode_bare(text[pos:end]), end
+
+
+def _parse_object(text: str, pos: int, partial: bool) -> tuple[dict, int]:
+    obj: dict = {}
+    while True:
+        pos = _skip_whitespace(text, pos)
+        if pos >= len(text):
+            raise _IncompleteValue
+        if text[pos] == "}":
+            return obj, pos + 1
+        if text[pos] == "]":
+            raise _MalformedValue
+        if text[pos] == ",":
+            pos += 1
+            continue
+        colon = text.find(":", pos)
+        if colon < 0:
+            raise _IncompleteValue
+        key = text[pos:colon].strip()
+        value, pos = _parse_value(text, colon + 1, partial)
+        obj[key] = value
+
+
+def _parse_array(text: str, pos: int, partial: bool) -> tuple[list, int]:
+    items: list = []
+    while True:
+        pos = _skip_whitespace(text, pos)
+        if pos >= len(text):
+            raise _IncompleteValue
+        if text[pos] == "]":
+            return items, pos + 1
+        if text[pos] == "}":
+            raise _MalformedValue
+        if text[pos] == ",":
+            pos += 1
+            continue
+        value, pos = _parse_value(text, pos, partial)
+        items.append(value)
+
 
 class FunctionGemmaToolParser(ToolParser):
     """Tool parser for Google's FunctionGemma model (google/functiongemma-270m-it).
@@ -43,10 +146,6 @@ class FunctionGemmaToolParser(ToolParser):
         r"|<start_function_call>call:(\w+)\{(.*)",
         re.DOTALL,
     )
-    arg_regex: re.Pattern = re.compile(
-        r"(\w+):<escape>(.*?)<escape>",
-        re.DOTALL,
-    )
 
     def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None):
         super().__init__(tokenizer, tools)
@@ -58,20 +157,32 @@ def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None):
         self.streamed_args_for_tool: list[str] = []
         self.buffered_delta_text = ""
 
-    def _parse_arguments(self, args_str: str) -> dict:
-        """Parse FunctionGemma argument string into a dictionary."""
-        arguments: dict = {}
-        if not args_str:
-            return arguments
+    def _parse_arguments(self, args_str: str, partial: bool = False) -> dict:
+        """Parse a FunctionGemma argument string into a dictionary.
 
-        matches = self.arg_regex.findall(args_str)
-        for key, value in matches:
+        The chat template writes strings as ``<escape>text<escape>`` and
+        everything else bare: ``true``/``false``, numbers, ``[item,...]`` and
+        ``{key:value,...}`` with bare keys. With ``partial`` the string is a
+        streaming prefix: parsing stops at the first incomplete value and the
+        complete arguments before it are returned.
+        """
+        arguments: dict = {}
+        pos = 0
+        while pos < len(args_str):
+            if args_str[pos] in ", \n":
+                pos += 1
+                continue
+            colon = args_str.find(":", pos)
+            if colon < 0:
+                break
+            key = args_str[pos:colon].strip()
+            if not key:
+                break
             try:
-                parsed_value = json.loads(value)
-                arguments[key] = parsed_value
-            except json.JSONDecodeError:
-                arguments[key] = value
-
+                value, pos = _parse_value(args_str, colon + 1, partial)
+            except (_IncompleteValue, _MalformedValue):
+                break
+            arguments[key] = value
         return arguments
 
     def adjust_request(
@@ -143,6 +254,46 @@ def extract_tool_calls(
                 tools_called=False, tool_calls=[], content=model_output
             )
 
+    def _stream_arguments(self, args: dict, final: bool) -> DeltaMessage | None:
+        """Stream the part of the arguments JSON not sent yet.
+
+        Arguments only grow by appended keys, so their JSON is prefix-stable
+        except for the closing brace, which is held back until ``final``.
+        """
+        self.prev_tool_call_arr[self.current_tool_id]["arguments"] = args
+        args_json = json.dumps(args, ensure_ascii=False)
+        if not final:
+            args_json = args_json[:-1]
+        prev_streamed = self.streamed_args_for_tool[self.current_tool_id]
+        if len(args_json) <= len(prev_streamed):
+            return None
+        diff = args_json[len(prev_streamed) :]
+        self.streamed_args_for_tool[self.current_tool_id] = args_json
+        return DeltaMessage(
+            tool_calls=[
+                DeltaToolCall(
+                    index=self.current_tool_id,
+                    function=DeltaFunctionCall(arguments=diff).model_dump(
+                        exclude_none=True
+                    ),
+                )
+            ]
+        )
+
+    def _close_tool_call(self, current_text: str) -> DeltaMessage | None:
+        """Stream what is left of the current call, including its closing brace."""
+        if not 0 <= self.current_tool_id < len(self.prev_tool_call_arr):
+            return None
+        all_calls = self.tool_call_regex.findall(current_text)
+        args = {}
+        if self.current_tool_id < len(all_calls):
+            match = all_calls[self.current_tool_id]
+            if match[0]:
+                args = self._parse_arguments(match[1])
+        if args or self.current_tool_name_sent:
+            return self._stream_arguments(args, final=True)
+        return None
+
     def _buffer_delta_text(self, delta_text: str) -> str:
         """Buffer incoming delta text to handle multi-token special sequences."""
         potential_start = "<start_function_call>"
@@ -192,12 +343,17 @@ def extract_tool_calls_streaming(
 
             # Starting a new function call
             if start_count > prev_start_count and start_count > end_count:
+                closing = (
+                    self._close_tool_call(current_text)
+                    if end_count > prev_end_count
+                    else None
+                )
                 self.current_tool_id += 1
                 self.current_tool_name_sent = False
                 self.streamed_args_for_tool.append("")
                 self.prev_tool_call_arr.append({})
                 logger.debug("Starting new tool call %d", self.current_tool_id)
-                return None
+                return closing
 
             # In the middle of a function call
             if start_count > end_count:
@@ -235,74 +391,17 @@ def extract_tool_calls_streaming(
                             )
 
                         if self.current_tool_name_sent and args_part:
-                            current_args = self._parse_arguments(args_part)
+                            current_args = self._parse_arguments(
+                                args_part, partial=True
+                            )
                             if current_args:
-                                current_args_json = json.dumps(
-                                    current_args, ensure_ascii=False
-                                )
-                                prev_streamed = self.streamed_args_for_tool[
-                                    self.current_tool_id
-                                ]
-
-                                if len(current_args_json) > len(prev_streamed):
-                                    diff = current_args_json[len(prev_streamed) :]
-                                    self.streamed_args_for_tool[
-                                        self.current_tool_id
-                                    ] = current_args_json
-                                    self.prev_tool_call_arr[self.current_tool_id][
-                                        "arguments"
-                                    ] = current_args
-
-                                    return DeltaMessage(
-                                        tool_calls=[
-                                            DeltaToolCall(
-                                                index=self.current_tool_id,
-                                                function=DeltaFunctionCall(
-                                                    arguments=diff
-                                                ).model_dump(exclude_none=True),
-                                            )
-                                        ]
-                                    )
+                                return self._stream_arguments(current_args, final=False)
 
                 return None
 
             # Function call just ended
             if end_count > prev_end_count:
-                if self.current_tool_id >= 0 and self.current_tool_id < len(
-                    self.prev_tool_call_arr
-                ):
-                    all_calls = self.tool_call_regex.findall(current_text)
-                    args = {}
-                    if self.current_tool_id < len(all_calls):
-                        match = all_calls[self.current_tool_id]
-                        if match[0]:
-                            args_str = match[1]
-                            args = self._parse_arguments(args_str)
-                            self.prev_tool_call_arr[self.current_tool_id][
-                                "arguments"
-                            ] = args
-
-                    if args:
-                        args_json = json.dumps(args, ensure_ascii=False)
-                        prev_streamed = self.streamed_args_for_tool[
-                            self.current_tool_id
-                        ]
-                        if len(args_json) > len(prev_streamed):
-                            diff = args_json[len(prev_streamed) :]
-                            self.streamed_args_for_tool[self.current_tool_id] = (
-                                args_json
-                            )
-                            return DeltaMessage(
-                                tool_calls=[
-                                    DeltaToolCall(
-                                        index=self.current_tool_id,
-                                        function=DeltaFunctionCall(
-                                            arguments=diff
-                                        ).model_dump(exclude_none=True),
-                                    )
-                                ]
-                            )
-                return None
+                return self._close_tool_call(current_text)
 
             if delta_text:
                 return DeltaMessage(content=delta_text)
diff --git a/vllm/tool_parsers/hy_v4_tool_parser.py b/vllm/tool_parsers/hy_v4_tool_parser.py
index de707c00228e..fc8b8f0ce14f 100644
--- a/vllm/tool_parsers/hy_v4_tool_parser.py
+++ b/vllm/tool_parsers/hy_v4_tool_parser.py
@@ -805,7 +805,7 @@ def _extract_streaming_incremental(
         # --- scan all fully closed kv pairs ---
         arg_pairs = self.func_args_regex.findall(args_text)
         for key, value in arg_pairs:
-            key = key.strip()
+            # Property names are exact schema keys, including whitespace.
             if key not in self._completed_args:
                 self._completed_args[key] = _parse_value(
                     value, self._streaming_tool_name or "", key, tools
@@ -827,9 +827,7 @@ def _extract_streaming_incremental(
                 ak_start + len(self.arg_key_start_token),
             )
             if ak_end != -1:
-                partial_key = tail[
-                    ak_start + len(self.arg_key_start_token) : ak_end
-                ].strip()
+                partial_key = tail[ak_start + len(self.arg_key_start_token) : ak_end]
                 self._current_arg_key = partial_key
                 self._current_arg_is_string = _is_only_string_type(
                     self._streaming_tool_name or "",
diff --git a/vllm/tool_parsers/structural_tag_registry.py b/vllm/tool_parsers/structural_tag_registry.py
index 7e5902b8c2f6..a9c8ee38119a 100644
--- a/vllm/tool_parsers/structural_tag_registry.py
+++ b/vllm/tool_parsers/structural_tag_registry.py
@@ -1008,6 +1008,14 @@ def single_tool_call(tool: FunctionToolParam) -> TagFormat:
 ]
 
 
+# Only scalar values are pinned as bare JSON. Array and object values stay raw
+# text: a JSON grammar there admits only JSON escapes after a backslash, so text
+# the model writes with bare backslashes (a Windows path) is forced into escapes
+# the model gives near-zero probability. The parser keeps the raw text when such
+# a value is not valid JSON.
+_GLM_4_7_JSON_SCALAR_TYPES = frozenset({"integer", "number", "boolean", "null"})
+
+
 def _glm_4_7_value_format(prop: dict[str, Any]) -> Any:
     if "enum" in prop:
         values = [
@@ -1016,16 +1024,11 @@ def _glm_4_7_value_format(prop: dict[str, Any]) -> Any:
         ]
         return values[0] if len(values) == 1 else OrFormat(elements=values)
     prop_type = prop.get("type")
-    if isinstance(prop_type, list):
-        formats = [
-            _glm_4_7_value_format({"type": t}) for t in prop_type if t != "string"
-        ]
-        if "string" in prop_type:
-            formats.append(AnyTextFormat(excludes=_GLM_4_7_VALUE_EXCLUDES))
-        return formats[0] if len(formats) == 1 else OrFormat(elements=formats)
-    if prop_type in (None, "string"):
+    types = prop_type if isinstance(prop_type, list) else [prop_type]
+    if not types or not all(t in _GLM_4_7_JSON_SCALAR_TYPES for t in types):
         return AnyTextFormat(excludes=_GLM_4_7_VALUE_EXCLUDES)
-    return JSONSchemaFormat(json_schema={"type": prop_type})
+    formats = [JSONSchemaFormat(json_schema={"type": t}) for t in types]
+    return formats[0] if len(formats) == 1 else OrFormat(elements=formats)
 
 
 def _glm_4_7_plain_properties(params: dict[str, Any]) -> dict[str, Any] | None:
@@ -1045,7 +1048,7 @@ def _glm_4_7_plain_properties(params: dict[str, Any]) -> dict[str, Any] | None:
 
 
 def _glm_4_7_tool_tag(tool: FunctionToolParam) -> TagFormat:
-    """Shallow tag for a non-strict tool: keys and basic value types pinned."""
+    """Shallow tag for a non-strict tool: keys and scalar value types pinned."""
     properties = _glm_4_7_plain_properties(tool.function.parameters or {})
     if properties is None:
         return TagFormat(
diff --git a/vllm/transformers_utils/chat_templates/registry.py b/vllm/transformers_utils/chat_templates/registry.py
index 67ea08efb2c0..ec0ff63953d4 100644
--- a/vllm/transformers_utils/chat_templates/registry.py
+++ b/vllm/transformers_utils/chat_templates/registry.py
@@ -35,6 +35,7 @@ def _get_minicpmv_chat_template_fallback(tokenizer_name_or_path: str) -> Path |
     "minicpmv": _get_minicpmv_chat_template_fallback,
     "minicpmv4_6": _get_minicpmv_chat_template_fallback,
     "minicpmv4_7": _get_minicpmv_chat_template_fallback,
+    "openvla": CHAT_TEMPLATES_DIR / "template_openvla.jinja",
     "paligemma": CHAT_TEMPLATES_DIR / "template_basic.jinja",
     "siglip": CHAT_TEMPLATES_DIR / "template_basic.jinja",
     "siglip2": CHAT_TEMPLATES_DIR / "template_basic.jinja",
diff --git a/vllm/transformers_utils/chat_templates/template_openvla.jinja b/vllm/transformers_utils/chat_templates/template_openvla.jinja
new file mode 100644
index 000000000000..3eeaa08b33ae
--- /dev/null
+++ b/vllm/transformers_utils/chat_templates/template_openvla.jinja
@@ -0,0 +1,19 @@
+{# Mirrors OpenVLA's training prompt, "In: What action should the robot take
+   to {instruction}?\nOut: " with the instruction lowercased. predict_action()
+   appends token 29871 (a trailing space) if missing. The multimodal processor 
+   inserts image tokens after BOS. #}
+{%- set ns = namespace(text='') -%}
+{%- for message in messages if message['role'] == 'user' -%}
+    {%- if message['content'] is string -%}
+        {%- set ns.text = message['content'] -%}
+    {%- else -%}
+        {%- set ns.text = message['content'] | selectattr('type', 'equalto', 'text') | map(attribute='text') | join('') -%}
+    {%- endif -%}
+{%- endfor -%}
+{%- set text = ns.text | trim -%}
+{{- bos_token -}}
+{%- if text.startswith('In:') -%}
+    {{- text ~ ' ' -}}
+{%- else -%}
+    {{- 'In: What action should the robot take to ' ~ (text.rstrip('.?!') | lower) ~ '?\nOut: ' -}}
+{%- endif -%}
diff --git a/vllm/transformers_utils/config.py b/vllm/transformers_utils/config.py
index e18fc2753cc0..dad35bca3253 100644
--- a/vllm/transformers_utils/config.py
+++ b/vllm/transformers_utils/config.py
@@ -154,8 +154,10 @@ def __getitem__(self, key):
     midashenglm="MiDashengLMConfig",
     minimax_m3_vl="MiniMaxM3Config",
     minimax_m3_mtp="MiniMaxM3MTPConfig",
+    ming_v3_asr="MingV3ASRConfig",
     moondream3="Moondream3Config",
     moss_transcribe_diarize="MossTranscribeDiarizeConfig",
+    nemo_speechlm="Nemotron3_5TranscribeConfig",
     eagle="EAGLEConfig",
     speculators="SpeculatorsConfig",
     openvla="OpenVLAConfig",
@@ -842,6 +844,9 @@ def maybe_override_with_speculators(
 
     # Speculators format detected - process overrides
     from vllm.transformers_utils.configs.speculators.base import SpeculatorsConfig
+    from vllm.v1.worker.gpu.spec_decode.speculator_config import (
+        get_custom_speculator_config_class,
+    )
 
     speculative_config = SpeculatorsConfig.extract_vllm_speculative_config(
         config_dict=config_dict
@@ -852,8 +857,17 @@ def maybe_override_with_speculators(
     if isinstance(vllm_speculative_config, dict):
         speculative_config.update(vllm_speculative_config)
 
-    # Lock fields dictated by the speculators format
-    speculative_config["method"] = speculators_method
+    # Lock fields dictated by the speculators format, except for a custom
+    # method based on the checkpoint's.
+    custom_config_cls = get_custom_speculator_config_class(speculative_config["method"])
+    if custom_config_cls is None:
+        speculative_config["method"] = speculators_method
+    elif custom_config_cls.base_method != speculators_method:
+        raise ValueError(
+            f"Speculative method '{speculative_config['method']}' is based on "
+            f"'{custom_config_cls.base_method}', but the speculators "
+            f"checkpoint {model} is for '{speculators_method}'."
+        )
     speculative_config["model"] = model
 
     # Override model and tokenizer with the verifier model from config
diff --git a/vllm/transformers_utils/configs/__init__.py b/vllm/transformers_utils/configs/__init__.py
index 3042e1f5aa10..0e4dcef8426e 100644
--- a/vllm/transformers_utils/configs/__init__.py
+++ b/vllm/transformers_utils/configs/__init__.py
@@ -41,6 +41,7 @@
     # `FalconConfig` class from the official HuggingFace transformers library.
     "RWConfig": "vllm.transformers_utils.configs.falcon",
     "MiDashengLMConfig": "vllm.transformers_utils.configs.midashenglm",
+    "MingV3ASRConfig": "vllm.transformers_utils.configs.ming_v3_asr",
     "MiniMaxM3Config": "vllm.transformers_utils.configs.minimax_m3",
     "MiniMaxM3MTPConfig": "vllm.transformers_utils.configs.minimax_m3",
     "MiniMaxM3TextConfig": "vllm.transformers_utils.configs.minimax_m3",
@@ -50,6 +51,9 @@
     "MossTranscribeDiarizeConfig": (
         "vllm.transformers_utils.configs.moss_transcribe_diarize"
     ),
+    "Nemotron3_5TranscribeConfig": (
+        "vllm.transformers_utils.configs.nemotron3_5_transcribe"
+    ),
     "MoonViTConfig": "vllm.transformers_utils.configs.moonvit",
     "KimiLinearConfig": "vllm.transformers_utils.configs.kimi_linear",
     "KimiVLConfig": "vllm.transformers_utils.configs.kimi_vl",
@@ -113,6 +117,7 @@
     "IsaacConfig",
     "RWConfig",
     "MiDashengLMConfig",
+    "MingV3ASRConfig",
     "MiniMaxM3Config",
     "MiniMaxM3MTPConfig",
     "MiniMaxM3TextConfig",
@@ -120,6 +125,7 @@
     "Moondream3TextConfig",
     "Moondream3VisionConfig",
     "MossTranscribeDiarizeConfig",
+    "Nemotron3_5TranscribeConfig",
     "MoonViTConfig",
     "KimiLinearConfig",
     "KimiVLConfig",
diff --git a/vllm/transformers_utils/configs/ming_v3_asr.py b/vllm/transformers_utils/configs/ming_v3_asr.py
new file mode 100644
index 000000000000..d057df8b2b55
--- /dev/null
+++ b/vllm/transformers_utils/configs/ming_v3_asr.py
@@ -0,0 +1,85 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+# Adapted from
+# https://huggingface.co/inclusionAI/Ming-3.0-asr-tiny/blob/ac95e7b1678c08ea6f4a2581a051e432ed63e9b7/configuration_ming_v3_asr.py
+# Copyright 2025 Antgroup and The HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+"""Typed configuration for Ming V3 ASR and its reusable audio tower."""
+
+from typing import Any
+
+from transformers import PreTrainedConfig
+
+from .bailing_moe_v3_vl import BailingMoeV3TextConfig
+
+
+class MingAudioConfig(PreTrainedConfig):
+    model_type = "ming_audio_encoder"
+
+    def __init__(
+        self,
+        n_mels: int = 128,
+        n_audio_ctx: int = 3000,
+        n_audio_state: int = 1280,
+        n_audio_head: int = 20,
+        n_audio_layer: int = 32,
+        downsampler_kernel_size_list: list[int] | None = None,
+        downsampler_stride_list: list[int] | None = None,
+        downsampler_padding_list: list[int] | None = None,
+        chunk_frames: int = 3000,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.n_mels = n_mels
+        self.n_audio_ctx = n_audio_ctx
+        self.n_audio_state = n_audio_state
+        self.n_audio_head = n_audio_head
+        self.n_audio_layer = n_audio_layer
+        self.downsampler_kernel_size_list = downsampler_kernel_size_list or [3]
+        self.downsampler_stride_list = downsampler_stride_list or [2]
+        self.downsampler_padding_list = (
+            downsampler_padding_list
+            if downsampler_padding_list is not None
+            else [kernel // 2 for kernel in self.downsampler_kernel_size_list]
+        )
+        self.chunk_frames = chunk_frames
+
+
+class MingV3ASRConfig(PreTrainedConfig):
+    model_type = "ming_v3_asr"
+    sub_configs = {
+        "text_config": BailingMoeV3TextConfig,
+        "audio_config": MingAudioConfig,
+    }
+
+    def __init__(
+        self,
+        text_config: dict | BailingMoeV3TextConfig | None = None,
+        audio_config: dict | MingAudioConfig | None = None,
+        audio_token_id: int = 157168,
+        **kwargs: Any,
+    ) -> None:
+        self.text_config = (
+            BailingMoeV3TextConfig(**text_config)
+            if isinstance(text_config, dict)
+            else text_config or BailingMoeV3TextConfig()
+        )
+        self.audio_config = (
+            MingAudioConfig(**audio_config)
+            if isinstance(audio_config, dict)
+            else audio_config or MingAudioConfig()
+        )
+        self.audio_token_id = audio_token_id
+        super().__init__(**kwargs)
diff --git a/vllm/transformers_utils/configs/nemotron3_5_transcribe.py b/vllm/transformers_utils/configs/nemotron3_5_transcribe.py
new file mode 100644
index 000000000000..f90207dd03e6
--- /dev/null
+++ b/vllm/transformers_utils/configs/nemotron3_5_transcribe.py
@@ -0,0 +1,334 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+"""Configuration for NVIDIA Nemotron-3.5-Transcribe.
+
+The checkpoint's ``config.json`` is a NeMo SpeechLM config: the decoder under
+``llm_config``, the speech front end under ``perception`` and the Parallel
+Expert Encoder (an ASR encoder plus a Sortformer speaker encoder) under
+``pe_encoder_config``. Everything the audio tower needs is resolved into one
+flat ``audio_config`` dict.
+"""
+
+import math
+from typing import Any
+
+from transformers import AutoConfig, PreTrainedConfig
+
+# The trained output vocabulary is followed by this many embedding rows. The
+# audio placeholder token lives in that range.
+SPEECHLM_EMBED_EXTRA_ROWS = 10
+
+_ARCHITECTURE = "Nemotron3_5TranscribeForConditionalGeneration"
+_DECODER_ARCHITECTURE = "NemotronHForCausalLM"
+_HYBRID_ARCHITECTURES = ("NemotronHForCausalLM", "NemotronHybridForCausalLM")
+_SUPPORTED_SPEAKER_FEATURE_MODES = ("thresholded", "continuous")
+
+
+def _encoder_config(cfg: dict[str, Any], name: str) -> dict[str, Any]:
+    """Resolve a NeMo ``TransformerEncoder`` config, applying its defaults."""
+    subsampling = cfg.get("subsampling", "feature_stacking")
+    if subsampling != "feature_stacking":
+        raise ValueError(
+            f"{name}: only feature_stacking subsampling is supported, "
+            f"got {subsampling!r}."
+        )
+    if cfg.get("self_attention_model", "rel_pos") != "rope":
+        raise ValueError(f"{name}: only rope self-attention is supported.")
+    if cfg.get("attn_mode", "full") != "full":
+        raise ValueError(f"{name}: only full attention is supported.")
+    if int(cfg.get("causal_tail_len", 0)) != 0:
+        raise ValueError(f"{name}: causal_tail_len must be 0 at inference.")
+    feat_out = int(cfg.get("feat_out", -1))
+    d_model = int(cfg["d_model"])
+    if feat_out > 0 and feat_out != d_model:
+        raise ValueError(
+            f"{name}: an output projection (feat_out={feat_out}) is not supported."
+        )
+    return {
+        "feat_in": int(cfg.get("feat_in", 128)),
+        "d_model": d_model,
+        "n_heads": int(cfg["n_heads"]),
+        "n_layers": int(cfg["n_layers"]),
+        "ff_expansion": float(cfg.get("ff_expansion", 4.0)),
+        "qkv_bias": bool(cfg.get("qkv_bias", False)),
+        "qk_norm": bool(cfg.get("qk_norm", False)),
+        "pre_block_norm": bool(cfg.get("pre_block_norm", True)),
+        "subsampling_factor": int(cfg.get("subsampling_factor", 4)),
+        "rope_base": float(cfg.get("rope_base", 10000.0)),
+        "rotary_fraction": float(cfg.get("rotary_fraction", 1.0)),
+        "xscaling": bool(cfg.get("xscaling", False)),
+    }
+
+
+def _preprocessor_config(cfg: dict[str, Any]) -> dict[str, Any]:
+    """Resolve a NeMo ``AudioToMelSpectrogramPreprocessor`` config."""
+    sample_rate = int(cfg.get("sample_rate", 16000))
+    # NeMo normalizes per feature when the key is missing.
+    if cfg.get("normalize", "per_feature") not in (None, "", "NA"):
+        raise ValueError(
+            "The perception preprocessor must not normalize; the encoders normalize "
+            "their inputs."
+        )
+    if bool(cfg.get("exact_pad", False)):
+        raise ValueError("exact_pad preprocessing is not supported.")
+    if int(cfg.get("frame_splicing", 1)) != 1:
+        raise ValueError("frame_splicing is not supported.")
+    if cfg.get("window", "hann") != "hann":
+        raise ValueError("Only the hann window is supported.")
+    if not bool(cfg.get("log", True)):
+        raise ValueError("Only log-mel features are supported (log=true).")
+    if cfg.get("log_zero_guard_type", "add") != "add":
+        raise ValueError("Only log_zero_guard_type='add' is supported.")
+    win_length = int(float(cfg.get("window_size", 0.02)) * sample_rate)
+    return {
+        "sample_rate": sample_rate,
+        "win_length": win_length,
+        "hop_length": int(float(cfg.get("window_stride", 0.01)) * sample_rate),
+        # NeMo's default: the next power of two at or above the window length.
+        "n_fft": int(cfg.get("n_fft") or 2 ** math.ceil(math.log2(win_length))),
+        "n_mels": int(cfg.get("features", 64)),
+        "preemph": float(cfg.get("preemph", 0.97)),
+        "log_zero_guard_value": float(cfg.get("log_zero_guard_value", 2**-24)),
+        "mag_power": float(cfg.get("mag_power", 2.0)),
+        "pad_value": float(cfg.get("pad_value", 0.0)),
+    }
+
+
+def _checked_online(online: dict[str, int]) -> dict[str, int]:
+    """Reject window settings the native encoders do not implement."""
+    if online["window_frames"] <= 0:
+        raise ValueError(
+            "The encoders run in windows; online_inference_length must be positive, "
+            f"got {online['window_frames']}."
+        )
+    for key in ("left_context_frames", "right_context_frames", "fifo_len"):
+        if online[key] < 0:
+            raise ValueError(f"{key} must not be negative, got {online[key]}.")
+    for key in ("spkcache_update_period", "spkcache_len"):
+        if online[key] <= 0:
+            raise ValueError(f"{key} must be positive, got {online[key]}.")
+    return online
+
+
+def resolve_audio_config(
+    perception: dict[str, Any], pe_encoder_config: dict[str, Any]
+) -> dict[str, Any]:
+    """Flatten the NeMo perception and Parallel Expert Encoder configs."""
+    if not perception:
+        raise ValueError("Nemotron-3.5-Transcribe config must declare perception.")
+    if not pe_encoder_config:
+        raise ValueError(
+            "Nemotron-3.5-Transcribe config must declare pe_encoder_config."
+        )
+    asr_type = pe_encoder_config.get("asr_encoder_type", "fastconformer")
+    if asr_type != "transformer":
+        raise ValueError(
+            f"Only a transformer ASR encoder is supported, got {asr_type!r}."
+        )
+
+    diar = pe_encoder_config["diarization_model_cfg"]
+    sortformer = diar.get("sortformer_modules", {})
+    asr_encoder = _encoder_config(
+        pe_encoder_config["asr_encoder_cfg"], "asr_encoder_cfg"
+    )
+    speaker_encoder = _encoder_config(diar["encoder"], "diarization_model_cfg.encoder")
+    if speaker_encoder["subsampling_factor"] != asr_encoder["subsampling_factor"]:
+        raise ValueError(
+            "The speaker and ASR encoders must share one subsampling factor."
+        )
+    if diar.get("transformer_encoder") is not None:
+        raise ValueError("A Sortformer transformer_encoder stage is not supported.")
+    if bool(diar.get("async_streaming", False)):
+        raise ValueError("Asynchronous Sortformer streaming is not supported.")
+    if not bool(diar.get("high_resolution", False)):
+        raise ValueError("The speaker encoder must be a high-resolution Sortformer.")
+    if not bool(sortformer.get("use_learnable_sil_emb", False)):
+        raise ValueError("The speaker encoder must use a learnable silence embedding.")
+    if float(sortformer.get("activity_weight", diar.get("activity_weight", 0.0))) > 0.0:
+        raise ValueError("A Sortformer activity head is not supported.")
+    diar_normalize = pe_encoder_config.get("diar_normalize_type")
+    if diar_normalize is None:
+        diar_normalize = diar.get("preprocessor", {}).get("normalize")
+    if diar_normalize not in (None, "", "NA"):
+        raise ValueError(
+            f"Speaker-encoder input normalization {diar_normalize!r} is not supported."
+        )
+    asr_normalize = pe_encoder_config.get("asr_normalize_type") or "per_feature"
+    if asr_normalize != "per_feature":
+        raise ValueError(
+            f"ASR-encoder input normalization {asr_normalize!r} is not supported."
+        )
+
+    mode = str(pe_encoder_config.get("speaker_feature_mode", "thresholded")).lower()
+    if mode not in _SUPPORTED_SPEAKER_FEATURE_MODES:
+        raise ValueError(f"Unsupported speaker_feature_mode {mode!r}.")
+    threshold = pe_encoder_config.get("speaker_activity_threshold")
+    if mode == "thresholded" and threshold is None:
+        raise ValueError(
+            "speaker_feature_mode='thresholded' requires speaker_activity_threshold."
+        )
+
+    modality_adapter = perception.get("modality_adapter", {})
+    if not str(modality_adapter.get("_target_", "")).endswith("IdentityConnector"):
+        raise ValueError("Only the identity modality adapter is supported.")
+    projection_in = int(modality_adapter.get("d_model", asr_encoder["d_model"]))
+    if projection_in != asr_encoder["d_model"]:
+        raise ValueError("The projection input must match the ASR encoder width.")
+
+    return {
+        "preprocessor": _preprocessor_config(perception["preprocessor"]),
+        "asr_encoder": asr_encoder,
+        "speaker_encoder": speaker_encoder,
+        "sortformer": {
+            "num_speakers": int(
+                sortformer.get("num_spks", diar.get("max_num_of_spks", 4))
+            ),
+            "fc_d_model": int(sortformer.get("fc_d_model", speaker_encoder["d_model"])),
+            "tf_d_model": int(sortformer.get("tf_d_model", 192)),
+            "spkcache_sil_frames_per_spk": int(
+                sortformer.get("spkcache_sil_frames_per_spk", 3)
+            ),
+            "pred_score_threshold": float(sortformer.get("pred_score_threshold", 0.25)),
+            "max_index": int(sortformer.get("max_index", 99999)),
+            "scores_boost_latest": float(sortformer.get("scores_boost_latest", 0.05)),
+            "strong_boost_rate": float(sortformer.get("strong_boost_rate", 0.75)),
+            "weak_boost_rate": float(sortformer.get("weak_boost_rate", 1.5)),
+            "min_pos_scores_rate": float(sortformer.get("min_pos_scores_rate", 0.5)),
+        },
+        "online": _checked_online(
+            {
+                "window_frames": int(
+                    pe_encoder_config.get("online_inference_length", 500)
+                ),
+                "left_context_frames": int(
+                    pe_encoder_config.get("chunk_left_context", 50)
+                ),
+                "right_context_frames": int(
+                    pe_encoder_config.get("chunk_right_context", 50)
+                ),
+                "fifo_len": int(pe_encoder_config.get("diar_fifo_len", 40)),
+                "spkcache_update_period": int(
+                    pe_encoder_config.get("diar_spkcache_update_period", 300)
+                ),
+                "spkcache_len": int(pe_encoder_config.get("diar_spkcache_len", 188)),
+            }
+        ),
+        "fusion": {
+            "speaker_feature_mode": mode,
+            "speaker_activity_threshold": None
+            if threshold is None
+            else float(threshold),
+            "spk_kernel_scale": float(pe_encoder_config.get("spk_kernel_scale", 1.0)),
+        },
+        "projection": {
+            "input_dim": projection_in,
+            "output_dim": int(perception["output_dim"]),
+        },
+    }
+
+
+def _build_text_config(llm_config: dict[str, Any] | None) -> PreTrainedConfig:
+    if llm_config is None:
+        return AutoConfig.for_model("nemotron_h")
+    if not isinstance(llm_config, dict):
+        raise ValueError(f"llm_config must be a dict, got {type(llm_config).__name__}.")
+    llm_config = dict(llm_config)
+    model_type = llm_config.pop("model_type", None)
+    if not model_type:
+        raise ValueError("llm_config must declare model_type.")
+    return AutoConfig.for_model(model_type, **llm_config)
+
+
+def _decoder_architectures(text_config: PreTrainedConfig) -> list[str]:
+    return getattr(text_config, "architectures", None) or [_DECODER_ARCHITECTURE]
+
+
+def _other_architecture_hint(
+    pe_encoder_config: dict[str, Any] | None, text_config: PreTrainedConfig
+) -> str:
+    # NeMo names checkpoints with this structure after this model; older exports
+    # carry NeMo's generic SpeechLM name.
+    if pe_encoder_config and set(_decoder_architectures(text_config)) & set(
+        _HYBRID_ARCHITECTURES
+    ):
+        return (
+            "This checkpoint has the Parallel Expert Encoder on a NemotronH decoder: "
+            f'set "architectures": ["{_ARCHITECTURE}"] in its config.json.'
+        )
+    return "Serve other NeMo SpeechLM models with the NeMo Speech vLLM plugin."
+
+
+class Nemotron3_5TranscribeConfig(PreTrainedConfig):
+    """Configuration for Nemotron-3.5-Transcribe.
+
+    Reads the checkpoint's NeMo SpeechLM ``config.json`` (see the module
+    docstring). The decoder config becomes ``text_config``; the audio tower
+    settings become ``audio_config`` (see :func:`resolve_audio_config`).
+
+    ``nemo_speechlm`` is NeMo's model type for all its SpeechLM checkpoints; the
+    architecture name selects Nemotron-3.5-Transcribe.
+    """
+
+    model_type = "nemo_speechlm"
+
+    def __init__(
+        self,
+        llm_config: dict[str, Any] | None = None,
+        perception: dict[str, Any] | None = None,
+        pe_encoder_config: dict[str, Any] | None = None,
+        text_config: dict[str, Any] | PreTrainedConfig | None = None,
+        audio_config: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        # Transformers may read the text config while the base class
+        # initializes; give it the real one before calling super().__init__.
+        if llm_config is not None or text_config is None:
+            self.text_config = _build_text_config(llm_config)
+        elif isinstance(text_config, PreTrainedConfig):
+            self.text_config = text_config
+        else:
+            self.text_config = _build_text_config(text_config)
+        super().__init__(**kwargs)
+        if self.architectures and self.architectures != [_ARCHITECTURE]:
+            raise ValueError(
+                f"Only {_ARCHITECTURE} is supported among NeMo SpeechLM checkpoints, "
+                f"got architectures={self.architectures!r}. "
+                + _other_architecture_hint(pe_encoder_config, self.text_config)
+            )
+
+        self.llm_config = llm_config
+        self.perception = perception or {}
+        self.pe_encoder_config = pe_encoder_config or {}
+
+        decoder_config = self.text_config
+        architectures = _decoder_architectures(decoder_config)
+        if len(architectures) != 1 or architectures[0] not in _HYBRID_ARCHITECTURES:
+            raise ValueError(
+                f"Expected a NemotronH decoder, got architectures={architectures!r}."
+            )
+        decoder_config.architectures = [_DECODER_ARCHITECTURE]
+
+        if (
+            llm_config is not None
+            and getattr(decoder_config, "speechlm_output_vocab_size", None) is None
+        ):
+            decoder_config.speechlm_output_vocab_size = int(decoder_config.vocab_size)
+            decoder_config.vocab_size = (
+                int(decoder_config.vocab_size) + SPEECHLM_EMBED_EXTRA_ROWS
+            )
+        self.output_vocab_size = int(
+            getattr(
+                decoder_config, "speechlm_output_vocab_size", decoder_config.vocab_size
+            )
+        )
+
+        if self.perception:
+            self.audio_config = resolve_audio_config(
+                self.perception, self.pe_encoder_config
+            )
+        else:
+            self.audio_config = audio_config or {}
+
+    def get_text_config(self, decoder: bool = False) -> PreTrainedConfig:
+        return self.text_config
diff --git a/vllm/transformers_utils/configs/radio.py b/vllm/transformers_utils/configs/radio.py
index 556e1b451a24..671985fec6e8 100644
--- a/vllm/transformers_utils/configs/radio.py
+++ b/vllm/transformers_utils/configs/radio.py
@@ -1,6 +1,6 @@
 # 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
 
@@ -37,6 +37,7 @@ class RadioConfig(PreTrainedConfig):
         norm_type: The normalization type to use.
         layer_norm_eps: The epsilon used by the layer normalization layers.
         initializer_factor: A factor for initializing all weight matrices.
+        layerscale_value: Initial value for the LayerScale ``lambda1`` parameters.
         hidden_act: The non-linear activation function in the encoder.
         cpe_max_size: Maximum image size for position embeddings.
         norm_mean: Mean values for image normalization (RGB channels).
@@ -44,6 +45,8 @@ class RadioConfig(PreTrainedConfig):
         norm_std: Standard deviation values for image normalization
             (RGB channels). Defaults to (0.26862954, 0.26130258, 0.27577711)).
         register_multiple: Number of register tokens to use.
+        cpe_num_registers: Explicit CPE register-token count, used by checkpoints
+            that set ``register_multiple=0`` (e.g. C-RADIOv3) instead.
         teachers: A list of teacher model configurations. Each teacher configuration is
             a dict with keys like "name" and some may have "use_summary".
         cls_token_per_teacher: Whether to use a separate CLS token for each teacher.
@@ -55,6 +58,11 @@ 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).
+        num_channels: Number of input image channels (RGB -> 3).
+
+    ``num_cls_tokens``, ``num_registers``, ``summary_idxs`` and ``max_img_size``
+    are derived from the fields above so the model reads a single canonical
+    schema (these were previously computed inside the model).
 
     """
 
@@ -70,15 +78,18 @@ def __init__(
         norm_type: str = "layer_norm",
         layer_norm_eps: float = 1e-6,
         initializer_factor: float = 1.0,
+        layerscale_value: float = 1.0,
         hidden_act: str = "gelu",
         cpe_max_size: int = 2048,
         norm_mean: tuple[float, float, float] | list = OPENAI_CLIP_MEAN,
         norm_std: tuple[float, float, float] | list = OPENAI_CLIP_STD,
         register_multiple: int | None = None,
+        cpe_num_registers: int | None = None,
         teachers: list[dict[str, Any]] | None = None,
         cls_token_per_teacher: bool = False,
         video_temporal_patch_size: int = 1,
         separate_video_embedder: bool = True,
+        num_channels: int = 3,
         **kwargs,
     ):
         self.model_name = model_name
@@ -95,6 +106,7 @@ def __init__(
         self.norm_type = norm_type
         self.layer_norm_eps = layer_norm_eps
         self.initializer_factor = initializer_factor
+        self.layerscale_value = layerscale_value
         self.hidden_act = hidden_act
         self.cpe_max_size = cpe_max_size
         self.norm_mean = (
@@ -104,8 +116,35 @@ def __init__(
             list(norm_std) if isinstance(norm_std, (tuple, list)) else norm_std
         )
         self.register_multiple = register_multiple
+        self.cpe_num_registers = cpe_num_registers
         self.teachers = teachers if teachers is not None else []
         self.cls_token_per_teacher = cls_token_per_teacher
         self.video_temporal_patch_size = video_temporal_patch_size
         self.separate_video_embedder = separate_video_embedder
+        self.num_channels = num_channels
+
+        # Fields derived from the fields above (previously computed inside the
+        # model), so the model consumes a single canonical schema.
+        unique_teachers = {t["name"] for t in self.teachers}
+        self.num_cls_tokens = len(unique_teachers) if self.cls_token_per_teacher else 1
+        if self.cpe_num_registers:
+            # An explicit register count wins, matching the reference ClsToken.
+            # C-RADIOv3 carries its count here and sets register_multiple=0.
+            self.num_registers = self.cpe_num_registers
+        elif self.register_multiple:
+            self.num_registers = self.register_multiple - (
+                self.num_cls_tokens % self.register_multiple
+            )
+        else:
+            self.num_registers = 0
+        # None (no teachers) means no explicit selection -> keep all class
+        # tokens; a list (possibly empty) gathers exactly those indices.
+        self.summary_idxs = (
+            [i for i, t in enumerate(self.teachers) if t.get("use_summary", True)]
+            if self.teachers
+            else None
+        )
+        self.max_img_size = int(
+            round(self.cpe_max_size / self.patch_size) * self.patch_size
+        )
         super().__init__(**kwargs)
diff --git a/vllm/transformers_utils/processors/__init__.py b/vllm/transformers_utils/processors/__init__.py
index 0736e385b928..b7d66a579a44 100644
--- a/vllm/transformers_utils/processors/__init__.py
+++ b/vllm/transformers_utils/processors/__init__.py
@@ -23,6 +23,7 @@
     "KimiK25Processor",
     "KimiK3Processor",
     "MiMoOmniProcessor",
+    "MingASRFeatureExtractor",
     "MiniCPMOProcessor",
     "MiniCPMVProcessor",
     "MiniMaxM3VLImageProcessor",
@@ -53,6 +54,7 @@
     "KimiK25Processor": "vllm.transformers_utils.processors.kimi_k25",
     "KimiK3Processor": "vllm.transformers_utils.processors.kimi_k3",
     "MiMoOmniProcessor": "vllm.transformers_utils.processors.mimo_v2_omni",
+    "MingASRFeatureExtractor": "vllm.transformers_utils.processors.ming_audio",
     "MiniCPMOProcessor": "vllm.transformers_utils.processors.minicpmo",
     "MiniCPMVProcessor": "vllm.transformers_utils.processors.minicpmv",
     "MiniMaxM3VLImageProcessor": "vllm.transformers_utils.processors.minimax_m3",
diff --git a/vllm/transformers_utils/processors/ming_audio.py b/vllm/transformers_utils/processors/ming_audio.py
new file mode 100644
index 000000000000..9226f6e9cf03
--- /dev/null
+++ b/vllm/transformers_utils/processors/ming_audio.py
@@ -0,0 +1,35 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Whisper feature extraction with Ming's mel coefficient precision."""
+
+import numpy as np
+from transformers import WhisperFeatureExtractor
+from transformers.audio_utils import hertz_to_mel, mel_filter_bank, mel_to_hertz
+
+
+class MingASRFeatureExtractor(WhisperFeatureExtractor):
+    def __init__(self, feature_size: int = 128) -> None:
+        if feature_size not in (80, 128):
+            raise ValueError("Ming audio requires 80 or 128 mel bins")
+        super().__init__(feature_size=feature_size)
+        max_frequency = self.sampling_rate / 2
+        filters = mel_filter_bank(
+            num_frequency_bins=1 + self.n_fft // 2,
+            num_mel_filters=feature_size,
+            min_frequency=0.0,
+            max_frequency=max_frequency,
+            sampling_rate=self.sampling_rate,
+            norm=None,
+            mel_scale="slaney",
+        )
+        centers = mel_to_hertz(
+            np.linspace(
+                0.0,
+                hertz_to_mel(max_frequency, mel_scale="slaney"),
+                feature_size + 2,
+            ),
+            mel_scale="slaney",
+        )
+        # Round before Slaney scaling and retain contiguous mel-major storage.
+        self.mel_filters = np.asfortranarray(filters, dtype=np.float32)
+        self.mel_filters *= (2.0 / (centers[2:] - centers[:-2]))[None, :]
diff --git a/vllm/transformers_utils/processors/minimax_m3.py b/vllm/transformers_utils/processors/minimax_m3.py
index 9d922f8e4c64..8d4212cee395 100644
--- a/vllm/transformers_utils/processors/minimax_m3.py
+++ b/vllm/transformers_utils/processors/minimax_m3.py
@@ -1,16 +1,17 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""MiniMax M3 VL HuggingFace-compatible Processor / ImageProcessor /
+"""MiniMax M3.1 VL HuggingFace-compatible Processor / ImageProcessor /
 VideoProcessor, vendored into vLLM so the model loads without
 ``--trust-remote-code`` (the released checkpoint only ships these classes as
 remote code via ``auto_map``).
 
-Adapted verbatim from the ``MiniMaxAI/Minimax-M3-preview`` repository files
-``image_processor.py``, ``video_processor.py`` and ``processing_minimax.py``
-(revision ``db01c0fe``). Both image and video processors use Qwen-style
-``smart_resize`` (bound by total pixels). The original async frame-sampling
-helpers are intentionally omitted: vLLM performs its own frame loading and
-feeds decoded frames to the processor.
+Both image and video processors use the MiniMax M3.1 aligned-size resize
+(``compute_aligned_size``): tiny inputs are upscaled to a minimum area /
+short side, then both dims are ceil-aligned to the patch factor and capped
+at a per-modality long side (selectable via the OpenAI-style ``detail``
+kwarg or an explicit ``max_long_side_pixel``). The original async
+frame-sampling helpers are intentionally omitted: vLLM performs its own
+frame loading and feeds decoded frames to the processor.
 """
 
 import math
@@ -35,125 +36,158 @@
 from transformers.video_processing_utils import BaseVideoProcessor
 from transformers.video_utils import group_videos_by_shape, reorder_videos
 
-# Maximum allowed aspect ratio before smart_resize rejects the input.
+# Maximum allowed aspect ratio before compute_aligned_size rejects the input.
 MAX_RATIO = 200
 
-# Fixed (non-configurable) bounds for the long-side resize logic, per the
-# MiniMax-M3 size spec. ``min_short_side_pixel`` is the floor the short edge is
-# enlarged to; ``*_MAX_TOTAL_PIXELS`` is the hard area cap that, once exceeded,
-# aborts processing instead of downscaling.
-MIN_SHORT_SIDE_PIXEL = 112
-IMAGE_MAX_TOTAL_PIXELS = 12_845_056  # 3584 ** 2 (width * height)
-VIDEO_MAX_TOTAL_PIXELS = 301_056_000  # width * height * frames
+# Fixed (non-configurable) bounds of the aligned-size resize, per the
+# MiniMax M3.1 size spec. ``MIN_SHORT_SIDE_PIXEL`` is the floor the short
+# edge is enlarged to; ``VIDEO_MAX_TOTAL_PIXELS`` is the hard volumetric cap
+# (width * height * frames) that, once exceeded, aborts video processing.
+MIN_SHORT_SIDE_PIXEL = 168
+MIN_PIXELS = MIN_SHORT_SIDE_PIXEL**2
+VIDEO_MAX_TOTAL_PIXELS = 301_056_000
 
+# OpenAI-style ``detail`` -> long-side pixel cap, per modality.
+IMAGE_DETAIL_MAX_LONG_SIDE_PIXEL = {"low": 672, "auto": 2016, "high": 3584}
+VIDEO_DETAIL_MAX_LONG_SIDE_PIXEL = {"low": 504, "auto": 672, "high": 2188}
 
-def round_by_factor(number: int | float, factor: int) -> int:
-    return round(number / factor) * factor
 
-
-def ceil_by_factor(number: int | float, factor: int) -> int:
-    return math.ceil(number / factor) * factor
-
-
-def floor_by_factor(number: int | float, factor: int) -> int:
-    return math.floor(number / factor) * factor
-
-
-def _smart_resize_by_long_side(
+def compute_aligned_size(
     height: int,
     width: int,
     factor: int,
+    *,
+    min_pixels: int = MIN_PIXELS,
     max_long_side_pixel: int,
-    min_short_side_pixel: int,
-    max_total_pixels: int | None,
+    min_short_side: int = MIN_SHORT_SIDE_PIXEL,
 ) -> tuple[int, int]:
-    """Long-side based resize (MiniMax-M3 size spec).
-
-    (a) if the long side exceeds ``max_long_side_pixel`` → shrink so the long
-        side equals ``max_long_side_pixel``;
-    (b) else if the short side is below ``min_short_side_pixel`` → enlarge so the
-        short side equals ``min_short_side_pixel``;
-    (c) if the resulting area still exceeds ``max_total_pixels`` → raise.
-
-    (a) and (b) are mutually exclusive (they branch on the *original* long side).
-    Both sides are then rounded to a multiple of ``factor``. For videos the
-    ``max_total_pixels`` cap is volumetric (width * height * frames) and is
-    enforced by the caller, so pass ``max_total_pixels=None`` here.
-    """
-    long_side = max(height, width)
-    short_side = min(height, width)
-
-    scaled_height: float = height
-    scaled_width: float = width
-    if long_side > max_long_side_pixel:
-        beta = max_long_side_pixel / long_side
-        scaled_height = height * beta
-        scaled_width = width * beta
-    elif short_side < min_short_side_pixel:
-        beta = min_short_side_pixel / short_side
-        scaled_height = height * beta
-        scaled_width = width * beta
-
-    h_bar = max(factor, round_by_factor(scaled_height, factor))
-    w_bar = max(factor, round_by_factor(scaled_width, factor))
-
-    if max_total_pixels is not None and h_bar * w_bar > max_total_pixels:
-        raise ValueError(
-            f"image area {h_bar * w_bar} exceeds max_total_pixels "
-            f"{max_total_pixels} after resizing"
-        )
-    return h_bar, w_bar
+    """Compute the resize target ``(resized_height, resized_width)``.
 
+    Three layers (mirrors sglang's M3.1 ``compute_aligned_size``):
 
-def smart_resize(
-    height: int,
-    width: int,
-    factor: int = 28,
-    min_pixels: int = 4 * 28 * 28,
-    max_pixels: int = 451584,
-    max_long_side_pixel: int | None = None,
-    min_short_side_pixel: int = MIN_SHORT_SIDE_PIXEL,
-    max_total_pixels: int | None = None,
-) -> tuple[int, int]:
-    """Rescale (height, width) so each side is a multiple of ``factor``.
+    1. Aspect-ratio guard: reject inputs whose ratio exceeds ``MAX_RATIO``.
+    2. Floor layers: upscale proportionally when the area is below
+       ``min_pixels``, or the short side is below ``min_short_side``.
+    3. Alignment layer: ``_align_to_factor`` ceil-aligns both dims to
+       ``factor`` and caps them at ``max_long_side_pixel``.
+
+    Args:
+        height: Original height in pixels, must be > 0.
+        width: Original width in pixels, must be > 0.
+        factor: Alignment factor, usually ``patch_size * merge_size``.
+        min_pixels: Area floor; smaller inputs are upscaled to at least
+            this area before alignment.
+        max_long_side_pixel: Per-dim cap applied by the alignment layer
+            (used for both dims).
+        min_short_side: Short-side pixel floor.
+
+    Returns:
+        ``(resized_height, resized_width)``, both multiples of ``factor``
+        and at most ``max_long_side_pixel``.
+
+    Raises:
+        ValueError: If the aspect ratio exceeds ``MAX_RATIO``.
 
-    When ``max_long_side_pixel`` is set, use the MiniMax-M3 long-side resize
-    spec (see :func:`_smart_resize_by_long_side`). Otherwise fall back to the
-    Qwen-VL area bound, keeping the total area within ``[min_pixels, max_pixels]``.
     """
     if max(height, width) / min(height, width) > MAX_RATIO:
         raise ValueError(
             f"absolute aspect ratio must be smaller than {MAX_RATIO}, "
             f"got {max(height, width) / min(height, width)}"
         )
-    if max_long_side_pixel is not None:
-        return _smart_resize_by_long_side(
-            height,
-            width,
-            factor=factor,
-            max_long_side_pixel=max_long_side_pixel,
-            min_short_side_pixel=min_short_side_pixel,
-            max_total_pixels=max_total_pixels,
-        )
-    h_bar = max(factor, round_by_factor(height, factor))
-    w_bar = max(factor, round_by_factor(width, factor))
-    if h_bar * w_bar > max_pixels:
-        beta = math.sqrt((height * width) / max_pixels)
-        h_bar = floor_by_factor(height / beta, factor)
-        w_bar = floor_by_factor(width / beta, factor)
-    elif h_bar * w_bar < min_pixels:
+    if height * width < min_pixels:
         beta = math.sqrt(min_pixels / (height * width))
-        h_bar = ceil_by_factor(height * beta, factor)
-        w_bar = ceil_by_factor(width * beta, factor)
-    return h_bar, w_bar
+        height = math.ceil(height * beta)
+        width = math.ceil(width * beta)
+    if min(height, width) < min_short_side:
+        beta = min_short_side / min(height, width)
+        height = math.ceil(height * beta)
+        width = math.ceil(width * beta)
+    return _align_to_factor(
+        height, width, factor, (max_long_side_pixel, max_long_side_pixel)
+    )
+
+
+def _align_to_factor(
+    height: int,
+    width: int,
+    factor: int,
+    max_size: tuple[int, int],
+) -> tuple[int, int]:
+    """Ceil-align both dims to ``factor``, then cap at per-dim ``max_size``.
+
+    When either aligned dim exceeds its cap, scale both dims back down
+    proportionally (integer arithmetic, preserving the aspect ratio) and
+    ceil-align again.
+
+    Returns:
+        ``(height, width)`` with both dims multiples of ``factor`` and at
+        most the corresponding ``max_size`` dim.
+
+    """
+    max_w, max_h = max_size
+    # Round the caps down to factor multiples so alignment never exceeds them.
+    max_w = max_w - (max_w % factor)
+    max_h = max_h - (max_h % factor)
+    new_w = width if width % factor == 0 else width + (factor - width % factor)
+    new_h = height if height % factor == 0 else height + (factor - height % factor)
+    if new_w > max_w or new_h > max_h:
+        new_w_ = min((new_w * max_w) // new_w, (new_w * max_h) // new_h)
+        new_h_ = min((new_h * max_w) // new_w, (new_h * max_h) // new_h)
+        new_w = new_w_
+        new_h = new_h_
+        new_w = new_w if new_w % factor == 0 else new_w + (factor - new_w % factor)
+        new_h = new_h if new_h % factor == 0 else new_h + (factor - new_h % factor)
+    assert new_w <= max_w and new_h <= max_h, (
+        f"_align_to_factor: result (h={new_h}, w={new_w}) exceeds max_size "
+        f"{max_size}; input (h={height}, w={width}, factor={factor})"
+    )
+    return new_h, new_w
+
+
+def resolve_max_long_side_pixel(
+    detail: "str | None",
+    max_long_side_pixel: "int | None",
+    *,
+    modality: str,
+) -> int:
+    """Resolve the long-side pixel cap from ``detail`` / explicit override.
+
+    ``detail`` (default ``"auto"``) maps to a per-modality cap. An explicit
+    ``max_long_side_pixel`` takes priority over a non-auto ``detail``; the
+    ``"auto"`` tier value itself is treated as "unset" so the processor class
+    defaults (2016 for images, 672 for videos) do not shadow ``detail``.
+
+    Raises:
+        ValueError: If ``detail`` or ``modality`` is not a known value.
+
+    """
+    if modality == "image":
+        detail_map = IMAGE_DETAIL_MAX_LONG_SIDE_PIXEL
+    elif modality == "video":
+        detail_map = VIDEO_DETAIL_MAX_LONG_SIDE_PIXEL
+    else:
+        raise ValueError(f"unknown modality {modality!r}")
+    detail = detail or "auto"
+    if detail not in detail_map:
+        raise ValueError(
+            f"invalid detail {detail!r} for {modality}; "
+            f"expected one of {sorted(detail_map)}"
+        )
+    if max_long_side_pixel is not None and max_long_side_pixel != detail_map["auto"]:
+        return max_long_side_pixel
+    return detail_map[detail]
 
 
 class MiniMaxM3VLImageProcessorKwargs(ImagesKwargs, total=False):  # type: ignore[call-arg]
     patch_size: int
     temporal_patch_size: int
     merge_size: int
+    detail: "str | None"
+    max_long_side_pixel: "int | None"
+    # Legacy keys from released preprocessor_config.json files; accepted for
+    # backward compatibility but ignored by the aligned-size resize.
+    min_pixels: int
     max_pixels: int
-    max_long_side_pixel: int
 
 
 class MiniMaxM3VLImageProcessor(TorchvisionBackend):
@@ -171,12 +205,11 @@ class MiniMaxM3VLImageProcessor(TorchvisionBackend):
     patch_size = 14
     temporal_patch_size = 2
     merge_size = 2
-    max_pixels = 451584  # 672 * 672
-    # Long-side resize spec (opt-in via ``max_long_side_pixel``). The latter two
-    # are fixed per the spec and are not exposed as configurable kwargs.
-    max_long_side_pixel = None
+    # Aligned-size resize bounds (M3.1). ``detail`` selects the long-side cap
+    # ("low"/"auto"/"high"); an explicit ``max_long_side_pixel`` wins.
+    detail = "auto"
+    max_long_side_pixel = 2016
     min_short_side_pixel = MIN_SHORT_SIDE_PIXEL
-    max_total_pixels = IMAGE_MAX_TOTAL_PIXELS
     valid_kwargs = MiniMaxM3VLImageProcessorKwargs
     model_input_names = ["pixel_values", "image_grid_thw"]
 
@@ -202,7 +235,7 @@ def _preprocess(
         patch_size: int,
         temporal_patch_size: int,
         merge_size: int,
-        max_pixels: int,
+        detail: str,
         max_long_side_pixel: "int | None",
         disable_grouping: "bool | None",
         return_tensors: "str | TensorType | None",
@@ -213,17 +246,18 @@ def _preprocess(
         )
         resized_images_grouped = {}
         factor = patch_size * merge_size
+        max_long_side = resolve_max_long_side_pixel(
+            detail, max_long_side_pixel, modality="image"
+        )
         for shape, stacked_images in grouped_images.items():
             height, width = stacked_images.shape[-2:]
             if do_resize:
-                resized_height, resized_width = smart_resize(
+                resized_height, resized_width = compute_aligned_size(
                     height,
                     width,
                     factor=factor,
-                    max_pixels=max_pixels,
-                    max_long_side_pixel=max_long_side_pixel,
-                    min_short_side_pixel=self.min_short_side_pixel,
-                    max_total_pixels=self.max_total_pixels,
+                    max_long_side_pixel=max_long_side,
+                    min_short_side=self.min_short_side_pixel,
                 )
                 stacked_images = self.resize(
                     stacked_images,
@@ -308,19 +342,19 @@ def get_number_of_image_patches(self, height: int, width: int, images_kwargs=Non
         images_kwargs = images_kwargs or {}
         patch_size = images_kwargs.get("patch_size", self.patch_size)
         merge_size = images_kwargs.get("merge_size", self.merge_size)
-        max_pixels = images_kwargs.get("max_pixels", self.max_pixels)
+        detail = images_kwargs.get("detail", self.detail)
         max_long_side_pixel = images_kwargs.get(
             "max_long_side_pixel", self.max_long_side_pixel
         )
-
-        resized_height, resized_width = smart_resize(
+        max_long_side = resolve_max_long_side_pixel(
+            detail, max_long_side_pixel, modality="image"
+        )
+        resized_height, resized_width = compute_aligned_size(
             height,
             width,
             factor=patch_size * merge_size,
-            max_pixels=max_pixels,
-            max_long_side_pixel=max_long_side_pixel,
-            min_short_side_pixel=self.min_short_side_pixel,
-            max_total_pixels=self.max_total_pixels,
+            max_long_side_pixel=max_long_side,
+            min_short_side=self.min_short_side_pixel,
         )
         grid_h, grid_w = resized_height // patch_size, resized_width // patch_size
         return grid_h * grid_w
@@ -330,13 +364,16 @@ class MiniMaxM3VLVideoProcessorKwargs(VideosKwargs, total=False):  # type: ignor
     patch_size: int
     temporal_patch_size: int
     merge_size: int
-    min_pixels: int
-    max_pixels: int
-    max_long_side_pixel: int
-    total_pixels: int
+    detail: "str | None"
+    max_long_side_pixel: "int | None"
     min_frames: int
     max_frames: int
     fps: "float | int"
+    # Legacy keys from released preprocessor_config.json files; accepted for
+    # backward compatibility but ignored by the aligned-size resize.
+    min_pixels: int
+    max_pixels: int
+    total_pixels: int
 
 
 class MiniMaxM3VLVideoProcessor(BaseVideoProcessor):
@@ -354,13 +391,12 @@ class MiniMaxM3VLVideoProcessor(BaseVideoProcessor):
     patch_size = 14
     temporal_patch_size = 2
     merge_size = 2
-    min_pixels = 4 * 28 * 28
-    max_pixels = 768 * 28 * 28  # 602,112
-    total_pixels = int(64000 * 28 * 28 * 0.9)  # ~45M, ~64k tokens budget
-    # Long-side resize spec (opt-in via ``max_long_side_pixel``). The video
+    # Aligned-size resize bounds (M3.1). ``detail`` selects the long-side cap
+    # ("low"/"auto"/"high"); an explicit ``max_long_side_pixel`` wins. The
     # ``max_total_pixels`` cap is volumetric (width * height * frames) and is
     # enforced in ``_preprocess`` once the frame count is known.
-    max_long_side_pixel = None
+    detail = "auto"
+    max_long_side_pixel = 672
     min_short_side_pixel = MIN_SHORT_SIDE_PIXEL
     max_total_pixels = VIDEO_MAX_TOTAL_PIXELS
     fps = 1.0
@@ -386,8 +422,7 @@ def _preprocess(
         patch_size: int,
         temporal_patch_size: int,
         merge_size: int,
-        min_pixels: int,
-        max_pixels: int,
+        detail: str = "auto",
         max_long_side_pixel: "int | None" = None,
         return_tensors: "str | TensorType | None" = None,
         **kwargs,
@@ -395,27 +430,21 @@ def _preprocess(
         grouped_videos, grouped_videos_index = group_videos_by_shape(videos)
         resized_videos_grouped = {}
         factor = patch_size * merge_size
+        max_long_side = resolve_max_long_side_pixel(
+            detail, max_long_side_pixel, modality="video"
+        )
         for shape, stacked_videos in grouped_videos.items():
             batch_size, num_frames, channels, height, width = stacked_videos.shape
             resized_height, resized_width = height, width
             if do_resize:
-                resized_height, resized_width = smart_resize(
+                resized_height, resized_width = compute_aligned_size(
                     height,
                     width,
                     factor=factor,
-                    min_pixels=min_pixels,
-                    max_pixels=max_pixels,
-                    max_long_side_pixel=max_long_side_pixel,
-                    min_short_side_pixel=self.min_short_side_pixel,
-                    # Per-frame raise disabled; the video cap is volumetric and
-                    # is enforced below once num_frames is known.
-                    max_total_pixels=None,
+                    max_long_side_pixel=max_long_side,
+                    min_short_side=self.min_short_side_pixel,
                 )
-                if (
-                    max_long_side_pixel is not None
-                    and resized_height * resized_width * num_frames
-                    > self.max_total_pixels
-                ):
+                if resized_height * resized_width * num_frames > self.max_total_pixels:
                     raise ValueError(
                         f"video area {resized_height * resized_width * num_frames} "
                         f"(width * height * frames) exceeds max_total_pixels "
@@ -424,8 +453,10 @@ def _preprocess(
                 stacked_videos = stacked_videos.view(
                     batch_size * num_frames, channels, height, width
                 )
+                # Resize in float32, matching the training/SFT pipeline (and
+                # sglang): uint8 resize rounds intermediates and shifts pixels.
                 stacked_videos = self.resize(
-                    stacked_videos,
+                    stacked_videos.float(),
                     size=SizeDict(height=resized_height, width=resized_width),
                     resample=resample,
                 )
diff --git a/vllm/transformers_utils/processors/nano_nemotron_vl.py b/vllm/transformers_utils/processors/nano_nemotron_vl.py
index 9572b7c4a76d..8e1adbda59e2 100644
--- a/vllm/transformers_utils/processors/nano_nemotron_vl.py
+++ b/vllm/transformers_utils/processors/nano_nemotron_vl.py
@@ -63,6 +63,7 @@ def _bicubic_resize_and_normalize(
     size: tuple[int, int] | None = None,
     norm_mean: torch.Tensor | None = None,
     norm_std: torch.Tensor | None = None,
+    do_cpu_normalize: bool = True,
     dtype: torch.dtype = torch.float32,
 ) -> torch.Tensor:
     """Permute NHWC→NCHW, optional bicubic resize, rescale + normalize.
@@ -70,17 +71,25 @@ def _bicubic_resize_and_normalize(
     Input must be a raw 4-D **NHWC** tensor.
 
     *size*: target ``(H, W)``; skips interpolation when ``None``.
-    *norm_mean* / *norm_std*: when both provided, fused
-    ``(x/255 - mean) / std`` + dtype cast; otherwise ``x/255`` + cast.
+    *norm_mean* / *norm_std*: when both provided, apply normalization.
+    When *do_cpu_normalize* is false, return the resized ``uint8`` tensor
+    unchanged so rescaling and normalization can run on the model device.
     """
-    tensor = tensor.permute(0, 3, 1, 2).to(dtype=torch.float32)
+    tensor = tensor.permute(0, 3, 1, 2)
+    if do_cpu_normalize:
+        tensor = tensor.to(dtype=torch.float32)
     if size is not None:
         tensor = torch.nn.functional.interpolate(
             tensor, size=size, mode="bicubic", align_corners=False, antialias=True
         )
+
+    if not do_cpu_normalize:
+        return tensor.contiguous()
+
+    tensor = tensor / 255.0
     if norm_mean is not None and norm_std is not None:
-        return ((tensor / 255.0 - norm_mean) / norm_std).to(dtype=dtype).contiguous()
-    return (tensor / 255.0).to(dtype=dtype).contiguous()
+        tensor = (tensor - norm_mean) / norm_std
+    return tensor.to(dtype=dtype).contiguous()
 
 
 def _pil_to_nhwc_tensor(image: Image.Image) -> torch.Tensor:
@@ -99,6 +108,7 @@ def dynamic_preprocess(
     use_thumbnail: bool = True,
     norm_mean: torch.Tensor | None = None,
     norm_std: torch.Tensor | None = None,
+    do_cpu_normalize: bool = True,
     dtype: torch.dtype = torch.float32,
 ) -> torch.Tensor:
     orig_width, orig_height = image.size
@@ -120,6 +130,7 @@ def dynamic_preprocess(
         size=(target_height, target_width),
         norm_mean=norm_mean,
         norm_std=norm_std,
+        do_cpu_normalize=do_cpu_normalize,
         dtype=dtype,
     )
     B, C, H, W = resized_img.shape
@@ -136,6 +147,7 @@ def dynamic_preprocess(
             size=(image_size, image_size),
             norm_mean=norm_mean,
             norm_std=norm_std,
+            do_cpu_normalize=do_cpu_normalize,
             dtype=dtype,
         )
         patches = torch.cat([patches, thumb], dim=0)
@@ -225,6 +237,7 @@ def video_to_pixel_values(
     downsample_ratio: float = 0.5,
     norm_mean: torch.Tensor | None = None,
     norm_std: torch.Tensor | None = None,
+    do_cpu_normalize: bool = True,
     dtype: torch.dtype = torch.float32,
 ) -> torch.Tensor:
     """Convert video ndarray (T, H, W, C) to normalized pixel tensor (T, C, H, W)."""
@@ -246,7 +259,14 @@ def video_to_pixel_values(
         size = (input_size, input_size)
 
     tensor = torch.from_numpy(video)
-    return _bicubic_resize_and_normalize(tensor, size, norm_mean, norm_std, dtype)
+    return _bicubic_resize_and_normalize(
+        tensor,
+        size=size,
+        norm_mean=norm_mean,
+        norm_std=norm_std,
+        do_cpu_normalize=do_cpu_normalize,
+        dtype=dtype,
+    )
 
 
 class DynamicResolutionImageTiler:
@@ -264,6 +284,7 @@ def __init__(
         downsample_ratio: int,
         norm_mean: Sequence[float],
         norm_std: Sequence[float],
+        do_cpu_normalize: bool = True,
         factor_max: float = 1.0,
         use_thumbnail: bool = False,
     ) -> None:
@@ -275,6 +296,7 @@ def __init__(
         self._factor_max = factor_max
         self.norm_mean = torch.tensor(norm_mean).reshape(3, 1, 1)
         self.norm_std = torch.tensor(norm_std).reshape(3, 1, 1)
+        self.do_cpu_normalize = do_cpu_normalize
         assert downsample_ratio < 1
         reduction_factor = 1 / downsample_ratio
         assert reduction_factor == 2.0
@@ -370,6 +392,7 @@ def apply_params(
             size=target_size,
             norm_mean=self.norm_mean,
             norm_std=self.norm_std,
+            do_cpu_normalize=self.do_cpu_normalize,
             dtype=dtype,
         )
         return list(resized_img)
@@ -589,6 +612,7 @@ def __init__(
         *args,
         max_model_len: int,
         max_num_tiles: int | None = None,
+        do_cpu_normalize: bool = True,
         **kwargs,
     ) -> None:
         super().__init__()
@@ -606,6 +630,7 @@ def __init__(
         )
         self.image_size = image_size
         self.use_thumbnail: bool = config.use_thumbnail
+        self.do_cpu_normalize = do_cpu_normalize
         self.norm_mean = torch.Tensor(config.norm_mean).reshape(1, 3, 1, 1)
         self.norm_std = torch.Tensor(config.norm_std).reshape(1, 3, 1, 1)
 
@@ -619,6 +644,7 @@ def __init__(
                 max_num_patches=config.vision_config.args["max_num_patches"],
                 norm_mean=config.norm_mean,
                 norm_std=config.norm_std,
+                do_cpu_normalize=do_cpu_normalize,
             )
         self.dtype: torch.dtype = getattr(config, "dtype", torch.float32)
 
@@ -671,6 +697,7 @@ def _images_to_pixel_values_lst(
                 use_thumbnail=self.use_thumbnail,
                 norm_mean=self.norm_mean,
                 norm_std=self.norm_std,
+                do_cpu_normalize=self.do_cpu_normalize,
                 dtype=self.dtype,
             )
             for image in images
@@ -779,12 +806,20 @@ def __init__(
         video_token: str | None = None,
         video_pruning_rate: float | None = None,
         use_audio_in_video: bool = False,
+        do_rescale: bool = True,
+        do_normalize: bool = True,
     ) -> None:
+        if do_rescale != do_normalize:
+            raise ValueError(
+                "NanoNemotronVLProcessor requires do_rescale and "
+                "do_normalize to have the same value"
+            )
         super().__init__(
             config=config,
             tokenizer=tokenizer,
             max_model_len=max_model_len,
             max_num_tiles=max_num_tiles,
+            do_cpu_normalize=do_normalize,
         )
         # add extra video token for video processing
         self.video_token = video_token
@@ -883,6 +918,7 @@ def _videos_to_pixel_values_lst(
                 downsample_ratio=self.config.downsample_ratio,
                 norm_mean=self.norm_mean,
                 norm_std=self.norm_std,
+                do_cpu_normalize=self.do_cpu_normalize,
                 dtype=dtype,
             )
             for video in videos
diff --git a/vllm/transformers_utils/repo_utils.py b/vllm/transformers_utils/repo_utils.py
index 43ef07676e3a..25b1abe4222d 100644
--- a/vllm/transformers_utils/repo_utils.py
+++ b/vllm/transformers_utils/repo_utils.py
@@ -304,6 +304,30 @@ def get_model_path(model: str | Path, revision: str | None = None):
     )
 
 
+def get_subfolder_path(
+    model: str,
+    subfolder: str,
+    revision: str | None = None,
+    cache_dir: str | None = None,
+) -> str:
+    """Return the local path of `subfolder` in a local or Hub checkpoint,
+    downloading only that subfolder from the Hub."""
+    if os.path.isdir(model):
+        return os.path.join(model, subfolder)
+    # Avoid circular import
+    from vllm.model_executor.model_loader.weight_utils import get_lock
+
+    with get_lock(model, cache_dir):
+        snapshot = hf_api().snapshot_download(
+            repo_id=model,
+            revision=revision,
+            cache_dir=cache_dir,
+            allow_patterns=f"{subfolder}/*",
+            local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE,
+        )
+    return os.path.join(snapshot, subfolder)
+
+
 _WEIGHT_FILE_PATTERNS = [
     "*.safetensors",
     "*.bin",
@@ -315,6 +339,14 @@ def get_model_path(model: str | Path, revision: str | None = None):
     "*.msgpack",
     "*.onnx",
     "*.onnx_data",
+    "*.ot",
+    "*.tflite",
+    "*.mlmodel",
+    "*.mlpackage/*",
+    "*.nemo",
+    "*.npy",
+    "*.npz",
+    "*.pdparams",
     # Large non-weight assets that tokenizers and processors never read
     "*.png",
     "*.jpg",
diff --git a/vllm/triton_utils/dispatcher.py b/vllm/triton_utils/dispatcher.py
index c7e2f9e46986..7d1c2e79aa0e 100644
--- a/vllm/triton_utils/dispatcher.py
+++ b/vllm/triton_utils/dispatcher.py
@@ -198,8 +198,22 @@ def lookup(value: Any) -> Callable[..., Any] | None:
             # kernel as an instance attribute; read the instance dict directly
             # so module-level __getattr__ hooks are never triggered, and so
             # this scan needs no dependency on the warmup machinery.
-            if hasattr(value, "__dict__"):
-                replacement = lookup(vars(value).get("kernel"))
+            #
+            # `hasattr(value, "__dict__")` is not a sufficient guard: proxy
+            # objects such as llvmlite's `binding.ffi.lib` synthesise every
+            # attribute through `__getattr__`, so `__dict__` resolves to a
+            # `_lib_fn_wrapper` and `.get()` would raise AttributeError. Only
+            # treat a real dict as an instance namespace.
+            try:
+                instance_dict = vars(value)
+            except TypeError:
+                # Slots-only objects have no instance dict.
+                instance_dict = None
+            replacement = (
+                lookup(instance_dict.get("kernel"))
+                if isinstance(instance_dict, dict)
+                else None
+            )
             if replacement is None:
                 continue
             value.kernel = replacement
diff --git a/vllm/utils/cache.py b/vllm/utils/cache.py
index 9fca5ec46541..57cbaf97ceb7 100644
--- a/vllm/utils/cache.py
+++ b/vllm/utils/cache.py
@@ -120,6 +120,16 @@ def touch(self, key: _K) -> None:
         if key in self:
             self._LRUCache__order.move_to_end(key)  # type: ignore
 
+    @overload
+    def peek(self, key: _K, /) -> _V | None: ...
+
+    @overload
+    def peek(self, key: _K, /, default: _V | _T) -> _V | _T: ...
+
+    def peek(self, key: _K, /, default: _V | _T | None = None) -> _V | _T | None:
+        """Look up a value without copying the cache or updating order or stats."""
+        return self._Cache__data.get(key, default)  # type: ignore
+
     @overload
     def get(self, key: _K, /) -> _V | None: ...
 
diff --git a/vllm/utils/cpu_resource_utils.py b/vllm/utils/cpu_resource_utils.py
index 79cdc01d215d..b8d37f645476 100644
--- a/vllm/utils/cpu_resource_utils.py
+++ b/vllm/utils/cpu_resource_utils.py
@@ -115,6 +115,16 @@ def get_cgroup_memory_usage() -> int | None:
     return None
 
 
+def get_available_memory() -> int:
+    """Return available host memory, honoring cgroup limits."""
+    available = psutil.virtual_memory().available
+    cgroup_limit = get_cgroup_memory_limit()
+    if cgroup_limit is None:
+        return available
+    cgroup_usage = get_cgroup_memory_usage()
+    return min(available, max(0, cgroup_limit - (cgroup_usage or 0)))
+
+
 def check_cgroup_memory_available(
     required_bytes: int,
     allocation_name: str,
diff --git a/vllm/utils/cuda_process_checkpoint.py b/vllm/utils/cuda_process_checkpoint.py
new file mode 100644
index 000000000000..57ad0559d591
--- /dev/null
+++ b/vllm/utils/cuda_process_checkpoint.py
@@ -0,0 +1,248 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Linux CUDA process checkpoint calls run in a separate interpreter."""
+
+import ctypes as C
+import time
+from collections.abc import Sequence
+from typing import Literal, TextIO, cast
+
+CheckpointOperation = Literal["Lock", "Checkpoint", "Restore", "Unlock"]
+_OPERATIONS: tuple[CheckpointOperation, ...] = (
+    "Lock",
+    "Checkpoint",
+    "Restore",
+    "Unlock",
+)
+
+
+class LockArgs(C.Structure):
+    _fields_ = [
+        ("timeoutMs", C.c_uint),
+        ("reserved0", C.c_uint),
+        ("reserved1", C.c_uint64 * 7),
+    ]
+
+
+class DirectDriver:
+    def __init__(self, emit):
+        self.emit = emit
+        self.lib = C.CDLL("libcuda.so.1")
+        self.lib.cuInit.argtypes = [C.c_uint]
+        rc = self.lib.cuInit(0)
+        if rc:
+            raise RuntimeError(f"cuInit: {rc}")
+        self.lib.cuCheckpointProcessGetState.argtypes = [C.c_int, C.POINTER(C.c_int)]
+        self.lib.cuGetErrorName.argtypes = [C.c_int, C.POINTER(C.c_char_p)]
+        for op in _OPERATIONS:
+            getattr(self.lib, "cuCheckpointProcess" + op).argtypes = [
+                C.c_int,
+                C.c_void_p,
+            ]
+
+    def state(self, pid):
+        state = C.c_int(-1)
+        rc = self.lib.cuCheckpointProcessGetState(pid, C.byref(state))
+        return {"rc": rc, "value": state.value}
+
+    def invoke(self, pid, op: CheckpointOperation):
+        """Return a result without writing to the helper protocol stream."""
+        before = self.state(pid)
+        args = LockArgs(10000) if op == "Lock" else None
+        start = time.perf_counter()
+        rc = getattr(self.lib, "cuCheckpointProcess" + op)(
+            pid, C.byref(args) if args is not None else None
+        )
+        elapsed = time.perf_counter() - start
+        name = C.c_char_p()
+        self.lib.cuGetErrorName(rc, C.byref(name))
+        after = self.state(pid)
+        return dict(
+            kind="control",
+            pid=pid,
+            op=op,
+            rc=rc,
+            error=name.value.decode() if name.value else None,
+            before=before,
+            after=after,
+            time_s=elapsed,
+        )
+
+    def __call__(self, pid, op: CheckpointOperation):
+        row = self.invoke(pid, op)
+        self.emit(**row)
+        if row["rc"] or row["after"] != {
+            "rc": 0,
+            "value": {"Lock": 1, "Checkpoint": 2, "Restore": 1, "Unlock": 0}[op],
+        }:
+            raise RuntimeError(str(row))
+
+
+def _run_parallel_stage(pids, op: CheckpointOperation, invoke):
+    """Run one stage in the helper; never emit partial protocol responses."""
+    from concurrent.futures import ThreadPoolExecutor
+
+    def run(pid):
+        try:
+            return invoke(pid, op)
+        except Exception as exc:
+            return {"pid": pid, "op": op, "error": repr(exc)}
+
+    # Waiting threads live only in the helper. The parent's stage deadline
+    # terminates this process if any driver call blocks.
+    with ThreadPoolExecutor(max_workers=len(pids)) as pool:
+        return list(pool.map(run, pids))
+
+
+class Driver:
+    """Persistent helper isolates potentially blocking driver calls."""
+
+    def __init__(self, emit):
+        import subprocess
+        import sys
+        from pathlib import Path
+
+        self.emit = emit
+        self.proc = subprocess.Popen(
+            [sys.executable, str(Path(__file__)), "--service"],
+            stdin=subprocess.PIPE,
+            stdout=subprocess.PIPE,
+            text=True,
+        )
+        try:
+            self._read(30)
+        except Exception:
+            if self.proc.poll() is None:
+                self.proc.kill()
+            self.proc.communicate(timeout=5)
+            raise
+
+    def _read(self, timeout):
+        import json
+        import selectors
+
+        sel = selectors.DefaultSelector()
+        try:
+            sel.register(cast(TextIO, self.proc.stdout), selectors.EVENT_READ)
+            if not sel.select(timeout):
+                self.proc.kill()
+                raise TimeoutError("CUDA controller deadline exceeded")
+            line = cast(TextIO, self.proc.stdout).readline()
+            if not line:
+                raise RuntimeError("CUDA controller exited unexpectedly")
+            return json.loads(line)
+        finally:
+            sel.close()
+
+    def __call__(self, pid, op: CheckpointOperation):
+        self.run_stage([pid], op)
+
+    def run_stage(self, pids, op: CheckpointOperation):
+        import json
+
+        for pid in pids:
+            self.emit(kind="control_start", pid=pid, op=op)
+        cast(TextIO, self.proc.stdin).write(json.dumps([pids, op]) + "\n")
+        cast(TextIO, self.proc.stdin).flush()
+        try:
+            response = self._read(30)
+        except Exception as exc:
+            self.emit(kind="control_failed", pids=pids, op=op, error=repr(exc))
+            raise
+        if not isinstance(response, dict):
+            raise RuntimeError(f"Invalid CUDA checkpoint stage response: {response}")
+        rows = response.get("results", [])
+        if (
+            response.get("kind") != "stage"
+            or response.get("op") != op
+            or not isinstance(rows, list)
+            or len(rows) != len(pids)
+            or any(
+                not isinstance(row, dict)
+                or row.get("pid") != pid
+                or row.get("op") != op
+                for pid, row in zip(pids, rows)
+            )
+        ):
+            raise RuntimeError(f"Invalid CUDA checkpoint stage response: {response}")
+        for row in rows:
+            self.emit(**row)
+        failed = [
+            row
+            for row in rows
+            if row.get("rc") != 0
+            or row.get("after")
+            != {
+                "rc": 0,
+                "value": {"Lock": 1, "Checkpoint": 2, "Restore": 1, "Unlock": 0}[op],
+            }
+        ]
+        if failed:
+            raise RuntimeError(str(failed))
+
+    def close(self):
+        if self.proc.poll() is None:
+            self.proc.kill()
+        self.proc.communicate(timeout=1)
+
+
+class CudaProcessCheckpoint:
+    """Suspend all local worker processes, or resume all before returning."""
+
+    def __init__(self, emit):
+        self.emit = emit
+        self.driver = Driver(emit)
+        self.pids = []
+        self.failed = False
+
+    def _run(self, operations: Sequence[CheckpointOperation]):
+        if self.failed:
+            raise RuntimeError("CUDA process checkpoint failed; restart required")
+        try:
+            for op in operations:
+                self.driver.run_stage(self.pids, op)
+                self.emit(kind="boundary", stage=op, pids=self.pids)
+        except Exception:
+            self.failed = True
+            raise
+
+    def suspend(self, pids):
+        if self.pids:
+            raise RuntimeError("CUDA processes are already checkpointed")
+        if not pids or len(set(pids)) != len(pids):
+            raise ValueError("Expected distinct worker PIDs")
+        self.pids = list(pids)
+        self._run(("Lock", "Checkpoint"))
+
+    def resume(self):
+        if not self.pids:
+            return
+        self._run(("Restore", "Unlock"))
+        self.pids = []
+
+    def close(self):
+        self.driver.close()
+
+
+if __name__ == "__main__":
+    import json
+    import sys
+
+    def output(**kw):
+        print(json.dumps(kw), flush=True)
+
+    driver = DirectDriver(output)
+    output(kind="controller_ready")
+    for line in sys.stdin:
+        try:
+            pids, op = json.loads(line)
+            if isinstance(pids, int):
+                driver(pids, op)
+            else:
+                output(
+                    kind="stage",
+                    op=op,
+                    results=_run_parallel_stage(pids, op, driver.invoke),
+                )
+        except Exception:
+            break
diff --git a/vllm/utils/extensible_tensor.py b/vllm/utils/extensible_tensor.py
new file mode 100644
index 000000000000..2a11b1032437
--- /dev/null
+++ b/vllm/utils/extensible_tensor.py
@@ -0,0 +1,441 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Growable device byte buffers backed by GPU virtual memory management."""
+
+from __future__ import annotations
+
+import ctypes
+import itertools
+import math
+import sys
+from collections.abc import Iterable, Sequence
+from contextlib import suppress
+from typing import Any
+
+import torch
+
+from vllm.utils.vmm_driver import get_vmm_driver
+
+
+def _round_up(value: int, multiple: int) -> int:
+    return ((value + multiple - 1) // multiple) * multiple
+
+
+def granule_block_alignment(block_strides: Iterable[int], granule: int) -> int:
+    """Block count multiple at which every stride's prefix is a granule multiple."""
+    alignment = 1
+    for stride in set(block_strides):
+        alignment = math.lcm(alignment, granule // math.gcd(granule, stride))
+    return alignment
+
+
+def granule_aligned_blocks(
+    num_blocks: int, block_strides: Iterable[int], granule: int
+) -> int:
+    """Largest count up to num_blocks at which, for every block stride, that
+    many blocks span a whole number of granules.
+
+    Segments hold num_blocks blocks of one stride each, so at such a count
+    every segment starts and ends on a granule boundary and one physical chunk
+    can back it exactly, as RDMA registration of a segment requires.
+    """
+    alignment = granule_block_alignment(block_strides, granule)
+    return num_blocks // alignment * alignment
+
+
+# ROCm reports a 4 KiB granularity (CUDA 2 MiB); commit in larger units to
+# keep the granule bookkeeping small.
+_MIN_GRANULARITY = 2 << 20
+
+
+def granule_size(device_index: int) -> int:
+    """Commit granule for device_index: driver granularity, floored at 2 MiB."""
+    driver_granularity = get_vmm_driver().granularity(device_index)
+    return _round_up(max(driver_granularity, _MIN_GRANULARITY), driver_granularity)
+
+
+class _VirtualBuffer:
+    """One device VA reservation and the physical chunks mapped into it.
+
+    Memory is committed in granules via `ensure_committed_range`; already-mapped
+    granules are skipped, so ranges may abut or overlap.
+    """
+
+    def __init__(self, max_bytes: int, device_index: int) -> None:
+        self._driver = get_vmm_driver()
+        self._driver.ensure_context(device_index)
+        self.device_index = device_index
+
+        self.granularity: int = granule_size(device_index)
+        self.reserved_size: int = _round_up(max(max_bytes, 1), self.granularity)
+        self.base_ptr: int = self._driver.reserve(self.reserved_size)
+
+        # Granule indices (VA offset // granularity) that have physical
+        # memory mapped.
+        self._mapped_granules: set[int] = set()
+        # Each entry is (handle, va_offset, size) for one mapped physical chunk.
+        self._handles: list[tuple[int, int, int]] = []
+        self._freed: bool = False
+
+    @property
+    def committed_bytes(self) -> int:
+        """Total physically mapped bytes (a multiple of the granularity)."""
+        return len(self._mapped_granules) * self.granularity
+
+    def ensure_committed_range(self, start: int, end: int) -> None:
+        """Map physical pages so that the byte range `[start, end)` is backed."""
+        if not 0 <= start <= end:
+            raise ValueError(f"Invalid range [{start}, {end}).")
+        if end > self.reserved_size:
+            raise ValueError(
+                f"Requested range end {end} exceeds reserved capacity "
+                f"{self.reserved_size}."
+            )
+        if start == end:
+            return
+        first = start // self.granularity
+        last = (end + self.granularity - 1) // self.granularity  # exclusive
+        run_start: int | None = None
+        for g in range(first, last + 1):
+            unmapped = g < last and g not in self._mapped_granules
+            if unmapped and run_start is None:
+                run_start = g
+            elif not unmapped and run_start is not None:
+                self._map_chunk_at(
+                    run_start * self.granularity, (g - run_start) * self.granularity
+                )
+                # Grant before recording the run as mapped: a failure in a
+                # later run must not leave these granules mapped but
+                # inaccessible, since later commits would skip them.
+                self._grant_access(*self._run_bounds(run_start, g))
+                self._mapped_granules.update(range(run_start, g))
+                run_start = None
+
+    def _run_bounds(self, first: int, last: int) -> tuple[int, int]:
+        """Bounds of the maximal contiguous mapped run covering `[first, last)`."""
+        while first - 1 in self._mapped_granules:
+            first -= 1
+        while last in self._mapped_granules:
+            last += 1
+        return first, last
+
+    def _grant_access(self, first: int, last: int) -> None:
+        """Grant device access over the granule run `[first, last)`.
+
+        Per run rather than per chunk: ROCm rejects a set-access range that starts
+        inside an already-mapped region. Re-granting is a cheap no-op.
+        """
+        self._driver.set_access(
+            self.base_ptr + first * self.granularity,
+            (last - first) * self.granularity,
+            self.device_index,
+        )
+
+    def _map_chunk_at(self, offset: int, size: int) -> None:
+        """Create one physical chunk of `size` bytes and map it at `offset`."""
+        driver = self._driver
+        driver.ensure_context(self.device_index)
+        try:
+            handle = driver.create(size, self.device_index)
+        except RuntimeError:
+            # The VMM allocator cannot reuse memory idling in torch's
+            # caching allocator; return it to the driver and retry once.
+            torch.accelerator.empty_cache()
+            handle = driver.create(size, self.device_index)
+
+        addr = self.base_ptr + offset
+        try:
+            driver.map(addr, size, handle)
+        except RuntimeError:
+            driver.release(handle)
+            raise
+        # Access is granted by the caller, per contiguous run.
+        self._handles.append((handle, offset, size))
+
+    def release_physical(self) -> None:
+        """Unmap and release all physical memory, keeping the VA reservation."""
+        driver = self._driver
+        driver.ensure_context(self.device_index)
+        if self._handles:
+            torch.accelerator.synchronize(self.device_index)
+        for handle, offset, size in self._handles:
+            driver.unmap(self.base_ptr + offset, size)
+            driver.release(handle)
+        self._handles = []
+        self._mapped_granules = set()
+
+    def free(self) -> None:
+        if self._freed:
+            return
+        self._freed = True
+        self.release_physical()
+        if self.base_ptr:
+            self._driver.free_reserved(self.base_ptr, self.reserved_size)
+        self.base_ptr = 0
+
+    def __del__(
+        self, _suppress: Any = suppress, _is_finalizing: Any = sys.is_finalizing
+    ) -> None:
+        # Bound as defaults: module globals may be gone at interpreter exit.
+        # The process frees its mappings on exit, and the HIP runtime may
+        # already be torn down, where a driver call aborts the process.
+        if _suppress is None or _is_finalizing():
+            return
+        with _suppress(Exception):
+            self.free()
+
+
+_K_DL_UINT = 1
+_UINT8_BITS = 8
+
+
+class _DLDevice(ctypes.Structure):
+    _fields_ = [("device_type", ctypes.c_int), ("device_id", ctypes.c_int)]
+
+
+class _DLDataType(ctypes.Structure):
+    _fields_ = [
+        ("code", ctypes.c_uint8),
+        ("bits", ctypes.c_uint8),
+        ("lanes", ctypes.c_uint16),
+    ]
+
+
+class _DLTensor(ctypes.Structure):
+    _fields_ = [
+        ("data", ctypes.c_void_p),
+        ("device", _DLDevice),
+        ("ndim", ctypes.c_int),
+        ("dtype", _DLDataType),
+        ("shape", ctypes.POINTER(ctypes.c_int64)),
+        ("strides", ctypes.POINTER(ctypes.c_int64)),
+        ("byte_offset", ctypes.c_uint64),
+    ]
+
+
+class _DLManagedTensor(ctypes.Structure):
+    pass
+
+
+_DLDeleter = ctypes.CFUNCTYPE(None, ctypes.POINTER(_DLManagedTensor))
+_DLManagedTensor._fields_ = [
+    ("dl_tensor", _DLTensor),
+    ("manager_ctx", ctypes.c_void_p),
+    ("deleter", _DLDeleter),
+]
+
+_PyCapsule_New = ctypes.pythonapi.PyCapsule_New
+_PyCapsule_New.restype = ctypes.py_object
+_PyCapsule_New.argtypes = [ctypes.c_void_p, ctypes.c_char_p, ctypes.c_void_p]
+_malloc = ctypes.CDLL(None).malloc
+_malloc.restype = ctypes.c_void_p
+_malloc.argtypes = [ctypes.c_size_t]
+
+
+def uint8_tensor_from_ptr(ptr: int, num_bytes: int, device_index: int) -> torch.Tensor:
+    """Wrap device memory at `ptr` as a uint8 tensor; the caller keeps it mapped.
+
+    The DLManagedTensor is malloc'd and never freed: torch reads it when the
+    view is destroyed, which can be after this module's globals are cleared at
+    interpreter exit. Its deleter is NULL, so nothing calls back into Python.
+    """
+    managed_size = ctypes.sizeof(_DLManagedTensor)
+    addr = _malloc(managed_size + ctypes.sizeof(ctypes.c_int64))
+    if not addr:
+        raise MemoryError("Failed to allocate a DLManagedTensor.")
+    ctypes.memset(addr, 0, managed_size + ctypes.sizeof(ctypes.c_int64))
+    shape_addr = addr + managed_size
+    ctypes.c_int64.from_address(shape_addr).value = num_bytes
+
+    managed = _DLManagedTensor.from_address(addr)
+    managed.dl_tensor.data = ctypes.c_void_p(ptr)
+    device_type = get_vmm_driver().dlpack_device_type
+    managed.dl_tensor.device = _DLDevice(device_type, device_index)
+    managed.dl_tensor.ndim = 1
+    managed.dl_tensor.dtype = _DLDataType(_K_DL_UINT, _UINT8_BITS, 1)
+    managed.dl_tensor.shape = ctypes.cast(shape_addr, ctypes.POINTER(ctypes.c_int64))
+
+    capsule = _PyCapsule_New(addr, b"dltensor", None)
+    return torch.from_dlpack(capsule)
+
+
+class ExtensibleTensor:
+    """A device byte buffer that grows without moving its base pointer.
+
+    The reservation is divided into contiguous segments whose committed bytes
+    form a prefix of each; num_segments equal segments by default, or the
+    explicit segment_capacities (in bytes, summing to max_num_bytes).
+    """
+
+    def __init__(
+        self,
+        max_num_bytes: int,
+        device: torch.device | str | int | None = None,
+        num_segments: int = 1,
+        segment_capacities: Sequence[int] | None = None,
+    ) -> None:
+        if max_num_bytes < 0:
+            raise ValueError("max_num_bytes must be non-negative.")
+        if segment_capacities is None:
+            if num_segments < 1:
+                raise ValueError(f"num_segments must be positive, got {num_segments}.")
+            if max_num_bytes % num_segments != 0:
+                raise ValueError(
+                    f"max_num_bytes ({max_num_bytes}) must be divisible by "
+                    f"num_segments ({num_segments})."
+                )
+            segment_capacities = [max_num_bytes // num_segments] * num_segments
+        elif not segment_capacities or any(c <= 0 for c in segment_capacities):
+            raise ValueError("segment_capacities must be non-empty and positive.")
+        elif sum(segment_capacities) != max_num_bytes:
+            raise ValueError(
+                f"segment_capacities sum to {sum(segment_capacities)}, not "
+                f"max_num_bytes ({max_num_bytes})."
+            )
+
+        if device is None:
+            device = torch.accelerator.current_device_index()
+        dev = device if isinstance(device, torch.device) else torch.device(device)
+        if dev.type != "cuda":
+            raise ValueError(f"ExtensibleTensor requires a cuda device, got {dev}.")
+        self._device_index: int = (
+            dev.index
+            if dev.index is not None
+            else torch.accelerator.current_device_index()
+        )
+
+        self._max_num_bytes: int = max_num_bytes
+        self._segment_capacities: list[int] = list(segment_capacities)
+        self._segment_offsets: list[int] = list(
+            itertools.accumulate([0, *self._segment_capacities[:-1]])
+        )
+        self._committed: list[int] = [0] * len(self._segment_capacities)
+        self._buffer: _VirtualBuffer = _VirtualBuffer(max_num_bytes, self._device_index)
+
+    def full_view(self) -> torch.Tensor:
+        """Uint8 view spanning the whole reservation."""
+        return uint8_tensor_from_ptr(
+            self._buffer.base_ptr, self._max_num_bytes, self._device_index
+        )
+
+    def segment_view(self, index: int) -> torch.Tensor:
+        """Uint8 tensor over one segment's committed prefix, sized exactly to it."""
+        if not 0 <= index < self.num_segments:
+            raise IndexError(f"Segment {index} out of range ({self.num_segments}).")
+        return uint8_tensor_from_ptr(
+            self._buffer.base_ptr + self._segment_offsets[index],
+            self._committed[index],
+            self._device_index,
+        )
+
+    def resize_per_segment_(
+        self, bytes_per_segment: int, zero_new: bool = False
+    ) -> None:
+        """Grow every segment's committed prefix to the same size."""
+        self.resize_segments_([bytes_per_segment] * self.num_segments, zero_new)
+
+    def resize_segments_(self, sizes: Sequence[int], zero_new: bool = False) -> None:
+        """Grow each segment's committed prefix; optionally zero the new bytes."""
+        if len(sizes) != self.num_segments:
+            raise ValueError(f"Expected {self.num_segments} sizes, got {len(sizes)}.")
+        for old, new, capacity in zip(self._committed, sizes, self._segment_capacities):
+            if new < old:
+                raise ValueError(
+                    f"ExtensibleTensor is grow-only: cannot resize from {old} "
+                    f"to {new} bytes per segment."
+                )
+            if new > capacity:
+                raise ValueError(
+                    f"Requested {new} bytes per segment exceeds the "
+                    f"segment capacity {capacity}."
+                )
+        grown = [
+            (start + old, start + new)
+            for start, old, new in zip(self._segment_offsets, self._committed, sizes)
+            if new > old
+        ]
+        for start, end in grown:
+            self._buffer.ensure_committed_range(start, end)
+        self._committed = list(sizes)
+        if zero_new and grown:
+            full = self.full_view()
+            for start, end in grown:
+                full[start:end].zero_()
+
+    def release_physical(self) -> None:
+        """Drop all physical pages but keep the reservation and its views."""
+        self._buffer.release_physical()
+        self._committed = [0] * self.num_segments
+
+    @property
+    def num_bytes(self) -> int:
+        """Current committed size in bytes, summed over all segments."""
+        return sum(self._committed)
+
+    @property
+    def bytes_per_segment(self) -> int:
+        """Current committed prefix size of each segment in bytes (equal segments)."""
+        return self._uniform(self._committed)
+
+    @property
+    def num_segments(self) -> int:
+        return len(self._segment_capacities)
+
+    @property
+    def segment_capacity_bytes(self) -> int:
+        """Maximum size of each segment (equal segments)."""
+        return self._uniform(self._segment_capacities)
+
+    @property
+    def segment_capacities(self) -> list[int]:
+        return list(self._segment_capacities)
+
+    @property
+    def segment_offsets(self) -> list[int]:
+        """Byte offset of each segment within the reservation."""
+        return list(self._segment_offsets)
+
+    @staticmethod
+    def _uniform(values: list[int]) -> int:
+        if any(v != values[0] for v in values):
+            raise ValueError(f"Segments differ: {values}.")
+        return values[0]
+
+    @property
+    def capacity_bytes(self) -> int:
+        return self._buffer.reserved_size
+
+    @property
+    def physical_bytes(self) -> int:
+        """Physically mapped bytes, including granule rounding."""
+        return self._buffer.committed_bytes
+
+    @property
+    def num_physical_chunks(self) -> int:
+        """Number of driver allocations currently mapped into the reservation."""
+        return len(self._buffer._handles)
+
+    def segments_backed_by_one_chunk(self, sizes: Sequence[int]) -> list[bool]:
+        """Whether each segment's first `sizes[i]` bytes lie within a single
+        driver allocation, as RDMA registration of the segment requires."""
+        chunks = [(offset, offset + size) for _, offset, size in self._buffer._handles]
+        return [
+            any(start <= offset and offset + size <= end for start, end in chunks)
+            for offset, size in zip(self._segment_offsets, sizes)
+        ]
+
+    @property
+    def granularity(self) -> int:
+        return self._buffer.granularity
+
+    @property
+    def base_ptr(self) -> int:
+        return self._buffer.base_ptr
+
+    @property
+    def device(self) -> torch.device:
+        return torch.device("cuda", self._device_index)
+
+    def free(self) -> None:
+        self._buffer.free()
+        self._committed = [0] * self.num_segments
diff --git a/vllm/utils/flashinfer.py b/vllm/utils/flashinfer.py
index 725e2f8039a6..1afc6d96dd8b 100644
--- a/vllm/utils/flashinfer.py
+++ b/vllm/utils/flashinfer.py
@@ -29,6 +29,13 @@
 logger = init_logger(__name__)
 
 
+def is_flashinfer_autotuning() -> bool:
+    """Whether a FlashInfer autotune(True) context is active."""
+    from flashinfer.autotuner import AutoTuner
+
+    return AutoTuner.get().is_tuning_mode
+
+
 def flashinfer_bf16_mm(
     a: torch.Tensor,
     b: torch.Tensor,
@@ -1201,7 +1208,8 @@ def flashinfer_scaled_fp8_mm(
         scale_a,
         scale_b,
         out_dtype,
-        "auto",
+        # Autotuned "auto" can switch backends between M buckets.
+        "cutlass" if envs.VLLM_BATCH_INVARIANT else "auto",
     ).view(a.shape[0], b.shape[1])
 
     if bias is not None:
@@ -1310,6 +1318,7 @@ def is_flashinfer_cudnn_fp8_prefill_attn_supported() -> bool:
 __all__ = [
     "has_flashinfer",
     "flashinfer_bf16_mm",
+    "is_flashinfer_autotuning",
     "has_flashinfer_bf16_gemm",
     "has_flashinfer_concat_mla_kv_quant_fp8",
     "is_flashinfer_bf16_gemm_supported",
diff --git a/vllm/utils/gc_utils.py b/vllm/utils/gc_utils.py
index be6ab4a23fa1..5a1fb73f1f70 100644
--- a/vllm/utils/gc_utils.py
+++ b/vllm/utils/gc_utils.py
@@ -73,11 +73,11 @@ def handle(self, phase: str, info: dict[str, int]) -> None:
         elif phase == "stop":
             # After GC finished, Record GC elapsed time and
             # optionally top collected objects
-            elpased_ms = (time.monotonic_ns() - self.start_time_ns) / 1e6
+            elapsed_ms = (time.monotonic_ns() - self.start_time_ns) / 1e6
             logger.info(
                 "GC took %.3fms to complete. "
                 "Collected %s objects (out of %d) in GC generation %d.%s",
-                elpased_ms,
+                elapsed_ms,
                 str(info.get("collected", "?")),
                 self.num_objects,
                 generation,
diff --git a/vllm/utils/import_utils.py b/vllm/utils/import_utils.py
index 277f987d4800..80374ec19356 100644
--- a/vllm/utils/import_utils.py
+++ b/vllm/utils/import_utils.py
@@ -639,6 +639,32 @@ def has_mori() -> bool:
     return _has_module("mori")
 
 
+# First release with single-node CCO windows (ROCm/mori#704).
+MORI_EPV2_MIN_VERSION = "1.2.4"
+
+
+def mori_epv2_unavailable_reason() -> str | None:
+    """Why the installed MoRI cannot run EPv2, or None if it can.
+
+    Requires the cco and dispatch_combine_v2 modules and MoRI >= 1.2.4.
+    """
+    if not _has_module("mori"):
+        return "MoRI is not installed"
+    if not (_has_module("mori.cco") and _has_module("mori.ops.dispatch_combine_v2")):
+        return "the installed MoRI has no EPv2 modules"
+    import mori  # type: ignore[import-not-found]
+    from packaging.version import InvalidVersion, Version
+
+    version = getattr(mori, "__version__", "unknown")
+    try:
+        too_old = Version(version) < Version(MORI_EPV2_MIN_VERSION)
+    except InvalidVersion:
+        too_old = True
+    if too_old:
+        return f"MoRI {version} is older than {MORI_EPV2_MIN_VERSION}"
+    return None
+
+
 def has_fbgemm_gpu() -> bool:
     """Whether the optional `fbgemm_gpu` package is available."""
     return _has_module("fbgemm_gpu")
diff --git a/vllm/utils/jit_monitor.py b/vllm/utils/jit_monitor.py
index 387189c793fb..5af373a4b4df 100644
--- a/vllm/utils/jit_monitor.py
+++ b/vllm/utils/jit_monitor.py
@@ -18,6 +18,7 @@
 - Triton ``@triton.jit`` first-time compilations
   (via ``knobs.runtime.jit_post_compile_hook``)
 - TileLang ``@tilelang.jit`` first-time compilations
+- FlyDSL ``MlirCompiler.compile`` calls
 """
 
 import contextlib
@@ -39,6 +40,7 @@
 _mode: JitMonitorMode = "warn"
 _verbose: bool = False
 _cutedsl_hook_installed: bool = False
+_flydsl_hook_installed: bool = False
 _tilelang_hook_installed: bool = False
 _tilelang_jitimpl_compile_depth: int = 0
 
@@ -74,6 +76,7 @@ def activate(*, mode: JitMonitorMode = "warn", verbose: bool = False) -> None:
     _setup_triton_autotuning_print()
     _setup_triton_jit_hook()
     _setup_cutedsl_jit_hook()
+    _setup_flydsl_jit_hook()
 
     # Refer to #51159. tilelang ships broken symbols
     # on rocm.
@@ -324,6 +327,71 @@ def _setup_cutedsl_jit_hook() -> None:
     _cutedsl_hook_installed = True
 
 
+# ------------------------------------------------------------------
+# FlyDSL JIT compilation hook
+# ------------------------------------------------------------------
+
+
+def _log_flydsl_jit_compile(fn_name: str, detail: str | None = None) -> None:
+    _handle_jit_event(
+        backend="FlyDSL",
+        event="JIT compilation",
+        fn_name=fn_name,
+        detail=detail,
+    )
+
+
+def _flydsl_kernel_symbols(module: object, infer_kernel_names) -> str | None:
+    """Name the ``gpu.func`` kernels in a module about to be compiled.
+
+    ``func_name`` is the Python launcher name, which many kernels share
+    verbatim (``launch``), so the ``gpu.func`` symbols are what identify the
+    kernel and the width that needs warmup.
+    """
+    if module is None:
+        return None
+    names = dict.fromkeys(infer_kernel_names(str(module)))
+    return "kernels=" + ",".join(names) if names else None
+
+
+def _setup_flydsl_jit_hook() -> None:
+    """Wrap ``MlirCompiler.compile`` to warn on compilation."""
+    global _flydsl_hook_installed
+    if _flydsl_hook_installed:
+        return
+
+    try:
+        from flydsl.compiler import jit_function
+    except ImportError:
+        logger.debug("FlyDSL is not available; skipping FlyDSL JIT monitor.")
+        return
+
+    try:
+        compiler_cls = jit_function.MlirCompiler
+        # MlirCompiler.compile is called only on a cache miss, so for now
+        # we don't monitor disk cache hits
+        original_compile = compiler_cls.compile
+        # TODO: use a public API
+        infer_kernel_names = jit_function._infer_kernel_names_from_asm
+    except AttributeError as exc:
+        logger.warning(
+            "FlyDSL JIT monitoring is disabled; its JIT internals moved (%s).", exc
+        )
+        return
+
+    @functools.wraps(original_compile)
+    def _compile_with_monitor(cls, *args, **kwargs):
+        module = args[0] if args else kwargs.get("module")
+        detail = None
+        with suppress(Exception):
+            detail = _flydsl_kernel_symbols(module, infer_kernel_names)
+        _log_flydsl_jit_compile(str(kwargs.get("func_name") or "<unknown>"), detail)
+        return original_compile(*args, **kwargs)
+
+    compiler_cls.compile = classmethod(_compile_with_monitor)
+    _flydsl_hook_installed = True
+
+
 # ------------------------------------------------------------------
 # TileLang JIT compilation hook
 # ------------------------------------------------------------------
diff --git a/vllm/utils/numa_utils.py b/vllm/utils/numa_utils.py
index 755d46e711fb..fefa30ea9ca7 100644
--- a/vllm/utils/numa_utils.py
+++ b/vllm/utils/numa_utils.py
@@ -26,6 +26,7 @@
 logger = logging.getLogger(__name__)
 _NUMACTL_ARGS_ENV = "_VLLM_INTERNAL_NUMACTL_ARGS"
 _NUMACTL_PYTHON_EXECUTABLE_ENV = "_VLLM_INTERNAL_NUMACTL_PYTHON_EXECUTABLE"
+_MPOL_DEFAULT = 0
 
 
 @cache
@@ -41,55 +42,84 @@ def get_libnuma():
     return libnuma
 
 
-def _can_set_mempolicy() -> bool:
-    """Check whether the current process can use NUMA memory policy syscalls."""
+def _get_mempolicy() -> int | None:
+    """Return the current memory-policy mode, or None if it cannot be read."""
     try:
         libnuma = get_libnuma()
         if libnuma is None or libnuma.numa_available() < 0:
-            return False
+            return None
         mode = ctypes.c_int()
         ret = libnuma.get_mempolicy(
             ctypes.byref(mode), None, ctypes.c_ulong(0), None, ctypes.c_ulong(0)
         )
-        return ret == 0
+        return mode.value if ret == 0 else None
     except Exception:
-        return False
+        return None
 
 
-def _is_auto_numa_available() -> bool:
-    """Check whether automatic GPU-to-NUMA detection should be attempted."""
+def _is_auto_numa_available(
+    verbose: bool = True, *, require_unbound_process: bool = False
+) -> bool:
+    """Check whether automatic GPU-to-NUMA detection should be attempted.
+
+    Args:
+        verbose: Log why detection is skipped at warning level; otherwise at
+            debug level, for the default mode where skipping is expected.
+        require_unbound_process: Require known unrestricted CPU affinity and
+            the default memory policy for default automatic binding. Explicit
+            topology discovery keeps its existing best-effort checks.
+
+    """
     from vllm.platforms import current_platform
 
+    log = logger.warning if verbose else logger.debug
+
     if not current_platform.is_cuda_alike():
+        log("Automatic NUMA binding requires a CUDA or ROCm platform.")
         return False
 
     if not os.path.isdir("/sys/devices/system/node/node1"):
+        log("Automatic NUMA binding requires multiple NUMA nodes.")
         return False
 
     try:
         process = psutil.Process(os.getpid())
         cpu_affinity = process.cpu_affinity()
         cpu_count = psutil.cpu_count()
-        if cpu_count is not None and cpu_affinity != list(range(cpu_count)):
-            logger.warning(
+        if cpu_count is not None and set(cpu_affinity) != set(range(cpu_count)):
+            log(
                 "CPU affinity is already constrained for this process. "
                 "Skipping automatic NUMA binding; pass --numa-bind-nodes "
                 "explicitly to override."
             )
             return False
+        if cpu_count is None and require_unbound_process:
+            log("Could not determine CPU availability for automatic NUMA binding.")
+            return False
     except (AttributeError, NotImplementedError, psutil.Error):
-        pass
+        if require_unbound_process:
+            log("Could not check CPU affinity for automatic NUMA binding.")
+            return False
 
-    if not _can_set_mempolicy():
-        logger.warning(
-            "User lacks permission to set NUMA memory policy. "
+    mempolicy = _get_mempolicy()
+    if mempolicy is None:
+        log(
+            "Could not read NUMA memory policy. "
             "Automatic NUMA detection may not work; if you are using Docker, "
             "try adding --cap-add SYS_NICE."
         )
         return False
+    if require_unbound_process and mempolicy != _MPOL_DEFAULT:
+        log(
+            "NUMA memory policy is already set for this process (mode=%d). "
+            "Skipping automatic NUMA binding; pass --numa-bind explicitly "
+            "to override.",
+            mempolicy,
+        )
+        return False
 
     if not hasattr(current_platform, "get_all_device_numa_nodes"):
-        logger.warning(
+        log(
             "Platform %s does not support automatic NUMA detection",
             type(current_platform).__name__,
         )
@@ -98,7 +128,6 @@ def _is_auto_numa_available() -> bool:
     return True
 
 
-@cache
 def get_auto_numa_nodes() -> list[int] | None:
     """Auto-detect NUMA nodes for all visible GPUs."""
     from vllm.platforms import current_platform
@@ -112,6 +141,76 @@ def get_auto_numa_nodes() -> list[int] | None:
     return numa_nodes
 
 
+def _can_auto_bind(
+    parallel_config, local_rank: int, dp_local_rank: int | None = None
+) -> bool:
+    """Validate and store this launcher's automatic GPU-to-NUMA topology.
+
+    Unlike ``--numa-bind``, this never fails startup or forces the ``spawn``
+    start method: any unmet requirement means starting unbound.
+    """
+    from shutil import which
+
+    if envs.VLLM_WORKER_MULTIPROC_METHOD != "spawn":
+        logger.debug("Automatic NUMA binding requires the spawn start method.")
+        return False
+    if (
+        parallel_config.distributed_executor_backend not in ("mp", "uni")
+        or parallel_config.data_parallel_backend != "mp"
+    ):
+        logger.debug("Automatic NUMA binding requires local mp/uni launchers.")
+        return False
+    # Topology is indexed in visible-device order, which --device-ids remaps.
+    if parallel_config.assigned_physical_gpu_ids is not None:
+        logger.debug("Automatic NUMA binding does not support --device-ids remapping.")
+        return False
+    if which("numactl") is None:
+        logger.debug("Automatic NUMA binding requires numactl on PATH.")
+        return False
+    try:
+        from vllm.platforms import current_platform
+        from vllm.utils.cpu_resource_utils import parse_id_list
+
+        if not _is_auto_numa_available(verbose=False, require_unbound_process=True):
+            return False
+        numa_nodes = get_auto_numa_nodes()
+        if not numa_nodes:
+            logger.debug("No GPU-to-NUMA topology available for automatic binding.")
+            return False
+        if len(numa_nodes) != current_platform.device_count() or any(
+            not isinstance(node, int) or node < 0 for node in numa_nodes
+        ):
+            logger.debug("Incomplete or invalid GPU-to-NUMA topology: %s", numa_nodes)
+            return False
+
+        last_local_rank = max(local_rank, parallel_config.local_world_size - 1)
+        last_dp_local_rank = (parallel_config.data_parallel_rank_local or 0) + max(
+            parallel_config.data_parallel_size_local - 1, 0
+        )
+        last_gpu_index = max(
+            parallel_config.get_local_gpu_index(last_local_rank, dp_local_rank),
+            parallel_config.get_local_gpu_index(last_local_rank, last_dp_local_rank),
+        )
+        if last_gpu_index >= len(numa_nodes):
+            logger.debug(
+                "GPU-to-NUMA topology does not cover all local workers: "
+                "GPU index %d, mapping size %d.",
+                last_gpu_index,
+                len(numa_nodes),
+            )
+            return False
+        for node in set(numa_nodes):
+            cpulist_path = Path(f"/sys/devices/system/node/node{node}/cpulist")
+            if not parse_id_list(cpulist_path.read_text().strip()):
+                logger.debug("NUMA node %d has no CPUs for automatic binding.", node)
+                return False
+        parallel_config.numa_bind_nodes = list(numa_nodes)
+        return True
+    except Exception:
+        logger.debug("Automatic GPU-to-NUMA topology validation failed", exc_info=True)
+        return False
+
+
 # PCT (Priority Core Turbo) auto-detection workaround for Granite Rapids
 # Xeon SKUs.
 #
@@ -326,24 +425,13 @@ def _get_numactl_worker_args(
     gpu_index = parallel_config.get_local_gpu_index(local_rank, dp_local_rank)
     numa_node = _get_numa_node(parallel_config, gpu_index)
     cpu_binding = _get_cpu_binding(parallel_config, gpu_index, [numa_node])
+    memory_policy = "preferred" if parallel_config._numa_bind_auto else "membind"
+    memory_arg = f"--{memory_policy}={numa_node}"
 
     if cpu_binding is not None:
-        logger.info(
-            "Binding worker subprocess (local_rank=%s, gpu_index=%s) to CPUs %s and NUMA node %s",  # noqa: E501
-            local_rank,
-            gpu_index,
-            cpu_binding,
-            numa_node,
-        )
-        return f"--physcpubind={cpu_binding} --membind={numa_node}"
+        return f"--physcpubind={cpu_binding} {memory_arg}"
 
-    logger.info(
-        "Binding worker subprocess (local_rank=%s, gpu_index=%s) to NUMA node %s",
-        local_rank,
-        gpu_index,
-        numa_node,
-    )
-    return f"--cpunodebind={numa_node} --membind={numa_node}"
+    return f"--cpunodebind={numa_node} {memory_arg}"
 
 
 def _get_enginecore_numa_nodes(
@@ -387,7 +475,11 @@ def _get_numactl_enginecore_args(
     is also a safe superset by construction).
     """
     shard_nodes = _get_enginecore_numa_nodes(parallel_config, dp_local_rank)
-    membind_arg = ",".join(str(n) for n in shard_nodes)
+    node_binding = ",".join(str(n) for n in shard_nodes)
+    memory_policy = "membind"
+    if parallel_config._numa_bind_auto:
+        memory_policy = "preferred" if len(shard_nodes) == 1 else "preferred-many"
+    memory_arg = f"--{memory_policy}={node_binding}"
 
     pct_cpus = (
         None
@@ -397,21 +489,9 @@ def _get_numactl_enginecore_args(
 
     if pct_cpus is not None:
         cpu_binding = ",".join(str(c) for c in pct_cpus)
-        logger.info(
-            "Binding EngineCore subprocess (local_rank=%s) to CPUs %s "
-            "and NUMA nodes %s",
-            local_rank,
-            cpu_binding,
-            membind_arg,
-        )
-        return f"--physcpubind={cpu_binding} --membind={membind_arg}"
+        return f"--physcpubind={cpu_binding} {memory_arg}"
 
-    logger.info(
-        "Binding EngineCore subprocess (local_rank=%s) to NUMA nodes %s",
-        local_rank,
-        membind_arg,
-    )
-    return f"--cpunodebind={membind_arg} --membind={membind_arg}"
+    return f"--cpunodebind={node_binding} {memory_arg}"
 
 
 def _log_numactl_show(label: str) -> bool:
@@ -455,20 +535,22 @@ def _probe_numactl_args(numactl_args: str) -> bool:
 
 
 def _resolve_numactl_args(numactl_args: str) -> str:
-    """Drop ``--membind`` if the container rejects it, keeping CPU binding."""
+    """Drop a rejected memory policy, keeping CPU binding when possible."""
     cpu_only = " ".join(
-        t for t in numactl_args.split() if not t.startswith("--membind=")
+        t
+        for t in numactl_args.split()
+        if not t.startswith(("--membind=", "--preferred=", "--preferred-many="))
     )
     for candidate in (numactl_args, cpu_only, ""):
         if _probe_numactl_args(candidate):
             if candidate != numactl_args:
                 logger.warning(
-                    "numactl args %r rejected; falling back to %r. Add "
-                    "--cap-add SYS_NICE for full NUMA binding.",
+                    "numactl args %r rejected; falling back to %r.",
                     numactl_args,
                     candidate or "no binding",
                 )
             return candidate
+    logger.warning("numactl probes failed; starting without NUMA binding.")
     return ""
 
 
@@ -481,6 +563,19 @@ def configure_subprocess(
 ):
     """Temporarily replace the multiprocessing executable with a numactl wrapper."""
     parallel_config = vllm_config.parallel_config
+    if parallel_config.numa_bind is None:
+        # Resolve once in the launching process; spawned children inherit the
+        # result through the pickled config instead of re-checking (a bound
+        # EngineCore would fail the unconstrained-affinity check).
+        parallel_config.numa_bind = _can_auto_bind(
+            parallel_config, local_rank, dp_local_rank
+        )
+        parallel_config._numa_bind_auto = parallel_config.numa_bind
+        if not parallel_config.numa_bind:
+            logger.info(
+                "Automatic NUMA binding unavailable on this host; starting "
+                "unbound. Enable debug logging for the skip reason."
+            )
     if not parallel_config.numa_bind:
         yield
         return
@@ -502,8 +597,20 @@ def configure_subprocess(
     numactl_args = _resolve_numactl_args(numactl_args)
     if not numactl_args:
         # No NUMA binding possible here; launch without the wrapper.
+        if parallel_config._numa_bind_auto:
+            parallel_config.numa_bind = False
+            parallel_config.numa_bind_nodes = None
+            parallel_config._numa_bind_auto = False
         yield
         return
+    if envs.VLLM_WORKER_MULTIPROC_METHOD == "spawn":
+        logger.info(
+            "Launching %s subprocess (local_rank=%s, dp_local_rank=%s) with numactl %s",
+            process_kind,
+            local_rank,
+            dp_local_rank,
+            numactl_args,
+        )
     python_executable = os.fsdecode(multiprocessing.spawn.get_executable())
     with (
         _set_numa_wrapper_env(numactl_args, python_executable),
diff --git a/vllm/utils/system_utils.py b/vllm/utils/system_utils.py
index 5ae13232b069..cbea3c862d25 100644
--- a/vllm/utils/system_utils.py
+++ b/vllm/utils/system_utils.py
@@ -137,9 +137,12 @@ def _maybe_force_spawn():
         os.environ["RAY_ADDRESS"] = ray.get_runtime_context().gcs_address
         reasons.append("In a Ray actor and can only be spawned")
 
-    # Force spawn if NUMA binding is enabled via --numa-bind.
-    # NUMA binding uses executable hijacking which requires spawn
-    if "--numa-bind" in sys.argv:
+    # NUMA binding uses executable hijacking, which requires spawn.
+    if any(
+        arg.split("=", 1)[0].replace("_", "-")
+        in ("--numa-bind", "--numa-bind-nodes", "--numa-bind-cpus")
+        for arg in sys.argv
+    ):
         reasons.append("NUMA binding requires spawn method")
 
     if cuda_is_initialized():
diff --git a/vllm/utils/vmm_driver.py b/vllm/utils/vmm_driver.py
new file mode 100644
index 000000000000..118f6fc76fe0
--- /dev/null
+++ b/vllm/utils/vmm_driver.py
@@ -0,0 +1,461 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""ctypes bindings for the CUDA (`cuMem*`) and HIP (`hipMem*`) virtual memory
+management APIs, used by :class:`vllm.utils.extensible_tensor.ExtensibleTensor`.
+"""
+
+from __future__ import annotations
+
+import ctypes
+import os
+from collections.abc import Iterable
+from functools import cache
+from typing import Any
+
+_SUCCESS = 0
+_MEM_ALLOCATION_TYPE_PINNED = 1
+_MEM_LOCATION_TYPE_DEVICE = 1
+_MEM_ALLOC_GRANULARITY_MINIMUM = 0
+_MEM_ACCESS_FLAGS_PROT_READWRITE = 3
+_MEM_ALLOCATION_COMP_NONE = 0
+_DEVICE_ATTRIBUTE_GPU_DIRECT_RDMA_WITH_CUDA_VMM_SUPPORTED = 110
+
+DevicePtr = ctypes.c_ulonglong
+MemHandle = ctypes.c_ulonglong
+_Context = ctypes.c_void_p
+
+
+class _MemLocation(ctypes.Structure):
+    _fields_ = [("type", ctypes.c_int), ("id", ctypes.c_int)]
+
+
+class _MemAllocFlags(ctypes.Structure):
+    _fields_ = [
+        ("compressionType", ctypes.c_ubyte),
+        ("gpuDirectRDMACapable", ctypes.c_ubyte),
+        ("usage", ctypes.c_ushort),
+        ("reserved", ctypes.c_ubyte * 4),
+    ]
+
+
+class _MemAllocationProp(ctypes.Structure):
+    # CUmemAllocationProp. hipMemAllocationProp matches it up to allocFlags,
+    # whose HIP variant lacks the trailing reserved bytes; the driver never
+    # reads past the end of its own layout, so the CUDA shape serves both.
+    _fields_ = [
+        ("type", ctypes.c_int),
+        ("requestedHandleTypes", ctypes.c_int),
+        ("location", _MemLocation),
+        ("win32HandleMetaData", ctypes.c_void_p),
+        ("allocFlags", _MemAllocFlags),
+    ]
+
+
+class _MemAccessDesc(ctypes.Structure):
+    _fields_ = [("location", _MemLocation), ("flags", ctypes.c_int)]
+
+
+def _find_loaded_library(
+    lib_name: str, maps_lines: Iterable[str] | None = None
+) -> str | None:
+    """Path of the mapped shared object named lib_name (e.g. libcuda),
+    matched on the file name so that libcudart does not stand in for it."""
+    try:
+        if maps_lines is None:
+            with open("/proc/self/maps") as f:
+                return _find_loaded_library(lib_name, f.readlines())
+        for line in maps_lines:
+            if "/" not in line:
+                continue
+            path = line[line.index("/") :].strip()
+            if os.path.basename(path).startswith(f"{lib_name}.so"):
+                return path
+    except (OSError, ValueError):
+        return None
+    return None
+
+
+class VmmDriver:
+    """Uniform interface over a GPU driver's VMM API.
+
+    Subclasses supply library names and symbols; signatures and structs are
+    shared between CUDA and HIP.
+    """
+
+    # DLPack device type for tensors viewing driver-mapped memory.
+    dlpack_device_type: int
+    _lib_candidates: tuple[str, ...]
+    _lib_search_name: str
+    # Logical name -> library symbol.
+    _symbols: dict[str, str]
+
+    def __init__(self) -> None:
+        self._lib = self._load_library()
+        self._fns: dict[str, Any] = {
+            logical: getattr(self._lib, symbol)
+            for logical, symbol in self._symbols.items()
+        }
+        self._configure_signatures()
+
+    def _load_library(self) -> ctypes.CDLL:
+        # Prefer the copy already in the process, e.g. the HIP runtime a ROCm
+        # torch wheel bundles: loading one by name could add a second runtime.
+        if path := _find_loaded_library(self._lib_search_name):
+            return ctypes.CDLL(path)
+        for name in self._lib_candidates:
+            try:
+                return ctypes.CDLL(name)
+            except OSError:
+                continue
+        raise RuntimeError(
+            f"Could not load {self._lib_candidates[0]}. The GPU driver "
+            "library is required for VMM-backed tensors."
+        )
+
+    def _configure_signatures(self) -> None:
+        pointer = ctypes.POINTER
+        fns = self._fns
+        fns["get_granularity"].argtypes = [
+            pointer(ctypes.c_size_t),
+            pointer(_MemAllocationProp),
+            ctypes.c_int,
+        ]
+        fns["address_reserve"].argtypes = [
+            pointer(DevicePtr),
+            ctypes.c_size_t,
+            ctypes.c_size_t,
+            DevicePtr,
+            ctypes.c_ulonglong,
+        ]
+        fns["create"].argtypes = [
+            pointer(MemHandle),
+            ctypes.c_size_t,
+            pointer(_MemAllocationProp),
+            ctypes.c_ulonglong,
+        ]
+        fns["map"].argtypes = [
+            DevicePtr,
+            ctypes.c_size_t,
+            ctypes.c_size_t,
+            MemHandle,
+            ctypes.c_ulonglong,
+        ]
+        fns["set_access"].argtypes = [
+            DevicePtr,
+            ctypes.c_size_t,
+            pointer(_MemAccessDesc),
+            ctypes.c_size_t,
+        ]
+        fns["unmap"].argtypes = [DevicePtr, ctypes.c_size_t]
+        fns["release"].argtypes = [MemHandle]
+        fns["address_free"].argtypes = [DevicePtr, ctypes.c_size_t]
+        for fn in fns.values():
+            fn.restype = ctypes.c_int
+
+    def error_string(self, code: int) -> str:
+        raise NotImplementedError
+
+    def unusable_runtime_reason(self, device_index: int) -> str | None:
+        """Why this driver's runtime cannot back a growable allocation, if so."""
+        return None
+
+    def ensure_context(self, device_index: int) -> None:
+        """Make sure a driver context for `device_index` is current."""
+        raise NotImplementedError
+
+    def _check(self, result: int) -> None:
+        if result == _SUCCESS:
+            return
+        raise RuntimeError(f"GPU driver error {result}: {self.error_string(result)}")
+
+    def _make_alloc_prop(self, device_index: int) -> _MemAllocationProp:
+        prop = _MemAllocationProp()
+        prop.type = _MEM_ALLOCATION_TYPE_PINNED
+        prop.location.type = _MEM_LOCATION_TYPE_DEVICE
+        prop.location.id = device_index
+        prop.allocFlags.compressionType = _MEM_ALLOCATION_COMP_NONE
+        return prop
+
+    def granularity(self, device_index: int) -> int:
+        prop = self._make_alloc_prop(device_index)
+        granularity = ctypes.c_size_t()
+        self._check(
+            self._fns["get_granularity"](
+                ctypes.byref(granularity),
+                ctypes.byref(prop),
+                _MEM_ALLOC_GRANULARITY_MINIMUM,
+            )
+        )
+        return granularity.value
+
+    def reserve(self, size: int) -> int:
+        """Reserve a virtual address range and return its base pointer."""
+        dptr = DevicePtr()
+        self._check(self._fns["address_reserve"](ctypes.byref(dptr), size, 0, 0, 0))
+        return dptr.value
+
+    def free_reserved(self, ptr: int, size: int) -> None:
+        self._check(self._fns["address_free"](ptr, size))
+
+    def create(self, size: int, device_index: int) -> int:
+        """Create a physical memory handle of `size` bytes."""
+        prop = self._make_alloc_prop(device_index)
+        handle = MemHandle()
+        self._check(
+            self._fns["create"](ctypes.byref(handle), size, ctypes.byref(prop), 0)
+        )
+        return handle.value
+
+    def map(self, ptr: int, size: int, handle: int) -> None:
+        self._check(self._fns["map"](ptr, size, 0, handle, 0))
+
+    def set_access(self, ptr: int, size: int, device_index: int) -> None:
+        desc = _MemAccessDesc()
+        desc.location.type = _MEM_LOCATION_TYPE_DEVICE
+        desc.location.id = device_index
+        desc.flags = _MEM_ACCESS_FLAGS_PROT_READWRITE
+        self._check(self._fns["set_access"](ptr, size, ctypes.byref(desc), 1))
+
+    def unmap(self, ptr: int, size: int) -> None:
+        self._check(self._fns["unmap"](ptr, size))
+
+    def release(self, handle: int) -> None:
+        self._check(self._fns["release"](handle))
+
+
+class CudaVmmDriver(VmmDriver):
+    dlpack_device_type = 2  # kDLCUDA
+    _lib_candidates = ("libcuda.so.1", "libcuda.so")
+    _lib_search_name = "libcuda"
+    _symbols = {
+        "get_granularity": "cuMemGetAllocationGranularity",
+        "address_reserve": "cuMemAddressReserve",
+        "create": "cuMemCreate",
+        "map": "cuMemMap",
+        "set_access": "cuMemSetAccess",
+        "unmap": "cuMemUnmap",
+        "release": "cuMemRelease",
+        "address_free": "cuMemAddressFree",
+    }
+
+    def __init__(self) -> None:
+        super().__init__()
+        lib = self._lib
+        lib.cuGetErrorString.argtypes = [
+            ctypes.c_int,
+            ctypes.POINTER(ctypes.c_char_p),
+        ]
+        lib.cuGetErrorString.restype = ctypes.c_int
+        lib.cuInit.argtypes = [ctypes.c_uint]
+        lib.cuInit.restype = ctypes.c_int
+        lib.cuCtxGetCurrent.argtypes = [ctypes.POINTER(_Context)]
+        lib.cuCtxGetCurrent.restype = ctypes.c_int
+        lib.cuCtxGetDevice.argtypes = [ctypes.POINTER(ctypes.c_int)]
+        lib.cuCtxGetDevice.restype = ctypes.c_int
+        lib.cuDevicePrimaryCtxRetain.argtypes = [
+            ctypes.POINTER(_Context),
+            ctypes.c_int,
+        ]
+        lib.cuDevicePrimaryCtxRetain.restype = ctypes.c_int
+        lib.cuCtxSetCurrent.argtypes = [_Context]
+        lib.cuCtxSetCurrent.restype = ctypes.c_int
+        lib.cuDeviceGetAttribute.argtypes = [
+            ctypes.POINTER(ctypes.c_int),
+            ctypes.c_int,
+            ctypes.c_int,
+        ]
+        lib.cuDeviceGetAttribute.restype = ctypes.c_int
+        self._rdma_capable: dict[int, bool] = {}
+
+    def _make_alloc_prop(self, device_index: int) -> _MemAllocationProp:
+        # Request GPUDirect RDMA support where available, as the sleep-mode
+        # allocator does: without it, KV connectors cannot register the
+        # memory with RDMA NICs (ibv_reg_mr fails with EFAULT).
+        prop = super()._make_alloc_prop(device_index)
+        if (capable := self._rdma_capable.get(device_index)) is None:
+            flag = ctypes.c_int()
+            result = self._lib.cuDeviceGetAttribute(
+                ctypes.byref(flag),
+                _DEVICE_ATTRIBUTE_GPU_DIRECT_RDMA_WITH_CUDA_VMM_SUPPORTED,
+                device_index,
+            )
+            capable = result == _SUCCESS and bool(flag.value)
+            self._rdma_capable[device_index] = capable
+        prop.allocFlags.gpuDirectRDMACapable = int(capable)
+        return prop
+
+    def error_string(self, code: int) -> str:
+        msg = ctypes.c_char_p()
+        self._lib.cuGetErrorString(code, ctypes.byref(msg))
+        return msg.value.decode() if msg.value else "unknown error"
+
+    def ensure_context(self, device_index: int) -> None:
+        """Make the primary context of device_index current, as torch does
+        for the device it runs on; usable before torch has touched the device."""
+        lib = self._lib
+        self._check(lib.cuInit(0))
+        pctx = _Context()
+        self._check(lib.cuCtxGetCurrent(ctypes.byref(pctx)))
+        if pctx.value:
+            device = ctypes.c_int()
+            self._check(lib.cuCtxGetDevice(ctypes.byref(device)))
+            if device.value == device_index:
+                return
+        self._check(lib.cuDevicePrimaryCtxRetain(ctypes.byref(pctx), device_index))
+        self._check(lib.cuCtxSetCurrent(pctx))
+
+
+class HipVmmDriver(VmmDriver):
+    """HIP mirrors the CUDA VMM API with identical signatures, structs and constants."""
+
+    dlpack_device_type = 10  # kDLROCM
+    _lib_candidates = (
+        "libamdhip64.so.7",
+        "libamdhip64.so",
+        "libamdhip64.so.6",
+        "libamdhip64.so.5",
+    )
+    _lib_search_name = "libamdhip64"
+    _symbols = {
+        "get_granularity": "hipMemGetAllocationGranularity",
+        "address_reserve": "hipMemAddressReserve",
+        "create": "hipMemCreate",
+        "map": "hipMemMap",
+        "set_access": "hipMemSetAccess",
+        "unmap": "hipMemUnmap",
+        "release": "hipMemRelease",
+        "address_free": "hipMemAddressFree",
+    }
+
+    def __init__(self) -> None:
+        super().__init__()
+        lib = self._lib
+        lib.hipGetErrorString.argtypes = [ctypes.c_int]
+        lib.hipGetErrorString.restype = ctypes.c_char_p
+        lib.hipGetDevice.argtypes = [ctypes.POINTER(ctypes.c_int)]
+        lib.hipGetDevice.restype = ctypes.c_int
+        lib.hipSetDevice.argtypes = [ctypes.c_int]
+        lib.hipSetDevice.restype = ctypes.c_int
+        lib.hipRuntimeGetVersion.argtypes = [ctypes.POINTER(ctypes.c_int)]
+        lib.hipRuntimeGetVersion.restype = ctypes.c_int
+        lib.hipGetLastError.argtypes = []
+        lib.hipGetLastError.restype = ctypes.c_int
+
+    def _check(self, result: int) -> None:
+        if result != _SUCCESS:
+            # A failed hipMem* call leaves the error sticky for the next
+            # torch launch on this thread; reading it clears it.
+            self._lib.hipGetLastError()
+        super()._check(result)
+
+    def runtime_version(self) -> tuple[int, int]:
+        """(major, minor) of the loaded HIP runtime."""
+        version = ctypes.c_int()
+        self._check(self._lib.hipRuntimeGetVersion(ctypes.byref(version)))
+        return divmod(version.value // 100000, 100)
+
+    def unusable_runtime_reason(self, device_index: int) -> str | None:
+        if not _set_access_fails_within_reservation(self, device_index):
+            return None
+        major, minor = self.runtime_version()
+        return (
+            f"the HIP runtime ({major}.{minor}) rejects hipMemSetAccess for an "
+            "allocation that is not the first in its virtual-address "
+            "reservation (ROCm/rocm-systems#2516, fixed by "
+            "ROCm/rocm-systems#2451 in TheRock 7.12 and later)"
+        )
+
+    def error_string(self, code: int) -> str:
+        msg = self._lib.hipGetErrorString(code)
+        return msg.decode() if msg else "unknown error"
+
+    def ensure_context(self, device_index: int) -> None:
+        # The HIP runtime manages contexts implicitly; just make sure the
+        # buffer's device is current on this thread.
+        device = ctypes.c_int()
+        self._check(self._lib.hipGetDevice(ctypes.byref(device)))
+        if device.value != device_index:
+            self._check(self._lib.hipSetDevice(device_index))
+
+
+def _set_access_fails_within_reservation(driver: VmmDriver, device_index: int) -> bool:
+    """Whether set-access rejects one of two allocations mapped side by side in
+    a reservation, as HIP runtimes before ROCm/rocm-systems#2451 do.
+
+    Those check the requested size against the reservation's allocations
+    counted from the first in their list rather than from the target. With
+    sizes of one and two granules, one of the two calls fails whichever comes
+    first in that list, so the check does not depend on allocation order.
+    """
+    granularity = driver.granularity(device_index)
+    sizes = (granularity, 2 * granularity)
+    base = driver.reserve(sum(sizes))
+    mapped: list[tuple[int, int, int]] = []
+    try:
+        ptr = base
+        for size in sizes:
+            handle = driver.create(size, device_index)
+            try:
+                driver.map(ptr, size, handle)
+            except RuntimeError:
+                driver.release(handle)
+                raise
+            mapped.append((ptr, size, handle))
+            ptr += size
+        try:
+            for ptr, size, _ in mapped:
+                driver.set_access(ptr, size, device_index)
+        except RuntimeError:
+            return True
+        return False
+    finally:
+        for ptr, size, handle in mapped:
+            driver.unmap(ptr, size)
+            driver.release(handle)
+        driver.free_reserved(base, sum(sizes))
+
+
+@cache
+def get_vmm_driver() -> VmmDriver:
+    import torch
+
+    if torch.version.hip is not None:
+        return HipVmmDriver()
+    return CudaVmmDriver()
+
+
+@cache
+def vmm_unavailable_reason() -> str | None:
+    """Why VMM cannot be used on the current device, or None if it can.
+
+    Checks library load, entry points, known runtime defects (by probing), and
+    one reserve/create/map/set-access/unmap/release round trip of a single
+    granule. Covers WSL2 and non-CUDA/ROCm builds.
+    """
+    try:
+        import torch
+
+        if not torch.accelerator.is_available():
+            return "no CUDA/ROCm device is available"
+        driver = get_vmm_driver()
+        device_index = torch.accelerator.current_device_index()
+        driver.ensure_context(device_index)
+        if reason := driver.unusable_runtime_reason(device_index):
+            return reason
+        granularity = driver.granularity(device_index)
+        ptr = driver.reserve(granularity)
+        try:
+            handle = driver.create(granularity, device_index)
+            try:
+                driver.map(ptr, granularity, handle)
+                try:
+                    driver.set_access(ptr, granularity, device_index)
+                finally:
+                    driver.unmap(ptr, granularity)
+            finally:
+                driver.release(handle)
+        finally:
+            driver.free_reserved(ptr, granularity)
+    except Exception as e:
+        return str(e)
+    return None
diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py
index 5c6ce28a956d..5f53f7699298 100644
--- a/vllm/v1/attention/backend.py
+++ b/vllm/v1/attention/backend.py
@@ -106,6 +106,11 @@ def supports_head_size(cls, head_size: int) -> bool:
         supported_head_sizes = cls.get_supported_head_sizes()
         return (not supported_head_sizes) or head_size in supported_head_sizes
 
+    @classmethod
+    def supports_head_size_v(cls, head_size: int, head_size_v: int) -> bool:
+        """Whether the backend can serve a V head size that differs from K's."""
+        return head_size_v == head_size
+
     @classmethod
     def supports_dtype(cls, dtype: torch.dtype) -> bool:
         return dtype in cls.supported_dtypes
@@ -1045,6 +1050,9 @@ class MLAAttentionImpl(AttentionImplBase[T], Generic[T]):
     hisparse_cache: "HiSparseCacheHandle | None" = None
     supports_pcp: bool = True
 
+    def get_fp8_ds_mla_row_bytes(self) -> int:
+        return 656
+
     @abstractmethod
     def __init__(
         self,
diff --git a/vllm/v1/attention/backends/cpu_attn.py b/vllm/v1/attention/backends/cpu_attn.py
index c2f822d067cf..3ca3d26aed00 100644
--- a/vllm/v1/attention/backends/cpu_attn.py
+++ b/vllm/v1/attention/backends/cpu_attn.py
@@ -337,7 +337,6 @@ def __init__(
         self.kv_cache_dtype = kv_cache_dtype
         self.num_queries_per_kv = self.num_heads // self.num_kv_heads
 
-        self.is_fp8_kv_cache = is_quantized_kv_cache(kv_cache_dtype)
         self.attn_type = attn_type
 
         self.sinks = sinks
diff --git a/vllm/v1/attention/backends/flash_attn.py b/vllm/v1/attention/backends/flash_attn.py
index c33def11a8ea..88fe210b9390 100755
--- a/vllm/v1/attention/backends/flash_attn.py
+++ b/vllm/v1/attention/backends/flash_attn.py
@@ -1179,6 +1179,25 @@ def __init__(
                 vllm_config.scheduler_config.max_num_batched_tokens
             )
 
+    def _descale(
+        self, scale: torch.Tensor, scale_float: float, shape: tuple[int, int]
+    ) -> torch.Tensor | None:
+        # FA4 reads a missing descale as 1.0, and it copies every expanded view
+        # into a contiguous tensor: three copy kernels per layer and step. The
+        # scales are fixed once weights are loaded, so a per-tensor scale of 1.0
+        # (FP8 KV cache without calibrated scales) needs none. Per-head scales
+        # keep their tensor: their float mirror holds only the maximum.
+        if (
+            self.vllm_flash_attn_version == 4
+            and scale.numel() == 1
+            and scale_float == 1.0
+        ):
+            return None
+        # A (1, 1) expanded view (one sequence, one KV head per rank) counts
+        # as contiguous, so FA4 skips its copy and its CuTe layout check then
+        # rejects the zero stride.
+        return canonicalize_singleton_dim_strides(scale.expand(shape))
+
     def forward(
         self,
         layer: torch.nn.Module,
@@ -1286,14 +1305,13 @@ def forward(
         scheduler_metadata = attn_metadata.scheduler_metadata
 
         descale_shape = (cu_seqlens_q.shape[0] - 1, self.num_kv_heads)
-
         q_descale = (
-            layer._q_scale.expand(descale_shape)
+            self._descale(layer._q_scale, layer._q_scale_float, descale_shape)
             if self.supports_quant_query_input
             else None
         )
-        k_descale = layer._k_scale.expand(descale_shape)
-        v_descale = layer._v_scale.expand(descale_shape)
+        k_descale = self._descale(layer._k_scale, layer._k_scale_float, descale_shape)
+        v_descale = self._descale(layer._v_scale, layer._v_scale_float, descale_shape)
 
         if self.dcp_world_size > 1:
             self._forward_with_dcp(
diff --git a/vllm/v1/attention/backends/flash_attn_diffkv.py b/vllm/v1/attention/backends/flash_attn_diffkv.py
index 758e2c563c60..e17488db81cf 100644
--- a/vllm/v1/attention/backends/flash_attn_diffkv.py
+++ b/vllm/v1/attention/backends/flash_attn_diffkv.py
@@ -35,6 +35,10 @@ class FlashAttentionDiffKVBackend(FlashAttentionBackend):
     # Default to 128 for this backend
     head_size_v: int = 128
 
+    @classmethod
+    def supports_head_size_v(cls, head_size: int, head_size_v: int) -> bool:
+        return True
+
     @classmethod
     def set_head_size_v(cls, head_size_v: int) -> None:
         cls.head_size_v = head_size_v
diff --git a/vllm/v1/attention/backends/flashinfer.py b/vllm/v1/attention/backends/flashinfer.py
index 1a8c6c377ab9..52372ade5a4e 100755
--- a/vllm/v1/attention/backends/flashinfer.py
+++ b/vllm/v1/attention/backends/flashinfer.py
@@ -78,6 +78,7 @@
 from vllm.v1.attention.backends.utils import (
     get_dcp_local_seq_lens,
     get_flashinfer_layout_string,
+    get_kv_cache_dtype_from_layers,
     get_num_attention_heads_from_layers,
     get_per_layer_parameters,
     infer_global_hyperparameters,
@@ -945,7 +946,7 @@ def __init__(
         self.page_size = self.kv_cache_spec.block_size
 
         if self.kv_cache_spec.kv_quant_mode != KVQuantMode.NONE:
-            self.cache_dtype = self.cache_config.cache_dtype
+            self.cache_dtype = self._resolve_cache_dtype()
             # Cannot use self.kv_cache_spec.dtype here because kv_cache_spec
             # storage dtype may not be the same as the op dtype (uint8 vs fp8_e4m3)
             self.is_kvcache_nvfp4 = self.cache_dtype.startswith("nvfp4")
@@ -1123,6 +1124,22 @@ def __init__(
             max_num_reqs, dtype=torch.int32, device=self.device, pin_memory=False
         )
 
+    def _resolve_cache_dtype(self) -> str:
+        cache_dtype = self.cache_config.cache_dtype
+        if cache_dtype != "auto":
+            return cache_dtype
+
+        layer_cache_dtype = get_kv_cache_dtype_from_layers(
+            self.vllm_config, self.layer_names
+        )
+        if layer_cache_dtype is not None and layer_cache_dtype != "auto":
+            return layer_cache_dtype
+
+        raise ValueError(
+            "FlashInfer requires a resolved KV cache dtype for quantized "
+            f"attention layers {self.layer_names}."
+        )
+
     @property
     def kv_cache_layout(self) -> KVCacheLayout:
         return self.cache_config.get_resolved_kv_cache_layout()
@@ -2043,7 +2060,6 @@ def __init__(
         self.nvfp4_fa2 = self.is_kvcache_nvfp4 and _nvfp4_kv_on_fa2()
         self.nvfp4_trtllm = self.is_kvcache_nvfp4 and not self.nvfp4_fa2
         self.kv_cache_dtype = "nvfp4" if self.is_kvcache_nvfp4 else kv_cache_dtype
-        self.fp4_data_dim = head_size // 2 if self.is_kvcache_nvfp4 else 0
         self.logits_soft_cap = logits_soft_cap
         self.kv_sharing_target_layer_name = kv_sharing_target_layer_name
 
diff --git a/vllm/v1/attention/backends/flashinfer_ts.py b/vllm/v1/attention/backends/flashinfer_ts.py
new file mode 100644
index 000000000000..248eb7f65013
--- /dev/null
+++ b/vllm/v1/attention/backends/flashinfer_ts.py
@@ -0,0 +1,512 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""FlashInfer task-scheduled (``prims_ts``) paged decode.
+
+Wraps ``flashinfer.attention.prims_ts.BatchDecodePagedTSWrapper`` for uniform
+query-block decode: every request contributes exactly ``max_query_len`` query
+tokens over a paged KV cache. That is the DFlash/DSpark drafter shape, and the
+reason this backend exists is that ``mask_type="dense"`` expresses a
+*non-causal* draft block directly -- the trtllm-gen paged decode kernel is
+bottom-right causal and silently ignores a supplied mask.
+
+Layout. No KV cache change: K and V are the two halves of vLLM's ordinary
+``[pages, Hkv, page_size, 2 * D]`` page, so their token stride is ``2 * D``.
+The kernel builds its TMA descriptors from the tensors' own strides
+(``_validate_native_hnd_tma_strides``), so the only requirement is a
+contiguous head_dim. Verified numerically equal to the same data repacked
+into compact planes.
+
+Sliding window. ``prims_ts`` rejects ``window_left >= 0`` unless the mask is
+causal, so a non-causal windowed layer cannot be expressed directly. The query
+block always sits at the end of the sequence, so the symmetric window reduces
+to a suffix, and we hand the kernel a block table restricted to the window's
+pages (see :func:`_restrict_to_window`). Because the restriction is page
+aligned, each query token sees up to ``page_size - 1`` extra tokens at the old
+edge, and all query tokens share one window start rather than each having its
+own (``max_query_len - 1`` tokens of drift).
+"""
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, ClassVar
+
+import torch
+
+from vllm.config import VllmConfig
+from vllm.logger import init_logger
+from vllm.platforms import current_platform
+from vllm.utils.math_utils import cdiv
+from vllm.v1.attention.backend import (
+    AttentionBackend,
+    AttentionCGSupport,
+    AttentionImpl,
+    AttentionLayer,
+    AttentionMetadata,
+    AttentionMetadataBuilder,
+    AttentionType,
+    CommonAttentionMetadata,
+    MultipleOf,
+    max_decode_query_len,
+)
+from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheSpec
+
+if TYPE_CHECKING:
+    from vllm.config.cache import CacheDType
+    from vllm.platforms.interface import DeviceCapability
+
+logger = init_logger(__name__)
+
+# flashinfer/attention/prims_ts/decode.py: _SUPPORTED_PAGE_SIZES,
+# _SUPPORTED_HEAD_DIMS, _SUPPORTED_COMPUTE_CAPABILITIES. (10, 7) is SM107a /
+# Rubin, which the kernel additionally gates on a CuTe DSL that can emit for
+# the sm_100f family (>= 4.8, else CUTE_DSL_ARCH=sm_100f in the environment).
+_TS_PAGE_SIZES = (16, 32, 64, 128)
+_TS_HEAD_DIMS = (64, 128, 256)
+_TS_CAPABILITIES = ((10, 0), (10, 3), (10, 7))
+_TS_MAX_HEAD_RATIO = 128
+
+
+def _require_prims_ts():
+    """Import the prims_ts wrapper, or explain which flashinfer is needed.
+
+    0.7.0.post1 ships prims_ts but calls ``Task._run_task_body_impl`` with a
+    ``context`` kwarg that cutlass-dsl 4.8.0 removed, so it raises TypeError
+    deep inside kernel codegen.
+    """
+    try:
+        from flashinfer.attention.prims_ts import BatchDecodePagedTSWrapper
+    except ImportError as exc:
+        raise RuntimeError(
+            "FLASHINFER_TS needs flashinfer-python>=0.7.1 "
+            "(flashinfer.attention.prims_ts is missing)."
+        ) from exc
+    import flashinfer
+
+    if tuple(int(p) for p in flashinfer.__version__.split(".")[:3]) < (0, 7, 1):
+        raise RuntimeError(
+            f"FLASHINFER_TS needs flashinfer-python>=0.7.1, found "
+            f"{flashinfer.__version__}: earlier builds of prims_ts are "
+            "incompatible with cutlass-dsl>=4.8."
+        )
+    return BatchDecodePagedTSWrapper
+
+
+def _ts_output_dtype(q_dtype: torch.dtype, model_dtype: torch.dtype) -> torch.dtype:
+    """The output dtype the kernel can write for this query dtype.
+
+    An FP8 query only yields FP16 or FP8 out, so a bf16 model stages through
+    FP16 and converts; every other combination writes the model dtype directly.
+    """
+    if q_dtype is torch.float8_e4m3fn and model_dtype is not torch.float16:
+        return torch.float16
+    return model_dtype
+
+
+def _kv_kernel_dtype(spec: AttentionSpec) -> torch.dtype:
+    """An FP8 cache is stored as uint8 and viewed as e4m3 at use."""
+    if spec.dtype is torch.uint8:
+        return current_platform.fp8_dtype()
+    return spec.dtype
+
+
+@dataclass
+class FlashInferTSMetadata(AttentionMetadata):
+    # Per-request K/V lengths the kernel reads, already restricted to the
+    # window when one applies.
+    seq_lens: torch.Tensor
+    # [B, C] int32 page ids, likewise restricted.
+    block_table: torch.Tensor
+    num_reqs: int
+    num_actual_tokens: int
+    query_len: int
+    # Plan bound at build time; the same object must be replayed by a captured
+    # graph, which is why plans are keyed only on (batch size, query length).
+    wrapper: Any
+
+
+class FlashInferTSMetadataBuilder(AttentionMetadataBuilder[FlashInferTSMetadata]):
+    _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
+
+    def __init__(
+        self,
+        kv_cache_spec: KVCacheSpec,
+        layer_names: list[str],
+        vllm_config: VllmConfig,
+        device: torch.device,
+    ):
+        super().__init__(kv_cache_spec, layer_names, vllm_config, device)
+        assert isinstance(kv_cache_spec, AttentionSpec)
+        self.page_size = kv_cache_spec.block_size
+        self.num_kv_heads = kv_cache_spec.num_kv_heads
+        self.head_size = kv_cache_spec.head_size
+        # Q and K/V dtypes are kept matched: FP8 query against the FP8 cache,
+        # or model-dtype against an unquantized one. The FP8 query is produced
+        # upstream (MiniMax-M3's fused QK-norm/RoPE kernel writes it through
+        # q_fp8_out), so the attention layer does not quantize it again.
+        self.kv_dtype = _kv_kernel_dtype(kv_cache_spec)
+        self.q_dtype = self.kv_dtype
+        self.out_dtype = vllm_config.model_config.dtype
+        self.num_qo_heads = vllm_config.model_config.get_num_attention_heads(
+            vllm_config.parallel_config
+        )
+        # FullAttentionSpec records the window when the hybrid KV cache manager
+        # is disabled; SlidingWindowSpec carries it directly.
+        self.sliding_window: int | None = getattr(kv_cache_spec, "sliding_window", None)
+        self.query_len = max_decode_query_len(vllm_config)
+        self._plans: dict[tuple[int, int, bool], Any] = {}
+        self._init_reorder_batch_threshold(1, supports_spec_as_decode=True)
+
+        # CUDA-graph-safe metadata. This builder advertises UNIFORM_BATCH, so
+        # the drafter's attention is captured and replayed; the kernel records
+        # the ADDRESS of seq_lens and the page table at capture time. Building
+        # fresh tensors per step would leave a replayed graph reading freed
+        # memory, so every build writes into these buffers in place.
+        max_reqs = max(
+            vllm_config.scheduler_config.max_num_seqs,
+            vllm_config.scheduler_config.max_num_batched_tokens
+            // max(self.query_len, 1),
+        )
+        self._buf_seq_lens = torch.zeros(max_reqs, dtype=torch.int32, device=device)
+        self._buf_block_table = torch.zeros(
+            (max_reqs, self._window_pages(self.query_len)),
+            dtype=torch.int32,
+            device=device,
+        )
+
+    def _window_pages(self, query_len: int) -> int:
+        """Page-table width a restricted (windowed) batch needs."""
+        if self.sliding_window is None:
+            return cdiv(self.vllm_config.model_config.max_model_len, self.page_size)
+        return cdiv(self.sliding_window + query_len - 1, self.page_size) + 1
+
+    def _max_kv_len(self) -> int:
+        """Constant K/V capacity every plan compiles for.
+
+        Keeping it constant is what makes one plan per (batch size, query
+        length) safe to replay from a CUDA graph: ``build()`` can never pick a
+        differently-compiled plan for a batch the graph already captured.
+        """
+        if self.sliding_window is not None:
+            # _restrict_to_window leaves at most one partial page in front of
+            # the window plus the query block itself.
+            return self.sliding_window + self.query_len + self.page_size
+        return self.vllm_config.model_config.max_model_len
+
+    def _get_plan(self, num_reqs: int, query_len: int, causal: bool) -> Any:
+        key = (num_reqs, query_len, causal)
+        plan = self._plans.get(key)
+        if plan is not None:
+            return plan
+
+        BatchDecodePagedTSWrapper = _require_prims_ts()
+
+        window = self.sliding_window
+        wrapper = BatchDecodePagedTSWrapper(kv_layout="HND")
+        wrapper.plan(
+            device=self.device,
+            batch_size=num_reqs,
+            num_qo_heads=self.num_qo_heads,
+            num_kv_heads=self.num_kv_heads,
+            head_dim=self.head_size,
+            page_size=self.page_size,
+            max_kv_len=self._max_kv_len(),
+            max_seq_len_q=query_len,
+            packed_query=False,
+            q_data_type=self.q_dtype,
+            kv_data_type=self.kv_dtype,
+            o_data_type=_ts_output_dtype(self.q_dtype, self.out_dtype),
+            # A dense plan is what makes a non-causal draft block expressible;
+            # it cannot carry a window, which _restrict_to_window applies to
+            # the page table instead.
+            mask_type="causal" if causal else "dense",
+            window_left=window - 1 if (causal and window is not None) else -1,
+        )
+        self._plans[key] = wrapper
+        return wrapper
+
+    def _restrict_to_window(
+        self, seq_lens: torch.Tensor, block_table: torch.Tensor, query_len: int
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Cut the page table down to the pages the window can reach.
+
+        The query block ends the sequence, so a symmetric window around it is
+        the suffix ``[L - w - query_len + 1, L)``. Round that start down to a
+        page so the slice is expressible as a page table, then renumber the
+        lengths relative to the new first page.
+        """
+        assert self.sliding_window is not None
+        span = self.sliding_window + query_len - 1
+        num_pages = self._buf_block_table.shape[1]
+        num_reqs = seq_lens.shape[0]
+
+        start_page = torch.clamp_min(
+            (seq_lens - span).div(self.page_size, rounding_mode="floor"), 0
+        )
+        cols = torch.arange(num_pages, device=block_table.device, dtype=torch.int32)
+        idx = start_page.unsqueeze(1).to(torch.int64) + cols.unsqueeze(0)
+        idx = idx.clamp_max(block_table.shape[1] - 1)
+
+        assert num_reqs <= self._buf_seq_lens.shape[0], (
+            f"FLASHINFER_TS metadata buffer holds {self._buf_seq_lens.shape[0]} "
+            f"requests, got {num_reqs}"
+        )
+        out_lens = self._buf_seq_lens[:num_reqs]
+        out_table = self._buf_block_table[:num_reqs]
+        out_lens.copy_(seq_lens - start_page.to(seq_lens.dtype) * self.page_size)
+        out_table.copy_(block_table.gather(1, idx))
+        return out_lens, out_table
+
+    def build(
+        self,
+        common_prefix_len: int,
+        common_attn_metadata: CommonAttentionMetadata,
+        fast_build: bool = False,
+    ) -> FlashInferTSMetadata:
+        num_reqs = common_attn_metadata.num_reqs
+        num_tokens = common_attn_metadata.num_actual_tokens
+        causal = common_attn_metadata.causal
+        if isinstance(causal, torch.Tensor):
+            raise NotImplementedError(
+                "FLASHINFER_TS has no per-sequence causality; it plans one "
+                "mask for the whole batch."
+            )
+        # The drafter contributes the same number of query tokens per request,
+        # so the block length is the quotient; max_query_len is only an upper
+        # bound. Startup profiling also builds ragged batches that no real
+        # draft step produces (e.g. 5 tokens over 2 requests). Rounding down
+        # covers a whole number of blocks and leaves the remainder untouched,
+        # which keeps those builds in bounds instead of failing the engine.
+        query_len = num_tokens // num_reqs if num_reqs else 0
+        num_tokens = num_reqs * query_len
+        # The model runner builds metadata for every KV cache group on every
+        # batch, so this is also called with the target's prefill shapes for
+        # the drafter's own group -- metadata the speculator never uses, since
+        # it builds its own. A drafter block is never longer than
+        # max_decode_query_len, so anything wider is not a batch this backend
+        # runs; planning it would compile a prefill-shaped plan and fail
+        # validation (seq_len_q > max_kv_len).
+        if query_len > self.query_len:
+            query_len = 0
+
+        seq_lens = common_attn_metadata.seq_lens[:num_reqs]
+        block_table = common_attn_metadata.block_table_tensor[:num_reqs]
+        if self.sliding_window is not None and not causal:
+            seq_lens, block_table = self._restrict_to_window(
+                seq_lens, block_table, query_len
+            )
+        else:
+            num_pages = min(block_table.shape[1], self._buf_block_table.shape[1])
+            out_lens = self._buf_seq_lens[:num_reqs]
+            out_table = self._buf_block_table[:num_reqs, :num_pages]
+            out_lens.copy_(seq_lens)
+            out_table.copy_(block_table[:, :num_pages])
+            seq_lens, block_table = out_lens, out_table
+
+        if query_len == 0:
+            return FlashInferTSMetadata(
+                seq_lens=common_attn_metadata.seq_lens[:num_reqs],
+                block_table=common_attn_metadata.block_table_tensor[:num_reqs],
+                num_reqs=num_reqs,
+                num_actual_tokens=0,
+                query_len=0,
+                wrapper=None,
+            )
+
+        return FlashInferTSMetadata(
+            seq_lens=seq_lens,
+            block_table=block_table,
+            num_reqs=num_reqs,
+            num_actual_tokens=num_tokens,
+            query_len=query_len,
+            wrapper=self._get_plan(num_reqs, query_len, causal),
+        )
+
+
+class FlashInferTSImpl(AttentionImpl[FlashInferTSMetadata]):
+    forward_includes_kv_cache_update: bool = False
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        scale: float,
+        num_kv_heads: int | None = None,
+        alibi_slopes: list[float] | None = None,
+        sliding_window: int | None = None,
+        kv_cache_dtype: str = "auto",
+        logits_soft_cap: float | None = None,
+        attn_type: str = AttentionType.DECODER,
+        kv_sharing_target_layer_name: str | None = None,
+        sinks: torch.Tensor | None = None,
+    ) -> None:
+        if sinks is not None:
+            raise NotImplementedError("FLASHINFER_TS does not support attention sinks.")
+        if alibi_slopes is not None:
+            raise NotImplementedError("FLASHINFER_TS does not support alibi.")
+        if logits_soft_cap:
+            raise NotImplementedError("FLASHINFER_TS does not support softcap.")
+        if kv_sharing_target_layer_name is not None:
+            raise NotImplementedError("FLASHINFER_TS does not support KV sharing.")
+        if attn_type != AttentionType.DECODER:
+            raise NotImplementedError(
+                f"FLASHINFER_TS does not support attention type {attn_type}."
+            )
+        self.num_heads = num_heads
+        self.head_size = head_size
+        self.scale = float(scale)
+        self.num_kv_heads = num_kv_heads or num_heads
+        self.sliding_window = sliding_window
+        self.kv_cache_dtype = kv_cache_dtype
+        # The kernel wants the query in the cache's dtype. For FP8 the caller
+        # supplies it already quantized (query_pre_quantized on the layer), so
+        # advertise support rather than making the layer quantize a second time.
+        self.supports_quant_query_input = kv_cache_dtype.startswith("fp8")
+
+    def _kv_views(self, kv_cache: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        """HND K/V views of vLLM's ``[pages, Hkv, page, 2 * D]`` page.
+
+        Their token stride is ``2 * D`` rather than ``D``; the kernel builds its
+        TMA descriptors from the tensor's own strides, so no repacking and no
+        layout change are needed.
+        """
+        if self.kv_cache_dtype.startswith("fp8"):
+            kv_cache = kv_cache.view(current_platform.fp8_dtype())
+        return kv_cache[..., : self.head_size], kv_cache[..., self.head_size :]
+
+    def do_kv_cache_update(
+        self,
+        layer: AttentionLayer,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        kv_cache: torch.Tensor,
+        slot_mapping: torch.Tensor,
+    ) -> None:
+        from vllm._custom_ops import reshape_and_cache_flash
+
+        key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1)
+        reshape_and_cache_flash(
+            key,
+            value,
+            key_cache,
+            value_cache,
+            slot_mapping,
+            self.kv_cache_dtype,
+            layer._k_scale,
+            layer._v_scale,
+        )
+
+    def forward(
+        self,
+        layer: AttentionLayer,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        kv_cache: torch.Tensor,
+        attn_metadata: FlashInferTSMetadata,
+        output: torch.Tensor,
+        output_scale: torch.Tensor | None = None,
+        output_block_scale: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        if attn_metadata is None:
+            # Profiling run: the KV cache is not bound yet.
+            return output
+        if output_scale is not None or output_block_scale is not None:
+            raise NotImplementedError(
+                "FLASHINFER_TS does not support fused output quantization."
+            )
+
+        n = attn_metadata.num_actual_tokens
+        if n == 0:
+            return output
+        b, sq = attn_metadata.num_reqs, attn_metadata.query_len
+        k_cache, v_cache = self._kv_views(kv_cache)
+        q = query[:n].view(b, sq, self.num_heads, self.head_size)
+
+        out = output[:n].view(b, sq, self.num_heads, self.head_size)
+        kernel_dtype = _ts_output_dtype(q.dtype, out.dtype)
+        needs_cast = kernel_dtype is not out.dtype
+        kernel_out = torch.empty_like(out, dtype=kernel_dtype) if needs_cast else out
+
+        attn_metadata.wrapper.run(
+            q=q,
+            paged_kv_cache=(k_cache, v_cache),
+            seq_lens=attn_metadata.seq_lens,
+            block_tables=attn_metadata.block_table,
+            bmm1_scale=self.scale * layer._q_scale_float * layer._k_scale_float,
+            bmm2_scale=layer._v_scale_float,
+            out=kernel_out,
+            validate=False,
+        )
+        if needs_cast:
+            out.copy_(kernel_out)
+        return output
+
+
+class FlashInferTSBackend(AttentionBackend):
+    supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
+    supported_kv_cache_dtypes: ClassVar[list["CacheDType"]] = [
+        "auto",
+        "float16",
+        "bfloat16",
+        "fp8",
+        "fp8_e4m3",
+    ]
+    forward_includes_kv_cache_update: bool = False
+
+    @staticmethod
+    def get_name() -> str:
+        return "FLASHINFER_TS"
+
+    @staticmethod
+    def get_impl_cls() -> type[FlashInferTSImpl]:
+        return FlashInferTSImpl
+
+    @staticmethod
+    def get_builder_cls() -> type[FlashInferTSMetadataBuilder]:
+        return FlashInferTSMetadataBuilder
+
+    @classmethod
+    def get_supported_head_sizes(cls) -> list[int]:
+        return list(_TS_HEAD_DIMS)
+
+    @staticmethod
+    def get_supported_kernel_block_sizes(
+        kv_cache_spec=None,
+    ) -> list[int | MultipleOf]:
+        return list(_TS_PAGE_SIZES)
+
+    @classmethod
+    def supports_sliding_window(cls) -> bool:
+        # Causal layers use the kernel's own window_left; non-causal ones get a
+        # page-restricted block table (see the module docstring).
+        return True
+
+    @classmethod
+    def supports_non_causal(cls) -> bool:
+        return True
+
+    @classmethod
+    def supports_attn_type(cls, attn_type: str) -> bool:
+        return attn_type == AttentionType.DECODER
+
+    @classmethod
+    def supports_compute_capability(cls, capability: "DeviceCapability") -> bool:
+        return (capability.major, capability.minor) in _TS_CAPABILITIES
+
+    @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:
+        if use_mla or use_sparse or use_mm_prefix:
+            return "FLASHINFER_TS supports dense non-MLA attention only"
+        if has_sink:
+            return "FLASHINFER_TS does not support attention sinks"
+        return None
diff --git a/vllm/v1/attention/backends/gdn_attn.py b/vllm/v1/attention/backends/gdn_attn.py
index 660ef5387b17..effd2a3838b3 100644
--- a/vllm/v1/attention/backends/gdn_attn.py
+++ b/vllm/v1/attention/backends/gdn_attn.py
@@ -20,6 +20,7 @@
     AttentionCGSupport,
     AttentionMetadataBuilder,
     CommonAttentionMetadata,
+    max_decode_query_len,
 )
 from vllm.v1.attention.backends.utils import (
     NULL_BLOCK_ID,
@@ -27,7 +28,7 @@
     mamba_get_block_table_tensor,
     split_decodes_and_prefills,
 )
-from vllm.v1.kv_cache_interface import MambaSpec
+from vllm.v1.kv_cache_interface import KVCacheSpec, MambaSpec
 
 # FlashInfer's GDN prefill emits states every N tokens, N a multiple of its
 # 64-token chunk, so checkpoint offsets must be multiples of 64.
@@ -84,6 +85,10 @@ def get_builder_cls() -> type["GDNAttentionMetadataBuilder"]:
     def is_ssm(cls) -> bool:
         return True
 
+    @classmethod
+    def supports_device_cpu_query_lens_mismatch(cls) -> bool:
+        return True
+
 
 @dataclass
 class GDNAttentionMetadata:
@@ -135,6 +140,15 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
     kv_cache_spec: MambaSpec
     _cudagraph_support = AttentionCGSupport.UNIFORM_BATCH
 
+    @classmethod
+    def get_varlen_cudagraph_max_query_len(
+        cls,
+        vllm_config: VllmConfig,
+        kv_cache_spec: KVCacheSpec,
+    ) -> int | None:
+        del kv_cache_spec
+        return max_decode_query_len(vllm_config)
+
     reorder_batch_threshold: int = 1
 
     def __init__(
@@ -147,6 +161,10 @@ def __init__(
         super().__init__(kv_cache_spec, layer_names, vllm_config, device)
         self.compilation_config = vllm_config.compilation_config
         self.speculative_config = vllm_config.speculative_config
+        self._adaptive_verification = bool(
+            self.speculative_config
+            and self.speculative_config.enable_adaptive_verification
+        )
         self.checkpoint_builder = MambaPrefillCheckpointBuilder(
             vllm_config, kv_cache_spec
         )
@@ -398,7 +416,7 @@ def build(  # type: ignore[override]
             non_spec_sequence_masks_cpu = ~spec_sequence_masks_cpu
             query_lens_cpu = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]
             spec_query_lens_cpu = query_lens_cpu[spec_sequence_masks_cpu]
-            if spec_query_lens_cpu.numel() > 0:
+            if spec_query_lens_cpu.numel() > 0 and not self._adaptive_verification:
                 first_spec_sequence_length = int(spec_query_lens_cpu[0])
                 if first_spec_sequence_length > 0 and bool(
                     torch.all(spec_query_lens_cpu == first_spec_sequence_length)
diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla.py b/vllm/v1/attention/backends/mla/flashinfer_mla.py
index 91492859d5d6..2179ab2a6ded 100644
--- a/vllm/v1/attention/backends/mla/flashinfer_mla.py
+++ b/vllm/v1/attention/backends/mla/flashinfer_mla.py
@@ -341,6 +341,8 @@ def forward_mqa(
         cum_seq_lens_q: torch.Tensor | None = None
         max_q_len: int | None = None
         row_req: torch.Tensor | None = None
+        # Per split row: query tokens of its request after it.
+        causal_rows_after: torch.Tensor | None = None
 
         if not attn_metadata.causal:
             # FlashInfer decode has no causal flag. Flatten each non-causal
@@ -357,15 +359,23 @@ def forward_mqa(
             # 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(
-                """FlashInferMLAImpl got a query of uneven length.
-                This usually indicates an issue in batch reordering
-                or incorrect setup in dummy_run."""
-            )
+        elif (
+            attn_metadata.num_decode_tokens
+            != attn_metadata.num_decodes * attn_metadata.decode.max_query_len
+        ):
+            # Ragged causal DCP decode, e.g. requests with fewer drafts: split
+            # into single-token rows, each bounded at its own causal position.
             q = q.unsqueeze(1)
+            block_table, seq_lens, row_req = self._flattened_decode_metadata(
+                attn_metadata, q.shape[0]
+            )
+            query_start_loc = attn_metadata.decode.query_start_loc
+            assert query_start_loc is not None
+            query_end = query_start_loc[1:]
+            rows = torch.arange(q.shape[0], device=q.device, dtype=query_end.dtype)
+            causal_rows_after = query_end[row_req] - 1 - rows
         else:
+            # trtllm API requires extra dimension q_len_per_request for MTP
             q = q.view(attn_metadata.num_decodes, -1, q.shape[-2], q.shape[-1])
 
         if self.bmm1_scale is None:
@@ -389,6 +399,8 @@ def forward_mqa(
             assert causal_seqlens_kv_global is not None
             if row_req is not None:
                 causal_seqlens_kv_global = causal_seqlens_kv_global[row_req]
+            if causal_rows_after is not None:
+                causal_seqlens_kv_global = causal_seqlens_kv_global - causal_rows_after
             extra_kwargs.update(
                 enable_dcp=True,
                 cp_world=self.dcp_world_size,
diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py
index 1f8338dd1de3..3231a9249c24 100644
--- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py
+++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py
@@ -42,6 +42,9 @@
 
 logger = init_logger(__name__)
 
+# CUDA's gridDim.z limit; trtllm-gen MLA decode launches one z-slice per request.
+_TRTLLM_GEN_MAX_REQUESTS = 65535
+
 
 class _FlashInferMLASparseBackendBase(AttentionBackend):
     """Common metadata for concrete FlashInfer sparse MLA backends."""
@@ -366,14 +369,16 @@ def compute_trtllm_sparse_mla_workspace_bytes(
     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.
+    returned unchanged. A launch covers at most ``_TRTLLM_GEN_MAX_REQUESTS``
+    tokens, and launches run in order on one stream, so larger batches reuse
+    the same slab.
     """
     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
+        * min(max_num_batched_tokens, _TRTLLM_GEN_MAX_REQUESTS)
         * _TRTLLM_GEN_SOFTMAX_SLOTS_PER_TOKEN
         + _TRTLLM_GEN_SOFTMAX_GUARD_BYTES
     )
@@ -612,7 +617,6 @@ def _forward_mqa_kernel(
         # deferred until MTP is validated end-to-end for this backend.
         query = q.unsqueeze(1)
         block_tables = topk_indices.unsqueeze(1)
-        seq_lens_arg = seq_lens
 
         # page_table width = topk buffer width, which kpool widens past
         # index_topk (topk_tokens) and rounds up to a multiple of 128. The
@@ -623,45 +627,52 @@ def _forward_mqa_kernel(
         # mismatches the page_table when index_kpool > 1.
         sparse_topk_capacity = topk_indices.shape[1]
 
-        extra_kwargs: dict[str, torch.Tensor] = {}
         needs_empty_query_guard = self.need_to_return_lse_for_decode or isinstance(
             self.index_group, HiSparseMLAIndexGroup
         )
+        topk_lens = None
         if self.is_nope_mla:
             # 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_cache.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=block_tables,
-            seq_lens=seq_lens_arg,
-            max_seq_len=sparse_topk_capacity,
-            bmm1_scale=self.bmm1_scale,
-            bmm2_scale=self.bmm2_scale,
-            sparse_mla_top_k=sparse_topk_capacity,
-            return_lse=self.need_to_return_lse_for_decode,
-            **extra_kwargs,
+
+        num_tokens, num_heads = q.shape[0], q.shape[1]
+        out = q.new_empty(
+            (num_tokens, num_heads, self.kv_lora_rank), dtype=torch.bfloat16
         )
+        lse = None
         if self.need_to_return_lse_for_decode:
-            assert isinstance(kernel_out, tuple)
-            o, lse = kernel_out
-        else:
-            assert isinstance(kernel_out, torch.Tensor)
-            o = kernel_out
-            lse = None
-
-        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])
+            lse = q.new_empty((num_tokens, num_heads), dtype=torch.float32)
+
+        # One trtllm-gen request per token, and requests map to gridDim.z: a
+        # launch over more than 65535 tokens fails without raising and leaves
+        # ``out`` unwritten (flashinfer-ai/flashinfer#3125). Launch row chunks.
+        for start in range(0, num_tokens, _TRTLLM_GEN_MAX_REQUESTS):
+            end = min(start + _TRTLLM_GEN_MAX_REQUESTS, num_tokens)
+            extra_kwargs: dict[str, torch.Tensor] = {}
+            if topk_lens is not None:
+                extra_kwargs["sparse_mla_top_k_lens"] = topk_lens[start:end]
+            if lse is not None:
+                extra_kwargs["lse"] = lse[start:end]
+            trtllm_batch_decode_with_kv_cache_mla(
+                query=query[start:end],
+                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,
+                qk_rope_head_dim=self.qk_rope_head_dim,
+                block_tables=block_tables[start:end],
+                seq_lens=seq_lens[start:end],
+                max_seq_len=sparse_topk_capacity,
+                out=out[start:end].unsqueeze(1),
+                bmm1_scale=self.bmm1_scale,
+                bmm2_scale=self.bmm2_scale,
+                sparse_mla_top_k=sparse_topk_capacity,
+                return_lse=lse is not None,
+                **extra_kwargs,
+            )
         if self.is_nope_mla and needs_empty_query_guard:
             empty_queries = seq_lens == 0
             if lse is not None:
@@ -674,27 +685,3 @@ def _forward_mqa_kernel(
             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
-    def _normalize_lse(
-        lse: torch.Tensor,
-        num_tokens: int,
-        num_heads: int,
-    ) -> torch.Tensor:
-        # FlashInfer returns the decode LSE either as 2D (num_tokens, num_heads)
-        # or 3D ((num_tokens, num_heads, 1) / (num_tokens, 1, num_heads)).
-        # Collapse all of these to the (num_tokens, num_heads) the shared DCP
-        # reducer expects.
-        if lse.dim() == 3:
-            if lse.shape[-1] == 1:
-                lse = lse.squeeze(-1)
-            elif lse.shape[1] == 1:
-                lse = lse.squeeze(1)
-            elif lse.shape[0] * lse.shape[1] == num_tokens:
-                lse = lse.reshape(num_tokens, lse.shape[-1])
-        if lse.shape != (num_tokens, num_heads):
-            raise RuntimeError(
-                "Unexpected FlashInfer sparse MLA LSE shape: "
-                f"{tuple(lse.shape)}, expected ({num_tokens}, {num_heads})."
-            )
-        return lse
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 04b6edf1256b..7fc81ded2299 100644
--- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py
+++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py
@@ -6,6 +6,7 @@
 
 import torch
 
+from vllm.config import get_current_vllm_config
 from vllm.model_executor.layers.attention.sparse_mla_attention import (
     SparseMLACommonImpl,
 )
@@ -36,6 +37,11 @@ class FlashInferMLASparseSM120Impl(SparseMLACommonImpl[FlashInferMLASparseMetada
     is_sparse = True
     supports_dense_mha_prefill = False
 
+    def get_fp8_ds_mla_row_bytes(self) -> int:
+        if self.kv_lora_rank == 512 and self.qk_rope_head_dim == 0:
+            return 528
+        return super().get_fp8_ds_mla_row_bytes()
+
     def __init__(
         self,
         num_heads: int,
@@ -83,8 +89,6 @@ def __init__(
             topk_indices_buffer=topk_indices_buffer,
             **mla_args,
         )
-        from vllm.config import get_current_vllm_config
-
         vllm_config = get_current_vllm_config()
         model_type = None
         if vllm_config.model_config is not None:
diff --git a/vllm/v1/attention/backends/mla/index_group.py b/vllm/v1/attention/backends/mla/index_group.py
index ba787b38da55..60a966e8ebf8 100644
--- a/vllm/v1/attention/backends/mla/index_group.py
+++ b/vllm/v1/attention/backends/mla/index_group.py
@@ -17,21 +17,12 @@
     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."""
@@ -45,11 +36,14 @@ class SparseMLAIndexGroup:
     has_indexer: bool
     num_layers: int = 0
 
+    def _current_stream(self) -> Any:
+        return current_stream()
+
     def register_layer(
         self,
         vllm_config: VllmConfig | None = None,
         *,
-        head_size: int | None = None,
+        row_width: int | None = None,
         kv_cache_dtype: str | None = None,
     ) -> int:
         layer_index = self.num_layers
@@ -66,7 +60,7 @@ def set_logical_topk_ready(self, layer_index: int) -> None:
         # touches this event), so skip the record as well.
         if in_piecewise_cudagraph():
             return
-        self.logical_topk_ready.record(current_stream())
+        self.logical_topk_ready.record(self._current_stream())
 
     def prepare_for_batch(self, layer_index: int, attn_metadata: Any | None) -> None:
         pass
@@ -121,6 +115,7 @@ def _convert_once(
         *,
         block_stride_rows: int | None,
         return_valid_counts: bool,
+        idx_mapping: torch.Tensor | None = None,
     ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
         num_tokens = logical_topk_indices.shape[0]
         if num_tokens > self.physical_topk_indices.shape[0]:
@@ -132,11 +127,12 @@ def _convert_once(
                 BLOCK_STRIDE_ROWS=block_stride_rows,
                 NUM_TOPK_TOKENS=logical_topk_indices.shape[1],
                 return_valid_counts=return_valid_counts,
+                idx_mapping=idx_mapping,
             )
 
         physical_topk_indices = self.physical_topk_indices[:num_tokens]
         valid_topk_counts = self.valid_topk_counts[:num_tokens]
-        compute_stream = current_stream()
+        compute_stream = self._current_stream()
         if layer_index == 0:
             if self.has_indexer and not in_piecewise_cudagraph():
                 self.side_stream.wait_event(self.logical_topk_ready)
@@ -155,6 +151,7 @@ def _convert_once(
                     valid_counts_out=(
                         valid_topk_counts if return_valid_counts else None
                     ),
+                    idx_mapping=idx_mapping,
                 )
                 self.physical_topk_ready.record(self.side_stream)
         compute_stream.wait_event(self.physical_topk_ready)
@@ -173,26 +170,27 @@ class HiSparseMLAIndexGroup(SparseMLAIndexGroup):
     prefill_stream: torch.Stream | None = None
     prefill_ready_events: list[torch.Event] = field(default_factory=list)
 
+    def _current_stream(self) -> Any:
+        return torch.accelerator.current_stream()
+
     def register_layer(
         self,
         vllm_config: VllmConfig | None = None,
         *,
-        head_size: int | None = None,
+        row_width: int | None = None,
         kv_cache_dtype: str | None = None,
     ) -> int:
         assert vllm_config is not None
-        assert head_size is not None
+        assert row_width 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)
+            self.prefill_stream = torch.Stream(device=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
             )
@@ -207,7 +205,7 @@ def register_layer(
         assert cache is not None
         self.caches.append(cache)
         cache.index_group_caches = self.caches
-        self.prefill_ready_events.append(_create_event())
+        self.prefill_ready_events.append(torch.Event())
         return layer_index
 
     def cache(self, layer_index: int) -> HiSparseCacheHandle:
@@ -246,10 +244,11 @@ def convert_logical_to_physical_topk(
                 layer_index,
                 logical_topk_indices,
                 req_id_per_token,
-                leader.batch_block_table(),
+                leader.block_table,
                 leader.view.block_size,
                 block_stride_rows=leader.view.attention_block_stride,
                 return_valid_counts=return_valid_counts,
+                idx_mapping=leader.runtime.request_state_indices,
             )
         source_block_table = cache.source_block_table
         assert source_block_table is not None
@@ -325,7 +324,7 @@ def gather_fp8_prefill(
         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())
+        prefill_stream.wait_stream(torch.accelerator.current_stream())
         with prefill_stream:
             ops.cp_gather_and_upconvert_fp8_kv_cache(
                 cache.view.cache,
@@ -361,7 +360,7 @@ def register_layer(
         is_index_producing_layer: bool,
         vllm_config: VllmConfig | None = None,
         *,
-        head_size: int | None = None,
+        row_width: int | None = None,
         kv_cache_dtype: str | None = None,
     ) -> tuple[SparseMLAIndexGroup, int]:
         if is_index_producing_layer or self.current_group is None:
@@ -385,15 +384,15 @@ def register_layer(
                     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(),
+                side_stream=torch.Stream(device=self.logical_topk_indices.device),
+                logical_topk_ready=torch.Event(),
+                physical_topk_ready=torch.Event(),
                 has_indexer=is_index_producing_layer,
             )
         group = self.current_group
         group_index = group.register_layer(
             vllm_config,
-            head_size=head_size,
+            row_width=row_width,
             kv_cache_dtype=kv_cache_dtype,
         )
         return group, group_index
diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py
index c353a1f2fbb8..c03d19d98cbb 100644
--- a/vllm/v1/attention/backends/mla/indexer.py
+++ b/vllm/v1/attention/backends/mla/indexer.py
@@ -10,6 +10,11 @@
 from vllm.config import CUDAGraphMode, VllmConfig, get_current_vllm_config
 from vllm.distributed import get_dcp_group, get_pcp_group
 from vllm.logger import init_logger
+from vllm.model_executor.kernels.attention.dsa.fused_mqa_topk import (
+    FUSED_MQA_TOPK_MAX_KEYS,
+    FUSED_MQA_TOPK_MAX_ROWS,
+    use_fused_mqa_topk,
+)
 from vllm.model_executor.warmup.jit_warmup import kernel_launcher, zip_inputs
 from vllm.model_executor.warmup.jit_warmup_triton_helper import (
     LaunchSpec,
@@ -927,6 +932,19 @@ def get_max_prefill_buffer_size(vllm_config: VllmConfig):
     return max_model_len * (576 * 2 // 132) * factor
 
 
+def dense_logits_rows(max_seqlen_k: int, max_logits_bytes: int | None = None) -> int:
+    """Query rows per dense MQA-logits call whose fp32 [rows, max_seqlen_k]
+    logits fit max_logits_bytes (default: VLLM_SPARSE_INDEXER_MAX_LOGITS_MB).
+
+    The prefill chunk planner sub-chunks requests with it; the indexers split
+    the rows of chunks planned for fused_mqa_topk (not bounded by the budget)
+    that take the dense path.
+    """
+    if max_logits_bytes is None:
+        max_logits_bytes = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024
+    return max(1, max_logits_bytes // (4 * max(1, max_seqlen_k)))
+
+
 def _supports_varlen_paged_mqa_logits() -> bool:
     return (
         current_platform.is_cuda()
@@ -1025,7 +1043,29 @@ def __init__(self, *args, block_table_width: int, **kwargs) -> None:
             if self.vllm_config.speculative_config
             else 0
         )
+        self.use_adaptive_verification = bool(
+            self.vllm_config.speculative_config
+            and self.vllm_config.speculative_config.enable_adaptive_verification
+        )
         self.indexer_uses_fp4 = dsa_indexer_uses_fp4(self.vllm_config)
+        # Plan prefill chunks for the fused MQA+top-k kernel (not bounded by
+        # the logits budget). PCP + DCP has its own planner. v4.1 candidate
+        # selection is not excluded: its pre-candidate indexers take
+        # fused_mqa_topk and the others row-split their dense logits. Indexers
+        # that pad their heads to 32 run fused_mqa_topk on budget-planned
+        # chunks.
+        model_config = self.vllm_config.model_config
+        cfg = model_config.hf_text_config if model_config is not None else None
+        self.fused_mqa_topk_eligible = cfg is not None and use_fused_mqa_topk(
+            query_shape=(
+                getattr(cfg, "index_n_heads", 0),
+                getattr(cfg, "index_head_dim", 0),
+            ),
+            topk=getattr(cfg, "index_topk", 0)
+            // (getattr(cfg, "index_kpool", None) or 1),
+            fp4_cache=self.indexer_uses_fp4,
+            full_context=self.dcp_world_size == 1,
+        )
 
         next_n = self.num_speculative_tokens + 1
         self.decode_threshold = next_n
@@ -1089,8 +1129,14 @@ def __init__(self, *args, block_table_width: int, **kwargs) -> None:
         self.dcp_rank_tensor = torch.tensor(
             self.dcp_rank, dtype=torch.int32, device=self.device
         )
+        # Only decode tokens are expanded: at most next_n per request. Allocated
+        # after the memory profile, so it must not scale with the token budget.
+        max_decode_tokens = min(
+            scheduler_config.max_num_seqs * next_n,
+            scheduler_config.max_num_batched_tokens,
+        )
         self.expanded_block_table_buffer = torch.zeros(
-            (scheduler_config.max_num_batched_tokens, block_table_width),
+            (max_decode_tokens, block_table_width),
             dtype=torch.int32,
             device=self.device,
         )
@@ -1362,6 +1408,7 @@ def _split_indexer_prefill_chunks(
         workspace_size: int,
         max_logits_bytes: int,
         request_offset: int = 0,
+        fused_mqa_topk_eligible: bool = False,
     ) -> list[tuple[slice, slice]]:
         """Split this step's prefill requests into chunks, respecting:
         - N constraint: total_seq_lens <= workspace_size (existing O(N)
@@ -1371,6 +1418,13 @@ def _split_indexer_prefill_chunks(
         When a single request-level chunk still exceeds the logits budget,
         sub-chunks on the query dimension (M) to bound peak memory.
 
+        With ``fused_mqa_topk_eligible`` (chunks for the fused MQA+top-k kernel,
+        which needs no logits buffer) chunks whose gathered keys fit
+        FUSED_MQA_TOPK_MAX_KEYS are bounded by FUSED_MQA_TOPK_MAX_ROWS rows
+        instead of the logits budget. Requests are never packed past that key
+        limit, so only a single request longer than it takes the logits-budget
+        path. Chunks the dispatch keeps dense are row-split by the indexer.
+
         Returns list of (req_slice, query_slice) tuples.
         """
         chunks: list[tuple[slice, slice]] = []
@@ -1387,7 +1441,15 @@ def _split_indexer_prefill_chunks(
                     compressed_seq_lens_cpu[end].item(),
                 )
                 new_m, new_n = chunk_m + q, chunk_n + s
-                if new_n <= workspace_size and new_m * new_n <= max_logits_elems:
+                if fused_mqa_topk_eligible and new_n <= FUSED_MQA_TOPK_MAX_KEYS:
+                    fits = new_m <= FUSED_MQA_TOPK_MAX_ROWS
+                else:
+                    # fused_mqa_topk planning never packs past the key limit.
+                    fits = (
+                        not fused_mqa_topk_eligible
+                        and new_m * new_n <= max_logits_elems
+                    )
+                if new_n <= workspace_size and fits:
                     chunk_m, chunk_n = new_m, new_n
                     end += 1
                 else:
@@ -1404,8 +1466,12 @@ def _split_indexer_prefill_chunks(
 
             req_slice = slice(start + request_offset, end + request_offset)
             max_q = (
-                max(1, max_logits_elems // chunk_n) if chunk_n > 0 else max(1, chunk_m)
+                dense_logits_rows(chunk_n, max_logits_bytes)
+                if chunk_n > 0
+                else max(1, chunk_m)
             )
+            if fused_mqa_topk_eligible and chunk_n <= FUSED_MQA_TOPK_MAX_KEYS:
+                max_q = FUSED_MQA_TOPK_MAX_ROWS
             for q_off in range(0, chunk_m, max_q):
                 sub_m = min(max_q, chunk_m - q_off)
                 chunks.append((req_slice, slice(q_off, q_off + sub_m)))
@@ -1541,6 +1607,7 @@ def build(
                     self.max_prefill_buffer_size,
                     max_logits_bytes,
                     request_offset=num_decodes,
+                    fused_mqa_topk_eligible=self.fused_mqa_topk_eligible,
                 )
 
             chunks = []
@@ -1598,6 +1665,7 @@ def build(
 
         decode_metadata = None
         if num_decodes > 0:
+            assert num_decode_tokens <= self.expanded_block_table_buffer.shape[0]
             if not self.supports_varlen:
                 torch.diff(
                     common_attn_metadata.query_start_loc[: num_decodes + 1],
@@ -1629,6 +1697,16 @@ def build(
             min_decode_len = int(decode_lens_cpu.min().item())
             write_is_uniform = min_decode_len == max_decode_len
             next_n = 1 + self.num_speculative_tokens
+            if self.use_adaptive_verification:
+                # Varlen FULL graphs are captured by distributing the graph's
+                # token bucket over as many one-token dummy requests as
+                # possible. Runtime adaptive budgets can put up to next_n
+                # tokens in one request, so the capture-time scatter width and
+                # branch must describe that upper bound rather than the dummy
+                # partition. Device per-request lengths still select the live
+                # rows at replay.
+                max_decode_len = next_n
+                write_is_uniform = False
             # The kernel sees max_decode_len Q rows, not the configured next_n,
             # so legality is per-step: on SM90 a uniformly 3-deep batch has no
             # native kernel. max_decode_len <= 1 always has one.
diff --git a/vllm/v1/attention/backends/mla/prefill/flash_attn.py b/vllm/v1/attention/backends/mla/prefill/flash_attn.py
index b6cf96e5b738..8f2dab7321ee 100644
--- a/vllm/v1/attention/backends/mla/prefill/flash_attn.py
+++ b/vllm/v1/attention/backends/mla/prefill/flash_attn.py
@@ -449,6 +449,11 @@ def _flash_attn_varlen_diff_headdims(
         if isinstance(attn_out, tuple):
             attn_out, lse = attn_out[0], attn_out[1]
 
+        # Upstream flash_attn on ROCm may return a head-major output.
+        # merge_attn_states needs token-major partials, so normalize it here.
+        if attn_out.stride(0) != attn_out.shape[1] * attn_out.stride(1):
+            attn_out = attn_out.contiguous()
+
         # Remain consistent with old `flash_attn_varlen_func` where there
         # is only one output tensor if `return_softmax_lse` is False.
         if return_softmax_lse:
diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py
index 1cfb2ae215fd..44003604e447 100644
--- a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py
+++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py
@@ -2,7 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import functools
-from collections.abc import Callable
+from collections.abc import Callable, Iterable
 from dataclasses import dataclass
 from enum import Enum
 from pathlib import Path
@@ -102,6 +102,19 @@ def _segmented_mla_page_size(block_size: int) -> int:
     return min(128, largest_power_of_2_divisor(block_size))
 
 
+def _max_mla_metadata_info(
+    infos: Iterable[tuple[tuple[int | tuple[int, ...], torch.dtype], ...]],
+) -> list[tuple[tuple[int, ...], torch.dtype]]:
+    """Per-dimension maximum of several ``get_mla_metadata_info_v1`` results."""
+    merged = []
+    for entries in zip(*infos):
+        shapes = [
+            shape if isinstance(shape, tuple) else (shape,) for shape, _ in entries
+        ]
+        merged.append((tuple(max(dims) for dims in zip(*shapes)), entries[0][1]))
+    return merged
+
+
 @functools.lru_cache(maxsize=1)
 def _get_mla_gluon():
     """Load the small-head Gluon MLA entry point."""
@@ -486,13 +499,40 @@ def _aiter_mla_small_head_mode() -> str:
     return mode
 
 
+def _adaptive_dcp_verify_unsupported_reason(vllm_config: VllmConfig) -> str | None:
+    parallel_config = vllm_config.parallel_config
+    dcp_world_size = parallel_config.decode_context_parallel_size
+    if dcp_world_size == 1:
+        return None
+    cp_interleave = parallel_config.cp_kv_cache_interleave_size
+    if cp_interleave != 1:
+        return (
+            "adaptive verification with decode context parallelism requires "
+            f"cp_kv_cache_interleave_size 1, got {cp_interleave}"
+        )
+    asm_selected = (
+        _asm_dcp_verify_configured(dcp_world_size, cp_interleave, True)
+        and max_decode_query_len(vllm_config) >= _MIN_CPRR_QLEN
+    )
+    if not (
+        asm_selected
+        or _segmented_dcp_verify_supported(dcp_world_size, cp_interleave)
+        or _triton_dcp_verify_supported(dcp_world_size, cp_interleave)
+    ):
+        return (
+            "adaptive verification with decode context parallelism needs the "
+            "round-robin asm decode, segmented MLA or the gfx942 Triton verify"
+        )
+    return None
+
+
 def _adaptive_varlen_verify_enabled(vllm_config: VllmConfig) -> bool:
     speculative_config = vllm_config.speculative_config
     return (
         speculative_config is not None
         and bool(getattr(speculative_config, "enable_adaptive_verification", False))
-        and vllm_config.parallel_config.decode_context_parallel_size == 1
         and max_decode_query_len(vllm_config) > 1
+        and _adaptive_dcp_verify_unsupported_reason(vllm_config) is None
     )
 
 
@@ -587,6 +627,20 @@ def validate_configuration(
             and kv_cache_dtype in (None, "auto", "float16")
         ):
             invalid_reasons.append("non-causal fp16 MLA decode not supported")
+        if use_adaptive_verification and use_dcp:
+            from vllm.config import get_current_vllm_config_or_none
+
+            vllm_config = get_current_vllm_config_or_none()
+            reason = (
+                _adaptive_dcp_verify_unsupported_reason(vllm_config)
+                if vllm_config is not None
+                else None
+            )
+            if reason is not None:
+                raise ValueError(
+                    f"{reason}. Disable enable_adaptive_verification in the "
+                    "speculative config or set decode_context_parallel_size=1."
+                )
         return invalid_reasons
 
     @staticmethod
@@ -764,6 +818,10 @@ def _uniform_padded_mtp_qo_len(
 
         return uniform_qo_len
 
+    @property
+    def _varlen_right_align(self) -> bool:
+        return self.dcp_world_size > 1
+
     @classmethod
     def get_varlen_cudagraph_max_query_len(
         cls: type["AiterMLAMetadataBuilder"],
@@ -922,7 +980,9 @@ def __init__(
         max_num_reqs = vllm_config.scheduler_config.max_num_seqs
         max_num_pages = max_num_reqs * max_num_pages_per_req
         speculative_config = vllm_config.speculative_config
-        if _adaptive_varlen_verify_enabled(vllm_config):
+        if _adaptive_varlen_verify_enabled(vllm_config) and not (
+            self._varlen_right_align
+        ):
             assert speculative_config is not None
             max_num_pages += max_num_reqs * (
                 speculative_config.num_speculative_tokens or 0
@@ -994,6 +1054,9 @@ def __init__(
         # wrong split/reduce metadata for the gfx950 fp8 nhead=32 fold path.
         self._mla_q_dtype = q_dtype
         self._mla_kv_dtype = kv_dtype
+        # The metadata is built for every qlen up to _mtp_decode_qlen, and
+        # aiter's buffer needs are not monotonic in qlen: on gfx950 a bf16
+        # block of fewer than 128 packed rows is tiled by 64 instead of 128.
         (
             (work_meta_data_size, work_meta_data_type),
             (work_indptr_size, work_indptr_type),
@@ -1001,21 +1064,24 @@ def __init__(
             (reduce_indptr_size, reduce_indptr_type),
             (reduce_final_map_size, reduce_final_map_type),
             (reduce_partial_map_size, reduce_partial_map_type),
-        ) = get_mla_metadata_info_v1(
-            max_num_reqs,
-            self._mtp_decode_qlen,
-            self._num_attention_heads,
-            q_dtype,
-            kv_dtype,
-            is_sparse=False,
-            fast_mode=True,
-            # Must match max_split_per_batch in get_mla_metadata_v1; without it
-            # aiter sizes the reduce scratch from capacity (~9 GiB).
-            **(
-                dict(max_split_per_batch=self._mla_max_split_per_batch)
-                if self._asm_dcp_verify
-                else {}
-            ),
+        ) = _max_mla_metadata_info(
+            get_mla_metadata_info_v1(
+                max_num_reqs,
+                qlen,
+                self._num_attention_heads,
+                q_dtype,
+                kv_dtype,
+                is_sparse=False,
+                fast_mode=True,
+                # Must match max_split_per_batch in get_mla_metadata_v1; without
+                # it aiter sizes the reduce scratch from capacity (~9 GiB).
+                **(
+                    dict(max_split_per_batch=self._mla_max_split_per_batch)
+                    if self._asm_dcp_verify
+                    else {}
+                ),
+            )
+            for qlen in range(1, self._mtp_decode_qlen + 1)
         )
         self._mla_work_meta_data = torch.empty(
             work_meta_data_size, dtype=work_meta_data_type, device=device
@@ -1564,8 +1630,10 @@ def _build_decode(
                 )
 
         varlen_q_rows = varlen_o_rows = None
-        if varlen_decode:
+        varlen_left_align = varlen_decode and not self._varlen_right_align
+        if varlen_left_align:
             seq_lens_for_kernel = seq_lens_device + (max_qo_len - qo_lens_device)
+        if varlen_decode:
             varlen_q_rows, varlen_o_rows = self._fill_varlen_rows(
                 query_start_loc_device, qo_lens_device, num_reqs
             )
@@ -1630,7 +1698,7 @@ def _build_decode(
             # device synchronisation. Programs whose chunk lies past
             # num_tokens mask out entirely.
             max_tokens_per_req = block_table_tensor.shape[1] * self.kernel_block_size
-            if varlen_decode:
+            if varlen_left_align:
                 max_tokens_per_req += max_qo_len - 1
             num_chunks = max(1, cdiv(max_tokens_per_req, 1024))
             _expand_page_indices_kernel[(num_reqs, num_chunks)](
@@ -1642,7 +1710,7 @@ def _build_decode(
                 seq_lens_device,
                 KERNEL_BLOCK_SIZE=self.kernel_block_size,
                 BLOCK_SIZE=1024,
-                CLAMP_TAIL=varlen_decode,
+                CLAMP_TAIL=varlen_left_align,
             )
             paged_kv_indices = self.paged_kv_indices
 
@@ -1886,7 +1954,9 @@ def _fill_varlen_rows(
         row_pos = self._varlen_row_pos[:num_rows]
         row_qlen = qo_lens_device.to(torch.int64).index_select(0, row_req)
         row_start = query_start_loc_device.to(torch.int64).index_select(0, row_req)
-        is_real = row_pos < row_qlen
+        if self._varlen_right_align:
+            row_pos = row_pos - (width - row_qlen)
+        is_real = (row_pos >= 0) & (row_pos < row_qlen)
         packed_row = row_start + row_pos
         q_rows = self._varlen_q_rows[:num_rows]
         q_rows.copy_(torch.where(is_real, packed_row, 0))
@@ -2994,8 +3064,11 @@ def forward_mqa(
             num_tokens = q.shape[0]
             q = q.index_select(0, q_rows)
         o, lse = self._forward_mqa(q, kv_c_and_k_pe_cache, attn_metadata, layer)
-        assert lse is None
-        return o.index_select(0, decode.varlen_o_rows[:num_tokens]), None
+        o_rows = decode.varlen_o_rows[:num_tokens]
+        assert lse is None or self.dcp_world_size > 1
+        if lse is not None:
+            lse = lse.index_select(0, o_rows)
+        return o.index_select(0, o_rows), lse
 
     def _forward_mqa(
         self,
diff --git a/vllm/v1/attention/backends/mla/rocm_paged_mxfp4_indexer.py b/vllm/v1/attention/backends/mla/rocm_paged_mxfp4_indexer.py
index 77c9eba68ca4..837f7ae964e3 100644
--- a/vllm/v1/attention/backends/mla/rocm_paged_mxfp4_indexer.py
+++ b/vllm/v1/attention/backends/mla/rocm_paged_mxfp4_indexer.py
@@ -392,6 +392,7 @@ def _split_indexer_prefill_chunks(  # type: ignore[override]
         workspace_size: int,
         max_logits_bytes: int,
         request_offset: int = 0,
+        fused_mqa_topk_eligible: bool = False,
     ) -> list[tuple[slice, slice]]:
         return split_prefill_chunks(
             compressed_seq_lens_cpu,
diff --git a/vllm/v1/attention/backends/mla/sparse_utils.py b/vllm/v1/attention/backends/mla/sparse_utils.py
index 2447df5ceb71..9af7ad8d21b0 100644
--- a/vllm/v1/attention/backends/mla/sparse_utils.py
+++ b/vllm/v1/attention/backends/mla/sparse_utils.py
@@ -81,6 +81,7 @@ class CompileKey:
         block_stride_rows: int
         block_n: int
         has_prefill_workspace: bool
+        has_idx_mapping: bool
         count_valid: bool
         single_tile: bool
         compact_to_front: bool
@@ -99,6 +100,7 @@ def sparse_mla_index_remap_kernel(
         token_indices_ptr,  # int32 [num_tokens, NUM_TOPK_TOKENS]
         out_ptr,  # int32 [num_tokens, NUM_TOPK_TOKENS]
         valid_count_ptr,  # int32 [num_tokens] - output valid count per row
+        idx_mapping_ptr,  # int32 [num_requests] - block table row per request
         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
@@ -109,6 +111,7 @@ def sparse_mla_index_remap_kernel(
         BLOCK_N: tl.constexpr,  # tile width along columns
         NUM_TOPK_TOKENS: tl.constexpr,
         HAS_PREFILL: tl.constexpr,
+        HAS_IDX_MAPPING: tl.constexpr,
         COUNT_VALID: tl.constexpr,  # whether to count valid indices
         # BLOCK_N >= NUM_TOPK_TOKENS: one program owns the query, so its count
         # is an in-register reduction and needs no atomic.
@@ -141,6 +144,8 @@ def sparse_mla_index_remap_kernel(
 
         # Load request id for this token (no mask: grid is exact)
         req = tl.load(req_id_ptr + token_id)
+        if HAS_IDX_MAPPING:
+            req = tl.maximum(tl.load(idx_mapping_ptr + req), 0)
 
         # Load token indices for this tile
         ti_ptr = token_indices_ptr + token_id * ti_stride0 + indice_id * ti_stride1
@@ -239,6 +244,7 @@ def dispatch(  # type: ignore[override]
         BLOCK_N: int,
         NUM_TOPK_TOKENS: int,
         HAS_PREFILL_WORKSPACE: bool,
+        HAS_IDX_MAPPING: bool,
         COUNT_VALID: bool,
         COMPACT_TO_FRONT: bool,
         DCP_SIZE: int,
@@ -252,6 +258,7 @@ def dispatch(  # type: ignore[override]
             block_stride_rows=BLOCK_STRIDE_ROWS,
             block_n=tiling[1],
             has_prefill_workspace=HAS_PREFILL_WORKSPACE,
+            has_idx_mapping=HAS_IDX_MAPPING,
             count_valid=COUNT_VALID,
             single_tile=tiling[0],
             compact_to_front=COMPACT_TO_FRONT,
@@ -279,6 +286,20 @@ def get_warmup_keys(
             block_size * dcp_size,
         )
         max_num_blocks_per_req = get_block_table_width(max_num_blocks, block_size)
+        hisparse_combos: list[dict[str, Any]] = []
+        if vllm_config.attention_config.hisparse_config is not None:
+            hisparse_combos.extend(
+                dict(
+                    HAS_PREFILL_WORKSPACE=False,
+                    HAS_IDX_MAPPING=True,
+                    COUNT_VALID=count_valid,
+                    COMPACT_TO_FRONT=count_valid,
+                    DCP_SIZE=1,
+                    DCP_RANK=0,
+                    DCP_INTERLEAVE=1,
+                )
+                for count_valid in (False, True)
+            )
         pcp_dcp_combos = []
         if (
             vllm_config.parallel_config.prefill_context_parallel_size > 1
@@ -288,6 +309,7 @@ def get_warmup_keys(
             pcp_dcp_combos.append(
                 dict(
                     HAS_PREFILL_WORKSPACE=True,
+                    HAS_IDX_MAPPING=False,
                     COUNT_VALID=True,
                     COMPACT_TO_FRONT=True,
                     DCP_SIZE=dcp_size,
@@ -299,6 +321,7 @@ def get_warmup_keys(
             zip_inputs(
                 dict(
                     HAS_PREFILL_WORKSPACE=False,
+                    HAS_IDX_MAPPING=False,
                     COUNT_VALID=False,
                     COMPACT_TO_FRONT=False,
                     DCP_SIZE=1,
@@ -307,6 +330,7 @@ def get_warmup_keys(
                 ),
                 dict(
                     HAS_PREFILL_WORKSPACE=False,
+                    HAS_IDX_MAPPING=False,
                     COUNT_VALID=True,
                     COMPACT_TO_FRONT=True,
                     DCP_SIZE=1,
@@ -315,6 +339,7 @@ def get_warmup_keys(
                 ),
                 dict(
                     HAS_PREFILL_WORKSPACE=True,
+                    HAS_IDX_MAPPING=False,
                     COUNT_VALID=True,
                     COMPACT_TO_FRONT=True,
                     DCP_SIZE=1,
@@ -323,6 +348,7 @@ def get_warmup_keys(
                 ),
                 dict(
                     HAS_PREFILL_WORKSPACE=False,
+                    HAS_IDX_MAPPING=False,
                     COUNT_VALID=True,
                     COMPACT_TO_FRONT=False,
                     DCP_SIZE=dcp_size,
@@ -331,6 +357,7 @@ def get_warmup_keys(
                 ),
                 dict(
                     HAS_PREFILL_WORKSPACE=False,
+                    HAS_IDX_MAPPING=False,
                     COUNT_VALID=True,
                     COMPACT_TO_FRONT=dcp_size > 1,
                     DCP_SIZE=dcp_size,
@@ -338,6 +365,7 @@ def get_warmup_keys(
                     DCP_INTERLEAVE=dcp_interleave,
                 ),
                 *pcp_dcp_combos,
+                *hisparse_combos,
             ),
             BLOCK_SIZE=block_size,
             BLOCK_STRIDE_ROWS=block_stride_rows,
@@ -373,6 +401,7 @@ def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
             prefill_workspace_starts=(
                 int32_ptr if compile_key.has_prefill_workspace else None
             ),
+            idx_mapping=(int32_ptr if compile_key.has_idx_mapping else None),
             max_num_blocks_per_req=1,
             workspace_rank_stride=1,
             BLOCK_SIZE=compile_key.block_size,
@@ -380,6 +409,7 @@ def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
             BLOCK_N=compile_key.block_n,
             NUM_TOPK_TOKENS=compile_key.num_topk_tokens,
             HAS_PREFILL_WORKSPACE=compile_key.has_prefill_workspace,
+            HAS_IDX_MAPPING=compile_key.has_idx_mapping,
             COUNT_VALID=compile_key.count_valid,
             COMPACT_TO_FRONT=compile_key.compact_to_front,
             DCP_SIZE=compile_key.dcp_size,
@@ -397,6 +427,7 @@ def __call__(
         valid_counts: torch.Tensor | None,
         prefill_workspace_request_ids: torch.Tensor | None,
         prefill_workspace_starts: torch.Tensor | None,
+        idx_mapping: torch.Tensor | None,
         *,
         workspace_rank_stride: int,
         max_num_blocks_per_req: int,
@@ -405,6 +436,7 @@ def __call__(
         BLOCK_N: int,
         NUM_TOPK_TOKENS: int,
         HAS_PREFILL_WORKSPACE: bool,
+        HAS_IDX_MAPPING: bool,
         COUNT_VALID: bool,
         COMPACT_TO_FRONT: bool,
         DCP_SIZE: int,
@@ -419,6 +451,7 @@ def __call__(
             valid_count_ptr=valid_counts,
             prefill_request_id_ptr=prefill_workspace_request_ids,
             workspace_starts_ptr=prefill_workspace_starts,
+            idx_mapping_ptr=idx_mapping,
             BLOCK_N=block_n,
             HAS_PREFILL=HAS_PREFILL_WORKSPACE,
             SINGLE_TILE=single_tile,
@@ -475,6 +508,7 @@ def triton_convert_req_index_to_global_index(
     return_valid_counts: bool = False,
     out: torch.Tensor | None = None,
     valid_counts_out: torch.Tensor | None = None,
+    idx_mapping: torch.Tensor | None = None,
 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
     """out[token_id, indice_id] =
         block_table[req_id[token_id],
@@ -494,11 +528,17 @@ def triton_convert_req_index_to_global_index(
     prefill_workspace_starts: int32 [num_prefills], 0-indexed workspace
         starts for each prefill request
 
+    When idx_mapping (int32 [num_requests]) is given, block_table is
+    indexed by idx_mapping[req_id[token_id]] instead of req_id[token_id].
+
     When return_valid_counts is True, also returns the count of valid (non -1)
     indices per row, computed during the same kernel pass (no extra overhead).
     """
     assert req_id.dtype == torch.int32
     assert block_table.dtype == torch.int32
+    assert idx_mapping is None or (
+        idx_mapping.dtype == torch.int32 and idx_mapping.is_contiguous()
+    )
     assert token_indices.dtype == torch.int32
     assert req_id.shape[0] == token_indices.shape[0], (
         f"req_id ({req_id.shape[0]}) and token_indices ({token_indices.shape[0]}) "
@@ -572,6 +612,7 @@ def triton_convert_req_index_to_global_index(
         valid_counts,
         prefill_workspace_request_ids,
         prefill_workspace_starts,
+        idx_mapping,
         # shapes / constexprs
         workspace_rank_stride=prefill_workspace_rank_stride or 0,
         max_num_blocks_per_req=max_num_blocks_per_req,
@@ -582,6 +623,7 @@ def triton_convert_req_index_to_global_index(
         BLOCK_N=BLOCK_N,
         NUM_TOPK_TOKENS=NUM_TOPK_TOKENS,
         HAS_PREFILL_WORKSPACE=HAS_PREFILL_WORKSPACE,
+        HAS_IDX_MAPPING=idx_mapping is not None,
         COUNT_VALID=return_valid_counts,
         COMPACT_TO_FRONT=return_valid_counts,
         DCP_SIZE=dcp_size,
@@ -680,6 +722,7 @@ def triton_filter_and_convert_dcp_index(
         # No prefill workspace on the DCP decode path.
         None,
         None,
+        None,
         workspace_rank_stride=0,
         max_num_blocks_per_req=max_num_blocks_per_req,
         BLOCK_SIZE=BLOCK_SIZE,
@@ -689,6 +732,7 @@ def triton_filter_and_convert_dcp_index(
         BLOCK_N=BLOCK_N,
         NUM_TOPK_TOKENS=NUM_TOPK_TOKENS,
         HAS_PREFILL_WORKSPACE=False,
+        HAS_IDX_MAPPING=False,
         COUNT_VALID=count_valid,
         COMPACT_TO_FRONT=compact_valid_to_front,
         DCP_SIZE=dcp_size,
diff --git a/vllm/v1/attention/backends/mla/xpu_mla_sparse.py b/vllm/v1/attention/backends/mla/xpu_mla_sparse.py
index 9fe9a27c8d89..463313f74d1c 100644
--- a/vllm/v1/attention/backends/mla/xpu_mla_sparse.py
+++ b/vllm/v1/attention/backends/mla/xpu_mla_sparse.py
@@ -220,6 +220,12 @@ def __init__(
         self.softmax_scale = scale
         self.init_topk_indices_buffer(indexer, topk_indices_buffer)
 
+    def record_logical_topk_ready(self) -> None:
+        # Called unconditionally by MLAAttention.forward for sparse layers. This
+        # impl shares the top-k buffer via SharedTopkIndicesBuffer but does not
+        # participate in sparse-MLA index groups (same as ROCm).
+        pass
+
     def _forward_bf16_kv(
         self,
         q: torch.Tensor,  # [sq, heads, d_qk]
diff --git a/vllm/v1/attention/backends/registry.py b/vllm/v1/attention/backends/registry.py
index 1e40298c7d69..eca068d40b34 100644
--- a/vllm/v1/attention/backends/registry.py
+++ b/vllm/v1/attention/backends/registry.py
@@ -69,6 +69,7 @@ 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"
+    FLASHINFER_TS = "vllm.v1.attention.backends.flashinfer_ts.FlashInferTSBackend"
     TRITON_FLASHINFER = (
         "vllm.v1.attention.backends.triton_flashinfer.TritonFlashInferBackend"
     )
diff --git a/vllm/v1/attention/backends/rocm_aiter_unified_attn.py b/vllm/v1/attention/backends/rocm_aiter_unified_attn.py
index 26aea0095635..3944d17cdaa3 100644
--- a/vllm/v1/attention/backends/rocm_aiter_unified_attn.py
+++ b/vllm/v1/attention/backends/rocm_aiter_unified_attn.py
@@ -87,6 +87,12 @@ def supports_block_size(cls, block_size: int | None) -> bool:
     def supports_head_size(cls, head_size: int) -> bool:
         return head_size >= 32
 
+    @classmethod
+    def supports_head_size_v(cls, head_size: int, head_size_v: int) -> bool:
+        # The cache is split with `split(head_size)` over a 2*head_size page, so
+        # V can be narrower than K but not wider.
+        return head_size_v <= head_size
+
     @classmethod
     def supports_mm_prefix(cls) -> bool:
         # Not implemented
diff --git a/vllm/v1/attention/backends/rocm_attn.py b/vllm/v1/attention/backends/rocm_attn.py
index 3cd25f2d58f0..2a0cb47f3507 100644
--- a/vllm/v1/attention/backends/rocm_attn.py
+++ b/vllm/v1/attention/backends/rocm_attn.py
@@ -154,9 +154,8 @@ def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleO
         # ROCM paged attention native C++ kernel only supports block sizes 16 and 32
         # due to shared memory (LDS) constraints on AMD GPUs.
         # The (head_size, block_size) pairs that have a native kernel are enumerated
-        # by csrc/rocm/attention_paged_attn_<head>_<dtype>_<kv>.cu, each of which
-        # instantiates ATTN_LAUNCH from csrc/rocm/attention_kernels.cuh; the runtime
-        # gate is use_rocm_custom_paged_attention() in vllm/platforms/rocm.py.
+        # by BLOCK_SIZES and VARIANTS in csrc/rocm/generate_paged_attn.py; the
+        # runtime gate is use_rocm_custom_paged_attention() in vllm/platforms/rocm.py.
         # However, vLLM allows support for any multiple of 16 via the Triton path.
         # As addressed in PR: https://github.com/vllm-project/vllm/pull/31380,
         # non-standard models (like qwen3-next with block_size 544, or qwen3_5
diff --git a/vllm/v1/attention/backends/triton_attn.py b/vllm/v1/attention/backends/triton_attn.py
index 232943dc38ba..ea6870ff2e03 100644
--- a/vllm/v1/attention/backends/triton_attn.py
+++ b/vllm/v1/attention/backends/triton_attn.py
@@ -456,6 +456,9 @@ def to_f32_units(elements: int) -> int:
         k_scale_off_f32 = base_off_f32 + to_f32_units(hs)
         v_scale_off_f32 = base_off_f32 + to_f32_units(padded_hs + hs)
 
+        # Not initialized: each scale is written with its token before any
+        # unmasked read, and an extensible KV cache backs only a prefix of
+        # num_blocks during warmup.
         # K scales (first content half)
         self._k_scale_cache = torch.as_strided(
             base_f32,
@@ -463,7 +466,6 @@ def to_f32_units(elements: int) -> int:
             stride=(block_f32, slot_f32, head_f32),
             storage_offset=k_scale_off_f32,
         )
-        self._k_scale_cache.fill_(1.0)
 
         # V scales (second content half)
         self._v_scale_cache = torch.as_strided(
@@ -472,7 +474,6 @@ def to_f32_units(elements: int) -> int:
             stride=(block_f32, slot_f32, head_f32),
             storage_offset=v_scale_off_f32,
         )
-        self._v_scale_cache.fill_(1.0)
 
     def fused_output_quant_supported(self, quant_key: QuantKey):
         return quant_key == kFp8StaticTensorSym
diff --git a/vllm/v1/attention/backends/triton_attn_diffkv.py b/vllm/v1/attention/backends/triton_attn_diffkv.py
index 2fff3226d67d..c4c9d083be12 100644
--- a/vllm/v1/attention/backends/triton_attn_diffkv.py
+++ b/vllm/v1/attention/backends/triton_attn_diffkv.py
@@ -69,6 +69,10 @@ class TritonAttentionDiffKVBackend(TritonAttentionBackend):
     # V head dim — set per layer via ``set_head_size_v`` before instantiation.
     head_size_v: int = 128
 
+    @classmethod
+    def supports_head_size_v(cls, head_size: int, head_size_v: int) -> bool:
+        return True
+
     supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
         "auto",
         "bfloat16",
diff --git a/vllm/v1/attention/backends/turboquant_attn.py b/vllm/v1/attention/backends/turboquant_attn.py
index 15633316b12e..afd08b1d3df6 100644
--- a/vllm/v1/attention/backends/turboquant_attn.py
+++ b/vllm/v1/attention/backends/turboquant_attn.py
@@ -346,7 +346,6 @@ def __init__(
             else head_size
         )
         self._val_data_bytes = math.ceil(head_size * cfg.effective_value_quant_bits / 8)
-        self._n_centroids = cfg.n_centroids if not cfg.key_fp8 else 1
 
         # Detect flash-attn version (FA2/3/4) for prefill paths.
         self.fa_version = get_flash_attn_version(head_size=head_size)
diff --git a/vllm/v1/attention/backends/utils.py b/vllm/v1/attention/backends/utils.py
index 40b283a3f1c5..5853b10094c8 100644
--- a/vllm/v1/attention/backends/utils.py
+++ b/vllm/v1/attention/backends/utils.py
@@ -28,6 +28,7 @@
     KVCacheLayout,
     KVCacheSpec,
     MambaSpec,
+    get_kv_quant_mode,
 )
 
 if TYPE_CHECKING:
@@ -251,12 +252,21 @@ def get_supported_kv_cache_layouts(
     their order; otherwise the layout the most backends put first wins, ties
     keeping the enum order. An empty intersection is a hard error.
     """
-    supported_layouts_lists: list[Sequence[KVCacheLayout]] = [
-        layouts
-        for backend in backends
-        if (layouts := backend.supported_kv_cache_layouts()) is not None
-    ] or [_DEFAULT_LAYOUT_PREFERENCE]
+    return _merge_supported_layouts(
+        [
+            layouts
+            for backend in backends
+            if (layouts := backend.supported_kv_cache_layouts()) is not None
+        ]
+        or [_DEFAULT_LAYOUT_PREFERENCE]
+    )
+
 
+def _merge_supported_layouts(
+    supported_layouts_lists: Sequence[Sequence[KVCacheLayout]],
+) -> list[KVCacheLayout]:
+    """Layouts in every supported set, merged as in
+    ``get_supported_kv_cache_layouts``."""
     first = supported_layouts_lists[0]
     if all(layouts == first for layouts in supported_layouts_lists[1:]):
         return list(first)
@@ -298,8 +308,10 @@ def resolve_kv_cache_layout(
     """Resolve one KV cache layout for the whole model.
 
     Runs once in the engine core. Every worker reports the layouts its backends
-    support, most preferred first (``get_supported_kv_cache_layouts``); all
-    ranks run the same backends, so their lists must agree. Specs mixing HNC
+    support, most preferred first (``get_supported_kv_cache_layouts``). Pipeline
+    stages hold different layers, so their lists may differ; the candidates are
+    the layouts every worker supports, and the layout the most workers put first
+    wins. Workers whose backends declare nothing do not vote. Specs mixing HNC
     shapes narrow the candidates to block-compact layouts. An explicit
     ``VLLM_KV_CACHE_LAYOUT`` must be one of the candidates or resolution fails,
     with the legacy ``NHD``/``HND`` names as aliases for ``LBNHC``/``LBHNC``; the
@@ -316,10 +328,15 @@ def resolve_kv_cache_layout(
     assert supported_layouts and all(supported_layouts), (
         "No worker reported supported KV cache layouts."
     )
-    assert all(names == supported_layouts[0] for names in supported_layouts[1:]), (
-        f"Workers disagree on supported KV cache layouts: {supported_layouts}."
+    reported = [
+        tuple(_layout_from_name(name) for name in names) for names in supported_layouts
+    ]
+    # Workers where no backend declares a set report the default preference;
+    # like undeclared backends of one worker, they don't vote.
+    candidates = _merge_supported_layouts(
+        [layouts for layouts in reported if layouts != _DEFAULT_LAYOUT_PREFERENCE]
+        or reported
     )
-    candidates = [_layout_from_name(name) for name in supported_layouts[0]]
 
     # A block-compact layout means the block is densely packed in memory, so any mix of
     # specs can re-interpret HNC with different sizes as long as the total number of
@@ -359,6 +376,19 @@ def resolve_kv_cache_layout(
                 f"a block-outer KV cache layout; supported sets: {supported_layouts}."
             )
 
+    # An extensible KV cache commits one contiguous run per layer only under a
+    # block-compact layout, which KV connectors need to register its memory.
+    if (
+        cache_config.enable_extensible_kv_cache
+        and vllm_config.kv_transfer_config is not None
+    ):
+        candidates = [m for m in candidates if m.is_block_compact]
+        if not candidates:
+            raise ValueError(
+                "KV connectors with the extensible KV cache need a block-compact "
+                f"layout, but none is in every supported set: {supported_layouts}."
+            )
+
     if (requested := envs.VLLM_KV_CACHE_LAYOUT) is not None:
         layout = _layout_from_name(requested)
         if layout not in candidates:
@@ -431,6 +461,43 @@ def get_per_layer_parameters(
     return per_layer_params
 
 
+def _kv_cache_read_encoding(cache_dtype: str | None) -> str | None:
+    # NVFP4 scale-search variants only differ in the store kernel.
+    if cache_dtype == "fp8":
+        return "fp8_e4m3"
+    if cache_dtype is not None and get_kv_quant_mode(cache_dtype).is_nvfp4:
+        return "nvfp4"
+    return cache_dtype
+
+
+def get_kv_cache_dtype_from_layers(
+    vllm_config: VllmConfig, layer_names: list[str]
+) -> str | None:
+    """Resolve the KV read encoding shared by an attention group.
+
+    Attention groups are split by spec, which separates quantization modes but
+    not encodings that share a storage layout (e.g. fp8_e4m3 vs fp8_e5m2).
+    Those must not be mixed within one group.
+    """
+    layers = get_layers_from_vllm_config(
+        vllm_config,
+        AttentionLayerBase,  # type: ignore[type-abstract]
+        layer_names,
+    )
+    cache_dtypes = {
+        _kv_cache_read_encoding(getattr(layer, "kv_cache_dtype", None))
+        for layer in layers.values()
+    }
+    if not cache_dtypes:
+        return None
+    if len(cache_dtypes) != 1:
+        raise ValueError(
+            "All layers in one attention group must share the same KV cache "
+            f"dtype; got {cache_dtypes} for {layer_names}."
+        )
+    return cache_dtypes.pop()
+
+
 def get_num_attention_heads_from_layers(
     vllm_config: VllmConfig, layer_names: list[str]
 ) -> int | None:
diff --git a/vllm/v1/attention/ops/flydsl_kernels/__init__.py b/vllm/v1/attention/ops/flydsl_kernels/__init__.py
index 208f01a7cb5e..74addfc33e6d 100644
--- a/vllm/v1/attention/ops/flydsl_kernels/__init__.py
+++ b/vllm/v1/attention/ops/flydsl_kernels/__init__.py
@@ -1,2 +1,20 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+
+def _run_compiled(launcher, *args) -> None:
+    """Reuse the compiled callable of a factory-style host launcher."""
+    import flydsl.compiler as flyc
+    from flydsl.utils import env
+
+    if env.compile.compile_only:
+        flyc.compile(launcher, *args)
+        return
+
+    compiled = getattr(launcher, "_cf", None)
+    if compiled is not None:
+        compiled(*args)
+        return
+
+    # flyc.compile executes once; do not launch its result again on this call.
+    launcher._cf = flyc.compile(launcher, *args)
diff --git a/vllm/v1/attention/ops/flydsl_turboquant_decode.py b/vllm/v1/attention/ops/flydsl_turboquant_decode.py
index e7e9ab3b9944..25523091bd55 100644
--- a/vllm/v1/attention/ops/flydsl_turboquant_decode.py
+++ b/vllm/v1/attention/ops/flydsl_turboquant_decode.py
@@ -29,6 +29,7 @@
 
 from vllm.logger import init_logger
 from vllm.triton_utils import tl, triton
+from vllm.v1.attention.ops.flydsl_kernels import _run_compiled
 
 logger = init_logger(__name__)
 
@@ -698,7 +699,8 @@ def flydsl_turboquant_decode_attention(
             num_partitions * tile_groups_per_partition * kv_compute_block,
             worst_case_max_seq_len,
         )
-    launch(
+    _run_compiled(
+        launch,
         segm_out,
         segm_sum,
         segm_max,
diff --git a/vllm/v1/attention/ops/flydsl_ultraquant_decode.py b/vllm/v1/attention/ops/flydsl_ultraquant_decode.py
index d2da66ffc8c4..022fff8a1243 100644
--- a/vllm/v1/attention/ops/flydsl_ultraquant_decode.py
+++ b/vllm/v1/attention/ops/flydsl_ultraquant_decode.py
@@ -15,6 +15,7 @@
 
 from vllm.logger import init_logger
 from vllm.triton_utils import triton
+from vllm.v1.attention.ops.flydsl_kernels import _run_compiled
 
 logger = init_logger(__name__)
 
@@ -773,7 +774,8 @@ def flydsl_ultraquant_decode_attention(
             num_partitions * tile_groups_per_partition * kv_compute_block,
             worst_case_max_seq_len,
         )
-    launch(
+    _run_compiled(
+        launch,
         segm_out,
         segm_sum,
         segm_max,
diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py
index 0db22e03c0bf..93f509240849 100644
--- a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py
+++ b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py
@@ -23,13 +23,23 @@
 from vllm.v1.worker.workspace import current_workspace_manager
 
 if current_platform.is_rocm():
-    from vllm.platforms.rocm import _ON_GFX942, _ON_GFX950, _ON_GFX1151, _ON_RDNA
+    from vllm.platforms.rocm import (
+        _ON_GFX942,
+        _ON_GFX950,
+        _ON_GFX1151,
+        _ON_RDNA,
+        get_cdna_version,
+    )
 else:
     _ON_GFX942 = False
     _ON_GFX950 = False
     _ON_GFX1151 = False
     _ON_RDNA = False
 
+    def get_cdna_version() -> int:
+        return 0
+
+
 logger = init_logger(__name__)
 
 FP8_DTYPE = current_platform.fp8_dtype()
@@ -52,6 +62,24 @@ def _get_aiter_sparse_prefill_opus() -> Callable[..., torch.Tensor] | None:
 # 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
+# gfx942 gate against the in-tree ragged prefill, not the gfx950 OPUS
+# crossover. MI325, H=16, D=512, bf16: from 512 rows AITER is ahead at
+# top-k 128, 512, and 2048. At 256 rows and top-k 128 the ratio is 1.04,
+# on a flat ~0.025 ms.
+_GFX942_AITER_PA_PREFILL_SPARSE_MIN_QUERIES = 512
+
+
+def _can_use_aiter_pa_prefill_sparse(q: torch.Tensor, kv: torch.Tensor) -> bool:
+    """Cheap gate. The AITER import runs only after this passes."""
+    from vllm._aiter_ops import rocm_aiter_ops
+
+    return (
+        get_cdna_version() == 3
+        and rocm_aiter_ops.is_mla_enabled()
+        and q.shape[0] >= _GFX942_AITER_PA_PREFILL_SPARSE_MIN_QUERIES
+        and q.dtype in (torch.bfloat16, torch.float16)
+        and kv.dtype == q.dtype
+    )
 
 
 def _indexer_k_is_c4a_block_flat(compress_ratio: int) -> bool:
@@ -3843,6 +3871,39 @@ def _rocm_sparse_attn_prefill_ragged_aiter_opus(
     return True
 
 
+def _rocm_sparse_attn_prefill_aiter(
+    pa_prefill_sparse: Callable[..., torch.Tensor],
+    q: torch.Tensor,
+    kv: torch.Tensor,
+    indices: torch.Tensor,
+    indptr: torch.Tensor,
+    scale: float,
+    attn_sink: torch.Tensor | None,
+    output: torch.Tensor,
+) -> None:
+    """gfx942 path. One KV pool; the extend sources stay unset."""
+    logger.info_once("Using AITER pa_prefill_sparse for sparse MLA prefill on gfx942")
+    indices = _as_int32_contiguous_1d(indices)
+    indptr = _as_int32_contiguous_1d(indptr)
+    inplace = (
+        output.shape == q.shape and output.dtype == q.dtype and output.is_contiguous()
+    )
+    written = pa_prefill_sparse(
+        q,
+        kv,
+        indices,
+        indptr,
+        None,
+        None,
+        None,
+        attn_sink,
+        float(scale),
+        out=output if inplace else None,
+    )
+    if written.data_ptr() != output.data_ptr():
+        output.copy_(written[..., : output.shape[-1]])
+
+
 @functools.lru_cache
 def _decode_cu_count() -> int:
     try:
@@ -4608,6 +4669,31 @@ def rocm_sparse_attn_prefill(
         ):
             return
 
+    if _can_use_aiter_pa_prefill_sparse(q, kv):
+        from aiter.ops.triton.attention.pa_prefill_sparse import pa_prefill_sparse
+
+        if ragged_indices is None or ragged_indptr is None:
+            assert indices is not None
+            indices_2d = indices.reshape(indices.shape[0], -1)
+            ragged_indices, ragged_indptr = build_ragged_indices_from_dense(
+                indices_2d,
+                topk_length
+                if topk_length is not None
+                else (indices_2d >= 0).sum(dim=-1, dtype=torch.int32),
+                num_rows=kv.shape[0],
+            )
+        _rocm_sparse_attn_prefill_aiter(
+            pa_prefill_sparse,
+            q=q,
+            kv=kv.squeeze(1),
+            indices=ragged_indices,
+            indptr=ragged_indptr,
+            scale=scale,
+            attn_sink=opus_attn_sink,
+            output=output,
+        )
+        return
+
     if ragged_indices is not None and ragged_indptr is not None:
         _rocm_sparse_attn_prefill_ragged_triton(
             q=q,
diff --git a/vllm/v1/core/kv_cache_manager.py b/vllm/v1/core/kv_cache_manager.py
index cf27161cd472..fd0781de14aa 100644
--- a/vllm/v1/core/kv_cache_manager.py
+++ b/vllm/v1/core/kv_cache_manager.py
@@ -696,11 +696,6 @@ def reset_prefix_cache(self) -> bool:
             self.prefix_cache_stats.reset = True
         return True
 
-    def get_num_common_prefix_blocks(self, running_request_id: str) -> list[int]:
-        """Deprecated: cascade attention has been removed. Always returns 0
-        for each KV cache group. Kept for out-of-tree schedulers."""
-        return [0] * self.num_kv_cache_groups
-
     def take_events(self) -> list[KVCacheEvent]:
         """Take the KV cache events from the block pool.
 
diff --git a/vllm/v1/core/kv_cache_planning.py b/vllm/v1/core/kv_cache_planning.py
index 279a40c0c036..ecb1cd44dd14 100644
--- a/vllm/v1/core/kv_cache_planning.py
+++ b/vllm/v1/core/kv_cache_planning.py
@@ -7,15 +7,19 @@
 from collections.abc import Callable, Iterable, Sequence
 from dataclasses import replace
 from functools import partial, reduce
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 from vllm.config import SpeculativeConfig, VllmConfig
 from vllm.logger import init_logger
+from vllm.utils.extensible_tensor import granule_aligned_blocks
 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.core.kv_cache_config_builder import KVCacheConfigBuilder
-from vllm.v1.core.kv_cache_utils import _kv_cache_device_name
+from vllm.v1.core.kv_cache_utils import (
+    _kv_cache_device_name,
+    update_kv_cache_capacity,
+)
 from vllm.v1.hisparse.layout import (
     get_hisparse_host_pool_bytes,
     get_hisparse_kv_cache_config,
@@ -44,21 +48,38 @@
 )
 from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry
 
+if TYPE_CHECKING:
+    from vllm.v1.worker.worker_base import CompilationTimes
+
+
 logger = init_logger(__name__)
 
 
+def _more_kv_cache_memory_hint(vllm_config: VllmConfig) -> str:
+    """How the user can make room for the KV cache, for error messages."""
+    if vllm_config.cache_config.enable_extensible_kv_cache:
+        return (
+            "The extensible KV cache already takes all device memory left after "
+            "the model and its activations; decrease `max_model_len` or "
+            "`max_num_seqs`, or the model needs a device with more memory."
+        )
+    return (
+        "Try increasing `gpu_memory_utilization` when initializing the engine "
+        "(this flag also controls CPU memory reservation on the CPU backend, "
+        "despite its name)."
+    )
+
+
 def _check_enough_kv_cache_memory(
     available_memory: int,
     get_needed_memory: Callable[[], int],
     max_model_len: int,
     estimate_max_model_len: Callable[[int], int],
+    hint: str,
 ):
     if available_memory <= 0:
         raise ValueError(
-            "No available memory for the cache blocks. "
-            "Try increasing `gpu_memory_utilization` when initializing the engine "
-            "(this flag also controls CPU memory reservation on the CPU "
-            "backend, despite its name). "
+            f"No available memory for the cache blocks. {hint} "
             "See https://docs.vllm.ai/en/latest/configuration/conserving_memory/ "
             "for more details."
         )
@@ -79,9 +100,7 @@ def _check_enough_kv_cache_memory(
             f"({max_model_len}), ({format_gib(needed_memory)} GiB KV "
             f"cache is needed, which is larger than the available KV cache "
             f"memory ({format_gib(available_memory)} GiB). {estimated_msg}"
-            f"Try increasing `gpu_memory_utilization` (which also controls "
-            f"CPU memory on the CPU backend) or decreasing `max_model_len` "
-            f"when initializing the engine. "
+            f"{hint} Decreasing `max_model_len` also reduces what is needed. "
             f"See https://docs.vllm.ai/en/latest/configuration/conserving_memory/ "
             f"for more details."
         )
@@ -181,6 +200,7 @@ def check_enough_kv_cache_memory(
             lambda: max_memory_usage_bytes(vllm_config, kv_cache_spec.values()),
             vllm_config.model_config.max_model_len,
             lambda am: estimate_max_model_len(vllm_config, kv_cache_spec, am),
+            _more_kv_cache_memory_hint(vllm_config),
         )
 
 
@@ -1376,7 +1396,7 @@ def _auto_fit_max_model_len(
     if auto_fit_max <= 0:
         raise ValueError(
             "Cannot auto-fit max_model_len: not enough GPU memory available "
-            "to serve even a single token. Try increasing `gpu_memory_utilization`."
+            f"to serve even a single token. {_more_kv_cache_memory_hint(vllm_config)}"
         )
 
     kv_cache_memory = (
@@ -1552,6 +1572,7 @@ def get_kv_cache_configs(
             partial(_max_memory_usage_bytes_from_groups, builder, vllm_config, groups),
             vllm_config.model_config.max_model_len,
             partial(_estimate_max_model_len_from_groups, builder, vllm_config, groups),
+            _more_kv_cache_memory_hint(vllm_config),
         )
 
     kv_cache_configs: list[KVCacheConfig] = []
@@ -1712,3 +1733,122 @@ def get_kv_cache_config_from_groups(
         return get_kv_cache_config_from_groups(
             vllm_config, kv_cache_groups, num_blocks * bytes_per_block
         )
+
+
+def shrink_kv_cache_configs(
+    vllm_config: VllmConfig, kv_cache_configs: list[KVCacheConfig], num_blocks: int
+) -> None:
+    """Lower num_blocks of every config in place after measured sizing. The
+    tensors' size follows (their strides and offsets keep describing the
+    reservation). Raises if max_model_len no longer fits.
+    """
+    from vllm.v1.core.kv_cache_config_builder import get_kv_cache_config_builder
+
+    builder = get_kv_cache_config_builder(vllm_config)
+    for kv_cache_config in kv_cache_configs:
+        assert num_blocks <= kv_cache_config.num_blocks
+        kv_cache_config.kv_cache_tensors = [
+            replace(tensor, size=num_blocks * tensor.size // kv_cache_config.num_blocks)
+            for tensor in kv_cache_config.kv_cache_tensors
+        ]
+        groups = kv_cache_config.kv_cache_groups
+        if groups:
+            # The null block is held back by the BlockPool, as in the capacity
+            # check of `get_kv_cache_configs`.
+            _check_enough_kv_cache_memory(
+                (num_blocks - 1)
+                * _pool_bytes_per_block_from_builder(builder, vllm_config, groups),
+                partial(
+                    _max_memory_usage_bytes_from_groups, builder, vllm_config, groups
+                ),
+                vllm_config.model_config.max_model_len,
+                partial(
+                    _estimate_max_model_len_from_groups, builder, vllm_config, groups
+                ),
+                _more_kv_cache_memory_hint(vllm_config),
+            )
+        kv_cache_config.num_blocks = num_blocks
+
+
+def granule_aligned_kv_cache_blocks(
+    num_blocks: int, kv_cache_configs: list[KVCacheConfig], commit_granule: int
+) -> int:
+    """Round num_blocks down so every worker's KV cache segments (one per
+    layer or one overall, num_blocks blocks of one tensor's block stride)
+    span whole commit granules. A KV connector registers each committed
+    segment for RDMA, which cannot span several physical chunks."""
+    strides = {
+        tensor.block_stride
+        for kv_cache_config in kv_cache_configs
+        for tensor in kv_cache_config.kv_cache_tensors
+    }
+    return granule_aligned_blocks(num_blocks, strides, commit_granule)
+
+
+def align_extensible_kv_cache_capacity(
+    vllm_config: VllmConfig,
+    kv_cache_configs: list[KVCacheConfig],
+    scheduler_kv_cache_config: KVCacheConfig,
+    commit_granule: int,
+) -> None:
+    """Lower the reserved capacity to a granule-aligned block count, so that
+    segment offsets (multiples of the capacity) sit on granule boundaries."""
+    num_blocks = granule_aligned_kv_cache_blocks(
+        scheduler_kv_cache_config.num_blocks, kv_cache_configs, commit_granule
+    )
+    if num_blocks >= scheduler_kv_cache_config.num_blocks:
+        return
+    for kv_cache_config in kv_cache_configs + [scheduler_kv_cache_config]:
+        old_num_blocks = kv_cache_config.num_blocks
+        shrink_kv_cache_configs(vllm_config, [kv_cache_config], num_blocks)
+        # Unlike the post-warmup shrink, this happens before allocation: the
+        # placements must describe the smaller reservation. Layer-outermost
+        # tensors place layers `num_blocks` blocks apart.
+        kv_cache_config.kv_cache_tensors = [
+            replace(
+                tensor,
+                layer_stride=num_blocks * tensor.block_stride,
+                offset=tensor.offset // old_num_blocks * num_blocks,
+            )
+            if tensor.layer_stride == old_num_blocks * tensor.block_stride
+            else tensor
+            for tensor in kv_cache_config.kv_cache_tensors
+        ]
+    vllm_config.cache_config.num_gpu_blocks = num_blocks
+
+
+def finalize_extensible_kv_cache(
+    vllm_config: VllmConfig,
+    kv_cache_configs: list[KVCacheConfig],
+    scheduler_kv_cache_config: KVCacheConfig,
+    compilation_times: list["CompilationTimes"],
+    commit_granule: int | None = None,
+) -> int:
+    """Settle the KV cache size from the memory workers measured after warmup.
+
+    Every rank must agree on one block count, so the minimum becomes the
+    scheduler's num_blocks; the caller commits it on the workers. With
+    commit_granule (set when a KV connector is configured) the count is
+    also granule-aligned, see `granule_aligned_kv_cache_blocks`.
+    """
+    # No measurements: the elastic-EP scale-up path skips warmup. Elastic EP
+    # leaves the extensible KV cache off, so this branch is not reached; it
+    # keeps the capacity.
+    block_counts = (
+        times.num_kv_blocks
+        for times in compilation_times
+        if times.num_kv_blocks is not None
+    )
+    num_blocks = min(block_counts, default=scheduler_kv_cache_config.num_blocks)
+    num_blocks = min(num_blocks, scheduler_kv_cache_config.num_blocks)
+    if commit_granule is not None:
+        num_blocks = granule_aligned_kv_cache_blocks(
+            num_blocks, kv_cache_configs, commit_granule
+        )
+    shrink_kv_cache_configs(
+        vllm_config, kv_cache_configs + [scheduler_kv_cache_config], num_blocks
+    )
+    vllm_config.cache_config.num_gpu_blocks = num_blocks
+    if scheduler_kv_cache_config.kv_cache_groups:
+        update_kv_cache_capacity(vllm_config, scheduler_kv_cache_config)
+    return num_blocks
diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py
index d24674e48318..9f239e82fb87 100644
--- a/vllm/v1/core/kv_cache_utils.py
+++ b/vllm/v1/core/kv_cache_utils.py
@@ -15,6 +15,7 @@
     NewType,
     TypeAlias,
     TypeVar,
+    cast,
     overload,
 )
 
@@ -39,6 +40,7 @@
 if TYPE_CHECKING:
     from vllm.v1.core.block_pool import BlockPool
     from vllm.v1.core.single_type_kv_cache_manager import SingleTypeKVCacheManager
+    from vllm.v1.engine import EngineCoreRequest, PrefixHashConfig
 
 
 # BlockHash represents the hash of a single KV-cache block used for
@@ -1036,6 +1038,28 @@ def request_block_hasher(request: Request) -> list[BlockHash]:
     return request_block_hasher
 
 
+def frontend_block_hashes(
+    request: "EngineCoreRequest", prefix_hash_config: "PrefixHashConfig | None"
+) -> list[BlockHash] | None:
+    """Return the frontend's prompt block hashes, unless their count doesn't
+    match the prompt, in which case the engine hashes the prompt itself."""
+    block_hashes = request.block_hashes
+    if block_hashes is None:
+        return None
+    prompt_len = len(request.prompt_token_ids) if request.prompt_token_ids else 0
+    # The engine hashes a block once its lookahead tokens exist too.
+    if prefix_hash_config is None or len(block_hashes) != (
+        max(prompt_len - prefix_hash_config.num_lookahead_tokens, 0)
+        // prefix_hash_config.hash_block_size
+    ):
+        logger.warning_once(
+            "Ignoring frontend-computed prompt block hashes that don't match the "
+            "engine's prefix hash config; hashing prompts in the engine instead."
+        )
+        return None
+    return cast(list[BlockHash], block_hashes)
+
+
 def generate_lookahead_block_hash_extra_keys(
     request: Request,
     start_token_idx: int,
diff --git a/vllm/v1/core/sched/async_scheduler.py b/vllm/v1/core/sched/async_scheduler.py
index 42906e2f19e5..e7b87592b458 100644
--- a/vllm/v1/core/sched/async_scheduler.py
+++ b/vllm/v1/core/sched/async_scheduler.py
@@ -2,22 +2,16 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import vllm.envs as envs
-from vllm.logger import init_logger
 from vllm.platforms import current_platform
 from vllm.v1.core.sched.output import SchedulerOutput
 from vllm.v1.core.sched.scheduler import Scheduler
 
-logger = init_logger(__name__)
-
 
 class AsyncScheduler(Scheduler):
     def __init__(self, *args, **kwargs) -> None:
         super().__init__(*args, **kwargs)
-        # reusable read-only placeholder list for speculative decoding.
-        self._spec_token_placeholders: list[int] = [-1] * self.num_spec_tokens
-        self.pp_size = self.parallel_config.pipeline_parallel_size
         # Only XPU can disable PP decode microbatching via VLLM_XPU_PP_MICROBATCH.
-        self.decode_stagger = self.pp_size
+        self.decode_stagger = self.parallel_config.pipeline_parallel_size
         if current_platform.is_xpu() and not envs.VLLM_XPU_PP_MICROBATCH:
             self.decode_stagger = 1
 
@@ -25,9 +19,7 @@ def _update_after_schedule(self, scheduler_output: SchedulerOutput) -> None:
         super()._update_after_schedule(scheduler_output)
         spec_decode_tokens = scheduler_output.scheduled_spec_decode_tokens
         # Use the latest num of scheduled draft tokens in next step as placeholder.
-        self._spec_token_placeholders = [
-            -1
-        ] * scheduler_output.num_spec_tokens_to_schedule
+        spec_token_placeholders = [-1] * scheduler_output.num_spec_tokens_to_schedule
         for req_id in scheduler_output.num_scheduled_tokens:
             request = self.requests[req_id]
             if request.is_prefill_chunk:
@@ -46,7 +38,7 @@ def _update_after_schedule(self, scheduler_output: SchedulerOutput) -> None:
             )
             # Add placeholders for the new draft/spec tokens.
             # We will update the actual spec token ids in the worker process.
-            request.spec_token_ids = self._spec_token_placeholders
+            request.spec_token_ids = spec_token_placeholders
 
             if self.use_v2_model_runner:
                 # Set the next step index in which this request is eligible to be
diff --git a/vllm/v1/core/sched/output.py b/vllm/v1/core/sched/output.py
index 7482f32d4fcb..bd2b93b93a48 100644
--- a/vllm/v1/core/sched/output.py
+++ b/vllm/v1/core/sched/output.py
@@ -315,10 +315,6 @@ class SchedulerOutput:
     # Number of spec tokens to schedule for the next step.
     num_spec_tokens_to_schedule: int = 0
 
-    # Deprecated: cascade attention has been removed, so this is always empty.
-    # Kept for out-of-tree schedulers and workers.
-    num_common_prefix_blocks: list[int] = field(default_factory=list)
-
     # Whether no scheduled request can use the drafts proposed after this step
     # (non-final prefill chunks, or requests sure to reach max_tokens), so the
     # drafter can skip proposing them.
@@ -344,8 +340,3 @@ class GrammarOutput:
     structured_output_request_ids: list[str]
     # Bitmask ordered as structured_output_request_ids.
     grammar_bitmask: "npt.NDArray[np.int32]"
-    # Per request, ordered as structured_output_request_ids: how many leading
-    # drafts the bitmask constrained. `grammar_bitmask` fills every row after
-    # the first -1 placeholder with the all-permissive `_full_mask`, so drafts
-    # from this index on must not be accepted. None invalidates every draft.
-    num_acceptable_drafts: list[int] | None = None
diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py
index cb8ee19ae3be..39d8a6bd5a49 100644
--- a/vllm/v1/core/sched/scheduler.py
+++ b/vllm/v1/core/sched/scheduler.py
@@ -77,7 +77,6 @@
 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 StructuredOutputManager
-from vllm.v1.structured_output.utils import strip_speculative_padding
 from vllm.v1.utils import record_function_or_nullcontext
 
 logger = init_logger(__name__)
@@ -167,6 +166,7 @@ def __init__(
         self.connector_prefix_cache_stats: PrefixCacheStats | None = None
         self.recompute_kv_load_failures = True
         self.defer_block_free = False
+        self.defer_async_loads = False
         # Whether a preempted request's in-flight output must be dropped; see
         # KVConnectorBase_V1.requires_kv_delivery.
         self.requires_kv_delivery = False
@@ -192,6 +192,12 @@ def __init__(
             multiple_inflight_batches = self.vllm_config.max_concurrent_batches > 1
             if multiple_inflight_batches and kv_transfer_config.is_kv_consumer:
                 self.defer_block_free = True
+            # HiSparse keeps reading a released page until the pool reuses it,
+            # so a load into a reused page must also wait for in-flight steps.
+            self.defer_async_loads = (
+                self.defer_block_free
+                and kv_cache_config.hisparse_host_num_blocks is not None
+            )
 
             self.requires_kv_delivery = self.connector.requires_kv_delivery
 
@@ -320,7 +326,7 @@ def __init__(
                     vllm_max_batch_size=self.scheduler_config.max_num_seqs,
                     vllm_num_speculative_tokens=self.num_spec_tokens,
                 )
-            self.use_eagle = speculative_config.use_eagle()
+            self.use_eagle = speculative_config.uses_target_hidden_states
             self.use_eagle_block_drop = speculative_config.use_eagle_block_drop()
             if (
                 speculative_config.needs_eagle_block_drop()
@@ -463,6 +469,10 @@ def __init__(
         # FIFO of (fence_seq, blocks): blocks become safe to free once
         # processed_step_seq >= fence_seq.
         self.deferred_frees: deque[tuple[int, list[KVCacheBlock]]] = deque()
+        # Request id -> (fence_seq, num_external_tokens, connector prefix cache
+        # queries, hits): async KV loads started once processed_step_seq >=
+        # fence_seq.
+        self.deferred_loads: dict[str, tuple[int, int, int, int]] = {}
 
         self.perf_metrics: ModelMetrics | None = None
         if self.log_stats and vllm_config.observability_config.enable_mfu_metrics:
@@ -1454,14 +1464,29 @@ def skip_request(from_queue: RequestQueue) -> None:
                 # This information is used to determine if a load is
                 # needed for this request.
                 if self.connector is not None:
-                    self.connector.update_state_after_alloc(
-                        request,
-                        self.kv_cache_manager.get_blocks(request_id),
-                        num_external_computed_tokens,
+                    defer_load = (
+                        load_kv_async
+                        and self.defer_async_loads
+                        and self.sched_step_seq > self.processed_step_seq
                     )
+                    if defer_load:
+                        request.last_sched_seq = self.sched_step_seq
+                        self.deferred_loads[request_id] = (
+                            self.sched_step_seq,
+                            num_external_computed_tokens,
+                            connector_prefix_cache_queries,
+                            connector_prefix_cache_hits,
+                        )
+                    else:
+                        self.connector.update_state_after_alloc(
+                            request,
+                            self.kv_cache_manager.get_blocks(request_id),
+                            num_external_computed_tokens,
+                        )
                     if (
                         self.connector_prefix_cache_stats is not None
                         and connector_prefix_cache_queries != 0
+                        and not defer_load
                     ):
                         preempted = request.num_preemptions > 0
                         self.connector_prefix_cache_stats.record(
@@ -2222,14 +2247,7 @@ def get_grammar_bitmask(
             structured_output_request_ids,
             scheduler_output.scheduled_spec_decode_tokens,
         )
-        spec_tokens = scheduler_output.scheduled_spec_decode_tokens
-        num_acceptable_drafts = [
-            len(strip_speculative_padding(spec_tokens.get(req_id, [])))
-            for req_id in structured_output_request_ids
-        ]
-        return GrammarOutput(
-            structured_output_request_ids, bitmask, num_acceptable_drafts
-        )
+        return GrammarOutput(structured_output_request_ids, bitmask)
 
     def update_from_output(
         self,
@@ -2254,6 +2272,7 @@ def update_from_output(
         if self.defer_block_free and scheduler_output.total_num_scheduled_tokens > 0:
             self.processed_step_seq += 1
             self._drain_deferred_frees()
+            self._drain_deferred_loads()
 
         perf_stats: PerfStats | None = None
         if self.perf_metrics and self.perf_metrics.is_enabled():
@@ -2850,7 +2869,7 @@ def update_draft_token_ids(self, draft_token_ids: DraftTokenIds) -> None:
                 continue
 
             # Add newly generated spec token ids to the request.
-            request.spec_token_ids = self.structured_output_manager.validate_tokens(
+            request.spec_token_ids = self.structured_output_manager.trim_spec_tokens(
                 request, spec_token_ids
             )
 
@@ -2877,10 +2896,7 @@ def update_draft_token_ids_in_output(
             # Trim drafts to scheduled number of spec tokens
             # (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.
-            spec_token_ids = self.structured_output_manager.validate_tokens(
-                request, spec_token_ids
-            )
+            # Preserve draft slots; rejection sampling handles grammar violations.
             # Pad to original number of spec tokens.
             num_invalid_tokens = orig_num_spec_tokens - len(spec_token_ids)
             if num_invalid_tokens:
@@ -2989,8 +3005,10 @@ def finish_requests(
         for request in valid_requests:
             delay_free_blocks = False
             if request.status == RequestStatus.WAITING_FOR_REMOTE_KVS:
+                # A deferred load never started, so it will not report done.
                 delay_free_blocks = (
-                    request.request_id not in self.finished_recving_kv_req_ids
+                    self.deferred_loads.pop(request.request_id, None) is None
+                    and request.request_id not in self.finished_recving_kv_req_ids
                 )
                 self.finished_recving_kv_req_ids.discard(request.request_id)
                 self.failed_recving_kv_req_ids.discard(request.request_id)
@@ -3090,6 +3108,32 @@ def _drain_deferred_frees(self):
             # Free in reverse order so that the tail blocks are evicted first.
             self.kv_cache_manager.block_pool.free_blocks(reversed(blocks))
 
+    def _drain_deferred_loads(self) -> None:
+        """Start deferred async loads whose fence step has completed."""
+        for request_id, (fence, num_tokens, queries, hits) in list(
+            self.deferred_loads.items()
+        ):
+            if fence > self.processed_step_seq:
+                break
+            del self.deferred_loads[request_id]
+            assert self.connector is not None
+            request = self.requests[request_id]
+            self.connector.update_state_after_alloc(
+                request, self.kv_cache_manager.get_blocks(request_id), num_tokens
+            )
+            if self.connector_prefix_cache_stats is not None and queries != 0:
+                preempted = request.num_preemptions > 0
+                self.connector_prefix_cache_stats.record(
+                    num_tokens=queries,
+                    num_hits=hits,
+                    preempted=preempted,
+                    hits_by_source=(
+                        self.connector.get_external_cache_hit_sources(request, hits)
+                        if hits and not preempted
+                        else None
+                    ),
+                )
+
     def get_num_unfinished_requests(self) -> int:
         if self._pause_state == PauseState.PAUSED_ALL:
             return 0
diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py
index dd2ba5aa73e5..6bce870265be 100644
--- a/vllm/v1/core/single_type_kv_cache_manager.py
+++ b/vllm/v1/core/single_type_kv_cache_manager.py
@@ -2278,6 +2278,12 @@ def get_num_blocks_to_allocate(
         apply_admission_cap: bool = False,
         prefill_end: int = 0,
     ) -> int:
+        assert self.coordinator is not None
+        available = self.block_pool.get_num_free_blocks()
+        if apply_admission_cap and total_computed_tokens > num_local_computed_tokens:
+            num_tokens += self.coordinator.import_step_tokens
+            available -= self.coordinator.reserved_import_pages
+            self.coordinator.size_import(request_id, cdiv(num_tokens, self.block_size))
         required = super().get_num_blocks_to_allocate(
             request_id,
             num_tokens,
@@ -2288,8 +2294,7 @@ def get_num_blocks_to_allocate(
         )
         # Host blocks come from a private pool, so they never count against
         # the GPU pool. Use the same admission sentinel as Mamba instead.
-        if required > self.block_pool.get_num_free_blocks():
-            assert self.coordinator is not None
+        if required > available:
             assert self.coordinator.gpu_pool is not None
             return self.coordinator.gpu_pool.num_gpu_blocks + 1
         return 0
@@ -2315,6 +2320,22 @@ def allocate_external_computed_blocks(
             num_external_computed_tokens,
             record_for_zeroing,
         )
+        self.coordinator.update_import_reservation(
+            request_id, len(self.req_to_blocks[request_id])
+        )
+
+    def allocate_new_blocks(
+        self, request_id: str, num_tokens: int, num_tokens_main_model: int
+    ) -> list[KVCacheBlock]:
+        new_blocks = super().allocate_new_blocks(
+            request_id, num_tokens, num_tokens_main_model
+        )
+        if new_blocks:
+            assert self.coordinator is not None
+            self.coordinator.update_import_reservation(
+                request_id, len(self.req_to_blocks[request_id])
+            )
+        return new_blocks
 
     def cache_blocks(
         self,
diff --git a/vllm/v1/engine/__init__.py b/vllm/v1/engine/__init__.py
index 346126355412..a7420981f145 100644
--- a/vllm/v1/engine/__init__.py
+++ b/vllm/v1/engine/__init__.py
@@ -90,6 +90,22 @@ def __str__(self):
         return FINISH_REASON_STRINGS[self.value]
 
 
+@dataclass
+class PrefixHashConfig:
+    """What a frontend needs to compute a request's prompt block hashes itself
+    (see `EngineCoreRequest.block_hashes`)."""
+
+    # The --prefix-caching-hash-algo name, e.g. "sha256".
+    hash_algo: str
+    # Tokens per hashed block.
+    hash_block_size: int
+    # Parent hash of the first block.
+    none_hash: bytes
+    # Tokens after each block that its hash also covers (as extra keys), so a
+    # block is hashed only once they exist.
+    num_lookahead_tokens: int = 0
+
+
 @dataclass
 class EngineCoreReadyResponse:
     """Sent from EngineCore to each frontend at the end of engine startup.
@@ -132,6 +148,9 @@ class EngineCoreReadyResponse:
     # `/server_info` without holding the config themselves (the Rust frontend).
     # Sent only with VLLM_SERVER_DEV_MODE, the same gate as `/server_info`.
     vllm_config: dict[str, Any] | None = None
+    # Set when the engine hashes request blocks (prefix caching or a KV
+    # connector), so that a frontend can send prompt block hashes.
+    prefix_hash_config: PrefixHashConfig | None = None
 
 
 class EngineCoreRequest(
@@ -188,6 +207,10 @@ class EngineCoreRequest(
     session_id: str | None = None
     kv_hints: KvHintsEnvelope | None = None
 
+    # Prompt block hashes computed by the frontend with the engine's
+    # PrefixHashConfig, including any extra keys (LoRA, multimodal, cache salt).
+    block_hashes: list[bytes] | None = None
+
     @property
     def params(self) -> SamplingParams | PoolingParams:
         """Return the processed params (sampling or pooling)."""
@@ -357,3 +380,5 @@ class EngineStatusType(enum.IntEnum):
     HEALTHY = 0
     DEAD = 1
     UNHEALTHY = 2
+    # Faulted; querying the workers' all2all mask before going UNHEALTHY.
+    DIAGNOSING = 3
diff --git a/vllm/v1/engine/async_llm.py b/vllm/v1/engine/async_llm.py
index b3dbf62e824a..5d95df0fa872 100644
--- a/vllm/v1/engine/async_llm.py
+++ b/vllm/v1/engine/async_llm.py
@@ -28,6 +28,7 @@
     set_scaling_elastic_ep,
 )
 from vllm.exceptions import (
+    EngineFaultedError,
     GenerationError,
     GracefulHTTPError,
     MaxQueuedTokensError,
@@ -345,10 +346,23 @@ def check_admission(self, n: int = 1, request_id: str | None = None) -> None:
             request_id: Request id, used for logging only.
 
         Raises:
+            EngineFaultedError: If the engine has faulted and is awaiting
+                FT recovery.
             QueueOverflowError: If ``max_num_queued_reqs`` would be exceeded.
             MaxQueuedTokensError: If ``max_num_queued_tokens`` would be exceeded.
 
         """
+        if (
+            self.vllm_config.parallel_config.enable_fault_tolerance
+            and self.engine_core.engine_status.get("status") != "healthy"
+        ):
+            logger.info(
+                "Engine %s - rejecting request %s.",
+                self.engine_core.engine_status.get("status"),
+                request_id,
+            )
+            raise EngineFaultedError()
+
         max_num_reqs = self.scheduler_config.max_num_queued_reqs
         if max_num_reqs is not None:
             current_requests = (
diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py
index c0510396315b..05f9fda5c5f6 100644
--- a/vllm/v1/engine/core.py
+++ b/vllm/v1/engine/core.py
@@ -54,12 +54,16 @@
 from vllm.utils.system_utils import decorate_logs, set_process_title
 from vllm.utils.watch_dog import get_watch_dog, start_watch_dog
 from vllm.v1.attention.backends.utils import resolve_kv_cache_layout
+from vllm.v1.core import kv_cache_utils
 from vllm.v1.core.kv_cache_planning import (
+    align_extensible_kv_cache_capacity,
+    finalize_extensible_kv_cache,
     get_global_kv_cache_groups,
     get_kv_cache_configs,
 )
 from vllm.v1.core.kv_cache_utils import (
     BlockHash,
+    frontend_block_hashes,
     generate_scheduler_kv_cache_config,
     get_request_block_hasher,
     init_none_hash,
@@ -81,6 +85,7 @@
     EngineCoreRequestType,
     FinishReason,
     PauseMode,
+    PrefixHashConfig,
     ReconfigureDistributedRequest,
     ReconfigureRankType,
     UtilityOutput,
@@ -127,6 +132,7 @@ class EngineCore:
     # object.__new__.
     _checkpoint_state = CheckpointState.RUNNING
     _checkpoint_error: str | None = None
+    _kv_cache_reclaimable = False
 
     def __init__(
         self,
@@ -224,6 +230,8 @@ def __init__(
         self.is_pooling_model = vllm_config.model_config.runner_type == "pooling"
 
         self.request_block_hasher: Callable[[Request], list[BlockHash]] | None = None
+        # Lets frontends hash prompt blocks; see EngineCoreRequest.block_hashes.
+        self.prefix_hash_config: PrefixHashConfig | None = None
         if vllm_config.cache_config.enable_prefix_caching or kv_connector is not None:
             caching_hash_fn = get_hash_fn_by_name(
                 vllm_config.cache_config.prefix_caching_hash_algo
@@ -235,6 +243,12 @@ def __init__(
                 caching_hash_fn,
                 vllm_config.num_hash_lookahead_tokens,
             )
+            self.prefix_hash_config = PrefixHashConfig(
+                hash_algo=vllm_config.cache_config.prefix_caching_hash_algo,
+                hash_block_size=hash_block_size,
+                none_hash=kv_cache_utils.NONE_HASH,
+                num_lookahead_tokens=vllm_config.num_hash_lookahead_tokens,
+            )
 
         self.step_fn = (
             self.step if self.batch_queue is None else self.step_with_batch_queue
@@ -290,6 +304,12 @@ def _initialize_kv_caches(self, vllm_config: VllmConfig) -> KVCacheConfig:
         # capture full cudagraphs initialize a minimal KV cache during it.
         # Attention-free models resolve the default so layout reads never precede
         # resolution.
+        if vllm_config.cache_config.enable_extensible_kv_cache:
+            # Disable it, everywhere, where it cannot be used.
+            executor = self.model_executor
+            reason = executor.extensible_kv_cache_unsupported_reason(kv_cache_specs)
+            if reason is not None:
+                executor.disable_extensible_kv_cache(reason)
         layout = resolve_kv_cache_layout(
             vllm_config,
             self.model_executor.get_supported_kv_cache_layouts(),
@@ -356,6 +376,18 @@ def _initialize_kv_caches(self, vllm_config: VllmConfig) -> KVCacheConfig:
                 if participating
                 else [g.kv_cache_spec.block_size for g in kv_cache_groups]
             )
+        # With an extensible KV cache, `num_blocks` is the reserved capacity
+        # until the post-warmup measurement below.
+        extensible = vllm_config.cache_config.enable_extensible_kv_cache
+        # A KV connector may register each committed segment for RDMA, which
+        # cannot span physical chunks: keep segment bounds granule-aligned.
+        commit_granule: int | None = None
+        if extensible and kv_cache_groups and vllm_config.kv_transfer_config:
+            commit_granule = max(self.collective_rpc("kv_cache_commit_granule"))
+            align_extensible_kv_cache_capacity(
+                vllm_config, kv_cache_configs, scheduler_kv_cache_config, commit_granule
+            )
+        if kv_cache_groups and not extensible:
             update_kv_cache_capacity(vllm_config, scheduler_kv_cache_config)
 
         vllm_config.validate_block_size()
@@ -377,6 +409,20 @@ def _initialize_kv_caches(self, vllm_config: VllmConfig) -> KVCacheConfig:
             kv_cache_config.cache_hit_alignment_tokens = cache_hit_alignment_tokens
 
         self.model_executor.initialize_from_config(kv_cache_configs)
+        compilation_times = []
+        if not envs.VLLM_ELASTIC_EP_SCALE_UP_LAUNCH:
+            compilation_times = self.model_executor.compile_or_warm_up_model()
+        if extensible:
+            num_blocks = finalize_extensible_kv_cache(
+                vllm_config,
+                kv_cache_configs,
+                scheduler_kv_cache_config,
+                compilation_times,
+                commit_granule,
+            )
+            self.model_executor.extend_kv_cache(num_blocks)
+        # Under the extensible KV cache, workers create the connector in
+        # `extend_kv_cache`, so ask for its block count only afterwards.
         kv_transfer_config = vllm_config.kv_transfer_config
         if (
             kv_transfer_config is not None
@@ -392,8 +438,6 @@ def _initialize_kv_caches(self, vllm_config: VllmConfig) -> KVCacheConfig:
             vllm_config.cache_config.simple_cpu_offload_num_blocks = (
                 aligned_num_cpu_blocks
             )
-        if not envs.VLLM_ELASTIC_EP_SCALE_UP_LAUNCH:
-            self.model_executor.compile_or_warm_up_model()
 
         elapsed = time.time() - start
         compile_time = vllm_config.compilation_config.compilation_time
@@ -807,14 +851,13 @@ def step_with_batch_queue(
         # in a field and do it immediately once step_with_batch_queue is
         # re-called. The latter slightly favors TTFT over TPOT/throughput.
         if deferred_scheduler_output:
-            # When draft tokens are used with structured output, validate them
-            # before computing the grammar bitmask for the deferred request.
+            # Update the deferred scheduler output before computing its grammar
+            # bitmask.
             if self.check_for_draft_tokens:
                 draft_token_ids = self.model_executor.take_draft_token_ids()
                 if draft_token_ids is not None:
-                    # Update the draft token ids in the scheduler output to
-                    # filter out the invalid spec tokens, which will be padded
-                    # with -1 and skipped by the grammar bitmask computation.
+                    # Update the draft token ids in the scheduler output while
+                    # preserving the worker-provided draft slots.
                     self.scheduler.update_draft_token_ids_in_output(
                         draft_token_ids, deferred_scheduler_output
                     )
@@ -962,6 +1005,7 @@ def _finish_pause(self, clear_cache: bool) -> None:
             self.model_executor.collective_rpc("synchronize_device")
             if clear_cache:
                 self._reset_caches()
+                self._set_kv_cache_reclaimable(True)
             if preparing:
                 self.release_kv_cache_memory()
                 self.collective_rpc("checkpoint_prepare")
@@ -1009,10 +1053,22 @@ def resume_scheduler(self) -> None:
                 f"Cannot resume scheduling while the checkpoint state is "
                 f"{self._checkpoint_state}"
             )
+        self._set_kv_cache_reclaimable(False)
         self.scheduler.set_pause_state(PauseState.UNPAUSED)
         if self._checkpoint_state == CheckpointState.RESTORING:
             self._checkpoint_state = CheckpointState.RUNNING
 
+    def _set_kv_cache_reclaimable(self, reclaimable: bool) -> None:
+        """Tell the workers whether their extensible KV cache holds nothing to
+        keep, as from a pause that cleared it until the resume, so that a weight
+        reload may borrow its memory."""
+        if reclaimable == self._kv_cache_reclaimable:
+            return
+        self._kv_cache_reclaimable = reclaimable
+        executor = self.model_executor
+        if executor.vllm_config.cache_config.enable_extensible_kv_cache:
+            executor.collective_rpc("set_kv_cache_reclaimable", args=(reclaimable,))
+
     def is_scheduler_paused(self) -> bool:
         """Return whether the scheduler is in any pause state."""
         return self.scheduler.pause_state != PauseState.UNPAUSED
@@ -1122,6 +1178,14 @@ def is_sleeping(self) -> bool:
         """Check if engine is sleeping at any level."""
         return self.is_scheduler_paused() or self.model_executor.is_sleeping
 
+    def is_inference_ready(self) -> bool:
+        """Whether the engine can accept and process new requests: it is not
+        paused (including for a checkpoint) or asleep, and no weight update has
+        left the weights incomplete."""
+        return not self.is_sleeping() and not any(
+            self.collective_rpc("has_incomplete_weight_update")
+        )
+
     def execute_dummy_batch(self):
         self.model_executor.execute_dummy_batch()
 
@@ -1180,7 +1244,10 @@ def preprocess_add_request(self, request: EngineCoreRequest) -> tuple[Request, i
                 request.mm_features
             )
 
-        req = Request.from_engine_core_request(request, self.request_block_hasher)
+        frontend_hashes = frontend_block_hashes(request, self.prefix_hash_config)
+        req = Request.from_engine_core_request(
+            request, self.request_block_hasher, frontend_hashes
+        )
         if req.use_structured_output:
             # Note on thread safety: no race condition.
             # `grammar_init` is only invoked in input processing thread. For
@@ -1226,6 +1293,7 @@ def __init__(
     ):
         self.input_queue = queue.Queue[tuple[EngineCoreRequestType, Any]]()
         self.output_queue = queue.Queue[tuple[int, EngineCoreOutputs] | bytes]()
+
         executor_fail_callback = lambda: self.input_queue.put_nowait(
             (EngineCoreRequestType.EXECUTOR_FAILED, b"")
         )
@@ -1986,6 +2054,7 @@ def _make_ready_response(self) -> EngineCoreReadyResponse:
                 if envs.VLLM_SERVER_DEV_MODE
                 else None
             ),
+            prefix_hash_config=self.prefix_hash_config,
         )
 
     def process_input_sockets(
@@ -2509,7 +2578,11 @@ def resume_scheduler(self):
             self.dp_group, self.scheduler.has_unfinished_requests()
         )
 
-        if has_global_unfinished:
+        # Once a scale_down has disabled the DP coordinator, resume into
+        # always-running — no coordinator wake-up is available.
+        if has_global_unfinished or (
+            self.enable_fault_tolerance and self.ft_sentinel.coordinator_disabled
+        ):
             self.engines_running = True
 
     def barrier(self):
@@ -2671,6 +2744,9 @@ def _has_global_unfinished_reqs(self, local_unfinished: bool) -> bool:
             self.ignore_start_dp_wave = True
             self.pending_pause = False
             logger.debug("DP pause consensus reached, ignoring START_DP_WAVE.")
+        elif self.enable_fault_tolerance and self.ft_sentinel.coordinator_disabled:
+            # Coordinator disabled: never idle-pause; engines keep stepping.
+            return True
 
         return has_unfinished
 
diff --git a/vllm/v1/engine/core_client.py b/vllm/v1/engine/core_client.py
index 924045c80a3a..b56deccaa949 100644
--- a/vllm/v1/engine/core_client.py
+++ b/vllm/v1/engine/core_client.py
@@ -408,6 +408,8 @@ def __init__(
     def get_output(self) -> EngineCoreOutputs:
         outputs, model_executed = self.engine_core.step_fn()
         self.engine_core.post_step(model_executed=model_executed)
+        # Read by LLMEngine.step() to keep the DP ranks' collectives aligned.
+        self.last_step_model_executed = model_executed
         return outputs and outputs.get(0) or EngineCoreOutputs()
 
     def get_supported_tasks(self) -> tuple[SupportedTask, ...]:
@@ -1469,15 +1471,40 @@ async def compute_weight_checksums_async(self) -> list[dict[str, str]]:
     async def handle_fault(
         self, ft_request: FaultToleranceRequest
     ) -> FaultToleranceResult:
-        res = await self.call_utility_async(FT_UTILITY_METHOD, ft_request)
-        result = msgspec.convert(res, FaultToleranceResult)
+        status = self._engine_status[self.engine_ranks_managed[0]]
+
+        # Reject a second recovery while one is in flight
+        if status.get("ft_state") == "recovering":
+            return FaultToleranceResult(
+                request_id=ft_request.request_id,
+                success=False,
+                reason="another recovery is in progress",
+            )
+
+        status["ft_state"] = "recovering"
+        status["last_ft_request_id"] = ft_request.request_id
+        status.pop("ft_error", None)
+
+        try:
+            res = await self.call_utility_async(FT_UTILITY_METHOD, ft_request)
+            result = msgspec.convert(res, FaultToleranceResult)
+        except Exception as e:
+            result = FaultToleranceResult(ft_request.request_id, False, str(e))
+
         if not result.success:
-            status = self._engine_status.get(self.engine_ranks_managed[0])
-            if status is not None:
-                status["last_ft_request_id"] = result.request_id
-                status["ft_error"] = result.reason
+            # Status pushes replace the cached entry; re-read before recording.
+            status = self._engine_status[self.engine_ranks_managed[0]]
+            status["ft_state"] = "failed"
+            status["last_ft_request_id"] = result.request_id
+            status["ft_error"] = result.reason
+
         return result
 
+    @property
+    def engine_status(self) -> dict:
+        """Cached FT status entry of the managed engine."""
+        return self._engine_status.get(self.engine_ranks_managed[0], {})
+
     async def get_status(self):
         return {
             "schema_version": 1,
diff --git a/vllm/v1/engine/llm_engine.py b/vllm/v1/engine/llm_engine.py
index 1ec14aaa4d0b..4875b13fd102 100644
--- a/vllm/v1/engine/llm_engine.py
+++ b/vllm/v1/engine/llm_engine.py
@@ -325,6 +325,12 @@ def step(self) -> list[RequestOutput | PoolingRequestOutput]:
         # 1) Get EngineCoreOutput from the EngineCore.
         with record_function_or_nullcontext("llm_engine step: get_output"):
             outputs = self.engine_core.get_output()
+            if self.dp_group is not None and not getattr(
+                self.engine_core, "last_step_model_executed", True
+            ):
+                # Each DP rank joins one forward's DP sync per step; a step that
+                # ran no forward joins with a dummy batch, as DPEngineCoreProc.
+                self.engine_core.execute_dummy_batch()
 
         # 2) Process EngineCoreOutputs.
         with record_function_or_nullcontext("llm_engine step: process_outputs"):
diff --git a/vllm/v1/engine/tensor_ipc.py b/vllm/v1/engine/tensor_ipc.py
index 2b8fbadd97a5..774ddc6c0960 100644
--- a/vllm/v1/engine/tensor_ipc.py
+++ b/vllm/v1/engine/tensor_ipc.py
@@ -54,13 +54,6 @@ def __init__(self, queue: TensorIpcQueue):
         self._message_counter = 0
         self._sender_id = uuid.uuid4().hex[:8]
 
-    def set_target_engine(self, target_engine: int) -> None:
-        if target_engine != 0:
-            raise IndexError(
-                "TensorIpcSender only supports a single queue; "
-                f"got target engine {target_engine}"
-            )
-
     def new_message(self) -> None:
         self._message_counter += 1
         self._tensor_id_counter = 0
diff --git a/vllm/v1/executor/abstract.py b/vllm/v1/executor/abstract.py
index dc18b4375e06..e10270986d16 100644
--- a/vllm/v1/executor/abstract.py
+++ b/vllm/v1/executor/abstract.py
@@ -130,7 +130,7 @@ def initialize_from_config(self, kv_cache_configs: list[KVCacheConfig]) -> None:
         """Initialize the KV caches on the underlying workers."""
         self.collective_rpc("initialize_from_config", args=(kv_cache_configs,))
 
-    def compile_or_warm_up_model(self) -> None:
+    def compile_or_warm_up_model(self) -> list[CompilationTimes]:
         """Compile/warm up the model and capture cudagraphs on workers."""
         compilation_times: list[CompilationTimes] = self.collective_rpc(
             "compile_or_warm_up_model"
@@ -146,6 +146,32 @@ def compile_or_warm_up_model(self) -> None:
             self.vllm_config.compilation_config.encoder_compilation_time = max(
                 t.encoder for t in compilation_times
             )
+        return compilation_times
+
+    def extensible_kv_cache_unsupported_reason(
+        self, kv_cache_specs: list[dict[str, KVCacheSpec]]
+    ) -> str | None:
+        """Why the extensible KV cache cannot be used on some worker, if so."""
+        if not any(kv_cache_specs):
+            return "the model has no KV cache"
+        unsupported: list[str | None] = self.collective_rpc(
+            "extensible_kv_cache_unsupported_reason"
+        )
+        return next((r for r in unsupported if r), None)
+
+    def disable_extensible_kv_cache(self, reason: str) -> None:
+        """Turn the extensible KV cache off here and on every worker."""
+        logger.warning(
+            "Disabling the extensible KV cache: %s. The KV cache will be sized "
+            "from profiling estimates instead of measured memory.",
+            reason,
+        )
+        self.vllm_config.cache_config.enable_extensible_kv_cache = False
+        self.collective_rpc("disable_extensible_kv_cache")
+
+    def extend_kv_cache(self, num_blocks: int) -> None:
+        """Commit the final size of an extensible KV cache on the workers."""
+        self.collective_rpc("extend_kv_cache", args=(num_blocks,))
 
     def register_failure_callback(self, callback: FailureCallback):  # noqa: B027
         """Register a function to be called if the executor enters a permanent
diff --git a/vllm/v1/executor/cuda_checkpoint_executor.py b/vllm/v1/executor/cuda_checkpoint_executor.py
new file mode 100644
index 000000000000..9b743e49fd29
--- /dev/null
+++ b/vllm/v1/executor/cuda_checkpoint_executor.py
@@ -0,0 +1,297 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Opt-in, single-node sleep executor preserving complete CUDA worker state."""
+
+import json
+import os
+import sys
+import threading
+import time
+from concurrent.futures import Future
+from functools import wraps
+from multiprocessing.process import BaseProcess
+from typing import Any, cast
+
+from vllm.logger import init_logger
+from vllm.platforms import current_platform
+from vllm.utils.cuda_process_checkpoint import CudaProcessCheckpoint
+from vllm.v1.executor.multiproc_executor import MultiprocExecutor
+
+logger = init_logger(__name__)
+
+
+def _serialized(method):
+    @wraps(method)
+    def call(self, *args, **kwargs):
+        with self._checkpoint_lock:
+            return method(self, *args, **kwargs)
+
+    return call
+
+
+def _validate_nccl_checkpoint(worker, suspended):
+    from vllm.distributed import get_tp_group
+
+    comm = cast(Any, get_tp_group().device_communicator).pynccl_comm
+    if (
+        comm is None
+        or not comm.available
+        or comm.disabled
+        or not comm.nccl.has_symbol("ncclCommSuspend")
+        or not comm.nccl.has_symbol("ncclCommResume")
+        or comm._suspended != suspended
+    ):
+        raise RuntimeError(
+            "NCCL checkpoint requires an active suspend/resume implementation"
+        )
+
+
+class CUDACheckpointExecutor(MultiprocExecutor):
+    """Use dedicated workers even at TP=1 so failed restoration is terminable."""
+
+    def __init__(self, vllm_config, monitor_workers=True):
+        self._checkpoint_lock = threading.RLock()
+        self._checkpoint = None
+        self._checkpoint_active = False
+        parallel = vllm_config.parallel_config
+        if sys.platform != "linux" or not current_platform.is_cuda():
+            raise ValueError("CUDA checkpoint requires Linux and NVIDIA CUDA")
+        if (
+            not vllm_config.use_v2_model_runner
+            or vllm_config.model_config.sleep_mode_backend != "cumem"
+            or parallel.enable_expert_parallel
+            or parallel.worker_cls not in ("auto", "vllm.v1.worker.gpu_worker.Worker")
+        ):
+            raise ValueError(
+                "CUDA checkpoint requires the V2 GPU model runner, cumem sleep "
+                "backend and standard GPU worker, without expert parallelism"
+            )
+        compilation = getattr(vllm_config, "compilation_config", None)
+        passes = getattr(compilation, "pass_config", None)
+        fused_allreduce = bool(getattr(compilation, "mode", 0)) and bool(
+            getattr(passes, "fuse_allreduce_rms", False)
+        )
+        standalone_flashinfer = os.environ.get("VLLM_ALLREDUCE_USE_FLASHINFER") == "1"
+        # The compiler fusion allocates FlashInfer workspaces independently
+        # of the standalone all-reduce environment switch.
+        self._checkpoint_flashinfer = parallel.tensor_parallel_size > 1 and (
+            standalone_flashinfer or fused_allreduce
+        )
+        self._use_nccl_suspend = (
+            parallel.tensor_parallel_size > 1
+            and vllm_config.model_config.enable_nccl_comm_suspend
+        )
+        if (
+            parallel.tensor_parallel_size < 1
+            or parallel.nnodes != 1
+            or parallel.decode_context_parallel_size != 1
+            or parallel.prefill_context_parallel_size != 1
+            or parallel.pipeline_parallel_size != 1
+            or parallel.data_parallel_size != 1
+            or parallel.world_size != parallel.tensor_parallel_size
+        ):
+            raise ValueError("CUDA checkpoint requires single-node tensor parallelism")
+        if (
+            vllm_config.kv_transfer_config
+            or vllm_config.ec_transfer_config
+            or vllm_config.weight_transfer_config
+        ):
+            raise ValueError(
+                "CUDA checkpoint does not support external cache or weight transfer"
+            )
+        if not vllm_config.model_config.enable_sleep_mode:
+            raise ValueError("CUDA checkpoint requires enable_sleep_mode=True")
+        if parallel.tensor_parallel_size > 1:
+            required = {
+                "NCCL_P2P_DISABLE": "1",
+                "NCCL_SHM_DISABLE": "1",
+                "NCCL_IB_DISABLE": "1",
+                "NCCL_CUMEM_ENABLE": "0",
+                "NCCL_CUMEM_HOST_ENABLE": "0",
+                "VLLM_ALLREDUCE_USE_FLASHINFER": "0",
+                "VLLM_ALLREDUCE_USE_SYMM_MEM": "0",
+                "VLLM_ALLREDUCE_USE_FLASHINFER_PCIE_IPC": "0",
+                "VLLM_USE_NCCL_SYMM_MEM": "0",
+            }
+            if vllm_config.model_config.enable_nccl_comm_suspend:
+                # Suspension releases P2P mappings while retaining addresses.
+                # Legacy SHM also survives CUDA checkpoint; let NCCL select
+                # P2P or SHM rather than requiring one transport topology.
+                required.pop("NCCL_P2P_DISABLE")
+                required.pop("NCCL_SHM_DISABLE")
+                # Live multicast allocations can block checkpoint even on
+                # newer user-mode drivers. Keep NVLS disabled until the sleep
+                # path can release those allocations before checkpoint.
+                if os.environ.get("NCCL_DISABLE_MEM_MANAGER", "0") != "0":
+                    raise ValueError("NCCL checkpoint requires the memory manager")
+                required.update(
+                    {
+                        "NCCL_CUMEM_ENABLE": "1",
+                        "NCCL_NVLS_ENABLE": "0",
+                        "VLLM_ALLREDUCE_USE_FLASHINFER": "1"
+                        if standalone_flashinfer
+                        else "0",
+                    }
+                )
+            if self._checkpoint_flashinfer:
+                required["VLLM_FLASHINFER_ALLREDUCE_BACKEND"] = "trtllm"
+            if not parallel.disable_custom_all_reduce or any(
+                os.environ.get(k) != v for k, v in required.items()
+            ):
+                raise ValueError(
+                    "CUDA checkpoint currently requires disabled custom "
+                    f"all-reduce and NCCL checkpoint configuration: {required}"
+                )
+        super().__init__(vllm_config, monitor_workers)
+
+    def _record(self, **event):
+        event["monotonic_s"] = time.monotonic()
+        logger.info("CUDA_PROCESS_CHECKPOINT %s", json.dumps(event))
+
+    def collective_rpc(
+        self,
+        method,
+        timeout=None,
+        args=(),
+        kwargs=None,
+        non_block=False,
+        unique_reply_rank=None,
+        kv_output_aggregator=None,
+        ec_output_aggregator=None,
+    ):
+        if self.is_failed:
+            raise RuntimeError("CUDA checkpoint executor failed")
+        if self._checkpoint_active:
+            # EngineCore's first pause synchronized and cleared these caches.
+            # No new worker work can run while checkpointed, so a repeated
+            # pause can acknowledge that state without calling locked CUDA APIs.
+            if (
+                method
+                in ("synchronize_device", "reset_mm_cache", "reset_encoder_cache")
+                and not args
+                and not kwargs
+                and unique_reply_rank is None
+                and kv_output_aggregator is None
+                and ec_output_aggregator is None
+            ):
+                result = [None] * len(self.workers)
+                if non_block:
+                    future: Future[Any] = Future()
+                    future.set_result(result)
+                    return future
+                return result
+            raise RuntimeError(
+                "Workers are checkpointed; wake them before issuing RPCs"
+            )
+        if method in ("sleep", "wake_up") and timeout is None:
+            timeout = 120
+        return super().collective_rpc(
+            method,
+            timeout=timeout,
+            args=args,
+            kwargs=kwargs,
+            non_block=non_block,
+            unique_reply_rank=unique_reply_rank,
+            kv_output_aggregator=kv_output_aggregator,
+            ec_output_aggregator=ec_output_aggregator,
+        )
+
+    def check_health(self):
+        if self.is_failed:
+            raise RuntimeError("CUDA checkpoint executor failed")
+        if not self._checkpoint_active:
+            super().check_health()
+
+    @staticmethod
+    def _ensure_worker_termination(worker_procs: list[BaseProcess]):
+        # Dedicated workers can be inside locked CUDA calls. Destroy their
+        # processes without restoring GPU memory or waiting for CUDA cleanup.
+        # Reap within the outer process manager's shutdown window.
+        for proc in worker_procs:
+            if proc.is_alive():
+                proc.kill()
+        deadline = time.monotonic() + 3
+        for proc in worker_procs:
+            proc.join(timeout=max(0, deadline - time.monotonic()))
+            if proc.exitcode is None:
+                logger.error(
+                    "Worker pid=%s did not exit within cleanup deadline", proc.pid
+                )
+
+    def _close_checkpoint(self):
+        checkpoint, self._checkpoint = self._checkpoint, None
+        if checkpoint is not None:
+            try:
+                checkpoint.close()
+            except Exception:
+                logger.exception("Failed to close CUDA checkpoint helper")
+
+    def _stop_failed(self):
+        self.is_failed = True
+        # Terminate worker processes even if a CUDA API is blocked. Never try
+        # to continue serving a partially checkpointed rank set.
+        super().shutdown()
+        self._close_checkpoint()
+        callback = self.failure_callback
+        if callback is not None:
+            self.failure_callback = None
+            callback()
+
+    @_serialized
+    def sleep(self, level=1):
+        if self.is_failed:
+            raise RuntimeError("CUDA checkpoint executor failed")
+        if level != 1:
+            raise ValueError("CUDA checkpoint currently supports sleep level 1 only")
+        if self._checkpoint_active:
+            return
+        # Match the parent's no-op after KV-first partial wake: weights are
+        # still asleep, but CUDA/NCCL have resumed and must not be assumed
+        # suspended. Fully wake all tags before starting another checkpoint.
+        if "weights" in self.sleeping_tags:
+            return
+        if self._checkpoint is None:
+            self._checkpoint = CudaProcessCheckpoint(self._record)
+        try:
+            if self._use_nccl_suspend:
+                self.collective_rpc(
+                    _validate_nccl_checkpoint, args=(False,), timeout=30
+                )
+            super().sleep(level)
+            if self._use_nccl_suspend:
+                self.collective_rpc(_validate_nccl_checkpoint, args=(True,), timeout=30)
+            if self._checkpoint_flashinfer:
+                self.collective_rpc("checkpoint_prepare", timeout=120)
+            self._checkpoint.suspend([w.proc.pid for w in self.workers])
+            self._checkpoint_active = True
+        except Exception:
+            self._stop_failed()
+            raise
+
+    @_serialized
+    def wake_up(self, tags=None):
+        if self.is_failed:
+            raise RuntimeError("CUDA checkpoint executor failed")
+        if tags and any(tag not in self.sleeping_tags for tag in tags):
+            raise ValueError("Invalid or already awake sleep tag")
+        try:
+            if self._checkpoint_active:
+                cast(CudaProcessCheckpoint, self._checkpoint).resume()
+                self._checkpoint_active = False
+                if self._checkpoint_flashinfer:
+                    self.collective_rpc("checkpoint_restore", timeout=120)
+            super().wake_up(tags)
+            if self._use_nccl_suspend:
+                self.collective_rpc(
+                    _validate_nccl_checkpoint, args=(False,), timeout=30
+                )
+        except Exception:
+            self._stop_failed()
+            raise
+
+    @_serialized
+    def shutdown(self):
+        try:
+            super().shutdown()
+        finally:
+            self._close_checkpoint()
diff --git a/vllm/v1/fault_tolerance/engine_core_sentinel.py b/vllm/v1/fault_tolerance/engine_core_sentinel.py
index ca859ce30fb1..4f5672469523 100644
--- a/vllm/v1/fault_tolerance/engine_core_sentinel.py
+++ b/vllm/v1/fault_tolerance/engine_core_sentinel.py
@@ -2,17 +2,25 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """EngineCoreSentinel and fault_tolerant_wrapper for the engine core."""
 
-import json
 import threading
 from collections.abc import Callable
+from datetime import timedelta
 from typing import TYPE_CHECKING, cast
 
 import msgspec
+import torch
+import torch.distributed as dist
+from torch.distributed import PrefixStore, TCPStore
+from torch.distributed.distributed_c10d import Backend, _get_default_timeout
 
 from vllm.distributed import stateless_destroy_torch_distributed_process_group
-from vllm.distributed.utils import stateless_init_torch_distributed_process_group
+from vllm.distributed.utils import (
+    create_tcp_store,
+    enter_steady_state,
+    init_gloo_process_group,
+    set_gloo_backend_timeout,
+)
 from vllm.logger import init_logger
-from vllm.utils.network_utils import get_open_port
 from vllm.v1.core.sched.output import SchedulerOutput
 from vllm.v1.core.sched.scheduler import Scheduler
 from vllm.v1.engine import (
@@ -21,18 +29,27 @@
     EngineStatusType,
     UtilityOutput,
 )
-from vllm.v1.fault_tolerance.utils import FaultToleranceRequest, FaultToleranceResult
+from vllm.v1.fault_tolerance.utils import (
+    ALLOWED_FT_INSTRUCTIONS,
+    FaultToleranceRequest,
+    FaultToleranceResult,
+)
 from vllm.v1.outputs import ModelRunnerOutput
 from vllm.v1.request import RequestStatus
 from vllm.v1.serial_utils import UtilityResult, run_method
 
 if TYPE_CHECKING:
-    from vllm.v1.engine.core import EngineCoreProc
+    from vllm.v1.engine.core import DPEngineCoreProc, EngineCoreProc
 
 logger = init_logger(__name__)
 
 FT_UTILITY_METHOD = "handle_fault_tolerance"
 
+# How long an engine waits for the other alive engines to check in for the
+# same recovery round. This covers orchestrator dispatch skew (the
+# apply POSTs are fired concurrently),.
+_ROUND_CHECKIN_TIMEOUT_SEC = 10
+
 
 class EngineCoreSentinel:
     """Manages fault tolerance state for a single engine core."""
@@ -48,12 +65,54 @@ def __init__(self, engine: "EngineCoreProc", parallel_config):
         self.resumed.set()
         self.status_type = EngineStatusType.HEALTHY
         self.fault_info: str | None = None
-        self._dp_reinit_epoch = 0
+        self._initial_dp_size = parallel_config.data_parallel_size
+        self._dead_dp_ranks: set[int] = set()
+        self._recovery_store: TCPStore | None = None
+        self._install_executor_failure_callback()
+
+    @property
+    def coordinator_disabled(self) -> bool:
+        """True once any scale_down has committed dead ranks. The DP wave
+        wake-up path depends on the DP coordinator, which has no failover;
+        after a scale_down the engine never idle-pauses and keeps stepping
+        dummy batches instead."""
+        return bool(self._dead_dp_ranks)
+
+    def activate_cpu_group_timeouts(self) -> None:
+        """Give startup-created gloo groups the configured cpu timeout."""
+        enter_steady_state()
+        timeout_seconds = self.parallel_config.cpu_distributed_timeout_seconds
+        if timeout_seconds is None:
+            return
+        if self._initial_dp_size > 1:
+            set_gloo_backend_timeout(
+                cast("DPEngineCoreProc", self.engine).dp_group,
+                timedelta(seconds=timeout_seconds),
+            )
+        self.engine.model_executor.collective_rpc(
+            "handle_ft_command",
+            args=(
+                FaultToleranceRequest(
+                    instruction="activate_cpu_group_timeouts", params={}
+                ),
+            ),
+        )
 
     def handle_command(self, client_idx: int, call_id: int, ft_args: dict):
         """Dispatch an FT command by instruction name."""
         ft_request = FaultToleranceRequest(**ft_args)
-        if self.status_type != EngineStatusType.UNHEALTHY:
+        if ft_request.instruction not in ALLOWED_FT_INSTRUCTIONS:
+            reason = (
+                f"[FT] Rejecting unknown instruction "
+                f"'{ft_request.instruction}' on engine {self.engine_index}"
+            )
+            logger.warning(reason)
+            result = FaultToleranceResult(
+                request_id=ft_request.request_id,
+                success=False,
+                reason=reason,
+            )
+        elif self.status_type != EngineStatusType.UNHEALTHY:
             reason = (
                 f"[FT] Rejecting {ft_request.instruction} on engine "
                 f"{self.engine_index}: status is {self.status_type.name}"
@@ -95,8 +154,19 @@ def on_fault(self, exc: Exception):
             and engine.model_executor.is_failed
         ):
             self.status_type = EngineStatusType.DEAD
+            mask = None
         else:
-            self.status_type = EngineStatusType.UNHEALTHY
+            # Push DIAGNOSING first. Commands are rejected until UNHEALTHY.
+            self.status_type = EngineStatusType.DIAGNOSING
+            self._push_status()
+            try:
+                mask = self._query_mask()
+            except Exception:
+                logger.warning("[FT] Failed to query mask for status push")
+                mask = None
+            # on_executor_failed may have set DEAD concurrently.
+            if self.status_type == EngineStatusType.DIAGNOSING:
+                self.status_type = EngineStatusType.UNHEALTHY
         self.fault_info = f"{type(exc).__name__}"
         logger.info(
             "[FT] Engine %d status -> %s:",
@@ -104,7 +174,7 @@ def on_fault(self, exc: Exception):
             self.status_type.name,
             exc_info=exc,
         )
-        self._push_status()
+        self._push_status(mask)
 
     def _drain_batch_queue(self) -> None:
         """Drain the batch queue, feeding leftover kv_connector_outputs to
@@ -139,11 +209,36 @@ def _drain_batch_queue(self) -> None:
             )
         engine.batch_queue.clear()
 
-    def _push_status(self):
+    def _install_executor_failure_callback(self) -> None:
+        """Report executor failure to the sentinel as well.
+
+        The busy loop cannot consume EXECUTOR_FAILED while it waits for
+        recovery, so the sentinel must observe the failure directly.
+        """
+        executor = self.engine.model_executor
+        upstream_callback = getattr(executor, "failure_callback", None)
+        if upstream_callback is None:
+            return
+
+        def callback() -> None:
+            upstream_callback()
+            self.on_executor_failed()
+
+        executor.register_failure_callback(callback)
+
+    def on_executor_failed(self):
+        """Notify the client about the executor failure"""
+        if self.status_type != EngineStatusType.HEALTHY:
+            self.status_type = EngineStatusType.DEAD
+            self._push_status()
+
+    def _push_status(self, mask: list[int] | None = None):
         """Push current health to the client so it can refresh its cache."""
         payload = {"id": self.engine_index, "status": self.status_type.name.lower()}
         if self.status_type == EngineStatusType.UNHEALTHY:
             payload["fault_info"] = self.fault_info
+            if mask is not None:
+                payload["mask"] = mask
         outputs = EngineCoreOutputs(
             utility_output=UtilityOutput(
                 call_id=FT_STATUS_CALL_ID,
@@ -153,61 +248,271 @@ def _push_status(self):
         outputs.engine_index = self.engine_index
         self.engine.output_queue.put_nowait((0, outputs))
 
+    def _query_mask(self) -> list[int]:
+        """Union of all workers' all2all masks.
+
+        A rank is excluded if any worker suspects it.
+        """
+        ft_request = FaultToleranceRequest(instruction="query_mask", params={})
+        results = self.engine.model_executor.collective_rpc(
+            "handle_ft_command",
+            args=(ft_request,),
+            timeout=self.engine_recovery_timeout_sec,
+        )
+        return [max(bits) for bits in zip(*(r["mask"] for r in results))]
+
     def retry(self, ft_request: FaultToleranceRequest) -> FaultToleranceResult:
-        engine = self.engine
-        executor = engine.model_executor
+        # Workers replay masks for the cumulative dead set.
+        ft_request.params["dead_dp_ranks"] = sorted(self._dead_dp_ranks)
+        self._prepare_round(ft_request)
+        self._recover_and_vote(ft_request, terminal=False)
+        return self._mark_healthy(ft_request)
 
-        ft_request.params.update(self._reinit_dp_group())
+    def scale_down(self, ft_request: FaultToleranceRequest) -> FaultToleranceResult:
+        if self._initial_dp_size == 1:
+            raise ValueError(
+                "scale_down requires data_parallel_size > 1; a dp=1 "
+                "engine has no ranks to remove"
+            )
+        removed_set = set(ft_request.params["removed_dp_ranks"])
+        my_rank = self.parallel_config.data_parallel_rank
+        newly_dead = removed_set - self._dead_dp_ranks
+        if (
+            not removed_set
+            or not removed_set <= set(range(self._initial_dp_size))
+            or my_rank in removed_set
+            or set(range(self._initial_dp_size)) <= self._dead_dp_ranks | newly_dead
+            or not newly_dead
+        ):
+            raise ValueError(
+                f"Invalid removed_dp_ranks {sorted(removed_set)} for engine "
+                f"{self.engine_index} (dp_rank={my_rank}, "
+                f"dead_dp_ranks={sorted(self._dead_dp_ranks)})"
+            )
+
+        # Removing the store master (the lowest alive rank) requires the new
+        # store location in the request.
+        alive_before = sorted(set(range(self._initial_dp_size)) - self._dead_dp_ranks)
+        if alive_before[0] in newly_dead and (
+            ft_request.params.get("dp_store_port") is None
+            or ft_request.params.get("dp_master_ip") is None
+        ):
+            raise ValueError(
+                "dp_store_port and dp_master_ip required when the store "
+                f"master (rank {alive_before[0]}) is removed"
+            )
+
+        new_dead = self._dead_dp_ranks | newly_dead
+        ft_request.params["dead_dp_ranks"] = sorted(new_dead)
+        self._prepare_round(ft_request, dead_ranks=new_dead)
+        self._recover_and_vote(ft_request, terminal=True)
+        # Commit the dead set only after the full recovery succeeded.
+        self._dead_dp_ranks = new_dead
+        logger.info(
+            "[FT] Engine %d scale_down complete: removed %s, "
+            "cumulative dead_dp_ranks=%s",
+            self.engine_index,
+            sorted(newly_dead),
+            sorted(new_dead),
+        )
+        return self._mark_healthy(ft_request)
+
+    def _prepare_round(
+        self, ft_request: FaultToleranceRequest, dead_ranks: set[int] | None = None
+    ) -> None:
+        """Fill the round's coordination params.
+
+        The round's rendezvous keys are namespaced by request_id, so every
+        engine must carry the same one.
+        """
+        params = ft_request.params
+        dead = self._dead_dp_ranks if dead_ranks is None else dead_ranks
+        # The rebuilt gloo group contains only alive members; its internal
+        # rank/size are dense over sorted(alive), while parallel_config keeps
+        # the frozen original values.
+        alive = sorted(set(range(self._initial_dp_size)) - dead)
+        master_ip = self.parallel_config.data_parallel_master_ip
+        if self._initial_dp_size > 1 and dead_ranks is not None:
+            # The lowest alive rank hosts the TCPStore master; rebuild the
+            # store if that rank was just removed.
+            alive_before = sorted(
+                set(range(self._initial_dp_size)) - self._dead_dp_ranks
+            )
+            if alive_before[0] in dead_ranks - self._dead_dp_ranks:
+                master_ip = params["dp_master_ip"]
+                self._rebuild_dp_store(
+                    master_ip,
+                    params["dp_store_port"],
+                    is_master=(self.parallel_config.data_parallel_rank == alive[0]),
+                    num_clients=len(alive),
+                )
+        params["dp_master_ip"] = master_ip
+        params["dp_group_rank"] = alive.index(self.parallel_config.data_parallel_rank)
+        params["dp_group_size"] = len(alive)
+        recovery_round = ft_request.request_id
+        params["recovery_round"] = recovery_round
+        if self._initial_dp_size == 1:
+            # dp=1 has no dp_store/dp_group; the engine still hosts the
+            # recovery store for its own workers' TP group reinit.
+            params["recovery_store_port"] = self._coordinate_recovery_store_port(
+                master_ip, recovery_round, is_master=True
+            )
+            return
+
+        store = cast("DPEngineCoreProc", self.engine).dp_store
+        member_prefix = f"ft_round_member_{recovery_round}"
+        store.set(f"{member_prefix}_{params['dp_group_rank']}", b"1")
+        try:
+            store.wait(
+                [f"{member_prefix}_{r}" for r in range(len(alive))],
+                timedelta(seconds=_ROUND_CHECKIN_TIMEOUT_SEC),
+            )
+        except Exception as e:
+            raise RuntimeError(
+                f"[FT] Not all engines joined recovery round "
+                f"'{recovery_round}' within {_ROUND_CHECKIN_TIMEOUT_SEC}s; "
+                f"the same request_id must be sent to every engine."
+            ) from e
+
+    def _rebuild_dp_store(
+        self,
+        host: str,
+        port: int,
+        is_master: bool,
+        num_clients: int,
+    ) -> None:
+        """Rebuild dp_store when the old store master was removed."""
+        timeout_seconds = self.parallel_config.cpu_distributed_timeout_seconds
+        if timeout_seconds is not None:
+            timeout = timedelta(seconds=timeout_seconds)
+        else:
+            timeout = _get_default_timeout(Backend.GLOO)
+        engine = cast("DPEngineCoreProc", self.engine)
+
+        if hasattr(engine, "dp_store"):
+            del engine.dp_store
+        engine.dp_store = TCPStore(
+            host,
+            port,
+            num_clients,
+            is_master=is_master,
+            timeout=timeout,
+        )
+
+    def _reinit_groups(self, ft_request: FaultToleranceRequest) -> None:
+        """Reinit the engine DP group (dp>1) and host the recovery store."""
+        if self._initial_dp_size == 1:
+            return
+        params = ft_request.params
+        recovery_round = params["recovery_round"]
+        master_ip = params["dp_master_ip"]
+        params["recovery_store_port"] = self._coordinate_recovery_store_port(
+            master_ip, recovery_round, is_master=params["dp_group_rank"] == 0
+        )
+        self._reinit_engine_groups(
+            params["dp_group_rank"], params["dp_group_size"], recovery_round
+        )
+        # Commit the master IP only after the group reinit succeeded, so a
+        # failed recovery leaves a consistent state that can be retried.
+        self.parallel_config.data_parallel_master_ip = master_ip
+
+    def _recover_and_vote(
+        self,
+        ft_request: FaultToleranceRequest,
+        terminal: bool,
+    ) -> None:
+        """Reinit the engine groups, dispatch the command to workers, then
+        all surviving engines vote one success bit over the rebuilt DP group
+        so they commit or fail together. Any failure is terminal when
+        `terminal` is set; otherwise the engine stays UNHEALTHY, retryable."""
+        engine = self.engine
         if hasattr(engine, "step_counter"):
             engine.step_counter = 0
 
-        executor.collective_rpc("handle_ft_command", args=(ft_request,))
-        self._drain_batch_queue()
+        local_error: Exception | None = None
+        try:
+            self._reinit_groups(ft_request)
+            engine.model_executor.collective_rpc(
+                "handle_ft_command", args=(ft_request,)
+            )
+            self._drain_batch_queue()
+        except Exception as exc:
+            local_error = exc
+
+        success = torch.tensor([local_error is None], dtype=torch.int32)
+        if self._initial_dp_size > 1:
+            try:
+                dist.all_reduce(
+                    success,
+                    op=dist.ReduceOp.MIN,
+                    group=cast("DPEngineCoreProc", engine).dp_group,
+                )
+            except Exception:
+                success.zero_()
+        if success.item():
+            return
+
+        if terminal:
+            self.status_type = EngineStatusType.DEAD
+            self.fault_info = "Recovery failed; full restart required"
+            self._push_status()
+        if local_error is not None:
+            raise local_error
+        raise RuntimeError("[FT] Recovery vote failed")
+
+    def _mark_healthy(self, ft_request: FaultToleranceRequest) -> FaultToleranceResult:
         self.status_type = EngineStatusType.HEALTHY
         logger.info("[FT] Engine %d status -> HEALTHY", self.engine_index)
         self.resumed.set()
         self._push_status()
         return FaultToleranceResult(request_id=ft_request.request_id, success=True)
 
-    def _reinit_dp_group(self) -> dict:
-        """Reinit DP process group if in DP mode. Returns worker params."""
-        engine = self.engine
-        if not hasattr(engine, "dp_group") or not hasattr(engine, "dp_store"):
-            return {}
-
-        parallel_config = engine.vllm_config.parallel_config
-        worker_key = f"ft_worker_dp_ports_{self._dp_reinit_epoch}"
-        engine_key = f"ft_engine_dp_port_{self._dp_reinit_epoch}"
-        self._dp_reinit_epoch += 1
-
-        if parallel_config.data_parallel_rank == 0:
-            worker_ports = [get_open_port() for _ in range(parallel_config.world_size)]
-            engine_port = get_open_port()
-            engine.dp_store.set(worker_key, json.dumps(worker_ports).encode())
-            engine.dp_store.set(engine_key, str(engine_port).encode())
-        else:
-            worker_ports = json.loads(engine.dp_store.get(worker_key).decode())
-            engine_port = int(engine.dp_store.get(engine_key).decode())
+    def _reinit_engine_groups(
+        self, dense_rank: int, dense_size: int, recovery_round: str
+    ) -> None:
+        """Reinit the engine DP group (dp>1). Worker group ports are
+        coordinated by the workers themselves via the recovery store."""
+        engine = cast("DPEngineCoreProc", self.engine)
 
+        prefix_store = PrefixStore(f"ft_engine_dp_{recovery_round}", engine.dp_store)
+        timeout = self.parallel_config.cpu_distributed_timeout
+        if timeout is None:
+            timeout = _get_default_timeout(Backend.GLOO)
+        new_group = init_gloo_process_group(
+            prefix_store=prefix_store,
+            group_rank=dense_rank,
+            group_size=dense_size,
+            timeout=timeout,
+        )
         stateless_destroy_torch_distributed_process_group(engine.dp_group)
-        engine.dp_group, engine.dp_store = (
-            stateless_init_torch_distributed_process_group(
-                parallel_config.data_parallel_master_ip,
-                engine_port,
-                parallel_config.data_parallel_rank,
-                parallel_config.data_parallel_size,
-                backend="gloo",
-                return_store=True,
-                timeout=parallel_config.cpu_distributed_timeout,
+        engine.dp_group = new_group
+
+    def _coordinate_recovery_store_port(
+        self, master_ip: str, recovery_round: str, is_master: bool
+    ) -> int:
+        """Create (is_master) or look up (others) the per-round TCPStore
+        that workers use to coordinate reinit ports; return its port."""
+        if is_master:
+            store = create_tcp_store(
+                master_ip, 0, is_master=True, world_size=-1, wait_for_workers=False
             )
-        )
-        return {"new_stateless_dp_group_ports": worker_ports}
+            self._recovery_store = store
+            if self._initial_dp_size > 1:
+                key = f"ft_recovery_store_port_{recovery_round}"
+                cast("DPEngineCoreProc", self.engine).dp_store.set(
+                    key, str(store.port).encode()
+                )
+            return store.port
+        key = f"ft_recovery_store_port_{recovery_round}"
+        return int(cast("DPEngineCoreProc", self.engine).dp_store.get(key).decode())
 
 
 def fault_tolerant_wrapper(busy_loop_func: Callable):
     """Wrap the busy loop to catch faults and delegate recovery."""
 
     def run_with_fault_tolerance(self: "EngineCoreProc"):
+        if self.enable_fault_tolerance:
+            self.ft_sentinel.activate_cpu_group_timeouts()
         while True:
             try:
                 busy_loop_func(self)
diff --git a/vllm/v1/fault_tolerance/utils.py b/vllm/v1/fault_tolerance/utils.py
index 0c1b1689b01a..6256522d726e 100644
--- a/vllm/v1/fault_tolerance/utils.py
+++ b/vllm/v1/fault_tolerance/utils.py
@@ -4,6 +4,11 @@
 
 import msgspec
 
+# Instructions accepted from external callers. The API router and the engine
+# both validate against this; keep it the single source of truth so a new
+# instruction cannot be allowed at one layer and rejected at the other.
+ALLOWED_FT_INSTRUCTIONS = frozenset({"retry", "scale_down"})
+
 
 class FaultToleranceResult(msgspec.Struct):
     request_id: str
diff --git a/vllm/v1/hisparse/backend.py b/vllm/v1/hisparse/backend.py
new file mode 100644
index 000000000000..2ed099e32219
--- /dev/null
+++ b/vllm/v1/hisparse/backend.py
@@ -0,0 +1,142 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Device operations used by the common HiSparse runtime and connector."""
+
+from __future__ import annotations
+
+from functools import cache
+from typing import TYPE_CHECKING, Any
+
+import torch
+
+from vllm.platforms import current_platform
+from vllm.utils.import_utils import resolve_obj_by_qualname
+
+if TYPE_CHECKING:
+    from vllm.v1.hisparse.runtime import HiSparseRuntime
+    from vllm.v1.kv_cache_interface import KVCacheConfig
+
+
+class HiSparseBackend:
+    supports_shared_host_pool = False
+    stats_dtype = torch.uint64
+
+    def check_available(self) -> None:
+        raise NotImplementedError
+
+    def pin_tensor(self, tensor: torch.Tensor) -> None:
+        raise NotImplementedError
+
+    def unpin_tensor(self, tensor: torch.Tensor) -> None:
+        raise NotImplementedError
+
+    def create_cache_ops(self, runtime: HiSparseRuntime) -> Any:
+        raise NotImplementedError
+
+    def validate_kv_cache_config(self, kv_cache_config: KVCacheConfig) -> None:
+        """Validate device cache formats before allocating or pinning storage."""
+
+    def prepare_kv_cache_views(
+        self,
+        kv_cache_config: KVCacheConfig,
+        kv_caches: dict[str, Any],
+        forward_context: dict[str, Any],
+    ) -> None:
+        """Adapt device views before runtime binding; retain allocation ownership.
+
+        Called once per cache initialization, inside the common rollback scope.
+        Backends must not put per-step metadata or transfers in this hook.
+        """
+
+    def create_ipc_event(self) -> Any:
+        """Create an event that can be exported to another process."""
+        raise NotImplementedError
+
+    def export_event(self, event: Any) -> Any:
+        raise NotImplementedError
+
+    def import_event(self, device: torch.device, handle: Any) -> Any:
+        raise NotImplementedError
+
+    def gather_resident_rows(
+        self, source: torch.Tensor, destination: torch.Tensor, row_ids: torch.Tensor
+    ) -> None:
+        """Copy resident rows into staging without changing their payload format."""
+        rows = row_ids.to(torch.long).clamp_min(0)
+        block_size = source.shape[1]
+        payload = source[rows // block_size, rows % block_size]
+        if payload.dtype != destination.dtype:
+            payload = payload.contiguous().view(destination.dtype)
+        destination.copy_(payload)
+
+    def is_current_stream_capturing(self) -> bool:
+        raise NotImplementedError
+
+    def swap_blocks_batch(
+        self, src: torch.Tensor, dst: torch.Tensor, sizes: torch.Tensor
+    ) -> None:
+        """Enqueue device-to-host copies on the current stream."""
+        raise NotImplementedError
+
+
+class CUDAHiSparseBackend(HiSparseBackend):
+    supports_shared_host_pool = True
+
+    def check_available(self) -> None:
+        names = (
+            "hisparse_resolve_residency",
+            "hisparse_invalidate_written_slots",
+            "hisparse_gather_plan",
+            "hisparse_gather_compact",
+        )
+        if not all(hasattr(torch.ops._C_cache_ops, name) for name in names):
+            raise RuntimeError(
+                "HiSparse requires its compiled _C_cache_ops CUDA kernels. "
+                "Rebuild vLLM from source with csrc/hisparse_kernels.cu."
+            )
+
+    def pin_tensor(self, tensor: torch.Tensor) -> None:
+        from vllm.v1.simple_kv_offload.cuda_mem_ops import pin_tensor
+
+        pin_tensor(tensor)
+
+    def unpin_tensor(self, tensor: torch.Tensor) -> None:
+        cudart = torch.cuda.cudart()
+        error = cudart.cudaHostUnregister(tensor.data_ptr())
+        if error.value != 0:
+            cudart.cudaGetLastError()
+            raise RuntimeError(f"cudaHostUnregister failed (code={error.value}).")
+
+    def create_cache_ops(self, runtime: HiSparseRuntime) -> Any:
+        return torch.ops._C_cache_ops
+
+    def create_ipc_event(self) -> Any:
+        return torch.cuda.Event(interprocess=True)
+
+    def export_event(self, event: Any) -> Any:
+        return event.ipc_handle()
+
+    def import_event(self, device: torch.device, handle: Any) -> Any:
+        return torch.cuda.Event.from_ipc_handle(device, handle)
+
+    def is_current_stream_capturing(self) -> bool:
+        return torch.cuda.is_current_stream_capturing()
+
+    def swap_blocks_batch(
+        self, src: torch.Tensor, dst: torch.Tensor, sizes: torch.Tensor
+    ) -> None:
+        from vllm import _custom_ops as ops
+
+        ops.swap_blocks_batch(src, dst, sizes)
+
+
+def get_hisparse_backend() -> HiSparseBackend:
+    backend_cls = current_platform.get_hisparse_backend_cls()
+    if backend_cls is None:
+        raise RuntimeError("The current platform does not support HiSparse.")
+    return _get_hisparse_backend(backend_cls)
+
+
+@cache
+def _get_hisparse_backend(backend_cls: str) -> HiSparseBackend:
+    return resolve_obj_by_qualname(backend_cls)()
diff --git a/vllm/v1/hisparse/binding.py b/vllm/v1/hisparse/binding.py
index 6a02dd571cc3..9dd061ef97aa 100644
--- a/vllm/v1/hisparse/binding.py
+++ b/vllm/v1/hisparse/binding.py
@@ -9,6 +9,7 @@
 import torch
 
 from vllm.config import VllmConfig
+from vllm.v1.hisparse.backend import get_hisparse_backend
 from vllm.v1.hisparse.layout import (
     HISPARSE_HOT_SUFFIX,
     HISPARSE_RESIDENT_SUFFIX,
@@ -126,8 +127,10 @@ def init_hisparse_kv_cache(
     vllm_config: VllmConfig,
     forward_context: dict[str, Any],
     block_tables: "BlockTables",
-) -> dict[str, torch.Tensor]:
+) -> dict[str, Any]:
     """Allocate and bind HiSparse caches within the caller's allocation context."""
+    backend = get_hisparse_backend()
+    backend.validate_kv_cache_config(kv_cache_config)
     host_pool = get_hisparse_host_pool(vllm_config, kv_cache_config)
     try:
         kv_caches = allocate_hisparse_kv_caches(
@@ -137,6 +140,7 @@ def init_hisparse_kv_cache(
             kernel_block_sizes,
             host_pool,
         )
+        backend.prepare_kv_cache_views(kv_cache_config, kv_caches, forward_context)
         cache_handles = bind_hisparse_kv_caches(
             forward_context=forward_context,
             kv_cache_config=kv_cache_config,
@@ -150,10 +154,7 @@ def init_hisparse_kv_cache(
         )
         return kv_caches
     except Exception:
-        if host_pool.shared_region is not None:
-            host_pool.shared_region.cleanup()
-        elif host_pool.registered is not None:
-            release_pinned_state([], [host_pool.registered])
+        release_pinned_state([], host_pool.pinned, host_pool.shared_region)
         drop_hisparse_host_pool(vllm_config)
         raise
 
@@ -210,7 +211,7 @@ def bind_hisparse_kv_caches(
         assert len(table_shapes) == 1
         max_num_reqs, max_num_pages = table_shapes.pop()
         residency = torch.zeros(
-            (max_num_reqs, len(resident_groups), max_num_pages),
+            (len(resident_groups), max_num_reqs, max_num_pages),
             dtype=torch.int32,
             device=block_tables.input_block_tables[resident_groups[0][0]].device,
         )
@@ -228,7 +229,7 @@ def bind_hisparse_kv_caches(
                 block_stride=tensor_config.block_stride,
                 num_blocks=kv_cache_config.num_blocks,
                 block_size=group.kv_cache_spec.block_size,
-                block_table=residency[:, resident_source_index],
+                block_table=residency[resident_source_index],
                 slot_mapping=block_tables.slot_mappings[group_id],
             )
             cache_handle.residency = residency
@@ -277,7 +278,7 @@ def bind_hisparse_kv_caches(
                 host_pool.registered.untyped_storage().data_ptr()
             )
             cache_handle.draft_layer = forward_context[layer_name].is_draft_layer
-            cache_handle.runtime.shared_host_region = host_pool.shared_region
+            cache_handle.runtime.host_pool = host_pool
             cache_handle.runtime.bind_source_cache(
                 source_cache,
                 registered_host_pool=host_pool.registered,
diff --git a/vllm/v1/hisparse/coordinator.py b/vllm/v1/hisparse/coordinator.py
index 374f100bc889..fddb57cba4f6 100644
--- a/vllm/v1/hisparse/coordinator.py
+++ b/vllm/v1/hisparse/coordinator.py
@@ -104,9 +104,12 @@ def __init__(
         max_model_len: int,
         *,
         num_reprefillable_tokens: int,
+        import_step_tokens: int = 0,
     ) -> None:
         self.managers = managers
         self.max_model_len = max_model_len
+        # Spec decode step tokens an import needs beyond its sequence.
+        self.import_step_tokens = import_step_tokens
         self.num_reprefillable_tokens = num_reprefillable_tokens
         groups = kv_cache_config.kv_cache_groups
 
@@ -173,6 +176,14 @@ def __init__(
         self._retained_copies: dict[int, tuple[KVCacheBlock, KVCacheBlock]] = {}
         # request -> prefix length an external load is still filling in.
         self._pending_imports: dict[str, int] = {}
+        # request -> full-sequence host pages of an admitted import. A waiting
+        # import holds host pages and is never preempted, so the pages it has
+        # yet to allocate stay reserved against later admissions.
+        self._import_pages: dict[str, int] = {}
+        self._import_reservations: dict[str, int] = {}
+        self._num_reserved_import_pages = 0
+        # The import the admission gate last sized, kept for its allocation.
+        self._sized_import: tuple[str, int] | None = None
         # Requests whose resident pages were all non-null when last scanned.
         # Pages only become null at admission or import (before the first
         # scan) and when the pool reuses an unpinned page (_lose_page).
@@ -540,6 +551,39 @@ 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._resident_requests.discard(request_id)
         self._pending_imports[request_id] = num_tokens
+        if self._sized_import is not None and self._sized_import[0] == request_id:
+            self._import_pages[request_id] = self._sized_import[1]
+        self._sized_import = None
+
+    def size_import(self, request_id: str, num_pages: int) -> None:
+        """Note the full-sequence host pages of an import being admitted.
+
+        The admission gate runs in the same ``allocate_slots`` call as the
+        allocation, so the next ``record_pending_host_import`` is this
+        request's if it is admitted; a size left by a rejected admission is
+        discarded.
+        """
+        self._sized_import = (request_id, num_pages)
+
+    @property
+    def reserved_import_pages(self) -> int:
+        """Host pages admitted imports still need for their sequence."""
+        return self._num_reserved_import_pages
+
+    def update_import_reservation(self, request_id: str, num_held_pages: int) -> None:
+        pages = self._import_pages.get(request_id)
+        if pages is None:
+            return
+        self._set_import_reservation(request_id, max(pages - num_held_pages, 0))
+        if num_held_pages >= pages:
+            del self._import_pages[request_id]
+
+    def _set_import_reservation(self, request_id: str, num_pages: int) -> None:
+        self._num_reserved_import_pages += num_pages - self._import_reservations.pop(
+            request_id, 0
+        )
+        if num_pages > 0:
+            self._import_reservations[request_id] = num_pages
 
     def finish_host_import(self, request_id: str, *, failed: bool) -> None:
         num_tokens = self._pending_imports.pop(request_id, None)
@@ -853,6 +897,8 @@ def free(self, request_id: str) -> None:
         """Detach the request; its clean pages stay readable copies in the pool."""
         self._resident_requests.discard(request_id)
         self._pending_imports.pop(request_id, None)
+        self._import_pages.pop(request_id, None)
+        self._set_import_reservation(request_id, 0)
         state = self.request_states.pop(request_id, None)
         if state is None:
             return
@@ -899,6 +945,8 @@ def free(self, request_id: str) -> None:
 
 def get_hisparse_coordinator(
     kv_cache_manager: "KVCacheManager",
+    *,
+    import_step_tokens: int | None = None,
 ) -> HiSparseCoordinator:
     """Return the coordinator shared by a KV cache manager's HiSparse groups.
 
@@ -910,12 +958,14 @@ def get_hisparse_coordinator(
         coordinator = getattr(manager, "coordinator", None)
         if coordinator is not None:
             assert isinstance(coordinator, HiSparseCoordinator)
+            assert import_step_tokens in (None, coordinator.import_step_tokens)
             return coordinator
     coordinator = HiSparseCoordinator(
         kv_cache_manager.kv_cache_config,
         managers,
         kv_cache_manager.max_model_len,
         num_reprefillable_tokens=kv_cache_manager.coordinator.num_reprefillable_tokens,
+        import_step_tokens=import_step_tokens or 0,
     )
     if coordinator.host_manager is None:
         raise ValueError("No HiSparse cache group is configured.")
diff --git a/vllm/v1/hisparse/runtime.py b/vllm/v1/hisparse/runtime.py
index 85866ac0e9b7..b334eaf3f6eb 100644
--- a/vllm/v1/hisparse/runtime.py
+++ b/vllm/v1/hisparse/runtime.py
@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""NVIDIA HiSparse resident, host, and hot-cache data plane."""
+"""HiSparse resident, host, and hot-cache data plane."""
 
 from __future__ import annotations
 
@@ -22,11 +22,14 @@
 from vllm.platforms import current_platform
 from vllm.utils.math_utils import cdiv, round_up
 from vllm.utils.mem_utils import get_max_shared_memory_bytes
-from vllm.utils.torch_utils import async_tensor_h2d, current_stream
+from vllm.utils.torch_utils import async_tensor_h2d
 from vllm.v1.attention.backend import max_decode_query_len
+from vllm.v1.hisparse.backend import get_hisparse_backend
 from vllm.v1.hisparse.types import SparseKVResidencyUpdate
-from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion
-from vllm.v1.simple_kv_offload.cuda_mem_ops import pin_tensor
+from vllm.v1.kv_offload.cpu.shared_offload_region import (
+    REGION_SEGMENTS,
+    SharedOffloadRegion,
+)
 
 logger = init_logger(__name__)
 
@@ -36,9 +39,10 @@
     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]
+# trtllm-gen sparse MLA reads up to ~250 rows from the start of its KV cache
+# whatever the cache's size: https://github.com/flashinfer-ai/flashinfer/issues/6484
+_MIN_STAGED_PREFILL_KV_ROWS = 256
 
 
 def is_hisparse_decode_batch(attn_metadata: Any | None) -> bool:
@@ -212,8 +216,12 @@ def allocate_pinned_host_pool(size: int) -> tuple[torch.Tensor, torch.Tensor]:
     return registered[:size], registered
 
 
+def pin_tensor(tensor: torch.Tensor) -> None:
+    get_hisparse_backend().pin_tensor(tensor)
+
+
 class HiSparseHostPool:
-    """Keep the single host backing and its exact CUDA registration range."""
+    """Keep the single host backing and its exact registration range."""
 
     def __init__(self, vllm_config: VllmConfig, kv_cache_config: KVCacheConfig) -> None:
         self.vllm_config = vllm_config
@@ -221,6 +229,7 @@ def __init__(self, vllm_config: VllmConfig, kv_cache_config: KVCacheConfig) -> N
         self.backing: torch.Tensor | None = None
         self.registered: torch.Tensor | None = None
         self.shared_region: SharedOffloadRegion | None = None
+        self.pinned: list[torch.Tensor] = []
 
     def allocate(self, size: int) -> torch.Tensor:
         if self.backing is not None:
@@ -239,7 +248,7 @@ def allocate(self, size: int) -> torch.Tensor:
             ]
             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.pinned, self.shared_region = allocate_hisparse_host_pools(
                 self.vllm_config,
                 tensor_sizes,
                 config.hisparse_host_num_blocks,
@@ -252,6 +261,7 @@ def allocate(self, size: int) -> torch.Tensor:
         else:
             check_hisparse_host_memory(size)
             self.backing, self.registered = allocate_pinned_host_pool(size)
+            self.pinned = [self.registered]
         return self.backing
 
 
@@ -276,7 +286,7 @@ 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()
+        get_hisparse_backend().supports_shared_host_pool
         and parallel.tensor_parallel_size > 1
         and parallel.pipeline_parallel_size == 1
         and parallel.prefill_context_parallel_size == 1
@@ -399,22 +409,23 @@ def allocate_hisparse_host_pools(
         barrier=tp_group.barrier,
         creator_memory_check=check_hisparse_host_memory,
         populate_shard=(tp_group.rank_in_group, tp_group.world_size),
+        num_segments=REGION_SEGMENTS,
     )
+    pinned: list[torch.Tensor] = []
     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
+            pinned.append(tensor)
         pools = [
             region.create_next_canonical_view(size).view(-1) for size in tensor_sizes
         ]
     except Exception:
-        region.cleanup()
+        release_pinned_state([], pinned, region)
         raise
-    return pools, [], region
+    return pools, pinned, region
 
 
 def release_pinned_state(
@@ -424,31 +435,31 @@ def release_pinned_state(
 ) -> None:
     """Synchronize and release registered host KV pools."""
     if pinned_host_pools or shared_host_region is not None:
+        backend = get_hisparse_backend()
         try:
             torch.accelerator.synchronize()
         except RuntimeError as e:
             logger.warning(
-                "HiSparse: CUDA context unusable at teardown (%s); leaving "
+                "HiSparse: device 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),
+                len(pinned_host_pools),
             )
             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:
+            try:
+                backend.unpin_tensor(tensor)
+            except RuntimeError as error:
                 logger.warning(
-                    "HiSparse: cudaHostUnregister failed (code=%d); leaving "
+                    "HiSparse: host unregister failed (%s); leaving "
                     "%d host-pool tensors pinned for kernel exit reclaim.",
-                    err.value,
+                    error,
                     len(pinned_host_pools),
                 )
-                cudart.cudaGetLastError()
                 break
             freed_bytes += tensor.nbytes
             pinned_host_pools.pop()
@@ -460,6 +471,8 @@ def release_pinned_state(
             )
 
     for runtime in runtimes:
+        runtime.cache_ops = None
+        runtime.host_pool = None
         del runtime._host_cache
         del runtime.registered_host_pool
         del runtime.hot_backing
@@ -630,14 +643,11 @@ def build_hisparse_prefill_staging_plan(
 
 
 def _has_hisparse_ops() -> bool:
-    if not hasattr(torch.ops, "_C_cache_ops"):
+    try:
+        get_hisparse_backend().check_available()
+    except RuntimeError:
         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")
-    )
+    return True
 
 
 class _SharedTopKState:
@@ -687,21 +697,21 @@ def __init__(
             dtype=torch.int32,
             device=device,
         )
-        self.lru_init = torch.arange(region_stride, dtype=torch.int16, device=device)
+        # NPU arange does not support int16 output, so generate int32 then cast.
+        self.lru_init = torch.arange(
+            region_stride, dtype=torch.int32, device=device
+        ).to(torch.int16)
         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)
-        # Request ids the resolver reads, built on the copy stream: a
-        # compute-stream write would race with its logical_topk_ready wait.
-        self.row_indices = torch.arange(max_swap_rows, dtype=torch.int32, device=device)
-        self.request_ids = torch.empty_like(self.row_indices)
         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)
+        stats_dtype = get_hisparse_backend().stats_dtype
+        self.swap_stats = torch.zeros(2, dtype=stats_dtype, device=device)
         self.swap_stats_host = torch.empty(
-            2, dtype=torch.uint64, device="cpu", pin_memory=True
+            2, dtype=stats_dtype, device="cpu", pin_memory=True
         )
         self.stats_row_bytes = 0
 
@@ -727,13 +737,9 @@ def __init__(
         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/hisparse_kernels.cu is included."
-            )
+        self.backend = get_hisparse_backend()
+        self.backend.check_available()
+        self.cache_ops: Any = None
         self.max_num_reqs = max_num_reqs
         self.row_width = row_width
         self.kv_dtype = kv_dtype
@@ -774,7 +780,7 @@ def __init__(
         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
+        self.host_pool: HiSparseHostPool | None = None
 
     @property
     def host_cache(self) -> torch.Tensor:
@@ -833,6 +839,7 @@ def bind_source_cache(
         self.registered_host_pool = (
             registered_host_pool if registered_host_pool is not None else kv_cache
         )
+        self.cache_ops = self.backend.create_cache_ops(self)
 
     def gather_prefill_cache(
         self,
@@ -848,24 +855,18 @@ def gather_prefill_cache(
             )
         row_width = kv_cache.shape[-1]
 
+        num_blocks = plan.row_ids.shape[1] // plan.block_size
+        min_blocks = cdiv(_MIN_STAGED_PREFILL_KV_ROWS, plan.block_size)
         staged = torch.empty(
-            (
-                plan.row_ids.shape[1] // plan.block_size,
-                plan.block_size,
-                row_width,
-            ),
+            (max(num_blocks, min_blocks), plan.block_size, row_width),
             dtype=kv_cache.dtype,
             device=plan.block_table.device,
-        )
+        )[:num_blocks]
         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(
+            self.backend.gather_resident_rows(
+                resident_cache, staged.view(-1, row_width), plan.gpu_row_ids[0]
+            )
+        self.cache_ops.hisparse_gather_plan(
             kv_cache.view(-1, row_width),
             staged,
             plan.row_ids,
@@ -880,7 +881,7 @@ def gather_prefill_cache(
     def reset_hot_state(self) -> None:
         """Drop all hot-buffer bookkeeping (hits become misses)."""
         group = self.index_group
-        compute_stream = current_stream()
+        compute_stream = torch.accelerator.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))
@@ -903,22 +904,6 @@ def invalidate_sorted_slots(
         )
         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
 
@@ -942,6 +927,7 @@ def _resolve_residency(
         *,
         resident: HiSparseCacheHandle | None = None,
         req_id_per_token: torch.Tensor,
+        num_valid_rows: torch.Tensor,
         block_table: torch.Tensor,
         topk_indices: torch.Tensor,
         block_size: int,
@@ -977,9 +963,9 @@ def _resolve_residency(
         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(
+        # Rows past num_valid_rows and rows without request state come out as
+        # -1 so the attention kernel masks them.
+        self.cache_ops.hisparse_resolve_residency(
             host_cache,
             hot.cache,
             self.hot_block_table,
@@ -1007,12 +993,13 @@ def _resolve_residency(
             if resident is not None and resident.view is not None
             else 0,
             0,
+            num_valid_rows,
         )
 
     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.cache_ops.hisparse_gather_compact(
             self.host_cache,
             hot.cache,
             shared.swap_host_physical_rows[shared_rows],
@@ -1035,23 +1022,16 @@ def _resolve_and_stage_group(
         num_valid_rows: torch.Tensor,
     ) -> None:
         group = self.index_group
-        compute_stream = current_stream()
+        compute_stream = torch.accelerator.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:
-            # CUDA-graph padding rows past the batch's tokens map to request 0;
-            # mark them -1 so residency resolution skips them.
-            num_tokens = logical_topk_indices.shape[0]
-            request_ids = group.request_ids[:num_tokens]
-            request_ids.copy_(req_id_per_token)
-            request_ids.masked_fill_(
-                group.row_indices[:num_tokens] >= num_valid_rows, -1
-            )
             self._resolve_residency(
                 resident=resident,
-                req_id_per_token=request_ids,
+                req_id_per_token=req_id_per_token,
+                num_valid_rows=num_valid_rows,
                 block_table=block_table,
                 topk_indices=logical_topk_indices,
                 block_size=block_size,
@@ -1072,7 +1052,7 @@ def _resolve_and_stage_group(
 
     def _stage_this_layer(self, shared_rows: slice) -> None:
         group = self.index_group
-        compute_stream = current_stream()
+        compute_stream = torch.accelerator.current_stream()
         group.copy_stream.wait_stream(compute_stream)
         with group.copy_stream:
             self._swap_rows(shared_rows)
@@ -1084,7 +1064,7 @@ def _stage_this_layer(self, shared_rows: slice) -> None:
     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)
+        torch.accelerator.current_stream().wait_event(self._layer_ready_event)
         self._swap_staged = False
 
     def swap_in(
@@ -1135,6 +1115,43 @@ def swap_in(
         return physical_topk_indices
 
 
+class HiSparseSlotInvalidator:
+    """Drop hot copies of rewritten host slots for many index groups at once."""
+
+    def __init__(self, request_state_indices: torch.Tensor) -> None:
+        self.request_state_indices = request_state_indices
+        self._group_tables: dict[tuple[int, ...], torch.Tensor] = {}
+
+    def __call__(
+        self,
+        leaders: list[HiSparseRuntime],
+        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 or not leaders:
+            return
+        tables = [leader.index_group.device_global_indices for leader in leaders]
+        key = tuple(table.data_ptr() for table in tables)
+        group_table = self._group_tables.get(key)
+        if group_table is None:
+            if any(table.shape != tables[0].shape for table in tables):
+                raise ValueError("HiSparse index groups must share one hot layout.")
+            group_table = async_tensor_h2d(
+                list(key), self.request_state_indices.device, torch.int64
+            )
+            self._group_tables[key] = group_table
+        num_state_rows, region_stride = tables[0].shape
+        leaders[0].cache_ops.hisparse_invalidate_written_slots(
+            group_table,
+            self.request_state_indices,
+            req_id_per_token[:num_tokens],
+            written_slots[:num_tokens],
+            num_state_rows,
+            region_stride,
+        )
+
+
 def update_hisparse_residency(
     residency: torch.Tensor,
     updates: Mapping[str, SparseKVResidencyUpdate],
@@ -1143,13 +1160,13 @@ def update_hisparse_residency(
 ) -> None:
     """Apply scheduled requests' residency changes.
 
-    ``residency`` is ``[state rows, resident groups, pages]``: the GPU block of
+    ``residency`` is ``[resident groups, state rows, pages]``: the GPU block of
     each page, or block 0 when the page is read from the host.
     """
     if not updates:
         return
     batch_rows = {request_id: row for row, request_id in enumerate(request_ids)}
-    num_groups = residency.shape[1]
+    num_groups = residency.shape[0]
     rows: list[int] = []
     pages: list[int] = []
     block_ids: list[list[int]] = [[] for _ in range(num_groups)]
@@ -1170,7 +1187,7 @@ def update_hisparse_residency(
         host_values, device=request_state_indices.device, dtype=torch.int32
     )
     state_rows = request_state_indices[values[:, 0]]
-    residency[state_rows, :, values[:, 1]] = values[:, 2:]
+    residency[:, state_rows, values[:, 1]] = values[:, 2:].T
 
 
 class HiSparseCacheHandle:
@@ -1195,7 +1212,6 @@ def __init__(self, runtime: HiSparseRuntime) -> None:
         # Speculator layers write their rows after the target forward.
         self.draft_layer = False
         self.index_group_caches: list[HiSparseCacheHandle] = [self]
-        self._batch_block_table: torch.Tensor | None = None
 
     def prepare_group_for_batch(self, attn_metadata: Any | None) -> None:
         assert self.runtime.is_group_leader
@@ -1204,7 +1220,6 @@ def prepare_group_for_batch(self, attn_metadata: Any | None) -> None:
 
     def _prepare_for_batch(self, attn_metadata: Any | None) -> None:
         self.dummy_batch = attn_metadata is None
-        self._batch_block_table = None
         self.runtime.begin_forward()
         self.num_actual_tokens = (
             attn_metadata.num_actual_tokens if attn_metadata is not None else 0
@@ -1220,17 +1235,6 @@ def _prepare_for_batch(self, attn_metadata: Any | None) -> None:
             not self.decode_batch or self.runtime.eager_host_mirror
         )
 
-    def batch_block_table(self) -> torch.Tensor:
-        """Resident rows by this step's batch row, for the prefill paths."""
-        if self._batch_block_table is None:
-            assert self.block_table is not None
-            indices = self.runtime.request_state_indices
-            assert indices is not None
-            self._batch_block_table = self.block_table.index_select(
-                0, indices.clamp(min=0)
-            )
-        return self._batch_block_table
-
     def write_target(
         self, num_input_rows: int, num_slot_rows: int
     ) -> tuple[torch.Tensor, torch.Tensor, int]:
diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py
index 7a6650ce5025..ff4f08f6bb32 100644
--- a/vllm/v1/kv_cache_interface.py
+++ b/vllm/v1/kv_cache_interface.py
@@ -106,17 +106,6 @@ def is_ultraquant(self) -> bool:
         """True for the UltraQuant 4-bit KV-cache format."""
         return self == KVQuantMode.ULTRAQUANT_4BIT
 
-    @property
-    def is_custom_mode(self) -> bool:
-        """True when a platform backend fully self-manages kernel dispatch.
-
-        The attention selector should defer to the backend's own kernel
-        selection rather than a generic per-token-head / fp8 / turboquant
-        path. Distinct from NONE (which means non-quantized): CUSTOM is
-        quantized, just with no generic kernel mode.
-        """
-        return self == KVQuantMode.CUSTOM
-
 
 def get_kv_quant_mode(kv_cache_dtype: str) -> KVQuantMode:
     """Map a ``kv_cache_dtype`` string to a :class:`KVQuantMode`."""
@@ -1068,6 +1057,7 @@ def merge(cls, specs: list[Self]) -> Self:
             dtype=specs[0].dtype,
             dcp_sharded=specs[0].dcp_sharded,
             max_tp_shards=specs[0].max_tp_shards,
+            kv_quant_mode=specs[0].kv_quant_mode,
             page_size_padded=specs[0].page_size_padded,
             num_head_slots=specs[0].num_head_slots,
             state_content_bytes=specs[0].state_content_bytes,
diff --git a/vllm/v1/kv_offload/base.py b/vllm/v1/kv_offload/base.py
index a5d9995527e5..51cdf7b72b6d 100644
--- a/vllm/v1/kv_offload/base.py
+++ b/vllm/v1/kv_offload/base.py
@@ -200,6 +200,8 @@ class OffloadingEvent:
     removal_expected: bool = False
     # Failed allocations release metadata without announcing a public eviction.
     metadata_only: bool = False
+    # Originating request context for a stored event.
+    req_context: ReqContext | None = field(default=None, repr=False, compare=False)
 
 
 """
@@ -214,11 +216,6 @@ class OffloadingEvent:
         The given blocks will be protected from eviction.
         This function returns a LoadSpec which encapsulates
         information required for performing the load.
-    touch() - marks the give blocks as recently used. Can be used
-        to track block's LRU. This function is separated from the
-        prepare_load function to allow setting block recency even
-        for blocks which do not need reading from the cache, such as
-        blocks that are cached by the GPU prefix cache.
     complete_load() - mark blocks which were previously prepared to be
         loaded as done loading. This is to re-allow their eviction.
     prepare_store() - prepare the given blocks to be written.
@@ -376,17 +373,6 @@ def prepare_load(
         """
         pass
 
-    def touch(self, keys: Collection[OffloadKey], req_context: ReqContext):
-        """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.
 
diff --git a/vllm/v1/kv_offload/cpu/gpu_worker.py b/vllm/v1/kv_offload/cpu/gpu_worker.py
index 9d20b7825af8..199632d053a3 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
@@ -36,6 +37,16 @@
 logger = init_logger(__name__)
 
 
+def _clear_pending_device_error() -> None:
+    """Clear the error a failed copy left pending for the next kernel launch."""
+    if not current_platform.is_cuda_alike():
+        return
+    try:
+        CudaRTLibrary().cudaGetLastError()
+    except (AssertionError, AttributeError, OSError):
+        logger.warning("Could not clear the pending device error", exc_info=True)
+
+
 def _select_swap_blocks_fn(
     layer_refs_per_group: list[list[CanonicalKVCacheRef]],
     gpu_to_cpu: bool,
@@ -75,6 +86,7 @@ class Transfer:
     batch_src: torch.Tensor
     batch_dst: torch.Tensor
     batch_sizes: torch.Tensor
+    failed: bool = False
 
 
 def compute_sub_block_ptrs(
@@ -704,16 +716,40 @@ def transfer_async(
         # writing; we must keep STREAM ordering so source reads are gated
         # by the transfer stream's wait_stream(compute) barrier.
         is_src_access_order_any = not self.gpu_to_cpu
-        with current_platform.stream(stream):
-            start_event.record(stream)
-            if op_idx > 0:
-                self._swap_blocks_batch(
-                    src,
-                    dst,
-                    sizes,
-                    is_src_access_order_any=is_src_access_order_any,
-                )
+        try:
+            with current_platform.stream(stream):
+                start_event.record(stream)
+                if op_idx > 0:
+                    self._swap_blocks_batch(
+                        src,
+                        dst,
+                        sizes,
+                        is_src_access_order_any=is_src_access_order_any,
+                    )
+                end_event.record(stream)
+        except RuntimeError:
+            _clear_pending_device_error()
+            logger.exception(
+                "KV offload %s transfer failed for job %d",
+                "GPU->CPU" if self.gpu_to_cpu else "CPU->GPU",
+                job_id,
+            )
             end_event.record(stream)
+            self._transfer_events[job_id] = end_event
+            self._transfers.append(
+                Transfer(
+                    job_id=job_id,
+                    stream=stream,
+                    start_event=start_event,
+                    end_event=end_event,
+                    num_bytes=0,
+                    batch_src=batch_src,
+                    batch_dst=batch_dst,
+                    batch_sizes=batch_sizes,
+                    failed=True,
+                )
+            )
+            return True
 
         self._transfer_events[job_id] = end_event
         self._transfers.append(
@@ -734,33 +770,63 @@ def transfer_async(
 
     def get_finished(self) -> list[TransferResult]:
         results: list[TransferResult] = []
-        while self._transfers and self._transfers[0].end_event.query():
-            transfer = self._transfers.popleft()
-            transfer_time = (
-                transfer.start_event.elapsed_time(transfer.end_event) * 1e-3
-            )  # elapsed_time is in milliseconds
-            result = TransferResult(
-                job_id=transfer.job_id,
-                success=True,
-                transfer_size=transfer.num_bytes,
-                transfer_time=transfer_time,
-            )
+        while self._transfers:
+            transfer = self._transfers[0]
+            try:
+                if not transfer.end_event.query():
+                    break
+                transfer_time = (
+                    None
+                    if transfer.failed
+                    # elapsed_time is in milliseconds
+                    else transfer.start_event.elapsed_time(transfer.end_event) * 1e-3
+                )
+            except RuntimeError:
+                logger.exception(
+                    "KV offload %s transfer failed for job %d while polling "
+                    "for completion",
+                    "GPU->CPU" if self.gpu_to_cpu else "CPU->GPU",
+                    transfer.job_id,
+                )
+                self._transfers.popleft()
+                self._transfer_events.pop(transfer.job_id, None)
+                results.append(TransferResult(job_id=transfer.job_id, success=False))
+                continue
+
+            self._transfers.popleft()
+            self._transfer_events.pop(transfer.job_id, None)
+            if transfer.failed:
+                results.append(TransferResult(job_id=transfer.job_id, success=False))
+                continue
 
-            results.append(result)
+            results.append(
+                TransferResult(
+                    job_id=transfer.job_id,
+                    success=True,
+                    transfer_size=transfer.num_bytes,
+                    transfer_time=transfer_time,
+                )
+            )
             self._stream_pool.append(transfer.stream)
             self._event_pool.append(transfer.end_event)
             self._event_pool.append(transfer.start_event)
             self._buffer_pool.append(
                 (transfer.batch_src, transfer.batch_dst, transfer.batch_sizes)
             )
-            del self._transfer_events[transfer.job_id]
         return results
 
     def wait(self, job_ids: set[int]):
         for job_id in job_ids:
             event = self._transfer_events.get(job_id)
             if event is not None:
-                event.synchronize()
+                try:
+                    event.synchronize()
+                except RuntimeError:
+                    logger.exception(
+                        "KV offload transfer for job %d failed while waiting "
+                        "for it to complete",
+                        job_id,
+                    )
 
     def shutdown(self) -> None:
         """Drain this direction and release its transfer-side resources."""
diff --git a/vllm/v1/kv_offload/cpu/manager.py b/vllm/v1/kv_offload/cpu/manager.py
index 4aa8df48e762..4ceaf3fc7d0a 100644
--- a/vllm/v1/kv_offload/cpu/manager.py
+++ b/vllm/v1/kv_offload/cpu/manager.py
@@ -240,10 +240,6 @@ def _prepare_load(
             self._record_request_cache_access(keys, req_context, reused_keys=keys)
         return self._get_load_store_spec(keys, chunks)
 
-    @override
-    def touch(self, keys: Collection[OffloadKey], req_context: ReqContext) -> None:
-        self._policy.touch(keys, req_context)
-
     @override
     def complete_load(
         self, keys: Collection[OffloadKey], req_context: ReqContext
@@ -415,6 +411,7 @@ def complete_store(
                     keys=stored_keys,
                     medium=self.medium,
                     removed=False,
+                    req_context=req_context,
                 )
             )
 
diff --git a/vllm/v1/kv_offload/cpu/policies/arc.py b/vllm/v1/kv_offload/cpu/policies/arc.py
index 292c82ac8a4c..e7b2a91ac8de 100644
--- a/vllm/v1/kv_offload/cpu/policies/arc.py
+++ b/vllm/v1/kv_offload/cpu/policies/arc.py
@@ -95,24 +95,6 @@ def on_store_miss(
         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)):
-            if key in self.t1:
-                chunk = self.t1.pop(key)
-                if not chunk.is_ready:
-                    # chunk was just prepared to be stored, not really touched
-                    # twice — keep it in T1 and mark as most recently used
-                    self.t1[key] = chunk
-                else:
-                    self.t2[key] = chunk
-
-            elif key in self.t2:
-                self.t2.move_to_end(key)
-
-            else:
-                self._adapt_to_ghost_hit(key)
-
     @override
     def on_request_finished(
         self,
diff --git a/vllm/v1/kv_offload/cpu/policies/base.py b/vllm/v1/kv_offload/cpu/policies/base.py
index eb05d4d45cec..5d5792218a82 100644
--- a/vllm/v1/kv_offload/cpu/policies/base.py
+++ b/vllm/v1/kv_offload/cpu/policies/base.py
@@ -72,26 +72,15 @@ def insert(
     def remove(self, key: OffloadKey) -> None:
         """Remove a chunk (used to clean up after a failed store)."""
 
-    @abstractmethod
-    def touch(self, keys: Iterable[OffloadKey], req_context: ReqContext) -> None:
-        """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.
+        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)
+        return
 
     def on_request_access(
         self,
@@ -106,6 +95,7 @@ def on_request_access(
         """
         return
 
+    @abstractmethod
     def on_request_finished(
         self,
         key_groups: Sequence[Sequence[OffloadKey]],
@@ -120,15 +110,7 @@ def on_request_finished(
         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(
diff --git a/vllm/v1/kv_offload/cpu/policies/lru.py b/vllm/v1/kv_offload/cpu/policies/lru.py
index f5b070e13a29..c36417e95f92 100644
--- a/vllm/v1/kv_offload/cpu/policies/lru.py
+++ b/vllm/v1/kv_offload/cpu/policies/lru.py
@@ -125,10 +125,6 @@ def remove(self, key: OffloadKey) -> None:
         self._ranks.pop(key, None)
         self._evictable.discard(key)
 
-    @override
-    def touch(self, keys: Iterable[OffloadKey], req_context: ReqContext) -> None:
-        self._update_recency(reversed(list(keys)))
-
     @override
     def on_request_access(
         self,
diff --git a/vllm/v1/kv_offload/cpu/shared_offload_region.py b/vllm/v1/kv_offload/cpu/shared_offload_region.py
index 3457a34ff023..8eb45e23370b 100644
--- a/vllm/v1/kv_offload/cpu/shared_offload_region.py
+++ b/vllm/v1/kv_offload/cpu/shared_offload_region.py
@@ -28,6 +28,9 @@
 # Page-cache inserts are serialized per shm file, capping pre-faulting of one
 # file near 10 GiB/s; smaller segments are not worth a file.
 _MIN_SEGMENT_BYTES = 1 << 30
+# Split the shared region into this many shm files so workers pre-fault them
+# in parallel (page-cache inserts are serialized per file).
+REGION_SEGMENTS = 32
 
 
 def _wait_for_file_size(fd: int, expected_size: int, timeout: float = 30.0) -> None:
@@ -127,15 +130,18 @@ def __init__(
 
         self.mmap_path = f"/dev/shm/vllm_offload_{engine_id}.mmap"
         num_segments = min(num_segments, self.total_size_bytes // _MIN_SEGMENT_BYTES)
-        rows = max(1, cdiv(num_chunks, max(1, num_segments)))
+        alignment = 2 << 20 if num_chunks == 1 else self._row_stride
+        segment_bytes = round_up(
+            cdiv(self.total_size_bytes, max(1, num_segments)), alignment
+        )
         # (path, offset, length) per file.
         self._segments = [
             (
                 f"{self.mmap_path}.{i}" if i else self.mmap_path,
-                row * self._row_stride,
-                min(rows, num_chunks - row) * self._row_stride,
+                offset,
+                min(segment_bytes, self.total_size_bytes - offset),
             )
-            for i, row in enumerate(range(0, max(num_chunks, 1), rows))
+            for i, offset in enumerate(range(0, self.total_size_bytes, segment_bytes))
         ]
         self._creator = False  # set True only if this worker creates the file
         self.rank = rank
@@ -181,13 +187,26 @@ def __init__(
                 )
 
             single = len(self._segments) == 1
-            self.mmap_obj: mmap.mmap | None = mmap.mmap(
-                self.fd if single else -1,
-                self.total_size_bytes,
-                flags=mmap.MAP_SHARED if single else mmap.MAP_PRIVATE,
-                prot=mmap.PROT_READ | mmap.PROT_WRITE,
-            )
-            if not single:
+            self.mmap_obj: mmap.mmap | None = None
+            if single:
+                self.mmap_obj = mmap.mmap(
+                    self.fd,
+                    self.total_size_bytes,
+                    flags=mmap.MAP_SHARED,
+                    prot=mmap.PROT_READ | mmap.PROT_WRITE,
+                )
+            else:
+                reservation = os.memfd_create("vllm_offload_reservation")
+                try:
+                    os.ftruncate(reservation, self.total_size_bytes)
+                    self.mmap_obj = mmap.mmap(
+                        reservation,
+                        self.total_size_bytes,
+                        flags=mmap.MAP_SHARED,
+                        prot=mmap.PROT_READ | mmap.PROT_WRITE,
+                    )
+                finally:
+                    os.close(reservation)
                 self._map_segments()
             if barrier is not None:
                 barriers_remaining -= 1
diff --git a/vllm/v1/kv_offload/cpu/spec.py b/vllm/v1/kv_offload/cpu/spec.py
index 23b4f610de2c..7d47bbebe11a 100644
--- a/vllm/v1/kv_offload/cpu/spec.py
+++ b/vllm/v1/kv_offload/cpu/spec.py
@@ -20,11 +20,10 @@
 from vllm.v1.kv_offload.cpu.common import CPUOffloadingInfo, CPUOffloadingMetrics
 from vllm.v1.kv_offload.cpu.gpu_worker import CPUOffloadingWorker
 from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager
-from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion
-
-# Split the shared region into this many shm files so workers pre-fault them
-# in parallel (page-cache inserts are serialized per file).
-_REGION_SEGMENTS = 32
+from vllm.v1.kv_offload.cpu.shared_offload_region import (
+    REGION_SEGMENTS,
+    SharedOffloadRegion,
+)
 
 
 def _all_workers_barrier() -> None:
@@ -197,7 +196,7 @@ def create_worker(self, kv_caches: CanonicalKVCaches) -> CPUOffloadingWorker:
                 kv_bytes_per_chunk=self.kv_bytes_per_chunk,
                 cpu_page_size=self.cpu_page_size_per_worker,
                 barrier=_all_workers_barrier,
-                num_segments=_REGION_SEGMENTS,
+                num_segments=REGION_SEGMENTS,
                 populate_shard=(
                     (worker_index, world_size) if self.replicated_layout else None
                 ),
diff --git a/vllm/v1/kv_offload/tiering/async_lookup.py b/vllm/v1/kv_offload/tiering/async_lookup.py
index 69a513f55de8..080c9139433e 100644
--- a/vllm/v1/kv_offload/tiering/async_lookup.py
+++ b/vllm/v1/kv_offload/tiering/async_lookup.py
@@ -33,7 +33,8 @@
 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.
+the same probe; a successful store drops such an orphan instead (see
+mark_present()).
 """
 
 import queue
@@ -61,8 +62,9 @@ class LookupPhase(Enum):
 class LookupState:
     generation: int
     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.
+    # None while pending/in flight; True if the key is known to exist (a probe
+    # found it, or this instance stored it); 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
 
@@ -229,6 +231,34 @@ def mark_miss(self, keys: Collection[OffloadKey]) -> None:
             if state is not None and state.phase is LookupPhase.RESOLVED:
                 state.result = False
 
+    def mark_present(self, keys: Collection[OffloadKey]) -> None:
+        """Record that keys were successfully stored to this tier.
+
+        Unlike mark_miss(), which leaves a newer probe alone (#59096), a store
+        completion is newer information than any probe started before it, so
+        this supersedes them: advancing the generation makes a stale pending
+        batch entry fail flush()'s generation check and a stale in-flight
+        result fail drain_results()'s, before either phase assert is reached.
+
+        Keys with no cached entry are skipped, since a verdict with no owning
+        request could never be reclaimed. An in-flight entry whose requests all
+        finished is dropped; cleanup() already removed the reverse-index
+        entries, so its late result is simply ignored.
+        """
+        for key in keys:
+            state = self._lookup_state.get(key)
+            if state is None or state.result is True:
+                continue
+            if not state.request_ids:
+                del self._lookup_state[key]
+                continue
+            # Generations are opaque identity tokens compared for equality
+            # only, so bumping one key cannot affect another key's results.
+            state.generation = self._next_generation
+            self._next_generation += 1
+            state.result = True
+            state.phase = LookupPhase.RESOLVED
+
     def cleanup(self, req_id: str) -> None:
         """Release request references, retaining in-flight lookups.
 
diff --git a/vllm/v1/kv_offload/tiering/base.py b/vllm/v1/kv_offload/tiering/base.py
index 3e9fccf4c740..e4dded94fc0f 100644
--- a/vllm/v1/kv_offload/tiering/base.py
+++ b/vllm/v1/kv_offload/tiering/base.py
@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Abstract interfaces and data types for the secondary tiering layer."""
 
+import threading
 import time
 from abc import ABC, abstractmethod
 from collections.abc import Collection, Iterable
@@ -66,6 +67,8 @@ class TieringOffloadingMetrics:
     BACKPRESSURE_STORES_DROPPED = "vllm:kv_offload_tiering_backpressure_stores_dropped"
     BACKPRESSURE_BLOCKS_DROPPED = "vllm:kv_offload_tiering_backpressure_blocks_dropped"
     PROMOTION_LATENCY = "vllm:kv_offload_tiering_promotion_latency_seconds"
+    DEFERRED_RELEASE_CHUNKS = "vllm:kv_offload_tiering_deferred_release_chunks"
+    DEFERRED_RELEASE_HOLD = "vllm:kv_offload_tiering_deferred_release_hold_seconds"
 
 
 @dataclass
@@ -93,6 +96,13 @@ class JobResult:
     successful_keys: Collection[OffloadKey] | None = None
     transfer_time: float | None = None
     transfer_bytes: int | None = None
+    # Only applicable to failed promotion jobs. Set when the tier's I/O may
+    # still write into the primary chunks of the failed keys: the manager fails
+    # those keys now but defers releasing their chunks to the free pool until
+    # the tier sets the event. The tier may set it from any thread, once its
+    # I/O can no longer touch those chunks; the manager applies the release on
+    # its own polling path.
+    deferred_release: threading.Event | None = None
 
 
 class ParentManager(ABC):
@@ -302,16 +312,6 @@ def take_events(self) -> Iterable[OffloadingEvent]:
         """Take KV events for storage state owned by this tier."""
         return ()
 
-    def touch(self, keys: Collection[OffloadKey], req_context: ReqContext):
-        """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.
@@ -329,7 +329,7 @@ def on_request_finished(self, req_context: ReqContext) -> None:
         """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
+        (submit_store, submit_load) have already been issued, and none
         will follow. Note this does NOT imply the tier's transfers have
         completed: jobs already submitted may still be in flight and will
         report via get_finished_jobs(). This is the right place to release
@@ -384,6 +384,10 @@ def drain_jobs(self) -> None:
         would corrupt either the primary memoryview or the secondary backing
         store. Queued (not-yet-started) transfers may be cancelled, but their
         failure result must still appear in `get_finished_jobs()`.
+
+        Likewise, every ``JobResult.deferred_release`` event must be set
+        before this returns, so the tier must bound how long it can defer a
+        release.
         """
         pass
 
diff --git a/vllm/v1/kv_offload/tiering/fs/manager.py b/vllm/v1/kv_offload/tiering/fs/manager.py
index 6522cfd23646..91a1dd755138 100644
--- a/vllm/v1/kv_offload/tiering/fs/manager.py
+++ b/vllm/v1/kv_offload/tiering/fs/manager.py
@@ -150,8 +150,8 @@ def __init__(
                     "emit events.",
                     tier_type,
                 )
-        # Keys of in-flight store jobs, tracked only when events are enabled.
-        self._store_job_keys: dict[JobId, list[OffloadKey]] = {}
+        # Track in-flight store jobs for lookup invalidation and KV events.
+        self._store_jobs: dict[JobId, TransferJob] = {}
         # 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]] = {}
@@ -222,8 +222,7 @@ def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult:
     @override
     def submit_store(self, job_metadata: TransferJob) -> None:
         keys = list(job_metadata.keys)
-        if self.events is not None:
-            self._store_job_keys[job_metadata.job_id] = keys
+        self._store_jobs[job_metadata.job_id] = job_metadata
         task = functools.partial(
             batch_store_block,
             [self.file_mapper.get_file_name(key) for key in keys],
@@ -283,15 +282,18 @@ def get_finished_jobs(self) -> Iterable[JobResult]:
         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:
+            job = self._store_jobs.pop(job_id, None)
+            if job is not None and success:
+                # Successful store makes previously absent keys known-present.
+                self._lookup_manager.mark_present(job.keys)
+                if self.events is not None and job.keys:
                     self.events.append(
                         OffloadingEvent(
-                            keys=keys,
+                            keys=list(job.keys),
                             medium=self.medium,
                             removed=False,
                             locality=self.locality,
+                            req_context=job.req_context,
                         )
                     )
             load_keys = self._load_job_keys.pop(job_id, None)
diff --git a/vllm/v1/kv_offload/tiering/kvcr/manager.py b/vllm/v1/kv_offload/tiering/kvcr/manager.py
index 1f3d9638b493..d85f6403203b 100644
--- a/vllm/v1/kv_offload/tiering/kvcr/manager.py
+++ b/vllm/v1/kv_offload/tiering/kvcr/manager.py
@@ -11,7 +11,6 @@
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, ClassVar, cast
 
-import msgspec
 from kvcr import (
     DURATION_METRIC,
     KVCR,
@@ -20,6 +19,7 @@
     TRANSFER_BLOCKS_METRIC,
     TRANSFER_BYTES_METRIC,
     KVCRBindings,
+    ZmqPeerControlChannel,
 )
 from kvcr.config import (
     G3Options,
@@ -30,7 +30,6 @@
     LocalDramOptions,
     RemoteFWDramOptions,
 )
-from kvcr.control_channels import ZmqPeerControlChannel
 from kvcr.policy import (
     FIFOPolicy,
     G3FIFOPolicy,
@@ -201,20 +200,17 @@ def _resolve_pin(
 
             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):
+            if len(job_keys) != len(job.chunk_ids) or (
+                job_keys is not hit_keys and set(job_keys) != set(hit_keys)
+            ):
                 return None
 
-            descriptors: dict[BlockKey, list[MemoryRef] | 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)
+            descriptors: dict[BlockKey, list[MemoryRef] | None] = dict.fromkeys(
+                cast(tuple[BlockKey, ...], offload_keys)
             )
+            make_descriptor = self._tier._make_descriptor
+            for key, chunk_id in zip(job_keys, job.chunk_ids.tolist()):
+                descriptors[BlockKey(key)] = [make_descriptor(chunk_id)]
             pin_handle = f"kvcr-job-{job.job_id}"
             self._pin_jobs[pin_handle] = job.job_id
             return pin_handle, descriptors
@@ -476,6 +472,7 @@ def __init__(
             if local_mapping is not None:
                 local_mapping.close()
             raise
+        self._nixl_agent_name = self._kvcr.config.nixl_agent_name
         self._local_dram_mmap = local_mapping
         self._finished_jobs: list[JobResult] = []
         self._jobs_by_op: dict[OpHandle, _JobState] = {}
@@ -493,9 +490,9 @@ def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult:
     @override
     def submit_load(self, job_metadata: TransferJob) -> None:
         blocks = {
-            self._key_adapter.encode(key): [self._make_descriptor(int(chunk_id))]
+            self._key_adapter.encode(key): [self._make_descriptor(chunk_id)]
             for key, chunk_id in zip(
-                job_metadata.keys, job_metadata.chunk_ids, strict=True
+                job_metadata.keys, job_metadata.chunk_ids.tolist(), strict=True
             )
         }
         if not blocks:
@@ -517,9 +514,9 @@ def submit_load(self, job_metadata: TransferJob) -> None:
     @override
     def submit_store(self, job_metadata: TransferJob) -> None:
         blocks = {
-            self._key_adapter.encode(key): [self._make_descriptor(int(chunk_id))]
+            self._key_adapter.encode(key): [self._make_descriptor(chunk_id)]
             for key, chunk_id in zip(
-                job_metadata.keys, job_metadata.chunk_ids, strict=True
+                job_metadata.keys, job_metadata.chunk_ids.tolist(), strict=True
             )
         }
         if not blocks:
@@ -566,7 +563,7 @@ def _poll_finished_jobs(self) -> list[JobResult]:
             job, 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:
+            if successful_keys is not None and (remaining > 0 or not success):
                 successful_keys.update(
                     OffloadKey(bytes(key))
                     for key, entry in entries.items()
@@ -620,14 +617,10 @@ 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),
+                hints=hint,
             )
         return RequestOffloadingContext()
 
-    @override
-    def touch(self, keys: Collection[OffloadKey], req_context: ReqContext) -> None:
-        self._align_sequence(keys, req_context, use_current_time=True)
-
     def _align_sequence(
         self,
         keys: Collection[OffloadKey],
@@ -666,7 +659,7 @@ def on_request_finished(self, req_context: ReqContext) -> None:
             for key, position in req_context._offload_key_positions.items()
             if processed_tokens is None or position <= processed_tokens
         ]
-        self.touch(keys, req_context)
+        self._align_sequence(keys, req_context, use_current_time=True)
         self._kvcr.discard_hint(req_context.req_id)
 
     @override
@@ -706,6 +699,6 @@ def _record_inventory(self, event: InventoryEvent) -> None:
 
     def _make_descriptor(self, chunk_id: int) -> MemoryRef:
         return MemoryRef(
-            end_point_name=self._kvcr.config.nixl_agent_name,
+            end_point_name=self._nixl_agent_name,
             element_index=chunk_id,
         )
diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py
index 74b6bd7c6b66..77850e480ad9 100644
--- a/vllm/v1/kv_offload/tiering/manager.py
+++ b/vllm/v1/kv_offload/tiering/manager.py
@@ -75,15 +75,13 @@
 # Warn after this much continuous failure to acquire the lock. Above the 5 s
 # warning drain_jobs() emits, so a slow reset_cache() does not warn twice.
 _TIER_POLL_STALL_WARN_S = 30.0
-# Give up after rounds have failed continuously for this long. Counted in time,
-# not rounds: at the default 1 ms interval a round budget is spent in a blink,
-# so a transient fault (a peer reconnecting) would turn the thread off for
-# good. on_schedule_end() still serves the tiers, so the fallback is per-step
-# servicing rather than none.
-_TIER_POLL_GIVE_UP_AFTER_S = 30.0
-# Minimum spacing between logged tracebacks for a run of failing rounds, which
-# would otherwise log one per interval.
-_TIER_POLL_ERROR_LOG_INTERVAL_S = 1.0
+# Back off after this much continuous failure (time, not rounds, so a brief
+# fault doesn't trip it). Never exit: an idle engine would leave peers unserved.
+_TIER_POLL_BACK_OFF_AFTER_S = 30.0
+# Minimum wait between rounds while backing off.
+_TIER_POLL_BACKOFF_S = 1.0
+# Minimum spacing between failure logs.
+_TIER_POLL_ERROR_LOG_INTERVAL_S = 30.0
 # How long shutdown() waits for the thread to finish its round.
 _TIER_POLL_JOIN_TIMEOUT_S = 5.0
 
@@ -111,6 +109,13 @@ class JobMetadata(NamedTuple):
     tier_idx: int
 
 
+class _DeferredRelease(NamedTuple):
+    event: threading.Event
+    tier_idx: int
+    chunk_ids: list[int]
+    deferred_at: float
+
+
 class CPUPrimaryTierOffloadingManager(CPUOffloadingManager):
     """CPUOffloadingManager with a primary/secondary transfer interface.
 
@@ -144,6 +149,9 @@ def __init__(
         self.complete_read = self.complete_load
         self.prepare_write = self.prepare_store
         self.complete_write = self.complete_store
+        # Chunks of failed writes the writer may still be filling; their
+        # release to the free pool is deferred until release_deferred_write().
+        self._deferred_release_chunk_ids: set[int] = set()
 
         self._kv_memoryview = mmap_region.create_kv_memoryview()
 
@@ -157,6 +165,54 @@ def prepare_read(
         """
         return self._prepare_load(keys, req_context, record_access=False)
 
+    def defer_write_release(self, keys: Collection[OffloadKey]) -> list[int]:
+        """Fail pending writes like ``complete_write(False)``, deferring the release.
+
+        The keys leave the cache (lookups miss and a later store allocates a
+        fresh chunk), while the chunks stay counted as write-pending and out
+        of the free pool. Returns the deferred chunk ids, to be passed to
+        ``release_deferred_write`` once the writer can no longer touch them.
+        """
+        chunk_ids: list[int] = []
+        for key in keys:
+            chunk = self._policy.get(key)
+            # Every key must still be pending, or a chunk the tier may still
+            # write to would go unheld.
+            assert chunk is not None and not chunk.is_ready, (
+                f"defer_write_release: key {key!r} is not a pending write"
+            )
+            self._policy.remove(key)
+            self._deferred_release_chunk_ids.add(chunk.chunk_id)
+            chunk_ids.append(chunk.chunk_id)
+        if chunk_ids and self.events is not None:
+            self.events.append(
+                OffloadingEvent(
+                    keys=list(keys),
+                    medium=self.medium,
+                    removed=True,
+                    metadata_only=True,
+                )
+            )
+        return chunk_ids
+
+    def release_deferred_write(self, chunk_ids: Iterable[int]) -> None:
+        """Complete a deferred release: return the chunks to the free pool."""
+        for chunk_id in chunk_ids:
+            self._deferred_release_chunk_ids.remove(chunk_id)
+            self._num_write_pending_chunks -= 1
+            self._free_list.append(chunk_id)
+
+    @property
+    def num_deferred_release_chunks(self) -> int:
+        return len(self._deferred_release_chunk_ids)
+
+    @override
+    def reset_cache(self) -> None:
+        assert not self._deferred_release_chunk_ids, (
+            "reset_cache with deferred-release chunks: a tier may still write to them"
+        )
+        super().reset_cache()
+
     def get_kv_memoryview(self) -> memoryview:
         """Return the memoryview over the primary tier's KV cache buffer.
 
@@ -250,6 +306,9 @@ def __init__(
         #   True:  secondary → primary (promotion)
         #   False: primary → secondary (cascade)
         self._jobs: dict[JobId, JobMetadata] = {}
+        # Finished promotions whose failed chunks await release until the tier
+        # sets their JobResult.deferred_release event.
+        self._deferred_releases: dict[JobId, _DeferredRelease] = {}
         # Set by reset_cache() until on_reset_flushed(); see lookup().
         self._promotions_held = False
         primary_view = self.primary_tier.get_kv_memoryview()
@@ -346,9 +405,12 @@ def _stop_tier_poll_thread(self) -> None:
 
     def _tier_poll_loop(self) -> None:
         stalled_since: float | None = None
-        # Start of the current run of failed rounds, and when one last logged.
+        # Current failure run: start, count, and last log time.
         failing_since: float | None = None
+        failures = 0
         last_error_log = 0.0
+        backing_off = False
+        backoff_s = max(self._tier_poll_interval_s, _TIER_POLL_BACKOFF_S)
         while not self._tier_poll_stop.is_set():
             if not self._lock.acquire(timeout=_TIER_POLL_LOCK_TIMEOUT_S):
                 now = time.monotonic()
@@ -367,28 +429,45 @@ def _tier_poll_loop(self) -> None:
             try:
                 self.poll_tiers()
                 failing_since = None
+                failures = 0
             except Exception:
-                # Never let a round kill the thread silently: on_schedule_end()
-                # still serves these tiers, so persistent failure degrades to
-                # per-step servicing instead of stopping the tier-poll thread.
+                # Keep retrying so peers are served while idle. Rate-limit
+                # logs: traceback on the first failure, a summary after.
                 now = time.monotonic()
                 if failing_since is None:
                     failing_since = now
+                failures += 1
                 if now - last_error_log >= _TIER_POLL_ERROR_LOG_INTERVAL_S:
                     last_error_log = now
-                    logger.exception("KV offload tier-poll round failed")
+                    if failures == 1:
+                        logger.exception("KV offload tier-poll round failed")
+                    else:
+                        logger.warning(
+                            "KV offload tier-poll thread has failed %d "
+                            "consecutive rounds (%.0fs); still retrying so "
+                            "peers are served while the engine is idle.",
+                            failures,
+                            now - failing_since,
+                        )
             finally:
                 self._lock.release()
             if (
                 failing_since is not None
-                and time.monotonic() - failing_since >= _TIER_POLL_GIVE_UP_AFTER_S
+                and time.monotonic() - failing_since >= _TIER_POLL_BACK_OFF_AFTER_S
             ):
-                logger.error(
-                    "KV offload tier-poll thread stopping after %.0fs of "
-                    "consecutive failures; falling back to per-step servicing.",
-                    _TIER_POLL_GIVE_UP_AFTER_S,
-                )
-                return
+                if not backing_off:
+                    logger.error(
+                        "KV offload tier-poll thread backing off to %.1fs "
+                        "after %.0fs of consecutive failures; it keeps "
+                        "retrying so peers are still served while the engine "
+                        "is idle.",
+                        backoff_s,
+                        _TIER_POLL_BACK_OFF_AFTER_S,
+                    )
+                    backing_off = True
+                self._tier_poll_stop.wait(backoff_s)
+                continue
+            backing_off = False
             # Yield unconditionally, and only after release(). Python locks are
             # not fair, so a release-then-reacquire loop could starve the
             # scheduler thread; this pause is what bounds its wait to one round.
@@ -491,7 +570,18 @@ def _complete_promotion(
                 transfer_job.req_context,
                 True,
             )
-        if failed_keys:
+        if not failed_keys:
+            assert completed_job.deferred_release is None
+        elif completed_job.deferred_release is not None:
+            chunk_ids = self.primary_tier.defer_write_release(failed_keys)
+            self._deferred_releases[completed_job.job_id] = _DeferredRelease(
+                event=completed_job.deferred_release,
+                tier_idx=job_metadata.tier_idx,
+                chunk_ids=chunk_ids,
+                deferred_at=time.monotonic(),
+            )
+            self._metrics.on_release_deferred(job_metadata.tier_idx, len(chunk_ids))
+        else:
             self.primary_tier.complete_write(
                 failed_keys,
                 transfer_job.req_context,
@@ -534,6 +624,9 @@ def _process_finished_jobs(
                     f" but belongs to tier #{job_metadata.tier_idx}"
                 )
                 transfer_job = job_metadata.transfer_job
+                assert (
+                    transfer_job.is_promotion or completed_job.deferred_release is None
+                )
                 self._metrics.on_job_finished(job_metadata, completed_job)
 
                 if transfer_job.is_promotion:
@@ -547,6 +640,22 @@ def _process_finished_jobs(
                         transfer_job.keys, transfer_job.req_context
                     )
                     self._update_backpressure(tier, job_metadata, completed_job)
+            if self._deferred_releases:
+                self._complete_deferred_releases(i)
+
+    def _complete_deferred_releases(self, tier_idx: int) -> None:
+        released = [
+            job_id
+            for job_id, deferred in self._deferred_releases.items()
+            if deferred.tier_idx == tier_idx and deferred.event.is_set()
+        ]
+        now = time.monotonic()
+        for job_id in released:
+            deferred = self._deferred_releases.pop(job_id)
+            self.primary_tier.release_deferred_write(deferred.chunk_ids)
+            self._metrics.on_deferred_release_done(
+                tier_idx, len(deferred.chunk_ids), now - deferred.deferred_at
+            )
 
     def _should_store_to_tier(
         self, tier: SecondaryTierManager, num_blocks: int
@@ -784,19 +893,6 @@ def prepare_load(
         """
         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.
-
-        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:
-            tier.touch(keys, req_context)
-
     @override
     def complete_load(self, keys: Collection[OffloadKey], req_context: ReqContext):
         """Mark chunks as done loading from primary tier to GPU.
@@ -1116,11 +1212,21 @@ def on_schedule_end(self, context: ScheduleEndContext) -> None:
 
     @override
     def has_pending_work(self) -> bool:
-        # In-flight primary<->secondary transfers (pending promotions are
-        # translated to transfer jobs in on_schedule_end), plus any work the
-        # secondary tiers themselves still have outstanding.
+        """Report whether outstanding work needs an engine step to finish.
+
+        Work a live tier-poll thread finishes on its own is excluded, so an
+        engine that only serves peers can go idle.
+        """
+        if self._polled_tiers:
+            # Without a live thread, or with another tier a peer lookup can
+            # promote from, only a step makes progress.
+            thread = self._tier_poll_thread
+            if thread is None or not thread.is_alive() or len(self.secondary_tiers) > 1:
+                return True
+            return any(state.pending_cascade_keys for state in self._req_state.values())
         return (
             bool(self._jobs)
+            or bool(self._deferred_releases)
             or any(state.pending_cascade_keys for state in self._req_state.values())
             or any(tier.has_pending_work() for tier in self.secondary_tiers)
         )
@@ -1167,6 +1273,9 @@ def reset_cache(self) -> None:
         # so manager bookkeeping is consistent before the primary reset.
         self._process_finished_jobs()
         assert not self._jobs
+        assert not self._deferred_releases, (
+            "drain_jobs() returned with deferred releases still pending"
+        )
 
         # Deferred promotion submissions reserve primary slots that the
         # reset below invalidates; their submit_load() has not yet been
diff --git a/vllm/v1/kv_offload/tiering/metrics.py b/vllm/v1/kv_offload/tiering/metrics.py
index 555077ec4305..1b28f67987d2 100644
--- a/vllm/v1/kv_offload/tiering/metrics.py
+++ b/vllm/v1/kv_offload/tiering/metrics.py
@@ -42,6 +42,7 @@ class _TierState:
     active_cascade_count: int = 0
     primary_write_chunk_count: int = 0
     primary_read_chunk_count: int = 0
+    deferred_release_chunk_count: int = 0
 
 
 class TieringMetricsTracker:
@@ -128,6 +129,21 @@ def on_job_finished(
             self._observe_promotion_latency(job_metadata, result)
         self._decrement_tier_state(job_metadata)
 
+    def on_release_deferred(self, tier_idx: int, chunk_count: int) -> None:
+        self._tier_states[tier_idx].deferred_release_chunk_count += chunk_count
+
+    def on_deferred_release_done(
+        self, tier_idx: int, chunk_count: int, hold_duration: float
+    ) -> None:
+        state = self._tier_states[tier_idx]
+        state.deferred_release_chunk_count -= chunk_count
+        assert state.deferred_release_chunk_count >= 0
+        self._stats.observe_histogram(
+            TieringOffloadingMetrics.DEFERRED_RELEASE_HOLD,
+            hold_duration,
+            self.tier_label(tier_idx),
+        )
+
     def on_promotion_allocation_failure(self) -> None:
         self._stats.increase_counter(
             TieringOffloadingMetrics.PROMOTION_ALLOCATION_FAILURES
@@ -184,6 +200,7 @@ def assert_idle(self) -> None:
             and state.active_cascade_count == 0
             and state.primary_write_chunk_count == 0
             and state.primary_read_chunk_count == 0
+            and state.deferred_release_chunk_count == 0
             for state in self._tier_states
         )
 
@@ -290,6 +307,11 @@ def _observe_active_transfer_stats(self, stats: OffloadingConnectorStats) -> Non
                 state.active_cascade_count,
                 labelvalues,
             )
+            stats.set_gauge(
+                TieringOffloadingMetrics.DEFERRED_RELEASE_CHUNKS,
+                state.deferred_release_chunk_count,
+                labelvalues,
+            )
 
     def _observe_resolved_lookup(
         self,
diff --git a/vllm/v1/kv_offload/tiering/obj/manager.py b/vllm/v1/kv_offload/tiering/obj/manager.py
index 37ce4ca3e244..fb18502332d9 100644
--- a/vllm/v1/kv_offload/tiering/obj/manager.py
+++ b/vllm/v1/kv_offload/tiering/obj/manager.py
@@ -149,8 +149,8 @@ def __init__(
                     "emit events.",
                     tier_type,
                 )
-        # Keys of in-flight store jobs, tracked only when events are enabled.
-        self._store_job_keys: dict[JobId, list[OffloadKey]] = {}
+        # Track in-flight store jobs for lookup invalidation and KV events.
+        self._store_jobs: dict[JobId, TransferJob] = {}
         # Keys of in-flight load (promotion) jobs, so a failed download can
         # mark its own cached lookup verdicts False (see get_finished_jobs).
         self._load_job_keys: dict[JobId, list[OffloadKey]] = {}
@@ -294,8 +294,7 @@ def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult:
         return LookupResult.HIT if result else LookupResult.MISS
 
     def submit_store(self, job_metadata: TransferJob) -> None:
-        if self.events is not None:
-            self._store_job_keys[job_metadata.job_id] = list(job_metadata.keys)
+        self._store_jobs[job_metadata.job_id] = job_metadata
         obj_keys = (self._file_mapper.get_file_name(k) for k in job_metadata.keys)
         self._submit_transfer(
             job_metadata.job_id, job_metadata.chunk_ids, obj_keys, NIXL_WRITE
@@ -337,8 +336,18 @@ def _poll_active_transfers(self) -> None:
 
             transfer_time = None
             if success:
-                telemetry = self._agent.get_xfer_telemetry(entry.xfer_handle)
-                transfer_time = telemetry.xferDuration / 1e6
+                # Telemetry is optional: some NIXL builds are compiled without
+                # it and raise here. A metrics failure must not fail a transfer
+                # that already completed.
+                try:
+                    telemetry = self._agent.get_xfer_telemetry(entry.xfer_handle)
+                except Exception:
+                    logger.warning_once(
+                        "NIXL transfer telemetry is unavailable; offload "
+                        "transfer timings will be omitted."
+                    )
+                else:
+                    transfer_time = telemetry.xferDuration / 1e6
 
             try:
                 self._agent.release_xfer_handle(entry.xfer_handle)
@@ -379,15 +388,17 @@ def get_finished_jobs(self) -> Iterable[JobResult]:
         results = self._pending_results
         self._pending_results = []
         for result in results:
-            if self.events is not None:
-                keys = self._store_job_keys.pop(result.job_id, None)
-                if result.success and keys:
+            job = self._store_jobs.pop(result.job_id, None)
+            if result.success and job is not None:
+                self._lookup_manager.mark_present(job.keys)
+                if self.events is not None and job.keys:
                     self.events.append(
                         OffloadingEvent(
-                            keys=keys,
+                            keys=list(job.keys),
                             medium=self.medium,
                             removed=False,
                             locality=self.locality,
+                            req_context=job.req_context,
                         )
                     )
             # Mark only the keys that did not load as a miss; the request
diff --git a/vllm/v1/kv_offload/tiering/p2p/control/zmq.py b/vllm/v1/kv_offload/tiering/p2p/control/zmq.py
index ee0fa819f3db..a64f0463844f 100644
--- a/vllm/v1/kv_offload/tiering/p2p/control/zmq.py
+++ b/vllm/v1/kv_offload/tiering/p2p/control/zmq.py
@@ -132,6 +132,10 @@ def close(self) -> None:
         self._closed = True
         self._dead = True
         logger.info("ZmqConnection: closing connection to %s", self.peer_id)
+        # libzmq sends monitor events with a blocking send from its I/O
+        # thread, so an event raised after the reader is closed (EVENT_CLOSED
+        # for a connect still in progress) blocks every socket in the context.
+        self._sockets.dealer.disable_monitor()
         self._sockets.monitor.close()
         self._sockets.dealer.close()
 
diff --git a/vllm/v1/kv_offload/tiering/p2p/data/base.py b/vllm/v1/kv_offload/tiering/p2p/data/base.py
index 30456ca2e22d..4a8ca7c10101 100644
--- a/vllm/v1/kv_offload/tiering/p2p/data/base.py
+++ b/vllm/v1/kv_offload/tiering/p2p/data/base.py
@@ -65,6 +65,14 @@
   state check failed stays inflight unless the failure proves it can no
   longer write, so the peer's teardown is what bounds it.
 
+remove_remote_peer() does not stop a transfer that is still running: the
+transport retires it instead. A retired transfer leaves every poll(), so
+a reconnected peer under the same id never sees it; poll_retired()
+reports it once it has stopped, and only then are its handle and the
+removed peer's registration released. The caller keeps the transfer's
+source pins until then, or until it gives up and calls
+force_release_retired().
+
 The transport never recovers a peer itself. The session turns a peer
 fault into the same connection death as a missed heartbeat, so the
 manager's reap is the only teardown: DisconnectMsg makes the peer tear
@@ -261,10 +269,27 @@ def add_remote_peer_async(
     def remove_remote_peer(self, peer_id: str) -> None:
         """Unregister a remote peer and release associated resources.
 
-        Inflight transfers to this peer should be cancelled first.
+        Transfers to this peer still inflight are retired, not released
+        (see "Failure model"); the peer's registration is released once
+        they have stopped.
         """
         ...
 
+    def poll_retired(self) -> list[int]:
+        """Release retired transfers that have stopped; return their ids.
+
+        Must be called periodically while any transfer is retired. The
+        default transport retires nothing.
+        """
+        return []
+
+    def force_release_retired(self, transfer_ids: Iterable[int]) -> None:
+        """Release retired transfers that have not stopped.
+
+        Unsafe while one may still be writing; for a bounded hold only.
+        """
+        self.cancel(transfer_ids, mode="immediate")
+
     @abstractmethod
     def write_blocks(
         self,
diff --git a/vllm/v1/kv_offload/tiering/p2p/data/nixl.py b/vllm/v1/kv_offload/tiering/p2p/data/nixl.py
index 382ed7576527..c62dab63eb22 100644
--- a/vllm/v1/kv_offload/tiering/p2p/data/nixl.py
+++ b/vllm/v1/kv_offload/tiering/p2p/data/nixl.py
@@ -105,11 +105,18 @@ class _Inflight(NamedTuple):
     transport is shared across all peer sessions of the engine.
     ``struck`` is set once the transfer has counted against its peer, so a
     state check that keeps raising charges it once, not once per poll.
+    ``retired`` is set once its peer is removed while it still runs; only
+    poll_retired() sees it from then on.
     """
 
     peer_id: str
     handle: object
     struck: bool = False
+    retired: bool = False
+
+
+# A removed peer's (peer_id, dlist, nixl_name, retired transfer_ids).
+_DeferredRemote = tuple[str, object | None, str | None, set[int]]
 
 
 class NixlTransport(DataTransport):
@@ -144,6 +151,9 @@ def __init__(
         self._peer_strikes: dict[str, int] = {}
         # peer_id → fault cause; reported by poll(peer_id) until removal.
         self._peer_faults: dict[str, str] = {}
+        # Remote handles of removed peers that retired transfers still
+        # reference. Released by poll_retired() once those transfers are gone.
+        self._deferred_remote: list[_DeferredRemote] = []
         # Registration is O(num_blocks) NIXL work — add_remote_agent plus a
         # prep_xfer_dlist over every block of the peer's region — so it runs
         # off the scheduler thread, mirroring the NIXL connector's
@@ -263,6 +273,17 @@ def remove_remote_peer(self, peer_id: str) -> None:
             dlist = self._remote_dlists.pop(peer_id, None)
         self._peer_strikes.pop(peer_id, None)
         self._peer_faults.pop(peer_id, None)
+        # Transfers still inflight here may still be writing. Retire them so
+        # a reconnected session under the same peer_id never sees them, and
+        # keep the remote handles they were prepped against until they end.
+        retired: set[int] = set()
+        for tid, entry in self._inflight.items():
+            if entry.peer_id == peer_id and not entry.retired:
+                self._inflight[tid] = entry._replace(retired=True)
+                retired.add(tid)
+        if retired:
+            self._deferred_remote.append((peer_id, dlist, nixl_name, retired))
+            return
         if self._agent is not None:
             self._release_remote_handles(peer_id, dlist, nixl_name)
 
@@ -446,7 +467,7 @@ def poll(self, peer_id: str | None = None) -> PollResult:
         failures: list[tuple[str, str]] | None = None
 
         for transfer_id, entry in self._inflight.items():
-            if peer_id is not None and entry.peer_id != peer_id:
+            if entry.retired or (peer_id is not None and entry.peer_id != peer_id):
                 continue
             try:
                 state = self._agent.check_xfer_state(entry.handle)
@@ -515,6 +536,48 @@ def _poll_result(
             peer_fault=fault,
         )
 
+    def poll_retired(self) -> list[int]:
+        """Release retired transfers that have stopped; return their ids.
+
+        A transfer stopped once NIXL reports it terminal, by the same rule
+        as ``cancel(mode="wait")``, but without charging its peer: the
+        peer is gone, and a reconnected one under the same id is not at
+        fault. A removed peer's remote handles are released once none of
+        its retired transfers remain.
+        """
+        stopped: list[int] = []
+        for tid, entry in self._inflight.items():
+            if not entry.retired:
+                continue
+            try:
+                state = self._agent.check_xfer_state(entry.handle)
+            except Exception as exc:
+                if not _is_terminal_state_error(exc):
+                    continue
+                state = "ERR"
+            if state not in ("PROC", "PEND"):
+                stopped.append(tid)
+        if not stopped:
+            return []
+        self._release_handles([self._inflight.pop(tid).handle for tid in stopped])
+        self._release_deferred_remote()
+        return stopped
+
+    def force_release_retired(self, transfer_ids: Iterable[int]) -> None:
+        """Release retired transfers that never stopped (bounded hold)."""
+        self.cancel(transfer_ids, mode="immediate")
+        self._release_deferred_remote()
+
+    def _release_deferred_remote(self) -> None:
+        pending: list[_DeferredRemote] = []
+        for peer_id, dlist, nixl_name, tids in self._deferred_remote:
+            tids &= self._inflight.keys()
+            if tids:
+                pending.append((peer_id, dlist, nixl_name, tids))
+            elif self._agent is not None:
+                self._release_remote_handles(peer_id, dlist, nixl_name)
+        self._deferred_remote = pending
+
     def cancel(
         self,
         transfer_ids: Iterable[int],
@@ -594,6 +657,7 @@ def close(self) -> None:
             return
         self._release_handles([entry.handle for entry in self._inflight.values()])
         self._inflight.clear()
+        self._release_deferred_remote()
         for peer_id in list(self._remote_dlists):
             self.remove_remote_peer(peer_id)
         if self._local_dlist is not None:
diff --git a/vllm/v1/kv_offload/tiering/p2p/manager.py b/vllm/v1/kv_offload/tiering/p2p/manager.py
index 52aab4bfd266..2fe14d65e6d9 100644
--- a/vllm/v1/kv_offload/tiering/p2p/manager.py
+++ b/vllm/v1/kv_offload/tiering/p2p/manager.py
@@ -71,6 +71,13 @@
 # enough to avoid busy-spinning the scheduler thread.
 _DRAIN_SLEEP_S = 0.001
 
+# Longest a reaped session's still-running transfer keeps its source pins.
+# NIXL may never report a wedged transfer terminal, and a hold without a
+# bound would look like a full cache; past this the handle is released.
+# The peer cannot reconnect while it holds retained transfers (see
+# _retained_peer), so this also bounds how long it stays unreachable.
+_RETAINED_HOLD_CAP_S = 30.0
+
 # Idle sessions are only reaped when remaining capacity is at most this
 # many slots. Session setup is expensive; below this headroom, keep idle
 # peers so a later burst does not pay handshake again.
@@ -222,8 +229,11 @@ class P2PSecondaryTierManager(SecondaryTierManager):
     not driven by the engine, so this tier sets ``tolerates_off_thread_polling``
     to be serviced from that thread too; without it a peer's lookup or fetch
     waits for a step boundary, which on a saturated rank can be seconds.
-    ``has_pending_work()`` keeps the engine ticking so the control transport
-    and existing sessions are polled even when no requests are scheduled.
+
+    ``has_pending_work()`` reports peer-originated work, not idle polling,
+    so a peer-only engine can pause and shut down. The tiering manager keeps
+    the engine ticking when the tier-poll thread is off or another secondary
+    tier exists.
     """
 
     cache_hit_source: ClassVar[CacheHitSource] = CacheHitSource.P2P
@@ -400,6 +410,10 @@ def __init__(
         # at the top of the next ``serve_external_requests`` where the
         # handle is valid.
         self._failed_serve_ctxs: list[ReqContext] = []
+        # transfer_id → (peer_id, store job_ids, retained_at) for transfers
+        # of reaped sessions that may still be writing; see _poll_retained.
+        self._retained: dict[int, tuple[str, frozenset[int], float]] = {}
+        self._retained_force_released = 0
 
     # ------------------------------------------------------------------
     # SecondaryTierManager interface
@@ -419,9 +433,11 @@ def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult:
         # in on_schedule_end batches the LookupMsg; a later step's
         # lookup() returns HIT/MISS once LookupRespMsg has arrived.
         # PD path (``remote_prefiller`` sub-dict only) keeps the eager HIT.
+        # Until the peer acks the handshake, prefill locally rather than wait
+        # up to VLLM_P2P_HANDSHAKE_TIMEOUT_S on a peer that may be gone.
         if source.do_probe:
             session = self._sessions.get(source.peer_id)
-            if session is None:
+            if session is None or not session.ready:
                 return LookupResult.MISS
             result = session.register_lookup(source.kv_request_id, key)
             if result is True:
@@ -646,28 +662,44 @@ def get_finished_jobs(self) -> Iterable[JobResult]:
 
     @override
     def has_pending_work(self) -> bool:
-        # The engine tick is the only driver of _control.poll() and
-        # session.poll(); without it we miss new peer connects and
-        # inbound fetch messages on existing sessions. Keep the engine
-        # ticking for the lifetime of this manager.
-        return True
+        # Idle polling is excluded, or the engine could never go idle; the
+        # tiering manager keeps it ticking when the tier-poll thread is off.
+        return bool(
+            self._finished_jobs
+            or self._failed_serve_ctxs
+            or self._unbound_stores
+            or any(s.has_pending_work for s in self._sessions.values())
+        )
 
     @override
     def drain_jobs(self) -> None:
         """Block until every submitted load/store job has completed or failed.
 
-        Loops calling ``_poll_once()`` until no session has outstanding
-        inbound loads or in-flight outbound stores. Mid-flight transfers
-        are NOT cancelled — the caller (``TieringOffloadingManager.reset_cache``)
-        needs the primary memoryview to be quiescent, not aborted. Results
+        Fail stores still waiting for a fetch, then poll until no session has
+        outstanding inbound loads or in-flight outbound stores, and no reaped session's
+        transfer is still retained (bounded by its hold cap). Mid-flight
+        transfers are NOT cancelled — the caller
+        (``TieringOffloadingManager.reset_cache``) needs the primary
+        memoryview to be quiescent, not aborted. Results
         accumulate in ``_finished_jobs`` and are surfaced by the next
         ``get_finished_jobs()`` call.
+
+        A load already failed with a deferred release stays pending until
+        its AbortAck, or until a missed AbortAck disconnects the session and
+        the reap's close() sets the event, so every event is set on return.
         """
         start = time.monotonic()
+        self._reap_unbound_stores(force=True)
+        self._kv_to_session.clear()
+        for session in self._sessions.values():
+            self._drain_session_stores(session)
+
         warned = False
         while True:
-            self._poll_once()
-            pending = any(s.has_pending_work for s in self._sessions.values())
+            self._poll_once(draining=True)
+            pending = bool(self._retained) or any(
+                s.has_pending_work for s in self._sessions.values()
+            )
             if not pending:
                 return
             if not warned and time.monotonic() - start > 5.0:
@@ -678,6 +710,18 @@ def drain_jobs(self) -> None:
                 warned = True
             time.sleep(_DRAIN_SLEEP_S)
 
+    def _drain_session_stores(
+        self, session: P2PSession, kv_request_ids: Iterable[str] | None = None
+    ) -> None:
+        result = session.drain_stores(kv_request_ids)
+        self._finished_jobs.extend(
+            JobResult(store.job_id, success=store.success) for store in result.stores
+        )
+        now = time.monotonic()
+        for kid in result.reaped_pd_ids:
+            self._reaped_stores[kid] = now
+            self._kv_to_session.pop(kid, None)
+
     @override
     def serve_external_requests(self, parent: ParentManager) -> None:
         """Serve inbound peer lookups against the tiering manager.
@@ -730,7 +774,7 @@ def _get_or_create_session(self, peer_id: str) -> P2PSession | None:
         session = self._sessions.get(peer_id)
         if session is not None:
             return session
-        if self._reconnect_wait(peer_id) > 0:
+        if self._reconnect_wait(peer_id) > 0 or self._retained_peer(peer_id):
             return None
 
         max_peers = envs.VLLM_P2P_MAX_PEERS
@@ -803,6 +847,17 @@ def _accept_new_peers(self, new_connections: Sequence[ControlConnection]) -> Non
                     conn.send({TYPE_KEY: DisconnectMsg.TYPE})
                 conn.close()
                 continue
+            if self._retained_peer(conn.peer_id):
+                logger.info(
+                    "P2P %s: rejecting inbound peer %s — its reaped "
+                    "session's transfers are still running",
+                    self._local_id,
+                    conn.peer_id,
+                )
+                with contextlib.suppress(Exception):
+                    conn.send({TYPE_KEY: DisconnectMsg.TYPE})
+                conn.close()
+                continue
             try:
                 existing = self._sessions.get(conn.peer_id)
                 if existing is not None:
@@ -852,6 +907,7 @@ def _reap_dead_sessions(self) -> None:
                 self._finished_jobs.append(JobResult(job_id=job_id, success=False))
             for job_id in close_result.failed_stores:
                 self._finished_jobs.append(JobResult(job_id=job_id, success=False))
+            self._retain(pid, close_result.retained)
             # Fail every client-side request (in-flight loads plus unresolved
             # symmetric-P2P probes) toward the dead peer so lookup() returns
             # MISS (local prefill) instead of RETRY forever — even if a fresh
@@ -877,6 +933,70 @@ def _reap_dead_sessions(self) -> None:
                 fault_info,
             )
 
+    def _retain(self, peer_id: str, retained: dict[int, frozenset[int]]) -> None:
+        """Hold a reaped session's still-running transfers and their pins."""
+        if not retained:
+            return
+        now = time.monotonic()
+        for tid, job_ids in retained.items():
+            self._retained[tid] = (peer_id, job_ids, now)
+        logger.warning(
+            "P2P %s: peer %s reaped with %d transfers still running; holding "
+            "their source pins until they stop (at most %.0fs)",
+            self._local_id,
+            peer_id,
+            len(retained),
+            _RETAINED_HOLD_CAP_S,
+        )
+
+    def _poll_retained(self) -> None:
+        """Unpin retained transfers once they stop, or at the hold cap.
+
+        A store job is failed only once no retained transfer reads it.
+        """
+        if not self._retained:
+            return
+        released = list(self._data.poll_retired())
+        deadline = time.monotonic() - _RETAINED_HOLD_CAP_S
+        expired = [
+            tid
+            for tid, (_, _, retained_at) in self._retained.items()
+            if retained_at <= deadline and tid not in released
+        ]
+        if expired:
+            self._data.force_release_retired(expired)
+            self._retained_force_released += len(expired)
+            logger.error(
+                "P2P %s: %d transfers of reaped peers still running after "
+                "%.0fs; releasing their handles and source pins anyway "
+                "(%d so far)",
+                self._local_id,
+                len(expired),
+                _RETAINED_HOLD_CAP_S,
+                self._retained_force_released,
+            )
+        jobs: set[int] = set()
+        for tid in (*released, *expired):
+            entry = self._retained.pop(tid, None)
+            if entry is not None:
+                jobs |= entry[1]
+        for job_id in jobs - self._retained_job_ids():
+            self._finished_jobs.append(JobResult(job_id=job_id, success=False))
+
+    def _retained_job_ids(self) -> set[int]:
+        return {jid for _, job_ids, _ in self._retained.values() for jid in job_ids}
+
+    def _retained_peer(self, peer_id: str) -> bool:
+        """Whether a reaped session of *peer_id* still has running transfers.
+
+        Its registration stays until they stop. A new session would load
+        the peer's NIXL agent again under the same name: NIXL keeps the old
+        UCX endpoints, possibly failed, and the old registration's release
+        removes the agent under the new one. So the peer stays out until
+        the transfers stop, at most _RETAINED_HOLD_CAP_S.
+        """
+        return any(pid == peer_id for pid, _, _ in self._retained.values())
+
     def _schedule_reconnect_backoff(self, peer_id: str) -> float:
         """Record a data-plane fault for *peer_id*; return its reconnect wait.
 
@@ -901,7 +1021,7 @@ def _reconnect_wait(self, peer_id: str) -> float:
         entry = self._reconnect_backoff.get(peer_id)
         return 0.0 if entry is None else entry[1] - time.monotonic()
 
-    def _reap_unbound_stores(self) -> None:
+    def _reap_unbound_stores(self, *, force: bool = False) -> None:
         """Time out submit_store batches that no peer has ever fetched.
 
         Walks `_unbound_stores` for entries whose oldest batch is older
@@ -929,7 +1049,7 @@ def _reap_unbound_stores(self) -> None:
         expired: list[str] | None = None
         for kid, batches in self._unbound_stores.items():
             # Batches are appended in arrival order, so the head is oldest.
-            if batches and batches[0].submitted_at <= deadline:
+            if batches and (force or batches[0].submitted_at <= deadline):
                 if expired is None:
                     expired = []
                 expired.append(kid)
@@ -942,14 +1062,15 @@ def _reap_unbound_stores(self) -> None:
                 self._finished_jobs.append(
                     JobResult(job_id=batch.job_id, success=False)
                 )
-            logger.warning(
-                "P2P %s: unbound store kv_request_id=%s timed out after %.0fs "
-                "without a fetch — failing %d job(s)",
-                self._local_id,
-                kid,
-                self._unbound_store_timeout_s,
-                len(batches),
-            )
+            if not force:
+                logger.warning(
+                    "P2P %s: unbound store kv_request_id=%s timed out after %.0fs "
+                    "without a fetch — failing %d job(s)",
+                    self._local_id,
+                    kid,
+                    self._unbound_store_timeout_s,
+                    len(batches),
+                )
 
     def _reap_idle_sessions(self) -> None:
         """Evict idle sessions when remaining peer slots are scarce.
@@ -988,6 +1109,7 @@ def _reap_idle_sessions(self) -> None:
                 self._finished_jobs.append(JobResult(job_id=job_id, success=False))
             for job_id in close_result.failed_stores:
                 self._finished_jobs.append(JobResult(job_id=job_id, success=False))
+            self._retain(pid, close_result.retained)
             self._failed_req_ids.update(close_result.failed_req_ids)
             self._failed_serve_ctxs.extend(close_result.failed_serves)
             self._data.remove_remote_peer(pid)
@@ -1002,7 +1124,7 @@ def _reap_idle_sessions(self) -> None:
     # Polling
     # ------------------------------------------------------------------
 
-    def _poll_once(self) -> None:
+    def _poll_once(self, *, draining: bool = False) -> None:
         """One sweep of the polling work.
 
         Drains the control transport, polls every session, accumulates
@@ -1025,14 +1147,22 @@ def _poll_once(self) -> None:
             result = session.poll()
             for lr in result.loads:
                 self._finished_jobs.append(
-                    JobResult(job_id=lr.job_id, success=lr.success)
+                    JobResult(
+                        job_id=lr.job_id,
+                        success=lr.success,
+                        deferred_release=lr.deferred_release,
+                    )
                 )
-                if not lr.success:
+                if not lr.success and not lr.request_finished:
                     self._failed_req_ids.add(lr.kv_request_id)
             for sr in result.stores:
                 self._finished_jobs.append(
                     JobResult(job_id=sr.job_id, success=sr.success)
                 )
+            if draining:
+                if result.new_fetch_ids:
+                    self._drain_session_stores(session, result.new_fetch_ids)
+                continue
             # Bind kv_request_id → session for any FetchMsg this tick and
             # replay any submit_store batches parked while no peer was
             # asking. ServerRole.on_fetch already recorded the demand
@@ -1066,6 +1196,7 @@ def _poll_once(self) -> None:
         self._reap_dead_sessions()
         self._reap_unbound_stores()
         self._reap_idle_sessions()
+        self._poll_retained()
 
     # ------------------------------------------------------------------
     # Lifecycle
@@ -1090,6 +1221,11 @@ def shutdown(self) -> None:
                 )
         self._unbound_stores.clear()
         self._reaped_stores.clear()
+        # The drain above already waited on retained transfers too, and
+        # _data.close() releases whatever is left.
+        for job_id in self._retained_job_ids():
+            self._finished_jobs.append(JobResult(job_id=job_id, success=False))
+        self._retained.clear()
         self._control.close()
         self._data.close()
 
@@ -1104,6 +1240,7 @@ def _drain_inflight_for_shutdown(self) -> None:
         wedged peer must not hang shutdown.
         """
         ids = [tid for s in self._sessions.values() for tid in s._server._inflight]
+        ids.extend(self._retained)
         if not ids:
             return
         deadline = time.monotonic() + _SHUTDOWN_DRAIN_TIMEOUT_S
diff --git a/vllm/v1/kv_offload/tiering/p2p/session/client.py b/vllm/v1/kv_offload/tiering/p2p/session/client.py
index b545be47aeae..9bb4b9f5780c 100644
--- a/vllm/v1/kv_offload/tiering/p2p/session/client.py
+++ b/vllm/v1/kv_offload/tiering/p2p/session/client.py
@@ -11,6 +11,7 @@
 
 from __future__ import annotations
 
+import threading
 import time
 from collections.abc import Callable, Sequence
 from dataclasses import dataclass, field
@@ -54,6 +55,13 @@ class _InboundLoadState:
     # _P2P_LOAD_TIMEOUT_S or _PD_LOAD_TIMEOUT_S, fixed at submit time.
     timeout_s: float
     aborted_at: float | None = None
+    # Set together with aborted_at, when the load is reported failed. The
+    # peer may still be writing into its blocks, so their release waits on
+    # this event: set on AbortAck/TransferDone, or when the session closes.
+    release: threading.Event | None = None
+    # Its AbortAck timed out; only close() can release it now.
+    ack_expired: bool = False
+    request_finished: bool = False
 
 
 @dataclass
@@ -99,14 +107,19 @@ class LoadResult(NamedTuple):
     job_id: int
     kv_request_id: str
     success: bool
+    # A failed load whose blocks the peer may still write: release them
+    # once this is set.
+    deferred_release: threading.Event | None = None
+    request_finished: bool = False
 
 
 class ClientCloseResult(NamedTuple):
     """What the manager must fail when a peer session is torn down.
 
-    ``failed_jobs`` is the ``job_id`` of every load still in flight; each
-    becomes a ``JobResult(success=False)``. ``failed_req_ids`` is every
-    kv_request_id whose lookup() would otherwise defer forever on the dead
+    ``failed_jobs`` is the ``job_id`` of every load still in flight and not
+    yet reported; each becomes a ``JobResult(success=False)``.
+    ``failed_req_ids`` is every kv_request_id whose lookup() would otherwise
+    defer forever on the dead
     peer — the in-flight loads plus any request holding an unresolved
     symmetric-P2P probe. ``failed_jobs`` is the subset of ``failed_req_ids``
     that had a load job.
@@ -146,6 +159,9 @@ def __init__(self, peer_id: str, send: Callable[[dict], None]) -> None:
         # Per-request progress deadlines. Only newly resolved probes extend
         # them; additional keys and repeated scheduler polls do not.
         self._lookup_deadlines: dict[str, float] = {}
+        # First load whose AbortAck never arrived. Sticky: the peer has
+        # stopped confirming that its WRITEs ended, so the session faults it.
+        self._expired_abort_ack: tuple[str, int] | None = None
 
     # ------------------------------------------------------------------
     # State helpers
@@ -177,6 +193,30 @@ def _on_load_terminal(self, kv_request_id: str, st: _ClientRequestState) -> None
         self._active_loads.discard(kv_request_id)
         self._maybe_prune(kv_request_id)
 
+    def _abort_load(
+        self, kv_request_id: str, round_seq: int, load: _InboundLoadState
+    ) -> None:
+        """Send AbortFetch and report the load failed now, deferring the
+        release of its blocks until the peer confirms its WRITEs ended."""
+        self._send(
+            {
+                TYPE_KEY: AbortFetchMsg.TYPE,
+                AbortFetchMsg.KV_REQUEST_ID: kv_request_id,
+                AbortFetchMsg.ROUND_SEQ: round_seq,
+            }
+        )
+        load.aborted_at = time.monotonic()
+        load.release = threading.Event()
+        self._completed_loads.append(
+            LoadResult(
+                job_id=load.job_id,
+                kv_request_id=kv_request_id,
+                success=False,
+                deferred_release=load.release,
+                request_finished=load.request_finished,
+            )
+        )
+
     @property
     def has_active_loads(self) -> bool:
         """True if any kv_request_id has a fetch in flight."""
@@ -190,6 +230,11 @@ def get_expired_lookup(self) -> tuple[str, float] | None:
                 return req_id, now - deadline + _LOOKUP_TIMEOUT_S
         return None
 
+    @property
+    def expired_abort_ack(self) -> tuple[str, int] | None:
+        """The first ``(kv_request_id, round_seq)`` whose AbortAck timed out."""
+        return self._expired_abort_ack
+
     # ------------------------------------------------------------------
     # Public API
     # ------------------------------------------------------------------
@@ -261,24 +306,23 @@ def finish(self, kv_request_id: str) -> None:
         while an earlier round's load is still in flight, so both can be
         owed at once.
 
-        Then drop all probe/lookup state. Keep aborted loads until their abort
-        acknowledgement or timeout produces a terminal result.
+        Then drop all probe/lookup state. Aborted loads are reported failed
+        at once and kept until their AbortAck (or the session's close)
+        releases their blocks.
         """
+        self._completed_loads = [
+            result._replace(request_finished=True)
+            if result.kv_request_id == kv_request_id
+            else result
+            for result in self._completed_loads
+        ]
         st = self._requests.get(kv_request_id)
         if st is None:
             return
-        if st.loads:
-            for round_seq, load in st.loads.items():
-                if load.aborted_at is not None:
-                    continue
-                self._send(
-                    {
-                        TYPE_KEY: AbortFetchMsg.TYPE,
-                        AbortFetchMsg.KV_REQUEST_ID: kv_request_id,
-                        AbortFetchMsg.ROUND_SEQ: round_seq,
-                    }
-                )
-                load.aborted_at = time.monotonic()
+        for round_seq, load in st.loads.items():
+            load.request_finished = True
+            if load.aborted_at is None:
+                self._abort_load(kv_request_id, round_seq, load)
         if st.peer_lookup_open:
             st.peer_lookup_open = False
             self._send(
@@ -302,12 +346,17 @@ def on_transfer_done(
         """Handle a TransferDoneMsg from the peer."""
         st = self._requests.get(kv_request_id)
         load = st.loads.pop(round_seq, None) if st is not None else None
-        if st is not None and load is not None:
+        if st is not None and load is not None and load.release is not None:
+            # Already reported failed; the peer's WRITEs are over.
+            load.release.set()
+            self._on_load_terminal(kv_request_id, st)
+        elif st is not None and load is not None:
             self._completed_loads.append(
                 LoadResult(
                     job_id=load.job_id,
                     kv_request_id=kv_request_id,
                     success=success,
+                    request_finished=load.request_finished,
                 )
             )
             self._on_load_terminal(kv_request_id, st)
@@ -331,20 +380,10 @@ def on_abort_ack(self, kv_request_id: str, round_seq: int) -> None:
         st = self._requests.get(kv_request_id)
         load = st.loads.pop(round_seq, None) if st is not None else None
         if st is not None and load is not None:
-            logger.warning(
-                "P2PSession %s: load request %s (job_id=%d) "
-                "completed with failure after abort",
-                self._peer_id,
-                kv_request_id,
-                load.job_id,
-            )
-            self._completed_loads.append(
-                LoadResult(
-                    job_id=load.job_id,
-                    kv_request_id=kv_request_id,
-                    success=False,
-                )
-            )
+            # Every abort reported the load failed when it was sent; the
+            # ack only releases its blocks.
+            assert load.release is not None
+            load.release.set()
             self._on_load_terminal(kv_request_id, st)
         else:
             # See on_transfer_done: same ambiguity (duplicate ack
@@ -482,54 +521,40 @@ def on_lookup_resp(
     def collect_results(self) -> list[LoadResult]:
         """Walk load timeouts and drain completed loads.
 
-        Loads past their load timeout (``_P2P_LOAD_TIMEOUT_S`` or
-        ``_PD_LOAD_TIMEOUT_S``) get an AbortFetchMsg sent and
-        enter the aborting phase. Aborting loads past
-        ``_ABORT_ACK_TIMEOUT_S`` are surfaced as failed.
+        A load past its load timeout (``_P2P_LOAD_TIMEOUT_S`` or
+        ``_PD_LOAD_TIMEOUT_S``) gets an AbortFetchMsg and is reported failed
+        at once, with its blocks' release deferred. One whose AbortAck then
+        misses ``_ABORT_ACK_TIMEOUT_S`` is recorded in ``expired_abort_ack``,
+        so the session disconnects and close() releases it.
         """
         now = time.monotonic()
-        to_remove: list[tuple[str, int]] = []
         for req_id in self._active_loads:
             st = self._requests[req_id]
             assert st.loads
             for round_seq, load in st.loads.items():
                 if load.aborted_at is None:
                     if now - load.submitted_at >= load.timeout_s:
-                        load.aborted_at = now
                         logger.warning(
                             "P2PSession %s: %s round=%s timed out, sending abort",
                             self._peer_id,
                             req_id,
                             round_seq,
                         )
-                        self._send(
-                            {
-                                TYPE_KEY: AbortFetchMsg.TYPE,
-                                AbortFetchMsg.KV_REQUEST_ID: req_id,
-                                AbortFetchMsg.ROUND_SEQ: round_seq,
-                            }
-                        )
-                else:
-                    if now - load.aborted_at >= _ABORT_ACK_TIMEOUT_S:
-                        to_remove.append((req_id, round_seq))
-                        self._completed_loads.append(
-                            LoadResult(
-                                job_id=load.job_id,
-                                kv_request_id=req_id,
-                                success=False,
-                            )
-                        )
-                        logger.warning(
-                            "P2PSession %s: abort_ack timed out for "
-                            "kv_request_id=%s round=%s",
-                            self._peer_id,
-                            req_id,
-                            round_seq,
-                        )
-        for req_id, round_seq in to_remove:
-            st = self._requests[req_id]
-            st.loads.pop(round_seq)
-            self._on_load_terminal(req_id, st)
+                        self._abort_load(req_id, round_seq, load)
+                elif (
+                    not load.ack_expired
+                    and now - load.aborted_at >= _ABORT_ACK_TIMEOUT_S
+                ):
+                    load.ack_expired = True
+                    if self._expired_abort_ack is None:
+                        self._expired_abort_ack = (req_id, round_seq)
+                    logger.warning(
+                        "P2PSession %s: abort_ack timed out for "
+                        "kv_request_id=%s round=%s",
+                        self._peer_id,
+                        req_id,
+                        round_seq,
+                    )
 
         results = self._completed_loads
         self._completed_loads = []
@@ -541,14 +566,22 @@ def close(self) -> ClientCloseResult:
         A request is failed if it has a load in flight (its job fails) or
         holds an unresolved symmetric-P2P probe (its lookup() would defer
         forever on an answer that can never arrive). See ``ClientCloseResult``.
+        The end of the session is the release point for every load already
+        reported failed: their deferred releases are set here.
         """
-        failed_jobs = [
-            load.job_id for st in self._requests.values() for load in st.loads.values()
-        ]
+        # Results never collected still owe their job an outcome.
+        failed_jobs = [r.job_id for r in self._completed_loads]
+        for st in self._requests.values():
+            for load in st.loads.values():
+                if load.release is not None:
+                    load.release.set()
+                else:
+                    failed_jobs.append(load.job_id)
         failed_req_ids = [
             req_id
             for req_id, st in self._requests.items()
-            if st.loads or any(hit is None for hit in st.probes.values())
+            if any(not load.request_finished for load in st.loads.values())
+            or any(hit is None for hit in st.probes.values())
         ]
         self._requests.clear()
         self._flush_pending.clear()
diff --git a/vllm/v1/kv_offload/tiering/p2p/session/server.py b/vllm/v1/kv_offload/tiering/p2p/session/server.py
index eff19222ea29..0e0bdca42a0c 100644
--- a/vllm/v1/kv_offload/tiering/p2p/session/server.py
+++ b/vllm/v1/kv_offload/tiering/p2p/session/server.py
@@ -22,6 +22,7 @@
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, NamedTuple
 
+import vllm.envs as envs
 from vllm.logger import init_logger
 from vllm.v1.kv_offload.base import LookupResult, OffloadKey, ReqContext
 from vllm.v1.kv_offload.tiering.p2p.session.protocol import (
@@ -57,6 +58,13 @@ class StoreResult(NamedTuple):
     success: bool
 
 
+class StoreDrainResult(NamedTuple):
+    """Failed queued stores and PD ids whose handoffs cannot be resumed."""
+
+    stores: list[StoreResult]
+    reaped_pd_ids: set[str]
+
+
 class _MatchResult(NamedTuple):
     """Result of block matching: pairs ready for transfer."""
 
@@ -166,6 +174,7 @@ class _InflightXfer:
     # Dummy default for test-seeded entries.
     round: _OutboundRequestState = field(default_factory=_OutboundRequestState)
     round_key: int = 0
+    submitted_at: float = field(default_factory=lambda: time.monotonic())
 
 
 @dataclass
@@ -288,6 +297,9 @@ def __init__(
         self._slow_aborts: set[tuple[str, int]] = set()
         # PollResult.peer_fault, once reported; no new WRITE after that.
         self.peer_fault: str | None = None
+        # True once peer_fault came from _check_stalled_transfers.
+        self.transfer_stalled = False
+        self._stall_timeout_s = envs.VLLM_P2P_TRANSFER_STALL_TIMEOUT_S
 
     # ------------------------------------------------------------------
     # State helpers
@@ -754,6 +766,30 @@ def finish(self, kv_request_id: str) -> None:
                 continue
             self._finalize_outbound(kv_request_id, key)
 
+    def drain_stores(
+        self, kv_request_ids: Iterable[str] | None = None
+    ) -> StoreDrainResult:
+        """Withdraw supply, keeping empty lookup rounds for their late fetch."""
+        reaped_pd_ids: set[str] = set()
+        ids = list(self._requests) if kv_request_ids is None else kv_request_ids
+        for kv_request_id in ids:
+            st = self._requests.get(kv_request_id)
+            if st is None:
+                continue
+            for round_key, rnd in list(st.outbound.items()):
+                if not rnd.lookup_supplied:
+                    reaped_pd_ids.add(kv_request_id)
+                rnd.available.clear()
+                rnd.finishing = True
+                self._pending_store_results.extend(self._fail_unwritten_round_jobs(rnd))
+                if not rnd.inflight and (
+                    rnd.demand_received or not rnd.lookup_supplied
+                ):
+                    self._finalize_cancelled(kv_request_id, round_key)
+        results = self._pending_store_results
+        self._pending_store_results = []
+        return StoreDrainResult(results, reaped_pd_ids)
+
     def collect_results(self) -> list[StoreResult]:
         """Drain timeouts, deferred results, and transport completions.
 
@@ -839,6 +875,10 @@ def collect_results(self) -> list[StoreResult]:
                 self._finalize_cancelled(xfer.kv_request_id, xfer.round_key)
             self._maybe_prune(xfer.kv_request_id)
 
+        # After the drain, so a transfer this poll resolved is not "stalled".
+        if self.peer_fault is None:
+            self._check_stalled_transfers()
+
         # Finalizing transfers can queue results for supplied but unfetched blocks.
         if self._pending_store_results:
             results.extend(self._pending_store_results)
@@ -860,19 +900,28 @@ def drain_pending_aborts(self) -> None:
         for kv_request_id, round_seq in list(self._pending_aborts):
             self._drain_abort(kv_request_id, round_seq)
 
-    def close(self) -> tuple[list[int], list[ReqContext]]:
-        """Tear down. Cancels inflight.
+    def close(
+        self,
+    ) -> tuple[list[int], list[ReqContext], dict[int, frozenset[int]]]:
+        """Tear down. Cancels inflight, retaining what may still write.
 
-        Returns ``(failed_store_job_ids, failed_serves)`` where
+        Returns ``(failed_store_job_ids, failed_serves, retained)``.
         ``failed_serves`` are synthetic lookup ctxs still owing a
         ``parent.on_request_finished``. The session is going away with no
         parent handle in hand, so the manager flushes these in its next
-        ``serve_external_requests``.
+        ``serve_external_requests``. ``retained`` maps each transfer NIXL
+        does not report terminal to the store jobs it still reads: like the
+        abort drain, close never releases such a transfer, so those jobs
+        stay pinned and are left out of ``failed_store_job_ids`` until the
+        transport reports the transfer stopped.
         """
-        failed_stores = list(self._store_jobs.keys())
-        self._store_jobs.clear()
+        retained: dict[int, frozenset[int]] = {}
         if self._inflight:
-            self._transport.cancel(list(self._inflight.keys()))
+            still = self._transport.cancel(list(self._inflight.keys()), mode="wait")
+            retained = {tid: frozenset(self._inflight[tid].job_ids) for tid in still}
+        held = set().union(*retained.values())
+        failed_stores = [jid for jid in self._store_jobs if jid not in held]
+        self._store_jobs.clear()
         self._inflight.clear()
         self._pending_store_results.clear()
         # Surface every synthetic ctx still owing on_request_finished so
@@ -888,7 +937,7 @@ def close(self) -> tuple[list[int], list[ReqContext]]:
         self._pending_aborts.clear()
         self._slow_aborts.clear()
         self._finished_lookup_ctxs.clear()
-        return failed_stores, failed_serves
+        return failed_stores, failed_serves, retained
 
     @property
     def has_inflight_transfers(self) -> bool:
@@ -1224,6 +1273,26 @@ def _submit_transfer(
             ):
                 self._finalize_outbound(kv_request_id, round_key, success=False)
 
+    def _check_stalled_transfers(self) -> None:
+        """Fault the peer once a transfer outlives the stall timeout.
+
+        An RDMA WRITE posted to a live NIC completes or errors within
+        milliseconds; one NIXL still reports PROC/PEND long after points at
+        a wedged lane or endpoint, which nothing else would ever surface.
+        The fault stops new WRITEs to this peer and ends the session, but
+        releases nothing: the transfer and its pins are retained across
+        teardown until the transport reports it stopped.
+        """
+        if not self._inflight or self._stall_timeout_s <= 0:
+            return
+        now = time.monotonic()
+        oldest = min(x.submitted_at for x in self._inflight.values())
+        age = now - oldest
+        if age < self._stall_timeout_s:
+            return
+        self.peer_fault = f"transfer stalled for {age:.0f}s"
+        self.transfer_stalled = True
+
     def _timeout_pending_store_jobs(self) -> list[StoreResult]:
         if not self._store_jobs:
             return []
diff --git a/vllm/v1/kv_offload/tiering/p2p/session/session.py b/vllm/v1/kv_offload/tiering/p2p/session/session.py
index 3c435bb7d334..d7bfa99a629e 100644
--- a/vllm/v1/kv_offload/tiering/p2p/session/session.py
+++ b/vllm/v1/kv_offload/tiering/p2p/session/session.py
@@ -22,7 +22,7 @@
 
 import contextlib
 import time
-from collections.abc import Sequence
+from collections.abc import Iterable, Sequence
 from concurrent.futures import Future
 from typing import TYPE_CHECKING, NamedTuple
 
@@ -45,6 +45,7 @@
 )
 from vllm.v1.kv_offload.tiering.p2p.session.server import (
     ServerRole,
+    StoreDrainResult,
     StoreResult,
 )
 
@@ -91,13 +92,17 @@ class SessionCloseResult(NamedTuple):
     server stores) the manager must mark failed. `failed_req_ids` is every
     client-side kv_request_id whose lookup() must fail (in-flight loads plus
     unresolved probes); `failed_serves` is the server-side lookup state the
-    dead peer can no longer resolve.
+    dead peer can no longer resolve. `retained` maps each server transfer
+    that may still be writing to the store jobs it reads: they are not in
+    `failed_stores`, and stay pinned until the transport reports it stopped.
     """
 
     failed_jobs: list[int]  # client load job_ids
     failed_req_ids: list[str]  # client kv_request_ids (loads + probes)
     failed_stores: list[int]  # server store job_ids
     failed_serves: list[ReqContext]  # server-side lookup ctxs needing release
+    # transfer_id → store job_ids it still reads
+    retained: dict[int, frozenset[int]]
 
 
 class P2PSession:
@@ -144,6 +149,8 @@ def __init__(
         # First cause recorded by _mark_dead; see death_reason.
         self._death_reason: str | None = None
         self._peer_fault_at: float | None = None
+        self._abort_ack_fault: str | None = None
+        self._registration_fault: str | None = None
 
         now = time.monotonic()
         self._connected_at: float | None = None
@@ -211,8 +218,15 @@ def death_reason(self) -> str:
 
     @property
     def data_plane_fault(self) -> str | None:
-        """The transport's peer fault, if one was reported (see server role)."""
-        return self._server.peer_fault
+        """The data-plane fault that ended this session, if any.
+
+        The transport's peer fault (see server role), an AbortAck that never
+        arrived, which leaves this side unable to tell whether the peer's
+        WRITE stopped, or a failed registration of the peer's NIXL agent.
+        """
+        if self._server.peer_fault is not None:
+            return self._server.peer_fault
+        return self._abort_ack_fault or self._registration_fault
 
     @property
     def last_activity_at(self) -> float:
@@ -263,6 +277,12 @@ def add_stored_blocks(
         """New blocks stored locally — match against pending fetch demand."""
         self._server.add_stored_blocks(kv_request_id, keys, block_ids, job_id)
 
+    def drain_stores(
+        self, kv_request_ids: Iterable[str] | None = None
+    ) -> StoreDrainResult:
+        """Withdraw queued stores, optionally for fetches arriving during drain."""
+        return self._server.drain_stores(kv_request_ids)
+
     def finish_request(self, kv_request_id: str) -> None:
         """Called when the request is finishing locally.
 
@@ -351,6 +371,7 @@ def poll(self) -> SessionPollResult:
         stores = self._server.collect_results()
         self._server.drain_pending_aborts()
         self._escalate_peer_fault()
+        self._escalate_abort_ack_timeout()
 
         # Drain buffered replies and completed transfers before expiring a peer.
         expired_lookup = self._client.get_expired_lookup()
@@ -384,6 +405,8 @@ def close(self) -> SessionCloseResult:
         failed_serves: synthetic lookup ctxs still owing
             ``parent.on_request_finished`` (the manager flushes these on
             its next ``serve_external_requests``).
+        retained: server transfers that may still be writing, with the
+            store jobs they read (held until they stop).
         """
         # Abandon any in-flight registration: nothing will ack this peer
         # now. The manager's remove_remote_peer bumps the transport's
@@ -394,7 +417,7 @@ def close(self) -> SessionCloseResult:
             self._pending_registration = None
 
         client_result = self._client.close()
-        failed_stores, failed_serves = self._server.close()
+        failed_stores, failed_serves, retained = self._server.close()
 
         if self._conn is not None:
             with contextlib.suppress(Exception):
@@ -407,6 +430,7 @@ def close(self) -> SessionCloseResult:
             failed_req_ids=client_result.failed_req_ids,
             failed_stores=failed_stores,
             failed_serves=failed_serves,
+            retained=retained,
         )
 
     # ------------------------------------------------------------------
@@ -477,13 +501,37 @@ def _escalate_peer_fault(self) -> None:
                 self.peer_id,
                 fault,
             )
+        # A stalled transfer would never drain, and teardown retains
+        # whatever is still running, so waiting buys nothing.
         if (
             self._server.has_inflight_transfers
+            and not self._server.transfer_stalled
             and now - self._peer_fault_at < _PEER_FAULT_DRAIN_S
         ):
             return
         self._disconnect(f"data plane: {fault}")
 
+    def _escalate_abort_ack_timeout(self) -> None:
+        """Fault the peer once one of its AbortAcks times out.
+
+        The timed-out load already failed in this poll's results. Without
+        this, the session stays up and every later fetch to the same stuck
+        peer waits out the full load and abort-ack timeouts again.
+        """
+        expired = self._client.expired_abort_ack
+        if expired is None or self._abort_ack_fault is not None:
+            return
+        req_id, round_seq = expired
+        self._abort_ack_fault = "abort ack timed out"
+        logger.warning(
+            "P2PSession %s: abort_ack for kv_request_id=%s round=%s timed out; "
+            "disconnecting peer session",
+            self.peer_id,
+            req_id,
+            round_seq,
+        )
+        self._disconnect(self._abort_ack_fault)
+
     def _disconnect(self, reason: str) -> None:
         """Best-effort DisconnectMsg, then mark the connection dead."""
         if self._conn is not None:
@@ -631,7 +679,8 @@ def _poll_registration(self) -> None:
         Called from poll() after message dispatch. A successful
         registration releases ConnectAckMsg — the peer treats that as
         permission to send, so it must not go out earlier. A failed one
-        marks the connection dead so the manager reaps the session.
+        sends a best-effort disconnect and marks the local connection dead
+        for the manager to reap.
         """
         registration = self._pending_registration
         if registration is None or not registration.done():
@@ -645,7 +694,11 @@ def _poll_registration(self) -> None:
                 self.peer_id,
                 exc,
             )
-            self._mark_dead(f"peer registration failed: {exc}")
+            # The peer may have registered us and be waiting for our ack. A
+            # session it keeps would ack our next ConnectMsg without sending
+            # its own, so we would look it up without ever registering it.
+            self._registration_fault = f"peer registration failed: {exc}"
+            self._disconnect(self._registration_fault)
             return
         if self._conn is not None and self._conn.alive:
             self._do_send(
diff --git a/vllm/v1/kv_offload/tiering/spec.py b/vllm/v1/kv_offload/tiering/spec.py
index 3aec216885b1..2e8bb2a02137 100644
--- a/vllm/v1/kv_offload/tiering/spec.py
+++ b/vllm/v1/kv_offload/tiering/spec.py
@@ -247,6 +247,27 @@ def build_metric_definitions(
             ),
             labelnames=(TIER_LABEL,),
         )
+        metrics[TieringOffloadingMetrics.DEFERRED_RELEASE_CHUNKS] = (
+            OffloadingGaugeMetadata(
+                documentation=(
+                    "Number of primary-tier chunks of failed promotions whose release "
+                    "to the free pool is deferred until the tier confirms its I/O can "
+                    "no longer write to them, labeled by tier."
+                ),
+                labelnames=(TIER_LABEL,),
+            )
+        )
+        metrics[TieringOffloadingMetrics.DEFERRED_RELEASE_HOLD] = (
+            OffloadingHistogramMetadata(
+                documentation=(
+                    "Histogram of how long a failed promotion's primary-tier "
+                    "chunks were held before their deferred release completed, "
+                    "observed at release, labeled by tier, in seconds."
+                ),
+                labelnames=(TIER_LABEL,),
+                buckets=(0.01, 0.05, 0.1, 0.5, 1, 2.5, 5, 10, 20, 30, 60, 120),
+            )
+        )
         metrics[TieringOffloadingMetrics.PROMOTION_LATENCY] = (
             OffloadingHistogramMetadata(
                 documentation=(
@@ -519,6 +540,8 @@ def _uses_shared_region(self) -> bool:
 
     @override
     def create_worker(self, kv_caches: CanonicalKVCaches) -> CPUOffloadingWorker:
+        if self.secondary_tier_configs and self.config.parallel.world_size > 1:
+            _reject_workers_across_nodes()
         rank = 0 if self.replicated_layout else self.config.parallel.rank
         worker_mmap = SharedOffloadRegion(
             engine_id=self._engine_id,
@@ -578,3 +601,31 @@ def _validate_canonical_refs(self, kv_caches: CanonicalKVCaches) -> None:
             "Canonical KV layout enabled (all_layers_portable=%s)",
             self.all_layers_portable,
         )
+
+
+def _reject_workers_across_nodes() -> None:
+    """Reject secondary tiers on an engine whose workers span nodes.
+
+    Secondary tiers run in the engine core and read only its node's CPU
+    region, which lacks the slots of workers on other nodes (#54415). An
+    engine's workers form a grid of TP, PP and PCP groups, so they share one
+    node if every such group does. Collective over each group.
+    """
+    from vllm.distributed.parallel_state import (
+        get_pcp_group,
+        get_pp_group,
+        get_tp_group,
+        in_the_same_node_as,
+    )
+
+    # Every rank must join every collective before any of them raises.
+    spans = [
+        group.world_size > 1
+        and not all(in_the_same_node_as(group.cpu_group, source_rank=0))
+        for group in (get_tp_group(), get_pp_group(), get_pcp_group())
+    ]
+    if any(spans):
+        raise NotImplementedError(
+            "TieringOffloadingSpec secondary tiers are not supported on "
+            "engines whose workers span several nodes."
+        )
diff --git a/vllm/v1/metrics/ray_wrappers.py b/vllm/v1/metrics/ray_wrappers.py
index 5475f2ee9d72..8db746c96f4b 100644
--- a/vllm/v1/metrics/ray_wrappers.py
+++ b/vllm/v1/metrics/ray_wrappers.py
@@ -205,8 +205,3 @@ class RayPrometheusStatLogger(PrometheusStatLogger):
     _spec_decoding_cls = RaySpecDecodingProm
     _kv_connector_cls = RayKVConnectorProm
     _perf_metrics_cls = RayPerfMetricsProm
-
-    @staticmethod
-    def _unregister_vllm_metrics():
-        # No-op on purpose
-        pass
diff --git a/vllm/v1/request.py b/vllm/v1/request.py
index c27a8c4034f9..e5d5a0043510 100644
--- a/vllm/v1/request.py
+++ b/vllm/v1/request.py
@@ -80,6 +80,7 @@ def __init__(
         reasoning_parser_kwargs: dict[str, Any] | None = None,
         abort_immediately: bool = False,
         kv_hints: KvHintsEnvelope | None = None,
+        block_hashes: list["BlockHash"] | None = None,
     ) -> None:
         self.request_id = request_id
         self.client_index = client_index
@@ -222,7 +223,9 @@ def __init__(
         # add_request, then observed each verify step); stays None otherwise.
         self.spec_decode_metrics: RequestSpecDecodeMetrics | None = None
 
-        self.block_hashes: list[BlockHash] = []
+        # Prompt block hashes precomputed by the frontend; the hasher continues
+        # from the first block they don't cover.
+        self.block_hashes: list[BlockHash] = block_hashes or []
         # Store the block hasher without binding self to avoid creating a
         # reference cycle (Request -> partial -> Request) that prevents
         # immediate garbage collection via reference counting.
@@ -245,6 +248,7 @@ def from_engine_core_request(
         cls,
         request: EngineCoreRequest,
         block_hasher: Callable[["Request"], list["BlockHash"]] | None,
+        block_hashes: list["BlockHash"] | None = None,
     ) -> "Request":
         return cls(
             request_id=request.request_id,
@@ -267,6 +271,7 @@ def from_engine_core_request(
             reasoning_ended=request.reasoning_ended,
             reasoning_parser_kwargs=request.reasoning_parser_kwargs,
             abort_immediately=request.abort_immediately,
+            block_hashes=block_hashes,
         )
 
     def append_output_token_ids(
diff --git a/vllm/v1/simple_kv_offload/manager.py b/vllm/v1/simple_kv_offload/manager.py
index 51e6473edf41..a877f0a53002 100644
--- a/vllm/v1/simple_kv_offload/manager.py
+++ b/vllm/v1/simple_kv_offload/manager.py
@@ -245,6 +245,9 @@ def __init__(
         self._pending_cpu_hits: dict[
             str, tuple[tuple[list[KVCacheBlock], ...], int, int]
         ] = {}
+        # Eager mode: local prefix length from get_num_new_matched_tokens(),
+        # whose CPU copy update_state_after_alloc() refreshes.
+        self._pending_local_hits: dict[str, int] = {}
 
         # Store metadata
         self._lazy_mode = lazy_offload
@@ -342,10 +345,15 @@ 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.prefix_cacheable_groups:
+        for g in kv_cache_config.kv_cache_groups:
             spec = g.kv_cache_spec
             block_size = resolve_dcp_kv_block_size(spec, cp_world_size)
-            if isinstance(spec, MambaSpec):
+            if not spec.prefix_cacheable:
+                # Bounded-replay SWA (DSV4.1) is never offloaded, but its blocks for
+                # each scheduled token come from the same free queue.
+                if isinstance(spec, SlidingWindowSpec):
+                    target += cdiv(max_num_batched_tokens, block_size)
+            elif isinstance(spec, MambaSpec):
                 target += 2
             elif isinstance(spec, SlidingWindowSpec):
                 target += cdiv(spec.sliding_window, block_size) + 1
@@ -368,10 +376,16 @@ def get_num_new_matched_tokens(
         # is dropped first.
         if stale := self._pending_cpu_hits.pop(request.request_id, None):
             self._free_pending_cpu_hit(stale)
+        self._pending_local_hits.pop(request.request_id, None)
 
         if request.skip_reading_external_kv:
             return 0, False
 
+        if num_computed_tokens and not self._lazy_mode:
+            # Refreshed at admission: a request waiting for blocks repeats
+            # this lookup every step.
+            self._pending_local_hits[request.request_id] = num_computed_tokens
+
         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
@@ -444,6 +458,16 @@ def update_state_after_alloc(
                 num_stored_blocks=[0] * num_groups,
             )
 
+        if num_local_tokens := self._pending_local_hits.pop(req_id, None):
+            # Eager CPU is inclusive: a GPU prefix hit also uses the CPU copy,
+            # but eager stores skip it, so move it to the MRU end here.
+            prefix_hit, _, _ = self.cpu_coordinator.find_longest_cache_hit(
+                request.block_hashes, num_local_tokens
+            )
+            prefix_blocks = [b for grp in prefix_hit for b in grp if not b.is_null]
+            self.cpu_block_pool.touch(prefix_blocks)
+            self.cpu_block_pool.free_blocks(reversed(prefix_blocks))
+
         # Pop the CPU hit cached by get_num_new_matched_tokens(). The
         # found blocks were pinned there to survive LRU eviction in the window
         # between get_num_new_matched_tokens() and this matching call.
@@ -1227,6 +1251,7 @@ def request_finished(
         pending = self._pending_cpu_hits.pop(req_id, None)
         if pending is not None:
             self._free_pending_cpu_hit(pending)
+        self._pending_local_hits.pop(req_id, None)
 
         # Handle load: defer cleanup if load is in-flight
         load_state = self._reqs_to_load.get(req_id)
@@ -1430,6 +1455,7 @@ def reset(self) -> bool:
         for pending in self._pending_cpu_hits.values():
             self._free_pending_cpu_hit(pending)
         self._pending_cpu_hits.clear()
+        self._pending_local_hits.clear()
 
         self._abandoned_store_event_to_blocks.update(self._store_event_to_blocks)
         for transfer in self._pending_finished_stores:
diff --git a/vllm/v1/simple_kv_offload/worker.py b/vllm/v1/simple_kv_offload/worker.py
index 6a43c67ed133..af3f3312f5eb 100644
--- a/vllm/v1/simple_kv_offload/worker.py
+++ b/vllm/v1/simple_kv_offload/worker.py
@@ -244,7 +244,8 @@ def _init_cpu_mode(
             # Allocate non-pinned first, then pin via cudaHostRegister to
             # bypass PyTorch's CUDACachingHostAllocator which rounds up to
             # the next power of 2 (e.g. 100 GB -> 128 GB).
-            tensor = torch.zeros(cpu_shape, dtype=gpu_tensor.dtype, device="cpu")
+            # CPU blocks are loaded only after a completed GPU-to-CPU store.
+            tensor = torch.empty(cpu_shape, dtype=gpu_tensor.dtype, device="cpu")
             if pin_memory:
                 pin_tensor(tensor)
             self.cpu_kv_caches[name] = tensor
diff --git a/vllm/v1/spec_decode/draft_model.py b/vllm/v1/spec_decode/draft_model.py
index 17be2732e67c..0714fcab917d 100644
--- a/vllm/v1/spec_decode/draft_model.py
+++ b/vllm/v1/spec_decode/draft_model.py
@@ -9,6 +9,7 @@
 from vllm.config.utils import replace
 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.tokenizers.registry import get_tokenizer
 from vllm.v1.spec_decode.llm_base_proposer import SpecDecodeBaseProposer
 from vllm.v1.spec_decode.vocab_mapping import VocabMapping
@@ -108,6 +109,7 @@ def _get_model(self) -> nn.Module:
             model = get_model(
                 vllm_config=draft_vllm_config,
                 prefix="draft_model",
+                load_config=get_draft_load_config(draft_vllm_config),
             )
         return model
 
diff --git a/vllm/v1/structured_output/__init__.py b/vllm/v1/structured_output/__init__.py
index 69e34514728c..9911fa06add8 100644
--- a/vllm/v1/structured_output/__init__.py
+++ b/vllm/v1/structured_output/__init__.py
@@ -3,6 +3,7 @@
 import multiprocessing
 from collections.abc import Iterable, Sequence
 from concurrent.futures import Future, ThreadPoolExecutor
+from itertools import chain
 from typing import TYPE_CHECKING, NamedTuple
 
 import numpy as np
@@ -359,12 +360,15 @@ def _get_constraint_bounds(
                 return ConstraintBounds(i + 1, i + 1)
         return ConstraintBounds(1, 1)
 
-    def validate_tokens(self, request: "Request", spec_tokens: list[int]) -> list[int]:
-        """Return the longest unconstrained or grammar-valid prefix of `spec_tokens`."""
+    def trim_spec_tokens(self, request: "Request", spec_tokens: list[int]) -> list[int]:
+        """Trim structured-output padding and grammar-invalid suffixes."""
         if not request.use_structured_output:
+            # Synchronous unstructured decoding uses -1 values as scheduler-side
+            # markers for speculative slots, so preserve their length.
             return spec_tokens
 
         spec_tokens = strip_speculative_padding(spec_tokens)
+
         grammar_start, _ = self._get_constraint_bounds(request, spec_tokens)
         if grammar_start >= len(spec_tokens):
             return spec_tokens
@@ -375,9 +379,11 @@ def validate_tokens(self, request: "Request", spec_tokens: list[int]) -> list[in
         grammar = structured_req.grammar
         if TYPE_CHECKING:
             assert isinstance(grammar, StructuredOutputGrammar)
+        if grammar.is_terminated():
+            return spec_tokens[:grammar_start]
         prefix = spec_tokens[:grammar_start]
-        validated = grammar.validate_tokens(spec_tokens[grammar_start:])
-        return prefix + validated
+        len_validated = len(grammar.validate_tokens(spec_tokens[grammar_start:]))
+        return prefix + spec_tokens[grammar_start : grammar_start + len_validated + 1]
 
     def grammar_bitmask(
         self,
@@ -449,7 +455,7 @@ def grammar_bitmask(
         else:
             # Fallback to serial filling of bitmasks for small-batch-size cases
             for grammar_idx, req_id in enumerate(structured_output_request_ids):
-                cumulative_index = grammar_idx * mask_stride
+                base_index = grammar_idx * mask_stride
                 request = requests[req_id]
                 structured_output_request = request.structured_output_request
 
@@ -460,62 +466,43 @@ def grammar_bitmask(
                     assert isinstance(grammar, StructuredOutputGrammar)
 
                 req_tokens = scheduled_spec_decode_tokens.get(req_id, list())
+                unpadded_tokens = strip_speculative_padding(req_tokens)
                 grammar_start, constraint_start = self._get_constraint_bounds(
-                    request, strip_speculative_padding(req_tokens)
+                    request, unpadded_tokens
                 )
                 # Rows before constraint_start are still reasoning.
                 masked_ids = (
                     self._masked_reasoning_ends(request) if constraint_start else ()
                 )
+                # Diffusion LLMs don't sample a bonus token after the
+                # scheduled positions, so skip its bitmask in that case.
+                tokens_to_mask = (
+                    chain(req_tokens, (-1,))
+                    if not (is_diffusion and req_tokens)
+                    else req_tokens
+                )
+
                 state_advancements = 0
-                seen_padding = False
-                # Row filled from the last valid grammar state before a draft
-                # was rejected; later rows reuse it rather than unconstraining.
-                failed_index: int | None = None
-                bitmask = self._grammar_bitmask
-                for i, token in enumerate(req_tokens):
-                    if failed_index is not None:
-                        bitmask[cumulative_index].copy_(bitmask[failed_index])
-                        cumulative_index += 1
+                padding_or_grammar_invalid = False
+                for i, token in enumerate(tokens_to_mask):
+                    if padding_or_grammar_invalid:
+                        self._fill_bitmasks(((grammar, base_index + i, False, ()),))
                         continue
-                    apply_bitmask = not seen_padding and i >= constraint_start
-                    reasoning = not seen_padding and i < constraint_start
-                    row_masked = masked_ids if reasoning else ()
+
+                    apply_bitmask = i >= constraint_start
+                    row_masked = masked_ids if i < constraint_start else ()
                     self._fill_bitmasks(
-                        ((grammar, cumulative_index, apply_bitmask, row_masked),)
+                        ((grammar, base_index + i, apply_bitmask, row_masked),)
                     )
                     if token == -1:
-                        seen_padding = True
-                    elif not seen_padding and i >= grammar_start:
+                        padding_or_grammar_invalid = True
+                    elif i >= grammar_start:
                         if not grammar.is_terminated() and grammar.accept_tokens(
                             req_id, [token]
                         ):
                             state_advancements += 1
                         else:
-                            failed_index = cumulative_index
-                            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 (is_diffusion and req_tokens):
-                    if failed_index is not None:
-                        bitmask[cumulative_index].copy_(bitmask[failed_index])
-                    else:
-                        bonus_apply = not seen_padding and constraint_start <= len(
-                            req_tokens
-                        )
-                        bonus_reasoning = not seen_padding and not bonus_apply
-                        bonus_masked = masked_ids if bonus_reasoning else ()
-                        self._fill_bitmasks(
-                            ((grammar, cumulative_index, bonus_apply, bonus_masked),)
-                        )
-                    cumulative_index += 1
+                            padding_or_grammar_invalid = True
 
                 if state_advancements > 0:
                     grammar.rollback(state_advancements)
diff --git a/vllm/v1/structured_output/backend_guidance.py b/vllm/v1/structured_output/backend_guidance.py
index 08aab576df1c..95c34e477df9 100644
--- a/vllm/v1/structured_output/backend_guidance.py
+++ b/vllm/v1/structured_output/backend_guidance.py
@@ -236,12 +236,15 @@ 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
-            self.terminated = True
+        eos_token = self.ll_tokenizer.eos_token
+        eos_index = next(
+            (index for index, token in enumerate(tokens) if token == eos_token),
+            None,
+        )
+        if eos_index is not None and eos_index != len(tokens) - 1:
+            return False
 
-        if self.ll_matcher.is_stopped():
+        if self.terminated:
             return True
 
         # TODO - Add jump decoding support in the future:
@@ -251,11 +254,31 @@ def accept_tokens(self, request_id: str, tokens: list[int]) -> bool:
         # For conversion between the two, see
         # https://github.com/guidance-ai/llguidance/blob/main/docs/fast_forward.md
 
-        r = self.ll_matcher.consume_tokens(tokens)
+        num_tokens = self.ll_matcher.try_consume_tokens(tokens)
+        if num_tokens != len(tokens):
+            if (
+                eos_index is not None
+                and num_tokens == eos_index
+                and self.ll_matcher.is_stopped()
+                and self.ll_matcher.is_accepting()
+            ):
+                # LLMatcher stops at an accepting grammar state and leaves the
+                # wrapper-level EOS token unconsumed.
+                self.rollback_lag = 1
+                self.terminated = True
+                self.check_error()
+                return True
+            if num_tokens > 0:
+                self.ll_matcher.rollback(num_tokens)
+            self.check_error()
+            return False
+
+        if eos_index is not None:
+            self.terminated = True
 
         self.check_error()
 
-        return r
+        return True
 
     def validate_tokens(self, tokens: list[int]) -> list[int]:
         """Checks if the list of tokens are accepted by the parser in sequence.
diff --git a/vllm/v1/structured_output/backend_types.py b/vllm/v1/structured_output/backend_types.py
index b3695efd009c..6291e4a32ad0 100644
--- a/vllm/v1/structured_output/backend_types.py
+++ b/vllm/v1/structured_output/backend_types.py
@@ -36,6 +36,8 @@ def accept_tokens(self, request_id: str, tokens: list[int]) -> bool:
         """Determines whether the provided tokens are accepted for the
         given request.
 
+        A false result must leave the grammar state unchanged.
+
         Args:
             request_id (str): The unique identifier for the request.
             tokens (list[int]): A list of token IDs to evaluate.
diff --git a/vllm/v1/structured_output/backend_xgrammar.py b/vllm/v1/structured_output/backend_xgrammar.py
index aae72873b2e6..3772d371ad01 100644
--- a/vllm/v1/structured_output/backend_xgrammar.py
+++ b/vllm/v1/structured_output/backend_xgrammar.py
@@ -270,14 +270,12 @@ def accept_tokens(self, request_id: str, tokens: list[int]) -> bool:
         """
         if self._is_terminated:
             return True
-        for token in tokens:
+        for num_accepted, token in enumerate(tokens):
             if not self.matcher.accept_token(token):
-                logger.error(
-                    "Failed to advance FSM for request %s "
-                    "for tokens %s. Please file an issue.",
-                    request_id,
-                    token,
-                )
+                if num_accepted:
+                    self.matcher.rollback(num_accepted)
+                    self.num_processed_tokens -= num_accepted
+                    self._is_terminated = self.matcher.is_terminated()
                 return False
             self.num_processed_tokens += 1
             self._is_terminated = self.matcher.is_terminated()
diff --git a/vllm/v1/utils.py b/vllm/v1/utils.py
index fc6f3a52a1d6..4d7978dabc67 100644
--- a/vllm/v1/utils.py
+++ b/vllm/v1/utils.py
@@ -740,7 +740,11 @@ def report_usage_stats(
             # Common configuration
             "dtype": str(model_config.dtype),
             "block_size": vllm_config.cache_config.block_size,
-            "gpu_memory_utilization": vllm_config.cache_config.gpu_memory_utilization,
+            # Reported at worker init: the extensible KV cache default (1.0) is
+            # still shown if the driver probe later falls back to 0.92.
+            "gpu_memory_utilization": (
+                vllm_config.cache_config.resolved_gpu_memory_utilization
+            ),
             "kv_cache_memory_bytes": vllm_config.cache_config.kv_cache_memory_bytes,
             # Quantization
             "quantization": model_config.quantization,
diff --git a/vllm/v1/watermarking/spec_decode.py b/vllm/v1/watermarking/spec_decode.py
index 6c432db959f3..05b8c255c9e3 100644
--- a/vllm/v1/watermarking/spec_decode.py
+++ b/vllm/v1/watermarking/spec_decode.py
@@ -14,6 +14,7 @@
     Watermarker,
 )
 from vllm.v1.worker.gpu.sample.watermark import draft_watermarking_mask
+from vllm.v1.worker.gpu.spec_decode.draft_logits import DraftLogitsCache
 from vllm.v1.worker.gpu.spec_decode.rejection_sampler_utils import rejection_sample
 
 
@@ -293,7 +294,7 @@ def speculative_target_watermark_key(
 
 def watermarked_rejection_sample(
     target_logits: torch.Tensor,
-    draft_logits: torch.Tensor | None,
+    draft_logits_cache: DraftLogitsCache | None,
     draft_sampled: torch.Tensor,
     cu_num_logits: torch.Tensor,
     pos: torch.Tensor,
@@ -313,7 +314,7 @@ def watermarked_rejection_sample(
     assert contexts.shape == (target_logits.shape[0], watermarker.context_width)
     return rejection_sample(
         target_logits,
-        draft_logits,
+        draft_logits_cache,
         draft_sampled,
         cu_num_logits,
         pos,
diff --git a/vllm/v1/worker/cpu/attn_utils.py b/vllm/v1/worker/cpu/attn_utils.py
new file mode 100644
index 000000000000..15dbf5c48591
--- /dev/null
+++ b/vllm/v1/worker/cpu/attn_utils.py
@@ -0,0 +1,357 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+from collections.abc import Iterable, Sequence
+from contextlib import AbstractContextManager, nullcontext
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, cast
+
+import numpy as np
+import torch
+
+import vllm.envs as envs
+from vllm.config import VllmConfig, get_layers_from_vllm_config
+from vllm.model_executor.layers.attention import Attention
+from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
+from vllm.v1.attention.backend import CommonAttentionMetadata
+from vllm.v1.attention.backends.utils import (
+    create_fast_prefill_custom_backend,
+)
+from vllm.v1.kv_cache_interface import (
+    AttentionSpec,
+    KVCacheConfig,
+    KVCacheSpec,
+    MambaSpec,
+    UniformTypeKVCacheSpecs,
+)
+from vllm.v1.worker.utils import (
+    AttentionGroup,
+    add_kv_sharing_layers_to_kv_cache_groups,
+    allocate_kv_cache,
+    bind_kv_cache_to_layers,
+    customize_attention_spec,
+    map_kv_caches_to_kernel_blocks,
+    prepare_kernel_block_sizes,
+)
+
+if TYPE_CHECKING:
+    from vllm.v1.worker.cpu.model_states.interface import ModelSpecificAttnMetadata
+
+
+@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."""
+
+    def prepare(
+        self,
+        logits_indices: torch.Tensor,
+        cu_num_logits_np: np.ndarray,
+        has_prefill: bool,
+    ) -> FastPrefillBatchMetadata | None:
+        if not has_prefill:
+            return None
+
+        num_logits = logits_indices.shape[0]
+        assert num_logits > 0
+        return FastPrefillBatchMetadata(
+            logits_indices_padded=logits_indices,
+            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)
+    attn_layers = get_layers_from_vllm_config(vllm_config, layer_type)
+    for layer_name, attn_module in attn_layers.items():
+        if getattr(attn_module, "kv_sharing_target_layer_name", None):
+            # This layer will use KV cache of the sharing target layer.
+            continue
+        # Skip modules that don't need KV cache (eg encoder-only attention)
+        if spec := attn_module.get_kv_cache_spec(vllm_config):
+            if isinstance(spec, AttentionSpec):
+                spec = customize_attention_spec(attn_module.get_attn_backend(), spec)
+            kv_cache_spec[layer_name] = spec
+    return kv_cache_spec
+
+
+def get_shared_kv_cache_layers(vllm_config: VllmConfig):
+    attn_layers = get_layers_from_vllm_config(vllm_config, Attention)
+    return {
+        layer_name: kv_tgt_layer
+        for layer_name, attn_module in attn_layers.items()
+        if (kv_tgt_layer := attn_module.kv_sharing_target_layer_name)
+    }
+
+
+def get_kv_sharing_fast_prefill_eligible_layers(
+    vllm_config: VllmConfig,
+) -> 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.
+    """
+    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 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,
+) -> tuple[list[list[AttentionGroup]], list[int]]:
+    # Phase 1: discover attention groups for each kv cache group.
+    attn_groups: list[list[AttentionGroup]] = []
+
+    # Add KV-sharing layers to their target's kv cache group so they are
+    # discovered alongside the target layer in Phase 1 below.
+    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
+    )
+
+    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
+
+        layer_type = cast(type[Any], AttentionLayerBase)
+        attn_layers = get_layers_from_vllm_config(vllm_config, layer_type, layer_names)
+        group_map: dict[tuple[tuple[str, str], KVCacheSpec, int], AttentionGroup] = {}
+        group_order: list[tuple[tuple[str, str], KVCacheSpec, int]] = []
+
+        for layer_name in layer_names:
+            attn_backend = attn_layers[layer_name].get_attn_backend()
+            if (
+                envs.VLLM_BATCH_INVARIANT
+                and not attn_backend.supports_batch_invariance()
+            ):
+                raise RuntimeError(
+                    "VLLM batch_invariant mode is not supported for "
+                    f"{attn_backend.get_name()}."
+                )
+            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):
+                layer_kv_cache_spec = layer_kv_cache_spec.kv_cache_specs[layer_name]
+
+            # Split on per-rank num_heads_q so layers with different Q-head
+            # counts get separate metadata builders.
+            num_heads_q = getattr(attn_layers[layer_name], "num_heads", 0)
+            key = (attn_backend.full_cls_name(), layer_kv_cache_spec, num_heads_q)
+            if key not in group_map:
+                group_map[key] = AttentionGroup(
+                    attn_backend, [layer_name], layer_kv_cache_spec, kv_cache_group_id
+                )
+                group_order.append(key)
+            else:
+                group_map[key].layer_names.append(layer_name)
+
+        attn_groups.append([group_map[key] for key in group_order])
+
+    # Phase 2: pick a kernel block size per kv cache group that is supported
+    # by all backends within that group.
+    kernel_block_sizes = prepare_kernel_block_sizes(kv_cache_config, attn_groups)
+
+    # Phase 3: create metadata builders.
+    for kv_cache_group_id, groups in enumerate(attn_groups):
+        kernel_block_size: int | None = None
+        if kv_cache_group_id < len(kernel_block_sizes):
+            kernel_block_size = kernel_block_sizes[kv_cache_group_id]
+        for group in groups:
+            group.create_metadata_builders(
+                vllm_config=vllm_config,
+                device=device,
+                kernel_block_size=kernel_block_size,
+            )
+
+    return attn_groups, kernel_block_sizes
+
+
+def init_kv_cache(
+    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,
+    *,
+    attn_groups: Iterable[AttentionGroup] = (),
+) -> 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,
+        )
+    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
+    # decoder layer, so a layer name carries two integers (layer + module index).
+    num_attn_module = (
+        2 if vllm_config.model_config.hf_config.model_type == "longcat_flash" else 1
+    )
+    bindable_caches = map_kv_caches_to_kernel_blocks(
+        {name: cache for name, cache in kv_caches.items() if name in forward_context},
+        attn_groups,
+    )
+    bind_kv_cache_to_layers(
+        bindable_caches,
+        forward_context,
+        num_attn_module,
+        kv_cache_groups=kv_cache_config.kv_cache_groups,
+    )
+    return kv_caches
+
+
+def build_slot_mappings_by_layer(
+    slot_mappings: torch.Tensor, kv_cache_config: KVCacheConfig
+) -> dict[str, torch.Tensor]:
+    slot_mappings_by_layer: dict[str, torch.Tensor] = {}
+    kv_cache_groups = kv_cache_config.kv_cache_groups
+    for slot_mapping, kv_cache_group in zip(slot_mappings, kv_cache_groups):
+        for layer_name in kv_cache_group.layer_names:
+            slot_mappings_by_layer[layer_name] = slot_mapping
+    return slot_mappings_by_layer
+
+
+def build_attn_metadata(
+    attn_groups: list[list[AttentionGroup]],
+    num_reqs: int,
+    num_tokens: int,
+    query_start_loc: torch.Tensor,
+    max_query_len: int,
+    seq_lens: torch.Tensor,
+    max_seq_len: int,
+    block_tables: Sequence[torch.Tensor],
+    slot_mappings: torch.Tensor,
+    kv_cache_config: KVCacheConfig,
+    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,
+    model_specific_attn_metadata: ModelSpecificAttnMetadata | None = None,
+    causal: bool | torch.Tensor = True,
+    rswa_prefix_lens: torch.Tensor | None = None,
+    fast_prefill: FastPrefillBatchMetadata | None = None,
+    req_idx: np.ndarray | None = None,
+) -> dict[str, Any]:
+    seq_lens = seq_lens[: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
+    # Mamba groups with the same spec and builder differ only in their state
+    # indices, so later groups re-gather those from the first group's metadata.
+    cached_metadata: dict[tuple[KVCacheSpec, type], Any] = {}
+    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]
+        common_attn_metadata_extra_kwargs = (
+            model_specific_attn_metadata.get_extra_common_attn_kwargs(i, num_reqs)
+            if model_specific_attn_metadata is not None
+            else {}
+        )
+        # Model-specific metadata (e.g. Mamba hybrid) may supply its own
+        # padding-aware is_prefilling, which takes precedence over the default.
+        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,
+            query_start_loc_cpu=query_start_loc,
+            seq_lens=seq_lens,
+            seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,
+            max_seq_len=max_seq_len,
+            num_reqs=num_reqs,
+            num_actual_tokens=num_tokens,
+            max_query_len=max_query_len,
+            block_table_tensor=block_table,
+            slot_mapping=slot_mapping,
+            causal=causal,
+            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,
+        )
+
+        for attn_group in attn_groups[i]:
+            attn_metadata_builder = attn_group.get_metadata_builder()
+            reuse_key = None
+            if attn_metadata_builder.supports_update_block_table and isinstance(
+                attn_group.kv_cache_spec, MambaSpec
+            ):
+                reuse_key = (attn_group.kv_cache_spec, type(attn_metadata_builder))
+            if reuse_key in cached_metadata:
+                metadata = attn_metadata_builder.update_block_table(
+                    cached_metadata[reuse_key], block_table, slot_mapping
+                )
+            else:
+                attn_metadata_extra_kwargs = (
+                    model_specific_attn_metadata.get_extra_attn_kwargs(
+                        attn_metadata_builder,
+                        num_reqs,
+                    )
+                    if model_specific_attn_metadata is not None
+                    else {}
+                )
+                metadata = attn_group.build_metadata(
+                    common_attn_metadata, **attn_metadata_extra_kwargs
+                )
+            if reuse_key is not None and reuse_key not in cached_metadata:
+                cached_metadata[reuse_key] = 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/cpu/block_table.py b/vllm/v1/worker/cpu/block_table.py
new file mode 100644
index 000000000000..c184f9267381
--- /dev/null
+++ b/vllm/v1/worker/cpu/block_table.py
@@ -0,0 +1,185 @@
+# 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.v1.attention.backends.utils import PAD_SLOT_ID
+
+
+class CPUBlockTables:
+    # Set by the elastic EP warmup so KV writes land in the null block.
+    redirect_writes_to_null_block = False
+
+    def __init__(
+        self,
+        block_sizes: list[int],
+        max_num_reqs: int,
+        max_num_batched_tokens: int,
+        max_num_blocks_per_group: list[int],
+        device: torch.device,
+        slot_mapping_enabled: list[bool] | None = None,
+        sliding_window_reach: list[int] | None = None,
+    ):
+        self.block_sizes = block_sizes
+        self.max_num_reqs = max_num_reqs
+        self.max_num_batched_tokens = max_num_batched_tokens
+        self.device = device
+        self.num_kv_cache_groups = len(self.block_sizes)
+        assert len(max_num_blocks_per_group) == self.num_kv_cache_groups
+        if slot_mapping_enabled is None:
+            slot_mapping_enabled = [True] * self.num_kv_cache_groups
+        assert len(slot_mapping_enabled) == self.num_kv_cache_groups
+        self._slot_mapping_enabled = slot_mapping_enabled
+        if sliding_window_reach is None:
+            sliding_window_reach = [0] * self.num_kv_cache_groups
+        assert len(sliding_window_reach) == self.num_kv_cache_groups
+        self._sliding_window_reach = sliding_window_reach
+        self.has_sliding_window_groups = any(r > 0 for r in sliding_window_reach)
+
+        # num_kv_cache_groups x [max_num_reqs, max_num_blocks]
+        self.block_tables: list[torch.Tensor] = []
+        for max_num_blocks in max_num_blocks_per_group:
+            block_table = torch.zeros(
+                (self.max_num_reqs, max_num_blocks),
+                dtype=torch.int32,
+                device=device,
+            )
+            self.block_tables.append(block_table)
+        self.block_tables_np = [table.numpy() for table in self.block_tables]
+        self.num_blocks = torch.zeros(
+            (self.num_kv_cache_groups, self.max_num_reqs),
+            dtype=torch.int32,
+            device=device,
+        )
+        self.num_blocks_np = self.num_blocks.numpy()
+
+        # Block tables used for model's forward pass.
+        # num_kv_cache_groups x [max_num_reqs, max_num_blocks]
+        self.input_block_tables: list[torch.Tensor] = [
+            torch.zeros_like(block_table) for block_table in self.block_tables
+        ]
+        self._num_dirty_rows = 0
+        self.slot_mappings = torch.zeros(
+            self.num_kv_cache_groups,
+            self.max_num_batched_tokens,
+            dtype=torch.int64,
+            device=self.device,
+        )
+
+    def append_block_ids(
+        self,
+        req_index: int,
+        new_block_ids: tuple[list[int], ...],
+        overwrite: bool,
+    ) -> None:
+        for i in range(self.num_kv_cache_groups):
+            start = self.num_blocks_np[i, req_index] if not overwrite else 0
+            block_ids = new_block_ids[i]
+            if self.redirect_writes_to_null_block:
+                block_ids = [0] * len(block_ids)
+            end = start + len(block_ids)
+            row_capacity = self.block_tables[i].shape[1]
+            if end > row_capacity:
+                raise RuntimeError(
+                    f"Block table write for request {req_index}, group {i} exceeds "
+                    f"row capacity ({end} > {row_capacity})"
+                )
+            self.block_tables_np[i][req_index, start:end] = block_ids
+            self.num_blocks_np[i, req_index] = end
+
+    def gather_block_tables(
+        self,
+        idx_mapping: torch.Tensor,
+        num_reqs_padded: int,
+        out: tuple[torch.Tensor, ...] | None = None,
+        num_computed_tokens: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, ...]:
+        """Gather the batch's block-table rows into the forward-pass tables.
+
+        When num_computed_tokens ([max_num_reqs], indexed like the block tables)
+        is given, sliding-window groups get every entry wholly below the window
+        of the request's first scheduled token replaced by the null block. Those
+        blocks may already have been freed by the sliding-window manager and
+        reused, while attention kernels can still load them as part of a KV
+        tile that straddles the window start: masked-out NaN in such a tile
+        propagates into the output.
+        """
+        if self.num_kv_cache_groups == 0:
+            return ()
+        mask_sliding_window = (
+            self.has_sliding_window_groups and num_computed_tokens is not None
+        )
+        num_reqs = idx_mapping.shape[0]
+        num_rows = num_reqs_padded
+        if out is None:
+            out = tuple(self.input_block_tables)
+            num_rows = max(num_rows, self._num_dirty_rows)
+            self._num_dirty_rows = num_reqs
+        else:
+            assert len(out) == self.num_kv_cache_groups
+
+        for group_id, output in enumerate(out):
+            output[:num_rows].zero_()
+            ops.gather_block_table_cpu(
+                self.block_tables[group_id],
+                output,
+                idx_mapping,
+                self.num_blocks[group_id],
+                num_rows,
+            )
+            reach = self._sliding_window_reach[group_id]
+            if mask_sliding_window and reach > 0 and num_reqs > 0:
+                assert num_computed_tokens is not None
+                first_positions = num_computed_tokens[idx_mapping].sub(reach)
+                num_masked_blocks = first_positions.clamp_min_(0).div(
+                    self.block_sizes[group_id], rounding_mode="floor"
+                )
+                block_indices = torch.arange(
+                    output.shape[1], dtype=torch.int32, device=output.device
+                )
+                output[:num_reqs].masked_fill_(
+                    block_indices.unsqueeze(0) < num_masked_blocks.unsqueeze(1), 0
+                )
+        return tuple(block_table[:num_reqs_padded] for block_table in out)
+
+    def get_dummy_block_tables(self, num_reqs: int) -> tuple[torch.Tensor, ...]:
+        # Zero the rows so dummy runs write mamba state to the reserved null
+        # block rather than through the previous real step's (stale) block
+        # ids, which may point at blocks since freed and reallocated.
+        num_rows = max(num_reqs, self._num_dirty_rows)
+        for block_table in self.input_block_tables:
+            block_table[:num_rows].zero_()
+        # Profiling may populate the returned rows with valid block IDs.
+        self._num_dirty_rows = num_reqs
+        return tuple(block_table[:num_reqs] for block_table in self.input_block_tables)
+
+    def compute_slot_mappings(
+        self,
+        idx_mapping: torch.Tensor,
+        query_start_loc: torch.Tensor,
+        positions: torch.Tensor,
+        num_tokens_padded: int,
+        out: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        if self.num_kv_cache_groups == 0:
+            return (self.slot_mappings if out is None else out)[:, :num_tokens_padded]
+        slot_mappings = self.slot_mappings if out is None else out
+        for group_id in range(self.num_kv_cache_groups):
+            ops.compute_slot_mapping_cpu(
+                self.block_tables[group_id],
+                slot_mappings[group_id],
+                idx_mapping,
+                query_start_loc,
+                positions,
+                self.block_sizes[group_id],
+                self._slot_mapping_enabled[group_id],
+                PAD_SLOT_ID,
+            )
+        return slot_mappings[:, :num_tokens_padded]
+
+    def get_dummy_slot_mappings(self, num_tokens: int) -> torch.Tensor:
+        # Fill the entire slot_mappings tensor, not just the first `num_tokens` entries.
+        # This is because the padding logic is complex and kernels may access beyond
+        # the requested range.
+        self.slot_mappings.fill_(PAD_SLOT_ID)
+        return self.slot_mappings[:, :num_tokens]
diff --git a/vllm/v1/worker/cpu/ec_connector.py b/vllm/v1/worker/cpu/ec_connector.py
new file mode 100644
index 000000000000..5c5e596741fa
--- /dev/null
+++ b/vllm/v1/worker/cpu/ec_connector.py
@@ -0,0 +1,106 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from collections.abc import Generator
+from contextlib import contextmanager
+from typing import TYPE_CHECKING
+
+import torch
+
+from vllm.config import VllmConfig
+from vllm.distributed.ec_transfer import get_ec_transfer, has_ec_transfer
+from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorBase
+from vllm.v1.outputs import (
+    EMPTY_MODEL_RUNNER_OUTPUT,
+    ECConnectorOutput,
+    ModelRunnerOutput,
+)
+
+if TYPE_CHECKING:
+    from vllm.v1.core.sched.output import SchedulerOutput
+    from vllm.v1.worker.cpu.encoder_cache import EncoderCache
+
+
+class ECConnector:
+    """EC connector interface used by the V2 CPU model runner."""
+
+    @contextmanager
+    def maybe_get_output(
+        self, scheduler_output: "SchedulerOutput"
+    ) -> Generator[ECConnectorOutput | None, None, None]:
+        yield None
+
+    def no_forward(self, scheduler_output: "SchedulerOutput") -> ModelRunnerOutput:
+        return EMPTY_MODEL_RUNNER_OUTPUT
+
+
+class ActiveECConnector(ECConnector):
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        encoder_cache: dict[str, torch.Tensor],
+    ) -> None:
+        self.encoder_cache = encoder_cache
+        self.ec_connector = get_ec_transfer()
+        assert isinstance(self.ec_connector, ECConnectorBase)
+        # Every producer offloads freshly computed encoder outputs, including
+        # an ec_both node that also reloads them.
+        self.save_new_caches = self.ec_connector.is_producer
+
+    @contextmanager
+    def maybe_get_output(
+        self, scheduler_output: "SchedulerOutput"
+    ) -> Generator[ECConnectorOutput | None, None, None]:
+        if scheduler_output.ec_connector_metadata is None:
+            yield None
+            return
+
+        output = ECConnectorOutput()
+        ec_connector = self.ec_connector
+        assert scheduler_output.ec_connector_metadata is not None
+        ec_connector.bind_connector_metadata(scheduler_output.ec_connector_metadata)
+
+        if ec_connector.is_producer:
+            ec_connector.start_save_caches(encoder_cache=self.encoder_cache)
+
+        if ec_connector.is_consumer:
+            ec_connector.start_load_caches(self.encoder_cache)
+
+        cached_hashes = set(self.encoder_cache) if self.save_new_caches else None
+        try:
+            yield output
+            if cached_hashes is not None:
+                for mm_hash in self.encoder_cache.keys() - cached_hashes:
+                    ec_connector.save_caches(
+                        encoder_cache=self.encoder_cache, mm_hash=mm_hash
+                    )
+        finally:
+            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:
+        # EC send/recv even if no work to do.
+        with self.maybe_get_output(scheduler_output) as ec_connector_output:
+            pass
+
+        return ModelRunnerOutput.with_ec_conn_output_only(ec_connector_output)
+
+
+NO_OP_EC_CONNECTOR = ECConnector()
+
+
+def get_ec_connector(
+    vllm_config: VllmConfig, encoder_cache: "EncoderCache | None"
+) -> ECConnector:
+    if (
+        not has_ec_transfer()
+        or vllm_config.model_config.is_encoder_decoder
+        or encoder_cache is None
+    ):
+        return NO_OP_EC_CONNECTOR
+
+    return ActiveECConnector(vllm_config, encoder_cache.encoder_outputs)
diff --git a/vllm/v1/worker/cpu/encoder_cache.py b/vllm/v1/worker/cpu/encoder_cache.py
new file mode 100644
index 000000000000..6ef641b9d9cb
--- /dev/null
+++ b/vllm/v1/worker/cpu/encoder_cache.py
@@ -0,0 +1,54 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from dataclasses import replace
+
+import torch
+
+from vllm.multimodal.inputs import MultiModalFeatureSpec, MultiModalKwargsItem
+
+
+class EncoderCache:
+    def __init__(self):
+        # req_id -> MM features
+        self.mm_features: dict[str, list[MultiModalFeatureSpec]] = {}
+        # MM hash -> encoder outputs
+        self.encoder_outputs: dict[str, torch.Tensor] = {}
+
+    def __len__(self) -> int:
+        return len(self.encoder_outputs)
+
+    def add_request(
+        self, req_id: str, mm_features: list[MultiModalFeatureSpec]
+    ) -> None:
+        # The uniprocess executor can share this list with EngineCore.
+        self.mm_features[req_id] = list(mm_features)
+
+    def remove_request(self, req_id: str) -> None:
+        self.mm_features.pop(req_id, None)
+
+    def update_encoder_inputs(
+        self, req_id: str, updates: dict[int, MultiModalKwargsItem | None]
+    ) -> None:
+        features = self.mm_features[req_id]
+        for input_id, data in updates.items():
+            # Replace entries so EngineCore keeps its original payload for replay.
+            features[input_id] = replace(features[input_id], data=data)
+
+    def reset_mm_cache(self) -> None:
+        """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
+        pass
+
+    def reset_encoder_cache(self) -> None:
+        """Clear the encoder cache storing vision embeddings.
+
+        This should be called when model weights are updated to ensure
+        stale embeddings computed with old weights are not reused.
+        """
+        self.encoder_outputs.clear()
+
+    def free_encoder_cache(self, mm_hash: str) -> None:
+        self.encoder_outputs.pop(mm_hash, None)
diff --git a/vllm/v1/worker/cpu/encoder_runner.py b/vllm/v1/worker/cpu/encoder_runner.py
new file mode 100644
index 000000000000..997f457da988
--- /dev/null
+++ b/vllm/v1/worker/cpu/encoder_runner.py
@@ -0,0 +1,247 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import threading
+import time
+from collections.abc import Collection
+from contextlib import contextmanager
+
+import numpy as np
+import torch
+
+from vllm.logger import init_logger
+from vllm.model_executor.models.interfaces import SupportsMultiModal, supports_realtime
+from vllm.multimodal.encoder_budget import MultiModalBudget
+from vllm.multimodal.inputs import MultiModalKwargsItem
+from vllm.multimodal.utils import (
+    get_mm_features_in_window,
+    group_and_batch_mm_kwargs,
+    set_mm_embedding_modality,
+)
+from vllm.v1.worker.cpu.encoder_cache import EncoderCache
+from vllm.v1.worker.utils import (
+    EncoderTimingStats,
+    sanity_check_mm_encoder_outputs,
+)
+
+logger = init_logger(__name__)
+
+
+class EncoderRunner:
+    def __init__(
+        self,
+        model: SupportsMultiModal,
+        max_num_tokens: int,
+        hidden_size: int,
+        encoder_cache: EncoderCache,
+        dtype: torch.dtype,
+        device: torch.device,
+        enable_timing: bool = False,
+    ):
+        self.model = model
+        self.max_num_tokens = max_num_tokens
+        self.hidden_size = hidden_size
+        self.encoder_cache = encoder_cache
+        self.dtype = dtype
+        self.device = device
+        self.is_realtime = supports_realtime(model)
+        self.enable_timing = enable_timing
+        self.encoder_timing_registry: dict[str, EncoderTimingStats] = {}
+        self._timing_lock = threading.Lock()
+
+        self.inputs_embeds = torch.zeros(
+            max_num_tokens, hidden_size, dtype=dtype, device=device
+        )
+
+    def prepare_mm_inputs(
+        self, scheduled_encoder_inputs: dict[str, list[int]]
+    ) -> tuple[list[str], list[tuple[str, MultiModalKwargsItem]]]:
+        mm_hashes: list[str] = []
+        mm_kwargs: list[tuple[str, MultiModalKwargsItem]] = []
+        for req_id, encoder_input_ids in scheduled_encoder_inputs.items():
+            mm_features = self.encoder_cache.mm_features[req_id]
+            for mm_input_id in encoder_input_ids:
+                mm_feature = mm_features[mm_input_id]
+                if mm_feature.data is None:
+                    continue
+                if mm_feature.identifier in self.encoder_cache.encoder_outputs:
+                    continue
+                mm_hashes.append(mm_feature.identifier)
+                mm_kwargs.append((mm_feature.modality, mm_feature.data))
+
+        return mm_hashes, mm_kwargs
+
+    @torch.inference_mode()
+    def profile_encoder_cache(
+        self,
+        dummy_mm_inputs: list[tuple[str, MultiModalKwargsItem]],
+        budget: MultiModalBudget,
+    ) -> None:
+        """Profile multimodal encoder and temporary encoder cache memory."""
+        if (encoder_budget := budget.get_encoder_budget()) <= 0:
+            return
+
+        if not budget.mm_max_toks_per_item:
+            logger.info(
+                "Skipping encoder profiling for embedding-only mode "
+                "(all modality limits=0 with enable_mm_embeds=True).",
+            )
+            return
+
+        assert dummy_mm_inputs, "Dummy inputs should be generated for encoder profiling"
+        dummy_modality = dummy_mm_inputs[0][0]
+        max_mm_items_per_batch = len(dummy_mm_inputs)
+
+        logger.info_once(
+            "Encoder cache will be initialized with a budget of %s tokens, "
+            "and profiled with %s %s items of the maximum feature size.",
+            encoder_budget,
+            max_mm_items_per_batch,
+            dummy_modality,
+        )
+
+        dummy_encoder_outputs = self.execute_mm_encoder(dummy_mm_inputs)
+
+        sanity_check_mm_encoder_outputs(
+            dummy_encoder_outputs,
+            expected_num_items=max_mm_items_per_batch,
+        )
+        self.encoder_cache.encoder_outputs.update(
+            (f"tmp_{i}", output) for i, output in enumerate(dummy_encoder_outputs)
+        )
+
+    @torch.inference_mode()
+    def execute_mm_encoder(
+        self, mm_kwargs: list[tuple[str, MultiModalKwargsItem]]
+    ) -> list[torch.Tensor]:
+        # Stable-sort items by modality so each modality is encoded in as few
+        # batches as possible. The returned outputs are restored to input order.
+        order = sorted(range(len(mm_kwargs)), key=lambda i: mm_kwargs[i][0])
+        encoder_outputs: list[torch.Tensor] = []
+        for _, num_items, mm_kwargs_batch in group_and_batch_mm_kwargs(
+            [mm_kwargs[i] for i in order], device=self.device
+        ):
+            batch_outputs = self.model.embed_multimodal(**mm_kwargs_batch)
+            sanity_check_mm_encoder_outputs(batch_outputs, expected_num_items=num_items)
+            encoder_outputs.extend(batch_outputs)
+        return [encoder_outputs[i] for i in np.argsort(order)]
+
+    @contextmanager
+    def timed_encoder_operation(self, request_ids: Collection[str]):
+        if not (self.enable_timing and request_ids):
+            yield
+            return
+
+        start_time = time.perf_counter()
+        try:
+            yield
+        finally:
+            per_request_time = (time.perf_counter() - start_time) / len(request_ids)
+            with self._timing_lock:
+                for req_id in request_ids:
+                    stats = self.encoder_timing_registry.setdefault(
+                        req_id, EncoderTimingStats()
+                    )
+                    stats.encoder_forward_secs += per_request_time
+                    stats.num_encoder_calls += 1
+
+    def get_encoder_timing_stats(self) -> dict[str, dict[str, float | int]]:
+        with self._timing_lock:
+            stats = {
+                req_id: stats_obj.to_dict()
+                for req_id, stats_obj in self.encoder_timing_registry.items()
+            }
+            self.encoder_timing_registry.clear()
+            return stats
+
+    def gather_mm_embeddings(
+        self,
+        req_ids: list[str],
+        total_num_scheduled_tokens: int,
+        num_scheduled_tokens: np.ndarray,
+        query_start_loc: np.ndarray,
+        prefill_lens: np.ndarray,
+        num_computed_tokens: np.ndarray,
+    ) -> tuple[list[torch.Tensor], torch.Tensor]:
+        is_mm_embed = torch.zeros(
+            total_num_scheduled_tokens,
+            dtype=torch.bool,
+            device=self.device,
+        )
+
+        # Whether to gather media embeddings this step.
+        exclude_embeddings: list[bool] | None = None
+        if not self.is_realtime:
+            # Non-realtime models only have media embeddings within the prompt.
+            is_decode = num_computed_tokens >= prefill_lens
+            if is_decode.all():
+                # All decode requests, so no need to gather any embeddings.
+                return [], is_mm_embed
+            exclude_embeddings = is_decode.tolist()
+
+        query_start = num_computed_tokens.tolist()
+        query_end = (num_computed_tokens + num_scheduled_tokens).tolist()
+
+        mm_embeds: list[torch.Tensor] = []
+        for i, req_id in enumerate(req_ids):
+            if exclude_embeddings is not None and exclude_embeddings[i]:
+                continue
+
+            cur_query_start = query_start[i]
+            cur_query_end = query_end[i]
+
+            mm_features = self.encoder_cache.mm_features[req_id]
+            lo, hi = get_mm_features_in_window(
+                mm_features, start=cur_query_start, end=cur_query_end
+            )
+            for idx in range(lo, hi):
+                mm_feature = mm_features[idx]
+                pos_info = mm_feature.mm_position
+                start_pos = pos_info.offset
+                num_encoder_tokens = pos_info.length
+
+                start_idx = max(cur_query_start - start_pos, 0)
+                end_idx = min(cur_query_end - start_pos, num_encoder_tokens)
+                assert start_idx < end_idx
+                curr_embeds_start, curr_embeds_end = (
+                    pos_info.get_embeds_indices_in_range(start_idx, end_idx)
+                )
+                # If there are no embeddings in the current range, we skip
+                # gathering the embeddings.
+                if curr_embeds_start == curr_embeds_end:
+                    continue
+
+                mm_hash = mm_feature.identifier
+                encoder_output = self.encoder_cache.encoder_outputs.get(mm_hash, None)
+                if encoder_output is None:
+                    raise RuntimeError(f"Encoder cache miss for {mm_hash}.")
+
+                if (is_embed := pos_info.is_embed) is not None:
+                    is_embed = is_embed[start_idx:end_idx]
+                    mm_embeds_item = encoder_output[curr_embeds_start:curr_embeds_end]
+                else:
+                    mm_embeds_item = encoder_output[start_idx:end_idx]
+
+                # Attach modality for Omni interleaved merge (collected on demand).
+                set_mm_embedding_modality(mm_embeds_item, mm_feature.modality)
+
+                req_start_pos = query_start_loc[i] + start_pos - cur_query_start
+                is_mm_embed[req_start_pos + start_idx : req_start_pos + end_idx] |= (
+                    True if is_embed is None else is_embed
+                )
+                mm_embeds.append(mm_embeds_item)
+
+        return mm_embeds, is_mm_embed
+
+    @torch.inference_mode()
+    def get_inputs_embeds(
+        self,
+        input_ids: torch.Tensor,
+        mm_embeds: list[torch.Tensor],
+        is_mm_embed: torch.Tensor,
+    ) -> torch.Tensor:
+        x = self.model.embed_input_ids(
+            input_ids, multimodal_embeddings=mm_embeds, is_multimodal=is_mm_embed
+        )
+        self.inputs_embeds[: x.shape[0]] = x
+        return self.inputs_embeds
diff --git a/vllm/v1/worker/cpu/gumbel.py b/vllm/v1/worker/cpu/gumbel.py
new file mode 100644
index 000000000000..118c2ac233ed
--- /dev/null
+++ b/vllm/v1/worker/cpu/gumbel.py
@@ -0,0 +1,40 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+
+from vllm import _custom_ops as ops
+
+
+def gumbel_sample(
+    logits: torch.Tensor,
+    expanded_idx_mapping: torch.Tensor,
+    temperature: torch.Tensor,
+    seed: torch.Tensor,
+    pos: torch.Tensor,
+    apply_temperature: bool,
+    is_drafting: bool,
+    logits_cache: torch.Tensor | None = None,
+    logits_cache_col: torch.Tensor | None = None,
+    use_fp64: bool = False,
+    logits_cache_source: torch.Tensor | None = None,
+) -> torch.Tensor:
+    if (
+        logits_cache is not None
+        or logits_cache_col is not None
+        or logits_cache_source is not None
+    ):
+        raise NotImplementedError(
+            "CPU Gumbel sampling does not yet support the draft-logits cache."
+        )
+
+    return ops.gumbel_sample_cpu(
+        logits,
+        expanded_idx_mapping.contiguous(),
+        temperature,
+        seed,
+        pos.contiguous(),
+        apply_temperature,
+        is_drafting,
+        use_fp64,
+    )
diff --git a/vllm/v1/worker/cpu/input_batch.py b/vllm/v1/worker/cpu/input_batch.py
new file mode 100644
index 000000000000..9cd969346d1c
--- /dev/null
+++ b/vllm/v1/worker/cpu/input_batch.py
@@ -0,0 +1,343 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""CPU implementations for Model Runner V2 input-batch kernels."""
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, NamedTuple
+
+import numpy as np
+import torch
+
+from vllm import _custom_ops as ops
+from vllm.utils import random_uuid
+
+if TYPE_CHECKING:
+    from vllm.v1.worker.cpu.attn_utils import FastPrefillBatchMetadata
+
+
+class InputBuffers:
+    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
+
+        self.input_ids = torch.zeros(max_num_tokens, dtype=torch.int32, device=device)
+        self.positions = torch.zeros(max_num_tokens, dtype=torch.int64, device=device)
+        self.is_padding = torch.zeros(max_num_tokens, dtype=torch.bool, device=device)
+        self.query_start_loc = torch.zeros(
+            max_num_reqs + 1, dtype=torch.int32, device=device
+        )
+        self.seq_lens = torch.zeros(max_num_reqs, dtype=torch.int32, device=device)
+
+
+@dataclass
+class InputBatch:
+    # batch_idx -> req_id
+    req_ids: list[str]
+    num_reqs: int
+    num_reqs_after_padding: int
+
+    # batch_idx -> req_state_idx
+    idx_mapping: torch.Tensor
+    idx_mapping_np: np.ndarray
+    # Identical to idx_mapping except for spec decoding.
+    expanded_idx_mapping: torch.Tensor
+    # [total_num_logits] position within request for each logit
+    expanded_local_pos: torch.Tensor
+
+    # [num_reqs]
+    # batch_idx -> num_scheduled_tokens
+    num_scheduled_tokens: np.ndarray
+    # number of tokens in the batch
+    num_tokens: int
+    num_tokens_after_padding: int
+
+    # [num_reqs + 1]
+    query_start_loc: torch.Tensor
+    query_start_loc_np: np.ndarray
+    # [num_reqs]
+    seq_lens: torch.Tensor
+    # [num_reqs] CPU upper bound on seq_lens (see CommonAttentionMetadata).
+    seq_lens_cpu_upper_bound: torch.Tensor
+    # [num_reqs]
+    num_computed_tokens_np: np.ndarray
+    # [num_reqs]
+    prefill_len_np: np.ndarray
+    # [num_reqs]
+    num_computed_prefill_tokens_np: np.ndarray
+    # [num_reqs] CPU bool array == (num_computed_prefill_tokens_np < prefill_len_np).
+    is_prefilling_np: np.ndarray
+    # == np.any(is_prefilling_np)
+    has_prefill: bool
+    # No prefills, or only prefill rows that run as decodes.
+    decode_graph_eligible: bool
+
+    # [num_tokens_after_padding]
+    input_ids: torch.Tensor
+    # [num_tokens_after_padding]
+    positions: torch.Tensor
+    # [num_tokens_after_padding]
+    is_padding: torch.Tensor
+
+    # [total_num_logits]
+    logits_indices: torch.Tensor
+    # [num_reqs + 1]
+    cu_num_logits: torch.Tensor
+    cu_num_logits_np: np.ndarray
+
+    # Whether any requests in batch use structured output.
+    has_structured_output_reqs: bool
+
+    # [num_reqs] per-request prompt length, only populated for R-SWA.
+    prompt_lens: torch.Tensor | None
+
+    # [num_reqs] only populated for non-dummy pipeline-parallel batches.
+    max_seq_len_np: np.ndarray | None = None
+
+    # Longest query the batch may contain.
+    max_query_len: int | None = None
+
+    # Arms the KV-sharing fast prefill path for this step.
+    fast_prefill: "FastPrefillBatchMetadata | None" = None
+
+    # [num_reqs] prefilling rows that schedule one new prompt token over
+    # existing context and so compute exactly like decodes.
+    prefill_runs_as_decode_np: np.ndarray | None = None
+
+    @classmethod
+    def make_dummy(
+        cls,
+        num_reqs: int,
+        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)
+        expanded_idx_mapping = idx_mapping
+        expanded_local_pos = torch.zeros(num_reqs, dtype=torch.int32, device=device)
+
+        # Distribute the remainder evenly so that no dummy request exceeds
+        # ceil(num_tokens / num_reqs) <= max_model_len tokens.
+        base_tokens = num_tokens // num_reqs
+        num_extra = num_tokens % num_reqs
+        assert max_query_len is None or base_tokens + (num_extra > 0) <= max_query_len
+        num_scheduled_tokens = np.full(num_reqs, base_tokens, dtype=np.int32)
+        if num_extra > 0:
+            num_scheduled_tokens[-num_extra:] += 1
+        assert int(num_scheduled_tokens.sum()) == num_tokens
+
+        # seq_len equals to query_len
+        input_buffers.seq_lens[: num_reqs - num_extra] = base_tokens
+        input_buffers.seq_lens[num_reqs - num_extra : num_reqs] = base_tokens + 1
+        # Pad unused rows.
+        input_buffers.seq_lens[num_reqs:] = 0
+        seq_lens = input_buffers.seq_lens[:num_reqs]
+
+        query_start_loc_np = np.empty(num_reqs + 1, dtype=np.int32)
+        query_start_loc_np[0] = 0
+        np.cumsum(num_scheduled_tokens, out=query_start_loc_np[1:])
+        input_buffers.query_start_loc[:1] = 0
+        torch.cumsum(
+            seq_lens, dim=0, out=input_buffers.query_start_loc[1 : num_reqs + 1]
+        )
+        # Pad unused rows.
+        input_buffers.query_start_loc[num_reqs + 1 :] = num_tokens
+        query_start_loc = input_buffers.query_start_loc[: num_reqs + 1]
+
+        input_ids = input_buffers.input_ids[:num_tokens].zero_()
+        positions = input_buffers.positions[:num_tokens].zero_()
+
+        input_buffers.is_padding[:num_tokens].fill_(is_padding)
+        is_padding = input_buffers.is_padding[:num_tokens]
+
+        logits_indices = query_start_loc[1:] - 1
+        cu_num_logits = torch.arange(num_reqs + 1, device=device, dtype=torch.int32)
+        cu_num_logits_np = np.arange(num_reqs + 1, dtype=np.int32)
+        # Copy so dummy context can be added without touching
+        # num_scheduled_tokens.
+        seq_lens_cpu_upper_bound = torch.from_numpy(num_scheduled_tokens.copy())
+        return cls(
+            req_ids=req_ids,
+            num_reqs=num_reqs,
+            num_reqs_after_padding=num_reqs,
+            idx_mapping=idx_mapping,
+            idx_mapping_np=idx_mapping_np,
+            expanded_idx_mapping=expanded_idx_mapping,
+            expanded_local_pos=expanded_local_pos,
+            num_scheduled_tokens=num_scheduled_tokens,
+            num_tokens=num_tokens,
+            num_tokens_after_padding=num_tokens,
+            query_start_loc=query_start_loc,
+            query_start_loc_np=query_start_loc_np,
+            seq_lens=seq_lens,
+            seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,
+            num_computed_tokens_np=np.zeros(num_reqs, dtype=np.int32),
+            prefill_len_np=np.zeros(num_reqs, dtype=np.int32),
+            num_computed_prefill_tokens_np=np.zeros(num_reqs, dtype=np.int32),
+            is_prefilling_np=np.zeros(num_reqs, dtype=np.bool_),
+            has_prefill=False,
+            max_seq_len_np=None,
+            decode_graph_eligible=True,
+            input_ids=input_ids,
+            positions=positions,
+            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=False,
+            prompt_lens=None,
+            max_query_len=max_query_len,
+        )
+
+
+class BatchReqState(NamedTuple):
+    """CPU request state for a scheduled batch, in batch (sorted) order."""
+
+    req_ids: list[str]
+    num_scheduled_tokens: np.ndarray  # [num_reqs]
+    num_tokens: int
+    idx_mapping_np: np.ndarray  # [num_reqs]
+    prefill_len_np: np.ndarray  # [num_reqs]
+    num_computed_prefill_tokens_np: np.ndarray  # [num_reqs]
+    is_prefilling_np: np.ndarray  # [num_reqs]
+    max_seq_len_np: np.ndarray | None  # [num_reqs], None in non-pp case
+    has_prefill: bool
+    prefill_runs_as_decode_np: np.ndarray | None  # [num_reqs]
+    decode_graph_eligible: bool
+
+
+def prepare_prefill_inputs(
+    input_ids: torch.Tensor,
+    next_prefill_tokens: torch.Tensor,
+    idx_mapping: torch.Tensor,
+    query_start_loc: torch.Tensor,
+    all_token_ids: torch.Tensor,
+    prefill_len: torch.Tensor,
+    num_computed_tokens: torch.Tensor,
+) -> None:
+    ops.prepare_prefill_inputs_cpu(
+        input_ids,
+        next_prefill_tokens,
+        idx_mapping,
+        query_start_loc,
+        all_token_ids,
+        prefill_len,
+        num_computed_tokens,
+    )
+
+
+def prepare_pos_seq_lens(
+    idx_mapping: torch.Tensor,
+    query_start_loc: torch.Tensor,
+    num_computed_tokens: torch.Tensor,
+    pos: torch.Tensor,
+    seq_lens: torch.Tensor,
+) -> None:
+    ops.prepare_pos_seq_lens_cpu(
+        idx_mapping,
+        query_start_loc,
+        num_computed_tokens,
+        pos,
+        seq_lens,
+    )
+
+
+def combine_sampled_and_draft_tokens(
+    input_ids: torch.Tensor,
+    idx_mapping: torch.Tensor,
+    last_sampled_tokens: torch.Tensor,
+    query_start_loc: torch.Tensor,
+    seq_lens: torch.Tensor,
+    prefill_len: torch.Tensor,
+    draft_tokens: torch.Tensor,
+    cu_num_logits: torch.Tensor,
+    num_logits: int,
+    num_new_sampled_tokens: int = 1,  # excl accepted draft tokens, a.k.a bonus tokens
+) -> torch.Tensor:
+    assert num_new_sampled_tokens in (0, 1), (
+        f"num_new_sampled_tokens must be 0 or 1, got {num_new_sampled_tokens}"
+    )
+    return ops.combine_sampled_and_draft_tokens_cpu(
+        input_ids,
+        idx_mapping,
+        last_sampled_tokens,
+        query_start_loc,
+        seq_lens,
+        prefill_len,
+        draft_tokens,
+        cu_num_logits,
+        num_logits,
+        num_new_sampled_tokens,
+    )
+
+
+def get_num_sampled_and_rejected(
+    num_sampled: torch.Tensor,
+    seq_lens: torch.Tensor,
+    cu_num_logits: torch.Tensor,
+    idx_mapping: torch.Tensor,
+    prefill_len: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    num_reqs = idx_mapping.shape[0]
+    num_rejected = torch.empty_like(num_sampled)
+    is_chunked_prefilling = seq_lens[:num_reqs] < prefill_len[idx_mapping]
+    num_sampled.masked_fill_(is_chunked_prefilling, 0)
+    num_rejected.copy_(cu_num_logits[1 : num_reqs + 1] - cu_num_logits[:num_reqs])
+    num_rejected.sub_(num_sampled)
+    num_rejected.masked_fill_(is_chunked_prefilling, 0)
+    return num_sampled, num_rejected
+
+
+def post_update(
+    # [num_reqs] batch_idx -> req_state_idx; negative index means skip.
+    idx_mapping: torch.Tensor,
+    # [max_num_reqs]
+    num_computed_tokens: torch.Tensor,
+    # [max_num_reqs]
+    last_sampled_tokens: torch.Tensor,
+    # [max_num_reqs, vocab_size]
+    output_bin_counts: torch.Tensor | None,
+    # [num_reqs, num_speculative_steps + 1]
+    sampled_tokens: torch.Tensor,
+    # [num_reqs]
+    num_sampled: torch.Tensor,
+    # [num_reqs]
+    num_rejected: torch.Tensor,
+    # [num_reqs + 1]
+    query_start_loc: torch.Tensor | None,
+    # [max_num_reqs, max_model_len]
+    all_token_ids: torch.Tensor,
+    # [max_num_reqs]
+    total_len: torch.Tensor,
+) -> None:
+    ops.post_update_cpu(
+        idx_mapping,
+        num_computed_tokens,
+        last_sampled_tokens,
+        output_bin_counts,
+        sampled_tokens,
+        num_sampled,
+        num_rejected,
+        query_start_loc,
+        all_token_ids,
+        total_len,
+    )
+
+
+def post_update_num_computed_tokens(
+    # [num_reqs]
+    idx_mapping: torch.Tensor,
+    # [max_num_reqs]
+    num_computed_tokens: torch.Tensor,
+    # [num_reqs + 1]
+    query_start_loc: torch.Tensor,
+) -> None:
+    query_lens = query_start_loc[1:] - query_start_loc[:-1]
+    num_computed_tokens.index_add_(0, idx_mapping, query_lens)
diff --git a/vllm/v1/worker/cpu/kv_connector.py b/vllm/v1/worker/cpu/kv_connector.py
new file mode 100644
index 000000000000..9ea7862343d3
--- /dev/null
+++ b/vllm/v1/worker/cpu/kv_connector.py
@@ -0,0 +1,173 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from typing import TYPE_CHECKING, Any
+
+import torch
+
+from vllm.config import VllmConfig
+from vllm.distributed.kv_transfer import (
+    get_kv_transfer_group,
+    has_kv_transfer_group,
+    kv_transfer_state,
+)
+from vllm.distributed.kv_transfer.kv_connector.utils import copy_kv_blocks
+from vllm.forward_context import (
+    get_forward_context,
+    is_forward_context_available,
+    set_forward_context,
+)
+from vllm.v1.outputs import (
+    EMPTY_MODEL_RUNNER_OUTPUT,
+    KVConnectorOutput,
+    ModelRunnerOutput,
+)
+
+if TYPE_CHECKING:
+    from vllm.v1.core.sched.output import SchedulerOutput
+    from vllm.v1.worker.cpu.input_batch import InputBatch
+
+
+class KVConnector:
+    """KV-transfer interface used by the CPU model runner."""
+
+    def pre_forward(
+        self,
+        scheduler_output: "SchedulerOutput",
+        input_batch: "InputBatch | None" = None,
+        attn_metadata: dict[str, Any] | None = None,
+    ) -> None:
+        pass
+
+    def finish_forward(self) -> None:
+        pass
+
+    def post_forward(self, finished_req_ids: set[str]) -> KVConnectorOutput | None:
+        return None
+
+    def no_forward(self, scheduler_output: "SchedulerOutput") -> ModelRunnerOutput:
+        return EMPTY_MODEL_RUNNER_OUTPUT
+
+    def set_disabled(self, disabled: bool) -> None:
+        pass
+
+
+class ActiveKVConnector(KVConnector):
+    def __init__(
+        self, vllm_config: VllmConfig, kv_caches_dict: dict[str, torch.Tensor]
+    ):
+        self.vllm_config = vllm_config
+        self.kv_connector = get_kv_transfer_group()
+        # Register kv caches with KV Connector if applicable.
+        self.kv_connector.register_kv_caches(kv_caches_dict)
+        self.kv_connector.set_host_xfer_buffer_ops(copy_kv_blocks)
+
+        self._pending_load_kwargs: dict[str, Any] | None = None
+        self._disabled = False
+
+    def pre_forward(
+        self,
+        scheduler_output: "SchedulerOutput",
+        input_batch: "InputBatch | None" = None,
+        attn_metadata: dict[str, Any] | None = None,
+    ) -> None:
+        if self._disabled:
+            return
+
+        kv_connector_metadata = scheduler_output.kv_connector_metadata
+        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 = self._build_load_kwargs(
+            scheduler_output, input_batch, attn_metadata
+        )
+
+        if scheduler_output.has_sync_kv_loads:
+            # Sync loads need to run before this step's forward.
+            self._start_load_kv()
+        # Otherwise start the async load after forward to keep submission
+        # off the critical path.
+
+    @staticmethod
+    def _build_load_kwargs(
+        scheduler_output: "SchedulerOutput",
+        input_batch: "InputBatch | None",
+        attn_metadata: dict[str, Any] | None,
+    ) -> dict[str, Any]:
+        """Keyword args for the connector's start_load_kv()."""
+        if input_batch is None:
+            # No forward this step.
+            return {}
+        load_kwargs: dict[str, Any] = dict(
+            scheduler_output=scheduler_output,
+            request_state_indices=input_batch.idx_mapping,
+            request_ids=input_batch.req_ids,
+            num_tokens=input_batch.num_tokens,
+        )
+        if attn_metadata is not None:
+            load_kwargs["attn_metadata"] = attn_metadata
+        return load_kwargs
+
+    def _start_load_kv(self) -> None:
+        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(), **load_kwargs)
+        else:
+            with set_forward_context(None, self.vllm_config):
+                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 post_forward(self, finished_req_ids: set[str]) -> KVConnectorOutput | None:
+        if self._disabled:
+            return None
+
+        if self._pending_load_kwargs is not None:
+            self._start_load_kv()
+
+        output = KVConnectorOutput()
+        self.kv_connector.wait_for_save()
+        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()
+        output.kv_connector_worker_meta = (
+            self.kv_connector.build_connector_worker_meta()
+        )
+        self.kv_connector.clear_connector_metadata()
+        return output
+
+    def no_forward(self, scheduler_output: "SchedulerOutput") -> ModelRunnerOutput:
+        if self._disabled:
+            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)
+        return ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output)
+
+    def set_disabled(self, disabled: bool) -> None:
+        # Ensure that layer-wise connector hooks aren't called when disabled.
+        kv_transfer_state._KV_CONNECTOR_AGENT = None if disabled else self.kv_connector
+        self._disabled = disabled
+
+
+NO_OP_KV_CONNECTOR = KVConnector()
+
+
+def get_kv_connector(
+    vllm_config: VllmConfig, kv_caches_dict: dict[str, torch.Tensor]
+) -> KVConnector:
+    if not has_kv_transfer_group():
+        # No-op connector.
+        return NO_OP_KV_CONNECTOR
+
+    return ActiveKVConnector(vllm_config, kv_caches_dict)
diff --git a/vllm/v1/worker/cpu/logits_processors.py b/vllm/v1/worker/cpu/logits_processors.py
new file mode 100644
index 000000000000..dbac63152a7e
--- /dev/null
+++ b/vllm/v1/worker/cpu/logits_processors.py
@@ -0,0 +1,355 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+import numpy as np
+import torch
+
+import vllm.envs as envs
+from vllm import _custom_ops as ops
+from vllm.sampling_params import SamplingParams
+
+if TYPE_CHECKING:
+    from vllm.config import VllmConfig
+    from vllm.v1.worker.cpu.states import CPURequestState
+
+MAX_NUM_ALLOWED_TOKEN_IDS = 1024
+MAX_NUM_LOGIT_BIAS_TOKENS = 1024
+MAX_NUM_STOP_TOKEN_IDS = 128
+
+
+@dataclass(frozen=True)
+class LogitsProcRequestState:
+    device: torch.device
+    max_num_reqs: int
+    vocab_size: int
+    all_token_ids: torch.Tensor
+    prompt_len: torch.Tensor
+    prompt_len_np: np.ndarray
+    prefill_len: torch.Tensor
+    prefill_len_np: np.ndarray
+    total_len: torch.Tensor
+
+    @classmethod
+    def from_request_state(cls, req_states: CPURequestState) -> 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,
+            prompt_len_np=req_states.prompt_len_np,
+            prefill_len=req_states.prefill_len,
+            prefill_len_np=req_states.prefill_len_np,
+            total_len=req_states.total_len,
+        )
+
+
+@dataclass(frozen=True)
+class LogitsContext:
+    expanded_idx_mapping: torch.Tensor
+    idx_mapping: torch.Tensor
+    idx_mapping_np: np.ndarray
+    expanded_local_pos: torch.Tensor
+    input_ids: torch.Tensor
+    pos: torch.Tensor
+    seq_lens_upper_bound_np: np.ndarray
+
+
+class LogitsProcessor(ABC):
+    def __init__(  # noqa: B027
+        self, vllm_config: VllmConfig, req_states: LogitsProcRequestState
+    ) -> None:
+        pass
+
+    @classmethod  # noqa: B027
+    def validate_params(cls, sampling_params: SamplingParams) -> None:
+        pass
+
+    def add_request(self, req_idx: int, sampling_params: SamplingParams) -> bool:
+        return True
+
+    @abstractmethod
+    def apply(self, logits: torch.Tensor, ctx: LogitsContext) -> torch.Tensor:
+        raise NotImplementedError
+
+
+class CPULogitBiasState(LogitsProcessor):
+    def __init__(self, vllm_config: VllmConfig, req_states: LogitsProcRequestState):
+        max_num_reqs = req_states.max_num_reqs
+        device = req_states.device
+        self.prompt_len_np = req_states.prompt_len_np
+
+        self.num_allowed_token_ids = torch.zeros(
+            max_num_reqs, dtype=torch.int32, device=device
+        )
+        self.num_allowed_token_ids_np = self.num_allowed_token_ids.numpy()
+        self.allowed_token_ids = torch.zeros(
+            (max_num_reqs, MAX_NUM_ALLOWED_TOKEN_IDS),
+            dtype=torch.int32,
+            device=device,
+        )
+        self.num_logit_bias = torch.zeros(
+            max_num_reqs, dtype=torch.int32, device=device
+        )
+        self.num_logit_bias_np = self.num_logit_bias.numpy()
+        self.logit_bias_token_ids = torch.zeros(
+            (max_num_reqs, MAX_NUM_LOGIT_BIAS_TOKENS),
+            dtype=torch.int32,
+            device=device,
+        )
+        self.logit_bias = torch.zeros(
+            (max_num_reqs, MAX_NUM_LOGIT_BIAS_TOKENS),
+            dtype=torch.float32,
+            device=device,
+        )
+        self.min_lens = torch.zeros(max_num_reqs, dtype=torch.int32, device=device)
+        self.min_lens_np = self.min_lens.numpy()
+        self.num_stop_token_ids = torch.zeros(
+            max_num_reqs, dtype=torch.int32, device=device
+        )
+        self.num_stop_token_ids_np = self.num_stop_token_ids.numpy()
+        self.restore_when_all_masked = torch.zeros(
+            max_num_reqs, dtype=torch.int32, device=device
+        )
+        self.restore_when_all_masked_np = self.restore_when_all_masked.numpy()
+        self.stop_token_ids = torch.zeros(
+            (max_num_reqs, MAX_NUM_STOP_TOKEN_IDS),
+            dtype=torch.int32,
+            device=device,
+        )
+        self.use_logit_bias = np.zeros(max_num_reqs, dtype=bool)
+
+    def add_request(self, req_idx: int, sampling_params: SamplingParams) -> bool:
+        prompt_len = int(self.prompt_len_np[req_idx])
+        use_logit_bias = False
+
+        allowed_token_ids = sampling_params.allowed_token_ids
+        if allowed_token_ids:
+            num_allowed_token_ids = len(allowed_token_ids)
+            if num_allowed_token_ids > MAX_NUM_ALLOWED_TOKEN_IDS:
+                raise ValueError(
+                    f"Too many allowed token IDs: {num_allowed_token_ids}. "
+                    f"The max size is {MAX_NUM_ALLOWED_TOKEN_IDS}."
+                )
+            self.num_allowed_token_ids_np[req_idx] = num_allowed_token_ids
+            self.allowed_token_ids[req_idx, :num_allowed_token_ids] = torch.tensor(
+                allowed_token_ids,
+                dtype=torch.int32,
+                device=self.allowed_token_ids.device,
+            )
+            use_logit_bias = True
+        else:
+            self.num_allowed_token_ids_np[req_idx] = 0
+
+        logit_bias = sampling_params.logit_bias
+        if logit_bias:
+            num_logit_bias = len(logit_bias)
+            if num_logit_bias > MAX_NUM_LOGIT_BIAS_TOKENS:
+                raise ValueError(
+                    f"Too many logit bias tokens: {num_logit_bias}. "
+                    f"The max size is {MAX_NUM_LOGIT_BIAS_TOKENS}."
+                )
+            self.num_logit_bias_np[req_idx] = num_logit_bias
+            self.logit_bias_token_ids[req_idx, :num_logit_bias] = torch.tensor(
+                list(logit_bias),
+                dtype=torch.int32,
+                device=self.logit_bias_token_ids.device,
+            )
+            self.logit_bias[req_idx, :num_logit_bias] = torch.tensor(
+                list(logit_bias.values()),
+                dtype=torch.float32,
+                device=self.logit_bias.device,
+            )
+            use_logit_bias = True
+        else:
+            self.num_logit_bias_np[req_idx] = 0
+
+        min_tokens = sampling_params.min_tokens
+        min_len = prompt_len + min_tokens
+        self.min_lens_np[req_idx] = min_len
+        stop_token_ids = sampling_params.all_stop_token_ids
+        if min_tokens > 0 and stop_token_ids:
+            num_stop_token_ids = len(stop_token_ids)
+            if num_stop_token_ids > MAX_NUM_STOP_TOKEN_IDS:
+                raise ValueError(
+                    f"Too many stop tokens: {num_stop_token_ids}. "
+                    f"The max size is {MAX_NUM_STOP_TOKEN_IDS}."
+                )
+            self.num_stop_token_ids_np[req_idx] = num_stop_token_ids
+            self.stop_token_ids[req_idx, :num_stop_token_ids] = torch.tensor(
+                list(stop_token_ids),
+                dtype=torch.int32,
+                device=self.stop_token_ids.device,
+            )
+            self.restore_when_all_masked_np[req_idx] = int(
+                sampling_params.structured_outputs is not None
+            )
+            use_logit_bias = True
+        else:
+            self.num_stop_token_ids_np[req_idx] = 0
+            self.restore_when_all_masked_np[req_idx] = 0
+
+        self.use_logit_bias[req_idx] = use_logit_bias
+        return use_logit_bias
+
+    def apply(self, logits: torch.Tensor, ctx: LogitsContext) -> torch.Tensor:
+        if not np.any(self.use_logit_bias[ctx.idx_mapping_np]):
+            return logits
+
+        enable_stop_token_restore = bool(
+            np.any(self.restore_when_all_masked_np[ctx.idx_mapping_np])
+        )
+        ops.apply_logit_bias_cpu(
+            logits,
+            ctx.expanded_idx_mapping,
+            ctx.pos,
+            self.num_allowed_token_ids,
+            self.allowed_token_ids,
+            self.num_logit_bias,
+            self.logit_bias_token_ids,
+            self.logit_bias,
+            self.min_lens,
+            self.num_stop_token_ids,
+            self.restore_when_all_masked,
+            self.stop_token_ids,
+            enable_stop_token_restore,
+        )
+        return logits
+
+
+class _CPULogitsProcessor(LogitsProcessor):
+    def __init__(self, vllm_config: VllmConfig, req_states: LogitsProcRequestState):
+        self.req_states = req_states
+
+    def _request_rows(self, ctx: LogitsContext):
+        for row, req_idx in enumerate(ctx.expanded_idx_mapping.tolist()):
+            if req_idx >= 0:
+                yield row, req_idx
+
+    def _output_tokens(
+        self, req_idx: int, row: int, ctx: LogitsContext
+    ) -> torch.Tensor:
+        prompt_len = int(self.req_states.prompt_len_np[req_idx])
+        total_len = int(self.req_states.total_len[req_idx])
+        first_row = row - int(ctx.expanded_local_pos[row])
+        return torch.cat(
+            (
+                self.req_states.all_token_ids[req_idx, prompt_len:total_len],
+                ctx.input_ids[first_row + 1 : row + 1].to(torch.int32),
+            )
+        )
+
+
+class CPUPenalties(_CPULogitsProcessor):
+    output_bin_counts = None
+
+    def __init__(self, vllm_config: VllmConfig, req_states: LogitsProcRequestState):
+        super().__init__(vllm_config, req_states)
+        self.repetition_penalty = torch.ones(
+            req_states.max_num_reqs, dtype=torch.float32, device=req_states.device
+        )
+        self.frequency_penalty = torch.zeros_like(self.repetition_penalty)
+        self.presence_penalty = torch.zeros_like(self.repetition_penalty)
+
+    def add_request(self, req_idx: int, sampling_params: SamplingParams) -> bool:
+        repetition = sampling_params.repetition_penalty
+        frequency = sampling_params.frequency_penalty
+        presence = sampling_params.presence_penalty
+        self.repetition_penalty[req_idx] = repetition
+        self.frequency_penalty[req_idx] = frequency
+        self.presence_penalty[req_idx] = presence
+        return repetition != 1.0 or frequency != 0.0 or presence != 0.0
+
+    def apply(self, logits: torch.Tensor, ctx: LogitsContext) -> torch.Tensor:
+        req_indices = ctx.expanded_idx_mapping.long()
+        repetition_by_row = self.repetition_penalty[req_indices]
+        frequency_by_row = self.frequency_penalty[req_indices]
+        presence_by_row = self.presence_penalty[req_indices]
+        active_rows = (
+            (repetition_by_row != 1.0)
+            | (frequency_by_row != 0.0)
+            | (presence_by_row != 0.0)
+        ).nonzero(as_tuple=True)[0]
+        if active_rows.numel() == 0:
+            return logits
+
+        vocab_size = logits.shape[1]
+        prompt_ids = []
+        output_ids = []
+        req_indices_list = req_indices.tolist()
+        for row in active_rows.tolist():
+            req_idx = req_indices_list[row]
+            offset = row * vocab_size
+            prompt_len = int(self.req_states.prompt_len_np[req_idx])
+            prompt = self.req_states.all_token_ids[req_idx, :prompt_len]
+            prompt = prompt[(prompt >= 0) & (prompt < vocab_size)]
+            prompt_ids.append(prompt + offset)
+            output_ids.append(self._output_tokens(req_idx, row, ctx) + offset)
+
+        flat_logits = logits.view(-1)
+        output_ids_flat, counts = torch.unique(
+            torch.cat(output_ids).long(), return_counts=True
+        )
+        seen_ids = torch.unique(torch.cat((*prompt_ids, output_ids_flat)).long())
+
+        repetition = repetition_by_row[seen_ids // vocab_size]
+        seen_logits = flat_logits[seen_ids]
+        flat_logits[seen_ids] = torch.where(
+            seen_logits > 0,
+            seen_logits / repetition,
+            seen_logits * repetition,
+        )
+        output_rows = output_ids_flat // vocab_size
+        frequency = frequency_by_row[output_rows]
+        presence = presence_by_row[output_rows]
+        flat_logits[output_ids_flat] -= frequency * counts + presence
+        return logits
+
+
+class CPUBadWords(_CPULogitsProcessor):
+    def __init__(self, vllm_config: VllmConfig, req_states: LogitsProcRequestState):
+        super().__init__(vllm_config, req_states)
+        self.bad_words: list[tuple[tuple[int, ...], ...] | None] = [
+            None
+        ] * req_states.max_num_reqs
+
+    def add_request(self, req_idx: int, sampling_params: SamplingParams) -> bool:
+        bad_words = sampling_params.bad_words_token_ids
+        if not bad_words:
+            self.bad_words[req_idx] = None
+            return False
+
+        if len(bad_words) > envs.VLLM_MAX_NUM_BAD_WORDS:
+            raise ValueError(
+                f"Too many bad words: {len(bad_words)}. "
+                f"The max number is {envs.VLLM_MAX_NUM_BAD_WORDS}."
+            )
+        num_tokens = sum(map(len, bad_words))
+        if num_tokens > envs.VLLM_MAX_BAD_WORDS_TOTAL_TOKENS:
+            raise ValueError(
+                f"Too many total bad word tokens: {num_tokens}. "
+                f"The max is {envs.VLLM_MAX_BAD_WORDS_TOTAL_TOKENS}."
+            )
+
+        self.bad_words[req_idx] = tuple(tuple(word) for word in bad_words)
+        return True
+
+    def apply(self, logits: torch.Tensor, ctx: LogitsContext) -> torch.Tensor:
+        for row, req_idx in self._request_rows(ctx):
+            bad_words = self.bad_words[req_idx]
+            if bad_words is None:
+                continue
+            history = self._output_tokens(req_idx, row, ctx).tolist()
+            for bad_word in bad_words:
+                prefix = bad_word[:-1]
+                if len(prefix) <= len(history) and (
+                    not prefix or history[-len(prefix) :] == list(prefix)
+                ):
+                    logits[row, bad_word[-1]] = -torch.inf
+        return logits
diff --git a/vllm/v1/worker/cpu/logprob.py b/vllm/v1/worker/cpu/logprob.py
new file mode 100644
index 000000000000..ca87710fc896
--- /dev/null
+++ b/vllm/v1/worker/cpu/logprob.py
@@ -0,0 +1,416 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from collections.abc import Callable
+from dataclasses import dataclass
+
+import numpy as np
+import torch
+
+from vllm import _custom_ops as ops
+from vllm.config.model import LogprobsMode
+from vllm.platforms import current_platform
+from vllm.sampling_params import MAX_LOGPROB_TOKEN_IDS, SamplingParams
+from vllm.v1.outputs import LogprobsTensors
+from vllm.v1.worker.cpu.input_batch import InputBatch
+
+CHUNK_SIZE = 1024
+
+
+@dataclass
+class _TokenIdScores:
+    token_ids: torch.Tensor  # [num_rows, num_ids], on device
+    pad: torch.Tensor  # [num_rows, num_ids] bool, True at padded columns
+    scores: torch.Tensor | None = None  # [num_rows, num_ids], filled per chunk
+
+
+class LogprobTokenIdsState:
+    def __init__(self, max_num_reqs: int, device: torch.device):
+        self.max_num_reqs = max_num_reqs
+        self.num_token_ids = torch.zeros(max_num_reqs, dtype=torch.int32, device=device)
+        self.num_token_ids_np = self.num_token_ids.numpy()
+        self.token_ids = torch.zeros(
+            (max_num_reqs, MAX_LOGPROB_TOKEN_IDS),
+            dtype=torch.int32,
+            device=device,
+        )
+
+    def add_request(self, req_idx: int, sampling_params: SamplingParams) -> None:
+        token_ids = sampling_params.logprob_token_ids
+        if not token_ids:
+            self.num_token_ids_np[req_idx] = 0
+            return
+        n = len(token_ids)
+        if n > MAX_LOGPROB_TOKEN_IDS:
+            raise ValueError(
+                f"Too many logprob_token_ids: {n}. The max is {MAX_LOGPROB_TOKEN_IDS}."
+            )
+        self.num_token_ids_np[req_idx] = n
+        self.token_ids[req_idx, :n] = torch.tensor(
+            token_ids, dtype=torch.int32, device=self.token_ids.device
+        )
+
+    def max_num_token_ids(self, idx_mapping_np: np.ndarray) -> int:
+        return int(self.num_token_ids_np[idx_mapping_np].max(initial=0))
+
+
+@torch.compile(
+    dynamic=True,
+    fullgraph=True,
+    backend=current_platform.simple_compile_backend,
+)
+def compute_token_logprobs(
+    logits: torch.Tensor, token_ids: torch.Tensor
+) -> torch.Tensor:
+    token_ids = token_ids.to(torch.int64)
+    selected = logits.gather(-1, token_ids).to(torch.float32)
+    # Accumulate the normalizer in FP32. Compilation fuses this cast into the
+    # reduction instead of materializing a full FP32 logits tensor.
+    normalizer = torch.logsumexp(logits.to(torch.float32), dim=-1, keepdim=True)
+    return selected - normalizer
+
+
+@torch.compile(
+    dynamic=True,
+    fullgraph=True,
+    backend=current_platform.simple_compile_backend,
+)
+def _compute_token_ranks(logits: torch.Tensor, token_ids: torch.Tensor) -> torch.Tensor:
+    selected = logits.gather(-1, token_ids.unsqueeze(-1))
+    return (logits >= selected).sum(-1)
+
+
+def compute_topk_scores(
+    logits: torch.Tensor,
+    num_logprobs: int,
+    sampled_token_ids: torch.Tensor,
+    cu_num_logits: list[int] | torch.Tensor | None = None,
+    num_per_req_token_ids: torch.Tensor | None = None,
+    per_req_token_ids: torch.Tensor | None = None,
+    expanded_idx_mapping: torch.Tensor | None = None,
+    max_per_req_token_ids: int = 0,
+    logits_mode: bool = False,
+) -> LogprobsTensors:
+    assert num_logprobs >= 0
+    sampled_token_ids = sampled_token_ids.to(torch.int64)
+
+    if max_per_req_token_ids == 0:
+        logprob_token_ids = sampled_token_ids.unsqueeze(-1)
+        if num_logprobs > 0:
+            topk_token_ids = torch.topk(logits, num_logprobs, dim=-1).indices
+            logprob_token_ids = torch.cat((logprob_token_ids, topk_token_ids), dim=1)
+        valid_mask = None
+    else:
+        assert num_per_req_token_ids is not None
+        assert per_req_token_ids is not None
+        assert expanded_idx_mapping is not None
+        batch_size = logits.shape[0]
+        num_cols = max(num_logprobs, max_per_req_token_ids)
+        selected_token_ids = sampled_token_ids.new_zeros((batch_size, num_cols))
+        valid_mask = torch.zeros_like(selected_token_ids, dtype=torch.bool)
+
+        if num_logprobs > 0:
+            topk_token_ids = torch.topk(logits, num_logprobs, dim=-1).indices
+            selected_token_ids[:, :num_logprobs] = topk_token_ids
+        valid_mask[:, :num_logprobs] = True
+
+        req_indices = expanded_idx_mapping.to(torch.int64)
+        num_custom = num_per_req_token_ids[req_indices].to(torch.int64)
+        custom_token_ids = torch.zeros_like(selected_token_ids)
+        custom_token_ids[:, :max_per_req_token_ids] = per_req_token_ids[
+            req_indices, :max_per_req_token_ids
+        ].to(torch.int64)
+        custom_valid_mask = torch.arange(num_cols, device=logits.device).unsqueeze(
+            0
+        ) < num_custom.unsqueeze(1)
+        use_custom = num_custom.unsqueeze(1) > 0
+        selected_token_ids = torch.where(
+            use_custom, custom_token_ids, selected_token_ids
+        )
+        valid_mask = torch.where(use_custom, custom_valid_mask, valid_mask)
+
+        logprob_token_ids = torch.cat(
+            (sampled_token_ids.unsqueeze(-1), selected_token_ids), dim=1
+        )
+        valid_mask = torch.cat(
+            (
+                torch.ones((batch_size, 1), dtype=torch.bool, device=logits.device),
+                valid_mask,
+            ),
+            dim=1,
+        )
+
+    if logits_mode:
+        logprobs = logits.gather(-1, logprob_token_ids).to(torch.float32)
+    else:
+        logprobs = compute_token_logprobs(logits, logprob_token_ids)
+    if valid_mask is not None:
+        logprobs.masked_fill_(~valid_mask, float("-inf"))
+
+    sampled_ranks = _compute_token_ranks(logits, sampled_token_ids)
+
+    is_tensor = isinstance(cu_num_logits, torch.Tensor)
+    return LogprobsTensors(
+        logprob_token_ids=logprob_token_ids,
+        logprobs=logprobs,
+        selected_token_ranks=sampled_ranks,
+        cu_num_generated_tokens=None if is_tensor else cu_num_logits,
+        cu_num_generated_tokens_tensor=cu_num_logits if is_tensor else None,
+    )
+
+
+def get_prompt_logprobs_token_ids(
+    num_tokens: int,
+    query_start_loc: torch.Tensor,
+    idx_mapping: torch.Tensor,
+    num_computed_tokens: torch.Tensor,
+    all_token_ids: torch.Tensor,
+) -> torch.Tensor:
+    return ops.get_prompt_logprobs_token_ids_cpu(
+        num_tokens,
+        query_start_loc,
+        idx_mapping,
+        num_computed_tokens,
+        all_token_ids,
+    )
+
+
+def compute_prompt_logprobs_with_chunking(
+    prompt_token_ids: torch.Tensor,
+    prompt_hidden_states: torch.Tensor,
+    logits_fn: Callable[[torch.Tensor], torch.Tensor],
+    num_prompt_logprobs: int,
+    logprobs_mode: LogprobsMode = "raw_logprobs",
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    # Since materializing the full prompt logits can take too much memory,
+    # we compute it in chunks.
+    token_ids = []
+    scores = []
+    ranks = []
+    logits_mode = logprobs_mode in ("raw_logits", "processed_logits")
+    prompt_token_ids = prompt_token_ids.to(torch.int64)
+    for start_idx in range(0, prompt_token_ids.shape[0], CHUNK_SIZE):
+        end_idx = start_idx + CHUNK_SIZE
+        # NOTE(woosuk): logits_fn can be slow because it involves all-gather.
+        prompt_logits = logits_fn(prompt_hidden_states[start_idx:end_idx])
+        requested_num = (
+            prompt_logits.shape[-1]
+            if num_prompt_logprobs == -1
+            else num_prompt_logprobs
+        )
+        result = compute_topk_scores(
+            prompt_logits,
+            requested_num,
+            prompt_token_ids[start_idx:end_idx],
+            logits_mode=logits_mode,
+        )
+        token_ids.append(result.logprob_token_ids)
+        scores.append(result.logprobs)
+        ranks.append(result.selected_token_ranks)
+
+    token_ids = torch.cat(token_ids, dim=0) if len(token_ids) > 1 else token_ids[0]
+    scores = torch.cat(scores, dim=0) if len(scores) > 1 else scores[0]
+    ranks = torch.cat(ranks, dim=0) if len(ranks) > 1 else ranks[0]
+    return token_ids, scores, ranks
+
+
+class CPUPromptLogprobsWorker:
+    def __init__(
+        self,
+        max_num_reqs: int,
+        device: torch.device,
+        logprobs_mode: LogprobsMode = "raw_logprobs",
+    ):
+        self.max_num_reqs = max_num_reqs
+        self.device = device
+        self.logprobs_mode = logprobs_mode
+
+        self.uses_prompt_logprobs = np.zeros(self.max_num_reqs, dtype=bool)
+        self.num_prompt_logprobs = np.zeros(self.max_num_reqs, dtype=np.int32)
+        self.prompt_logprob_start = np.zeros(self.max_num_reqs, dtype=np.int32)
+        # req_idx -> list of in-progress LogprobsTensors
+        self.in_progress_prompt_logprobs: dict[str, list[LogprobsTensors]] = {}
+        # req_id -> fixed-ID scoring state
+        self.token_id_scores: dict[str, _TokenIdScores] = {}
+
+    def add_request(
+        self, req_id: str, req_idx: int, sampling_params: SamplingParams
+    ) -> None:
+        uses_prompt_logprobs = sampling_params.prompt_logprobs is not None
+        self.uses_prompt_logprobs[req_idx] = uses_prompt_logprobs
+        self.num_prompt_logprobs[req_idx] = sampling_params.prompt_logprobs or 0
+        self.prompt_logprob_start[req_idx] = sampling_params.prompt_logprob_start or 0
+        if uses_prompt_logprobs:
+            self.in_progress_prompt_logprobs[req_id] = []
+        if sampling_params.prompt_logprob_token_ids is not None:
+            ids = sampling_params.prompt_logprob_token_ids
+            # Decoded arrays may be read-only; torch.as_tensor warns on those.
+            if not ids.flags.writeable:
+                ids = ids.copy()
+            ids = torch.as_tensor(ids, dtype=torch.int64, device=self.device)
+            self.token_id_scores[req_id] = _TokenIdScores(ids.clamp_min(0), ids < 0)
+
+    def remove_request(self, req_id: str) -> None:
+        self.in_progress_prompt_logprobs.pop(req_id, None)
+        self.token_id_scores.pop(req_id, None)
+
+    def compute_prompt_token_id_logprobs(
+        self,
+        logits_fn: Callable[[torch.Tensor], torch.Tensor],
+        hidden_states: torch.Tensor,
+        input_batch: InputBatch,
+        prompt_lens: np.ndarray,
+    ) -> dict[str, torch.Tensor]:
+        """Fill scores row (prompt row - start) per chunk; emit on the last."""
+        if not self.token_id_scores:
+            return {}
+        logits_mode = self.logprobs_mode in ("raw_logits", "processed_logits")
+        out: dict[str, torch.Tensor] = {}
+        for i, req_id in enumerate(input_batch.req_ids):
+            req = self.token_id_scores.get(req_id)
+            if req is None:
+                continue
+            req_idx = input_batch.idx_mapping_np[i]
+            prompt_len = int(prompt_lens[req_idx])
+            start = int(self.prompt_logprob_start[req_idx])
+            chunk_start = int(input_batch.num_computed_prefill_tokens_np[i])
+            chunk_end = chunk_start + int(input_batch.num_scheduled_tokens[i])
+            # Skip decode steps and, as compute_prompt_logprobs does, requests
+            # resumed after preemption: their scores were already emitted.
+            if chunk_start >= prompt_len or prompt_len < input_batch.prefill_len_np[i]:
+                self.token_id_scores.pop(req_id, None)
+                continue
+            if req.scores is None:
+                if chunk_start > start:
+                    # This prefill starts past the first scored row, so those
+                    # rows will never be written; drop the request instead of
+                    # emitting a buffer with unwritten rows.
+                    self.token_id_scores.pop(req_id, None)
+                    continue
+                req.scores = torch.empty_like(req.token_ids, dtype=torch.float32)
+            # The last prompt row predicts the first decode token; skip it.
+            lo = max(start, chunk_start)
+            hi = min(chunk_end, prompt_len - 1)
+            base = int(input_batch.query_start_loc_np[i]) - chunk_start
+            for a in range(lo, hi, CHUNK_SIZE):
+                b = min(a + CHUNK_SIZE, hi)
+                logits = logits_fn(hidden_states[base + a : base + b])
+                ids = req.token_ids[a - start : b - start]
+                req.scores[a - start : b - start] = (
+                    logits.gather(-1, ids)
+                    if logits_mode
+                    else compute_token_logprobs(logits, ids)
+                )
+            if chunk_end >= prompt_len:
+                req.scores.masked_fill_(req.pad, float("-inf"))
+                out[req_id] = req.scores
+                self.token_id_scores.pop(req_id, None)
+        return out
+
+    def compute_prompt_logprobs(
+        self,
+        logits_fn: Callable[[torch.Tensor], torch.Tensor],
+        hidden_states: torch.Tensor,
+        input_batch: InputBatch,
+        all_token_ids: torch.Tensor,
+        num_computed_tokens: torch.Tensor,
+        prompt_lens: np.ndarray,
+    ) -> dict[str, LogprobsTensors]:
+        idx_mapping_np = input_batch.idx_mapping_np
+        needs_prompt_logprobs = self.uses_prompt_logprobs[idx_mapping_np]
+        if not np.any(needs_prompt_logprobs):
+            # Common case: No request asks for prompt logprobs.
+            return {}
+
+        num_prompt_logprobs = self.num_prompt_logprobs[idx_mapping_np]
+        prompt_logprob_start = self.prompt_logprob_start[idx_mapping_np]
+        prompt_lens = prompt_lens[idx_mapping_np]
+        computed_prefill = input_batch.num_computed_prefill_tokens_np
+        includes_prompt = computed_prefill < prompt_lens
+        # NOTE(woosuk): If the request was resumed after preemption, its prompt
+        # logprobs must have been computed before preemption. Skip.
+        resumed_after_prompt = prompt_lens < input_batch.prefill_len_np
+        needs_prompt_logprobs &= includes_prompt & ~resumed_after_prompt
+        if not np.any(needs_prompt_logprobs):
+            return {}
+
+        # get the maximum number in this batch
+        requested_num_prompt_logprobs = num_prompt_logprobs[needs_prompt_logprobs]
+        max_num_prompt_logprobs = (
+            -1
+            if np.any(requested_num_prompt_logprobs == -1)
+            else int(requested_num_prompt_logprobs.max())
+        )
+
+        # Get the prompt logprobs token_ids.
+        prompt_logprobs_token_ids = get_prompt_logprobs_token_ids(
+            input_batch.num_tokens,
+            input_batch.query_start_loc,
+            input_batch.idx_mapping,
+            num_computed_tokens,
+            all_token_ids,
+        )
+        prompt_token_ids, prompt_logprobs, prompt_ranks = (
+            compute_prompt_logprobs_with_chunking(
+                prompt_logprobs_token_ids,
+                hidden_states[: input_batch.num_tokens],
+                logits_fn,
+                max_num_prompt_logprobs,
+                self.logprobs_mode,
+            )
+        )
+
+        pos_after_step = computed_prefill + input_batch.num_scheduled_tokens
+        is_prompt_chunked = pos_after_step < prompt_lens
+
+        query_start_loc_np = input_batch.query_start_loc_np
+        prompt_logprobs_dict: dict[str, LogprobsTensors] = {}
+        for i, req_id in enumerate(input_batch.req_ids):
+            if not needs_prompt_logprobs[i]:
+                continue
+
+            req_is_prompt_chunked = is_prompt_chunked[i]
+            req_num_prompt_logprobs = int(num_prompt_logprobs[i])
+            start_idx = query_start_loc_np[i]
+            end_idx = query_start_loc_np[i + 1]
+            assert start_idx < end_idx, (
+                f"start_idx ({start_idx}) >= end_idx ({end_idx})"
+            )
+            if not req_is_prompt_chunked:
+                end_idx -= 1
+            # Skip rows before prompt_logprob_start.
+            start_idx += max(int(prompt_logprob_start[i]) - int(computed_prefill[i]), 0)
+
+            width = (
+                prompt_logprobs.shape[1]
+                if req_num_prompt_logprobs == -1
+                else req_num_prompt_logprobs + 1
+            )
+            # no logprobs if start_idx >= end_idx
+            logprobs = (
+                None
+                if start_idx >= end_idx
+                else LogprobsTensors(
+                    logprob_token_ids=prompt_token_ids[start_idx:end_idx, :width],
+                    logprobs=prompt_logprobs[start_idx:end_idx, :width],
+                    selected_token_ranks=prompt_ranks[start_idx:end_idx],
+                )
+            )
+
+            prompt_logprobs_list = self.in_progress_prompt_logprobs[req_id]
+            if logprobs is not None and (req_is_prompt_chunked or prompt_logprobs_list):
+                prompt_logprobs_list.append(logprobs)
+            if req_is_prompt_chunked:
+                # Prompt is chunked. Do not return the logprobs yet.
+                continue
+
+            if prompt_logprobs_list:
+                # Merge the in-progress logprobs.
+                logprobs = LogprobsTensors.cat(prompt_logprobs_list)
+                prompt_logprobs_list.clear()
+
+            if logprobs is None:
+                continue
+
+            prompt_logprobs_dict[req_id] = logprobs
+        return prompt_logprobs_dict
diff --git a/vllm/v1/worker/cpu/lora_utils.py b/vllm/v1/worker/cpu/lora_utils.py
new file mode 100644
index 000000000000..58ea79b2a315
--- /dev/null
+++ b/vllm/v1/worker/cpu/lora_utils.py
@@ -0,0 +1,50 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import numpy as np
+
+from vllm.lora.request import LoRARequest
+
+NO_LORA_ID = 0
+
+
+class LoraState:
+    def __init__(self, max_num_reqs: int):
+        self.lora_ids = np.zeros(max_num_reqs, dtype=np.int32)
+        self.lora_ids.fill(NO_LORA_ID)
+        # req_id -> lora_request
+        self.lora_requests: dict[str, LoRARequest] = {}
+
+    def add_request(
+        self, req_id: str, req_index: int, lora_request: LoRARequest | None
+    ) -> None:
+        if lora_request is not None:
+            self.lora_requests[req_id] = lora_request
+            self.lora_ids[req_index] = lora_request.lora_int_id
+        else:
+            self.lora_ids[req_index] = NO_LORA_ID
+
+    def remove_request(self, req_id: str) -> None:
+        self.lora_requests.pop(req_id, None)
+
+    def make_lora_inputs(
+        self,
+        req_ids: list[str],
+        idx_mapping: np.ndarray,
+        num_scheduled_tokens: np.ndarray,
+        num_logits: np.ndarray,
+    ) -> tuple[tuple[int, ...], tuple[int, ...], set[LoRARequest]]:
+        lora_ids = self.lora_ids[idx_mapping]
+        # There is one lm_head LoRA ID per logits row.
+        prompt_lora_mapping = tuple(lora_ids.repeat(num_logits))
+        token_lora_mapping = tuple(lora_ids.repeat(num_scheduled_tokens))
+        active_lora_requests: set[LoRARequest] = self.get_activate_loras(req_ids)
+        return prompt_lora_mapping, token_lora_mapping, active_lora_requests
+
+    def get_activate_loras(self, req_ids: list[str]) -> set[LoRARequest]:
+        active_lora_requests: set[LoRARequest] = set()
+        for req_id in req_ids:
+            lora_request = self.lora_requests.get(req_id)
+            if lora_request is not None:
+                active_lora_requests.add(lora_request)
+        return active_lora_requests
diff --git a/vllm/v1/worker/cpu/mm_pruning.py b/vllm/v1/worker/cpu/mm_pruning.py
new file mode 100644
index 000000000000..691418da539b
--- /dev/null
+++ b/vllm/v1/worker/cpu/mm_pruning.py
@@ -0,0 +1,131 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import torch
+import torch.nn as nn
+
+from vllm.config import ModelConfig
+from vllm.model_executor.models.interfaces import supports_multimodal_pruning
+from vllm.multimodal.utils import (
+    copy_mm_embedding_modality,
+    get_mm_features_in_window,
+)
+from vllm.v1.worker.cpu.encoder_cache import EncoderCache
+from vllm.v1.worker.cpu.input_batch import InputBatch
+from vllm.v1.worker.cpu.rope import RopeState
+from vllm.v1.worker.cpu.states import CPURequestState
+
+
+class MultiModalPruner:
+    """Recomputes M-RoPE positions for multimodal models that prune embeddings
+    (e.g. Qwen2.5-VL / Qwen3-VL / Nemotron-Nano-VL Efficient Video Sampling).
+
+    Pruning models append their mrope-position channels to the (variable-count)
+    media embeddings from `embed_multimodal`. Those channels must be split off and
+    used to recompute mrope positions before the embeddings are merged.
+    """
+
+    def __init__(
+        self,
+        model: nn.Module,
+        rope_state: RopeState,
+        encoder_cache: EncoderCache,
+        inputs_embeds_size: int,
+    ) -> None:
+        self.model = model
+        self.rope_state = rope_state
+        self.encoder_cache = encoder_cache
+        # The cleaned embedding width: pruning models append their mrope-position
+        # channels as trailing columns, so embeds[:, :inputs_embeds_size] strips them.
+        self.inputs_embeds_size = inputs_embeds_size
+
+    def recompute(
+        self,
+        mm_embeds: list[torch.Tensor],
+        input_batch: InputBatch,
+        req_states: CPURequestState,
+    ) -> list[torch.Tensor]:
+        """Target forward: split the appended mrope-position channels off each
+        request's media embeddings, recompute the corrected mrope positions, and
+        stage them back into RopeState. Returns the cleaned, flattened embeddings.
+        """
+        cleaned: list[torch.Tensor] = []
+        pos = 0
+        req_idx_list = input_batch.idx_mapping_np.tolist()
+        prefill_lens_list = input_batch.prefill_len_np.tolist()
+        num_computed_list = input_batch.num_computed_prefill_tokens_np.tolist()
+        num_scheduled_list = input_batch.num_scheduled_tokens.tolist()
+        for batch_idx, req_id in enumerate(input_batch.req_ids):
+            num_computed = num_computed_list[batch_idx]
+            query_end = num_computed + num_scheduled_list[batch_idx]
+            num_req_embeds = self._num_window_embeds(req_id, num_computed, query_end)
+            if num_req_embeds == 0:
+                continue
+            req_embeds = mm_embeds[pos : pos + num_req_embeds]
+            pos += num_req_embeds
+
+            req_idx = req_idx_list[batch_idx]
+            prefill_len = prefill_lens_list[batch_idx]
+            input_ids = req_states.all_token_ids[req_idx, :prefill_len]
+            mrope_positions = self.rope_state.read_prefill_positions(
+                req_idx, prefill_len
+            ).long()
+            req_cleaned, new_positions, delta = self.model.recompute_mrope_positions(
+                input_ids=input_ids,
+                multimodal_embeddings=req_embeds,
+                mrope_positions=mrope_positions,
+                num_computed_tokens=num_computed,
+            )
+            self.rope_state.update_prefill_positions(req_idx, new_positions, delta)
+            cleaned.extend(
+                copy_mm_embedding_modality(src, dst)
+                for src, dst in zip(req_embeds, req_cleaned)
+            )
+
+        assert pos == len(mm_embeds)
+        return cleaned
+
+    def _num_window_embeds(self, req_id: str, query_start: int, query_end: int) -> int:
+        """Count the media items contributing embeddings to [query_start,
+        query_end), mirroring EncoderRunner.gather_mm_embeddings' per-request
+        windowing so the flat mm_embeds list can be re-segmented per request.
+
+        Note: This logic is intentionally duplicated here rather than being emitted
+        from gather_mm_embeddings, to keep the main path cleaner, since this is a niche
+        feature.
+        """
+        mm_features = self.encoder_cache.mm_features[req_id]
+        lo, hi = get_mm_features_in_window(
+            mm_features, start=query_start, end=query_end
+        )
+        count = 0
+        for mm_feature in mm_features[lo:hi]:
+            pos_info = mm_feature.mm_position
+            start_idx = max(query_start - pos_info.offset, 0)
+            end_idx = min(query_end - pos_info.offset, pos_info.length)
+            embeds_start, embeds_end = pos_info.get_embeds_indices_in_range(
+                start_idx, end_idx
+            )
+            if embeds_start != embeds_end:
+                count += 1
+        return count
+
+
+def maybe_create_mm_pruner(
+    model_config: ModelConfig,
+    model: nn.Module,
+    rope_state: RopeState | None,
+    encoder_cache: EncoderCache | None,
+) -> MultiModalPruner | None:
+    """Create a MultiModalPruner if the model prunes embeddings and uses M-RoPE."""
+    if (
+        rope_state is None
+        or encoder_cache is None
+        or model_config.multimodal_config is None
+        or not model_config.multimodal_config.is_multimodal_pruning_enabled()
+        or not supports_multimodal_pruning(model)
+    ):
+        return None
+
+    return MultiModalPruner(
+        model, rope_state, encoder_cache, model_config.get_inputs_embeds_size()
+    )
diff --git a/vllm/v1/worker/cpu/model_states/__init__.py b/vllm/v1/worker/cpu/model_states/__init__.py
new file mode 100644
index 000000000000..a571f07f1075
--- /dev/null
+++ b/vllm/v1/worker/cpu/model_states/__init__.py
@@ -0,0 +1,57 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import torch
+import torch.nn as nn
+
+from vllm.config import VllmConfig, get_layers_from_vllm_config
+from vllm.model_executor.layers.attention import Attention, CrossAttention
+from vllm.v1.attention.backend import AttentionType
+from vllm.v1.worker.cpu.encoder_cache import EncoderCache
+from vllm.v1.worker.cpu.model_states.interface import ModelState
+
+
+def init_model_state(
+    vllm_config: VllmConfig,
+    model: nn.Module,
+    encoder_cache: EncoderCache | None,
+    device: torch.device,
+) -> ModelState:
+    cls = resolve_model_state_cls(vllm_config, model)
+    return cls(vllm_config, model, encoder_cache, device)
+
+
+def resolve_model_state_cls(
+    vllm_config: VllmConfig, model: nn.Module
+) -> type[ModelState]:
+    # Custom model states are GPU-owned until a CPU implementation is provided.
+    if hasattr(model, "get_model_state_cls"):
+        raise ValueError(
+            f"Custom model state from {type(model).__name__} is not supported "
+            "by CPU Model Runner V2 without Triton."
+        )
+
+    # Cross-attention encoder-decoder models (Whisper, CohereASR, ...).
+    if any(isinstance(module, CrossAttention) for module in model.modules()):
+        from vllm.v1.worker.cpu.model_states.encoder_decoder import (
+            EncoderDecoderModelState,
+        )
+
+        return EncoderDecoderModelState
+
+    # Encoder-only attention is non-causal and needs no KV cache.
+    if any(
+        layer.attn_type == AttentionType.ENCODER_ONLY
+        for layer in get_layers_from_vllm_config(vllm_config, Attention).values()
+    ):
+        from vllm.v1.worker.cpu.model_states.encoder_only import EncoderOnlyModelState
+
+        return EncoderOnlyModelState
+
+    if vllm_config.model_config.is_hybrid or vllm_config.model_config.is_attention_free:
+        from vllm.v1.worker.cpu.model_states.mamba_hybrid import MambaHybridModelState
+
+        return MambaHybridModelState
+
+    from vllm.v1.worker.cpu.model_states.default import DefaultModelState
+
+    return DefaultModelState
diff --git a/vllm/v1/worker/cpu/model_states/default.py b/vllm/v1/worker/cpu/model_states/default.py
new file mode 100644
index 000000000000..0299d5b6a89a
--- /dev/null
+++ b/vllm/v1/worker/cpu/model_states/default.py
@@ -0,0 +1,166 @@
+# 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.v1.attention.backends.utils import compute_mm_prefix_ranges
+from vllm.v1.core.sched.output import NewRequestData
+from vllm.v1.kv_cache_interface import KVCacheConfig
+from vllm.v1.worker.cpu.attn_utils import build_attn_metadata
+from vllm.v1.worker.cpu.encoder_cache import EncoderCache
+from vllm.v1.worker.cpu.input_batch import InputBatch
+from vllm.v1.worker.cpu.mm_pruning import maybe_create_mm_pruner
+from vllm.v1.worker.cpu.model_states.interface import (
+    ModelSpecificAttnMetadata,
+    ModelState,
+)
+from vllm.v1.worker.cpu.rope import get_rope_state
+from vllm.v1.worker.cpu.states import CPURequestState
+from vllm.v1.worker.utils import AttentionGroup
+
+
+class DefaultModelState(ModelState):
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        model: nn.Module,
+        encoder_cache: EncoderCache | None,
+        device: torch.device,
+    ) -> None:
+        super().__init__(vllm_config, model, encoder_cache, device)
+
+        self.rope_state = get_rope_state(
+            self.model_config,
+            model,
+            max_num_reqs=self.max_num_reqs,
+            max_num_tokens=self.max_num_tokens,
+            max_model_len=self.max_model_len,
+            device=self.device,
+        )
+        self.mm_pruner = maybe_create_mm_pruner(
+            self.model_config, model, self.rope_state, encoder_cache
+        )
+
+    def add_request(self, req_index: int, new_req_data: NewRequestData) -> None:
+        if self.rope_state is not None:
+            assert new_req_data.prefill_token_ids is not None
+            # `prompt_embeds` is a passthrough modality with no grid info, but
+            # M-RoPE assumes per-feature grids. Filter it out.
+            mm_features = [
+                f for f in new_req_data.mm_features if f.modality != "prompt_embeds"
+            ]
+            self.rope_state.init_prefill_positions(
+                req_index,
+                self.model,
+                new_req_data.prefill_token_ids,
+                mm_features=mm_features,
+            )
+
+    def dummy_inputs_embeds(self, num_tokens: int) -> torch.Tensor:
+        """Pre-allocated inputs_embeds buffer for dummy runs (contents unused)."""
+        return self.encoder_runner.inputs_embeds[:num_tokens]
+
+    def prepare_inputs_embeds(
+        self,
+        scheduled_encoder_inputs: dict[str, list[int]],
+        input_batch: InputBatch,
+        req_states: CPURequestState,
+    ) -> torch.Tensor:
+        # Use unpadded input_ids to match is_mm_embed size (num_tokens).
+        input_ids_unpadded = input_batch.input_ids[: input_batch.num_tokens]
+
+        self.execute_mm_encoder(scheduled_encoder_inputs)
+
+        mm_embeds, is_mm_embed = super().gather_mm_embeddings(input_batch)
+        if self.mm_pruner is not None and mm_embeds:
+            # EVS: recompute mrope positions for pruned media.
+            mm_embeds = self.mm_pruner.recompute(mm_embeds, input_batch, req_states)
+
+        inputs_embeds = self.encoder_runner.get_inputs_embeds(
+            input_ids_unpadded, mm_embeds, is_mm_embed
+        )
+        return inputs_embeds[: input_batch.num_tokens_after_padding]
+
+    def prepare_inputs(
+        self, input_batch: InputBatch, req_states: CPURequestState
+    ) -> dict[str, torch.Tensor | None]:
+        if self.rope_state is None:
+            return {}  # Common case (1D positions).
+
+        self.rope_state.prepare_positions(
+            input_batch.idx_mapping,
+            input_batch.query_start_loc,
+            req_states.prefill_len,
+            req_states.num_computed_tokens,
+        )
+        positions = self.rope_state.get_positions(input_batch.num_tokens_after_padding)
+        return {"positions": positions}
+
+    def prepare_dummy_inputs(self, num_reqs: int, num_tokens: int) -> dict[str, Any]:
+        model_inputs = {}
+        if self.supports_mm_inputs:
+            model_inputs["inputs_embeds"] = self.dummy_inputs_embeds(num_tokens)
+        if self.rope_state is not None:
+            model_inputs["positions"] = self.rope_state.get_positions(num_tokens)
+        return model_inputs
+
+    def prepare_attn(
+        self,
+        input_batch: InputBatch,
+        block_tables: tuple[torch.Tensor, ...],
+        slot_mappings: torch.Tensor,
+        attn_groups: list[list[AttentionGroup]],
+        kv_cache_config: KVCacheConfig,
+        model_specific_attn_metadata: ModelSpecificAttnMetadata | None = None,
+    ) -> dict[str, Any]:
+        num_reqs = input_batch.num_reqs
+        num_tokens = input_batch.num_tokens
+        query_start_loc = input_batch.query_start_loc[: num_reqs + 1]
+        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
+        max_seq_len = seq_lens_cpu_upper_bound[:num_reqs].max().item()
+
+        req_doc_ranges: dict[int, list[tuple[int, int]]] | None = None
+        if (
+            self.supports_mm_inputs
+            and self.encoder_cache is not None
+            and self.model_config.is_mm_prefix_lm
+        ):
+            mm_features = self.encoder_cache.mm_features
+            hf_text_config = self.model_config.hf_text_config
+            req_doc_ranges = compute_mm_prefix_ranges(
+                [mm_features.get(req_id, ()) for req_id in input_batch.req_ids],
+                sliding_window=self.model_config.get_sliding_window(),
+                clamp_sliding_window=getattr(
+                    self.model, "mm_prefix_clamp_sliding_window", False
+                )
+                or getattr(hf_text_config, "mm_prefix_clamp_sliding_window", False),
+                span_pad=getattr(self.model, "mm_prefix_span_leading_pad_modulus", 0),
+            )
+
+        attn_metadata = build_attn_metadata(
+            attn_groups=attn_groups,
+            num_reqs=num_reqs,
+            num_tokens=num_tokens,
+            query_start_loc=query_start_loc,
+            max_query_len=max_query_len,
+            seq_lens=input_batch.seq_lens,
+            max_seq_len=max_seq_len,
+            block_tables=block_tables,
+            slot_mappings=slot_mappings,
+            kv_cache_config=kv_cache_config,
+            seq_lens_cpu_upper_bound=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,
+            rswa_prefix_lens=input_batch.prompt_lens,
+            req_idx=input_batch.idx_mapping_np,
+            fast_prefill=input_batch.fast_prefill,
+            model_specific_attn_metadata=model_specific_attn_metadata,
+        )
+        return attn_metadata
diff --git a/vllm/v1/worker/cpu/model_states/encoder_decoder.py b/vllm/v1/worker/cpu/model_states/encoder_decoder.py
new file mode 100644
index 000000000000..18f9584bf60b
--- /dev/null
+++ b/vllm/v1/worker/cpu/model_states/encoder_decoder.py
@@ -0,0 +1,167 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from dataclasses import dataclass
+from typing import Any
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+from vllm.config import VllmConfig
+from vllm.v1.kv_cache_interface import CrossAttentionSpec, KVCacheConfig
+from vllm.v1.worker.cpu.attn_utils import build_attn_metadata
+from vllm.v1.worker.cpu.encoder_cache import EncoderCache
+from vllm.v1.worker.cpu.input_batch import InputBatch
+from vllm.v1.worker.cpu.model_states.interface import (
+    ModelSpecificAttnMetadata,
+    ModelState,
+)
+from vllm.v1.worker.cpu.states import CPURequestState
+from vllm.v1.worker.utils import AttentionGroup
+
+
+@dataclass
+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
+    ) -> dict[str, Any]:
+        encoder_seq_lens = self.encoder_seq_lens.get(kv_cache_group_id)
+        if encoder_seq_lens is None:
+            return {}
+        encoder_seq_lens_tensor, encoder_seq_lens_np = encoder_seq_lens
+        return {
+            "encoder_seq_lens": encoder_seq_lens_tensor[:num_reqs],
+            "encoder_seq_lens_cpu": encoder_seq_lens_np[:num_reqs],
+        }
+
+
+class EncoderDecoderModelState(ModelState):
+    """ModelState for cross-attention encoder-decoder models
+    (Whisper, CohereASR, NemotronParse, ...)
+    """
+
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        model: nn.Module,
+        encoder_cache: EncoderCache | None,
+        device: torch.device,
+    ) -> None:
+        assert encoder_cache is not None
+        if vllm_config.model_config.enable_prompt_embeds:
+            raise ValueError(
+                "--enable-prompt-embeds is not supported with encoder-decoder models."
+            )
+        super().__init__(vllm_config, model, encoder_cache, device)
+
+        self.max_encoder_len = getattr(
+            self.model_config.hf_config,
+            "max_source_positions",
+            self.max_model_len,
+        )
+        self.encoder_seq_lens = torch.zeros(
+            self.max_num_reqs, dtype=torch.int32, device=self.device
+        )
+        self.encoder_seq_lens_np = self.encoder_seq_lens.numpy()
+        self.encoder_outputs: list[torch.Tensor] = []
+
+    def prepare_inputs_embeds(
+        self,
+        scheduled_encoder_inputs: dict[str, list[int]],
+        input_batch: InputBatch,
+        req_states: CPURequestState,
+    ) -> None:
+        # Ensure encoder inputs are ordered consistently with input_batch.req_ids.
+        encoder_inputs: dict[str, list[int]] = {}
+        for req_id in input_batch.req_ids:
+            req_encoder_inputs = scheduled_encoder_inputs.get(req_id, [])
+            if req_encoder_inputs:
+                encoder_inputs[req_id] = req_encoder_inputs
+
+        _, mm_kwargs = self.encoder_runner.prepare_mm_inputs(encoder_inputs)
+        if mm_kwargs:
+            # Encoder-decoder models consume encoder outputs through the
+            # `encoder_outputs` forward kwarg, not `inputs_embeds`. Single modality
+            # so execute_mm_encoder preserves request order; use its return value
+            # directly. No need to store in encoder_cache: cross-attention K/V are
+            # written to the KV cache on the first step; decode steps use the cache.
+            with self.encoder_runner.timed_encoder_operation(encoder_inputs.keys()):
+                self.encoder_outputs = self.encoder_runner.execute_mm_encoder(mm_kwargs)
+        else:
+            # Decode steps: encoder K/V are in cross-attention KV cache.
+            self.encoder_outputs = []
+        return None
+
+    def prepare_inputs(
+        self, input_batch: InputBatch, req_states: CPURequestState
+    ) -> dict[str, Any]:
+        model_inputs = {"encoder_outputs": self.encoder_outputs}
+        self.encoder_outputs = []
+        return model_inputs
+
+    def prepare_dummy_inputs(self, num_reqs: int, num_tokens: int) -> dict[str, Any]:
+        return {"encoder_outputs": []}
+
+    def prepare_attn(
+        self,
+        input_batch: InputBatch,
+        block_tables: tuple[torch.Tensor, ...],
+        slot_mappings: torch.Tensor,
+        attn_groups: list[list[AttentionGroup]],
+        kv_cache_config: KVCacheConfig,
+    ) -> dict[str, Any]:
+        num_reqs = input_batch.num_reqs
+        num_tokens = input_batch.num_tokens
+        enc_dec_attn_metadata = EncoderDecoderAttnMetadata(
+            self._get_encoder_seq_lens(input_batch.req_ids, attn_groups)
+        )
+
+        max_query_len = input_batch.num_scheduled_tokens.max().item()
+        seq_lens_cpu_upper_bound = input_batch.seq_lens_cpu_upper_bound
+        max_seq_len = int(seq_lens_cpu_upper_bound[:num_reqs].max())
+        attn_metadata = build_attn_metadata(
+            attn_groups=attn_groups,
+            num_reqs=num_reqs,
+            num_tokens=num_tokens,
+            query_start_loc=input_batch.query_start_loc,
+            max_query_len=max_query_len,
+            seq_lens=input_batch.seq_lens,
+            max_seq_len=max_seq_len,
+            block_tables=block_tables,
+            slot_mappings=slot_mappings,
+            kv_cache_config=kv_cache_config,
+            seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,
+            model_specific_attn_metadata=enc_dec_attn_metadata,
+            rswa_prefix_lens=input_batch.prompt_lens,
+        )
+        return attn_metadata
+
+    def _get_encoder_seq_lens(
+        self,
+        req_ids: list[str],
+        attn_groups: list[list[AttentionGroup]],
+    ) -> dict[int, tuple[torch.Tensor, np.ndarray]]:
+        num_reqs = len(req_ids)
+        encoder_seq_lens_np = self.encoder_seq_lens_np[:num_reqs]
+        encoder_seq_lens_np.fill(0)
+        for i, req_id in enumerate(req_ids):
+            mm_features = self.encoder_cache.mm_features.get(req_id, [])
+            encoder_seq_lens_np[i] = sum(
+                feature.mm_position.get_num_embeds() for feature in mm_features
+            )
+
+        seq_lens_by_group: dict[int, tuple[torch.Tensor, np.ndarray]] = {}
+        encoder_seq_lens = self.encoder_seq_lens[:num_reqs]
+        for kv_cache_group_idx, groups in enumerate(attn_groups):
+            has_cross_attn = any(
+                isinstance(attn_group.kv_cache_spec, CrossAttentionSpec)
+                for attn_group in groups
+            )
+            if has_cross_attn:
+                seq_lens_by_group[kv_cache_group_idx] = (
+                    encoder_seq_lens,
+                    encoder_seq_lens_np,
+                )
+        return seq_lens_by_group
diff --git a/vllm/v1/worker/cpu/model_states/encoder_only.py b/vllm/v1/worker/cpu/model_states/encoder_only.py
new file mode 100644
index 000000000000..f3f2f5b9113e
--- /dev/null
+++ b/vllm/v1/worker/cpu/model_states/encoder_only.py
@@ -0,0 +1,189 @@
+# 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, get_layers_from_vllm_config
+from vllm.model_executor.layers.attention import Attention
+from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
+from vllm.v1.attention.backend import AttentionType, CommonAttentionMetadata
+from vllm.v1.core.sched.output import NewRequestData
+from vllm.v1.kv_cache_interface import EncoderOnlyAttentionSpec, KVCacheConfig
+from vllm.v1.worker.cpu.encoder_cache import EncoderCache
+from vllm.v1.worker.cpu.input_batch import InputBatch
+from vllm.v1.worker.cpu.model_states.default import DefaultModelState
+from vllm.v1.worker.cpu.model_states.interface import ModelSpecificAttnMetadata
+from vllm.v1.worker.cpu.states import CPURequestState
+from vllm.v1.worker.utils import AttentionGroup
+
+
+class EncoderOnlyModelState(DefaultModelState):
+    """ModelState for encoder-only (BERT/RoBERTa) models.
+
+    Encoder attention needs no KV cache: it runs full, bidirectional
+    self-attention over each request's tokens in a single forward. Such layers
+    return no ``get_kv_cache_spec`` and therefore never enter
+    ``kv_cache_config.kv_cache_groups``, so the KV-backed attention path never
+    builds their metadata. We build their (non-causal) metadata here, keeping
+    the normal KV-backed path untouched.
+    """
+
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        model: nn.Module,
+        encoder_cache: EncoderCache | None,
+        device: torch.device,
+    ) -> None:
+        super().__init__(vllm_config, model, encoder_cache, device)
+
+        cache_config = vllm_config.cache_config
+        if cache_config.cache_dtype == "auto":
+            kv_cache_dtype = self.dtype
+        else:
+            kv_cache_dtype = STR_DTYPE_TO_TORCH_DTYPE[cache_config.cache_dtype]
+
+        # Build an attention group (and its non-causal metadata builder) for the
+        # encoder-only layers, grouped by backend + query/KV head config. Models
+        # are typically uniform, yielding a single group.
+        attn_layers = get_layers_from_vllm_config(vllm_config, Attention)
+        groups: dict[tuple[tuple[str, str], int, int, int], AttentionGroup] = {}
+        for name, layer in attn_layers.items():
+            if layer.attn_type != AttentionType.ENCODER_ONLY:
+                continue
+            # No KV cache is bound for these layers; give them an (empty)
+            # device tensor so the attention forward context is well-formed.
+            layer.kv_cache = torch.empty(0, dtype=kv_cache_dtype, device=device)
+            backend = layer.get_attn_backend()
+            # Metadata builders require a uniform query-head count per group.
+            key = (
+                backend.full_cls_name(),
+                layer.num_heads,
+                layer.num_kv_heads,
+                layer.head_size,
+            )
+            group = groups.get(key)
+            if group is None:
+                spec = EncoderOnlyAttentionSpec(
+                    block_size=cache_config.block_size,
+                    num_kv_heads=layer.num_kv_heads,
+                    head_size=layer.head_size,
+                    dtype=kv_cache_dtype,
+                )
+                group = AttentionGroup(backend, [], spec, kv_cache_group_id=len(groups))
+                groups[key] = group
+            group.layer_names.append(name)
+
+        self.encoder_attn_groups = list(groups.values())
+        for group in self.encoder_attn_groups:
+            group.create_metadata_builders(vllm_config, device)
+
+        # Encoder attention reads neither the block table nor the slot mapping
+        # (full varlen self-attention over the batch's q/k/v), but the metadata
+        # builder expects both tensors to be present.
+        self._dummy_block_table = torch.zeros(
+            self.max_num_reqs, 1, dtype=torch.int32, device=device
+        )
+        self._dummy_slot_mapping = torch.zeros(
+            self.max_num_tokens, dtype=torch.int64, device=device
+        )
+        self.token_type_ids: dict[str, torch.Tensor] = {}
+
+    def add_request(self, req_index: int, new_req_data: NewRequestData) -> None:
+        super().add_request(req_index, new_req_data)
+        pooling_params = new_req_data.pooling_params
+        if pooling_params is None or pooling_params.extra_kwargs is None:
+            return
+
+        token_type_range = pooling_params.extra_kwargs.get("compressed_token_type_ids")
+        if token_type_range is not None:
+            assert new_req_data.prompt_token_ids is not None
+            token_type_start, token_type_end = token_type_range
+            positions = torch.arange(
+                len(new_req_data.prompt_token_ids), dtype=torch.int32
+            )
+            self.token_type_ids[new_req_data.req_id] = (
+                (positions >= token_type_start) & (positions < token_type_end)
+            ).to(torch.int32)
+
+    def remove_request(self, req_id: str) -> None:
+        super().remove_request(req_id)
+        self.token_type_ids.pop(req_id, None)
+
+    def prepare_inputs(
+        self, input_batch: InputBatch, req_states: CPURequestState
+    ) -> dict[str, torch.Tensor | None]:
+        model_inputs = super().prepare_inputs(input_batch, req_states)
+        if not self.token_type_ids:
+            return model_inputs
+
+        token_type_ids = torch.zeros(
+            input_batch.num_tokens_after_padding,
+            dtype=torch.int32,
+            device=self.device,
+        )
+        offset = 0
+        for i, req_id in enumerate(input_batch.req_ids):
+            num_tokens = int(input_batch.num_scheduled_tokens[i])
+            request_token_type_ids = self.token_type_ids.get(req_id)
+            if request_token_type_ids is not None:
+                start = int(input_batch.num_computed_tokens_np[i])
+                token_type_ids[offset : offset + num_tokens].copy_(
+                    request_token_type_ids[start : start + num_tokens]
+                )
+            offset += num_tokens
+        model_inputs["token_type_ids"] = token_type_ids
+        return model_inputs
+
+    def prepare_attn(
+        self,
+        input_batch: InputBatch,
+        block_tables: tuple[torch.Tensor, ...],
+        slot_mappings: torch.Tensor,
+        attn_groups: list[list[AttentionGroup]],
+        kv_cache_config: KVCacheConfig,
+        model_specific_attn_metadata: ModelSpecificAttnMetadata | None = None,
+    ) -> dict[str, Any]:
+        attn_metadata = super().prepare_attn(
+            input_batch,
+            block_tables,
+            slot_mappings,
+            attn_groups,
+            kv_cache_config,
+            model_specific_attn_metadata=model_specific_attn_metadata,
+        )
+        attn_metadata.update(self._build_encoder_attn_metadata(input_batch))
+        return attn_metadata
+
+    def _build_encoder_attn_metadata(self, input_batch: InputBatch) -> dict[str, Any]:
+        num_reqs = input_batch.num_reqs
+        num_tokens = input_batch.num_tokens
+        max_query_len = int(input_batch.num_scheduled_tokens.max())
+        max_seq_len = int(input_batch.seq_lens_cpu_upper_bound[:num_reqs].max())
+
+        # The encoder builder forces ``causal=False`` regardless of this value.
+        common_attn_metadata = CommonAttentionMetadata(
+            query_start_loc=input_batch.query_start_loc,
+            query_start_loc_cpu=input_batch.query_start_loc,
+            seq_lens=input_batch.seq_lens[:num_reqs],
+            num_reqs=num_reqs,
+            num_actual_tokens=num_tokens,
+            max_query_len=max_query_len,
+            max_seq_len=max_seq_len,
+            block_table_tensor=self._dummy_block_table[:num_reqs],
+            slot_mapping=self._dummy_slot_mapping[:num_tokens],
+            seq_lens_cpu_upper_bound=input_batch.seq_lens_cpu_upper_bound[:num_reqs],
+            positions=input_batch.positions,
+        )
+
+        attn_metadata: dict[str, Any] = {}
+        for group in self.encoder_attn_groups:
+            metadata = group.get_metadata_builder(0).build(
+                common_prefix_len=0,
+                common_attn_metadata=common_attn_metadata,
+            )
+            for name in group.layer_names:
+                attn_metadata[name] = metadata
+        return attn_metadata
diff --git a/vllm/v1/worker/cpu/model_states/interface.py b/vllm/v1/worker/cpu/model_states/interface.py
new file mode 100644
index 000000000000..f19e78e28e72
--- /dev/null
+++ b/vllm/v1/worker/cpu/model_states/interface.py
@@ -0,0 +1,186 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from abc import ABC, abstractmethod
+from typing import Any
+
+import torch
+import torch.nn as nn
+
+from vllm.config import VllmConfig
+from vllm.tasks import GenerationTask
+from vllm.v1.core.sched.output import NewRequestData
+from vllm.v1.kv_cache_interface import KVCacheConfig
+from vllm.v1.worker.cpu.block_table import CPUBlockTables
+from vllm.v1.worker.cpu.encoder_cache import EncoderCache
+from vllm.v1.worker.cpu.encoder_runner import EncoderRunner
+from vllm.v1.worker.cpu.input_batch import InputBatch
+from vllm.v1.worker.cpu.states import CPURequestState
+from vllm.v1.worker.utils import AttentionGroup
+
+
+class ModelSpecificAttnMetadata:
+    """Base class for model-specific attention metadata."""
+
+    def get_extra_common_attn_kwargs(
+        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
+    ) -> dict[str, Any]:
+        return {}
+
+
+class ModelState(ABC):
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        model: nn.Module,
+        encoder_cache: EncoderCache | None,
+        device: torch.device,
+    ) -> None:
+        self.vllm_config = vllm_config
+        self.model_config = vllm_config.model_config
+        self.scheduler_config = vllm_config.scheduler_config
+        self.model = model
+        self.device = device
+
+        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()
+        self.dtype = self.model_config.dtype
+
+        self.supports_mm_inputs = encoder_cache is not None
+        if encoder_cache is not None:
+            self.encoder_cache = encoder_cache
+            observability_config = vllm_config.observability_config
+            self.encoder_runner = EncoderRunner(
+                model=self.model,
+                max_num_tokens=self.max_num_tokens,
+                hidden_size=self.inputs_embeds_size,
+                encoder_cache=encoder_cache,
+                dtype=self.dtype,
+                device=self.device,
+                enable_timing=bool(
+                    observability_config
+                    and observability_config.enable_mm_processor_stats
+                ),
+            )
+
+    @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,
+            supports_transcription,
+        )
+        from vllm.model_executor.models.interfaces_base import is_text_generation_model
+
+        supported_tasks = list[GenerationTask]()
+        if is_text_generation_model(self.model):
+            supported_tasks.append("generate")
+        if supports_transcription(self.model):
+            if self.model.supports_transcription_only:
+                return ("transcription",)
+            supported_tasks.append("transcription")
+        if supports_realtime(self.model):
+            supported_tasks.append("realtime")
+        return tuple(supported_tasks)
+
+    def add_request(self, req_index: int, new_req_data: NewRequestData) -> None:
+        return None
+
+    def remove_request(self, req_id: str) -> None:
+        return None
+
+    def initialize_kv_cache(
+        self, kv_cache_config: KVCacheConfig, block_tables: CPUBlockTables
+    ) -> None:
+        """Hook run after the KV cache tensors are allocated and bound."""
+        return None
+
+    def preprocess_state(
+        self,
+        input_batch: InputBatch,
+        block_tables: tuple[torch.Tensor, ...],
+        kv_cache_config: KVCacheConfig,
+        num_computed_tokens: torch.Tensor,
+    ) -> None:
+        """Hook run on real batches before the forward pass (after block tables
+        are gathered). Used by mamba "align" prefix caching to pre-copy state
+        across block boundaries. No-op by default."""
+        return None
+
+    @abstractmethod
+    def prepare_inputs_embeds(
+        self,
+        scheduled_encoder_inputs: dict[str, list[int]],
+        input_batch: InputBatch,
+        req_states: CPURequestState,
+    ) -> torch.Tensor | None:
+        """Prepare the ``inputs_embeds`` tensor for the current forward pass."""
+        raise NotImplementedError
+
+    def dummy_inputs_embeds(self, num_tokens: int) -> torch.Tensor | None:
+        """Pre-allocated inputs_embeds buffer for dummy runs (contents unused)."""
+        return None
+
+    def execute_mm_encoder(
+        self, scheduled_encoder_inputs: dict[str, list[int]]
+    ) -> None:
+        """Run the multi-modal encoder and cache its outputs by `mm_hash`.
+
+        The encode half of `get_mm_embeddings`, without the gather, for callers
+        that run no language model.
+        """
+        mm_hashes, mm_kwargs = self.encoder_runner.prepare_mm_inputs(
+            scheduled_encoder_inputs
+        )
+        if mm_kwargs:
+            with self.encoder_runner.timed_encoder_operation(
+                scheduled_encoder_inputs.keys()
+            ):
+                encoder_outputs = self.encoder_runner.execute_mm_encoder(mm_kwargs)
+            self.encoder_cache.encoder_outputs.update(zip(mm_hashes, encoder_outputs))
+
+    def gather_mm_embeddings(
+        self, input_batch: InputBatch
+    ) -> tuple[list[torch.Tensor], torch.Tensor]:
+        """Gather cached multimodal embeddings."""
+        return self.encoder_runner.gather_mm_embeddings(
+            input_batch.req_ids,
+            input_batch.num_tokens,
+            input_batch.num_scheduled_tokens,
+            input_batch.query_start_loc_np,
+            input_batch.prefill_len_np,
+            input_batch.num_computed_tokens_np,
+        )
+
+    @abstractmethod
+    def prepare_inputs(
+        self, input_batch: InputBatch, req_states: CPURequestState
+    ) -> dict[str, Any]:
+        raise NotImplementedError
+
+    @abstractmethod
+    def prepare_dummy_inputs(self, num_reqs: int, num_tokens: int) -> dict[str, Any]:
+        raise NotImplementedError
+
+    @abstractmethod
+    def prepare_attn(
+        self,
+        input_batch: InputBatch,
+        block_tables: tuple[torch.Tensor, ...],
+        slot_mappings: torch.Tensor,
+        attn_groups: list[list[AttentionGroup]],
+        kv_cache_config: KVCacheConfig,
+    ) -> dict[str, Any]:
+        raise NotImplementedError
+
+    num_new_sampled_tokens_per_step: int = 1
+    """New tokens sampled on each decode step
+    (excluding accepted draft tokens, a.k.a num bonus tokens)."""
diff --git a/vllm/v1/worker/cpu/model_states/mamba_hybrid.py b/vllm/v1/worker/cpu/model_states/mamba_hybrid.py
new file mode 100644
index 000000000000..3ccaf501ba02
--- /dev/null
+++ b/vllm/v1/worker/cpu/model_states/mamba_hybrid.py
@@ -0,0 +1,149 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import ctypes
+
+import torch
+import torch.nn as nn
+
+from vllm.config import VllmConfig
+from vllm.model_executor.layers.mamba.mamba_utils import MambaStateCopyFuncsByType
+from vllm.v1.core.sched.output import NewRequestData
+from vllm.v1.kv_cache_interface import (
+    KVCacheConfig,
+    KVCacheGroupSpec,
+    MambaSpec,
+    UniformTypeKVCacheSpecs,
+)
+from vllm.v1.worker.cpu.block_table import CPUBlockTables
+from vllm.v1.worker.cpu.encoder_cache import EncoderCache
+from vllm.v1.worker.cpu.input_batch import InputBatch
+from vllm.v1.worker.cpu.model_states.default import DefaultModelState
+from vllm.v1.worker.mamba_utils import (
+    get_mamba_group_ids,
+    get_mamba_groups,
+    validate_mamba_state_copy_funcs,
+)
+
+
+def _get_mamba_spec(kv_cache_group: KVCacheGroupSpec, layer_name: str) -> MambaSpec:
+    spec = kv_cache_group.kv_cache_spec
+    if isinstance(spec, UniformTypeKVCacheSpecs):
+        spec = spec.kv_cache_specs[layer_name]
+    assert isinstance(spec, MambaSpec)
+    return spec
+
+
+class MambaHybridModelState(DefaultModelState):
+    """Model state for hybrid attention and recurrent CPU models."""
+
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        model: nn.Module,
+        encoder_cache: EncoderCache | None,
+        device: torch.device,
+    ) -> None:
+        super().__init__(vllm_config, model, encoder_cache, device)
+        self.cache_config = vllm_config.cache_config
+        self._align_mode = self.cache_config.mamba_cache_mode == "align"
+        if self._align_mode:
+            self._mamba_state_idx = torch.zeros(
+                self.max_num_reqs, dtype=torch.int32, device=self.device
+            )
+            self._mamba_group_ids: list[int] = []
+            self._mamba_spec: MambaSpec | None = None
+            self._mamba_state_copy_funcs: MambaStateCopyFuncsByType | None = None
+
+    def add_request(self, req_index: int, new_req_data: NewRequestData) -> None:
+        super().add_request(req_index, new_req_data)
+        if self._align_mode:
+            mamba_block_size = self.cache_config.mamba_block_size
+            assert mamba_block_size is not None
+            self._mamba_state_idx[req_index] = (
+                new_req_data.num_computed_tokens - 1
+            ) // mamba_block_size
+
+    def _get_mamba_group_info(
+        self, kv_cache_config: KVCacheConfig
+    ) -> tuple[list[int], MambaSpec]:
+        if self._mamba_spec is None:
+            mamba_groups = get_mamba_groups(kv_cache_config)
+            mamba_spec = next(iter(mamba_groups))
+            assert all(
+                spec.block_size == mamba_spec.block_size
+                and spec.num_speculative_blocks == mamba_spec.num_speculative_blocks
+                and spec.mamba_cache_mode == mamba_spec.mamba_cache_mode
+                for spec in mamba_groups
+            ), "all mamba groups must share cache scheduling parameters"
+            self._mamba_group_ids = get_mamba_group_ids(mamba_groups)
+            self._mamba_spec = mamba_spec
+        return self._mamba_group_ids, self._mamba_spec
+
+    def initialize_kv_cache(
+        self, kv_cache_config: KVCacheConfig, block_tables: CPUBlockTables
+    ) -> None:
+        if not self._align_mode:
+            return
+        self._get_mamba_group_info(kv_cache_config)
+        if self._mamba_state_copy_funcs is None:
+            mamba_groups = get_mamba_groups(kv_cache_config)
+            mamba_types = {spec.mamba_type for spec in mamba_groups}
+            copy_funcs = self.model.get_mamba_state_copy_funcs(mamba_types)
+            validate_mamba_state_copy_funcs(mamba_groups, copy_funcs)
+            self._mamba_state_copy_funcs = copy_funcs
+
+    def preprocess_state(
+        self,
+        input_batch: InputBatch,
+        block_tables: tuple[torch.Tensor, ...],
+        kv_cache_config: KVCacheConfig,
+        num_computed_tokens: torch.Tensor,
+    ) -> None:
+        if not self._align_mode or input_batch.num_reqs == 0:
+            return
+
+        mamba_group_ids, mamba_spec = self._get_mamba_group_info(kv_cache_config)
+        req_indices = input_batch.idx_mapping
+        previous_state_indices = self._mamba_state_idx[req_indices]
+        query_lens = input_batch.query_start_loc[1:] - input_batch.query_start_loc[:-1]
+        computed_after = num_computed_tokens[req_indices] + query_lens
+        current_state_indices = (
+            computed_after + mamba_spec.block_size - 1
+        ) // mamba_spec.block_size - 1
+        self._mamba_state_idx[req_indices] = current_state_indices
+
+        crossed_rows = (
+            (previous_state_indices >= 0)
+            & (previous_state_indices != current_state_indices)
+        ).nonzero(as_tuple=True)[0]
+        if crossed_rows.numel() == 0:
+            return
+
+        copy_funcs_by_type = self._mamba_state_copy_funcs
+        assert copy_funcs_by_type is not None
+        forward_context = self.vllm_config.compilation_config.static_forward_context
+        for batch_idx in crossed_rows.tolist():
+            src_col = int(previous_state_indices[batch_idx])
+            dst_col = int(current_state_indices[batch_idx])
+            for group_id in mamba_group_ids:
+                kv_cache_group = kv_cache_config.kv_cache_groups[group_id]
+                block_ids = block_tables[group_id][batch_idx].tolist()
+                dst_block_id = block_ids[dst_col]
+                for layer_name in kv_cache_group.layer_names:
+                    mamba_type = _get_mamba_spec(kv_cache_group, layer_name).mamba_type
+                    copy_funcs = copy_funcs_by_type[mamba_type]
+                    states = forward_context[layer_name].kv_cache
+                    if len(states) < len(copy_funcs):
+                        raise ValueError(
+                            f"Expected at least {len(copy_funcs)} Mamba state "
+                            f"tensors, got {len(states)}"
+                        )
+                    for state, copy_func in zip(states, copy_funcs):
+                        # This runner rejects speculative decoding, so the
+                        # accepted-token offset is always zero (1 - 1).
+                        copy_spec = copy_func(state, block_ids, src_col, 1)
+                        ctypes.memmove(
+                            state[dst_block_id].data_ptr(),
+                            copy_spec.start_addr,
+                            copy_spec.num_elements * state.element_size(),
+                        )
diff --git a/vllm/v1/worker/cpu/pooling_runner.py b/vllm/v1/worker/cpu/pooling_runner.py
new file mode 100644
index 000000000000..7d454db64251
--- /dev/null
+++ b/vllm/v1/worker/cpu/pooling_runner.py
@@ -0,0 +1,208 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from typing import cast
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+from vllm.config import VllmConfig
+from vllm.model_executor.models import VllmModelForPooling, is_pooling_model
+from vllm.pooling_params import PoolingParams
+from vllm.tasks import PoolingTask
+from vllm.v1.outputs import PoolerOutput
+from vllm.v1.pool.late_interaction_runner import LateInteractionRunner
+from vllm.v1.pool.metadata import PoolingMetadata, PoolingStates
+from vllm.v1.worker.cpu.input_batch import InputBatch
+from vllm.v1.worker.cpu.states import CPURequestState
+
+_SUPPORTED_TASKS: frozenset[PoolingTask] = frozenset(
+    {"embed", "classify", "token_embed", "token_classify", "embed&token_classify"}
+)
+
+
+class PoolingRunner:
+    def __init__(self, model: nn.Module, vllm_config: VllmConfig):
+        self.model = cast(VllmModelForPooling, model)
+        self.model_config = vllm_config.model_config
+        self.max_num_reqs = vllm_config.scheduler_config.max_num_seqs
+        model_tasks = tuple(sorted(self.model.pooler.get_supported_tasks()))
+        selected_task = self.model_config.get_pooling_task(model_tasks)
+        if selected_task not in _SUPPORTED_TASKS:
+            hint = (
+                "Set an explicitly supported task or VLLM_USE_V2_MODEL_RUNNER=0."
+                if _SUPPORTED_TASKS.intersection(model_tasks)
+                else "Set VLLM_USE_V2_MODEL_RUNNER=0 to use this model."
+            )
+            raise ValueError(
+                "Model Runner V2 supports pooling tasks "
+                f"{sorted(_SUPPORTED_TASKS)}, but this model selects "
+                f"{selected_task!r} from {list(model_tasks)}. {hint}"
+            )
+        self.supported_tasks = frozenset(self.get_supported_tasks(model))
+        if not self.supported_tasks:
+            raise ValueError(
+                "Model Runner V2 supports pooling tasks "
+                f"{sorted(_SUPPORTED_TASKS)}, but this model supports "
+                f"{list(model_tasks)}. "
+                "Set VLLM_USE_V2_MODEL_RUNNER=0 to use this model."
+            )
+        self.pooling_params: dict[int, PoolingParams] = {}
+        self.pooling_states: dict[int, PoolingStates] = {}
+        self.prompt_token_ids: dict[int, torch.Tensor] = {}
+        pooler_config = self.model_config.pooler_config
+        self.late_interaction_runner = LateInteractionRunner(
+            enable_flash=pooler_config is None
+            or getattr(pooler_config, "enable_flash_late_interaction", True)
+        )
+
+    @staticmethod
+    def get_supported_tasks(model: nn.Module) -> list[PoolingTask]:
+        if not is_pooling_model(model):
+            return []
+        return sorted(model.pooler.get_supported_tasks() & _SUPPORTED_TASKS)
+
+    def add_request(
+        self,
+        req_id: str,
+        req_index: int,
+        pooling_params: PoolingParams,
+        prompt_token_ids: list[int] | None,
+    ) -> None:
+        task = pooling_params.task
+        if task not in self.supported_tasks:
+            raise ValueError(
+                f"Unsupported task: {task!r}. "
+                f"Supported tasks: {sorted(self.supported_tasks)}"
+            )
+        self.model.pooler.get_pooling_updates(task).apply(pooling_params)
+        self.late_interaction_runner.register_request(req_id, pooling_params)
+        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
+            )
+
+    def remove_request(self, req_index: int) -> None:
+        self.pooling_params.pop(req_index, None)
+        if state := self.pooling_states.pop(req_index, None):
+            state.clean()
+        self.prompt_token_ids.pop(req_index, None)
+
+    def on_requests_finished(self, req_ids: set[str]) -> None:
+        self.late_interaction_runner.on_requests_finished(req_ids)
+
+    def clear(self) -> None:
+        self.late_interaction_runner.clear()
+
+    def release_late_interaction_query_cache(self, query_keys: list[str]) -> None:
+        self.late_interaction_runner.release_queries(query_keys)
+
+    def _get_pooling_metadata(
+        self,
+        input_batch: InputBatch,
+        req_states: CPURequestState,
+    ) -> PoolingMetadata:
+        req_indices = input_batch.idx_mapping_np.tolist()
+        pooling_params = [self.pooling_params[i] for i in req_indices]
+        pooling_states = [self.pooling_states[i] for i in req_indices]
+        prompt_lens = req_states.prompt_len[input_batch.idx_mapping]
+
+        prompt_token_ids_cpu = None
+        if any(params.requires_token_ids for params in pooling_params):
+            max_prompt_len = int(prompt_lens.max())
+            prompt_token_ids_cpu = torch.zeros(
+                (input_batch.num_reqs, max_prompt_len),
+                dtype=torch.int64,
+                device=req_states.device,
+            )
+            for i, (req_index, params) in enumerate(zip(req_indices, pooling_params)):
+                if not params.requires_token_ids:
+                    continue
+                token_ids = self.prompt_token_ids[req_index]
+                prompt_token_ids_cpu[i, : token_ids.numel()] = token_ids
+
+        return PoolingMetadata(
+            prompt_lens=prompt_lens,
+            prompt_token_ids=None,
+            prompt_token_ids_cpu=prompt_token_ids_cpu,
+            pooling_params=pooling_params,
+            pooling_states=pooling_states,
+        )
+
+    def pool(
+        self,
+        hidden_states: torch.Tensor,
+        input_batch: InputBatch,
+        req_states: CPURequestState,
+    ) -> tuple[PoolerOutput, list[bool]]:
+        hidden_states = hidden_states[: input_batch.num_tokens]
+        pooling_metadata = self._get_pooling_metadata(input_batch, req_states)
+        num_reqs = input_batch.num_reqs
+        # Pooling has no speculative tokens, so this CPU upper bound is exact.
+        seq_lens_cpu = input_batch.seq_lens_cpu_upper_bound[:num_reqs]
+        pooling_metadata.build_pooling_cursor(
+            input_batch.num_scheduled_tokens,
+            seq_lens_cpu,
+            device=hidden_states.device,
+            query_start_loc_gpu=input_batch.query_start_loc[: num_reqs + 1],
+        )
+        pooler_output = self.model.pooler(hidden_states, pooling_metadata)
+        finished_mask = pooling_metadata.get_pooling_cursor().get_finished_mask()
+        pooler_output = self.late_interaction_runner.postprocess_pooler_output(
+            raw_pooler_output=pooler_output,
+            pooling_params=pooling_metadata.pooling_params,
+            req_ids=input_batch.req_ids,
+            finished_mask=finished_mask,
+        )
+        return pooler_output, finished_mask
+
+    def _dummy_pooler_run_task(
+        self, hidden_states: torch.Tensor, task: PoolingTask
+    ) -> PoolerOutput:
+        num_tokens = hidden_states.shape[0]
+        num_reqs = min(num_tokens, self.max_num_reqs)
+        base_tokens = num_tokens // num_reqs
+        num_extra = num_tokens % num_reqs
+        num_scheduled_tokens = np.full(num_reqs, base_tokens, dtype=np.int32)
+        if num_extra > 0:
+            num_scheduled_tokens[-num_extra:] += 1
+        prompt_lens = torch.from_numpy(num_scheduled_tokens)
+
+        pooling_params = PoolingParams(task=task)
+        pooling_params.verify(self.model_config)
+        self.model.pooler.get_pooling_updates(task).apply(pooling_params)
+        prompt_token_ids_cpu = None
+        if pooling_params.requires_token_ids:
+            prompt_token_ids_cpu = torch.zeros(
+                (num_reqs, int(prompt_lens.max())),
+                dtype=torch.int64,
+            )
+        pooling_metadata = PoolingMetadata(
+            prompt_lens=prompt_lens,
+            prompt_token_ids=None,
+            prompt_token_ids_cpu=prompt_token_ids_cpu,
+            pooling_params=[pooling_params] * num_reqs,
+            pooling_states=[PoolingStates() for _ in range(num_reqs)],
+        )
+        pooling_metadata.build_pooling_cursor(
+            num_scheduled_tokens,
+            seq_lens_cpu=prompt_lens,
+            device=hidden_states.device,
+        )
+        try:
+            return self.model.pooler(hidden_states, pooling_metadata)
+        except RuntimeError as e:
+            if "out of memory" not in str(e):
+                raise
+            raise RuntimeError(
+                "Out of memory while warming up the pooler "
+                f"({task=}) with {num_reqs} dummy requests. "
+                "Lower `max_num_seqs`."
+            ) from e
+
+    def dummy_pooler_run(self, hidden_states: torch.Tensor) -> None:
+        for task in sorted(self.supported_tasks):
+            self._dummy_pooler_run_task(hidden_states, task)
diff --git a/vllm/v1/worker/cpu/pp_utils.py b/vllm/v1/worker/cpu/pp_utils.py
new file mode 100644
index 000000000000..1dcc39d8085b
--- /dev/null
+++ b/vllm/v1/worker/cpu/pp_utils.py
@@ -0,0 +1,146 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Pipeline-parallel sampled-token synchronization for the CPU runner."""
+
+from collections import deque
+from dataclasses import dataclass
+
+import numpy as np
+import torch
+
+from vllm.distributed.parallel_state import get_pp_group
+from vllm.v1.worker.cpu.input_batch import InputBatch
+
+
+@dataclass
+class PendingRecv:
+    sampled_tokens: torch.Tensor  # [num_reqs, 1]
+    num_sampled: torch.Tensor  # [num_reqs]
+    num_rejected: torch.Tensor  # [num_reqs]
+    idx_mapping: torch.Tensor  # [num_reqs]
+    idx_mapping_np: np.ndarray  # [num_reqs]
+    # Records which rows need a deferred postprocess (bool).
+    need_sampled_mask: np.ndarray  # [num_reqs]
+    # Snapshot of slot generation counters at receive time, used to
+    # detect requests aborted since then.
+    gen_at_receive_np: np.ndarray  # [num_reqs]
+
+
+def compute_need_sampled_mask(input_batch: InputBatch) -> np.ndarray | None:
+    """Return a bool array of shape `[input_batch.num_reqs]` marking requests
+    with outputs that might be needed in a subsequent (decode) step.
+    Returns None if no sampled outputs are needed in the requests' next step."""
+    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).
+    need_sampled_mask = old_computed + input_batch.num_scheduled_tokens >= prefill_len
+    if input_batch.max_seq_len_np is not None:
+        # Also exclude final prefill chunks whose single sampled token reaches
+        # the request's length cap.
+        finished_prefill = prefill_len + 1 >= input_batch.max_seq_len_np
+        finished_prefill &= input_batch.is_prefilling_np
+        need_sampled_mask &= ~finished_prefill
+    return need_sampled_mask if need_sampled_mask.any() else None
+
+
+class PPHandler:
+    """Synchronize sampled tokens from the last pipeline rank on CPU."""
+
+    def __init__(self, max_num_reqs: int, device: torch.device):
+        pp_group = get_pp_group()
+        self.is_last_rank = pp_group.is_last_rank
+        self.last_rank = pp_group.last_rank
+        self.device = device
+
+        self.queue: deque[PendingRecv | None] = (
+            deque() if self.is_last_rank else deque([None] * pp_group.world_size)
+        )
+        self.req_idx_gen_np = np.zeros(max_num_reqs, dtype=np.int32)
+        self.broadcast_group = pp_group.make_sibling_device_group(
+            group_desc="pp_broadcast"
+        )
+
+    def on_req_idx_freed(self, req_idx: int) -> None:
+        self.req_idx_gen_np[req_idx] += 1
+
+    def get_prev_sampled_outputs(self) -> dict[str, torch.Tensor] | None:
+        if not self.queue:
+            return None
+
+        slot = self.queue.popleft()
+        self.queue.append(None)
+        if slot is None:
+            return None
+
+        freed = self.req_idx_gen_np[slot.idx_mapping_np] != slot.gen_at_receive_np
+        exclude_mask = freed | ~slot.need_sampled_mask
+        idx_mapping = slot.idx_mapping
+        if exclude_mask.any():
+            if exclude_mask.all():
+                return None
+            idx_mapping_np = np.where(exclude_mask, -1, slot.idx_mapping_np)
+            idx_mapping = torch.from_numpy(idx_mapping_np).to(device=self.device)
+
+        return {
+            "sampled_tokens": slot.sampled_tokens,
+            "num_sampled": slot.num_sampled,
+            "num_rejected": slot.num_rejected,
+            "idx_mapping": idx_mapping,
+        }
+
+    def receive(self, input_batch: InputBatch) -> bool:
+        """Receive this step's sampled-token data on a non-last rank."""
+        assert not self.is_last_rank
+        need_sampled_mask = compute_need_sampled_mask(input_batch)
+        if need_sampled_mask is None:
+            return False
+
+        # Snapshot the per-slot generation counter so a later free of any of
+        # these request indices is detectable at consume time.
+        gen_at_receive_np = self.req_idx_gen_np[input_batch.idx_mapping_np]
+
+        num_reqs = input_batch.num_reqs
+        sampled_tokens = torch.empty(num_reqs, 1, dtype=torch.int64, device=self.device)
+        combined = torch.empty(2, num_reqs, dtype=torch.int32, device=self.device)
+        torch.distributed.broadcast(
+            sampled_tokens, src=self.last_rank, group=self.broadcast_group
+        )
+        torch.distributed.broadcast(
+            combined, src=self.last_rank, group=self.broadcast_group
+        )
+        num_sampled, num_rejected = combined.unbind(dim=0)
+
+        self.queue[-1] = PendingRecv(
+            sampled_tokens=sampled_tokens,
+            num_sampled=num_sampled,
+            num_rejected=num_rejected,
+            idx_mapping=input_batch.idx_mapping,
+            idx_mapping_np=input_batch.idx_mapping_np,
+            need_sampled_mask=need_sampled_mask,
+            gen_at_receive_np=gen_at_receive_np,
+        )
+        return bool(need_sampled_mask.all())
+
+    def broadcast(
+        self,
+        sampled_token_ids: torch.Tensor,
+        num_sampled: torch.Tensor,
+        num_rejected: torch.Tensor,
+        input_batch: InputBatch,
+    ) -> None:
+        """Broadcast this step's sampled-token data from the last rank."""
+        assert self.is_last_rank
+        if compute_need_sampled_mask(input_batch) is None:
+            return
+
+        assert sampled_token_ids.dtype == torch.int64
+        assert sampled_token_ids.shape[-1] == 1
+        torch.distributed.broadcast(
+            sampled_token_ids.contiguous(),
+            src=self.last_rank,
+            group=self.broadcast_group,
+        )
+        combined = torch.stack((num_sampled, num_rejected), dim=0)
+        torch.distributed.broadcast(
+            combined, src=self.last_rank, group=self.broadcast_group
+        )
diff --git a/vllm/v1/worker/cpu/rope.py b/vllm/v1/worker/cpu/rope.py
new file mode 100644
index 000000000000..7ae405899f65
--- /dev/null
+++ b/vllm/v1/worker/cpu/rope.py
@@ -0,0 +1,126 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from typing import cast
+
+import torch
+import torch.nn as nn
+
+from vllm import _custom_ops as ops
+from vllm.config import ModelConfig
+from vllm.model_executor.models.interfaces import SupportsMRoPE
+
+
+class RopeState:
+    """State for multi-dimensional (M-RoPE) positions.
+
+    `num_dims` is the number of position channels the model consumes, which
+    the model config derives from its M-RoPE sections.
+
+    NOTE: `positions` is implemented with one additional dummy position on
+    purpose to make it non-contiguous so that it can work with torch compile.
+    See detailed explanation in
+    https://github.com/vllm-project/vllm/pull/12128#discussion_r1926431923
+
+    NOTE: For text-only inputs, each dimension has identical position IDs,
+    making M-RoPE functionally equivalent to 1D-RoPE.
+    See page 5 of https://arxiv.org/abs/2409.12191
+    """
+
+    def __init__(
+        self,
+        num_dims: int,
+        max_num_reqs: int,
+        max_num_tokens: int,
+        max_model_len: int,
+        device: torch.device,
+    ):
+        self.num_dims = num_dims
+        self.max_num_reqs = max_num_reqs
+        self.max_num_tokens = max_num_tokens
+        self.max_model_len = max_model_len
+        self.device = device
+
+        self.prefill_positions = torch.zeros(
+            (max_num_reqs * num_dims, max_model_len),
+            dtype=torch.int32,
+            device=device,
+        )
+        self.positions = torch.zeros(
+            (num_dims, max_num_tokens + 1), dtype=torch.int64, device=device
+        )
+        self.prefill_delta = torch.zeros(max_num_reqs, dtype=torch.int32, device=device)
+
+    def init_prefill_positions(
+        self,
+        req_idx: int,
+        model: nn.Module,
+        prefill_token_ids: list[int],
+        mm_features: list,
+    ) -> None:
+        mrope_model = cast(SupportsMRoPE, model)
+        prefill_positions, delta = mrope_model.get_mrope_input_positions(
+            prefill_token_ids, mm_features
+        )
+        self.prefill_delta[req_idx] = delta
+
+        for i in range(self.num_dims):
+            pos = prefill_positions[i]
+            self.prefill_positions[self.num_dims * req_idx + i, : pos.shape[0]].copy_(
+                pos
+            )
+
+    def get_positions(self, num_tokens: int) -> torch.Tensor:
+        return self.positions[:, :num_tokens]
+
+    def read_prefill_positions(self, req_idx: int, length: int) -> torch.Tensor:
+        """Return per-request prefill positions as [num_dims, length]."""
+        base = self.num_dims * req_idx
+        return self.prefill_positions[base : base + self.num_dims, :length]
+
+    def update_prefill_positions(
+        self, req_idx: int, positions: torch.Tensor, delta: int
+    ) -> None:
+        """Overwrite a request's prefill positions with recomputed values."""
+        base = self.num_dims * req_idx
+        length = positions.shape[1]
+        self.prefill_positions[base : base + self.num_dims, :length].copy_(positions)
+        self.prefill_delta[req_idx] = delta
+
+    def prepare_positions(
+        self,
+        idx_mapping: torch.Tensor,
+        query_start_loc: torch.Tensor,
+        prefill_lens: torch.Tensor,
+        num_computed_tokens: torch.Tensor,
+    ) -> None:
+        ops.prepare_mrope_positions_cpu(
+            self.positions,
+            self.prefill_positions,
+            self.prefill_delta,
+            idx_mapping,
+            query_start_loc,
+            prefill_lens,
+            num_computed_tokens,
+        )
+
+
+def get_rope_state(
+    model_config: ModelConfig,
+    model: nn.Module,
+    max_num_reqs: int,
+    max_num_tokens: int,
+    max_model_len: int,
+    device: torch.device,
+) -> RopeState | None:
+    """Create a RopeState if the model uses multi-dimensional RoPE."""
+    if not model_config.uses_mrope:
+        return None
+
+    assert isinstance(model, SupportsMRoPE)
+    return RopeState(
+        num_dims=model_config.mrope_num_dims,
+        max_num_reqs=max_num_reqs,
+        max_num_tokens=max_num_tokens,
+        max_model_len=max_model_len,
+        device=device,
+    )
diff --git a/vllm/v1/worker/cpu/sampler.py b/vllm/v1/worker/cpu/sampler.py
new file mode 100644
index 000000000000..95deaa9d2de7
--- /dev/null
+++ b/vllm/v1/worker/cpu/sampler.py
@@ -0,0 +1,997 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Literal, NamedTuple
+
+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.logger import init_logger
+from vllm.platforms import current_platform
+from vllm.sampling_params import SamplingParams
+from vllm.v1.outputs import LogprobsTensors, SamplingMaskLists
+from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p
+from vllm.v1.watermarking.watermarker import RandomSampler, Watermarker
+from vllm.v1.worker.cpu.gumbel import gumbel_sample
+from vllm.v1.worker.cpu.input_batch import InputBatch, get_num_sampled_and_rejected
+from vllm.v1.worker.cpu.logits_processors import (
+    CPUBadWords,
+    CPULogitBiasState,
+    CPUPenalties,
+    LogitsContext,
+    LogitsProcessor,
+    LogitsProcRequestState,
+)
+from vllm.v1.worker.cpu.logprob import LogprobTokenIdsState, compute_topk_scores
+from vllm.v1.worker.cpu.states import CPURequestState
+
+if TYPE_CHECKING:
+    from vllm.config.reasoning import ReasoningConfig
+
+logger = init_logger(__name__)
+
+_INT32_MAX = np.iinfo(np.int32).max
+_NP_INT64_MIN = np.iinfo(np.int64).min
+_NP_INT64_MAX = np.iinfo(np.int64).max
+NO_LOGPROBS = -1
+MAX_COMPACT_SUPPORT = 2048
+
+
+def get_num_nans(logits: torch.Tensor) -> torch.Tensor:
+    return torch.isnan(logits).sum(dim=1, dtype=torch.int32)
+
+
+class CPUThinkingBudgetState:
+    """Model Runner V2 state for per-request thinking token budgets."""
+
+    def __init__(
+        self,
+        req_states: CPURequestState,
+        reasoning_config: ReasoningConfig | None,
+    ) -> None:
+        self.req_states = req_states
+        self.max_num_reqs = req_states.max_num_reqs
+        self.device = req_states.device
+        start_ids = (
+            []
+            if reasoning_config is None
+            else reasoning_config.reasoning_start_token_ids or []
+        )
+        end_ids = (
+            []
+            if reasoning_config is None
+            else reasoning_config.reasoning_end_token_ids or []
+        )
+        natural_end_ids = (
+            []
+            if reasoning_config is None
+            else reasoning_config.natural_reasoning_end_token_ids or []
+        )
+        self.enabled = bool(start_ids and end_ids and natural_end_ids)
+        if not self.enabled:
+            return
+
+        self.thinking_token_budget = np.full(self.max_num_reqs, -1, dtype=np.int32)
+        self.use_thinking_budget = np.zeros(self.max_num_reqs, dtype=bool)
+        self.cached_last_start = np.full(self.max_num_reqs, -1, dtype=np.int32)
+        self.cached_last_end = np.full(self.max_num_reqs, -1, dtype=np.int32)
+        self.cached_scan_pos = np.zeros(self.max_num_reqs, dtype=np.int32)
+        self.reasoning_start_token_ids = tuple(start_ids)
+        self.natural_reasoning_end_token_ids = tuple(natural_end_ids)
+        self.reasoning_end_token_ids = tuple(end_ids)
+
+    def add_request(self, req_idx: int, sampling_params: SamplingParams) -> bool:
+        if not self.enabled:
+            return False
+        budget = sampling_params.thinking_token_budget
+        use_thinking_budget = budget is not None
+        self.use_thinking_budget[req_idx] = use_thinking_budget
+        if budget is None:
+            budget = -1
+        else:
+            budget = min(budget, _INT32_MAX)
+            self.cached_last_start[req_idx] = -1
+            self.cached_last_end[req_idx] = -1
+            self.cached_scan_pos[req_idx] = 0
+        self.thinking_token_budget[req_idx] = budget
+        return use_thinking_budget
+
+    @staticmethod
+    def _last_match(
+        tokens: list[int],
+        marker: tuple[int, ...],
+        offset: int,
+        current: int,
+    ) -> int:
+        marker_list = list(marker)
+        for pos in range(len(tokens) - len(marker) + 1):
+            if tokens[pos : pos + len(marker)] == marker_list:
+                current = offset + pos
+        return current
+
+    def _update_committed_cache(self, req_idx: int) -> int:
+        total_len = int(self.req_states.total_len[req_idx])
+        scan_pos = int(self.cached_scan_pos[req_idx])
+        if scan_pos > total_len:
+            scan_pos = 0
+            self.cached_last_start[req_idx] = -1
+            self.cached_last_end[req_idx] = -1
+
+        tokens = self.req_states.all_token_ids[req_idx, scan_pos:total_len].tolist()
+        self.cached_last_start[req_idx] = self._last_match(
+            tokens,
+            self.reasoning_start_token_ids,
+            scan_pos,
+            int(self.cached_last_start[req_idx]),
+        )
+        self.cached_last_end[req_idx] = self._last_match(
+            tokens,
+            self.natural_reasoning_end_token_ids,
+            scan_pos,
+            int(self.cached_last_end[req_idx]),
+        )
+        overlap = (
+            max(
+                len(self.reasoning_start_token_ids),
+                len(self.natural_reasoning_end_token_ids),
+            )
+            - 1
+        )
+        self.cached_scan_pos[req_idx] = max(0, total_len - overlap)
+        return total_len
+
+    def apply(self, logits: torch.Tensor, ctx: LogitsContext) -> None:
+        if not self.enabled or not np.any(self.use_thinking_budget[ctx.idx_mapping_np]):
+            return
+
+        total_lens = {
+            req_idx: self._update_committed_cache(int(req_idx))
+            for req_idx in ctx.idx_mapping_np
+            if self.use_thinking_budget[req_idx]
+        }
+        for row, req_idx in enumerate(ctx.expanded_idx_mapping.tolist()):
+            if req_idx < 0 or not self.use_thinking_budget[req_idx]:
+                continue
+
+            total_len = total_lens[req_idx]
+            local_pos = int(ctx.expanded_local_pos[row])
+            first_row = row - local_pos
+            overlap = (
+                max(
+                    len(self.reasoning_start_token_ids),
+                    len(self.natural_reasoning_end_token_ids),
+                    len(self.reasoning_end_token_ids),
+                )
+                - 1
+            )
+            tail_start = max(0, total_len - overlap)
+            tokens = self.req_states.all_token_ids[
+                req_idx, tail_start:total_len
+            ].tolist()
+            tokens.extend(ctx.input_ids[first_row + 1 : row + 1].tolist())
+
+            last_start = self._last_match(
+                tokens,
+                self.reasoning_start_token_ids,
+                tail_start,
+                int(self.cached_last_start[req_idx]),
+            )
+            last_end = self._last_match(
+                tokens,
+                self.natural_reasoning_end_token_ids,
+                tail_start,
+                int(self.cached_last_end[req_idx]),
+            )
+            if last_start < 0 or last_start <= last_end:
+                continue
+            effective_len = total_len + local_pos
+            if (
+                effective_len - last_start - len(self.reasoning_start_token_ids)
+                < self.thinking_token_budget[req_idx]
+            ):
+                continue
+
+            end_prefix_len = 0
+            for prefix_len in range(
+                1, min(len(tokens), len(self.reasoning_end_token_ids) - 1) + 1
+            ):
+                if tokens[-prefix_len:] == list(
+                    self.reasoning_end_token_ids[:prefix_len]
+                ):
+                    end_prefix_len = prefix_len
+            logits[row, self.reasoning_end_token_ids[end_prefix_len]] = 1.0e9
+
+
+class CPUTraceReplayState:
+    """Per-request state for inference trace-replay.
+
+    When a request carries ``SamplingParams.trace_decode_token_ids``, the
+    sampler overwrites the sampled token at each decode step with the
+    predetermined trace token, while real logprobs and ranks are still computed
+    from the unmodified logit distribution.
+    """
+
+    def __init__(self, req_states: CPURequestState):
+        self.max_num_reqs = req_states.max_num_reqs
+        self.device = req_states.device
+        self.req_states = req_states
+        self.trace_token_ids = torch.zeros(
+            (self.max_num_reqs, req_states.max_model_len),
+            dtype=torch.int32,
+            device=self.device,
+        )
+        self.trace_len = torch.zeros(
+            self.max_num_reqs, dtype=torch.int32, device=self.device
+        )
+
+    def add_request(self, req_idx: int, sampling_params: SamplingParams) -> None:
+        trace = sampling_params.trace_decode_token_ids
+        if trace is not None:
+            self.trace_len[req_idx] = len(trace)
+            self.trace_token_ids[req_idx, : len(trace)] = torch.tensor(
+                trace, dtype=torch.int32, device=self.trace_token_ids.device
+            )
+        else:
+            self.trace_len[req_idx] = 0
+
+    def apply_trace(self, sampled: torch.Tensor, idx_mapping: torch.Tensor) -> None:
+        req_indices = idx_mapping.long()
+        valid = req_indices >= 0
+        safe_indices = req_indices.clamp_min(0)
+        steps = (
+            self.req_states.total_len[safe_indices]
+            - self.req_states.prompt_len[safe_indices]
+        ).long()
+        trace_lens = self.trace_len[safe_indices].long()
+        active = valid & (steps >= 0) & (steps < trace_lens)
+        if torch.any(active):
+            sampled[active] = self.trace_token_ids[
+                safe_indices[active], steps[active]
+            ].to(sampled.dtype)
+
+
+class CPURandomSampler(RandomSampler):
+    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=self.is_drafting,
+            logits_cache=self.logits_cache,
+            logits_cache_col=self.logits_cache_col,
+            logits_cache_source=self.logits_cache_source,
+            use_fp64=self.use_fp64,
+        )
+
+
+@torch.compile(
+    dynamic=True,
+    fullgraph=True,
+    backend=current_platform.simple_compile_backend,
+)
+def _apply_min_p(
+    logits: torch.Tensor,
+    expanded_idx_mapping: torch.Tensor,
+    min_p_values: torch.Tensor,
+) -> None:
+    min_p = min_p_values[expanded_idx_mapping].unsqueeze(1)
+    threshold = logits.amax(dim=1, keepdim=True) + min_p.log()
+    logits.masked_fill_((min_p != 0.0) & (logits < threshold), float("-inf"))
+
+
+class CPUSamplingStates:
+    def __init__(self, max_num_reqs: int, vocab_size: int, device: torch.device):
+        self.max_num_reqs = max_num_reqs
+        self.vocab_size = vocab_size
+
+        self.temperature = torch.zeros(max_num_reqs, dtype=torch.float32, device=device)
+        self.temperature_np = self.temperature.numpy()
+        self.top_k = torch.full(
+            (max_num_reqs,), vocab_size, dtype=torch.int32, device=device
+        )
+        self.top_k_np = self.top_k.numpy()
+        self.top_p = torch.ones(max_num_reqs, dtype=torch.float32, device=device)
+        self.top_p_np = self.top_p.numpy()
+        self.min_p = torch.zeros(max_num_reqs, dtype=torch.float32, device=device)
+        self.min_p_np = self.min_p.numpy()
+        self.seeds = torch.zeros(max_num_reqs, dtype=torch.int64, device=device)
+        self.seeds_np = self.seeds.numpy()
+        # Tracks whether `seed` was set explicitly by the user, so callers
+        # can fall back from RNG paths that don't honor per-request seeds.
+        self.seeds_set = np.zeros(max_num_reqs, dtype=bool)
+
+        self.num_logprobs = np.empty(self.max_num_reqs, dtype=np.int32)
+        # -1 means no logprobs are requested.
+        self.num_logprobs.fill(NO_LOGPROBS)
+
+    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] = min_p
+
+        seed = sampling_params.seed
+        self.seeds_set[req_idx] = seed is not None
+        if seed is None:
+            seed = np.random.randint(_NP_INT64_MIN, _NP_INT64_MAX)
+        self.seeds_np[req_idx] = seed
+
+        num_logprobs = sampling_params.logprobs
+        if num_logprobs is None:
+            num_logprobs = NO_LOGPROBS
+        elif num_logprobs == -1:
+            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_temperature(
+        self,
+        logits: torch.Tensor,
+        expanded_idx_mapping: torch.Tensor,
+        idx_mapping_np: np.ndarray,
+    ) -> None:
+        temp_np = self.temperature_np[idx_mapping_np]
+        if np.all((temp_np == 0.0) | (temp_np == 1.0)):
+            # No request requires temperature.
+            return
+
+        temperature = self.temperature[expanded_idx_mapping].unsqueeze(1)
+        logits.div_(torch.where(temperature == 0.0, 1.0, temperature))
+
+    def apply_min_p(
+        self,
+        logits: torch.Tensor,
+        expanded_idx_mapping: torch.Tensor,
+        idx_mapping_np: np.ndarray,
+    ) -> None:
+        if np.all(self.min_p_np[idx_mapping_np] == 0.0):
+            # No request uses min_p.
+            return
+
+        _apply_min_p(logits, expanded_idx_mapping, self.min_p)
+
+    def get_top_k_top_p(
+        self,
+        expanded_idx_mapping: torch.Tensor,
+        idx_mapping_np: np.ndarray,
+        vocab_size: int,
+    ) -> tuple[torch.Tensor | None, torch.Tensor | None, bool]:
+        """Returns (top_k, top_p, has_topp_only_rows)."""
+        top_k_np = self.top_k_np[idx_mapping_np]
+        top_p_np = self.top_p_np[idx_mapping_np]
+        do_top_k = np.any(top_k_np != self.vocab_size)
+        do_top_p = np.any(top_p_np != 1.0)
+        top_k = self.top_k[expanded_idx_mapping] if do_top_k else None
+        top_p = self.top_p[expanded_idx_mapping] if do_top_p else None
+        has_topp_only_rows = bool(
+            do_top_p and np.any((top_k_np >= vocab_size) & (top_p_np < 1.0))
+        )
+        return top_k, top_p, has_topp_only_rows
+
+    def apply_top_k_top_p(
+        self,
+        logits: torch.Tensor,
+        expanded_idx_mapping: torch.Tensor,
+        idx_mapping_np: np.ndarray,
+    ) -> torch.Tensor:
+        top_k, top_p, has_topp_only_rows = self.get_top_k_top_p(
+            expanded_idx_mapping, idx_mapping_np, logits.shape[-1]
+        )
+        if top_k is None and top_p is None:
+            return logits
+        return apply_top_k_top_p(
+            logits, top_k, top_p, has_topp_only_rows=has_topp_only_rows
+        )
+
+    def any_greedy(self, idx_mapping_np: np.ndarray) -> bool:
+        return bool(np.any(self.temperature_np[idx_mapping_np] == 0.0))
+
+    def any_explicit_seed(self, idx_mapping_np: np.ndarray) -> bool:
+        return bool(np.any(self.seeds_set[idx_mapping_np]))
+
+    def max_num_logprobs(self, idx_mapping_np: np.ndarray) -> int:
+        return int(np.max(self.num_logprobs[idx_mapping_np]))
+
+
+@dataclass
+class SamplerOutput:
+    sampled_token_ids: torch.Tensor
+    logprobs_tensors: LogprobsTensors | None
+    num_nans: torch.Tensor | None
+    num_sampled: torch.Tensor
+    num_rejected: torch.Tensor
+    sampling_mask_tensors: CPUSamplingMaskTensors | None = None
+
+
+class CPUSamplingMaskTensors(NamedTuple):
+    token_ids: torch.Tensor
+    packed_mask: torch.Tensor
+    counts: torch.Tensor
+    vocab_size: int
+    rows_per_request: int = 1
+
+    @classmethod
+    def from_logits(
+        cls,
+        logits: torch.Tensor,
+        cu_num_logits: torch.Tensor,
+        num_sampled_tokens: torch.Tensor,
+        max_num_kept: int,
+        rows_per_request: int = 1,
+    ) -> CPUSamplingMaskTensors:
+        num_reqs, vocab_size = num_sampled_tokens.shape[0], logits.shape[1]
+        max_num_kept = min(max_num_kept, vocab_size, MAX_COMPACT_SUPPORT)
+        num_output_rows = num_reqs * rows_per_request
+        token_ids = torch.empty(
+            (num_output_rows, max_num_kept),
+            dtype=torch.int32,
+            device=logits.device,
+        )
+        packed_mask = torch.zeros(
+            (num_output_rows, (vocab_size + 7) // 8),
+            dtype=torch.uint8,
+            device=logits.device,
+        )
+        counts = torch.zeros(num_output_rows, dtype=torch.int32, device=logits.device)
+        bit_weights = 1 << torch.arange(8, dtype=torch.int16, device=logits.device)
+
+        for output_row in range(num_output_rows):
+            req_idx, slot_idx = divmod(output_row, rows_per_request)
+            if slot_idx >= int(num_sampled_tokens[req_idx]):
+                continue
+            source_row = int(cu_num_logits[req_idx]) + slot_idx
+            support = torch.isfinite(logits[source_row])
+            kept_ids = torch.nonzero(support).flatten()
+            counts[output_row] = kept_ids.numel()
+            num_compact = min(kept_ids.numel(), max_num_kept)
+            token_ids[output_row, :num_compact] = kept_ids[:num_compact]
+
+            padded = torch.zeros(
+                packed_mask.shape[1] * 8,
+                dtype=torch.bool,
+                device=logits.device,
+            )
+            padded[:vocab_size] = support
+            packed_mask[output_row] = (
+                (padded.view(-1, 8).to(torch.int16) * bit_weights)
+                .sum(dim=1)
+                .to(torch.uint8)
+            )
+
+        return cls(
+            token_ids,
+            packed_mask,
+            counts,
+            vocab_size,
+            rows_per_request,
+        )
+
+    @classmethod
+    def cat(cls, chunks: Sequence[CPUSamplingMaskTensors]) -> CPUSamplingMaskTensors:
+        first = chunks[0]
+        assert len({(c.vocab_size, c.rows_per_request) for c in chunks}) == 1
+        return cls(
+            torch.cat([chunk.token_ids for chunk in chunks]),
+            torch.cat([chunk.packed_mask for chunk in chunks]),
+            torch.cat([chunk.counts for chunk in chunks]),
+            first.vocab_size,
+            first.rows_per_request,
+        )
+
+    def to_cpu_nonblocking(self) -> CPUSamplingMaskTensors:
+        return self
+
+    def tolists(
+        self, num_sampled_tokens: np.ndarray | None = None
+    ) -> SamplingMaskLists:
+        counts = self.counts.numpy()
+        token_ids = self.token_ids.numpy()
+        packed_mask = self.packed_mask.numpy()
+        width = token_ids.shape[1]
+
+        cu_num_generated_tokens = None
+        if self.rows_per_request == 1:
+            sampled_rows = np.arange(len(counts))
+        else:
+            assert num_sampled_tokens is not None
+            active_slots = (
+                np.arange(self.rows_per_request)[None, :] < num_sampled_tokens[:, None]
+            )
+            sampled_rows = np.flatnonzero(active_slots)
+            cu_num_generated_tokens = np.cumsum(
+                np.concatenate(([0], num_sampled_tokens))
+            ).tolist()
+
+        def support(row: int) -> np.ndarray:
+            if counts[row] <= width:
+                return token_ids[row, : counts[row]]
+            bits = np.unpackbits(
+                packed_mask[row], count=self.vocab_size, bitorder="little"
+            )
+            return np.flatnonzero(bits).astype(np.int32, copy=False)
+
+        supports = [support(row) for row in sampled_rows]
+        offsets = np.zeros(len(supports) + 1, dtype=np.int64)
+        np.cumsum([len(s) for s in supports], out=offsets[1:])
+        flat = np.concatenate([np.empty(0, dtype=np.int32), *supports])
+        return SamplingMaskLists(flat, offsets, cu_num_generated_tokens)
+
+
+class CPUSampler:
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        max_num_reqs: int,
+        vocab_size: int,
+        device: torch.device,
+        req_states: CPURequestState,
+        logprobs_mode: LogprobsMode = "raw_logprobs",
+        use_fp64_gumbel: bool = False,
+        enable_trace_replay: bool = False,
+        return_sampling_mask: bool = False,
+    ):
+        self.logprobs_mode = logprobs_mode
+        self.compute_nans = envs.VLLM_COMPUTE_NANS_IN_LOGITS  # False by default.
+        self.use_fp64_gumbel = use_fp64_gumbel
+
+        self.req_states = req_states
+        self.sampling_states = CPUSamplingStates(max_num_reqs, vocab_size, device)
+
+        lp_req_state = LogitsProcRequestState.from_request_state(req_states)
+        self.penalties_state = CPUPenalties(vllm_config, lp_req_state)
+        logit_bias_state = CPULogitBiasState(vllm_config, lp_req_state)
+        bad_words_state = CPUBadWords(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,
+        ]
+
+        self.logprob_token_ids_state = LogprobTokenIdsState(max_num_reqs, device)
+        self.thinking_budget_state = CPUThinkingBudgetState(
+            req_states, vllm_config.reasoning_config
+        )
+        self.trace_replay_state = (
+            CPUTraceReplayState(req_states) if enable_trace_replay else None
+        )
+        self.needs_logits_processing = np.zeros(max_num_reqs, dtype=bool)
+        self.return_sampling_mask = return_sampling_mask
+
+    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
+
+        self.logprob_token_ids_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)
+
+    def get_logprobs_dims(
+        self, idx_mapping_np: np.ndarray, include_token_ids: bool = True
+    ) -> tuple[int, int] | None:
+        """(num_logprobs, max_per_req_token_ids) for the given requests, or
+        None when none of them want logprobs."""
+        max_num_logprobs = self.sampling_states.max_num_logprobs(idx_mapping_np)
+        max_token_ids = (
+            self.logprob_token_ids_state.max_num_token_ids(idx_mapping_np)
+            if include_token_ids
+            else 0
+        )
+        if max_num_logprobs == NO_LOGPROBS and max_token_ids == 0:
+            return None
+        num_logprobs = max_num_logprobs if max_num_logprobs != NO_LOGPROBS else 0
+        return num_logprobs, max_token_ids
+
+    def __call__(self, logits: torch.Tensor, input_batch: InputBatch) -> SamplerOutput:
+        expanded_idx_mapping = input_batch.expanded_idx_mapping
+        idx_mapping = input_batch.idx_mapping
+        idx_mapping_np = input_batch.idx_mapping_np
+        cu_num_logits_np = input_batch.cu_num_logits_np
+        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.
+        num_nans = get_num_nans(logits) if self.compute_nans else None
+
+        logprobs_dims = self.get_logprobs_dims(idx_mapping_np)
+
+        sampled, processed_logits = self.sample(
+            logits,
+            expanded_idx_mapping,
+            idx_mapping,
+            idx_mapping_np,
+            pos,
+            input_ids,
+            expanded_local_pos,
+            seq_lens_upper_bound_np,
+            return_logprobs=logprobs_dims is not None,
+        )
+
+        if self.trace_replay_state is not None:
+            self.trace_replay_state.apply_trace(sampled, idx_mapping)
+
+        if logprobs_dims is not None:
+            num_logprobs, max_per_req_token_ids = logprobs_dims
+            if self.logprobs_mode in PROCESSED_LOGPROBS_MODES:
+                logits = processed_logits
+            expanded_logits = logits.shape[0] != idx_mapping_np.shape[0]
+            cu_num_logits = cu_num_logits_np.tolist() if expanded_logits else None
+            logprobs_tensors = compute_topk_scores(
+                logits,
+                num_logprobs,
+                sampled,
+                cu_num_logits,
+                num_per_req_token_ids=self.logprob_token_ids_state.num_token_ids,
+                per_req_token_ids=self.logprob_token_ids_state.token_ids,
+                expanded_idx_mapping=input_batch.expanded_idx_mapping,
+                max_per_req_token_ids=max_per_req_token_ids,
+                logits_mode=self.logprobs_mode in ("raw_logits", "processed_logits"),
+            )
+        else:
+            logprobs_tensors = None
+
+        # 1 sampled token per request, except chunked-prefill requests
+        # (seq_len < prefill_len) which aren't done prefilling and produce no
+        # output token. num_rejected is always 0 here (one logit per request).
+        num_sampled, num_rejected = get_num_sampled_and_rejected(
+            input_batch.seq_lens.new_ones(input_batch.num_reqs),
+            input_batch.seq_lens,
+            input_batch.cu_num_logits,
+            input_batch.idx_mapping,
+            self.req_states.prefill_len,
+        )
+
+        sampling_mask_tensors = None
+        if self.return_sampling_mask:
+            # Size by the validated top_k batch max; wider supports use the bitmask.
+            max_num_kept = int(np.max(self.sampling_states.top_k_np[idx_mapping_np]))
+            sampling_mask_tensors = CPUSamplingMaskTensors.from_logits(
+                processed_logits,
+                input_batch.cu_num_logits,
+                num_sampled,
+                max_num_kept,
+            )
+
+        sampler_output = SamplerOutput(
+            # The sampled tokens are expanded to 2D tensor with shape
+            # [num_requests, 1], where each row represents one generated
+            # token per request.
+            sampled_token_ids=sampled.view(-1, 1),
+            logprobs_tensors=logprobs_tensors,
+            num_nans=num_nans,
+            num_sampled=num_sampled,
+            num_rejected=num_rejected,
+            sampling_mask_tensors=sampling_mask_tensors,
+        )
+        return sampler_output
+
+    def apply_sampling_params(
+        self,
+        logits: torch.Tensor,
+        expanded_idx_mapping: torch.Tensor,
+        idx_mapping: torch.Tensor,
+        idx_mapping_np: np.ndarray,
+        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]):
+            return logits
+
+        # Copy logits to a new FP32 tensor.
+        logits = torch.empty_like(logits, dtype=torch.float32).copy_(logits)
+        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 logits processors (native + any custom).
+        for processor in self.logits_processors:
+            logits = processor.apply(logits, ctx)
+
+        # 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(
+            logits, expanded_idx_mapping, idx_mapping_np
+        )
+
+        # Apply min_p in place.
+        self.sampling_states.apply_min_p(logits, expanded_idx_mapping, idx_mapping_np)
+
+        if skip_top_k_top_p:
+            return logits
+        return self.sampling_states.apply_top_k_top_p(
+            logits, expanded_idx_mapping, idx_mapping_np
+        )
+
+    def sample(
+        self,
+        logits: torch.Tensor,
+        expanded_idx_mapping: torch.Tensor,
+        idx_mapping: torch.Tensor,
+        idx_mapping_np: np.ndarray,
+        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]:
+        del return_logprobs
+        processed_logits = self.apply_sampling_params(
+            logits,
+            expanded_idx_mapping,
+            idx_mapping,
+            idx_mapping_np,
+            pos,
+            input_ids,
+            expanded_local_pos,
+            seq_lens_upper_bound_np,
+            skip_top_k_top_p=True,
+        )
+        top_k, top_p, has_topp_only_rows = self.sampling_states.get_top_k_top_p(
+            expanded_idx_mapping, idx_mapping_np, logits.shape[-1]
+        )
+        return self._sample_random(
+            processed_logits,
+            expanded_idx_mapping,
+            idx_mapping_np,
+            pos,
+            top_k,
+            top_p,
+            has_topp_only_rows,
+            use_fused_sampler=False,
+        )
+
+    def _sample_random(
+        self,
+        processed_logits: torch.Tensor,
+        expanded_idx_mapping: torch.Tensor,
+        idx_mapping_np: np.ndarray,
+        pos: torch.Tensor,
+        top_k: torch.Tensor | None,
+        top_p: torch.Tensor | None,
+        has_topp_only_rows: bool,
+        use_fused_sampler: bool,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        assert not use_fused_sampler
+
+        # TODO (fadara01): Add C++ fused ops for top-p / top-k sampling.
+        processed_logits = apply_top_k_top_p(
+            processed_logits,
+            top_k,
+            top_p,
+            has_topp_only_rows=has_topp_only_rows,
+        )
+
+        sampled = gumbel_sample(
+            processed_logits,
+            expanded_idx_mapping,
+            self.sampling_states.temperature,
+            self.sampling_states.seeds,
+            pos,
+            apply_temperature=False,
+            is_drafting=False,
+            use_fp64=self.use_fp64_gumbel,
+        )
+        return sampled, processed_logits
+
+
+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:
+    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(0, 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
+
+
+class CPUWatermarkSampler(CPUSampler):
+    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 = torch.ones(
+            self.sampling_states.max_num_reqs,
+            dtype=torch.bool,
+            device=self.req_states.device,
+        )
+        self.watermarking_np = self.watermarking.numpy()
+
+    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(
+                "Watermarking is enabled, but greedy decoding "
+                "(temperature=0) cannot be watermarked. This request will use "
+                "ordinary greedy sampling."
+            )
+
+    def _sample_random(
+        self,
+        processed_logits: torch.Tensor,
+        expanded_idx_mapping: torch.Tensor,
+        idx_mapping_np: np.ndarray,
+        pos: torch.Tensor,
+        top_k: torch.Tensor | None,
+        top_p: torch.Tensor | None,
+        has_topp_only_rows: 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
+        )
+        if not np.any(enabled):
+            return super()._sample_random(
+                processed_logits,
+                expanded_idx_mapping,
+                idx_mapping_np,
+                pos,
+                top_k,
+                top_p,
+                has_topp_only_rows,
+                use_fused_sampler,
+            )
+
+        processed_logits = apply_top_k_top_p(
+            processed_logits,
+            top_k,
+            top_p,
+            has_topp_only_rows=has_topp_only_rows,
+        )
+        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
+            )
+
+        temperatures = self.sampling_states.temperature[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[expanded_idx_mapping] & (temperatures != 0)
+            if repeated_contexts is not None:
+                watermarking &= ~repeated_contexts
+            skip_mask = ~watermarking
+
+        random_sampler = CPURandomSampler(
+            expanded_idx_mapping=expanded_idx_mapping,
+            temperatures=self.sampling_states.temperature,
+            seeds=self.sampling_states.seeds,
+            positions=pos,
+            use_fp64=self.use_fp64_gumbel,
+        )
+        output = self.watermarker.sample(
+            processed_logits,
+            contexts,
+            random_sampler,
+            skip_mask=skip_mask,
+        )
+        sampled = output.token_ids
+        if self.sampling_states.any_greedy(idx_mapping_np):
+            sampled = torch.where(
+                temperatures == 0,
+                processed_logits.argmax(dim=-1),
+                sampled,
+            )
+        return sampled, output.logits
+
+    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,
+            expanded_idx_mapping,
+            self.req_states.prompt_len,
+            self.req_states.total_len,
+            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) -> torch.Tensor:
+        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[safe_req_indices].to(torch.int64)
+        prompt_lens = self.req_states.prompt_len[safe_req_indices].to(torch.int64)
+        offsets = torch.arange(
+            -context_width, 0, dtype=torch.int64, device=req_indices.device
+        )
+        positions = total_lens.unsqueeze(-1) + offsets
+        valid = (
+            valid_reqs.unsqueeze(-1)
+            & (positions >= prompt_lens.unsqueeze(-1))
+            & (positions < total_lens.unsqueeze(-1))
+        )
+        contexts = self.req_states.all_token_ids[
+            safe_req_indices.unsqueeze(-1),
+            positions.clamp(0, self.req_states.all_token_ids.shape[1] - 1),
+        ]
+        return torch.where(valid, contexts, -1)
diff --git a/vllm/v1/worker/cpu/shm.py b/vllm/v1/worker/cpu/shm.py
index 961654ddaee9..474e6e6263e8 100644
--- a/vllm/v1/worker/cpu/shm.py
+++ b/vllm/v1/worker/cpu/shm.py
@@ -11,6 +11,7 @@
 
 # Patch torch APIs
 import torch
+from vllm.triton_utils import HAS_TRITON
 
 
 def noop(*args: Any, **kwargs: Any) -> None:
@@ -91,22 +92,19 @@ def async_tensor_h2d(
 
 torch_utils.async_tensor_h2d = async_tensor_h2d
 
-# Patch model runner APIs
-import vllm.v1.worker.gpu.buffer_utils as gpu_buffer_utils
-import vllm.v1.worker.cpu.buffer_utils as cpu_buffer_utils
-
-gpu_buffer_utils.UvaBuffer = cpu_buffer_utils.UvaBuffer
-
-import vllm.v1.worker.gpu.sample.penalties as gpu_penalties
-import vllm.v1.worker.gpu.structured_outputs as gpu_structured_outputs
-
-# triton-cpu lowers the bitmask unpack in these kernels to IR quadratic in the
-# block size: ~70 s per cold compile at 8192, ~2 s at 1024, at no runtime cost.
-gpu_penalties.PENALTIES_BLOCK_SIZE = 1024
-gpu_structured_outputs.GRAMMAR_BITMASK_BLOCK_SIZE = 1024
-
-# Patch Triton
-from vllm.triton_utils import HAS_TRITON, tl
-
 if HAS_TRITON:
+    # Patch model runner APIs used by the Triton CPU runner.
+    import vllm.v1.worker.cpu.buffer_utils as cpu_buffer_utils
+    import vllm.v1.worker.gpu.buffer_utils as gpu_buffer_utils
+    import vllm.v1.worker.gpu.sample.penalties as gpu_penalties
+    import vllm.v1.worker.gpu.structured_outputs as gpu_structured_outputs
+    from vllm.triton_utils import tl
+
+    gpu_buffer_utils.UvaBuffer = cpu_buffer_utils.UvaBuffer
+
+    # triton-cpu lowers the bitmask unpack in these kernels to IR quadratic in
+    # the block size: ~70 s per cold compile at 8192, ~2 s at 1024, at no
+    # runtime cost.
+    gpu_penalties.PENALTIES_BLOCK_SIZE = 1024
+    gpu_structured_outputs.GRAMMAR_BITMASK_BLOCK_SIZE = 1024
     tl.debug_barrier = noop
diff --git a/vllm/v1/worker/cpu/standalone_model_runner.py b/vllm/v1/worker/cpu/standalone_model_runner.py
new file mode 100644
index 000000000000..a6926359b966
--- /dev/null
+++ b/vllm/v1/worker/cpu/standalone_model_runner.py
@@ -0,0 +1,1345 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import gc
+import time
+from collections.abc import Iterable
+from contextlib import AbstractContextManager
+from copy import deepcopy
+from typing import Any, NamedTuple
+
+import numpy as np
+import torch
+
+import vllm.envs as envs
+from vllm.compilation.wrapper import compile_model_with_stock_torch
+from vllm.config import VllmConfig
+from vllm.config.compilation import CompilationMode
+from vllm.distributed.parallel_state import get_pp_group
+from vllm.forward_context import BatchDescriptor, set_forward_context
+from vllm.logger import init_logger
+from vllm.model_executor.layers.mamba.ops.ssu_dispatch import (
+    initialize_mamba_ssu_backend,
+)
+from vllm.model_executor.model_loader import get_model_loader
+from vllm.model_executor.model_loader.reload.inplace import reload_weights
+from vllm.model_executor.models.interfaces import requires_raw_input_tokens
+from vllm.model_executor.offloader import (
+    create_offloader,
+    get_offloader,
+    set_offloader,
+)
+from vllm.model_executor.warmup.jit_warmup import JitWarmupRegistry
+from vllm.multimodal import MULTIMODAL_REGISTRY
+from vllm.multimodal.encoder_budget import MultiModalBudget
+from vllm.sequence import IntermediateTensors
+from vllm.tasks import SupportedTask
+from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib
+from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
+from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
+from vllm.v1.kv_cache_interface import (
+    KVCacheConfig,
+    KVCacheSpec,
+    SlidingWindowSpec,
+    UniformTypeKVCacheSpecs,
+)
+from vllm.v1.outputs import ECConnectorOutput, ModelRunnerOutput
+from vllm.v1.watermarking import create_watermarker
+from vllm.v1.worker.block_table import get_block_table_width
+from vllm.v1.worker.cpu.attn_utils import (
+    FastPrefillHelper,
+    build_slot_mappings_by_layer,
+    get_kv_cache_spec,
+    init_attn_backend,
+    init_kv_cache,
+)
+from vllm.v1.worker.cpu.block_table import CPUBlockTables
+from vllm.v1.worker.cpu.ec_connector import get_ec_connector
+from vllm.v1.worker.cpu.encoder_cache import EncoderCache
+from vllm.v1.worker.cpu.input_batch import (
+    BatchReqState,
+    InputBatch,
+    InputBuffers,
+    combine_sampled_and_draft_tokens,
+    post_update,
+    post_update_num_computed_tokens,
+    prepare_pos_seq_lens,
+    prepare_prefill_inputs,
+)
+from vllm.v1.worker.cpu.kv_connector import (
+    NO_OP_KV_CONNECTOR,
+    KVConnector,
+    get_kv_connector,
+)
+from vllm.v1.worker.cpu.logprob import CPUPromptLogprobsWorker
+from vllm.v1.worker.cpu.lora_utils import LoraState
+from vllm.v1.worker.cpu.model_states import init_model_state
+from vllm.v1.worker.cpu.pooling_runner import PoolingRunner
+from vllm.v1.worker.cpu.pp_utils import PPHandler
+from vllm.v1.worker.cpu.sampler import (
+    CPUSampler,
+    CPUWatermarkSampler,
+    SamplerOutput,
+)
+from vllm.v1.worker.cpu.states import CPURequestState
+from vllm.v1.worker.cpu.structured_outputs import CPUStructuredOutputsWorker
+from vllm.v1.worker.cpu.utils import CPUKVBlockZeroer, free_before_shutdown
+from vllm.v1.worker.dp_utils import coordinate_batch_across_dp
+from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin
+from vllm.v1.worker.utils import (
+    clear_layer_kv_caches,
+    copy_kv_cache_blocks_inplace,
+    get_uniform_decode_token_count,
+    raise_if_nan_logits,
+)
+
+logger = init_logger(__name__)
+
+
+class ExecuteModelState(NamedTuple):
+    input_batch: InputBatch
+    hidden_states: torch.Tensor | None
+    finished_req_ids: set[str]
+    ec_connector_output: ECConnectorOutput | None
+
+
+def sort_batch_req_ids(
+    num_tokens_per_req: dict[str, int], decode_query_len: int
+) -> list[str]:
+    key = lambda r: (
+        (num := num_tokens_per_req[r]) != decode_query_len,
+        num,
+    )
+    return sorted(num_tokens_per_req, key=key)
+
+
+class StandaloneCPUModelRunner(LoRAModelRunnerMixin):
+    """Standalone MRV2 implementation for CPUs without mature Triton support.
+
+    It follows the GPU MRV2 architecture with a smaller feature set. Batch
+    preparation, request bookkeeping, KV-cache metadata, and sampling use C++
+    and PyTorch operations instead of Triton.
+    """
+
+    def __init__(self, vllm_config: VllmConfig, device: torch.device):
+        self.vllm_config = vllm_config
+        self.model_config = vllm_config.model_config
+        self.cache_config = vllm_config.cache_config
+        self.compilation_config = vllm_config.compilation_config
+        self.lora_config = vllm_config.lora_config
+        self.load_config = vllm_config.load_config
+        self.parallel_config = vllm_config.parallel_config
+        self.scheduler_config = vllm_config.scheduler_config
+        self.jit_warmup_registry = JitWarmupRegistry(vllm_config)
+
+        self.device = device
+        self.model: torch.nn.Module
+        self.dtype = self.model_config.dtype
+        self.kv_cache_dtype = self.dtype
+        if self.cache_config.cache_dtype != "auto":
+            # Quantized KV cache.
+            self.kv_cache_dtype = STR_DTYPE_TO_TORCH_DTYPE[
+                self.cache_config.cache_dtype
+            ]
+
+        # Lazily initialized in _init_kv_zero_meta() when the KV cache needs
+        # zeroing (e.g. hybrid models with fp8 KV cache).
+        self.kv_block_zeroer: CPUKVBlockZeroer | None = None
+
+        self.vocab_size = self.model_config.get_vocab_size()
+        self.max_model_len = self.model_config.max_model_len
+        self.max_num_tokens = self.scheduler_config.max_num_batched_tokens
+        self.max_num_reqs = self.scheduler_config.max_num_seqs
+        self.is_encoder_decoder = self.model_config.is_encoder_decoder
+
+        # Pipeline parallelism.
+        self.use_pp = self.parallel_config.pipeline_parallel_size > 1
+        self.is_first_pp_rank = get_pp_group().is_first_rank
+        self.is_last_pp_rank = get_pp_group().is_last_rank
+
+        # PP broadcast/receive helper.
+        self.pp_handler: PPHandler | None = None
+
+        # Persistent buffer for intermediate tensors (non-first PP ranks).
+        self.intermediate_tensors: IntermediateTensors | None = None
+
+        # Multimodal
+        self.mm_registry = MULTIMODAL_REGISTRY
+        self.supports_mm_inputs = self.model_config.supports_multimodal_inputs
+        self.uses_inputs_embeds = self.supports_mm_inputs
+        self.encoder_cache: EncoderCache | None = None
+        if self.supports_mm_inputs and self.is_first_pp_rank:
+            self.encoder_cache = EncoderCache()
+        self.ec_connector = get_ec_connector(vllm_config, self.encoder_cache)
+
+        # General request states.
+        self.req_states = CPURequestState(
+            max_num_reqs=self.max_num_reqs,
+            max_model_len=self.max_model_len,
+            max_num_batched_tokens=self.max_num_tokens,
+            vocab_size=self.vocab_size,
+            device=self.device,
+        )
+
+        # Pooling models.
+        self.is_pooling_model = self.model_config.runner_type == "pooling"
+        self.pooling_runner: PoolingRunner | None = None
+
+        self.input_buffers = InputBuffers(
+            max_num_reqs=self.max_num_reqs,
+            max_num_tokens=self.max_num_tokens,
+            device=device,
+        )
+        self.fast_prefill: FastPrefillHelper | None = None
+        if self.use_pp:
+            self.pp_handler = PPHandler(
+                max_num_reqs=self.max_num_reqs,
+                device=device,
+            )
+
+        # Created in load_model() after the model state is available.
+        self.sampler: CPUSampler | None = None
+        self.prompt_logprobs_worker: CPUPromptLogprobsWorker | None = None
+        self.structured_outputs_worker: CPUStructuredOutputsWorker | None = None
+
+        # LoRA-related workers.
+        self.lora_state = LoraState(max_num_reqs=self.max_num_reqs)
+
+        # KV Connector if configured.
+        self.kv_connector: KVConnector = NO_OP_KV_CONNECTOR
+        self.execute_model_state: ExecuteModelState | None = None
+
+        set_offloader(create_offloader(self.vllm_config.offload_config))
+
+    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 get_supported_tasks(self) -> tuple[SupportedTask, ...]:
+        tasks: list[SupportedTask] = []
+        if self.model_config.runner_type == "generate":
+            tasks.extend(self.model_state.get_supported_generation_tasks())
+        if self.is_pooling_model:
+            tasks.extend(PoolingRunner.get_supported_tasks(self.model))
+        return tuple(tasks)
+
+    def load_model(
+        self,
+        load_dummy_weights: bool = False,
+        *args,
+        model: torch.nn.Module | None = None,
+        **kwargs,
+    ) -> None:
+        time_before_load = time.perf_counter()
+        if model is None and load_dummy_weights:
+            self.load_config.load_format = "dummy"
+
+        with DeviceMemoryProfiler() as m:
+            model_loader = get_model_loader(self.vllm_config.load_config)
+            logger.info_once("Loading model from scratch...")
+
+            # Capture warmup providers selected while constructing the model.
+            with self.jit_warmup_registry.activate():
+                self.model = model or model_loader.load_model(
+                    vllm_config=self.vllm_config,
+                    model_config=self.vllm_config.model_config,
+                )
+            if self.lora_config:
+                self.model = self.load_lora_model(
+                    self.model, self.vllm_config, self.device
+                )
+        time_after_load = time.perf_counter()
+
+        self.model_memory_usage = m.consumed_memory
+        logger.info(
+            "Model loading took %s GiB memory and %.6f seconds",
+            format_gib(m.consumed_memory),
+            time_after_load - time_before_load,
+        )
+
+        self.model_state = init_model_state(
+            self.vllm_config, self.model, self.encoder_cache, self.device
+        )
+        self.decode_query_len = self.model_state.num_new_sampled_tokens_per_step
+        if self.decode_query_len != 1:
+            raise NotImplementedError(
+                "CPU Model Runner V2 without Triton supports one sampled "
+                "token per decode step."
+            )
+
+        if self.is_last_pp_rank and not self.is_pooling_model:
+            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,
+                "req_states": self.req_states,
+                "logprobs_mode": self.model_config.logprobs_mode,
+                "use_fp64_gumbel": self.model_config.use_fp64_gumbel,
+                "enable_trace_replay": self.model_config.enable_trace_replay,
+                "return_sampling_mask": self.model_config.return_sampling_mask,
+            }
+            if self.vllm_config.watermark_config is None:
+                self.sampler = CPUSampler(**sampler_kwargs)
+            else:
+                wm_config = self.vllm_config.watermark_config
+                watermarker = create_watermarker(wm_config)
+                self.sampler = CPUWatermarkSampler(
+                    watermarker,
+                    deduplicate_contexts=wm_config.deduplicate_contexts,
+                    deduplicate_contexts_max_history=(
+                        wm_config.deduplicate_contexts_max_history
+                    ),
+                    **sampler_kwargs,
+                )
+            self.prompt_logprobs_worker = CPUPromptLogprobsWorker(
+                self.max_num_reqs,
+                self.device,
+                logprobs_mode=self.model_config.logprobs_mode,
+            )
+            self.structured_outputs_worker = CPUStructuredOutputsWorker(
+                max_num_logits=self.max_num_reqs * self.decode_query_len,
+                vocab_size=self.vocab_size,
+                device=self.device,
+                mask_stride=self.decode_query_len,
+            )
+
+        if self.is_pooling_model and self.is_last_pp_rank:
+            self.pooling_runner = PoolingRunner(self.model, self.vllm_config)
+
+        if not self.is_first_pp_rank:
+            self.intermediate_tensors = self.model.make_empty_intermediate_tensors(
+                batch_size=self.max_num_tokens,
+                dtype=self.model_config.dtype,
+                device=self.device,
+            )
+
+        get_offloader().post_init()
+
+        if self.compilation_config.mode == CompilationMode.STOCK_TORCH_COMPILE:
+            compile_model_with_stock_torch(self.model, self.vllm_config)
+
+    @property
+    def dp_size(self) -> int:
+        return self.parallel_config.data_parallel_size
+
+    @property
+    def dp_rank(self) -> int:
+        return self.parallel_config.data_parallel_rank
+
+    def get_model(self) -> torch.nn.Module:
+        return self.model
+
+    def get_draft_model(self) -> torch.nn.Module | None:
+        return None
+
+    def reload_weights(
+        self,
+        weights_iterator: Iterable[tuple[str, torch.Tensor]] | None = None,
+        weights_path: str | None = None,
+        is_checkpoint_format: bool = True,
+    ) -> None:
+        """Reload weights in place."""
+        vllm_config = self.vllm_config
+        model = self.get_model()
+        reload_weights(
+            vllm_config, model, weights_iterator, weights_path, is_checkpoint_format
+        )
+        self.reset_lora_state()
+        self.reset_encoder_cache()
+        self.reset_mm_cache()
+
+    def get_encoder_timing_stats(self) -> dict[str, dict[str, float | int]]:
+        encoder_runner = getattr(self.model_state, "encoder_runner", None)
+        return encoder_runner.get_encoder_timing_stats() if encoder_runner else {}
+
+    def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]:
+        return get_kv_cache_spec(self.vllm_config)
+
+    def initialize_kv_cache(
+        self,
+        kv_cache_config: KVCacheConfig,
+        is_profiling: bool = False,
+        kv_cache_allocation_context: AbstractContextManager | None = None,
+    ) -> None:
+        kv_cache_config = deepcopy(kv_cache_config)
+        self.kv_cache_config = kv_cache_config
+
+        block_table_max_model_len = self.max_model_len
+        if self.is_encoder_decoder:
+            block_table_max_model_len = max(
+                block_table_max_model_len,
+                self.scheduler_config.max_num_encoder_input_tokens,
+                getattr(self.model_config.hf_config, "max_source_positions", 0),
+            )
+
+        block_sizes = []
+        max_num_blocks_per_group = []
+        slot_mapping_enabled = []
+        sliding_window_reach = []
+        for kv_cache_group in kv_cache_config.kv_cache_groups:
+            spec = kv_cache_group.kv_cache_spec
+            block_sizes.append(spec.block_size)
+            layer_spec = (
+                spec.first_spec if isinstance(spec, UniformTypeKVCacheSpecs) else spec
+            )
+            slot_mapping_enabled.append(layer_spec.uses_slot_mapping)
+            sliding_window_reach.append(
+                layer_spec.num_retained_tokens
+                if isinstance(layer_spec, SlidingWindowSpec)
+                else 0
+            )
+            max_num_blocks = spec.max_num_blocks_per_req(
+                self.vllm_config, block_table_max_model_len
+            )
+            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)
+
+        with self.jit_warmup_registry.activate():
+            self.attn_groups, self.kernel_block_sizes = init_attn_backend(
+                self.kv_cache_config,
+                self.vllm_config,
+                self.device,
+            )
+
+        self.block_tables = CPUBlockTables(
+            block_sizes=block_sizes,
+            max_num_reqs=self.max_num_reqs,
+            max_num_batched_tokens=self.max_num_tokens,
+            max_num_blocks_per_group=max_num_blocks_per_group,
+            device=self.device,
+            slot_mapping_enabled=slot_mapping_enabled,
+            sliding_window_reach=sliding_window_reach,
+        )
+        initialize_mamba_ssu_backend(
+            self.vllm_config.mamba_config,
+            self.kv_cache_config,
+            use_replayssm=self.cache_config.use_replayssm,
+        )
+        if self.cache_config.kv_sharing_fast_prefill:
+            self.fast_prefill = FastPrefillHelper()
+
+        with self.jit_warmup_registry.activate():
+            kv_caches_dict = init_kv_cache(
+                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,
+                attn_groups=(group for groups in self.attn_groups for group in groups),
+            )
+        self.kv_caches = [
+            cache for cache in kv_caches_dict.values() if cache.device == self.device
+        ]
+        self.model_state.initialize_kv_cache(self.kv_cache_config, self.block_tables)
+        if is_profiling:
+            self.kv_connector = NO_OP_KV_CONNECTOR
+        else:
+            self.kv_connector = get_kv_connector(self.vllm_config, kv_caches_dict)
+
+    def _init_kv_zero_meta(self) -> None:
+        self.kv_block_zeroer = CPUKVBlockZeroer(
+            self.device,
+            attn_groups_iter=(group for groups in self.attn_groups for group in groups),
+            kernel_block_sizes=self.kernel_block_sizes,
+            static_forward_context=self.compilation_config.static_forward_context,
+            num_blocks=self.kv_cache_config.num_blocks,
+        )
+
+    @torch.inference_mode()
+    def _dummy_run(
+        self,
+        num_tokens: int,
+        *,
+        skip_attn: bool = False,
+        uniform_decode: bool = False,
+        is_profile: bool = False,
+        randomize_inputs: bool = False,
+    ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
+        if skip_attn and not is_profile:
+            raise ValueError(
+                "skip_attn must only be True for initial memory profiling."
+            )
+
+        # Create a dummy scheduler output.
+        num_reqs = min(num_tokens, self.max_num_reqs)
+        if uniform_decode:
+            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
+
+        assert num_tokens <= self.max_num_tokens, (
+            f"Dummy run with {num_tokens} tokens exceeds max_num_batched_tokens "
+            f"({self.max_num_tokens})"
+        )
+
+        # Distribute the remainder evenly so no dummy request exceeds
+        # ceil(num_tokens / num_reqs) <= max_model_len tokens.
+        num_tokens_per_request = [
+            num_tokens // num_reqs + (i >= num_reqs - num_tokens % num_reqs)
+            for i in range(num_reqs)
+        ]
+
+        assert sum(num_tokens_per_request) == num_tokens
+        num_scheduled_tokens = {
+            f"_dummy_req_{i}": n for i, n in enumerate(num_tokens_per_request)
+        }
+        dummy_scheduler_output = SchedulerOutput.make_empty()
+        dummy_scheduler_output.total_num_scheduled_tokens = num_tokens
+        dummy_scheduler_output.num_scheduled_tokens = num_scheduled_tokens
+
+        # Disable any use of KVConnector for dummy runs.
+        self.kv_connector.set_disabled(True)
+
+        # Get the intermediate tensors for the dummy run.
+        intermediate_tensors = None
+        if not self.is_first_pp_rank:
+            assert self.intermediate_tensors is not None
+            intermediate_tensors = self.intermediate_tensors[:num_tokens]
+
+        max_loras = self.lora_config.max_loras if self.lora_config is not None else 0
+        with self.maybe_dummy_run_with_lora(
+            self.lora_config,
+            num_scheduled_tokens=np.array(num_tokens_per_request, dtype=np.int32),
+            num_sampled_tokens=None,
+            remove_lora=True,
+            num_active_loras=max_loras,
+        ):
+            # Execute the model.
+            self.execute_model(
+                dummy_scheduler_output,
+                intermediate_tensors=intermediate_tensors,
+                dummy_run=True,
+                skip_attn_for_dummy_run=skip_attn,
+                is_profile=is_profile,
+                randomize_inputs=randomize_inputs,
+            )
+        self.kv_connector.set_disabled(False)
+
+        # Non-last PP ranks don't produce output for sampling.
+        if not self.is_last_pp_rank:
+            return None, None
+
+        assert self.execute_model_state is not None
+        input_batch = self.execute_model_state.input_batch
+        hidden_states = self.execute_model_state.hidden_states
+        self.execute_model_state = None
+
+        assert hidden_states is not None  # Last PP rank always has hidden_states
+        sample_hidden_states = hidden_states[input_batch.logits_indices]
+        return hidden_states, sample_hidden_states
+
+    @torch.inference_mode()
+    def _dummy_sampler_run(self, hidden_states: torch.Tensor) -> None:
+        num_reqs = hidden_states.shape[0]
+        logits = self.model.compute_logits(hidden_states)
+        input_batch = InputBatch.make_dummy(num_reqs, num_reqs, self.input_buffers)
+        assert self.sampler is not None
+        self.sampler(logits, input_batch)
+
+    @torch.inference_mode()
+    def _dummy_pooler_run(self, hidden_states: torch.Tensor) -> None:
+        assert self.pooling_runner is not None
+        self.pooling_runner.dummy_pooler_run(hidden_states)
+
+    @torch.inference_mode()
+    def profile_run(self, randomize_inputs: bool = False) -> None:
+        if self.supports_mm_inputs and self.is_first_pp_rank:
+            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 = mm_budget.get_dummy_encoder_profile_inputs()
+                self.model_state.encoder_runner.profile_encoder_cache(
+                    dummy_mm_inputs, mm_budget
+                )
+
+        hidden_states, sample_hidden_states = self._dummy_run(
+            self.max_num_tokens,
+            skip_attn=True,
+            is_profile=True,
+            randomize_inputs=randomize_inputs,
+        )
+
+        if self.is_last_pp_rank:
+            assert sample_hidden_states is not None
+            if self.pooling_runner is None:
+                self._dummy_sampler_run(sample_hidden_states)
+            else:
+                assert hidden_states is not None
+                self._dummy_pooler_run(hidden_states)
+
+        del hidden_states, sample_hidden_states
+        self.reset_encoder_cache()
+        gc.collect()
+
+    def warming_up_model(self) -> None:
+        logger.info("Warming up model for the compilation...")
+        self.profile_run()
+        logger.info("Warming up done.")
+
+    def reset_mm_cache(self) -> None:
+        if self.encoder_cache is not None:
+            self.encoder_cache.reset_mm_cache()
+
+    def reset_encoder_cache(self) -> None:
+        if self.encoder_cache is not None:
+            self.encoder_cache.reset_encoder_cache()
+        if self.pooling_runner is not None:
+            self.pooling_runner.clear()
+
+    def release_late_interaction_query_cache(self, query_keys: list[str]) -> None:
+        if self.pooling_runner is not None:
+            self.pooling_runner.release_late_interaction_query_cache(query_keys)
+
+    def _remove_request(self, req_id: str) -> bool:
+        self.model_state.remove_request(req_id)
+        req_idx = self.req_states.remove_request(req_id)
+        if req_idx is None:
+            return False
+        if self.pooling_runner is not None:
+            self.pooling_runner.remove_request(req_idx)
+        if self.pp_handler is not None:
+            self.pp_handler.on_req_idx_freed(req_idx)
+        if self.encoder_cache is not None:
+            self.encoder_cache.remove_request(req_id)
+        if self.prompt_logprobs_worker is not None:
+            self.prompt_logprobs_worker.remove_request(req_id)
+        self.lora_state.remove_request(req_id)
+        return True
+
+    def finish_requests(self, scheduler_output: SchedulerOutput) -> None:
+        finished_req_ids = scheduler_output.finished_req_ids
+        if self.pooling_runner is not None:
+            # Preempted docs keep their query-use reservation until rescheduled.
+            self.pooling_runner.on_requests_finished(finished_req_ids)
+        preempted_req_ids = scheduler_output.preempted_req_ids
+        if preempted_req_ids:
+            finished_req_ids = finished_req_ids.union(preempted_req_ids)
+        # Sorted so every TP rank frees request slots in the same order.
+        # Features like batch-sharded sampling derive rank request ownership
+        # from the slot index.
+        for req_id in sorted(finished_req_ids):
+            self._remove_request(req_id)
+
+    def free_states(self, scheduler_output: SchedulerOutput) -> None:
+        if self.encoder_cache is not None:
+            for mm_hash in scheduler_output.free_encoder_mm_hashes:
+                self.encoder_cache.free_encoder_cache(mm_hash)
+
+    def update_pp_decode_requests(self) -> None:
+        if self.pp_handler is not None:
+            outputs = self.pp_handler.get_prev_sampled_outputs()
+            if outputs is not None:
+                self.postprocess_sampled(**outputs)
+
+    def add_requests(self, scheduler_output: SchedulerOutput) -> None:
+        for new_req_data in scheduler_output.scheduled_new_reqs:
+            assert new_req_data.prefill_token_ids is not None
+            req_id = new_req_data.req_id
+
+            # Streaming input update: request already exists from a prior
+            # chunk. Remove old state so it can be cleanly re-added below
+            # with the updated prompt_token_ids and mm_features.
+            self._remove_request(req_id)
+
+            prompt_len = new_req_data.prompt_len
+            sampling_params = new_req_data.sampling_params
+            self.req_states.add_request(
+                req_id=req_id,
+                prompt_len=prompt_len,
+                all_token_ids=new_req_data.prefill_token_ids,
+                num_computed_tokens=new_req_data.num_computed_tokens,
+                max_tokens=sampling_params.max_tokens if sampling_params else 1,  # type: ignore[arg-type]
+            )
+            req_index = self.req_states.req_id_to_index[req_id]
+
+            if self.pooling_runner is not None:
+                assert new_req_data.pooling_params is not None
+                self.pooling_runner.add_request(
+                    req_id,
+                    req_index,
+                    new_req_data.pooling_params,
+                    new_req_data.prompt_token_ids,
+                )
+
+            if self.encoder_cache is not None:
+                self.encoder_cache.add_request(req_id, new_req_data.mm_features)
+
+            self.model_state.add_request(req_index, new_req_data)
+            self.block_tables.append_block_ids(
+                req_index, new_req_data.block_ids, overwrite=True
+            )
+            self.lora_state.add_request(req_id, req_index, new_req_data.lora_request)
+
+            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, 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
+                )
+
+    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
+        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:
+                self.block_tables.append_block_ids(
+                    req_index, req_new_block_ids, overwrite=False
+                )
+
+        # Update CPU num_computed_prefill_tokens.
+        np.minimum(
+            self.req_states.num_computed_tokens_np,
+            self.req_states.prefill_len_np,
+            out=self.req_states.num_computed_prefill_tokens,
+        )
+
+        # Zero memory for freshly allocated cache blocks to prevent stale
+        # NaN/data from corrupting attention or SSM computation.
+        if scheduler_output.new_block_ids_to_zero:
+            assert self.kv_block_zeroer is not None
+            self.kv_block_zeroer.zero_block_ids(scheduler_output.new_block_ids_to_zero)
+
+        # Apply copy-on-write block copies for partial prefix-cache hits, after
+        # zeroing new blocks and before the forward pass reads them.
+        if scheduler_output.kv_cache_block_copies:
+            copy_kv_cache_blocks_inplace(
+                self.kv_caches,
+                self.kv_cache_config.num_blocks,
+                scheduler_output.kv_cache_block_copies,
+            )
+
+    def gather_batch_req_state(
+        self, scheduler_output: SchedulerOutput, dummy_run: bool
+    ) -> tuple[BatchReqState | None, int | None]:
+        """Gather CPU request state for the scheduled batch, in batch order.
+        Returns (batch_state, uniform_decode_token_count)
+        """
+        num_tokens_per_req = scheduler_output.num_scheduled_tokens
+        num_reqs = len(num_tokens_per_req)
+        num_toks = scheduler_output.total_num_scheduled_tokens
+        max_query_len = max(scheduler_output.num_scheduled_tokens.values())
+
+        if dummy_run:
+            # Dummy batches are uniform by construction.
+            return None, get_uniform_decode_token_count(
+                num_reqs, num_toks, max_query_len, decode_graph_eligible=True
+            )
+
+        # batch_idx -> req_id
+        req_ids = sort_batch_req_ids(num_tokens_per_req, self.decode_query_len)
+
+        numtoks_iter = map(num_tokens_per_req.__getitem__, req_ids)
+        num_scheduled_tokens = np.fromiter(numtoks_iter, dtype=np.int32, count=num_reqs)
+
+        idx_mapping_iter = map(self.req_states.req_id_to_index.__getitem__, req_ids)
+        idx_mapping_np = np.fromiter(idx_mapping_iter, dtype=np.intp, count=num_reqs)
+        prefill_len_np = self.req_states.prefill_len_np[idx_mapping_np]
+        num_computed_prefill_tokens_np = self.req_states.num_computed_prefill_tokens[
+            idx_mapping_np
+        ]
+        is_prefilling_np = num_computed_prefill_tokens_np < prefill_len_np
+        has_prefill = bool(is_prefilling_np.any())
+        max_seq_len_np = None
+        if self.use_pp and has_prefill:
+            # consumed by the PP `compute_need_sampled_mask` for batches with prefill.
+            max_seq_len_np = self.req_states.max_seq_len[idx_mapping_np]
+
+        prefill_runs_as_decode_np = None
+        decode_graph_eligible = not has_prefill
+        if has_prefill:
+            # One new prompt token over existing context runs as a decode.
+            num_new_tokens = num_scheduled_tokens
+            one_token_ext = (num_new_tokens == 1) & (num_computed_prefill_tokens_np > 0)
+            prefill_runs_as_decode_np = is_prefilling_np & one_token_ext
+            decode_graph_eligible = bool(
+                (prefill_runs_as_decode_np | ~is_prefilling_np).all()
+            )
+
+        batch_state = BatchReqState(
+            req_ids=req_ids,
+            num_scheduled_tokens=num_scheduled_tokens,
+            num_tokens=num_toks,
+            idx_mapping_np=idx_mapping_np,
+            prefill_len_np=prefill_len_np,
+            num_computed_prefill_tokens_np=num_computed_prefill_tokens_np,
+            is_prefilling_np=is_prefilling_np,
+            has_prefill=has_prefill,
+            prefill_runs_as_decode_np=prefill_runs_as_decode_np,
+            decode_graph_eligible=decode_graph_eligible,
+            max_seq_len_np=max_seq_len_np,
+        )
+        return batch_state, get_uniform_decode_token_count(
+            num_reqs, num_toks, max_query_len, decode_graph_eligible
+        )
+
+    def prepare_inputs(
+        self,
+        scheduler_output: SchedulerOutput,
+        batch_req_state: BatchReqState,
+    ) -> InputBatch:
+        num_tokens = batch_req_state.num_tokens
+        assert num_tokens > 0
+        req_ids = batch_req_state.req_ids
+        num_reqs = len(req_ids)
+
+        is_padding = self.input_buffers.is_padding[:num_tokens]
+        if envs.VLLM_MOE_SKIP_PADDING:
+            is_padding.fill_(False)
+
+        num_scheduled_tokens = batch_req_state.num_scheduled_tokens
+        idx_mapping_np = batch_req_state.idx_mapping_np
+        idx_mapping = torch.from_numpy(idx_mapping_np)
+        expanded_idx_mapping = idx_mapping
+        expanded_local_pos = torch.zeros(
+            num_reqs, dtype=torch.int32, device=self.device
+        )
+
+        cu_num_logits = torch.arange(
+            num_reqs + 1, dtype=torch.int32, device=self.device
+        )
+        cu_num_logits_np = cu_num_logits.numpy()
+
+        query_start_loc = self.input_buffers.query_start_loc[: num_reqs + 1]
+        query_start_loc_np = query_start_loc.numpy()
+        query_start_loc_np[0] = 0
+        np.cumsum(num_scheduled_tokens, out=query_start_loc_np[1:])
+
+        if batch_req_state.has_prefill:
+            prepare_prefill_inputs(
+                self.input_buffers.input_ids,
+                self.req_states.next_prefill_tokens,
+                idx_mapping,
+                query_start_loc,
+                self.req_states.all_token_ids,
+                self.req_states.prefill_len,
+                self.req_states.num_computed_tokens,
+            )
+
+        prepare_pos_seq_lens(
+            idx_mapping,
+            query_start_loc,
+            self.req_states.num_computed_tokens,
+            self.input_buffers.positions,
+            self.input_buffers.seq_lens,
+        )
+        seq_lens = self.input_buffers.seq_lens[:num_reqs]
+        logits_indices = combine_sampled_and_draft_tokens(
+            self.input_buffers.input_ids,
+            idx_mapping,
+            self.req_states.last_sampled_tokens,
+            query_start_loc,
+            seq_lens,
+            self.req_states.prefill_len,
+            self.req_states.last_sampled_tokens[:, :0],
+            cu_num_logits,
+            num_reqs,
+            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,
+                cu_num_logits_np,
+                batch_req_state.has_prefill,
+            )
+
+        num_computed_tokens_np = self.req_states.num_computed_tokens_np[idx_mapping_np]
+        prompt_lens = None
+        if self.model_config.rswa_window is not None:
+            prompt_lens = self.req_states.prompt_len[idx_mapping]
+
+        return InputBatch(
+            req_ids=req_ids,
+            num_reqs=num_reqs,
+            num_reqs_after_padding=num_reqs,
+            idx_mapping=idx_mapping,
+            idx_mapping_np=idx_mapping_np,
+            expanded_idx_mapping=expanded_idx_mapping,
+            expanded_local_pos=expanded_local_pos,
+            num_scheduled_tokens=num_scheduled_tokens,
+            num_tokens=num_tokens,
+            num_tokens_after_padding=num_tokens,
+            query_start_loc=query_start_loc,
+            query_start_loc_np=query_start_loc_np,
+            seq_lens=seq_lens,
+            seq_lens_cpu_upper_bound=seq_lens,
+            num_computed_tokens_np=num_computed_tokens_np,
+            prefill_len_np=batch_req_state.prefill_len_np,
+            num_computed_prefill_tokens_np=(
+                batch_req_state.num_computed_prefill_tokens_np
+            ),
+            is_prefilling_np=batch_req_state.is_prefilling_np,
+            has_prefill=batch_req_state.has_prefill,
+            prefill_runs_as_decode_np=batch_req_state.prefill_runs_as_decode_np,
+            decode_graph_eligible=batch_req_state.decode_graph_eligible,
+            max_seq_len_np=batch_req_state.max_seq_len_np,
+            input_ids=self.input_buffers.input_ids[:num_tokens],
+            positions=self.input_buffers.positions[:num_tokens],
+            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,
+        )
+
+    def prepare_attn(
+        self, input_batch: InputBatch
+    ) -> tuple[tuple[torch.Tensor, ...], torch.Tensor]:
+        block_tables = self.block_tables.gather_block_tables(
+            input_batch.idx_mapping,
+            num_reqs_padded=input_batch.num_reqs_after_padding,
+            num_computed_tokens=self.req_states.num_computed_tokens,
+        )
+        slot_mappings = self.block_tables.compute_slot_mappings(
+            input_batch.idx_mapping,
+            input_batch.query_start_loc,
+            input_batch.positions,
+            num_tokens_padded=input_batch.num_tokens_after_padding,
+        )
+        return block_tables, slot_mappings
+
+    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)
+        if valid_state_slots:
+            state_slots = torch.arange(
+                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 = self.block_tables.get_dummy_slot_mappings(
+            input_batch.num_tokens
+        )
+        return block_tables, slot_mappings
+
+    def sample(
+        self,
+        hidden_states: torch.Tensor,
+        input_batch: InputBatch,
+        grammar_output: GrammarOutput | None,
+    ) -> tuple[SamplerOutput, torch.Tensor, torch.Tensor]:
+        sample_hidden_states = hidden_states[input_batch.logits_indices]
+        logits = self.model.compute_logits(sample_hidden_states)
+
+        if grammar_output is not None and logits.shape[0] > 0:
+            assert self.structured_outputs_worker is not None
+            self.structured_outputs_worker.apply_grammar_bitmask(
+                logits,
+                input_batch,
+                grammar_output.structured_output_request_ids,
+                grammar_output.grammar_bitmask,
+            )
+
+        assert self.sampler is not None
+        sampler_output = self.sampler(logits, input_batch)
+        return sampler_output, sampler_output.num_sampled, sampler_output.num_rejected
+
+    def postprocess_sampled(
+        self,
+        idx_mapping: torch.Tensor,
+        sampled_tokens: torch.Tensor,
+        num_sampled: torch.Tensor,
+        num_rejected: torch.Tensor,
+        query_start_loc: torch.Tensor | None = None,
+    ) -> None:
+        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
+
+        post_update(
+            idx_mapping,
+            self.req_states.num_computed_tokens,
+            self.req_states.last_sampled_tokens,
+            output_bin_counts,
+            sampled_tokens,
+            num_sampled,
+            num_rejected,
+            query_start_loc,
+            self.req_states.all_token_ids,
+            self.req_states.total_len,
+        )
+
+    def postprocess_num_computed_tokens(self, input_batch: InputBatch) -> None:
+        post_update_num_computed_tokens(
+            input_batch.idx_mapping,
+            self.req_states.num_computed_tokens,
+            input_batch.query_start_loc,
+        )
+
+    def _merge_ec_connector_no_forward(
+        self, scheduler_output: SchedulerOutput, output: ModelRunnerOutput
+    ) -> ModelRunnerOutput:
+        return ModelRunnerOutput.with_ec_conn_output(
+            output,
+            self.ec_connector.no_forward(scheduler_output).ec_connector_output,
+        )
+
+    @torch.inference_mode()
+    def execute_model(
+        self,
+        scheduler_output: SchedulerOutput,
+        intermediate_tensors: IntermediateTensors | None = None,
+        dummy_run: bool = False,
+        skip_attn_for_dummy_run: bool = False,
+        is_profile: bool = False,
+        valid_dummy_state_slots: bool = False,
+        randomize_inputs: bool = False,
+    ) -> ModelRunnerOutput | IntermediateTensors | None:
+        if not dummy_run:
+            self.update_pp_decode_requests()
+            self.finish_requests(scheduler_output)
+            self.free_states(scheduler_output)
+            self.add_requests(scheduler_output)
+            self.update_requests(scheduler_output)
+            if self.encoder_cache is not None:
+                # add_requests initializes RoPE positions from CPU metadata first.
+                for req_id, updates in scheduler_output.encoder_input_updates.items():
+                    self.encoder_cache.update_encoder_inputs(req_id, updates)
+
+            if scheduler_output.total_num_scheduled_tokens == 0:
+                empty_output = self.kv_connector.no_forward(scheduler_output)
+                return self._merge_ec_connector_no_forward(
+                    scheduler_output, empty_output
+                )
+
+        num_reqs = len(scheduler_output.num_scheduled_tokens)
+        num_tokens = scheduler_output.total_num_scheduled_tokens
+        _, num_tokens_across_dp, _ = coordinate_batch_across_dp(
+            num_tokens_unpadded=num_tokens,
+            num_tokens_padded=num_tokens,
+            allow_microbatching=False,
+            parallel_config=self.parallel_config,
+            cudagraph_mode=0,
+        )
+        num_active_loras = 0
+        if dummy_run:
+            input_batch = InputBatch.make_dummy(
+                num_reqs,
+                num_tokens,
+                self.input_buffers,
+                is_padding=not is_profile,
+            )
+            if randomize_inputs:
+                input_batch.input_ids.random_(0, self.vocab_size)
+            if skip_attn_for_dummy_run:
+                block_tables = None
+                slot_mappings = None
+            else:
+                block_tables, slot_mappings = self.prepare_dummy_attn(
+                    input_batch, valid_dummy_state_slots
+                )
+            if self.lora_config is not None:
+                num_active_loras = self.lora_config.max_loras
+        else:
+            batch_req_state, _ = self.gather_batch_req_state(
+                scheduler_output, dummy_run=False
+            )
+            assert batch_req_state is not None
+            input_batch = self.prepare_inputs(scheduler_output, batch_req_state)
+            block_tables, slot_mappings = self.prepare_attn(input_batch)
+            self.model_state.preprocess_state(
+                input_batch,
+                block_tables,
+                self.kv_cache_config,
+                self.req_states.num_computed_tokens,
+            )
+
+            if self.lora_config:
+                lora_inputs = self.lora_state.make_lora_inputs(
+                    input_batch.req_ids,
+                    input_batch.idx_mapping_np,
+                    input_batch.num_scheduled_tokens,
+                    np.diff(input_batch.cu_num_logits_np),
+                )
+                num_active_loras = len(lora_inputs[2])
+                self._set_active_loras(*lora_inputs)
+
+        attn_metadata = None
+        slot_mappings_by_layer = None
+        if not (dummy_run and skip_attn_for_dummy_run):
+            assert block_tables is not None and slot_mappings is not None
+            slot_mappings_by_layer = build_slot_mappings_by_layer(
+                slot_mappings, self.kv_cache_config
+            )
+            attn_metadata = self.model_state.prepare_attn(
+                input_batch,
+                block_tables,
+                slot_mappings,
+                self.attn_groups,
+                self.kv_cache_config,
+            )
+
+        input_ids = input_batch.input_ids
+        inputs_embeds = None
+        ec_connector_output = None
+        if self.uses_inputs_embeds and self.is_first_pp_rank:
+            if dummy_run:
+                inputs_embeds = self.model_state.dummy_inputs_embeds(
+                    input_batch.num_tokens_after_padding
+                )
+            else:
+                scheduled_encoder_inputs = scheduler_output.scheduled_encoder_inputs
+                with self.ec_connector.maybe_get_output(
+                    scheduler_output
+                ) as ec_connector_output:
+                    inputs_embeds = self.model_state.prepare_inputs_embeds(
+                        scheduled_encoder_inputs, input_batch, self.req_states
+                    )
+            if inputs_embeds is not None and not requires_raw_input_tokens(self.model):
+                input_ids = None
+
+        model_inputs: dict[str, Any] = {
+            "input_ids": input_ids,
+            "positions": input_batch.positions,
+            "inputs_embeds": inputs_embeds,
+            "intermediate_tensors": None,
+            **self.model_state.prepare_inputs(input_batch, self.req_states),
+        }
+        if not self.is_first_pp_rank:
+            if not requires_raw_input_tokens(self.model):
+                model_inputs["input_ids"] = None
+            model_inputs["inputs_embeds"] = None
+
+            assert intermediate_tensors is not None
+            assert self.intermediate_tensors is not None
+            n = input_batch.num_tokens_after_padding
+            new_tensors = {
+                k: v[:n]
+                if dummy_run
+                else v[:n].copy_(intermediate_tensors.tensors[k][:n])
+                for k, v in self.intermediate_tensors.tensors.items()
+            }
+            model_inputs["intermediate_tensors"] = IntermediateTensors(new_tensors)
+            del intermediate_tensors
+
+        skip_compiled = bool(
+            self.is_encoder_decoder and scheduler_output.scheduled_encoder_inputs
+        )
+        batch_descriptor = BatchDescriptor(
+            num_tokens=input_batch.num_tokens_after_padding,
+            has_lora=self.lora_config is not None,
+            num_active_loras=num_active_loras,
+        )
+        with set_forward_context(
+            attn_metadata,
+            self.vllm_config,
+            num_tokens=input_batch.num_tokens_after_padding,
+            num_tokens_across_dp=num_tokens_across_dp,
+            batch_descriptor=batch_descriptor,
+            slot_mapping=slot_mappings_by_layer,
+            skip_compiled=skip_compiled,
+            is_padding=input_batch.is_padding,
+        ):
+            self.kv_connector.pre_forward(scheduler_output, input_batch)
+            model_output = self.model(**model_inputs)
+        self.kv_connector.finish_forward()
+
+        if self.is_last_pp_rank:
+            assert isinstance(model_output, torch.Tensor)
+            hidden_states = model_output
+            output_intermediate_tensors = None
+        else:
+            assert isinstance(model_output, IntermediateTensors)
+            hidden_states = None
+            output_intermediate_tensors = model_output
+
+        self.execute_model_state = ExecuteModelState(
+            input_batch=input_batch,
+            hidden_states=hidden_states,
+            finished_req_ids=scheduler_output.finished_req_ids,
+            ec_connector_output=ec_connector_output,
+        )
+
+        return output_intermediate_tensors
+
+    @torch.inference_mode()
+    def sample_tokens(
+        self, grammar_output: GrammarOutput | None
+    ) -> ModelRunnerOutput | None:
+        if self.execute_model_state is None:
+            return None
+
+        input_batch = self.execute_model_state.input_batch
+        hidden_states = self.execute_model_state.hidden_states
+        finished_req_ids = self.execute_model_state.finished_req_ids
+        ec_connector_output = self.execute_model_state.ec_connector_output
+        self.execute_model_state = None
+
+        if not self.is_last_pp_rank:
+            assert self.pp_handler is not None
+            self.pp_handler.receive(input_batch)
+            self.postprocess_num_computed_tokens(input_batch)
+
+            kv_connector_output = self.kv_connector.post_forward(finished_req_ids)
+            output = ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output)
+            return ModelRunnerOutput.with_ec_conn_output(output, ec_connector_output)
+
+        assert hidden_states is not None
+        sampler_output, num_sampled, num_rejected = self.sample(
+            hidden_states, input_batch, grammar_output
+        )
+
+        if self.pp_handler is not None:
+            self.pp_handler.broadcast(
+                sampler_output.sampled_token_ids,
+                num_sampled,
+                num_rejected,
+                input_batch,
+            )
+
+        assert self.prompt_logprobs_worker is not None
+        prompt_logprobs_dict = self.prompt_logprobs_worker.compute_prompt_logprobs(
+            self.model.compute_logits,
+            hidden_states,
+            input_batch,
+            self.req_states.all_token_ids,
+            self.req_states.num_computed_tokens,
+            self.req_states.prompt_len_np,
+        )
+        prompt_token_id_logprobs_dict = (
+            self.prompt_logprobs_worker.compute_prompt_token_id_logprobs(
+                self.model.compute_logits,
+                hidden_states,
+                input_batch,
+                self.req_states.prompt_len_np,
+            )
+        )
+
+        sampled_token_ids = sampler_output.sampled_token_ids.tolist()
+        num_sampled_list = num_sampled.tolist()
+        for token_ids, num_tokens in zip(
+            sampled_token_ids, num_sampled_list, strict=True
+        ):
+            del token_ids[num_tokens:]
+
+        model_runner_output = ModelRunnerOutput(
+            req_ids=input_batch.req_ids,
+            req_id_to_index={req_id: i for i, req_id in enumerate(input_batch.req_ids)},
+            sampled_token_ids=sampled_token_ids,
+            prompt_logprobs_dict=prompt_logprobs_dict,  # type: ignore[arg-type]
+            prompt_token_id_logprobs_dict=prompt_token_id_logprobs_dict,
+        )
+        if sampler_output.logprobs_tensors is not None:
+            model_runner_output.logprobs = sampler_output.logprobs_tensors.tolists()
+        if sampler_output.sampling_mask_tensors is not None:
+            model_runner_output.sampling_masks = (
+                sampler_output.sampling_mask_tensors.tolists(num_sampled.numpy())
+            )
+        if sampler_output.num_nans is not None:
+            model_runner_output.num_nans_in_logits = dict(
+                zip(
+                    input_batch.req_ids,
+                    sampler_output.num_nans.tolist(),
+                    strict=True,
+                )
+            )
+
+        model_runner_output.kv_connector_output = self.kv_connector.post_forward(
+            finished_req_ids
+        )
+
+        self.postprocess_sampled(
+            input_batch.idx_mapping,
+            sampler_output.sampled_token_ids,
+            num_sampled,
+            num_rejected,
+            input_batch.query_start_loc,
+        )
+
+        model_runner_output.ec_connector_output = ec_connector_output
+        if (
+            envs.VLLM_RAISE_ON_LOGIT_NANS
+            and model_runner_output.num_nans_in_logits is not None
+        ):
+            raise_if_nan_logits(model_runner_output.num_nans_in_logits)
+        return model_runner_output
+
+    @torch.inference_mode()
+    def pool(self) -> ModelRunnerOutput | None:
+        if self.execute_model_state is None:
+            return None
+
+        input_batch = self.execute_model_state.input_batch
+        hidden_states = self.execute_model_state.hidden_states
+        finished_req_ids = self.execute_model_state.finished_req_ids
+        ec_connector_output = self.execute_model_state.ec_connector_output
+        self.execute_model_state = None
+
+        kv_connector_output = self.kv_connector.post_forward(finished_req_ids)
+        if not self.is_last_pp_rank:
+            self.postprocess_num_computed_tokens(input_batch)
+            output = ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output)
+            return ModelRunnerOutput.with_ec_conn_output(output, ec_connector_output)
+
+        assert hidden_states is not None
+        assert self.pooling_runner is not None
+        pooler_output, finished_mask = self.pooling_runner.pool(
+            hidden_states, input_batch, self.req_states
+        )
+
+        outputs = (
+            pooler_output.unbind()
+            if isinstance(pooler_output, torch.Tensor)
+            else pooler_output
+        )
+        model_runner_output = ModelRunnerOutput(
+            req_ids=input_batch.req_ids,
+            req_id_to_index={req_id: i for i, req_id in enumerate(input_batch.req_ids)},
+            pooler_output=[
+                None if output is None or not is_finished else output
+                for output, is_finished in zip(outputs, finished_mask, strict=True)
+            ],
+            kv_connector_output=kv_connector_output,
+            ec_connector_output=ec_connector_output,
+        )
+        self.postprocess_num_computed_tokens(input_batch)
+        return model_runner_output
+
+    def shutdown(self) -> None:
+        """Release CPU tensors (model weights, KV caches, workspace) so that
+        memory is reclaimable when running in the same process."""
+        set_offloader(None)
+        self.fast_prefill = None
+        if self.pooling_runner is not None:
+            self.pooling_runner.clear()
+        self.pooling_runner = None
+        self.kv_block_zeroer = None
+        if hasattr(self, "kv_caches"):
+            self.kv_caches.clear()
+        if hasattr(self, "attn_groups"):
+            self.attn_groups.clear()
+        if hasattr(self, "kv_cache_config"):
+            del self.kv_cache_config
+        if self.encoder_cache is not None:
+            self.encoder_cache.reset_encoder_cache()
+        free_before_shutdown(self.vllm_config)
+        if hasattr(self, "model_state"):
+            del self.model_state
+        # Detach the layer-level KV/state cache tensors before dropping the
+        # model; the model object can outlive this runner.
+        if hasattr(self, "model"):
+            clear_layer_kv_caches(self.model.modules())
+            del self.model
+
+        gc.collect()
+        logger.debug("Cleaned up model weights, KV caches, and workspace")
diff --git a/vllm/v1/worker/cpu/states.py b/vllm/v1/worker/cpu/states.py
new file mode 100644
index 000000000000..5cbc10a5d57d
--- /dev/null
+++ b/vllm/v1/worker/cpu/states.py
@@ -0,0 +1,113 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import numpy as np
+import torch
+
+
+class CPURequestState:
+    def __init__(
+        self,
+        max_num_reqs: int,
+        max_model_len: int,
+        max_num_batched_tokens: int,
+        vocab_size: int,
+        device: torch.device,
+    ):
+        self.max_num_reqs = max_num_reqs
+        self.max_model_len = max_model_len
+        self.max_num_batched_tokens = max_num_batched_tokens
+        self.vocab_size = vocab_size
+        self.device = device
+
+        self.req_id_to_index: dict[str, int] = {}
+        self.index_to_req_id: dict[int, str] = {}
+        self.free_indices = list(range(max_num_reqs))
+
+        # CPU keeps this tensor dense because it is already host-resident.
+        self.all_token_ids = torch.zeros(
+            (self.max_num_reqs, self.max_model_len),
+            dtype=torch.int32,
+            device=device,
+        )
+
+        # NOTE(woosuk): Distinguish clearly between prompt_len and prefill_len:
+        # - prompt_len: Number of tokens in the user-provided prompt.
+        # - prefill_len: Number of tokens passed into the model runner.
+        #   This can include the prompt and additional partial output tokens,
+        #   so prefill_len >= prompt_len.
+        # Usually, prefill_len equals prompt_len, but in cases such as resumption after
+        # preemption, prefill_len may be greater. Differentiating between these values
+        # is crucial, as certain features such as prompt logprobs or frequency penalties
+        # must treat prompt and output tokens separately.
+        self.prompt_len = torch.zeros(
+            self.max_num_reqs, dtype=torch.int32, device=device
+        )
+        self.prompt_len_np = self.prompt_len.numpy()
+        self.prefill_len = torch.zeros(
+            self.max_num_reqs, dtype=torch.int32, device=device
+        )
+        self.prefill_len_np = self.prefill_len.numpy()
+        # total_len = prompt_len + output_len. It grows as the request progresses.
+        self.total_len = torch.zeros(
+            self.max_num_reqs, dtype=torch.int32, device=device
+        )
+
+        # Number of computed tokens.
+        self.num_computed_prefill_tokens = np.zeros(self.max_num_reqs, dtype=np.int32)
+        self.num_computed_tokens = torch.zeros(
+            self.max_num_reqs, dtype=torch.int32, device=device
+        )
+        self.num_computed_tokens_np = self.num_computed_tokens.numpy()
+
+        # Last sampled tokens.
+        self.last_sampled_tokens = torch.zeros(
+            self.max_num_reqs, 1, dtype=torch.int64, device=device
+        )
+        # Max total seq length (prompt_len + max_tokens).
+        self.max_seq_len = np.zeros(self.max_num_reqs, dtype=np.int32)
+        self.next_prefill_tokens = torch.zeros(
+            1,
+            self.max_num_reqs,
+            dtype=torch.int32,
+            device=device,
+        )
+
+    @property
+    def num_reqs(self) -> int:
+        return len(self.req_id_to_index)
+
+    def add_request(
+        self,
+        req_id: str,
+        prompt_len: int,
+        all_token_ids: list[int],
+        num_computed_tokens: int,
+        max_tokens: int,
+    ) -> None:
+        assert len(self.free_indices) > 0, "No free indices"
+        req_idx = self.free_indices.pop()
+        self.req_id_to_index[req_id] = req_idx
+        self.index_to_req_id[req_idx] = req_id
+
+        self.max_seq_len[req_idx] = prompt_len + max_tokens
+        self.prompt_len_np[req_idx] = prompt_len
+        prefill_len = len(all_token_ids)
+        assert prefill_len >= prompt_len, (
+            f"prefill_len {prefill_len} < prompt_len {prompt_len}"
+        )
+        self.prefill_len_np[req_idx] = prefill_len
+        self.total_len[req_idx] = prefill_len
+        self.all_token_ids[req_idx, :prefill_len] = torch.tensor(
+            all_token_ids, dtype=torch.int32, device=self.device
+        )
+        self.num_computed_prefill_tokens[req_idx] = num_computed_tokens
+        self.num_computed_tokens_np[req_idx] = num_computed_tokens
+
+    def remove_request(self, req_id: str) -> int | None:
+        """Return the freed slot index, or None if the request was not found."""
+        req_idx = self.req_id_to_index.pop(req_id, None)
+        if req_idx is None:
+            return None
+        self.index_to_req_id.pop(req_idx, None)
+        self.free_indices.append(req_idx)
+        return req_idx
diff --git a/vllm/v1/worker/cpu/structured_outputs.py b/vllm/v1/worker/cpu/structured_outputs.py
new file mode 100644
index 000000000000..7fd2c6f867b7
--- /dev/null
+++ b/vllm/v1/worker/cpu/structured_outputs.py
@@ -0,0 +1,59 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import numpy as np
+import torch
+
+from vllm.v1.worker.cpu.input_batch import InputBatch
+
+
+class CPUStructuredOutputsWorker:
+    def __init__(
+        self,
+        max_num_logits: int,
+        vocab_size: int,
+        device: torch.device,
+        mask_stride: int,
+    ) -> None:
+        self.mask_stride = mask_stride
+
+    def apply_grammar_bitmask(
+        self,
+        logits: torch.Tensor,
+        input_batch: InputBatch,
+        grammar_req_ids: list[str],
+        grammar_bitmask: np.ndarray,
+    ) -> None:
+        if not grammar_req_ids:
+            return
+
+        # Each grammar request owns `mask_stride` bitmask rows. Key them by
+        # (request, position); positions beyond cu_num_logits are skipped.
+        req_id_to_idx = {req_id: i for i, req_id in enumerate(input_batch.req_ids)}
+        mapping = [
+            req_id_to_idx[req_id] * self.mask_stride + position
+            for req_id in grammar_req_ids
+            for position in range(self.mask_stride)
+        ]
+        assert grammar_bitmask.shape[0] == len(mapping)
+
+        mapping_tensor = torch.tensor(mapping, dtype=torch.int64, device=logits.device)
+        req_indices = mapping_tensor // self.mask_stride
+        positions = mapping_tensor % self.mask_stride
+        starts = input_batch.cu_num_logits[req_indices].long()
+        ends = input_batch.cu_num_logits[req_indices + 1].long()
+        active = positions < ends - starts
+        if not torch.any(active):
+            return
+
+        logit_rows = starts[active] + positions[active]
+        packed = torch.from_numpy(grammar_bitmask).to(logits.device)[active]
+        token_ids = torch.arange(
+            logits.shape[1], dtype=torch.int64, device=logits.device
+        )
+        words = packed[:, token_ids // 32]
+        allowed = ((words >> (token_ids % 32)) & 1).bool()
+
+        masked_logits = logits.index_select(0, logit_rows)
+        masked_logits.masked_fill_(~allowed, -torch.inf)
+        logits.index_copy_(0, logit_rows, masked_logits)
diff --git a/vllm/v1/worker/cpu/utils.py b/vllm/v1/worker/cpu/utils.py
new file mode 100644
index 000000000000..75e9f6fe1fd1
--- /dev/null
+++ b/vllm/v1/worker/cpu/utils.py
@@ -0,0 +1,88 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from collections.abc import Iterable
+from typing import TYPE_CHECKING, Any
+
+import torch
+
+from vllm.config import VllmConfig
+from vllm.v1.kv_cache_interface import AttentionSpec
+from vllm.v1.worker.utils import KVBlockZeroer
+
+if TYPE_CHECKING:
+    from vllm.v1.worker.utils import AttentionGroup
+
+
+def free_before_shutdown(vllm_config: VllmConfig) -> None:
+    from vllm.model_executor.layers.rotary_embedding import _ROPE_DICT
+    from vllm.v1.worker.workspace import reset_workspace_manager
+
+    cache_config = vllm_config.cache_config
+    cache_config.num_gpu_blocks = None
+
+    compilation_config = vllm_config.compilation_config
+    compilation_config.static_forward_context.clear()
+
+    _ROPE_DICT.clear()
+    reset_workspace_manager()
+
+
+class CPUKVBlockZeroer(KVBlockZeroer):
+    """Zero newly allocated logical KV blocks through their tensor views."""
+
+    def __init__(
+        self,
+        device: torch.device,
+        attn_groups_iter: Iterable["AttentionGroup"],
+        kernel_block_sizes: list[int],
+        static_forward_context: dict[str, Any],
+        num_blocks: int,
+        runner_only_attn_layers: set[str] | None = None,
+    ) -> None:
+        self.device = device
+        self._cache_views: list[tuple[torch.Tensor, int]] = []
+
+        if runner_only_attn_layers is None:
+            runner_only_attn_layers = set()
+        for group in attn_groups_iter:
+            spec = group.kv_cache_spec
+            if not isinstance(spec, AttentionSpec):
+                continue
+            if group.kv_cache_group_id >= len(kernel_block_sizes):
+                continue
+            kernel_block_size = kernel_block_sizes[group.kv_cache_group_id]
+            assert spec.block_size % kernel_block_size == 0
+            for layer_name in group.layer_names:
+                if layer_name in runner_only_attn_layers:
+                    continue
+                kv_cache = static_forward_context[layer_name].kv_cache
+                if not isinstance(kv_cache, torch.Tensor):
+                    continue
+                if group.kernel_block_stride:
+                    kv_cache = group.map_kv_cache_to_manager_blocks(
+                        kv_cache, num_blocks
+                    )
+                if kv_cache.device.type != self.device.type:
+                    continue
+                assert kv_cache.shape[0] % num_blocks == 0, (
+                    f"{layer_name}: {kv_cache.shape[0]} kernel blocks is not a "
+                    f"multiple of {num_blocks} logical blocks"
+                )
+                self._cache_views.append((kv_cache, kv_cache.shape[0] // num_blocks))
+
+    def zero_block_ids(self, block_ids: list[int]) -> None:
+        if not block_ids or not self._cache_views:
+            return
+
+        logical_block_ids = torch.tensor(block_ids, dtype=torch.int64)
+        indices_by_ratio = {1: logical_block_ids}
+        for kv_cache, blocks_per_logical_block in self._cache_views:
+            kernel_block_ids = indices_by_ratio.get(blocks_per_logical_block)
+            if kernel_block_ids is None:
+                offsets = torch.arange(blocks_per_logical_block, dtype=torch.int64)
+                kernel_block_ids = (
+                    logical_block_ids[:, None] * blocks_per_logical_block + offsets
+                ).flatten()
+                indices_by_ratio[blocks_per_logical_block] = kernel_block_ids
+            kv_cache[kernel_block_ids] = 0
diff --git a/vllm/v1/worker/cpu/warmup.py b/vllm/v1/worker/cpu/warmup.py
new file mode 100644
index 000000000000..b83869f161f2
--- /dev/null
+++ b/vllm/v1/worker/cpu/warmup.py
@@ -0,0 +1,227 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+import torch
+
+from vllm import PoolingParams, SamplingParams
+from vllm.multimodal.inputs import MultiModalFeatureSpec, PlaceholderRange
+from vllm.utils.math_utils import cdiv
+from vllm.v1.core.sched.output import (
+    CachedRequestData,
+    GrammarOutput,
+    NewRequestData,
+    SchedulerOutput,
+)
+from vllm.v1.kv_cache_interface import (
+    CircularBufferSpec,
+    CrossAttentionSpec,
+    KVCacheSpec,
+    MambaSpec,
+    UniformTypeKVCacheSpecs,
+)
+from vllm.v1.request import Request
+
+if TYPE_CHECKING:
+    from vllm.v1.worker.cpu.standalone_model_runner import StandaloneCPUModelRunner
+
+
+def _reserved_block_count(
+    num_tokens: int,
+    kvcache_spec: KVCacheSpec,
+    *,
+    num_lookahead_tokens: int,
+    max_model_len: int,
+    max_encoder_len: int,
+) -> int:
+    if isinstance(kvcache_spec, UniformTypeKVCacheSpecs):
+        kvcache_spec = kvcache_spec.first_spec
+    if isinstance(kvcache_spec, CircularBufferSpec):
+        return 1
+    if isinstance(kvcache_spec, CrossAttentionSpec):
+        return cdiv(max_encoder_len, kvcache_spec.block_size)
+    num_speculative_blocks = 0
+    if isinstance(kvcache_spec, MambaSpec):
+        num_speculative_blocks = kvcache_spec.num_speculative_blocks
+        if kvcache_spec.mamba_cache_mode == "align":
+            return cdiv(num_tokens, kvcache_spec.block_size) + num_speculative_blocks
+    num_tokens = min(num_tokens + num_lookahead_tokens, max_model_len)
+    return cdiv(num_tokens, kvcache_spec.block_size) + num_speculative_blocks
+
+
+def _warmup_block_counter(
+    model_runner: "StandaloneCPUModelRunner",
+) -> Callable[[int, KVCacheSpec], int]:
+    num_lookahead_tokens = model_runner.vllm_config.num_lookahead_tokens
+    max_model_len = model_runner.max_model_len
+    max_encoder_len = getattr(model_runner.model_state, "max_encoder_len", 0)
+
+    def block_count(num_tokens: int, kvcache_spec: KVCacheSpec) -> int:
+        return _reserved_block_count(
+            num_tokens,
+            kvcache_spec,
+            num_lookahead_tokens=num_lookahead_tokens,
+            max_model_len=max_model_len,
+            max_encoder_len=max_encoder_len,
+        )
+
+    return block_count
+
+
+@torch.inference_mode()
+def warmup_kernels(
+    model_runner: "StandaloneCPUModelRunner",
+    worker_execute_model: Callable[[SchedulerOutput], Any],
+    worker_sample_tokens: Callable[[GrammarOutput | None], Any],
+) -> None:
+    if model_runner.vllm_config.is_mm_encoder_only:
+        return
+
+    decode_query_len = model_runner.decode_query_len
+    prompt_len = decode_query_len + 1
+    prompt_token_ids = list(range(prompt_len))
+    num_decode_steps = 1 if model_runner.is_pooling_model else 3
+    decode_len = prompt_len + num_decode_steps * decode_query_len
+
+    kv_cache_groups = model_runner.kv_cache_config.kv_cache_groups
+    max_encoder_len = getattr(model_runner.model_state, "max_encoder_len", 0)
+    warmup_mm_features: list[MultiModalFeatureSpec] = []
+    if model_runner.is_encoder_decoder and max_encoder_len:
+        warmup_mm_features = [
+            MultiModalFeatureSpec(
+                data=None,
+                modality="",
+                identifier="_warmup_encoder",
+                mm_position=PlaceholderRange(offset=0, length=max_encoder_len),
+            )
+        ]
+
+    block_count = _warmup_block_counter(model_runner)
+    kv_cache_specs = [group.kv_cache_spec for group in kv_cache_groups]
+    prefill_block_counts = [block_count(prompt_len, spec) for spec in kv_cache_specs]
+    decode_block_counts = [block_count(decode_len, spec) for spec in kv_cache_specs]
+    max_blocks_per_req = sum(decode_block_counts)
+
+    num_reqs = min(
+        model_runner.scheduler_config.max_num_seqs,
+        model_runner.scheduler_config.max_num_batched_tokens
+        // max(prompt_len, decode_query_len),
+    )
+    null_blocks = model_runner.block_tables.redirect_writes_to_null_block
+    if max_blocks_per_req > 0 and not null_blocks:
+        num_reqs = min(
+            num_reqs,
+            max(1, (model_runner.kv_cache_config.num_blocks - 1) // max_blocks_per_req),
+        )
+
+    req_ids = [f"_warmup_{i}_" for i in range(num_reqs)]
+    if model_runner.is_pooling_model:
+        sampling_params = None
+        pooling_task = model_runner.model_config.get_pooling_task(
+            model_runner.get_supported_tasks()
+        )
+        pooling_params = PoolingParams(task=pooling_task)
+        pooling_params.verify(model_runner.model_config)
+    else:
+        sampling_params = SamplingParams.for_sampler_warmup()
+        pooling_params = None
+
+    next_block_id = 1
+
+    def _alloc_blocks(num_blocks: int) -> list[int]:
+        nonlocal next_block_id
+        return list(range(next_block_id, next_block_id := next_block_id + num_blocks))
+
+    if model_runner.kv_block_zeroer is not None:
+        model_runner.kv_block_zeroer.warmup(model_runner.kv_cache_config.num_blocks)
+
+    new_reqs = [
+        NewRequestData.from_request(
+            Request(
+                req_ids[i],
+                prompt_token_ids,
+                sampling_params,
+                pooling_params,
+                mm_features=warmup_mm_features,
+            ),
+            block_ids=tuple(_alloc_blocks(n) for n in prefill_block_counts),
+            prefill_token_ids=prompt_token_ids,
+        )
+        for i in range(num_reqs)
+    ]
+
+    prefill_output = SchedulerOutput.make_empty()
+    prefill_output.scheduled_new_reqs = new_reqs
+    prefill_output.num_scheduled_tokens = {req_id: prompt_len for req_id in req_ids}
+    prefill_output.total_num_scheduled_tokens = prompt_len * num_reqs
+
+    model_runner.kv_connector.set_disabled(True)
+    try:
+        worker_execute_model(prefill_output)
+
+        if not model_runner.is_pooling_model:
+            grammar_output = None
+            if model_runner.is_last_pp_rank:
+                vocab_size = model_runner.model_config.get_vocab_size()
+                bitmask_width = (vocab_size + 31) // 32
+                grammar_output = GrammarOutput(
+                    structured_output_request_ids=req_ids,
+                    grammar_bitmask=np.full(
+                        (len(req_ids) * decode_query_len, bitmask_width),
+                        fill_value=-1,
+                        dtype=np.int32,
+                    ),
+                )
+            worker_sample_tokens(grammar_output)
+
+            req_computed = [prompt_len] * num_reqs
+            req_blocks = [list(prefill_block_counts) for _ in range(num_reqs)]
+
+            def _run_decode_step(indices: list[int]) -> None:
+                cached_req_data = CachedRequestData.make_empty()
+                cached_req_data.req_ids = [req_ids[i] for i in indices]
+                cached_req_data.num_computed_tokens = [req_computed[i] for i in indices]
+                cached_req_data.num_output_tokens = [1] * len(indices)
+                cached_req_data.new_block_ids = []
+
+                num_scheduled_tokens: dict[str, int] = {}
+                for i in indices:
+                    after = req_computed[i] + 1
+                    deltas = [
+                        block_count(after, spec) - held
+                        for spec, held in zip(kv_cache_specs, req_blocks[i])
+                    ]
+                    cached_req_data.new_block_ids.append(
+                        tuple(_alloc_blocks(n) for n in deltas) if any(deltas) else None
+                    )
+                    req_blocks[i] = [
+                        held + delta for held, delta in zip(req_blocks[i], deltas)
+                    ]
+                    num_scheduled_tokens[req_ids[i]] = 1
+
+                decode_output = SchedulerOutput.make_empty()
+                decode_output.scheduled_cached_reqs = cached_req_data
+                decode_output.num_scheduled_tokens = num_scheduled_tokens
+                decode_output.total_num_scheduled_tokens = len(indices)
+                worker_execute_model(decode_output)
+                worker_sample_tokens(None)
+
+                for i in indices:
+                    req_computed[i] += 1
+
+            all_indices = list(range(num_reqs))
+            _run_decode_step(all_indices)
+            if num_reqs > 1:
+                _run_decode_step([0])
+
+        cleanup_output = SchedulerOutput.make_empty()
+        cleanup_output.finished_req_ids = set(req_ids)
+        worker_execute_model(cleanup_output)
+    finally:
+        model_runner.kv_connector.set_disabled(False)
+
+    if model_runner.kv_block_zeroer is not None:
+        model_runner.kv_block_zeroer.zero_block_ids([0])
diff --git a/vllm/v1/worker/cpu_worker.py b/vllm/v1/worker/cpu_worker.py
index a8cd0b3b7b50..5b7dc061fcc5 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 TYPE_CHECKING, cast
 
 import psutil
 import torch
@@ -27,9 +26,6 @@
 from vllm.v1.worker.gpu_worker import Worker, init_worker_distributed_environment
 from vllm.v1.worker.worker_base import CompilationTimes
 
-if TYPE_CHECKING:
-    from vllm.v1.worker.gpu.model_runner import GPUModelRunner
-
 logger = init_logger(__name__)
 
 
@@ -76,7 +72,7 @@ def __init__(
         torch.ops._C.init_cpu_memory_env([memory_node])
 
         memory_status = get_memory_node_info(memory_node)
-        memory_fraction = vllm_config.cache_config.gpu_memory_utilization
+        memory_fraction = vllm_config.cache_config.resolved_gpu_memory_utilization
         self.requested_cpu_memory = math.ceil(
             memory_status.total_memory * memory_fraction
         )
@@ -153,13 +149,22 @@ def skip_set_num_threads(x: int):
 
         # Construct the model runner
         if self.use_v2_model_runner:
-            from vllm.v1.worker.cpu.model_runner import (
-                CPUModelRunner as CPUModelRunnerV2,
-            )
+            if current_platform.uses_standalone_cpu_model_runner():
+                from vllm.v1.worker.cpu.standalone_model_runner import (
+                    StandaloneCPUModelRunner,
+                )
 
-            self.model_runner: CPUModelRunner = CPUModelRunnerV2(  # type: ignore
-                self.vllm_config, self.device
-            )
+                self.model_runner: CPUModelRunner = StandaloneCPUModelRunner(  # type: ignore
+                    self.vllm_config, self.device
+                )
+            else:
+                from vllm.v1.worker.cpu.model_runner import (
+                    CPUModelRunner as CPUModelRunnerV2,
+                )
+
+                self.model_runner = CPUModelRunnerV2(  # type: ignore
+                    self.vllm_config, self.device
+                )
         else:
             self.model_runner = CPUModelRunner(self.vllm_config, torch.device("cpu"))
 
@@ -271,10 +276,13 @@ def compile_or_warm_up_model(self) -> CompilationTimes:
         # their Triton JIT lands in the first request's TTFT.
         kernel_config = self.vllm_config.kernel_config
         if self.use_v2_model_runner and kernel_config.enable_jit_warmup:
-            from vllm.v1.worker.gpu.warmup import warmup_kernels
+            if current_platform.uses_standalone_cpu_model_runner():
+                from vllm.v1.worker.cpu.warmup import warmup_kernels
+            else:
+                from vllm.v1.worker.gpu.warmup import warmup_kernels
 
             warmup_kernels(
-                cast("GPUModelRunner", self.model_runner),
+                self.model_runner,
                 self.execute_model,
                 self.sample_tokens,
             )
diff --git a/vllm/v1/worker/extensible_kv_cache.py b/vllm/v1/worker/extensible_kv_cache.py
new file mode 100644
index 000000000000..05a0ba746073
--- /dev/null
+++ b/vllm/v1/worker/extensible_kv_cache.py
@@ -0,0 +1,633 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""KV cache backing allocation that grows after warmup."""
+
+import copy
+import gc
+from bisect import bisect_right
+from collections.abc import Iterator
+from contextlib import contextmanager
+from dataclasses import replace
+from typing import TYPE_CHECKING
+
+import torch
+
+from vllm.distributed.kv_transfer import (
+    ensure_kv_transfer_initialized,
+    get_kv_transfer_group,
+    has_kv_transfer_group,
+)
+from vllm.logger import init_logger
+from vllm.model_executor.model_loader.reload.layerwise import max_layer_reload_bytes
+from vllm.multimodal.gpu_ipc_memory import reserve_mm_ipc_gpu_memory
+from vllm.platforms import current_platform
+from vllm.utils.extensible_tensor import (
+    ExtensibleTensor,
+    granule_aligned_blocks,
+    granule_block_alignment,
+)
+from vllm.utils.mem_utils import format_gib
+from vllm.v1.kv_cache_interface import (
+    KVCacheConfig,
+    KVCacheTensor,
+    create_kv_cache_views,
+)
+from vllm.v1.kv_cache_layout import KVCacheLayout
+from vllm.v1.worker.gpu.attn_utils import bind_kv_caches
+from vllm.v1.worker.gpu.kv_connector import get_kv_connector
+from vllm.v1.worker.utils import layer_kernel_block_size, layer_spec_for_tensor
+
+if TYPE_CHECKING:
+    from vllm.config import VllmConfig
+    from vllm.v1.worker.gpu.model_runner import GPUModelRunner
+
+logger = init_logger(__name__)
+
+# Headroom left after sizing from measured memory, beyond the measured transient
+# peak, for what the allocated peak does not show: the caching allocator's
+# rounding and fragmentation around it, real steps combining shapes warmup
+# exercises separately (a prompt-logprobs chunk alongside a live decode, the
+# rejection sampler over a full spec-decode batch), and inputs profiling only
+# approximates (a multimodal encoder batch of a modality other than the one
+# with the largest feature size). The share scales with the model's working
+# set; the floor covers small models whose measured peak is a few hundred MiB
+# while their encoder or sampler transients are not. Measured on GB200: the
+# first prefill of an 8B model needed over 25% beyond its allocated peak before
+# the allocator could serve it, and adversarial workloads reached up to 45%.
+# ROCr also maps kernel scratch outside the allocator when a kernel first needs
+# it, so ROCm keeps a larger floor.
+KV_CACHE_MARGIN_FLOOR_BYTES = (2 if current_platform.is_rocm() else 1) << 30
+KV_CACHE_MARGIN_FRACTION = 0.75
+# Part of the margin kept for allocations outside torch's caching allocator
+# (kernel scratch, NCCL/RCCL buffers, lazily loaded code objects); the
+# allocator is capped to leave it free. The rest absorbs torch's fragmentation,
+# which scales with the transient where these do not: serving grew them by at
+# most 0.17 GiB on MI355 and not at all on GB200.
+KV_CACHE_NON_TORCH_RESERVE_BYTES = (2 if current_platform.is_rocm() else 1) << 29
+# Share of the headroom warmup leaves uncommitted, on top of the profiled
+# activation peak, for transients the profiling run does not exercise.
+KV_CACHE_WARMUP_RESERVE_FRACTION = 0.1
+
+
+class ExtensibleKVCache:
+    """Growable backing allocation for the KV cache.
+
+    Reserves address space for `kv_cache_config.num_blocks` and commits a
+    prefix of blocks on demand. Blocks keep fixed offsets, so views over the
+    reservation and captured CUDA graphs stay valid as more are committed. The
+    allocation splits into segments holding num_blocks blocks each: one per
+    layer for layer-outermost layouts (sized by that layer's block stride), one
+    overall for block-outermost.
+    """
+
+    def __init__(self, kv_cache_config: KVCacheConfig, device: torch.device):
+        tensors = kv_cache_config.kv_cache_tensors
+        assert tensors, "Extensible KV cache requires at least one KV cache tensor."
+        sizes = {tensor.size for tensor in tensors}
+        assert len(sizes) == 1, "KV cache tensors must share one backing allocation."
+        self.size = sizes.pop()
+        self.capacity_blocks = kv_cache_config.num_blocks
+        if self.capacity_blocks <= 0 or self.size % self.capacity_blocks != 0:
+            raise ValueError(
+                f"KV cache allocation of {self.size} bytes is not a whole number "
+                f"of {self.capacity_blocks} blocks."
+            )
+        self.bytes_per_block = self.size // self.capacity_blocks
+        self.segment_offsets, self.segment_strides = self._segments(tensors)
+        self.num_committed_blocks = 0
+        # Memory to keep free while committing before the final sizing, e.g.
+        # the profiled activation peak the warmup steps are about to hit.
+        self.reserved_headroom_bytes = 0
+        # Blocks warmup may commit: the minimum of `committable_blocks()` over
+        # ranks when the cache is created.
+        self.warmup_committable_blocks = 0
+        # Bytes torch's caching allocator may reserve, set once the cache is
+        # sized; None leaves it unlimited.
+        self.torch_memory_limit: int | None = None
+        self.buffer = ExtensibleTensor(
+            self.size,
+            device=device,
+            segment_capacities=[
+                self.capacity_blocks * stride for stride in self.segment_strides
+            ],
+        )
+
+    def _segments(self, tensors: list[KVCacheTensor]) -> tuple[list[int], list[int]]:
+        """Segment offsets and block strides tiling the allocation.
+
+        With one block stride the segments are implied by it; with several, each
+        tensor must give its layers contiguous num_blocks-block regions.
+        """
+        block_strides = {tensor.block_stride for tensor in tensors}
+        if len(block_strides) == 1:
+            stride = block_strides.pop()
+            segment_bytes = self.capacity_blocks * stride
+            num_segments = self.size // segment_bytes
+            return [i * segment_bytes for i in range(num_segments)], [
+                stride
+            ] * num_segments
+        segments: dict[int, int] = {}
+        for tensor in tensors:
+            if tensor.layer_stride != self.capacity_blocks * tensor.block_stride:
+                raise ValueError(
+                    "KV cache tensors with different block strides need a "
+                    "layer-outermost layout for the extensible KV cache."
+                )
+            for layer_idx in range(len(tensor.layers)):
+                offset = tensor.offset + layer_idx * tensor.layer_stride
+                stride = segments.setdefault(offset, tensor.block_stride)
+                if stride != tensor.block_stride:
+                    raise ValueError(
+                        f"KV cache tensors alias offset {offset} with different "
+                        "block strides."
+                    )
+        offsets = sorted(segments)
+        strides = [segments[offset] for offset in offsets]
+        end = 0
+        for offset, stride in zip(offsets, strides):
+            if offset != end:
+                raise ValueError("KV cache tensors do not tile the allocation.")
+            end += self.capacity_blocks * stride
+        if end != self.size:
+            raise ValueError("KV cache tensors do not tile the allocation.")
+        return offsets, strides
+
+    def allocate(self, size: int) -> torch.Tensor:
+        """Allocator hook for allocate_kv_cache; commits only the null block."""
+        assert size == self.size, f"Expected {self.size} bytes, got {size}."
+        self.commit(1)
+        return self.buffer.full_view()
+
+    def commit(
+        self, num_blocks: int, defragment: bool = False, shrink: bool = False
+    ) -> None:
+        """Back the first num_blocks blocks; new blocks are zeroed.
+
+        Grow-only unless shrink is set, which remaps a smaller prefix and
+        discards contents. defragment remaps each segment as one driver
+        allocation (discarding contents): UCX cannot RDMA a range spanning
+        several allocations.
+        """
+        num_blocks = min(num_blocks, self.capacity_blocks)
+        sizes = [num_blocks * stride for stride in self.segment_strides]
+        remap = (
+            defragment and not all(self.buffer.segments_backed_by_one_chunk(sizes))
+        ) or (shrink and num_blocks < self.num_committed_blocks)
+        if remap:
+            self.buffer.release_physical()
+            self.num_committed_blocks = 0
+        if num_blocks <= self.num_committed_blocks:
+            return
+        self.buffer.resize_segments_(sizes, zero_new=True)
+        self.num_committed_blocks = num_blocks
+        logger.debug(
+            "Committed %d of %d KV cache blocks (%.2f GiB physical).",
+            num_blocks,
+            self.capacity_blocks,
+            self.physical_bytes / (1 << 30),
+        )
+
+    def committed_views(
+        self,
+        kv_cache_config: KVCacheConfig,
+        layout: KVCacheLayout,
+        kernel_block_sizes: list[int] | None = None,
+    ) -> dict[str, torch.Tensor] | None:
+        """Per-layer views whose storage spans exactly the committed blocks.
+
+        `kv_cache_config.num_blocks` must equal the committed count. Connectors
+        derive extents from `untyped_storage().nbytes()`, so they must not see the
+        reservation. None for layouts whose layers span several segments.
+        """
+        if not layout.is_block_compact:
+            return None
+        assert kv_cache_config.num_blocks == self.num_committed_blocks
+        kv_caches: dict[str, torch.Tensor] = {}
+        for tensor in kv_cache_config.kv_cache_tensors:
+            group_id, spec = layer_spec_for_tensor(kv_cache_config, tensor)
+            if not spec.has_layer_views:
+                # Laid out by the layer itself over the raw buffer; as at
+                # allocation, it keeps the whole reservation.
+                full_view = self.buffer.full_view()
+                kv_caches.update((name, full_view) for name in tensor.layers)
+                continue
+            kernel_block_size = layer_kernel_block_size(group_id, kernel_block_sizes)
+            for layer_idx, layer_name in enumerate(tensor.layers):
+                segment, offset = self._locate(
+                    tensor.offset + layer_idx * tensor.layer_stride
+                )
+                (view,) = create_kv_cache_views(
+                    self.buffer.segment_view(segment),
+                    spec,
+                    kv_cache_config.num_blocks,
+                    layout,
+                    replace(tensor, layers=[layer_name], offset=offset),
+                    kernel_block_size=kernel_block_size,
+                )
+                kv_caches[layer_name] = view
+        return kv_caches
+
+    def committable_blocks(self) -> int:
+        """Blocks that can be committed now without eating into the headroom.
+
+        Bounds what warmup may commit: free memory plus the committed prefix,
+        less the sizing margin and the larger of reserved_headroom_bytes and
+        a share of the headroom, never below what is already committed nor
+        above the capacity.
+        """
+        free_memory, _ = torch.accelerator.get_memory_info(self.buffer.device)
+        headroom = free_memory + self.physical_bytes
+        reserve = max(
+            self.reserved_headroom_bytes,
+            int(headroom * KV_CACHE_WARMUP_RESERVE_FRACTION),
+        )
+        margin = max(
+            KV_CACHE_MARGIN_FLOOR_BYTES, int(reserve * KV_CACHE_MARGIN_FRACTION)
+        )
+        budget = headroom - margin - reserve
+        return max(self.blocks_within(budget), self.num_committed_blocks)
+
+    @property
+    def physical_bytes(self) -> int:
+        """Physically mapped bytes, including granule rounding."""
+        return self.buffer.physical_bytes
+
+    def physical_bytes_for(self, num_blocks: int) -> int:
+        """Bytes a commit of num_blocks maps, each segment rounded to granules."""
+        granule = self.buffer.granularity
+        return sum(
+            -(-num_blocks * stride // granule) * granule
+            for stride in self.segment_strides
+        )
+
+    @property
+    def block_alignment(self) -> int:
+        """Block count multiple at which every segment ends on a granule."""
+        return granule_block_alignment(self.segment_strides, self.buffer.granularity)
+
+    def blocks_within(self, num_bytes: int, aligned: bool = False) -> int:
+        """Most blocks whose physical footprint fits in num_bytes.
+
+        aligned restricts the count to multiples of block_alignment, so
+        that a defragmenting commit backs each segment with one chunk exactly.
+        """
+        num_blocks = min(num_bytes // self.bytes_per_block, self.capacity_blocks)
+        if aligned:
+            return granule_aligned_blocks(
+                num_blocks, self.segment_strides, self.buffer.granularity
+            )
+        while num_blocks > 0 and self.physical_bytes_for(num_blocks) > num_bytes:
+            num_blocks -= 1
+        return num_blocks
+
+    def _locate(self, start: int) -> tuple[int, int]:
+        """Segment index and offset within it of allocation byte start."""
+        segment = bisect_right(self.segment_offsets, start) - 1
+        return segment, start - self.segment_offsets[segment]
+
+    def committed_kv_cache_tensors(
+        self, kv_cache_config: KVCacheConfig, layout: KVCacheLayout, num_blocks: int
+    ) -> list[KVCacheTensor]:
+        """Tensor placements matching committed_views for num_blocks.
+
+        size is the committed total, as every tensor describes the one
+        allocation. Under block-outermost layouts the committed blocks are one
+        prefix of it, so the capacity placements still hold. Under layer-outermost
+        layouts each layer's view has its own storage over its committed prefix:
+        describe it as a one-layer tensor placed at the start of that storage.
+        """
+        size = num_blocks * self.bytes_per_block
+        tensors = [
+            replace(tensor, size=size) for tensor in kv_cache_config.kv_cache_tensors
+        ]
+        if layout.is_block_outermost:
+            return tensors
+        committed: list[KVCacheTensor] = []
+        for tensor in tensors:
+            _, spec = layer_spec_for_tensor(kv_cache_config, tensor)
+            for layer_idx, layer_name in enumerate(tensor.layers):
+                start = tensor.offset + layer_idx * tensor.layer_stride
+                segment, offset = self._locate(start)
+                if not spec.has_layer_views:
+                    # The layer keeps the whole reservation as its storage.
+                    offset = start
+                committed.append(
+                    KVCacheTensor(
+                        size=size,
+                        layers=[layer_name],
+                        layer_stride=num_blocks * tensor.block_stride,
+                        block_stride=tensor.block_stride,
+                        offset=offset,
+                    )
+                )
+        return committed
+
+    def release_physical(self) -> None:
+        """Drop physical pages for sleep; addresses and views stay valid."""
+        self.buffer.release_physical()
+
+    def recommit(self) -> None:
+        """Map fresh zeroed pages for the blocks committed before release."""
+        num_blocks = self.num_committed_blocks
+        self.num_committed_blocks = 0
+        self.commit(num_blocks)
+
+    def set_torch_memory_limit(self, num_bytes: int | None) -> None:
+        """Cap what torch's caching allocator reserves on this device, or lift
+        the cap (None). At the cap it releases cached segments before growing,
+        rather than taking memory other allocations need."""
+        self.torch_memory_limit = num_bytes
+        total = torch.accelerator.get_memory_info(self.buffer.device)[1]
+        fraction = 1.0 if num_bytes is None else min(num_bytes / total, 1.0)
+        # torch.accelerator has no per-process memory fraction.
+        torch.cuda.set_per_process_memory_fraction(fraction, self.buffer.device)
+
+    def free(self) -> None:
+        # The process may start another engine, which must not inherit the cap.
+        if self.torch_memory_limit is not None:
+            self.set_torch_memory_limit(None)
+        self.buffer.free()
+        self.num_committed_blocks = 0
+
+
+def num_committable_kv_blocks(runner: "GPUModelRunner") -> int:
+    """Blocks warmup may address: all of them, or, for an extensible cache,
+    the count agreed across ranks so that every rank warms up the same shapes."""
+    kv_cache = getattr(runner, "extensible_kv_cache", None)
+    if kv_cache is not None:
+        return kv_cache.warmup_committable_blocks
+    return runner.kv_cache_config.num_blocks
+
+
+def ensure_kv_cache_blocks(runner: "GPUModelRunner", num_blocks: int) -> None:
+    """Commit enough of an extensible KV cache to address num_blocks."""
+    kv_cache = getattr(runner, "extensible_kv_cache", None)
+    if kv_cache is not None:
+        kv_cache.commit(num_blocks)
+
+
+def extend_kv_cache(
+    runner: "GPUModelRunner", kv_cache_config: KVCacheConfig, num_blocks: int
+) -> None:
+    """Commit the final size once post-warmup memory is known.
+
+    Afterwards the runner matches one that allocated num_blocks up front:
+    views are rebuilt over the committed bytes (same addresses) and the KV
+    connector is created against final memory. kv_cache_config is the
+    engine's config the cache was allocated from.
+    """
+    kv_cache = runner.extensible_kv_cache
+    assert kv_cache is not None
+    kv_cache_config = copy.copy(kv_cache_config)
+    kv_cache_config.num_blocks = num_blocks
+    # Connectors that derive geometry from the config must see the committed
+    # storages the rebuilt views will have.
+    kv_cache_config.kv_cache_tensors = kv_cache.committed_kv_cache_tensors(
+        kv_cache_config,
+        runner.cache_config.get_resolved_kv_cache_layout(),
+        num_blocks,
+    )
+    ensure_kv_transfer_initialized(runner.vllm_config, kv_cache_config)
+    # Config resolution rejects `custom_mem_pool`; a connector that wants the
+    # KV cache in its own pool some other way cannot have it either.
+    if (
+        has_kv_transfer_group()
+        and get_kv_transfer_group().get_mem_pool_context() is not None
+    ):
+        raise ValueError(
+            "The KV connector allocates the KV cache from its own memory pool, "
+            "which the extensible KV cache cannot use. Disable it with "
+            "--no-enable-extensible-kv-cache."
+        )
+    if num_blocks < kv_cache.num_committed_blocks:
+        # Warmup committed more than the measured budget fits (its transient
+        # peak exceeded the reserve); the contents are warmup garbage.
+        logger.info(
+            "Remapping the KV cache from %d to %d blocks to fit the measured budget.",
+            kv_cache.num_committed_blocks,
+            num_blocks,
+        )
+    kv_cache.commit(
+        num_blocks,
+        defragment=runner.vllm_config.kv_transfer_config is not None,
+        shrink=True,
+    )
+    runner.kv_cache_config.num_blocks = kv_cache.num_committed_blocks
+    logger.info(
+        "Extended KV cache to %d blocks (%.2f GiB).",
+        kv_cache.num_committed_blocks,
+        kv_cache.physical_bytes / (1 << 30),
+    )
+    forward_context = runner.compilation_config.static_forward_context
+    kv_caches = kv_cache.committed_views(
+        runner.kv_cache_config,
+        runner.cache_config.get_resolved_kv_cache_layout(),
+        runner.kernel_block_sizes,
+    )
+    if kv_caches is not None:
+        # As in `initialize_kv_cache`, draft groups map their views too.
+        attn_groups = [
+            group
+            for groups in runner.attn_groups
+            + getattr(runner.speculator, "attn_groups", [])
+            for group in groups
+        ]
+        bind_kv_caches(
+            kv_caches,
+            forward_context,
+            runner.kv_cache_config,
+            runner.vllm_config,
+            attn_groups,
+        )
+        runner.kv_caches = [
+            cache for cache in kv_caches.values() if cache.device == runner.device
+        ]
+    else:
+        kv_caches = {
+            layer_name: layer.kv_cache
+            for layer_name, layer in forward_context.items()
+            if isinstance(getattr(layer, "kv_cache", None), torch.Tensor)
+        }
+    runner.kv_connector = get_kv_connector(runner.vllm_config, kv_caches)
+
+
+def measure_kv_cache_bytes(
+    *,
+    init_free_memory: int,
+    free_memory: int,
+    committed_bytes: int,
+    requested_memory: int,
+    transient_peak_bytes: int = 0,
+    margin_floor_bytes: int = KV_CACHE_MARGIN_FLOOR_BYTES,
+    margin_fraction: float = KV_CACHE_MARGIN_FRACTION,
+) -> int:
+    """Bytes the KV cache may occupy given the memory measured after warmup.
+
+    Everything resident except the committed KV prefix is needed by the engine;
+    the cache gets the rest of the budget, capped by the device headroom (free
+    memory plus the committed prefix) less a margin of
+    `max(margin_floor_bytes, margin_fraction * transient_peak_bytes)`. The
+    measured transient_peak_bytes are then kept free in full. The margin
+    guards physical memory, so it never reduces an explicit budget that already
+    leaves that much free.
+    """
+    non_kv_used = init_free_memory - free_memory - committed_bytes
+    headroom = free_memory + committed_bytes
+    margin = max(margin_floor_bytes, int(transient_peak_bytes * margin_fraction))
+    available = min(requested_memory - non_kv_used, headroom - margin)
+    available -= transient_peak_bytes
+    return max(available, 0)
+
+
+def measure_kv_cache_blocks(
+    kv_cache: ExtensibleKVCache,
+    vllm_config: "VllmConfig",
+    device: torch.device,
+    *,
+    init_free_memory: int,
+    requested_memory: int,
+) -> int:
+    """Size an extensible KV cache from the memory free after warmup.
+
+    Keeps free reserved_headroom_bytes, the larger of the profiled peak and
+    warmup's allocated (not reserved) transient peak. Torch's cache is emptied
+    first: its blocks may not fit the large workspaces the first real steps
+    allocate. Then caps torch's caching allocator to leave
+    KV_CACHE_NON_TORCH_RESERVE_BYTES free for allocations outside it.
+    """
+    torch.accelerator.synchronize()
+    gc.collect()
+    torch.accelerator.empty_cache()
+    free_memory, _ = torch.accelerator.get_memory_info(device)
+    transient_peak = kv_cache.reserved_headroom_bytes
+    available = measure_kv_cache_bytes(
+        init_free_memory=init_free_memory,
+        free_memory=free_memory,
+        committed_bytes=kv_cache.physical_bytes,
+        requested_memory=requested_memory,
+        transient_peak_bytes=transient_peak,
+    )
+    kv_cache_bytes = reserve_mm_ipc_gpu_memory(
+        available,
+        vllm_config.model_config.multimodal_config,
+        getattr(vllm_config.parallel_config, "_api_process_count", 1),
+    )
+    num_blocks = kv_cache.blocks_within(
+        kv_cache_bytes, aligned=vllm_config.kv_transfer_config is not None
+    )
+    # Torch may grow into what stays free once the cache is extended, less the
+    # memory kept for multimodal IPC and for allocations outside torch.
+    left_free = (
+        free_memory
+        - (kv_cache.physical_bytes_for(num_blocks) - kv_cache.physical_bytes)
+        - (available - kv_cache_bytes)
+        - KV_CACHE_NON_TORCH_RESERVE_BYTES
+    )
+    torch_memory_limit = torch.accelerator.memory_reserved(device) + max(left_free, 0)
+    kv_cache.set_torch_memory_limit(torch_memory_limit)
+    logger.info(
+        "Memory after warmup: %s GiB free, %s GiB committed to the KV cache, "
+        "%s GiB kept for transient activations; %d KV cache blocks (%s GiB) "
+        "fit within the requested %s GiB. Torch may reserve up to %s GiB.",
+        format_gib(free_memory),
+        format_gib(kv_cache.physical_bytes),
+        format_gib(transient_peak),
+        num_blocks,
+        format_gib(num_blocks * kv_cache.bytes_per_block),
+        format_gib(requested_memory),
+        format_gib(torch_memory_limit),
+    )
+    return num_blocks
+
+
+# A layerwise weight update holds a layer's checkpoint-format copy while weight
+# processing (format conversion, padding, shuffles) allocates around as much
+# again, or twice that, before the result is copied into place.
+WEIGHT_UPDATE_PEAK_LAYER_FACTOR = 3
+
+
+@contextmanager
+def kv_cache_lent_for_weight_update(
+    kv_cache: ExtensibleKVCache | None,
+    model: torch.nn.Module,
+    *,
+    reclaimable: bool,
+    fail_if_short: bool,
+) -> Iterator[None]:
+    """Make room for a layerwise weight update of model within the context.
+
+    Only an extensible KV cache (kv_cache, None for a standard one) leaves
+    the free memory short, since it keeps free just the transients warmup
+    measured. When short, release the cache's physical memory for the
+    context if it holds nothing to keep (reclaimable); else fail before any
+    weight changes, when not even the largest layer fits and fail_if_short,
+    or warn. Torch's memory cap is lifted for the context.
+    """
+    if kv_cache is None:
+        yield
+        return
+    lent = _release_for_weight_update(
+        kv_cache, model, reclaimable=reclaimable, fail_if_short=fail_if_short
+    )
+    # Weight processing allocates through torch, beyond the cap sized for
+    # serving.
+    torch_memory_limit = kv_cache.torch_memory_limit
+    if torch_memory_limit is not None:
+        kv_cache.set_torch_memory_limit(None)
+    try:
+        yield
+    finally:
+        if lent or torch_memory_limit is not None:
+            torch.accelerator.empty_cache()
+        if lent:
+            kv_cache.recommit()
+            if has_kv_transfer_group():
+                get_kv_transfer_group().restore_kv_caches()
+        if torch_memory_limit is not None:
+            kv_cache.set_torch_memory_limit(torch_memory_limit)
+
+
+def _release_for_weight_update(
+    kv_cache: ExtensibleKVCache,
+    model: torch.nn.Module,
+    *,
+    reclaimable: bool,
+    fail_if_short: bool,
+) -> bool:
+    """Whether kv_cache was released to fit the update; see the caller."""
+    if not kv_cache.buffer.physical_bytes:
+        # Released by sleep or release_kv_cache_memory(), which also restores it.
+        return False
+    layer_bytes = max_layer_reload_bytes(model)
+    needed = WEIGHT_UPDATE_PEAK_LAYER_FACTOR * layer_bytes
+    free = (
+        torch.accelerator.get_memory_info()[0]
+        + torch.accelerator.memory_reserved()
+        - torch.accelerator.memory_allocated()
+    )
+    if free >= needed:
+        return False
+    if reclaimable:
+        logger.info(
+            "Releasing the emptied KV cache's memory for a weight update: "
+            "%s GiB free, about %s GiB needed.",
+            format_gib(free),
+            format_gib(needed),
+        )
+        if has_kv_transfer_group():
+            get_kv_transfer_group().unregister_kv_caches()
+        kv_cache.release_physical()
+        return True
+    message = (
+        f"Updating weights layer by layer needs about {format_gib(needed)} GiB "
+        f"of free device memory, at least {format_gib(layer_bytes)} GiB for the "
+        f"largest layer alone, but only {format_gib(free)} GiB is free beside "
+        "the extensible KV cache. Pause generation with the cache cleared "
+        "first, or release it with release_kv_cache_memory(), so that the "
+        "update can borrow its memory; or leave headroom by setting "
+        "gpu_memory_utilization."
+    )
+    if fail_if_short and free < layer_bytes:
+        raise RuntimeError(message)
+    logger.warning(message)
+    return False
diff --git a/vllm/v1/worker/gpu/attn_utils.py b/vllm/v1/worker/gpu/attn_utils.py
index 03a81881b755..bcfaa7585c65 100644
--- a/vllm/v1/worker/gpu/attn_utils.py
+++ b/vllm/v1/worker/gpu/attn_utils.py
@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-from collections.abc import Iterable, Mapping, Sequence
+from collections.abc import Callable, Iterable, Mapping, Sequence
 from contextlib import AbstractContextManager, nullcontext
 from dataclasses import dataclass, replace
 from typing import TYPE_CHECKING, Any, cast
@@ -404,6 +404,7 @@ def init_kv_cache(
     *,
     block_tables: "BlockTables | None" = None,
     attn_groups: Iterable[AttentionGroup] = (),
+    allocate: Callable[[int], torch.Tensor] | None = None,
 ) -> dict[str, Any]:
     allocation_context = kv_cache_allocation_context or nullcontext()
     with allocation_context:
@@ -423,7 +424,21 @@ def init_kv_cache(
                 device,
                 vllm_config.cache_config.get_resolved_kv_cache_layout(),
                 kernel_block_sizes,
+                allocate=allocate,
             )
+    return bind_kv_caches(
+        kv_caches, forward_context, kv_cache_config, vllm_config, attn_groups
+    )
+
+
+def bind_kv_caches(
+    kv_caches: dict[str, torch.Tensor],
+    forward_context: dict[str, Any],
+    kv_cache_config: KVCacheConfig,
+    vllm_config: VllmConfig,
+    attn_groups: Iterable[AttentionGroup] = (),
+) -> dict[str, Any]:
+    """Bind per-layer KV cache views to the attention layers."""
     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
diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py
index d29fd3341bb5..995657462bf2 100644
--- a/vllm/v1/worker/gpu/cudagraph_utils.py
+++ b/vllm/v1/worker/gpu/cudagraph_utils.py
@@ -179,6 +179,9 @@ def __init__(
 
         self.graphs: dict[BatchExecutionDescriptor, torch.cuda.CUDAGraph] = {}
         self.pool = current_platform.get_global_graph_pool() if cudagraph_mode else None
+        # The stream of the last capture; graphs replayed inside these ones (in
+        # their eager breaks) can capture on it to share the pool's free blocks.
+        self.capture_stream: torch.cuda.Stream | None = None
 
         self._graphs_captured = False
 
@@ -425,6 +428,7 @@ def capture(
         self,
         create_forward_fn: CreateForwardFn,
         progress_bar_desc: str = "Capturing CUDA graphs",
+        stream: torch.cuda.Stream | None = None,
     ) -> None:
         """Capture CUDA graphs.
 
@@ -435,9 +439,11 @@ def capture(
                 because attention backends may mutate or lazily initialize
                 metadata during warmup.
             progress_bar_desc: Description shown on the capture progress bar.
+            stream: Capture on this stream instead of a new one.
 
         """
-        with graph_capture(device=self.device), ExitStack() as stack:
+        with graph_capture(self.device, stream) as context, ExitStack() as stack:
+            self.capture_stream = context.stream
             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)
@@ -715,7 +721,7 @@ def create_forward_fn(
                     input_batch,
                     input_block_tables,
                     block_tables.get_dummy_slot_mappings(num_tokens),
-                    cg_mode=CUDAGraphMode.FULL,
+                    desc,
                     for_capture=True,
                 )
                 # Capture with dummy rows marked as padding.
diff --git a/vllm/v1/worker/gpu/dp_utils.py b/vllm/v1/worker/gpu/dp_utils.py
index fb9f4c02d3e4..008c3cafd28d 100644
--- a/vllm/v1/worker/gpu/dp_utils.py
+++ b/vllm/v1/worker/gpu/dp_utils.py
@@ -5,10 +5,12 @@
 from dataclasses import dataclass, replace
 
 import torch
+import torch.distributed as dist
 
 from vllm.config import ParallelConfig
 from vllm.config.compilation import CUDAGraphMode
-from vllm.distributed.parallel_state import get_dp_group
+from vllm.config.fault_tolerance import ft_tp_barrier_required
+from vllm.distributed.parallel_state import get_dp_group, get_tp_group
 from vllm.v1.worker.dp_utils import should_skip_dp_coordination
 from vllm.v1.worker.gpu.cudagraph_utils import (
     BatchExecutionDescriptor,
@@ -32,6 +34,8 @@ class DPSyncState:
     # Agreed uniform decode length. None means the ranks disagreed, which is an
     # answer, not "unknown".
     uniform_token_count: int | None
+    # Agreed query-length bound. None disables varlen decode graphs.
+    max_query_len: int | None
     # Whether the ranks agreed to run eager. A dispatch reusing this must run
     # eager too; no shape was agreed, so picking a graph per rank would diverge.
     eager: bool
@@ -59,9 +63,6 @@ def sync_cudagraph_and_dp_padding(
 ) -> tuple[BatchExecutionDescriptor, DPSyncState | None]:
     """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.
-
     Returns (synced_batch_desc, sync). `sync` is None when no rank has work.
     """
     assert dp_size > 1, "DP size must be greater than 1"
@@ -78,6 +79,24 @@ def sync_cudagraph_and_dp_padding(
     else:
         get_dp_group().cpu_all_reduce_(tensor)
 
+    # Full width, dead slots included: returned in DPSyncState, whose
+    # consumers index it by original dp_rank.
+    num_tokens_across_dp_full = tensor[0]
+
+    if parallel_config is not None and parallel_config.enable_fault_tolerance:
+        # Per-step barrier over the TP cpu group: a faulted sibling stops
+        # arriving, so survivors fail here on the host instead of leaving
+        # an orphaned TP collective running on device.
+        if ft_tp_barrier_required(parallel_config):
+            dist.barrier(group=get_tp_group().cpu_group)
+
+        if dead_dp_ranks := get_dp_group().dead_dp_ranks:
+            # Drop the failed ranks' columns so the min / all(==1) agreements
+            # below only see ranks that are still running.
+            tensor = tensor[:, [r for r in range(dp_size) if r not in dead_dp_ranks]]
+
+    # With dead DP ranks (FT), the dead ranks' columns were dropped above:
+    # dim 1 (per-rank) of these rows covers live ranks only.
     num_tokens_across_dp = tensor[0]
     cg_mode_across_dp = tensor[1]
     uniform_token_counts_across_dp = tensor[2]
@@ -94,6 +113,11 @@ def sync_cudagraph_and_dp_padding(
     ):
         synced_uniform_token_count = None
 
+    # Varlen decode graphs need the same query-length bound on every rank.
+    synced_max_query_len: int | None = None
+    if bool(torch.all(max_query_lens_across_dp != -1)):
+        synced_max_query_len = int(max_query_lens_across_dp.max())
+
     if torch.all(num_tokens_across_dp == 0).item():
         synced_desc = BatchExecutionDescriptor(
             cg_mode=CUDAGraphMode.NONE, num_tokens=0, num_reqs=0
@@ -145,9 +169,10 @@ def sync_cudagraph_and_dp_padding(
             ubatch_num_tokens = ubatch_desc.num_tokens
             return ubatch_desc, DPSyncState(
                 num_tokens_across_dp=torch.full_like(
-                    num_tokens_across_dp, ubatch_num_tokens
+                    num_tokens_across_dp_full, ubatch_num_tokens
                 ),
                 uniform_token_count=synced_uniform_token_count,
+                max_query_len=synced_max_query_len,
                 eager=ubatch_desc.cg_mode == CUDAGraphMode.NONE,
                 num_reqs=int(num_reqs_across_dp.max()),
                 skip_drafts=skip_drafts,
@@ -165,8 +190,9 @@ def sync_cudagraph_and_dp_padding(
                 num_active_loras=desired_batch_desc.num_active_loras,
             ),
             DPSyncState(
-                num_tokens_across_dp=num_tokens_across_dp,
+                num_tokens_across_dp=num_tokens_across_dp_full,
                 uniform_token_count=synced_uniform_token_count,
+                max_query_len=synced_max_query_len,
                 eager=True,
                 num_reqs=int(num_reqs_across_dp.max()),
                 skip_drafts=skip_drafts,
@@ -179,12 +205,6 @@ def sync_cudagraph_and_dp_padding(
     )
     synced_num_tokens = int(num_tokens_across_dp.max())
 
-    # Varlen decode graphs are selected by the query-length bound, so ranks must agree
-    # on it or they pad to different token counts below.
-    synced_max_query_len: int | None = None
-    if bool(torch.all(max_query_lens_across_dp != -1)):
-        synced_max_query_len = int(max_query_lens_across_dp.max())
-
     # Dispatch for the final synced values, use num_reqs instead of synced_num_reqs
     # so we don't perform request padding for PIECEWISE graphs.
     # num_active_loras is per-rank and doesn't need cross-rank agreement.
@@ -197,11 +217,12 @@ def sync_cudagraph_and_dp_padding(
     )
 
     # Update num_tokens_across_dp to reflect padded size.
-    num_tokens_across_dp[:] = synced_desc.num_tokens
+    num_tokens_across_dp_full[:] = synced_desc.num_tokens
 
     return synced_desc, DPSyncState(
-        num_tokens_across_dp=num_tokens_across_dp,
+        num_tokens_across_dp=num_tokens_across_dp_full,
         uniform_token_count=synced_uniform_token_count,
+        max_query_len=synced_max_query_len,
         eager=False,
         num_reqs=int(num_reqs_across_dp.max()),
         skip_drafts=skip_drafts,
@@ -243,12 +264,13 @@ def dispatch_cg_and_sync_dp(
         dp_rank: This rank's index in the DP group.
         max_query_len: Upper bound on per-request query length, for selecting
             varlen decode graphs. None, as for a batch with a prefill, keeps
-            the batch out of graphs that constrain it.
+            the batch out of graphs that constrain it. Reusing a sync replaces
+            this with `dp_sync_state.max_query_len`, including when it is None.
         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.
+        parallel_config: Only needed for FT coordination and to decide whether
+            to microbatch, so callers that do neither 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
@@ -286,7 +308,9 @@ def dispatch_cg_and_sync_dp(
             num_tokens,
             dp_sync_state.uniform_token_count if dp_sync_state else uniform_token_count,
             num_active_loras=num_active_loras,
-            max_query_len=max_query_len,
+            max_query_len=(
+                dp_sync_state.max_query_len if dp_sync_state else max_query_len
+            ),
         )
 
     if dp_size == 1:
@@ -320,8 +344,8 @@ def dispatch_cg_and_sync_dp(
         dp_size,
         dp_rank,
         max_query_len=max_query_len,
-        num_active_loras=num_active_loras,
         parallel_config=parallel_config,
+        num_active_loras=num_active_loras,
         allow_ubatching=allow_ubatching,
         uniform_decode=uniform_decode,
         allow_skip_drafts=allow_skip_drafts,
diff --git a/vllm/v1/worker/gpu/input_batch.py b/vllm/v1/worker/gpu/input_batch.py
index 6d9732b67ab1..b2040891615d 100644
--- a/vllm/v1/worker/gpu/input_batch.py
+++ b/vllm/v1/worker/gpu/input_batch.py
@@ -577,6 +577,9 @@ def _post_update_kernel(
     all_token_ids_ptr,
     all_token_ids_stride,
     total_len_ptr,
+    draft_tokens_ptr,
+    draft_tokens_out_ptr,
+    num_draft_tokens,
 ):
     req_id = tl.program_id(0)
     req_state_idx = tl.load(idx_mapping_ptr + req_id)
@@ -584,6 +587,13 @@ def _post_update_kernel(
         # Filter rows with negative index entries.
         return
 
+    if draft_tokens_ptr is not None:
+        for i in range(num_draft_tokens):
+            token_id = tl.load(draft_tokens_ptr + req_id * num_draft_tokens + i)
+            tl.store(
+                draft_tokens_out_ptr + req_state_idx * num_draft_tokens + i, token_id
+            )
+
     total_len = tl.load(total_len_ptr + req_state_idx)
     num_sampled = tl.load(num_sampled_ptr + req_id)
     if num_sampled > 0:
@@ -644,8 +654,17 @@ def post_update(
     all_token_ids: torch.Tensor,
     # [max_num_reqs]
     total_len: torch.Tensor,
+    # [num_reqs, num_speculative_steps], contiguous
+    draft_tokens: torch.Tensor | None = None,
+    # [max_num_reqs, num_speculative_steps], contiguous; required with draft_tokens
+    draft_tokens_out: torch.Tensor | None = None,
 ) -> None:
     num_reqs = idx_mapping.shape[0]
+    if draft_tokens is not None:
+        # The kernel uses num_speculative_steps as the row stride of both.
+        assert draft_tokens_out is not None
+        assert draft_tokens.shape[1] == draft_tokens_out.shape[1]
+        assert draft_tokens.is_contiguous() and draft_tokens_out.is_contiguous()
     _post_update_kernel[(num_reqs,)](
         idx_mapping,
         num_computed_tokens,
@@ -660,6 +679,9 @@ def post_update(
         all_token_ids,
         all_token_ids.stride(0),
         total_len,
+        draft_tokens,
+        draft_tokens_out,
+        draft_tokens.shape[1] if draft_tokens is not None else 0,
         num_warps=1,
     )
 
diff --git a/vllm/v1/worker/gpu/lora_utils.py b/vllm/v1/worker/gpu/lora_utils.py
index eb81a83662a1..738221317ef3 100644
--- a/vllm/v1/worker/gpu/lora_utils.py
+++ b/vllm/v1/worker/gpu/lora_utils.py
@@ -96,9 +96,12 @@ def make_lora_inputs(
         req_ids: list[str],
         idx_mapping: np.ndarray,
         num_scheduled_tokens: np.ndarray,
+        num_logits: np.ndarray,
     ) -> tuple[tuple[int, ...], tuple[int, ...], set[LoRARequest]]:
         lora_ids = self.lora_ids[idx_mapping]
-        prompt_lora_mapping = tuple(lora_ids)
+        # One lm_head id per logits row: 1 + num_drafts rows per request
+        # under spec decode.
+        prompt_lora_mapping = tuple(lora_ids.repeat(num_logits))
         token_lora_mapping = tuple(lora_ids.repeat(num_scheduled_tokens))
         active_lora_requests: set[LoRARequest] = self.get_activate_loras(req_ids)
         return prompt_lora_mapping, token_lora_mapping, active_lora_requests
diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py
index 025fc08567d4..665e8347329a 100644
--- a/vllm/v1/worker/gpu/model_runner.py
+++ b/vllm/v1/worker/gpu/model_runner.py
@@ -91,6 +91,7 @@
 )
 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.extensible_kv_cache import ExtensibleKVCache
 from vllm.v1.worker.gpu import pcp_manager as pcp
 from vllm.v1.worker.gpu.async_utils import (
     AsyncOutput,
@@ -148,7 +149,7 @@
 from vllm.v1.worker.gpu.pp_utils import PPHandler
 from vllm.v1.worker.gpu.sample.batch_shard import (
     BatchSharder,
-    all_to_all_logits,
+    all_to_all_vocab_shards,
     batch_sharded_sampling_enabled,
     gather_sampler_output,
 )
@@ -164,9 +165,7 @@
     resolve_adaptive_cudagraph_mode,
 )
 from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import (
-    requires_aux_hidden_states,
     set_eagle3_aux_hidden_state_layers,
-    verify_supports_aux_hidden_states_over_pp,
 )
 from vllm.v1.worker.gpu.spec_decode.rejection_sampler import (
     RejectionSampler,
@@ -178,10 +177,7 @@
     get_drafter_hidden_states,
 )
 from vllm.v1.worker.gpu.states import RequestState
-from vllm.v1.worker.gpu.structured_outputs import (
-    StructuredOutputsWorker,
-    grammar_invalid_drafts,
-)
+from vllm.v1.worker.gpu.structured_outputs import StructuredOutputsWorker
 from vllm.v1.worker.gpu.ubatch_utils import (
     UBatchRunner,
     UBatchState,
@@ -211,7 +207,8 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device):
         self.scheduler_config = vllm_config.scheduler_config
         self.speculative_config = vllm_config.speculative_config
         self._draft_workspace_lane = int(
-            self.speculative_config is not None and self.speculative_config.use_dspark()
+            self.speculative_config is not None
+            and self.speculative_config.uses_separate_workspace_lane
         )
         self.observability_config = vllm_config.observability_config
         self.jit_warmup_registry = JitWarmupRegistry(vllm_config)
@@ -282,7 +279,8 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device):
         self.num_speculative_steps = vllm_config.num_speculative_tokens
         num_prefill_lookahead = max(1, vllm_config.num_prefill_lookahead_tokens)
         use_dense_all_token_ids = (
-            self.speculative_config is not None and self.speculative_config.use_ngram()
+            self.speculative_config is not None
+            and self.speculative_config.speculator_config.prefer_dense_token_history
         )
 
         # General request states.
@@ -306,8 +304,8 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device):
                     self.vllm_config, self.device, self.req_states
                 )
 
-            self.use_aux_hidden_state_outputs = requires_aux_hidden_states(
-                self.speculative_config
+            self.use_aux_hidden_state_outputs = (
+                self.speculative_config.uses_aux_hidden_states
             )
 
         # Draft tokens propagation - for spec-dec + struct outputs.
@@ -414,13 +412,10 @@ def load_model(
 
             if self.use_aux_hidden_state_outputs:
                 assert self.speculative_config is not None
-                set_eagle3_aux_hidden_state_layers(self.model, self.speculative_config)
-                if self.use_pp:
-                    assert self.speculative_config.method is not None
-                    verify_supports_aux_hidden_states_over_pp(
-                        self.model, self.speculative_config.method
-                    )
-                    assert self.pp_handler is not None
+                set_eagle3_aux_hidden_state_layers(
+                    self.model, self.speculative_config, self.use_pp
+                )
+                if self.pp_handler is not None:
                     self.pp_handler.configure_aux_hidden_state_relay(self.model)
             if isinstance(self.speculator, DraftModelSpeculator):
                 with use_workspace_lane(self._draft_workspace_lane):
@@ -454,6 +449,11 @@ def load_model(
             + self.model_state.num_new_sampled_tokens_per_step
         )
 
+        draft_logits = (
+            self.speculator.draft_logits
+            if isinstance(self.speculator, DraftModelSpeculator)
+            else None
+        )
         if self.parallel_config.enable_batch_sharded_sampling:
             if batch_sharded_sampling_enabled(self.vllm_config, self.model):
                 self.batch_sharder = BatchSharder(
@@ -464,6 +464,7 @@ def load_model(
                         self.speculative_config is not None
                         and self.speculative_config.enable_adaptive_verification
                     ),
+                    draft_logits=draft_logits,
                 )
                 logger.info("Batch-sharded sampling enabled.")
             else:
@@ -521,13 +522,8 @@ def load_model(
                 )
             custom = self.model_state.custom_sampler(self.sampler)
 
-            draft_logits_num_shards = (
-                self.speculator.draft_logits_num_shards
-                if isinstance(self.speculator, DraftModelSpeculator)
-                else 1
-            )
             if custom:
-                assert draft_logits_num_shards == 1, (
+                assert draft_logits is None or draft_logits.row_mapping is None, (
                     "A custom rejection sampler cannot read batch-sharded draft logits"
                 )
                 self.vllm_config._check_supports_watermarking(custom_sampler=True)
@@ -540,7 +536,6 @@ def load_model(
                     watermark_key=speculative_target_watermark_key(
                         self.vllm_config.watermark_config
                     ),
-                    draft_logits_num_shards=draft_logits_num_shards,
                 )
             self.prompt_logprobs_worker = PromptLogprobsWorker(
                 self.max_num_reqs,
@@ -647,11 +642,20 @@ def initialize_kv_cache(
         kv_cache_config: KVCacheConfig,
         is_profiling: bool = False,
         kv_cache_allocation_context: AbstractContextManager | None = None,
+        extensible: bool = False,
     ) -> None:
+        """Initialize the KV cache and everything that depends on its layout.
+
+        With `extensible=True`, num_blocks is a capacity: address space is
+        reserved for it but only the null block is committed before capture.
+        """
         # GPUWorker finalizes the PD interleave before KV cache initialization.
         self.cp_interleave = self.parallel_config.cp_kv_cache_interleave_size
         kv_cache_config = deepcopy(kv_cache_config)
         self.kv_cache_config = kv_cache_config
+        self.extensible_kv_cache: ExtensibleKVCache | None = None
+        if extensible:
+            self.extensible_kv_cache = ExtensibleKVCache(kv_cache_config, self.device)
 
         block_table_max_model_len = self.max_model_len
         if self.is_encoder_decoder:
@@ -813,6 +817,10 @@ def initialize_kv_cache(
             # to its own attention support.
             self.speculator.init_cudagraph_manager(cudagraph_mode)
 
+        allocate_fn = None
+        if self.extensible_kv_cache is not None:
+            allocate_fn = self.extensible_kv_cache.allocate
+
         # Capture warmup providers that depend on allocated KV-cache strides.
         with self.jit_warmup_registry.activate():
             kv_caches_dict = init_kv_cache(
@@ -830,6 +838,7 @@ def initialize_kv_cache(
                     + getattr(self.speculator, "attn_groups", [])
                     for g in groups
                 ),
+                allocate=allocate_fn,
             )
         self.kv_caches = [
             cache for cache in kv_caches_dict.values() if cache.device == self.device
@@ -842,9 +851,11 @@ def initialize_kv_cache(
 
             # 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
-                )
+                if self.aux_output_connector is None:
+                    self.aux_output_connector = get_aux_output_connector(
+                        self.model, self.vllm_config
+                    )
+                self.aux_output_connector.initialize_output_storage(kv_cache_config)
 
     def _init_kv_zero_meta(self) -> None:
         """Build KV-block zeroing metadata; invoked from gpu_worker."""
@@ -1013,13 +1024,15 @@ def _dummy_run(
 
     @torch.inference_mode()
     def _dummy_sampler_run(self, hidden_states: torch.Tensor) -> None:
+        assert self.sampler is not None
+        if self.model_state.dummy_sampler_run(self.sampler, hidden_states):
+            return
         num_reqs = hidden_states.shape[0]
         logits = self.model.compute_logits(hidden_states)
         dummy_input_batch = InputBatch.make_dummy(
             num_reqs, num_reqs, self.input_buffers
         )
 
-        assert self.sampler is not None
         # Default pass: compiles the kernels the first real step uses (logits in
         # the model dtype, no per-request processing).
         self.sampler(logits, dummy_input_batch)
@@ -1042,6 +1055,10 @@ def _dummy_pooler_run(self, hidden_states: torch.Tensor) -> None:
 
     @torch.inference_mode()
     def profile_run(self, randomize_inputs: bool = False) -> None:
+        if self.vllm_config.aux_output_config.enabled:
+            self.aux_output_connector = get_aux_output_connector(
+                self.model, self.vllm_config
+            )
         if self.supports_mm_inputs and self.is_first_pp_rank:
             mm_config = self.model_config.multimodal_config
             if mm_config is not None and not mm_config.skip_mm_profiling:
@@ -1090,6 +1107,9 @@ def release_late_interaction_query_cache(self, query_keys: list[str]) -> None:
     @torch.inference_mode()
     def profile_cudagraph_memory(self) -> int:
         """Estimate the GPU memory required to capture CUDA graphs."""
+        if self.cache_config.enable_extensible_kv_cache:
+            # The extensible KV cache measures the capture footprint after warmup.
+            return 0
         # Runs before the KV cache exists, outside compile_or_warm_up_model.
         with self.eplb.suppress():
             return _profile_cudagraph_memory(self)
@@ -1155,6 +1175,7 @@ def capture_model(self, *, profile_only: bool = False) -> int:
                         self.block_tables,
                         self.attn_groups,
                         self.kv_cache_config,
+                        self.cudagraph_manager.capture_stream,
                     )
                     if self.speculator is not None:
                         with use_workspace_lane(self._draft_workspace_lane):
@@ -1196,6 +1217,7 @@ def _remove_request(self, req_id: str) -> bool:
         # Call model_state.remove_request *before* req_states.remove_request
         # so the model_state can still look up the slot index.
         self.model_state.remove_request(req_id)
+        self.draft_tokens_handler.discard(req_id)
         req_idx = self.req_states.remove_request(req_id)
         if req_idx is None:
             return False
@@ -1233,9 +1255,7 @@ def update_pp_decode_requests(self):
         # For non-last PP ranks, update decode requests with sampler output from
         # the prior step in which they were scheduled (pp_size steps ago).
         if self.pp_handler is not None:
-            outputs = self.pp_handler.get_prev_sampled_outputs(
-                self.req_states.draft_tokens
-            )
+            outputs = self.pp_handler.get_prev_sampled_outputs()
             if outputs is not None:
                 self.postprocess_sampled(**outputs)
 
@@ -1291,8 +1311,8 @@ def add_requests(self, scheduler_output: SchedulerOutput) -> None:
         if scheduler_output.scheduled_new_reqs:
             self.req_states.apply_staged_writes()
             self.model_state.apply_staged_writes()
-        if self.sampler is not None:
-            self.sampler.apply_staged_writes()
+            if self.sampler is not None:
+                self.sampler.apply_staged_writes()
 
     def update_requests(self, scheduler_output: SchedulerOutput) -> None:
         # Add new blocks and update num_computed_tokens for the existing requests.
@@ -1676,13 +1696,16 @@ def sample(
             # produces full-vocab logits for just the locally-owned requests.
             sample_hidden_states = hidden_states[sorted_logits_indices]
             local_logits = self.model.compute_logits_local(sample_hidden_states)
-            logits = all_to_all_logits(local_logits, shard_metadata)
+            logits = all_to_all_vocab_shards(
+                local_logits,
+                shard_metadata.num_logits_per_rank,
+                shard_metadata.num_local_logits,
+            )
             logits = logits[:, : self.vocab_size]
         else:
             sample_hidden_states = hidden_states[input_batch.logits_indices]
-            logits = self.model.compute_logits(sample_hidden_states)
+            logits = self.model_state.compute_sample_logits(sample_hidden_states)
 
-        invalid_drafts = 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:
@@ -1694,11 +1717,6 @@ def sample(
                 grammar_output.structured_output_request_ids,
                 grammar_output.grammar_bitmask,
             )
-            invalid_drafts = grammar_invalid_drafts(
-                input_batch,
-                grammar_output.structured_output_request_ids,
-                grammar_output.num_acceptable_drafts,
-            )
 
         sampler_output: SamplerOutput | None
         if input_batch.num_reqs == 0:
@@ -1717,7 +1735,6 @@ def sample(
                 input_batch,
                 # Draft logits are needed for probabilistic rejection sampling.
                 self.speculator.draft_logits,
-                invalid_drafts,
             )
 
         if shard_metadata is not None:
@@ -1732,12 +1749,7 @@ def sample(
                 local_batch=input_batch,
                 gather_num_nans=self.sampler.compute_nans,
                 logprobs_dims=self.sampler.get_logprobs_dims(
-                    global_input_batch.idx_mapping_np,
-                    # Rejection sampler does not return logprob token ids.
-                    include_token_ids=(
-                        global_input_batch.num_draft_tokens == 0
-                        or self.rejection_sampler is None
-                    ),
+                    global_input_batch.idx_mapping_np
                 ),
             )
 
@@ -1751,6 +1763,8 @@ def postprocess_sampled(
         num_sampled: torch.Tensor,
         num_rejected: torch.Tensor,
         query_start_loc: torch.Tensor | None = None,
+        # Drafts received from the last PP rank (non-last ranks only).
+        draft_tokens: torch.Tensor | None = None,
     ) -> None:
         # Update the number of computed tokens.
         output_bin_counts = None
@@ -1768,6 +1782,8 @@ def postprocess_sampled(
             query_start_loc,
             self.req_states.all_token_ids.gpu,
             self.req_states.total_len.gpu,
+            draft_tokens,
+            self.req_states.draft_tokens,
         )
 
         self.model_state.postprocess_state(
@@ -1796,6 +1812,7 @@ def execute_model(
         randomize_inputs: bool = False,
     ) -> ModelRunnerOutput | IntermediateTensors | None:
         if not dummy_run:
+            self.draft_tokens_handler.begin_step(scheduler_output)
             # Update the request states.
             self.update_pp_decode_requests()
             self.finish_requests(scheduler_output)
@@ -1903,6 +1920,7 @@ def execute_model(
                     input_batch.req_ids,
                     input_batch.idx_mapping_np,
                     input_batch.num_scheduled_tokens,
+                    np.diff(input_batch.cu_num_logits_np),
                 )
                 self._set_active_loras(*lora_inputs)
         else:
@@ -1935,11 +1953,15 @@ def execute_model(
                     input_batch, valid_dummy_state_slots
                 )
                 if context_len:
+                    if self.extensible_kv_cache is not None:
+                        num_blocks = self.extensible_kv_cache.num_committed_blocks
+                    else:
+                        num_blocks = self.kv_cache_config.num_blocks
                     set_dummy_context(
                         input_batch,
                         self.block_tables,
                         context_len,
-                        self.kv_cache_config.num_blocks,
+                        num_blocks,
                         self.max_model_len,
                         input_block_tables=block_tables,
                     )
@@ -1968,7 +1990,8 @@ def execute_model(
                 input_batch,
                 block_tables,
                 slot_mappings,
-                cg_mode=batch_desc.cg_mode,
+                batch_desc,
+                dp_sync_state,
                 for_capture=dummy_run and batch_desc.cg_mode == CUDAGraphMode.FULL,
             )
         elif not (dummy_run and skip_attn_for_dummy_run):
@@ -2195,7 +2218,9 @@ def sample_tokens(
             if self.pcp_manager is not None:
                 input_batch = self.pcp_manager.global_batch
             assert self.pp_handler is not None
-            all_decode_next = self.pp_handler.receive(input_batch)
+            all_decode_next = self.pp_handler.receive(
+                input_batch, recv_drafts=num_spec_tokens > 0
+            )
             # Optimistically update num_computed_tokens for entire batch here.
             # Will be adjusted for rejections if necessary in update_requests.
             self.postprocess_num_computed_tokens(input_batch)
@@ -2315,16 +2340,12 @@ 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
+                self.speculator.prepare_propose(
+                    input_batch, num_sampled, num_rejected, self.sampler
                 )
             spec_hidden_states = get_drafter_hidden_states(
                 self.model, draft_hidden_states
             )
-            if isinstance(self.sampler, GPUWatermarkSampler):
-                self.speculator.prepare_watermarking(
-                    self.sampler, input_batch.idx_mapping
-                )
             with use_workspace_lane(self._draft_workspace_lane):
                 draft_tokens = self.speculator.propose(
                     input_batch,
@@ -2357,14 +2378,13 @@ def sample_tokens(
         if self.num_speculative_steps > 0:
             # Spec-decode and diffusion LLMs both use draft tokens but the latter does
             # not have a speculator (i.e. self.speculator is None)
+            draft_tokens = self.req_states.draft_tokens[input_batch.idx_mapping]
             self.draft_tokens_handler.set_draft_tokens(
-                input_batch,
-                self.req_states.draft_tokens[input_batch.idx_mapping, :num_spec_tokens],
+                input_batch, draft_tokens[:, :num_spec_tokens]
             )
-            if self.pp_handler is not None:
-                self.pp_handler.broadcast_drafts(
-                    self.req_states.draft_tokens, input_batch
-                )
+            # Drafts aren't used when none were proposed this step (e.g. skip_drafts).
+            if self.pp_handler is not None and num_spec_tokens > 0:
+                self.pp_handler.broadcast_drafts(draft_tokens)
 
         if defer_kv_finalize:
             self.kv_connector.finalize_step()
@@ -2463,6 +2483,11 @@ def shutdown(self) -> None:
         if hasattr(self, "model"):
             clear_layer_kv_caches(self.model.modules())
             del self.model
+        # Its views do not keep it alive, and the runner can outlive shutdown
+        # (an in-process engine), so unmap the extensible KV cache explicitly.
+        if (kv_cache := getattr(self, "extensible_kv_cache", None)) is not None:
+            kv_cache.free()
+            self.extensible_kv_cache = None
 
         gc.collect()
         torch.accelerator.empty_cache()
diff --git a/vllm/v1/worker/gpu/model_states/interface.py b/vllm/v1/worker/gpu/model_states/interface.py
index 3db3fc9e7ae6..fa49f470b1e3 100644
--- a/vllm/v1/worker/gpu/model_states/interface.py
+++ b/vllm/v1/worker/gpu/model_states/interface.py
@@ -139,8 +139,10 @@ def capture_inner_cudagraphs(
         block_tables: BlockTables,
         attn_groups: list[list[AttentionGroup]],
         kv_cache_config: KVCacheConfig,
+        stream: torch.cuda.Stream | None = None,
     ) -> None:
-        """Capture the CUDA graphs this state runs inside the model's forward."""
+        """Capture the CUDA graphs this state runs inside the model's forward,
+        on ``stream``: the model graphs' capture stream."""
         return None
 
     def get_additional_cg_support(self) -> tuple[AttentionCGSupport, str | None]:
@@ -251,6 +253,17 @@ def custom_sampler(self, sampler: Any) -> tuple[Any, Any] | None:
         """
         return None
 
+    def compute_sample_logits(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        """Project the sampled hidden states to the logits fed to the sampler."""
+        return cast(Any, self.model).compute_logits(hidden_states)
+
+    def dummy_sampler_run(self, sampler: Any, hidden_states: torch.Tensor) -> bool:
+        """Run a model-specific sampler profiling pass.
+
+        Return ``True`` to skip the default dummy sampler run.
+        """
+        return False
+
     num_new_sampled_tokens_per_step: int = 1
     """New tokens sampled on each decode step 
     (excluding accepted draft tokens, a.k.a num bonus tokens)."""
diff --git a/vllm/v1/worker/gpu/pcp_manager.py b/vllm/v1/worker/gpu/pcp_manager.py
index a906993ec2bf..4e54eb312bb5 100644
--- a/vllm/v1/worker/gpu/pcp_manager.py
+++ b/vllm/v1/worker/gpu/pcp_manager.py
@@ -20,6 +20,7 @@
 
 if TYPE_CHECKING:
     from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor
+    from vllm.v1.worker.gpu.spec_decode.draft_logits import DraftLogitsCache
 
 logger = init_logger(__name__)
 
@@ -174,25 +175,12 @@ def validate_config(
         if vllm_config.lora_config is not None:
             raise NotImplementedError("MRV2 PCP does not support LoRA 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)
-                    and speculative_config.draft_model_config.use_mla
-                ):
-                    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."
-                )
+        if speculative_config is not None and (
+            reason := speculative_config.speculator_config.pcp_unsupported_reason(
+                parallel_config
+            )
+        ):
+            raise NotImplementedError(reason)
         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:
@@ -883,7 +871,7 @@ def sync_drafts(
         self,
         draft_tokens: torch.Tensor,
         idx_mapping: torch.Tensor,
-        draft_logits: torch.Tensor | None,
+        draft_logits: "DraftLogitsCache | None",
     ) -> None:
         """Broadcast PCP rank 0's drafts and draft logits.
 
@@ -894,9 +882,9 @@ def sync_drafts(
         pcp_group.broadcast(draft_tokens, src=0)
         if draft_logits is not None:
             # Draft logits are indexed by persistent request slots.
-            logits = draft_logits[idx_mapping]
+            logits = draft_logits.values[idx_mapping]
             pcp_group.broadcast(logits, src=0)
-            draft_logits[idx_mapping] = logits
+            draft_logits.values[idx_mapping] = logits
 
     def restore_for_sampling(
         self, hidden_states: torch.Tensor, aux_hidden_states: list[torch.Tensor] | None
diff --git a/vllm/v1/worker/gpu/pp_utils.py b/vllm/v1/worker/gpu/pp_utils.py
index 818d80759a7b..b70b30c18a16 100644
--- a/vllm/v1/worker/gpu/pp_utils.py
+++ b/vllm/v1/worker/gpu/pp_utils.py
@@ -22,7 +22,7 @@ class PendingRecv:
 
     event: torch.cuda.Event
 
-    sampled_tokens: torch.Tensor  # [num_reqs, max_sample_len]
+    sampled_tokens: torch.Tensor  # [num_reqs, sample_len]
     num_sampled: torch.Tensor  # [num_reqs]
     num_rejected: torch.Tensor  # [num_reqs]
     idx_mapping: torch.Tensor  # [num_reqs]
@@ -94,11 +94,17 @@ def __init__(
         self.broadcast_group = get_pp_group().make_sibling_device_group(
             group_desc="pp_broadcast"
         )
+        # Whether sampled tokens were broadcast this step.
+        self.sampled_tokens_broadcast = False
         self.aux_hidden_state_relay_keys: tuple[str, ...] = ()
 
     def on_req_idx_freed(self, req_idx: int) -> None:
         self.req_idx_gen_np[req_idx] += 1
 
+    def _sample_len(self, input_batch: InputBatch) -> int:
+        # Width of the sampled-token broadcast. Must be the same on all ranks.
+        return self.max_sample_len if input_batch.num_draft_tokens > 0 else 1
+
     def configure_aux_hidden_state_relay(self, model: torch.nn.Module) -> None:
         from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import (
             aux_hidden_state_relay_keys,
@@ -122,9 +128,7 @@ def relay_aux_hidden_states(
             }
         )
 
-    def get_prev_sampled_outputs(
-        self, draft_tokens_to_update: torch.Tensor | None = None
-    ) -> dict[str, torch.Tensor] | None:
+    def get_prev_sampled_outputs(self) -> dict[str, torch.Tensor | None] | None:
         """Consume the entry from pp_size steps ago and wait for its recv event,
         then filter out entries whose request was freed since `receive`.
         """
@@ -150,45 +154,29 @@ def get_prev_sampled_outputs(
             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:
-            draft_tokens = slot.draft_tokens
-            draft_idx_mapping = slot.idx_mapping
-            if exclude_mask.any():
-                keep = ~exclude_mask
-                keep_t = torch.as_tensor(keep, device=self.device)
-                draft_tokens = draft_tokens[keep_t]
-                draft_idx_mapping = async_tensor_h2d(
-                    slot.idx_mapping_np[keep], device=self.device
-                )
-            draft_tokens_to_update[draft_idx_mapping] = draft_tokens
-
         return dict(
             sampled_tokens=slot.sampled_tokens,
             num_sampled=slot.num_sampled,
             num_rejected=slot.num_rejected,
             idx_mapping=idx_mapping,
+            draft_tokens=slot.draft_tokens,
         )
 
-    def broadcast_drafts(
-        self, draft_tokens: torch.Tensor, input_batch: InputBatch
-    ) -> None:
+    def broadcast_drafts(self, draft_tokens: torch.Tensor) -> None:
         """Broadcast draft proposals so non-last ranks can embed real token ids."""
         assert self.is_last_rank
-        if compute_need_sampled_mask(input_batch) is None:
-            return
-        with torch.cuda.stream(self.broadcast_stream):
-            self.broadcast_stream.wait_stream(self.main_stream)
-            send = draft_tokens[input_batch.idx_mapping].contiguous()
-            # Must record the idx_mapping tensor since it was allocated
-            # on the main stream.
-            input_batch.idx_mapping.record_stream(self.broadcast_stream)
-            torch.distributed.broadcast(
-                send, src=self.last_rank, group=self.broadcast_group
-            )
+        if self.sampled_tokens_broadcast:
+            with torch.cuda.stream(self.broadcast_stream):
+                self.broadcast_stream.wait_stream(self.main_stream)
+                torch.distributed.broadcast(
+                    draft_tokens, src=self.last_rank, group=self.broadcast_group
+                )
+                draft_tokens.record_stream(self.broadcast_stream)
 
-    def receive(self, input_batch: InputBatch) -> bool:
+    def receive(self, input_batch: InputBatch, recv_drafts: bool) -> bool:
         """Returns True iff sampled tokens need to be gathered from *all*
-        requests in the batch."""
+        requests in the batch. `recv_drafts` must match whether the last rank
+        calls `broadcast_drafts` for this step."""
         assert not self.is_last_rank
         need_sampled_mask = compute_need_sampled_mask(input_batch)
         if need_sampled_mask is None:
@@ -202,8 +190,9 @@ def receive(self, input_batch: InputBatch) -> bool:
         num_reqs = input_batch.num_reqs
         with torch.cuda.stream(self.broadcast_stream):
             self.broadcast_stream.wait_stream(self.main_stream)
+            sample_len = self._sample_len(input_batch)
             sampled_tokens = torch.empty(
-                num_reqs, self.max_sample_len, dtype=torch.int64, device=self.device
+                num_reqs, sample_len, dtype=torch.int64, device=self.device
             )
             combined = torch.empty(2, num_reqs, dtype=torch.int32, device=self.device)
             torch.distributed.broadcast(
@@ -213,12 +202,9 @@ def receive(self, input_batch: InputBatch) -> bool:
                 combined, src=self.last_rank, group=self.broadcast_group
             )
             draft_tokens = None
-            if self.num_speculative_steps > 0:
-                draft_tokens = torch.empty(
-                    num_reqs,
-                    self.num_speculative_steps,
-                    dtype=torch.int64,
-                    device=self.device,
+            if recv_drafts:
+                draft_tokens = sampled_tokens.new_empty(
+                    num_reqs, self.num_speculative_steps
                 )
                 torch.distributed.broadcast(
                     draft_tokens, src=self.last_rank, group=self.broadcast_group
@@ -254,21 +240,21 @@ def broadcast(
         assert self.is_last_rank
         if compute_need_sampled_mask(input_batch) is None:
             # No request needs sampled outputs for a subsequent decode step.
+            self.sampled_tokens_broadcast = False
             return
 
         assert sampled_token_ids.dtype == torch.int64
+        self.sampled_tokens_broadcast = True
 
         if current_platform.is_xpu():
             self.main_stream.synchronize()
 
         with torch.cuda.stream(self.broadcast_stream):
             self.broadcast_stream.wait_stream(self.main_stream)
-            send_tokens = torch.nn.functional.pad(
-                sampled_token_ids,
-                (0, self.max_sample_len - sampled_token_ids.shape[-1]),
-            )
+            # receive() allocates the same width on the other ranks.
+            assert sampled_token_ids.shape[-1] == self._sample_len(input_batch)
             torch.distributed.broadcast(
-                send_tokens.contiguous(),
+                sampled_token_ids.contiguous(),
                 src=self.last_rank,
                 group=self.broadcast_group,
             )
diff --git a/vllm/v1/worker/gpu/sample/batch_shard.py b/vllm/v1/worker/gpu/sample/batch_shard.py
index daaf8806ff89..1ed99b1e5160 100644
--- a/vllm/v1/worker/gpu/sample/batch_shard.py
+++ b/vllm/v1/worker/gpu/sample/batch_shard.py
@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 from dataclasses import dataclass, replace
+from typing import TYPE_CHECKING
 
 import numpy as np
 import torch
@@ -17,6 +18,10 @@
 from vllm.v1.worker.gpu.input_batch import InputBatch
 from vllm.v1.worker.gpu.sample.output import SamplerOutput
 
+if TYPE_CHECKING:
+    from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
+    from vllm.v1.worker.gpu.spec_decode.draft_logits import DraftLogitsCache
+
 
 @dataclass
 class BatchShardMetadata:
@@ -60,19 +65,41 @@ def batch_sharded_sampling_enabled(
     )
 
 
+def get_tp_pynccl_comm() -> "PyNcclCommunicator | None":
+    """The TP group's PyNCCL communicator, or None if it is unavailable."""
+    pynccl_comm = getattr(get_tp_group().device_communicator, "pynccl_comm", None)
+    if pynccl_comm is None or pynccl_comm.disabled:
+        return None
+    return pynccl_comm
+
+
 @triton.jit
-def shard_owner(req_state_idx, NUM_SHARDS: tl.constexpr):
-    """TP rank that samples a request and holds its draft logits.
+def shard_owner(req_idx, NUM_SHARDS: tl.constexpr):
+    """TP rank that samples the request at a batch position.
 
-    Mirrored on the host in `BatchSharder.shard_sampler_inputs`.
+    Mirrored on the host by `shard_owners`.
     """
-    return req_state_idx % NUM_SHARDS
+    return req_idx % NUM_SHARDS
+
+
+def shard_owners(num_reqs: int, num_shards: int) -> np.ndarray:
+    """[num_reqs] The `shard_owner` of each batch position."""
+    return np.arange(num_reqs) % num_shards
 
 
 @triton.jit
-def shard_row(req_state_idx, NUM_SHARDS: tl.constexpr):
-    """Row of a request in its owner's batch-sharded draft logits cache."""
-    return req_state_idx // NUM_SHARDS
+def shard_row(req_idx, NUM_SHARDS: tl.constexpr):
+    """Row of the request at a batch position in its `shard_owner`'s sharded
+    draft logits cache.
+
+    Mirrored on the host by `shard_rows`.
+    """
+    return req_idx // NUM_SHARDS
+
+
+def shard_rows(num_reqs: int, num_shards: int) -> np.ndarray:
+    """[num_reqs] The `shard_row` of each batch position."""
+    return np.arange(num_reqs) // num_shards
 
 
 @triton.jit(
@@ -91,6 +118,7 @@ def _build_shard_plan_kernel(
     local_expanded_idx_mapping_ptr,
     local_expanded_local_pos_ptr,
     local_seq_lens_ptr,
+    owners_ptr,
     num_reqs,
     local_logits_start,
     max_num_reqs_per_rank,
@@ -101,16 +129,22 @@ def _build_shard_plan_kernel(
 ):
     req_idx = tl.program_id(0)
     req_state_idx = tl.load(idx_mapping_ptr + req_idx)
-    owner = shard_owner(req_state_idx, TP_SIZE)
     num_logits = tl.load(cu_num_logits_ptr + req_idx + 1) - tl.load(
         cu_num_logits_ptr + req_idx
     )
 
     req_block = tl.arange(0, PADDED_NUM_REQS)
     req_mask = req_block < num_reqs
-    owners = shard_owner(
-        tl.load(idx_mapping_ptr + req_block, mask=req_mask, other=0), TP_SIZE
-    )
+    if owners_ptr is None:
+        owner = shard_owner(req_idx, TP_SIZE)
+        owners = shard_owner(req_block, TP_SIZE)
+    else:
+        owner = tl.load(owners_ptr + req_state_idx)
+        owners = tl.load(
+            owners_ptr + tl.load(idx_mapping_ptr + req_block, mask=req_mask, other=0),
+            mask=req_mask,
+            other=0,
+        )
     num_logits_per_req = tl.load(
         cu_num_logits_ptr + req_block + 1, mask=req_mask, other=0
     ) - tl.load(cu_num_logits_ptr + req_block, mask=req_mask, other=0)
@@ -181,12 +215,14 @@ def __init__(
         max_num_logits_per_req: int,
         device: torch.device,
         adaptive_verification_enabled: bool,
+        draft_logits: "DraftLogitsCache | None" = None,
     ):
         tp_group = get_tp_group()
         self.tp_rank = tp_group.rank_in_group
         self.tp_size = tp_group.world_size
         self.device = device
         self.adaptive_verification_enabled = adaptive_verification_enabled
+        self.draft_logits = draft_logits
         self._padded_num_reqs = triton.next_power_of_2(max_num_reqs)
         self._padded_num_logits_per_req = triton.next_power_of_2(max_num_logits_per_req)
         self._num_warps = max(1, min(8, self._padded_num_reqs // 128))
@@ -208,11 +244,15 @@ def shard_sampler_inputs(
         tp_size = self.tp_size
         num_reqs = input_batch.idx_mapping_np.shape[0]
 
-        # Deterministically assign requests to ranks, round-robin over slot
-        # indices so ownership stays balanced when the slot allocator fills
-        # low slots first (partial occupancy).
+        # Deterministically assign requests to ranks: to the rank that cached
+        # their draft logits, or else round-robin over batch positions.
         # NOTE: This mirrors the assignment used in _build_shard_plan_kernel.
-        req_owner_np = input_batch.idx_mapping_np % tp_size
+        if self.draft_logits is None:
+            req_owner_np, req_owner = shard_owners(num_reqs, tp_size), None
+        else:
+            req_owner_np, req_owner = self.draft_logits.owners_of(
+                input_batch.idx_mapping_np
+            )
         local_req_indices_np = np.flatnonzero(req_owner_np == tp_rank)
         local_idx_mapping_np = input_batch.idx_mapping_np[local_req_indices_np]
         num_local_reqs = local_req_indices_np.shape[0]
@@ -282,6 +322,7 @@ def shard_sampler_inputs(
                 local_expanded_idx_mapping,
                 local_expanded_local_pos,
                 local_seq_lens,
+                req_owner,
                 num_reqs,
                 local_logits_start,
                 max_num_reqs_per_rank,
@@ -342,25 +383,39 @@ def shard_sampler_inputs(
         return local_batch, sorted_logits_indices, local_grammar_output, metadata
 
 
-def all_to_all_logits(
-    logits_shard: torch.Tensor,
-    metadata: BatchShardMetadata,
+def all_to_all_vocab_shards(
+    logits_shard: torch.Tensor, send_splits: list[int], num_local_rows: int
 ) -> torch.Tensor:
-    recv = logits_shard.new_empty(
-        metadata.tp_size * metadata.num_local_logits, logits_shard.shape[-1]
-    )
-    torch.distributed.all_to_all_single(
-        recv,
-        logits_shard.contiguous(),
-        output_split_sizes=[metadata.num_local_logits] * metadata.tp_size,
-        input_split_sizes=metadata.num_logits_per_rank,
-        group=get_tp_group().device_group,
-    )
+    """Exchange vocab-shard logits so each rank gets full-vocab logits for its
+    own rows. `send_splits[r]` rows of `logits_shard` go to rank r.
+
+    Returns:
+        [num_local_rows, tp_size * shard_width] logits of this rank's rows.
+
+    """
+    tp_size = len(send_splits)
+    recv = logits_shard.new_empty(tp_size * num_local_rows, logits_shard.shape[-1])
+    recv_splits = [num_local_rows] * tp_size
+    pynccl_comm = get_tp_pynccl_comm()
+    if pynccl_comm is not None:
+        # The TP collectives' communicator: a c10d all-to-all would make NCCL
+        # allocate a second set of point-to-point buffers.
+        pynccl_comm.all_to_all_single(
+            recv, logits_shard.contiguous(), recv_splits, send_splits
+        )
+    else:
+        torch.distributed.all_to_all_single(
+            recv,
+            logits_shard.contiguous(),
+            output_split_sizes=recv_splits,
+            input_split_sizes=send_splits,
+            group=get_tp_group().device_group,
+        )
     shard_width = recv.shape[-1]
     return (
-        recv.view(metadata.tp_size, metadata.num_local_logits, shard_width)
+        recv.view(tp_size, num_local_rows, shard_width)
         .permute(1, 0, 2)
-        .reshape(metadata.num_local_logits, metadata.tp_size * shard_width)
+        .reshape(num_local_rows, tp_size * shard_width)
     )
 
 
@@ -378,14 +433,9 @@ def _shard_grammar_output(
     bitmask = grammar_output.grammar_bitmask
     blocks = bitmask.reshape(len(grammar_req_ids), -1, bitmask.shape[-1])
     local_bitmask = blocks[keep].reshape(-1, bitmask.shape[-1])
-    # num_acceptable_drafts is ordered with structured_output_request_ids.
-    num_acceptable = grammar_output.num_acceptable_drafts
-    if num_acceptable is not None:
-        num_acceptable = [num_acceptable[i] for i in keep]
     return GrammarOutput(
         structured_output_request_ids=local_ids,
         grammar_bitmask=local_bitmask,
-        num_acceptable_drafts=num_acceptable,
     )
 
 
diff --git a/vllm/v1/worker/gpu/sample/greedy_argmax.py b/vllm/v1/worker/gpu/sample/greedy_argmax.py
new file mode 100644
index 000000000000..8809e0b32235
--- /dev/null
+++ b/vllm/v1/worker/gpu/sample/greedy_argmax.py
@@ -0,0 +1,66 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import torch
+
+from vllm.triton_utils import tl, triton
+
+
+@triton.jit
+def _argmax_partial(
+    X,
+    PV,
+    PI,
+    N: tl.constexpr,
+    SX: tl.constexpr,
+    SPLITS: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    row, part = tl.program_id(0), tl.program_id(1)
+    col = part * BLOCK + tl.arange(0, BLOCK)
+    value = tl.load(X + row.to(tl.int64) * SX + col, col < N, float("-inf"))
+    value = value.to(tl.float32)
+    # NaN breaks tl.argmax index bounds, as in the rejection sampler.
+    value = tl.where(value != value, float("-inf"), value)
+    index = tl.argmax(value, axis=0)
+    tl.store(PV + row * SPLITS + part, tl.max(value, axis=0))
+    tl.store(PI + row * SPLITS + part, part * BLOCK + index)
+
+
+@triton.jit
+def _argmax_finish(PV, PI, OUT, SPLITS: tl.constexpr, BLOCK: tl.constexpr):
+    row = tl.program_id(0)
+    col = tl.arange(0, BLOCK)
+    value = tl.load(PV + row * SPLITS + col, col < SPLITS, float("-inf"))
+    part = tl.argmax(value, axis=0)
+    tl.store(OUT + row, tl.load(PI + row * SPLITS + part).to(tl.int64))
+
+
+def greedy_argmax(logits: torch.Tensor) -> torch.Tensor:
+    """``logits.argmax(dim=-1)`` over a large vocabulary, split across CTAs.
+
+    Ties resolve to the first index. NaN logits are treated as -inf, so a row
+    with NaN may differ from ``torch.argmax``, which returns the NaN's index.
+    """
+    if (
+        not logits.is_cuda
+        or logits.dtype not in (torch.bfloat16, torch.float32)
+        or logits.ndim != 2
+        or logits.stride(1) != 1
+        or logits.shape[1] < 4096
+    ):
+        return logits.argmax(dim=-1)
+    rows, vocab = logits.shape
+    out = torch.empty(rows, dtype=torch.int64, device=logits.device)
+    if rows == 0:
+        return out
+    block = 4096
+    splits = triton.cdiv(vocab, block)
+    values = torch.empty((rows, splits), dtype=torch.float32, device=logits.device)
+    indices = torch.empty((rows, splits), dtype=torch.int32, device=logits.device)
+    _argmax_partial[(rows, splits)](
+        logits, values, indices, vocab, logits.stride(0), splits, block, num_warps=4
+    )
+    _argmax_finish[(rows,)](
+        values, indices, out, splits, triton.next_power_of_2(splits), num_warps=1
+    )
+    return out
diff --git a/vllm/v1/worker/gpu/sample/gumbel.py b/vllm/v1/worker/gpu/sample/gumbel.py
index 0c6503cbd6da..91ea3e1240f3 100644
--- a/vllm/v1/worker/gpu/sample/gumbel.py
+++ b/vllm/v1/worker/gpu/sample/gumbel.py
@@ -3,7 +3,7 @@
 import torch
 
 from vllm.triton_utils import tl, triton
-from vllm.v1.worker.gpu.sample.batch_shard import shard_owner, shard_row
+from vllm.v1.worker.gpu.spec_decode.draft_logits import draft_logits_row
 
 # Smallest positive value produced by Triton's fp32 `tl.rand`. Used by the
 # retained Philox helper for rejection sampling.
@@ -229,13 +229,12 @@ def gumbel_block_argmax(
     logits_cache_col_ptr,
     logits_cache_source_ptr,
     logits_cache_source_stride,
+    logits_cache_row_mapping_ptr,
     vocab_size,
     IS_DRAFTING: tl.constexpr,
     APPLY_TEMPERATURE: tl.constexpr,
     USE_FP64: tl.constexpr,
     PER_TOKEN_COL: tl.constexpr = False,
-    CACHE_NUM_SHARDS: tl.constexpr = 1,
-    CACHE_SHARD_IDX: tl.constexpr = 0,
 ):
     req_state_idx = tl.load(expanded_idx_mapping_ptr + token_idx).to(tl.int64)
     is_valid_req = req_state_idx >= 0
@@ -255,16 +254,14 @@ def gumbel_block_argmax(
             logits_cache_source_ptr + token_idx * logits_cache_source_stride + block,
             mask=mask,
         )
+        row = draft_logits_row(req_state_idx, logits_cache_row_mapping_ptr)
         tl.store(
             logits_cache_ptr
-            + shard_row(req_state_idx, CACHE_NUM_SHARDS) * logits_cache_stride_0
+            + row * logits_cache_stride_0
             + col * logits_cache_stride_1
             + block,
             cached_logits,
-            mask=mask
-            & is_valid_req
-            # A batch-sharded cache holds only this rank's requests, packed by row.
-            & (shard_owner(req_state_idx, CACHE_NUM_SHARDS) == CACHE_SHARD_IDX),
+            mask=mask & (row >= 0),
         )
 
     seed = tl.load(seeds_ptr + req_state_idx, mask=is_valid_req, other=0)
@@ -297,6 +294,7 @@ def _gumbel_sample_kernel(
     logits_cache_col_ptr,
     logits_cache_source_ptr,
     logits_cache_source_stride,
+    logits_cache_row_mapping_ptr,
     logits_ptr,
     logits_stride,
     expanded_idx_mapping_ptr,
@@ -309,8 +307,6 @@ def _gumbel_sample_kernel(
     APPLY_TEMPERATURE: tl.constexpr,
     USE_FP64: tl.constexpr,
     PER_TOKEN_COL: tl.constexpr,
-    CACHE_NUM_SHARDS: tl.constexpr,
-    CACHE_SHARD_IDX: tl.constexpr,
 ):
     token_idx = tl.program_id(0).to(tl.int64)
     block_idx = tl.program_id(1)
@@ -338,13 +334,12 @@ def _gumbel_sample_kernel(
         logits_cache_col_ptr,
         logits_cache_source_ptr,
         logits_cache_source_stride,
+        logits_cache_row_mapping_ptr,
         vocab_size,
         IS_DRAFTING=IS_DRAFTING,
         APPLY_TEMPERATURE=APPLY_TEMPERATURE,
         USE_FP64=USE_FP64,
         PER_TOKEN_COL=PER_TOKEN_COL,
-        CACHE_NUM_SHARDS=CACHE_NUM_SHARDS,
-        CACHE_SHARD_IDX=CACHE_SHARD_IDX,
     )
     token_id = block_idx * BLOCK_SIZE + idx
     tl.store(local_argmax_ptr + token_idx * local_argmax_stride + block_idx, token_id)
@@ -363,8 +358,7 @@ def gumbel_sample(
     logits_cache_col: torch.Tensor | None = None,  # scalar or [num_tokens]
     use_fp64: bool = False,
     logits_cache_source: torch.Tensor | None = None,
-    logits_cache_num_shards: int = 1,
-    logits_cache_shard_idx: int = 0,
+    logits_cache_row_mapping: torch.Tensor | None = None,  # [max_num_reqs]
 ) -> torch.Tensor:
     # Enforce contiguity on non-strided input tensors
     expanded_idx_mapping = expanded_idx_mapping.contiguous()
@@ -410,6 +404,7 @@ def gumbel_sample(
         logits_cache_col,
         logits_cache_source,
         logits_cache_source.stride(0) if logits_cache_source is not None else 0,
+        logits_cache_row_mapping,
         logits,
         logits.stride(0),
         expanded_idx_mapping,
@@ -422,8 +417,6 @@ def gumbel_sample(
         APPLY_TEMPERATURE=apply_temperature,
         USE_FP64=use_fp64,
         PER_TOKEN_COL=per_token_col,
-        CACHE_NUM_SHARDS=logits_cache_num_shards,
-        CACHE_SHARD_IDX=logits_cache_shard_idx,
     )
     # NOTE(woosuk): Use int64 for later indexing.
     max_block_idx = local_max.argmax(dim=-1, keepdim=True)
diff --git a/vllm/v1/worker/gpu/sample/logits_processor/interface.py b/vllm/v1/worker/gpu/sample/logits_processor/interface.py
index 9934492ec227..d1a8d7750d58 100644
--- a/vllm/v1/worker/gpu/sample/logits_processor/interface.py
+++ b/vllm/v1/worker/gpu/sample/logits_processor/interface.py
@@ -136,9 +136,9 @@ def add_request(self, req_idx: int, sampling_params: SamplingParams) -> bool:
     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.
+        Called before the forward pass of every step that admits new requests,
+        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
diff --git a/vllm/v1/worker/gpu/sample/penalties.py b/vllm/v1/worker/gpu/sample/penalties.py
index 4f0f003f05a8..232b86e0c634 100644
--- a/vllm/v1/worker/gpu/sample/penalties.py
+++ b/vllm/v1/worker/gpu/sample/penalties.py
@@ -130,7 +130,8 @@ def _penalties_kernel(
     use_rep_penalty = rep_penalty != 1.0
     use_freq_penalty = freq_penalty != 0.0
     use_pres_penalty = pres_penalty != 0.0
-    use_penalty = use_rep_penalty or use_freq_penalty or use_pres_penalty
+    # Parenthesize for Triton backends that reject chained `A or B or C`.
+    use_penalty = (use_rep_penalty or use_freq_penalty) or use_pres_penalty
     if not use_penalty:
         # Early return to avoid loading logits.
         return
diff --git a/vllm/v1/worker/gpu/sample/prompt_logprob.py b/vllm/v1/worker/gpu/sample/prompt_logprob.py
index adc785dbad49..9c5ca87e0c46 100644
--- a/vllm/v1/worker/gpu/sample/prompt_logprob.py
+++ b/vllm/v1/worker/gpu/sample/prompt_logprob.py
@@ -18,6 +18,9 @@
 )
 
 CHUNK_SIZE = 1024
+# Cap rows x k per torch.topk call, whose workspace grows with both (~5.5 GiB
+# for 1024 rows at a 152k vocab). Only splits chunks when k is very large.
+TOPK_MAX_ELEMENTS = 64 << 20
 
 
 @dataclass
@@ -293,12 +296,13 @@ def compute_prompt_logprobs_with_chunking(
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
     # Since materializing the full prompt logits can take too much memory,
     # we compute it in chunks.
-    token_ids = []
-    scores = []
-    ranks = []
+    num_tokens = prompt_token_ids.shape[0]
+    token_ids: torch.Tensor | None = None
+    scores: torch.Tensor | None = None
+    ranks: torch.Tensor | None = None
     logits_mode = logprobs_mode in ("raw_logits", "processed_logits")
     prompt_token_ids = prompt_token_ids.to(torch.int64)
-    for start_idx in range(0, prompt_token_ids.shape[0], CHUNK_SIZE):
+    for start_idx in range(0, num_tokens, CHUNK_SIZE):
         end_idx = start_idx + CHUNK_SIZE
         # NOTE(woosuk): logits_fn can be slow because it involves all-gather.
         prompt_logits = logits_fn(prompt_hidden_states[start_idx:end_idx])
@@ -307,17 +311,30 @@ def compute_prompt_logprobs_with_chunking(
             if num_prompt_logprobs == -1
             else num_prompt_logprobs
         )
-        result = compute_topk_scores(
-            prompt_logits,
-            requested_num,
-            prompt_token_ids[start_idx:end_idx],
-            logits_mode=logits_mode,
-        )
-        token_ids.append(result.logprob_token_ids)
-        scores.append(result.logprobs)
-        ranks.append(result.selected_token_ranks)
-
-    token_ids = torch.cat(token_ids, dim=0) if len(token_ids) > 1 else token_ids[0]
-    scores = torch.cat(scores, dim=0) if len(scores) > 1 else scores[0]
-    ranks = torch.cat(ranks, dim=0) if len(ranks) > 1 else ranks[0]
+        rows_per_topk = max(1, TOPK_MAX_ELEMENTS // (requested_num + 1))
+        for block_start in range(0, prompt_logits.shape[0], rows_per_topk):
+            logits = prompt_logits[block_start : block_start + rows_per_topk]
+            num_logits = len(logits)
+            rows = slice(start_idx + block_start, start_idx + block_start + num_logits)
+            result = compute_topk_scores(
+                logits, requested_num, prompt_token_ids[rows], logits_mode=logits_mode
+            )
+            # Return token_ids as int32 to reduce peak memory usage.
+            if num_logits == num_tokens:
+                logprob_token_ids = result.logprob_token_ids.to(torch.int32)
+                return logprob_token_ids, result.logprobs, result.selected_token_ranks
+            if token_ids is None:
+                size = num_tokens, result.logprobs.shape[1]
+                token_ids = prompt_token_ids.new_empty(size, dtype=torch.int32)
+                scores = result.logprobs.new_empty(size)
+                ranks = result.selected_token_ranks.new_empty(num_tokens)
+            assert scores is not None and ranks is not None
+            token_ids[rows] = result.logprob_token_ids
+            scores[rows] = result.logprobs
+            ranks[rows] = result.selected_token_ranks
+            # Free each block's results before ranking the next.
+            del result
+        del prompt_logits, logits
+
+    assert token_ids is not None and scores is not None and ranks is not None
     return token_ids, scores, ranks
diff --git a/vllm/v1/worker/gpu/sample/sampler.py b/vllm/v1/worker/gpu/sample/sampler.py
index 4ad56292e47b..11d5c5f33add 100644
--- a/vllm/v1/worker/gpu/sample/sampler.py
+++ b/vllm/v1/worker/gpu/sample/sampler.py
@@ -115,17 +115,11 @@ def apply_staged_writes(self) -> None:
         if self.trace_replay_state is not None:
             self.trace_replay_state.apply_staged_writes()
 
-    def get_logprobs_dims(
-        self, idx_mapping_np: np.ndarray, include_token_ids: bool = True
-    ) -> tuple[int, int] | None:
+    def get_logprobs_dims(self, idx_mapping_np: np.ndarray) -> tuple[int, int] | None:
         """(num_logprobs, max_per_req_token_ids) for the given requests, or
         None when none of them want logprobs."""
         max_num_logprobs = self.sampling_states.max_num_logprobs(idx_mapping_np)
-        max_token_ids = (
-            self.logprob_token_ids_state.max_num_token_ids(idx_mapping_np)
-            if include_token_ids
-            else 0
-        )
+        max_token_ids = self.logprob_token_ids_state.max_num_token_ids(idx_mapping_np)
         if max_num_logprobs == NO_LOGPROBS and max_token_ids == 0:
             return None
         num_logprobs = max_num_logprobs if max_num_logprobs != NO_LOGPROBS else 0
diff --git a/vllm/v1/worker/gpu/sample/trace_replay.py b/vllm/v1/worker/gpu/sample/trace_replay.py
index 9b58d4cf2923..7ecec14393ba 100644
--- a/vllm/v1/worker/gpu/sample/trace_replay.py
+++ b/vllm/v1/worker/gpu/sample/trace_replay.py
@@ -57,6 +57,26 @@ def apply_trace(
             self.req_states.prompt_len.gpu,
         )
 
+    def apply_trace_to_verified(
+        self,
+        sampled: torch.Tensor,
+        num_sampled: torch.Tensor,
+        draft_sampled: torch.Tensor,
+        cu_num_logits: torch.Tensor,
+        idx_mapping: torch.Tensor,
+    ) -> None:
+        apply_trace_tokens_to_verified(
+            sampled,
+            num_sampled,
+            draft_sampled,
+            cu_num_logits,
+            idx_mapping,
+            self.trace_token_ids.gpu,
+            self.trace_len.gpu,
+            self.req_states.total_len.gpu,
+            self.req_states.prompt_len.gpu,
+        )
+
 
 @triton.jit
 def _trace_replay_kernel(
@@ -111,3 +131,83 @@ def apply_trace_tokens(
         total_len,
         prompt_len,
     )
+
+
+@triton.jit
+def _trace_replay_verified_kernel(
+    sampled_ptr,
+    sampled_stride,
+    num_sampled_ptr,
+    draft_sampled_ptr,
+    cu_num_logits_ptr,
+    idx_mapping_ptr,
+    trace_token_ids_ptr,
+    trace_token_ids_stride,
+    trace_len_ptr,
+    total_len_ptr,
+    prompt_len_ptr,
+):
+    req_idx = tl.program_id(0)
+    req_state_idx = tl.load(idx_mapping_ptr + req_idx)
+    if req_state_idx < 0:
+        return
+
+    trace_len = tl.load(trace_len_ptr + req_state_idx)
+    if trace_len <= 0:
+        return
+
+    step = tl.load(total_len_ptr + req_state_idx) - tl.load(
+        prompt_len_ptr + req_state_idx
+    )
+    if step < 0 or step >= trace_len:
+        return
+
+    start = tl.load(cu_num_logits_ptr + req_idx)
+    num_draft_tokens = tl.load(cu_num_logits_ptr + req_idx + 1) - start - 1
+    num_replayed = tl.minimum(num_draft_tokens + 1, trace_len - step)
+    trace_row = (
+        trace_token_ids_ptr + req_state_idx.to(tl.int64) * trace_token_ids_stride + step
+    )
+    num_accepted = tl.zeros((), tl.int32)
+    matching = tl.full((), 1, tl.int1)
+    for i in range(num_replayed):
+        if matching:
+            token_id = tl.load(trace_row + i).to(tl.int64)
+            tl.store(sampled_ptr + req_idx * sampled_stride + i, token_id)
+            num_accepted = i + 1
+            draft_token_id = tl.load(
+                draft_sampled_ptr + start + i + 1,
+                mask=i < num_draft_tokens,
+                other=-1,
+            ).to(tl.int64)
+            matching = draft_token_id == token_id
+    tl.store(num_sampled_ptr + req_idx, num_accepted)
+
+
+def apply_trace_tokens_to_verified(
+    sampled: torch.Tensor,
+    num_sampled: torch.Tensor,
+    draft_sampled: torch.Tensor,
+    cu_num_logits: torch.Tensor,
+    idx_mapping: torch.Tensor,
+    trace_token_ids: torch.Tensor,
+    trace_len: torch.Tensor,
+    total_len: torch.Tensor,
+    prompt_len: torch.Tensor,
+) -> None:
+    """Replace verified tokens with the trace, accepting each draft only if it
+    equals the trace token at its position."""
+    num_reqs = idx_mapping.shape[0]
+    _trace_replay_verified_kernel[(num_reqs,)](
+        sampled,
+        sampled.stride(0),
+        num_sampled,
+        draft_sampled,
+        cu_num_logits,
+        idx_mapping,
+        trace_token_ids,
+        trace_token_ids.stride(0),
+        trace_len,
+        total_len,
+        prompt_len,
+    )
diff --git a/vllm/v1/worker/gpu/spec_decode/__init__.py b/vllm/v1/worker/gpu/spec_decode/__init__.py
index ccad8d5f953d..6ec9afdd4118 100644
--- a/vllm/v1/worker/gpu/spec_decode/__init__.py
+++ b/vllm/v1/worker/gpu/spec_decode/__init__.py
@@ -1,94 +1,49 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-from typing import TYPE_CHECKING
+import inspect
+from collections.abc import Callable
+from typing import TYPE_CHECKING, cast
 
 import torch
 
-from vllm.config import VllmConfig
-from vllm.config.speculative import XPRESS_ARCHITECTURES
+from vllm.utils.import_utils import resolve_obj_by_qualname
 
+# Keep this light: vllm.config imports the speculator configs in this package.
 if TYPE_CHECKING:
+    from vllm.config import VllmConfig
+    from vllm.v1.worker.gpu.spec_decode.speculator import BaseSpeculator
     from vllm.v1.worker.gpu.states import RequestState
 
 
 def init_speculator(
-    vllm_config: VllmConfig,
+    vllm_config: "VllmConfig",
     device: torch.device,
     req_states: "RequestState",
 ):
-    """Build the speculator for this config."""
-    speculative_config = vllm_config.speculative_config
-    assert speculative_config is not None
-    if speculative_config.method == "extract_hidden_states":
-        from vllm.v1.worker.gpu.spec_decode.extract_hidden_states import (
-            ExtractHiddenStatesSpeculator,
-        )
-
-        return ExtractHiddenStatesSpeculator(vllm_config, device)
-    elif speculative_config.method == "dflash":
-        architectures = speculative_config.draft_model_config.architectures
-        if any(arch in XPRESS_ARCHITECTURES for arch in architectures):
-            from vllm.v1.worker.gpu.spec_decode.xpress.speculator import (
-                XPressSpeculator,
-            )
-
-            return XPressSpeculator(vllm_config, device)
-        if "LiLiCorrDraftModel" in architectures:
-            from vllm.v1.worker.gpu.spec_decode.lilicorr.speculator import (
-                LiLiCorrSpeculator,
-            )
-
-            return LiLiCorrSpeculator(vllm_config, device)
-        if {"DFlash2DraftModel", "DFlash2BeamDraftModel"}.intersection(architectures):
-            from vllm.v1.worker.gpu.spec_decode.dflash2.speculator import (
-                DFlash2Speculator,
-            )
-
-            return DFlash2Speculator(vllm_config, device)
-        from vllm.v1.worker.gpu.spec_decode.dflash.speculator import (
-            DFlashSpeculator,
-        )
-
-        return DFlashSpeculator(vllm_config, device)
-    elif speculative_config.method == "dspark":
-        from vllm.v1.worker.gpu.spec_decode.dspark.speculator import (
-            DSparkSpeculator,
-        )
-
-        return DSparkSpeculator(vllm_config, device)
-    elif speculative_config.use_gemma4_mtp():
-        from vllm.v1.worker.gpu.spec_decode.gemma4.speculator import (
-            Gemma4Speculator,
-        )
-
-        return Gemma4Speculator(vllm_config, device)
-    elif speculative_config.use_multi_module_mtp():
-        from vllm.v1.worker.gpu.spec_decode.multi_module_mtp.speculator import (
-            MultiModuleMTPSpeculator,
-        )
+    """Build the speculator named by `speculative_config.speculator_config`.
 
-        return MultiModuleMTPSpeculator(vllm_config, device)
-    elif speculative_config.method == "mtp":
-        from vllm.v1.worker.gpu.spec_decode.mtp.speculator import MTPSpeculator
+    The speculator is constructed with `(vllm_config, device)`, plus
+    `req_states` if its constructor accepts it.
+    """
+    from vllm.v1.worker.gpu.spec_decode.speculator import BaseSpeculator
 
-        return MTPSpeculator(vllm_config, device)
-    elif speculative_config.use_eagle():
-        from vllm.v1.worker.gpu.spec_decode.eagle.speculator import (
-            EagleSpeculator,
-        )
-
-        return EagleSpeculator(vllm_config, device)
-    elif speculative_config.uses_draft_model():
-        from vllm.v1.worker.gpu.spec_decode.standalone_ar.speculator import (
-            StandaloneARSpeculator,
-        )
-
-        return StandaloneARSpeculator(vllm_config, device)
-    elif speculative_config.use_ngram():
-        from vllm.v1.worker.gpu.spec_decode.ngram.speculator import (
-            NgramGPUSpeculator,
-        )
-
-        return NgramGPUSpeculator(vllm_config, device, req_states)
-    else:
+    speculative_config = vllm_config.speculative_config
+    assert speculative_config is not None
+    speculator_config = speculative_config.speculator_config
+    qualname = speculator_config.speculator_class
+    if qualname is None:
         raise NotImplementedError(f"{speculative_config.method} is not supported yet.")
+    speculator_cls = resolve_obj_by_qualname(qualname)
+    if not (
+        isinstance(speculator_cls, type) and issubclass(speculator_cls, BaseSpeculator)
+    ):
+        raise TypeError(
+            f"speculator_class '{qualname}' of {type(speculator_config).__name__} "
+            f"must be a subclass of {BaseSpeculator.__module__}."
+            f"{BaseSpeculator.__qualname__}."
+        )
+    # BaseSpeculator doesn't define the constructor signature.
+    factory = cast(Callable[..., "BaseSpeculator"], speculator_cls)
+    if "req_states" in inspect.signature(factory).parameters:
+        return factory(vllm_config, device, req_states=req_states)
+    return factory(vllm_config, device)
diff --git a/vllm/v1/worker/gpu/spec_decode/acceptance_estimator.py b/vllm/v1/worker/gpu/spec_decode/acceptance_estimator.py
index f7a4a2f2c1e3..c210e23a1787 100644
--- a/vllm/v1/worker/gpu/spec_decode/acceptance_estimator.py
+++ b/vllm/v1/worker/gpu/spec_decode/acceptance_estimator.py
@@ -347,6 +347,7 @@ def __init__(
         self.device = device
         self._steps_since_refit = 0
         self._refits = 0
+        self._accumulate_block_r = max(256, triton.next_power_of_2(max_num_reqs))
 
         # The learned coefficients used for predicting acceptance.
         self.slope = torch.full((1,), _INIT_SLOPE, dtype=torch.float32, device=device)
@@ -422,8 +423,8 @@ def step(
             self.predictions.stride(0),
             self.counts,
             num_reqs,
-            BLOCK_R=max(256, triton.next_power_of_2(num_reqs)),
-            num_warps=1 if num_reqs <= 256 else 4,
+            BLOCK_R=self._accumulate_block_r,
+            num_warps=1 if self._accumulate_block_r <= 256 else 4,
         )
 
         self._steps_since_refit += 1
@@ -433,6 +434,19 @@ def step(
 
         # Fit the coefficients to the accumulated statistics gathered over the
         # course of the last REFIT_INTERVAL steps.
+        self._refit()
+        self._refits += 1
+
+    def warmup(self) -> None:
+        """Drop warmup statistics and compile the refit kernel. A refit on
+        zeroed statistics leaves the coefficients unchanged."""
+        self.info.zero_()
+        self.grad.zero_()
+        self.counts.zero_()
+        self._steps_since_refit = 0
+        self._refit()
+
+    def _refit(self) -> None:
         _refit_kernel[(1,)](
             self.slope,
             self.intercepts,
@@ -446,7 +460,6 @@ def step(
             DAMPING=self.DAMPING_OBSERVATIONS,
             BLOCK=triton.next_power_of_2(self.num_speculative_steps),
         )
-        self._refits += 1
 
     def predict(
         self,
diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/config.py b/vllm/v1/worker/gpu/spec_decode/dflash/config.py
new file mode 100644
index 000000000000..95538d8f5a26
--- /dev/null
+++ b/vllm/v1/worker/gpu/spec_decode/dflash/config.py
@@ -0,0 +1,43 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from typing import TYPE_CHECKING
+
+from vllm.v1.worker.gpu.spec_decode.eagle.config import wrap_eagle_hf_config
+from vllm.v1.worker.gpu.spec_decode.speculator_config import (
+    DraftModelSpeculatorConfig,
+)
+
+if TYPE_CHECKING:
+    from vllm.config import SpeculativeConfig
+
+
+class DFlashSpeculatorConfig(DraftModelSpeculatorConfig):
+    base_method = "dflash"
+    speculator_class = (
+        "vllm.v1.worker.gpu.spec_decode.dflash.speculator.DFlashSpeculator"
+    )
+    uses_target_hidden_states = True
+    supports_watermarking = True
+
+    @classmethod
+    def matches(cls, spec: "SpeculativeConfig") -> bool:
+        return "dflash" in spec.draft_model_config.model.lower()
+
+    @classmethod
+    def adapt_draft_model_config(cls, spec: "SpeculativeConfig") -> None:
+        assert cls.base_method is not None
+        wrap_eagle_hf_config(spec, cls.base_method)
+        spec.parallel_drafting = True
+
+    @property
+    def uses_aux_hidden_states(self) -> bool:
+        return True
+
+    @property
+    def num_lookahead_tokens(self) -> int:
+        # Fill-in drafting uses a bonus query plus one query per draft token.
+        return self.num_speculative_tokens + 1
+
+    def max_num_draft_seqs(self, max_num_batched_tokens: int) -> int:
+        # The draft pass has num_lookahead_tokens queries for every request.
+        return max_num_batched_tokens // self.num_lookahead_tokens
diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py
index c694d99a4335..8b369280abfb 100644
--- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py
+++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py
@@ -133,9 +133,9 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device):
     @property
     def attn_vllm_config(self) -> VllmConfig:
         # The draft's attention differs from the target's in causality.
-        config = copy.copy(super().attn_vllm_config)
+        config = super().attn_vllm_config
         config.attention_config = replace(
-            self.vllm_config.attention_config,
+            config.attention_config,
             use_non_causal=self.requires_non_causal,
         )
         return config
@@ -575,6 +575,7 @@ def propose(
             dp_rank=self.dp_rank,
             need_eager=is_profile,
             dp_sync_state=batch_sync,
+            parallel_config=self.vllm_config.parallel_config,
         )
         num_tokens_padded = batch_desc.num_tokens
         num_tokens_across_dp = (
@@ -888,7 +889,7 @@ def prepare_dflash_inputs(
     # per-request query length, not the total token count across the batch.
     max_target_query_len = int(input_batch.num_scheduled_tokens[:num_reqs].max())
     max_tokens_per_req = max_target_query_len + num_query_per_req
-    BLOCK_SIZE = min(256, triton.next_power_of_2(max(1, max_tokens_per_req)))
+    BLOCK_SIZE = 128
     num_blocks = triton.cdiv(max_tokens_per_req, BLOCK_SIZE)
     _prepare_dflash_inputs_kernel[(num_reqs, num_blocks)](
         input_buffers.input_ids,
diff --git a/vllm/v1/worker/gpu/spec_decode/dflash2/config.py b/vllm/v1/worker/gpu/spec_decode/dflash2/config.py
new file mode 100644
index 000000000000..595ee4b87b90
--- /dev/null
+++ b/vllm/v1/worker/gpu/spec_decode/dflash2/config.py
@@ -0,0 +1,11 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from vllm.v1.worker.gpu.spec_decode.dflash.config import DFlashSpeculatorConfig
+
+
+class DFlash2SpeculatorConfig(DFlashSpeculatorConfig):
+    speculator_class = (
+        "vllm.v1.worker.gpu.spec_decode.dflash2.speculator.DFlash2Speculator"
+    )
+    # Draft-side watermarking doesn't cover candidate-head drafting yet.
+    supports_watermarking = False
diff --git a/vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py
index 198da76a06cb..e0d12efcabb8 100644
--- a/vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py
+++ b/vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py
@@ -8,9 +8,12 @@
 from vllm.config import VllmConfig
 from vllm.config.compilation import CUDAGraphMode
 from vllm.triton_utils import tl, triton
-from vllm.v1.worker.gpu.sample.batch_shard import shard_owner, shard_row
 from vllm.v1.worker.gpu.sample.gumbel import gumbel_noised_argmax, spec_step_seed
 from vllm.v1.worker.gpu.spec_decode.dflash.speculator import DFlashSpeculator
+from vllm.v1.worker.gpu.spec_decode.draft_logits import (
+    DraftLogitsCache,
+    draft_logits_row,
+)
 
 
 @triton.jit
@@ -83,31 +86,24 @@ def _cache_draft_logits_kernel(
     candidate_ptr,
     scores_ptr,
     req_state_ptr,
+    draft_logits_row_mapping_ptr,
     draft_logits_stride_0,
     draft_logits_stride_1,
     num_steps: tl.constexpr,
     top_k: tl.constexpr,
     BLOCK_K: tl.constexpr,
-    CACHE_NUM_SHARDS: tl.constexpr,
-    CACHE_SHARD_IDX: tl.constexpr,
 ):
     flat = tl.program_id(0)
     req_state = tl.load(req_state_ptr + flat)
     step = flat % num_steps
     offsets = tl.arange(0, BLOCK_K)
-    mask = (
-        (req_state >= 0)
-        & (offsets < top_k)
-        # A batch-sharded cache holds only this rank's requests, packed by row.
-        & (shard_owner(req_state, CACHE_NUM_SHARDS) == CACHE_SHARD_IDX)
-    )
+    row = draft_logits_row(req_state, draft_logits_row_mapping_ptr)
+    mask = (row >= 0) & (offsets < top_k)
     candidate_base = flat * top_k
-    cache_base = (req_state * num_steps + step) * top_k
+    cache_base = (row * num_steps + step) * top_k
     old_token_ids = tl.load(cached_candidate_ptr + cache_base + offsets, mask=mask)
     logits_base = (
-        draft_logits_ptr
-        + shard_row(req_state, CACHE_NUM_SHARDS) * draft_logits_stride_0
-        + step * draft_logits_stride_1
+        draft_logits_ptr + row * draft_logits_stride_0 + step * draft_logits_stride_1
     )
     tl.store(logits_base + old_token_ids, -float("inf"), mask=mask)
     token_ids = tl.load(candidate_ptr + candidate_base + offsets, mask=mask)
@@ -141,10 +137,8 @@ def sample(
         temperature: torch.Tensor,
         seeds: torch.Tensor,
         draft_tokens: torch.Tensor,
-        draft_logits: torch.Tensor | None,
+        draft_logits: DraftLogitsCache | None,
         use_fp64: bool,
-        draft_logits_num_shards: int = 1,
-        draft_logits_shard_idx: int = 0,
     ) -> None:
         candidate_ids = candidate_ids.contiguous()
         sample_pos = sample_pos.contiguous()
@@ -173,8 +167,6 @@ def sample(
                 num_reqs * self.num_steps,
                 idx_mapping,
                 draft_logits,
-                draft_logits_num_shards,
-                draft_logits_shard_idx,
             )
 
     def _cache_draft_logits(
@@ -182,24 +174,21 @@ def _cache_draft_logits(
         candidate_ids: torch.Tensor,
         num_sample: int,
         idx_mapping: torch.Tensor,
-        draft_logits: torch.Tensor,
-        draft_logits_num_shards: int,
-        draft_logits_shard_idx: int,
+        draft_logits: DraftLogitsCache,
     ) -> None:
         block_k = triton.next_power_of_2(self.top_k)
         _cache_draft_logits_kernel[(num_sample,)](
-            draft_logits,
+            draft_logits.values,
             self.cached_candidate_ids,
             candidate_ids,
             self.scores,
             idx_mapping,
-            draft_logits.stride(0),
-            draft_logits.stride(1),
+            draft_logits.row_mapping,
+            draft_logits.values.stride(0),
+            draft_logits.values.stride(1),
             num_steps=self.num_steps,
             top_k=self.top_k,
             BLOCK_K=block_k,
-            CACHE_NUM_SHARDS=draft_logits_num_shards,
-            CACHE_SHARD_IDX=draft_logits_shard_idx,
             num_warps=1,
         )
 
@@ -261,8 +250,6 @@ def _generate_draft(
             self.draft_tokens,
             self.draft_logits,
             self.use_fp64_gumbel,
-            self.draft_logits_num_shards,
-            self.draft_logits_shard_idx,
         )
         if self.enable_adaptive_verification:
             self._maybe_predict_acceptance(
diff --git a/vllm/v1/worker/gpu/spec_decode/draft_logits.py b/vllm/v1/worker/gpu/spec_decode/draft_logits.py
new file mode 100644
index 000000000000..1bebcf69b75b
--- /dev/null
+++ b/vllm/v1/worker/gpu/spec_decode/draft_logits.py
@@ -0,0 +1,201 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import numpy as np
+import torch
+
+from vllm.distributed.parallel_state import get_tp_group
+from vllm.triton_utils import tl, triton
+from vllm.utils.math_utils import cdiv
+from vllm.v1.worker.gpu.input_batch import InputBatch
+from vllm.v1.worker.gpu.sample.batch_shard import (
+    shard_owner,
+    shard_owners,
+    shard_row,
+    shard_rows,
+)
+
+
+@triton.jit
+def draft_logits_row(req_state_idx, row_mapping_ptr):
+    """Row of a request in this rank's draft logits cache, or -1 if no row.
+
+    ``row_mapping_ptr`` is None for a replicated cache, which is indexed by request
+    state. Otherwise it is `DraftLogitsCache.row_mapping`.
+    """
+    row = req_state_idx
+    if row_mapping_ptr is not None:
+        row = tl.load(
+            row_mapping_ptr + req_state_idx, mask=req_state_idx >= 0, other=-1
+        )
+    return row.to(tl.int64)
+
+
+@triton.jit(do_not_specialize=["num_reqs"])
+def _assign_rows_kernel(
+    idx_mapping_ptr,
+    shard_mapping_ptr,
+    row_mapping_ptr,
+    inverse_row_mapping_ptr,
+    num_reqs,
+    NUM_SHARDS: tl.constexpr,
+    SHARD_IDX: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    req_block_idx = tl.program_id(0)
+    req_block = req_block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    mask = req_block < num_reqs
+    req_state_block = tl.load(idx_mapping_ptr + req_block, mask=mask, other=0)
+    owners = shard_owner(req_block, NUM_SHARDS)
+    rows = shard_row(req_block, NUM_SHARDS)
+    is_local = owners == SHARD_IDX
+    tl.store(shard_mapping_ptr + req_state_block, owners, mask=mask)
+    tl.store(row_mapping_ptr + req_state_block, tl.where(is_local, rows, -1), mask=mask)
+    tl.store(inverse_row_mapping_ptr + rows, req_state_block, mask=mask & is_local)
+
+
+@triton.jit(do_not_specialize=["num_logits"])
+def _stale_drafts_kernel(
+    stale_ptr,
+    expanded_idx_mapping_ptr,
+    expanded_local_pos_ptr,
+    row_mapping_ptr,
+    inverse_row_mapping_ptr,
+    num_logits,
+    BLOCK_SIZE: tl.constexpr,
+):
+    logit_block_idx = tl.program_id(0)
+    logit_block = logit_block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    mask = logit_block < num_logits
+    req_state_block = tl.load(
+        expanded_idx_mapping_ptr + logit_block, mask=mask, other=-1
+    )
+    local_pos = tl.load(expanded_local_pos_ptr + logit_block, mask=mask, other=0)
+    rows = draft_logits_row(req_state_block, row_mapping_ptr)
+    holders = tl.load(inverse_row_mapping_ptr + rows, mask=mask & (rows >= 0), other=-1)
+    # Position 0 is the bonus token, which has no draft.
+    stale = (local_pos > 0) & ((rows < 0) | (holders != req_state_block))
+    tl.store(stale_ptr + logit_block, stale, mask=mask)
+
+
+class DraftLogitsCache:
+    """Pre-temperature draft logits, cached by a draft step for the next
+    verification.
+
+    A replicated cache holds every request on every TP rank, at row
+    req_state_idx.
+
+    A sharded cache (batch-sharded sampling) holds on each TP rank only the
+    requests that rank verifies. Every draft step gives the request at batch
+    position i to rank i % tp, at row i // tp of its cache, so no rank caches
+    more than cdiv(num_reqs, tp) requests however the request slots are spread.
+
+    Rows are reused every step. A request that the scheduler skips for a step
+    keeps its drafts, but its row may since have been given to another
+    request. `inverse_row_mapping` detects this, and `stale_drafts` marks those
+    drafts invalid, so verification rejects them and samples from the target.
+    """
+
+    def __init__(
+        self,
+        max_num_reqs: int,
+        num_steps: int,
+        vocab_size: int,
+        *,
+        dtype: torch.dtype,
+        fill: float,
+        device: torch.device | str,
+        sharded: bool = False,
+    ):
+        num_rows = max_num_reqs
+        # [max_num_reqs] Each request's row in values, -1 if another rank
+        # caches it. None for a replicated cache.
+        self.row_mapping: torch.Tensor | None = None
+        if sharded:
+            tp_group = get_tp_group()
+            self._num_shards = tp_group.world_size
+            self._shard_idx = tp_group.rank_in_group
+            num_rows = cdiv(max_num_reqs, self._num_shards)
+            # Requests never drafted keep a slot-based owner. They have no
+            # drafts, so any owner works as long as every rank agrees.
+            slots = np.arange(max_num_reqs, dtype=np.int32)
+            # [max_num_reqs] The TP rank that caches, and verifies, each request.
+            self.shard_mapping_np = slots % self._num_shards
+            self.shard_mapping = torch.from_numpy(self.shard_mapping_np).to(device)
+            rows_np = np.where(
+                self.shard_mapping_np == self._shard_idx,
+                shard_rows(max_num_reqs, self._num_shards),
+                -1,
+            )
+            self.row_mapping = torch.from_numpy(rows_np.astype(np.int32)).to(device)
+            # [num_rows] The request last drafted into each row of this rank.
+            self.inverse_row_mapping = torch.full(
+                (num_rows,), -1, dtype=torch.int32, device=device
+            )
+        # [num_rows, num_steps, vocab_size]
+        self.values = torch.full(
+            (num_rows, num_steps, vocab_size), fill, dtype=dtype, device=device
+        )
+
+    def assign(self, idx_mapping: torch.Tensor, idx_mapping_np: np.ndarray) -> None:
+        """Give a sharded cache's rows out by batch position. Call after
+        verification and before the draft step writes the draft logits."""
+        num_reqs = idx_mapping_np.shape[0]
+        if self.row_mapping is None or num_reqs == 0:
+            return
+        self.shard_mapping_np[idx_mapping_np] = shard_owners(num_reqs, self._num_shards)
+        block_size = 1024
+        _assign_rows_kernel[(triton.cdiv(num_reqs, block_size),)](
+            idx_mapping,
+            self.shard_mapping,
+            self.row_mapping,
+            self.inverse_row_mapping,
+            num_reqs,
+            NUM_SHARDS=self._num_shards,
+            SHARD_IDX=self._shard_idx,
+            BLOCK_SIZE=block_size,
+        )
+
+    def owners_of(
+        self, idx_mapping_np: np.ndarray
+    ) -> tuple[np.ndarray, torch.Tensor | None]:
+        """The TP rank that verifies each of a batch's requests: the rank that
+        cached its draft logits, or for a replicated cache, round-robin over
+        batch positions.
+
+        Returns:
+            [num_reqs] each request's owner, and the [max_num_reqs] device owner
+            table, or None for a replicated cache.
+
+        """
+        if self.row_mapping is None:
+            num_reqs = idx_mapping_np.shape[0]
+            return shard_owners(num_reqs, get_tp_group().world_size), None
+        return self.shard_mapping_np[idx_mapping_np], self.shard_mapping
+
+    def stale_drafts(self, input_batch: InputBatch) -> torch.Tensor | None:
+        """Mark the drafts whose draft logits were overwritten since drafting.
+        Only valid for requests this rank owns.
+
+        Returns:
+            [num_logits] True for each stale draft, or None for a replicated
+            cache, whose rows are never reused.
+
+        """
+        if self.row_mapping is None:
+            return None
+        expanded_idx_mapping = input_batch.expanded_idx_mapping
+        num_logits = expanded_idx_mapping.shape[0]
+        stale = expanded_idx_mapping.new_empty(num_logits, dtype=torch.bool)
+        if num_logits == 0:
+            return stale
+        block_size = 1024
+        _stale_drafts_kernel[(triton.cdiv(num_logits, block_size),)](
+            stale,
+            expanded_idx_mapping,
+            input_batch.expanded_local_pos,
+            self.row_mapping,
+            self.inverse_row_mapping,
+            num_logits,
+            BLOCK_SIZE=block_size,
+        )
+        return stale
diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/config.py b/vllm/v1/worker/gpu/spec_decode/dspark/config.py
new file mode 100644
index 000000000000..232891054672
--- /dev/null
+++ b/vllm/v1/worker/gpu/spec_decode/dspark/config.py
@@ -0,0 +1,144 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from typing import TYPE_CHECKING, Any
+
+from vllm.config.speculative import (
+    _MINIMAX_M3_MODEL_TYPES,
+    _QWEN3_OMNI_DSPARK_ARCHITECTURE,
+    _get_minimax_m3_dspark_draft,
+    _validate_qwen3_omni_dspark,
+)
+from vllm.v1.worker.gpu.spec_decode.dflash.config import DFlashSpeculatorConfig
+from vllm.v1.worker.gpu.spec_decode.eagle.config import (
+    check_draft_lm_head_quantization,
+)
+
+if TYPE_CHECKING:
+    from vllm.config import ParallelConfig, SpeculativeConfig
+
+
+class DSparkSpeculatorConfig(DFlashSpeculatorConfig):
+    base_method = "dspark"
+    speculator_class = (
+        "vllm.v1.worker.gpu.spec_decode.dspark.speculator.DSparkSpeculator"
+    )
+    supported_options = frozenset({"dspark_draft_topk", "draft_lm_head_quantization"})
+    uses_separate_workspace_lane = True
+    # DSpark predicts acceptance with its own confidence head.
+    adaptive_verification_uses_draft_logits = False
+    supports_pcp = True
+
+    @classmethod
+    def matches(cls, spec: "SpeculativeConfig") -> bool:
+        # e.g. deepseek-ai/dspark_qwen3_8b_block7
+        return "dspark" in spec.draft_model_config.model.lower()
+
+    @classmethod
+    def default_model(cls, spec: "SpeculativeConfig") -> str | None:
+        target = spec.target_model_config
+        if target is None:
+            raise ValueError("target_model_config must be present for dspark")
+        if target.hf_config.model_type in _MINIMAX_M3_MODEL_TYPES:
+            load_config = spec.draft_load_config
+            return _get_minimax_m3_dspark_draft(
+                target, load_config.download_dir if load_config else None
+            )
+        # DeepSeek DSpark can ship the weights inside the target checkpoint
+        if not spec.quantization:
+            spec.quantization = target.quantization
+        return target.model
+
+    @classmethod
+    def adapt_draft_model_config(cls, spec: "SpeculativeConfig") -> None:
+        draft = spec.draft_model_config
+        target = spec.target_model_config
+        if draft.speculator_config_class is None:
+            # Not a self-contained DSpark draft: DeepSeek-V4(.1) DSpark reuses
+            # the target's config, and its weights ship in the target checkpoint.
+            is_v41 = target.hf_config.model_type == "deepseek_v41"
+            draft.hf_config.model_type = "deepseek_v41" if is_v41 else "deepseek_v4"
+            draft.hf_config.architectures = [
+                "DSparkV41DraftModel" if is_v41 else "DSparkDraftModel"
+            ]
+            draft.quantization = target.quantization
+            spec.update_arch_()
+        spec.parallel_drafting = True
+
+    @classmethod
+    def resolve_num_speculative_tokens(cls, spec: "SpeculativeConfig") -> None:
+        if spec.num_speculative_tokens is None:
+            # DSpark's parallel width is independent of MTP stages.
+            hf_config = spec.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)
+            spec.num_speculative_tokens = block_size
+
+    @classmethod
+    def apply_options(cls, spec: "SpeculativeConfig") -> None:
+        draft = spec.draft_model_config
+        hf_config = draft.hf_config
+        dspark_draft_topk = spec.dspark_draft_topk
+        if dspark_draft_topk is None:
+            dspark_draft_topk = getattr(hf_config, "dspark_draft_topk", None)
+        if dspark_draft_topk is not None:
+            draft_vocab_size = (
+                getattr(hf_config, "draft_vocab_size", None) or hf_config.vocab_size
+            )
+            if not 1 <= dspark_draft_topk <= draft_vocab_size:
+                raise ValueError(
+                    "dspark_draft_topk must be between 1 and the "
+                    f"draft vocabulary size ({draft_vocab_size})"
+                )
+            if (
+                "Qwen3DSparkModel" not in draft.architectures
+                and _QWEN3_OMNI_DSPARK_ARCHITECTURE not in draft.architectures
+            ):
+                raise ValueError(
+                    "dspark_draft_topk is only supported by "
+                    "Qwen3DSparkModel and Qwen3OmniDSparkModel"
+                )
+            hf_config.dspark_draft_topk = dspark_draft_topk
+
+        if spec.draft_lm_head_quantization is not None:
+            # Other drafters may read lm_head.weight directly (e.g. Muse Spark).
+            if not {"DSparkDraftModel", "DSparkV41DraftModel"} & set(
+                draft.architectures
+            ):
+                raise ValueError(
+                    "draft_lm_head_quantization is only supported by the "
+                    "DeepSeek-V4/V4.1 DSpark drafters"
+                )
+            check_draft_lm_head_quantization(spec)
+
+        assert spec.target_model_config is not None
+        _validate_qwen3_omni_dspark(
+            spec.target_model_config, draft, spec.num_speculative_tokens
+        )
+
+    @property
+    def num_lookahead_tokens(self) -> int:
+        hf_config = self.spec_config.draft_model_config.hf_config
+        if getattr(hf_config, "sample_from_anchor", True):
+            # Anchor-sampling DSpark drafts a block of num_speculative_tokens
+            # query tokens in which the anchor itself is the first prediction
+            # position (no separate bonus query).
+            return self.num_speculative_tokens
+        return super().num_lookahead_tokens
+
+    def hash_factors(self) -> list[Any]:
+        # Adaptive verification changes the DSpark draft graph.
+        return [self.spec_config.enable_adaptive_verification]
+
+    def pcp_unsupported_reason(self, parallel_config: "ParallelConfig") -> str | None:
+        dcp_size = parallel_config.decode_context_parallel_size
+        pcp_size = parallel_config.prefill_context_parallel_size
+        if (
+            dcp_size not in (1, pcp_size)
+            and self.spec_config.draft_model_config.use_mla
+        ):
+            return (
+                "MRV2 PCP DSpark requires DCP=1 or DCP=PCP; got "
+                f"DCP={dcp_size}, PCP={pcp_size}."
+            )
+        return super().pcp_unsupported_reason(parallel_config)
diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py b/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py
index d2425ad95a87..58b8c76f62cb 100644
--- a/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py
+++ b/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py
@@ -21,6 +21,10 @@
 
 CUDA graphs (FULL, mirroring DFlash) cover the whole draft step: the parallel
 backbone forward AND the sequential Markov sampling.
+
+Under batch-sharded sampling, draft steps also shard the sequential sampling
+across TP ranks: each rank runs the Markov loop only for the requests it owns
+(the same owners as verification), see ``_sample_sequential_sharded``.
 """
 
 from typing import Any
@@ -29,7 +33,16 @@
 
 from vllm.config import VllmConfig
 from vllm.config.compilation import CUDAGraphMode
+from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
+from vllm.distributed.parallel_state import get_tp_group
 from vllm.logger import init_logger
+from vllm.utils.math_utils import cdiv
+from vllm.v1.worker.gpu.sample.batch_shard import (
+    all_to_all_vocab_shards,
+    batch_sharded_sampling_enabled,
+    get_tp_pynccl_comm,
+)
+from vllm.v1.worker.gpu.sample.greedy_argmax import greedy_argmax
 from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample
 from vllm.v1.worker.gpu.spec_decode.dflash.speculator import DFlashSpeculator
 from vllm.v1.worker.gpu.spec_decode.dspark.utils import load_dspark_model
@@ -81,6 +94,53 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device):
 
         self.use_confidence_head: bool = False
 
+        # Set when the sequential sampling can be batch-sharded across TP.
+        self._sharded_sequential_comm: PyNcclCommunicator | None = None
+
+    def load_model(self, target_model: torch.nn.Module) -> None:
+        super().load_model(target_model)
+        tp_size = get_tp_group().world_size
+        if tp_size > 1 and batch_sharded_sampling_enabled(
+            self.vllm_config, target_model
+        ):
+            self._try_enable_sharded_sequential_sampling()
+
+    def _try_enable_sharded_sequential_sampling(self) -> None:
+        tp_size = get_tp_group().world_size
+        lm_head = getattr(self.model, "lm_head", None)
+        pynccl_comm = get_tp_pynccl_comm()
+        reason = None
+        if self.draft_logits is not None and self.draft_logits.row_mapping is None:
+            reason = "the draft logits cache is replicated"
+        elif self.use_acceptance_estimator:
+            reason = (
+                "the online acceptance estimator fits on every request's "
+                "draft logits on every rank"
+            )
+        elif self._draft_topk is not None:
+            reason = "top-k Markov sampling is not supported"
+        elif hasattr(self.model, "sample_sequential"):
+            reason = "the draft model fuses its own sequential sampling"
+        elif not hasattr(self.model, "compute_draft_logits_local") or (
+            getattr(lm_head, "tp_size", None) != tp_size
+        ):
+            reason = "the draft model cannot compute vocab-sharded draft logits"
+        elif pynccl_comm is None:
+            reason = "PyNCCL is unavailable"
+
+        if reason is not None:
+            logger.info(
+                "DSpark sequential sampling stays replicated across TP ranks: %s.",
+                reason,
+            )
+            return
+
+        self._sharded_sequential_comm = pynccl_comm
+        logger.info(
+            "DSpark sequential sampling is batch-sharded across %d TP ranks.",
+            tp_size,
+        )
+
     def load_draft_model(
         self,
         target_model: torch.nn.Module,
@@ -126,7 +186,7 @@ def _sample_logits(
     ) -> torch.Tensor:
         self._maybe_predict_acceptance(logits, idx_map, self._step_cols[step])
         if self.draft_logits is None:
-            draft_ids = logits.argmax(dim=-1)
+            draft_ids = greedy_argmax(logits)
             return self.model.map_draft_to_target(draft_ids)
 
         # Probabilistic sampling and rejection operate in target-vocabulary
@@ -148,7 +208,7 @@ def _sample_logits(
                 sample_pos - 1,
                 apply_temperature=True,
                 is_drafting=True,
-                logits_cache=self.draft_logits,
+                logits_cache=self.draft_logits.values,
                 logits_cache_col=self._step_cols[step],
                 use_fp64=self.use_fp64_gumbel,
             )
@@ -160,11 +220,10 @@ def _sample_logits(
             sample_pos - 1,
             apply_temperature=True,
             is_drafting=True,
-            logits_cache=self.draft_logits,
+            logits_cache=self.draft_logits.values,
             logits_cache_col=self._step_cols[step],
             use_fp64=self.use_fp64_gumbel,
-            logits_cache_num_shards=self.draft_logits_num_shards,
-            logits_cache_shard_idx=self.draft_logits_shard_idx,
+            logits_cache_row_mapping=self.draft_logits.row_mapping,
         )
 
     def _sample_sequential(self, num_reqs: int, head_hidden: torch.Tensor) -> None:
@@ -190,19 +249,42 @@ def _sample_sequential(self, num_reqs: int, head_hidden: torch.Tensor) -> None:
         vocab_size = base_logits.shape[-1]
         base_logits = base_logits.view(num_reqs, n_spec, vocab_size)
 
-        idx_map = self.sample_idx_mapping[:num_sample].view(num_reqs, n_spec)
-        sample_pos = self.sample_pos[:num_sample].view(num_reqs, n_spec)
-        confidence_markov_embeds = []
-
         # Anchor (bonus) token per request = the input id at query offset 0,
         # read via the precomputed persistent index (fixed buffer for capture).
-        prev = self.input_buffers.input_ids[self._anchor_idx[:num_reqs]]
+        confidence = self._sample_markov_chain(
+            base_logits,
+            sample_hidden,
+            self.input_buffers.input_ids[self._anchor_idx[:num_reqs]],
+            self.sample_idx_mapping[:num_sample].view(num_reqs, n_spec),
+            self.sample_pos[:num_sample].view(num_reqs, n_spec),
+            self.draft_tokens[:num_reqs],
+        )
+        if confidence is not None:
+            self.draft_token_confidence_probs[:num_reqs] = confidence
 
-        for i in range(n_spec):
+    def _sample_markov_chain(
+        self,
+        base_logits: torch.Tensor,
+        hidden: torch.Tensor,
+        prev: torch.Tensor,
+        idx_map: torch.Tensor,
+        sample_pos: torch.Tensor,
+        out_tokens: torch.Tensor,
+    ) -> torch.Tensor | None:
+        """Run the sequential Markov steps for a block of requests, writing the
+        drafts to `out_tokens`.
+
+        Returns:
+            [num_rows, num_speculative_steps] draft confidences, or None
+            without a confidence head.
+
+        """
+        markov_embeds = []
+        for i in range(self.num_speculative_steps):
             # Sequential stage: Markov bias from the previously sampled token.
             markov_embed = self.model.markov_embed(prev)
             if self.use_confidence_head:
-                confidence_markov_embeds.append(markov_embed)
+                markov_embeds.append(markov_embed)
             bias = self.model.markov_bias(markov_embed)
             logits_i = base_logits[:, i] + bias
             if hasattr(self.model, "apply_output_logit_transform"):
@@ -210,16 +292,88 @@ def _sample_sequential(self, num_reqs: int, head_hidden: torch.Tensor) -> None:
             draft_sampled_i = self._sample_logits(
                 logits_i, idx_map[:, i], sample_pos[:, i], i
             )
-            self.draft_tokens[:num_reqs, i] = draft_sampled_i
+            out_tokens[:, i] = draft_sampled_i
             prev = draft_sampled_i
 
+        confidence = None
         if self.use_confidence_head:
             confidence = self.model.compute_confidence(
-                sample_hidden,
-                torch.stack(confidence_markov_embeds, dim=1).flatten(0, 1),
+                hidden, torch.stack(markov_embeds, dim=1).flatten(0, 1)
+            ).view(out_tokens.shape)
+        return confidence
+
+    def _sample_sequential_sharded(
+        self, num_reqs: int, head_hidden: torch.Tensor
+    ) -> None:
+        """_sample_sequential with each TP rank sampling only its own requests.
+
+        Request i is owned by rank shard_owner(i), as in DraftLogitsCache, so no
+        rank owns more than cdiv(num_reqs, tp_size) requests. Each rank computes
+        vocab-shard base logits for every request, in owner-sorted order. An
+        all-to-all then gives each rank the full-vocab rows of the requests it
+        owns, the Markov loop runs on just those rows, and an all-gather returns
+        the drafts in batch order. All shapes depend only on num_reqs, so the
+        step can be captured.
+        """
+        assert self._sharded_sequential_comm is not None
+        tp_group = get_tp_group()
+        tp_size, tp_rank = tp_group.world_size, tp_group.rank_in_group
+        n_spec = self.num_speculative_steps
+        capacity = cdiv(num_reqs, tp_size)
+        # Rank r owns requests r, r + tp_size, r + 2 * tp_size, ...
+        send_reqs = torch.arange(capacity, device=self.device) * tp_size + torch.arange(
+            tp_size, device=self.device
+        ).unsqueeze(-1)
+        # Mark padding slots with -1.
+        send_reqs = torch.where(send_reqs < num_reqs, send_reqs, -1)
+        req_idxs = torch.arange(num_reqs, device=self.device)
+        # Inverse of send_reqs (ignoring padding), used for restoring the
+        # all-gathered drafts/confidences into the original request order.
+        gathered_src_idxs = (req_idxs % tp_size) * capacity + req_idxs // tp_size
+
+        # Re-shard the draft logits from the vocab dimension to the batch dimension.
+        # Afterwards, each rank holds the full-vocab logits of its own requests.
+        # Padding slots read request 0's rows, and their results are discarded.
+        send_rows = send_reqs.clamp(min=0).unsqueeze(-1) * n_spec + self._step_cols
+        send_hidden = head_hidden[self.sample_indices[send_rows.flatten()]]
+        local_logits = self.model.compute_draft_logits_local(send_hidden)
+        num_local_rows = capacity * n_spec
+        base_logits = all_to_all_vocab_shards(
+            local_logits, [num_local_rows] * tp_size, num_local_rows
+        ).view(capacity, n_spec, -1)[..., : self.model.logits_processor.org_vocab_size]
+
+        # Sample draft tokens for the locally assigned requests.
+        local_reqs = send_reqs[tp_rank]
+        local_rows = send_rows[tp_rank]
+        local_tokens = self.draft_tokens.new_empty(capacity, n_spec)
+        local_hidden = send_hidden.view(tp_size, capacity * n_spec, -1)[tp_rank]
+        confidence = self._sample_markov_chain(
+            base_logits,
+            local_hidden,
+            self.input_buffers.input_ids[self._anchor_idx[local_reqs.clamp(min=0)]],
+            # Empty slots are inert: no draft logits cache writes.
+            torch.where(
+                local_reqs.view(-1, 1) >= 0, self.sample_idx_mapping[local_rows], -1
+            ),
+            self.sample_pos[local_rows],
+            local_tokens,
+        )
+
+        # Gather the draft tokens and confidences sampled by each rank.
+        packed = local_tokens
+        if confidence is not None:
+            packed = torch.cat(
+                [local_tokens, confidence.float().view(torch.int32).long()], dim=1
             )
-            self.draft_token_confidence_probs[:num_reqs] = confidence.view(
-                num_reqs, n_spec
+        gathered = packed.new_empty(tp_size * capacity, packed.shape[1])
+        self._sharded_sequential_comm.all_gather(gathered, packed.contiguous())
+
+        # Write the output draft tokens and confidences.
+        packed = gathered[gathered_src_idxs]
+        self.draft_tokens[:num_reqs] = packed[:, :n_spec]
+        if confidence is not None:
+            self.draft_token_confidence_probs[:num_reqs] = (
+                packed[:, n_spec:].int().view(torch.float32)
             )
 
     def _sample_sequential_topk(self, num_reqs: int, head_hidden: torch.Tensor) -> None:
@@ -289,4 +443,7 @@ def _generate_draft(
             num_tokens_across_dp,
             cudagraph_runtime_mode,
         )
-        self._sample_sequential(num_reqs, head_hidden)
+        if self._sharded_sequential_comm is not None:
+            self._sample_sequential_sharded(num_reqs, head_hidden)
+        else:
+            self._sample_sequential(num_reqs, head_hidden)
diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/utils.py b/vllm/v1/worker/gpu/spec_decode/dspark/utils.py
index 38108d4e0414..dc03e318d8c2 100644
--- a/vllm/v1/worker/gpu/spec_decode/dspark/utils.py
+++ b/vllm/v1/worker/gpu/spec_decode/dspark/utils.py
@@ -79,9 +79,9 @@ def load_dspark_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo
     from vllm.model_executor.models.qwen3_dflash import dflash_has_any_non_causal
     from vllm.model_executor.models.utils import get_draft_quant_config
     from vllm.v1.worker.gpu.spec_decode.eagle.utils import (
-        _should_share,
         get_target_lm_head,
         maybe_share_target_embed,
+        share_target_lm_head,
     )
 
     draft_attention_backend = _resolve_dspark_attention_backend(
@@ -133,19 +133,15 @@ def load_dspark_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo
     if draft_model_config.get_vocab_size() <= target_vocab_size:
         maybe_share_target_embed(draft_model, draft_inner, target_inner)
 
-    target_lm_head = get_target_lm_head(target_model, target_language_model)
-    draft_lm_head = getattr(draft_model, "lm_head", None)
     draft_output_vocab_size = (
         getattr(draft_model_config.hf_config, "draft_vocab_size", None)
         or draft_model_config.get_vocab_size()
     )
-    if (
-        target_lm_head is not None
-        and draft_output_vocab_size == target_vocab_size
-        and _should_share(draft_model, "has_own_lm_head", draft_lm_head, target_lm_head)
-    ):
-        if draft_lm_head is not None:
-            del draft_model.lm_head
-        draft_model.lm_head = target_lm_head
+    share_target_lm_head(
+        draft_model,
+        get_target_lm_head(target_model, target_language_model),
+        speculative_config.draft_lm_head_quantization,
+        can_share=draft_output_vocab_size == target_vocab_size,
+    )
 
     return draft_model
diff --git a/vllm/v1/worker/gpu/spec_decode/eagle/config.py b/vllm/v1/worker/gpu/spec_decode/eagle/config.py
new file mode 100644
index 000000000000..8f4644543038
--- /dev/null
+++ b/vllm/v1/worker/gpu/spec_decode/eagle/config.py
@@ -0,0 +1,80 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from vllm.config import SpeculativeConfig
+from vllm.transformers_utils.configs.eagle import EAGLEConfig
+from vllm.transformers_utils.configs.speculators import SpeculatorsConfig
+from vllm.v1.worker.gpu.spec_decode.speculator_config import (
+    DraftModelSpeculatorConfig,
+)
+
+
+def wrap_eagle_hf_config(spec: SpeculativeConfig, method: str) -> None:
+    """Wrap the draft hf_config in an EAGLEConfig for the built-in `method`,
+    unless it is one already."""
+    draft = spec.draft_model_config
+    if isinstance(draft.hf_config, (EAGLEConfig, SpeculatorsConfig)):
+        return
+    draft.hf_config = EAGLEConfig(draft.hf_config, method=method, model_type="eagle")
+    spec.update_arch_()
+
+
+def check_draft_lm_head_quantization(spec: SpeculativeConfig) -> None:
+    draft = spec.draft_model_config
+    if (
+        spec.draft_lm_head_quantization is not None
+        and draft is not None
+        and draft.head_dtype != draft.dtype
+    ):
+        raise ValueError(
+            "draft_lm_head_quantization needs the lm_head to run in the "
+            f"model dtype, got head_dtype={draft.head_dtype}."
+        )
+
+
+class EagleSpeculatorConfig(DraftModelSpeculatorConfig):
+    base_method = "eagle"
+    speculator_class = "vllm.v1.worker.gpu.spec_decode.eagle.speculator.EagleSpeculator"
+    supported_options = frozenset({"draft_lm_head_quantization"})
+    uses_target_hidden_states = True
+    supports_watermarking = True
+
+    @classmethod
+    def matches(cls, spec: SpeculativeConfig) -> bool:
+        # e.g. yuhuili/EAGLE-LLaMA3-Instruct-8B
+        return "eagle-" in spec.draft_model_config.model.lower()
+
+    @classmethod
+    def adapt_draft_model_config(cls, spec: SpeculativeConfig) -> None:
+        assert cls.base_method is not None
+        wrap_eagle_hf_config(spec, cls.base_method)
+
+    @classmethod
+    def apply_options(cls, spec: SpeculativeConfig) -> None:
+        check_draft_lm_head_quantization(spec)
+
+    @property
+    def num_lookahead_tokens(self) -> int:
+        return self.num_speculative_tokens
+
+    @property
+    def max_num_new_slots_for_drafting(self) -> int:
+        if self.spec_config.parallel_drafting:
+            # P-EAGLE (V1 model runner only): the existing query is reused; only
+            # masked queries need new slots.
+            return self.num_speculative_tokens - 1
+        return 0
+
+
+class Eagle3SpeculatorConfig(EagleSpeculatorConfig):
+    base_method = "eagle3"
+
+    @classmethod
+    def matches(cls, spec: SpeculativeConfig) -> bool:
+        # e.g. yuhuili/EAGLE3-LLaMA3.1-Instruct-8B, AngelSlim/Qwen3-8B_eagle3
+        return "eagle3" in spec.draft_model_config.model.lower()
+
+    @property
+    def uses_aux_hidden_states(self) -> bool:
+        hf_config = self.spec_config.draft_model_config.hf_config
+        eagle_config = getattr(hf_config, "eagle_config", None) or {}
+        return eagle_config.get("use_aux_hidden_state", True)
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 e434e58a9e95..86377496d395 100644
--- a/vllm/v1/worker/gpu/spec_decode/eagle/eagle3_utils.py
+++ b/vllm/v1/worker/gpu/spec_decode/eagle/eagle3_utils.py
@@ -12,19 +12,8 @@
 logger = init_logger(__name__)
 
 
-def requires_aux_hidden_states(spec_config: SpeculativeConfig) -> bool:
-    if spec_config.method == "eagle3":
-        eagle_config = (
-            getattr(spec_config.draft_model_config.hf_config, "eagle_config", None)
-            or {}
-        )
-        return eagle_config.get("use_aux_hidden_state", True)
-    return spec_config.method in ("dflash", "dspark", "extract_hidden_states")
-
-
 def set_eagle3_aux_hidden_state_layers(
-    model: nn.Module,
-    spec_config: SpeculativeConfig,
+    model: nn.Module, spec_config: SpeculativeConfig, use_pp: bool = False
 ) -> None:
     if not supports_eagle3(model):
         raise RuntimeError("Model does not support EAGLE3 interface")
@@ -41,6 +30,10 @@ def set_eagle3_aux_hidden_state_layers(
     eagle3_model.set_aux_hidden_state_layers(aux_layers)
     reserve_aux_intermediate_tensor_slots(model)
 
+    if use_pp:
+        assert spec_config.method is not None
+        verify_supports_aux_hidden_states_over_pp(model, spec_config.method)
+
 
 def _inner_decoder(model: nn.Module) -> nn.Module | None:
     parent_ref = model
diff --git a/vllm/v1/worker/gpu/spec_decode/eagle/utils.py b/vllm/v1/worker/gpu/spec_decode/eagle/utils.py
index c71395ec3f4e..8cba9af86421 100644
--- a/vllm/v1/worker/gpu/spec_decode/eagle/utils.py
+++ b/vllm/v1/worker/gpu/spec_decode/eagle/utils.py
@@ -1,5 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from typing import Literal
+
 import torch
 import torch.nn as nn
 
@@ -40,6 +42,36 @@ def _should_share(eagle: nn.Module, flag: str, draft, target) -> bool:
     return torch.equal(w, target_w)
 
 
+def share_target_lm_head(
+    draft_model: nn.Module,
+    target_lm_head: nn.Module | None,
+    quantization: Literal["fp8", "nvfp4"] | None,
+    can_share: bool = True,
+) -> nn.Module | None:
+    """Install the target lm_head, or its quantized copy, on a drafter that
+    shares it. Returns the installed head, or None if the drafter keeps its own.
+    """
+    draft_lm_head = getattr(draft_model, "lm_head", None)
+    if not (
+        can_share
+        and target_lm_head is not None
+        and _should_share(draft_model, "has_own_lm_head", draft_lm_head, target_lm_head)
+    ):
+        if quantization is not None:
+            raise ValueError(
+                "draft_lm_head_quantization needs a drafter that shares the "
+                "target lm_head."
+            )
+        return None
+    head = target_lm_head
+    if quantization is not None:
+        head = quantized_lm_head_copy(target_lm_head, quantization)
+    if draft_lm_head is not None:
+        del draft_model.lm_head
+    draft_model.lm_head = head
+    return head
+
+
 def get_target_lm_head(target_model: nn.Module, target_language_model: nn.Module):
     """The target's lm_head — from get_language_model() for
     *ForConditionalGeneration targets, else the top-level module."""
@@ -110,20 +142,12 @@ def load_eagle_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mod
 
     maybe_share_target_embed(eagle_model, draft_inner, target_inner)
 
-    target_lm_head = get_target_lm_head(target_model, target_language_model)
-    draft_lm_head = getattr(eagle_model, "lm_head", None)
-    if target_lm_head is not None and _should_share(
-        eagle_model, "has_own_lm_head", draft_lm_head, target_lm_head
-    ):
-        draft_head = target_lm_head
-        if speculative_config.draft_lm_head_quantization is not None:
-            draft_head = quantized_lm_head_copy(
-                target_lm_head, speculative_config.draft_lm_head_quantization
-            )
-        if draft_lm_head is not None:
-            del eagle_model.lm_head
-        eagle_model.lm_head = draft_head
-
+    draft_head = share_target_lm_head(
+        eagle_model,
+        get_target_lm_head(target_model, target_language_model),
+        speculative_config.draft_lm_head_quantization,
+    )
+    if draft_head is not None:
         # MTP layers route logits through layer.shared_head.head, not
         # eagle_model.lm_head, so the per-layer copies need fixing up too.
         layers = getattr(draft_inner, "layers", None)
@@ -134,10 +158,6 @@ def load_eagle_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mod
                 if sh is not None and hasattr(sh, "head"):
                     del sh.head
                     sh.head = draft_head
-    elif speculative_config.draft_lm_head_quantization is not None:
-        raise ValueError(
-            "draft_lm_head_quantization needs a drafter that shares the target lm_head."
-        )
 
     # MTP shares topk_indices_buffer with the target model. We update
     # every module in the draft that holds a buffer reference so that
diff --git a/vllm/v1/worker/gpu/spec_decode/extract_hidden_states.py b/vllm/v1/worker/gpu/spec_decode/extract_hidden_states.py
index e02d7ee11d70..8ede3b30afd6 100644
--- a/vllm/v1/worker/gpu/spec_decode/extract_hidden_states.py
+++ b/vllm/v1/worker/gpu/spec_decode/extract_hidden_states.py
@@ -10,9 +10,11 @@
 from vllm.config.compilation import CUDAGraphMode
 from vllm.forward_context import set_forward_context
 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.dp_utils import DPSyncState
 from vllm.v1.worker.gpu.input_batch import InputBatch
 from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator
+from vllm.v1.worker.gpu.spec_decode.utils import get_pp_safe_draft_load_config
 
 
 class ExtractHiddenStatesSpeculator(DraftModelSpeculator):
@@ -68,6 +70,9 @@ def load_draft_model(
             return get_model(
                 vllm_config=self.vllm_config,
                 model_config=self.draft_model_config,
+                load_config=get_pp_safe_draft_load_config(
+                    get_draft_load_config(self.vllm_config)
+                ),
             )
 
     def load_model(self, target_model: nn.Module) -> None:
@@ -134,6 +139,10 @@ def propose(
         num_tokens = stacked_hidden_states.shape[0]
         self.hidden_states[:num_tokens].copy_(stacked_hidden_states)
 
+        attn_metadata, slot_mappings = self._build_drafter_attn_metadata(
+            input_batch, attn_metadata, slot_mappings
+        )
+
         draft_attn_metadata = {
             name: attn_metadata[name] for name in self.draft_attn_layer_names
         }
diff --git a/vllm/v1/worker/gpu/spec_decode/extract_hidden_states_config.py b/vllm/v1/worker/gpu/spec_decode/extract_hidden_states_config.py
new file mode 100644
index 000000000000..159827e644c0
--- /dev/null
+++ b/vllm/v1/worker/gpu/spec_decode/extract_hidden_states_config.py
@@ -0,0 +1,48 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import copy
+from typing import TYPE_CHECKING
+
+from vllm.transformers_utils.configs.extract_hidden_states import (
+    ExtractHiddenStatesConfig,
+)
+from vllm.v1.worker.gpu.spec_decode.speculator_config import SpeculatorConfig
+
+if TYPE_CHECKING:
+    from vllm.config import SpeculativeConfig
+
+
+class ExtractHiddenStatesSpeculatorConfig(SpeculatorConfig):
+    base_method = "extract_hidden_states"
+    speculator_class = (
+        "vllm.v1.worker.gpu.spec_decode.extract_hidden_states."
+        "ExtractHiddenStatesSpeculator"
+    )
+    proposes_draft_tokens = False
+
+    @classmethod
+    def configure(cls, spec: "SpeculativeConfig") -> None:
+        # ExtractHiddenStatesModel is instantiated manually in load_model()
+        # We just need to store the target model config for KV cache shape info
+        spec.model = "extract_hidden_states"
+        spec.prompt_lookup_max = 0
+        spec.prompt_lookup_min = 0
+
+        draft_model_config = spec.draft_model_config
+        if hasattr(draft_model_config, "hf_config"):
+            hf_config = draft_model_config.hf_config.to_dict()
+        elif isinstance(draft_model_config, dict) and "hf_config" in draft_model_config:
+            hf_config = draft_model_config["hf_config"]
+        else:
+            hf_config = {}
+
+        spec.draft_model_config = copy.copy(spec.target_model_config)
+        spec.draft_model_config.hf_config = ExtractHiddenStatesConfig(
+            spec.draft_model_config.hf_config, **hf_config
+        )
+        spec.update_arch_()
+        spec.draft_parallel_config = spec.target_parallel_config
+
+    @property
+    def uses_aux_hidden_states(self) -> bool:
+        return True
diff --git a/vllm/v1/worker/gpu/spec_decode/gemma4/config.py b/vllm/v1/worker/gpu/spec_decode/gemma4/config.py
new file mode 100644
index 000000000000..fa369aeae8fb
--- /dev/null
+++ b/vllm/v1/worker/gpu/spec_decode/gemma4/config.py
@@ -0,0 +1,13 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from vllm.v1.worker.gpu.spec_decode.mtp.config import MTPSpeculatorConfig
+
+
+class Gemma4MTPSpeculatorConfig(MTPSpeculatorConfig):
+    speculator_class = (
+        "vllm.v1.worker.gpu.spec_decode.gemma4.speculator.Gemma4Speculator"
+    )
+    # Gemma-4 MTP scores tied embedding rows directly; no lm_head to share.
+    supported_options = MTPSpeculatorConfig.supported_options - {
+        "draft_lm_head_quantization"
+    }
diff --git a/vllm/v1/worker/gpu/spec_decode/lilicorr/config.py b/vllm/v1/worker/gpu/spec_decode/lilicorr/config.py
new file mode 100644
index 000000000000..db0a4c90020a
--- /dev/null
+++ b/vllm/v1/worker/gpu/spec_decode/lilicorr/config.py
@@ -0,0 +1,11 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from vllm.v1.worker.gpu.spec_decode.dflash.config import DFlashSpeculatorConfig
+
+
+class LiLiCorrSpeculatorConfig(DFlashSpeculatorConfig):
+    speculator_class = (
+        "vllm.v1.worker.gpu.spec_decode.lilicorr.speculator.LiLiCorrSpeculator"
+    )
+    # Draft-side watermarking doesn't cover candidate-head drafting yet.
+    supports_watermarking = False
diff --git a/vllm/v1/worker/gpu/spec_decode/mtp/config.py b/vllm/v1/worker/gpu/spec_decode/mtp/config.py
new file mode 100644
index 000000000000..47139f1c08be
--- /dev/null
+++ b/vllm/v1/worker/gpu/spec_decode/mtp/config.py
@@ -0,0 +1,126 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from typing import TYPE_CHECKING, Any, get_args
+
+from vllm.config.speculative import MTPModelTypes
+from vllm.logger import init_logger
+from vllm.v1.worker.gpu.spec_decode.eagle.config import EagleSpeculatorConfig
+from vllm.v1.worker.gpu.spec_decode.speculator_config import (
+    SpeculatorConfig,
+)
+
+if TYPE_CHECKING:
+    from vllm.config import SpeculativeConfig
+
+logger = init_logger(__name__)
+
+
+class MTPSpeculatorConfig(EagleSpeculatorConfig):
+    base_method = "mtp"
+    speculator_class = "vllm.v1.worker.gpu.spec_decode.mtp.speculator.MTPSpeculator"
+    supported_options = EagleSpeculatorConfig.supported_options | {
+        "index_share_for_mtp_iteration"
+    }
+    supports_pcp = True
+
+    @classmethod
+    def matches(cls, spec: "SpeculativeConfig") -> bool:
+        model_type = spec.draft_model_config.hf_config.model_type
+        return model_type in get_args(MTPModelTypes)
+
+    @classmethod
+    def specialize(cls, spec: "SpeculativeConfig") -> type[SpeculatorConfig]:
+        if not spec.use_multi_module_mtp():
+            return cls
+        # Imported lazily since it imports this module.
+        from vllm.v1.worker.gpu.spec_decode.multi_module_mtp.config import (
+            MultiModuleMTPSpeculatorConfig,
+        )
+
+        return MultiModuleMTPSpeculatorConfig
+
+    @classmethod
+    def default_model(cls, spec: "SpeculativeConfig") -> str | None:
+        target = spec.target_model_config
+        if target is None:
+            raise ValueError("target_model_config must be present for mtp")
+        # Align the quantization of draft model for cases such as
+        # --quantization fp8 with a bf16 checkpoint.
+        if not spec.quantization:
+            spec.quantization = target.quantization
+        # Use the draft model from the same model. Use model_weights if set
+        # (e.g. runai_streamer replaces model with a local cache dir but keeps
+        # the original path in model_weights), so the draft model can load
+        # weights from the same source as the target model.
+        return target.model_weights or target.model
+
+    @classmethod
+    def transformers_backend_hf_overrides(
+        cls, spec: "SpeculativeConfig"
+    ) -> dict[str, Any] | None:
+        target = spec.target_model_config
+        if (
+            target is None
+            or not target.using_transformers_backend()
+            or getattr(target.hf_text_config, "num_mtp_layers", None) is None
+        ):
+            return None
+        if target.hf_config is not target.hf_text_config:
+            raise NotImplementedError(
+                "The Transformers modeling backend does not support MTP for "
+                "multimodal models yet."
+            )
+        moe = "MoE" if target.is_moe else ""
+        return {
+            "architectures": [f"Transformers{moe}ForMTP"],
+            "n_predict": target.hf_text_config.num_mtp_layers,
+        }
+
+    @classmethod
+    def adapt_draft_model_config(cls, spec: "SpeculativeConfig") -> None:
+        target = spec.target_model_config
+        if target is not None and target.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 (
+            spec.num_speculative_tokens > 1
+            and spec.draft_model_config.hf_config.model_type
+            not in ("step3p5_mtp", "step5_mtp", "inkling_mtp")
+        ):
+            logger.warning(
+                "Enabling num_speculative_tokens > 1 will run "
+                "multiple times of forward on same MTP layer"
+                ",which may result in lower acceptance rate"
+            )
+
+    @classmethod
+    def resolve_num_speculative_tokens(cls, spec: "SpeculativeConfig") -> None:
+        if (
+            "ParseAuxMTPModel" in spec.draft_model_config.architectures
+            and spec.num_speculative_tokens not in (None, 1)
+        ):
+            raise ValueError("Parse 3.5 has only one trained MTP step")
+        super().resolve_num_speculative_tokens(spec)
+
+    @classmethod
+    def apply_options(cls, spec: "SpeculativeConfig") -> None:
+        super().apply_options(spec)
+        if spec.index_share_for_mtp_iteration is None:
+            return
+        if spec.draft_model_config is None:
+            raise ValueError(
+                "index_share_for_mtp_iteration is only supported with method='mtp'"
+            )
+        spec.draft_model_config.hf_config.index_share_for_mtp_iteration = (
+            spec.index_share_for_mtp_iteration
+        )
+
+    def hash_factors(self) -> list[Any]:
+        draft = self.spec_config.draft_model_config
+        if draft is None:
+            return []
+        return [getattr(draft.hf_config, "index_share_for_mtp_iteration", False)]
diff --git a/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/config.py b/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/config.py
new file mode 100644
index 000000000000..102d87599f3f
--- /dev/null
+++ b/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/config.py
@@ -0,0 +1,18 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from vllm.v1.worker.gpu.spec_decode.mtp.config import MTPSpeculatorConfig
+
+
+class MultiModuleMTPSpeculatorConfig(MTPSpeculatorConfig):
+    speculator_class = (
+        "vllm.v1.worker.gpu.spec_decode.multi_module_mtp.speculator."
+        "MultiModuleMTPSpeculator"
+    )
+    supports_pcp = False
+
+    @property
+    def num_prefill_lookahead_tokens(self) -> int:
+        # 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
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 666193011bf6..5cc90722675a 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
@@ -121,7 +121,7 @@ def capture(self) -> None:
         self.last_token_indices.zero_()
         assert self.cudagraph_manager is not None
         if self.cudagraph_manager.use_breakable_cg:
-            num_layers = self.draft_model_config.hf_config.num_nextn_predict_layers
+            num_layers = self.speculative_config.num_mtp_modules()
             self.cudagraph_manager.init_breakable_cg_runner(
                 self.model,
                 key_fn=lambda **kwargs: kwargs["spec_step_idx"] % num_layers,
@@ -214,6 +214,7 @@ def propose(
             dp_rank=self.dp_rank,
             need_eager=is_profile,
             dp_sync_state=dp_sync_state,
+            parallel_config=self.vllm_config.parallel_config,
         )
         num_tokens_across_dp = (
             batch_sync.num_tokens_across_dp if batch_sync is not None else None
diff --git a/vllm/v1/worker/gpu/spec_decode/ngram/config.py b/vllm/v1/worker/gpu/spec_decode/ngram/config.py
new file mode 100644
index 000000000000..56f0e381f722
--- /dev/null
+++ b/vllm/v1/worker/gpu/spec_decode/ngram/config.py
@@ -0,0 +1,47 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from typing import TYPE_CHECKING
+
+from vllm.v1.worker.gpu.spec_decode.speculator_config import SpeculatorConfig
+
+if TYPE_CHECKING:
+    from vllm.config import SpeculativeConfig
+
+
+class NgramSpeculatorConfig(SpeculatorConfig):
+    base_method = "ngram"
+    speculator_class = (
+        "vllm.v1.worker.gpu.spec_decode.ngram.speculator.NgramGPUSpeculator"
+    )
+    prefer_dense_token_history = True
+
+    @classmethod
+    def configure(cls, spec: "SpeculativeConfig") -> None:
+        # Set default values if not provided
+        if spec.prompt_lookup_min is None and spec.prompt_lookup_max is None:
+            # TODO(woosuk): Tune these values. They are arbitrarily chosen.
+            spec.prompt_lookup_min = 5
+            spec.prompt_lookup_max = 5
+        elif spec.prompt_lookup_min is None:
+            spec.prompt_lookup_min = spec.prompt_lookup_max
+        elif spec.prompt_lookup_max is None:
+            spec.prompt_lookup_max = spec.prompt_lookup_min
+
+        # Validate values
+        assert spec.prompt_lookup_min is not None
+        assert spec.prompt_lookup_max is not None
+        if spec.prompt_lookup_min > spec.prompt_lookup_max:
+            raise ValueError(
+                f"prompt_lookup_min={spec.prompt_lookup_min} must "
+                f"be <= prompt_lookup_max={spec.prompt_lookup_max}"
+            )
+
+        # TODO: current we still need extract vocab_size from target model
+        # config, in future, we may try refactor it out, and set
+        # draft related config as None here.
+        spec.draft_model_config = spec.target_model_config
+        spec.draft_parallel_config = spec.target_parallel_config
+
+
+class NgramGPUSpeculatorConfig(NgramSpeculatorConfig):
+    base_method = "ngram_gpu"
diff --git a/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py b/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py
index b74a0ec4f984..29c771f0d9f1 100644
--- a/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py
+++ b/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py
@@ -151,9 +151,6 @@ def _ngram_finalize_kernel(
 class NgramGPUSpeculator(BaseSpeculator):
     """V2-compatible GPU n-gram speculator."""
 
-    supports_mm_inputs = False
-    draft_logits = None
-
     def __init__(
         self,
         vllm_config: VllmConfig,
diff --git a/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py b/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py
index fd9808063c4a..0e490c8e793d 100644
--- a/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py
+++ b/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py
@@ -20,7 +20,7 @@
 from vllm.v1.worker.gpu.sample.logprob import compute_topk_scores
 from vllm.v1.worker.gpu.sample.output import SamplerOutput, SamplingMaskTensors
 from vllm.v1.worker.gpu.sample.sampler import Sampler
-from vllm.v1.worker.gpu.sample.states import NO_LOGPROBS
+from vllm.v1.worker.gpu.spec_decode.draft_logits import DraftLogitsCache
 from vllm.v1.worker.gpu.spec_decode.rejection_sampler_utils import (
     rejection_sample,
 )
@@ -152,7 +152,6 @@ def __init__(
         device: torch.device,
         *,
         watermark_key: int | None = None,
-        draft_logits_num_shards: int = 1,
     ):
         self.sampler = sampler
         self.watermark_key = watermark_key
@@ -165,9 +164,6 @@ def __init__(
         self.fly_window_size = 0
         self.fly_entropy_threshold = spec_config.fly_entropy_threshold
         self.fly_entropy_top_k = spec_config.fly_entropy_top_k
-        # Above 1 when each TP rank caches only the draft logits of the
-        # requests it owns (see DraftModelSpeculator).
-        self.draft_logits_num_shards = draft_logits_num_shards
         self.synthetic_conditional_rates: torch.Tensor | None = None
         if rejection_sample_method == "synthetic":
             assert spec_config.synthetic_acceptance_rates is not None
@@ -191,10 +187,12 @@ def _get_logprobs_tensors(
         logits: torch.Tensor,
         cu_num_logits: torch.Tensor,
         cu_num_logits_np: np.ndarray,
-        max_num_logprobs: int,
+        expanded_idx_mapping: torch.Tensor,
+        logprobs_dims: tuple[int, int] | None,
     ) -> LogprobsTensors | None:
-        if max_num_logprobs == NO_LOGPROBS:
+        if logprobs_dims is None:
             return None
+        num_logprobs, max_token_ids = logprobs_dims
 
         num_reqs = cu_num_logits.shape[0] - 1
         num_logits = logits.shape[0]
@@ -220,9 +218,12 @@ def _get_logprobs_tensors(
                 cu_num_generated_tokens = cu_num_logits_np.tolist()
         return compute_topk_scores(
             logits,
-            max_num_logprobs,
+            num_logprobs,
             flat_sampled,
             cu_num_generated_tokens,
+            logprob_token_ids_state=self.sampler.logprob_token_ids_state,
+            expanded_idx_mapping=expanded_idx_mapping,
+            max_per_req_token_ids=max_token_ids,
             logits_mode=self.sampler.logprobs_mode
             in ("raw_logits", "processed_logits"),
         )
@@ -254,7 +255,7 @@ def _watermarking_kwargs(
     def _verify(
         self,
         logits: torch.Tensor,
-        draft_logits: torch.Tensor | None,
+        draft_logits: DraftLogitsCache | None,
         draft_sampled: torch.Tensor,
         pos: torch.Tensor,
         cu_num_logits: torch.Tensor,
@@ -263,7 +264,6 @@ def _verify(
         expanded_idx_mapping: torch.Tensor,
         expanded_local_pos: torch.Tensor,
         seq_lens_upper_bound_np: np.ndarray,
-        verify_draft_sampled: torch.Tensor,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         processed_logits = self.sampler.apply_sampling_params(
             logits,
@@ -278,7 +278,7 @@ def _verify(
         sampled, num_sampled = rejection_sample(
             processed_logits,
             draft_logits,
-            verify_draft_sampled,
+            draft_sampled,
             cu_num_logits,
             pos,
             idx_mapping,
@@ -293,31 +293,31 @@ def _verify(
             fly_window_size=self.fly_window_size,
             fly_entropy_threshold=self.fly_entropy_threshold,
             fly_entropy_top_k=self.fly_entropy_top_k,
-            draft_logits_num_shards=self.draft_logits_num_shards,
             **self._watermarking_kwargs(
                 draft_sampled, expanded_idx_mapping, expanded_local_pos
             ),
         )
+        if self.sampler.trace_replay_state is not None:
+            self.sampler.trace_replay_state.apply_trace_to_verified(
+                sampled, num_sampled, draft_sampled, cu_num_logits, idx_mapping
+            )
         return processed_logits, sampled, num_sampled
 
     def _verify_in_chunks(
         self,
         logits: torch.Tensor,
         input_batch: InputBatch,
-        draft_logits: torch.Tensor | None,
+        draft_logits: DraftLogitsCache | None,
         draft_sampled: torch.Tensor,
         pos: torch.Tensor,
         max_chunk_logits: int,
-        max_num_logprobs: int,
-        verify_draft_sampled: torch.Tensor | None = None,
+        logprobs_dims: tuple[int, int] | None,
     ) -> tuple[
         torch.Tensor,
         torch.Tensor,
         LogprobsTensors | None,
         SamplingMaskTensors | None,
     ]:
-        if verify_draft_sampled is None:
-            verify_draft_sampled = draft_sampled
         cu_num_logits_np = input_batch.cu_num_logits_np
         use_processed_logits = self.sampler.logprobs_mode in PROCESSED_LOGPROBS_MODES
         num_reqs = input_batch.num_reqs
@@ -361,7 +361,6 @@ def _verify_in_chunks(
                 input_batch.expanded_idx_mapping[lo:hi],
                 input_batch.expanded_local_pos[lo:hi],
                 input_batch.seq_lens_cpu_upper_bound.numpy()[start:end],
-                verify_draft_sampled[lo:hi],
             )
             chunk_logprobs = self._get_logprobs_tensors(
                 sampled,
@@ -369,7 +368,8 @@ def _verify_in_chunks(
                 processed_logits if use_processed_logits else logits[lo:hi],
                 chunk_cu_num_logits,
                 chunk_cu_num_logits_np,
-                max_num_logprobs,
+                input_batch.expanded_idx_mapping[lo:hi],
+                logprobs_dims,
             )
             if chunk_logprobs is not None:
                 logprobs_chunks.append(chunk_logprobs)
@@ -422,8 +422,7 @@ def __call__(
         self,
         logits: torch.Tensor,
         input_batch: InputBatch,
-        draft_logits: torch.Tensor | None = None,
-        invalid_drafts: torch.Tensor | None = None,
+        draft_logits: DraftLogitsCache | None = None,
     ) -> SamplerOutput:
         # NOTE(woosuk): We intentionally compute num_nans before sampling to make clear
         # that num_nans is computed before applying penalties and temperature.
@@ -438,17 +437,7 @@ def __call__(
             self.sampler.req_states.prefill_len.gpu,
         )
 
-        # Marking a draft invalid (`is_valid_draft = draft_sampled >= 0`) pins the
-        # accepted length before it reaches a permissive bitmask row. Only
-        # verification sees the copy: `apply_sampling_params` and watermarking
-        # keep the real draft ids.
-        verify_draft_sampled = draft_sampled
-        if invalid_drafts is not None:
-            verify_draft_sampled = draft_sampled.masked_fill(invalid_drafts, -1)
-
-        max_num_logprobs = self.sampler.sampling_states.max_num_logprobs(
-            input_batch.idx_mapping_np
-        )
+        logprobs_dims = self.sampler.get_logprobs_dims(input_batch.idx_mapping_np)
         chunk_logit_limit = get_max_chunk_logits(logits.shape[1])
         (
             sampled,
@@ -462,8 +451,7 @@ def __call__(
             draft_sampled,
             pos,
             chunk_logit_limit,
-            max_num_logprobs,
-            verify_draft_sampled,
+            logprobs_dims,
         )
 
         num_sampled, num_rejected = get_num_sampled_and_rejected(
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 0b0139fbba28..bf86dcf0ea5e 100644
--- a/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py
+++ b/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py
@@ -4,13 +4,16 @@
 
 from vllm.triton_utils import tl, tldevice, triton
 from vllm.v1.spec_decode.fly import compute_fly_entropy
-from vllm.v1.worker.gpu.sample.batch_shard import shard_row
 from vllm.v1.worker.gpu.sample.gumbel import (
     gumbel_block_argmax,
     spec_step_seed,
     tl_rand32,
 )
 from vllm.v1.worker.gpu.sample.watermark import philox_gumbel_block_argmax
+from vllm.v1.worker.gpu.spec_decode.draft_logits import (
+    DraftLogitsCache,
+    draft_logits_row,
+)
 
 
 @triton.jit
@@ -142,6 +145,7 @@ def _compute_global_logprobs_and_logsumexp(
     draft_logits_ptr,
     draft_logits_stride_0,
     draft_logits_stride_1,
+    draft_logits_row_mapping_ptr,
     # [num_logits, num_blocks]
     draft_local_max_ptr,
     draft_local_max_stride,
@@ -150,7 +154,6 @@ def _compute_global_logprobs_and_logsumexp(
     vocab_num_blocks,
     PADDED_VOCAB_NUM_BLOCKS: tl.constexpr,
     HAS_DRAFT_LOGITS: tl.constexpr,
-    DRAFT_LOGITS_NUM_SHARDS: tl.constexpr,
 ):
     target_logit = tl.load(
         target_logits_ptr + logit_idx * target_logits_stride + token,
@@ -172,7 +175,9 @@ def _compute_global_logprobs_and_logsumexp(
         draft_logit = (
             tl.load(
                 draft_logits_ptr
-                + shard_row(req_state_idx, DRAFT_LOGITS_NUM_SHARDS)
+                + tl.maximum(
+                    draft_logits_row(req_state_idx, draft_logits_row_mapping_ptr), 0
+                )
                 * draft_logits_stride_0
                 + draft_step * draft_logits_stride_1
                 + token,
@@ -208,7 +213,8 @@ def _is_bonus_logit(expanded_local_pos_ptr, logit_idx, num_logits):
     return next_local_pos == 0
 
 
-@triton.jit
+# num_logits only bounds the bonus-logit lookahead; one kernel serves every value.
+@triton.jit(do_not_specialize=["num_logits"])
 def _compute_local_logits_stats_kernel(
     # [num_logits, num_blocks]
     target_local_argmax_ptr,
@@ -232,6 +238,7 @@ def _compute_local_logits_stats_kernel(
     draft_logits_ptr,
     draft_logits_stride_0,
     draft_logits_stride_1,
+    draft_logits_row_mapping_ptr,
     # [num_logits]
     expanded_idx_mapping_ptr,
     # [num_logits]
@@ -242,7 +249,6 @@ def _compute_local_logits_stats_kernel(
     num_logits,
     BLOCK_SIZE: tl.constexpr,
     HAS_DRAFT_LOGITS: tl.constexpr,
-    DRAFT_LOGITS_NUM_SHARDS: tl.constexpr,
 ):
     logit_idx = tl.program_id(0).to(tl.int64)
 
@@ -301,7 +307,9 @@ def _compute_local_logits_stats_kernel(
             draft_logits = (
                 tl.load(
                     draft_logits_ptr
-                    + shard_row(req_state_idx, DRAFT_LOGITS_NUM_SHARDS)
+                    + tl.maximum(
+                        draft_logits_row(req_state_idx, draft_logits_row_mapping_ptr), 0
+                    )
                     * draft_logits_stride_0
                     + draft_step_idx * draft_logits_stride_1
                     + block_offsets,
@@ -342,6 +350,7 @@ def _compute_cumulative_log_p_kernel(
     draft_logits_ptr,
     draft_logits_stride_0,
     draft_logits_stride_1,
+    draft_logits_row_mapping_ptr,
     # [num_logits, num_blocks]
     draft_local_max_ptr,
     draft_local_max_stride,
@@ -357,7 +366,6 @@ def _compute_cumulative_log_p_kernel(
     vocab_num_blocks,
     PADDED_VOCAB_NUM_BLOCKS: tl.constexpr,
     HAS_DRAFT_LOGITS: tl.constexpr,
-    DRAFT_LOGITS_NUM_SHARDS: tl.constexpr,
 ):
     req_idx = tl.program_id(0)
     req_state_idx = tl.load(idx_mapping_ptr + req_idx).to(tl.int64)
@@ -392,6 +400,7 @@ def _compute_cumulative_log_p_kernel(
                     draft_logits_ptr,
                     draft_logits_stride_0,
                     draft_logits_stride_1,
+                    draft_logits_row_mapping_ptr,
                     draft_local_max_ptr,
                     draft_local_max_stride,
                     draft_local_sumexp_ptr,
@@ -399,14 +408,13 @@ def _compute_cumulative_log_p_kernel(
                     vocab_num_blocks,
                     PADDED_VOCAB_NUM_BLOCKS,
                     HAS_DRAFT_LOGITS,
-                    DRAFT_LOGITS_NUM_SHARDS,
                 )
             )
             log_p = tl.minimum(log_p + (target_logprob - draft_logprob), 0.0)
         tl.store(cumulative_log_p_ptr + logit_idx, log_p)
 
 
-@triton.jit
+@triton.jit(do_not_specialize=["num_logits"])
 def _compute_local_residual_mass_kernel(
     # [num_logits, num_blocks]
     local_residual_mass_ptr,
@@ -426,6 +434,7 @@ def _compute_local_residual_mass_kernel(
     draft_logits_ptr,
     draft_logits_stride_0,
     draft_logits_stride_1,
+    draft_logits_row_mapping_ptr,
     # [num_logits, num_blocks]
     draft_local_max_ptr,
     draft_local_max_stride,
@@ -445,7 +454,6 @@ def _compute_local_residual_mass_kernel(
     vocab_num_blocks,
     BLOCK_SIZE: tl.constexpr,
     PADDED_VOCAB_NUM_BLOCKS: tl.constexpr,
-    DRAFT_LOGITS_NUM_SHARDS: tl.constexpr,
 ):
     logit_idx = tl.program_id(0).to(tl.int64)
     # Adaptive verification can end before num_speculative_steps.
@@ -487,6 +495,7 @@ def _compute_local_residual_mass_kernel(
         draft_logits_ptr,
         draft_logits_stride_0,
         draft_logits_stride_1,
+        draft_logits_row_mapping_ptr,
         draft_local_max_ptr,
         draft_local_max_stride,
         draft_local_sumexp_ptr,
@@ -494,7 +503,6 @@ def _compute_local_residual_mass_kernel(
         vocab_num_blocks,
         PADDED_VOCAB_NUM_BLOCKS,
         True,  # HAS_DRAFT_LOGITS
-        DRAFT_LOGITS_NUM_SHARDS,
     )
 
     # Compute the residual mass: max(p_i * M_b(x|x_{<i}) - M_s(x|x_{<i}), 0)
@@ -537,6 +545,7 @@ def _rejection_kernel(
     draft_logits_ptr,
     draft_logits_stride_0,
     draft_logits_stride_1,
+    draft_logits_row_mapping_ptr,
     # [num_logits, num_blocks]
     draft_local_max_ptr,
     draft_local_max_stride,
@@ -566,7 +575,6 @@ def _rejection_kernel(
     vocab_num_blocks,
     PADDED_VOCAB_NUM_BLOCKS: tl.constexpr,
     HAS_DRAFT_LOGITS: tl.constexpr,
-    DRAFT_LOGITS_NUM_SHARDS: tl.constexpr,
     SYNTHETIC_MODE: tl.constexpr,
     USE_BLOCK_VERIFICATION: tl.constexpr,
     FLY_WINDOW_SIZE: tl.constexpr,
@@ -686,6 +694,7 @@ def _rejection_kernel(
                         draft_logits_ptr,
                         draft_logits_stride_0,
                         draft_logits_stride_1,
+                        draft_logits_row_mapping_ptr,
                         draft_local_max_ptr,
                         draft_local_max_stride,
                         draft_local_sumexp_ptr,
@@ -693,7 +702,6 @@ def _rejection_kernel(
                         vocab_num_blocks,
                         PADDED_VOCAB_NUM_BLOCKS,
                         HAS_DRAFT_LOGITS,
-                        DRAFT_LOGITS_NUM_SHARDS,
                     )
                 )
                 if FLY_WINDOW_SIZE == 0 or pending_rejection < 0:
@@ -804,6 +812,7 @@ def _seeded_resample_argmax(
         None,  # logits_cache_col_ptr
         None,  # logits_cache_source_ptr
         0,  # logits_cache_source_stride
+        None,  # logits_cache_row_mapping_ptr
         vocab_size,
         IS_DRAFTING=False,
         APPLY_TEMPERATURE=False,
@@ -828,6 +837,7 @@ def _resample_kernel(
     draft_logits_ptr,
     draft_logits_stride_0,
     draft_logits_stride_1,
+    draft_logits_row_mapping_ptr,
     # [num_reqs]
     draft_rejected_logsumexp_ptr,
     # [num_reqs]
@@ -857,7 +867,6 @@ def _resample_kernel(
     vocab_size,
     BLOCK_SIZE: tl.constexpr,
     HAS_DRAFT_LOGITS: tl.constexpr,
-    DRAFT_LOGITS_NUM_SHARDS: tl.constexpr,
     USE_FP64: tl.constexpr,
     USE_BLOCK_VERIFICATION: tl.constexpr,
     CONTEXT_WIDTH: tl.constexpr,
@@ -923,7 +932,9 @@ def _resample_kernel(
         draft_logits = (
             tl.load(
                 draft_logits_ptr
-                + shard_row(req_state_idx, DRAFT_LOGITS_NUM_SHARDS)
+                + tl.maximum(
+                    draft_logits_row(req_state_idx, draft_logits_row_mapping_ptr), 0
+                )
                 * draft_logits_stride_0
                 + resample_idx * draft_logits_stride_1
                 + block,
@@ -1098,10 +1109,7 @@ def _insert_resampled_kernel(
 def rejection_sample(
     # [num_logits, V]
     target_logits: torch.Tensor,
-    # [max_num_reqs, num_speculative_steps, V], or under batch-sharded
-    # sampling, [cdiv(max_num_reqs, draft_logits_num_shards), ...] holding this
-    # rank's requests at row req_state_idx // draft_logits_num_shards.
-    draft_logits: torch.Tensor | None,
+    draft_logits_cache: DraftLogitsCache | None,
     # [num_logits]
     draft_sampled: torch.Tensor,
     # [num_reqs + 1]
@@ -1130,9 +1138,13 @@ def rejection_sample(
     fly_window_size: int = 0,
     fly_entropy_threshold: float = 0.3,
     fly_entropy_top_k: int = 3,
-    draft_logits_num_shards: int = 1,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     assert target_logits.ndim == 2 and target_logits.stride(-1) == 1
+    draft_logits = None
+    draft_logits_row_mapping = None
+    if draft_logits_cache is not None:
+        draft_logits = draft_logits_cache.values
+        draft_logits_row_mapping = draft_logits_cache.row_mapping
     assert draft_logits is None or (
         draft_logits.ndim == 3 and draft_logits.stride(-1) == 1
     )
@@ -1225,6 +1237,7 @@ def rejection_sample(
         draft_logits,
         draft_logits_stride_0,
         draft_logits_stride_1,
+        draft_logits_row_mapping,
         expanded_idx_mapping,
         expanded_local_pos,
         temperature,
@@ -1232,7 +1245,6 @@ def rejection_sample(
         num_logits,
         BLOCK_SIZE=VOCAB_BLOCK_SIZE,
         HAS_DRAFT_LOGITS=has_draft_logits,
-        DRAFT_LOGITS_NUM_SHARDS=draft_logits_num_shards,
     )
 
     # Precompute the running joint ratio and residual mass for block
@@ -1258,6 +1270,7 @@ def rejection_sample(
             draft_logits,
             draft_logits_stride_0,
             draft_logits_stride_1,
+            draft_logits_row_mapping,
             draft_local_max,
             draft_local_max.stride(0),
             draft_local_sumexp,
@@ -1268,7 +1281,6 @@ def rejection_sample(
             vocab_num_blocks,
             PADDED_VOCAB_NUM_BLOCKS=padded_vocab_num_blocks,
             HAS_DRAFT_LOGITS=has_draft_logits,
-            DRAFT_LOGITS_NUM_SHARDS=draft_logits_num_shards,
             num_warps=1,
         )
 
@@ -1293,6 +1305,7 @@ def rejection_sample(
                 draft_logits,
                 draft_logits_stride_0,
                 draft_logits_stride_1,
+                draft_logits_row_mapping,
                 draft_local_max,
                 draft_local_max.stride(0),
                 draft_local_sumexp,
@@ -1306,7 +1319,6 @@ def rejection_sample(
                 vocab_num_blocks,
                 BLOCK_SIZE=VOCAB_BLOCK_SIZE,
                 PADDED_VOCAB_NUM_BLOCKS=padded_vocab_num_blocks,
-                DRAFT_LOGITS_NUM_SHARDS=draft_logits_num_shards,
             )
         else:
             local_residual_mass = None
@@ -1340,6 +1352,7 @@ def rejection_sample(
         draft_logits,
         draft_logits_stride_0,
         draft_logits_stride_1,
+        draft_logits_row_mapping,
         draft_local_max,
         draft_local_max.stride(0),
         draft_local_sumexp,
@@ -1358,7 +1371,6 @@ def rejection_sample(
         vocab_num_blocks,
         PADDED_VOCAB_NUM_BLOCKS=padded_vocab_num_blocks,
         HAS_DRAFT_LOGITS=has_draft_logits,
-        DRAFT_LOGITS_NUM_SHARDS=draft_logits_num_shards,
         SYNTHETIC_MODE=synthetic_conditional_rates is not None,
         USE_BLOCK_VERIFICATION=use_block_verification,
         FLY_WINDOW_SIZE=fly_window_size,
@@ -1388,6 +1400,7 @@ def rejection_sample(
         draft_logits,
         draft_logits_stride_0,
         draft_logits_stride_1,
+        draft_logits_row_mapping,
         draft_rejected_logsumexp,
         num_sampled,
         cu_num_logits,
@@ -1406,7 +1419,6 @@ def rejection_sample(
         vocab_size,
         BLOCK_SIZE=RESAMPLE_BLOCK_SIZE,
         HAS_DRAFT_LOGITS=has_draft_logits,
-        DRAFT_LOGITS_NUM_SHARDS=draft_logits_num_shards,
         USE_FP64=use_fp64,
         USE_BLOCK_VERIFICATION=use_block_verification,
         CONTEXT_WIDTH=context_width,
diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py
index 261c83a94727..f964dd55cd59 100644
--- a/vllm/v1/worker/gpu/spec_decode/speculator.py
+++ b/vllm/v1/worker/gpu/spec_decode/speculator.py
@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import copy
 from abc import ABC, abstractmethod
 from collections.abc import Mapping
 from dataclasses import replace
@@ -11,20 +12,21 @@
 
 from vllm.config import VllmConfig, get_layers_from_vllm_config
 from vllm.config.compilation import CUDAGraphMode
+from vllm.config.speculative import _DRAFT_VLLM_CONFIG_OVERRIDES
 from vllm.distributed.eplb.eplb_state import EplbState
-from vllm.distributed.parallel_state import get_tp_group
 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.utils.math_utils import cdiv
 from vllm.v1.kv_cache_interface import KVCacheConfig
 from vllm.v1.watermarking import create_watermarker
+from vllm.v1.watermarking.gpu_sampler import GPUWatermarkSampler
 from vllm.v1.watermarking.spec_decode import (
     DraftWatermarker,
     create_speculative_draft_watermarker,
 )
 from vllm.v1.worker.gpu.attn_utils import (
     build_attn_metadata,
+    build_slot_mappings_by_layer,
     init_attn_backend,
 )
 from vllm.v1.worker.gpu.block_table import BlockTables
@@ -34,15 +36,17 @@
 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.batch_shard import batch_sharded_sampling_enabled
+from vllm.v1.worker.gpu.sample.greedy_argmax import greedy_argmax
 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.gpu.spec_decode.draft_logits import DraftLogitsCache
 from vllm.v1.worker.utils import AttentionGroup
 
 if TYPE_CHECKING:
-    from vllm.v1.watermarking.gpu_sampler import GPUWatermarkSampler
     from vllm.v1.worker.gpu.pcp_manager import PCPManager
+    from vllm.v1.worker.gpu.sample.sampler import Sampler
 
 logger = init_logger(__name__)
 
@@ -61,9 +65,23 @@ def _target_feeds_hc_residual(vllm_config: VllmConfig) -> bool:
 
 
 class BaseSpeculator(ABC):
+    """Base class for V2 model runner speculators.
+
+    This is also the interface for custom speculators, named by the
+    `speculator_class` of their registered `SpeculatorConfig` (see
+    `register_speculative_method`). Speculators that run a draft model
+    should subclass `DraftModelSpeculator` (or `TargetDependentARSpeculator` if
+    drafting consumes target model state), which the model runner loads and
+    allocates draft KV cache for.
+    """
+
     num_query_per_req: int = 1
     # Extra query slots reserved per request outside the regular queries.
     num_extra_query_per_req: int = 0
+    supports_mm_inputs: bool = False
+    # Draft logits for probabilistic rejection sampling, or None for
+    # greedy-draft verification and for speculators without a draft model.
+    draft_logits: DraftLogitsCache | None = None
 
     @abstractmethod
     def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None:
@@ -104,8 +122,18 @@ def propose(
     ) -> torch.Tensor:
         pass
 
+    def prepare_watermarking(
+        self, sampler: "GPUWatermarkSampler", idx_mapping: torch.Tensor
+    ) -> None:
+        """Prepare draft-side watermarking before `propose`. No-op by default."""
+        return
+
 
 class DraftModelSpeculator(BaseSpeculator):
+    # Draft prefill runs on the target's batch, so it can reuse the target
+    # forward's attention metadata. Otherwise the drafter rebuilds it.
+    can_reuse_target_attn_metadata: bool = True
+
     def __init__(self, vllm_config: VllmConfig, device: torch.device):
         self.vllm_config = vllm_config
         self.device = device
@@ -145,6 +173,9 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device):
         self.dp_size = vllm_config.parallel_config.data_parallel_size
         self.dp_rank = vllm_config.parallel_config.data_parallel_rank
         self.dcp_size = vllm_config.parallel_config.decode_context_parallel_size
+        if vllm_config.parallel_config.use_ubatching:
+            # A microbatched target forward builds only per-microbatch metadata.
+            self.can_reuse_target_attn_metadata = False
 
         self.eplb_state: EplbState | None = None
 
@@ -180,12 +211,10 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device):
 
         # Pre-temperature logits, cached from the previous decode step.
         # Allocated in load_model. See _init_draft_logits.
-        self.draft_logits: torch.Tensor | None = None
+        self.draft_logits: DraftLogitsCache | None = None
         self.use_draft_logits = (
             self.speculative_config.draft_sample_method == "probabilistic"
         )
-        self.draft_logits_num_shards = 1
-        self.draft_logits_shard_idx = 0
 
         self.draft_watermarker: DraftWatermarker | None = None
         watermark_config = getattr(vllm_config, "watermark_config", None)
@@ -266,7 +295,7 @@ def load_model(self, target_model: nn.Module) -> None:
 
     def supports_batch_sharded_draft_logits(self) -> bool:
         """Whether every draft logits write on this speculator's path honors
-        draft_logits_num_shards and draft_logits_shard_idx."""
+        DraftLogitsCache.row_mapping."""
         return False
 
     def _init_draft_logits(self, target_model: nn.Module) -> None:
@@ -274,27 +303,25 @@ def _init_draft_logits(self, target_model: nn.Module) -> None:
 
         Under batch-sharded sampling, the rank that verifies a request is the
         only one that reads its draft logits, so each TP rank keeps just the
-        rows of the requests it owns, at row
-        req_state_idx // draft_logits_num_shards.
+        rows of the requests it owns. See DraftLogitsCache.
         """
-        if (
-            batch_sharded_sampling_enabled(self.vllm_config, target_model)
-            and self.draft_watermarker is None
-            and self.supports_batch_sharded_draft_logits()
-        ):
-            tp_group = get_tp_group()
-            self.draft_logits_num_shards = tp_group.world_size
-            self.draft_logits_shard_idx = tp_group.rank_in_group
+        parallel_config = self.vllm_config.parallel_config
         dtype, fill = self.draft_logits_spec(self.vllm_config)
-        self.draft_logits = torch.full(
-            (
-                cdiv(self.max_num_reqs, self.draft_logits_num_shards),
-                self.num_speculative_steps,
-                self.vocab_size,
-            ),
-            fill,
+        self.draft_logits = DraftLogitsCache(
+            self.max_num_reqs,
+            self.num_speculative_steps,
+            self.vocab_size,
             dtype=dtype,
+            fill=fill,
             device=self.device,
+            sharded=(
+                batch_sharded_sampling_enabled(self.vllm_config, target_model)
+                # PCPManager.sync_drafts indexes the draft logits by request
+                # state, so it must see every request's logits.
+                and parallel_config.prefill_context_parallel_size == 1
+                and self.draft_watermarker is None
+                and self.supports_batch_sharded_draft_logits()
+            ),
         )
 
     def set_eplb_state(self, eplb_state: EplbState) -> None:
@@ -313,7 +340,16 @@ def _prepare_eplb_forward(self, num_unpadded_tokens: int) -> None:
     def attn_vllm_config(self) -> VllmConfig:
         """Config for the draft's attention metadata builders. Overridden by
         speculators whose attention mode differs from the target's."""
-        return self.vllm_config
+        # Shallow copies instead of apply_draft_overrides: its replace()
+        # rebuilds each sub-config and drops engine-resolved init=False state
+        # such as cache_config.kv_cache_layout.
+        config = copy.copy(self.vllm_config)
+        for src, config_name, dst in _DRAFT_VLLM_CONFIG_OVERRIDES:
+            if (value := getattr(self.speculative_config, src)) is not None:
+                sub_config = copy.copy(getattr(config, config_name))
+                setattr(sub_config, dst, value)
+                setattr(config, config_name, sub_config)
+        return config
 
     def set_attn(
         self,
@@ -340,6 +376,41 @@ def set_attn(
         self.target_attn_groups = target_attn_groups
         self._configure_decode_metadata_update()
 
+    def _build_drafter_attn_metadata(
+        self,
+        input_batch: InputBatch,
+        attn_metadata: dict[str, Any],
+        slot_mappings: dict[str, torch.Tensor],
+        cg_mode: CUDAGraphMode = CUDAGraphMode.NONE,
+        for_capture: bool = False,
+    ) -> tuple[dict[str, Any], dict[str, torch.Tensor]]:
+        """Attention metadata and slot mappings for draft prefill on the target
+        batch: the target forward's when reusable, else rebuilt on the draft
+        builders from the block tables and slot mappings the target prepared.
+        Must run before draft decode steps overwrite the slot mappings.
+        """
+        if self.can_reuse_target_attn_metadata:
+            return attn_metadata, slot_mappings
+        block_tables = tuple(
+            block_table[: input_batch.num_reqs_after_padding]
+            for block_table in self.block_tables.input_block_tables
+        )
+        target_slot_mappings = self.block_tables.slot_mappings[
+            :, : input_batch.num_tokens_after_padding
+        ]
+        attn_metadata = self.model_state.prepare_attn(
+            input_batch,
+            cg_mode,
+            block_tables,
+            target_slot_mappings,
+            self.attn_groups,
+            self.kv_cache_config,
+            for_capture=for_capture,
+        )
+        return attn_metadata, build_slot_mappings_by_layer(
+            target_slot_mappings, self.kv_cache_config
+        )
+
     def _configure_decode_metadata_update(self) -> None:
         unsupported_backends = sorted(
             {
@@ -480,7 +551,7 @@ def sample_draft(
         temperature: torch.Tensor,
         seeds: torch.Tensor,
         draft_step: torch.Tensor,
-        draft_logits: torch.Tensor | None,
+        draft_logits: DraftLogitsCache | None,
         spec_step_idx: int = 0,
         num_steps: int = 1,
     ) -> torch.Tensor:
@@ -497,7 +568,7 @@ def sample_draft(
                 sample_src_positions,
                 apply_temperature=True,
                 is_drafting=True,
-                logits_cache=draft_logits,
+                logits_cache=draft_logits.values,
                 logits_cache_col=draft_step,
                 use_fp64=self.use_fp64_gumbel,
                 num_steps=num_steps,
@@ -511,14 +582,13 @@ def sample_draft(
                 sample_src_positions,
                 apply_temperature=True,
                 is_drafting=True,
-                logits_cache=draft_logits,
+                logits_cache=draft_logits.values,
                 logits_cache_col=draft_step,
                 use_fp64=self.use_fp64_gumbel,
-                logits_cache_num_shards=self.draft_logits_num_shards,
-                logits_cache_shard_idx=self.draft_logits_shard_idx,
+                logits_cache_row_mapping=draft_logits.row_mapping,
             )
         else:
-            sampled = logits.argmax(dim=-1)
+            sampled = greedy_argmax(logits)
         self._maybe_predict_acceptance(logits, idx_mapping, draft_step)
         return sampled
 
@@ -550,7 +620,7 @@ def observe_verification(
 
     def prepare_watermarking(
         self,
-        sampler: "GPUWatermarkSampler",
+        sampler: GPUWatermarkSampler,
         idx_mapping: torch.Tensor,
     ) -> None:
         if self.draft_watermarker is None:
@@ -563,6 +633,24 @@ def prepare_watermarking(
             sampler.req_states.total_len.gpu,
         )
 
+    def prepare_propose(
+        self,
+        input_batch: InputBatch,
+        num_sampled: torch.Tensor,
+        num_rejected: torch.Tensor,
+        sampler: "Sampler",
+    ) -> None:
+        """Per-step work between verifying the last drafts and the next
+        `propose`."""
+        self.observe_verification(input_batch.idx_mapping, num_sampled, num_rejected)
+        if isinstance(sampler, GPUWatermarkSampler):
+            self.prepare_watermarking(sampler, input_batch.idx_mapping)
+        if self.draft_logits is not None:
+            # Verification read the previous step's draft logits.
+            self.draft_logits.assign(
+                input_batch.idx_mapping, input_batch.idx_mapping_np
+            )
+
     def _copy_request_inputs(
         self,
         num_reqs: int,
@@ -608,6 +696,7 @@ def _build_uniform_batch_dp_sync(
                 target_dp_sync_state.num_tokens_across_dp, num_batch_tokens
             ),
             uniform_token_count=num_query_per_req,
+            max_query_len=num_query_per_req,
             eager=False,
         ), num_batch_tokens
 
diff --git a/vllm/v1/worker/gpu/spec_decode/speculator_config.py b/vllm/v1/worker/gpu/spec_decode/speculator_config.py
new file mode 100644
index 000000000000..8ae4d0ba41f9
--- /dev/null
+++ b/vllm/v1/worker/gpu/spec_decode/speculator_config.py
@@ -0,0 +1,439 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Per-speculator configuration.
+
+Each speculator has a `SpeculatorConfig` subclass, in a `config.py` next to its
+V2 model runner implementation, which sets up `SpeculativeConfig` for it,
+describes its behavior to code outside the speculator (queried through
+`SpeculativeConfig.speculator_config` rather than by method name), and names
+its implementation class. Config modules must not import the speculators, as
+they're used in every process. Custom speculators are registered with
+`register_speculative_method`.
+
+Parts only needed by the V1 model runner are marked "V1 model runner only", and
+will be removed with it.
+"""
+
+from typing import TYPE_CHECKING, Any, ClassVar, final, get_args
+
+from vllm.utils.import_utils import resolve_obj_by_qualname
+
+# vllm.config imports this module, so it imports vllm.config lazily.
+if TYPE_CHECKING:
+    from vllm.config import ParallelConfig, SpeculativeConfig
+
+_PKG = "vllm.v1.worker.gpu.spec_decode"
+
+
+class SpeculatorConfig:
+    """Configuration of a kind of speculator.
+
+    The defaults describe a speculator without a draft checkpoint or special
+    requirements. A custom speculator can subclass this, or the config of the
+    built-in speculator it behaves like (overriding `speculator_class`), whose
+    setup and behavior it then inherits.
+    """
+
+    # For in-tree speculators, the `SpeculativeConfig.method` that enables it.
+    # For out-of-tree custom speculators, the method of the built-in one it subclasses.
+    base_method: ClassVar[str | None] = None
+    # Fully-qualified name of the V2 `BaseSpeculator` subclass, if any.
+    speculator_class: ClassVar[str | None] = None
+    # Method-specific `SpeculativeConfig` options that this speculator accepts
+    # (e.g. "speculator_options" for a custom speculator's own options).
+    supported_options: ClassVar[frozenset[str]] = frozenset()
+
+    # Uses the target model's hidden states (EAGLE or EAGLE3-style).
+    uses_target_hidden_states: ClassVar[bool] = False
+    # Drafts in its own workspace lane.
+    uses_separate_workspace_lane: ClassVar[bool] = False
+    # For performance, prefers requests' full token history on the GPU instead of
+    # UVA memory (e.g. ngram).
+    prefer_dense_token_history: ClassVar[bool] = False
+    # Proposes draft tokens (unlike e.g. hidden state extraction).
+    proposes_draft_tokens: ClassVar[bool] = True
+    adaptive_verification_uses_draft_logits: ClassVar[bool] = True
+    # The draft model has its own vocabulary, which must match the target's
+    # unless use_heterogeneous_vocab is set.
+    draft_has_own_vocab: ClassVar[bool] = False
+
+    supports_watermarking: ClassVar[bool] = False
+    # Prefill context parallelism.
+    supports_pcp: ClassVar[bool] = False
+
+    # --- SpeculativeConfig setup ---
+
+    @classmethod
+    def matches(cls, spec: "SpeculativeConfig") -> bool:
+        """Whether the draft checkpoint is for this speculator, other than by
+        its model's declaration (e.g. by name). Used to detect the method."""
+        return False
+
+    @classmethod
+    def specialize(cls, spec: "SpeculativeConfig") -> type["SpeculatorConfig"]:
+        """This config, or a more specific subclass for the settings."""
+        return cls
+
+    @classmethod
+    def default_model(cls, spec: "SpeculativeConfig") -> str | None:
+        """The `model` to use when it isn't given."""
+        return cls.base_method
+
+    @classmethod
+    def configure(cls, spec: "SpeculativeConfig") -> None:
+        """Set up the draft model config (and related options)."""
+        spec.prompt_lookup_max = 0
+        spec.prompt_lookup_min = 0
+        # No draft model; the target's config provides the vocab size etc.
+        spec.draft_model_config = spec.target_model_config
+        spec.draft_parallel_config = spec.target_parallel_config
+
+    @classmethod
+    def apply_options(cls, spec: "SpeculativeConfig") -> None:
+        """Apply method-specific options to the draft model config."""
+        return
+
+    # --- Behavior ---
+
+    def __init__(self, spec_config: "SpeculativeConfig"):
+        # Read lazily, since the spec config may still be updated.
+        self.spec_config = spec_config
+
+    @property
+    @final
+    def num_speculative_tokens(self) -> int:
+        return self.spec_config.num_speculative_tokens
+
+    @property
+    def uses_aux_hidden_states(self) -> bool:
+        """Whether the target must also output intermediate hidden states."""
+        return False
+
+    @property
+    def num_lookahead_tokens(self) -> int:
+        """See `VllmConfig.num_lookahead_tokens`."""
+        return 0
+
+    @property
+    def max_model_len(self) -> int | None:
+        """The longest sequence the speculator can draft for, or None if
+        unlimited. A smaller limit lowers overall max_model_len."""
+        return None
+
+    @property
+    def num_prefill_lookahead_tokens(self) -> int:
+        """See `VllmConfig.num_prefill_lookahead_tokens`."""
+        # Hidden-state drafters read only the immediate next token.
+        return 1 if self.uses_target_hidden_states else 0
+
+    @property
+    def max_num_new_slots_for_drafting(self) -> int:
+        """See `SpeculativeConfig.max_num_new_slots_for_drafting`."""
+        return 0
+
+    def max_num_draft_seqs(self, max_num_batched_tokens: int) -> int | None:
+        """The most requests the draft forward pass fits in
+        max_num_batched_tokens, which limits max_num_active_seqs, or None if
+        unlimited (e.g. when it's sized by the target's scheduled tokens)."""
+        return None
+
+    def hash_factors(self) -> list[Any]:
+        """Extra factors of `SpeculativeConfig.compute_hash`."""
+        return []
+
+    def pcp_unsupported_reason(self, parallel_config: "ParallelConfig") -> str | None:
+        """Why prefill context parallelism can't be used, or None."""
+        if self.supports_pcp:
+            return None
+        return (
+            "MRV2 PCP only supports DSpark or single-module MTP speculative decoding."
+        )
+
+
+class DraftModelSpeculatorConfig(SpeculatorConfig):
+    """A speculator whose draft model is loaded from the `model` checkpoint."""
+
+    # The draft model takes the target's positions, so its
+    # max_position_embeddings is raised to the target's max_model_len where its
+    # RoPE type allows.
+    shares_target_positions: ClassVar[bool] = True
+
+    @property
+    def max_model_len(self) -> int | None:
+        from vllm.config.speculative import _rope_only_sizes_cache
+
+        draft = self.spec_config.draft_model_config
+        if draft is None or (
+            # Its max_position_embeddings was raised to the target's.
+            self.shares_target_positions and _rope_only_sizes_cache(draft.hf_config)
+        ):
+            return None
+        return draft.max_model_len
+
+    @classmethod
+    def default_model(cls, spec: "SpeculativeConfig") -> str | None:
+        raise ValueError(
+            "num_speculative_tokens was provided but without speculative model."
+        )
+
+    @classmethod
+    def configure(cls, spec: "SpeculativeConfig") -> None:
+        from vllm.config import ModelConfig, SpeculativeConfig
+
+        spec.prompt_lookup_max = 0
+        spec.prompt_lookup_min = 0
+        if spec.model is None:
+            return
+
+        target = spec.target_model_config
+        # Callable target hf_overrides also apply to the draft; dict ones don't.
+        draft_hf_overrides = SpeculativeConfig.compose_draft_hf_overrides(
+            target.hf_overrides, cls.transformers_backend_hf_overrides(spec)
+        )
+        spec.draft_model_config = ModelConfig(
+            model=spec.model,
+            runner="draft",
+            # use_heterogeneous_vocab is V1 model runner only.
+            tokenizer=(
+                spec.model if spec.use_heterogeneous_vocab else target.tokenizer
+            ),
+            tokenizer_mode=target.tokenizer_mode,
+            trust_remote_code=target.trust_remote_code,
+            allowed_local_media_path=target.allowed_local_media_path,
+            allowed_media_domains=target.allowed_media_domains,
+            dtype=target.dtype,
+            seed=target.seed,
+            revision=spec.revision,
+            code_revision=spec.code_revision,
+            tokenizer_revision=target.tokenizer_revision,
+            max_model_len=spec.max_model_len,  # type: ignore[arg-type]
+            spec_target_max_model_len=target.max_model_len,
+            quantization=spec.quantization,
+            enforce_eager=target.enforce_eager,
+            max_logprobs=target.max_logprobs,
+            hf_overrides=draft_hf_overrides,
+            config_format=target.config_format,
+        )
+
+        config_cls = _resolve_checkpoint_config_class(spec, cls)
+        if spec.method is None:
+            spec.method = config_cls.base_method  # type: ignore[assignment]
+
+        if config_cls.shares_target_positions:
+            SpeculativeConfig._maybe_override_draft_max_position_embeddings(
+                spec.draft_model_config.hf_config, target.max_model_len
+            )
+        config_cls.adapt_draft_model_config(spec)
+        config_cls.resolve_num_speculative_tokens(spec)
+        if spec.num_speculative_tokens is None:
+            raise ValueError(
+                "A speculative model was provided, but "
+                "`num_speculative_tokens` was not provided"
+            )
+
+        spec.draft_tensor_parallel_size = SpeculativeConfig._verify_and_get_draft_tp(
+            spec.target_parallel_config, spec.draft_tensor_parallel_size
+        )
+        spec.draft_model_config.max_model_len = (
+            SpeculativeConfig._maybe_override_draft_max_model_len(
+                spec.max_model_len,
+                spec.draft_model_config.max_model_len,
+                target.max_model_len,
+            )
+        )
+        # From the final draft config, so EP agrees with its expert count.
+        spec.draft_parallel_config = SpeculativeConfig.create_draft_parallel_config(
+            spec.target_parallel_config,
+            spec.draft_tensor_parallel_size,
+            draft_model_config=spec.draft_model_config,
+        )
+
+    @classmethod
+    def transformers_backend_hf_overrides(
+        cls, spec: "SpeculativeConfig"
+    ) -> dict[str, Any] | None:
+        """HF config updates for a draft built by the Transformers modeling
+        backend, which replace the usual draft config rewrites, or None."""
+        return None
+
+    @classmethod
+    def adapt_draft_model_config(cls, spec: "SpeculativeConfig") -> None:
+        return
+
+    @classmethod
+    def resolve_num_speculative_tokens(cls, spec: "SpeculativeConfig") -> None:
+        """Default or validate `num_speculative_tokens` from the checkpoint."""
+        n_predict = getattr(spec.draft_model_config.hf_config, "n_predict", None)
+        if n_predict is None:
+            return
+        if spec.num_speculative_tokens is None:
+            # Default to max value defined in draft model config.
+            spec.num_speculative_tokens = n_predict
+        elif (
+            spec.num_speculative_tokens > n_predict
+            and spec.num_speculative_tokens % n_predict != 0
+        ):
+            # Ensure divisibility for MTP module reuse.
+            raise ValueError(
+                f"num_speculative_tokens:{spec.num_speculative_tokens}"
+                f" must be divisible by {n_predict=}"
+            )
+
+
+def _builtin(module: str, name: str) -> str:
+    return f"{_PKG}.{module}.{name}"
+
+
+_CONFIG_BY_METHOD: dict[str, str] = {
+    "extract_hidden_states": _builtin(
+        "extract_hidden_states_config", "ExtractHiddenStatesSpeculatorConfig"
+    ),
+    "draft_model": _builtin("standalone_ar.config", "StandaloneARSpeculatorConfig"),
+    "ngram": _builtin("ngram.config", "NgramSpeculatorConfig"),
+    "ngram_gpu": _builtin("ngram.config", "NgramGPUSpeculatorConfig"),
+    "eagle": _builtin("eagle.config", "EagleSpeculatorConfig"),
+    "eagle3": _builtin("eagle.config", "Eagle3SpeculatorConfig"),
+    "mtp": _builtin("mtp.config", "MTPSpeculatorConfig"),
+    "dflash": _builtin("dflash.config", "DFlashSpeculatorConfig"),
+    "dspark": _builtin("dspark.config", "DSparkSpeculatorConfig"),
+}
+
+
+# V1 model runner only (see `v1_speculator_configs`).
+_V1_ONLY_METHODS = ("suffix", "custom_class")
+
+# Custom speculative methods' configs (or their qualified names), by name.
+_CUSTOM_CONFIGS: dict[str, "type[SpeculatorConfig] | str"] = {}
+
+
+def register_speculative_method(
+    name: str, config_cls: "type[SpeculatorConfig] | str"
+) -> None:
+    """Register a custom speculative method, used with `method=name`.
+
+    `config_cls` is its `SpeculatorConfig` subclass, or its qualified name to
+    import it lazily. Register it from a vLLM general plugin, so that it's
+    registered in every process. Custom methods require the V2 model runner.
+    This is an experimental interface that may change between releases.
+    """
+    from vllm.config.speculative import SpeculativeMethod
+
+    if name in get_args(SpeculativeMethod) or name in _CONFIG_BY_METHOD:
+        raise ValueError(f"'{name}' is a built-in speculative method.")
+    if name in _CUSTOM_CONFIGS and _CUSTOM_CONFIGS[name] != config_cls:
+        raise ValueError(f"Speculative method '{name}' is already registered.")
+    _CUSTOM_CONFIGS[name] = config_cls
+
+
+def registered_speculative_methods() -> list[str]:
+    """The names of the registered custom speculative methods."""
+    return sorted(_CUSTOM_CONFIGS)
+
+
+def get_custom_speculator_config_class(
+    method: str | None,
+) -> type[SpeculatorConfig] | None:
+    """The config of a registered custom speculative method, or None."""
+    if method not in _CUSTOM_CONFIGS:
+        return None
+    config_cls = _CUSTOM_CONFIGS[method]
+    if isinstance(config_cls, str):
+        config_cls = resolve_obj_by_qualname(config_cls)
+    if not (isinstance(config_cls, type) and issubclass(config_cls, SpeculatorConfig)):
+        raise TypeError(
+            f"The config of speculative method '{method}' must be a subclass of "
+            f"{SpeculatorConfig.__module__}.{SpeculatorConfig.__qualname__}."
+        )
+    if config_cls.speculator_class is None:
+        raise ValueError(
+            f"The config of speculative method '{method}' must set speculator_class."
+        )
+    return config_cls
+
+
+def get_method_config_class(spec_config: "SpeculativeConfig") -> type[SpeculatorConfig]:
+    """The config class for `spec_config.method`, regardless of the draft
+    checkpoint (for setting up `spec_config`)."""
+    if spec_config.custom_method is not None:
+        custom_cls = get_custom_speculator_config_class(spec_config.custom_method)
+        assert custom_cls is not None
+        return custom_cls
+    method = spec_config.method
+    if method is None:
+        # Not configured: a draft model, unless the checkpoint shows otherwise.
+        method = "draft_model"
+    if method in _V1_ONLY_METHODS:
+        from vllm.v1.worker.gpu.spec_decode.v1_speculator_configs import (
+            SuffixSpeculatorConfig,
+            V1CustomProposerConfig,
+        )
+
+        if method == "suffix":
+            return SuffixSpeculatorConfig
+        return V1CustomProposerConfig
+    if method not in _CONFIG_BY_METHOD:
+        raise ValueError(f"Unsupported speculative method: '{method}'")
+    return resolve_obj_by_qualname(_CONFIG_BY_METHOD[method])
+
+
+def get_speculator_config_class(
+    spec_config: "SpeculativeConfig",
+) -> type[SpeculatorConfig]:
+    """The config class for a set-up spec config."""
+    config_cls = get_method_config_class(spec_config)
+    if spec_config.custom_method is not None:
+        return config_cls
+    declared_cls = _get_draft_declared_config_class(spec_config)
+    if declared_cls is not None and issubclass(declared_cls, config_cls):
+        return declared_cls
+    return config_cls.specialize(spec_config)
+
+
+def _get_draft_declared_config_class(
+    spec_config: "SpeculativeConfig",
+) -> type[SpeculatorConfig] | None:
+    """The config declared by the draft model (see `SpeculativeDraftModel`)."""
+    draft = spec_config.draft_model_config
+    qualname = getattr(draft, "speculator_config_class", None)
+    return resolve_obj_by_qualname(qualname) if qualname is not None else None
+
+
+# Methods detected by `matches` (in order) if the draft model declares none.
+_CHECKPOINT_DETECTION_ORDER = ("eagle", "eagle3", "dflash", "dspark", "mtp")
+
+
+def _resolve_checkpoint_config_class(
+    spec_config: "SpeculativeConfig",
+    configured_cls: type[DraftModelSpeculatorConfig],
+) -> type[DraftModelSpeculatorConfig]:
+    """The built-in config for the draft checkpoint: detected if no method is
+    configured, else the configured one, checked against the speculator the
+    draft model declares (if any).
+    """
+    declared_cls = _get_draft_declared_config_class(spec_config)
+    if spec_config.method is not None:
+        if declared_cls is not None and declared_cls.base_method != spec_config.method:
+            draft = spec_config.draft_model_config
+            declared = declared_cls.base_method
+            requested = spec_config.custom_method or spec_config.method
+            raise ValueError(
+                f"speculative_config requested method={requested!r}, but "
+                f"the draft model {draft.model!r} is a {declared!r} checkpoint "
+                f"(model_type={getattr(draft.hf_config, 'model_type', None)!r}, "
+                f"architectures={getattr(draft, 'architectures', None)}). Pass a "
+                f"matching method or a matching draft checkpoint. Use "
+                f"method={declared!r} if that is what you intended."
+            )
+        return configured_cls
+    if spec_config.custom_method is not None:
+        # Not based on a built-in speculator, so there's nothing to check.
+        return configured_cls
+
+    if declared_cls is not None:
+        assert issubclass(declared_cls, DraftModelSpeculatorConfig)
+        return declared_cls
+    for method in _CHECKPOINT_DETECTION_ORDER:
+        config_cls = resolve_obj_by_qualname(_CONFIG_BY_METHOD[method])
+        if config_cls.matches(spec_config):
+            return config_cls
+    return configured_cls
diff --git a/vllm/v1/worker/gpu/spec_decode/standalone_ar/config.py b/vllm/v1/worker/gpu/spec_decode/standalone_ar/config.py
new file mode 100644
index 000000000000..2a71e26c8d9a
--- /dev/null
+++ b/vllm/v1/worker/gpu/spec_decode/standalone_ar/config.py
@@ -0,0 +1,30 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from vllm.v1.worker.gpu.spec_decode.speculator_config import (
+    DraftModelSpeculatorConfig,
+)
+
+
+class StandaloneARSpeculatorConfig(DraftModelSpeculatorConfig):
+    """A speculator that drafts autoregressively with its own draft model."""
+
+    base_method = "draft_model"
+    speculator_class = (
+        "vllm.v1.worker.gpu.spec_decode.standalone_ar.speculator.StandaloneARSpeculator"
+    )
+    draft_has_own_vocab = True
+    # An independent model with its own positional limit.
+    shares_target_positions = False
+
+    @property
+    def num_lookahead_tokens(self) -> int:
+        return self.num_speculative_tokens
+
+    @property
+    def max_num_new_slots_for_drafting(self) -> int:
+        if self.spec_config.parallel_drafting:
+            # PARD does not shift the existing input, so all K query positions
+            # require additional slots.
+            return self.num_speculative_tokens
+        # The autoregressive draft-model input retains one unsliced token.
+        return 1
diff --git a/vllm/v1/worker/gpu/spec_decode/standalone_ar/speculator.py b/vllm/v1/worker/gpu/spec_decode/standalone_ar/speculator.py
index 1df19a3f9693..f4ef1ed96c36 100644
--- a/vllm/v1/worker/gpu/spec_decode/standalone_ar/speculator.py
+++ b/vllm/v1/worker/gpu/spec_decode/standalone_ar/speculator.py
@@ -11,6 +11,7 @@
 from vllm.forward_context import BatchDescriptor, set_forward_context
 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.triton_utils import tl, triton
 from vllm.v1.kv_cache_interface import KVCacheConfig
 from vllm.v1.spec_decode.utils import next_power_of_2
@@ -21,6 +22,7 @@
 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.spec_decode.speculator import DraftModelSpeculator
+from vllm.v1.worker.gpu.spec_decode.utils import get_pp_safe_draft_load_config
 from vllm.v1.worker.utils import AttentionGroup
 
 logger = init_logger(__name__)
@@ -105,15 +107,10 @@ def load_draft_model(
     ) -> nn.Module:
         spec = self.speculative_config
         assert spec is not None
-        cache_config = self.vllm_config.cache_config
-        if spec.kv_cache_dtype is not None:
-            cache_config = replace(cache_config, cache_dtype=spec.kv_cache_dtype)
-
         draft_vllm_config = replace(
-            self.vllm_config,
+            spec.apply_draft_overrides(self.vllm_config),
             model_config=self.draft_model_config,
             quant_config=None,
-            cache_config=cache_config,
             parallel_config=replace(
                 spec.draft_parallel_config,
                 rank=self.vllm_config.parallel_config.rank,
@@ -127,6 +124,9 @@ def load_draft_model(
             return get_model(
                 vllm_config=draft_vllm_config,
                 prefix="draft_model",
+                load_config=get_pp_safe_draft_load_config(
+                    get_draft_load_config(draft_vllm_config)
+                ),
             )
 
     @torch.inference_mode()
diff --git a/vllm/v1/worker/gpu/spec_decode/target_dependent_ar/speculator.py b/vllm/v1/worker/gpu/spec_decode/target_dependent_ar/speculator.py
index 7b70fd0ee0bd..f17a820c4f1f 100644
--- a/vllm/v1/worker/gpu/spec_decode/target_dependent_ar/speculator.py
+++ b/vllm/v1/worker/gpu/spec_decode/target_dependent_ar/speculator.py
@@ -93,6 +93,8 @@ def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None:
             self.device,
             cudagraph_mode,
             self.num_speculative_steps + 1,
+            # Draft prefill shares the target's padded attention metadata.
+            varlen_decode=self.enable_adaptive_verification,
         )
 
         # PIECEWISE cudagraphs are not supported for draft decodes.
@@ -131,7 +133,9 @@ def capture(self) -> None:
             self.model_state,
             self.target_input_buffers,
             self.block_tables,
-            self.target_attn_groups,
+            self.target_attn_groups
+            if self.can_reuse_target_attn_metadata
+            else self.attn_groups,
             self.kv_cache_config,
             progress_bar_desc="Capturing prefill CUDA graphs",
         )
@@ -266,8 +270,10 @@ def propose(
             uniform_token_count,
             dp_size=self.dp_size,
             dp_rank=self.dp_rank,
+            max_query_len=None if input_batch.has_prefill else int(max_query_len),
             need_eager=is_profile,
             dp_sync_state=dp_sync_state,
+            parallel_config=self.vllm_config.parallel_config,
         )
         num_tokens_across_dp = (
             prefill_batch_sync.num_tokens_across_dp
@@ -277,6 +283,16 @@ def propose(
 
         self._prepare_eplb_forward(num_tokens)
 
+        if not skip_attn_for_dummy_run:
+            attn_metadata, slot_mappings = self._build_drafter_attn_metadata(
+                input_batch,
+                attn_metadata,
+                slot_mappings,
+                prefill_batch_desc.cg_mode,
+                for_capture=dummy_run
+                and prefill_batch_desc.cg_mode == CUDAGraphMode.FULL,
+            )
+
         self.on_prefill_begin(num_reqs)
         if prefill_batch_desc.cg_mode == CUDAGraphMode.FULL:
             # Replay the full graph for draft prefill.
@@ -344,6 +360,7 @@ def propose(
             dp_rank=self.dp_rank,
             need_eager=is_profile,
             dp_sync_state=decode_batch_sync,
+            parallel_config=self.vllm_config.parallel_config,
         )
         if self.decode_cudagraph_manager is not None:
             decode_batch_desc = self.decode_cudagraph_manager.specialize_spec_tokens(
diff --git a/vllm/v1/worker/gpu/spec_decode/utils.py b/vllm/v1/worker/gpu/spec_decode/utils.py
index b005f5c3f3ae..975ba32e2263 100644
--- a/vllm/v1/worker/gpu/spec_decode/utils.py
+++ b/vllm/v1/worker/gpu/spec_decode/utils.py
@@ -1,11 +1,14 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from collections.abc import Iterable
+
 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.core.sched.output import SchedulerOutput
 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
@@ -51,20 +54,34 @@ def __init__(self, device: torch.device | None = None):
         self.copy_stream = torch.cuda.Stream(device)
         # Blocking (sleep) event to avoid busy-polling the CUDA driver lock.
         self.copy_event = torch.cuda.Event(blocking=True)
-
-        self.req_ids: list[str] = []
-        self.draft_tokens_np: np.ndarray | None = None
-        self.num_draft_tokens: int = 0
+        self._copy_pending = False
+
+        # Latest drafts per live request: a row view into an async-copied
+        # pinned array (synced in get_draft_tokens), or just the draft width
+        # when the scheduler doesn't need the token ids.
+        self._draft_tokens_by_req: dict[str, np.ndarray | int] = {}
+        # Requests whose drafts the scheduler will fetch for the latest step's
+        # deferred grammar bitmask; None means all requests.
+        self._deferred_req_ids: Iterable[str] | None = None
+
+    def begin_step(self, scheduler_output: SchedulerOutput) -> None:
+        # get_draft_tokens is only called for the most recently executed
+        # scheduler output.
+        self._deferred_req_ids = (
+            scheduler_output.scheduled_spec_decode_tokens.keys()
+            if scheduler_output.pending_structured_output_tokens
+            else None
+        )
 
     def set_draft_tokens(
         self, input_batch: InputBatch, draft_tokens: torch.Tensor
     ) -> None:
-        self.req_ids = input_batch.req_ids
-        self.num_draft_tokens = draft_tokens.shape[1]
         if not input_batch.has_structured_output_reqs:
             # No draft token validation needs to be performed by
-            # the scheduler for this batch.
-            self.draft_tokens_np = None
+            # the scheduler for this batch. Record draft token widths only.
+            num_draft_tokens = draft_tokens.shape[1]
+            for req_id in input_batch.req_ids:
+                self._draft_tokens_by_req[req_id] = num_draft_tokens
             return
 
         # For spec decoding + structured outputs, we must transfer the
@@ -72,21 +89,46 @@ def set_draft_tokens(
         current_stream = torch.cuda.current_stream(self.device)
         self.copy_stream.wait_stream(current_stream)
         with torch.cuda.stream(self.copy_stream):
-            self.draft_tokens_np = async_copy_to_np(draft_tokens)
+            draft_tokens_np = async_copy_to_np(draft_tokens)
             # draft_tokens is a temporary allocation on the main stream and read here on
             # copy_stream; without record_stream, the caching allocator may reuse its
             # memory before the async copy executes.
             draft_tokens.record_stream(self.copy_stream)
             self.copy_event.record()
+            self._copy_pending = True
+
+        for req_id, row in zip(input_batch.req_ids, draft_tokens_np):
+            self._draft_tokens_by_req[req_id] = row
+
+    def discard(self, req_id: str) -> None:
+        self._draft_tokens_by_req.pop(req_id, None)
 
     def get_draft_tokens(self) -> DraftTokenIds | None:
-        if self.draft_tokens_np is not None:
+        """Return the latest drafts needed for the current step's deferred
+        grammar bitmask, or for all live requests if sampling isn't deferred.
+        Requests without retained drafts are omitted."""
+        if self._copy_pending:
             self.copy_event.synchronize()
-            draft_token_ids = self.draft_tokens_np.tolist()
+            self._copy_pending = False
+
+        drafts_by_req = self._draft_tokens_by_req
+        req_ids = self._deferred_req_ids
+        if req_ids is None:
+            out_req_ids = list(drafts_by_req)
         else:
-            # This case only happens when async scheduling is disabled.
-            draft_token_ids = [[-1] * self.num_draft_tokens for _ in self.req_ids]
-        return DraftTokenIds(self.req_ids, draft_token_ids)
+            out_req_ids = [req_id for req_id in req_ids if req_id in drafts_by_req]
+        draft_token_ids: list[list[int]] = []
+        for req_id in out_req_ids:
+            draft_tokens = drafts_by_req[req_id]
+            if isinstance(draft_tokens, int):
+                # Construct placeholders from draft token width
+                draft_tokens = [-1] * draft_tokens
+            else:
+                # Convert numpy view to list
+                draft_tokens = draft_tokens.tolist()
+            draft_token_ids.append(draft_tokens)
+
+        return DraftTokenIds(out_req_ids, draft_token_ids)
 
 
 def get_parallel_drafting_token_id(hf_config) -> int:
diff --git a/vllm/v1/worker/gpu/spec_decode/v1_speculator_configs.py b/vllm/v1/worker/gpu/spec_decode/v1_speculator_configs.py
new file mode 100644
index 000000000000..5e9c08ad6017
--- /dev/null
+++ b/vllm/v1/worker/gpu/spec_decode/v1_speculator_configs.py
@@ -0,0 +1,51 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Speculator configs that only the V1 model runner needs.
+
+These cover speculative methods that the V2 model runner doesn't implement.
+They will be removed together with the V1 model runner.
+"""
+
+from typing import TYPE_CHECKING
+
+from vllm.logger import init_logger
+from vllm.v1.worker.gpu.spec_decode.speculator_config import SpeculatorConfig
+
+if TYPE_CHECKING:
+    from vllm.config import SpeculativeConfig
+
+logger = init_logger(__name__)
+
+
+class UnsupportedSpeculatorConfig(SpeculatorConfig):
+    """A speculative method that the V2 model runner doesn't implement."""
+
+
+class SuffixSpeculatorConfig(UnsupportedSpeculatorConfig):
+    base_method = "suffix"
+
+    @classmethod
+    def configure(cls, spec: "SpeculativeConfig") -> None:
+        spec._validate_suffix_decoding()
+
+
+class V1CustomProposerConfig(UnsupportedSpeculatorConfig):
+    """A V1 custom proposer, loaded from the module path in `model`."""
+
+    base_method = "custom_class"
+
+    @classmethod
+    def default_model(cls, spec: "SpeculativeConfig") -> str | None:
+        raise ValueError(
+            "method='custom_class' requires 'model' to contain the custom "
+            "proposer module path (e.g., 'my_module.MyProposer')."
+        )
+
+    @classmethod
+    def configure(cls, spec: "SpeculativeConfig") -> None:
+        logger.warning_once(
+            "Using a custom class-based proposer backend. This is an "
+            "experimental feature and the proposer interface is subject to "
+            "breaking changes in future vLLM releases."
+        )
+        super().configure(spec)
diff --git a/vllm/v1/worker/gpu/spec_decode/xpress/config.py b/vllm/v1/worker/gpu/spec_decode/xpress/config.py
new file mode 100644
index 000000000000..21c242026f13
--- /dev/null
+++ b/vllm/v1/worker/gpu/spec_decode/xpress/config.py
@@ -0,0 +1,29 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from vllm.config import SpeculativeConfig
+from vllm.v1.worker.gpu.spec_decode.dflash.config import DFlashSpeculatorConfig
+
+
+class XPressSpeculatorConfig(DFlashSpeculatorConfig):
+    speculator_class = (
+        "vllm.v1.worker.gpu.spec_decode.xpress.speculator.XPressSpeculator"
+    )
+    supported_options = frozenset({"xpress_topc"})
+
+    @classmethod
+    def apply_options(cls, spec: SpeculativeConfig) -> None:
+        hf_config = spec.draft_model_config.hf_config
+        xpress_topc = spec.xpress_topc
+        if xpress_topc is None:
+            xpress_topc = getattr(hf_config, "xpress_topc", None)
+        if xpress_topc is None:
+            return
+        draft_vocab_size = (
+            getattr(hf_config, "draft_vocab_size", None) or hf_config.vocab_size
+        )
+        if not 0 <= xpress_topc <= draft_vocab_size:
+            raise ValueError(
+                "xpress_topc must be between 0 and the draft "
+                f"vocabulary size ({draft_vocab_size})"
+            )
+        hf_config.xpress_topc = xpress_topc
diff --git a/vllm/v1/worker/gpu/structured_outputs.py b/vllm/v1/worker/gpu/structured_outputs.py
index 46a276e7a4c1..dc5d23d85e5f 100644
--- a/vllm/v1/worker/gpu/structured_outputs.py
+++ b/vllm/v1/worker/gpu/structured_outputs.py
@@ -12,43 +12,6 @@
 GRAMMAR_BITMASK_BLOCK_SIZE = 8192
 
 
-def grammar_invalid_drafts(
-    input_batch: InputBatch,
-    grammar_req_ids: list[str],
-    num_acceptable_drafts: list[int] | None,
-) -> torch.Tensor | None:
-    """Mask over logit rows whose draft was verified against a permissive row.
-
-    Draft i sits at local position i + 1, and drafts from `num_acceptable_drafts`
-    on met `_full_mask`. Local positions come from the device `cu_num_logits`, so
-    this stays right when adaptive verification trims drafts on device.
-    Returns None when there is nothing to invalidate.
-    """
-    if not grammar_req_ids or input_batch.num_draft_tokens == 0:
-        return None
-    req_id_to_idx = {req_id: i for i, req_id in enumerate(input_batch.req_ids)}
-    limit = np.full(input_batch.num_reqs, np.iinfo(np.int32).max, dtype=np.int32)
-    for i, req_id in enumerate(grammar_req_ids):
-        req_idx = req_id_to_idx.get(req_id)
-        if req_idx is not None:
-            # None (an older scheduler, or warmup): invalidate the whole window.
-            limit[req_idx] = (
-                num_acceptable_drafts[i] if num_acceptable_drafts is not None else 0
-            )
-    # Adaptive verification only trims drafts, so this bound holds either way.
-    num_drafts = input_batch.num_draft_tokens_per_req
-    if num_drafts is not None and (limit >= num_drafts).all():
-        return None
-    local_pos = input_batch.expanded_local_pos
-    cu_num_logits = input_batch.cu_num_logits
-    row_limit = torch.repeat_interleave(
-        async_tensor_h2d(limit, device=local_pos.device),
-        cu_num_logits[1:] - cu_num_logits[:-1],
-        output_size=local_pos.shape[0],
-    )
-    return local_pos > row_limit
-
-
 class StructuredOutputsWorker:
     def __init__(
         self,
diff --git a/vllm/v1/worker/gpu/ubatch_utils.py b/vllm/v1/worker/gpu/ubatch_utils.py
index 608231a3ed42..ed01e5e243a8 100644
--- a/vllm/v1/worker/gpu/ubatch_utils.py
+++ b/vllm/v1/worker/gpu/ubatch_utils.py
@@ -6,7 +6,7 @@
 from collections.abc import Callable
 from contextlib import ExitStack
 from dataclasses import replace
-from typing import Any, NamedTuple
+from typing import TYPE_CHECKING, Any, NamedTuple
 
 import numpy as np
 import torch
@@ -20,6 +20,7 @@
     override_forward_context,
 )
 from vllm.logger import init_logger
+from vllm.platforms import current_platform
 from vllm.sequence import IntermediateTensors
 from vllm.utils.torch_utils import current_stream
 from vllm.v1.kv_cache_interface import KVCacheConfig
@@ -36,6 +37,10 @@
 from vllm.v1.worker.ubatching import make_ubatch_contexts
 from vllm.v1.worker.utils import AttentionGroup
 
+if TYPE_CHECKING:
+    from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor
+    from vllm.v1.worker.gpu.dp_utils import DPSyncState
+
 logger = init_logger(__name__)
 
 
@@ -46,7 +51,9 @@ class UBatchState(NamedTuple):
     forward_contexts: list[ForwardContext]
 
 
-def create_ubatch_slices(input_batch: InputBatch, num_ubatches: int) -> UBatchSlices:
+def create_ubatch_slices(
+    input_batch: InputBatch, num_ubatches: int, uniform_token_count: int | None = None
+) -> UBatchSlices:
     """Split a DP-padded batch into the slices its microbatches run on.
 
     Splitting at the midpoint of the DP-padded token count leaves the trailing
@@ -54,13 +61,24 @@ def create_ubatch_slices(input_batch: InputBatch, num_ubatches: int) -> UBatchSl
     request slice onto the last request leaves it holding that request with
     zero query tokens -- the shape a straddling request already takes -- so it
     stays well-formed with no work to do, like a dummy run.
+
+    A uniform batch of multi-token queries (spec decode verification) is split
+    on a request boundary instead, so each microbatch stays a uniform decode.
+    `uniform_token_count` must be agreed across DP ranks.
     """
+    split_point = None
+    if uniform_token_count is not None and uniform_token_count > 1:
+        num_tokens = input_batch.num_tokens_after_padding
+        reqs_per_ubatch = num_tokens // (uniform_token_count * num_ubatches)
+        if reqs_per_ubatch > 0:
+            split_point = reqs_per_ubatch * uniform_token_count
     _, ubatch_slices = maybe_create_ubatch_slices(
         True,
         input_batch.num_scheduled_tokens,
         input_batch.num_tokens_after_padding,
         input_batch.num_reqs_after_padding,
         num_ubatches,
+        split_point=split_point,
     )
     assert ubatch_slices is not None
     return [
@@ -263,8 +281,9 @@ def slice_model_inputs(
 def merge_ubatch_outputs(outputs: list[Any]) -> Any:
     """Reassemble the full-batch model output from the per-microbatch ones.
 
-    Models return hidden states, a tuple of those plus auxiliary hidden states
-    (EAGLE3), or IntermediateTensors on non-final pipeline ranks.
+    Models return hidden states, a tuple of those plus a list of auxiliary
+    hidden states (EAGLE3, DFlash), or IntermediateTensors on non-final
+    pipeline ranks.
     """
     first = outputs[0]
     if isinstance(first, IntermediateTensors):
@@ -275,7 +294,9 @@ def merge_ubatch_outputs(outputs: list[Any]) -> Any:
             }
         )
     if isinstance(first, tuple):
-        return tuple(torch.cat(parts, dim=0) for parts in zip(*outputs))
+        return tuple(merge_ubatch_outputs(list(parts)) for parts in zip(*outputs))
+    if isinstance(first, list):
+        return [merge_ubatch_outputs(list(parts)) for parts in zip(*outputs)]
     return torch.cat(outputs, dim=0)
 
 
@@ -337,7 +358,8 @@ def prepare(
         input_batch: InputBatch,
         block_tables: tuple[torch.Tensor, ...],
         slot_mappings: torch.Tensor,
-        cg_mode: CUDAGraphMode = CUDAGraphMode.NONE,
+        batch_desc: "BatchExecutionDescriptor",
+        dp_sync_state: "DPSyncState | None" = None,
         for_capture: bool = False,
     ) -> UBatchState:
         """Split the batch into the microbatches the step will run on.
@@ -346,9 +368,19 @@ def prepare(
         contexts the threads install, not in the caller's.
 
         FULL uses padded sizes; `for_capture` refreshes metadata from dummy
-        block tables.
+        block tables. Eager steps split by the DP-agreed uniform token count in
+        `dp_sync_state` (see `create_ubatch_slices`); FULL replays the split
+        its graph captured with.
         """
-        ubatch_slices = create_ubatch_slices(input_batch, self.num_ubatches)
+        cg_mode = batch_desc.cg_mode
+        uniform_token_count = (
+            batch_desc.uniform_token_count
+            if cg_mode == CUDAGraphMode.FULL or dp_sync_state is None
+            else dp_sync_state.uniform_token_count
+        )
+        ubatch_slices = create_ubatch_slices(
+            input_batch, self.num_ubatches, uniform_token_count
+        )
 
         attn_metadata = []
         slot_mappings_by_layer = []
@@ -467,6 +499,7 @@ def begin_capturable_run(
 
         outputs: dict[int, Any] = {}
         errors: dict[int, BaseException] = {}
+        dtype = self.vllm_config.model_config.dtype
 
         @torch.inference_mode()
         def run_ubatch(ubatch_context, inputs: dict[str, Any]) -> None:
@@ -478,6 +511,12 @@ def run_ubatch(ubatch_context, inputs: dict[str, Any]) -> None:
                     for stream in (compute_stream, self.comm_stream):
                         with torch.cuda.stream(stream):
                             torch.cuda.current_blas_handle()
+                            if current_platform.is_rocm():
+                                # hipBLASLt handles are per thread and stream and
+                                # only a GEMM creates one; creating it in capture
+                                # fails.
+                                x = torch.zeros(1, 1, dtype=dtype, device=self.device)
+                                torch.mm(x, x)
                 with ubatch_context:
                     outputs[ubatch_context.id] = model(**inputs)
             except BaseException as e:  # noqa: BLE001
diff --git a/vllm/v1/worker/gpu/warmup.py b/vllm/v1/worker/gpu/warmup.py
index 13de688b5d5d..6c279d209551 100644
--- a/vllm/v1/worker/gpu/warmup.py
+++ b/vllm/v1/worker/gpu/warmup.py
@@ -26,6 +26,10 @@
     UniformTypeKVCacheSpecs,
 )
 from vllm.v1.request import Request
+from vllm.v1.worker.extensible_kv_cache import (
+    ensure_kv_cache_blocks,
+    num_committable_kv_blocks,
+)
 from vllm.v1.worker.gpu.model_runner import GPUModelRunner
 
 logger = init_logger(__name__)
@@ -38,12 +42,14 @@ def _reserved_block_count(
     num_lookahead_tokens: int,
     max_model_len: int,
     max_encoder_len: int,
+    dcp_world_size: int = 1,
 ) -> int:
     """Number of blocks the scheduler would hold for a request of `num_tokens`.
 
     Warmup hand-builds its `SchedulerOutput`s, so it must reserve what
     `KVCacheManager.allocate_slots` reserves: the token range plus
     `num_lookahead_tokens`, where the speculator writes the KV of its drafts.
+    A DCP-sharded block holds `dcp_world_size` times its tokens.
     """
     if isinstance(kvcache_spec, UniformTypeKVCacheSpecs):
         kvcache_spec = kvcache_spec.first_spec
@@ -60,17 +66,29 @@ def _reserved_block_count(
         num_speculative_blocks = kvcache_spec.num_speculative_blocks
         if kvcache_spec.mamba_cache_mode == "align":
             return cdiv(num_tokens, kvcache_spec.block_size) + num_speculative_blocks
+    block_size = kvcache_spec.block_size
+    if kvcache_spec.dcp_sharded:
+        block_size *= dcp_world_size
     num_tokens = min(num_tokens + num_lookahead_tokens, max_model_len)
-    return cdiv(num_tokens, kvcache_spec.block_size) + num_speculative_blocks
+    return cdiv(num_tokens, block_size) + num_speculative_blocks
+
+
+def _warmup_token_ids(model_runner: GPUModelRunner, num_tokens: int) -> list[int]:
+    """Distinct-looking prompt token ids that stay within the vocabulary, which
+    can be smaller than a full-budget or full-length prompt."""
+    vocab_size = model_runner.model_config.get_vocab_size()
+    return [i % vocab_size for i in range(num_tokens)]
 
 
 def _warmup_block_counter(
     model_runner: GPUModelRunner,
 ) -> Callable[[int, KVCacheSpec], int]:
     """Bind `_reserved_block_count` to `model_runner`'s reservation policy."""
-    num_lookahead_tokens = model_runner.vllm_config.num_lookahead_tokens
+    vllm_config = model_runner.vllm_config
+    num_lookahead_tokens = vllm_config.num_lookahead_tokens
     max_model_len = model_runner.max_model_len
     max_encoder_len = getattr(model_runner.model_state, "max_encoder_len", 0)
+    dcp_world_size = vllm_config.parallel_config.decode_context_parallel_size
 
     def block_count(num_tokens: int, kvcache_spec: KVCacheSpec) -> int:
         return _reserved_block_count(
@@ -79,6 +97,7 @@ def block_count(num_tokens: int, kvcache_spec: KVCacheSpec) -> int:
             num_lookahead_tokens=num_lookahead_tokens,
             max_model_len=max_model_len,
             max_encoder_len=max_encoder_len,
+            dcp_world_size=dcp_world_size,
         )
 
     return block_count
@@ -92,8 +111,10 @@ def run_mixed_prefill_decode_warmup(
     *,
     mixed_step_context: AbstractContextManager[object] | None = None,
     req_id_prefix: str = "_v2_mixed_warmup",
+    prompt_logprobs: int | None = None,
 ) -> bool:
-    """Run a V2 mixed prefill+decode step through normal scheduler inputs."""
+    """Run a V2 mixed prefill+decode step through normal scheduler inputs,
+    requesting prompt_logprobs on the prefill request."""
     if model_runner.is_pooling_model or model_runner.max_num_reqs < 2 or num_tokens < 3:
         return False
 
@@ -102,8 +123,8 @@ def run_mixed_prefill_decode_warmup(
     decode_prompt_len = 2
     decode_scheduled_tokens = 1
     prefill_len = num_tokens - decode_scheduled_tokens
-    decode_token_ids = list(range(decode_prompt_len))
-    prefill_token_ids = list(range(prefill_len))
+    decode_token_ids = _warmup_token_ids(model_runner, decode_prompt_len)
+    prefill_token_ids = _warmup_token_ids(model_runner, prefill_len)
 
     kv_cache_groups = model_runner.kv_cache_config.kv_cache_groups
     block_count = _warmup_block_counter(model_runner)
@@ -121,14 +142,16 @@ def run_mixed_prefill_decode_warmup(
     ]
     prefill_block_counts = [block_count(prefill_len, s) for s in kv_cache_specs]
     required_blocks = sum(decode_block_counts) + sum(prefill_block_counts)
-    if model_runner.kv_cache_config.num_blocks <= required_blocks:
+    num_blocks = num_committable_kv_blocks(model_runner)
+    if num_blocks <= required_blocks:
         logger.warning(
             "Skipping V2 mixed prefill+decode warmup because only %d KV blocks "
             "are available for %d required warmup blocks.",
-            model_runner.kv_cache_config.num_blocks,
+            num_blocks,
             required_blocks,
         )
         return False
+    ensure_kv_cache_blocks(model_runner, 1 + required_blocks)
 
     next_block_id = 1
 
@@ -183,7 +206,12 @@ def _alloc_blocks(num_blocks: int) -> list[int]:
             req_id=prefill_req_id,
             prompt_token_ids=prefill_token_ids,
             mm_features=[],
-            sampling_params=sampling_params,
+            sampling_params=SamplingParams(
+                max_tokens=2,
+                temperature=0.0,
+                watermarking=False,
+                prompt_logprobs=prompt_logprobs,
+            ),
             pooling_params=None,
             block_ids=tuple(_alloc_blocks(n) for n in prefill_block_counts),
             num_computed_tokens=0,
@@ -214,6 +242,121 @@ def _alloc_blocks(num_blocks: int) -> list[int]:
     return True
 
 
+def run_context_prefill_warmup(
+    model_runner: GPUModelRunner,
+    worker_execute_model: Callable[[SchedulerOutput], Any],
+    worker_sample_tokens: Callable[[GrammarOutput | None], Any],
+    num_context_tokens: int,
+) -> bool:
+    """Run a prefill step over about `num_context_tokens` of cached context.
+
+    MLA up-projects a prefix cache hit's context chunk by chunk, a transient
+    that no step starting from an empty prefix reaches. The queries are kept
+    short: the full-budget step already measured theirs.
+    """
+    query_len = model_runner.decode_query_len + 1
+    # Leave room for a sampled token after the prompt; and no request needs
+    # more context than the step does.
+    context_len = min(model_runner.max_model_len - query_len - 1, num_context_tokens)
+    if model_runner.is_pooling_model or context_len <= 0:
+        return False
+
+    kv_cache_specs = [
+        g.kv_cache_spec for g in model_runner.kv_cache_config.kv_cache_groups
+    ]
+    block_count = _warmup_block_counter(model_runner)
+    block_counts = [block_count(context_len + query_len, s) for s in kv_cache_specs]
+    blocks_per_req = max(sum(block_counts), 1)
+    num_reqs = min(
+        cdiv(num_context_tokens, context_len),
+        model_runner.max_num_reqs,
+        model_runner.scheduler_config.max_num_batched_tokens // query_len,
+        (num_committable_kv_blocks(model_runner) - 1) // blocks_per_req,
+    )
+    if num_reqs < 1:
+        return False
+    ensure_kv_cache_blocks(model_runner, 1 + num_reqs * blocks_per_req)
+
+    next_block_id = 1
+
+    def _alloc_blocks(num_blocks: int) -> list[int]:
+        nonlocal next_block_id
+        return list(range(next_block_id, next_block_id := next_block_id + num_blocks))
+
+    req_ids = [f"_context_warmup_{i}_" for i in range(num_reqs)]
+    token_ids = _warmup_token_ids(model_runner, context_len + query_len)
+    sampling_params = SamplingParams(max_tokens=1, temperature=0.0)
+    prefill_output = SchedulerOutput.make_empty()
+    prefill_output.num_spec_tokens_to_schedule = model_runner.num_speculative_steps
+    prefill_output.scheduled_new_reqs = [
+        NewRequestData(
+            req_id=req_id,
+            prompt_token_ids=token_ids,
+            mm_features=[],
+            sampling_params=sampling_params,
+            pooling_params=None,
+            block_ids=tuple(_alloc_blocks(n) for n in block_counts),
+            num_computed_tokens=context_len,
+            lora_request=None,
+            prefill_token_ids=token_ids,
+        )
+        for req_id in req_ids
+    ]
+    prefill_output.num_scheduled_tokens = dict.fromkeys(req_ids, query_len)
+    prefill_output.total_num_scheduled_tokens = num_reqs * query_len
+
+    cleanup_output = SchedulerOutput.make_empty()
+    cleanup_output.finished_req_ids = set(req_ids)
+
+    model_runner.kv_connector.set_disabled(True)
+    try:
+        worker_execute_model(prefill_output)
+        worker_sample_tokens(None)
+        worker_execute_model(cleanup_output)
+    finally:
+        model_runner.kv_connector.set_disabled(False)
+    return True
+
+
+@torch.inference_mode()
+def run_kv_cache_sizing_warmup(
+    model_runner: GPUModelRunner,
+    worker_execute_model: Callable[[SchedulerOutput], Any],
+    worker_sample_tokens: Callable[[GrammarOutput | None], Any],
+) -> None:
+    """Run the steps an extensible KV cache measures its transient reserve on.
+    These size rather than compile, so they run without JIT warmup.
+    """
+    if model_runner.is_encoder_decoder:
+        return
+    # Profiling skipped attention; measure it on the largest possible prefill.
+    # Its prompt logprobs materialize full-vocab logits a chunk at a time.
+    scheduler_config = model_runner.scheduler_config
+    max_logprobs = model_runner.model_config.max_logprobs
+    run_mixed_prefill_decode_warmup(
+        model_runner,
+        worker_execute_model,
+        worker_sample_tokens,
+        min(scheduler_config.max_num_batched_tokens, model_runner.max_model_len),
+        prompt_logprobs=max_logprobs if max_logprobs != 0 else None,
+    )
+    if model_runner.model_config.use_mla:
+        from vllm.model_executor.layers.attention.mla_attention import (
+            MLACommonMetadataBuilder,
+        )
+
+        # Fill one chunked-context workspace, so the measured peak covers
+        # up-projecting a prefix cache hit's context.
+        workspace_size = (
+            MLACommonMetadataBuilder.determine_chunked_prefill_workspace_size(
+                model_runner.vllm_config
+            )
+        )
+        run_context_prefill_warmup(
+            model_runner, worker_execute_model, worker_sample_tokens, workspace_size
+        )
+
+
 @torch.inference_mode()
 def warmup_kernels(
     model_runner: GPUModelRunner,
@@ -252,6 +395,11 @@ def warmup_kernels(
             rejection_sampler.enable_adaptive_verification = True
     if adaptive_verification is not None:
         adaptive_verification.warmup_budget_kernel()
+    acceptance_estimator = getattr(
+        model_runner.speculator, "acceptance_estimator", None
+    )
+    if acceptance_estimator is not None:
+        acceptance_estimator.warmup()
 
 
 def _warmup_kernels(
@@ -317,8 +465,9 @@ def _warmup_kernels(
         # Encoder-only models allocate no KV blocks, so this cap doesn't apply.
         num_reqs = min(
             num_reqs,
-            max(1, (model_runner.kv_cache_config.num_blocks - 1) // max_blocks_per_req),
+            max(1, (num_committable_kv_blocks(model_runner) - 1) // max_blocks_per_req),
         )
+        ensure_kv_cache_blocks(model_runner, 1 + num_reqs * max_blocks_per_req)
 
     req_ids = [f"_warmup_{i}_" for i in range(num_reqs)]
 
diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py
index ac3ff3b65865..d98498a4f97b 100644
--- a/vllm/v1/worker/gpu_model_runner.py
+++ b/vllm/v1/worker/gpu_model_runner.py
@@ -3953,18 +3953,8 @@ def execute_model(
                     return make_empty_encoder_model_runner_output(scheduler_output)
 
             if not num_scheduled_tokens:
-                if (
-                    self.parallel_config.distributed_executor_backend
-                    == "external_launcher"
-                    and self.parallel_config.data_parallel_size > 1
-                ):
-                    # this is a corner case when both external launcher
-                    # and DP are enabled, num_scheduled_tokens could be
-                    # 0, and has_unfinished_requests in the outer loop
-                    # returns True. before returning early here we call
-                    # dummy run to ensure coordinate_batch_across_dp
-                    # is called into to avoid out of sync issues.
-                    self._dummy_run(1)
+                # With DP, LLMEngine.step() / DPEngineCoreProc run a dummy
+                # batch for this step so the DP collectives stay aligned.
                 if not has_kv_transfer_group():
                     # Return empty ModelRunnerOutput if no work to do.
                     return EMPTY_MODEL_RUNNER_OUTPUT
diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py
index 147e3b4e2ae7..6ae2160c701d 100644
--- a/vllm/v1/worker/gpu_worker.py
+++ b/vllm/v1/worker/gpu_worker.py
@@ -5,8 +5,8 @@
 import gc
 import os
 import time
-from collections.abc import Callable
-from contextlib import AbstractContextManager, contextmanager, nullcontext
+from collections.abc import Callable, Iterable
+from contextlib import AbstractContextManager, ExitStack, contextmanager, nullcontext
 from datetime import timedelta
 from fnmatch import filter as fnmatch_filter
 from types import NoneType
@@ -14,6 +14,7 @@
 
 import torch
 import torch.nn as nn
+from torch.distributed import ReduceOp
 
 import vllm.envs as envs
 from vllm.config import (
@@ -53,9 +54,11 @@
     get_pcp_group,
     get_pp_group,
     get_tp_group,
+    get_world_group,
     resume_device_comms,
     suspend_device_comms,
 )
+from vllm.distributed.utils import all_reduce_int
 from vllm.distributed.weight_transfer import (
     WeightTransferEngine,
     WeightTransferEngineFactory,
@@ -67,6 +70,7 @@
     share_source_lm_heads,
 )
 from vllm.model_executor.model_loader import get_model_loader
+from vllm.model_executor.model_loader.reload.direct import direct_reload_enabled
 from vllm.model_executor.model_loader.weight_checksum import (
     compute_tensor_digests,
     zero_weights,
@@ -83,6 +87,7 @@
 from vllm.tasks import SupportedTask
 from vllm.tracing import instrument
 from vllm.triton_utils import HAS_TRITON
+from vllm.utils.extensible_tensor import granule_size
 from vllm.utils.flashinfer import warn_if_flashinfer_kernels_missing
 from vllm.utils.gc_utils import freeze_gc_heap, maybe_attach_gc_debug_callback
 from vllm.utils.gpu_sync_debug import enable_gpu_sync_check, with_gpu_sync_check
@@ -94,6 +99,7 @@
     memory_profiling,
 )
 from vllm.utils.torch_utils import set_random_seed, set_torch_threads_for_runtime
+from vllm.utils.vmm_driver import vmm_unavailable_reason
 from vllm.v1.attention.backends.utils import record_kv_cache_layout
 from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
 from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec
@@ -103,6 +109,11 @@
     ModelRunnerOutput,
 )
 from vllm.v1.utils import compute_iteration_details, report_usage_stats
+from vllm.v1.worker.extensible_kv_cache import (
+    extend_kv_cache,
+    kv_cache_lent_for_weight_update,
+    measure_kv_cache_blocks,
+)
 from vllm.v1.worker.sentinel.gpu_worker_sentinel import WorkerSentinel
 from vllm.v1.worker.startup_plan import (
     maybe_apply_startup_plan,
@@ -113,17 +124,25 @@
 
 from ...model_executor.model_loader import TensorizerLoader
 from .gpu.cudagraph_utils import has_compiled_submodule
-from .gpu.warmup import warmup_kernels
+from .gpu.warmup import run_kv_cache_sizing_warmup, warmup_kernels
 from .utils import request_memory
 
 logger = init_logger(__name__)
 
 
+def _min_across_ranks(value: int) -> int:
+    """Minimum of value over the world group: every rank of this engine and of
+    those that run in lockstep with it (MoE data parallel, external launcher)."""
+    return all_reduce_int(get_world_group().cpu_group, value)
+
+
 def _num_workspace_lanes(vllm_config: VllmConfig, use_v2_model_runner: bool) -> int:
     spec_config = vllm_config.speculative_config
     return (
         2
-        if use_v2_model_runner and spec_config is not None and spec_config.use_dspark()
+        if use_v2_model_runner
+        and spec_config is not None
+        and spec_config.uses_separate_workspace_lane
         else 1
     )
 
@@ -204,7 +223,11 @@ def __init__(
         is_driver_worker: bool = False,
     ):
         self.use_v2_model_runner = vllm_config.use_v2_model_runner
-        if self.use_v2_model_runner and not HAS_TRITON:
+        if (
+            self.use_v2_model_runner
+            and not HAS_TRITON
+            and not current_platform.uses_standalone_cpu_model_runner()
+        ):
             raise ValueError("Model Runner V2 requires Triton.")
 
         super().__init__(
@@ -234,7 +257,13 @@ def __init__(
         # is available, since the engine needs a reference to the model.
         self.weight_transfer_engine: WeightTransferEngine | None = None
         self._weight_update_active = False
+        # Update targets (is_draft) whose weights may be incomplete.
+        self._incomplete_weight_updates: set[bool] = set()
         self._weight_update_is_draft = False
+        # Set by the engine while a pause has emptied the KV cache.
+        self._kv_cache_reclaimable = False
+        # Keeps the KV cache lent to the active weight update until it ends.
+        self._weight_update_kv_cache_loan = ExitStack()
 
         # Worker profiler. Enabled and configured through profiler_config.
         # Profiler wrapper is created lazily in profile() when start is called,
@@ -317,6 +346,11 @@ def sleep(self, level: int = 1) -> None:
 
         if has_kv_transfer_group():
             get_kv_transfer_group().unregister_kv_caches()
+        # An extensible KV cache lives outside the allocator pool; keep its
+        # reservation so views and graphs stay valid.
+        extensible_kv_cache = getattr(self.model_runner, "extensible_kv_cache", None)
+        if extensible_kv_cache is not None:
+            extensible_kv_cache.release_physical()
         self.sleep_mode_backend.suspend(level)
         if self.vllm_config.model_config.enable_nccl_comm_suspend:
             suspend_device_comms()
@@ -340,7 +374,12 @@ def sleep(self, level: int = 1) -> None:
 
     def wake_up(self, tags: list[str] | None = None) -> None:
         self.sleep_mode_backend.resume(tags)
-        if has_kv_transfer_group() and (tags is None or "kv_cache" in tags):
+        wake_kv_cache = tags is None or "kv_cache" in tags
+        # Back the extensible KV cache before the connector re-registers it.
+        extensible_kv_cache = getattr(self.model_runner, "extensible_kv_cache", None)
+        if extensible_kv_cache is not None and wake_kv_cache:
+            extensible_kv_cache.recommit()
+        if has_kv_transfer_group() and wake_kv_cache:
             get_kv_transfer_group().restore_kv_caches()
         if self.vllm_config.model_config.enable_nccl_comm_suspend:
             resume_device_comms()
@@ -369,6 +408,10 @@ def wake_up(self, tags: list[str] | None = None) -> None:
     def discard(self, tags: tuple[str, ...]) -> None:
         if has_kv_transfer_group() and "kv_cache" in tags:
             get_kv_transfer_group().unregister_kv_caches()
+        # Outside the allocator pool, like in sleep.
+        extensible_kv_cache = getattr(self.model_runner, "extensible_kv_cache", None)
+        if extensible_kv_cache is not None and "kv_cache" in tags:
+            extensible_kv_cache.release_physical()
         self.sleep_mode_backend.discard(tags)
 
     def checkpoint_prepare(self) -> None:
@@ -540,14 +583,14 @@ def init_device(self):
             torch.accelerator.empty_cache()
 
             # Weights from external model loader process
-            external_weight_memory = get_model_loader(
+            self.external_weight_memory = get_model_loader(
                 self.load_config
             ).get_external_weight_memory(self.vllm_config)
 
             # take current memory snapshot after external weights are ready
             self.init_snapshot = init_snapshot = MemorySnapshot(device=self.device)
             self.requested_memory = request_memory(
-                init_snapshot, self.cache_config, external_weight_memory
+                init_snapshot, self.cache_config, self.external_weight_memory
             )
             logger.debug("worker init memory snapshot: %r", self.init_snapshot)
             logger.debug(
@@ -640,9 +683,32 @@ def update_config(self, overrides: dict[str, Any]) -> None:
             setattr(self.model_runner, config_name, new_config)
             setattr(self, config_name, new_config)
 
-    def reload_weights(self, *args, **kwargs) -> None:
-        with set_current_vllm_config(self.vllm_config):
-            self.model_runner.reload_weights(*args, **kwargs)
+    def reload_weights(
+        self,
+        weights_iterator: Iterable[tuple[str, torch.Tensor]] | None = None,
+        weights_path: str | None = None,
+        is_checkpoint_format: bool = True,
+    ) -> None:
+        # Kernel-format weights are copied into place, needing no room.
+        kv_cache_loan = (
+            kv_cache_lent_for_weight_update(
+                getattr(self.model_runner, "extensible_kv_cache", None),
+                self.model_runner.get_model(),
+                reclaimable=self._kv_cache_reclaimable,
+                fail_if_short=True,
+            )
+            if is_checkpoint_format
+            else nullcontext()
+        )
+        with kv_cache_loan, set_current_vllm_config(self.vllm_config):
+            self.model_runner.reload_weights(
+                weights_iterator, weights_path, is_checkpoint_format
+            )
+
+    def set_kv_cache_reclaimable(self, reclaimable: bool) -> None:
+        """Whether the KV cache holds nothing to keep, so that a weight update
+        may borrow its memory; see `EngineCore._set_kv_cache_reclaimable`."""
+        self._kv_cache_reclaimable = reclaimable
 
     @torch.inference_mode()
     def determine_available_memory(self) -> int:
@@ -763,7 +829,7 @@ def determine_available_memory(self) -> int:
         logger.debug(
             "Initial free memory: %s GiB; Requested memory: %f (util), %s GiB",
             format_gib(self.init_snapshot.free_memory),
-            self.cache_config.gpu_memory_utilization,
+            self.cache_config.resolved_gpu_memory_utilization,
             format_gib(self.requested_memory),
         )
         logger.debug(
@@ -779,7 +845,7 @@ def determine_available_memory(self) -> int:
 
         if cudagraph_memory_estimate > 0:
             total_mem = self.init_snapshot.total_memory
-            current_util = self.cache_config.gpu_memory_utilization
+            current_util = self.cache_config.resolved_gpu_memory_utilization
             cg_util_delta = cudagraph_memory_estimate / total_mem
             if envs.VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS:
                 equiv_util = round(current_util - cg_util_delta, 4)
@@ -887,18 +953,20 @@ def initialize_from_config(self, kv_cache_config: KVCacheConfig) -> None:
         # NOTE(Kuntai): This need to be done before `initialize_kv_cache`,
         # because `initialize_kv_cache` will inject kv cache groups not
         # related to kv cache connector (e.g. kv cache sharing layers).
-        ensure_kv_transfer_initialized(self.vllm_config, kv_cache_config)
-
-        # If the connector provides a custom memory pool (e.g. Mooncake
-        # NVLink/BAREX), use it for KV cache allocation; otherwise fall
-        # back to the standard CuMem pool.
-        mem_pool_context = (
-            get_kv_transfer_group().get_mem_pool_context()
-            if has_kv_transfer_group()
-            else None
-        )
-        if mem_pool_context is None:
-            mem_pool_context = self._maybe_get_memory_pool_context(tag="kv_cache")
+        allocation_context = self._maybe_get_memory_pool_context(tag="kv_cache")
+        if self.cache_config.enable_extensible_kv_cache:
+            # The connector is created in `extend_kv_cache`, once the
+            # memory it registers is final.
+            self._kv_cache_config = kv_cache_config
+        else:
+            ensure_kv_transfer_initialized(self.vllm_config, kv_cache_config)
+            # If the connector provides a custom memory pool (e.g. Mooncake
+            # NVLink/BAREX), use it for KV cache allocation; otherwise fall
+            # back to the standard CuMem pool.
+            if has_kv_transfer_group():
+                mem_pool_context = get_kv_transfer_group().get_mem_pool_context()
+                if mem_pool_context is not None:
+                    allocation_context = mem_pool_context
 
         # Offload KV-init state, except on XPU, whose outermost pool would win.
         runtime_pool = (
@@ -907,10 +975,27 @@ def initialize_from_config(self, kv_cache_config: KVCacheConfig) -> None:
             else self._maybe_get_memory_pool_context(tag="runtime")
         )
         with runtime_pool:
-            self.model_runner.initialize_kv_cache(
-                kv_cache_config,
-                kv_cache_allocation_context=mem_pool_context,
-            )
+            if self.cache_config.enable_extensible_kv_cache:
+                self._v2_model_runner().initialize_kv_cache(
+                    kv_cache_config,
+                    kv_cache_allocation_context=allocation_context,
+                    extensible=True,
+                )
+                # Warmup commits blocks as it goes; keep the profiled activation
+                # peak free for the steps it runs.
+                kv_cache = self._v2_model_runner().extensible_kv_cache
+                assert kv_cache is not None
+                kv_cache.reserved_headroom_bytes = getattr(
+                    self, "peak_activation_memory", 0
+                )
+                # Every rank must warm up the same shapes.
+                kv_cache.warmup_committable_blocks = _min_across_ranks(
+                    kv_cache.committable_blocks()
+                )
+            else:
+                self.model_runner.initialize_kv_cache(
+                    kv_cache_config, kv_cache_allocation_context=allocation_context
+                )
 
         # Build KV-zero metadata outside the CuMem pool so the bookkeeping
         # GPU tensors (seg_addrs, block-id buffers) use the standard PyTorch
@@ -932,6 +1017,35 @@ def get_simple_cpu_offload_num_cpu_blocks(self) -> int | None:
             return None
         return connector.get_aligned_num_cpu_blocks()
 
+    def extensible_kv_cache_unsupported_reason(self) -> str | None:
+        # Every rank must reach the same decision: one that keeps the feature
+        # would wait in the block-count all-reduces for one that dropped it.
+        reason = vmm_unavailable_reason()
+        if all_reduce_int(
+            get_world_group().cpu_group, int(reason is not None), ReduceOp.MAX
+        ):
+            return reason or "another rank cannot use it"
+        return None
+
+    def kv_cache_commit_granule(self) -> int:
+        """Bytes one physical KV cache commit is rounded to on this device."""
+        return granule_size(torch.accelerator.current_device_index())
+
+    def disable_extensible_kv_cache(self) -> None:
+        self.cache_config.enable_extensible_kv_cache = False
+        self.requested_memory = request_memory(
+            self.init_snapshot, self.cache_config, self.external_weight_memory
+        )
+
+    def _v2_model_runner(self) -> "GPUModelRunnerV2":
+        assert self.use_v2_model_runner
+        return cast("GPUModelRunnerV2", self.model_runner)
+
+    def extend_kv_cache(self, num_blocks: int) -> None:
+        """Commit the final size of an extensible KV cache."""
+        extend_kv_cache(self._v2_model_runner(), self._kv_cache_config, num_blocks)
+        self.cache_config.num_gpu_blocks = num_blocks
+
     @instrument(span_name="Warmup (GPU)")
     def compile_or_warm_up_model(self) -> CompilationTimes:
         if not self.use_v2_model_runner:
@@ -941,6 +1055,11 @@ def compile_or_warm_up_model(self) -> CompilationTimes:
             return self._compile_or_warm_up_model()
 
     def _compile_or_warm_up_model(self) -> CompilationTimes:
+        if self.cache_config.enable_extensible_kv_cache:
+            # Track the transient peak of the warmup steps themselves; they
+            # reach shapes (spec-decode drafts, encoder batches) profiling
+            # does not, and `measure_kv_cache_blocks` keeps that peak free.
+            torch.accelerator.reset_peak_memory_stats(self.device)
         warmup_sizes: list[int] = []
 
         if (
@@ -984,12 +1103,31 @@ def _compile_or_warm_up_model(self) -> CompilationTimes:
         # cuda graph capture.
         kernel_warmup(self)
 
-        if self.use_v2_model_runner:  # noqa: SIM102
-            if self.vllm_config.kernel_config.enable_jit_warmup:
+        if self.use_v2_model_runner:
+            # An extensible KV cache is sized from the peak of these full-batch
+            # steps, so they also run when JIT warmup is off (enforce_eager).
+            if (
+                self.vllm_config.kernel_config.enable_jit_warmup
+                or self.cache_config.enable_extensible_kv_cache
+            ):
                 # A workspace resize after capture frees what the graphs point at.
                 warmup_kernels(
                     self.model_runner, self.execute_model, self.sample_tokens
                 )
+            if self.cache_config.enable_extensible_kv_cache:
+                run_kv_cache_sizing_warmup(
+                    self._v2_model_runner(), self.execute_model, self.sample_tokens
+                )
+                # Measure before capture, whose persistent allocations would
+                # hide part of the peak from a later peak - current.
+                kv_cache = self._v2_model_runner().extensible_kv_cache
+                assert kv_cache is not None
+                stats = torch.accelerator.memory_stats(self.device)
+                kv_cache.reserved_headroom_bytes = max(
+                    kv_cache.reserved_headroom_bytes,
+                    stats.get("allocated_bytes.all.peak", 0)
+                    - stats.get("allocated_bytes.all.current", 0),
+                )
 
         cuda_graph_memory_bytes = 0
         if not self.model_config.enforce_eager:
@@ -1049,7 +1187,7 @@ def _compile_or_warm_up_model(self) -> CompilationTimes:
                 f"({format_gib(self.init_snapshot.free_memory)}/"
                 f"{format_gib(self.init_snapshot.total_memory)} GiB) on startup. "
                 f"Desired GPU memory utilization is "
-                f"({self.cache_config.gpu_memory_utilization}, "
+                f"({self.cache_config.resolved_gpu_memory_utilization}, "
                 f"{format_gib(self.requested_memory)} GiB). "
                 f"Actual usage is {format_gib(self.total_consumed)} "
                 f"GiB for consumed memory (weights + non-torch), "
@@ -1092,6 +1230,27 @@ def _compile_or_warm_up_model(self) -> CompilationTimes:
             else:
                 self.model_runner._dummy_sampler_run(hidden_states=last_hidden_states)
 
+        if (
+            self.use_v2_model_runner
+            and get_pp_group().is_last_rank
+            and not self.model_runner.is_pooling_model
+        ):
+            # MRV2 counterpart of the V1 warm-up above: keep the sampler's
+            # max-batch buffers (e.g. the fp32 logits copy) cached after the
+            # empty_cache() in capture_model(), so later steps reuse them
+            # instead of reserving a new, larger block at each new batch size.
+            max_num_reqs = min(
+                self.scheduler_config.max_num_seqs,
+                self.scheduler_config.max_num_batched_tokens,
+            )
+            # Take the hidden states from a dummy forward pass, as V1 does, so they
+            # have the LM head's input size (e.g. OPT projects to a smaller size).
+            _, sample_hidden_states = self.model_runner._dummy_run(
+                num_tokens=max_num_reqs, skip_eplb=True
+            )
+            self.model_runner._dummy_sampler_run(hidden_states=sample_hidden_states)
+            del sample_hidden_states
+
         # Reset the seed to ensure that the random state is not affected by
         # the model initialization and profiling.
         set_random_seed(
@@ -1121,9 +1280,25 @@ def _compile_or_warm_up_model(self) -> CompilationTimes:
         # gate so subsequent `execute_model` / `sample_tokens` calls enforce it.
         enable_gpu_sync_check()
 
+        num_kv_blocks = None
+        if self.cache_config.enable_extensible_kv_cache:
+            kv_cache = self._v2_model_runner().extensible_kv_cache
+            assert kv_cache is not None
+            # Every rank must size the KV cache identically.
+            num_kv_blocks = _min_across_ranks(
+                measure_kv_cache_blocks(
+                    kv_cache,
+                    self.vllm_config,
+                    self.device,
+                    init_free_memory=self.init_snapshot.free_memory,
+                    requested_memory=int(self.requested_memory),
+                )
+            )
+
         return CompilationTimes(
             language_model=self.compilation_config.compilation_time,
             encoder=self.compilation_config.encoder_compilation_time,
+            num_kv_blocks=num_kv_blocks,
         )
 
     def _maybe_activate_jit_monitor(self) -> None:
@@ -1542,6 +1717,9 @@ def shutdown_weight_transfer_engine(self) -> None:
         if self.weight_transfer_engine is not None:
             self.weight_transfer_engine.shutdown()
 
+    def has_incomplete_weight_update(self) -> bool:
+        return bool(self._incomplete_weight_updates)
+
     def start_weight_update(self) -> None:
         """Start a new weight update session.
 
@@ -1583,12 +1761,25 @@ def _start_weight_update(self, is_draft: bool = False) -> None:
                 # Load through the unquantized target head, as without
                 # draft_lm_head_quantization; finish re-derives the copies.
                 share_source_lm_heads(draft_model)
+            # Starting may already move weights; only a successful finish
+            # restores them.
+            self._incomplete_weight_updates.add(is_draft)
             self.weight_transfer_engine.start_weight_update()
         except BaseException:
             self.weight_transfer_engine.reset_weight_update_target()
             if draft_model is not None:
                 restore_quantized_lm_heads(draft_model, refresh=False)
             raise
+        # Engines may receive kernel-format weights, needing no room; so warn
+        # rather than fail when short.
+        self._weight_update_kv_cache_loan.enter_context(
+            kv_cache_lent_for_weight_update(
+                getattr(self.model_runner, "extensible_kv_cache", None),
+                draft_model or self.model_runner.get_model(),
+                reclaimable=self._kv_cache_reclaimable,
+                fail_if_short=False,
+            )
+        )
         self._weight_update_active = True
         self._weight_update_is_draft = is_draft
 
@@ -1625,14 +1816,19 @@ def update_weights(self, update_info: dict | list[dict]) -> None:
                     local_update_info = update_info
                 self.weight_transfer_engine.update_weights(local_update_info)
             except BaseException:
-                self._weight_update_active = False
-                self.weight_transfer_engine.reset_weight_update_target()
-                if self._weight_update_is_draft and (
-                    (draft_model := self.get_draft_model()) is not None
-                ):
-                    restore_quantized_lm_heads(draft_model, refresh=False)
+                self._end_weight_update(refresh=False)
                 raise
 
+    def _end_weight_update(self, refresh: bool) -> None:
+        assert self.weight_transfer_engine is not None
+        self._weight_update_active = False
+        self._weight_update_kv_cache_loan.close()
+        self.weight_transfer_engine.reset_weight_update_target()
+        if self._weight_update_is_draft and (
+            (draft_model := self.get_draft_model()) is not None
+        ):
+            restore_quantized_lm_heads(draft_model, refresh=refresh)
+
     def finish_weight_update(self) -> None:
         """Finish the current weight update session."""
         self._check_weight_transfer_engine()
@@ -1644,13 +1840,18 @@ def finish_weight_update(self) -> None:
             )
 
         with set_current_vllm_config(self.vllm_config):
-            self.weight_transfer_engine.finish_weight_update()
-            self.weight_transfer_engine.reset_weight_update_target()
-            self._weight_update_active = False
-            if self._weight_update_is_draft and (
-                (draft_model := self.get_draft_model()) is not None
-            ):
-                restore_quantized_lm_heads(draft_model)
+            try:
+                self.weight_transfer_engine.finish_weight_update()
+            except BaseException:
+                if direct_reload_enabled():
+                    # a failed direct finish cannot be retried: end the
+                    # session so that a new full update can start
+                    self._end_weight_update(refresh=False)
+                raise
+            finally:
+                self._weight_update_kv_cache_loan.close()
+            self._end_weight_update(refresh=True)
+            self._incomplete_weight_updates.discard(self._weight_update_is_draft)
 
         # Weight transfer bypasses GPUModelRunner.reload_weights().
         if not self._weight_update_is_draft:
diff --git a/vllm/v1/worker/sentinel/eplb_redistribute.py b/vllm/v1/worker/sentinel/eplb_redistribute.py
new file mode 100644
index 000000000000..a74e699ee609
--- /dev/null
+++ b/vllm/v1/worker/sentinel/eplb_redistribute.py
@@ -0,0 +1,288 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""EPLB fault-tolerance helpers: expert redistribution and weight reload
+after peer death.
+
+All redistribution operations are deterministic with stable iteration
+order, so all surviving ranks running the same function with the same
+args produce bit-identical results — no cross-rank communication during
+recovery.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Generator
+from inspect import signature
+from itertools import zip_longest
+from typing import TYPE_CHECKING
+
+import torch
+
+from vllm.config import VllmConfig
+from vllm.logger import init_logger
+from vllm.v1.worker.sentinel.expert_reload import _placement_weight_loader
+
+if TYPE_CHECKING:
+    from vllm.v1.worker.gpu_model_runner import GPUModelRunner
+
+logger = init_logger(__name__)
+
+
+def compute_dead_ep_ranks(
+    dead_dp_ranks: set[int] | list[int],
+    tp_size: int,
+    pcp_size: int = 1,
+) -> set[int]:
+    """Expand dead DP ranks to the corresponding dead EP ranks.
+
+    Each DP rank owns a contiguous block of tp_size * pcp_size EP ranks
+    """
+    block = tp_size * pcp_size
+    return {
+        dp_rank * block + offset for dp_rank in dead_dp_ranks for offset in range(block)
+    }
+
+
+def mark_dead_expert_slots_inplace(
+    physical_to_logical_map: torch.Tensor,
+    dead_ep_ranks: set[int],
+    num_local_experts: int,
+) -> None:
+    """Mark dead EP ranks' physical expert slots as -1 in place.
+
+    physical_to_logical_map is [num_moe_layers, num_physical]; each EP
+    rank owns num_local_experts consecutive physical slots starting at
+    rank * num_local_experts. Shape stays constant (no topology change).
+    """
+    if not dead_ep_ranks:
+        return
+    slots_by_rank = physical_to_logical_map.view(
+        physical_to_logical_map.shape[0], -1, num_local_experts
+    )
+    slots_by_rank[:, sorted(dead_ep_ranks)] = -1
+
+
+def redistribute_expert_placement(
+    physical_to_logical_map: torch.Tensor,
+    num_logical: int,
+    num_local_experts: int,
+) -> set[tuple[int, int]]:
+    """Steals slots from over-replicated experts (donors with >1 replica)
+    to recover missing logical experts, interleaving across EP ranks
+    for balanced weight loading. Modifies physical_to_logical_map in place.
+
+    Returns:
+        Set of (moe_layer_idx, logical_id) pairs needing weight reload.
+
+    Raises:
+        RuntimeError: If not enough slots to cover missing experts.
+
+    """
+    num_layers = physical_to_logical_map.shape[0]
+    all_logical = set(range(num_logical))
+    reassignments: set[tuple[int, int]] = set()
+
+    for layer_idx in range(num_layers):
+        layer_p2l = physical_to_logical_map[layer_idx].cpu().tolist()
+
+        # Replica count per logical expert across the surviving slots.
+        replica_count: dict[int, int] = {}
+        for logical_id in layer_p2l:
+            if logical_id >= 0:
+                replica_count[logical_id] = replica_count.get(logical_id, 0) + 1
+
+        missing_logical = sorted(all_logical - set(replica_count.keys()))
+        if not missing_logical:
+            continue
+
+        # Spare slots by owning EP rank: an expert with >1 replica can give
+        # up a slot without losing coverage. Values: (replica_count, phys_idx).
+        spare_slots_by_rank: dict[int, list[tuple[int, int]]] = {}
+        for phys_idx, logical_id in enumerate(layer_p2l):
+            if logical_id >= 0 and replica_count.get(logical_id, 0) > 1:
+                ep_rank = phys_idx // num_local_experts
+                spare_slots_by_rank.setdefault(ep_rank, []).append(
+                    (replica_count[logical_id], phys_idx)
+                )
+        for bucket in spare_slots_by_rank.values():
+            bucket.sort(reverse=True)
+
+        # Interleave across ranks (round-robin) so reassignments
+        # spread evenly — each rank loads ~equal new weights.
+        interleaved: list[tuple[int, int]] = []
+        for items in zip_longest(
+            *[spare_slots_by_rank[r] for r in sorted(spare_slots_by_rank)]
+        ):
+            for item in items:
+                if item is not None:
+                    interleaved.append(item)
+        slot_iter = iter(interleaved)
+
+        for logical_id in missing_logical:
+            while True:
+                candidate = next(slot_iter, None)
+                if candidate is None:
+                    raise RuntimeError(
+                        f"Layer {layer_idx}: no redundant slot for "
+                        f"missing expert {logical_id}. "
+                        f"EPLB redundancy insufficient."
+                    )
+                _, spare_slot = candidate
+                prev_logical = layer_p2l[spare_slot]
+                # Re-check: an earlier reuse may have left prev_logical
+                # with a single replica, making this slot no longer spare.
+                if replica_count.get(prev_logical, 0) > 1:
+                    layer_p2l[spare_slot] = logical_id
+                    replica_count[prev_logical] -= 1
+                    reassignments.add((layer_idx, logical_id))
+                    break
+
+        physical_to_logical_map[layer_idx].copy_(
+            torch.tensor(layer_p2l, dtype=physical_to_logical_map.dtype)
+        )
+
+    return reassignments
+
+
+def reload_experts_from_disk(
+    model: torch.nn.Module,
+    vllm_config: VllmConfig,
+    reload_set: set[tuple[int, int]],
+    physical_to_logical_map: torch.Tensor,
+) -> None:
+    """Reload complete MoE layers affected by expert reassignment from disk.
+
+    Args:
+        model: The model whose FusedMoE layers receive the reloaded weights.
+        vllm_config: Engine config, used for the model loader and model config.
+        reload_set: {(moe_layer_idx, logical_expert_id), ...} from
+            redistribute_expert_placement.
+        physical_to_logical_map: Recovered placement, including all replicas.
+
+    """
+    if not reload_set:
+        return
+
+    from vllm.model_executor.model_loader import get_model_loader
+    from vllm.model_executor.model_loader.reload import (
+        finalize_layerwise_reload,
+        initialize_layerwise_reload,
+    )
+    from vllm.model_executor.model_loader.reload.utils import get_layer_tensors
+
+    moe_layers = list(model.moe_layers)
+    prefix_renames = getattr(
+        getattr(model, "hf_to_vllm_mapper", None), "orig_to_new_prefix", {}
+    )
+
+    def ckpt_prefix(layer_idx: int) -> str:
+        prefix = moe_layers[layer_idx].layer_name + "."
+        for orig, new in prefix_renames.items():
+            if new and prefix.startswith(new):
+                return orig + prefix[len(new) :]
+        return prefix
+
+    # Runtime-format expert reload: processing can merge scales or repack weights.
+    # Reload every expert in an affected layer so finalization never consumes
+    # uninitialized checkpoint storage for unchanged local experts.
+    layer_indices = sorted({i for i, _ in reload_set})
+    prefixes = tuple(ckpt_prefix(i) for i in layer_indices)
+    layers = [
+        getattr(moe_layers[i], "routed_experts", moe_layers[i]) for i in layer_indices
+    ]
+    loader = get_model_loader(vllm_config.load_config)
+    all_weights = loader.get_all_weights(vllm_config.model_config, model)
+    matched: set[int] = set()
+    rename_checkpoint_weight = getattr(
+        model, "_rename_weight_for_modelopt_checkpoint", lambda name: name
+    )
+
+    def filtered_iter() -> Generator[tuple[str, torch.Tensor], None, None]:
+        for name, tensor in all_weights:
+            # Honor model-specific checkpoint renaming before filtering, while
+            # passing the original name to load_weights for its normal conversion.
+            renamed = rename_checkpoint_weight(name)
+            for layer_idx, prefix in zip(layer_indices, prefixes):
+                runtime_prefix = moe_layers[layer_idx].layer_name + "."
+                if name.startswith(prefix) or renamed.startswith(runtime_prefix):
+                    matched.add(layer_idx)
+                    yield name, tensor
+                    break
+
+    # Reload recovered placement: remap before the layerwise loader buffers its
+    # arguments, so deferred processing replays physical IDs against runtime maps.
+    for layer_idx, layer in zip(layer_indices, layers):
+        if getattr(layer, "quant_config", None) is not None and (
+            layer.quant_config.get_name() == "gpt_oss_mxfp4"
+        ):
+            raise NotImplementedError(
+                "Scale-down reload of fused GPT-OSS MXFP4 weights"
+            )
+        initialize_layerwise_reload(layer, defer_processing=True)
+        quant_method = layer.quant_method
+        quant_method = (
+            getattr(quant_method, "requantization_source", None) or quant_method
+        )
+        local_ids: dict[int, list[int]] = {}
+        for physical_id, logical_id in enumerate(
+            physical_to_logical_map[layer_idx].cpu().tolist()
+        ):
+            if (
+                logical_id >= 0
+                and layer._map_global_expert_id_to_local_expert_id(physical_id) >= 0
+            ):
+                local_ids.setdefault(logical_id, []).append(physical_id)
+        for param in get_layer_tensors(layer).values():
+            weight_loader = getattr(param, "weight_loader", None)
+            if (
+                weight_loader is not None
+                and "expert_id" in signature(weight_loader).parameters
+            ):
+                param.weight_loader = _placement_weight_loader(
+                    weight_loader,
+                    local_ids,
+                    layer.moe_config.num_logical_experts,
+                    physical_to_logical_map.shape[1],
+                    use_global_sf=getattr(quant_method, "use_global_sf", False),
+                )
+
+    logger.info("[FT] Reloading %d (layer, expert) pair(s) from disk.", len(reload_set))
+
+    loaded = model.load_weights(filtered_iter())
+
+    unmatched = set(layer_indices) - matched
+    if unmatched:
+        raise RuntimeError(
+            f"[FT] MoE layers {sorted(unmatched)} had no matching checkpoint weights."
+        )
+    for layer in layers:
+        finalize_layerwise_reload(layer, vllm_config.model_config)
+
+    logger.info(
+        "[FT] Expert weight reload complete: loaded=%d tensors.",
+        len(loaded) if loaded else 0,
+    )
+
+
+def reset_eplb_async_state(model_runner: GPUModelRunner) -> None:
+    """Clear stale EPLB async state after fault or scale-down."""
+    eplb_state = getattr(model_runner, "eplb_state", None)
+    if eplb_state is None:
+        return
+
+    # Increment the epoch first, so an async worker that publishes a result
+    # after this point sees the change and discards it.
+    eplb_state.ft_reset_epoch += 1
+    for ms in eplb_state.model_states.values():
+        ms.rebalanced = False
+        result = ms.pending_result
+        ms.pending_result = None
+        if result is not None:
+            # Wake an async worker waiting on consumed_event; the in-flight
+            # layer result is discarded, not applied.
+            result.consumed_event.record()
+        ms.expert_load_pass.zero_()
+        ms.expert_load_window.zero_()
+
+    eplb_state.expert_rearrangement_step = 0
+    eplb_state.expert_load_window_step = 0
diff --git a/vllm/v1/worker/sentinel/expert_reload.py b/vllm/v1/worker/sentinel/expert_reload.py
new file mode 100644
index 000000000000..89b37f15c5e3
--- /dev/null
+++ b/vllm/v1/worker/sentinel/expert_reload.py
@@ -0,0 +1,40 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from collections.abc import Callable
+from functools import wraps
+from inspect import signature
+
+
+def _placement_weight_loader(
+    loader: Callable,
+    local_physical_ids: dict[int, list[int]],
+    num_logical: int,
+    num_physical: int,
+    use_global_sf: bool = False,
+) -> Callable:
+    """Translate initial checkpoint IDs before buffering layerwise reload calls."""
+    loader_signature = signature(loader)
+
+    @wraps(loader)
+    def online_process_loader(*args, **kwargs):
+        bound = loader_signature.bind(*args, **kwargs)
+        bound.apply_defaults()
+        source = bound.arguments["expert_id"]
+        # Preserve global activation scales: NVFP4 reduces all experts' input
+        # scales, including experts that are not hosted on this rank.
+        if use_global_sf and "input_scale" in bound.arguments.get("weight_name", ""):
+            return loader(*args, **kwargs)
+        if source >= num_physical:
+            return loader(*args, **kwargs)
+        # Reload every local replica: checkpoint mappings enumerate the initial
+        # placement, so dispatch its first copy to every recovered physical slot.
+        if source >= num_logical:
+            return False if bound.arguments.get("return_success") else None
+        success = False
+        for physical_id in local_physical_ids.get(source, []):
+            bound.arguments["expert_id"] = physical_id
+            success = bool(loader(*bound.args, **bound.kwargs)) or success
+        return success if bound.arguments.get("return_success") else None
+
+    return online_process_loader
diff --git a/vllm/v1/worker/sentinel/gpu_worker_sentinel.py b/vllm/v1/worker/sentinel/gpu_worker_sentinel.py
index 9d06a4099090..fd454486056c 100644
--- a/vllm/v1/worker/sentinel/gpu_worker_sentinel.py
+++ b/vllm/v1/worker/sentinel/gpu_worker_sentinel.py
@@ -1,21 +1,44 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from datetime import timedelta
 from typing import TYPE_CHECKING, cast
 
 import torch
+from torch.distributed import TCPStore
 
 from vllm.config import set_current_vllm_config
 from vllm.distributed import (
     get_dp_group,
+    get_ep_group,
+    get_eplb_group,
+    get_pp_group,
+    get_tp_group,
     stateless_destroy_torch_distributed_process_group,
     stateless_init_torch_distributed_process_group,
 )
+from vllm.distributed.eplb.eplb_state import _commit_eplb_maps
+from vllm.distributed.utils import (
+    allocate_group_ports,
+    enter_steady_state,
+    fetch_group_ports,
+    get_cached_tcp_store_client,
+    set_gloo_backend_timeout,
+)
 from vllm.logger import init_logger
 from vllm.model_executor.layers.fused_moe.all2all_utils import get_ep_all2all_manager
+from vllm.utils.network_utils import get_ip
 from vllm.v1.fault_tolerance.utils import FaultToleranceRequest
 from vllm.v1.serial_utils import run_method
+from vllm.v1.worker.sentinel.eplb_redistribute import (
+    compute_dead_ep_ranks,
+    mark_dead_expert_slots_inplace,
+    redistribute_expert_placement,
+    reload_experts_from_disk,
+    reset_eplb_async_state,
+)
 
 if TYPE_CHECKING:
+    from vllm.distributed.parallel_state import GroupCoordinator
     from vllm.v1.worker.gpu.model_runner import GPUModelRunner as GPUModelRunnerV2
     from vllm.v1.worker.gpu_worker import Worker
 
@@ -26,6 +49,38 @@
 FT_BACKEND_SET = frozenset({"deepep_low_latency", "nixl_ep"})
 
 
+def _reinit_cpu_group(
+    group: "GroupCoordinator",
+    store: TCPStore,
+    key: str,
+    rank: int,
+    size: int,
+    timeout: timedelta | None = None,
+) -> None:
+    """Rebuild a group's Gloo cpu_group over a port coordinated via ``store``"""
+    # Rendezvous recovery groups on their actual leader hosts: DP replicas and
+    # pipeline stages can span nodes independently of the deployment master.
+    if rank == 0:
+        store.set(f"{key}_host", get_ip())
+    host = store.get(f"{key}_host").decode()
+    if rank == 0:
+        ports, socks = allocate_group_ports(store, key, host, 1)
+        port, listen_socket = ports[0], socks[0]
+    else:
+        port, listen_socket = fetch_group_ports(store, key, 1)[0], None
+    new_group = stateless_init_torch_distributed_process_group(
+        host,
+        port,
+        rank,
+        size,
+        backend="gloo",
+        listen_socket=listen_socket,
+        timeout=timeout,
+    )
+    stateless_destroy_torch_distributed_process_group(group.cpu_group)
+    group.cpu_group = new_group
+
+
 class WorkerSentinel:
     """Holds FT state for a single worker (mask tensors, DP config).
 
@@ -34,57 +89,213 @@ class WorkerSentinel:
 
     def __init__(self, worker: "Worker"):
         self.worker = worker
-        self.dp_rank = worker.parallel_config.data_parallel_rank
-        self.dp_size = worker.parallel_config.data_parallel_size
-        self.data_parallel_master_ip = worker.parallel_config.data_parallel_master_ip
         all2all_backend = worker.parallel_config.all2all_backend
         if all2all_backend not in FT_BACKEND_SET:
             raise ValueError(
                 f"Fault tolerance requires an FT-capable all2all backend "
                 f"(one of {sorted(FT_BACKEND_SET)}), but got '{all2all_backend}'."
             )
+        # Recover registered draft experts: until draft reload is supported,
+        # reject MoE drafts before serving rather than leaving stale EP mappings.
+        spec_config = worker.vllm_config.speculative_config
+        if (
+            spec_config is not None
+            and spec_config.draft_model_config is not None
+            and spec_config.draft_model_config.is_moe
+        ):
+            raise ValueError("Fault tolerance does not support MoE draft models yet.")
+        if worker.parallel_config.prefill_context_parallel_size > 1:
+            raise ValueError("Fault tolerance does not support PCP yet.")
+        if not worker.vllm_config.use_v2_model_runner:
+            raise ValueError(
+                "Fault tolerance requires the v2 model runner "
+                "(VLLM_USE_V2_MODEL_RUNNER=1)."
+            )
 
     def handle_command(self, ft_request: FaultToleranceRequest):
         """Dispatch an FT command by instruction name."""
         with set_current_vllm_config(self.worker.vllm_config):
             return run_method(self, ft_request.instruction, (ft_request,), {})
 
+    def activate_cpu_group_timeouts(self, ft_request: FaultToleranceRequest):
+        enter_steady_state()
+        parallel_config = self.worker.parallel_config
+        timeout_seconds = parallel_config.cpu_distributed_timeout_seconds
+        if timeout_seconds is None:
+            return
+        timeout = timedelta(seconds=timeout_seconds)
+        groups = [get_dp_group(), get_tp_group(), get_pp_group()]
+        if parallel_config.enable_expert_parallel:
+            groups.append(get_ep_group())
+        if parallel_config.enable_eplb:
+            groups.append(get_eplb_group())
+        for group in groups:
+            if group.world_size > 1:
+                set_gloo_backend_timeout(group.cpu_group, timeout)
+                # GroupCoordinator caches cpu_timeout at init (None under FT
+                # startup); its cpu_all_reduce_ forwards it, so update both.
+                group.cpu_timeout = timeout
+
     def retry(self, ft_request: FaultToleranceRequest):
         torch.accelerator.synchronize()
         params = ft_request.params
+        # After a scale_down the DP master may be re-elected and the rebuilt
+        # gloo group is dense over the alive slots.
+        master_ip = params["dp_master_ip"]
+        dp_group_rank = params["dp_group_rank"]
+        dp_group_size = params["dp_group_size"]
+        round_key = params["recovery_round"]
+        store = get_cached_tcp_store_client(master_ip, params["recovery_store_port"])
         self._clean_worker_state()
-        if self.dp_size > 1:
+        reset_eplb_async_state(self.worker.model_runner)
+        if self.worker.parallel_config.data_parallel_size > 1:
             get_ep_all2all_manager().clean_buffers()
-            old_cpu_group = get_dp_group().cpu_group
-            stateless_destroy_torch_distributed_process_group(old_cpu_group)
             world_size = self.worker.parallel_config.world_size
-            port = params["new_stateless_dp_group_ports"][self.worker.rank % world_size]
-            get_dp_group().cpu_group = stateless_init_torch_distributed_process_group(
-                self.data_parallel_master_ip,
-                port,
-                self.dp_rank,
-                self.dp_size,
-                backend="gloo",
+            dp_index = self.worker.rank % world_size
+            _reinit_cpu_group(
+                get_dp_group(),
+                store,
+                f"ft_dp_port_{round_key}_{dp_index}",
+                dp_group_rank,
+                dp_group_size,
+                timeout=self.worker.parallel_config.cpu_distributed_timeout,
+            )
+            get_dp_group().dead_dp_ranks = set(params["dead_dp_ranks"])
+
+            # clean_buffers wiped the mask; replay masks for the cumulative dead set.
+            tp_size = self.worker.parallel_config.tensor_parallel_size
+            pcp_size = self.worker.parallel_config.prefill_context_parallel_size
+            all2all_manager = get_ep_all2all_manager()
+            for ep_rank in compute_dead_ep_ranks(
+                params["dead_dp_ranks"], tp_size, pcp_size
+            ):
+                all2all_manager.update_mask(ep_rank, masked=True)
+
+            if (
+                self.worker.parallel_config.enable_eplb
+                and not self.worker.model_runner.eep_eplb_suppressed
+            ):
+                self._reinit_eplb_groups(store, round_key)
+        if self.worker.parallel_config.tensor_parallel_size > 1:
+            # The per-step TP barrier (dp_utils) leaves the group in a
+            # timed-out state on fault; rebuild it for a clean slate.
+            tp_group = get_tp_group()
+            # Pipeline-stage recovery ports: each stage owns a distinct TP group.
+            pp_rank = get_pp_group().rank_in_group
+            _reinit_cpu_group(
+                tp_group,
+                store,
+                f"ft_tp_port_{round_key}_{dp_group_rank}_{pp_rank}",
+                tp_group.rank_in_group,
+                tp_group.world_size,
                 timeout=self.worker.parallel_config.cpu_distributed_timeout,
             )
 
+    def scale_down(self, ft_request: FaultToleranceRequest):
+        model_runner = self.worker.model_runner
+        eplb_config = self.worker.parallel_config.eplb_config
+        if model_runner.eplb_state is None or eplb_config.num_redundant_experts <= 0:
+            raise ValueError(
+                "[FT] scale_down requires EPLB with num_redundant_experts > 0 "
+                "to re-host the dead rank's experts."
+            )
+        tp_size = self.worker.parallel_config.tensor_parallel_size
+        pcp_size = self.worker.parallel_config.prefill_context_parallel_size
+        dead_dp_ranks = ft_request.params["dead_dp_ranks"]
+        dead_ep_ranks = compute_dead_ep_ranks(dead_dp_ranks, tp_size, pcp_size)
+        # Suppress EPLB async rebalancing before retry so it skips the EPLB
+        # group reinit; placement is fixed by the redistribution below.
+        model_runner.eep_eplb_suppressed = True
+        self.retry(ft_request)
+
+        self._redistribute_experts(dead_ep_ranks)
+
+        logger.info(
+            "[FT] Worker scale_down complete: dp_group_size=%d, "
+            "dp_group_rank=%d, dead_ep_ranks=%s, eplb_suppressed=True",
+            ft_request.params["dp_group_size"],
+            ft_request.params["dp_group_rank"],
+            sorted(dead_ep_ranks),
+        )
+
+    def query_mask(self, ft_request: FaultToleranceRequest) -> dict:
+        """Query the mask on a side stream."""
+        with torch.cuda.stream(torch.cuda.Stream()):
+            mask = get_ep_all2all_manager().query_active_mask()
+            return {"mask": mask.tolist()}
+
+    def _reinit_eplb_groups(self, store: TCPStore, round_key: str) -> None:
+        """Reinit the EP/EPLB Gloo groups and refresh the EPLB
+        communicator's cpu_group reference."""
+        # Pipeline-stage recovery ports: EP/EPLB leaders must not share keys.
+        pp_rank = get_pp_group().rank_in_group
+        for key_suffix, get_group in [
+            ("ep", get_ep_group),
+            ("eplb", get_eplb_group),
+        ]:
+            group = get_group()
+            _reinit_cpu_group(
+                group,
+                store,
+                f"ft_{key_suffix}_port_{round_key}_{pp_rank}",
+                group.rank_in_group,
+                group.world_size,
+                timeout=self.worker.parallel_config.cpu_distributed_timeout,
+            )
+            logger.info("[FT] Reinited %s Gloo group", key_suffix)
+
+        eplb_state = self.worker.model_runner.eplb_state
+        assert eplb_state is not None
+        eplb_group = get_eplb_group()
+        for ms in eplb_state.model_states.values():
+            if hasattr(ms.communicator, "_cpu_group"):
+                ms.communicator._cpu_group = eplb_group.cpu_group
+
+    def _eplb_model_state(self):
+        model_runner = self.worker.model_runner
+        assert model_runner.eplb_state is not None
+        return model_runner.eplb_state.model_states[
+            model_runner.model_config.compute_hash()
+        ]
+
+    def _redistribute_experts(self, dead_ep_ranks: set[int]) -> None:
+        """One-shot expert redistribution after scale-down."""
+        model_runner = self.worker.model_runner
+        eplb_model_state = self._eplb_model_state()
+
+        p2l = eplb_model_state.physical_to_logical_map
+        num_logical = eplb_model_state.logical_replica_count.shape[1]
+        ep_world_size = get_ep_group().world_size
+        num_local_experts = p2l.shape[1] // ep_world_size
+
+        mark_dead_expert_slots_inplace(p2l, dead_ep_ranks, num_local_experts)
+        reassignments = redistribute_expert_placement(
+            p2l, num_logical, num_local_experts
+        )
+        # p2l was updated in place; the commit derives l2p/lrc from it.
+        _commit_eplb_maps(eplb_model_state, p2l.cpu())
+        # Preserve physical expert IDs: slot ownership is unchanged by EPLB;
+        # only the logical expert stored in each surviving slot changes.
+
+        if reassignments:
+            reload_experts_from_disk(
+                model_runner.model,
+                self.worker.vllm_config,
+                reassignments,
+                p2l,
+            )
+
+        logger.info(
+            "[FT] Expert redistribution: num_logical=%d, "
+            "ep_world_size=%d, reassignments=%d",
+            num_logical,
+            ep_world_size,
+            len(reassignments),
+        )
+
     def _clean_worker_state(self):
         model_runner = self.worker.model_runner
         model_runner.execute_model_state = None
-        if self.worker.use_v2_model_runner:
-            runner = cast("GPUModelRunnerV2", model_runner)
-            for req_id in list(runner.req_states.req_id_to_index):
-                runner._remove_request(req_id)
-        else:
-            model_runner.kv_connector_output = None
-
-            input_batch = model_runner.input_batch
-            cached_req_ids = list(input_batch.req_id_to_index)
-            for req_id in cached_req_ids:
-                model_runner.requests.pop(req_id, None)
-                model_runner.num_prompt_logprobs.pop(req_id, None)
-                input_batch.remove_request(req_id)
-
-            input_batch.condense()
-            input_batch.refresh_metadata()
-            input_batch.req_prompt_embeds.clear()
+        runner = cast("GPUModelRunnerV2", model_runner)
+        for req_id in list(runner.req_states.req_id_to_index):
+            runner._remove_request(req_id)
diff --git a/vllm/v1/worker/startup_plan.py b/vllm/v1/worker/startup_plan.py
index 2c4942070705..a5812d8bf5f5 100644
--- a/vllm/v1/worker/startup_plan.py
+++ b/vllm/v1/worker/startup_plan.py
@@ -35,7 +35,7 @@
 
 logger = init_logger(__name__)
 
-PLAN_SCHEMA_VERSION = 1
+PLAN_SCHEMA_VERSION = 2
 
 
 def compute_plan_fingerprint(
diff --git a/vllm/v1/worker/utils.py b/vllm/v1/worker/utils.py
index 5a254c481eb0..ec9f85484660 100644
--- a/vllm/v1/worker/utils.py
+++ b/vllm/v1/worker/utils.py
@@ -2,7 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import math
 from collections import defaultdict
-from collections.abc import Iterable, Mapping, Sequence
+from collections.abc import Callable, Iterable, Mapping, Sequence
 from dataclasses import dataclass, field, replace
 from itertools import product as iprod
 from typing import Any
@@ -35,6 +35,7 @@
     KVCacheGroupSpec,
     KVCacheLayout,
     KVCacheSpec,
+    KVCacheTensor,
     MambaSpec,
     UniformTypeKVCacheSpecs,
     create_kv_cache_views,
@@ -566,12 +567,16 @@ def allocate_kv_cache(
     device: torch.device,
     layout: KVCacheLayout,
     kernel_block_sizes: list[int] | None = None,
+    allocate: Callable[[int], torch.Tensor] | None = None,
 ) -> dict[str, torch.Tensor]:
     """Allocate the KV cache and view it as ``[B, H, N, C]`` per layer.
 
     Every KVCacheTensor places its layers in the same backing allocation: layer ``l`` of
     block ``b`` starts at ``offset + l * layer_stride + b * block_stride``. Cache
     groups overlay each other, so tensors may address the same bytes.
+
+    allocate replaces the default zeroed backing allocation: it takes the
+    size in bytes and returns a one-byte-per-element tensor.
     """
     if not kv_cache_config.kv_cache_tensors:
         return {}
@@ -579,54 +584,74 @@ def allocate_kv_cache(
     sizes = {tensor.size for tensor in kv_cache_config.kv_cache_tensors}
     assert len(sizes) == 1, "KV cache tensors must share one backing allocation."
     raw_size = sizes.pop()
-    # wvSplitKrc's process-lifetime static workspaces (csrc/rocm/skinny_gemms.cu)
-    # are created lazily on the first qualifying GEMM. Force that now, before
-    # the giant backing allocation below: if one landed in this segment's
-    # rounding tail it would pin the whole segment at engine shutdown.
-    if current_platform.is_rocm():
-        warmup_rocm_skinny_gemm_workspaces(device)
-        # Pad to the page granularity MoRIIO needs to register the shared
-        # backing as a single RDMA memory region. Other platforms keep the
-        # exact-size allocation (see #53974), so anything reading
-        # storage.nbytes() has to tolerate the tail on ROCm alone.
-        page_size = 4096
-        buf_size = ((raw_size + page_size - 1) // page_size) * page_size
+    if allocate is not None:
+        buf = allocate(raw_size)
     else:
-        buf_size = raw_size
-    buf = torch.zeros(buf_size, dtype=torch.int8, device=device)
+        buf = _allocate_kv_cache_buffer(raw_size, device)
 
     kv_caches: dict[str, torch.Tensor] = {}
     for tensor in kv_cache_config.kv_cache_tensors:
-        layer_name = tensor.layers[0]
-        group_id, group = next(
-            (group_id, group)
-            for group_id, group in enumerate(kv_cache_config.kv_cache_groups)
-            if layer_name in group.layer_names
-        )
-        spec = group.kv_cache_spec
-        if isinstance(spec, UniformTypeKVCacheSpecs):
-            spec = spec.kv_cache_specs[layer_name]
-
+        group_id, spec = layer_spec_for_tensor(kv_cache_config, tensor)
         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]
         views = create_kv_cache_views(
             buf,
             spec,
-            num_blocks,
+            kv_cache_config.num_blocks_of(tensor),
             layout,
             tensor,
-            kernel_block_size=kernel_block_size,
+            kernel_block_size=layer_kernel_block_size(group_id, kernel_block_sizes),
         )
         kv_caches.update(zip(tensor.layers, views))
     return kv_caches
 
 
+def layer_spec_for_tensor(
+    kv_cache_config: KVCacheConfig, tensor: KVCacheTensor
+) -> tuple[int, KVCacheSpec]:
+    """Cache group index and per-layer spec of the layers tensor places."""
+    layer_name = tensor.layers[0]
+    group_id, group = next(
+        (group_id, group)
+        for group_id, group in enumerate(kv_cache_config.kv_cache_groups)
+        if layer_name in group.layer_names
+    )
+    spec = group.kv_cache_spec
+    if isinstance(spec, UniformTypeKVCacheSpecs):
+        spec = spec.kv_cache_specs[layer_name]
+    return group_id, spec
+
+
+def layer_kernel_block_size(
+    group_id: int, kernel_block_sizes: list[int] | None
+) -> int | None:
+    """Kernel block size to view the group's pages with, if it differs."""
+    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]
+    return kernel_block_size
+
+
+def _allocate_kv_cache_buffer(raw_size: int, device: torch.device) -> torch.Tensor:
+    # wvSplitKrc's process-lifetime static workspaces (csrc/rocm/skinny_gemms.cu)
+    # are created lazily on the first qualifying GEMM. Force that now, before
+    # the giant backing allocation below: if one landed in this segment's
+    # rounding tail it would pin the whole segment at engine shutdown.
+    if current_platform.is_rocm():
+        warmup_rocm_skinny_gemm_workspaces(device)
+        # Pad to the page granularity MoRIIO needs to register the shared
+        # backing as a single RDMA memory region. Other platforms keep the
+        # exact-size allocation (see #53974), so anything reading
+        # storage.nbytes() has to tolerate the tail on ROCm alone.
+        page_size = 4096
+        buf_size = ((raw_size + page_size - 1) // page_size) * page_size
+    else:
+        buf_size = raw_size
+    return torch.zeros(buf_size, dtype=torch.int8, device=device)
+
+
 def prepare_kernel_block_sizes(
     kv_cache_config: KVCacheConfig, attn_groups: list[list[AttentionGroup]]
 ) -> list[int]:
@@ -721,9 +746,8 @@ def request_memory(
     """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(
-        init_snapshot.total_memory * cache_config.gpu_memory_utilization
-    )
+    gpu_memory_utilization = cache_config.resolved_gpu_memory_utilization
+    requested_memory = math.ceil(init_snapshot.total_memory * gpu_memory_utilization)
 
     if external_weight_memory > 0:
         engine_memory = requested_memory - external_weight_memory
@@ -733,18 +757,23 @@ def request_memory(
                 f"({format_gib(external_weight_memory)}/"
                 f"{format_gib(init_snapshot.total_memory)} GiB) exceed the "
                 "desired GPU memory utilization "
-                f"({cache_config.gpu_memory_utilization}, "
+                f"({gpu_memory_utilization}, "
                 f"{format_gib(requested_memory)} GiB). Increase GPU memory "
                 "utilization or reduce GPU memory used by other processes."
             )
-        if init_snapshot.free_memory < engine_memory:
+        # With the extensible KV cache the budget is only a cap; exceeding free
+        # memory over-reserves address space rather than pages.
+        if (
+            init_snapshot.free_memory < engine_memory
+            and not cache_config.enable_extensible_kv_cache
+        ):
             raise ValueError(
                 f"Free memory on device {init_snapshot.device_} "
                 f"({format_gib(init_snapshot.free_memory)}/"
                 f"{format_gib(init_snapshot.total_memory)} GiB) on startup "
                 "is less than the engine's budget after excluding "
                 "external process's weights "
-                f"({cache_config.gpu_memory_utilization}, "
+                f"({gpu_memory_utilization}, "
                 f"{format_gib(engine_memory)} GiB). Decrease GPU memory "
                 "utilization or reduce GPU memory used by other processes."
             )
@@ -758,13 +787,18 @@ def request_memory(
         )
         return engine_memory
 
-    if init_snapshot.free_memory < requested_memory:
+    if (
+        init_snapshot.free_memory < requested_memory
+        and not cache_config.enable_extensible_kv_cache
+    ):
+        # With the extensible KV cache the budget is only a cap; exceeding free
+        # memory over-reserves address space rather than pages.
         raise ValueError(
             f"Free memory on device {init_snapshot.device_} "
             f"({format_gib(init_snapshot.free_memory)}/"
             f"{format_gib(init_snapshot.total_memory)} GiB) on startup "
             f"is less than desired GPU memory utilization "
-            f"({cache_config.gpu_memory_utilization}, "
+            f"({gpu_memory_utilization}, "
             f"{format_gib(requested_memory)} GiB). Decrease GPU memory "
             f"utilization or reduce GPU memory used by other processes."
         )
diff --git a/vllm/v1/worker/worker_base.py b/vllm/v1/worker/worker_base.py
index f5ba02d2e16d..23f6f2889536 100644
--- a/vllm/v1/worker/worker_base.py
+++ b/vllm/v1/worker/worker_base.py
@@ -39,6 +39,9 @@
 class CompilationTimes(NamedTuple):
     language_model: float
     encoder: float
+    num_kv_blocks: int | None = None
+    """KV cache blocks that fit in the memory measured after warmup, reported
+    by workers with an extensible KV cache."""
 
 
 class WorkerBase:
@@ -122,11 +125,25 @@ def compile_or_warm_up_model(self) -> CompilationTimes:
         """Prepare model for execution through compilation/warmup.
 
         Returns:
-            Compilation times (language_model, encoder) in seconds.
+            Compilation times (language_model, encoder) in seconds, and the
+            measured KV cache size for an extensible KV cache.
 
         """
         raise NotImplementedError
 
+    def extensible_kv_cache_unsupported_reason(self) -> str | None:
+        """Why this worker cannot back an extensible KV cache, or None."""
+        return f"{type(self).__name__} does not support the extensible KV cache"
+
+    def disable_extensible_kv_cache(self) -> None:
+        self.vllm_config.cache_config.enable_extensible_kv_cache = False
+
+    def extend_kv_cache(self, num_blocks: int) -> None:
+        """Commit the final size of an extensible KV cache."""
+        raise NotImplementedError(
+            f"{type(self).__name__} does not support the extensible KV cache."
+        )
+
     def check_health(self) -> None:
         """Basic health check (override for device-specific checks)."""
         return
@@ -162,6 +179,11 @@ def supports_draft_weight_updates(self) -> bool:
         """Whether this worker can update its configured speculative model."""
         return False
 
+    def has_incomplete_weight_update(self) -> bool:
+        """Whether a weight update started and has not finished successfully,
+        so the model weights may be incomplete."""
+        return False
+
     def apply_model(self, fn: Callable[[nn.Module], _R]) -> _R:
         """Apply a function on the model inside this worker."""
         return fn(self.get_model())
@@ -193,12 +215,6 @@ def sample_tokens(
         """Should be called immediately after execute_model iff it returned None."""
         raise NotImplementedError
 
-    def get_cache_block_size_bytes(self) -> int:
-        """Return the size of a single cache block, in bytes. Used in
-        speculative decoding.
-        """
-        raise NotImplementedError
-
     def add_lora(self, lora_request: LoRARequest) -> bool:
         raise NotImplementedError
 
@@ -358,6 +374,12 @@ def initialize_from_config(self, kv_cache_configs: list[Any]) -> None:
         with set_current_vllm_config(self.vllm_config):
             self.worker.initialize_from_config(kv_cache_config)  # type: ignore
 
+    def extend_kv_cache(self, num_blocks: int) -> None:
+        # The KV connector created here reads the current config.
+        assert self.vllm_config is not None
+        with set_current_vllm_config(self.vllm_config):
+            self.worker.extend_kv_cache(num_blocks)  # type: ignore
+
     def init_device(self):
         assert self.vllm_config is not None
         with set_current_vllm_config(self.vllm_config):
diff --git a/vllm/v1/worker/workspace.py b/vllm/v1/worker/workspace.py
index a05534281c11..7a625671b261 100644
--- a/vllm/v1/worker/workspace.py
+++ b/vllm/v1/worker/workspace.py
@@ -14,7 +14,7 @@
 
 import vllm.envs as envs
 from vllm.logger import init_logger
-from vllm.utils.math_utils import round_up
+from vllm.utils.math_utils import round_down, round_up
 from vllm.v1.worker.ubatching import dbo_current_ubatch_id
 
 logger = init_logger(__name__)
@@ -28,6 +28,10 @@ def _compute_bytes(shape: tuple[int, ...], dtype: torch.dtype) -> int:
 # Constants
 _MB = 1024**2
 _GiB = 1024**3
+# Every view starts at a multiple of this, the alignment cudaMalloc guarantees,
+# so kernels see the same alignment as on a standalone allocation. Head views
+# also end on it, so a retained span aligned down to it never overlaps them.
+_VIEW_ALIGNMENT = 256
 
 # Global workspace manager instance
 _manager: "WorkspaceManager | None" = None
@@ -73,6 +77,11 @@ def __init__(
         self._persistent_resources: list[dict[Hashable, object]] = [
             {} for _ in self._current_workspaces
         ]
+        # Per slot, (key, start) of the tail span get_retained last handed out;
+        # cleared when a scratch request reaches into it or on reallocation.
+        self._retained: list[tuple[Hashable, int] | None] = [
+            None for _ in self._current_workspaces
+        ]
         self._locked: bool = False
 
     @staticmethod
@@ -191,13 +200,17 @@ def get_simultaneous(
 
         """
         actual_bytes = [_compute_bytes(s, d) for s, d in shapes_and_dtypes]
-        aligned_bytes = [round_up(actual, 256) for actual in actual_bytes]
+        aligned_bytes = [round_up(actual, _VIEW_ALIGNMENT) for actual in actual_bytes]
         total_bytes = sum(aligned_bytes)
 
         # Calculate cumulative offsets using itertools.accumulate
         offsets = list(accumulate([0] + aligned_bytes[:-1]))
 
         current_workspace = self._ensure_workspace_size(total_bytes)
+        workspace_id = self._get_workspace_id()
+        retained = self._retained[workspace_id]
+        if retained is not None and total_bytes > retained[1]:
+            self._retained[workspace_id] = None
 
         return [
             current_workspace[offsets[i] : offsets[i] + actual_bytes[i]]
@@ -206,6 +219,37 @@ def get_simultaneous(
             for i in range(len(shapes_and_dtypes))
         ]
 
+    def get_retained(
+        self, key: Hashable, shape: tuple[int, ...], dtype: torch.dtype
+    ) -> tuple[torch.Tensor, bool]:
+        """Get a scratch view at the end of the workspace that can outlive requests.
+
+        ``get_simultaneous`` hands out views from the start of the workspace, so
+        requests smaller than the workspace leave its end alone. This view spans
+        the last bytes of the workspace, and the returned flag says whether they
+        still hold what the previous ``get_retained`` left there: that call used
+        the same key and size, no ``get_simultaneous`` request since reached into
+        the span, and the workspace was not reallocated in between. What the
+        bytes mean is the caller's to track; only the latest key is remembered.
+
+        Args:
+            key: Identifies the caller; another key's call takes the span over.
+            shape: Shape of the view.
+            dtype: Dtype of the view.
+
+        Returns:
+            The view and whether its contents are intact.
+
+        """
+        num_bytes = _compute_bytes(shape, dtype)
+        workspace_id = self._get_workspace_id()
+        current_workspace = self._ensure_workspace_size(num_bytes)
+        start = round_down(current_workspace.numel() - num_bytes, _VIEW_ALIGNMENT)
+        intact = self._retained[workspace_id] == (key, start)
+        self._retained[workspace_id] = (key, start)
+        view = current_workspace[start : start + num_bytes].view(dtype).reshape(shape)
+        return view, intact
+
     def _ensure_workspace_size(self, required_bytes: int) -> torch.Tensor:
         """Ensure workspace is allocated and large enough, return current workspace.
 
@@ -252,6 +296,7 @@ def get_caller_info() -> str:
             # Resizing all ubatches here would orphan the other ubatch's
             # old tensor when it still holds views into it (DBO leak).
             self._current_workspaces[workspace_id] = None
+            self._retained[workspace_id] = None
             del current_workspace
             # Release the freed segment back to CUDA so the caching
             # allocator can reuse the GPU memory for the larger